././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.8674593 apache_beam-2.74.0/0000755000175100017510000000000015206113160013452 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/MANIFEST.in0000644000175100017510000000161515206113032015211 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # include gen_protos.py include gen_xlang_wrappers.py include README.md include NOTICE include LICENSE include LICENSE.python ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.8674593 apache_beam-2.74.0/PKG-INFO0000644000175100017510000006037715206113160014564 0ustar00runnerrunnerMetadata-Version: 2.4 Name: apache-beam Version: 2.74.0 Summary: Apache Beam SDK for Python Home-page: https://beam.apache.org Download-URL: https://pypi.python.org/pypi/apache-beam Author: Apache Software Foundation Author-email: dev@beam.apache.org License: Apache License, Version 2.0 Keywords: apache beam Classifier: Intended Audience :: End Users/Desktop Classifier: License :: OSI Approved :: Apache Software License Classifier: Operating System :: POSIX :: Linux Classifier: Programming Language :: Python :: 3.10 Classifier: Programming Language :: Python :: 3.11 Classifier: Programming Language :: Python :: 3.12 Classifier: Programming Language :: Python :: 3.13 Classifier: Programming Language :: Python :: 3.14 Classifier: Topic :: Software Development :: Libraries Classifier: Topic :: Software Development :: Libraries :: Python Modules Requires-Python: >=3.10 Description-Content-Type: text/markdown Requires-Dist: cryptography<48.0.0,>=39.0.0 Requires-Dist: envoy-data-plane<2,>=1.0.3; python_version >= "3.11" Requires-Dist: envoy-data-plane<0.3.0; python_version < "3.11" Requires-Dist: fastavro<2,>=0.23.6 Requires-Dist: fasteners<1.0,>=0.3 Requires-Dist: grpcio!=1.48.0,!=1.59.*,!=1.60.*,!=1.61.*,!=1.62.0,!=1.62.1,!=1.66.*,!=1.67.*,!=1.68.*,!=1.69.*,!=1.70.*,<2,>=1.33.1 Requires-Dist: httplib2<0.32.0,>=0.8 Requires-Dist: jsonpickle<4.0.0,>=3.0.0 Requires-Dist: numpy<2.5.0,>=1.14.3 Requires-Dist: objsize<0.8.0,>=0.6.1 Requires-Dist: packaging>=22.0 Requires-Dist: pillow<13,>=12.1.1 Requires-Dist: pymongo<5.0.0,>=3.8.0 Requires-Dist: proto-plus<2,>=1.7.1 Requires-Dist: protobuf!=4.0.*,!=4.21.*,!=4.22.0,!=4.23.*,!=4.24.*,<7.0.0.dev0,>=3.20.3 Requires-Dist: python-dateutil<3,>=2.8.0 Requires-Dist: pytz>=2018.3 Requires-Dist: requests<3.0.0,>=2.32.4 Requires-Dist: sortedcontainers>=2.4.0 Requires-Dist: typing-extensions>=3.7.0 Requires-Dist: zstandard<1,>=0.18.0 Requires-Dist: pyyaml<7.0.0,>=3.12 Requires-Dist: beartype<0.23.0,>=0.21.0 Requires-Dist: pyarrow<24.0.0,>=6.0.1 Requires-Dist: pyarrow-hotfix<1 Provides-Extra: dev Requires-Dist: isort==7.0.0; extra == "dev" Requires-Dist: pyrefly==0.54.0; extra == "dev" Requires-Dist: ruff==0.15.7; extra == "dev" Requires-Dist: yapf==0.43.0; extra == "dev" Provides-Extra: dill Requires-Dist: dill<0.3.2,>=0.3.1.1; extra == "dill" Provides-Extra: docs Requires-Dist: jinja2<3.2,>=3.0; extra == "docs" Requires-Dist: Sphinx<8.0,>=7.0.0; extra == "docs" Requires-Dist: docstring-parser<1.0,>=0.15; extra == "docs" Requires-Dist: docutils>=0.18.1; extra == "docs" Requires-Dist: markdown; extra == "docs" Requires-Dist: pandas<2.3.0; extra == "docs" Requires-Dist: openai; extra == "docs" Requires-Dist: virtualenv-clone<1.0,>=0.5; extra == "docs" Provides-Extra: test Requires-Dist: cloud-sql-python-connector[pg8000]<2.0.0,>=1.0.0; extra == "test" Requires-Dist: docstring-parser<1.0,>=0.15; extra == "test" Requires-Dist: freezegun>=0.3.12; extra == "test" Requires-Dist: jinja2<3.2,>=3.0; extra == "test" Requires-Dist: joblib>=1.0.1; extra == "test" Requires-Dist: mock<6.0.0,>=1.0.1; extra == "test" Requires-Dist: pandas<2.3.0; extra == "test" Requires-Dist: parameterized<0.10.0,>=0.7.1; extra == "test" Requires-Dist: pydot<2,>=1.2.0; extra == "test" Requires-Dist: pyhamcrest!=1.10.0,<3.0.0,>=1.9; extra == "test" Requires-Dist: requests_mock<2.0,>=1.7; extra == "test" Requires-Dist: tenacity<10,>=8.0.0; extra == "test" Requires-Dist: pytest<10.0,>=7.1.2; extra == "test" Requires-Dist: pytest-xdist<4,>=2.5.0; extra == "test" Requires-Dist: pytest-timeout<3,>=2.1.0; extra == "test" Requires-Dist: scikit-learn<1.8.0,>=0.20.0; extra == "test" Requires-Dist: sqlalchemy<3.0,>=1.3; extra == "test" Requires-Dist: psycopg2-binary<3.0,>=2.8.5; extra == "test" Requires-Dist: testcontainers[kafka,milvus,mysql]<5.0.0,>=4.0.0; extra == "test" Requires-Dist: cryptography>=41.0.2; extra == "test" Requires-Dist: hypothesis<6.148.4,>5.0.0; extra == "test" Requires-Dist: virtualenv-clone<1.0,>=0.5; extra == "test" Requires-Dist: python-tds>=1.16.1; extra == "test" Requires-Dist: sqlalchemy-pytds>=1.0.2; extra == "test" Requires-Dist: pg8000>=1.31.5; extra == "test" Requires-Dist: PyMySQL>=1.1.0; extra == "test" Requires-Dist: oracledb>=3.1.1; extra == "test" Provides-Extra: gcp Requires-Dist: cachetools<7,>=3.1.0; extra == "gcp" Requires-Dist: google-api-core<3,>=2.0.0; extra == "gcp" Requires-Dist: google-apitools<0.5.32,>=0.5.31; python_version < "3.13" and extra == "gcp" Requires-Dist: google-apitools>=0.5.35; python_version >= "3.13" and extra == "gcp" Requires-Dist: google-auth<3,>=1.18.0; extra == "gcp" Requires-Dist: google-auth-httplib2<0.3.0,>=0.1.0; extra == "gcp" Requires-Dist: google-cloud-datastore<3,>=2.0.0; extra == "gcp" Requires-Dist: google-cloud-pubsub<3,>=2.1.0; extra == "gcp" Requires-Dist: google-cloud-storage<4,>=2.18.2; extra == "gcp" Requires-Dist: google-cloud-bigquery<4,>=2.0.0; extra == "gcp" Requires-Dist: google-cloud-bigquery-storage<3,>=2.6.3; extra == "gcp" Requires-Dist: google-cloud-core<3,>=2.0.0; extra == "gcp" Requires-Dist: google-cloud-bigtable<3,>=2.19.0; extra == "gcp" Requires-Dist: google-cloud-build<4,>=3.35.0; extra == "gcp" Requires-Dist: google-cloud-spanner<4,>=3.0.0; extra == "gcp" Requires-Dist: google-cloud-dlp<4,>=3.0.0; extra == "gcp" Requires-Dist: google-cloud-kms<4,>=3.0.0; extra == "gcp" Requires-Dist: google-cloud-language<3,>=2.0; extra == "gcp" Requires-Dist: google-cloud-secret-manager<3,>=2.0; extra == "gcp" Requires-Dist: google-cloud-videointelligence<3,>=2.0; extra == "gcp" Requires-Dist: google-cloud-vision<4,>=2; extra == "gcp" Requires-Dist: google-cloud-recommendations-ai<0.11.0,>=0.1.0; extra == "gcp" Requires-Dist: google-cloud-aiplatform<2.0,>=1.26.0; extra == "gcp" Requires-Dist: cloud-sql-python-connector<2.0.0,>=1.18.2; extra == "gcp" Requires-Dist: python-tds>=1.16.1; extra == "gcp" Requires-Dist: pg8000>=1.31.5; extra == "gcp" Requires-Dist: PyMySQL>=1.1.0; extra == "gcp" Requires-Dist: keyrings.google-artifactregistry-auth; extra == "gcp" Requires-Dist: orjson<4,>=3.9.7; extra == "gcp" Requires-Dist: regex>=2020.6.8; extra == "gcp" Provides-Extra: interactive Requires-Dist: facets-overview<2,>=1.1.0; extra == "interactive" Requires-Dist: google-cloud-dataproc<6,>=5.0.0; extra == "interactive" Requires-Dist: ipython<9,>=7; extra == "interactive" Requires-Dist: ipykernel<7,>=6; extra == "interactive" Requires-Dist: ipywidgets<9,>=8; extra == "interactive" Requires-Dist: jupyter-client!=6.1.13,<8.2.1,>=6.1.11; extra == "interactive" Requires-Dist: pydot<2,>=1.2.0; extra == "interactive" Requires-Dist: timeloop<2,>=1.0.2; extra == "interactive" Requires-Dist: nbformat<6,>=5.0.5; extra == "interactive" Requires-Dist: nbconvert<8,>=6.2.0; extra == "interactive" Requires-Dist: pandas!=1.5.0,!=1.5.1,<2.3,>=1.4.3; extra == "interactive" Provides-Extra: interactive-test Requires-Dist: needle<1,>=0.5.0; extra == "interactive-test" Requires-Dist: chromedriver-binary<118,>=117; extra == "interactive-test" Requires-Dist: pillow<10,>=7.1.1; extra == "interactive-test" Requires-Dist: urllib3<2,>=1.21.1; extra == "interactive-test" Provides-Extra: ml-test Requires-Dist: datatable; extra == "ml-test" Requires-Dist: dill; extra == "ml-test" Requires-Dist: tensorflow_transform<1.15.0,>=1.14.0; extra == "ml-test" Requires-Dist: numpy<2; extra == "ml-test" Requires-Dist: tf2onnx<1.17,>=1.16.1; extra == "ml-test" Requires-Dist: embeddings>=0.0.4; extra == "ml-test" Requires-Dist: onnxruntime; extra == "ml-test" Requires-Dist: onnx<2,>=1.14.1; extra == "ml-test" Requires-Dist: langchain; extra == "ml-test" Requires-Dist: sentence-transformers>=2.2.2; extra == "ml-test" Requires-Dist: skl2onnx; extra == "ml-test" Requires-Dist: pyod>=0.7.6; extra == "ml-test" Requires-Dist: tensorflow; extra == "ml-test" Requires-Dist: absl-py>=0.12.0; extra == "ml-test" Requires-Dist: tensorflow-hub; extra == "ml-test" Requires-Dist: tokenizers<0.23; extra == "ml-test" Requires-Dist: torch<2.12.0; extra == "ml-test" Requires-Dist: transformers<4.56.0,>=4.28.0; extra == "ml-test" Requires-Dist: tokenizers<0.22.0,>=0.13.3; extra == "ml-test" Provides-Extra: p310-ml-test Requires-Dist: datatable; extra == "p310-ml-test" Requires-Dist: embeddings>=0.0.4; extra == "p310-ml-test" Requires-Dist: onnxruntime; extra == "p310-ml-test" Requires-Dist: onnx<2,>=1.14.1; extra == "p310-ml-test" Requires-Dist: langchain; extra == "p310-ml-test" Requires-Dist: sentence-transformers>=2.2.2; extra == "p310-ml-test" Requires-Dist: skl2onnx; extra == "p310-ml-test" Requires-Dist: pyod>=0.7.6; extra == "p310-ml-test" Requires-Dist: tensorflow; extra == "p310-ml-test" Requires-Dist: absl-py>=0.12.0; extra == "p310-ml-test" Requires-Dist: tensorflow-hub; extra == "p310-ml-test" Requires-Dist: tokenizers<0.23; extra == "p310-ml-test" Requires-Dist: torch<2.12.0; extra == "p310-ml-test" Requires-Dist: transformers<4.56.0,>=4.28.0; extra == "p310-ml-test" Requires-Dist: tokenizers<0.22.0,>=0.13.3; extra == "p310-ml-test" Requires-Dist: google-adk==1.28.1; extra == "p310-ml-test" Requires-Dist: proto-plus<2,>=1.26.1; extra == "p310-ml-test" Requires-Dist: opentelemetry-api==1.37.0; extra == "p310-ml-test" Requires-Dist: opentelemetry-sdk==1.37.0; extra == "p310-ml-test" Requires-Dist: opentelemetry-exporter-otlp-proto-http==1.37.0; extra == "p310-ml-test" Requires-Dist: tf2onnx<1.18,>=1.17.0; extra == "p310-ml-test" Provides-Extra: p312-ml-test Requires-Dist: datatable; extra == "p312-ml-test" Requires-Dist: embeddings>=0.0.4; extra == "p312-ml-test" Requires-Dist: onnxruntime; extra == "p312-ml-test" Requires-Dist: onnx<2,>=1.14.1; extra == "p312-ml-test" Requires-Dist: langchain; extra == "p312-ml-test" Requires-Dist: sentence-transformers>=2.2.2; extra == "p312-ml-test" Requires-Dist: skl2onnx; extra == "p312-ml-test" Requires-Dist: pyod>=0.7.6; extra == "p312-ml-test" Requires-Dist: tensorflow; extra == "p312-ml-test" Requires-Dist: absl-py>=0.12.0; extra == "p312-ml-test" Requires-Dist: tensorflow-hub; extra == "p312-ml-test" Requires-Dist: tokenizers<0.23; extra == "p312-ml-test" Requires-Dist: torch<2.12.0; extra == "p312-ml-test" Requires-Dist: transformers<4.56.0,>=4.28.0; extra == "p312-ml-test" Requires-Dist: tokenizers<0.22.0,>=0.13.3; extra == "p312-ml-test" Requires-Dist: google-adk==1.28.1; extra == "p312-ml-test" Requires-Dist: proto-plus<2,>=1.26.1; extra == "p312-ml-test" Requires-Dist: opentelemetry-api==1.37.0; extra == "p312-ml-test" Requires-Dist: opentelemetry-sdk==1.37.0; extra == "p312-ml-test" Requires-Dist: opentelemetry-exporter-otlp-proto-http==1.37.0; extra == "p312-ml-test" Requires-Dist: tf2onnx<1.18,>=1.17.0; extra == "p312-ml-test" Provides-Extra: p313-ml-test Requires-Dist: embeddings>=0.0.4; extra == "p313-ml-test" Requires-Dist: onnxruntime; extra == "p313-ml-test" Requires-Dist: onnx<2,>=1.14.1; extra == "p313-ml-test" Requires-Dist: langchain; extra == "p313-ml-test" Requires-Dist: sentence-transformers>=2.2.2; extra == "p313-ml-test" Requires-Dist: skl2onnx; extra == "p313-ml-test" Requires-Dist: pyod>=0.7.6; extra == "p313-ml-test" Requires-Dist: tensorflow; extra == "p313-ml-test" Requires-Dist: absl-py>=0.12.0; extra == "p313-ml-test" Requires-Dist: tensorflow-hub; extra == "p313-ml-test" Requires-Dist: tokenizers<0.23; extra == "p313-ml-test" Requires-Dist: torch<2.12.0; extra == "p313-ml-test" Requires-Dist: transformers<4.56.0,>=4.28.0; extra == "p313-ml-test" Requires-Dist: tokenizers<0.22.0,>=0.13.3; extra == "p313-ml-test" Requires-Dist: google-adk==1.28.1; extra == "p313-ml-test" Requires-Dist: proto-plus<2,>=1.26.1; extra == "p313-ml-test" Requires-Dist: opentelemetry-api==1.37.0; extra == "p313-ml-test" Requires-Dist: opentelemetry-sdk==1.37.0; extra == "p313-ml-test" Requires-Dist: opentelemetry-exporter-otlp-proto-http==1.37.0; extra == "p313-ml-test" Requires-Dist: tf2onnx<1.18,>=1.17.0; extra == "p313-ml-test" Requires-Dist: pymilvus<3.0.0,>=2.5.10; extra == "p313-ml-test" Provides-Extra: aws Requires-Dist: boto3<2,>=1.9; extra == "aws" Provides-Extra: azure Requires-Dist: azure-storage-blob<13,>=12.3.2; extra == "azure" Requires-Dist: azure-core<2,>=1.7.0; extra == "azure" Requires-Dist: azure-identity<2,>=1.12.0; extra == "azure" Provides-Extra: dataframe Requires-Dist: pandas!=1.5.0,!=1.5.1,<2.3,>=1.4.3; extra == "dataframe" Provides-Extra: dask Requires-Dist: distributed>=2024.4.2; extra == "dask" Requires-Dist: dask>=2024.4.2; extra == "dask" Provides-Extra: hadoop Requires-Dist: hdfs<3.0.0,>=2.1.0; extra == "hadoop" Provides-Extra: yaml Requires-Dist: docstring-parser<1.0,>=0.15; extra == "yaml" Requires-Dist: jinja2<3.2,>=3.0; extra == "yaml" Requires-Dist: virtualenv-clone<1.0,>=0.5; extra == "yaml" Requires-Dist: js2py<1,>=0.74; python_version < "3.12" and extra == "yaml" Requires-Dist: jsonschema<5.0.0,>=4.0.0; extra == "yaml" Requires-Dist: pandas!=1.5.0,!=1.5.1,<2.3,>=1.4.3; extra == "yaml" Provides-Extra: torch Requires-Dist: torch<2.8.0,>=1.9.0; extra == "torch" Provides-Extra: tensorflow Requires-Dist: tensorflow<2.21,>=2.12rc1; extra == "tensorflow" Requires-Dist: absl-py>=0.12.0; extra == "tensorflow" Provides-Extra: transformers Requires-Dist: transformers<4.56.0,>=4.28.0; extra == "transformers" Requires-Dist: tensorflow>=2.12.0; extra == "transformers" Requires-Dist: torch>=1.9.0; extra == "transformers" Provides-Extra: ml-cpu Requires-Dist: tensorflow>=2.12.0; extra == "ml-cpu" Requires-Dist: torch==2.8.0+cpu; extra == "ml-cpu" Requires-Dist: transformers<4.56.0,>=4.28.0; extra == "ml-cpu" Requires-Dist: absl-py>=0.12.0; extra == "ml-cpu" Provides-Extra: redis Requires-Dist: redis<6,>=5.0.0; extra == "redis" Provides-Extra: tft Requires-Dist: tensorflow_transform<1.15.0,>=1.14.0; extra == "tft" Requires-Dist: numpy<2; extra == "tft" Requires-Dist: dill; extra == "tft" Provides-Extra: tfrecord Requires-Dist: crcmod<2.0,>=1.7; extra == "tfrecord" Provides-Extra: onnx Requires-Dist: onnxruntime==1.13.1; extra == "onnx" Requires-Dist: torch==1.13.1; extra == "onnx" Requires-Dist: tensorflow==2.11.0; extra == "onnx" Requires-Dist: tf2onnx==1.13.0; extra == "onnx" Requires-Dist: skl2onnx==1.13; extra == "onnx" Requires-Dist: transformers==4.25.1; extra == "onnx" Requires-Dist: absl-py>=0.12.0; extra == "onnx" Provides-Extra: xgboost Requires-Dist: xgboost<2.1.3,>=1.6.0; extra == "xgboost" Requires-Dist: datatable==1.0.0; extra == "xgboost" Provides-Extra: tensorflow-hub Requires-Dist: tensorflow-hub<0.16.0,>=0.14.0; extra == "tensorflow-hub" Provides-Extra: milvus Requires-Dist: pymilvus<3.0.0,>=2.5.10; extra == "milvus" Provides-Extra: vllm Requires-Dist: openai==1.107.1; extra == "vllm" Requires-Dist: vllm==0.10.1.1; extra == "vllm" Requires-Dist: triton==3.3.1; extra == "vllm" Dynamic: author Dynamic: author-email Dynamic: classifier Dynamic: description Dynamic: description-content-type Dynamic: download-url Dynamic: home-page Dynamic: keywords Dynamic: license Dynamic: provides-extra Dynamic: requires-dist Dynamic: requires-python Dynamic: summary # Apache Beam [Apache Beam](http://beam.apache.org/) is a unified model for defining both batch and streaming data-parallel processing pipelines, as well as a set of language-specific SDKs for constructing pipelines and Runners for executing them on distributed processing backends, including [Apache Flink](http://flink.apache.org/), [Apache Spark](http://spark.apache.org/), [Google Cloud Dataflow](http://cloud.google.com/dataflow/), and [Hazelcast Jet](https://jet.hazelcast.org/). ## Overview Beam provides a general approach to expressing [embarrassingly parallel](https://en.wikipedia.org/wiki/Embarrassingly_parallel) data processing pipelines and supports three categories of users, each of which have relatively disparate backgrounds and needs. 1. _End Users_: Writing pipelines with an existing SDK, running it on an existing runner. These users want to focus on writing their application logic and have everything else just work. 2. _SDK Writers_: Developing a Beam SDK targeted at a specific user community (Java, Python, Scala, Go, R, graphical, etc). These users are language geeks and would prefer to be shielded from all the details of various runners and their implementations. 3. _Runner Writers_: Have an execution environment for distributed processing and would like to support programs written against the Beam Model. Would prefer to be shielded from details of multiple SDKs. ### The Beam Model The model behind Beam evolved from several internal Google data processing projects, including [MapReduce](http://research.google.com/archive/mapreduce.html), [FlumeJava](http://research.google.com/pubs/pub35650.html), and [Millwheel](http://research.google.com/pubs/pub41378.html). This model was originally known as the “[Dataflow Model](http://www.vldb.org/pvldb/vol8/p1792-Akidau.pdf)”. To learn more about the Beam Model (though still under the original name of Dataflow), see the World Beyond Batch: [Streaming 101](https://www.oreilly.com/ideas/the-world-beyond-batch-streaming-101) and [Streaming 102](https://www.oreilly.com/ideas/the-world-beyond-batch-streaming-102) posts on O’Reilly’s Radar site, and the [VLDB 2015 paper](http://www.vldb.org/pvldb/vol8/p1792-Akidau.pdf). The key concepts in the Beam programming model are: * `PCollection`: represents a collection of data, which could be bounded or unbounded in size. * `PTransform`: represents a computation that transforms input PCollections into output PCollections. * `Pipeline`: manages a directed acyclic graph of PTransforms and PCollections that is ready for execution. * `PipelineRunner`: specifies where and how the pipeline should execute. ### Runners Beam supports executing programs on multiple distributed processing backends through PipelineRunners. Currently, the following PipelineRunners are available: - The `DirectRunner` runs the pipeline on your local machine. - The `PrismRunner` runs the pipeline on your local machine using Beam Portability. - The `DataflowRunner` submits the pipeline to the [Google Cloud Dataflow](http://cloud.google.com/dataflow/). - The `FlinkRunner` runs the pipeline on an Apache Flink cluster. The code has been donated from [dataArtisans/flink-dataflow](https://github.com/dataArtisans/flink-dataflow) and is now part of Beam. - The `SparkRunner` runs the pipeline on an Apache Spark cluster. - The `JetRunner` runs the pipeline on a Hazelcast Jet cluster. The code has been donated from [hazelcast/hazelcast-jet](https://github.com/hazelcast/hazelcast-jet) and is now part of Beam. - The `Twister2Runner` runs the pipeline on a Twister2 cluster. The code has been donated from [DSC-SPIDAL/twister2](https://github.com/DSC-SPIDAL/twister2) and is now part of Beam. Have ideas for new Runners? See the [runner-ideas label](https://github.com/apache/beam/issues?q=is%3Aopen+is%3Aissue+label%3Arunner-ideas). ## Get started with the Python SDK Get started with the [Beam Python SDK quickstart](/get-started/quickstart-py) to set up your Python development environment, get the Beam SDK for Python, and run an example pipeline. Then, read through the [Beam programming guide](/documentation/programming-guide) to learn the basic concepts that apply to all SDKs in Beam. The [Python Tips document](../../contributor-docs/python-tips.md) is also a useful resource for setting up a development environment and performing common processes. See the [Python API reference](https://beam.apache.org/releases/pydoc/) for more information on individual APIs. ## Python streaming pipelines Python [streaming pipeline execution](https://beam.apache.org/documentation/sdks/python-streaming) is available (with some [limitations](https://beam.apache.org//documentation/sdks/python-streaming/#unsupported-features)) starting with Beam SDK version 2.5.0. ## Python type safety Python is a dynamically-typed language with no static type checking. The Beam SDK for Python uses type hints during pipeline construction and runtime to try to emulate the correctness guarantees achieved by true static typing. [Ensuring Python Type Safety](https://beam.apache.org/documentation/sdks/python-type-safety) walks through how to use type hints, which help you to catch potential bugs up front with the [Direct Runner](https://beam.apache.org//documentation/runners/direct/). ## Managing Python pipeline dependencies When you run your pipeline locally, the packages that your pipeline depends on are available because they are installed on your local machine. However, when you want to run your pipeline remotely, you must make sure these dependencies are available on the remote machines. [Managing Python Pipeline Dependencies](https://beam.apache.org/documentation/sdks/python-pipeline-dependencies) shows you how to make your dependencies available to the remote workers. ## Developing new I/O connectors for Python The Beam SDK for Python provides an extensible API that you can use to create new I/O connectors. See the [Developing I/O connectors overview](https://beam.apache.org/documentation/io/developing-io-overview) for information about developing new I/O connectors and links to language-specific implementation guidance. ## Making machine learning inferences with Python To integrate machine learning models into your pipelines for making inferences, use the RunInference API for PyTorch and Scikit-learn models. If you are using TensorFlow models, you can make use of the [library from `tfx_bsl`](https://github.com/tensorflow/tfx-bsl/tree/master/tfx_bsl/beam). You can create multiple types of transforms using the RunInference API: the API takes multiple types of setup parameters from model handlers, and the parameter type determines the model implementation. For more information, see [About Beam ML](https://beam.apache.org/documentation/ml/about-ml). [TensorFlow Extended (TFX)](https://www.tensorflow.org/tfx) is an end-to-end platform for deploying production ML pipelines. TFX is integrated with Beam. For more information, see [TFX user guide](https://www.tensorflow.org/tfx/guide). ## Python multi-language pipelines quickstart Apache Beam lets you combine transforms written in any supported SDK language and use them in one multi-language pipeline. To learn how to create a multi-language pipeline using the Python SDK, see the [Python multi-language pipelines quickstart](https://beam.apache.org/documentation/sdks/python-multi-language-pipelines). ## Unrecoverable Errors in Beam Python Some common errors can occur during worker start-up and prevent jobs from starting. To learn about these errors and how to troubleshoot them in the Python SDK, see [Unrecoverable Errors in Beam Python](https://beam.apache.org/documentation/sdks/python-unrecoverable-errors). ## 📚 Learn More Here are some resources actively maintained by the Beam community to help you get started:
Resource Details
Apache Beam Website Our website discussing the project, and it's specifics.
Python Quickstart A guide to getting started with the Python SDK.
Tour of Beam A comprehensive, interactive learning experience covering Beam concepts in depth.
Beam Quest A certification granted by Google Cloud, certifying proficiency in Beam.
Community Metrics Beam's Git Community Metrics.
## Contribution Instructions for building and testing Beam itself are in the [contribution guide](https://beam.apache.org/contribute). ## Contact Us To get involved with Apache Beam: * [Subscribe to](https://beam.apache.org/community/contact-us/#:~:text=Subscribe%20and%20Unsubscribe) or e-mail the [user@beam.apache.org](http://mail-archives.apache.org/mod_mbox/beam-user/) list. * [Subscribe to](https://beam.apache.org/community/contact-us/#:~:text=Subscribe%20and%20Unsubscribe) or e-mail the [dev@beam.apache.org](http://mail-archives.apache.org/mod_mbox/beam-dev/) list. * [Join ASF Slack](https://s.apache.org/slack-invite) on [#beam channel](https://s.apache.org/beam-slack-channel) * [Report an issue](https://github.com/apache/beam/issues/new/choose). ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/README.md0000644000175100017510000002310615206113032014731 0ustar00runnerrunner# Apache Beam [Apache Beam](http://beam.apache.org/) is a unified model for defining both batch and streaming data-parallel processing pipelines, as well as a set of language-specific SDKs for constructing pipelines and Runners for executing them on distributed processing backends, including [Apache Flink](http://flink.apache.org/), [Apache Spark](http://spark.apache.org/), [Google Cloud Dataflow](http://cloud.google.com/dataflow/), and [Hazelcast Jet](https://jet.hazelcast.org/). ## Overview Beam provides a general approach to expressing [embarrassingly parallel](https://en.wikipedia.org/wiki/Embarrassingly_parallel) data processing pipelines and supports three categories of users, each of which have relatively disparate backgrounds and needs. 1. _End Users_: Writing pipelines with an existing SDK, running it on an existing runner. These users want to focus on writing their application logic and have everything else just work. 2. _SDK Writers_: Developing a Beam SDK targeted at a specific user community (Java, Python, Scala, Go, R, graphical, etc). These users are language geeks and would prefer to be shielded from all the details of various runners and their implementations. 3. _Runner Writers_: Have an execution environment for distributed processing and would like to support programs written against the Beam Model. Would prefer to be shielded from details of multiple SDKs. ### The Beam Model The model behind Beam evolved from several internal Google data processing projects, including [MapReduce](http://research.google.com/archive/mapreduce.html), [FlumeJava](http://research.google.com/pubs/pub35650.html), and [Millwheel](http://research.google.com/pubs/pub41378.html). This model was originally known as the “[Dataflow Model](http://www.vldb.org/pvldb/vol8/p1792-Akidau.pdf)”. To learn more about the Beam Model (though still under the original name of Dataflow), see the World Beyond Batch: [Streaming 101](https://www.oreilly.com/ideas/the-world-beyond-batch-streaming-101) and [Streaming 102](https://www.oreilly.com/ideas/the-world-beyond-batch-streaming-102) posts on O’Reilly’s Radar site, and the [VLDB 2015 paper](http://www.vldb.org/pvldb/vol8/p1792-Akidau.pdf). The key concepts in the Beam programming model are: * `PCollection`: represents a collection of data, which could be bounded or unbounded in size. * `PTransform`: represents a computation that transforms input PCollections into output PCollections. * `Pipeline`: manages a directed acyclic graph of PTransforms and PCollections that is ready for execution. * `PipelineRunner`: specifies where and how the pipeline should execute. ### Runners Beam supports executing programs on multiple distributed processing backends through PipelineRunners. Currently, the following PipelineRunners are available: - The `DirectRunner` runs the pipeline on your local machine. - The `PrismRunner` runs the pipeline on your local machine using Beam Portability. - The `DataflowRunner` submits the pipeline to the [Google Cloud Dataflow](http://cloud.google.com/dataflow/). - The `FlinkRunner` runs the pipeline on an Apache Flink cluster. The code has been donated from [dataArtisans/flink-dataflow](https://github.com/dataArtisans/flink-dataflow) and is now part of Beam. - The `SparkRunner` runs the pipeline on an Apache Spark cluster. - The `JetRunner` runs the pipeline on a Hazelcast Jet cluster. The code has been donated from [hazelcast/hazelcast-jet](https://github.com/hazelcast/hazelcast-jet) and is now part of Beam. - The `Twister2Runner` runs the pipeline on a Twister2 cluster. The code has been donated from [DSC-SPIDAL/twister2](https://github.com/DSC-SPIDAL/twister2) and is now part of Beam. Have ideas for new Runners? See the [runner-ideas label](https://github.com/apache/beam/issues?q=is%3Aopen+is%3Aissue+label%3Arunner-ideas). ## Get started with the Python SDK Get started with the [Beam Python SDK quickstart](/get-started/quickstart-py) to set up your Python development environment, get the Beam SDK for Python, and run an example pipeline. Then, read through the [Beam programming guide](/documentation/programming-guide) to learn the basic concepts that apply to all SDKs in Beam. The [Python Tips document](../../contributor-docs/python-tips.md) is also a useful resource for setting up a development environment and performing common processes. See the [Python API reference](https://beam.apache.org/releases/pydoc/) for more information on individual APIs. ## Python streaming pipelines Python [streaming pipeline execution](https://beam.apache.org/documentation/sdks/python-streaming) is available (with some [limitations](https://beam.apache.org//documentation/sdks/python-streaming/#unsupported-features)) starting with Beam SDK version 2.5.0. ## Python type safety Python is a dynamically-typed language with no static type checking. The Beam SDK for Python uses type hints during pipeline construction and runtime to try to emulate the correctness guarantees achieved by true static typing. [Ensuring Python Type Safety](https://beam.apache.org/documentation/sdks/python-type-safety) walks through how to use type hints, which help you to catch potential bugs up front with the [Direct Runner](https://beam.apache.org//documentation/runners/direct/). ## Managing Python pipeline dependencies When you run your pipeline locally, the packages that your pipeline depends on are available because they are installed on your local machine. However, when you want to run your pipeline remotely, you must make sure these dependencies are available on the remote machines. [Managing Python Pipeline Dependencies](https://beam.apache.org/documentation/sdks/python-pipeline-dependencies) shows you how to make your dependencies available to the remote workers. ## Developing new I/O connectors for Python The Beam SDK for Python provides an extensible API that you can use to create new I/O connectors. See the [Developing I/O connectors overview](https://beam.apache.org/documentation/io/developing-io-overview) for information about developing new I/O connectors and links to language-specific implementation guidance. ## Making machine learning inferences with Python To integrate machine learning models into your pipelines for making inferences, use the RunInference API for PyTorch and Scikit-learn models. If you are using TensorFlow models, you can make use of the [library from `tfx_bsl`](https://github.com/tensorflow/tfx-bsl/tree/master/tfx_bsl/beam). You can create multiple types of transforms using the RunInference API: the API takes multiple types of setup parameters from model handlers, and the parameter type determines the model implementation. For more information, see [About Beam ML](https://beam.apache.org/documentation/ml/about-ml). [TensorFlow Extended (TFX)](https://www.tensorflow.org/tfx) is an end-to-end platform for deploying production ML pipelines. TFX is integrated with Beam. For more information, see [TFX user guide](https://www.tensorflow.org/tfx/guide). ## Python multi-language pipelines quickstart Apache Beam lets you combine transforms written in any supported SDK language and use them in one multi-language pipeline. To learn how to create a multi-language pipeline using the Python SDK, see the [Python multi-language pipelines quickstart](https://beam.apache.org/documentation/sdks/python-multi-language-pipelines). ## Unrecoverable Errors in Beam Python Some common errors can occur during worker start-up and prevent jobs from starting. To learn about these errors and how to troubleshoot them in the Python SDK, see [Unrecoverable Errors in Beam Python](https://beam.apache.org/documentation/sdks/python-unrecoverable-errors). ## 📚 Learn More Here are some resources actively maintained by the Beam community to help you get started:
Resource Details
Apache Beam Website Our website discussing the project, and it's specifics.
Python Quickstart A guide to getting started with the Python SDK.
Tour of Beam A comprehensive, interactive learning experience covering Beam concepts in depth.
Beam Quest A certification granted by Google Cloud, certifying proficiency in Beam.
Community Metrics Beam's Git Community Metrics.
## Contribution Instructions for building and testing Beam itself are in the [contribution guide](https://beam.apache.org/contribute). ## Contact Us To get involved with Apache Beam: * [Subscribe to](https://beam.apache.org/community/contact-us/#:~:text=Subscribe%20and%20Unsubscribe) or e-mail the [user@beam.apache.org](http://mail-archives.apache.org/mod_mbox/beam-user/) list. * [Subscribe to](https://beam.apache.org/community/contact-us/#:~:text=Subscribe%20and%20Unsubscribe) or e-mail the [dev@beam.apache.org](http://mail-archives.apache.org/mod_mbox/beam-dev/) list. * [Join ASF Slack](https://s.apache.org/slack-invite) on [#beam channel](https://s.apache.org/beam-slack-channel) * [Report an issue](https://github.com/apache/beam/issues/new/choose). ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996270.8663535 apache_beam-2.74.0/apache_beam/0000755000175100017510000000000015206113157015665 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/__init__.py0000644000175100017510000000743315206113032017775 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Apache Beam SDK for Python ========================== `Apache Beam `_ provides a simple, powerful programming model for building both batch and streaming parallel data processing pipelines. The Apache Beam SDK for Python provides access to Apache Beam capabilities from the Python programming language. Overview -------- The key concepts in this programming model are * :class:`~apache_beam.pvalue.PCollection`: represents a collection of data, which could be bounded or unbounded in size. * :class:`~apache_beam.transforms.ptransform.PTransform`: represents a computation that transforms input PCollections into output PCollections. * :class:`~apache_beam.pipeline.Pipeline`: manages a directed acyclic graph of :class:`~apache_beam.transforms.ptransform.PTransform` s and :class:`~apache_beam.pvalue.PCollection` s that is ready for execution. * :class:`~apache_beam.runners.runner.PipelineRunner`: specifies where and how the pipeline should execute. * :class:`~apache_beam.io.iobase.Read`: read from an external source. * :class:`~apache_beam.io.iobase.Write`: write to an external data sink. Typical usage ------------- At the top of your source file:: import apache_beam as beam After this import statement * Transform classes are available as :class:`beam.FlatMap `, :class:`beam.GroupByKey `, etc. * Pipeline class is available as :class:`beam.Pipeline ` * Text read/write transforms are available as :class:`beam.io.ReadFromText `, :class:`beam.io.WriteToText `. Examples -------- The `examples subdirectory `_ has some examples. """ import sys import warnings if sys.version_info.major == 3: if sys.version_info.minor <= 9 or sys.version_info.minor >= 15: warnings.warn( 'This version of Apache Beam has not been sufficiently tested on ' 'Python %s.%s. You may encounter bugs or missing features.' % (sys.version_info.major, sys.version_info.minor)) pass else: raise RuntimeError( 'The Apache Beam SDK for Python is only supported on Python 3. ' 'It is not supported on Python [' + str(sys.version_info) + '].') # pylint: disable=wrong-import-position import apache_beam.internal.pickler from apache_beam import coders from apache_beam import io from apache_beam import metrics from apache_beam import typehints from apache_beam import version from apache_beam.pipeline import * from apache_beam.pvalue import PCollection from apache_beam.pvalue import Row from apache_beam.pvalue import TaggedOutput from apache_beam.transforms import * try: # Add mitigation for CVE-2023-47248 while Beam allows affected versions # of pyarrow. (https://github.com/apache/beam/issues/29392) import pyarrow_hotfix except ImportError: pass # pylint: enable=wrong-import-position __version__ = version.__version__ ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996270.8763547 apache_beam-2.74.0/apache_beam/coders/0000755000175100017510000000000015206113157017144 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/__init__.py0000644000175100017510000000162615206113032021252 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from apache_beam.coders.coders import * from apache_beam.coders.row_coder import * from apache_beam.coders.typecoders import registry ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/avro_record.py0000644000175100017510000000221415206113032022012 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """AvroRecord for AvroGenericCoder.""" # pytype: skip-file __all__ = ['AvroRecord'] class AvroRecord(object): """Simple wrapper class for dictionary records.""" def __init__(self, value): self.record = value def __eq__(self, other): return issubclass(type(other), AvroRecord) and self.record == other.record def __hash__(self): return hash(self.record) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/coder_impl.pxd0000644000175100017510000002344115206113032021772 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # cimport cython cimport cpython.ref cimport cpython.tuple cimport libc.stdint cimport libc.stdlib cimport libc.string cdef extern from "math.h": libc.stdint.int64_t abs "llabs"(libc.stdint.int64_t) ctypedef char* char_ptr from .stream cimport InputStream, OutputStream from apache_beam.utils cimport windowed_value cdef object loads, dumps, create_InputStream, create_OutputStream, ByteCountingOutputStream, get_varint_size, past_unicode # Temporarily untyped to allow monkeypatching on failed import. #cdef type WindowedValue cdef bint is_compiled cdef class CoderImpl(object): cpdef encode_to_stream(self, value, OutputStream stream, bint nested) cpdef decode_from_stream(self, InputStream stream, bint nested) cpdef bytes encode(self, value) cpdef decode(self, bytes encoded) cpdef bytes encode_nested(self, value) cpdef decode_nested(self, bytes encoded) cpdef estimate_size(self, value, bint nested=?) @cython.locals(varint_size=int, bits=libc.stdint.uint64_t) @cython.overflowcheck(False) cpdef int _get_nested_size(self, int inner_size, bint nested) cpdef get_estimated_size_and_observables(self, value, bint nested=?) cdef class SimpleCoderImpl(CoderImpl): pass cdef class StreamCoderImpl(CoderImpl): pass cdef class CallbackCoderImpl(CoderImpl): cdef object _encoder cdef object _decoder cdef object _size_estimator cdef object NoneType cdef unsigned char UNKNOWN_TYPE, NONE_TYPE, INT_TYPE, FLOAT_TYPE, BOOL_TYPE cdef unsigned char BYTES_TYPE, UNICODE_TYPE, LIST_TYPE, TUPLE_TYPE, DICT_TYPE cdef unsigned char SET_TYPE, ITERABLE_LIKE_TYPE cdef unsigned char PROTO_TYPE, DATACLASS_TYPE, NAMED_TUPLE_TYPE cdef set _ITERABLE_LIKE_TYPES cdef class FastPrimitivesCoderImpl(StreamCoderImpl): cdef CoderImpl fallback_coder_impl cdef CoderImpl iterable_coder_impl cdef object requires_deterministic_step_label cdef bint warn_deterministic_fallback cdef bint force_use_dill cdef bint use_relative_filepaths @cython.locals(dict_value=dict, int_value=libc.stdint.int64_t, unicode_value=unicode) cpdef encode_to_stream(self, value, OutputStream stream, bint nested) @cython.locals(t=int) cpdef decode_from_stream(self, InputStream stream, bint nested) cdef encode_special_deterministic(self, value, OutputStream stream) cdef encode_type_2_67_0(self, t, OutputStream stream) cdef encode_type(self, t, OutputStream stream) cdef decode_type(self, InputStream stream) cdef dict _pickled_types cdef dict _unpickled_types cdef class MapCoderImpl(StreamCoderImpl): cdef CoderImpl _key_coder cdef CoderImpl _value_coder cdef bint _is_deterministic cdef class BytesCoderImpl(CoderImpl): pass cdef class BooleanCoderImpl(CoderImpl): pass cdef class BigEndianShortCoderImpl(StreamCoderImpl): pass cdef class ByteCoderImpl(StreamCoderImpl): pass cdef class SinglePrecisionFloatCoderImpl(StreamCoderImpl): pass cdef class FloatCoderImpl(StreamCoderImpl): pass cdef class TimestampCoderImpl(StreamCoderImpl): cdef object timestamp_class cdef list small_ints cdef class VarIntCoderImpl(StreamCoderImpl): @cython.locals(ivalue=libc.stdint.int64_t) cpdef bytes encode(self, value) cdef class VarInt32CoderImpl(StreamCoderImpl): @cython.locals(ivalue=libc.stdint.int32_t) cpdef bytes encode(self, value) cdef class SingletonCoderImpl(CoderImpl): cdef object _value cdef class AbstractComponentCoderImpl(StreamCoderImpl): cdef tuple _coder_impls cpdef _extract_components(self, value) cpdef _construct_from_components(self, components) @cython.locals(c=CoderImpl) cpdef encode_to_stream(self, value, OutputStream stream, bint nested) @cython.locals(c=CoderImpl) cpdef decode_from_stream(self, InputStream stream, bint nested) @cython.locals(c=CoderImpl) cpdef get_estimated_size_and_observables(self, value, bint nested=?) cdef class TupleCoderImpl(AbstractComponentCoderImpl): pass cdef class SequenceCoderImpl(StreamCoderImpl): cdef CoderImpl _elem_coder cdef object _read_state cdef object _write_state cdef int _write_state_threshold cpdef _construct_from_sequence(self, values) @cython.locals(buffer=OutputStream, target_buffer_size=libc.stdint.int64_t, index=libc.stdint.int64_t, prev_index=libc.stdint.int64_t) cpdef encode_to_stream(self, value, OutputStream stream, bint nested) cdef class TupleSequenceCoderImpl(SequenceCoderImpl): pass cdef class _AbstractIterable: cdef object _contents cdef class IterableCoderImpl(SequenceCoderImpl): cdef bint _use_abstract_iterable cdef class ListCoderImpl(SequenceCoderImpl): pass cdef object IntervalWindow cdef class IntervalWindowCoderImpl(StreamCoderImpl): cdef libc.stdint.uint64_t _to_normal_time(self, libc.stdint.int64_t value) cdef libc.stdint.int64_t _from_normal_time(self, libc.stdint.uint64_t value) @cython.locals(typed_value=windowed_value._IntervalWindowBase, span_millis=libc.stdint.int64_t) cpdef encode_to_stream(self, value, OutputStream stream, bint nested) @cython.locals(typed_value=windowed_value._IntervalWindowBase) cpdef decode_from_stream(self, InputStream stream, bint nested) @cython.locals(typed_value=windowed_value._IntervalWindowBase, span_millis=libc.stdint.int64_t) cpdef estimate_size(self, value, bint nested=?) cdef int PaneInfoTiming_UNKNOWN cdef int PaneInfoEncoding_FIRST cdef class PaneInfoCoderImpl(StreamCoderImpl): cdef int _choose_encoding(self, windowed_value.PaneInfo value) @cython.locals(pane_info=windowed_value.PaneInfo, encoding_type=int) cpdef encode_to_stream(self, value, OutputStream stream, bint nested) @cython.locals(encoded_first_byte=int, encoding_type=int) cpdef decode_from_stream(self, InputStream stream, bint nested) cdef libc.stdint.uint64_t _TIME_SHIFT cdef libc.stdint.int64_t MIN_TIMESTAMP_micros cdef libc.stdint.int64_t MAX_TIMESTAMP_micros cdef class _OrderedUnionCoderImpl(StreamCoderImpl): cdef tuple _types cdef tuple _coder_impls cdef CoderImpl _fallback_coder_impl @cython.locals(ix=int, c=CoderImpl) cpdef encode_to_stream(self, value, OutputStream stream, bint nested) @cython.locals(ix=int, c=CoderImpl) cpdef decode_from_stream(self, InputStream stream, bint nested) cdef class WindowedValueCoderImpl(StreamCoderImpl): """A coder for windowed values.""" cdef CoderImpl _value_coder cdef CoderImpl _timestamp_coder cdef CoderImpl _windows_coder cdef CoderImpl _pane_info_coder cdef libc.stdint.uint64_t _to_normal_time(self, libc.stdint.int64_t value) cdef libc.stdint.int64_t _from_normal_time(self, libc.stdint.uint64_t value) @cython.locals(c=CoderImpl) cpdef get_estimated_size_and_observables(self, value, bint nested=?) @cython.locals(timestamp=libc.stdint.int64_t) cpdef decode_from_stream(self, InputStream stream, bint nested) @cython.locals(wv=windowed_value.WindowedValue, restore_sign=int) cpdef encode_to_stream(self, value, OutputStream stream, bint nested) cdef class ParamWindowedValueCoderImpl(WindowedValueCoderImpl): """A coder for windowed values with constant timestamp, windows and pane info.""" cdef readonly libc.stdint.int64_t _timestamp cdef readonly object _windows cdef readonly windowed_value.PaneInfo _pane_info cdef class LengthPrefixCoderImpl(StreamCoderImpl): cdef CoderImpl _value_coder cdef class RowColumnEncoder: cdef bint encode_to_stream(self, size_t index, OutputStream stream) except -1 cdef bint decode_from_stream(self, size_t index, InputStream stream) except -1 cdef class GenericRowColumnEncoder(RowColumnEncoder): cdef object column cdef CoderImpl coder_impl cdef class RowCoderImpl(StreamCoderImpl): cdef object schema cdef int num_fields cdef list field_names cdef list field_nullable cdef object constructor cdef list encoding_positions cdef list encoding_positions_argsort cdef bint encoding_positions_are_trivial cdef list components cdef bint has_nullable_fields @cython.locals(i=int, nvals=libc.stdint.int64_t, running=int, component_coder=CoderImpl, null_mask=bytes, null_mask_c=char_ptr) cpdef decode_from_stream(self, InputStream stream, bint nested) @cython.locals(i=int, nvals=libc.stdint.int64_t, k=size_t, n=size_t, null_mask=bytes, null_mask_c=char_ptr) cpdef decode_batch_from_stream(self, dict dest, InputStream stream) @cython.locals(i=int, running=int, component_coder=CoderImpl) cpdef encode_to_stream(self, value, OutputStream stream, bint nested) @cython.locals(i=int, k=size_t, n=size_t, null_flags=libc.stdint.uint8_t[:,::1], null_bits=libc.stdint.uint8_t[:,::1], has_null_bits=libc.stdint.uint8_t[::1]) cpdef encode_batch_to_stream(self, dict values, OutputStream stream) cdef class LogicalTypeCoderImpl(StreamCoderImpl): cdef object logical_type cdef CoderImpl representation_coder cpdef decode_from_stream(self, InputStream stream, bint nested) cpdef encode_to_stream(self, value, OutputStream stream, bint nested) cdef class BigIntegerCoderImpl(StreamCoderImpl): pass cdef class DecimalCoderImpl(StreamCoderImpl): pass ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/coder_impl.py0000644000175100017510000023100615206113032021625 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Coder implementations. The actual encode/decode implementations are split off from coders to allow conditional (compiled/pure) implementations, which can be used to encode many elements with minimal overhead. This module may be optionally compiled with Cython, using the corresponding coder_impl.pxd file for type hints. In particular, because CoderImpls are never pickled and sent across the wire (unlike Coders themselves) the workers can use compiled Impls even if the main program does not (or vice versa). For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file # ruff: noqa: UP006 import dataclasses import decimal import enum import itertools import json import logging import pickle from io import BytesIO from typing import TYPE_CHECKING from typing import Any from typing import Callable from typing import Dict from typing import Iterable from typing import Iterator from typing import List from typing import Optional from typing import Sequence from typing import Set from typing import Tuple from typing import Type import numpy as np from fastavro import parse_schema from fastavro import schemaless_reader from fastavro import schemaless_writer from apache_beam.coders import observable from apache_beam.coders.avro_record import AvroRecord from apache_beam.internal import cloudpickle_pickler from apache_beam.internal.cloudpickle import cloudpickle from apache_beam.typehints.schemas import named_tuple_from_schema from apache_beam.utils import proto_utils from apache_beam.utils import windowed_value from apache_beam.utils.sharded_key import ShardedKey from apache_beam.utils.timestamp import MAX_TIMESTAMP from apache_beam.utils.timestamp import MIN_TIMESTAMP from apache_beam.utils.timestamp import Timestamp try: import dill except ImportError: dill = None if TYPE_CHECKING: import proto from apache_beam.transforms import userstate from apache_beam.transforms.window import IntervalWindow try: from . import stream # pylint: disable=unused-import except ImportError: SLOW_STREAM = True else: SLOW_STREAM = False is_compiled = False fits_in_64_bits = lambda x: -(1 << 63) <= x <= (1 << 63) - 1 if TYPE_CHECKING or SLOW_STREAM: from .slow_stream import ByteCountingOutputStream from .slow_stream import InputStream as create_InputStream from .slow_stream import OutputStream as create_OutputStream from .slow_stream import get_varint_size try: import cython is_compiled = cython.compiled except ImportError: globals()['cython'] = type('fake_cython', (), {'cast': lambda typ, x: x}) else: # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports from .stream import ByteCountingOutputStream from .stream import InputStream as create_InputStream from .stream import OutputStream as create_OutputStream from .stream import get_varint_size # Make it possible to import create_InputStream and other cdef-classes # from apache_beam.coders.coder_impl when Cython codepath is used. globals()['create_InputStream'] = create_InputStream globals()['create_OutputStream'] = create_OutputStream globals()['ByteCountingOutputStream'] = ByteCountingOutputStream # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports is_compiled = True _LOGGER = logging.getLogger(__name__) _TIME_SHIFT = 1 << 63 MIN_TIMESTAMP_micros = MIN_TIMESTAMP.micros MAX_TIMESTAMP_micros = MAX_TIMESTAMP.micros IterableStateReader = Callable[[bytes, 'CoderImpl'], Iterable] IterableStateWriter = Callable[[Iterable, 'CoderImpl'], bytes] Observables = List[Tuple[observable.ObservableMixin, 'CoderImpl']] class CoderImpl(object): """For internal use only; no backwards-compatibility guarantees.""" def encode_to_stream(self, value, stream, nested): # type: (Any, create_OutputStream, bool) -> None """Reads object from potentially-nested encoding in stream.""" raise NotImplementedError def decode_from_stream(self, stream, nested): # type: (create_InputStream, bool) -> Any """Reads object from potentially-nested encoding in stream.""" raise NotImplementedError def encode(self, value): # type: (Any) -> bytes """Encodes an object to an unnested string.""" raise NotImplementedError def decode(self, encoded): # type: (bytes) -> Any """Decodes an object to an unnested string.""" raise NotImplementedError def encode_all(self, values): # type: (Iterable[Any]) -> bytes out = create_OutputStream() for value in values: self.encode_to_stream(value, out, True) return out.get() def decode_all(self, encoded): # type: (bytes) -> Iterator[Any] input_stream = create_InputStream(encoded) while input_stream.size() > 0: yield self.decode_from_stream(input_stream, True) def encode_nested(self, value): # type: (Any) -> bytes out = create_OutputStream() self.encode_to_stream(value, out, True) return out.get() def decode_nested(self, encoded): # type: (bytes) -> Any return self.decode_from_stream(create_InputStream(encoded), True) def estimate_size(self, value, nested=False): # type: (Any, bool) -> int """Estimates the encoded size of the given value, in bytes.""" out = ByteCountingOutputStream() self.encode_to_stream(value, out, nested) return out.get_count() def _get_nested_size(self, inner_size, nested): if not nested: return inner_size varint_size = get_varint_size(inner_size) return varint_size + inner_size def get_estimated_size_and_observables(self, value, nested=False): # type: (Any, bool) -> Tuple[int, Observables] """Returns estimated size of value along with any nested observables. The list of nested observables is returned as a list of 2-tuples of (obj, coder_impl), where obj is an instance of observable.ObservableMixin, and coder_impl is the CoderImpl that can be used to encode elements sent by obj to its observers. Arguments: value: the value whose encoded size is to be estimated. nested: whether the value is nested. Returns: The estimated encoded size of the given value and a list of observables whose elements are 2-tuples of (obj, coder_impl) as described above. """ return self.estimate_size(value, nested), [] class SimpleCoderImpl(CoderImpl): """For internal use only; no backwards-compatibility guarantees. Subclass of CoderImpl implementing stream methods using encode/decode.""" def encode_to_stream(self, value, stream, nested): # type: (Any, create_OutputStream, bool) -> None """Reads object from potentially-nested encoding in stream.""" stream.write(self.encode(value), nested) def decode_from_stream(self, stream, nested): # type: (create_InputStream, bool) -> Any """Reads object from potentially-nested encoding in stream.""" return self.decode(stream.read_all(nested)) class StreamCoderImpl(CoderImpl): """For internal use only; no backwards-compatibility guarantees. Subclass of CoderImpl implementing encode/decode using stream methods.""" def encode(self, value): # type: (Any) -> bytes out = create_OutputStream() self.encode_to_stream(value, out, False) return out.get() def decode(self, encoded): # type: (bytes) -> Any return self.decode_from_stream(create_InputStream(encoded), False) def estimate_size(self, value, nested=False): # type: (Any, bool) -> int """Estimates the encoded size of the given value, in bytes.""" out = ByteCountingOutputStream() self.encode_to_stream(value, out, nested) return out.get_count() class CallbackCoderImpl(CoderImpl): """For internal use only; no backwards-compatibility guarantees. A CoderImpl that calls back to the _impl methods on the Coder itself. This is the default implementation used if Coder._get_impl() is not overwritten. """ def __init__(self, encoder, decoder, size_estimator=None): self._encoder = encoder self._decoder = decoder self._size_estimator = size_estimator or self._default_size_estimator def _default_size_estimator(self, value): return len(self.encode(value)) def encode_to_stream(self, value, stream, nested): # type: (Any, create_OutputStream, bool) -> None return stream.write(self._encoder(value), nested) def decode_from_stream(self, stream, nested): # type: (create_InputStream, bool) -> Any return self._decoder(stream.read_all(nested)) def encode(self, value): return self._encoder(value) def decode(self, encoded): return self._decoder(encoded) def estimate_size(self, value, nested=False): # type: (Any, bool) -> int return self._get_nested_size(self._size_estimator(value), nested) def get_estimated_size_and_observables(self, value, nested=False): # type: (Any, bool) -> Tuple[int, Observables] # TODO(robertwb): Remove this once all coders are correct. if isinstance(value, observable.ObservableMixin): # CallbackCoderImpl can presumably encode the elements too. return 1, [(value, self)] return self.estimate_size(value, nested), [] def __repr__(self): return 'CallbackCoderImpl[encoder=%s, decoder=%s]' % ( self._encoder, self._decoder) class ProtoCoderImpl(SimpleCoderImpl): """For internal use only; no backwards-compatibility guarantees.""" def __init__(self, proto_message_type): self.proto_message_type = proto_message_type def encode(self, value): return value.SerializePartialToString() def decode(self, encoded): proto_message = self.proto_message_type() proto_message.ParseFromString(encoded) # This is in effect "ParsePartial". return proto_message class DeterministicProtoCoderImpl(ProtoCoderImpl): """For internal use only; no backwards-compatibility guarantees.""" def encode(self, value): return value.SerializePartialToString(deterministic=True) class ProtoPlusCoderImpl(SimpleCoderImpl): """For internal use only; no backwards-compatibility guarantees.""" def __init__(self, proto_plus_type): # type: (Type[proto.Message]) -> None self.proto_plus_type = proto_plus_type def encode(self, value): return value._pb.SerializePartialToString(deterministic=True) def decode(self, value): return self.proto_plus_type.deserialize(value) UNKNOWN_TYPE = 0xFF NONE_TYPE = 0 INT_TYPE = 1 FLOAT_TYPE = 2 BYTES_TYPE = 3 UNICODE_TYPE = 4 BOOL_TYPE = 9 LIST_TYPE = 5 TUPLE_TYPE = 6 DICT_TYPE = 7 SET_TYPE = 8 ITERABLE_LIKE_TYPE = 10 PROTO_TYPE = 100 DATACLASS_TYPE = 101 NAMED_TUPLE_TYPE = 102 ENUM_TYPE = 103 NESTED_STATE_TYPE = 104 DATACLASS_KW_ONLY_TYPE = 105 # Types that can be encoded as iterables, but are not literally # lists, etc. due to being lazy. The actual type is not preserved # through encoding, only the elements. This is particularly useful # for the value list types created in GroupByKey. _ITERABLE_LIKE_TYPES = set() # type: Set[Type] def _verify_dill_compat(): base_error = ( "This pipeline runs with the pipeline option " "--update_compatibility_version=2.67.0 or earlier. " "When running with this option on SDKs 2.68.0 or " "later, you must ensure dill==0.3.1.1 is installed.") if not dill: raise RuntimeError(base_error + ". Dill is not installed.") if dill.__version__ != "0.3.1.1": raise RuntimeError(base_error + f". Found dill version '{dill.__version__}") class FastPrimitivesCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees.""" def __init__( self, fallback_coder_impl, requires_deterministic_step_label=None, force_use_dill=False, use_relative_filepaths=True): self.fallback_coder_impl = fallback_coder_impl self.iterable_coder_impl = IterableCoderImpl(self) self.requires_deterministic_step_label = requires_deterministic_step_label self.warn_deterministic_fallback = True self.force_use_dill = force_use_dill self.use_relative_filepaths = use_relative_filepaths @staticmethod def register_iterable_like_type(t): _ITERABLE_LIKE_TYPES.add(t) def get_estimated_size_and_observables(self, value, nested=False): # type: (Any, bool) -> Tuple[int, Observables] if isinstance(value, observable.ObservableMixin): # FastPrimitivesCoderImpl can presumably encode the elements too. return 1, [(value, self)] out = ByteCountingOutputStream() self.encode_to_stream(value, out, nested) return out.get_count(), [] def encode_to_stream(self, value, stream, nested): # type: (Any, create_OutputStream, bool) -> None t = type(value) if value is None: stream.write_byte(NONE_TYPE) elif t is int: # In Python 3, an int may be larger than 64 bits. # We need to check whether value fits into a 64 bit integer before # writing the marker byte. try: # In Cython-compiled code this will throw an overflow error # when value does not fit into int64. int_value = value # If Cython is not used, we must do a (slower) check ourselves. if not TYPE_CHECKING and not is_compiled: if not fits_in_64_bits(value): raise OverflowError() stream.write_byte(INT_TYPE) stream.write_var_int64(int_value) except OverflowError: stream.write_byte(UNKNOWN_TYPE) self.fallback_coder_impl.encode_to_stream(value, stream, nested) elif t is float: stream.write_byte(FLOAT_TYPE) stream.write_bigendian_double(value) elif t is bytes: stream.write_byte(BYTES_TYPE) stream.write(value, nested) elif t is str: unicode_value = value # for typing stream.write_byte(UNICODE_TYPE) stream.write(unicode_value.encode('utf-8'), nested) elif t is list or t is tuple: stream.write_byte(LIST_TYPE if t is list else TUPLE_TYPE) stream.write_var_int64(len(value)) for e in value: self.encode_to_stream(e, stream, True) elif t is bool: stream.write_byte(BOOL_TYPE) stream.write_byte(value) elif t in _ITERABLE_LIKE_TYPES: stream.write_byte(ITERABLE_LIKE_TYPE) self.iterable_coder_impl.encode_to_stream(value, stream, nested) elif t is dict: dict_value = value # for typing stream.write_byte(DICT_TYPE) stream.write_var_int64(len(dict_value)) if self.requires_deterministic_step_label is not None: try: ordered_kvs = sorted(dict_value.items()) except Exception as exn: raise TypeError( "Unable to deterministically order keys of dict for '%s'" % self.requires_deterministic_step_label) from exn for k, v in ordered_kvs: self.encode_to_stream(k, stream, True) self.encode_to_stream(v, stream, True) else: # Loop over dict.items() is optimized by Cython. for k, v in dict_value.items(): self.encode_to_stream(k, stream, True) self.encode_to_stream(v, stream, True) elif t is set: stream.write_byte(SET_TYPE) stream.write_var_int64(len(value)) if self.requires_deterministic_step_label is not None: try: value = sorted(value) except Exception as exn: raise TypeError( "Unable to deterministically order element of set for '%s'" % self.requires_deterministic_step_label) from exn for e in value: self.encode_to_stream(e, stream, True) # All possibly deterministic encodings should be above this clause, # all non-deterministic ones below. elif self.requires_deterministic_step_label is not None: self.encode_special_deterministic(value, stream) else: stream.write_byte(UNKNOWN_TYPE) self.fallback_coder_impl.encode_to_stream(value, stream, nested) def encode_special_deterministic(self, value, stream): if self.warn_deterministic_fallback: _LOGGER.warning( "Using fallback deterministic coder for type '%s' in '%s'. ", type(value), self.requires_deterministic_step_label) self.warn_deterministic_fallback = False if isinstance(value, proto_utils.message_types): stream.write_byte(PROTO_TYPE) self.encode_type(type(value), stream) stream.write(value.SerializePartialToString(deterministic=True), True) elif dataclasses.is_dataclass(value): if not type(value).__dataclass_params__.frozen: raise TypeError( "Unable to deterministically encode non-frozen '%s' of type '%s' " "for the input of '%s'" % (value, type(value), self.requires_deterministic_step_label)) init_fields = [field for field in dataclasses.fields(value) if field.init] try: if any(field.kw_only for field in init_fields): stream.write_byte(DATACLASS_KW_ONLY_TYPE) self.encode_type(type(value), stream) stream.write_var_int64(len(init_fields)) for field in init_fields: stream.write(field.name.encode("utf-8"), True) self.encode_to_stream(getattr(value, field.name), stream, True) else: # Not using kw_only, we can pass parameters by position. stream.write_byte(DATACLASS_TYPE) self.encode_type(type(value), stream) values = [getattr(value, field.name) for field in init_fields] self.iterable_coder_impl.encode_to_stream(values, stream, True) except Exception as e: raise TypeError(self._deterministic_encoding_error_msg(value)) from e elif isinstance(value, tuple) and hasattr(type(value), '_fields'): stream.write_byte(NAMED_TUPLE_TYPE) self.encode_type(type(value), stream) try: self.iterable_coder_impl.encode_to_stream(value, stream, True) except Exception as e: raise TypeError(self._deterministic_encoding_error_msg(value)) from e elif isinstance(value, enum.Enum): stream.write_byte(ENUM_TYPE) self.encode_type(type(value), stream) # Enum values can be of any type. try: self.encode_to_stream(value.value, stream, True) except Exception as e: raise TypeError(self._deterministic_encoding_error_msg(value)) from e elif (hasattr(value, "__getstate__") and # https://github.com/apache/beam/issues/33020 type(value).__reduce__ == object.__reduce__): if not hasattr(value, "__setstate__"): raise TypeError( "Unable to deterministically encode '%s' of type '%s', " "for the input of '%s'. The object defines __getstate__ but not " "__setstate__." % (value, type(value), self.requires_deterministic_step_label)) stream.write_byte(NESTED_STATE_TYPE) self.encode_type(type(value), stream) state_value = value.__getstate__() try: self.encode_to_stream(state_value, stream, True) except Exception as e: raise TypeError(self._deterministic_encoding_error_msg(value)) from e else: raise TypeError(self._deterministic_encoding_error_msg(value)) def _deterministic_encoding_error_msg(self, value): return ( "Unable to deterministically encode '%s' of type '%s', " "please provide a type hint for the input of '%s'" % (value, type(value), self.requires_deterministic_step_label)) def encode_type_2_67_0(self, t, stream): """ Encode special type with <=2.67.0 compatibility. """ if t not in _pickled_types: _verify_dill_compat() _pickled_types[t] = dill.dumps(t) stream.write(_pickled_types[t], True) def encode_type(self, t, stream): if self.force_use_dill: return self.encode_type_2_67_0(t, stream) if t not in _pickled_types: config = cloudpickle.CloudPickleConfig( id_generator=None, skip_reset_dynamic_type_state=True, filepath_interceptor=cloudpickle.get_relative_path) if not self.use_relative_filepaths: config.filepath_interceptor = None _pickled_types[t] = cloudpickle_pickler.dumps(t, config=config) stream.write(_pickled_types[t], True) def decode_type(self, stream): if self.force_use_dill: return _unpickle_type_2_67_0(stream.read_all(True)) return _unpickle_type(stream.read_all(True)) def decode_from_stream(self, stream, nested): # type: (create_InputStream, bool) -> Any t = stream.read_byte() if t == NONE_TYPE: return None elif t == INT_TYPE: return stream.read_var_int64() elif t == FLOAT_TYPE: return stream.read_bigendian_double() elif t == BYTES_TYPE: return stream.read_all(nested) elif t == UNICODE_TYPE: return stream.read_all(nested).decode('utf-8') elif t == LIST_TYPE or t == TUPLE_TYPE or t == SET_TYPE: vlen = stream.read_var_int64() vlist = [self.decode_from_stream(stream, True) for _ in range(vlen)] if t == LIST_TYPE: return vlist elif t == TUPLE_TYPE: return tuple(vlist) return set(vlist) elif t == DICT_TYPE: vlen = stream.read_var_int64() v = {} for _ in range(vlen): k = self.decode_from_stream(stream, True) v[k] = self.decode_from_stream(stream, True) return v elif t == BOOL_TYPE: return not not stream.read_byte() elif t == ITERABLE_LIKE_TYPE: return self.iterable_coder_impl.decode_from_stream(stream, nested) elif t == PROTO_TYPE: cls = self.decode_type(stream) msg = cls() msg.ParseFromString(stream.read_all(True)) return msg elif t == DATACLASS_KW_ONLY_TYPE: cls = self.decode_type(stream) vlen = stream.read_var_int64() fields = {} for _ in range(vlen): field_name = stream.read_all(True).decode('utf-8') fields[field_name] = self.decode_from_stream(stream, True) return cls(**fields) elif t == DATACLASS_TYPE or t == NAMED_TUPLE_TYPE: cls = self.decode_type(stream) return cls(*self.iterable_coder_impl.decode_from_stream(stream, True)) elif t == ENUM_TYPE: cls = self.decode_type(stream) return cls(self.decode_from_stream(stream, True)) elif t == NESTED_STATE_TYPE: cls = self.decode_type(stream) state = self.decode_from_stream(stream, True) value = cls.__new__(cls) value.__setstate__(state) return value elif t == UNKNOWN_TYPE: return self.fallback_coder_impl.decode_from_stream(stream, nested) else: raise ValueError('Unknown type tag %x' % t) _pickled_types = {} # type: Dict[type, bytes] _unpickled_types = {} # type: Dict[bytes, type] def _unpickle_type_2_67_0(bs): """ Decode special type with <=2.67.0 compatibility. """ t = _unpickled_types.get(bs, None) if t is None: _verify_dill_compat() t = _unpickled_types[bs] = dill.loads(bs) # Fix unpicklable anonymous named tuples for Python 3.6. if t.__base__ is tuple and hasattr(t, '_fields'): try: pickle.loads(pickle.dumps(t)) except pickle.PicklingError: t.__reduce__ = lambda self: ( _unpickle_named_tuple_2_67_0, (bs, tuple(self))) return t def _unpickle_named_tuple_2_67_0(bs, items): return _unpickle_type_2_67_0(bs)(*items) def _unpickle_type(bs): if not _unpickled_types.get(bs, None): _unpickled_types[bs] = cloudpickle_pickler.loads(bs) return _unpickled_types[bs] def _unpickle_named_tuple(bs, items): return _unpickle_type(bs)(*items) class BytesCoderImpl(CoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder for bytes/str objects.""" def encode_to_stream(self, value, out, nested): # type: (bytes, create_OutputStream, bool) -> None # value might be of type np.bytes if passed from encode_batch, and cython # does not recognize it as bytes. if is_compiled and isinstance(value, np.bytes_): value = bytes(value) out.write(value, nested) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> bytes return in_stream.read_all(nested) def encode(self, value): assert isinstance(value, bytes), (value, type(value)) return value def decode(self, encoded): return encoded class BooleanCoderImpl(CoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder for bool objects.""" def encode_to_stream(self, value, out, nested): out.write_byte(1 if value else 0) def decode_from_stream(self, in_stream, nested): value = in_stream.read_byte() if value == 0: return False elif value == 1: return True raise ValueError("Expected 0 or 1, got %s" % value) def encode(self, value): return b'\x01' if value else b'\x00' def decode(self, encoded): value = ord(encoded) if value == 0: return False elif value == 1: return True raise ValueError("Expected 0 or 1, got %s" % value) def estimate_size(self, unused_value, nested=False): # Note that booleans are encoded the same way regardless of nesting. return 1 class MapCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. Note this implementation always uses nested context when encoding keys and values. This differs from Java's MapCoder, which uses nested=False if possible for the last value encoded. This difference is acceptable because MapCoder is not standard. It is only used in a standard context by RowCoder which always uses nested context for attribute values. A coder for typing.Mapping objects.""" def __init__( self, key_coder, # type: CoderImpl value_coder, # type: CoderImpl is_deterministic=False): self._key_coder = key_coder self._value_coder = value_coder self._is_deterministic = is_deterministic def encode_to_stream(self, dict_value, out, nested): out.write_bigendian_int32(len(dict_value)) # Note this implementation always uses nested context when encoding keys # and values which differs from Java. See note in docstring. if self._is_deterministic: for key, value in sorted(dict_value.items()): self._key_coder.encode_to_stream(key, out, True) self._value_coder.encode_to_stream(value, out, True) else: # This loop is separate from the above so the loop over dict.items() # will be optimized by Cython. for key, value in dict_value.items(): self._key_coder.encode_to_stream(key, out, True) self._value_coder.encode_to_stream(value, out, True) def decode_from_stream(self, in_stream, nested): size = in_stream.read_bigendian_int32() result = {} for _ in range(size): # Note this implementation always uses nested context when encoding keys # and values which differs from Java. See note in docstring. key = self._key_coder.decode_from_stream(in_stream, True) value = self._value_coder.decode_from_stream(in_stream, True) result[key] = value return result def estimate_size(self, unused_value, nested=False): estimate = 4 # 4 bytes for int32 size prefix for key, value in unused_value.items(): estimate += self._key_coder.estimate_size(key, True) estimate += self._value_coder.estimate_size(value, True) return estimate class NullableCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder for typing.Optional objects.""" ENCODE_NULL = 0 ENCODE_PRESENT = 1 def __init__( self, value_coder # type: CoderImpl ): self._value_coder = value_coder def encode_to_stream(self, value, out, nested): if value is None: out.write_byte(self.ENCODE_NULL) else: out.write_byte(self.ENCODE_PRESENT) self._value_coder.encode_to_stream(value, out, nested) def decode_from_stream(self, in_stream, nested): null_indicator = in_stream.read_byte() if null_indicator == self.ENCODE_NULL: return None elif null_indicator == self.ENCODE_PRESENT: return self._value_coder.decode_from_stream(in_stream, nested) else: raise ValueError( "Encountered unexpected value for null indicator: '%s'" % null_indicator) def estimate_size(self, unused_value, nested=False): return 1 + ( self._value_coder.estimate_size(unused_value) if unused_value is not None else 0) class BigEndianShortCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees.""" def encode_to_stream(self, value, out, nested): # type: (int, create_OutputStream, bool) -> None out.write_bigendian_int16(value) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> int return in_stream.read_bigendian_int16() def estimate_size(self, unused_value, nested=False): # type: (Any, bool) -> int # A short is encoded as 2 bytes, regardless of nesting. return 2 class ByteCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees.""" def encode_to_stream(self, value, out, nested): # type: (int, create_OutputStream, bool) -> None out.write_byte(value) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> int return in_stream.read_byte() def estimate_size(self, unused_value, nested=False): # type: (Any, bool) -> int # A byte is encoded as 1 byte, regardless of nesting. return 1 class SinglePrecisionFloatCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees.""" def encode_to_stream(self, value, out, nested): # type: (float, create_OutputStream, bool) -> None out.write_bigendian_float(value) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> float return in_stream.read_bigendian_float() def estimate_size(self, unused_value, nested=False): # type: (Any, bool) -> int # A float is encoded as 4 bytes, regardless of nesting. return 4 class FloatCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees.""" def encode_to_stream(self, value, out, nested): # type: (float, create_OutputStream, bool) -> None out.write_bigendian_double(value) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> float return in_stream.read_bigendian_double() def estimate_size(self, unused_value, nested=False): # type: (Any, bool) -> int # A double is encoded as 8 bytes, regardless of nesting. return 8 if not TYPE_CHECKING: IntervalWindow = None class IntervalWindowCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees.""" # TODO: Fn Harness only supports millis. Is this important enough to fix? def _to_normal_time(self, value): """Convert "lexicographically ordered unsigned" to signed.""" return value - _TIME_SHIFT def _from_normal_time(self, value): """Convert signed to "lexicographically ordered unsigned".""" return value + _TIME_SHIFT def encode_to_stream(self, value, out, nested): # type: (IntervalWindow, create_OutputStream, bool) -> None typed_value = value span_millis = ( typed_value._end_micros // 1000 - typed_value._start_micros // 1000) out.write_bigendian_uint64( self._from_normal_time(typed_value._end_micros // 1000)) out.write_var_int64(span_millis) def decode_from_stream(self, in_, nested): # type: (create_InputStream, bool) -> IntervalWindow if not TYPE_CHECKING: global IntervalWindow # pylint: disable=global-variable-not-assigned if IntervalWindow is None: from apache_beam.transforms.window import IntervalWindow # instantiating with None is not part of the public interface # pylint: disable=too-many-function-args typed_value = IntervalWindow(None, None) # type: ignore[arg-type] typed_value._end_micros = ( 1000 * self._to_normal_time(in_.read_bigendian_uint64())) typed_value._start_micros = ( typed_value._end_micros - 1000 * in_.read_var_int64()) return typed_value def estimate_size(self, value, nested=False): # type: (Any, bool) -> int # An IntervalWindow is context-insensitive, with a timestamp (8 bytes) # and a varint timespam. typed_value = value span_millis = ( typed_value._end_micros // 1000 - typed_value._start_micros // 1000) return 8 + get_varint_size(span_millis) class TimestampCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. TODO: SDK agnostic encoding For interoperability with Java SDK, encoding needs to match that of the Java SDK InstantCoder. https://github.com/apache/beam/blob/f5029b4f0dfff404310b2ef55e2632bbacc7b04f/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/InstantCoder.java#L79 """ def encode_to_stream(self, value, out, nested): # type: (Timestamp, create_OutputStream, bool) -> None millis = value.micros // 1000 if millis >= 0: millis = millis - _TIME_SHIFT else: millis = millis + _TIME_SHIFT out.write_bigendian_int64(millis) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> Timestamp millis = in_stream.read_bigendian_int64() if millis < 0: millis = millis + _TIME_SHIFT else: millis = millis - _TIME_SHIFT return Timestamp(micros=millis * 1000) def estimate_size(self, unused_value, nested=False): # A Timestamp is encoded as a 64-bit integer in 8 bytes, regardless of # nesting. return 8 class TimerCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees.""" def __init__(self, key_coder_impl, window_coder_impl): self._timestamp_coder_impl = TimestampCoderImpl() self._boolean_coder_impl = BooleanCoderImpl() self._pane_info_coder_impl = PaneInfoCoderImpl() self._key_coder_impl = key_coder_impl self._windows_coder_impl = TupleSequenceCoderImpl(window_coder_impl) from apache_beam.coders.coders import StrUtf8Coder self._tag_coder_impl = StrUtf8Coder().get_impl() def encode_to_stream(self, value, out, nested): # type: (userstate.Timer, create_OutputStream, bool) -> None self._key_coder_impl.encode_to_stream(value.user_key, out, True) self._tag_coder_impl.encode_to_stream(value.dynamic_timer_tag, out, True) self._windows_coder_impl.encode_to_stream(value.windows, out, True) self._boolean_coder_impl.encode_to_stream(value.clear_bit, out, True) if not value.clear_bit: self._timestamp_coder_impl.encode_to_stream( value.fire_timestamp, out, True) self._timestamp_coder_impl.encode_to_stream( value.hold_timestamp, out, True) self._pane_info_coder_impl.encode_to_stream(value.paneinfo, out, True) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> userstate.Timer from apache_beam.transforms import userstate user_key = self._key_coder_impl.decode_from_stream(in_stream, True) dynamic_timer_tag = self._tag_coder_impl.decode_from_stream(in_stream, True) windows = self._windows_coder_impl.decode_from_stream(in_stream, True) clear_bit = self._boolean_coder_impl.decode_from_stream(in_stream, True) if clear_bit: return userstate.Timer( user_key=user_key, dynamic_timer_tag=dynamic_timer_tag, windows=windows, clear_bit=clear_bit, fire_timestamp=None, hold_timestamp=None, paneinfo=None) return userstate.Timer( user_key=user_key, dynamic_timer_tag=dynamic_timer_tag, windows=windows, clear_bit=clear_bit, fire_timestamp=self._timestamp_coder_impl.decode_from_stream( in_stream, True), hold_timestamp=self._timestamp_coder_impl.decode_from_stream( in_stream, True), paneinfo=self._pane_info_coder_impl.decode_from_stream(in_stream, True)) small_ints = [chr(_).encode('latin-1') for _ in range(128)] class VarIntCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder for int objects.""" def encode_to_stream(self, value, out, nested): # type: (int, create_OutputStream, bool) -> None try: out.write_var_int64(value) except OverflowError as e: raise OverflowError( f"Integer value '{value}' is out of the encodable range for " f"VarIntCoder. This coder is limited to values that fit " f"within a 64-bit signed integer (-(2**63) to 2**63 - 1). " f"Original error: {e}") from e def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> int return in_stream.read_var_int64() def encode(self, value): ivalue = value # type cast if 0 <= ivalue < len(small_ints): return small_ints[ivalue] return StreamCoderImpl.encode(self, value) def decode(self, encoded): if len(encoded) == 1: i = ord(encoded) if 0 <= i < 128: return i return StreamCoderImpl.decode(self, encoded) def estimate_size(self, value, nested=False): # type: (Any, bool) -> int # Note that VarInts are encoded the same way regardless of nesting. try: return get_varint_size(value) except OverflowError as e: raise OverflowError( f"Cannot estimate size for integer value '{value}'. " f"Value is out of the range for VarIntCoder (64-bit signed integer). " f"Original error: {e}") from e class VarInt32CoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder for int32 objects.""" def encode_to_stream(self, value, out, nested): # type: (int, create_OutputStream, bool) -> None out.write_var_int32(value) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> int return in_stream.read_var_int32() def encode(self, value): ivalue = value # type cast if 0 <= ivalue < len(small_ints): return small_ints[ivalue] return StreamCoderImpl.encode(self, value) def decode(self, encoded): if len(encoded) == 1: i = ord(encoded) if 0 <= i < 128: return i return StreamCoderImpl.decode(self, encoded) def estimate_size(self, value, nested=False): # type: (Any, bool) -> int # Note that VarInts are encoded the same way regardless of nesting. return get_varint_size(int(value) & 0xFFFFFFFF) class SingletonCoderImpl(CoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder that always encodes exactly one value.""" def __init__(self, value): self._value = value def encode_to_stream(self, value, stream, nested): # type: (Any, create_OutputStream, bool) -> None pass def decode_from_stream(self, stream, nested): # type: (create_InputStream, bool) -> Any return self._value def encode(self, value): b = b'' # avoid byte vs str vs unicode error return b def decode(self, encoded): return self._value def estimate_size(self, value, nested=False): # type: (Any, bool) -> int return 0 class AbstractComponentCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. CoderImpl for coders that are comprised of several component coders.""" def __init__(self, coder_impls): for c in coder_impls: assert isinstance(c, CoderImpl), c self._coder_impls = tuple(coder_impls) def _extract_components(self, value): raise NotImplementedError def _construct_from_components(self, components): raise NotImplementedError def encode_to_stream(self, value, out, nested): # type: (Any, create_OutputStream, bool) -> None values = self._extract_components(value) if len(self._coder_impls) != len(values): raise ValueError('Number of components does not match number of coders.') for i in range(0, len(self._coder_impls)): c = self._coder_impls[i] # type cast c.encode_to_stream( values[i], out, nested or i + 1 < len(self._coder_impls)) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> Any return self._construct_from_components([ c.decode_from_stream( in_stream, nested or i + 1 < len(self._coder_impls)) for i, c in enumerate(self._coder_impls) ]) def estimate_size(self, value, nested=False): # type: (Any, bool) -> int """Estimates the encoded size of the given value, in bytes.""" # TODO(ccy): This ignores sizes of observable components. estimated_size, _ = (self.get_estimated_size_and_observables(value)) return estimated_size def get_estimated_size_and_observables(self, value, nested=False): # type: (Any, bool) -> Tuple[int, Observables] """Returns estimated size of value along with any nested observables.""" values = self._extract_components(value) estimated_size = 0 observables = [] # type: Observables for i in range(0, len(self._coder_impls)): c = self._coder_impls[i] # type cast child_size, child_observables = ( c.get_estimated_size_and_observables( values[i], nested=nested or i + 1 < len(self._coder_impls))) estimated_size += child_size observables += child_observables return estimated_size, observables class AvroCoderImpl(SimpleCoderImpl): """For internal use only; no backwards-compatibility guarantees.""" def __init__(self, schema): self.parsed_schema = parse_schema(json.loads(schema)) def encode(self, value): assert issubclass(type(value), AvroRecord) with BytesIO() as buf: schemaless_writer(buf, self.parsed_schema, value.record) return buf.getvalue() def decode(self, encoded): with BytesIO(encoded) as buf: return AvroRecord(schemaless_reader(buf, self.parsed_schema)) class TupleCoderImpl(AbstractComponentCoderImpl): """A coder for tuple objects.""" def _extract_components(self, value): return tuple(value) def _construct_from_components(self, components): return tuple(components) class _ConcatSequence(object): def __init__(self, head, tail): # type: (Iterable[Any], Iterable[Any]) -> None self._head = head self._tail = tail def __iter__(self): # type: () -> Iterator[Any] for elem in self._head: yield elem for elem in self._tail: yield elem def __eq__(self, other): return list(self) == list(other) def __hash__(self): raise NotImplementedError def __reduce__(self): return list, (list(self), ) FastPrimitivesCoderImpl.register_iterable_like_type(_ConcatSequence) class SequenceCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder for sequences. If the length of the sequence in known we encode the length as a 32 bit ``int`` followed by the encoded bytes. If the length of the sequence is unknown, we encode the length as ``-1`` followed by the encoding of elements buffered up to 64K bytes before prefixing the count of number of elements. A ``0`` is encoded at the end to indicate the end of stream. The resulting encoding would look like this:: -1 countA element(0) element(1) ... element(countA - 1) countB element(0) element(1) ... element(countB - 1) ... countX element(0) element(1) ... element(countX - 1) 0 If writing to state is enabled, the final terminating 0 will instead be repaced with:: varInt64(-1) len(state_token) state_token where state_token is a bytes object used to retrieve the remainder of the iterable via the state API. """ # Default buffer size of 64kB of handling iterables of unknown length. _DEFAULT_BUFFER_SIZE = 64 * 1024 def __init__( self, elem_coder, # type: CoderImpl read_state=None, # type: Optional[IterableStateReader] write_state=None, # type: Optional[IterableStateWriter] write_state_threshold=0 # type: int ): self._elem_coder = elem_coder self._read_state = read_state self._write_state = write_state self._write_state_threshold = write_state_threshold def _construct_from_sequence(self, values): raise NotImplementedError def encode_to_stream(self, value, out, nested): # type: (Sequence, create_OutputStream, bool) -> None # Compatible with Java's IterableLikeCoder. if hasattr(value, '__len__') and self._write_state is None: out.write_bigendian_int32(len(value)) for elem in value: self._elem_coder.encode_to_stream(elem, out, True) else: # We don't know the size without traversing it so use a fixed size buffer # and encode as many elements as possible into it before outputting # the size followed by the elements. # -1 to indicate that the length is not known. out.write_bigendian_int32(-1) buffer = create_OutputStream() if self._write_state is None: target_buffer_size = self._DEFAULT_BUFFER_SIZE else: target_buffer_size = min( self._DEFAULT_BUFFER_SIZE, self._write_state_threshold) prev_index = index = -1 # Don't want to miss out on fast list iteration optimization. value_iter = value if isinstance(value, (list, tuple)) else iter(value) start_size = out.size() for elem in value_iter: index += 1 self._elem_coder.encode_to_stream(elem, buffer, True) if buffer.size() > target_buffer_size: out.write_var_int64(index - prev_index) out.write(buffer.get()) prev_index = index buffer = create_OutputStream() if (self._write_state is not None and out.size() - start_size > self._write_state_threshold): tail = ( value_iter[index + 1:] if isinstance(value_iter, (list, tuple)) else value_iter) state_token = self._write_state(tail, self._elem_coder) out.write_var_int64(-1) out.write(state_token, True) break else: if index > prev_index: out.write_var_int64(index - prev_index) out.write(buffer.get()) out.write_var_int64(0) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> Sequence size = in_stream.read_bigendian_int32() if size >= 0: elements = [ self._elem_coder.decode_from_stream(in_stream, True) for _ in range(size) ] # type: Iterable[Any] else: elements = [] count = in_stream.read_var_int64() while count > 0: for _ in range(count): elements.append(self._elem_coder.decode_from_stream(in_stream, True)) count = in_stream.read_var_int64() if count == -1: if self._read_state is None: raise ValueError( 'Cannot read state-written iterable without state reader.') state_token = in_stream.read_all(True) elements = _ConcatSequence( elements, self._read_state(state_token, self._elem_coder)) return self._construct_from_sequence(elements) def estimate_size(self, value, nested=False): # type: (Any, bool) -> int """Estimates the encoded size of the given value, in bytes.""" # TODO(ccy): This ignores element sizes. estimated_size, _ = (self.get_estimated_size_and_observables(value)) return estimated_size def get_estimated_size_and_observables(self, value, nested=False): # type: (Any, bool) -> Tuple[int, Observables] """Returns estimated size of value along with any nested observables.""" estimated_size = 0 # Size of 32-bit integer storing number of elements. estimated_size += 4 if isinstance(value, observable.ObservableMixin): return estimated_size, [(value, self._elem_coder)] observables = [] # type: Observables for elem in value: child_size, child_observables = ( self._elem_coder.get_estimated_size_and_observables( elem, nested=True)) estimated_size += child_size observables += child_observables # TODO: (https://github.com/apache/beam/issues/18169) Update to use an # accurate count depending on size and count, currently we are # underestimating the size by up to 10 bytes per block of data since we are # not including the count prefix which occurs at most once per 64k of data # and is upto 10 bytes long. The upper bound of the underestimate is # 10 / 65536 ~= 0.0153% of the actual size. # TODO: More efficient size estimation in the case of state-backed # iterables. return estimated_size, observables class TupleSequenceCoderImpl(SequenceCoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder for homogeneous tuple objects.""" def _construct_from_sequence(self, components): return tuple(components) class _AbstractIterable(object): """Wraps an iterable hiding methods that might not always be available.""" def __init__(self, contents): self._contents = contents def __iter__(self): return iter(self._contents) def __repr__(self): head = [repr(e) for e in itertools.islice(self, 4)] if len(head) == 4: head[-1] = '...' return '_AbstractIterable([%s])' % ', '.join(head) # Mostly useful for tests. def __eq__(left, right): end = object() for a, b in itertools.zip_longest(left, right, fillvalue=end): if a != b: return False return True FastPrimitivesCoderImpl.register_iterable_like_type(_AbstractIterable) # TODO(https://github.com/apache/beam/issues/21167): Enable using abstract # iterables permanently _iterable_coder_uses_abstract_iterable_by_default = False class IterableCoderImpl(SequenceCoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder for homogeneous iterable objects.""" def __init__(self, *args, use_abstract_iterable=None, **kwargs): super().__init__(*args, **kwargs) if use_abstract_iterable is None: use_abstract_iterable = _iterable_coder_uses_abstract_iterable_by_default self._use_abstract_iterable = use_abstract_iterable def _construct_from_sequence(self, components): if self._use_abstract_iterable: return _AbstractIterable(components) else: return components class ListCoderImpl(SequenceCoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder for homogeneous list objects.""" def _construct_from_sequence(self, components): return components if isinstance(components, list) else list(components) class PaneInfoEncoding(object): """For internal use only; no backwards-compatibility guarantees. Encoding used to describe a PaneInfo descriptor. A PaneInfo descriptor can be encoded in three different ways: with a single byte (FIRST), with a single byte followed by a varint describing a single index (ONE_INDEX) or with a single byte followed by two varints describing two separate indices: the index and nonspeculative index. """ FIRST = 0 ONE_INDEX = 1 TWO_INDICES = 2 # These are cdef'd to ints to optimized the common case. PaneInfoTiming_UNKNOWN = windowed_value.PaneInfoTiming.UNKNOWN PaneInfoEncoding_FIRST = PaneInfoEncoding.FIRST class PaneInfoCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. Coder for a PaneInfo descriptor.""" def _choose_encoding(self, value): if ((value._index == 0 and value._nonspeculative_index == 0) or value._timing == PaneInfoTiming_UNKNOWN): return PaneInfoEncoding_FIRST elif (value._index == value._nonspeculative_index or value._timing == windowed_value.PaneInfoTiming.EARLY): return PaneInfoEncoding.ONE_INDEX else: return PaneInfoEncoding.TWO_INDICES def encode_to_stream(self, value, out, nested): # type: (windowed_value.PaneInfo, create_OutputStream, bool) -> None pane_info = value # cast encoding_type = self._choose_encoding(pane_info) out.write_byte(pane_info._encoded_byte | (encoding_type << 4)) if encoding_type == PaneInfoEncoding_FIRST: return elif encoding_type == PaneInfoEncoding.ONE_INDEX: out.write_var_int64(value.index) elif encoding_type == PaneInfoEncoding.TWO_INDICES: out.write_var_int64(value.index) out.write_var_int64(value.nonspeculative_index) else: raise NotImplementedError('Invalid PaneInfoEncoding: %s' % encoding_type) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> windowed_value.PaneInfo encoded_first_byte = in_stream.read_byte() base = windowed_value._BYTE_TO_PANE_INFO[encoded_first_byte & 0xF] assert base is not None encoding_type = encoded_first_byte >> 4 if encoding_type == PaneInfoEncoding_FIRST: return base elif encoding_type == PaneInfoEncoding.ONE_INDEX: index = in_stream.read_var_int64() if base.timing == windowed_value.PaneInfoTiming.EARLY: nonspeculative_index = -1 else: nonspeculative_index = index elif encoding_type == PaneInfoEncoding.TWO_INDICES: index = in_stream.read_var_int64() nonspeculative_index = in_stream.read_var_int64() else: raise NotImplementedError('Invalid PaneInfoEncoding: %s' % encoding_type) return windowed_value.PaneInfo( base.is_first, base.is_last, base.timing, index, nonspeculative_index) def estimate_size(self, value, nested=False): # type: (Any, bool) -> int """Estimates the encoded size of the given value, in bytes.""" size = 1 encoding_type = self._choose_encoding(value) if encoding_type == PaneInfoEncoding.ONE_INDEX: size += get_varint_size(value.index) elif encoding_type == PaneInfoEncoding.TWO_INDICES: size += get_varint_size(value.index) size += get_varint_size(value.nonspeculative_index) return size class _OrderedUnionCoderImpl(StreamCoderImpl): def __init__(self, coder_impl_types, fallback_coder_impl): assert len(coder_impl_types) < 128 self._types, self._coder_impls = zip(*coder_impl_types) self._fallback_coder_impl = fallback_coder_impl def encode_to_stream(self, value, out, nested): value_t = type(value) for (ix, t) in enumerate(self._types): if value_t is t: out.write_byte(ix) c = self._coder_impls[ix] # for typing c.encode_to_stream(value, out, nested) break else: if self._fallback_coder_impl is None: raise ValueError("No fallback.") out.write_byte(0xFF) self._fallback_coder_impl.encode_to_stream(value, out, nested) def decode_from_stream(self, in_stream, nested): ix = in_stream.read_byte() if ix == 0xFF: if self._fallback_coder_impl is None: raise ValueError("No fallback.") return self._fallback_coder_impl.decode_from_stream(in_stream, nested) else: c = self._coder_impls[ix] # for typing return c.decode_from_stream(in_stream, nested) class WindowedValueCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder for windowed values.""" # Ensure that lexicographic ordering of the bytes corresponds to # chronological order of timestamps. # TODO(https://github.com/apache/beam/issues/18190): Clean this up once we # have a BEAM wide consensus on byte representation of timestamps. def _to_normal_time(self, value): """Convert "lexicographically ordered unsigned" to signed.""" return value - _TIME_SHIFT def _from_normal_time(self, value): """Convert signed to "lexicographically ordered unsigned".""" return value + _TIME_SHIFT def __init__(self, value_coder, timestamp_coder, window_coder): # TODO(lcwik): Remove the timestamp coder field self._value_coder = value_coder self._timestamp_coder = timestamp_coder self._windows_coder = TupleSequenceCoderImpl(window_coder) self._pane_info_coder = PaneInfoCoderImpl() def encode_to_stream(self, value, out, nested): # type: (windowed_value.WindowedValue, create_OutputStream, bool) -> None wv = value # type cast # Avoid creation of Timestamp object. restore_sign = -1 if wv.timestamp_micros < 0 else 1 out.write_bigendian_uint64( # Convert to postive number and divide, since python rounds off to the # lower negative number. For ex: -3 / 2 = -2, but we expect it to be -1, # to be consistent across SDKs. # TODO(https://github.com/apache/beam/issues/18190): Clean this up once # we have a BEAM wide consensus on precision of timestamps. self._from_normal_time( restore_sign * ( abs( MIN_TIMESTAMP_micros if wv.timestamp_micros < MIN_TIMESTAMP_micros else wv.timestamp_micros) // 1000))) self._windows_coder.encode_to_stream(wv.windows, out, True) # Default PaneInfo encoded byte representing NO_FIRING. self._pane_info_coder.encode_to_stream(wv.pane_info, out, True) self._value_coder.encode_to_stream(wv.value, out, nested) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> windowed_value.WindowedValue timestamp = self._to_normal_time(in_stream.read_bigendian_uint64()) # Restore MIN/MAX timestamps to their actual values as encoding incurs loss # of precision while converting to millis. # Note: This is only a best effort here as there is no way to know if these # were indeed MIN/MAX timestamps. # TODO(https://github.com/apache/beam/issues/18190): Clean this up once we # have a BEAM wide consensus on precision of timestamps. if timestamp <= -(abs(MIN_TIMESTAMP_micros) // 1000): timestamp = MIN_TIMESTAMP_micros elif timestamp >= MAX_TIMESTAMP_micros // 1000: timestamp = MAX_TIMESTAMP_micros else: timestamp *= 1000 windows = self._windows_coder.decode_from_stream(in_stream, True) # Read PaneInfo encoded byte. pane_info = self._pane_info_coder.decode_from_stream(in_stream, True) value = self._value_coder.decode_from_stream(in_stream, nested) return windowed_value.create( value, timestamp, # Avoid creation of Timestamp object. windows, pane_info) def get_estimated_size_and_observables(self, value, nested=False): # type: (Any, bool) -> Tuple[int, Observables] """Returns estimated size of value along with any nested observables.""" if isinstance(value, observable.ObservableMixin): # Should never be here. # TODO(robertwb): Remove when coders are set correctly. return 0, [(value, self._value_coder)] estimated_size = 0 observables = [] # type: Observables value_estimated_size, value_observables = ( self._value_coder.get_estimated_size_and_observables( value.value, nested=nested)) estimated_size += value_estimated_size observables += value_observables estimated_size += ( self._timestamp_coder.estimate_size(value.timestamp, nested=True)) estimated_size += ( self._windows_coder.estimate_size(value.windows, nested=True)) estimated_size += ( self._pane_info_coder.estimate_size(value.pane_info, nested=True)) return estimated_size, observables class ParamWindowedValueCoderImpl(WindowedValueCoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder for windowed values with constant timestamp, windows and pane info. The coder drops timestamp, windows and pane info during encoding, and uses the supplied parameterized timestamp, windows and pane info values during decoding when reconstructing the windowed value.""" def __init__(self, value_coder, window_coder, payload): super().__init__(value_coder, TimestampCoderImpl(), window_coder) self._timestamp, self._windows, self._pane_info = self._from_proto( payload, window_coder) def _from_proto(self, payload, window_coder): windowed_value_coder = WindowedValueCoderImpl( BytesCoderImpl(), TimestampCoderImpl(), window_coder) wv = windowed_value_coder.decode(payload) return wv.timestamp_micros, wv.windows, wv.pane_info def encode_to_stream(self, value, out, nested): wv = value # type cast self._value_coder.encode_to_stream(wv.value, out, nested) def decode_from_stream(self, in_stream, nested): value = self._value_coder.decode_from_stream(in_stream, nested) return windowed_value.create( value, self._timestamp, self._windows, self._pane_info) def get_estimated_size_and_observables(self, value, nested=False): """Returns estimated size of value along with any nested observables.""" if isinstance(value, observable.ObservableMixin): # Should never be here. # TODO(robertwb): Remove when coders are set correctly. return 0, [(value, self._value_coder)] estimated_size = 0 observables = [] value_estimated_size, value_observables = ( self._value_coder.get_estimated_size_and_observables( value.value, nested=nested)) estimated_size += value_estimated_size observables += value_observables return estimated_size, observables class LengthPrefixCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. Coder which prefixes the length of the encoded object in the stream.""" def __init__(self, value_coder): # type: (CoderImpl) -> None self._value_coder = value_coder def encode_to_stream(self, value, out, nested): # type: (Any, create_OutputStream, bool) -> None encoded_value = self._value_coder.encode(value) out.write_var_int64(len(encoded_value)) out.write(encoded_value) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> Any value_length = in_stream.read_var_int64() return self._value_coder.decode(in_stream.read(value_length)) def estimate_size(self, value, nested=False): # type: (Any, bool) -> int value_size = self._value_coder.estimate_size(value) return get_varint_size(value_size) + value_size class ShardedKeyCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder for sharded user keys. The encoding and decoding should follow the order: shard id byte string encoded user key """ def __init__(self, key_coder_impl): self._shard_id_coder_impl = BytesCoderImpl() self._key_coder_impl = key_coder_impl def encode_to_stream(self, value, out, nested): # type: (ShardedKey, create_OutputStream, bool) -> None self._shard_id_coder_impl.encode_to_stream(value._shard_id, out, True) self._key_coder_impl.encode_to_stream(value.key, out, True) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> ShardedKey shard_id = self._shard_id_coder_impl.decode_from_stream(in_stream, True) key = self._key_coder_impl.decode_from_stream(in_stream, True) return ShardedKey(key=key, shard_id=shard_id) def estimate_size(self, value, nested=False): # type: (Any, bool) -> int estimated_size = 0 estimated_size += ( self._shard_id_coder_impl.estimate_size(value._shard_id, nested=True)) estimated_size += ( self._key_coder_impl.estimate_size(value.key, nested=True)) return estimated_size class TimestampPrefixingWindowCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder for custom window types, which prefix required max_timestamp to encoded original window. The coder encodes and decodes custom window types with following format: window's max_timestamp() encoded window using it's own coder. """ def __init__(self, window_coder_impl: CoderImpl) -> None: self._window_coder_impl = window_coder_impl def encode_to_stream(self, value, stream, nested): TimestampCoderImpl().encode_to_stream(value.max_timestamp(), stream, nested) self._window_coder_impl.encode_to_stream(value, stream, nested) def decode_from_stream(self, stream, nested): TimestampCoderImpl().decode_from_stream(stream, nested) return self._window_coder_impl.decode_from_stream(stream, nested) def estimate_size(self, value: Any, nested: bool = False) -> int: return ( TimestampCoderImpl().estimate_size(value.max_timestamp()) + self._window_coder_impl.estimate_size(value, nested)) _OpaqueWindow = None def _create_opaque_window(end, encoded_window): # This is lazy to avoid circular import issues. global _OpaqueWindow if _OpaqueWindow is None: from apache_beam.transforms.window import BoundedWindow class _OpaqueWindow(BoundedWindow): def __init__(self, end, encoded_window): super().__init__(end) self.encoded_window = encoded_window def __repr__(self): return 'OpaqueWindow(%s, %s)' % (self.end, self.encoded_window) def __hash__(self): return hash(self.encoded_window) def __eq__(self, other): return ( type(self) == type(other) and self.end == other.end and self.encoded_window == other.encoded_window) return _OpaqueWindow(end, encoded_window) class TimestampPrefixingOpaqueWindowCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. A coder for unknown window types, which prefix required max_timestamp to encoded original window. The coder encodes and decodes custom window types with following format: window's max_timestamp() length prefixed encoded window """ def __init__(self) -> None: pass def encode_to_stream(self, value, stream, nested): TimestampCoderImpl().encode_to_stream(value.max_timestamp(), stream, True) stream.write(value.encoded_window, True) def decode_from_stream(self, stream, nested): max_timestamp = TimestampCoderImpl().decode_from_stream(stream, True) return _create_opaque_window( max_timestamp.successor(), stream.read_all(True)) def estimate_size(self, value: Any, nested: bool = False) -> int: return ( TimestampCoderImpl().estimate_size(value.max_timestamp()) + len(value.encoded_window)) row_coders_registered = False class RowColumnEncoder: ROW_ENCODERS = {1: 12345} @classmethod def register(cls, field_type, coder_impl): cls.ROW_ENCODERS[field_type, coder_impl] = cls @classmethod def create(cls, field_type, coder_impl, column): global row_coders_registered if not row_coders_registered: try: # pylint: disable=unused-import from apache_beam.coders import coder_impl_row_encoders except ImportError: pass row_coders_registered = True return cls.ROW_ENCODERS.get((field_type, column.dtype), GenericRowColumnEncoder)(coder_impl, column) def null_flags(self): raise NotImplementedError(type(self)) def encode_to_stream(self, index, out): raise NotImplementedError(type(self)) def decode_from_stream(self, index, in_stream): raise NotImplementedError(type(self)) def finalize_write(self): raise NotImplementedError(type(self)) class GenericRowColumnEncoder(RowColumnEncoder): def __init__(self, coder_impl, column): self.coder_impl = coder_impl self.column = column def null_flags(self): # pylint: disable=singleton-comparison return self.column == None # This is an array. def encode_to_stream(self, index, out): self.coder_impl.encode_to_stream(self.column[index], out, True) def decode_from_stream(self, index, in_stream): self.column[index] = self.coder_impl.decode_from_stream(in_stream, True) def finalize_write(self): pass class RowCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees.""" def __init__(self, schema, components): self.schema = schema self.num_fields = len(self.schema.fields) self.field_names = [f.name for f in self.schema.fields] self.field_nullable = [field.type.nullable for field in self.schema.fields] self.constructor = named_tuple_from_schema(schema) self.encoding_positions = list(range(len(self.schema.fields))) if self.schema.encoding_positions_set: # should never be duplicate encoding positions. enc_posx = list( set(field.encoding_position for field in self.schema.fields)) if len(enc_posx) != len(self.schema.fields): names_no_pos = ", ".join( [f.name for f in self.schema.fields if f.encoding_position is None]) raise ValueError( f'''Schema with id {schema.id} has encoding_positions_set=True, but found fields without encoding_position set: {names_no_pos}''') self.encoding_positions = list( field.encoding_position for field in self.schema.fields) self.encoding_positions_argsort = list(np.argsort(self.encoding_positions)) self.encoding_positions_are_trivial = self.encoding_positions == list( range(len(self.encoding_positions))) self.components = list( components[self.encoding_positions.index(i)].get_impl() for i in self.encoding_positions) self.has_nullable_fields = any( field.type.nullable for field in self.schema.fields) def encode_to_stream(self, value, out, nested): out.write_var_int64(self.num_fields) attrs = [getattr(value, name) for name in self.field_names] if self.has_nullable_fields: any_nulls = False for attr in attrs: if attr is None: any_nulls = True break if any_nulls: out.write_var_int64((self.num_fields + 7) // 8) # Pack the bits, little-endian, in consecutive bytes. running = 0 for i, attr in enumerate(attrs): if i and i % 8 == 0: out.write_byte(running) running = 0 running |= (attr is None) << (i % 8) out.write_byte(running) else: out.write_byte(0) else: out.write_byte(0) for i in range(self.num_fields): if not self.encoding_positions_are_trivial: i = self.encoding_positions_argsort[i] attr = attrs[i] if attr is None: if not self.field_nullable[i]: raise ValueError( "Attempted to encode null for non-nullable field \"{}\".".format( self.schema.fields[i].name)) continue component_coder = self.components[i] # for typing component_coder.encode_to_stream(attr, out, True) def _row_column_encoders(self, columns): return [ RowColumnEncoder.create( self.schema.fields[i].type.atomic_type, self.components[i], columns[name]) for i, name in enumerate(self.field_names) ] def encode_batch_to_stream(self, columns: Dict[str, np.ndarray], out): attrs = self._row_column_encoders(columns) n = len(next(iter(columns.values()))) if self.has_nullable_fields: null_flags_py = np.zeros((n, self.num_fields), dtype=np.uint8) null_bits_len = (self.num_fields + 7) // 8 null_bits_py = np.zeros((n, null_bits_len), dtype=np.uint8) for i, attr in enumerate(attrs): attr_null_flags = attr.null_flags() if attr_null_flags is not None and attr_null_flags.any(): null_flags_py[:, i] = attr_null_flags null_bits_py[:, i // 8] |= attr_null_flags << np.uint8(i % 8) has_null_bits = (null_bits_py.sum(axis=1) != 0).astype(np.uint8) null_bits = null_bits_py null_flags = null_flags_py else: has_null_bits = np.zeros((n, ), dtype=np.uint8) for k in range(n): out.write_var_int64(self.num_fields) if has_null_bits[k]: out.write_byte(null_bits_len) for i in range(null_bits_len): out.write_byte(null_bits[k, i]) else: out.write_byte(0) for i in range(self.num_fields): if not self.encoding_positions_are_trivial: i = self.encoding_positions_argsort[i] if has_null_bits[k] and null_flags[k, i]: if not self.field_nullable[i]: raise ValueError( "Attempted to encode null for non-nullable field \"{}\".". format(self.schema.fields[i].name)) else: cython.cast(RowColumnEncoder, attrs[i]).encode_to_stream(k, out) def decode_from_stream(self, in_stream, nested): nvals = in_stream.read_var_int64() null_mask_len = in_stream.read_var_int64() if null_mask_len: # pylint: disable=unused-variable null_mask_c = null_mask_py = in_stream.read(null_mask_len) # Note that if this coder's schema has *fewer* attributes than the encoded # value, we just need to ignore the additional values, which will occur # here because we only decode as many values as we have coders for. sorted_components = [] for i in range(min(self.num_fields, nvals)): if not self.encoding_positions_are_trivial: i = self.encoding_positions_argsort[i] if (null_mask_len and i >> 3 < null_mask_len and null_mask_c[i >> 3] & (0x01 << (i & 0x07))): item = None else: component_coder = self.components[i] # for typing item = component_coder.decode_from_stream(in_stream, True) sorted_components.append(item) # If this coder's schema has more attributes than the encoded value, then # the schema must have changed. Populate the unencoded fields with nulls. while len(sorted_components) < self.num_fields: sorted_components.append(None) return self.constructor( *( sorted_components if self.encoding_positions_are_trivial else [sorted_components[i] for i in self.encoding_positions])) def decode_batch_from_stream(self, dest: Dict[str, np.ndarray], in_stream): attrs = self._row_column_encoders(dest) n = len(next(iter(dest.values()))) for k in range(n): if in_stream.size() == 0: break nvals = in_stream.read_var_int64() null_mask_len = in_stream.read_var_int64() if null_mask_len: # pylint: disable=unused-variable null_mask_c = null_mask = in_stream.read(null_mask_len) for i in range(min(self.num_fields, nvals)): if not self.encoding_positions_are_trivial: i = self.encoding_positions_argsort[i] if (null_mask_len and i >> 3 < null_mask_len and null_mask_c[i >> 3] & (0x01 << (i & 0x07))): continue else: cython.cast(RowColumnEncoder, attrs[i]).decode_from_stream(k, in_stream) else: # Loop variable will be n-1 on normal exit. k = n for attr in attrs: attr.finalize_write() return k class LogicalTypeCoderImpl(StreamCoderImpl): def __init__(self, logical_type, representation_coder): self.logical_type = logical_type self.representation_coder = representation_coder.get_impl() def encode_to_stream(self, value, out, nested): return self.representation_coder.encode_to_stream( self.logical_type.to_representation_type(value), out, nested) def decode_from_stream(self, in_stream, nested): return self.logical_type.to_language_type( self.representation_coder.decode_from_stream(in_stream, nested)) class BigIntegerCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. For interoperability with Java SDK, encoding needs to match that of the Java SDK BigIntegerCoder.""" def encode_to_stream(self, value, out, nested): # type: (int, create_OutputStream, bool) -> None if value < 0: byte_length = ((value + 1).bit_length() + 8) // 8 else: byte_length = (value.bit_length() + 8) // 8 encoded_value = value.to_bytes( length=byte_length, byteorder='big', signed=True) out.write(encoded_value, nested) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> int encoded_value = in_stream.read_all(nested) return int.from_bytes(encoded_value, byteorder='big', signed=True) class DecimalCoderImpl(StreamCoderImpl): """For internal use only; no backwards-compatibility guarantees. For interoperability with Java SDK, encoding needs to match that of the Java SDK BigDecimalCoder.""" BIG_INT_CODER_IMPL = BigIntegerCoderImpl() def encode_to_stream(self, value, out, nested): # type: (decimal.Decimal, create_OutputStream, bool) -> None scale = -value.as_tuple().exponent # type: ignore[operator] int_value = int(value.scaleb(scale)) out.write_var_int64(scale) self.BIG_INT_CODER_IMPL.encode_to_stream(int_value, out, nested) def decode_from_stream(self, in_stream, nested): # type: (create_InputStream, bool) -> decimal.Decimal scale = in_stream.read_var_int64() int_value = self.BIG_INT_CODER_IMPL.decode_from_stream(in_stream, nested) value = decimal.Decimal(int_value).scaleb(-scale) return value ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/coder_impl_row_encoders.pyx0000644000175100017510000001357315206113032024575 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Optimized implementations of various schema columner types.""" # pytype: skip-file import numpy as np cimport numpy as np from apache_beam.coders import coder_impl from apache_beam.coders.coder_impl cimport RowColumnEncoder, OutputStream, InputStream from apache_beam.portability.api import schema_pb2 cdef class AtomicTypeRowColumnEncoder(RowColumnEncoder): cdef original cdef contiguous def __init__(self, column): self.original = column self.contiguous = np.ascontiguousarray(column) def null_flags(self): return None def finalize_write(self): if self.original is not self.contiguous: self.original[:] = self.contiguous cdef class FloatFloat32RowColumnEncoder(AtomicTypeRowColumnEncoder): cdef np.float32_t* data def __init__(self, unused_coder, column): super(FloatFloat32RowColumnEncoder, self).__init__(column) cdef np.float32_t[::1] view = self.contiguous self.data = &view[0] cdef bint encode_to_stream(self, size_t index, OutputStream stream) except -1: stream.write_bigendian_float(self.data[index]) cdef bint decode_from_stream(self, size_t index, InputStream stream) except -1: self.data[index] = stream.read_bigendian_float() FloatFloat32RowColumnEncoder.register(schema_pb2.FLOAT, np.float32().dtype) cdef class FloatFloat64RowColumnEncoder(AtomicTypeRowColumnEncoder): cdef np.float64_t* data def __init__(self, unused_coder, column): super(FloatFloat64RowColumnEncoder, self).__init__(column) cdef np.float64_t[::1] view = self.contiguous self.data = &view[0] cdef bint encode_to_stream(self, size_t index, OutputStream stream) except -1: stream.write_bigendian_float(self.data[index]) cdef bint decode_from_stream(self, size_t index, InputStream stream) except -1: self.data[index] = stream.read_bigendian_float() FloatFloat64RowColumnEncoder.register(schema_pb2.FLOAT, np.float64().dtype) cdef class DoubleFloat32RowColumnEncoder(AtomicTypeRowColumnEncoder): cdef np.float32_t* data def __init__(self, unused_coder, column): super(DoubleFloat32RowColumnEncoder, self).__init__(column) cdef np.float32_t[::1] view = self.contiguous self.data = &view[0] cdef bint encode_to_stream(self, size_t index, OutputStream stream) except -1: stream.write_bigendian_double(self.data[index]) cdef bint decode_from_stream(self, size_t index, InputStream stream) except -1: self.data[index] = stream.read_bigendian_double() DoubleFloat32RowColumnEncoder.register(schema_pb2.DOUBLE, np.float32().dtype) cdef class DoubleFloat64RowColumnEncoder(AtomicTypeRowColumnEncoder): cdef np.float64_t* data def __init__(self, unused_coder, column): super(DoubleFloat64RowColumnEncoder, self).__init__(column) cdef np.float64_t[::1] view = self.contiguous self.data = &view[0] cdef bint encode_to_stream(self, size_t index, OutputStream stream) except -1: stream.write_bigendian_double(self.data[index]) cdef bint decode_from_stream(self, size_t index, InputStream stream) except -1: self.data[index] = stream.read_bigendian_double() DoubleFloat64RowColumnEncoder.register(schema_pb2.DOUBLE, np.float64().dtype) cdef class Int32Int32RowColumnEncoder(AtomicTypeRowColumnEncoder): cdef np.int32_t* data def __init__(self, unused_coder, column): super(Int32Int32RowColumnEncoder, self).__init__(column) cdef np.int32_t[::1] view = self.contiguous self.data = &view[0] cdef bint encode_to_stream(self, size_t index, OutputStream stream) except -1: stream.write_var_int64(self.data[index]) cdef bint decode_from_stream(self, size_t index, InputStream stream) except -1: self.data[index] = stream.read_var_int64() Int32Int32RowColumnEncoder.register(schema_pb2.INT32, np.int32().dtype) Int32Int32RowColumnEncoder.register(schema_pb2.INT32, np.int64().dtype) cdef class Int64Int64RowColumnEncoder(AtomicTypeRowColumnEncoder): cdef np.int64_t* data def __init__(self, unused_coder, column): super(Int64Int64RowColumnEncoder, self).__init__(column) cdef np.int64_t[::1] view = self.contiguous self.data = &view[0] cdef bint encode_to_stream(self, size_t index, OutputStream stream) except -1: stream.write_var_int64(self.data[index]) cdef bint decode_from_stream(self, size_t index, InputStream stream) except -1: self.data[index] = stream.read_var_int64() Int64Int64RowColumnEncoder.register(schema_pb2.INT64, np.int64().dtype) cdef class BoolRowColumnEncoder(AtomicTypeRowColumnEncoder): cdef np.uint8_t* data def __init__(self, unused_coder, column): super(BoolRowColumnEncoder, self).__init__(column) self.contiguous = self.contiguous.astype(np.uint8) cdef np.uint8_t[::1] view = self.contiguous self.data = &view[0] cdef bint encode_to_stream(self, size_t index, OutputStream stream) except -1: stream.write_byte(self.data[index]) cdef bint decode_from_stream(self, size_t index, InputStream stream) except -1: self.data[index] = stream.read_byte() BoolRowColumnEncoder.register(schema_pb2.BOOLEAN, np.int8().dtype) BoolRowColumnEncoder.register(schema_pb2.BOOLEAN, np.uint8().dtype) BoolRowColumnEncoder.register(schema_pb2.BOOLEAN, np.bool_().dtype) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/coders.py0000644000175100017510000015547415206113032021005 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Collection of useful coders. Only those coders listed in __all__ are part of the public API of this module. ## On usage of `pickle`, `dill` and `pickler` in Beam In Beam, we generally we use `pickle` for pipeline elements and `dill` for more complex types, like user functions. `pickler` is Beam's own wrapping of dill + compression + error handling. It serves also as an API to mask the actual encoding layer (so we can change it from `dill` if necessary). We created `_MemoizingPickleCoder` to improve performance when serializing complex user types for the execution of SDF. Specifically to address BEAM-12781, where many identical `BoundedSource` instances are being encoded. """ # pytype: skip-file import base64 import decimal import pickle from functools import lru_cache from typing import TYPE_CHECKING from typing import Any from typing import Callable from typing import Iterable from typing import Optional from typing import Sequence from typing import TypeVar from typing import overload import google.protobuf.wrappers_pb2 import proto from google.protobuf import message from apache_beam.coders import coder_impl from apache_beam.coders.avro_record import AvroRecord from apache_beam.internal import cloudpickle_pickler from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.typehints import typehints from apache_beam.utils import proto_utils from apache_beam.utils import windowed_value if TYPE_CHECKING: from apache_beam.runners.pipeline_context import PipelineContext # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from .stream import get_varint_size except ImportError: from .slow_stream import get_varint_size # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports # pylint: disable=wrong-import-order, wrong-import-position # Avoid dependencies on the full SDK. try: # Import dill from the pickler module to make sure our monkey-patching of dill # occurs. from apache_beam.internal.dill_pickler import dill except ImportError: dill = None __all__ = [ 'Coder', 'AvroGenericCoder', 'BooleanCoder', 'ByteCoder', 'BytesCoder', 'CloudpickleCoder', 'DillCoder', 'FastPrimitivesCoder', 'FloatCoder', 'IterableCoder', 'ListCoder', 'MapCoder', 'NullableCoder', 'PickleCoder', 'ProtoCoder', 'ProtoPlusCoder', 'ShardedKeyCoder', 'SinglePrecisionFloatCoder', 'SingletonCoder', 'StrUtf8Coder', 'TimestampCoder', 'TupleCoder', 'TupleSequenceCoder', 'VarIntCoder', 'WindowedValueCoder', 'ParamWindowedValueCoder', 'BigIntegerCoder', 'DecimalCoder', 'PaneInfoCoder' ] T = TypeVar('T') CoderT = TypeVar('CoderT', bound='Coder') ProtoCoderT = TypeVar('ProtoCoderT', bound='ProtoCoder') ConstructorFn = Callable[[Optional[Any], list['Coder'], 'PipelineContext'], Any] def serialize_coder(coder): from apache_beam.internal import pickler return b'%s$%s' % ( coder.__class__.__name__.encode('utf-8'), pickler.dumps(coder, use_zlib=True)) def deserialize_coder(serialized): from apache_beam.internal import pickler return pickler.loads(serialized.split(b'$', 1)[1], use_zlib=True) # pylint: enable=wrong-import-order, wrong-import-position class Coder(object): """Base class for coders.""" def encode(self, value): # type: (Any) -> bytes """Encodes the given object into a byte string.""" raise NotImplementedError('Encode not implemented: %s.' % self) def decode(self, encoded): """Decodes the given byte string into the corresponding object.""" raise NotImplementedError('Decode not implemented: %s.' % self) def encode_nested(self, value): """Uses the underlying implementation to encode in nested format.""" return self.get_impl().encode_nested(value) def decode_nested(self, encoded): """Uses the underlying implementation to decode in nested format.""" return self.get_impl().decode_nested(encoded) def is_deterministic(self): # type: () -> bool """Whether this coder is guaranteed to encode values deterministically. A deterministic coder is required for key coders in GroupByKey operations to produce consistent results. For example, note that the default coder, the PickleCoder, is not deterministic: the ordering of picked entries in maps may vary across executions since there is no defined order, and such a coder is not in general suitable for usage as a key coder in GroupByKey operations, since each instance of the same key may be encoded differently. Returns: Whether coder is deterministic. """ return False def as_deterministic_coder(self, step_label, error_message=None): """Returns a deterministic version of self, if possible. Otherwise raises a value error. """ if self.is_deterministic(): return self else: raise ValueError( error_message or "%s cannot be made deterministic for '%s'." % (self, step_label)) def estimate_size(self, value): """Estimates the encoded size of the given value, in bytes. Dataflow estimates the encoded size of a PCollection processed in a pipeline step by using the estimated size of a random sample of elements in that PCollection. The default implementation encodes the given value and returns its byte size. If a coder can provide a fast estimate of the encoded size of a value (e.g., if the encoding has a fixed size), it can provide its estimate here to improve performance. Arguments: value: the value whose encoded size is to be estimated. Returns: The estimated encoded size of the given value. """ return len(self.encode(value)) # =========================================================================== # Methods below are internal SDK details that don't need to be modified for # user-defined coders. # =========================================================================== def _create_impl(self): # type: () -> coder_impl.CoderImpl """Creates a CoderImpl to do the actual encoding and decoding. """ return coder_impl.CallbackCoderImpl( self.encode, self.decode, self.estimate_size) def get_impl(self): """For internal use only; no backwards-compatibility guarantees. Returns the CoderImpl backing this Coder. """ if not hasattr(self, '_impl'): self._impl = self._create_impl() assert isinstance(self._impl, coder_impl.CoderImpl) return self._impl def __getstate__(self): return self._dict_without_impl() def _dict_without_impl(self): if hasattr(self, '_impl'): d = dict(self.__dict__) del d['_impl'] return d return self.__dict__ def to_type_hint(self): # TODO: After https://github.com/apache/beam/issues/18490 we should be # able to infer the type hint rather than require every subclass define # it raise NotImplementedError @classmethod def from_type_hint(cls, unused_typehint, unused_registry): # type: (Type[CoderT], Any, CoderRegistry) -> CoderT # If not overridden, just construct the coder without arguments. return cls() def is_kv_coder(self): # type: () -> bool return False def key_coder(self): # type: () -> Coder if self.is_kv_coder(): raise NotImplementedError('key_coder: %s' % self) else: raise ValueError('Not a KV coder: %s.' % self) def value_coder(self): # type: () -> Coder if self.is_kv_coder(): raise NotImplementedError('value_coder: %s' % self) else: raise ValueError('Not a KV coder: %s.' % self) def _get_component_coders(self): # type: () -> Sequence[Coder] """For internal use only; no backwards-compatibility guarantees. Returns the internal component coders of this coder.""" # This is an internal detail of the Coder API and does not need to be # refined in user-defined Coders. return [] def __repr__(self): return self.__class__.__name__ # pylint: disable=protected-access def __eq__(self, other): return ( self.__class__ == other.__class__ and self._dict_without_impl() == other._dict_without_impl()) # pylint: enable=protected-access def __hash__(self): return hash(type(self)) _known_urns = {} # type: Dict[str, Tuple[type, ConstructorFn]] @classmethod @overload def register_urn( cls, urn, # type: str parameter_type, # type: Optional[Type[T]] ): # type: (...) -> Callable[[Callable[[T, List[Coder], PipelineContext], Any]], Callable[[T, List[Coder], PipelineContext], Any]] pass @classmethod @overload def register_urn( cls, urn, # type: str parameter_type, # type: Optional[Type[T]] fn # type: Callable[[T, List[Coder], PipelineContext], Any] ): # type: (...) -> None pass @classmethod def register_urn(cls, urn, parameter_type, fn=None): """Registers a urn with a constructor. For example, if 'beam:fn:foo' had parameter type FooPayload, one could write `RunnerApiFn.register_urn('bean:fn:foo', FooPayload, foo_from_proto)` where foo_from_proto took as arguments a FooPayload and a PipelineContext. This function can also be used as a decorator rather than passing the callable in as the final parameter. A corresponding to_runner_api_parameter method would be expected that returns the tuple ('beam:fn:foo', FooPayload) """ def register(fn): cls._known_urns[urn] = parameter_type, fn return fn if fn: # Used as a statement. register(fn) else: # Used as a decorator. return register def to_runner_api(self, context): # type: (PipelineContext) -> beam_runner_api_pb2.Coder urn, typed_param, components = self.to_runner_api_parameter(context) return beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec( urn=urn, payload=typed_param if isinstance(typed_param, (bytes, type(None))) else typed_param.SerializeToString()), component_coder_ids=[context.coders.get_id(c) for c in components]) @classmethod def from_runner_api(cls, coder_proto, context): # type: (Type[CoderT], beam_runner_api_pb2.Coder, PipelineContext) -> CoderT """Converts from an FunctionSpec to a Fn object. Prefer registering a urn with its parameter type and constructor. """ parameter_type, constructor = cls._known_urns[coder_proto.spec.urn] return constructor( proto_utils.parse_Bytes(coder_proto.spec.payload, parameter_type), [context.coders.get_by_id(c) for c in coder_proto.component_coder_ids], context) def to_runner_api_parameter(self, context): # type: (Optional[PipelineContext]) -> Tuple[str, Any, Sequence[Coder]] return ( python_urns.PICKLED_CODER, google.protobuf.wrappers_pb2.BytesValue(value=serialize_coder(self)), ()) @staticmethod def register_structured_urn(urn, cls): # type: (str, Type[Coder]) -> None """Register a coder that's completely defined by its urn and its component(s), if any, which are passed to construct the instance. """ setattr( cls, 'to_runner_api_parameter', lambda self, unused_context: (urn, None, self._get_component_coders())) # pylint: disable=unused-variable @Coder.register_urn(urn, None) def from_runner_api_parameter(unused_payload, components, unused_context): if components: return cls(*components) else: return cls() def version_tag(self) -> str: """For internal use. Appends a version tag to the coder key in the pipeline proto. Some runners (e.g. DataflowRunner) use coder key/id to verify if a pipeline is update compatible. If the implementation of a coder changed in an update incompatible way a version tag can be added to fail compatibility checks. """ return "" @Coder.register_urn( python_urns.PICKLED_CODER, google.protobuf.wrappers_pb2.BytesValue) def _pickle_from_runner_api_parameter(payload, components, context): return deserialize_coder(payload.value) class StrUtf8Coder(Coder): """A coder used for reading and writing strings as UTF-8.""" def encode(self, value): return value.encode('utf-8') def decode(self, value): return value.decode('utf-8') def is_deterministic(self): # type: () -> bool return True def to_type_hint(self): return str Coder.register_structured_urn(common_urns.coders.STRING_UTF8.urn, StrUtf8Coder) class ToBytesCoder(Coder): """A default string coder used if no sink coder is specified.""" def encode(self, value): return value if isinstance(value, bytes) else str(value).encode('utf-8') def decode(self, _): raise NotImplementedError('ToBytesCoder cannot be used for decoding.') def is_deterministic(self): # type: () -> bool return True # alias to the old class name for a courtesy to users who reference it ToStringCoder = ToBytesCoder class FastCoder(Coder): """Coder subclass used when a (faster) CoderImpl is supplied directly. The Coder class defines _create_impl in terms of encode() and decode(); this class inverts that by defining encode() and decode() in terms of _create_impl(). """ def encode(self, value): """Encodes the given object into a byte string.""" return self.get_impl().encode(value) def decode(self, encoded): """Decodes the given byte string into the corresponding object.""" return self.get_impl().decode(encoded) def estimate_size(self, value): return self.get_impl().estimate_size(value) def _create_impl(self): raise NotImplementedError class BytesCoder(FastCoder): """Byte string coder.""" def _create_impl(self): return coder_impl.BytesCoderImpl() def is_deterministic(self): # type: () -> bool return True def to_type_hint(self): return bytes def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) Coder.register_structured_urn(common_urns.coders.BYTES.urn, BytesCoder) class BooleanCoder(FastCoder): def _create_impl(self): return coder_impl.BooleanCoderImpl() def is_deterministic(self): # type: () -> bool return True def to_type_hint(self): return bool def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) Coder.register_structured_urn(common_urns.coders.BOOL.urn, BooleanCoder) class MapCoder(FastCoder): def __init__(self, key_coder, value_coder): # type: (Coder, Coder) -> None self._key_coder = key_coder self._value_coder = value_coder def _create_impl(self): return coder_impl.MapCoderImpl( self._key_coder.get_impl(), self._value_coder.get_impl()) @classmethod def from_type_hint(cls, typehint, registry): # type: (typehints.DictConstraint, CoderRegistry) -> MapCoder return cls( registry.get_coder(typehint.key_type), registry.get_coder(typehint.value_type)) def to_type_hint(self): return typehints.Dict[self._key_coder.to_type_hint(), self._value_coder.to_type_hint()] def is_deterministic(self): # type: () -> bool # Map ordering is non-deterministic return False def as_deterministic_coder(self, step_label, error_message=None): return DeterministicMapCoder( self._key_coder.as_deterministic_coder(step_label, error_message), self._value_coder.as_deterministic_coder(step_label, error_message)) def __eq__(self, other): return ( type(self) == type(other) and self._key_coder == other._key_coder and self._value_coder == other._value_coder) def __hash__(self): return hash(type(self)) + hash(self._key_coder) + hash(self._value_coder) def __repr__(self): return 'MapCoder[%s, %s]' % (self._key_coder, self._value_coder) # This is a separate class from MapCoder as the former is a standard coder with # no way to carry the is_deterministic bit. class DeterministicMapCoder(FastCoder): def __init__(self, key_coder, value_coder): # type: (Coder, Coder) -> None assert key_coder.is_deterministic() assert value_coder.is_deterministic() self._key_coder = key_coder self._value_coder = value_coder def _create_impl(self): return coder_impl.MapCoderImpl( self._key_coder.get_impl(), self._value_coder.get_impl(), True) def is_deterministic(self): return True def __eq__(self, other): return ( type(self) == type(other) and self._key_coder == other._key_coder and self._value_coder == other._value_coder) def __hash__(self): return hash(type(self)) + hash(self._key_coder) + hash(self._value_coder) def __repr__(self): return 'DeterministicMapCoder[%s, %s]' % ( self._key_coder, self._value_coder) class NullableCoder(FastCoder): def __init__(self, value_coder): # type: (Coder) -> None self._value_coder = value_coder def _create_impl(self): return coder_impl.NullableCoderImpl(self._value_coder.get_impl()) def to_type_hint(self): return typehints.Optional[self._value_coder.to_type_hint()] def _get_component_coders(self): # type: () -> List[Coder] return [self._value_coder] @classmethod def from_type_hint(cls, typehint, registry): if typehints.is_nullable(typehint): return cls( registry.get_coder( typehints.get_concrete_type_from_nullable(typehint))) else: raise TypeError( 'Typehint is not of nullable type, ' 'and cannot be converted to a NullableCoder', typehint) def is_deterministic(self): # type: () -> bool return self._value_coder.is_deterministic() def as_deterministic_coder(self, step_label, error_message=None): if self.is_deterministic(): return self else: deterministic_value_coder = self._value_coder.as_deterministic_coder( step_label, error_message) return NullableCoder(deterministic_value_coder) def __eq__(self, other): return ( type(self) == type(other) and self._value_coder == other._value_coder) def __hash__(self): return hash(type(self)) + hash(self._value_coder) def __repr__(self): return 'NullableCoder[%s]' % self._value_coder Coder.register_structured_urn(common_urns.coders.NULLABLE.urn, NullableCoder) class VarIntCoder(FastCoder): """Variable-length integer coder matches Java SDK's VarLongCoder.""" def _create_impl(self): return coder_impl.VarIntCoderImpl() def is_deterministic(self): # type: () -> bool return True def to_type_hint(self): return int def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) Coder.register_structured_urn(common_urns.coders.VARINT.urn, VarIntCoder) class VarInt32Coder(FastCoder): """Variable-length integer coder matches Java SDK's VarIntCoder.""" def _create_impl(self): return coder_impl.VarInt32CoderImpl() def is_deterministic(self): # type: () -> bool return True def to_type_hint(self): return int def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) class BigEndianShortCoder(FastCoder): """A coder used for big-endian int16 values.""" def _create_impl(self): return coder_impl.BigEndianShortCoderImpl() def is_deterministic(self): # type: () -> bool return True def to_type_hint(self): return int def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) class ByteCoder(FastCoder): """A coder used for single byte values""" def _create_impl(self): return coder_impl.ByteCoderImpl() def is_deterministic(self): # type: () -> bool return True def to_type_hint(self): return int def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) class SinglePrecisionFloatCoder(FastCoder): """A coder used for single-precision floating-point values.""" def _create_impl(self): return coder_impl.SinglePrecisionFloatCoderImpl() def is_deterministic(self): # type: () -> bool return True def to_type_hint(self): return float def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) class FloatCoder(FastCoder): """A coder used for **double-precision** floating-point values. Note that the name "FloatCoder" is in reference to Python's ``float`` built-in which is generally implemented using C doubles. See :class:`SinglePrecisionFloatCoder` for a single-precision version of this coder. """ def _create_impl(self): return coder_impl.FloatCoderImpl() def is_deterministic(self): # type: () -> bool return True def to_type_hint(self): return float def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) Coder.register_structured_urn(common_urns.coders.DOUBLE.urn, FloatCoder) class TimestampCoder(FastCoder): """A coder used for timeutil.Timestamp values.""" def _create_impl(self): return coder_impl.TimestampCoderImpl() def is_deterministic(self): # type: () -> bool return True def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) class _TimerCoder(FastCoder): """A coder used for timer values. For internal use.""" def __init__(self, key_coder, window_coder): # type: (Coder, Coder) -> None self._key_coder = key_coder self._window_coder = window_coder def _get_component_coders(self): # type: () -> List[Coder] return [self._key_coder, self._window_coder] def _create_impl(self): return coder_impl.TimerCoderImpl( self._key_coder.get_impl(), self._window_coder.get_impl()) def is_deterministic(self): # type: () -> bool return ( self._key_coder.is_deterministic() and self._window_coder.is_deterministic()) def __eq__(self, other): return ( type(self) == type(other) and self._key_coder == other._key_coder and self._window_coder == other._window_coder) def __hash__(self): return hash(type(self)) + hash(self._key_coder) + hash(self._window_coder) Coder.register_structured_urn(common_urns.coders.TIMER.urn, _TimerCoder) class SingletonCoder(FastCoder): """A coder that always encodes exactly one value.""" def __init__(self, value): self._value = value def _create_impl(self): return coder_impl.SingletonCoderImpl(self._value) def is_deterministic(self): # type: () -> bool return True def __eq__(self, other): return type(self) == type(other) and self._value == other._value def __hash__(self): return hash(self._value) def maybe_dill_dumps(o): """Pickle using cPickle or the Dill pickler as a fallback.""" # We need to use the dill pickler for objects of certain custom classes, # including, for example, ones that contain lambdas. try: return pickle.dumps(o, pickle.HIGHEST_PROTOCOL) except Exception: # pylint: disable=broad-except return dill.dumps(o) def maybe_dill_loads(o): """Unpickle using cPickle or the Dill pickler as a fallback.""" try: return pickle.loads(o) except Exception: # pylint: disable=broad-except return dill.loads(o) class _PickleCoderBase(FastCoder): """Base class for pickling coders.""" def is_deterministic(self): # type: () -> bool # Note that the default coder, the PickleCoder, is not deterministic (for # example, the ordering of picked entries in maps may vary across # executions), and so is not in general suitable for usage as a key coder in # GroupByKey operations. return False # We allow .key_coder() and .value_coder() to be called on PickleCoder since # we can't always infer the return values of lambdas in ParDo operations, the # result of which may be used in a GroupBykey. def is_kv_coder(self): # type: () -> bool return True def key_coder(self): return self def value_coder(self): return self def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) class _MemoizingPickleCoder(_PickleCoderBase): """Coder using Python's pickle functionality with memoization.""" def __init__(self, cache_size=16): super().__init__() self.cache_size = cache_size def _create_impl(self): from apache_beam.internal import pickler dumps = pickler.dumps mdumps = lru_cache(maxsize=self.cache_size, typed=True)(dumps) def _nonhashable_dumps(x): try: return mdumps(x) except TypeError: return dumps(x) return coder_impl.CallbackCoderImpl(_nonhashable_dumps, pickler.loads) def as_deterministic_coder(self, step_label, error_message=None): return FastPrimitivesCoder(self, requires_deterministic=step_label) def to_type_hint(self): return Any class PickleCoder(_PickleCoderBase): """Coder using Python's pickle functionality.""" def _create_impl(self): dumps = pickle.dumps protocol = pickle.HIGHEST_PROTOCOL return coder_impl.CallbackCoderImpl( lambda x: dumps(x, protocol), pickle.loads) def as_deterministic_coder(self, step_label, error_message=None): return FastPrimitivesCoder(self, requires_deterministic=step_label) def to_type_hint(self): return Any class DillCoder(_PickleCoderBase): """Coder using dill's pickle functionality.""" def __init__(self): if not dill: raise RuntimeError( "This pipeline contains a DillCoder which requires " "the dill package. Install the dill package with the dill extra " "e.g. apache-beam[dill]") def _create_impl(self): return coder_impl.CallbackCoderImpl(maybe_dill_dumps, maybe_dill_loads) class CloudpickleCoder(_PickleCoderBase): """Coder using Apache Beam's vendored Cloudpickle pickler.""" def _create_impl(self): return coder_impl.CallbackCoderImpl( cloudpickle_pickler.dumps, cloudpickle_pickler.loads) class DeterministicFastPrimitivesCoderV2(FastCoder): """Throws runtime errors when encoding non-deterministic values.""" def __init__(self, coder, step_label): self._underlying_coder = coder self._step_label = step_label self._use_relative_filepaths = True self._version_tag = "v2_69" # Versions prior to 2.69.0 did not use relative filepaths. from apache_beam.options.pipeline_options_context import get_pipeline_options opts = get_pipeline_options() if opts and opts.is_compat_version_prior_to("2.69.0"): self._version_tag = "" self._use_relative_filepaths = False def _create_impl(self): return coder_impl.FastPrimitivesCoderImpl( self._underlying_coder.get_impl(), requires_deterministic_step_label=self._step_label, force_use_dill=False, use_relative_filepaths=self._use_relative_filepaths) def is_deterministic(self): # type: () -> bool return True def is_kv_coder(self): # type: () -> bool return True def key_coder(self): return self def value_coder(self): return self def to_type_hint(self): return Any def to_runner_api_parameter(self, context): # type: (Optional[PipelineContext]) -> Tuple[str, Any, Sequence[Coder]] return ( python_urns.PICKLED_CODER, google.protobuf.wrappers_pb2.BytesValue(value=serialize_coder(self)), ()) def version_tag(self): return self._version_tag class DeterministicFastPrimitivesCoder(FastCoder): """Throws runtime errors when encoding non-deterministic values.""" def __init__(self, coder, step_label): self._underlying_coder = coder self._step_label = step_label def _create_impl(self): return coder_impl.FastPrimitivesCoderImpl( self._underlying_coder.get_impl(), requires_deterministic_step_label=self._step_label, force_use_dill=True) def is_deterministic(self): # type: () -> bool return True def is_kv_coder(self): # type: () -> bool return True def key_coder(self): return self def value_coder(self): return self def to_type_hint(self): return Any def _should_force_use_dill(): from apache_beam.options.pipeline_options_context import get_pipeline_options opts = get_pipeline_options() if opts is None or not opts.is_compat_version_prior_to("2.68.0"): return False try: import dill assert dill.__version__ == "0.3.1.1" except Exception as e: raise RuntimeError("This pipeline runs with the pipeline option " \ "--update_compatibility_version=2.67.0 or earlier. When running with " \ "this option on SDKs 2.68.0 or later, you must ensure dill==0.3.1.1 " \ f"is installed. Error {e}") return True def _update_compatible_deterministic_fast_primitives_coder(coder, step_label): """ Returns the update compatible version of DeterministicFastPrimitivesCoder The differences are in how "special types" e.g. NamedTuples, Dataclasses are deterministically encoded. - In SDK version <= 2.67.0 dill is used to encode "special types" - In SDK version 2.68.0 cloudpickle is used to encode "special types" with absolute filepaths in code objects and dynamic functions. - In SDK version 2.69.0 cloudpickle is used to encode "special types" with relative filepaths in code objects and dynamic functions. """ if _should_force_use_dill(): return DeterministicFastPrimitivesCoder(coder, step_label) return DeterministicFastPrimitivesCoderV2(coder, step_label) class FastPrimitivesCoder(FastCoder): """Encodes simple primitives (e.g. str, int) efficiently. For unknown types, falls back to another coder (e.g. PickleCoder). """ def __init__(self, fallback_coder=PickleCoder()): # type: (Coder) -> None self._fallback_coder = fallback_coder def _create_impl(self): return coder_impl.FastPrimitivesCoderImpl(self._fallback_coder.get_impl()) def is_deterministic(self): # type: () -> bool return self._fallback_coder.is_deterministic() def as_deterministic_coder(self, step_label, error_message=None): if self.is_deterministic(): return self else: return _update_compatible_deterministic_fast_primitives_coder( self, step_label) def to_type_hint(self): return Any # We allow .key_coder() and .value_coder() to be called on FastPrimitivesCoder # since we can't always infer the return values of lambdas in ParDo # operations, the result of which may be used in a GroupBykey. def is_kv_coder(self): # type: () -> bool return True def key_coder(self): return self def value_coder(self): return self def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) class FakeDeterministicFastPrimitivesCoder(FastPrimitivesCoder): """A FastPrimitivesCoder that claims to be deterministic. This can be registered as a fallback coder to go back to the behavior before deterministic encoding was enforced (BEAM-11719). """ def is_deterministic(self): return True class Base64PickleCoder(Coder): """Coder of objects by Python pickle, then base64 encoding.""" # TODO(robertwb): Do base64 encoding where it's needed (e.g. in json) rather # than via a special Coder. def encode(self, value): return base64.b64encode(pickle.dumps(value, pickle.HIGHEST_PROTOCOL)) def decode(self, encoded): return pickle.loads(base64.b64decode(encoded)) def is_deterministic(self): # type: () -> bool # Note that the Base64PickleCoder is not deterministic. See the # corresponding comments for PickleCoder above. return False # We allow .key_coder() and .value_coder() to be called on Base64PickleCoder # since we can't always infer the return values of lambdas in ParDo # operations, the result of which may be used in a GroupBykey. # # TODO(ccy): this is currently only used for KV values from Create transforms. # Investigate a way to unify this with PickleCoder. def is_kv_coder(self): # type: () -> bool return True def key_coder(self): return self def value_coder(self): return self class ProtoCoder(FastCoder): """A Coder for Google Protocol Buffers. It supports both Protocol Buffers syntax versions 2 and 3. However, the runtime version of the python protobuf library must exactly match the version of the protoc compiler what was used to generate the protobuf messages. ProtoCoder is registered in the global CoderRegistry as the default coder for any protobuf Message object. """ def __init__(self, proto_message_type): # type: (Type[google.protobuf.message.Message]) -> None self.proto_message_type = proto_message_type def _create_impl(self): return coder_impl.ProtoCoderImpl(self.proto_message_type) def is_deterministic(self): # type: () -> bool # TODO(vikasrk): A proto message can be deterministic if it does not contain # a Map. return False def as_deterministic_coder(self, step_label, error_message=None): return DeterministicProtoCoder(self.proto_message_type) def __eq__(self, other): return ( type(self) == type(other) and self.proto_message_type == other.proto_message_type) def __hash__(self): return hash(self.proto_message_type) @classmethod def from_type_hint(cls, typehint, unused_registry): # The typehint must be a strict subclass of google.protobuf.message.Message. # ProtoCoder cannot work with message.Message itself, as deserialization of # a serialized proto requires knowledge of the desired concrete proto # subclass which is not stored in the encoded bytes themselves. If this # occurs, an error is raised and the system defaults to other fallback # coders. if (issubclass(typehint, proto_utils.message_types) and typehint != message.Message): return cls(typehint) else: raise ValueError(( 'Expected a strict subclass of google.protobuf.message.Message' ', but got a %s' % typehint)) def to_type_hint(self): return self.proto_message_type class DeterministicProtoCoder(ProtoCoder): """A deterministic Coder for Google Protocol Buffers. It supports both Protocol Buffers syntax versions 2 and 3. However, the runtime version of the python protobuf library must exactly match the version of the protoc compiler what was used to generate the protobuf messages. """ def _create_impl(self): return coder_impl.DeterministicProtoCoderImpl(self.proto_message_type) def is_deterministic(self): # type: () -> bool return True def as_deterministic_coder(self, step_label, error_message=None): return self class ProtoPlusCoder(FastCoder): """A Coder for Google Protocol Buffers wrapped using the proto-plus library. ProtoPlusCoder is registered in the global CoderRegistry as the default coder for any proto.Message object. """ def __init__(self, proto_plus_message_type): # type: (Type[proto.Message]) -> None self.proto_plus_message_type = proto_plus_message_type def _create_impl(self): return coder_impl.ProtoPlusCoderImpl(self.proto_plus_message_type) def is_deterministic(self): return True def __eq__(self, other): return ( type(self) == type(other) and self.proto_plus_message_type == other.proto_plus_message_type) def __hash__(self): return hash(self.proto_plus_message_type) @classmethod def from_type_hint(cls, typehint, unused_registry): if issubclass(typehint, proto.Message): return cls(typehint) else: raise ValueError( 'Expected a subclass of proto.Message, but got a %s' % typehint) def to_type_hint(self): return self.proto_plus_message_type AVRO_GENERIC_CODER_URN = "beam:coder:avro:generic:v1" class AvroGenericCoder(FastCoder): """A coder used for AvroRecord values.""" def __init__(self, schema): self.schema = schema def _create_impl(self): return coder_impl.AvroCoderImpl(self.schema) def is_deterministic(self): # TODO(https://github.com/apache/beam/issues/19628): need to confirm if # it's deterministic return False def __eq__(self, other): return type(self) == type(other) and self.schema == other.schema def __hash__(self): return hash(self.schema) def to_type_hint(self): return AvroRecord def to_runner_api_parameter(self, context): return AVRO_GENERIC_CODER_URN, self.schema.encode('utf-8'), () @staticmethod @Coder.register_urn(AVRO_GENERIC_CODER_URN, bytes) def from_runner_api_parameter(payload, unused_components, unused_context): return AvroGenericCoder(payload.decode('utf-8')) class TupleCoder(FastCoder): """Coder of tuple objects.""" def __init__(self, components): # type: (Iterable[Coder]) -> None self._coders = tuple(components) def _create_impl(self): return coder_impl.TupleCoderImpl([c.get_impl() for c in self._coders]) def is_deterministic(self): # type: () -> bool return all(c.is_deterministic() for c in self._coders) def as_deterministic_coder(self, step_label, error_message=None): if self.is_deterministic(): return self else: return TupleCoder([ c.as_deterministic_coder(step_label, error_message) for c in self._coders ]) def to_type_hint(self): return typehints.Tuple[tuple(c.to_type_hint() for c in self._coders)] @classmethod def from_type_hint(cls, typehint, registry): # type: (typehints.TupleConstraint, CoderRegistry) -> TupleCoder return cls([registry.get_coder(t) for t in typehint.tuple_types]) def _get_component_coders(self): # type: () -> Tuple[Coder, ...] return self.coders() def coders(self): # type: () -> Tuple[Coder, ...] return self._coders def is_kv_coder(self): # type: () -> bool return len(self._coders) == 2 def key_coder(self): # type: () -> Coder if len(self._coders) != 2: raise ValueError('TupleCoder does not have exactly 2 components.') return self._coders[0] def value_coder(self): # type: () -> Coder if len(self._coders) != 2: raise ValueError('TupleCoder does not have exactly 2 components.') return self._coders[1] def __repr__(self): return 'TupleCoder[%s]' % ', '.join(str(c) for c in self._coders) def __eq__(self, other): return type(self) == type(other) and self._coders == other.coders() def __hash__(self): return hash(self._coders) def to_runner_api_parameter(self, context): if self.is_kv_coder(): return common_urns.coders.KV.urn, None, self.coders() else: return python_urns.TUPLE_CODER, None, self.coders() @staticmethod @Coder.register_urn(common_urns.coders.KV.urn, None) @Coder.register_urn(python_urns.TUPLE_CODER, None) def from_runner_api_parameter(unused_payload, components, unused_context): return TupleCoder(components) class TupleSequenceCoder(FastCoder): """Coder of homogeneous tuple objects.""" def __init__(self, elem_coder): # type: (Coder) -> None self._elem_coder = elem_coder def value_coder(self): return self._elem_coder def _create_impl(self): return coder_impl.TupleSequenceCoderImpl(self._elem_coder.get_impl()) def is_deterministic(self): # type: () -> bool return self._elem_coder.is_deterministic() def as_deterministic_coder(self, step_label, error_message=None): if self.is_deterministic(): return self else: return TupleSequenceCoder( self._elem_coder.as_deterministic_coder(step_label, error_message)) @classmethod def from_type_hint(cls, typehint, registry): # type: (Any, CoderRegistry) -> TupleSequenceCoder return cls(registry.get_coder(typehint.inner_type)) def _get_component_coders(self): # type: () -> Tuple[Coder, ...] return (self._elem_coder, ) def __repr__(self): return 'TupleSequenceCoder[%r]' % self._elem_coder def __eq__(self, other): return ( type(self) == type(other) and self._elem_coder == other.value_coder()) def __hash__(self): return hash((type(self), self._elem_coder)) class ListLikeCoder(FastCoder): """Coder of iterables of homogeneous objects.""" def __init__(self, elem_coder): # type: (Coder) -> None self._elem_coder = elem_coder def _create_impl(self): return coder_impl.IterableCoderImpl(self._elem_coder.get_impl()) def is_deterministic(self): # type: () -> bool return self._elem_coder.is_deterministic() def as_deterministic_coder(self, step_label, error_message=None): if self.is_deterministic(): return self else: return type(self)( self._elem_coder.as_deterministic_coder(step_label, error_message)) def value_coder(self): return self._elem_coder @classmethod def from_type_hint(cls, typehint, registry): # type: (Any, CoderRegistry) -> ListLikeCoder return cls(registry.get_coder(typehint.inner_type)) def _get_component_coders(self): # type: () -> Tuple[Coder, ...] return (self._elem_coder, ) def __repr__(self): return '%s[%r]' % (self.__class__.__name__, self._elem_coder) def __eq__(self, other): return ( type(self) == type(other) and self._elem_coder == other.value_coder()) def __hash__(self): return hash((type(self), self._elem_coder)) class IterableCoder(ListLikeCoder): """Coder of iterables of homogeneous objects.""" def to_type_hint(self): return typehints.Iterable[self._elem_coder.to_type_hint()] Coder.register_structured_urn(common_urns.coders.ITERABLE.urn, IterableCoder) class ListCoder(ListLikeCoder): """Coder of Python lists.""" def to_type_hint(self): return typehints.List[self._elem_coder.to_type_hint()] def _create_impl(self): return coder_impl.ListCoderImpl(self._elem_coder.get_impl()) class GlobalWindowCoder(SingletonCoder): """Coder for global windows.""" def __init__(self): from apache_beam.transforms import window super().__init__(window.GlobalWindow()) Coder.register_structured_urn( common_urns.coders.GLOBAL_WINDOW.urn, GlobalWindowCoder) class IntervalWindowCoder(FastCoder): """Coder for an window defined by a start timestamp and a duration.""" def _create_impl(self): return coder_impl.IntervalWindowCoderImpl() def is_deterministic(self): # type: () -> bool return True def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) Coder.register_structured_urn( common_urns.coders.INTERVAL_WINDOW.urn, IntervalWindowCoder) class _OrderedUnionCoder(FastCoder): def __init__( self, *coder_types: tuple[type, Coder], fallback_coder: Optional[Coder]): self._coder_types = coder_types self._fallback_coder = fallback_coder def _create_impl(self): return coder_impl._OrderedUnionCoderImpl( [(t, c.get_impl()) for t, c in self._coder_types], fallback_coder_impl=self._fallback_coder.get_impl() if self._fallback_coder else None) def is_deterministic(self) -> bool: return ( all(c.is_deterministic() for _, c in self._coder_types) and ( self._fallback_coder is None or self._fallback_coder.is_deterministic())) def to_type_hint(self): return Any def __eq__(self, other): return ( type(self) == type(other) and self._coder_types == other._coder_types and self._fallback_coder == other._fallback_coder) def __hash__(self): return hash((type(self), tuple(self._coder_types), self._fallback_coder)) class WindowedValueCoder(FastCoder): """Coder for windowed values.""" def __init__(self, wrapped_value_coder, window_coder=None): # type: (Coder, Optional[Coder]) -> None if not window_coder: # Avoid circular imports. from apache_beam.transforms import window window_coder = _OrderedUnionCoder( (window.GlobalWindow, GlobalWindowCoder()), (window.IntervalWindow, IntervalWindowCoder()), fallback_coder=PickleCoder()) self.wrapped_value_coder = wrapped_value_coder self.timestamp_coder = TimestampCoder() self.window_coder = window_coder def _create_impl(self): return coder_impl.WindowedValueCoderImpl( self.wrapped_value_coder.get_impl(), self.timestamp_coder.get_impl(), self.window_coder.get_impl()) def is_deterministic(self): # type: () -> bool return all( c.is_deterministic() for c in [self.wrapped_value_coder, self.timestamp_coder, self.window_coder]) def _get_component_coders(self): # type: () -> List[Coder] return [self.wrapped_value_coder, self.window_coder] def is_kv_coder(self): # type: () -> bool return self.wrapped_value_coder.is_kv_coder() def key_coder(self): # type: () -> Coder return self.wrapped_value_coder.key_coder() def value_coder(self): # type: () -> Coder return self.wrapped_value_coder.value_coder() def __repr__(self): return ( f'WindowedValueCoder[window_coder={self.window_coder}, ' f'wrapped_value_coder={self.wrapped_value_coder}]') def __eq__(self, other): return ( type(self) == type(other) and self.wrapped_value_coder == other.wrapped_value_coder and self.timestamp_coder == other.timestamp_coder and self.window_coder == other.window_coder) def __hash__(self): return hash( (self.wrapped_value_coder, self.timestamp_coder, self.window_coder)) @classmethod def from_type_hint(cls, typehint, registry): # type: (Any, CoderRegistry) -> WindowedValueCoder # Ideally this'd take two parameters so that one could hint at # the window type as well instead of falling back to the # pickle coders. return cls(registry.get_coder(typehint.inner_type)) def to_type_hint(self): return typehints.WindowedValue[self.wrapped_value_coder.to_type_hint()] Coder.register_structured_urn( common_urns.coders.WINDOWED_VALUE.urn, WindowedValueCoder) class ParamWindowedValueCoder(WindowedValueCoder): """A coder used for parameterized windowed values.""" def __init__(self, payload, components): super().__init__(components[0], components[1]) self.payload = payload def _create_impl(self): return coder_impl.ParamWindowedValueCoderImpl( self.wrapped_value_coder.get_impl(), self.window_coder.get_impl(), self.payload) def is_deterministic(self): # type: () -> bool return self.wrapped_value_coder.is_deterministic() def __repr__(self): return 'ParamWindowedValueCoder[%s]' % self.wrapped_value_coder def __eq__(self, other): return ( type(self) == type(other) and self.wrapped_value_coder == other.wrapped_value_coder and self.window_coder == other.window_coder and self.payload == other.payload) def __hash__(self): return hash((self.wrapped_value_coder, self.window_coder, self.payload)) @staticmethod @Coder.register_urn(common_urns.coders.PARAM_WINDOWED_VALUE.urn, bytes) def from_runner_api_parameter(payload, components, unused_context): return ParamWindowedValueCoder(payload, components) def to_runner_api_parameter(self, context): return ( common_urns.coders.PARAM_WINDOWED_VALUE.urn, self.payload, (self.wrapped_value_coder, self.window_coder)) class LengthPrefixCoder(FastCoder): """For internal use only; no backwards-compatibility guarantees. Coder which prefixes the length of the encoded object in the stream.""" def __init__(self, value_coder): # type: (Coder) -> None self._value_coder = value_coder def _create_impl(self): return coder_impl.LengthPrefixCoderImpl(self._value_coder.get_impl()) def is_deterministic(self): # type: () -> bool return self._value_coder.is_deterministic() def estimate_size(self, value): value_size = self._value_coder.estimate_size(value) return get_varint_size(value_size) + value_size def value_coder(self): return self._value_coder def _get_component_coders(self): # type: () -> Tuple[Coder, ...] return (self._value_coder, ) def __repr__(self): return 'LengthPrefixCoder[%r]' % self._value_coder def __eq__(self, other): return ( type(self) == type(other) and self._value_coder == other._value_coder) def __hash__(self): return hash((type(self), self._value_coder)) def to_type_hint(length_prefix_coder): return length_prefix_coder.value_coder().to_type_hint() Coder.register_structured_urn( common_urns.coders.LENGTH_PREFIX.urn, LengthPrefixCoder) class StateBackedIterableCoder(FastCoder): DEFAULT_WRITE_THRESHOLD = 1 def __init__( self, element_coder, # type: Coder read_state=None, # type: Optional[coder_impl.IterableStateReader] write_state=None, # type: Optional[coder_impl.IterableStateWriter] write_state_threshold=DEFAULT_WRITE_THRESHOLD): self._element_coder = element_coder self._read_state = read_state self._write_state = write_state self._write_state_threshold = write_state_threshold def _create_impl(self): return coder_impl.IterableCoderImpl( self._element_coder.get_impl(), self._read_state, self._write_state, self._write_state_threshold) def is_deterministic(self): # type: () -> bool return False def _get_component_coders(self): # type: () -> Tuple[Coder, ...] return (self._element_coder, ) def __repr__(self): return 'StateBackedIterableCoder[%r]' % self._element_coder def __eq__(self, other): return ( type(self) == type(other) and self._element_coder == other._element_coder and self._write_state_threshold == other._write_state_threshold) def __hash__(self): return hash((type(self), self._element_coder, self._write_state_threshold)) def to_runner_api_parameter(self, context): # type: (Optional[PipelineContext]) -> Tuple[str, Any, Sequence[Coder]] return ( common_urns.coders.STATE_BACKED_ITERABLE.urn, str(self._write_state_threshold).encode('ascii'), self._get_component_coders()) @staticmethod @Coder.register_urn(common_urns.coders.STATE_BACKED_ITERABLE.urn, bytes) def from_runner_api_parameter(payload, components, context): return StateBackedIterableCoder( components[0], read_state=context.iterable_state_read, write_state=context.iterable_state_write, write_state_threshold=int(payload) if payload else StateBackedIterableCoder.DEFAULT_WRITE_THRESHOLD) class ShardedKeyCoder(FastCoder): """A coder for sharded key.""" def __init__(self, key_coder): # type: (Coder) -> None self._key_coder = key_coder def _get_component_coders(self): # type: () -> List[Coder] return [self._key_coder] def _create_impl(self): return coder_impl.ShardedKeyCoderImpl(self._key_coder.get_impl()) def is_deterministic(self): # type: () -> bool return self._key_coder.is_deterministic() def to_type_hint(self): from apache_beam.typehints import sharded_key_type return sharded_key_type.ShardedKeyTypeConstraint( self._key_coder.to_type_hint()) @classmethod def from_type_hint(cls, typehint, registry): from apache_beam.typehints import sharded_key_type if isinstance(typehint, sharded_key_type.ShardedKeyTypeConstraint): return cls(registry.get_coder(typehint.key_type)) else: raise ValueError(( 'Expected an instance of ShardedKeyTypeConstraint' ', but got a %s' % typehint)) def __eq__(self, other): return type(self) == type(other) and self._key_coder == other._key_coder def __hash__(self): return hash(type(self)) + hash(self._key_coder) def __repr__(self): return 'ShardedKeyCoder[%s]' % self._key_coder Coder.register_structured_urn( common_urns.coders.SHARDED_KEY.urn, ShardedKeyCoder) class TimestampPrefixingWindowCoder(FastCoder): """For internal use only; no backwards-compatibility guarantees. Coder which prefixes the max timestamp of arbitrary window to its encoded form.""" def __init__(self, window_coder: Coder) -> None: self._window_coder = window_coder def _create_impl(self): return coder_impl.TimestampPrefixingWindowCoderImpl( self._window_coder.get_impl()) def to_type_hint(self): return self._window_coder.to_type_hint() def _get_component_coders(self) -> list[Coder]: return [self._window_coder] def is_deterministic(self) -> bool: return self._window_coder.is_deterministic() def __repr__(self): return 'TimestampPrefixingWindowCoder[%r]' % self._window_coder def __eq__(self, other): return ( type(self) == type(other) and self._window_coder == other._window_coder) def __hash__(self): return hash((type(self), self._window_coder)) Coder.register_structured_urn( common_urns.coders.CUSTOM_WINDOW.urn, TimestampPrefixingWindowCoder) class TimestampPrefixingOpaqueWindowCoder(FastCoder): """For internal use only; no backwards-compatibility guarantees. Coder which decodes windows as bytes.""" def __init__(self) -> None: pass def _create_impl(self): return coder_impl.TimestampPrefixingOpaqueWindowCoderImpl() def is_deterministic(self) -> bool: return True def __repr__(self): return 'TimestampPrefixingOpaqueWindowCoder' def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash((type(self))) Coder.register_structured_urn( python_urns.TIMESTAMP_PREFIXED_OPAQUE_WINDOW_CODER, TimestampPrefixingOpaqueWindowCoder) class BigIntegerCoder(FastCoder): def _create_impl(self): return coder_impl.BigIntegerCoderImpl() def is_deterministic(self): # type: () -> bool return True def to_type_hint(self): return int def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) class PaneInfoCoder(FastCoder): def _create_impl(self): return coder_impl.PaneInfoCoderImpl() def is_deterministic(self): # type: () -> bool return True def to_type_hint(self): return windowed_value.PaneInfo def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) class DecimalCoder(FastCoder): def _create_impl(self): return coder_impl.DecimalCoderImpl() def is_deterministic(self): # type: () -> bool return True def to_type_hint(self): return decimal.Decimal def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/coders_property_based_test.py0000644000175100017510000001335215206113032025132 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Property tests for coders in the Python SDK. The tests in this file utilize the hypothesis library to generate random test cases and run them against Beam's coder implementations. These tests are similar to fuzzing, except they test invariant properties of code. """ import keyword import math import typing import unittest # TODO(pabloem): Include other categories of characters from datetime import datetime from string import ascii_letters from string import digits import numpy as np from hypothesis import assume from hypothesis import given from hypothesis import settings from hypothesis import strategies as st from pytz import utc from apache_beam.coders import FloatCoder from apache_beam.coders import RowCoder from apache_beam.coders import StrUtf8Coder from apache_beam.coders.typecoders import registry as coders_registry from apache_beam.typehints.schemas import PRIMITIVE_TO_ATOMIC_TYPE from apache_beam.typehints.schemas import typing_to_runner_api from apache_beam.utils.timestamp import Timestamp SCHEMA_TYPES_TO_STRATEGY = { str: st.text(), bytes: st.binary(), typing.ByteString: st.binary(), # Maximum datetime on year 3000 to conform to Windows OS limits. Timestamp: st.datetimes( min_value=datetime(1970, 1, 1, 1, 1), max_value=datetime( 3000, 1, 1, 0, 0)).map(lambda dt: Timestamp.from_utc_datetime(dt.astimezone(utc))), int: st.integers(min_value=-(1 << 63 - 1), max_value=1 << 63 - 1), # INT8/BYTE not yet supported by RowCoder. # np.int8: st.binary(min_size=1, max_size=1), # INT16 not yet supported by RowCoder. # np.int16: st.integers(min_value=-(1 << 15 - 1), max_value=1 << 15 - 1), np.int32: st.integers(min_value=-(1 << 31 - 1), max_value=1 << 31 - 1), np.int64: st.integers(min_value=-(1 << 63 - 1), max_value=1 << 63 - 1), np.uint32: st.integers(min_value=0, max_value=1 << 32 - 1), np.uint64: st.integers(min_value=0, max_value=1 << 64 - 1), np.float32: st.floats(width=32, allow_nan=False), np.float64: st.floats(width=64, allow_nan=False), float: st.floats(width=64, allow_nan=False), bool: st.booleans() } # TODO(https://github.com/apache/beam/issues/23003): Support logical types. SCHEMA_TYPES = list(SCHEMA_TYPES_TO_STRATEGY.keys()) # A hypothesis strategy that generates schemas. # A schema is a list containing tuples of strings (field names), types (field # types) and boolean (nullable or not). # This strategy currently generates rows with simple types (i.e. non-list, and # non-map fields). SCHEMA_GENERATOR_STRATEGY = st.lists( st.tuples( st.text(ascii_letters + digits + '_', min_size=1), st.sampled_from(SCHEMA_TYPES), st.booleans())) TYPES_UNSUPPORTED_BY_ROW_CODER = {np.int8, np.int16} class TypesAreAllTested(unittest.TestCase): def test_all_types_are_tested(self): # Verify that all types among Beam's defined types are being tested self.assertEqual( set(SCHEMA_TYPES).intersection(PRIMITIVE_TO_ATOMIC_TYPE.keys()), set(PRIMITIVE_TO_ATOMIC_TYPE.keys()).difference( TYPES_UNSUPPORTED_BY_ROW_CODER)) class ProperyTestingCoders(unittest.TestCase): @given(st.text()) def test_string_coder(self, txt: str): coder = StrUtf8Coder() self.assertEqual(coder.decode(coder.encode(txt)), txt) @given(st.floats()) def test_float_coder(self, num: float): coder = FloatCoder() test_num = coder.decode(coder.encode(num)) if math.isnan(num): # nan != nan. self.assertTrue(math.isnan(test_num)) else: self.assertEqual(coder.decode(coder.encode(num)), num) @settings(deadline=None, print_blob=True) @given(st.data()) def test_row_coder(self, data: st.DataObject): """Generate rows and schemas, and test their encoding/decoding. The schemas are generated based on the SCHEMA_GENERATOR_STRATEGY. """ schema = data.draw(SCHEMA_GENERATOR_STRATEGY) # Assume that the cardinality of the set of names is the same # as the length of the schema. This means there's no duplicate # names for fields. # If this condition does not hold, then we must not continue the # test. assume(len({name for name, _, _ in schema}) == len(schema)) assume(all(not keyword.iskeyword(name) for name, _, _ in schema)) assume( len({n[0] for n, _, _ in schema}.intersection(set(digits + '_'))) == 0) RowType = typing.NamedTuple( # type: ignore 'RandomRowType', [(name, type_ if not nullable else typing.Optional[type_]) for name, type_, nullable in schema]) coders_registry.register_coder(RowType, RowCoder) # TODO(https://github.com/apache/beam/issues/23002): Apply nulls for these row = RowType( **{ name: data.draw(SCHEMA_TYPES_TO_STRATEGY[type_]) for name, type_, nullable in schema }) coder = RowCoder(typing_to_runner_api(RowType).row_type.schema) self.assertEqual(coder.decode(coder.encode(row)), row) if __name__ == "__main__": unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/coders_test.py0000644000175100017510000002376315206113032022037 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import base64 import logging import unittest import proto import pytest from google.protobuf import message import apache_beam as beam from apache_beam import typehints from apache_beam.coders import coders from apache_beam.coders import proto2_coder_test_messages_pb2 as test_message from apache_beam.coders.avro_record import AvroRecord from apache_beam.coders.typecoders import registry as coders_registry from apache_beam.testing.test_pipeline import TestPipeline class PickleCoderTest(unittest.TestCase): def test_basics(self): v = ('a' * 10, 'b' * 90) pickler = coders.PickleCoder() self.assertEqual(v, pickler.decode(pickler.encode(v))) pickler = coders.Base64PickleCoder() self.assertEqual(v, pickler.decode(pickler.encode(v))) self.assertEqual( coders.Base64PickleCoder().encode(v), base64.b64encode(coders.PickleCoder().encode(v))) def test_equality(self): self.assertEqual(coders.PickleCoder(), coders.PickleCoder()) self.assertEqual(coders.Base64PickleCoder(), coders.Base64PickleCoder()) self.assertNotEqual(coders.Base64PickleCoder(), coders.PickleCoder()) self.assertNotEqual(coders.Base64PickleCoder(), object()) class CodersTest(unittest.TestCase): def test_str_utf8_coder(self): real_coder = coders_registry.get_coder(bytes) expected_coder = coders.BytesCoder() self.assertEqual(real_coder.encode(b'abc'), expected_coder.encode(b'abc')) self.assertEqual(b'abc', real_coder.decode(real_coder.encode(b'abc'))) # The test proto message file was generated by running the following: # # `cd ` # `cp sdks/java/extensions/protobuf/src/test/proto/\ # proto2_coder_test_messages.proto sdks/python/apache_beam/coders/` # `cd sdks/python` # `protoc apache_beam/coders/proto2_coder_test_messages.proto # --python_out=. # `rm apache_beam/coders/proto2_coder_test_message.proto` # # Note: The protoc version should match the protobuf library version specified # in setup.py. # # TODO(https://github.com/apache/beam/issues/22319): The proto file should be # placed in a common directory that can be shared between java and python. class ProtoCoderTest(unittest.TestCase): def test_proto_coder(self): ma = test_message.MessageA() mb = ma.field2.add() mb.field1 = True ma.field1 = 'hello world' expected_coder = coders.ProtoCoder(ma.__class__) real_coder = coders_registry.get_coder(ma.__class__) self.assertEqual(expected_coder, real_coder) self.assertEqual(real_coder.encode(ma), expected_coder.encode(ma)) self.assertEqual(ma, real_coder.decode(real_coder.encode(ma))) self.assertEqual(ma.__class__, real_coder.to_type_hint()) def test_proto_coder_on_protobuf_message_subclasses(self): # This replicates a scenario where users provide message.Message as the # output typehint for a Map function, even though the actual output messages # are subclasses of message.Message. ma = test_message.MessageA() mb = ma.field2.add() mb.field1 = True ma.field1 = 'hello world' coder = coders_registry.get_coder(message.Message) # For messages of google.protobuf.message.Message, the fallback coder will # be FastPrimitivesCoder rather than ProtoCoder. # See the comment on ProtoCoder.from_type_hint() for further details. self.assertEqual(coder, coders.FastPrimitivesCoder()) self.assertEqual(ma, coder.decode(coder.encode(ma))) class DeterministicProtoCoderTest(unittest.TestCase): def test_deterministic_proto_coder(self): ma = test_message.MessageA() mb = ma.field2.add() mb.field1 = True ma.field1 = 'hello world' expected_coder = coders.DeterministicProtoCoder(ma.__class__) real_coder = ( coders_registry.get_coder( ma.__class__).as_deterministic_coder(step_label='unused')) self.assertTrue(real_coder.is_deterministic()) self.assertEqual(expected_coder, real_coder) self.assertEqual(real_coder.encode(ma), expected_coder.encode(ma)) self.assertEqual(ma, real_coder.decode(real_coder.encode(ma))) def test_deterministic_proto_coder_determinism(self): for _ in range(10): keys = list(range(20)) mm_forward = test_message.MessageWithMap() for key in keys: mm_forward.field1[str(key)].field1 = str(key) mm_reverse = test_message.MessageWithMap() for key in reversed(keys): mm_reverse.field1[str(key)].field1 = str(key) coder = coders.DeterministicProtoCoder(mm_forward.__class__) self.assertEqual(coder.encode(mm_forward), coder.encode(mm_reverse)) class ProtoPlusMessageB(proto.Message): field1 = proto.Field(proto.BOOL, number=1) class ProtoPlusMessageA(proto.Message): field1 = proto.Field(proto.STRING, number=1) field2 = proto.RepeatedField(ProtoPlusMessageB, number=2) class ProtoPlusMessageWithMap(proto.Message): field1 = proto.MapField(proto.STRING, ProtoPlusMessageA, number=1) class ProtoPlusCoderTest(unittest.TestCase): def test_proto_plus_coder(self): ma = ProtoPlusMessageA() ma.field2 = [ProtoPlusMessageB(field1=True)] ma.field1 = 'hello world' expected_coder = coders.ProtoPlusCoder(ma.__class__) real_coder = coders_registry.get_coder(ma.__class__) self.assertTrue(issubclass(ma.__class__, proto.Message)) self.assertEqual(expected_coder, real_coder) self.assertTrue(real_coder.is_deterministic()) self.assertEqual(real_coder.encode(ma), expected_coder.encode(ma)) self.assertEqual(ma, real_coder.decode(real_coder.encode(ma))) def test_proto_plus_coder_determinism(self): for _ in range(10): keys = list(range(20)) mm_forward = ProtoPlusMessageWithMap() for key in keys: mm_forward.field1[str(key)] = ProtoPlusMessageA(field1=str(key)) # pylint: disable=E1137 mm_reverse = ProtoPlusMessageWithMap() for key in reversed(keys): mm_reverse.field1[str(key)] = ProtoPlusMessageA(field1=str(key)) # pylint: disable=E1137 coder = coders.ProtoPlusCoder(ProtoPlusMessageWithMap) self.assertEqual(coder.encode(mm_forward), coder.encode(mm_reverse)) class AvroTestCoder(coders.AvroGenericCoder): SCHEMA = """ { "type": "record", "name": "testrecord", "fields": [ {"name": "name", "type": "string"}, {"name": "age", "type": "int"} ] } """ def __init__(self): super().__init__(self.SCHEMA) class AvroTestRecord(AvroRecord): pass coders_registry.register_coder(AvroTestRecord, AvroTestCoder) class AvroCoderTest(unittest.TestCase): def test_avro_record_coder(self): real_coder = coders_registry.get_coder(AvroTestRecord) expected_coder = AvroTestCoder() self.assertEqual( real_coder.encode( AvroTestRecord({ "name": "Daenerys targaryen", "age": 23 })), expected_coder.encode( AvroTestRecord({ "name": "Daenerys targaryen", "age": 23 }))) self.assertEqual( AvroTestRecord({ "name": "Jon Snow", "age": 23 }), real_coder.decode( real_coder.encode(AvroTestRecord({ "name": "Jon Snow", "age": 23 })))) class DummyClass(object): """A class with no registered coder.""" def __init__(self): pass def __eq__(self, other): if isinstance(other, self.__class__): return True return False def __hash__(self): return hash(type(self)) class FallbackCoderTest(unittest.TestCase): def test_default_fallback_path(self): """Test fallback path picks a matching coder if no coder is registered.""" coder = coders_registry.get_coder(DummyClass) # No matching coder, so picks the last fallback coder which is a # FastPrimitivesCoder. self.assertEqual(coder, coders.FastPrimitivesCoder()) self.assertEqual(DummyClass(), coder.decode(coder.encode(DummyClass()))) class NullableCoderTest(unittest.TestCase): def test_determinism(self): deterministic = coders_registry.get_coder(typehints.Optional[int]) deterministic.as_deterministic_coder('label') complex_deterministic = coders_registry.get_coder( typehints.Optional[DummyClass]) complex_deterministic.as_deterministic_coder('label') nondeterministic = coders.NullableCoder(coders.Base64PickleCoder()) with pytest.raises(ValueError): nondeterministic.as_deterministic_coder('label') class LengthPrefixCoderTest(unittest.TestCase): def test_to_type_hint(self): coder = coders.LengthPrefixCoder(coders.BytesCoder()) assert coder.to_type_hint() is bytes class NumpyIntAsKeyTest(unittest.TestCase): def test_numpy_int(self): # this type is not supported as the key import numpy as np with self.assertRaisesRegex(Exception, "Unable to deterministically"): with TestPipeline() as p: indata = p | "Create" >> beam.Create([(a, int(a)) for a in np.arange(3)]) # Apply CombinePerkey to sum values for each key. _ = indata | "CombinePerKey" >> beam.CombinePerKey(sum) class WindowedValueCoderTest(unittest.TestCase): def test_to_type_hint(self): coder = coders.WindowedValueCoder(coders.VarIntCoder()) self.assertEqual(coder.to_type_hint(), typehints.WindowedValue[int]) # type: ignore[misc] if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/coders_test_common.py0000644000175100017510000011702415206113032023401 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests common to all coder implementations.""" # pytype: skip-file import base64 import collections import enum import logging import math import os import pickle import subprocess import sys import textwrap import unittest from decimal import Decimal from typing import Any from typing import NamedTuple import pytest from parameterized import param from parameterized import parameterized from apache_beam.coders import coders from apache_beam.coders import proto2_coder_test_messages_pb2 as test_message from apache_beam.coders import typecoders from apache_beam.internal import pickler from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options_context import scoped_pipeline_options from apache_beam.runners import pipeline_context from apache_beam.transforms import userstate from apache_beam.transforms import window from apache_beam.transforms.window import GlobalWindow from apache_beam.typehints import sharded_key_type from apache_beam.typehints import typehints from apache_beam.utils import timestamp from apache_beam.utils import windowed_value from apache_beam.utils.sharded_key import ShardedKey from apache_beam.utils.timestamp import MIN_TIMESTAMP from . import observable try: import dataclasses except ImportError: dataclasses = None # type: ignore try: import dill except ImportError: dill = None MyNamedTuple = collections.namedtuple('A', ['x', 'y']) # type: ignore[name-match] AnotherNamedTuple = collections.namedtuple('AnotherNamedTuple', ['x', 'y']) MyTypedNamedTuple = NamedTuple('MyTypedNamedTuple', [('f1', int), ('f2', str)]) class MyEnum(enum.Enum): E1 = 5 E2 = enum.auto() E3 = 'abc' MyIntEnum = enum.IntEnum('MyIntEnum', 'I1 I2 I3') MyIntFlag = enum.IntFlag('MyIntFlag', 'F1 F2 F3') MyFlag = enum.Flag('MyFlag', 'F1 F2 F3') # pylint: disable=too-many-function-args class DefinesGetState: def __init__(self, value): self.value = value def __getstate__(self): return self.value def __eq__(self, other): return type(other) is type(self) and other.value == self.value class DefinesGetAndSetState(DefinesGetState): def __setstate__(self, value): self.value = value # Defined out of line for picklability. class CustomCoder(coders.Coder): def encode(self, x): return str(x + 1).encode('utf-8') def decode(self, encoded): return int(encoded) - 1 if dataclasses is not None: @dataclasses.dataclass(frozen=True) class FrozenDataClass: a: Any b: int @dataclasses.dataclass(frozen=True, kw_only=True) class FrozenKwOnlyDataClass: c: int d: int @dataclasses.dataclass class UnFrozenDataClass: x: int y: int @dataclasses.dataclass(frozen=True, kw_only=True) class FrozenUnInitKwOnlyDataClass: side: int area: int = dataclasses.field(init=False) def __post_init__(self): # Hack to update an attribute in a frozen dataclass. object.__setattr__(self, 'area', self.side**2) # These tests need to all be run in the same process due to the asserts # in tearDownClass. @pytest.mark.no_xdist @pytest.mark.uses_dill class CodersTest(unittest.TestCase): # These class methods ensure that we test each defined coder in both # nested and unnested context. # Common test values representing Python's built-in types. test_values_deterministic: list[Any] = [ None, 1, -1, 1.5, b'str\0str', 'unicode\0\u0101', (), (1, 2, 3), [], [1, 2, 3], True, False, ] test_values = test_values_deterministic + [ {}, { 'a': 'b' }, { 0: {}, 1: len }, set(), {'a', 'b'}, len, ] @classmethod def setUpClass(cls): cls.seen = set() cls.seen_nested = set() @classmethod def tearDownClass(cls): standard = set( c for c in coders.__dict__.values() if isinstance(c, type) and issubclass(c, coders.Coder) and 'Base' not in c.__name__) standard -= set([ coders.Coder, coders.AvroGenericCoder, coders.DeterministicProtoCoder, coders.FastCoder, coders.ListLikeCoder, coders.ProtoCoder, coders.ProtoPlusCoder, coders.BigEndianShortCoder, coders.ByteCoder, coders.SinglePrecisionFloatCoder, coders.ToBytesCoder, coders.BigIntegerCoder, # tested in DecimalCoder coders.TimestampPrefixingOpaqueWindowCoder, ]) if not dill: standard -= set( [coders.DillCoder, coders.DeterministicFastPrimitivesCoder]) cls.seen_nested -= set( [coders.ProtoCoder, coders.ProtoPlusCoder, CustomCoder]) assert not standard - cls.seen, str(standard - cls.seen) assert not cls.seen_nested - standard, str(cls.seen_nested - standard) @classmethod def _observe(cls, coder): cls.seen.add(type(coder)) cls._observe_nested(coder) @classmethod def _observe_nested(cls, coder): if isinstance(coder, coders.TupleCoder): for c in coder.coders(): cls.seen_nested.add(type(c)) cls._observe_nested(c) def check_coder(self, coder, *values, **kwargs): context = kwargs.pop('context', pipeline_context.PipelineContext()) test_size_estimation = kwargs.pop('test_size_estimation', True) assert not kwargs self._observe(coder) for v in values: self.assertEqual(v, coder.decode(coder.encode(v))) if test_size_estimation: self.assertEqual(coder.estimate_size(v), len(coder.encode(v))) self.assertEqual( coder.estimate_size(v), coder.get_impl().estimate_size(v)) self.assertEqual( coder.get_impl().get_estimated_size_and_observables(v), (coder.get_impl().estimate_size(v), [])) copy1 = pickler.loads(pickler.dumps(coder)) copy2 = coders.Coder.from_runner_api(coder.to_runner_api(context), context) for v in values: self.assertEqual(v, copy1.decode(copy2.encode(v))) if coder.is_deterministic(): self.assertEqual(copy1.encode(v), copy2.encode(v)) def test_custom_coder(self): self.check_coder(CustomCoder(), 1, -10, 5) self.check_coder( coders.TupleCoder((CustomCoder(), coders.BytesCoder())), (1, b'a'), (-10, b'b'), (5, b'c')) def test_pickle_coder(self): coder = coders.PickleCoder() self.check_coder(coder, *self.test_values) def test_cloudpickle_pickle_coder(self): cell_value = (lambda x: lambda: x)(0).__closure__[0] self.check_coder(coders.CloudpickleCoder(), 'a', 1, cell_value) self.check_coder( coders.TupleCoder((coders.VarIntCoder(), coders.CloudpickleCoder())), (1, cell_value)) def test_memoizing_pickle_coder(self): coder = coders._MemoizingPickleCoder() self.check_coder(coder, *self.test_values) @parameterized.expand([ param(compat_version=None), param(compat_version="2.67.0"), param(compat_version="2.68.0"), ]) def test_deterministic_coder(self, compat_version): """ Test in process determinism for all special deterministic types - In SDK version <= 2.67.0 dill is used to encode "special types" - In SDK version 2.68.0 cloudpickle is used to encode "special types" with absolute filepaths in code objects and dynamic functions. - In SDK version >=2.69.0 cloudpickle is used to encode "special types" with relative filepaths in code objects and dynamic functions. """ with scoped_pipeline_options( PipelineOptions(update_compatibility_version=compat_version)): coder = coders.FastPrimitivesCoder() if not dill and compat_version == "2.67.0": with self.assertRaises(RuntimeError): coder.as_deterministic_coder(step_label="step") self.skipTest('Dill not installed') deterministic_coder = coder.as_deterministic_coder(step_label="step") self.check_coder(deterministic_coder, *self.test_values_deterministic) for v in self.test_values_deterministic: self.check_coder(coders.TupleCoder((deterministic_coder, )), (v, )) self.check_coder( coders.TupleCoder( (deterministic_coder, ) * len(self.test_values_deterministic)), tuple(self.test_values_deterministic)) self.check_coder(deterministic_coder, {}) self.check_coder(deterministic_coder, {2: 'x', 1: 'y'}) with self.assertRaises(TypeError): self.check_coder(deterministic_coder, {1: 'x', 'y': 2}) self.check_coder(deterministic_coder, [1, {}]) with self.assertRaises(TypeError): self.check_coder(deterministic_coder, [1, {1: 'x', 'y': 2}]) self.check_coder( coders.TupleCoder((deterministic_coder, coder)), (1, {}), ('a', [{}])) self.check_coder( deterministic_coder, test_message.MessageA(field1='value')) # Skip this test during cloudpickle. Dill monkey patches the __reduce__ # method for anonymous named tuples (MyNamedTuple) which is not # pickleable. Since the test is parameterized the type gets colbbered. if compat_version == "2.67.0": self.check_coder( deterministic_coder, [MyNamedTuple(1, 2), MyTypedNamedTuple(1, 'a')]) self.check_coder( deterministic_coder, [AnotherNamedTuple(1, 2), MyTypedNamedTuple(1, 'a')]) if dataclasses is not None: self.check_coder(deterministic_coder, FrozenDataClass(1, 2)) self.check_coder(deterministic_coder, FrozenKwOnlyDataClass(c=1, d=2)) self.check_coder( deterministic_coder, FrozenUnInitKwOnlyDataClass(side=11)) with self.assertRaises(TypeError): self.check_coder(deterministic_coder, UnFrozenDataClass(1, 2)) with self.assertRaises(TypeError): self.check_coder( deterministic_coder, FrozenDataClass(UnFrozenDataClass(1, 2), 3)) with self.assertRaises(TypeError): self.check_coder( deterministic_coder, AnotherNamedTuple(UnFrozenDataClass(1, 2), 3)) self.check_coder(deterministic_coder, list(MyEnum)) self.check_coder(deterministic_coder, list(MyIntEnum)) self.check_coder(deterministic_coder, list(MyIntFlag)) self.check_coder(deterministic_coder, list(MyFlag)) self.check_coder( deterministic_coder, [DefinesGetAndSetState(1), DefinesGetAndSetState((1, 2, 3))]) with self.assertRaises(TypeError): self.check_coder(deterministic_coder, DefinesGetState(1)) with self.assertRaises(TypeError): self.check_coder( deterministic_coder, DefinesGetAndSetState({ 1: 'x', 'y': 2 })) @parameterized.expand([ param(compat_version=None), param(compat_version="2.67.0"), param(compat_version="2.68.0"), ]) def test_deterministic_map_coder_is_update_compatible(self, compat_version): """ Test in process determinism for map coder including when a component coder uses DeterministicFastPrimitivesCoder for "special types". - In SDK version <= 2.67.0 dill is used to encode "special types" - In SDK version 2.68.0 cloudpickle is used to encode "special types" with absolute filepaths in code objects and dynamic functions. - In SDK version >=2.69.0 cloudpickle is used to encode "special types" with relative file. """ with scoped_pipeline_options( PipelineOptions(update_compatibility_version=compat_version)): values = [{ MyTypedNamedTuple(i, 'a'): MyTypedNamedTuple('a', i) for i in range(10) }] coder = coders.MapCoder( coders.FastPrimitivesCoder(), coders.FastPrimitivesCoder()) if not dill and compat_version == "2.67.0": with self.assertRaises(RuntimeError): coder.as_deterministic_coder(step_label="step") self.skipTest('Dill not installed') deterministic_coder = coder.as_deterministic_coder(step_label="step") assert isinstance( deterministic_coder._key_coder, coders.DeterministicFastPrimitivesCoderV2 if compat_version in (None, "2.68.0") else coders.DeterministicFastPrimitivesCoder) self.check_coder(deterministic_coder, *values) def test_dill_coder(self): if not dill: with self.assertRaises(RuntimeError): coders.DillCoder() self.skipTest('Dill not installed') cell_value = (lambda x: lambda: x)(0).__closure__[0] self.check_coder(coders.DillCoder(), 'a', 1, cell_value) self.check_coder( coders.TupleCoder((coders.VarIntCoder(), coders.DillCoder())), (1, cell_value)) def test_fast_primitives_coder(self): coder = coders.FastPrimitivesCoder(coders.SingletonCoder(len)) self.check_coder(coder, *self.test_values) for v in self.test_values: self.check_coder(coders.TupleCoder((coder, )), (v, )) def test_fast_primitives_coder_large_int(self): coder = coders.FastPrimitivesCoder() self.check_coder(coder, 10**100) def test_fake_deterministic_fast_primitives_coder(self): coder = coders.FakeDeterministicFastPrimitivesCoder(coders.PickleCoder()) self.check_coder(coder, *self.test_values) for v in self.test_values: self.check_coder(coders.TupleCoder((coder, )), (v, )) def test_bytes_coder(self): self.check_coder(coders.BytesCoder(), b'a', b'\0', b'z' * 1000) def test_bool_coder(self): self.check_coder(coders.BooleanCoder(), True, False) def test_varint_coder(self): # Small ints. self.check_coder(coders.VarIntCoder(), *range(-10, 10)) # Multi-byte encoding starts at 128 self.check_coder(coders.VarIntCoder(), *range(120, 140)) # Large values MAX_64_BIT_INT = 0x7fffffffffffffff self.check_coder( coders.VarIntCoder(), *[ int(math.pow(-1, k) * math.exp(k)) for k in range(0, int(math.log(MAX_64_BIT_INT))) ]) def test_varint32_coder(self): # Small ints. self.check_coder(coders.VarInt32Coder(), *range(-10, 10)) # Multi-byte encoding starts at 128 self.check_coder(coders.VarInt32Coder(), *range(120, 140)) # Large values MAX_32_BIT_INT = 0x7fffffff self.check_coder( coders.VarIntCoder(), *[ int(math.pow(-1, k) * math.exp(k)) for k in range(0, int(math.log(MAX_32_BIT_INT))) ]) def test_float_coder(self): self.check_coder( coders.FloatCoder(), *[float(0.1 * x) for x in range(-100, 100)]) self.check_coder( coders.FloatCoder(), *[float(2**(0.1 * x)) for x in range(-100, 100)]) self.check_coder(coders.FloatCoder(), float('-Inf'), float('Inf')) self.check_coder( coders.TupleCoder((coders.FloatCoder(), coders.FloatCoder())), (0, 1), (-100, 100), (0.5, 0.25)) def test_singleton_coder(self): a = 'anything' b = 'something else' self.check_coder(coders.SingletonCoder(a), a) self.check_coder(coders.SingletonCoder(b), b) self.check_coder( coders.TupleCoder((coders.SingletonCoder(a), coders.SingletonCoder(b))), (a, b)) def test_interval_window_coder(self): self.check_coder( coders.IntervalWindowCoder(), *[ window.IntervalWindow(x, y) for x in [-2**52, 0, 2**52] for y in range(-100, 100) ]) self.check_coder( coders.TupleCoder((coders.IntervalWindowCoder(), )), (window.IntervalWindow(0, 10), )) def test_paneinfo_window_coder(self): self.check_coder( coders.PaneInfoCoder(), *[ windowed_value.PaneInfo( is_first=y == 0, is_last=y == 9, timing=windowed_value.PaneInfoTiming.EARLY, index=y, nonspeculative_index=-1) for y in range(0, 10) ]) def test_timestamp_coder(self): self.check_coder( coders.TimestampCoder(), *[timestamp.Timestamp(micros=x) for x in (-1000, 0, 1000)]) self.check_coder( coders.TimestampCoder(), timestamp.Timestamp(micros=-1234567000), timestamp.Timestamp(micros=1234567000)) self.check_coder( coders.TimestampCoder(), timestamp.Timestamp(micros=-1234567890123456000), timestamp.Timestamp(micros=1234567890123456000)) self.check_coder( coders.TupleCoder((coders.TimestampCoder(), coders.BytesCoder())), (timestamp.Timestamp.of(27), b'abc')) def test_timer_coder(self): self.check_coder( coders._TimerCoder(coders.StrUtf8Coder(), coders.GlobalWindowCoder()), *[ userstate.Timer( user_key="key", dynamic_timer_tag="tag", windows=(GlobalWindow(), ), clear_bit=True, fire_timestamp=None, hold_timestamp=None, paneinfo=None), userstate.Timer( user_key="key", dynamic_timer_tag="tag", windows=(GlobalWindow(), ), clear_bit=False, fire_timestamp=timestamp.Timestamp.of(123), hold_timestamp=timestamp.Timestamp.of(456), paneinfo=windowed_value.PANE_INFO_UNKNOWN) ]) def test_tuple_coder(self): kv_coder = coders.TupleCoder((coders.VarIntCoder(), coders.BytesCoder())) # Test binary representation self.assertEqual(b'\x04abc', kv_coder.encode((4, b'abc'))) # Test unnested self.check_coder(kv_coder, (1, b'a'), (-2, b'a' * 100), (300, b'abc\0' * 5)) # Test nested self.check_coder( coders.TupleCoder(( coders.TupleCoder((coders.PickleCoder(), coders.VarIntCoder())), coders.StrUtf8Coder(), coders.BooleanCoder())), ((1, 2), 'a', True), ((-2, 5), 'a\u0101' * 100, False), ((300, 1), 'abc\0' * 5, True)) def test_tuple_sequence_coder(self): int_tuple_coder = coders.TupleSequenceCoder(coders.VarIntCoder()) self.check_coder(int_tuple_coder, (1, -1, 0), (), tuple(range(1000))) self.check_coder( coders.TupleCoder((coders.VarIntCoder(), int_tuple_coder)), (1, (1, 2, 3))) def test_base64_pickle_coder(self): self.check_coder(coders.Base64PickleCoder(), 'a', 1, 1.5, (1, 2, 3)) def test_utf8_coder(self): self.check_coder(coders.StrUtf8Coder(), 'a', 'ab\u00FF', '\u0101\0') def test_iterable_coder(self): iterable_coder = coders.IterableCoder(coders.VarIntCoder()) # Test unnested self.check_coder(iterable_coder, [1], [-1, 0, 100]) # Test nested self.check_coder( coders.TupleCoder( (coders.VarIntCoder(), coders.IterableCoder(coders.VarIntCoder()))), (1, [1, 2, 3])) def test_iterable_coder_unknown_length(self): # Empty self._test_iterable_coder_of_unknown_length(0) # Single element self._test_iterable_coder_of_unknown_length(1) # Multiple elements self._test_iterable_coder_of_unknown_length(100) # Multiple elements with underlying stream buffer overflow. self._test_iterable_coder_of_unknown_length(80000) def _test_iterable_coder_of_unknown_length(self, count): def iter_generator(count): for i in range(count): yield i iterable_coder = coders.IterableCoder(coders.VarIntCoder()) self.assertCountEqual( list(iter_generator(count)), iterable_coder.decode(iterable_coder.encode(iter_generator(count)))) def test_list_coder(self): list_coder = coders.ListCoder(coders.VarIntCoder()) # Test unnested self.check_coder(list_coder, [1], [-1, 0, 100]) # Test nested self.check_coder( coders.TupleCoder((coders.VarIntCoder(), list_coder)), (1, [1, 2, 3])) def test_windowedvalue_coder_paneinfo(self): coder = coders.WindowedValueCoder( coders.VarIntCoder(), coders.GlobalWindowCoder()) test_paneinfo_values = [ windowed_value.PANE_INFO_UNKNOWN, windowed_value.PaneInfo( True, True, windowed_value.PaneInfoTiming.EARLY, 0, -1), windowed_value.PaneInfo( True, False, windowed_value.PaneInfoTiming.ON_TIME, 0, 0), windowed_value.PaneInfo( True, False, windowed_value.PaneInfoTiming.ON_TIME, 10, 0), windowed_value.PaneInfo( False, True, windowed_value.PaneInfoTiming.ON_TIME, 0, 23), windowed_value.PaneInfo( False, True, windowed_value.PaneInfoTiming.ON_TIME, 12, 23), windowed_value.PaneInfo( False, False, windowed_value.PaneInfoTiming.LATE, 0, 123), ] test_values = [ windowed_value.WindowedValue(123, 234, (GlobalWindow(), ), p) for p in test_paneinfo_values ] # Test unnested. self.check_coder( coder, windowed_value.WindowedValue( 123, 234, (GlobalWindow(), ), windowed_value.PANE_INFO_UNKNOWN)) for value in test_values: self.check_coder(coder, value) # Test nested. for value1 in test_values: for value2 in test_values: self.check_coder(coders.TupleCoder((coder, coder)), (value1, value2)) def test_windowed_value_coder(self): coder = coders.WindowedValueCoder( coders.VarIntCoder(), coders.GlobalWindowCoder()) # Test binary representation self.assertEqual( b'\x7f\xdf;dZ\x1c\xac\t\x00\x00\x00\x01\x0f\x01', coder.encode(window.GlobalWindows.windowed_value(1))) # Test decoding large timestamp self.assertEqual( coder.decode(b'\x7f\xdf;dZ\x1c\xac\x08\x00\x00\x00\x01\x0f\x00'), windowed_value.create(0, MIN_TIMESTAMP.micros, (GlobalWindow(), ))) # Test unnested self.check_coder( coders.WindowedValueCoder(coders.VarIntCoder()), windowed_value.WindowedValue(3, -100, ()), windowed_value.WindowedValue(-1, 100, (1, 2, 3))) # Test Global Window self.check_coder( coders.WindowedValueCoder( coders.VarIntCoder(), coders.GlobalWindowCoder()), window.GlobalWindows.windowed_value(1)) # Test nested self.check_coder( coders.TupleCoder(( coders.WindowedValueCoder(coders.FloatCoder()), coders.WindowedValueCoder(coders.StrUtf8Coder()))), ( windowed_value.WindowedValue(1.5, 0, ()), windowed_value.WindowedValue("abc", 10, ('window', )))) def test_param_windowed_value_coder(self): from apache_beam.transforms.window import IntervalWindow from apache_beam.utils.windowed_value import PaneInfo # pylint: disable=too-many-function-args wv = windowed_value.create( b'', # Milliseconds to microseconds 1000 * 1000, (IntervalWindow(11, 21), ), PaneInfo(True, False, 1, 2, 3)) windowed_value_coder = coders.WindowedValueCoder( coders.BytesCoder(), coders.IntervalWindowCoder()) payload = windowed_value_coder.encode(wv) coder = coders.ParamWindowedValueCoder( payload, [coders.VarIntCoder(), coders.IntervalWindowCoder()]) # Test binary representation self.assertEqual( b'\x01', coder.encode(window.GlobalWindows.windowed_value(1))) # Test unnested self.check_coder( coders.ParamWindowedValueCoder( payload, [coders.VarIntCoder(), coders.IntervalWindowCoder()]), windowed_value.WindowedValue( 3, 1, (window.IntervalWindow(11, 21), ), PaneInfo(True, False, 1, 2, 3)), windowed_value.WindowedValue( 1, 1, (window.IntervalWindow(11, 21), ), PaneInfo(True, False, 1, 2, 3))) # Test nested self.check_coder( coders.TupleCoder(( coders.ParamWindowedValueCoder( payload, [coders.FloatCoder(), coders.IntervalWindowCoder()]), coders.ParamWindowedValueCoder( payload, [coders.StrUtf8Coder(), coders.IntervalWindowCoder()]))), ( windowed_value.WindowedValue( 1.5, 1, (window.IntervalWindow(11, 21), ), PaneInfo(True, False, 1, 2, 3)), windowed_value.WindowedValue( "abc", 1, (window.IntervalWindow(11, 21), ), PaneInfo(True, False, 1, 2, 3)))) @parameterized.expand([ param(compat_version=None), param(compat_version="2.67.0"), param(compat_version="2.68.0"), ]) def test_cross_process_encoding_of_special_types_is_deterministic( self, compat_version): """Test cross-process determinism for all special deterministic types - In SDK version <= 2.67.0 dill is used to encode "special types" - In SDK version 2.68.0 cloudpickle is used to encode "special types" with absolute filepaths in code objects and dynamic functions. - In SDK version 2.69.0 cloudpickle is used to encode "special types" with relative filepaths in code objects and dynamic functions. """ is_using_dill = compat_version == "2.67.0" if is_using_dill: pytest.importorskip("dill") if sys.executable is None: self.skipTest('No Python interpreter found') # pylint: disable=line-too-long script = textwrap.dedent( f'''\ import pickle import sys import collections import enum import logging from apache_beam.coders import coders from apache_beam.options.pipeline_options_context import scoped_pipeline_options from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.coders.coders_test_common import MyNamedTuple from apache_beam.coders.coders_test_common import MyTypedNamedTuple from apache_beam.coders.coders_test_common import MyEnum from apache_beam.coders.coders_test_common import MyIntEnum from apache_beam.coders.coders_test_common import MyIntFlag from apache_beam.coders.coders_test_common import MyFlag from apache_beam.coders.coders_test_common import DefinesGetState from apache_beam.coders.coders_test_common import DefinesGetAndSetState from apache_beam.coders.coders_test_common import FrozenDataClass from apache_beam.coders.coders_test_common import FrozenKwOnlyDataClass from apache_beam.coders.coders_test_common import FrozenUnInitKwOnlyDataClass from apache_beam.coders import proto2_coder_test_messages_pb2 as test_message logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', stream=sys.stderr, force=True ) # Test cases for all special deterministic types # NOTE: When this script run in a subprocess the module is considered # __main__. Dill cannot pickle enums in __main__ because it # needs to define a way to create the type if it does not exist # in the session, and reaches recursion depth limits. test_cases = [ ("proto_message", test_message.MessageA(field1='value')), ("named_tuple_simple", MyNamedTuple(1, 2)), ("typed_named_tuple", MyTypedNamedTuple(1, 'a')), ("named_tuple_list", [MyNamedTuple(1, 2), MyTypedNamedTuple(1, 'a')]), ("enum_single", MyEnum.E1), ("enum_list", list(MyEnum)), ("int_enum_list", list(MyIntEnum)), ("int_flag_list", list(MyIntFlag)), ("flag_list", list(MyFlag)), ("getstate_setstate_simple", DefinesGetAndSetState(1)), ("getstate_setstate_complex", DefinesGetAndSetState((1, 2, 3))), ("getstate_setstate_list", [DefinesGetAndSetState(1), DefinesGetAndSetState((1, 2, 3))]), ] test_cases.extend([ ("frozen_dataclass", FrozenDataClass(1, 2)), ("frozen_dataclass_list", [FrozenDataClass(1, 2), FrozenDataClass(3, 4)]), ("frozen_kwonly_dataclass", FrozenKwOnlyDataClass(c=1, d=2)), ("frozen_kwonly_dataclass_list", [FrozenKwOnlyDataClass(c=1, d=2), FrozenUnInitKwOnlyDataClass(side=3)]), ]) compat_version = {'"'+ compat_version +'"' if compat_version else None} with scoped_pipeline_options(PipelineOptions(update_compatibility_version=compat_version)): coder = coders.FastPrimitivesCoder() deterministic_coder = coder.as_deterministic_coder("step") results = dict() for test_name, value in test_cases: try: encoded = deterministic_coder.encode(value) results[test_name] = encoded except Exception as e: logging.warning("Encoding failed with %s", e) sys.exit(1) sys.stdout.buffer.write(pickle.dumps(results)) ''') def run_subprocess(): result = subprocess.run([sys.executable, '-c', script], capture_output=True, timeout=30, check=False) self.assertEqual( 0, result.returncode, f"Subprocess failed: {result.stderr}") return pickle.loads(result.stdout) results1 = run_subprocess() results2 = run_subprocess() coder = coders.FastPrimitivesCoder() deterministic_coder = coder.as_deterministic_coder("step") for test_name in results1: data1 = results1[test_name] data2 = results2[test_name] self.assertEqual( data1, data2, f"Cross-process encoding differs for {test_name}") self.assertGreater(len(data1), 1) try: decoded1 = deterministic_coder.decode(data1) decoded2 = deterministic_coder.decode(data2) except Exception as e: logging.warning("Could not decode %s data due to %s", test_name, e) continue if test_name == "named_tuple_simple" and not is_using_dill: # The absense of a compat_version means we are using the most recent # implementation of the coder, which uses relative paths. should_have_relative_path = not compat_version named_tuple_type = type(decoded1) self.assertEqual( os.path.isabs(named_tuple_type._make.__code__.co_filename), not should_have_relative_path) self.assertEqual( os.path.isabs( named_tuple_type.__getnewargs__.__globals__['__file__']), not should_have_relative_path) self.assertEqual( decoded1, decoded2, f"Cross-process decoding differs for {test_name}") self.assertIsInstance( decoded1, type(decoded2), f"Cross-process decoding differs for {test_name}") def test_proto_coder(self): # For instructions on how these test proto message were generated, # see coders_test.py ma = test_message.MessageA() mab = ma.field2.add() mab.field1 = True ma.field1 = 'hello world' mb = test_message.MessageA() mb.field1 = 'beam' proto_coder = coders.ProtoCoder(ma.__class__) self.check_coder(proto_coder, ma) self.check_coder( coders.TupleCoder((proto_coder, coders.BytesCoder())), (ma, b'a'), (mb, b'b')) def test_global_window_coder(self): coder = coders.GlobalWindowCoder() value = window.GlobalWindow() # Test binary representation self.assertEqual(b'', coder.encode(value)) self.assertEqual(value, coder.decode(b'')) # Test unnested self.check_coder(coder, value) # Test nested self.check_coder(coders.TupleCoder((coder, coder)), (value, value)) def test_length_prefix_coder(self): coder = coders.LengthPrefixCoder(coders.BytesCoder()) # Test binary representation self.assertEqual(b'\x00', coder.encode(b'')) self.assertEqual(b'\x01a', coder.encode(b'a')) self.assertEqual(b'\x02bc', coder.encode(b'bc')) self.assertEqual(b'\xff\x7f' + b'z' * 16383, coder.encode(b'z' * 16383)) # Test unnested self.check_coder(coder, b'', b'a', b'bc', b'def') # Test nested self.check_coder( coders.TupleCoder((coder, coder)), (b'', b'a'), (b'bc', b'def')) def test_nested_observables(self): class FakeObservableIterator(observable.ObservableMixin): def __iter__(self): return iter([1, 2, 3]) # Coder for elements from the observable iterator. elem_coder = coders.VarIntCoder() iter_coder = coders.TupleSequenceCoder(elem_coder) # Test nested WindowedValue observable. coder = coders.WindowedValueCoder(iter_coder) observ = FakeObservableIterator() value = windowed_value.WindowedValue(observ, 0, ()) self.assertEqual( coder.get_impl().get_estimated_size_and_observables(value)[1], [(observ, elem_coder.get_impl())]) # Test nested tuple observable. coder = coders.TupleCoder((coders.StrUtf8Coder(), iter_coder)) value = ('123', observ) self.assertEqual( coder.get_impl().get_estimated_size_and_observables(value)[1], [(observ, elem_coder.get_impl())]) def test_state_backed_iterable_coder(self): # pylint: disable=global-variable-undefined # required for pickling by reference global state state = {} def iterable_state_write(values, element_coder_impl): token = b'state_token_%d' % len(state) state[token] = [element_coder_impl.encode(e) for e in values] return token def iterable_state_read(token, element_coder_impl): return [element_coder_impl.decode(s) for s in state[token]] coder = coders.StateBackedIterableCoder( coders.VarIntCoder(), read_state=iterable_state_read, write_state=iterable_state_write, write_state_threshold=1) # Note: do not use check_coder # see https://github.com/cloudpipe/cloudpickle/issues/452 self._observe(coder) self.assertEqual([1, 2, 3], coder.decode(coder.encode([1, 2, 3]))) # Ensure that state was actually used. self.assertNotEqual(state, {}) tupleCoder = coders.TupleCoder((coder, coder)) self._observe(tupleCoder) self.assertEqual(([1], [2, 3]), tupleCoder.decode(tupleCoder.encode(([1], [2, 3])))) def test_nullable_coder(self): self.check_coder(coders.NullableCoder(coders.VarIntCoder()), None, 2 * 64) def test_map_coder(self): values = [ { 1: "one", 300: "three hundred" }, # force yapf to be nice {}, { i: str(i) for i in range(5000) }, ] map_coder = coders.MapCoder(coders.VarIntCoder(), coders.StrUtf8Coder()) self.check_coder(map_coder, *values) self.check_coder(map_coder.as_deterministic_coder("label"), *values) def test_sharded_key_coder(self): key_and_coders = [(b'', b'\x00', coders.BytesCoder()), (b'key', b'\x03key', coders.BytesCoder()), ('key', b'\03\x6b\x65\x79', coders.StrUtf8Coder()), (('k', 1), b'\x01\x6b\x01', coders.TupleCoder( (coders.StrUtf8Coder(), coders.VarIntCoder())))] for key, bytes_repr, key_coder in key_and_coders: coder = coders.ShardedKeyCoder(key_coder) # Test str repr self.assertEqual('%s' % coder, 'ShardedKeyCoder[%s]' % key_coder) self.assertEqual(b'\x00' + bytes_repr, coder.encode(ShardedKey(key, b''))) self.assertEqual( b'\x03123' + bytes_repr, coder.encode(ShardedKey(key, b'123'))) # Test unnested self.check_coder(coder, ShardedKey(key, b'')) self.check_coder(coder, ShardedKey(key, b'123')) # Test type hints self.assertTrue( isinstance( coder.to_type_hint(), sharded_key_type.ShardedKeyTypeConstraint)) key_type = coder.to_type_hint().key_type if isinstance(key_type, typehints.TupleConstraint): self.assertEqual(key_type.tuple_types, (type(key[0]), type(key[1]))) else: self.assertEqual(key_type, type(key)) self.assertEqual( coders.ShardedKeyCoder.from_type_hint( coder.to_type_hint(), typecoders.CoderRegistry()), coder) for other_key, _, other_key_coder in key_and_coders: other_coder = coders.ShardedKeyCoder(other_key_coder) # Test nested self.check_coder( coders.TupleCoder((coder, other_coder)), (ShardedKey(key, b''), ShardedKey(other_key, b''))) self.check_coder( coders.TupleCoder((coder, other_coder)), (ShardedKey(key, b'123'), ShardedKey(other_key, b''))) def test_timestamp_prefixing_window_coder(self): self.check_coder( coders.TimestampPrefixingWindowCoder(coders.IntervalWindowCoder()), *[ window.IntervalWindow(x, y) for x in [-2**52, 0, 2**52] for y in range(-100, 100) ]) self.check_coder( coders.TupleCoder(( coders.TimestampPrefixingWindowCoder( coders.IntervalWindowCoder()), )), (window.IntervalWindow(0, 10), )) def test_timestamp_prefixing_opaque_window_coder(self): sdk_coder = coders.TimestampPrefixingWindowCoder( coders.LengthPrefixCoder(coders.PickleCoder())) safe_coder = coders.TimestampPrefixingOpaqueWindowCoder() for w in [window.IntervalWindow(1, 123), window.GlobalWindow()]: round_trip = sdk_coder.decode( safe_coder.encode(safe_coder.decode(sdk_coder.encode(w)))) self.assertEqual(w, round_trip) def test_decimal_coder(self): test_coder = coders.DecimalCoder() test_values = [ Decimal("-10.5"), Decimal("-1"), Decimal(), Decimal("1"), Decimal("13.258"), ] test_encodings = ("AZc", "AP8", "AAA", "AAE", "AzPK") self.check_coder(test_coder, *test_values) for idx, value in enumerate(test_values): self.assertEqual( test_encodings[idx], base64.b64encode(test_coder.encode(value)).decode().rstrip("=")) def test_byte_coder(self): test_coder = coders.ByteCoder() test_values = [0, 80, 127, 128, 255] test_encodings = ("AA", "UA", "fw", "gA", "/w") self.check_coder(test_coder, *test_values) for idx, value in enumerate(test_values): self.assertEqual( test_encodings[idx], base64.b64encode(test_coder.encode(value)).decode().rstrip("=")) def test_OrderedUnionCoder(self): test_coder = coders._OrderedUnionCoder((str, coders.StrUtf8Coder()), (int, coders.VarIntCoder()), fallback_coder=coders.FloatCoder()) self.check_coder(test_coder, 's') self.check_coder(test_coder, 123) self.check_coder(test_coder, 1.5) def test_OrderedUnionCoderDeterministic(self): # CustomCoder is not deterministic therefore test_coder is not # deterministic test_coder = coders._OrderedUnionCoder((str, coders.StrUtf8Coder()), (int, CustomCoder()), fallback_coder=coders.FloatCoder()) self.assertFalse(test_coder.is_deterministic()) test_coder = coders._OrderedUnionCoder((str, coders.StrUtf8Coder()), (int, coders.VarIntCoder()), fallback_coder=coders.FloatCoder()) self.assertTrue(test_coder.is_deterministic()) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/fast_coders_test.py0000644000175100017510000000267715206113032023055 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for compiled implementation of coder impls.""" # pytype: skip-file import logging import unittest # Run all the standard coder test cases. from apache_beam.coders.coders_test_common import * from apache_beam.tools import utils class FastCoders(unittest.TestCase): def test_using_fast_impl(self): try: utils.check_compiled('apache_beam.coders.coder_impl') except RuntimeError: self.skipTest('Cython is not installed') # pylint: disable=wrong-import-order, wrong-import-position # pylint: disable=unused-import import apache_beam.coders.stream if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/observable.py0000644000175100017510000000261015206113032021631 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Observable base class for iterables. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file class ObservableMixin(object): """For internal use only; no backwards-compatibility guarantees. An observable iterable. Subclasses need to call self.notify_observers with any object yielded. """ def __init__(self): self.observers = [] def register_observer(self, callback): self.observers.append(callback) def notify_observers(self, value, **kwargs): # self.observers is almost always empty if self.observers: for o in self.observers: o(value, **kwargs) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/observable_test.py0000644000175100017510000000340415206113032022672 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for the Observable mixin class.""" # pytype: skip-file import logging import unittest from typing import Optional from apache_beam.coders import observable class ObservableMixinTest(unittest.TestCase): observed_count = 0 observed_sum = 0 observed_keys: list[Optional[str]] = [] def observer(self, value, key=None): self.observed_count += 1 self.observed_sum += value self.observed_keys.append(key) def test_observable(self): class Watched(observable.ObservableMixin): def __iter__(self): for i in (1, 4, 3): self.notify_observers(i, key='a%d' % i) yield i watched = Watched() watched.register_observer(lambda v, key: self.observer(v, key=key)) for _ in watched: pass self.assertEqual(3, self.observed_count) self.assertEqual(8, self.observed_sum) self.assertEqual(['a1', 'a3', 'a4'], sorted(self.observed_keys)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/proto2_coder_test_messages_pb2.py0000644000175100017510000001011115206113032025572 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # NO CHECKED-IN PROTOBUF GENCODE # source: apache_beam/coders/proto2_coder_test_messages.proto # Protobuf Python Version: 5.28.0 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( b'\n3apache_beam/coders/proto2_coder_test_messages.proto\x12\x1aproto2_coder_test_messages\"P\n\x08MessageA\x12\x0e\n\x06\x66ield1\x18\x01 \x01(\t\x12\x34\n\x06\x66ield2\x18\x02 \x03(\x0b\x32$.proto2_coder_test_messages.MessageB\"\x1a\n\x08MessageB\x12\x0e\n\x06\x66ield1\x18\x01 \x01(\x08\"\x10\n\x08MessageC*\x04\x08\x64\x10j\"\xad\x01\n\x0eMessageWithMap\x12\x46\n\x06\x66ield1\x18\x01 \x03(\x0b\x32\x36.proto2_coder_test_messages.MessageWithMap.Field1Entry\x1aS\n\x0b\x46ield1Entry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x33\n\x05value\x18\x02 \x01(\x0b\x32$.proto2_coder_test_messages.MessageA:\x02\x38\x01\"V\n\x18ReferencesMessageWithMap\x12:\n\x06\x66ield1\x18\x01 \x03(\x0b\x32*.proto2_coder_test_messages.MessageWithMap\";\n\x08MessageD\"/\n\nNestedEnum\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x07\n\x03ONE\x10\x01\x12\x07\n\x03TWO\x10\x02*1\n\x0cTopLevelEnum\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x07\n\x03ONE\x10\x01\x12\x07\n\x03TWO\x10\x02:Z\n\x06\x66ield1\x12$.proto2_coder_test_messages.MessageC\x18\x65 \x01(\x0b\x32$.proto2_coder_test_messages.MessageA:Z\n\x06\x66ield2\x12$.proto2_coder_test_messages.MessageC\x18\x66 \x01(\x0b\x32$.proto2_coder_test_messages.MessageBB)\n\'org.apache.beam.sdk.extensions.protobuf' ) _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages( DESCRIPTOR, 'apache_beam.coders.proto2_coder_test_messages_pb2', _globals) if not _descriptor._USE_C_DESCRIPTORS: _globals['DESCRIPTOR']._loaded_options = None _globals[ 'DESCRIPTOR']._serialized_options = b'\n\'org.apache.beam.sdk.extensions.protobuf' _globals['_MESSAGEWITHMAP_FIELD1ENTRY']._loaded_options = None _globals['_MESSAGEWITHMAP_FIELD1ENTRY']._serialized_options = b'8\001' _globals['_TOPLEVELENUM']._serialized_start = 536 _globals['_TOPLEVELENUM']._serialized_end = 585 _globals['_MESSAGEA']._serialized_start = 83 _globals['_MESSAGEA']._serialized_end = 163 _globals['_MESSAGEB']._serialized_start = 165 _globals['_MESSAGEB']._serialized_end = 191 _globals['_MESSAGEC']._serialized_start = 193 _globals['_MESSAGEC']._serialized_end = 209 _globals['_MESSAGEWITHMAP']._serialized_start = 212 _globals['_MESSAGEWITHMAP']._serialized_end = 385 _globals['_MESSAGEWITHMAP_FIELD1ENTRY']._serialized_start = 302 _globals['_MESSAGEWITHMAP_FIELD1ENTRY']._serialized_end = 385 _globals['_REFERENCESMESSAGEWITHMAP']._serialized_start = 387 _globals['_REFERENCESMESSAGEWITHMAP']._serialized_end = 473 _globals['_MESSAGED']._serialized_start = 475 _globals['_MESSAGED']._serialized_end = 534 _globals['_MESSAGED_NESTEDENUM']._serialized_start = 487 _globals['_MESSAGED_NESTEDENUM']._serialized_end = 534 # @@protoc_insertion_point(module_scope) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/row_coder.py0000644000175100017510000001776215206113032021506 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file from apache_beam.coders import typecoders from apache_beam.coders.coder_impl import LogicalTypeCoderImpl from apache_beam.coders.coder_impl import RowCoderImpl from apache_beam.coders.coders import BigEndianShortCoder from apache_beam.coders.coders import BooleanCoder from apache_beam.coders.coders import ByteCoder from apache_beam.coders.coders import BytesCoder from apache_beam.coders.coders import Coder from apache_beam.coders.coders import DecimalCoder from apache_beam.coders.coders import FastCoder from apache_beam.coders.coders import FloatCoder from apache_beam.coders.coders import IterableCoder from apache_beam.coders.coders import MapCoder from apache_beam.coders.coders import NullableCoder from apache_beam.coders.coders import SinglePrecisionFloatCoder from apache_beam.coders.coders import StrUtf8Coder from apache_beam.coders.coders import TimestampCoder from apache_beam.coders.coders import VarInt32Coder from apache_beam.coders.coders import VarIntCoder from apache_beam.portability import common_urns from apache_beam.portability.api import schema_pb2 from apache_beam.typehints import row_type from apache_beam.typehints.schemas import PYTHON_ANY_URN from apache_beam.typehints.schemas import LogicalType from apache_beam.typehints.schemas import named_tuple_from_schema from apache_beam.typehints.schemas import schema_from_element_type from apache_beam.utils import proto_utils __all__ = ["RowCoder"] class RowCoder(FastCoder): """ Coder for `typing.NamedTuple` instances. Implements the beam:coder:row:v1 standard coder spec. """ def __init__(self, schema, force_deterministic=False): """Initializes a :class:`RowCoder`. Args: schema (apache_beam.portability.api.schema_pb2.Schema): The protobuf representation of the schema of the data that the RowCoder will be used to encode/decode. """ self.schema = schema # Eagerly generate type hint to escalate any issues with the Schema proto self._type_hint = named_tuple_from_schema(self.schema) # Use non-null coders because null values are represented separately self.components = [ _nonnull_coder_from_type(field.type) for field in self.schema.fields ] if force_deterministic: self.components = [ c.as_deterministic_coder(force_deterministic) for c in self.components ] self.forced_deterministic = bool(force_deterministic) def _create_impl(self): return RowCoderImpl(self.schema, self.components) def is_deterministic(self): return all(c.is_deterministic() for c in self.components) def as_deterministic_coder(self, step_label, error_message=None): if self.is_deterministic(): return self else: return RowCoder(self.schema, error_message or step_label) def to_type_hint(self): return self._type_hint def __hash__(self): return hash(self.schema.SerializeToString()) def __eq__(self, other): return ( type(self) == type(other) and self.schema == other.schema and self.forced_deterministic == other.forced_deterministic) def to_runner_api_parameter(self, unused_context): return (common_urns.coders.ROW.urn, self.schema, []) @staticmethod @Coder.register_urn(common_urns.coders.ROW.urn, schema_pb2.Schema) def from_runner_api_parameter(schema, components, unused_context): return RowCoder(schema) @classmethod def from_type_hint(cls, type_hint, registry): # TODO(https://github.com/apache/beam/issues/21541): Remove once all # runners are portable. if isinstance(type_hint, str): import importlib main_module = importlib.import_module('__main__') type_hint = getattr(main_module, type_hint, type_hint) schema = schema_from_element_type(type_hint) return cls(schema) @staticmethod def from_payload(payload: bytes) -> 'RowCoder': return RowCoder(proto_utils.parse_Bytes(payload, schema_pb2.Schema)) def __reduce__(self): # when pickling, use bytes representation of the schema. schema_pb2.Schema # objects cannot be pickled. return (RowCoder.from_payload, (self.schema.SerializeToString(), )) typecoders.registry.register_coder(row_type.RowTypeConstraint, RowCoder) typecoders.registry.register_coder( row_type.GeneratedClassRowTypeConstraint, RowCoder) def _coder_from_type(field_type): coder = _nonnull_coder_from_type(field_type) if field_type.nullable: return NullableCoder(coder) else: return coder def _nonnull_coder_from_type(field_type): type_info = field_type.WhichOneof("type_info") if type_info == "atomic_type": if field_type.atomic_type == schema_pb2.INT64: return VarIntCoder() elif field_type.atomic_type == schema_pb2.INT32: return VarInt32Coder() elif field_type.atomic_type == schema_pb2.INT16: return BigEndianShortCoder() elif field_type.atomic_type == schema_pb2.BYTE: return ByteCoder() elif field_type.atomic_type == schema_pb2.FLOAT: return SinglePrecisionFloatCoder() elif field_type.atomic_type == schema_pb2.DOUBLE: return FloatCoder() elif field_type.atomic_type == schema_pb2.STRING: return StrUtf8Coder() elif field_type.atomic_type == schema_pb2.BOOLEAN: return BooleanCoder() elif field_type.atomic_type == schema_pb2.BYTES: return BytesCoder() elif type_info == "array_type": return IterableCoder(_coder_from_type(field_type.array_type.element_type)) elif type_info == "iterable_type": return IterableCoder( _coder_from_type(field_type.iterable_type.element_type)) elif type_info == "map_type": return MapCoder( _coder_from_type(field_type.map_type.key_type), _coder_from_type(field_type.map_type.value_type)) elif type_info == "logical_type": if field_type.logical_type.urn == PYTHON_ANY_URN: # Special case for the Any logical type. Just use the default coder for an # unknown Python object. return typecoders.registry.get_coder(object) elif field_type.logical_type.urn == common_urns.millis_instant.urn: # Special case for millis instant logical type used to handle Java sdk's # millis Instant. It explicitly uses TimestampCoder which deals with fix # length 8-bytes big-endian-long instead of VarInt coder. return TimestampCoder() elif field_type.logical_type.urn == 'beam:logical_type:decimal:v1': return DecimalCoder() logical_type = LogicalType.from_runner_api(field_type.logical_type) return LogicalTypeCoder( logical_type, _coder_from_type(field_type.logical_type.representation)) elif type_info == "row_type": return RowCoder(field_type.row_type.schema) # The Java SDK supports several more types, but the coders are not yet # standard, and are not implemented in Python. raise ValueError( "Encountered a type that is not currently supported by RowCoder: %s" % field_type) class LogicalTypeCoder(FastCoder): def __init__(self, logical_type, representation_coder): self.logical_type = logical_type self.representation_coder = representation_coder def _create_impl(self): return LogicalTypeCoderImpl(self.logical_type, self.representation_coder) def is_deterministic(self): return self.representation_coder.is_deterministic() def to_type_hint(self): return self.logical_type.language_type() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/row_coder_test.py0000644000175100017510000004046515206113032022541 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import typing import unittest from itertools import chain import numpy as np from numpy.testing import assert_array_equal import apache_beam as beam from apache_beam.coders import RowCoder from apache_beam.coders import coder_impl from apache_beam.coders.typecoders import registry as coders_registry from apache_beam.internal import pickler from apache_beam.portability.api import schema_pb2 from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.typehints.schemas import named_tuple_from_schema from apache_beam.typehints.schemas import typing_to_runner_api from apache_beam.utils.timestamp import Timestamp Person = typing.NamedTuple( "Person", [ ("name", str), ("age", np.int32), ("address", typing.Optional[str]), ("aliases", list[str]), ("knows_javascript", bool), ("payload", typing.Optional[bytes]), ("custom_metadata", typing.Mapping[str, int]), ("favorite_time", Timestamp), ]) NullablePerson = typing.NamedTuple( "NullablePerson", [("name", typing.Optional[str]), ("age", np.int32), ("address", typing.Optional[str]), ("aliases", list[str]), ("knows_javascript", bool), ("payload", typing.Optional[bytes]), ("custom_metadata", typing.Mapping[str, int]), ("favorite_time", typing.Optional[Timestamp]), ("one_more_field", typing.Optional[str])]) @coders_registry.register_row class People(typing.NamedTuple): primary: Person partner: typing.Optional[Person] coders_registry.register_row(Person) class RowCoderTest(unittest.TestCase): JON_SNOW = Person( name="Jon Snow", age=np.int32(23), address=None, aliases=["crow", "wildling"], knows_javascript=False, payload=None, custom_metadata={}, favorite_time=Timestamp.from_rfc3339('2016-03-18T23:22:59.123456Z'), ) PEOPLE = [ JON_SNOW, Person( "Daenerys Targaryen", np.int32(25), "Westeros", ["Mother of Dragons"], False, None, {"dragons": 3}, Timestamp.from_rfc3339('1970-04-26T17:46:40Z'), ), Person( "Michael Bluth", np.int32(30), None, [], True, b"I've made a huge mistake", {}, Timestamp.from_rfc3339('2020-08-12T15:51:00.032Z')) ] def test_row_accepts_trailing_zeros_truncated(self): expected_coder = RowCoder( typing_to_runner_api(NullablePerson).row_type.schema) person = NullablePerson( None, np.int32(25), "Westeros", ["Mother of Dragons"], False, None, {"dragons": 3}, None, "NotNull") out = expected_coder.encode(person) # 9 fields, 1 null byte, field 0, 5, 7 are null new_payload = bytes([9, 1, 1 | 1 << 5 | 1 << 7]) + out[4:] new_value = expected_coder.decode(new_payload) self.assertEqual(person, new_value) def test_create_row_coder_from_named_tuple(self): expected_coder = RowCoder(typing_to_runner_api(Person).row_type.schema) real_coder = coders_registry.get_coder(Person) for test_case in self.PEOPLE: self.assertEqual( expected_coder.encode(test_case), real_coder.encode(test_case)) self.assertEqual( test_case, real_coder.decode(real_coder.encode(test_case))) def test_create_row_coder_from_nested_named_tuple(self): expected_coder = RowCoder(typing_to_runner_api(People).row_type.schema) real_coder = coders_registry.get_coder(People) for primary in self.PEOPLE: for other in self.PEOPLE + [None]: test_case = People(primary=primary, partner=other) self.assertEqual( expected_coder.encode(test_case), real_coder.encode(test_case)) self.assertEqual( test_case, real_coder.decode(real_coder.encode(test_case))) def test_create_row_coder_from_schema(self): schema = schema_pb2.Schema( id="person", fields=[ schema_pb2.Field( name="name", type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), schema_pb2.Field( name="age", type=schema_pb2.FieldType(atomic_type=schema_pb2.INT32)), schema_pb2.Field( name="address", type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING, nullable=True)), schema_pb2.Field( name="aliases", type=schema_pb2.FieldType( array_type=schema_pb2.ArrayType( element_type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING)))), schema_pb2.Field( name="knows_javascript", type=schema_pb2.FieldType(atomic_type=schema_pb2.BOOLEAN)), schema_pb2.Field( name="payload", type=schema_pb2.FieldType( atomic_type=schema_pb2.BYTES, nullable=True)), schema_pb2.Field( name="custom_metadata", type=schema_pb2.FieldType( map_type=schema_pb2.MapType( key_type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING), value_type=schema_pb2.FieldType( atomic_type=schema_pb2.INT64), ))), schema_pb2.Field( name="favorite_time", type=schema_pb2.FieldType( logical_type=schema_pb2.LogicalType( urn="beam:logical_type:micros_instant:v1", representation=schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( id="micros_instant", fields=[ schema_pb2.Field( name="seconds", type=schema_pb2.FieldType( atomic_type=schema_pb2.INT64)), schema_pb2.Field( name="micros", type=schema_pb2.FieldType( atomic_type=schema_pb2.INT64)), ])))))), ]) coder = RowCoder(schema) for test_case in self.PEOPLE: self.assertEqual(test_case, coder.decode(coder.encode(test_case))) def test_row_coder_negative_varint(self): schema = schema_pb2.Schema( id="negative", fields=[ schema_pb2.Field( name="i64", type=schema_pb2.FieldType(atomic_type=schema_pb2.INT64)), schema_pb2.Field( name="i32", type=schema_pb2.FieldType(atomic_type=schema_pb2.INT32)) ]) coder = RowCoder(schema) Negative = typing.NamedTuple( "Negative", [ ("i64", np.int64), ("i32", np.int32), ]) test_cases = [ Negative(-1, -1023), Negative(-1023, -1), Negative(-2**63, -2**31) ] for test_case in test_cases: self.assertEqual(test_case, coder.decode(coder.encode(test_case))) @unittest.skip( "https://github.com/apache/beam/issues/19696 - Overflow behavior in " "VarIntCoder is currently inconsistent") def test_overflows(self): IntTester = typing.NamedTuple( 'IntTester', [ # TODO(https://github.com/apache/beam/issues/19815): Test int8 and # int16 here as well when those types are supported # ('i8', typing.Optional[np.int8]), # ('i16', typing.Optional[np.int16]), ('i32', typing.Optional[np.int32]), ('i64', typing.Optional[np.int64]), ]) c = RowCoder.from_type_hint(IntTester, None) no_overflow = chain( (IntTester(i32=i, i64=None) for i in (-2**31, 2**31 - 1)), (IntTester(i32=None, i64=i) for i in (-2**63, 2**63 - 1)), ) # Encode max/min ints to make sure they don't throw any error for case in no_overflow: c.encode(case) overflow = chain( (IntTester(i32=i, i64=None) for i in (-2**31 - 1, 2**31)), (IntTester(i32=None, i64=i) for i in (-2**63 - 1, 2**63)), ) # Encode max+1/min-1 ints to make sure they DO throw an error # pylint: disable=cell-var-from-loop for case in overflow: self.assertRaises(OverflowError, lambda: c.encode(case)) def test_none_in_non_nullable_field_throws(self): Test = typing.NamedTuple('Test', [('foo', str)]) c = RowCoder.from_type_hint(Test, None) self.assertRaises(ValueError, lambda: c.encode(Test(foo=None))) def test_schema_remove_column(self): fields = [("field1", str), ("field2", str)] # new schema is missing one field that was in the old schema Old = typing.NamedTuple('Old', fields) New = typing.NamedTuple('New', fields[:-1]) old_coder = RowCoder.from_type_hint(Old, None) new_coder = RowCoder.from_type_hint(New, None) self.assertEqual( New("foo"), new_coder.decode(old_coder.encode(Old("foo", "bar")))) def test_schema_add_column(self): fields = [("field1", str), ("field2", typing.Optional[str])] # new schema has one (optional) field that didn't exist in the old schema Old = typing.NamedTuple('Old', fields[:-1]) New = typing.NamedTuple('New', fields) old_coder = RowCoder.from_type_hint(Old, None) new_coder = RowCoder.from_type_hint(New, None) self.assertEqual( New("bar", None), new_coder.decode(old_coder.encode(Old("bar")))) def test_schema_add_column_with_null_value(self): fields = [("field1", typing.Optional[str]), ("field2", str), ("field3", typing.Optional[str])] # new schema has one (optional) field that didn't exist in the old schema Old = typing.NamedTuple('Old', fields[:-1]) New = typing.NamedTuple('New', fields) old_coder = RowCoder.from_type_hint(Old, None) new_coder = RowCoder.from_type_hint(New, None) self.assertEqual( New(None, "baz", None), new_coder.decode(old_coder.encode(Old(None, "baz")))) def test_row_coder_picklable(self): # occasionally coders can get pickled, RowCoder should be able to handle it coder = coders_registry.get_coder(Person) roundtripped = pickler.loads(pickler.dumps(coder)) self.assertEqual(roundtripped, coder) def test_row_coder_in_pipeine(self): with TestPipeline() as p: res = ( p | beam.Create(self.PEOPLE) | beam.Filter(lambda person: person.name == "Jon Snow")) assert_that(res, equal_to([self.JON_SNOW])) def test_row_coder_nested_struct(self): Pair = typing.NamedTuple('Pair', [('left', Person), ('right', Person)]) value = Pair(self.PEOPLE[0], self.PEOPLE[1]) coder = RowCoder(typing_to_runner_api(Pair).row_type.schema) self.assertEqual(value, coder.decode(coder.encode(value))) def test_encoding_position_reorder_fields(self): schema1 = schema_pb2.Schema( id="reorder_test_schema1", fields=[ schema_pb2.Field( name="f_int32", type=schema_pb2.FieldType(atomic_type=schema_pb2.INT32), ), schema_pb2.Field( name="f_str", type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING), ), ]) schema2 = schema_pb2.Schema( id="reorder_test_schema2", encoding_positions_set=True, fields=[ schema_pb2.Field( name="f_str", type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING), encoding_position=1, ), schema_pb2.Field( name="f_int32", type=schema_pb2.FieldType(atomic_type=schema_pb2.INT32), encoding_position=0, ), ]) RowSchema1 = named_tuple_from_schema(schema1) RowSchema2 = named_tuple_from_schema(schema2) roundtripped = RowCoder(schema2).decode( RowCoder(schema1).encode(RowSchema1(42, "Hello World!"))) self.assertEqual(RowSchema2(f_int32=42, f_str="Hello World!"), roundtripped) def test_encoding_position_add_fields_and_reorder(self): old_schema = schema_pb2.Schema( id="add_test_old", fields=[ schema_pb2.Field( name="f_int32", type=schema_pb2.FieldType(atomic_type=schema_pb2.INT32), ), schema_pb2.Field( name="f_str", type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING), ), ]) new_schema = schema_pb2.Schema( encoding_positions_set=True, id="add_test_new", fields=[ schema_pb2.Field( name="f_new_str", type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING, nullable=True), encoding_position=2, ), schema_pb2.Field( name="f_int32", type=schema_pb2.FieldType(atomic_type=schema_pb2.INT32), encoding_position=0, ), schema_pb2.Field( name="f_str", type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING), encoding_position=1, ), ]) Old = named_tuple_from_schema(old_schema) New = named_tuple_from_schema(new_schema) roundtripped = RowCoder(new_schema).decode( RowCoder(old_schema).encode(Old(42, "Hello World!"))) self.assertEqual( New(f_new_str=None, f_int32=42, f_str="Hello World!"), roundtripped) def test_row_coder_fail_early_bad_schema(self): schema_proto = schema_pb2.Schema( fields=[ schema_pb2.Field( name="type_with_no_typeinfo", type=schema_pb2.FieldType()) ], id='bad-schema') # Should raise an exception referencing the problem field self.assertRaisesRegex( ValueError, "type_with_no_typeinfo", lambda: RowCoder(schema_proto)) def test_batch_encode_decode(self): coder = RowCoder(typing_to_runner_api(Person).row_type.schema).get_impl() seq_out = coder_impl.create_OutputStream() for person in self.PEOPLE: coder.encode_to_stream(person, seq_out, False) batch_out = coder_impl.create_OutputStream() columnar = { field: np.array([getattr(person, field) for person in self.PEOPLE], ndmin=1, dtype=object) for field in Person._fields } coder.encode_batch_to_stream(columnar, batch_out) if seq_out.get() != batch_out.get(): a, b = seq_out.get(), batch_out.get() N = 25 for k in range(0, max(len(a), len(b)), N): print(k, a[k:k + N] == b[k:k + N]) print(a[k:k + N]) print(b[k:k + N]) self.assertEqual(seq_out.get(), batch_out.get()) for size in [len(self.PEOPLE) - 1, len(self.PEOPLE), len(self.PEOPLE) + 1]: dest = { field: np.ndarray((size, ), dtype=a.dtype) for field, a in columnar.items() } n = min(size, len(self.PEOPLE)) self.assertEqual( n, coder.decode_batch_from_stream( dest, coder_impl.create_InputStream(seq_out.get()))) for field, a in columnar.items(): assert_array_equal(a[:n], dest[field][:n]) if __name__ == "__main__": logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/slow_coders_test.py0000644000175100017510000000332715206113032023075 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for uncompiled implementation of coder impls.""" # pytype: skip-file import logging import unittest # Run all the standard coder test cases. from apache_beam.coders.coders_test_common import * @unittest.skip( 'Remove non-cython tests.' 'https://github.com/apache/beam/issues/28307') class SlowCoders(unittest.TestCase): def test_using_slow_impl(self): try: # pylint: disable=wrong-import-position # pylint: disable=unused-import from Cython.Build import cythonize self.skipTest('Found cython, cannot test non-compiled implementation.') except ImportError: # Assert that we are not using the compiled implementation. with self.assertRaises(ImportError): # pylint: disable=wrong-import-order, wrong-import-position # pylint: disable=unused-import import apache_beam.coders.stream if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/slow_stream.py0000644000175100017510000001222415206113032022046 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A pure Python implementation of stream.pyx. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file import struct class OutputStream(object): """For internal use only; no backwards-compatibility guarantees. A pure Python implementation of stream.OutputStream.""" def __init__(self): self.data: list[bytes] = [] self.byte_count = 0 def write(self, b: bytes, nested: bool = False) -> None: assert isinstance(b, bytes) if nested: self.write_var_int64(len(b)) self.data.append(b) self.byte_count += len(b) def write_byte(self, val): self.data.append(chr(val).encode('latin-1')) self.byte_count += 1 def write_var_int64(self, v: int) -> None: if v < 0: v += 1 << 64 if v <= 0: raise ValueError('Value too large (negative).') while True: bits = v & 0x7F v >>= 7 if v: bits |= 0x80 self.write_byte(bits) if not v: break def write_var_int32(self, v: int) -> None: self.write_var_int64(int(v) & 0xFFFFFFFF) def write_bigendian_int64(self, v): self.write(struct.pack('>q', v)) def write_bigendian_uint64(self, v): self.write(struct.pack('>Q', v)) def write_bigendian_int32(self, v): self.write(struct.pack('>i', v)) def write_bigendian_int16(self, v): self.write(struct.pack('>h', v)) def write_bigendian_double(self, v): self.write(struct.pack('>d', v)) def write_bigendian_float(self, v): self.write(struct.pack('>f', v)) def get(self) -> bytes: return b''.join(self.data) def size(self) -> int: return self.byte_count def _clear(self) -> None: self.data = [] self.byte_count = 0 class ByteCountingOutputStream(OutputStream): """For internal use only; no backwards-compatibility guarantees. A pure Python implementation of stream.ByteCountingOutputStream.""" def __init__(self): # Note that we don't actually use any of the data initialized by our super. super().__init__() self.count = 0 def write(self, byte_array: bytes, nested: bool = False) -> None: blen = len(byte_array) if nested: self.write_var_int64(blen) self.count += blen def write_byte(self, _): self.count += 1 def get_count(self): return self.count def get(self): raise NotImplementedError def __str__(self): return '<%s %s>' % (self.__class__.__name__, self.count) class InputStream(object): """For internal use only; no backwards-compatibility guarantees. A pure Python implementation of stream.InputStream.""" def __init__(self, data: bytes) -> None: self.data = data self.pos = 0 def size(self): return len(self.data) - self.pos def read(self, size: int) -> bytes: self.pos += size return self.data[self.pos - size:self.pos] def read_all(self, nested: bool) -> bytes: return self.read(self.read_var_int64() if nested else self.size()) def read_byte(self) -> int: self.pos += 1 return self.data[self.pos - 1] def read_var_int64(self): shift = 0 result = 0 while True: byte = self.read_byte() if byte < 0: raise RuntimeError('VarLong not terminated.') bits = byte & 0x7F if shift >= 64 or (shift >= 63 and bits > 1): raise RuntimeError('VarLong too long.') result |= bits << shift shift += 7 if not byte & 0x80: break if result >= 1 << 63: result -= 1 << 64 return result def read_var_int32(self): v = self.read_var_int64() return struct.unpack('q', self.read(8))[0] def read_bigendian_uint64(self): return struct.unpack('>Q', self.read(8))[0] def read_bigendian_int32(self): return struct.unpack('>i', self.read(4))[0] def read_bigendian_int16(self): return struct.unpack('>h', self.read(2))[0] def read_bigendian_double(self): return struct.unpack('>d', self.read(8))[0] def read_bigendian_float(self): return struct.unpack('>f', self.read(4))[0] def get_varint_size(v): """For internal use only; no backwards-compatibility guarantees. Returns the size of the given integer value when encode as a VarInt.""" if v < 0: v += 1 << 64 if v <= 0: raise ValueError('Value too large (negative).') varint_size = 0 while True: varint_size += 1 v >>= 7 if not v: break return varint_size ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/standard_coders_test.py0000644000175100017510000003132315206113032023706 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for coders that must be consistent across all Beam SDKs. """ # pytype: skip-file import json import logging import math import os.path import sys import unittest from copy import deepcopy import numpy as np import yaml from numpy.testing import assert_array_equal from apache_beam.coders import coder_impl from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.portability.api import schema_pb2 from apache_beam.runners import pipeline_context from apache_beam.transforms import userstate from apache_beam.transforms import window from apache_beam.transforms.window import IntervalWindow from apache_beam.typehints import schemas from apache_beam.utils import windowed_value from apache_beam.utils.sharded_key import ShardedKey from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.windowed_value import PaneInfo from apache_beam.utils.windowed_value import PaneInfoTiming STANDARD_CODERS_YAML = os.path.normpath( os.path.join( os.path.dirname(__file__), '../portability/api/standard_coders.yaml')) def _load_test_cases(test_yaml): """Load test data from yaml file and return an iterable of test cases. See ``standard_coders.yaml`` for more details. """ if not os.path.exists(test_yaml): raise ValueError('Could not find the test spec: %s' % test_yaml) with open(test_yaml, 'rb') as coder_spec: for ix, spec in enumerate( yaml.load_all(coder_spec, Loader=yaml.SafeLoader)): spec['index'] = ix name = spec.get('name', spec['coder']['urn'].split(':')[-2]) yield [name, spec] def parse_float(s): x = float(s) if math.isnan(x): # In Windows, float('NaN') has opposite sign from other platforms. # For the purpose of this test, we just need consistency. x = abs(x) return x def value_parser_from_schema(schema): def attribute_parser_from_type(type_): parser = nonnull_attribute_parser_from_type(type_) if type_.nullable: return lambda x: None if x is None else parser(x) else: return parser def nonnull_attribute_parser_from_type(type_): # TODO: This should be exhaustive type_info = type_.WhichOneof("type_info") if type_info == "atomic_type": if type_.atomic_type == schema_pb2.BYTES: return lambda x: x.encode("utf-8") else: return schemas.ATOMIC_TYPE_TO_PRIMITIVE[type_.atomic_type] elif type_info == "array_type": element_parser = attribute_parser_from_type(type_.array_type.element_type) return lambda x: list(map(element_parser, x)) elif type_info == "map_type": key_parser = attribute_parser_from_type(type_.map_type.key_type) value_parser = attribute_parser_from_type(type_.map_type.value_type) return lambda x: dict( (key_parser(k), value_parser(v)) for k, v in x.items()) elif type_info == "row_type": return value_parser_from_schema(type_.row_type.schema) elif type_info == "logical_type": # In YAML logical types are represented with their representation types. to_language_type = schemas.LogicalType.from_runner_api( type_.logical_type).to_language_type parse_representation = attribute_parser_from_type( type_.logical_type.representation) return lambda x: to_language_type(parse_representation(x)) parsers = [(field.name, attribute_parser_from_type(field.type)) for field in schema.fields] constructor = schemas.named_tuple_from_schema(schema) def value_parser(x): result = [] x = deepcopy(x) for name, parser in parsers: value = x.pop(name) result.append(None if value is None else parser(value)) if len(x): raise ValueError( "Test data contains attributes that don't exist in the schema: {}". format(', '.join(x.keys()))) return constructor(*result) return value_parser class StandardCodersTest(unittest.TestCase): _urn_to_json_value_parser = { 'beam:coder:bytes:v1': lambda x: x.encode('utf-8'), 'beam:coder:bool:v1': lambda x: x, 'beam:coder:string_utf8:v1': lambda x: x, 'beam:coder:varint:v1': lambda x: x, 'beam:coder:kv:v1': lambda x, key_parser, value_parser: (key_parser(x['key']), value_parser(x['value'])), 'beam:coder:interval_window:v1': lambda x: IntervalWindow( start=Timestamp(micros=(x['end'] - x['span']) * 1000), end=Timestamp( micros=x['end'] * 1000)), 'beam:coder:iterable:v1': lambda x, parser: list(map(parser, x)), 'beam:coder:state_backed_iterable:v1': lambda x, parser: list( map(parser, x)), 'beam:coder:global_window:v1': lambda x: window.GlobalWindow(), 'beam:coder:windowed_value:v1': lambda x, value_parser, window_parser: windowed_value.create( value_parser(x['value']), x['timestamp'] * 1000, tuple( window_parser(w) for w in x['windows'])), 'beam:coder:param_windowed_value:v1': lambda x, value_parser, window_parser: windowed_value.create( value_parser(x['value']), x['timestamp'] * 1000, tuple( window_parser(w) for w in x['windows']), PaneInfo( x['pane']['is_first'], x['pane']['is_last'], PaneInfoTiming. from_string(x['pane']['timing']), x['pane']['index'], x[ 'pane']['on_time_index'])), 'beam:coder:timer:v1': lambda x, value_parser, window_parser: userstate. Timer( user_key=value_parser(x['userKey']), dynamic_timer_tag=x[ 'dynamicTimerTag'], clear_bit=x['clearBit'], windows=tuple( window_parser(w) for w in x['windows']), fire_timestamp=None, hold_timestamp=None, paneinfo=None) if x['clearBit'] else userstate.Timer( user_key=value_parser(x['userKey']), dynamic_timer_tag=x[ 'dynamicTimerTag'], clear_bit=x['clearBit'], fire_timestamp= Timestamp(micros=x['fireTimestamp'] * 1000), hold_timestamp=Timestamp( micros=x['holdTimestamp'] * 1000), windows=tuple( window_parser(w) for w in x['windows']), paneinfo=PaneInfo( x['pane']['is_first'], x['pane']['is_last'], PaneInfoTiming.from_string(x['pane']['timing']), x[ 'pane']['index'], x['pane']['on_time_index'])), 'beam:coder:double:v1': parse_float, 'beam:coder:sharded_key:v1': lambda x, value_parser: ShardedKey( key=value_parser(x['key']), shard_id=x['shardId'].encode('utf-8')), 'beam:coder:custom_window:v1': lambda x, window_parser: window_parser( x['window']), 'beam:coder:nullable:v1': lambda x, value_parser: x.encode('utf-8') if x else None } def test_standard_coders(self): for name, spec in _load_test_cases(STANDARD_CODERS_YAML): logging.info('Executing %s test.', name) self._run_standard_coder(name, spec) def _run_standard_coder(self, name, spec): def assert_equal(actual, expected): """Handle nan values which self.assertEqual fails on.""" if (isinstance(actual, float) and isinstance(expected, float) and math.isnan(actual) and math.isnan(expected)): return self.assertEqual(actual, expected) coder = self.parse_coder(spec['coder']) parse_value = self.json_value_parser(spec['coder']) nested_list = [spec['nested']] if 'nested' in spec else [True, False] for nested in nested_list: for expected_encoded, json_value in spec['examples'].items(): value = parse_value(json_value) expected_encoded = expected_encoded.encode('latin1') if not spec['coder'].get('non_deterministic', False): actual_encoded = encode_nested(coder, value, nested) if self.fix and actual_encoded != expected_encoded: self.to_fix[spec['index'], expected_encoded] = actual_encoded else: self.assertEqual(expected_encoded, actual_encoded) decoded = decode_nested(coder, expected_encoded, nested) assert_equal(decoded, value) else: # Only verify decoding for a non-deterministic coder self.assertEqual( decode_nested(coder, expected_encoded, nested), value) if spec['coder']['urn'] == 'beam:coder:row:v1': # Test batch encoding/decoding as well. values = [ parse_value(json_value) for json_value in spec['examples'].values() ] columnar = { field.name: np.array([getattr(value, field.name) for value in values]) for field in coder.schema.fields } dest = { field: np.empty_like(values) for field, values in columnar.items() } for column in dest.values(): column[:] = 0 if 'int' in column.dtype.name else None expected_encoded = ''.join(spec['examples'].keys()).encode('latin1') actual_encoded = encode_batch(coder, columnar) assert_equal(expected_encoded, actual_encoded) decoded_count = decode_batch(coder, expected_encoded, dest) assert_equal(len(spec['examples']), decoded_count) for field, values in dest.items(): assert_array_equal(columnar[field], dest[field]) def parse_coder(self, spec): context = pipeline_context.PipelineContext() coder_id = str(hash(str(spec))) component_ids = [ context.coders.get_id(self.parse_coder(c)) for c in spec.get('components', ()) ] if spec.get('state'): def iterable_state_read(state_token, elem_coder): state = spec.get('state').get(state_token.decode('latin1')) if state is None: state = '' input_stream = coder_impl.create_InputStream(state.encode('latin1')) while input_stream.size() > 0: yield elem_coder.decode_from_stream(input_stream, True) context.iterable_state_read = iterable_state_read context.coders.put_proto( coder_id, beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec( urn=spec['urn'], payload=spec.get('payload', '').encode('latin1')), component_coder_ids=component_ids)) return context.coders.get_by_id(coder_id) def json_value_parser(self, coder_spec): # TODO: integrate this with the logic for the other parsers if coder_spec['urn'] == 'beam:coder:row:v1': schema = schema_pb2.Schema.FromString( coder_spec['payload'].encode('latin1')) return value_parser_from_schema(schema) component_parsers = [ self.json_value_parser(c) for c in coder_spec.get('components', ()) ] return lambda x: self._urn_to_json_value_parser[coder_spec['urn']]( x, *component_parsers) # Used when --fix is passed. fix = False to_fix: dict[tuple[int, bytes], bytes] = {} @classmethod def tearDownClass(cls): if cls.fix and cls.to_fix: print("FIXING", len(cls.to_fix), "TESTS") doc_sep = '\n---\n' docs = open(STANDARD_CODERS_YAML).read().split(doc_sep) def quote(s): return json.dumps(s.decode('latin1')).replace(r'\u0000', r'\0') for (doc_ix, expected_encoded), actual_encoded in cls.to_fix.items(): print(quote(expected_encoded), "->", quote(actual_encoded)) docs[doc_ix] = docs[doc_ix].replace( quote(expected_encoded) + ':', quote(actual_encoded) + ':') open(STANDARD_CODERS_YAML, 'w').write(doc_sep.join(docs)) def encode_nested(coder, value, nested=True): out = coder_impl.create_OutputStream() coder.get_impl().encode_to_stream(value, out, nested) return out.get() def decode_nested(coder, encoded, nested=True): return coder.get_impl().decode_from_stream( coder_impl.create_InputStream(encoded), nested) def encode_batch(row_coder, values): out = coder_impl.create_OutputStream() row_coder.get_impl().encode_batch_to_stream(values, out) return out.get() def decode_batch(row_coder, encoded, dest): return row_coder.get_impl().decode_batch_from_stream( dest, coder_impl.create_InputStream(encoded)) if __name__ == '__main__': if '--fix' in sys.argv: StandardCodersTest.fix = True sys.argv.remove('--fix') unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/stream.pxd0000644000175100017510000000561415206113032021152 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # cimport libc.stdint cdef class OutputStream(object): cdef char* data cdef size_t buffer_size cdef size_t pos cpdef write(self, bytes b, bint nested=*) cpdef write_byte(self, unsigned char val) cpdef write_var_int64(self, libc.stdint.int64_t v) cpdef write_var_int32(self, libc.stdint.int64_t v) cpdef write_bigendian_int64(self, libc.stdint.int64_t signed_v) cpdef write_bigendian_uint64(self, libc.stdint.uint64_t signed_v) cpdef write_bigendian_int32(self, libc.stdint.int32_t signed_v) cpdef write_bigendian_int16(self, libc.stdint.int16_t signed_v) cpdef write_bigendian_double(self, double d) cpdef write_bigendian_float(self, float d) cpdef bytes get(self) cpdef size_t size(self) except? -1 cdef extend(self, size_t missing) cpdef _clear(self) cdef class ByteCountingOutputStream(OutputStream): cdef size_t count cpdef write(self, bytes b, bint nested=*) cpdef write_var_int64(self, libc.stdint.int64_t val) cpdef write_var_int32(self, libc.stdint.int64_t val) cpdef write_byte(self, unsigned char val) cpdef write_bigendian_int64(self, libc.stdint.int64_t val) cpdef write_bigendian_uint64(self, libc.stdint.uint64_t val) cpdef write_bigendian_int32(self, libc.stdint.int32_t val) cpdef write_bigendian_int16(self, libc.stdint.int16_t val) cpdef size_t get_count(self) cpdef bytes get(self) cdef class InputStream(object): cdef size_t pos cdef bytes all cdef char* allc cpdef ssize_t size(self) except? -1 cpdef bytes read(self, size_t len) cpdef long read_byte(self) except? -1 cpdef libc.stdint.int64_t read_var_int64(self) except? -1 cpdef libc.stdint.int32_t read_var_int32(self) except? -1 cpdef libc.stdint.int64_t read_bigendian_int64(self) except? -1 cpdef libc.stdint.uint64_t read_bigendian_uint64(self) except? -1 cpdef libc.stdint.int32_t read_bigendian_int32(self) except? -1 cpdef libc.stdint.int16_t read_bigendian_int16(self) except? -1 cpdef double read_bigendian_double(self) except? -1 cpdef float read_bigendian_float(self) except? -1 cpdef bytes read_all(self, bint nested=*) cpdef libc.stdint.int64_t get_varint_size(libc.stdint.int64_t value) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/stream.pyx0000644000175100017510000002315215206113032021174 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Compiled version of the Stream objects used by CoderImpl. For internal use only; no backwards-compatibility guarantees. """ cimport libc.stdlib cimport libc.string cdef class OutputStream(object): """An output string stream implementation supporting write() and get().""" #TODO(robertwb): Consider using raw C++ streams. def __cinit__(self): self.buffer_size = 1024 self.pos = 0 self.data = libc.stdlib.malloc(self.buffer_size) assert self.data, "OutputStream malloc failed." def __dealloc__(self): if self.data: libc.stdlib.free(self.data) cpdef write(self, bytes b, bint nested=False): cdef size_t blen = len(b) if nested: self.write_var_int64(blen) if self.buffer_size < self.pos + blen: self.extend(blen) libc.string.memcpy(self.data + self.pos, b, blen) self.pos += blen cpdef write_byte(self, unsigned char val): if self.buffer_size < self.pos + 1: self.extend(1) self.data[self.pos] = val self.pos += 1 cpdef write_var_int64(self, libc.stdint.int64_t signed_v): """Encode a long using variable-length encoding to a stream.""" cdef libc.stdint.uint64_t v = signed_v # Inline common case. if v <= 0x7F and self.pos < self.buffer_size - 1: self.data[self.pos] = v self.pos += 1 return cdef long bits while True: bits = v & 0x7F v >>= 7 if v: bits |= 0x80 self.write_byte(bits) if not v: break cpdef write_var_int32(self, libc.stdint.int64_t signed_v): """Encode an int using variable-length encoding to a stream.""" # for backward compatibility, input type is int64_t thus tolerates overflow cdef libc.stdint.int64_t v = signed_v & 0xFFFFFFFF self.write_var_int64(v) cpdef write_bigendian_int64(self, libc.stdint.int64_t signed_v): self.write_bigendian_uint64(signed_v) cpdef write_bigendian_uint64(self, libc.stdint.uint64_t v): if self.buffer_size < self.pos + 8: self.extend(8) self.data[self.pos ] = (v >> 56) self.data[self.pos + 1] = (v >> 48) self.data[self.pos + 2] = (v >> 40) self.data[self.pos + 3] = (v >> 32) self.data[self.pos + 4] = (v >> 24) self.data[self.pos + 5] = (v >> 16) self.data[self.pos + 6] = (v >> 8) self.data[self.pos + 7] = (v ) self.pos += 8 cpdef write_bigendian_int32(self, libc.stdint.int32_t signed_v): cdef libc.stdint.uint32_t v = signed_v if self.buffer_size < self.pos + 4: self.extend(4) self.data[self.pos ] = (v >> 24) self.data[self.pos + 1] = (v >> 16) self.data[self.pos + 2] = (v >> 8) self.data[self.pos + 3] = (v ) self.pos += 4 cpdef write_bigendian_int16(self, libc.stdint.int16_t signed_v): cdef libc.stdint.uint16_t v = signed_v if self.buffer_size < self.pos + 2: self.extend(2) self.data[self.pos ] = (v >> 8) self.data[self.pos + 1] = (v ) self.pos += 2 cpdef write_bigendian_double(self, double d): self.write_bigendian_int64((&d)[0]) cpdef write_bigendian_float(self, float f): self.write_bigendian_int32((&f)[0]) cpdef bytes get(self): return self.data[:self.pos] cpdef size_t size(self) except? -1: return self.pos cdef extend(self, size_t missing): while missing > self.buffer_size - self.pos: self.buffer_size *= 2 self.data = libc.stdlib.realloc(self.data, self.buffer_size) assert self.data, "OutputStream realloc failed." cpdef _clear(self): self.pos = 0 cdef class ByteCountingOutputStream(OutputStream): """An output string stream implementation that only counts the bytes. This implementation counts the number of bytes it "writes" but doesn't actually write them anyway. Thus it has write() but not get(). get_count() returns how many bytes were written. This is useful for sizing an encoding. """ def __cinit__(self): self.count = 0 cpdef write(self, bytes b, bint nested=False): cdef size_t blen = len(b) if nested: self.write_var_int64(blen) self.count += blen cpdef write_var_int64(self, libc.stdint.int64_t signed_v): self.count += get_varint_size(signed_v) cpdef write_var_int32(self, libc.stdint.int64_t signed_v): if signed_v < 0: self.count += 5 else: self.count += get_varint_size(signed_v) cpdef write_byte(self, unsigned char _): self.count += 1 cpdef write_bigendian_int64(self, libc.stdint.int64_t _): self.count += 8 cpdef write_bigendian_uint64(self, libc.stdint.uint64_t _): self.count += 8 cpdef write_bigendian_int32(self, libc.stdint.int32_t _): self.count += 4 cpdef write_bigendian_int16(self, libc.stdint.int16_t _): self.count += 2 cpdef size_t get_count(self): return self.count cpdef bytes get(self): raise NotImplementedError def __str__(self): return '<%s %s>' % (self.__class__.__name__, self.count) cdef class InputStream(object): """An input string stream implementation supporting read() and size().""" def __init__(self, all): self.allc = self.all = all cpdef bytes read(self, size_t size): self.pos += size return self.allc[self.pos - size : self.pos] cpdef long read_byte(self) except? -1: self.pos += 1 # Note: Some C++ compilers treats the char array below as a signed char. # This causes incorrect coder behavior unless explicitly cast to an # unsigned char here. return ( self.allc[self.pos - 1]) cpdef ssize_t size(self) except? -1: return len(self.all) - self.pos cpdef bytes read_all(self, bint nested=False): return self.read(self.read_var_int64() if nested else self.size()) cpdef libc.stdint.int64_t read_var_int64(self) except? -1: """Decode a variable-length encoded long from a stream.""" # Inline common case. cdef long byte = self.allc[self.pos] self.pos += 1 if byte <= 0x7F: return byte cdef libc.stdint.int64_t bits cdef long shift = 0 cdef libc.stdint.int64_t result = 0 while True: bits = byte & 0x7F if (shift >= sizeof(libc.stdint.int64_t) * 8 or (shift >= (sizeof(libc.stdint.int64_t) * 8 - 1) and bits > 1)): raise RuntimeError('VarLong too long.') result |= bits << shift shift += 7 if not (byte & 0x80): break byte = self.read_byte() if byte < 0: raise RuntimeError('VarInt not terminated.') return result cpdef libc.stdint.int32_t read_var_int32(self) except? -1: """Decode a variable-length encoded int32 from a stream.""" cdef libc.stdint.int64_t v = self.read_var_int64() return (v); cpdef libc.stdint.int64_t read_bigendian_int64(self) except? -1: return self.read_bigendian_uint64() cpdef libc.stdint.uint64_t read_bigendian_uint64(self) except? -1: self.pos += 8 return (self.allc[self.pos - 1] | self.allc[self.pos - 2] << 8 | self.allc[self.pos - 3] << 16 | self.allc[self.pos - 4] << 24 | self.allc[self.pos - 5] << 32 | self.allc[self.pos - 6] << 40 | self.allc[self.pos - 7] << 48 | self.allc[self.pos - 8] << 56) cpdef libc.stdint.int32_t read_bigendian_int32(self) except? -1: self.pos += 4 return (self.allc[self.pos - 1] | self.allc[self.pos - 2] << 8 | self.allc[self.pos - 3] << 16 | self.allc[self.pos - 4] << 24) cpdef libc.stdint.int16_t read_bigendian_int16(self) except? -1: self.pos += 2 return (self.allc[self.pos - 1] | self.allc[self.pos - 2] << 8) cpdef double read_bigendian_double(self) except? -1: cdef libc.stdint.int64_t as_long = self.read_bigendian_int64() return (&as_long)[0] cpdef float read_bigendian_float(self) except? -1: cdef libc.stdint.int32_t as_int = self.read_bigendian_int32() return (&as_int)[0] cpdef libc.stdint.int64_t get_varint_size(libc.stdint.int64_t value): """Returns the size of the given integer value when encode as a VarInt.""" cdef libc.stdint.int64_t varint_size = 0 cdef libc.stdint.uint64_t bits = value while True: varint_size += 1 bits >>= 7 if not bits: break return varint_size ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/stream_test.py0000644000175100017510000001632615206113032022050 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for the stream implementations.""" # pytype: skip-file import logging import math import unittest import numpy as np from apache_beam.coders import slow_stream class StreamTest(unittest.TestCase): # pylint: disable=invalid-name InputStream = slow_stream.InputStream OutputStream = slow_stream.OutputStream ByteCountingOutputStream = slow_stream.ByteCountingOutputStream # pylint: enable=invalid-name def test_read_write(self): out_s = self.OutputStream() out_s.write(b'abc') out_s.write(b'\0\t\n') out_s.write(b'xyz', True) out_s.write(b'', True) in_s = self.InputStream(out_s.get()) self.assertEqual(b'abc\0\t\n', in_s.read(6)) self.assertEqual(b'xyz', in_s.read_all(True)) self.assertEqual(b'', in_s.read_all(True)) def test_read_all(self): out_s = self.OutputStream() out_s.write(b'abc') in_s = self.InputStream(out_s.get()) self.assertEqual(b'abc', in_s.read_all(False)) def test_read_write_byte(self): out_s = self.OutputStream() out_s.write_byte(1) out_s.write_byte(0) out_s.write_byte(0xFF) in_s = self.InputStream(out_s.get()) self.assertEqual(1, in_s.read_byte()) self.assertEqual(0, in_s.read_byte()) self.assertEqual(0xFF, in_s.read_byte()) def test_read_write_large(self): values = range(4 * 1024) out_s = self.OutputStream() for v in values: out_s.write_bigendian_int64(v) in_s = self.InputStream(out_s.get()) for v in values: self.assertEqual(v, in_s.read_bigendian_int64()) def run_read_write_var_int64(self, values): out_s = self.OutputStream() for v in values: out_s.write_var_int64(v) in_s = self.InputStream(out_s.get()) for v in values: self.assertEqual(v, in_s.read_var_int64()) def test_small_var_int64(self): self.run_read_write_var_int64(range(-10, 30)) def test_medium_var_int64(self): base = -1.7 self.run_read_write_var_int64([ int(base**pow) for pow in range(1, int(63 * math.log(2) / math.log(-base))) ]) def test_large_var_int64(self): self.run_read_write_var_int64([0, 2**63 - 1, -2**63, 2**63 - 3]) def run_read_write_var_int32(self, values): out_s = self.OutputStream() for v in values: out_s.write_var_int32(v) in_s = self.InputStream(out_s.get()) for v in values: self.assertEqual(v, in_s.read_var_int32()) def test_small_var_int32(self): self.run_read_write_var_int32(range(-10, 30)) def test_medium_var_int32(self): base = -1.7 self.run_read_write_var_int32([ int(base**pow) for pow in range(1, int(31 * math.log(2) / math.log(-base))) ]) def test_large_var_int32(self): self.run_read_write_var_int32([0, 2**31 - 1, -2**31, 2**31 - 3]) def test_read_write_double(self): values = 0, 1, -1, 1e100, 1.0 / 3, math.pi, float('inf') out_s = self.OutputStream() for v in values: out_s.write_bigendian_double(v) in_s = self.InputStream(out_s.get()) for v in values: self.assertEqual(v, in_s.read_bigendian_double()) def test_read_write_float(self): values = 0, 1, -1, 1e20, 1.0 / 3, math.pi, float('inf') # Restrict to single precision before coder roundtrip values = tuple(float(np.float32(v)) for v in values) out_s = self.OutputStream() for v in values: out_s.write_bigendian_float(v) in_s = self.InputStream(out_s.get()) for v in values: self.assertEqual(v, in_s.read_bigendian_float()) def test_read_write_bigendian_int64(self): values = 0, 1, -1, 2**63 - 1, -2**63, int(2**61 * math.pi) out_s = self.OutputStream() for v in values: out_s.write_bigendian_int64(v) in_s = self.InputStream(out_s.get()) for v in values: self.assertEqual(v, in_s.read_bigendian_int64()) def test_read_write_bigendian_uint64(self): values = 0, 1, 2**64 - 1, int(2**61 * math.pi) out_s = self.OutputStream() for v in values: out_s.write_bigendian_uint64(v) in_s = self.InputStream(out_s.get()) for v in values: self.assertEqual(v, in_s.read_bigendian_uint64()) def test_read_write_bigendian_int32(self): values = 0, 1, -1, 2**31 - 1, -2**31, int(2**29 * math.pi) out_s = self.OutputStream() for v in values: out_s.write_bigendian_int32(v) in_s = self.InputStream(out_s.get()) for v in values: self.assertEqual(v, in_s.read_bigendian_int32()) def test_read_write_bigendian_int16(self): values = 0, 1, -1, 2**15 - 1, -2**15, int(2**13 * math.pi) out_s = self.OutputStream() for v in values: out_s.write_bigendian_int16(v) in_s = self.InputStream(out_s.get()) for v in values: self.assertEqual(v, in_s.read_bigendian_int16()) def test_byte_counting(self): bc_s = self.ByteCountingOutputStream() self.assertEqual(0, bc_s.get_count()) bc_s.write(b'def') self.assertEqual(3, bc_s.get_count()) bc_s.write(b'') self.assertEqual(3, bc_s.get_count()) bc_s.write_byte(10) self.assertEqual(4, bc_s.get_count()) # "nested" also writes the length of the string, which should # cause 1 extra byte to be counted. bc_s.write(b'2345', nested=True) self.assertEqual(9, bc_s.get_count()) bc_s.write_var_int64(63) self.assertEqual(10, bc_s.get_count()) bc_s.write_bigendian_int64(42) self.assertEqual(18, bc_s.get_count()) bc_s.write_bigendian_int32(36) self.assertEqual(22, bc_s.get_count()) bc_s.write_bigendian_double(6.25) self.assertEqual(30, bc_s.get_count()) bc_s.write_bigendian_uint64(47) self.assertEqual(38, bc_s.get_count()) try: # pylint: disable=wrong-import-position from apache_beam.coders import stream class FastStreamTest(StreamTest): """Runs the test with the compiled stream classes.""" InputStream = stream.InputStream OutputStream = stream.OutputStream ByteCountingOutputStream = stream.ByteCountingOutputStream class SlowFastStreamTest(StreamTest): """Runs the test with compiled and uncompiled stream classes.""" InputStream = stream.InputStream OutputStream = slow_stream.OutputStream ByteCountingOutputStream = slow_stream.ByteCountingOutputStream class FastSlowStreamTest(StreamTest): """Runs the test with uncompiled and compiled stream classes.""" InputStream = slow_stream.InputStream OutputStream = stream.OutputStream ByteCountingOutputStream = stream.ByteCountingOutputStream except ImportError: pass if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/typecoders.py0000644000175100017510000002433715206113032021700 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Type coders registration. This module contains functionality to define and use coders for custom classes. Let's say we have a class Xyz and we are processing a PCollection with elements of type Xyz. If we do not register a coder for Xyz, a default pickle-based fallback coder will be used. This can be undesirable for two reasons. First, we may want a faster coder or a more space efficient one. Second, the pickle-based coder is not deterministic in the sense that objects like dictionaries or sets are not guaranteed to be encoded in the same way every time (elements are not really ordered). Two (sometimes three) steps are needed to define and use a custom coder: - define the coder class - associate the code with the class (a.k.a. coder registration) - typehint DoFns or transforms with the new class or composite types using the class. A coder class is defined by subclassing from CoderBase and defining the encode_to_bytes and decode_from_bytes methods. The framework uses duck-typing for coders so it is not strictly required to subclass from CoderBase as long as the encode/decode methods are defined. Registering a coder class is made with a register_coder() call:: from apache_beam import coders ... coders.registry.register_coder(Xyz, XyzCoder) Additionally, DoFns and PTransforms may need type hints. This is not always necessary since there is functionality to infer the return types of DoFns by analyzing the code. For instance, for the function below the return type of 'Xyz' will be inferred:: def MakeXyzs(v): return Xyz(v) If Xyz is inferred then its coder will be used whenever the framework needs to serialize data (e.g., writing to the shuffler subsystem responsible for group by key operations). If a typehint is needed it can be specified by decorating the DoFns or using with_input_types/with_output_types methods on PTransforms. For example, the above function can be decorated:: @with_output_types(Xyz) def MakeXyzs(v): return complex_operation_returning_Xyz(v) See apache_beam.typehints.decorators module for more details. """ # pytype: skip-file from typing import Any from typing import Iterable from apache_beam.coders import coders from apache_beam.typehints import typehints from apache_beam.utils import windowed_value __all__ = ['registry'] class CoderRegistry(object): """A coder registry for typehint/coder associations.""" def __init__(self, fallback_coder=None): self._coders: dict[Any, type[coders.Coder]] = {} self.custom_types: list[Any] = [] self.register_standard_coders(fallback_coder) def register_standard_coders(self, fallback_coder): """Register coders for all basic and composite types.""" # Coders without subclasses. self._register_coder_internal(int, coders.VarIntCoder) self._register_coder_internal(float, coders.FloatCoder) self._register_coder_internal(bytes, coders.BytesCoder) self._register_coder_internal(bool, coders.BooleanCoder) self._register_coder_internal(str, coders.StrUtf8Coder) self._register_coder_internal(windowed_value.PaneInfo, coders.PaneInfoCoder) self._register_coder_internal(typehints.TupleConstraint, coders.TupleCoder) self._register_coder_internal(typehints.DictConstraint, coders.MapCoder) self._register_coder_internal( typehints.WindowedTypeConstraint, coders.WindowedValueCoder) # Default fallback coders applied in that order until the first matching # coder found. default_fallback_coders = [ coders.ProtoCoder, coders.ProtoPlusCoder, coders.FastPrimitivesCoder ] self._fallback_coder = fallback_coder or FirstOf(default_fallback_coders) def register_fallback_coder(self, fallback_coder): self._fallback_coder = FirstOf([fallback_coder, self._fallback_coder]) def _register_coder_internal( self, typehint_type: Any, typehint_coder_class: type[coders.Coder]) -> None: self._coders[typehint_type] = typehint_coder_class @staticmethod def _normalize_typehint_type(typehint_type): if typehint_type.__module__ == '__main__': # See https://github.com/apache/beam/issues/21541 # TODO(robertwb): Remove once all runners are portable. return getattr(typehint_type, '__name__', str(typehint_type)) return typehint_type def register_coder( self, typehint_type: Any, typehint_coder_class: type[coders.Coder]) -> None: """ Register a user type with a coder. Typical usage:: class MyCustomType: pass coders.registry.register_coder(MyCustomType, MyCustomCoder) To register a supported user type (data class or named tuple) with Beam Row, use :meth:`register_row` instead, as it registers both coder and schema. """ if not isinstance(typehint_coder_class, type): raise TypeError( 'Coder registration requires a coder class object. ' 'Received %r instead.' % typehint_coder_class) if typehint_type not in self.custom_types: self.custom_types.append(typehint_type) self._register_coder_internal( self._normalize_typehint_type(typehint_type), typehint_coder_class) def register_row(self, typehint_type: type[Any]) -> type[Any]: """ Register a user type with a Beam Row. This registers the type with a RowCoder and register its schema. Register a dataclass:: @coders.registry.register_row @dataclass class MyDataClass: id: int name: str Register a named tuple:: coders.registry.register_row(MyNamedTuple) """ from apache_beam.coders import RowCoder from apache_beam.typehints.schemas import typing_to_runner_api # Register with row coder self.register_coder(typehint_type, RowCoder) # This call generated a schema id for the type and register it with # schema registry typing_to_runner_api(typehint_type) return typehint_type def get_coder(self, typehint: Any) -> coders.Coder: if typehint and typehint.__module__ == '__main__': # See https://github.com/apache/beam/issues/21541 # TODO(robertwb): Remove once all runners are portable. typehint = getattr(typehint, '__name__', str(typehint)) coder = self._coders.get( typehint.__class__ if isinstance(typehint, typehints.TypeConstraint) else typehint, None) if isinstance(typehint, typehints.TypeConstraint) and coder is not None: return coder.from_type_hint(typehint, self) if coder is None: # We use the fallback coder when there is no coder registered for a # typehint. For example a user defined class with no coder specified. if not hasattr(self, '_fallback_coder'): raise RuntimeError( 'Coder registry has no fallback coder. This can happen if the ' 'fast_coders module could not be imported.') if isinstance(typehint, typehints.IterableTypeConstraint): return coders.IterableCoder.from_type_hint(typehint, self) elif isinstance(typehint, typehints.ListConstraint): return coders.ListCoder.from_type_hint(typehint, self) elif typehints.is_nullable(typehint): return coders.NullableCoder.from_type_hint(typehint, self) elif typehint is None: # In some old code, None is used for Any. # TODO(robertwb): Clean this up. pass elif typehint is object or typehint == typehints.Any: # We explicitly want the fallback coder. pass elif isinstance(typehint, typehints.TypeVariable): # TODO(robertwb): Clean this up when type inference is fully enabled. pass else: # TODO(robertwb): Re-enable this warning when it's actionable. # warnings.warn('Using fallback coder for typehint: %r.' % typehint) pass coder = self._fallback_coder return coder.from_type_hint(typehint, self) def get_custom_type_coder_tuples(self, types=None): """Returns type/coder tuples for all custom types passed in.""" return [(t, self._coders[self._normalize_typehint_type(t)]) for t in self.custom_types if (types is None or t in types)] def load_custom_type_coder_tuples(self, type_coder): """Load type/coder tuples into coder registry.""" for t, c in type_coder: self.register_coder(t, c) def verify_deterministic(self, key_coder, op_name, silent=True): if not key_coder.is_deterministic(): error_msg = ( 'The key coder "%s" for %s ' 'is not deterministic. This may result in incorrect ' 'pipeline output. This can be fixed by adding a type ' 'hint to the operation preceding the GroupByKey step, ' 'and for custom key classes, by writing a ' 'deterministic custom Coder. Please see the ' 'documentation for more details.' % (key_coder, op_name)) return key_coder.as_deterministic_coder(op_name, error_msg) else: return key_coder class FirstOf(object): """For internal use only; no backwards-compatibility guarantees. A class used to get the first matching coder from a list of coders.""" def __init__(self, coders: Iterable[type[coders.Coder]]) -> None: self._coders = coders def from_type_hint(self, typehint, registry): messages = [] for coder in self._coders: try: return coder.from_type_hint(typehint, registry) except Exception as e: msg = ( '%s could not provide a Coder for type %s: %s' % (coder, typehint, e)) messages.append(msg) raise ValueError( 'Cannot provide coder for %s: %s' % (typehint, ';'.join(messages))) registry = CoderRegistry() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/coders/typecoders_test.py0000644000175100017510000001460215206113032022731 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the typecoders module.""" # pytype: skip-file import unittest from apache_beam.coders import coders from apache_beam.coders import typecoders from apache_beam.internal import pickler from apache_beam.typehints import typehints from apache_beam.utils import windowed_value class CustomClass(object): def __init__(self, n): self.number = n def __eq__(self, other): return self.number == other.number def __hash__(self): return self.number class CustomCoder(coders.Coder): def encode(self, value): return str(value.number).encode('ASCII') def decode(self, encoded): return CustomClass(int(encoded)) def is_deterministic(self): # This coder is deterministic. Though we don't use need this coder to be # deterministic for this test, we annotate this as such to follow best # practices. return True class TypeCodersTest(unittest.TestCase): def test_register_non_type_coder(self): coder = CustomCoder() with self.assertRaisesRegex( TypeError, ('Coder registration requires a coder class object. ' 'Received %r instead.' % coder)): # When registering a coder the coder class must be specified. typecoders.registry.register_coder(CustomClass, coder) def test_get_coder_with_custom_coder(self): typecoders.registry.register_coder(CustomClass, CustomCoder) self.assertEqual( CustomCoder, typecoders.registry.get_coder(CustomClass).__class__) def test_get_coder_with_composite_custom_coder(self): typecoders.registry.register_coder(CustomClass, CustomCoder) coder = typecoders.registry.get_coder(typehints.KV[CustomClass, str]) revived_coder = pickler.loads(pickler.dumps(coder)) self.assertEqual( (CustomClass(123), 'abc'), revived_coder.decode(revived_coder.encode((CustomClass(123), 'abc')))) def test_get_coder_with_standard_coder(self): self.assertEqual( coders.BytesCoder, typecoders.registry.get_coder(bytes).__class__) def test_fallbackcoder(self): coder = typecoders.registry.get_coder(typehints.Any) self.assertEqual(('abc', 123), coder.decode(coder.encode(('abc', 123)))) def test_get_coder_can_be_pickled(self): coder = typecoders.registry.get_coder(typehints.Tuple[str, int]) revived_coder = pickler.loads(pickler.dumps(coder)) self.assertEqual(('abc', 123), revived_coder.decode(revived_coder.encode(('abc', 123)))) def test_standard_int_coder(self): real_coder = typecoders.registry.get_coder(int) expected_coder = coders.VarIntCoder() self.assertEqual(real_coder.encode(0x0404), expected_coder.encode(0x0404)) self.assertEqual(0x0404, real_coder.decode(real_coder.encode(0x0404))) self.assertEqual( real_coder.encode(0x040404040404), expected_coder.encode(0x040404040404)) self.assertEqual( 0x040404040404, real_coder.decode(real_coder.encode(0x040404040404))) def test_standard_str_coder(self): real_coder = typecoders.registry.get_coder(bytes) expected_coder = coders.BytesCoder() self.assertEqual(real_coder.encode(b'abc'), expected_coder.encode(b'abc')) self.assertEqual(b'abc', real_coder.decode(real_coder.encode(b'abc'))) def test_standard_bool_coder(self): real_coder = typecoders.registry.get_coder(bool) expected_coder = coders.BooleanCoder() self.assertEqual(real_coder.encode(True), expected_coder.encode(True)) self.assertEqual(True, real_coder.decode(real_coder.encode(True))) self.assertEqual(real_coder.encode(False), expected_coder.encode(False)) self.assertEqual(False, real_coder.decode(real_coder.encode(False))) def test_iterable_coder(self): real_coder = typecoders.registry.get_coder(typehints.Iterable[bytes]) expected_coder = coders.IterableCoder(coders.BytesCoder()) values = [b'abc', b'xyz'] self.assertEqual(expected_coder, real_coder) self.assertEqual(real_coder.encode(values), expected_coder.encode(values)) @unittest.skip('https://github.com/apache/beam/issues/21658') def test_list_coder(self): real_coder = typecoders.registry.get_coder(typehints.List[bytes]) expected_coder = coders.IterableCoder(coders.BytesCoder()) values = [b'abc', b'xyz'] self.assertEqual(expected_coder, real_coder) self.assertEqual(real_coder.encode(values), expected_coder.encode(values)) # IterableCoder.decode() always returns a list. Its implementation, # IterableCoderImpl, *can* return a non-list if it is provided a read_state # object, but this is not possible using the atomic IterableCoder interface. self.assertIs( list, type(expected_coder.decode(expected_coder.encode(values)))) def test_nullable_coder(self): expected_coder = coders.NullableCoder(coders.BytesCoder()) real_coder = typecoders.registry.get_coder(typehints.Optional[bytes]) self.assertEqual(expected_coder, real_coder) self.assertEqual(expected_coder.encode(None), real_coder.encode(None)) self.assertEqual(expected_coder.encode(b'abc'), real_coder.encode(b'abc')) def test_paneinfo_coder(self): expected_coder = coders.PaneInfoCoder() real_coder = typecoders.registry.get_coder(windowed_value.PaneInfo) self.assertEqual(expected_coder, real_coder) for i in range(10): pane_info = windowed_value.PaneInfo( is_first=i == 0, is_last=i == 9, timing=windowed_value.PaneInfoTiming.EARLY, # 0 index=i, nonspeculative_index=-1) encoded = real_coder.encode(pane_info) self.assertEqual(expected_coder.encode(pane_info), encoded) self.assertEqual(pane_info, real_coder.decode(encoded)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996270.8853557 apache_beam-2.74.0/apache_beam/dataframe/0000755000175100017510000000000015206113157017611 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/__init__.py0000644000175100017510000000267615206113032021725 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Beam DataFrame API - For high-level documentation see https://beam.apache.org/documentation/dsls/dataframes/overview/ - :mod:`apache_beam.dataframe.io`: DataFrame I/Os - :mod:`apache_beam.dataframe.frames`: DataFrame operations - :mod:`apache_beam.dataframe.convert`: Conversion between :class:`~apache_beam.pvalue.PCollection` and :class:`~apache_beam.dataframe.frames.DeferredDataFrame`. - :mod:`apache_beam.dataframe.schemas`: Mapping from Beam Schemas to pandas dtypes. - :mod:`apache_beam.dataframe.transforms`: Use DataFrame operations within a Beam pipeline with `DataframeTransform`. """ from apache_beam.dataframe.expressions import allow_non_parallel_operations ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/convert.py0000644000175100017510000002622015206113032021635 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import inspect import warnings import weakref from collections.abc import Iterable from typing import Any from typing import Optional from typing import Union import pandas as pd import apache_beam as beam from apache_beam import pvalue from apache_beam.dataframe import expressions from apache_beam.dataframe import frame_base from apache_beam.dataframe import transforms from apache_beam.dataframe.schemas import element_typehint_from_dataframe_proxy from apache_beam.dataframe.schemas import generate_proxy from apache_beam.typehints.pandas_type_compatibility import dtype_to_fieldtype # TODO: Or should this be called as_dataframe? def to_dataframe( pcoll: pvalue.PCollection, proxy: Optional[pd.core.generic.NDFrame] = None, label: Optional[str] = None, ) -> frame_base.DeferredFrame: """Converts a PCollection to a deferred dataframe-like object, which can manipulated with pandas methods like `filter` and `groupby`. For example, one might write:: pcoll = ... df = to_dataframe(pcoll, proxy=...) result = df.groupby('col').sum() pcoll_result = to_pcollection(result) A proxy object must be given if the schema for the PCollection is not known. """ if proxy is None: if pcoll.element_type is None: raise ValueError( "Cannot infer a proxy because the input PCollection does not have a " "schema defined. Please make sure a schema type is specified for " "the input PCollection, or provide a proxy.") # If no proxy is given, assume this is an element-wise schema-aware # PCollection that needs to be batched. if label is None: # Attempt to come up with a reasonable, stable label by retrieving # the name of these variables in the calling context. label = 'BatchElements(%s)' % _var_name(pcoll, 2) proxy = generate_proxy(pcoll.element_type) shim_dofn: beam.DoFn if isinstance(proxy, pd.DataFrame): shim_dofn = RowsToDataFrameFn() elif isinstance(proxy, pd.Series): shim_dofn = ElementsToSeriesFn() else: raise AssertionError("Unknown proxy type: %s" % proxy) pcoll = pcoll | label >> beam.ParDo(shim_dofn) return frame_base.DeferredFrame.wrap( expressions.PlaceholderExpression(proxy, pcoll)) # PCollections generated by to_pcollection are memoized, keyed by expression id. # WeakValueDictionary is used so the caches are cleaned up with the parent # pipelines # Note that the pipeline (indirectly) holds references to the transforms which # keeps both the PCollections and expressions alive. This ensures the # expression's ids are never accidentally re-used. TO_PCOLLECTION_CACHE: 'weakref.WeakValueDictionary[str, pvalue.PCollection]' = ( weakref.WeakValueDictionary()) UNBATCHED_CACHE: 'weakref.WeakValueDictionary[str, pvalue.PCollection]' = ( weakref.WeakValueDictionary()) class RowsToDataFrameFn(beam.DoFn): @beam.DoFn.yields_elements def process_batch(self, batch: pd.DataFrame) -> Iterable[pd.DataFrame]: yield batch class ElementsToSeriesFn(beam.DoFn): @beam.DoFn.yields_elements def process_batch(self, batch: pd.Series) -> Iterable[pd.Series]: yield batch def _make_unbatched_pcoll( pc: pvalue.PCollection, expr: expressions.Expression, include_indexes: bool): label = f"Unbatch '{expr._id}'" if include_indexes: label += " with indexes" if label not in UNBATCHED_CACHE: proxy = expr.proxy() shim_dofn: beam.DoFn if isinstance(proxy, pd.DataFrame): shim_dofn = DataFrameToRowsFn(proxy, include_indexes) elif isinstance(proxy, pd.Series): if include_indexes: warnings.warn( "Pipeline is converting a DeferredSeries to PCollection " "with include_indexes=True. Note that this parameter is " "_not_ respected for DeferredSeries conversion. To " "include the index with your data, produce a" "DeferredDataFrame instead.") shim_dofn = SeriesToElementsFn(proxy) else: raise TypeError(f"Proxy '{proxy}' has unsupported type '{type(proxy)}'") UNBATCHED_CACHE[label] = pc | label >> beam.ParDo(shim_dofn) # Note unbatched cache is keyed by the expression id as well as parameters # for the unbatching (i.e. include_indexes) return UNBATCHED_CACHE[label] class DataFrameToRowsFn(beam.DoFn): def __init__(self, proxy, include_indexes): self._proxy = proxy self._include_indexes = include_indexes @beam.DoFn.yields_batches def process(self, element: pd.DataFrame) -> Iterable[pd.DataFrame]: yield element def infer_output_type(self, input_element_type): return element_typehint_from_dataframe_proxy( self._proxy, self._include_indexes) class SeriesToElementsFn(beam.DoFn): def __init__(self, proxy): self._proxy = proxy @beam.DoFn.yields_batches def process(self, element: pd.Series) -> Iterable[pd.Series]: yield element def infer_output_type(self, input_element_type): return dtype_to_fieldtype(self._proxy.dtype) # TODO: Or should this be called from_dataframe? def to_pcollection( *dataframes: Union[frame_base.DeferredFrame, pd.DataFrame, pd.Series], label=None, always_return_tuple=False, yield_elements='schemas', include_indexes=False, pipeline=None) -> Union[pvalue.PCollection, tuple[pvalue.PCollection, ...]]: """Converts one or more deferred dataframe-like objects back to a PCollection. This method creates and applies the actual Beam operations that compute the given deferred dataframes, returning a PCollection of their results. By default the resulting PCollections are schema-aware PCollections where each element is one row from the output dataframes, excluding indexes. This behavior can be modified with the `yield_elements` and `include_indexes` arguments. Also accepts non-deferred pandas dataframes, which are converted to deferred, schema'd PCollections. In this case the contents of the entire dataframe are serialized into the graph, so for large amounts of data it is preferable to write them to disk and read them with one of the read methods. If more than one (related) result is desired, it can be more efficient to pass them all at the same time to this method. Args: label: (optional, default "ToPCollection(...)"") the label to use for the conversion transform. always_return_tuple: (optional, default: False) If true, always return a tuple of PCollections, even if there's only one output. yield_elements: (optional, default: "schemas") If set to "pandas", return PCollections containing the raw Pandas objects (DataFrames or Series), if set to "schemas", return an element-wise PCollection, where DataFrame and Series instances are expanded to one element per row. DataFrames are converted to schema-aware PCollections, where column values can be accessed by attribute. include_indexes: (optional, default: False) When yield_elements="schemas", if include_indexes=True, attempt to include index columns in the output schema for expanded DataFrames. Raises an error if any of the index levels are unnamed (name=None), or if any of the names are not unique among all column and index names. pipeline: (optional, unless non-deferred dataframes are passed) Used when creating a PCollection from a non-deferred dataframe. """ if not yield_elements in ("pandas", "schemas"): raise ValueError( "Invalid value for yield_elements argument, '%s'. " "Allowed values are 'pandas' and 'schemas'" % yield_elements) if label is None: # Attempt to come up with a reasonable, stable label by retrieving the name # of these variables in the calling context. label = 'ToPCollection(%s)' % ', '.join(_var_name(e, 3) for e in dataframes) # Support for non-deferred dataframes. deferred_dataframes = [] for ix, df in enumerate(dataframes): if isinstance(df, frame_base.DeferredBase): # TODO(robertwb): Maybe extract pipeline object? deferred_dataframes.append(df) elif isinstance(df, (pd.Series, pd.DataFrame)): if pipeline is None: raise ValueError( 'Pipeline keyword required for non-deferred dataframe conversion.') deferred = pipeline | '%s_Defer%s' % (label, ix) >> beam.Create([df]) deferred_dataframes.append( frame_base.DeferredFrame.wrap( expressions.PlaceholderExpression(df.iloc[:0], deferred))) else: raise TypeError( 'Unable to convert objects of type %s to a PCollection' % type(df)) dataframes = tuple(deferred_dataframes) def extract_input(placeholder): if not isinstance(placeholder._reference, pvalue.PCollection): raise TypeError( 'Expression roots must have been created with to_dataframe.') return placeholder._reference placeholders = frozenset.union( frozenset(), *[df._expr.placeholders() for df in dataframes]) # Exclude any dataframes that have already been converted to PCollections. # We only want to convert each DF expression once, then re-use. new_dataframes = [ df for df in dataframes if df._expr._id not in TO_PCOLLECTION_CACHE ] if len(new_dataframes): new_results: dict[Any, pvalue.PCollection] = { p: extract_input(p) for p in placeholders } | label >> transforms._DataframeExpressionsTransform( {ix: df._expr for (ix, df) in enumerate(new_dataframes)}) TO_PCOLLECTION_CACHE.update({ new_dataframes[ix]._expr._id: pc for ix, pc in new_results.items() }) raw_results = { ix: TO_PCOLLECTION_CACHE[df._expr._id] for ix, df in enumerate(dataframes) } if yield_elements == "schemas": def maybe_unbatch(pc, value): if isinstance(value, frame_base._DeferredScalar): return pc else: return _make_unbatched_pcoll(pc, value._expr, include_indexes) results = { ix: maybe_unbatch(pc, dataframes[ix]) for (ix, pc) in raw_results.items() } else: results = raw_results if len(results) == 1 and not always_return_tuple: return results[0] else: return tuple(value for key, value in sorted(results.items())) def _var_name(obj, level): frame = inspect.currentframe() for _ in range(level): if frame is None: return '...' frame = frame.f_back for key, value in frame.f_locals.items(): if obj is value: return key for key, value in frame.f_globals.items(): if obj is value: return key return '...' ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/convert_test.py0000644000175100017510000001650015206113032022674 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import unittest import pandas as pd import apache_beam as beam from apache_beam.dataframe import convert from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to def equal_to_unordered_series(expected): def check(actual): actual = pd.concat(actual) if sorted(expected) != sorted(actual): raise AssertionError('Series not equal: \n%s\n%s\n' % (expected, actual)) return check class ConvertTest(unittest.TestCase): def test_convert_yield_pandas(self): with beam.Pipeline() as p: a = pd.Series([1, 2, 3]) b = pd.Series([100, 200, 300]) pc_a = p | 'A' >> beam.Create([a]) pc_b = p | 'B' >> beam.Create([b]) df_a = convert.to_dataframe(pc_a, proxy=a[:0]) df_b = convert.to_dataframe(pc_b, proxy=b[:0]) df_2a = 2 * df_a df_3a = 3 * df_a df_ab = df_a * df_b # Converting multiple results at a time can be more efficient. pc_2a, pc_ab = convert.to_pcollection(df_2a, df_ab, yield_elements='pandas') # But separate conversions can be done as well. pc_3a = convert.to_pcollection(df_3a, yield_elements='pandas') assert_that(pc_2a, equal_to_unordered_series(2 * a), label='Check2a') assert_that(pc_3a, equal_to_unordered_series(3 * a), label='Check3a') assert_that(pc_ab, equal_to_unordered_series(a * b), label='Checkab') def test_convert(self): with beam.Pipeline() as p: a = pd.Series([1, 2, 3]) b = pd.Series([100, 200, 300]) pc_a = p | 'A' >> beam.Create(a) pc_b = p | 'B' >> beam.Create(b) df_a = convert.to_dataframe(pc_a) df_b = convert.to_dataframe(pc_b) df_2a = 2 * df_a df_3a = 3 * df_a df_ab = df_a * df_b # Converting multiple results at a time can be more efficient. pc_2a, pc_ab = convert.to_pcollection(df_2a, df_ab) # But separate conversions can be done as well. pc_3a = convert.to_pcollection(df_3a) assert_that(pc_2a, equal_to(list(2 * a)), label='Check2a') assert_that(pc_3a, equal_to(list(3 * a)), label='Check3a') assert_that(pc_ab, equal_to(list(a * b)), label='Checkab') def test_convert_with_none(self): # Ensure the logical Any type allows (nullable) None, see BEAM-12587. df = pd.DataFrame({'A': ['str', 10, None], 'B': [None, 'str', 20]}) with beam.Pipeline() as p: res = convert.to_pcollection(df, pipeline=p) | beam.Map(tuple) assert_that(res, equal_to([(row.A, row.B) for _, row in df.iterrows()])) def test_convert_scalar(self): with beam.Pipeline() as p: pc = p | 'A' >> beam.Create([1, 2, 3]) s = convert.to_dataframe(pc) pc_sum = convert.to_pcollection(s.sum()) assert_that(pc_sum, equal_to([6])) def test_convert_non_deferred(self): with beam.Pipeline() as p: s1 = pd.Series([1, 2, 3]) s2 = convert.to_dataframe(p | beam.Create([100, 200, 300])) pc1, pc2 = convert.to_pcollection(s1, s2, pipeline=p) assert_that(pc1, equal_to([1, 2, 3]), label='CheckNonDeferred') assert_that(pc2, equal_to([100, 200, 300]), label='CheckDeferred') def test_convert_memoization(self): with beam.Pipeline() as p: a = pd.Series([1, 2, 3]) b = pd.Series([100, 200, 300]) pc_a = p | 'A' >> beam.Create([a]) pc_b = p | 'B' >> beam.Create([b]) df_a = convert.to_dataframe(pc_a, proxy=a[:0]) df_b = convert.to_dataframe(pc_b, proxy=b[:0]) df_2a = 2 * df_a df_3a = 3 * df_a df_ab = df_a * df_b # Two calls to to_pcollection with the same Dataframe should produce the # same PCollection(s) pc_2a_, pc_ab_ = convert.to_pcollection(df_2a, df_ab) pc_3a, pc_2a, pc_ab = convert.to_pcollection(df_3a, df_2a, df_ab) self.assertIs(pc_2a, pc_2a_) self.assertIs(pc_ab, pc_ab_) self.assertIsNot(pc_3a, pc_2a) self.assertIsNot(pc_3a, pc_ab) # The same conversions without the unbatching transform should also cache # PCollections pc_2a_pandas_, pc_ab_pandas_ = convert.to_pcollection(df_2a, df_ab, yield_elements='pandas') pc_3a_pandas, pc_2a_pandas, pc_ab_pandas = convert.to_pcollection(df_3a, df_2a, df_ab, yield_elements='pandas') self.assertIs(pc_2a_pandas, pc_2a_pandas_) self.assertIs(pc_ab_pandas, pc_ab_pandas_) self.assertIsNot(pc_3a_pandas, pc_2a_pandas) self.assertIsNot(pc_3a_pandas, pc_ab_pandas) # .. but the cached PCollections should be different self.assertIsNot(pc_2a_pandas, pc_2a) self.assertIsNot(pc_ab_pandas, pc_ab) self.assertIsNot(pc_3a_pandas, pc_3a) def test_convert_memoization_clears_cache(self): # This test re-runs the other memoization test, and makes sure that the # cache is cleaned up with the pipeline. Otherwise there would be concerns # of it growing without bound. import gc # Make sure cache is clear gc.collect() self.assertEqual(len(convert.TO_PCOLLECTION_CACHE), 0) # Disable GC so it doesn't run pre-emptively, confounding assertions about # cache size gc.disable() # Also disable logging, as some implementations may artificially extend # the life of objects. import logging logging.disable(logging.INFO) try: self.test_convert_memoization() self.assertEqual(len(convert.TO_PCOLLECTION_CACHE), 3) gc.collect() # PCollections should be removed from cache after pipelines go out of # scope and are GC'd self.assertEqual(len(convert.TO_PCOLLECTION_CACHE), 0) finally: # Always re-enable GC and logging gc.enable() logging.disable(logging.NOTSET) def test_auto_convert(self): class MySchemaTransform(beam.PTransform): def expand(self, pcoll): return pcoll | beam.Map( lambda x: beam.Row( a=x.n**2 - x.m**2, b=2 * x.m * x.n, c=x.n**2 + x.m**2)) with beam.Pipeline() as p: pc_mn = p | beam.Create([ (1, 2), (2, 3), (3, 10) ]) | beam.MapTuple(lambda m, n: beam.Row(m=m, n=n)) df_mn = convert.to_dataframe(pc_mn) # Apply a transform directly to a dataframe to get another dataframe. df_abc = df_mn | MySchemaTransform() pc_abc = convert.to_pcollection(df_abc) | beam.Map(tuple) assert_that(pc_abc, equal_to([(3, 4, 5), (5, 12, 13), (91, 60, 109)])) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/doctests.py0000644000175100017510000006102015206113032022002 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """A module that allows running existing pandas doctests with Beam dataframes. This module hooks into the doctesting framework by providing a custom runner and, in particular, an OutputChecker, as well as providing a fake object for mocking out the pandas module. The (novel) sequence of events when running a doctest is as follows. 1. The test invokes `pd.DataFrame(...)` (or similar) and an actual dataframe is computed and stashed but a Beam deferred dataframe is returned in its place. 2. Computations are done on these "dataframes," resulting in new objects, but as these are actually deferred, only expression trees are built. In the background, a mapping of id -> deferred dataframe is stored for each newly created dataframe. 3. When any dataframe is printed out, the repr has been overwritten to print `Dataframe[id]`. The aforementened mapping is used to map this back to the actual dataframe object, which is then computed via Beam, and its the (stringified) result plugged into the actual output for comparison. 4. The comparison is then done on the sorted lines of the expected and actual values. """ import collections import contextlib import doctest import re import sys import traceback from io import StringIO from typing import Any import numpy as np import pandas as pd import apache_beam as beam from apache_beam.dataframe import expressions from apache_beam.dataframe import frames # pylint: disable=unused-import from apache_beam.dataframe import pandas_top_level_functions from apache_beam.dataframe import transforms from apache_beam.dataframe.frame_base import DeferredBase class FakePandasObject(object): """A stand-in for the wrapped pandas objects. """ def __init__(self, pandas_obj, test_env): self._pandas_obj = pandas_obj self._test_env = test_env def __call__(self, *args, **kwargs): result = self._pandas_obj(*args, **kwargs) if type(result) in DeferredBase._pandas_type_map: placeholder = expressions.PlaceholderExpression(result.iloc[0:0]) self._test_env._inputs[placeholder] = result return DeferredBase.wrap(placeholder) else: return result def __getattr__(self, name): attr = getattr(self._pandas_obj, name) if callable(attr): result = FakePandasObject(attr, self._test_env) else: result = attr # Cache this so two lookups return the same object. setattr(self, name, result) return result def __reduce__(self): return lambda: pd, () class TestEnvironment(object): """A class managing the patching (of methods, inputs, and outputs) needed to run and validate tests. These classes are patched to be able to recognize and retrieve inputs and results, stored in `self._inputs` and `self._all_frames` respectively. """ def __init__(self): self._inputs = {} self._all_frames = {} def fake_pandas_module(self): return FakePandasObject(pandas_top_level_functions.pd_wrapper, self) @contextlib.contextmanager def _monkey_patch_type(self, deferred_type): """Monkey-patch __init__ to record a pointer to all created frames, and __repr__ to be able to recognize them in the doctest output. """ try: old_init, old_repr = deferred_type.__init__, deferred_type.__repr__ def new_init(df, *args, **kwargs): old_init(df, *args, **kwargs) self._all_frames[id(df)] = df deferred_type.__init__ = new_init deferred_type.__repr__ = lambda self: 'DeferredBase[%s]' % id(self) self._recorded_results = collections.defaultdict(list) yield finally: deferred_type.__init__, deferred_type.__repr__ = old_init, old_repr @contextlib.contextmanager def context(self): """Creates a context within which DeferredBase types are monkey patched to record ids.""" with contextlib.ExitStack() as stack: for deferred_type in DeferredBase._pandas_type_map.values(): stack.enter_context(self._monkey_patch_type(deferred_type)) yield class _InMemoryResultRecorder(object): """Helper for extracting computed results from a Beam pipeline. Used as follows:: with _InMemoryResultRecorder() as recorder: with beam.Pipeline() as p: ... pcoll | beam.Map(recorder.record_fn(name)) seen = recorder.get_recorded(name) """ # Class-level value to survive pickling. _ALL_RESULTS: dict[str, list[Any]] = {} def __init__(self): self._id = id(self) def __enter__(self): self._ALL_RESULTS[self._id] = collections.defaultdict(list) return self def __exit__(self, *unused_args): del self._ALL_RESULTS[self._id] def record_fn(self, name): def record(value): self._ALL_RESULTS[self._id][name].append(value) return record def get_recorded(self, name): return self._ALL_RESULTS[self._id][name] WONT_IMPLEMENT = 'apache_beam.dataframe.frame_base.WontImplementError' NOT_IMPLEMENTED = 'NotImplementedError' class _DeferrredDataframeOutputChecker(doctest.OutputChecker): """Validates output by replacing DeferredBase[...] with computed values. """ def __init__(self, env, use_beam): self._env = env if use_beam: self.compute = self.compute_using_beam else: self.compute = self.compute_using_session self.reset() def reset(self): self._last_error = None def compute_using_session(self, to_compute): session = expressions.PartitioningSession(self._env._inputs) return { name: session.evaluate(frame._expr) for name, frame in to_compute.items() } def compute_using_beam(self, to_compute): with _InMemoryResultRecorder() as recorder: with beam.Pipeline() as p: input_pcolls = { placeholder: p | 'Create%s' % placeholder >> beam.Create([input[::2], input[1::2]]) for placeholder, input in self._env._inputs.items() } output_pcolls = ( input_pcolls | transforms._DataframeExpressionsTransform({ name: frame._expr for name, frame in to_compute.items() })) for name, output_pcoll in output_pcolls.items(): _ = output_pcoll | 'Record%s' % name >> beam.FlatMap( recorder.record_fn(name)) # pipeline runs, side effects recorded def concat(values): if len(values) > 1: return pd.concat(values) else: return values[0] return { name: concat(recorder.get_recorded(name)) for name in to_compute.keys() } def fix(self, want, got): if 'DeferredBase' in got: # When we have a tuple of Dataframes, pandas prints each from a new line. got = re.sub(r'DeferredBase\[(\d+)\],', '\\g<0>\n', got) try: to_compute = { m.group(0): self._env._all_frames[int(m.group(1))] for m in re.finditer(r'DeferredBase\[(\d+)\]', got) } computed = self.compute(to_compute) for name, frame in computed.items(): got = got.replace(name, repr(frame)) # If a multiindex is used, compensate for it if any(isinstance(frame, pd.core.generic.NDFrame) and frame.index.nlevels > 1 for frame in computed.values()): def fill_multiindex(text): """An awful hack to work around the fact that pandas omits repeated elements in a multi-index. For example: Series name Row ID s1 0 a 1 b s2 0 c 1 d dtype: object The s1 and s2 are implied for the 2nd and 4th rows. However if we re-order this Series it might be printed this way: Series name Row ID s1 0 a s2 1 d s2 0 c s1 1 b dtype: object In our model these are equivalent, but when we sort the lines and check equality they are not. This method fills in any omitted multiindex values, so that we can successfully sort and compare.""" lines = [list(line) for line in text.split('\n')] for prev, line in zip(lines[:-1], lines[1:]): if all(l == ' ' for l in line): continue for i, l in enumerate(line): if l != ' ': break line[i] = prev[i] return '\n'.join(''.join(line) for line in lines) got = fill_multiindex(got) want = fill_multiindex(want) def sort_and_normalize(text): return '\n'.join( sorted( [line.rstrip() for line in text.split('\n') if line.strip()], key=str.strip)) + '\n' got = sort_and_normalize(got) want = sort_and_normalize(want) except Exception: got = traceback.format_exc() if 'NumpyExtensionArray' in got: # Work around formatting differences (np.int32(1) vs 1). got = re.sub('np.(int32|str_)[(]([^()]+)[)]', r'\2', got) got = re.sub('np.complex128[(]([^()]+)[)]', r'(\1)', got) return want, got @property def _seen_error(self): return self._last_error is not None def check_output(self, want, got, optionflags): # When an error occurs check_output is called with want=example.exc_msg, # and got=exc_msg # First check if `want` is a special string indicating wont_implement_ok # and/or not_implemented_ok allowed_exceptions = want.split('|') if all(exc in (WONT_IMPLEMENT, NOT_IMPLEMENTED) for exc in allowed_exceptions): # If it is, check for WontImplementError and NotImplementedError if WONT_IMPLEMENT in allowed_exceptions and got.startswith( WONT_IMPLEMENT): self._last_error = WONT_IMPLEMENT return True elif NOT_IMPLEMENTED in allowed_exceptions and got.startswith( NOT_IMPLEMENTED): self._last_error = NOT_IMPLEMENTED return True elif got.startswith('NameError') and self._seen_error: # This allows us to gracefully skip tests like # >>> res = df.unsupported_operation() # >>> check(res) return True self.reset() want, got = self.fix(want, got) return super().check_output(want, got, optionflags) def output_difference(self, example, got, optionflags): want, got = self.fix(example.want, got) if want != example.want: example = doctest.Example( example.source, want, example.exc_msg, example.lineno, example.indent, example.options) return super().output_difference(example, got, optionflags) class BeamDataframeDoctestRunner(doctest.DocTestRunner): """A Doctest runner suitable for replacing the `pd` module with one backed by beam. """ def __init__( self, env, use_beam=True, wont_implement_ok=None, not_implemented_ok=None, skip=None, **kwargs): self._test_env = env def to_callable(cond): if cond == '*': return lambda example: True else: return lambda example: example.source.strip() == cond self._wont_implement_ok = { test: [to_callable(cond) for cond in examples] for test, examples in (wont_implement_ok or {}).items() } self._not_implemented_ok = { test: [to_callable(cond) for cond in examples] for test, examples in (not_implemented_ok or {}).items() } self._skip = { test: [to_callable(cond) for cond in examples] for test, examples in (skip or {}).items() } super().__init__( checker=_DeferrredDataframeOutputChecker(self._test_env, use_beam), **kwargs) self.success = 0 self.skipped = 0 self._reasons = collections.defaultdict(list) self._skipped_set = set() def _is_wont_implement_ok(self, example, test): always_wont_implement = self._wont_implement_ok.get('*', []) return any( wont_implement(example) for wont_implement in ( self._wont_implement_ok.get(test.name, []) + always_wont_implement)) def _is_not_implemented_ok(self, example, test): always_not_impl = self._not_implemented_ok.get('*', []) return any( not_implemented(example) for not_implemented in ( self._not_implemented_ok.get(test.name, []) + always_not_impl)) def run(self, test, **kwargs): self._checker.reset() always_skip = self._skip.get('*', []) for example in test.examples: if any(should_skip(example) for should_skip in self._skip.get(test.name, []) + always_skip): self._skipped_set.add(example) example.source = 'pass' example.want = '' self.skipped += 1 elif example.exc_msg is None: allowed_exceptions = [] if self._is_not_implemented_ok(example, test): allowed_exceptions.append(NOT_IMPLEMENTED) if self._is_wont_implement_ok(example, test): allowed_exceptions.append(WONT_IMPLEMENT) if len(allowed_exceptions): # Don't fail doctests that raise this error. example.exc_msg = '|'.join(allowed_exceptions) with self._test_env.context(): result = super().run(test, **kwargs) # Can't add attributes to builtin result. result = AugmentedTestResults(result.failed, result.attempted) result.summary = self.summary() return result def report_success(self, out, test, example, got): def extract_concise_reason(got, expected_exc): m = re.search(r"Implement(?:ed)?Error:\s+(.*)\n$", got) if m: return m.group(1) elif "NameError" in got: return "NameError following %s" % expected_exc elif re.match(r"DeferredBase\[\d+\]\n", got): return "DeferredBase[*]" else: return got.replace("\n", "\\n") if self._checker._last_error is not None: self._reasons[self._checker._last_error].append( extract_concise_reason(got, self._checker._last_error)) if self._checker._seen_error: m = re.search('^([a-zA-Z0-9_, ]+)=', example.source) if m: for var in m.group(1).split(','): var = var.strip() if var in test.globs: # More informative to get a NameError than # use the wrong previous value. del test.globs[var] return super().report_success(out, test, example, got) def fake_pandas_module(self): return self._test_env.fake_pandas_module() def summarize(self): super().summarize() self.summary().summarize() def summary(self): return Summary(self.failures, self.tries, self.skipped, self._reasons) class AugmentedTestResults(doctest.TestResults): pass class Summary(object): def __init__(self, failures=0, tries=0, skipped=0, error_reasons=None): self.failures = failures self.tries = tries self.skipped = skipped self.error_reasons = error_reasons or collections.defaultdict(list) def result(self): res = AugmentedTestResults(self.failures, self.tries) res.summary = self return res def __add__(self, other): merged_reasons = { key: self.error_reasons.get(key, []) + other.error_reasons.get(key, []) for key in set(self.error_reasons.keys()).union( other.error_reasons.keys()) } return Summary( self.failures + other.failures, self.tries + other.tries, self.skipped + other.skipped, merged_reasons) def summarize(self): def print_partition(indent, desc, n, total): print("%s%d %s (%.1f%%)" % (" " * indent, n, desc, n / total * 100)) print() print("%d total test cases:" % self.tries) if not self.tries: return print_partition(1, "skipped", self.skipped, self.tries) for error, reasons in self.error_reasons.items(): print_partition(1, error, len(reasons), self.tries) reason_counts = sorted( collections.Counter(reasons).items(), key=lambda x: x[1], reverse=True) for desc, count in reason_counts: print_partition(2, desc, count, len(reasons)) print_partition(1, "failed", self.failures, self.tries) print_partition( 1, "passed", self.tries - self.skipped - sum(len(reasons) for reasons in self.error_reasons.values()) - self.failures, self.tries) print() def parse_rst_ipython_tests(rst, name, extraglobs=None, optionflags=None): """Extracts examples from an rst file and produce a test suite by running them through pandas to get the expected outputs. """ # Optional dependency. import IPython from traitlets.config import Config def get_indent(line): return len(line) - len(line.lstrip()) def is_example_line(line): line = line.strip() return line and not line.startswith('#') and not line[0] == line[-1] == ':' IMPORT_PANDAS = 'import pandas as pd' example_srcs = [] lines = iter([(lineno, line.rstrip()) for lineno, line in enumerate(rst.split('\n')) if is_example_line(line)] + [(None, 'END')]) # https://ipython.readthedocs.io/en/stable/sphinxext.html lineno, line = next(lines) while True: if line == 'END': break if line.startswith('.. ipython::'): lineno, line = next(lines) indent = get_indent(line) example = [] example_srcs.append((lineno, example)) while get_indent(line) >= indent: if '@verbatim' in line or ':verbatim:' in line or '@savefig' in line: example_srcs.pop() break line = re.sub(r'In \[\d+\]: ', '', line) line = re.sub(r'\.\.\.+:', '', line) example.append(line[indent:]) lineno, line = next(lines) if get_indent(line) == indent and line[indent] not in ')]}': example = [] example_srcs.append((lineno, example)) else: lineno, line = next(lines) # TODO(robertwb): Would it be better to try and detect/compare the actual # objects in two parallel sessions than make (stringified) doctests? examples = [] config = Config() config.HistoryManager.hist_file = ':memory:' config.InteractiveShell.autocall = False config.InteractiveShell.autoindent = False config.InteractiveShell.colors = 'NoColor' set_pandas_options() IP = IPython.InteractiveShell.instance(config=config) IP.run_cell(IMPORT_PANDAS + '\n') IP.run_cell('import numpy as np\n') try: stdout = sys.stdout for lineno, src in example_srcs: sys.stdout = cout = StringIO() src = '\n'.join(src) if src == IMPORT_PANDAS: continue IP.run_cell(src + '\n') output = cout.getvalue() if output: # Strip the prompt. # TODO(robertwb): Figure out how to suppress this. output = re.sub(r'^Out\[\d+\]:[ \t]*\n?', '', output) examples.append(doctest.Example(src, output, lineno=lineno)) finally: sys.stdout = stdout return doctest.DocTest( examples, dict(extraglobs or {}, np=np), name, name, None, None) def test_rst_ipython( rst, name, report=False, wont_implement_ok=(), not_implemented_ok=(), skip=(), **kwargs): """Extracts examples from an rst file and run them through pandas to get the expected output, and then compare them against our dataframe implementation. """ def run_tests(extraglobs, optionflags, **kwargs): # The patched one. tests = parse_rst_ipython_tests(rst, name, extraglobs, optionflags) runner = doctest.DocTestRunner(optionflags=optionflags) set_pandas_options() result = runner.run(tests, **kwargs) if report: runner.summarize() return result result = _run_patched( run_tests, wont_implement_ok={name: wont_implement_ok}, not_implemented_ok={name: not_implemented_ok}, skip={name: skip}, **kwargs) return result def teststring(text, wont_implement_ok=None, not_implemented_ok=None, **kwargs): return teststrings( {'': text}, wont_implement_ok={'': ['*']} if wont_implement_ok else None, not_implemented_ok={'': ['*']} if not_implemented_ok else None, **kwargs) def teststrings(texts, report=False, **runner_kwargs): optionflags = runner_kwargs.pop('optionflags', 0) optionflags |= ( doctest.NORMALIZE_WHITESPACE | doctest.IGNORE_EXCEPTION_DETAIL) parser = doctest.DocTestParser() runner = BeamDataframeDoctestRunner( TestEnvironment(), optionflags=optionflags, **runner_kwargs) globs = { 'pd': runner.fake_pandas_module(), 'np': np, 'option_context': pd.option_context, } with expressions.allow_non_parallel_operations(): for name, text in texts.items(): test = parser.get_doctest(text, globs, name, name, 0) runner.run(test) if report: runner.summarize() return runner.summary().result() def set_pandas_options(): # See # https://github.com/pandas-dev/pandas/blob/a00202d12d399662b8045a8dd3fdac04f18e1e55/doc/source/conf.py#L319 np.random.seed(123456) legacy = None if np.version.version.startswith('2'): legacy = '1.25' np.set_printoptions(precision=4, suppress=True, legacy=legacy) pd.options.display.max_rows = 15 def _run_patched(func, *args, **kwargs): set_pandas_options() # https://github.com/pandas-dev/pandas/blob/1.0.x/setup.cfg#L63 optionflags = kwargs.pop('optionflags', 0) optionflags |= ( doctest.NORMALIZE_WHITESPACE | doctest.IGNORE_EXCEPTION_DETAIL) env = TestEnvironment() use_beam = kwargs.pop('use_beam', True) skip = kwargs.pop('skip', {}) wont_implement_ok = kwargs.pop('wont_implement_ok', {}) not_implemented_ok = kwargs.pop('not_implemented_ok', {}) extraglobs = dict(kwargs.pop('extraglobs', {})) extraglobs['pd'] = env.fake_pandas_module() try: # Unfortunately the runner is not injectable. original_doc_test_runner = doctest.DocTestRunner doctest.DocTestRunner = lambda **kwargs: BeamDataframeDoctestRunner( env, use_beam=use_beam, wont_implement_ok=wont_implement_ok, not_implemented_ok=not_implemented_ok, skip=skip, **kwargs) with expressions.allow_non_parallel_operations(): return func( *args, extraglobs=extraglobs, optionflags=optionflags, **kwargs) finally: doctest.DocTestRunner = original_doc_test_runner def with_run_patched_docstring(target=None): assert target is not None def wrapper(fn): fn.__doc__ = f"""Run all pandas doctests in the specified {target}. Arguments `skip`, `wont_implement_ok`, `not_implemented_ok` are all in the format:: {{ "module.Class.method": ['*'], "module.Class.other_method": [ 'instance.other_method(bad_input)', 'observe_result_of_bad_input()', ], }} `'*'` indicates all examples should be matched, otherwise the list is a list of specific input strings that should be matched. All arguments are kwargs. Args: optionflags (int): Passed through to doctests. extraglobs (dict[str,Any]): Passed through to doctests. use_beam (bool): If true, run a Beam pipeline with partitioned input to verify the examples, else use PartitioningSession to simulate distributed execution. skip (dict[str,str]): A set of examples to skip entirely. If a key is '*', an example will be skipped in all test scenarios. wont_implement_ok (dict[str,str]): A set of examples that are allowed to raise WontImplementError. not_implemented_ok (dict[str,str]): A set of examples that are allowed to raise NotImplementedError. Returns: ~doctest.TestResults: A doctest result describing the passed/failed tests. """ return fn return wrapper @with_run_patched_docstring(target="file") def testfile(*args, **kwargs): return _run_patched(doctest.testfile, *args, **kwargs) @with_run_patched_docstring(target="module") def testmod(*args, **kwargs): return _run_patched(doctest.testmod, *args, **kwargs) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/doctests_test.py0000644000175100017510000001524015206113032023044 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import doctest import os import tempfile import unittest from apache_beam.dataframe import doctests SAMPLE_DOCTEST = ''' >>> df = pd.DataFrame({'Animal': ['Falcon', 'Falcon', ... 'Parrot', 'Parrot'], ... 'Max Speed': [380., 370., 24., 26.]}) >>> df Animal Max Speed 0 Falcon 380.0 1 Falcon 370.0 2 Parrot 24.0 3 Parrot 26.0 >>> df.groupby(['Animal']).mean() Max Speed Animal Falcon 375.0 Parrot 25.0 ''' CHECK_USES_DEFERRED_DATAFRAMES = ''' >>> type(pd).__name__ 'FakePandasObject' >>> type(pd.DataFrame([])) >>> type(pd.DataFrame.from_dict({'a': [1, 2], 'b': [3, 4]})) >>> pd.Index(range(10)) RangeIndex(start=0, stop=10, step=1) ''' WONT_IMPLEMENT_RAISING_TESTS = ''' >>> import apache_beam >>> raise apache_beam.dataframe.frame_base.WontImplementError('anything') ignored exception >>> pd.Series(range(10)).__array__() ignored result ''' ERROR_RAISING_NAME_ERROR_TESTS = ''' >>> import apache_beam >>> raise %s('anything') ignored exception >>> raise NameError ignored exception >>> undefined_name ignored exception >>> 2 + 2 4 >>> raise NameError failed exception ''' WONT_IMPLEMENT_RAISING_NAME_ERROR_TESTS = ERROR_RAISING_NAME_ERROR_TESTS % ( 'apache_beam.dataframe.frame_base.WontImplementError', ) NOT_IMPLEMENTED_RAISING_TESTS = ''' >>> import apache_beam >>> raise NotImplementedError('anything') ignored exception ''' NOT_IMPLEMENTED_RAISING_NAME_ERROR_TESTS = ERROR_RAISING_NAME_ERROR_TESTS % ( 'NotImplementedError', ) FAILED_ASSIGNMENT = ''' >>> def foo(): raise NotImplementedError() >>> res = 'old_value' >>> res = foo() >>> print(res) ignored NameError ''' RST_IPYTHON = ''' Here is an example .. ipython:: 2 + 2 some multi-line examples .. ipython:: def foo(x): return x * x foo(4) foo( 4 ) In [100]: def foo(x): ....: return x * x * x ....: foo(5) history is preserved foo(3) foo(4) and finally an example with pandas .. ipython:: pd.Series([1, 2, 3]).max() This one should be skipped: .. ipython:: @verbatim not run or tested and someting that'll fail (due to fake vs. real pandas) .. ipython:: type(pd) ''' class DoctestTest(unittest.TestCase): def test_good(self): result = doctests.teststring(SAMPLE_DOCTEST, report=False) self.assertEqual(result.attempted, 3) self.assertEqual(result.failed, 0) def test_failure(self): result = doctests.teststring( SAMPLE_DOCTEST.replace('25.0', '25.00001'), report=False) self.assertEqual(result.attempted, 3) self.assertEqual(result.failed, 1) def test_uses_beam_dataframes(self): result = doctests.teststring(CHECK_USES_DEFERRED_DATAFRAMES, report=False) self.assertNotEqual(result.attempted, 0) self.assertEqual(result.failed, 0) def test_file(self): with tempfile.TemporaryDirectory() as dir: filename = os.path.join(dir, 'tests.py') with open(filename, 'w') as fout: fout.write(SAMPLE_DOCTEST) result = doctests.testfile(filename, module_relative=False, report=False) self.assertEqual(result.attempted, 3) self.assertEqual(result.failed, 0) def test_file_uses_beam_dataframes(self): with tempfile.TemporaryDirectory() as dir: filename = os.path.join(dir, 'tests.py') with open(filename, 'w') as fout: fout.write(CHECK_USES_DEFERRED_DATAFRAMES) result = doctests.testfile(filename, module_relative=False, report=False) self.assertNotEqual(result.attempted, 0) self.assertEqual(result.failed, 0) def test_wont_implement(self): result = doctests.teststring( WONT_IMPLEMENT_RAISING_TESTS, optionflags=doctest.ELLIPSIS, wont_implement_ok=True) self.assertNotEqual(result.attempted, 0) self.assertEqual(result.failed, 0) result = doctests.teststring( WONT_IMPLEMENT_RAISING_TESTS, optionflags=doctest.IGNORE_EXCEPTION_DETAIL, wont_implement_ok=True) self.assertNotEqual(result.attempted, 0) self.assertEqual(result.failed, 0) def test_wont_implement_followed_by_name_error(self): result = doctests.teststring( WONT_IMPLEMENT_RAISING_NAME_ERROR_TESTS, optionflags=doctest.ELLIPSIS, wont_implement_ok=True) self.assertEqual(result.attempted, 6) self.assertEqual(result.failed, 1) # Only the very last one. def test_not_implemented(self): result = doctests.teststring( NOT_IMPLEMENTED_RAISING_TESTS, optionflags=doctest.ELLIPSIS, not_implemented_ok=True) self.assertNotEqual(result.attempted, 0) self.assertEqual(result.failed, 0) result = doctests.teststring( NOT_IMPLEMENTED_RAISING_TESTS, optionflags=doctest.IGNORE_EXCEPTION_DETAIL, not_implemented_ok=True) self.assertNotEqual(result.attempted, 0) self.assertEqual(result.failed, 0) def test_not_implemented_followed_by_name_error(self): result = doctests.teststring( NOT_IMPLEMENTED_RAISING_NAME_ERROR_TESTS, optionflags=doctest.ELLIPSIS, not_implemented_ok=True) self.assertEqual(result.attempted, 6) self.assertEqual(result.failed, 1) # Only the very last one. def test_failed_assignment(self): result = doctests.teststring( FAILED_ASSIGNMENT, optionflags=doctest.ELLIPSIS, not_implemented_ok=True) self.assertNotEqual(result.attempted, 0) self.assertEqual(result.failed, 0) def test_rst_ipython(self): try: # pylint: disable=unused-import import IPython except ImportError: raise unittest.SkipTest('IPython not available') result = doctests.test_rst_ipython(RST_IPYTHON, 'test_rst_ipython') self.assertEqual(result.attempted, 8) self.assertEqual(result.failed, 1) # Only the very last one. if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/expressions.py0000644000175100017510000003514115206113032022541 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import contextlib import random import threading from collections.abc import Callable from collections.abc import Iterable from typing import Any from typing import Generic from typing import Optional from typing import TypeVar from apache_beam.dataframe import partitionings class Session(object): """A session represents a mapping of expressions to concrete values. The bindings typically include required placeholders, but may be any intermediate expression as well. """ def __init__(self, bindings=None): self._bindings = dict(bindings or {}) def evaluate(self, expr: 'Expression') -> Any: if expr not in self._bindings: self._bindings[expr] = expr.evaluate_at(self) return self._bindings[expr] def lookup(self, expr: 'Expression') -> Any: return self._bindings[expr] class PartitioningSession(Session): """An extension of Session that enforces actual partitioning of inputs. Each expression is evaluated multiple times for various supported partitionings determined by its `requires_partition_by` specification. For each tested partitioning, the input is partitioned and the expression is evaluated on each partition separately, as if this were actually executed in a parallel manner. For each input partitioning, the results are verified to be partitioned appropriately according to the expression's `preserves_partition_by` specification. For testing only. """ def evaluate(self, expr): import collections import pandas as pd def is_scalar(expr): return not isinstance(expr.proxy(), pd.core.generic.NDFrame) if expr not in self._bindings: if is_scalar(expr) or not expr.args(): result = super().evaluate(expr) else: scaler_args = [arg for arg in expr.args() if is_scalar(arg)] def evaluate_with(input_partitioning): parts = collections.defaultdict( lambda: Session({arg: self.evaluate(arg) for arg in scaler_args})) for arg in expr.args(): if not is_scalar(arg): input = self.evaluate(arg) for key, part in input_partitioning.test_partition_fn(input): parts[key]._bindings[arg] = part if not parts: parts[None] # Create at least one entry. results = [] for session in parts.values(): if any(len(session.lookup(arg)) for arg in expr.args() if not is_scalar(arg)): results.append(session.evaluate(expr)) expected_output_partitioning = output_partitioning( expr, input_partitioning) if not expected_output_partitioning.check(results): raise AssertionError( f"""Expression does not preserve partitioning! Expression: {expr} Requires: {expr.requires_partition_by()} Preserves: {expr.preserves_partition_by()} Input partitioning: {input_partitioning} Expected output partitioning: {expected_output_partitioning} """) if results: return pd.concat(results) else: # Choose any single session. return next(iter(parts.values())).evaluate(expr) # Store random state so it can be re-used for each execution, in case # the expression is part of a test that relies on the random seed. random_state = random.getstate() result = None # Run with all supported partitionings s.t. the smallest subpartitioning # is used last. This way the final result is computed with the most # challenging partitioning. Avoids heisenbugs where sometimes the result # is computed trivially with Singleton partitioning and passes. for input_partitioning in sorted(set([expr.requires_partition_by(), partitionings.Arbitrary(), partitionings.JoinIndex(), partitionings.Index(), partitionings.Singleton()])): if not expr.requires_partition_by().is_subpartitioning_of( input_partitioning): continue random.setstate(random_state) result = evaluate_with(input_partitioning) assert result is not None self._bindings[expr] = result return self._bindings[expr] # The return type of an Expression T = TypeVar('T') def output_partitioning(expr, input_partitioning): """ Return the expected output partitioning for `expr` when it's input is partitioned by `input_partitioning`. For internal use only; No backward compatibility guarantees """ assert expr.requires_partition_by().is_subpartitioning_of(input_partitioning) if expr.preserves_partition_by().is_subpartitioning_of(input_partitioning): return min(input_partitioning, expr.preserves_partition_by()) else: return partitionings.Arbitrary() class Expression(Generic[T]): """An expression is an operation bound to a set of arguments. An expression represents a deferred tree of operations, which can be evaluated at a specific bindings of root expressions to values. requires_partition_by indicates the upper bound of a set of partitionings that are acceptable inputs to this expression. The expression should be able to produce the correct result when given input(s) partitioned by its requires_partition_by attribute, or by any partitoning that is _not_ a subpartitioning of it. preserves_partition_by indicates the upper bound of a set of partitionings that can be preserved by this expression. When the input(s) to this expression are partitioned by preserves_partition_by, or by any partitioning that is _not_ a subpartitioning of it, this expression should produce output(s) partitioned by the same partitioning. However, if the partitioning of an expression's input is a subpartitioning of the partitioning that it preserves, the output is presumed to have no particular partitioning (i.e. Arbitrary()). For example, let's look at an "element-wise operation", that has no partitioning requirement, and preserves any partitioning given to it:: requires_partition_by = Arbitrary() -----------------------------+ | +-----------+-------------+---------- ... ----+---------| | | | | | Singleton() < Index([i]) < Index([i, j]) < ... < Index() < Arbitrary() | | | | | +-----------+-------------+---------- ... ----+---------| | preserves_partition_by = Arbitrary() ----------------------------+ As a more interesting example, consider this expression, which requires Index partitioning, and preserves just Singleton partitioning:: requires_partition_by = Index() -----------------------+ | +-----------+-------------+---------- ... ----| | | | | Singleton() < Index([i]) < Index([i, j]) < ... < Index() < Arbitrary() | | preserves_partition_by = Singleton() Note that any non-Arbitrary partitioning is an acceptable input for this expression. However, unless the inputs are Singleton-partitioned, the expression makes no guarantees about the partitioning of the output. """ def __init__(self, name: str, proxy: T, _id: Optional[str] = None): self._name = name self._proxy = proxy # Store for preservation through pickling. self._id = _id or '%s_%s_%s' % (name, type(proxy).__name__, id(self)) def proxy(self) -> T: return self._proxy def __hash__(self): return hash(self._id) def __eq__(self, other): return self._id == other._id def __repr__(self): return '%s[%s]' % (self.__class__.__name__, self._id) def placeholders(self): """Returns all the placeholders that self depends on.""" raise NotImplementedError(type(self)) def evaluate_at(self, session: Session) -> T: """Returns the result of self with the bindings given in session.""" raise NotImplementedError(type(self)) def requires_partition_by(self) -> partitionings.Partitioning: """Returns the partitioning, if any, require to evaluate this expression. Returns partitioning.Arbitrary() to require no partitioning is required. """ raise NotImplementedError(type(self)) def preserves_partition_by(self) -> partitionings.Partitioning: """Returns the partitioning, if any, preserved by this expression. This gives an upper bound on the partitioning of its ouput. The actual partitioning of the output may be less strict (e.g. if the input was less partitioned). """ raise NotImplementedError(type(self)) class PlaceholderExpression(Expression): """An expression whose value must be explicitly bound in the session.""" def __init__( self, proxy: T, reference: Any = None, ): """Initialize a placeholder expression. Args: proxy: A proxy object with the type expected to be bound to this expression. Used for type checking at pipeline construction time. """ super().__init__('placeholder', proxy) self._reference = reference def placeholders(self): return frozenset([self]) def args(self): return () def evaluate_at(self, session): return session.lookup(self) def requires_partition_by(self): return partitionings.Arbitrary() def preserves_partition_by(self): return partitionings.Index() class ConstantExpression(Expression): """An expression whose value is known at pipeline construction time.""" def __init__(self, value: T, proxy: Optional[T] = None): """Initialize a constant expression. Args: value: The constant value to be produced by this expression. proxy: (Optional) a proxy object with same type as `value` to use for rapid type checking at pipeline construction time. If not provided, `value` will be used directly. """ if proxy is None: proxy = value super().__init__('constant', proxy) self._value = value def placeholders(self): return frozenset() def args(self): return () def evaluate_at(self, session): return self._value def requires_partition_by(self): return partitionings.Arbitrary() def preserves_partition_by(self): return partitionings.Arbitrary() class ComputedExpression(Expression): """An expression whose value must be computed at pipeline execution time.""" def __init__( self, name: str, func: Callable[..., T], args: Iterable[Expression], proxy: Optional[T] = None, _id: Optional[str] = None, requires_partition_by: partitionings.Partitioning = partitionings.Index(), preserves_partition_by: partitionings.Partitioning = partitionings. Singleton(), ): """Initialize a computed expression. Args: name: The name of this expression. func: The function that will be used to compute the value of this expression. Should accept arguments of the types returned when evaluating the `args` expressions. args: The list of expressions that will be used to produce inputs to `func`. proxy: (Optional) a proxy object with same type as the objects that this ComputedExpression will produce at execution time. If not provided, a proxy will be generated using `func` and the proxies of `args`. _id: (Optional) a string to uniquely identify this expression. requires_partition_by: The required (common) partitioning of the args. preserves_partition_by: The level of partitioning preserved. """ if (not _get_allow_non_parallel() and isinstance(requires_partition_by, partitionings.Singleton)): reason = requires_partition_by.reason or ( f"Encountered non-parallelizable form of {name!r}.") raise NonParallelOperation( f"{reason}\n" "Consider using an allow_non_parallel_operations block if you're " "sure you want to do this. See " "https://s.apache.org/dataframe-non-parallel-operations for more " "information.") args = tuple(args) if proxy is None: proxy = func(*(arg.proxy() for arg in args)) super().__init__(name, proxy, _id) self._func = func self._args = args self._requires_partition_by = requires_partition_by self._preserves_partition_by = preserves_partition_by def placeholders(self): if not hasattr(self, '_placeholders'): self._placeholders = frozenset.union( frozenset(), *[arg.placeholders() for arg in self.args()]) return self._placeholders def args(self): return self._args def evaluate_at(self, session): return self._func(*(session.evaluate(arg) for arg in self._args)) def requires_partition_by(self): return self._requires_partition_by def preserves_partition_by(self): return self._preserves_partition_by def elementwise_expression(name, func, args): return ComputedExpression( name, func, args, requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary()) _ALLOW_NON_PARALLEL = threading.local() _ALLOW_NON_PARALLEL.value = False def _get_allow_non_parallel(): return _ALLOW_NON_PARALLEL.value @contextlib.contextmanager def allow_non_parallel_operations(allow=True): if allow is None: yield else: old_value, _ALLOW_NON_PARALLEL.value = _ALLOW_NON_PARALLEL.value, allow try: yield finally: _ALLOW_NON_PARALLEL.value = old_value class NonParallelOperation(Exception): def __init__(self, msg): super().__init__(self, msg) self.msg = msg ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/expressions_test.py0000644000175100017510000001140215206113032023572 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import unittest import pandas as pd from apache_beam.dataframe import expressions from apache_beam.dataframe import partitionings class ExpressionTest(unittest.TestCase): def test_placeholder_expression(self): a = expressions.PlaceholderExpression(None) b = expressions.PlaceholderExpression(None) session = expressions.Session({a: 1, b: 2}) self.assertEqual(session.evaluate(a), 1) self.assertEqual(session.evaluate(b), 2) def test_constant_expresion(self): two = expressions.ConstantExpression(2) session = expressions.Session({}) self.assertEqual(session.evaluate(two), 2) def test_computed_expression(self): a = expressions.PlaceholderExpression(0) b = expressions.PlaceholderExpression(0) a_plus_b = expressions.ComputedExpression('add', lambda a, b: a + b, [a, b]) session = expressions.Session({a: 1, b: 2}) self.assertEqual(session.evaluate(a_plus_b), 3) def test_expression_proxy(self): a = expressions.PlaceholderExpression(1) b = expressions.PlaceholderExpression(2) a_plus_b = expressions.ComputedExpression('add', lambda a, b: a + b, [a, b]) self.assertEqual(a_plus_b.proxy(), 3) def test_expression_proxy_error(self): a = expressions.PlaceholderExpression(1) b = expressions.PlaceholderExpression('s') with self.assertRaises(TypeError): expressions.ComputedExpression('add', lambda a, b: a + b, [a, b]) def test_preserves_singleton_output_partitioning(self): # Empty DataFrame with one column and two index levels input_expr = expressions.ConstantExpression( pd.DataFrame(columns=["column"], index=[[], []])) preserves_only_singleton = expressions.ComputedExpression( 'preserves_only_singleton', # index is replaced with an entirely new one, so # if we were partitioned by Index we're not anymore. lambda df: df.set_index('column'), [input_expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) for partitioning in (partitionings.Singleton(), ): self.assertEqual( expressions.output_partitioning( preserves_only_singleton, partitioning), partitioning, f"Should preserve {partitioning}") for partitioning in (partitionings.Index([0]), partitionings.Index(), partitionings.Arbitrary()): self.assertEqual( expressions.output_partitioning( preserves_only_singleton, partitioning), partitionings.Arbitrary(), f"Should NOT preserve {partitioning}") def test_preserves_index_output_partitioning(self): # Empty DataFrame with two columns and two index levels input_expr = expressions.ConstantExpression( pd.DataFrame(columns=["foo", "bar"], index=[[], []])) preserves_partial_index = expressions.ComputedExpression( 'preserves_partial_index', # This adds an additional index level, so we'd only preserve # partitioning on the two index levels that existed before. lambda df: df.set_index('foo', append=True), [input_expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Index([0, 1])) for partitioning in ( partitionings.Singleton(), partitionings.Index([0]), partitionings.Index([1]), partitionings.Index([0, 1]), ): self.assertEqual( expressions.output_partitioning( preserves_partial_index, partitioning), partitioning, f"Should preserve {partitioning}") for partitioning in (partitionings.Index([0, 1, 2]), partitionings.Index(), partitionings.Arbitrary()): self.assertEqual( expressions.output_partitioning( preserves_partial_index, partitioning), partitionings.Arbitrary(), f"Should NOT preserve {partitioning}") if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/frame_base.py0000644000175100017510000006150415206113032022245 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import functools import operator import re from collections.abc import Callable from inspect import cleandoc from inspect import getfullargspec from inspect import isclass from inspect import ismodule from inspect import unwrap from typing import Any from typing import Optional from typing import Union import pandas as pd from apache_beam.dataframe import expressions from apache_beam.dataframe import partitionings class DeferredBase(object): _pandas_type_map: dict[Union[type, None], type] = {} def __init__(self, expr): self._expr = expr @classmethod def _register_for(cls, pandas_type): def wrapper(deferred_type): cls._pandas_type_map[pandas_type] = deferred_type return deferred_type return wrapper @classmethod def wrap(cls, expr, split_tuples=True): proxy_type = type(expr.proxy()) if proxy_type is tuple and split_tuples: def get(ix): return expressions.ComputedExpression( # yapf: disable 'get_%d' % ix, lambda t: t[ix], [expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) return tuple(cls.wrap(get(ix)) for ix in range(len(expr.proxy()))) elif proxy_type in cls._pandas_type_map: wrapper_type = cls._pandas_type_map[proxy_type] else: if expr.requires_partition_by() != partitionings.Singleton(): raise ValueError( 'Scalar expression %s of type %s partitoned by non-singleton %s' % (expr, proxy_type, expr.requires_partition_by())) wrapper_type = _DeferredScalar return wrapper_type(expr) def _elementwise( self, func, name=None, other_args=(), other_kwargs=None, inplace=False): other_kwargs = other_kwargs or {} return _elementwise_function( func, name, inplace=inplace)(self, *other_args, **other_kwargs) def __reduce__(self): return UnusableUnpickledDeferredBase, (str(self), ) class UnusableUnpickledDeferredBase(object): """Placeholder object used to break the transitive pickling chain in case a DeferredBase accidentially gets pickled (e.g. as part of globals). Trying to use this object after unpickling is a bug and will result in an error. """ def __init__(self, name): self._name = name def __repr__(self): return 'UnusablePickledDeferredBase(%r)' % self.name class DeferredFrame(DeferredBase): pass class _DeferredScalar(DeferredBase): def apply(self, func, name=None, args=()): if name is None: name = func.__name__ with expressions.allow_non_parallel_operations( all(isinstance(arg, _DeferredScalar) for arg in args) or None): return DeferredFrame.wrap( expressions.ComputedExpression( name, func, [self._expr] + [arg._expr for arg in args], requires_partition_by=partitionings.Singleton())) def __neg__(self): return self.apply(operator.neg) def __pos__(self): return self.apply(operator.pos) def __invert__(self): return self.apply(operator.invert) def __repr__(self): return f"DeferredScalar[type={type(self._expr.proxy())}]" def __bool__(self): # TODO(BEAM-11951): Link to documentation raise TypeError( "Testing the truth value of a deferred scalar is not " "allowed. It's not possible to branch on the result of " "deferred operations.") def _scalar_binop(op): def binop(self, other): if not isinstance(other, DeferredBase): return self.apply(lambda left: getattr(left, op)(other), name=op) elif isinstance(other, _DeferredScalar): return self.apply( lambda left, right: getattr(left, op)(right), name=op, args=[other]) else: return NotImplemented return binop for op in ['__add__', '__sub__', '__mul__', '__div__', '__truediv__', '__floordiv__', '__mod__', '__divmod__', '__pow__', '__and__', '__or__']: setattr(_DeferredScalar, op, _scalar_binop(op)) DeferredBase._pandas_type_map[None] = _DeferredScalar def name_and_func(method: Union[str, Callable]) -> tuple[str, Callable]: """For the given method name or method, return the method name and the method itself. For internal use only. No backwards compatibility guarantees.""" if isinstance(method, str): method_str = method func = lambda df, *args, **kwargs: getattr(df, method_str)(*args, **kwargs) return method, func else: return method.__name__, method def _elementwise_method( func, name=None, restrictions=None, inplace=False, base=None): return _proxy_method( func, name, restrictions, inplace, base, requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary()) def _proxy_method( func, name=None, restrictions=None, inplace=False, base=None, *, requires_partition_by: partitionings.Partitioning, preserves_partition_by: partitionings.Partitioning, ): if name is None: name, func = name_and_func(func) if base is None: raise ValueError("base is required for _proxy_method") return _proxy_function( func, name, restrictions, inplace, base, requires_partition_by=requires_partition_by, preserves_partition_by=preserves_partition_by) def _elementwise_function( func, name=None, restrictions=None, inplace=False, base=None): return _proxy_function( func, name, restrictions, inplace, base, requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary()) def _proxy_function( func: Union[Callable, str], name: Optional[str] = None, restrictions: Optional[dict[str, Union[Any, list[Any]]]] = None, inplace: bool = False, base: Optional[type] = None, *, requires_partition_by: partitionings.Partitioning, preserves_partition_by: partitionings.Partitioning, ): if name is None: if isinstance(func, str): name = func else: name = func.__name__ if restrictions is None: restrictions = {} def wrapper(*args, **kwargs): for key, values in restrictions.items(): if key in kwargs: value = kwargs[key] else: try: ix = getfullargspec(func).args.index(key) except ValueError: # TODO: fix for delegation? continue if len(args) <= ix: continue value = args[ix] if callable(values): check = values elif isinstance(values, list): check = lambda x, values=values: x in values else: check = lambda x, value=value: x == value if not check(value): raise NotImplementedError( '%s=%s not supported for %s' % (key, value, name)) deferred_arg_indices = [] deferred_arg_exprs = [] constant_args = [None] * len(args) from apache_beam.dataframe.frames import _DeferredIndex for ix, arg in enumerate(args): if isinstance(arg, DeferredBase): deferred_arg_indices.append(ix) deferred_arg_exprs.append(arg._expr) elif isinstance(arg, _DeferredIndex): # TODO(robertwb): Consider letting indices pass through as indices. # This would require updating the partitioning code, as indices don't # have indices. deferred_arg_indices.append(ix) deferred_arg_exprs.append( expressions.ComputedExpression( 'index_as_series', lambda ix: ix.index.to_series(), # yapf break [arg._frame._expr], preserves_partition_by=partitionings.Singleton(), requires_partition_by=partitionings.Arbitrary())) elif isinstance(arg, pd.core.generic.NDFrame): deferred_arg_indices.append(ix) deferred_arg_exprs.append(expressions.ConstantExpression(arg, arg[0:0])) else: constant_args[ix] = arg deferred_kwarg_keys = [] deferred_kwarg_exprs = [] constant_kwargs = {key: None for key in kwargs} for key, arg in kwargs.items(): if isinstance(arg, DeferredBase): deferred_kwarg_keys.append(key) deferred_kwarg_exprs.append(arg._expr) elif isinstance(arg, pd.core.generic.NDFrame): deferred_kwarg_keys.append(key) deferred_kwarg_exprs.append( expressions.ConstantExpression(arg, arg[0:0])) else: constant_kwargs[key] = arg deferred_exprs = deferred_arg_exprs + deferred_kwarg_exprs if inplace: actual_func = _copy_and_mutate(func) else: actual_func = func def apply(*actual_args): actual_args, actual_kwargs = (actual_args[:len(deferred_arg_exprs)], actual_args[len(deferred_arg_exprs):]) full_args = list(constant_args) for ix, arg in zip(deferred_arg_indices, actual_args): full_args[ix] = arg full_kwargs = dict(constant_kwargs) for key, arg in zip(deferred_kwarg_keys, actual_kwargs): full_kwargs[key] = arg return actual_func(*full_args, **full_kwargs) if (requires_partition_by.is_subpartitioning_of(partitionings.Index()) and sum(isinstance(arg.proxy(), pd.core.generic.NDFrame) for arg in deferred_exprs) > 1): # Implicit join on index if there is more than one indexed input. actual_requires_partition_by = partitionings.JoinIndex() else: actual_requires_partition_by = requires_partition_by result_expr = expressions.ComputedExpression( name, apply, deferred_exprs, requires_partition_by=actual_requires_partition_by, preserves_partition_by=preserves_partition_by) if inplace: args[0]._expr = result_expr else: return DeferredFrame.wrap(result_expr) wrapper.__name__ = name if restrictions: wrapper.__doc__ = "\n".join( f"Only {kw}={value!r} is supported" for (kw, value) in restrictions.items()) if base is not None: return with_docs_from(base)(wrapper) else: return wrapper def _prettify_pandas_type(pandas_type): if pandas_type in (pd.DataFrame, pd.Series): return f'pandas.{pandas_type.__name__}' elif isclass(pandas_type): return f'{pandas_type.__module__}.{pandas_type.__name__}' elif ismodule(pandas_type): return pandas_type.__name__ else: raise TypeError(pandas_type) def wont_implement_method(base_type, name, reason=None, explanation=None): """Generate a stub method that raises WontImplementError. Note either reason or explanation must be specified. If both are specified, explanation is ignored. Args: base_type: The pandas type of the method that this is trying to replicate. name: The name of the method that this is aiming to replicate. reason: If specified, use data from the corresponding entry in ``_WONT_IMPLEMENT_REASONS`` to generate a helpful exception message and docstring for the method. explanation: If specified, use this string as an explanation for why this operation is not supported when generating an exception message and docstring. """ if reason is not None: if reason not in _WONT_IMPLEMENT_REASONS: raise AssertionError( f"reason must be one of {list(_WONT_IMPLEMENT_REASONS.keys())}, " f"got {reason!r}") reason_data = _WONT_IMPLEMENT_REASONS[reason] elif explanation is not None: reason_data = {'explanation': explanation} else: raise ValueError("One of (reason, explanation) must be specified") def wrapper(*args, **kwargs): raise WontImplementError( f"'{name}' is not yet supported {reason_data['explanation']}", reason=reason) wrapper.__name__ = name wrapper.__doc__ = ( f":meth:`{_prettify_pandas_type(base_type)}.{name}` is not yet supported " f"in the Beam DataFrame API {reason_data['explanation']}") if 'url' in reason_data: wrapper.__doc__ += f"\n\n For more information see {reason_data['url']}." return wrapper def not_implemented_method(op, issue='20318', base_type=None): """Generate a stub method for ``op`` that simply raises a NotImplementedError. For internal use only. No backwards compatibility guarantees.""" assert base_type is not None, "base_type must be specified" issue_url = f"https://issues.apache.org/jira/{issue}." if issue.startswith( "BEAM-") else f"https://github.com/apache/beam/issues/{issue}" def wrapper(*args, **kwargs): raise NotImplementedError( f"{op!r} is not implemented yet. " f"If support for {op!r} is important to you, please let the Beam " "community know by writing to user@beam.apache.org " "(see https://beam.apache.org/community/contact-us/) or commenting on " f"{issue_url}") wrapper.__name__ = op wrapper.__doc__ = ( f":meth:`{_prettify_pandas_type(base_type)}.{op}` is not implemented yet " "in the Beam DataFrame API.\n\n" f"If support for {op!r} is important to you, please let the Beam " "community know by `writing to user@beam.apache.org " "`_ or commenting on " f"`{issue} <{issue_url}>`_.") return wrapper def _copy_and_mutate(func): def wrapper(self, *args, **kwargs): copy = self.copy() func(copy, *args, **kwargs) return copy return wrapper def maybe_inplace(func): """Handles the inplace= kwarg available in many pandas operations. This decorator produces a new function handles the inplace kwarg. When `inplace=False`, the new function simply yields the result of `func` directly. When `inplace=True`, the output of `func` is used to replace this instances expression. The result is that any operations applied to this instance after the inplace operation will refernce the updated expression. For internal use only. No backwards compatibility guarantees.""" @functools.wraps(func) def wrapper(self, inplace=False, **kwargs): result = func(self, **kwargs) if inplace: self._expr = result._expr else: return result return wrapper def args_to_kwargs(base_type, removed_method=False, removed_args=None): """Convert all args to kwargs before calling the decorated function. When applied to a function, this decorator creates a new function that always calls the wrapped function with *only* keyword arguments. It inspects the argspec for the identically-named method on `base_type` to determine the name to use for arguments that are converted to keyword arguments. For internal use only. No backwards compatibility guarantees. Args: base_type: The pandas type of the method that this is trying to replicate. removed_method: Whether this method has been removed in the running Pandas version. removed_args: If not empty, which arguments have been dropped in the running Pandas version. """ def wrap(func): if removed_method: # Do no processing, let Beam function itself raise the error if called. return func removed_arg_names = removed_args if removed_args is not None else [] # We would need to add position only arguments if they ever become a thing # in Pandas (as of 2.1 currently they aren't). base_arg_spec = getfullargspec(unwrap(getattr(base_type, func.__name__))) base_arg_names = base_arg_spec.args # Some arguments are keyword only and we still want to check against those. all_possible_base_arg_names = base_arg_names + base_arg_spec.kwonlyargs beam_arg_names = getfullargspec(func).args if not_found := (set(beam_arg_names) - set(all_possible_base_arg_names) - set(removed_arg_names)): raise TypeError( f"Beam definition of {func.__name__} has arguments that are not found" f" in the base version of the function: {not_found}") @functools.wraps(func) def wrapper(*args, **kwargs): if len(args) > len(base_arg_names): raise TypeError(f"{func.__name__} got too many positioned arguments.") for name, value in zip(base_arg_names, args): if name in kwargs: raise TypeError( "%s() got multiple values for argument '%s'" % (func.__name__, name)) kwargs[name] = value # Still have to populate these for the Beam function signature. if removed_args: for name in removed_args: if name not in kwargs: kwargs[name] = None return func(**kwargs) return wrapper return wrap BEAM_SPECIFIC = "Differences from pandas" SECTION_ORDER = [ 'Parameters', 'Returns', 'Raises', BEAM_SPECIFIC, 'See Also', 'Notes', 'Examples' ] EXAMPLES_DISCLAIMER = ( "**NOTE:** These examples are pulled directly from the pandas " "documentation for convenience. Usage of the Beam DataFrame API will look " "different because it is a deferred API.") EXAMPLES_DIFFERENCES = EXAMPLES_DISCLAIMER + ( " In addition, some arguments shown here may not be supported, see " f"**{BEAM_SPECIFIC!r}** for details.") def with_docs_from(base_type, name=None, removed_method=False): """Decorator that updates the documentation from the wrapped function to duplicate the documentation from the identically-named method in `base_type`. Any docstring on the original function will be included in the new function under a "Differences from pandas" heading. removed_method used in cases where a method has been removed in a later version of Pandas. """ def wrap(func): if removed_method: func.__doc__ = ( "This method has been removed in the current version of Pandas.") return func fn_name = name or func.__name__ orig_doc = getattr(base_type, fn_name).__doc__ if orig_doc is None: return func orig_doc = cleandoc(orig_doc) section_splits = re.split(r'^(.*)$\n^-+$\n', orig_doc, flags=re.MULTILINE) intro = section_splits[0].strip() sections = dict(zip(section_splits[1::2], section_splits[2::2])) beam_has_differences = bool(func.__doc__) for header, content in sections.items(): content = content.strip() # Replace references to version numbers so its clear they reference # *pandas* versions content = re.sub(r'([Vv]ersion\s+[\d\.]+)', r'pandas \1', content) if header == "Examples": content = '\n\n'.join([ ( EXAMPLES_DIFFERENCES if beam_has_differences else EXAMPLES_DISCLAIMER), # Indent the examples under a doctest heading, # add skipif option. This makes sure our doctest # framework doesn't run these pandas tests. (".. doctest::\n" " :skipif: True"), re.sub(r"^", " ", content, flags=re.MULTILINE), ]) elif "Examples" in content and ">>>" in content: # some new examples don't have the correct heading # this catches those examples split_content = content.split("Examples") content = '\n\n'.join([ split_content[0], "Examples\n", # Indent the code snippet under a doctest heading, # add skipif option. This makes sure our doctest # framework doesn't run these pandas tests. (".. doctest::\n" " :skipif: True"), re.sub(r"^", " ", content, flags=re.MULTILINE), split_content[1] ]) else: content = content.replace('DataFrame', 'DeferredDataFrame').replace( 'Series', 'DeferredSeries') sections[header] = content if beam_has_differences: sections[BEAM_SPECIFIC] = cleandoc(func.__doc__) else: sections[BEAM_SPECIFIC] = ( "This operation has no known divergences from the " "pandas API.") def format_section(header): return '\n'.join([header, ''.join('-' for _ in header), sections[header]]) func.__doc__ = '\n\n'.join([intro] + [ format_section(header) for header in SECTION_ORDER if header in sections ]) return func return wrap def populate_defaults(base_type, removed_method=False, removed_args=None): """Populate default values for keyword arguments in decorated function. When applied to a function, this decorator creates a new function with default values for all keyword arguments, based on the default values for the identically-named method on `base_type`. For internal use only. No backwards compatibility guarantees. Args: base_type: The pandas type of the method that this is trying to replicate. removed_method: Whether this method has been removed in the running Pandas version. removed_args: If not empty, which arguments have been dropped in the running Pandas version. """ def wrap(func): if removed_method: return func base_argspec = getfullargspec(unwrap(getattr(base_type, func.__name__))) if not base_argspec.defaults and not base_argspec.kwonlydefaults: return func arg_to_default = {} if base_argspec.defaults: arg_to_default.update( zip( base_argspec.args[-len(base_argspec.defaults):], base_argspec.defaults)) if base_argspec.kwonlydefaults: arg_to_default.update(base_argspec.kwonlydefaults) unwrapped_func = unwrap(func) # args that do not have defaults in func, but do have defaults in base func_argspec = getfullargspec(unwrapped_func) num_non_defaults = len(func_argspec.args) - len(func_argspec.defaults or ()) defaults_to_populate = set( func_argspec.args[:num_non_defaults]).intersection( arg_to_default.keys()) if removed_args: defaults_to_populate -= set(removed_args) # In pandas 2, many methods rely on the default copy=None # to mean that copy is the value of copy_on_write. Since # copy_on_write will always be true for Beam, just fill it # in here. In pandas 1, the default was True anyway. if 'copy' in arg_to_default and arg_to_default['copy'] is None: arg_to_default['copy'] = True @functools.wraps(func) def wrapper(**kwargs): for name in defaults_to_populate: if name not in kwargs: kwargs[name] = arg_to_default[name] return func(**kwargs) return wrapper return wrap _WONT_IMPLEMENT_REASONS = { 'order-sensitive': { 'explanation': "because it is sensitive to the order of the data.", 'url': 'https://s.apache.org/dataframe-order-sensitive-operations', }, 'non-deferred-columns': { 'explanation': ( "because the columns in the output DataFrame depend " "on the data."), 'url': 'https://s.apache.org/dataframe-non-deferred-columns', }, 'non-deferred-result': { 'explanation': ( "because it produces an output type that is not " "deferred."), 'url': 'https://s.apache.org/dataframe-non-deferred-result', }, 'plotting-tools': { 'explanation': "because it is a plotting tool.", 'url': 'https://s.apache.org/dataframe-plotting-tools', }, 'event-time-semantics': { 'explanation': ( "because implementing it would require integrating with Beam " "event-time semantics"), 'url': 'https://s.apache.org/dataframe-event-time-semantics', }, 'deprecated': { 'explanation': "because it is deprecated in pandas.", }, 'experimental': { 'explanation': "because it is experimental in pandas.", }, } class WontImplementError(NotImplementedError): """An subclass of NotImplementedError to raise indicating that implementing the given method is not planned. Raising this error will also prevent this doctests from being validated when run with the beam dataframe validation doctest runner. """ def __init__(self, msg, reason=None): if reason is not None: if reason not in _WONT_IMPLEMENT_REASONS: raise AssertionError( f"reason must be one of {list(_WONT_IMPLEMENT_REASONS.keys())}, " f"got {reason!r}") reason_data = _WONT_IMPLEMENT_REASONS[reason] if 'url' in reason_data: msg = f"{msg}\nFor more information see {reason_data['url']}." super().__init__(msg) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/frame_base_test.py0000644000175100017510000001637715206113032023314 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import unittest import pandas as pd from apache_beam.dataframe import expressions from apache_beam.dataframe import frame_base from apache_beam.dataframe import frames class FrameBaseTest(unittest.TestCase): def test_elementwise_func(self): a = pd.Series([1, 2, 3]) b = pd.Series([100, 200, 300]) empty_proxy = a[:0] x = frames.DeferredSeries(expressions.PlaceholderExpression(empty_proxy)) y = frames.DeferredSeries(expressions.PlaceholderExpression(empty_proxy)) sub = frame_base._elementwise_function(lambda x, y: x - y) session = expressions.Session({x._expr: a, y._expr: b}) self.assertTrue(sub(x, y)._expr.evaluate_at(session).equals(a - b)) self.assertTrue(sub(x, 1)._expr.evaluate_at(session).equals(a - 1)) self.assertTrue(sub(1, x)._expr.evaluate_at(session).equals(1 - a)) self.assertTrue(sub(x, b)._expr.evaluate_at(session).equals(a - b)) self.assertTrue(sub(a, y)._expr.evaluate_at(session).equals(a - b)) def test_elementwise_func_kwarg(self): a = pd.Series([1, 2, 3]) b = pd.Series([100, 200, 300]) empty_proxy = a[:0] x = frames.DeferredSeries(expressions.PlaceholderExpression(empty_proxy)) y = frames.DeferredSeries(expressions.PlaceholderExpression(empty_proxy)) sub = frame_base._elementwise_function(lambda x, y=1: x - y) session = expressions.Session({x._expr: a, y._expr: b}) self.assertTrue(sub(x, y=y)._expr.evaluate_at(session).equals(a - b)) self.assertTrue(sub(x)._expr.evaluate_at(session).equals(a - 1)) self.assertTrue(sub(1, y=x)._expr.evaluate_at(session).equals(1 - a)) self.assertTrue(sub(x, y=b)._expr.evaluate_at(session).equals(a - b)) self.assertTrue(sub(a, y=y)._expr.evaluate_at(session).equals(a - b)) self.assertTrue(sub(x, y)._expr.evaluate_at(session).equals(a - b)) def test_maybe_inplace(self): @frame_base.maybe_inplace def add_one(frame): return frame + 1 frames.DeferredSeries.add_one = add_one original_expr = expressions.PlaceholderExpression(pd.Series([1, 2, 3])) x = frames.DeferredSeries(original_expr) x.add_one() self.assertIs(x._expr, original_expr) x.add_one(inplace=False) self.assertIs(x._expr, original_expr) x.add_one(inplace=True) self.assertIsNot(x._expr, original_expr) def test_args_to_kwargs(self): class Base(object): def func(self, a=1, b=2, c=3, *, kw_only=4): pass class Proxy(object): @frame_base.args_to_kwargs(Base) def func(self, **kwargs): return kwargs proxy = Proxy() # pylint: disable=too-many-function-args self.assertEqual(proxy.func(), {}) self.assertEqual(proxy.func(100), {'a': 100}) self.assertEqual(proxy.func(2, 4, 6), {'a': 2, 'b': 4, 'c': 6}) self.assertEqual(proxy.func(2, c=6), {'a': 2, 'c': 6}) self.assertEqual(proxy.func(c=6, a=2), {'a': 2, 'c': 6}) self.assertEqual(proxy.func(2, kw_only=20), {'a': 2, 'kw_only': 20}) with self.assertRaises(TypeError): # got too many positioned arguments proxy.func(2, 4, 6, 8) def test_args_to_kwargs_populates_defaults(self): class Base(object): def func(self, a=1, b=2, c=3): pass def func_removed_args(self, a): pass class Proxy(object): @frame_base.args_to_kwargs(Base) @frame_base.populate_defaults(Base) def func(self, a, c=1000, **kwargs): return dict(kwargs, a=a, c=c) @frame_base.args_to_kwargs(Base, removed_method=True) @frame_base.populate_defaults(Base, removed_method=True) def func_removed_method(self, a, **kwargs): return dict(kwargs, a=a) @frame_base.args_to_kwargs(Base, removed_args=['c']) @frame_base.populate_defaults(Base, removed_args=['c']) def func_removed_args(self, a, c, **kwargs): return dict(kwargs, a=a) proxy = Proxy() # pylint: disable=too-many-function-args,no-value-for-parameter self.assertEqual(proxy.func(), {'a': 1, 'c': 1000}) self.assertEqual(proxy.func(100), {'a': 100, 'c': 1000}) self.assertEqual(proxy.func(2, 4, 6), {'a': 2, 'b': 4, 'c': 6}) self.assertEqual(proxy.func(2, c=6), {'a': 2, 'c': 6}) self.assertEqual(proxy.func(c=6, a=2), {'a': 2, 'c': 6}) self.assertEqual(proxy.func(c=6), {'a': 1, 'c': 6}) with self.assertRaises(TypeError): # missing 1 required positional argument proxy.func_removed_method() self.assertEqual(proxy.func_removed_method(12, c=100), {'a': 12, 'c': 100}) with self.assertRaises(TypeError): # missing 1 required positional argument proxy.func_removed_args() self.assertEqual(proxy.func_removed_args(12, d=100), {'a': 12, 'd': 100}) def test_args_to_kwargs_populates_default_handles_kw_only(self): class Base(object): def func(self, a, b=2, c=3, *, kw_only=4): pass class ProxyUsesKwOnly(object): @frame_base.args_to_kwargs(Base) @frame_base.populate_defaults(Base) def func(self, a, kw_only, **kwargs): return dict(kwargs, a=a, kw_only=kw_only) proxy = ProxyUsesKwOnly() # pylint: disable=too-many-function-args,no-value-for-parameter with self.assertRaises(TypeError): # missing 1 required positional argument proxy.func() self.assertEqual(proxy.func(100), {'a': 100, 'kw_only': 4}) self.assertEqual( proxy.func(2, 4, 6, kw_only=8), { 'a': 2, 'b': 4, 'c': 6, 'kw_only': 8 }) with self.assertRaises(TypeError): proxy.func(2, 4, 6, 8) # got too many positioned arguments class ProxyDoesntUseKwOnly(object): @frame_base.args_to_kwargs(Base) @frame_base.populate_defaults(Base) def func(self, a, **kwargs): return dict(kwargs, a=a) proxy = ProxyDoesntUseKwOnly() # pylint: disable=too-many-function-args,no-value-for-parameter with self.assertRaises(TypeError): # missing 1 required positional argument proxy.func() self.assertEqual(proxy.func(100), {'a': 100}) self.assertEqual( proxy.func(2, 4, 6, kw_only=8), { 'a': 2, 'b': 4, 'c': 6, 'kw_only': 8 }) def test_populate_defaults_overwrites_copy(self): class Base(object): def func(self, a=1, b=2, c=3, *, copy=None): pass class Proxy(object): @frame_base.args_to_kwargs(Base) @frame_base.populate_defaults(Base) def func(self, a, copy, **kwargs): return dict(kwargs, a=a, copy=copy) proxy = Proxy() self.assertEqual(proxy.func(), {'a': 1, 'copy': True}) self.assertEqual(proxy.func(copy=False), {'a': 1, 'copy': False}) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/frames.py0000644000175100017510000065447415206113032021454 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Analogs for :class:`pandas.DataFrame` and :class:`pandas.Series`: :class:`DeferredDataFrame` and :class:`DeferredSeries`. These classes are effectively wrappers around a `schema-aware`_ :class:`~apache_beam.pvalue.PCollection` that provide a set of operations compatible with the `pandas`_ API. Note that we aim for the Beam DataFrame API to be completely compatible with the pandas API, but there are some features that are currently unimplemented for various reasons. Pay particular attention to the **'Differences from pandas'** section for each operation to understand where we diverge. .. _schema-aware: https://beam.apache.org/documentation/programming-guide/#what-is-a-schema .. _pandas: https://pandas.pydata.org/ """ import collections import inspect import itertools import math import re import warnings from typing import Optional import numpy as np import pandas as pd from pandas._libs import lib from pandas.api.types import is_float_dtype from pandas.api.types import is_int64_dtype from pandas.api.types import is_list_like from pandas.core.groupby.generic import DataFrameGroupBy from apache_beam.dataframe import convert from apache_beam.dataframe import expressions from apache_beam.dataframe import frame_base from apache_beam.dataframe import io from apache_beam.dataframe import partitionings from apache_beam.transforms import PTransform __all__ = [ 'DeferredSeries', 'DeferredDataFrame', ] # Get major, minor version PD_VERSION = tuple(map(int, pd.__version__.split('.')[0:2])) def populate_not_implemented(pd_type): def wrapper(deferred_type): for attr in dir(pd_type): # Don't auto-define hidden methods or dunders if attr.startswith('_'): continue if not hasattr(deferred_type, attr): pd_value = getattr(pd_type, attr) if isinstance(pd_value, property) or inspect.isclass(pd_value): # Some of the properties on pandas types (cat, dt, sparse), are # actually attributes with class values, not properties setattr( deferred_type, attr, property( frame_base.not_implemented_method(attr, base_type=pd_type))) elif callable(pd_value): setattr( deferred_type, attr, frame_base.not_implemented_method(attr, base_type=pd_type)) return deferred_type return wrapper def _fillna_alias(method): def wrapper(self, *args, **kwargs): return self.fillna(*args, method=method, **kwargs) wrapper.__name__ = method wrapper.__doc__ = ( f'{method} is only supported for axis="columns". ' 'axis="index" is order-sensitive.') return frame_base.with_docs_from(pd.DataFrame)( frame_base.args_to_kwargs(pd.DataFrame)( frame_base.populate_defaults(pd.DataFrame)(wrapper))) # These aggregations are commutative and associative, they can be trivially # "lifted" (i.e. we can pre-aggregate on partitions, group, then post-aggregate) LIFTABLE_AGGREGATIONS = ['all', 'any', 'max', 'min', 'prod', 'sum'] # These aggregations can be lifted if post-aggregated with "sum" LIFTABLE_WITH_SUM_AGGREGATIONS = ['size', 'count'] UNLIFTABLE_AGGREGATIONS = [ 'mean', 'median', 'quantile', 'describe', 'sem', 'skew', 'kurt', 'kurtosis', 'std', 'var', 'corr', 'cov', 'nunique', ] # mad was removed in Pandas 2.0. if PD_VERSION < (2, 0): UNLIFTABLE_AGGREGATIONS.append('mad') ALL_AGGREGATIONS = ( LIFTABLE_AGGREGATIONS + LIFTABLE_WITH_SUM_AGGREGATIONS + UNLIFTABLE_AGGREGATIONS) # These aggregations have specialized distributed implementations on # DeferredSeries, which are re-used in DeferredFrame. Note they are *not* used # for grouped aggregations, since they generally require tracking multiple # intermediate series, which is difficult to lift in groupby. HAND_IMPLEMENTED_GLOBAL_AGGREGATIONS = { 'quantile', 'std', 'var', 'mean', 'nunique', 'corr', 'cov', 'skew', 'kurt', 'kurtosis' } UNLIFTABLE_GLOBAL_AGGREGATIONS = ( set(UNLIFTABLE_AGGREGATIONS) - set(HAND_IMPLEMENTED_GLOBAL_AGGREGATIONS)) def _agg_method(base, func): def wrapper(self, *args, **kwargs): return self.agg(func, *args, **kwargs) if func in UNLIFTABLE_GLOBAL_AGGREGATIONS: wrapper.__doc__ = ( f"``{func}`` cannot currently be parallelized. It will " "require collecting all data on a single node.") wrapper.__name__ = func return frame_base.with_docs_from(base)(wrapper) # Docstring to use for head and tail (commonly used to peek at datasets) _PEEK_METHOD_EXPLANATION = ( "because it is `order-sensitive " "`_.\n\n" "If you want to peek at a large dataset consider using interactive Beam's " ":func:`ib.collect " "` " "with ``n`` specified, or :meth:`sample`. If you want to find the " "N largest elements, consider using :meth:`DeferredDataFrame.nlargest`.") class DeferredDataFrameOrSeries(frame_base.DeferredFrame): def _render_indexes(self): if self.index.nlevels == 1: return 'index=' + ( '' if self.index.name is None else repr(self.index.name)) else: return 'indexes=[' + ', '.join( '' if ix is None else repr(ix) for ix in self.index.names) + ']' __array__ = frame_base.wont_implement_method( pd.Series, '__array__', reason="non-deferred-result") @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) @frame_base.maybe_inplace def drop(self, labels, axis, index, columns, errors, **kwargs): """drop is not parallelizable when dropping from the index and ``errors="raise"`` is specified. It requires collecting all data on a single node in order to detect if one of the index values is missing.""" if labels is not None: if index is not None or columns is not None: raise ValueError("Cannot specify both 'labels' and 'index'/'columns'") if axis in (0, 'index'): index = labels columns = None elif axis in (1, 'columns'): index = None columns = labels else: raise ValueError( "axis must be one of (0, 1, 'index', 'columns'), " "got '%s'" % axis) if columns is not None: # Compute the proxy based on just the columns that are dropped. proxy = self._expr.proxy().drop(columns=columns, errors=errors) else: proxy = self._expr.proxy() if index is not None and errors == 'raise': # In order to raise an error about missing index values, we'll # need to collect the entire dataframe. # TODO: This could be parallelized by putting index values in a # ConstantExpression and partitioning by index. requires = partitionings.Singleton( reason=( "drop(errors='raise', axis='index') is not currently " "parallelizable. This requires collecting all data on a single " f"node in order to detect if one of {index!r} is missing.")) else: requires = partitionings.Arbitrary() return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'drop', lambda df: df.drop( axis=axis, index=index, columns=columns, errors=errors, **kwargs ), [self._expr], proxy=proxy, requires_partition_by=requires)) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def droplevel(self, level, axis): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'droplevel', lambda df: df.droplevel(level, axis=axis), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary() if axis in (1, 'column') else partitionings.Singleton())) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) def swaplevel(self, **kwargs): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'swaplevel', lambda df: df.swaplevel(**kwargs), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary())) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) @frame_base.maybe_inplace def fillna(self, value, method, axis, limit, **kwargs): """When ``axis="index"``, both ``method`` and ``limit`` must be ``None``. otherwise this operation is order-sensitive.""" # Default value is None, but is overriden with index. axis = axis or 'index' if axis in (0, 'index'): if method is not None: raise frame_base.WontImplementError( f"fillna(method={method!r}, axis={axis!r}) is not supported " "because it is order-sensitive. Only fillna(method=None) is " f"supported with axis={axis!r}.", reason="order-sensitive") if limit is not None: raise frame_base.WontImplementError( f"fillna(limit={method!r}, axis={axis!r}) is not supported because " "it is order-sensitive. Only fillna(limit=None) is supported with " f"axis={axis!r}.", reason="order-sensitive") if isinstance(self, DeferredDataFrame) and isinstance(value, DeferredSeries): # If self is a DataFrame and value is a Series we want to broadcast value # to all partitions of self. # This is OK, as its index must be the same size as the columns set of # self, so cannot be too large. class AsScalar(object): def __init__(self, value): self.value = value with expressions.allow_non_parallel_operations(): value_expr = expressions.ComputedExpression( 'as_scalar', lambda df: AsScalar(df), [value._expr], requires_partition_by=partitionings.Singleton()) get_value = lambda x: x.value requires = partitionings.Arbitrary() elif isinstance(value, frame_base.DeferredBase): # For other DeferredBase combinations, use Index partitioning to # co-locate on the Index value_expr = value._expr get_value = lambda x: x requires = partitionings.Index() else: # Default case, pass value through as a constant, no particular # partitioning requirement value_expr = expressions.ConstantExpression(value) get_value = lambda x: x requires = partitionings.Arbitrary() return frame_base.DeferredFrame.wrap( # yapf: disable expressions.ComputedExpression( 'fillna', lambda df, value: df.fillna( get_value(value), method=method, axis=axis, limit=limit, ** kwargs), [self._expr, value_expr], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=requires)) if hasattr(pd.DataFrame, 'ffill'): ffill = _fillna_alias('ffill') if hasattr(pd.DataFrame, 'bfill'): bfill = _fillna_alias('bfill') if hasattr(pd.DataFrame, 'backfill'): backfill = _fillna_alias('backfill') if hasattr(pd.DataFrame, 'pad'): pad = _fillna_alias('pad') @frame_base.with_docs_from(pd.DataFrame) def first(self, offset): per_partition = expressions.ComputedExpression( 'first-per-partition', lambda df: df.sort_index().first(offset=offset), [self._expr], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Arbitrary()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'first', lambda df: df.sort_index().first(offset=offset), [per_partition], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.DataFrame) def last(self, offset): per_partition = expressions.ComputedExpression( 'last-per-partition', lambda df: df.sort_index().last(offset=offset), [self._expr], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Arbitrary()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'last', lambda df: df.sort_index().last(offset=offset), [per_partition], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def groupby(self, by, level, axis, as_index, group_keys, **kwargs): """``as_index`` must be ``True``. Aggregations grouping by a categorical column with ``observed=False`` set are not currently parallelizable (`Issue 21827 `_). """ if not as_index: raise NotImplementedError('groupby(as_index=False)') if axis in (1, 'columns'): return _DeferredGroupByCols( expressions.ComputedExpression( 'groupbycols', lambda df: df.groupby( by, axis=axis, group_keys=group_keys, **kwargs), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary()), group_keys=group_keys) if level is None and by is None: raise TypeError("You have to supply one of 'by' and 'level'") elif level is not None: if isinstance(level, (list, tuple)): grouping_indexes = level else: grouping_indexes = [level] grouping_columns = [] index = self._expr.proxy().index # Translate to level numbers only grouping_indexes = [ l if isinstance(l, int) else index.names.index(l) for l in grouping_indexes ] if index.nlevels == 1: to_group_with_index = self._expr to_group = self._expr else: levels_to_drop = [ i for i in range(index.nlevels) if i not in grouping_indexes ] # Reorder so the grouped indexes are first to_group_with_index = self.reorder_levels( grouping_indexes + levels_to_drop) grouping_indexes = list(range(len(grouping_indexes))) levels_to_drop = list(range(len(grouping_indexes), index.nlevels)) if levels_to_drop: to_group = to_group_with_index.droplevel(levels_to_drop)._expr else: to_group = to_group_with_index._expr to_group_with_index = to_group_with_index._expr elif callable(by): def map_index(df): df = df.copy() df.index = df.index.map(by) return df to_group = expressions.ComputedExpression( 'map_index', map_index, [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) orig_nlevels = self._expr.proxy().index.nlevels def prepend_mapped_index(df): df = df.copy() index = df.index.to_frame() index.insert(0, None, df.index.map(by)) df.index = pd.MultiIndex.from_frame( index, names=[None] + list(df.index.names)) return df to_group_with_index = expressions.ComputedExpression( 'map_index_keep_orig', prepend_mapped_index, [self._expr], requires_partition_by=partitionings.Arbitrary(), # Partitioning by the original indexes is preserved preserves_partition_by=partitionings.Index( list(range(1, orig_nlevels + 1)))) grouping_columns = [] # The index we need to group by is the last one grouping_indexes = [0] elif isinstance(by, DeferredSeries): if isinstance(self, DeferredSeries): def set_index(s, by): df = pd.DataFrame(s) df, by = df.align(by, axis=0, join='inner') return df.set_index(by).iloc[:, 0] def prepend_index(s, by): df = pd.DataFrame(s) df, by = df.align(by, axis=0, join='inner') return df.set_index([by, df.index]).iloc[:, 0] else: def set_index(df, by): # type: ignore df, by = df.align(by, axis=0, join='inner') return df.set_index(by) def prepend_index(df, by): # type: ignore df, by = df.align(by, axis=0, join='inner') return df.set_index([by, df.index]) to_group = expressions.ComputedExpression( 'set_index', set_index, [self._expr, by._expr], requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Singleton()) orig_nlevels = self._expr.proxy().index.nlevels to_group_with_index = expressions.ComputedExpression( 'prependindex', prepend_index, [self._expr, by._expr], requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Index( list(range(1, orig_nlevels + 1)))) grouping_columns = [] grouping_indexes = [0] elif isinstance(by, np.ndarray): raise frame_base.WontImplementError( "Grouping by a concrete ndarray is order sensitive.", reason="order-sensitive") elif isinstance(self, DeferredDataFrame): if not isinstance(by, list): by = [by] # Find the columns that we need to move into the index so we can group by # them column_names = self._expr.proxy().columns grouping_columns = list(set(by).intersection(column_names)) index_names = self._expr.proxy().index.names for label in by: if label not in index_names and label not in self._expr.proxy().columns: raise KeyError(label) grouping_indexes = list(set(by).intersection(index_names)) if grouping_indexes: if set(by) == set(index_names): to_group = self._expr elif set(by).issubset(index_names): to_group = self.droplevel(index_names.difference(by))._expr else: to_group = self.reset_index(grouping_indexes).set_index(by)._expr else: to_group = self.set_index(by)._expr if grouping_columns: # TODO(https://github.com/apache/beam/issues/20759): # It should be possible to do this without creating # an expression manually, by using DeferredDataFrame.set_index, i.e.: # to_group_with_index = self.set_index([self.index] + # grouping_columns)._expr to_group_with_index = expressions.ComputedExpression( 'move_grouped_columns_to_index', lambda df: df.set_index([df.index] + grouping_columns, drop=False), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Index( list(range(self._expr.proxy().index.nlevels)))) else: to_group_with_index = self._expr else: raise NotImplementedError(by) return DeferredGroupBy( expressions.ComputedExpression( 'groupbyindex', lambda df: df.groupby( level=list(range(df.index.nlevels)), group_keys=group_keys, ** kwargs), [to_group], requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Arbitrary()), kwargs, to_group, to_group_with_index, grouping_columns=grouping_columns, grouping_indexes=grouping_indexes, group_keys=group_keys) @property # type: ignore @frame_base.with_docs_from(pd.DataFrame) def loc(self): return _DeferredLoc(self) @property # type: ignore @frame_base.with_docs_from(pd.DataFrame) def iloc(self): """Position-based indexing with `iloc` is order-sensitive in almost every case. Beam DataFrame users should prefer label-based indexing with `loc`. """ return _DeferredILoc(self) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) @frame_base.maybe_inplace def reset_index(self, level=None, **kwargs): """Dropping the entire index (e.g. with ``reset_index(level=None)``) is not parallelizable. It is also only guaranteed that the newly generated index values will be unique. The Beam DataFrame API makes no guarantee that the same index values as the equivalent pandas operation will be generated, because that implementation is order-sensitive.""" if level is not None and not isinstance(level, (tuple, list)): level = [level] if level is None or len(level) == self._expr.proxy().index.nlevels: # TODO(https://github.com/apache/beam/issues/20859): # Could do distributed re-index with offsets. requires_partition_by = partitionings.Singleton( reason=( f"reset_index(level={level!r}) drops the entire index and " "creates a new one, so it cannot currently be parallelized " "(https://github.com/apache/beam/issues/20859).")) else: requires_partition_by = partitionings.Arbitrary() return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'reset_index', lambda df: df.reset_index(level=level, **kwargs), [self._expr], preserves_partition_by=partitionings.Singleton(), requires_partition_by=requires_partition_by)) abs = frame_base._elementwise_method('abs', base=pd.core.generic.NDFrame) @frame_base.with_docs_from(pd.core.generic.NDFrame) @frame_base.args_to_kwargs(pd.core.generic.NDFrame) @frame_base.populate_defaults(pd.core.generic.NDFrame) def astype(self, dtype, copy, errors): """astype is not parallelizable when ``errors="ignore"`` is specified. ``copy=False`` is not supported because it relies on memory-sharing semantics. ``dtype="category`` is not supported because the type of the output column depends on the data. Please use ``pd.CategoricalDtype`` with explicit categories instead. """ requires = partitionings.Arbitrary() if errors == "ignore": # We need all data in order to ignore errors and propagate the original # data. requires = partitionings.Singleton( reason=( f"astype(errors={errors!r}) is currently not parallelizable, " "because all data must be collected on one node to determine if " "the original data should be propagated instead.")) if not copy: raise frame_base.WontImplementError( f"astype(copy={copy!r}) is not supported because it relies on " "memory-sharing semantics that are not compatible with the Beam " "model.") # An instance of CategoricalDtype is actualy considered equal to the string # 'category', so we have to explicitly check if dtype is an instance of # CategoricalDtype, and allow it. # See https://github.com/apache/beam/issues/23276 if dtype == 'category' and not isinstance(dtype, pd.CategoricalDtype): raise frame_base.WontImplementError( "astype(dtype='category') is not supported because the type of the " "output column depends on the data. Please use pd.CategoricalDtype " "with explicit categories instead.", reason="non-deferred-columns") return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'astype', lambda df: df.astype(dtype=dtype, copy=copy, errors=errors), [self._expr], requires_partition_by=requires, preserves_partition_by=partitionings.Arbitrary())) at_time = frame_base._elementwise_method( 'at_time', base=pd.core.generic.NDFrame) between_time = frame_base._elementwise_method( 'between_time', base=pd.core.generic.NDFrame) copy = frame_base._elementwise_method('copy', base=pd.core.generic.NDFrame) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) @frame_base.maybe_inplace def replace(self, to_replace, value, limit, method, **kwargs): """``method`` is not supported in the Beam DataFrame API because it is order-sensitive. It cannot be specified. If ``limit`` is specified this operation is not parallelizable.""" # pylint: disable-next=c-extension-no-member value_compare = None if PD_VERSION < (1, 4) else lib.no_default if method is not None and not isinstance(to_replace, dict) and value is value_compare: # pandas only relies on method if to_replace is not a dictionary, and # value is the value. This is different than # if ``None`` is explicitly passed for ``value``. In this case, it will be # respected raise frame_base.WontImplementError( f"replace(method={method!r}) is not supported because it is " "order sensitive. Only replace(method=None) is supported.", reason="order-sensitive") if limit is None: requires_partition_by = partitionings.Arbitrary() else: requires_partition_by = partitionings.Singleton( reason=( f"replace(limit={limit!r}) cannot currently be parallelized. It " "requires collecting all data on a single node.")) return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'replace', lambda df: df.replace( to_replace=to_replace, value=value, limit=limit, method=method, **kwargs), [self._expr], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=requires_partition_by)) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def tz_localize(self, ambiguous, **kwargs): """``ambiguous`` cannot be set to ``"infer"`` as its semantics are order-sensitive. Similarly, specifying ``ambiguous`` as an :class:`~numpy.ndarray` is order-sensitive, but you can achieve similar functionality by specifying ``ambiguous`` as a Series.""" if isinstance(ambiguous, np.ndarray): raise frame_base.WontImplementError( "tz_localize(ambiguous=ndarray) is not supported because it makes " "this operation sensitive to the order of the data. Please use a " "DeferredSeries instead.", reason="order-sensitive") elif isinstance(ambiguous, frame_base.DeferredFrame): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'tz_localize', lambda df, ambiguous: df.tz_localize( ambiguous=ambiguous, **kwargs), [self._expr, ambiguous._expr], requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Singleton())) elif ambiguous == 'infer': # infer attempts to infer based on the order of the timestamps raise frame_base.WontImplementError( f"tz_localize(ambiguous={ambiguous!r}) is not allowed because it " "makes this operation sensitive to the order of the data.", reason="order-sensitive") return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'tz_localize', lambda df: df.tz_localize(ambiguous=ambiguous, **kwargs), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton())) @property # type: ignore @frame_base.with_docs_from(pd.DataFrame) def size(self): sizes = expressions.ComputedExpression( 'get_sizes', # Wrap scalar results in a Series for easier concatenation later lambda df: pd.Series(df.size), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'sum_sizes', lambda sizes: sizes.sum(), [sizes], requires_partition_by=partitionings.Singleton(), preserves_partition_by=partitionings.Singleton())) def length(self): """Alternative to ``len(df)`` which returns a deferred result that can be used in arithmetic with :class:`DeferredSeries` or :class:`DeferredDataFrame` instances.""" lengths = expressions.ComputedExpression( 'get_lengths', # Wrap scalar results in a Series for easier concatenation later lambda df: pd.Series(len(df)), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'sum_lengths', lambda lengths: lengths.sum(), [lengths], requires_partition_by=partitionings.Singleton(), preserves_partition_by=partitionings.Singleton())) def __len__(self): raise frame_base.WontImplementError( "len(df) is not currently supported because it produces a non-deferred " "result. Consider using df.length() instead.", reason="non-deferred-result") @property # type: ignore @frame_base.with_docs_from(pd.DataFrame) def empty(self): empties = expressions.ComputedExpression( 'get_empties', # Wrap scalar results in a Series for easier concatenation later lambda df: pd.Series(df.empty), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'check_all_empty', lambda empties: empties.all(), [empties], requires_partition_by=partitionings.Singleton(), preserves_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.DataFrame) def bool(self): # TODO: Documentation about DeferredScalar # Will throw if any partition has >1 element bools = expressions.ComputedExpression( 'get_bools', # Wrap scalar results in a Series for easier concatenation later lambda df: pd.Series([], dtype=bool) if df.empty else pd.Series([df.bool()]), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) with expressions.allow_non_parallel_operations(True): # Will throw if overall dataset has != 1 element return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'combine_all_bools', lambda bools: bools.bool(), [bools], proxy=bool(), requires_partition_by=partitionings.Singleton(), preserves_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.DataFrame) def equals(self, other): intermediate = expressions.ComputedExpression( 'equals_partitioned', # Wrap scalar results in a Series for easier concatenation later lambda df, other: pd.Series(df.equals(other)), [self._expr, other._expr], requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Singleton()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'aggregate_equals', lambda df: df.all(), [intermediate], requires_partition_by=partitionings.Singleton(), preserves_partition_by=partitionings.Singleton())) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def sort_values(self, axis, **kwargs): """``sort_values`` is not implemented. It is not implemented for ``axis=index`` because it imposes an ordering on the dataset, and it likely will not be maintained (see https://s.apache.org/dataframe-order-sensitive-operations). It is not implemented for ``axis=columns`` because it makes the order of the columns depend on the data (see https://s.apache.org/dataframe-non-deferred-columns).""" if axis in (0, 'index'): # axis=index imposes an ordering on the DataFrame rows which we do not # support raise frame_base.WontImplementError( "sort_values(axis=index) is not supported because it imposes an " "ordering on the dataset which likely will not be preserved.", reason="order-sensitive") else: # axis=columns will reorder the columns based on the data raise frame_base.WontImplementError( "sort_values(axis=columns) is not supported because the order of the " "columns in the result depends on the data.", reason="non-deferred-columns") @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) @frame_base.maybe_inplace def sort_index(self, axis, **kwargs): """``axis=index`` is not allowed because it imposes an ordering on the dataset, and we cannot guarantee it will be maintained (see https://s.apache.org/dataframe-order-sensitive-operations). Only ``axis=columns`` is allowed.""" if axis in (0, 'index'): # axis=rows imposes an ordering on the DataFrame which we do not support raise frame_base.WontImplementError( "sort_index(axis=index) is not supported because it imposes an " "ordering on the dataset which we cannot guarantee will be " "preserved.", reason="order-sensitive") # axis=columns reorders the columns by name return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'sort_index', lambda df: df.sort_index(axis=axis, **kwargs), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary(), )) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs( pd.DataFrame, removed_args=["errors"] if PD_VERSION >= (2, 0) else None) @frame_base.populate_defaults( pd.DataFrame, removed_args=["errors"] if PD_VERSION >= (2, 0) else None) @frame_base.maybe_inplace def where(self, cond, other, errors, **kwargs): """where is not parallelizable when ``errors="ignore"`` is specified.""" requires = partitionings.Arbitrary() deferred_args = {} actual_args = {} # TODO(bhulette): This is very similar to the logic in # frame_base.elementwise_method, can we unify it? if isinstance(cond, frame_base.DeferredFrame): deferred_args['cond'] = cond requires = partitionings.Index() else: actual_args['cond'] = cond if isinstance(other, frame_base.DeferredFrame): deferred_args['other'] = other requires = partitionings.Index() else: actual_args['other'] = other # For Pandas 2.0, errors was removed as an argument. if PD_VERSION < (2, 0): if "errors" in kwargs and kwargs['errors'] == "ignore": # We need all data in order to ignore errors and propagate the original # data. requires = partitionings.Singleton( reason=( f"where(errors={kwargs['errors']!r}) is currently not " "parallelizable, because all data must be collected on one " "node to determine if the original data should be propagated " "instead.")) actual_args['errors'] = kwargs['errors'] if 'errors' in kwargs else None def where_execution(df, *args): runtime_values = { name: value for (name, value) in zip(deferred_args.keys(), args) } return df.where(**runtime_values, **actual_args, **kwargs) return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( "where", where_execution, [self._expr] + [df._expr for df in deferred_args.values()], requires_partition_by=requires, preserves_partition_by=partitionings.Index(), )) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) @frame_base.maybe_inplace def mask(self, cond, **kwargs): """mask is not parallelizable when ``errors="ignore"`` is specified.""" return self.where(~cond, **kwargs) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def truncate(self, before, after, axis): if axis in (None, 0, 'index'): def truncate(df): return df.sort_index().truncate(before=before, after=after, axis=axis) else: def truncate(df): return df.truncate(before=before, after=after, axis=axis) return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'truncate', truncate, [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary())) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def unstack(self, **kwargs): level = kwargs.get('level', -1) if self._expr.proxy().index.nlevels == 1: if PD_VERSION < (1, 2): raise frame_base.WontImplementError( "unstack() is not supported when using pandas < 1.2.0\n" "Please upgrade to pandas 1.2.0 or higher to use this operation.") return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'unstack', lambda s: s.unstack(**kwargs), [self._expr], requires_partition_by=partitionings.Index())) else: # Unstacking MultiIndex objects idx = self._expr.proxy().index # Converting level (int, str, or combination) to a list of number levels level_list = level if isinstance(level, list) else [level] level_number_list = [idx._get_level_number(l) for l in level_list] # Checking if levels provided are of CategoricalDtype if not all(isinstance(idx.levels[l].dtype, (pd.CategoricalDtype, pd.BooleanDtype)) for l in level_number_list): raise frame_base.WontImplementError( "unstack() is only supported on DataFrames if unstacked level " "is a categorical or boolean column", reason="non-deferred-columns") else: tmp = self._expr.proxy().unstack(**kwargs) if isinstance(tmp.columns, pd.MultiIndex): levels = [] for i in range(tmp.columns.nlevels): level = tmp.columns.levels[i] levels.append(level) col_idx = pd.MultiIndex.from_product(levels) else: if tmp.columns.dtype == 'boolean': col_idx = pd.Index([False, True], dtype='boolean') else: col_idx = pd.CategoricalIndex(tmp.columns.categories) if isinstance(self._expr.proxy(), pd.Series): proxy_dtype = self._expr.proxy().dtypes else: # Set dtype to object if more than one value dtypes = [d for d in self._expr.proxy().dtypes] proxy_dtype = object if np.int64 in dtypes: proxy_dtype = np.int64 if np.float64 in dtypes: proxy_dtype = np.float64 if object in dtypes: proxy_dtype = object proxy = pd.DataFrame( columns=col_idx, dtype=proxy_dtype, index=tmp.index) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'unstack', lambda s: pd.concat([proxy, s.unstack(**kwargs)]), [self._expr], proxy=proxy, requires_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def xs(self, key, axis, level, **kwargs): """Note that ``xs(axis='index')`` will raise a ``KeyError`` at execution time if the key does not exist in the index.""" if axis in ('columns', 1): # Special case for axis=columns. This is a simple project that raises a # KeyError at construction time for missing columns. return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'xs', lambda df: df.xs(key, axis=axis, **kwargs), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary())) elif axis not in ('index', 0): # Make sure that user's axis is valid raise ValueError( "axis must be one of ('index', 0, 'columns', 1). " f"got {axis!r}.") if not isinstance(key, tuple): key_size = 1 key_series = pd.Series([key], index=[key]) else: key_size = len(key) key_series = pd.Series([key], pd.MultiIndex.from_tuples([key])) key_expr = expressions.ConstantExpression( key_series, proxy=key_series.iloc[:0]) if level is None: reindexed = self else: if not isinstance(level, list): level = [level] # If user specifed levels, reindex so those levels are at the beginning. # Keep the others and preserve their order. level = [ l if isinstance(l, int) else list(self.index.names).index(l) for l in level ] reindexed = self.reorder_levels( level + [i for i in range(self.index.nlevels) if i not in level]) def xs_partitioned(frame, key): if not len(key): # key is not in this partition, return empty dataframe result = frame.iloc[:0] if key_size < frame.index.nlevels: return result.droplevel(list(range(key_size))) else: return result # key should be in this partition, call xs. Will raise KeyError if not # present. return frame.xs(key.item()) return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'xs', xs_partitioned, [reindexed._expr, key_expr], requires_partition_by=partitionings.Index(list(range(key_size))), # Drops index levels, so partitioning is not preserved preserves_partition_by=partitionings.Singleton())) @property def dtype(self): return self._expr.proxy().dtype isin = frame_base._elementwise_method('isin', base=pd.DataFrame) combine_first = frame_base._elementwise_method( 'combine_first', base=pd.DataFrame) combine = frame_base._proxy_method( 'combine', base=pd.DataFrame, requires_partition_by=expressions.partitionings.Singleton( reason="combine() is not parallelizable because func might operate " "on the full dataset."), preserves_partition_by=expressions.partitionings.Singleton()) @property # type: ignore @frame_base.with_docs_from(pd.DataFrame) def ndim(self): return self._expr.proxy().ndim @property # type: ignore @frame_base.with_docs_from(pd.DataFrame) def index(self): return _DeferredIndex(self) @index.setter def _set_index(self, value): # TODO: assigning the index is generally order-sensitive, but we could # support it in some rare cases, e.g. when assigning the index from one # of a DataFrame's columns raise NotImplementedError( "Assigning an index is not yet supported. " "Consider using set_index() instead.") reindex = frame_base.wont_implement_method( pd.DataFrame, 'reindex', reason="order-sensitive") hist = frame_base.wont_implement_method( pd.DataFrame, 'hist', reason="plotting-tools") attrs = property( fget=frame_base.wont_implement_method( pd.DataFrame, 'attrs', reason='experimental'), fset=frame_base.wont_implement_method( pd.DataFrame, 'attrs', reason='experimental'), ) reorder_levels = frame_base._proxy_method( 'reorder_levels', base=pd.DataFrame, requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) resample = frame_base.wont_implement_method( pd.DataFrame, 'resample', reason='event-time-semantics') rolling = frame_base.wont_implement_method( pd.DataFrame, 'rolling', reason='event-time-semantics') to_xarray = frame_base.wont_implement_method( pd.DataFrame, 'to_xarray', reason='non-deferred-result') to_clipboard = frame_base.wont_implement_method( pd.DataFrame, 'to_clipboard', reason="non-deferred-result") swapaxes = frame_base.wont_implement_method( pd.Series, 'swapaxes', reason="non-deferred-columns") infer_object = frame_base.wont_implement_method( pd.Series, 'infer_objects', reason="non-deferred-columns") ewm = frame_base.wont_implement_method( pd.Series, 'ewm', reason="event-time-semantics") expanding = frame_base.wont_implement_method( pd.Series, 'expanding', reason="event-time-semantics") sparse = property( frame_base.not_implemented_method( 'sparse', '20902', base_type=pd.DataFrame)) transform = frame_base._elementwise_method('transform', base=pd.DataFrame) tz_convert = frame_base._proxy_method( 'tz_convert', base=pd.DataFrame, requires_partition_by=partitionings.Arbitrary(), # Manipulates index, partitioning is not preserved preserves_partition_by=partitionings.Singleton()) @frame_base.with_docs_from(pd.DataFrame) def pipe(self, func, *args, **kwargs): if isinstance(func, tuple): func, data = func kwargs[data] = self return func(*args, **kwargs) return func(self, *args, **kwargs) @populate_not_implemented(pd.Series) @frame_base.DeferredFrame._register_for(pd.Series) class DeferredSeries(DeferredDataFrameOrSeries): def __repr__(self): return ( f'DeferredSeries(name={self.name!r}, dtype={self.dtype}, ' f'{self._render_indexes()})') @property # type: ignore @frame_base.with_docs_from(pd.Series) def name(self): return self._expr.proxy().name @name.setter def name(self, value): def fn(s): s = s.copy() s.name = value return s self._expr = expressions.ComputedExpression( 'series_set_name', fn, [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary()) @property # type: ignore @frame_base.with_docs_from(pd.Series) def hasnans(self): has_nans = expressions.ComputedExpression( 'hasnans', lambda s: pd.Series(s.hasnans), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) with expressions.allow_non_parallel_operations(): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'combine_hasnans', lambda s: s.any(), [has_nans], requires_partition_by=partitionings.Singleton(), preserves_partition_by=partitionings.Singleton())) @property # type: ignore @frame_base.with_docs_from(pd.Series) def dtype(self): return self._expr.proxy().dtype dtypes = dtype def __getitem__(self, key): if _is_null_slice(key) or key is Ellipsis: return self elif (isinstance(key, int) or _is_integer_slice(key) ) and self._expr.proxy().index._should_fallback_to_positional(): raise frame_base.WontImplementError( "Accessing an item by an integer key is order sensitive for this " "Series.", reason="order-sensitive") elif isinstance(key, slice) or callable(key): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( # yapf: disable 'getitem', lambda df: df[key], [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary())) elif isinstance(key, DeferredSeries) and key._expr.proxy().dtype == bool: return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( # yapf: disable 'getitem', lambda df, indexer: df[indexer], [self._expr, key._expr], requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Arbitrary())) elif pd.core.series.is_iterator(key) or pd.core.common.is_bool_indexer(key): raise frame_base.WontImplementError( "Accessing a DeferredSeries with an iterator is sensitive to the " "order of the data.", reason="order-sensitive") else: # We could consider returning a deferred scalar, but that might # be more surprising than a clear error. raise frame_base.WontImplementError( f"Indexing a series with key of type {type(key)} is not supported " "because it produces a non-deferred result.", reason="non-deferred-result") @frame_base.with_docs_from(pd.Series) def keys(self): return self.index # Series.T == transpose. Both are a no-op T = frame_base._elementwise_method('T', base=pd.Series) transpose = frame_base._elementwise_method('transpose', base=pd.Series) shape = property( frame_base.wont_implement_method( pd.Series, 'shape', reason="non-deferred-result")) @frame_base.with_docs_from(pd.Series, removed_method=PD_VERSION >= (2, 0)) @frame_base.args_to_kwargs(pd.Series, removed_method=PD_VERSION >= (2, 0)) @frame_base.populate_defaults(pd.Series, removed_method=PD_VERSION >= (2, 0)) def append(self, to_append, ignore_index, verify_integrity, **kwargs): """``ignore_index=True`` is not supported, because it requires generating an order-sensitive index.""" if PD_VERSION >= (2, 0): raise frame_base.WontImplementError('append() was removed in Pandas 2.0.') if not isinstance(to_append, DeferredSeries): raise frame_base.WontImplementError( "append() only accepts DeferredSeries instances, received " + str(type(to_append))) if ignore_index: raise frame_base.WontImplementError( "append(ignore_index=True) is order sensitive because it requires " "generating a new index based on the order of the data.", reason="order-sensitive") if verify_integrity: # We can verify the index is non-unique within index partitioned data. requires = partitionings.Index() else: requires = partitionings.Arbitrary() return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'append', lambda s, to_append: s.append( to_append, verify_integrity=verify_integrity, **kwargs), [self._expr, to_append._expr], requires_partition_by=requires, preserves_partition_by=partitionings.Arbitrary())) @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) def align(self, other, join, axis, level, method, **kwargs): """Aligning per-level is not yet supported. Only the default, ``level=None``, is allowed. Filling NaN values via ``method`` is not supported, because it is `order-sensitive `_. Only the default, ``method=None``, is allowed.""" if level is not None: raise NotImplementedError('per-level align') if method is not None and method != lib.no_default: raise frame_base.WontImplementError( f"align(method={method!r}) is not supported because it is " "order sensitive. Only align(method=None) is supported.", reason="order-sensitive") # We're using pd.concat here as expressions don't yet support # multiple return values. aligned = frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'align', lambda x, y: pd.concat([x, y], axis=1, join='inner'), [self._expr, other._expr], requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Arbitrary())) return aligned.iloc[:, 0], aligned.iloc[:, 1] argsort = frame_base.wont_implement_method( pd.Series, 'argsort', reason="order-sensitive") array = property( frame_base.wont_implement_method( pd.Series, 'array', reason="non-deferred-result")) # We can't reliably predict the output type, it depends on whether `key` is: # - not in the index (default_value) # - in the index once (constant) # - in the index multiple times (Series) get = frame_base.wont_implement_method( pd.Series, 'get', reason="non-deferred-columns") ravel = frame_base.wont_implement_method( pd.Series, 'ravel', reason="non-deferred-result") slice_shift = frame_base.wont_implement_method( pd.Series, 'slice_shift', reason="deprecated") tshift = frame_base.wont_implement_method( pd.Series, 'tshift', reason="deprecated") rename = frame_base._proxy_method( 'rename', base=pd.Series, requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) between = frame_base._elementwise_method('between', base=pd.Series) add_suffix = frame_base._proxy_method( 'add_suffix', base=pd.DataFrame, requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) add_prefix = frame_base._proxy_method( 'add_prefix', base=pd.DataFrame, requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) info = frame_base.wont_implement_method( pd.Series, 'info', reason="non-deferred-result") def _idxmaxmin_helper(self, op, **kwargs): if op == 'idxmax': func = pd.Series.idxmax elif op == 'idxmin': func = pd.Series.idxmin else: raise ValueError( "op must be one of ('idxmax', 'idxmin'). " f"got {op!r}.") def compute_idx(s): index = func(s, **kwargs) if pd.isna(index): return s else: return s.loc[[index]] # Avoids empty Series error when evaluating proxy index_dtype = self._expr.proxy().index.dtype index = pd.Index([], dtype=index_dtype) proxy = self._expr.proxy().copy() proxy.index = index proxy = pd.concat([ proxy, pd.Series([1], index=np.asarray(['0']).astype(proxy.index.dtype)) ]) idx_func = expressions.ComputedExpression( 'idx_func', compute_idx, [self._expr], proxy=proxy, requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'idx_combine', lambda s: func(s, **kwargs), [idx_func], requires_partition_by=partitionings.Singleton(), preserves_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) def idxmin(self, **kwargs): return self._idxmaxmin_helper('idxmin', **kwargs) @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) def idxmax(self, **kwargs): return self._idxmaxmin_helper('idxmax', **kwargs) @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) def explode(self, ignore_index): # ignoring the index will not preserve it preserves = ( partitionings.Singleton() if ignore_index else partitionings.Index()) return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'explode', lambda s: s.explode(ignore_index), [self._expr], preserves_partition_by=preserves, requires_partition_by=partitionings.Arbitrary())) @frame_base.with_docs_from(pd.DataFrame) def dot(self, other): """``other`` must be a :class:`DeferredDataFrame` or :class:`DeferredSeries` instance. Computing the dot product with an array-like is not supported because it is order-sensitive.""" left = self._expr if isinstance(other, DeferredSeries): right = expressions.ComputedExpression( 'to_dataframe', pd.DataFrame, [other._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary()) right_is_series = True elif isinstance(other, DeferredDataFrame): right = other._expr right_is_series = False else: raise frame_base.WontImplementError( "other must be a DeferredDataFrame or DeferredSeries instance. " "Passing a concrete list or numpy array is not supported. Those " "types have no index and must be joined based on the order of the " "data.", reason="order-sensitive") dots = expressions.ComputedExpression( 'dot', # Transpose so we can sum across rows. (lambda left, right: pd.DataFrame(left @ right).T), [left, right], requires_partition_by=partitionings.Index()) with expressions.allow_non_parallel_operations(True): sums = expressions.ComputedExpression( 'sum', lambda dots: dots.sum(), # [dots], requires_partition_by=partitionings.Singleton()) if right_is_series: result = expressions.ComputedExpression( 'extract', lambda df: df[0], [sums], requires_partition_by=partitionings.Singleton()) else: result = sums return frame_base.DeferredFrame.wrap(result) __matmul__ = dot @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) def nunique(self, **kwargs): return self.drop_duplicates(keep="any").size @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) def quantile(self, q, **kwargs): """quantile is not parallelizable. See `Issue 20933 `_ tracking the possible addition of an approximate, parallelizable implementation of quantile.""" # TODO(https://github.com/apache/beam/issues/20933): Provide an option for # approximate distributed quantiles requires = partitionings.Singleton( reason=( "Computing quantiles across index cannot currently be " "parallelized. See https://github.com/apache/beam/issues/20933 " "tracking the possible addition of an approximate, parallelizable " "implementation of quantile.")) return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'quantile', lambda df: df.quantile(q=q, **kwargs), [self._expr], requires_partition_by=requires, preserves_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.Series) def std(self, *args, **kwargs): # Compute variance (deferred scalar) with same args, then sqrt it return self.var(*args, **kwargs).apply(lambda var: math.sqrt(var)) @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) def mean(self, skipna, **kwargs): if skipna: size = self.count() else: size = self.length() return self.sum(skipna=skipna, **kwargs) / size @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs( pd.Series, removed_args=["level"] if PD_VERSION >= (2, 0) else None) @frame_base.populate_defaults( pd.Series, removed_args=["level"] if PD_VERSION >= (2, 0) else None) def var(self, axis, skipna, level, ddof, **kwargs): if skipna is None or skipna: self = self.dropna() # pylint: disable=self-cls-assignment # See the online, numerically stable formulae at # https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm # and # https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm def compute_moments(x): n = len(x) m = x.std(ddof=0)**2 * n s = x.sum() return pd.DataFrame(dict(m=[m], s=[s], n=[n])) def combine_moments(data): m = s = n = 0.0 for datum in data.itertuples(): if datum.n == 0: continue elif n == 0: m, s, n = datum.m, datum.s, datum.n else: delta = s / n - datum.s / datum.n m += datum.m + delta**2 * n * datum.n / (n + datum.n) s += datum.s n += datum.n if n <= ddof: return float('nan') else: return m / (n - ddof) moments = expressions.ComputedExpression( 'compute_moments', compute_moments, [self._expr], requires_partition_by=partitionings.Arbitrary()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'combine_moments', combine_moments, [moments], requires_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) def corr(self, other, method, min_periods): """Only ``method='pearson'`` is currently parallelizable.""" if method == 'pearson': # Note that this is the default. x, y = self.dropna().align(other.dropna(), 'inner') return x._corr_aligned(y, min_periods) else: reason = ( f"Encountered corr(method={method!r}) which cannot be " "parallelized. Only corr(method='pearson') is currently " "parallelizable.") # The rank-based correlations are not obviously parallelizable, though # perhaps an approximation could be done with a knowledge of quantiles # and custom partitioning. return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'corr', lambda df, other: df.corr( other, method=method, min_periods=min_periods), [self._expr, other._expr], requires_partition_by=partitionings.Singleton(reason=reason))) @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs( pd.Series, removed_args=["level"] if PD_VERSION >= (2, 0) else None) @frame_base.populate_defaults( pd.Series, removed_args=["level"] if PD_VERSION >= (2, 0) else None) def skew(self, axis, skipna, level, numeric_only, **kwargs): if skipna is None or skipna: self = self.dropna() # pylint: disable=self-cls-assignment # See the online, numerically stable formulae at # https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Higher-order_statistics # Note that we are calculating the unbias (sample) version of skew here. # See https://en.wikipedia.org/wiki/Skewness#Sample_skewness # for more details. def compute_moments(x): n = len(x) if n == 0: m2, sum, m3 = 0, 0, 0 else: m2 = x.std(ddof=0)**2 * n sum = x.sum() m3 = (((x - x.mean())**3).sum()) return pd.DataFrame(dict(m2=[m2], sum=[sum], n=[n], m3=[m3])) def combine_moments(data): m2 = sum = n = m3 = 0.0 for datum in data.itertuples(): if datum.n == 0: continue elif n == 0: m2, sum, n, m3 = datum.m2, datum.sum, datum.n, datum.m3 else: n_a, n_b = datum.n, n sum_a, sum_b = datum.sum, sum m2_a, m2_b = datum.m2, m2 mean_a, mean_b = sum_a / n_a, sum_b / n_b delta = mean_b - mean_a combined_n = n_a + n_b m3 += datum.m3 + ( (delta**3 * ((n_a * n_b) * (n_a - n_b)) / ((combined_n)**2)) + ((3 * delta) * ((n_a * m2_b) - (n_b * m2_a)) / (combined_n))) m2 += datum.m2 + delta**2 * n_b * n_a / combined_n sum += datum.sum n += datum.n if n < 3: return float('nan') elif m2 == 0: return float(0) else: return combined_n * math.sqrt(combined_n - 1) / (combined_n - 2) * m3 / ( m2**(3 / 2)) moments = expressions.ComputedExpression( 'compute_moments', compute_moments, [self._expr], requires_partition_by=partitionings.Arbitrary()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'combine_moments', combine_moments, [moments], requires_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs( pd.Series, removed_args=["level"] if PD_VERSION >= (2, 0) else None) @frame_base.populate_defaults( pd.Series, removed_args=["level"] if PD_VERSION >= (2, 0) else None) def kurtosis(self, axis, skipna, level, numeric_only, **kwargs): if skipna is None or skipna: self = self.dropna() # pylint: disable=self-cls-assignment # See the online, numerically stable formulae at # https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Higher-order_statistics # kurtosis here calculated as sample kurtosis # https://en.wikipedia.org/wiki/Kurtosis#Sample_kurtosis def compute_moments(x): n = len(x) if n == 0: m2, sum, m3, m4 = 0, 0, 0, 0 else: m2 = x.std(ddof=0)**2 * n sum = x.sum() m3 = (((x - x.mean())**3).sum()) m4 = (((x - x.mean())**4).sum()) return pd.DataFrame(dict(m2=[m2], sum=[sum], n=[n], m3=[m3], m4=[m4])) def combine_moments(data): m2 = sum = n = m3 = m4 = 0.0 for datum in data.itertuples(): if datum.n == 0: continue elif n == 0: m2, sum, n, m3, m4 = datum.m2, datum.sum, datum.n, datum.m3, datum.m4 else: n_a, n_b = datum.n, n m2_a, m2_b = datum.m2, m2 m3_a, m3_b = datum.m3, m3 sum_a, sum_b = datum.sum, sum mean_a, mean_b = sum_a / n_a, sum_b / n_b delta = mean_b - mean_a combined_n = n_a + n_b m4 += datum.m4 + ((delta**4) * (n_a * n_b) * ( (n_a**2) - (n_a * n_b) + (n_b**2)) / combined_n**3) + ((6 * delta**2) * ((n_a**2 * m2_b) + (n_b**2 * m2_a)) / (combined_n**2)) + ((4 * delta) * ((n_a * m3_b) - (n_b * m3_a)) / (combined_n)) m3 += datum.m3 + ( (delta**3 * ((n_a * n_b) * (n_a - n_b)) / ((combined_n)**2)) + ((3 * delta) * ((n_a * m2_b) - (n_b * m2_a)) / (combined_n))) m2 += datum.m2 + delta**2 * n_b * n_a / combined_n sum += datum.sum n += datum.n if n < 4: return float('nan') elif m2 == 0: return float(0) else: return (((combined_n + 1) * (combined_n) * (combined_n - 1)) / ((combined_n - 2) * (combined_n - 3))) * (m4 / (m2)**2) - ((3 * (combined_n - 1)**2) / ((combined_n - 2) * (combined_n - 3))) moments = expressions.ComputedExpression( 'compute_moments', compute_moments, [self._expr], requires_partition_by=partitionings.Arbitrary()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'combine_moments', combine_moments, [moments], requires_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.Series) def kurt(self, *args, **kwargs): # Compute Kurtosis as kurt is an alias for kurtosis. return self.kurtosis(*args, **kwargs) def _corr_aligned(self, other, min_periods): std_x = self.std() std_y = other.std() cov = self._cov_aligned(other, min_periods) return cov.apply( lambda cov, std_x, std_y: cov / (std_x * std_y), args=[std_x, std_y]) @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) def cov(self, other, min_periods, ddof): x, y = self.dropna().align(other.dropna(), 'inner') return x._cov_aligned(y, min_periods, ddof) def _cov_aligned(self, other, min_periods, ddof=1): # Use the formulae from # https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Covariance def compute_co_moments(x, y): n = len(x) if n <= 1: c = 0 else: c = x.cov(y) * (n - 1) sx = x.sum() sy = y.sum() return pd.DataFrame(dict(c=[c], sx=[sx], sy=[sy], n=[n])) def combine_co_moments(data): c = sx = sy = n = 0.0 for datum in data.itertuples(): if datum.n == 0: continue elif n == 0: c, sx, sy, n = datum.c, datum.sx, datum.sy, datum.n else: c += ( datum.c + (sx / n - datum.sx / datum.n) * (sy / n - datum.sy / datum.n) * n * datum.n / (n + datum.n)) sx += datum.sx sy += datum.sy n += datum.n if n < max(2, ddof, min_periods or 0): return float('nan') else: return c / (n - ddof) moments = expressions.ComputedExpression( 'compute_co_moments', compute_co_moments, [self._expr, other._expr], requires_partition_by=partitionings.Index()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'combine_co_moments', combine_co_moments, [moments], requires_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) @frame_base.maybe_inplace def dropna(self, **kwargs): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'dropna', lambda df: df.dropna(**kwargs), [self._expr], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Arbitrary())) @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults( pd.Series, removed_args=['inplace'] if PD_VERSION >= (2, 0) else None) @frame_base.maybe_inplace def set_axis(self, labels, **kwargs): # TODO: assigning the index is generally order-sensitive, but we could # support it in some rare cases, e.g. when assigning the index from one # of a DataFrame's columns raise NotImplementedError( "Assigning an index is not yet supported. " "Consider using set_index() instead.") isnull = isna = frame_base._elementwise_method('isna', base=pd.Series) notnull = notna = frame_base._elementwise_method('notna', base=pd.Series) items = frame_base.wont_implement_method( pd.Series, 'items', reason="non-deferred-result") iteritems = frame_base.wont_implement_method( pd.Series, 'iteritems', reason="non-deferred-result") tolist = frame_base.wont_implement_method( pd.Series, 'tolist', reason="non-deferred-result") to_numpy = frame_base.wont_implement_method( pd.Series, 'to_numpy', reason="non-deferred-result") to_string = frame_base.wont_implement_method( pd.Series, 'to_string', reason="non-deferred-result") def _wrap_in_df(self): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'wrap_in_df', lambda s: pd.DataFrame(s), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary(), )) @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) @frame_base.maybe_inplace def duplicated(self, keep): """Only ``keep=False`` and ``keep="any"`` are supported. Other values of ``keep`` make this an order-sensitive operation. Note ``keep="any"`` is a Beam-specific option that guarantees only one duplicate will be kept, but unlike ``"first"`` and ``"last"`` it makes no guarantees about _which_ duplicate element is kept.""" # Re-use the DataFrame based duplcated, extract the series back out df = self._wrap_in_df() return df.duplicated(keep=keep)[df.columns[0]] @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) @frame_base.maybe_inplace def drop_duplicates(self, keep): """Only ``keep=False`` and ``keep="any"`` are supported. Other values of ``keep`` make this an order-sensitive operation. Note ``keep="any"`` is a Beam-specific option that guarantees only one duplicate will be kept, but unlike ``"first"`` and ``"last"`` it makes no guarantees about _which_ duplicate element is kept.""" # Re-use the DataFrame based drop_duplicates, extract the series back out df = self._wrap_in_df() return df.drop_duplicates(keep=keep)[df.columns[0]] @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) @frame_base.maybe_inplace def sample(self, **kwargs): """Only ``n`` and/or ``weights`` may be specified. ``frac``, ``random_state``, and ``replace=True`` are not yet supported. See `Issue 21010 `_. Note that pandas will raise an error if ``n`` is larger than the length of the dataset, while the Beam DataFrame API will simply return the full dataset in that case.""" # Re-use the DataFrame based sample, extract the series back out df = self._wrap_in_df() return df.sample(**kwargs)[df.columns[0]] @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) def aggregate(self, func, axis, *args, **kwargs): """Some aggregation methods cannot be parallelized, and computing them will require collecting all data on a single machine.""" if kwargs.get('skipna', False): # Eagerly generate a proxy to make sure skipna is a valid argument # for this aggregation method _ = self._expr.proxy().aggregate(func, axis, *args, **kwargs) kwargs.pop('skipna') return self.dropna().aggregate(func, axis, *args, **kwargs) if isinstance(func, list) and len(func) > 1: # level arg is ignored for multiple aggregations _ = kwargs.pop('level', None) # Aggregate with each method separately, then stick them all together. rows = [self.agg([f], *args, **kwargs) for f in func] return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'join_aggregate', lambda *rows: pd.concat(rows), [row._expr for row in rows])) else: # We're only handling a single column. It could be 'func' or ['func'], # which produce different results. 'func' produces a scalar, ['func'] # produces a single element Series. base_func = func[0] if isinstance(func, list) else func if (_is_numeric(base_func) and not pd.core.dtypes.common.is_numeric_dtype(self.dtype)): warnings.warn( f"Performing a numeric aggregation, {base_func!r}, on " f"Series {self._expr.proxy().name!r} with non-numeric type " f"{self.dtype!r}. This can result in runtime errors or surprising " "results.") if 'level' in kwargs: # Defer to groupby.agg for level= mode return self.groupby( level=kwargs.pop('level'), axis=axis).agg(func, *args, **kwargs) singleton_reason = None if 'min_count' in kwargs: # Eagerly generate a proxy to make sure min_count is a valid argument # for this aggregation method _ = self._expr.proxy().agg(func, axis, *args, **kwargs) singleton_reason = ( "Aggregation with min_count= requires collecting all data on a " "single node.") # We have specialized distributed implementations for these if base_func in HAND_IMPLEMENTED_GLOBAL_AGGREGATIONS: result = getattr(self, base_func)(*args, **kwargs) if isinstance(func, list): with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( f'wrap_aggregate_{base_func}', lambda x: pd.Series(x, index=[base_func]), [result._expr], requires_partition_by=partitionings.Singleton(), preserves_partition_by=partitionings.Singleton())) else: return result agg_kwargs = kwargs.copy() if ((_is_associative(base_func) or _is_liftable_with_sum(base_func)) and singleton_reason is None): intermediate = expressions.ComputedExpression( f'pre_aggregate_{base_func}', # Coerce to a Series, if the result is scalar we still want a Series # so we can combine and do the final aggregation next. lambda s: pd.Series(s.agg(func, *args, **kwargs)), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) allow_nonparallel_final = True if _is_associative(base_func): agg_func = func else: agg_func = ['sum'] if isinstance(func, list) else 'sum' else: intermediate = self._expr allow_nonparallel_final = None # i.e. don't change the value agg_func = func singleton_reason = ( f"Aggregation function {func!r} cannot currently be " "parallelized. It requires collecting all data for " "this Series on a single node.") with expressions.allow_non_parallel_operations(allow_nonparallel_final): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( f'post_aggregate_{base_func}', lambda s: s.agg(agg_func, *args, **agg_kwargs), [intermediate], preserves_partition_by=partitionings.Singleton(), requires_partition_by=partitionings.Singleton( reason=singleton_reason))) agg = aggregate @property # type: ignore @frame_base.with_docs_from(pd.Series) def axes(self): return [self.index] clip = frame_base._elementwise_method('clip', base=pd.Series) all = _agg_method(pd.Series, 'all') any = _agg_method(pd.Series, 'any') # TODO(BEAM-12074): Document that Series.count(level=) will drop NaN's count = _agg_method(pd.Series, 'count') describe = _agg_method(pd.Series, 'describe') min = _agg_method(pd.Series, 'min') max = _agg_method(pd.Series, 'max') prod = product = _agg_method(pd.Series, 'prod') sum = _agg_method(pd.Series, 'sum') median = _agg_method(pd.Series, 'median') sem = _agg_method(pd.Series, 'sem') # mad was removed in Pandas 2.0. if PD_VERSION < (2, 0): mad = _agg_method(pd.Series, 'mad') argmax = frame_base.wont_implement_method( pd.Series, 'argmax', reason='order-sensitive') argmin = frame_base.wont_implement_method( pd.Series, 'argmin', reason='order-sensitive') cummax = frame_base.wont_implement_method( pd.Series, 'cummax', reason='order-sensitive') cummin = frame_base.wont_implement_method( pd.Series, 'cummin', reason='order-sensitive') cumprod = frame_base.wont_implement_method( pd.Series, 'cumprod', reason='order-sensitive') cumsum = frame_base.wont_implement_method( pd.Series, 'cumsum', reason='order-sensitive') diff = frame_base.wont_implement_method( pd.Series, 'diff', reason='order-sensitive') interpolate = frame_base.wont_implement_method( pd.Series, 'interpolate', reason='order-sensitive') searchsorted = frame_base.wont_implement_method( pd.Series, 'searchsorted', reason='order-sensitive') shift = frame_base.wont_implement_method( pd.Series, 'shift', reason='order-sensitive') pct_change = frame_base.wont_implement_method( pd.Series, 'pct_change', reason='order-sensitive') is_monotonic = frame_base.wont_implement_method( pd.Series, 'is_monotonic', reason='order-sensitive') is_monotonic_increasing = frame_base.wont_implement_method( pd.Series, 'is_monotonic_increasing', reason='order-sensitive') is_monotonic_decreasing = frame_base.wont_implement_method( pd.Series, 'is_monotonic_decreasing', reason='order-sensitive') asof = frame_base.wont_implement_method( pd.Series, 'asof', reason='order-sensitive') first_valid_index = frame_base.wont_implement_method( pd.Series, 'first_valid_index', reason='order-sensitive') last_valid_index = frame_base.wont_implement_method( pd.Series, 'last_valid_index', reason='order-sensitive') autocorr = frame_base.wont_implement_method( pd.Series, 'autocorr', reason='order-sensitive') iat = property( frame_base.wont_implement_method( pd.Series, 'iat', reason='order-sensitive')) head = frame_base.wont_implement_method( pd.Series, 'head', explanation=_PEEK_METHOD_EXPLANATION) tail = frame_base.wont_implement_method( pd.Series, 'tail', explanation=_PEEK_METHOD_EXPLANATION) filter = frame_base._elementwise_method('filter', base=pd.Series) memory_usage = frame_base.wont_implement_method( pd.Series, 'memory_usage', reason="non-deferred-result") nbytes = frame_base.wont_implement_method( pd.Series, 'nbytes', reason="non-deferred-result") to_list = frame_base.wont_implement_method( pd.Series, 'to_list', reason="non-deferred-result") factorize = frame_base.wont_implement_method( pd.Series, 'factorize', reason="non-deferred-columns") # In Series __contains__ checks the index __contains__ = frame_base.wont_implement_method( pd.Series, '__contains__', reason="non-deferred-result") @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) def nlargest(self, keep, **kwargs): """Only ``keep=False`` and ``keep="any"`` are supported. Other values of ``keep`` make this an order-sensitive operation. Note ``keep="any"`` is a Beam-specific option that guarantees only one duplicate will be kept, but unlike ``"first"`` and ``"last"`` it makes no guarantees about _which_ duplicate element is kept.""" # TODO(robertwb): Document 'any' option. # TODO(robertwb): Consider (conditionally) defaulting to 'any' if no # explicit keep parameter is requested. if keep == 'any': keep = 'first' elif keep != 'all': raise frame_base.WontImplementError( f"nlargest(keep={keep!r}) is not supported because it is " "order sensitive. Only keep=\"all\" is supported.", reason="order-sensitive") kwargs['keep'] = keep per_partition = expressions.ComputedExpression( 'nlargest-per-partition', lambda df: df.nlargest(**kwargs), [self._expr], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Arbitrary()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'nlargest', lambda df: df.nlargest(**kwargs), [per_partition], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) def nsmallest(self, keep, **kwargs): """Only ``keep=False`` and ``keep="any"`` are supported. Other values of ``keep`` make this an order-sensitive operation. Note ``keep="any"`` is a Beam-specific option that guarantees only one duplicate will be kept, but unlike ``"first"`` and ``"last"`` it makes no guarantees about _which_ duplicate element is kept.""" if keep == 'any': keep = 'first' elif keep != 'all': raise frame_base.WontImplementError( f"nsmallest(keep={keep!r}) is not supported because it is " "order sensitive. Only keep=\"all\" is supported.", reason="order-sensitive") kwargs['keep'] = keep per_partition = expressions.ComputedExpression( 'nsmallest-per-partition', lambda df: df.nsmallest(**kwargs), [self._expr], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Arbitrary()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'nsmallest', lambda df: df.nsmallest(**kwargs), [per_partition], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Singleton())) @property # type: ignore @frame_base.with_docs_from(pd.Series) def is_unique(self): def set_index(s): s = s[:] s.index = s return s self_index = expressions.ComputedExpression( 'set_index', set_index, [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) is_unique_distributed = expressions.ComputedExpression( 'is_unique_distributed', lambda s: pd.Series(s.is_unique), [self_index], requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Singleton()) with expressions.allow_non_parallel_operations(): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'combine', lambda s: s.all(), [is_unique_distributed], requires_partition_by=partitionings.Singleton(), preserves_partition_by=partitionings.Singleton())) plot = frame_base.wont_implement_method( pd.Series, 'plot', reason="plotting-tools") pop = frame_base.wont_implement_method( pd.Series, 'pop', reason="non-deferred-result") rename_axis = frame_base._elementwise_method('rename_axis', base=pd.Series) round = frame_base._elementwise_method('round', base=pd.Series) take = frame_base.wont_implement_method( pd.Series, 'take', reason='deprecated') to_dict = frame_base.wont_implement_method( pd.Series, 'to_dict', reason="non-deferred-result") to_frame = frame_base._elementwise_method('to_frame', base=pd.Series) @frame_base.with_docs_from(pd.Series) def unique(self, as_series=False): """unique is not supported by default because it produces a non-deferred result: an :class:`~numpy.ndarray`. You can use the Beam-specific argument ``unique(as_series=True)`` to get the result as a :class:`DeferredSeries`""" if not as_series: raise frame_base.WontImplementError( "unique() is not supported by default because it produces a " "non-deferred result: a numpy array. You can use the Beam-specific " "argument unique(as_series=True) to get the result as a " "DeferredSeries", reason="non-deferred-result") return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'unique', lambda df: pd.Series(df.unique()), [self._expr], preserves_partition_by=partitionings.Singleton(), requires_partition_by=partitionings.Singleton( reason="unique() cannot currently be parallelized."))) @frame_base.with_docs_from(pd.Series) def update(self, other): self._expr = expressions.ComputedExpression( 'update', lambda df, other: df.update(other) or df, [self._expr, other._expr], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Index()) @frame_base.with_docs_from(pd.Series) def value_counts( self, sort=False, normalize=False, ascending=False, bins=None, dropna=True): """``sort`` is ``False`` by default, and ``sort=True`` is not supported because it imposes an ordering on the dataset which likely will not be preserved. When ``bin`` is specified this operation is not parallelizable. See [Issue 20903](https://github.com/apache/beam/issues/20903) tracking the possible addition of a distributed implementation.""" if sort: raise frame_base.WontImplementError( "value_counts(sort=True) is not supported because it imposes an " "ordering on the dataset which likely will not be preserved.", reason="order-sensitive") if bins is not None: return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'value_counts', lambda s: s.value_counts( normalize=normalize, bins=bins, dropna=dropna)[self._expr], requires_partition_by=partitionings.Singleton( reason=( "value_counts with bin specified requires collecting " "the entire dataset to identify the range.")), preserves_partition_by=partitionings.Singleton(), )) if dropna: column = self.dropna() else: column = self result = column.groupby(column, dropna=dropna).size() # Pandas 2 introduces new naming for the results. if PD_VERSION >= (2, 0): result.index.name = getattr(self, "name", None) result.name = "proportion" if normalize else "count" else: # groupby.size() names the index, which we don't need result.index.name = None if normalize: return result / column.length() else: return result values = property( frame_base.wont_implement_method( pd.Series, 'values', reason="non-deferred-result")) view = frame_base.wont_implement_method( pd.Series, 'view', explanation=( "because it relies on memory-sharing semantics that are " "not compatible with the Beam model.")) @property # type: ignore @frame_base.with_docs_from(pd.Series) def str(self): return _DeferredStringMethods(self._expr) @property # type: ignore @frame_base.with_docs_from(pd.Series) def cat(self): return _DeferredCategoricalMethods(self._expr) @property # type: ignore @frame_base.with_docs_from(pd.Series) def dt(self): return _DeferredDatetimeMethods(self._expr) @frame_base.with_docs_from(pd.Series) def mode(self, *args, **kwargs): """mode is not currently parallelizable. An approximate, parallelizable implementation of mode may be added in the future (`Issue 20946 `_).""" return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'mode', lambda df: df.mode(*args, **kwargs), [self._expr], #TODO(https://github.com/apache/beam/issues/20946): # Can we add an approximate implementation? requires_partition_by=partitionings.Singleton( reason=( "mode cannot currently be parallelized. See " "https://github.com/apache/beam/issues/20946 tracking the " "possble addition of an approximate, parallelizable " "implementation of mode.")), preserves_partition_by=partitionings.Singleton())) apply = frame_base._elementwise_method('apply', base=pd.Series) map = frame_base._elementwise_method('map', base=pd.Series) # TODO(https://github.com/apache/beam/issues/20764): Implement transform # using type inference to determine the proxy #transform = frame_base._elementwise_method('transform', base=pd.Series) @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) def repeat(self, repeats, axis): """``repeats`` must be an ``int`` or a :class:`DeferredSeries`. Lists are not supported because they make this operation order-sensitive.""" if isinstance(repeats, int): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'repeat', lambda series: series.repeat(repeats), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary())) elif isinstance(repeats, frame_base.DeferredBase): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'repeat', lambda series, repeats_series: series.repeat(repeats_series), [self._expr, repeats._expr], requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Arbitrary())) elif isinstance(repeats, list): raise frame_base.WontImplementError( "repeat(repeats=) repeats must be an int or a DeferredSeries. " "Lists are not supported because they make this operation sensitive " "to the order of the data.", reason="order-sensitive") else: raise TypeError( "repeat(repeats=) value must be an int or a " f"DeferredSeries (encountered {type(repeats)}).") if hasattr(pd.Series, 'compare'): @frame_base.with_docs_from(pd.Series) @frame_base.args_to_kwargs(pd.Series) @frame_base.populate_defaults(pd.Series) def compare(self, other, align_axis, **kwargs): if align_axis in ('index', 0): preserves_partition = partitionings.Singleton() elif align_axis in ('columns', 1): preserves_partition = partitionings.Arbitrary() else: raise ValueError( "align_axis must be one of ('index', 0, 'columns', 1). " f"got {align_axis!r}.") return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'compare', lambda s, other: s.compare(other, align_axis, **kwargs), [self._expr, other._expr], requires_partition_by=partitionings.Index(), preserves_partition_by=preserves_partition)) @populate_not_implemented(pd.DataFrame) @frame_base.DeferredFrame._register_for(pd.DataFrame) class DeferredDataFrame(DeferredDataFrameOrSeries): def __repr__(self): return ( f'DeferredDataFrame(columns={list(self.columns)}, ' f'{self._render_indexes()})') @property # type: ignore @frame_base.with_docs_from(pd.DataFrame) def columns(self): return self._expr.proxy().columns @columns.setter def columns(self, columns): def set_columns(df): df = df.copy() df.columns = columns return df return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'set_columns', set_columns, [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary())) @frame_base.with_docs_from(pd.DataFrame) def keys(self): return self.columns def __getattr__(self, name): # Column attribute access. if name in self._expr.proxy().columns: return self[name] else: return object.__getattribute__(self, name) def __getitem__(self, key): # TODO: Replicate pd.DataFrame.__getitem__ logic if isinstance(key, DeferredSeries) and key._expr.proxy().dtype == bool: return self.loc[key] elif isinstance(key, frame_base.DeferredBase): # Fail early if key is a DeferredBase as it interacts surprisingly with # key in self._expr.proxy().columns raise NotImplementedError( "Indexing with a non-bool deferred frame is not yet supported. " "Consider using df.loc[...]") elif isinstance(key, slice): if _is_null_slice(key): return self elif _is_integer_slice(key): # This depends on the contents of the index. raise frame_base.WontImplementError( "Integer slices are not supported as they are ambiguous. Please " "use iloc or loc with integer slices.") else: return self.loc[key] elif ( (isinstance(key, list) and all(key_column in self._expr.proxy().columns for key_column in key)) or key in self._expr.proxy().columns): return self._elementwise(lambda df: df[key], 'get_column') else: raise NotImplementedError(key) def __contains__(self, key): # Checks if proxy has the given column return self._expr.proxy().__contains__(key) def __setitem__(self, key, value): if isinstance( key, str) or (isinstance(key, list) and all(isinstance(c, str) for c in key)) or (isinstance(key, DeferredSeries) and key._expr.proxy().dtype == bool): # yapf: disable return self._elementwise( lambda df, key, value: df.__setitem__(key, value), 'set_column', (key, value), inplace=True) else: raise NotImplementedError(key) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def align(self, other, join, axis, copy, level, method, **kwargs): """Aligning per level is not yet supported. Only the default, ``level=None``, is allowed. Filling NaN values via ``method`` is not supported, because it is `order-sensitive `_. Only the default, ``method=None``, is allowed. ``copy=False`` is not supported because its behavior (whether or not it is an inplace operation) depends on the data.""" if not copy: raise frame_base.WontImplementError( "align(copy=False) is not supported because it might be an inplace " "operation depending on the data. Please prefer the default " "align(copy=True).") if method is not None and method != lib.no_default: raise frame_base.WontImplementError( f"align(method={method!r}) is not supported because it is " "order sensitive. Only align(method=None) is supported.", reason="order-sensitive") if kwargs: raise NotImplementedError('align(%s)' % ', '.join(kwargs.keys())) # In Pandas 2.0, all aggregations lost the level keyword. if PD_VERSION < (2, 0) and level is not None: # Could probably get by partitioning on the used levels. requires_partition_by = partitionings.Singleton(reason=( f"align(level={level}) is not currently parallelizable. Only " "align(level=None) can be parallelized.")) elif axis in ('columns', 1): requires_partition_by = partitionings.Arbitrary() else: requires_partition_by = partitionings.Index() return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'align', lambda df, other: df.align(other, join=join, axis=axis), [self._expr, other._expr], requires_partition_by=requires_partition_by, preserves_partition_by=partitionings.Arbitrary())) @frame_base.with_docs_from(pd.DataFrame, removed_method=PD_VERSION >= (2, 0)) @frame_base.args_to_kwargs(pd.DataFrame, removed_method=PD_VERSION >= (2, 0)) @frame_base.populate_defaults(pd.DataFrame, removed_method=PD_VERSION >= (2, 0)) def append(self, other, ignore_index, verify_integrity, sort, **kwargs): """``ignore_index=True`` is not supported, because it requires generating an order-sensitive index.""" if PD_VERSION >= (2, 0): raise frame_base.WontImplementError('append() was removed in Pandas 2.0.') if not isinstance(other, DeferredDataFrame): raise frame_base.WontImplementError( "append() only accepts DeferredDataFrame instances, received " + str(type(other))) if ignore_index: raise frame_base.WontImplementError( "append(ignore_index=True) is order sensitive because it requires " "generating a new index based on the order of the data.", reason="order-sensitive") if verify_integrity: # We can verify the index is non-unique within index partitioned data. requires = partitionings.Index() else: requires = partitionings.Arbitrary() return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'append', lambda s, other: s.append(other, sort=sort, verify_integrity=verify_integrity, **kwargs), [self._expr, other._expr], requires_partition_by=requires, preserves_partition_by=partitionings.Arbitrary() ) ) # If column name exists this is a simple project, otherwise it is a constant # (default_value) @frame_base.with_docs_from(pd.DataFrame) def get(self, key, default_value=None): if key in self.columns: return self[key] else: return default_value @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) @frame_base.maybe_inplace def set_index(self, keys, **kwargs): """``keys`` must be a ``str`` or ``list[str]``. Passing an Index or Series is not yet supported (`Issue 20759 `_).""" if isinstance(keys, str): keys = [keys] if any(isinstance(k, (_DeferredIndex, frame_base.DeferredFrame)) for k in keys): raise NotImplementedError("set_index with Index or Series instances is " "not yet supported " "(https://github.com/apache/beam/issues/20759)" ".") return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'set_index', lambda df: df.set_index(keys, **kwargs), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults( pd.DataFrame, removed_args=['inplace'] if PD_VERSION >= (2, 0) else None) @frame_base.maybe_inplace def set_axis(self, labels, axis, **kwargs): if axis in ('index', 0): # TODO: assigning the index is generally order-sensitive, but we could # support it in some rare cases, e.g. when assigning the index from one # of a DataFrame's columns raise NotImplementedError( "Assigning an index is not yet supported. " "Consider using set_index() instead.") else: return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'set_axis', lambda df: df.set_axis(labels, axis=axis, **kwargs), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary())) @property # type: ignore @frame_base.with_docs_from(pd.DataFrame) def axes(self): return (self.index, self.columns) @property # type: ignore @frame_base.with_docs_from(pd.DataFrame) def dtypes(self): return self._expr.proxy().dtypes @frame_base.with_docs_from(pd.DataFrame) def assign(self, **kwargs): """``value`` must be a ``callable`` or :class:`DeferredSeries`. Other types make this operation order-sensitive.""" for name, value in kwargs.items(): if not callable(value) and not isinstance(value, DeferredSeries): raise frame_base.WontImplementError( f"Unsupported value for new column '{name}': '{value}'. Only " "callables and DeferredSeries instances are supported. Other types " "make this operation sensitive to the order of the data", reason="order-sensitive") return self._elementwise( lambda df, *args, **kwargs: df.assign(*args, **kwargs), 'assign', other_kwargs=kwargs) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def explode(self, column, ignore_index): # ignoring the index will not preserve it preserves = (partitionings.Singleton() if ignore_index else partitionings.Index()) return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'explode', lambda df: df.explode(column, ignore_index), [self._expr], preserves_partition_by=preserves, requires_partition_by=partitionings.Arbitrary())) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def insert(self, value, **kwargs): """``value`` cannot be a ``List`` because aligning it with this DeferredDataFrame is order-sensitive.""" if isinstance(value, list): raise frame_base.WontImplementError( "insert(value=list) is not supported because it joins the input " "list to the deferred DataFrame based on the order of the data.", reason="order-sensitive") if isinstance(value, pd.core.generic.NDFrame): value = frame_base.DeferredFrame.wrap( expressions.ConstantExpression(value)) if isinstance(value, frame_base.DeferredFrame): def func_zip(df, value): df = df.copy() df.insert(value=value, **kwargs) return df inserted = frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'insert', func_zip, [self._expr, value._expr], requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Arbitrary())) else: def func_elementwise(df): df = df.copy() df.insert(value=value, **kwargs) return df inserted = frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'insert', func_elementwise, [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary())) self._expr = inserted._expr @staticmethod @frame_base.with_docs_from(pd.DataFrame) def from_dict(*args, **kwargs): return frame_base.DeferredFrame.wrap( expressions.ConstantExpression(pd.DataFrame.from_dict(*args, **kwargs))) @staticmethod @frame_base.with_docs_from(pd.DataFrame) def from_records(*args, **kwargs): return frame_base.DeferredFrame.wrap( expressions.ConstantExpression(pd.DataFrame.from_records(*args, **kwargs))) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) @frame_base.maybe_inplace def duplicated(self, keep, subset): """Only ``keep=False`` and ``keep="any"`` are supported. Other values of ``keep`` make this an order-sensitive operation. Note ``keep="any"`` is a Beam-specific option that guarantees only one duplicate will be kept, but unlike ``"first"`` and ``"last"`` it makes no guarantees about _which_ duplicate element is kept.""" # TODO(BEAM-12074): Document keep="any" if keep == 'any': keep = 'first' elif keep is not False: raise frame_base.WontImplementError( f"duplicated(keep={keep!r}) is not supported because it is " "sensitive to the order of the data. Only keep=False and " "keep=\"any\" are supported.", reason="order-sensitive") by = subset or list(self.columns) return self.groupby(by).apply( lambda df: pd.DataFrame(df.duplicated(keep=keep, subset=subset), columns=[None]))[None].droplevel(by) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) @frame_base.maybe_inplace def drop_duplicates(self, keep, subset, ignore_index): """Only ``keep=False`` and ``keep="any"`` are supported. Other values of ``keep`` make this an order-sensitive operation. Note ``keep="any"`` is a Beam-specific option that guarantees only one duplicate will be kept, but unlike ``"first"`` and ``"last"`` it makes no guarantees about _which_ duplicate element is kept.""" # TODO(BEAM-12074): Document keep="any" if keep == 'any': keep = 'first' elif keep is not False: raise frame_base.WontImplementError( f"drop_duplicates(keep={keep!r}) is not supported because it is " "sensitive to the order of the data. Only keep=False and " "keep=\"any\" are supported.", reason="order-sensitive") if ignore_index is not False: raise frame_base.WontImplementError( "drop_duplicates(ignore_index=False) is not supported because it " "requires generating a new index that is sensitive to the order of " "the data.", reason="order-sensitive") by = subset or list(self.columns) return self.groupby(by).apply( lambda df: df.drop_duplicates(keep=keep, subset=subset)).droplevel(by) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def aggregate(self, func, axis, *args, **kwargs): # We have specialized implementations for these. if func in ('quantile',): return getattr(self, func)(*args, axis=axis, **kwargs) # In pandas<1.3.0, maps to a property, args are ignored if func in ('size',) and PD_VERSION < (1, 3): return getattr(self, func) # We also have specialized distributed implementations for these. They only # support axis=0 (implicitly) though. axis=1 should fall through if func in ('corr', 'cov') and axis in (0, 'index'): return getattr(self, func)(*args, **kwargs) if axis is None: # Aggregate across all elements by first aggregating across columns, # then across rows. return self.agg(func, *args, **dict(kwargs, axis=1)).agg( func, *args, **dict(kwargs, axis=0)) elif axis in (1, 'columns'): # This is an easy elementwise aggregation. return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'aggregate', lambda df: df.agg(func, axis=1, *args, **kwargs), [self._expr], requires_partition_by=partitionings.Arbitrary())) elif len(self._expr.proxy().columns) == 0: # For this corner case, just colocate everything. return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'aggregate', lambda df: df.agg(func, *args, **kwargs), [self._expr], requires_partition_by=partitionings.Singleton())) else: # In the general case, we will compute the aggregation of each column # separately, then recombine. # First, handle any kwargs that cause a projection, by eagerly generating # the proxy, and only including the columns that are in the output. PROJECT_KWARGS = ('numeric_only', 'bool_only', 'include', 'exclude') proxy = self._expr.proxy().agg(func, axis, *args, **kwargs) if isinstance(proxy, pd.DataFrame): projected = self[list(proxy.columns)] elif isinstance(proxy, pd.Series): projected = self[list(proxy.index)] else: projected = self nonnumeric_columns = [name for (name, dtype) in projected.dtypes.items() if not pd.core.dtypes.common.is_numeric_dtype(dtype)] if _is_numeric(func) and nonnumeric_columns: if 'numeric_only' in kwargs and kwargs['numeric_only'] is False: # User has opted in to execution with non-numeric columns, they # will accept runtime errors pass else: raise frame_base.WontImplementError( f"Numeric aggregation ({func!r}) on a DataFrame containing " f"non-numeric columns ({*nonnumeric_columns,!r} is not " "supported, unless `numeric_only=` is specified.\n" "Use `numeric_only=True` to only aggregate over numeric " "columns.\nUse `numeric_only=False` to aggregate over all " "columns. Note this is not recommended, as it could result in " "execution time errors.") for key in PROJECT_KWARGS: if key in kwargs: kwargs.pop(key) if not isinstance(func, dict): col_names = list(projected._expr.proxy().columns) func_by_col = {col: func for col in col_names} else: func_by_col = func col_names = list(func.keys()) aggregated_cols = [] has_lists = any(isinstance(f, list) for f in func_by_col.values()) for col in col_names: funcs = func_by_col[col] if has_lists and not isinstance(funcs, list): # If any of the columns do multiple aggregations, they all must use # "list" style output funcs = [funcs] aggregated_cols.append(projected[col].agg(funcs, *args, **kwargs)) # The final shape is different depending on whether any of the columns # were aggregated by a list of aggregators. with expressions.allow_non_parallel_operations(): if isinstance(proxy, pd.Series): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'join_aggregate', lambda *cols: pd.Series( {col: value for col, value in zip(col_names, cols)}), [col._expr for col in aggregated_cols], requires_partition_by=partitionings.Singleton())) elif isinstance(proxy, pd.DataFrame): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'join_aggregate', lambda *cols: pd.DataFrame( {col: value for col, value in zip(col_names, cols)}), [col._expr for col in aggregated_cols], requires_partition_by=partitionings.Singleton())) else: raise AssertionError("Unexpected proxy type for " f"DataFrame.aggregate!: proxy={proxy!r}, " f"type(proxy)={type(proxy)!r}") agg = aggregate applymap = frame_base._elementwise_method('applymap', base=pd.DataFrame) if PD_VERSION >= (2, 1): map = frame_base._elementwise_method('map', base=pd.DataFrame) add_prefix = frame_base._elementwise_method('add_prefix', base=pd.DataFrame) add_suffix = frame_base._elementwise_method('add_suffix', base=pd.DataFrame) memory_usage = frame_base.wont_implement_method( pd.DataFrame, 'memory_usage', reason="non-deferred-result") info = frame_base.wont_implement_method( pd.DataFrame, 'info', reason="non-deferred-result") @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) @frame_base.maybe_inplace def clip(self, axis, **kwargs): """``lower`` and ``upper`` must be :class:`DeferredSeries` instances, or constants. Array-like arguments are not supported because they are order-sensitive.""" if any(isinstance(kwargs.get(arg, None), frame_base.DeferredFrame) for arg in ('upper', 'lower')) and axis not in (0, 'index'): raise frame_base.WontImplementError( "axis must be 'index' when upper and/or lower are a DeferredFrame", reason='order-sensitive') return frame_base._elementwise_method('clip', base=pd.DataFrame)(self, axis=axis, **kwargs) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def corr(self, method, min_periods): """Only ``method="pearson"`` can be parallelized. Other methods require collecting all data on a single worker (see https://s.apache.org/dataframe-non-parallel-operations for details). """ if method == 'pearson': proxy = self._expr.proxy().corr() columns = list(proxy.columns) args = [] arg_indices = [] for col1, col2 in itertools.combinations(columns, 2): arg_indices.append((col1, col2)) args.append(self[col1].corr(self[col2], method=method, min_periods=min_periods)) def fill_matrix(*args): data = collections.defaultdict(dict) for col in columns: data[col][col] = 1.0 for ix, (col1, col2) in enumerate(arg_indices): data[col1][col2] = data[col2][col1] = args[ix] return pd.DataFrame(data, columns=columns, index=columns) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'fill_matrix', fill_matrix, [arg._expr for arg in args], requires_partition_by=partitionings.Singleton(), proxy=proxy)) else: reason = (f"Encountered corr(method={method!r}) which cannot be " "parallelized. Only corr(method='pearson') is currently " "parallelizable.") return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'corr', lambda df: df.corr(method=method, min_periods=min_periods), [self._expr], requires_partition_by=partitionings.Singleton(reason=reason))) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def cov(self, min_periods, ddof): proxy = self._expr.proxy().corr() columns = list(proxy.columns) args = [] arg_indices = [] for col in columns: arg_indices.append((col, col)) std = self[col].std(ddof) args.append(std.apply(lambda x: x*x, 'square')) for ix, col1 in enumerate(columns): for col2 in columns[ix+1:]: arg_indices.append((col1, col2)) # Note that this set may be different for each pair. no_na = self.loc[self[col1].notna() & self[col2].notna()] args.append(no_na[col1]._cov_aligned(no_na[col2], min_periods, ddof)) def fill_matrix(*args): data = collections.defaultdict(dict) for ix, (col1, col2) in enumerate(arg_indices): data[col1][col2] = data[col2][col1] = args[ix] return pd.DataFrame(data, columns=columns, index=columns) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'fill_matrix', fill_matrix, [arg._expr for arg in args], requires_partition_by=partitionings.Singleton(), proxy=proxy)) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def corrwith(self, other, axis, drop, method): if axis in (1, 'columns'): return self._elementwise( lambda df, other: df.corrwith(other, axis=axis, drop=drop, method=method), 'corrwith', other_args=(other,)) if not isinstance(other, frame_base.DeferredFrame): other = frame_base.DeferredFrame.wrap( expressions.ConstantExpression(other)) if isinstance(other, DeferredSeries): proxy = self._expr.proxy().corrwith(other._expr.proxy(), axis=axis, drop=drop, method=method) self, other = self.align(other, axis=0, join='inner') col_names = proxy.index other_cols = [other] * len(col_names) elif isinstance(other, DeferredDataFrame): proxy = self._expr.proxy().corrwith( other._expr.proxy(), axis=axis, method=method, drop=drop) self, other = self.align(other, axis=0, join='inner') col_names = list( set(self.columns) .intersection(other.columns) .intersection(proxy.index)) other_cols = [other[col_name] for col_name in col_names] else: # Raise the right error. self._expr.proxy().corrwith(other._expr.proxy(), axis=axis, drop=drop, method=method) # Just in case something else becomes valid. raise NotImplementedError('corrwith(%s)' % type(other._expr.proxy)) # Generate expressions to compute the actual correlations. corrs = [ self[col_name].corr(other_col, method) for col_name, other_col in zip(col_names, other_cols)] # Combine the results def fill_dataframe(*args): result = proxy.copy(deep=True) for col, value in zip(proxy.index, args): result[col] = value return result with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'fill_dataframe', fill_dataframe, [corr._expr for corr in corrs], requires_partition_by=partitionings.Singleton(), proxy=proxy)) cummax = frame_base.wont_implement_method(pd.DataFrame, 'cummax', reason='order-sensitive') cummin = frame_base.wont_implement_method(pd.DataFrame, 'cummin', reason='order-sensitive') cumprod = frame_base.wont_implement_method(pd.DataFrame, 'cumprod', reason='order-sensitive') cumsum = frame_base.wont_implement_method(pd.DataFrame, 'cumsum', reason='order-sensitive') # TODO(BEAM-12071): Consider adding an order-insensitive implementation for # diff that relies on the index diff = frame_base.wont_implement_method(pd.DataFrame, 'diff', reason='order-sensitive') interpolate = frame_base.wont_implement_method(pd.DataFrame, 'interpolate', reason='order-sensitive') pct_change = frame_base.wont_implement_method( pd.DataFrame, 'pct_change', reason='order-sensitive') asof = frame_base.wont_implement_method( pd.DataFrame, 'asof', reason='order-sensitive') first_valid_index = frame_base.wont_implement_method( pd.DataFrame, 'first_valid_index', reason='order-sensitive') last_valid_index = frame_base.wont_implement_method( pd.DataFrame, 'last_valid_index', reason='order-sensitive') iat = property(frame_base.wont_implement_method( pd.DataFrame, 'iat', reason='order-sensitive')) lookup = frame_base.wont_implement_method( pd.DataFrame, 'lookup', reason='deprecated') head = frame_base.wont_implement_method(pd.DataFrame, 'head', explanation=_PEEK_METHOD_EXPLANATION) tail = frame_base.wont_implement_method(pd.DataFrame, 'tail', explanation=_PEEK_METHOD_EXPLANATION) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def sample(self, n, frac, replace, weights, random_state, axis): """When ``axis='index'``, only ``n`` and/or ``weights`` may be specified. ``frac``, ``random_state``, and ``replace=True`` are not yet supported. See `Issue 21010 `_. Note that pandas will raise an error if ``n`` is larger than the length of the dataset, while the Beam DataFrame API will simply return the full dataset in that case. sample is fully supported for axis='columns'.""" if axis in (1, 'columns'): # Sampling on axis=columns just means projecting random columns # Eagerly generate proxy to determine the set of columns at construction # time proxy = self._expr.proxy().sample(n=n, frac=frac, replace=replace, weights=weights, random_state=random_state, axis=axis) # Then do the projection return self[list(proxy.columns)] # axis='index' if frac is not None or random_state is not None or replace: raise NotImplementedError( f"When axis={axis!r}, only n and/or weights may be specified. " "frac, random_state, and replace=True are not yet supported " f"(got frac={frac!r}, random_state={random_state!r}, " f"replace={replace!r}). See " "https://github.com/apache/beam/issues/21010.") if n is None: n = 1 if isinstance(weights, str): weights = self[weights] tmp_weight_column_name = "___Beam_DataFrame_weights___" if weights is None: self_with_randomized_weights = frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'randomized_weights', lambda df: df.assign(**{tmp_weight_column_name: np.random.rand(len(df))}), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary())) else: # See "Fast Parallel Weighted Random Sampling" by Efraimidis and Spirakis # https://www.cti.gr/images_gr/reports/99-06-02.ps def assign_randomized_weights(df, weights): non_zero_weights = (weights > 0) | pd.Series(dtype=bool, index=df.index) df = df.loc[non_zero_weights] weights = weights.loc[non_zero_weights] random_weights = np.log(np.random.rand(len(weights))) / weights return df.assign(**{tmp_weight_column_name: random_weights}) self_with_randomized_weights = frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'randomized_weights', assign_randomized_weights, [self._expr, weights._expr], requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Arbitrary())) return self_with_randomized_weights.nlargest( n=n, columns=tmp_weight_column_name, keep='any').drop( tmp_weight_column_name, axis=1) @frame_base.with_docs_from(pd.DataFrame) def dot(self, other): # We want to broadcast the right hand side to all partitions of the left. # This is OK, as its index must be the same size as the columns set of self, # so cannot be too large. class AsScalar(object): def __init__(self, value): self.value = value if isinstance(other, frame_base.DeferredFrame): proxy = other._expr.proxy() with expressions.allow_non_parallel_operations(): side = expressions.ComputedExpression( 'as_scalar', lambda df: AsScalar(df), [other._expr], requires_partition_by=partitionings.Singleton()) else: proxy = pd.DataFrame(columns=range(len(other[0]))) side = expressions.ConstantExpression(AsScalar(other)) return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'dot', lambda left, right: left @ right.value, [self._expr, side], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary(), proxy=proxy)) __matmul__ = dot @frame_base.with_docs_from(pd.DataFrame) def mode(self, axis=0, *args, **kwargs): """mode with axis="columns" is not implemented because it produces non-deferred columns. mode with axis="index" is not currently parallelizable. An approximate, parallelizable implementation of mode may be added in the future (`Issue 20946 `_).""" if axis == 1 or axis == 'columns': # Number of columns is max(number mode values for each row), so we can't # determine how many there will be before looking at the data. raise frame_base.WontImplementError( "mode(axis=columns) is not supported because it produces a variable " "number of columns depending on the data.", reason="non-deferred-columns") return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'mode', lambda df: df.mode(*args, **kwargs), [self._expr], #TODO(https://github.com/apache/beam/issues/20946): # Can we add an approximate implementation? requires_partition_by=partitionings.Singleton(reason=( "mode(axis='index') cannot currently be parallelized. See " "https://github.com/apache/beam/issues/20946 tracking the " "possble addition of an approximate, parallelizable " "implementation of mode." )), preserves_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) @frame_base.maybe_inplace def dropna(self, axis, **kwargs): """dropna with axis="columns" specified cannot be parallelized.""" # TODO(robertwb): This is a common pattern. Generalize? if axis in (1, 'columns'): requires_partition_by = partitionings.Singleton(reason=( "dropna(axis=1) cannot currently be parallelized. It requires " "checking all values in each column for NaN values, to determine " "if that column should be dropped." )) else: requires_partition_by = partitionings.Arbitrary() return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'dropna', lambda df: df.dropna(axis=axis, **kwargs), [self._expr], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=requires_partition_by)) def _eval_or_query(self, name, expr, inplace, **kwargs): for key in ('local_dict', 'global_dict', 'level', 'target', 'resolvers'): if key in kwargs: raise NotImplementedError(f"Setting '{key}' is not yet supported") # look for '@' if re.search(r'\@[^\d\W]\w*', expr, re.UNICODE): raise NotImplementedError("Accessing locals with @ is not yet supported " "(https://github.com/apache/beam/issues/20626)" ) result_expr = expressions.ComputedExpression( name, lambda df: getattr(df, name)(expr, **kwargs), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary()) if inplace: self._expr = result_expr else: return frame_base.DeferredFrame.wrap(result_expr) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def eval(self, expr, inplace, **kwargs): """Accessing local variables with ``@`` is not yet supported (`Issue 20626 `_). Arguments ``local_dict``, ``global_dict``, ``level``, ``target``, and ``resolvers`` are not yet supported.""" return self._eval_or_query('eval', expr, inplace, **kwargs) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def query(self, expr, inplace, **kwargs): """Accessing local variables with ``@`` is not yet supported (`Issue 20626 `_). Arguments ``local_dict``, ``global_dict``, ``level``, ``target``, and ``resolvers`` are not yet supported.""" return self._eval_or_query('query', expr, inplace, **kwargs) isnull = isna = frame_base._elementwise_method('isna', base=pd.DataFrame) notnull = notna = frame_base._elementwise_method('notna', base=pd.DataFrame) items = frame_base.wont_implement_method(pd.DataFrame, 'items', reason="non-deferred-result") itertuples = frame_base.wont_implement_method(pd.DataFrame, 'itertuples', reason="non-deferred-result") iterrows = frame_base.wont_implement_method(pd.DataFrame, 'iterrows', reason="non-deferred-result") iteritems = frame_base.wont_implement_method(pd.DataFrame, 'iteritems', reason="non-deferred-result") def _cols_as_temporary_index(self, cols, suffix=''): original_index_names = list(self._expr.proxy().index.names) new_index_names = [ '__apache_beam_temp_%d_%s' % (ix, suffix) for (ix, _) in enumerate(original_index_names)] def reindex(df): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'reindex', lambda df: df.rename_axis(index=new_index_names, copy=False) .reset_index().set_index(cols), [df._expr], preserves_partition_by=partitionings.Singleton(), requires_partition_by=partitionings.Arbitrary())) def revert(df): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'join_restoreindex', lambda df: df.reset_index().set_index(new_index_names) .rename_axis(index=original_index_names, copy=False), [df._expr], preserves_partition_by=partitionings.Singleton(), requires_partition_by=partitionings.Arbitrary())) return reindex, revert @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def join(self, other, on, **kwargs): if on is not None: reindex, revert = self._cols_as_temporary_index(on) return revert(reindex(self).join(other, **kwargs)) if isinstance(other, list): other_is_list = True else: other = [other] other_is_list = False placeholder = object() other_exprs = [ df._expr for df in other if isinstance(df, frame_base.DeferredFrame)] const_others = [ placeholder if isinstance(df, frame_base.DeferredFrame) else df for df in other] def fill_placeholders(values): values = iter(values) filled = [ next(values) if df is placeholder else df for df in const_others] if other_is_list: return filled else: return filled[0] return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'join', lambda df, *deferred_others: df.join( fill_placeholders(deferred_others), **kwargs), [self._expr] + other_exprs, preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Index())) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def merge( self, right, on, left_on, right_on, left_index, right_index, suffixes, **kwargs): """merge is not parallelizable unless ``left_index`` or ``right_index`` is ``True`, because it requires generating an entirely new unique index. See notes on :meth:`DeferredDataFrame.reset_index`. It is recommended to move the join key for one of your columns to the index to avoid this issue. For an example see the enrich pipeline in :mod:`apache_beam.examples.dataframe.taxiride`. ``how="cross"`` is not yet supported. """ self_proxy = self._expr.proxy() right_proxy = right._expr.proxy() # Validate with a pandas call. _ = self_proxy.merge( right_proxy, on=on, left_on=left_on, right_on=right_on, left_index=left_index, right_index=right_index, **kwargs) if kwargs.get('how', None) == 'cross': raise NotImplementedError( "cross join is not yet implemented " "(https://github.com/apache/beam/issues/20318)") if not any([on, left_on, right_on, left_index, right_index]): on = [col for col in self_proxy.columns if col in right_proxy.columns] if not left_on: left_on = on if left_on and not isinstance(left_on, list): left_on = [left_on] if not right_on: right_on = on if right_on and not isinstance(right_on, list): right_on = [right_on] if left_index: indexed_left = self else: indexed_left = self.set_index(left_on, drop=False) if right_index: indexed_right = right else: indexed_right = right.set_index(right_on, drop=False) if left_on and right_on: common_cols = set(left_on).intersection(right_on) if len(common_cols): # When merging on the same column name from both dfs, we need to make # sure only one df has the column. Otherwise we end up with # two duplicate columns, one with lsuffix and one with rsuffix. # It's safe to drop from either because the data has already been duped # to the index. indexed_right = indexed_right.drop(columns=common_cols) merged = frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'merge', lambda left, right: left.merge(right, left_index=True, right_index=True, suffixes=suffixes, **kwargs), [indexed_left._expr, indexed_right._expr], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Index())) if left_index or right_index: return merged else: return merged.reset_index(drop=True) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def nlargest(self, keep, **kwargs): """Only ``keep=False`` and ``keep="any"`` are supported. Other values of ``keep`` make this an order-sensitive operation. Note ``keep="any"`` is a Beam-specific option that guarantees only one duplicate will be kept, but unlike ``"first"`` and ``"last"`` it makes no guarantees about _which_ duplicate element is kept.""" if keep == 'any': keep = 'first' elif keep != 'all': raise frame_base.WontImplementError( f"nlargest(keep={keep!r}) is not supported because it is " "order sensitive. Only keep=\"all\" is supported.", reason="order-sensitive") kwargs['keep'] = keep per_partition = expressions.ComputedExpression( 'nlargest-per-partition', lambda df: df.nlargest(**kwargs), [self._expr], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Arbitrary()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'nlargest', lambda df: df.nlargest(**kwargs), [per_partition], preserves_partition_by=partitionings.Singleton(), requires_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def nsmallest(self, keep, **kwargs): """Only ``keep=False`` and ``keep="any"`` are supported. Other values of ``keep`` make this an order-sensitive operation. Note ``keep="any"`` is a Beam-specific option that guarantees only one duplicate will be kept, but unlike ``"first"`` and ``"last"`` it makes no guarantees about _which_ duplicate element is kept.""" if keep == 'any': keep = 'first' elif keep != 'all': raise frame_base.WontImplementError( f"nsmallest(keep={keep!r}) is not supported because it is " "order sensitive. Only keep=\"all\" is supported.", reason="order-sensitive") kwargs['keep'] = keep per_partition = expressions.ComputedExpression( 'nsmallest-per-partition', lambda df: df.nsmallest(**kwargs), [self._expr], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Arbitrary()) with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'nsmallest', lambda df: df.nsmallest(**kwargs), [per_partition], preserves_partition_by=partitionings.Singleton(), requires_partition_by=partitionings.Singleton())) plot = frame_base.wont_implement_method(pd.DataFrame, 'plot', reason="plotting-tools") @frame_base.with_docs_from(pd.DataFrame) def pop(self, item): result = self[item] self._expr = expressions.ComputedExpression( 'popped', lambda df: df.drop(columns=[item]), [self._expr], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Arbitrary()) return result @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def quantile(self, q, axis, **kwargs): """``quantile(axis="index")`` is not parallelizable. See `Issue 20933 `_ tracking the possible addition of an approximate, parallelizable implementation of quantile. When using quantile with ``axis="columns"`` only a single ``q`` value can be specified.""" if axis in (1, 'columns'): if isinstance(q, list): raise frame_base.WontImplementError( "quantile(axis=columns) with multiple q values is not supported " "because it transposes the input DataFrame. Note computing " "an individual quantile across columns (e.g. " f"df.quantile(q={q[0]!r}, axis={axis!r}) is supported.", reason="non-deferred-columns") else: requires = partitionings.Arbitrary() else: # axis='index' # TODO(https://github.com/apache/beam/issues/20933): Provide an option # for approximate distributed quantiles requires = partitionings.Singleton(reason=( "Computing quantiles across index cannot currently be parallelized. " "See https://github.com/apache/beam/issues/20933 tracking the " "possible addition of an approximate, parallelizable implementation " "of quantile." )) return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'quantile', lambda df: df.quantile(q=q, axis=axis, **kwargs), [self._expr], requires_partition_by=requires, preserves_partition_by=partitionings.Singleton())) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.maybe_inplace def rename(self, **kwargs): """rename is not parallelizable when ``axis="index"`` and ``errors="raise"``. It requires collecting all data on a single node in order to detect if one of the index values is missing.""" rename_index = ( 'index' in kwargs or kwargs.get('axis', None) in (0, 'index') or ('columns' not in kwargs and 'axis' not in kwargs)) rename_columns = ( 'columns' in kwargs or kwargs.get('axis', None) in (1, 'columns')) if rename_index: # Technically, it's still partitioned by index, but it's no longer # partitioned by the hash of the index. preserves_partition_by = partitionings.Singleton() else: preserves_partition_by = partitionings.Index() if kwargs.get('errors', None) == 'raise' and rename_index: # TODO: We could do this in parallel by creating a ConstantExpression # with a series created from the mapper dict. Then Index() partitioning # would co-locate the necessary index values and we could raise # individually within each partition. Execution time errors are # discouraged anyway so probably not worth the effort. requires_partition_by = partitionings.Singleton(reason=( "rename(errors='raise', axis='index') requires collecting all " "data on a single node in order to detect missing index values." )) else: requires_partition_by = partitionings.Arbitrary() proxy = None if rename_index: # The proxy can't be computed by executing rename, it will error # renaming the index. if rename_columns: # Note if both are being renamed, index and columns must be specified # (not axis) proxy = self._expr.proxy().rename(**{k: v for (k, v) in kwargs.items() if not k == 'index'}) else: # No change in columns, reuse proxy proxy = self._expr.proxy() return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'rename', lambda df: df.rename(**kwargs), [self._expr], proxy=proxy, preserves_partition_by=preserves_partition_by, requires_partition_by=requires_partition_by)) rename_axis = frame_base._elementwise_method('rename_axis', base=pd.DataFrame) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def round(self, decimals, *args, **kwargs): if isinstance(decimals, frame_base.DeferredFrame): # Disallow passing a deferred Series in, our current partitioning model # prevents us from using it correctly. raise NotImplementedError("Passing a deferred series to round() is not " "supported, please use a concrete pd.Series " "instance or a dictionary") return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'round', lambda df: df.round(decimals, *args, **kwargs), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Index() ) ) select_dtypes = frame_base._elementwise_method('select_dtypes', base=pd.DataFrame) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def shift(self, axis, freq, **kwargs): """shift with ``axis="index" is only supported with ``freq`` specified and ``fill_value`` undefined. Other configurations make this operation order-sensitive.""" if axis in (1, 'columns'): preserves = partitionings.Arbitrary() proxy = None else: if freq is None or 'fill_value' in kwargs: fill_value = kwargs.get('fill_value', 'NOT SET') raise frame_base.WontImplementError( f"shift(axis={axis!r}) is only supported with freq defined, and " f"fill_value undefined (got freq={freq!r}," f"fill_value={fill_value!r}). Other configurations are sensitive " "to the order of the data because they require populating shifted " "rows with `fill_value`.", reason="order-sensitive") # proxy generation fails in pandas <1.2 # Seems due to https://github.com/pandas-dev/pandas/issues/14811, # bug with shift on empty indexes. # Fortunately the proxy should be identical to the input. proxy = self._expr.proxy().copy() # index is modified, so no partitioning is preserved. preserves = partitionings.Singleton() return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'shift', lambda df: df.shift(axis=axis, freq=freq, **kwargs), [self._expr], proxy=proxy, preserves_partition_by=preserves, requires_partition_by=partitionings.Arbitrary())) shape = property(frame_base.wont_implement_method( pd.DataFrame, 'shape', reason="non-deferred-result")) stack = frame_base._proxy_method( 'stack', base=pd.DataFrame, requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton()) all = _agg_method(pd.DataFrame, 'all') any = _agg_method(pd.DataFrame, 'any') count = _agg_method(pd.DataFrame, 'count') describe = _agg_method(pd.DataFrame, 'describe') max = _agg_method(pd.DataFrame, 'max') min = _agg_method(pd.DataFrame, 'min') @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def pivot(self, index=None, columns=None, values=None, **kwargs): """Because pivot is a non-deferred method, any columns specified in ``columns`` must be CategoricalDType so we can determine the output column names.""" def verify_all_categorical(all_cols_are_categorical): if not all_cols_are_categorical: message = "pivot() of non-categorical type is not supported because " \ "the type of the output column depends on the data. Please use " \ "pd.CategoricalDtype with explicit categories." raise frame_base.WontImplementError( message, reason="non-deferred-columns") # If values not provided, take all remaining columns of dataframe if not values: tmp = self._expr.proxy() if index: tmp = tmp.drop(index, axis=1) if columns: tmp = tmp.drop(columns, axis=1) values = tmp.columns.values # Construct column index if is_list_like(columns) and len(columns) <= 1: columns = columns[0] selected_cols = self._expr.proxy()[columns] if isinstance(selected_cols, pd.Series): all_cols_are_categorical = isinstance( selected_cols.dtype, pd.CategoricalDtype ) verify_all_categorical(all_cols_are_categorical) if is_list_like(values) and len(values) > 1: # If more than one value provided, don't create a None level values_in_col_index = values names = [None, columns] col_index = pd.MultiIndex.from_product( [values_in_col_index, selected_cols.dtypes.categories.astype('category')], names=names ) else: col_index = pd.CategoricalIndex( selected_cols.dtype.categories, name=columns ) else: all_cols_are_categorical = all( isinstance(c, pd.CategoricalDtype) for c in selected_cols.dtypes ) verify_all_categorical(all_cols_are_categorical) if is_list_like(values) and len(values) > 1: # If more than one value provided, don't create a None level values_in_col_index = values names = [None, *columns] categories = [ c.categories.astype('category') for c in selected_cols.dtypes ] col_index = pd.MultiIndex.from_product( [values_in_col_index, *categories], names=names ) else: # If one value provided, don't create a None level names = columns categories = [ c.categories.astype('category') for c in selected_cols.dtypes ] col_index = pd.MultiIndex.from_product( categories, names=names ) # Construct row index if index: if PD_VERSION < (1, 4) and is_list_like(index) and len(index) > 1: raise frame_base.WontImplementError( "pivot() is not supported when pandas<1.4 and index is a MultiIndex") per_partition = expressions.ComputedExpression( 'pivot-per-partition', lambda df: df.set_index(keys=index), [self._expr], preserves_partition_by=partitionings.Singleton(), requires_partition_by=partitionings.Arbitrary() ) tmp = per_partition.proxy().pivot( columns=columns, values=values, **kwargs) row_index = tmp.index else: per_partition = self._expr row_index = self._expr.proxy().index if PD_VERSION < (1, 4) and isinstance(row_index, pd.MultiIndex): raise frame_base.WontImplementError( "pivot() is not supported when pandas<1.4 and index is a MultiIndex") selected_values = self._expr.proxy()[values] if isinstance(selected_values, pd.Series): value_dtype = selected_values.dtype else: # Set dtype to object if more than one value dtypes = [d for d in selected_values.dtypes] value_dtype = object if any((is_int64_dtype(x) for x in dtypes)): value_dtype = np.int64 if any((is_float_dtype(x) for x in dtypes)): value_dtype = np.float64 if object in dtypes: value_dtype = object # Construct proxy proxy = pd.DataFrame( columns=col_index, dtype=value_dtype, index=row_index ) def pivot_helper(df): result = pd.concat( [proxy, df.pivot(columns=columns, values=values, **kwargs)] ) result.columns = col_index result = result.astype(value_dtype) return result return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'pivot', pivot_helper, [per_partition], proxy=proxy, preserves_partition_by=partitionings.Index(), requires_partition_by=partitionings.Index())) prod = product = _agg_method(pd.DataFrame, 'prod') sum = _agg_method(pd.DataFrame, 'sum') mean = _agg_method(pd.DataFrame, 'mean') median = _agg_method(pd.DataFrame, 'median') nunique = _agg_method(pd.DataFrame, 'nunique') std = _agg_method(pd.DataFrame, 'std') var = _agg_method(pd.DataFrame, 'var') sem = _agg_method(pd.DataFrame, 'sem') skew = _agg_method(pd.DataFrame, 'skew') kurt = _agg_method(pd.DataFrame, 'kurt') kurtosis = _agg_method(pd.DataFrame, 'kurtosis') # mad was removed in Pandas 2.0. if PD_VERSION < (2, 0): mad = _agg_method(pd.DataFrame, 'mad') take = frame_base.wont_implement_method(pd.DataFrame, 'take', reason='deprecated') to_records = frame_base.wont_implement_method(pd.DataFrame, 'to_records', reason="non-deferred-result") to_dict = frame_base.wont_implement_method(pd.DataFrame, 'to_dict', reason="non-deferred-result") to_numpy = frame_base.wont_implement_method(pd.DataFrame, 'to_numpy', reason="non-deferred-result") to_string = frame_base.wont_implement_method(pd.DataFrame, 'to_string', reason="non-deferred-result") to_sparse = frame_base.wont_implement_method(pd.DataFrame, 'to_sparse', reason="non-deferred-result") transpose = frame_base.wont_implement_method( pd.DataFrame, 'transpose', reason='non-deferred-columns') T = property(frame_base.wont_implement_method( pd.DataFrame, 'T', reason='non-deferred-columns')) update = frame_base._proxy_method( 'update', inplace=True, base=pd.DataFrame, requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Arbitrary()) values = property(frame_base.wont_implement_method( pd.DataFrame, 'values', reason="non-deferred-result")) style = property(frame_base.wont_implement_method( pd.DataFrame, 'style', reason="non-deferred-result")) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def melt(self, ignore_index, **kwargs): """``ignore_index=True`` is not supported, because it requires generating an order-sensitive index.""" if ignore_index: raise frame_base.WontImplementError( "melt(ignore_index=True) is order sensitive because it requires " "generating a new index based on the order of the data.", reason="order-sensitive") return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'melt', lambda df: df.melt(ignore_index=False, **kwargs), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton())) if hasattr(pd.DataFrame, 'value_counts'): @frame_base.with_docs_from(pd.DataFrame) def value_counts(self, subset=None, sort=False, normalize=False, ascending=False, dropna=True): """``sort`` is ``False`` by default, and ``sort=True`` is not supported because it imposes an ordering on the dataset which likely will not be preserved.""" if sort: raise frame_base.WontImplementError( "value_counts(sort=True) is not supported because it imposes an " "ordering on the dataset which likely will not be preserved.", reason="order-sensitive") columns = subset or list(self.columns) if dropna: # Must include subset here because otherwise we spuriously drop NAs due # to columns outside our subset. dropped = self.dropna(subset=subset) else: dropped = self result = dropped.groupby(columns, dropna=dropna).size() # Pandas 2 introduces new naming for the results. if PD_VERSION >= (2,0): result.name = "proportion" if normalize else "count" if normalize: return result/dropped.length() else: return result if hasattr(pd.DataFrame, 'compare'): @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def compare(self, other, align_axis, keep_shape, **kwargs): """The default values ``align_axis=1 and ``keep_shape=False`` are not supported, because the output columns depend on the data. To use ``align_axis=1``, please specify ``keep_shape=True``.""" preserve_partition = None if align_axis in (1, 'columns') and not keep_shape: raise frame_base.WontImplementError( f"compare(align_axis={align_axis!r}, keep_shape={keep_shape!r}) " "is not allowed because the output columns depend on the data, " "please specify keep_shape=True.", reason='non-deferred-columns' ) if align_axis in (1, 'columns'): preserve_partition = partitionings.Arbitrary() elif align_axis in (0, 'index'): preserve_partition = partitionings.Singleton() else: raise ValueError( "align_axis must be one of ('index', 0, 'columns', 1). " f"got {align_axis!r}.") return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'compare', lambda df, other: df.compare(other, align_axis, keep_shape, **kwargs), [self._expr, other._expr], requires_partition_by=partitionings.Index(), preserves_partition_by=preserve_partition ) ) def _idxmaxmin_helper(self, op, **kwargs): if op == 'idxmax': func = pd.DataFrame.idxmax elif op == 'idxmin': func = pd.DataFrame.idxmin else: raise ValueError("op must be one of ('idxmax', 'idxmin'). " f"got {op!r}.") axis = kwargs.get('axis', 0) index_dtype = self._expr.proxy().index.dtype columns_dtype = self._expr.proxy().columns.dtype def compute_idx(df): indexes = func(df, **kwargs).unique() if pd.isna(indexes).any(): return df else: return df.loc[indexes] if axis in ('index', 0): requires_partition = partitionings.Singleton() proxy_index = pd.Index([], dtype=columns_dtype) proxy = pd.Series([], index=proxy_index, dtype=index_dtype) partition_proxy = self._expr.proxy().copy() idx_per_partition = expressions.ComputedExpression( 'idx-per-partition', compute_idx, [self._expr], proxy=partition_proxy, requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary() ) elif axis in ('columns', 1): requires_partition = partitionings.Index() proxy_index = pd.Index([], dtype=index_dtype) proxy = pd.Series([], index=proxy_index, dtype=columns_dtype) idx_per_partition = self._expr else: raise ValueError("axis must be one of ('index', 0, 'columns', 1). " f"got {axis!r}.") with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'idx', lambda df: func(df, **kwargs), [idx_per_partition], proxy=proxy, requires_partition_by=requires_partition, preserves_partition_by=partitionings.Singleton() ) ) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def idxmin(self, **kwargs): return self._idxmaxmin_helper('idxmin', **kwargs) @frame_base.with_docs_from(pd.DataFrame) @frame_base.args_to_kwargs(pd.DataFrame) @frame_base.populate_defaults(pd.DataFrame) def idxmax(self, **kwargs): return self._idxmaxmin_helper('idxmax', **kwargs) for io_func in dir(io): if io_func.startswith('to_'): setattr(DeferredDataFrame, io_func, getattr(io, io_func)) setattr(DeferredSeries, io_func, getattr(io, io_func)) for meth in ('filter', ): setattr(DeferredDataFrame, meth, frame_base._elementwise_method(meth, base=pd.DataFrame)) @populate_not_implemented(DataFrameGroupBy) class DeferredGroupBy(frame_base.DeferredFrame): def __init__(self, expr, kwargs, ungrouped: expressions.Expression[pd.core.generic.NDFrame], ungrouped_with_index: expressions.Expression[pd.core.generic.NDFrame], # pylint: disable=line-too-long grouping_columns, grouping_indexes, group_keys, projection=None): """This object represents the result of:: ungrouped.groupby(level=[grouping_indexes + grouping_columns], **kwargs)[projection] :param expr: An expression to compute a pandas GroupBy object. Convenient for unliftable aggregations. :param ungrouped: An expression to compute the DataFrame pre-grouping, the (Multi)Index contains only the grouping columns/indexes. :param ungrouped_with_index: Same as ungrouped, except the index includes all of the original indexes as well as any grouping columns. This is important for operations that expose the original index, e.g. .apply(), but we only use it when necessary to avoid unnessary data transfer and GBKs. :param grouping_columns: list of column labels that were in the original groupby(..) ``by`` parameter. Only relevant for grouped DataFrames. :param grouping_indexes: list of index names (or index level numbers) to be grouped. :param kwargs: Keywords args passed to the original groupby(..) call.""" super().__init__(expr) self._ungrouped = ungrouped self._ungrouped_with_index = ungrouped_with_index self._projection = projection self._grouping_columns = grouping_columns self._grouping_indexes = grouping_indexes self._group_keys = group_keys self._kwargs = kwargs if (self._kwargs.get('dropna', True) is False and self._ungrouped.proxy().index.nlevels > 1): raise NotImplementedError( "dropna=False does not work as intended in the Beam DataFrame API " "when grouping on multiple columns or indexes (See " "https://github.com/apache/beam/issues/21014).") def __getattr__(self, name): return DeferredGroupBy( expressions.ComputedExpression( 'groupby_project', lambda gb: getattr(gb, name), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary()), self._kwargs, self._ungrouped, self._ungrouped_with_index, self._grouping_columns, self._grouping_indexes, self._group_keys, projection=name) def __getitem__(self, name): return DeferredGroupBy( expressions.ComputedExpression( 'groupby_project', lambda gb: gb[name], [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary()), self._kwargs, self._ungrouped, self._ungrouped_with_index, self._grouping_columns, self._grouping_indexes, self._group_keys, projection=name) @frame_base.with_docs_from(DataFrameGroupBy) def agg(self, fn=None, *args, **kwargs): if fn is None: return _agg_with_no_function(self, *args, **kwargs) else: return _handle_agg_function(self, fn, "agg", *args, **kwargs) @property def ndim(self): return self._expr.proxy().ndim @frame_base.with_docs_from(DataFrameGroupBy) def apply(self, func, *args, **kwargs): """Note that ``func`` will be called once during pipeline construction time with an empty pandas object, so take care if ``func`` has a side effect. When called with an empty pandas object, ``func`` is expected to return an object of the same type as what will be returned when the pipeline is processing actual data. If the result is a pandas object it should have the same type and name (for a Series) or column types and names (for a DataFrame) as the actual results. Note that in pandas, ``apply`` attempts to detect if the index is unmodified in ``func`` (indicating ``func`` is a transform) and drops the duplicate index in the output. To determine this, pandas tests the indexes for equality. However, Beam cannot do this since it is sensitive to the input data; instead this implementation tests if the indexes are equivalent with ``is``. See the `pandas 1.4.0 release notes `_ for a good explanation of the distinction between these approaches. In practice, this just means that in some cases the Beam result will have a duplicate index, whereas pandas would have dropped it.""" project = _maybe_project_func(self._projection) grouping_indexes = self._grouping_indexes grouping_columns = self._grouping_columns group_keys = self._group_keys # Unfortunately pandas does not execute func to determine the right proxy. # We run user func on a proxy here to detect the return type and generate # the proxy. fn_input = project(self._ungrouped_with_index.proxy().reset_index( grouping_columns, drop=True)) result = func(fn_input) def index_to_arrays(index): return [index.get_level_values(level) for level in range(index.nlevels)] # By default do_apply will just call pandas apply() # We override it below if necessary do_apply = lambda gb: gb.apply(func, *args, **kwargs) if (isinstance(result, pd.core.generic.NDFrame) and result.index is fn_input.index): # Special case where apply fn is a transform # Note we trust that if the user fn produces a proxy with the identical # index, it will produce results with identical indexes at execution # time too proxy = result elif isinstance(result, pd.DataFrame): # apply fn is not a transform, we need to make sure the original index # values are prepended to the result's index proxy = result[:0] # First adjust proxy proxy.index = pd.MultiIndex.from_arrays( index_to_arrays(self._ungrouped.proxy().index) + index_to_arrays(proxy.index), names=self._ungrouped.proxy().index.names + proxy.index.names) # Then override do_apply function new_index_names = self._ungrouped.proxy().index.names if len(new_index_names) > 1: def add_key_index(key, df): # df is a dataframe or Series representing the result of func for # a single key # key is a tuple with the MultiIndex values for this key df.index = pd.MultiIndex.from_arrays( [[key[i]] * len(df) for i in range(len(new_index_names))] + index_to_arrays(df.index), names=new_index_names + df.index.names) return df else: def add_key_index(key, df): # df is a dataframe or Series representing the result of func for # a single key df.index = pd.MultiIndex.from_arrays( [[key] * len(df)] + index_to_arrays(df.index), names=new_index_names + df.index.names) return df do_apply = lambda gb: pd.concat([ add_key_index(k, func(gb.get_group(k), *args, **kwargs)) for k in gb.groups.keys()]) elif isinstance(result, pd.Series): if isinstance(fn_input, pd.DataFrame): # DataFrameGroupBy # In this case pandas transposes the Series result, s.t. the Series # index values are the columns, and the grouping keys are the new index # values. dtype = pd.Series([result]).dtype proxy = pd.DataFrame(columns=result.index, dtype=result.dtype, index=self._ungrouped.proxy().index) elif isinstance(fn_input, pd.Series): # SeriesGroupBy # In this case the output is still a Series, but with an additional # index with the grouping keys. proxy = pd.Series(dtype=result.dtype, name=result.name, index=index_to_arrays(self._ungrouped.proxy().index) + index_to_arrays(result[:0].index)) else: # The user fn returns some non-pandas type. The expected result is a # Series where each element is the result of one user fn call. dtype = pd.Series([result]).dtype proxy = pd.Series([], dtype=dtype, index=self._ungrouped.proxy().index) def do_partition_apply(df): # Remove columns from index, we only needed them there for partitioning df = df.reset_index(grouping_columns, drop=True) gb = df.groupby(level=grouping_indexes or None, by=grouping_columns or None, group_keys=group_keys) gb = project(gb) return do_apply(gb) return DeferredDataFrame( expressions.ComputedExpression( 'apply', do_partition_apply, [self._ungrouped_with_index], proxy=proxy, requires_partition_by=partitionings.Index(grouping_indexes + grouping_columns), preserves_partition_by=partitionings.Index(grouping_indexes))) @frame_base.with_docs_from(DataFrameGroupBy) def transform(self, fn, *args, **kwargs): """Note that ``func`` will be called once during pipeline construction time with an empty pandas object, so take care if ``func`` has a side effect. When called with an empty pandas object, ``func`` is expected to return an object of the same type as what will be returned when the pipeline is processing actual data. The result should have the same type and name (for a Series) or column types and names (for a DataFrame) as the actual results.""" if not callable(fn): raise NotImplementedError( "String functions are not yet supported in transform.") if self._grouping_columns and not self._projection: grouping_columns = self._grouping_columns def fn_wrapper(x, *args, **kwargs): x = x.droplevel(grouping_columns) return fn(x, *args, **kwargs) else: fn_wrapper = fn project = _maybe_project_func(self._projection) group_keys = self._group_keys # pandas cannot execute fn to determine the right proxy. # We run user fn on a proxy here to detect the return type and generate the # proxy. result = fn_wrapper(project(self._ungrouped_with_index.proxy())) parent_frame = self._ungrouped.args()[0].proxy() if isinstance(result, pd.core.generic.NDFrame): proxy = result[:0] else: # The user fn returns some non-pandas type. The expected result is a # Series where each element is the result of one user fn call. dtype = pd.Series([result]).dtype proxy = pd.Series([], dtype=dtype, name=project(parent_frame).name) if not isinstance(self._projection, list): proxy.name = self._projection # The final result will have the original indexes proxy.index = parent_frame.index levels = self._grouping_indexes + self._grouping_columns return DeferredDataFrame( expressions.ComputedExpression( 'transform', lambda df: project( df.groupby(level=levels, group_keys=group_keys) ).transform( fn_wrapper, *args, **kwargs).droplevel(self._grouping_columns), [self._ungrouped_with_index], proxy=proxy, requires_partition_by=partitionings.Index(levels), preserves_partition_by=partitionings.Index(self._grouping_indexes))) @frame_base.with_docs_from(DataFrameGroupBy) def pipe(self, func, *args, **kwargs): if isinstance(func, tuple): func, data = func kwargs[data] = self return func(*args, **kwargs) return func(self, *args, **kwargs) @frame_base.with_docs_from(DataFrameGroupBy) def filter(self, func=None, dropna=True): if func is None or not callable(func): raise TypeError("func must be specified and it must be callable") def apply_fn(df): if func(df): return df elif not dropna: result = df.copy() result.iloc[:, :] = np.nan return result else: return df.iloc[:0] return self.apply(apply_fn).droplevel(self._grouping_columns) @property # type: ignore @frame_base.with_docs_from(DataFrameGroupBy) def dtypes(self): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'dtypes', lambda gb: gb.dtypes, [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary() ) ) if hasattr(DataFrameGroupBy, 'value_counts'): @frame_base.with_docs_from(DataFrameGroupBy) def value_counts(self, **kwargs): """ DataFrameGroupBy.value_counts() is the same as DataFrame.value_counts() """ return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'value_counts', lambda df: df.value_counts(**kwargs), [self._expr], preserves_partition_by=partitionings.Arbitrary(), requires_partition_by=partitionings.Arbitrary()) ) fillna = frame_base.wont_implement_method( DataFrameGroupBy, 'fillna', explanation=( "df.fillna() should be used instead. Only method=None is supported " "because other methods are order-sensitive. df.groupby(..).fillna() " "without a method is equivalent to df.fillna().")) ffill = frame_base.wont_implement_method(DataFrameGroupBy, 'ffill', reason="order-sensitive") bfill = frame_base.wont_implement_method(DataFrameGroupBy, 'bfill', reason="order-sensitive") pad = frame_base.wont_implement_method(DataFrameGroupBy, 'pad', reason="order-sensitive") backfill = frame_base.wont_implement_method(DataFrameGroupBy, 'backfill', reason="order-sensitive") aggregate = agg hist = frame_base.wont_implement_method(DataFrameGroupBy, 'hist', reason="plotting-tools") plot = frame_base.wont_implement_method(DataFrameGroupBy, 'plot', reason="plotting-tools") boxplot = frame_base.wont_implement_method(DataFrameGroupBy, 'boxplot', reason="plotting-tools") head = frame_base.wont_implement_method( DataFrameGroupBy, 'head', explanation=_PEEK_METHOD_EXPLANATION) tail = frame_base.wont_implement_method( DataFrameGroupBy, 'tail', explanation=_PEEK_METHOD_EXPLANATION) first = frame_base.not_implemented_method('first', base_type=DataFrameGroupBy) last = frame_base.not_implemented_method('last', base_type=DataFrameGroupBy) nth = property(frame_base.wont_implement_method( DataFrameGroupBy, 'nth', reason='order-sensitive')) cumcount = frame_base.wont_implement_method( DataFrameGroupBy, 'cumcount', reason='order-sensitive') cummax = frame_base.wont_implement_method( DataFrameGroupBy, 'cummax', reason='order-sensitive') cummin = frame_base.wont_implement_method( DataFrameGroupBy, 'cummin', reason='order-sensitive') cumsum = frame_base.wont_implement_method( DataFrameGroupBy, 'cumsum', reason='order-sensitive') cumprod = frame_base.wont_implement_method( DataFrameGroupBy, 'cumprod', reason='order-sensitive') diff = frame_base.wont_implement_method(DataFrameGroupBy, 'diff', reason='order-sensitive') shift = frame_base.wont_implement_method(DataFrameGroupBy, 'shift', reason='order-sensitive') pct_change = frame_base.wont_implement_method(DataFrameGroupBy, 'pct_change', reason='order-sensitive') ohlc = frame_base.wont_implement_method(DataFrameGroupBy, 'ohlc', reason='order-sensitive') # TODO(https://github.com/apache/beam/issues/20958): Consider allowing this # for categorical keys. __len__ = frame_base.wont_implement_method( DataFrameGroupBy, '__len__', reason="non-deferred-result") groups = property(frame_base.wont_implement_method( DataFrameGroupBy, 'groups', reason="non-deferred-result")) indices = property(frame_base.wont_implement_method( DataFrameGroupBy, 'indices', reason="non-deferred-result")) resample = frame_base.wont_implement_method( DataFrameGroupBy, 'resample', reason='event-time-semantics') rolling = frame_base.wont_implement_method( DataFrameGroupBy, 'rolling', reason='event-time-semantics') ewm = frame_base.wont_implement_method( DataFrameGroupBy, 'ewm', reason="event-time-semantics") expanding = frame_base.wont_implement_method( DataFrameGroupBy, 'expanding', reason="event-time-semantics") tshift = frame_base.wont_implement_method( DataFrameGroupBy, 'tshift', reason="deprecated") def _maybe_project_func(projection: Optional[list[str]]): """ Returns identity func if projection is empty or None, else returns a function that projects the specified columns. """ if projection: return lambda df: df[projection] else: return lambda x: x def _liftable_agg(meth, postagg_meth=None): agg_name, _ = frame_base.name_and_func(meth) if postagg_meth is None: post_agg_name = agg_name else: post_agg_name, _ = frame_base.name_and_func(postagg_meth) @frame_base.with_docs_from(DataFrameGroupBy, name=agg_name) def wrapper(self, *args, **kwargs): assert isinstance(self, DeferredGroupBy) if 'min_count' in kwargs: return _unliftable_agg(meth)(self, *args, **kwargs) to_group = self._ungrouped.proxy().index is_categorical_grouping = any( isinstance(to_group.get_level_values(i).dtype, pd.CategoricalDtype) for i in self._grouping_indexes) groupby_kwargs = self._kwargs group_keys = self._group_keys # Don't include un-observed categorical values in the preagg preagg_groupby_kwargs = groupby_kwargs.copy() preagg_groupby_kwargs['observed'] = True project = _maybe_project_func(self._projection) pre_agg = expressions.ComputedExpression( 'pre_combine_' + agg_name, lambda df: getattr( project( df.groupby(level=list(range(df.index.nlevels)), group_keys=group_keys, **preagg_groupby_kwargs) ), agg_name)(**kwargs), [self._ungrouped], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary()) post_agg = expressions.ComputedExpression( 'post_combine_' + post_agg_name, lambda df: getattr( df.groupby(level=list(range(df.index.nlevels)), group_keys=group_keys, **groupby_kwargs), post_agg_name)(**kwargs), [pre_agg], requires_partition_by=(partitionings.Singleton(reason=( "Aggregations grouped by a categorical column are not currently " "parallelizable (https://github.com/apache/beam/issues/21827)." )) if is_categorical_grouping else partitionings.Index()), preserves_partition_by=partitionings.Arbitrary()) return frame_base.DeferredFrame.wrap(post_agg) return wrapper def _unliftable_agg(meth): agg_name, _ = frame_base.name_and_func(meth) @frame_base.with_docs_from(DataFrameGroupBy, name=agg_name) def wrapper(self, *args, **kwargs): assert isinstance(self, DeferredGroupBy) to_group = self._ungrouped.proxy().index group_keys = self._group_keys is_categorical_grouping = any( isinstance(to_group.get_level_values(i).dtype, pd.CategoricalDtype) for i in self._grouping_indexes) groupby_kwargs = self._kwargs project = _maybe_project_func(self._projection) post_agg = expressions.ComputedExpression( agg_name, lambda df: getattr(project( df.groupby(level=list(range(df.index.nlevels)), group_keys=group_keys, **groupby_kwargs), ), agg_name)(**kwargs), [self._ungrouped], requires_partition_by=(partitionings.Singleton(reason=( "Aggregations grouped by a categorical column are not currently " "parallelizable (https://github.com/apache/beam/issues/21827)." )) if is_categorical_grouping else partitionings.Index()), # Some aggregation methods (e.g. corr/cov) add additional index levels. # We only preserve the ones that existed _before_ the groupby. preserves_partition_by=partitionings.Index( list(range(self._ungrouped.proxy().index.nlevels)))) return frame_base.DeferredFrame.wrap(post_agg) return wrapper for meth in LIFTABLE_AGGREGATIONS: setattr(DeferredGroupBy, meth, _liftable_agg(meth)) for meth in LIFTABLE_WITH_SUM_AGGREGATIONS: setattr(DeferredGroupBy, meth, _liftable_agg(meth, postagg_meth='sum')) for meth in UNLIFTABLE_AGGREGATIONS: if meth in ('kurt', 'kurtosis'): # pandas doesn't currently allow kurtosis on GroupBy: # https://github.com/pandas-dev/pandas/issues/40139 continue setattr(DeferredGroupBy, meth, _unliftable_agg(meth)) def _check_str_or_np_builtin(agg_func, func_list): return agg_func in func_list or ( getattr(agg_func, '__name__', None) in func_list and agg_func.__module__ in ('numpy', 'builtins')) def _agg_with_no_function(gb, *args, **kwargs): """ Applies aggregation functions to the grouped data based on the provided keyword arguments and combines the results into a single DataFrame. Args: gb: The groupby instance (DeferredGroupBy). *args: Additional positional arguments passed to the aggregation funcs. **kwargs: A dictionary where each key is the column name to aggregate, the value is a tuple containing the input column name and the aggregation function to apply. Returns: DeferredDataFrame: A DataFrame that contains the aggregated results of all specified columns. Raises: ValueError: If no aggregation functions are provided in the `kwargs`. NotImplementedError: If the aggregation function type is unsupported. """ if not kwargs: raise ValueError("No aggregation functions specified") # Handle dictionary-like input for aggregation. result_columns, result_frames = [], [] for col_name, (input_col, agg_fn) in kwargs.items(): frame = _handle_agg_function( gb[input_col], agg_fn, f"agg_{col_name}", *args ) result_frames.append(frame) result_columns.append(col_name) # Combine all the resulting DeferredDataFrames into a single DataFrame. return DeferredDataFrame( expressions.ComputedExpression( "agg", lambda *results: pd.concat(results, axis=1, keys=result_columns), [frame._expr for frame in result_frames], requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Singleton(), ) ) def _handle_agg_function(gb, agg_func, agg_name, *args, **kwargs): """ Handles the aggregation logic based on the function type passed. Args: gb: The groupby instance (DeferredGroupBy). agg_name: The name/label of the aggregation function. fn: The aggregation function to apply. *args: Additional arguments to pass to the aggregation function. **kwargs: Keyword arguments to pass to the aggregation function. Returns: A DeferredDataFrame or the result of the aggregation function. Raises: NotImplementedError: If the aggregation function type is unsupported. """ if _is_associative(agg_func): return _liftable_agg(agg_func)(gb, *args, **kwargs) elif _is_liftable_with_sum(agg_func): return _liftable_agg(agg_func, postagg_meth='sum')(gb, *args, **kwargs) elif _is_unliftable(agg_func): return _unliftable_agg(agg_func)(gb, *args, **kwargs) elif callable(agg_func): return DeferredDataFrame( expressions.ComputedExpression( agg_name, lambda gb_val: gb_val.agg(agg_func, *args, **kwargs), [gb._expr], requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Singleton())) else: raise NotImplementedError(f"GroupBy.agg(func={agg_func!r})") def _is_associative(agg_func): return _check_str_or_np_builtin(agg_func, LIFTABLE_AGGREGATIONS) def _is_liftable_with_sum(agg_func): return _check_str_or_np_builtin(agg_func, LIFTABLE_WITH_SUM_AGGREGATIONS) def _is_unliftable(agg_func): return _check_str_or_np_builtin(agg_func, UNLIFTABLE_AGGREGATIONS) NUMERIC_AGGREGATIONS = ['max', 'min', 'prod', 'sum', 'mean', 'median', 'std', 'var', 'sem', 'skew', 'kurt', 'kurtosis'] # mad was removed in Pandas 2.0. if PD_VERSION < (2, 0): NUMERIC_AGGREGATIONS.append('mad') def _is_numeric(agg_func): return _check_str_or_np_builtin(agg_func, NUMERIC_AGGREGATIONS) @populate_not_implemented(DataFrameGroupBy) class _DeferredGroupByCols(frame_base.DeferredFrame): # It's not clear that all of these make sense in Pandas either... agg = aggregate = frame_base._elementwise_method('agg', base=DataFrameGroupBy) any = frame_base._elementwise_method('any', base=DataFrameGroupBy) all = frame_base._elementwise_method('all', base=DataFrameGroupBy) boxplot = frame_base.wont_implement_method( DataFrameGroupBy, 'boxplot', reason="plotting-tools") describe = frame_base.not_implemented_method('describe', base_type=DataFrameGroupBy) diff = frame_base._elementwise_method('diff', base=DataFrameGroupBy) fillna = frame_base._elementwise_method('fillna', base=DataFrameGroupBy) filter = frame_base._elementwise_method('filter', base=DataFrameGroupBy) first = frame_base._elementwise_method('first', base=DataFrameGroupBy) get_group = frame_base._elementwise_method('get_group', base=DataFrameGroupBy) head = frame_base.wont_implement_method( DataFrameGroupBy, 'head', explanation=_PEEK_METHOD_EXPLANATION) hist = frame_base.wont_implement_method( DataFrameGroupBy, 'hist', reason="plotting-tools") idxmax = frame_base._elementwise_method('idxmax', base=DataFrameGroupBy) idxmin = frame_base._elementwise_method('idxmin', base=DataFrameGroupBy) last = frame_base._elementwise_method('last', base=DataFrameGroupBy) max = frame_base._elementwise_method('max', base=DataFrameGroupBy) mean = frame_base._elementwise_method('mean', base=DataFrameGroupBy) median = frame_base._elementwise_method('median', base=DataFrameGroupBy) min = frame_base._elementwise_method('min', base=DataFrameGroupBy) nunique = frame_base._elementwise_method('nunique', base=DataFrameGroupBy) plot = frame_base.wont_implement_method( DataFrameGroupBy, 'plot', reason="plotting-tools") prod = frame_base._elementwise_method('prod', base=DataFrameGroupBy) quantile = frame_base._elementwise_method('quantile', base=DataFrameGroupBy) shift = frame_base._elementwise_method('shift', base=DataFrameGroupBy) size = frame_base._elementwise_method('size', base=DataFrameGroupBy) skew = frame_base._elementwise_method('skew', base=DataFrameGroupBy) std = frame_base._elementwise_method('std', base=DataFrameGroupBy) sum = frame_base._elementwise_method('sum', base=DataFrameGroupBy) tail = frame_base.wont_implement_method( DataFrameGroupBy, 'tail', explanation=_PEEK_METHOD_EXPLANATION) take = frame_base.wont_implement_method( DataFrameGroupBy, 'take', reason='deprecated') var = frame_base._elementwise_method('var', base=DataFrameGroupBy) # These already deprecated methods were removed in Pandas 2.0 if PD_VERSION < (2, 0): mad = frame_base._elementwise_method('mad', base=DataFrameGroupBy) tshift = frame_base._elementwise_method('tshift', base=DataFrameGroupBy) @property # type: ignore @frame_base.with_docs_from(DataFrameGroupBy) def groups(self): return self._expr.proxy().groups @property # type: ignore @frame_base.with_docs_from(DataFrameGroupBy) def indices(self): return self._expr.proxy().indices @property # type: ignore @frame_base.with_docs_from(DataFrameGroupBy) def ndim(self): return self._expr.proxy().ndim @property # type: ignore @frame_base.with_docs_from(DataFrameGroupBy) def ngroups(self): return self._expr.proxy().ngroups @populate_not_implemented(pd.core.indexes.base.Index) class _DeferredIndex(object): def __init__(self, frame): self._frame = frame @property def names(self): return self._frame._expr.proxy().index.names @names.setter def names(self, value): def set_index_names(df): df = df.copy() df.index.names = value return df self._frame._expr = expressions.ComputedExpression( 'set_index_names', set_index_names, [self._frame._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary()) @property def name(self): return self._frame._expr.proxy().index.name @name.setter def name(self, value): self.names = [value] @property def ndim(self): return self._frame._expr.proxy().index.ndim @property def dtype(self): return self._frame._expr.proxy().index.dtype @property def nlevels(self): return self._frame._expr.proxy().index.nlevels def __getattr__(self, name): raise NotImplementedError('index.%s' % name) @populate_not_implemented(pd.core.indexing._LocIndexer) class _DeferredLoc(object): def __init__(self, frame): self._frame = frame def __getitem__(self, key): if isinstance(key, tuple): rows, cols = key return self[rows][cols] elif isinstance(key, list) and key and isinstance(key[0], bool): # Aligned by numerical key. raise NotImplementedError(type(key)) elif isinstance(key, list): # Select rows, but behaves poorly on missing values. raise NotImplementedError(type(key)) elif isinstance(key, slice): args = [self._frame._expr] func = lambda df: df.loc[key] elif isinstance(key, frame_base.DeferredFrame): func = lambda df, key: df.loc[key] if pd.core.dtypes.common.is_bool_dtype(key._expr.proxy()): # Boolean indexer, just pass it in as-is args = [self._frame._expr, key._expr] else: # Likely a DeferredSeries of labels, overwrite the key's index with it's # values so we can colocate them with the labels they're selecting def data_to_index(s): s = s.copy() s.index = s return s reindexed_expr = expressions.ComputedExpression( 'data_to_index', data_to_index, [key._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton(), ) args = [self._frame._expr, reindexed_expr] elif callable(key): def checked_callable_key(df): computed_index = key(df) if isinstance(computed_index, tuple): row_index, _ = computed_index else: row_index = computed_index if isinstance(row_index, list) and row_index and isinstance( row_index[0], bool): raise NotImplementedError(type(row_index)) elif not isinstance(row_index, (slice, pd.Series)): raise NotImplementedError(type(row_index)) return computed_index args = [self._frame._expr] func = lambda df: df.loc[checked_callable_key] else: raise NotImplementedError(type(key)) return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'loc', func, args, requires_partition_by=( partitionings.JoinIndex() if len(args) > 1 else partitionings.Arbitrary()), preserves_partition_by=partitionings.Arbitrary())) __setitem__ = frame_base.not_implemented_method( 'loc.setitem', base_type=pd.core.indexing._LocIndexer) @populate_not_implemented(pd.core.indexing._iLocIndexer) class _DeferredILoc(object): def __init__(self, frame): self._frame = frame def __getitem__(self, index): if isinstance(index, tuple): rows, _ = index if rows != slice(None, None, None): raise frame_base.WontImplementError( "Using iloc to select rows is not supported because it's " "position-based indexing is sensitive to the order of the data.", reason="order-sensitive") return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'iloc', lambda df: df.iloc[index], [self._frame._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary())) else: raise frame_base.WontImplementError( "Using iloc to select rows is not supported because it's " "position-based indexing is sensitive to the order of the data.", reason="order-sensitive") def __setitem__(self, index, value): raise frame_base.WontImplementError( "Using iloc to mutate a frame is not supported because it's " "position-based indexing is sensitive to the order of the data.", reason="order-sensitive") class _DeferredStringMethods(frame_base.DeferredBase): @frame_base.with_docs_from(pd.Series.str) @frame_base.args_to_kwargs(pd.Series.str) @frame_base.populate_defaults(pd.Series.str) def cat(self, others, join, **kwargs): """If defined, ``others`` must be a :class:`DeferredSeries` or a ``list`` of ``DeferredSeries``.""" if others is None: # Concatenate series into a single String requires = partitionings.Singleton(reason=( "cat(others=None) concatenates all data in a Series into a single " "string, so it requires collecting all data on a single node." )) func = lambda df: df.str.cat(join=join, **kwargs) args = [self._expr] elif (isinstance(others, frame_base.DeferredBase) or (isinstance(others, list) and all(isinstance(other, frame_base.DeferredBase) for other in others))): if isinstance(others, frame_base.DeferredBase): others = [others] requires = partitionings.Index() def func(*args): return args[0].str.cat(others=args[1:], join=join, **kwargs) args = [self._expr] + [other._expr for other in others] else: raise frame_base.WontImplementError( "others must be None, DeferredSeries, or list[DeferredSeries] " f"(encountered {type(others)}). Other types are not supported " "because they make this operation sensitive to the order of the " "data.", reason="order-sensitive") return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'cat', func, args, requires_partition_by=requires, preserves_partition_by=partitionings.Arbitrary())) @frame_base.with_docs_from(pd.Series.str) @frame_base.args_to_kwargs(pd.Series.str) def repeat(self, repeats): """``repeats`` must be an ``int`` or a :class:`DeferredSeries`. Lists are not supported because they make this operation order-sensitive.""" if isinstance(repeats, int): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'repeat', lambda series: series.str.repeat(repeats), [self._expr], # TODO(https://github.com/apache/beam/issues/20573): Defer to # pandas to compute this proxy. Currently it incorrectly infers # dtype bool, may require upstream fix. proxy=self._expr.proxy(), requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary())) elif isinstance(repeats, frame_base.DeferredBase): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'repeat', lambda series, repeats_series: series.str.repeat(repeats_series), [self._expr, repeats._expr], # TODO(https://github.com/apache/beam/issues/20573): Defer to # pandas to compute this proxy. Currently it incorrectly infers # dtype bool, may require upstream fix. proxy=self._expr.proxy(), requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Arbitrary())) elif isinstance(repeats, list): raise frame_base.WontImplementError( "str.repeat(repeats=) repeats must be an int or a DeferredSeries. " "Lists are not supported because they make this operation sensitive " "to the order of the data.", reason="order-sensitive") else: raise TypeError("str.repeat(repeats=) value must be an int or a " f"DeferredSeries (encountered {type(repeats)}).") @frame_base.with_docs_from(pd.Series.str) @frame_base.args_to_kwargs(pd.Series.str) def get_dummies(self, **kwargs): """ Series must be categorical dtype. Please cast to ``CategoricalDtype`` to ensure correct categories. """ dtype = self._expr.proxy().dtype if not isinstance(dtype, pd.CategoricalDtype): raise frame_base.WontImplementError( "get_dummies() of non-categorical type is not supported because " "the type of the output column depends on the data. Please use " "pd.CategoricalDtype with explicit categories.", reason="non-deferred-columns") split_cats = [ cat.split(sep=kwargs.get('sep', '|')) for cat in dtype.categories ] # Adding the nan category because there could be the case that # the data includes NaNs, which is not valid to be casted as a Category, # but nevertheless would be broadcasted as a column in get_dummies() columns = sorted(set().union(*split_cats)) columns = columns + ['nan'] if 'nan' not in columns else columns proxy = pd.DataFrame(columns=columns).astype(int) return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'get_dummies', lambda series: pd.concat( [proxy, series.str.get_dummies(**kwargs)] ).fillna(value=0, method=None).astype('int64'), [self._expr], proxy=proxy, requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary())) def _split_helper(self, rsplit=False, **kwargs): expand = kwargs.get('expand', False) if not expand: # Not creating separate columns proxy = self._expr.proxy() if not rsplit: func = lambda s: pd.concat([proxy, s.str.split(**kwargs)]) else: func = lambda s: pd.concat([proxy, s.str.rsplit(**kwargs)]) else: # Creating separate columns, so must be more strict on dtype dtype = self._expr.proxy().dtype if not isinstance(dtype, pd.CategoricalDtype): method_name = 'rsplit' if rsplit else 'split' raise frame_base.WontImplementError( f"{method_name}() of non-categorical type is not supported because " "the type of the output column depends on the data. Please use " "pd.CategoricalDtype with explicit categories.", reason="non-deferred-columns") # Split the categories split_cats = dtype.categories.str.split(**kwargs) # Count the number of new columns to create for proxy max_splits = len(max(split_cats, key=len)) proxy = pd.DataFrame(columns=range(max_splits)) def func(s): if not rsplit: result = s.str.split(**kwargs) else: result = s.str.rsplit(**kwargs) result[~result.isna()].replace(np.nan, value=None) return result return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'split', func, [self._expr], proxy=proxy, requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary())) @frame_base.with_docs_from(pd.Series.str) @frame_base.args_to_kwargs(pd.Series.str) @frame_base.populate_defaults(pd.Series.str) def split(self, **kwargs): """ Like other non-deferred methods, dtype must be CategoricalDtype. One exception is when ``expand`` is ``False``. Because we are not creating new columns at construction time, dtype can be `str`. """ return self._split_helper(rsplit=False, **kwargs) @frame_base.with_docs_from(pd.Series.str) @frame_base.args_to_kwargs(pd.Series.str) @frame_base.populate_defaults(pd.Series.str) def rsplit(self, **kwargs): """ Like other non-deferred methods, dtype must be CategoricalDtype. One exception is when ``expand`` is ``False``. Because we are not creating new columns at construction time, dtype can be `str`. """ return self._split_helper(rsplit=True, **kwargs) ELEMENTWISE_STRING_METHODS = [ 'capitalize', 'casefold', 'center', 'contains', 'count', 'decode', 'encode', 'endswith', 'extract', 'find', 'findall', 'fullmatch', 'get', 'index', 'isalnum', 'isalpha', 'isdecimal', 'isdigit', 'islower', 'isnumeric', 'isspace', 'istitle', 'isupper', 'join', 'len', 'lfind', 'ljust', 'lower', 'lstrip', 'match', 'normalize', 'pad', 'partition', 'removeprefix', 'removesuffix', 'replace', 'rpartition', 'rfind', 'rindex', 'rjust', 'rstrip', 'slice', 'slice_replace', 'startswith', 'strip', 'swapcase', 'title', 'translate', 'upper', 'wrap', 'zfill', '__getitem__', ] NON_ELEMENTWISE_STRING_METHODS = [ 'extractall', ] def make_str_func(method): def func(df, *args, **kwargs): try: df_str = df.str except AttributeError: # If there's a non-string value in a Series passed to .str method, pandas # will generally just replace it with NaN in the result. However if # there are _only_ non-string values, pandas will raise: # # AttributeError: Can only use .str accessor with string values! # # This can happen to us at execution time if we split a partition that is # only non-strings. This branch just replaces all those values with NaN # in that case. return df.map(lambda _: np.nan) else: return getattr(df_str, method)(*args, **kwargs) return func for method in ELEMENTWISE_STRING_METHODS: if not hasattr(pd.Series.str, method): # older versions (1.0.x) don't support some of these methods continue setattr(_DeferredStringMethods, method, frame_base._elementwise_method(make_str_func(method), name=method, base=pd.Series.str)) for method in NON_ELEMENTWISE_STRING_METHODS: if not hasattr(pd.Series.str, method): # older versions (1.0.x) don't support some of these methods continue setattr(_DeferredStringMethods, method, frame_base._proxy_method( make_str_func(method), name=method, base=pd.Series.str, requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Singleton())) def make_cat_func(method): def func(df, *args, **kwargs): return getattr(df.cat, method)(*args, **kwargs) return func class _DeferredCategoricalMethods(frame_base.DeferredBase): @property # type: ignore @frame_base.with_docs_from(pd.core.arrays.categorical.CategoricalAccessor) def categories(self): return self._expr.proxy().cat.categories @property # type: ignore @frame_base.with_docs_from(pd.core.arrays.categorical.CategoricalAccessor) def ordered(self): return self._expr.proxy().cat.ordered @property # type: ignore @frame_base.with_docs_from(pd.core.arrays.categorical.CategoricalAccessor) def codes(self): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'codes', lambda s: s.cat.codes, [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary(), ) ) remove_unused_categories = frame_base.wont_implement_method( pd.core.arrays.categorical.CategoricalAccessor, 'remove_unused_categories', reason="non-deferred-columns") ELEMENTWISE_CATEGORICAL_METHODS = [ 'add_categories', 'as_ordered', 'as_unordered', 'remove_categories', 'rename_categories', 'reorder_categories', 'set_categories', ] for method in ELEMENTWISE_CATEGORICAL_METHODS: setattr(_DeferredCategoricalMethods, method, frame_base._elementwise_method( make_cat_func(method), name=method, base=pd.core.arrays.categorical.CategoricalAccessor)) class _DeferredDatetimeMethods(frame_base.DeferredBase): @property # type: ignore @frame_base.with_docs_from(pd.core.indexes.accessors.DatetimeProperties) def tz(self): return self._expr.proxy().dt.tz @property # type: ignore @frame_base.with_docs_from(pd.core.indexes.accessors.DatetimeProperties) def freq(self): return self._expr.proxy().dt.freq @frame_base.with_docs_from(pd.core.indexes.accessors.DatetimeProperties) def tz_localize(self, *args, ambiguous='infer', **kwargs): """``ambiguous`` cannot be set to ``"infer"`` as its semantics are order-sensitive. Similarly, specifying ``ambiguous`` as an :class:`~numpy.ndarray` is order-sensitive, but you can achieve similar functionality by specifying ``ambiguous`` as a Series.""" if isinstance(ambiguous, np.ndarray): raise frame_base.WontImplementError( "tz_localize(ambiguous=ndarray) is not supported because it makes " "this operation sensitive to the order of the data. Please use a " "DeferredSeries instead.", reason="order-sensitive") elif isinstance(ambiguous, frame_base.DeferredFrame): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'tz_localize', lambda s, ambiguous: s.dt.tz_localize(*args, ambiguous=ambiguous, **kwargs), [self._expr, ambiguous._expr], requires_partition_by=partitionings.Index(), preserves_partition_by=partitionings.Arbitrary())) elif ambiguous == 'infer': # infer attempts to infer based on the order of the timestamps raise frame_base.WontImplementError( f"tz_localize(ambiguous={ambiguous!r}) is not allowed because it " "makes this operation sensitive to the order of the data.", reason="order-sensitive") return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( 'tz_localize', lambda s: s.dt.tz_localize(*args, ambiguous=ambiguous, **kwargs), [self._expr], requires_partition_by=partitionings.Arbitrary(), preserves_partition_by=partitionings.Arbitrary())) to_period = frame_base.wont_implement_method( pd.core.indexes.accessors.DatetimeProperties, 'to_period', reason="event-time-semantics") to_pydatetime = frame_base.wont_implement_method( pd.core.indexes.accessors.DatetimeProperties, 'to_pydatetime', reason="non-deferred-result") to_pytimedelta = frame_base.wont_implement_method( pd.core.indexes.accessors.DatetimeProperties, 'to_pytimedelta', reason="non-deferred-result") def make_dt_property(method): def func(df): return getattr(df.dt, method) return func def make_dt_func(method): def func(df, *args, **kwargs): return getattr(df.dt, method)(*args, **kwargs) return func ELEMENTWISE_DATETIME_METHODS = [ 'ceil', 'day_name', 'month_name', 'floor', 'isocalendar', 'round', 'normalize', 'strftime', 'tz_convert', ] for method in ELEMENTWISE_DATETIME_METHODS: if not hasattr(pd.core.indexes.accessors.DatetimeProperties, method): # older versions (1.0.x) don't support some of these methods continue setattr(_DeferredDatetimeMethods, method, frame_base._elementwise_method( make_dt_func(method), name=method, base=pd.core.indexes.accessors.DatetimeProperties)) ELEMENTWISE_DATETIME_PROPERTIES = [ 'date', 'day', 'dayofweek', 'dayofyear', 'days_in_month', 'daysinmonth', 'hour', 'is_leap_year', 'is_month_end', 'is_month_start', 'is_quarter_end', 'is_quarter_start', 'is_year_end', 'is_year_start', 'microsecond', 'minute', 'month', 'nanosecond', 'quarter', 'second', 'time', 'timetz', 'weekday', 'year', ] # Pandas 2 removed these. if PD_VERSION < (2, 0): ELEMENTWISE_DATETIME_PROPERTIES += ['week', 'weekofyear'] for method in ELEMENTWISE_DATETIME_PROPERTIES: setattr(_DeferredDatetimeMethods, method, property(frame_base._elementwise_method( make_dt_property(method), name=method, base=pd.core.indexes.accessors.DatetimeProperties))) for base in ['add', 'sub', 'mul', 'div', 'truediv', 'floordiv', 'mod', 'divmod', 'pow', 'and', 'or']: for p in ['%s', 'r%s', '__%s__', '__r%s__']: # TODO: non-trivial level? name = p % base if hasattr(pd.Series, name): setattr( DeferredSeries, name, frame_base._elementwise_method(name, restrictions={'level': None}, base=pd.Series)) if hasattr(pd.DataFrame, name): setattr( DeferredDataFrame, name, frame_base._elementwise_method(name, restrictions={'level': None}, base=pd.DataFrame)) inplace_name = '__i%s__' % base if hasattr(pd.Series, inplace_name): setattr( DeferredSeries, inplace_name, frame_base._elementwise_method(inplace_name, inplace=True, base=pd.Series)) if hasattr(pd.DataFrame, inplace_name): setattr( DeferredDataFrame, inplace_name, frame_base._elementwise_method(inplace_name, inplace=True, base=pd.DataFrame)) # Allow dataframe | SchemaTransform def _create_maybe_elementwise_or(base): elementwise = frame_base._elementwise_method( '__or__', restrictions={'level': None}, base=base) def _maybe_elementwise_or(self, right): if isinstance(right, PTransform): return convert.to_dataframe(convert.to_pcollection(self) | right) else: return elementwise(self, right) return _maybe_elementwise_or DeferredSeries.__or__ = _create_maybe_elementwise_or(pd.Series) # type: ignore DeferredDataFrame.__or__ = _create_maybe_elementwise_or(pd.DataFrame) # type: ignore for name in ['lt', 'le', 'gt', 'ge', 'eq', 'ne']: for p in '%s', '__%s__': # Note that non-underscore name is used for both as the __xxx__ methods are # order-sensitive. setattr(DeferredSeries, p % name, frame_base._elementwise_method(name, base=pd.Series)) setattr(DeferredDataFrame, p % name, frame_base._elementwise_method(name, base=pd.DataFrame)) for name in ['__neg__', '__pos__', '__invert__']: setattr(DeferredSeries, name, frame_base._elementwise_method(name, base=pd.Series)) setattr(DeferredDataFrame, name, frame_base._elementwise_method(name, base=pd.DataFrame)) DeferredSeries.multiply = DeferredSeries.mul # type: ignore DeferredDataFrame.multiply = DeferredDataFrame.mul # type: ignore DeferredSeries.subtract = DeferredSeries.sub # type: ignore DeferredDataFrame.subtract = DeferredDataFrame.sub # type: ignore DeferredSeries.divide = DeferredSeries.div # type: ignore DeferredDataFrame.divide = DeferredDataFrame.div # type: ignore def _slice_parts(s): yield s.start yield s.stop yield s.step def _is_null_slice(s): return isinstance(s, slice) and all(x is None for x in _slice_parts(s)) def _is_integer_slice(s): return isinstance(s, slice) and all( x is None or isinstance(x, int) for x in _slice_parts(s)) and not _is_null_slice(s) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/frames_test.py0000644000175100017510000035452515206113032022505 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import re import sys import unittest import warnings import numpy as np import pandas as pd import pytest from parameterized import parameterized import apache_beam as beam from apache_beam.dataframe import expressions from apache_beam.dataframe import frame_base from apache_beam.dataframe import frames from apache_beam.dataframe.convert import to_dataframe from apache_beam.dataframe.doctests import teststring from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.interactive_runner import InteractiveRunner from apache_beam.runners.interactive.testing.mock_env import isolated_env # Get major, minor version PD_VERSION = tuple(map(int, pd.__version__.split('.')[0:2])) GROUPBY_DF = pd.DataFrame({ 'group': ['a' if i % 5 == 0 or i % 3 == 0 else 'b' for i in range(100)], 'foo': [None if i % 11 == 0 else i for i in range(100)], 'bar': [None if i % 7 == 0 else 99 - i for i in range(100)], 'baz': [None if i % 13 == 0 else i * 2 for i in range(100)], 'bool': [i % 17 == 0 for i in range(100)], 'str': [str(i) for i in range(100)], }) if PD_VERSION < (2, 0): # All these are things that are fixed in the Pandas 2 transition. pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") def _get_deferred_args(*args): return [ frame_base.DeferredFrame.wrap( expressions.ConstantExpression(arg, arg[0:0])) for arg in args ] class _AbstractFrameTest(unittest.TestCase): """Test sub-class with utilities for verifying DataFrame operations.""" def _run_error_test( self, func, *args, construction_time=True, distributed=True): """Verify that func(*args) raises the same exception in pandas and in Beam. Note that by default this only checks for exceptions that the Beam DataFrame API raises during expression generation (i.e. construction time). Exceptions raised while the pipeline is executing are less helpful, but are sometimes unavoidable (e.g. data validation exceptions), to check for these exceptions use construction_time=False.""" deferred_args = _get_deferred_args(*args) # Get expected error try: expected = func(*args) except Exception as e: expected_error = e else: raise AssertionError( "Expected an error, but executing with pandas successfully " f"returned:\n{expected}") # Get actual error if construction_time: try: _ = func(*deferred_args)._expr except Exception as e: actual = e else: raise AssertionError( f"Expected an error:\n{expected_error}\nbut Beam successfully " f"generated an expression.") else: # not construction_time # Check for an error raised during pipeline execution expr = func(*deferred_args)._expr session_type = ( expressions.PartitioningSession if distributed else expressions.Session) try: result = session_type({}).evaluate(expr) except Exception as e: actual = e else: raise AssertionError( f"Expected an error:\n{expected_error}\nbut Beam successfully " f"Computed the result:\n{result}.") # Verify if (not isinstance(actual, type(expected_error)) or str(expected_error) not in str(actual)): raise AssertionError( f'Expected {expected_error!r} to be raised, but got {actual!r}' ) from actual def _run_inplace_test(self, func, arg, **kwargs): """Verify an inplace operation performed by func. Checks that func performs the same inplace operation on arg, in pandas and in Beam.""" def wrapper(df): df = df.copy() func(df) return df self._run_test(wrapper, arg, **kwargs) def _run_test( self, func, *args, distributed=True, nonparallel=False, check_proxy=True, lenient_dtype_check=False): """Verify that func(*args) produces the same result in pandas and in Beam. Args: distributed (bool): Whether or not to use PartitioningSession to simulate parallel execution. nonparallel (bool): Whether or not this function contains a non-parallelizable operation. If True, the expression will be generated twice, once outside of an allow_non_parallel_operations block (to verify NonParallelOperation is raised), and again inside of an allow_non_parallel_operations block to actually generate an expression to verify. check_proxy (bool): Whether or not to check that the proxy of the generated expression matches the actual result, defaults to True. This option should NOT be set to False in tests added for new operations if at all possible. Instead make sure the new operation produces the correct proxy. This flag only exists as an escape hatch until existing failures can be addressed (https://github.com/apache/beam/issues/20926). lenient_dtype_check (bool): Whether or not to check that numeric columns are still numeric between actual and proxy. i.e. verify that they are at least int64 or float64, and not necessarily have the exact same dtype. This may need to be set to True for some non-deferred operations, where the dtype of the values in the proxy are not known ahead of time, causing int64 to float64 coercion issues. """ # Compute expected value expected = func(*args) # Compute actual value deferred_args = _get_deferred_args(*args) if nonparallel: # First run outside a nonparallel block to confirm this raises as expected with self.assertRaises(expressions.NonParallelOperation) as raised: func(*deferred_args) if raised.exception.msg.startswith( "Encountered non-parallelizable form of"): raise AssertionError( "Default NonParallelOperation raised, please specify a reason in " "the Singleton() partitioning requirement for this operation." ) from raised.exception # Re-run in an allow non parallel block to get an expression to verify with beam.dataframe.allow_non_parallel_operations(): expr = func(*deferred_args)._expr else: expr = func(*deferred_args)._expr # Compute the result of the generated expression session_type = ( expressions.PartitioningSession if distributed else expressions.Session) actual = session_type({}).evaluate(expr) # Verify if isinstance(expected, pd.core.generic.NDFrame): if distributed: if expected.index.is_unique: expected = expected.sort_index() actual = actual.sort_index() elif isinstance(expected, pd.Series): expected = expected.sort_values() actual = actual.sort_values() else: expected = expected.sort_values(list(expected.columns)) actual = actual.sort_values(list(actual.columns)) if isinstance(expected, pd.Series): if lenient_dtype_check: pd.testing.assert_series_equal( expected.astype('Float64'), actual.astype('Float64')) else: pd.testing.assert_series_equal(expected, actual) elif isinstance(expected, pd.DataFrame): if lenient_dtype_check: pd.testing.assert_frame_equal( expected.astype('Float64'), actual.astype('Float64')) else: pd.testing.assert_frame_equal(expected, actual) else: raise ValueError( f"Expected value is a {type(expected)}," "not a Series or DataFrame.") else: # Expectation is not a pandas object if isinstance(expected, float): if np.isnan(expected): cmp = np.isnan else: cmp = lambda x: np.isclose(expected, x) else: cmp = lambda x: x == expected self.assertTrue( cmp(actual), 'Expected:\n\n%r\n\nActual:\n\n%r' % (expected, actual)) if check_proxy: # Verify that the actual result agrees with the proxy proxy = expr.proxy() if type(actual) in (np.float32, np.float64): self.assertTrue(type(actual) == type(proxy) or np.isnan(proxy)) else: self.assertEqual(type(actual), type(proxy)) if isinstance(expected, pd.core.generic.NDFrame): if isinstance(expected, pd.Series): if lenient_dtype_check: self.assertEqual( actual.astype('Float64').dtype, proxy.astype('Float64').dtype) else: self.assertEqual(actual.dtype, proxy.dtype) self.assertEqual(actual.name, proxy.name) elif isinstance(expected, pd.DataFrame): if lenient_dtype_check: pd.testing.assert_series_equal( actual.astype('Float64').dtypes, proxy.astype('Float64').dtypes) else: pd.testing.assert_series_equal(actual.dtypes, proxy.dtypes) else: raise ValueError( f"Expected value is a {type(expected)}," "not a Series or DataFrame.") self.assertEqual(actual.index.names, proxy.index.names) for i in range(actual.index.nlevels): if lenient_dtype_check: self.assertEqual( actual.astype('Float64').index.get_level_values(i).dtype, proxy.astype('Float64').index.get_level_values(i).dtype) else: self.assertEqual( actual.index.get_level_values(i).dtype, proxy.index.get_level_values(i).dtype) class DeferredFrameTest(_AbstractFrameTest): """Miscellaneous tessts for DataFrame operations.""" def test_series_arithmetic(self): a = pd.Series([1, 2, 3]) b = pd.Series([100, 200, 300]) self._run_test(lambda a, b: a - 2 * b, a, b) self._run_test(lambda a, b: a.subtract(2).multiply(b).divide(a), a, b) def test_dataframe_arithmetic(self): df = pd.DataFrame({'a': [1, 2, 3], 'b': [100, 200, 300]}) df2 = pd.DataFrame({'a': [3000, 1000, 2000], 'b': [7, 11, 13]}) self._run_test(lambda df, df2: df - 2 * df2, df, df2) self._run_test( lambda df, df2: df.subtract(2).multiply(df2).divide(df), df, df2) @unittest.skipIf(PD_VERSION < (1, 3), "dropna=False is new in pandas 1.3") def test_value_counts_dropna_false(self): df = pd.DataFrame({ 'first_name': ['John', 'Anne', 'John', 'Beth'], 'middle_name': ['Smith', pd.NA, pd.NA, 'Louise'] }) # TODO(https://github.com/apache/beam/issues/21014): Remove the # assertRaises this when the underlying bug in # https://github.com/pandas-dev/pandas/issues/36470 is fixed. with self.assertRaises(NotImplementedError): self._run_test(lambda df: df.value_counts(dropna=False), df) def test_get_column(self): df = pd.DataFrame({ 'Animal': ['Falcon', 'Falcon', 'Parrot', 'Parrot'], 'Speed': [380., 370., 24., 26.] }) self._run_test(lambda df: df['Animal'], df) self._run_test(lambda df: df.Speed, df) self._run_test(lambda df: df.get('Animal'), df) self._run_test(lambda df: df.get('FOO', df.Animal), df) def test_series_xs(self): # pandas doctests only verify DataFrame.xs, here we verify Series.xs as well d = { 'num_legs': [4, 4, 2, 2], 'num_wings': [0, 0, 2, 2], 'class': ['mammal', 'mammal', 'mammal', 'bird'], 'animal': ['cat', 'dog', 'bat', 'penguin'], 'locomotion': ['walks', 'walks', 'flies', 'walks'] } df = pd.DataFrame(data=d) df = df.set_index(['class', 'animal', 'locomotion']) self._run_test(lambda df: df.num_legs.xs('mammal'), df) self._run_test(lambda df: df.num_legs.xs(('mammal', 'dog')), df) self._run_test(lambda df: df.num_legs.xs('cat', level=1), df) self._run_test( lambda df: df.num_legs.xs(('bird', 'walks'), level=[0, 'locomotion']), df) def test_dataframe_xs(self): # Test cases reported in BEAM-13421 df = pd.DataFrame( np.array([ ['state', 'day1', 12], ['state', 'day1', 1], ['state', 'day2', 14], ['county', 'day1', 9], ]), columns=['provider', 'time', 'value']) self._run_test(lambda df: df.xs('state'), df.set_index(['provider'])) self._run_test( lambda df: df.xs('state'), df.set_index(['provider', 'time'])) def test_set_column(self): def new_column(df): df['NewCol'] = df['Speed'] df = pd.DataFrame({ 'Animal': ['Falcon', 'Falcon', 'Parrot', 'Parrot'], 'Speed': [380., 370., 24., 26.] }) self._run_inplace_test(new_column, df) def test_set_column_from_index(self): def new_column(df): df['NewCol'] = df.index df = pd.DataFrame({ 'Animal': ['Falcon', 'Falcon', 'Parrot', 'Parrot'], 'Speed': [380., 370., 24., 26.] }) self._run_inplace_test(new_column, df) def test_tz_with_utc_zone_set_explicitly(self): test = """ >>> s = pd.Series(["1/1/2020 10:00:00+00:00", "2/1/2020 11:00:00+03:00"]) >>> s = pd.to_datetime(s, utc=True) >>> s 0 2020-01-01 10:00:00+00:00 1 2020-02-01 08:00:00+00:00 dtype: datetime64[ns, UTC] >>> s.dt.tz datetime.timezone.utc """ teststring(test) def test_tz_localize_ambiguous_series(self): # This replicates a tz_localize doctest: # s.tz_localize('CET', ambiguous=np.array([True, True, False])) # But using a DeferredSeries instead of a np array s = pd.Series( range(3), index=pd.DatetimeIndex([ '2018-10-28 01:20:00', '2018-10-28 02:36:00', '2018-10-28 03:46:00' ])) ambiguous = pd.Series([True, True, False], index=s.index) self._run_test( lambda s, ambiguous: s.tz_localize('CET', ambiguous=ambiguous), s, ambiguous) def test_tz_convert(self): # This replicates a tz_localize doctest: # s.tz_localize('CET', ambiguous=np.array([True, True, False])) # But using a DeferredSeries instead of a np array s = pd.Series( range(3), index=pd.DatetimeIndex([ '2018-10-27 01:20:00', '2018-10-27 02:36:00', '2018-10-27 03:46:00' ], tz='Europe/Berlin')) self._run_test(lambda s: s.tz_convert('America/Los_Angeles'), s) def test_sort_index_columns(self): df = pd.DataFrame({ 'c': range(10), 'a': range(10), 'b': range(10), np.nan: range(10), }) self._run_test(lambda df: df.sort_index(axis=1), df) self._run_test(lambda df: df.sort_index(axis=1, ascending=False), df) self._run_test(lambda df: df.sort_index(axis=1, na_position='first'), df) def test_where_callable_args(self): df = pd.DataFrame( np.arange(10, dtype=np.int64).reshape(-1, 2), columns=['A', 'B']) self._run_test( lambda df: df.where(lambda df: df % 2 == 0, lambda df: df * 10), df) def test_where_concrete_args(self): df = pd.DataFrame( np.arange(10, dtype=np.int64).reshape(-1, 2), columns=['A', 'B']) self._run_test( lambda df: df.where( df % 2 == 0, pd.Series({ 'A': 123, 'B': 456 }), axis=1), df) def test_combine_dataframe(self): df = pd.DataFrame({'A': [0, 0], 'B': [4, 4]}) df2 = pd.DataFrame({'A': [1, 1], 'B': [3, 3]}) take_smaller = lambda s1, s2: s1 if s1.sum() < s2.sum() else s2 self._run_test( lambda df, df2: df.combine(df2, take_smaller), df, df2, nonparallel=True) def test_combine_dataframe_fill(self): df1 = pd.DataFrame({'A': [0, 0], 'B': [None, 4]}) df2 = pd.DataFrame({'A': [1, 1], 'B': [3, 3]}) take_smaller = lambda s1, s2: s1 if s1.sum() < s2.sum() else s2 self._run_test( lambda df1, df2: df1.combine(df2, take_smaller, fill_value=-5), df1, df2, nonparallel=True) def test_combine_Series(self): s1 = pd.Series({'falcon': 330.0, 'eagle': 160.0}) s2 = pd.Series({'falcon': 345.0, 'eagle': 200.0, 'duck': 30.0}) self._run_test( lambda s1, s2: s1.combine(s2, max), s1, s2, nonparallel=True, check_proxy=False) def test_combine_first_dataframe(self): df1 = pd.DataFrame({'A': [None, 0], 'B': [None, 4]}) df2 = pd.DataFrame({'A': [1, 1], 'B': [3, 3]}) self._run_test(lambda df1, df2: df1.combine_first(df2), df1, df2) def test_combine_first_series(self): s1 = pd.Series([1, np.nan]) s2 = pd.Series([3, 4]) self._run_test(lambda s1, s2: s1.combine_first(s2), s1, s2) def test_add_prefix(self): df = pd.DataFrame({'A': [1, 2, 3, 4], 'B': [3, 4, 5, 6]}) s = pd.Series([1, 2, 3, 4]) self._run_test(lambda df: df.add_prefix('col_'), df) self._run_test(lambda s: s.add_prefix('col_'), s) def test_add_suffix(self): df = pd.DataFrame({'A': [1, 2, 3, 4], 'B': [3, 4, 5, 6]}) s = pd.Series([1, 2, 3, 4]) self._run_test(lambda df: df.add_suffix('_col'), df) self._run_test(lambda s: s.add_prefix('_col'), s) def test_set_index(self): df = pd.DataFrame({ # [19, 18, ..] 'index1': reversed(range(20)), # [15, 16, .., 0, 1, .., 13, 14] 'index2': np.roll(range(20), 5), # ['', 'a', 'bb', ...] 'values': [chr(ord('a') + i) * i for i in range(20)], }) self._run_test(lambda df: df.set_index(['index1', 'index2']), df) self._run_test(lambda df: df.set_index(['index1', 'index2'], drop=True), df) self._run_test(lambda df: df.set_index('values'), df) self._run_error_test(lambda df: df.set_index('bad'), df) self._run_error_test( lambda df: df.set_index(['index2', 'bad', 'really_bad']), df) def test_set_axis(self): df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]}, index=['X', 'Y', 'Z']) self._run_test(lambda df: df.set_axis(['I', 'II'], axis='columns'), df) self._run_test(lambda df: df.set_axis([0, 1], axis=1), df) self._run_inplace_test( lambda df: df.set_axis(['i', 'ii'], axis='columns'), df) with self.assertRaises(NotImplementedError): self._run_test(lambda df: df.set_axis(['a', 'b', 'c'], axis='index'), df) self._run_test(lambda df: df.set_axis([0, 1, 2], axis=0), df) def test_series_set_axis(self): s = pd.Series(list(range(3)), index=['X', 'Y', 'Z']) with self.assertRaises(NotImplementedError): self._run_test(lambda s: s.set_axis(['a', 'b', 'c']), s) self._run_test(lambda s: s.set_axis([1, 2, 3]), s) def test_series_drop_ignore_errors(self): midx = pd.MultiIndex( levels=[['lama', 'cow', 'falcon'], ['speed', 'weight', 'length']], codes=[[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]]) s = pd.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=midx) # drop() requires singleton partitioning unless errors are ignored # Add some additional tests here to make sure the implementation works in # non-singleton partitioning. self._run_test(lambda s: s.drop('lama', level=0, errors='ignore'), s) self._run_test(lambda s: s.drop(('cow', 'speed'), errors='ignore'), s) self._run_test(lambda s: s.drop('falcon', level=0, errors='ignore'), s) def test_dataframe_drop_ignore_errors(self): midx = pd.MultiIndex( levels=[['lama', 'cow', 'falcon'], ['speed', 'weight', 'length']], codes=[[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]]) df = pd.DataFrame( index=midx, columns=['big', 'small'], data=[[45, 30], [200, 100], [1.5, 1], [30, 20], [250, 150], [1.5, 0.8], [320, 250], [1, 0.8], [0.3, 0.2]]) # drop() requires singleton partitioning unless errors are ignored # Add some additional tests here to make sure the implementation works in # non-singleton partitioning. self._run_test( lambda df: df.drop(index='lama', level=0, errors='ignore'), df) self._run_test( lambda df: df.drop(index=('cow', 'speed'), errors='ignore'), df) self._run_test( lambda df: df.drop(index='falcon', level=0, errors='ignore'), df) self._run_test( lambda df: df.drop(index='cow', columns='small', errors='ignore'), df) def test_merge(self): # This is from the pandas doctests, but fails due to re-indexing being # order-sensitive. df1 = pd.DataFrame({ 'lkey': ['foo', 'bar', 'baz', 'foo'], 'value': [1, 2, 3, 5] }) df2 = pd.DataFrame({ 'rkey': ['foo', 'bar', 'baz', 'foo'], 'value': [5, 6, 7, 8] }) self._run_test( lambda df1, df2: df1.merge(df2, left_on='lkey', right_on='rkey').rename( index=lambda x: '*'), df1, df2, nonparallel=True, check_proxy=False) self._run_test( lambda df1, df2: df1.merge( df2, left_on='lkey', right_on='rkey', suffixes=('_left', '_right')). rename(index=lambda x: '*'), df1, df2, nonparallel=True, check_proxy=False) def test_merge_left_join(self): # This is from the pandas doctests, but fails due to re-indexing being # order-sensitive. df1 = pd.DataFrame({'a': ['foo', 'bar'], 'b': [1, 2]}) df2 = pd.DataFrame({'a': ['foo', 'baz'], 'c': [3, 4]}) self._run_test( lambda df1, df2: df1.merge(df2, how='left', on='a').rename( index=lambda x: '*'), df1, df2, nonparallel=True, check_proxy=False) def test_merge_on_index(self): # This is from the pandas doctests, but fails due to re-indexing being # order-sensitive. df1 = pd.DataFrame({ 'lkey': ['foo', 'bar', 'baz', 'foo'], 'value': [1, 2, 3, 5] }).set_index('lkey') df2 = pd.DataFrame({ 'rkey': ['foo', 'bar', 'baz', 'foo'], 'value': [5, 6, 7, 8] }).set_index('rkey') self._run_test( lambda df1, df2: df1.merge(df2, left_index=True, right_index=True), df1, df2, check_proxy=False) def test_merge_same_key(self): df1 = pd.DataFrame({ 'key': ['foo', 'bar', 'baz', 'foo'], 'value': [1, 2, 3, 5] }) df2 = pd.DataFrame({ 'key': ['foo', 'bar', 'baz', 'foo'], 'value': [5, 6, 7, 8] }) self._run_test( lambda df1, df2: df1.merge(df2, on='key').rename(index=lambda x: '*'), df1, df2, nonparallel=True, check_proxy=False) self._run_test( lambda df1, df2: df1.merge(df2, on='key', suffixes=('_left', '_right')). rename(index=lambda x: '*'), df1, df2, nonparallel=True, check_proxy=False) def test_merge_same_key_doctest(self): df1 = pd.DataFrame({'a': ['foo', 'bar'], 'b': [1, 2]}) df2 = pd.DataFrame({'a': ['foo', 'baz'], 'c': [3, 4]}) self._run_test( lambda df1, df2: df1.merge(df2, how='left', on='a').rename( index=lambda x: '*'), df1, df2, nonparallel=True, check_proxy=False) # Test without specifying 'on' self._run_test( lambda df1, df2: df1.merge(df2, how='left').rename(index=lambda x: '*'), df1, df2, nonparallel=True, check_proxy=False) def test_merge_same_key_suffix_collision(self): df1 = pd.DataFrame({'a': ['foo', 'bar'], 'b': [1, 2], 'a_lsuffix': [5, 6]}) df2 = pd.DataFrame({'a': ['foo', 'baz'], 'c': [3, 4], 'a_rsuffix': [7, 8]}) self._run_test( lambda df1, df2: df1.merge( df2, how='left', on='a', suffixes=('_lsuffix', '_rsuffix')).rename( index=lambda x: '*'), df1, df2, nonparallel=True, check_proxy=False) # Test without specifying 'on' self._run_test( lambda df1, df2: df1.merge( df2, how='left', suffixes=('_lsuffix', '_rsuffix')).rename( index=lambda x: '*'), df1, df2, nonparallel=True, check_proxy=False) def test_swaplevel(self): df = pd.DataFrame( {"Grade": ["A", "B", "A", "C"]}, index=[ ["Final exam", "Final exam", "Coursework", "Coursework"], ["History", "Geography", "History", "Geography"], ["January", "February", "March", "April"], ]) self._run_test(lambda df: df.swaplevel(), df) def test_value_counts_with_nans(self): # similar to doctests that verify value_counts, but include nan values to # make sure we handle them correctly. df = pd.DataFrame({ 'num_legs': [2, 4, 4, 6, np.nan, np.nan], 'num_wings': [2, 0, 0, 0, np.nan, 2] }, index=['falcon', 'dog', 'cat', 'ant', 'car', 'plane']) self._run_test(lambda df: df.value_counts(), df) self._run_test(lambda df: df.value_counts(normalize=True), df) # Ensure we don't drop rows due to nan values in unused columns. self._run_test(lambda df: df.value_counts('num_wings'), df) if PD_VERSION >= (1, 3): # dropna=False is new in pandas 1.3 # TODO(https://github.com/apache/beam/issues/21014): Remove the # assertRaises this when the underlying bug in # https://github.com/pandas-dev/pandas/issues/36470 is fixed. with self.assertRaises(NotImplementedError): self._run_test(lambda df: df.value_counts(dropna=False), df) # Test the defaults. self._run_test(lambda df: df.num_wings.value_counts(), df) self._run_test(lambda df: df.num_wings.value_counts(normalize=True), df) self._run_test(lambda df: df.num_wings.value_counts(dropna=False), df) # Test the combination interactions. for normalize in (True, False): for dropna in (True, False): self._run_test( lambda df, dropna=dropna, normalize=normalize: df.num_wings. value_counts(dropna=dropna, normalize=normalize), df) def test_value_counts_does_not_support_sort(self): df = pd.DataFrame({ 'num_legs': [2, 4, 4, 6, np.nan, np.nan], 'num_wings': [2, 0, 0, 0, np.nan, 2] }, index=['falcon', 'dog', 'cat', 'ant', 'car', 'plane']) with self.assertRaisesRegex(frame_base.WontImplementError, r"value_counts\(sort\=True\)"): self._run_test(lambda df: df.value_counts(sort=True), df) with self.assertRaisesRegex(frame_base.WontImplementError, r"value_counts\(sort\=True\)"): self._run_test(lambda df: df.num_wings.value_counts(sort=True), df) def test_series_getitem(self): s = pd.Series([x**2 for x in range(10)]) self._run_test(lambda s: s[...], s) self._run_test(lambda s: s[:], s) self._run_test(lambda s: s[s < 10], s) self._run_test(lambda s: s[lambda s: s < 10], s) s.index = s.index.map(float) self._run_test(lambda s: s[1.5:6], s) def test_series_truncate(self): s = pd.Series(['a', 'b', 'c', 'd', 'e', 'f']) self._run_test(lambda s: s.truncate(before=1, after=3), s) def test_dataframe_truncate(self): df = pd.DataFrame({ 'C': list('abcde'), 'B': list('fghij'), 'A': list('klmno') }, index=[1, 2, 3, 4, 5]) self._run_test(lambda df: df.truncate(before=1, after=3), df) self._run_test(lambda df: df.truncate(before='A', after='B', axis=1), df) self._run_test(lambda df: df['A'].truncate(before=2, after=4), df) @parameterized.expand([ (pd.Series(range(10)), ), # unique (pd.Series(list(range(100)) + [0]), ), # non-unique int (pd.Series(list(range(100)) + [0]) / 100, ), # non-unique flt (pd.Series(['a', 'b', 'c', 'd']), ), # unique str (pd.Series(['a', 'b', 'a', 'c', 'd']), ), # non-unique str ]) def test_series_is_unique(self, series): self._run_test(lambda s: s.is_unique, series) @parameterized.expand([ (pd.Series(range(10)), ), # False (pd.Series([1, 2, np.nan, 3, np.nan]), ), # True (pd.Series(['a', 'b', 'c', 'd', 'e']), ), # False (pd.Series(['a', 'b', None, 'c', None]), ), # True ]) def test_series_hasnans(self, series): self._run_test(lambda s: s.hasnans, series) def test_dataframe_getitem(self): df = pd.DataFrame({'A': [x**2 for x in range(6)], 'B': list('abcdef')}) self._run_test(lambda df: df['A'], df) self._run_test(lambda df: df[['A', 'B']], df) self._run_test(lambda df: df[:], df) self._run_test(lambda df: df[df.A < 10], df) df.index = df.index.map(float) self._run_test(lambda df: df[1.5:4], df) def test_loc(self): dates = pd.date_range('1/1/2000', periods=8) # TODO(https://github.com/apache/beam/issues/20765): # We do not preserve the freq attribute on a DateTime index dates.freq = None df = pd.DataFrame( np.arange(32).reshape((8, 4)), index=dates, columns=['A', 'B', 'C', 'D']) self._run_test(lambda df: df.loc[:], df) self._run_test(lambda df: df.loc[:, 'A'], df) self._run_test(lambda df: df.loc[:dates[3]], df) self._run_test(lambda df: df.loc[df.A > 10], df) self._run_test(lambda df: df.loc[lambda df: df.A > 10], df) self._run_test(lambda df: df.C.loc[df.A > 10], df) self._run_test(lambda df, s: df.loc[s.loc[1:3]], df, pd.Series(dates)) @unittest.skipIf(PD_VERSION >= (2, 0), 'append removed in Pandas 2.0') def test_append_sort(self): # yapf: disable df1 = pd.DataFrame({'int': [1, 2, 3], 'str': ['a', 'b', 'c']}, columns=['int', 'str'], index=[1, 3, 5]) df2 = pd.DataFrame({'int': [4, 5, 6], 'str': ['d', 'e', 'f']}, columns=['str', 'int'], index=[2, 4, 6]) # yapf: enable self._run_test(lambda df1, df2: df1.append(df2, sort=True), df1, df2) self._run_test(lambda df1, df2: df1.append(df2, sort=False), df1, df2) self._run_test(lambda df1, df2: df2.append(df1, sort=True), df1, df2) self._run_test(lambda df1, df2: df2.append(df1, sort=False), df1, df2) def test_smallest_largest(self): df = pd.DataFrame({'A': [1, 1, 2, 2], 'B': [2, 3, 5, 7]}) self._run_test(lambda df: df.nlargest(1, 'A', keep='all'), df) self._run_test(lambda df: df.nsmallest(3, 'A', keep='all'), df) self._run_test(lambda df: df.nlargest(3, ['A', 'B'], keep='all'), df) def test_series_cov_corr(self): for s in [pd.Series([1, 2, 3]), pd.Series(range(100)), pd.Series([x**3 for x in range(-50, 50)])]: self._run_test(lambda s: s.std(), s) self._run_test(lambda s: s.var(), s) self._run_test(lambda s: s.corr(s), s) self._run_test(lambda s: s.corr(s + 1), s) self._run_test(lambda s: s.corr(s * s), s) self._run_test(lambda s: s.cov(s * s), s) self._run_test(lambda s: s.skew(), s) self._run_test(lambda s: s.kurtosis(), s) self._run_test(lambda s: s.kurt(), s) def test_dataframe_cov_corr(self): df = pd.DataFrame(np.random.randn(20, 3), columns=['a', 'b', 'c']) df.loc[df.index[:5], 'a'] = np.nan df.loc[df.index[5:10], 'b'] = np.nan self._run_test(lambda df: df.corr(), df) self._run_test(lambda df: df.cov(), df) self._run_test(lambda df: df.corr(min_periods=12), df) self._run_test(lambda df: df.cov(min_periods=12), df) self._run_test(lambda df: df.corrwith(df.a), df) self._run_test(lambda df: df[['a', 'b']].corrwith(df[['b', 'c']]), df) df2 = pd.DataFrame(np.random.randn(20, 3), columns=['a', 'b', 'c']) self._run_test( lambda df, df2: df.corrwith(df2, axis=1), df, df2, check_proxy=False) def test_corrwith_bad_axis(self): df = pd.DataFrame({'a': range(3), 'b': range(3, 6), 'c': range(6, 9)}) self._run_error_test(lambda df: df.corrwith(df.a, axis=2), df) self._run_error_test(lambda df: df.corrwith(df, axis=5), df) @unittest.skipIf(PD_VERSION < (1, 2), "na_action added in pandas 1.2.0") @pytest.mark.filterwarnings( "ignore:The default of observed=False is deprecated:FutureWarning") def test_applymap_na_action(self): # Replicates a doctest for na_action which is incompatible with # doctest framework df = pd.DataFrame([[pd.NA, 2.12], [3.356, 4.567]]) self._run_test( lambda df: df.applymap(lambda x: len(str(x)), na_action='ignore'), df, # TODO: generate proxy using naive type inference on fn check_proxy=False) @unittest.skipIf(PD_VERSION < (2, 1), "map added in 2.1.0") def test_map_na_action(self): # Replicates a doctest for na_action which is incompatible with # doctest framework df = pd.DataFrame([[pd.NA, 2.12], [3.356, 4.567]]) self._run_test( lambda df: df.map(lambda x: len(str(x)), na_action='ignore'), df, # TODO: generate proxy using naive type inference on fn check_proxy=False) def test_dataframe_eval_query(self): df = pd.DataFrame(np.random.randn(20, 3), columns=['a', 'b', 'c']) self._run_test(lambda df: df.eval('foo = a + b - c'), df) self._run_test(lambda df: df.query('a > b + c'), df) self._run_inplace_test(lambda df: df.eval('foo = a + b - c'), df) # Verify that attempting to access locals raises a useful error deferred_df = frame_base.DeferredFrame.wrap( expressions.ConstantExpression(df, df[0:0])) self.assertRaises( NotImplementedError, lambda: deferred_df.eval('foo = a + @b - c')) self.assertRaises( NotImplementedError, lambda: deferred_df.query('a > @b + c')) def test_index_name_assignment(self): df = pd.DataFrame({'a': ['foo', 'bar'], 'b': [1, 2]}) df = df.set_index(['a', 'b'], drop=False) def change_index_names(df): df.index.names = ['A', None] self._run_inplace_test(change_index_names, df) def test_quantile(self): df = pd.DataFrame( np.array([[1, 1], [2, 10], [3, 100], [4, 100]]), columns=['a', 'b']) self._run_test( lambda df: df.quantile(0.1, axis='columns'), df, check_proxy=False) self._run_test( lambda df: df.quantile(0.1, axis='columns'), df, check_proxy=False) with self.assertRaisesRegex(frame_base.WontImplementError, r"df\.quantile\(q=0\.1, axis='columns'\)"): self._run_test(lambda df: df.quantile([0.1, 0.5], axis='columns'), df) def test_dataframe_melt(self): df = pd.DataFrame({ 'A': { 0: 'a', 1: 'b', 2: 'c' }, 'B': { 0: 1, 1: 3, 2: 5 }, 'C': { 0: 2, 1: 4, 2: 6 } }) self._run_test( lambda df: df.melt(id_vars=['A'], value_vars=['B'], ignore_index=False), df) self._run_test( lambda df: df.melt( id_vars=['A'], value_vars=['B', 'C'], ignore_index=False), df) self._run_test( lambda df: df.melt( id_vars=['A'], value_vars=['B'], var_name='myVarname', value_name= 'myValname', ignore_index=False), df) self._run_test( lambda df: df.melt( id_vars=['A'], value_vars=['B', 'C'], ignore_index=False), df) df.columns = [list('ABC'), list('DEF')] self._run_test( lambda df: df.melt( col_level=0, id_vars=['A'], value_vars=['B'], ignore_index=False), df) self._run_test( lambda df: df.melt( id_vars=[('A', 'D')], value_vars=[('B', 'E')], ignore_index=False), df) def test_fillna_columns(self): df = pd.DataFrame( [[np.nan, 2, np.nan, 0], [3, 4, np.nan, 1], [np.nan, np.nan, np.nan, 5], [np.nan, 3, np.nan, 4], [3, np.nan, np.nan, 4]], columns=list('ABCD')) self._run_test(lambda df: df.fillna(method='ffill', axis='columns'), df) self._run_test( lambda df: df.fillna(method='ffill', axis='columns', limit=1), df) self._run_test( lambda df: df.fillna(method='bfill', axis='columns', limit=1), df) # Intended behavior is unclear here. See # https://github.com/pandas-dev/pandas/issues/40989 # self._run_test(lambda df: df.fillna(axis='columns', value=100, # limit=2), df) def test_dataframe_fillna_dataframe_as_value(self): df = pd.DataFrame([[np.nan, 2, np.nan, 0], [3, 4, np.nan, 1], [np.nan, np.nan, np.nan, 5], [np.nan, 3, np.nan, 4]], columns=list("ABCD")) df2 = pd.DataFrame(np.zeros((4, 4)), columns=list("ABCE")) self._run_test(lambda df, df2: df.fillna(df2), df, df2) def test_dataframe_fillna_series_as_value(self): df = pd.DataFrame([[np.nan, 2, np.nan, 0], [3, 4, np.nan, 1], [np.nan, np.nan, np.nan, 5], [np.nan, 3, np.nan, 4]], columns=list("ABCD")) s = pd.Series(range(4), index=list("ABCE")) self._run_test(lambda df, s: df.fillna(s), df, s) def test_series_fillna_series_as_value(self): df = pd.DataFrame([[np.nan, 2, np.nan, 0], [3, 4, np.nan, 1], [np.nan, np.nan, np.nan, 5], [np.nan, 3, np.nan, 4]], columns=list("ABCD")) df2 = pd.DataFrame(np.zeros((4, 4)), columns=list("ABCE")) self._run_test(lambda df, df2: df.A.fillna(df2.A), df, df2) def test_dataframe_column_fillna_constant_as_value(self): from apache_beam.dataframe import convert from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to with beam.Pipeline(None) as p: pcoll = ( p | beam.Create([1.0, np.nan, -1.0]) | beam.Select(x=lambda x: x)) df = convert.to_dataframe(pcoll) df_new = df['x'].fillna(0) assert_that(convert.to_pcollection(df_new), equal_to([1.0, 0.0, -1.0])) @unittest.skipIf(PD_VERSION >= (2, 0), 'append removed in Pandas 2.0') def test_append_verify_integrity(self): df1 = pd.DataFrame({'A': range(10), 'B': range(10)}, index=range(10)) df2 = pd.DataFrame({'A': range(10), 'B': range(10)}, index=range(9, 19)) self._run_error_test( lambda s1, s2: s1.append(s2, verify_integrity=True), df1['A'], df2['A'], construction_time=False) self._run_error_test( lambda df1, df2: df1.append(df2, verify_integrity=True), df1, df2, construction_time=False) def test_categorical_groupby(self): df = pd.DataFrame({'A': np.arange(6), 'B': list('aabbca')}) df['B'] = df['B'].astype(pd.CategoricalDtype(list('cab'))) df = df.set_index('B') # TODO(BEAM-11190): These aggregations can be done in index partitions, but # it will require a little more complex logic self._run_test( lambda df: df.groupby(level=0, observed=False).sum(), df, nonparallel=True) self._run_test( lambda df: df.groupby(level=0, observed=False).mean(), df, nonparallel=True) def test_astype_categorical(self): df = pd.DataFrame({'A': np.arange(6), 'B': list('aabbca')}) categorical_dtype = pd.CategoricalDtype(df.B.unique()) self._run_test(lambda df: df.B.astype(categorical_dtype), df) @unittest.skipIf( PD_VERSION < (1, 2), "DataFrame.unstack not supported in pandas <1.2.x") def test_astype_categorical_with_unstack(self): df = pd.DataFrame({ 'index1': ['one', 'one', 'two', 'two'], 'index2': ['a', 'b', 'a', 'b'], 'data': np.arange(1.0, 5.0), }) def with_categorical_index(df): df.index1 = df.index1.astype(pd.CategoricalDtype(['one', 'two'])) df.index2 = df.index2.astype(pd.CategoricalDtype(['a', 'b'])) df.set_index(['index1', 'index2'], drop=True) return df self._run_test( lambda df: with_categorical_index(df).unstack(level=-1), df, check_proxy=False) def test_dataframe_sum_nonnumeric_raises(self): # Attempting a numeric aggregation with the str column present should # raise, and suggest the numeric_only argument with self.assertRaisesRegex(frame_base.WontImplementError, 'numeric_only'): self._run_test(lambda df: df.sum(), GROUPBY_DF) # numeric_only=True should work self._run_test(lambda df: df.sum(numeric_only=True), GROUPBY_DF) # projecting only numeric columns should too self._run_test(lambda df: df[['foo', 'bar']].sum(), GROUPBY_DF) def test_insert(self): df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]}) self._run_inplace_test(lambda df: df.insert(1, 'C', df.A * 2), df) self._run_inplace_test( lambda df: df.insert(0, 'foo', pd.Series([8], index=[1])), df, check_proxy=False) self._run_inplace_test(lambda df: df.insert(2, 'bar', value='q'), df) def test_insert_does_not_support_list_value(self): df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]}) with self.assertRaisesRegex(frame_base.WontImplementError, r"insert\(value=list\)"): self._run_inplace_test(lambda df: df.insert(1, 'C', [7, 8, 9]), df) def test_drop_duplicates(self): df = pd.DataFrame({ 'brand': ['Yum Yum', 'Yum Yum', 'Indomie', 'Indomie', 'Indomie'], 'style': ['cup', 'cup', 'cup', 'pack', 'pack'], 'rating': [4, 4, 3.5, 15, 5] }) self._run_test(lambda df: df.drop_duplicates(keep=False), df) self._run_test( lambda df: df.drop_duplicates(subset=['brand'], keep=False), df) self._run_test( lambda df: df.drop_duplicates(subset=['brand', 'style'], keep=False), df) @parameterized.expand([ ( lambda base: base.from_dict({ 'col_1': [3, 2, 1, 0], 'col_2': ['a', 'b', 'c', 'd'] }), ), ( lambda base: base.from_dict({ 'row_1': [3, 2, 1, 0], 'row_2': ['a', 'b', 'c', 'd'] }, orient='index'), ), ( lambda base: base.from_records( np.array([(3, 'a'), (2, 'b'), (1, 'c'), (0, 'd')], dtype=[('col_1', 'i4'), ('col_2', 'U1')])), ), ]) def test_create_methods(self, func): expected = func(pd.DataFrame) deferred_df = func(frames.DeferredDataFrame) actual = expressions.Session({}).evaluate(deferred_df._expr) pd.testing.assert_frame_equal(actual, expected) def test_replace(self): # verify a replace() doctest case that doesn't quite work in Beam as it uses # the default method='pad' df = pd.DataFrame({'A': ['bat', 'foo', 'bait'], 'B': ['abc', 'bar', 'xyz']}) self._run_test( lambda df: df.replace( regex={ r'^ba.$': 'new', 'foo': 'xyz' }, method=None), df) def test_sample_columns(self): df = pd.DataFrame({ 'brand': ['Yum Yum', 'Yum Yum', 'Indomie', 'Indomie', 'Indomie'], 'style': ['cup', 'cup', 'cup', 'pack', 'pack'], 'rating': [4, 4, 3.5, 15, 5] }) self._run_test(lambda df: df.sample(axis=1, n=2, random_state=1), df) self._run_error_test(lambda df: df.sample(axis=1, n=10, random_state=2), df) self._run_test( lambda df: df.sample(axis=1, n=10, random_state=3, replace=True), df) def test_cat(self): # Replicate the doctests from CategorigcalAccessor # These tests don't translate into pandas_doctests_test.py because it # tries to use astype("category") in Beam, which makes a non-deferred # column type. s = pd.Series(list("abbccc")).astype("category") self._run_test(lambda s: s.cat.rename_categories(list("cba")), s) self._run_test(lambda s: s.cat.reorder_categories(list("cba")), s) self._run_test(lambda s: s.cat.add_categories(["d", "e"]), s) self._run_test(lambda s: s.cat.remove_categories(["a", "c"]), s) self._run_test(lambda s: s.cat.set_categories(list("abcde")), s) self._run_test(lambda s: s.cat.as_ordered(), s) self._run_test(lambda s: s.cat.as_unordered(), s) self._run_test(lambda s: s.cat.codes, s) @parameterized.expand(frames.ELEMENTWISE_DATETIME_PROPERTIES) def test_dt_property(self, prop_name): # Generate a series with a lot of unique timestamps s = pd.Series( pd.date_range('1/1/2000', periods=100, freq='m') + pd.timedelta_range(start='0 days', end='70 days', periods=100)) self._run_test(lambda s: getattr(s.dt, prop_name), s) @parameterized.expand([ ('month_name', {}), ('day_name', {}), ('normalize', {}), ( 'strftime', { 'date_format': '%B %d, %Y, %r' }, ), ('tz_convert', { 'tz': 'Europe/Berlin' }), ]) def test_dt_method(self, op, kwargs): # Generate a series with a lot of unique timestamps s = pd.Series( pd.date_range( '1/1/2000', periods=100, freq='m', tz='America/Los_Angeles') + pd.timedelta_range(start='0 days', end='70 days', periods=100)) self._run_test(lambda s: getattr(s.dt, op)(**kwargs), s) def test_dt_tz_localize_ambiguous_series(self): # This replicates a dt.tz_localize doctest: # s.tz_localize('CET', ambiguous=np.array([True, True, False])) # But using a DeferredSeries instead of a np array s = pd.to_datetime( pd.Series([ '2018-10-28 01:20:00', '2018-10-28 02:36:00', '2018-10-28 03:46:00' ])) ambiguous = pd.Series([True, True, False], index=s.index) self._run_test( lambda s, ambiguous: s.dt.tz_localize('CET', ambiguous=ambiguous), s, ambiguous) def test_dt_tz_localize_nonexistent(self): # This replicates dt.tz_localize doctests that exercise `nonexistent`. # However they specify ambiguous='NaT' because the default, # ambiguous='infer', is not supported. s = pd.to_datetime( pd.Series(['2015-03-29 02:30:00', '2015-03-29 03:30:00'])) self._run_test( lambda s: s.dt.tz_localize( 'Europe/Warsaw', ambiguous='NaT', nonexistent='shift_forward'), s) self._run_test( lambda s: s.dt.tz_localize( 'Europe/Warsaw', ambiguous='NaT', nonexistent='shift_backward'), s) self._run_test( lambda s: s.dt.tz_localize( 'Europe/Warsaw', ambiguous='NaT', nonexistent=pd.Timedelta('1H')), s) def test_compare_series(self): s1 = pd.Series(["a", "b", "c", "d", "e"]) s2 = pd.Series(["a", "a", "c", "b", "e"]) self._run_test(lambda s1, s2: s1.compare(s2), s1, s2) self._run_test(lambda s1, s2: s1.compare(s2, align_axis=0), s1, s2) self._run_test(lambda s1, s2: s1.compare(s2, keep_shape=True), s1, s2) self._run_test( lambda s1, s2: s1.compare(s2, keep_shape=True, keep_equal=True), s1, s2) def test_compare_dataframe(self): df1 = pd.DataFrame( { "col1": ["a", "a", "b", "b", "a"], "col2": [1.0, 2.0, 3.0, np.nan, 5.0], "col3": [1.0, 2.0, 3.0, 4.0, 5.0] }, columns=["col1", "col2", "col3"], ) df2 = df1.copy() df2.loc[0, 'col1'] = 'c' df2.loc[2, 'col3'] = 4.0 # Skipped because keep_shape=False won't be implemented with self.assertRaisesRegex( frame_base.WontImplementError, r"compare\(align_axis\=1, keep_shape\=False\) is not allowed"): self._run_test(lambda df1, df2: df1.compare(df2), df1, df2) self._run_test( lambda df1, df2: df1.compare(df2, align_axis=0), df1, df2, check_proxy=False) self._run_test(lambda df1, df2: df1.compare(df2, keep_shape=True), df1, df2) self._run_test( lambda df1, df2: df1.compare(df2, align_axis=0, keep_shape=True), df1, df2) self._run_test( lambda df1, df2: df1.compare(df2, keep_shape=True, keep_equal=True), df1, df2) self._run_test( lambda df1, df2: df1.compare( df2, align_axis=0, keep_shape=True, keep_equal=True), df1, df2) def test_idxmin(self): df = pd.DataFrame({ 'consumption': [10.51, 103.11, 55.48], 'co2_emissions': [37.2, 19.66, 1712] }, index=['Pork', 'Wheat Products', 'Beef']) df2 = df.copy() df2.loc['Pork', 'co2_emissions'] = None df2.loc['Wheat Products', 'consumption'] = None df2.loc['Beef', 'co2_emissions'] = None df3 = pd.DataFrame({ 'consumption': [1.1, 2.2, 3.3], 'co2_emissions': [3.3, 2.2, 1.1] }, index=[0, 1, 2]) s = pd.Series(data=[4, 3, None, 1], index=['A', 'B', 'C', 'D']) s2 = pd.Series(data=[1, 2, 3], index=[1, 2, 3]) self._run_test(lambda df: df.idxmin(), df) self._run_test(lambda df: df.idxmin(skipna=False), df) self._run_test(lambda df: df.idxmin(axis=1), df) self._run_test(lambda df: df.idxmin(axis=1, skipna=False), df) self._run_test(lambda df2: df2.idxmin(), df2) self._run_test(lambda df2: df2.idxmin(axis=1), df2) self._run_test(lambda df2: df2.idxmin(skipna=False), df2, check_proxy=False) self._run_test( lambda df2: df2.idxmin(axis=1, skipna=False), df2, check_proxy=False) self._run_test(lambda df3: df3.idxmin(), df3) self._run_test(lambda df3: df3.idxmin(axis=1), df3) self._run_test(lambda df3: df3.idxmin(skipna=False), df3) self._run_test(lambda df3: df3.idxmin(axis=1, skipna=False), df3) self._run_test(lambda s: s.idxmin(), s) self._run_test(lambda s: s.idxmin(skipna=False), s, check_proxy=False) self._run_test(lambda s2: s2.idxmin(), s2) self._run_test(lambda s2: s2.idxmin(skipna=False), s2) def test_idxmax(self): df = pd.DataFrame({ 'consumption': [10.51, 103.11, 55.48], 'co2_emissions': [37.2, 19.66, 1712] }, index=['Pork', 'Wheat Products', 'Beef']) df2 = df.copy() df2.loc['Pork', 'co2_emissions'] = None df2.loc['Wheat Products', 'consumption'] = None df2.loc['Beef', 'co2_emissions'] = None df3 = pd.DataFrame({ 'consumption': [1.1, 2.2, 3.3], 'co2_emissions': [3.3, 2.2, 1.1] }, index=[0, 1, 2]) s = pd.Series(data=[1, None, 4, 1], index=['A', 'B', 'C', 'D']) s2 = pd.Series(data=[1, 2, 3], index=[1, 2, 3]) self._run_test(lambda df: df.idxmax(), df) self._run_test(lambda df: df.idxmax(skipna=False), df) self._run_test(lambda df: df.idxmax(axis=1), df) self._run_test(lambda df: df.idxmax(axis=1, skipna=False), df) self._run_test(lambda df2: df2.idxmax(), df2) self._run_test(lambda df2: df2.idxmax(axis=1), df2) self._run_test( lambda df2: df2.idxmax(axis=1, skipna=False), df2, check_proxy=False) self._run_test(lambda df2: df2.idxmax(skipna=False), df2, check_proxy=False) self._run_test(lambda df3: df3.idxmax(), df3) self._run_test(lambda df3: df3.idxmax(axis=1), df3) self._run_test(lambda df3: df3.idxmax(skipna=False), df3) self._run_test(lambda df3: df3.idxmax(axis=1, skipna=False), df3) self._run_test(lambda s: s.idxmax(), s) self._run_test(lambda s: s.idxmax(skipna=False), s, check_proxy=False) self._run_test(lambda s2: s2.idxmax(), s2) self._run_test(lambda s2: s2.idxmax(skipna=False), s2) def test_pipe(self): def df_times(df, column, times): df[column] = df[column] * times return df def df_times_shuffled(column, times, df): return df_times(df, column, times) def s_times(s, times): return s * times def s_times_shuffled(times, s): return s_times(s, times) df = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]}, index=[0, 1, 2]) s = pd.Series([1, 2, 3, 4, 5], index=[0, 1, 2, 3, 4]) self._run_inplace_test(lambda df: df.pipe(df_times, 'A', 2), df) self._run_inplace_test( lambda df: df.pipe((df_times_shuffled, 'df'), 'A', 2), df) self._run_test(lambda s: s.pipe(s_times, 2), s) self._run_test(lambda s: s.pipe((s_times_shuffled, 's'), 2), s) def test_unstack_pandas_series_not_multiindex(self): # Pandas should throw a ValueError if performing unstack # on a Series without MultiIndex s = pd.Series([1, 2, 3, 4], index=['one', 'two', 'three', 'four']) with self.assertRaises((AttributeError, ValueError)): self._run_test(lambda s: s.unstack(), s) def test_unstack_non_categorical_index(self): index = pd.MultiIndex.from_tuples([('one', 'a'), ('one', 'b'), ('two', 'a'), ('two', 'b')]) index = index.set_levels( index.levels[0].astype(pd.CategoricalDtype(['one', 'two'])), level=0) s = pd.Series(np.arange(1.0, 5.0), index=index) with self.assertRaisesRegex( frame_base.WontImplementError, r"unstack\(\) is only supported on DataFrames if"): self._run_test(lambda s: s.unstack(level=-1), s) def _unstack_get_categorical_index(self): index = pd.MultiIndex.from_tuples([('one', 'a'), ('one', 'b'), ('two', 'a'), ('two', 'b')]) index = index.set_levels( index.levels[0].astype(pd.CategoricalDtype(['one', 'two'])), level=0) index = index.set_levels( index.levels[1].astype(pd.CategoricalDtype(['a', 'b'])), level=1) return index def test_unstack_pandas_example1(self): index = self._unstack_get_categorical_index() s = pd.Series(np.arange(1.0, 5.0), index=index) self._run_test(lambda s: s.unstack(level=-1), s) def test_unstack_pandas_example2(self): index = self._unstack_get_categorical_index() s = pd.Series(np.arange(1.0, 5.0), index=index) self._run_test(lambda s: s.unstack(level=0), s) @unittest.skipIf( sys.version_info >= (3, 12) and PD_VERSION < (2, 3), 'https://github.com/pandas-dev/pandas/issues/58604') def test_unstack_pandas_example3(self): index = self._unstack_get_categorical_index() s = pd.Series(np.arange(1.0, 5.0), index=index) df = s.unstack(level=0) if PD_VERSION < (1, 2): with self.assertRaisesRegex( frame_base.WontImplementError, r"unstack\(\) is not supported when using pandas < 1.2.0"): self._run_test(lambda df: df.unstack(), df) else: self._run_test(lambda df: df.unstack(), df) @unittest.skipIf( PD_VERSION < (1, 4), "Cannot set dtype of index to boolean for pandas<1.4") def test_unstack_bool(self): index = pd.MultiIndex.from_tuples([(True, 'a'), (True, 'b'), (False, 'a'), (False, 'b')]) index = index.set_levels(index.levels[0].astype('boolean'), level=0) index = index.set_levels( index.levels[1].astype(pd.CategoricalDtype(['a', 'b'])), level=1) s = pd.Series(np.arange(1.0, 5.0), index=index) self._run_test(lambda s: s.unstack(level=0), s) def test_unstack_series_multiple_index_levels(self): tuples = list( zip( *[ ["bar", "bar", "bar", "bar", "baz", "baz", "baz", "baz"], ["one", "one", "two", "two", "one", "one", "two", "two"], ["A", "B", "A", "B", "A", "B", "A", "B"], ])) index = pd.MultiIndex.from_tuples( tuples, names=["first", "second", "third"]) index = index.set_levels( index.levels[0].astype(pd.CategoricalDtype(['bar', 'baz'])), level=0) index = index.set_levels( index.levels[1].astype(pd.CategoricalDtype(['one', 'two'])), level=1) index = index.set_levels( index.levels[2].astype(pd.CategoricalDtype(['A', 'B'])), level=2) df = pd.Series(np.random.randn(8), index=index) self._run_test(lambda df: df.unstack(level=['first', 'third']), df) def test_unstack_series_multiple_index_and_column_levels(self): columns = pd.MultiIndex.from_tuples( [ ("A", "cat", "long"), ("B", "cat", "long"), ("A", "dog", "short"), ("B", "dog", "short"), ], names=["exp", "animal", "hair_length"], ) index = pd.MultiIndex.from_product( [['one', 'two'], ['a', 'b'], ['bar', 'baz']], names=["first", "second", "third"]) index = index.set_levels( index.levels[0].astype(pd.CategoricalDtype(['one', 'two'])), level=0) index = index.set_levels( index.levels[1].astype(pd.CategoricalDtype(['a', 'b'])), level=1) index = index.set_levels( index.levels[2].astype(pd.CategoricalDtype(['bar', 'baz'])), level=2) df = pd.DataFrame(np.random.randn(8, 4), index=index, columns=columns) df = df.stack(level=["animal", "hair_length"]) self._run_test(lambda df: df.unstack(level=['second', 'third']), df) self._run_test(lambda df: df.unstack(level=['second']), df) def test_pivot_non_categorical(self): df = pd.DataFrame({ 'foo': ['one', 'one', 'one', 'two', 'two', 'two'], 'bar': ['A', 'B', 'C', 'A', 'B', 'C'], 'baz': [1, 2, 3, 4, 5, 6], 'zoo': ['x', 'y', 'z', 'q', 'w', 't'] }) with self.assertRaisesRegex( frame_base.WontImplementError, r"pivot\(\) of non-categorical type is not supported"): self._run_test( lambda df: df.pivot(index='foo', columns='bar', values='baz'), df) def test_pivot_pandas_example1(self): # Simple test 1 df = pd.DataFrame({ 'foo': ['one', 'one', 'one', 'two', 'two', 'two'], 'bar': ['A', 'B', 'C', 'A', 'B', 'C'], 'baz': [1, 2, 3, 4, 5, 6], 'zoo': ['x', 'y', 'z', 'q', 'w', 't'] }) df['bar'] = df['bar'].astype( pd.CategoricalDtype(categories=['A', 'B', 'C'])) self._run_test( lambda df: df.pivot(index='foo', columns='bar', values='baz'), df) self._run_test( lambda df: df.pivot(index=['foo'], columns='bar', values='baz'), df) def test_pivot_pandas_example3(self): # Multiple values df = pd.DataFrame({ 'foo': ['one', 'one', 'one', 'two', 'two', 'two'], 'bar': ['A', 'B', 'C', 'A', 'B', 'C'], 'baz': [1, 2, 3, 4, 5, 6], 'zoo': ['x', 'y', 'z', 'q', 'w', 't'] }) df['bar'] = df['bar'].astype( pd.CategoricalDtype(categories=['A', 'B', 'C'])) self._run_test( lambda df: df.pivot(index='foo', columns='bar', values=['baz', 'zoo']), df) self._run_test( lambda df: df.pivot( index='foo', columns=['bar'], values=['baz', 'zoo']), df) def test_pivot_pandas_example4(self): # Multiple columns df = pd.DataFrame({ "lev1": [1, 1, 1, 2, 2, 2], "lev2": [1, 1, 2, 1, 1, 2], "lev3": [1, 2, 1, 2, 1, 2], "lev4": [1, 2, 3, 4, 5, 6], "values": [0, 1, 2, 3, 4, 5] }) df['lev2'] = df['lev2'].astype(pd.CategoricalDtype(categories=[1, 2])) df['lev3'] = df['lev3'].astype(pd.CategoricalDtype(categories=[1, 2])) df['values'] = df['values'].astype('Int64') self._run_test( lambda df: df.pivot( index="lev1", columns=["lev2", "lev3"], values="values"), df) def test_pivot_pandas_example5(self): # Multiple index df = pd.DataFrame({ "lev1": [1, 1, 1, 2, 2, 2], "lev2": [1, 1, 2, 1, 1, 2], "lev3": [1, 2, 1, 2, 1, 2], "lev4": [1, 2, 3, 4, 5, 6], "values": [0, 1, 2, 3, 4, 5] }) df['lev3'] = df['lev3'].astype(pd.CategoricalDtype(categories=[1, 2])) # Cast to nullable Int64 because Beam doesn't do the correct conversion to # float64 df['values'] = df['values'].astype('Int64') if PD_VERSION < (1, 4): with self.assertRaisesRegex( frame_base.WontImplementError, r"pivot\(\) is not supported when pandas<1.4 and index is a Multi"): self._run_test( lambda df: df.pivot( index=["lev1", "lev2"], columns=["lev3"], values="values"), df) else: self._run_test( lambda df: df.pivot( index=["lev1", "lev2"], columns=["lev3"], values="values"), df) def test_pivot_pandas_example6(self): # Value error when there are duplicates df = pd.DataFrame({ "foo": ['one', 'one', 'two', 'two'], "bar": ['A', 'A', 'B', 'C'], "baz": [1, 2, 3, 4] }) df['bar'] = df['bar'].astype( pd.CategoricalDtype(categories=['A', 'B', 'C'])) self._run_error_test( lambda df: df.pivot(index='foo', columns='bar', values='baz'), df, construction_time=False) def test_pivot_no_index_provided_on_single_level_index(self): # Multiple columns, no index value provided df = pd.DataFrame({ "lev1": [1, 1, 1, 2, 2, 2], "lev2": [1, 1, 2, 1, 1, 2], "lev3": [1, 2, 1, 2, 1, 2], "lev4": [1, 2, 3, 4, 5, 6], "values": [0, 1, 2, 3, 4, 5] }) df['lev2'] = df['lev2'].astype(pd.CategoricalDtype(categories=[1, 2])) df['lev3'] = df['lev3'].astype(pd.CategoricalDtype(categories=[1, 2])) df['values'] = df['values'].astype('Int64') self._run_test( lambda df: df.pivot(columns=["lev2", "lev3"], values="values"), df) def test_pivot_no_index_provided_on_multiindex(self): # Multiple columns, no index value provided tuples = list( zip( *[ ["bar", "bar", "bar", "baz", "baz", "baz"], [ "one", "two", "three", "one", "two", "three", ], ])) index = pd.MultiIndex.from_tuples(tuples, names=["first", "second"]) df = pd.DataFrame({ "lev1": [1, 1, 1, 2, 2, 2], "lev2": [1, 1, 2, 1, 1, 2], "lev3": [1, 2, 1, 2, 1, 2], "lev4": [1, 2, 3, 4, 5, 6], "values": [0, 1, 2, 3, 4, 5] }, index=index) df['lev2'] = df['lev2'].astype(pd.CategoricalDtype(categories=[1, 2])) df['lev3'] = df['lev3'].astype(pd.CategoricalDtype(categories=[1, 2])) df['values'] = df['values'].astype('float64') df['lev1'] = df['lev1'].astype('int64') df['lev4'] = df['lev4'].astype('int64') if PD_VERSION < (1, 4): with self.assertRaisesRegex( frame_base.WontImplementError, r"pivot\(\) is not supported when pandas<1.4 and index is a Multi"): self._run_test(lambda df: df.pivot(columns=["lev2", "lev3"]), df) else: self._run_test( lambda df: df.pivot(columns=["lev2", "lev3"]), df, lenient_dtype_check=True) # pandas doesn't support kurtosis on GroupBys: # https://github.com/pandas-dev/pandas/issues/40139 ALL_GROUPING_AGGREGATIONS = sorted( set(frames.ALL_AGGREGATIONS) - set(('kurt', 'kurtosis'))) AGGREGATIONS_WHERE_NUMERIC_ONLY_DEFAULTS_TO_TRUE_IN_PANDAS_1 = set( frames.ALL_AGGREGATIONS) - set(( 'nunique', 'size', 'count', 'idxmin', 'idxmax', 'mode', 'rank', 'all', 'any', 'describe')) def numeric_only_kwargs_for_pandas_2(agg_type: str) -> dict[str, bool]: """Get proper arguments for numeric_only. Behavior for numeric_only in these methods changed in Pandas 2 to default to False instead of True, so explicitly make it True in Pandas 2.""" if PD_VERSION >= (2, 0) and ( agg_type in AGGREGATIONS_WHERE_NUMERIC_ONLY_DEFAULTS_TO_TRUE_IN_PANDAS_1): return {'numeric_only': True} else: return {} class GroupByTest(_AbstractFrameTest): """Tests for DataFrame/Series GroupBy operations.""" @staticmethod def median_sum_fn(x): with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="Mean of empty slice") return (x.foo + x.bar).median() @parameterized.expand(ALL_GROUPING_AGGREGATIONS) def test_groupby_agg(self, agg_type): if agg_type == 'describe' and PD_VERSION < (1, 2): self.skipTest( "https://github.com/apache/beam/issues/20967: proxy generation of " "DataFrameGroupBy.describe fails in pandas < 1.2") kwargs = numeric_only_kwargs_for_pandas_2(agg_type) self._run_test( lambda df: df.groupby('group').agg(agg_type, **kwargs), GROUPBY_DF, check_proxy=False) @parameterized.expand(ALL_GROUPING_AGGREGATIONS) def test_groupby_with_filter(self, agg_type): if agg_type == 'describe' and PD_VERSION < (1, 2): self.skipTest( "https://github.com/apache/beam/issues/20967: proxy generation of " "DataFrameGroupBy.describe fails in pandas < 1.2") kwargs = numeric_only_kwargs_for_pandas_2(agg_type) self._run_test( lambda df: getattr(df[df.foo > 30].groupby('group'), agg_type) (**kwargs), GROUPBY_DF, check_proxy=False) @parameterized.expand(ALL_GROUPING_AGGREGATIONS) def test_groupby(self, agg_type): if agg_type == 'describe' and PD_VERSION < (1, 2): self.skipTest( "https://github.com/apache/beam/issues/20967: proxy generation of " "DataFrameGroupBy.describe fails in pandas < 1.2") kwargs = numeric_only_kwargs_for_pandas_2(agg_type) self._run_test( lambda df: getattr(df.groupby('group'), agg_type)(**kwargs), GROUPBY_DF, check_proxy=False) @parameterized.expand(ALL_GROUPING_AGGREGATIONS) def test_groupby_series(self, agg_type): if agg_type == 'describe' and PD_VERSION < (1, 2): self.skipTest( "https://github.com/apache/beam/issues/20967: proxy generation of " "DataFrameGroupBy.describe fails in pandas < 1.2") kwargs = numeric_only_kwargs_for_pandas_2(agg_type) self._run_test( lambda df: getattr(df[df.foo > 40].groupby(df.group), agg_type) (**kwargs), GROUPBY_DF, check_proxy=False) def test_groupby_user_guide(self): # Example from https://pandas.pydata.org/docs/user_guide/groupby.html arrays = [['bar', 'bar', 'baz', 'baz', 'foo', 'foo', 'qux', 'qux'], ['one', 'two', 'one', 'two', 'one', 'two', 'one', 'two']] index = pd.MultiIndex.from_arrays(arrays, names=['first', 'second']) df = pd.DataFrame({ 'A': [1, 1, 1, 1, 2, 2, 3, 3], 'B': np.arange(8) }, index=index) self._run_test(lambda df: df.groupby(['second', 'A']).sum(), df) @parameterized.expand(ALL_GROUPING_AGGREGATIONS) def test_groupby_project_series(self, agg_type): df = GROUPBY_DF if agg_type == 'describe': self.skipTest( "https://github.com/apache/beam/issues/20967: proxy generation of " "SeriesGroupBy.describe fails") if agg_type in ('corr', 'cov'): self.skipTest( "https://github.com/apache/beam/issues/20895: " "SeriesGroupBy.{corr, cov} do not raise the expected error.") kwargs = numeric_only_kwargs_for_pandas_2(agg_type) self._run_test( lambda df: getattr(df.groupby('group').foo, agg_type)(**kwargs), df) self._run_test( lambda df: getattr(df.groupby('group').bar, agg_type)(**kwargs), df) self._run_test( lambda df: getattr(df.groupby('group')['foo'], agg_type)(**kwargs), df) self._run_test( lambda df: getattr(df.groupby('group')['bar'], agg_type)(**kwargs), df) @parameterized.expand(ALL_GROUPING_AGGREGATIONS) def test_groupby_project_dataframe(self, agg_type): if agg_type == 'describe' and PD_VERSION < (1, 2): self.skipTest( "https://github.com/apache/beam/issues/20967: proxy generation of " "DataFrameGroupBy.describe fails in pandas < 1.2") kwargs = numeric_only_kwargs_for_pandas_2(agg_type) self._run_test( lambda df: getattr(df.groupby('group')[['bar', 'baz']], agg_type) (**kwargs), GROUPBY_DF, check_proxy=False) def test_groupby_errors_bad_projection(self): df = GROUPBY_DF # non-existent projection column self._run_error_test( lambda df: df.groupby('group')[['bar', 'baz']].bar.median(), df) self._run_error_test(lambda df: df.groupby('group')[['bad']].median(), df) self._run_error_test(lambda df: df.groupby('group').bad.median(), df) self._run_error_test( lambda df: df.groupby('group')[['bar', 'baz']].bar.sum(), df) self._run_error_test(lambda df: df.groupby('group')[['bat']].sum(), df) self._run_error_test(lambda df: df.groupby('group').bat.sum(), df) def test_groupby_errors_non_existent_label(self): df = GROUPBY_DF # non-existent grouping label self._run_error_test( lambda df: df.groupby(['really_bad', 'foo', 'bad']).foo.sum(), df) self._run_error_test(lambda df: df.groupby('bad').foo.sum(), df) def test_groupby_callable(self): df = GROUPBY_DF kwargs = numeric_only_kwargs_for_pandas_2('sum') self._run_test(lambda df: df.groupby(lambda x: x % 2).foo.sum(**kwargs), df) kwargs = numeric_only_kwargs_for_pandas_2('median') self._run_test(lambda df: df.groupby(lambda x: x % 5).median(**kwargs), df) def test_groupby_apply(self): df = GROUPBY_DF # Note this is the same as DataFrameGroupBy.describe. Using it here is # just a convenient way to test apply() with a user fn that returns a Series describe = lambda df: df.describe() self._run_test(lambda df: df.groupby('group').foo.apply(describe), df) self._run_test( lambda df: df.groupby('group')[['foo', 'bar']].apply(describe), df) self._run_test(lambda df: df.groupby('group').apply(self.median_sum_fn), df) self._run_test( lambda df: df.set_index('group').foo.groupby(level=0).apply(describe), df) self._run_test(lambda df: df.groupby(level=0).apply(self.median_sum_fn), df) self._run_test(lambda df: df.groupby(lambda x: x % 3).apply(describe), df) self._run_test( lambda df: df.bar.groupby(lambda x: x % 3).apply(describe), df) self._run_test( lambda df: df.set_index(['str', 'group', 'bool']).groupby( level='group').apply(self.median_sum_fn), df) def test_groupby_apply_preserves_column_order(self): df = GROUPBY_DF self._run_test( lambda df: df[['foo', 'group', 'bar']].groupby( 'group', group_keys=False).apply(lambda x: x), df) def test_groupby_transform(self): df = pd.DataFrame({ "Date": [ "2015-05-08", "2015-05-07", "2015-05-06", "2015-05-05", "2015-05-08", "2015-05-07", "2015-05-06", "2015-05-05" ], "Data": [5, 8, 6, 1, 50, 100, 60, 120], }) self._run_test(lambda df: df.groupby('Date')['Data'].transform(np.sum), df) self._run_test( lambda df: df.groupby('Date')['Data'].transform( lambda x: (x - x.mean()) / x.std()), df) def test_groupby_pipe(self): df = GROUPBY_DF kwargs = numeric_only_kwargs_for_pandas_2('sum') self._run_test( lambda df: df.groupby('group').pipe(lambda x: x.sum(**kwargs)), df) self._run_test( lambda df: df.groupby('group')['bool'].pipe(lambda x: x.any()), df) self._run_test( lambda df: df.groupby(['group', 'foo']).pipe( (lambda a, x: x.sum(numeric_only=a), 'x'), False), df, check_proxy=False) def test_groupby_apply_modified_index(self): df = GROUPBY_DF # If apply fn modifies the index then the output will include the grouped # index self._run_test( lambda df: df.groupby('group').apply( lambda x: x[x.foo > x.foo.median()]), df) @unittest.skip('https://github.com/apache/beam/issues/20762') def test_groupby_aggregate_grouped_column(self): df = pd.DataFrame({ 'group': ['a' if i % 5 == 0 or i % 3 == 0 else 'b' for i in range(100)], 'foo': [None if i % 11 == 0 else i for i in range(100)], 'bar': [None if i % 7 == 0 else 99 - i for i in range(100)], 'baz': [None if i % 13 == 0 else i * 2 for i in range(100)], }) self._run_test(lambda df: df.groupby('group').group.count(), df) self._run_test(lambda df: df.groupby('group')[['group', 'bar']].count(), df) self._run_test( lambda df: df.groupby('group')[['group', 'bar']].apply( lambda x: x.describe()), df) @parameterized.expand((x, ) for x in [ 0, [1], 3, [0, 3], [2, 1], ['foo', 0], [1, 'str'], [3, 0, 2, 1], ]) def test_groupby_level_agg(self, level): df = GROUPBY_DF.set_index(['group', 'foo', 'bar', 'str'], drop=False) self._run_test(lambda df: df.groupby(level=level).bar.max(), df) self._run_test( lambda df: df.groupby(level=level).sum(numeric_only=True), df) self._run_test( lambda df: df.groupby(level=level).apply(self.median_sum_fn), df) @unittest.skipIf(PD_VERSION < (1, 1), "drop_na added in pandas 1.1.0") def test_groupby_count_na(self): # Verify we can do a groupby.count() that doesn't drop NaN values self._run_test( lambda df: df.groupby('foo', dropna=True).bar.count(), GROUPBY_DF) self._run_test( lambda df: df.groupby('foo', dropna=False).bar.count(), GROUPBY_DF) def test_groupby_sum_min_count(self): df = pd.DataFrame({ 'good': [1, 2, 3, np.nan], 'bad': [np.nan, np.nan, np.nan, 4], 'group': ['a', 'b', 'a', 'b'] }) self._run_test(lambda df: df.groupby('group').sum(min_count=2), df) @unittest.skipIf( PD_VERSION >= (2, 0), "dtypes on groups is deprecated in Pandas 2.") def test_groupby_dtypes(self): self._run_test( lambda df: df.groupby('group').dtypes, GROUPBY_DF, check_proxy=False) self._run_test( lambda df: df.groupby(level=0).dtypes, GROUPBY_DF, check_proxy=False) @parameterized.expand(ALL_GROUPING_AGGREGATIONS) def test_dataframe_groupby_series(self, agg_type): if agg_type == 'describe' and PD_VERSION < (1, 2): self.skipTest( "https://github.com/apache/beam/issues/20967: proxy generation of " "DataFrameGroupBy.describe fails in pandas < 1.2") def agg(df, group_by): kwargs = numeric_only_kwargs_for_pandas_2(agg_type) return df[df.foo > 40].groupby(group_by).agg(agg_type, **kwargs) self._run_test(lambda df: agg(df, df.group), GROUPBY_DF, check_proxy=False) self._run_test( lambda df: agg(df, df.foo % 3), GROUPBY_DF, check_proxy=False) @parameterized.expand(ALL_GROUPING_AGGREGATIONS) def test_series_groupby_series(self, agg_type): if agg_type == 'describe': self.skipTest( "https://github.com/apache/beam/issues/20967: proxy generation of " "SeriesGroupBy.describe fails") if agg_type in ('corr', 'cov'): self.skipTest( "https://github.com/apache/beam/issues/20895: " "SeriesGroupBy.{corr, cov} do not raise the expected error.") self._run_test( lambda df: df[df.foo < 40].bar.groupby(df.group).agg(agg_type), GROUPBY_DF) self._run_test( lambda df: df[df.foo < 40].bar.groupby(df.foo % 3).agg(agg_type), GROUPBY_DF) def test_groupby_series_apply(self): df = GROUPBY_DF # Note this is the same as DataFrameGroupBy.describe. Using it here is # just a convenient way to test apply() with a user fn that returns a Series describe = lambda df: df.describe() self._run_test(lambda df: df.groupby(df.group).foo.apply(describe), df) self._run_test( lambda df: df.groupby(df.group)[['foo', 'bar']].apply(describe), df) self._run_test( lambda df: df.groupby(df.group).apply(self.median_sum_fn), df) def test_groupby_multiindex_keep_nans(self): # Due to https://github.com/pandas-dev/pandas/issues/36470 # groupby(dropna=False) doesn't work with multiple columns with self.assertRaisesRegex(NotImplementedError, "https://github.com/apache/beam/issues/21014"): self._run_test( lambda df: df.groupby(['foo', 'bar'], dropna=False).sum(), GROUPBY_DF) NONPARALLEL_METHODS = ['quantile', 'describe', 'median', 'sem'] # mad was removed in pandas 2 if PD_VERSION < (2, 0): NONPARALLEL_METHODS.append('mad') class AggregationTest(_AbstractFrameTest): """Tests for global aggregation methods on DataFrame/Series.""" # corr, cov on Series require an other argument @parameterized.expand( sorted(set(frames.ALL_AGGREGATIONS) - set(['corr', 'cov']))) def test_series_agg(self, agg_method): s = pd.Series(list(range(16))) nonparallel = agg_method in NONPARALLEL_METHODS # TODO(https://github.com/apache/beam/issues/20926): max and min produce # the wrong proxy check_proxy = agg_method not in ('max', 'min') self._run_test( lambda s: s.agg(agg_method), s, nonparallel=nonparallel, check_proxy=check_proxy) # corr, cov on Series require an other argument # Series.size is a property @parameterized.expand( sorted(set(frames.ALL_AGGREGATIONS) - set(['corr', 'cov', 'size']))) def test_series_agg_method(self, agg_method): s = pd.Series(list(range(16))) nonparallel = agg_method in NONPARALLEL_METHODS # TODO(https://github.com/apache/beam/issues/20926): max and min produce # the wrong proxy check_proxy = agg_method not in ('max', 'min') self._run_test( lambda s: getattr(s, agg_method)(), s, nonparallel=nonparallel, check_proxy=check_proxy) @parameterized.expand(frames.ALL_AGGREGATIONS) def test_dataframe_agg(self, agg_method): df = pd.DataFrame({'A': [1, 2, 3, 4], 'B': [2, 3, 5, 7]}) nonparallel = agg_method in NONPARALLEL_METHODS # TODO(https://github.com/apache/beam/issues/20926): max and min produce # the wrong proxy check_proxy = agg_method not in ('max', 'min') self._run_test( lambda df: df.agg(agg_method), df, nonparallel=nonparallel, check_proxy=check_proxy) # DataFrame.size is a property @parameterized.expand(sorted(set(frames.ALL_AGGREGATIONS) - set(['size']))) def test_dataframe_agg_method(self, agg_method): df = pd.DataFrame({'A': [1, 2, 3, 4], 'B': [2, 3, 5, 7]}) nonparallel = agg_method in NONPARALLEL_METHODS # TODO(https://github.com/apache/beam/issues/20926): max and min produce # the wrong proxy check_proxy = agg_method not in ('max', 'min') self._run_test( lambda df: getattr(df, agg_method)(), df, nonparallel=nonparallel, check_proxy=check_proxy) def test_series_agg_modes(self): s = pd.Series(list(range(16))) self._run_test(lambda s: s.agg('sum'), s) self._run_test(lambda s: s.agg(['sum']), s) self._run_test(lambda s: s.agg(['sum', 'mean']), s) self._run_test(lambda s: s.agg(['mean']), s) self._run_test(lambda s: s.agg('mean'), s) def test_dataframe_agg_modes(self): df = pd.DataFrame({'A': [1, 2, 3, 4], 'B': [2, 3, 5, 7]}) self._run_test(lambda df: df.agg('sum'), df) self._run_test(lambda df: df.agg(['sum', 'mean']), df) self._run_test(lambda df: df.agg({'A': 'sum', 'B': 'sum'}), df) self._run_test(lambda df: df.agg({'A': 'sum', 'B': 'mean'}), df) self._run_test(lambda df: df.agg({'A': ['sum', 'mean']}), df) self._run_test(lambda df: df.agg({'A': ['sum', 'mean'], 'B': 'min'}), df) @unittest.skipIf(PD_VERSION >= (2, 0), "level argument removed in Pandas 2") def test_series_agg_level(self): self._run_test( lambda df: df.set_index(['group', 'foo']).bar.count(level=0), GROUPBY_DF) self._run_test( lambda df: df.set_index(['group', 'foo']).bar.max(level=0), GROUPBY_DF) self._run_test( lambda df: df.set_index(['group', 'foo']).bar.median(level=0), GROUPBY_DF) self._run_test( lambda df: df.set_index(['foo', 'group']).bar.count(level=1), GROUPBY_DF) self._run_test( lambda df: df.set_index(['group', 'foo']).bar.max(level=1), GROUPBY_DF) self._run_test( lambda df: df.set_index(['group', 'foo']).bar.max(level='foo'), GROUPBY_DF) self._run_test( lambda df: df.set_index(['group', 'foo']).bar.median(level=1), GROUPBY_DF) @unittest.skipIf(PD_VERSION >= (2, 0), "level argument removed in Pandas 2") def test_dataframe_agg_level(self): self._run_test( lambda df: df.set_index(['group', 'foo']).count(level=0), GROUPBY_DF) self._run_test( lambda df: df.set_index(['group', 'foo']).max( level=0, numeric_only=False), GROUPBY_DF, check_proxy=False) # pandas implementation doesn't respect numeric_only argument here # (https://github.com/pandas-dev/pandas/issues/40788), it # always acts as if numeric_only=True. Our implmentation respects it so we # need to make it explicit. self._run_test( lambda df: df.set_index(['group', 'foo']).sum( level=0, numeric_only=True), GROUPBY_DF) self._run_test( lambda df: df.set_index(['group', 'foo'])[['bar']].count(level=1), GROUPBY_DF) self._run_test( lambda df: df.set_index(['group', 'foo']).count(level=1), GROUPBY_DF) self._run_test( lambda df: df.set_index(['group', 'foo']).max( level=1, numeric_only=False), GROUPBY_DF, check_proxy=False) # sum with str columns is order-sensitive self._run_test( lambda df: df.set_index(['group', 'foo']).sum( level=1, numeric_only=True), GROUPBY_DF) self._run_test( lambda df: df.set_index(['group', 'foo']).median( level=0, numeric_only=True), GROUPBY_DF) self._run_test( lambda df: df.drop('str', axis=1).set_index(['foo', 'group']).median( level=1, numeric_only=True), GROUPBY_DF) @unittest.skipIf(PD_VERSION >= (2, 0), "level argument removed in Pandas 2") def test_series_agg_multifunc_level(self): # level= is ignored for multiple agg fns self._run_test( lambda df: df.set_index(['group', 'foo']).bar.agg(['min', 'max'], level= 0), GROUPBY_DF) def test_series_mean_skipna(self): df = pd.DataFrame({ 'one': [i if i % 8 == 0 else np.nan for i in range(8)], 'two': [i if i % 4 == 0 else np.nan for i in range(8)], 'three': [i if i % 2 == 0 else np.nan for i in range(8)], }) self._run_test(lambda df: df.one.mean(skipna=False), df) self._run_test(lambda df: df.two.mean(skipna=False), df) self._run_test(lambda df: df.three.mean(skipna=False), df) self._run_test(lambda df: df.one.mean(skipna=True), df) self._run_test(lambda df: df.two.mean(skipna=True), df) self._run_test(lambda df: df.three.mean(skipna=True), df) @unittest.skipIf(PD_VERSION >= (2, 0), "level argument removed in Pandas 2") def test_dataframe_agg_multifunc_level(self): # level= is ignored for multiple agg fns self._run_test( lambda df: df.set_index(['group', 'foo']).agg(['min', 'max'], level=0), GROUPBY_DF, check_proxy=False) @parameterized.expand([(True, ), (False, )]) @unittest.skipIf( PD_VERSION < (1, 2), "pandas 1.1.0 produces different dtypes for these examples") def test_dataframe_agg_numeric_only(self, numeric_only): # Note other aggregation functions can fail on this input with # numeric_only={False,None}. These are the only ones that actually work for # the string inputs. self._run_test( lambda df: df.max(numeric_only=numeric_only), GROUPBY_DF, check_proxy=False) self._run_test( lambda df: df.min(numeric_only=numeric_only), GROUPBY_DF, check_proxy=False) @unittest.skip( "pandas implementation doesn't respect numeric_only= with " "level= (https://github.com/pandas-dev/pandas/issues/40788)") def test_dataframe_agg_level_numeric_only(self): self._run_test( lambda df: df.set_index('foo').sum(level=0, numeric_only=True), GROUPBY_DF) self._run_test( lambda df: df.set_index('foo').max(level=0, numeric_only=True), GROUPBY_DF) self._run_test( lambda df: df.set_index('foo').mean(level=0, numeric_only=True), GROUPBY_DF) self._run_test( lambda df: df.set_index('foo').median(level=0, numeric_only=True), GROUPBY_DF) def test_dataframe_agg_bool_only(self): df = pd.DataFrame({ 'all': [True for i in range(10)], 'any': [i % 3 == 0 for i in range(10)], 'int': range(10) }) self._run_test(lambda df: df.all(), df) self._run_test(lambda df: df.any(), df) self._run_test(lambda df: df.all(bool_only=True), df) self._run_test(lambda df: df.any(bool_only=True), df) @unittest.skip( "pandas doesn't implement bool_only= with level= " "(https://github.com/pandas-dev/pandas/blob/" "v1.2.3/pandas/core/generic.py#L10573)") def test_dataframe_agg_level_bool_only(self): df = pd.DataFrame({ 'all': [True for i in range(10)], 'any': [i % 3 == 0 for i in range(10)], 'int': range(10) }) self._run_test(lambda df: df.set_index('int', drop=False).all(level=0), df) self._run_test(lambda df: df.set_index('int', drop=False).any(level=0), df) self._run_test( lambda df: df.set_index('int', drop=False).all(level=0, bool_only=True), df) self._run_test( lambda df: df.set_index('int', drop=False).any(level=0, bool_only=True), df) def test_series_agg_np_size(self): self._run_test( lambda df: df.set_index(['group', 'foo']).agg(np.size), GROUPBY_DF, check_proxy=False) def test_df_agg_invalid_kwarg_raises(self): self._run_error_test(lambda df: df.agg('mean', bool_only=True), GROUPBY_DF) self._run_error_test( lambda df: df.agg('any', numeric_only=True), GROUPBY_DF) self._run_error_test( lambda df: df.agg('median', min_count=3, numeric_only=True), GROUPBY_DF) def test_series_agg_method_invalid_kwarg_raises(self): self._run_error_test(lambda df: df.foo.median(min_count=3), GROUPBY_DF) self._run_error_test( lambda df: df.foo.agg('median', min_count=3), GROUPBY_DF) @unittest.skipIf( PD_VERSION < (1, 3), ( "DataFrame.agg raises a different exception from the " "aggregation methods. Fixed in " "https://github.com/pandas-dev/pandas/pull/40543.")) def test_df_agg_method_invalid_kwarg_raises(self): self._run_error_test(lambda df: df.mean(bool_only=True), GROUPBY_DF) self._run_error_test(lambda df: df.any(numeric_only=True), GROUPBY_DF) self._run_error_test( lambda df: df.median(min_count=3, numeric_only=True), GROUPBY_DF) @unittest.skipIf(PD_VERSION >= (2, 0), "level argument removed in Pandas 2") def test_agg_min_count(self): df = pd.DataFrame({ 'good': [1, 2, 3, np.nan], 'bad': [np.nan, np.nan, np.nan, 4], }, index=['a', 'b', 'a', 'b']) self._run_test(lambda df: df.sum(level=0, min_count=2), df) self._run_test(lambda df: df.sum(min_count=3), df, nonparallel=True) self._run_test(lambda df: df.sum(min_count=1), df, nonparallel=True) self._run_test(lambda df: df.good.sum(min_count=2), df, nonparallel=True) self._run_test(lambda df: df.bad.sum(min_count=2), df, nonparallel=True) def test_series_agg_std(self): s = pd.Series(range(10)) self._run_test(lambda s: s.agg('std'), s) self._run_test(lambda s: s.agg('var'), s) self._run_test(lambda s: s.agg(['std', 'sum']), s) self._run_test(lambda s: s.agg(['var']), s) def test_std_all_na(self): s = pd.Series([np.nan] * 10) self._run_test(lambda s: s.agg('std'), s) self._run_test(lambda s: s.std(), s) def test_df_agg_operations_on_columns(self): self._run_test( lambda df: df.groupby('group').agg( mean_foo=('foo', lambda x: np.mean(x)), median_bar=('bar', lambda x: np.median(x)), sum_baz=('baz', 'sum'), count_bool=('bool', 'count'), ), GROUPBY_DF) def test_std_mostly_na_with_ddof(self): df = pd.DataFrame({ 'one': [i if i % 8 == 0 else np.nan for i in range(8)], 'two': [i if i % 4 == 0 else np.nan for i in range(8)], 'three': [i if i % 2 == 0 else np.nan for i in range(8)], }, index=pd.MultiIndex.from_arrays( [list(range(8)), list(reversed(range(8)))], names=['forward', None])) self._run_test(lambda df: df.std(), df) # ddof=1 self._run_test(lambda df: df.std(ddof=0), df) self._run_test(lambda df: df.std(ddof=2), df) self._run_test(lambda df: df.std(ddof=3), df) self._run_test(lambda df: df.std(ddof=4), df) def test_dataframe_std(self): self._run_test(lambda df: df.std(numeric_only=True), GROUPBY_DF) self._run_test(lambda df: df.var(numeric_only=True), GROUPBY_DF) def test_dataframe_mode(self): self._run_test( lambda df: df.mode(), GROUPBY_DF, nonparallel=True, check_proxy=False) self._run_test( lambda df: df.mode(numeric_only=True), GROUPBY_DF, nonparallel=True, check_proxy=False) self._run_test( lambda df: df.mode(dropna=True, numeric_only=True), GROUPBY_DF, nonparallel=True, check_proxy=False) def test_series_mode(self): self._run_test(lambda df: df.foo.mode(), GROUPBY_DF, nonparallel=True) self._run_test( lambda df: df.baz.mode(dropna=True), GROUPBY_DF, nonparallel=True) class BeamSpecificTest(unittest.TestCase): """Tests for functionality that's specific to the Beam DataFrame API. These features don't exist in pandas so we must verify them independently.""" def assert_frame_data_equivalent( self, actual, expected, check_column_subset=False, extra_col_value=0): """Verify that actual is the same as expected, ignoring the index and order of the data. Note: In order to perform non-deferred column operations in Beam, we have to enumerate all possible categories of data, even if they are ultimately unobserved. The default Pandas implementation on the other hand does not produce unobserved columns. This means when conducting tests, we need to account for the fact that the Beam result may be a superset of that of the Pandas result. If ``check_column_subset`` is `True`, we verify that all of the columns in the Dataframe returned from the Pandas implementation is contained in the Dataframe created from the Beam implementation. We also check if all columns that exist in the Beam implementation but not in the Pandas implementation are all equal to the ``extra_col_value`` to ensure that they were not erroneously populated. """ if check_column_subset: if isinstance(expected, pd.DataFrame): expected_cols = set(expected.columns) actual_cols = set(actual.columns) # Verifying that expected columns is a subset of the actual columns if not set(expected_cols).issubset(set(actual_cols)): raise AssertionError( f"Expected columns:\n{expected.columns}\n is not a" f"subset of {actual.columns}.") # Verifying that columns that don't exist in expected # but do in actual, are all equal to `extra_col_value` (default of 0) extra_columns = actual_cols - expected_cols if extra_columns: actual_extra_only = actual[list(extra_columns)] if np.isnan(extra_col_value): extra_cols_all_match = actual_extra_only.isna().all().all() else: extra_cols_all_match = actual_extra_only.eq( extra_col_value).all().all() if not extra_cols_all_match: raise AssertionError( f"Extra columns:{extra_columns}\n should all " f"be {extra_col_value}, but got \n{actual_extra_only}.") # Filtering actual to contain only columns in expected actual = actual[expected.columns] def sort_and_drop_index(df): if isinstance(df, pd.Series): df = df.sort_values() elif isinstance(df, pd.DataFrame): df = df.sort_values(by=list(df.columns)) return df.reset_index(drop=True) actual = sort_and_drop_index(actual) expected = sort_and_drop_index(expected) if isinstance(expected, pd.Series): pd.testing.assert_series_equal(actual, expected) elif isinstance(expected, pd.DataFrame): pd.testing.assert_frame_equal(actual, expected) def _evaluate(self, func, *args, distributed=True): deferred_args = [ frame_base.DeferredFrame.wrap( expressions.ConstantExpression(arg, arg[0:0])) for arg in args ] session_type = ( expressions.PartitioningSession if distributed else expressions.Session) return session_type({}).evaluate(func(*deferred_args)._expr) def test_drop_duplicates_keep_any(self): df = pd.DataFrame({ 'brand': ['Yum Yum', 'Yum Yum', 'Indomie', 'Indomie', 'Indomie'], 'style': ['cup', 'cup', 'cup', 'pack', 'pack'], 'rating': [4, 4, 3.5, 15, 5] }) result = self._evaluate(lambda df: df.drop_duplicates(keep='any'), df) # Verify that the result is the same as conventional drop_duplicates self.assert_frame_data_equivalent(result, df.drop_duplicates()) def test_drop_duplicates_keep_any_subset(self): df = pd.DataFrame({ 'brand': ['Yum Yum', 'Yum Yum', 'Indomie', 'Indomie', 'Indomie'], 'style': ['cup', 'cup', 'cup', 'pack', 'pack'], 'rating': [4, 4, 3.5, 15, 5] }) result = self._evaluate( lambda df: df.drop_duplicates(keep='any', subset=['brand']), df) self.assertTrue(result.brand.unique) self.assert_frame_data_equivalent( result.brand, df.drop_duplicates(subset=['brand']).brand) def test_series_drop_duplicates_keep_any(self): df = pd.DataFrame({ 'brand': ['Yum Yum', 'Yum Yum', 'Indomie', 'Indomie', 'Indomie'], 'style': ['cup', 'cup', 'cup', 'pack', 'pack'], 'rating': [4, 4, 3.5, 15, 5] }) result = self._evaluate(lambda df: df.brand.drop_duplicates(keep='any'), df) self.assert_frame_data_equivalent(result, df.brand.drop_duplicates()) def test_duplicated_keep_any(self): df = pd.DataFrame({ 'brand': ['Yum Yum', 'Yum Yum', 'Indomie', 'Indomie', 'Indomie'], 'style': ['cup', 'cup', 'cup', 'pack', 'pack'], 'rating': [4, 4, 3.5, 15, 5] }) result = self._evaluate(lambda df: df.duplicated(keep='any'), df) # Verify that the result is the same as conventional duplicated self.assert_frame_data_equivalent(result, df.duplicated()) def test_get_dummies_not_categoricaldtype(self): # Should not work because series is not a CategoricalDtype type with self.assertRaisesRegex( frame_base.WontImplementError, r"get_dummies\(\) of non-categorical type is not supported"): s = pd.Series(['a ,b', 'a', 'a, d']) self._evaluate(lambda s: s.str.get_dummies(','), s) # bool series do not work because they are not a CategoricalDtype type with self.assertRaisesRegex( frame_base.WontImplementError, r"get_dummies\(\) of non-categorical type is not supported"): s = pd.Series([True, False, False, True]) self._evaluate(lambda s: s.str.get_dummies(), s) def test_get_dummies_comma_separator(self): s = pd.Series(['a ,b', 'a', 'a, d', 'c']) s = s.astype(pd.CategoricalDtype(categories=['a ,b', 'c', 'b', 'a,d'])) result = self._evaluate(lambda s: s.str.get_dummies(','), s) self.assert_frame_data_equivalent( result, s.str.get_dummies(','), check_column_subset=True) def test_get_dummies_pandas_doc_example1(self): s = pd.Series(['a|b', 'a', 'a|c']) s = s.astype(pd.CategoricalDtype(categories=['a|b', 'a', 'a|c'])) result = self._evaluate(lambda s: s.str.get_dummies(), s) self.assert_frame_data_equivalent( result, s.str.get_dummies(), check_column_subset=True) def test_get_dummies_pandas_doc_example2(self): # Shouldn't still work even though np.nan is not considered a category # because we automatically create a nan column s = pd.Series(['a|b', np.nan, 'a|c']) s = s.astype(pd.CategoricalDtype(categories=['a|b', 'a|c'])) result = self._evaluate(lambda s: s.str.get_dummies(), s) self.assert_frame_data_equivalent( result, s.str.get_dummies(), check_column_subset=True) def test_get_dummies_pass_nan_as_category(self): # Explicitly pass 'nan' as a category s = pd.Series(['a|b', 'b|c', 'a|c', 'c', 'd']) s = s.astype(pd.CategoricalDtype(categories=['a', 'b', 'c', 'nan'])) result = self._evaluate(lambda s: s.str.get_dummies(), s) self.assert_frame_data_equivalent( result, s.str.get_dummies(), check_column_subset=True) def test_get_dummies_bools_casted_to_string(self): s = pd.Series([True, False, False, True]).astype('str') s = s.astype(pd.CategoricalDtype(categories=['True', 'False'])) result = self._evaluate(lambda s: s.str.get_dummies(), s) self.assert_frame_data_equivalent( result, s.str.get_dummies(), check_column_subset=True) def test_nsmallest_any(self): df = pd.DataFrame({ 'population': [ 59000000, 65000000, 434000, 434000, 434000, 337000, 337000, 11300, 11300 ], 'GDP': [1937894, 2583560, 12011, 4520, 12128, 17036, 182, 38, 311], 'alpha-2': ["IT", "FR", "MT", "MV", "BN", "IS", "NR", "TV", "AI"] }, index=[ "Italy", "France", "Malta", "Maldives", "Brunei", "Iceland", "Nauru", "Tuvalu", "Anguilla" ]) result = self._evaluate( lambda df: df.population.nsmallest(3, keep='any'), df) # keep='any' should produce the same result as keep='first', # but not necessarily with the same index self.assert_frame_data_equivalent(result, df.population.nsmallest(3)) def test_nlargest_any(self): df = pd.DataFrame({ 'population': [ 59000000, 65000000, 434000, 434000, 434000, 337000, 337000, 11300, 11300 ], 'GDP': [1937894, 2583560, 12011, 4520, 12128, 17036, 182, 38, 311], 'alpha-2': ["IT", "FR", "MT", "MV", "BN", "IS", "NR", "TV", "AI"] }, index=[ "Italy", "France", "Malta", "Maldives", "Brunei", "Iceland", "Nauru", "Tuvalu", "Anguilla" ]) result = self._evaluate( lambda df: df.population.nlargest(3, keep='any'), df) # keep='any' should produce the same result as keep='first', # but not necessarily with the same index self.assert_frame_data_equivalent(result, df.population.nlargest(3)) def test_pivot_pandas_example2(self): # Simple test 2 df = pd.DataFrame({ 'foo': ['one', 'one', 'one', 'two', 'two', 'two'], 'bar': ['A', 'B', 'C', 'A', 'B', 'C'], 'baz': [1, 2, 3, 4, 5, 6], 'zoo': ['x', 'y', 'z', 'q', 'w', 't'] }) df['bar'] = df['bar'].astype( pd.CategoricalDtype(categories=['A', 'B', 'C'])) result = self._evaluate(lambda df: df.pivot(index='foo', columns='bar'), df) # When there are multiple values, dtypes default to object. # Thus, need to convert to numeric with pd.to_numeric self.assert_frame_data_equivalent( result['baz'].apply(pd.to_numeric), df.pivot(index='foo', columns='bar')['baz']) def test_sample(self): df = pd.DataFrame({ 'population': [ 59000000, 65000000, 434000, 434000, 434000, 337000, 337000, 11300, 11300 ], 'GDP': [1937894, 2583560, 12011, 4520, 12128, 17036, 182, 38, 311], 'alpha-2': ["IT", "FR", "MT", "MV", "BN", "IS", "NR", "TV", "AI"] }, index=[ "Italy", "France", "Malta", "Maldives", "Brunei", "Iceland", "Nauru", "Tuvalu", "Anguilla" ]) result = self._evaluate(lambda df: df.sample(n=3), df) self.assertEqual(len(result), 3) series_result = self._evaluate(lambda df: df.GDP.sample(n=3), df) self.assertEqual(len(series_result), 3) self.assertEqual(series_result.name, "GDP") def test_sample_with_weights(self): df = pd.DataFrame({ 'population': [ 59000000, 65000000, 434000, 434000, 434000, 337000, 337000, 11300, 11300 ], 'GDP': [1937894, 2583560, 12011, 4520, 12128, 17036, 182, 38, 311], 'alpha-2': ["IT", "FR", "MT", "MV", "BN", "IS", "NR", "TV", "AI"] }, index=[ "Italy", "France", "Malta", "Maldives", "Brunei", "Iceland", "Nauru", "Tuvalu", "Anguilla" ]) weights = pd.Series([0, 0, 0, 0, 0, 0, 0, 1, 1], index=df.index) result = self._evaluate( lambda df, weights: df.sample(n=2, weights=weights), df, weights) self.assertEqual(len(result), 2) self.assertEqual(set(result.index), set(["Tuvalu", "Anguilla"])) series_result = self._evaluate( lambda df, weights: df.GDP.sample(n=2, weights=weights), df, weights) self.assertEqual(len(series_result), 2) self.assertEqual(series_result.name, "GDP") self.assertEqual(set(series_result.index), set(["Tuvalu", "Anguilla"])) def test_sample_with_missing_weights(self): df = pd.DataFrame({ 'population': [ 59000000, 65000000, 434000, 434000, 434000, 337000, 337000, 11300, 11300 ], 'GDP': [1937894, 2583560, 12011, 4520, 12128, 17036, 182, 38, 311], 'alpha-2': ["IT", "FR", "MT", "MV", "BN", "IS", "NR", "TV", "AI"] }, index=[ "Italy", "France", "Malta", "Maldives", "Brunei", "Iceland", "Nauru", "Tuvalu", "Anguilla" ]) # Missing weights are treated as 0 weights = pd.Series([.1, .01, np.nan, 0], index=["Nauru", "Iceland", "Anguilla", "Italy"]) result = self._evaluate( lambda df, weights: df.sample(n=2, weights=weights), df, weights) self.assertEqual(len(result), 2) self.assertEqual(set(result.index), set(["Nauru", "Iceland"])) series_result = self._evaluate( lambda df, weights: df.GDP.sample(n=2, weights=weights), df, weights) self.assertEqual(len(series_result), 2) self.assertEqual(series_result.name, "GDP") self.assertEqual(set(series_result.index), set(["Nauru", "Iceland"])) def test_sample_with_weights_distribution(self): target_prob = 0.25 num_samples = 100 num_targets = 200 num_other_elements = 10000 target_weight = target_prob / num_targets other_weight = (1 - target_prob) / num_other_elements self.assertTrue(target_weight > other_weight * 10, "weights too close") result = self._evaluate( lambda s, weights: s.sample(n=num_samples, weights=weights).sum(), # The first elements are 1, the rest are all 0. This means that when # we sum all the sampled elements (above), the result should be the # number of times the first elements (aka targets) were sampled. pd.Series([1] * num_targets + [0] * num_other_elements), pd.Series([target_weight] * num_targets + [other_weight] * num_other_elements)) # With the above constants, the probability of violating this invariant # (as computed using the Bernoulli distribution) is about 0.0012%. expected = num_samples * target_prob self.assertTrue(expected / 3 < result < expected * 2, (expected, result)) def test_split_pandas_examples_no_expand(self): # if expand=False (default), then no need to cast dtype to be # CategoricalDtype. s = pd.Series([ "this is a regular sentence", "https://docs.python.org/3/tutorial/index.html", np.nan ]) result = self._evaluate(lambda s: s.str.split(), s) self.assert_frame_data_equivalent(result, s.str.split()) result = self._evaluate(lambda s: s.str.rsplit(), s) self.assert_frame_data_equivalent(result, s.str.rsplit()) result = self._evaluate(lambda s: s.str.split(n=2), s) self.assert_frame_data_equivalent(result, s.str.split(n=2)) result = self._evaluate(lambda s: s.str.rsplit(n=2), s) self.assert_frame_data_equivalent(result, s.str.rsplit(n=2)) result = self._evaluate(lambda s: s.str.split(pat="/"), s) self.assert_frame_data_equivalent(result, s.str.split(pat="/")) def test_split_pandas_examples_expand_not_categorical(self): # When expand=True, there is exception because series is not categorical s = pd.Series([ "this is a regular sentence", "https://docs.python.org/3/tutorial/index.html", np.nan ]) with self.assertRaisesRegex( frame_base.WontImplementError, r"split\(\) of non-categorical type is not supported"): self._evaluate(lambda s: s.str.split(expand=True), s) with self.assertRaisesRegex( frame_base.WontImplementError, r"rsplit\(\) of non-categorical type is not supported"): self._evaluate(lambda s: s.str.rsplit(expand=True), s) def test_split_pandas_examples_expand_pat_is_string_literal1(self): # When expand=True and pattern is treated as a string literal s = pd.Series([ "this is a regular sentence", "https://docs.python.org/3/tutorial/index.html", np.nan ]) s = s.astype( pd.CategoricalDtype( categories=[ 'this is a regular sentence', 'https://docs.python.org/3/tutorial/index.html' ])) result = self._evaluate(lambda s: s.str.split(expand=True), s) self.assert_frame_data_equivalent(result, s.str.split(expand=True)) result = self._evaluate(lambda s: s.str.rsplit("/", n=1, expand=True), s) self.assert_frame_data_equivalent( result, s.str.rsplit("/", n=1, expand=True)) @unittest.skipIf(PD_VERSION < (1, 4), "regex arg is new in pandas 1.4") def test_split_pandas_examples_expand_pat_is_string_literal2(self): # when regex is None (default) regex pat is string literal if len(pat) == 1 s = pd.Series(['foojpgbar.jpg']).astype('category') s = s.astype(pd.CategoricalDtype(categories=["foojpgbar.jpg"])) result = self._evaluate(lambda s: s.str.split(r".", expand=True), s) self.assert_frame_data_equivalent(result, s.str.split(r".", expand=True)) # When regex=False, pat is interpreted as the string itself result = self._evaluate( lambda s: s.str.split(r"\.jpg", regex=False, expand=True), s) self.assert_frame_data_equivalent( result, s.str.split(r"\.jpg", regex=False, expand=True)) @unittest.skipIf(PD_VERSION < (1, 4), "regex arg is new in pandas 1.4") def test_split_pandas_examples_expand_pat_is_regex(self): # when regex is None (default) regex pat is compiled if len(pat) != 1 s = pd.Series(["foo and bar plus baz"]) s = s.astype(pd.CategoricalDtype(categories=["foo and bar plus baz"])) result = self._evaluate(lambda s: s.str.split(r"and|plus", expand=True), s) self.assert_frame_data_equivalent( result, s.str.split(r"and|plus", expand=True)) s = pd.Series(['foojpgbar.jpg']).astype('category') s = s.astype(pd.CategoricalDtype(categories=["foojpgbar.jpg"])) result = self._evaluate(lambda s: s.str.split(r"\.jpg", expand=True), s) self.assert_frame_data_equivalent( result, s.str.split(r"\.jpg", expand=True)) # When regex=True, pat is interpreted as a regex result = self._evaluate( lambda s: s.str.split(r"\.jpg", regex=True, expand=True), s) self.assert_frame_data_equivalent( result, s.str.split(r"\.jpg", regex=True, expand=True)) # A compiled regex can be passed as pat result = self._evaluate( lambda s: s.str.split(re.compile(r"\.jpg"), expand=True), s) self.assert_frame_data_equivalent( result, s.str.split(re.compile(r"\.jpg"), expand=True)) @unittest.skipIf(PD_VERSION < (1, 4), "regex arg is new in pandas 1.4") def test_split_pat_is_regex(self): # regex=True, but expand=False s = pd.Series(['foojpgbar.jpg']).astype('category') s = s.astype(pd.CategoricalDtype(categories=["foojpgbar.jpg"])) result = self._evaluate( lambda s: s.str.split(r"\.jpg", regex=True, expand=False), s) self.assert_frame_data_equivalent( result, s.str.split(r"\.jpg", regex=True, expand=False)) def test_astype_categorical_rejected(self): df = pd.DataFrame({'A': np.arange(6), 'B': list('aabbca')}) with self.assertRaisesRegex(frame_base.WontImplementError, r"astype\(dtype='category'\)"): self._evaluate(lambda df: df.B.astype('category'), df) class AllowNonParallelTest(unittest.TestCase): def _use_non_parallel_operation(self): _ = frame_base.DeferredFrame.wrap( expressions.PlaceholderExpression(pd.Series([1, 2, 3]))).replace( 'a', 'b', limit=1) def test_disallow_non_parallel(self): with self.assertRaises(expressions.NonParallelOperation): self._use_non_parallel_operation() def test_allow_non_parallel_in_context(self): with beam.dataframe.allow_non_parallel_operations(): self._use_non_parallel_operation() def test_allow_non_parallel_nesting(self): # disallowed with beam.dataframe.allow_non_parallel_operations(): # allowed self._use_non_parallel_operation() with beam.dataframe.allow_non_parallel_operations(False): # disallowed again with self.assertRaises(expressions.NonParallelOperation): self._use_non_parallel_operation() # allowed self._use_non_parallel_operation() # disallowed with self.assertRaises(expressions.NonParallelOperation): self._use_non_parallel_operation() class ConstructionTimeTest(unittest.TestCase): """Tests for operations that can be executed eagerly.""" DF = pd.DataFrame({ 'str_col': ['foo', 'bar'] * 3, 'int_col': [1, 2] * 3, 'flt_col': [1.1, 2.2] * 3, 'cat_col': pd.Series(list('aabbca'), dtype="category"), 'datetime_col': pd.Series( pd.date_range( '1/1/2000', periods=6, freq='m', tz='America/Los_Angeles')) }) DEFERRED_DF = frame_base.DeferredFrame.wrap( expressions.PlaceholderExpression(DF.iloc[:0])) def _run_test(self, fn): expected = fn(self.DF) actual = fn(self.DEFERRED_DF) if isinstance(expected, pd.Index): pd.testing.assert_index_equal(expected, actual) elif isinstance(expected, pd.Series): pd.testing.assert_series_equal(expected, actual) elif isinstance(expected, pd.DataFrame): pd.testing.assert_frame_equal(expected, actual) else: self.assertEqual(expected, actual) @parameterized.expand(DF.columns) def test_series_name(self, col_name): self._run_test(lambda df: df[col_name].name) @parameterized.expand(DF.columns) def test_series_dtype(self, col_name): self._run_test(lambda df: df[col_name].dtype) self._run_test(lambda df: df[col_name].dtypes) def test_dataframe_columns(self): self._run_test(lambda df: list(df.columns)) def test_dataframe_dtypes(self): self._run_test(lambda df: list(df.dtypes)) def test_categories(self): self._run_test(lambda df: df.cat_col.cat.categories) def test_categorical_ordered(self): self._run_test(lambda df: df.cat_col.cat.ordered) def test_groupby_ndim(self): self._run_test(lambda df: df.groupby('int_col').ndim) def test_groupby_project_ndim(self): self._run_test(lambda df: df.groupby('int_col').flt_col.ndim) self._run_test( lambda df: df.groupby('int_col')[['flt_col', 'str_col']].ndim) def test_get_column_default_None(self): # .get just returns default_value=None at construction time if the column # doesn't exist self._run_test(lambda df: df.get('FOO')) def test_datetime_tz(self): self._run_test(lambda df: df.datetime_col.dt.tz) class DocstringTest(unittest.TestCase): @parameterized.expand([ (frames.DeferredDataFrame, pd.DataFrame), (frames.DeferredSeries, pd.Series), #(frames._DeferredIndex, pd.Index), (frames._DeferredStringMethods, pd.Series.str), ( frames._DeferredCategoricalMethods, pd.core.arrays.categorical.CategoricalAccessor), (frames.DeferredGroupBy, pd.core.groupby.generic.DataFrameGroupBy), (frames._DeferredGroupByCols, pd.core.groupby.generic.DataFrameGroupBy), ( frames._DeferredDatetimeMethods, pd.core.indexes.accessors.DatetimeProperties), ]) def test_docs_defined(self, beam_type, pd_type): beam_attrs = set(dir(beam_type)) pd_attrs = set(dir(pd_type)) docstring_required = sorted([ attr for attr in beam_attrs.intersection(pd_attrs) if getattr(pd_type, attr).__doc__ and not attr.startswith('_') ]) docstring_missing = [ attr for attr in docstring_required if not getattr(beam_type, attr).__doc__ ] self.assertTrue( len(docstring_missing) == 0, f'{beam_type.__name__} is missing a docstring for ' f'{len(docstring_missing)}/{len(docstring_required)} ' f'({len(docstring_missing)/len(docstring_required):%}) ' f'operations:\n{docstring_missing}') class ReprTest(unittest.TestCase): def test_basic_dataframe(self): df = frame_base.DeferredFrame.wrap( expressions.ConstantExpression(GROUPBY_DF)) self.assertEqual( repr(df), ( "DeferredDataFrame(columns=['group', 'foo', 'bar', 'baz', 'bool', " "'str'], index=)")) def test_dataframe_with_named_index(self): df = frame_base.DeferredFrame.wrap( expressions.ConstantExpression(GROUPBY_DF.set_index('group'))) self.assertEqual( repr(df), ( "DeferredDataFrame(columns=['foo', 'bar', 'baz', 'bool', 'str'], " "index='group')")) def test_dataframe_with_partial_named_index(self): df = frame_base.DeferredFrame.wrap( expressions.ConstantExpression( GROUPBY_DF.set_index([GROUPBY_DF.index, 'group']))) self.assertEqual( repr(df), ( "DeferredDataFrame(columns=['foo', 'bar', 'baz', 'bool', 'str'], " "indexes=[, 'group'])")) def test_dataframe_with_named_multi_index(self): df = frame_base.DeferredFrame.wrap( expressions.ConstantExpression(GROUPBY_DF.set_index(['str', 'group']))) self.assertEqual( repr(df), ( "DeferredDataFrame(columns=['foo', 'bar', 'baz', 'bool'], " "indexes=['str', 'group'])")) def test_dataframe_with_multiple_column_levels(self): df = pd.DataFrame({ 'foofoofoo': ['one', 'one', 'one', 'two', 'two', 'two'], 'barbar': ['A', 'B', 'C', 'A', 'B', 'C'], 'bazzy': [1, 2, 3, 4, 5, 6], 'zoop': ['x', 'y', 'z', 'q', 'w', 't'] }) df = df.pivot(index='foofoofoo', columns='barbar') df = frame_base.DeferredFrame.wrap(expressions.ConstantExpression(df)) self.assertEqual( repr(df), ( "DeferredDataFrame(columns=[('bazzy', 'A'), ('bazzy', 'B'), " "('bazzy', 'C'), ('zoop', 'A'), ('zoop', 'B'), ('zoop', 'C')], " "index='foofoofoo')")) def test_dataframe_with_multiple_column_and_multiple_index_levels(self): df = pd.DataFrame({ 'foofoofoo': ['one', 'one', 'one', 'two', 'two', 'two'], 'barbar': ['A', 'B', 'C', 'A', 'B', 'C'], 'bazzy': [1, 2, 3, 4, 5, 6], 'zoop': ['x', 'y', 'z', 'q', 'w', 't'] }) df = df.pivot(index='foofoofoo', columns='barbar') df.index = [['a', 'b'], df.index] # pandas repr displays this: # bazzy zoop # barbar A B C A B C # foofoofoo # a one 1 2 3 x y z # b two 4 5 6 q w t df = frame_base.DeferredFrame.wrap(expressions.ConstantExpression(df)) self.assertEqual( repr(df), ( "DeferredDataFrame(columns=[('bazzy', 'A'), ('bazzy', 'B'), " "('bazzy', 'C'), ('zoop', 'A'), ('zoop', 'B'), ('zoop', 'C')], " "indexes=[, 'foofoofoo'])")) def test_basic_series(self): df = frame_base.DeferredFrame.wrap( expressions.ConstantExpression(GROUPBY_DF['bool'])) self.assertEqual( repr(df), "DeferredSeries(name='bool', dtype=bool, index=)") def test_series_with_named_index(self): df = frame_base.DeferredFrame.wrap( expressions.ConstantExpression(GROUPBY_DF.set_index('group')['str'])) self.assertEqual( repr(df), "DeferredSeries(name='str', dtype=object, index='group')") def test_series_with_partial_named_index(self): df = frame_base.DeferredFrame.wrap( expressions.ConstantExpression( GROUPBY_DF.set_index([GROUPBY_DF.index, 'group'])['bar'])) self.assertEqual( repr(df), ( "DeferredSeries(name='bar', dtype=float64, " "indexes=[, 'group'])")) def test_series_with_named_multi_index(self): df = frame_base.DeferredFrame.wrap( expressions.ConstantExpression( GROUPBY_DF.set_index(['str', 'group'])['baz'])) self.assertEqual( repr(df), "DeferredSeries(name='baz', dtype=float64, indexes=['str', 'group'])") @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') @isolated_env class InteractiveDataFrameTest(unittest.TestCase): def test_collect_merged_dataframes(self): p = beam.Pipeline(InteractiveRunner()) pcoll_1 = ( p | 'Create data 1' >> beam.Create([(1, 'a'), (2, 'b'), (3, 'c'), (4, 'd')]) | 'To rows 1' >> beam.Select(col_1=lambda x: x[0], col_2=lambda x: x[1])) df_1 = to_dataframe(pcoll_1) pcoll_2 = ( p | 'Create data 2' >> beam.Create([(5, 'e'), (6, 'f'), (7, 'g'), (8, 'h')]) | 'To rows 2' >> beam.Select(col_3=lambda x: x[0], col_4=lambda x: x[1])) df_2 = to_dataframe(pcoll_2) df_merged = df_1.merge(df_2, left_index=True, right_index=True) pd_df = ib.collect(df_merged).sort_values(by='col_1') self.assertEqual(pd_df.shape, (4, 4)) self.assertEqual(list(pd_df['col_1']), [1, 2, 3, 4]) self.assertEqual(list(pd_df['col_2']), ['a', 'b', 'c', 'd']) self.assertEqual(list(pd_df['col_3']), [5, 6, 7, 8]) self.assertEqual(list(pd_df['col_4']), ['e', 'f', 'g', 'h']) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/io.py0000644000175100017510000007103315206113032020566 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Sources and sinks for the Beam DataFrame API. Sources ####### This module provides analogs for pandas ``read`` methods, like :func:`pandas.read_csv`. However Beam sources like :func:`read_csv` create a Beam :class:`~apache_beam.PTransform`, and return a :class:`~apache_beam.dataframe.frames.DeferredDataFrame` or :class:`~apache_beam.dataframe.frames.DeferredSeries` representing the contents of the referenced file(s) or data source. The result of these methods must be applied to a :class:`~apache_beam.Pipeline` object, for example:: df = p | beam.dataframe.io.read_csv(...) Sinks ##### This module also defines analogs for pandas sink, or ``to``, methods that generate a Beam :class:`~apache_beam.PTransform`. Users should prefer calling these operations from :class:`~apache_beam.dataframe.frames.DeferredDataFrame` instances (for example with :meth:`DeferredDataFrame.to_csv `). """ import itertools import math import re from io import BytesIO from io import StringIO from io import TextIOWrapper import pandas as pd import apache_beam as beam from apache_beam import io from apache_beam.dataframe import frame_base from apache_beam.io import fileio _DEFAULT_LINES_CHUNKSIZE = 10_000 _DEFAULT_BYTES_CHUNKSIZE = 1 << 20 def read_gbq( table, dataset=None, project_id=None, use_bqstorage_api=False, **kwargs): """This function reads data from a BigQuery table and produces a :class:`~apache_beam.dataframe.frames.DeferredDataFrame. Args: table (str): Please specify a table. This can be done in the format 'PROJECT:dataset.table' if one would not wish to utilize the parameters below. dataset (str): Please specify the dataset (can omit if table was specified as 'PROJECT:dataset.table'). project_id (str): Please specify the project ID (can omit if table was specified as 'PROJECT:dataset.table'). use_bqstorage_api (bool): If you would like to utilize the BigQuery Storage API in ReadFromBigQuery, please set this flag to true. Otherwise, please set flag to false or leave it unspecified. """ if table is None: raise ValueError("Please specify a BigQuery table to read from.") elif len(kwargs) > 0: raise ValueError( f"Encountered unsupported parameter(s) in read_gbq: {kwargs.keys()!r}" "") return _ReadGbq(table, dataset, project_id, use_bqstorage_api) @frame_base.with_docs_from(pd) def read_csv(path, *args, splittable=False, binary=True, **kwargs): """If your files are large and records do not contain quoted newlines, you may pass the extra argument ``splittable=True`` to enable dynamic splitting for this read on newlines. Using this option for records that do contain quoted newlines may result in partial records and data corruption.""" if 'nrows' in kwargs: raise ValueError('nrows not yet supported') filename_column = kwargs.pop('filename_column', None) return _ReadFromPandas( pd.read_csv, path, args, kwargs, incremental=True, binary=binary, splitter=_TextFileSplitter(args, kwargs) if splittable else None, filename_column=filename_column) def _as_pc(df, label=None): from apache_beam.dataframe import convert # avoid circular import # TODO(roberwb): Amortize the computation for multiple writes? return convert.to_pcollection(df, yield_elements='pandas', label=label) @frame_base.with_docs_from(pd.DataFrame) def to_csv(df, path, transform_label=None, *args, **kwargs): label_pc = f"{transform_label} - ToPCollection" if transform_label \ else f"ToPCollection(df) - {path}" label_pd = f"{transform_label} - ToPandasDataFrame" if transform_label \ else f"WriteToPandas(df) - {path}" return _as_pc(df, label_pc) | label_pd >> _WriteToPandas( 'to_csv', path, args, kwargs, incremental=True, binary=False) @frame_base.with_docs_from(pd) def read_fwf(path, *args, **kwargs): return _ReadFromPandas( pd.read_fwf, path, args, kwargs, incremental=True, binary=False, splitter=_TextFileSplitter(args, kwargs)) @frame_base.with_docs_from(pd) def read_json(path, *args, **kwargs): if 'nrows' in kwargs: raise NotImplementedError('nrows not yet supported') elif kwargs.get('lines', False): # Work around https://github.com/pandas-dev/pandas/issues/34548. kwargs = dict(kwargs, nrows=1 << 63) return _ReadFromPandas( pd.read_json, path, args, kwargs, incremental=kwargs.get('lines', False), splitter=_DelimSplitter(b'\n', _DEFAULT_BYTES_CHUNKSIZE) if kwargs.get( 'lines', False) else None, binary=False) @frame_base.with_docs_from(pd.DataFrame) def to_json(df, path, orient=None, *args, **kwargs): if orient is None: if isinstance(df._expr.proxy(), pd.DataFrame): orient = 'columns' elif isinstance(df._expr.proxy(), pd.Series): orient = 'index' else: raise frame_base.WontImplementError('not dataframes or series') kwargs['orient'] = orient return _as_pc(df) | _WriteToPandas( 'to_json', path, args, kwargs, incremental=orient in ('index', 'records', 'values'), binary=False) @frame_base.with_docs_from(pd) def read_html(path, *args, **kwargs): return _ReadFromPandas( lambda *args, **kwargs: pd.read_html(*args, **kwargs)[0], path, args, kwargs) @frame_base.with_docs_from(pd.DataFrame) def to_html(df, path, *args, **kwargs): return _as_pc(df) | _WriteToPandas( 'to_html', path, args, kwargs, incremental=( df._expr.proxy().index.nlevels == 1 or not kwargs.get('sparsify', True)), binary=False) def _binary_reader(format): func = getattr(pd, 'read_%s' % format) result = lambda path, *args, **kwargs: _ReadFromPandas( func, path, args, kwargs) result.__name__ = f'read_{format}' return result def _binary_writer(format): result = ( lambda df, path, *args, **kwargs: _as_pc(df) | _WriteToPandas( f'to_{format}', path, args, kwargs)) result.__name__ = f'to_{format}' return result for format in ('excel', 'feather', 'parquet', 'stata'): globals()['read_%s' % format] = frame_base.with_docs_from(pd)( _binary_reader(format)) globals()['to_%s' % format] = frame_base.with_docs_from(pd.DataFrame)( _binary_writer(format)) for format in ('sas', 'spss'): if hasattr(pd, 'read_%s' % format): # Depends on pandas version. globals()['read_%s' % format] = frame_base.with_docs_from(pd)( _binary_reader(format)) read_clipboard = frame_base.not_implemented_method( 'read_clipboard', base_type=pd) to_clipboard = frame_base.not_implemented_method( 'to_clipboard', base_type=pd.DataFrame) read_msgpack = frame_base.wont_implement_method( pd, 'read_msgpack', reason="deprecated") to_msgpack = frame_base.wont_implement_method( pd.DataFrame, 'to_msgpack', reason="deprecated") read_hdf = frame_base.wont_implement_method( pd, 'read_hdf', explanation="because HDF5 is a random access file format") to_hdf = frame_base.wont_implement_method( pd.DataFrame, 'to_hdf', explanation="because HDF5 is a random access file format") for name in dir(pd): if name.startswith('read_') and name not in globals(): globals()[name] = frame_base.not_implemented_method(name, base_type=pd) def _shift_range_index(offset, df): if isinstance(df.index, pd.RangeIndex): return df.set_index(df.index + offset) else: return df class _ReadFromPandas(beam.PTransform): def __init__( self, reader, path, args, kwargs, binary=True, incremental=False, splitter=False, filename_column=None): if 'compression' in kwargs: raise NotImplementedError('compression') if not isinstance(path, str): raise frame_base.WontImplementError('non-deferred') self.reader = reader self.path = path self.args = args self.kwargs = kwargs self.binary = binary self.incremental = incremental self.splitter = splitter self.filename_column = filename_column def expand(self, root): paths_pcoll = root | beam.Create([self.path]) match = io.filesystems.FileSystems.match([self.path], limits=[1])[0] if not match.metadata_list: # TODO(https://github.com/apache/beam/issues/20858): This should be # allowed for streaming pipelines if user provides an explicit schema. raise FileNotFoundError(f"Found no files that match {self.path!r}") first_path = match.metadata_list[0].path with io.filesystems.FileSystems.open(first_path) as handle: if not self.binary: handle = TextIOWrapper( handle, encoding=self.kwargs.get("encoding", None)) if self.incremental: with self.reader(handle, *self.args, **dict(self.kwargs, chunksize=100)) as stream: sample = next(stream) else: sample = self.reader(handle, *self.args, **self.kwargs) if self.filename_column: sample[self.filename_column] = '' matches_pcoll = paths_pcoll | fileio.MatchAll() indices_pcoll = ( matches_pcoll.pipeline | 'DoOnce' >> beam.Create([None]) | beam.Map( lambda _, paths: { path: ix for ix, path in enumerate(sorted(paths)) }, paths=beam.pvalue.AsList( matches_pcoll | beam.Map(lambda match: match.path)))) pcoll = ( matches_pcoll | beam.Reshuffle() | fileio.ReadMatches() | beam.ParDo( _ReadFromPandasDoFn( self.reader, self.args, self.kwargs, self.binary, self.incremental, self.splitter, self.filename_column), path_indices=beam.pvalue.AsSingleton(indices_pcoll))) from apache_beam.dataframe import convert return convert.to_dataframe(pcoll, proxy=sample[:0]) class _Splitter: def empty_buffer(self): """Returns an empty buffer of the right type (string or bytes). """ raise NotImplementedError(self) def read_header(self, handle): """Reads the header from handle, which points to the start of the file. Returns the pair (header, buffer) where buffer contains any part of the file that was "overread" from handle while seeking the end of header. """ raise NotImplementedError(self) def read_to_record_boundary(self, buffered, handle): """Reads the given handle up to the end of the current record. The buffer argument represents bytes that were read previously; logically it's as if these were pushed back into handle for reading. If the record end is within buffered, it's possible that no more bytes will be read from handle at all. Returns the pair (remaining_record_bytes, buffer) where buffer contains any part of the file that was "overread" from handle while seeking the end of the record. """ raise NotImplementedError(self) class _DelimSplitter(_Splitter): """A _Splitter that splits on delimiters between records. This delimiter is assumed ot never occur within a record. """ def __init__(self, delim, read_chunk_size=_DEFAULT_BYTES_CHUNKSIZE): # Multi-char delimiters would require more care across chunk boundaries. assert len(delim) == 1 self._delim = delim self._empty = delim[:0] self._read_chunk_size = read_chunk_size def empty_buffer(self): return self._empty def read_header(self, handle): return self._empty, self._empty def read_to_record_boundary(self, buffered, handle): if self._delim in buffered: ix = buffered.index(self._delim) + len(self._delim) return buffered[:ix], buffered[ix:] else: while True: chunk = handle.read(self._read_chunk_size) if self._delim in chunk: ix = chunk.index(self._delim) + len(self._delim) return buffered + chunk[:ix], chunk[ix:] elif not chunk: return buffered, self._empty else: buffered += chunk def _maybe_encode(str_or_bytes): if isinstance(str_or_bytes, str): return str_or_bytes.encode('utf-8') else: return str_or_bytes class _TextFileSplitter(_DelimSplitter): """Splitter for dynamically sharding CSV files and newline record boundaries. Currently does not handle quoted newlines, so is off by default, but such support could be added in the future. """ def __init__(self, args, kwargs, read_chunk_size=_DEFAULT_BYTES_CHUNKSIZE): if args: # TODO(robertwb): Automatically populate kwargs as we do for df methods. raise ValueError( 'Non-path arguments must be passed by keyword ' 'for splittable csv reads.') if kwargs.get('skipfooter', 0): raise ValueError('Splittablility incompatible with skipping footers.') super().__init__( _maybe_encode(kwargs.get('lineterminator', b'\n')), _DEFAULT_BYTES_CHUNKSIZE) self._kwargs = kwargs def read_header(self, handle): if self._kwargs.get('header', 'infer') == 'infer': if 'names' in self._kwargs: header = None else: header = 0 else: header = self._kwargs['header'] if header is None: return self._empty, self._empty if isinstance(header, int): max_header = header else: max_header = max(header) skiprows = self._kwargs.get('skiprows', 0) if isinstance(skiprows, int): is_skiprow = lambda ix: ix < skiprows elif callable(skiprows): is_skiprow = skiprows elif skiprows is None: is_skiprow = lambda ix: False else: is_skiprow = lambda ix: ix in skiprows comment = _maybe_encode(self._kwargs.get('comment', None)) if comment: is_comment = lambda line: line.startswith(comment) else: is_comment = lambda line: False skip_blank_lines = self._kwargs.get('skip_blank_lines', True) if skip_blank_lines: is_blank = lambda line: re.match(rb'^\s*$', line) else: is_blank = lambda line: False text_header = b'' rest = b'' skipped = 0 for ix in itertools.count(): line, rest = self.read_to_record_boundary(rest, handle) text_header += line if is_skiprow(ix) or is_blank(line) or is_comment(line): skipped += 1 continue if ix - skipped == max_header: return text_header, rest class _TruncatingFileHandle(object): """A wrapper of a file-like object representing the restriction of the underling handle according to the given SDF restriction tracker, breaking the file only after the given delimiter. For example, if the underling restriction is [103, 607) and each line were exactly 10 characters long (i.e. every 10th charcter was a newline), then this would give a view of a 500-byte file consisting of bytes bytes 110 to 609 (inclusive) of the underlying file. As with all SDF trackers, the endpoint may change dynamically during reading. """ def __init__(self, underlying, tracker, splitter): self._underlying = underlying self._tracker = tracker self._splitter = splitter self._empty = self._splitter.empty_buffer() self._done = False self._header, self._buffer = self._splitter.read_header(self._underlying) self._buffer_start_pos = len(self._header) self._iterator = None start = self._tracker.current_restriction().start # Seek to first delimiter after the start position. if start > len(self._header): if start > len(self._header) + len(self._buffer): self._buffer_start_pos = start self._buffer = self._empty self._underlying.seek(start) else: self._buffer_start_pos = start self._buffer = self._buffer[start - len(self._header):] skip, self._buffer = self._splitter.read_to_record_boundary( self._buffer, self._underlying) self._buffer_start_pos += len(skip) @property def mode(self): return getattr(self._underlying, "mode", "r") def readable(self): return True def writable(self): return False def seekable(self): return False @property def closed(self): return False def __iter__(self): # For pandas is_file_like. return self def __next__(self): if self._iterator is None: self._iterator = self._line_iterator() return next(self._iterator) def readline(self): # This attribute is checked, but unused, by pandas. return next(self) def _line_iterator(self): line_start = 0 chunk = self._read() while True: line_end = chunk.find(self._splitter._delim, line_start) while line_end == -1: more = self._read() if not more: if line_start < len(chunk): yield chunk[line_start:] return chunk = chunk[line_start:] + more line_start = 0 line_end = chunk.find(self._splitter._delim, line_start) yield chunk[line_start:line_end + 1] line_start = line_end + 1 def read(self, size=-1): if self._iterator: raise NotImplementedError('Cannot call read after iterating.') return self._read(size) def _read(self, size=-1): if self._header: res = self._header self._header = None return res elif self._done: return self._empty elif size == -1: self._buffer += self._underlying.read() elif not self._buffer: self._buffer = self._underlying.read(size) if not self._buffer: self._done = True return self._empty if self._tracker.try_claim(self._buffer_start_pos + len(self._buffer)): res = self._buffer self._buffer = self._empty self._buffer_start_pos += len(res) else: offset = self._tracker.current_restriction().stop - self._buffer_start_pos if offset <= 0: res = self._empty else: rest, _ = self._splitter.read_to_record_boundary( self._buffer[offset:], self._underlying) res = self._buffer[:offset] + rest self._done = True return res def flush(self): self._underlying.flush() class _ReadFromPandasDoFn(beam.DoFn, beam.RestrictionProvider): def __init__( self, reader, args, kwargs, binary, incremental, splitter, filename_column=None): # avoid pickling issues if reader.__module__.startswith('pandas.'): reader = reader.__name__ self.reader = reader self.args = args self.kwargs = kwargs self.binary = binary self.incremental = incremental self.splitter = splitter self.filename_column = filename_column def initial_restriction(self, readable_file): return beam.io.restriction_trackers.OffsetRange( 0, readable_file.metadata.size_in_bytes) def restriction_size(self, readable_file, restriction): return restriction.size() def create_tracker(self, restriction): tracker = beam.io.restriction_trackers.OffsetRestrictionTracker(restriction) if self.splitter: return tracker else: return beam.io.restriction_trackers.UnsplittableRestrictionTracker( tracker) def process( self, readable_file, path_indices, tracker=beam.DoFn.RestrictionParam()): reader = self.reader if isinstance(reader, str): reader = getattr(pd, self.reader) indices_per_file = 10**int(math.log(2**63 // len(path_indices), 10)) if readable_file.metadata.size_in_bytes > indices_per_file: raise RuntimeError( f'Cannot safely index records from {len(path_indices)} files ' f'of size {readable_file.metadata.size_in_bytes} ' f'as their product is greater than 2^63.') start_index = ( tracker.current_restriction().start + path_indices[readable_file.metadata.path] * indices_per_file) with readable_file.open() as handle: if self.incremental: # TODO(robertwb): We could consider trying to get progress for # non-incremental sources that are read linearly, as long as they # don't try to seek. This could be deceptive as progress would # advance to 100% the instant the (large) read was done, discounting # any downstream processing. handle = _TruncatingFileHandle( handle, tracker, splitter=self.splitter or _DelimSplitter(b'\n', _DEFAULT_BYTES_CHUNKSIZE)) if not self.binary: handle = TextIOWrapper( handle, encoding=self.kwargs.get("encoding", None)) if self.incremental: if 'chunksize' not in self.kwargs: self.kwargs['chunksize'] = _DEFAULT_LINES_CHUNKSIZE frames = reader(handle, *self.args, **self.kwargs) else: frames = [reader(handle, *self.args, **self.kwargs)] for df in frames: if self.filename_column: df[self.filename_column] = readable_file.metadata.path yield _shift_range_index(start_index, df) if not self.incremental: # Satisfy the SDF contract by claiming the whole range. # Do this after emitting the frames to avoid advancing progress to 100% # prior to that. tracker.try_claim(tracker.current_restriction().stop) class _WriteToPandas(beam.PTransform): def __init__( self, writer, path, args, kwargs, incremental=False, binary=True): self.writer = writer self.path = path self.args = args self.kwargs = kwargs self.incremental = incremental self.binary = binary def expand(self, pcoll): if 'file_naming' in self.kwargs: dir, name = self.path, '' else: dir, name = io.filesystems.FileSystems.split(self.path) return pcoll | fileio.WriteToFiles( path=dir, shards=self.kwargs.pop('num_shards', None), file_naming=self.kwargs.pop( 'file_naming', fileio.default_file_naming(name)), sink=lambda _: _WriteToPandasFileSink( self.writer, self.args, self.kwargs, self.incremental, self.binary)) class _WriteToPandasFileSink(fileio.FileSink): def __init__(self, writer, args, kwargs, incremental, binary): if 'compression' in kwargs: raise NotImplementedError('compression') self.writer = writer self.args = args self.kwargs = kwargs self.incremental = incremental self.binary = binary self.StringOrBytesIO = BytesIO if binary else StringIO if incremental: self.write = self.write_record_incremental self.flush = self.close_incremental else: self.write = self.buffer_record self.flush = self.flush_buffer def open(self, file_handle): self.buffer = [] self.empty = self.header = self.footer = None if not self.binary: file_handle = TextIOWrapper( file_handle, encoding=self.kwargs.get("encoding", None)) self.file_handle = file_handle def write_to(self, df, file_handle=None): non_none_handle = file_handle or self.StringOrBytesIO() getattr(df, self.writer)(non_none_handle, *self.args, **self.kwargs) if file_handle is None: return non_none_handle.getvalue() def write_record_incremental(self, value): if self.empty is None: self.empty = self.write_to(value[:0]) if self.header is None and len(value): def new_value(ix): if isinstance(ix, tuple): return (new_value(ix[0]), ) + ix[1:] else: return str('x') + '_again' def change_index(df): df.index = df.index.map(new_value) return df one_row = self.write_to(value[:1]) another_row = self.write_to(change_index(value[:1])) two_rows = self.write_to(pd.concat([value[:1], change_index(value[:1])])) for ix, c in enumerate(self.empty): if one_row[ix] != c: break else: ix = len(self.empty) self.header = self.empty[:ix] self.footer = self.empty[ix:] self.delimiter = two_rows[len(one_row) - len(self.footer):-( len(another_row) - len(self.header)) or None] self.file_handle.write(self.header) self.first = True if len(value): if self.first: self.first = False else: self.file_handle.write(self.delimiter) # IDEA(robertwb): Construct a "truncating" stream wrapper to avoid the # in-memory copy. rows = self.write_to(value) self.file_handle.write(rows[len(self.header):-len(self.footer) or None]) def close_incremental(self): if self.footer is not None: self.file_handle.write(self.footer) elif self.empty is not None: self.file_handle.write(self.empty) self.file_handle.flush() def buffer_record(self, value): self.buffer.append(value) def flush_buffer(self): if self.buffer: self.write_to(pd.concat(self.buffer), self.file_handle) self.file_handle.flush() class ReadViaPandas(beam.PTransform): def __init__( self, format, *args, include_indexes=False, objects_as_strings=True, filename_column=None, **kwargs): if format == 'csv': kwargs['filename_column'] = filename_column self._reader = globals()['read_%s' % format](*args, **kwargs) self._include_indexes = include_indexes self._objects_as_strings = objects_as_strings self._filename_column = filename_column def expand(self, p): from apache_beam.dataframe import convert # avoid circular import df = p | self._reader if self._objects_as_strings: for col, t in zip(df.columns, df.dtypes): if t == object: df[col] = df[col].astype(pd.StringDtype()) return convert.to_pcollection(df, include_indexes=self._include_indexes) class WriteViaPandas(beam.PTransform): def __init__(self, format, *args, **kwargs): self._writer_func = globals()['to_%s' % format] self._args = args self._kwargs = kwargs def expand(self, pcoll): from apache_beam.dataframe import convert # avoid circular import return { 'files_written': self._writer_func( convert.to_dataframe(pcoll), *self._args, **self._kwargs) | beam.Map(lambda file_result: file_result.file_name).with_output_types( str) } class _ReadGbq(beam.PTransform): """Read data from BigQuery with output type 'BEAM_ROW', then convert it into a deferred dataframe. This PTransform wraps the Python ReadFromBigQuery PTransform, and sets the output_type as 'BEAM_ROW' to convert into a Beam Schema. Once applied to a pipeline object, it is passed into the to_dataframe() function to convert the PCollection into a deferred dataframe. This PTransform currently does not support queries. Args: table (str): The ID of the table. The ID must contain only letters ``a-z``, ``A-Z``, numbers ``0-9``, underscores ``_`` or white spaces. Note that the table argument must contain the entire table reference specified as: ``'PROJECT:DATASET.TABLE'``. use_bq_storage_api (bool): The method to use to read from BigQuery. It may be 'EXPORT' or 'DIRECT_READ'. EXPORT invokes a BigQuery export request (https://cloud.google.com/bigquery/docs/exporting-data). 'DIRECT_READ' reads directly from BigQuery storage using the BigQuery Read API (https://cloud.google.com/bigquery/docs/reference/storage). If unspecified or set to false, the default is currently utilized (EXPORT). If the flag is set to true, 'DIRECT_READ' will be utilized.""" def __init__( self, table=None, dataset_id=None, project_id=None, use_bqstorage_api=None): self.table = table self.dataset_id = dataset_id self.project_id = project_id self.use_bqstorage_api = use_bqstorage_api def expand(self, root): from apache_beam.dataframe import convert # avoid circular import if self.use_bqstorage_api: method = 'DIRECT_READ' else: method = 'EXPORT' return convert.to_dataframe( root | '_DataFrame_Read_From_BigQuery' >> beam.io.ReadFromBigQuery( table=self.table, dataset=self.dataset_id, project=self.project_id, method=method, output_type='BEAM_ROW')) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/io_it_test.py0000644000175100017510000001061315206113032022316 0ustar00runnerrunner#!/usr/bin/env python # -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration tests for Dataframe sources and sinks.""" # pytype: skip-file import logging import unittest import pytest import apache_beam.io.gcp.bigquery from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to _LOGGER = logging.getLogger(__name__) try: from google.api_core.exceptions import GoogleAPICallError except ImportError: GoogleAPICallError = None @unittest.skipIf( GoogleAPICallError is None, 'GCP dependencies are not installed') class ReadUsingReadGbqTests(unittest.TestCase): @pytest.mark.it_postcommit def test_ReadGbq(self): from apache_beam.dataframe import convert with TestPipeline(is_integration_test=True) as p: actual_df = p | apache_beam.dataframe.io.read_gbq( table="apache-beam-testing:beam_bigquery_io_test." "dfsqltable_3c7d6fd5_16e0460dfd0", use_bqstorage_api=False) assert_that( convert.to_pcollection(actual_df), equal_to([(3, 'customer1', 'test'), (1, 'customer1', 'test'), (2, 'customer2', 'test'), (4, 'customer2', 'test')])) @pytest.mark.it_postcommit def test_ReadGbq_export_with_project(self): from apache_beam.dataframe import convert with TestPipeline(is_integration_test=True) as p: actual_df = p | apache_beam.dataframe.io.read_gbq( table="dfsqltable_3c7d6fd5_16e0460dfd0", dataset="beam_bigquery_io_test", project_id="apache-beam-testing", use_bqstorage_api=False) assert_that( convert.to_pcollection(actual_df), equal_to([(3, 'customer1', 'test'), (1, 'customer1', 'test'), (2, 'customer2', 'test'), (4, 'customer2', 'test')])) @pytest.mark.it_postcommit def test_ReadGbq_direct_read(self): from apache_beam.dataframe import convert with TestPipeline(is_integration_test=True) as p: actual_df = p | apache_beam.dataframe.io.\ read_gbq( table= "apache-beam-testing:beam_bigquery_io_test." "dfsqltable_3c7d6fd5_16e0460dfd0", use_bqstorage_api=True) assert_that( convert.to_pcollection(actual_df), equal_to([(3, 'customer1', 'test'), (1, 'customer1', 'test'), (2, 'customer2', 'test'), (4, 'customer2', 'test')])) @pytest.mark.it_postcommit def test_ReadGbq_direct_read_with_project(self): from apache_beam.dataframe import convert with TestPipeline(is_integration_test=True) as p: actual_df = p | apache_beam.dataframe.io.read_gbq( table="dfsqltable_3c7d6fd5_16e0460dfd0", dataset="beam_bigquery_io_test", project_id="apache-beam-testing", use_bqstorage_api=True) assert_that( convert.to_pcollection(actual_df), equal_to([(3, 'customer1', 'test'), (1, 'customer1', 'test'), (2, 'customer2', 'test'), (4, 'customer2', 'test')])) @pytest.mark.it_postcommit def test_ReadGbq_with_computation(self): from apache_beam.dataframe import convert with TestPipeline(is_integration_test=True) as p: beam_df = p | apache_beam.dataframe.io.read_gbq( table="dfsqltable_3c7d6fd5_16e0460dfd0", dataset="beam_bigquery_io_test", project_id="apache-beam-testing") actual_df = beam_df.groupby('id').count() assert_that( convert.to_pcollection(actual_df, include_indexes=True), equal_to([(1, 1, 1), (2, 1, 1), (3, 1, 1), (4, 1, 1)])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/io_test.py0000644000175100017510000004250315206113032021625 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import glob import importlib import math import os import platform import shutil import tempfile import typing import unittest from datetime import datetime from io import BytesIO from io import StringIO import mock import pandas as pd import pandas.testing import pyarrow import pytest from pandas.testing import assert_frame_equal from parameterized import parameterized import apache_beam as beam import apache_beam.io.gcp.bigquery from apache_beam.dataframe import convert from apache_beam.dataframe import io from apache_beam.io import fileio from apache_beam.io import restriction_trackers from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to try: from google.api_core.exceptions import GoogleAPICallError except ImportError: GoogleAPICallError = None # Get major, minor version PD_VERSION = tuple(map(int, pd.__version__.split('.')[0:2])) PYARROW_VERSION = tuple(map(int, pyarrow.__version__.split('.')[0:2])) class SimpleRow(typing.NamedTuple): value: int class MyRow(typing.NamedTuple): timestamp: int value: int @unittest.skipIf( platform.system() == 'Windows', 'https://github.com/apache/beam/issues/20642') class IOTest(unittest.TestCase): def setUp(self): self._temp_roots = [] def tearDown(self): for root in self._temp_roots: shutil.rmtree(root) def temp_dir(self, files=None): dir = tempfile.mkdtemp(prefix='beam-test') self._temp_roots.append(dir) if files: for name, contents in files.items(): with open(os.path.join(dir, name), 'w') as fout: fout.write(contents) return dir + os.path.sep def read_all_lines(self, pattern, delete=False): for path in glob.glob(pattern): with open(path) as fin: # TODO(Py3): yield from for line in fin: yield line.rstrip('\n') if delete: os.remove(path) def test_read_fwf(self): input = self.temp_dir( {'all.fwf': ''' A B 11a 0 37a 1 389a 2 '''.strip()}) with beam.Pipeline() as p: df = p | io.read_fwf(input + 'all.fwf') rows = convert.to_pcollection(df) | beam.Map(tuple) assert_that(rows, equal_to([('11a', 0), ('37a', 1), ('389a', 2)])) def test_read_write_csv(self): input = self.temp_dir({'1.csv': 'a,b\n1,2\n', '2.csv': 'a,b\n3,4\n'}) output = self.temp_dir() with beam.Pipeline() as p: df = p | io.read_csv(input + '*.csv') df['c'] = df.a + df.b df.to_csv(output + 'out.csv', index=False) self.assertCountEqual(['a,b,c', '1,2,3', '3,4,7'], set(self.read_all_lines(output + 'out.csv*'))) def test_wide_csv_with_dtypes(self): # Verify https://github.com/apache/beam/issues/31152 is resolved. cols = ','.join(f'col{ix}' for ix in range(123)) data = ','.join(str(ix) for ix in range(123)) input = self.temp_dir({'tmp.csv': f'{cols}\n{data}'}) with beam.Pipeline() as p: pcoll = p | beam.io.ReadFromCsv(f'{input}tmp.csv', dtype=str) assert_that(pcoll | beam.Map(max), equal_to(['99'])) def test_sharding_parameters(self): data = pd.DataFrame({'label': ['11a', '37a', '389a'], 'rank': [0, 1, 2]}) output = self.temp_dir() with beam.Pipeline() as p: df = convert.to_dataframe(p | beam.Create([data]), proxy=data[:0]) df.to_csv( output, num_shards=1, file_naming=fileio.single_file_naming('out.csv')) self.assertEqual(glob.glob(output + '*'), [output + 'out.csv']) @pytest.mark.uses_pyarrow @unittest.skipIf( PD_VERSION >= (1, 4) and PYARROW_VERSION < (1, 0), "pandas 1.4 requires at least pyarrow 1.0.1") def test_read_write_parquet(self): self._run_read_write_test( 'parquet', {}, {}, dict(check_index=False), ['pyarrow']) @parameterized.expand([ ('csv', dict(index_col=0)), ('csv', dict(index_col=0, splittable=True)), ('json', dict(orient='index'), dict(orient='index')), ('json', dict(orient='columns'), dict(orient='columns')), ('json', dict(orient='split'), dict(orient='split')), ( 'json', dict(orient='values'), dict(orient='values'), dict(check_index=False, check_names=False)), ( 'json', dict(orient='records'), dict(orient='records'), dict(check_index=False)), ( 'json', dict(orient='records', lines=True), dict(orient='records', lines=True), dict(check_index=False)), ('html', dict(index_col=0), {}, {}, ['lxml']), ('excel', dict(index_col=0), {}, {}, ['openpyxl', 'xlrd']), ]) # pylint: disable=dangerous-default-value def test_read_write( self, format, read_kwargs={}, write_kwargs={}, check_options={}, requires=()): self._run_read_write_test( format, read_kwargs, write_kwargs, check_options, requires) # pylint: disable=dangerous-default-value def _run_read_write_test( self, format, read_kwargs={}, write_kwargs={}, check_options={}, requires=()): for module in requires: try: importlib.import_module(module) except ImportError: raise unittest.SkipTest('Missing dependency: %s' % module) small = pd.DataFrame({'label': ['11a', '37a', '389a'], 'rank': [0, 1, 2]}) big = pd.DataFrame({'number': list(range(1000))}) big['float'] = big.number.map(math.sqrt) big['text'] = big.number.map(lambda n: 'f' + 'o' * n) def frame_equal_to(expected_, check_index=True, check_names=True): def check(actual): expected = expected_ try: actual = pd.concat(actual) if not check_index: expected = expected.sort_values(list( expected.columns)).reset_index(drop=True) actual = actual.sort_values(list( actual.columns)).reset_index(drop=True) if not check_names: actual = actual.rename( columns=dict(zip(actual.columns, expected.columns))) return assert_frame_equal(expected, actual, check_like=True) except: print("EXPECTED") print(expected) print("ACTUAL") print(actual) raise return check for df in (small, big): with tempfile.TemporaryDirectory() as dir: dest = os.path.join(dir, 'out') try: with beam.Pipeline() as p: deferred_df = convert.to_dataframe( p | beam.Create([df[::3], df[1::3], df[2::3]]), proxy=df[:0]) # This does the write. getattr(deferred_df, 'to_%s' % format)(dest, **write_kwargs) with beam.Pipeline() as p: # Now do the read. # TODO(robertwb): Allow reading from pcoll of paths to do it all in # one pipeline. result = convert.to_pcollection( p | getattr(io, 'read_%s' % format)(dest + '*', **read_kwargs), yield_elements='pandas') assert_that(result, frame_equal_to(df, **check_options)) except: os.system('head -n 100 ' + dest + '*') raise def _run_truncating_file_handle_test( self, s, splits, delim=' ', chunk_size=10): split_results = [] next_range = restriction_trackers.OffsetRange(0, len(s)) for split in list(splits) + [None]: tracker = restriction_trackers.OffsetRestrictionTracker(next_range) handle = io._TruncatingFileHandle( StringIO(s), tracker, splitter=io._DelimSplitter(delim, chunk_size)) data = '' chunk = handle.read(1) if split is not None: _, next_range = tracker.try_split(split) while chunk: data += chunk chunk = handle.read(7) split_results.append(data) return split_results def test_truncating_filehandle(self): self.assertEqual( self._run_truncating_file_handle_test('a b c d e', [0.5]), ['a b c ', 'd e']) self.assertEqual( self._run_truncating_file_handle_test('aaaaaaaaaaaaaaXaaa b', [0.5]), ['aaaaaaaaaaaaaaXaaa ', 'b']) self.assertEqual( self._run_truncating_file_handle_test( 'aa bbbbbbbbbbbbbbbbbbbbbbbbbb ccc ', [0.01, 0.5]), ['aa ', 'bbbbbbbbbbbbbbbbbbbbbbbbbb ', 'ccc ']) numbers = 'x'.join(str(k) for k in range(1000)) splits = self._run_truncating_file_handle_test( numbers, [0.1] * 20, delim='x') self.assertEqual(numbers, ''.join(splits)) self.assertTrue(s.endswith('x') for s in splits[:-1]) self.assertLess(max(len(s) for s in splits), len(numbers) * 0.9 + 10) self.assertGreater( min(len(s) for s in splits), len(numbers) * 0.9**20 * 0.1) def _run_truncating_file_handle_iter_test(self, s, delim=' ', chunk_size=10): tracker = restriction_trackers.OffsetRestrictionTracker( restriction_trackers.OffsetRange(0, len(s))) handle = io._TruncatingFileHandle( StringIO(s), tracker, splitter=io._DelimSplitter(delim, chunk_size)) self.assertEqual(s, ''.join(list(handle))) def test_truncating_filehandle_iter(self): self._run_truncating_file_handle_iter_test('a b c') self._run_truncating_file_handle_iter_test('aaaaaaaaaaaaaaaaaaaa b ccc') self._run_truncating_file_handle_iter_test('aaa b cccccccccccccccccccc') self._run_truncating_file_handle_iter_test('aaa b ccccccccccccccccc ') @parameterized.expand([ ('defaults', {}), ('header', dict(header=1)), ('multi_header', dict(header=[0, 1])), ('multi_header', dict(header=[0, 1, 4])), ('names', dict(names=('m', 'n', 'o'))), ('names_and_header', dict(names=('m', 'n', 'o'), header=0)), ('skip_blank_lines', dict(header=4, skip_blank_lines=True)), ('skip_blank_lines', dict(header=4, skip_blank_lines=False)), ('comment', dict(comment='X', header=4)), ('comment', dict(comment='X', header=[0, 3])), ('skiprows', dict(skiprows=0, header=[0, 1])), ('skiprows', dict(skiprows=[1], header=[0, 3], skip_blank_lines=False)), ('skiprows', dict(skiprows=[0, 1], header=[0, 1], comment='X')), ]) def test_csv_splitter(self, name, kwargs): def assert_frame_equal(expected, actual): try: pandas.testing.assert_frame_equal(expected, actual) except AssertionError: print("Expected:\n", expected) print("Actual:\n", actual) raise def read_truncated_csv(start, stop): return pd.read_csv( io._TruncatingFileHandle( BytesIO(contents.encode('ascii')), restriction_trackers.OffsetRestrictionTracker( restriction_trackers.OffsetRange(start, stop)), splitter=io._TextFileSplitter((), kwargs, read_chunk_size=7)), index_col=0, **kwargs) contents = ''' a0, a1, a2 b0, b1, b2 X , c1, c2 e0, e1, e2 f0, f1, f2 w, daaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaata, w x, daaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaata, x y, daaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaata, y z, daaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaata, z '''.strip() expected = pd.read_csv(StringIO(contents), index_col=0, **kwargs) one_shard = read_truncated_csv(0, len(contents)) assert_frame_equal(expected, one_shard) equal_shards = pd.concat([ read_truncated_csv(0, len(contents) // 2), read_truncated_csv(len(contents) // 2, len(contents)), ]) assert_frame_equal(expected, equal_shards) three_shards = pd.concat([ read_truncated_csv(0, len(contents) // 3), read_truncated_csv(len(contents) // 3, len(contents) * 2 // 3), read_truncated_csv(len(contents) * 2 // 3, len(contents)), ]) assert_frame_equal(expected, three_shards) # https://github.com/pandas-dev/pandas/issues/38292 if not isinstance(kwargs.get('header'), list): split_in_header = pd.concat([ read_truncated_csv(0, 1), read_truncated_csv(1, len(contents)), ]) assert_frame_equal(expected, split_in_header) if not kwargs: # Make sure we're correct as we cross the header boundary. # We don't need to do this for every permutation. header_end = contents.index('a2') + 3 for split in range(header_end - 2, header_end + 2): split_at_header = pd.concat([ read_truncated_csv(0, split), read_truncated_csv(split, len(contents)), ]) assert_frame_equal(expected, split_at_header) def test_file_not_found(self): with self.assertRaisesRegex(FileNotFoundError, r'/tmp/fake_dir/\*\*'): with beam.Pipeline() as p: _ = p | io.read_csv('/tmp/fake_dir/**') def test_windowed_write(self): output = self.temp_dir() with beam.Pipeline() as p: pc = ( p | beam.Create([MyRow(timestamp=i, value=i % 3) for i in range(20)]) | beam.Map(lambda v: beam.window.TimestampedValue(v, v.timestamp)). with_output_types(MyRow) | beam.WindowInto( beam.window.FixedWindows(10)).with_output_types(MyRow)) deferred_df = convert.to_dataframe(pc) deferred_df.to_csv(output + 'out.csv', index=False) first_window_files = ( f'{output}out.csv-' f'{datetime.utcfromtimestamp(0).isoformat()}*') self.assertCountEqual( ['timestamp,value'] + [f'{i},{i % 3}' for i in range(10)], set(self.read_all_lines(first_window_files, delete=True))) second_window_files = ( f'{output}out.csv-' f'{datetime.utcfromtimestamp(10).isoformat()}*') self.assertCountEqual( ['timestamp,value'] + [f'{i},{i%3}' for i in range(10, 20)], set(self.read_all_lines(second_window_files, delete=True))) # Check that we've read (and removed) every output file self.assertEqual(len(glob.glob(f'{output}out.csv*')), 0) def test_double_write(self): output = self.temp_dir() with beam.Pipeline() as p: pc1 = p | 'create pc1' >> beam.Create( [SimpleRow(value=i) for i in [1, 2]]) pc2 = p | 'create pc2' >> beam.Create( [SimpleRow(value=i) for i in [3, 4]]) deferred_df1 = convert.to_dataframe(pc1) deferred_df2 = convert.to_dataframe(pc2) deferred_df1.to_csv( f'{output}out1.csv', transform_label="Writing to csv PC1", index=False) deferred_df2.to_csv( f'{output}out2.csv', transform_label="Writing to csv PC2", index=False) self.assertCountEqual(['value', '1', '2'], set(self.read_all_lines(output + 'out1.csv*'))) self.assertCountEqual(['value', '3', '4'], set(self.read_all_lines(output + 'out2.csv*'))) @unittest.skipIf( GoogleAPICallError is None, 'GCP dependencies are not installed') class ReadGbqTransformTests(unittest.TestCase): @mock.patch.object(BigQueryWrapper, 'get_table') def test_bad_schema_public_api_direct_read(self, get_table): try: bigquery.TableFieldSchema except AttributeError: raise ValueError('Please install GCP Dependencies.') fields = [ bigquery.TableFieldSchema(name='stn', type='DOUBLE', mode="NULLABLE"), bigquery.TableFieldSchema(name='temp', type='FLOAT64', mode="REPEATED"), bigquery.TableFieldSchema(name='count', type='INTEGER', mode=None) ] schema = bigquery.TableSchema(fields=fields) table = apache_beam.io.gcp.internal.clients.bigquery. \ bigquery_v2_messages.Table( schema=schema) get_table.return_value = table with self.assertRaisesRegex(ValueError, "Encountered an unsupported type: 'DOUBLE'"): p = apache_beam.Pipeline() _ = p | apache_beam.dataframe.io.read_gbq( table="dataset.sample_table", use_bqstorage_api=True) def test_unsupported_callable(self): def filterTable(table): if table is not None: return table res = filterTable with self.assertRaisesRegex(TypeError, 'ReadFromBigQuery: table must be of type string' '; got a callable instead'): p = beam.Pipeline() _ = p | beam.dataframe.io.read_gbq(table=res) def test_ReadGbq_unsupported_param(self): with self.assertRaisesRegex(ValueError, r"""Encountered unsupported parameter\(s\) """ r"""in read_gbq: dict_keys\(\['reauth']\)"""): p = beam.Pipeline() _ = p | beam.dataframe.io.read_gbq( table="table", use_bqstorage_api=False, reauth="true_config") if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/pandas_docs_test.py0000644000175100017510000001074515206113032023477 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """A module for running the pandas docs (such as the users guide) against our dataframe implementation. Run as python -m apache_beam.dataframe.pandas_docs_test [getting_started ...] """ import argparse import contextlib import io import multiprocessing import os import sys import time import urllib.request import zipfile from apache_beam.dataframe import doctests PANDAS_VERSION = '1.1.1' PANDAS_DIR = os.path.expanduser("~/.apache_beam/cache/pandas-" + PANDAS_VERSION) PANDAS_DOCS_SOURCE = os.path.join(PANDAS_DIR, 'doc', 'source') def main(): parser = argparse.ArgumentParser() parser.add_argument( '-p', '--parallel', type=int, default=0, help='Number of tests to run in parallel. ' 'Defaults to 0, meaning the number of cores on the machine.') parser.add_argument('docs', nargs='*') args = parser.parse_args() if not os.path.exists(PANDAS_DIR): # Download the pandas source. os.makedirs(os.path.dirname(PANDAS_DIR), exist_ok=True) zip = os.path.join(PANDAS_DIR + '.zip') if not os.path.exists(zip): url = ( 'https://github.com/pandas-dev/pandas/archive/v%s.zip' % PANDAS_VERSION) print('Downloading', url) with urllib.request.urlopen(url) as fin: with open(zip + '.tmp', 'wb') as fout: fout.write(fin.read()) os.rename(zip + '.tmp', zip) print('Extracting', zip) with zipfile.ZipFile(zip, 'r') as handle: handle.extractall(os.path.dirname(PANDAS_DIR)) tests = args.docs or ['getting_started', 'user_guide'] paths = [] filters = [] # Explicit paths. for test in tests: if os.path.exists(test): paths.append(test) else: filters.append(test) # Names of pandas source files. for root, _, files in os.walk(PANDAS_DOCS_SOURCE): for name in files: if name.endswith('.rst'): path = os.path.join(root, name) if any(filter in path for filter in filters): paths.append(path) # Using a global here is a bit hacky, but avoids pickling issues when used # with multiprocessing. parallelism = max(args.parallel or multiprocessing.cpu_count(), len(paths)) if parallelism > 1: pool_map = multiprocessing.pool.Pool(parallelism).imap_unordered run_tests = run_tests_capturing_stdout # Make sure slow tests get started first. paths.sort( key=lambda path: ('enhancingperf' in path, os.path.getsize(path)), reverse=True) else: pool_map = map run_tests = run_tests_streaming_stdout # Now run all the tests. running_summary = doctests.Summary() for count, (summary, stdout) in enumerate(pool_map(run_tests, paths)): running_summary += summary if stdout: print(stdout) print(count, '/', len(paths), 'done.') print('*' * 72) print("Final summary:") running_summary.summarize() def run_tests_capturing_stdout(path): with deferred_stdout() as stdout: return run_tests(path), stdout() def run_tests_streaming_stdout(path): return run_tests(path), None def run_tests(path): # Optionally capture the stdout as interleaved test errors are painful # to debug. On the other hand, if there is no parallelism, let the # output be streamed. start = time.time() with open(path) as f: rst = f.read() res = doctests.test_rst_ipython( rst, path, report=True, wont_implement_ok=['*'], not_implemented_ok=['*'], use_beam=False).summary print("Total time for {}: {:.2f} secs".format(path, time.time() - start)) return res @contextlib.contextmanager def deferred_stdout(): captured = io.StringIO() old_stdout, sys.stdout = sys.stdout, captured try: yield captured.getvalue finally: sys.stdout = old_stdout if __name__ == '__main__': main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/pandas_doctests_test.py0000644000175100017510000014155115206113032024377 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import sys import unittest import pandas as pd from apache_beam.dataframe import doctests from apache_beam.dataframe.frames import PD_VERSION from apache_beam.dataframe.pandas_top_level_functions import _is_top_level_function # pylint: skip-file @unittest.skipIf( sys.platform == 'win32', '[https://github.com/apache/beam/issues/20361]') class DoctestTest(unittest.TestCase): def test_ndframe_tests(self): # IO methods are tested in io_test.py skip_writes = { f'pandas.core.generic.NDFrame.{name}': ['*'] for name in dir(pd.core.generic.NDFrame) if name.startswith('to_') } result = doctests.testmod( pd.core.generic, use_beam=False, report=True, wont_implement_ok={ 'pandas.core.generic.NDFrame.attrs': ['*'], 'pandas.core.generic.NDFrame.bfill': ['*'], 'pandas.core.generic.NDFrame.ffill': ['*'], 'pandas.core.generic.NDFrame.first_valid_index': ['*'], 'pandas.core.generic.NDFrame.head': ['*'], 'pandas.core.generic.NDFrame.last_valid_index': ['*'], 'pandas.core.generic.NDFrame.shift': [ 'df.shift(periods=3)', 'df.shift(periods=3, fill_value=0)', "df['Col1'].shift(periods=[0, 1, 2])", ], 'pandas.core.generic.NDFrame.tail': ['*'], 'pandas.core.generic.NDFrame.take': ['*'], 'pandas.core.generic.NDFrame.values': ['*'], 'pandas.core.generic.NDFrame.tz_localize': [ "s.tz_localize('CET', ambiguous='infer')", # np.array is not a deferred object. This use-case is possible # with a deferred Series though, which is tested in # frames_test.py "s.tz_localize('CET', ambiguous=np.array([True, True, False]))", ], 'pandas.core.generic.NDFrame.truncate': [ # These inputs rely on tail (wont implement, order # sensitive) for verification "df.tail()", "df.truncate(before=pd.Timestamp('2016-01-05'),\n" " after=pd.Timestamp('2016-01-10')).tail()", "df.truncate('2016-01-05', '2016-01-10').tail()", "df.loc['2016-01-05':'2016-01-10', :].tail()" ], 'pandas.core.generic.NDFrame.replace': [ "s.replace([1, 2], method='bfill')", # Relies on method='pad' "s.replace('a')", # Relies on method='pad' # value=None is not valid for pandas < 1.4 "s.replace('a', None)", # Implicitly uses method='pad', but output doesn't rely on that # behavior. Verified indepently in # frames_test.py::DeferredFrameTest::test_replace "df.replace(regex={r'^ba.$': 'new', 'foo': 'xyz'})" ], 'pandas.core.generic.NDFrame.fillna': [ 'df.fillna(method=\'ffill\')', 'df.fillna(method="ffill")', 'df.fillna(value=values, limit=1)', ], 'pandas.core.generic.NDFrame.sort_values': ['*'], 'pandas.core.generic.NDFrame.mask': [ 'df.where(m, -df) == np.where(m, df, -df)' ], 'pandas.core.generic.NDFrame.where': [ 'df.where(m, -df) == np.where(m, df, -df)' ], 'pandas.core.generic.NDFrame.interpolate': ['*'], 'pandas.core.generic.NDFrame.resample': ['*'], 'pandas.core.generic.NDFrame.rolling': ['*'], # argsort wont implement 'pandas.core.generic.NDFrame.abs': [ 'df.loc[(df.c - 43).abs().argsort()]', ], 'pandas.core.generic.NDFrame.reindex': ['*'], 'pandas.core.generic.NDFrame.pct_change': ['*'], 'pandas.core.generic.NDFrame.asof': ['*'], 'pandas.core.generic.NDFrame.infer_objects': ['*'], 'pandas.core.generic.NDFrame.ewm': ['*'], 'pandas.core.generic.NDFrame.expanding': ['*'], 'pandas.core.generic.NDFrame.get': ['*'], }, not_implemented_ok={ 'pandas.core.generic.NDFrame.__iter__': ['*'], 'pandas.core.generic.NDFrame.asof': ['*'], 'pandas.core.generic.NDFrame.at_time': ['*'], 'pandas.core.generic.NDFrame.between_time': ['*'], 'pandas.core.generic.NDFrame.ewm': ['*'], 'pandas.core.generic.NDFrame.expanding': ['*'], 'pandas.core.generic.NDFrame.flags': ['*'], 'pandas.core.generic.NDFrame.rank': ['*'], 'pandas.core.generic.NDFrame.reindex_like': ['*'], 'pandas.core.generic.NDFrame.replace': ['*'], 'pandas.core.generic.NDFrame.sample': ['*'], 'pandas.core.generic.NDFrame.set_flags': ['*'], 'pandas.core.generic.NDFrame.squeeze': ['*'], 'pandas.core.generic.NDFrame.truncate': ['*'], }, skip={ # Internal test 'pandas.core.generic.NDFrame._set_axis_name': ['*'], # Fails to construct test series. asfreq is not implemented anyway. 'pandas.core.generic.NDFrame.asfreq': ['*'], 'pandas.core.generic.NDFrame.astype': ['*'], 'pandas.core.generic.NDFrame.convert_dtypes': ['*'], 'pandas.core.generic.NDFrame.copy': ['*'], 'pandas.core.generic.NDFrame.droplevel': ['*'], 'pandas.core.generic.NDFrame.get': ['*'], 'pandas.core.generic.NDFrame.rank': ['*'], 'pandas.core.generic.NDFrame.rename': [ # Seems to be an upstream bug. The actual error has a different # message: # TypeError: Index(...) must be called with a collection of # some kind, 2 was passed # pandas doctests only verify the type of exception 'df.rename(2)' ], # For pandas >= 1.4, rename is changed to _rename 'pandas.core.generic.NDFrame._rename': [ # Seems to be an upstream bug. The actual error has a different # message: # TypeError: Index(...) must be called with a collection of # some kind, 2 was passed # pandas doctests only verify the type of exception 'df.rename(2)' ], # Tests rely on setting index 'pandas.core.generic.NDFrame.rename_axis': ['*'], # Raises right exception, but testing framework has matching issues. 'pandas.core.generic.NDFrame.replace': [ "df.replace({'a string': 'new value', True: False}) # raises" ], 'pandas.core.generic.NDFrame.squeeze': ['*'], # NameError 'pandas.core.generic.NDFrame.resample': ['df'], # Skipped so we don't need to install natsort 'pandas.core.generic.NDFrame.sort_values': [ 'from natsort import index_natsorted', 'df.sort_values(\n' ' by="time",\n' ' key=lambda x: np.argsort(index_natsorted(df["time"]))\n' ')' ], # TODO(https://github.com/apache/beam/issues/28559): Re-enable when # bug is fixed. 'pandas.core.generic.NDFrame.xs': ['*'], **skip_writes }) self.assertEqual(result.failed, 0) def test_dataframe_tests(self): result = doctests.testmod( pd.core.frame, use_beam=False, report=True, wont_implement_ok={ 'pandas.core.frame.DataFrame.T': ['*'], 'pandas.core.frame.DataFrame.cummax': ['*'], 'pandas.core.frame.DataFrame.cummin': ['*'], 'pandas.core.frame.DataFrame.cumsum': ['*'], 'pandas.core.frame.DataFrame.cumprod': ['*'], 'pandas.core.frame.DataFrame.diff': ['*'], 'pandas.core.frame.DataFrame.fillna': [ 'df.fillna(method=\'ffill\')', 'df.fillna(method="ffill")', 'df.fillna(value=values, limit=1)', ], 'pandas.core.frame.DataFrame.items': ['*'], 'pandas.core.frame.DataFrame.itertuples': ['*'], 'pandas.core.frame.DataFrame.iterrows': ['*'], 'pandas.core.frame.DataFrame.iteritems': ['*'], # default keep is 'first' 'pandas.core.frame.DataFrame.nlargest': [ "df.nlargest(3, 'population')", "df.nlargest(3, ['population', 'GDP'])", "df.nlargest(3, 'population', keep='last')" ], 'pandas.core.frame.DataFrame.nsmallest': [ "df.nsmallest(3, 'population')", "df.nsmallest(3, ['population', 'GDP'])", "df.nsmallest(3, 'population', keep='last')", ], 'pandas.core.frame.DataFrame.replace': [ "s.replace([1, 2], method='bfill')", # Relies on method='pad' "s.replace('a')", # Relies on method='pad' # value=None is not valid for pandas < 1.4 "s.replace('a', None)", # Implicitly uses method='pad', but output doesn't rely on that # behavior. Verified indepently in # frames_test.py::DeferredFrameTest::test_replace "df.replace(regex={r'^ba.$': 'new', 'foo': 'xyz'})" ], 'pandas.core.frame.DataFrame.to_records': ['*'], 'pandas.core.frame.DataFrame.to_dict': ['*'], 'pandas.core.frame.DataFrame.to_numpy': ['*'], 'pandas.core.frame.DataFrame.to_string': ['*'], 'pandas.core.frame.DataFrame.transpose': ['*'], 'pandas.core.frame.DataFrame.shape': ['*'], 'pandas.core.frame.DataFrame.shift': [ 'df.shift(periods=3)', 'df.shift(periods=3, fill_value=0)', "df['Col1'].shift(periods=[0, 1, 2])", ], 'pandas.core.frame.DataFrame.unstack': ['*'], 'pandas.core.frame.DataFrame.memory_usage': ['*'], 'pandas.core.frame.DataFrame.info': ['*'], # Not equal to df.agg('mode', axis='columns', numeric_only=True) # because there can be multiple columns if a row has more than one # mode 'pandas.core.frame.DataFrame.mode': [ "df.mode(axis='columns', numeric_only=True)" ], 'pandas.core.frame.DataFrame.append': [ 'df.append(df2, ignore_index=True)', "for i in range(5):\n" + " df = df.append({'A': i}, ignore_index=True)", ], 'pandas.core.frame.DataFrame.sort_index': ['*'], 'pandas.core.frame.DataFrame.sort_values': ['*'], 'pandas.core.frame.DataFrame.melt': [ "df.melt(id_vars=['A'], value_vars=['B'])", "df.melt(id_vars=['A'], value_vars=['B', 'C'])", "df.melt(col_level=0, id_vars=['A'], value_vars=['B'])", "df.melt(id_vars=[('A', 'D')], value_vars=[('B', 'E')])", "df.melt(id_vars=['A'], value_vars=['B'],\n" + " var_name='myVarname', value_name='myValname')" ], # Most keep= options are order-sensitive 'pandas.core.frame.DataFrame.drop_duplicates': ['*'], 'pandas.core.frame.DataFrame.duplicated': [ 'df.duplicated()', "df.duplicated(keep='last')", "df.duplicated(subset=['brand'])", ], 'pandas.core.frame.DataFrame.reindex': ['*'], 'pandas.core.frame.DataFrame.dot': [ # reindex not supported 's2 = s.reindex([1, 0, 2, 3])', ], 'pandas.core.frame.DataFrame.resample': ['*'], 'pandas.core.frame.DataFrame.values': ['*'], }, not_implemented_ok={ 'pandas.core.frame.DataFrame.transform': [ # str arg not supported. Tested with np.sum in # frames_test.py::DeferredFrameTest::test_groupby_transform_sum "df.groupby('Date')['Data'].transform('sum')", ], 'pandas.core.frame.DataFrame.melt': ['*'], 'pandas.core.frame.DataFrame.reindex_axis': ['*'], 'pandas.core.frame.DataFrame.round': [ 'df.round(decimals)', ], # Trivially elementwise for axis=columns. Relies on global indexing # for axis=rows. # Difficult to determine proxy, need to inspect function 'pandas.core.frame.DataFrame.apply': ['*'], # Cross-join not implemented 'pandas.core.frame.DataFrame.merge': [ "df1.merge(df2, how='cross')" ], # TODO(https://github.com/apache/beam/issues/20759) 'pandas.core.frame.DataFrame.set_index': [ "df.set_index([s, s**2])", ], 'pandas.core.frame.DataFrame.set_axis': [ "df.set_axis(range(0,2), axis='index')", ], # TODO(https://github.com/apache/beam/issues/21014) 'pandas.core.frame.DataFrame.value_counts': [ 'df.value_counts(dropna=False)' ], 'pandas.core.frame.DataFrame.to_timestamp': ['*'] }, skip={ # These examples occur in docstrings for several ops. '*': [ # mul doesn't work in Beam with axis='index'. "df.mul({'circle': 0, 'triangle': 2, 'rectangle': 3}, " "axis='index')", # eq doesn't work with axis='index'. "df.eq([250, 250, 100], axis='index')", # New test in Pandas 2.1 that uses indexes. 'df != pd.Series([100, 250], index=["cost", "revenue"])', # New test in Pandas 2.1 that uses indexes. 'df.le(df_multindex, level=1)' ], # DeferredDataFrame doesn't implement the DF interchange protocol. 'pandas.core.frame.DataFrame.__dataframe__': ['*'], # DataFrame construction from a dictionary, Series, or other # DataFrame requires using the len() function, which is a # non-deferred operation that we do not allow 'pandas.core.frame.DataFrame': [ 'pd.DataFrame(data=d, index=[0, 1, 2, 3])', 'df = pd.DataFrame(data=ser, index=["a", "c"])', 'df', 'df2 = pd.DataFrame(data=df1, index=["a", "c"])', 'df2', ], # s2 created with reindex 'pandas.core.frame.DataFrame.dot': [ 'df.dot(s2)', ], 'pandas.core.frame.DataFrame.resample': ['df'], 'pandas.core.frame.DataFrame.asfreq': ['*'], # Throws NotImplementedError when modifying df 'pandas.core.frame.DataFrame.axes': [ # Returns deferred index. 'df.axes', ], # Skipped because the relies on loc to set cells in df2 'pandas.core.frame.DataFrame.compare': ['*'], 'pandas.core.frame.DataFrame.cov': [ # Relies on setting entries ahead of time. "df.loc[df.index[:5], 'a'] = np.nan", "df.loc[df.index[5:10], 'b'] = np.nan", 'df.cov(min_periods=12)', ], 'pandas.core.frame.DataFrame.rename': [ # Returns deferred index. 'df.index', 'df.rename(index=str).index', ], 'pandas.core.frame.DataFrame.set_index': [ # TODO(https://github.com/apache/beam/issues/20759): This could # pass in the index as a DeferredIndex, and we should fail it # as order-sensitive. "df.set_index([pd.Index([1, 2, 3, 4]), 'year'])", ], 'pandas.core.frame.DataFrame.set_axis': [ # This should pass as set_axis(axis='columns') # and fail with set_axis(axis='index') "df.set_axis(['a', 'b', 'c'], axis='index')" ], # Beam's implementation takes a filepath as an argument. 'pandas.core.frame.DataFrame.to_html': ['*'], 'pandas.core.frame.DataFrame.to_markdown': ['*'], 'pandas.core.frame.DataFrame.to_parquet': ['*'], # Raises right exception, but testing framework has matching issues. # Tested in `frames_test.py`. 'pandas.core.frame.DataFrame.insert': [ 'df', 'df.insert(1, "newcol", [99, 99])', 'df.insert(0, "col1", [100, 100], allow_duplicates=True)' ], 'pandas.core.frame.DataFrame.to_records': [ 'df.index = df.index.rename("I")', 'index_dtypes = f" 100).mean()', ], 'pandas.core.series.Series.asfreq': ['*'], # error formatting 'pandas.core.series.Series.append': [ 's1.append(s2, verify_integrity=True)', ], 'pandas.core.series.Series.cov': [ # Differs in LSB on jenkins. "s1.cov(s2)", ], # Test framework doesn't materialze DeferredIndex. 'pandas.core.series.Series.keys': ['s.keys()'], # Skipped idxmax/idxmin due an issue with the test framework 'pandas.core.series.Series.idxmin': ['s.idxmin()'], 'pandas.core.series.Series.idxmax': ['s.idxmax()'], 'pandas.core.series.Series.duplicated': ['*'], # Relies on setting index. 'pandas.core.series.Series.rename_axis': ['*'], 'pandas.core.series.Series.set_axis': ['*'], 'pandas.core.series.Series.nonzero': ['*'], 'pandas.core.series.Series.pop': ['ser'], # testing side effect # Raises right exception, but testing framework has matching issues. 'pandas.core.series.Series.replace': [ "df.replace({'a string': 'new value', True: False}) # raises" ], 'pandas.core.series.Series.searchsorted': [ # This doctest seems to be incorrectly parsed. "x = pd.Categorical(['apple', 'bread', 'bread'," ], 'pandas.core.series.Series.to_csv': ['*'], 'pandas.core.series.Series.to_markdown': ['*'], 'pandas.core.series.Series.update': ['*'], 'pandas.core.series.Series.view': [ # Inspection after modification. 's' ], 'pandas.core.series.Series.resample': ['df'], # Fails when result is a singleton: # https://github.com/apache/beam/issues/28559 'pandas.core.series.Series.kurt': [ 'df.kurt(axis=None).round(6)', 's.kurt()' ], # Fails when result is a singleton: # https://github.com/apache/beam/issues/28559 'pandas.core.series.Series.sem': [ 'df.sem().round(6)', 's.sem().round(6)' ], }) self.assertEqual(result.failed, 0) def test_string_tests(self): if PD_VERSION < (1, 2): module = pd.core.strings else: # Definitions were moved to accessor in pandas 1.2.0 module = pd.core.strings.accessor module_name = module.__name__ result = doctests.testmod( module, use_beam=False, wont_implement_ok={ # These methods can accept deferred series objects, but not lists f'{module_name}.StringMethods.cat': [ "s.str.cat(['A', 'B', 'C', 'D'], sep=',')", "s.str.cat(['A', 'B', 'C', 'D'], sep=',', na_rep='-')", "s.str.cat(['A', 'B', 'C', 'D'], na_rep='-')" ], f'{module_name}.StringMethods.repeat': [ 's.str.repeat(repeats=[1, 2, 3])' ], f'{module_name}.str_repeat': ['s.str.repeat(repeats=[1, 2, 3])'], # get_dummies pandas examples are not casted to CategoricalDtype # Must be CategoricalDtype to work in Beam f'{module_name}.StringMethods.get_dummies': ['*'], f'{module_name}.str_get_dummies': ['*'], f'{module_name}.StringMethods': ['s.str.split("_")'], }, skip={ # count() on Series with a NaN produces mismatched type if we # have a NaN-only partition. f'{module_name}.StringMethods.count': ["s.str.count('a')"], f'{module_name}.str_count': ["s.str.count('a')"], # Bad test strings in pandas 1.1.x f'{module_name}.str_replace': [ "pd.Series(['foo', 'fuz', np.nan]).str.replace('f', repr)" ], f'{module_name}.StringMethods.replace': [ "pd.Series(['foo', 'fuz', np.nan]).str.replace('f', repr)" ], # output has incorrect formatting in 1.2.x f'{module_name}.StringMethods.extractall': ['*'], # For split and rsplit, if expand=True, then the series # must be of CategoricalDtype, which pandas doesn't convert to f'{module_name}.StringMethods.rsplit': [ 's.str.split(r"\\+|=", expand=True)', # for pandas<1.4 's.str.split(expand=True)', 's.str.rsplit("/", n=1, expand=True)', 's.str.split(r"and|plus", expand=True)', 's.str.split(r".", expand=True)', 's.str.split(r"\\.jpg", expand=True)', 's.str.split(r"\\.jpg", regex=True, expand=True)', 's.str.split(re.compile(r"\\.jpg"), expand=True)', 's.str.split(r"\\.jpg", regex=False, expand=True)' ], f'{module_name}.StringMethods.split': [ 's.str.split(r"\\+|=", expand=True)', # for pandas<1.4 's.str.split(expand=True)', 's.str.rsplit("/", n=1, expand=True)', 's.str.split(r"and|plus", expand=True)', 's.str.split(r".", expand=True)', 's.str.split(r"\\.jpg", expand=True)', 's.str.split(r"\\.jpg", regex=True, expand=True)', 's.str.split(re.compile(r"\\.jpg"), expand=True)', 's.str.split(r"\\.jpg", regex=False, expand=True)' ] }) self.assertEqual(result.failed, 0) def test_datetime_tests(self): # TODO(BEAM-10721) indexes_accessors_result = doctests.testmod( pd.core.indexes.accessors, use_beam=False, skip={ '*': ["ser = pd.Series(pd.to_timedelta([1, 2, 3], unit='d'))"], 'pandas.core.indexes.accessors.TimedeltaProperties': [ # Seems like an upstream bug. The property is 'second' 'seconds_series.dt.seconds' ], # TODO(https://github.com/apache/beam/issues/21013): Test data # creation fails for these # s = pd.Series(pd.to_timedelta(np.arange(5), unit="d")) # pylint: disable=line-too-long 'pandas.core.indexes.accessors.DatetimeProperties.to_pydatetime': [ '*' ], # Beam doesn't have a deferred version of PeriodIndex. # PeriodIndex instance is created in the test scenario. 'pandas.core.indexes.accessors.PeriodProperties.end_time': ['*'], 'pandas.core.indexes.accessors.TimedeltaProperties.components': [ '*' ], 'pandas.core.indexes.accessors.TimedeltaProperties.days': ['*'], 'pandas.core.indexes.accessors.TimedeltaProperties.seconds': ['*'], 'pandas.core.indexes.accessors.TimedeltaProperties.microseconds': [ '*' ], 'pandas.core.indexes.accessors.TimedeltaProperties.nanoseconds': [ '*' ], 'pandas.core.indexes.accessors.TimedeltaProperties.to_pytimedelta': [ '*' ], # pylint: enable=line-too-long # Test uses to_datetime. Beam calls to_datetime element-wise, and # therefore the .tz attribute is not evaluated on entire Series. # Hence, .tz becomes None, unless explicitly set. # See: see test_tz_with_utc_zone_set_explicitly 'pandas.core.indexes.accessors.DatetimeProperties.tz': ['*'], }) datetimelike_result = doctests.testmod( pd.core.arrays.datetimelike, use_beam=False, not_implemented_ok={ # Beam Dataframes don't implement a deferred to_timedelta operation. # Top-level issue: https://github.com/apache/beam/issues/20318 '*': [ "ser = pd.Series(pd.to_timedelta([1, 2, 3], unit='d'))", "tdelta_idx = pd.to_timedelta([1, 2, 3], unit='D')", 'tdelta_idx = pd.to_timedelta(["0 days", "10 days", "20 days"])', # pylint: disable=line-too-long "tdelta_idx", "tdelta_idx.inferred_freq", "tdelta_idx.mean()", ], }) datetime_result = doctests.testmod( pd.core.arrays.datetimes, use_beam=False, wont_implement_ok={ 'pandas.core.arrays.datetimes.DatetimeArray.to_period': ['*'], # All tz_localize tests use unsupported values for ambiguous= # Verified seperately in # frames_test.py::DeferredFrameTest::test_dt_tz_localize_* 'pandas.core.arrays.datetimes.DatetimeArray.tz_localize': ['*'], }, not_implemented_ok={ # Beam Dataframes don't implement a deferred to_timedelta operation. # Top-level issue: https://github.com/apache/beam/issues/20318 '*': [ "ser = pd.Series(pd.to_timedelta([1, 2, 3], unit='d'))", "tdelta_idx = pd.to_timedelta([1, 2, 3], unit='D')", 'tdelta_idx = pd.to_timedelta(["0 days", "10 days", "20 days"])' ], # pylint: disable=line-too-long # Verifies index version of this method 'pandas.core.arrays.datetimes.DatetimeArray.to_period': [ 'df.index.to_period("M")' ], }, skip={ # Test uses to_datetime. Beam calls to_datetime element-wise, and # therefore the .tz attribute is not evaluated on entire Series. # Hence, .tz becomes None, unless explicitly set. # See: see test_tz_with_utc_zone_set_explicitly 'pandas.core.arrays.datetimes.DatetimeArray.tz': ['*'], }) self.assertEqual(indexes_accessors_result.failed, 0) self.assertEqual(datetimelike_result.failed, 0) self.assertEqual(datetime_result.failed, 0) def test_indexing_tests(self): result = doctests.testmod( pd.core.indexing, use_beam=False, skip={ 'pandas.core.indexing._IndexSlice': ['*'], 'pandas.core.indexing.IndexingMixin.at': ['*'], 'pandas.core.indexing.IndexingMixin.iat': ['*'], 'pandas.core.indexing.IndexingMixin.iloc': ['*'], 'pandas.core.indexing.IndexingMixin.loc': ['*'], 'pandas.core.indexing._AtIndexer': ['*'], 'pandas.core.indexing._LocIndexer': ['*'], 'pandas.core.indexing._iAtIndexer': ['*'], 'pandas.core.indexing._iLocIndexer': ['*'], }) self.assertEqual(result.failed, 0) def test_groupby_tests(self): result = doctests.testmod( pd.core.groupby.groupby, use_beam=False, verbose=True, wont_implement_ok={ '*': [ # resample is WontImpl. "ser.resample('MS').count()", "ser.resample('MS').median()", "ser.resample('MS').sem()", "ser.resample('MS').size()", ], 'pandas.core.groupby.groupby.BaseGroupBy.indices': ['*'], 'pandas.core.groupby.groupby.GroupBy.bfill': ['*'], 'pandas.core.groupby.groupby.GroupBy.ffill': ['*'], 'pandas.core.groupby.groupby.GroupBy.diff': ['*'], 'pandas.core.groupby.groupby.GroupBy.first': ['*'], 'pandas.core.groupby.groupby.GroupBy.head': ['*'], 'pandas.core.groupby.groupby.GroupBy.last': ['*'], 'pandas.core.groupby.groupby.GroupBy.ohlc': ['*'], 'pandas.core.groupby.groupby.GroupBy.pct_change': ['*'], 'pandas.core.groupby.groupby.GroupBy.tail': ['*'], 'pandas.core.groupby.groupby.GroupBy.nth': ['*'], 'pandas.core.groupby.groupby.GroupBy.cumcount': ['*'], 'pandas.core.groupby.groupby.GroupBy.cummax': ['*'], 'pandas.core.groupby.groupby.GroupBy.cummin': ['*'], 'pandas.core.groupby.groupby.GroupBy.cumprod': ['*'], 'pandas.core.groupby.groupby.GroupBy.cumsum': ['*'], 'pandas.core.groupby.groupby.GroupBy.plot': ['*'], 'pandas.core.groupby.groupby.GroupBy.resample': ['*'], 'pandas.core.groupby.groupby.GroupBy.rolling': ['*'], 'pandas.core.groupby.groupby.GroupBy.shift': ['*'], }, not_implemented_ok={ 'pandas.core.groupby.groupby.GroupBy.first': ['*'], 'pandas.core.groupby.groupby.GroupBy.last': ['*'], 'pandas.core.groupby.groupby.GroupBy.ngroup': ['*'], 'pandas.core.groupby.groupby.GroupBy.sample': ['*'], 'pandas.core.groupby.groupby.GroupBy.rank': ['*'], 'pandas.core.groupby.groupby.GroupBy.nth': [ "df.groupby('A', as_index=False).nth(1)", ], }, skip={ # New test that didn't pass on Pandas 1.5.x. 'pandas.core.groupby.groupby.BaseGroupBy.__iter__': ['*'], # Not implemented; some tests also use resample (won't implement) 'pandas.core.groupby.groupby.BaseGroupBy.get_group': ['*'], 'pandas.core.groupby.groupby.BaseGroupBy.groups': ['*'], # uses resample, which is WontImplement atm. # Uses iloc to mutate a DataFrame 'pandas.core.groupby.groupby.GroupBy.resample': [ 'df.iloc[2, 0] = 5', 'df', ], # df is reassigned 'pandas.core.groupby.groupby.GroupBy.rank': ['df'], # TODO: Raise wont implement for list passed as a grouping column # Currently raises unhashable type: list 'pandas.core.groupby.groupby.GroupBy.ngroup': [ 'df.groupby(["A", [1,1,2,3,2,1]]).ngroup()' ], }) self.assertEqual(result.failed, 0) result = doctests.testmod( pd.core.groupby.generic, use_beam=False, wont_implement_ok={ '*': [ # resample is WontImpl. "ser.resample('MS').nunique()", ], # TODO: Is take actually deprecated? 'pandas.core.groupby.generic.DataFrameGroupBy.take': ['*'], 'pandas.core.groupby.generic.SeriesGroupBy.take': ['*'], 'pandas.core.groupby.generic.SeriesGroupBy.nsmallest': [ "s.nsmallest(3, keep='last')", "s.nsmallest(3)", "s.nsmallest()", ], 'pandas.core.groupby.generic.SeriesGroupBy.nlargest': [ "s.nlargest(3, keep='last')", "s.nlargest(3)", "s.nlargest()", ], 'pandas.core.groupby.generic.DataFrameGroupBy.diff': ['*'], 'pandas.core.groupby.generic.SeriesGroupBy.diff': ['*'], 'pandas.core.groupby.generic.DataFrameGroupBy.hist': ['*'], 'pandas.core.groupby.generic.SeriesGroupBy.hist': ['*'], 'pandas.core.groupby.generic.DataFrameGroupBy.plot': ['*'], 'pandas.core.groupby.generic.SeriesGroupBy.plot': ['*'], 'pandas.core.groupby.generic.DataFrameGroupBy.fillna': [ 'df.fillna(method=\'ffill\')', 'df.fillna(method="ffill")', 'df.fillna(value=values, limit=1)', 'df.groupby("key").fillna(method="ffill")', 'df.groupby("key").fillna(method="bfill")', 'df.groupby("key").fillna(method="ffill", limit=1)', ], 'pandas.core.groupby.generic.SeriesGroupBy.fillna': [ 'df.fillna(method=\'ffill\')', 'df.fillna(method="ffill")', 'df.fillna(value=values, limit=1)', ], 'pandas.core.groupby.groupby.GroupBy.tail': ['*'], }, not_implemented_ok={ 'pandas.core.groupby.generic.DataFrameGroupBy.idxmax': ['*'], 'pandas.core.groupby.generic.DataFrameGroupBy.idxmin': ['*'], 'pandas.core.groupby.generic.DataFrameGroupBy.transform': ['*'], 'pandas.core.groupby.generic.SeriesGroupBy.transform': ['*'], 'pandas.core.groupby.generic.SeriesGroupBy.idxmax': ['*'], 'pandas.core.groupby.generic.SeriesGroupBy.idxmin': ['*'], 'pandas.core.groupby.generic.SeriesGroupBy.apply': ['*'], }, skip={ 'pandas.core.groupby.generic.SeriesGroupBy.cov': [ # Floating point comparison fails 's1.cov(s2)', ], 'pandas.core.groupby.generic.DataFrameGroupBy.cov': [ # Mutates input DataFrame with loc # TODO: Replicate in frames_test.py "df.loc[df.index[:5], 'a'] = np.nan", "df.loc[df.index[5:10], 'b'] = np.nan", "df.cov(min_periods=12)", ], # These examples rely on grouping by a list 'pandas.core.groupby.generic.SeriesGroupBy.aggregate': ['*'], 'pandas.core.groupby.generic.DataFrameGroupBy.aggregate': ['*'], # Skipped idxmax/idxmin due an issue with the test framework 'pandas.core.groupby.generic.SeriesGroupBy.idxmin': ['s.idxmin()'], 'pandas.core.groupby.generic.SeriesGroupBy.idxmax': ['s.idxmax()'], # Order-sensitive operations. TODO: Return a better error message. 'pandas.core.groupby.generic.SeriesGroupBy.is_monotonic_increasing': [ '*' ], # pylint: disable=line-too-long 'pandas.core.groupby.generic.SeriesGroupBy.is_monotonic_decreasing': [ '*' ], # pylint: disable=line-too-long # Uses as_index, which is currently not_implemented 'pandas.core.groupby.generic.DataFrameGroupBy.value_counts': [ "df.groupby('gender', as_index=False).value_counts()", # pylint: disable=line-too-long "df.groupby('gender', as_index=False).value_counts(normalize=True)", ], # These examples rely on grouping by a list 'pandas.core.groupby.generic.SeriesGroupBy.fillna': ['*'], # These examples rely on grouping by a list 'pandas.core.groupby.generic.DataFrameGroupBy.fillna': ['*'], # These examples rely on grouping by a list 'pandas.core.groupby.generic.SeriesGroupBy.take': ['*'], # These examples rely on grouping by a list 'pandas.core.groupby.generic.DataFrameGroupBy.take': ['*'], # Named aggregation not supported yet. 'pandas.core.groupby.generic.NamedAgg': [ 'df.groupby("key").agg(result_a=agg_a, result_1=agg_1)' ], # These examples rely on grouping by a list 'pandas.core.groupby.generic.DataFrameGroupBy.transform': ['*'], # These examples rely on grouping by a list 'pandas.core.groupby.generic.SeriesGroupBy.transform': ['*'], # Returns an array by default, not a Series. WontImplement # (non-deferred) 'pandas.core.groupby.generic.SeriesGroupBy.unique': ['*'], }, ) self.assertEqual(result.failed, 0) def test_top_level(self): tests = { name: func.__doc__ for (name, func) in pd.__dict__.items() if _is_top_level_function(func) and getattr(func, '__doc__', None) } # IO methods are tested in io_test.py skip_reads = {name: ['*'] for name in dir(pd) if name.startswith('read_')} result = doctests.teststrings( tests, use_beam=False, report=True, not_implemented_ok={ 'concat': ['pd.concat([s1, s2], ignore_index=True)'], 'crosstab': ['*'], 'cut': ['*'], 'eval': ['*'], 'from_dummies': ['*'], 'get_dummies': ['*'], 'infer_freq': ['*'], 'lreshape': ['*'], 'melt': ['*'], 'merge': ["df1.merge(df2, how='cross')"], 'merge_asof': ['*'], 'pivot_table': ['*'], 'qcut': ['*'], 'reset_option': ['*'], 'set_option': ['*'], 'to_numeric': ['*'], 'to_timedelta': ['*'], 'unique': ['*'], 'wide_to_long': ['*'], }, wont_implement_ok={ 'factorize': ['*'], 'pivot': ['*'], 'to_datetime': ['s.head()'], 'to_pickle': ['*'], 'unique': [ 'pd.unique(pd.Series([("a", "b"), ("b", "a"), ("a", "c"), ("b", "a")]).values)' ], # pylint: disable=line-too-long 'melt': [ "pd.melt(df, id_vars=['A'], value_vars=['B'])", "pd.melt(df, id_vars=['A'], value_vars=['B', 'C'])", "pd.melt(df, col_level=0, id_vars=['A'], value_vars=['B'])", "pd.melt(df, id_vars=[('A', 'D')], value_vars=[('B', 'E')])", "pd.melt(df, id_vars=['A'], value_vars=['B'],\n" + " var_name='myVarname', value_name='myValname')" ], }, skip={ # error formatting 'concat': [ 'pd.concat([df5, df6], verify_integrity=True)', 'pd.concat([df7, new_row.to_frame().T], ignore_index=True)' ], # doctest DeprecationWarning 'melt': ['df'], # Order-sensitive re-indexing. 'merge': [ "df1.merge(df2, left_on='lkey', right_on='rkey')", "df1.merge(df2, left_on='lkey', right_on='rkey',\n" " suffixes=('_left', '_right'))", "df1.merge(df2, how='left', on='a')", ], # Not an actual test. 'option_context': ['*'], 'factorize': ['codes', 'uniques'], # Bad top-level use of un-imported function. 'merge_ordered': [ 'merge_ordered(df1, df2, fill_method="ffill", left_by="group")' ], # Expected error. 'pivot': [ "df.pivot(index='foo', columns='bar', values='baz')", "df.pivot(index='foo', columns='bar')['baz']", "df.pivot(index='foo', columns='bar', values=['baz', 'zoo'])", # pylint: disable=line-too-long 'df.pivot(index="lev1", columns=["lev2", "lev3"],values="values")', # pylint: disable=line-too-long 'df.pivot(index=["lev1", "lev2"], columns=["lev3"],values="values")' ], # Never written. 'to_pickle': ['os.remove("./dummy.pkl")'], **skip_reads }) self.assertEqual(result.failed, 0) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/pandas_top_level_functions.py0000644000175100017510000001374315206113032025572 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """A module providing various functionality from the top-level pandas namespace. """ import re from collections.abc import Mapping import pandas as pd from apache_beam.dataframe import expressions from apache_beam.dataframe import frame_base from apache_beam.dataframe import partitionings def _call_on_first_arg(name): def wrapper(target, *args, **kwargs): if isinstance(target, frame_base.DeferredBase): return getattr(target, name)(*args, **kwargs) else: return getattr(pd, name)(target, *args, **kwargs) return staticmethod(wrapper) def _maybe_wrap_constant_expr(res): if type(res) in frame_base.DeferredBase._pandas_type_map: return frame_base.DeferredBase.wrap( expressions.ConstantExpression(res, res[0:0])) else: return res def _defer_to_pandas(name): func = getattr(pd, name) def wrapper(*args, **kwargs): res = func(*args, **kwargs) return _maybe_wrap_constant_expr(res) return staticmethod(wrapper) def _defer_to_pandas_maybe_elementwise(name): """ Same as _defer_to_pandas, except it handles DeferredBase args, assuming the function can be processed elementwise. """ func = getattr(pd, name) def wrapper(*args, **kwargs): if any(isinstance(arg, frame_base.DeferredBase) for arg in args + tuple(kwargs.values())): return frame_base._elementwise_function(func, name)(*args, **kwargs) res = func(*args, **kwargs) return _maybe_wrap_constant_expr(res) return staticmethod(wrapper) def _is_top_level_function(o): return ( callable(o) and not isinstance(o, type) and hasattr(o, '__name__') and re.match('[a-z].*', o.__name__)) class DeferredPandasModule(object): array = _defer_to_pandas('array') bdate_range = _defer_to_pandas('bdate_range') @staticmethod @frame_base.args_to_kwargs(pd) @frame_base.populate_defaults(pd) def concat( objs, axis, join, ignore_index, keys, levels, names, verify_integrity, sort, copy): if ignore_index: raise NotImplementedError('concat(ignore_index)') if levels: raise NotImplementedError('concat(levels)') if isinstance(objs, Mapping): if keys is None: keys = list(objs.keys()) objs = [objs[k] for k in keys] else: objs = list(objs) if keys is None: preserves_partitioning = partitionings.Arbitrary() else: # Index 0 will be a new index for keys, only partitioning by the original # indexes (1 to N) will be preserved. nlevels = min(o._expr.proxy().index.nlevels for o in objs) preserves_partitioning = partitionings.Index( [i for i in range(1, nlevels + 1)]) deferred_none = expressions.ConstantExpression(None) exprs = [deferred_none if o is None else o._expr for o in objs] if axis in (1, 'columns'): required_partitioning = partitionings.Index() elif verify_integrity: required_partitioning = partitionings.Index() else: required_partitioning = partitionings.Arbitrary() return frame_base.DeferredBase.wrap( expressions.ComputedExpression( 'concat', lambda *objs: pd.concat( objs, axis=axis, join=join, ignore_index=ignore_index, keys= keys, levels=levels, names=names, verify_integrity= verify_integrity), # yapf break exprs, requires_partition_by=required_partitioning, preserves_partition_by=preserves_partitioning)) date_range = _defer_to_pandas('date_range') describe_option = _defer_to_pandas('describe_option') factorize = _call_on_first_arg('factorize') get_option = _defer_to_pandas('get_option') interval_range = _defer_to_pandas('interval_range') isna = _call_on_first_arg('isna') isnull = _call_on_first_arg('isnull') json_normalize = _defer_to_pandas('json_normalize') melt = _call_on_first_arg('melt') merge = _call_on_first_arg('merge') melt = _call_on_first_arg('melt') merge_ordered = frame_base.wont_implement_method( pd, 'merge_ordered', reason='order-sensitive') notna = _call_on_first_arg('notna') notnull = _call_on_first_arg('notnull') option_context = _defer_to_pandas('option_context') period_range = _defer_to_pandas('period_range') pivot = _call_on_first_arg('pivot') pivot_table = _call_on_first_arg('pivot_table') set_eng_float_format = _defer_to_pandas('set_eng_float_format') show_versions = _defer_to_pandas('show_versions') test = frame_base.wont_implement_method( pd, 'test', explanation="because it is an internal pandas testing utility.") timedelta_range = _defer_to_pandas('timedelta_range') to_pickle = frame_base.wont_implement_method( pd, 'to_pickle', reason='order-sensitive') to_datetime = _defer_to_pandas_maybe_elementwise('to_datetime') notna = _call_on_first_arg('notna') def __getattr__(self, name): if name.startswith('read_'): def func(*args, **kwargs): raise frame_base.WontImplementError( 'Use p | apache_beam.dataframe.io.%s' % name) return func res = getattr(pd, name) if _is_top_level_function(res): return frame_base.not_implemented_method(name, base_type=pd) else: return res pd_wrapper = DeferredPandasModule() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/partitionings.py0000644000175100017510000001571115206113032023052 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import random from collections.abc import Iterable from typing import Any from typing import TypeVar import numpy as np import pandas as pd Frame = TypeVar('Frame', bound=pd.core.generic.NDFrame) class Partitioning(object): """A class representing a (consistent) partitioning of dataframe objects. """ def __repr__(self): return self.__class__.__name__ def is_subpartitioning_of(self, other: 'Partitioning') -> bool: """Returns whether self is a sub-partition of other. Specifically, returns whether something partitioned by self is necissarily also partitioned by other. """ raise NotImplementedError def __lt__(self, other): return self != other and self <= other def __le__(self, other): return not self.is_subpartitioning_of(other) def partition_fn(self, df: Frame, num_partitions: int) -> Iterable[tuple[Any, Frame]]: """A callable that actually performs the partitioning of a Frame df. This will be invoked via a FlatMap in conjunction with a GroupKey to achieve the desired partitioning. """ raise NotImplementedError def test_partition_fn(self, df): return self.partition_fn(df, 5) class Index(Partitioning): """A partitioning by index (either fully or partially). If the set of "levels" of the index to consider is not specified, the entire index is used. These form a partial order, given by Singleton() < Index([i]) < Index([i, j]) < ... < Index() < Arbitrary() The ordering is implemented via the is_subpartitioning_of method, where the examples on the right are subpartitionings of the examples on the left above. """ def __init__(self, levels=None): self._levels = levels def __repr__(self): if self._levels: return 'Index%s' % self._levels else: return 'Index' def __eq__(self, other): return type(self) == type(other) and self._levels == other._levels def __hash__(self): if self._levels: return hash(tuple(sorted(self._levels))) else: return hash(type(self)) def is_subpartitioning_of(self, other): if isinstance(other, Singleton): return True elif isinstance(other, Index): if self._levels is None: return True elif other._levels is None: return False else: return all(level in self._levels for level in other._levels) elif isinstance(other, (Arbitrary, JoinIndex)): return False else: raise ValueError(f"Encountered unknown type {other!r}") def _hash_index(self, df): if self._levels is None: levels = list(range(df.index.nlevels)) else: levels = self._levels return sum( pd.util.hash_array(np.asarray(df.index.get_level_values(level))) for level in levels) def partition_fn(self, df, num_partitions): hashes = self._hash_index(df) for key in range(num_partitions): yield key, df[hashes % num_partitions == key] def check(self, dfs): # Drop empty DataFrames dfs = [df for df in dfs if len(df)] if not len(dfs): return True def apply_consistent_order(dfs): # Apply consistent order between dataframes by using sum of the index's # hash. # Apply consistent order within dataframe with sort_index() # Also drops any empty dataframes. return sorted((df.sort_index() for df in dfs if len(df)), key=lambda df: sum(self._hash_index(df))) dfs = apply_consistent_order(dfs) repartitioned_dfs = apply_consistent_order( df for _, df in self.test_partition_fn(pd.concat(dfs))) # Assert that each index is identical for df, repartitioned_df in zip(dfs, repartitioned_dfs): if not df.index.equals(repartitioned_df.index): return False return True class Singleton(Partitioning): """A partitioning of all the data into a single partition. """ def __init__(self, reason=None): self._reason = reason @property def reason(self): return self._reason def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) def is_subpartitioning_of(self, other): return isinstance(other, Singleton) def partition_fn(self, df, num_partitions): yield None, df def check(self, dfs): return len(dfs) <= 1 class JoinIndex(Partitioning): """A partitioning that lets two frames be joined. This can either be a hash partitioning on the full index, or a common ancestor with no intervening re-indexing/re-partitioning. It fits into the partial ordering as Index() < JoinIndex(x) < JoinIndex() < Arbitrary() with JoinIndex(x) and JoinIndex(y) being incomparable for nontrivial x != y. Expressions desiring to make use of this index should simply declare a requirement of JoinIndex(). """ def __init__(self, ancestor=None): self._ancestor = ancestor def __repr__(self): if self._ancestor: return 'JoinIndex[%s]' % self._ancestor else: return 'JoinIndex' def __eq__(self, other): if type(self) != type(other): return False elif self._ancestor is None: return other._ancestor is None elif other._ancestor is None: return False else: return self._ancestor == other._ancestor def __hash__(self): return hash((type(self), self._ancestor)) def is_subpartitioning_of(self, other): if isinstance(other, Arbitrary): return False elif isinstance(other, JoinIndex): return self._ancestor is None or self == other else: return True def test_partition_fn(self, df): return Index().test_partition_fn(df) def check(self, dfs): return True class Arbitrary(Partitioning): """A partitioning imposing no constraints on the actual partitioning. """ def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) def is_subpartitioning_of(self, other): return True def test_partition_fn(self, df): num_partitions = 10 def shuffled(seq): seq = list(seq) random.shuffle(seq) return seq part = pd.Series(shuffled(range(len(df))), index=df.index) % num_partitions for k in range(num_partitions): yield k, df[part == k] def check(self, dfs): return True ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/partitionings_test.py0000644000175100017510000000700715206113032024110 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import unittest import pandas as pd from apache_beam.dataframe.partitionings import Arbitrary from apache_beam.dataframe.partitionings import Index from apache_beam.dataframe.partitionings import JoinIndex from apache_beam.dataframe.partitionings import Singleton class PartitioningsTest(unittest.TestCase): # pylint: disable=bad-option-value multi_index_df = pd.DataFrame({ 'shape': ['dodecahedron', 'icosahedron'] * 12, 'color': ['red', 'yellow', 'blue'] * 8, 'size': range(24), 'value': range(24) }).set_index(['shape', 'color', 'size']) def test_index_is_subpartition(self): ordered_list = [ Singleton(), Index([3]), Index([1, 3]), Index(), JoinIndex('ref'), JoinIndex(), Arbitrary() ] for loose, strict in zip(ordered_list[:-1], ordered_list[1:]): self.assertTrue(strict.is_subpartitioning_of(loose), (strict, loose)) self.assertFalse(loose.is_subpartitioning_of(strict), (loose, strict)) # Incomparable. self.assertFalse(Index([1, 2]).is_subpartitioning_of(Index([1, 3]))) self.assertFalse(Index([1, 3]).is_subpartitioning_of(Index([1, 2]))) self.assertFalse(JoinIndex('a').is_subpartitioning_of(JoinIndex('b'))) self.assertFalse(JoinIndex('b').is_subpartitioning_of(JoinIndex('a'))) def _check_partition(self, partitioning, min_non_empty, max_non_empty=None): num_partitions = 1000 if max_non_empty is None: max_non_empty = min_non_empty parts = list(partitioning.partition_fn(self.multi_index_df, num_partitions)) self.assertEqual(num_partitions, len(parts)) self.assertGreaterEqual(len([p for _, p in parts if len(p)]), min_non_empty) self.assertLessEqual(len([p for _, p in parts if len(p)]), max_non_empty) self.assertEqual( sorted(self.multi_index_df.value), sorted(sum((list(p.value) for _, p in parts), []))) def test_index_partition(self): self._check_partition(Index([0]), 2) self._check_partition(Index([0, 1]), 6) self._check_partition(Index([1]), 3) self._check_partition(Index([2]), 7, 24) self._check_partition(Index([0, 2]), 7, 24) self._check_partition(Index(), 7, 24) def test_nothing_subpartition(self): for p in [Index([1]), Index([1, 2]), Index(), Singleton()]: self.assertTrue(Arbitrary().is_subpartitioning_of(p), p) def test_singleton_subpartition(self): self.assertTrue(Singleton().is_subpartitioning_of(Singleton())) for p in [Arbitrary(), Index([1]), Index([1, 2]), Index()]: self.assertFalse(Singleton().is_subpartitioning_of(p), p) def test_singleton_partition(self): parts = list(Singleton().partition_fn(pd.Series(range(10)), 1000)) self.assertEqual(1, len(parts)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/schemas.py0000644000175100017510000001726015206113032021604 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utilities for relating schema-aware PCollections and DataFrame transforms. The utilities here enforce the type mapping defined in :mod:`apache_beam.typehints.pandas_type_compatibility`. """ # pytype: skip-file import warnings from collections.abc import Sequence from typing import Any from typing import NamedTuple from typing import Optional from typing import TypeVar from typing import Union import pandas as pd import apache_beam as beam from apache_beam import typehints from apache_beam.transforms.util import BatchElements from apache_beam.typehints.pandas_type_compatibility import INDEX_OPTION_NAME from apache_beam.typehints.pandas_type_compatibility import create_pandas_batch_converter from apache_beam.typehints.pandas_type_compatibility import dtype_from_typehint from apache_beam.typehints.pandas_type_compatibility import dtype_to_fieldtype from apache_beam.typehints.row_type import RowTypeConstraint from apache_beam.typehints.schemas import named_fields_from_element_type from apache_beam.typehints.typehints import normalize __all__ = ( 'BatchRowsAsDataFrame', 'generate_proxy', 'UnbatchPandas', 'element_type_from_dataframe') T = TypeVar('T', bound=NamedTuple) @typehints.with_input_types(T) @typehints.with_output_types(pd.DataFrame) class BatchRowsAsDataFrame(beam.PTransform): """A transform that batches schema-aware PCollection elements into DataFrames Batching parameters are inherited from :class:`~apache_beam.transforms.util.BatchElements`. """ def __init__(self, *args, proxy=None, **kwargs): self._batch_elements_transform = BatchElements(*args, **kwargs) self._proxy = proxy def expand(self, pcoll): if self._proxy is not None: # Generate typehint proxy = self._proxy element_typehint = _element_typehint_from_proxy(proxy) else: # Generate proxy proxy = generate_proxy(pcoll.element_type) element_typehint = pcoll.element_type converter = create_pandas_batch_converter( element_type=element_typehint, batch_type=type(proxy)) return ( pcoll | self._batch_elements_transform | beam.Map(converter.produce_batch)) def generate_proxy(element_type: type) -> pd.DataFrame: """Generate a proxy pandas object for the given PCollection element_type. Currently only supports generating a DataFrame proxy from a schema-aware PCollection or a Series proxy from a primitively typed PCollection. """ dtype = dtype_from_typehint(element_type) if dtype is not object: return pd.Series(dtype=dtype) else: fields = named_fields_from_element_type(element_type) proxy = pd.DataFrame(columns=[name for name, _ in fields]) for name, typehint in fields: dtype = dtype_from_typehint(typehint) proxy[name] = proxy[name].astype(dtype) return proxy def element_type_from_dataframe( proxy: pd.DataFrame, include_indexes: bool = False) -> type: """Generate an element_type for an element-wise PCollection from a proxy pandas object. Currently only supports converting the element_type for a schema-aware PCollection to a proxy DataFrame. Currently only supports generating a DataFrame proxy from a schema-aware PCollection. """ return element_typehint_from_dataframe_proxy(proxy, include_indexes).user_type def _element_typehint_from_proxy( proxy: pd.core.generic.NDFrame, include_indexes: bool = False): if isinstance(proxy, pd.DataFrame): return element_typehint_from_dataframe_proxy( proxy, include_indexes=include_indexes) elif isinstance(proxy, pd.Series): if include_indexes: warnings.warn( "include_indexes=True for a Series input. Note that this " "parameter is _not_ respected for DeferredSeries " "conversion.") return dtype_to_fieldtype(proxy.dtype) else: raise TypeError(f"Proxy '{proxy}' has unsupported type '{type(proxy)}'") def element_typehint_from_dataframe_proxy( proxy: pd.DataFrame, include_indexes: bool = False) -> RowTypeConstraint: output_columns = [] if include_indexes: remaining_index_names = list(proxy.index.names) i = 0 while len(remaining_index_names): index_name = remaining_index_names.pop(0) if index_name is None: raise ValueError( "Encountered an unnamed index. Cannot convert to a " "schema-aware PCollection with include_indexes=True. " "Please name all indexes or consider not including " "indexes.") elif index_name in remaining_index_names: raise ValueError( "Encountered multiple indexes with the name '%s'. " "Cannot convert to a schema-aware PCollection with " "include_indexes=True. Please ensure all indexes have " "unique names or consider not including indexes." % index_name) elif index_name in proxy.columns: raise ValueError( "Encountered an index that has the same name as one " "of the columns, '%s'. Cannot convert to a " "schema-aware PCollection with include_indexes=True. " "Please ensure all indexes have unique names or " "consider not including indexes." % index_name) else: # its ok! output_columns.append( (index_name, proxy.index.get_level_values(i).dtype)) i += 1 output_columns.extend(zip(proxy.columns, proxy.dtypes)) fields = [(column, dtype_to_fieldtype(dtype)) for (column, dtype) in output_columns] field_options: Optional[dict[str, Sequence[tuple[str, Any]]]] if include_indexes: field_options = { index_name: [(INDEX_OPTION_NAME, None)] for index_name in proxy.index.names } else: field_options = None return RowTypeConstraint.from_fields(fields, field_options=field_options) def _unbatch_transform(proxy, include_indexes): element_typehint = normalize( _element_typehint_from_proxy(proxy, include_indexes=include_indexes)) converter = create_pandas_batch_converter( element_type=element_typehint, batch_type=type(proxy)) return beam.FlatMap( converter.explode_batch).with_output_types(element_typehint) @typehints.with_input_types(Union[pd.DataFrame, pd.Series]) class UnbatchPandas(beam.PTransform): """A transform that explodes a PCollection of DataFrame or Series. DataFrame is converterd to a schema-aware PCollection, while Series is converted to its underlying type. Args: include_indexes: (optional, default: False) When unbatching a DataFrame if include_indexes=True, attempt to include index columns in the output schema for expanded DataFrames. Raises an error if any of the index levels are unnamed (name=None), or if any of the names are not unique among all column and index names. """ def __init__(self, proxy, include_indexes=False): self._proxy = proxy self._include_indexes = include_indexes def expand(self, pcoll): return pcoll | _unbatch_transform(self._proxy, self._include_indexes) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/schemas_test.py0000644000175100017510000003055615206113032022646 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for schemas.""" # pytype: skip-file import typing import unittest import numpy as np import pandas as pd from parameterized import parameterized import apache_beam as beam from apache_beam.coders import RowCoder from apache_beam.coders.typecoders import registry as coders_registry from apache_beam.dataframe import schemas from apache_beam.dataframe import transforms from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.typehints import row_type from apache_beam.typehints import typehints from apache_beam.typehints.native_type_compatibility import match_is_named_tuple Simple = typing.NamedTuple( 'Simple', [('name', str), ('id', int), ('height', float)]) coders_registry.register_coder(Simple, RowCoder) Animal = typing.NamedTuple( 'Animal', [('animal', str), ('max_speed', typing.Optional[float])]) coders_registry.register_coder(Animal, RowCoder) def matches_df(expected): def check_df_pcoll_equal(actual): actual = pd.concat(actual) sorted_actual = actual.sort_values(by=list(actual.columns)).reset_index( drop=True) sorted_expected = expected.sort_values( by=list(expected.columns)).reset_index(drop=True) pd.testing.assert_frame_equal(sorted_actual, sorted_expected) return check_df_pcoll_equal # Test data for all supported types that can be easily tested. # Excludes bytes because it's difficult to create a series and dataframe bytes # dtype. For example: # pd.Series([b'abc'], dtype=bytes).dtype != 'S' # pd.Series([b'abc'], dtype=bytes).astype(bytes).dtype == 'S' # (test data, pandas_type, column_name, beam_type) COLUMNS: list[tuple[list[typing.Any], typing.Any, str, typing.Any]] = [ ([375, 24, 0, 10, 16], np.int32, 'i32', np.int32), ([375, 24, 0, 10, 16], np.int64, 'i64', np.int64), ([375, 24, None, 10, 16], pd.Int32Dtype(), 'i32_nullable', typing.Optional[np.int32]), ([375, 24, None, 10, 16], pd.Int64Dtype(), 'i64_nullable', typing.Optional[np.int64]), ([375., 24., None, 10., 16.], np.float64, 'f64', typing.Optional[np.float64]), ([375., 24., None, 10., 16.], np.float32, 'f32', typing.Optional[np.float32]), ([True, False, True, True, False], bool, 'bool', bool), (['Falcon', 'Ostrich', None, 3.14, 0], object, 'any', typing.Any), ([True, False, True, None, False], pd.BooleanDtype(), 'bool_nullable', typing.Optional[bool]), (['Falcon', 'Ostrich', None, 'Aardvark', 'Elephant'], pd.StringDtype(), 'strdtype', typing.Optional[str]), ] NICE_TYPES_DF = pd.DataFrame(columns=[name for _, _, name, _ in COLUMNS]) for arr, dtype, name, _ in COLUMNS: NICE_TYPES_DF[name] = pd.Series(arr, dtype=dtype, name=name).astype(dtype) NICE_TYPES_PROXY = NICE_TYPES_DF[:0] SERIES_TESTS = [(pd.Series(arr, dtype=dtype, name=name), arr, beam_type) for (arr, dtype, name, beam_type) in COLUMNS] _TEST_ARRAYS: list[list[typing.Any]] = [arr for (arr, _, _, _) in COLUMNS] DF_RESULT = list(zip(*_TEST_ARRAYS)) BEAM_SCHEMA = typing.NamedTuple( # type: ignore 'BEAM_SCHEMA', [(name, beam_type) for _, _, name, beam_type in COLUMNS]) INDEX_DF_TESTS = [( NICE_TYPES_DF.set_index([name for _, _, name, _ in COLUMNS[:i]]), DF_RESULT, BEAM_SCHEMA) for i in range(1, len(COLUMNS) + 1)] NOINDEX_DF_TESTS = [(NICE_TYPES_DF, DF_RESULT, BEAM_SCHEMA)] # Get major, minor, bugfix version PD_VERSION = tuple(map(int, pd.__version__.split('.')[0:3])) def test_name_func(testcase_func, param_num, params): df_or_series, _, _ = params.args if isinstance(df_or_series, pd.Series): return f"{testcase_func.__name__}_Series[{df_or_series.dtype}]" elif isinstance(df_or_series, pd.DataFrame): return ( f"{testcase_func.__name__}_DataFrame" f"[{','.join(str(dtype) for dtype in df_or_series.dtypes)}]") else: raise ValueError( f"Encountered unsupported param in {testcase_func.__name__}. " "Expected Series or DataFrame, got:\n" + str(df_or_series)) class SchemasTest(unittest.TestCase): def test_simple_df(self): expected = pd.DataFrame({ 'name': list(str(i) for i in range(5)), 'id': list(range(5)), 'height': list(float(i) for i in range(5)) }, columns=['name', 'id', 'height']) expected.name = expected.name.astype(pd.StringDtype()) with TestPipeline() as p: res = ( p | beam.Create( [Simple(name=str(i), id=i, height=float(i)) for i in range(5)]) | schemas.BatchRowsAsDataFrame(min_batch_size=10, max_batch_size=10)) assert_that(res, matches_df(expected)) def test_simple_df_with_beam_row(self): expected = pd.DataFrame({ 'name': list(str(i) for i in range(5)), 'id': list(range(5)), 'height': list(float(i) for i in range(5)) }, columns=['name', 'id', 'height']) expected.name = expected.name.astype(pd.StringDtype()) with TestPipeline() as p: res = ( p | beam.Create([(str(i), i, float(i)) for i in range(5)]) | beam.Select( name=lambda r: str(r[0]), id=lambda r: int(r[1]), height=lambda r: float(r[2])) | schemas.BatchRowsAsDataFrame(min_batch_size=10, max_batch_size=10)) assert_that(res, matches_df(expected)) def test_generate_proxy(self): expected = pd.DataFrame({ 'animal': pd.Series(dtype=pd.StringDtype()), 'max_speed': pd.Series(dtype=np.float64) }) pd.testing.assert_frame_equal(schemas.generate_proxy(Animal), expected) def test_generate_proxy_beam_typehint(self): expected = pd.Series(dtype=pd.Int32Dtype()) actual = schemas.generate_proxy(typehints.Optional[np.int32]) pd.testing.assert_series_equal(actual, expected) def test_nice_types_proxy_roundtrip(self): roundtripped = schemas.generate_proxy( schemas.element_type_from_dataframe(NICE_TYPES_PROXY)) self.assertTrue(roundtripped.equals(NICE_TYPES_PROXY)) @unittest.skipIf( PD_VERSION == (1, 2, 1), "Can't roundtrip bytes in pandas 1.2.1" "https://github.com/pandas-dev/pandas/issues/39474") def test_bytes_proxy_roundtrip(self): proxy = pd.DataFrame({'bytes': []}) proxy.bytes = proxy.bytes.astype(bytes) roundtripped = schemas.generate_proxy( schemas.element_type_from_dataframe(proxy)) self.assertEqual(roundtripped.bytes.dtype.kind, 'S') def test_batch_with_df_transform(self): with TestPipeline() as p: res = ( p | beam.Create([ Animal('Falcon', 380.0), Animal('Falcon', 370.0), Animal('Parrot', 24.0), Animal('Parrot', 26.0) ]) | schemas.BatchRowsAsDataFrame() | transforms.DataframeTransform( lambda df: df.groupby('animal').mean(), # TODO: Generate proxy in this case as well proxy=schemas.generate_proxy(Animal), include_indexes=True)) assert_that(res, equal_to([('Falcon', 375.), ('Parrot', 25.)])) # Do the same thing, but use reset_index() to make sure 'animal' is included with TestPipeline() as p: with beam.dataframe.allow_non_parallel_operations(): res = ( p | beam.Create([ Animal('Falcon', 380.0), Animal('Falcon', 370.0), Animal('Parrot', 24.0), Animal('Parrot', 26.0) ]) | schemas.BatchRowsAsDataFrame() | transforms.DataframeTransform( lambda df: df.groupby('animal').mean().reset_index(), # TODO: Generate proxy in this case as well proxy=schemas.generate_proxy(Animal))) assert_that(res, equal_to([('Falcon', 375.), ('Parrot', 25.)])) def assert_typehints_equal(self, left, right): def maybe_drop_rowtypeconstraint(typehint): if isinstance(typehint, row_type.RowTypeConstraint): return typehint.user_type else: return typehint left = maybe_drop_rowtypeconstraint(typehints.normalize(left)) right = maybe_drop_rowtypeconstraint(typehints.normalize(right)) if match_is_named_tuple(left): self.assertTrue(match_is_named_tuple(right)) self.assertEqual(left.__annotations__, right.__annotations__) else: self.assertEqual(left, right) @parameterized.expand( SERIES_TESTS + NOINDEX_DF_TESTS, name_func=test_name_func) def test_unbatch_no_index(self, df_or_series, rows, beam_type): proxy = df_or_series[:0] with TestPipeline() as p: res = ( p | beam.Create([df_or_series[::2], df_or_series[1::2]]) | schemas.UnbatchPandas(proxy)) # Verify that the unbatched PCollection has the expected typehint # TODO(https://github.com/apache/beam/issues/19923): typehints should # support NamedTuple so we can use typehints.is_consistent_with here # instead self.assert_typehints_equal(res.element_type, beam_type) assert_that(res, equal_to(rows)) @parameterized.expand(SERIES_TESTS + INDEX_DF_TESTS, name_func=test_name_func) def test_unbatch_with_index(self, df_or_series, rows, _): proxy = df_or_series[:0] if (PD_VERSION < (1, 2) and set(['i32_nullable', 'i64_nullable']).intersection(proxy.index.names)): self.skipTest( "pandas<1.2 incorrectly changes Int64Dtype to int64 when " "moved to index.") with TestPipeline() as p: res = ( p | beam.Create([df_or_series[::2], df_or_series[1::2]]) | schemas.UnbatchPandas(proxy, include_indexes=True)) assert_that(res, equal_to(rows)) @parameterized.expand(SERIES_TESTS, name_func=test_name_func) def test_unbatch_series_with_index_warns( self, series, unused_rows, unused_type): proxy = series[:0] with TestPipeline() as p: input_pc = p | beam.Create([series[::2], series[1::2]]) with self.assertWarns(UserWarning): _ = input_pc | schemas.UnbatchPandas(proxy, include_indexes=True) def test_unbatch_include_index_unnamed_index_raises(self): df = pd.DataFrame({'foo': [1, 2, 3, 4]}) proxy = df[:0] with TestPipeline() as p: pc = p | beam.Create([df[::2], df[1::2]]) with self.assertRaisesRegex(ValueError, 'unnamed'): _ = pc | schemas.UnbatchPandas(proxy, include_indexes=True) def test_unbatch_include_index_nonunique_index_raises(self): df = pd.DataFrame({'foo': [1, 2, 3, 4]}) df.index = pd.MultiIndex.from_arrays([[1, 2, 3, 4], [4, 3, 2, 1]], names=['bar', 'bar']) proxy = df[:0] with TestPipeline() as p: pc = p | beam.Create([df[::2], df[1::2]]) with self.assertRaisesRegex(ValueError, 'bar'): _ = pc | schemas.UnbatchPandas(proxy, include_indexes=True) def test_unbatch_include_index_column_conflict_raises(self): df = pd.DataFrame({'foo': [1, 2, 3, 4]}) df.index = pd.Index([4, 3, 2, 1], name='foo') proxy = df[:0] with TestPipeline() as p: pc = p | beam.Create([df[::2], df[1::2]]) with self.assertRaisesRegex(ValueError, 'foo'): _ = pc | schemas.UnbatchPandas(proxy, include_indexes=True) def test_unbatch_datetime(self): s = pd.Series( pd.date_range( '1/1/2000', periods=100, freq='m', tz='America/Los_Angeles')) proxy = s[:0] with TestPipeline() as p: res = ( p | beam.Create([s[::2], s[1::2]]) | schemas.UnbatchPandas(proxy, include_indexes=True)) assert_that(res, equal_to(list(s))) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/transforms.py0000644000175100017510000005447015206113032022363 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import collections import logging from collections.abc import Mapping from typing import Any from typing import TypeVar from typing import Union import pandas as pd import apache_beam as beam from apache_beam import transforms from apache_beam.dataframe import expressions from apache_beam.dataframe import frame_base from apache_beam.dataframe import frames # pylint: disable=unused-import from apache_beam.dataframe import partitionings from apache_beam.pvalue import PCollection from apache_beam.utils import windowed_value __all__ = [ 'DataframeTransform', ] T = TypeVar('T') TARGET_PARTITION_SIZE = 1 << 23 # 8M MIN_PARTITION_SIZE = 1 << 19 # 0.5M MAX_PARTITIONS = 1000 DEFAULT_PARTITIONS = 100 MIN_PARTITIONS = 10 PER_COL_OVERHEAD = 1000 class DataframeTransform(transforms.PTransform): """A PTransform for applying function that takes and returns dataframes to one or more PCollections. :class:`DataframeTransform` will accept a PCollection with a `schema`_ and batch it into :class:`~pandas.DataFrame` instances if necessary:: (pcoll | beam.Select(key=..., foo=..., bar=...) | DataframeTransform(lambda df: df.group_by('key').sum())) It is also possible to process a PCollection of :class:`~pandas.DataFrame` instances directly, in this case a "proxy" must be provided. For example, if ``pcoll`` is a PCollection of DataFrames, one could write:: pcoll | DataframeTransform(lambda df: df.group_by('key').sum(), proxy=...) To pass multiple PCollections, pass a tuple of PCollections wich will be passed to the callable as positional arguments, or a dictionary of PCollections, in which case they will be passed as keyword arguments. Args: yield_elements: (optional, default: "schemas") If set to ``"pandas"``, return PCollection(s) containing the raw Pandas objects (:class:`~pandas.DataFrame` or :class:`~pandas.Series` as appropriate). If set to ``"schemas"``, return an element-wise PCollection, where DataFrame and Series instances are expanded to one element per row. DataFrames are converted to `schema-aware`_ PCollections, where column values can be accessed by attribute. include_indexes: (optional, default: False) When ``yield_elements="schemas"``, if ``include_indexes=True``, attempt to include index columns in the output schema for expanded DataFrames. Raises an error if any of the index levels are unnamed (name=None), or if any of the names are not unique among all column and index names. proxy: (optional) An empty :class:`~pandas.DataFrame` or :class:`~pandas.Series` instance with the same ``dtype`` and ``name`` as the elements of the input PCollection. Required when input PCollection :class:`~pandas.DataFrame` or :class:`~pandas.Series` elements. Ignored when input PCollection has a `schema`_. .. _schema: https://beam.apache.org/documentation/programming-guide/#what-is-a-schema .. _schema-aware: https://beam.apache.org/documentation/programming-guide/#what-is-a-schema """ def __init__( self, func, proxy=None, yield_elements="schemas", include_indexes=False): self._func = func self._proxy = proxy self._yield_elements = yield_elements self._include_indexes = include_indexes def expand(self, input_pcolls): # Avoid circular import. from apache_beam.dataframe import convert # Convert inputs to a flat dict. input_dict: dict[Any, PCollection] = _flatten(input_pcolls) proxies = _flatten(self._proxy) if self._proxy is not None else { tag: None for tag in input_dict } input_frames: dict[Any, frame_base.DeferredFrame] = { k: convert.to_dataframe(pc, proxies[k], str(k)) for k, pc in input_dict.items() } # noqa: F821 # Apply the function. frames_input = _substitute(input_pcolls, input_frames) if isinstance(frames_input, dict): result_frames = self._func(**frames_input) elif isinstance(frames_input, tuple): result_frames = self._func(*frames_input) else: result_frames = self._func(frames_input) # Compute results as a tuple. result_frames_dict = _flatten(result_frames) keys = list(result_frames_dict.keys()) result_frames_tuple = tuple(result_frames_dict[key] for key in keys) result_pcolls_tuple = convert.to_pcollection( *result_frames_tuple, label='Eval', always_return_tuple=True, yield_elements=self._yield_elements, include_indexes=self._include_indexes) # Convert back to the structure returned by self._func. result_pcolls_dict = dict(zip(keys, result_pcolls_tuple)) return _substitute(result_frames, result_pcolls_dict) class _DataframeExpressionsTransform(transforms.PTransform): def __init__(self, outputs): self._outputs = outputs def expand(self, inputs): return self._apply_deferred_ops(inputs, self._outputs) def _apply_deferred_ops( self, inputs: dict[expressions.Expression, PCollection], outputs: dict[Any, expressions.Expression], ) -> dict[Any, PCollection]: """Construct a Beam graph that evaluates a set of expressions on a set of input PCollections. :param inputs: A mapping of placeholder expressions to PCollections. :param outputs: A mapping of keys to expressions defined in terms of the placeholders of inputs. Returns a dictionary whose keys are those of outputs, and whose values are PCollections corresponding to the values of outputs evaluated at the values of inputs. Logically, `_apply_deferred_ops({x: a, y: b}, {f: F(x, y), g: G(x, y)})` returns `{f: F(a, b), g: G(a, b)}`. """ class ComputeStage(beam.PTransform): """A helper transform that computes a single stage of operations. """ def __init__(self, stage): self.stage = stage def default_label(self): return '%s:%s' % (self.stage.ops, id(self)) def expand(self, pcolls): logging.info('Computing dataframe stage %s for %s', self, self.stage) scalar_inputs = [expr for expr in self.stage.inputs if is_scalar(expr)] tabular_inputs = [ expr for expr in self.stage.inputs if not is_scalar(expr) ] if len(tabular_inputs) == 0: partitioned_pcoll = next(iter( pcolls.values())).pipeline | beam.Create([{}]) elif self.stage.partitioning != partitionings.Arbitrary(): # Partitioning required for these operations. # Compute the number of partitions to use for the inputs based on # the estimated size of the inputs. if self.stage.partitioning == partitionings.Singleton(): # Always a single partition, don't waste time computing sizes. num_partitions = 1 else: # Estimate the sizes from the outputs of a *previous* stage such # that using these estimates will not cause a fusion break. input_sizes = [ estimate_size(input, same_stage_ok=False) for input in tabular_inputs ] if None in input_sizes: # We were unable to (cheaply) compute the size of one or more # inputs. num_partitions = DEFAULT_PARTITIONS else: num_partitions = beam.pvalue.AsSingleton( input_sizes | 'FlattenSizes' >> beam.Flatten() | 'SumSizes' >> beam.CombineGlobally(sum) | 'NumPartitions' >> beam.Map( lambda size: max( MIN_PARTITIONS, min( MAX_PARTITIONS, size // TARGET_PARTITION_SIZE)))) partition_fn = self.stage.partitioning.partition_fn class Partition(beam.PTransform): def expand(self, pcoll): return ( pcoll # Attempt to create batches of reasonable size. | beam.ParDo(_PreBatch()) # Actually partition. | beam.FlatMap(partition_fn, num_partitions) # Don't bother shuffling empty partitions. | beam.Filter(lambda k_df: len(k_df[1]))) # Arrange such that partitioned_pcoll is properly partitioned. main_pcolls = { expr._id: pcolls[expr._id] | 'Partition_%s_%s' % (self.stage.partitioning, expr._id) >> Partition() for expr in tabular_inputs } | beam.CoGroupByKey() partitioned_pcoll = main_pcolls | beam.ParDo(_ReBatch()) else: # Already partitioned, or no partitioning needed. assert len(tabular_inputs) == 1 tag = tabular_inputs[0]._id partitioned_pcoll = pcolls[tag] | beam.Map(lambda df: {tag: df}) side_pcolls = { expr._id: beam.pvalue.AsSingleton(pcolls[expr._id]) for expr in scalar_inputs } # Actually evaluate the expressions. def evaluate(partition, stage=self.stage, **side_inputs): def lookup(expr): # Use proxy if there's no data in this partition return expr.proxy().iloc[:0] if partition[ expr._id] is None else partition[expr._id] session = expressions.Session( dict([(expr, lookup(expr)) for expr in tabular_inputs] + [(expr, side_inputs[expr._id]) for expr in scalar_inputs])) for expr in stage.outputs: yield beam.pvalue.TaggedOutput(expr._id, expr.evaluate_at(session)) return partitioned_pcoll | beam.FlatMap(evaluate, ** side_pcolls).with_outputs() class Stage(object): """Used to build up a set of operations that can be fused together. Note that these Dataframe "stages" contain a CoGBK and hence are often split across multiple "executable" stages. """ def __init__(self, inputs, partitioning): self.inputs = set(inputs) if (len(self.inputs) > 1 and partitioning.is_subpartitioning_of(partitionings.Index())): # We have to shuffle to co-locate, might as well partition. self.partitioning = partitionings.Index() elif isinstance(partitioning, partitionings.JoinIndex): # Not an actionable partitioning, use index. self.partitioning = partitionings.Index() else: self.partitioning = partitioning self.ops = [] self.outputs = set() def __repr__(self, indent=0): if indent: sep = '\n' + ' ' * indent else: sep = '' return ( "Stage[%sinputs=%s, %spartitioning=%s, %sops=%s, %soutputs=%s]" % ( sep, self.inputs, sep, self.partitioning, sep, self.ops, sep, self.outputs)) # First define some helper functions. @_memoize def output_partitioning_in_stage(expr, stage): """Return the output partitioning of expr when computed in stage, or returns None if the expression cannot be computed in this stage. """ def maybe_upgrade_to_join_index(partitioning): if partitioning.is_subpartitioning_of(partitionings.JoinIndex()): return partitionings.JoinIndex(expr) else: return partitioning if expr in stage.inputs or expr in inputs: # Inputs are all partitioned by stage.partitioning. return maybe_upgrade_to_join_index(stage.partitioning) # Anything that's not an input must have arguments assert len(expr.args()) arg_partitionings = set( output_partitioning_in_stage(arg, stage) for arg in expr.args() if not is_scalar(arg)) if len(arg_partitionings) == 0: # All inputs are scalars, output partitioning isn't dependent on the # input. return maybe_upgrade_to_join_index(expr.preserves_partition_by()) if len(arg_partitionings) > 1: # Arguments must be identically partitioned, can't compute this # expression here. return None arg_partitioning = arg_partitionings.pop() if not expr.requires_partition_by().is_subpartitioning_of( arg_partitioning): # Arguments aren't partitioned sufficiently for this expression return None return maybe_upgrade_to_join_index( expressions.output_partitioning(expr, arg_partitioning)) def is_computable_in_stage(expr, stage): return output_partitioning_in_stage(expr, stage) is not None def common_stages(stage_lists): # Set intersection, with a preference for earlier items in the list. if stage_lists: for stage in stage_lists[0]: if all(stage in other for other in stage_lists[1:]): yield stage @_memoize def is_scalar(expr): return not isinstance(expr.proxy(), pd.core.generic.NDFrame) @_memoize def expr_to_stages(expr): if expr in inputs: # Don't create a stage for each input, but it is still useful to record # what which stages inputs are available from. return [] # First attempt to compute this expression as part of an existing stage, # if possible. if all(arg in inputs for arg in expr.args()): # All input arguments; try to pick a stage that already has as many # of the inputs, correctly partitioned, as possible. inputs_by_stage = collections.defaultdict(int) for arg in expr.args(): for stage in expr_to_stages(arg): if is_computable_in_stage(expr, stage): inputs_by_stage[stage] += 1 + 100 * ( expr.requires_partition_by() == stage.partitioning) if inputs_by_stage: # Take the stage with the largest count. stage = max(inputs_by_stage.items(), key=lambda kv: kv[1])[0] else: stage = None else: # Try to pick a stage that has all the available non-input expressions. # TODO(robertwb): Baring any that have all of them, we could try and # pick one that has the most, but we need to ensure it is not a # predecessor of any of the missing argument's stages. for stage in common_stages([expr_to_stages(arg) for arg in expr.args() if arg not in inputs]): if is_computable_in_stage(expr, stage): break else: stage = None if stage is None: # No stage available, compute this expression as part of a new stage. stage = Stage([ arg for arg in expr.args() if not isinstance(arg, expressions.ConstantExpression) ], expr.requires_partition_by()) for arg in expr.args(): # For each argument, declare that it is also available in # this new stage. expr_to_stages(arg).append(stage) # It also must be declared as an output of the producing stage. expr_to_stage(arg).outputs.add(arg) stage.ops.append(expr) # Ensure that any inputs for the overall transform are added # in downstream stages. for arg in expr.args(): if arg in inputs: stage.inputs.add(arg) # This is a list as given expression may be available in many stages. return [stage] def expr_to_stage(expr): # Any will do; the first requires the fewest intermediate stages. return expr_to_stages(expr)[0] # Ensure each output is computed. for expr in outputs.values(): if expr not in inputs: expr_to_stage(expr).outputs.add(expr) @_memoize def stage_to_result(stage): return { expr._id: expr_to_pcoll(expr) for expr in stage.inputs } | ComputeStage(stage) @_memoize def expr_to_pcoll(expr): if expr in inputs: return inputs[expr] else: return stage_to_result(expr_to_stage(expr))[expr._id] @_memoize def estimate_size(expr, same_stage_ok): # Returns a pcollection of ints whose sum is the estimated size of the # given expression. pipeline = next(iter(inputs.values())).pipeline label = 'Size[%s, %s]' % (expr._id, same_stage_ok) if is_scalar(expr): return pipeline | label >> beam.Create([0]) elif same_stage_ok: return expr_to_pcoll(expr) | label >> beam.Map(_total_memory_usage) elif expr in inputs: return None else: # This is the stage to avoid. expr_stage = expr_to_stage(expr) # If the stage doesn't start with a shuffle, it's not safe to fuse # the computation into its parent either. has_shuffle = expr_stage.partitioning != partitionings.Arbitrary() # We assume the size of an expression is the sum of the size of its # inputs, which may be off by quite a bit, but the goal is to get # within an order of magnitude or two. arg_sizes = [] for arg in expr.args(): if is_scalar(arg): continue elif arg in inputs: return None arg_size = estimate_size( arg, same_stage_ok=has_shuffle and expr_to_stage(arg) != expr_stage) if arg_size is None: return None arg_sizes.append(arg_size) return arg_sizes | label >> beam.Flatten(pipeline=pipeline) # Now we can compute and return the result. return {k: expr_to_pcoll(expr) for k, expr in outputs.items()} def _total_memory_usage(frame): assert isinstance(frame, (pd.core.generic.NDFrame, pd.Index)) try: size = frame.memory_usage() if not isinstance(size, int): size = size.sum() + PER_COL_OVERHEAD * len(size) else: size += PER_COL_OVERHEAD return size except AttributeError: # Don't know, assume it's really big. float('inf') class _PreBatch(beam.DoFn): def __init__( self, target_size=TARGET_PARTITION_SIZE, min_size=MIN_PARTITION_SIZE): self._target_size = target_size self._min_size = min_size def start_bundle(self): self._parts = collections.defaultdict(list) self._running_size = 0 def process( self, part, window=beam.DoFn.WindowParam, timestamp=beam.DoFn.TimestampParam): part_size = _total_memory_usage(part) if part_size >= self._min_size: yield part else: self._running_size += part_size self._parts[window, timestamp].append(part) if self._running_size >= self._target_size: yield from self.finish_bundle() def finish_bundle(self): for (window, timestamp), parts in self._parts.items(): yield windowed_value.WindowedValue(_concat(parts), timestamp, (window, )) self.start_bundle() class _ReBatch(beam.DoFn): """Groups all the parts from various workers into the same dataframe. Also groups across partitions, up to a given data size, to recover some efficiency in the face of over-partitioning. """ def __init__( self, target_size=TARGET_PARTITION_SIZE, min_size=MIN_PARTITION_SIZE): self._target_size = target_size self._min_size = min_size def start_bundle(self): self._parts = collections.defaultdict(lambda: collections.defaultdict(list)) self._running_size = 0 def process( self, element, window=beam.DoFn.WindowParam, timestamp=beam.DoFn.TimestampParam): _, tagged_parts = element for tag, parts in tagged_parts.items(): for part in parts: self._running_size += _total_memory_usage(part) self._parts[window, timestamp][tag].extend(parts) if self._running_size >= self._target_size: yield from self.finish_bundle() def finish_bundle(self): for (window, timestamp), tagged_parts in self._parts.items(): yield windowed_value.WindowedValue( # yapf break { tag: _concat(parts) if parts else None for (tag, parts) in tagged_parts.items() }, timestamp, (window, )) self.start_bundle() def _memoize(f): cache = {} def wrapper(*args, **kwargs): key = args, tuple(sorted(kwargs.items())) if key not in cache: cache[key] = f(*args, **kwargs) return cache[key] return wrapper def _dict_union(dicts): result = {} for d in dicts: result.update(d) return result def _concat(parts): if len(parts) == 1: return parts[0] else: return pd.concat(parts) def _flatten( valueish: Union[T, list[T], tuple[T], dict[Any, T]], root: tuple[Any, ...] = (), ) -> Mapping[tuple[Any, ...], T]: """Given a nested structure of dicts, tuples, and lists, return a flat dictionary where the values are the leafs and the keys are the "paths" to these leaves. For example `{a: x, b: (y, z)}` becomes `{(a,): x, (b, 0): y, (b, 1): c}`. """ if isinstance(valueish, dict): return _dict_union(_flatten(v, root + (k, )) for k, v in valueish.items()) elif isinstance(valueish, (tuple, list)): return _dict_union( _flatten(v, root + (ix, )) for ix, v in enumerate(valueish)) else: return {root: valueish} def _substitute(valueish, replacements, root=()): """Substitutes the values in valueish with those in replacements where the keys are as in _flatten. For example, ``` _substitute( {a: x, b: (y, z)}, {(a,): X, (b, 0): Y, (b, 1): Z}) ``` returns `{a: X, b: (Y, Z)}`. """ if isinstance(valueish, dict): return type(valueish)({ k: _substitute(v, replacements, root + (k, )) for (k, v) in valueish.items() }) elif isinstance(valueish, (tuple, list)): return type(valueish)(( _substitute(v, replacements, root + (ix, )) for (ix, v) in enumerate(valueish))) else: return replacements[root] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/dataframe/transforms_test.py0000644000175100017510000004130615206113032023414 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import typing import unittest import warnings import pandas as pd import apache_beam as beam from apache_beam import coders from apache_beam import metrics from apache_beam.dataframe import convert from apache_beam.dataframe import expressions from apache_beam.dataframe import frame_base from apache_beam.dataframe import transforms from apache_beam.runners.portability.fn_api_runner import fn_runner from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to def check_correct(expected, actual): if actual is None: raise AssertionError('Empty frame but expected: \n\n%s' % (expected)) if isinstance(expected, pd.core.generic.NDFrame): expected = expected.sort_index() actual = actual.sort_index() if isinstance(expected, pd.Series): pd.testing.assert_series_equal(expected, actual) elif isinstance(expected, pd.DataFrame): pd.testing.assert_frame_equal(expected, actual) else: raise ValueError( f"Expected value is a {type(expected)}," "not a Series or DataFrame.") else: if actual != expected: raise AssertionError('Scalars not equal: %s != %s' % (actual, expected)) def concat(parts): if len(parts) > 1: return pd.concat(parts) elif len(parts) == 1: return parts[0] else: return None def df_equal_to(expected): return lambda actual: check_correct(expected, concat(actual)) AnimalSpeed = typing.NamedTuple( 'AnimalSpeed', [('Animal', str), ('Speed', int)]) coders.registry.register_coder(AnimalSpeed, coders.RowCoder) Nested = typing.NamedTuple( 'Nested', [('id', int), ('animal_speed', AnimalSpeed)]) coders.registry.register_coder(Nested, coders.RowCoder) class TransformTest(unittest.TestCase): def run_scenario(self, input, func): expected = func(input) empty = input.iloc[0:0] input_placeholder = expressions.PlaceholderExpression(empty) input_deferred = frame_base.DeferredFrame.wrap(input_placeholder) actual_deferred = func(input_deferred)._expr.evaluate_at( expressions.Session({input_placeholder: input})) check_correct(expected, actual_deferred) with beam.Pipeline() as p: input_pcoll = p | beam.Create([input.iloc[::2], input.iloc[1::2]]) input_df = convert.to_dataframe(input_pcoll, proxy=empty) output_df = func(input_df) output_proxy = output_df._expr.proxy() if isinstance(output_proxy, pd.core.generic.NDFrame): self.assertTrue( output_proxy.iloc[:0].equals(expected.iloc[:0]), ( 'Output proxy is incorrect:\n' f'Expected:\n{expected.iloc[:0]}\n\n' f'Actual:\n{output_proxy.iloc[:0]}')) else: self.assertEqual(type(output_proxy), type(expected)) output_pcoll = convert.to_pcollection(output_df, yield_elements='pandas') assert_that( output_pcoll, lambda actual: check_correct(expected, concat(actual))) def test_identity(self): df = pd.DataFrame({ 'Animal': ['Falcon', 'Falcon', 'Parrot', 'Parrot'], 'Speed': [380., 370., 24., 26.] }) self.run_scenario(df, lambda x: x) def test_groupby_sum_mean(self): df = pd.DataFrame({ 'Animal': ['Falcon', 'Falcon', 'Parrot', 'Parrot'], 'Speed': [380., 370., 24., 26.] }) self.run_scenario(df, lambda df: df.groupby('Animal').sum()) with expressions.allow_non_parallel_operations(): self.run_scenario(df, lambda df: df.groupby('Animal').mean()) self.run_scenario( df, lambda df: df.loc[df.Speed > 25].groupby('Animal').sum()) def test_groupby_apply(self): df = pd.DataFrame({ 'group': ['a' if i % 5 == 0 or i % 3 == 0 else 'b' for i in range(100)], 'foo': [None if i % 11 == 0 else i for i in range(100)], 'bar': [None if i % 7 == 0 else 99 - i for i in range(100)], 'baz': [None if i % 13 == 0 else i * 2 for i in range(100)], }) def median_sum_fn(x): with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="Mean of empty slice") return (x.foo + x.bar).median() describe = lambda df: df.describe() self.run_scenario(df, lambda df: df.groupby('group').foo.apply(describe)) self.run_scenario( df, lambda df: df.groupby('group')[['foo', 'bar']].apply(describe)) self.run_scenario(df, lambda df: df.groupby('group').apply(median_sum_fn)) self.run_scenario( df, lambda df: df.set_index('group').foo.groupby(level=0).apply(describe)) self.run_scenario(df, lambda df: df.groupby(level=0).apply(median_sum_fn)) self.run_scenario( df, lambda df: df.groupby(lambda x: x % 3).apply(describe)) def test_filter(self): df = pd.DataFrame({ 'Animal': ['Aardvark', 'Ant', 'Elephant', 'Zebra'], 'Speed': [5, 2, 35, 40] }) self.run_scenario(df, lambda df: df.filter(items=['Animal'])) self.run_scenario(df, lambda df: df.filter(regex='Anim.*')) self.run_scenario( df, lambda df: df.set_index('Animal').filter(regex='F.*', axis='index')) with expressions.allow_non_parallel_operations(): a = pd.DataFrame({'col': [1, 2, 3]}) self.run_scenario(a, lambda a: a.agg(sum)) self.run_scenario(a, lambda a: a.agg(['mean', 'min', 'max'])) def test_scalar(self): with expressions.allow_non_parallel_operations(): a = pd.Series([1, 2, 6]) self.run_scenario(a, lambda a: a.agg(sum)) self.run_scenario(a, lambda a: a / a.agg(sum)) self.run_scenario(a, lambda a: a / (a.max() - a.min())) self.run_scenario(a, lambda a: a / (a.sum() - 1)) # Tests scalar being used as an input to a downstream stage. df = pd.DataFrame({'key': ['a', 'a', 'b'], 'val': [1, 2, 6]}) self.run_scenario( df, lambda df: df.groupby('key').sum().val / df.val.agg(sum)) def test_getitem_projection(self): df = pd.DataFrame({ 'Animal': ['Aardvark', 'Ant', 'Elephant', 'Zebra'], 'Speed': [5, 2, 35, 40], 'Size': ['Small', 'Extra Small', 'Large', 'Medium'] }) self.run_scenario(df, lambda df: df[['Speed', 'Size']]) def test_offset_elementwise(self): s = pd.Series(range(10)).astype(float) df = pd.DataFrame({'value': s, 'square': s * s, 'cube': s * s * s}) # Only those values that are both squares and cubes will intersect. self.run_scenario( df, lambda df: df.set_index('square').value + df.set_index('cube').value) def test_batching_named_tuple_input(self): with beam.Pipeline() as p: result = ( p | beam.Create([ AnimalSpeed('Aardvark', 5), AnimalSpeed('Ant', 2), AnimalSpeed('Elephant', 35), AnimalSpeed('Zebra', 40) ]).with_output_types(AnimalSpeed) | transforms.DataframeTransform(lambda df: df.filter(regex='Anim.*'))) assert_that( result, equal_to([('Aardvark', ), ('Ant', ), ('Elephant', ), ('Zebra', )])) def test_batching_beam_row_input(self): with beam.Pipeline() as p: result = ( p | beam.Create([('Falcon', 380.), ('Falcon', 370.), ('Parrot', 24.), ('Parrot', 26.)]) | beam.Map(lambda tpl: beam.Row(Animal=tpl[0], Speed=tpl[1])) | transforms.DataframeTransform( lambda df: df.groupby('Animal').mean(), include_indexes=True)) assert_that(result, equal_to([('Falcon', 375.), ('Parrot', 25.)])) def test_batching_beam_row_to_dataframe(self): with beam.Pipeline() as p: df = convert.to_dataframe( p | beam.Create([('Falcon', 380.), ('Falcon', 370.), ('Parrot', 24.), ( 'Parrot', 26.)]) | beam.Map(lambda tpl: beam.Row(Animal=tpl[0], Speed=tpl[1]))) result = convert.to_pcollection( df.groupby('Animal').mean(), include_indexes=True) assert_that(result, equal_to([('Falcon', 375.), ('Parrot', 25.)])) def test_batching_passthrough_nested_schema(self): with beam.Pipeline() as p: nested_schema_pc = ( p | beam.Create([Nested(1, AnimalSpeed('Aardvark', 5)) ]).with_output_types(Nested)) result = nested_schema_pc | transforms.DataframeTransform( # pylint: disable=expression-not-assigned lambda df: df.filter(items=['animal_speed'])) assert_that(result, equal_to([(('Aardvark', 5), )])) def test_batching_passthrough_nested_array(self): Array = typing.NamedTuple( 'Array', [('id', int), ('business_numbers', typing.Sequence[int])]) coders.registry.register_coder(Array, coders.RowCoder) with beam.Pipeline() as p: array_schema_pc = (p | beam.Create([Array(1, [7, 8, 9])])) result = array_schema_pc | transforms.DataframeTransform( # pylint: disable=expression-not-assigned lambda df: df.filter(items=['business_numbers'])) assert_that(result, equal_to([([7, 8, 9], )])) def test_unbatching_series(self): with beam.Pipeline() as p: result = ( p | beam.Create([('Falcon', 380.), ('Falcon', 370.), ('Parrot', 24.), ('Parrot', 26.)]) | beam.Map(lambda tpl: beam.Row(Animal=tpl[0], Speed=tpl[1])) | transforms.DataframeTransform(lambda df: df.Animal)) assert_that(result, equal_to(['Falcon', 'Falcon', 'Parrot', 'Parrot'])) def test_input_output_polymorphism(self): one_series = pd.Series([1]) two_series = pd.Series([2]) three_series = pd.Series([3]) proxy = one_series[:0] def equal_to_series(expected): def check(actual): actual = pd.concat(actual) if not expected.equals(actual): raise AssertionError( 'Series not equal: \n%s\n%s\n' % (expected, actual)) return check with beam.Pipeline() as p: one = p | 'One' >> beam.Create([one_series]) two = p | 'Two' >> beam.Create([two_series]) assert_that( one | 'PcollInPcollOut' >> transforms.DataframeTransform( lambda x: 3 * x, proxy=proxy, yield_elements='pandas'), equal_to_series(three_series), label='CheckPcollInPcollOut') assert_that( (one, two) | 'TupleIn' >> transforms.DataframeTransform( lambda x, y: (x + y), (proxy, proxy), yield_elements='pandas'), equal_to_series(three_series), label='CheckTupleIn') assert_that( dict(x=one, y=two) | 'DictIn' >> transforms.DataframeTransform( lambda x, y: (x + y), proxy=dict(x=proxy, y=proxy), yield_elements='pandas'), equal_to_series(three_series), label='CheckDictIn') double, triple = one | 'TupleOut' >> transforms.DataframeTransform( lambda x: (2*x, 3*x), proxy, yield_elements='pandas') assert_that(double, equal_to_series(two_series), 'CheckTupleOut0') assert_that(triple, equal_to_series(three_series), 'CheckTupleOut1') res = one | 'DictOut' >> transforms.DataframeTransform( lambda x: {'res': 3 * x}, proxy, yield_elements='pandas') assert_that(res['res'], equal_to_series(three_series), 'CheckDictOut') def test_multiple_dataframes_transforms(self): expected_output = ["Bryan", "DKER2"] def transform_func(a, b): b["name"] = "DKER2" return a, b with beam.Pipeline() as p: pcol1 = p | "Create1" >> beam.Create([beam.Row(name="Bryan")]) pcol2 = p | "Create2" >> beam.Create([beam.Row(name="common")]) result = ({ "a": pcol1, "b": pcol2 } | "TransformDF" >> transforms.DataframeTransform(transform_func) | "Flatten" >> beam.Flatten() | transforms.DataframeTransform(lambda df: df.name)) assert_that(result, equal_to(expected_output)) def test_cat(self): # verify that cat works with a List[Series] since this is # missing from doctests df = pd.DataFrame({ 'one': ['A', 'B', 'C'], 'two': ['BB', 'CC', 'A'], 'three': ['CCC', 'AA', 'B'], }) self.run_scenario(df, lambda df: df.two.str.cat([df.three], join='outer')) self.run_scenario( df, lambda df: df.one.str.cat([df.two, df.three], join='outer')) def test_repeat(self): # verify that repeat works with a Series since this is # missing from doctests df = pd.DataFrame({ 'strings': ['A', 'B', 'C', 'D', 'E'], 'repeats': [3, 1, 4, 5, 2], }) self.run_scenario(df, lambda df: df.strings.str.repeat(df.repeats)) def test_rename(self): df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]}) self.run_scenario( df, lambda df: df.rename(columns={'B': 'C'}, index={ 0: 2, 2: 0 })) with expressions.allow_non_parallel_operations(): self.run_scenario( df, lambda df: df.rename( columns={'B': 'C'}, index={ 0: 2, 2: 0 }, errors='raise')) class FusionTest(unittest.TestCase): @staticmethod def fused_stages(p): return p.result.monitoring_metrics().query( metrics.MetricsFilter().with_name( fn_runner.FnApiRunner.NUM_FUSED_STAGES_COUNTER) )['counters'][0].result @staticmethod def create_animal_speed_input(p): return p | beam.Create([ AnimalSpeed('Aardvark', 5), AnimalSpeed('Ant', 2), AnimalSpeed('Elephant', 35), AnimalSpeed('Zebra', 40) ], reshuffle=False) def test_loc_filter(self): # TODO(https://github.com/apache/beam/issues/34549): This test relies on # monitoring_metrics property of the FnApiRunner which does not exist on # other runners like Prism. # https://github.com/apache/beam/blob/5f9cd73b7c9a2f37f83971ace3a399d633201dd1/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py#L1590 with beam.Pipeline('FnApiRunner') as p: _ = ( self.create_animal_speed_input(p) | transforms.DataframeTransform(lambda df: df[df.Speed > 10])) self.assertEqual(self.fused_stages(p), 1) def test_column_manipulation(self): def set_column(df, name, s): df[name] = s return df # TODO(https://github.com/apache/beam/issues/34549): This test relies on # monitoring_metrics property of the FnApiRunner which does not exist on # other runners like Prism. # https://github.com/apache/beam/blob/5f9cd73b7c9a2f37f83971ace3a399d633201dd1/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py#L1590 with beam.Pipeline('FnApiRunner') as p: _ = ( self.create_animal_speed_input(p) | transforms.DataframeTransform( lambda df: set_column(df, 'x', df.Speed + df.Animal.str.len()))) self.assertEqual(self.fused_stages(p), 1) class TransformPartsTest(unittest.TestCase): def test_rebatch(self): with beam.Pipeline() as p: sA = pd.Series(range(1000)) sB = sA * sA pcA = p | 'CreatePCollA' >> beam.Create([('k0', sA[::3]), ('k1', sA[1::3]), ('k2', sA[2::3])]) pcB = p | 'CreatePCollB' >> beam.Create([('k0', sB[::3]), ('k1', sB[1::3]), ('k2', sB[2::3])]) input = {'A': pcA, 'B': pcB} | beam.CoGroupByKey() output = input | beam.ParDo( transforms._ReBatch(target_size=sA.memory_usage())) # There should be exactly two elements, as the target size will be # hit when 2/3 of pcA and 2/3 of pcB is seen, but not before. assert_that(output | beam.combiners.Count.Globally(), equal_to([2])) # Sanity check that we got all the right values. assert_that( output | beam.Map(lambda x: x['A'].sum()) | 'SumA' >> beam.CombineGlobally(sum), equal_to([sA.sum()]), label='CheckValuesA') assert_that( output | beam.Map(lambda x: x['B'].sum()) | 'SumB' >> beam.CombineGlobally(sum), equal_to([sB.sum()]), label='CheckValuesB') if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/error.py0000644000175100017510000000301615206113032017360 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Python Dataflow error classes.""" # pytype: skip-file class BeamError(Exception): """Base class for all Beam errors.""" class PipelineError(BeamError): """An error in the pipeline object (e.g. a PValue not linked to it).""" class PValueError(BeamError): """An error related to a PValue object (e.g. value is not computed).""" class RunnerError(BeamError): """An error related to a Runner object (e.g. cannot find a runner to run).""" class RuntimeValueProviderError(RuntimeError): """An error related to a ValueProvider object raised during runtime.""" class SideInputError(BeamError): """An error related to a side input to a parallel Do operation.""" class TransformError(BeamError): """An error related to a PTransform object.""" ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996270.8963568 apache_beam-2.74.0/apache_beam/examples/0000755000175100017510000000000015206113157017503 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/__init__.py0000644000175100017510000000142015206113032021601 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/avro_nyc_trips.py0000644000175100017510000002224515206113032023113 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """An example that reads New York City for hire vehicle trips. An example that reads New York City for hire vehicle trips data from Google Cloud Storage and calculates various statistics including the passenger's price per trip, price per mile, and price per minute. Data originally downloaded as parquet files from the New York City government `website `_ before being converted to avro format. Here is the `documentation `_ for the data format and schema. To execute this pipeline locally using the DirectRunner, specify an output prefix on GCS::: --output gs://YOUR_OUTPUT_PREFIX To execute this pipeline using the Google Cloud Dataflow service, specify pipeline configuration in addition to the above::: --job_name NAME_FOR_YOUR_JOB --project YOUR_PROJECT_ID --region GCE_REGION --staging_location gs://YOUR_STAGING_DIRECTORY --temp_location gs://YOUR_TEMPORARY_DIRECTORY --runner DataflowRunner The default input is ``gs://apache-beam-samples/nyc_trip/avro/fhvhv_tripdata_2023-02.avro`` and takes about 15 minutes with 6 workers running on Dataflow. The default input can be overridden with --input argument. More data can be accessed at ``gs://apache-beam-samples/nyc_trip/avro/*``. Additionally, the original parquet files can be found at ``gs://apache-beam-samples/nyc_trip/parquet/*``. """ # pytype: skip-file import argparse import datetime import logging import pytz import apache_beam as beam from apache_beam.io import ReadFromAvro from apache_beam.io import WriteToAvro from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions SCHEMA = { 'fields': [ { 'name': 'service', 'type': 'string' }, { 'name': 'day', 'type': 'string' }, { 'name': 'total_price', 'type': 'double' }, { 'name': 'total_driver_pay', 'type': 'double' }, { 'name': 'total_trip_miles', 'type': 'double' }, { 'name': 'total_trip_minutes', 'type': 'double' }, { 'name': 'total_number_of_trips', 'type': 'long' }, { 'name': 'price_per_trip', 'type': 'double' }, { 'name': 'price_per_mile', 'type': 'double' }, { 'name': 'price_per_minute', 'type': 'double' }, { 'name': 'driver_pay_per_trip', 'type': 'double' }, { 'name': 'driver_pay_per_mile', 'type': 'double' }, { 'name': 'driver_pay_per_minute', 'type': 'double' }, { 'name': 'miles_per_hour', 'type': 'double' }, ], 'name': 'nyc_trip_prices', 'type': 'record', } class CreateKeyWithServiceAndDay(beam.DoFn): """Creates a key, value group. The key is the combination of the HVFHS Code converted to the ride-sharing company name and the day of the week of the ride. The value is the original record dictionary object. """ def process(self, record: dict): options = { 'HV0002': 'Juno', 'HV0003': 'Uber', 'HV0004': 'Via', 'HV0005': 'Lyft' } service = options.get(record['hvfhs_license_num']) if service: timestamp = None for k in ('request_datetime', 'on_scene_datetime', 'pickup_datetime', 'dropoff_datetime'): timestamp = timestamp or record[k] if timestamp: break day_of_the_week = datetime.datetime.fromtimestamp( timestamp / 1000.0, tz=pytz.timezone('America/New_York')).strftime('%a') yield (service, day_of_the_week), record class CalculatePricePerAttribute(beam.CombineFn): """Calculates the price for attribute. Calculates the total driver pay, price, miles, minutes, and trips per for hire vehicle service. And calculates the price per mile, minute, and trip for both the driver and passenger. """ def create_accumulator(self): total_price = 0.0 total_driver_pay = 0.0 total_trip_miles = 0.0 total_trip_time = 0.0 total_number_of_trips = 0 accumulator = ( total_price, total_driver_pay, total_trip_miles, total_trip_time, total_number_of_trips, ) return accumulator def add_input(self, accumulator, record): ( total_price, total_driver_pay, total_trip_miles, total_trip_time, total_number_of_trips, ) = accumulator return ( total_price + sum( record[name] for name in ( 'base_passenger_fare', 'tolls', 'bcf', 'sales_tax', 'congestion_surcharge', 'airport_fee', 'tips') if record[name] is not None), total_driver_pay + record['driver_pay'] + record['tips'], total_trip_miles + record['trip_miles'], total_trip_time + record['trip_time'], total_number_of_trips + 1, ) def merge_accumulators(self, accumulators): # calculates the sum for each item in the accumulator return tuple(sum(item) for item in zip(*accumulators)) def extract_output(self, accumulator): ( total_price, total_driver_pay, total_trip_miles, total_trip_time, total_number_of_trips, ) = accumulator total_trip_minutes = total_trip_time / 60 return { 'total_driver_pay': total_driver_pay, 'total_price': total_price, 'total_trip_miles': total_trip_miles, 'total_trip_minutes': total_trip_minutes, 'total_number_of_trips': total_number_of_trips, 'price_per_trip': total_price / total_number_of_trips, 'price_per_mile': total_price / total_trip_miles, 'price_per_minute': total_price / total_trip_minutes, 'driver_pay_per_trip': total_driver_pay / total_number_of_trips, 'driver_pay_per_mile': total_driver_pay / total_trip_miles, 'driver_pay_per_minute': total_driver_pay / total_trip_minutes, 'miles_per_hour': total_trip_miles / (total_trip_minutes / 60), } def flatten_group(element): """Flattens the key, value pair to a single record dictionary.""" key, record = element service, day_of_the_week = key record['service'] = service record['day'] = day_of_the_week return record def run(argv=None): """Runs the New York City trips pipeline. Args: argv: Pipeline options as a list of arguments. """ parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', default= 'gs://apache-beam-samples/nyc_trip/avro/fhvhv_tripdata_2023-02.avro', help='Input file of NYC FHV data to process. Larger dataset can be found ' 'here: gs://apache-beam-samples/nyc_trip/avro/*', ) parser.add_argument( '--output', dest='output', help='Output file to write results to.', required=True, ) known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = True with beam.Pipeline(options=pipeline_options) as p: ( # pylint: disable=expression-not-assigned p | ReadFromAvro(known_args.input) # filter records that have keys that are not None # and at least one timestamp that is not None | beam.Filter( lambda record: all( record[k] is not None for k in ( 'hvfhs_license_num', 'trip_miles', 'trip_time', 'base_passenger_fare', 'tips', 'driver_pay')) and any( record[k] is not None for k in ( 'request_datetime', 'on_scene_datetime', 'pickup_datetime', 'dropoff_datetime'))) | beam.ParDo(CreateKeyWithServiceAndDay()) | beam.CombinePerKey(CalculatePricePerAttribute()) | beam.Map(flatten_group) | WriteToAvro(known_args.output, SCHEMA, file_name_suffix='.avro')) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/avro_nyc_trips_it_test.py0000644000175100017510000001456015206113032024647 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for the avro New York City Trips example.""" # pytype: skip-file import logging import unittest import uuid from io import BytesIO import fastavro import pytest from apache_beam.examples import avro_nyc_trips from apache_beam.io.filesystems import FileSystems from apache_beam.testing.test_pipeline import TestPipeline class AvroNycTripsIT(unittest.TestCase): SCHEMA = { 'fields': [ { 'name': 'hvfhs_license_num', 'type': ['null', 'string'] }, { 'name': 'request_datetime', 'logicalType': 'timestamp-millis', 'type': ['null', 'long'] }, { 'name': 'trip_miles', 'type': ['null', 'double'] }, { 'name': 'trip_time', 'type': ['null', 'long'] }, { 'name': 'base_passenger_fare', 'type': ['null', 'double'] }, { 'name': 'tolls', 'type': ['null', 'double'] }, { 'name': 'bcf', 'type': ['null', 'double'] }, { 'name': 'sales_tax', 'type': ['null', 'double'] }, { 'name': 'congestion_surcharge', 'type': ['null', 'double'] }, { 'name': 'airport_fee', 'type': ['null', 'long'] }, { 'name': 'tips', 'type': ['null', 'double'] }, { 'name': 'driver_pay', 'type': ['null', 'double'] }, ], 'name': 'nyc_fhv_trips', 'type': 'record', } RECORDS = [ { 'hvfhs_license_num': 'HV0003', 'request_datetime': 1549008086000, 'trip_miles': 2.45, 'trip_time': 579, 'base_passenger_fare': 9.35, 'tolls': 0.0, 'bcf': 0.23, 'sales_tax': 0.83, 'congestion_surcharge': 0.0, 'airport_fee': None, 'tips': 0.0, 'driver_pay': 7.48 }, { 'hvfhs_license_num': 'HV0003', 'request_datetime': 1549009568000, 'trip_miles': 1.71, 'trip_time': 490, 'base_passenger_fare': 7.91, 'tolls': 0.0, 'bcf': 0.2, 'sales_tax': 0.7, 'congestion_surcharge': 0.0, 'airport_fee': None, 'tips': 2.0, 'driver_pay': 7.93 }, { 'hvfhs_license_num': 'HV0005', 'request_datetime': 1549010613000, 'trip_miles': 11.24, 'trip_time': 1739, 'base_passenger_fare': 29.77, 'tolls': 0.72, 'bcf': 0.76, 'sales_tax': 2.71, 'congestion_surcharge': 0.0, 'airport_fee': None, 'tips': 0.0, 'driver_pay': 22.09 }, { 'hvfhs_license_num': 'HV0005', 'request_datetime': 1549010420000, 'trip_miles': 5.71, 'trip_time': 1559, 'base_passenger_fare': 21.69, 'tolls': 0.24, 'bcf': 0.55, 'sales_tax': 1.95, 'congestion_surcharge': 0.0, 'airport_fee': None, 'tips': 0.0, 'driver_pay': 14.87 }, ] EXPECTED = [ { 'service': 'Lyft', 'day': 'Fri', 'total_price': 58.39, 'total_driver_pay': 36.96, 'total_trip_miles': 16.95, 'total_trip_minutes': 54.96666666666667, 'total_number_of_trips': 2, 'price_per_trip': 29.195, 'price_per_mile': 3.4448377581120946, 'price_per_minute': 1.0622801697998787, 'driver_pay_per_trip': 18.48, 'driver_pay_per_mile': 2.1805309734513276, 'driver_pay_per_minute': 0.6724075197089144, 'miles_per_hour': 18.502122498483928 }, { 'service': 'Uber', 'day': 'Fri', 'total_price': 21.22, 'total_driver_pay': 17.41, 'total_trip_miles': 4.16, 'total_trip_minutes': 17.816666666666666, 'total_number_of_trips': 2, 'price_per_trip': 10.61, 'price_per_mile': 5.100961538461538, 'price_per_minute': 1.1910196445275958, 'driver_pay_per_trip': 8.705, 'driver_pay_per_mile': 4.185096153846153, 'driver_pay_per_minute': 0.9771749298409729, 'miles_per_hour': 14.00935453695042 }, ] @pytest.mark.sickbay_dataflow @pytest.mark.no_xdist @pytest.mark.examples_postcommit def test_avro_nyc_trips_output_files_on_small_input(self): test_pipeline = TestPipeline(is_integration_test=True) # set up the files with expected content. temp_location = test_pipeline.get_option('temp_location') test_output = '/'.join((temp_location, str(uuid.uuid4()), 'result')) test_input = '/'.join((temp_location, str(uuid.uuid4()), 'input.avro')) # create avro data fo = BytesIO() fastavro.writer(fo, self.SCHEMA, self.RECORDS) fo.seek(0) # write avro test case with FileSystems.create(test_input) as f: f.write(fo.read()) extra_opts = {'input': test_input, 'output': test_output} avro_nyc_trips.run(test_pipeline.get_full_options_as_args(**extra_opts)) # load result avro file and compare metadata_list = FileSystems.match([f'{test_output}*'])[0].metadata_list result = [] for metadata in metadata_list: with FileSystems.open(metadata.path) as f: avro_reader = fastavro.reader(f) result.extend(avro_reader) result.sort(key=lambda x: x['service']) self.assertEqual(self.EXPECTED, result) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/avro_nyc_trips_test.py0000644000175100017510000001145415206113032024152 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for the avro New York City Trips example.""" # pytype: skip-file import unittest import apache_beam as beam from apache_beam.examples import avro_nyc_trips from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class AvroNycTripsTest(unittest.TestCase): def test_create_key_with_service_and_day(self): RECORDS = [ { 'hvfhs_license_num': 'HV0002', 'request_datetime': 1557705616000, 'on_scene_datetime': 1557792016000, 'pickup_datetime': 1557878416000, 'dropoff_datetime': 1557964816000, 'trip_miles': 2.45, }, { 'hvfhs_license_num': 'HV0003', 'request_datetime': None, 'on_scene_datetime': 1557792016000, 'pickup_datetime': 1557878416000, 'dropoff_datetime': 1557964816000, 'trip_miles': 3.45, }, { 'hvfhs_license_num': 'HV0004', 'request_datetime': None, 'on_scene_datetime': None, 'pickup_datetime': 1557878416000, 'dropoff_datetime': 1557964816000, 'trip_miles': 4.45, }, { 'hvfhs_license_num': 'HV0005', 'request_datetime': None, 'on_scene_datetime': None, 'pickup_datetime': None, 'dropoff_datetime': 1557964816000, 'trip_miles': 5.45, }, ] EXPECTED = [ (('Juno', 'Sun'), RECORDS[0]), (('Uber', 'Mon'), RECORDS[1]), (('Via', 'Tue'), RECORDS[2]), (('Lyft', 'Wed'), RECORDS[3]), ] with TestPipeline() as p: records = p | beam.Create(RECORDS) result = records | beam.ParDo(avro_nyc_trips.CreateKeyWithServiceAndDay()) assert_that(result, equal_to(EXPECTED)) def test_calculate_price_per_attribute(self): RECORDS = [ (('Uber', 'Fri'), { 'hvfhs_license_num': 'HV0003', 'request_datetime': 1549008086000, 'trip_miles': 2.45, 'trip_time': 579, 'base_passenger_fare': 9.35, 'tolls': 0.0, 'bcf': 0.23, 'sales_tax': 0.83, 'congestion_surcharge': 0.0, 'airport_fee': None, 'tips': 0.0, 'driver_pay': 7.48 }), (('Uber', 'Fri'), { 'hvfhs_license_num': 'HV0003', 'request_datetime': 1549009568000, 'trip_miles': 1.71, 'trip_time': 490, 'base_passenger_fare': 7.91, 'tolls': 0.0, 'bcf': 0.2, 'sales_tax': 0.7, 'congestion_surcharge': 0.0, 'airport_fee': None, 'tips': 2.0, 'driver_pay': 7.93 }), ] EXPECTED = [ (('Uber', 'Fri'), { 'total_price': 21.22, 'total_driver_pay': 17.41, 'total_trip_miles': 4.16, 'total_trip_minutes': 17.816666666666666, 'total_number_of_trips': 2, 'price_per_trip': 10.61, 'price_per_mile': 5.100961538461538, 'price_per_minute': 1.1910196445275958, 'driver_pay_per_trip': 8.705, 'driver_pay_per_mile': 4.185096153846153, 'driver_pay_per_minute': 0.9771749298409729, 'miles_per_hour': 14.00935453695042 }), ] with TestPipeline() as p: records = p | beam.Create(RECORDS) result = records | beam.CombinePerKey( avro_nyc_trips.CalculatePricePerAttribute()) assert_that(result, equal_to(EXPECTED)) def test_flatten_group(self): record = {'total_driver_pay': 123.54} element = (('Uber', 'Fri'), record) expected_record = { 'service': 'Uber', 'day': 'Fri', 'total_driver_pay': 123.54 } output_record = avro_nyc_trips.flatten_group(element) self.assertEqual(expected_record, output_record) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/california_housing_clustering.py0000644000175100017510000001012415206113032026145 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A pipeline that uses OnlineClustering transform to group houses with a similar value together. This example uses the California Housing Prices dataset from kaggle. https://www.kaggle.com/datasets/camnugent/california-housing-prices In the first step of the pipeline, the clustering model is trained using the OnlineKMeans transform, then the AssignClusterLabels transform assigns a cluster to each record in the dataset. This transform makes use of the RunInference API under the hood. In order to run this example: 1. Download the data from kaggle as csv 2. Run `python california_housing_clustering.py --input --checkpoints_path ` # pylint: disable=line-too-long """ import argparse import numpy as np import apache_beam as beam from apache_beam import pvalue from apache_beam.dataframe.convert import to_pcollection from apache_beam.dataframe.io import read_csv from apache_beam.examples.online_clustering import AssignClusterLabelsInMemoryModel from apache_beam.examples.online_clustering import OnlineClustering from apache_beam.examples.online_clustering import OnlineKMeans from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', required=True, help='A csv file containing the data that needs to be clustered.') parser.add_argument( '--checkpoints_path', dest='checkpoints_path', required=True, help='A path to a directory where model checkpoints can be stored.') return parser.parse_known_args(argv) def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) data = pipeline | read_csv(known_args.input) features = ['longitude', 'latitude', 'median_income'] housing_features = to_pcollection(data[features]) # 1. Calculate clustering centers and save model to persistent storage model = ( housing_features | beam.Map(lambda record: list(record)) | "Train clustering model" >> OnlineClustering( OnlineKMeans, n_clusters=6, batch_size=256, cluster_args={}, checkpoints_path=known_args.checkpoints_path)) # 2. Calculate labels for all records in the dataset # using the trained clustering model using in memory model _ = ( housing_features | beam.Map(lambda sample: np.array(sample)) | "RunInference" >> AssignClusterLabelsInMemoryModel( model=pvalue.AsSingleton(model), model_id="kmeans", n_clusters=6, batch_size=512) | beam.Map(print)) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': run() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996270.9023573 apache_beam-2.74.0/apache_beam/examples/complete/0000755000175100017510000000000015206113157021313 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/__init__.py0000644000175100017510000000142015206113032023411 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/autocomplete.py0000644000175100017510000000612715206113032024364 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A workflow emitting the top k most common words for each prefix.""" # pytype: skip-file import argparse import logging import re import apache_beam as beam from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions def run(argv=None): parser = argparse.ArgumentParser() parser.add_argument('--input', required=True, help='Input file to process.') parser.add_argument( '--output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = True with beam.Pipeline(options=pipeline_options) as p: def format_result(prefix_candidates): (prefix, candidates) = prefix_candidates return '%s: %s' % (prefix, candidates) ( # pylint: disable=expression-not-assigned p | 'read' >> ReadFromText(known_args.input) | 'split' >> beam.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x)) | 'TopPerPrefix' >> TopPerPrefix(5) | 'format' >> beam.Map(format_result) | 'write' >> WriteToText(known_args.output)) class TopPerPrefix(beam.PTransform): def __init__(self, count): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.PTransform.__init__(self) self._count = count def expand(self, words): """Compute the most common words for each possible prefixes. Args: words: a PCollection of strings Returns: A PCollection of most common words with each prefix, in the form (prefix, [(count, word), (count, word), ...]) """ return ( words | beam.combiners.Count.PerElement() | beam.FlatMap(extract_prefixes) | beam.combiners.Top.LargestPerKey(self._count)) def extract_prefixes(element): word, count = element for k in range(1, len(word) + 1): prefix = word[:k] yield prefix, (count, word) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/autocomplete_it_test.py0000644000175100017510000000650115206113032026113 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for Autocomplete example.""" # pytype: skip-file import logging import re import unittest import uuid import pytest from apache_beam.examples.complete import autocomplete from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import create_file from apache_beam.testing.test_utils import read_files_from_pattern def format_output_file(output_string): def extract_prefix_topk_words_tuples(line): match = re.match(r'(.*): \[(.*)\]', line) prefix = match.group(1) topK_words_string = extract_top_k_words_tuples(match.group(2)) return prefix, topK_words_string def extract_top_k_words_tuples(top_k_words_string): top_k_list = top_k_words_string.split("), (") return tuple( map( lambda top_k_string: tuple(format_top_k_tuples(top_k_string)), top_k_list)) def format_top_k_tuples(top_k_string): (frequency, words) = top_k_string.replace('(', '').replace(')', '').replace( '\"', '').replace('\'', '').replace(' ', '').split(',') return int(frequency), words return list( map( lambda line: extract_prefix_topk_words_tuples(line), output_string.split('\n'))) class AutocompleteIT(unittest.TestCase): WORDS = ['this', 'this', 'that', 'to', 'to', 'to'] EXPECTED_PREFIXES = [ ('t', ((3, 'to'), (2, 'this'), (1, 'that'))), ('to', ((3, 'to'), )), ('th', ((2, 'this'), (1, 'that'))), ('thi', ((2, 'this'), )), ('this', ((2, 'this'), )), ('tha', ((1, 'that'), )), ('that', ((1, 'that'), )), ] @pytest.mark.no_xdist @pytest.mark.examples_postcommit def test_autocomplete_output_files_on_small_input(self): test_pipeline = TestPipeline(is_integration_test=True) # Setup the files with expected content. OUTPUT_FILE_DIR = \ 'gs://temp-storage-for-end-to-end-tests/py-it-cloud/output' output = '/'.join([OUTPUT_FILE_DIR, str(uuid.uuid4()), 'result']) INPUT_FILE_DIR = \ 'gs://temp-storage-for-end-to-end-tests/py-it-cloud/input' input = '/'.join([INPUT_FILE_DIR, str(uuid.uuid4()), 'input.txt']) create_file(input, ' '.join(self.WORDS)) extra_opts = {'input': input, 'output': output} autocomplete.run(test_pipeline.get_full_options_as_args(**extra_opts)) # Load result file and compare. result = read_files_from_pattern('%s*' % output).strip() self.assertEqual( sorted(self.EXPECTED_PREFIXES), sorted(format_output_file(result))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/autocomplete_test.py0000644000175100017510000000505415206113032025421 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the autocomplete example.""" # pytype: skip-file import unittest import pytest import apache_beam as beam from apache_beam.examples.complete import autocomplete from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import compute_hash from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class AutocompleteTest(unittest.TestCase): WORDS = ['this', 'this', 'that', 'to', 'to', 'to'] KINGLEAR_HASH_SUM = 268011785062540 KINGLEAR_INPUT = 'gs://dataflow-samples/shakespeare/kinglear.txt' EXPECTED_PREFIXES = [ ('t', ((3, 'to'), (2, 'this'), (1, 'that'))), ('to', ((3, 'to'), )), ('th', ((2, 'this'), (1, 'that'))), ('thi', ((2, 'this'), )), ('this', ((2, 'this'), )), ('tha', ((1, 'that'), )), ('that', ((1, 'that'), )), ] def test_top_prefixes(self): with TestPipeline() as p: words = p | beam.Create(self.WORDS) result = words | autocomplete.TopPerPrefix(5) # values must be hashable for now result = result | beam.Map(lambda k_vs: (k_vs[0], tuple(k_vs[1]))) assert_that(result, equal_to(self.EXPECTED_PREFIXES)) @pytest.mark.it_postcommit def test_autocomplete_it(self): with TestPipeline(is_integration_test=True) as p: words = p | beam.io.ReadFromText(self.KINGLEAR_INPUT) result = words | autocomplete.TopPerPrefix(10) # values must be hashable for now result = result | beam.Map( lambda k_vs: [k_vs[0], k_vs[1][0][0], k_vs[1][0][1]]) checksum = ( result | beam.Map(lambda x: int(compute_hash(x)[:8], 16)) | beam.CombineGlobally(sum)) assert_that(checksum, equal_to([self.KINGLEAR_HASH_SUM])) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/distribopt.py0000644000175100017510000003315515206113032024047 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Example illustrating the use of Apache Beam for solving distributing optimization tasks. This example solves an optimization problem which consists of distributing a number of crops to grow in several greenhouses. The decision where to grow the crop has an impact on the production parameters associated with the greenhouse, which affects the total cost of production at the greenhouse. Additionally, each crop needs to be transported to a customer so the decision where to grow the crop has an impact on the transportation costs as well. This type of optimization problems are known as mixed-integer programs as they exist of discrete parameters (do we produce a crop in greenhouse A, B or C?) and continuous parameters (the greenhouse production parameters). Running this example requires NumPy and SciPy. The input consists of a CSV file with the following columns (Tx representing the transporation cost/unit if the crop is produced in greenhouse x): Crop name, Quantity, Ta, Tb, Tc, .... Example input file with 5 crops and 3 greenhouses (a transporation cost of 0 forbids production of the crop in a greenhouse): OP01,8,12,0,12 OP02,30,14,3,12 OP03,25,7,3,14 OP04,87,7,2,2 OP05,19,1,7,10 The pipeline consists of three phases: - Creating a grid of mappings (assignment of each crop to a greenhouse) - For each mapping and each greenhouse, optimization of the production parameters for cost, addition of the transporation costs, and aggregation of the costs for each mapping. - Selecting the mapping with the lowest cost. """ # pytype: skip-file import argparse import logging import string import uuid from collections import defaultdict import numpy as np from scipy.optimize import minimize import apache_beam as beam from apache_beam import pvalue from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions class Simulator(object): """Greenhouse simulation for the optimization of greenhouse parameters.""" def __init__(self, quantities): self.quantities = np.atleast_1d(quantities) self.A = np.array([[3.0, 10, 30], [0.1, 10, 35], [3.0, 10, 30], [0.1, 10, 35]]) self.P = 1e-4 * np.array([[3689, 1170, 2673], [4699, 4387, 7470], [1091, 8732, 5547], [381, 5743, 8828]]) a0 = np.array([[1.0, 1.2, 3.0, 3.2]]) coeff = np.sum(np.cos(np.dot(a0.T, self.quantities[None, :])), axis=1) self.alpha = coeff / np.sum(coeff) def simulate(self, xc): # Map the input parameter to a cost for each crop. weighted_distance = np.sum(self.A * np.square(xc - self.P), axis=1) f = -np.sum(self.alpha * np.exp(-weighted_distance)) return np.square(f) * np.log(self.quantities) class CreateGrid(beam.PTransform): """A transform for generating the mapping grid. Input: Formatted records of the input file, e.g., { 'crop': 'OP009', 'quantity': 102, 'transport_costs': [('A', None), ('B', 3), ('C', 8)] } Output: tuple (mapping_identifier, {crop -> greenhouse}) """ class PreGenerateMappings(beam.DoFn): """ParDo implementation forming based on two elements a small sub grid. This facilitates parallellization of the grid generation. Emits two PCollections: the subgrid represented as collection of lists of two tuples, and a list of remaining records. Both serve as an input to GenerateMappings. """ def process(self, element): records = list(element[1]) # Split of 2 crops and pre-generate the subgrid. # Select the crop with highest number of possible greenhouses: # in case two crops with only a single possible greenhouse were selected # the subgrid would consist of only 1 element. best_split = np.argsort([-len(r['transport_costs']) for r in records])[:2] rec1 = records[best_split[0]] rec2 = records[best_split[1]] # Generate & emit all combinations for a in rec1['transport_costs']: if a[1]: for b in rec2['transport_costs']: if b[1]: combination = [(rec1['crop'], a[0]), (rec2['crop'], b[0])] yield pvalue.TaggedOutput('splitted', combination) # Pass on remaining records remaining = [rec for i, rec in enumerate(records) if i not in best_split] yield pvalue.TaggedOutput('combine', remaining) class GenerateMappings(beam.DoFn): """ParDo implementation to generate all possible mappings. Input: output of PreGenerateMappings Output: tuples of the form (mapping_identifier, {crop -> greenhouse}) """ @staticmethod def _coordinates_to_greenhouse(coordinates, greenhouses, crops): # Map the grid coordinates back to greenhouse labels arr = [] for coord in coordinates: arr.append(greenhouses[coord]) return dict(zip(crops, np.array(arr))) def process(self, element, records): # Generate available greenhouses and grid coordinates for each crop. grid_coordinates = [] for rec in records: # Get indices for available greenhouses (w.r.t crops) filtered = [i for i, av in enumerate(rec['transport_costs']) if av[1]] grid_coordinates.append(filtered) # Generate all mappings grid = np.vstack(list(map(np.ravel, np.meshgrid(*grid_coordinates)))).T crops = [rec['crop'] for rec in records] greenhouses = [rec[0] for rec in records[0]['transport_costs']] for point in grid: # translate back to greenhouse label mapping = self._coordinates_to_greenhouse(point, greenhouses, crops) assert all(rec[0] not in mapping for rec in element) # include the incomplete mapping of 2 crops mapping.update(element) # include identifier yield (uuid.uuid4().hex, mapping) def expand(self, records): o = ( records | 'pair one' >> beam.Map(lambda x: (1, x)) | 'group all records' >> beam.GroupByKey() | 'split one of' >> beam.ParDo(self.PreGenerateMappings()).with_outputs( 'splitted', 'combine')) # Create mappings, and prevent fusion (this limits the parallelization # in the optimization step) mappings = ( o.splitted | 'create mappings' >> beam.ParDo( self.GenerateMappings(), pvalue.AsSingleton(o.combine)) | 'prevent fusion' >> beam.Reshuffle()) return mappings class OptimizeGrid(beam.PTransform): """A transform for optimizing all greenhouses of the mapping grid.""" class CreateOptimizationTasks(beam.DoFn): """ Create tasks for optimization. Input: (mapping_identifier, {crop -> greenhouse}) Output: ((mapping_identifier, greenhouse), [(crop, quantity),...]) """ def process(self, element, quantities): mapping_identifier, mapping = element # Create (crop, quantity) lists for each greenhouse greenhouses = defaultdict(list) for crop, greenhouse in mapping.items(): quantity = quantities[crop] greenhouses[greenhouse].append((crop, quantity)) # Create input for OptimizeProductParameters for greenhouse, crops in greenhouses.items(): key = (mapping_identifier, greenhouse) yield (key, crops) class OptimizeProductParameters(beam.DoFn): """Solve the optimization task to determine optimal production parameters. Input: ((mapping_identifier, greenhouse), [(crop, quantity),...]) Two outputs: - solution: (mapping_identifier, (greenhouse, [production parameters])) - costs: (crop, greenhouse, mapping_identifier, cost) """ @staticmethod def _optimize_production_parameters(sim): # setup initial starting point & bounds x0 = 0.5 * np.ones(3) bounds = list(zip(np.zeros(3), np.ones(3))) # Run L-BFGS-B optimizer result = minimize(lambda x: np.sum(sim.simulate(x)), x0, bounds=bounds) # Ensure result.x is always a list, regardless of NumPy version x_values = result.x if isinstance(result.x, list) else result.x.tolist() # Ensure simulation output is also properly converted costs = sim.simulate(result.x) costs = costs if isinstance(costs, list) else costs.tolist() return x_values, costs def process(self, element): mapping_identifier, greenhouse = element[0] crops, quantities = zip(*element[1]) sim = Simulator(quantities) optimum, costs = self._optimize_production_parameters(sim) # Ensure NumPy arrays are converted to lists before yielding if isinstance(optimum, np.ndarray): optimum = optimum.tolist() if isinstance(costs, np.ndarray): costs = costs.tolist() solution = (mapping_identifier, (greenhouse, optimum)) yield pvalue.TaggedOutput('solution', solution) for crop, cost, quantity in zip(crops, costs, quantities): costs = (crop, greenhouse, mapping_identifier, cost * quantity) yield pvalue.TaggedOutput('costs', costs) def expand(self, inputs): mappings, quantities = inputs opt = ( mappings | 'optimization tasks' >> beam.ParDo( self.CreateOptimizationTasks(), pvalue.AsDict(quantities)) | 'optimize' >> beam.ParDo( self.OptimizeProductParameters()).with_outputs('costs', 'solution')) return opt class CreateTransportData(beam.DoFn): """Transform records to pvalues ((crop, greenhouse), transport_cost)""" def process(self, record): crop = record['crop'] for greenhouse, transport_cost in record['transport_costs']: yield ((crop, greenhouse), transport_cost) def add_transport_costs(element, transport, quantities): """Adds the transport cost for the crop to the production cost. elements are of the form (crop, greenhouse, mapping, cost), the cost only corresponds to the production cost. Return the same format, but including the transport cost. """ crop = element[0] cost = element[3] # lookup & compute cost transport_key = element[:2] transport_cost = transport[transport_key] * quantities[crop] return element[:3] + (cost + transport_cost, ) def parse_input(line): # Process each line of the input file to a dict representing each crop # and the transport costs columns = line.split(',') # Assign each greenhouse a character transport_costs = [] for greenhouse, cost in zip(string.ascii_uppercase, columns[2:]): info = (greenhouse, int(cost) if cost else None) transport_costs.append(info) return { 'crop': columns[0], 'quantity': int(columns[1]), 'transport_costs': transport_costs } def format_output(element): """Transforms the datastructure (unpack lists introduced by CoGroupByKey) before writing the result to file. """ result = element[1] result['cost'] = result['cost'][0] result['production'] = dict(result['production']) result['mapping'] = result['mapping'][0] return result def run(argv=None, save_main_session=True): parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', required=True, help='Input description to process.') parser.add_argument( '--output', dest='output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session with beam.Pipeline(options=pipeline_options) as p: # Parse input file records = ( p | 'read' >> beam.io.ReadFromText(known_args.input) | 'process input' >> beam.Map(parse_input)) # Create two pcollections, used as side inputs transport = ( records | 'create transport' >> beam.ParDo(CreateTransportData())) quantities = ( records | 'create quantities' >> beam.Map(lambda r: (r['crop'], r['quantity']))) # Generate all mappings and optimize greenhouse production parameters mappings = records | CreateGrid() opt = (mappings, quantities) | OptimizeGrid() # Then add the transport costs and sum costs per crop. costs = ( opt.costs | 'include transport' >> beam.Map( add_transport_costs, pvalue.AsDict(transport), pvalue.AsDict(quantities)) | 'drop crop and greenhouse' >> beam.Map(lambda x: (x[2], x[3])) | 'aggregate crops' >> beam.CombinePerKey(sum)) # Join cost, mapping and production settings solution on mapping identifier. # Then select best. join_operands = { 'cost': costs, 'production': opt.solution, 'mapping': mappings } best = ( join_operands | 'join' >> beam.CoGroupByKey() | 'select best' >> beam.CombineGlobally( min, key=lambda x: x[1]['cost']).without_defaults() | 'format output' >> beam.Map(format_output)) # pylint: disable=expression-not-assigned best | 'write optimum' >> beam.io.WriteToText(known_args.output) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/distribopt_test.py0000644000175100017510000000636515206113032025111 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the distrib_optimization example.""" # pytype: skip-file import logging import unittest import uuid from ast import literal_eval as make_tuple import numpy as np import pytest from mock import MagicMock from mock import patch from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import create_file from apache_beam.testing.test_utils import read_files_from_pattern FILE_CONTENTS = 'OP01,8,12,0,12\n' \ 'OP02,30,14,3,12\n' \ 'OP03,25,7,3,14\n' \ 'OP04,87,7,2,2\n' \ 'OP05,19,1,7,10' EXPECTED_MAPPING = { 'OP01': 'A', 'OP02': 'B', 'OP03': 'B', 'OP04': 'C', 'OP05': 'A' } class DistribOptimizationTest(unittest.TestCase): #TODO(https://github.com/apache/beam/issues/23606) Fix and enable @pytest.mark.sickbay_dataflow @pytest.mark.examples_postcommit def test_basics(self): test_pipeline = TestPipeline(is_integration_test=True) # Setup the files with expected content. temp_location = test_pipeline.get_option('temp_location') input = '/'.join([temp_location, str(uuid.uuid4()), 'input.txt']) output = '/'.join([temp_location, str(uuid.uuid4()), 'result']) create_file(input, FILE_CONTENTS) extra_opts = {'input': input, 'output': output} # Run pipeline # Avoid dependency on SciPy scipy_mock = MagicMock() result_mock = MagicMock(x=np.ones(3)) scipy_mock.optimize.minimize = MagicMock(return_value=result_mock) modules = {'scipy': scipy_mock, 'scipy.optimize': scipy_mock.optimize} with patch.dict('sys.modules', modules): from apache_beam.examples.complete import distribopt distribopt.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) # Load result file and compare. lines = read_files_from_pattern('%s*' % output).splitlines() # Only 1 result self.assertEqual(len(lines), 1) # Handle NumPy string representation before parsing cleaned_line = lines[0].replace("np.str_('", "'").replace("')", "'") # parse result line and verify optimum optimum = make_tuple(cleaned_line) self.assertAlmostEqual(optimum['cost'], 454.39597, places=3) self.assertDictEqual(optimum['mapping'], EXPECTED_MAPPING) production = optimum['production'] for plant in ['A', 'B', 'C']: np.testing.assert_almost_equal(production[plant], np.ones(3)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/estimate_pi.py0000644000175100017510000001034115206113032024157 0ustar00runnerrunner# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A workflow that uses a simple Monte Carlo method to estimate π. The algorithm computes the fraction of points drawn uniformly within the unit square that also fall in the quadrant of the unit circle that overlaps the square. A simple area calculation shows that this fraction should be π/4, so we multiply our counts ratio by four to estimate π. """ # pytype: skip-file import argparse import json import logging import random from collections.abc import Iterable from typing import Any import apache_beam as beam from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions @beam.typehints.with_output_types(tuple[int, int, int]) @beam.typehints.with_input_types(int) def run_trials(runs): """Run trials and return a 3-tuple representing the results. Args: runs: Number of trial runs to be executed. Returns: A 3-tuple (total trials, inside trials, 0). The final zero is needed solely to make sure that the combine_results function has same type for inputs and outputs (a requirement for combiner functions). """ inside_runs = 0 for _ in range(runs): x = random.uniform(0, 1) y = random.uniform(0, 1) inside_runs += 1 if x * x + y * y <= 1.0 else 0 return runs, inside_runs, 0 @beam.typehints.with_output_types(tuple[int, int, float]) @beam.typehints.with_input_types(Iterable[tuple[int, int, Any]]) def combine_results(results): """Combiner function to sum up trials and compute the estimate. Args: results: An iterable of 3-tuples (total trials, inside trials, ignored). Returns: A 3-tuple containing the sum of total trials, sum of inside trials, and the probability computed from the two numbers. """ # TODO(silviuc): Do we guarantee that argument can be iterated repeatedly? # Should document one way or the other. total, inside = sum(r[0] for r in results), sum(r[1] for r in results) return total, inside, 4 * float(inside) / total class JsonCoder(object): """A JSON coder used to format the final result.""" def encode(self, x): return json.dumps(x).encode('utf-8') class EstimatePiTransform(beam.PTransform): """Runs 10M trials, and combine the results to estimate pi.""" def __init__(self, tries_per_work_item=100000): self.tries_per_work_item = tries_per_work_item def expand(self, pcoll): # A hundred work items of a hundred thousand tries each. return ( pcoll | 'Initialize' >> beam.Create( [self.tries_per_work_item] * 100).with_output_types(int) | 'Run trials' >> beam.Map(run_trials) | 'Sum' >> beam.CombineGlobally(combine_results).without_defaults()) def run(argv=None): parser = argparse.ArgumentParser() parser.add_argument( '--output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = True with beam.Pipeline(options=pipeline_options) as p: ( # pylint: disable=expression-not-assigned p | EstimatePiTransform() | WriteToText(known_args.output, coder=JsonCoder())) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/estimate_pi_it_test.py0000644000175100017510000000374315206113032025722 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for Estimate Pi example.""" # pytype: skip-file import json import logging import unittest import uuid import pytest from apache_beam.examples.complete import estimate_pi from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import read_files_from_pattern class EstimatePiIT(unittest.TestCase): @pytest.mark.no_xdist @pytest.mark.examples_postcommit def test_estimate_pi_output_file(self): test_pipeline = TestPipeline(is_integration_test=True) OUTPUT_FILE = \ 'gs://temp-storage-for-end-to-end-tests/py-it-cloud/output' output = '/'.join([OUTPUT_FILE, str(uuid.uuid4()), 'result']) extra_opts = {'output': output} estimate_pi.run(test_pipeline.get_full_options_as_args(**extra_opts)) # Load result file and compare. result = read_files_from_pattern('%s*' % output) [_, _, estimated_pi] = json.loads(result.strip()) # Note: Probabilistically speaking this test can fail with a probability # that is very small (VERY) given that we run at least 100 thousand # trials. self.assertTrue(3.125 <= estimated_pi <= 3.155) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/estimate_pi_test.py0000644000175100017510000000346515206113032025227 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the estimate_pi example.""" # pytype: skip-file import logging import unittest from apache_beam.examples.complete import estimate_pi from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import BeamAssertException from apache_beam.testing.util import assert_that def in_between(lower, upper): def _in_between(actual): _, _, estimate = actual[0] if estimate < lower or estimate > upper: raise BeamAssertException( 'Failed assert: %f not in [%f, %f]' % (estimate, lower, upper)) return _in_between class EstimatePiTest(unittest.TestCase): def test_basics(self): with TestPipeline() as p: result = p | 'Estimate' >> estimate_pi.EstimatePiTransform(5000) # Note: Probabilistically speaking this test can fail with a probability # that is very small (VERY) given that we run at least 500 thousand # trials. assert_that(result, in_between(3.125, 3.155)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996270.907358 apache_beam-2.74.0/apache_beam/examples/complete/game/0000755000175100017510000000000015206113157022224 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/game/__init__.py0000644000175100017510000000142015206113032024322 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/game/game_stats.py0000644000175100017510000003531515206113032024724 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Fourth in a series of four pipelines that tell a story in a 'gaming' domain. New concepts: session windows and finding session duration; use of both singleton and non-singleton side inputs. This pipeline builds on the {@link LeaderBoard} functionality, and adds some "business intelligence" analysis: abuse detection and usage patterns. The pipeline derives the Mean user score sum for a window, and uses that information to identify likely spammers/robots. (The robots have a higher click rate than the human users). The 'robot' users are then filtered out when calculating the team scores. Additionally, user sessions are tracked: that is, we find bursts of user activity using session windows. Then, the mean session duration information is recorded in the context of subsequent fixed windowing. (This could be used to tell us what games are giving us greater user retention). Run injector.Injector to generate pubsub data for this pipeline. The Injector documentation provides more detail on how to do this. The injector is currently implemented in Java only, it can be used from the Java SDK. The PubSub topic you specify should be the same topic to which the Injector is publishing. To run the Java injector: /examples/java$ mvn compile exec:java \ -Dexec.mainClass=org.apache.beam.examples.complete.game.injector.Injector \ -Dexec.args="$PROJECT_ID $PUBSUB_TOPIC none" For a description of the usage and options, use -h or --help. To specify a different runner: --runner YOUR_RUNNER NOTE: When specifying a different runner, additional runner-specific options may have to be passed in as well EXAMPLES -------- # DirectRunner python game_stats.py \ --project $PROJECT_ID \ --topic projects/$PROJECT_ID/topics/$PUBSUB_TOPIC \ --dataset $BIGQUERY_DATASET # DataflowRunner python game_stats.py \ --project $PROJECT_ID \ --region $REGION_ID \ --topic projects/$PROJECT_ID/topics/$PUBSUB_TOPIC \ --dataset $BIGQUERY_DATASET \ --runner DataflowRunner \ --temp_location gs://$BUCKET/user_score/temp """ # pytype: skip-file import argparse import csv import logging import sys import time from datetime import datetime import apache_beam as beam from apache_beam.metrics.metric import Metrics from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions def timestamp2str(t, fmt='%Y-%m-%d %H:%M:%S.000'): """Converts a unix timestamp into a formatted string.""" return datetime.fromtimestamp(t).strftime(fmt) class ParseGameEventFn(beam.DoFn): """Parses the raw game event info into a Python dictionary. Each event line has the following format: username,teamname,score,timestamp_in_ms,readable_time e.g.: user2_AsparagusPig,AsparagusPig,10,1445230923951,2015-11-02 09:09:28.224 The human-readable time string is not used here. """ def __init__(self): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.DoFn.__init__(self) self.num_parse_errors = Metrics.counter(self.__class__, 'num_parse_errors') def process(self, elem): try: row = list(csv.reader([elem]))[0] yield { 'user': row[0], 'team': row[1], 'score': int(row[2]), 'timestamp': int(row[3]) / 1000.0, } except: # pylint: disable=bare-except # Log and count parse errors self.num_parse_errors.inc() logging.error('Parse error on "%s"', elem) class ExtractAndSumScore(beam.PTransform): """A transform to extract key/score information and sum the scores. The constructor argument `field` determines whether 'team' or 'user' info is extracted. """ def __init__(self, field): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.PTransform.__init__(self) self.field = field def expand(self, pcoll): return ( pcoll | beam.Map(lambda elem: (elem[self.field], elem['score'])) | beam.CombinePerKey(sum)) class TeamScoresDict(beam.DoFn): """Formats the data into a dictionary of BigQuery columns with their values Receives a (team, score) pair, extracts the window start timestamp, and formats everything together into a dictionary. The dictionary is in the format {'bigquery_column': value} """ def process(self, team_score, window=beam.DoFn.WindowParam): team, score = team_score start = timestamp2str(int(window.start)) yield { 'team': team, 'total_score': score, 'window_start': start, 'processing_time': timestamp2str(int(time.time())) } class WriteToBigQuery(beam.PTransform): """Generate, format, and write BigQuery table row information.""" def __init__(self, table_name, dataset, schema, project): """Initializes the transform. Args: table_name: Name of the BigQuery table to use. dataset: Name of the dataset to use. schema: Dictionary in the format {'column_name': 'bigquery_type'} project: Name of the Cloud project containing BigQuery table. """ # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.PTransform.__init__(self) self.table_name = table_name self.dataset = dataset self.schema = schema self.project = project def get_schema(self): """Build the output table schema.""" return ', '.join('%s:%s' % (col, self.schema[col]) for col in self.schema) def expand(self, pcoll): return ( pcoll | 'ConvertToRow' >> beam.Map(lambda elem: {col: elem[col] for col in self.schema}) | beam.io.WriteToBigQuery( self.table_name, self.dataset, self.project, self.get_schema())) # [START abuse_detect] class CalculateSpammyUsers(beam.PTransform): """Filter out all but those users with a high clickrate, which we will consider as 'spammy' users. We do this by finding the mean total score per user, then using that information as a side input to filter out all but those user scores that are larger than (mean * SCORE_WEIGHT). """ SCORE_WEIGHT = 2.5 def expand(self, user_scores): # Get the sum of scores for each user. sum_scores = (user_scores | 'SumUsersScores' >> beam.CombinePerKey(sum)) # Extract the score from each element, and use it to find the global mean. global_mean_score = ( sum_scores | beam.Values() | beam.CombineGlobally(beam.combiners.MeanCombineFn())\ .as_singleton_view()) # Filter the user sums using the global mean. filtered = ( sum_scores # Use the derived mean total score (global_mean_score) as a side input. | 'ProcessAndFilter' >> beam.Filter( lambda key_score, global_mean:\ key_score[1] > global_mean * self.SCORE_WEIGHT, global_mean_score)) return filtered # [END abuse_detect] class UserSessionActivity(beam.DoFn): """Calculate and output an element's session duration, in seconds.""" def process(self, elem, window=beam.DoFn.WindowParam): yield (window.end.micros - window.start.micros) // 1000000 def run(argv=None, save_main_session=True): """Main entry point; defines and runs the hourly_team_score pipeline.""" parser = argparse.ArgumentParser() parser.add_argument('--topic', type=str, help='Pub/Sub topic to read from') parser.add_argument( '--subscription', type=str, help='Pub/Sub subscription to read from') parser.add_argument( '--dataset', type=str, required=True, help='BigQuery Dataset to write tables to. ' 'Must already exist.') parser.add_argument( '--table_name', type=str, default='game_stats', help='The BigQuery table name. Should not already exist.') parser.add_argument( '--fixed_window_duration', type=int, default=60, help='Numeric value of fixed window duration for user ' 'analysis, in minutes') parser.add_argument( '--session_gap', type=int, default=5, help='Numeric value of gap between user sessions, ' 'in minutes') parser.add_argument( '--user_activity_window_duration', type=int, default=30, help='Numeric value of fixed window for finding mean of ' 'user session duration, in minutes') args, pipeline_args = parser.parse_known_args(argv) if args.topic is None and args.subscription is None: parser.print_usage() print(sys.argv[0] + ': error: one of --topic or --subscription is required') sys.exit(1) options = PipelineOptions(pipeline_args) # We also require the --project option to access --dataset if options.view_as(GoogleCloudOptions).project is None: parser.print_usage() print(sys.argv[0] + ': error: argument --project is required') sys.exit(1) fixed_window_duration = args.fixed_window_duration * 60 session_gap = args.session_gap * 60 user_activity_window_duration = args.user_activity_window_duration * 60 # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). options.view_as(SetupOptions).save_main_session = save_main_session # Enforce that this pipeline is always run in streaming mode options.view_as(StandardOptions).streaming = True with beam.Pipeline(options=options) as p: # Read game events from Pub/Sub using custom timestamps, which # are extracted from the data elements, and parse the data. if args.subscription: scores = p | 'ReadPubSub' >> beam.io.ReadFromPubSub( subscription=args.subscription) else: scores = p | 'ReadPubSub' >> beam.io.ReadFromPubSub(topic=args.topic) raw_events = ( scores | 'DecodeString' >> beam.Map(lambda b: b.decode('utf-8')) | 'ParseGameEventFn' >> beam.ParDo(ParseGameEventFn()) | 'AddEventTimestamps' >> beam.Map( lambda elem: beam.window.TimestampedValue(elem, elem['timestamp']))) # Extract username/score pairs from the event stream user_events = ( raw_events | 'ExtractUserScores' >> beam.Map(lambda elem: (elem['user'], elem['score']))) # Calculate the total score per user over fixed windows, and cumulative # updates for late data spammers_view = ( user_events | 'UserFixedWindows' >> beam.WindowInto( beam.window.FixedWindows(fixed_window_duration)) # Filter out everyone but those with (SCORE_WEIGHT * avg) clickrate. # These might be robots/spammers. | 'CalculateSpammyUsers' >> CalculateSpammyUsers() # Derive a view from the collection of spammer users. It will be used as # a side input in calculating the team score sums, below | 'CreateSpammersView' >> beam.CombineGlobally( beam.combiners.ToDictCombineFn()).as_singleton_view()) # [START filter_and_calc] # Calculate the total score per team over fixed windows, and emit cumulative # updates for late data. Uses the side input derived above --the set of # suspected robots-- to filter out scores from those users from the sum. # Write the results to BigQuery. ( # pylint: disable=expression-not-assigned raw_events | 'WindowIntoFixedWindows' >> beam.WindowInto( beam.window.FixedWindows(fixed_window_duration)) # Filter out the detected spammer users, using the side input derived # above | 'FilterOutSpammers' >> beam.Filter( lambda elem, spammers: elem['user'] not in spammers, spammers_view) # Extract and sum teamname/score pairs from the event data. | 'ExtractAndSumScore' >> ExtractAndSumScore('team') # [END filter_and_calc] | 'TeamScoresDict' >> beam.ParDo(TeamScoresDict()) | 'WriteTeamScoreSums' >> WriteToBigQuery( args.table_name + '_teams', args.dataset, { 'team': 'STRING', 'total_score': 'INTEGER', 'window_start': 'STRING', 'processing_time': 'STRING', }, options.view_as(GoogleCloudOptions).project)) # [START session_calc] # Detect user sessions-- that is, a burst of activity separated by a gap # from further activity. Find and record the mean session lengths. # This information could help the game designers track the changing user # engagement as their set of game changes. ( # pylint: disable=expression-not-assigned user_events | 'WindowIntoSessions' >> beam.WindowInto( beam.window.Sessions(session_gap), timestamp_combiner=beam.window.TimestampCombiner.OUTPUT_AT_EOW) # For this use, we care only about the existence of the session, not any # particular information aggregated over it, so we can just group by key # and assign a "dummy value" of None. | beam.CombinePerKey(lambda _: None) # Get the duration of the session | 'UserSessionActivity' >> beam.ParDo(UserSessionActivity()) # [END session_calc] # [START rewindow] # Re-window to process groups of session sums according to when the # sessions complete | 'WindowToExtractSessionMean' >> beam.WindowInto( beam.window.FixedWindows(user_activity_window_duration)) # Find the mean session duration in each window | beam.CombineGlobally( beam.combiners.MeanCombineFn()).without_defaults() | 'FormatAvgSessionLength' >> beam.Map(lambda elem: {'mean_duration': float(elem)}) | 'WriteAvgSessionLength' >> WriteToBigQuery( args.table_name + '_sessions', args.dataset, { 'mean_duration': 'FLOAT', }, options.view_as(GoogleCloudOptions).project)) # [END rewindow] if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/game/game_stats_it_test.py0000644000175100017510000001311015206113032026444 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for the game stats example. Code: beam/sdks/python/apache_beam/examples/complete/game/game_stats.py Usage: pytest --test-pipeline-options=" \ --runner=TestDataflowRunner \ --project=... \ --region=... \ --staging_location=gs://... \ --temp_location=gs://... \ --output=gs://... \ --sdk_location=... \ """ # pytype: skip-file import logging import time import unittest import uuid import pytest from hamcrest.core.core.allof import all_of from apache_beam.examples.complete.game import game_stats from apache_beam.io.gcp.tests import utils from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryMatcher from apache_beam.runners.runner import PipelineState from apache_beam.testing import test_utils from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline class GameStatsIT(unittest.TestCase): # Input events containing user, team, score, processing time, window start. INPUT_EVENT = 'user1,teamA,10,%d,2015-11-02 09:09:28.224' INPUT_TOPIC = 'game_stats_it_input_topic' INPUT_SUB = 'game_stats_it_input_subscription' # SHA-1 hash generated from sorted rows reading from BigQuery table DEFAULT_EXPECTED_CHECKSUM = '5288ccaab77d347c8460d77c15a0db234ef5eb4f' OUTPUT_DATASET = 'game_stats_it_dataset' OUTPUT_TABLE_SESSIONS = 'game_stats_sessions' OUTPUT_TABLE_TEAMS = 'game_stats_teams' DEFAULT_INPUT_COUNT = 500 WAIT_UNTIL_FINISH_DURATION = 12 * 60 * 1000 # in milliseconds def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.project = self.test_pipeline.get_option('project') _unique_id = str(uuid.uuid4()) # Set up PubSub environment. from google.cloud import pubsub self.pub_client = pubsub.PublisherClient() self.input_topic = self.pub_client.create_topic( name=self.pub_client.topic_path( self.project, self.INPUT_TOPIC + _unique_id)) self.sub_client = pubsub.SubscriberClient() self.input_sub = self.sub_client.create_subscription( name=self.sub_client.subscription_path( self.project, self.INPUT_SUB + _unique_id), topic=self.input_topic.name) # Set up BigQuery environment self.dataset_ref = utils.create_bq_dataset( self.project, self.OUTPUT_DATASET) self._test_timestamp = int(time.time() * 1000) def _inject_pubsub_game_events(self, topic, message_count): """Inject game events as test data to PubSub.""" logging.debug( 'Injecting %d game events to topic %s', message_count, topic.name) for _ in range(message_count): self.pub_client.publish( topic.name, (self.INPUT_EVENT % self._test_timestamp).encode('utf-8')) def _cleanup_pubsub(self): test_utils.cleanup_subscriptions(self.sub_client, [self.input_sub]) test_utils.cleanup_topics(self.pub_client, [self.input_topic]) @pytest.mark.it_postcommit @pytest.mark.examples_postcommit # TODO(https://github.com/apache/beam/issues/21300) This example only works in # Dataflow, remove mark to enable for other runners when fixed @pytest.mark.sickbay_direct @pytest.mark.sickbay_spark @pytest.mark.sickbay_flink def test_game_stats_it(self): state_verifier = PipelineStateMatcher(PipelineState.RUNNING) success_condition = 'mean_duration=300 LIMIT 1' sessions_query = ( 'SELECT mean_duration FROM `%s.%s.%s` ' 'WHERE %s' % ( self.project, self.dataset_ref.dataset_id, self.OUTPUT_TABLE_SESSIONS, success_condition)) bq_sessions_verifier = BigqueryMatcher( self.project, sessions_query, self.DEFAULT_EXPECTED_CHECKSUM) # TODO(mariagh): Add teams table verifier once game_stats.py is fixed. extra_opts = { 'subscription': self.input_sub.name, 'dataset': self.dataset_ref.dataset_id, 'topic': self.input_topic.name, 'fixed_window_duration': 1, 'user_activity_window_duration': 1, 'wait_until_finish_duration': self.WAIT_UNTIL_FINISH_DURATION, 'on_success_matcher': all_of(state_verifier, bq_sessions_verifier) } # Register cleanup before pipeline execution. # Note that actual execution happens in reverse order. self.addCleanup(self._cleanup_pubsub) self.addCleanup(utils.delete_bq_dataset, self.project, self.dataset_ref) # Generate input data and inject to PubSub. self._inject_pubsub_game_events(self.input_topic, self.DEFAULT_INPUT_COUNT) # Get pipeline options from command argument: --test-pipeline-options, # and start pipeline job by calling pipeline main function. game_stats.run( self.test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/game/game_stats_test.py0000644000175100017510000000634515206113032025764 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the game_stats example.""" # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam.examples.complete.game import game_stats from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class GameStatsTest(unittest.TestCase): SAMPLE_DATA = [ 'user1_team1,team1,18,1447686663000,2015-11-16 15:11:03.921', 'user1_team1,team1,18,1447690263000,2015-11-16 16:11:03.921', 'user2_team2,team2,2,1447690263000,2015-11-16 16:11:03.955', 'user3_team3,team3,8,1447690263000,2015-11-16 16:11:03.955', 'user4_team3,team3,5,1447690263000,2015-11-16 16:11:03.959', 'user1_team1,team1,14,1447697463000,2015-11-16 18:11:03.955', 'robot1_team1,team1,9000,1447697463000,2015-11-16 18:11:03.955', 'robot2_team2,team2,1,1447697463000,2015-11-16 20:11:03.955', 'robot2_team2,team2,9000,1447697463000,2015-11-16 21:11:03.955', ] def create_data(self, p): return (p | beam.Create(GameStatsTest.SAMPLE_DATA) | beam.ParDo(game_stats.ParseGameEventFn()) | beam.Map(lambda elem:\ beam.window.TimestampedValue(elem, elem['timestamp']))) def test_spammy_users(self): with TestPipeline() as p: result = ( self.create_data(p) | beam.Map(lambda elem: (elem['user'], elem['score'])) | game_stats.CalculateSpammyUsers()) assert_that( result, equal_to([('robot1_team1', 9000), ('robot2_team2', 9001)])) def test_game_stats_sessions(self): session_gap = 5 * 60 user_activity_window_duration = 30 * 60 with TestPipeline() as p: result = ( self.create_data(p) | beam.Map(lambda elem: (elem['user'], elem['score'])) | 'WindowIntoSessions' >> beam.WindowInto( beam.window.Sessions(session_gap), timestamp_combiner=beam.window.TimestampCombiner.OUTPUT_AT_EOW) | beam.CombinePerKey(lambda _: None) | beam.ParDo(game_stats.UserSessionActivity()) | 'WindowToExtractSessionMean' >> beam.WindowInto( beam.window.FixedWindows(user_activity_window_duration)) | beam.CombineGlobally(beam.combiners.MeanCombineFn())\ .without_defaults()) assert_that(result, equal_to([300.0, 300.0, 300.0])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/game/hourly_team_score.py0000644000175100017510000002677715206113032026334 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Second in a series of four pipelines that tell a story in a 'gaming' domain. In addition to the concepts introduced in `user_score`, new concepts include: windowing and element timestamps; use of `Filter`; using standalone DoFns. This pipeline processes data collected from gaming events in batch, building on `user_score` but using fixed windows. It calculates the sum of scores per team, for each window, optionally allowing specification of two timestamps before and after which data is filtered out. This allows a model where late data collected after the intended analysis window can be included, and any late-arriving data prior to the beginning of the analysis window can be removed as well. By using windowing and adding element timestamps, we can do finer-grained analysis than with the `user_score` pipeline. However, our batch processing is high-latency, in that we don't get results from plays at the beginning of the batch's time period until the batch is processed. Optionally include the `--input` argument to specify a batch input file. To indicate a time after which the data should be filtered out, include the `--stop_min` arg. E.g., `--stop_min=2015-10-18-23-59` indicates that any data timestamped after 23:59 PST on 2015-10-18 should not be included in the analysis. To indicate a time before which data should be filtered out, include the `--start_min` arg. If you're using the default input "gs://dataflow-samples/game/gaming_data*.csv", then `--start_min=2015-11-16-16-10 --stop_min=2015-11-17-16-10` are good values. For a description of the usage and options, use -h or --help. To specify a different runner: --runner YOUR_RUNNER NOTE: When specifying a different runner, additional runner-specific options may have to be passed in as well EXAMPLES -------- # DirectRunner python hourly_team_score.py \ --project $PROJECT_ID \ --dataset $BIGQUERY_DATASET # DataflowRunner python hourly_team_score.py \ --project $PROJECT_ID \ --region $REGION_ID \ --dataset $BIGQUERY_DATASET \ --runner DataflowRunner \ --temp_location gs://$BUCKET/user_score/temp """ # pytype: skip-file import argparse import csv import logging import sys import time from datetime import datetime import apache_beam as beam from apache_beam.metrics.metric import Metrics from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions def str2timestamp(s, fmt='%Y-%m-%d-%H-%M'): """Converts a string into a unix timestamp.""" dt = datetime.strptime(s, fmt) epoch = datetime.utcfromtimestamp(0) return (dt - epoch).total_seconds() def timestamp2str(t, fmt='%Y-%m-%d %H:%M:%S.000'): """Converts a unix timestamp into a formatted string.""" return datetime.fromtimestamp(t).strftime(fmt) class ParseGameEventFn(beam.DoFn): """Parses the raw game event info into a Python dictionary. Each event line has the following format: username,teamname,score,timestamp_in_ms,readable_time e.g.: user2_AsparagusPig,AsparagusPig,10,1445230923951,2015-11-02 09:09:28.224 The human-readable time string is not used here. """ def __init__(self): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.DoFn.__init__(self) self.num_parse_errors = Metrics.counter(self.__class__, 'num_parse_errors') def process(self, elem): try: row = list(csv.reader([elem]))[0] yield { 'user': row[0], 'team': row[1], 'score': int(row[2]), 'timestamp': int(row[3]) / 1000.0, } except: # pylint: disable=bare-except # Log and count parse errors self.num_parse_errors.inc() logging.error('Parse error on "%s"', elem) class ExtractAndSumScore(beam.PTransform): """A transform to extract key/score information and sum the scores. The constructor argument `field` determines whether 'team' or 'user' info is extracted. """ def __init__(self, field): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.PTransform.__init__(self) self.field = field def expand(self, pcoll): return ( pcoll | beam.Map(lambda elem: (elem[self.field], elem['score'])) | beam.CombinePerKey(sum)) class TeamScoresDict(beam.DoFn): """Formats the data into a dictionary of BigQuery columns with their values Receives a (team, score) pair, extracts the window start timestamp, and formats everything together into a dictionary. The dictionary is in the format {'bigquery_column': value} """ def process(self, team_score, window=beam.DoFn.WindowParam): team, score = team_score start = timestamp2str(int(window.start)) yield { 'team': team, 'total_score': score, 'window_start': start, 'processing_time': timestamp2str(int(time.time())) } class WriteToBigQuery(beam.PTransform): """Generate, format, and write BigQuery table row information.""" def __init__(self, table_name, dataset, schema, project): """Initializes the transform. Args: table_name: Name of the BigQuery table to use. dataset: Name of the dataset to use. schema: Dictionary in the format {'column_name': 'bigquery_type'} project: Name of the Cloud project containing BigQuery table. """ # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.PTransform.__init__(self) self.table_name = table_name self.dataset = dataset self.schema = schema self.project = project def get_schema(self): """Build the output table schema.""" return ', '.join('%s:%s' % (col, self.schema[col]) for col in self.schema) def expand(self, pcoll): return ( pcoll | 'ConvertToRow' >> beam.Map(lambda elem: {col: elem[col] for col in self.schema}) | beam.io.WriteToBigQuery( self.table_name, self.dataset, self.project, self.get_schema())) # [START main] class HourlyTeamScore(beam.PTransform): def __init__(self, start_min, stop_min, window_duration): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.PTransform.__init__(self) self.start_timestamp = str2timestamp(start_min) self.stop_timestamp = str2timestamp(stop_min) self.window_duration_in_seconds = window_duration * 60 def expand(self, pcoll): return ( pcoll | 'ParseGameEventFn' >> beam.ParDo(ParseGameEventFn()) # Filter out data before and after the given times so that it is not # included in the calculations. As we collect data in batches (say, by # day), the batch for the day that we want to analyze could potentially # include some late-arriving data from the previous day. If so, we want # to weed it out. Similarly, if we include data from the following day # (to scoop up late-arriving events from the day we're analyzing), we # need to weed out events that fall after the time period we want to # analyze. # [START filter_by_time_range] | 'FilterStartTime' >> beam.Filter(lambda elem: elem['timestamp'] > self.start_timestamp) | 'FilterEndTime' >> beam.Filter(lambda elem: elem['timestamp'] < self.stop_timestamp) # [END filter_by_time_range] # [START add_timestamp_and_window] # Add an element timestamp based on the event log, and apply fixed # windowing. | 'AddEventTimestamps' >> beam.Map( lambda elem: beam.window.TimestampedValue(elem, elem['timestamp'])) | 'FixedWindowsTeam' >> beam.WindowInto( beam.window.FixedWindows(self.window_duration_in_seconds)) # [END add_timestamp_and_window] # Extract and sum teamname/score pairs from the event data. | 'ExtractAndSumScore' >> ExtractAndSumScore('team')) def run(argv=None, save_main_session=True): """Main entry point; defines and runs the hourly_team_score pipeline.""" parser = argparse.ArgumentParser() # The default maps to two large Google Cloud Storage files (each ~12GB) # holding two subsequent day's worth (roughly) of data. parser.add_argument( '--input', type=str, default='gs://apache-beam-samples/game/gaming_data*.csv', help='Path to the data file(s) containing game data.') parser.add_argument( '--dataset', type=str, required=True, help='BigQuery Dataset to write tables to. ' 'Must already exist.') parser.add_argument( '--table_name', default='leader_board', help='The BigQuery table name. Should not already exist.') parser.add_argument( '--window_duration', type=int, default=60, help='Numeric value of fixed window duration, in minutes') parser.add_argument( '--start_min', type=str, default='1970-01-01-00-00', help='String representation of the first minute after ' 'which to generate results in the format: ' 'yyyy-MM-dd-HH-mm. Any input data timestamped ' 'prior to that minute won\'t be included in the ' 'sums.') parser.add_argument( '--stop_min', type=str, default='2100-01-01-00-00', help='String representation of the first minute for ' 'which to generate results in the format: ' 'yyyy-MM-dd-HH-mm. Any input data timestamped ' 'after to that minute won\'t be included in the ' 'sums.') args, pipeline_args = parser.parse_known_args(argv) options = PipelineOptions(pipeline_args) # We also require the --project option to access --dataset if options.view_as(GoogleCloudOptions).project is None: parser.print_usage() print(sys.argv[0] + ': error: argument --project is required') sys.exit(1) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). options.view_as(SetupOptions).save_main_session = save_main_session with beam.Pipeline(options=options) as p: ( # pylint: disable=expression-not-assigned p | 'ReadInputText' >> beam.io.ReadFromText(args.input) | 'HourlyTeamScore' >> HourlyTeamScore( args.start_min, args.stop_min, args.window_duration) | 'TeamScoresDict' >> beam.ParDo(TeamScoresDict()) | 'WriteTeamScoreSums' >> WriteToBigQuery( args.table_name, args.dataset, { 'team': 'STRING', 'total_score': 'INTEGER', 'window_start': 'STRING', }, options.view_as(GoogleCloudOptions).project)) # [END main] if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/game/hourly_team_score_it_test.py0000644000175100017510000001101415206113032030041 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for the hourly team score example. Code: beam/sdks/python/apache_beam/examples/complete/game/hourly_team_score.py Usage: pytest --test-pipeline-options=" \ --runner=TestDataflowRunner \ --project=... \ --region=... \ --staging_location=gs://... \ --temp_location=gs://... \ --output=gs://... \ --sdk_location=... \ """ # pytype: skip-file import logging import unittest import pytest from hamcrest.core.core.allof import all_of from apache_beam.examples.complete.game import hourly_team_score from apache_beam.io.gcp.tests import utils from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryMatcher from apache_beam.runners.runner import PipelineState from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline class HourlyTeamScoreIT(unittest.TestCase): DEFAULT_INPUT_FILE = 'gs://dataflow-samples/game/gaming_data*' # SHA-1 hash generated from sorted rows reading from BigQuery table DEFAULT_EXPECTED_CHECKSUM = '4fa761fb5c3341ec573d5d12c6ab75e3b2957a25' OUTPUT_DATASET = 'hourly_team_score_it_dataset' OUTPUT_TABLE = 'leader_board' def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.project = self.test_pipeline.get_option('project') # Set up BigQuery environment self.dataset_ref = utils.create_bq_dataset( self.project, self.OUTPUT_DATASET) @pytest.mark.it_postcommit def test_hourly_team_score_it(self): state_verifier = PipelineStateMatcher(PipelineState.DONE) query = ( 'SELECT COUNT(*) FROM `%s.%s.%s`' % (self.project, self.dataset_ref.dataset_id, self.OUTPUT_TABLE)) bigquery_verifier = BigqueryMatcher( self.project, query, self.DEFAULT_EXPECTED_CHECKSUM) extra_opts = { 'input': self.DEFAULT_INPUT_FILE, 'dataset': self.dataset_ref.dataset_id, 'window_duration': 1, 'on_success_matcher': all_of(state_verifier, bigquery_verifier) } # Register clean up before pipeline execution # Note that actual execution happens in reverse order. self.addCleanup(utils.delete_bq_dataset, self.project, self.dataset_ref) # Get pipeline options from command argument: --test-pipeline-options, # and start pipeline job by calling pipeline main function. hourly_team_score.run( self.test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) @pytest.mark.no_xdist @pytest.mark.examples_postcommit def test_hourly_team_score_output_checksum_on_small_input(self): # Small dataset to prevent Out of Memory when running in local runners INPUT_FILE = 'gs://apache-beam-samples/game/small/gaming_data.csv' EXPECTED_CHECKSUM = '91143e81622aa391eb62eaa3f3a5123401edb07d' state_verifier = PipelineStateMatcher(PipelineState.DONE) query = ( 'SELECT COUNT(*) FROM `%s.%s.%s`' % (self.project, self.dataset_ref.dataset_id, self.OUTPUT_TABLE)) bigquery_verifier = BigqueryMatcher(self.project, query, EXPECTED_CHECKSUM) extra_opts = { 'input': INPUT_FILE, 'dataset': self.dataset_ref.dataset_id, 'window_duration': 1, 'on_success_matcher': all_of(state_verifier, bigquery_verifier) } # Register clean up before pipeline execution # Note that actual execution happens in reverse order. self.addCleanup(utils.delete_bq_dataset, self.project, self.dataset_ref) # Get pipeline options from command argument: --test-pipeline-options, # and start pipeline job by calling pipeline main function. hourly_team_score.run( self.test_pipeline.get_full_options_as_args(**extra_opts)) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/game/hourly_team_score_test.py0000644000175100017510000000407315206113032027354 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the user_score example.""" # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam.examples.complete.game import hourly_team_score from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class HourlyTeamScoreTest(unittest.TestCase): SAMPLE_DATA = [ 'user1_team1,team1,18,1447686663000,2015-11-16 15:11:03.921', 'user1_team1,team1,18,1447690263000,2015-11-16 16:11:03.921', 'user2_team2,team2,2,1447690263000,2015-11-16 16:11:03.955', 'user3_team3,team3,8,1447690263000,2015-11-16 16:11:03.955', 'user4_team3,team3,5,1447690263000,2015-11-16 16:11:03.959', 'user1_team1,team1,14,1447697463000,2015-11-16 18:11:03.955', ] def test_hourly_team_score(self): with TestPipeline() as p: result = ( p | beam.Create(HourlyTeamScoreTest.SAMPLE_DATA) | hourly_team_score.HourlyTeamScore( start_min='2015-11-16-15-20', stop_min='2015-11-16-17-20', window_duration=60)) assert_that( result, equal_to([('team1', 18), ('team2', 2), ('team3', 13)])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/game/leader_board.py0000644000175100017510000003305315206113032025175 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Third in a series of four pipelines that tell a story in a 'gaming' domain. Concepts include: processing unbounded data using fixed windows; use of custom timestamps and event-time processing; generation of early/speculative results; using AccumulationMode.ACCUMULATING to do cumulative processing of late-arriving data. This pipeline processes an unbounded stream of 'game events'. The calculation of the team scores uses fixed windowing based on event time (the time of the game play event), not processing time (the time that an event is processed by the pipeline). The pipeline calculates the sum of scores per team, for each window. By default, the team scores are calculated using one-hour windows. In contrast-- to demo another windowing option-- the user scores are calculated using a global window, which periodically (every ten minutes) emits cumulative user score sums. In contrast to the previous pipelines in the series, which used static, finite input data, here we're using an unbounded data source, which lets us provide speculative results, and allows handling of late data, at much lower latency. We can use the early/speculative results to keep a 'leaderboard' updated in near-realtime. Our handling of late data lets us generate correct results, e.g. for 'team prizes'. We're now outputting window results as they're calculated, giving us much lower latency than with the previous batch examples. Run injector.Injector to generate pubsub data for this pipeline. The Injector documentation provides more detail on how to do this. The injector is currently implemented in Java only, it can be used from the Java SDK. The PubSub topic you specify should be the same topic to which the Injector is publishing. To run the Java injector: /examples/java$ mvn compile exec:java \ -Dexec.mainClass=org.apache.beam.examples.complete.game.injector.Injector \ -Dexec.args="$PROJECT_ID $PUBSUB_TOPIC none" For a description of the usage and options, use -h or --help. To specify a different runner: --runner YOUR_RUNNER NOTE: When specifying a different runner, additional runner-specific options may have to be passed in as well EXAMPLES -------- # DirectRunner python leader_board.py \ --project $PROJECT_ID \ --topic projects/$PROJECT_ID/topics/$PUBSUB_TOPIC \ --dataset $BIGQUERY_DATASET # DataflowRunner python leader_board.py \ --project $PROJECT_ID \ --region $REGION_ID \ --topic projects/$PROJECT_ID/topics/$PUBSUB_TOPIC \ --dataset $BIGQUERY_DATASET \ --runner DataflowRunner \ --temp_location gs://$BUCKET/user_score/temp """ # pytype: skip-file import argparse import csv import logging import sys import time from datetime import datetime import apache_beam as beam from apache_beam.metrics.metric import Metrics from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.transforms import trigger def timestamp2str(t, fmt='%Y-%m-%d %H:%M:%S.000'): """Converts a unix timestamp into a formatted string.""" return datetime.fromtimestamp(t).strftime(fmt) class ParseGameEventFn(beam.DoFn): """Parses the raw game event info into a Python dictionary. Each event line has the following format: username,teamname,score,timestamp_in_ms,readable_time e.g.: user2_AsparagusPig,AsparagusPig,10,1445230923951,2015-11-02 09:09:28.224 The human-readable time string is not used here. """ def __init__(self): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.DoFn.__init__(self) self.num_parse_errors = Metrics.counter(self.__class__, 'num_parse_errors') def process(self, elem): try: row = list(csv.reader([elem]))[0] yield { 'user': row[0], 'team': row[1], 'score': int(row[2]), 'timestamp': int(row[3]) / 1000.0, } except: # pylint: disable=bare-except # Log and count parse errors self.num_parse_errors.inc() logging.error('Parse error on "%s"', elem) class ExtractAndSumScore(beam.PTransform): """A transform to extract key/score information and sum the scores. The constructor argument `field` determines whether 'team' or 'user' info is extracted. """ def __init__(self, field): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.PTransform.__init__(self) self.field = field def expand(self, pcoll): return ( pcoll | beam.Map(lambda elem: (elem[self.field], elem['score'])) | beam.CombinePerKey(sum)) class TeamScoresDict(beam.DoFn): """Formats the data into a dictionary of BigQuery columns with their values Receives a (team, score) pair, extracts the window start timestamp, and formats everything together into a dictionary. The dictionary is in the format {'bigquery_column': value} """ def process(self, team_score, window=beam.DoFn.WindowParam): team, score = team_score start = timestamp2str(int(window.start)) yield { 'team': team, 'total_score': score, 'window_start': start, 'processing_time': timestamp2str(int(time.time())) } class WriteToBigQuery(beam.PTransform): """Generate, format, and write BigQuery table row information.""" def __init__(self, table_name, dataset, schema, project): """Initializes the transform. Args: table_name: Name of the BigQuery table to use. dataset: Name of the dataset to use. schema: Dictionary in the format {'column_name': 'bigquery_type'} project: Name of the Cloud project containing BigQuery table. """ # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.PTransform.__init__(self) self.table_name = table_name self.dataset = dataset self.schema = schema self.project = project def get_schema(self): """Build the output table schema.""" return ', '.join('%s:%s' % (col, self.schema[col]) for col in self.schema) def expand(self, pcoll): return ( pcoll | 'ConvertToRow' >> beam.Map(lambda elem: {col: elem[col] for col in self.schema}) | beam.io.WriteToBigQuery( self.table_name, self.dataset, self.project, self.get_schema())) # [START window_and_trigger] class CalculateTeamScores(beam.PTransform): """Calculates scores for each team within the configured window duration. Extract team/score pairs from the event stream, using hour-long windows by default. """ def __init__(self, team_window_duration, allowed_lateness): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.PTransform.__init__(self) self.team_window_duration = team_window_duration * 60 self.allowed_lateness_seconds = allowed_lateness * 60 def expand(self, pcoll): # NOTE: the behavior does not exactly match the Java example # TODO: allowed_lateness not implemented yet in FixedWindows # TODO: AfterProcessingTime not implemented yet, replace AfterCount return ( pcoll # We will get early (speculative) results as well as cumulative # processing of late data. | 'LeaderboardTeamFixedWindows' >> beam.WindowInto( beam.window.FixedWindows(self.team_window_duration), trigger=trigger.AfterWatermark( trigger.AfterCount(10), trigger.AfterCount(20)), accumulation_mode=trigger.AccumulationMode.ACCUMULATING) # Extract and sum teamname/score pairs from the event data. | 'ExtractAndSumScore' >> ExtractAndSumScore('team')) # [END window_and_trigger] # [START processing_time_trigger] class CalculateUserScores(beam.PTransform): """Extract user/score pairs from the event stream using processing time, via global windowing. Get periodic updates on all users' running scores. """ def __init__(self, allowed_lateness): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.PTransform.__init__(self) self.allowed_lateness_seconds = allowed_lateness * 60 def expand(self, pcoll): # NOTE: the behavior does not exactly match the Java example # TODO: allowed_lateness not implemented yet in FixedWindows # TODO: AfterProcessingTime not implemented yet, replace AfterCount return ( pcoll # Get periodic results every ten events. | 'LeaderboardUserGlobalWindows' >> beam.WindowInto( beam.window.GlobalWindows(), trigger=trigger.Repeatedly(trigger.AfterCount(10)), accumulation_mode=trigger.AccumulationMode.ACCUMULATING) # Extract and sum username/score pairs from the event data. | 'ExtractAndSumScore' >> ExtractAndSumScore('user')) # [END processing_time_trigger] def run(argv=None, save_main_session=True): """Main entry point; defines and runs the hourly_team_score pipeline.""" parser = argparse.ArgumentParser() parser.add_argument('--topic', type=str, help='Pub/Sub topic to read from') parser.add_argument( '--subscription', type=str, help='Pub/Sub subscription to read from') parser.add_argument( '--dataset', type=str, required=True, help='BigQuery Dataset to write tables to. ' 'Must already exist.') parser.add_argument( '--table_name', default='leader_board', help='The BigQuery table name. Should not already exist.') parser.add_argument( '--team_window_duration', type=int, default=60, help='Numeric value of fixed window duration for team ' 'analysis, in minutes') parser.add_argument( '--allowed_lateness', type=int, default=120, help='Numeric value of allowed data lateness, in minutes') args, pipeline_args = parser.parse_known_args(argv) if args.topic is None and args.subscription is None: parser.print_usage() print(sys.argv[0] + ': error: one of --topic or --subscription is required') sys.exit(1) options = PipelineOptions(pipeline_args) # We also require the --project option to access --dataset if options.view_as(GoogleCloudOptions).project is None: parser.print_usage() print(sys.argv[0] + ': error: argument --project is required') sys.exit(1) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). options.view_as(SetupOptions).save_main_session = save_main_session # Enforce that this pipeline is always run in streaming mode options.view_as(StandardOptions).streaming = True with beam.Pipeline(options=options) as p: # Read game events from Pub/Sub using custom timestamps, which are extracted # from the pubsub data elements, and parse the data. # Read from PubSub into a PCollection. if args.subscription: scores = p | 'ReadPubSub' >> beam.io.ReadFromPubSub( subscription=args.subscription) else: scores = p | 'ReadPubSub' >> beam.io.ReadFromPubSub(topic=args.topic) events = ( scores | 'DecodeString' >> beam.Map(lambda b: b.decode('utf-8')) | 'ParseGameEventFn' >> beam.ParDo(ParseGameEventFn()) | 'AddEventTimestamps' >> beam.Map( lambda elem: beam.window.TimestampedValue(elem, elem['timestamp']))) # Get team scores and write the results to BigQuery ( # pylint: disable=expression-not-assigned events | 'CalculateTeamScores' >> CalculateTeamScores( args.team_window_duration, args.allowed_lateness) | 'TeamScoresDict' >> beam.ParDo(TeamScoresDict()) | 'WriteTeamScoreSums' >> WriteToBigQuery( args.table_name + '_teams', args.dataset, { 'team': 'STRING', 'total_score': 'INTEGER', 'window_start': 'STRING', 'processing_time': 'STRING', }, options.view_as(GoogleCloudOptions).project)) def format_user_score_sums(user_score): (user, score) = user_score return {'user': user, 'total_score': score} # Get user scores and write the results to BigQuery ( # pylint: disable=expression-not-assigned events | 'CalculateUserScores' >> CalculateUserScores(args.allowed_lateness) | 'FormatUserScoreSums' >> beam.Map(format_user_score_sums) | 'WriteUserScoreSums' >> WriteToBigQuery( args.table_name + '_users', args.dataset, { 'user': 'STRING', 'total_score': 'INTEGER', }, options.view_as(GoogleCloudOptions).project)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/game/leader_board_it_test.py0000644000175100017510000001354415206113032026733 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for the leader board example. Code: beam/sdks/python/apache_beam/examples/complete/game/leader_board.py Usage: pytest --test-pipeline-options=" \ --runner=TestDataflowRunner \ --project=... \ --region=... \ --staging_location=gs://... \ --temp_location=gs://... \ --output=gs://... \ --sdk_location=... \ """ # pytype: skip-file import logging import time import unittest import uuid import pytest from hamcrest.core.core.allof import all_of from apache_beam.examples.complete.game import leader_board from apache_beam.io.gcp.tests import utils from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryMatcher from apache_beam.runners.runner import PipelineState from apache_beam.testing import test_utils from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline class LeaderBoardIT(unittest.TestCase): # Input event containing user, team, score, processing time, window start. INPUT_EVENT = 'user1,teamA,10,%d,2015-11-02 09:09:28.224' INPUT_TOPIC = 'leader_board_it_input_topic' INPUT_SUB = 'leader_board_it_input_subscription' # SHA-1 hash generated from sorted rows reading from BigQuery table DEFAULT_EXPECTED_CHECKSUM = 'de00231fe6730b972c0ff60a99988438911cda53' OUTPUT_DATASET = 'leader_board_it_dataset' OUTPUT_TABLE_USERS = 'leader_board_users' OUTPUT_TABLE_TEAMS = 'leader_board_teams' DEFAULT_INPUT_COUNT = 500 WAIT_UNTIL_FINISH_DURATION = 10 * 60 * 1000 # in milliseconds def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.project = self.test_pipeline.get_option('project') _unique_id = str(uuid.uuid4()) # Set up PubSub environment. from google.cloud import pubsub self.pub_client = pubsub.PublisherClient() self.input_topic = self.pub_client.create_topic( name=self.pub_client.topic_path( self.project, self.INPUT_TOPIC + _unique_id)) self.sub_client = pubsub.SubscriberClient() self.input_sub = self.sub_client.create_subscription( name=self.sub_client.subscription_path( self.project, self.INPUT_SUB + _unique_id), topic=self.input_topic.name) # Set up BigQuery environment self.dataset_ref = utils.create_bq_dataset( self.project, self.OUTPUT_DATASET) self._test_timestamp = int(time.time() * 1000) def _inject_pubsub_game_events(self, topic, message_count): """Inject game events as test data to PubSub.""" logging.debug( 'Injecting %d game events to topic %s', message_count, topic.name) for _ in range(message_count): self.pub_client.publish( topic.name, (self.INPUT_EVENT % self._test_timestamp).encode('utf-8')) def _cleanup_pubsub(self): test_utils.cleanup_subscriptions(self.sub_client, [self.input_sub]) test_utils.cleanup_topics(self.pub_client, [self.input_topic]) @pytest.mark.it_postcommit @pytest.mark.examples_postcommit # TODO(https://github.com/apache/beam/issues/21300) This example only works # in Dataflow, remove mark to enable for other runners when fixed @pytest.mark.sickbay_direct @pytest.mark.sickbay_spark @pytest.mark.sickbay_flink def test_leader_board_it(self): state_verifier = PipelineStateMatcher(PipelineState.RUNNING) success_condition = 'total_score=5000 LIMIT 1' users_query = ( 'SELECT total_score FROM `%s.%s.%s` ' 'WHERE %s' % ( self.project, self.dataset_ref.dataset_id, self.OUTPUT_TABLE_USERS, success_condition)) bq_users_verifier = BigqueryMatcher( self.project, users_query, self.DEFAULT_EXPECTED_CHECKSUM) teams_query = ( 'SELECT total_score FROM `%s.%s.%s` ' 'WHERE %s' % ( self.project, self.dataset_ref.dataset_id, self.OUTPUT_TABLE_TEAMS, success_condition)) bq_teams_verifier = BigqueryMatcher( self.project, teams_query, self.DEFAULT_EXPECTED_CHECKSUM) extra_opts = { 'allow_unsafe_triggers': True, 'subscription': self.input_sub.name, 'dataset': self.dataset_ref.dataset_id, 'topic': self.input_topic.name, 'team_window_duration': 1, 'wait_until_finish_duration': self.WAIT_UNTIL_FINISH_DURATION, 'on_success_matcher': all_of( state_verifier, bq_users_verifier, bq_teams_verifier) } # Register cleanup before pipeline execution. # Note that actual execution happens in reverse order. self.addCleanup(self._cleanup_pubsub) self.addCleanup(utils.delete_bq_dataset, self.project, self.dataset_ref) # Generate input data and inject to PubSub. self._inject_pubsub_game_events(self.input_topic, self.DEFAULT_INPUT_COUNT) # Get pipeline options from command argument: --test-pipeline-options, # and start pipeline job by calling pipeline main function. leader_board.run( self.test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/game/leader_board_test.py0000644000175100017510000000522715206113032026236 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the leader_board example.""" # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam.examples.complete.game import leader_board from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class LeaderBoardTest(unittest.TestCase): SAMPLE_DATA = [ 'user1_team1,team1,18,1447686663000,2015-11-16 15:11:03.921', 'user1_team1,team1,18,1447690263000,2015-11-16 16:11:03.921', 'user2_team2,team2,2,1447690263000,2015-11-16 16:11:03.955', 'user3_team3,team3,8,1447690263000,2015-11-16 16:11:03.955', 'user4_team3,team3,5,1447690263000,2015-11-16 16:11:03.959', 'user1_team1,team1,14,1447697463000,2015-11-16 18:11:03.955', ] def create_data(self, p): return (p | beam.Create(LeaderBoardTest.SAMPLE_DATA) | beam.ParDo(leader_board.ParseGameEventFn()) | beam.Map(lambda elem:\ beam.window.TimestampedValue(elem, elem['timestamp']))) def test_leader_board_teams(self): with TestPipeline() as p: result = ( self.create_data(p) | leader_board.CalculateTeamScores( team_window_duration=60, allowed_lateness=120)) assert_that( result, equal_to([('team1', 14), ('team1', 18), ('team1', 18), ('team2', 2), ('team3', 13)])) def test_leader_board_users(self): test_options = PipelineOptions(flags=['--allow_unsafe_triggers']) with TestPipeline(options=test_options) as p: result = ( self.create_data(p) | leader_board.CalculateUserScores(allowed_lateness=120)) assert_that(result, equal_to([])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/game/user_score.py0000644000175100017510000001423015206113032024737 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """First in a series of four pipelines that tell a story in a 'gaming' domain. Concepts: batch processing; reading input from Google Cloud Storage or a from a local text file, and writing output to a text file; using standalone DoFns; use of the CombinePerKey transform. In this gaming scenario, many users play, as members of different teams, over the course of a day, and their actions are logged for processing. Some of the logged game events may be late-arriving, if users play on mobile devices and go transiently offline for a period of time. This pipeline does batch processing of data collected from gaming events. It calculates the sum of scores per user, over an entire batch of gaming data (collected, say, for each day). The batch processing will not include any late data that arrives after the day's cutoff point. For a description of the usage and options, use -h or --help. To specify a different runner: --runner YOUR_RUNNER NOTE: When specifying a different runner, additional runner-specific options may have to be passed in as well EXAMPLES -------- # DirectRunner python user_score.py \ --output /local/path/user_score/output # DataflowRunner python user_score.py \ --output gs://$BUCKET/user_score/output \ --runner DataflowRunner \ --project $PROJECT_ID \ --region $REGION_ID \ --temp_location gs://$BUCKET/user_score/temp """ # pytype: skip-file # beam-playground: # name: UserScore # description: batch processing; reading input from Google Cloud Storage or a # from a local text file, and writing output to a text file; using # standalone DoFns; use of the CombinePerKey transform. # multifile: false # pipeline_options: --output output.txt # context_line: 174 # categories: # - Batch # - Combiners # - Options # complexity: ADVANCED # tags: # - batch # - combine # - io # - strings import argparse import csv import logging import apache_beam as beam from apache_beam.metrics.metric import Metrics from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions class ParseGameEventFn(beam.DoFn): """Parses the raw game event info into a Python dictionary. Each event line has the following format: username,teamname,score,timestamp_in_ms,readable_time e.g.: user2_AsparagusPig,AsparagusPig,10,1445230923951,2015-11-02 09:09:28.224 The human-readable time string is not used here. """ def __init__(self): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.DoFn.__init__(self) self.num_parse_errors = Metrics.counter(self.__class__, 'num_parse_errors') def process(self, elem): try: row = list(csv.reader([elem]))[0] yield { 'user': row[0], 'team': row[1], 'score': int(row[2]), 'timestamp': int(row[3]) / 1000.0, } except: # pylint: disable=bare-except # Log and count parse errors self.num_parse_errors.inc() logging.error('Parse error on "%s"', elem) # [START extract_and_sum_score] class ExtractAndSumScore(beam.PTransform): """A transform to extract key/score information and sum the scores. The constructor argument `field` determines whether 'team' or 'user' info is extracted. """ def __init__(self, field): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.PTransform.__init__(self) self.field = field def expand(self, pcoll): return ( pcoll | beam.Map(lambda elem: (elem[self.field], elem['score'])) | beam.CombinePerKey(sum)) # [END extract_and_sum_score] class UserScore(beam.PTransform): def expand(self, pcoll): return ( pcoll | 'ParseGameEventFn' >> beam.ParDo(ParseGameEventFn()) # Extract and sum username/score pairs from the event data. | 'ExtractAndSumScore' >> ExtractAndSumScore('user')) # [START main] def run(argv=None, save_main_session=True): """Main entry point; defines and runs the user_score pipeline.""" parser = argparse.ArgumentParser() # The default maps to two large Google Cloud Storage files (each ~12GB) # holding two subsequent day's worth (roughly) of data. parser.add_argument( '--input', type=str, default='gs://apache-beam-samples/game/small/gaming_data.csv', help='Path to the data file(s) containing game data.') parser.add_argument( '--output', type=str, required=True, help='Path to the output file(s).') args, pipeline_args = parser.parse_known_args(argv) options = PipelineOptions(pipeline_args) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). options.view_as(SetupOptions).save_main_session = save_main_session with beam.Pipeline(options=options) as p: def format_user_score_sums(user_score): (user, score) = user_score return 'user: %s, total_score: %s' % (user, score) ( # pylint: disable=expression-not-assigned p | 'ReadInputText' >> beam.io.ReadFromText(args.input) | 'UserScore' >> UserScore() | 'FormatUserScoreSums' >> beam.Map(format_user_score_sums) | 'WriteUserScoreSums' >> beam.io.WriteToText(args.output)) # [END main] if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/game/user_score_it_test.py0000644000175100017510000001027115206113032026473 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for the user score example. Code: beam/sdks/python/apache_beam/examples/complete/game/user_score.py Usage: pytest --test-pipeline-options=" \ --runner=TestDataflowRunner \ --project=... \ --region=... \ --staging_location=gs://... \ --temp_location=gs://... \ --output=gs://... \ --sdk_location=... \ """ # pytype: skip-file import logging import unittest import uuid import pytest from hamcrest.core.core.allof import all_of from apache_beam.examples.complete.game import user_score from apache_beam.runners.runner import PipelineState from apache_beam.testing.pipeline_verifiers import FileChecksumMatcher from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import delete_files class UserScoreIT(unittest.TestCase): DEFAULT_INPUT_FILE = 'gs://dataflow-samples/game/gaming_data*' DEFAULT_OUTPUT_FILE = \ 'gs://temp-storage-for-end-to-end-tests/py-it-cloud/output' DEFAULT_EXPECTED_CHECKSUM = '9f3bd81669607f0d98ec80ddd477f3277cfba0a2' def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.uuid = str(uuid.uuid4()) self.output = '/'.join([self.DEFAULT_OUTPUT_FILE, self.uuid, 'results']) @pytest.mark.it_postcommit def test_user_score_it(self): state_verifier = PipelineStateMatcher(PipelineState.DONE) arg_sleep_secs = self.test_pipeline.get_option('sleep_secs') sleep_secs = int(arg_sleep_secs) if arg_sleep_secs is not None else None file_verifier = FileChecksumMatcher( self.output + '/*-of-*', self.DEFAULT_EXPECTED_CHECKSUM, sleep_secs) extra_opts = { 'input': self.DEFAULT_INPUT_FILE, 'output': self.output + '/user-score', 'on_success_matcher': all_of(state_verifier, file_verifier) } # Register clean up before pipeline execution self.addCleanup(delete_files, [self.output + '*']) # Get pipeline options from command argument: --test-pipeline-options, # and start pipeline job by calling pipeline main function. user_score.run( self.test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) @pytest.mark.no_xdist @pytest.mark.examples_postcommit def test_userscore_output_checksum_on_small_input(self): # Small dataset to prevent Out of Memory when running in local runners INPUT_FILE = 'gs://apache-beam-samples/game/small/gaming_data.csv' EXPECTED_CHECKSUM = '5b1bc0e8080e3c0f162809ac4c0f49acab23854e' state_verifier = PipelineStateMatcher(PipelineState.DONE) arg_sleep_secs = self.test_pipeline.get_option('sleep_secs') sleep_secs = int(arg_sleep_secs) if arg_sleep_secs is not None else None file_verifier = FileChecksumMatcher( self.output + '/*-of-*', EXPECTED_CHECKSUM, sleep_secs) extra_opts = { 'input': INPUT_FILE, 'output': self.output + '/user-score', 'on_success_matcher': all_of(state_verifier, file_verifier) } # Register clean up before pipeline execution self.addCleanup(delete_files, [self.output + '*']) # Get pipeline options from command argument: --test-pipeline-options, # and start pipeline job by calling pipeline main function. user_score.run(self.test_pipeline.get_full_options_as_args(**extra_opts)) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/game/user_score_test.py0000644000175100017510000000371315206113032026002 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the user_score example.""" # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam.examples.complete.game import user_score from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class UserScoreTest(unittest.TestCase): SAMPLE_DATA = [ 'user1_team1,team1,18,1447686663000,2015-11-16 15:11:03.921', 'user1_team1,team1,18,1447690263000,2015-11-16 16:11:03.921', 'user2_team2,team2,2,1447690263000,2015-11-16 16:11:03.955', 'user3_team3,team3,8,1447690263000,2015-11-16 16:11:03.955', 'user4_team3,team3,5,1447690263000,2015-11-16 16:11:03.959', 'user1_team1,team1,14,1447697463000,2015-11-16 18:11:03.955', ] def test_user_score(self): with TestPipeline() as p: result = ( p | beam.Create(UserScoreTest.SAMPLE_DATA) | user_score.UserScore()) assert_that( result, equal_to([('user1_team1', 50), ('user2_team2', 2), ('user3_team3', 8), ('user4_team3', 5)])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996270.908358 apache_beam-2.74.0/apache_beam/examples/complete/juliaset/0000755000175100017510000000000015206113157023133 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/juliaset/__init__.py0000644000175100017510000000142015206113032025231 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996270.9093583 apache_beam-2.74.0/apache_beam/examples/complete/juliaset/juliaset/0000755000175100017510000000000015206113157024753 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/juliaset/juliaset/__init__.py0000644000175100017510000000142015206113032027051 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/juliaset/juliaset/juliaset.py0000644000175100017510000001044615206113032027142 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A Julia set computing workflow: https://en.wikipedia.org/wiki/Julia_set. We use the quadratic polinomial f(z) = z*z + c, with c = -.62772 +.42193i """ # pytype: skip-file import argparse import apache_beam as beam from apache_beam.io import WriteToText def from_pixel(x, y, n): """Converts a NxN pixel position to a (-1..1, -1..1) complex number.""" return complex(2.0 * x / n - 1.0, 2.0 * y / n - 1.0) def get_julia_set_point_color(element, c, n, max_iterations): """Given an pixel, convert it into a point in our julia set.""" x, y = element z = from_pixel(x, y, n) for i in range(max_iterations): if z.real * z.real + z.imag * z.imag > 2.0: break z = z * z + c return x, y, i # pylint: disable=undefined-loop-variable def generate_julia_set_colors(pipeline, c, n, max_iterations): """Compute julia set coordinates for each point in our set.""" def point_set(n): for x in range(n): for y in range(n): yield (x, y) julia_set_colors = ( pipeline | 'add points' >> beam.Create(point_set(n)) | beam.Map(get_julia_set_point_color, c, n, max_iterations)) return julia_set_colors def generate_julia_set_visualization(data, n, max_iterations): """Generate the pixel matrix for rendering the julia set as an image.""" import numpy as np # pylint: disable=wrong-import-order, wrong-import-position colors = [] for r in range(0, 256, 16): for g in range(0, 256, 16): for b in range(0, 256, 16): colors.append((r, g, b)) xy = np.zeros((n, n, 3), dtype=np.uint8) for x, y, iteration in data: xy[x, y] = colors[iteration * len(colors) // max_iterations] return xy def save_julia_set_visualization(out_file, image_array): """Save the fractal image of our julia set as a png.""" from matplotlib import pyplot as plt # pylint: disable=wrong-import-order, wrong-import-position plt.imsave(out_file, image_array, format='png') def run(argv=None): # pylint: disable=missing-docstring parser = argparse.ArgumentParser() parser.add_argument( '--grid_size', dest='grid_size', default=1000, help='Size of the NxN matrix') parser.add_argument( '--coordinate_output', dest='coordinate_output', required=True, help='Output file to write the color coordinates of the image to.') parser.add_argument( '--image_output', dest='image_output', default=None, help='Output file to write the resulting image to.') known_args, pipeline_args = parser.parse_known_args(argv) with beam.Pipeline(argv=pipeline_args) as p: n = int(known_args.grid_size) coordinates = generate_julia_set_colors(p, complex(-.62772, .42193), n, 100) def x_coord_key(x_y_i): (x, y, i) = x_y_i return (x, (x, y, i)) # Group each coordinate triplet by its x value, then write the coordinates # to the output file with an x-coordinate grouping per line. # pylint: disable=expression-not-assigned ( coordinates | 'x coord key' >> beam.Map(x_coord_key) | 'x coord' >> beam.GroupByKey() | 'format' >> beam.Map( lambda k_coords: ' '.join('(%s, %s, %s)' % c for c in k_coords[1])) | WriteToText(known_args.coordinate_output)) # Optionally render the image and save it to a file. # TODO(silviuc): Add this functionality. # if p.options.image_output is not None: # julia_set_image = generate_julia_set_visualization( # file_with_coordinates, n, 100) # save_julia_set_visualization(p.options.image_output, julia_set_image) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/juliaset/juliaset/juliaset_test.py0000644000175100017510000000573115206113032030202 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the juliaset example.""" # pytype: skip-file import logging import os import re import tempfile import unittest import pytest from apache_beam.examples.complete.juliaset.juliaset import juliaset from apache_beam.testing.util import open_shards @pytest.mark.examples_postcommit class JuliaSetTest(unittest.TestCase): def setUp(self): self.test_files = {} self.test_files['output_coord_file_name'] = self.generate_temp_file() self.test_files['output_image_file_name'] = self.generate_temp_file() def tearDown(self): for test_file in self.test_files.values(): if os.path.exists(test_file): os.remove(test_file) def generate_temp_file(self): with tempfile.NamedTemporaryFile(delete=False) as temp_file: return temp_file.name def run_example(self, grid_size, image_file_name=None): args = [ '--coordinate_output=%s' % self.test_files['output_coord_file_name'], '--grid_size=%s' % grid_size, ] if image_file_name is not None: args.append('--image_output=%s' % image_file_name) juliaset.run(args) def test_output_file_format(self): grid_size = 5 self.run_example(grid_size) # Parse the results from the file, and ensure it was written in the proper # format. with open_shards(self.test_files['output_coord_file_name'] + '-*-of-*') as result_file: output_lines = result_file.readlines() # Should have a line for each x-coordinate. self.assertEqual(grid_size, len(output_lines)) for line in output_lines: coordinates = re.findall(r'(\(\d+, \d+, \d+\))', line) # Should have 5 coordinates on each line. self.assertTrue(coordinates) self.assertEqual(grid_size, len(coordinates)) @unittest.skip( 'TODO(silviuc): Reactivate the test when --image_output is ' 'supported.') def test_generate_fractal_image(self): temp_image_file = self.test_files['output_image_file_name'] self.run_example(10, image_file_name=temp_image_file) # Ensure that the image was saved properly. self.assertTrue(os.stat(temp_image_file).st_size > 0) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/juliaset/juliaset/juliaset_test_it.py0000644000175100017510000000406115206113032030671 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration test for the juliaset example.""" # pytype: skip-file import logging import os import unittest import uuid import pytest from hamcrest.core.core.allof import all_of from apache_beam.examples.complete.juliaset.juliaset import juliaset from apache_beam.io.filesystems import FileSystems from apache_beam.runners.runner import PipelineState from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline @pytest.mark.it_postcommit class JuliaSetTestIT(unittest.TestCase): GRID_SIZE = 1000 def test_run_example_with_setup_file(self): pipeline = TestPipeline(is_integration_test=True) coordinate_output = FileSystems.join( pipeline.get_option('output'), 'juliaset-{}'.format(str(uuid.uuid4())), 'coordinates.txt') extra_args = { 'coordinate_output': coordinate_output, 'grid_size': self.GRID_SIZE, 'setup_file': os.path.normpath( os.path.join(os.path.dirname(__file__), '..', 'setup.py')), 'on_success_matcher': all_of(PipelineStateMatcher(PipelineState.DONE)), } args = pipeline.get_full_options_as_args(**extra_args) juliaset.run(args) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/juliaset/juliaset_main.py0000644000175100017510000000440615206113032026325 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A Julia set computing workflow: https://en.wikipedia.org/wiki/Julia_set. This example has in the juliaset/ folder all the code needed to execute the workflow. It is organized in this way so that it can be packaged as a Python package and later installed in the VM workers executing the job. The root directory for the example contains just a "driver" script to launch the job and the setup.py file needed to create a package. The advantages for organizing the code is that large projects will naturally evolve beyond just one module and you will have to make sure the additional modules are present in the worker. In Python Dataflow, using the --setup_file option when submitting a job, will trigger creating a source distribution (as if running python setup.py sdist) and then staging the resulting tarball in the staging area. The workers, upon startup, will install the tarball. Below is a complete command line for running the juliaset workflow remotely as an example: python juliaset_main.py \ --job_name juliaset-$USER \ --project YOUR-PROJECT \ --region GCE-REGION \ --runner DataflowRunner \ --setup_file ./setup.py \ --staging_location gs://YOUR-BUCKET/juliaset/staging \ --temp_location gs://YOUR-BUCKET/juliaset/temp \ --coordinate_output gs://YOUR-BUCKET/juliaset/out \ --grid_size 20 """ # pytype: skip-file import logging from apache_beam.examples.complete.juliaset.juliaset import juliaset if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) juliaset.run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/juliaset/setup.py0000644000175100017510000001123115206113032024633 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Setup.py module for the workflow's worker utilities. All the workflow related code is gathered in a package that will be built as a source distribution, staged in the staging area for the workflow being run and then installed in the workers when they start running. This behavior is triggered by specifying the --setup_file command line option when running the workflow for remote execution. """ # pytype: skip-file import subprocess import setuptools from setuptools.command.build import build as _build # isort:skip # This class handles the pip install mechanism. class build(_build): # pylint: disable=invalid-name """A build command class that will be invoked during package install. The package built using the current setup.py will be staged and later installed in the worker using `pip install package'. This class will be instantiated during install for this specific scenario and will trigger running the custom commands specified. """ sub_commands = _build.sub_commands + [('CustomCommands', None)] # Some custom command to run during setup. The command is not essential for this # workflow. It is used here as an example. Each command will spawn a child # process. Typically, these commands will include steps to install non-Python # packages. For instance, to install a C++-based library libjpeg62 the following # two commands will have to be added: # # ['apt-get', 'update'], # ['apt-get', '--assume-yes', 'install', 'libjpeg62'], # # First, note that there is no need to use the sudo command because the setup # script runs with appropriate access. # Second, if apt-get tool is used then the first command needs to be 'apt-get # update' so the tool refreshes itself and initializes links to download # repositories. Without this initial step the other apt-get install commands # will fail with package not found errors. Note also --assume-yes option which # shortcuts the interactive confirmation. # # Note that in this example custom commands will run after installing required # packages. If you have a PyPI package that depends on one of the custom # commands, move installation of the dependent package to the list of custom # commands, e.g.: # # ['pip', 'install', 'my_package'], # # TODO(https://github.com/apache/beam/issues/18568): Output from the custom # commands are missing from the logs. The output of custom commands (including # failures) will be logged in the worker-startup log. CUSTOM_COMMANDS = [['echo', 'Custom command worked!']] class CustomCommands(setuptools.Command): """A setuptools Command class able to run arbitrary commands.""" def initialize_options(self): pass def finalize_options(self): pass def RunCustomCommand(self, command_list): print('Running command: %s' % command_list) p = subprocess.Popen( command_list, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) # Can use communicate(input='y\n'.encode()) if the command run requires # some confirmation. stdout_data, _ = p.communicate() print('Command output: %s' % stdout_data) if p.returncode != 0: raise RuntimeError( 'Command %s failed: exit code: %s' % (command_list, p.returncode)) def run(self): for command in CUSTOM_COMMANDS: self.RunCustomCommand(command) # Configure the required packages and scripts to install. # Note that the Python Dataflow containers come with numpy already installed # so this dependency will not trigger anything to be installed unless a version # restriction is specified. REQUIRED_PACKAGES = [ 'numpy', ] setuptools.setup( name='juliaset', version='0.0.1', description='Julia set workflow package.', install_requires=REQUIRED_PACKAGES, packages=setuptools.find_packages(), cmdclass={ # Command class instantiated and run during pip install scenarios. 'build': build, 'CustomCommands': CustomCommands, }) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/tfidf.py0000644000175100017510000002110515206113032022750 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A TF-IDF workflow (term frequency - inverse document frequency). For an explanation of the TF-IDF algorithm see the following link: http://en.wikipedia.org/wiki/Tf-idf """ # pytype: skip-file import argparse import math import re import apache_beam as beam from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.io.filesystems import FileSystems from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.pvalue import AsSingleton def read_documents(pipeline, uris): """Read the documents at the provided uris and returns (uri, line) pairs.""" pcolls = [] for uri in uris: pcolls.append( pipeline | 'Read: %s' % uri >> ReadFromText(uri) | 'WithKey: %s' % uri >> beam.Map(lambda v, uri: (uri, v), uri)) return pcolls | 'FlattenReadPColls' >> beam.Flatten() class TfIdf(beam.PTransform): """A transform containing a basic TF-IDF pipeline. The input consists of KV objects where the key is the document's URI and the value is a piece of the document's content. The output is mapping from terms to scores for each document URI. """ def expand(self, uri_to_content): # Compute the total number of documents, and prepare a singleton # PCollection to use as side input. total_documents = ( uri_to_content | 'GetUris 1' >> beam.Keys() | 'GetUniqueUris' >> beam.Distinct() | 'CountUris' >> beam.combiners.Count.Globally()) # Create a collection of pairs mapping a URI to each of the words # in the document associated with that that URI. def split_into_words(uri_line): (uri, line) = uri_line return [(uri, w.lower()) for w in re.findall(r'[A-Za-z\']+', line)] uri_to_words = ( uri_to_content | 'SplitWords' >> beam.FlatMap(split_into_words)) # Compute a mapping from each word to the total number of documents # in which it appears. word_to_doc_count = ( uri_to_words | 'GetUniqueWordsPerDoc' >> beam.Distinct() | 'GetWords' >> beam.Values() | 'CountDocsPerWord' >> beam.combiners.Count.PerElement()) # Compute a mapping from each URI to the total number of words in the # document associated with that URI. uri_to_word_total = ( uri_to_words | 'GetUris 2' >> beam.Keys() | 'CountWordsInDoc' >> beam.combiners.Count.PerElement()) # Count, for each (URI, word) pair, the number of occurrences of that word # in the document associated with the URI. uri_and_word_to_count = ( uri_to_words | 'CountWord-DocPairs' >> beam.combiners.Count.PerElement()) # Adjust the above collection to a mapping from (URI, word) pairs to counts # into an isomorphic mapping from URI to (word, count) pairs, to prepare # for a join by the URI key. def shift_keys(uri_word_count): return (uri_word_count[0][0], (uri_word_count[0][1], uri_word_count[1])) uri_to_word_and_count = ( uri_and_word_to_count | 'ShiftKeys' >> beam.Map(shift_keys)) # Perform a CoGroupByKey (a sort of pre-join) on the prepared # uri_to_word_total and uri_to_word_and_count tagged by 'word totals' and # 'word counts' strings. This yields a mapping from URI to a dictionary # that maps the above mentioned tag strings to an iterable containing the # word total for that URI and word and count respectively. # # A diagram (in which '[]' just means 'iterable'): # # URI: {'word totals': [count], # Total words within this URI's document. # 'word counts': [(word, count), # Counts of specific words # (word, count), # within this URI's document. # ... ]} uri_to_word_and_count_and_total = ({ 'word totals': uri_to_word_total, 'word counts': uri_to_word_and_count } | 'CoGroupByUri' >> beam.CoGroupByKey()) # Compute a mapping from each word to a (URI, term frequency) pair for each # URI. A word's term frequency for a document is simply the number of times # that word occurs in the document divided by the total number of words in # the document. def compute_term_frequency(uri_count_and_total): (uri, count_and_total) = uri_count_and_total word_and_count = count_and_total['word counts'] # We have an iterable for one element that we want extracted. [word_total] = count_and_total['word totals'] for word, count in word_and_count: yield word, (uri, float(count) / word_total) word_to_uri_and_tf = ( uri_to_word_and_count_and_total | 'ComputeTermFrequencies' >> beam.FlatMap(compute_term_frequency)) # Compute a mapping from each word to its document frequency. # A word's document frequency in a corpus is the number of # documents in which the word appears divided by the total # number of documents in the corpus. # # This calculation uses a side input, a Dataflow-computed auxiliary value # presented to each invocation of our MapFn lambda. The second argument to # the function (called total---note that the first argument is a tuple) # receives the value we listed after the lambda in Map(). Additional side # inputs (and ordinary Python values, too) can be provided to MapFns and # DoFns in this way. def div_word_count_by_total(word_count, total): (word, count) = word_count return (word, float(count) / total) word_to_df = ( word_to_doc_count | 'ComputeDocFrequencies' >> beam.Map( div_word_count_by_total, AsSingleton(total_documents))) # Join the term frequency and document frequency collections, # each keyed on the word. word_to_uri_and_tf_and_df = ({ 'tf': word_to_uri_and_tf, 'df': word_to_df } | 'CoGroupWordsByTf-df' >> beam.CoGroupByKey()) # Compute a mapping from each word to a (URI, TF-IDF) score for each URI. # There are a variety of definitions of TF-IDF # ("term frequency - inverse document frequency") score; here we use a # basic version that is the term frequency divided by the log of the # document frequency. def compute_tf_idf(word_tf_and_df): (word, tf_and_df) = word_tf_and_df [docf] = tf_and_df['df'] for uri, tf in tf_and_df['tf']: yield word, (uri, tf * math.log(1 / docf)) word_to_uri_and_tfidf = ( word_to_uri_and_tf_and_df | 'ComputeTf-idf' >> beam.FlatMap(compute_tf_idf)) return word_to_uri_and_tfidf def run(argv=None, save_main_session=True): """Main entry point; defines and runs the tfidf pipeline.""" parser = argparse.ArgumentParser() parser.add_argument('--uris', required=True, help='URIs to process.') parser.add_argument( '--output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session with beam.Pipeline(options=pipeline_options) as p: # Read documents specified by the uris command line option. metadata_list = FileSystems.match([known_args.uris])[0].metadata_list uris = [metadata.path for metadata in metadata_list] pcoll = read_documents(p, uris) # Compute TF-IDF information for each word. output = pcoll | TfIdf() # Write the output using a "Write" transform that has side effects. # pylint: disable=expression-not-assigned output | 'write' >> WriteToText(known_args.output) # Execute the pipeline and wait until it is completed. if __name__ == '__main__': run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/tfidf_it_test.py0000644000175100017510000000533015206113032024505 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for TF-IDF example.""" # pytype: skip-file import logging import re import unittest import uuid import pytest from apache_beam.examples.complete import tfidf from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import create_file from apache_beam.testing.test_utils import read_files_from_pattern EXPECTED_RESULTS = set([ ('ghi', '1.txt', 0.3662040962227032), ('abc', '1.txt', 0.0), ('abc', '3.txt', 0.0), ('abc', '2.txt', 0.0), ('def', '1.txt', 0.13515503603605478), ('def', '2.txt', 0.2027325540540822) ]) EXPECTED_LINE_RE = r'\(u?\'([a-z]*)\', \(\'.*([0-9]\.txt)\', (.*)\)\)' class TfIdfIT(unittest.TestCase): @pytest.mark.examples_postcommit @pytest.mark.sickbay_flink def test_basics(self): test_pipeline = TestPipeline(is_integration_test=True) # Setup the files with expected content. temp_location = test_pipeline.get_option('temp_location') input_folder = '/'.join([temp_location, str(uuid.uuid4())]) create_file('/'.join([input_folder, '1.txt']), 'abc def ghi') create_file('/'.join([input_folder, '2.txt']), 'abc def') create_file('/'.join([input_folder, '3.txt']), 'abc') output = '/'.join([temp_location, str(uuid.uuid4()), 'result']) extra_opts = {'uris': '%s/**' % input_folder, 'output': output} tfidf.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) # Parse result file and compare. results = [] lines = read_files_from_pattern('%s*' % output).splitlines() for line in lines: match = re.search(EXPECTED_LINE_RE, line) logging.info('Result line: %s', line) if match is not None: results.append((match.group(1), match.group(2), float(match.group(3)))) logging.info('Computed results: %s', set(results)) self.assertEqual(set(results), EXPECTED_RESULTS) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/tfidf_test.py0000644000175100017510000000414115206113032024010 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the TF-IDF example.""" # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam.examples.complete import tfidf from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to EXPECTED_RESULTS = set([ ('ghi', '1.txt', 0.3662040962227032), ('abc', '1.txt', 0.0), ('abc', '3.txt', 0.0), ('abc', '2.txt', 0.0), ('def', '1.txt', 0.13515503603605478), ('def', '2.txt', 0.2027325540540822) ]) class TfIdfTest(unittest.TestCase): def test_tfidf_transform(self): with TestPipeline() as p: def re_key(word_uri_tfidf): (word, (uri, tfidf)) = word_uri_tfidf return (word, uri, tfidf) uri_to_line = p | 'create sample' >> beam.Create( [('1.txt', 'abc def ghi'), ('2.txt', 'abc def'), ('3.txt', 'abc')]) result = (uri_to_line | tfidf.TfIdf() | beam.Map(re_key)) assert_that(result, equal_to(EXPECTED_RESULTS)) # Run the pipeline. Note that the assert_that above adds to the pipeline # a check that the result PCollection contains expected values. # To actually trigger the check the pipeline must be run (e.g. by # exiting the with context). if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/top_wikipedia_sessions.py0000644000175100017510000001333315206113032026436 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """An example that reads Wikipedia edit data and computes strings of edits. An example that reads Wikipedia edit data from Cloud Storage and computes the user with the longest string of edits separated by no more than an hour within each 30 day period. To execute this pipeline locally using the DirectRunner, specify an output prefix on GCS::: --output gs://YOUR_OUTPUT_PREFIX To execute this pipeline using the Google Cloud Dataflow service, specify pipeline configuration in addition to the above::: --job_name NAME_FOR_YOUR_JOB --project YOUR_PROJECT_ID --region GCE_REGION --staging_location gs://YOUR_STAGING_DIRECTORY --temp_location gs://YOUR_TEMPORARY_DIRECTORY --runner DataflowRunner The default input is ``gs://dataflow-samples/wikipedia_edits/*.json`` and can be overridden with --input. """ # pytype: skip-file import argparse import json import logging import apache_beam as beam from apache_beam import combiners from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import Sessions from apache_beam.transforms.window import TimestampedValue ONE_HOUR_IN_SECONDS = 3600 THIRTY_DAYS_IN_SECONDS = 30 * 24 * ONE_HOUR_IN_SECONDS MAX_TIMESTAMP = 0x7fffffffffffffff def extract_user_and_timestamp(element): """Extracts user and timestamp representing a Wikipedia edit.""" table_row = json.loads(element) if 'contributor_username' in table_row: user_name = table_row['contributor_username'] timestamp = table_row['timestamp'] return TimestampedValue(user_name, timestamp) class ComputeSessions(beam.PTransform): """Computes the number of edits in each user session. A session is defined as a string of edits where each is separated from the next by less than an hour. """ def expand(self, pcoll): return ( pcoll | 'ComputeSessionsWindow' >> beam.WindowInto( Sessions(gap_size=ONE_HOUR_IN_SECONDS)) | combiners.Count.PerElement()) class TopPerMonth(beam.PTransform): """Computes the longest session ending in each month.""" def expand(self, pcoll): return ( pcoll | 'TopPerMonthWindow' >> beam.WindowInto( FixedWindows(size=THIRTY_DAYS_IN_SECONDS)) | 'Top' >> combiners.core.CombineGlobally( combiners.TopCombineFn( n=10, key=lambda sessions_count: sessions_count[1])). without_defaults()) def sessions_to_strings(element, window=beam.DoFn.WindowParam): """Adds the session information to be part of the key.""" return (element[0] + ' : ' + str(window), element[1]) def format_output(element, window=beam.DoFn.WindowParam): """Formats a string containing the user, count, and session.""" for session, count in element: yield session + ' : ' + str(count) + ' : ' + str(window) class ComputeTopSessions(beam.PTransform): """Computes the top user sessions for each month.""" def __init__(self, sampling_threshold): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.PTransform.__init__(self) self.sampling_threshold = sampling_threshold def expand(self, pcoll): return ( pcoll | 'ExtractUserAndTimestamp' >> beam.Map(extract_user_and_timestamp) | beam.Filter( lambda x: (abs(hash(x)) <= MAX_TIMESTAMP * self.sampling_threshold)) | ComputeSessions() | 'SessionsToStrings' >> beam.Map(sessions_to_strings) | TopPerMonth() | 'FormatOutput' >> beam.FlatMap(format_output)) def run(argv=None): """Runs the Wikipedia top edits pipeline. Args: argv: Pipeline options as a list of arguments. """ parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', default='gs://dataflow-samples/wikipedia_edits/*.json', help='Input specified as a GCS path containing a BigQuery table exported ' 'as json.') parser.add_argument( '--output', required=True, help='Output file to write results to.') parser.add_argument( '--sampling_threshold', type=float, default=0.1, help='Fraction of entries used for session tracking') known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = True with beam.Pipeline(options=pipeline_options) as p: ( # pylint: disable=expression-not-assigned p | ReadFromText(known_args.input) | ComputeTopSessions(known_args.sampling_threshold) | WriteToText(known_args.output)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/top_wikipedia_sessions_it_test.py0000644000175100017510000000662115206113032030173 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for Top Wikipedia Sessions example.""" # pytype: skip-file import json import logging import unittest import uuid import pytest from apache_beam.examples.complete import top_wikipedia_sessions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import create_file from apache_beam.testing.test_utils import read_files_from_pattern class ComputeTopSessionsIT(unittest.TestCase): EDITS = [ json.dumps({ 'timestamp': 0.0, 'contributor_username': 'user1' }), json.dumps({ 'timestamp': 0.001, 'contributor_username': 'user1' }), json.dumps({ 'timestamp': 0.002, 'contributor_username': 'user1' }), json.dumps({ 'timestamp': 0.0, 'contributor_username': 'user2' }), json.dumps({ 'timestamp': 0.001, 'contributor_username': 'user2' }), json.dumps({ 'timestamp': 3.601, 'contributor_username': 'user2' }), json.dumps({ 'timestamp': 3.602, 'contributor_username': 'user2' }), json.dumps({ 'timestamp': 2 * 3600.0, 'contributor_username': 'user2' }), json.dumps({ 'timestamp': 35 * 24 * 3.600, 'contributor_username': 'user3' }) ] EXPECTED = [ 'user1 : [0.0, 3600.002) : 3 : [0.0, 2592000.0)', 'user2 : [0.0, 3603.602) : 4 : [0.0, 2592000.0)', 'user2 : [7200.0, 10800.0) : 1 : [0.0, 2592000.0)', 'user3 : [3024.0, 6624.0) : 1 : [0.0, 2592000.0)', ] # TODO Enable when fixed this tests for Dataflow runner @pytest.mark.sickbay_dataflow @pytest.mark.no_xdist @pytest.mark.examples_postcommit def test_top_wikipedia_sessions_output_files_on_small_input(self): test_pipeline = TestPipeline(is_integration_test=True) # Setup the files with expected content. OUTPUT_FILE_DIR = \ 'gs://temp-storage-for-end-to-end-tests/py-it-cloud/output' output = '/'.join([OUTPUT_FILE_DIR, str(uuid.uuid4()), 'result']) INPUT_FILE_DIR = \ 'gs://temp-storage-for-end-to-end-tests/py-it-cloud/input' input = '/'.join([INPUT_FILE_DIR, str(uuid.uuid4()), 'input.txt']) create_file(input, '\n'.join(self.EDITS)) extra_opts = {'input': input, 'output': output, 'sampling_threshold': '1.0'} top_wikipedia_sessions.run( test_pipeline.get_full_options_as_args(**extra_opts)) # Load result file and compare. result = read_files_from_pattern('%s*' % output).strip().splitlines() self.assertEqual(self.EXPECTED, sorted(result, key=lambda x: x.split()[0])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/complete/top_wikipedia_sessions_test.py0000644000175100017510000000511115206113032027470 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the top wikipedia sessions example.""" # pytype: skip-file import json import unittest import apache_beam as beam from apache_beam.examples.complete import top_wikipedia_sessions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # TODO: Unit test top_wikipedia_sessions.extract_user_and_timestamp, etc. class ComputeTopSessionsTest(unittest.TestCase): EDITS = [ json.dumps({ 'timestamp': 0.0, 'contributor_username': 'user1' }), json.dumps({ 'timestamp': 0.001, 'contributor_username': 'user1' }), json.dumps({ 'timestamp': 0.002, 'contributor_username': 'user1' }), json.dumps({ 'timestamp': 0.0, 'contributor_username': 'user2' }), json.dumps({ 'timestamp': 0.001, 'contributor_username': 'user2' }), json.dumps({ 'timestamp': 3.601, 'contributor_username': 'user2' }), json.dumps({ 'timestamp': 3.602, 'contributor_username': 'user2' }), json.dumps({ 'timestamp': 2 * 3600.0, 'contributor_username': 'user2' }), json.dumps({ 'timestamp': 35 * 24 * 3.600, 'contributor_username': 'user3' }) ] EXPECTED = [ 'user1 : [0.0, 3600.002) : 3 : [0.0, 2592000.0)', 'user2 : [0.0, 3603.602) : 4 : [0.0, 2592000.0)', 'user2 : [7200.0, 10800.0) : 1 : [0.0, 2592000.0)', 'user3 : [3024.0, 6624.0) : 1 : [0.0, 2592000.0)', ] def test_compute_top_sessions(self): with TestPipeline() as p: edits = p | beam.Create(self.EDITS) result = edits | top_wikipedia_sessions.ComputeTopSessions(1.0) assert_that(result, equal_to(self.EXPECTED)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.0443723 apache_beam-2.74.0/apache_beam/examples/cookbook/0000755000175100017510000000000015206113157021311 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/__init__.py0000644000175100017510000000142015206113032023407 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/bigquery_schema.py0000644000175100017510000001057615206113032025033 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A workflow that writes to a BigQuery table with nested and repeated fields. Demonstrates how to build a bigquery.TableSchema object with nested and repeated fields. Also, shows how to generate data to be written to a BigQuery table with nested and repeated fields. """ # pytype: skip-file import argparse import logging import apache_beam as beam def run(argv=None): """Run the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--output', required=True, help=( 'Output BigQuery table for results specified as: ' 'PROJECT:DATASET.TABLE or DATASET.TABLE.')) known_args, pipeline_args = parser.parse_known_args(argv) with beam.Pipeline(argv=pipeline_args) as p: from apache_beam.io.gcp.internal.clients import bigquery # pylint: disable=wrong-import-order, wrong-import-position table_schema = bigquery.TableSchema() # Fields that use standard types. kind_schema = bigquery.TableFieldSchema() kind_schema.name = 'kind' kind_schema.type = 'string' kind_schema.mode = 'nullable' table_schema.fields.append(kind_schema) full_name_schema = bigquery.TableFieldSchema() full_name_schema.name = 'fullName' full_name_schema.type = 'string' full_name_schema.mode = 'required' table_schema.fields.append(full_name_schema) age_schema = bigquery.TableFieldSchema() age_schema.name = 'age' age_schema.type = 'integer' age_schema.mode = 'nullable' table_schema.fields.append(age_schema) gender_schema = bigquery.TableFieldSchema() gender_schema.name = 'gender' gender_schema.type = 'string' gender_schema.mode = 'nullable' table_schema.fields.append(gender_schema) # A nested field phone_number_schema = bigquery.TableFieldSchema() phone_number_schema.name = 'phoneNumber' phone_number_schema.type = 'record' phone_number_schema.mode = 'nullable' area_code = bigquery.TableFieldSchema() area_code.name = 'areaCode' area_code.type = 'integer' area_code.mode = 'nullable' phone_number_schema.fields.append(area_code) number = bigquery.TableFieldSchema() number.name = 'number' number.type = 'integer' number.mode = 'nullable' phone_number_schema.fields.append(number) table_schema.fields.append(phone_number_schema) # A repeated field. children_schema = bigquery.TableFieldSchema() children_schema.name = 'children' children_schema.type = 'string' children_schema.mode = 'repeated' table_schema.fields.append(children_schema) def create_random_record(record_id): return { 'kind': 'kind' + record_id, 'fullName': 'fullName' + record_id, 'age': int(record_id) * 10, 'gender': 'male', 'phoneNumber': { 'areaCode': int(record_id) * 100, 'number': int(record_id) * 100000 }, 'children': [ 'child' + record_id + '1', 'child' + record_id + '2', 'child' + record_id + '3' ] } # pylint: disable=expression-not-assigned record_ids = p | 'CreateIDs' >> beam.Create(['1', '2', '3', '4', '5']) records = record_ids | 'CreateRecords' >> beam.Map(create_random_record) records | 'write' >> beam.io.WriteToBigQuery( known_args.output, schema=table_schema, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE) # Run the pipeline (all operations are deferred until run() is called). if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/bigquery_side_input.py0000644000175100017510000001036415206113032025731 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A Dataflow job that uses BigQuery sources as a side inputs. Illustrates how to insert side-inputs into transforms in three different forms, as a singleton, as a iterator, and as a list. This workflow generate a set of tuples of the form (groupId, corpus, word) where groupId is a generated identifier for the group and corpus and word are randomly selected from corresponding rows in BQ dataset 'publicdata:samples.shakespeare'. Users should specify the number of groups to form and optionally a corpus and/or a word that should be ignored when forming groups. """ # pytype: skip-file import argparse import logging from random import randrange import apache_beam as beam from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.pvalue import AsList from apache_beam.pvalue import AsSingleton def create_groups(group_ids, corpus, word, ignore_corpus, ignore_word): """Generate groups given the input PCollections.""" def attach_corpus_fn(group, corpus, ignore): selected = None len_corpus = len(corpus) while not selected: c = list(corpus[randrange(0, len_corpus)].values())[0] if c != ignore: selected = c yield (group, selected) def attach_word_fn(group, words, ignore): selected = None len_words = len(words) while not selected: c = list(words[randrange(0, len_words)].values())[0] if c != ignore: selected = c yield group + (selected, ) return ( group_ids | 'attach corpus' >> beam.FlatMap( attach_corpus_fn, AsList(corpus), AsSingleton(ignore_corpus)) | 'attach word' >> beam.FlatMap( attach_word_fn, AsList(word), AsSingleton(ignore_word))) def run(argv=None): """Run the workflow.""" parser = argparse.ArgumentParser() parser.add_argument('--output') parser.add_argument('--ignore_corpus', default='') parser.add_argument('--ignore_word', default='') parser.add_argument('--num_groups') known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = True with beam.Pipeline(options=pipeline_options) as p: group_ids = [] for i in range(0, int(known_args.num_groups)): group_ids.append('id' + str(i)) query_corpus = 'select UNIQUE(corpus) from publicdata:samples.shakespeare' query_word = 'select UNIQUE(word) from publicdata:samples.shakespeare' ignore_corpus = known_args.ignore_corpus ignore_word = known_args.ignore_word pcoll_corpus = p | 'read corpus' >> beam.io.ReadFromBigQuery( query=query_corpus) pcoll_word = p | 'read_words' >> beam.io.ReadFromBigQuery(query=query_word) pcoll_ignore_corpus = p | 'create_ignore_corpus' >> beam.Create( [ignore_corpus]) pcoll_ignore_word = p | 'create_ignore_word' >> beam.Create([ignore_word]) pcoll_group_ids = p | 'create groups' >> beam.Create(group_ids) pcoll_groups = create_groups( pcoll_group_ids, pcoll_corpus, pcoll_word, pcoll_ignore_corpus, pcoll_ignore_word) # pylint:disable=expression-not-assigned pcoll_groups | WriteToText(known_args.output) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/bigquery_side_input_it_test.py0000644000175100017510000000454215206113032027465 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration test for the BigQuery side input example.""" # pytype: skip-file import logging import unittest import uuid import pytest from hamcrest.core.core.allof import all_of from apache_beam.examples.cookbook import bigquery_side_input from apache_beam.runners.runner import PipelineState from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import delete_files class BigQuerySideInputIT(unittest.TestCase): DEFAULT_OUTPUT_FILE = \ 'gs://temp-storage-for-end-to-end-tests/py-it-cloud/output' def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.uuid = str(uuid.uuid4()) self.output = '/'.join([self.DEFAULT_OUTPUT_FILE, self.uuid, 'results']) @pytest.mark.no_xdist @pytest.mark.examples_postcommit @pytest.mark.sickbay_flink def test_bigquery_side_input_it(self): state_verifier = PipelineStateMatcher(PipelineState.DONE) NUM_GROUPS = 3 extra_opts = { 'output': self.output, 'num_groups': str(NUM_GROUPS), 'on_success_matcher': all_of(state_verifier) } # Register clean up before pipeline execution self.addCleanup(delete_files, [self.output + '*']) # Get pipeline options from command argument: --test-pipeline-options, # and start pipeline job by calling pipeline main function. bigquery_side_input.run( self.test_pipeline.get_full_options_as_args(**extra_opts)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/bigquery_side_input_test.py0000644000175100017510000000416615206113032026773 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the BigQuery side input example.""" # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam.examples.cookbook import bigquery_side_input from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class BigQuerySideInputTest(unittest.TestCase): def test_create_groups(self): with TestPipeline() as p: group_ids_pcoll = p | 'CreateGroupIds' >> beam.Create(['A', 'B', 'C']) corpus_pcoll = p | 'CreateCorpus' >> beam.Create([{ 'f': 'corpus1' }, { 'f': 'corpus2' }]) words_pcoll = p | 'CreateWords' >> beam.Create([{ 'f': 'word1' }, { 'f': 'word2' }]) ignore_corpus_pcoll = p | 'CreateIgnoreCorpus' >> beam.Create(['corpus1']) ignore_word_pcoll = p | 'CreateIgnoreWord' >> beam.Create(['word1']) groups = bigquery_side_input.create_groups( group_ids_pcoll, corpus_pcoll, words_pcoll, ignore_corpus_pcoll, ignore_word_pcoll) assert_that( groups, equal_to([('A', 'corpus2', 'word2'), ('B', 'corpus2', 'word2'), ('C', 'corpus2', 'word2')])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/bigquery_tornadoes.py0000644000175100017510000000714715206113032025571 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A workflow using BigQuery sources and sinks. The workflow will read from a table that has the 'month' and 'tornado' fields as part of the table schema (other additional fields are ignored). The 'month' field is a number represented as a string (e.g., '23') and the 'tornado' field is a boolean field. The workflow will compute the number of tornadoes in each month and output the results to a table (created if needed) with the following schema: - month: number - tornado_count: number This example uses the default behavior for BigQuery source and sinks that represents table rows as plain Python dictionaries. """ # pytype: skip-file import argparse import logging import apache_beam as beam def count_tornadoes(input_data): """Workflow computing the number of tornadoes for each month that had one. Args: input_data: a PCollection of dictionaries representing table rows. Each dictionary will have a 'month' and a 'tornado' key as described in the module comment. Returns: A PCollection of dictionaries containing 'month' and 'tornado_count' keys. Months without tornadoes are skipped. """ return ( input_data | 'months with tornadoes' >> beam.FlatMap( lambda row: [(int(row['month']), 1)] if row['tornado'] else []) | 'monthly count' >> beam.CombinePerKey(sum) | 'format' >> beam.Map(lambda k_v: { 'month': k_v[0], 'tornado_count': k_v[1] })) def run(argv=None): parser = argparse.ArgumentParser() parser.add_argument( '--input', default='apache-beam-testing.samples.weather_stations', help=( 'Input BigQuery table to process specified as: ' 'PROJECT:DATASET.TABLE or DATASET.TABLE.')) parser.add_argument( '--output', required=True, help=( 'Output BigQuery table for results specified as: ' 'PROJECT:DATASET.TABLE or DATASET.TABLE.')) parser.add_argument( '--gcs_location', required=False, help=('GCS Location to store files to load ' 'data into Bigquery')) known_args, pipeline_args = parser.parse_known_args(argv) with beam.Pipeline(argv=pipeline_args) as p: # Read the table rows into a PCollection. rows = p | 'read' >> beam.io.ReadFromBigQuery(table=known_args.input) counts = count_tornadoes(rows) # Write the output using a "Write" transform that has side effects. # pylint: disable=expression-not-assigned counts | 'Write' >> beam.io.WriteToBigQuery( known_args.output, schema='month:INTEGER, tornado_count:INTEGER', create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE) # Run the pipeline (all operations are deferred until run() is called). if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/bigquery_tornadoes_it_test.py0000644000175100017510000000530115206113032027312 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for Bigquery tornadoes example.""" # pytype: skip-file import logging import time import unittest import pytest from hamcrest.core.core.allof import all_of from apache_beam.examples.cookbook import bigquery_tornadoes from apache_beam.io.gcp.tests import utils from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryMatcher from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline class BigqueryTornadoesIT(unittest.TestCase): # The default checksum is a SHA-1 hash generated from sorted rows reading # from expected Bigquery table. DEFAULT_CHECKSUM = 'd860e636050c559a16a791aff40d6ad809d4daf0' @pytest.mark.examples_postcommit @pytest.mark.it_postcommit def test_bigquery_tornadoes_it(self): test_pipeline = TestPipeline(is_integration_test=True) # Set extra options to the pipeline for test purpose project = test_pipeline.get_option('project') dataset = 'BigQueryTornadoesIT' table = 'monthly_tornadoes_%s' % int(round(time.time() * 1000)) output_table = '.'.join([dataset, table]) query = 'SELECT month, tornado_count FROM `%s`' % output_table pipeline_verifiers = [ PipelineStateMatcher(), BigqueryMatcher( project=project, query=query, checksum=self.DEFAULT_CHECKSUM) ] extra_opts = { 'output': output_table, 'on_success_matcher': all_of(*pipeline_verifiers) } # Register cleanup before pipeline execution. # Note that actual execution happens in reverse order. self.addCleanup(utils.delete_bq_table, project, dataset, table) # Get pipeline options from command argument: --test-pipeline-options, # and start pipeline job by calling pipeline main function. bigquery_tornadoes.run(test_pipeline.get_full_options_as_args(**extra_opts)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/bigquery_tornadoes_test.py0000644000175100017510000000353715206113032026627 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the BigQuery tornadoes example.""" # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam.examples.cookbook import bigquery_tornadoes from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class BigQueryTornadoesTest(unittest.TestCase): def test_basics(self): with TestPipeline() as p: rows = ( p | 'create' >> beam.Create([{ 'month': 1, 'day': 1, 'tornado': False }, { 'month': 1, 'day': 2, 'tornado': True }, { 'month': 1, 'day': 3, 'tornado': True }, { 'month': 2, 'day': 1, 'tornado': True }])) results = bigquery_tornadoes.count_tornadoes(rows) assert_that( results, equal_to([{ 'month': 1, 'tornado_count': 2 }, { 'month': 2, 'tornado_count': 1 }])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/bigtableio_it_test.py0000644000175100017510000001561715206113032025521 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration tests for bigtableio.""" # pytype: skip-file import datetime import logging import random import string import unittest import uuid from typing import TYPE_CHECKING import pytest import pytz import apache_beam as beam from apache_beam.io.gcp.bigtableio import WriteToBigTable from apache_beam.metrics.metric import MetricsFilter from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.runners.runner import PipelineState from apache_beam.testing.test_pipeline import TestPipeline # Protect against environments where bigtable library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from google.cloud._helpers import _datetime_from_microseconds from google.cloud._helpers import _microseconds_from_datetime from google.cloud.bigtable import row, column_family, Client except ImportError: Client = None UTC = pytz.utc _microseconds_from_datetime = lambda label_stamp: label_stamp _datetime_from_microseconds = lambda micro: micro if TYPE_CHECKING: import google.cloud.bigtable.instance EXISTING_INSTANCES: list['google.cloud.bigtable.instance.Instance'] = [] LABEL_KEY = 'python-bigtable-beam' label_stamp = datetime.datetime.now(datetime.timezone.utc) label_stamp_micros = _microseconds_from_datetime(label_stamp) LABELS = {LABEL_KEY: str(label_stamp_micros)} class GenerateTestRows(beam.PTransform): """ A transform test to run write to the Bigtable Table. A PTransform that generate a list of `DirectRow` to write it in Bigtable Table. """ def __init__(self, number, project_id=None, instance_id=None, table_id=None): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.PTransform.__init__(self) self.number = number self.rand = random.choice(string.ascii_letters + string.digits) self.column_family_id = 'cf1' self.beam_options = { 'project_id': project_id, 'instance_id': instance_id, 'table_id': table_id } def _generate(self): value = ''.join(self.rand for i in range(100)) for index in range(self.number): key = "beam_key%s" % ('{0:07}'.format(index)) direct_row = row.DirectRow(row_key=key) for column_id in range(10): direct_row.set_cell( self.column_family_id, ('field%s' % column_id).encode('utf-8'), value, datetime.datetime.now()) yield direct_row def expand(self, pvalue): beam_options = self.beam_options return ( pvalue | beam.Create(self._generate()) | WriteToBigTable( beam_options['project_id'], beam_options['instance_id'], beam_options['table_id'])) @unittest.skipIf(Client is None, 'GCP Bigtable dependencies are not installed') class BigtableIOWriteTest(unittest.TestCase): """ Bigtable Write Connector Test """ DEFAULT_TABLE_PREFIX = "python-test" instance_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8] cluster_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8] table_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8] number = 500 LOCATION_ID = "us-east1-b" def setUp(self): try: from google.cloud.bigtable import enums self.STORAGE_TYPE = enums.StorageType.HDD self.INSTANCE_TYPE = enums.Instance.Type.DEVELOPMENT except ImportError: self.STORAGE_TYPE = 2 self.INSTANCE_TYPE = 2 self.test_pipeline = TestPipeline(is_integration_test=True) self.runner_name = type(self.test_pipeline.runner).__name__ self.project = self.test_pipeline.get_option('project') self.client = Client(project=self.project, admin=True) self._delete_old_instances() self.instance = self.client.instance( self.instance_id, instance_type=self.INSTANCE_TYPE, labels=LABELS) if not self.instance.exists(): cluster = self.instance.cluster( self.cluster_id, self.LOCATION_ID, default_storage_type=self.STORAGE_TYPE) operation = self.instance.create(clusters=[cluster]) operation.result(timeout=300) # Wait up to 5 min. self.table = self.instance.table(self.table_id) if not self.table.exists(): max_versions_rule = column_family.MaxVersionsGCRule(2) column_family_id = 'cf1' column_families = {column_family_id: max_versions_rule} self.table.create(column_families=column_families) def _delete_old_instances(self): instances = self.client.list_instances() EXISTING_INSTANCES[:] = instances def age_in_hours(micros): return ( datetime.datetime.now(datetime.timezone.utc) - (_datetime_from_microseconds(micros))).total_seconds() // 3600 CLEAN_INSTANCE = [ i for instance in EXISTING_INSTANCES for i in instance if ( LABEL_KEY in i.labels.keys() and (age_in_hours(int(i.labels[LABEL_KEY])) >= 2)) ] if CLEAN_INSTANCE: for instance in CLEAN_INSTANCE: instance.delete() def tearDown(self): if self.instance.exists(): self.instance.delete() @pytest.mark.it_postcommit def test_bigtable_write(self): number = self.number pipeline_args = self.test_pipeline.options_list pipeline_options = PipelineOptions(pipeline_args) with beam.Pipeline(options=pipeline_options) as pipeline: config_data = { 'project_id': self.project, 'instance_id': self.instance_id, 'table_id': self.table_id } _ = ( pipeline | 'Generate Direct Rows' >> GenerateTestRows(number, **config_data)) assert pipeline.result.state == PipelineState.DONE read_rows = self.table.read_rows() assert len([_ for _ in read_rows]) == number if not hasattr(pipeline.result, 'has_job') or pipeline.result.has_job: read_filter = MetricsFilter().with_name('Written Row') query_result = pipeline.result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] logging.info('Number of Rows: %d', read_counter.committed) assert read_counter.committed == number if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/coders.py0000644000175100017510000000637315206113032023143 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A workflow using custom JSON-based coders for text sources and sinks. The input file contains a JSON string on each line describing a match record using the following schema: {'guest': [TEAM_NAME, GOALS], 'host': [TEAM_NAME, GOALS]} The output file will contain the computed points for each team with one team per line in the following format: [TEAM_NAME, POINTS] """ # pytype: skip-file import argparse import json import logging import apache_beam as beam from apache_beam.coders import Coder from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions class JsonCoder(Coder): """A JSON coder interpreting each line as a JSON string.""" def encode(self, x): return json.dumps(x).encode('utf-8') def decode(self, x): return json.loads(x) def compute_points(record): """Compute points based on the record containing the match result. The function assigns 3 points for a win, 1 point for a draw, and 0 points for a loss (see http://en.wikipedia.org/wiki/Three_points_for_a_win). """ host_name, host_goals = record['host'] guest_name, guest_goals = record['guest'] if host_goals == guest_goals: yield host_name, 1 yield guest_name, 1 elif host_goals > guest_goals: yield host_name, 3 yield guest_name, 0 else: yield host_name, 0 yield guest_name, 3 def run(argv=None): """Runs the workflow computing total points from a collection of matches.""" parser = argparse.ArgumentParser() parser.add_argument('--input', required=True, help='Input file to process.') parser.add_argument( '--output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = True with beam.Pipeline(options=pipeline_options) as p: ( # pylint: disable=expression-not-assigned p | 'read' >> ReadFromText(known_args.input, coder=JsonCoder()) | 'points' >> beam.FlatMap(compute_points) | beam.CombinePerKey(sum) | 'write' >> WriteToText(known_args.output, coder=JsonCoder())) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/coders_it_test.py0000644000175100017510000000561615206113032024675 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for Coders example.""" # pytype: skip-file import json import logging import unittest import uuid import pytest from apache_beam.examples.cookbook import coders from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import create_file from apache_beam.testing.test_utils import read_files_from_pattern def format_result(result_string): def format_tuple(result_elem_list): [country, counter] = result_elem_list return country, int(counter.strip()) result_list = list( map( lambda result_elem: format_tuple(result_elem.split(',')), result_string.replace('\'', '').replace('[', '').replace(']', '').replace( '\"', '').split('\n'))) return result_list class CodersIT(unittest.TestCase): SAMPLE_RECORDS = [{ 'host': ['Germany', 1], 'guest': ['Italy', 0] }, { 'host': ['Germany', 1], 'guest': ['Brasil', 3] }, { 'host': ['Brasil', 1], 'guest': ['Italy', 0] }] EXPECTED_RESULT = [('Italy', 0), ('Brasil', 6), ('Germany', 3)] @pytest.mark.no_xdist @pytest.mark.examples_postcommit def test_coders_output_files_on_small_input(self): test_pipeline = TestPipeline(is_integration_test=True) # Setup the files with expected content. OUTPUT_FILE_DIR = \ 'gs://temp-storage-for-end-to-end-tests/py-it-cloud/output' output = '/'.join([OUTPUT_FILE_DIR, str(uuid.uuid4()), 'result']) INPUT_FILE_DIR = \ 'gs://temp-storage-for-end-to-end-tests/py-it-cloud/input' input = '/'.join([INPUT_FILE_DIR, str(uuid.uuid4()), 'input.txt']) create_file(input, '\n'.join(map(json.dumps, self.SAMPLE_RECORDS))) extra_opts = {'input': input, 'output': output} coders.run(test_pipeline.get_full_options_as_args(**extra_opts)) # Load result file and compare. result = read_files_from_pattern('%s*' % output).strip() self.assertEqual( sorted(self.EXPECTED_RESULT), sorted(format_result(result))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/coders_test.py0000644000175100017510000000343315206113032024174 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the coders example.""" # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam.examples.cookbook import coders from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class CodersTest(unittest.TestCase): SAMPLE_RECORDS = [{ 'host': ['Germany', 1], 'guest': ['Italy', 0] }, { 'host': ['Germany', 1], 'guest': ['Brasil', 3] }, { 'host': ['Brasil', 1], 'guest': ['Italy', 0] }] EXPECTED_RESULT = [('Italy', 0), ('Brasil', 6), ('Germany', 3)] def test_compute_points(self): with TestPipeline() as p: records = p | 'create' >> beam.Create(self.SAMPLE_RECORDS) result = ( records | 'points' >> beam.FlatMap(coders.compute_points) | beam.CombinePerKey(sum)) assert_that(result, equal_to(self.EXPECTED_RESULT)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/combiners_test.py0000644000175100017510000000514015206113032024673 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Simple tests to showcase combiners. The tests are meant to be "copy/paste" code snippets for the topic they address (combiners in this case). Most examples use neither sources nor sinks. The input data is generated simply with a Create transform and the output is checked directly on the last PCollection produced. """ # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class CombinersTest(unittest.TestCase): """Tests showcasing Dataflow combiners.""" SAMPLE_DATA = [('a', 1), ('b', 10), ('a', 2), ('a', 3), ('b', 20), ('c', 100)] def test_combine_per_key_with_callable(self): """CombinePerKey using a standard callable reducing iterables. A common case for Dataflow combiners is to sum (or max or min) over the values of each key. Such standard functions can be used directly as combiner functions. In fact, any function "reducing" an iterable to a single value can be used. """ with TestPipeline() as p: result = ( p | beam.Create(CombinersTest.SAMPLE_DATA) | beam.CombinePerKey(sum)) assert_that(result, equal_to([('a', 6), ('b', 30), ('c', 100)])) def test_combine_per_key_with_custom_callable(self): """CombinePerKey using a custom function reducing iterables.""" def multiply(values): result = 1 for v in values: result *= v return result with TestPipeline() as p: result = ( p | beam.Create(CombinersTest.SAMPLE_DATA) | beam.CombinePerKey(multiply)) assert_that(result, equal_to([('a', 6), ('b', 200), ('c', 100)])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/custom_ptransform.py0000644000175100017510000000760115206113032025444 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Various implementations of a Count custom PTransform. These example show the different ways you can write custom PTransforms. """ # pytype: skip-file import argparse import logging import apache_beam as beam from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions # pylint doesn't understand our pipeline syntax: # pylint:disable=expression-not-assigned class Count1(beam.PTransform): """Count as a subclass of PTransform, with an apply method.""" def expand(self, pcoll): return ( pcoll | 'ParWithOne' >> beam.Map(lambda v: (v, 1)) | beam.CombinePerKey(sum)) def run_count1(known_args, options): """Runs the first example pipeline.""" logging.info('Running first pipeline') with beam.Pipeline(options=options) as p: ( p | beam.io.ReadFromText(known_args.input) | Count1() | beam.io.WriteToText(known_args.output)) @beam.ptransform_fn def Count2(pcoll): # pylint: disable=invalid-name """Count as a decorated function.""" return ( pcoll | 'PairWithOne' >> beam.Map(lambda v: (v, 1)) | beam.CombinePerKey(sum)) def run_count2(known_args, options): """Runs the second example pipeline.""" logging.info('Running second pipeline') with beam.Pipeline(options=options) as p: ( p | ReadFromText(known_args.input) | Count2() # pylint: disable=no-value-for-parameter | WriteToText(known_args.output)) @beam.ptransform_fn def Count3(pcoll, factor=1): # pylint: disable=invalid-name """Count as a decorated function with a side input. Args: pcoll: the PCollection passed in from the previous transform factor: the amount by which to count Returns: A PCollection counting the number of times each unique element occurs. """ return ( pcoll | 'PairWithOne' >> beam.Map(lambda v: (v, factor)) | beam.CombinePerKey(sum)) def run_count3(known_args, options): """Runs the third example pipeline.""" logging.info('Running third pipeline') with beam.Pipeline(options=options) as p: ( p | ReadFromText(known_args.input) | Count3(2) # pylint: disable=no-value-for-parameter | WriteToText(known_args.output)) def get_args(argv): """Determines user specified arguments from the given list of arguments. Args: argv: all arguments. Returns: A pair of argument lists containing known and remaining arguments. """ parser = argparse.ArgumentParser() parser.add_argument('--input', required=True, help='Input file to process.') parser.add_argument( '--output', required=True, help='Output file to write results to.') return parser.parse_known_args(argv) def run(argv=None): known_args, pipeline_args = get_args(argv) # pipeline initialization may modify PipelineOptions object. # Create instances for each. run_count1(known_args, PipelineOptions(pipeline_args)) run_count2(known_args, PipelineOptions(pipeline_args)) run_count3(known_args, PipelineOptions(pipeline_args)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/custom_ptransform_it_test.py0000644000175100017510000000507115206113032027176 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for Custom PTransform example.""" # pytype: skip-file import logging import unittest import uuid import pytest from apache_beam.examples.cookbook import custom_ptransform from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import create_file from apache_beam.testing.test_utils import read_files_from_pattern def format_result(result_string): def format_tuple(result_elem_list): [country, counter] = result_elem_list return country, int(counter.strip()) result_list = list( map( lambda result_elem: format_tuple(result_elem.split(',')), result_string.replace('\'', '').replace('[', '').replace(']', '').replace( '\"', '').split('\n'))) return result_list class CustomPTransformIT(unittest.TestCase): WORDS = ['CAT', 'DOG', 'CAT', 'CAT', 'DOG'] EXPECTED_RESULT = "('CAT DOG CAT CAT DOG', 2)" @pytest.mark.examples_postcommit def test_custom_ptransform_output_files_on_small_input(self): test_pipeline = TestPipeline(is_integration_test=True) # Setup the files with expected content. temp_location = test_pipeline.get_option('temp_location') input = '/'.join([temp_location, str(uuid.uuid4()), 'input.txt']) output = '/'.join([temp_location, str(uuid.uuid4()), 'result']) create_file(input, ' '.join(self.WORDS)) extra_opts = {'input': input, 'output': output} custom_ptransform.run(test_pipeline.get_full_options_as_args(**extra_opts)) # Load result file and compare. result = read_files_from_pattern('%s*' % output).strip() self.assertEqual(result, self.EXPECTED_RESULT) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/custom_ptransform_test.py0000644000175100017510000000374615206113032026511 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for the various custom Count implementation examples.""" # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam.examples.cookbook import custom_ptransform from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class CustomCountTest(unittest.TestCase): WORDS = ['CAT', 'DOG', 'CAT', 'CAT', 'DOG'] def create_content_input_file(self, path, contents): logging.info('Creating temp file: %s', path) with open(path, 'w') as f: f.write(contents) def test_count1(self): self.run_pipeline(custom_ptransform.Count1()) def test_count2(self): self.run_pipeline(custom_ptransform.Count2()) def test_count3(self): factor = 2 self.run_pipeline(custom_ptransform.Count3(factor), factor=factor) def run_pipeline(self, count_implementation, factor=1): with TestPipeline() as p: words = p | beam.Create(self.WORDS) result = words | count_implementation assert_that( result, equal_to([('CAT', (3 * factor)), ('DOG', (2 * factor))])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/datastore_wordcount.py0000644000175100017510000002275215206113032025755 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A word-counting workflow that uses Google Cloud Datastore. This example shows how to use ``datastoreio`` to read from and write to Google Cloud Datastore. Note that running this example may incur charge for Cloud Datastore operations. See https://developers.google.com/datastore/ for more details on Google Cloud Datastore. See https://beam.apache.org/get-started/quickstart on how to run a Beam pipeline. Read-only Mode: In this mode, this example reads Cloud Datastore entities using the ``datastoreio.ReadFromDatastore`` transform, extracts the words, counts them and write the output to a set of files. The following options must be provided to run this pipeline in read-only mode: `` --project GCP_PROJECT --kind YOUR_DATASTORE_KIND --output [YOUR_LOCAL_FILE *or* gs://YOUR_OUTPUT_PATH] --read_only `` Read-write Mode: In this mode, this example reads words from an input file, converts them to Beam ``Entity`` objects and writes them to Cloud Datastore using the ``datastoreio.WriteToDatastore`` transform. The second pipeline will then read these Cloud Datastore entities using the ``datastoreio.ReadFromDatastore`` transform, extract the words, count them and write the output to a set of files. The following options must be provided to run this pipeline in read-write mode: `` --project GCP_PROJECT --kind YOUR_DATASTORE_KIND --output [YOUR_LOCAL_FILE *or* gs://YOUR_OUTPUT_PATH] `` """ # pytype: skip-file import argparse import logging import re import sys from collections.abc import Iterable from typing import Optional from typing import Text import uuid import apache_beam as beam from apache_beam.io import ReadFromText from apache_beam.io.gcp.datastore.v1new.datastoreio import ReadFromDatastore from apache_beam.io.gcp.datastore.v1new.datastoreio import WriteToDatastore from apache_beam.io.gcp.datastore.v1new.types import Entity from apache_beam.io.gcp.datastore.v1new.types import Key from apache_beam.io.gcp.datastore.v1new.types import Query from apache_beam.metrics import Metrics from apache_beam.metrics.metric import MetricsFilter from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions @beam.typehints.with_input_types(Entity) @beam.typehints.with_output_types(Text) class WordExtractingDoFn(beam.DoFn): """Parse each line of input text into words.""" def __init__(self): self.empty_line_counter = Metrics.counter('main', 'empty_lines') self.word_length_counter = Metrics.counter('main', 'word_lengths') self.word_counter = Metrics.counter('main', 'total_words') self.word_lengths_dist = Metrics.distribution('main', 'word_len_dist') def process(self, element: Entity) -> Optional[Iterable[Text]]: """Extract words from the 'content' property of Cloud Datastore entities. The element is a line of text. If the line is blank, note that, too. Args: element: the input entity to be processed Returns: A list of words found. """ text_line = element.properties.get('content', '') if not text_line: self.empty_line_counter.inc() return None words = re.findall(r'[A-Za-z\']+', text_line) for w in words: self.word_length_counter.inc(len(w)) self.word_lengths_dist.update(len(w)) self.word_counter.inc() return words class EntityWrapper(object): """Create a Cloud Datastore entity from the given string.""" def __init__(self, project, namespace, kind, ancestor): self._project = project self._namespace = namespace self._kind = kind self._ancestor = ancestor def make_entity(self, content): ancestor_key = Key([self._kind, self._ancestor], namespace=self._namespace, project=self._project) # Namespace and project are inherited from parent key. key = Key([self._kind, str(uuid.uuid4())], parent=ancestor_key) entity = Entity(key) entity.set_properties({'content': content}) return entity def write_to_datastore(project, user_options, pipeline_options): """Creates a pipeline that writes entities to Cloud Datastore.""" with beam.Pipeline(options=pipeline_options) as p: _ = ( p | 'read' >> ReadFromText(user_options.input) | 'create entity' >> beam.Map( EntityWrapper( project, user_options.namespace, user_options.kind, user_options.ancestor).make_entity) | 'write to datastore' >> WriteToDatastore(project)) def make_ancestor_query(project, kind, namespace, ancestor): """Creates a Cloud Datastore ancestor query. The returned query will fetch all the entities that have the parent key name set to the given `ancestor`. """ ancestor_key = Key([kind, ancestor], project=project, namespace=namespace) return Query(kind, project, namespace, ancestor_key) def read_from_datastore(project, user_options, pipeline_options): """Creates a pipeline that reads entities from Cloud Datastore.""" p = beam.Pipeline(options=pipeline_options) # Create a query to read entities from datastore. query = make_ancestor_query( project, user_options.kind, user_options.namespace, user_options.ancestor) # Read entities from Cloud Datastore into a PCollection. lines = p | 'read from datastore' >> ReadFromDatastore(query) # Count the occurrences of each word. def count_ones(word_ones): (word, ones) = word_ones return word, sum(ones) counts = ( lines | 'split' >> beam.ParDo(WordExtractingDoFn()) | 'pair_with_one' >> beam.Map(lambda x: (x, 1)) | 'group' >> beam.GroupByKey() | 'count' >> beam.Map(count_ones)) # Format the counts into a PCollection of strings. def format_result(word_count): (word, count) = word_count return '%s: %s' % (word, count) output = counts | 'format' >> beam.Map(format_result) # Write the output using a "Write" transform that has side effects. # pylint: disable=expression-not-assigned output | 'write' >> beam.io.WriteToText( file_path_prefix=user_options.output, num_shards=user_options.num_shards) result = p.run() # Wait until completion, main thread would access post-completion job results. result.wait_until_finish() return result def run(argv=None): """Main entry point; defines and runs the wordcount pipeline.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='Input file to process.') parser.add_argument( '--kind', dest='kind', required=True, help='Datastore Kind') parser.add_argument( '--namespace', dest='namespace', help='Datastore Namespace') parser.add_argument( '--ancestor', dest='ancestor', default='root', help='The ancestor key name for all entities.') parser.add_argument( '--output', dest='output', required=True, help='Output file to write results to.') parser.add_argument( '--read_only', action='store_true', help='Read an existing dataset, do not write first') parser.add_argument( '--num_shards', dest='num_shards', type=int, # If the system should choose automatically. default=0, help='Number of output shards') known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) project = pipeline_options.view_as(GoogleCloudOptions).project if project is None: parser.print_usage() print(sys.argv[0] + ': error: argument --project is required') sys.exit(1) # Write to Datastore if `read_only` options is not specified. if not known_args.read_only: write_to_datastore(project, known_args, pipeline_options) # Read entities from Datastore. result = read_from_datastore(project, known_args, pipeline_options) empty_lines_filter = MetricsFilter().with_name('empty_lines') query_result = result.metrics().query(empty_lines_filter) if query_result['counters']: empty_lines_counter = query_result['counters'][0] logging.info('number of empty lines: %d', empty_lines_counter.committed) else: logging.warning('unable to retrieve counter metrics from runner') word_lengths_filter = MetricsFilter().with_name('word_len_dist') query_result = result.metrics().query(word_lengths_filter) if query_result['distributions']: word_lengths_dist = query_result['distributions'][0] logging.info('average word length: %d', word_lengths_dist.committed.mean) else: logging.warning('unable to retrieve distribution metrics from runner') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/datastore_wordcount_it_test.py0000644000175100017510000000505315206113032027503 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for Datastore Wordcount example.""" # pytype: skip-file import logging import time import unittest import pytest from hamcrest.core.core.allof import all_of from apache_beam.testing.pipeline_verifiers import FileChecksumMatcher from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline # Don't fail during pytest collection scan. try: from apache_beam.examples.cookbook import datastore_wordcount except ImportError: datastore_wordcount = None class DatastoreWordCountIT(unittest.TestCase): DATASTORE_WORDCOUNT_KIND = "DatastoreWordCount" EXPECTED_CHECKSUM = '826f69ed0275858c2e098f1e8407d4e3ba5a4b3f' @pytest.mark.it_postcommit def test_datastore_wordcount_it(self): test_pipeline = TestPipeline(is_integration_test=True) kind = self.DATASTORE_WORDCOUNT_KIND output = '/'.join([ test_pipeline.get_option('output'), str(int(time.time() * 1000)), 'datastore_wordcount_results' ]) arg_sleep_secs = test_pipeline.get_option('sleep_secs') sleep_secs = int(arg_sleep_secs) if arg_sleep_secs is not None else None pipeline_verifiers = [ PipelineStateMatcher(), FileChecksumMatcher( output + '*-of-*', self.EXPECTED_CHECKSUM, sleep_secs) ] extra_opts = { 'kind': kind, 'output': output, # Comment this out to regenerate input data on Datastore (delete # existing data first using the bulk delete Dataflow template). 'read_only': True, 'on_success_matcher': all_of(*pipeline_verifiers) } datastore_wordcount.run( test_pipeline.get_full_options_as_args(**extra_opts)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/filters.py0000644000175100017510000000752515206113032023334 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """An example workflow that demonstrates filters and other features. - Reading and writing data from BigQuery. - Manipulating BigQuery rows (as Python dicts) in memory. - Global aggregates. - Filtering PCollections using both user-specified parameters as well as global aggregates computed during pipeline execution. """ # pytype: skip-file import argparse import logging import apache_beam as beam from apache_beam.pvalue import AsSingleton def filter_cold_days(input_data, month_filter): """Workflow computing rows in a specific month with low temperatures. Args: input_data: a PCollection of dictionaries representing table rows. Each dictionary must have the keys ['year', 'month', 'day', and 'mean_temp']. month_filter: an int representing the month for which colder-than-average days should be returned. Returns: A PCollection of dictionaries with the same keys described above. Each row represents a day in the specified month where temperatures were colder than the global mean temperature in the entire dataset. """ # Project to only the desired fields from a complete input row. # E.g., SELECT f1, f2, f3, ... FROM InputTable. projection_fields = ['year', 'month', 'day', 'mean_temp'] fields_of_interest = ( input_data | 'Projected' >> beam.Map(lambda row: {f: row[f] for f in projection_fields})) # Compute the global mean temperature. global_mean = AsSingleton( fields_of_interest | 'ExtractMean' >> beam.Map(lambda row: row['mean_temp']) | 'GlobalMean' >> beam.combiners.Mean.Globally()) # Filter to the rows representing days in the month of interest # in which the mean daily temperature is below the global mean. return ( fields_of_interest | 'DesiredMonth' >> beam.Filter(lambda row: row['month'] == month_filter) | 'BelowMean' >> beam.Filter( lambda row, mean: row['mean_temp'] < mean, global_mean)) def run(argv=None): """Constructs and runs the example filtering pipeline.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', help='BigQuery table to read from.', default='apache-beam-testing.samples.weather_stations') parser.add_argument( '--output', required=True, help='BigQuery table to write to.') parser.add_argument( '--month_filter', default=7, help='Numeric value of month to filter on.') known_args, pipeline_args = parser.parse_known_args(argv) with beam.Pipeline(argv=pipeline_args) as p: input_data = p | beam.io.ReadFromBigQuery(table=known_args.input) # pylint: disable=expression-not-assigned ( filter_cold_days(input_data, known_args.month_filter) | 'SaveToBQ' >> beam.io.WriteToBigQuery( known_args.output, schema='year:INTEGER,month:INTEGER,day:INTEGER,mean_temp:FLOAT', create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/filters_test.py0000644000175100017510000001047415206113032024370 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the filters example.""" # pytype: skip-file import logging import time import unittest import pytest from hamcrest.core.core.allof import all_of import apache_beam as beam from apache_beam.examples.cookbook import filters from apache_beam.io.gcp.tests import utils from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryMatcher from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class FiltersTest(unittest.TestCase): # The default checksum is a SHA-1 hash generated from sorted rows reading # from expected Bigquery table. DEFAULT_CHECKSUM = '813b6da1624334732fad4467d74a7c8a62559c6b' # Note that 'removed' should be projected away by the pipeline input_data = [ { 'year': 2010, 'month': 1, 'day': 1, 'mean_temp': 3, 'removed': 'a' }, { 'year': 2012, 'month': 1, 'day': 2, 'mean_temp': 3, 'removed': 'a' }, { 'year': 2011, 'month': 1, 'day': 3, 'mean_temp': 5, 'removed': 'a' }, { 'year': 2013, 'month': 2, 'day': 1, 'mean_temp': 3, 'removed': 'a' }, { 'year': 2011, 'month': 3, 'day': 3, 'mean_temp': 5, 'removed': 'a' }, ] def _get_result_for_month(self, pipeline, month): rows = (pipeline | 'create' >> beam.Create(self.input_data)) results = filters.filter_cold_days(rows, month) return results def test_basics(self): """Test that the correct result is returned for a simple dataset.""" with TestPipeline() as p: results = self._get_result_for_month(p, 1) assert_that( results, equal_to([{ 'year': 2010, 'month': 1, 'day': 1, 'mean_temp': 3 }, { 'year': 2012, 'month': 1, 'day': 2, 'mean_temp': 3 }])) def test_basic_empty(self): """Test that the correct empty result is returned for a simple dataset.""" with TestPipeline() as p: results = self._get_result_for_month(p, 3) assert_that(results, equal_to([])) def test_basic_empty_missing(self): """Test that the correct empty result is returned for a missing month.""" with TestPipeline() as p: results = self._get_result_for_month(p, 4) assert_that(results, equal_to([])) @pytest.mark.examples_postcommit def test_filters_output_bigquery_matcher(self): test_pipeline = TestPipeline(is_integration_test=True) # Set extra options to the pipeline for test purpose project = test_pipeline.get_option('project') dataset = 'FiltersTestIT' table = 'cold_days_%s' % int(round(time.time() * 1000)) output_table = '.'.join([dataset, table]) query = 'SELECT year, month, day, mean_temp FROM `%s`' % output_table pipeline_verifiers = [ PipelineStateMatcher(), BigqueryMatcher( project=project, query=query, checksum=self.DEFAULT_CHECKSUM) ] extra_opts = { 'output': output_table, 'on_success_matcher': all_of(*pipeline_verifiers) } # Register cleanup before pipeline execution. # Note that actual execution happens in reverse order. self.addCleanup(utils.delete_bq_table, project, dataset, table) # Get pipeline options from command argument: --test-pipeline-options, # and start pipeline job by calling pipeline main function. filters.run(test_pipeline.get_full_options_as_args(**extra_opts)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/group_with_coder.py0000644000175100017510000001113715206113032025221 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """An example of using custom classes and coder for grouping operations. This workflow demonstrates registration and usage of a custom coder for a user- defined class. A deterministic custom coder is needed to use a class as a key in a combine or group operation. This example assumes an input file with, on each line, a comma-separated name and score. """ # pytype: skip-file import argparse import logging import sys import apache_beam as beam from apache_beam import coders from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.typehints.decorators import with_output_types class Player(object): """A custom class used as a key in combine/group transforms.""" def __init__(self, name): self.name = name class PlayerCoder(coders.Coder): """A custom coder for the Player class.""" def encode(self, o): """Encode to bytes with a trace that coder was used.""" # Our encoding prepends an 'x:' prefix. return b'x:%s' % str(o.name).encode('utf-8') def decode(self, s): # To decode, we strip off the prepended 'x:' prefix. s = s.decode('utf-8') assert s[0:2] == 'x:' return Player(s[2:]) def is_deterministic(self): # Since coded Player objects are used as keys below with # beam.CombinePerKey(sum), we require that this coder is deterministic # (i.e., two equivalent instances of the classes are encoded into the same # byte string) in order to guarantee consistent results. return True # Annotate the get_players function so that the typehint system knows that the # input to the CombinePerKey operation is a key-value pair of a Player object # and an integer. @with_output_types(tuple[Player, int]) def get_players(descriptor): name, points = descriptor.split(',') return Player(name), int(points) def run(args=None, save_main_session=True): """Runs the workflow computing total points from a collection of matches.""" if args is None: args = sys.argv[1:] parser = argparse.ArgumentParser() parser.add_argument('--input', required=True, help='Input file to process.') parser.add_argument( '--output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(args) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session with beam.Pipeline(options=pipeline_options) as p: # Register the custom coder for the Player class, so that it will be used in # the computation. coders.registry.register_coder(Player, PlayerCoder) ( # pylint: disable=expression-not-assigned p | ReadFromText(known_args.input) # The get_players function is annotated with a type hint above, # so the type system knows the output type of the following operation # is a key-value pair of a Player and an int. Please see the # documentation for details on types that are inferred automatically # as well as other ways to specify type hints. | beam.Map(get_players) # The output type hint of the previous step is used to infer that the # key type of the following operation is the Player type. Since a # custom coder is registered for the Player class above, # a PlayerCoder will be used to encode Player objects as keys for this # combine operation. | beam.CombinePerKey(sum) | beam.Map(lambda k_v: '%s,%d' % (k_v[0].name, k_v[1])) | WriteToText(known_args.output)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/group_with_coder_test.py0000644000175100017510000001055215206113032026260 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the custom coders example.""" # pytype: skip-file import logging import unittest import uuid import pytest from apache_beam.examples.cookbook import group_with_coder from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import read_files_from_pattern # Protect against environments where gcsio library is not available. try: from apache_beam.io.gcp import gcsio except ImportError: gcsio = None # Patch group_with_coder.PlayerCoder.decode(). To test that the PlayerCoder was # used, we do not strip the prepended 'x:' string when decoding a Player object. group_with_coder.PlayerCoder.decode = lambda self, s: group_with_coder.Player( # type: ignore[assignment] s.decode('utf-8')) def create_content_input_file(path, records): logging.info('Creating file: %s', path) gcs = gcsio.GcsIO() with gcs.open(path, 'w') as f: for record in records: f.write(b'%s\n' % record.encode('utf-8')) return path @unittest.skipIf(gcsio is None, 'GCP dependencies are not installed') @pytest.mark.examples_postcommit class GroupWithCoderTest(unittest.TestCase): SAMPLE_RECORDS = [ 'joe,10', 'fred,3', 'mary,7', 'joe,20', 'fred,6', 'ann,5', 'joe,30', 'ann,10', 'mary,1' ] def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) # Setup the file with expected content. self.temp_location = self.test_pipeline.get_option('temp_location') self.input_file = create_content_input_file( '/'.join([self.temp_location, str(uuid.uuid4()), 'input.txt']), self.SAMPLE_RECORDS) #TODO(https://github.com/apache/beam/issues/23608) Fix and enable @pytest.mark.sickbay_dataflow def test_basics_with_type_check(self): # Run the workflow with pipeline_type_check option. This will make sure # the typehints associated with all transforms will have non-default values # and therefore any custom coders will be used. In our case we want to make # sure the coder for the Player class will be used. output = '/'.join([self.temp_location, str(uuid.uuid4()), 'result']) extra_opts = {'input': self.input_file, 'output': output} group_with_coder.run( self.test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) # Parse result file and compare. results = [] lines = read_files_from_pattern('%s*' % output).splitlines() for line in lines: name, points = line.split(',') results.append((name, int(points))) logging.info('result: %s', results) self.assertEqual( sorted(results), sorted([('x:ann', 15), ('x:fred', 9), ('x:joe', 60), ('x:mary', 8)])) def test_basics_without_type_check(self): # Run the workflow without pipeline_type_check option. This will make sure # the typehints associated with all transforms will have default values and # therefore any custom coders will not be used. The default coder (pickler) # will be used instead. output = '/'.join([self.temp_location, str(uuid.uuid4()), 'result']) extra_opts = {'input': self.input_file, 'output': output} with self.assertRaises(Exception) as context: # yapf: disable group_with_coder.run( self.test_pipeline.get_full_options_as_args(**extra_opts) + ['--no_pipeline_type_check'], save_main_session=False) self.assertIn('Unable to deterministically encode', str(context.exception)) self.assertIn('CombinePerKey(sum)/GroupByKey', str(context.exception)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/mergecontacts.py0000644000175100017510000001357615206113032024525 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Merge phone, email, and mailing address information. A Dataflow pipeline that merges phone, email, and address information associated with the same names. Each input "database" is a tab-delimited text file pairing names with one phone number/email address/mailing address; multiple entries associated with the same name are allowed. Outputs are a tab-delimited text file with the merged information and another file containing some simple statistics. See mergecontacts_test.py for example inputs and outputs. A demonstration of: CoGroupByKey Non-linear pipelines (i.e., pipelines with branches) """ # pytype: skip-file import argparse import logging import re import apache_beam as beam from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to def run(argv=None, assert_results=None, save_main_session=True): parser = argparse.ArgumentParser() parser.add_argument( '--input_email', required=True, help='Email database, with each line formatted as "nameemail".') parser.add_argument( '--input_phone', required=True, help='Phonebook, with each line formatted as "namephone number".') parser.add_argument( '--input_snailmail', required=True, help='Address database, with each line formatted as "nameaddress".') parser.add_argument( '--output_tsv', required=True, help='Tab-delimited output file.') parser.add_argument( '--output_stats', required=True, help='Output file for statistics about the input.') known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session with beam.Pipeline(options=pipeline_options) as p: # Helper: read a tab-separated key-value mapping from a text file, # escape all quotes/backslashes, and convert it a PCollection of # (key, value) pairs. def read_kv_textfile(label, textfile): return ( p | 'Read: %s' % label >> ReadFromText(textfile) | 'Backslash: %s' % label >> beam.Map(lambda x: re.sub(r'\\', r'\\\\', x)) | 'EscapeQuotes: %s' % label >> beam.Map(lambda x: re.sub(r'"', r'\"', x)) | 'Split: %s' % label >> beam.Map(lambda x: re.split(r'\t+', x, 1))) # Read input databases. email = read_kv_textfile('email', known_args.input_email) phone = read_kv_textfile('phone', known_args.input_phone) snailmail = read_kv_textfile('snailmail', known_args.input_snailmail) # Group together all entries under the same name. grouped = (email, phone, snailmail) | 'group_by_name' >> beam.CoGroupByKey() # Prepare tab-delimited output; something like this: # "name""email_1,email_2""phone""first_snailmail_only" def format_as_tsv(name_email_phone_snailmail): (name, (email, phone, snailmail)) = name_email_phone_snailmail return '\t'.join([ '"%s"' % name, '"%s"' % ','.join(email), '"%s"' % ','.join(phone), '"%s"' % next(iter(snailmail), '') ]) tsv_lines = grouped | beam.Map(format_as_tsv) # Compute some stats about our database of people. def without_email(name_email_phone_snailmail): (_, (email, _, _)) = name_email_phone_snailmail return not next(iter(email), None) def without_phones(name_email_phone_snailmail): (_, (_, phone, _)) = name_email_phone_snailmail return not next(iter(phone), None) def without_address(name_email_phone_snailmail): (_, (_, _, snailmail)) = name_email_phone_snailmail return not next(iter(snailmail), None) luddites = grouped | beam.Filter(without_email) # People without email. writers = grouped | beam.Filter(without_phones) # People without phones. nomads = grouped | beam.Filter(without_address) # People without addresses. num_luddites = luddites | 'Luddites' >> beam.combiners.Count.Globally() num_writers = writers | 'Writers' >> beam.combiners.Count.Globally() num_nomads = nomads | 'Nomads' >> beam.combiners.Count.Globally() # Write tab-delimited output. # pylint: disable=expression-not-assigned tsv_lines | 'WriteTsv' >> WriteToText(known_args.output_tsv) # TODO(silviuc): Move the assert_results logic to the unit test. if assert_results is not None: expected_luddites, expected_writers, expected_nomads = assert_results assert_that( num_luddites, equal_to([expected_luddites]), label='assert:luddites') assert_that( num_writers, equal_to([expected_writers]), label='assert:writers') assert_that( num_nomads, equal_to([expected_nomads]), label='assert:nomads') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/mergecontacts_test.py0000644000175100017510000001341315206113032025552 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the mergecontacts example.""" # pytype: skip-file import logging import unittest import uuid import pytest from apache_beam.examples.cookbook import mergecontacts from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import create_file from apache_beam.testing.test_utils import read_files_from_pattern class MergeContactsTest(unittest.TestCase): CONTACTS_EMAIL = '\n'.join([ 'Nathan Nomad\tnathan@example.com', 'Nicky Nomad\tnicky@example.com', 'Noreen Nomad\tnoreen@example.com', 'Noreen Nomad\tnomad@example.com', 'Robert B\trobert@example.com', 'Silviu C\tsilviu@example.com', 'Tom S\ttom@example.com', 'Wally Writer\twally@example.com', '' ]) CONTACTS_PHONE = '\n'.join([ 'Larry Luddite\t724-228-3529', 'Lisa Luddite\t304-277-3504', 'Nathan Nomad\t412-466-8968', 'Nicky Nomad\t724-379-5815', 'Noreen Nomad\t412-472-0145', 'Robert B\t814-865-8799', 'Silviu C\t724-537-0671', 'Tom S\t570-368-3420', 'Tom S\t814-793-9655', '' ]) CONTACTS_SNAILMAIL = '\n'.join([ 'Larry Luddite\t1949 Westcott St, Detroit, MI 48222', 'Lisa Luddite\t1949 Westcott St, Detroit, MI 48222', 'Robert B\t601 N 34th St, Seattle, WA 98103', 'Silviu C\t1600 Amphitheatre Pkwy, Mountain View, CA 94043', 'Tom S\t6425 Penn Ave Ste 700, Pittsburgh, PA 15206', 'Wally Writer\t101 Ridge Rd, Chalkyitsik, AK 99788', '' ]) EXPECTED_TSV = '\n'.join([ '\t'.join([ '"Larry Luddite"', '""', '"724-228-3529"', '"1949 Westcott St, Detroit, MI 48222"' ]), '\t'.join([ '"Lisa Luddite"', '""', '"304-277-3504"', '"1949 Westcott St, Detroit, MI 48222"' ]), '\t'.join( ['"Nathan Nomad"', '"nathan@example.com"', '"412-466-8968"', '""']), '\t'.join( ['"Nicky Nomad"', '"nicky@example.com"', '"724-379-5815"', '""']), '\t'.join([ '"Noreen Nomad"', '"nomad@example.com,noreen@example.com"', '"412-472-0145"', '""' ]), '\t'.join([ '"Robert B"', '"robert@example.com"', '"814-865-8799"', '"601 N 34th St, Seattle, WA 98103"' ]), '\t'.join([ '"Silviu C"', '"silviu@example.com"', '"724-537-0671"', '"1600 Amphitheatre Pkwy, Mountain View, CA 94043"' ]), '\t'.join([ '"Tom S"', '"tom@example.com"', '"570-368-3420,814-793-9655"', '"6425 Penn Ave Ste 700, Pittsburgh, PA 15206"' ]), '\t'.join([ '"Wally Writer"', '"wally@example.com"', '""', '"101 Ridge Rd, Chalkyitsik, AK 99788"' ]), '' ]) EXPECTED_STATS = '\n'.join(['2 luddites', '1 writers', '3 nomads', '']) def normalize_tsv_results(self, tsv_data): """Sort .tsv file data so we can compare it with expected output.""" lines_in = tsv_data.strip().split('\n') lines_out = [] for line in lines_in: name, email, phone, snailmail = line.split('\t') lines_out.append( '\t'.join([ name, '"%s"' % ','.join(sorted(email.strip('"').split(','))), '"%s"' % ','.join(sorted(phone.strip('"').split(','))), snailmail ])) return '\n'.join(sorted(lines_out)) + '\n' @pytest.mark.examples_postcommit @pytest.mark.sickbay_flink def test_mergecontacts(self): test_pipeline = TestPipeline(is_integration_test=True) # Setup the files with expected content. temp_location = test_pipeline.get_option('temp_location') input_folder = '/'.join([temp_location, str(uuid.uuid4())]) path_email = create_file( '/'.join([input_folder, 'path_email.txt']), self.CONTACTS_EMAIL) path_phone = create_file( '/'.join([input_folder, 'path_phone.txt']), self.CONTACTS_PHONE) path_snailmail = create_file( '/'.join([input_folder, 'path_snailmail.txt']), self.CONTACTS_SNAILMAIL) result_prefix = '/'.join([temp_location, str(uuid.uuid4()), 'result']) extra_opts = { 'input_email': path_email, 'input_phone': path_phone, 'input_snailmail': path_snailmail, 'output_tsv': '%s.tsv' % result_prefix, 'output_stats': '%s.stats' % result_prefix } pipeline_opts = test_pipeline.get_full_options_as_args(**extra_opts) # Prevent ambiguous option error between output in # args and expected output_tsv and output_stats output_arg = [i for i in pipeline_opts if i.startswith('--output=')] if output_arg: pipeline_opts.remove(output_arg[0]) mergecontacts.run( pipeline_opts, assert_results=(2, 1, 3), save_main_session=False) contents = read_files_from_pattern('%s*' % result_prefix) self.assertEqual(self.EXPECTED_TSV, self.normalize_tsv_results(contents)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/multiple_output_pardo.py0000644000175100017510000001634615206113032026325 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A workflow demonstrating a DoFn with multiple outputs. DoFns may produce multiple outputs. Outputs that are not the default ("main") output are marked with a tag at output time and later the same tag will be used to get the corresponding result (a PCollection) for that output. This is a slightly modified version of the basic wordcount example. In this example words are divided into 2 buckets as shorts words (3 characters in length or less) and words (all other words). There will be 3 output files::: [OUTPUT]-chars : Character count for the input. [OUTPUT]-short-words : Word count for short words only. [OUTPUT]-words : Word count for all other words. To execute this pipeline locally, specify a local output file or output prefix on GCS::: --output [YOUR_LOCAL_FILE | gs://YOUR_OUTPUT_PREFIX] To execute this pipeline using the Google Cloud Dataflow service, specify pipeline configuration::: --project YOUR_PROJECT_ID --region GCE_REGION --staging_location gs://YOUR_STAGING_DIRECTORY --temp_location gs://YOUR_TEMP_DIRECTORY --job_name YOUR_JOB_NAME --runner DataflowRunner and an output prefix on GCS::: --output gs://YOUR_OUTPUT_PREFIX """ # pytype: skip-file # beam-playground: # name: MultipleOutputPardo # description: This is a slightly modified version # of the basic wordcount example. In this example words # are divided into 2 buckets as shorts # words (3 characters in length or less) and words (other). # multifile: false # pipeline_options: --output output.txt # context_line: 172 # categories: # - IO # - Options # - Multiple Outputs # complexity: MEDIUM # tags: # - count # - split # - strings import argparse import logging import re import apache_beam as beam from apache_beam import pvalue from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions class SplitLinesToWordsFn(beam.DoFn): """A transform to split a line of text into individual words. This transform will have 3 outputs: - main output: all words that are longer than 3 characters. - short words output: all other words. - character count output: Number of characters in each processed line. """ # These tags will be used to tag the outputs of this DoFn. OUTPUT_TAG_SHORT_WORDS = 'tag_short_words' OUTPUT_TAG_CHARACTER_COUNT = 'tag_character_count' def process(self, element): """Receives a single element (a line) and produces words and character counts. Important things to note here: - For a single element you may produce multiple main outputs: words of a single line. - For that same input you may produce multiple outputs, potentially across multiple PCollections - Outputs may have different types (count) or may share the same type (words) as with the main output. Args: element: processing element. Yields: words as main output, short words as tagged output, line character count as tagged output. """ # yield a count (integer) to the OUTPUT_TAG_CHARACTER_COUNT tagged # collection. yield pvalue.TaggedOutput(self.OUTPUT_TAG_CHARACTER_COUNT, len(element)) words = re.findall(r'[A-Za-z\']+', element) for word in words: if len(word) <= 3: # yield word as an output to the OUTPUT_TAG_SHORT_WORDS tagged # collection. yield pvalue.TaggedOutput(self.OUTPUT_TAG_SHORT_WORDS, word) else: # yield word to add it to the main collection. yield word class CountWords(beam.PTransform): """A transform to count the occurrences of each word. A PTransform that converts a PCollection containing words into a PCollection of "word: count" strings. """ def expand(self, pcoll): def count_ones(word_ones): (word, ones) = word_ones return (word, sum(ones)) def format_result(word_count): (word, count) = word_count return '%s: %s' % (word, count) return ( pcoll | 'pair_with_one' >> beam.Map(lambda x: (x, 1)) | 'group' >> beam.GroupByKey() | 'count' >> beam.Map(count_ones) | 'format' >> beam.Map(format_result)) def run(argv=None, save_main_session=True): """Runs the workflow counting the long words and short words separately.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='Input file to process.') parser.add_argument( '--output', required=True, help='Output prefix for files to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session with beam.Pipeline(options=pipeline_options) as p: lines = p | ReadFromText(known_args.input) # with_outputs allows accessing the explicitly tagged outputs of a DoFn. split_lines_result = ( lines | beam.ParDo(SplitLinesToWordsFn()).with_outputs( SplitLinesToWordsFn.OUTPUT_TAG_SHORT_WORDS, SplitLinesToWordsFn.OUTPUT_TAG_CHARACTER_COUNT, main='words')) # split_lines_result is an object of type DoOutputsTuple. It supports # accessing result in alternative ways. words, _, _ = split_lines_result short_words = split_lines_result[SplitLinesToWordsFn.OUTPUT_TAG_SHORT_WORDS] character_count = split_lines_result.tag_character_count # pylint: disable=expression-not-assigned ( character_count | 'pair_with_key' >> beam.Map(lambda x: ('chars_temp_key', x)) | beam.GroupByKey() | 'count chars' >> beam.Map(lambda char_counts: sum(char_counts[1])) | 'write chars' >> WriteToText(known_args.output + '-chars')) # pylint: disable=expression-not-assigned ( short_words | 'count short words' >> CountWords() | 'write short words' >> WriteToText(known_args.output + '-short-words')) # pylint: disable=expression-not-assigned ( words | 'count words' >> CountWords() | 'write words' >> WriteToText(known_args.output + '-words')) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/multiple_output_pardo_test.py0000644000175100017510000000572215206113032027360 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the multiple_output_pardo example.""" # pytype: skip-file import logging import re import unittest import uuid import pytest from apache_beam.examples.cookbook import multiple_output_pardo from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import create_file from apache_beam.testing.test_utils import read_files_from_pattern class MultipleOutputParDo(unittest.TestCase): SAMPLE_TEXT = 'A whole new world\nA new fantastic point of view' EXPECTED_SHORT_WORDS = [('A', 2), ('new', 2), ('of', 1)] EXPECTED_WORDS = [('whole', 1), ('world', 1), ('fantastic', 1), ('point', 1), ('view', 1)] def get_wordcount_results(self, result_path): results = [] lines = read_files_from_pattern(result_path).splitlines() for line in lines: match = re.search(r'([A-Za-z]+): ([0-9]+)', line) if match is not None: results.append((match.group(1), int(match.group(2)))) return results @pytest.mark.examples_postcommit @pytest.mark.sickbay_flink def test_multiple_output_pardo(self): test_pipeline = TestPipeline(is_integration_test=True) # Setup the files with expected content. temp_location = test_pipeline.get_option('temp_location') input_folder = '/'.join([temp_location, str(uuid.uuid4())]) input = create_file('/'.join([input_folder, 'input.txt']), self.SAMPLE_TEXT) result_prefix = '/'.join([temp_location, str(uuid.uuid4()), 'result']) extra_opts = {'input': input, 'output': result_prefix} multiple_output_pardo.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) expected_char_count = len(''.join(self.SAMPLE_TEXT.split('\n'))) contents = read_files_from_pattern(result_prefix + '-chars*') self.assertEqual(expected_char_count, int(contents)) short_words = self.get_wordcount_results(result_prefix + '-short-words*') self.assertEqual(sorted(short_words), sorted(self.EXPECTED_SHORT_WORDS)) words = self.get_wordcount_results(result_prefix + '-words*') self.assertEqual(sorted(words), sorted(self.EXPECTED_WORDS)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.0463727 apache_beam-2.74.0/apache_beam/examples/cookbook/ordered_window_elements/0000755000175100017510000000000015206113157026220 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/ordered_window_elements/__init__.py0000644000175100017510000000142015206113032030316 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/ordered_window_elements/batch.py0000644000175100017510000004717615206113032027662 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging from enum import Enum from typing import Any from typing import Callable from typing import Optional import apache_beam as beam from apache_beam.coders import BooleanCoder from apache_beam.coders import PickleCoder from apache_beam.pvalue import AsDict from apache_beam.transforms.combiners import ToListCombineFn from apache_beam.transforms.ptransform import PTransform from apache_beam.transforms.timeutil import TimeDomain from apache_beam.transforms.userstate import OrderedListStateSpec from apache_beam.transforms.userstate import ReadModifyWriteStateSpec from apache_beam.transforms.userstate import TimerSpec from apache_beam.transforms.userstate import on_timer from apache_beam.transforms.window import GlobalWindow from apache_beam.utils.timestamp import MIN_TIMESTAMP from apache_beam.utils.timestamp import DurationTypes # pylint: disable=unused-import from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.timestamp import TimestampTypes # pylint: disable=unused-import class FanOutToWindows(beam.DoFn): """ Assigns each element to all the windows that contain it. This DoFn is used to expand a single element into multiple elements, each associated with a specific window. Args: duration: The duration of each window in seconds. slide_interval: The interval at which windows slide in seconds. offset: The offset for window alignment in seconds. """ def __init__(self, duration, slide_interval, offset): self.duration = duration self.slide_interval = slide_interval self.offset = offset def process(self, element): """ Processes an element and assigns it to relevant windows. Args: element: A tuple (timestamp, value) where timestamp is a Timestamp object and value is the actual element data. Yields: A tuple ((window_start, window_end), element) for each window the element belongs to. """ timestamp = element[0] timestamp_secs = timestamp.micros / 1e6 # Align the timestamp with the windowing scheme. aligned_timestamp = timestamp_secs - self.offset # Calculate the start of the last window that could contain this timestamp. last_window_start_aligned = ((aligned_timestamp // self.slide_interval) * self.slide_interval) last_window_start = last_window_start_aligned + self.offset # To find out the start of the first possible window that covers this # timestamp, we start with the last window and assume we slide backward n # times: # first_possible_start = last_window_start - n * slide_interval # first_possible_end = last_window_start - n * slide_interval + duration # The conditions hold: # first_possible_end > timestamp. # first_possible_end - slide_interval <= timestamp # Therefore, # n < (last_window_start + duration - timestamp) / slide_interval # n >= (last_window_start + duration - timestamp) / slide_interval - 1 # The worst case is that the element is at the beginning of the slide: # i.e. timestamp = last_window_start # And n is an integer satisfies # duration / slide_interval - 1 <= n < duration / slide_interval # Case 1: if duration is divisible by slide_interval, # then n = duration / slide_interval - 1 # Case 2: if duration is not divisible by slide_interval, # then n = duration // slide_interval # A unified solution is n = (duration - 1) // slide_interval n = (self.duration - 1) // self.slide_interval first_possible_start = last_window_start - n * self.slide_interval # We iterate from the first possible window start up to the last one. current_start = first_possible_start while current_start <= last_window_start: # An element is in a window [start, start + duration) if: # start <= timestamp < start + duration if current_start <= timestamp_secs < current_start + self.duration: yield (current_start, current_start + self.duration), element current_start += self.slide_interval class FanOutToSlideBoundaries(beam.DoFn): """ Assigns each element to a window representing its slide. This DoFn is used to group elements by the start of the slide they belong to. This is a preliminary step for generating context information for window gaps. Args: slide_interval: The interval at which windows slide in seconds. offset: The offset for window alignment in seconds. """ def __init__(self, slide_interval, offset): self.slide_interval = slide_interval self.offset = offset def process(self, element): """ Processes an element and assigns it to its corresponding slide boundary. Args: element: A tuple (timestamp, value) where timestamp is a Timestamp object and value is the actual element data. Yields: A tuple (slide_start, element) where slide_start is the beginning timestamp of the slide the element belongs to. """ timestamp = element[0] timestamp_secs = timestamp.micros / 1e6 # Align the timestamp with the windowing scheme. aligned_timestamp = timestamp_secs - self.offset # Calculate the start of the slide containing this timestamp. slide_start_aligned = ((aligned_timestamp // self.slide_interval) * self.slide_interval) slide_start = slide_start_aligned + self.offset # slide_end = slide_start + self.slide_interval yield slide_start, element class GenerateContextDoFn(beam.DoFn): """ Generates context information for filling gaps in windows. This DoFn uses Beam's state and timer features to collect elements within slides and emit a "context" value for each slide. This context value is typically the element with the maximum timestamp within that slide, which can then be used to forward-fill empty windows or gaps at the start of windows. Args: duration: The duration of each window in seconds. slide_interval: The interval at which windows slide in seconds. offset: The offset for window alignment in seconds. default: The default value to use when no context is available. """ ORDERED_BUFFER_STATE = OrderedListStateSpec('ordered_buffer', PickleCoder()) WINDOW_TIMER = TimerSpec('window_timer', TimeDomain.WATERMARK) TIMER_STATE = ReadModifyWriteStateSpec('timer_state', BooleanCoder()) def __init__(self, duration, slide_interval, offset, default): self.duration = duration self.slide_interval = slide_interval self.offset = offset self.default = default def process( self, element=beam.DoFn.ElementParam, timestamp=beam.DoFn.TimestampParam, window_timer=beam.DoFn.TimerParam(WINDOW_TIMER), timer_state=beam.DoFn.StateParam(TIMER_STATE), ordered_buffer=beam.DoFn.StateParam(ORDERED_BUFFER_STATE), ): """ Buffers elements and sets a timer to process them when the window closes. Args: element: The input element, expected to be (key, (slide_start, value)). timestamp: The timestamp of the element. window_timer: The timer for the current window. timer_state: State to track if the timer has been started. ordered_buffer: Ordered list state to buffer elements. """ _, (slide_start, value) = element ordered_buffer.add((Timestamp.of(slide_start), value)) timer_started = timer_state.read() if not timer_started: window_timer.set(GlobalWindow().end) timer_state.write(True) return [] @on_timer(WINDOW_TIMER) def on_timer( self, ordered_buffer=beam.DoFn.StateParam(ORDERED_BUFFER_STATE), ): """ Emits context results when the window timer fires. This method processes the buffered elements, identifies the maximum timestamp element for each slide, and yields context values to fill potential gaps in subsequent windows. Args: ordered_buffer: Ordered list state containing buffered elements. Yields: A tuple (timestamp, element) representing the context for a slide. """ # Emit the context result once we collect all elements prev_max_timestamp_element = None prev_max_timestamp = MIN_TIMESTAMP prev_slide_start = None for slide_start, max_timestamp_event in ordered_buffer.read(): event_ts = max_timestamp_event[0] if prev_slide_start != slide_start: # a new slide starts if prev_max_timestamp_element is not None: # Use the last available max timestamp element for slide between # the last seen slide and the current slide (which includes # empty slides in the middle). start = prev_slide_start while start < slide_start: yield (start + self.slide_interval, prev_max_timestamp_element) start += self.slide_interval else: yield (slide_start, (MIN_TIMESTAMP, self.default)) prev_slide_start = slide_start if prev_max_timestamp < event_ts < slide_start + self.slide_interval: prev_max_timestamp = event_ts prev_max_timestamp_element = max_timestamp_event class WindowGapStrategy(Enum): """ Defines strategies for handling gaps in windows. Attributes: IGNORE: Do nothing for empty windows or gaps. DISCARD: Discard the window. Only applied to empty windows. FORWARD_FILL: Fill empty windows or gaps with the last known value. """ IGNORE = 1 DISCARD = 2 FORWARD_FILL = 3 class WindowGapFillingDoFn(beam.DoFn): """ On-demand filling the start gaps of a window or empty windows. This DoFn takes windowed data and a side input containing context information (e.g., the last element from a previous slide). It uses this context to fill gaps at the beginning of windows or to generate entire empty windows based on the configured gap filling strategies. Args: duration: The duration of each window in seconds. slide_interval: The interval at which windows slide in seconds. default: The default value to use for filling gaps. empty_window_strategy: The strategy for handling completely empty windows. window_start_gap_strategy: The strategy for handling gaps at the start of non-empty windows. """ def __init__( self, duration, slide_interval, default, empty_window_strategy, window_start_gap_strategy): self.duration = duration self.slide_interval = slide_interval self.default = default self.empty_window_strategy = empty_window_strategy self.window_start_gap_strategy = window_start_gap_strategy def process(self, element, context_side): """ Processes a window of elements and fills gaps according to strategies. Args: element: A tuple (window, values) where window is (start_ts, end_ts) and values is a list of elements within that window. context_side: A side input (AsDict) containing context information (slide_start -> max_timestamp_element) for previous slides. Yields: A tuple ((window_start, window_end), filled_values) where filled_values is the list of elements for the window, potentially with gaps filled. """ window, values = element window_start_ts = Timestamp.of(window[0]) # Part 1: Handle the current, non-empty window. # We get the value that should be used to fill gaps at the start of this # window. This value is the element with the max timestamp from the # *previous* slide, provided as a side input. context_for_current_window = context_side.get( window_start_ts, (window_start_ts, self.default)) sorted_values = sorted(values, key=lambda x: x[0]) first_element_ts = sorted_values[0][0] if self.window_start_gap_strategy == WindowGapStrategy.FORWARD_FILL: # If the first element is not at the very beginning of the window, # prepend the context value to fill the gap. if first_element_ts > window_start_ts: _, fill_val = context_for_current_window sorted_values.insert(0, (window_start_ts, fill_val)) yield (Timestamp.of(window[0]), Timestamp.of(window[1])), sorted_values if self.empty_window_strategy == WindowGapStrategy.DISCARD: # We won't emit empty windows prior to the current window under this # strategy return [] # Part 2: Fill completely empty windows that preceded this one. # We iterate backwards from the current window's start time, slide by # slide, to find and fill any empty windows. prev_window_start_ts = window_start_ts - self.slide_interval while True: # Get the context for the preceding window. context_for_prev_window = context_side.get( prev_window_start_ts, (prev_window_start_ts, self.default)) # A preceding window was empty if two conditions are met: # 1. Its context is the same as the current window's context. This implies # that no new elements arrived in the slide(s) between them. # 2. The first element of the current window appeared *after* the end # of the preceding window we are considering. is_empty = ( context_for_prev_window == context_for_current_window and first_element_ts > prev_window_start_ts + self.duration) if is_empty: if self.empty_window_strategy == WindowGapStrategy.FORWARD_FILL: _, fill_val = context_for_prev_window fill_ts = prev_window_start_ts filled_window_values = [(fill_ts, fill_val)] else: assert (self.empty_window_strategy == WindowGapStrategy.IGNORE) filled_window_values = [] yield (prev_window_start_ts, prev_window_start_ts + self.duration), filled_window_values else: # Stop when we find a non-empty window. break prev_window_start_ts -= self.slide_interval return [] def max_timestamp_element(elements): """ Finds the element with the maximum timestamp from a list of elements. Args: elements: A list of elements, where each element is a tuple (timestamp, value). Returns: The element with the maximum timestamp, or None if the list is empty. """ max_timestamp = MIN_TIMESTAMP ret = None for e in elements: if max_timestamp <= e[0]: max_timestamp = e[0] ret = e return ret class OrderedWindowElements(PTransform): """ A PTransform that orders elements within windows and fills gaps. This transform takes a PCollection of elements, assigns them to windows, and then processes these windows to ensure elements are ordered and to fill any gaps (empty windows or gaps at the start of windows) based on specified strategies. Args: duration: The duration of each window. slide_interval: The interval at which windows slide. Defaults to `duration`. offset: The offset for window alignment. default_start_value: The default value to use for filling gaps at the start of windows. empty_window_strategy: The strategy for handling completely empty windows. window_start_gap_strategy: The strategy for handling gaps at the start of non-empty windows. timestamp: An optional callable to extract a timestamp from an element. If not provided, elements are assumed to be (timestamp, value) tuples. """ def __init__( self, duration: DurationTypes, slide_interval: Optional[DurationTypes] = None, offset: TimestampTypes = 0, default_start_value=None, empty_window_strategy: WindowGapStrategy = WindowGapStrategy.IGNORE, window_start_gap_strategy: WindowGapStrategy = WindowGapStrategy.IGNORE, timestamp: Optional[Callable[[Any], Timestamp]] = None): self.duration = duration self.slide_interval = duration if slide_interval is None else slide_interval self.offset = offset self.default_start_value = default_start_value self.empty_window_strategy = empty_window_strategy self.window_start_gap_strategy = window_start_gap_strategy self.timestamp_func = timestamp if self.window_start_gap_strategy == WindowGapStrategy.DISCARD: raise ValueError( "Using DISCARD on windows with start gap is not allowed " "due to potential data loss.") def key_with_timestamp(self, element) -> tuple[Timestamp, Any]: """ Extracts the timestamp from an element and keys it with the element. Args: element: The input element. Returns: A tuple (timestamp, element). """ return self.timestamp_func(element), element def expand(self, input): """ Applies the PTransform to the input PCollection. Args: input: The input PCollection of elements. Returns: A PCollection of ((window_start, window_end), [ordered_elements]) where ordered_elements are sorted by timestamp and gaps are filled according to the specified strategies. """ if self.timestamp_func: input = input | beam.Map(self.key_with_timestamp) # PCollection[((window_start, window_end), [element...])] windowed_data = ( input | "FanOutToWindows" >> beam.ParDo( FanOutToWindows(self.duration, self.slide_interval, self.offset)) | beam.CombinePerKey(ToListCombineFn()) | "LogWindowedData" >> beam.LogElements( prefix="windowed=", level=logging.WARNING)) if (self.empty_window_strategy == WindowGapStrategy.DISCARD and self.window_start_gap_strategy == WindowGapStrategy.IGNORE): # A shortcut for doing nothing on empty window and window start gap. # PCollection[((window_start, window_end), [element...])] return ( windowed_data | beam.MapTuple( lambda window, elements: ((Timestamp.of(window[0]), Timestamp.of(window[1])), sorted( elements))) | "LogReturn" >> beam.LogElements( prefix="return=", level=logging.WARNING)) # PCollection[(slide_start, max_timestamp_element)] fanout_data = ( input | "FanOutToSlideBoundaries" >> beam.ParDo( FanOutToSlideBoundaries(self.slide_interval, self.offset)) | beam.CombinePerKey(max_timestamp_element)) # PCollection[(slide_start, element_to_fill_missing_start)] context = ( fanout_data | beam.WithKeys(0) | "GenerateContextDoFn" >> beam.ParDo( GenerateContextDoFn( self.duration, self.slide_interval, self.offset, self.default_start_value), ) | "LogContext" >> beam.LogElements( prefix="context=", level=logging.WARNING)) # PCollection[((window_start, window_end), [element...])] return ( windowed_data | beam.ParDo( WindowGapFillingDoFn( self.duration, self.slide_interval, self.default_start_value, self.empty_window_strategy, self.window_start_gap_strategy), context_side=AsDict(context)) | "LogReturn" >> beam.LogElements( prefix="return=", level=logging.WARNING)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/ordered_window_elements/batch_test.py0000644000175100017510000003135115206113032030705 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import random import shutil import sys import unittest from parameterized import param from parameterized import parameterized import apache_beam as beam from apache_beam.examples.cookbook.ordered_window_elements.batch import OrderedWindowElements from apache_beam.examples.cookbook.ordered_window_elements.batch import WindowGapStrategy from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.utils.timestamp import Timestamp logging.basicConfig(level=logging.INFO) #logging.basicConfig(level=logging.WARNING) options = PipelineOptions([ "--environment_type=LOOPBACK", "--runner=PrismRunner", #"--runner=FnApiRunner", "--prism_log_kind=dev", # "--runner=PortableRunner", # "--job_endpoint=localhost:8073", ]) ENABLE_LOGGING = False WINDOW_SIZE = 3 def _maybe_log_elements(pcoll, prefix="result="): if ENABLE_LOGGING: return pcoll | beam.LogElements( prefix=prefix, level=logging.WARNING, with_timestamp=True, with_window=True, use_epoch_time=True) else: return pcoll def _create_input_batch(elements: list[int], shuffle_data=True): if shuffle_data: random.shuffle(elements) return beam.Create([(Timestamp.of(e), e) for e in elements]) def _create_input_batch_without_timestamp( elements: list[int], shuffle_data=True): if shuffle_data: random.shuffle(elements) return beam.Create(elements) def _convert_timestamp_to_int(): return beam.MapTuple( lambda window, elements: ((int(window[0].micros // 1e6), int(window[1].micros // 1e6)), [(int(t.micros // 1e6), v) for t, v in elements])) _go_installed = shutil.which('go') is not None _in_windows = sys.platform == "win32" @unittest.skipUnless(_go_installed, 'Go is not installed.') # TODO: Go environments is not configured correctly on Windows test boxes. @unittest.skipIf(_in_windows, reason="Not supported on Windows") class OrderedWindowElementsTest(unittest.TestCase): def setUp(self) -> None: self.options = PipelineOptions([ "--environment_type=LOOPBACK", "--runner=PrismRunner", "--prism_log_kind=dev", # # run on an external Portable Runner for debugging # "--runner=PortableRunner", # "--job_endpoint=localhost:8073", ]) # # dataflow runner option # self.options = PipelineOptions([ # "--runner=DataflowRunner", # "--temp_location=gs://shunping-test/anomaly-temp", # "--staging_location=gs://shunping-test/anomaly-temp", # "--project=apache-beam-testing", # "--region=us-central1", # "--sdk_location=dist/apache_beam-2.70.0.dev0.tar.gz", # #"--pickle_library=dill", # #"--save_main_session", # ]) def test_default(self): with TestPipeline(options=self.options) as p: result = ( p | _create_input_batch([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) | OrderedWindowElements(WINDOW_SIZE)) result = _maybe_log_elements(result) | _convert_timestamp_to_int() assert_that( result, equal_to([ ((0, 3), [(0, 0), (1, 1), (2, 2)]), ((3, 6), [(3, 3), (4, 4), (5, 5)]), ((6, 9), [(6, 6), (7, 7), (8, 8)]), ((9, 12), [(9, 9)]), ])) def test_timestamp_func(self): with TestPipeline(options=self.options) as p: result = ( p | _create_input_batch_without_timestamp( [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) | OrderedWindowElements( WINDOW_SIZE, timestamp=lambda x: Timestamp.of(x))) result = _maybe_log_elements(result) | _convert_timestamp_to_int() assert_that( result, equal_to([ ((0, 3), [(0, 0), (1, 1), (2, 2)]), ((3, 6), [(3, 3), (4, 4), (5, 5)]), ((6, 9), [(6, 6), (7, 7), (8, 8)]), ((9, 12), [(9, 9)]), ])) def test_offset(self): with TestPipeline(options=self.options) as p: result = ( p | _create_input_batch([2, 3, 4, 5, 6, 7, 8, 9]) | OrderedWindowElements(WINDOW_SIZE, offset=2)) result = _maybe_log_elements(result) | _convert_timestamp_to_int() assert_that( result, equal_to([ ((2, 5), [(2, 2), (3, 3), (4, 4)]), # window start at 2 ((5, 8), [(5, 5), (6, 6), (7, 7)]), ((8, 11), [(8, 8), (9, 9)]) ])) def test_slide_interval(self): with TestPipeline(options=self.options) as p: result = ( p | _create_input_batch([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) | OrderedWindowElements(WINDOW_SIZE, slide_interval=1)) result = _maybe_log_elements(result) | _convert_timestamp_to_int() assert_that( result, equal_to([ ((-2, 1), [(0, 0)]), ((-1, 2), [(0, 0), (1, 1)]), ((0, 3), [(0, 0), (1, 1), (2, 2)]), ((1, 4), [(1, 1), (2, 2), (3, 3)]), ((2, 5), [(2, 2), (3, 3), (4, 4)]), ((3, 6), [(3, 3), (4, 4), (5, 5)]), ((4, 7), [(4, 4), (5, 5), (6, 6)]), ((5, 8), [(5, 5), (6, 6), (7, 7)]), ((6, 9), [(6, 6), (7, 7), (8, 8)]), ((7, 10), [(7, 7), (8, 8), (9, 9)]), ((8, 11), [(8, 8), (9, 9)]), ((9, 12), [(9, 9)]), ])) @parameterized.expand([ param( empty_window_strategy=WindowGapStrategy.DISCARD, window_start_gap_strategy=WindowGapStrategy.IGNORE), param( empty_window_strategy=WindowGapStrategy.DISCARD, window_start_gap_strategy=WindowGapStrategy.FORWARD_FILL), param( empty_window_strategy=WindowGapStrategy.IGNORE, window_start_gap_strategy=WindowGapStrategy.IGNORE), param( empty_window_strategy=WindowGapStrategy.IGNORE, window_start_gap_strategy=WindowGapStrategy.FORWARD_FILL), param( empty_window_strategy=WindowGapStrategy.FORWARD_FILL, window_start_gap_strategy=WindowGapStrategy.IGNORE), param( empty_window_strategy=WindowGapStrategy.FORWARD_FILL, window_start_gap_strategy=WindowGapStrategy.FORWARD_FILL), ]) def test_gaps(self, empty_window_strategy, window_start_gap_strategy): if empty_window_strategy == WindowGapStrategy.DISCARD: if window_start_gap_strategy == WindowGapStrategy.IGNORE: expected = [ ((0, 3), [(1, 1), (2, 2)]), ((3, 6), [(3, 3), (4, 4)]), # empty windows (6, 9), (9, 12), (12, 15) are discarded ((15, 18), [(16, 16)]), ((18, 21), [(20, 20)]), ] elif window_start_gap_strategy == WindowGapStrategy.FORWARD_FILL: expected = [ # fill the beginning of (0, 3) with default value `None` ((0, 3), [(0, None), (1, 1), (2, 2)]), ((3, 6), [(3, 3), (4, 4)]), # fill the beginning of (15, 18) with 4 from Timestamp(4) ((15, 18), [(15, 4), (16, 16)]), # fill the beginning of (18, 21) with 16 from Timestamp(16) ((18, 21), [(18, 16), (20, 20)]), ] elif empty_window_strategy == WindowGapStrategy.IGNORE: if window_start_gap_strategy == WindowGapStrategy.IGNORE: expected = [ ((0, 3), [(1, 1), (2, 2)]), ((3, 6), [(3, 3), (4, 4)]), ((6, 9), []), # empty windows are kept ((9, 12), []), ((12, 15), []), ((15, 18), [(16, 16)]), ((18, 21), [(20, 20)]), ] elif window_start_gap_strategy == WindowGapStrategy.FORWARD_FILL: expected = [ ((0, 3), [(0, None), (1, 1), (2, 2)]), ((3, 6), [(3, 3), (4, 4)]), ((6, 9), []), ((9, 12), []), ((12, 15), []), ((15, 18), [(15, 4), (16, 16)]), ((18, 21), [(18, 16), (20, 20)]), ] elif empty_window_strategy == WindowGapStrategy.FORWARD_FILL: if window_start_gap_strategy == WindowGapStrategy.IGNORE: expected = [ ((0, 3), [(1, 1), (2, 2)]), ((3, 6), [(3, 3), (4, 4)]), ((6, 9), [(6, 4)]), # empty windows are forward filled ((9, 12), [(9, 4)]), ((12, 15), [(12, 4)]), ((15, 18), [(16, 16)]), ((18, 21), [(20, 20)]), ] elif window_start_gap_strategy == WindowGapStrategy.FORWARD_FILL: expected = [ ((0, 3), [(0, None), (1, 1), (2, 2)]), ((3, 6), [(3, 3), (4, 4)]), ((6, 9), [(6, 4)]), ((9, 12), [(9, 4)]), ((12, 15), [(12, 4)]), ((15, 18), [(15, 4), (16, 16)]), ((18, 21), [(18, 16), (20, 20)]), ] with TestPipeline(options=self.options) as p: result = ( p | _create_input_batch([1, 2, 3, 4, 16, 20]) | OrderedWindowElements( WINDOW_SIZE, empty_window_strategy=empty_window_strategy, window_start_gap_strategy=window_start_gap_strategy)) result = _maybe_log_elements(result) | _convert_timestamp_to_int() assert_that(result, equal_to(expected)) @parameterized.expand([ param( empty_window_strategy=WindowGapStrategy.DISCARD, window_start_gap_strategy=WindowGapStrategy.IGNORE), param( empty_window_strategy=WindowGapStrategy.DISCARD, window_start_gap_strategy=WindowGapStrategy.FORWARD_FILL), param( empty_window_strategy=WindowGapStrategy.IGNORE, window_start_gap_strategy=WindowGapStrategy.IGNORE), param( empty_window_strategy=WindowGapStrategy.IGNORE, window_start_gap_strategy=WindowGapStrategy.FORWARD_FILL), param( empty_window_strategy=WindowGapStrategy.FORWARD_FILL, window_start_gap_strategy=WindowGapStrategy.IGNORE), param( empty_window_strategy=WindowGapStrategy.FORWARD_FILL, window_start_gap_strategy=WindowGapStrategy.FORWARD_FILL), ]) def test_long_slide(self, empty_window_strategy, window_start_gap_strategy): if empty_window_strategy == WindowGapStrategy.DISCARD: if window_start_gap_strategy == WindowGapStrategy.IGNORE: expected = [((0, 2), [(0, 0)]), ((5, 7), [(6, 6)]), ((15, 17), [(16, 16)])] elif window_start_gap_strategy == WindowGapStrategy.FORWARD_FILL: expected = [((0, 2), [(0, 0)]), ((5, 7), [(5, 4), (6, 6)]), ((15, 17), [(15, 7), (16, 16)])] elif empty_window_strategy == WindowGapStrategy.IGNORE: if window_start_gap_strategy == WindowGapStrategy.IGNORE: expected = [((0, 2), [(0, 0)]), ((5, 7), [(6, 6)]), ((10, 12), []), ((15, 17), [(16, 16)])] elif window_start_gap_strategy == WindowGapStrategy.FORWARD_FILL: expected = [((0, 2), [(0, 0)]), ((5, 7), [(5, 4), (6, 6)]), ((10, 12), []), ((15, 17), [(15, 7), (16, 16)])] elif empty_window_strategy == WindowGapStrategy.FORWARD_FILL: if window_start_gap_strategy == WindowGapStrategy.IGNORE: expected = [((0, 2), [(0, 0)]), ((5, 7), [(6, 6)]), ((10, 12), [(10, 7)]), ((15, 17), [(16, 16)])] elif window_start_gap_strategy == WindowGapStrategy.FORWARD_FILL: expected = [((0, 2), [(0, 0)]), ((5, 7), [(5, 4), (6, 6)]), ((10, 12), [(10, 7)]), ((15, 17), [(15, 7), (16, 16)])] with TestPipeline(options=self.options) as p: result = ( p | _create_input_batch([0, 2, 4, 6, 7, 16]) | OrderedWindowElements( 2, 5, 0, -100, empty_window_strategy=empty_window_strategy, window_start_gap_strategy=window_start_gap_strategy) ) # window size < slide interval result = _maybe_log_elements(result) | _convert_timestamp_to_int() assert_that(result, equal_to(expected)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/ordered_window_elements/streaming.py0000644000175100017510000006307315206113032030564 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import enum import logging from typing import Optional import apache_beam as beam from apache_beam.coders import BooleanCoder from apache_beam.coders import PickleCoder from apache_beam.coders import TimestampCoder from apache_beam.transforms.ptransform import PTransform from apache_beam.transforms.timeutil import TimeDomain from apache_beam.transforms.userstate import BagStateSpec from apache_beam.transforms.userstate import OrderedListStateSpec from apache_beam.transforms.userstate import ReadModifyWriteStateSpec from apache_beam.transforms.userstate import TimerSpec from apache_beam.transforms.userstate import on_timer from apache_beam.transforms.window import GlobalWindows from apache_beam.transforms.window import TimestampedValue from apache_beam.typehints.typehints import TupleConstraint from apache_beam.utils.timestamp import MAX_TIMESTAMP from apache_beam.utils.timestamp import MIN_TIMESTAMP from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import DurationTypes # pylint: disable=unused-import from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.timestamp import TimestampTypes # pylint: disable=unused-import _LOGGER = logging.getLogger("ordered_window_elements") """An example putting elements into window in time order on a streaming setting. The PTransform is a turn-key transform that can handle different input window settings and element types. Not only does it buffer elements, it can also prepend a window with the last seen element if the window is empty or there is a gap between the beginning of the window and the timestamp of its first element. """ class OrderedWindowElementsDoFn(beam.DoFn): """A Stateful DoFn that buffers and emits elements in time-ordered windows. This DoFn uses Beam's stateful processing capabilities to buffer elements and emit them in order within sliding windows. It handles out-of-order data, late data, and can fill starting gaps in windows by leveraging states and timers. Attributes: BUFFER_STATE: A `StateSpec` for storing incoming elements (timestamp, value) in a time-ordered buffer. WINDOW_TIMER: A `TimerSpec` set to the watermark time domain, used to trigger the emission of windowed elements. TIMER_STATE: A `ReadModifyWriteStateSpec` (BooleanCoder) to track whether the window timer has been initialized and set for the current key. LAST_VALUE: A `ReadModifyWriteStateSpec` (PickleCoder) to store the last emitted value for a key, used to fill the start of a window if there is a gap. BUFFER_MIN_TS_STATE: A `ReadModifyWriteStateSpec` (TimestampCoder) to keep track of the minimum timestamp currently present in the `buffer_state` for efficient clearing. ESTIMATED_WM_STATE: A `ReadModifyWriteStateSpec` (TimestampCoder) to store the highest observed timestamp for a key, used as an estimated watermark to detect and filter excessively late data. """ BUFFER_STATE = OrderedListStateSpec('buffer', PickleCoder()) WINDOW_TIMER = TimerSpec('window_timer', TimeDomain.WATERMARK) TIMER_STATE = ReadModifyWriteStateSpec('timer_state', BooleanCoder()) LAST_VALUE = ReadModifyWriteStateSpec('last_value', PickleCoder()) BUFFER_MIN_TS_STATE = ReadModifyWriteStateSpec( 'buffer_min_ts', TimestampCoder()) ESTIMATED_WM_STATE = ReadModifyWriteStateSpec( 'estimated_wm', TimestampCoder()) def __init__( self, duration: DurationTypes, slide_interval: DurationTypes, offset: DurationTypes, allowed_lateness: DurationTypes, default_start_value, fill_start_if_missing: bool, stop_timestamp: Optional[TimestampTypes]): """Initializes the OrderedWindowElementsFn. Args: duration: The duration of each window. slide_interval: The interval at which windows slide. offset: The offset of the window boundaries. Windows start at `offset` past each `duration` interval. allowed_lateness: The duration for which late data is still processed after the window's end. default_start_value: The default value to prepend or emit if a window is empty and `fill_start_if_missing` is true. fill_start_if_missing: A boolean indicating whether to prepend the last seen value to a window that has missing values at its start. stop_timestamp: An optional `Timestamp` at which to stop processing and firing timers for this key. """ self.duration = duration self.slide_interval = slide_interval self.offset = offset self.allowed_lateness = allowed_lateness self.default_start_value = default_start_value self.fill_start_if_missing = fill_start_if_missing self.stop_timestamp = stop_timestamp def start_bundle(self): _LOGGER.info("start bundle") def finish_bundle(self): _LOGGER.info("finish bundle") def _state_add(self, buffer_state, timestamp, value): """Add a timestamped-value into the buffer state.""" buffer_state.add((timestamp, value)) def _state_read_range(self, buffer_state, range_lo, range_hi): """Retrieves a specified range of elements from the buffer state.""" return list(buffer_state.read_range(range_lo, range_hi)) def _state_clear_range(self, buffer_state, range_lo, range_hi): """Clears a specified range of elements from the buffer state.""" # TODO: Dataflow runner v2 gets stuck when MIN_TIMESTAMP is used # as the lower bound for clear_range. Investigate this further. buffer_state.clear_range(range_lo, range_hi) def process( self, element, timestamp=beam.DoFn.TimestampParam, buffer_state=beam.DoFn.StateParam(BUFFER_STATE), window_timer=beam.DoFn.TimerParam(WINDOW_TIMER), timer_state=beam.DoFn.StateParam(TIMER_STATE), last_value_state=beam.DoFn.StateParam(LAST_VALUE), buffer_min_ts_state=beam.DoFn.StateParam(BUFFER_MIN_TS_STATE), estimated_wm_state=beam.DoFn.StateParam(ESTIMATED_WM_STATE), ): """Processes incoming elements, buffering them and setting timers. This method receives elements, updates the estimated watermark, buffers the element in `buffer_state`, and sets an initial window timer if one hasn't been set yet for the current key. It also handles the `fill_start_if_missing` logic for the `last_value_state`. Args: element: A `(key, value)` tuple representing the input element. timestamp: The event-time timestamp of the element. buffer_state: The `State` instance for buffering elements. window_timer: The `Timer` instance for scheduling window firings. timer_state: The `ReadModifyWriteState` instance to check/set if the timer has been initialized. last_value_state: The `ReadModifyWriteState` instance for the last emitted value, used for filling gaps. buffer_min_ts_state: The `ReadModifyWriteState` instance for the minimum timestamp in the buffer. estimated_wm_state: The `ReadModifyWriteState` instance for the estimated watermark. Returns: An empty list, as elements are emitted by the `on_timer` method, not directly by `process`. """ _, value = element _LOGGER.info( "[process] received element %s at timestamp %s", element, timestamp) estimated_wm = estimated_wm_state.read() if not estimated_wm or estimated_wm < timestamp: estimated_wm = timestamp estimated_wm_state.write(estimated_wm) else: # If the element is too late for the current watermark, drop it. if estimated_wm > timestamp + self.allowed_lateness: _LOGGER.info( "[process] data %s at %s is too late for watermark %s; dropping.", element, timestamp, estimated_wm) return [] buffer_min_ts = buffer_min_ts_state.read() if not buffer_min_ts or timestamp < buffer_min_ts: buffer_min_ts_state.write(timestamp) self._state_add(buffer_state, timestamp, value) timer_started = timer_state.read() if not timer_started: offset_duration = Duration.of(self.offset) slide_duration = Duration.of(self.slide_interval) duration_duration = Duration.of(self.duration) # Align the timestamp with the windowing scheme. aligned_micros = (timestamp - offset_duration).micros # Calculate the start of the last window that could contain this timestamp last_window_start_aligned_micros = ( (aligned_micros // slide_duration.micros) * slide_duration.micros) last_window_start = Timestamp( micros=last_window_start_aligned_micros) + offset_duration n = (duration_duration.micros - 1) // slide_duration.micros # Calculate the start of the first sliding window. first_slide_start_ts = last_window_start - Duration( micros=n * slide_duration.micros) # Set the initial timer to fire at the end of the first window plus # allowed lateness. first_window_end_ts = first_slide_start_ts + self.duration _LOGGER.info( "[process] setting initial timer to %s", first_window_end_ts + self.allowed_lateness) if (self.stop_timestamp is None or first_window_end_ts + self.allowed_lateness < self.stop_timestamp): window_timer.set(first_window_end_ts + self.allowed_lateness) timer_state.write(True) if self.fill_start_if_missing: last_value = last_value_state.read() if not last_value: last_value_state.write((MIN_TIMESTAMP, self.default_start_value)) return [] def _get_windowed_values_from_state( self, buffer_state, window_start_ts, window_end_ts, last_value_state): """Retrieves values for a window from the state, handling missing data. This helper method reads elements within a given window range from the buffer state. If `fill_start_if_missing` is enabled, it prepends the `last_value` if the window is initially empty or if there's a gap between the window start and the first element. It also updates the `last_value_state` with the last relevant element for the next window. Args: buffer_state: The state instance containing buffered elements. window_start_ts: The start timestamp of the window. window_end_ts: The end timestamp of the window. last_value_state: The `ReadModifyWriteState` instance storing the last emitted value. Returns: A list of `(timestamp, value)` tuples for the current window, potentially including a prepended last value if `fill_start_if_missing` is true. """ windowed_values = self._state_read_range( buffer_state, window_start_ts, window_end_ts) _LOGGER.info( "[on_timer] windowed data in buffer (%d): %s", len(windowed_values), windowed_values) if self.fill_start_if_missing: if not windowed_values: # If the window is empty, use the last value. last_value = last_value_state.read() value_to_insert = (window_start_ts, last_value[1]) windowed_values.append(value_to_insert) else: first_timestamp = windowed_values[0][0] last_value = last_value_state.read() if first_timestamp > window_start_ts and last_value: # Prepend the last value if there's a gap between the first element # in the window and the start of the window. value_to_insert = (window_start_ts, last_value[1]) windowed_values = [value_to_insert] + windowed_values # Find the last element before the beginning of the next window to update # last_value_state. i = 0 for v in windowed_values: if v[0] >= window_start_ts + self.slide_interval: break i += 1 if i > 0: last_value = windowed_values[i - 1] last_value_state.write(last_value) return windowed_values @on_timer(WINDOW_TIMER) def on_timer( self, key=beam.DoFn.KeyParam, fire_ts=beam.DoFn.TimestampParam, buffer_state=beam.DoFn.StateParam(BUFFER_STATE), window_timer=beam.DoFn.TimerParam(WINDOW_TIMER), last_value_state=beam.DoFn.StateParam(LAST_VALUE), buffer_min_ts_state=beam.DoFn.StateParam(BUFFER_MIN_TS_STATE), ): """Handles timer firings to emit windowed elements. When the `WINDOW_TIMER` fires, this method extracts elements for the current window from the `buffer_state`, handles late-firing windows (if `allowed_lateness` > 0), and emits them as a windowed `PCollection`. It also clears processed elements from the buffer and sets the next timer. Args: key: The key for which the timer fired. fire_ts: The event-time timestamp at which the timer fired. buffer_state: The `State` instance containing buffered elements. window_timer: The `Timer` instance for scheduling subsequent timers. last_value_state: The `ReadModifyWriteState` instance for the last emitted value. buffer_min_ts_state: The `ReadModifyWriteState` instance for the minimum timestamp in the buffer. Yields: `TimestampedValue`: A tuple `((key, window_start_ts, window_end_ts), list_of_values)` where `list_of_values` are the elements windowed and ordered, timestamped at `window_end_ts - 1`. """ _LOGGER.info("[on_timer] timer fired at %s", fire_ts) window_end_ts = fire_ts - self.allowed_lateness window_start_ts = window_end_ts - self.duration buffer_min_ts = buffer_min_ts_state.read() if not buffer_min_ts or buffer_min_ts > window_start_ts: buffer_min_ts = window_start_ts if self.allowed_lateness > 0: # Emit late windows that occurred prior to the current window. late_start_ts = window_start_ts while late_start_ts > buffer_min_ts: late_start_ts -= self.slide_interval while late_start_ts < window_start_ts: late_end_ts = late_start_ts + self.duration _LOGGER.info( "[on_timer] emitting late window: start=%s, end=%s", late_start_ts, late_end_ts) windowed_values = self._get_windowed_values_from_state( buffer_state, late_start_ts, late_end_ts, last_value_state) yield TimestampedValue( (key, ((late_start_ts, late_end_ts), windowed_values)), late_end_ts - 1) late_start_ts += self.slide_interval # Read and emit elements for the on-time window. _LOGGER.info( "[on_timer] emitting on-time window: start=%s, end=%s", window_start_ts, window_end_ts) windowed_values = self._get_windowed_values_from_state( buffer_state, window_start_ts, window_end_ts, last_value_state) yield TimestampedValue( (key, ((window_start_ts, window_end_ts), windowed_values)), window_end_ts - 1) # Post-emit actions for the current window: # - Compute the next window's start and end timestamps. # - Clean up states for expired windows. # - Set a new timer for the next window. next_window_end_ts = fire_ts - self.allowed_lateness + self.slide_interval next_window_start_ts = window_start_ts + self.slide_interval _LOGGER.info( "[on_timer] clearing timestamp range [%s, %s]", buffer_min_ts, next_window_start_ts) self._state_clear_range(buffer_state, buffer_min_ts, next_window_start_ts) buffer_min_ts_state.write(next_window_start_ts) _LOGGER.info( "[on_timer] setting follow-up timer to %s", next_window_end_ts + self.allowed_lateness) if (self.stop_timestamp is None or next_window_end_ts + self.allowed_lateness < self.stop_timestamp): window_timer.set(next_window_end_ts + self.allowed_lateness) class OrderedWindowElementsDoFnWithBag(OrderedWindowElementsDoFn): """The implementation of stateful Dofn with BagState as buffer state""" BUFFER_STATE = BagStateSpec('buffer', PickleCoder()) WINDOW_TIMER = TimerSpec('window_timer', TimeDomain.WATERMARK) def _state_add(self, buffer_state, timestamp, value): """Add a timestamped-value into the buffer state.""" buffer_state.add((timestamp, value)) def _state_read_range(self, buffer_state, range_lo, range_hi): """Retrieves a specified range of elements from the buffer state.""" all_elements = list(buffer_state.read()) filtered_elements = [(ts, val) for ts, val in all_elements if range_lo <= ts < range_hi] filtered_elements.sort(key=lambda x: x[0]) return filtered_elements def _state_clear_range(self, buffer_state, range_lo, range_hi): """Clears a specified range of elements from the buffer state.""" remaining_elements = self._state_read_range( buffer_state, range_hi, MAX_TIMESTAMP) buffer_state.clear() for e in remaining_elements: buffer_state.add(e) def process( self, element, timestamp=beam.DoFn.TimestampParam, buffer_state=beam.DoFn.StateParam(BUFFER_STATE), window_timer=beam.DoFn.TimerParam(WINDOW_TIMER), timer_state=beam.DoFn.StateParam(OrderedWindowElementsDoFn.TIMER_STATE), last_value_state=beam.DoFn.StateParam( OrderedWindowElementsDoFn.LAST_VALUE), buffer_min_ts_state=beam.DoFn.StateParam( OrderedWindowElementsDoFn.BUFFER_MIN_TS_STATE), estimated_wm_state=beam.DoFn.StateParam( OrderedWindowElementsDoFn.ESTIMATED_WM_STATE), ): yield from super().process( element, timestamp, buffer_state, window_timer, timer_state, last_value_state, buffer_min_ts_state, estimated_wm_state) @on_timer(WINDOW_TIMER) def on_timer( self, key=beam.DoFn.KeyParam, fire_ts=beam.DoFn.TimestampParam, buffer_state=beam.DoFn.StateParam(BUFFER_STATE), window_timer=beam.DoFn.TimerParam(WINDOW_TIMER), last_value_state=beam.DoFn.StateParam( OrderedWindowElementsDoFn.LAST_VALUE), buffer_min_ts_state=beam.DoFn.StateParam( OrderedWindowElementsDoFn.BUFFER_MIN_TS_STATE), ): yield from super().on_timer( key=key, fire_ts=fire_ts, buffer_state=buffer_state, window_timer=window_timer, last_value_state=last_value_state, buffer_min_ts_state=buffer_min_ts_state) class OrderedWindowElementsDoFnWithValue(OrderedWindowElementsDoFn): """The implementation of stateful Dofn with ValueState as buffer state""" BUFFER_STATE = ReadModifyWriteStateSpec('buffer', PickleCoder()) WINDOW_TIMER = TimerSpec('window_timer', TimeDomain.WATERMARK) def _state_add(self, buffer_state, timestamp, value): """Add a timestamped-value into the buffer state.""" buffer = buffer_state.read() or [] buffer.append((timestamp, value)) buffer_state.write(buffer) def _state_read_range(self, buffer_state, range_lo, range_hi): """Retrieves a specified range of elements from the buffer state.""" all_elements = buffer_state.read() filtered_elements = [(ts, val) for ts, val in all_elements if range_lo <= ts < range_hi] filtered_elements.sort(key=lambda x: x[0]) return filtered_elements def _state_clear_range(self, buffer_state, range_lo, range_hi): """Clears a specified range of elements from the buffer state.""" remaining_elements = self._state_read_range( buffer_state, range_hi, MAX_TIMESTAMP) buffer_state.write(remaining_elements) def process( self, element, timestamp=beam.DoFn.TimestampParam, buffer_state=beam.DoFn.StateParam(BUFFER_STATE), window_timer=beam.DoFn.TimerParam(WINDOW_TIMER), timer_state=beam.DoFn.StateParam(OrderedWindowElementsDoFn.TIMER_STATE), last_value_state=beam.DoFn.StateParam( OrderedWindowElementsDoFn.LAST_VALUE), buffer_min_ts_state=beam.DoFn.StateParam( OrderedWindowElementsDoFn.BUFFER_MIN_TS_STATE), estimated_wm_state=beam.DoFn.StateParam( OrderedWindowElementsDoFn.ESTIMATED_WM_STATE), ): yield from super().process( element, timestamp, buffer_state, window_timer, timer_state, last_value_state, buffer_min_ts_state, estimated_wm_state) @on_timer(WINDOW_TIMER) def on_timer( self, key=beam.DoFn.KeyParam, fire_ts=beam.DoFn.TimestampParam, buffer_state=beam.DoFn.StateParam(BUFFER_STATE), window_timer=beam.DoFn.TimerParam(WINDOW_TIMER), last_value_state=beam.DoFn.StateParam( OrderedWindowElementsDoFn.LAST_VALUE), buffer_min_ts_state=beam.DoFn.StateParam( OrderedWindowElementsDoFn.BUFFER_MIN_TS_STATE), ): yield from super().on_timer( key=key, fire_ts=fire_ts, buffer_state=buffer_state, window_timer=window_timer, last_value_state=last_value_state, buffer_min_ts_state=buffer_min_ts_state) class BufferStateType(enum.Enum): ORDERED_LIST = 0 BAG = 1 VALUE = 2 class OrderedWindowElements(PTransform): """A PTransform that batches elements into ordered, sliding windows. This transform processes elements with timestamps, buffering them and emitting them in fixed or sliding windows. It supports late data handling and can fill missing initial values in a window. """ def __init__( self, duration: DurationTypes, slide_interval: Optional[DurationTypes] = None, offset: DurationTypes = 0, allowed_lateness: DurationTypes = 0, default_start_value=None, fill_start_if_missing: bool = False, stop_timestamp: Optional[TimestampTypes] = None, buffer_state_type: BufferStateType = BufferStateType.ORDERED_LIST, ): """Initializes the OrderedWindowElements transform. Args: duration: The duration of each window. slide_interval: The interval at which windows slide. Defaults to `duration` if not provided (i.e., fixed windows). offset: The offset of the window boundaries. allowed_lateness: The maximum amount of time an element can be late and still be processed. default_start_value: The default value to use if `fill_start_if_missing` is true and a window is empty at its start. fill_start_if_missing: If true, the transform will attempt to fill the beginning of a window with the last known value if no elements are present at the window's start. stop_timestamp: An optional timestamp to stop processing and firing timers. buffer_state_type: An optional enum to control what backend state to use to store buffered elements. By default, it is using ordered list state. """ self.duration = duration self.slide_interval = duration if slide_interval is None else slide_interval self.offset = offset self.allowed_lateness = allowed_lateness self.default_start_value = default_start_value self.fill_start_if_missing = fill_start_if_missing self.stop_timestamp = stop_timestamp self.buffer_state_type = buffer_state_type def expand(self, input): """Applies the OrderedWindowElements transform to the input PCollection. The input PCollection is first ensured to be in `GlobalWindows`. If it's unkeyed, a default key is added. The `OrderedWindowElementsFn` is then applied. If the input was originally unkeyed, the default key is removed. Args: input: The input `PCollection`. Can be keyed (e.g., `PCollection[Tuple[K, V]]`) or unkeyed (e.g., `PCollection[V]`). Returns: A `PCollection` of `((key, window_start, window_end), list_of_values)` (if input was keyed) or `list_of_values` (if input was unkeyed), where `list_of_values` are the elements windowed and ordered. """ windowing = input.windowing if not isinstance(windowing.windowfn, GlobalWindows): _LOGGER.warning( 'Input PCollection is not in GlobalWindows. Overwriting windowing ' 'function with GlobalWindows.') input = input | "ToGlobalWindows" >> beam.WindowInto(GlobalWindows()) if isinstance(input.element_type, TupleConstraint): keyed_input = input else: # Add a default key (0) if the input PCollection is unkeyed. keyed_input = input | beam.WithKeys(0) if self.buffer_state_type == BufferStateType.ORDERED_LIST: dofn = OrderedWindowElementsDoFn elif self.buffer_state_type == BufferStateType.BAG: dofn = OrderedWindowElementsDoFnWithBag elif self.buffer_state_type == BufferStateType.VALUE: dofn = OrderedWindowElementsDoFnWithValue else: raise ValueError("Unknown buffer_state_type: " + self.buffer_state_type) keyed_output = ( keyed_input | 'Ordered Sliding Window' >> beam.ParDo( dofn( self.duration, self.slide_interval, self.offset, self.allowed_lateness, self.default_start_value, self.fill_start_if_missing, self.stop_timestamp))) if isinstance(input.element_type, TupleConstraint): ret = keyed_output else: # Remove the default key if the input PCollection was originally unkeyed. ret = keyed_output | beam.Values() return ret ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/cookbook/ordered_window_elements/streaming_test.py0000644000175100017510000003455315206113032031624 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import shutil import sys import unittest from parameterized import param from parameterized import parameterized from parameterized import parameterized_class import apache_beam as beam from apache_beam.examples.cookbook.ordered_window_elements.streaming import BufferStateType from apache_beam.examples.cookbook.ordered_window_elements.streaming import OrderedWindowElements from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.periodicsequence import PeriodicImpulse from apache_beam.transforms.periodicsequence import RebaseMode from apache_beam.utils.timestamp import Timestamp logging.basicConfig(level=logging.WARNING) ENABLE_LOGGING = False WINDOW_SIZE = 3 FIRE_INTERVAL = 0.5 def _maybe_log_elements(pcoll, prefix="result="): if ENABLE_LOGGING: return pcoll | beam.LogElements( prefix=prefix, level=logging.WARNING, with_timestamp=True, with_window=True, use_epoch_time=True) else: return pcoll # Creates an unbounded source via `PeriodicImpulse`, simulating a continuous # stream of elements fired at a fixed interval. This method is closer to # real-world streaming but is sensitive to system load and can cause test # flakiness. # If the test runner is slow or under heavy load, elements may be delayed and # processed in a single large bundle. This can defeat the purpose of testing # time-based logic, as the elements will not arrive distributed over time as # intended. def _create_periodic_impulse_stream(elements: list[int]): now = Timestamp.now() length = len(elements) fire_interval = FIRE_INTERVAL return PeriodicImpulse( data=[(Timestamp.of(e), e) for e in elements], fire_interval=fire_interval, start_timestamp=now, stop_timestamp=now + length * fire_interval, rebase=RebaseMode.REBASE_ALL, ) # Creates an unbounded source via `TestStream`, allowing precise control over # watermarks and element emission for deterministic testing scenarios. However, # it is an instantaneous data stream and it is less realistic than the stream # from `PeriodicImpulse`. def _create_test_stream(elements: list[int]): test_stream = TestStream() wm = None for e in elements: test_stream.add_elements([e], event_timestamp=e) if wm is None or wm < e: wm = e test_stream.advance_watermark_to(wm) test_stream.advance_watermark_to_infinity() return test_stream def _convert_timestamp_to_int(has_key=False): if has_key: return beam.MapTuple( lambda key, value: ( key, ((int(value[0][0].micros // 1e6), int(value[0][1].micros // 1e6)), [(int(t.micros // 1e6), v) for t, v in value[1]]))) return beam.MapTuple( lambda window, elements: ((int(window[0].micros // 1e6), int(window[1].micros // 1e6)), [(int(t.micros // 1e6), v) for t, v in elements])) _go_installed = shutil.which('go') is not None _in_windows = sys.platform == "win32" @unittest.skipUnless(_go_installed, 'Go is not installed.') # TODO: Go environments is not configured correctly on Windows test boxes. @unittest.skipIf(_in_windows, reason="Not supported on Windows") @parameterized_class( 'buffer_state_type', [ (BufferStateType.ORDERED_LIST, ), (BufferStateType.BAG, ), (BufferStateType.VALUE, ), ]) class OrderedWindowElementsTest(unittest.TestCase): def setUp(self) -> None: self.options = PipelineOptions([ "--streaming", "--environment_type=LOOPBACK", "--runner=PrismRunner", "--prism_log_kind=dev", # # run on an external Portable Runner for debugging # "--runner=PortableRunner", # "--job_endpoint=localhost:8073", ]) # # dataflow runner option # self.options = PipelineOptions([ # "--streaming", # "--runner=DataflowRunner", # "--temp_location=gs://shunping-test/anomaly-temp", # "--staging_location=gs://shunping-test/anomaly-temp", # "--project=apache-beam-testing", # "--region=us-central1", # "--sdk_location=dist/apache_beam-2.69.0.dev0.tar.gz", # #"--pickle_library=dill", # #"--save_main_session", # ]) def test_default(self): with TestPipeline(options=self.options) as p: result = ( p | _create_test_stream([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) | OrderedWindowElements( WINDOW_SIZE, stop_timestamp=13, buffer_state_type=self.buffer_state_type)) result = _maybe_log_elements(result) | _convert_timestamp_to_int() assert_that( result, equal_to([ ((0, 3), [(0, 0), (1, 1), (2, 2)]), ((3, 6), [(3, 3), (4, 4), (5, 5)]), ((6, 9), [(6, 6), (7, 7), (8, 8)]), ((9, 12), [(9, 9)]), ])) def test_offset(self): with TestPipeline(options=self.options) as p: result = ( p | _create_test_stream([2, 3, 4, 5, 6, 7, 8, 9]) | OrderedWindowElements(WINDOW_SIZE, stop_timestamp=13, offset=2)) result = _maybe_log_elements(result) | _convert_timestamp_to_int() assert_that( result, equal_to([ ((2, 5), [(2, 2), (3, 3), (4, 4)]), # window start at 2 ((5, 8), [(5, 5), (6, 6), (7, 7)]), ((8, 11), [(8, 8), (9, 9)]) ])) def test_slide_interval(self): with TestPipeline(options=self.options) as p: result = ( p | _create_test_stream([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) | OrderedWindowElements(WINDOW_SIZE, 1, stop_timestamp=13)) result = _maybe_log_elements(result) assert_that( result, equal_to([ ((-2, 1), [(0, 0)]), ((-1, 2), [(0, 0), (1, 1)]), ((0, 3), [(0, 0), (1, 1), (2, 2)]), ((1, 4), [(1, 1), (2, 2), (3, 3)]), ((2, 5), [(2, 2), (3, 3), (4, 4)]), ((3, 6), [(3, 3), (4, 4), (5, 5)]), ((4, 7), [(4, 4), (5, 5), (6, 6)]), ((5, 8), [(5, 5), (6, 6), (7, 7)]), ((6, 9), [(6, 6), (7, 7), (8, 8)]), ((7, 10), [(7, 7), (8, 8), (9, 9)]), ((8, 11), [(8, 8), (9, 9)]), ((9, 12), [(9, 9)]), ])) def test_keyed_input(self): with TestPipeline(options=self.options) as p: result = ( p | _create_test_stream([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) | beam.WithKeys("my_key") # key is present in the output | OrderedWindowElements(WINDOW_SIZE, stop_timestamp=13)) result = _maybe_log_elements(result) | _convert_timestamp_to_int( has_key=True) assert_that( result, equal_to([ ("my_key", ((0, 3), [(1, 1), (2, 2)])), ("my_key", ((3, 6), [(3, 3), (4, 4), (5, 5)])), ("my_key", ((6, 9), [(6, 6), (7, 7), (8, 8)])), ("my_key", ((9, 12), [(9, 9), (10, 10)])), ])) @parameterized.expand([ param(fill_window_start=False), param(fill_window_start=True), ]) def test_non_zero_offset_and_default_value(self, fill_window_start): if fill_window_start: expected = [ # window [-2, 1), and the start is filled with default value ((-2, 1), [(-2, -100), (0, 0)]), ((1, 4), [(1, 1), (2, 2), (3, 3)]), # window [1, 4) ((4, 7), [(4, 4), (5, 5), (6, 6)]), ((7, 10), [(7, 7), (8, 8), (9, 9)]), ] else: expected = [ ((-2, 1), [(0, 0)]), # window [-2, 1) ((1, 4), [(1, 1), (2, 2), (3, 3)]), # window [1, 4) ((4, 7), [(4, 4), (5, 5), (6, 6)]), ((7, 10), [(7, 7), (8, 8), (9, 9)]), ] with TestPipeline(options=self.options) as p: result = ( p | _create_test_stream([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) | OrderedWindowElements( WINDOW_SIZE, offset=1, default_start_value=-100, fill_start_if_missing=fill_window_start, stop_timestamp=13)) result = _maybe_log_elements(result) | _convert_timestamp_to_int() assert_that(result, equal_to(expected)) @parameterized.expand([ param(fill_window_start=False), param(fill_window_start=True), ]) def test_ordered_data_with_gap(self, fill_window_start): if fill_window_start: expected = [ ((0, 3), [(0, 0), (1, 1), (2, 2)]), ((3, 6), [(3, 3), (4, 4)]), # window [6, 9) is empty, so the start is filled with last value. ((6, 9), [(6, 4)]), # window [9, 12) is empty, so the start is filled with last value. ((9, 12), [(9, 4)]), # window [12, 15) is empty, so the start is filled with last value. ((12, 15), [(12, 4)]), ((15, 18), [(15, 4), (16, 16), (17, 17)]), ((18, 21), [(18, 18), (19, 19), (20, 20)]) ] else: expected = [ ((0, 3), [(0, 0), (1, 1), (2, 2)]), ((3, 6), [(3, 3), (4, 4)]), ((6, 9), []), # window [6, 9) is empty ((9, 12), []), # window [9, 12) is empty ((12, 15), []), # window [12, 15) is empty ((15, 18), [(16, 16), (17, 17)]), ((18, 21), [(18, 18), (19, 19), (20, 20)]) ] with TestPipeline(options=self.options) as p: result = ( p | _create_test_stream([0, 1, 2, 3, 4, 16, 17, 18, 19, 20]) | OrderedWindowElements( WINDOW_SIZE, fill_start_if_missing=fill_window_start, stop_timestamp=23)) result = _maybe_log_elements(result) | _convert_timestamp_to_int() assert_that(result, equal_to(expected)) def test_single_late_data_with_no_allowed_lateness(self): with TestPipeline(options=self.options) as p: result = ( p | _create_test_stream([0, 1, 2, 3, 4, 6, 7, 8, 9, 5]) | OrderedWindowElements(WINDOW_SIZE, stop_timestamp=13)) result = _maybe_log_elements(result) | _convert_timestamp_to_int() assert_that( result, equal_to([ ((0, 3), [(0, 0), (1, 1), (2, 2)]), ((3, 6), [(3, 3), (4, 4)]), # 5 is late and discarded ((6, 9), [(6, 6), (7, 7), (8, 8)]), ((9, 12), [(9, 9)]), ])) def test_single_late_data_with_allowed_lateness(self): with TestPipeline(options=self.options) as p: result = ( p | _create_test_stream([0, 1, 2, 3, 4, 6, 7, 8, 9, 5]) | OrderedWindowElements( WINDOW_SIZE, allowed_lateness=4, stop_timestamp=17)) result = _maybe_log_elements(result) | _convert_timestamp_to_int() assert_that( result, equal_to([ ((0, 3), [(0, 0), (1, 1), (2, 2)]), # allow late data up to: # 9 (watermark before late data) - 4 (allowed lateness) = 5 ((3, 6), [(3, 3), (4, 4), (5, 5)]), ((6, 9), [(6, 6), (7, 7), (8, 8)]), ((9, 12), [(9, 9)]), ])) @parameterized.expand([ param(fill_start=False), param(fill_start=True), ]) def test_reversed_ordered_data_with_allowed_lateness(self, fill_start): if fill_start: expected = [ # allow late data up to: # 9 (watermark before late data) - 5 (allowed lateness) = 4 ((3, 6), [(3, None), (4, 4), (5, 5)]), ((6, 9), [(6, 6), (7, 7), (8, 8)]), ((9, 12), [(9, 9)]), ((12, 15), [(12, 9)]), ((15, 18), [(15, 9)]), ] else: expected = [ ((3, 6), [(4, 4), (5, 5)]), ((6, 9), [(6, 6), (7, 7), (8, 8)]), ((9, 12), [(9, 9)]), ((12, 15), []), ((15, 18), []), ] with TestPipeline(options=self.options) as p: result = ( p | _create_test_stream([9, 8, 7, 6, 5, 4, 3, 2, 1, 0]) | OrderedWindowElements( WINDOW_SIZE, fill_start_if_missing=fill_start, allowed_lateness=5, stop_timestamp=25)) result = _maybe_log_elements(result) | _convert_timestamp_to_int() assert_that(result, equal_to(expected)) def test_multiple_late_data_with_allowed_lateness(self): with TestPipeline(options=self.options) as p: result = ( p | _create_test_stream([1, 2, 9, 3, 14, 7, 5, 12, 16, 17]) | OrderedWindowElements( WINDOW_SIZE, 1, allowed_lateness=6, fill_start_if_missing=True, stop_timestamp=28)) result = _maybe_log_elements(result) | _convert_timestamp_to_int() # yapf: disable assert_that( result, equal_to([ ((-1, 2), [(-1, None), (1, 1)]), ((0, 3), [(0, None), (1, 1), (2, 2)]), ((1, 4), [(1, 1), (2, 2), (3, 3)]), ((2, 5), [(2, 2), (3, 3)]), ((3, 6), [(3, 3)]), ((4, 7), [(4, 3)]), ((5, 8), [(5, 3)]), ((6, 9), [(6, 3)]), ((7, 10), [(7, 3), (9, 9)]), ((8, 11), [(8, 3), (9, 9)]), ((9, 12), [(9, 9)]), ((10, 13), [(10, 9), (12, 12)]), ((11, 14), [(11, 9), (12, 12)]), ((12, 15), [(12, 12), (14, 14)]), ((13, 16), [(13, 12), (14, 14)]), ((14, 17), [(14, 14), (16, 16)]), ((15, 18), [(15, 14), (16, 16),(17, 17)]), ((16, 19), [(16, 16), (17, 17)]), ((17, 20), [(17, 17)]), ((18, 21), [(18, 17)]) ])) # yapf: enable if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.049373 apache_beam-2.74.0/apache_beam/examples/dataframe/0000755000175100017510000000000015206113157021427 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/dataframe/__init__.py0000644000175100017510000000142015206113032023525 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/dataframe/flight_delays.py0000644000175100017510000001126615206113032024615 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A pipeline using dataframes to compute typical flight delay times.""" # pytype: skip-file from __future__ import absolute_import import argparse import logging import apache_beam as beam from apache_beam.dataframe.convert import to_dataframe from apache_beam.options.pipeline_options import PipelineOptions def get_mean_delay_at_top_airports(airline_df): arr = airline_df.rename(columns={ 'arrival_airport': 'airport' }).airport.value_counts() dep = airline_df.rename(columns={ 'departure_airport': 'airport' }).airport.value_counts() total = arr + dep # Note we keep all to include duplicates - this ensures the result is # deterministic. # NaNs can be included in the output in pandas 1.4.0 and above, so we # explicitly drop them. top_airports = total.nlargest(10, keep='all').dropna() at_top_airports = airline_df['arrival_airport'].isin( top_airports.index.values) return airline_df[at_top_airports].mean(numeric_only=True) def input_date(date): import datetime parsed = datetime.datetime.strptime(date, '%Y-%m-%d') if (parsed > datetime.datetime(2012, 12, 31) or parsed < datetime.datetime(2002, 1, 1)): raise ValueError("There's only data from 2002-01-01 to 2012-12-31") return date def run_flight_delay_pipeline( pipeline, start_date=None, end_date=None, output=None): query = f""" SELECT FlightDate AS date, IATA_CODE_Reporting_Airline AS airline, Origin AS departure_airport, Dest AS arrival_airport, DepDelay AS departure_delay, ArrDelay AS arrival_delay FROM `apache-beam-testing.airline_ontime_data.flights` WHERE FlightDate >= '{start_date}' AND FlightDate <= '{end_date}' AND DepDelay IS NOT NULL AND ArrDelay IS NOT NULL """ # Import this here to avoid pickling the main session. import time from apache_beam import window def to_unixtime(s): return time.mktime(s.timetuple()) # The pipeline will be run on exiting the with block. with pipeline as p: tbl = ( p | 'read table' >> beam.io.ReadFromBigQuery( query=query, use_standard_sql=True) | 'assign timestamp' >> beam.Map(lambda x: window.TimestampedValue(x, to_unixtime(x['date']))) # Use beam.Select to make sure data has a schema # The casts in lambdas ensure data types are properly inferred | 'set schema' >> beam.Select( date=lambda x: str(x['date']), airline=lambda x: str(x['airline']), departure_airport=lambda x: str(x['departure_airport']), arrival_airport=lambda x: str(x['arrival_airport']), departure_delay=lambda x: float(x['departure_delay']), arrival_delay=lambda x: float(x['arrival_delay']))) daily = tbl | 'daily windows' >> beam.WindowInto( beam.window.FixedWindows(60 * 60 * 24)) # group the flights data by carrier df = to_dataframe(daily) result = df.groupby('airline').apply(get_mean_delay_at_top_airports) result.to_csv(output) def run(argv=None): """Main entry point; defines and runs the flight delay pipeline.""" parser = argparse.ArgumentParser() parser.add_argument( '--start_date', dest='start_date', type=input_date, default='2012-12-22', help='YYYY-MM-DD lower bound (inclusive) for input dataset.') parser.add_argument( '--end_date', dest='end_date', type=input_date, default='2012-12-26', help='YYYY-MM-DD upper bound (inclusive) for input dataset.') parser.add_argument( '--output', dest='output', required=True, help='Location to write the output.') known_args, pipeline_args = parser.parse_known_args(argv) run_flight_delay_pipeline( beam.Pipeline(options=PipelineOptions(pipeline_args)), start_date=known_args.start_date, end_date=known_args.end_date, output=known_args.output) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/dataframe/flight_delays_it_test.py0000644000175100017510000001255115206113032026346 0ustar00runnerrunner# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the flight delay example.""" # pytype: skip-file from __future__ import absolute_import import logging import os import unittest import uuid import pandas as pd import pytest from apache_beam.examples.dataframe import flight_delays from apache_beam.io.filesystems import FileSystems from apache_beam.testing.test_pipeline import TestPipeline class FlightDelaysTest(unittest.TestCase): EXPECTED = { '2012-12-23': [ ('AA', 20.082559339525282, 12.825593395252838), ('AS', 5.0456273764258555, 1.0722433460076046), ('B6', 20.646569646569645, 16.405405405405407), ('DL', 5.241148325358852, -3.2401913875598085), ('EV', 9.982053838484546, 4.40777666999003), ('F9', 23.67883211678832, 25.27007299270073), ('FL', 4.4602272727272725, -0.8352272727272727), ('HA', -1.0829015544041452, 0.010362694300518135), ('MQ', 8.912912912912914, 3.6936936936936937), ('OO', 30.526699029126213, 31.17961165048544), ('UA', 19.142555438225976, 11.07180570221753), ('US', 3.092541436464088, -2.350828729281768), ('VX', 62.755102040816325, 62.61224489795919), ('WN', 12.05824508320726, 6.713313161875946), ('YV', 16.155844155844157, 13.376623376623376), ], '2012-12-24': [ ('AA', 7.049086757990867, -1.5970319634703196), ('AS', 0.5917602996254682, -2.2659176029962547), ('B6', 8.070993914807302, 2.73630831643002), ('DL', 3.700745473908413, -2.2396166134185305), ('EV', 7.322115384615385, 2.3653846153846154), ('F9', 13.786764705882351, 15.5), ('FL', 2.416909620991253, 2.224489795918368), ('HA', -2.6785714285714284, -2.4744897959183674), ('MQ', 15.818181818181818, 9.935828877005347), ('OO', 10.902374670184695, 10.08575197889182), ('UA', 10.935406698564593, -1.3337320574162679), ('US', 1.369281045751634, -1.4101307189542485), ('VX', 3.841666666666667, -2.4166666666666665), ('WN', 7.3715753424657535, 0.348458904109589), ('YV', 0.32, 0.78), ], '2012-12-25': [ ('AA', 23.551581843191197, 35.62585969738652), ('AS', 3.4816326530612245, 0.27346938775510204), ('B6', 9.10590631364562, 3.989816700610998), ('DL', 2.2863795110593714, -3.668218859138533), ('EV', 17.35576923076923, 16.414835164835164), ('F9', 19.38, 21.786666666666665), ('FL', 1.3823529411764706, 0.9205882352941176), ('HA', -4.725806451612903, -3.9946236559139785), ('MQ', 32.527716186252775, 44.148558758314856), ('OO', 15.788595271210012, 16.617524339360223), ('UA', 16.663145539906104, 10.772300469483568), ('US', 2.7953216374269005, 0.2236842105263158), ('VX', 23.62878787878788, 23.636363636363637), ('WN', 14.423791821561338, 10.142193308550183), ('YV', 11.256302521008404, 11.659663865546218), ], } def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.outdir = ( self.test_pipeline.get_option('temp_location') + '/flight_delays_it-' + str(uuid.uuid4())) self.output_path = os.path.join(self.outdir, 'output.csv') def tearDown(self): FileSystems.delete([self.outdir + '/']) @pytest.mark.examples_postcommit @pytest.mark.it_postcommit def test_flight_delays(self): flight_delays.run_flight_delay_pipeline( self.test_pipeline, start_date='2012-12-23', end_date='2012-12-25', output=self.output_path) def read_csv(path): with FileSystems.open(path) as fp: return pd.read_csv(fp) # Parse result file and compare. for date, expectation in self.EXPECTED.items(): result_df = pd.concat( read_csv(metadata.path) for metadata in FileSystems.match( [f'{self.output_path}-{date}*'])[0].metadata_list) result_df = result_df.sort_values('airline').reset_index(drop=True) expected_df = pd.DataFrame( expectation, columns=['airline', 'departure_delay', 'arrival_delay']) expected_df = expected_df.sort_values('airline').reset_index(drop=True) try: pd.testing.assert_frame_equal(result_df, expected_df) except AssertionError as e: raise AssertionError( f"date={date!r} result DataFrame:\n\n" f"{result_df}\n\n" "Differs from Expectation:\n\n" f"{expected_df}") from e if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/dataframe/taxiride.py0000644000175100017510000001041615206113032023604 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Pipelines that use the DataFrame API to process NYC taxiride CSV data.""" # pytype: skip-file from __future__ import absolute_import import argparse import logging import apache_beam as beam from apache_beam.dataframe.io import read_csv from apache_beam.options.pipeline_options import PipelineOptions ZONE_LOOKUP_PATH = ( "gs://apache-beam-samples/nyc_taxi/misc/taxi+_zone_lookup.csv") def run_aggregation_pipeline(pipeline, input_path, output_path): # The pipeline will be run on exiting the with block. # [START DataFrame_taxiride_aggregation] with pipeline as p: rides = p | read_csv(input_path) # Count the number of passengers dropped off per LocationID agg = rides.groupby('DOLocationID').passenger_count.sum() agg.to_csv(output_path) # [END DataFrame_taxiride_aggregation] def run_enrich_pipeline( pipeline, input_path, output_path, zone_lookup_path=ZONE_LOOKUP_PATH): """Enrich taxi ride data with zone lookup table and perform a grouped aggregation.""" # The pipeline will be run on exiting the with block. # [START DataFrame_taxiride_enrich] with pipeline as p: rides = p | "Read taxi rides" >> read_csv(input_path) zones = p | "Read zone lookup" >> read_csv(zone_lookup_path) # Enrich taxi ride data with boroughs from zone lookup table # Joins on zones.LocationID and rides.DOLocationID, by first making the # former the index for zones. rides = rides.merge( zones.set_index('LocationID').Borough, right_index=True, left_on='DOLocationID', how='left') # Sum passengers dropped off per Borough agg = rides.groupby('Borough').passenger_count.sum() agg.to_csv(output_path) # [END DataFrame_taxiride_enrich] # A more intuitive alternative to the above merge call, but this option # doesn't preserve index, thus requires non-parallel execution. #rides = rides.merge(zones[['LocationID','Borough']], # how="left", # left_on='DOLocationID', # right_on='LocationID') def run(argv=None): """Main entry point.""" parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument( '--input', dest='input', default='gs://apache-beam-samples/nyc_taxi/misc/sample.csv', help='Input file to process.') parser.add_argument( '--output', dest='output', required=True, help='Output file to write results to.') parser.add_argument( '--zone_lookup', dest='zone_lookup_path', default=ZONE_LOOKUP_PATH, help='Location for taxi zone lookup CSV.') parser.add_argument( '--pipeline', dest='pipeline', default='location_id_agg', help=( "Choice of pipeline to run. Must be one of " "(location_id_agg, borough_enrich).")) known_args, pipeline_args = parser.parse_known_args(argv) pipeline = beam.Pipeline(options=PipelineOptions(pipeline_args)) if known_args.pipeline == 'location_id_agg': run_aggregation_pipeline(pipeline, known_args.input, known_args.output) elif known_args.pipeline == 'borough_enrich': run_enrich_pipeline( pipeline, known_args.input, known_args.output, known_args.zone_lookup_path) else: raise ValueError( f"Unrecognized value for --pipeline: {known_args.pipeline!r}. " "Must be one of ('location_id_agg', 'borough_enrich')") if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/dataframe/taxiride_it_test.py0000644000175100017510000000660415206113032025343 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end tests for the taxiride examples.""" # pytype: skip-file import logging import os import unittest import uuid import pandas as pd import pytest from apache_beam.examples.dataframe import taxiride from apache_beam.io.filesystems import FileSystems from apache_beam.options.pipeline_options import WorkerOptions from apache_beam.testing.test_pipeline import TestPipeline class TaxirideIT(unittest.TestCase): def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.outdir = ( self.test_pipeline.get_option('temp_location') + '/taxiride_it-' + str(uuid.uuid4())) self.output_path = os.path.join(self.outdir, 'output.csv') def tearDown(self): FileSystems.delete([self.outdir + '/']) @pytest.mark.it_postcommit def test_aggregation(self): taxiride.run_aggregation_pipeline( self.test_pipeline, 'gs://apache-beam-samples/nyc_taxi/2018/*.csv', self.output_path) # Verify expected = pd.read_csv( os.path.join( os.path.dirname(__file__), 'data', 'taxiride_2018_aggregation_truth.csv'), comment='#') expected = expected.sort_values('DOLocationID').reset_index(drop=True) def read_csv(path): with FileSystems.open(path) as fp: return pd.read_csv(fp) result = pd.concat( read_csv(metadata.path) for metadata in FileSystems.match( [f'{self.output_path}*'])[0].metadata_list) result = result.sort_values('DOLocationID').reset_index(drop=True) pd.testing.assert_frame_equal(expected, result) @pytest.mark.it_postcommit def test_enrich(self): # Standard workers OOM with the enrich pipeline self.test_pipeline.get_pipeline_options().view_as( WorkerOptions).machine_type = 'e2-highmem-2' taxiride.run_enrich_pipeline( self.test_pipeline, 'gs://apache-beam-samples/nyc_taxi/2018/*.csv', self.output_path) # Verify expected = pd.read_csv( os.path.join( os.path.dirname(__file__), 'data', 'taxiride_2018_enrich_truth.csv'), comment='#') expected = expected.sort_values('Borough').reset_index(drop=True) def read_csv(path): with FileSystems.open(path) as fp: return pd.read_csv(fp) result = pd.concat( read_csv(metadata.path) for metadata in FileSystems.match( [f'{self.output_path}*'])[0].metadata_list) result = result.sort_values('Borough').reset_index(drop=True) pd.testing.assert_frame_equal(expected, result) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/dataframe/taxiride_test.py0000644000175100017510000001262515206113032024647 0ustar00runnerrunner# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the taxiride example pipelines.""" # pytype: skip-file from __future__ import absolute_import import glob import logging import os import re import tempfile import unittest import pandas as pd import apache_beam as beam from apache_beam.examples.dataframe import taxiride from apache_beam.testing.util import open_shards class TaxiRideExampleTest(unittest.TestCase): # First 10 lines from gs://apache-beam-samples/nyc_taxi/misc/sample.csv # pylint: disable=line-too-long SAMPLE_RIDES = """VendorID,tpep_pickup_datetime,tpep_dropoff_datetime,passenger_count,trip_distance,RatecodeID,store_and_fwd_flag,PULocationID,DOLocationID,payment_type,fare_amount,extra,mta_tax,tip_amount,tolls_amount,improvement_surcharge,total_amount,congestion_surcharge 1,2019-01-01 00:46:40,2019-01-01 00:53:20,1,1.50,1,N,151,239,1,7,0.5,0.5,1.65,0,0.3,9.95, 1,2019-01-01 00:59:47,2019-01-01 01:18:59,1,2.60,1,N,239,246,1,14,0.5,0.5,1,0,0.3,16.3, 2,2018-12-21 13:48:30,2018-12-21 13:52:40,3,.00,1,N,236,236,1,4.5,0.5,0.5,0,0,0.3,5.8, 2,2018-11-28 15:52:25,2018-11-28 15:55:45,5,.00,1,N,193,193,2,3.5,0.5,0.5,0,0,0.3,7.55, 2,2018-11-28 15:56:57,2018-11-28 15:58:33,5,.00,2,N,193,193,2,52,0,0.5,0,0,0.3,55.55, 2,2018-11-28 16:25:49,2018-11-28 16:28:26,5,.00,1,N,193,193,2,3.5,0.5,0.5,0,5.76,0.3,13.31, 2,2018-11-28 16:29:37,2018-11-28 16:33:43,5,.00,2,N,193,193,2,52,0,0.5,0,0,0.3,55.55, 1,2019-01-01 00:21:28,2019-01-01 00:28:37,1,1.30,1,N,163,229,1,6.5,0.5,0.5,1.25,0,0.3,9.05, 1,2019-01-01 00:32:01,2019-01-01 00:45:39,1,3.70,1,N,229,7,1,13.5,0.5,0.5,3.7,0,0.3,18.5 """ # pylint: enable=line-too-long SAMPLE_ZONE_LOOKUP = """"LocationID","Borough","Zone","service_zone" 7,"Queens","Astoria","Boro Zone" 193,"Queens","Queensbridge/Ravenswood","Boro Zone" 229,"Manhattan","Sutton Place/Turtle Bay North","Yellow Zone" 236,"Manhattan","Upper East Side North","Yellow Zone" 239,"Manhattan","Upper West Side South","Yellow Zone" 246,"Manhattan","West Chelsea/Hudson Yards","Yellow Zone" """ def setUp(self): self.tmpdir = tempfile.TemporaryDirectory() self.input_path = os.path.join(self.tmpdir.name, 'rides*.csv') self.lookup_path = os.path.join(self.tmpdir.name, 'lookup.csv') self.output_path = os.path.join(self.tmpdir.name, 'output.csv') # Duplicate sample data in 100 different files to replicate multi-file read for i in range(100): with open(os.path.join(self.tmpdir.name, f'rides{i}.csv'), 'w') as fp: fp.write(self.SAMPLE_RIDES) with open(self.lookup_path, 'w') as fp: fp.write(self.SAMPLE_ZONE_LOOKUP) def tearDown(self): self.tmpdir.cleanup() def test_aggregation(self): # Compute expected result rides = pd.concat(pd.read_csv(path) for path in glob.glob(self.input_path)) expected_counts = rides.groupby('DOLocationID').passenger_count.sum() taxiride.run_aggregation_pipeline( beam.Pipeline(), self.input_path, self.output_path) # Parse result file and compare. # TODO(https://github.com/apache/beam/issues/20926): taxiride examples # should produce int sums, not floats results = [] with open_shards(f'{self.output_path}-*') as result_file: for line in result_file: match = re.search(r'(\S+),([0-9\.]+)', line) if match is not None: results.append((int(match.group(1)), int(float(match.group(2))))) elif line.strip(): self.assertEqual(line.strip(), 'DOLocationID,passenger_count') self.assertEqual(sorted(results), sorted(expected_counts.items())) def test_enrich(self): # Compute expected result rides = pd.concat(pd.read_csv(path) for path in glob.glob(self.input_path)) zones = pd.read_csv(self.lookup_path) rides = rides.merge( zones.set_index('LocationID').Borough, right_index=True, left_on='DOLocationID', how='left') expected_counts = rides.groupby('Borough').passenger_count.sum() taxiride.run_enrich_pipeline( beam.Pipeline(), self.input_path, self.output_path, self.lookup_path) # Parse result file and compare. # TODO(BEAM-XXXX): taxiride examples should produce int sums, not floats results = [] with open_shards(f'{self.output_path}-*') as result_file: for line in result_file: match = re.search(r'(\S+),([0-9\.]+)', line) if match is not None: results.append((match.group(1), int(float(match.group(2))))) elif line.strip(): self.assertEqual(line.strip(), 'Borough,passenger_count') self.assertEqual(sorted(results), sorted(expected_counts.items())) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/dataframe/wordcount.py0000644000175100017510000000555515206113032024027 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A word-counting workflow using the DataFrame API.""" # pytype: skip-file import argparse import logging import apache_beam as beam from apache_beam.dataframe.convert import to_dataframe from apache_beam.dataframe.convert import to_pcollection from apache_beam.io import ReadFromText from apache_beam.options.pipeline_options import PipelineOptions def run(argv=None): """Main entry point; defines and runs the wordcount pipeline.""" parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument( '--input', dest='input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='Input file to process.') parser.add_argument( '--output', dest='output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) # Import this here to avoid pickling the main session. import re # The pipeline will be run on exiting the with block. with beam.Pipeline(options=PipelineOptions(pipeline_args)) as p: # [START DataFrame_wordcount] # Read the text file[pattern] into a PCollection. lines = p | 'Read' >> ReadFromText(known_args.input) words = ( lines | 'Split' >> beam.FlatMap( lambda line: re.findall(r'[\w]+', line)).with_output_types(str) # Map to Row objects to generate a schema suitable for conversion # to a dataframe. | 'ToRows' >> beam.Map(lambda word: beam.Row(word=word))) df = to_dataframe(words) df['count'] = 1 counted = df.groupby('word').sum() counted.to_csv(known_args.output) # Deferred DataFrames can also be converted back to schema'd PCollections counted_pc = to_pcollection(counted, include_indexes=True) # [END DataFrame_wordcount] # Print out every word that occurred >50 times _ = ( counted_pc | beam.Filter(lambda row: row.count > 50) | beam.Map(lambda row: f'{row.word}: {row.count}') | beam.Map(print)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/dataframe/wordcount_test.py0000644000175100017510000000465015206113032025061 0ustar00runnerrunner# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the wordcount example.""" # pytype: skip-file import collections import logging import re import unittest import uuid import pytest from apache_beam.examples.dataframe import wordcount from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import create_file from apache_beam.testing.test_utils import read_files_from_pattern class WordCountTest(unittest.TestCase): SAMPLE_TEXT = """ a a b a b c loooooonger words """ @pytest.mark.examples_postcommit def test_basics(self): test_pipeline = TestPipeline(is_integration_test=True) # Setup the files with expected content. temp_location = test_pipeline.get_option('temp_location') temp_path = '/'.join([temp_location, str(uuid.uuid4())]) input = create_file('/'.join([temp_path, 'input.txt']), self.SAMPLE_TEXT) expected_words = collections.defaultdict(int) for word in re.findall(r'[\w]+', self.SAMPLE_TEXT): expected_words[word] += 1 extra_opts = {'input': input, 'output': '%s.result' % temp_path} wordcount.run(test_pipeline.get_full_options_as_args(**extra_opts)) # Parse result file and compare. results = [] lines = read_files_from_pattern(temp_path + '.result*').splitlines() for line in lines: match = re.search(r'(\S+),([0-9]+)', line) if match is not None: results.append((match.group(1), int(match.group(2)))) elif line.strip(): self.assertEqual(line.strip(), 'word,count') self.assertEqual(sorted(results), sorted(expected_words.items())) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/fastavro_it_test.py0000644000175100017510000001333315206113032023430 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for Avro IO's fastavro support. Writes a configurable number of records to a temporary location with fastavro, then reads them back in from source, joins the generated records and records that are read from the source, and verifies they have the same elements. Usage: DataFlowRunner: pytest apache_beam/examples/fastavro_it_test.py \ --test-pipeline-options=" --runner=TestDataflowRunner --project=... --region=... --staging_location=gs://... --temp_location=gs://... --output=gs://... --sdk_location=... " DirectRunner: pytest apache_beam/examples/fastavro_it_test.py \ --test-pipeline-options=" --output=/tmp --records=5000 " """ # pytype: skip-file import json import logging import unittest import uuid import pytest from fastavro import parse_schema from apache_beam.io.avroio import ReadAllFromAvro from apache_beam.io.avroio import WriteToAvro from apache_beam.runners.runner import PipelineState from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import delete_files from apache_beam.testing.util import BeamAssertException from apache_beam.transforms.core import Create from apache_beam.transforms.core import FlatMap from apache_beam.transforms.core import Map from apache_beam.transforms.util import CoGroupByKey LABELS = ['abc', 'def', 'ghi', 'jkl', 'mno', 'pqr', 'stu', 'vwx'] COLORS = ['RED', 'ORANGE', 'YELLOW', 'GREEN', 'BLUE', 'PURPLE', None] def record(i): return { 'label': LABELS[i % len(LABELS)], 'number': i, 'number_str': str(i), 'color': COLORS[i % len(COLORS)] } def assertEqual(l, r): if l != r: raise BeamAssertException('Assertion failed: %s == %s' % (l, r)) def check(element): assert element['color'] in COLORS assert element['label'] in LABELS assertEqual( sorted(element.keys()), ['color', 'label', 'number', 'number_str']) class FastavroIT(unittest.TestCase): SCHEMA_STRING = ''' {"namespace": "example.avro", "type": "record", "name": "User", "fields": [ {"name": "label", "type": "string"}, {"name": "number", "type": ["int", "null"]}, {"name": "number_str", "type": ["string", "null"]}, {"name": "color", "type": ["string", "null"]} ] } ''' def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.uuid = str(uuid.uuid4()) self.output = '/'.join([self.test_pipeline.get_option('output'), self.uuid]) @pytest.mark.it_postcommit def test_avro_it(self): num_records = self.test_pipeline.get_option('records') num_records = int(num_records) if num_records else 1000000 fastavro_output = '/'.join([self.output, 'fastavro']) # Seed a `PCollection` with indices that will each be FlatMap'd into # `batch_size` records, to avoid having a too-large list in memory at # the outset batch_size = self.test_pipeline.get_option('batch-size') batch_size = int(batch_size) if batch_size else 10000 # pylint: disable=bad-option-value batches = range(int(num_records / batch_size)) def batch_indices(start): # pylint: disable=bad-option-value return range(start * batch_size, (start + 1) * batch_size) # A `PCollection` with `num_records` avro records records_pcoll = \ self.test_pipeline \ | 'create-batches' >> Create(batches) \ | 'expand-batches' >> FlatMap(batch_indices) \ | 'create-records' >> Map(record) # pylint: disable=expression-not-assigned records_pcoll \ | 'write_fastavro' >> WriteToAvro( fastavro_output, parse_schema(json.loads(self.SCHEMA_STRING)), ) result = self.test_pipeline.run() result.wait_until_finish() fastavro_pcoll = self.test_pipeline \ | 'create-fastavro' >> Create(['%s*' % fastavro_output]) \ | 'read-fastavro' >> ReadAllFromAvro() mapped_fastavro_pcoll = fastavro_pcoll | "map_fastavro" >> Map( lambda x: (x['number'], x)) mapped_record_pcoll = records_pcoll | "map_record" >> Map( lambda x: (x['number'], x)) def validate_record(elem): v = elem[1] def assertEqual(l, r): if l != r: raise BeamAssertException('Assertion failed: %s == %s' % (l, r)) assertEqual(sorted(v.keys()), ['fastavro', 'record_pcoll']) record_pcoll_values = v['record_pcoll'] fastavro_values = v['fastavro'] assertEqual(record_pcoll_values, fastavro_values) assertEqual(len(record_pcoll_values), 1) { "record_pcoll": mapped_record_pcoll, "fastavro": mapped_fastavro_pcoll } | CoGroupByKey() | Map(validate_record) result = self.test_pipeline.run() result.wait_until_finish() self.addCleanup(delete_files, [self.output]) assert result.state == PipelineState.DONE if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.050373 apache_beam-2.74.0/apache_beam/examples/flink/0000755000175100017510000000000015206113157020606 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/flink/__init__.py0000644000175100017510000000142015206113032022704 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/flink/flink_streaming_impulse.py0000644000175100017510000000560015206113032026063 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A streaming workflow that uses a synthetic streaming source. This can only be used with the Flink portable runner. """ # pytype: skip-file import argparse import logging import sys import apache_beam as beam from apache_beam.io.flink.flink_streaming_impulse_source import FlinkStreamingImpulseSource from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.transforms import window from apache_beam.transforms.trigger import AccumulationMode from apache_beam.transforms.trigger import AfterProcessingTime from apache_beam.transforms.trigger import Repeatedly def split(s): a = s.split("-") return a[0], int(a[1]) def count(x): return x[0], sum(x[1]) def apply_timestamp(element): import time yield window.TimestampedValue(element, time.time()) def run(argv=None): """Build and run the pipeline.""" args = [ "--runner=PortableRunner", "--job_endpoint=localhost:8099", "--streaming" ] if argv: args.extend(argv) parser = argparse.ArgumentParser() parser.add_argument( '--count', dest='count', default=0, help='Number of triggers to generate ' '(0 means emit forever).') parser.add_argument( '--interval_ms', dest='interval_ms', default=500, help='Interval between records per parallel ' 'Flink subtask.') known_args, pipeline_args = parser.parse_known_args(args) pipeline_options = PipelineOptions(pipeline_args) with beam.Pipeline(options=pipeline_options) as p: messages = ( p | FlinkStreamingImpulseSource().set_message_count( known_args.count).set_interval_ms(known_args.interval_ms)) _ = ( messages | 'decode' >> beam.Map(lambda x: ('', 1)) | 'window' >> beam.WindowInto( window.GlobalWindows(), trigger=Repeatedly(AfterProcessingTime(5 * 1000)), accumulation_mode=AccumulationMode.DISCARDING) | 'group' >> beam.GroupByKey() | 'count' >> beam.Map(count) | 'log' >> beam.Map(lambda x: logging.info("%d" % x[1]))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run(sys.argv[1:]) ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.107379 apache_beam-2.74.0/apache_beam/examples/inference/0000755000175100017510000000000015206113157021441 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/__init__.py0000644000175100017510000000142015206113032023537 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/gemini_image_generation.py0000644000175100017510000001077715206113032026644 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A sample pipeline using the RunInference API to classify text using an LLM. This pipeline creates a set of prompts and sends it to a Gemini service then returns the predictions from the classifier model. This example uses the gemini-2.0-flash-001 model. """ import argparse import logging from collections.abc import Iterable from io import BytesIO import apache_beam as beam from apache_beam.io.fileio import FileSink from apache_beam.io.fileio import WriteToFiles from apache_beam.io.fileio import default_file_naming from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.gemini_inference import GeminiModelHandler from apache_beam.ml.inference.gemini_inference import generate_image_from_strings_and_images from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from PIL import Image def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--output', dest='output', type=str, required=True, help='Path to save output predictions.') parser.add_argument( '--api_key', dest='api_key', type=str, required=False, help='Gemini Developer API key.') parser.add_argument( '--cloud_project', dest='project', type=str, required=False, help='GCP Project') parser.add_argument( '--cloud_region', dest='location', type=str, required=False, help='GCP location for the Endpoint') return parser.parse_known_args(argv) class PostProcessor(beam.DoFn): def process(self, element: PredictionResult) -> Iterable[Image.Image]: try: response = element.inference for part in response.parts: if part.text is not None: print(part.text) elif part.inline_data is not None: image = Image.open(BytesIO(part.inline_data.data)) yield image except Exception as e: print(f"Can't decode inference for element: {element.example}, got {e}") raise e class ImageSink(FileSink): def open(self, fh) -> None: self._fh = fh def write(self, record): record.save(self._fh, format='PNG') def flush(self): self._fh.flush() def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session model_handler = GeminiModelHandler( model_name='gemini-2.5-flash-image', request_fn=generate_image_from_strings_and_images, api_key=known_args.api_key, project=known_args.project, location=known_args.location) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) prompts = [ "Create a picture of a pineapple in the sand at a beach.", ] read_prompts = pipeline | "Get prompt" >> beam.Create(prompts) predictions = read_prompts | "RunInference" >> RunInference(model_handler) processed = predictions | "PostProcess" >> beam.ParDo(PostProcessor()) _ = processed | "WriteOutput" >> WriteToFiles( path=known_args.output, file_naming=default_file_naming("gemini-image", ".png"), sink=ImageSink()) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/gemini_text_classification.py0000644000175100017510000001010715206113032027371 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A sample pipeline using the RunInference API to classify text using an LLM. This pipeline creates a set of prompts and sends it to a Gemini service then returns the predictions from the classifier model. This example uses the gemini-2.0-flash-001 model. """ import argparse import logging from collections.abc import Iterable import apache_beam as beam from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.gemini_inference import GeminiModelHandler from apache_beam.ml.inference.gemini_inference import generate_from_string from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--output', dest='output', type=str, required=True, help='Path to save output predictions.') parser.add_argument( '--api_key', dest='api_key', type=str, required=False, help='Gemini Developer API key.') parser.add_argument( '--cloud_project', dest='project', type=str, required=False, help='GCP Project') parser.add_argument( '--cloud_region', dest='location', type=str, required=False, help='GCP location for the Endpoint') return parser.parse_known_args(argv) class PostProcessor(beam.DoFn): def process(self, element: PredictionResult) -> Iterable[str]: for part in element.inference.parts: try: output_text = part.text yield f"Input: {element.example}, Output: {output_text}" except Exception as e: print(f"Can't decode inference for element: {element.example}, got {e}") raise e def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session model_handler = GeminiModelHandler( model_name='gemini-2.0-flash-001', request_fn=generate_from_string, api_key=known_args.api_key, project=known_args.project, location=known_args.location) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) prompts = [ "What is 5+2?", "Who is the protagonist of Lord of the Rings?", "What is the air-speed velocity of a laden swallow?" ] read_prompts = pipeline | "Get prompt" >> beam.Create(prompts) predictions = read_prompts | "RunInference" >> RunInference(model_handler) processed = predictions | "PostProcess" >> beam.ParDo(PostProcessor()) _ = processed | "PrintOutput" >> beam.Map(print) _ = processed | "WriteOutput" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/huggingface_language_modeling.py0000644000175100017510000001515615206113032030003 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A pipeline that uses RunInference to perform Language Modeling with masked language model from Hugging Face. This pipeline takes sentences from a custom text file, converts the last word of the sentence into a token, and then uses the AutoModelForMaskedLM from Hugging Face to predict the best word for the masked token given all the words already in the sentence. The pipeline then writes the prediction to an output file in which users can then compare against the original sentence. """ import argparse import logging from collections.abc import Iterable from collections.abc import Iterator import apache_beam as beam import torch from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.huggingface_inference import HuggingFaceModelHandlerKeyedTensor from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from transformers import AutoModelForMaskedLM from transformers import AutoTokenizer def add_mask_to_last_word(text: str) -> tuple[str, str]: text_list = text.split() return text, ' '.join(text_list[:-2] + ['', text_list[-1]]) def tokenize_sentence( text_and_mask: tuple[str, str], tokenizer: AutoTokenizer) -> tuple[str, dict[str, torch.Tensor]]: text, masked_text = text_and_mask tokenized_sentence = tokenizer.encode_plus(masked_text, return_tensors="pt") # Workaround to manually remove batch dim until we have the feature to # add optional batching flag. # TODO(https://github.com/apache/beam/issues/21863): Remove once optional # batching flag added return text, { k: torch.squeeze(v) for k, v in dict(tokenized_sentence).items() } def filter_empty_lines(text: str) -> Iterator[str]: if len(text.strip()) > 0: yield text class PostProcessor(beam.DoFn): """Processes the PredictionResult to get the predicted word. The logits are the output of the Model. We can get the word with the highest probability of being a candidate replacement word by taking the argmax. """ def __init__(self, tokenizer: AutoTokenizer): super().__init__() self.tokenizer = tokenizer def process(self, element: tuple[str, PredictionResult]) -> Iterable[str]: text, prediction_result = element inputs = prediction_result.example logits = prediction_result.inference['logits'] mask_token_index = torch.where( inputs["input_ids"] == self.tokenizer.mask_token_id)[0] predicted_token_id = logits[mask_token_index].argmax(axis=-1) decoded_word = self.tokenizer.decode(predicted_token_id) yield text + ';' + decoded_word def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', help='Path to the text file containing sentences.') parser.add_argument( '--output', dest='output', required=True, help='Path of file in which to save the output predictions.') parser.add_argument( '--model_name', dest='model_name', required=True, help='bert uncased model. This can be base model or large model') parser.add_argument( '--model_class', dest='model_class', default=AutoModelForMaskedLM, help='Name of the model from Hugging Face') parser.add_argument( '--large_model', action='store_true', dest='large_model', default=False, help='Set to true if your model is large enough to run into memory ' 'pressure if you load multiple copies.') return parser.parse_known_args(argv) def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) tokenizer = AutoTokenizer.from_pretrained(known_args.model_name) model_handler = HuggingFaceModelHandlerKeyedTensor( model_uri=known_args.model_name, model_class=known_args.model_class, framework='pt', max_batch_size=1, large_model=known_args.large_model) if not known_args.input: text = ( pipeline | 'CreateSentences' >> beam.Create([ 'The capital of France is Paris .', 'It is raining cats and dogs .', 'Today is Monday and tomorrow is Tuesday .', 'There are 5 coconuts on this palm tree .', 'The strongest person in the world is not famous .', 'The secret ingredient to his wonderful life was gratitude .', 'The biggest animal in the world is the whale .', ])) else: text = ( pipeline | 'ReadSentences' >> beam.io.ReadFromText(known_args.input)) text_and_tokenized_text_tuple = ( text | 'FilterEmptyLines' >> beam.ParDo(filter_empty_lines) | 'AddMask' >> beam.Map(add_mask_to_last_word) | 'TokenizeSentence' >> beam.Map(lambda x: tokenize_sentence(x, tokenizer))) output = ( text_and_tokenized_text_tuple | 'RunInference' >> RunInference(KeyedModelHandler(model_handler)) | 'ProcessOutput' >> beam.ParDo(PostProcessor(tokenizer=tokenizer))) _ = output | "WriteOutput" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/huggingface_question_answering.py0000644000175100017510000001354715206113032030270 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """"A pipeline that uses RunInference to perform Question Answering using the model from Hugging Face Models Hub. This pipeline takes questions and context from a custom text file separated by a semicolon. These are converted to SquadExamples by using the utility provided by transformers.QuestionAnsweringPipeline and passed to the model handler. We just provide the model name here because the model repository specifies the task that it will do. The pipeline then writes the prediction to an output file in which users can then compare against the original context. """ import argparse import logging from collections.abc import Iterable import apache_beam as beam from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.huggingface_inference import HuggingFacePipelineModelHandler from apache_beam.ml.inference.huggingface_inference import PipelineTask from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from transformers import QuestionAnsweringPipeline class PostProcessor(beam.DoFn): """Processes the PredictionResult to get the predicted answer. Hugging Face Pipeline for Question Answering returns a dictionary with score, start and end index of answer and the answer. """ def process(self, result: tuple[str, PredictionResult]) -> Iterable[str]: text, prediction = result predicted_answer = prediction.inference['answer'] yield text + ';' + predicted_answer def preprocess(text): """ preprocess separates the text into question and context by splitting on semi-colon. Args: text (str): string with question and context separated by semi-colon. Yields: (str, str): yields question and context from text. """ if len(text.strip()) > 0: question, context = text.split(';') yield (question, context) def create_squad_example(text): """Creates SquadExample objects to be fed to QuestionAnsweringPipeline supported by Hugging Face. Check out https://huggingface.co/docs/transformers/main_classes/pipelines#transformers.QuestionAnsweringPipeline.__call__.X #pylint: disable=line-too-long to learn about valid input types for QuestionAnswering Pipeline. Args: text (Tuple[str,str]): a tuple of question and context. """ question, context = text yield question, QuestionAnsweringPipeline.create_sample(question, context) def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', help='Path of file containing question and context separated by semicolon' ) parser.add_argument( '--output', dest='output', required=True, help='Path of file in which to save the output predictions.') parser.add_argument( '--model_name', dest='model_name', default="deepset/roberta-base-squad2", help='Model repository-id from Hugging Face Models Hub.') parser.add_argument( '--revision', dest='revision', help= 'Specific model version to use - branch name, tag name, or a commit-id.') return parser.parse_known_args(argv) def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) model_handler = HuggingFacePipelineModelHandler( task=PipelineTask.QuestionAnswering, model=known_args.model_name, load_model_args={ 'framework': 'pt', 'revision': known_args.revision }) if not known_args.input: text = ( pipeline | 'CreateSentences' >> beam.Create([ "What does Apache Beam do?;" "Apache Beam enables batch and streaming data processing.", "What is the capital of France?;The capital of France is Paris .", "Where was beam summit?;Apache Beam Summit 2023 was in NYC.", ])) else: text = ( pipeline | 'ReadSentences' >> beam.io.ReadFromText(known_args.input)) processed_text = ( text | 'PreProcess' >> beam.ParDo(preprocess) | 'SquadExample' >> beam.ParDo(create_squad_example)) output = ( processed_text | 'RunInference' >> RunInference(KeyedModelHandler(model_handler)) | 'ProcessOutput' >> beam.ParDo(PostProcessor())) _ = output | "WriteOutput" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/milk_quality_prediction_windowing.py0000644000175100017510000002056115206113032031020 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A streaming pipeline that uses RunInference API and windowing that classifies the quality of milk as good, bad or medium based on pH, temperature, taste, odor, fat, turbidity and color. Each minute new measurements come in and a sliding window aggregates the number of good, bad and medium samples. This example uses the milk quality prediction dataset from kaggle. https://www.kaggle.com/datasets/cpluzshrijayan/milkquality In order to set this example up, you will need two things. 1. Download the data in csv format from kaggle and host it. 2. Split the dataset in a training set and test set (preprocess_data function). 3. Train the classifier. """ import argparse import logging import time from typing import NamedTuple import pandas from sklearn.model_selection import train_test_split import apache_beam as beam import xgboost from apache_beam import window from apache_beam.ml.inference import RunInference from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.xgboost_inference import XGBoostModelHandlerPandas from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from apache_beam.testing.test_stream import TestStream def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--dataset', dest='dataset', required=True, help='Path to the csv containing Kaggle Milk Quality dataset.') parser.add_argument( '--pipeline_input_data', dest='pipeline_input_data', required=True, help='Path to store the csv containing input data for the pipeline.' 'This will be generated as part of preprocessing the data.') parser.add_argument( '--training_set', dest='training_set', required=True, help='Path to store the csv containing the training set.' 'This will be generated as part of preprocessing the data.') parser.add_argument( '--labels', dest='labels', required=True, help='Path to store the csv containing the labels used in training.' 'This will be generated as part of preprocessing the data.') parser.add_argument( '--model_state', dest='model_state', required=True, help='Path to the state of the XGBoost model loaded for Inference.') return parser.parse_known_args(argv) def preprocess_data( dataset_path: str, training_set_path: str, labels_path: str, test_set_path: str): """ Helper function to split the dataset into a training set and its labels and a test set. The training set and its labels are used to train a lightweight model. The test set is used to create a test streaming pipeline. Args: dataset_path: path to csv file containing the Kaggle milk quality dataset training_set_path: path to output the training samples labels_path: path to output the labels for the training set test_set_path: path to output the test samples """ df = pandas.read_csv(dataset_path) df['Grade'].replace(['low', 'medium', 'high'], [0, 1, 2], inplace=True) x = df.drop(columns=['Grade']) y = df['Grade'] x_train, x_test, y_train, _ = \ train_test_split(x, y, test_size=0.60, random_state=99) x_train.to_csv(training_set_path, index=False) y_train.to_csv(labels_path, index=False) x_test.to_csv(test_set_path, index=False) def train_model( samples_path: str, labels_path: str, model_state_output_path: str): """Function to train the XGBoost model. Args: samples_path: path to csv file containing the training data labels_path: path to csv file containing the labels for the training data model_state_output_path: Path to store the trained model """ samples = pandas.read_csv(samples_path) labels = pandas.read_csv(labels_path) xgb = xgboost.XGBClassifier(max_depth=3) xgb.fit(samples, labels) xgb.save_model(model_state_output_path) return xgb class MilkQualityAggregation(NamedTuple): bad_quality_measurements: int medium_quality_measurements: int high_quality_measurements: int class AggregateMilkQualityResults(beam.CombineFn): """Simple aggregation to keep track of the number of samples with good, bad and medium quality milk.""" def create_accumulator(self): return MilkQualityAggregation(0, 0, 0) def add_input( self, accumulator: MilkQualityAggregation, element: PredictionResult): quality = element.inference[0] if quality == 0: return MilkQualityAggregation( accumulator.bad_quality_measurements + 1, accumulator.medium_quality_measurements, accumulator.high_quality_measurements) elif quality == 1: return MilkQualityAggregation( accumulator.bad_quality_measurements, accumulator.medium_quality_measurements + 1, accumulator.high_quality_measurements) else: return MilkQualityAggregation( accumulator.bad_quality_measurements, accumulator.medium_quality_measurements, accumulator.high_quality_measurements + 1) def merge_accumulators(self, accumulators: MilkQualityAggregation): return MilkQualityAggregation( sum( aggregation.bad_quality_measurements for aggregation in accumulators), sum( aggregation.medium_quality_measurements for aggregation in accumulators), sum( aggregation.high_quality_measurements for aggregation in accumulators), ) def extract_output(self, accumulator: MilkQualityAggregation): return accumulator def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session milk_quality_data = pandas.read_csv(known_args.pipeline_input_data) start = time.mktime(time.strptime('2023/06/29 10:00:00', '%Y/%m/%d %H:%M:%S')) # Create a test stream test_stream = TestStream() # Watermark is set to 10:00:00 test_stream.advance_watermark_to(start) # Split the dataframe in individual samples samples = [ milk_quality_data.iloc[i:i + 1] for i in range(len(milk_quality_data)) ] for watermark_offset, sample in enumerate(samples, 1): test_stream.advance_watermark_to(start + watermark_offset) test_stream.add_elements([sample]) test_stream.advance_watermark_to_infinity() model_handler = XGBoostModelHandlerPandas( model_class=xgboost.XGBClassifier, model_state=known_args.model_state) with beam.Pipeline() as p: _ = ( p | test_stream | 'window' >> beam.WindowInto(window.SlidingWindows(30, 5)) | "RunInference" >> RunInference(model_handler) | 'Count number of elements in window' >> beam.CombineGlobally( AggregateMilkQualityResults()).without_defaults() | 'Print' >> beam.Map(print)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) known_args, _ = parse_known_args(None) preprocess_data( known_args.dataset, training_set_path=known_args.training_set, labels_path=known_args.labels, test_set_path=known_args.pipeline_input_data) train_model( samples_path=known_args.training_set, labels_path=known_args.labels, model_state_output_path=known_args.model_state) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/onnx_sentiment_classification.py0000644000175100017510000001240315206113032030126 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """"A pipeline that uses RunInference to perform sentiment classification using RoBERTa. This pipeline takes sentences from a custom text file, and then uses RoBERTa from Hugging Face to predict the sentiment of a given review. The pipeline then writes the prediction to an output file in which users can then compare against true labels. Model is fine-tuned RoBERTa from https://github.com/SeldonIO/seldon-models/blob/master/pytorch/moviesentiment_roberta/pytorch-roberta-onnx.ipynb # pylint: disable=line-too-long """ import argparse import logging from collections.abc import Iterable from collections.abc import Iterator import numpy as np import apache_beam as beam import torch from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.onnx_inference import OnnxModelHandlerNumpy from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from transformers import RobertaTokenizer def tokenize_sentence(text: str, tokenizer: RobertaTokenizer) -> tuple[str, torch.Tensor]: tokenized_sentence = tokenizer.encode(text, add_special_tokens=True) # Workaround to manually remove batch dim until we have the feature to # add optional batching flag. # TODO(https://github.com/apache/beam/issues/21863): Remove once optional # batching flag added return text, torch.tensor(tokenized_sentence).numpy() def filter_empty_lines(text: str) -> Iterator[str]: if len(text.strip()) > 0: yield text class PostProcessor(beam.DoFn): def process(self, element: tuple[str, PredictionResult]) -> Iterable[str]: filename, prediction_result = element prediction = np.argmax(prediction_result.inference, axis=0) yield filename + ';' + str(prediction) def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', help='Path to the text file containing sentences.') parser.add_argument( '--output', dest='output', required=True, help='Path of file in which to save the output predictions.') parser.add_argument( '--model_uri', dest='model_uri', required=True, help="Path to the model's uri.") return parser.parse_known_args(argv) def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session # TODO: Remove once nested tensors https://github.com/pytorch/nestedtensor # is officially released. class OnnxNoBatchModelHandler(OnnxModelHandlerNumpy): """Wrapper to OnnxModelHandlerNumpy to limit batch size to 1. The tokenized strings generated from RobertaTokenizer may have different lengths, which doesn't work with torch.stack() in current RunInference implementation since stack() requires tensors to be the same size. Restricting max_batch_size to 1 means there is only 1 example per `batch` in the run_inference() call. """ def batch_elements_kwargs(self): return {'max_batch_size': 1} model_handler = OnnxNoBatchModelHandler(model_uri=known_args.model_uri) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) tokenizer = RobertaTokenizer.from_pretrained('roberta-base') text = (pipeline | 'ReadSentences' >> beam.io.ReadFromText(known_args.input)) text_and_tokenized_text_tuple = ( text | 'FilterEmptyLines' >> beam.ParDo(filter_empty_lines) | 'TokenizeSentence' >> beam.Map(lambda x: tokenize_sentence(x, tokenizer))) output = ( text_and_tokenized_text_tuple | 'PyTorchRunInference' >> RunInference(KeyedModelHandler(model_handler)) | 'ProcessOutput' >> beam.ParDo(PostProcessor())) _ = output | "WriteOutput" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/pytorch_image_classification.py0000644000175100017510000001360615206113032027716 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A pipeline that uses RunInference API to perform image classification.""" import argparse import io import logging import os from collections.abc import Iterator from typing import Optional import apache_beam as beam import torch from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerTensor from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from PIL import Image from torchvision import models from torchvision import transforms def read_image(image_file_name: str, path_to_dir: Optional[str] = None) -> tuple[str, Image.Image]: if path_to_dir is not None: image_file_name = os.path.join(path_to_dir, image_file_name) with FileSystems().open(image_file_name, 'r') as file: data = Image.open(io.BytesIO(file.read())).convert('RGB') return image_file_name, data def preprocess_image(data: Image.Image) -> torch.Tensor: image_size = (224, 224) # Pre-trained PyTorch models expect input images normalized with the # below values (see: https://pytorch.org/vision/stable/models.html) normalize = transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = transforms.Compose([ transforms.Resize(image_size), transforms.ToTensor(), normalize, ]) return transform(data) def filter_empty_lines(text: str) -> Iterator[str]: if len(text.strip()) > 0: yield text def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', required=True, help='Path to the text file containing image names.') parser.add_argument( '--output', dest='output', required=True, help='Path where to save output predictions.' ' text file.') parser.add_argument( '--model_state_dict_path', dest='model_state_dict_path', required=True, help="Path to the model's state_dict.") parser.add_argument( '--images_dir', default=None, help='Path to the directory where images are stored.' 'Not required if image names in the input file have absolute path.') return parser.parse_known_args(argv) def run( argv=None, model_class=None, model_params=None, save_main_session=True, device='CPU', test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. model_class: Reference to the class definition of the model. model_params: Parameters passed to the constructor of the model_class. These will be used to instantiate the model object in the RunInference API. save_main_session: Used for internal testing. device: Device to be used on the Runner. Choices are (CPU, GPU). test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session if not model_class: # default model class will be mobilenet with pretrained weights. model_class = models.mobilenet_v2 model_params = {'num_classes': 1000} def preprocess(image_name: str) -> tuple[str, torch.Tensor]: image_name, image = read_image( image_file_name=image_name, path_to_dir=known_args.images_dir) return (image_name, preprocess_image(image)) def postprocess(element: tuple[str, PredictionResult]) -> str: filename, prediction_result = element prediction = torch.argmax(prediction_result.inference, dim=0) return filename + ',' + str(prediction.item()) # In this example we pass keyed inputs to RunInference transform. # Therefore, we use KeyedModelHandler wrapper over PytorchModelHandler. model_handler = KeyedModelHandler( PytorchModelHandlerTensor( state_dict_path=known_args.model_state_dict_path, model_class=model_class, model_params=model_params, device=device, min_batch_size=10, max_batch_size=100)).with_preprocess_fn( preprocess).with_postprocess_fn(postprocess) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) filename_value_pair = ( pipeline | 'ReadImageNames' >> beam.io.ReadFromText(known_args.input) | 'FilterEmptyLines' >> beam.ParDo(filter_empty_lines)) predictions = ( filename_value_pair | 'PyTorchRunInference' >> RunInference(model_handler)) predictions | "WriteOutputToGCS" >> beam.io.WriteToText( # pylint: disable=expression-not-assigned known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/pytorch_image_classification_with_side_inputs.py0000644000175100017510000001745315206113032033363 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A pipeline that uses RunInference PTransform to perform image classification and uses WatchFilePattern as side input to the RunInference PTransform. WatchFilePattern is used to watch for a file updates matching the file_pattern based on timestamps and emits latest model metadata, which is used in RunInference API for the dynamic model updates without the need for stopping the beam pipeline. This pipeline follows the pattern from https://beam.apache.org/documentation/patterns/side-inputs/ To use the PubSub reading from a topic in the pipeline as source, you can publish a path to the model(resnet152 used in the pipeline from torchvision.models.resnet152) to the PubSub topic. Then pass that topic via command line arg --topic. The published path(str) should be UTF-8 encoded. To run the example on DataflowRunner, python apache_beam/examples/inference/pytorch_image_classification_with_side_inputs.py # pylint: disable=line-too-long --project= --re= --temp_location= --staging_location= --runner=DataflowRunner --streaming --interval=10 --num_workers=5 --requirements_file=apache_beam/ml/inference/torch_tests_requirements.txt --topic= --file_pattern= file_pattern is path(can contain glob characters), which will be passed to WatchContinuously transform for model updates. WatchContinuously watches the file_pattern and emits a latest file path, sorted by timestamp. Files that are read before and updated with same name will be ignored as an update. The pipeline expects there is at least one file present to match the file_pattern before pipeline startup. Presumably, this would be the `initial_model_path`. If there is no file matching before pipeline startup time, the pipeline would fail. """ import argparse import io import logging import os from collections.abc import Iterable from collections.abc import Iterator from typing import Optional import apache_beam as beam import torch from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerTensor from apache_beam.ml.inference.utils import WatchFilePattern from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from PIL import Image from torchvision import models from torchvision import transforms def read_image(image_file_name: str, path_to_dir: Optional[str] = None) -> tuple[str, Image.Image]: if path_to_dir is not None: image_file_name = os.path.join(path_to_dir, image_file_name) with FileSystems().open(image_file_name, 'r') as file: data = Image.open(io.BytesIO(file.read())).convert('RGB') return image_file_name, data def preprocess_image(data: Image.Image) -> torch.Tensor: image_size = (224, 224) # Pre-trained PyTorch models expect input images normalized with the # below values (see: https://pytorch.org/vision/stable/models.html) normalize = transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = transforms.Compose([ transforms.Resize(image_size), transforms.ToTensor(), normalize, ]) return transform(data) def filter_empty_lines(text: str) -> Iterator[str]: if len(text.strip()) > 0: yield text class PostProcessor(beam.DoFn): """ Return filename, prediction and the model id used to perform the prediction """ def process(self, element: tuple[str, PredictionResult]) -> Iterable[str]: filename, prediction_result = element prediction = torch.argmax(prediction_result.inference, dim=0) yield filename, prediction, prediction_result.model_id def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--topic', dest='topic', help='PubSub topic emitting absolute path to the images.' 'Path must be accessible by the pipeline.') parser.add_argument( '--model_path', '--initial_model_path', dest='model_path', default='gs://apache-beam-samples/run_inference/resnet152.pth', help="Path to the initial model's state_dict. " "This will be used until the first model update occurs.") parser.add_argument( '--file_pattern', help='Glob pattern to watch for an update.') parser.add_argument( '--interval', default=10, type=int, help='Interval used to check for file updates.') return parser.parse_known_args(argv) def run( argv=None, model_class=None, model_params=None, save_main_session=True, device='CPU', test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. model_class: Reference to the class definition of the model. model_params: Parameters passed to the constructor of the model_class. These will be used to instantiate the model object in the RunInference PTransform. save_main_session: Used for internal testing. device: Device to be used on the Runner. Choices are (CPU, GPU). test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session if not model_class: model_class = models.resnet152 model_params = {'num_classes': 1000} # In this example we pass keyed inputs to RunInference transform. # Therefore, we use KeyedModelHandler wrapper over PytorchModelHandler. model_handler = KeyedModelHandler( PytorchModelHandlerTensor( state_dict_path=known_args.model_path, model_class=model_class, model_params=model_params, device=device, min_batch_size=10, max_batch_size=100)) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) side_input = pipeline | WatchFilePattern( interval=known_args.interval, file_pattern=known_args.file_pattern) filename_value_pair = ( pipeline | 'ReadImageNamesFromPubSub' >> beam.io.ReadFromPubSub(known_args.topic) | 'DecodeBytes' >> beam.Map(lambda x: x.decode('utf-8')) | 'ReadImageData' >> beam.Map(lambda image_name: read_image(image_file_name=image_name)) | 'PreprocessImages' >> beam.MapTuple( lambda file_name, data: (file_name, preprocess_image(data)))) predictions = ( filename_value_pair | 'PyTorchRunInference' >> RunInference( model_handler, model_metadata_pcoll=side_input) | 'ProcessOutput' >> beam.ParDo(PostProcessor())) _ = predictions | beam.Map(logging.info) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/pytorch_image_segmentation.py0000644000175100017510000001556515206113032027426 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A pipeline that uses RunInference API to perform image segmentation.""" import argparse import io import logging import os from collections.abc import Iterable from collections.abc import Iterator from typing import Optional import apache_beam as beam import torch from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerTensor from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from PIL import Image from torchvision import transforms from torchvision.models.detection import maskrcnn_resnet50_fpn COCO_INSTANCE_CLASSES = [ '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush' ] CLASS_ID_TO_NAME = dict(enumerate(COCO_INSTANCE_CLASSES)) def read_image(image_file_name: str, path_to_dir: Optional[str] = None) -> tuple[str, Image.Image]: if path_to_dir is not None: image_file_name = os.path.join(path_to_dir, image_file_name) with FileSystems().open(image_file_name, 'r') as file: data = Image.open(io.BytesIO(file.read())).convert('RGB') return image_file_name, data def preprocess_image(data: Image.Image) -> torch.Tensor: image_size = (224, 224) transform = transforms.Compose([ transforms.Resize(image_size), transforms.ToTensor(), ]) return transform(data) def filter_empty_lines(text: str) -> Iterator[str]: if len(text.strip()) > 0: yield text class PostProcessor(beam.DoFn): def process(self, element: tuple[str, PredictionResult]) -> Iterable[str]: filename, prediction_result = element prediction_labels = prediction_result.inference['labels'] classes = [CLASS_ID_TO_NAME[label.item()] for label in prediction_labels] yield filename + ';' + str(classes) def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', required=True, help='Path to the text file containing image names.') parser.add_argument( '--output', dest='output', required=True, help='Path where to save output predictions.' ' text file.') parser.add_argument( '--model_state_dict_path', dest='model_state_dict_path', required=True, help="Path to the model's state_dict. " "Default state_dict would be maskrcnn_resnet50_fpn.") parser.add_argument( '--images_dir', help='Path to the directory where images are stored.' 'Not required if image names in the input file have absolute path.') return parser.parse_known_args(argv) def run( argv=None, model_class=None, model_params=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. model_class: Reference to the class definition of the model. If None, maskrcnn_resnet50_fpn will be used as default . model_params: Parameters passed to the constructor of the model_class. These will be used to instantiate the model object in the RunInference API. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session if not model_class: model_class = maskrcnn_resnet50_fpn model_params = {'num_classes': 91} model_handler = PytorchModelHandlerTensor( state_dict_path=known_args.model_state_dict_path, model_class=model_class, model_params=model_params) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) filename_value_pair = ( pipeline | 'ReadImageNames' >> beam.io.ReadFromText(known_args.input) | 'FilterEmptyLines' >> beam.ParDo(filter_empty_lines) | 'ReadImageData' >> beam.Map( lambda image_name: read_image( image_file_name=image_name, path_to_dir=known_args.images_dir)) | 'PreprocessImages' >> beam.MapTuple( lambda file_name, data: (file_name, preprocess_image(data)))) predictions = ( filename_value_pair | 'PyTorchRunInference' >> RunInference(KeyedModelHandler(model_handler)) | 'ProcessOutput' >> beam.ParDo(PostProcessor())) _ = predictions | "WriteOutput" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/pytorch_language_modeling.py0000644000175100017510000002053115206113032027215 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A pipeline that uses RunInference to perform Language Modeling with Bert. This pipeline takes sentences from a custom text file, converts the last word of the sentence into a [MASK] token, and then uses the BertForMaskedLM from Hugging Face to predict the best word for the masked token given all the words already in the sentence. The pipeline then writes the prediction to an output file in which users can then compare against the original sentence. """ import argparse import logging from collections.abc import Iterable from collections.abc import Iterator import apache_beam as beam import torch from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerKeyedTensor from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from transformers import BertConfig from transformers import BertForMaskedLM from transformers import BertTokenizer def add_mask_to_last_word(text: str) -> tuple[str, str]: text_list = text.split() return text, ' '.join(text_list[:-2] + ['[MASK]', text_list[-1]]) def tokenize_sentence( text_and_mask: tuple[str, str], bert_tokenizer: BertTokenizer) -> tuple[str, dict[str, torch.Tensor]]: text, masked_text = text_and_mask tokenized_sentence = bert_tokenizer(masked_text, return_tensors="pt") # Workaround to manually remove batch dim until we have the feature to # add optional batching flag. # TODO(https://github.com/apache/beam/issues/21863): Remove once optional # batching flag added return text, { k: torch.squeeze(v) for k, v in dict(tokenized_sentence).items() } def filter_empty_lines(text: str) -> Iterator[str]: if len(text.strip()) > 0: yield text class PostProcessor(beam.DoFn): """Processes the PredictionResult to get the predicted word. The logits are the output of the BERT Model. After applying a softmax activation function to the logits, we get probabilistic distributions for each of the words in BERT’s vocabulary. We can get the word with the highest probability of being a candidate replacement word by taking the argmax. """ def __init__(self, bert_tokenizer: BertTokenizer): super().__init__() self.bert_tokenizer = bert_tokenizer def process(self, element: tuple[str, PredictionResult]) -> Iterable[str]: text, prediction_result = element inputs = prediction_result.example logits = prediction_result.inference['logits'] mask_token_index = ( inputs['input_ids'] == self.bert_tokenizer.mask_token_id).nonzero( as_tuple=True)[0] predicted_token_id = logits[mask_token_index].argmax(axis=-1) decoded_word = self.bert_tokenizer.decode(predicted_token_id) yield text + ';' + decoded_word def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', help='Path to the text file containing sentences.') parser.add_argument( '--output', dest='output', required=True, help='Path of file in which to save the output predictions.') parser.add_argument( '--bert_tokenizer', dest='bert_tokenizer', default='bert-base-uncased', help='bert uncased model. This can be base model or large model') parser.add_argument( '--model_state_dict_path', dest='model_state_dict_path', required=True, help="Path to the model's state_dict.") parser.add_argument( '--large_model', action='store_true', dest='large_model', default=False, help='Set to true if your model is large enough to run into memory ' 'pressure if you load multiple copies.') return parser.parse_known_args(argv) def run( argv=None, model_class=None, model_params=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. model_class: Reference to the class definition of the model. If None, BertForMaskedLM will be used as default . model_params: Parameters passed to the constructor of the model_class. These will be used to instantiate the model object in the RunInference API. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session if not model_class: model_config = BertConfig.from_pretrained( known_args.bert_tokenizer, is_decoder=False, return_dict=True) model_class = BertForMaskedLM model_params = {'config': model_config} # TODO: Remove once nested tensors https://github.com/pytorch/nestedtensor # is officially released. class PytorchNoBatchModelHandler(PytorchModelHandlerKeyedTensor): """Wrapper to PytorchModelHandler to limit batch size to 1. The tokenized strings generated from BertTokenizer may have different lengths, which doesn't work with torch.stack() in current RunInference implementation since stack() requires tensors to be the same size. Restricting max_batch_size to 1 means there is only 1 example per `batch` in the run_inference() call. """ def batch_elements_kwargs(self): return {'max_batch_size': 1} model_handler = PytorchNoBatchModelHandler( state_dict_path=known_args.model_state_dict_path, model_class=model_class, model_params=model_params, large_model=known_args.large_model) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) bert_tokenizer = BertTokenizer.from_pretrained(known_args.bert_tokenizer) if not known_args.input: text = ( pipeline | 'CreateSentences' >> beam.Create([ 'The capital of France is Paris .', 'It is raining cats and dogs .', 'He looked up and saw the sun and stars .', 'Today is Monday and tomorrow is Tuesday .', 'There are 5 coconuts on this palm tree .', 'The richest person in the world is not here .', 'Malls are amazing places to shop because you can find everything you need under one roof .', # pylint: disable=line-too-long 'This audiobook is sure to liquefy your brain .', 'The secret ingredient to his wonderful life was gratitude .', 'The biggest animal in the world is the whale .', ])) else: text = ( pipeline | 'ReadSentences' >> beam.io.ReadFromText(known_args.input)) text_and_tokenized_text_tuple = ( text | 'FilterEmptyLines' >> beam.ParDo(filter_empty_lines) | 'AddMask' >> beam.Map(add_mask_to_last_word) | 'TokenizeSentence' >> beam.Map(lambda x: tokenize_sentence(x, bert_tokenizer))) output = ( text_and_tokenized_text_tuple | 'PyTorchRunInference' >> RunInference(KeyedModelHandler(model_handler)) | 'ProcessOutput' >> beam.ParDo( PostProcessor(bert_tokenizer=bert_tokenizer))) output | "WriteOutput" >> beam.io.WriteToText( # pylint: disable=expression-not-assigned known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/pytorch_model_per_key_image_segmentation.py0000644000175100017510000002203115206113032032306 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A pipeline that uses RunInference API to perform image segmentation using multiple different models. """ import argparse import io import logging import os from collections.abc import Iterable from collections.abc import Iterator from typing import Optional import apache_beam as beam import torch from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import KeyModelMapping from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerTensor from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from PIL import Image from torchvision import transforms from torchvision.models.detection import maskrcnn_resnet50_fpn from torchvision.models.detection import maskrcnn_resnet50_fpn_v2 COCO_INSTANCE_CLASSES = [ '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush' ] CLASS_ID_TO_NAME = dict(enumerate(COCO_INSTANCE_CLASSES)) def read_image(image_file_name: str, path_to_dir: Optional[str] = None) -> tuple[str, Image.Image]: if path_to_dir is not None: image_file_name = os.path.join(path_to_dir, image_file_name) with FileSystems().open(image_file_name, 'r') as file: data = Image.open(io.BytesIO(file.read())).convert('RGB') return image_file_name, data def preprocess_image(data: Image.Image) -> torch.Tensor: image_size = (224, 224) transform = transforms.Compose([ transforms.Resize(image_size), transforms.ToTensor(), ]) return transform(data) def filter_empty_lines(text: str) -> Iterator[str]: if len(text.strip()) > 0: yield text class KeyExamplesForEachModelType(beam.DoFn): """Duplicate data to run against each model type""" def process( self, element: tuple[torch.Tensor, str]) -> Iterable[tuple[str, torch.Tensor]]: yield 'v1', element[0] yield 'v2', element[0] class PostProcessor(beam.DoFn): def process( self, element: tuple[str, PredictionResult]) -> tuple[torch.Tensor, str]: model, prediction_result = element prediction_labels = prediction_result.inference['labels'] classes = [CLASS_ID_TO_NAME[label.item()] for label in prediction_labels] yield prediction_result.example, f'{model} predictions: {str(classes)}' class FormatResults(beam.DoFn): def process(self, element): _, filename_prediction = element predictions = filename_prediction['predictions'] v1_predictions = next(p for p in predictions if 'v1 predictions' in p) v2_predictions = next(p for p in predictions if 'v2 predictions' in p) yield ( f"{filename_prediction['image_names'][0]} --- " f"{v1_predictions} --- " f"{v2_predictions}") def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', required=True, help='Path to the text file containing image names.') parser.add_argument( '--output', dest='output', required=True, help='Path where to save output predictions.' ' text file.') parser.add_argument( '--model_state_dict_paths', dest='model_state_dict_paths', required=True, help="Comma separated paths to the models' state_dicts. " "For this example, should include exactly 2 state_dicts corresponding " "to maskrcnn_resnet50_fpn and maskrcnn_resnet50_fpn_v2 classes.") parser.add_argument( '--images_dir', help='Path to the directory where images are stored.' 'Not required if image names in the input file have absolute path.') return parser.parse_known_args(argv) def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. model_params: Parameters passed to the constructor of the model_class. These will be used to instantiate the model object in the RunInference API. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session state_dicts = known_args.model_state_dict_paths.split(',') if len(state_dicts) != 2: raise AssertionError( f'Expected exactly 2 state_dicts to be supplied, got {len(state_dicts)}' ) mh1 = PytorchModelHandlerTensor( state_dict_path=state_dicts[0], model_class=maskrcnn_resnet50_fpn, model_params={'num_classes': 91}) mh2 = PytorchModelHandlerTensor( state_dict_path=state_dicts[1], model_class=maskrcnn_resnet50_fpn_v2, model_params={'num_classes': 91}) # We'll use v1 and v2 as our keys to point to our model handlers. # Note that multiple keys can also point to a single model handler, # unlike this example. model_handler = KeyedModelHandler( [KeyModelMapping(['v1'], mh1), KeyModelMapping(['v2'], mh2)]) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) value_filename_pair = ( pipeline | 'ReadImageNames' >> beam.io.ReadFromText(known_args.input) | 'FilterEmptyLines' >> beam.ParDo(filter_empty_lines) | 'ReadImageData' >> beam.Map( lambda image_name: read_image( image_file_name=image_name, path_to_dir=known_args.images_dir)) | 'PreprocessImages' >> beam.MapTuple( lambda file_name, data: (preprocess_image(data), file_name))) predictions = ( value_filename_pair | 'DuplicateData' >> beam.ParDo(KeyExamplesForEachModelType()) | 'PyTorchRunInference' >> RunInference(model_handler) | 'ProcessOutput' >> beam.ParDo(PostProcessor())) # We now have our set of (example, prediction) and (example, original # filename) tuples. We can use CoGroupByKey to join them by original example, # converting the Tensors to lists first so that Beam can compare them without # a custom coder. If all you care about is the model used for inference, you # can also get that from the PredictionResult returned from RunInference. results = ({ 'image_names': ( value_filename_pair | beam.MapTuple( lambda example, filename: (example.tolist(), filename))), 'predictions': ( predictions | beam.MapTuple( lambda example, prediction: (example.tolist(), prediction))) } | beam.CoGroupByKey()) _ = ( results | 'FormatResults' >> beam.ParDo(FormatResults()) | "WriteOutput" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True)) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/pytorch_sentiment.py0000644000175100017510000002451615206113032025571 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """A Streaming pipeline that uses RunInference to perform Sentiment analysis with DistilBERT model. It reads input lines from a Pub/Sub topic (populated from a GCS file), performs sentiment classification using a DistilBERT model via RunInference, and writes results to BigQuery. Resources like Pub/Sub topic/subscription cleanup is handled programmatically. """ import argparse import logging import threading import time from collections.abc import Iterable from google.cloud import pubsub_v1 import apache_beam as beam import torch import torch.nn.functional as F from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerKeyedTensor from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.runners.runner import PipelineResult from transformers import DistilBertConfig from transformers import DistilBertForSequenceClassification from transformers import DistilBertTokenizerFast class SentimentPostProcessor(beam.DoFn): """Processes PredictionResult to extract sentiment label and confidence.""" def __init__(self, tokenizer: DistilBertTokenizerFast): self.tokenizer = tokenizer def process(self, element: tuple[str, PredictionResult]) -> Iterable[dict]: text, prediction_result = element logits = prediction_result.inference['logits'] probs = F.softmax(logits, dim=-1) predicted_class = torch.argmax(probs).item() confidence = probs[predicted_class].item() sentiment = 'POSITIVE' if predicted_class == 1 else 'NEGATIVE' yield { 'text': text, 'sentiment': sentiment, 'confidence': float(confidence) } def tokenize_text(text: str, tokenizer: DistilBertTokenizerFast) -> tuple[str, dict]: """Tokenizes input text using the specified tokenizer.""" tokenized = tokenizer( text, padding='max_length', truncation=True, max_length=128, return_tensors="pt") return text, {k: torch.squeeze(v) for k, v in tokenized.items()} class RateLimitDoFn(beam.DoFn): def __init__(self, rate_per_sec: float): self.delay = 1.0 / rate_per_sec def process(self, element): time.sleep(self.delay) yield element def parse_known_args(argv): """Parses command-line arguments for pipeline execution.""" parser = argparse.ArgumentParser() parser.add_argument( '--output_table', required=True, help='Path to output BigQuery table') parser.add_argument( '--model_path', default='distilbert-base-uncased-finetuned-sst-2-english', help='Hugging Face model name or local path') parser.add_argument( '--model_state_dict_path', dest='model_state_dict_path', required=True, help="Path to the model's state_dict.") parser.add_argument( '--input', required=True, help='Path to input file on GCS') parser.add_argument( '--pubsub_topic', default='projects/apache-beam-testing/topics/test_sentiment_topic', help='Pub/Sub topic to publish text lines to') parser.add_argument( '--pubsub_subscription', default='projects/apache-beam-testing/subscriptions/' 'test_sentiment_subscription', help='Pub/Sub subscription to read from') parser.add_argument( '--project', default='apache-beam-testing', help='GCP project ID') parser.add_argument( '--mode', default='streaming', choices=['streaming', 'batch']) parser.add_argument( '--rate_limit', type=float, default=None, help='Elements per second to send to Pub/Sub') return parser.parse_known_args(argv) def ensure_pubsub_resources( project: str, topic_path: str, subscription_path: str): """Ensures that the required Pub/Sub topic and subscription exist.""" publisher = pubsub_v1.PublisherClient() subscriber = pubsub_v1.SubscriberClient() topic_name = topic_path.split("/")[-1] subscription_name = subscription_path.split("/")[-1] full_topic_path = publisher.topic_path(project, topic_name) full_subscription_path = subscriber.subscription_path( project, subscription_name) try: publisher.get_topic(request={"topic": full_topic_path}) except Exception: publisher.create_topic(name=full_topic_path) try: subscriber.get_subscription( request={"subscription": full_subscription_path}) except Exception: subscriber.create_subscription( name=full_subscription_path, topic=full_topic_path) def cleanup_pubsub_resources( project: str, topic_path: str, subscription_path: str): """Deletes Pub/Sub topic and subscription created for the pipeline.""" publisher = pubsub_v1.PublisherClient() subscriber = pubsub_v1.SubscriberClient() topic_name = topic_path.split("/")[-1] subscription_name = subscription_path.split("/")[-1] full_topic_path = publisher.topic_path(project, topic_name) full_subscription_path = subscriber.subscription_path( project, subscription_name) try: subscriber.delete_subscription( request={"subscription": full_subscription_path}) print(f"Deleted subscription: {subscription_name}") except Exception as e: print(f"Failed to delete subscription: {e}") try: publisher.delete_topic(request={"topic": full_topic_path}) print(f"Deleted topic: {topic_name}") except Exception as e: print(f"Failed to delete topic: {e}") def override_or_add(args, flag, value): if flag in args: idx = args.index(flag) args[idx + 1] = str(value) else: args.extend([flag, str(value)]) def run_load_pipeline(known_args, pipeline_args): """Load data pipeline: read lines from GCS file and send to Pub/Sub.""" override_or_add(pipeline_args, '--device', 'CPU') override_or_add(pipeline_args, '--num_workers', '5') override_or_add(pipeline_args, '--max_num_workers', '10') override_or_add( pipeline_args, '--job_name', f"sentiment-load-pubsub-{int(time.time())}") override_or_add(pipeline_args, '--project', known_args.project) pipeline_options = PipelineOptions(pipeline_args) pipeline = beam.Pipeline(options=pipeline_options) lines = ( pipeline | 'ReadGCSFile' >> beam.io.ReadFromText(known_args.input) | 'FilterEmpty' >> beam.Filter(lambda line: line.strip())) if known_args.rate_limit: lines = ( lines | 'RateLimit' >> beam.ParDo( RateLimitDoFn(rate_per_sec=known_args.rate_limit))) _ = ( lines | 'ToBytes' >> beam.Map(lambda line: line.encode('utf-8')) | 'PublishToPubSub' >> beam.io.WriteToPubSub(topic=known_args.pubsub_topic)) return pipeline.run() def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: known_args, pipeline_args = parse_known_args(argv) ensure_pubsub_resources( project=known_args.project, topic_path=known_args.pubsub_topic, subscription_path=known_args.pubsub_subscription) if known_args.mode == 'streaming': threading.Thread( target=lambda: (time.sleep(100), run_load_pipeline(known_args, pipeline_args)), daemon=True).start() pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session method = beam.io.WriteToBigQuery.Method.FILE_LOADS pipeline_options.view_as(StandardOptions).streaming = False if known_args.mode == 'streaming': method = beam.io.WriteToBigQuery.Method.STREAMING_INSERTS pipeline_options.view_as(StandardOptions).streaming = True model_handler = PytorchModelHandlerKeyedTensor( model_class=DistilBertForSequenceClassification, model_params={'config': DistilBertConfig(num_labels=2)}, state_dict_path=known_args.model_state_dict_path, device='GPU') tokenizer = DistilBertTokenizerFast.from_pretrained(known_args.model_path) pipeline = test_pipeline or beam.Pipeline(options=pipeline_options) # Main pipeline: read, process, write result to BigQuery output table if known_args.mode == 'batch': input = ( pipeline | 'ReadGCSFile' >> beam.io.ReadFromText(known_args.input) | 'FilterEmpty' >> beam.Filter(lambda line: line.strip())) else: input = ( pipeline | 'ReadFromPubSub' >> beam.io.ReadFromPubSub(subscription=known_args.pubsub_subscription) | 'DecodeText' >> beam.Map(lambda x: x.decode('utf-8')) | 'WindowedOutput' >> beam.WindowInto( beam.window.FixedWindows(60), trigger=beam.trigger.AfterProcessingTime(30), accumulation_mode=beam.trigger.AccumulationMode.DISCARDING, allowed_lateness=0)) _ = ( input | 'Tokenize' >> beam.Map(lambda text: tokenize_text(text, tokenizer)) | 'RunInference' >> RunInference(KeyedModelHandler(model_handler)) | 'PostProcess' >> beam.ParDo(SentimentPostProcessor(tokenizer)) | 'WriteToBigQuery' >> beam.io.WriteToBigQuery( known_args.output_table, schema='text:STRING, sentiment:STRING, confidence:FLOAT', write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, method=method)) result = pipeline.run() result.wait_until_finish(duration=1800000) # 30 min result.cancel() cleanup_pubsub_resources( project=known_args.project, topic_path=known_args.pubsub_topic, subscription_path=known_args.pubsub_subscription) return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/rate_limiter_vertex_ai.py0000644000175100017510000000576315206113032026544 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A simple example demonstrating usage of the EnvoyRateLimiter with Vertex AI. """ import argparse import logging import apache_beam as beam from apache_beam.io.components.rate_limiter import EnvoyRateLimiter from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.vertex_ai_inference import VertexAIModelHandlerJSON from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions def run(argv=None): parser = argparse.ArgumentParser() parser.add_argument( '--project', dest='project', help='The Google Cloud project ID for Vertex AI.') parser.add_argument( '--location', dest='location', help='The Google Cloud location (e.g. us-central1) for Vertex AI.') parser.add_argument( '--endpoint_id', dest='endpoint_id', help='The ID of the Vertex AI endpoint.') parser.add_argument( '--rls_address', dest='rls_address', help='The address of the Envoy Rate Limit Service (e.g. localhost:8081).') known_args, pipeline_args = parser.parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = True # [START RateLimiterVertexPy] # Initialize the EnvoyRateLimiter rate_limiter = EnvoyRateLimiter( service_address=known_args.rls_address, domain="mongo_cps", descriptors=[{ "database": "users" }], namespace='example_pipeline') # Initialize the VertexAIModelHandler with the rate limiter model_handler = VertexAIModelHandlerJSON( endpoint_id=known_args.endpoint_id, project=known_args.project, location=known_args.location, rate_limiter=rate_limiter) # [END RateLimiterVertexPy] # Input features for the model features = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0], [10.0, 11.0, 12.0], [13.0, 14.0, 15.0]] with beam.Pipeline(options=pipeline_options) as p: _ = ( p | 'CreateInputs' >> beam.Create(features) | 'RunInference' >> RunInference(model_handler) | 'PrintPredictions' >> beam.Map(logging.info)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/run_inference_side_inputs.py0000644000175100017510000001311015206113032027227 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Used for internal testing. No backwards compatibility. """ import argparse import logging import time from collections.abc import Iterable from collections.abc import Sequence from typing import Optional import apache_beam as beam from apache_beam.ml.inference import base from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.transforms.periodicsequence import PeriodicImpulse from apache_beam.transforms.userstate import CombiningValueStateSpec # create some fake models which returns different inference results. class FakeModelDefault: def predict(self, example: int) -> int: return example class FakeModelAdd(FakeModelDefault): def predict(self, example: int) -> int: return example + 1 class FakeModelSub(FakeModelDefault): def predict(self, example: int) -> int: return example - 1 class FakeModelHandlerReturnsPredictionResult( base.ModelHandler[int, base.PredictionResult, FakeModelDefault]): def __init__(self, clock=None, model_id='model_default'): self.model_id = model_id self._fake_clock = clock def load_model(self): if self._fake_clock: self._fake_clock.current_time_ns += 500_000_000 # 500ms if self.model_id == 'model_add.pkl': return FakeModelAdd() elif self.model_id == 'model_sub.pkl': return FakeModelSub() return FakeModelDefault() def run_inference( self, batch: Sequence[int], model: FakeModelDefault, inference_args=None) -> Iterable[base.PredictionResult]: for example in batch: yield base.PredictionResult( model_id=self.model_id, example=example, inference=model.predict(example)) def update_model_path(self, model_path: Optional[str] = None): self.model_id = model_path if model_path else self.model_id def run(argv=None, save_main_session=True): parser = argparse.ArgumentParser() first_ts = time.time() side_input_interval = 60 main_input_interval = 20 # give some time for dataflow to start. last_ts = first_ts + 1200 mid_ts = (first_ts + last_ts) / 2 _, pipeline_args = parser.parse_known_args(argv) options = PipelineOptions(pipeline_args) options.view_as(SetupOptions).save_main_session = save_main_session class GetModel(beam.DoFn): def process(self, element) -> Iterable[base.ModelMetadata]: if time.time() > mid_ts: yield base.ModelMetadata( model_id='model_add.pkl', model_name='model_add') else: yield base.ModelMetadata( model_id='model_sub.pkl', model_name='model_sub') class _EmitSingletonSideInput(beam.DoFn): COUNT_STATE = CombiningValueStateSpec('count', combine_fn=sum) def process(self, element, count_state=beam.DoFn.StateParam(COUNT_STATE)): _, path = element counter = count_state.read() if counter == 0: count_state.add(1) yield path def validate_prediction_result(x: base.PredictionResult): model_id = x.model_id if model_id == 'model_sub.pkl': assert (x.example == 1 and x.inference == 0) if model_id == 'model_add.pkl': assert (x.example == 1 and x.inference == 2) if model_id == 'model_default': assert (x.example == 1 and x.inference == 1) with beam.Pipeline(options=options) as pipeline: side_input = ( pipeline | "SideInputPColl" >> PeriodicImpulse( first_ts, last_ts, fire_interval=side_input_interval) | "GetModelId" >> beam.ParDo(GetModel()) | "AttachKey" >> beam.Map(lambda x: (x, x)) # due to periodic impulse, which has a start timestamp before # Dataflow pipeline process data, it can trigger in multiple # firings, causing an Iterable instead of singleton. So, using # the _EmitSingletonSideInput DoFn will ensure unique path will be # fired only once. | "GetSingleton" >> beam.ParDo(_EmitSingletonSideInput()) | "ApplySideInputWindow" >> beam.WindowInto( window.GlobalWindows(), trigger=trigger.Repeatedly(trigger.AfterProcessingTime(1)), accumulation_mode=trigger.AccumulationMode.DISCARDING)) model_handler = FakeModelHandlerReturnsPredictionResult() inference_pcoll = ( pipeline | "MainInputPColl" >> PeriodicImpulse( first_ts, last_ts, fire_interval=main_input_interval, apply_windowing=True) | beam.Map(lambda x: 1) | base.RunInference( model_handler=model_handler, model_metadata_pcoll=side_input)) _ = inference_pcoll | "AssertPredictionResult" >> beam.Map( validate_prediction_result) _ = inference_pcoll | "Logging" >> beam.Map(logging.info) if __name__ == '__main__': run() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.1103795 apache_beam-2.74.0/apache_beam/examples/inference/runinference_metrics/0000755000175100017510000000000015206113157025652 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/runinference_metrics/__init__.py0000644000175100017510000000142015206113032027750 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/runinference_metrics/config.py0000644000175100017510000000235515206113032027466 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """The file defines global variables.""" PROJECT_ID = "" REGION = "us-central1" JOB_NAME = "benchmarking-runinference" NUM_WORKERS = 1 TOKENIZER_NAME = "distilbert-base-uncased-finetuned-sst-2-english" MODEL_STATE_DICT_PATH = ( f"gs://{PROJECT_ID}-ml-examples/{TOKENIZER_NAME}/pytorch_model.bin") MODEL_CONFIG_PATH = TOKENIZER_NAME IMG_NAME = "kfp-components-preprocessing/pytorch-gpu" TAG = "latest" DOCKER_IMG = f"{REGION}-docker.pkg.dev/{PROJECT_ID}/{IMG_NAME}:{TAG}" ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/runinference_metrics/main.py0000644000175100017510000001040215206113032027135 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """This file contains the pipeline for loading a ML model, and exploring the different RunInference metrics.""" import argparse import logging import sys import apache_beam as beam import config as cfg from apache_beam.ml.inference import RunInference from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerKeyedTensor from pipeline.options import get_pipeline_options from pipeline.transformations import CustomPytorchModelHandlerKeyedTensor from pipeline.transformations import HuggingFaceStripBatchingWrapper from pipeline.transformations import PostProcessor from pipeline.transformations import Tokenize from transformers import DistilBertConfig def parse_arguments(argv): """ Parses the arguments passed to the command line and returns them as an object Args: argv: The arguments passed to the command line. Returns: The arguments that are being passed in. """ parser = argparse.ArgumentParser(description="benchmark-runinference") parser.add_argument( "-m", "--mode", help="Mode to run pipeline in.", choices=["local", "cloud"], default="local", ) parser.add_argument( "-p", "--project", help="GCP project to run pipeline on.", default=cfg.PROJECT_ID, ) parser.add_argument( "-d", "--device", help="Device to run the dataflow job on", choices=["CPU", "GPU"], default="CPU", ) args, _ = parser.parse_known_args(args=argv) return args def run(): """ Runs the pipeline that loads a transformer based text classification model and does inference on a list of sentences. At the end of pipeline, different metrics like latency, throughput and others are printed. """ args = parse_arguments(sys.argv) inputs = [ "This is the worst food I have ever eaten", "In my soul and in my heart, I’m convinced I’m wrong!", "Be with me always—take any form—drive me mad!"\ "only do not leave me in this abyss, where I cannot find you!", "Do I want to live? Would you like to live with your soul in the grave?", "Honest people don’t hide their deeds.", "Nelly, I am Heathcliff! He’s always,"\ "always in my mind: not as a pleasure,"\ "any more than I am always a pleasure to myself, but as my own being.", ] * 1000 pipeline_options = get_pipeline_options( job_name=cfg.JOB_NAME, num_workers=cfg.NUM_WORKERS, project=args.project, mode=args.mode, device=args.device, ) model_handler_class = ( PytorchModelHandlerKeyedTensor if args.device == "GPU" else CustomPytorchModelHandlerKeyedTensor) device = "cuda:0" if args.device == "GPU" else args.device model_handler = model_handler_class( state_dict_path=cfg.MODEL_STATE_DICT_PATH, model_class=HuggingFaceStripBatchingWrapper, model_params={ "config": DistilBertConfig.from_pretrained(cfg.MODEL_CONFIG_PATH) }, device=device, ) with beam.Pipeline(options=pipeline_options) as pipeline: _ = ( pipeline | "Create inputs" >> beam.Create(inputs) | "Tokenize" >> beam.ParDo(Tokenize(cfg.TOKENIZER_NAME)) | "Inference" >> RunInference(model_handler=KeyedModelHandler(model_handler)) | "Decode Predictions" >> beam.ParDo(PostProcessor())) metrics = pipeline.result.metrics().query(beam.metrics.MetricsFilter()) logging.info(metrics) if __name__ == "__main__": run() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.1113794 apache_beam-2.74.0/apache_beam/examples/inference/runinference_metrics/pipeline/0000755000175100017510000000000015206113157027457 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/runinference_metrics/pipeline/__init__.py0000644000175100017510000000142015206113032031555 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/runinference_metrics/pipeline/options.py0000644000175100017510000000474415206113032031525 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """This file contains the pipeline options to configure the Dataflow pipeline.""" from datetime import datetime from typing import Any import config as cfg from apache_beam.options.pipeline_options import PipelineOptions def get_pipeline_options( project: str, job_name: str, mode: str, device: str, num_workers: int = cfg.NUM_WORKERS, **kwargs: Any, ) -> PipelineOptions: """Function to retrieve the pipeline options. Args: project: GCP project to run on mode: Indicator to run local, cloud or template num_workers: Number of Workers for running the job parallely Returns: Dataflow pipeline options """ job_name = f'{job_name}-{datetime.now().strftime("%Y%m%d%H%M%S")}' staging_bucket = f"gs://{cfg.PROJECT_ID}-ml-examples" # For a list of available options, check: # https://cloud.google.com/dataflow/docs/guides/specifying-exec-params#setting-other-cloud-dataflow-pipeline-options dataflow_options = { "runner": "DirectRunner" if mode == "local" else "DataflowRunner", "job_name": job_name, "project": project, "region": cfg.REGION, "staging_location": f"{staging_bucket}/dflow-staging", "temp_location": f"{staging_bucket}/dflow-temp", "setup_file": "./setup.py", } flags = [] if device == "GPU": flags = [ "--experiment=worker_accelerator=type:nvidia-tesla-p4;count:1;"\ "install-nvidia-driver", ] dataflow_options.update({ "sdk_container_image": cfg.DOCKER_IMG, "machine_type": "n1-standard-4", }) # Optional parameters if num_workers: dataflow_options.update({"num_workers": num_workers}) return PipelineOptions(flags=flags, **dataflow_options) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/runinference_metrics/pipeline/transformations.py0000644000175100017510000000766215206113032033265 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """This file contains the transformations and utility functions for the pipeline.""" import apache_beam as beam import torch from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerKeyedTensor from transformers import DistilBertForSequenceClassification from transformers import DistilBertTokenizer class CustomPytorchModelHandlerKeyedTensor(PytorchModelHandlerKeyedTensor): """Wrapper around PytorchModelHandlerKeyedTensor to load a model on CPU.""" def load_model(self) -> torch.nn.Module: """Loads and initializes a Pytorch model for processing.""" model = self._model_class(**self._model_params) model.to(self._device) file = FileSystems.open(self._state_dict_path, "rb") model.load_state_dict(torch.load(file, map_location=self._device)) model.eval() return model # Can be removed once https://github.com/apache/beam/issues/21863 is fixed class HuggingFaceStripBatchingWrapper(DistilBertForSequenceClassification): """Wrapper around HuggingFace model because RunInference requires a batch as a list of dicts instead of a dict of lists. Another workaround can be found here where they disable batching instead. https://github.com/apache/beam/blob/master/sdks/python/apache_beam/examples/inference/pytorch_language_modeling.py""" def forward(self, **kwargs): output = super().forward(**kwargs) return [dict(zip(output, v)) for v in zip(*output.values())] class Tokenize(beam.DoFn): """A DoFn for tokenizing texts""" def __init__(self, model_name: str): """Initialises a tokenizer based on the model_name""" self._model_name = model_name def setup(self): """Loads the tokenizer""" self._tokenizer = DistilBertTokenizer.from_pretrained(self._model_name) def process(self, text_input: str): """Prepocesses the text using the tokenizer""" # We need to pad the tokens tensors to max length to make sure # that all the tensors are of the same length and hence # stack-able by the RunInference API, normally you would batch first # and tokenize the batch after and pad each tensor # the the max length in the batch. tokens = self._tokenizer( text_input, return_tensors="pt", padding="max_length", max_length=512) # squeeze because tokenization add an extra dimension, which is empty # in this case because we're tokenizing one element at a time. tokens = {key: torch.squeeze(val) for key, val in tokens.items()} return [(text_input, tokens)] class PostProcessor(beam.DoFn): """Postprocess the RunInference output""" def process(self, element): """ Takes the input text and the prediction result, and returns a dictionary with the input text and the softmax probabilities Args: element: The tuple of input text and the prediction result Returns: A list of dictionaries, each containing the input text and the softmax output. """ text_input, prediction_result = element softmax = ( torch.nn.Softmax(dim=-1)( prediction_result.inference["logits"]).detach().numpy()) return [{"input": text_input, "softmax": softmax}] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/runinference_metrics/setup.py0000644000175100017510000000306215206113032027355 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Setup.py module for the workflow's worker utilities. All the workflow related code is gathered in a package that will be built as a source distribution, staged in the staging area for the workflow being run and then installed in the workers when they start running. This behavior is triggered by specifying the --setup_file command line option when running the workflow for remote execution. """ import setuptools from setuptools import find_packages REQUIREMENTS = [ "apache-beam[gcp]==2.41.0", "transformers==4.38.0", "torch==1.13.1" ] setuptools.setup( name="write-to-pubsub-pipeline", version="1.1.1", install_requires=REQUIREMENTS, packages=find_packages(), author="Apache Software Foundation", author_email="dev@beam.apache.org", py_modules=["config"], ) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/sklearn_japanese_housing_regression.py0000644000175100017510000001477615206113032031323 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A pipeline that uses RunInference API on a regression about housing prices. This example uses the japanese housing data from kaggle. https://www.kaggle.com/datasets/nishiodens/japan-real-estate-transaction-prices Since the data has missing fields, this example illustrates how to split data and assign it to the models that are trained on different subsets of features. The predictions are then recombined. In order to set this example up, you will need two things. 1. Build models (or use ours) and reference those via the model directory. 2. Download the data from kaggle and host it. """ import argparse import os from collections.abc import Iterable import pandas import apache_beam as beam from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.sklearn_inference import ModelFileType from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerPandas from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult # yapf: disable MODELS = [ { 'name': 'all_features', 'required_features': [ 'Area', 'Year', 'MinTimeToNearestStation', 'MaxTimeToNearestStation', 'TotalFloorArea', 'Frontage', 'Breadth', 'BuildingYear'] }, { 'name': 'floor_area', 'required_features': ['Area', 'Year', 'TotalFloorArea'] }, { 'name': 'stations', 'required_features': [ 'Area', 'Year', 'MinTimeToNearestStation', 'MaxTimeToNearestStation'] }, { 'name': 'no_features', 'required_features': ['Area', 'Year'] } ] # yapf: enable def sort_by_features(dataframe, max_size): """ Returns an index to a model, based on available data.""" for i, model in enumerate(MODELS): required_features = dataframe[model['required_features']] # A model can only make a prediction if all required features # are present. # required_features is 2D single row, so all() must be called twice. if required_features.notnull().all().all(): return i return -1 class LoadDataframe(beam.DoFn): def process(self, file_name: str) -> Iterable[pandas.DataFrame]: """ Loads data files as a pandas dataframe.""" file = FileSystems.open(file_name, 'rb') dataframe = pandas.read_csv(file) for i in range(dataframe.shape[0]): yield dataframe.iloc[[i]] def report_predictions(prediction_result): true_result = prediction_result.example['TradePrice'].values[0] inference = prediction_result.inference return 'True Price %.0f, Predicted Price %.0f' % (true_result, inference) def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', required=True, help='A single or comma-separated list of files or uris.') parser.add_argument( '--model_path', dest='model_path', required=True, help='A path from where all models can be read.') parser.add_argument( '--output', dest='output', required=True, help='Path to save output predictions.') return parser.parse_known_args(argv) def inference_transform(model_name, model_path): """Returns a RunInference transform.""" model_filename = model_path + model_name + '.pickle' model_loader = SklearnModelHandlerPandas( model_file_type=ModelFileType.PICKLE, model_uri=model_filename) transform_name = 'RunInference ' + model_name return transform_name >> RunInference(model_loader) def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """Entry point. Defines and runs the pipeline.""" known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session requirements_dir = os.path.dirname(os.path.realpath(__file__)) # Pin to the version that we trained the model on. # Sklearn doesn't guarantee compatability between versions. pipeline_options.view_as( SetupOptions ).requirements_file = f'{requirements_dir}/sklearn_examples_requirements.txt' pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) # Input may be a single file or a comma separated list of files. file_names = pipeline | 'FileNames' >> beam.Create( known_args.input.split(',')) loaded_data = file_names | beam.ParDo(LoadDataframe()) # Some examples don't have all features. Pipelines # that expect those fields will fail. There are many ways to deal with # missing data. This example illustrates how to assign predictions to # different models depending upon what data is available. [all, floor_area, stations, no_features] = ( loaded_data | 'Partition' >> beam.Partition(sort_by_features, len(MODELS))) model_path = known_args.model_path prediction_1 = all | inference_transform('all_features', model_path) prediction_2 = floor_area | inference_transform('floor_area', model_path) prediction_3 = stations | inference_transform('stations', model_path) prediction_4 = no_features | inference_transform('no_features', model_path) all_predictions = (prediction_1, prediction_2, prediction_3, prediction_4) flattened_predictions = all_predictions | 'Flatten' >> beam.Flatten() prediction_report = ( flattened_predictions | 'AllPredictions' >> beam.Map(report_predictions)) _ = prediction_report | "WriteOutput" >> beam.io.WriteToText( known_args.output, append_trailing_newlines=True, shard_name_template='') result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/sklearn_mnist_classification.py0000644000175100017510000001164715206113032027740 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A pipeline that uses RunInference API to classify MNIST data. This pipeline takes a text file in which data is comma separated ints. The first column would be the true label and the rest would be the pixel values. The data is processed and then a model trained on the MNIST data would be used to perform the inference. The pipeline writes the prediction to an output file in which users can then compare against the true label. """ import argparse import logging import os from collections.abc import Iterable import apache_beam as beam from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.sklearn_inference import ModelFileType from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerNumpy from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult def process_input(row: str) -> tuple[int, list[int]]: data = row.split(',') label, pixels = int(data[0]), data[1:] pixels = [int(pixel) for pixel in pixels] return label, pixels class PostProcessor(beam.DoFn): """Process the PredictionResult to get the predicted label. Returns a comma separated string with true label and predicted label. """ def process(self, element: tuple[int, PredictionResult]) -> Iterable[str]: label, prediction_result = element prediction = prediction_result.inference yield '{},{}'.format(label, prediction) def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', required=True, help='text file with comma separated int values.') parser.add_argument( '--output', dest='output', required=True, help='Path to save output predictions.') parser.add_argument( '--model_path', dest='model_path', required=True, help='Path to load the Sklearn model for Inference.') parser.add_argument( '--large_model', action='store_true', dest='large_model', default=False, help='Set to true if your model is large enough to run into memory ' 'pressure if you load multiple copies.') return parser.parse_known_args(argv) def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session requirements_dir = os.path.dirname(os.path.realpath(__file__)) # Pin to the version that we trained the model on. # Sklearn doesn't guarantee compatability between versions. pipeline_options.view_as( SetupOptions ).requirements_file = f'{requirements_dir}/sklearn_examples_requirements.txt' # In this example we pass keyed inputs to RunInference transform. # Therefore, we use KeyedModelHandler wrapper over SklearnModelHandlerNumpy. model_loader = KeyedModelHandler( SklearnModelHandlerNumpy( model_file_type=ModelFileType.PICKLE, model_uri=known_args.model_path, large_model=known_args.large_model)) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) label_pixel_tuple = ( pipeline | "ReadFromInput" >> beam.io.ReadFromText(known_args.input) | "PreProcessInputs" >> beam.Map(process_input)) predictions = ( label_pixel_tuple | "RunInference" >> RunInference(model_loader) | "PostProcessOutputs" >> beam.ParDo(PostProcessor())) _ = predictions | "WriteOutput" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/table_row_batch_example.py0000644000175100017510000001400715206113032026637 0ustar00runnerrunner#!/usr/bin/env python # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Complete batch inference example with sample data generation. This script demonstrates how to use the batch table row inference pipeline with automatically generated sample data and model. Usage: # Run complete local example python table_row_batch_example.py # Run with custom parameters python table_row_batch_example.py --num_rows=1000 --num_features=5 """ import argparse import json import logging import os import pickle import sys import tempfile logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) def create_sample_data_and_model(tmpdir, num_rows=100, num_features=3): """Create sample model and data for testing. Args: tmpdir: Temporary directory path num_rows: Number of data rows to generate num_features: Number of features per row Returns: Tuple of (model_path, data_path, feature_columns) """ try: import numpy as np from sklearn.linear_model import LinearRegression except ImportError: logger.error( 'sklearn and numpy are required. ' 'Install with: pip install scikit-learn numpy') sys.exit(1) logger.info('Creating sample model with %s features...', num_features) model = LinearRegression() X_train = np.random.randn(100, num_features) y_train = np.sum(X_train, axis=1) + np.random.randn(100) * 0.1 model.fit(X_train, y_train) model_path = os.path.join(tmpdir, 'model.pkl') with open(model_path, 'wb') as f: pickle.dump(model, f) logger.info(' ✓ Model saved to %s', model_path) logger.info('Generating %s sample data rows...', num_rows) data_path = os.path.join(tmpdir, 'input_data.jsonl') feature_columns = [f'feature{i+1}' for i in range(num_features)] with open(data_path, 'w') as f: for i in range(num_rows): row = {'id': f'row_{i}'} for col in feature_columns: row[col] = float(np.random.randn()) f.write(json.dumps(row) + '\n') logger.info(' ✓ Data saved to %s', data_path) return model_path, data_path, feature_columns def run_example(num_rows=100, num_features=3): """Run complete batch inference example. Args: num_rows: Number of data rows to generate num_features: Number of features per row """ logger.info('=' * 70) logger.info('BATCH TABLE ROW INFERENCE - COMPLETE EXAMPLE') logger.info('=' * 70) with tempfile.TemporaryDirectory() as tmpdir: logger.info('\n[Step 1/4] Creating sample model and data...') model_path, data_path, feature_columns = create_sample_data_and_model( tmpdir, num_rows, num_features) logger.info('\n[Step 2/4] Setting up output paths...') output_file = os.path.join(tmpdir, 'predictions') logger.info(' Output file: %s.jsonl', output_file) logger.info('\n[Step 3/4] Running batch inference pipeline...') logger.info(' Features: %s', feature_columns) cmd = [ sys.executable, 'table_row_inference.py', '--mode=batch', f'--input_file={data_path}', f'--model_path={model_path}', f'--feature_columns={",".join(feature_columns)}', f'--output_file={output_file}', '--runner=DirectRunner' ] logger.info(' Command: %s', ' '.join(cmd)) import subprocess script_dir = os.path.dirname(os.path.abspath(__file__)) result = subprocess.run( cmd, capture_output=True, text=True, check=False, cwd=script_dir) if result.returncode != 0: logger.error('Pipeline failed!') logger.error(result.stderr) sys.exit(1) logger.info(' ✓ Pipeline completed successfully!') logger.info('\n[Step 4/4] Viewing results...') output_path = '%s.jsonl' % output_file if os.path.exists(output_path): with open(output_path, 'r') as f: lines = f.readlines() logger.info(' Total predictions: %s', len(lines)) logger.info('\n Sample predictions (first 5):') for i, line in enumerate(lines[:5]): prediction = json.loads(line) feats = [prediction[f'input_{f}'] for f in feature_columns][:3] logger.info( ' %s. Key: %10s | Prediction: %8.4f | Features: %s...', i + 1, prediction['row_key'], prediction['prediction'], feats) logger.info('\n Full output saved to: %s', output_path) else: logger.warning(' Output file not found: %s', output_path) logger.info('\n' + '=' * 70) logger.info('EXAMPLE COMPLETED SUCCESSFULLY!') logger.info('=' * 70) logger.info('\n📖 Next steps:') logger.info(' 1. Review the generated predictions above') logger.info( ' 2. Try with your own data: ' 'python table_row_inference.py --help') logger.info(' 3. Deploy to Dataflow with --runner=DataflowRunner') logger.info('\n✨ You now understand batch table row inference!') def main(): parser = argparse.ArgumentParser( description='Run batch table row inference example') parser.add_argument( '--num_rows', type=int, default=100, help='Number of data rows to generate (default: 100)') parser.add_argument( '--num_features', type=int, default=3, help='Number of features per row (default: 3)') args = parser.parse_args() run_example(num_rows=args.num_rows, num_features=args.num_features) if __name__ == '__main__': main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/table_row_inference.py0000644000175100017510000003104715206113032026004 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A pipeline that uses RunInference to perform inference on table rows. This pipeline demonstrates ML Pipelines #18: handling continuous new table rows with RunInference using table input models. It reads structured data (table rows) from a streaming source, performs inference while preserving the table schema, and writes results to a table output. The pipeline supports both streaming and batch modes: - Streaming: Reads from Pub/Sub, applies windowing, writes via streaming inserts - Batch: Reads from file, processes all data, writes via file loads Example usage for streaming: python table_row_inference.py \ --mode=streaming \ --input_subscription=projects/PROJECT/subscriptions/SUBSCRIPTION \ --output_table=PROJECT:DATASET.TABLE \ --model_path=gs://BUCKET/model.pkl \ --feature_columns=feature1,feature2,feature3 \ --runner=DataflowRunner \ --project=PROJECT \ --region=REGION \ --temp_location=gs://BUCKET/temp Example usage for batch: python table_row_inference.py \ --mode=batch \ --input_file=gs://BUCKET/input.jsonl \ --output_table=PROJECT:DATASET.TABLE \ --model_path=gs://BUCKET/model.pkl \ --feature_columns=feature1,feature2,feature3 # Batch with file output python table_row_inference.py \ --mode=batch \ --input_file=data.jsonl \ --output_file=predictions.jsonl \ --model_path=model.pkl \ --feature_columns=feature1,feature2,feature3 """ import argparse import hashlib import json import logging from collections.abc import Iterable from typing import Any from typing import Optional import apache_beam as beam import numpy as np from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.sklearn_inference import ModelFileType from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerNumpy from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.runners.runner import PipelineResult class TableRowModelHandler(SklearnModelHandlerNumpy): """ModelHandler that processes table rows (beam.Row objects) for inference. This handler extends SklearnModelHandlerNumpy to work with structured table data represented as beam.Row objects. It extracts specified feature columns from the row and converts them to numpy arrays for model input. Attributes: feature_columns: List of column names to extract as features from input rows """ def __init__( self, model_uri: str, feature_columns: list[str], model_file_type: ModelFileType = ModelFileType.PICKLE): """Initialize the TableRowModelHandler. Args: model_uri: Path to the saved model file (local or GCS) feature_columns: List of column names to use as model features model_file_type: Type of model file (PICKLE or JOBLIB) """ super().__init__(model_uri=model_uri, model_file_type=model_file_type) self.feature_columns = feature_columns def run_inference( self, batch: list[beam.Row], model: Any, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """Run inference on a batch of beam.Row objects. Args: batch: List of beam.Row objects containing input features model: Loaded sklearn model inference_args: Optional additional arguments for inference Yields: PredictionResult containing the original row and prediction """ features_array = [] for row in batch: row_dict = row._asdict() features = [row_dict.get(col, 0.0) for col in self.feature_columns] features_array.append(features) features_array = np.array(features_array, dtype=np.float32) predictions = model.predict(features_array) for row, prediction in zip(batch, predictions): yield PredictionResult( example=row, inference=float(prediction), model_id=self._model_uri) class FormatTableOutput(beam.DoFn): """DoFn that formats inference results into table output schema. Takes PredictionResult objects from KeyedModelHandler and formats them into dictionaries suitable for writing to BigQuery or other table outputs. """ def __init__(self, feature_columns: list[str]): self.feature_columns = feature_columns def process( self, element: tuple[str, PredictionResult]) -> Iterable[dict[str, Any]]: """Process a keyed inference result into table output format. Args: element: Tuple of (row_key, PredictionResult) Yields: Dictionary with all input fields plus prediction and metadata """ key, prediction = element row = prediction.example row_dict = row._asdict() output = {'row_key': key, 'prediction': prediction.inference} if prediction.model_id: output['model_id'] = prediction.model_id for field_name in self.feature_columns: output[f'input_{field_name}'] = row_dict.get(field_name, 0.0) yield output def parse_json_to_table_row( message: bytes, schema_fields: Optional[list[str]] = None) -> tuple[str, beam.Row]: """Parse JSON message to (key, beam.Row) format for KeyedModelHandler. Args: message: JSON-encoded bytes schema_fields: Optional list of expected field names Returns: Tuple of (unique_key, beam.Row with parsed data) """ data = json.loads(message.decode('utf-8')) row_key = data.get('id', hashlib.sha256(message).hexdigest()) row_fields = {} for key, value in data.items(): if key != 'id' and (schema_fields is None or key in schema_fields): if isinstance(value, (int, float)): row_fields[key] = float(value) else: row_fields[key] = value table_row = beam.Row(**row_fields) return row_key, table_row def build_output_schema(feature_columns: list[str]) -> str: """Build BigQuery schema string for output table. Args: feature_columns: List of feature column names Returns: BigQuery schema string """ schema_parts = ['row_key:STRING', 'prediction:FLOAT', 'model_id:STRING'] for col in feature_columns: schema_parts.append(f'input_{col}:FLOAT') return ','.join(schema_parts) def parse_known_args(argv): """Parse command-line arguments for the pipeline.""" parser = argparse.ArgumentParser() parser.add_argument( '--mode', default='batch', choices=['streaming', 'batch'], help='Pipeline mode: streaming or batch') parser.add_argument( '--input_subscription', help='Pub/Sub subscription for streaming mode ' '(format: projects/PROJECT/subscriptions/SUBSCRIPTION)') parser.add_argument( '--input_file', help='Input file path for batch mode (e.g., gs://bucket/input.jsonl)') parser.add_argument( '--output_table', help='BigQuery output table (format: PROJECT:DATASET.TABLE)') parser.add_argument( '--output_file', help='Output file path (JSONL format) for batch mode. ' 'Alternative to or in addition to output_table.') parser.add_argument('--model_path', help='Path to saved model file') parser.add_argument( '--feature_columns', required=True, help='Comma-separated list of feature column names') parser.add_argument( '--window_size_sec', type=int, default=60, help='Window size in seconds for streaming mode (default: 60)') parser.add_argument( '--trigger_interval_sec', type=int, default=30, help='Trigger interval in seconds for streaming mode (default: 30)') parser.add_argument( '--input_expand_factor', type=int, default=1, help='In batch mode: repeat each input line this many times to scale up ' 'volume (e.g. 100k lines × 100 = 10M rows). Default 1 = no expansion.') return parser.parse_known_args(argv) def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """Main pipeline execution function. Args: argv: Command-line arguments save_main_session: Whether to save main session for workers test_pipeline: Optional test pipeline (for testing) Returns: PipelineResult from pipeline execution """ known_args, pipeline_args = parse_known_args(argv) if known_args.mode == 'streaming' and not known_args.input_subscription: raise ValueError('input_subscription is required for streaming mode') if known_args.mode == 'batch' and not known_args.input_file: raise ValueError('input_file is required for batch mode') if known_args.mode == 'streaming' and not known_args.output_table: raise ValueError('output_table is required for streaming mode') if (known_args.mode == 'batch' and not known_args.output_table and not known_args.output_file): raise ValueError( 'In batch mode, specify at least one of --output_table or --output_file' ) feature_columns = [ col.strip() for col in known_args.feature_columns.split(',') ] pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session pipeline_options.view_as(StandardOptions).streaming = ( known_args.mode == 'streaming') model_handler = TableRowModelHandler( model_uri=known_args.model_path, feature_columns=feature_columns) output_schema = build_output_schema(feature_columns) pipeline = test_pipeline or beam.Pipeline(options=pipeline_options) if known_args.mode == 'streaming': input_data = ( pipeline | 'ReadFromPubSub' >> beam.io.ReadFromPubSub(subscription=known_args.input_subscription) | 'ParseToTableRows' >> beam.Map(lambda msg: parse_json_to_table_row(msg, feature_columns)) | 'WindowedData' >> beam.WindowInto( beam.window.FixedWindows(known_args.window_size_sec), trigger=beam.trigger.AfterProcessingTime( known_args.trigger_interval_sec), accumulation_mode=beam.trigger.AccumulationMode.DISCARDING, allowed_lateness=0)) write_method = beam.io.WriteToBigQuery.Method.STREAMING_INSERTS else: read_lines = ( pipeline | 'ReadFromFile' >> beam.io.ReadFromText(known_args.input_file)) expand_factor = getattr(known_args, 'input_expand_factor', 1) or 1 if expand_factor > 1: read_lines = ( read_lines | 'ExpandInput' >> beam.FlatMap(lambda line: [line] * expand_factor)) input_data = ( read_lines | 'ParseToTableRows' >> beam.Map( lambda line: parse_json_to_table_row( line.encode('utf-8'), feature_columns))) write_method = beam.io.WriteToBigQuery.Method.FILE_LOADS write_disposition = ( beam.io.BigQueryDisposition.WRITE_APPEND if known_args.mode == 'streaming' else beam.io.BigQueryDisposition.WRITE_TRUNCATE) formatted = ( input_data | 'RunInference' >> RunInference(KeyedModelHandler(model_handler)) | 'FormatOutput' >> beam.ParDo(FormatTableOutput(feature_columns))) if known_args.output_table: _ = ( formatted | 'WriteToBigQuery' >> beam.io.WriteToBigQuery( known_args.output_table, schema=output_schema, write_disposition=write_disposition, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, method=write_method)) if known_args.mode == 'batch' and known_args.output_file: _ = ( formatted | 'FormatJSON' >> beam.Map(json.dumps) | 'WriteToFile' >> beam.io.WriteToText( known_args.output_file, file_name_suffix='.jsonl', shard_name_template='')) result = pipeline.run() if known_args.mode == 'batch' and not test_pipeline: result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/table_row_inference_test.py0000644000175100017510000001443615206113032027046 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for table row inference pipeline.""" import json import os import pickle import tempfile import unittest import apache_beam as beam import numpy as np from apache_beam.examples.inference.table_row_inference import FormatTableOutput from apache_beam.examples.inference.table_row_inference import TableRowModelHandler from apache_beam.examples.inference.table_row_inference import build_output_schema from apache_beam.examples.inference.table_row_inference import ( parse_json_to_table_row) from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that # Module-level matcher for assert_that (must be picklable; no closure over # self). REQUIRED_OUTPUT_KEYS = ( 'row_key', 'prediction', 'input_feature1', 'input_feature2', 'input_feature3') def _assert_table_inference_outputs(outputs): """Asserts pipeline output has expected structure. Used in assert_that.""" if len(outputs) != 2: raise AssertionError(f'Expected 2 outputs, got {len(outputs)}') for output in outputs: for key in REQUIRED_OUTPUT_KEYS: if key not in output: raise AssertionError(f'Missing key {key!r} in output {output}') try: from sklearn.linear_model import LinearRegression SKLEARN_AVAILABLE = True except ImportError: SKLEARN_AVAILABLE = False class SimpleLinearModel: """Simple model for testing without sklearn dependency.""" def predict(self, X): return np.sum(X, axis=1) @unittest.skipIf(not SKLEARN_AVAILABLE, 'sklearn is not available') class TableRowInferenceTest(unittest.TestCase): def setUp(self): self.tmp_dir = tempfile.mkdtemp() self.model_path = os.path.join(self.tmp_dir, 'test_model.pkl') model = LinearRegression() X_train = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) y_train = np.array([6, 15, 24]) model.fit(X_train, y_train) with open(self.model_path, 'wb') as f: pickle.dump(model, f) def test_parse_json_to_table_row(self): json_data = json.dumps({ 'id': 'test_1', 'feature1': 1.0, 'feature2': 2.0, 'feature3': 3.0 }).encode('utf-8') key, row = parse_json_to_table_row( json_data, schema_fields=['feature1', 'feature2', 'feature3']) self.assertEqual(key, 'test_1') self.assertEqual(row.feature1, 1.0) self.assertEqual(row.feature2, 2.0) self.assertEqual(row.feature3, 3.0) def test_build_output_schema(self): feature_cols = ['feature1', 'feature2', 'feature3'] schema = build_output_schema(feature_cols) expected_fields = [ 'row_key', 'prediction', 'model_id', 'input_feature1', 'input_feature2', 'input_feature3' ] for field in expected_fields: self.assertIn(field, schema) def test_table_row_model_handler(self): model_handler = TableRowModelHandler( model_uri=self.model_path, feature_columns=['f1', 'f2', 'f3']) model = model_handler.load_model() test_rows = [ beam.Row(f1=1.0, f2=2.0, f3=3.0), beam.Row(f1=4.0, f2=5.0, f3=6.0), ] results = list(model_handler.run_inference(test_rows, model)) self.assertEqual(len(results), 2) self.assertIsInstance(results[0], PredictionResult) self.assertEqual(results[0].example, test_rows[0]) self.assertIsNotNone(results[0].inference) def test_format_table_output(self): row = beam.Row(feature1=1.0, feature2=2.0, feature3=3.0) prediction_result = PredictionResult( example=row, inference=6.0, model_id='test_model') keyed_result = ('test_key', prediction_result) feature_columns = ['feature1', 'feature2', 'feature3'] formatter = FormatTableOutput(feature_columns=feature_columns) outputs = list(formatter.process(keyed_result)) self.assertEqual(len(outputs), 1) output = outputs[0] self.assertEqual(output['row_key'], 'test_key') self.assertEqual(output['prediction'], 6.0) self.assertEqual(output['model_id'], 'test_model') self.assertEqual(output['input_feature1'], 1.0) self.assertEqual(output['input_feature2'], 2.0) self.assertEqual(output['input_feature3'], 3.0) def test_pipeline_integration(self): test_data = [ json.dumps({ 'id': 'row_1', 'feature1': 1.0, 'feature2': 2.0, 'feature3': 3.0 }), json.dumps({ 'id': 'row_2', 'feature1': 4.0, 'feature2': 5.0, 'feature3': 6.0 }), ] feature_columns = ['feature1', 'feature2', 'feature3'] model_handler = TableRowModelHandler( model_uri=self.model_path, feature_columns=feature_columns) with TestPipeline() as p: input_data = ( p | beam.Create(test_data) | beam.Map( lambda line: parse_json_to_table_row( line.encode('utf-8'), feature_columns))) predictions = ( input_data | RunInference(KeyedModelHandler(model_handler)) | beam.ParDo(FormatTableOutput(feature_columns=feature_columns))) assert_that(predictions, _assert_table_inference_outputs) class TableRowInferenceNoSklearnTest(unittest.TestCase): """Tests that don't require sklearn.""" def test_parse_json_without_schema(self): json_data = json.dumps({'id': 'test', 'value': 123}).encode('utf-8') key, row = parse_json_to_table_row(json_data) self.assertEqual(key, 'test') self.assertTrue(hasattr(row, 'value')) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/table_row_inference_utils.py0000644000175100017510000002335615206113032027230 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utility functions for table row inference pipeline. This module provides helper functions for testing and deploying the table row inference pipeline, including data generation, model creation, and Pub/Sub resource management. """ import json import logging import pickle import time from typing import Optional try: from google.api_core import exceptions as api_exceptions from google.cloud import pubsub_v1 PUBSUB_AVAILABLE = True except ImportError: PUBSUB_AVAILABLE = False api_exceptions = None logging.warning('google-cloud-pubsub not available') try: import numpy as np from sklearn.linear_model import LinearRegression SKLEARN_AVAILABLE = True except ImportError: SKLEARN_AVAILABLE = False logging.warning('sklearn not available') def create_sample_model(output_path: str, num_features: int = 3): """Create and save a simple linear regression model for testing. Args: output_path: Path where to save the model (local or GCS) num_features: Number of input features """ if not SKLEARN_AVAILABLE: raise ImportError('sklearn is required to create sample models') model = LinearRegression() X_train = np.random.randn(100, num_features) y_train = np.sum(X_train, axis=1) + np.random.randn(100) * 0.1 model.fit(X_train, y_train) with open(output_path, 'wb') as f: pickle.dump(model, f) logging.info('Sample model saved to %s', output_path) def generate_sample_data(num_rows: int = 100, num_features: int = 3) -> list[dict]: """Generate sample table row data for testing. Args: num_rows: Number of rows to generate num_features: Number of features per row Returns: List of dictionaries representing table rows """ if not SKLEARN_AVAILABLE: raise ImportError('numpy is required to generate sample data') data = [] for i in range(num_rows): row = {'id': f'row_{i}', 'timestamp': time.time()} for j in range(num_features): row[f'feature{j+1}'] = float(np.random.randn()) data.append(row) return data def write_data_to_file(data: list[dict], output_path: str): """Write sample data to JSONL file. Args: data: List of data dictionaries output_path: Output file path """ with open(output_path, 'w') as f: for row in data: f.write(json.dumps(row) + '\n') logging.info('Wrote %d rows to %s', len(data), output_path) def publish_to_pubsub( project: str, topic_name: str, data: list[dict], rate_limit: Optional[float] = None): """Publish sample data to Pub/Sub topic. Args: project: GCP project ID topic_name: Pub/Sub topic name data: List of data dictionaries to publish rate_limit: Optional rate limit (rows per second) """ if not PUBSUB_AVAILABLE: raise ImportError('google-cloud-pubsub is required for Pub/Sub operations') publisher = pubsub_v1.PublisherClient() topic_path = publisher.topic_path(project, topic_name) delay = 1.0 / rate_limit if rate_limit else 0 for i, row in enumerate(data): message = json.dumps(row).encode('utf-8') publisher.publish(topic_path, message) if (i + 1) % 100 == 0: logging.info('Published %d messages', i + 1) if delay > 0: time.sleep(delay) logging.info('Published %d messages to %s', len(data), topic_path) def ensure_pubsub_topic(project: str, topic_name: str) -> str: """Create Pub/Sub topic if it doesn't exist. Args: project: GCP project ID topic_name: Pub/Sub topic name Returns: Full topic path """ if not PUBSUB_AVAILABLE: raise ImportError('google-cloud-pubsub is required for Pub/Sub operations') publisher = pubsub_v1.PublisherClient() topic_path = publisher.topic_path(project, topic_name) try: publisher.get_topic(request={'topic': topic_path}) logging.info('Topic %s already exists', topic_name) except api_exceptions.NotFound: try: publisher.create_topic(name=topic_path) logging.info('Created topic %s', topic_name) except api_exceptions.AlreadyExists: logging.info('Topic %s was created by another process', topic_name) return topic_path def ensure_pubsub_subscription( project: str, topic_name: str, subscription_name: str) -> str: """Create Pub/Sub subscription if it doesn't exist. Args: project: GCP project ID topic_name: Pub/Sub topic name subscription_name: Subscription name Returns: Full subscription path """ if not PUBSUB_AVAILABLE: raise ImportError('google-cloud-pubsub is required for Pub/Sub operations') publisher = pubsub_v1.PublisherClient() subscriber = pubsub_v1.SubscriberClient() topic_path = publisher.topic_path(project, topic_name) subscription_path = subscriber.subscription_path(project, subscription_name) try: subscriber.get_subscription(request={'subscription': subscription_path}) logging.info('Subscription %s already exists', subscription_name) except api_exceptions.NotFound: try: subscriber.create_subscription(name=subscription_path, topic=topic_path) logging.info('Created subscription %s', subscription_name) except api_exceptions.AlreadyExists: logging.info( 'Subscription %s was created by another process', subscription_name) return subscription_path def cleanup_pubsub_resources( project: str, topic_name: str, subscription_name: Optional[str] = None): """Delete Pub/Sub topic and optionally subscription. Args: project: GCP project ID topic_name: Pub/Sub topic name subscription_name: Optional subscription name to delete """ if not PUBSUB_AVAILABLE: raise ImportError('google-cloud-pubsub is required for Pub/Sub operations') publisher = pubsub_v1.PublisherClient() subscriber = pubsub_v1.SubscriberClient() if subscription_name: subscription_path = subscriber.subscription_path(project, subscription_name) try: subscriber.delete_subscription( request={'subscription': subscription_path}) logging.info('Deleted subscription %s', subscription_name) except api_exceptions.NotFound: logging.info( 'Subscription %s not found or already deleted', subscription_name) except Exception as e: logging.warning('Failed to delete subscription: %s', e) topic_path = publisher.topic_path(project, topic_name) try: publisher.delete_topic(request={'topic': topic_path}) logging.info('Deleted topic %s', topic_name) except api_exceptions.NotFound: logging.info('Topic %s not found or already deleted', topic_name) except Exception as e: logging.warning('Failed to delete topic: %s', e) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) import argparse parser = argparse.ArgumentParser( description='Utility for table row inference pipeline') parser.add_argument( '--action', required=True, choices=[ 'create_model', 'generate_data', 'publish_data', 'create_topic', 'create_subscription', 'cleanup' ], help='Action to perform') parser.add_argument('--project', help='GCP project ID') parser.add_argument('--topic', help='Pub/Sub topic name') parser.add_argument('--subscription', help='Pub/Sub subscription name') parser.add_argument('--output_path', help='Output path for model or data') parser.add_argument( '--num_rows', type=int, default=100, help='Number of rows to generate') parser.add_argument( '--num_features', type=int, default=3, help='Number of features per row') parser.add_argument( '--rate_limit', type=float, help='Rate limit for publishing (rows/sec)') args = parser.parse_args() if args.action == 'create_model': if not args.output_path: raise ValueError('--output_path required for create_model') create_sample_model(args.output_path, args.num_features) elif args.action == 'generate_data': if not args.output_path: raise ValueError('--output_path required for generate_data') data = generate_sample_data(args.num_rows, args.num_features) write_data_to_file(data, args.output_path) elif args.action == 'publish_data': if not args.project or not args.topic: raise ValueError('--project and --topic required for publish_data') data = generate_sample_data(args.num_rows, args.num_features) publish_to_pubsub(args.project, args.topic, data, args.rate_limit) elif args.action == 'create_topic': if not args.project or not args.topic: raise ValueError('--project and --topic required for create_topic') ensure_pubsub_topic(args.project, args.topic) elif args.action == 'create_subscription': if not args.project or not args.topic or not args.subscription: raise ValueError( '--project, --topic, and --subscription required for ' 'create_subscription') ensure_pubsub_subscription(args.project, args.topic, args.subscription) elif args.action == 'cleanup': if not args.project or not args.topic: raise ValueError('--project and --topic required for cleanup') cleanup_pubsub_resources(args.project, args.topic, args.subscription) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/tensorflow_imagenet_segmentation.py0000644000175100017510000001053215206113032030634 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import argparse import logging from collections.abc import Iterable from collections.abc import Iterator import numpy import apache_beam as beam import tensorflow as tf from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.tensorflow_inference import TFModelHandlerTensor from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from PIL import Image class PostProcessor(beam.DoFn): """Process the PredictionResult to get the predicted label. Returns predicted label. """ def setup(self): labels_path = tf.keras.utils.get_file( 'ImageNetLabels.txt', 'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt' # pylint: disable=line-too-long ) self._imagenet_labels = numpy.array(open(labels_path).read().splitlines()) def process(self, element: PredictionResult) -> Iterable[str]: predicted_class = numpy.argmax(element.inference, axis=-1) predicted_class_name = self._imagenet_labels[predicted_class] yield predicted_class_name.title() def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', required=True, help='Path to the text file containing image names.') parser.add_argument( '--output', dest='output', required=True, help='Path to save output predictions.') parser.add_argument( '--model_path', dest='model_path', required=True, help='Path to load the Tensorflow model for Inference.') parser.add_argument( '--image_dir', help='Path to the directory where images are stored.') return parser.parse_known_args(argv) def filter_empty_lines(text: str) -> Iterator[str]: if len(text.strip()) > 0: yield text def read_image(image_name, image_dir): img = tf.keras.utils.get_file(image_name, image_dir + image_name) img = Image.open(img).resize((224, 224)) img = numpy.array(img) / 255.0 img_tensor = tf.cast(tf.convert_to_tensor(img[...]), dtype=tf.float32) return img_tensor def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session # In this example we will use the TensorflowHub model URL. model_loader = TFModelHandlerTensor( model_uri=known_args.model_path).with_preprocess_fn( lambda image_name: read_image(image_name, known_args.image_dir)) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) image = ( pipeline | 'ReadImageNames' >> beam.io.ReadFromText(known_args.input) | 'FilterEmptyLines' >> beam.ParDo(filter_empty_lines)) predictions = ( image | "RunInference" >> RunInference(model_loader) | "PostProcessOutputs" >> beam.ParDo(PostProcessor())) _ = predictions | "WriteOutput" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/tensorflow_mnist_classification.py0000644000175100017510000001047215206113032030476 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import argparse import logging from collections.abc import Iterable import numpy import apache_beam as beam from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.tensorflow_inference import ModelType from apache_beam.ml.inference.tensorflow_inference import TFModelHandlerNumpy from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult def process_input(row: str) -> tuple[int, numpy.ndarray]: data = row.split(',') label, pixels = int(data[0]), data[1:] pixels = [int(pixel) for pixel in pixels] # the trained model accepts the input of shape 28x28 pixels = numpy.array(pixels).reshape((28, 28, 1)) return label, pixels class PostProcessor(beam.DoFn): """Process the PredictionResult to get the predicted label. Returns a comma separated string with true label and predicted label. """ def process(self, element: tuple[int, PredictionResult]) -> Iterable[str]: label, prediction_result = element prediction = numpy.argmax(prediction_result.inference, axis=0) yield '{},{}'.format(label, prediction) def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', required=True, help='text file with comma separated int values.') parser.add_argument( '--output', dest='output', required=True, help='Path to save output predictions.') parser.add_argument( '--model_path', dest='model_path', required=True, help='Path to load the Tensorflow model for Inference.') parser.add_argument( '--large_model', action='store_true', dest='large_model', default=False, help='Set to true if your model is large enough to run into memory ' 'pressure if you load multiple copies.') return parser.parse_known_args(argv) def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session # In this example we pass keyed inputs to RunInference transform. # Therefore, we use KeyedModelHandler wrapper over TFModelHandlerNumpy. model_loader = KeyedModelHandler( TFModelHandlerNumpy( model_uri=known_args.model_path, model_type=ModelType.SAVED_MODEL, large_model=known_args.large_model)) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) label_pixel_tuple = ( pipeline | "ReadFromInput" >> beam.io.ReadFromText(known_args.input) | "PreProcessInputs" >> beam.Map(process_input)) predictions = ( label_pixel_tuple | "RunInference" >> RunInference(model_loader) | "PostProcessOutputs" >> beam.ParDo(PostProcessor())) _ = predictions | "WriteOutput" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/tensorflow_mnist_with_weights.py0000644000175100017510000000707615206113032030216 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import apache_beam as beam import tensorflow as tf from apache_beam.examples.inference.tensorflow_mnist_classification import PostProcessor from apache_beam.examples.inference.tensorflow_mnist_classification import parse_known_args from apache_beam.examples.inference.tensorflow_mnist_classification import process_input from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.tensorflow_inference import ModelType from apache_beam.ml.inference.tensorflow_inference import TFModelHandlerNumpy from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult def get_model(): inputs = tf.keras.layers.Input(shape=(28, 28, 1)) x = tf.keras.layers.Conv2D(32, 3, activation="relu")(inputs) x = tf.keras.layers.Conv2D(32, 3, activation="relu")(x) x = tf.keras.layers.MaxPooling2D(2)(x) x = tf.keras.layers.Conv2D(64, 3, activation="relu")(x) x = tf.keras.layers.Conv2D(64, 3, activation="relu")(x) x = tf.keras.layers.MaxPooling2D(2)(x) x = tf.keras.layers.Flatten()(x) x = tf.keras.layers.Dropout(0.2)(x) outputs = tf.keras.layers.Dense(10, activation='softmax')(x) model = tf.keras.Model(inputs, outputs) return model def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session # In this example we pass keyed inputs to RunInference transform. # Therefore, we use KeyedModelHandler wrapper over TFModelHandlerNumpy. model_loader = KeyedModelHandler( TFModelHandlerNumpy( model_uri=known_args.model_path, model_type=ModelType.SAVED_WEIGHTS, create_model_fn=get_model)) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) label_pixel_tuple = ( pipeline | "ReadFromInput" >> beam.io.ReadFromText(known_args.input) | "PreProcessInputs" >> beam.Map(process_input)) predictions = ( label_pixel_tuple | "RunInference" >> RunInference(model_loader) | "PostProcessOutputs" >> beam.ParDo(PostProcessor())) _ = predictions | "WriteOutput" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/tensorrt_object_detection.py0000644000175100017510000001620115206113032027247 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A pipeline that uses RunInference API to perform object detection with TensorRT. """ import argparse import io import os from collections.abc import Iterable from typing import Optional import numpy as np import apache_beam as beam from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.tensorrt_inference import TensorRTEngineHandlerNumPy # pylint: disable=line-too-long from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from PIL import Image COCO_OBJ_DET_CLASSES = [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'street sign', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'hat', 'backpack', 'umbrella', 'shoe', 'eye glasses', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'plate', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'mirror', 'dining table', 'window', 'desk', 'toilet', 'door', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'blender', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush', 'hair brush', ] def attach_im_size_to_key( data: tuple[str, Image.Image]) -> tuple[tuple[str, int, int], Image.Image]: filename, image = data width, height = image.size return ((filename, width, height), image) def read_image(image_file_name: str, path_to_dir: Optional[str] = None) -> tuple[str, Image.Image]: if path_to_dir is not None: image_file_name = os.path.join(path_to_dir, image_file_name) with FileSystems().open(image_file_name, 'r') as file: data = Image.open(io.BytesIO(file.read())).convert('RGB') return image_file_name, data def preprocess_image(image: Image.Image) -> np.ndarray: ssd_mobilenet_v2_320x320_input_dims = (300, 300) image = image.resize( ssd_mobilenet_v2_320x320_input_dims, resample=Image.Resampling.BILINEAR) image = np.expand_dims(np.asarray(image, dtype=np.float32), axis=0) return image class PostProcessor(beam.DoFn): """Processes the PredictionResult that consists of number of detections per image, box coordinates, scores and classes. We loop over all detections to organize attributes on a per detection basis. Box coordinates are normalized, hence we have to scale them according to original image dimensions. Score is a floating point number that provides probability percentage of a particular object. Class is an integer that we can transform into actual string class using COCO_OBJ_DET_CLASSES as reference. """ def process(self, element: tuple[str, PredictionResult]) -> Iterable[str]: key, prediction_result = element filename, im_width, im_height = key num_detections = prediction_result.inference[0] boxes = prediction_result.inference[1] scores = prediction_result.inference[2] classes = prediction_result.inference[3] detections = [] for i in range(int(num_detections[0])): detections.append({ 'ymin': str(boxes[i][0] * im_height), 'xmin': str(boxes[i][1] * im_width), 'ymax': str(boxes[i][2] * im_height), 'xmax': str(boxes[i][3] * im_width), 'score': str(scores[i]), 'class': COCO_OBJ_DET_CLASSES[int(classes[i])] }) yield filename + ',' + str(detections) def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', required=True, help='Path to the text file containing image names.') parser.add_argument( '--output', dest='output', required=True, help='Path where to save output predictions.' ' text file.') parser.add_argument( '--engine_path', dest='engine_path', required=True, help='Path to the pre-built TFOD ssd_mobilenet_v2_320x320_coco17_tpu-8' 'TensorRT engine.') parser.add_argument( '--images_dir', default=None, help='Path to the directory where images are stored.' 'Not required if image names in the input file have absolute path.') return parser.parse_known_args(argv) def run(argv=None, save_main_session=True): """ Args: argv: Command line arguments defined for this example. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session engine_handler = KeyedModelHandler( TensorRTEngineHandlerNumPy( min_batch_size=1, max_batch_size=1, engine_path=known_args.engine_path)) with beam.Pipeline(options=pipeline_options) as p: filename_value_pair = ( p | 'ReadImageNames' >> beam.io.ReadFromText(known_args.input) | 'ReadImageData' >> beam.Map( lambda image_name: read_image( image_file_name=image_name, path_to_dir=known_args.images_dir)) | 'AttachImageSizeToKey' >> beam.Map(attach_im_size_to_key) | 'PreprocessImages' >> beam.MapTuple( lambda file_name, data: (file_name, preprocess_image(data)))) predictions = ( filename_value_pair | 'TensorRTRunInference' >> RunInference(engine_handler) | 'ProcessOutput' >> beam.ParDo(PostProcessor())) _ = ( predictions | "WriteOutputToGCS" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True)) if __name__ == '__main__': run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/tensorrt_text_classification.py0000644000175100017510000001024715206113032030006 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A pipeline to demonstrate usage of TensorRT with RunInference for a text classification model. This pipeline reads data from a text file, preprocesses the data, and then uses RunInference to generate predictions from the text classification TensorRT engine. Next, it postprocesses the RunInference outputs to print the input and the predicted class label. It also prints metrics provided by RunInference. """ import argparse import logging import numpy as np import apache_beam as beam from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.tensorrt_inference import TensorRTEngineHandlerNumPy from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from transformers import AutoTokenizer class Preprocess(beam.DoFn): """Processes the input sentences to tokenize them. The input sentences are tokenized because the model is expecting tokens. """ def __init__(self, tokenizer: AutoTokenizer): self._tokenizer = tokenizer def process(self, element): inputs = self._tokenizer( element, return_tensors="np", padding="max_length", max_length=128) return inputs.input_ids class Postprocess(beam.DoFn): """Processes the PredictionResult to get the predicted class. The logits are the output of the TensorRT engine. We can get the class label by getting the index of maximum logit using argmax. """ def __init__(self, tokenizer: AutoTokenizer): self._tokenizer = tokenizer def process(self, element): decoded_input = self._tokenizer.decode( element.example, skip_special_tokens=True) logits = element.inference[0] argmax = np.argmax(logits) output = "Positive" if argmax == 1 else "Negative" yield decoded_input, output def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', required=True, help='Path to the text file containing sentences.') parser.add_argument( '--trt_model_path', dest='trt_model_path', required=True, help='Path to the pre-built textattack/bert-base-uncased-SST-2' 'TensorRT engine.') parser.add_argument( '--model_id', dest='model_id', default="textattack/bert-base-uncased-SST-2", help="name of model.") return parser.parse_known_args(argv) def run( argv=None, save_main_session=True, ): known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session model_handler = TensorRTEngineHandlerNumPy( min_batch_size=1, max_batch_size=1, engine_path=known_args.trt_model_path, ) tokenizer = AutoTokenizer.from_pretrained(known_args.model_id) with beam.Pipeline(options=pipeline_options) as pipeline: _ = ( pipeline | "ReadSentences" >> beam.io.ReadFromText(known_args.input) | "Preprocess" >> beam.ParDo(Preprocess(tokenizer=tokenizer)) | "RunInference" >> RunInference(model_handler=model_handler) | "PostProcess" >> beam.ParDo(Postprocess(tokenizer=tokenizer)) | "LogResult" >> beam.Map(logging.info)) metrics = pipeline.result.metrics().query(beam.metrics.MetricsFilter()) logging.info(metrics) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.1133797 apache_beam-2.74.0/apache_beam/examples/inference/tfx_bsl/0000755000175100017510000000000015206113157023102 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/tfx_bsl/__init__.py0000644000175100017510000000142015206113032025200 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/tfx_bsl/build_tensorflow_model.py0000644000175100017510000001356415206113032030216 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # Intended only for internal testing. from typing import Optional import tensorflow as tf class TFModelWrapperWithSignature(tf.keras.Model): """ Helper class used to wrap a based tf.keras.Model object with a serving signature that can passed to the tfx_bsl RunInference transform. A TF model saved using this helper class expects inputs as images serialized to tf.string using tf.io.parse_tensor and then passing serialized images to the RunInference transform in the tf.train.Example. More about tf.train.Example at https://www.tensorflow.org/api_docs/python/tf/train/Example Usage: Step 1: # Save the base TF model with modified signature . signature_model = TFModelWrapperWithSignature( model=model, preprocess_input=preprocess_input, input_dtype=input_dtype, feature_description=feature_description, **kwargs ) tf.saved_model.save(signature_model, path) Step 2: # Load the saved_model in the beam pipeline to create ModelHandler. saved_model_spec = model_spec_pb2.SavedModelSpec( model_path=known_args.model_path) inferece_spec_type = model_spec_pb2.InferenceSpecType( saved_model_spec=saved_model_spec) model_handler = CreateModelHandler(inferece_spec_type) """ def __init__( self, model, preprocess_input=None, input_dtype=None, feature_description=None, **kwargs): """ model: model: Base tensorflow model used for TFX-BSL RunInference transform. preprocess_input: Preprocess method to be included as part of the model's serving signature. input_dtype: tf dtype of the inputs passed to the model. For eg: tf.int32, tf.uint8. feature_description: Feature spec to parse inputs from tf.train.Example using tf.parse_example(). For more details, please take a look at https://www.tensorflow.org/api_docs/python/tf/io/parse_example If there are extra arguments(for eg: training=False) that should be passed to the base tf model during inference, please pass them in kwargs. """ super().__init__() self.model = model self.preprocess_input = preprocess_input self.input_dtype = input_dtype self.feature_description = feature_description if not feature_description: self.feature_description = {'image': tf.io.FixedLenFeature((), tf.string)} self._kwargs = kwargs @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)]) def call(self, serialized_examples): features = tf.io.parse_example( serialized_examples, features=self.feature_description) # Initialize a TensorArray to store the deserialized values. # For more details, please look at # https://github.com/tensorflow/tensorflow/issues/39323#issuecomment-627586602 num_batches = len(features['image']) deserialized_vectors = tf.TensorArray( self.input_dtype, size=num_batches, dynamic_size=True) # Vectorized version of tf.io.parse_tensor is not available. # Use for loop to vectorize the tensor. For more details, refer # https://github.com/tensorflow/tensorflow/issues/43706 for i in range(num_batches): deserialized_value = tf.io.parse_tensor( features['image'][i], out_type=self.input_dtype) # In Graph mode, return value must get assigned in order to # update the array. More details at # http://github.com/tensorflow/tensorflow/issues/30409#issuecomment-508962873 deserialized_vectors = deserialized_vectors.write(i, deserialized_value) deserialized_tensor = deserialized_vectors.stack() if self.preprocess_input: deserialized_tensor = self.preprocess_input(deserialized_tensor) return self.model(deserialized_tensor, **self._kwargs) def save_tf_model_with_signature( path_to_save_model, model=None, preprocess_input=None, input_dtype=tf.float32, feature_description: Optional[dict] = None, **kwargs, ): """ Helper function used to save the Tensorflow Model with a serving signature. This is intended only for internal testing. Args: path_to_save_model: Path to save the model with modified signature. model: model: Base tensorflow model used for TFX-BSL RunInference transform. preprocess_input: Preprocess method to be included as part of the model's serving signature. input_dtype: tf dtype of the inputs passed to the model. For eg: tf.int32, tf.uint8. feature_description: Feature spec to parse inputs from tf.train.Example using tf.parse_example(). For more details, please take a look at https://www.tensorflow.org/api_docs/python/tf/io/parse_example If there are extra arguments(for eg: training=False) that should be passed to the base tf model during inference, please pass them in kwargs. """ if not model: model = tf.keras.applications.MobileNetV2(weights='imagenet') preprocess_input = tf.keras.applications.mobilenet_v2.preprocess_input signature_model = TFModelWrapperWithSignature( model=model, preprocess_input=preprocess_input, input_dtype=input_dtype, feature_description=feature_description, **kwargs) tf.saved_model.save(signature_model, path_to_save_model) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/tfx_bsl/tensorflow_image_classification.py0000644000175100017510000001556615206113032032100 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A sample pipeline illustrating how to use Apache Beam RunInference with TFX_BSL CreateModelHandler API. For more details, please look at https://www.tensorflow.org/tfx/tfx_bsl/api_docs/python/tfx_bsl/public/beam/run_inference/CreateModelHandler. Note: A Tensorflow Model needs to be updated with a @tf.function signature in order to accept bytes as inputs, and should have logic to decode bytes to Tensors that is acceptable by the TensorFlow model. Please take a look at TFModelWrapperWithSignature class in build_tensorflow_model.py on how to modify TF Model's signature and the logic to decode the image tensor. """ import argparse import io import logging import os from collections.abc import Iterable from collections.abc import Iterator from typing import Optional import apache_beam as beam import tensorflow as tf from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import RunInference from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from PIL import Image from tfx_bsl.public.beam.run_inference import CreateModelHandler from tfx_bsl.public.beam.run_inference import prediction_log_pb2 from tfx_bsl.public.proto import model_spec_pb2 _IMG_SIZE = (224, 224) def filter_empty_lines(text: str) -> Iterator[str]: if len(text.strip()) > 0: yield text def read_and_process_image( image_file_name: str, path_to_dir: Optional[str] = None) -> tuple[str, tf.Tensor]: if path_to_dir is not None: image_file_name = os.path.join(path_to_dir, image_file_name) with FileSystems().open(image_file_name, 'r') as file: data = Image.open(io.BytesIO(file.read())).convert('RGB') # Note: Converts the image dtype from uint8 to float32 # https://www.tensorflow.org/api_docs/python/tf/image/resize image = tf.keras.preprocessing.image.img_to_array(data) image = tf.image.resize(image, _IMG_SIZE) return image_file_name, image def convert_image_to_example_proto(tensor: tf.Tensor) -> tf.train.Example: """ This method performs the following: 1. Accepts the tensor as input 2. Serializes the tensor into bytes and pass it through tf.train.Feature 3. Pass the serialized tensor feature using tf.train.Example Proto to the RunInference transform. Args: tensor: A TF tensor. Returns: example_proto: A tf.train.Example containing serialized tensor. """ serialized_non_scalar = tf.io.serialize_tensor(tensor) feature_of_bytes = tf.train.Feature( bytes_list=tf.train.BytesList(value=[serialized_non_scalar.numpy()])) features_for_example = {'image': feature_of_bytes} example_proto = tf.train.Example( features=tf.train.Features(feature=features_for_example)) return example_proto class ProcessInferenceToString(beam.DoFn): def process( self, element: tuple[str, prediction_log_pb2.PredictionLog]) -> Iterable[str]: """ Args: element: Tuple of str, and PredictionLog. Inference can be parsed from prediction_log returns: str of filename and inference. """ filename, predict_log = element[0], element[1].predict_log output_value = predict_log.response.outputs output_tensor = ( tf.io.decode_raw( output_value['output_0'].tensor_content, out_type=tf.float32)) max_index_output_tensor = tf.math.argmax(output_tensor, axis=0) yield filename + ',' + str(tf.get_static_value(max_index_output_tensor)) def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', required=True, help='Path to the text file containing image names.') parser.add_argument( '--output', dest='output', required=True, help='Path to save output predictions text file.') parser.add_argument( '--model_path', dest='model_path', required=True, help="Path to the model.") parser.add_argument( '--images_dir', default=None, help='Path to the directory where images are stored.' 'Not required if image names in the input file have absolute path.') return parser.parse_known_args(argv) def run(argv=None, save_main_session=True, pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session saved_model_spec = model_spec_pb2.SavedModelSpec( model_path=known_args.model_path) inferece_spec_type = model_spec_pb2.InferenceSpecType( saved_model_spec=saved_model_spec) model_handler = CreateModelHandler(inferece_spec_type) # create a KeyedModelHandler to accommodate image names as keys. keyed_model_handler = KeyedModelHandler(model_handler) if not pipeline: pipeline = beam.Pipeline(options=pipeline_options) filename_value_pair = ( pipeline | 'ReadImageNames' >> beam.io.ReadFromText(known_args.input) | 'FilterEmptyLines' >> beam.ParDo(filter_empty_lines) | 'ProcessImageData' >> beam.Map( lambda image_name: read_and_process_image( image_file_name=image_name, path_to_dir=known_args.images_dir))) predictions = ( filename_value_pair | 'ConvertToExampleProto' >> beam.Map(lambda x: (x[0], convert_image_to_example_proto(x[1]))) | 'TFXRunInference' >> RunInference(keyed_model_handler) | 'PostProcess' >> beam.ParDo(ProcessInferenceToString())) _ = ( predictions | "WriteOutputToGCS" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True)) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/tfx_bsl/tfx_bsl_inference_it_test.py0000644000175100017510000001037315206113032030662 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest import uuid import pytest from apache_beam.io.filesystems import FileSystems from apache_beam.testing.test_pipeline import TestPipeline # pylint: disable=ungrouped-imports try: import tfx_bsl import tensorflow as tf from apache_beam.examples.inference.tfx_bsl import tensorflow_image_classification from apache_beam.examples.inference.tfx_bsl.build_tensorflow_model import save_tf_model_with_signature except ImportError as e: tfx_bsl = None # pylint: disable=line-too-long _EXPECTED_OUTPUTS = { 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005001.JPEG': '681', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005002.JPEG': '333', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005003.JPEG': '711', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005004.JPEG': '286', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005005.JPEG': '445', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005006.JPEG': '288', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005007.JPEG': '880', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005008.JPEG': '534', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005009.JPEG': '888', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005010.JPEG': '996', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005011.JPEG': '327', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005012.JPEG': '573' } def process_outputs(filepath): with FileSystems().open(filepath) as f: lines = f.readlines() lines = [l.decode('utf-8').strip('\n') for l in lines] return lines @unittest.skipIf( tfx_bsl is None, 'Missing dependencies. ' 'Test depends on tfx_bsl') class TFXRunInferenceTests(unittest.TestCase): @pytest.mark.uses_tensorflow @pytest.mark.it_postcommit def test_tfx_run_inference_mobilenetv2(self): test_pipeline = TestPipeline(is_integration_test=True) # Save the Tensorflow model with TF Signature model = tf.keras.applications.MobileNetV2(weights='imagenet') preprocess_input = tf.keras.applications.mobilenet_v2.preprocess_input path_to_save_model = 'gs://apache-beam-ml/models/tensorflow_models/mobilenet_v2' save_tf_model_with_signature( path_to_save_model, model, preprocess_input, training=False) # Use the saved model to run the test file_of_image_names = ( 'gs://apache-beam-ml/testing/inputs/it_mobilenetv2_imagenet_validation_inputs.txt' ) output_file_dir = 'gs://apache-beam-ml/testing/predictions' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) extra_opts = { 'input': file_of_image_names, 'output': output_file, 'model_path': path_to_save_model, } tensorflow_image_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) predictions = process_outputs(filepath=output_file) for prediction in predictions: filename, prediction = prediction.split(',') self.assertEqual(_EXPECTED_OUTPUTS[filename], prediction) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/vertex_ai_image_classification.py0000644000175100017510000001362215206113032030212 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A sample pipeline using the RunInference API to classify images of flowers. This pipeline reads an already-processes representation of an image of sunflowers and sends it to a deployed Vertex AI model endpoint, then returns the predictions from the classifier model. The model and image are from the Hello Image Data Vertex AI tutorial (see https://cloud.google.com/vertex-ai/docs/tutorials/image-recognition-custom for more information.) """ import argparse import io import logging from collections.abc import Iterable import apache_beam as beam import tensorflow as tf from apache_beam.io import fileio from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.vertex_ai_inference import VertexAIModelHandlerJSON from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() # TODO: Update this to accept a glob of files. parser.add_argument( '--input', dest='input', type=str, required=True, help='File glob to read images from.') parser.add_argument( '--output', dest='output', type=str, required=True, help='Path to save output predictions.') parser.add_argument( '--endpoint_id', dest='endpoint', type=str, required=True, help='Vertex AI Endpoint resource ID to query (string).') parser.add_argument( '--endpoint_project', dest='project', required=True, help='GCP Project') parser.add_argument( '--endpoint_region', dest='location', type=str, required=True, help='GCP location for the Endpoint') parser.add_argument( '--endpoint_network', dest='vpc_network', type=str, required=False, help='GCP network the endpoint is peered to') parser.add_argument( '--experiment', dest='experiment', type=str, required=False, help='Vertex AI experiment label to apply to queries') parser.add_argument( '--private', dest='private', type=bool, default=False, help="True if the Vertex AI endpoint is a private endpoint") return parser.parse_known_args(argv) # Image height and width expected by the model IMG_WIDTH = 128 # Column labels for the output probabilities. COLUMNS = ['dandelion', 'daisy', 'tulips', 'sunflowers', 'roses'] def read_image(image_file_name: str) -> tuple[str, bytes]: with FileSystems().open(image_file_name, 'r') as file: data = io.BytesIO(file.read()).getvalue() return image_file_name, data def preprocess_image(data: bytes) -> list[float]: """Preprocess the image, resizing it and normalizing it before converting to a list. """ image = tf.io.decode_jpeg(data, channels=3) image = tf.image.resize_with_pad(image, IMG_WIDTH, IMG_WIDTH) image = image / 255 return image.numpy().tolist() class PostProcessor(beam.DoFn): def process(self, element: tuple[str, PredictionResult]) -> Iterable[str]: img_name, prediction_result = element prediction_vals = prediction_result.inference index = prediction_vals.index(max(prediction_vals)) yield img_name + ": " + str(COLUMNS[index]) + " (" + str( max(prediction_vals)) + ")" def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session model_handler = VertexAIModelHandlerJSON( endpoint_id=known_args.endpoint, project=known_args.project, location=known_args.location, experiment=known_args.experiment, network=known_args.vpc_network, private=known_args.private) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) read_glob = pipeline | "Get glob" >> beam.Create([known_args.input]) read_image_name = read_glob | "Get Image Paths" >> fileio.MatchAll() load_image = read_image_name | "Read Image" >> beam.Map( lambda image_name: read_image(image_name.path)) preprocess = load_image | "Preprocess Image" >> beam.MapTuple( lambda img_name, img: (img_name, preprocess_image(img))) predictions = preprocess | "RunInference" >> RunInference( KeyedModelHandler(model_handler)) process_output = predictions | "Process Predictions" >> beam.ParDo( PostProcessor()) _ = process_output | "WriteOutput" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/vertex_ai_llm_text_classification.py0000644000175100017510000001224515206113032030760 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A sample pipeline using the RunInference API to classify text using an LLM. This pipeline creates a set of prompts and sends it to a deployed Vertex AI model endpoint, then returns the predictions from the classifier model. The model can be generated by fine tuning the text-bison model or another similar model (see https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-models#supervised-fine-tuning) **NOTE** Google has deprecated PaLM LLMs and no longer supports querying them on Vertex AI endpoints. Separately, the use of the Vertex AI Predict API is not supported for Gemini models in favor of use of the google-genai API. As a result, this example no longer works as-written. To perform inference with Gemini models deployed on Google infrastructure, please see the `GeminiModelHandler` (in `apache_beam.ml.inference.gemini_inference`) and the `gemini_text_classification.py` example. For custom LLMs, you may still follow this design pattern. """ import argparse import logging import apache_beam as beam from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.vertex_ai_inference import VertexAIModelHandlerJSON from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--output', dest='output', type=str, required=True, help='Path to save output predictions.') parser.add_argument( '--endpoint_id', dest='endpoint', type=str, required=True, help='Vertex AI Endpoint resource ID to query (string).') parser.add_argument( '--endpoint_project', dest='project', required=True, help='GCP Project') parser.add_argument( '--endpoint_region', dest='location', type=str, required=True, help='GCP location for the Endpoint') parser.add_argument( '--endpoint_network', dest='vpc_network', type=str, required=False, help='GCP network the endpoint is peered to') parser.add_argument( '--experiment', dest='experiment', type=str, required=False, help='Vertex AI experiment label to apply to queries') parser.add_argument( '--private', dest='private', type=bool, default=False, help="True if the Vertex AI endpoint is a private endpoint") return parser.parse_known_args(argv) def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session model_handler = VertexAIModelHandlerJSON( endpoint_id=known_args.endpoint, project=known_args.project, location=known_args.location, experiment=known_args.experiment, network=known_args.vpc_network, private=known_args.private) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) # For this example, use the default parameters from # https://cloud.google.com/vertex-ai/docs/generative-ai/start/quickstarts/api-quickstart#parameter_definitions parameters = { "temperature": 0.2, "maxOutputTokens": 256, "topK": 40, "topP": 0.95 } prompts = [ "What is 5+2?", "Who is the president?", "Write me a business plan for a cookie shop." ] read_prompts = pipeline | "Get prompt" >> beam.Create(prompts) preprocess = read_prompts | "Format prompt" >> beam.Map( lambda data: (data, { "prompt": data })) predictions = preprocess | "RunInference" >> RunInference( KeyedModelHandler(model_handler), inference_args=parameters) _ = predictions | "PrintOutput" >> beam.Map(print) _ = predictions | "WriteOutput" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/vllm_gemma_batch.py0000644000175100017510000001117415206113032025270 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from __future__ import annotations import logging import os import tempfile import apache_beam as beam from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.vllm_inference import VLLMCompletionsModelHandler from apache_beam.ml.inference.vllm_inference import _VLLMModelServer from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions class GemmaVLLMOptions(PipelineOptions): """Custom pipeline options for the Gemma vLLM batch inference job.""" @classmethod def _add_argparse_args(cls, parser): parser.add_argument( "--input", dest="input_file", required=True, help="Input file gs://path containing prompts.", ) parser.add_argument( "--output_table", required=True, help="BigQuery table to write to in the form project:dataset.table.", ) parser.add_argument( "--model_gcs_path", required=True, help="GCS path to the directory containing model files.", ) class FormatOutput(beam.DoFn): def process(self, element): prompt = element.example comp = element.inference if hasattr(comp, 'choices'): completion = comp.choices[0].text # fallback to a single .text field elif hasattr(comp, 'text'): completion = comp.text # final fallback else: completion = str(comp) yield {'prompt': prompt, 'completion': completion} class GcsVLLMCompletionsModelHandler(VLLMCompletionsModelHandler): def __init__(self, model_name, vllm_server_kwargs=None): super().__init__(model_name, vllm_server_kwargs) self._local_model_dir = None def _download_gcs_directory(self, gcs_path: str, local_path: str): logging.info("Downloading model from %s to %s…", gcs_path, local_path) matches = FileSystems.match([os.path.join(gcs_path, "**")])[0].metadata_list for md in matches: rel = os.path.relpath(md.path, gcs_path) dst = os.path.join(local_path, rel) os.makedirs(os.path.dirname(dst), exist_ok=True) with FileSystems.open(md.path) as src, open(dst, "wb") as dstf: dstf.write(src.read()) logging.info("Download complete.") def load_model(self) -> _VLLMModelServer: uri = self._model_name if uri.startswith("gs://"): self._local_model_dir = tempfile.mkdtemp(prefix="vllm_model_") self._download_gcs_directory(uri, self._local_model_dir) logging.info("Loading vLLM from local dir %s", self._local_model_dir) return _VLLMModelServer(self._local_model_dir, self._vllm_server_kwargs) else: logging.info("Loading vLLM from HF hub: %s", uri) return super().load_model() def run(argv=None, save_main_session=True, test_pipeline=None): # Build pipeline options opts = PipelineOptions(argv) gem = opts.view_as(GemmaVLLMOptions) opts.view_as(SetupOptions).save_main_session = save_main_session logging.info("Pipeline starting with model path: %s", gem.model_gcs_path) handler = GcsVLLMCompletionsModelHandler( model_name=gem.model_gcs_path, vllm_server_kwargs={"served-model-name": gem.model_gcs_path}) with (test_pipeline or beam.Pipeline(options=opts)) as p: _ = ( p | "Read" >> beam.io.ReadFromText(gem.input_file) | "InferBatch" >> RunInference(handler, inference_batch_size=32) | "FormatForBQ" >> beam.ParDo(FormatOutput()) | "WriteToBQ" >> beam.io.WriteToBigQuery( gem.output_table, schema="prompt:STRING,completion:STRING", write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, method=beam.io.WriteToBigQuery.Method.FILE_LOADS, )) return p.result if __name__ == "__main__": logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/vllm_text_completion.py0000644000175100017510000001662715206113032026266 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A sample pipeline using the RunInference API to interface with an LLM using vLLM. Takes in a set of prompts or lists of previous messages and produces responses using a model of choice. Requires a GPU runtime with vllm, openai, and apache-beam installed to run correctly. """ import argparse import logging from collections.abc import Iterable from typing import Optional import apache_beam as beam from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.vllm_inference import OpenAIChatMessage from apache_beam.ml.inference.vllm_inference import VLLMChatModelHandler from apache_beam.ml.inference.vllm_inference import VLLMCompletionsModelHandler from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult # Defaults avoid CUDA OOM on ~16GB GPUs (e.g. NVIDIA T4) with vLLM V1: the # engine warms the sampler with many dummy sequences unless max_num_seqs is # reduced, and the default gpu_memory_utilization can leave no free VRAM for # that step. _DEFAULT_VLLM_MAX_NUM_SEQS = 32 _DEFAULT_VLLM_GPU_MEMORY_UTILIZATION = 0.72 COMPLETION_EXAMPLES = [ "Hello, my name is", "The president of the United States is", "The capital of France is", "The future of AI is", "John cena is", ] CHAT_EXAMPLES = [ [ OpenAIChatMessage( role='user', content='What is an example of a type of penguin?'), OpenAIChatMessage( role='assistant', content='Emperor penguin is a type of penguin.'), OpenAIChatMessage(role='user', content='Tell me about them') ], [ OpenAIChatMessage( role='user', content='What colors are in the rainbow?'), OpenAIChatMessage( role='assistant', content='Red, orange, yellow, green, blue, indigo, and violet.'), OpenAIChatMessage(role='user', content='Do other colors ever appear?') ], [ OpenAIChatMessage( role='user', content='Who is the president of the United States?') ], [ OpenAIChatMessage(role='user', content='What state is Fargo in?'), OpenAIChatMessage(role='assistant', content='It is in North Dakota.'), OpenAIChatMessage(role='user', content='How many people live there?'), OpenAIChatMessage( role='assistant', content='Approximately 130,000 people live in Fargo, North Dakota.' ), OpenAIChatMessage(role='user', content='What is Fargo known for?'), ], [ OpenAIChatMessage( role='user', content='How many fish are in the ocean?'), ], ] def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--model', dest='model', type=str, required=False, default='facebook/opt-125m', help='LLM to use for task') parser.add_argument( '--output', dest='output', type=str, required=True, help='Path to save output predictions.') parser.add_argument( '--chat', dest='chat', type=bool, required=False, default=False, help='Whether to use chat model handler and examples') parser.add_argument( '--chat_template', dest='chat_template', type=str, required=False, default=None, help='Chat template to use for chat example.') parser.add_argument( '--vllm_max_num_seqs', dest='vllm_max_num_seqs', type=int, default=_DEFAULT_VLLM_MAX_NUM_SEQS, help=( 'Passed to the vLLM OpenAI server as --max-num-seqs. ' 'Lower values use less GPU memory during startup and inference; ' 'required for many ~16GB GPUs (see --vllm_gpu_memory_utilization).')) parser.add_argument( '--vllm_gpu_memory_utilization', dest='vllm_gpu_memory_utilization', type=float, default=_DEFAULT_VLLM_GPU_MEMORY_UTILIZATION, help=( 'Passed to the vLLM OpenAI server as --gpu-memory-utilization ' '(fraction of total GPU memory for KV cache). Lower this if the ' 'engine fails to start with CUDA out of memory.')) return parser.parse_known_args(argv) def build_vllm_server_kwargs(known_args) -> dict[str, str]: """Returns vllm_server_kwargs for ``VLLMCompletionsModelHandler``.""" return { 'max-num-seqs': str(known_args.vllm_max_num_seqs), 'gpu-memory-utilization': str(known_args.vllm_gpu_memory_utilization), } class PostProcessor(beam.DoFn): def process(self, element: PredictionResult) -> Iterable[str]: yield str(element.example) + ": " + str(element.inference) def run( argv=None, save_main_session=True, test_pipeline=None, vllm_server_kwargs: Optional[dict[str, str]] = None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. vllm_server_kwargs: Optional override for vLLM server options. When None, options are taken from argv (``--vllm_max_num_seqs``, ``--vllm_gpu_memory_utilization``). When set, argv tuning flags for the server are ignored in favor of this dict (e.g. for programmatic use). """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session effective_vllm_kwargs = ( vllm_server_kwargs if vllm_server_kwargs is not None else build_vllm_server_kwargs(known_args)) model_handler = VLLMCompletionsModelHandler( model_name=known_args.model, vllm_server_kwargs=effective_vllm_kwargs) input_examples = COMPLETION_EXAMPLES if known_args.chat: model_handler = VLLMChatModelHandler( model_name=known_args.model, chat_template_path=known_args.chat_template, vllm_server_kwargs=dict(effective_vllm_kwargs)) input_examples = CHAT_EXAMPLES pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) examples = pipeline | "Create examples" >> beam.Create(input_examples) predictions = examples | "RunInference" >> RunInference(model_handler) process_output = predictions | "Process Predictions" >> beam.ParDo( PostProcessor()) _ = process_output | "WriteOutput" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/inference/xgboost_iris_classification.py0000644000175100017510000001313615206113032027575 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import argparse import logging from collections.abc import Callable from collections.abc import Iterable from typing import Union import numpy import pandas import scipy from sklearn.datasets import load_iris from sklearn.model_selection import train_test_split import apache_beam as beam import datatable import xgboost from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.xgboost_inference import XGBoostModelHandlerDatatable from apache_beam.ml.inference.xgboost_inference import XGBoostModelHandlerNumpy from apache_beam.ml.inference.xgboost_inference import XGBoostModelHandlerPandas from apache_beam.ml.inference.xgboost_inference import XGBoostModelHandlerSciPy from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult class PostProcessor(beam.DoFn): """Process the PredictionResult to get the predicted label. Returns a comma separated string with true label and predicted label. """ def process(self, element: tuple[int, PredictionResult]) -> Iterable[str]: label, prediction_result = element prediction = prediction_result.inference yield '{},{}'.format(label, prediction) def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input_type', dest='input_type', required=True, choices=['numpy', 'pandas', 'scipy', 'datatable'], help='Datatype of the input data.') parser.add_argument( '--output', dest='output', required=True, help='Path to save output predictions.') parser.add_argument( '--model_state', dest='model_state', required=True, help='Path to the state of the XGBoost model loaded for Inference.') parser.add_argument( '--large_model', action='store_true', dest='large_model', default=False, help='Set to true if your model is large enough to run into memory ' 'pressure if you load multiple copies.') group = parser.add_mutually_exclusive_group(required=True) group.add_argument('--split', action='store_true', dest='split') group.add_argument('--no_split', action='store_false', dest='split') return parser.parse_known_args(argv) def load_sklearn_iris_test_data( data_type: Callable, split: bool = True, seed: int = 999) -> list[Union[numpy.array, pandas.DataFrame]]: """ Loads test data from the sklearn Iris dataset in a given format, either in a single or multiple batches. Args: data_type: Datatype of the iris test dataset. split: Split the dataset in different batches or return single batch. seed: Random state for splitting the train and test set. """ dataset = load_iris() _, x_test, _, _ = train_test_split( dataset['data'], dataset['target'], test_size=.2, random_state=seed) if split: return [(index, data_type(sample.reshape(1, -1))) for index, sample in enumerate(x_test)] return [(0, data_type(x_test))] def run( argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult: """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. test_pipeline: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session data_types = { 'numpy': (numpy.array, XGBoostModelHandlerNumpy), 'pandas': (pandas.DataFrame, XGBoostModelHandlerPandas), 'scipy': (scipy.sparse.csr_matrix, XGBoostModelHandlerSciPy), 'datatable': (datatable.Frame, XGBoostModelHandlerDatatable), } input_data_type, model_handler = data_types[known_args.input_type] xgboost_model_handler = KeyedModelHandler( model_handler( model_class=xgboost.XGBClassifier, model_state=known_args.model_state, large_model=known_args.large_model)) input_data = load_sklearn_iris_test_data( data_type=input_data_type, split=known_args.split) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=pipeline_options) predictions = ( pipeline | "ReadInputData" >> beam.Create(input_data) | "RunInference" >> RunInference(xgboost_model_handler) | "PostProcessOutputs" >> beam.ParDo(PostProcessor())) _ = predictions | "WriteOutput" >> beam.io.WriteToText( known_args.output, shard_name_template='', append_trailing_newlines=True) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000003200000000000010210 xustar0026 mtime=1779996271.11638 apache_beam-2.74.0/apache_beam/examples/kafkataxi/0000755000175100017510000000000015206113157021446 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/kafkataxi/__init__.py0000644000175100017510000000142015206113032023544 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/kafkataxi/kafka_taxi.py0000644000175100017510000001343215206113032024115 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """An example that writes to and reads from Kafka. This example reads from the PubSub NYC Taxi stream described in https://github.com/googlecodelabs/cloud-dataflow-nyc-taxi-tycoon, writes to a given Kafka topic and reads back from the same Kafka topic. """ # pytype: skip-file import logging import sys import apache_beam as beam from apache_beam.io.kafka import ReadFromKafka from apache_beam.io.kafka import WriteToKafka from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions def run( bootstrap_servers, topic, with_metadata, bq_dataset, bq_table_name, project, pipeline_options): # bootstrap_servers = '123.45.67.89:123:9092' # topic = 'kafka_taxirides_realtime' # pipeline_args = ['--project', 'my-project', # '--runner', 'DataflowRunner', # '--temp_location', 'my-temp-location', # '--region', 'my-region', # '--num_workers', 'my-num-workers'] window_size = 15 # size of the Window in seconds. def log_ride(ride): if 'timestamp' in ride: logging.info( 'Found ride at latitude %r and longitude %r with %r ' 'passengers at timestamp %r', ride['latitude'], ride['longitude'], ride['passenger_count'], ride['timestamp']) else: logging.info( 'Found ride at latitude %r and longitude %r with %r ' 'passengers', ride['latitude'], ride['longitude'], ride['passenger_count']) def convert_kafka_record_to_dictionary(record): # the records have 'value' attribute when --with_metadata is given if hasattr(record, 'value'): ride_bytes = record.value elif isinstance(record, tuple): ride_bytes = record[1] else: raise RuntimeError('unknown record type: %s' % type(record)) # Converting bytes record from Kafka to a dictionary. import ast ride = ast.literal_eval(ride_bytes.decode("UTF-8")) output = { key: ride[key] for key in ['latitude', 'longitude', 'passenger_count'] } if hasattr(record, 'timestamp'): # timestamp is read from Kafka metadata output['timestamp'] = record.timestamp return output with beam.Pipeline(options=pipeline_options) as pipeline: _ = ( pipeline | beam.io.ReadFromPubSub( topic='projects/pubsub-public-data/topics/taxirides-realtime'). with_output_types(bytes) | beam.Map(lambda x: (b'', x)).with_output_types( tuple[bytes, bytes]) # Kafka write transforms expects KVs. | beam.WindowInto(beam.window.FixedWindows(window_size)) | WriteToKafka( producer_config={'bootstrap.servers': bootstrap_servers}, topic=topic)) ride_col = ( pipeline | ReadFromKafka( consumer_config={'bootstrap.servers': bootstrap_servers}, topics=[topic], with_metadata=with_metadata) | beam.Map(lambda record: convert_kafka_record_to_dictionary(record))) if bq_dataset: schema = 'latitude:STRING,longitude:STRING,passenger_count:INTEGER' if with_metadata: schema += ',timestamp:STRING' _ = ( ride_col | beam.io.WriteToBigQuery(bq_table_name, bq_dataset, project, schema)) else: _ = ride_col | beam.FlatMap(lambda ride: log_ride(ride)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) import argparse parser = argparse.ArgumentParser() parser.add_argument( '--bootstrap_servers', dest='bootstrap_servers', required=True, help='Bootstrap servers for the Kafka cluster. Should be accessible by ' 'the runner') parser.add_argument( '--topic', dest='topic', default='kafka_taxirides_realtime', help='Kafka topic to write to and read from') parser.add_argument( '--with_metadata', default=False, action='store_true', help='If set, also reads metadata from the Kafka broker.') parser.add_argument( '--bq_dataset', type=str, default='', help='BigQuery Dataset to write tables to. ' 'If set, export data to a BigQuery table instead of just logging. ' 'Must already exist.') parser.add_argument( '--bq_table_name', default='xlang_kafka_taxi', help='The BigQuery table name. Should not already exist.') known_args, pipeline_args = parser.parse_known_args() pipeline_options = PipelineOptions( pipeline_args, save_main_session=True, streaming=True) # We also require the --project option to access --bq_dataset project = pipeline_options.view_as(GoogleCloudOptions).project if project is None: parser.print_usage() print(sys.argv[0] + ': error: argument --project is required') sys.exit(1) run( known_args.bootstrap_servers, known_args.topic, known_args.with_metadata, known_args.bq_dataset, known_args.bq_table_name, project, pipeline_options) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/matrix_power.py0000644000175100017510000000675515206113032022602 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """An example that computes the matrix power y = A^m * v. A is square matrix and v is a given vector with appropriate dimension. In this computation, each element of the matrix is represented by ((i,j), a) where a is the element in the i-th row and j-th column. Each element of the vector is computed as a PCollection (i, v) where v is the element of the i-th row. For multiplication, the vector is converted into a dict side input. """ import argparse import logging import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline def extract_matrix(line): tokens = line.split(':') row = int(tokens[0]) numbers = tokens[1].strip().split() for column, number in enumerate(numbers): yield ((row, column), float(number)) def extract_vector(line): return enumerate(map(float, line.split())) def multiply_elements(element, vector): ((row, col), value) = element return (row, value * vector[col]) def run(argv=None): parser = argparse.ArgumentParser() parser.add_argument( '--input_matrix', required=True, help='Input file containing the matrix.') parser.add_argument( '--input_vector', required=True, help='Input file containing initial vector.') parser.add_argument( '--output', required=True, help='Output file to write results to.') parser.add_argument( '--exponent', required=True, type=int, help='Exponent of input square matrix.') known_args, pipeline_args = parser.parse_known_args(argv) p = TestPipeline(options=PipelineOptions(pipeline_args)) # Read the matrix from the input file and extract into the ((i,j), a) format. matrix = ( p | 'read matrix' >> beam.io.ReadFromText(known_args.input_matrix) | 'extract matrix' >> beam.FlatMap(extract_matrix)) # Read and extract the vector from its input file. vector = ( p | 'read vector' >> beam.io.ReadFromText(known_args.input_vector) | 'extract vector' >> beam.FlatMap(extract_vector)) for i in range(known_args.exponent): # Multiply the matrix by the current vector once, # and keep the resulting vector for the next iteration. vector = ( matrix # Convert vector into side-input dictionary, compute the product. | 'multiply elements %d' % i >> beam.Map( multiply_elements, beam.pvalue.AsDict(vector)) | 'sum element products %d' % i >> beam.CombinePerKey(sum)) # Format and output final vector. _ = ( vector # pylint: disable=expression-not-assigned | 'format' >> beam.Map(repr) | 'write' >> beam.io.WriteToText(known_args.output)) p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/matrix_power_test.py0000644000175100017510000000356415206113032023634 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the matrix power integration test.""" import logging import tempfile import unittest from apache_beam.examples import matrix_power class MatrixPowerTest(unittest.TestCase): MATRIX_INPUT = b''' 0:1 1 1 1:1 1 1 2:1 1 1 '''.strip() VECTOR_INPUT = b'1 2 3' EXPONENT = 3 EXPECTED_OUTPUT = ''' (0, 54.0) (1, 54.0) (2, 54.0) '''.lstrip() def create_temp_file(self, contents): with tempfile.NamedTemporaryFile(delete=False) as f: f.write(contents) return f.name def test_basics(self): matrix_path = self.create_temp_file(self.MATRIX_INPUT) vector_path = self.create_temp_file(self.VECTOR_INPUT) matrix_power.run(( '--input_matrix=%s --input_vector=%s --exponent=%d --output=%s.result' % (matrix_path, vector_path, self.EXPONENT, vector_path)).split()) # Parse result file and compare. with open(vector_path + '.result-00000-of-00001') as result_file: results = result_file.read() self.assertEqual(sorted(self.EXPECTED_OUTPUT), sorted(results)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003200000000000010210 xustar0026 mtime=1779996271.11838 apache_beam-2.74.0/apache_beam/examples/ml_transform/0000755000175100017510000000000015206113157022206 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/ml_transform/__init__.py0000644000175100017510000000142015206113032024304 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/ml_transform/ml_transform_basic.py0000644000175100017510000001602615206113032026421 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ This example demonstrates how to use MLTransform. MLTransform is a PTransform that applies multiple data transformations on the incoming data. This example computes the vocabulary on the incoming data. Then, it computes the TF-IDF of the incoming data using the vocabulary computed in the previous step. 1. ComputeAndApplyVocabulary computes the vocabulary on the incoming data and overrides the incoming data with the vocabulary indices. 2. TFIDF computes the TF-IDF of the incoming data using the vocabulary and provides vocab_index and tf-idf weights. vocab_index is suffixed with '_vocab_index' and tf-idf weights are suffixed with '_tfidf' to the original column name(which is the output of ComputeAndApplyVocabulary). MLTransform produces artifacts, for example: ComputeAndApplyVocabulary produces a text file that contains vocabulary which is saved in `artifact_location`. ComputeAndApplyVocabulary outputs vocab indices associated with the saved vocab file. This mode of MLTransform is called artifact `produce` mode. This will be useful when the data is preprocessed before ML model training. The second mode of MLTransform is artifact `consume` mode. In this mode, the transformations are applied on the incoming data using the artifacts produced by the previous run of MLTransform. This mode will be useful when the data is preprocessed before ML model inference. """ import argparse import logging import tempfile import apache_beam as beam from apache_beam.ml.transforms.base import MLTransform from apache_beam.ml.transforms.tft import TFIDF from apache_beam.ml.transforms.tft import ComputeAndApplyVocabulary from apache_beam.ml.transforms.utils import ArtifactsFetcher def parse_args(): parser = argparse.ArgumentParser() parser.add_argument('--artifact_location', type=str, default='') return parser.parse_known_args() def preprocess_data_for_ml_training(train_data, args): """ Preprocess the data for ML training. This method runs a pipeline to preprocess the data needed for ML training. It produces artifacts that can be used for ML inference later. """ with beam.Pipeline() as p: train_data_pcoll = (p | "CreateData" >> beam.Create(train_data)) # When using write_artifact_location, the ComputeAndApplyVocabulary # function generates a vocabulary file. This file, stored in # 'write_artifact_location', contains the vocabulary of the entire dataset. # This is considered as an artifact of ComputeAndApplyVocabulary transform. # The indices of the vocabulary in this file are returned as # the output of MLTransform. transformed_data_pcoll = ( train_data_pcoll | 'MLTransform' >> MLTransform( write_artifact_location=args.artifact_location, ).with_transform(ComputeAndApplyVocabulary( columns=['x'])).with_transform(TFIDF(columns=['x']))) _ = transformed_data_pcoll | beam.Map(logging.info) # output for the element dict(x=["Let's", "go", "to", "the", "park"]) # will be: # Row(x=array([21, 5, 0, 2, 3]), # x_tfidf_weight=array([0.28109303, 0.36218604, 0.36218604, 0.41972247, # 0.5008155 ], dtype=float32), x_vocab_index=array([ 0, 2, 3, 5, 21])) def preprocess_data_for_ml_inference(test_data, args): """ Preprocess the data for ML inference. This method runs a pipeline to preprocess the data needed for ML inference. It consumes the artifacts produced during the preprocessing stage for ML training. """ with beam.Pipeline() as p: test_data_pcoll = (p | beam.Create(test_data)) # Here, the previously saved vocabulary from an MLTransform run is used by # ComputeAndApplyVocabulary to access and apply the stored artifacts to the # test data. transformed_data_pcoll = ( test_data_pcoll | "MLTransformOnTestData" >> MLTransform( read_artifact_location=args.artifact_location, # ww don't need to specify transforms as they are already saved in # in the artifacts. )) _ = transformed_data_pcoll | beam.Map(logging.info) # output for dict(x=['I', 'love', 'books']) will be: # Row(x=array([1, 4, 7]), # x_tfidf_weight=array([0.4684884 , 0.6036434 , 0.69953746], dtype=float32) # , x_vocab_index=array([1, 4, 7])) def run(args): """ This example demonstrates how to use MLTransform in ML workflow. 1. Preprocess the data for ML training. 2. Do some ML model training. 3. Preprocess the data for ML inference. training and inference on ML modes are not shown in this example. This example only shows how to use MLTransform for preparing data for ML training and inference. """ train_data = [ dict(x=["Let's", "go", "to", "the", "park"]), dict(x=["I", "enjoy", "going", "to", "the", "park"]), dict(x=["I", "enjoy", "reading", "books"]), dict(x=["Beam", "can", "be", "fun"]), dict(x=["The", "weather", "is", "really", "nice", "today"]), dict(x=["I", "love", "to", "go", "to", "the", "park"]), dict(x=["I", "love", "to", "read", "books"]), dict(x=["I", "love", "to", "program"]), ] test_data = [ dict(x=['I', 'love', 'books']), dict(x=['I', 'love', 'Apache', 'Beam']) ] # Preprocess the data for ML training. # For the data going into the ML model training, we want to produce the # artifacts. preprocess_data_for_ml_training(train_data, args=args) # Do some ML model training here. # Preprocess the data for ML inference. # For the data going into the ML model inference, we want to consume the # artifacts produced during the stage where we preprocessed the data for ML # training. preprocess_data_for_ml_inference(test_data, args=args) # To fetch the artifacts produced in MLTransform, you can use # ArtifactsFetcher for fetching vocab related artifacts. For # others such as TFIDF weight, they can be accessed directly # from the output of MLTransform. artifacts_fetcher = ArtifactsFetcher(artifact_location=args.artifact_location) vocab_list = artifacts_fetcher.get_vocab_list() assert vocab_list[22] == 'Beam' if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) args, pipeline_args = parse_args() # for this example, create a temp artifact location if not provided. if args.artifact_location == '': args.artifact_location = tempfile.mkdtemp() run(args) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/ml_transform/ml_transform_it_test.py0000644000175100017510000000763515206113032027021 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import os import time import unittest import uuid import pytest try: from apache_beam.examples.ml_transform import vocab_tfidf_processing from apache_beam.io.filesystems import FileSystems from apache_beam.ml.transforms.utils import ArtifactsFetcher from apache_beam.testing.load_tests.load_test_metrics_utils import InfluxDBMetricsPublisherOptions from apache_beam.testing.load_tests.load_test_metrics_utils import MetricsReader from apache_beam.testing.test_pipeline import TestPipeline except ImportError: # pylint: disable=bare-except raise unittest.SkipTest('tensorflow_transform is not installed.') _OUTPUT_GCS_BUCKET_ROOT = 'gs://temp-storage-for-end-to-end-tests/tft/' def _publish_metrics(pipeline, metric_value, metrics_table, metric_name): influx_options = InfluxDBMetricsPublisherOptions( metrics_table, pipeline.get_option('influx_db_name'), pipeline.get_option('influx_hostname'), os.getenv('INFLUXDB_USER'), os.getenv('INFLUXDB_USER_PASSWORD'), ) metric_reader = MetricsReader( project_name=pipeline.get_option('project'), bq_table=metrics_table, bq_dataset=pipeline.get_option('metrics_dataset'), publish_to_bq=True, influxdb_options=influx_options, ) metric_reader.publish_values([( metric_name, metric_value, )]) @pytest.mark.uses_tft class LargeMovieReviewDatasetProcessTest(unittest.TestCase): def test_process_large_movie_review_dataset(self): input_data_dir = 'gs://apache-beam-ml/datasets/aclImdb' artifact_location = os.path.join(_OUTPUT_GCS_BUCKET_ROOT, uuid.uuid4().hex) output_dir = os.path.join(_OUTPUT_GCS_BUCKET_ROOT, uuid.uuid4().hex) extra_opts = { 'input_data_dir': input_data_dir, 'output_dir': output_dir, 'artifact_location': artifact_location, } extra_opts['job_name'] = 'mltransform-large-movie-review-dataset-{}'.format( uuid.uuid4().hex) test_pipeline = TestPipeline(is_integration_test=True) start_time = time.time() vocab_tfidf_processing.run( test_pipeline.get_full_options_as_args( **extra_opts, save_main_session=False), ) end_time = time.time() metrics_table = 'ml_transform_large_movie_review_dataset_process_metrics' _publish_metrics( pipeline=test_pipeline, metric_value=end_time - start_time, metrics_table=metrics_table, metric_name='runtime_sec') artifacts_fetcher = ArtifactsFetcher(artifact_location=artifact_location) vocab_filename = f'vocab_{vocab_tfidf_processing.REVIEW_COLUMN}' actual_vocab_list = artifacts_fetcher.get_vocab_list( vocab_filename=vocab_filename) expected_artifact_filepath = 'gs://apache-beam-ml/testing/expected_outputs/compute_and_apply_vocab' # pylint: disable=line-too-long with FileSystems.open(expected_artifact_filepath, 'r') as f: expected_vocab_list = f.readlines() expected_vocab_list = [ s.decode('utf-8').rstrip('\n') for s in expected_vocab_list ] self.assertListEqual(actual_vocab_list, expected_vocab_list) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/ml_transform/vocab_tfidf_processing.py0000644000175100017510000001446015206113032027257 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ This example uses Large Movie Review Dataset: http://ai.stanford.edu/~amaas/data/sentiment/ # pylint:disable=line-too-long to preprocess the input text data and generate TF-IDF scores for each word of the input text. Workflow: 1. The input text is split into words using a delimiter. 2. The words are then converted to a vocabulary index using ComputeAndApplyVocabulary. 3. The vocabulary index is then converted to TF-IDF scores using TFIDF. 4. The output of the pipeline is a Tuple of (input_text, [(vocab_index, tfidf_score)]. To run this pipeline, download the Large Movie Review Dataset and place it in a directory. Pass the directory path to --input_data_dir. The pipeline will read the data from the directory and write the transformed data to --output_dir. To save the artifacts, such as the vocabulary file generated by ComputeAndApplyVocabulary will be saved to the --artifact_location. In this pipeline, we only preprocess the train data(25000 samples). """ import argparse import logging import os import apache_beam as beam from apache_beam.ml.transforms.base import MLTransform from apache_beam.ml.transforms.tft import TFIDF from apache_beam.ml.transforms.tft import ComputeAndApplyVocabulary RAW_DATA_KEY = 'raw_data' REVIEW_COLUMN = 'review' LABEL_COLUMN = 'label' DELIMITERS = '.,!?() ' VOCAB_SIZE = 20000 # pylint: disable=invalid-name @beam.ptransform_fn def Shuffle(pcoll): """Shuffles a PCollection. Collection should not contain duplicates.""" return ( pcoll | 'PairWithHash' >> beam.Map(lambda x: (hash(x), x)) | 'GroupByHash' >> beam.GroupByKey() | 'DropHash' >> beam.FlatMap(lambda hash_and_values: hash_and_values[1])) class ReadAndShuffleData(beam.PTransform): def __init__(self, pos_file_pattern, neg_file_pattern): self.pos_file_pattern = pos_file_pattern self.neg_file_pattern = neg_file_pattern def expand(self, pcoll): negative_examples = ( pcoll | "ReadNegativeExample" >> beam.io.ReadFromText(self.neg_file_pattern) | 'PairWithZero' >> beam.Map(lambda review: (review, 0)) | 'DistinctNeg' >> beam.Distinct()) positive_examples = ( pcoll | "ReadPositiveExample" >> beam.io.ReadFromText(self.pos_file_pattern) | 'PairWithOne' >> beam.Map(lambda review: (review, 1)) | 'DistinctPos' >> beam.Distinct()) all_examples = ((negative_examples, positive_examples) | 'FlattenPColls' >> beam.Flatten()) shuffled_examples = (all_examples | 'Shuffle' >> Shuffle()) # tag with column names for MLTransform return ( shuffled_examples | beam.Map( lambda label_review: { REVIEW_COLUMN: label_review[0], LABEL_COLUMN: label_review[1], RAW_DATA_KEY: label_review[0] })) def preprocess_data( file_patterns, pipeline_args, artifact_location, output_dir, test_pipeline=None # used for testing purposes. ): positive_pattern, negative_pattern = file_patterns options = beam.options.pipeline_options.PipelineOptions(pipeline_args) pipeline = test_pipeline if not test_pipeline: pipeline = beam.Pipeline(options=options) data_pcoll = ( pipeline | 'ReadTrainData' >> ReadAndShuffleData(positive_pattern, negative_pattern)) ml_transform = MLTransform( write_artifact_location=artifact_location, ).with_transform( ComputeAndApplyVocabulary( top_k=VOCAB_SIZE, frequency_threshold=10, columns=[REVIEW_COLUMN], vocab_filename='vocab', split_string_by_delimiter=DELIMITERS)).with_transform( TFIDF( columns=[REVIEW_COLUMN], vocab_size=VOCAB_SIZE, )) data_pcoll = data_pcoll | 'MLTransform' >> ml_transform data_pcoll = ( data_pcoll | beam.ParDo(MapTFIDFScoreToVocab(artifact_location))) _ = (data_pcoll | beam.io.WriteToText(output_dir)) _ = data_pcoll | beam.Map(logging.info) result = pipeline.run() result.wait_until_finish() return result class MapTFIDFScoreToVocab(beam.DoFn): def __init__(self, artifact_location): self.artifact_location = artifact_location def process(self, element): index_column_name = REVIEW_COLUMN + '_vocab_index' weight_column_name = REVIEW_COLUMN + '_tfidf_weight' element = element.as_dict() raw_data = element[RAW_DATA_KEY] vocab_index = element[index_column_name] weights = element[weight_column_name] vocabs_with_weights = [(vocab_index[i], weights[i]) for i in range(len(vocab_index))] return [( raw_data, vocabs_with_weights, )] def parse_known_args(argv): parser = argparse.ArgumentParser() parser.add_argument( '--input_data_dir', help='path to directory containing input data.') parser.add_argument( '--artifact_location', help='path to directory to hold artifacts such as vocab files.') parser.add_argument( '--output_dir', help='path to directory to hold transformed data.') return parser.parse_known_args(argv) def run(argv=None): args, pipeline_args = parse_known_args(argv) neg_filepatterm = os.path.join(args.input_data_dir, 'train/neg/*') pos_filepattern = os.path.join(args.input_data_dir, 'train/pos/*') artifact_location = args.artifact_location _ = preprocess_data((pos_filepattern, neg_filepatterm), pipeline_args, artifact_location, output_dir=args.output_dir) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/online_clustering.py0000644000175100017510000002745415206113032023604 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import datetime from operator import itemgetter import joblib import numpy as np import pandas as pd from sklearn.cluster import MiniBatchKMeans import apache_beam as beam from apache_beam.coders import PickleCoder from apache_beam.coders import VarIntCoder from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.sklearn_inference import ModelFileType from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerNumpy from apache_beam.transforms import core from apache_beam.transforms import ptransform from apache_beam.transforms.userstate import ReadModifyWriteStateSpec class SaveModel(core.DoFn): """Saves trained clustering model to persistent storage""" def __init__(self, checkpoints_path: str): self.checkpoints_path = checkpoints_path def process(self, model): # generate ISO 8601 iso_timestamp = datetime.datetime.now( datetime.timezone.utc).strftime("%Y%m%dT%H%M%SZ") checkpoint_name = f'{self.checkpoints_path}/{iso_timestamp}.checkpoint' latest_checkpoint = f'{self.checkpoints_path}/latest.checkpoint' # rename previous checkpoint if FileSystems.exists(latest_checkpoint): FileSystems.rename([latest_checkpoint], [checkpoint_name]) file = FileSystems.create(latest_checkpoint, 'wb') if not joblib: raise ImportError( 'Could not import joblib in this execution environment. ' 'For help with managing dependencies on Python workers.' 'see https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/' # pylint: disable=line-too-long ) joblib.dump(model, file) yield checkpoint_name class AssignClusterLabelsFn(core.DoFn): """Takes a trained model and input data and labels all data instances using the trained model.""" def process(self, batch, model, model_id): cluster_labels = model.predict(batch) for e, i in zip(batch, cluster_labels): yield PredictionResult(example=e, inference=i, model_id=model_id) class SelectLatestModelState(core.CombineFn): """Selects that latest version of a model after training""" def create_accumulator(self): # create and initialise accumulator return None, 0 def add_input(self, accumulator, element): # accumulates each element from input in accumulator if element[1] > accumulator[1]: return element return accumulator def merge_accumulators(self, accumulators): # Multiple accumulators could be processed in parallel, # this function merges them return max(accumulators, key=itemgetter(1)) def extract_output(self, accumulator): # Only output the tracker return accumulator[0] class ClusteringAlgorithm(core.DoFn): """Abstract class with the interface that clustering algorithms need to follow.""" MODEL_SPEC = ReadModifyWriteStateSpec("clustering_model", PickleCoder()) ITERATION_SPEC = ReadModifyWriteStateSpec( 'training_iterations', VarIntCoder()) MODEL_ID = 'ClusteringAlgorithm' def __init__( self, n_clusters: int, checkpoints_path: str, cluster_args: dict): super().__init__() self.n_clusters = n_clusters self.checkpoints_path = checkpoints_path self.cluster_args = cluster_args self.clustering_algorithm = None def process( self, keyed_batch, model_state=core.DoFn.StateParam(MODEL_SPEC), iteration_state=core.DoFn.StateParam(ITERATION_SPEC), *args, **kwargs): raise NotImplementedError def load_model_checkpoint(self): latest_checkpoint = f'{self.checkpoints_path}/latest.checkpoint' if FileSystems.exists(latest_checkpoint): file = FileSystems.open(latest_checkpoint, 'rb') if not joblib: raise ImportError( 'Could not import joblib in this execution environment. ' 'For help with managing dependencies on Python workers.' 'see https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/' # pylint: disable=line-too-long ) return joblib.load(file) return self.clustering_algorithm( n_clusters=self.n_clusters, **self.cluster_args) class OnlineKMeans(ClusteringAlgorithm): """Online K-Means function. Used the MiniBatchKMeans from sklearn More information: https://scikit-learn.org/stable/modules/generated/sklearn.cluster.MiniBatchKMeans.html""" # pylint: disable=line-too-long MODEL_SPEC = ReadModifyWriteStateSpec("clustering_model", PickleCoder()) ITERATION_SPEC = ReadModifyWriteStateSpec( 'training_iterations', VarIntCoder()) MODEL_ID = 'OnlineKmeans' def __init__( self, n_clusters: int, checkpoints_path: str, cluster_args: dict): super().__init__(n_clusters, checkpoints_path, cluster_args) self.clustering_algorithm = MiniBatchKMeans def process( self, keyed_batch, model_state=core.DoFn.StateParam(MODEL_SPEC), iteration_state=core.DoFn.StateParam(ITERATION_SPEC), *args, **kwargs): # 1. Initialise or load states clustering = model_state.read() or self.load_model_checkpoint() iteration = iteration_state.read() or 0 iteration += 1 # 2. Remove the temporary assigned keys _, batch = keyed_batch # 3. Calculate cluster centroids clustering.partial_fit(batch) # 4. Store the training set and model model_state.write(clustering) iteration_state.write(iteration) # checkpoint = joblib.dump(clustering, f'kmeans_checkpoint_{iteration}') yield clustering, iteration class ConvertToNumpyArray(core.DoFn): """Helper function to convert incoming data to numpy arrays that are accepted by sklearn""" def process(self, element, *args, **kwargs): if isinstance(element, (tuple, list)): yield np.array(element) elif isinstance(element, np.ndarray): yield element elif isinstance(element, (pd.DataFrame, pd.Series)): yield element.to_numpy() else: raise ValueError(f"Unsupported type: {type(element)}") class ClusteringPreprocessing(ptransform.PTransform): def __init__( self, n_clusters: int, batch_size: int, is_batched: bool = False): """ Preprocessing for Clustering Transformation The clustering transform expects batches for performance reasons, therefore this batches the data and converts it to numpy arrays, which are accepted by sklearn. This transform also adds the same key to all batches, such that only 1 state is created and updated during clustering updates. Example Usage:: pcoll | ClusteringPreprocessing( n_clusters=8, batch_size=1024, is_batched=False) Args: n_clusters: number of clusters used by the algorithm batch_size: size of the data batches is_batched: boolean value that marks if the collection is already batched and thus doesn't need to be batched by this transform """ super().__init__() self.n_clusters = n_clusters self.batch_size = batch_size self.is_batched = is_batched def expand(self, pcoll): pcoll = ( pcoll | "Convert element to numpy arrays" >> beam.ParDo(ConvertToNumpyArray())) if not self.is_batched: pcoll = ( pcoll | "Create batches of elements" >> beam.BatchElements( min_batch_size=self.n_clusters, max_batch_size=self.batch_size) | "Covert to 2d numpy array" >> beam.Map(lambda record: np.array(record))) return pcoll class OnlineClustering(ptransform.PTransform): def __init__( self, clustering_algorithm, n_clusters: int, cluster_args: dict, checkpoints_path: str, batch_size: int = 1024, is_batched: bool = False): """ Clustering transformation itself, it first preprocesses the data, then it applies the clustering transformation step by step on each of the batches. Example Usage:: pcoll | OnlineClustering( clustering_algorithm=OnlineKMeansClustering batch_size=1024, n_clusters=6 cluster_args={})) Args: clustering_algorithm: Clustering algorithm (DoFn) n_clusters: Number of clusters cluster_args: Arguments for the sklearn clustering algorithm (check sklearn documentation for more information) batch_size: size of the data batches is_batched: boolean value that marks if the collection is already batched and thus doesn't need to be batched by this transform """ super().__init__() self.clustering_algorithm = clustering_algorithm self.n_clusters = n_clusters self.batch_size = batch_size self.cluster_args = cluster_args self.checkpoints_path = checkpoints_path self.is_batched = is_batched def expand(self, pcoll): # 1. Preprocess data for more efficient clustering data = ( pcoll | 'Batch data for faster processing' >> ClusteringPreprocessing( n_clusters=self.n_clusters, batch_size=self.batch_size, is_batched=self.is_batched) | "Add a key for stateful processing" >> beam.Map(lambda record: (1, record))) # 2. Calculate cluster centers model = ( data | 'Cluster' >> core.ParDo( self.clustering_algorithm( n_clusters=self.n_clusters, cluster_args=self.cluster_args, checkpoints_path=self.checkpoints_path)) | 'Select latest model state' >> core.CombineGlobally( SelectLatestModelState()).without_defaults()) # 3. Save the trained model checkpoint to persistent storage, # so it can be loaded for further training in the next window # or loaded into an sklearn modelhandler for inference _ = (model | core.ParDo(SaveModel(checkpoints_path=self.checkpoints_path))) return model class AssignClusterLabelsRunInference(ptransform.PTransform): def __init__(self, checkpoints_path): super().__init__() self.clustering_model = SklearnModelHandlerNumpy( model_uri=f'{checkpoints_path}/latest.checkpoint', model_file_type=ModelFileType.JOBLIB) def expand(self, pcoll): predictions = ( pcoll | "RunInference" >> RunInference(self.clustering_model)) return predictions class AssignClusterLabelsInMemoryModel(ptransform.PTransform): def __init__( self, model, n_clusters, batch_size, is_batched=False, model_id=None): self.model = model self.n_clusters = n_clusters self.batch_size = batch_size self.is_batched = is_batched self.model_id = model_id def expand(self, pcoll): return ( pcoll | "Preprocess data for faster prediction" >> ClusteringPreprocessing( n_clusters=self.n_clusters, batch_size=self.batch_size, is_batched=self.is_batched) | "Assign cluster labels" >> core.ParDo( AssignClusterLabelsFn(), model=self.model, model_id=self.model_id)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/per_entity_training.py0000644000175100017510000001222115206113032024120 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A pipeline to demonstrate per-entity training. This pipeline reads data from a CSV file, that contains information about 15 different attributes like salary >=50k, education level, native country, age, occupation and others. The pipeline does some filtering by selecting certain education level, discarding missing values and empty rows. The pipeline then groups the rows based on education level and trains Decision Trees for each group and finally saves them. """ import argparse import logging import pickle import pandas as pd from sklearn.compose import ColumnTransformer from sklearn.pipeline import Pipeline from sklearn.preprocessing import LabelEncoder from sklearn.preprocessing import MinMaxScaler from sklearn.preprocessing import OneHotEncoder from sklearn.tree import DecisionTreeClassifier import apache_beam as beam from apache_beam.io import fileio from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions class CreateKey(beam.DoFn): def process(self, element, *args, **kwargs): # 3rd column of the dataset is Education idx = 3 key = element.pop(idx) yield (key, element) def custom_filter(element): """Discard data point if contains ?, doesn't have all features, or doesn't have Bachelors, Masters or a Doctorate Degree""" return len(element) == 15 and '?' not in element \ and ' Bachelors' in element or ' Masters' in element \ or ' Doctorate' in element class PrepareDataforTraining(beam.DoFn): """Preprocess data in a format suitable for training.""" def process(self, element, *args, **kwargs): key, values = element #Convert to dataframe df = pd.DataFrame(values) last_ix = len(df.columns) - 1 X, y = df.drop(last_ix, axis=1), df[last_ix] # select categorical and numerical features cat_ix = X.select_dtypes(include=['object', 'bool']).columns num_ix = X.select_dtypes(include=['int64', 'float64']).columns # label encode the target variable to have the classes 0 and 1 y = LabelEncoder().fit_transform(y) yield (X, y, cat_ix, num_ix, key) class TrainModel(beam.DoFn): """Takes preprocessed data as input, transforms categorical columns using OneHotEncoder, normalizes numerical columns and then fits a decision tree classifier. """ def process(self, element, *args, **kwargs): X, y, cat_ix, num_ix, key = element steps = [('c', OneHotEncoder(handle_unknown='ignore'), cat_ix), ('n', MinMaxScaler(), num_ix)] # one hot encode categorical, normalize numerical ct = ColumnTransformer(steps) # wrap the model in a pipeline pipeline = Pipeline(steps=[('t', ct), ('m', DecisionTreeClassifier())]) pipeline.fit(X, y) yield (key, pipeline) class ModelSink(fileio.FileSink): def open(self, fh): self._fh = fh def write(self, record): _, trained_model = record pickled_model = pickle.dumps(trained_model) self._fh.write(pickled_model) def flush(self): self._fh.flush() def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', help='Path to the text file containing sentences.') parser.add_argument( '--output-dir', dest='output', required=True, help='Path of directory for saving trained models.') return parser.parse_known_args(argv) def run( argv=None, save_main_session=True, ): """ Args: argv: Command line arguments defined for this example. save_main_session: Used for internal testing. """ known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session with beam.Pipeline(options=pipeline_options) as pipeline: _ = ( pipeline | "Read Data" >> beam.io.ReadFromText(known_args.input) | "Split data to make List" >> beam.Map(lambda x: x.split(',')) | "Filter rows" >> beam.Filter(custom_filter) | "Create Key" >> beam.ParDo(CreateKey()) | "Group by education" >> beam.GroupByKey() | "Prepare Data" >> beam.ParDo(PrepareDataforTraining()) | "Train Model" >> beam.ParDo(TrainModel()) | "Save" >> fileio.WriteToFiles(path=known_args.output, sink=ModelSink())) if __name__ == "__main__": logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/rate_limiter_simple.py0000644000175100017510000000600415206113032024076 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A simple example demonstrating usage of the EnvoyRateLimiter in a Beam pipeline. """ import argparse import logging import time import apache_beam as beam from apache_beam.io.components.rate_limiter import EnvoyRateLimiter from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.utils import shared # [START RateLimiterSimplePython] class SampleApiDoFn(beam.DoFn): """A DoFn that simulates calling an external API with rate limiting.""" def __init__(self, rls_address, domain, descriptors): self.rls_address = rls_address self.domain = domain self.descriptors = descriptors self._shared = shared.Shared() self.rate_limiter = None def setup(self): # Initialize the rate limiter in setup() # We use shared.Shared() to ensure only one RateLimiter instance is created # per worker and shared across threads. def init_limiter(): logging.info("Connecting to Envoy RLS at %s", self.rls_address) return EnvoyRateLimiter( service_address=self.rls_address, domain=self.domain, descriptors=self.descriptors, namespace='example_pipeline') self.rate_limiter = self._shared.acquire(init_limiter) def process(self, element): self.rate_limiter.allow() # Process the element mock API call logging.info("Processing element: %s", element) time.sleep(0.1) yield element # [END RateLimiterSimplePython] def parse_known_args(argv): """Parses args for the workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--rls_address', default='localhost:8081', help='Address of the Envoy Rate Limit Service') return parser.parse_known_args(argv) def run(argv=None): known_args, pipeline_args = parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) with beam.Pipeline(options=pipeline_options) as p: _ = ( p | 'Create' >> beam.Create(range(100)) | 'RateLimit' >> beam.ParDo( SampleApiDoFn( rls_address=known_args.rls_address, domain="mongo_cps", descriptors=[{ "database": "users" }]))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.1203804 apache_beam-2.74.0/apache_beam/examples/sinks/0000755000175100017510000000000015206113157020632 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/sinks/__init__.py0000644000175100017510000000142015206113032022730 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/sinks/generate_event.py0000644000175100017510000001524415206113032024175 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from datetime import datetime import pytz import apache_beam as beam from apache_beam.testing.test_stream import TestStream class GenerateEvent(beam.PTransform): # pylint: disable=line-too-long """This class simulates streaming data. It leverages [TestStream](https://beam.apache.org/releases/javadoc/current/org/apache/beam/sdk/testing/TestStream.html), a method where you can control when data arrives and how watermark advances. This is especially useful in unit tests.""" # noqa @staticmethod def sample_data(): return GenerateEvent() def expand(self, input): # these are the elements that will arrive in the simulated TestStream # at multiple timestamps elem = [{'age': 10}, {'age': 20}, {'age': 30}] # The simulated TestStream adds elements at specific timestamps # using add_elements and advances the watermark after 1 or more # elements are arrive using advance_watermark_to return ( input | TestStream().add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 1, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 2, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 3, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 4, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 5, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 5, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 6, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 7, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 8, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 9, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 10, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 10, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 11, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 12, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 13, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 14, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 15, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 15, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 16, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 17, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 18, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 19, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 20, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 20, 0, tzinfo=pytz.UTC).timestamp()).advance_watermark_to( datetime( 2021, 3, 1, 0, 0, 25, 0, tzinfo=pytz.UTC). timestamp()).advance_watermark_to_infinity()) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/sinks/test_periodicimpulse.py0000644000175100017510000000456715206113032025444 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # # To run the pipelines locally: # python -m apache_beam.examples.sinks.test_periodicimpulse # This file contains examples of writing unbounded PCollection using # PeriodicImpulse to files import argparse import logging import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from apache_beam.transforms.window import FixedWindows def run(argv=None, save_main_session=True) -> PipelineResult: """Main entry point; defines and runs the wordcount pipeline.""" parser = argparse.ArgumentParser() _, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session p = beam.Pipeline(options=pipeline_options) _ = ( p | "Create elements" >> beam.transforms.periodicsequence.PeriodicImpulse( start_timestamp=1, stop_timestamp=100, fire_interval=10, apply_windowing=False) | 'ApplyWindowing' >> beam.WindowInto(FixedWindows(20)) | beam.io.WriteToText( file_path_prefix="__output__/ouput_WriteToText", file_name_suffix=".txt")) # Execute the pipeline and return the result. result = p.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/sinks/test_write_bounded.py0000644000175100017510000000674515206113032025101 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # # To run the pipelines locally: # python -m apache_beam.examples.sinks.test_write_bounded # This file contains multiple examples of writing bounded PCollection to files import argparse import json import logging import pyarrow import apache_beam as beam from apache_beam.io.fileio import WriteToFiles from apache_beam.io.textio import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from apache_beam.transforms.util import LogElements def run(argv=None, save_main_session=True) -> PipelineResult: """Main entry point; defines and runs the wordcount pipeline.""" parser = argparse.ArgumentParser() _, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session p = beam.Pipeline(options=pipeline_options) output = ( p | beam.Create([{ 'age': 10 }, { 'age': 20 }, { 'age': 30 }]) | beam.LogElements( prefix='before write ', with_window=False, level=logging.INFO)) #TextIO output2 = output | 'Write to text' >> WriteToText( file_path_prefix="__output_batch__/ouput_WriteToText", file_name_suffix=".txt", shard_name_template='-U-SSSSS-of-NNNNN') _ = output2 | 'LogElements after WriteToText' >> LogElements( prefix='after WriteToText ', with_window=False, level=logging.INFO) #FileIO output3 = ( output | 'Serialize' >> beam.Map(json.dumps) | 'Write to files' >> WriteToFiles(path="__output_batch__/output_WriteToFiles")) _ = output3 | 'LogElements after WriteToFiles' >> LogElements( prefix='after WriteToFiles ', with_window=False, level=logging.INFO) #ParquetIO output4 = output | 'Write' >> beam.io.WriteToParquet( file_path_prefix="__output_batch__/output_parquet", schema=pyarrow.schema([('age', pyarrow.int64())])) _ = output4 | 'LogElements after WriteToParquet' >> LogElements( prefix='after WriteToParquet ', with_window=False, level=logging.INFO) _ = output | 'Write parquet' >> beam.io.WriteToParquet( file_path_prefix="__output_batch__/output_WriteToParquet", schema=pyarrow.schema([('age', pyarrow.int64())]), record_batch_size=10, num_shards=0) # Execute the pipeline and return the result. result = p.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/sinks/test_write_unbounded.py0000644000175100017510000001420015206113032025425 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # # To run the pipelines locally: # python -m apache_beam.examples.sinks.test_write_unbounded # This file contains multiple examples of writing unbounded PCollection to files import argparse import json import logging import pyarrow import apache_beam as beam from apache_beam.examples.sinks.generate_event import GenerateEvent from apache_beam.io.fileio import WriteToFiles from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult from apache_beam.transforms.trigger import AccumulationMode from apache_beam.transforms.trigger import AfterWatermark from apache_beam.transforms.util import LogElements from apache_beam.transforms.window import FixedWindows from apache_beam.utils.timestamp import Duration def run(argv=None, save_main_session=True) -> PipelineResult: """Main entry point; defines and runs the wordcount pipeline.""" parser = argparse.ArgumentParser() _, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session p = beam.Pipeline(options=pipeline_options) output = p | GenerateEvent.sample_data() #TextIO output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( file_path_prefix="__output__/ouput_WriteToText", file_name_suffix=".txt", #shard_name_template='-V-SSSSS-of-NNNNN', num_shards=2, triggering_frequency=5, ) _ = output2 | 'LogElements after WriteToText' >> LogElements( prefix='after WriteToText ', with_window=True, level=logging.INFO) #FileIO _ = ( output | 'FileIO window' >> beam.WindowInto( FixedWindows(5), trigger=AfterWatermark(), accumulation_mode=AccumulationMode.DISCARDING, allowed_lateness=Duration(seconds=0)) | 'Serialize' >> beam.Map(json.dumps) | 'FileIO WriteToFiles' >> WriteToFiles(path="__output__/output_WriteToFiles")) #ParquetIO pyschema = pyarrow.schema([('age', pyarrow.int64())]) output4a = output | 'WriteToParquet' >> beam.io.WriteToParquet( file_path_prefix="__output__/output_parquet", #shard_name_template='-V-SSSSS-of-NNNNN', file_name_suffix=".parquet", num_shards=2, triggering_frequency=5, schema=pyschema) _ = output4a | 'LogElements after WriteToParquet' >> LogElements( prefix='after WriteToParquet 4a ', with_window=True, level=logging.INFO) output4aw = ( output | 'ParquetIO window' >> beam.WindowInto( FixedWindows(20), trigger=AfterWatermark(), accumulation_mode=AccumulationMode.DISCARDING, allowed_lateness=Duration(seconds=0)) | 'WriteToParquet windowed' >> beam.io.WriteToParquet( file_path_prefix="__output__/output_parquet", shard_name_template='-W-SSSSS-of-NNNNN', file_name_suffix=".parquet", num_shards=2, schema=pyschema)) _ = output4aw | 'LogElements after WriteToParquet windowed' >> LogElements( prefix='after WriteToParquet 4aw ', with_window=True, level=logging.INFO) output4b = ( output | 'To PyArrow Table' >> beam.Map(lambda x: pyarrow.Table.from_pylist([x], schema=pyschema)) | 'WriteToParquetBatched to parquet' >> beam.io.WriteToParquetBatched( file_path_prefix="__output__/output_parquet_batched", shard_name_template='-V-SSSSS-of-NNNNN', file_name_suffix=".parquet", num_shards=2, triggering_frequency=5, schema=pyschema)) _ = output4b | 'LogElements after WriteToParquetBatched' >> LogElements( prefix='after WriteToParquetBatched 4b ', with_window=True, level=logging.INFO) #AvroIO avroschema = { 'name': 'dummy', # your supposed to be file name with .avro extension 'type': 'record', # type of avro serilazation, there are more (see above # docs) but as per me this will do most of the time 'fields': [ # this defines actual keys & their types {'name': 'age', 'type': 'int'}, ], } output5 = output | 'WriteToAvro' >> beam.io.WriteToAvro( file_path_prefix="__output__/output_avro", #shard_name_template='-V-SSSSS-of-NNNNN', file_name_suffix=".avro", num_shards=2, triggering_frequency=5, schema=avroschema) _ = output5 | 'LogElements after WriteToAvro' >> LogElements( prefix='after WriteToAvro 5 ', with_window=True, level=logging.INFO) #TFrecordIO output6 = ( output | "encode" >> beam.Map(lambda s: json.dumps(s).encode('utf-8')) | 'WriteToTFRecord' >> beam.io.WriteToTFRecord( file_path_prefix="__output__/output_tfrecord", #shard_name_template='-V-SSSSS-of-NNNNN', file_name_suffix=".tfrecord", num_shards=2, triggering_frequency=5)) _ = output6 | 'LogElements after WriteToTFRecord' >> LogElements( prefix='after WriteToTFRecord 6 ', with_window=True, level=logging.INFO) # Execute the pipeline and return the result. result = p.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.125381 apache_beam-2.74.0/apache_beam/examples/snippets/0000755000175100017510000000000015206113157021350 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/__init__.py0000644000175100017510000000142015206113032023446 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/snippets.py0000644000175100017510000017547715206113032023604 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Code snippets used in webdocs. The examples here are written specifically to read well with the accompanying web docs. Do not rewrite them until you make sure the webdocs still read well and the rewritten code supports the concept being described. For example, there are snippets that could be shorter but they are written like this to make a specific point in the docs. The code snippets are all organized as self contained functions. Parts of the function body delimited by [START tag] and [END tag] will be included automatically in the web docs. The naming convention for the tags is to have as prefix the PATH_TO_HTML where they are included followed by a descriptive string. The tags can contain only letters, digits and _. """ # pytype: skip-file import argparse import base64 import json from decimal import Decimal import mock import apache_beam as beam from apache_beam.io import iobase from apache_beam.io.range_trackers import OffsetRangeTracker from apache_beam.metrics import Metrics from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.core import PTransform # Protect against environments where Google Cloud Natural Language client is # not available. try: from apache_beam.ml.gcp import naturallanguageml as nlp except ImportError: nlp = None # Quiet some pylint warnings that happen because of the somewhat special # format for the code snippets. # pylint:disable=invalid-name # pylint:disable=expression-not-assigned # pylint:disable=redefined-outer-name # pylint:disable=reimported # pylint:disable=unused-variable # pylint:disable=wrong-import-order, wrong-import-position class SnippetUtils(object): from apache_beam.pipeline import PipelineVisitor class RenameFiles(PipelineVisitor): """RenameFiles will rewire read/write paths for unit testing. RenameFiles will replace the GCS files specified in the read and write transforms to local files so the pipeline can be run as a unit test. This assumes that read and write transforms defined in snippets have already been replaced by transforms 'DummyReadForTesting' and 'DummyReadForTesting' (see snippets_test.py). This is as close as we can get to have code snippets that are executed and are also ready to presented in webdocs. """ def __init__(self, renames): self.renames = renames def visit_transform(self, transform_node): if transform_node.full_label.find('DummyReadForTesting') >= 0: transform_node.transform.fn.file_to_read = self.renames['read'] elif transform_node.full_label.find('DummyWriteForTesting') >= 0: transform_node.transform.fn.file_to_write = self.renames['write'] @mock.patch('apache_beam.Pipeline', TestPipeline) def construct_pipeline(renames): """A reverse words snippet as an example for constructing a pipeline.""" import re # This is duplicate of the import statement in # pipelines_constructing_creating tag below, but required to avoid # Unresolved reference in ReverseWords class import apache_beam as beam @beam.ptransform_fn @beam.typehints.with_input_types(str) @beam.typehints.with_output_types(str) def ReverseWords(pcoll): """A PTransform that reverses individual elements in a PCollection.""" return pcoll | beam.Map(lambda word: word[::-1]) def filter_words(unused_x): """Pass through filter to select everything.""" return True # [START pipelines_constructing_creating] import apache_beam as beam with beam.Pipeline() as pipeline: pass # build your pipeline here # [END pipelines_constructing_creating] # [START pipelines_constructing_reading] lines = pipeline | 'ReadMyFile' >> beam.io.ReadFromText( 'gs://some/inputData.txt') # [END pipelines_constructing_reading] # [START pipelines_constructing_applying] words = lines | beam.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x)) reversed_words = words | ReverseWords() # [END pipelines_constructing_applying] # [START pipelines_constructing_writing] filtered_words = reversed_words | 'FilterWords' >> beam.Filter(filter_words) filtered_words | 'WriteMyFile' >> beam.io.WriteToText( 'gs://some/outputData.txt') # [END pipelines_constructing_writing] pipeline.visit(SnippetUtils.RenameFiles(renames)) def model_pipelines(): """A wordcount snippet as a simple pipeline example.""" # [START model_pipelines] import argparse import re import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions parser = argparse.ArgumentParser() parser.add_argument( '--input-file', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='The file path for the input text to process.') parser.add_argument( '--output-path', required=True, help='The path prefix for output files.') args, beam_args = parser.parse_known_args() beam_options = PipelineOptions(beam_args) with beam.Pipeline(options=beam_options) as pipeline: ( pipeline | beam.io.ReadFromText(args.input_file) | beam.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x)) | beam.Map(lambda x: (x, 1)) | beam.combiners.Count.PerKey() | beam.io.WriteToText(args.output_path)) # [END model_pipelines] def model_pcollection(output_path): """Creating a PCollection from data in local memory.""" # [START model_pcollection] import apache_beam as beam with beam.Pipeline() as pipeline: lines = ( pipeline | beam.Create([ 'To be, or not to be: that is the question: ', "Whether 'tis nobler in the mind to suffer ", 'The slings and arrows of outrageous fortune, ', 'Or to take arms against a sea of troubles, ', ])) # [END model_pcollection] lines | beam.io.WriteToText(output_path) def pipeline_options_remote(): """Creating a Pipeline using a PipelineOptions object for remote execution.""" # [START pipeline_options_create] from apache_beam.options.pipeline_options import PipelineOptions beam_options = PipelineOptions() # [END pipeline_options_create] # [START pipeline_options_define_custom] from apache_beam.options.pipeline_options import PipelineOptions class MyOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--input') parser.add_argument('--output') # [END pipeline_options_define_custom] @mock.patch('apache_beam.Pipeline') def dataflow_options(mock_pipeline): # [START pipeline_options_dataflow_service] import argparse import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions parser = argparse.ArgumentParser() # parser.add_argument('--my-arg', help='description') args, beam_args = parser.parse_known_args() # Create and set your PipelineOptions. # For Cloud execution, specify DataflowRunner and set the Cloud Platform # project, job name, temporary files location, and region. # For more information about regions, check: # https://cloud.google.com/dataflow/docs/concepts/regional-endpoints beam_options = PipelineOptions( beam_args, runner='DataflowRunner', project='my-project-id', job_name='unique-job-name', temp_location='gs://my-bucket/temp', region='us-central1') # Note: Repeatable options like dataflow_service_options or experiments must # be specified as a list of string(s). # e.g. dataflow_service_options=['enable_prime'] # Create the Pipeline with the specified options. with beam.Pipeline(options=beam_options) as pipeline: pass # build your pipeline here. # [END pipeline_options_dataflow_service] return beam_options beam_options = dataflow_options() args = beam_options.view_as(MyOptions) with TestPipeline() as pipeline: # Use TestPipeline for testing. lines = pipeline | beam.io.ReadFromText(args.input) lines | beam.io.WriteToText(args.output) @mock.patch('apache_beam.Pipeline', TestPipeline) def pipeline_options_local(): """Creating a Pipeline using a PipelineOptions object for local execution.""" # [START pipeline_options_define_custom_with_help_and_default] from apache_beam.options.pipeline_options import PipelineOptions class MyOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='The file path for the input text to process.') parser.add_argument( '--output', required=True, help='The path prefix for output files.') # [END pipeline_options_define_custom_with_help_and_default] # [START pipeline_options_local] import argparse import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions parser = argparse.ArgumentParser() # parser.add_argument('--my-arg') args, beam_args = parser.parse_known_args() # Create and set your Pipeline Options. beam_options = PipelineOptions(beam_args) args = beam_options.view_as(MyOptions) with beam.Pipeline(options=beam_options) as pipeline: lines = ( pipeline | beam.io.ReadFromText(args.input) | beam.io.WriteToText(args.output)) # [END pipeline_options_local] @mock.patch('apache_beam.Pipeline', TestPipeline) def pipeline_options_command_line(): """Creating a Pipeline by passing a list of arguments.""" # [START pipeline_options_command_line] # Use Python argparse module to parse custom arguments import argparse import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions # For more details on how to use argparse, take a look at: # https://docs.python.org/3/library/argparse.html parser = argparse.ArgumentParser() parser.add_argument( '--input-file', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='The file path for the input text to process.') parser.add_argument( '--output-path', required=True, help='The path prefix for output files.') args, beam_args = parser.parse_known_args() # Create the Pipeline with remaining arguments. beam_options = PipelineOptions(beam_args) with beam.Pipeline(options=beam_options) as pipeline: lines = ( pipeline | 'Read files' >> beam.io.ReadFromText(args.input_file) | 'Write files' >> beam.io.WriteToText(args.output_path)) # [END pipeline_options_command_line] def pipeline_logging(lines, output): """Logging Pipeline Messages.""" import re import apache_beam as beam # [START pipeline_logging] # import Python logging module. import logging class ExtractWordsFn(beam.DoFn): def process(self, element): words = re.findall(r'[A-Za-z\']+', element) for word in words: yield word if word.lower() == 'love': # Log using the root logger at info or higher levels logging.info('Found : %s', word.lower()) # Remaining WordCount example code ... # [END pipeline_logging] with TestPipeline() as pipeline: # Use TestPipeline for testing. ( pipeline | beam.Create(lines) | beam.ParDo(ExtractWordsFn()) | beam.io.WriteToText(output)) def pipeline_monitoring(): """Using monitoring interface snippets.""" import argparse import re import apache_beam as beam class ExtractWordsFn(beam.DoFn): def process(self, element): words = re.findall(r'[A-Za-z\']+', element) for word in words: yield word class FormatCountsFn(beam.DoFn): def process(self, element): word, count = element yield '%s: %s' % (word, count) # [START pipeline_monitoring_composite] # The CountWords Composite Transform inside the WordCount pipeline. @beam.ptransform_fn def CountWords(pcoll): return ( pcoll # Convert lines of text into individual words. | 'ExtractWords' >> beam.ParDo(ExtractWordsFn()) # Count the number of times each word occurs. | beam.combiners.Count.PerElement() # Format each word and count into a printable string. | 'FormatCounts' >> beam.ParDo(FormatCountsFn())) # [END pipeline_monitoring_composite] parser = argparse.ArgumentParser() parser.add_argument( '--input-file', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='The file path for the input text to process.') parser.add_argument( '--output-path', required=True, help='The path prefix for output files.') args, _ = parser.parse_known_args() with TestPipeline() as pipeline: # Use TestPipeline for testing. # [START pipeline_monitoring_execution] ( pipeline # Read the lines of the input text. | 'ReadLines' >> beam.io.ReadFromText(args.input_file) # Count the words. | CountWords() # Write the formatted word counts to output. | 'WriteCounts' >> beam.io.WriteToText(args.output_path)) # [END pipeline_monitoring_execution] def examples_wordcount_templated(): """Templated WordCount example snippet.""" import re import apache_beam as beam from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions # [START example_wordcount_templated] class WordcountTemplatedOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): # Use add_value_provider_argument for arguments to be templatable # Use add_argument as usual for non-templatable arguments parser.add_value_provider_argument( '--input-file', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='The file path for the input text to process.') parser.add_argument( '--output-path', required=True, help='The path prefix for output files.') beam_options = PipelineOptions() args = beam_options.view_as(WordcountTemplatedOptions) with beam.Pipeline(options=beam_options) as pipeline: lines = pipeline | 'Read' >> ReadFromText(args.input_file.get()) # [END example_wordcount_templated] def format_result(word_count): (word, count) = word_count return '%s: %s' % (word, count) ( lines | 'ExtractWords' >> beam.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x)) | 'PairWithOnes' >> beam.Map(lambda x: (x, 1)) | 'Group' >> beam.GroupByKey() | 'Sum' >> beam.Map(lambda word_ones: (word_ones[0], sum(word_ones[1]))) | 'Format' >> beam.Map(format_result) | 'Write' >> WriteToText(args.output_path)) def examples_wordcount_streaming(): import apache_beam as beam from apache_beam import window from apache_beam.options.pipeline_options import PipelineOptions # Parse out arguments. parser = argparse.ArgumentParser() parser.add_argument( '--output_topic', required=True, help=( 'Output PubSub topic of the form ' '"projects//topic/".')) group = parser.add_mutually_exclusive_group(required=True) group.add_argument( '--input_topic', help=( 'Input PubSub topic of the form ' '"projects//topics/".')) group.add_argument( '--input_subscription', help=( 'Input PubSub subscription of the form ' '"projects//subscriptions/."')) args, beam_args = parser.parse_known_args() beam_options = PipelineOptions(beam_args, streaming=True) with TestPipeline(options=beam_options) as pipeline: # [START example_wordcount_streaming_read] # Read from Pub/Sub into a PCollection. if args.input_subscription: lines = pipeline | beam.io.ReadFromPubSub( subscription=args.input_subscription) else: lines = pipeline | beam.io.ReadFromPubSub(topic=args.input_topic) # [END example_wordcount_streaming_read] output = ( lines | 'DecodeUnicode' >> beam.Map(lambda encoded: encoded.decode('utf-8')) | 'ExtractWords' >> beam.FlatMap(lambda x: __import__('re').findall(r'[A-Za-z\']+', x)) | 'PairWithOnes' >> beam.Map(lambda x: (x, 1)) | beam.WindowInto(window.FixedWindows(15, 0)) | 'Group' >> beam.GroupByKey() | 'Sum' >> beam.Map(lambda word_ones: (word_ones[0], sum(word_ones[1]))) | 'Format' >> beam.MapTuple(lambda word, count: f'{word}: {count}'.encode('utf-8'))) # [START example_wordcount_streaming_write] # Write to Pub/Sub output | beam.io.WriteToPubSub(args.output_topic) # [END example_wordcount_streaming_write] def examples_ptransforms_templated(renames): # [START examples_ptransforms_templated] import apache_beam as beam from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.value_provider import StaticValueProvider class TemplatedUserOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument('--templated_int', type=int) class MySumFn(beam.DoFn): def __init__(self, templated_int): self.templated_int = templated_int def process(self, an_int): yield self.templated_int.get() + an_int beam_options = PipelineOptions() args = beam_options.view_as(TemplatedUserOptions) with beam.Pipeline(options=beam_options) as pipeline: my_sum_fn = MySumFn(args.templated_int) sum = ( pipeline | 'ReadCollection' >> beam.io.ReadFromText('gs://some/integer_collection') | 'StringToInt' >> beam.Map(lambda w: int(w)) | 'AddGivenInt' >> beam.ParDo(my_sum_fn) | 'WriteResultingCollection' >> WriteToText('some/output_path')) # [END examples_ptransforms_templated] # Templates are not supported by DirectRunner (only by DataflowRunner) # so a value must be provided at graph-construction time my_sum_fn.templated_int = StaticValueProvider(int, 10) pipeline.visit(SnippetUtils.RenameFiles(renames)) # Defining a new source. # [START model_custom_source_new_source] class CountingSource(iobase.BoundedSource): def __init__(self, count): self.records_read = Metrics.counter(self.__class__, 'recordsRead') self._count = count def estimate_size(self): return self._count def get_range_tracker(self, start_position, stop_position): if start_position is None: start_position = 0 if stop_position is None: stop_position = self._count return OffsetRangeTracker(start_position, stop_position) def read(self, range_tracker): for i in range(range_tracker.start_position(), range_tracker.stop_position()): if not range_tracker.try_claim(i): return self.records_read.inc() yield i def split(self, desired_bundle_size, start_position=None, stop_position=None): if start_position is None: start_position = 0 if stop_position is None: stop_position = self._count bundle_start = start_position while bundle_start < stop_position: bundle_stop = min(stop_position, bundle_start + desired_bundle_size) yield iobase.SourceBundle( weight=(bundle_stop - bundle_start), source=self, start_position=bundle_start, stop_position=bundle_stop) bundle_start = bundle_stop # [END model_custom_source_new_source] # We recommend users to start Source classes with an underscore to discourage # using the Source class directly when a PTransform for the source is # available. We simulate that here by simply extending the previous Source # class. class _CountingSource(CountingSource): pass # [START model_custom_source_new_ptransform] class ReadFromCountingSource(PTransform): def __init__(self, count): super().__init__() self._count = count def expand(self, pcoll): return pcoll | iobase.Read(_CountingSource(self._count)) # [END model_custom_source_new_ptransform] def model_custom_source(count): """Demonstrates creating a new custom source and using it in a pipeline. Defines a new source ``CountingSource`` that produces integers starting from 0 up to a given size. Uses the new source in an example pipeline. Additionally demonstrates how a source should be implemented using a ``PTransform``. This is the recommended way to develop sources that are to distributed to a large number of end users. This method runs two pipelines. (1) A pipeline that uses ``CountingSource`` directly using the ``df.Read`` transform. (2) A pipeline that uses a custom ``PTransform`` that wraps ``CountingSource``. Args: count: the size of the counting source to be used in the pipeline demonstrated in this method. """ # Using the source in an example pipeline. # [START model_custom_source_use_new_source] with beam.Pipeline() as pipeline: numbers = pipeline | 'ProduceNumbers' >> beam.io.Read(CountingSource(count)) # [END model_custom_source_use_new_source] lines = numbers | beam.core.Map(lambda number: 'line %d' % number) assert_that( lines, equal_to(['line ' + str(number) for number in range(0, count)])) # [START model_custom_source_use_ptransform] with beam.Pipeline() as pipeline: numbers = pipeline | 'ProduceNumbers' >> ReadFromCountingSource(count) # [END model_custom_source_use_ptransform] lines = numbers | beam.core.Map(lambda number: 'line %d' % number) assert_that( lines, equal_to(['line ' + str(number) for number in range(0, count)])) # Defining the new sink. # # Defines a new sink ``SimpleKVSink`` that demonstrates writing to a simple # key-value based storage system which has following API. # # simplekv.connect(url) - # connects to the storage system and returns an access token which can be # used to perform further operations # simplekv.open_table(access_token, table_name) - # creates a table named 'table_name'. Returns a table object. # simplekv.write_to_table(access_token, table, key, value) - # writes a key-value pair to the given table. # simplekv.rename_table(access_token, old_name, new_name) - # renames the table named 'old_name' to 'new_name'. # # [START model_custom_sink_new_sink] class SimpleKVSink(iobase.Sink): def __init__(self, simplekv, url, final_table_name): self._simplekv = simplekv self._url = url self._final_table_name = final_table_name def initialize_write(self): access_token = self._simplekv.connect(self._url) return access_token def open_writer(self, access_token, uid): table_name = 'table' + uid return SimpleKVWriter(self._simplekv, access_token, table_name) def pre_finalize(self, init_result, writer_results): pass def finalize_write(self, access_token, table_names, pre_finalize_result): for i, table_name in enumerate(table_names): self._simplekv.rename_table( access_token, table_name, self._final_table_name + str(i)) # [END model_custom_sink_new_sink] # Defining a writer for the new sink. # [START model_custom_sink_new_writer] class SimpleKVWriter(iobase.Writer): def __init__(self, simplekv, access_token, table_name): self._simplekv = simplekv self._access_token = access_token self._table_name = table_name self._table = self._simplekv.open_table(access_token, table_name) def write(self, record): key, value = record self._simplekv.write_to_table(self._access_token, self._table, key, value) def close(self): return self._table_name # [END model_custom_sink_new_writer] # [START model_custom_sink_new_ptransform] class WriteToKVSink(PTransform): def __init__(self, simplekv, url, final_table_name): self._simplekv = simplekv super().__init__() self._url = url self._final_table_name = final_table_name def expand(self, pcoll): return pcoll | iobase.Write( _SimpleKVSink(self._simplekv, self._url, self._final_table_name)) # [END model_custom_sink_new_ptransform] # We recommend users to start Sink class names with an underscore to # discourage using the Sink class directly when a PTransform for the sink is # available. We simulate that here by simply extending the previous Sink # class. class _SimpleKVSink(SimpleKVSink): pass def model_custom_sink( simplekv, KVs, final_table_name_no_ptransform, final_table_name_with_ptransform): """Demonstrates creating a new custom sink and using it in a pipeline. Uses the new sink in an example pipeline. Additionally demonstrates how a sink should be implemented using a ``PTransform``. This is the recommended way to develop sinks that are to be distributed to a large number of end users. This method runs two pipelines. (1) A pipeline that uses ``SimpleKVSink`` directly using the ``df.Write`` transform. (2) A pipeline that uses a custom ``PTransform`` that wraps ``SimpleKVSink``. Args: simplekv: an object that mocks the key-value storage. KVs: the set of key-value pairs to be written in the example pipeline. final_table_name_no_ptransform: the prefix of final set of tables to be created by the example pipeline that uses ``SimpleKVSink`` directly. final_table_name_with_ptransform: the prefix of final set of tables to be created by the example pipeline that uses a ``PTransform`` that wraps ``SimpleKVSink``. """ final_table_name = final_table_name_no_ptransform # Using the new sink in an example pipeline. # [START model_custom_sink_use_new_sink] with beam.Pipeline(options=PipelineOptions()) as pipeline: kvs = pipeline | 'CreateKVs' >> beam.Create(KVs) kvs | 'WriteToSimpleKV' >> beam.io.Write( SimpleKVSink(simplekv, 'http://url_to_simple_kv/', final_table_name)) # [END model_custom_sink_use_new_sink] final_table_name = final_table_name_with_ptransform # [START model_custom_sink_use_ptransform] with beam.Pipeline(options=PipelineOptions()) as pipeline: kvs = pipeline | 'CreateKVs' >> beam.core.Create(KVs) kvs | 'WriteToSimpleKV' >> WriteToKVSink( simplekv, 'http://url_to_simple_kv/', final_table_name) # [END model_custom_sink_use_ptransform] def model_textio(renames): """Using a Read and Write transform to read/write text files.""" def filter_words(x): import re return re.findall(r'[A-Za-z\']+', x) # [START model_textio_read] with beam.Pipeline(options=PipelineOptions()) as pipeline: # [START model_pipelineio_read] lines = pipeline | 'ReadFromText' >> beam.io.ReadFromText( 'path/to/input-*.csv') # [END model_pipelineio_read] # [END model_textio_read] # [START model_textio_write] filtered_words = lines | 'FilterWords' >> beam.FlatMap(filter_words) # [START model_pipelineio_write] filtered_words | 'WriteToText' >> beam.io.WriteToText( '/path/to/numbers', file_name_suffix='.csv') # [END model_pipelineio_write] # [END model_textio_write] pipeline.visit(SnippetUtils.RenameFiles(renames)) def model_textio_compressed(renames, expected): """Using a Read Transform to read compressed text files.""" with TestPipeline() as pipeline: # [START model_textio_write_compressed] lines = pipeline | 'ReadFromText' >> beam.io.ReadFromText( '/path/to/input-*.csv.gz', compression_type=beam.io.filesystem.CompressionTypes.GZIP) # [END model_textio_write_compressed] assert_that(lines, equal_to(expected)) pipeline.visit(SnippetUtils.RenameFiles(renames)) def model_datastoreio(): """Using a Read and Write transform to read/write to Cloud Datastore.""" import uuid import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.io.gcp.datastore.v1new.datastoreio import ReadFromDatastore from apache_beam.io.gcp.datastore.v1new.datastoreio import WriteToDatastore from apache_beam.io.gcp.datastore.v1new.types import Entity from apache_beam.io.gcp.datastore.v1new.types import Key from apache_beam.io.gcp.datastore.v1new.types import Query project = 'my_project' kind = 'my_kind' query = Query(kind, project) # [START model_datastoreio_read] pipeline = beam.Pipeline(options=PipelineOptions()) entities = pipeline | 'Read From Datastore' >> ReadFromDatastore(query) # [END model_datastoreio_read] # [START model_datastoreio_write] pipeline = beam.Pipeline(options=PipelineOptions()) musicians = pipeline | 'Musicians' >> beam.Create( ['Mozart', 'Chopin', 'Beethoven', 'Vivaldi']) def to_entity(content): key = Key([kind, str(uuid.uuid4())]) entity = Entity(key) entity.set_properties({'content': content}) return entity entities = musicians | 'To Entity' >> beam.Map(to_entity) entities | 'Write To Datastore' >> WriteToDatastore(project) # [END model_datastoreio_write] def model_bigqueryio( pipeline, write_project='', write_dataset='', write_table=''): """Using a Read and Write transform to read/write from/to BigQuery.""" # [START model_bigqueryio_table_spec] # project-id:dataset_id.table_id table_spec = 'apache-beam-testing.samples.weather_stations' # [END model_bigqueryio_table_spec] # [START model_bigqueryio_table_spec_without_project] # dataset_id.table_id table_spec = 'samples.weather_stations' # [END model_bigqueryio_table_spec_without_project] # [START model_bigqueryio_table_spec_object] from apache_beam.io.gcp.internal.clients import bigquery table_spec = bigquery.TableReference( projectId='clouddataflow-readonly', datasetId='samples', tableId='weather_stations') # [END model_bigqueryio_table_spec_object] # [START model_bigqueryio_data_types] bigquery_data = [{ 'string': 'abc', 'bytes': base64.b64encode(b'\xab\xac'), 'integer': 5, 'float': 0.5, 'numeric': Decimal('5'), 'boolean': True, 'timestamp': '2018-12-31 12:44:31.744957 UTC', 'date': '2018-12-31', 'time': '12:44:31', 'datetime': '2018-12-31T12:44:31', 'geography': 'POINT(30 10)' }] # [END model_bigqueryio_data_types] # [START model_bigqueryio_read_table] max_temperatures = ( pipeline | 'ReadTable' >> beam.io.ReadFromBigQuery(table=table_spec) # Each row is a dictionary where the keys are the BigQuery columns | beam.Map(lambda elem: elem['max_temperature'])) # [END model_bigqueryio_read_table] # [START model_bigqueryio_read_query] max_temperatures = ( pipeline | 'QueryTable' >> beam.io.ReadFromBigQuery( query='SELECT max_temperature FROM '\ '[apache-beam-testing.samples.weather_stations]') # Each row is a dictionary where the keys are the BigQuery columns | beam.Map(lambda elem: elem['max_temperature'])) # [END model_bigqueryio_read_query] # [START model_bigqueryio_read_query_std_sql] max_temperatures = ( pipeline | 'QueryTableStdSQL' >> beam.io.ReadFromBigQuery( query='SELECT max_temperature FROM '\ '`clouddataflow-readonly.samples.weather_stations`', use_standard_sql=True) # Each row is a dictionary where the keys are the BigQuery columns | beam.Map(lambda elem: elem['max_temperature'])) # [END model_bigqueryio_read_query_std_sql] # [START model_bigqueryio_read_table_with_storage_api] max_temperatures = ( pipeline | 'ReadTableWithStorageAPI' >> beam.io.ReadFromBigQuery( table=table_spec, method=beam.io.ReadFromBigQuery.Method.DIRECT_READ) | beam.Map(lambda elem: elem['max_temperature'])) # [END model_bigqueryio_read_table_with_storage_api] # [START model_bigqueryio_schema] # column_name:BIGQUERY_TYPE, ... table_schema = 'source:STRING, quote:STRING' # [END model_bigqueryio_schema] # [START model_bigqueryio_schema_object] table_schema = { 'fields': [{ 'name': 'source', 'type': 'STRING', 'mode': 'NULLABLE' }, { 'name': 'quote', 'type': 'STRING', 'mode': 'REQUIRED' }] } # [END model_bigqueryio_schema_object] if write_project and write_dataset and write_table: table_spec = '{}:{}.{}'.format(write_project, write_dataset, write_table) # [START model_bigqueryio_write_input] quotes = pipeline | beam.Create([ { 'source': 'Mahatma Gandhi', 'quote': 'My life is my message.' }, { 'source': 'Yoda', 'quote': "Do, or do not. There is no 'try'." }, ]) # [END model_bigqueryio_write_input] # [START model_bigqueryio_write] quotes | beam.io.WriteToBigQuery( table_spec, schema=table_schema, write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED) # [END model_bigqueryio_write] # [START model_bigqueryio_write_dynamic_destinations] fictional_characters_view = beam.pvalue.AsDict( pipeline | 'CreateCharacters' >> beam.Create([('Yoda', True), ('Obi Wan Kenobi', True)])) def table_fn(element, fictional_characters): if element in fictional_characters: return 'my_dataset.fictional_quotes' else: return 'my_dataset.real_quotes' quotes | 'WriteWithDynamicDestination' >> beam.io.WriteToBigQuery( table_fn, schema=table_schema, table_side_inputs=(fictional_characters_view, ), write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED) # [END model_bigqueryio_write_dynamic_destinations] # [START model_bigqueryio_time_partitioning] quotes | 'WriteWithTimePartitioning' >> beam.io.WriteToBigQuery( table_spec, schema=table_schema, write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, additional_bq_parameters={'timePartitioning': { 'type': 'HOUR' }}) # [END model_bigqueryio_time_partitioning] def model_bigqueryio_xlang( pipeline, write_project='', write_dataset='', write_table=''): """Examples for cross-language BigQuery sources and sinks.""" # to avoid a validation error(input data schema and the table schema) # use a table that does not exist import uuid never_exists_table = str(uuid.uuid4()) table_spec = 'apache-beam-testing.samples.{}'.format(never_exists_table) if write_project and write_dataset and write_table: table_spec = '{}:{}.{}'.format(write_project, write_dataset, write_table) # [START model_bigqueryio_write_schema] table_schema = { 'fields': [{ 'name': 'source', 'type': 'STRING', 'mode': 'NULLABLE' }, { 'name': 'quote', 'type': 'STRING', 'mode': 'REQUIRED' }] } # [END model_bigqueryio_write_schema] quotes = pipeline | beam.Create([ { 'source': 'Mahatma Gandhi', 'quote': 'My life is my message.' }, { 'source': 'Yoda', 'quote': "Do, or do not. There is no 'try'." }, ]) # [START model_bigqueryio_storage_write_api_with_frequency] # The Python SDK doesn't currently support setting the number of write streams quotes | "StorageWriteAPIWithFrequency" >> beam.io.WriteToBigQuery( table_spec, schema=table_schema, method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, triggering_frequency=5) # [END model_bigqueryio_storage_write_api_with_frequency] # [START model_bigqueryio_write_with_storage_write_api] quotes | "WriteTableWithStorageAPI" >> beam.io.WriteToBigQuery( table_spec, schema=table_schema, method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API) # [END model_bigqueryio_write_with_storage_write_api] def model_managed_iceberg(): """Examples for Managed Iceberg sources and sinks.""" # [START hadoop_catalog_config] hadoop_catalog_props = { 'type': 'hadoop', 'warehouse': 'file:///tmp/beam-iceberg-local-quickstart' } # [END hadoop_catalog_config] # [START managed_iceberg_config] managed_config = { 'table': 'my_db.my_table', 'catalog_properties': hadoop_catalog_props } # Note: The table will get created when inserting data (see below) # [END managed_iceberg_config] # [START managed_iceberg_insert] with beam.Pipeline() as p: ( p | beam.Create([ beam.Row(id=1, name="Mark", age=32), beam.Row(id=2, name="Omar", age=24), beam.Row(id=3, name="Rachel", age=27) ]) | beam.managed.Write("iceberg", config=managed_config)) # [END managed_iceberg_insert] # [START managed_iceberg_read] with beam.Pipeline() as p: ( p | beam.managed.Read("iceberg", config=managed_config) | beam.LogElements()) # [END managed_iceberg_read] BUCKET_NAME = "" PROJECT_ID = "" # [START biglake_public_catalog_props] biglake_catalog_props = { 'type': 'rest', 'uri': 'https://biglake.googleapis.com/iceberg/v1/restcatalog', 'warehouse': 'gs://biglake-public-nyc-taxi-iceberg', 'header.x-goog-user-project': PROJECT_ID, 'rest.auth.type': 'google', 'io-impl': 'org.apache.iceberg.gcp.gcs.GCSFileIO', 'header.X-Iceberg-Access-Delegation': 'vended-credentials' } # [END biglake_public_catalog_props] # [START biglake_public_query] from statistics import mean config = { "table": "public_data.nyc_taxicab", "catalog_properties": biglake_catalog_props, "filter": "data_file_year = 2021 AND tip_amount > 100", "keep": ["passenger_count", "total_amount", "trip_distance"] } with beam.Pipeline() as p: rows = p | beam.managed.Read("iceberg", config=config) result = ( rows | beam.Select(num_trips=lambda x: 1, *rows.element_type._fields) | beam.GroupBy('passenger_count').aggregate_field( 'num_trips', sum, 'total_trips').aggregate_field( 'total_amount', mean, 'avg_fare').aggregate_field( 'trip_distance', mean, 'avg_distance')) result | beam.Map(print) # [END biglake_public_query] # [START biglake_catalog_props] biglake_catalog_config = { 'type': 'rest', 'uri': 'https://biglake.googleapis.com/iceberg/v1/restcatalog', 'warehouse': 'gs://' + BUCKET_NAME, 'header.x-goog-user-project': PROJECT_ID, 'rest.auth.type': 'google', 'io-impl': 'org.apache.iceberg.gcp.gcs.GCSFileIO', 'header.X-Iceberg-Access-Delegation': 'vended-credentials' } # [END biglake_catalog_props] # [START model_managed_iceberg_data_types] import apache_beam as beam from apache_beam.utils.timestamp import Timestamp row = beam.Row( boolean_field=True, int_field=1, long_field=2, float_field=3.45, bytes_field=b'value', string_field="value", timestamptz_field=Timestamp(4, 5), array_field=[1, 2, 3], map_field={ "a": 1, "b": 2 }, struct_field=beam.Row(nested_field="nested_value", nested_field2=123)) import numpy as np from apache_beam.typehints.row_type import RowTypeConstraint from typing import Sequence # Override data schema by adding `with_output_types` to the transform: beam.Create(row).with_output_types( RowTypeConstraint.from_fields([ ('boolean_field', bool), ('int_field', int), ('long_field', np.int64), ('float_field', float), ('bytes_field', bytes), ('string_field', str), ('timestamptz_field', Timestamp), ('array_field', Sequence[int]), ('map_field', dict[str, int]), ( 'struct_field', RowTypeConstraint.from_fields([('nested_field', str), ('nested_field2', int)])) ])) # [END model_managed_iceberg_data_types] def model_composite_transform_example(contents, output_path): """Example of a composite transform. To declare a composite transform, define a subclass of PTransform. To override the apply method, define a method "apply" that takes a PCollection as its only parameter and returns a PCollection. """ import re import apache_beam as beam # [START composite_transform_example] # [START composite_ptransform_apply_method] # [START composite_ptransform_declare] class CountWords(beam.PTransform): # [END composite_ptransform_declare] def expand(self, pcoll): return ( pcoll | beam.FlatMap(lambda x: re.findall(r'\w+', x)) | beam.combiners.Count.PerElement() | beam.Map(lambda word_c: '%s: %s' % (word_c[0], word_c[1]))) # [END composite_ptransform_apply_method] # [END composite_transform_example] with TestPipeline() as pipeline: # Use TestPipeline for testing. ( pipeline | beam.Create(contents) | CountWords() | beam.io.WriteToText(output_path)) def model_multiple_pcollections_flatten(contents, output_path): """Merging a PCollection with Flatten.""" some_hash_fn = lambda s: ord(s[0]) partition_fn = lambda element, partitions: some_hash_fn(element) % partitions import apache_beam as beam with TestPipeline() as pipeline: # Use TestPipeline for testing. # Partition into deciles partitioned = pipeline | beam.Create(contents) | beam.Partition( partition_fn, 3) pcoll1 = partitioned[0] pcoll2 = partitioned[1] pcoll3 = partitioned[2] # Flatten them back into 1 # A collection of PCollection objects can be represented simply # as a tuple (or list) of PCollections. # (The SDK for Python has no separate type to store multiple # PCollection objects, whether containing the same or different # types.) # [START model_multiple_pcollections_flatten] merged = ( (pcoll1, pcoll2, pcoll3) # A list of tuples can be "piped" directly into a Flatten transform. | beam.Flatten()) # [END model_multiple_pcollections_flatten] merged | beam.io.WriteToText(output_path) def model_multiple_pcollections_flatten_with(contents, output_path): """Merging a PCollection with FlattenWith.""" some_hash_fn = lambda s: ord(s[0]) partition_fn = lambda element, partitions: some_hash_fn(element) % partitions import apache_beam as beam with TestPipeline() as pipeline: # Use TestPipeline for testing. # Partition into deciles partitioned = pipeline | beam.Create(contents) | beam.Partition( partition_fn, 3) pcoll1 = partitioned[0] pcoll2 = partitioned[1] pcoll3 = partitioned[2] SomeTransform = lambda: beam.Map(lambda x: x) SomeOtherTransform = lambda: beam.Map(lambda x: x) # Flatten them back into 1 # A collection of PCollection objects can be represented simply # as a tuple (or list) of PCollections. # (The SDK for Python has no separate type to store multiple # PCollection objects, whether containing the same or different # types.) # [START model_multiple_pcollections_flatten_with] merged = ( pcoll1 | SomeTransform() | beam.FlattenWith(pcoll2, pcoll3) | SomeOtherTransform()) # [END model_multiple_pcollections_flatten_with] merged | beam.io.WriteToText(output_path) def model_multiple_pcollections_flatten_with_transform(contents, output_path): """Merging output of PTransform with FlattenWith.""" some_hash_fn = lambda s: ord(s[0]) partition_fn = lambda element, partitions: some_hash_fn(element) % partitions import apache_beam as beam with TestPipeline() as pipeline: # Use TestPipeline for testing. pcoll = pipeline | beam.Create(contents) SomeTransform = lambda: beam.Map(lambda x: x) SomeOtherTransform = lambda: beam.Map(lambda x: x) # [START model_multiple_pcollections_flatten_with_transform] merged = ( pcoll | SomeTransform() | beam.FlattenWith(beam.Create(['x', 'y', 'z'])) | SomeOtherTransform()) # [END model_multiple_pcollections_flatten_with_transform] merged | beam.io.WriteToText(output_path) def model_multiple_pcollections_partition(contents, output_path): """Splitting a PCollection with Partition.""" some_hash_fn = lambda s: ord(s[0]) def get_percentile(i): """Assume i in [0,100).""" return i import apache_beam as beam with TestPipeline() as pipeline: # Use TestPipeline for testing. students = pipeline | beam.Create(contents) # [START model_multiple_pcollections_partition] def partition_fn(student, num_partitions): return int(get_percentile(student) * num_partitions / 100) by_decile = students | beam.Partition(partition_fn, 10) # [END model_multiple_pcollections_partition] # [START model_multiple_pcollections_partition_40th] fortieth_percentile = by_decile[4] # [END model_multiple_pcollections_partition_40th] ([by_decile[d] for d in range(10) if d != 4] + [fortieth_percentile] | beam.Flatten() | beam.io.WriteToText(output_path)) def model_group_by_key(contents, output_path): """Applying a GroupByKey Transform.""" import re import apache_beam as beam with TestPipeline() as pipeline: # Use TestPipeline for testing. def count_ones(word_ones): (word, ones) = word_ones return (word, sum(ones)) words_and_counts = ( pipeline | beam.Create(contents) | beam.FlatMap(lambda x: re.findall(r'\w+', x)) | 'one word' >> beam.Map(lambda w: (w, 1))) # GroupByKey accepts a PCollection of (w, 1) and # outputs a PCollection of (w, (1, 1, ...)). # (A key/value pair is just a tuple in Python.) # This is a somewhat forced example, since one could # simply use beam.combiners.Count.PerElement here. # [START model_group_by_key_transform] grouped_words = words_and_counts | beam.GroupByKey() # [END model_group_by_key_transform] ( grouped_words | 'count words' >> beam.Map(count_ones) | beam.io.WriteToText(output_path)) def model_co_group_by_key_tuple(emails, phones, output_path): """Applying a CoGroupByKey Transform to a tuple.""" import apache_beam as beam # [START model_group_by_key_cogroupbykey_tuple] # The result PCollection contains one key-value element for each key in the # input PCollections. The key of the pair will be the key from the input and # the value will be a dictionary with two entries: 'emails' - an iterable of # all values for the current key in the emails PCollection and 'phones': an # iterable of all values for the current key in the phones PCollection. results = ({'emails': emails, 'phones': phones} | beam.CoGroupByKey()) def join_info(name_info): (name, info) = name_info return '%s; %s; %s' %\ (name, sorted(info['emails']), sorted(info['phones'])) contact_lines = results | beam.Map(join_info) # [END model_group_by_key_cogroupbykey_tuple] contact_lines | beam.io.WriteToText(output_path) def model_join_using_side_inputs( name_list, email_list, phone_list, output_path): """Joining PCollections using side inputs.""" import apache_beam as beam from apache_beam.pvalue import AsIter with TestPipeline() as pipeline: # Use TestPipeline for testing. # [START model_join_using_side_inputs] # This code performs a join by receiving the set of names as an input and # passing PCollections that contain emails and phone numbers as side inputs # instead of using CoGroupByKey. names = pipeline | 'names' >> beam.Create(name_list) emails = pipeline | 'email' >> beam.Create(email_list) phones = pipeline | 'phone' >> beam.Create(phone_list) def join_info(name, emails, phone_numbers): filtered_emails = [] for name_in_list, email in emails: if name_in_list == name: filtered_emails.append(email) filtered_phone_numbers = [] for name_in_list, phone_number in phone_numbers: if name_in_list == name: filtered_phone_numbers.append(phone_number) return '; '.join([ '%s' % name, '%s' % ','.join(filtered_emails), '%s' % ','.join(filtered_phone_numbers) ]) contact_lines = names | 'CreateContacts' >> beam.core.Map( join_info, AsIter(emails), AsIter(phones)) # [END model_join_using_side_inputs] contact_lines | beam.io.WriteToText(output_path) # [START model_library_transforms_keys] class Keys(beam.PTransform): def expand(self, pcoll): return pcoll | 'Keys' >> beam.Map(lambda k_v: k_v[0]) # [END model_library_transforms_keys] # pylint: enable=invalid-name # [START model_library_transforms_count] class Count(beam.PTransform): def expand(self, pcoll): return ( pcoll | 'PairWithOne' >> beam.Map(lambda v: (v, 1)) | beam.CombinePerKey(sum)) # [END model_library_transforms_count] def file_process_pattern_access_metadata(): import apache_beam as beam from apache_beam.io import fileio # [START FileProcessPatternAccessMetadataSnip1] with beam.Pipeline() as pipeline: readable_files = ( pipeline | fileio.MatchFiles('hdfs://path/to/*.txt') | fileio.ReadMatches() | beam.Reshuffle()) files_and_contents = ( readable_files | beam.Map(lambda x: (x.metadata.path, x.read_utf8()))) # [END FileProcessPatternAccessMetadataSnip1] def accessing_valueprovider_info_after_run(): # [START AccessingValueProviderInfoAfterRunSnip1] import logging import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.value_provider import RuntimeValueProvider class MyOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument('--string_value', type=str) class LogValueProvidersFn(beam.DoFn): def __init__(self, string_vp): self.string_vp = string_vp # Define the DoFn that logs the ValueProvider value. # The DoFn is called when creating the pipeline branch. # This example logs the ValueProvider value, but # you could store it by pushing it to an external database. def process(self, an_int): logging.info('The string_value is %s' % self.string_vp.get()) # Another option (where you don't need to pass the value at all) is: logging.info( 'The string value is %s' % RuntimeValueProvider.get_value('string_value', str, '')) beam_options = PipelineOptions() args = beam_options.view_as(MyOptions) # Create pipeline. with beam.Pipeline(options=beam_options) as pipeline: # Add a branch for logging the ValueProvider value. _ = ( pipeline | beam.Create([None]) | 'LogValueProvs' >> beam.ParDo(LogValueProvidersFn(args.string_value))) # The main pipeline. result_pc = ( pipeline | "main_pc" >> beam.Create([1, 2, 3]) | beam.CombineGlobally(sum)) # [END AccessingValueProviderInfoAfterRunSnip1] def side_input_slow_update( src_file_pattern, first_timestamp, last_timestamp, interval, sample_main_input_elements, main_input_windowing_interval): # [START SideInputSlowUpdateSnip1] from apache_beam.transforms.periodicsequence import PeriodicImpulse from apache_beam.transforms.window import TimestampedValue from apache_beam.transforms import window # from apache_beam.utils.timestamp import MAX_TIMESTAMP # last_timestamp = MAX_TIMESTAMP to go on indefninitely # Any user-defined function. # cross join is used as an example. def cross_join(left, rights): for x in rights: yield (left, x) # Create pipeline. pipeline = beam.Pipeline() side_input = ( pipeline | 'PeriodicImpulse' >> PeriodicImpulse( first_timestamp, last_timestamp, interval, True) | 'MapToFileName' >> beam.Map(lambda x: src_file_pattern + str(x)) | 'ReadFromFile' >> beam.io.ReadAllFromText()) main_input = ( pipeline | 'MpImpulse' >> beam.Create(sample_main_input_elements) | 'MapMpToTimestamped' >> beam.Map(lambda src: TimestampedValue(src, src)) | 'WindowMpInto' >> beam.WindowInto( window.FixedWindows(main_input_windowing_interval))) result = ( main_input | 'ApplyCrossJoin' >> beam.FlatMap( cross_join, rights=beam.pvalue.AsIter(side_input))) # [END SideInputSlowUpdateSnip1] return pipeline, result def bigqueryio_deadletter(): # [START BigQueryIODeadLetter] # Create pipeline. schema = ({'fields': [{'name': 'a', 'type': 'STRING', 'mode': 'REQUIRED'}]}) pipeline = beam.Pipeline() errors = ( pipeline | 'Data' >> beam.Create([1, 2]) | 'CreateBrokenData' >> beam.Map(lambda src: {'a': src} if src == 2 else {'a': None}) | 'WriteToBigQuery' >> beam.io.WriteToBigQuery( "> beam.FlatMap(lambda err: print("Error Found {}".format(err)))) # [END BigQueryIODeadLetter] return result def extract_sentiments(response): # [START nlp_extract_sentiments] return { 'sentences': [{ sentence.text.content: sentence.sentiment.score } for sentence in response.sentences], 'document_sentiment': response.document_sentiment.score, } # [END nlp_extract_sentiments] def extract_entities(response): # [START nlp_extract_entities] return [{ 'name': entity.name, 'type': nlp.enums.Entity.Type(entity.type).name, } for entity in response.entities] # [END nlp_extract_entities] def analyze_dependency_tree(response): # [START analyze_dependency_tree] from collections import defaultdict adjacency_lists = [] index = 0 for sentence in response.sentences: adjacency_list = defaultdict(list) sentence_begin = sentence.text.begin_offset sentence_end = sentence_begin + len(sentence.text.content) - 1 while index < len(response.tokens) and \ response.tokens[index].text.begin_offset <= sentence_end: token = response.tokens[index] head_token_index = token.dependency_edge.head_token_index head_token_text = response.tokens[head_token_index].text.content adjacency_list[head_token_text].append(token.text.content) index += 1 adjacency_lists.append(adjacency_list) # [END analyze_dependency_tree] return adjacency_lists def nlp_analyze_text(): # [START nlp_analyze_text] features = nlp.types.AnnotateTextRequest.Features( extract_entities=True, extract_document_sentiment=True, extract_entity_sentiment=True, extract_syntax=True, ) with beam.Pipeline() as pipeline: responses = ( pipeline | beam.Create([ 'My experience so far has been fantastic! ' 'I\'d really recommend this product.' ]) | beam.Map(lambda x: nlp.Document(x, type='PLAIN_TEXT')) | nlp.AnnotateText(features)) _ = ( responses | beam.Map(extract_sentiments) | 'Parse sentiments to JSON' >> beam.Map(json.dumps) | 'Write sentiments' >> beam.io.WriteToText('sentiments.txt')) _ = ( responses | beam.Map(extract_entities) | 'Parse entities to JSON' >> beam.Map(json.dumps) | 'Write entities' >> beam.io.WriteToText('entities.txt')) _ = ( responses | beam.Map(analyze_dependency_tree) | 'Parse adjacency list to JSON' >> beam.Map(json.dumps) | 'Write adjacency list' >> beam.io.WriteToText('adjancency_list.txt')) # [END nlp_analyze_text] def sdf_basic_example(): import os from apache_beam.io.restriction_trackers import OffsetRange read_next_record = None # [START SDF_BasicExample] class FileToWordsRestrictionProvider(beam.transforms.core.RestrictionProvider ): def initial_restriction(self, file_name): return OffsetRange(0, os.stat(file_name).st_size) def create_tracker(self, restriction): return beam.io.restriction_trackers.OffsetRestrictionTracker() class FileToWordsFn(beam.DoFn): def process( self, file_name, # Alternatively, we can let FileToWordsFn itself inherit from # RestrictionProvider, implement the required methods and let # tracker=beam.DoFn.RestrictionParam() which will use self as # the provider. tracker=beam.DoFn.RestrictionParam(FileToWordsRestrictionProvider())): with open(file_name) as file_handle: file_handle.seek(tracker.current_restriction.start()) while tracker.try_claim(file_handle.tell()): yield read_next_record(file_handle) # Providing the coder is only necessary if it can not be inferred at # runtime. def restriction_coder(self): return ... # [END SDF_BasicExample] def sdf_basic_example_with_splitting(): from apache_beam.io.restriction_trackers import OffsetRange # [START SDF_BasicExampleWithSplitting] class FileToWordsRestrictionProvider(beam.transforms.core.RestrictionProvider ): def split(self, file_name, restriction): # Compute and output 64 MiB size ranges to process in parallel split_size = 64 * (1 << 20) i = restriction.start while i < restriction.end - split_size: yield OffsetRange(i, i + split_size) i += split_size yield OffsetRange(i, restriction.end) # [END SDF_BasicExampleWithSplitting] def sdf_sdk_initiated_checkpointing(): timestamp = None external_service = None class MyRestrictionProvider(object): pass # [START SDF_UserInitiatedCheckpoint] class MySplittableDoFn(beam.DoFn): def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam( MyRestrictionProvider())): current_position = restriction_tracker.current_restriction.start() while True: # Pull records from an external service. try: records = external_service.fetch(current_position) if records.empty(): # Set a shorter delay in case we are being throttled. restriction_tracker.defer_remainder(timestamp.Duration(second=10)) return for record in records: if restriction_tracker.try_claim(record.position): current_position = record.position yield record else: return except TimeoutError: # Set a longer delay in case we are being throttled. restriction_tracker.defer_remainder(timestamp.Duration(seconds=60)) return # [END SDF_UserInitiatedCheckpoint] def sdf_get_size(): # [START SDF_GetSize] # The RestrictionProvider is responsible for calculating the size of given # restriction. class MyRestrictionProvider(beam.transforms.core.RestrictionProvider): def restriction_size(self, file_name, restriction): weight = 2 if "expensiveRecords" in file_name else 1 return restriction.size() * weight # [END SDF_GetSize] def sdf_bad_try_claim_loop(): class FileToWordsRestrictionProvider(object): pass read_next_record = None # [START SDF_BadTryClaimLoop] class BadTryClaimLoop(beam.DoFn): def process( self, file_name, tracker=beam.DoFn.RestrictionParam(FileToWordsRestrictionProvider())): with open(file_name) as file_handle: file_handle.seek(tracker.current_restriction.start()) # The restriction tracker can be modified by another thread in parallel # so storing state locally is ill advised. end = tracker.current_restriction.end() while file_handle.tell() < end: # Only after successfully claiming should we produce any output and/or # perform side effects. tracker.try_claim(file_handle.tell()) yield read_next_record(file_handle) # [END SDF_BadTryClaimLoop] def sdf_custom_watermark_estimator(): from apache_beam.io.iobase import WatermarkEstimator from apache_beam.transforms.core import WatermarkEstimatorProvider current_watermark = None class MyRestrictionProvider(object): pass # [START SDF_CustomWatermarkEstimator] # (Optional) Define a custom watermark state type to save information between # bundle processing rounds. class MyCustomerWatermarkEstimatorState(object): def __init__(self, element, restriction): # Store data necessary for future watermark computations pass # Define a WatermarkEstimator class MyCustomWatermarkEstimator(WatermarkEstimator): def __init__(self, estimator_state): self.state = estimator_state def observe_timestamp(self, timestamp): # Will be invoked on each output from the SDF pass def current_watermark(self): # Return a monotonically increasing value return current_watermark def get_estimator_state(self): # Return state to resume future watermark estimation after a # checkpoint/split return self.state # Then, a WatermarkEstimatorProvider needs to be created for this # WatermarkEstimator class MyWatermarkEstimatorProvider(WatermarkEstimatorProvider): def initial_estimator_state(self, element, restriction): return MyCustomerWatermarkEstimatorState(element, restriction) def create_watermark_estimator(self, estimator_state): return MyCustomWatermarkEstimator(estimator_state) # Finally, define the SDF using your estimator. class MySplittableDoFn(beam.DoFn): def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam(MyRestrictionProvider()), watermark_estimator=beam.DoFn.WatermarkEstimatorParam( MyWatermarkEstimatorProvider())): # The current watermark can be inspected. watermark_estimator.current_watermark() # [END SDF_CustomWatermarkEstimator] def sdf_truncate(): # [START SDF_Truncate] class MyRestrictionProvider(beam.transforms.core.RestrictionProvider): def truncate(self, file_name, restriction): if "optional" in file_name: # Skip optional files return None return restriction # [END SDF_Truncate] def bundle_finalize(): my_callback_func = None # [START BundleFinalize] class MySplittableDoFn(beam.DoFn): def process(self, element, bundle_finalizer=beam.DoFn.BundleFinalizerParam): # ... produce output ... # Register callback function for this bundle that performs the side # effect. bundle_finalizer.register(my_callback_func) # [END BundleFinalize] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/snippets_examples_wordcount_debugging.py0000644000175100017510000001112615206113032031575 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # Quiet some pylint warnings that happen because of the somewhat special # format for the code snippets. # pylint:disable=invalid-name # pylint:disable=expression-not-assigned # pylint:disable=redefined-outer-name # pylint:disable=reimported # pylint:disable=unused-variable # pylint:disable=wrong-import-order, wrong-import-position # beam-playground: # name: WordCountDebuggingSnippet # description: An example that counts words in Shakespeare's works. # includes regex filter("Flourish|stomach"). # multifile: false # pipeline_options: --output output.txt # context_line: 57 # categories: # - Flatten # - Debugging # - Options # - Combiners # - Filtering # - Quickstart # complexity: ADVANCED # tags: # - count # - debug # - string import apache_beam as beam from apache_beam.examples.snippets.snippets import SnippetUtils from apache_beam.metrics import Metrics from apache_beam.testing.test_pipeline import TestPipeline def examples_wordcount_debugging(renames): """DebuggingWordCount example snippets.""" import re # [START example_wordcount_debugging_logging] # [START example_wordcount_debugging_aggregators] import logging class FilterTextFn(beam.DoFn): """A DoFn that filters for a specific key based on a regular expression.""" def __init__(self, pattern): self.pattern = pattern # A custom metric can track values in your pipeline as it runs. Create # custom metrics matched_word and unmatched_words. self.matched_words = Metrics.counter(self.__class__, 'matched_words') self.umatched_words = Metrics.counter(self.__class__, 'umatched_words') def process(self, element): word, _ = element if re.match(self.pattern, word): # Log at INFO level each element we match. When executing this pipeline # using the Dataflow service, these log lines will appear in the Cloud # Logging UI. logging.info('Matched %s', word) # Add 1 to the custom metric counter matched_words self.matched_words.inc() yield element else: # Log at the "DEBUG" level each element that is not matched. Different # log levels can be used to control the verbosity of logging providing # an effective mechanism to filter less important information. Note # currently only "INFO" and higher level logs are emitted to the Cloud # Logger. This log message will not be visible in the Cloud Logger. logging.debug('Did not match %s', word) # Add 1 to the custom metric counter umatched_words self.umatched_words.inc() # [END example_wordcount_debugging_logging] # [END example_wordcount_debugging_aggregators] with TestPipeline() as pipeline: # Use TestPipeline for testing. # assert False filtered_words = ( pipeline | beam.io.ReadFromText('gs://dataflow-samples/shakespeare/kinglear.txt') | 'ExtractWords' >> beam.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x)) | beam.combiners.Count.PerElement() | 'FilterText' >> beam.ParDo(FilterTextFn('Flourish|stomach'))) # [START example_wordcount_debugging_assert] beam.testing.util.assert_that( filtered_words, beam.testing.util.equal_to([('Flourish', 3), ('stomach', 1)])) # [END example_wordcount_debugging_assert] def format_result(word_count): (word, count) = word_count return '%s: %s' % (word, count) output = ( filtered_words | 'format' >> beam.Map(format_result) | 'Write' >> beam.io.WriteToText('output.txt')) pipeline.visit(SnippetUtils.RenameFiles(renames)) if __name__ == '__main__': import glob examples_wordcount_debugging(None) for file_name in glob.glob('output.txt*'): with open(file_name) as f: print(f.read()) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/snippets_examples_wordcount_minimal.py0000644000175100017510000000712015206113032031267 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # Quiet some pylint warnings that happen because of the somewhat special # format for the code snippets. # pylint:disable=invalid-name # pylint:disable=expression-not-assigned # pylint:disable=redefined-outer-name # pylint:disable=reimported # pylint:disable=unused-variable # pylint:disable=wrong-import-order, wrong-import-position # pylint:disable=line-too-long # beam-playground: # name: WordCountMinimalSnippet # description: An example that counts words in Shakespeare's works. # multifile: false # pipeline_options: --input-file gs://dataflow-samples/shakespeare/kinglear.txt --output-path output.txt # context_line: 85 # categories: # - IO # - Core Transforms # - Flatten # - Options # - Combiners # - Quickstart # complexity: MEDIUM # tags: # - count # - strings # - hellobeam def examples_wordcount_minimal(): """MinimalWordCount example snippets.""" import re import apache_beam as beam # [START examples_wordcount_minimal_options] from apache_beam.options.pipeline_options import PipelineOptions input_file = 'gs://dataflow-samples/shakespeare/kinglear.txt' output_path = 'gs://my-bucket/counts.txt' beam_options = PipelineOptions( runner='DataflowRunner', project='my-project-id', job_name='unique-job-name', temp_location='gs://my-bucket/temp', ) # [END examples_wordcount_minimal_options] # Run it locally for testing. import argparse parser = argparse.ArgumentParser() parser.add_argument('--input-file') parser.add_argument('--output-path') args, beam_args = parser.parse_known_args() input_file = args.input_file output_path = args.output_path beam_options = PipelineOptions(beam_args) # [START examples_wordcount_minimal_create] pipeline = beam.Pipeline(options=beam_options) # [END examples_wordcount_minimal_create] ( # [START examples_wordcount_minimal_read] pipeline | beam.io.ReadFromText(input_file) # [END examples_wordcount_minimal_read] # [START examples_wordcount_minimal_pardo] | 'ExtractWords' >> beam.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x)) # [END examples_wordcount_minimal_pardo] # [START examples_wordcount_minimal_count] | beam.combiners.Count.PerElement() # [END examples_wordcount_minimal_count] # [START examples_wordcount_minimal_map] | beam.MapTuple(lambda word, count: '%s: %s' % (word, count)) # [END examples_wordcount_minimal_map] # [START examples_wordcount_minimal_write] | beam.io.WriteToText(output_path) # [END examples_wordcount_minimal_write] ) # [START examples_wordcount_minimal_run] result = pipeline.run() # [END examples_wordcount_minimal_run] result.wait_until_finish() if __name__ == '__main__': examples_wordcount_minimal() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/snippets_examples_wordcount_wordcount.py0000644000175100017510000000623415206113032031672 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # Quiet some pylint warnings that happen because of the somewhat special # format for the code snippets. # pylint:disable=invalid-name # pylint:disable=expression-not-assigned # pylint:disable=redefined-outer-name # pylint:disable=reimported # pylint:disable=unused-variable # pylint:disable=wrong-import-order, wrong-import-position # beam-playground: # name: WordCountSnippet # description: An example that counts words in Shakespeare's works. # multifile: false # pipeline_options: --output-path output.txt # context_line: 50 # categories: # - Combiners # - Options # - Quickstart # complexity: MEDIUM # tags: # - options # - count # - combine # - strings def examples_wordcount_wordcount(): """WordCount example snippets.""" import re import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions # [START examples_wordcount_wordcount_options] import argparse parser = argparse.ArgumentParser() parser.add_argument( '--input-file', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='The file path for the input text to process.') parser.add_argument( '--output-path', required=True, help='The path prefix for output files.') args, beam_args = parser.parse_known_args() beam_options = PipelineOptions(beam_args) with beam.Pipeline(options=beam_options) as pipeline: lines = pipeline | beam.io.ReadFromText(args.input_file) # [END examples_wordcount_wordcount_options] # [START examples_wordcount_wordcount_composite] @beam.ptransform_fn def CountWords(pcoll): return ( pcoll # Convert lines of text into individual words. | 'ExtractWords' >> beam.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x)) # Count the number of times each word occurs. | beam.combiners.Count.PerElement()) counts = lines | CountWords() # [END examples_wordcount_wordcount_composite] # [START examples_wordcount_wordcount_dofn] class FormatAsTextFn(beam.DoFn): def process(self, element): word, count = element yield '%s: %s' % (word, count) formatted = counts | beam.ParDo(FormatAsTextFn()) # [END examples_wordcount_wordcount_dofn] formatted | beam.io.WriteToText(args.output_path) if __name__ == '__main__': examples_wordcount_wordcount() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/snippets_test.py0000644000175100017510000015217015206113032024624 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for all code snippets used in public docs.""" # pytype: skip-file import gc import glob import gzip import logging import math import os import sys import tempfile import time import unittest import uuid import mock import parameterized import pytest import apache_beam as beam from apache_beam import WindowInto from apache_beam import coders from apache_beam import pvalue from apache_beam import typehints from apache_beam.coders.coders import ToBytesCoder from apache_beam.examples.snippets import snippets from apache_beam.examples.snippets import snippets_examples_wordcount_debugging from apache_beam.examples.snippets import snippets_examples_wordcount_minimal from apache_beam.examples.snippets import snippets_examples_wordcount_wordcount from apache_beam.metrics import Metrics from apache_beam.metrics.metric import MetricsFilter from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import combiners from apache_beam.transforms.trigger import AccumulationMode from apache_beam.transforms.trigger import AfterAny from apache_beam.transforms.trigger import AfterCount from apache_beam.transforms.trigger import AfterProcessingTime from apache_beam.transforms.trigger import AfterWatermark from apache_beam.transforms.trigger import Repeatedly from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import TimestampedValue from apache_beam.utils.windowed_value import WindowedValue # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apitools.base.py import base_api except ImportError: base_api = None # pylint: enable=wrong-import-order, wrong-import-position # Protect against environments where datastore library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from google.cloud.datastore import client as datastore_client except ImportError: datastore_client = None # pylint: enable=wrong-import-order, wrong-import-position # Protect against environments where the PubSub library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from google.cloud import pubsub except ImportError: pubsub = None # pylint: enable=wrong-import-order, wrong-import-position class ParDoTest(unittest.TestCase): """Tests for model/par-do.""" def test_pardo(self): # Note: "words" and "ComputeWordLengthFn" are referenced by name in # the text of the doc. words = ['aa', 'bbb', 'c'] # [START model_pardo_pardo] class ComputeWordLengthFn(beam.DoFn): def process(self, element): return [len(element)] # [END model_pardo_pardo] # [START model_pardo_apply] # Apply a ParDo to the PCollection "words" to compute lengths for each word. word_lengths = words | beam.ParDo(ComputeWordLengthFn()) # [END model_pardo_apply] self.assertEqual({2, 3, 1}, set(word_lengths)) def test_pardo_yield(self): words = ['aa', 'bbb', 'c'] # [START model_pardo_yield] class ComputeWordLengthFn(beam.DoFn): def process(self, element): yield len(element) # [END model_pardo_yield] word_lengths = words | beam.ParDo(ComputeWordLengthFn()) self.assertEqual({2, 3, 1}, set(word_lengths)) def test_pardo_using_map(self): words = ['aa', 'bbb', 'c'] # [START model_pardo_using_map] word_lengths = words | beam.Map(len) # [END model_pardo_using_map] self.assertEqual({2, 3, 1}, set(word_lengths)) def test_pardo_using_flatmap(self): words = ['aa', 'bbb', 'c'] # [START model_pardo_using_flatmap] word_lengths = words | beam.FlatMap(lambda word: [len(word)]) # [END model_pardo_using_flatmap] self.assertEqual({2, 3, 1}, set(word_lengths)) def test_pardo_using_flatmap_yield(self): words = ['aA', 'bbb', 'C'] # [START model_pardo_using_flatmap_yield] def capitals(word): for letter in word: if 'A' <= letter <= 'Z': yield letter all_capitals = words | beam.FlatMap(capitals) # [END model_pardo_using_flatmap_yield] self.assertEqual({'A', 'C'}, set(all_capitals)) def test_pardo_with_label(self): words = ['aa', 'bbc', 'defg'] # [START model_pardo_with_label] result = words | 'CountUniqueLetters' >> beam.Map( lambda word: len(set(word))) # [END model_pardo_with_label] self.assertEqual({1, 2, 4}, set(result)) def test_pardo_side_input(self): # pylint: disable=line-too-long with TestPipeline() as p: words = p | 'start' >> beam.Create(['a', 'bb', 'ccc', 'dddd']) # [START model_pardo_side_input] # Callable takes additional arguments. def filter_using_length(word, lower_bound, upper_bound=float('inf')): if lower_bound <= len(word) <= upper_bound: yield word # Construct a deferred side input. avg_word_len = ( words | beam.Map(len) | beam.CombineGlobally(beam.combiners.MeanCombineFn())) # Call with explicit side inputs. small_words = words | 'small' >> beam.FlatMap(filter_using_length, 0, 3) # A single deferred side input. larger_than_average = ( words | 'large' >> beam.FlatMap( filter_using_length, lower_bound=pvalue.AsSingleton(avg_word_len)) ) # Mix and match. small_but_nontrivial = words | beam.FlatMap( filter_using_length, lower_bound=2, upper_bound=pvalue.AsSingleton(avg_word_len)) # [END model_pardo_side_input] assert_that(small_words, equal_to(['a', 'bb', 'ccc'])) assert_that( larger_than_average, equal_to(['ccc', 'dddd']), label='larger_than_average') assert_that( small_but_nontrivial, equal_to(['bb']), label='small_but_not_trivial') def test_pardo_side_input_dofn(self): words = ['a', 'bb', 'ccc', 'dddd'] # [START model_pardo_side_input_dofn] class FilterUsingLength(beam.DoFn): def process(self, element, lower_bound, upper_bound=float('inf')): if lower_bound <= len(element) <= upper_bound: yield element small_words = words | beam.ParDo(FilterUsingLength(), 0, 3) # [END model_pardo_side_input_dofn] self.assertEqual({'a', 'bb', 'ccc'}, set(small_words)) def test_pardo_with_tagged_outputs(self): # [START model_pardo_emitting_values_on_tagged_outputs] class ProcessWords(beam.DoFn): def process(self, element, cutoff_length, marker): if len(element) <= cutoff_length: # Emit this short word to the main output. yield element else: # Emit this word's long length to the 'above_cutoff_lengths' output. yield pvalue.TaggedOutput('above_cutoff_lengths', len(element)) if element.startswith(marker): # Emit this word to a different output with the 'marked strings' tag. yield pvalue.TaggedOutput('marked strings', element) # [END model_pardo_emitting_values_on_tagged_outputs] words = ['a', 'an', 'the', 'music', 'xyz'] # [START model_pardo_with_tagged_outputs] results = ( words | beam.ParDo(ProcessWords(), cutoff_length=2, marker='x').with_outputs( 'above_cutoff_lengths', 'marked strings', main='below_cutoff_strings')) below = results.below_cutoff_strings above = results.above_cutoff_lengths marked = results['marked strings'] # indexing works as well # [END model_pardo_with_tagged_outputs] self.assertEqual({'a', 'an'}, set(below)) self.assertEqual({3, 5}, set(above)) self.assertEqual({'xyz'}, set(marked)) # [START model_pardo_with_tagged_outputs_iter] below, above, marked = (words | beam.ParDo( ProcessWords(), cutoff_length=2, marker='x') .with_outputs('above_cutoff_lengths', 'marked strings', main='below_cutoff_strings')) # [END model_pardo_with_tagged_outputs_iter] self.assertEqual({'a', 'an'}, set(below)) self.assertEqual({3, 5}, set(above)) self.assertEqual({'xyz'}, set(marked)) def test_pardo_with_undeclared_outputs(self): # Note: the use of undeclared outputs is currently not supported in eager # execution mode. with TestPipeline() as p: numbers = p | beam.Create([1, 2, 3, 4, 5, 10, 20]) # [START model_pardo_with_undeclared_outputs] def even_odd(x): yield pvalue.TaggedOutput('odd' if x % 2 else 'even', x) if x % 10 == 0: yield x results = numbers | beam.FlatMap(even_odd).with_outputs() evens = results.even odds = results.odd tens = results[None] # the undeclared main output # [END model_pardo_with_undeclared_outputs] assert_that(evens, equal_to([2, 4, 10, 20]), label='assert_even') assert_that(odds, equal_to([1, 3, 5]), label='assert_odds') assert_that(tens, equal_to([10, 20]), label='assert_tens') class TypeHintsTest(unittest.TestCase): def test_bad_types(self): # [START type_hints_missing_define_numbers] p = TestPipeline() numbers = p | beam.Create(['1', '2', '3']) # [END type_hints_missing_define_numbers] # Consider the following code. # pylint: disable=expression-not-assigned # pylint: disable=unused-variable # [START type_hints_missing_apply] evens = numbers | beam.Filter(lambda x: x % 2 == 0) # [END type_hints_missing_apply] # Now suppose numbers was defined as [snippet above]. # When running this pipeline, you'd get a runtime error, # possibly on a remote machine, possibly very late. with self.assertRaisesRegex(Exception, "not all arguments converted"): p.run().wait_until_finish() # To catch this early, we can assert what types we expect. with self.assertRaises(typehints.TypeCheckError): # [START type_hints_takes] evens = numbers | beam.Filter(lambda x: x % 2 == 0).with_input_types(int) # [END type_hints_takes] # Type hints can be declared on DoFns and callables as well, rather # than where they're used, to be more self contained. with self.assertRaises(typehints.TypeCheckError): # [START type_hints_do_fn] @beam.typehints.with_input_types(int) class FilterEvensDoFn(beam.DoFn): def process(self, element): if element % 2 == 0: yield element evens = numbers | beam.ParDo(FilterEvensDoFn()) # [END type_hints_do_fn] words = p | 'words' >> beam.Create(['a', 'bb', 'c']) # One can assert outputs and apply them to transforms as well. # Helps document the contract and checks it at pipeline construction time. # [START type_hints_transform] from typing import TypeVar T = TypeVar('T') @beam.typehints.with_input_types(T) @beam.typehints.with_output_types(tuple[int, T]) class MyTransform(beam.PTransform): def expand(self, pcoll): return pcoll | beam.Map(lambda x: (len(x), x)) words_with_lens = words | MyTransform() # [END type_hints_transform] # Given an input of str, the inferred output type would be Tuple[int, str]. self.assertEqual(typehints.Tuple[int, str], words_with_lens.element_type) # pylint: disable=expression-not-assigned with self.assertRaises(typehints.TypeCheckError): words_with_lens | beam.Map(lambda x: x).with_input_types(tuple[int, int]) def test_bad_types_annotations(self): p = TestPipeline(options=PipelineOptions(pipeline_type_check=True)) numbers = p | beam.Create(['1', '2', '3']) # Consider the following code. # pylint: disable=expression-not-assigned # pylint: disable=unused-variable class FilterEvensDoFn(beam.DoFn): def process(self, element): if element % 2 == 0: yield element evens = numbers | 'Untyped Filter' >> beam.ParDo(FilterEvensDoFn()) # Now suppose numbers was defined as [snippet above]. # When running this pipeline, you'd get a runtime error, # possibly on a remote machine, possibly very late. with self.assertRaisesRegex(Exception, "not all arguments converted"): p.run().wait_until_finish() # To catch this early, we can annotate process() with the expected types. # Beam will then use these as type hints and perform type checking before # the pipeline starts. with self.assertRaises(typehints.TypeCheckError): # [START type_hints_do_fn_annotations] from typing import Iterable class TypedFilterEvensDoFn(beam.DoFn): def process(self, element: int) -> Iterable[int]: if element % 2 == 0: yield element evens = numbers | 'filter_evens' >> beam.ParDo(TypedFilterEvensDoFn()) # [END type_hints_do_fn_annotations] # Another example, using a list output type. Notice that the output # annotation has an additional Optional for the else clause. with self.assertRaises(typehints.TypeCheckError): # [START type_hints_do_fn_annotations_optional] from typing import Optional class FilterEvensDoubleDoFn(beam.DoFn): def process(self, element: int) -> Optional[list[int]]: if element % 2 == 0: return [element, element] return None evens = numbers | 'double_evens' >> beam.ParDo(FilterEvensDoubleDoFn()) # [END type_hints_do_fn_annotations_optional] # Example using an annotated function. with self.assertRaises(typehints.TypeCheckError): # [START type_hints_map_annotations] def my_fn(element: int) -> str: return 'id_' + str(element) ids = numbers | 'to_id' >> beam.Map(my_fn) # [END type_hints_map_annotations] # Example using an annotated PTransform. with self.assertRaises(typehints.TypeCheckError): # [START type_hints_ptransforms] from apache_beam.pvalue import PCollection class IntToStr(beam.PTransform): def expand(self, pcoll: PCollection[int]) -> PCollection[str]: return pcoll | beam.Map(lambda elem: str(elem)) ids = numbers | 'convert to str' >> IntToStr() # [END type_hints_ptransforms] def test_runtime_checks_off(self): # We do not run the following pipeline, as it has incorrect type # information, and may fail with obscure errors, depending on the runner # implementation. # pylint: disable=expression-not-assigned # [START type_hints_runtime_off] p = TestPipeline() p | beam.Create(['a']) | beam.Map(lambda x: 3).with_output_types(str) # [END type_hints_runtime_off] def test_runtime_checks_on(self): # pylint: disable=expression-not-assigned with self.assertRaisesRegex(Exception, "According to type-hint"): # [START type_hints_runtime_on] p = TestPipeline(options=PipelineOptions(runtime_type_check=True)) p | beam.Create(['a']) | beam.Map(lambda x: 3).with_output_types(str) result = p.run() # [END type_hints_runtime_on] result.wait_until_finish() def test_deterministic_key(self): with TestPipeline() as p: lines = ( p | beam.Create([ 'banana,fruit,3', 'kiwi,fruit,2', 'kiwi,fruit,2', 'zucchini,veg,3' ])) # For pickling. global Player # pylint: disable=global-variable-not-assigned # [START type_hints_deterministic_key] class Player(object): def __init__(self, team, name): self.team = team self.name = name class PlayerCoder(beam.coders.Coder): def encode(self, player): return ('%s:%s' % (player.team, player.name)).encode('utf-8') def decode(self, s): return Player(*s.decode('utf-8').split(':')) def is_deterministic(self): return True beam.coders.registry.register_coder(Player, PlayerCoder) def parse_player_and_score(csv): name, team, score = csv.split(',') return Player(team, name), int(score) totals = ( lines | beam.Map(parse_player_and_score) | beam.CombinePerKey(sum).with_input_types(tuple[Player, int])) # [END type_hints_deterministic_key] assert_that( totals | beam.Map(lambda k_v: (k_v[0].name, k_v[1])), equal_to([('banana', 3), ('kiwi', 4), ('zucchini', 3)])) class SnippetsTest(unittest.TestCase): # Replacing text read/write transforms with dummy transforms for testing. class DummyReadTransform(beam.PTransform): """A transform that will replace iobase.ReadFromText. To be used for testing. """ def __init__(self, file_to_read=None, compression_type=None): self.file_to_read = file_to_read self.compression_type = compression_type class ReadDoFn(beam.DoFn): def __init__(self, file_to_read, compression_type): self.file_to_read = file_to_read self.compression_type = compression_type self.coder = coders.StrUtf8Coder() def process(self, element): pass def finish_bundle(self): from apache_beam.transforms import window assert self.file_to_read for file_name in glob.glob(self.file_to_read): if self.compression_type is None: with open(file_name, 'rb') as file: for record in file: value = self.coder.decode(record.rstrip(b'\n')) yield WindowedValue(value, -1, [window.GlobalWindow()]) else: with gzip.open(file_name, 'rb') as file: for record in file: value = self.coder.decode(record.rstrip(b'\n')) yield WindowedValue(value, -1, [window.GlobalWindow()]) def expand(self, pcoll): return pcoll | beam.Create([None]) | 'DummyReadForTesting' >> beam.ParDo( SnippetsTest.DummyReadTransform.ReadDoFn( self.file_to_read, self.compression_type)) class DummyWriteTransform(beam.PTransform): """A transform that will replace iobase.WriteToText. To be used for testing. """ def __init__(self, file_to_write=None, file_name_suffix=''): self.file_to_write = file_to_write class WriteDoFn(beam.DoFn): def __init__(self, file_to_write): self.file_to_write = file_to_write self.file_obj = None self.coder = ToBytesCoder() def start_bundle(self): assert self.file_to_write # Appending a UUID to create a unique file object per invocation. self.file_obj = open(self.file_to_write + str(uuid.uuid4()), 'wb') def process(self, element): assert self.file_obj self.file_obj.write(self.coder.encode(element) + b'\n') def finish_bundle(self): assert self.file_obj self.file_obj.close() def expand(self, pcoll): return pcoll | 'DummyWriteForTesting' >> beam.ParDo( SnippetsTest.DummyWriteTransform.WriteDoFn(self.file_to_write)) def setUp(self): self.old_read_from_text = beam.io.ReadFromText self.old_write_to_text = beam.io.WriteToText # Monkey patching to allow testing pipelines defined in snippets.py using # real data. beam.io.ReadFromText = SnippetsTest.DummyReadTransform beam.io.WriteToText = SnippetsTest.DummyWriteTransform self.temp_files = [] def tearDown(self): beam.io.ReadFromText = self.old_read_from_text beam.io.WriteToText = self.old_write_to_text # Cleanup all the temporary files created in the test. map(os.remove, self.temp_files) # Ensure that PipelineOptions subclasses have been cleaned up between tests gc.collect() def create_temp_file(self, contents=''): with tempfile.NamedTemporaryFile(delete=False) as f: f.write(contents.encode('utf-8')) self.temp_files.append(f.name) return f.name def get_output(self, path, sorted_output=True, suffix=''): all_lines = [] for file_name in glob.glob(path + '*'): with open(file_name) as f: lines = f.readlines() all_lines.extend([s.rstrip('\n') for s in lines]) if sorted_output: return sorted(s.rstrip('\n') for s in all_lines) return all_lines def test_model_pipelines(self): temp_path = self.create_temp_file('aa bb cc\n bb cc\n cc') result_path = temp_path + '.result' test_argv = [ "unused_argv[0]", f"--input-file={temp_path}*", f"--output-path={result_path}", ] with mock.patch.object(sys, 'argv', test_argv): snippets.model_pipelines() self.assertEqual( self.get_output(result_path), [str(s) for s in [('aa', 1), ('bb', 2), ('cc', 3)]]) def test_model_pcollection(self): temp_path = self.create_temp_file() snippets.model_pcollection(temp_path) self.assertEqual( self.get_output(temp_path), [ 'Or to take arms against a sea of troubles, ', 'The slings and arrows of outrageous fortune, ', 'To be, or not to be: that is the question: ', 'Whether \'tis nobler in the mind to suffer ', ]) def test_construct_pipeline(self): temp_path = self.create_temp_file('abc def ghi\n jkl mno pqr\n stu vwx yz') result_path = self.create_temp_file() snippets.construct_pipeline({'read': temp_path, 'write': result_path}) self.assertEqual( self.get_output(result_path), ['cba', 'fed', 'ihg', 'lkj', 'onm', 'rqp', 'uts', 'xwv', 'zy']) def test_model_custom_source(self): snippets.model_custom_source(100) def test_model_custom_sink(self): tempdir_name = tempfile.mkdtemp() class SimpleKV(object): def __init__(self, tmp_dir): self._dummy_token = 'dummy_token' self._tmp_dir = tmp_dir def connect(self, url): return self._dummy_token def open_table(self, access_token, table_name): assert access_token == self._dummy_token file_name = self._tmp_dir + os.sep + table_name assert not os.path.exists(file_name) open(file_name, 'wb').close() return table_name def write_to_table(self, access_token, table_name, key, value): assert access_token == self._dummy_token file_name = self._tmp_dir + os.sep + table_name assert os.path.exists(file_name) with open(file_name, 'ab') as f: content = (key + ':' + value + os.linesep).encode('utf-8') f.write(content) def rename_table(self, access_token, old_name, new_name): assert access_token == self._dummy_token old_file_name = self._tmp_dir + os.sep + old_name new_file_name = self._tmp_dir + os.sep + new_name assert os.path.isfile(old_file_name) assert not os.path.exists(new_file_name) os.rename(old_file_name, new_file_name) snippets.model_custom_sink( SimpleKV(tempdir_name), [('key' + str(i), 'value' + str(i)) for i in range(100)], 'final_table_no_ptransform', 'final_table_with_ptransform') expected_output = [ 'key' + str(i) + ':' + 'value' + str(i) for i in range(100) ] glob_pattern = tempdir_name + os.sep + 'final_table_no_ptransform*' output_files = glob.glob(glob_pattern) assert output_files received_output = [] for file_name in output_files: with open(file_name) as f: for line in f: received_output.append(line.rstrip(os.linesep)) self.assertCountEqual(expected_output, received_output) glob_pattern = tempdir_name + os.sep + 'final_table_with_ptransform*' output_files = glob.glob(glob_pattern) assert output_files received_output = [] for file_name in output_files: with open(file_name) as f: for line in f: received_output.append(line.rstrip(os.linesep)) self.assertCountEqual(expected_output, received_output) def test_model_textio(self): temp_path = self.create_temp_file('aa bb cc\n bb cc\n cc') result_path = temp_path + '.result' snippets.model_textio({'read': temp_path, 'write': result_path}) self.assertEqual(['aa', 'bb', 'bb', 'cc', 'cc', 'cc'], self.get_output(result_path, suffix='.csv')) def test_model_textio_compressed(self): temp_path = self.create_temp_file('aa\nbb\ncc') gzip_file_name = temp_path + '.gz' with open(temp_path, 'rb') as src, gzip.open(gzip_file_name, 'wb') as dst: dst.writelines(src) # Add the temporary gzip file to be cleaned up as well. self.temp_files.append(gzip_file_name) snippets.model_textio_compressed({'read': gzip_file_name}, ['aa', 'bb', 'cc']) @unittest.skipIf( datastore_client is None, 'GCP dependencies are not installed') def test_model_datastoreio(self): # We cannot test DatastoreIO functionality in unit tests, therefore we limit # ourselves to making sure the pipeline containing Datastore read and write # transforms can be built. # TODO(vikasrk): Expore using Datastore Emulator. snippets.model_datastoreio() @unittest.skipIf(base_api is None, 'GCP dependencies are not installed') def test_model_bigqueryio(self): # We cannot test BigQueryIO functionality in unit tests, therefore we limit # ourselves to making sure the pipeline containing BigQuery sources and # sinks can be built. # # To run locally, set `run_locally` to `True`. You will also have to set # `project`, `dataset` and `table` to the BigQuery table the test will write # to. run_locally = False if run_locally: project = 'my-project' dataset = 'samples' # this must already exist table = 'model_bigqueryio' # this will be created if needed options = PipelineOptions().view_as(GoogleCloudOptions) options.project = project with beam.Pipeline(options=options) as p: snippets.model_bigqueryio(p, project, dataset, table) else: p = TestPipeline() p.options.view_as(GoogleCloudOptions).temp_location = 'gs://mylocation' snippets.model_bigqueryio(p) @pytest.mark.uses_gcp_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_PORT'), "EXPANSION_PORT environment var is not provided.") def test_model_bigqueryio_xlang(self): p = TestPipeline() p.options.view_as(GoogleCloudOptions).temp_location = 'gs://mylocation' snippets.model_bigqueryio_xlang(p) def _run_test_pipeline_for_options(self, fn): temp_path = self.create_temp_file('aa\nbb\ncc') result_path = temp_path + '.result' test_argv = [ "unused_argv[0]", f"--input={temp_path}*", f"--output={result_path}", ] with mock.patch.object(sys, 'argv', test_argv): fn() self.assertEqual(['aa', 'bb', 'cc'], self.get_output(result_path)) def test_pipeline_options_local(self): self._run_test_pipeline_for_options(snippets.pipeline_options_local) def test_pipeline_options_remote(self): self._run_test_pipeline_for_options(snippets.pipeline_options_remote) def test_pipeline_options_command_line(self): self._run_test_pipeline_for_options(snippets.pipeline_options_command_line) def test_pipeline_logging(self): result_path = self.create_temp_file() lines = [ 'we found love right where we are', 'we found love right from the start', 'we found love in a hopeless place' ] snippets.pipeline_logging(lines, result_path) self.assertEqual( sorted(' '.join(lines).split(' ')), self.get_output(result_path)) @parameterized.parameterized.expand([ [snippets_examples_wordcount_minimal.examples_wordcount_minimal], [snippets_examples_wordcount_wordcount.examples_wordcount_wordcount], [snippets.pipeline_monitoring], [snippets.examples_wordcount_templated], ]) def test_examples_wordcount(self, pipeline): temp_path = self.create_temp_file('abc def ghi\n abc jkl') result_path = self.create_temp_file() test_argv = [ "unused_argv[0]", f"--input-file={temp_path}*", f"--output-path={result_path}", ] with mock.patch.object(sys, 'argv', test_argv): pipeline() self.assertEqual( self.get_output(result_path), ['abc: 2', 'def: 1', 'ghi: 1', 'jkl: 1']) def test_examples_ptransforms_templated(self): pipelines = [snippets.examples_ptransforms_templated] for pipeline in pipelines: temp_path = self.create_temp_file('1\n 2\n 3') result_path = self.create_temp_file() pipeline({'read': temp_path, 'write': result_path}) self.assertEqual(self.get_output(result_path), ['11', '12', '13']) def test_examples_wordcount_debugging(self): temp_path = self.create_temp_file( 'Flourish Flourish Flourish stomach abc def') result_path = self.create_temp_file() snippets_examples_wordcount_debugging.examples_wordcount_debugging({ 'read': temp_path, 'write': result_path }) self.assertEqual( self.get_output(result_path), ['Flourish: 3', 'stomach: 1']) @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') @mock.patch('apache_beam.io.ReadFromPubSub') @mock.patch('apache_beam.io.WriteToPubSub') def test_examples_wordcount_streaming(self, *unused_mocks): def FakeReadFromPubSub(topic=None, subscription=None, values=None): expected_topic = topic expected_subscription = subscription def _inner(topic=None, subscription=None): assert topic == expected_topic assert subscription == expected_subscription return TestStream().add_elements(values) return _inner class AssertTransform(beam.PTransform): def __init__(self, matcher): self.matcher = matcher def expand(self, pcoll): assert_that(pcoll, self.matcher) def FakeWriteToPubSub(topic=None, values=None): expected_topic = topic def _inner(topic=None, subscription=None): assert topic == expected_topic return AssertTransform(equal_to(values)) return _inner # Test basic execution. input_topic = 'projects/fake-beam-test-project/topic/intopic' input_values = [ TimestampedValue(b'a a b', 1), TimestampedValue('🤷 ¯\\_(ツ)_/¯ b b '.encode('utf-8'), 12), TimestampedValue(b'a b c c c', 20) ] output_topic = 'projects/fake-beam-test-project/topic/outtopic' output_values = [b'a: 1', b'a: 2', b'b: 1', b'b: 3', b'c: 3'] beam.io.ReadFromPubSub = ( FakeReadFromPubSub(topic=input_topic, values=input_values)) beam.io.WriteToPubSub = ( FakeWriteToPubSub(topic=output_topic, values=output_values)) test_argv = [ 'unused_argv[0]', '--input_topic', 'projects/fake-beam-test-project/topic/intopic', '--output_topic', 'projects/fake-beam-test-project/topic/outtopic' ] with mock.patch.object(sys, 'argv', test_argv): snippets.examples_wordcount_streaming() # Test with custom subscription. input_sub = 'projects/fake-beam-test-project/subscriptions/insub' beam.io.ReadFromPubSub = FakeReadFromPubSub( subscription=input_sub, values=input_values) test_argv = [ 'unused_argv[0]', '--input_subscription', 'projects/fake-beam-test-project/subscriptions/insub', '--output_topic', 'projects/fake-beam-test-project/topic/outtopic' ] with mock.patch.object(sys, 'argv', test_argv): snippets.examples_wordcount_streaming() def test_model_composite_transform_example(self): contents = ['aa bb cc', 'bb cc', 'cc'] result_path = self.create_temp_file() snippets.model_composite_transform_example(contents, result_path) self.assertEqual(['aa: 1', 'bb: 2', 'cc: 3'], self.get_output(result_path)) def test_model_multiple_pcollections_flatten(self): contents = ['a', 'b', 'c', 'd', 'e', 'f'] result_path = self.create_temp_file() snippets.model_multiple_pcollections_flatten(contents, result_path) self.assertEqual(contents, self.get_output(result_path)) def test_model_multiple_pcollections_flatten_with(self): contents = ['a', 'b', 'c', 'd', 'e', 'f'] result_path = self.create_temp_file() snippets.model_multiple_pcollections_flatten_with(contents, result_path) self.assertEqual(contents, self.get_output(result_path)) def test_model_multiple_pcollections_flatten_with_transform(self): contents = ['a', 'b', 'c', 'd', 'e', 'f'] result_path = self.create_temp_file() snippets.model_multiple_pcollections_flatten_with_transform( contents, result_path) self.assertEqual(contents + ['x', 'y', 'z'], self.get_output(result_path)) def test_model_multiple_pcollections_partition(self): contents = [17, 42, 64, 32, 0, 99, 53, 89] result_path = self.create_temp_file() snippets.model_multiple_pcollections_partition(contents, result_path) self.assertEqual(['0', '17', '32', '42', '53', '64', '89', '99'], self.get_output(result_path)) def test_model_group_by_key(self): contents = ['a bb ccc bb bb a'] result_path = self.create_temp_file() snippets.model_group_by_key(contents, result_path) expected = [('a', 2), ('bb', 3), ('ccc', 1)] self.assertEqual([str(s) for s in expected], self.get_output(result_path)) def test_model_co_group_by_key_tuple(self): with TestPipeline() as p: # [START model_group_by_key_cogroupbykey_tuple_inputs] emails_list = [ ('amy', 'amy@example.com'), ('carl', 'carl@example.com'), ('julia', 'julia@example.com'), ('carl', 'carl@email.com'), ] phones_list = [ ('amy', '111-222-3333'), ('james', '222-333-4444'), ('amy', '333-444-5555'), ('carl', '444-555-6666'), ] emails = p | 'CreateEmails' >> beam.Create(emails_list) phones = p | 'CreatePhones' >> beam.Create(phones_list) # [END model_group_by_key_cogroupbykey_tuple_inputs] result_path = self.create_temp_file() snippets.model_co_group_by_key_tuple(emails, phones, result_path) # [START model_group_by_key_cogroupbykey_tuple_outputs] results = [ ( 'amy', { 'emails': ['amy@example.com'], 'phones': ['111-222-3333', '333-444-5555'] }), ( 'carl', { 'emails': ['carl@email.com', 'carl@example.com'], 'phones': ['444-555-6666'] }), ('james', { 'emails': [], 'phones': ['222-333-4444'] }), ('julia', { 'emails': ['julia@example.com'], 'phones': [] }), ] # [END model_group_by_key_cogroupbykey_tuple_outputs] # [START model_group_by_key_cogroupbykey_tuple_formatted_outputs] formatted_results = [ "amy; ['amy@example.com']; ['111-222-3333', '333-444-5555']", "carl; ['carl@email.com', 'carl@example.com']; ['444-555-6666']", "james; []; ['222-333-4444']", "julia; ['julia@example.com']; []", ] # [END model_group_by_key_cogroupbykey_tuple_formatted_outputs] expected_results = [ '%s; %s; %s' % (name, info['emails'], info['phones']) for name, info in results ] self.assertEqual(expected_results, formatted_results) self.assertEqual(formatted_results, self.get_output(result_path)) def test_model_use_and_query_metrics(self): """DebuggingWordCount example snippets.""" import re p = TestPipeline() # Use TestPipeline for testing. words = p | beam.Create( ['albert', 'sam', 'mark', 'sarah', 'swati', 'daniel', 'andrea']) # pylint: disable=unused-variable # [START metrics_usage_example] class FilterTextFn(beam.DoFn): """A DoFn that filters for a specific key based on a regex.""" def __init__(self, pattern): self.pattern = pattern # A custom metric can track values in your pipeline as it runs. Create # custom metrics to count unmatched words, and know the distribution of # word lengths in the input PCollection. self.word_len_dist = Metrics.distribution( self.__class__, 'word_len_dist') self.unmatched_words = Metrics.counter( self.__class__, 'unmatched_words') def process(self, element): word = element self.word_len_dist.update(len(word)) if re.match(self.pattern, word): yield element else: self.unmatched_words.inc() filtered_words = (words | 'FilterText' >> beam.ParDo(FilterTextFn('s.*'))) # [END metrics_usage_example] # pylint: enable=unused-variable # [START metrics_check_values_example] result = p.run() result.wait_until_finish() custom_distribution = result.metrics().query( MetricsFilter().with_name('word_len_dist'))['distributions'] custom_counter = result.metrics().query( MetricsFilter().with_name('unmatched_words'))['counters'] if custom_distribution: logging.info( 'The average word length was %d', custom_distribution[0].committed.mean) if custom_counter: logging.info( 'There were %d words that did not match the filter.', custom_counter[0].committed) # [END metrics_check_values_example] # There should be 4 words that did not match self.assertEqual(custom_counter[0].committed, 4) # The shortest word is 3 characters, the longest is 6 self.assertEqual(custom_distribution[0].committed.min, 3) self.assertEqual(custom_distribution[0].committed.max, 6) def test_model_join_using_side_inputs(self): name_list = ['a', 'b'] email_list = [['a', 'a@example.com'], ['b', 'b@example.com']] phone_list = [['a', 'x4312'], ['b', 'x8452']] result_path = self.create_temp_file() snippets.model_join_using_side_inputs( name_list, email_list, phone_list, result_path) expect = ['a; a@example.com; x4312', 'b; b@example.com; x8452'] self.assertEqual(expect, self.get_output(result_path)) def test_model_early_late_triggers(self): pipeline_options = PipelineOptions() pipeline_options.view_as(StandardOptions).streaming = True with TestPipeline(options=pipeline_options) as p: test_stream = ( TestStream().advance_watermark_to(10).add_elements([ 'a', 'a', 'a', 'b', 'b' ]).add_elements([ TimestampedValue('a', 10) ]).advance_watermark_to(20).advance_processing_time(60).add_elements( [TimestampedValue('a', 10)])) trigger = ( # [START model_early_late_triggers] AfterWatermark( early=AfterProcessingTime(delay=1 * 60), late=AfterCount(1)) # [END model_early_late_triggers] ) counts = ( p | test_stream | 'pair_with_one' >> beam.Map(lambda x: (x, 1)) | WindowInto( FixedWindows(15), trigger=trigger, allowed_lateness=20, accumulation_mode=AccumulationMode.DISCARDING) | 'group' >> beam.GroupByKey() | 'count' >> beam.Map(lambda word_ones: (word_ones[0], sum(word_ones[1])))) assert_that(counts, equal_to([('a', 4), ('b', 2), ('a', 1)])) def test_model_setting_trigger(self): pipeline_options = PipelineOptions( flags=['--streaming', '--allow_unsafe_triggers']) with TestPipeline(options=pipeline_options) as p: test_stream = ( TestStream().advance_watermark_to(10).add_elements( ['a', 'a', 'a', 'b', 'b']).advance_watermark_to(70).advance_processing_time(600)) pcollection = ( p | test_stream | 'pair_with_one' >> beam.Map(lambda x: (x, 1))) counts = ( pcollection | WindowInto( FixedWindows(1 * 60), trigger=AfterProcessingTime(10 * 60), accumulation_mode=AccumulationMode.DISCARDING) | 'group' >> beam.GroupByKey() | 'count' >> beam.Map(lambda word_ones: (word_ones[0], sum(word_ones[1])))) assert_that(counts, equal_to([('a', 3), ('b', 2)])) def test_model_composite_triggers(self): pipeline_options = PipelineOptions() pipeline_options.view_as(StandardOptions).streaming = True with TestPipeline(options=pipeline_options) as p: test_stream = ( TestStream().advance_watermark_to(10).add_elements( ['a', 'a', 'a', 'b', 'b']).advance_watermark_to(70).add_elements([ TimestampedValue('a', 10), TimestampedValue('a', 10), TimestampedValue('c', 10), TimestampedValue('c', 10) ]).advance_processing_time(600)) pcollection = ( p | test_stream | 'pair_with_one' >> beam.Map(lambda x: (x, 1))) counts = ( # [START model_composite_triggers] pcollection | WindowInto( FixedWindows(1 * 60), trigger=AfterWatermark(late=AfterProcessingTime(10 * 60)), allowed_lateness=2 * 24 * 60 * 60, accumulation_mode=AccumulationMode.DISCARDING) # [END model_composite_triggers] | 'group' >> beam.GroupByKey() | 'count' >> beam.Map(lambda word_ones: (word_ones[0], sum(word_ones[1])))) assert_that(counts, equal_to([('a', 3), ('b', 2), ('a', 2), ('c', 2)])) def test_model_other_composite_triggers(self): pipeline_options = PipelineOptions( flags=['--streaming', '--allow_unsafe_triggers']) with TestPipeline(options=pipeline_options) as p: test_stream = ( TestStream().advance_watermark_to(10).add_elements( ['a', 'a']).add_elements( ['a', 'b', 'b']).advance_processing_time(60).add_elements(['a'] * 100)) pcollection = ( p | test_stream | 'pair_with_one' >> beam.Map(lambda x: (x, 1))) counts = ( # [START model_other_composite_triggers] pcollection | WindowInto( FixedWindows(1 * 60), trigger=Repeatedly( AfterAny(AfterCount(100), AfterProcessingTime(1 * 60))), accumulation_mode=AccumulationMode.DISCARDING) # [END model_other_composite_triggers] | 'group' >> beam.GroupByKey() | 'count' >> beam.Map(lambda word_ones: (word_ones[0], sum(word_ones[1])))) assert_that(counts, equal_to([('a', 3), ('b', 2), ('a', 100)])) class CombineTest(unittest.TestCase): """Tests for model/combine.""" def test_global_sum(self): pc = [1, 2, 3] # [START global_sum] result = pc | beam.CombineGlobally(sum) # [END global_sum] self.assertEqual([6], result) def test_combine_values(self): occurences = [('cat', 1), ('cat', 5), ('cat', 9), ('dog', 5), ('dog', 2)] # [START combine_values] first_occurences = occurences | beam.GroupByKey() | beam.CombineValues(min) # [END combine_values] self.assertEqual({('cat', 1), ('dog', 2)}, set(first_occurences)) def test_combine_per_key(self): player_accuracies = [('cat', 1), ('cat', 5), ('cat', 9), ('cat', 1), ('dog', 5), ('dog', 2)] # [START combine_per_key] avg_accuracy_per_player = ( player_accuracies | beam.CombinePerKey(beam.combiners.MeanCombineFn())) # [END combine_per_key] self.assertEqual({('cat', 4.0), ('dog', 3.5)}, set(avg_accuracy_per_player)) def test_combine_concat(self): pc = ['a', 'b'] # [START combine_concat] def concat(values, separator=', '): return separator.join(values) with_commas = pc | beam.CombineGlobally(concat) with_dashes = pc | beam.CombineGlobally(concat, separator='-') # [END combine_concat] self.assertEqual(1, len(with_commas)) self.assertTrue(with_commas[0] in {'a, b', 'b, a'}) self.assertEqual(1, len(with_dashes)) self.assertTrue(with_dashes[0] in {'a-b', 'b-a'}) def test_bounded_sum(self): # [START combine_bounded_sum] pc = [1, 10, 100, 1000] def bounded_sum(values, bound=500): return min(sum(values), bound) small_sum = pc | beam.CombineGlobally(bounded_sum) # [500] large_sum = pc | beam.CombineGlobally(bounded_sum, bound=5000) # [1111] # [END combine_bounded_sum] self.assertEqual([500], small_sum) self.assertEqual([1111], large_sum) def test_combine_reduce(self): factors = [2, 3, 5, 7] # [START combine_reduce] import functools import operator product = factors | beam.CombineGlobally( functools.partial(functools.reduce, operator.mul), 1) # [END combine_reduce] self.assertEqual([210], product) def test_custom_average(self): pc = [2, 3, 5, 7] # [START combine_custom_average_define] class AverageFn(beam.CombineFn): def create_accumulator(self): return (0.0, 0) def add_input(self, sum_count, input): (sum, count) = sum_count return sum + input, count + 1 def merge_accumulators(self, accumulators): sums, counts = zip(*accumulators) return sum(sums), sum(counts) def extract_output(self, sum_count): (sum, count) = sum_count return sum / count if count else float('NaN') def compact(self, accumulator): # No-op return accumulator # [END combine_custom_average_define] # [START combine_custom_average_execute] average = pc | beam.CombineGlobally(AverageFn()) # [END combine_custom_average_execute] self.assertEqual([4.25], average) def test_keys(self): occurrences = [('cat', 1), ('cat', 5), ('dog', 5), ('cat', 9), ('dog', 2)] unique_keys = occurrences | snippets.Keys() self.assertEqual({'cat', 'dog'}, set(unique_keys)) def test_count(self): occurrences = ['cat', 'dog', 'cat', 'cat', 'dog'] perkey_counts = occurrences | snippets.Count() self.assertEqual({('cat', 3), ('dog', 2)}, set(perkey_counts)) def test_setting_fixed_windows(self): with TestPipeline() as p: unkeyed_items = p | beam.Create([22, 33, 55, 100, 115, 120]) items = ( unkeyed_items | 'key' >> beam.Map(lambda x: beam.window.TimestampedValue(('k', x), x))) # [START setting_fixed_windows] from apache_beam import window fixed_windowed_items = ( items | 'window' >> beam.WindowInto(window.FixedWindows(60))) # [END setting_fixed_windows] summed = ( fixed_windowed_items | 'group' >> beam.GroupByKey() | 'combine' >> beam.CombineValues(sum)) unkeyed = summed | 'unkey' >> beam.Map(lambda x: x[1]) assert_that(unkeyed, equal_to([110, 215, 120])) def test_setting_sliding_windows(self): with TestPipeline() as p: unkeyed_items = p | beam.Create([2, 16, 23]) items = ( unkeyed_items | 'key' >> beam.Map(lambda x: beam.window.TimestampedValue(('k', x), x))) # [START setting_sliding_windows] from apache_beam import window sliding_windowed_items = ( items | 'window' >> beam.WindowInto(window.SlidingWindows(30, 5))) # [END setting_sliding_windows] summed = ( sliding_windowed_items | 'group' >> beam.GroupByKey() | 'combine' >> beam.CombineValues(sum)) unkeyed = summed | 'unkey' >> beam.Map(lambda x: x[1]) assert_that(unkeyed, equal_to([2, 2, 2, 18, 23, 39, 39, 39, 41, 41])) def test_setting_session_windows(self): with TestPipeline() as p: unkeyed_items = p | beam.Create([2, 11, 16, 27]) items = ( unkeyed_items | 'key' >> beam.Map(lambda x: beam.window.TimestampedValue(('k', x), x * 60))) # [START setting_session_windows] from apache_beam import window session_windowed_items = ( items | 'window' >> beam.WindowInto(window.Sessions(10 * 60))) # [END setting_session_windows] summed = ( session_windowed_items | 'group' >> beam.GroupByKey() | 'combine' >> beam.CombineValues(sum)) unkeyed = summed | 'unkey' >> beam.Map(lambda x: x[1]) assert_that(unkeyed, equal_to([29, 27])) def test_setting_global_window(self): with TestPipeline() as p: unkeyed_items = p | beam.Create([2, 11, 16, 27]) items = ( unkeyed_items | 'key' >> beam.Map(lambda x: beam.window.TimestampedValue(('k', x), x))) # [START setting_global_window] from apache_beam import window global_windowed_items = ( items | 'window' >> beam.WindowInto(window.GlobalWindows())) # [END setting_global_window] summed = ( global_windowed_items | 'group' >> beam.GroupByKey() | 'combine' >> beam.CombineValues(sum)) unkeyed = summed | 'unkey' >> beam.Map(lambda x: x[1]) assert_that(unkeyed, equal_to([56])) def test_setting_timestamp(self): with TestPipeline() as p: unkeyed_items = p | beam.Create([12, 30, 60, 61, 66]) items = (unkeyed_items | 'key' >> beam.Map(lambda x: ('k', x))) def extract_timestamp_from_log_entry(entry): return entry[1] # [START setting_timestamp] class AddTimestampDoFn(beam.DoFn): def process(self, element): # Extract the numeric Unix seconds-since-epoch timestamp to be # associated with the current log entry. unix_timestamp = extract_timestamp_from_log_entry(element) # Wrap and emit the current entry and new timestamp in a # TimestampedValue. yield beam.window.TimestampedValue(element, unix_timestamp) timestamped_items = items | 'timestamp' >> beam.ParDo(AddTimestampDoFn()) # [END setting_timestamp] fixed_windowed_items = ( timestamped_items | 'window' >> beam.WindowInto(beam.window.FixedWindows(60))) summed = ( fixed_windowed_items | 'group' >> beam.GroupByKey() | 'combine' >> beam.CombineValues(sum)) unkeyed = summed | 'unkey' >> beam.Map(lambda x: x[1]) assert_that(unkeyed, equal_to([42, 187])) class PTransformTest(unittest.TestCase): """Tests for PTransform.""" def test_composite(self): # [START model_composite_transform] class ComputeWordLengths(beam.PTransform): def expand(self, pcoll): # Transform logic goes here. return pcoll | beam.Map(lambda x: len(x)) # [END model_composite_transform] with TestPipeline() as p: lengths = p | beam.Create(["a", "ab", "abc"]) | ComputeWordLengths() assert_that(lengths, equal_to([1, 2, 3])) class SlowlyChangingSideInputsTest(unittest.TestCase): """Tests for PTransform.""" def test_side_input_slow_update(self): temp_file = tempfile.NamedTemporaryFile(delete=True) src_file_pattern = temp_file.name temp_file.close() first_ts = math.floor(time.time()) - 30 interval = 5 main_input_windowing_interval = 7 # aligning timestamp to get persistent results first_ts = first_ts - ( first_ts % (interval * main_input_windowing_interval)) last_ts = first_ts + 45 for i in range(-1, 10, 1): count = i + 2 idstr = str(first_ts + interval * i) with open(src_file_pattern + idstr, "w") as f: for j in range(count): f.write('f' + idstr + 'a' + str(j) + '\n') sample_main_input_elements = ([ first_ts - 2, # no output due to no SI first_ts + 1, # First window first_ts + 8, # Second window first_ts + 15, # Third window first_ts + 22, # Fourth window ]) pipeline, pipeline_result = snippets.side_input_slow_update( src_file_pattern, first_ts, last_ts, interval, sample_main_input_elements, main_input_windowing_interval) try: with pipeline: pipeline_result = ( pipeline_result | 'AddKey' >> beam.Map(lambda v: ('key', v)) | combiners.Count.PerKey()) assert_that( pipeline_result, equal_to([('key', 3), ('key', 4), ('key', 6), ('key', 7)])) finally: for i in range(-1, 10, 1): os.unlink(src_file_pattern + str(first_ts + interval * i)) class ValueProviderInfoTest(unittest.TestCase): """Tests for accessing value provider info after run.""" def test_accessing_valueprovider_info_after_run(self): with self.assertLogs(level='INFO') as log_capture: snippets.accessing_valueprovider_info_after_run() expected_log_message = "The string value is" self.assertTrue( any(expected_log_message in log for log in log_capture.output), f"Expected log message '{expected_log_message}' not found in logs: " f"{log_capture.output}") if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.125381 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/0000755000175100017510000000000015206113157023546 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/__init__.py0000644000175100017510000000142015206113032025644 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.1943882 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/0000755000175100017510000000000015206113157026035 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/__init__.py0000644000175100017510000000142015206113032030133 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/approximatequantiles.py0000644000175100017510000000310215206113032032652 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: ApproximateQuantiles # description: Demonstration of ApproximateQuantiles transform usage. # multifile: false # default_example: false # context_line: 37 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - integers def approximatequantiles(test=None): # [START approximatequantiles] import apache_beam as beam with beam.Pipeline() as pipeline: quantiles = ( pipeline | 'Create data' >> beam.Create(list(range(1001))) | 'Compute quantiles' >> beam.ApproximateQuantiles.Globally(5) | beam.Map(print)) # [END approximatequantiles] if test: test(quantiles) if __name__ == '__main__': approximatequantiles() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/approximatequantiles_test.py0000644000175100017510000000275415206113032033725 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import mock from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from . import approximatequantiles @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.' 'approximatequantiles.print', lambda x: x) class ApproximateQuantilesTest(unittest.TestCase): def test_approximatequantiles(self): def check_result(quantiles): assert_that(quantiles, equal_to([[0, 250, 500, 750, 1000]])) approximatequantiles.approximatequantiles(test=check_result) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/approximateunique.py0000644000175100017510000000305215206113032032157 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: ApproximateUnique # description: Demonstration of ApproximateUnique transform usage. # multifile: false # default_example: false # context_line: 37 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - integers def approximateunique(test=None): # [START approximateunique] import random import apache_beam as beam with beam.Pipeline() as pipeline: data = list(range(1000)) random.shuffle(data) result = ( pipeline | 'create' >> beam.Create(data) | 'get_estimate' >> beam.ApproximateUnique.Globally(size=16) | beam.Map(print)) # [END approximateunique] if test: test(result) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/approximateunique_test.py0000644000175100017510000000332715206113032033223 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import math import unittest import mock import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from . import approximateunique # pytype: skip-file @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.' 'approximateunique.print', lambda x: x) class ApproximateUniqueTest(unittest.TestCase): def test_approximateunique(self): def check_result(approx_count): actual_count = 1000 sample_size = 16 error = 2 / math.sqrt(sample_size) assert_that( approx_count | 'compare' >> beam.FlatMap( lambda x: [abs(x - actual_count) * 1.0 / actual_count <= error]), equal_to([True])) approximateunique.approximateunique(test=check_result) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/batchelements.py0000644000175100017510000000334015206113032031215 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: BatchElements # description: Demonstration of BatchElements transform usage. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - group def batchelements(test=None): # [START batchelements] import apache_beam as beam with beam.Pipeline() as pipeline: batches = ( pipeline | 'Create produce' >> beam.Create([ '🍓', '🥕', '🍆', '🍅', '🥕', '🍅', '🌽', '🥕', '🍅', '🍆', ]) | beam.BatchElements(min_batch_size=3, max_batch_size=5) | beam.Map(print)) # [END batchelements] if test: test(batches) if __name__ == '__main__': batchelements() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/batchelements_test.py0000644000175100017510000000407115206113032032256 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest import mock import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from . import batchelements def check_batches(actual): # How the elements are grouped is not guaranteed, so we just # check that all elements are lists and then count the elements # to make sure none are lost. all_elements_are_lists = ( actual | 'Check type' >> beam.Map(lambda x: isinstance(x, list)) | 'All elements are lists' >> beam.CombineGlobally(all)) assert_that(all_elements_are_lists, equal_to([True])) assert_that( actual | beam.FlatMap(lambda x: x) | 'Count' >> beam.combiners.Count.PerElement(), equal_to([('🍓', 1), ('🥕', 3), ('🍆', 2), ('🍅', 3), ('🌽', 1)]), label='Check counts') def identity(x): return x @mock.patch('apache_beam.Pipeline', TestPipeline) # pylint: disable=line-too-long @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.batchelements.print', identity) # pylint: enable=line-too-long class BatchElementsTest(unittest.TestCase): def test_batchelements(self): batchelements.batchelements(check_batches) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/cogroupbykey.py0000644000175100017510000000367615206113032031135 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CoGroupByKeyMerge # description: Demonstration of CoGroupByKey transform usage. # multifile: false # default_example: false # context_line: 43 # categories: # - Core Transforms # - Joins # complexity: BASIC # tags: # - transforms # - strings # - integers # - tuples # - pairs # - group def cogroupbykey(test=None): # [START cogroupbykey] import apache_beam as beam with beam.Pipeline() as pipeline: icon_pairs = pipeline | 'Create icons' >> beam.Create([ ('Apple', '🍎'), ('Apple', '🍏'), ('Eggplant', '🍆'), ('Tomato', '🍅'), ]) duration_pairs = pipeline | 'Create durations' >> beam.Create([ ('Apple', 'perennial'), ('Carrot', 'biennial'), ('Tomato', 'perennial'), ('Tomato', 'annual'), ]) plants = (({ 'icons': icon_pairs, 'durations': duration_pairs }) | 'Merge' >> beam.CoGroupByKey() | beam.Map(print)) # [END cogroupbykey] if test: test(plants) if __name__ == '__main__': cogroupbykey() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/cogroupbykey_test.py0000644000175100017510000000366715206113032032174 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import cogroupbykey def check_plants(actual): expected = '''[START plants] ('Apple', {'icons': ['🍎', '🍏'], 'durations': ['perennial']}) ('Carrot', {'icons': [], 'durations': ['biennial']}) ('Tomato', {'icons': ['🍅'], 'durations': ['perennial', 'annual']}) ('Eggplant', {'icons': ['🍆'], 'durations': []}) [END plants]'''.splitlines()[1:-1] # Make it deterministic by sorting all sublists in each element. def normalize_element(elem): name, details = elem details['icons'] = sorted(details['icons']) details['durations'] = sorted(details['durations']) return name, details assert_matches_stdout(actual, expected, normalize_element) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.cogroupbykey.print', str) class CoGroupByKeyTest(unittest.TestCase): def test_cogroupbykey(self): cogroupbykey.cogroupbykey(check_plants) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineglobally_combinefn.py0000644000175100017510000000544015206113032033564 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombineGloballyCombineFn # description: Demonstration of CombineGlobally transform usage with a CombineFn. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - group def combineglobally_combinefn(test=None): # [START combineglobally_combinefn] import apache_beam as beam class PercentagesFn(beam.CombineFn): def create_accumulator(self): return {} def add_input(self, accumulator, input): # accumulator == {} # input == '🥕' if input not in accumulator: accumulator[input] = 0 # {'🥕': 0} accumulator[input] += 1 # {'🥕': 1} return accumulator def merge_accumulators(self, accumulators): # accumulators == [ # {'🥕': 1, '🍅': 2}, # {'🥕': 1, '🍅': 1, '🍆': 1}, # {'🥕': 1, '🍅': 3}, # ] merged = {} for accum in accumulators: for item, count in accum.items(): if item not in merged: merged[item] = 0 merged[item] += count # merged == {'🥕': 3, '🍅': 6, '🍆': 1} return merged def extract_output(self, accumulator): # accumulator == {'🥕': 3, '🍅': 6, '🍆': 1} total = sum(accumulator.values()) # 10 percentages = {item: count / total for item, count in accumulator.items()} # percentages == {'🥕': 0.3, '🍅': 0.6, '🍆': 0.1} return percentages with beam.Pipeline() as pipeline: percentages = ( pipeline | 'Create produce' >> beam.Create( ['🥕', '🍅', '🍅', '🥕', '🍆', '🍅', '🍅', '🍅', '🥕', '🍅']) | 'Get percentages' >> beam.CombineGlobally(PercentagesFn()) | beam.Map(print)) # [END combineglobally_combinefn] if test: test(percentages) if __name__ == '__main__': combineglobally_combinefn() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineglobally_function.py0000644000175100017510000000425315206113032033452 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombineGloballyFunction # description: Demonstration of CombineGlobally transform usage with a function. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - group def combineglobally_function(test=None): # [START combineglobally_function] import apache_beam as beam def get_common_items(sets): # set.intersection() takes multiple sets as separete arguments. # We unpack the `sets` list into multiple arguments with the * operator. # The combine transform might give us an empty list of `sets`, # so we use a list with an empty set as a default value. return set.intersection(*(sets or [set()])) with beam.Pipeline() as pipeline: common_items = ( pipeline | 'Create produce' >> beam.Create([ {'🍓', '🥕', '🍌', '🍅', '🌶️'}, {'🍇', '🥕', '🥝', '🍅', '🥔'}, {'🍉', '🥕', '🍆', '🍅', '🍍'}, {'🥑', '🥕', '🌽', '🍅', '🥥'}, ]) | 'Get common items' >> beam.CombineGlobally(get_common_items) | beam.Map(print)) # [END combineglobally_function] if test: test(common_items) if __name__ == '__main__': combineglobally_function() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineglobally_lambda.py0000644000175100017510000000356115206113032033046 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombineGloballyLambda # description: Demonstration of CombineGlobally transform usage with a lambda function. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - group def combineglobally_lambda(test=None): # [START combineglobally_lambda] import apache_beam as beam with beam.Pipeline() as pipeline: common_items = ( pipeline | 'Create produce' >> beam.Create([ {'🍓', '🥕', '🍌', '🍅', '🌶️'}, {'🍇', '🥕', '🥝', '🍅', '🥔'}, {'🍉', '🥕', '🍆', '🍅', '🍍'}, {'🥑', '🥕', '🌽', '🍅', '🥥'}, ]) | 'Get common items' >> beam.CombineGlobally(lambda sets: set.intersection(*(sets or [set()]))) | beam.Map(print)) # [END combineglobally_lambda] if test: test(common_items) if __name__ == '__main__': combineglobally_lambda() ././@PaxHeader0000000000000000000000000000021500000000000010213 xustar00119 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineglobally_multiple_arguments.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineglobally_multiple_arg0000644000175100017510000000405215206113032033657 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombineGloballyMultipleArguments # description: Demonstration of CombineGlobally transform usage with multiple arguments. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - group def combineglobally_multiple_arguments(test=None): # [START combineglobally_multiple_arguments] import apache_beam as beam with beam.Pipeline() as pipeline: common_items_with_exceptions = ( pipeline | 'Create produce' >> beam.Create([ {'🍓', '🥕', '🍌', '🍅', '🌶️'}, {'🍇', '🥕', '🥝', '🍅', '🥔'}, {'🍉', '🥕', '🍆', '🍅', '🍍'}, {'🥑', '🥕', '🌽', '🍅', '🥥'}, ]) | 'Get common items with exceptions' >> beam.CombineGlobally( lambda sets, exclude: \ set.intersection(*(sets or [set()])) - exclude, exclude={'🥕'}) | beam.Map(print) ) # [END combineglobally_multiple_arguments] if test: test(common_items_with_exceptions) if __name__ == '__main__': combineglobally_multiple_arguments() ././@PaxHeader0000000000000000000000000000021300000000000010211 xustar00117 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineglobally_side_inputs_dict.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineglobally_side_inputs_0000644000175100017510000000371115206113032033661 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file def combineglobally_side_inputs_dict(test=None): # [START combineglobally_side_inputs_dict] import apache_beam as beam def get_custom_common_items(sets, options): sets = sets or [set()] common_items = set.intersection(*sets) common_items |= options['include'] # union common_items &= options['exclude'] # intersection return common_items with beam.Pipeline() as pipeline: options = pipeline | 'Create options' >> beam.Create([ ('exclude', {'🥕'}), ('include', {'🍇', '🌽'}), ]) custom_common_items = ( pipeline | 'Create produce' >> beam.Create([ {'🍓', '🥕', '🍌', '🍅', '🌶️'}, {'🍇', '🥕', '🥝', '🍅', '🥔'}, {'🍉', '🥕', '🍆', '🍅', '🍍'}, {'🥑', '🥕', '🌽', '🍅', '🥥'}, ]) | 'Get common items' >> beam.CombineGlobally( get_custom_common_items, options=beam.pvalue.AsDict(options)) | beam.Map(print)) # [END combineglobally_side_inputs_dict] if test: test(custom_common_items) if __name__ == '__main__': combineglobally_side_inputs_dict() ././@PaxHeader0000000000000000000000000000021300000000000010211 xustar00117 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineglobally_side_inputs_iter.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineglobally_side_inputs_0000644000175100017510000000340715206113032033663 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file def combineglobally_side_inputs_iter(test=None): # [START combineglobally_side_inputs_iter] import apache_beam as beam with beam.Pipeline() as pipeline: exclude = pipeline | 'Create exclude' >> beam.Create(['🥕']) common_items_with_exceptions = ( pipeline | 'Create produce' >> beam.Create([ {'🍓', '🥕', '🍌', '🍅', '🌶️'}, {'🍇', '🥕', '🥝', '🍅', '🥔'}, {'🍉', '🥕', '🍆', '🍅', '🍍'}, {'🥑', '🥕', '🌽', '🍅', '🥥'}, ]) | 'Get common items with exceptions' >> beam.CombineGlobally( lambda sets, exclude: \ set.intersection(*(sets or [set()])) - set(exclude), exclude=beam.pvalue.AsIter(exclude)) | beam.Map(print) ) # [END combineglobally_side_inputs_iter] if test: test(common_items_with_exceptions) if __name__ == '__main__': combineglobally_side_inputs_iter() ././@PaxHeader0000000000000000000000000000022000000000000010207 xustar00122 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineglobally_side_inputs_singleton.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineglobally_side_inputs_0000644000175100017510000000350715206113032033664 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file def combineglobally_side_inputs_singleton(test=None): # [START combineglobally_side_inputs_singleton] import apache_beam as beam with beam.Pipeline() as pipeline: single_exclude = pipeline | 'Create single_exclude' >> beam.Create(['🥕']) common_items_with_exceptions = ( pipeline | 'Create produce' >> beam.Create([ {'🍓', '🥕', '🍌', '🍅', '🌶️'}, {'🍇', '🥕', '🥝', '🍅', '🥔'}, {'🍉', '🥕', '🍆', '🍅', '🍍'}, {'🥑', '🥕', '🌽', '🍅', '🥥'}, ]) | 'Get common items with exceptions' >> beam.CombineGlobally( lambda sets, single_exclude: \ set.intersection(*(sets or [set()])) - {single_exclude}, single_exclude=beam.pvalue.AsSingleton(single_exclude)) | beam.Map(print) ) # [END combineglobally_side_inputs_singleton] if test: test(common_items_with_exceptions) if __name__ == '__main__': combineglobally_side_inputs_singleton() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineglobally_test.py0000644000175100017510000000760315206113032032606 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline # from . import combineglobally_side_inputs_dict # from . import combineglobally_side_inputs_iter from . import combineglobally_combinefn from . import combineglobally_function from . import combineglobally_lambda from . import combineglobally_multiple_arguments from . import combineglobally_side_inputs_singleton def check_common_items(actual): expected = '''[START common_items] {'🍅', '🥕'} [END common_items]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_common_items_with_exceptions(actual): expected = '''[START common_items_with_exceptions] {'🍅'} [END common_items_with_exceptions]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_custom_common_items(actual): expected = '''[START custom_common_items] {'🍅', '🍇', '🌽'} [END custom_common_items]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_percentages(actual): expected = '''[START percentages] {'🥕': 0.3, '🍅': 0.6, '🍆': 0.1} [END percentages]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) # pylint: disable=line-too-long @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combineglobally_function.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combineglobally_lambda.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combineglobally_multiple_arguments.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combineglobally_side_inputs_singleton.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combineglobally_combinefn.print', str) # pylint: enable=line-too-long class CombineGloballyTest(unittest.TestCase): def test_combineglobally_function(self): combineglobally_function.combineglobally_function(check_common_items) def test_combineglobally_lambda(self): combineglobally_lambda.combineglobally_lambda(check_common_items) def test_combineglobally_multiple_arguments(self): combineglobally_multiple_arguments.combineglobally_multiple_arguments( check_common_items_with_exceptions) def test_combineglobally_side_inputs_singleton(self): combineglobally_side_inputs_singleton.combineglobally_side_inputs_singleton( check_common_items_with_exceptions) # TODO: enable side inputs tests after # [https://github.com/apache/beam/issues/19851] is fixed. # def test_combineglobally_side_inputs_iter(self): # combineglobally.combineglobally_side_inputs_iter( # check_common_items_with_exceptions) # def test_combineglobally_side_inputs_dict(self): # combineglobally.combineglobally_side_inputs_dict( # check_custom_common_items) def test_combineglobally_combinefn(self): combineglobally_combinefn.combineglobally_combinefn(check_percentages) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineperkey_combinefn.py0000644000175100017510000000465715206113032033267 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombinePerKeyCombineFn # description: Demonstration of CombinePerKey transform usage with side inputs. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - group def combineperkey_combinefn(test=None): # [START combineperkey_combinefn] import apache_beam as beam class AverageFn(beam.CombineFn): def create_accumulator(self): sum = 0.0 count = 0 accumulator = sum, count return accumulator def add_input(self, accumulator, input): sum, count = accumulator return sum + input, count + 1 def merge_accumulators(self, accumulators): # accumulators = [(sum1, count1), (sum2, count2), (sum3, count3), ...] sums, counts = zip(*accumulators) # sums = [sum1, sum2, sum3, ...] # counts = [count1, count2, count3, ...] return sum(sums), sum(counts) def extract_output(self, accumulator): sum, count = accumulator if count == 0: return float('NaN') return sum / count with beam.Pipeline() as pipeline: average = ( pipeline | 'Create plant counts' >> beam.Create([ ('🥕', 3), ('🥕', 2), ('🍆', 1), ('🍅', 4), ('🍅', 5), ('🍅', 3), ]) | 'Average' >> beam.CombinePerKey(AverageFn()) | beam.Map(print)) # [END combineperkey_combinefn] if test: test(average) if __name__ == '__main__': combineperkey_combinefn() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineperkey_function.py0000644000175100017510000000352515206113032033145 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombinePerKeyFunction # description: Demonstration of CombinePerKey transform usage with side inputs. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - group def combineperkey_function(test=None): # [START combineperkey_function] import apache_beam as beam def saturated_sum(values): max_value = 8 return min(sum(values), max_value) with beam.Pipeline() as pipeline: saturated_total = ( pipeline | 'Create plant counts' >> beam.Create([ ('🥕', 3), ('🥕', 2), ('🍆', 1), ('🍅', 4), ('🍅', 5), ('🍅', 3), ]) | 'Saturated sum' >> beam.CombinePerKey(saturated_sum) | beam.Map(print)) # [END combineperkey_function] if test: test(saturated_total) if __name__ == '__main__': combineperkey_function() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineperkey_lambda.py0000644000175100017510000000342115206113032032533 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombinePerKeyLambda # description: Demonstration of CombinePerKey transform usage with side inputs. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - group def combineperkey_lambda(test=None): # [START combineperkey_lambda] import apache_beam as beam with beam.Pipeline() as pipeline: saturated_total = ( pipeline | 'Create plant counts' >> beam.Create([ ('🥕', 3), ('🥕', 2), ('🍆', 1), ('🍅', 4), ('🍅', 5), ('🍅', 3), ]) | 'Saturated sum' >> beam.CombinePerKey(lambda values: min(sum(values), 8)) | beam.Map(print)) # [END combineperkey_lambda] if test: test(saturated_total) if __name__ == '__main__': combineperkey_lambda() ././@PaxHeader0000000000000000000000000000021300000000000010211 xustar00117 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineperkey_multiple_arguments.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineperkey_multiple_argum0000644000175100017510000000356115206113032033717 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombinePerKeyMultipleArguments # description: Demonstration of CombinePerKey transform usage with side inputs. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - group def combineperkey_multiple_arguments(test=None): # [START combineperkey_multiple_arguments] import apache_beam as beam with beam.Pipeline() as pipeline: saturated_total = ( pipeline | 'Create plant counts' >> beam.Create([ ('🥕', 3), ('🥕', 2), ('🍆', 1), ('🍅', 4), ('🍅', 5), ('🍅', 3), ]) | 'Saturated sum' >> beam.CombinePerKey( lambda values, max_value: min(sum(values), max_value), max_value=8) | beam.Map(print)) # [END combineperkey_multiple_arguments] if test: test(saturated_total) if __name__ == '__main__': combineperkey_multiple_arguments() ././@PaxHeader0000000000000000000000000000021100000000000010207 xustar00115 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineperkey_side_inputs_dict.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineperkey_side_inputs_di0000644000175100017510000000352715206113032033675 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file def combineperkey_side_inputs_dict(test=None): # [START combineperkey_side_inputs_dict] import apache_beam as beam def bounded_sum(values, data_range): min_value = data_range['min'] result = sum(values) if result < min_value: return min_value max_value = data_range['max'] if result > max_value: return max_value return result with beam.Pipeline() as pipeline: data_range = pipeline | 'Create data_range' >> beam.Create([ ('min', 2), ('max', 8), ]) bounded_total = ( pipeline | 'Create plant counts' >> beam.Create([ ('🥕', 3), ('🥕', 2), ('🍆', 1), ('🍅', 4), ('🍅', 5), ('🍅', 3), ]) | 'Bounded sum' >> beam.CombinePerKey( bounded_sum, data_range=beam.pvalue.AsDict(data_range)) | beam.Map(print)) # [END combineperkey_side_inputs_dict] if test: test(bounded_total) if __name__ == '__main__': combineperkey_side_inputs_dict() ././@PaxHeader0000000000000000000000000000021100000000000010207 xustar00115 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineperkey_side_inputs_iter.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineperkey_side_inputs_it0000644000175100017510000000345515206113032033715 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file def combineperkey_side_inputs_iter(test=None): # [START combineperkey_side_inputs_iter] import apache_beam as beam def bounded_sum(values, data_range): min_value = min(data_range) result = sum(values) if result < min_value: return min_value max_value = max(data_range) if result > max_value: return max_value return result with beam.Pipeline() as pipeline: data_range = pipeline | 'Create data_range' >> beam.Create([2, 4, 8]) bounded_total = ( pipeline | 'Create plant counts' >> beam.Create([ ('🥕', 3), ('🥕', 2), ('🍆', 1), ('🍅', 4), ('🍅', 5), ('🍅', 3), ]) | 'Bounded sum' >> beam.CombinePerKey( bounded_sum, data_range=beam.pvalue.AsIter(data_range)) | beam.Map(print)) # [END combineperkey_side_inputs_iter] if test: test(bounded_total) if __name__ == '__main__': combineperkey_side_inputs_iter() ././@PaxHeader0000000000000000000000000000021600000000000010214 xustar00120 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineperkey_side_inputs_singleton.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineperkey_side_inputs_si0000644000175100017510000000320115206113032033701 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file def combineperkey_side_inputs_singleton(test=None): # [START combineperkey_side_inputs_singleton] import apache_beam as beam with beam.Pipeline() as pipeline: max_value = pipeline | 'Create max_value' >> beam.Create([8]) saturated_total = ( pipeline | 'Create plant counts' >> beam.Create([ ('🥕', 3), ('🥕', 2), ('🍆', 1), ('🍅', 4), ('🍅', 5), ('🍅', 3), ]) | 'Saturated sum' >> beam.CombinePerKey( lambda values, max_value: min(sum(values), max_value), max_value=beam.pvalue.AsSingleton(max_value)) | beam.Map(print)) # [END combineperkey_side_inputs_singleton] if test: test(saturated_total) if __name__ == '__main__': combineperkey_side_inputs_singleton() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineperkey_simple.py0000644000175100017510000000327315206113032032611 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombinePerKeySimple # description: Demonstration of CombinePerKey transform usage. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - group def combineperkey_simple(test=None): # [START combineperkey_simple] import apache_beam as beam with beam.Pipeline() as pipeline: total = ( pipeline | 'Create plant counts' >> beam.Create([ ('🥕', 3), ('🥕', 2), ('🍆', 1), ('🍅', 4), ('🍅', 5), ('🍅', 3), ]) | 'Sum' >> beam.CombinePerKey(sum) | beam.Map(print)) # [END combineperkey_simple] if test: test(total) if __name__ == '__main__': combineperkey_simple() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combineperkey_test.py0000644000175100017510000001033315206113032032272 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import combineperkey_combinefn from . import combineperkey_function from . import combineperkey_lambda from . import combineperkey_multiple_arguments from . import combineperkey_side_inputs_dict from . import combineperkey_side_inputs_iter from . import combineperkey_side_inputs_singleton from . import combineperkey_simple def check_total(actual): expected = '''[START total] ('🥕', 5) ('🍆', 1) ('🍅', 12) [END total]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_saturated_total(actual): expected = '''[START saturated_total] ('🥕', 5) ('🍆', 1) ('🍅', 8) [END saturated_total]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_bounded_total(actual): expected = '''[START bounded_total] ('🥕', 5) ('🍆', 2) ('🍅', 8) [END bounded_total]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_average(actual): expected = '''[START average] ('🥕', 2.5) ('🍆', 1.0) ('🍅', 4.0) [END average]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combineperkey_simple.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combineperkey_function.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combineperkey_lambda.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combineperkey_multiple_arguments.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combineperkey_multiple_arguments.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combineperkey_side_inputs_singleton.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combineperkey_side_inputs_iter.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combineperkey_side_inputs_dict.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combineperkey_combinefn.print', str) class CombinePerKeyTest(unittest.TestCase): def test_combineperkey_simple(self): combineperkey_simple.combineperkey_simple(check_total) def test_combineperkey_function(self): combineperkey_function.combineperkey_function(check_saturated_total) def test_combineperkey_lambda(self): combineperkey_lambda.combineperkey_lambda(check_saturated_total) def test_combineperkey_multiple_arguments(self): combineperkey_multiple_arguments.combineperkey_multiple_arguments( check_saturated_total) def test_combineperkey_side_inputs_singleton(self): combineperkey_side_inputs_singleton.combineperkey_side_inputs_singleton( check_saturated_total) def test_combineperkey_side_inputs_iter(self): combineperkey_side_inputs_iter.combineperkey_side_inputs_iter( check_bounded_total) def test_combineperkey_side_inputs_dict(self): combineperkey_side_inputs_dict.combineperkey_side_inputs_dict( check_bounded_total) def test_combineperkey_combinefn(self): combineperkey_combinefn.combineperkey_combinefn(check_average) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combinevalues_combinefn.py0000644000175100017510000000560415206113032033260 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombineValuesCombineFn # description: Demonstration of CombineValue transform usage with CombineFn. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - pairs # - combine def combinevalues_combinefn(test=None): # [START combinevalues_combinefn] import apache_beam as beam class AverageFn(beam.CombineFn): def create_accumulator(self): return {} def add_input(self, accumulator, input): # accumulator == {} # input == '🥕' if input not in accumulator: accumulator[input] = 0 # {'🥕': 0} accumulator[input] += 1 # {'🥕': 1} return accumulator def merge_accumulators(self, accumulators): # accumulators == [ # {'🥕': 1, '🍅': 1}, # {'🥕': 1, '🍅': 1, '🍆': 1}, # ] merged = {} for accum in accumulators: for item, count in accum.items(): if item not in merged: merged[item] = 0 merged[item] += count # merged == {'🥕': 2, '🍅': 2, '🍆': 1} return merged def extract_output(self, accumulator): # accumulator == {'🥕': 2, '🍅': 2, '🍆': 1} total = sum(accumulator.values()) # 5 percentages = {item: count / total for item, count in accumulator.items()} # percentages == {'🥕': 0.4, '🍅': 0.4, '🍆': 0.2} return percentages with beam.Pipeline() as pipeline: percentages_per_season = ( pipeline | 'Create produce' >> beam.Create([ ('spring', ['🥕', '🍅', '🥕', '🍅', '🍆']), ('summer', ['🥕', '🍅', '🌽', '🍅', '🍅']), ('fall', ['🥕', '🥕', '🍅', '🍅']), ('winter', ['🍆', '🍆']), ]) | 'Average' >> beam.CombineValues(AverageFn()) | beam.Map(print)) # [END combinevalues_combinefn] if test: test(percentages_per_season) if __name__ == '__main__': combinevalues_combinefn() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combinevalues_function.py0000644000175100017510000000342715206113032033146 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombineValuesFunction # description: Demonstration of CombineValue transform usage with a function. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - pairs # - combine def combinevalues_function(test=None): # [START combinevalues_function] import apache_beam as beam def saturated_sum(values): max_value = 8 return min(sum(values), max_value) with beam.Pipeline() as pipeline: saturated_total = ( pipeline | 'Create plant counts' >> beam.Create([ ('🥕', [3, 2]), ('🍆', [1]), ('🍅', [4, 5, 3]), ]) | 'Saturated sum' >> beam.CombineValues(saturated_sum) | beam.Map(print)) # [END combinevalues_function] if test: test(saturated_total) if __name__ == '__main__': combinevalues_function() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combinevalues_lambda.py0000644000175100017510000000333215206113032032534 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombineValuesLambda # description: Demonstration of CombineValue transform usage with a lambda function. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - pairs # - combine def combinevalues_lambda(test=None): # [START combinevalues_lambda] import apache_beam as beam with beam.Pipeline() as pipeline: saturated_total = ( pipeline | 'Create plant counts' >> beam.Create([ ('🥕', [3, 2]), ('🍆', [1]), ('🍅', [4, 5, 3]), ]) | 'Saturated sum' >> beam.CombineValues(lambda values: min(sum(values), 8)) | beam.Map(print)) # [END combinevalues_lambda] if test: test(saturated_total) if __name__ == '__main__': combinevalues_lambda() ././@PaxHeader0000000000000000000000000000021300000000000010211 xustar00117 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combinevalues_multiple_arguments.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combinevalues_multiple_argum0000644000175100017510000000352215206113032033714 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombineValuesMultipleArguments # description: Demonstration of CombineValue transform usage with a lambda function with multiple arguments. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - pairs # - combine def combinevalues_multiple_arguments(test=None): # [START combinevalues_multiple_arguments] import apache_beam as beam with beam.Pipeline() as pipeline: saturated_total = ( pipeline | 'Create plant counts' >> beam.Create([ ('🥕', [3, 2]), ('🍆', [1]), ('🍅', [4, 5, 3]), ]) | 'Saturated sum' >> beam.CombineValues( lambda values, max_value: min(sum(values), max_value), max_value=8) | beam.Map(print)) # [END combinevalues_multiple_arguments] if test: test(saturated_total) if __name__ == '__main__': combinevalues_multiple_arguments() ././@PaxHeader0000000000000000000000000000021100000000000010207 xustar00115 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combinevalues_side_inputs_dict.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combinevalues_side_inputs_di0000644000175100017510000000423215206113032033667 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombineValuesSideInputsDict # description: Demonstration of CombineValue transform usage with side inputs in a dictionary format. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - pairs # - combine def combinevalues_side_inputs_dict(test=None): # [START combinevalues_side_inputs_dict] import apache_beam as beam def bounded_sum(values, data_range): min_value = data_range['min'] result = sum(values) if result < min_value: return min_value max_value = data_range['max'] if result > max_value: return max_value return result with beam.Pipeline() as pipeline: data_range = pipeline | 'Create data_range' >> beam.Create([ ('min', 2), ('max', 8), ]) bounded_total = ( pipeline | 'Create plant counts' >> beam.Create([ ('🥕', [3, 2]), ('🍆', [1]), ('🍅', [4, 5, 3]), ]) | 'Bounded sum' >> beam.CombineValues( bounded_sum, data_range=beam.pvalue.AsDict(data_range)) | beam.Map(print)) # [END combinevalues_side_inputs_dict] if test: test(bounded_total) if __name__ == '__main__': combinevalues_side_inputs_dict() ././@PaxHeader0000000000000000000000000000021100000000000010207 xustar00115 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combinevalues_side_inputs_iter.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combinevalues_side_inputs_it0000644000175100017510000000414115206113032033706 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombineValuesSideInputsIter # description: Demonstration of CombineValue transform usage with side inputs as Iter. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - pairs # - combine def combinevalues_side_inputs_iter(test=None): # [START combinevalues_side_inputs_iter] import apache_beam as beam def bounded_sum(values, data_range): min_value = min(data_range) result = sum(values) if result < min_value: return min_value max_value = max(data_range) if result > max_value: return max_value return result with beam.Pipeline() as pipeline: data_range = pipeline | 'Create data_range' >> beam.Create([2, 4, 8]) bounded_total = ( pipeline | 'Create plant counts' >> beam.Create([ ('🥕', [3, 2]), ('🍆', [1]), ('🍅', [4, 5, 3]), ]) | 'Bounded sum' >> beam.CombineValues( bounded_sum, data_range=beam.pvalue.AsIter(data_range)) | beam.Map(print)) # [END combinevalues_side_inputs_iter] if test: test(bounded_total) if __name__ == '__main__': combinevalues_side_inputs_iter() ././@PaxHeader0000000000000000000000000000021600000000000010214 xustar00120 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combinevalues_side_inputs_singleton.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combinevalues_side_inputs_si0000644000175100017510000000370315206113032033710 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombineValuesSideInputsSingleton # description: Demonstration of CombineValue transform usage with side inputs from a singleton. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - pairs # - combine def combinevalues_side_inputs_singleton(test=None): # [START combinevalues_side_inputs_singleton] import apache_beam as beam with beam.Pipeline() as pipeline: max_value = pipeline | 'Create max_value' >> beam.Create([8]) saturated_total = ( pipeline | 'Create plant counts' >> beam.Create([ ('🥕', [3, 2]), ('🍆', [1]), ('🍅', [4, 5, 3]), ]) | 'Saturated sum' >> beam.CombineValues( lambda values, max_value: min(sum(values), max_value), max_value=beam.pvalue.AsSingleton(max_value)) | beam.Map(print)) # [END combinevalues_side_inputs_singleton] if test: test(saturated_total) if __name__ == '__main__': combinevalues_side_inputs_singleton() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combinevalues_simple.py0000644000175100017510000000320015206113032032577 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CombineValuesSimple # description: Demonstration of CombineValue transform usage. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - pairs # - combine def combinevalues_simple(test=None): # [START combinevalues_simple] import apache_beam as beam with beam.Pipeline() as pipeline: total = ( pipeline | 'Create produce counts' >> beam.Create([ ('🥕', [3, 2]), ('🍆', [1]), ('🍅', [4, 5, 3]), ]) | 'Sum' >> beam.CombineValues(sum) | beam.Map(print)) # [END combinevalues_simple] if test: test(total) if __name__ == '__main__': combinevalues_simple() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/combinevalues_test.py0000644000175100017510000001006115206113032032270 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import combinevalues_combinefn from . import combinevalues_function from . import combinevalues_lambda from . import combinevalues_multiple_arguments from . import combinevalues_side_inputs_dict from . import combinevalues_side_inputs_iter from . import combinevalues_side_inputs_singleton def check_total(actual): expected = '''[START total] ('🥕', 5) ('🍆', 1) ('🍅', 12) [END total]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_saturated_total(actual): expected = '''[START saturated_total] ('🥕', 5) ('🍆', 1) ('🍅', 8) [END saturated_total]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_bounded_total(actual): expected = '''[START bounded_total] ('🥕', 5) ('🍆', 2) ('🍅', 8) [END bounded_total]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_percentages_per_season(actual): expected = '''[START percentages_per_season] ('spring', {'🥕': 0.4, '🍅': 0.4, '🍆': 0.2}) ('summer', {'🥕': 0.2, '🍅': 0.6, '🌽': 0.2}) ('fall', {'🥕': 0.5, '🍅': 0.5}) ('winter', {'🍆': 1.0}) [END percentages_per_season]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combinevalues_function.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combinevalues_lambda.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combinevalues_multiple_arguments.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combinevalues_side_inputs_singleton.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combinevalues_side_inputs_iter.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combinevalues_side_inputs_dict.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.combinevalues_combinefn.print', str) class CombineValuesTest(unittest.TestCase): def test_combinevalues_function(self): combinevalues_function.combinevalues_function(check_saturated_total) def test_combinevalues_lambda(self): combinevalues_lambda.combinevalues_lambda(check_saturated_total) def test_combinevalues_multiple_arguments(self): combinevalues_multiple_arguments.combinevalues_multiple_arguments( check_saturated_total) def test_combinevalues_side_inputs_singleton(self): combinevalues_side_inputs_singleton.combinevalues_side_inputs_singleton( check_saturated_total) def test_combinevalues_side_inputs_iter(self): combinevalues_side_inputs_iter.combinevalues_side_inputs_iter( check_bounded_total) def test_combinevalues_side_inputs_dict(self): combinevalues_side_inputs_dict.combinevalues_side_inputs_dict( check_bounded_total) def test_combinevalues_combinefn(self): combinevalues_combinefn.combinevalues_combinefn( check_percentages_per_season) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/count_globally.py0000644000175100017510000000314615206113032031420 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CountGlobally # description: Demonstration of Count transform usage. # multifile: false # default_example: false # context_line: 38 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings def count_globally(test=None): # [START count_globally] import apache_beam as beam with beam.Pipeline() as pipeline: total_elements = ( pipeline | 'Create plants' >> beam.Create( ['🍓', '🥕', '🥕', '🥕', '🍆', '🍆', '🍅', '🍅', '🍅', '🌽']) | 'Count all elements' >> beam.combiners.Count.Globally() | beam.Map(print)) # [END count_globally] if test: test(total_elements) if __name__ == '__main__': count_globally() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/count_per_element.py0000644000175100017510000000325715206113032032115 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CountPerElement # description: Demonstration of Count transform usage to count the number of unique elements. # multifile: false # default_example: false # context_line: 38 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings def count_per_element(test=None): # [START count_per_element] import apache_beam as beam with beam.Pipeline() as pipeline: total_unique_elements = ( pipeline | 'Create produce' >> beam.Create( ['🍓', '🥕', '🥕', '🥕', '🍆', '🍆', '🍅', '🍅', '🍅', '🌽']) | 'Count unique elements' >> beam.combiners.Count.PerElement() | beam.Map(print)) # [END count_per_element] if test: test(total_unique_elements) if __name__ == '__main__': count_per_element() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/count_per_key.py0000644000175100017510000000356615206113032031257 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: CountPerKey # description: Demonstration of Count transform usage to count elements per key. # multifile: false # default_example: false # context_line: 38 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - pairs def count_per_key(test=None): # [START count_per_key] import apache_beam as beam with beam.Pipeline() as pipeline: total_elements_per_keys = ( pipeline | 'Create plants' >> beam.Create([ ('spring', '🍓'), ('spring', '🥕'), ('summer', '🥕'), ('fall', '🥕'), ('spring', '🍆'), ('winter', '🍆'), ('spring', '🍅'), ('summer', '🍅'), ('fall', '🍅'), ('summer', '🌽'), ]) | 'Count elements per key' >> beam.combiners.Count.PerKey() | beam.Map(print)) # [END count_per_key] if test: test(total_elements_per_keys) if __name__ == '__main__': count_per_key() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/count_test.py0000644000175100017510000000465615206113032030601 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import count_globally from . import count_per_element from . import count_per_key def check_total_elements(actual): expected = '''[START total_elements] 10 [END total_elements]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_total_elements_per_key(actual): expected = '''[START total_elements_per_key] ('spring', 4) ('summer', 3) ('fall', 2) ('winter', 1) [END total_elements_per_key]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_total_unique_elements(actual): expected = '''[START total_unique_elements] ('🍓', 1) ('🥕', 3) ('🍆', 2) ('🍅', 3) ('🌽', 1) [END total_unique_elements]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.count_globally.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.count_per_key.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.count_per_element.print', str) class CountTest(unittest.TestCase): def test_count_globally(self): count_globally.count_globally(check_total_elements) def test_count_per_key(self): count_per_key.count_per_key(check_total_elements_per_key) def test_count_per_element(self): count_per_element.count_per_element(check_total_unique_elements) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/distinct.py0000644000175100017510000000317015206113032030221 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: Distinct # description: Demonstration of Distinct transform usage. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - numbers # - distinct def distinct(test=None): # [START distinct] import apache_beam as beam with beam.Pipeline() as pipeline: unique_elements = ( pipeline | 'Create produce' >> beam.Create([ '🥕', '🥕', '🍆', '🍅', '🍅', '🍅', ]) | 'Deduplicate elements' >> beam.Distinct() | beam.Map(print)) # [END distinct] if test: test(unique_elements) if __name__ == '__main__': distinct() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/distinct_test.py0000644000175100017510000000271415206113032031263 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import distinct def check_unique_elements(actual): expected = '''[START unique_elements] 🥕 🍆 🍅 [END unique_elements]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.distinct.print', str) class DistinctTest(unittest.TestCase): def test_distinct(self): distinct.distinct(check_unique_elements) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/groupby_attr.py0000644000175100017510000000400615206113032031120 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # Wrapping hurts the readability of the docs. # pylint: disable=line-too-long # beam-playground: # name: GroupByAttr # description: Demonstration of GroupBy transform usage by an attribute. # multifile: false # default_example: false # context_line: 51 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - group import apache_beam as beam # [START groupby_table] GROCERY_LIST = [ beam.Row(recipe='pie', fruit='raspberry', quantity=1, unit_price=3.50), beam.Row(recipe='pie', fruit='blackberry', quantity=1, unit_price=4.00), beam.Row(recipe='pie', fruit='blueberry', quantity=1, unit_price=2.00), beam.Row(recipe='muffin', fruit='blueberry', quantity=2, unit_price=2.00), beam.Row(recipe='muffin', fruit='banana', quantity=3, unit_price=1.00), beam.Row(recipe='pie', fruit='strawberry', quantity=3, unit_price=1.50), ] # [END groupby_table] def groupby_attr(test=None): with beam.Pipeline() as p: # [START groupby_attr] grouped = (p | beam.Create(GROCERY_LIST) | beam.GroupBy('recipe')) # [END groupby_attr] if test: test(grouped) else: _ = grouped | beam.Map(print) if __name__ == '__main__': groupby_attr() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/groupby_attr_expr.py0000644000175100017510000000416215206113032032161 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # Wrapping hurts the readability of the docs. # pylint: disable=line-too-long # beam-playground: # name: GroupByAttrExpr # description: Demonstration of GroupBy transform usage by an attribute using an expression. # multifile: false # default_example: false # context_line: 51 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - group import apache_beam as beam # [START groupby_table] GROCERY_LIST = [ beam.Row(recipe='pie', fruit='raspberry', quantity=1, unit_price=3.50), beam.Row(recipe='pie', fruit='blackberry', quantity=1, unit_price=4.00), beam.Row(recipe='pie', fruit='blueberry', quantity=1, unit_price=2.00), beam.Row(recipe='muffin', fruit='blueberry', quantity=2, unit_price=2.00), beam.Row(recipe='muffin', fruit='banana', quantity=3, unit_price=1.00), beam.Row(recipe='pie', fruit='strawberry', quantity=3, unit_price=1.50), ] # [END groupby_table] def groupby_attr_expr(test=None): with beam.Pipeline() as p: # [START groupby_attr_expr] grouped = ( p | beam.Create(GROCERY_LIST) | beam.GroupBy('recipe', is_berry=lambda x: 'berry' in x.fruit)) # [END groupby_attr_expr] if test: test(grouped) else: _ = grouped | beam.Map(print) if __name__ == '__main__': groupby_attr_expr() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/groupby_expr.py0000644000175100017510000000313215206113032031123 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # Wrapping hurts the readability of the docs. # pylint: disable=line-too-long # beam-playground: # name: GroupByExpr # description: Demonstration of GroupBy transform usage using an expression. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - group import apache_beam as beam def groupby_expr(test=None): with beam.Pipeline() as p: # [START groupby_expr] grouped = ( p | beam.Create( ['strawberry', 'raspberry', 'blueberry', 'blackberry', 'banana']) | beam.GroupBy(lambda s: s[0])) # [END groupby_expr] if test: test(grouped) else: _ = grouped | beam.Map(print) if __name__ == '__main__': groupby_expr() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/groupby_expr_aggregate.py0000644000175100017510000000432715206113032033140 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # Wrapping hurts the readability of the docs. # pylint: disable=line-too-long # beam-playground: # name: GroupByExprAggregate # description: Demonstration of GroupBy transform usage using an expression with aggregate. # multifile: false # default_example: false # context_line: 51 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - group import apache_beam as beam # [START groupby_table] GROCERY_LIST = [ beam.Row(recipe='pie', fruit='raspberry', quantity=1, unit_price=3.50), beam.Row(recipe='pie', fruit='blackberry', quantity=1, unit_price=4.00), beam.Row(recipe='pie', fruit='blueberry', quantity=1, unit_price=2.00), beam.Row(recipe='muffin', fruit='blueberry', quantity=2, unit_price=2.00), beam.Row(recipe='muffin', fruit='banana', quantity=3, unit_price=1.00), beam.Row(recipe='pie', fruit='strawberry', quantity=3, unit_price=1.50), ] # [END groupby_table] def expr_aggregate(test=None): with beam.Pipeline() as p: # [START expr_aggregate] grouped = ( p | beam.Create(GROCERY_LIST) | beam.GroupBy('recipe').aggregate_field( 'quantity', sum, 'total_quantity').aggregate_field( lambda x: x.quantity * x.unit_price, sum, 'price')) # [END expr_aggregate] if test: test(grouped) else: _ = grouped | beam.Map(print) if __name__ == '__main__': expr_aggregate() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/groupby_global_aggregate.py0000644000175100017510000000452615206113032033423 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # Wrapping hurts the readability of the docs. # pylint: disable=line-too-long # beam-playground: # name: GroupByGlobalAggregate # description: Demonstration of GroupBy transform usage using an expression with global aggregate. # multifile: false # default_example: false # context_line: 52 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - group import apache_beam as beam from apache_beam.transforms.combiners import MeanCombineFn # [START groupby_table] GROCERY_LIST = [ beam.Row(recipe='pie', fruit='raspberry', quantity=1, unit_price=3.50), beam.Row(recipe='pie', fruit='blackberry', quantity=1, unit_price=4.00), beam.Row(recipe='pie', fruit='blueberry', quantity=1, unit_price=2.00), beam.Row(recipe='muffin', fruit='blueberry', quantity=2, unit_price=2.00), beam.Row(recipe='muffin', fruit='banana', quantity=3, unit_price=1.00), beam.Row(recipe='pie', fruit='strawberry', quantity=3, unit_price=1.50), ] # [END groupby_table] def global_aggregate(test=None): with beam.Pipeline() as p: # [START global_aggregate] grouped = ( p | beam.Create(GROCERY_LIST) | beam.GroupBy().aggregate_field( 'unit_price', min, 'min_price').aggregate_field( 'unit_price', MeanCombineFn(), 'mean_price').aggregate_field( 'unit_price', max, 'max_price')) # [END global_aggregate] if test: test(grouped) else: _ = grouped | beam.Map(print) if __name__ == '__main__': global_aggregate() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/groupby_simple_aggregate.py0000644000175100017510000000412315206113032033445 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # Wrapping hurts the readability of the docs. # pylint: disable=line-too-long # beam-playground: # name: GroupBySimpleAggregate # description: Demonstration of GroupBy transform usage with a simple aggregate. # multifile: false # default_example: false # context_line: 51 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - group import apache_beam as beam # [START groupby_table] GROCERY_LIST = [ beam.Row(recipe='pie', fruit='raspberry', quantity=1, unit_price=3.50), beam.Row(recipe='pie', fruit='blackberry', quantity=1, unit_price=4.00), beam.Row(recipe='pie', fruit='blueberry', quantity=1, unit_price=2.00), beam.Row(recipe='muffin', fruit='blueberry', quantity=2, unit_price=2.00), beam.Row(recipe='muffin', fruit='banana', quantity=3, unit_price=1.00), beam.Row(recipe='pie', fruit='strawberry', quantity=3, unit_price=1.50), ] # [END groupby_table] def simple_aggregate(test=None): with beam.Pipeline() as p: # [START simple_aggregate] grouped = ( p | beam.Create(GROCERY_LIST) | beam.GroupBy('fruit').aggregate_field( 'quantity', sum, 'total_quantity')) # [END simple_aggregate] if test: test(grouped) if __name__ == '__main__': simple_aggregate() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/groupby_test.py0000644000175100017510000002061715206113032031133 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # Wrapping hurts the readability of the docs. # pylint: disable=line-too-long import typing import unittest import mock import apache_beam as beam from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from .groupby_attr import groupby_attr from .groupby_attr_expr import groupby_attr_expr from .groupby_expr import groupby_expr from .groupby_expr_aggregate import expr_aggregate from .groupby_global_aggregate import global_aggregate from .groupby_simple_aggregate import simple_aggregate from .groupby_two_exprs import groupby_two_exprs class UnorderedList(object): def __init__(self, contents): self._contents = list(contents) def __eq__(self, other): try: return sorted(self._contents) == sorted(other) except TypeError: return sorted(self._contents, key=str) == sorted(other, key=str) def __hash__(self): return hash(tuple(sorted(self._contents))) def __repr__(self): return 'UnorderedList(%r)' % self._contents def normalize(x): if isinstance(x, tuple) and hasattr(x, '_fields'): # A named tuple. return beam.Row(**dict(zip(x._fields, x))) elif isinstance(x, typing.Iterable) and not isinstance(x, (str, beam.Row)): return UnorderedList(normalize(e) for e in x) else: return x def normalize_kv(k, v): return normalize(k), normalize(v) # For documentation. NamedTuple = beam.Row def check_groupby_expr_result(grouped): assert_that( grouped | beam.MapTuple(normalize_kv), equal_to([ #[START groupby_expr_result] ('s', ['strawberry']), ('r', ['raspberry']), ('b', ['banana', 'blackberry', 'blueberry']), #[END groupby_expr_result] ])) def check_groupby_two_exprs_result(grouped): assert_that( grouped | beam.MapTuple(normalize_kv), equal_to([ #[START groupby_two_exprs_result] (NamedTuple(letter='s', is_berry=True), ['strawberry']), (NamedTuple(letter='r', is_berry=True), ['raspberry']), (NamedTuple(letter='b', is_berry=True), ['blackberry', 'blueberry']), (NamedTuple(letter='b', is_berry=False), ['banana']), #[END groupby_two_exprs_result] ])) def check_groupby_attr_result(grouped): assert_that( grouped | beam.MapTuple(normalize_kv), equal_to([ #[START groupby_attr_result] ( 'pie', [ beam.Row( recipe='pie', fruit='strawberry', quantity=3, unit_price=1.50), beam.Row( recipe='pie', fruit='raspberry', quantity=1, unit_price=3.50), beam.Row( recipe='pie', fruit='blackberry', quantity=1, unit_price=4.00), beam.Row( recipe='pie', fruit='blueberry', quantity=1, unit_price=2.00), ]), ( 'muffin', [ beam.Row( recipe='muffin', fruit='blueberry', quantity=2, unit_price=2.00), beam.Row( recipe='muffin', fruit='banana', quantity=3, unit_price=1.00), ]), #[END groupby_attr_result] ])) def check_groupby_attr_expr_result(grouped): assert_that( grouped | beam.MapTuple(normalize_kv), equal_to([ #[START groupby_attr_expr_result] ( NamedTuple(recipe='pie', is_berry=True), [ beam.Row( recipe='pie', fruit='strawberry', quantity=3, unit_price=1.50), beam.Row( recipe='pie', fruit='raspberry', quantity=1, unit_price=3.50), beam.Row( recipe='pie', fruit='blackberry', quantity=1, unit_price=4.00), beam.Row( recipe='pie', fruit='blueberry', quantity=1, unit_price=2.00), ]), ( NamedTuple(recipe='muffin', is_berry=True), [ beam.Row( recipe='muffin', fruit='blueberry', quantity=2, unit_price=2.00), ]), ( NamedTuple(recipe='muffin', is_berry=False), [ beam.Row( recipe='muffin', fruit='banana', quantity=3, unit_price=1.00), ]), #[END groupby_attr_expr_result] ])) def check_simple_aggregate_result(grouped): assert_that( grouped, equal_to([ #[START simple_aggregate_result] NamedTuple(fruit='strawberry', total_quantity=3), NamedTuple(fruit='raspberry', total_quantity=1), NamedTuple(fruit='blackberry', total_quantity=1), NamedTuple(fruit='blueberry', total_quantity=3), NamedTuple(fruit='banana', total_quantity=3), #[END simple_aggregate_result] ])) def check_expr_aggregate_result(grouped): assert_that( grouped | beam.Map(normalize), equal_to([ #[START expr_aggregate_result] NamedTuple(recipe='pie', total_quantity=6, price=14.00), NamedTuple(recipe='muffin', total_quantity=5, price=7.00), #[END expr_aggregate_result] ])) def check_global_aggregate_result(grouped): assert_that( grouped | beam.Map(normalize), equal_to([ #[START global_aggregate_result] NamedTuple(min_price=1.00, mean_price=7 / 3, max_price=4.00), #[END global_aggregate_result] ])) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.groupby_expr.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.groupby_two_exprs.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.groupby_attr.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.groupby_attr_expr.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.groupby_simple_aggregate.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.groupby_expr_aggregate.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.groupby_global_aggregate.print', str) class GroupByTest(unittest.TestCase): def test_groupby_expr(self): groupby_expr(check_groupby_expr_result) def test_groupby_two_exprs(self): groupby_two_exprs(check_groupby_two_exprs_result) def test_group_by_attr(self): groupby_attr(check_groupby_attr_result) def test_group_by_attr_expr(self): groupby_attr_expr(check_groupby_attr_expr_result) def test_simple_aggregate(self): simple_aggregate(check_simple_aggregate_result) def test_expr_aggregate(self): expr_aggregate(check_expr_aggregate_result) def test_global_aggregate(self): global_aggregate(check_global_aggregate_result) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/groupby_two_exprs.py0000644000175100017510000000323415206113032032202 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # Wrapping hurts the readability of the docs. # pylint: disable=line-too-long # beam-playground: # name: GroupByTwoExprs # description: Demonstration of GroupBy transform usage with two expressions. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - group import apache_beam as beam def groupby_two_exprs(test=None): with beam.Pipeline() as p: # [START groupby_two_exprs] grouped = ( p | beam.Create( ['strawberry', 'raspberry', 'blueberry', 'blackberry', 'banana']) | beam.GroupBy(letter=lambda s: s[0], is_berry=lambda s: 'berry' in s)) # [END groupby_two_exprs] if test: test(grouped) else: _ = grouped | beam.Map(print) if __name__ == '__main__': groupby_two_exprs() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/groupbykey.py0000644000175100017510000000370715206113032030606 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: GroupByKeySort # description: Demonstration of GroupByKey transform usage with per-key sorting. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - pairs # - group def groupbykey(test=None): # [START groupbykey] import apache_beam as beam with beam.Pipeline() as pipeline: produce_per_season = ( pipeline | 'Create produce list' >> beam.Create([ ('spring', '🍓'), ('spring', '🥕'), ('spring', '🍆'), ('spring', '🍅'), ('summer', '🥕'), ('summer', '🍅'), ('summer', '🌽'), ('fall', '🥕'), ('fall', '🍅'), ('winter', '🍆'), ]) | 'Group produce per season' >> beam.GroupByKey() | beam.MapTuple(lambda k, vs: (k, sorted(vs))) # sort and format | beam.Map(print)) # [END groupbykey] if test: test(produce_per_season) if __name__ == '__main__': groupbykey() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/groupbykey_test.py0000644000175100017510000000327415206113032031644 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import groupbykey def check_produce_counts(actual): expected = '''[START produce_counts] ('spring', ['🍓', '🥕', '🍆', '🍅']) ('summer', ['🥕', '🍅', '🌽']) ('fall', ['🥕', '🍅']) ('winter', ['🍆']) [END produce_counts]'''.splitlines()[1:-1] # The elements order is non-deterministic, so sort them first. assert_matches_stdout( actual, expected, lambda pair: (pair[0], sorted(pair[1]))) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.groupbykey.print', str) class GroupByKeyTest(unittest.TestCase): def test_groupbykey(self): groupbykey.groupbykey(check_produce_counts) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/groupintobatches.py0000644000175100017510000000356415206113032031767 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: GroupIntoBatches # description: Demonstration of GroupIntoBatches transform usage. # multifile: false # default_example: false # context_line: 39 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - group def groupintobatches(test=None): # [START groupintobatches] import apache_beam as beam with beam.Pipeline() as pipeline: batches_with_keys = ( pipeline | 'Create produce' >> beam.Create([ ('spring', '🍓'), ('spring', '🥕'), ('spring', '🍆'), ('spring', '🍅'), ('summer', '🥕'), ('summer', '🍅'), ('summer', '🌽'), ('fall', '🥕'), ('fall', '🍅'), ('winter', '🍆'), ]) | 'Group into batches' >> beam.GroupIntoBatches(3) | beam.Map(print)) # [END groupintobatches] if test: test(batches_with_keys) if __name__ == '__main__': groupintobatches() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/groupintobatches_test.py0000644000175100017510000000334215206113032033020 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import groupintobatches def check_batches_with_keys(actual): expected = '''[START batches_with_keys] ('spring', ['🍓', '🥕', '🍆']) ('summer', ['🥕', '🍅', '🌽']) ('spring', ['🍅']) ('fall', ['🥕', '🍅']) ('winter', ['🍆']) [END batches_with_keys]'''.splitlines()[1:-1] assert_matches_stdout( actual, expected, lambda batch: (batch[0], len(batch[1]))) @mock.patch('apache_beam.Pipeline', TestPipeline) # pylint: disable=line-too-long @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.groupintobatches.print', str) # pylint: enable=line-too-long class GroupIntoBatchesTest(unittest.TestCase): def test_groupintobatches(self): groupintobatches.groupintobatches(check_batches_with_keys) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/latest_globally.py0000644000175100017510000000444115206113032031563 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: LatestGlobally # description: Demonstration of Latest transform usage. # multifile: false # default_example: false # context_line: 45 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - timestamps # - latest def latest_globally(test=None): # [START latest_globally] import apache_beam as beam import time def to_unix_time(time_str, format='%Y-%m-%d %H:%M:%S'): return time.mktime(time.strptime(time_str, format)) with beam.Pipeline() as pipeline: latest_element = ( pipeline | 'Create crops' >> beam.Create([ { 'item': '🥬', 'harvest': '2020-02-24 00:00:00' }, { 'item': '🍓', 'harvest': '2020-06-16 00:00:00' }, { 'item': '🥕', 'harvest': '2020-07-17 00:00:00' }, { 'item': '🍆', 'harvest': '2020-10-26 00:00:00' }, { 'item': '🍅', 'harvest': '2020-10-01 00:00:00' }, ]) | 'With timestamps' >> beam.Map( lambda crop: beam.window.TimestampedValue( crop['item'], to_unix_time(crop['harvest']))) | 'Get latest element' >> beam.combiners.Latest.Globally() | beam.Map(print)) # [END latest_globally] if test: test(latest_element) if __name__ == '__main__': latest_globally() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/latest_per_key.py0000644000175100017510000000600515206113032031412 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: LatestPerKey # description: Demonstration of Latest transform usage per key. # multifile: false # default_example: false # context_line: 45 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - timestamps # - latest def latest_per_key(test=None): # [START latest_per_key] import apache_beam as beam import time def to_unix_time(time_str, format='%Y-%m-%d %H:%M:%S'): return time.mktime(time.strptime(time_str, format)) with beam.Pipeline() as pipeline: latest_elements_per_key = ( pipeline | 'Create crops' >> beam.Create([ ('spring', { 'item': '🥕', 'harvest': '2020-06-28 00:00:00' }), ('spring', { 'item': '🍓', 'harvest': '2020-06-16 00:00:00' }), ('summer', { 'item': '🥕', 'harvest': '2020-07-17 00:00:00' }), ('summer', { 'item': '🍓', 'harvest': '2020-08-26 00:00:00' }), ('summer', { 'item': '🍆', 'harvest': '2020-09-04 00:00:00' }), ('summer', { 'item': '🥬', 'harvest': '2020-09-18 00:00:00' }), ('summer', { 'item': '🍅', 'harvest': '2020-09-22 00:00:00' }), ('autumn', { 'item': '🍅', 'harvest': '2020-10-01 00:00:00' }), ('autumn', { 'item': '🥬', 'harvest': '2020-10-20 00:00:00' }), ('autumn', { 'item': '🍆', 'harvest': '2020-10-26 00:00:00' }), ('winter', { 'item': '🥬', 'harvest': '2020-02-24 00:00:00' }), ]) | 'With timestamps' >> beam.Map( lambda pair: beam.window.TimestampedValue( (pair[0], pair[1]['item']), to_unix_time(pair[1]['harvest']))) | 'Get latest elements per key' >> beam.combiners.Latest.PerKey() | beam.Map(print)) # [END latest_per_key] if test: test(latest_elements_per_key) if __name__ == '__main__': latest_per_key() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/latest_test.py0000644000175100017510000000376015206113032030740 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import latest_globally from . import latest_per_key def check_latest_element(actual): expected = '''[START latest_element] 🍆 [END latest_element]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_latest_elements_per_key(actual): expected = '''[START latest_elements_per_key] ('spring', '🥕') ('summer', '🍅') ('autumn', '🍆') ('winter', '🥬') [END latest_elements_per_key]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.latest_globally.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.latest_per_key.print', str) class LatestTest(unittest.TestCase): def test_latest_globally(self): latest_globally.latest_globally(check_latest_element) def test_latest_per_key(self): latest_per_key.latest_per_key(check_latest_elements_per_key) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/max_globally.py0000644000175100017510000000304615206113032031054 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MaxGlobally # description: Demonstration of Max transform usage. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - numbers def max_globally(test=None): # [START max_globally] import apache_beam as beam with beam.Pipeline() as pipeline: max_element = ( pipeline | 'Create numbers' >> beam.Create([3, 4, 1, 2]) | 'Get max value' >> beam.CombineGlobally(lambda elements: max(elements or [None])) | beam.Map(print)) # [END max_globally] if test: test(max_element) if __name__ == '__main__': max_globally() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/max_per_key.py0000644000175100017510000000327615206113032030712 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MaxPerKey # description: Demonstration of Max transform usage per key. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - numbers def max_per_key(test=None): # [START max_per_key] import apache_beam as beam with beam.Pipeline() as pipeline: elements_with_max_value_per_key = ( pipeline | 'Create produce' >> beam.Create([ ('🥕', 3), ('🥕', 2), ('🍆', 1), ('🍅', 4), ('🍅', 5), ('🍅', 3), ]) | 'Get max value per key' >> beam.CombinePerKey(max) | beam.Map(print)) # [END max_per_key] if test: test(elements_with_max_value_per_key) if __name__ == '__main__': max_per_key() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/max_test.py0000644000175100017510000000371415206113032030230 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import max_globally as beam_max_globally from . import max_per_key as beam_max_per_key def check_max_element(actual): expected = '''[START max_element] 4 [END max_element]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_elements_with_max_value_per_key(actual): expected = '''[START elements_with_max_value_per_key] ('🥕', 3) ('🍆', 1) ('🍅', 5) [END elements_with_max_value_per_key]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.max_globally.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.max_per_key.print', str) class MaxTest(unittest.TestCase): def test_max_globally(self): beam_max_globally.max_globally(check_max_element) def test_max_per_key(self): beam_max_per_key.max_per_key(check_elements_with_max_value_per_key) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/mean_globally.py0000644000175100017510000000302515206113032031204 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MeanGlobally # description: Demonstration of Mean transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - numbers # - pairs def mean_globally(test=None): # [START mean_globally] import apache_beam as beam with beam.Pipeline() as pipeline: mean_element = ( pipeline | 'Create numbers' >> beam.Create([3, 4, 1, 2]) | 'Get mean value' >> beam.combiners.Mean.Globally() | beam.Map(print)) # [END mean_globally] if test: test(mean_element) if __name__ == '__main__': mean_globally() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/mean_per_key.py0000644000175100017510000000333215206113032031036 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MeanPerKey # description: Demonstration of Mean transform usage per key. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - numbers # - pairs def mean_per_key(test=None): # [START mean_per_key] import apache_beam as beam with beam.Pipeline() as pipeline: elements_with_mean_value_per_key = ( pipeline | 'Create produce' >> beam.Create([ ('🥕', 3), ('🥕', 2), ('🍆', 1), ('🍅', 4), ('🍅', 5), ('🍅', 3), ]) | 'Get mean value per key' >> beam.combiners.Mean.PerKey() | beam.Map(print)) # [END mean_per_key] if test: test(elements_with_mean_value_per_key) if __name__ == '__main__': mean_per_key() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/mean_test.py0000644000175100017510000000366415206113032030367 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import mean_globally from . import mean_per_key def check_mean_element(actual): expected = '''[START mean_element] 2.5 [END mean_element]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_elements_with_mean_value_per_key(actual): expected = '''[START elements_with_mean_value_per_key] ('🥕', 2.5) ('🍆', 1.0) ('🍅', 4.0) [END elements_with_mean_value_per_key]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.mean_globally.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.mean_per_key.print', str) class MeanTest(unittest.TestCase): def test_mean_globally(self): mean_globally.mean_globally(check_mean_element) def test_mean_per_key(self): mean_per_key.mean_per_key(check_elements_with_mean_value_per_key) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/min_globally.py0000644000175100017510000000304415206113032031050 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MinGlobally # description: Demonstration of Min transform usage. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - numbers def min_globally(test=None): # [START min_globally] import apache_beam as beam with beam.Pipeline() as pipeline: min_element = ( pipeline | 'Create numbers' >> beam.Create([3, 4, 1, 2]) | 'Get min value' >> beam.CombineGlobally(lambda elements: min(elements or [-1])) | beam.Map(print)) # [END min_globally] if test: test(min_element) if __name__ == '__main__': min_globally() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/min_per_key.py0000644000175100017510000000327615206113032030710 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MinPerKey # description: Demonstration of Max transform usage per key. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - numbers def min_per_key(test=None): # [START min_per_key] import apache_beam as beam with beam.Pipeline() as pipeline: elements_with_min_value_per_key = ( pipeline | 'Create produce' >> beam.Create([ ('🥕', 3), ('🥕', 2), ('🍆', 1), ('🍅', 4), ('🍅', 5), ('🍅', 3), ]) | 'Get min value per key' >> beam.CombinePerKey(min) | beam.Map(print)) # [END min_per_key] if test: test(elements_with_min_value_per_key) if __name__ == '__main__': min_per_key() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/min_test.py0000644000175100017510000000366715206113032030235 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import min_globally as beam_min_globally from . import min_per_key as beam_min_per_key def check_min_element(actual): expected = '''[START min_element] 1 [END min_element]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_elements_with_min_value_per_key(actual): expected = '''[START elements_with_min_value_per_key] ('🥕', 2) ('🍆', 1) ('🍅', 3) [END elements_with_min_value_per_key]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.min_globally.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.min_per_key.print', str) class MinTest(unittest.TestCase): def test_min_globally(self): beam_min_globally.min_globally(check_min_element) def test_min_per_key(self): beam_min_per_key.min_per_key(check_elements_with_min_value_per_key) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000020500000000000010212 xustar00111 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/sample_fixed_size_globally.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/sample_fixed_size_globally.p0000644000175100017510000000336315206113032033572 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: SampleFixedSizeGlobally # description: Demonstration of Sample transform usage with fixed size. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - pairs # - group def sample_fixed_size_globally(test=None): # [START sample_fixed_size_globally] import apache_beam as beam with beam.Pipeline() as pipeline: sample = ( pipeline | 'Create produce' >> beam.Create([ '🍓 Strawberry', '🥕 Carrot', '🍆 Eggplant', '🍅 Tomato', '🥔 Potato', ]) | 'Sample N elements' >> beam.combiners.Sample.FixedSizeGlobally(3) | beam.Map(print)) # [END sample_fixed_size_globally] if test: test(sample) if __name__ == '__main__': sample_fixed_size_globally() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/sample_fixed_size_per_key.py0000644000175100017510000000366215206113032033616 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: SampleFixedSizePerKey # description: Demonstration of Sample transform usage with fixed size per key. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - pairs # - group def sample_fixed_size_per_key(test=None): # [START sample_fixed_size_per_key] import apache_beam as beam with beam.Pipeline() as pipeline: samples_per_key = ( pipeline | 'Create produce' >> beam.Create([ ('spring', '🍓'), ('spring', '🥕'), ('spring', '🍆'), ('spring', '🍅'), ('summer', '🥕'), ('summer', '🍅'), ('summer', '🌽'), ('fall', '🥕'), ('fall', '🍅'), ('winter', '🍆'), ]) | 'Samples per key' >> beam.combiners.Sample.FixedSizePerKey(3) | beam.Map(print)) # [END sample_fixed_size_per_key] if test: test(samples_per_key) if __name__ == '__main__': sample_fixed_size_per_key() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/sample_test.py0000644000175100017510000000452115206113032030721 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import sample_fixed_size_globally from . import sample_fixed_size_per_key def check_sample(actual): expected = '''[START sample] ['🥕 Carrot', '🍆 Eggplant', '🍅 Tomato'] [END sample]'''.splitlines()[1:-1] # The sampled elements are non-deterministic, so check the sample size. assert_matches_stdout(actual, expected, lambda elements: len(elements)) def check_samples_per_key(actual): expected = '''[START samples_per_key] ('spring', ['🍓', '🥕', '🍆']) ('summer', ['🥕', '🍅', '🌽']) ('fall', ['🥕', '🍅']) ('winter', ['🍆']) [END samples_per_key]'''.splitlines()[1:-1] # The sampled elements are non-deterministic, so check the sample size. assert_matches_stdout(actual, expected, lambda pair: (pair[0], len(pair[1]))) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.sample_fixed_size_globally.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.sample_fixed_size_per_key.print', str) class SampleTest(unittest.TestCase): def test_sample_fixed_size_globally(self): sample_fixed_size_globally.sample_fixed_size_globally(check_sample) def test_sample_fixed_size_per_key(self): sample_fixed_size_per_key.sample_fixed_size_per_key(check_samples_per_key) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/sum_globally.py0000644000175100017510000000277715206113032031105 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: SumGlobally # description: Demonstration of CombineGlobally transform usage with Sum. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - numbers def sum_globally(test=None): # [START sum_globally] import apache_beam as beam with beam.Pipeline() as pipeline: total = ( pipeline | 'Create numbers' >> beam.Create([3, 4, 1, 2]) | 'Sum values' >> beam.CombineGlobally(sum) | beam.Map(print)) # [END sum_globally] if test: test(total) if __name__ == '__main__': sum_globally() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/sum_per_key.py0000644000175100017510000000324415206113032030724 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: SumPerKey # description: Demonstration of CombinePerKey transform usage with Sum. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - numbers def sum_per_key(test=None): # [START sum_per_key] import apache_beam as beam with beam.Pipeline() as pipeline: totals_per_key = ( pipeline | 'Create produce' >> beam.Create([ ('🥕', 3), ('🥕', 2), ('🍆', 1), ('🍅', 4), ('🍅', 5), ('🍅', 3), ]) | 'Sum values per key' >> beam.CombinePerKey(sum) | beam.Map(print)) # [END sum_per_key] if test: test(totals_per_key) if __name__ == '__main__': sum_per_key() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/sum_test.py0000644000175100017510000000353515206113032030250 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import sum_globally as beam_sum_globally from . import sum_per_key as beam_sum_per_key def check_total(actual): expected = '''[START total] 10 [END total]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_totals_per_key(actual): expected = '''[START totals_per_key] ('🥕', 5) ('🍆', 1) ('🍅', 12) [END totals_per_key]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.sum_globally.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.sum_per_key.print', str) class SumTest(unittest.TestCase): def test_sum_globally(self): beam_sum_globally.sum_globally(check_total) def test_sum_per_key(self): beam_sum_per_key.sum_per_key(check_totals_per_key) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/tolist.py0000644000175100017510000000272415206113032027722 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: ToList # description: Demonstration of ToList transform usage. # multifile: false # default_example: false # context_line: 37 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms def tolist(test=None): # [START tolist] import apache_beam as beam with beam.Pipeline() as pipeline: listed_produce = ( pipeline | 'Create produce' >> beam.Create(['🍓', '🥕', '🍆', '🍅']) | beam.combiners.ToList() | beam.Map(print)) # [END tolist] if test: test(listed_produce) if __name__ == '__main__': tolist() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/tolist_test.py0000644000175100017510000000345615206113032030764 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest import mock import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from . import tolist def identity(x): return x @mock.patch('apache_beam.Pipeline', TestPipeline) # pylint: disable=line-too-long @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.tolist.print', identity) # pylint: enable=line-too-long class BatchElementsTest(unittest.TestCase): def test_tolist(self): def check(result): assert_that( result | 'Flatten' >> beam.FlatMap(identity) | 'Count' >> beam.combiners.Count.PerElement(), equal_to([('🍓', 1), ('🥕', 1), ('🍆', 1), ('🍅', 1)]), label='Check counts') assert_that( result | beam.Map(lambda x: isinstance(x, list)), equal_to([True]), label='Check type') tolist.tolist(check) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/top_largest.py0000644000175100017510000000303515206113032030723 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: TopLargest # description: Demonstration of Top transform usage with the largest values. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - numbers def top_largest(test=None): # [START top_largest] import apache_beam as beam with beam.Pipeline() as pipeline: largest_elements = ( pipeline | 'Create numbers' >> beam.Create([3, 4, 1, 2]) | 'Largest N values' >> beam.combiners.Top.Largest(2) | beam.Map(print)) # [END top_largest] if test: test(largest_elements) if __name__ == '__main__': top_largest() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/top_largest_per_key.py0000644000175100017510000000337615206113032032451 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: TopLargestPerKey # description: Demonstration of Top transform usage with the largest values per key. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - numbers def top_largest_per_key(test=None): # [START top_largest_per_key] import apache_beam as beam with beam.Pipeline() as pipeline: largest_elements_per_key = ( pipeline | 'Create produce' >> beam.Create([ ('🥕', 3), ('🥕', 2), ('🍆', 1), ('🍅', 4), ('🍅', 5), ('🍅', 3), ]) | 'Largest N values per key' >> beam.combiners.Top.LargestPerKey(2) | beam.Map(print)) # [END top_largest_per_key] if test: test(largest_elements_per_key) if __name__ == '__main__': top_largest_per_key() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/top_of.py0000644000175100017510000000347315206113032027674 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: TopOf # description: Demonstration of Top.Of transform usage. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings def top_of(test=None): # [START top_of] import apache_beam as beam with beam.Pipeline() as pipeline: shortest_elements = ( pipeline | 'Create produce names' >> beam.Create([ '🍓 Strawberry', '🥕 Carrot', '🍏 Green apple', '🍆 Eggplant', '🌽 Corn', ]) | 'Shortest names' >> beam.combiners.Top.Of( 2, # number of elements key=len, # optional, defaults to the element itself reverse=True, # optional, defaults to False (largest/descending) ) | beam.Map(print)) # [END top_of] if test: test(shortest_elements) if __name__ == '__main__': top_of() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/top_per_key.py0000644000175100017510000000404315206113032030720 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: TopPerKey # description: Demonstration of Top.PerKey transform usage. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings def top_per_key(test=None): # [START top_per_key] import apache_beam as beam with beam.Pipeline() as pipeline: shortest_elements_per_key = ( pipeline | 'Create produce names' >> beam.Create([ ('spring', '🥕 Carrot'), ('spring', '🍓 Strawberry'), ('summer', '🥕 Carrot'), ('summer', '🌽 Corn'), ('summer', '🍏 Green apple'), ('fall', '🥕 Carrot'), ('fall', '🍏 Green apple'), ('winter', '🍆 Eggplant'), ]) | 'Shortest names per key' >> beam.combiners.Top.PerKey( 2, # number of elements key=len, # optional, defaults to the value itself reverse=True, # optional, defaults to False (largest/descending) ) | beam.Map(print)) # [END top_per_key] if test: test(shortest_elements_per_key) if __name__ == '__main__': top_per_key() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/top_smallest.py0000644000175100017510000000304715206113032031111 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: TopSmallest # description: Demonstration of Top transform usage with the smallest values. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - numbers def top_smallest(test=None): # [START top_smallest] import apache_beam as beam with beam.Pipeline() as pipeline: smallest_elements = ( pipeline | 'Create numbers' >> beam.Create([3, 4, 1, 2]) | 'Smallest N values' >> beam.combiners.Top.Smallest(2) | beam.Map(print)) # [END top_smallest] if test: test(smallest_elements) if __name__ == '__main__': top_smallest() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/top_smallest_per_key.py0000644000175100017510000000341015206113032032621 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: TopSmallestPerKey # description: Demonstration of Top transform usage with the smallest values per key. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - numbers def top_smallest_per_key(test=None): # [START top_smallest_per_key] import apache_beam as beam with beam.Pipeline() as pipeline: smallest_elements_per_key = ( pipeline | 'Create produce' >> beam.Create([ ('🥕', 3), ('🥕', 2), ('🍆', 1), ('🍅', 4), ('🍅', 5), ('🍅', 3), ]) | 'Smallest N values per key' >> beam.combiners.Top.SmallestPerKey(2) | beam.Map(print)) # [END top_smallest_per_key] if test: test(smallest_elements_per_key) if __name__ == '__main__': top_smallest_per_key() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/aggregation/top_test.py0000644000175100017510000000741115206113032030243 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pylint:disable=line-too-long import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import top_largest from . import top_largest_per_key from . import top_of from . import top_per_key from . import top_smallest from . import top_smallest_per_key def check_largest_elements(actual): expected = '''[START largest_elements] [4, 3] [END largest_elements]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_largest_elements_per_key(actual): expected = '''[START largest_elements_per_key] ('🥕', [3, 2]) ('🍆', [1]) ('🍅', [5, 4]) [END largest_elements_per_key]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_smallest_elements(actual): expected = '''[START smallest_elements] [1, 2] [END smallest_elements]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_smallest_elements_per_key(actual): expected = '''[START smallest_elements_per_key] ('🥕', [2, 3]) ('🍆', [1]) ('🍅', [3, 4]) [END smallest_elements_per_key]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_shortest_elements(actual): expected = '''[START shortest_elements] ['🌽 Corn', '🥕 Carrot'] [END shortest_elements]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_shortest_elements_per_key(actual): expected = '''[START shortest_elements_per_key] ('spring', ['🥕 Carrot', '🍓 Strawberry']) ('summer', ['🌽 Corn', '🥕 Carrot']) ('fall', ['🥕 Carrot', '🍏 Green apple']) ('winter', ['🍆 Eggplant']) [END shortest_elements_per_key]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.top_largest.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.top_largest_per_key.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.top_smallest.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.top_smallest_per_key.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.top_of.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.aggregation.top_per_key.print', str) class TopTest(unittest.TestCase): def test_top_largest(self): top_largest.top_largest(check_largest_elements) def test_top_largest_per_key(self): top_largest_per_key.top_largest_per_key(check_largest_elements_per_key) def test_top_smallest(self): top_smallest.top_smallest(check_smallest_elements) def test_top_smallest_per_key(self): top_smallest_per_key.top_smallest_per_key(check_smallest_elements_per_key) def test_top_of(self): top_of.top_of(check_shortest_elements) def test_top_per_key(self): top_per_key.top_per_key(check_shortest_elements_per_key) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.2363927 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/0000755000175100017510000000000015206113157026067 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/__init__.py0000644000175100017510000000142015206113032030165 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/enrichment.py0000644000175100017510000003320115206113032030564 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint: disable=line-too-long def enrichment_with_bigtable(): # [START enrichment_with_bigtable] import apache_beam as beam from apache_beam.transforms.enrichment import Enrichment from apache_beam.transforms.enrichment_handlers.bigtable import BigTableEnrichmentHandler project_id = 'apache-beam-testing' instance_id = 'beam-test' table_id = 'bigtable-enrichment-test' row_key = 'product_id' data = [ beam.Row(sale_id=1, customer_id=1, product_id=1, quantity=1), beam.Row(sale_id=3, customer_id=3, product_id=2, quantity=3), beam.Row(sale_id=5, customer_id=5, product_id=4, quantity=2) ] bigtable_handler = BigTableEnrichmentHandler( project_id=project_id, instance_id=instance_id, table_id=table_id, row_key=row_key) with beam.Pipeline() as p: _ = ( p | "Create" >> beam.Create(data) | "Enrich W/ BigTable" >> Enrichment(bigtable_handler) | "Print" >> beam.Map(print)) # [END enrichment_with_bigtable] def enrichment_with_vertex_ai(): # [START enrichment_with_vertex_ai] import apache_beam as beam from apache_beam.transforms.enrichment import Enrichment from apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store \ import VertexAIFeatureStoreEnrichmentHandler project_id = 'apache-beam-testing' location = 'us-central1' api_endpoint = f"{location}-aiplatform.googleapis.com" data = [ beam.Row(user_id='2963', product_id=14235, sale_price=15.0), beam.Row(user_id='21422', product_id=11203, sale_price=12.0), beam.Row(user_id='20592', product_id=8579, sale_price=9.0), ] vertex_ai_handler = VertexAIFeatureStoreEnrichmentHandler( project=project_id, location=location, api_endpoint=api_endpoint, feature_store_name="vertexai_enrichment_example", feature_view_name="users", row_key="user_id", ) with beam.Pipeline() as p: _ = ( p | "Create" >> beam.Create(data) | "Enrich W/ Vertex AI" >> Enrichment(vertex_ai_handler) | "Print" >> beam.Map(print)) # [END enrichment_with_vertex_ai] def enrichment_with_vertex_ai_legacy(): # [START enrichment_with_vertex_ai_legacy] import apache_beam as beam from apache_beam.transforms.enrichment import Enrichment from apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store \ import VertexAIFeatureStoreLegacyEnrichmentHandler project_id = 'apache-beam-testing' location = 'us-central1' api_endpoint = f"{location}-aiplatform.googleapis.com" data = [ beam.Row(entity_id="movie_01", title='The Shawshank Redemption'), beam.Row(entity_id="movie_02", title="The Shining"), beam.Row(entity_id="movie_04", title='The Dark Knight'), ] vertex_ai_handler = VertexAIFeatureStoreLegacyEnrichmentHandler( project=project_id, location=location, api_endpoint=api_endpoint, entity_type_id='movies', feature_store_id="movie_prediction_unique", feature_ids=["title", "genres"], row_key="entity_id", ) with beam.Pipeline() as p: _ = ( p | "Create" >> beam.Create(data) | "Enrich W/ Vertex AI" >> Enrichment(vertex_ai_handler) | "Print" >> beam.Map(print)) # [END enrichment_with_vertex_ai_legacy] def enrichment_with_google_cloudsql_pg(): # [START enrichment_with_google_cloudsql_pg] import apache_beam as beam from apache_beam.transforms.enrichment import Enrichment from apache_beam.transforms.enrichment_handlers.cloudsql import ( CloudSQLEnrichmentHandler, DatabaseTypeAdapter, TableFieldsQueryConfig, CloudSQLConnectionConfig) import os database_adapter = DatabaseTypeAdapter.POSTGRESQL database_uri = os.environ.get("GOOGLE_CLOUD_SQL_DB_URI") database_user = os.environ.get("GOOGLE_CLOUD_SQL_DB_USER") database_password = os.environ.get("GOOGLE_CLOUD_SQL_DB_PASSWORD") database_id = os.environ.get("GOOGLE_CLOUD_SQL_DB_ID") table_id = os.environ.get("GOOGLE_CLOUD_SQL_DB_TABLE_ID") where_clause_template = "product_id = :pid" where_clause_fields = ["product_id"] data = [ beam.Row(product_id=1, name='A'), beam.Row(product_id=2, name='B'), beam.Row(product_id=3, name='C'), ] connection_config = CloudSQLConnectionConfig( db_adapter=database_adapter, instance_connection_uri=database_uri, user=database_user, password=database_password, db_id=database_id) query_config = TableFieldsQueryConfig( table_id=table_id, where_clause_template=where_clause_template, where_clause_fields=where_clause_fields) handler = CloudSQLEnrichmentHandler( connection_config=connection_config, table_id=table_id, query_config=query_config) with beam.Pipeline() as p: _ = ( p | "Create" >> beam.Create(data) | "Enrich W/ Google CloudSQL PostgreSQL" >> Enrichment(handler) | "Print" >> beam.Map(print)) # [END enrichment_with_google_cloudsql_pg] def enrichment_with_external_pg(): # [START enrichment_with_external_pg] import apache_beam as beam from apache_beam.transforms.enrichment import Enrichment from apache_beam.transforms.enrichment_handlers.cloudsql import ( CloudSQLEnrichmentHandler, DatabaseTypeAdapter, TableFieldsQueryConfig, ExternalSQLDBConnectionConfig) import os database_adapter = DatabaseTypeAdapter.POSTGRESQL database_host = os.environ.get("EXTERNAL_SQL_DB_HOST") database_port = int(os.environ.get("EXTERNAL_SQL_DB_PORT")) database_user = os.environ.get("EXTERNAL_SQL_DB_USER") database_password = os.environ.get("EXTERNAL_SQL_DB_PASSWORD") database_id = os.environ.get("EXTERNAL_SQL_DB_ID") table_id = os.environ.get("EXTERNAL_SQL_DB_TABLE_ID") where_clause_template = "product_id = :pid" where_clause_fields = ["product_id"] data = [ beam.Row(product_id=1, name='A'), beam.Row(product_id=2, name='B'), beam.Row(product_id=3, name='C'), ] connection_config = ExternalSQLDBConnectionConfig( db_adapter=database_adapter, host=database_host, port=database_port, user=database_user, password=database_password, db_id=database_id) query_config = TableFieldsQueryConfig( table_id=table_id, where_clause_template=where_clause_template, where_clause_fields=where_clause_fields) cloudsql_handler = CloudSQLEnrichmentHandler( connection_config=connection_config, table_id=table_id, query_config=query_config) with beam.Pipeline() as p: _ = ( p | "Create" >> beam.Create(data) | "Enrich W/ Unmanaged PostgreSQL" >> Enrichment(cloudsql_handler) | "Print" >> beam.Map(print)) # [END enrichment_with_external_pg] def enrichment_with_external_mysql(): # [START enrichment_with_external_mysql] import apache_beam as beam from apache_beam.transforms.enrichment import Enrichment from apache_beam.transforms.enrichment_handlers.cloudsql import ( CloudSQLEnrichmentHandler, DatabaseTypeAdapter, TableFieldsQueryConfig, ExternalSQLDBConnectionConfig) import os database_adapter = DatabaseTypeAdapter.MYSQL database_host = os.environ.get("EXTERNAL_SQL_DB_HOST") database_port = int(os.environ.get("EXTERNAL_SQL_DB_PORT")) database_user = os.environ.get("EXTERNAL_SQL_DB_USER") database_password = os.environ.get("EXTERNAL_SQL_DB_PASSWORD") database_id = os.environ.get("EXTERNAL_SQL_DB_ID") table_id = os.environ.get("EXTERNAL_SQL_DB_TABLE_ID") where_clause_template = "product_id = :pid" where_clause_fields = ["product_id"] data = [ beam.Row(product_id=1, name='A'), beam.Row(product_id=2, name='B'), beam.Row(product_id=3, name='C'), ] connection_config = ExternalSQLDBConnectionConfig( db_adapter=database_adapter, host=database_host, port=database_port, user=database_user, password=database_password, db_id=database_id) query_config = TableFieldsQueryConfig( table_id=table_id, where_clause_template=where_clause_template, where_clause_fields=where_clause_fields) cloudsql_handler = CloudSQLEnrichmentHandler( connection_config=connection_config, table_id=table_id, query_config=query_config) with beam.Pipeline() as p: _ = ( p | "Create" >> beam.Create(data) | "Enrich W/ Unmanaged MySQL" >> Enrichment(cloudsql_handler) | "Print" >> beam.Map(print)) # [END enrichment_with_external_mysql] def enrichment_with_external_sqlserver(): # [START enrichment_with_external_sqlserver] import apache_beam as beam from apache_beam.transforms.enrichment import Enrichment from apache_beam.transforms.enrichment_handlers.cloudsql import ( CloudSQLEnrichmentHandler, DatabaseTypeAdapter, TableFieldsQueryConfig, ExternalSQLDBConnectionConfig) import os database_adapter = DatabaseTypeAdapter.SQLSERVER database_host = os.environ.get("EXTERNAL_SQL_DB_HOST") database_port = int(os.environ.get("EXTERNAL_SQL_DB_PORT")) database_user = os.environ.get("EXTERNAL_SQL_DB_USER") database_password = os.environ.get("EXTERNAL_SQL_DB_PASSWORD") database_id = os.environ.get("EXTERNAL_SQL_DB_ID") table_id = os.environ.get("EXTERNAL_SQL_DB_TABLE_ID") where_clause_template = "product_id = :pid" where_clause_fields = ["product_id"] data = [ beam.Row(product_id=1, name='A'), beam.Row(product_id=2, name='B'), beam.Row(product_id=3, name='C'), ] connection_config = ExternalSQLDBConnectionConfig( db_adapter=database_adapter, host=database_host, port=database_port, user=database_user, password=database_password, db_id=database_id) query_config = TableFieldsQueryConfig( table_id=table_id, where_clause_template=where_clause_template, where_clause_fields=where_clause_fields) cloudsql_handler = CloudSQLEnrichmentHandler( connection_config=connection_config, table_id=table_id, query_config=query_config) with beam.Pipeline() as p: _ = ( p | "Create" >> beam.Create(data) | "Enrich W/ Unmanaged SQL Server" >> Enrichment(cloudsql_handler) | "Print" >> beam.Map(print)) # [END enrichment_with_external_sqlserver] def enrichment_with_milvus(): # [START enrichment_with_milvus] import os import apache_beam as beam from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Embedding from apache_beam.transforms.enrichment import Enrichment from apache_beam.ml.rag.enrichment.milvus_search import ( MilvusSearchEnrichmentHandler, MilvusConnectionParameters, MilvusSearchParameters, MilvusCollectionLoadParameters, VectorSearchParameters, VectorSearchMetrics) uri = os.environ.get("MILVUS_VECTOR_DB_URI") user = os.environ.get("MILVUS_VECTOR_DB_USER") password = os.environ.get("MILVUS_VECTOR_DB_PASSWORD") db_id = os.environ.get("MILVUS_VECTOR_DB_ID") token = os.environ.get("MILVUS_VECTOR_DB_TOKEN") collection_name = os.environ.get("MILVUS_VECTOR_DB_COLLECTION_NAME") data = [ Chunk( id="query1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content()) ] connection_parameters = MilvusConnectionParameters( uri, user, password, db_id, token) # The first condition (language == "en") excludes documents in other # languages. Initially, this gives us two documents. After applying the second # condition (cost < 50), only the first document returns in search results. filter_expr = 'metadata["language"] == "en" AND cost < 50' search_params = {"metric_type": VectorSearchMetrics.COSINE.value, "nprobe": 1} vector_search_params = VectorSearchParameters( anns_field="dense_embedding_cosine", limit=3, filter=filter_expr, search_params=search_params) search_parameters = MilvusSearchParameters( collection_name=collection_name, search_strategy=vector_search_params, output_fields=["id", "content", "domain", "cost", "metadata"], round_decimal=2) # The collection load parameters are optional. They provide fine-graine # control over how collections are loaded into memory. For simple use cases or # when getting started, this parameter can be omitted to use default loading # behavior. Consider using it in resource-constrained environments to optimize # memory usage and query performance. collection_load_parameters = MilvusCollectionLoadParameters() milvus_search_handler = MilvusSearchEnrichmentHandler( connection_parameters=connection_parameters, search_parameters=search_parameters, collection_load_parameters=collection_load_parameters) with beam.Pipeline() as p: _ = ( p | "Create" >> beam.Create(data) | "Enrich W/ Milvus" >> Enrichment(milvus_search_handler) | "Print" >> beam.Map(print)) # [END enrichment_with_milvus] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py0000644000175100017510000004034015206113032031625 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint: disable=line-too-long import os import unittest import uuid from collections.abc import Callable from contextlib import contextmanager from dataclasses import dataclass from io import StringIO from typing import Optional import mock import pytest from sqlalchemy.engine import Connection as DBAPIConnection # pylint: disable=unused-import try: from sqlalchemy import ( Column, Integer, VARCHAR, Engine, MetaData, create_engine) from apache_beam.examples.snippets.transforms.elementwise.enrichment import ( enrichment_with_bigtable, enrichment_with_vertex_ai_legacy) from apache_beam.examples.snippets.transforms.elementwise.enrichment import ( enrichment_with_vertex_ai, enrichment_with_google_cloudsql_pg, enrichment_with_external_pg, enrichment_with_external_mysql, enrichment_with_external_sqlserver, enrichment_with_milvus) from apache_beam.transforms.enrichment_handlers.cloudsql import ( DatabaseTypeAdapter) from apache_beam.transforms.enrichment_handlers.cloudsql_it_test import ( SQLEnrichmentTestHelper, SQLDBContainerInfo, ConnectionConfig, CloudSQLConnectionConfig, ExternalSQLDBConnectionConfig) from apache_beam.ml.rag.enrichment.milvus_search import MilvusConnectionParameters from apache_beam.ml.rag.test_utils import MilvusTestHelpers from apache_beam.ml.rag.test_utils import VectorDBContainerInfo from apache_beam.ml.rag.test_utils import MilvusTestHelpers from apache_beam.ml.rag.utils import parse_chunk_strings from apache_beam.io.requestresponse import RequestResponseIO except ImportError as e: raise unittest.SkipTest(f'Examples dependencies are not installed: {str(e)}') class TestContainerStartupError(Exception): """Raised when any test container fails to start.""" pass class TestContainerTeardownError(Exception): """Raised when any test container fails to teardown.""" pass def validate_enrichment_with_bigtable(): expected = '''[START enrichment_with_bigtable] Row(sale_id=1, customer_id=1, product_id=1, quantity=1, product={'product_id': '1', 'product_name': 'pixel 5', 'product_stock': '2'}) Row(sale_id=3, customer_id=3, product_id=2, quantity=3, product={'product_id': '2', 'product_name': 'pixel 6', 'product_stock': '4'}) Row(sale_id=5, customer_id=5, product_id=4, quantity=2, product={'product_id': '4', 'product_name': 'pixel 8', 'product_stock': '10'}) [END enrichment_with_bigtable]'''.splitlines()[1:-1] return expected def validate_enrichment_with_vertex_ai(): expected = '''[START enrichment_with_vertex_ai] Row(user_id='2963', product_id=14235, sale_price=15.0, age=12.0, state='1', gender='1', country='1') Row(user_id='21422', product_id=11203, sale_price=12.0, age=12.0, state='0', gender='0', country='0') Row(user_id='20592', product_id=8579, sale_price=9.0, age=12.0, state='2', gender='1', country='2') [END enrichment_with_vertex_ai]'''.splitlines()[1:-1] return expected def validate_enrichment_with_vertex_ai_legacy(): expected = '''[START enrichment_with_vertex_ai_legacy] Row(entity_id='movie_01', title='The Shawshank Redemption', genres='Drama') Row(entity_id='movie_02', title='The Shining', genres='Horror') Row(entity_id='movie_04', title='The Dark Knight', genres='Action') [END enrichment_with_vertex_ai_legacy]'''.splitlines()[1:-1] return expected def validate_enrichment_with_google_cloudsql_pg(): expected = '''[START enrichment_with_google_cloudsql_pg] Row(product_id=1, name='A', quantity=2, region_id=3) Row(product_id=2, name='B', quantity=3, region_id=1) Row(product_id=3, name='C', quantity=10, region_id=4) [END enrichment_with_google_cloudsql_pg]'''.splitlines()[1:-1] return expected def validate_enrichment_with_external_pg(): expected = '''[START enrichment_with_external_pg] Row(product_id=1, name='A', quantity=2, region_id=3) Row(product_id=2, name='B', quantity=3, region_id=1) Row(product_id=3, name='C', quantity=10, region_id=4) [END enrichment_with_external_pg]'''.splitlines()[1:-1] return expected def validate_enrichment_with_external_mysql(): expected = '''[START enrichment_with_external_mysql] Row(product_id=1, name='A', quantity=2, region_id=3) Row(product_id=2, name='B', quantity=3, region_id=1) Row(product_id=3, name='C', quantity=10, region_id=4) [END enrichment_with_external_mysql]'''.splitlines()[1:-1] return expected def validate_enrichment_with_external_sqlserver(): expected = '''[START enrichment_with_external_sqlserver] Row(product_id=1, name='A', quantity=2, region_id=3) Row(product_id=2, name='B', quantity=3, region_id=1) Row(product_id=3, name='C', quantity=10, region_id=4) [END enrichment_with_external_sqlserver]'''.splitlines()[1:-1] return expected def validate_enrichment_with_milvus(): expected = '''[START enrichment_with_milvus] Chunk(content=Content(text=None), id='query1', index=0, metadata={'enrichment_data': defaultdict(, {'id': [1], 'distance': [1.0], 'fields': [{'content': 'This is a test document', 'cost': 49, 'domain': 'medical', 'id': 1, 'metadata': {'language': 'en'}}]})}, embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3], sparse_embedding=None)) [END enrichment_with_milvus]'''.splitlines()[1:-1] return expected @mock.patch('sys.stdout', new_callable=StringIO) @pytest.mark.uses_testcontainer class EnrichmentTest(unittest.TestCase): def test_enrichment_with_bigtable(self, mock_stdout): enrichment_with_bigtable() output = mock_stdout.getvalue().splitlines() expected = validate_enrichment_with_bigtable() self.assertEqual(sorted(output), sorted(expected)) def test_enrichment_with_vertex_ai(self, mock_stdout): enrichment_with_vertex_ai() output = sorted(mock_stdout.getvalue().splitlines()) expected = sorted(validate_enrichment_with_vertex_ai()) for i in range(len(expected)): self.assertEqual( set(output[i][4:-1].split(',')), set(expected[i][4:-1].split(','))) def test_enrichment_with_vertex_ai_legacy(self, mock_stdout): enrichment_with_vertex_ai_legacy() output = mock_stdout.getvalue().splitlines() expected = validate_enrichment_with_vertex_ai_legacy() self.maxDiff = None self.assertEqual(sorted(output), sorted(expected)) @unittest.skipUnless( os.environ.get('ALLOYDB_PASSWORD'), "ALLOYDB_PASSWORD environment var is not provided") def test_enrichment_with_google_cloudsql_pg(self, mock_stdout): try: db_adapter = DatabaseTypeAdapter.POSTGRESQL with EnrichmentTestHelpers.sql_test_context(True, db_adapter): enrichment_with_google_cloudsql_pg() output = mock_stdout.getvalue().splitlines() expected = validate_enrichment_with_google_cloudsql_pg() self.assertEqual(sorted(output), sorted(expected)) except Exception as e: self.fail(f"Test failed with unexpected error: {e}") def test_enrichment_with_external_pg(self, mock_stdout): try: db_adapter = DatabaseTypeAdapter.POSTGRESQL with EnrichmentTestHelpers.sql_test_context(False, db_adapter): enrichment_with_external_pg() output = mock_stdout.getvalue().splitlines() expected = validate_enrichment_with_external_pg() self.assertEqual(sorted(output), sorted(expected)) except (TestContainerStartupError, TestContainerTeardownError) as e: raise unittest.SkipTest(str(e)) except Exception as e: self.fail(f"Test failed with unexpected error: {e}") def test_enrichment_with_external_mysql(self, mock_stdout): try: db_adapter = DatabaseTypeAdapter.MYSQL with EnrichmentTestHelpers.sql_test_context(False, db_adapter): enrichment_with_external_mysql() output = mock_stdout.getvalue().splitlines() expected = validate_enrichment_with_external_mysql() self.assertEqual(sorted(output), sorted(expected)) except (TestContainerStartupError, TestContainerTeardownError) as e: raise unittest.SkipTest(str(e)) except Exception as e: self.fail(f"Test failed with unexpected error: {e}") def test_enrichment_with_external_sqlserver(self, mock_stdout): try: db_adapter = DatabaseTypeAdapter.SQLSERVER with EnrichmentTestHelpers.sql_test_context(False, db_adapter): enrichment_with_external_sqlserver() output = mock_stdout.getvalue().splitlines() expected = validate_enrichment_with_external_sqlserver() self.assertEqual(sorted(output), sorted(expected)) except (TestContainerStartupError, TestContainerTeardownError) as e: raise unittest.SkipTest(str(e)) except Exception as e: self.fail(f"Test failed with unexpected error: {e}") def test_enrichment_with_milvus(self, mock_stdout): try: with EnrichmentTestHelpers.milvus_test_context(): enrichment_with_milvus() output = mock_stdout.getvalue().splitlines() expected = validate_enrichment_with_milvus() self.maxDiff = None output = parse_chunk_strings(output) expected = parse_chunk_strings(expected) MilvusTestHelpers.assert_chunks_equivalent(output, expected) except (TestContainerStartupError, TestContainerTeardownError) as e: raise unittest.SkipTest(str(e)) except Exception as e: self.fail(f"Test failed with unexpected error: {e}") @dataclass class CloudSQLEnrichmentTestDataConstruct: client_handler: Callable[[], DBAPIConnection] engine: Engine metadata: MetaData db: SQLDBContainerInfo = None class EnrichmentTestHelpers: @staticmethod @contextmanager def sql_test_context(is_cloudsql: bool, db_adapter: DatabaseTypeAdapter): result: Optional[CloudSQLEnrichmentTestDataConstruct] = None try: result = EnrichmentTestHelpers.pre_sql_enrichment_test( is_cloudsql, db_adapter) yield finally: if result: EnrichmentTestHelpers.post_sql_enrichment_test(result) @staticmethod @contextmanager def milvus_test_context(): db: Optional[VectorDBContainerInfo] = None try: db = EnrichmentTestHelpers.pre_milvus_enrichment() yield finally: if db: EnrichmentTestHelpers.post_milvus_enrichment(db) @staticmethod def pre_sql_enrichment_test( is_cloudsql: bool, db_adapter: DatabaseTypeAdapter) -> CloudSQLEnrichmentTestDataConstruct: unique_suffix = str(uuid.uuid4())[:8] table_id = f"products_{unique_suffix}" columns = [ Column("product_id", Integer, primary_key=True), Column("name", VARCHAR(255), nullable=False), Column("quantity", Integer, nullable=False), Column("region_id", Integer, nullable=False), ] table_data = [ { "product_id": 1, "name": "A", 'quantity': 2, 'region_id': 3 }, { "product_id": 2, "name": "B", 'quantity': 3, 'region_id': 1 }, { "product_id": 3, "name": "C", 'quantity': 10, 'region_id': 4 }, ] metadata = MetaData() connection_config: ConnectionConfig db = None if is_cloudsql: gcp_project_id = "apache-beam-testing" region = "us-central1" instance_name = "beam-integration-tests" instance_connection_uri = f"{gcp_project_id}:{region}:{instance_name}" db_id = "postgres" user = "postgres" password = os.getenv("ALLOYDB_PASSWORD") os.environ['GOOGLE_CLOUD_SQL_DB_URI'] = instance_connection_uri os.environ['GOOGLE_CLOUD_SQL_DB_ID'] = db_id os.environ['GOOGLE_CLOUD_SQL_DB_USER'] = user os.environ['GOOGLE_CLOUD_SQL_DB_PASSWORD'] = password os.environ['GOOGLE_CLOUD_SQL_DB_TABLE_ID'] = table_id connection_config = CloudSQLConnectionConfig( db_adapter=db_adapter, instance_connection_uri=instance_connection_uri, user=user, password=password, db_id=db_id) else: try: db = SQLEnrichmentTestHelper.start_sql_db_container(db_adapter) os.environ['EXTERNAL_SQL_DB_HOST'] = db.host os.environ['EXTERNAL_SQL_DB_PORT'] = str(db.port) os.environ['EXTERNAL_SQL_DB_ID'] = db.id os.environ['EXTERNAL_SQL_DB_USER'] = db.user os.environ['EXTERNAL_SQL_DB_PASSWORD'] = db.password os.environ['EXTERNAL_SQL_DB_TABLE_ID'] = table_id connection_config = ExternalSQLDBConnectionConfig( db_adapter=db_adapter, host=db.host, port=db.port, user=db.user, password=db.password, db_id=db.id) except Exception as e: db_name = db_adapter.value.lower() raise TestContainerStartupError( f"{db_name} container failed to start: {str(e)}") conenctor = connection_config.get_connector_handler() engine = create_engine( url=connection_config.get_db_url(), creator=conenctor) SQLEnrichmentTestHelper.create_table( table_id=table_id, engine=engine, columns=columns, table_data=table_data, metadata=metadata) result = CloudSQLEnrichmentTestDataConstruct( db=db, client_handler=conenctor, engine=engine, metadata=metadata) return result @staticmethod def post_sql_enrichment_test(res: CloudSQLEnrichmentTestDataConstruct): # Clean up the data inserted previously. res.metadata.drop_all(res.engine) res.engine.dispose(close=True) # Check if the test used a container-based external SQL database. if res.db: SQLEnrichmentTestHelper.stop_sql_db_container(res.db) os.environ.pop('EXTERNAL_SQL_DB_HOST', None) os.environ.pop('EXTERNAL_SQL_DB_PORT', None) os.environ.pop('EXTERNAL_SQL_DB_ID', None) os.environ.pop('EXTERNAL_SQL_DB_USER', None) os.environ.pop('EXTERNAL_SQL_DB_PASSWORD', None) os.environ.pop('EXTERNAL_SQL_DB_TABLE_ID', None) else: os.environ.pop('GOOGLE_CLOUD_SQL_DB_URI', None) os.environ.pop('GOOGLE_CLOUD_SQL_DB_ID', None) os.environ.pop('GOOGLE_CLOUD_SQL_DB_USER', None) os.environ.pop('GOOGLE_CLOUD_SQL_DB_PASSWORD', None) os.environ.pop('GOOGLE_CLOUD_SQL_DB_TABLE_ID', None) @staticmethod def pre_milvus_enrichment() -> VectorDBContainerInfo: try: db = MilvusTestHelpers.start_db_container() connection_params = MilvusConnectionParameters( uri=db.uri, user=db.user, password=db.password, db_id=db.id, token=db.token) collection_name = MilvusTestHelpers.initialize_db_with_data( connection_params) except Exception as e: raise TestContainerStartupError( f"Milvus container failed to start: {str(e)}") # Setup environment variables for db and collection configuration. This will # be used downstream by the milvus enrichment handler. os.environ['MILVUS_VECTOR_DB_URI'] = db.uri os.environ['MILVUS_VECTOR_DB_USER'] = db.user os.environ['MILVUS_VECTOR_DB_PASSWORD'] = db.password os.environ['MILVUS_VECTOR_DB_ID'] = db.id os.environ['MILVUS_VECTOR_DB_TOKEN'] = db.token os.environ['MILVUS_VECTOR_DB_COLLECTION_NAME'] = collection_name return db @staticmethod def post_milvus_enrichment(db: VectorDBContainerInfo): try: MilvusTestHelpers.stop_db_container(db) except Exception as e: raise TestContainerTeardownError( f"Milvus container failed to tear down: {str(e)}") os.environ.pop('MILVUS_VECTOR_DB_URI', None) os.environ.pop('MILVUS_VECTOR_DB_USER', None) os.environ.pop('MILVUS_VECTOR_DB_PASSWORD', None) os.environ.pop('MILVUS_VECTOR_DB_ID', None) os.environ.pop('MILVUS_VECTOR_DB_TOKEN', None) os.environ.pop('MILVUS_VECTOR_DB_COLLECTION_NAME', None) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/filter_function.py0000644000175100017510000000417215206113032031627 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FilterFunction # description: Demonstration of Filter transform usage with a function. # multifile: false # default_example: false # context_line: 44 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - filter def filter_function(test=None): # [START filter_function] import apache_beam as beam def is_perennial(plant): return plant['duration'] == 'perennial' with beam.Pipeline() as pipeline: perennials = ( pipeline | 'Gardening plants' >> beam.Create([ { 'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial' }, { 'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial' }, { 'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial' }, { 'icon': '🍅', 'name': 'Tomato', 'duration': 'annual' }, { 'icon': '🥔', 'name': 'Potato', 'duration': 'perennial' }, ]) | 'Filter perennials' >> beam.Filter(is_perennial) | beam.Map(print)) # [END filter_function] if test: test(perennials) if __name__ == '__main__': filter_function() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/filter_lambda.py0000644000175100017510000000413115206113032031215 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FilterLambda # description: Demonstration of Filter transform usage with a lambda function. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - filter def filter_lambda(test=None): # [START filter_lambda] import apache_beam as beam with beam.Pipeline() as pipeline: perennials = ( pipeline | 'Gardening plants' >> beam.Create([ { 'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial' }, { 'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial' }, { 'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial' }, { 'icon': '🍅', 'name': 'Tomato', 'duration': 'annual' }, { 'icon': '🥔', 'name': 'Potato', 'duration': 'perennial' }, ]) | 'Filter perennials' >> beam.Filter(lambda plant: plant['duration'] == 'perennial') | beam.Map(print)) # [END filter_lambda] if test: test(perennials) if __name__ == '__main__': filter_lambda() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/filter_multiple_arguments.py0000644000175100017510000000433615206113032033724 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FilterMultipleArguments # description: Demonstration of Filter transform usage with a lambda function with multiple arguments. # multifile: false # default_example: false # context_line: 44 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - filter def filter_multiple_arguments(test=None): # [START filter_multiple_arguments] import apache_beam as beam def has_duration(plant, duration): return plant['duration'] == duration with beam.Pipeline() as pipeline: perennials = ( pipeline | 'Gardening plants' >> beam.Create([ { 'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial' }, { 'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial' }, { 'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial' }, { 'icon': '🍅', 'name': 'Tomato', 'duration': 'annual' }, { 'icon': '🥔', 'name': 'Potato', 'duration': 'perennial' }, ]) | 'Filter perennials' >> beam.Filter(has_duration, 'perennial') | beam.Map(print)) # [END filter_multiple_arguments] if test: test(perennials) if __name__ == '__main__': filter_multiple_arguments() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_dict.py0000644000175100017510000000463015206113032033332 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FilterSideInputsDict # description: Demonstration of Filter transform usage with side inputs from a dictionary. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - filter def filter_side_inputs_dict(test=None): # [START filter_side_inputs_dict] import apache_beam as beam with beam.Pipeline() as pipeline: keep_duration = pipeline | 'Duration filters' >> beam.Create([ ('annual', False), ('biennial', False), ('perennial', True), ]) perennials = ( pipeline | 'Gardening plants' >> beam.Create([ { 'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial' }, { 'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial' }, { 'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial' }, { 'icon': '🍅', 'name': 'Tomato', 'duration': 'annual' }, { 'icon': '🥔', 'name': 'Potato', 'duration': 'perennial' }, ]) | 'Filter plants by duration' >> beam.Filter( lambda plant, keep_duration: keep_duration[plant['duration']], keep_duration=beam.pvalue.AsDict(keep_duration), ) | beam.Map(print)) # [END filter_side_inputs_dict] if test: test(perennials) if __name__ == '__main__': filter_side_inputs_dict() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_iter.py0000644000175100017510000000460415206113032033353 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FilterSideInputsIter # description: Demonstration of Filter transform usage with side inputs as an iterator. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - filter def filter_side_inputs_iter(test=None): # [START filter_side_inputs_iter] import apache_beam as beam with beam.Pipeline() as pipeline: valid_durations = pipeline | 'Valid durations' >> beam.Create([ 'annual', 'biennial', 'perennial', ]) valid_plants = ( pipeline | 'Gardening plants' >> beam.Create([ { 'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial' }, { 'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial' }, { 'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial' }, { 'icon': '🍅', 'name': 'Tomato', 'duration': 'annual' }, { 'icon': '🥔', 'name': 'Potato', 'duration': 'PERENNIAL' }, ]) | 'Filter valid plants' >> beam.Filter( lambda plant, valid_durations: plant['duration'] in valid_durations, valid_durations=beam.pvalue.AsIter(valid_durations), ) | beam.Map(print)) # [END filter_side_inputs_iter] if test: test(valid_plants) if __name__ == '__main__': filter_side_inputs_iter() ././@PaxHeader0000000000000000000000000000020700000000000010214 xustar00113 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_singleton.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_singleton0000644000175100017510000000447715206113032033773 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FilterSideInputsSingleton # description: Demonstration of Filter transform usage with side inputs as a singleton # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - filter def filter_side_inputs_singleton(test=None): # [START filter_side_inputs_singleton] import apache_beam as beam with beam.Pipeline() as pipeline: perennial = pipeline | 'Perennial' >> beam.Create(['perennial']) perennials = ( pipeline | 'Gardening plants' >> beam.Create([ { 'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial' }, { 'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial' }, { 'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial' }, { 'icon': '🍅', 'name': 'Tomato', 'duration': 'annual' }, { 'icon': '🥔', 'name': 'Potato', 'duration': 'perennial' }, ]) | 'Filter perennials' >> beam.Filter( lambda plant, duration: plant['duration'] == duration, duration=beam.pvalue.AsSingleton(perennial), ) | beam.Map(print)) # [END filter_side_inputs_singleton] if test: test(perennials) if __name__ == '__main__': filter_side_inputs_singleton() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/filter_test.py0000644000175100017510000000651115206113032030760 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import filter_function from . import filter_lambda from . import filter_multiple_arguments from . import filter_side_inputs_dict from . import filter_side_inputs_iter from . import filter_side_inputs_singleton def check_perennials(actual): expected = '''[START perennials] {'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'} {'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'} {'icon': '🥔', 'name': 'Potato', 'duration': 'perennial'} [END perennials]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_valid_plants(actual): expected = '''[START valid_plants] {'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'} {'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial'} {'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'} {'icon': '🍅', 'name': 'Tomato', 'duration': 'annual'} [END valid_plants]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.filter_function.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.filter_lambda.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.filter_multiple_arguments.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.filter_side_inputs_singleton.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.filter_side_inputs_iter.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.filter_side_inputs_dict.print', str) class FilterTest(unittest.TestCase): def test_filter_function(self): filter_function.filter_function(check_perennials) def test_filter_lambda(self): filter_lambda.filter_lambda(check_perennials) def test_filter_multiple_arguments(self): filter_multiple_arguments.filter_multiple_arguments(check_perennials) def test_filter_side_inputs_singleton(self): filter_side_inputs_singleton.filter_side_inputs_singleton(check_perennials) def test_filter_side_inputs_iter(self): filter_side_inputs_iter.filter_side_inputs_iter(check_valid_plants) def test_filter_side_inputs_dict(self): filter_side_inputs_dict.filter_side_inputs_dict(check_perennials) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/flatmap_function.py0000644000175100017510000000326215206113032031765 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FlatMapFunction # description: Demonstration of FlatMap transform usage with a function. # multifile: false # default_example: false # context_line: 44 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def flatmap_function(test=None): # [START flatmap_function] import apache_beam as beam def split_words(text): return text.split(',') with beam.Pipeline() as pipeline: plants = ( pipeline | 'Gardening plants' >> beam.Create([ '🍓Strawberry,🥕Carrot,🍆Eggplant', '🍅Tomato,🥔Potato', ]) | 'Split words' >> beam.FlatMap(split_words) | beam.Map(print)) # [END flatmap_function] if test: test(plants) if __name__ == '__main__': flatmap_function() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/flatmap_generator.py0000644000175100017510000000335515206113032032131 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FlatMapGenerator # description: Demonstration of FlatMap transform usage with a generator. # multifile: false # default_example: false # context_line: 45 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def flatmap_generator(test=None): # [START flatmap_generator] import apache_beam as beam def generate_elements(elements): for element in elements: yield element with beam.Pipeline() as pipeline: plants = ( pipeline | 'Gardening plants' >> beam.Create([ ['🍓Strawberry', '🥕Carrot', '🍆Eggplant'], ['🍅Tomato', '🥔Potato'], ]) | 'Flatten lists' >> beam.FlatMap(generate_elements) | beam.Map(print)) # [END flatmap_generator] if test: test(plants) if __name__ == '__main__': flatmap_generator() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/flatmap_lambda.py0000644000175100017510000000322715206113032031361 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FlatMapLambda # description: Demonstration of FlatMap transform usage with a lambda function. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def flatmap_lambda(test=None): # [START flatmap_lambda] import apache_beam as beam with beam.Pipeline() as pipeline: plants = ( pipeline | 'Gardening plants' >> beam.Create([ ['🍓Strawberry', '🥕Carrot', '🍆Eggplant'], ['🍅Tomato', '🥔Potato'], ]) | 'Flatten lists' >> beam.FlatMap(lambda elements: elements) | beam.Map(print)) # [END flatmap_lambda] if test: test(plants) if __name__ == '__main__': flatmap_lambda() ././@PaxHeader0000000000000000000000000000020500000000000010212 xustar00111 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/flatmap_multiple_arguments.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/flatmap_multiple_arguments.p0000644000175100017510000000344615206113032033673 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FlatMapMultipleArguments # description: Demonstration of FlatMap transform usage with a function that takes multiple arguments. # multifile: false # default_example: false # context_line: 44 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def flatmap_multiple_arguments(test=None): # [START flatmap_multiple_arguments] import apache_beam as beam def split_words(text, delimiter=None): return text.split(delimiter) with beam.Pipeline() as pipeline: plants = ( pipeline | 'Gardening plants' >> beam.Create([ '🍓Strawberry,🥕Carrot,🍆Eggplant', '🍅Tomato,🥔Potato', ]) | 'Split words' >> beam.FlatMap(split_words, delimiter=',') | beam.Map(print)) # [END flatmap_multiple_arguments] if test: test(plants) if __name__ == '__main__': flatmap_multiple_arguments() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/flatmap_nofunction.py0000644000175100017510000000316615206113032032325 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FlatMapNoFunction # description: Demonstration of FlatMap transform usage without a function. # multifile: false # default_example: false # context_line: 44 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def flatmap_nofunction(test=None): # [START flatmap_nofunction] import apache_beam as beam with beam.Pipeline() as pipeline: plants = ( pipeline | 'Gardening plants' >> beam.Create( [['🍓Strawberry', '🥕Carrot', '🍆Eggplant'], ['🍅Tomato', '🥔Potato']]) | 'Split words' >> beam.FlatMap() | beam.Map(print)) # [END flatmap_nofunction] if test: test(plants) if __name__ == '__main__': flatmap_nofunction() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/flatmap_side_inputs_dict.py0000644000175100017510000000473315206113032033475 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FlatMapSideInputsDict # description: Demonstration of FlatMap transform usage with a side input from a dictionary. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def flatmap_side_inputs_dict(test=None): # [START flatmap_side_inputs_dict] import apache_beam as beam def replace_duration_if_valid(plant, durations): if plant['duration'] in durations: plant['duration'] = durations[plant['duration']] yield plant with beam.Pipeline() as pipeline: durations = pipeline | 'Durations dict' >> beam.Create([ (0, 'annual'), (1, 'biennial'), (2, 'perennial'), ]) valid_plants = ( pipeline | 'Gardening plants' >> beam.Create([ { 'icon': '🍓', 'name': 'Strawberry', 'duration': 2 }, { 'icon': '🥕', 'name': 'Carrot', 'duration': 1 }, { 'icon': '🍆', 'name': 'Eggplant', 'duration': 2 }, { 'icon': '🍅', 'name': 'Tomato', 'duration': 0 }, { 'icon': '🥔', 'name': 'Potato', 'duration': -1 }, ]) | 'Replace duration if valid' >> beam.FlatMap( replace_duration_if_valid, durations=beam.pvalue.AsDict(durations), ) | beam.Map(print)) # [END flatmap_side_inputs_dict] if test: test(valid_plants) if __name__ == '__main__': flatmap_side_inputs_dict() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/flatmap_side_inputs_iter.py0000644000175100017510000000504515206113032033512 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FlatMapSideInputsIter # description: Demonstration of FlatMap transform usage with a side input from an iterator. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def flatmap_side_inputs_iter(test=None): # [START flatmap_side_inputs_iter] import apache_beam as beam def normalize_and_validate_durations(plant, valid_durations): plant['duration'] = plant['duration'].lower() if plant['duration'] in valid_durations: yield plant with beam.Pipeline() as pipeline: valid_durations = pipeline | 'Valid durations' >> beam.Create([ 'annual', 'biennial', 'perennial', ]) valid_plants = ( pipeline | 'Gardening plants' >> beam.Create([ { 'icon': '🍓', 'name': 'Strawberry', 'duration': 'Perennial' }, { 'icon': '🥕', 'name': 'Carrot', 'duration': 'BIENNIAL' }, { 'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial' }, { 'icon': '🍅', 'name': 'Tomato', 'duration': 'annual' }, { 'icon': '🥔', 'name': 'Potato', 'duration': 'unknown' }, ]) | 'Normalize and validate durations' >> beam.FlatMap( normalize_and_validate_durations, valid_durations=beam.pvalue.AsIter(valid_durations), ) | beam.Map(print)) # [END flatmap_side_inputs_iter] if test: test(valid_plants) if __name__ == '__main__': flatmap_side_inputs_iter() ././@PaxHeader0000000000000000000000000000021000000000000010206 xustar00114 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/flatmap_side_inputs_singleton.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/flatmap_side_inputs_singleto0000644000175100017510000000360415206113032033743 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FlatMapSideInputSingleton # description: Demonstration of FlatMap transform usage with a side input as a singleton. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def flatmap_side_inputs_singleton(test=None): # [START flatmap_side_inputs_singleton] import apache_beam as beam with beam.Pipeline() as pipeline: delimiter = pipeline | 'Create delimiter' >> beam.Create([',']) plants = ( pipeline | 'Gardening plants' >> beam.Create([ '🍓Strawberry,🥕Carrot,🍆Eggplant', '🍅Tomato,🥔Potato', ]) | 'Split words' >> beam.FlatMap( lambda text, delimiter: text.split(delimiter), delimiter=beam.pvalue.AsSingleton(delimiter), ) | beam.Map(print)) # [END flatmap_side_inputs_singleton] if test: test(plants) if __name__ == '__main__': flatmap_side_inputs_singleton() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/flatmap_simple.py0000644000175100017510000000314115206113032031425 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FlatMapSimple # description: Demonstration of FlatMap transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def flatmap_simple(test=None): # [START flatmap_simple] import apache_beam as beam with beam.Pipeline() as pipeline: plants = ( pipeline | 'Gardening plants' >> beam.Create([ '🍓Strawberry 🥕Carrot 🍆Eggplant', '🍅Tomato 🥔Potato', ]) | 'Split words' >> beam.FlatMap(str.split) | beam.Map(print)) # [END flatmap_simple] if test: test(plants) if __name__ == '__main__': flatmap_simple() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/flatmap_test.py0000644000175100017510000001010715206113032031113 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import flatmap_function from . import flatmap_generator from . import flatmap_lambda from . import flatmap_multiple_arguments from . import flatmap_nofunction from . import flatmap_side_inputs_dict from . import flatmap_side_inputs_iter from . import flatmap_side_inputs_singleton from . import flatmap_simple from . import flatmap_tuple def check_plants(actual): expected = '''[START plants] 🍓Strawberry 🥕Carrot 🍆Eggplant 🍅Tomato 🥔Potato [END plants]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_valid_plants(actual): expected = '''[START valid_plants] {'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'} {'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial'} {'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'} {'icon': '🍅', 'name': 'Tomato', 'duration': 'annual'} [END valid_plants]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.flatmap_simple.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.flatmap_function.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.flatmap_nofunction.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.flatmap_lambda.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.flatmap_generator.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.flatmap_multiple_arguments.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.flatmap_tuple.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.flatmap_side_inputs_singleton.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.flatmap_side_inputs_iter.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.flatmap_side_inputs_dict.print', str) class FlatMapTest(unittest.TestCase): def test_flatmap_simple(self): flatmap_simple.flatmap_simple(check_plants) def test_flatmap_function(self): flatmap_function.flatmap_function(check_plants) def test_flatmap_nofunction(self): flatmap_nofunction.flatmap_nofunction(check_plants) def test_flatmap_lambda(self): flatmap_lambda.flatmap_lambda(check_plants) def test_flatmap_generator(self): flatmap_generator.flatmap_generator(check_plants) def test_flatmap_multiple_arguments(self): flatmap_multiple_arguments.flatmap_multiple_arguments(check_plants) def test_flatmap_tuple(self): flatmap_tuple.flatmap_tuple(check_plants) def test_flatmap_side_inputs_singleton(self): flatmap_side_inputs_singleton.flatmap_side_inputs_singleton(check_plants) def test_flatmap_side_inputs_iter(self): flatmap_side_inputs_iter.flatmap_side_inputs_iter(check_valid_plants) def test_flatmap_side_inputs_dict(self): flatmap_side_inputs_dict.flatmap_side_inputs_dict(check_valid_plants) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/flatmap_tuple.py0000644000175100017510000000344415206113032031273 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FlatMapTuple # description: Demonstration of FlatMapTuple transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def flatmap_tuple(test=None): # [START flatmap_tuple] import apache_beam as beam def format_plant(icon, plant): if icon: yield '{}{}'.format(icon, plant) with beam.Pipeline() as pipeline: plants = ( pipeline | 'Gardening plants' >> beam.Create([ ('🍓', 'Strawberry'), ('🥕', 'Carrot'), ('🍆', 'Eggplant'), ('🍅', 'Tomato'), ('🥔', 'Potato'), (None, 'Invalid'), ]) | 'Format' >> beam.FlatMapTuple(format_plant) | beam.Map(print)) # [END flatmap_tuple] if test: test(plants) if __name__ == '__main__': flatmap_tuple() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/keys.py0000644000175100017510000000312215206113032027402 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: Keys # description: Demonstration of Keys transform usage. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - pairs def keys(test=None): # [START keys] import apache_beam as beam with beam.Pipeline() as pipeline: icons = ( pipeline | 'Garden plants' >> beam.Create([ ('🍓', 'Strawberry'), ('🥕', 'Carrot'), ('🍆', 'Eggplant'), ('🍅', 'Tomato'), ('🥔', 'Potato'), ]) | 'Keys' >> beam.Keys() | beam.Map(print)) # [END keys] if test: test(icons) if __name__ == '__main__': keys() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/keys_test.py0000644000175100017510000000262615206113032030451 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import keys def check_icons(actual): expected = '''[START icons] 🍓 🥕 🍆 🍅 🥔 [END icons]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.keys.print', str) class KeysTest(unittest.TestCase): def test_keys(self): keys.keys(check_icons) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/kvswap.py0000644000175100017510000000315415206113032027747 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: KvSwap # description: Demonstration of KvSwap transform usage. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - pairs def kvswap(test=None): # [START kvswap] import apache_beam as beam with beam.Pipeline() as pipeline: plants = ( pipeline | 'Garden plants' >> beam.Create([ ('🍓', 'Strawberry'), ('🥕', 'Carrot'), ('🍆', 'Eggplant'), ('🍅', 'Tomato'), ('🥔', 'Potato'), ]) | 'Key-Value swap' >> beam.KvSwap() | beam.Map(print)) # [END kvswap] if test: test(plants) if __name__ == '__main__': kvswap() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/kvswap_test.py0000644000175100017510000000276215206113032031012 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import kvswap def check_plants(actual): expected = '''[START plants] ('Strawberry', '🍓') ('Carrot', '🥕') ('Eggplant', '🍆') ('Tomato', '🍅') ('Potato', '🥔') [END plants]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.kvswap.print', str) class KvSwapTest(unittest.TestCase): def test_kvswap(self): kvswap.kvswap(check_plants) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/map_context.py0000644000175100017510000000476215206113032030763 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MapContext # description: Demonstration of Map transform usage with a context parameters. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def map_context(test=None): # [START map_context] import apache_beam as beam import random from contextlib import contextmanager @contextmanager def random_nonce(): """A context manager which typically would be used with a `with` statement. See https://peps.python.org/pep-0343/ """ # Setup code here runs once. # Here we just create a (shared) random value, but in practice one # could use this to set up network connections, caches, or other # re-usable state that could require cleanup. nonce = random.randrange(0, 100) # This (re-used) value is passed on every element call. yield nonce # Cleanup/closing code, if any, would go here to be called once the # processing is done. with beam.Pipeline() as pipeline: plants = ( pipeline | 'Gardening plants' >> beam.Create([ '🍓Strawberry', '🥕Carrot', '🍆Eggplant', '🍅Tomato', '🥔Potato', ], reshuffle=False) | 'Strip header' >> beam.Map( lambda text, a=beam.DoFn.SetupContextParam(random_nonce), b=beam. DoFn.BundleContextParam(random_nonce): f"{text} {a} {b}") | beam.Map(print)) # [END map_context] if test: test(plants) if __name__ == '__main__': map_context() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/map_function.py0000644000175100017510000000336515206113032031122 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MapFunction # description: Demonstration of Map transform usage with a function. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def map_function(test=None): # [START map_function] import apache_beam as beam def strip_header_and_newline(text): return text.strip('# \n') with beam.Pipeline() as pipeline: plants = ( pipeline | 'Gardening plants' >> beam.Create([ '# 🍓Strawberry\n', '# 🥕Carrot\n', '# 🍆Eggplant\n', '# 🍅Tomato\n', '# 🥔Potato\n', ]) | 'Strip header' >> beam.Map(strip_header_and_newline) | beam.Map(print)) # [END map_function] if test: test(plants) if __name__ == '__main__': map_function() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/map_lambda.py0000644000175100017510000000326415206113032030513 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MapLambda # description: Demonstration of Map transform usage with a lambda function. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def map_lambda(test=None): # [START map_lambda] import apache_beam as beam with beam.Pipeline() as pipeline: plants = ( pipeline | 'Gardening plants' >> beam.Create([ '# 🍓Strawberry\n', '# 🥕Carrot\n', '# 🍆Eggplant\n', '# 🍅Tomato\n', '# 🥔Potato\n', ]) | 'Strip header' >> beam.Map(lambda text: text.strip('# \n')) | beam.Map(print)) # [END map_lambda] if test: test(plants) if __name__ == '__main__': map_lambda() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/map_multiple_arguments.py0000644000175100017510000000347015206113032033212 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MapMultipleArguments # description: Demonstration of Map transform usage with a lambda function with multiple arguments. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def map_multiple_arguments(test=None): # [START map_multiple_arguments] import apache_beam as beam def strip(text, chars=None): return text.strip(chars) with beam.Pipeline() as pipeline: plants = ( pipeline | 'Gardening plants' >> beam.Create([ '# 🍓Strawberry\n', '# 🥕Carrot\n', '# 🍆Eggplant\n', '# 🍅Tomato\n', '# 🥔Potato\n', ]) | 'Strip header' >> beam.Map(strip, chars='# \n') | beam.Map(print)) # [END map_multiple_arguments] if test: test(plants) if __name__ == '__main__': map_multiple_arguments() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/map_side_inputs_dict.py0000644000175100017510000000456515206113032032631 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MapSideInputsDict # description: Demonstration of Map transform usage with side inputs from a dictionary. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def map_side_inputs_dict(test=None): # [START map_side_inputs_dict] import apache_beam as beam def replace_duration(plant, durations): plant['duration'] = durations[plant['duration']] return plant with beam.Pipeline() as pipeline: durations = pipeline | 'Durations' >> beam.Create([ (0, 'annual'), (1, 'biennial'), (2, 'perennial'), ]) plant_details = ( pipeline | 'Gardening plants' >> beam.Create([ { 'icon': '🍓', 'name': 'Strawberry', 'duration': 2 }, { 'icon': '🥕', 'name': 'Carrot', 'duration': 1 }, { 'icon': '🍆', 'name': 'Eggplant', 'duration': 2 }, { 'icon': '🍅', 'name': 'Tomato', 'duration': 0 }, { 'icon': '🥔', 'name': 'Potato', 'duration': 2 }, ]) | 'Replace duration' >> beam.Map( replace_duration, durations=beam.pvalue.AsDict(durations), ) | beam.Map(print)) # [END map_side_inputs_dict] if test: test(plant_details) if __name__ == '__main__': map_side_inputs_dict() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/map_side_inputs_iter.py0000644000175100017510000000361015206113032032637 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MapSideInputsIter # description: Demonstration of Map transform usage with side inputs as an iterable. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def map_side_inputs_iter(test=None): # [START map_side_inputs_iter] import apache_beam as beam with beam.Pipeline() as pipeline: chars = pipeline | 'Create chars' >> beam.Create(['#', ' ', '\n']) plants = ( pipeline | 'Gardening plants' >> beam.Create([ '# 🍓Strawberry\n', '# 🥕Carrot\n', '# 🍆Eggplant\n', '# 🍅Tomato\n', '# 🥔Potato\n', ]) | 'Strip header' >> beam.Map( lambda text, chars: text.strip(''.join(chars)), chars=beam.pvalue.AsIter(chars), ) | beam.Map(print)) # [END map_side_inputs_iter] if test: test(plants) if __name__ == '__main__': map_side_inputs_iter() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/map_side_inputs_singleton.py0000644000175100017510000000362515206113032033704 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MapSideInputsSingleton # description: Demonstration of Map transform usage with side inputs as a singleton. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def map_side_inputs_singleton(test=None): # [START map_side_inputs_singleton] import apache_beam as beam with beam.Pipeline() as pipeline: chars = pipeline | 'Create chars' >> beam.Create(['# \n']) plants = ( pipeline | 'Gardening plants' >> beam.Create([ '# 🍓Strawberry\n', '# 🥕Carrot\n', '# 🍆Eggplant\n', '# 🍅Tomato\n', '# 🥔Potato\n', ]) | 'Strip header' >> beam.Map( lambda text, chars: text.strip(chars), chars=beam.pvalue.AsSingleton(chars), ) | beam.Map(print)) # [END map_side_inputs_singleton] if test: test(plants) if __name__ == '__main__': map_side_inputs_singleton() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/map_simple.py0000644000175100017510000000322415206113032030560 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MapSimple # description: Demonstration of Map transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def map_simple(test=None): # [START map_simple] import apache_beam as beam with beam.Pipeline() as pipeline: plants = ( pipeline | 'Gardening plants' >> beam.Create([ ' 🍓Strawberry \n', ' 🥕Carrot \n', ' 🍆Eggplant \n', ' 🍅Tomato \n', ' 🥔Potato \n', ]) | 'Strip' >> beam.Map(str.strip) | beam.Map(print)) # [END map_simple] if test: test(plants) if __name__ == '__main__': map_simple() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/map_test.py0000644000175100017510000000775215206113032030260 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from . import map_context from . import map_function from . import map_lambda from . import map_multiple_arguments from . import map_side_inputs_dict from . import map_side_inputs_iter from . import map_side_inputs_singleton from . import map_simple from . import map_tuple def check_plants(actual): expected = '''[START plants] 🍓Strawberry 🥕Carrot 🍆Eggplant 🍅Tomato 🥔Potato [END plants]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_plant_details(actual): expected = '''[START plant_details] {'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'} {'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial'} {'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'} {'icon': '🍅', 'name': 'Tomato', 'duration': 'annual'} {'icon': '🥔', 'name': 'Potato', 'duration': 'perennial'} [END plant_details]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.map_simple.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.map_function.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.map_lambda.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.map_multiple_arguments.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.map_tuple.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.map_side_inputs_singleton.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.map_side_inputs_iter.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.map_side_inputs_dict.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.map_context.print', str) class MapTest(unittest.TestCase): def test_map_simple(self): map_simple.map_simple(check_plants) def test_map_function(self): map_function.map_function(check_plants) def test_map_lambda(self): map_lambda.map_lambda(check_plants) def test_map_multiple_arguments(self): map_multiple_arguments.map_multiple_arguments(check_plants) def test_map_tuple(self): map_tuple.map_tuple(check_plants) def test_map_side_inputs_singleton(self): map_side_inputs_singleton.map_side_inputs_singleton(check_plants) def test_map_side_inputs_iter(self): map_side_inputs_iter.map_side_inputs_iter(check_plants) def test_map_side_inputs_dict(self): map_side_inputs_dict.map_side_inputs_dict(check_plant_details) def test_map_context(self): import re def check_nonces(output): def shares_same_nonces(elements): s = set(re.search(r'\d+ \d+', e).group(0) for e in elements) assert len(s) == 1, s assert_that(output, shares_same_nonces) map_context.map_context(check_nonces) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/map_tuple.py0000644000175100017510000000327515206113032030426 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: MapTuple # description: Demonstration of MapTuple transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - map def map_tuple(test=None): # [START map_tuple] import apache_beam as beam with beam.Pipeline() as pipeline: plants = ( pipeline | 'Gardening plants' >> beam.Create([ ('🍓', 'Strawberry'), ('🥕', 'Carrot'), ('🍆', 'Eggplant'), ('🍅', 'Tomato'), ('🥔', 'Potato'), ]) | 'Format' >> beam.MapTuple(lambda icon, plant: '{}{}'.format(icon, plant)) | beam.Map(print)) # [END map_tuple] if test: test(plants) if __name__ == '__main__': map_tuple() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/mltransform.py0000644000175100017510000000667115206113032031007 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint: disable=reimported # pylint:disable=line-too-long def mltransform_scale_to_0_1(test=None): # [START mltransform_scale_to_0_1] import apache_beam as beam from apache_beam.ml.transforms.base import MLTransform from apache_beam.ml.transforms.tft import ScaleTo01 import tempfile data = [ { 'x': [1, 5, 3] }, { 'x': [4, 2, 8] }, ] artifact_location = tempfile.mkdtemp() scale_to_0_1_fn = ScaleTo01(columns=['x']) with beam.Pipeline() as p: transformed_data = ( p | beam.Create(data) | MLTransform(write_artifact_location=artifact_location).with_transform( scale_to_0_1_fn) | beam.Map(print)) # [END mltransform_scale_to_0_1] if test: test(transformed_data) def mltransform_compute_and_apply_vocabulary(test=None): # [START mltransform_compute_and_apply_vocabulary] import apache_beam as beam from apache_beam.ml.transforms.base import MLTransform from apache_beam.ml.transforms.tft import ComputeAndApplyVocabulary import tempfile artifact_location = tempfile.mkdtemp() data = [ { 'x': ['I', 'love', 'Beam'] }, { 'x': ['Beam', 'is', 'awesome'] }, ] compute_and_apply_vocabulary_fn = ComputeAndApplyVocabulary(columns=['x']) with beam.Pipeline() as p: transformed_data = ( p | beam.Create(data) | MLTransform(write_artifact_location=artifact_location).with_transform( compute_and_apply_vocabulary_fn) | beam.Map(print)) # [END mltransform_compute_and_apply_vocabulary] if test: test(transformed_data) def mltransform_compute_and_apply_vocabulary_with_scalar(test=None): # [START mltransform_compute_and_apply_vocabulary_with_scalar] import apache_beam as beam from apache_beam.ml.transforms.base import MLTransform from apache_beam.ml.transforms.tft import ComputeAndApplyVocabulary import tempfile data = [ { 'x': 'I' }, { 'x': 'love' }, { 'x': 'Beam' }, { 'x': 'Beam' }, { 'x': 'is' }, { 'x': 'awesome' }, ] artifact_location = tempfile.mkdtemp() compute_and_apply_vocabulary_fn = ComputeAndApplyVocabulary(columns=['x']) with beam.Pipeline() as p: transformed_data = ( p | beam.Create(data) | MLTransform(write_artifact_location=artifact_location).with_transform( compute_and_apply_vocabulary_fn) | beam.Map(print)) # [END mltransform_compute_and_apply_vocabulary_with_scalar] if test: test(transformed_data) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/mltransform_test.py0000644000175100017510000000653115206113032032041 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint: disable=ungrouped-imports import unittest from io import StringIO import mock from apache_beam.testing.test_pipeline import TestPipeline try: # fail when tft is not installed. import tensorflow_transform as tft # pylint: disable=unused-import from apache_beam.examples.snippets.transforms.elementwise.mltransform import mltransform_scale_to_0_1 from apache_beam.examples.snippets.transforms.elementwise.mltransform import mltransform_compute_and_apply_vocabulary from apache_beam.examples.snippets.transforms.elementwise.mltransform import mltransform_compute_and_apply_vocabulary_with_scalar except ImportError: raise unittest.SkipTest('tensorflow_transform is not installed.') def check_mltransform_compute_and_apply_vocab(): expected = '''[START mltransform_compute_and_apply_vocab] Row(x=array([4, 1, 0])) Row(x=array([0, 2, 3])) [END mltransform_compute_and_apply_vocab] '''.splitlines()[1:-1] return expected def check_mltransform_scale_to_0_1(): expected = '''[START mltransform_scale_to_0_1] Row(x=array([0. , 0.5714286, 0.2857143], dtype=float32)) Row(x=array([0.42857143, 0.14285715, 1. ], dtype=float32)) [END mltransform_scale_to_0_1] '''.splitlines()[1:-1] return expected def check_mltransform_compute_and_apply_vocabulary_with_scalar(): expected = '''[START mltransform_compute_and_apply_vocabulary_with_scalar] Row(x=array([4])) Row(x=array([1])) Row(x=array([0])) Row(x=array([0])) Row(x=array([2])) Row(x=array([3])) [END mltransform_compute_and_apply_vocabulary_with_scalar] '''.splitlines( )[1:-1] return expected @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch('sys.stdout', new_callable=StringIO) class MLTransformStdOutTest(unittest.TestCase): def test_mltransform_compute_and_apply_vocab(self, mock_stdout): mltransform_compute_and_apply_vocabulary() predicted = mock_stdout.getvalue().splitlines() expected = check_mltransform_compute_and_apply_vocab() self.assertEqual(predicted, expected) def test_mltransform_scale_to_0_1(self, mock_stdout): mltransform_scale_to_0_1() predicted = mock_stdout.getvalue().splitlines() expected = check_mltransform_scale_to_0_1() self.assertEqual(predicted, expected) def test_mltransform_compute_and_apply_vocab_scalar(self, mock_stdout): mltransform_compute_and_apply_vocabulary_with_scalar() predicted = mock_stdout.getvalue().splitlines() expected = check_mltransform_compute_and_apply_vocabulary_with_scalar() self.assertEqual(predicted, expected) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn.py0000644000175100017510000000342415206113032030547 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: ParDoDoFn # description: Demonstration of ParDo transform usage with DoFn. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings def pardo_dofn(test=None): # [START pardo_dofn] import apache_beam as beam class SplitWords(beam.DoFn): def __init__(self, delimiter=','): self.delimiter = delimiter def process(self, text): for word in text.split(self.delimiter): yield word with beam.Pipeline() as pipeline: plants = ( pipeline | 'Gardening plants' >> beam.Create([ '🍓Strawberry,🥕Carrot,🍆Eggplant', '🍅Tomato,🥔Potato', ]) | 'Split words' >> beam.ParDo(SplitWords(',')) | beam.Map(print)) # [END pardo_dofn] if test: test(plants) if __name__ == '__main__': pardo_dofn() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn_methods.py0000644000175100017510000000506015206113032032270 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: ParDoDoFnMethods # description: Demonstration of ParDo transform usage with a DoFn with methods. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings def pardo_dofn_methods(test=None): # Portable runners do not guarantee that teardown will be executed, so we # use FnApiRunner instead of prism. runner = 'FnApiRunner' # [START pardo_dofn_methods] import apache_beam as beam class DoFnMethods(beam.DoFn): def __init__(self): print('__init__') self.window = beam.transforms.window.GlobalWindow() def setup(self): print('setup') def start_bundle(self): print('start_bundle') def process(self, element, window=beam.DoFn.WindowParam): self.window = window yield '* process: ' + element def finish_bundle(self): yield beam.utils.windowed_value.WindowedValue( value='* finish_bundle: 🌱🌳🌍', timestamp=0, windows=[self.window], ) def teardown(self): # Teardown is best effort and not guaranteed to be executed by all # runners in all cases (for example, it may be skipped if the pipeline # can otherwise complete). It should be used for best effort resource # cleanup. print('teardown') with beam.Pipeline(runner) as pipeline: results = ( pipeline | 'Create inputs' >> beam.Create(['🍓', '🥕', '🍆', '🍅', '🥔']) | 'DoFn methods' >> beam.ParDo(DoFnMethods()) | beam.Map(print)) # [END pardo_dofn_methods] if test: return test(results) if __name__ == '__main__': pardo_dofn_methods() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn_params.py0000644000175100017510000000535715206113032032121 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: ParDoDoFnParams # description: Demonstration of ParDo transform usage with a DoFn with parameters. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings def pardo_dofn_params(test=None): # pylint: disable=line-too-long # [START pardo_dofn_params] import apache_beam as beam class AnalyzeElement(beam.DoFn): def process( self, elem, timestamp=beam.DoFn.TimestampParam, window=beam.DoFn.WindowParam): yield '\n'.join([ '# timestamp', 'type(timestamp) -> ' + repr(type(timestamp)), 'timestamp.micros -> ' + repr(timestamp.micros), 'timestamp.to_rfc3339() -> ' + repr(timestamp.to_rfc3339()), 'timestamp.to_utc_datetime() -> ' + repr(timestamp.to_utc_datetime()), '', '# window', 'type(window) -> ' + repr(type(window)), 'window.start -> {} ({})'.format( window.start, window.start.to_utc_datetime()), 'window.end -> {} ({})'.format( window.end, window.end.to_utc_datetime()), 'window.max_timestamp() -> {} ({})'.format( window.max_timestamp(), window.max_timestamp().to_utc_datetime()), ]) with beam.Pipeline() as pipeline: dofn_params = ( pipeline | 'Create a single test element' >> beam.Create([':)']) | 'Add timestamp (Spring equinox 2020)' >> beam.Map(lambda elem: beam.window.TimestampedValue(elem, 1584675660)) | 'Fixed 30sec windows' >> beam.WindowInto(beam.window.FixedWindows(30)) | 'Analyze element' >> beam.ParDo(AnalyzeElement()) | beam.Map(print)) # [END pardo_dofn_params] # pylint: enable=line-too-long if test: test(dofn_params) if __name__ == '__main__': pardo_dofn_params() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/pardo_test.py0000644000175100017510000000742015206113032030600 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long import unittest from io import StringIO import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from . import pardo_dofn from . import pardo_dofn_methods from . import pardo_dofn_params def check_plants(actual): expected = '''[START plants] 🍓Strawberry 🥕Carrot 🍆Eggplant 🍅Tomato 🥔Potato [END plants]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_dofn_params(actual): # pylint: disable=line-too-long expected = '\n'.join( '''[START dofn_params] # timestamp type(timestamp) -> timestamp.micros -> 1584675660000000 timestamp.to_rfc3339() -> '2020-03-20T03:41:00Z' timestamp.to_utc_datetime() -> datetime.datetime(2020, 3, 20, 3, 41) # window type(window) -> window.start -> Timestamp(1584675660) (2020-03-20 03:41:00) window.end -> Timestamp(1584675690) (2020-03-20 03:41:30) window.max_timestamp() -> Timestamp(1584675689.999999) (2020-03-20 03:41:29.999999) [END dofn_params]'''.splitlines()[1:-1]) # pylint: enable=line-too-long assert_that(actual, equal_to([expected])) def check_dofn_methods(actual): # Return the expected stdout to check the ordering of the called methods. return '''[START results] __init__ setup start_bundle * process: 🍓 * process: 🥕 * process: 🍆 * process: 🍅 * process: 🥔 * finish_bundle: 🌱🌳🌍 teardown [END results]'''.splitlines()[1:-1] @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.pardo_dofn.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.pardo_dofn_params.print', str) class ParDoTest(unittest.TestCase): def test_pardo_dofn(self): pardo_dofn.pardo_dofn(check_plants) def test_pardo_dofn_params(self): pardo_dofn_params.pardo_dofn_params(check_dofn_params) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch('sys.stdout', new_callable=StringIO) class ParDoStdoutTest(unittest.TestCase): def test_pardo_dofn_methods(self, mock_stdout): expected = pardo_dofn_methods.pardo_dofn_methods(check_dofn_methods) actual = mock_stdout.getvalue().splitlines() # For the stdout, check the ordering of the methods, not of the elements. actual_stdout = [line.split(':')[0] for line in actual] expected_stdout = [line.split(':')[0] for line in expected] self.assertEqual(actual_stdout, expected_stdout) # For the elements, ignore the stdout and just make sure all elements match. actual_elements = {line for line in actual if line.startswith('*')} expected_elements = {line for line in expected if line.startswith('*')} self.assertEqual(actual_elements, expected_elements) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/partition_function.py0000644000175100017510000000460015206113032032347 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: PartitionFunction # description: Demonstration of Partition transform usage with a function. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: MEDIUM # tags: # - transforms # - partitions def partition_function(test=None): # pylint:disable=expression-not-assigned # [START partition_function] import apache_beam as beam durations = ['annual', 'biennial', 'perennial'] def by_duration(plant, num_partitions): return durations.index(plant['duration']) with beam.Pipeline() as pipeline: annuals, biennials, perennials = ( pipeline | 'Gardening plants' >> beam.Create([ {'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'}, {'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial'}, {'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'}, {'icon': '🍅', 'name': 'Tomato', 'duration': 'annual'}, {'icon': '🥔', 'name': 'Potato', 'duration': 'perennial'}, ]) | 'Partition' >> beam.Partition(by_duration, len(durations)) ) annuals | 'Annuals' >> beam.Map(lambda x: print('annual: {}'.format(x))) biennials | 'Biennials' >> beam.Map( lambda x: print('biennial: {}'.format(x))) perennials | 'Perennials' >> beam.Map( lambda x: print('perennial: {}'.format(x))) # [END partition_function] if test: test(annuals, biennials, perennials) if __name__ == '__main__': partition_function() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/partition_lambda.py0000644000175100017510000000467015206113032031751 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: PartitionLambda # description: Demonstration of Partition transform usage with a lambda function. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: MEDIUM # tags: # - transforms # - partitions def partition_lambda(test=None): # pylint:disable=expression-not-assigned # [START partition_lambda] import apache_beam as beam durations = ['annual', 'biennial', 'perennial'] with beam.Pipeline() as pipeline: annuals, biennials, perennials = ( pipeline | 'Gardening plants' >> beam.Create([ {'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'}, {'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial'}, {'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'}, {'icon': '🍅', 'name': 'Tomato', 'duration': 'annual'}, {'icon': '🥔', 'name': 'Potato', 'duration': 'perennial'}, ]) | 'Partition' >> beam.Partition( lambda plant, num_partitions: durations.index(plant['duration']), len(durations), ) ) annuals | 'Annuals' >> beam.Map(lambda x: print('annual: {}'.format(x))) biennials | 'Biennials' >> beam.Map( lambda x: print('biennial: {}'.format(x))) perennials | 'Perennials' >> beam.Map( lambda x: print('perennial: {}'.format(x))) # [END partition_lambda] # pylint: enable=line-too-long, expression-not-assigned if test: test(annuals, biennials, perennials) if __name__ == '__main__': partition_lambda() ././@PaxHeader0000000000000000000000000000020700000000000010214 xustar00113 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/partition_multiple_arguments.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/partition_multiple_arguments0000644000175100017510000000506415206113032034020 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: PartitionMultipleArguments # description: Demonstration of Partition transform usage with a lambda function with multiple arguments. # multifile: false # default_example: false # context_line: 42 # categories: # - Core Transforms # complexity: MEDIUM # tags: # - transforms # - partitions def partition_multiple_arguments(test=None): # pylint: disable=expression-not-assigned # [START partition_multiple_arguments] import apache_beam as beam import json def split_dataset(plant, num_partitions, ratio): assert num_partitions == len(ratio) bucket = sum(map(ord, json.dumps(plant))) % sum(ratio) total = 0 for i, part in enumerate(ratio): total += part if bucket < total: return i return len(ratio) - 1 with beam.Pipeline() as pipeline: train_dataset, test_dataset = ( pipeline | 'Gardening plants' >> beam.Create([ {'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'}, {'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial'}, {'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'}, {'icon': '🍅', 'name': 'Tomato', 'duration': 'annual'}, {'icon': '🥔', 'name': 'Potato', 'duration': 'perennial'}, ]) | 'Partition' >> beam.Partition(split_dataset, 2, ratio=[8, 2]) ) train_dataset | 'Train' >> beam.Map(lambda x: print('train: {}'.format(x))) test_dataset | 'Test' >> beam.Map(lambda x: print('test: {}'.format(x))) # [END partition_multiple_arguments] # pylint: enable=expression-not-assigned if test: test(train_dataset, test_dataset) if __name__ == '__main__': partition_multiple_arguments() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/partition_test.py0000644000175100017510000000736615206113032031515 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import partition_function from . import partition_lambda from . import partition_multiple_arguments def check_partitions(actual1, actual2, actual3): expected = '''[START partitions] perennial: {'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'} biennial: {'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial'} perennial: {'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'} annual: {'icon': '🍅', 'name': 'Tomato', 'duration': 'annual'} perennial: {'icon': '🥔', 'name': 'Potato', 'duration': 'perennial'} [END partitions]'''.splitlines()[1:-1] annuals = [ line.split(':', 1)[1].strip() for line in expected if line.split(':', 1)[0] == 'annual' ] biennials = [ line.split(':', 1)[1].strip() for line in expected if line.split(':', 1)[0] == 'biennial' ] perennials = [ line.split(':', 1)[1].strip() for line in expected if line.split(':', 1)[0] == 'perennial' ] assert_matches_stdout(actual1, annuals, label='annuals') assert_matches_stdout(actual2, biennials, label='biennials') assert_matches_stdout(actual3, perennials, label='perennials') def check_split_datasets(actual1, actual2): expected = '''[START train_test] train: {'icon': '🍓', 'name': 'Strawberry', 'duration': 'perennial'} train: {'icon': '🥕', 'name': 'Carrot', 'duration': 'biennial'} test: {'icon': '🍆', 'name': 'Eggplant', 'duration': 'perennial'} test: {'icon': '🍅', 'name': 'Tomato', 'duration': 'annual'} train: {'icon': '🥔', 'name': 'Potato', 'duration': 'perennial'} [END train_test]'''.splitlines()[1:-1] train_dataset = [ line.split(':', 1)[1].strip() for line in expected if line.split(':', 1)[0] == 'train' ] test_dataset = [ line.split(':', 1)[1].strip() for line in expected if line.split(':', 1)[0] == 'test' ] assert_matches_stdout(actual1, train_dataset, label='train_dataset') assert_matches_stdout(actual2, test_dataset, label='test_dataset') @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.partition_function.print', lambda elem: elem) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.partition_lambda.print', lambda elem: elem) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.partition_multiple_arguments.print', lambda elem: elem) class PartitionTest(unittest.TestCase): def test_partition_function(self): partition_function.partition_function(check_partitions) def test_partition_lambda(self): partition_lambda.partition_lambda(check_partitions) def test_partition_multiple_arguments(self): partition_multiple_arguments.partition_multiple_arguments( check_split_datasets) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/regex_all_matches.py0000644000175100017510000000373315206113032032105 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: RegexAllMatches # description: Demonstration of Regex.all_matches transform usage. # multifile: false # default_example: false # context_line: 43 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - regex def regex_all_matches(test=None): # [START regex_all_matches] import apache_beam as beam # Matches a named group 'icon', and then two comma-separated groups. regex = r'(?P[^\s,]+), *(\w+), *(\w+)' with beam.Pipeline() as pipeline: plants_all_matches = ( pipeline | 'Garden plants' >> beam.Create([ '🍓, Strawberry, perennial', '🥕, Carrot, biennial ignoring trailing words', '🍆, Eggplant, perennial', '🍅, Tomato, annual', '🥔, Potato, perennial', '# 🍌, invalid, format', 'invalid, 🍉, format', ]) | 'Parse plants' >> beam.Regex.all_matches(regex) | beam.Map(print)) # [END regex_all_matches] if test: test(plants_all_matches) if __name__ == '__main__': regex_all_matches() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/regex_find.py0000644000175100017510000000362615206113032030552 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: RegexFind # description: Demonstration of Regex.find transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - regex def regex_find(test=None): # [START regex_find] import apache_beam as beam # Matches a named group 'icon', and then two comma-separated groups. regex = r'(?P[^\s,]+), *(\w+), *(\w+)' with beam.Pipeline() as pipeline: plants_matches = ( pipeline | 'Garden plants' >> beam.Create([ '# 🍓, Strawberry, perennial', '# 🥕, Carrot, biennial ignoring trailing words', '# 🍆, Eggplant, perennial - 🍌, Banana, perennial', '# 🍅, Tomato, annual - 🍉, Watermelon, annual', '# 🥔, Potato, perennial', ]) | 'Parse plants' >> beam.Regex.find(regex) | beam.Map(print)) # [END regex_find] if test: test(plants_matches) if __name__ == '__main__': regex_find() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/regex_find_all.py0000644000175100017510000000366415206113032031404 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: RegexFindAll # description: Demonstration of Regex.finda_all transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - regex def regex_find_all(test=None): # [START regex_find_all] import apache_beam as beam # Matches a named group 'icon', and then two comma-separated groups. regex = r'(?P[^\s,]+), *(\w+), *(\w+)' with beam.Pipeline() as pipeline: plants_find_all = ( pipeline | 'Garden plants' >> beam.Create([ '# 🍓, Strawberry, perennial', '# 🥕, Carrot, biennial ignoring trailing words', '# 🍆, Eggplant, perennial - 🍌, Banana, perennial', '# 🍅, Tomato, annual - 🍉, Watermelon, annual', '# 🥔, Potato, perennial', ]) | 'Parse plants' >> beam.Regex.find_all(regex) | beam.Map(print)) # [END regex_find_all] if test: test(plants_find_all) if __name__ == '__main__': regex_find_all() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/regex_find_kv.py0000644000175100017510000000370115206113032031244 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: RegexFindKV # description: Demonstration of Regex.find_kv transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - regex def regex_find_kv(test=None): # [START regex_find_kv] import apache_beam as beam # Matches a named group 'icon', and then two comma-separated groups. regex = r'(?P[^\s,]+), *(\w+), *(\w+)' with beam.Pipeline() as pipeline: plants_matches_kv = ( pipeline | 'Garden plants' >> beam.Create([ '# 🍓, Strawberry, perennial', '# 🥕, Carrot, biennial ignoring trailing words', '# 🍆, Eggplant, perennial - 🍌, Banana, perennial', '# 🍅, Tomato, annual - 🍉, Watermelon, annual', '# 🥔, Potato, perennial', ]) | 'Parse plants' >> beam.Regex.find_kv(regex, keyGroup='icon') | beam.Map(print)) # [END regex_find_kv] if test: test(plants_matches_kv) if __name__ == '__main__': regex_find_kv() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/regex_matches.py0000644000175100017510000000367015206113032031255 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: RegexMatches # description: Demonstration of Regex.matches transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - regex def regex_matches(test=None): # [START regex_matches] import apache_beam as beam # Matches a named group 'icon', and then two comma-separated groups. regex = r'(?P[^\s,]+), *(\w+), *(\w+)' with beam.Pipeline() as pipeline: plants_matches = ( pipeline | 'Garden plants' >> beam.Create([ '🍓, Strawberry, perennial', '🥕, Carrot, biennial ignoring trailing words', '🍆, Eggplant, perennial', '🍅, Tomato, annual', '🥔, Potato, perennial', '# 🍌, invalid, format', 'invalid, 🍉, format', ]) | 'Parse plants' >> beam.Regex.matches(regex) | beam.Map(print)) # [END regex_matches] if test: test(plants_matches) if __name__ == '__main__': regex_matches() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/regex_matches_kv.py0000644000175100017510000000374315206113032031756 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: RegexMatchesKV # description: Demonstration of Regex.matches_kv transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - regex def regex_matches_kv(test=None): # [START regex_matches_kv] import apache_beam as beam # Matches a named group 'icon', and then two comma-separated groups. regex = r'(?P[^\s,]+), *(\w+), *(\w+)' with beam.Pipeline() as pipeline: plants_matches_kv = ( pipeline | 'Garden plants' >> beam.Create([ '🍓, Strawberry, perennial', '🥕, Carrot, biennial ignoring trailing words', '🍆, Eggplant, perennial', '🍅, Tomato, annual', '🥔, Potato, perennial', '# 🍌, invalid, format', 'invalid, 🍉, format', ]) | 'Parse plants' >> beam.Regex.matches_kv(regex, keyGroup='icon') | beam.Map(print)) # [END regex_matches_kv] if test: test(plants_matches_kv) if __name__ == '__main__': regex_matches_kv() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/regex_replace_all.py0000644000175100017510000000342315206113032032070 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: RegexReplaceAll # description: Demonstration of Regex.replace_all transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - regex def regex_replace_all(test=None): # [START regex_replace_all] import apache_beam as beam with beam.Pipeline() as pipeline: plants_replace_all = ( pipeline | 'Garden plants' >> beam.Create([ '🍓 : Strawberry : perennial', '🥕 : Carrot : biennial', '🍆\t:\tEggplant\t:\tperennial', '🍅 : Tomato : annual', '🥔 : Potato : perennial', ]) | 'To CSV' >> beam.Regex.replace_all(r'\s*:\s*', ',') | beam.Map(print)) # [END regex_replace_all] if test: test(plants_replace_all) if __name__ == '__main__': regex_replace_all() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/regex_replace_first.py0000644000175100017510000000344015206113032032446 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: RegexReplaceFirst # description: Demonstration of Regex.replace_first transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - regex def regex_replace_first(test=None): # [START regex_replace_first] import apache_beam as beam with beam.Pipeline() as pipeline: plants_replace_first = ( pipeline | 'Garden plants' >> beam.Create([ '🍓, Strawberry, perennial', '🥕, Carrot, biennial', '🍆,\tEggplant, perennial', '🍅, Tomato, annual', '🥔, Potato, perennial', ]) | 'As dictionary' >> beam.Regex.replace_first(r'\s*,\s*', ': ') | beam.Map(print)) # [END regex_replace_first] if test: test(plants_replace_first) if __name__ == '__main__': regex_replace_first() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/regex_split.py0000644000175100017510000000333515206113032030762 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: RegexSplit # description: Demonstration of Regex.split transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - regex def regex_split(test=None): # [START regex_split] import apache_beam as beam with beam.Pipeline() as pipeline: plants_split = ( pipeline | 'Garden plants' >> beam.Create([ '🍓 : Strawberry : perennial', '🥕 : Carrot : biennial', '🍆\t:\tEggplant : perennial', '🍅 : Tomato : annual', '🥔 : Potato : perennial', ]) | 'Parse plants' >> beam.Regex.split(r'\s*:\s*') | beam.Map(print)) # [END regex_split] if test: test(plants_split) if __name__ == '__main__': regex_split() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/regex_test.py0000644000175100017510000001323115206113032030602 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import regex_all_matches from . import regex_find from . import regex_find_all from . import regex_find_kv from . import regex_matches from . import regex_matches_kv from . import regex_replace_all from . import regex_replace_first from . import regex_split def check_matches(actual): expected = '''[START plants_matches] 🍓, Strawberry, perennial 🥕, Carrot, biennial 🍆, Eggplant, perennial 🍅, Tomato, annual 🥔, Potato, perennial [END plants_matches]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_all_matches(actual): expected = '''[START plants_all_matches] ['🍓, Strawberry, perennial', '🍓', 'Strawberry', 'perennial'] ['🥕, Carrot, biennial', '🥕', 'Carrot', 'biennial'] ['🍆, Eggplant, perennial', '🍆', 'Eggplant', 'perennial'] ['🍅, Tomato, annual', '🍅', 'Tomato', 'annual'] ['🥔, Potato, perennial', '🥔', 'Potato', 'perennial'] [END plants_all_matches]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_matches_kv(actual): expected = '''[START plants_matches_kv] ('🍓', '🍓, Strawberry, perennial') ('🥕', '🥕, Carrot, biennial') ('🍆', '🍆, Eggplant, perennial') ('🍅', '🍅, Tomato, annual') ('🥔', '🥔, Potato, perennial') [END plants_matches_kv]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_find_all(actual): expected = '''[START plants_find_all] ['🍓, Strawberry, perennial'] ['🥕, Carrot, biennial'] ['🍆, Eggplant, perennial', '🍌, Banana, perennial'] ['🍅, Tomato, annual', '🍉, Watermelon, annual'] ['🥔, Potato, perennial'] [END plants_find_all]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_find_kv(actual): expected = '''[START plants_find_kv] ('🍓', '🍓, Strawberry, perennial') ('🥕', '🥕, Carrot, biennial') ('🍆', '🍆, Eggplant, perennial') ('🍌', '🍌, Banana, perennial') ('🍅', '🍅, Tomato, annual') ('🍉', '🍉, Watermelon, annual') ('🥔', '🥔, Potato, perennial') [END plants_find_kv]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_replace_all(actual): expected = '''[START plants_replace_all] 🍓,Strawberry,perennial 🥕,Carrot,biennial 🍆,Eggplant,perennial 🍅,Tomato,annual 🥔,Potato,perennial [END plants_replace_all]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_replace_first(actual): expected = '''[START plants_replace_first] 🍓: Strawberry, perennial 🥕: Carrot, biennial 🍆: Eggplant, perennial 🍅: Tomato, annual 🥔: Potato, perennial [END plants_replace_first]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_split(actual): expected = '''[START plants_split] ['🍓', 'Strawberry', 'perennial'] ['🥕', 'Carrot', 'biennial'] ['🍆', 'Eggplant', 'perennial'] ['🍅', 'Tomato', 'annual'] ['🥔', 'Potato', 'perennial'] [END plants_split]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.regex_matches.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.regex_all_matches.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.regex_matches_kv.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.regex_find.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.regex_find_all.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.regex_find_kv.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.regex_replace_all.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.regex_replace_first.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.regex_split.print', str) class RegexTest(unittest.TestCase): def test_matches(self): regex_matches.regex_matches(check_matches) def test_all_matches(self): regex_all_matches.regex_all_matches(check_all_matches) def test_matches_kv(self): regex_matches_kv.regex_matches_kv(check_matches_kv) def test_find(self): regex_find.regex_find(check_matches) def test_find_all(self): regex_find_all.regex_find_all(check_find_all) def test_find_kv(self): regex_find_kv.regex_find_kv(check_find_kv) def test_replace_all(self): regex_replace_all.regex_replace_all(check_replace_all) def test_replace_first(self): regex_replace_first.regex_replace_first(check_replace_first) def test_split(self): regex_split.regex_split(check_split) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/runinference.py0000644000175100017510000000665215206113032031125 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint: disable=reimported import torch class LinearRegression(torch.nn.Module): def __init__(self, input_dim=1, output_dim=1): super().__init__() self.linear = torch.nn.Linear(input_dim, output_dim) def forward(self, x): out = self.linear(x) return out def torch_unkeyed_model_handler(test=None): # [START torch_unkeyed_model_handler] import apache_beam as beam import numpy import torch from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerTensor model_state_dict_path = 'gs://apache-beam-samples/run_inference/five_times_table_torch.pt' # pylint: disable=line-too-long model_class = LinearRegression model_params = {'input_dim': 1, 'output_dim': 1} model_handler = PytorchModelHandlerTensor( model_class=model_class, model_params=model_params, state_dict_path=model_state_dict_path) unkeyed_data = numpy.array([10, 40, 60, 90], dtype=numpy.float32).reshape(-1, 1) with beam.Pipeline() as p: predictions = ( p | 'InputData' >> beam.Create(unkeyed_data) | 'ConvertNumpyToTensor' >> beam.Map(torch.Tensor) | 'PytorchRunInference' >> RunInference(model_handler=model_handler) | beam.Map(print)) # [END torch_unkeyed_model_handler] if test: test(predictions) def torch_keyed_model_handler(test=None): # [START torch_keyed_model_handler] import apache_beam as beam import torch from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerTensor model_state_dict_path = 'gs://apache-beam-samples/run_inference/five_times_table_torch.pt' # pylint: disable=line-too-long model_class = LinearRegression model_params = {'input_dim': 1, 'output_dim': 1} keyed_model_handler = KeyedModelHandler( PytorchModelHandlerTensor( model_class=model_class, model_params=model_params, state_dict_path=model_state_dict_path)) keyed_data = [("first_question", 105.00), ("second_question", 108.00), ("third_question", 1000.00), ("fourth_question", 1013.00)] with beam.Pipeline() as p: predictions = ( p | 'KeyedInputData' >> beam.Create(keyed_data) | "ConvertIntToTensor" >> beam.Map(lambda x: (x[0], torch.Tensor([x[1]]))) | 'PytorchRunInference' >> RunInference(model_handler=keyed_model_handler) | beam.Map(print)) # [END torch_keyed_model_handler] if test: test(predictions) ././@PaxHeader0000000000000000000000000000022300000000000010212 xustar00125 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/runinference_sklearn_keyed_model_handler.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/runinference_sklearn_keyed_m0000644000175100017510000000470315206113032033705 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint: disable=reimported # pylint:disable=line-too-long # beam-playground: # name: RunInferenceSklearnKeyed # description: Demonstration of RunInference transform usage with Sklearn keyed model handler. # multifile: false # default_example: false # context_line: 46 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - inference # - sklearn def sklearn_keyed_model_handler(test=None): # [START sklearn_keyed_model_handler] import apache_beam as beam from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.sklearn_inference import ModelFileType from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerNumpy sklearn_model_filename = 'gs://apache-beam-samples/run_inference/five_times_table_sklearn.pkl' # pylint: disable=line-too-long sklearn_model_handler = KeyedModelHandler( SklearnModelHandlerNumpy( model_uri=sklearn_model_filename, model_file_type=ModelFileType.PICKLE)) keyed_data = [("first_question", 105.00), ("second_question", 108.00), ("third_question", 1000.00), ("fourth_question", 1013.00)] with beam.Pipeline() as p: predictions = ( p | "ReadInputs" >> beam.Create(keyed_data) | "ConvertDataToList" >> beam.Map(lambda x: (x[0], [x[1]])) | "RunInferenceSklearn" >> RunInference(model_handler=sklearn_model_handler) | beam.Map(print)) # [END sklearn_keyed_model_handler] if test: test(predictions) if __name__ == '__main__': sklearn_keyed_model_handler() ././@PaxHeader0000000000000000000000000000022500000000000010214 xustar00127 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/runinference_sklearn_unkeyed_model_handler.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/runinference_sklearn_unkeyed0000644000175100017510000000442115206113032033731 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint: disable=reimported # pylint:disable=line-too-long # beam-playground: # name: RunInferenceSklearnUnkeyed # description: Demonstration of RunInference transform usage with Sklearn unkeyed model handler. # multifile: false # default_example: false # context_line: 46 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - inference # - sklearn def sklearn_unkeyed_model_handler(test=None): # [START sklearn_unkeyed_model_handler] import apache_beam as beam import numpy from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.sklearn_inference import ModelFileType from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerNumpy sklearn_model_filename = 'gs://apache-beam-samples/run_inference/five_times_table_sklearn.pkl' # pylint: disable=line-too-long sklearn_model_handler = SklearnModelHandlerNumpy( model_uri=sklearn_model_filename, model_file_type=ModelFileType.PICKLE) unkeyed_data = numpy.array([20, 40, 60, 90], dtype=numpy.float32).reshape(-1, 1) with beam.Pipeline() as p: predictions = ( p | "ReadInputs" >> beam.Create(unkeyed_data) | "RunInferenceSklearn" >> RunInference(model_handler=sklearn_model_handler) | beam.Map(print)) # [END sklearn_unkeyed_model_handler] if test: test(predictions) if __name__ == '__main__': sklearn_unkeyed_model_handler() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/runinference_test.py0000644000175100017510000001444215206113032032160 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest from io import StringIO import mock import pytest from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import runinference_sklearn_keyed_model_handler from . import runinference_sklearn_unkeyed_model_handler # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports, unused-import try: import torch from . import runinference except ImportError: raise unittest.SkipTest('PyTorch dependencies are not installed') # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports, unused-import try: from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem except ImportError: raise unittest.SkipTest('GCP dependencies are not installed') def check_torch_keyed_model_handler(): expected = '''[START torch_keyed_model_handler] ('first_question', PredictionResult(example=tensor([105.]), inference=tensor([523.6982]), model_id='gs://apache-beam-samples/run_inference/five_times_table_torch.pt')) ('second_question', PredictionResult(example=tensor([108.]), inference=tensor([538.5867]), model_id='gs://apache-beam-samples/run_inference/five_times_table_torch.pt')) ('third_question', PredictionResult(example=tensor([1000.]), inference=tensor([4965.4019]), model_id='gs://apache-beam-samples/run_inference/five_times_table_torch.pt')) ('fourth_question', PredictionResult(example=tensor([1013.]), inference=tensor([5029.9180]), model_id='gs://apache-beam-samples/run_inference/five_times_table_torch.pt')) [END torch_keyed_model_handler] '''.splitlines()[1:-1] return expected def check_sklearn_keyed_model_handler(actual): expected = '''[START sklearn_keyed_model_handler] ('first_question', PredictionResult(example=[105.0], inference=array([525.]), model_id='gs://apache-beam-samples/run_inference/five_times_table_sklearn.pkl')) ('second_question', PredictionResult(example=[108.0], inference=array([540.]), model_id='gs://apache-beam-samples/run_inference/five_times_table_sklearn.pkl')) ('third_question', PredictionResult(example=[1000.0], inference=array([5000.]), model_id='gs://apache-beam-samples/run_inference/five_times_table_sklearn.pkl')) ('fourth_question', PredictionResult(example=[1013.0], inference=array([5065.]), model_id='gs://apache-beam-samples/run_inference/five_times_table_sklearn.pkl')) [END sklearn_keyed_model_handler] '''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_torch_unkeyed_model_handler(): expected = '''[START torch_unkeyed_model_handler] PredictionResult(example=tensor([10.]), inference=tensor([52.2325]), model_id='gs://apache-beam-samples/run_inference/five_times_table_torch.pt') PredictionResult(example=tensor([40.]), inference=tensor([201.1165]), model_id='gs://apache-beam-samples/run_inference/five_times_table_torch.pt') PredictionResult(example=tensor([60.]), inference=tensor([300.3724]), model_id='gs://apache-beam-samples/run_inference/five_times_table_torch.pt') PredictionResult(example=tensor([90.]), inference=tensor([449.2563]), model_id='gs://apache-beam-samples/run_inference/five_times_table_torch.pt') [END torch_unkeyed_model_handler] '''.splitlines()[1:-1] return expected def check_sklearn_unkeyed_model_handler(actual): expected = '''[START sklearn_unkeyed_model_handler] PredictionResult(example=array([20.], dtype=float32), inference=array([100.], dtype=float32), model_id='gs://apache-beam-samples/run_inference/five_times_table_sklearn.pkl') PredictionResult(example=array([40.], dtype=float32), inference=array([200.], dtype=float32), model_id='gs://apache-beam-samples/run_inference/five_times_table_sklearn.pkl') PredictionResult(example=array([60.], dtype=float32), inference=array([300.], dtype=float32), model_id='gs://apache-beam-samples/run_inference/five_times_table_sklearn.pkl') PredictionResult(example=array([90.], dtype=float32), inference=array([450.], dtype=float32), model_id='gs://apache-beam-samples/run_inference/five_times_table_sklearn.pkl') [END sklearn_unkeyed_model_handler] '''.splitlines()[1:-1] assert_matches_stdout(actual, expected) # pylint:disable=line-too-long @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.runinference_sklearn_unkeyed_model_handler.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.runinference_sklearn_keyed_model_handler.print', str) class RunInferenceTest(unittest.TestCase): def test_sklearn_unkeyed_model_handler(self): runinference_sklearn_unkeyed_model_handler.sklearn_unkeyed_model_handler( check_sklearn_unkeyed_model_handler) def test_sklearn_keyed_model_handler(self): runinference_sklearn_keyed_model_handler.sklearn_keyed_model_handler( check_sklearn_keyed_model_handler) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch('sys.stdout', new_callable=StringIO) class RunInferenceStdoutTest(unittest.TestCase): @pytest.mark.uses_pytorch def test_check_torch_keyed_model_handler(self, mock_stdout): runinference.torch_keyed_model_handler() predicted = mock_stdout.getvalue().splitlines() expected = check_torch_keyed_model_handler() self.assertEqual(predicted, expected) @pytest.mark.uses_pytorch def test_check_torch_unkeyed_model_handler(self, mock_stdout): runinference.torch_unkeyed_model_handler() predicted = mock_stdout.getvalue().splitlines() expected = check_torch_unkeyed_model_handler() self.assertEqual(predicted, expected) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/tostring_element.py0000644000175100017510000000340315206113032032013 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: ToStringElement # description: Demonstration of ToString.Element transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - lists def tostring_element(test=None): # [START tostring_element] import apache_beam as beam with beam.Pipeline() as pipeline: plant_lists = ( pipeline | 'Garden plants' >> beam.Create([ ['🍓', 'Strawberry', 'perennial'], ['🥕', 'Carrot', 'biennial'], ['🍆', 'Eggplant', 'perennial'], ['🍅', 'Tomato', 'annual'], ['🥔', 'Potato', 'perennial'], ]) | 'To string' >> beam.ToString.Element() | beam.Map(print)) # [END tostring_element] if test: test(plant_lists) if __name__ == '__main__': tostring_element() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/tostring_iterables.py0000644000175100017510000000341715206113032032341 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: ToStringIterables # description: Demonstration of ToString.Iterables transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - lists def tostring_iterables(test=None): # [START tostring_iterables] import apache_beam as beam with beam.Pipeline() as pipeline: plants_csv = ( pipeline | 'Garden plants' >> beam.Create([ ['🍓', 'Strawberry', 'perennial'], ['🥕', 'Carrot', 'biennial'], ['🍆', 'Eggplant', 'perennial'], ['🍅', 'Tomato', 'annual'], ['🥔', 'Potato', 'perennial'], ]) | 'To string' >> beam.ToString.Iterables() | beam.Map(print)) # [END tostring_iterables] if test: test(plants_csv) if __name__ == '__main__': tostring_iterables() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/tostring_kvs.py0000644000175100017510000000324015206113032031164 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: ToStringKvs # description: Demonstration of ToString.Kvs transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings # - pairs def tostring_kvs(test=None): # [START tostring_kvs] import apache_beam as beam with beam.Pipeline() as pipeline: plants = ( pipeline | 'Garden plants' >> beam.Create([ ('🍓', 'Strawberry'), ('🥕', 'Carrot'), ('🍆', 'Eggplant'), ('🍅', 'Tomato'), ('🥔', 'Potato'), ]) | 'To string' >> beam.ToString.Kvs() | beam.Map(print)) # [END tostring_kvs] if test: test(plants) if __name__ == '__main__': tostring_kvs() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/tostring_test.py0000644000175100017510000000504215206113032031342 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import tostring_element from . import tostring_iterables from . import tostring_kvs def check_plants(actual): expected = '''[START plants] 🍓,Strawberry 🥕,Carrot 🍆,Eggplant 🍅,Tomato 🥔,Potato [END plants]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_plant_lists(actual): expected = '''[START plant_lists] ['🍓', 'Strawberry', 'perennial'] ['🥕', 'Carrot', 'biennial'] ['🍆', 'Eggplant', 'perennial'] ['🍅', 'Tomato', 'annual'] ['🥔', 'Potato', 'perennial'] [END plant_lists]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_plants_csv(actual): expected = '''[START plants_csv] 🍓,Strawberry,perennial 🥕,Carrot,biennial 🍆,Eggplant,perennial 🍅,Tomato,annual 🥔,Potato,perennial [END plants_csv]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.tostring_kvs.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.tostring_element.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.tostring_iterables.print', str) class ToStringTest(unittest.TestCase): def test_tostring_kvs(self): tostring_kvs.tostring_kvs(check_plants) def test_tostring_element(self): tostring_element.tostring_element(check_plant_lists) def test_tostring_iterables(self): tostring_iterables.tostring_iterables(check_plants_csv) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/values.py0000644000175100017510000000314415206113032027732 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: Values # description: Demonstration of Values transform usage. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - pairs def values(test=None): # [START values] import apache_beam as beam with beam.Pipeline() as pipeline: plants = ( pipeline | 'Garden plants' >> beam.Create([ ('🍓', 'Strawberry'), ('🥕', 'Carrot'), ('🍆', 'Eggplant'), ('🍅', 'Tomato'), ('🥔', 'Potato'), ]) | 'Values' >> beam.Values() | beam.Map(print)) # [END values] if test: test(plants) if __name__ == '__main__': values() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/values_test.py0000644000175100017510000000266615206113032031001 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import values def check_plants(actual): expected = '''[START plants] Strawberry Carrot Eggplant Tomato Potato [END plants]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.values.print', str) class ValuesTest(unittest.TestCase): def test_values(self): values.values(check_plants) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/withtimestamps.py0000644000175100017510000000242415206113032031515 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file def time_tuple2unix_time(): # [START time_tuple2unix_time] import time time_tuple = time.strptime('2020-03-19 20:50:00', '%Y-%m-%d %H:%M:%S') unix_time = time.mktime(time_tuple) # [END time_tuple2unix_time] return unix_time def datetime2unix_time(): # [START datetime2unix_time] import time import datetime now = datetime.datetime.now() time_tuple = now.timetuple() unix_time = time.mktime(time_tuple) # [END datetime2unix_time] return unix_time ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_event_time.py0000644000175100017510000000455515206113032033743 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: WithTimestampsEventTime # description: Demonstration of WithTimestamps transform usage. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - timestamps def withtimestamps_event_time(test=None): # [START withtimestamps_event_time] import apache_beam as beam class GetTimestamp(beam.DoFn): def process(self, plant, timestamp=beam.DoFn.TimestampParam): yield '{} - {}'.format(timestamp.to_utc_datetime(), plant['name']) with beam.Pipeline() as pipeline: plant_timestamps = ( pipeline | 'Garden plants' >> beam.Create([ { 'name': 'Strawberry', 'season': 1585699200 }, # April, 2020 { 'name': 'Carrot', 'season': 1590969600 }, # June, 2020 { 'name': 'Artichoke', 'season': 1583020800 }, # March, 2020 { 'name': 'Tomato', 'season': 1588291200 }, # May, 2020 { 'name': 'Potato', 'season': 1598918400 }, # September, 2020 ]) | 'With timestamps' >> beam.Map( lambda plant: beam.window.TimestampedValue(plant, plant['season'])) | 'Get timestamp' >> beam.ParDo(GetTimestamp()) | beam.Map(print)) # [END withtimestamps_event_time] if test: test(plant_timestamps) if __name__ == '__main__': withtimestamps_event_time() ././@PaxHeader0000000000000000000000000000020700000000000010214 xustar00113 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_logical_clock.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_logical_clock0000644000175100017510000000427515206113032033741 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: WithTimestampsLogicalClock # description: Demonstration of WithTimestamps transform usage with logical clock. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - timestamps def withtimestamps_logical_clock(test=None): # [START withtimestamps_logical_clock] import apache_beam as beam class GetTimestamp(beam.DoFn): def process(self, plant, timestamp=beam.DoFn.TimestampParam): event_id = int(timestamp.micros / 1e6) # equivalent to seconds yield '{} - {}'.format(event_id, plant['name']) with beam.Pipeline() as pipeline: plant_events = ( pipeline | 'Garden plants' >> beam.Create([ {'name': 'Strawberry', 'event_id': 1}, {'name': 'Carrot', 'event_id': 4}, {'name': 'Artichoke', 'event_id': 2}, {'name': 'Tomato', 'event_id': 3}, {'name': 'Potato', 'event_id': 5}, ]) | 'With timestamps' >> beam.Map(lambda plant: \ beam.window.TimestampedValue(plant, plant['event_id'])) | 'Get timestamp' >> beam.ParDo(GetTimestamp()) | beam.Map(print) ) # [END withtimestamps_logical_clock] if test: test(plant_events) if __name__ == '__main__': withtimestamps_logical_clock() ././@PaxHeader0000000000000000000000000000021100000000000010207 xustar00115 path=apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_processing_time.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_processing_ti0000644000175100017510000000412215206113032034013 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: WithTimestampsProcessingTime # description: Demonstration of WithTimestamps transform usage. # multifile: false # default_example: false # context_line: 41 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - timestamps def withtimestamps_processing_time(test=None): # [START withtimestamps_processing_time] import apache_beam as beam import time class GetTimestamp(beam.DoFn): def process(self, plant, timestamp=beam.DoFn.TimestampParam): yield '{} - {}'.format(timestamp.to_utc_datetime(), plant['name']) with beam.Pipeline() as pipeline: plant_processing_times = ( pipeline | 'Garden plants' >> beam.Create([ {'name': 'Strawberry'}, {'name': 'Carrot'}, {'name': 'Artichoke'}, {'name': 'Tomato'}, {'name': 'Potato'}, ]) | 'With timestamps' >> beam.Map(lambda plant: \ beam.window.TimestampedValue(plant, time.time())) | 'Get timestamp' >> beam.ParDo(GetTimestamp()) | beam.Map(print) ) # [END withtimestamps_processing_time] if test: test(plant_processing_times) if __name__ == '__main__': withtimestamps_processing_time() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_test.py0000644000175100017510000000666615206113032032570 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long import unittest import mock import apache_beam as beam from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import withtimestamps from . import withtimestamps_event_time from . import withtimestamps_logical_clock from . import withtimestamps_processing_time def check_plant_timestamps(actual): expected = '''[START plant_timestamps] 2020-04-01 00:00:00 - Strawberry 2020-06-01 00:00:00 - Carrot 2020-03-01 00:00:00 - Artichoke 2020-05-01 00:00:00 - Tomato 2020-09-01 00:00:00 - Potato [END plant_timestamps]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_plant_events(actual): expected = '''[START plant_events] 1 - Strawberry 4 - Carrot 2 - Artichoke 3 - Tomato 5 - Potato [END plant_events]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) def check_plant_processing_times(actual): expected = '''[START plant_processing_times] 2020-03-20 20:12:42.145594 - Strawberry 2020-03-20 20:12:42.145827 - Carrot 2020-03-20 20:12:42.145962 - Artichoke 2020-03-20 20:12:42.146093 - Tomato 2020-03-20 20:12:42.146216 - Potato [END plant_processing_times]'''.splitlines()[1:-1] # Since `time.time()` will always give something different, we'll # simply strip the timestamp information before testing the results. actual = actual | beam.Map(lambda row: row.split('-')[-1].strip()) expected = [row.split('-')[-1].strip() for row in expected] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.withtimestamps_event_time.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.withtimestamps_logical_clock.print', str) @mock.patch( 'apache_beam.examples.snippets.transforms.elementwise.withtimestamps_processing_time.print', str) class WithTimestampsTest(unittest.TestCase): def test_event_time(self): withtimestamps_event_time.withtimestamps_event_time(check_plant_timestamps) def test_logical_clock(self): withtimestamps_logical_clock.withtimestamps_logical_clock( check_plant_events) def test_processing_time(self): withtimestamps_processing_time.withtimestamps_processing_time( check_plant_processing_times) def test_time_tuple2unix_time(self): unix_time = withtimestamps.time_tuple2unix_time() self.assertIsInstance(unix_time, float) def test_datetime2unix_time(self): unix_time = withtimestamps.datetime2unix_time() self.assertIsInstance(unix_time, float) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.239393 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/other/0000755000175100017510000000000015206113157024667 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/other/__init__.py0000644000175100017510000000142015206113032026765 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/other/create.py0000644000175100017510000000271615206113032026502 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: Create # description: Demonstration of Create transform usage. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings def create(test=None): # [START create] import apache_beam as beam with beam.Pipeline() as pipeline: produce = ( pipeline | 'Garden plants' >> beam.Create(['🍓', '🥕', '🍆', '🍅', '🥔']) | 'Print results' >> beam.Map(print)) # [END create] if test: test(produce) if __name__ == '__main__': create() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/other/create_test.py0000644000175100017510000000271215206113032027535 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import create def check_create(actual): expected = '''[START flatten_result] 🍓 🥕 🍆 🍅 🥔 [END flatten_result]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch('apache_beam.examples.snippets.transforms.other.create.print', str) class CreateTest(unittest.TestCase): def test_create(self): create.create(check_create) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/other/flatten.py0000644000175100017510000000314715206113032026673 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: FlattenTransform # description: Demonstration of Flatten transform usage. # multifile: false # default_example: false # context_line: 40 # categories: # - Core Transforms # complexity: BASIC # tags: # - transforms # - strings def flatten(test=None): # [START flatten] import apache_beam as beam with beam.Pipeline() as pipeline: pc1 = pipeline | 'Create produce 1' >> beam.Create(['🍓', '🥕', '🍆']) pc2 = pipeline | 'Create produce 2' >> beam.Create(['🍅', '🥔']) pc3 = pipeline | 'Create produce 3' >> beam.Create(['🍎', '🍐', '🍊']) icons = ((pc1, pc2, pc3) | beam.Flatten() | beam.Map(print)) # [END flatten] if test: test(icons) if __name__ == '__main__': flatten() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/other/flatten_test.py0000644000175100017510000000270215206113032027726 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import flatten def check_flatten(actual): expected = '''[START flatten_result] 🍓 🥕 🍆 🍅 🥔 🍎 🍐 🍊 [END flatten_result]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch('apache_beam.examples.snippets.transforms.other.flatten.print', str) class FlattenTest(unittest.TestCase): def test_flatten(self): flatten.flatten(check_flatten) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/other/window.py0000644000175100017510000000477415206113032026554 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # pylint:disable=line-too-long # beam-playground: # name: Window # description: Demonstration of Window transform usage. # multifile: false # default_example: false # context_line: 43 # categories: # - Core Transforms # - Windowing # complexity: BASIC # tags: # - transforms # - strings # - timestamps # - windows def window(test=None): # [START window] import apache_beam as beam with beam.Pipeline() as pipeline: produce = ( pipeline | 'Garden plants' >> beam.Create([ { 'name': 'Strawberry', 'season': 1585699200 }, # April, 2020 { 'name': 'Strawberry', 'season': 1588291200 }, # May, 2020 { 'name': 'Carrot', 'season': 1590969600 }, # June, 2020 { 'name': 'Artichoke', 'season': 1583020800 }, # March, 2020 { 'name': 'Artichoke', 'season': 1585699200 }, # April, 2020 { 'name': 'Tomato', 'season': 1588291200 }, # May, 2020 { 'name': 'Potato', 'season': 1598918400 }, # September, 2020 ]) | 'With timestamps' >> beam.Map( lambda plant: beam.window.TimestampedValue( plant['name'], plant['season'])) | 'Window into fixed 2-month windows' >> beam.WindowInto( beam.window.FixedWindows(2 * 30 * 24 * 60 * 60)) | 'Count per window' >> beam.combiners.Count.PerElement() | 'Print results' >> beam.Map(print)) # [END window] if test: test(produce) if __name__ == '__main__': window() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/transforms/other/window_test.py0000644000175100017510000000276115206113032027605 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import mock from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.testing.test_pipeline import TestPipeline from . import window def check_window(actual): expected = '''[START flatten_result] ('Strawberry', 1) ('Strawberry', 1) ('Carrot', 1) ('Artichoke', 2) ('Tomato', 1) ('Potato', 1) [END flatten_result]'''.splitlines()[1:-1] assert_matches_stdout(actual, expected) @mock.patch('apache_beam.Pipeline', TestPipeline) @mock.patch('apache_beam.examples.snippets.transforms.other.window.print', str) class WindowTest(unittest.TestCase): def test_window(self): window.window(check_window) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/util.py0000644000175100017510000000441515206113032022673 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import ast import shlex import subprocess as sp import apache_beam as beam from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to def assert_matches_stdout( actual, expected_stdout, normalize_fn=lambda elem: elem, label=''): """Asserts a PCollection of strings matches the expected stdout elements. Args: actual (beam.PCollection): A PCollection. expected (List[str]): A list of stdout elements, one line per element. normalize_fn (Function[any]): A function to normalize elements before comparing them. Can be used to sort lists before comparing. label (str): [optional] Label to make transform names unique. """ def stdout_to_python_object(elem_str): try: elem = ast.literal_eval(elem_str) except (SyntaxError, ValueError): elem = elem_str return normalize_fn(elem) actual = actual | label >> beam.Map(stdout_to_python_object) expected = list(map(stdout_to_python_object, expected_stdout)) assert_that(actual, equal_to(expected), 'assert ' + label) def run_shell_commands(commands, **kwargs): """Runs a list of Notebook-like shell commands. Lines starting with `#` are ignored as comments. Lines starting with `!` are run as commands. Variables like `{variable}` are substituted with **kwargs. """ for cmd in commands: cmd = cmd.strip().lstrip('!').format(**kwargs) sp_cmd = shlex.split(cmd, comments=True, posix=True) if sp_cmd: sp.call(sp_cmd) yield sp_cmd ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/snippets/util_test.py0000644000175100017510000000623015206113032023727 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest from mock import patch import apache_beam as beam from apache_beam.examples.snippets import util from apache_beam.testing.test_pipeline import TestPipeline class UtilTest(unittest.TestCase): def test_assert_matches_stdout_object(self): expected = [ "{'a': '🍓', 'b': True}", "{'a': '🥕', 'b': 42}", "{'a': '🍆', 'b': '\"hello\"'}", "{'a': '🍅', 'b': [1, 2, 3]}", "{'b': 'B', 'a': '🥔'}", ] with TestPipeline() as pipeline: actual = ( pipeline | beam.Create([ { 'a': '🍓', 'b': True }, { 'a': '🥕', 'b': 42 }, { 'a': '🍆', 'b': '"hello"' }, { 'a': '🍅', 'b': [1, 2, 3] }, { 'a': '🥔', 'b': 'B' }, ]) | beam.Map(str)) util.assert_matches_stdout(actual, expected) def test_assert_matches_stdout_string(self): expected = ['🍓', '🥕', '🍆', '🍅', '🥔'] with TestPipeline() as pipeline: actual = ( pipeline | beam.Create(['🍓', '🥕', '🍆', '🍅', '🥔']) | beam.Map(str)) util.assert_matches_stdout(actual, expected) def test_assert_matches_stdout_sorted_keys(self): expected = [{'list': [1, 2]}, {'list': [3, 4]}] with TestPipeline() as pipeline: actual = ( pipeline | beam.Create([{ 'list': [2, 1] }, { 'list': [4, 3] }]) | beam.Map(str)) util.assert_matches_stdout( actual, expected, lambda elem: {'sorted': sorted(elem['list'])}) @patch('subprocess.call', lambda cmd: None) def test_run_shell_commands(self): commands = [ ' # this is a comment ', ' ! echo this is a shell command ', ' !echo {variable} ', ' echo "quoted arguments work" # trailing comment ', ] actual = list(util.run_shell_commands(commands, variable='hello world')) expected = [ ['echo', 'this', 'is', 'a', 'shell', 'command'], ['echo', 'hello', 'world'], ['echo', 'quoted arguments work'], ] self.assertEqual(actual, expected) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/sql_taxi.py0000644000175100017510000000746415206113032021704 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """An example that processes streaming NYC Taxi data with SqlTransform. This example reads from the PubSub NYC Taxi stream described in https://github.com/googlecodelabs/cloud-dataflow-nyc-taxi-tycoon, aggregates the data in 15s windows using SqlTransform, and writes the output to a user-defined PubSub topic. A Java version supported by Beam must be installed locally to run this pipeline. Additionally, Docker must also be available to run this pipeline locally. """ # pytype: skip-file import json import logging import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.transforms.sql import SqlTransform def run(output_topic, pipeline_args): pipeline_options = PipelineOptions( pipeline_args, save_main_session=True, streaming=True) with beam.Pipeline(options=pipeline_options) as pipeline: _ = ( pipeline | beam.io.ReadFromPubSub( topic='projects/pubsub-public-data/topics/taxirides-realtime', timestamp_attribute="ts").with_output_types(bytes) | "Parse JSON payload" >> beam.Map(json.loads) # Use beam.Row to create a schema-aware PCollection | "Create beam Row" >> beam.Map( lambda x: beam.Row( ride_status=str(x['ride_status']), passenger_count=int( x['passenger_count']))) # SqlTransform will computes result within an existing window | "15s fixed windows" >> beam.WindowInto(beam.window.FixedWindows(15)) # Aggregate drop offs and pick ups that occur within each 15s window | SqlTransform( """ SELECT ride_status, COUNT(*) AS num_rides, SUM(passenger_count) AS total_passengers FROM PCOLLECTION WHERE NOT ride_status = 'enroute' GROUP BY ride_status""") # SqlTransform yields python objects with attributes corresponding to # the outputs of the query. # Collect those attributes, as well as window information, into a dict | "Assemble Dictionary" >> beam.Map( lambda row, window=beam.DoFn.WindowParam: { "ride_status": row.ride_status, "num_rides": row.num_rides, "total_passengers": row.total_passengers, "window_start": window .start.to_rfc3339(), "window_end": window.end.to_rfc3339() }) | "Convert to JSON" >> beam.Map(json.dumps) | "UTF-8 encode" >> beam.Map(lambda s: s.encode("utf-8")) | beam.io.WriteToPubSub(topic=output_topic)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) import argparse parser = argparse.ArgumentParser() parser.add_argument( '--output_topic', dest='output_topic', required=True, help=( 'Cloud PubSub topic to write to (e.g. ' 'projects/my-project/topics/my-topic), must be created prior to ' 'running the pipeline.')) known_args, pipeline_args = parser.parse_known_args() run(known_args.output_topic, pipeline_args) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/streaming_wordcount.py0000644000175100017510000000733115206113032024146 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A streaming word-counting workflow. """ # pytype: skip-file import argparse import logging import apache_beam as beam from apache_beam.examples.wordcount_with_metrics import WordExtractingDoFn from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.transforms import window def run(argv=None, save_main_session=True): """Build and run the pipeline.""" parser = argparse.ArgumentParser() parser.add_argument( '--output_topic', required=True, help=( 'Output PubSub topic of the form ' '"projects//topics/".')) group = parser.add_mutually_exclusive_group(required=True) group.add_argument( '--input_topic', help=( 'Input PubSub topic of the form ' '"projects//topics/".')) group.add_argument( '--input_subscription', help=( 'Input PubSub subscription of the form ' '"projects//subscriptions/."')) known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session pipeline_options.view_as(StandardOptions).streaming = True with beam.Pipeline(options=pipeline_options) as p: # Read from PubSub into a PCollection. if known_args.input_subscription: messages = ( p | beam.io.ReadFromPubSub(subscription=known_args.input_subscription). with_output_types(bytes)) else: messages = ( p | beam.io.ReadFromPubSub( topic=known_args.input_topic).with_output_types(bytes)) lines = messages | 'decode' >> beam.Map(lambda x: x.decode('utf-8')) # Count the occurrences of each word. def count_ones(word_ones): (word, ones) = word_ones return (word, sum(ones)) counts = ( lines | 'split' >> (beam.ParDo(WordExtractingDoFn()).with_output_types(str)) | 'pair_with_one' >> beam.Map(lambda x: (x, 1)) | beam.WindowInto(window.FixedWindows(15, 0)) | 'group' >> beam.GroupByKey() | 'count' >> beam.Map(count_ones)) # Format the counts into a PCollection of strings. def format_result(word_count): (word, count) = word_count return '%s: %d' % (word, count) output = ( counts | 'format' >> beam.Map(format_result) | 'encode' >> beam.Map(lambda x: x.encode('utf-8')).with_output_types(bytes)) # Write to PubSub. # pylint: disable=expression-not-assigned output | beam.io.WriteToPubSub(known_args.output_topic) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/streaming_wordcount_debugging.py0000644000175100017510000001601415206113032026157 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A streaming wordcount example with debugging capabilities. It demonstrate the use of logging and assert_that in streaming mode. This workflow only works with the DirectRunner (https://github.com/apache/beam/issues/18709). Usage: python streaming_wordcount_debugging.py --input_topic projects/$PROJECT_ID/topics/$PUBSUB_INPUT_TOPIC --output_topic projects/$PROJECT_ID/topics/$PUBSUB_OUTPUT_TOPIC --streaming To publish messages: gcloud alpha pubsub topics publish $PUBSUB_INPUT_TOPIC --message '210 213 151' """ # pytype: skip-file import argparse import logging import re import time import apache_beam as beam from apache_beam.examples.wordcount import WordExtractingDoFn from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to_per_window from apache_beam.transforms import window from apache_beam.transforms.core import ParDo class PrintFn(beam.DoFn): """A DoFn that prints label, element, its window, and its timstamp. """ def __init__(self, label): self.label = label def process( self, element, timestamp=beam.DoFn.TimestampParam, window=beam.DoFn.WindowParam): # Log at INFO level each element processed. logging.info('[%s]: %s %s %s', self.label, element, window, timestamp) yield element class AddTimestampFn(beam.DoFn): """A DoFn that attaches timestamps to its elements. It takes an element and attaches a timestamp of its same value for integer and current timestamp in other cases. For example, 120 and Sometext will result in: (120, Timestamp(120) and (Sometext, Timestamp(1234567890). """ def process(self, element): logging.info('Adding timestamp to: %s', element) try: timestamp = int(element) except ValueError: timestamp = int(time.time()) yield beam.window.TimestampedValue(element, timestamp) def run(argv=None, save_main_session=True): """Build and run the pipeline.""" parser = argparse.ArgumentParser() parser.add_argument( '--output_topic', required=True, help=( 'Output PubSub topic of the form ' '"projects//topic/".')) group = parser.add_mutually_exclusive_group(required=True) group.add_argument( '--input_topic', help=( 'Input PubSub topic of the form ' '"projects//topics/".')) group.add_argument( '--input_subscription', help=( 'Input PubSub subscription of the form ' '"projects//subscriptions/."')) known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session pipeline_options.view_as(StandardOptions).streaming = True with beam.Pipeline(options=pipeline_options) as p: # Read from PubSub into a PCollection. if known_args.input_subscription: messages = p | beam.io.ReadFromPubSub( subscription=known_args.input_subscription) else: messages = p | beam.io.ReadFromPubSub(topic=known_args.input_topic) lines = messages | 'decode' >> beam.Map(lambda x: x.decode('utf-8')) # Count the occurrences of each word. def count_ones(word_ones): (word, ones) = word_ones return (word, sum(ones)) counts = ( lines | 'Split' >> (beam.ParDo(WordExtractingDoFn()).with_output_types(str)) | 'AddTimestampFn' >> beam.ParDo(AddTimestampFn()) | 'After AddTimestampFn' >> ParDo(PrintFn('After AddTimestampFn')) | 'PairWithOne' >> beam.Map(lambda x: (x, 1)) | beam.WindowInto(window.FixedWindows(5, 0)) | 'GroupByKey' >> beam.GroupByKey() | 'CountOnes' >> beam.Map(count_ones)) # Format the counts into a PCollection of strings. def format_result(word_count): (word, count) = word_count return '%s: %d' % (word, count) output = ( counts | 'format' >> beam.Map(format_result) | 'encode' >> beam.Map(lambda x: x.encode('utf-8')).with_output_types(bytes)) # Write to PubSub. # pylint: disable=expression-not-assigned output | beam.io.WriteToPubSub(known_args.output_topic) def check_gbk_format(): # A matcher that checks that the output of GBK is of the form word: count. def matcher(elements): # pylint: disable=unused-variable actual_elements_in_window, window = elements for elm in actual_elements_in_window: assert re.match(r'\S+:\s+\d+', elm.decode('utf-8')) is not None return matcher # Check that the format of the output is correct. assert_that( output, check_gbk_format(), use_global_window=False, label='Assert word:count format.') # Check also that elements are ouput in the right window. # This expects exactly 1 occurrence of any subset of the elements # 150, 151, 152, 153, 154 in the window [150, 155) # or exactly 1 occurrence of any subset of the elements # 210, 211, 212, 213, 214 in the window [210, 215). first_window_val = [ '150: 1', '151: 1', '152: 1', '153: 1', '154: 1', ] second_window_val = [ '210: 1', '211: 1', '212: 1', '213: 1', '214: 1', ] expected_window_to_elements = { window.IntervalWindow(150, 155): [ x.encode('utf-8') for x in first_window_val ], window.IntervalWindow(210, 215): [ x.encode('utf-8') for x in second_window_val ], } # To pass, publish numbers in [150-155) or [210-215) with no repeats. # To fail, publish a repeated number in the range above range. # For example: '210 213 151 213' assert_that( output, equal_to_per_window(expected_window_to_elements), use_global_window=False, label='Assert correct streaming windowing.') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/streaming_wordcount_debugging_it_test.py0000644000175100017510000001100615206113032027706 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for the streaming wordcount example with debug.""" # pytype: skip-file import logging import unittest import uuid import pytest from hamcrest.core.core.allof import all_of from apache_beam.examples import streaming_wordcount_debugging from apache_beam.io.gcp.tests.pubsub_matcher import PubSubMessageMatcher from apache_beam.runners.runner import PipelineState from apache_beam.testing import test_utils from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline INPUT_TOPIC = 'wc_topic_input' OUTPUT_TOPIC = 'wc_topic_output' INPUT_SUB = 'wc_subscription_input' OUTPUT_SUB = 'wc_subscription_output' SAMPLE_MESSAGES = [ '150', '151', '152', '153', '154', '210', '211', '212', '213', '214' ] EXPECTED_MESSAGE = [ '150: 1', '151: 1', '152: 1', '153: 1', '154: 1', '210: 1', '211: 1', '212: 1', '213: 1', '214: 1' ] WAIT_UNTIL_FINISH_DURATION = 6 * 60 * 1000 # in milliseconds class StreamingWordcountDebuggingIT(unittest.TestCase): def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.project = self.test_pipeline.get_option('project') self.setup_pubsub() def setup_pubsub(self): self.uuid = str(uuid.uuid4()) # Set up PubSub environment. from google.cloud import pubsub self.pub_client = pubsub.PublisherClient() self.input_topic = self.pub_client.create_topic( name=self.pub_client.topic_path(self.project, INPUT_TOPIC + self.uuid)) self.output_topic = self.pub_client.create_topic( name=self.pub_client.topic_path(self.project, OUTPUT_TOPIC + self.uuid)) self.sub_client = pubsub.SubscriberClient() self.input_sub = self.sub_client.create_subscription( name=self.sub_client.subscription_path( self.project, INPUT_SUB + self.uuid), topic=self.input_topic.name) self.output_sub = self.sub_client.create_subscription( name=self.sub_client.subscription_path( self.project, OUTPUT_SUB + self.uuid), topic=self.output_topic.name, ack_deadline_seconds=60) def _inject_data(self, topic, data): """Inject numbers as test data to PubSub.""" logging.debug('Injecting test data to topic %s', topic.name) for n in data: self.pub_client.publish(self.input_topic.name, str(n).encode('utf-8')) def tearDown(self): test_utils.cleanup_subscriptions( self.sub_client, [self.input_sub, self.output_sub]) test_utils.cleanup_topics( self.pub_client, [self.input_topic, self.output_topic]) @pytest.mark.it_postcommit @unittest.skip( "Skipped due to [https://github.com/apache/beam/issues/18709]: " "assert_that not working for streaming") def test_streaming_wordcount_debugging_it(self): # Set extra options to the pipeline for test purpose state_verifier = PipelineStateMatcher(PipelineState.RUNNING) pubsub_msg_verifier = PubSubMessageMatcher( self.project, self.output_sub.name, EXPECTED_MESSAGE, timeout=400) extra_opts = { 'input_subscription': self.input_sub.name, 'output_topic': self.output_topic.name, 'wait_until_finish_duration': WAIT_UNTIL_FINISH_DURATION, 'on_success_matcher': all_of(state_verifier, pubsub_msg_verifier) } # Generate input data and inject to PubSub. self._inject_data(self.input_topic, SAMPLE_MESSAGES) # Get pipeline options from command argument: --test-pipeline-options, # and start pipeline job by calling pipeline main function. streaming_wordcount_debugging.run( self.test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/streaming_wordcount_debugging_test.py0000644000175100017510000000674615206113032027231 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit test for the streaming wordcount example with debug.""" # pytype: skip-file import unittest import mock import apache_beam as beam from apache_beam.examples import streaming_wordcount_debugging from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # Protect against environments where the PubSub library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from google.cloud import pubsub except ImportError: pubsub = None # pylint: enable=wrong-import-order, wrong-import-position class StreamingWordcountDebugging(unittest.TestCase): @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') @mock.patch('apache_beam.io.ReadFromPubSub') @mock.patch('apache_beam.io.WriteToPubSub') def test_streaming_wordcount_debugging(self, *unused_mocks): def FakeReadFromPubSub(topic=None, subscription=None, values=None): expected_topic = topic expected_subscription = subscription def _inner(topic=None, subscription=None): assert topic == expected_topic assert subscription == expected_subscription return TestStream().add_elements(values) return _inner class AssertTransform(beam.PTransform): def __init__(self, matcher): self.matcher = matcher def expand(self, pcoll): assert_that(pcoll, self.matcher) def FakeWriteToPubSub(topic=None, values=None): expected_topic = topic def _inner(topic=None, subscription=None): assert topic == expected_topic return AssertTransform(equal_to(values)) return _inner input_topic = 'projects/fake-beam-test-project/topic/intopic' input_values = [ '150', '151', '152', '153', '154', '210', '211', '212', '213', '214' ] output_topic = 'projects/fake-beam-test-project/topic/outtopic' output_values = [ '150: 1', '151: 1', '152: 1', '153: 1', '154: 1', '210: 1', '211: 1', '212: 1', '213: 1', '214: 1' ] beam.io.ReadFromPubSub = ( FakeReadFromPubSub( topic=input_topic, values=list(x.encode('utf-8') for x in input_values))) beam.io.WriteToPubSub = ( FakeWriteToPubSub( topic=output_topic, values=list(x.encode('utf-8') for x in output_values))) streaming_wordcount_debugging.run([ '--input_topic', 'projects/fake-beam-test-project/topic/intopic', '--output_topic', 'projects/fake-beam-test-project/topic/outtopic' ], save_main_session=False) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/streaming_wordcount_it_test.py0000644000175100017510000001036315206113032025700 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for the streaming wordcount example.""" # pytype: skip-file import logging import unittest import uuid import pytest from hamcrest.core.core.allof import all_of from apache_beam.examples import streaming_wordcount from apache_beam.io.gcp.tests.pubsub_matcher import PubSubMessageMatcher from apache_beam.runners.runner import PipelineState from apache_beam.testing import test_utils from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline INPUT_TOPIC = 'wc_topic_input' OUTPUT_TOPIC = 'wc_topic_output' INPUT_SUB = 'wc_subscription_input' OUTPUT_SUB = 'wc_subscription_output' DEFAULT_INPUT_NUMBERS = 500 WAIT_UNTIL_FINISH_DURATION = 20 * 60 * 1000 # in milliseconds class StreamingWordCountIT(unittest.TestCase): def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.project = self.test_pipeline.get_option('project') self.uuid = str(uuid.uuid4()) # Set up PubSub environment. from google.cloud import pubsub self.pub_client = pubsub.PublisherClient() self.input_topic = self.pub_client.create_topic( name=self.pub_client.topic_path(self.project, INPUT_TOPIC + self.uuid)) self.output_topic = self.pub_client.create_topic( name=self.pub_client.topic_path(self.project, OUTPUT_TOPIC + self.uuid)) self.sub_client = pubsub.SubscriberClient() self.input_sub = self.sub_client.create_subscription( name=self.sub_client.subscription_path( self.project, INPUT_SUB + self.uuid), topic=self.input_topic.name) self.output_sub = self.sub_client.create_subscription( name=self.sub_client.subscription_path( self.project, OUTPUT_SUB + self.uuid), topic=self.output_topic.name, ack_deadline_seconds=60) def _inject_numbers(self, topic, num_messages): """Inject numbers as test data to PubSub.""" logging.debug('Injecting %d numbers to topic %s', num_messages, topic.name) for n in range(num_messages): self.pub_client.publish(self.input_topic.name, str(n).encode('utf-8')) def tearDown(self): test_utils.cleanup_subscriptions( self.sub_client, [self.input_sub, self.output_sub]) test_utils.cleanup_topics( self.pub_client, [self.input_topic, self.output_topic]) @pytest.mark.it_postcommit def test_streaming_wordcount_it(self): # Build expected dataset. expected_msg = [('%d: 1' % num).encode('utf-8') for num in range(DEFAULT_INPUT_NUMBERS)] # Set extra options to the pipeline for test purpose state_verifier = PipelineStateMatcher(PipelineState.RUNNING) pubsub_msg_verifier = PubSubMessageMatcher( self.project, self.output_sub.name, expected_msg, timeout=400) extra_opts = { 'input_subscription': self.input_sub.name, 'output_topic': self.output_topic.name, 'wait_until_finish_duration': WAIT_UNTIL_FINISH_DURATION, 'on_success_matcher': all_of(state_verifier, pubsub_msg_verifier) } # Generate input data and inject to PubSub. self._inject_numbers(self.input_topic, DEFAULT_INPUT_NUMBERS) # Get pipeline options from command argument: --test-pipeline-options, # and start pipeline job by calling pipeline main function. streaming_wordcount.run( self.test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/windowed_wordcount.py0000644000175100017510000000622115206113032023772 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A streaming word-counting workflow. Important: streaming pipeline support in Python Dataflow is in development and is not yet available for use. """ # pytype: skip-file import argparse import logging import apache_beam as beam from apache_beam.transforms import window TABLE_SCHEMA = ( 'word:STRING, count:INTEGER, ' 'window_start:TIMESTAMP, window_end:TIMESTAMP') def find_words(element): import re return re.findall(r'[A-Za-z\']+', element) class FormatDoFn(beam.DoFn): def process(self, element, window=beam.DoFn.WindowParam): ts_format = '%Y-%m-%d %H:%M:%S.%f UTC' window_start = window.start.to_utc_datetime().strftime(ts_format) window_end = window.end.to_utc_datetime().strftime(ts_format) return [{ 'word': element[0], 'count': element[1], 'window_start': window_start, 'window_end': window_end }] def main(argv=None): """Build and run the pipeline.""" parser = argparse.ArgumentParser() parser.add_argument( '--input_topic', required=True, help='Input PubSub topic of the form "/topics//".') parser.add_argument( '--output_table', required=True, help=( 'Output BigQuery table for results specified as: ' 'PROJECT:DATASET.TABLE or DATASET.TABLE.')) known_args, pipeline_args = parser.parse_known_args(argv) with beam.Pipeline(argv=pipeline_args) as p: # Read the text from PubSub messages. lines = p | beam.io.ReadFromPubSub(known_args.input_topic) # Get the number of appearances of a word. def count_ones(word_ones): (word, ones) = word_ones return (word, sum(ones)) transformed = ( lines | 'Split' >> (beam.FlatMap(find_words).with_output_types(str)) | 'PairWithOne' >> beam.Map(lambda x: (x, 1)) | beam.WindowInto(window.FixedWindows(2 * 60, 0)) | 'Group' >> beam.GroupByKey() | 'Count' >> beam.Map(count_ones) | 'Format' >> beam.ParDo(FormatDoFn())) # Write to BigQuery. # pylint: disable=expression-not-assigned transformed | 'Write' >> beam.io.WriteToBigQuery( known_args.output_table, schema=TABLE_SCHEMA, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/wordcount.py0000644000175100017510000000713215206113032022074 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A word-counting workflow.""" # pytype: skip-file # beam-playground: # name: WordCount # description: An example that counts words in Shakespeare's works. # multifile: false # pipeline_options: --output output.txt # context_line: 87 # categories: # - Combiners # - Options # - Quickstart # complexity: MEDIUM # tags: # - options # - count # - combine # - strings import argparse import logging import re import apache_beam as beam from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.runner import PipelineResult class WordExtractingDoFn(beam.DoFn): """Parse each line of input text into words.""" def process(self, element): """Returns an iterator over the words of this element. The element is a line of text. If the line is blank, note that, too. Args: element: the element being processed Returns: The processed element. """ return re.findall(r'[\w\']+', element, re.UNICODE) def run(argv=None, save_main_session=True) -> PipelineResult: """Main entry point; defines and runs the wordcount pipeline.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='Input file to process.') parser.add_argument( '--output', dest='output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session pipeline = beam.Pipeline(options=pipeline_options) # Read the text file[pattern] into a PCollection. lines = pipeline | 'Read' >> ReadFromText(known_args.input) counts = ( lines | 'Split' >> (beam.ParDo(WordExtractingDoFn()).with_output_types(str)) | 'PairWithOne' >> beam.Map(lambda x: (x, 1)) | 'GroupAndSum' >> beam.CombinePerKey(sum)) # Format the counts into a PCollection of strings. def format_result(word, count): return '%s: %d' % (word, count) output = counts | 'Format' >> beam.MapTuple(format_result) # Write the output using a "Write" transform that has side effects. # pylint: disable=expression-not-assigned output | 'Write' >> WriteToText(known_args.output) # Execute the pipeline and return the result. result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/wordcount_dataframe.py0000644000175100017510000000205515206113032024077 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Alias for apache_beam.examples.dataframe.wordcount, a word-counting workflow using the DataFrame API.""" # pytype: skip-file import logging from apache_beam.examples.dataframe.wordcount import run if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/wordcount_debugging.py0000644000175100017510000001567515206113032024122 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """An example that verifies the counts and includes best practices. On top of the basic concepts in the wordcount example, this workflow introduces logging to Cloud Logging, and using assertions in a Dataflow pipeline. To execute this pipeline locally, specify a local output file or output prefix on GCS:: --output [YOUR_LOCAL_FILE | gs://YOUR_OUTPUT_PREFIX] To execute this pipeline using the Google Cloud Dataflow service, specify pipeline configuration:: --project YOUR_PROJECT_ID --staging_location gs://YOUR_STAGING_DIRECTORY --temp_location gs://YOUR_TEMP_DIRECTORY --region GCE_REGION --job_name YOUR_JOB_NAME --runner DataflowRunner and an output prefix on GCS:: --output gs://YOUR_OUTPUT_PREFIX """ # pytype: skip-file # beam-playground: # name: WordCountDebugging # description: An example that counts words in Shakespeare's works. # includes regex filter("Flourish|stomach"). # multifile: false # pipeline_options: --output output.txt # context_line: 152 # categories: # - Flatten # - Debugging # - Options # - Combiners # - Filtering # - Quickstart # complexity: ADVANCED # tags: # - count # - debug # - strings import argparse import logging import re import apache_beam as beam from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.metrics import Metrics from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class FilterTextFn(beam.DoFn): """A DoFn that filters for a specific key based on a regular expression.""" def __init__(self, pattern): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.DoFn.__init__(self) self.pattern = pattern # A custom metric can track values in your pipeline as it runs. Those # values will be available in the monitoring system of the runner used # to run the pipeline. These metrics below track the number of # matched and unmatched words. self.matched_words = Metrics.counter(self.__class__, 'matched_words') self.umatched_words = Metrics.counter(self.__class__, 'umatched_words') def process(self, element): word, _ = element if re.match(self.pattern, word): # Log at INFO level each element we match. When executing this pipeline # using the Dataflow service, these log lines will appear in the Cloud # Logging UI. logging.info('Matched %s', word) self.matched_words.inc() yield element else: # Log at the "DEBUG" level each element that is not matched. Different log # levels can be used to control the verbosity of logging providing an # effective mechanism to filter less important information. # Note currently only "INFO" and higher level logs are emitted to the # Cloud Logger. This log message will not be visible in the Cloud Logger. logging.debug('Did not match %s', word) self.umatched_words.inc() class CountWords(beam.PTransform): """A transform to count the occurrences of each word. A PTransform that converts a PCollection containing lines of text into a PCollection of (word, count) tuples. """ def expand(self, pcoll): def count_ones(word_ones): (word, ones) = word_ones return (word, sum(ones)) return ( pcoll | 'split' >> ( beam.FlatMap( lambda x: re.findall(r'[A-Za-z\']+', x)).with_output_types(str)) | 'pair_with_one' >> beam.Map(lambda x: (x, 1)) | 'group' >> beam.GroupByKey() | 'count' >> beam.Map(count_ones)) def run(argv=None, save_main_session=True): """Runs the debugging wordcount pipeline.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='Input file to process.') parser.add_argument( '--output', dest='output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session with beam.Pipeline(options=pipeline_options) as p: # Read the text file[pattern] into a PCollection, count the occurrences of # each word and filter by a list of words. filtered_words = ( p | 'read' >> ReadFromText(known_args.input) | CountWords() | 'FilterText' >> beam.ParDo(FilterTextFn('Flourish|stomach'))) # assert_that is a convenient PTransform that checks a PCollection has an # expected value. Asserts are best used in unit tests with small data sets # but is demonstrated here as a teaching tool. # # Note assert_that does not provide any output and that successful # completion of the Pipeline implies that the expectations were met. Learn # more at https://cloud.google.com/dataflow/pipelines/testing-your-pipeline # on how to best test your pipeline. assert_that(filtered_words, equal_to([('Flourish', 3), ('stomach', 1)])) # Format the counts into a PCollection of strings and write the output using # a "Write" transform that has side effects. # pylint: disable=unused-variable def format_result(word_count): (word, count) = word_count return '%s: %s' % (word, count) output = ( filtered_words | 'format' >> beam.Map(format_result) | 'write' >> WriteToText(known_args.output)) if __name__ == '__main__': # Cloud Logging would contain only logging.INFO and higher level logs logged # by the root logger. All log statements emitted by the root logger will be # visible in the Cloud Logging UI. Learn more at # https://cloud.google.com/logging about the Cloud Logging UI. # # You can set the default logging level to a different level when running # locally. logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/wordcount_debugging_test.py0000644000175100017510000000462315206113032025150 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the debugging wordcount example.""" # pytype: skip-file import logging import re import unittest import uuid import pytest from apache_beam.examples import wordcount_debugging from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import create_file from apache_beam.testing.test_utils import read_files_from_pattern @pytest.mark.examples_postcommit class WordCountDebuggingTest(unittest.TestCase): SAMPLE_TEXT = 'xx yy Flourish\n zz Flourish Flourish stomach\n aa\n bb cc dd' def get_results(self, temp_path): results = [] lines = read_files_from_pattern(temp_path + '.result*').splitlines() for line in lines: match = re.search(r'([A-Za-z]+): ([0-9]+)', line) if match is not None: results.append((match.group(1), int(match.group(2)))) return results def test_basics(self): test_pipeline = TestPipeline(is_integration_test=True) # Setup the files with expected content. temp_location = test_pipeline.get_option('temp_location') temp_path = '/'.join([temp_location, str(uuid.uuid4())]) input = create_file('/'.join([temp_path, 'input.txt']), self.SAMPLE_TEXT) extra_opts = {'input': input, 'output': '%s.result' % temp_path} expected_words = [('Flourish', 3), ('stomach', 1)] wordcount_debugging.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) # Parse result file and compare. results = self.get_results(temp_path) self.assertEqual(sorted(results), sorted(expected_words)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/wordcount_it_test.py0000644000175100017510000001510715206113032023630 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test for the wordcount example.""" # pytype: skip-file import logging import os import time import unittest import pytest from hamcrest.core.core.allof import all_of from apache_beam.examples import wordcount from apache_beam.internal.gcp import auth from apache_beam.testing.load_tests.load_test_metrics_utils import InfluxDBMetricsPublisherOptions from apache_beam.testing.load_tests.load_test_metrics_utils import MetricsReader from apache_beam.testing.pipeline_verifiers import FileChecksumMatcher from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import delete_files class WordCountIT(unittest.TestCase): # The default checksum is a SHA-1 hash generated from a sorted list of # lines read from expected output. This value corresponds to the default # input of WordCount example. DEFAULT_CHECKSUM = '33535a832b7db6d78389759577d4ff495980b9c0' @pytest.mark.it_postcommit @pytest.mark.it_validatescontainer def test_wordcount_it(self): self._run_wordcount_it(wordcount.run) @pytest.mark.it_postcommit @pytest.mark.sickbay_direct @pytest.mark.sickbay_spark @pytest.mark.sickbay_flink def test_wordcount_impersonation_it(self): """Tests impersonation on dataflow. For testing impersonation, we use three ingredients: - a principal to impersonate - a dataflow service account that only that principal is allowed to launch jobs as - a temp root that only the above two accounts have access to Jenkins and Dataflow workers both run as GCE default service account. So we remove that account from all the above. """ # Credentials need to be reset or this test will fail and credentials # from a previous test will be used. with auth._Credentials._credentials_lock: auth._Credentials._credentials_init = False try: ACCOUNT_TO_IMPERSONATE = ( 'allows-impersonation@apache-' 'beam-testing.iam.gserviceaccount.com') RUNNER_ACCOUNT = ( 'impersonation-dataflow-worker@' 'apache-beam-testing.iam.gserviceaccount.com') TEMP_DIR = 'gs://impersonation-test-bucket/temp-it' STAGING_LOCATION = 'gs://impersonation-test-bucket/staging-it' extra_options = { 'impersonate_service_account': ACCOUNT_TO_IMPERSONATE, 'service_account_email': RUNNER_ACCOUNT, 'temp_location': TEMP_DIR, 'staging_location': STAGING_LOCATION } self._run_wordcount_it(wordcount.run, **extra_options) finally: # Reset credentials for future tests. with auth._Credentials._credentials_lock: auth._Credentials._credentials_init = False @pytest.mark.it_validatescontainer def test_wordcount_it_with_prebuilt_sdk_container_local_docker(self): self._run_wordcount_it( wordcount.run, experiment='beam_fn_api', prebuild_sdk_container_engine='local_docker') @pytest.mark.it_validatescontainer def test_wordcount_it_with_prebuilt_sdk_container_cloud_build(self): self._run_wordcount_it( wordcount.run, experiment='beam_fn_api', prebuild_sdk_container_engine='cloud_build') def _run_wordcount_it(self, run_wordcount, **opts): test_pipeline = TestPipeline(is_integration_test=True) extra_opts = {} if (test_pipeline.get_option('machine_type') == 't2a-standard-1' and 'prebuild_sdk_container_engine' in opts): # TODO(https://github.com/apache/beam/issues/28340) pytest.skip('prebuild_sdk_container_engine not supported on ARM') # Set extra options to the pipeline for test purpose test_output = '/'.join([ test_pipeline.get_option('output'), str(int(time.time() * 1000)), 'results' ]) extra_opts['output'] = test_output test_input = test_pipeline.get_option('input') if test_input: extra_opts['input'] = test_input arg_sleep_secs = test_pipeline.get_option('sleep_secs') sleep_secs = int(arg_sleep_secs) if arg_sleep_secs is not None else None expect_checksum = ( test_pipeline.get_option('expect_checksum') or self.DEFAULT_CHECKSUM) pipeline_verifiers = [ PipelineStateMatcher(), FileChecksumMatcher( test_output + '*-of-*', expect_checksum, sleep_secs) ] extra_opts['on_success_matcher'] = all_of(*pipeline_verifiers) extra_opts.update(opts) # Register clean up before pipeline execution self.addCleanup(delete_files, [test_output + '*']) publish_to_bq = bool(test_pipeline.get_option('publish_to_big_query')) # Start measure time for performance test start_time = time.time() # Get pipeline options from command argument: --test-pipeline-options, # and start pipeline job by calling pipeline main function. run_wordcount( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False, ) end_time = time.time() run_time = end_time - start_time if publish_to_bq: self._publish_metrics(test_pipeline, run_time) def _publish_metrics(self, pipeline, metric_value): influx_options = InfluxDBMetricsPublisherOptions( pipeline.get_option('influx_measurement'), pipeline.get_option('influx_db_name'), pipeline.get_option('influx_hostname'), os.getenv('INFLUXDB_USER'), os.getenv('INFLUXDB_USER_PASSWORD'), ) metric_reader = MetricsReader( project_name=pipeline.get_option('project'), bq_table=pipeline.get_option('metrics_table'), bq_dataset=pipeline.get_option('metrics_dataset'), publish_to_bq=True, influxdb_options=influx_options, ) metric_reader.publish_values([( 'runtime', metric_value, )]) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/wordcount_minimal.py0000644000175100017510000001253215206113032023602 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A minimalist word-counting workflow that counts words in Shakespeare. This is the first in a series of successively more detailed 'word count' examples. Next, see the wordcount pipeline, then the wordcount_debugging pipeline, for more detailed examples that introduce additional concepts. Concepts: 1. Reading data from text files 2. Specifying 'inline' transforms 3. Counting a PCollection 4. Writing data to Cloud Storage as text files To execute this pipeline locally, first edit the code to specify the output location. Output location could be a local file path or an output prefix on GCS. (Only update the output location marked with the first CHANGE comment.) To execute this pipeline remotely, first edit the code to set your project ID, runner type, the staging location, the temp location, and the output location. The specified GCS bucket(s) must already exist. (Update all the places marked with a CHANGE comment.) Then, run the pipeline as described in the README. It will be deployed and run using the Google Cloud Dataflow Service. No args are required to run the pipeline. You can see the results in your output bucket in the GCS browser. """ # pytype: skip-file # beam-playground: # name: WordCountMinimal # description: An example that counts words in Shakespeare's works. # multifile: false # pipeline_options: --output output.txt # context_line: 120 # categories: # - IO # - Core Transforms # - Flatten # - Options # - Combiners # - Quickstart # complexity: MEDIUM # tags: # - count # - strings # - hellobeam import argparse import logging import re import apache_beam as beam from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions def main(argv=None, save_main_session=True): """Main entry point; defines and runs the wordcount pipeline.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='Input file to process.') parser.add_argument( '--output', dest='output', # CHANGE 1/6: (OPTIONAL) The Google Cloud Storage path is required # for outputting the results. default='gs://YOUR_OUTPUT_BUCKET/AND_OUTPUT_PREFIX', help='Output file to write results to.') # If you use DataflowRunner, below options can be passed: # CHANGE 2/6: (OPTIONAL) Change this to DataflowRunner to # run your pipeline on the Google Cloud Dataflow Service. # '--runner=DirectRunner', # CHANGE 3/6: (OPTIONAL) Your project ID is required in order to # run your pipeline on the Google Cloud Dataflow Service. # '--project=SET_YOUR_PROJECT_ID_HERE', # CHANGE 4/6: (OPTIONAL) The Google Cloud region (e.g. us-central1) # is required in order to run your pipeline on the Google Cloud # Dataflow Service. # '--region=SET_REGION_HERE', # CHANGE 5/6: Your Google Cloud Storage path is required for staging local # files. # '--staging_location=gs://YOUR_BUCKET_NAME/AND_STAGING_DIRECTORY', # CHANGE 6/6: Your Google Cloud Storage path is required for temporary # files. # '--temp_location=gs://YOUR_BUCKET_NAME/AND_TEMP_DIRECTORY', # '--job_name=your-wordcount-job', known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session with beam.Pipeline(options=pipeline_options) as p: # Read the text file[pattern] into a PCollection. lines = p | ReadFromText(known_args.input) # Count the occurrences of each word. counts = ( lines | 'Split' >> ( beam.FlatMap( lambda x: re.findall(r'[A-Za-z\']+', x)).with_output_types(str)) | 'PairWithOne' >> beam.Map(lambda x: (x, 1)) | 'GroupAndSum' >> beam.CombinePerKey(sum)) # Format the counts into a PCollection of strings. def format_result(word_count): (word, count) = word_count return '%s: %s' % (word, count) output = counts | 'Format' >> beam.Map(format_result) # Write the output using a "Write" transform that has side effects. # pylint: disable=expression-not-assigned output | WriteToText(known_args.output) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/wordcount_minimal_test.py0000644000175100017510000000467315206113032024650 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the minimal wordcount example.""" # pytype: skip-file import collections import logging import re import unittest import uuid import pytest from apache_beam.examples import wordcount_minimal from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import create_file from apache_beam.testing.test_utils import read_files_from_pattern @pytest.mark.examples_postcommit class WordCountMinimalTest(unittest.TestCase): """Unit test for wordcount_minimal example with direct runner.""" SAMPLE_TEXT = 'a b c a b a\n aa bb cc aa bb aa' def test_basics(self): test_pipeline = TestPipeline(is_integration_test=True) # Setup the files with expected content. temp_location = test_pipeline.get_option('temp_location') temp_path = '/'.join([temp_location, str(uuid.uuid4())]) input = create_file('/'.join([temp_path, 'input.txt']), self.SAMPLE_TEXT) extra_opts = {'input': input, 'output': '%s.result' % temp_path} expected_words = collections.defaultdict(int) for word in re.findall(r'\w+', self.SAMPLE_TEXT): expected_words[word] += 1 wordcount_minimal.main( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) # Parse result file and compare. results = [] lines = read_files_from_pattern(temp_path + '.result*').splitlines() for line in lines: match = re.search(r'([a-z]+): ([0-9]+)', line) if match is not None: results.append((match.group(1), int(match.group(2)))) self.assertEqual(sorted(results), sorted(expected_words.items())) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/wordcount_test.py0000644000175100017510000000464415206113032023140 0ustar00runnerrunner# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for the wordcount example.""" # pytype: skip-file import collections import logging import re import unittest import uuid import pytest from apache_beam.examples import wordcount from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_utils import create_file from apache_beam.testing.test_utils import read_files_from_pattern @pytest.mark.examples_postcommit class WordCountTest(unittest.TestCase): SAMPLE_TEXT = ( 'a b c a b a\nacento gráfico\nJuly 30, 2018\n\n aa bb cc aa bb aa') def test_basics(self): test_pipeline = TestPipeline(is_integration_test=True) # Setup the files with expected content. temp_location = test_pipeline.get_option('temp_location') temp_path = '/'.join([temp_location, str(uuid.uuid4())]) input = create_file('/'.join([temp_path, 'input.txt']), self.SAMPLE_TEXT) extra_opts = {'input': input, 'output': '%s.result' % temp_path} expected_words = collections.defaultdict(int) for word in re.findall(r'[\w\']+', self.SAMPLE_TEXT, re.UNICODE): expected_words[word] += 1 wordcount.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) # Parse result file and compare. results = [] lines = read_files_from_pattern(temp_path + '.result*').splitlines() for line in lines: match = re.search(r'(\S+): ([0-9]+)', line) if match is not None: results.append((match.group(1), int(match.group(2)))) self.assertEqual(sorted(results), sorted(expected_words.items())) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/wordcount_with_metrics.py0000644000175100017510000001231015206113032024647 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A word-counting workflow.""" # pytype: skip-file # beam-playground: # name: WordCountWithMetrics # description: A word-counting workflow with metrics. # multifile: false # default_example: true # pipeline_options: --output output.txt # context_line: 48 # categories: # - Combiners # - Options # - Metrics # - Quickstart # complexity: MEDIUM # tags: # - count # - metrics # - strings import argparse import logging import re import apache_beam as beam from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.metrics import Metrics from apache_beam.metrics.metric import MetricsFilter from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions class WordExtractingDoFn(beam.DoFn): """Parse each line of input text into words.""" def __init__(self): # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3. # super().__init__() beam.DoFn.__init__(self) self.words_counter = Metrics.counter(self.__class__, 'words') self.word_lengths_counter = Metrics.counter(self.__class__, 'word_lengths') self.word_lengths_dist = Metrics.distribution( self.__class__, 'word_len_dist') self.empty_line_counter = Metrics.counter(self.__class__, 'empty_lines') def process(self, element): """Returns an iterator over the words of this element. The element is a line of text. If the line is blank, note that, too. Args: element: the element being processed Returns: The processed element. """ text_line = element.strip() if not text_line: self.empty_line_counter.inc(1) words = re.findall(r'[\w\']+', text_line, re.UNICODE) for w in words: self.words_counter.inc() self.word_lengths_counter.inc(len(w)) self.word_lengths_dist.update(len(w)) return words def main(argv=None, save_main_session=True): """Main entry point; defines and runs the wordcount pipeline.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='Input file to process.') parser.add_argument( '--output', dest='output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args(argv) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session p = beam.Pipeline(options=pipeline_options) # Read the text file[pattern] into a PCollection. lines = p | 'read' >> ReadFromText(known_args.input) # Count the occurrences of each word. def count_ones(word_ones): (word, ones) = word_ones return (word, sum(ones)) counts = ( lines | 'split' >> (beam.ParDo(WordExtractingDoFn()).with_output_types(str)) | 'pair_with_one' >> beam.Map(lambda x: (x, 1)) | 'group' >> beam.GroupByKey() | 'count' >> beam.Map(count_ones)) # Format the counts into a PCollection of strings. def format_result(word_count): (word, count) = word_count return '%s: %d' % (word, count) output = counts | 'format' >> beam.Map(format_result) # Write the output using a "Write" transform that has side effects. # pylint: disable=expression-not-assigned output | 'write' >> WriteToText(known_args.output) result = p.run() result.wait_until_finish() # Do not query metrics when creating a template which doesn't run if (not hasattr(result, 'has_job') # direct runner or result.has_job): # not just a template creation empty_lines_filter = MetricsFilter().with_name('empty_lines') query_result = result.metrics().query(empty_lines_filter) if query_result['counters']: empty_lines_counter = query_result['counters'][0] logging.info('number of empty lines: %d', empty_lines_counter.result) word_lengths_filter = MetricsFilter().with_name('word_len_dist') query_result = result.metrics().query(word_lengths_filter) if query_result['distributions']: word_lengths_dist = query_result['distributions'][0] logging.info('average word length: %d', word_lengths_dist.result.mean) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/wordcount_xlang.py0000644000175100017510000000754615206113032023276 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A cross-language word-counting workflow.""" # pytype: skip-file import argparse import logging import re import subprocess import grpc import apache_beam as beam from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions # avoid possible conflict with job-server embedded expansion service at 8097 EXPANSION_SERVICE_PORT = '8096' EXPANSION_SERVICE_ADDR = 'localhost:%s' % EXPANSION_SERVICE_PORT class WordExtractingDoFn(beam.DoFn): """Parse each line of input text into words.""" def process(self, element): """Returns an iterator over the words of this element. The element is a line of text. If the line is blank, note that, too. Args: element: the element being processed Returns: The processed element. """ text_line = element.strip() return re.findall(r'[\w\']+', text_line) def build_pipeline(p, input_file, output_file): # Read the text file[pattern] into a PCollection. lines = p | 'read' >> ReadFromText(input_file) counts = ( lines | 'split' >> (beam.ParDo(WordExtractingDoFn()).with_output_types(str)) | 'count' >> beam.ExternalTransform( 'beam:transforms:xlang:count', None, EXPANSION_SERVICE_ADDR)) # Format the counts into a PCollection of strings. def format_result(word_count): (word, count) = word_count return '%s: %d' % (word, count) output = counts | 'format' >> beam.Map(format_result) # Write the output using a "Write" transform that has side effects. # pylint: disable=expression-not-assigned output | 'write' >> WriteToText(output_file) def main(): logging.getLogger().setLevel(logging.INFO) parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='Input file to process.') parser.add_argument( '--output', dest='output', required=True, help='Output file to write results to.') parser.add_argument( '--expansion_service_jar', dest='expansion_service_jar', required=True, help='Jar file for expansion service') known_args, pipeline_args = parser.parse_known_args() pipeline_options = PipelineOptions(pipeline_args) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options.view_as(SetupOptions).save_main_session = True try: server = subprocess.Popen([ 'java', '-jar', known_args.expansion_service_jar, EXPANSION_SERVICE_PORT ]) with grpc.insecure_channel(EXPANSION_SERVICE_ADDR) as channel: grpc.channel_ready_future(channel).result() with beam.Pipeline(options=pipeline_options) as p: # Preemptively start due to BEAM-6666. p.runner.create_job_service(pipeline_options) build_pipeline(p, known_args.input, known_args.output) finally: server.kill() if __name__ == '__main__': main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/examples/wordcount_xlang_sql.py0000644000175100017510000000770715206113032024154 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A word-counting workflow that uses the SQL transform. A Java version supported by Beam must be installed locally to run this pipeline. Additionally, Docker must also be available to run this pipeline locally. """ import argparse import logging import re from typing import NamedTuple import apache_beam as beam from apache_beam import coders from apache_beam.io import ReadFromText from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.portability import portable_runner from apache_beam.transforms.sql import SqlTransform # The input to SqlTransform must be a PCollection(s) of known schema. # One way to create such a PCollection is to produce a PCollection of # NamedTuple registered with the RowCoder. # # Here we create and register a simple NamedTuple with a single str typed # field named 'word' which we will use below. MyRow = NamedTuple('MyRow', [('word', str)]) coders.registry.register_coder(MyRow, coders.RowCoder) def run(p, input_file, output_file): #pylint: disable=expression-not-assigned ( p # Read the lines from a text file. | 'Read' >> ReadFromText(input_file) # Split the line into individual words. | 'Split' >> beam.FlatMap(lambda line: re.split(r'\W+', line)) # Map each word to an instance of MyRow. | 'ToRow' >> beam.Map(MyRow).with_output_types(MyRow) # SqlTransform yields a PCollection containing elements with attributes # based on the output of the query. | 'Sql!!' >> SqlTransform( """ SELECT word as key, COUNT(*) as `count` FROM PCOLLECTION GROUP BY word""") | 'Format' >> beam.Map(lambda row: '{}: {}'.format(row.key, row.count)) | 'Write' >> WriteToText(output_file)) def main(): logging.getLogger().setLevel(logging.INFO) parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', default='gs://dataflow-samples/shakespeare/kinglear.txt', help='Input file to process.') parser.add_argument( '--output', dest='output', required=True, help='Output file to write results to.') known_args, pipeline_args = parser.parse_known_args() pipeline_options = PipelineOptions(pipeline_args) # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). pipeline_options.view_as(SetupOptions).save_main_session = True with beam.Pipeline(options=pipeline_options) as p: if isinstance(p.runner, portable_runner.PortableRunner): # Preemptively start due to BEAM-6666. p.runner.create_job_service(pipeline_options) run(p, known_args.input, known_args.output) # Some more fun queries: # ------ # SELECT # word as key, # COUNT(*) as `count` # FROM PCOLLECTION # GROUP BY word # ORDER BY `count` DESC # LIMIT 100 # ------ # SELECT # len as key, # COUNT(*) as `count` # FROM ( # SELECT # LENGTH(word) AS len # FROM PCOLLECTION # ) # GROUP BY len if __name__ == '__main__': main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.2473938 apache_beam-2.74.0/apache_beam/internal/0000755000175100017510000000000015206113157017501 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/__init__.py0000644000175100017510000000152515206113032021605 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """For internal use only; no backwards-compatibility guarantees.""" ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.248394 apache_beam-2.74.0/apache_beam/internal/azure/0000755000175100017510000000000015206113157020627 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/azure/__init__.py0000644000175100017510000000152515206113032022733 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """For internal use only; no backwards-compatibility guarantees.""" ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/azure/auth.py0000644000175100017510000000527115206113032022137 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Azure credentials and authentication.""" # pytype: skip-file import logging import threading from apache_beam.options.pipeline_options import AzureOptions try: from azure.identity import DefaultAzureCredential _AZURE_AUTH_AVAILABLE = True except ImportError: _AZURE_AUTH_AVAILABLE = False _LOGGER = logging.getLogger(__name__) def get_service_credentials(pipeline_options): """For internal use only; no backwards-compatibility guarantees. Get credentials to access Azure services. Args: pipeline_options: Pipeline options, used in creating credentials like managed identity credentials. Returns: A ``azure.identity.*Credential`` object or None if credentials not found. Returned object is thread-safe. """ return _Credentials.get_service_credentials(pipeline_options) class _Credentials(object): _credentials_lock = threading.Lock() _credentials_init = False _credentials = None @classmethod def get_service_credentials(cls, pipeline_options): with cls._credentials_lock: if cls._credentials_init: return cls._credentials cls._credentials = cls._get_service_credentials(pipeline_options) cls._credentials_init = True return cls._credentials @staticmethod def _get_service_credentials(pipeline_options): if not _AZURE_AUTH_AVAILABLE: _LOGGER.warning( 'Unable to find default credentials because the azure.identity ' 'library is not available. Install the azure.identity library to use ' 'Azure default credentials.') return None try: credentials = DefaultAzureCredential( managed_identity_client_id=pipeline_options.view_as(AzureOptions)\ .azure_managed_identity_client_id) _LOGGER.debug('Connecting using Azure Default Credentials.') return credentials except Exception as e: _LOGGER.warning('Unable to find Azure credentials to use: %s\n', e) return None ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.249394 apache_beam-2.74.0/apache_beam/internal/cloudpickle/0000755000175100017510000000000015206113157021777 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/cloudpickle/__init__.py0000644000175100017510000000046415206113032024104 0ustar00runnerrunnerfrom . import cloudpickle from .cloudpickle import * # noqa __doc__ = cloudpickle.__doc__ __version__ = "3.1.1" __all__ = [ # noqa "__version__", "Pickler", "CloudPickler", "dumps", "loads", "dump", "load", "register_pickle_by_value", "unregister_pickle_by_value", ] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/cloudpickle/cloudpickle.py0000644000175100017510000017431015206113032024645 0ustar00runnerrunner"""Pickler class to extend the standard pickle.Pickler functionality The main objective is to make it natural to perform distributed computing on clusters (such as PySpark, Dask, Ray...) with interactively defined code (functions, classes, ...) written in notebooks or console. In particular this pickler adds the following features: - serialize interactively-defined or locally-defined functions, classes, enums, typevars, lambdas and nested functions to compiled byte code; - deal with some other non-serializable objects in an ad-hoc manner where applicable. This pickler is therefore meant to be used for the communication between short lived Python processes running the same version of Python and libraries. In particular, it is not meant to be used for long term storage of Python objects. It does not include an unpickler, as standard Python unpickling suffices. This module was extracted from the `cloud` package, developed by `PiCloud, Inc. `_. Copyright (c) 2012-now, CloudPickle developers and contributors. Copyright (c) 2012, Regents of the University of California. Copyright (c) 2009 `PiCloud, Inc. `_. All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. * Neither the name of the University of California, Berkeley nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. """ # mypy: ignore-errors # pylint: skip-file import _collections_abc from collections import ChainMap, OrderedDict import abc import builtins import copyreg import dataclasses import dis from enum import Enum import functools import hashlib import io import itertools import logging import opcode import os import pickle from pickle import _getattribute as _pickle_getattribute import platform import struct import sys import threading import types import typing import uuid import warnings import weakref # The following import is required to be imported in the cloudpickle # namespace to be able to load pickle files generated with older versions of # cloudpickle. See: tests/test_backward_compat.py from types import CellType # noqa: F401 # cloudpickle is meant for inter process communication: we expect all # communicating processes to run the same Python version hence we favor # communication speed over compatibility: DEFAULT_PROTOCOL = pickle.HIGHEST_PROTOCOL # Names of modules whose resources should be treated as dynamic. _PICKLE_BY_VALUE_MODULES = set() # Track the provenance of reconstructed dynamic classes to make it possible to # reconstruct instances from the matching singleton class definition when # appropriate and preserve the usual "isinstance" semantics of Python objects. _DYNAMIC_CLASS_TRACKER_BY_CLASS = weakref.WeakKeyDictionary() _DYNAMIC_CLASS_TRACKER_BY_ID = weakref.WeakValueDictionary() _DYNAMIC_CLASS_STATE_TRACKER_BY_CLASS = weakref.WeakKeyDictionary() _DYNAMIC_CLASS_TRACKER_LOCK = threading.RLock() PYPY = platform.python_implementation() == "PyPy" def uuid_generator(_): return uuid.uuid4().hex @dataclasses.dataclass class GetCodeObjectParams: """Parameters for enabling stable code object pickling. Attributes: get_code_object_identifier: This function should take a Python callable (e.g., a function or lambda) and return a unique string identifier. This identifier represents a stable "path" to locate the code object within a module, rather than depending on the exact bytecode. If no stable identifier can be generated, it should return None. (See code_object_pickler.get_code_object_identifier). get_code_from_identifier: This function takes an identifier string generated by get_code_object_identifier and returns the corresponding types.CodeType object from the currently loaded modules. It should raise an AttributeError or ValueError if the code object cannot be found or reconstructed from the identifier. (See code_object_pickler.get_code_from_identifier). """ get_code_object_identifier: typing.Optional[callable] get_code_from_identifier: typing.Optional[callable] @dataclasses.dataclass class CloudPickleConfig: """Configuration for cloudpickle behavior. This class controls various aspects of how cloudpickle serializes objects. Attributes: id_generator: Callable that generates unique identifiers for dynamic types. Controls isinstance semantics preservation. If None, disables type tracking and isinstance relationships are not preserved across pickle/unpickle cycles. If callable, generates unique IDs to maintain object identity. Default: uuid_generator (generates UUID hex strings). skip_reset_dynamic_type_state: Whether to skip resetting state when reconstructing dynamic types. If True, skips state reset for already-reconstructed types. filepath_interceptor: Used to modify filepaths in `co_filename` and function.__globals__['__file__']. get_code_object_params: An optional `GetCodeObjectParams` instance. If provided, cloudpickle will use identifiers derived from code location when pickling dynamic functions (e.g. lambdas). Enabling this setting results in pickled payloads becoming more stable to code changes: when a particular lambda function is slightly modified but the location of the function in the codebase has not changed, the pickled representation might stay the same. pickle_main_by_ref: An optional boolean. If provided, cloudpickle will pickle main by reference instead of by value. """ id_generator: typing.Optional[callable] = uuid_generator skip_reset_dynamic_type_state: bool = False filepath_interceptor: typing.Optional[callable] = None get_code_object_params: typing.Optional[GetCodeObjectParams] = None pickle_main_by_ref: bool = False DEFAULT_CONFIG = CloudPickleConfig() _GENERATING_SENTINEL = object() builtin_code_type = None if PYPY: # builtin-code objects only exist in pypy builtin_code_type = type(float.__new__.__code__) _extract_code_globals_cache = weakref.WeakKeyDictionary() def _get_or_create_tracker_id(class_def, id_generator): with _DYNAMIC_CLASS_TRACKER_LOCK: class_tracker_id = _DYNAMIC_CLASS_TRACKER_BY_CLASS.get(class_def) if class_tracker_id is _GENERATING_SENTINEL and id_generator: raise RuntimeError( f"Recursive ID generation detected for {class_def}. " f"The id_generator cannot recursively request an ID for the same class." ) if class_tracker_id is None and id_generator is not None: _DYNAMIC_CLASS_TRACKER_BY_CLASS[class_def] = _GENERATING_SENTINEL try: class_tracker_id = id_generator(class_def) _DYNAMIC_CLASS_TRACKER_BY_CLASS[class_def] = class_tracker_id _DYNAMIC_CLASS_TRACKER_BY_ID[class_tracker_id] = class_def except Exception: _DYNAMIC_CLASS_TRACKER_BY_CLASS.pop(class_def, None) raise return class_tracker_id def _lookup_class_or_track(class_tracker_id, class_def): if class_tracker_id is not None: with _DYNAMIC_CLASS_TRACKER_LOCK: class_def = _DYNAMIC_CLASS_TRACKER_BY_ID.setdefault( class_tracker_id, class_def) _DYNAMIC_CLASS_TRACKER_BY_CLASS[class_def] = class_tracker_id return class_def def register_pickle_by_value(module): """Register a module to make its functions and classes picklable by value. By default, functions and classes that are attributes of an importable module are to be pickled by reference, that is relying on re-importing the attribute from the module at load time. If `register_pickle_by_value(module)` is called, all its functions and classes are subsequently to be pickled by value, meaning that they can be loaded in Python processes where the module is not importable. This is especially useful when developing a module in a distributed execution environment: restarting the client Python process with the new source code is enough: there is no need to re-install the new version of the module on all the worker nodes nor to restart the workers. Note: this feature is considered experimental. See the cloudpickle README.md file for more details and limitations. """ if not isinstance(module, types.ModuleType): raise ValueError( f"Input should be a module object, got {str(module)} instead") # In the future, cloudpickle may need a way to access any module registered # for pickling by value in order to introspect relative imports inside # functions pickled by value. (see # https://github.com/cloudpipe/cloudpickle/pull/417#issuecomment-873684633). # This access can be ensured by checking that module is present in # sys.modules at registering time and assuming that it will still be in # there when accessed during pickling. Another alternative would be to # store a weakref to the module. Even though cloudpickle does not implement # this introspection yet, in order to avoid a possible breaking change # later, we still enforce the presence of module inside sys.modules. if module.__name__ not in sys.modules: raise ValueError( f"{module} was not imported correctly, have you used an " "`import` statement to access it?") _PICKLE_BY_VALUE_MODULES.add(module.__name__) def unregister_pickle_by_value(module): """Unregister that the input module should be pickled by value.""" if not isinstance(module, types.ModuleType): raise ValueError( f"Input should be a module object, got {str(module)} instead") if module.__name__ not in _PICKLE_BY_VALUE_MODULES: raise ValueError(f"{module} is not registered for pickle by value") else: _PICKLE_BY_VALUE_MODULES.remove(module.__name__) def list_registry_pickle_by_value(): return _PICKLE_BY_VALUE_MODULES.copy() def _is_registered_pickle_by_value(module): module_name = module.__name__ if module_name in _PICKLE_BY_VALUE_MODULES: return True while True: parent_name = module_name.rsplit(".", 1)[0] if parent_name == module_name: break if parent_name in _PICKLE_BY_VALUE_MODULES: return True module_name = parent_name return False if sys.version_info >= (3, 14): def _getattribute(obj, name): return _pickle_getattribute(obj, name.split('.')) else: def _getattribute(obj, name): return _pickle_getattribute(obj, name)[0] def _whichmodule(obj, name): """Find the module an object belongs to. This function differs from ``pickle.whichmodule`` in two ways: - it does not mangle the cases where obj's module is __main__ and obj was not found in any module. - Errors arising during module introspection are ignored, as those errors are considered unwanted side effects. """ module_name = getattr(obj, "__module__", None) if module_name is not None: return module_name # Protect the iteration by using a copy of sys.modules against dynamic # modules that trigger imports of other modules upon calls to getattr or # other threads importing at the same time. for module_name, module in sys.modules.copy().items(): # Some modules such as coverage can inject non-module objects inside # sys.modules if (module_name == "__main__" or module_name == "__mp_main__" or module is None or not isinstance(module, types.ModuleType)): continue try: if _getattribute(module, name) is obj: return module_name except Exception: pass return None def _should_pickle_by_reference(obj, name=None, config=DEFAULT_CONFIG): """Test whether an function or a class should be pickled by reference Pickling by reference means by that the object (typically a function or a class) is an attribute of a module that is assumed to be importable in the target Python environment. Loading will therefore rely on importing the module and then calling `getattr` on it to access the function or class. Pickling by reference is the only option to pickle functions and classes in the standard library. In cloudpickle the alternative option is to pickle by value (for instance for interactively or locally defined functions and classes or for attributes of modules that have been explicitly registered to be pickled by value. """ if isinstance(obj, types.FunctionType) or issubclass(type(obj), type): module_and_name = _lookup_module_and_qualname(obj, name=name, config=config) if module_and_name is None: return False module, name = module_and_name return not _is_registered_pickle_by_value(module) elif isinstance(obj, types.ModuleType): # We assume that sys.modules is primarily used as a cache mechanism for # the Python import machinery. Checking if a module has been added in # is sys.modules therefore a cheap and simple heuristic to tell us # whether we can assume that a given module could be imported by name # in another Python process. if _is_registered_pickle_by_value(obj): return False return obj.__name__ in sys.modules else: raise TypeError( "cannot check importability of {} instances".format(type(obj).__name__)) def _lookup_module_and_qualname(obj, name=None, config=DEFAULT_CONFIG): if name is None: name = getattr(obj, "__qualname__", None) if name is None: # pragma: no cover # This used to be needed for Python 2.7 support but is probably not # needed anymore. However we keep the __name__ introspection in case # users of cloudpickle rely on this old behavior for unknown reasons. name = getattr(obj, "__name__", None) module_name = _whichmodule(obj, name) if module_name is None: # In this case, obj.__module__ is None AND obj was not found in any # imported module. obj is thus treated as dynamic. return None if module_name == "__main__" and not config.pickle_main_by_ref: return None # Note: if module_name is in sys.modules, the corresponding module is # assumed importable at unpickling time. See #357 module = sys.modules.get(module_name, None) if module is None: # The main reason why obj's module would not be imported is that this # module has been dynamically created, using for example # types.ModuleType. The other possibility is that module was removed # from sys.modules after obj was created/imported. But this case is not # supported, as the standard pickle does not support it either. return None try: obj2 = _getattribute(module, name) except AttributeError: # obj was not found inside the module it points to return None if obj2 is not obj: return None return module, name def _extract_code_globals(co): """Find all globals names read or written to by codeblock co.""" out_names = _extract_code_globals_cache.get(co) if out_names is None: # We use a dict with None values instead of a set to get a # deterministic order and avoid introducing non-deterministic pickle # bytes as a results. out_names = {name: None for name in _walk_global_ops(co)} # Declaring a function inside another one using the "def ..." syntax # generates a constant code object corresponding to the one of the # nested function's As the nested function may itself need global # variables, we need to introspect its code, extract its globals, (look # for code object in it's co_consts attribute..) and add the result to # code_globals if co.co_consts: for const in co.co_consts: if isinstance(const, types.CodeType): out_names.update(_extract_code_globals(const)) _extract_code_globals_cache[co] = out_names return out_names def _find_imported_submodules(code, top_level_dependencies): """Find currently imported submodules used by a function. Submodules used by a function need to be detected and referenced for the function to work correctly at depickling time. Because submodules can be referenced as attribute of their parent package (``package.submodule``), we need a special introspection technique that does not rely on GLOBAL-related opcodes to find references of them in a code object. Example: ``` import concurrent.futures import cloudpickle def func(): x = concurrent.futures.ThreadPoolExecutor if __name__ == '__main__': cloudpickle.dumps(func) ``` The globals extracted by cloudpickle in the function's state include the concurrent package, but not its submodule (here, concurrent.futures), which is the module used by func. Find_imported_submodules will detect the usage of concurrent.futures. Saving this module alongside with func will ensure that calling func once depickled does not fail due to concurrent.futures not being imported """ subimports = [] # check if any known dependency is an imported package for x in top_level_dependencies: if (isinstance(x, types.ModuleType) and hasattr(x, "__package__") and x.__package__): # check if the package has any currently loaded sub-imports prefix = x.__name__ + "." # A concurrent thread could mutate sys.modules, # make sure we iterate over a copy to avoid exceptions for name in list(sys.modules): # Older versions of pytest will add a "None" module to # sys.modules. if name is not None and name.startswith(prefix): # check whether the function can address the sub-module tokens = set(name[len(prefix):].split(".")) if not tokens - set(code.co_names): subimports.append(sys.modules[name]) return subimports def get_relative_path(path): """Returns the path of a filename relative to the longest matching directory in sys.path. Args: path: The path to the file. """ abs_path = os.path.abspath(path) longest_match = "" for dir_path in sys.path: if not dir_path.endswith(os.path.sep): dir_path += os.path.sep if abs_path.startswith(dir_path) and len(dir_path) > len(longest_match): longest_match = dir_path if not longest_match: return path return os.path.relpath(abs_path, longest_match) # relevant opcodes STORE_GLOBAL = opcode.opmap["STORE_GLOBAL"] DELETE_GLOBAL = opcode.opmap["DELETE_GLOBAL"] LOAD_GLOBAL = opcode.opmap["LOAD_GLOBAL"] GLOBAL_OPS = (STORE_GLOBAL, DELETE_GLOBAL, LOAD_GLOBAL) HAVE_ARGUMENT = dis.HAVE_ARGUMENT EXTENDED_ARG = dis.EXTENDED_ARG _BUILTIN_TYPE_NAMES = {} for k, v in types.__dict__.items(): if type(v) is type: _BUILTIN_TYPE_NAMES[v] = k def _builtin_type(name): if name == "ClassType": # pragma: no cover # Backward compat to load pickle files generated with cloudpickle # < 1.3 even if loading pickle files from older versions is not # officially supported. return type return getattr(types, name) def _walk_global_ops(code): """Yield referenced name for global-referencing instructions in code.""" for instr in dis.get_instructions(code): op = instr.opcode if op in GLOBAL_OPS: yield instr.argval def _extract_class_dict(cls): """Retrieve a copy of the dict of a class without the inherited method.""" # Hack to circumvent non-predictable memoization caused by string interning. # See the inline comment in _class_setstate for details. clsdict = {"".join(k): cls.__dict__[k] for k in sorted(cls.__dict__)} if len(cls.__bases__) == 1: inherited_dict = cls.__bases__[0].__dict__ else: inherited_dict = {} for base in reversed(cls.__bases__): inherited_dict.update(base.__dict__) to_remove = [] for name, value in clsdict.items(): try: base_value = inherited_dict[name] if value is base_value: to_remove.append(name) except KeyError: pass for name in to_remove: clsdict.pop(name) return clsdict def is_tornado_coroutine(func): """Return whether `func` is a Tornado coroutine function. Running coroutines are not supported. """ warnings.warn( "is_tornado_coroutine is deprecated in cloudpickle 3.0 and will be " "removed in cloudpickle 4.0. Use tornado.gen.is_coroutine_function " "directly instead.", category=DeprecationWarning, ) if "tornado.gen" not in sys.modules: return False gen = sys.modules["tornado.gen"] if not hasattr(gen, "is_coroutine_function"): # Tornado version is too old return False return gen.is_coroutine_function(func) def subimport(name): # We cannot do simply: `return __import__(name)`: Indeed, if ``name`` is # the name of a submodule, __import__ will return the top-level root module # of this submodule. For instance, __import__('os.path') returns the `os` # module. __import__(name) return sys.modules[name] def dynamic_subimport(name, vars): mod = types.ModuleType(name) mod.__dict__.update(vars) mod.__dict__["__builtins__"] = builtins.__dict__ return mod def _get_cell_contents(cell): try: return cell.cell_contents except ValueError: # Handle empty cells explicitly with a sentinel value. return _empty_cell_value def instance(cls): """Create a new instance of a class. Parameters ---------- cls : type The class to create an instance of. Returns ------- instance : cls A new instance of ``cls``. """ return cls() @instance class _empty_cell_value: """Sentinel for empty closures.""" @classmethod def __reduce__(cls): return cls.__name__ def _make_function(code, globals, name, argdefs, closure): # Setting __builtins__ in globals is needed for nogil CPython. globals["__builtins__"] = __builtins__ return types.FunctionType(code, globals, name, argdefs, closure) def _make_function_from_identifier( get_code_from_identifier, code_path, globals, name, argdefs): fcode = get_code_from_identifier(code_path) expected_closure_len = len(fcode.co_freevars) closure = tuple(types.CellType() for _ in range(expected_closure_len)) return _make_function(fcode, globals, name, argdefs, closure) def _make_empty_cell(): if False: # trick the compiler into creating an empty cell in our lambda cell = None raise AssertionError("this route should not be executed") return (lambda: cell).__closure__[0] def _make_cell(value=_empty_cell_value): cell = _make_empty_cell() if value is not _empty_cell_value: cell.cell_contents = value return cell def _make_skeleton_class( type_constructor, name, bases, type_kwargs, class_tracker_id, extra): """Build dynamic class with an empty __dict__ to be filled once memoized If class_tracker_id is not None, try to lookup an existing class definition matching that id. If none is found, track a newly reconstructed class definition under that id so that other instances stemming from the same class id will also reuse this class definition. The "extra" variable is meant to be a dict (or None) that can be used for forward compatibility shall the need arise. """ # We need to intern the keys of the type_kwargs dict to avoid having # different pickles for the same dynamic class depending on whether it was # dynamically created or reconstructed from a pickled stream. type_kwargs = {sys.intern(k): v for k, v in type_kwargs.items()} skeleton_class = types.new_class( name, bases, {"metaclass": type_constructor}, lambda ns: ns.update(type_kwargs)) return _lookup_class_or_track(class_tracker_id, skeleton_class) def _make_skeleton_enum( bases, name, qualname, members, module, class_tracker_id, extra): """Build dynamic enum with an empty __dict__ to be filled once memoized The creation of the enum class is inspired by the code of EnumMeta._create_. If class_tracker_id is not None, try to lookup an existing enum definition matching that id. If none is found, track a newly reconstructed enum definition under that id so that other instances stemming from the same class id will also reuse this enum definition. The "extra" variable is meant to be a dict (or None) that can be used for forward compatibility shall the need arise. """ # enums always inherit from their base Enum class at the last position in # the list of base classes: enum_base = bases[-1] metacls = enum_base.__class__ classdict = metacls.__prepare__(name, bases) for member_name, member_value in members.items(): classdict[member_name] = member_value enum_class = metacls.__new__(metacls, name, bases, classdict) enum_class.__module__ = module enum_class.__qualname__ = qualname return _lookup_class_or_track(class_tracker_id, enum_class) def _make_typevar( name, bound, constraints, covariant, contravariant, class_tracker_id): tv = typing.TypeVar( name, *constraints, bound=bound, covariant=covariant, contravariant=contravariant, ) return _lookup_class_or_track(class_tracker_id, tv) def _decompose_typevar(obj, config: CloudPickleConfig): return ( obj.__name__, obj.__bound__, obj.__constraints__, obj.__covariant__, obj.__contravariant__, _get_or_create_tracker_id(obj, config.id_generator), ) def _typevar_reduce(obj, config: CloudPickleConfig): # TypeVar instances require the module information hence why we # are not using the _should_pickle_by_reference directly module_and_name = _lookup_module_and_qualname( obj, name=obj.__name__, config=config) if module_and_name is None: return (_make_typevar, _decompose_typevar(obj, config)) elif _is_registered_pickle_by_value(module_and_name[0]): return (_make_typevar, _decompose_typevar(obj, config)) return (getattr, module_and_name) def _get_bases(typ): if "__orig_bases__" in getattr(typ, "__dict__", {}): # For generic types (see PEP 560) # Note that simply checking `hasattr(typ, '__orig_bases__')` is not # correct. Subclasses of a fully-parameterized generic class does not # have `__orig_bases__` defined, but `hasattr(typ, '__orig_bases__')` # will return True because it's defined in the base class. bases_attr = "__orig_bases__" else: # For regular class objects bases_attr = "__bases__" return getattr(typ, bases_attr) def _make_dict_keys(obj, is_ordered=False): if is_ordered: return OrderedDict.fromkeys(obj).keys() else: return dict.fromkeys(obj).keys() def _make_dict_values(obj, is_ordered=False): if is_ordered: return OrderedDict((i, _) for i, _ in enumerate(obj)).values() else: return {i: _ for i, _ in enumerate(obj)}.values() def _make_dict_items(obj, is_ordered=False): if is_ordered: return OrderedDict(obj).items() else: return obj.items() # COLLECTION OF OBJECTS __getnewargs__-LIKE METHODS # ------------------------------------------------- def _class_getnewargs(obj, config: CloudPickleConfig): type_kwargs = {} if "__module__" in obj.__dict__: type_kwargs["__module__"] = obj.__module__ __dict__ = obj.__dict__.get("__dict__", None) if isinstance(__dict__, property): type_kwargs["__dict__"] = __dict__ return ( type(obj), obj.__name__, _get_bases(obj), type_kwargs, _get_or_create_tracker_id(obj, config.id_generator), None, ) def _enum_getnewargs(obj, config: CloudPickleConfig): members = {e.name: e.value for e in obj} return ( obj.__bases__, obj.__name__, obj.__qualname__, members, obj.__module__, _get_or_create_tracker_id(obj, config.id_generator), None, ) # COLLECTION OF OBJECTS RECONSTRUCTORS # ------------------------------------ def _file_reconstructor(retval): return retval # COLLECTION OF OBJECTS STATE GETTERS # ----------------------------------- def _function_getstate(func): # - Put func's dynamic attributes (stored in func.__dict__) in state. These # attributes will be restored at unpickling time using # f.__dict__.update(state) # - Put func's members into slotstate. Such attributes will be restored at # unpickling time by iterating over slotstate and calling setattr(func, # slotname, slotvalue) slotstate = { # Hack to circumvent non-predictable memoization caused by string interning. # See the inline comment in _class_setstate for details. "__name__": "".join(func.__name__), "__qualname__": "".join(func.__qualname__), "__annotations__": func.__annotations__, "__kwdefaults__": func.__kwdefaults__, "__defaults__": func.__defaults__, "__module__": func.__module__, "__doc__": func.__doc__, "__closure__": func.__closure__, } f_globals_ref = _extract_code_globals(func.__code__) f_globals = { k: func.__globals__[k] for k in f_globals_ref if k in func.__globals__ } if func.__closure__ is not None: closure_values = list(map(_get_cell_contents, func.__closure__)) else: closure_values = () # Extract currently-imported submodules used by func. Storing these modules # in a smoke _cloudpickle_subimports attribute of the object's state will # trigger the side effect of importing these modules at unpickling time # (which is necessary for func to work correctly once depickled) slotstate["_cloudpickle_submodules"] = _find_imported_submodules( func.__code__, itertools.chain(f_globals.values(), closure_values)) slotstate["__globals__"] = f_globals # Hack to circumvent non-predictable memoization caused by string interning. # See the inline comment in _class_setstate for details. state = {"".join(k): v for k, v in func.__dict__.items()} return state, slotstate def _class_getstate(obj): clsdict = _extract_class_dict(obj) clsdict.pop("__weakref__", None) if issubclass(type(obj), abc.ABCMeta): # If obj is an instance of an ABCMeta subclass, don't pickle the # cache/negative caches populated during isinstance/issubclass # checks, but pickle the list of registered subclasses of obj. clsdict.pop("_abc_cache", None) clsdict.pop("_abc_negative_cache", None) clsdict.pop("_abc_negative_cache_version", None) registry = clsdict.pop("_abc_registry", None) if registry is None: # The abc caches and registered subclasses of a # class are bundled into the single _abc_impl attribute clsdict.pop("_abc_impl", None) (registry, _, _, _) = abc._get_dump(obj) clsdict["_abc_impl"] = [ subclass_weakref() for subclass_weakref in registry ] else: # In the above if clause, registry is a set of weakrefs -- in # this case, registry is a WeakSet clsdict["_abc_impl"] = [type_ for type_ in registry] if "__slots__" in clsdict: # pickle string length optimization: member descriptors of obj are # created automatically from obj's __slots__ attribute, no need to # save them in obj's state if isinstance(obj.__slots__, str): clsdict.pop(obj.__slots__) else: for k in obj.__slots__: clsdict.pop(k, None) clsdict.pop("__dict__", None) # unpicklable property object return (clsdict, {}) def _enum_getstate(obj): clsdict, slotstate = _class_getstate(obj) members = {e.name: e.value for e in obj} # Cleanup the clsdict that will be passed to _make_skeleton_enum: # Those attributes are already handled by the metaclass. for attrname in [ "_generate_next_value_", "_member_names_", "_member_map_", "_member_type_", "_value2member_map_", ]: clsdict.pop(attrname, None) for member in members: clsdict.pop(member) # Special handling of Enum subclasses return clsdict, slotstate # COLLECTIONS OF OBJECTS REDUCERS # ------------------------------- # A reducer is a function taking a single argument (obj), and that returns a # tuple with all the necessary data to re-construct obj. Apart from a few # exceptions (list, dict, bytes, int, etc.), a reducer is necessary to # correctly pickle an object. # While many built-in objects (Exceptions objects, instances of the "object" # class, etc), are shipped with their own built-in reducer (invoked using # obj.__reduce__), some do not. The following methods were created to "fill # these holes". def _code_reduce(obj, config: CloudPickleConfig): """code object reducer.""" # If you are not sure about the order of arguments, take a look at help # of the specific type from types, for example: # >>> from types import CodeType # >>> help(CodeType) # Hack to circumvent non-predictable memoization caused by string interning. # See the inline comment in _class_setstate for details. co_name = "".join(obj.co_name) # Create shallow copies of these tuple to make cloudpickle payload deterministic. # When creating a code object during load, copies of these four tuples are # created, while in the main process, these tuples can be shared. # By always creating copies, we make sure the resulting payload is deterministic. co_names = tuple(name for name in obj.co_names) co_varnames = tuple(name for name in obj.co_varnames) co_freevars = tuple(name for name in obj.co_freevars) co_cellvars = tuple(name for name in obj.co_cellvars) co_filename = obj.co_filename if (config and config.filepath_interceptor): co_filename = config.filepath_interceptor(co_filename) if hasattr(obj, "co_exceptiontable"): # Python 3.11 and later: there are some new attributes # related to the enhanced exceptions. args = ( obj.co_argcount, obj.co_posonlyargcount, obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize, obj.co_flags, obj.co_code, obj.co_consts, co_names, co_varnames, co_filename, co_name, obj.co_qualname, obj.co_firstlineno, obj.co_linetable, obj.co_exceptiontable, co_freevars, co_cellvars, ) elif hasattr(obj, "co_linetable"): # Python 3.10 and later: obj.co_lnotab is deprecated and constructor # expects obj.co_linetable instead. args = ( obj.co_argcount, obj.co_posonlyargcount, obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize, obj.co_flags, obj.co_code, obj.co_consts, co_names, co_varnames, co_filename, co_name, obj.co_firstlineno, obj.co_linetable, co_freevars, co_cellvars, ) elif hasattr(obj, "co_nmeta"): # pragma: no cover # "nogil" Python: modified attributes from 3.9 args = ( obj.co_argcount, obj.co_posonlyargcount, obj.co_kwonlyargcount, obj.co_nlocals, obj.co_framesize, obj.co_ndefaultargs, obj.co_nmeta, obj.co_flags, obj.co_code, obj.co_consts, co_varnames, co_filename, co_name, obj.co_firstlineno, obj.co_lnotab, obj.co_exc_handlers, obj.co_jump_table, co_freevars, co_cellvars, obj.co_free2reg, obj.co_cell2reg, ) else: # Backward compat for 3.8 and 3.9 args = ( obj.co_argcount, obj.co_posonlyargcount, obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize, obj.co_flags, obj.co_code, obj.co_consts, co_names, co_varnames, co_filename, co_name, obj.co_firstlineno, obj.co_lnotab, co_freevars, co_cellvars, ) return types.CodeType, args def _cell_reduce(obj): """Cell (containing values of a function's free variables) reducer.""" try: obj.cell_contents except ValueError: # cell is empty return _make_empty_cell, () else: return _make_cell, (obj.cell_contents, ) def _classmethod_reduce(obj): orig_func = obj.__func__ return type(obj), (orig_func, ) def _file_reduce(obj): """Save a file.""" import io if not hasattr(obj, "name") or not hasattr(obj, "mode"): raise pickle.PicklingError( "Cannot pickle files that do not map to an actual file") if obj is sys.stdout: return getattr, (sys, "stdout") if obj is sys.stderr: return getattr, (sys, "stderr") if obj is sys.stdin: raise pickle.PicklingError("Cannot pickle standard input") if obj.closed: raise pickle.PicklingError("Cannot pickle closed files") if hasattr(obj, "isatty") and obj.isatty(): raise pickle.PicklingError("Cannot pickle files that map to tty objects") if "r" not in obj.mode and "+" not in obj.mode: raise pickle.PicklingError( "Cannot pickle files that are not opened for reading: %s" % obj.mode) name = obj.name retval = io.StringIO() try: # Read the whole file curloc = obj.tell() obj.seek(0) contents = obj.read() obj.seek(curloc) except OSError as e: raise pickle.PicklingError( "Cannot pickle file %s as it cannot be read" % name) from e retval.write(contents) retval.seek(curloc) retval.name = name return _file_reconstructor, (retval, ) def _getset_descriptor_reduce(obj): return getattr, (obj.__objclass__, obj.__name__) def _mappingproxy_reduce(obj): return types.MappingProxyType, (dict(obj), ) def _memoryview_reduce(obj): return bytes, (obj.tobytes(), ) def _module_reduce(obj): if _should_pickle_by_reference(obj): return subimport, (obj.__name__, ) else: # Some external libraries can populate the "__builtins__" entry of a # module's `__dict__` with unpicklable objects (see #316). For that # reason, we do not attempt to pickle the "__builtins__" entry, and # restore a default value for it at unpickling time. state = obj.__dict__.copy() state.pop("__builtins__", None) return dynamic_subimport, (obj.__name__, state) def _method_reduce(obj): return (types.MethodType, (obj.__func__, obj.__self__)) def _logger_reduce(obj): return logging.getLogger, (obj.name, ) def _root_logger_reduce(obj): return logging.getLogger, () def _property_reduce(obj): return property, (obj.fget, obj.fset, obj.fdel, obj.__doc__) def _weakset_reduce(obj): return weakref.WeakSet, (list(obj), ) def _dynamic_class_reduce(obj, config: CloudPickleConfig): """Save a class that can't be referenced as a module attribute. This method is used to serialize classes that are defined inside functions, or that otherwise can't be serialized as attribute lookups from importable modules. """ if Enum is not None and issubclass(obj, Enum): return ( _make_skeleton_enum, _enum_getnewargs(obj, config), _enum_getstate(obj), None, None, functools.partial( _class_setstate, skip_reset_dynamic_type_state=config.skip_reset_dynamic_type_state), ) else: return ( _make_skeleton_class, _class_getnewargs(obj, config), _class_getstate(obj), None, None, functools.partial( _class_setstate, skip_reset_dynamic_type_state=config.skip_reset_dynamic_type_state), ) def _class_reduce(obj, config: CloudPickleConfig): """Select the reducer depending on the dynamic nature of the class obj.""" if obj is type(None): # noqa return type, (None, ) elif obj is type(Ellipsis): return type, (Ellipsis, ) elif obj is type(NotImplemented): return type, (NotImplemented, ) elif obj in _BUILTIN_TYPE_NAMES: return _builtin_type, (_BUILTIN_TYPE_NAMES[obj], ) elif not _should_pickle_by_reference(obj, config=config): return _dynamic_class_reduce(obj, config) return NotImplemented def _dict_keys_reduce(obj): # Safer not to ship the full dict as sending the rest might # be unintended and could potentially cause leaking of # sensitive information return _make_dict_keys, (list(obj), ) def _dict_values_reduce(obj): # Safer not to ship the full dict as sending the rest might # be unintended and could potentially cause leaking of # sensitive information return _make_dict_values, (list(obj), ) def _dict_items_reduce(obj): return _make_dict_items, (dict(obj), ) def _odict_keys_reduce(obj): # Safer not to ship the full dict as sending the rest might # be unintended and could potentially cause leaking of # sensitive information return _make_dict_keys, (list(obj), True) def _odict_values_reduce(obj): # Safer not to ship the full dict as sending the rest might # be unintended and could potentially cause leaking of # sensitive information return _make_dict_values, (list(obj), True) def _odict_items_reduce(obj): return _make_dict_items, (dict(obj), True) def _dataclass_field_base_reduce(obj): return _get_dataclass_field_type_sentinel, (obj.name, ) # COLLECTIONS OF OBJECTS STATE SETTERS # ------------------------------------ # state setters are called at unpickling time, once the object is created and # it has to be updated to how it was at unpickling time. def _function_setstate(obj, state): """Update the state of a dynamic function. As __closure__ and __globals__ are readonly attributes of a function, we cannot rely on the native setstate routine of pickle.load_build, that calls setattr on items of the slotstate. Instead, we have to modify them inplace. """ state, slotstate = state obj.__dict__.update(state) obj_globals = slotstate.pop("__globals__") obj_closure = slotstate.pop("__closure__") # _cloudpickle_subimports is a set of submodules that must be loaded for # the pickled function to work correctly at unpickling time. Now that these # submodules are depickled (hence imported), they can be removed from the # object's state (the object state only served as a reference holder to # these submodules) slotstate.pop("_cloudpickle_submodules") obj.__globals__.update(obj_globals) obj.__globals__["__builtins__"] = __builtins__ if obj_closure is not None: for i, cell in enumerate(obj_closure): try: value = cell.cell_contents except ValueError: # cell is empty continue obj.__closure__[i].cell_contents = value for k, v in slotstate.items(): setattr(obj, k, v) def _class_setstate(obj, state, skip_reset_dynamic_type_state=False): # Lock while potentially modifying class state. with _DYNAMIC_CLASS_TRACKER_LOCK: if skip_reset_dynamic_type_state and obj in _DYNAMIC_CLASS_STATE_TRACKER_BY_CLASS: return obj _DYNAMIC_CLASS_STATE_TRACKER_BY_CLASS[obj] = True state, slotstate = state registry = None for attrname, attr in state.items(): if attrname == "_abc_impl": registry = attr else: # Note: setting attribute names on a class automatically triggers their # interning in CPython: # https://github.com/python/cpython/blob/v3.12.0/Objects/object.c#L957 # # This means that to get deterministic pickling for a dynamic class that # was initially defined in a different Python process, the pickler # needs to ensure that dynamic class and function attribute names are # systematically copied into a non-interned version to avoid # unpredictable pickle payloads. # # Indeed the Pickler's memoizer relies on physical object identity to break # cycles in the reference graph of the object being serialized. setattr(obj, attrname, attr) if sys.version_info >= (3, 13) and "__firstlineno__" in state: # Set the Python 3.13+ only __firstlineno__ attribute one more time, as it # will be automatically deleted by the `setattr(obj, attrname, attr)` call # above when `attrname` is "__firstlineno__". We assume that preserving this # information might be important for some users and that it not stale in the # context of cloudpickle usage, hence legitimate to propagate. Furthermore it # is necessary to do so to keep deterministic chained pickling as tested in # test_deterministic_str_interning_for_chained_dynamic_class_pickling. obj.__firstlineno__ = state["__firstlineno__"] if registry is not None: for subclass in registry: obj.register(subclass) return obj # COLLECTION OF DATACLASS UTILITIES # --------------------------------- # There are some internal sentinel values whose identity must be preserved when # unpickling dataclass fields. Each sentinel value has a unique name that we can # use to retrieve its identity at unpickling time. _DATACLASSE_FIELD_TYPE_SENTINELS = { dataclasses._FIELD.name: dataclasses._FIELD, dataclasses._FIELD_CLASSVAR.name: dataclasses._FIELD_CLASSVAR, dataclasses._FIELD_INITVAR.name: dataclasses._FIELD_INITVAR, } def _get_dataclass_field_type_sentinel(name): return _DATACLASSE_FIELD_TYPE_SENTINELS[name] class Pickler(pickle.Pickler): # set of reducers defined and used by cloudpickle (private) _dispatch_table = {} _dispatch_table[classmethod] = _classmethod_reduce _dispatch_table[io.TextIOWrapper] = _file_reduce _dispatch_table[logging.Logger] = _logger_reduce _dispatch_table[logging.RootLogger] = _root_logger_reduce _dispatch_table[memoryview] = _memoryview_reduce _dispatch_table[property] = _property_reduce _dispatch_table[staticmethod] = _classmethod_reduce _dispatch_table[CellType] = _cell_reduce _dispatch_table[types.GetSetDescriptorType] = _getset_descriptor_reduce _dispatch_table[types.ModuleType] = _module_reduce _dispatch_table[types.MethodType] = _method_reduce _dispatch_table[types.MappingProxyType] = _mappingproxy_reduce _dispatch_table[weakref.WeakSet] = _weakset_reduce _dispatch_table[_collections_abc.dict_keys] = _dict_keys_reduce _dispatch_table[_collections_abc.dict_values] = _dict_values_reduce _dispatch_table[_collections_abc.dict_items] = _dict_items_reduce _dispatch_table[type(OrderedDict().keys())] = _odict_keys_reduce _dispatch_table[type(OrderedDict().values())] = _odict_values_reduce _dispatch_table[type(OrderedDict().items())] = _odict_items_reduce _dispatch_table[abc.abstractmethod] = _classmethod_reduce _dispatch_table[abc.abstractclassmethod] = _classmethod_reduce _dispatch_table[abc.abstractstaticmethod] = _classmethod_reduce _dispatch_table[abc.abstractproperty] = _property_reduce _dispatch_table[dataclasses._FIELD_BASE] = _dataclass_field_base_reduce dispatch_table = ChainMap(_dispatch_table, copyreg.dispatch_table) def _stable_identifier_function_reduce(self, func): code_object_params = self.config.get_code_object_params if code_object_params is None: return self._dynamic_function_reduce(func) code_path = code_object_params.get_code_object_identifier(func) if not code_path: return self._dynamic_function_reduce(func) base_globals = self.globals_ref.setdefault(id(func.__globals__), {}) if base_globals == {}: if "__file__" in func.__globals__: # Apply normalization ONLY to the __file__ attribute file_path = func.__globals__["__file__"] if self.config.filepath_interceptor: file_path = self.config.filepath_interceptor(file_path) base_globals["__file__"] = file_path # Add module attributes used to resolve relative imports # instructions inside func. for k in ["__package__", "__name__", "__path__"]: if k in func.__globals__: base_globals[k] = func.__globals__[k] newargs = (code_path, base_globals, func.__name__, func.__defaults__) state = _function_getstate(func) return ( functools.partial( _make_function_from_identifier, code_object_params.get_code_from_identifier), newargs, state, None, None, _function_setstate) # function reducers are defined as instance methods of cloudpickle.Pickler # objects, as they rely on a cloudpickle.Pickler attribute (globals_ref) def _dynamic_function_reduce(self, func): """Reduce a function that is not pickleable via attribute lookup.""" newargs = self._function_getnewargs(func) state = _function_getstate(func) return (_make_function, newargs, state, None, None, _function_setstate) def _function_reduce(self, obj): """Reducer for function objects. If obj is a top-level attribute of a file-backed module, this reducer returns NotImplemented, making the cloudpickle.Pickler fall back to traditional pickle.Pickler routines to save obj. Otherwise, it reduces obj using a custom cloudpickle reducer designed specifically to handle dynamic functions. """ if _should_pickle_by_reference(obj, config=self.config): return NotImplemented elif self.config.get_code_object_params is not None: return self._stable_identifier_function_reduce(obj) else: return self._dynamic_function_reduce(obj) def _function_getnewargs(self, func): code = func.__code__ # base_globals represents the future global namespace of func at # unpickling time. Looking it up and storing it in # cloudpickle.Pickler.globals_ref allow functions sharing the same # globals at pickling time to also share them once unpickled, at one # condition: since globals_ref is an attribute of a cloudpickle.Pickler # instance, and that a new cloudpickle.Pickler is created each time # cloudpickle.dump or cloudpickle.dumps is called, functions also need # to be saved within the same invocation of # cloudpickle.dump/cloudpickle.dumps (for example: # cloudpickle.dumps([f1, f2])). There is no such limitation when using # cloudpickle.Pickler.dump, as long as the multiple invocations are # bound to the same cloudpickle.Pickler instance. base_globals = self.globals_ref.setdefault(id(func.__globals__), {}) if base_globals == {}: if "__file__" in func.__globals__: # Apply normalization ONLY to the __file__ attribute file_path = func.__globals__["__file__"] if self.config.filepath_interceptor: file_path = self.config.filepath_interceptor(file_path) base_globals["__file__"] = file_path # Add module attributes used to resolve relative imports # instructions inside func. for k in ["__package__", "__name__", "__path__"]: if k in func.__globals__: base_globals[k] = func.__globals__[k] # Do not bind the free variables before the function is created to # avoid infinite recursion. if func.__closure__ is None: closure = None else: closure = tuple(_make_empty_cell() for _ in range(len(code.co_freevars))) return code, base_globals, None, None, closure def dump(self, obj): try: return super().dump(obj) except RecursionError as e: msg = "Could not pickle object as excessively deep recursion required." raise pickle.PicklingError(msg) from e def __init__( self, file, protocol=None, buffer_callback=None, config: CloudPickleConfig = DEFAULT_CONFIG): if protocol is None: protocol = DEFAULT_PROTOCOL super().__init__(file, protocol=protocol, buffer_callback=buffer_callback) # map functions __globals__ attribute ids, to ensure that functions # sharing the same global namespace at pickling time also share # their global namespace at unpickling time. self.globals_ref = {} self.proto = int(protocol) self.config = config if not PYPY: # pickle.Pickler is the C implementation of the CPython pickler and # therefore we rely on reduce_override method to customize the pickler # behavior. # `cloudpickle.Pickler.dispatch` is only left for backward # compatibility - note that when using protocol 5, # `cloudpickle.Pickler.dispatch` is not an extension of # `pickle._Pickler.dispatch` dictionary, because `cloudpickle.Pickler` # subclasses the C-implemented `pickle.Pickler`, which does not expose # a `dispatch` attribute. Earlier versions of `cloudpickle.Pickler` # used `cloudpickle.Pickler.dispatch` as a class-level attribute # storing all reducers implemented by cloudpickle, but the attribute # name was not a great choice given because it would collide with a # similarly named attribute in the pure-Python `pickle._Pickler` # implementation in the standard library. dispatch = dispatch_table # Implementation of the reducer_override callback, in order to # efficiently serialize dynamic functions and classes by subclassing # the C-implemented `pickle.Pickler`. # TODO: decorrelate reducer_override (which is tied to CPython's # implementation - would it make sense to backport it to pypy? - and # pickle's protocol 5 which is implementation agnostic. Currently, the # availability of both notions coincide on CPython's pickle, but it may # not be the case anymore when pypy implements protocol 5. def reducer_override(self, obj): """Type-agnostic reducing callback for function and classes. For performance reasons, subclasses of the C `pickle.Pickler` class cannot register custom reducers for functions and classes in the dispatch_table attribute. Reducers for such types must instead implemented via the special `reducer_override` method. Note that this method will be called for any object except a few builtin-types (int, lists, dicts etc.), which differs from reducers in the Pickler's dispatch_table, each of them being invoked for objects of a specific type only. This property comes in handy for classes: although most classes are instances of the ``type`` metaclass, some of them can be instances of other custom metaclasses (such as enum.EnumMeta for example). In particular, the metaclass will likely not be known in advance, and thus cannot be special-cased using an entry in the dispatch_table. reducer_override, among other things, allows us to register a reducer that will be called for any class, independently of its type. Notes: * reducer_override has the priority over dispatch_table-registered reducers. * reducer_override can be used to fix other limitations of cloudpickle for other types that suffered from type-specific reducers, such as Exceptions. See https://github.com/cloudpipe/cloudpickle/issues/248 """ t = type(obj) try: is_anyclass = issubclass(t, type) except TypeError: # t is not a class (old Boost; see SF #502085) is_anyclass = False if is_anyclass: return _class_reduce(obj, self.config) elif isinstance(obj, typing.TypeVar): # Add this check return _typevar_reduce(obj, self.config) elif isinstance(obj, types.CodeType): return _code_reduce(obj, self.config) elif isinstance(obj, types.FunctionType): return self._function_reduce(obj) else: # fallback to save_global, including the Pickler's # dispatch_table return NotImplemented else: # When reducer_override is not available, hack the pure-Python # Pickler's types.FunctionType and type savers. Note: the type saver # must override Pickler.save_global, because pickle.py contains a # hard-coded call to save_global when pickling meta-classes. dispatch = pickle.Pickler.dispatch.copy() def _save_reduce_pickle5( self, func, args, state=None, listitems=None, dictitems=None, state_setter=None, obj=None, ): save = self.save write = self.write self.save_reduce( func, args, state=None, listitems=listitems, dictitems=dictitems, obj=obj, ) # backport of the Python 3.8 state_setter pickle operations save(state_setter) save(obj) # simple BINGET opcode as obj is already memoized. save(state) write(pickle.TUPLE2) # Trigger a state_setter(obj, state) function call. write(pickle.REDUCE) # The purpose of state_setter is to carry-out an # inplace modification of obj. We do not care about what the # method might return, so its output is eventually removed from # the stack. write(pickle.POP) def save_global(self, obj, name=None, pack=struct.pack): """Main dispatch method. The name of this method is somewhat misleading: all types get dispatched here. """ if obj is type(None): # noqa return self.save_reduce( type, (None, ), obj=obj) elif obj is type(Ellipsis): return self.save_reduce( type, (Ellipsis, ), obj=obj) elif obj is type(NotImplemented): return self.save_reduce( type, (NotImplemented, ), obj=obj) elif obj in _BUILTIN_TYPE_NAMES: return self.save_reduce( _builtin_type, (_BUILTIN_TYPE_NAMES[obj], ), obj=obj) if name is not None: super().save_global(obj, name=name) elif not _should_pickle_by_reference(obj, name=name, config=self.config): self._save_reduce_pickle5( *_dynamic_class_reduce(obj, self.config), obj=obj) else: super().save_global(obj, name=name) dispatch[type] = save_global def save_typevar(self, obj, name=None): """Handle TypeVar objects with access to config.""" return self.save_reduce(*_typevar_reduce(obj, self.config), obj=obj) dispatch[typing.TypeVar] = save_typevar def save_code(self, obj, name=None): return self.save_reduce(*_code_reduce(obj, self.config), obj=obj) dispatch[types.CodeType] = save_code def save_function(self, obj, name=None): """Registered with the dispatch to handle all function types. Determines what kind of function obj is (e.g. lambda, defined at interactive prompt, etc) and handles the pickling appropriately. """ if _should_pickle_by_reference(obj, name=name, config=self.config): return super().save_global(obj, name=name) elif PYPY and isinstance(obj.__code__, builtin_code_type): return self.save_pypy_builtin_func(obj) else: return self._save_reduce_pickle5( *self._dynamic_function_reduce(obj), obj=obj) def save_pypy_builtin_func(self, obj): """Save pypy equivalent of builtin functions. PyPy does not have the concept of builtin-functions. Instead, builtin-functions are simple function instances, but with a builtin-code attribute. Most of the time, builtin functions should be pickled by attribute. But PyPy has flaky support for __qualname__, so some builtin functions such as float.__new__ will be classified as dynamic. For this reason only, we created this special routine. Because builtin-functions are not expected to have closure or globals, there is no additional hack (compared the one already implemented in pickle) to protect ourselves from reference cycles. A simple (reconstructor, newargs, obj.__dict__) tuple is save_reduced. Note also that PyPy improved their support for __qualname__ in v3.6, so this routing should be removed when cloudpickle supports only PyPy 3.6 and later. """ rv = ( types.FunctionType, (obj.__code__, {}, obj.__name__, obj.__defaults__, obj.__closure__), obj.__dict__, ) self.save_reduce(*rv, obj=obj) dispatch[types.FunctionType] = save_function # Shorthands similar to pickle.dump/pickle.dumps def dump( obj, file, protocol=None, buffer_callback=None, config: CloudPickleConfig = DEFAULT_CONFIG): """Serialize obj as bytes streamed into file protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to pickle.HIGHEST_PROTOCOL. This setting favors maximum communication speed between processes running the same Python version. Set protocol=pickle.DEFAULT_PROTOCOL instead if you need to ensure compatibility with older versions of Python (although this is not always guaranteed to work because cloudpickle relies on some internal implementation details that can change from one Python version to the next). """ Pickler( file, protocol=protocol, buffer_callback=buffer_callback, config=config).dump(obj) def dumps( obj, protocol=None, buffer_callback=None, config: CloudPickleConfig = DEFAULT_CONFIG): """Serialize obj as a string of bytes allocated in memory protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to pickle.HIGHEST_PROTOCOL. This setting favors maximum communication speed between processes running the same Python version. Set protocol=pickle.DEFAULT_PROTOCOL instead if you need to ensure compatibility with older versions of Python (although this is not always guaranteed to work because cloudpickle relies on some internal implementation details that can change from one Python version to the next). """ with io.BytesIO() as file: cp = Pickler( file, protocol=protocol, buffer_callback=buffer_callback, config=config) cp.dump(obj) return file.getvalue() # Include pickles unloading functions in this namespace for convenience. load, loads = pickle.load, pickle.loads # Backward compat alias. CloudPickler = Pickler def hash_dynamic_classdef(classdef): """Generates a deterministic ID by hashing the pickled class definition.""" hexdigest = hashlib.sha256( dumps(classdef, config=CloudPickleConfig(id_generator=None))).hexdigest() return hexdigest ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/cloudpickle/cloudpickle_fast.py0000644000175100017510000000050115206113032025650 0ustar00runnerrunner"""Compatibility module. It can be necessary to load files generated by previous versions of cloudpickle that rely on symbols being defined under the `cloudpickle.cloudpickle_fast` namespace. See: tests/test_backward_compat.py """ from . import cloudpickle def __getattr__(name): return getattr(cloudpickle, name) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/cloudpickle_pickler.py0000644000175100017510000002234715206113032024062 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Pickler for values, functions, and classes. For internal use only. No backwards compatibility guarantees. Uses the cloudpickle library to pickle data, functions, lambdas and classes. dump_session and load_session are no-ops. """ # pytype: skip-file import base64 import bz2 import io import logging import sys import threading import zlib from apache_beam.internal import code_object_pickler from apache_beam.internal.cloudpickle import cloudpickle from apache_beam.internal.code_object_pickler import get_normalized_path DEFAULT_CONFIG = cloudpickle.CloudPickleConfig( skip_reset_dynamic_type_state=True, filepath_interceptor=get_normalized_path) STABLE_CODE_IDENTIFIER_CONFIG = cloudpickle.CloudPickleConfig( skip_reset_dynamic_type_state=True, filepath_interceptor=get_normalized_path, get_code_object_params=cloudpickle.GetCodeObjectParams( get_code_object_identifier=code_object_pickler. get_code_object_identifier, get_code_from_identifier=code_object_pickler.get_code_from_identifier)) try: from absl import flags except (ImportError, ModuleNotFoundError): pass def _get_proto_enum_descriptor_class(): try: from google.protobuf.internal import api_implementation except ImportError: return None implementation_type = api_implementation.Type() if implementation_type == 'upb': try: from google._upb._message import EnumDescriptor return EnumDescriptor except ImportError: pass elif implementation_type == 'cpp': try: from google.protobuf.pyext._message import EnumDescriptor return EnumDescriptor except ImportError: pass elif implementation_type == 'python': try: from google.protobuf.internal.python_message import EnumDescriptor return EnumDescriptor except ImportError: pass return None EnumDescriptor = _get_proto_enum_descriptor_class() # Pickling, especially unpickling, causes broken module imports on Python 3 # if executed concurrently, see: BEAM-8651, http://bugs.python.org/issue38884. _pickle_lock = threading.RLock() RLOCK_TYPE = type(_pickle_lock) LOCK_TYPE = type(threading.Lock()) _LOGGER = logging.getLogger(__name__) # Helper to return an object directly during unpickling. def _return_obj(obj): return obj # Optional import for Python 3.12 TypeAliasType try: # pragma: no cover - dependent on Python version from typing import TypeAliasType as _TypeAliasType # type: ignore[attr-defined] except Exception: _TypeAliasType = None def _typealias_reduce(obj): # Unwrap typing.TypeAliasType to its underlying value for robust pickling. underlying = getattr(obj, '__value__', None) if underlying is None: # Fallback: return the object itself; lets default behavior handle it. return _return_obj, (obj, ) return _return_obj, (underlying, ) def _reconstruct_enum_descriptor(full_name): for _, module in list(sys.modules.items()): if not hasattr(module, 'DESCRIPTOR'): continue if hasattr(module.DESCRIPTOR, 'enum_types_by_name'): for (_, enum_desc) in module.DESCRIPTOR.enum_types_by_name.items(): if enum_desc.full_name == full_name: return enum_desc for _, attr_value in vars(module).items(): if not hasattr(attr_value, 'DESCRIPTOR'): continue if hasattr(attr_value.DESCRIPTOR, 'enum_types_by_name'): for (_, enum_desc) in attr_value.DESCRIPTOR.enum_types_by_name.items(): if enum_desc.full_name == full_name: return enum_desc raise ImportError(f'Could not find enum descriptor: {full_name}') def _pickle_enum_descriptor(obj): full_name = obj.full_name return _reconstruct_enum_descriptor, (full_name, ) def dumps( o, enable_trace=True, use_zlib=False, enable_best_effort_determinism=False, enable_stable_code_identifier_pickling=False, config: cloudpickle.CloudPickleConfig = DEFAULT_CONFIG) -> bytes: """For internal use only; no backwards-compatibility guarantees.""" s = _dumps( o, enable_best_effort_determinism, enable_stable_code_identifier_pickling, config) # Compress as compactly as possible (compresslevel=9) to decrease peak memory # usage (of multiple in-memory copies) and to avoid hitting protocol buffer # limits. # WARNING: Be cautious about compressor change since it can lead to pipeline # representation change, and can break streaming job update compatibility on # runners such as Dataflow. if use_zlib: c = zlib.compress(s, 9) else: c = bz2.compress(s, compresslevel=9) del s # Free up some possibly large and no-longer-needed memory. return base64.b64encode(c) def _dumps( o, enable_best_effort_determinism=False, enable_stable_code_identifier_pickling=False, config: cloudpickle.CloudPickleConfig = DEFAULT_CONFIG) -> bytes: if enable_best_effort_determinism: # TODO: Add support once https://github.com/cloudpipe/cloudpickle/pull/563 # is merged in. _LOGGER.warning( 'Ignoring unsupported option: enable_best_effort_determinism. ' 'This has only been implemented for dill.') with _pickle_lock: with io.BytesIO() as file: if enable_stable_code_identifier_pickling: config = STABLE_CODE_IDENTIFIER_CONFIG pickler = cloudpickle.CloudPickler(file, config=config) try: pickler.dispatch_table[type(flags.FLAGS)] = _pickle_absl_flags except NameError: pass # Register Python 3.12 `type` alias reducer to unwrap to underlying value. if _TypeAliasType is not None: pickler.dispatch_table[_TypeAliasType] = _typealias_reduce try: pickler.dispatch_table[RLOCK_TYPE] = _pickle_rlock except NameError: pass try: pickler.dispatch_table[LOCK_TYPE] = _lock_reducer except NameError: pass if EnumDescriptor is not None: pickler.dispatch_table[EnumDescriptor] = _pickle_enum_descriptor pickler.dump(o) return file.getvalue() def loads(encoded, enable_trace=True, use_zlib=False): """For internal use only; no backwards-compatibility guarantees.""" c = base64.b64decode(encoded) if use_zlib: s = zlib.decompress(c) else: s = bz2.decompress(c) del c # Free up some possibly large and no-longer-needed memory. return _loads(s) def _loads(s): with _pickle_lock: unpickled = cloudpickle.loads(s) return unpickled def roundtrip(o): """Internal utility for testing round-trip pickle serialization.""" return _loads(_dumps(o)) def _pickle_absl_flags(obj): return _create_absl_flags, tuple([]) def _create_absl_flags(): return flags.FLAGS def _pickle_rlock(obj): return RLOCK_TYPE, tuple([]) def _lock_reducer(obj): return threading.Lock, tuple([]) def dump_session(file_path): # Since References are saved (https://s.apache.org/beam-picklers), we only # dump supported Beam Registries (currently only logical type registry) from apache_beam.coders import typecoders from apache_beam.typehints import schemas from apache_beam.typehints.schema_registry import SCHEMA_REGISTRY with _pickle_lock, open(file_path, 'wb') as file: coder_reg = typecoders.registry.get_custom_type_coder_tuples() logical_type_reg = schemas.LogicalType._known_logical_types.copy_custom() schema_reg = SCHEMA_REGISTRY.get_registered_typings() pickler = cloudpickle.CloudPickler(file) # TODO(https://github.com/apache/beam/issues/18500) add file system registry # once implemented pickler.dump({ "coder": coder_reg, "logical_type": logical_type_reg, "schema": schema_reg }) def load_session(file_path): from apache_beam.coders import typecoders from apache_beam.typehints import schemas from apache_beam.typehints.schema_registry import SCHEMA_REGISTRY with _pickle_lock, open(file_path, 'rb') as file: registries = cloudpickle.load(file) if type(registries) != dict: raise ValueError( "Faled loading session: expected dict, got {}", type(registries)) if "coder" in registries: typecoders.registry.load_custom_type_coder_tuples(registries["coder"]) else: _LOGGER.warning('No coder registry found in saved session') if "logical_type" in registries: schemas.LogicalType._known_logical_types.load(registries["logical_type"]) else: _LOGGER.warning('No logical type registry found in saved session') if "schema" in registries: SCHEMA_REGISTRY.load_registered_typings(registries["schema"]) else: _LOGGER.warning('No schema registry found in saved session') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/cloudpickle_pickler_test.py0000644000175100017510000002242715206113032025120 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the cloudpickle_pickler module.""" # pytype: skip-file import os import tempfile import threading import types import unittest from unittest import mock from apache_beam.coders import proto2_coder_test_messages_pb2 from apache_beam.internal import cloudpickle_pickler as beam_cloudpickle from apache_beam.internal import code_object_pickler from apache_beam.internal import module_test from apache_beam.internal.cloudpickle_pickler import dumps from apache_beam.internal.cloudpickle_pickler import loads from apache_beam.typehints.schemas import LogicalTypeRegistry from apache_beam.utils import shared GLOBAL_DICT_REF = module_test.GLOBAL_DICT # Allow weakref to dict class DictWrapper(dict): pass MAIN_MODULE_DICT = DictWrapper() def acquire_dict(): return DictWrapper() class PicklerTest(unittest.TestCase): NO_MAPPINGPROXYTYPE = not hasattr(types, "MappingProxyType") def test_globals_main_are_pickled_by_value(self): self.assertIsNot(MAIN_MODULE_DICT, loads(dumps(lambda: MAIN_MODULE_DICT))()) def test_globals_shared_are_pickled_by_reference(self): shared_handler = shared.Shared() original_dict = shared_handler.acquire(acquire_dict) unpickled_dict = loads( dumps(lambda: shared_handler.acquire(acquire_dict)))() self.assertIs(original_dict, unpickled_dict) def test_module_globals_are_pickled_by_value_when_directly_referenced(self): global_dict = loads(dumps(module_test.GLOBAL_DICT)) self.assertIsNot(module_test.GLOBAL_DICT, global_dict) def test_function_main_with_explicit_module_reference_pickles_by_reference( self): def returns_global_dict(): return module_test.GLOBAL_DICT self.assertIs(module_test.GLOBAL_DICT, loads(dumps(returns_global_dict))()) def test_function_main_with_indirect_module_reference_pickles_by_value(self): def returns_global_dict(): return GLOBAL_DICT_REF self.assertIsNot( module_test.GLOBAL_DICT, loads(dumps(returns_global_dict))()) def test_function_referencing_unpicklable_object_works_when_imported(self): self.assertEqual( module_test.UNPICKLABLE_INSTANCE, loads(dumps(module_test.fn_returns_unpicklable))()) def test_closure_with_unpicklable_object_fails_when_imported(self): # The entire closure is pickled by value, and therefore module_test is # not imported. Requires the module global to be pickled by value. with self.assertRaises(Exception): loads(dumps(module_test.closure_contains_unpicklable())) def test_closure_with_explicit_self_import_can_reference_unpicklable_objects( self): # The closure imports module_test within the function definition # and returns self.UNPICKLABLE_INSTANCE. This allows cloudpickle # to use submimort to reference module_test.UNPICKLABLE_INSTANCE self.assertIs( module_test.UNPICKLABLE_INSTANCE, loads(dumps(module_test.closure_contains_unpicklable_imports_self()))()) def test_closure_main_can_reference_unpicklable_module_objects(self): def outer(): def inner(): return module_test.UNPICKLABLE_INSTANCE return inner # Uses subimport to reference module_test.UNPICKLABLE_INSTANCE rather than # recreate. self.assertIs(module_test.UNPICKLABLE_INSTANCE, loads(dumps(outer()))()) def test_pickle_nested_enum_descriptor(self): NestedEnum = proto2_coder_test_messages_pb2.MessageD.NestedEnum def fn(): return NestedEnum.TWO self.assertEqual(fn(), loads(dumps(fn))()) def test_pickle_top_level_enum_descriptor(self): TopLevelEnum = proto2_coder_test_messages_pb2.TopLevelEnum def fn(): return TopLevelEnum.ONE self.assertEqual(fn(), loads(dumps(fn))()) def test_basics(self): self.assertEqual([1, 'a', ('z', )], loads(dumps([1, 'a', ('z', )]))) fun = lambda x: 'xyz-%s' % x self.assertEqual('xyz-abc', loads(dumps(fun))('abc')) def test_lambda_with_globals(self): """Tests that the globals of a function are preserved.""" # The point of the test is that the lambda being called after unpickling # relies on having the re module being loaded. self.assertEqual(['abc', 'def'], loads(dumps( module_test.get_lambda_with_globals()))('abc def')) def test_lambda_with_main_globals(self): self.assertEqual(unittest, loads(dumps(lambda: unittest))()) def test_lambda_with_closure(self): """Tests that the closure of a function is preserved.""" self.assertEqual( 'closure: abc', loads(dumps(module_test.get_lambda_with_closure('abc')))()) def test_class_object_pickled(self): self.assertEqual(['abc', 'def'], loads(dumps(module_test.Xyz))().foo('abc def')) def test_class_instance_pickled(self): self.assertEqual(['abc', 'def'], loads(dumps(module_test.XYZ_OBJECT)).foo('abc def')) def test_pickling_preserves_closure_of_a_function(self): self.assertEqual( 'X:abc', loads(dumps(module_test.TopClass.NestedClass('abc'))).datum) self.assertEqual( 'Y:abc', loads(dumps(module_test.TopClass.MiddleClass.NestedClass('abc'))).datum) def test_pickle_dynamic_class(self): self.assertEqual( 'Z:abc', loads(dumps(module_test.create_class('abc'))).get()) def test_generators(self): with self.assertRaises(TypeError): dumps((_ for _ in range(10))) def test_recursive_class(self): self.assertEqual( 'RecursiveClass:abc', loads(dumps(module_test.RecursiveClass('abc').datum))) def test_function_with_external_reference(self): out_of_scope_var = 'expected_value' def foo(): return out_of_scope_var self.assertEqual('expected_value', loads(dumps(foo))()) def test_pickle_rlock(self): rlock_instance = threading.RLock() rlock_type = type(rlock_instance) self.assertIsInstance(loads(dumps(rlock_instance)), rlock_type) def test_pickle_lock(self): lock_instance = threading.Lock() lock_type = type(lock_instance) self.assertIsInstance(loads(dumps(lock_instance)), lock_type) @unittest.skipIf(NO_MAPPINGPROXYTYPE, 'test if MappingProxyType introduced') def test_dump_and_load_mapping_proxy(self): self.assertEqual( 'def', loads(dumps(types.MappingProxyType({'abc': 'def'})))['abc']) self.assertEqual( types.MappingProxyType, type(loads(dumps(types.MappingProxyType({}))))) # pylint: disable=exec-used def test_dataclass(self): exec( ''' from apache_beam.internal.module_test import DataClass self.assertEqual(DataClass(datum='abc'), loads(dumps(DataClass(datum='abc')))) ''') def test_best_effort_determinism_not_implemented(self): with self.assertLogs('apache_beam.internal.cloudpickle_pickler', "WARNING") as l: dumps(123, enable_best_effort_determinism=True) self.assertIn( 'Ignoring unsupported option: enable_best_effort_determinism', '\n'.join(l.output)) @mock.patch.object( beam_cloudpickle.DEFAULT_CONFIG, 'filepath_interceptor', autospec=True) def test_default_config_interceptor(self, mock_filepath_interceptor): """Tests config.filepath_interceptor is called for CodeType pickling.""" mock_filepath_interceptor.side_effect = ( code_object_pickler.get_normalized_path) def sample_func(): return "Beam" code_obj = sample_func.__code__ original_filename = os.path.abspath(code_obj.co_filename) pickled_code = beam_cloudpickle.dumps(code_obj) unpickled_code = beam_cloudpickle.loads(pickled_code) mock_filepath_interceptor.assert_called() unpickled_filename = os.path.abspath(unpickled_code.co_filename) self.assertEqual(unpickled_filename, original_filename) @mock.patch( "apache_beam.coders.typecoders.registry.load_custom_type_coder_tuples") @mock.patch( "apache_beam.typehints.schemas.LogicalType._known_logical_types.load") def test_dump_load_session(self, logicaltype_mock, coder_mock): session_file = 'pickled' with tempfile.TemporaryDirectory() as tmp_dirname: pickled_session_file = os.path.join(tmp_dirname, session_file) beam_cloudpickle.dump_session(pickled_session_file) beam_cloudpickle.load_session(pickled_session_file) load_logical_types = logicaltype_mock.call_args.args load_coders = coder_mock.call_args.args self.assertEqual(len(load_logical_types), 1) self.assertEqual(len(load_coders), 1) self.assertTrue(isinstance(load_logical_types[0], LogicalTypeRegistry)) self.assertTrue(isinstance(load_coders[0], list)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/code_object_pickler.py0000644000175100017510000004354715206113032024031 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Customizations to how Python code objects are pickled. This module provides helper functions to improve pickling code objects, especially lambdas, in a consistent way by using code object identifiers. These helper functions are used to patch pickler implementations used by Beam (e.g. Cloudpickle). A code object identifier is a unique identifier for a code object that provides a unique reference to the code object in the context where the code is defined and is invariant to small changes in the surrounding code. The code object identifiers consists of a sequence of the following parts separated by periods: - Module names - The name of the module the code object is in - Class names - The name of a class containing the code object. There can be multiple of these in the same identifier in the case of nested classes. - Function names - The name of the function containing the code object. There can be multiple of these in the case of nested functions. - __code__ - Attribute indicating that we are entering the code object of a function/method. - __co_consts__[] - The name of the local variable containing the code object. In the case of lambdas, the name is created by using the signature of the lambda and hashing the bytecode, as shown below. Examples: - __main__.top_level_function.__code__ - __main__.ClassWithNestedFunction.process.__code__.co_consts[nested_function] - __main__.ClassWithNestedLambda.process.__code__.co_consts[ get_lambda_from_dictionary].co_consts[, ('x',)] - __main__.ClassWithNestedLambda.process.__code__.co_consts[ , ('x',), 1234567890] """ import collections import hashlib import inspect import re import sys import types from typing import Optional from typing import Union def get_normalized_path(path): """Returns a normalized path. This function is intended to be overridden.""" return path def get_code_object_identifier(callable: types.FunctionType): """Returns the code object identifier for a given callable. Args: callable: The callable object to search for. Returns: The code object identifier. Examples: - __main__.top_level_function.__code__ - __main__.ClassWithNestedFunction.process.__code__.co_consts[ nested_function] - __main__.ClassWithNestedLambda.process.__code__.co_consts[ get_lambda_from_dictionary].co_consts[, ('x',)] - __main__.ClassWithNestedLambda.process.__code__.co_consts[ , ('x',), 1234567890] """ if (not hasattr(callable, '__module__') or not hasattr(callable, '__qualname__') or not callable.__module__ or callable.__module__ not in sys.modules): return None code_path: str = _extend_path( callable.__module__, _search( callable, sys.modules[callable.__module__], callable.__qualname__.split('.'), ), ) return code_path def _extend_path(prefix: str, current_path: Optional[str]): """Extends the path to the code object. Args: prefix: The prefix of the path. current_path: The rest of the path. Returns: The extended path. """ if current_path is None: return None if not current_path: return prefix return prefix + '.' + current_path def _search( callable: types.FunctionType, node: Union[types.ModuleType, types.FunctionType, types.CodeType], qual_name_parts: list[str]): """Searches an object to create a code object identifier. Recursively searches the tree of objects starting from node to find the callable's code object. It navigates through the attributes by using the first element of qual_name_parts to indicate what object it is currently at, then recursively passes through the rest of the list until the callable is found. Special components like '' and '' direct the search within nested code objects. Example of qual_name_parts: ['MyClass', 'process', '', ''] Args: callable: The callable object to search for. node: The object to search within. qual_name_parts: A list of strings representing the qualified name of the callable object. Returns: The code object identifier, or None if not found. """ if node is None: return None if not qual_name_parts: if (hasattr(node, '__code__') and hasattr(callable, '__code__') and node.__code__ == callable.__code__): return '__code__' else: return None if inspect.ismodule(node) or inspect.isclass(node): return _search_module_or_class(callable, node, qual_name_parts) elif inspect.isfunction(node): return _search_function(callable, node, qual_name_parts) elif inspect.iscode(node): return _search_code(callable, node, qual_name_parts) def _search_module_or_class( callable: types.FunctionType, node: types.ModuleType, qual_name_parts: list[str]): """Searches a module or class to create a code object identifier. Args: callable: The callable object to search for. node: The module or class to search within. qual_name_parts: The list of qual name parts. Returns: The code object identifier, or None if not found. """ # Functions/methods have a name that is unique within a given module or class # so the traversal can directly lookup function object identified by the name. # Lambdas don't have a name so we need to search all the attributes of the # node. first_part = qual_name_parts[0] rest = qual_name_parts[1:] if first_part == '': for name in dir(node): value = getattr(node, name) if (hasattr(callable, '__code__') and isinstance(value, type(callable)) and value.__code__ == callable.__code__): return name + '.__code__' elif (isinstance(value, types.FunctionType) and value.__defaults__ is not None): # Python functions can have other functions as default parameters which # might contain the code object so we have to search them. for i, default_param_value in enumerate(value.__defaults__): path = _search(callable, default_param_value, rest) if path is not None: return _extend_path(name, _extend_path(f'__defaults__[{i}]', path)) else: if not hasattr(node, first_part): return None return _extend_path( first_part, _search(callable, getattr(node, first_part), rest)) def _search_function( callable: types.FunctionType, node: types.FunctionType, qual_name_parts: list[str]): """Searches a function to create a code object identifier. Args: callable: The callable object to search for. node: The function to search within. qual_name_parts: The list of qual name parts. Returns: The code object identifier, or None if not found. """ first_part = qual_name_parts[0] if (node.__code__ == callable.__code__): if len(qual_name_parts) > 1: raise ValueError('Qual name parts too long') return '__code__' # If first part is '' then the code object is in a local variable # so we should add __code__ to the path to indicate that we are entering # the code object of the function. if first_part == '': return _extend_path( '__code__', _search(callable, node.__code__, qual_name_parts)) def _search_code( callable: types.FunctionType, node: types.CodeType, qual_name_parts: list[str]): """Searches a code object to create a code object identifier. Args: callable: The callable to search for. node: The code object to search within. qual_name_parts: The list of qual name parts. Returns: The code object identifier, or None if not found. Raises: ValueError: If the qual name parts are too long. """ first_part = qual_name_parts[0] rest = qual_name_parts[1:] if hasattr(callable, '__code__') and node == callable.__code__: if len(qual_name_parts) > 1: raise ValueError('Qual name parts too long') return '' elif first_part == '': code_objects_by_name = collections.defaultdict(list) for co_const in node.co_consts: if inspect.iscode(co_const): code_objects_by_name[co_const.co_name].append(co_const) num_lambdas = len(code_objects_by_name.get('', [])) # If there is only one lambda, we can use the default path # 'co_consts[]'. This is the most common case and it is # faster than calculating the signature and the hash. if num_lambdas == 1: path = _search(callable, code_objects_by_name[''][0], rest) if path is not None: return _extend_path('co_consts[]', path) else: return _search_lambda(callable, code_objects_by_name, rest) elif node.co_name == first_part: return _search(callable, node, rest) def _search_lambda( callable: types.FunctionType, code_objects_by_name: dict[str, list[types.CodeType]], qual_name_parts: list[str]): """Searches a lambda to create a code object identifier. Args: callable: The callable to search for. code_objects_by_name: The code objects to search within, keyed by name. qual_name_parts: The rest of the qual_name_parts. Returns: The code object identifier, or None if not found. """ # There are multiple lambdas in the code object, so we need to calculate # the signature and the hash to identify the correct lambda. lambda_code_objects_by_name = collections.defaultdict(list) name = qual_name_parts[0] code_objects = code_objects_by_name[name] if not code_objects: return None if name == '': for code_object in code_objects: lambda_name = f', {_signature(code_object)}' lambda_code_objects_by_name[lambda_name].append(code_object) # Check if there are any lambdas with the same signature. # If there are, we need to calculate the hash to identify the correct # lambda. for lambda_name, lambda_objects in lambda_code_objects_by_name.items(): if len(lambda_objects) > 1: for lambda_object in lambda_objects: path = _search(callable, lambda_object, qual_name_parts) if path is not None: return _extend_path( f'co_consts[{lambda_name},' f' {_create_bytecode_hash(lambda_object)}]', path, ) else: # If there is only one lambda with this signature, we can # use the signature to identify the correct lambda. path = _search(callable, code_objects[0], qual_name_parts) if path is not None: return _extend_path(f'co_consts[{lambda_name}]', path) else: # For non lambda objects, we can use the name to identify the object. path = _search(callable, code_objects[0], qual_name_parts) if path is not None: return _extend_path(f'co_consts[{name}]', path) # Matches a path like: co_consts[my_function] _SINGLE_NAME_PATTERN = re.compile(r'co_consts\[([a-zA-Z0-9\<\>_-]+)]') # Matches a path like: co_consts[, ('x',)] _LAMBDA_WITH_ARGS_PATTERN = re.compile( r"co_consts\[(<.*?>),\s(\('[^']+'(?:,\s*'[^']+')*,?\))\]") # Matches a path like: co_consts[, ('x',), 1234567890] _LAMBDA_WITH_HASH_PATTERN = re.compile( r"co_consts\[(<[^>]+>),\s*(\([^\)]*\)),?\s*(.*)\]") # Matches a path like: __defaults__[0] _DEFAULT_PATTERN = re.compile(r'(__defaults__)\[(\d+)\]') # Matches an argument like: 'x' _ARGUMENT_PATTERN = re.compile(r"'([^']*)'") def _get_code_object_from_single_name_pattern( obj: types.ModuleType, name_result: re.Match[str], path: str): """Returns the code object from a name pattern. Args: obj: The object to search within. name_result: The result of the name pattern search. path: The path to the code object. Returns: The code object. Raises: ValueError: If the pattern is invalid. AttributeError: If the code object is not found. """ if len(name_result.groups()) > 1: raise ValueError(f'Invalid pattern for single name: {name_result.group(0)}') # Groups are indexed starting at 1, group(0) is the entire match. name = name_result.group(1) if hasattr(obj, 'co_consts'): for co_const in obj.co_consts: if inspect.iscode(co_const) and co_const.co_name == name: return co_const raise AttributeError(f'Could not find code object with path: {path}') def _get_code_object_from_lambda_with_args_pattern( obj: types.ModuleType, lambda_with_args_result: re.Match[str], path: str): """Returns the code object from a lambda with args pattern. Args: obj: The object to search within. lambda_with_args_result: The result of the lambda with args pattern search. path: The path to the code object. Returns: The code object. Raises: AttributeError: If the code object is not found. """ name = lambda_with_args_result.group(1) code_objects = collections.defaultdict(list) if hasattr(obj, 'co_consts'): for co_const in obj.co_consts: if inspect.iscode(co_const) and co_const.co_name == name: code_objects[co_const.co_name].append(co_const) for name, objects in code_objects.items(): for obj_ in objects: args = tuple( re.findall(_ARGUMENT_PATTERN, lambda_with_args_result.group(2))) if obj_.co_varnames[:_get_arg_count(obj_)] == args: return obj_ raise AttributeError(f'Could not find code object with path: {path}') def _get_code_object_from_lambda_with_hash_pattern( obj: types.ModuleType, lambda_with_hash_result: re.Match[str], path: str): """Returns the code object from a lambda with hash pattern. Args: obj: The object to search within. lambda_with_hash_result: The result of the lambda with hash pattern search. path: The path to the code object. Returns: The code object. Raises: AttributeError: If the code object is not found. """ name = lambda_with_hash_result.group(1) code_objects = collections.defaultdict(list) if hasattr(obj, 'co_consts'): for co_const in obj.co_consts: if inspect.iscode(co_const) and co_const.co_name == name: code_objects[co_const.co_name].append(co_const) for name, objects in code_objects.items(): for obj_ in objects: args = tuple( re.findall(_ARGUMENT_PATTERN, lambda_with_hash_result.group(2))) if obj_.co_varnames[:_get_arg_count(obj_)] == args: hash_value = lambda_with_hash_result.group(3) if hash_value == str(_create_bytecode_hash(obj_)): return obj_ raise AttributeError(f'Could not find code object with path: {path}') def get_code_from_identifier(code_object_identifier: str): """Returns the code object corresponding to the code object identifier. Args: code_object_identifier: A string representing the code object identifier. Returns: The code object. Raises: ValueError: If the path is empty or invalid. AttributeError: If the attribute is not found. """ if not code_object_identifier: raise ValueError('Path must not be empty.') parts = code_object_identifier.split('.') if parts[0] not in sys.modules: raise AttributeError(f'Module {parts[0]} not found in sys.modules') obj = sys.modules[parts[0]] for part in parts[1:]: if name_result := _SINGLE_NAME_PATTERN.fullmatch(part): obj = _get_code_object_from_single_name_pattern( obj, name_result, code_object_identifier) elif lambda_with_args_result := _LAMBDA_WITH_ARGS_PATTERN.fullmatch(part): obj = _get_code_object_from_lambda_with_args_pattern( obj, lambda_with_args_result, code_object_identifier) elif lambda_with_hash_result := _LAMBDA_WITH_HASH_PATTERN.fullmatch(part): obj = _get_code_object_from_lambda_with_hash_pattern( obj, lambda_with_hash_result, code_object_identifier) elif default_result := _DEFAULT_PATTERN.fullmatch(part): index = int(default_result.group(2)) if index >= len(obj.__defaults__): raise ValueError( f'Index {index} is out of bounds for obj.__defaults__' f' {len(obj.__defaults__)} in path {code_object_identifier}') obj = getattr(obj, '__defaults__')[index] else: obj = getattr(obj, part) if isinstance(obj, types.CodeType): return obj else: raise AttributeError( f'Could not find code object with path: {code_object_identifier}') def _signature(obj: types.CodeType): """Returns the signature of a code object. The signature is the names of the arguments of the code object. This is used to unique identify lambdas. Args: obj: A code object, function, method, or cell. Returns: A tuple of the names of the arguments of the code object. """ return obj.co_varnames[:_get_arg_count(obj)] def _get_arg_count(obj: types.CodeType): """Returns the number of arguments of a code object. Args: obj: A code object, function, method, or cell. Returns: The number of arguments of the code object, or None if the object is not a code object. """ return ( obj.co_argcount + obj.co_kwonlyargcount + (obj.co_flags & 4 == 4) # PyCF_VARARGS + (obj.co_flags & 8 == 8) # PyCF_VARKEYWORDS ) def _create_bytecode_hash(code_object: types.CodeType): """Returns the hash of a code object. Args: code_object: A code object. Returns: The hash of the code object. """ return hashlib.md5(code_object.co_code).hexdigest() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/code_object_pickler_test.py0000644000175100017510000005273215206113032025064 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for generating stable identifiers to use for Pickle serialization.""" import hashlib import unittest from parameterized import parameterized # pylint: disable=unused-import from apache_beam.internal import code_object_pickler from apache_beam.internal.test_data import module_1 from apache_beam.internal.test_data import module_1_class_added from apache_beam.internal.test_data import module_1_function_added from apache_beam.internal.test_data import module_1_global_variable_added from apache_beam.internal.test_data import module_1_lambda_variable_added from apache_beam.internal.test_data import module_1_local_variable_added from apache_beam.internal.test_data import module_1_local_variable_removed from apache_beam.internal.test_data import module_1_nested_function_2_added from apache_beam.internal.test_data import module_1_nested_function_added from apache_beam.internal.test_data import module_2 from apache_beam.internal.test_data import module_2_modified from apache_beam.internal.test_data import module_3 from apache_beam.internal.test_data import module_3_modified from apache_beam.internal.test_data import module_with_default_argument def top_level_function(): return 1 top_level_lambda = lambda x: 1 def get_nested_function(): def nested_function(): return 1 return nested_function def get_lambda_from_dictionary(): d = {"a": lambda x: 1, "b": lambda y: 2} return d["a"] def get_lambda_from_dictionary_same_args(): d = {"a": lambda x: 1, "b": lambda x: x + 1} return d["a"] def function_with_lambda_default_argument(fn=lambda x: 1): return fn def function_with_function_default_argument(fn=top_level_function): return fn def function_decorator(f): return lambda x: f(f(x)) @function_decorator def add_one(x): return x + 1 class ClassWithFunction: def process(self): return 1 class ClassWithStaticMethod: @staticmethod def static_method(): return 1 class ClassWithClassMethod: @classmethod def class_method(cls): return 1 class ClassWithNestedFunction: def process(self): def nested_function(): return 1 return nested_function class ClassWithLambda: def process(self): return lambda: 1 class ClassWithNestedClass: class InnerClass: def process(self): return 1 class ClassWithNestedLambda: def process(self): def get_lambda_from_dictionary(): d = {"a": lambda x: 1, "b": lambda y: 2} return d["a"] return get_lambda_from_dictionary() prefix = __name__ test_cases = [ (top_level_function, f"{prefix}.top_level_function" ".__code__"), (top_level_lambda, f"{prefix}.top_level_lambda" ".__code__"), ( get_nested_function(), ( f"{prefix}.get_nested_function" ".__code__.co_consts[nested_function]")), ( get_lambda_from_dictionary(), ( f"{prefix}" ".get_lambda_from_dictionary.__code__.co_consts[, ('x',)]") ), ( get_lambda_from_dictionary_same_args(), ( f"{prefix}" ".get_lambda_from_dictionary_same_args.__code__.co_consts" "[, ('x',), " + hashlib.md5( get_lambda_from_dictionary_same_args().__code__.co_code). hexdigest() + "]")), ( function_with_lambda_default_argument(), ( f"{prefix}" ".function_with_lambda_default_argument.__defaults__[0].__code__")), ( function_with_function_default_argument(), f"{prefix}.top_level_function" ".__code__"), (add_one, f"{prefix}.function_decorator" ".__code__.co_consts[]"), ( ClassWithFunction.process, f"{prefix}.ClassWithFunction" ".process.__code__"), ( ClassWithStaticMethod.static_method, f"{prefix}.ClassWithStaticMethod" ".static_method.__code__"), ( ClassWithClassMethod.class_method, f"{prefix}.ClassWithClassMethod" ".class_method.__code__"), ( ClassWithNestedFunction().process(), ( f"{prefix}.ClassWithNestedFunction.process.__code__.co_consts" "[nested_function]")), ( ClassWithLambda().process(), f"{prefix}.ClassWithLambda.process.__code__.co_consts[]"), ( ClassWithNestedClass.InnerClass().process, f"{prefix}.ClassWithNestedClass.InnerClass.process.__code__"), ( ClassWithNestedLambda().process(), ( f"{prefix}" ".ClassWithNestedLambda.process.__code__.co_consts" "[get_lambda_from_dictionary].co_consts[, ('x',)]")), ( ClassWithNestedLambda.process, f"{prefix}.ClassWithNestedLambda.process.__code__"), ] class CodeObjectIdentifierGenerationTest(unittest.TestCase): @parameterized.expand(test_cases) def test_get_code_object_identifier(self, callable, expected_path): actual = code_object_pickler.get_code_object_identifier(callable) self.assertEqual(actual, expected_path) @parameterized.expand(test_cases) def test_get_code_from_identifier(self, expected_callable, path): actual = code_object_pickler.get_code_from_identifier(path) self.assertEqual(actual, expected_callable.__code__) @parameterized.expand(test_cases) def test_roundtrip(self, callable, unused_path): path = code_object_pickler.get_code_object_identifier(callable) actual = code_object_pickler.get_code_from_identifier(path) self.assertEqual(actual, callable.__code__) class GetCodeFromCodeObjectIdentifierTest(unittest.TestCase): def test_empty_path_raises_exception(self): with self.assertRaisesRegex(ValueError, "Path must not be empty"): code_object_pickler.get_code_from_identifier("") def test_invalid_default_index_raises_exception(self): with self.assertRaisesRegex(ValueError, "out of bounds"): code_object_pickler.get_code_from_identifier( "apache_beam.internal.test_data.module_with_default_argument." "function_with_lambda_default_argument.__defaults__[1]") def test_invalid_single_name_path_raises_exception(self): with self.assertRaisesRegex(AttributeError, "Could not find code object with path"): code_object_pickler.get_code_from_identifier( "apache_beam.internal.test_data.module_3." "my_function.__code__.co_consts[something]") def test_invalid_lambda_with_args_path_raises_exception(self): with self.assertRaisesRegex(AttributeError, "Could not find code object with path"): code_object_pickler.get_code_from_identifier( "apache_beam.internal.test_data.module_3." "my_function.__code__.co_consts[, ('x',)]") def test_invalid_lambda_with_hash_path_raises_exception(self): with self.assertRaisesRegex(AttributeError, "Could not find code object with path"): code_object_pickler.get_code_from_identifier( "apache_beam.internal.test_data.module_3." "my_function.__code__.co_consts[, ('',), 1234567890]") def test_adding_local_variable_in_class_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_2.AddLocalVariable.my_method(self)).replace( "module_2", "module_2_modified")), module_2_modified.AddLocalVariable.my_method(self).__code__, ) def test_removing_local_variable_in_class_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_2.RemoveLocalVariable.my_method(self)).replace( "module_2", "module_2_modified")), module_2_modified.RemoveLocalVariable.my_method(self).__code__, ) def test_adding_lambda_variable_in_class_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_2.AddLambdaVariable.my_method(self)).replace( "module_2", "module_2_modified")), module_2_modified.AddLambdaVariable.my_method(self).__code__, ) def test_removing_lambda_variable_in_class_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_2.RemoveLambdaVariable.my_method(self)).replace( "module_2", "module_2_modified")), module_2_modified.RemoveLambdaVariable.my_method(self).__code__, ) def test_adding_nested_function_in_class_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_2.ClassWithNestedFunction.my_method(self)).replace( "module_2", "module_2_modified")), module_2_modified.ClassWithNestedFunction.my_method(self).__code__, ) def test_adding_nested_function_2_in_class_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_2.ClassWithNestedFunction2.my_method(self)).replace( "module_2", "module_2_modified")), module_2_modified.ClassWithNestedFunction2.my_method(self).__code__, ) def test_adding_new_function_in_class_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_2.ClassWithTwoMethods.my_method(self)).replace( "module_2", "module_2_modified")), module_2_modified.ClassWithTwoMethods.my_method(self).__code__, ) def test_removing_method_in_class_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_2.RemoveMethod.my_method(self)).replace( "module_2", "module_2_modified")), module_2_modified.RemoveMethod.my_method(self).__code__, ) def test_adding_global_variable_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace( "module_1", "module_1_global_variable_added", )), module_1_global_variable_added.my_function().__code__, ) def test_adding_top_level_function_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace( "module_1", "module_1_function_added")), module_1_function_added.my_function().__code__, ) def test_adding_local_variable_in_function_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace( "module_1", "module_1_local_variable_added")), module_1_local_variable_added.my_function().__code__, ) def test_removing_local_variable_in_function_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace( "module_1", "module_1_local_variable_removed")), module_1_local_variable_removed.my_function().__code__, ) def test_adding_nested_function_in_function_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace( "module_1", "module_1_nested_function_added")), module_1_nested_function_added.my_function().__code__, ) def test_adding_nested_function_2_in_function_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace( "module_1", "module_1_nested_function_2_added")), module_1_nested_function_2_added.my_function().__code__, ) def test_adding_class_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace( "module_1", "module_1_class_added")), module_1_class_added.my_function().__code__, ) def test_adding_lambda_variable_in_function_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace( "module_1", "module_1_lambda_variable_added")), module_1_lambda_variable_added.my_function().__code__, ) def test_removing_lambda_variable_in_function_preserves_object(self): self.assertEqual( code_object_pickler.get_code_from_identifier( code_object_pickler.get_code_object_identifier( module_3.my_function()).replace( "module_3", "module_3_modified")), module_3_modified.my_function().__code__, ) class CodePathStabilityTest(unittest.TestCase): def test_adding_local_variable_in_class_preserves_path(self): self.assertEqual( code_object_pickler.get_code_object_identifier( module_2.AddLocalVariable.my_method(self)).replace( "module_2", "module_name"), code_object_pickler.get_code_object_identifier( module_2_modified.AddLocalVariable.my_method(self)).replace( "module_2_modified", "module_name"), ) def test_removing_local_variable_in_class_preserves_path(self): self.assertEqual( code_object_pickler.get_code_object_identifier( module_2.RemoveLocalVariable.my_method(self)).replace( "module_2", "module_name"), code_object_pickler.get_code_object_identifier( module_2_modified.RemoveLocalVariable.my_method(self)).replace( "module_2_modified", "module_name"), ) def test_adding_lambda_variable_in_class_changes_path(self): self.assertNotEqual( code_object_pickler.get_code_object_identifier( module_2.AddLambdaVariable.my_method(self)).replace( "module_2", "module_name"), code_object_pickler.get_code_object_identifier( module_2_modified.AddLambdaVariable.my_method(self)).replace( "module_2_modified", "module_name"), ) def test_removing_lambda_variable_in_class_changes_path(self): self.assertNotEqual( code_object_pickler.get_code_object_identifier( module_2.RemoveLambdaVariable.my_method(self)).replace( "module_2", "module_name"), code_object_pickler.get_code_object_identifier( module_2_modified.RemoveLambdaVariable.my_method(self)).replace( "module_2_modified", "module_name"), ) def test_adding_nested_function_in_class_preserves_path(self): self.assertEqual( code_object_pickler.get_code_object_identifier( module_2.ClassWithNestedFunction.my_method(self)).replace( "module_2", "module_name"), code_object_pickler.get_code_object_identifier( module_2_modified.ClassWithNestedFunction.my_method(self)).replace( "module_2_modified", "module_name"), ) def test_adding_nested_function_2_in_class_preserves_path(self): self.assertEqual( code_object_pickler.get_code_object_identifier( module_2.ClassWithNestedFunction2.my_method(self)).replace( "module_2", "module_name"), code_object_pickler.get_code_object_identifier( module_2_modified.ClassWithNestedFunction2.my_method(self)).replace( "module_2_modified", "module_name"), ) def test_adding_new_function_in_class_preserves_path(self): self.assertEqual( code_object_pickler.get_code_object_identifier( module_2.ClassWithTwoMethods.my_method(self)).replace( "module_2", "module_name"), code_object_pickler.get_code_object_identifier( module_2_modified.ClassWithTwoMethods.my_method(self)).replace( "module_2_modified", "module_name"), ) def test_removing_function_in_class_preserves_path(self): self.assertEqual( code_object_pickler.get_code_object_identifier( module_2.RemoveMethod.my_method(self)).replace( "module_2", "module_name"), code_object_pickler.get_code_object_identifier( module_2_modified.RemoveMethod.my_method(self)).replace( "module_2_modified", "module_name"), ) def test_adding_global_variable_preserves_path(self): self.assertEqual( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace("module_1", "module_name"), code_object_pickler.get_code_object_identifier( module_1_global_variable_added.my_function()).replace( "module_1_global_variable_added", "module_name"), ) def test_adding_top_level_function_preserves_path(self): self.assertEqual( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace("module_1", "module_name"), code_object_pickler.get_code_object_identifier( module_1_function_added.my_function()).replace( "module_1_function_added", "module_name"), ) def test_adding_local_variable_in_function_preserves_path(self): self.assertEqual( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace("module_1", "module_name"), code_object_pickler.get_code_object_identifier( module_1_local_variable_added.my_function()).replace( "module_1_local_variable_added", "module_name"), ) def test_removing_local_variable_in_function_preserves_path(self): self.assertEqual( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace("module_1", "module_name"), code_object_pickler.get_code_object_identifier( module_1_local_variable_removed.my_function()).replace( "module_1_local_variable_removed", "module_name"), ) def test_adding_nested_function_in_function_preserves_path(self): self.assertEqual( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace("module_1", "module_name"), code_object_pickler.get_code_object_identifier( module_1_nested_function_added.my_function()).replace( "module_1_nested_function_added", "module_name"), ) def test_adding_nested_function_2_in_function_preserves_path(self): self.assertEqual( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace("module_1", "module_name"), code_object_pickler.get_code_object_identifier( module_1_nested_function_2_added.my_function()).replace( "module_1_nested_function_2_added", "module_name"), ) def test_adding_class_preserves_path(self): self.assertEqual( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace("module_1", "module_name"), code_object_pickler.get_code_object_identifier( module_1_class_added.my_function()).replace( "module_1_class_added", "module_name"), ) def test_adding_lambda_variable_in_function_changes_path(self): self.assertNotEqual( code_object_pickler.get_code_object_identifier( module_1.my_function()).replace("module_1", "module_name"), code_object_pickler.get_code_object_identifier( module_1_lambda_variable_added.my_function()).replace( "module_1_lambda_variable_added", "module_name"), ) def test_removing_lambda_variable_in_function_changes_path(self): self.assertNotEqual( code_object_pickler.get_code_object_identifier( module_3.my_function()).replace("module_3", "module_name"), code_object_pickler.get_code_object_identifier( module_3_modified.my_function()).replace( "module_3_modified", "module_name"), ) if __name__ == "__main__": unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/dill_pickler.py0000644000175100017510000003702015206113032022502 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Pickler for values, functions, and classes. For internal use only. No backwards compatibility guarantees. Pickles created by the pickling library contain non-ASCII characters, so we base64-encode the results so that we can put them in a JSON objects. The pickler is used to embed FlatMap callable objects into the workflow JSON description. The pickler module should be used to pickle functions and modules; for values, the coders.*PickleCoder classes should be used instead. """ # pytype: skip-file import base64 import bz2 import logging import sys import threading import traceback import types import zlib from typing import Any import dill from apache_beam.internal.code_object_pickler import get_normalized_path from apache_beam.internal.set_pickler import save_frozenset from apache_beam.internal.set_pickler import save_set settings = {'dill_byref': None} patch_save_code = dill.__version__ == "0.3.1.1" if patch_save_code: # The following function is based on 'save_code' from 'dill' # Author: Mike McKerns (mmckerns @caltech and @uqfoundation) # Copyright (c) 2008-2015 California Institute of Technology. # Copyright (c) 2016-2023 The Uncertainty Quantification Foundation. # License: 3-clause BSD. The full license text is available at: # - https://github.com/uqfoundation/dill/blob/master/LICENSE # The following function is also based on 'save_codeobject' from 'cloudpickle' # Copyright (c) 2012, Regents of the University of California. # Copyright (c) 2009 `PiCloud, Inc. `_. # License: 3-clause BSD. The full license text is available at: # - https://github.com/cloudpipe/cloudpickle/blob/master/LICENSE from types import CodeType @dill.register(CodeType) def save_code(pickler, obj): co_filename = get_normalized_path(obj.co_filename) if hasattr(obj, "co_endlinetable"): # python 3.11a (20 args) args = ( obj.co_argcount, obj.co_posonlyargcount, obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize, obj.co_flags, obj.co_code, obj.co_consts, obj.co_names, obj.co_varnames, co_filename, obj.co_name, obj.co_qualname, obj.co_firstlineno, obj.co_linetable, obj.co_endlinetable, obj.co_columntable, obj.co_exceptiontable, obj.co_freevars, obj.co_cellvars) elif hasattr(obj, "co_exceptiontable"): # python 3.11 (18 args) args = ( obj.co_argcount, obj.co_posonlyargcount, obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize, obj.co_flags, obj.co_code, obj.co_consts, obj.co_names, obj.co_varnames, co_filename, obj.co_name, obj.co_qualname, obj.co_firstlineno, obj.co_linetable, obj.co_exceptiontable, obj.co_freevars, obj.co_cellvars) elif hasattr(obj, "co_linetable"): # python 3.10 (16 args) args = ( obj.co_argcount, obj.co_posonlyargcount, obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize, obj.co_flags, obj.co_code, obj.co_consts, obj.co_names, obj.co_varnames, co_filename, obj.co_name, obj.co_firstlineno, obj.co_linetable, obj.co_freevars, obj.co_cellvars) elif hasattr(obj, "co_posonlyargcount"): # python 3.8 (16 args) args = ( obj.co_argcount, obj.co_posonlyargcount, obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize, obj.co_flags, obj.co_code, obj.co_consts, obj.co_names, obj.co_varnames, co_filename, obj.co_name, obj.co_firstlineno, obj.co_lnotab, obj.co_freevars, obj.co_cellvars) else: # python 3.7 (15 args) args = ( obj.co_argcount, obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize, obj.co_flags, obj.co_code, obj.co_consts, obj.co_names, obj.co_varnames, obj.co_filename, obj.co_name, obj.co_firstlineno, obj.co_lnotab, obj.co_freevars, obj.co_cellvars) pickler.save_reduce(CodeType, args, obj=obj) dill._dill.save_code = save_code class _NoOpContextManager(object): def __enter__(self): pass def __exit__(self, *unused_exc_info): pass # Pickling, especially unpickling, causes broken module imports on Python 3 # if executed concurrently, see: BEAM-8651, http://bugs.python.org/issue38884. _pickle_lock = threading.RLock() # Dill 0.28.0 renamed dill.dill to dill._dill: # https://github.com/uqfoundation/dill/commit/f0972ecc7a41d0b8acada6042d557068cac69baa # TODO: Remove this once Beam depends on dill >= 0.2.8 if not getattr(dill, 'dill', None): dill.dill = dill._dill sys.modules['dill.dill'] = dill._dill # TODO: Remove once Dataflow has containers with a preinstalled dill >= 0.2.8 if not getattr(dill, '_dill', None): dill._dill = dill.dill sys.modules['dill._dill'] = dill.dill dill_log = getattr(dill.dill, 'log', None) # dill v0.3.6 changed the attribute name from 'log' to 'logger' if not dill_log: dill_log = getattr(dill.dill, 'logger') def _is_nested_class(cls): """Returns true if argument is a class object that appears to be nested.""" return ( isinstance(cls, type) and cls.__module__ is not None and cls.__module__ != 'builtins' and cls.__name__ not in sys.modules[cls.__module__].__dict__) def _find_containing_class(nested_class): """Finds containing class of a nested class passed as argument.""" seen = set() def _find_containing_class_inner(outer): if outer in seen: return None seen.add(outer) for k, v in outer.__dict__.items(): if v is nested_class: return outer, k elif isinstance(v, type) and hasattr(v, '__dict__'): res = _find_containing_class_inner(v) if res: return res return _find_containing_class_inner(sys.modules[nested_class.__module__]) def _dict_from_mappingproxy(mp): d = mp.copy() d.pop('__dict__', None) d.pop('__prepare__', None) d.pop('__weakref__', None) return d def _nested_type_wrapper(fun): """A wrapper for the standard pickler handler for class objects. Args: fun: Original pickler handler for type objects. Returns: A wrapper for type objects that handles nested classes. The wrapper detects if an object being pickled is a nested class object. For nested class object only it will save the containing class object so the nested structure is recreated during unpickle. """ def wrapper(pickler, obj): # When the nested class is defined in the __main__ module we do not have to # do anything special because the pickler itself will save the constituent # parts of the type (i.e., name, base classes, dictionary) and then # recreate it during unpickling. if _is_nested_class(obj) and obj.__module__ != '__main__': containing_class_and_name = _find_containing_class(obj) if containing_class_and_name is not None: return pickler.save_reduce(getattr, containing_class_and_name, obj=obj) try: return fun(pickler, obj) except dill.dill.PicklingError: # pylint: disable=protected-access return pickler.save_reduce( dill.dill._create_type, ( type(obj), obj.__name__, obj.__bases__, _dict_from_mappingproxy(obj.__dict__)), obj=obj) # pylint: enable=protected-access return wrapper # Monkey patch the standard pickler dispatch table entry for type objects. # Dill, for certain types, defers to the standard pickler (including type # objects). We wrap the standard handler using type_wrapper() because # for nested class we want to pickle the actual enclosing class object so we # can recreate it during unpickling. # TODO(silviuc): Make sure we submit the fix upstream to GitHub dill project. dill.dill.Pickler.dispatch[type] = _nested_type_wrapper( dill.dill.Pickler.dispatch[type]) # Dill pickles generators objects without complaint, but unpickling produces # TypeError: object.__new__(generator) is not safe, use generator.__new__() # on some versions of Python. def _reject_generators(unused_pickler, unused_obj): raise TypeError("can't (safely) pickle generator objects") dill.dill.Pickler.dispatch[types.GeneratorType] = _reject_generators # This if guards against dill not being full initialized when generating docs. if 'save_module' in dir(dill.dill): # Always pickle non-main modules by name. old_save_module = dill.dill.save_module @dill.dill.register(dill.dill.ModuleType) def save_module(pickler, obj): if dill.dill.is_dill(pickler) and obj is pickler._main: return old_save_module(pickler, obj) else: dill_log.info('M2: %s' % obj) # pylint: disable=protected-access pickler.save_reduce( dill.dill._import_module, (obj.__name__, ), obj=obj) # pylint: enable=protected-access dill_log.info('# M2') # Pickle module dictionaries (commonly found in lambda's globals) # by referencing their module. old_save_module_dict = dill.dill.save_module_dict known_module_dicts: dict[int, tuple[types.ModuleType, dict[str, Any]]] = {} @dill.dill.register(dict) def new_save_module_dict(pickler, obj): obj_id = id(obj) if not known_module_dicts or '__file__' in obj or '__package__' in obj: if obj_id not in known_module_dicts: # Trigger loading of lazily loaded modules (such as pytest vendored # modules). # This pass over sys.modules needs to iterate on a copy of sys.modules # since lazy loading modifies the dictionary, hence the use of list(). for m in list(sys.modules.values()): try: _ = m.__dict__ except AttributeError: pass for m in list(sys.modules.values()): try: if (m and m.__name__ != '__main__' and isinstance(m, dill.dill.ModuleType)): d = m.__dict__ known_module_dicts[id(d)] = m, d except AttributeError: # Skip modules that do not have the __name__ attribute. pass if obj_id in known_module_dicts and dill.dill.is_dill(pickler): m = known_module_dicts[obj_id][0] try: # pylint: disable=protected-access dill.dill._import_module(m.__name__) return pickler.save_reduce( getattr, (known_module_dicts[obj_id][0], '__dict__'), obj=obj) except (ImportError, AttributeError): return old_save_module_dict(pickler, obj) else: return old_save_module_dict(pickler, obj) dill.dill.save_module_dict = new_save_module_dict def _nest_dill_logging(): """Prefix all dill logging with its depth in the callstack. Useful for debugging pickling of deeply nested structures. """ old_log_info = dill_log.info def new_log_info(msg, *args, **kwargs): old_log_info( ('1 2 3 4 5 6 7 8 9 0 ' * 10)[:len(traceback.extract_stack())] + msg, *args, **kwargs) dill_log.info = new_log_info # Turn off verbose logging from the dill pickler. logging.getLogger('dill').setLevel(logging.WARN) def dumps( o, enable_trace=True, use_zlib=False, enable_best_effort_determinism=False) -> bytes: """For internal use only; no backwards-compatibility guarantees.""" s = _dumps(o, enable_trace, enable_best_effort_determinism) # Compress as compactly as possible (compresslevel=9) to decrease peak memory # usage (of multiple in-memory copies) and to avoid hitting protocol buffer # limits. # WARNING: Be cautious about compressor change since it can lead to pipeline # representation change, and can break streaming job update compatibility on # runners such as Dataflow. if use_zlib: c = zlib.compress(s, 9) else: c = bz2.compress(s, compresslevel=9) del s # Free up some possibly large and no-longer-needed memory. return base64.b64encode(c) def _dumps(o, enable_trace=True, enable_best_effort_determinism=False) -> bytes: """For internal use only; no backwards-compatibility guarantees.""" with _pickle_lock: if enable_best_effort_determinism: old_save_set = dill.dill.Pickler.dispatch[set] old_save_frozenset = dill.dill.Pickler.dispatch[frozenset] dill.dill.pickle(set, save_set) dill.dill.pickle(frozenset, save_frozenset) try: s = dill.dumps(o, byref=settings['dill_byref']) except Exception: # pylint: disable=broad-except if enable_trace: dill.dill._trace(True) # pylint: disable=protected-access s = dill.dumps(o, byref=settings['dill_byref']) else: raise finally: dill.dill._trace(False) # pylint: disable=protected-access if enable_best_effort_determinism: dill.dill.pickle(set, old_save_set) dill.dill.pickle(frozenset, old_save_frozenset) return s def loads(encoded, enable_trace=True, use_zlib=False): """For internal use only; no backwards-compatibility guarantees.""" c = base64.b64decode(encoded) if use_zlib: s = zlib.decompress(c) else: s = bz2.decompress(c) del c # Free up some possibly large and no-longer-needed memory. return _loads(s, enable_trace) def _loads(s, enable_trace=True): with _pickle_lock: try: return dill.loads(s) except Exception: # pylint: disable=broad-except if enable_trace: dill.dill._trace(True) # pylint: disable=protected-access return dill.loads(s) else: raise finally: dill.dill._trace(False) # pylint: disable=protected-access def roundtrip(o): """Internal utility for testing round-trip pickle serialization.""" return _loads(_dumps(o)) def dump_session(file_path): """For internal use only; no backwards-compatibility guarantees. Pickle the current python session to be used in the worker. Note: Due to the inconsistency in the first dump of dill dump_session we create and load the dump twice to have consistent results in the worker and the running session. Check: https://github.com/uqfoundation/dill/issues/195 """ with _pickle_lock: dill.dump_session(file_path) dill.load_session(file_path) return dill.dump_session(file_path) def load_session(file_path): with _pickle_lock: return dill.load_session(file_path) def override_pickler_hooks(extend=True): """ Extends the dill library hooks into that of the standard pickler library. If false all hooks that dill overrides will be removed. If true dill hooks will be injected into the pickler library dispatch_table. """ dill.extend(extend) ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.2533944 apache_beam-2.74.0/apache_beam/internal/gcp/0000755000175100017510000000000015206113157020252 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/gcp/__init__.py0000644000175100017510000000152515206113032022356 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """For internal use only; no backwards-compatibility guarantees.""" ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/gcp/auth.py0000644000175100017510000001554415206113032021566 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Dataflow credentials and authentication.""" # pytype: skip-file import logging import socket import threading from typing import Optional from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.utils import retry # google.auth is only available when Beam is installed with the gcp extra. try: import google.auth import google_auth_httplib2 from google.auth import impersonated_credentials _GOOGLE_AUTH_AVAILABLE = True except ImportError: _GOOGLE_AUTH_AVAILABLE = False # When we are running in GCE, we can authenticate with VM credentials. is_running_in_gce = False # When we are running in GCE, this value is set based on worker startup # information. executing_project = None _LOGGER = logging.getLogger(__name__) def set_running_in_gce(worker_executing_project): """For internal use only; no backwards-compatibility guarantees. Informs the authentication library that we are running in GCE. When we are running in GCE, we have the option of using the VM metadata credentials for authentication to Google services. Args: worker_executing_project: The project running the workflow. This information comes from worker startup information. """ global is_running_in_gce global executing_project is_running_in_gce = True executing_project = worker_executing_project def get_service_credentials(pipeline_options): # type: (PipelineOptions) -> Optional[_ApitoolsCredentialsAdapter] """For internal use only; no backwards-compatibility guarantees. Get credentials to access Google services. Args: pipeline_options: Pipeline options, used in creating credentials like impersonated credentials. Returns: A ``_ApitoolsCredentialsAdapter`` object or None if credentials not found. Returned object is thread-safe. """ return _Credentials.get_service_credentials(pipeline_options) if _GOOGLE_AUTH_AVAILABLE: class _ApitoolsCredentialsAdapter: """For internal use only; no backwards-compatibility guarantees. Adapter allowing use of google-auth credentials with apitools, which normally expects credentials from the oauth2client library. This allows upgrading the auth library used by Beam without simultaneously upgrading all the GCP client libraries (a much larger change). """ def __init__(self, google_auth_credentials): self._google_auth_credentials = google_auth_credentials def authorize(self, http): """Return an http client authorized with the google-auth credentials. Args: http: httplib2.Http, an http object to be used to make the refresh request. Returns: google_auth_httplib2.AuthorizedHttp: An authorized http client. """ return google_auth_httplib2.AuthorizedHttp( self._google_auth_credentials, http=http) def __getattr__(self, attr): """Delegate attribute access to underlying google-auth credentials.""" return getattr(self._google_auth_credentials, attr) def get_google_auth_credentials(self): return self._google_auth_credentials class _Credentials(object): _credentials_lock = threading.Lock() _credentials_init = False _credentials = None @classmethod def get_service_credentials(cls, pipeline_options): # type: (PipelineOptions) -> Optional[_ApitoolsCredentialsAdapter] with cls._credentials_lock: if cls._credentials_init: return cls._credentials # apitools use urllib with the global timeout. Set it to 60 seconds # to prevent network related stuckness issues. if not socket.getdefaulttimeout(): _LOGGER.debug("Setting socket default timeout to 60 seconds.") socket.setdefaulttimeout(60) _LOGGER.debug( "socket default timeout is %s seconds.", socket.getdefaulttimeout()) cls._credentials = cls._get_service_credentials(pipeline_options) cls._credentials_init = True return cls._credentials @staticmethod def _get_service_credentials(pipeline_options): # type: (PipelineOptions) -> Optional[_ApitoolsCredentialsAdapter] if not _GOOGLE_AUTH_AVAILABLE: _LOGGER.warning( 'Unable to find default credentials because the google-auth library ' 'is not available. Install the gcp extra (apache_beam[gcp]) to use ' 'Google default credentials. Connecting anonymously.') return None try: # pylint: disable=c-extension-no-member credentials = _Credentials._get_credentials_with_retrys(pipeline_options) credentials = _Credentials._add_impersonation_credentials( credentials, pipeline_options) credentials = _ApitoolsCredentialsAdapter(credentials) logging.debug( 'Connecting using Google Application Default ' 'Credentials.') return credentials except Exception as e: _LOGGER.warning( 'Unable to find default credentials to use: %s\n' 'Connecting anonymously. This is expected if no ' 'credentials are needed to access GCP resources.', e) return None @staticmethod @retry.with_exponential_backoff(num_retries=4, initial_delay_secs=2) def _get_credentials_with_retrys(pipeline_options): credentials, _ = google.auth.default( scopes=pipeline_options.view_as(GoogleCloudOptions).gcp_oauth_scopes) return credentials @staticmethod def _add_impersonation_credentials(credentials, pipeline_options): gcs_options = pipeline_options.view_as(GoogleCloudOptions) impersonate_service_account = gcs_options.impersonate_service_account scopes = gcs_options.gcp_oauth_scopes if impersonate_service_account: _LOGGER.info('Impersonating: %s', impersonate_service_account) impersonate_accounts = impersonate_service_account.split(',') target_principal = impersonate_accounts[-1] delegate_to = impersonate_accounts[0:-1] credentials = impersonated_credentials.Credentials( source_credentials=credentials, target_principal=target_principal, delegates=delegate_to, target_scopes=scopes, ) return credentials ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/gcp/auth_test.py0000644000175100017510000001070115206113032022613 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest import mock from apache_beam.internal.gcp import auth from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions try: import google.auth as gauth import google_auth_httplib2 # pylint: disable=unused-import except ImportError: gauth = None # type: ignore class MockLoggingHandler(logging.Handler): """Mock logging handler to check for expected logs.""" def __init__(self, *args, **kwargs): self.reset() logging.Handler.__init__(self, *args, **kwargs) def emit(self, record): self.messages[record.levelname.lower()].append(record.getMessage()) def reset(self): self.messages = { 'debug': [], 'info': [], 'warning': [], 'error': [], 'critical': [], } @unittest.skipIf(gauth is None, 'Google Auth dependencies are not installed') class AuthTest(unittest.TestCase): @mock.patch('google.auth.default') def test_auth_with_retrys(self, unused_mock_arg): pipeline_options = PipelineOptions() pipeline_options.view_as( GoogleCloudOptions).impersonate_service_account = False credentials = ('creds', 1) self.is_called = False def side_effect(scopes=None): if self.is_called: return credentials else: self.is_called = True raise IOError('Failed') google_auth_mock = mock.MagicMock() gauth.default = google_auth_mock google_auth_mock.side_effect = side_effect # _Credentials caches the actual credentials. # This resets it for idempotent tests. if auth._Credentials._credentials_init: auth._Credentials._credentials_init = False auth._Credentials._credentials = None returned_credentials = auth.get_service_credentials(pipeline_options) # _Credentials caches the actual credentials. # This resets it for idempotent tests. if auth._Credentials._credentials_init: auth._Credentials._credentials_init = False auth._Credentials._credentials = None self.assertEqual('creds', returned_credentials._google_auth_credentials) @mock.patch( 'apache_beam.internal.gcp.auth._Credentials._get_credentials_with_retrys') def test_auth_with_retrys_always_fail(self, unused_mock_arg): pipeline_options = PipelineOptions() pipeline_options.view_as( GoogleCloudOptions).impersonate_service_account = False loggerHandler = MockLoggingHandler() auth._LOGGER.addHandler(loggerHandler) #Remove call to retrying method, as otherwise test takes ~10 minutes to run def raise_(scopes=None): raise IOError('Failed') retry_auth_mock = mock.MagicMock() auth._Credentials._get_credentials_with_retrys = retry_auth_mock retry_auth_mock.side_effect = raise_ # _Credentials caches the actual credentials. # This resets it for idempotent tests. if auth._Credentials._credentials_init: auth._Credentials._credentials_init = False auth._Credentials._credentials = None returned_credentials = auth.get_service_credentials(pipeline_options) self.assertEqual(None, returned_credentials) self.assertEqual([ 'Unable to find default credentials to use: Failed\n' 'Connecting anonymously. This is expected if no credentials are ' 'needed to access GCP resources.' ], loggerHandler.messages.get('warning')) # _Credentials caches the actual credentials. # This resets it for idempotent tests. if auth._Credentials._credentials_init: auth._Credentials._credentials_init = False auth._Credentials._credentials = None auth._LOGGER.removeHandler(loggerHandler) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/gcp/json_value.py0000644000175100017510000001353215206113032022765 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """JSON conversion utility functions.""" # pytype: skip-file from apache_beam.options.value_provider import ValueProvider # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apitools.base.py import extra_types except ImportError: extra_types = None # pylint: enable=wrong-import-order, wrong-import-position _MAXINT64 = (1 << 63) - 1 _MININT64 = -(1 << 63) def get_typed_value_descriptor(obj): """For internal use only; no backwards-compatibility guarantees. Converts a basic type into a @type/value dictionary. Args: obj: A bytes, unicode, bool, int, or float to be converted. Returns: A dictionary containing the keys ``@type`` and ``value`` with the value for the ``@type`` of appropriate type. Raises: TypeError: if the Python object has a type that is not supported. """ if isinstance(obj, (bytes, str)): type_name = 'Text' elif isinstance(obj, bool): type_name = 'Boolean' elif isinstance(obj, int): type_name = 'Integer' elif isinstance(obj, float): type_name = 'Float' else: raise TypeError('Cannot get a type descriptor for %s.' % repr(obj)) return {'@type': 'http://schema.org/%s' % type_name, 'value': obj} def to_json_value(obj, with_type=False): """For internal use only; no backwards-compatibility guarantees. Converts Python objects into extra_types.JsonValue objects. Args: obj: Python object to be converted. Can be :data:`None`. with_type: If true then the basic types (``bytes``, ``unicode``, ``int``, ``float``, ``bool``) will be wrapped in ``@type:value`` dictionaries. Otherwise the straight value is encoded into a ``JsonValue``. Returns: A ``JsonValue`` object using ``JsonValue``, ``JsonArray`` and ``JsonObject`` types for the corresponding values, lists, or dictionaries. Raises: TypeError: if the Python object contains a type that is not supported. The types supported are ``str``, ``bool``, ``list``, ``tuple``, ``dict``, and ``None``. The Dataflow API requires JsonValue(s) in many places, and it is quite convenient to be able to specify these hierarchical objects using Python syntax. """ if obj is None: return extra_types.JsonValue(is_null=True) elif isinstance(obj, (list, tuple)): return extra_types.JsonValue( array_value=extra_types.JsonArray( entries=[to_json_value(o, with_type=with_type) for o in obj])) elif isinstance(obj, dict): json_object = extra_types.JsonObject() for k, v in obj.items(): json_object.properties.append( extra_types.JsonObject.Property( key=k, value=to_json_value(v, with_type=with_type))) return extra_types.JsonValue(object_value=json_object) elif with_type: return to_json_value(get_typed_value_descriptor(obj), with_type=False) elif isinstance(obj, str): return extra_types.JsonValue(string_value=obj) elif isinstance(obj, bytes): return extra_types.JsonValue(string_value=obj.decode('utf8')) elif isinstance(obj, bool): return extra_types.JsonValue(boolean_value=obj) elif isinstance(obj, int): if _MININT64 <= obj <= _MAXINT64: return extra_types.JsonValue(integer_value=obj) else: raise TypeError('Can not encode {} as a 64-bit integer'.format(obj)) elif isinstance(obj, float): return extra_types.JsonValue(double_value=obj) elif isinstance(obj, ValueProvider): if obj.is_accessible(): return to_json_value(obj.get()) return extra_types.JsonValue(is_null=True) else: raise TypeError('Cannot convert %s to a JSON value.' % repr(obj)) def from_json_value(v): """For internal use only; no backwards-compatibility guarantees. Converts ``extra_types.JsonValue`` objects into Python objects. Args: v: ``JsonValue`` object to be converted. Returns: A Python object structured as values, lists, and dictionaries corresponding to ``JsonValue``, ``JsonArray`` and ``JsonObject`` types. Raises: TypeError: if the ``JsonValue`` object contains a type that is not supported. The types supported are ``str``, ``bool``, ``list``, ``dict``, and ``None``. The Dataflow API returns JsonValue(s) in many places and it is quite convenient to be able to convert these hierarchical objects to much simpler Python objects. """ if isinstance(v, extra_types.JsonValue): if v.string_value is not None: return v.string_value elif v.boolean_value is not None: return v.boolean_value elif v.integer_value is not None: return v.integer_value elif v.double_value is not None: return v.double_value elif v.array_value is not None: return from_json_value(v.array_value) elif v.object_value is not None: return from_json_value(v.object_value) elif v.is_null: return None elif isinstance(v, extra_types.JsonArray): return [from_json_value(e) for e in v.entries] elif isinstance(v, extra_types.JsonObject): return {p.key: from_json_value(p.value) for p in v.properties} raise TypeError('Cannot convert %s from a JSON value.' % repr(v)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/gcp/json_value_test.py0000644000175100017510000000742115206113032024024 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the json_value module.""" # pytype: skip-file import unittest from apache_beam.internal.gcp.json_value import from_json_value from apache_beam.internal.gcp.json_value import to_json_value from apache_beam.options.value_provider import RuntimeValueProvider from apache_beam.options.value_provider import StaticValueProvider # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apitools.base.py.extra_types import JsonValue except ImportError: JsonValue = None # pylint: enable=wrong-import-order, wrong-import-position @unittest.skipIf(JsonValue is None, 'GCP dependencies are not installed') class JsonValueTest(unittest.TestCase): def test_string_to(self): self.assertEqual(JsonValue(string_value='abc'), to_json_value('abc')) def test_bytes_to(self): self.assertEqual(JsonValue(string_value='abc'), to_json_value(b'abc')) def test_true_to(self): self.assertEqual(JsonValue(boolean_value=True), to_json_value(True)) def test_false_to(self): self.assertEqual(JsonValue(boolean_value=False), to_json_value(False)) def test_int_to(self): self.assertEqual(JsonValue(integer_value=14), to_json_value(14)) def test_float_to(self): self.assertEqual(JsonValue(double_value=2.75), to_json_value(2.75)) def test_static_value_provider_to(self): svp = StaticValueProvider(str, 'abc') self.assertEqual(JsonValue(string_value=svp.value), to_json_value(svp)) def test_runtime_value_provider_to(self): RuntimeValueProvider.set_runtime_options(None) rvp = RuntimeValueProvider('arg', 123, int) self.assertEqual(JsonValue(is_null=True), to_json_value(rvp)) # Reset runtime options to avoid side-effects in other tests. RuntimeValueProvider.set_runtime_options(None) def test_none_to(self): self.assertEqual(JsonValue(is_null=True), to_json_value(None)) def test_string_from(self): self.assertEqual('WXYZ', from_json_value(to_json_value('WXYZ'))) def test_true_from(self): self.assertEqual(True, from_json_value(to_json_value(True))) def test_false_from(self): self.assertEqual(False, from_json_value(to_json_value(False))) def test_int_from(self): self.assertEqual(-27, from_json_value(to_json_value(-27))) def test_float_from(self): self.assertEqual(4.5, from_json_value(to_json_value(4.5))) def test_with_type(self): rt = from_json_value(to_json_value('abcd', with_type=True)) self.assertEqual('http://schema.org/Text', rt['@type']) self.assertEqual('abcd', rt['value']) def test_none_from(self): self.assertIsNone(from_json_value(to_json_value(None))) def test_large_integer(self): num = 1 << 35 self.assertEqual(num, from_json_value(to_json_value(num))) def test_long_value(self): num = 1 << 63 - 1 self.assertEqual(num, from_json_value(to_json_value(num))) def test_too_long_value(self): with self.assertRaises(TypeError): to_json_value(1 << 64) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/http_client.py0000644000175100017510000000462315206113032022365 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utility functions used for creating a common Http client from httplib2. For internal use only. No backwards compatibility guarantees. """ # pytype: skip-file import logging import os import re import httplib2 # This is the number of seconds the library will wait for GCS operations to # complete. DEFAULT_HTTP_TIMEOUT_SECONDS = 60 def proxy_info_from_environment_var(proxy_env_var): """Reads proxy info from the environment and converts to httplib2.ProxyInfo. Args: proxy_env_var: environment variable string to read, http_proxy or https_proxy (in lower case). Example: http://myproxy.domain.com:8080 Returns: httplib2.ProxyInfo constructed from the environment string. """ proxy_url = os.environ.get(proxy_env_var) if not proxy_url: return None proxy_protocol = proxy_env_var.lower().split('_')[0] if not re.match('^https?://', proxy_url, flags=re.IGNORECASE): logging.warning( "proxy_info_from_url requires a protocol, which is always " "http or https.") proxy_url = proxy_protocol + '://' + proxy_url return httplib2.proxy_info_from_url(proxy_url, method=proxy_protocol) def get_new_http(timeout_secs=DEFAULT_HTTP_TIMEOUT_SECONDS): """Creates and returns a new httplib2.Http instance. Returns: An initialized httplib2.Http instance. """ proxy_info = None for proxy_env_var in ['http_proxy', 'https_proxy']: if os.environ.get(proxy_env_var): proxy_info = proxy_info_from_environment_var(proxy_env_var) break # Use a non-infinite SSL timeout to avoid hangs during network flakiness. return httplib2.Http(proxy_info=proxy_info, timeout=timeout_secs) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/http_client_test.py0000644000175100017510000001106515206113032023422 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the http_client module.""" # pytype: skip-file import os import unittest import mock from httplib2 import ProxyInfo from apache_beam.internal.http_client import DEFAULT_HTTP_TIMEOUT_SECONDS from apache_beam.internal.http_client import get_new_http from apache_beam.internal.http_client import proxy_info_from_environment_var class HttpClientTest(unittest.TestCase): def test_proxy_from_env_http_with_port(self): with mock.patch.dict(os.environ, http_proxy='http://localhost:9000'): proxy_info = proxy_info_from_environment_var('http_proxy') expected = ProxyInfo(3, 'localhost', 9000) self.assertEqual(str(expected), str(proxy_info)) def test_proxy_from_env_https_with_port(self): with mock.patch.dict(os.environ, https_proxy='https://localhost:9000'): proxy_info = proxy_info_from_environment_var('https_proxy') expected = ProxyInfo(3, 'localhost', 9000) self.assertEqual(str(expected), str(proxy_info)) def test_proxy_from_env_http_without_port(self): with mock.patch.dict(os.environ, http_proxy='http://localhost'): proxy_info = proxy_info_from_environment_var('http_proxy') expected = ProxyInfo(3, 'localhost', 80) self.assertEqual(str(expected), str(proxy_info)) def test_proxy_from_env_https_without_port(self): with mock.patch.dict(os.environ, https_proxy='https://localhost'): proxy_info = proxy_info_from_environment_var('https_proxy') expected = ProxyInfo(3, 'localhost', 443) self.assertEqual(str(expected), str(proxy_info)) def test_proxy_from_env_http_without_method(self): with mock.patch.dict(os.environ, http_proxy='localhost:8000'): proxy_info = proxy_info_from_environment_var('http_proxy') expected = ProxyInfo(3, 'localhost', 8000) self.assertEqual(str(expected), str(proxy_info)) def test_proxy_from_env_https_without_method(self): with mock.patch.dict(os.environ, https_proxy='localhost:8000'): proxy_info = proxy_info_from_environment_var('https_proxy') expected = ProxyInfo(3, 'localhost', 8000) self.assertEqual(str(expected), str(proxy_info)) def test_proxy_from_env_http_without_port_without_method(self): with mock.patch.dict(os.environ, http_proxy='localhost'): proxy_info = proxy_info_from_environment_var('http_proxy') expected = ProxyInfo(3, 'localhost', 80) self.assertEqual(str(expected), str(proxy_info)) def test_proxy_from_env_https_without_port_without_method(self): with mock.patch.dict(os.environ, https_proxy='localhost'): proxy_info = proxy_info_from_environment_var('https_proxy') expected = ProxyInfo(3, 'localhost', 443) self.assertEqual(str(expected), str(proxy_info)) def test_proxy_from_env_invalid_var(self): proxy_info = proxy_info_from_environment_var('http_proxy_host') expected = None self.assertEqual(str(expected), str(proxy_info)) def test_proxy_from_env_wrong_method_in_var_name(self): with mock.patch.dict(os.environ, smtp_proxy='localhost'): with self.assertRaises(KeyError): proxy_info_from_environment_var('smtp_proxy') def test_proxy_from_env_wrong_method_in_url(self): with mock.patch.dict(os.environ, http_proxy='smtp://localhost:8000'): proxy_info = proxy_info_from_environment_var('http_proxy') expected = ProxyInfo(3, 'smtp', 80) # wrong proxy info generated self.assertEqual(str(expected), str(proxy_info)) def test_get_new_http_proxy_info(self): with mock.patch.dict(os.environ, http_proxy='localhost'): http = get_new_http() expected = ProxyInfo(3, 'localhost', 80) self.assertEqual(str(http.proxy_info), str(expected)) def test_get_new_http_timeout(self): http = get_new_http() self.assertEqual(http.timeout, DEFAULT_HTTP_TIMEOUT_SECONDS) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.2533944 apache_beam-2.74.0/apache_beam/internal/metrics/0000755000175100017510000000000015206113157021147 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/metrics/__init__.py0000644000175100017510000000152515206113032023253 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """For internal use only; no backwards-compatibility guarantees.""" ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/metrics/metric.py0000644000175100017510000001706115206113032023001 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Metrics API classes for internal use only. Users should use apache_beam.metrics.metric package instead. For internal use only. No backwards compatibility guarantees. """ # pytype: skip-file # mypy: disallow-untyped-defs import datetime import logging import threading import time from typing import TYPE_CHECKING from typing import Optional from typing import Union from apache_beam.metrics import monitoring_infos from apache_beam.metrics.metric import Metrics as UserMetrics from apache_beam.metrics.metricbase import MetricName if TYPE_CHECKING: from apache_beam.metrics.cells import MetricCell from apache_beam.metrics.cells import MetricCellFactory # Protect against environments where bigquery library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apitools.base.py.exceptions import HttpError except ImportError: pass __all__ = ['Metrics'] _LOGGER = logging.getLogger(__name__) class Metrics(object): @staticmethod def counter( urn: str, labels: Optional[dict[str, str]] = None, process_wide: bool = False) -> UserMetrics.DelegatingCounter: """Obtains or creates a Counter metric. Args: namespace: A class or string that gives the namespace to a metric name: A string that gives a unique name to a metric urn: URN to populate on a MonitoringInfo, when sending to RunnerHarness. labels: Labels to populate on a MonitoringInfo process_wide: Whether or not the metric is specific to the current bundle or should be calculated for the entire process. Returns: A Counter object. """ return UserMetrics.DelegatingCounter( MetricName(namespace=None, name=None, urn=urn, labels=labels), process_wide=process_wide) class MetricLogger(object): """Simple object to locally aggregate and log metrics.""" def __init__(self) -> None: self._metric: dict[MetricName, 'MetricCell'] = {} self._lock = threading.Lock() self._last_logging_millis = int(time.time() * 1000) self.minimum_logging_frequency_msec = 180000 def update( self, cell_type: Union[type['MetricCell'], 'MetricCellFactory'], metric_name: MetricName, value: object) -> None: cell = self._get_metric_cell(cell_type, metric_name) cell.update(value) def _get_metric_cell( self, cell_type: Union[type['MetricCell'], 'MetricCellFactory'], metric_name: MetricName) -> 'MetricCell': with self._lock: if metric_name not in self._metric: self._metric[metric_name] = cell_type() return self._metric[metric_name] def log_metrics(self, reset_after_logging: bool = False) -> None: if self._lock.acquire(False): try: current_millis = int(time.time() * 1000) if ((current_millis - self._last_logging_millis) > self.minimum_logging_frequency_msec): logging_metric_info = [ '[Locally aggregated metrics since %s]' % datetime.datetime.fromtimestamp( self._last_logging_millis / 1000.0) ] for name, cell in self._metric.items(): logging_metric_info.append('%s: %s' % (name, cell.get_cumulative())) _LOGGER.info('\n'.join(logging_metric_info)) if reset_after_logging: self._metric = {} self._last_logging_millis = current_millis finally: self._lock.release() class ServiceCallMetric(object): """Metric class which records Service API call metrics. This class will capture a request count metric for the specified request_count_urn and base_labels. When call() is invoked the status must be provided, which will be converted to a canonical GCP status code, if possible. TODO(ajamato): Add Request latency metric. """ def __init__( self, request_count_urn: str, base_labels: Optional[dict[str, str]] = None) -> None: self.base_labels = base_labels if base_labels else {} self.request_count_urn = request_count_urn def call(self, status: Union[int, str, 'HttpError']) -> None: """Record the status of the call into appropriate metrics.""" canonical_status = self.convert_to_canonical_status_string(status) additional_labels = {monitoring_infos.STATUS_LABEL: canonical_status} labels = dict( list(self.base_labels.items()) + list(additional_labels.items())) request_counter = Metrics.counter( urn=self.request_count_urn, labels=labels, process_wide=True) request_counter.inc() def convert_to_canonical_status_string( self, status: Union[int, str, 'HttpError']) -> str: """Converts a status to a canonical GCP status cdoe string.""" http_status_code = None if isinstance(status, int): http_status_code = status elif isinstance(status, str): return status.lower() elif isinstance(status, HttpError): http_status_code = int(status.status_code) http_to_canonical_gcp_status = { 200: 'ok', 400: 'out_of_range', 401: 'unauthenticated', 403: 'permission_denied', 404: 'not_found', 409: 'already_exists', 429: 'resource_exhausted', 499: 'cancelled', 500: 'internal', 501: 'not_implemented', 503: 'unavailable', 504: 'deadline_exceeded' } if (http_status_code is not None and http_status_code in http_to_canonical_gcp_status): return http_to_canonical_gcp_status[http_status_code] return str(http_status_code) @staticmethod def bigtable_error_code_to_grpc_status_string( grpc_status_code: Optional[int]) -> str: """ Converts the bigtable error code to a canonical GCP status code string. This Bigtable client library is not using the canonical http status code values (i.e. https://cloud.google.com/apis/design/errors)" Instead they are numbered using an enum with these values corresponding to each status code: https://cloud.google.com/bigtable/docs/status-codes Args: grpc_status_code: An int that corresponds to an enum of status codes Returns: A GCP status code string """ grpc_to_canonical_gcp_status = { 0: 'ok', 1: 'cancelled', 2: 'unknown', 3: 'invalid_argument', 4: 'deadline_exceeded', 5: 'not_found', 6: 'already_exists', 7: 'permission_denied', 8: 'resource_exhausted', 9: 'failed_precondition', 10: 'aborted', 11: 'out_of_range', 12: 'unimplemented', 13: 'internal', 14: 'unavailable' } if grpc_status_code is None: # Bigtable indicates this can be retried but itself has exhausted retry # timeout or there is no retry policy set for bigtable. return grpc_to_canonical_gcp_status[4] return grpc_to_canonical_gcp_status.get( grpc_status_code, str(grpc_status_code)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/module_test.py0000644000175100017510000000515015206113032022370 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module used to define functions and classes used by the coder unit tests.""" # pytype: skip-file import re import sys from typing import Any GLOBAL_DICT = {} def mutable_test_function(): def dynamic_function(): return 'version1' return dynamic_function class UnPicklable: def __init__(self, x): self.x = x def __reduce__(self): return NotImplemented UNPICKLABLE_INSTANCE = UnPicklable(5) def closure_contains_unpicklable(): def inner(): return UNPICKLABLE_INSTANCE return inner def closure_contains_unpicklable_imports_self(): def inner(): import apache_beam.internal.module_test as self_module return self_module.UNPICKLABLE_INSTANCE return inner def fn_returns_unpicklable(): return UNPICKLABLE_INSTANCE class TopClass(object): class NestedClass(object): def __init__(self, datum): self.datum = 'X:%s' % datum class MiddleClass(object): class NestedClass(object): def __init__(self, datum): self.datum = 'Y:%s' % datum def get_lambda_with_globals(): return lambda s: re.findall(r'\w+', s) def get_lambda_with_closure(message): return lambda: 'closure: %s' % message class Xyz(object): """A class to be pickled.""" def foo(self, s): return re.findall(r'\w+', s) def create_class(datum): """Creates an unnamable class to be pickled.""" class Z(object): def get(self): return 'Z:%s' % datum return Z() XYZ_OBJECT = Xyz() class RecursiveClass(object): """A class that contains a reference to itself.""" SELF_TYPE: Any = None def __init__(self, datum): self.datum = 'RecursiveClass:%s' % datum RecursiveClass.SELF_TYPE = RecursiveClass # pylint: disable=exec-used if sys.version_info >= (3, 7): # create dataclass to be pickled exec( ''' from dataclasses import dataclass @dataclass class DataClass: datum: str''') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/pickler.py0000644000175100017510000001034415206113032021476 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Pickler for values, functions, and classes. For internal use only. No backwards compatibility guarantees. Pickles created by the pickling library contain non-ASCII characters, so we base64-encode the results so that we can put them in a JSON objects. The pickler is used to embed FlatMap callable objects into the workflow JSON description. The pickler module should be used to pickle functions and modules; for values, the coders.*PickleCoder classes should be used instead. """ from apache_beam.internal import cloudpickle_pickler try: from apache_beam.internal import dill_pickler except ImportError: dill_pickler = None # type: ignore[assignment] USE_CLOUDPICKLE = 'cloudpickle' USE_DILL = 'dill' USE_DILL_UNSAFE = 'dill_unsafe' DEFAULT_PICKLE_LIB = USE_CLOUDPICKLE desired_pickle_lib = cloudpickle_pickler def dumps( o, enable_trace=True, use_zlib=False, enable_best_effort_determinism=False, enable_stable_code_identifier_pickling=False) -> bytes: if (desired_pickle_lib == cloudpickle_pickler): return cloudpickle_pickler.dumps( o, enable_trace=enable_trace, use_zlib=use_zlib, enable_best_effort_determinism=enable_best_effort_determinism, enable_stable_code_identifier_pickling= enable_stable_code_identifier_pickling, ) return desired_pickle_lib.dumps( o, enable_trace=enable_trace, use_zlib=use_zlib, enable_best_effort_determinism=enable_best_effort_determinism) def loads(encoded, enable_trace=True, use_zlib=False): """For internal use only; no backwards-compatibility guarantees.""" return desired_pickle_lib.loads( encoded, enable_trace=enable_trace, use_zlib=use_zlib) def roundtrip(o): """Internal utility for testing round-trip pickle serialization.""" return desired_pickle_lib.roundtrip(o) def dump_session(file_path): """For internal use only; no backwards-compatibility guarantees. Pickle the current python session to be used in the worker. """ return desired_pickle_lib.dump_session(file_path) def load_session(file_path): return desired_pickle_lib.load_session(file_path) def is_currently_dill(): return desired_pickle_lib == dill_pickler def is_currently_cloudpickle(): return desired_pickle_lib == cloudpickle_pickler def set_library(selected_library=DEFAULT_PICKLE_LIB): """ Sets pickle library that will be used. """ global desired_pickle_lib if selected_library == 'default': selected_library = DEFAULT_PICKLE_LIB if selected_library == USE_DILL and not dill_pickler: raise ImportError( "Pipeline option pickle_library=dill is set, but dill is not " "installed. Install apache-beam with the dill extras package " "e.g. apache-beam[dill].") if selected_library == USE_DILL_UNSAFE and not dill_pickler: raise ImportError( "Pipeline option pickle_library=dill_unsafe is set, but dill is not " "installed. Install dill in job submission and runtime environments.") dill_is_requested = ( selected_library == USE_DILL or selected_library == USE_DILL_UNSAFE) # If switching to or from dill, update the pickler hook overrides. if is_currently_dill() != dill_is_requested: dill_pickler.override_pickler_hooks(selected_library == USE_DILL) if dill_is_requested: desired_pickle_lib = dill_pickler elif selected_library == USE_CLOUDPICKLE: desired_pickle_lib = cloudpickle_pickler else: raise ValueError(f'Unknown pickler library: {selected_library}') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/pickler_test.py0000644000175100017510000002665015206113032022544 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the pickler module.""" # pytype: skip-file import random import sys import textwrap import threading import types import unittest import pytest from parameterized import param from parameterized import parameterized from apache_beam.internal import module_test from apache_beam.internal import pickler from apache_beam.internal.pickler import dumps from apache_beam.internal.pickler import loads def maybe_skip_if_no_dill(pickle_library): if pickle_library == 'dill': pytest.importorskip("dill") @pytest.mark.uses_dill class PicklerTest(unittest.TestCase): NO_MAPPINGPROXYTYPE = not hasattr(types, "MappingProxyType") @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_basics(self, pickle_lib): maybe_skip_if_no_dill(pickle_lib) pickler.set_library(pickle_lib) self.assertEqual([1, 'a', ('z', )], loads(dumps([1, 'a', ('z', )]))) fun = lambda x: 'xyz-%s' % x self.assertEqual('xyz-abc', loads(dumps(fun))('abc')) @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_lambda_with_globals(self, pickle_lib): """Tests that the globals of a function are preserved.""" maybe_skip_if_no_dill(pickle_lib) pickler.set_library(pickle_lib) # The point of the test is that the lambda being called after unpickling # relies on having the re module being loaded. self.assertEqual(['abc', 'def'], loads(dumps( module_test.get_lambda_with_globals()))('abc def')) @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_lambda_with_main_globals(self, pickle_lib): maybe_skip_if_no_dill(pickle_lib) pickler.set_library(pickle_lib) self.assertEqual(unittest, loads(dumps(lambda: unittest))()) @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_lambda_with_closure(self, pickle_lib): """Tests that the closure of a function is preserved.""" maybe_skip_if_no_dill(pickle_lib) pickler.set_library(pickle_lib) self.assertEqual( 'closure: abc', loads(dumps(module_test.get_lambda_with_closure('abc')))()) @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_class(self, pickle_lib): """Tests that a class object is pickled correctly.""" maybe_skip_if_no_dill(pickle_lib) pickler.set_library(pickle_lib) self.assertEqual(['abc', 'def'], loads(dumps(module_test.Xyz))().foo('abc def')) @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_object(self, pickle_lib): """Tests that a class instance is pickled correctly.""" maybe_skip_if_no_dill(pickle_lib) pickler.set_library(pickle_lib) self.assertEqual(['abc', 'def'], loads(dumps(module_test.XYZ_OBJECT)).foo('abc def')) @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_nested_class(self, pickle_lib): """Tests that a nested class object is pickled correctly.""" maybe_skip_if_no_dill(pickle_lib) pickler.set_library(pickle_lib) self.assertEqual( 'X:abc', loads(dumps(module_test.TopClass.NestedClass('abc'))).datum) self.assertEqual( 'Y:abc', loads(dumps(module_test.TopClass.MiddleClass.NestedClass('abc'))).datum) @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_dynamic_class(self, pickle_lib): """Tests that a nested class object is pickled correctly.""" maybe_skip_if_no_dill(pickle_lib) pickler.set_library(pickle_lib) self.assertEqual( 'Z:abc', loads(dumps(module_test.create_class('abc'))).get()) @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_generators(self, pickle_lib): maybe_skip_if_no_dill(pickle_lib) pickler.set_library(pickle_lib) with self.assertRaises(TypeError): dumps((_ for _ in range(10))) @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_recursive_class(self, pickle_lib): maybe_skip_if_no_dill(pickle_lib) pickler.set_library(pickle_lib) self.assertEqual( 'RecursiveClass:abc', loads(dumps(module_test.RecursiveClass('abc').datum))) @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_pickle_rlock(self, pickle_lib): maybe_skip_if_no_dill(pickle_lib) pickler.set_library(pickle_lib) rlock_instance = threading.RLock() rlock_type = type(rlock_instance) self.assertIsInstance(loads(dumps(rlock_instance)), rlock_type) @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_save_paths(self, pickle_lib): maybe_skip_if_no_dill(pickle_lib) pickler.set_library(pickle_lib) f = loads(dumps(lambda x: x)) co_filename = f.__code__.co_filename self.assertTrue(co_filename.endswith('pickler_test.py')) @unittest.skipIf(NO_MAPPINGPROXYTYPE, 'test if MappingProxyType introduced') @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_dump_and_load_mapping_proxy(self, pickle_lib): maybe_skip_if_no_dill(pickle_lib) pickler.set_library(pickle_lib) self.assertEqual( 'def', loads(dumps(types.MappingProxyType({'abc': 'def'})))['abc']) self.assertEqual( types.MappingProxyType, type(loads(dumps(types.MappingProxyType({}))))) # pylint: disable=exec-used @unittest.skipIf(sys.version_info < (3, 7), 'Python 3.7 or above only') @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_dataclass(self, pickle_lib): maybe_skip_if_no_dill(pickle_lib) exec( ''' from apache_beam.internal.module_test import DataClass self.assertEqual(DataClass(datum='abc'), loads(dumps(DataClass(datum='abc')))) ''') @parameterized.expand([ param(pickle_lib='dill'), param(pickle_lib='cloudpickle'), ]) def test_class_states_not_changed_at_subsequent_loading(self, pickle_lib): maybe_skip_if_no_dill(pickle_lib) pickler.set_library(pickle_lib) class Local: def fun(self): pass obj = Local() serialized = dumps(obj) obj2 = loads(serialized) obj2_type_id = id(type(obj2)) obj2_func_id = id(obj2.fun.__func__) # The obj and obj2 may have different classes (dill) or the same class ( # cloudpickle). The cloudpickle tracks the weak references of # dumped/loaded dynamic classes with class_tracker_id, reusing the same # class as possible if the original class is the same. Besides, # dill creates a new class for each loaded object. loads(serialized) # obj3 # The obj2 and obj3 may have different classes (dill) or the same class ( # cloudpickle) However, the obj2's class states must not be changed after # obj3 loading. # https://github.com/apache/beam/issues/35062 self.assertEqual( obj2_type_id, id(type(obj2)), 'Class must not changed after subsequent loading.') self.assertEqual( obj2_func_id, id(obj2.fun.__func__), 'Class states must not changed after subsequent loading.') def maybe_get_sets_with_different_iteration_orders(self): # Use a mix of types in an attempt to create sets with the same elements # whose iteration order is different. elements = [ 100, 'hello', 3.14159, True, None, -50, 'world', False, (1, 2), (4, 3), ('hello', 'world') ] set1 = set(elements) # Try random addition orders until finding an order that works. for _ in range(100): set2 = set() random.shuffle(elements) for e in elements: set2.add(e) if list(set1) != list(set2): break return set1, set2 def test_best_effort_determinism(self): maybe_skip_if_no_dill('dill') pickler.set_library('dill') set1, set2 = self.maybe_get_sets_with_different_iteration_orders() self.assertEqual( dumps(set1, enable_best_effort_determinism=True), dumps(set2, enable_best_effort_determinism=True)) # The test relies on the sets having different iteration orders for the # elements. Iteration order is implementation dependent and undefined, # meaning the test won't always be able to setup these conditions. if list(set1) == list(set2): self.skipTest('Set iteration orders matched. Test results inconclusive.') def test_disable_best_effort_determinism(self): maybe_skip_if_no_dill('dill') pickler.set_library('dill') set1, set2 = self.maybe_get_sets_with_different_iteration_orders() # The test relies on the sets having different iteration orders for the # elements. Iteration order is implementation dependent and undefined, # meaning the test won't always be able to setup these conditions. if list(set1) == list(set2): self.skipTest('Set iteration orders matched. Unable to complete test.') self.assertNotEqual( dumps(set1, enable_best_effort_determinism=False), dumps(set2, enable_best_effort_determinism=False)) def test_stable_identifier_uses_current_code(self): pickler.set_library('cloudpickle') # Get original dynamic function func_v1 = module_test.mutable_test_function() pickled_stable = pickler.dumps( func_v1, enable_stable_code_identifier_pickling=True) pickled_frozen = pickler.dumps( func_v1, enable_stable_code_identifier_pickling=False) # Save original function for cleanup original_function = module_test.mutable_test_function try: # Monkey patch: Replace the entire outer function with v2 code_v2 = textwrap.dedent( """ def mutable_test_function(): def dynamic_function(): return "version2" return dynamic_function """) namespace = {} exec(code_v2, namespace) module_test.mutable_test_function = namespace['mutable_test_function'] # Unpickle both func_stable = pickler.loads(pickled_stable) func_frozen = pickler.loads(pickled_frozen) # Stable identifier resolves to NEW code (version2) self.assertEqual('version2', func_stable()) # Frozen bytecode uses OLD code (version1) self.assertEqual('version1', func_frozen()) finally: # Restore original function module_test.mutable_test_function = original_function if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/set_pickler.py0000644000175100017510000001317415206113032022355 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Custom pickling logic for sets to make the serialization semi-deterministic. To make set serialization semi-deterministic, we must pick an order for the set elements. Sets may contain elements of types not defining a comparison "<" operator. To provide an order, we define our own custom comparison function which supports elements of near-arbitrary types and use that to sort the contents of each set during serialization. Attempts at determinism are made on a best-effort basis to improve hit rates for cached workflows and the ordering does not define a total order for all values. """ import enum import functools def compare(lhs, rhs): """Returns -1, 0, or 1 depending on whether lhs <, =, or > rhs.""" if lhs < rhs: return -1 elif lhs > rhs: return 1 else: return 0 def generic_object_comparison(lhs, rhs, lhs_path, rhs_path, max_depth): """Identifies which object goes first in an (almost) total order of objects. Args: lhs: An arbitrary Python object or built-in type. rhs: An arbitrary Python object or built-in type. lhs_path: Traversal path from the root lhs object up to, but not including, lhs. The original contents of lhs_path are restored before the function returns. rhs_path: Same as lhs_path except for the rhs. max_depth: Maximum recursion depth. Returns: -1, 0, or 1 depending on whether lhs or rhs goes first in the total order. 0 if max_depth is exhausted. 0 if lhs is in lhs_path or rhs is in rhs_path (there is a cycle). """ if id(lhs) == id(rhs): # Fast path return 0 if type(lhs) != type(rhs): return compare(str(type(lhs)), str(type(rhs))) if type(lhs) in [int, float, bool, str, bool, bytes, bytearray]: return compare(lhs, rhs) if isinstance(lhs, enum.Enum): # Enums can have values with arbitrary types. The names are strings. return compare(lhs.name, rhs.name) # To avoid exceeding the recursion depth limit, set a limit on recursion. max_depth -= 1 if max_depth < 0: return 0 # Check for cycles in the traversal path to avoid getting stuck in a loop. if id(lhs) in lhs_path or id(rhs) in rhs_path: return 0 lhs_path.append(id(lhs)) rhs_path.append(id(rhs)) # The comparison logic is split across two functions to simplifying updating # and restoring the traversal paths. result = _generic_object_comparison_recursive_path( lhs, rhs, lhs_path, rhs_path, max_depth) lhs_path.pop() rhs_path.pop() return result def _generic_object_comparison_recursive_path( lhs, rhs, lhs_path, rhs_path, max_depth): if type(lhs) == tuple or type(lhs) == list: result = compare(len(lhs), len(rhs)) if result != 0: return result for i in range(len(lhs)): result = generic_object_comparison( lhs[i], rhs[i], lhs_path, rhs_path, max_depth) if result != 0: return result return 0 if type(lhs) == frozenset or type(lhs) == set: return generic_object_comparison( tuple(sort_if_possible(lhs, lhs_path, rhs_path, max_depth)), tuple(sort_if_possible(rhs, lhs_path, rhs_path, max_depth)), lhs_path, rhs_path, max_depth) if type(lhs) == dict: lhs_keys = list(lhs.keys()) rhs_keys = list(rhs.keys()) result = compare(len(lhs_keys), len(rhs_keys)) if result != 0: return result lhs_keys = sort_if_possible(lhs_keys, lhs_path, rhs_path, max_depth) rhs_keys = sort_if_possible(rhs_keys, lhs_path, rhs_path, max_depth) for lhs_key, rhs_key in zip(lhs_keys, rhs_keys): result = generic_object_comparison( lhs_key, rhs_key, lhs_path, rhs_path, max_depth) if result != 0: return result result = generic_object_comparison( lhs[lhs_key], rhs[rhs_key], lhs_path, rhs_path, max_depth) if result != 0: return result lhs_fields = dir(lhs) rhs_fields = dir(rhs) result = compare(len(lhs_fields), len(rhs_fields)) if result != 0: return result for i in range(len(lhs_fields)): result = compare(lhs_fields[i], rhs_fields[i]) if result != 0: return result result = generic_object_comparison( getattr(lhs, lhs_fields[i], None), getattr(rhs, rhs_fields[i], None), lhs_path, rhs_path, max_depth) if result != 0: return result return 0 def sort_if_possible(obj, lhs_path=None, rhs_path=None, max_depth=4): def cmp(lhs, rhs): if lhs_path is None: # Start the traversal at the root call to cmp. return generic_object_comparison(lhs, rhs, [], [], max_depth) else: # Continue the existing traversal path for recursive calls to cmp. return generic_object_comparison(lhs, rhs, lhs_path, rhs_path, max_depth) return sorted(obj, key=functools.cmp_to_key(cmp)) def save_set(pickler, obj): pickler.save_set(sort_if_possible(obj)) def save_frozenset(pickler, obj): pickler.save_frozenset(sort_if_possible(obj)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/set_pickler_test.py0000644000175100017510000001573715206113032023423 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for best-effort deterministic pickling of sets.""" import enum import unittest from apache_beam.internal.set_pickler import sort_if_possible class A: def __init__(self, x, y): self.x = x self.y = y class B: pass class C: pass class SortIfPossibleTest(unittest.TestCase): def test_order_by_type(self): a = A(1, 2) self.assertEqual( sort_if_possible([123, "abc", True, a]), # A, bool, int, str [a, True, 123, "abc"], ) def test_sorts_ints(self): self.assertEqual(sort_if_possible({5, 2, 3, 1, 4}), [1, 2, 3, 4, 5]) def test_sorts_booleans(self): self.assertEqual(sort_if_possible({False, True}), [False, True]) def test_sorts_floats(self): self.assertEqual(sort_if_possible({-18.0, 3.14, 2.71}), [-18.0, 2.71, 3.14]) def test_sorts_strings(self): self.assertEqual( sort_if_possible({"when", "in", "the", "course", "of"}), ["course", "in", "of", "the", "when"], ) def test_sorts_bytes(self): self.assertEqual( sort_if_possible({b"when", b"in", b"the", b"course", b"of"}), [b"course", b"in", b"of", b"the", b"when"], ) def test_sorts_bytearrays(self): f = bytearray self.assertEqual( sort_if_possible( [f(b"when"), f(b"in"), f(b"the"), f(b"course"), f(b"of")]), [f(b"course"), f(b"in"), f(b"of"), f(b"the"), f(b"when")], ) def test_sort_tuples_by_length(self): self.assertEqual( sort_if_possible({(1, 1, 1), (1, 1), (1, )}), [(1, ), (1, 1), (1, 1, 1)]) def test_sort_tuples_by_element_values(self): self.assertEqual( sort_if_possible({(0, 0), (1, 1), (0, 1), (1, 0)}), [(0, 0), (0, 1), (1, 0), (1, 1)], ) def test_sort_nested_tuples(self): self.assertEqual( sort_if_possible({(1, (4, )), (1, (1, )), (1, (3, )), (1, (2, ))}), [(1, (1, )), (1, (2, )), (1, (3, )), (1, (4, ))], ) def test_sort_lists_by_length(self): self.assertEqual( sort_if_possible([[1, 1, 1], [1, 1], [ 1, ]]), [[ 1, ], [1, 1], [1, 1, 1]]) def test_sort_lists_by_element_values(self): self.assertEqual( sort_if_possible([[0, 0], [1, 1], [0, 1], [1, 0]]), [[0, 0], [0, 1], [1, 0], [1, 1]], ) def test_sort_frozenset_like_sorted_tuple(self): self.assertEqual( sort_if_possible( {frozenset([1, 2, 3]), frozenset([1]), frozenset([1, 2, 4])}), [frozenset([1]), frozenset([1, 2, 3]), frozenset([1, 2, 4])], ) def test_sort_set_like_sorted_tuple(self): self.assertEqual( sort_if_possible([set([1, 2, 3]), set([1]), set([1, 2, 4])]), [set([1]), set([1, 2, 3]), set([1, 2, 4])], ) def test_order_objects_by_class_name(self): a = A(1, 2) b = B() c = C() self.assertEqual(sort_if_possible({b, c, a}), [a, b, c]) def test_order_objects_by_number_of_fields(self): o1 = C() o2 = C() setattr(o2, "f1", 1) o3 = C() setattr(o3, "f1", 1) setattr(o3, "f2", 2) self.assertEqual(sort_if_possible({o3, o2, o1}), [o1, o2, o3]) def test_order_objects_by_field_name(self): o1 = C() setattr(o1, "aaa", 1) o2 = C() setattr(o2, "bbb", 1) o3 = C() setattr(o3, "ccc", 1) self.assertEqual(sort_if_possible({o3, o2, o1}), [o1, o2, o3]) def test_order_objects_by_field_value(self): a1_1 = A(1, 1) a1_2 = A(1, 2) a2_1 = A(2, 1) a2_2 = A(2, 2) self.assertEqual( sort_if_possible({a2_1, a1_1, a2_2, a1_2}), [a1_1, a1_2, a2_1, a2_2]) def test_cyclic_data(self): def create_tuple_with_cycles(): o = C() t = (o, ) setattr(t[0], "t", t) return t t1 = create_tuple_with_cycles() t2 = create_tuple_with_cycles() t3 = create_tuple_with_cycles() actual = {hash(t) for t in sort_if_possible({t1, t2, t3})} expected = {hash(t1), hash(t2), hash(t3)} self.assertEqual(actual, expected) def test_order_dict_by_length(self): self.assertEqual( sort_if_possible([{ 'a': 1, 'b': 2 }, { 'a': 1 }, { 'a': 1, 'b': 2, 'c': 3 }]), [{ 'a': 1 }, { 'a': 1, 'b': 2 }, { 'a': 1, 'b': 2, 'c': 3 }]) def test_order_dict_by_key(self): self.assertEqual( sort_if_possible([{ 'b': 1 }, { 'a': 1 }, { 'c': 1 }]), [{ 'a': 1 }, { 'b': 1 }, { 'c': 1 }]) def test_order_dict_by_value(self): self.assertEqual( sort_if_possible([{ 'a': 2 }, { 'a': 1 }, { 'a': 3 }]), [{ 'a': 1 }, { 'a': 2 }, { 'a': 3 }]) def test_dict_keys_do_not_have_lt(self): self.assertEqual( sort_if_possible([{(1, 1): 1}, {(1, ): 1}, {(1, 1, 1): 1}]), [{(1, ): 1}, {(1, 1): 1}, {(1, 1, 1): 1}]) def test_dict_values_do_not_have_lt(self): self.assertEqual( sort_if_possible([{ 'a': (1, 1) }, { 'a': (1, ) }, { 'a': (1, 1, 1) }]), [{ 'a': (1, ) }, { 'a': (1, 1) }, { 'a': (1, 1, 1) }]) def test_order_enums_by_name(self): class CardinalDirection(enum.Enum): NORTH = 1 EAST = 2 SOUTH = 3 WEST = 4 self.assertEqual( sort_if_possible({ CardinalDirection.NORTH, CardinalDirection.SOUTH, CardinalDirection.EAST, CardinalDirection.WEST, }), [ CardinalDirection.EAST, CardinalDirection.NORTH, CardinalDirection.SOUTH, CardinalDirection.WEST, ]) def test_enum_with_many_values(self): MyEnum = enum.Enum('MyEnum', ' '.join(f'N{i}' for i in range(10000))) self.assertEqual( sort_if_possible({ MyEnum.N789, MyEnum.N123, MyEnum.N456, }), [ MyEnum.N123, MyEnum.N456, MyEnum.N789, ]) if __name__ == "__main__": unittest.main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.258395 apache_beam-2.74.0/apache_beam/internal/test_data/0000755000175100017510000000000015206113157021451 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/__init__.py0000644000175100017510000000155715206113032023562 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test data to validate that code identifiers are invariant to small modifications. """ ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/module_1.py0000644000175100017510000000201315206113032023514 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for testing code path generation with functions. Counterpart to after_module_with_functions and is used as a test case for various code changes. """ def my_function(): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/module_1_class_added.py0000644000175100017510000000220115206113032026021 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for testing code path generation with functions. Counterpart to before_module_with_functions and is used as a test case for adding a class. """ class MyClass: def another_function(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b def my_function(): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/module_1_function_added.py0000644000175100017510000000215115206113032026545 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for testing code path generation with functions. Counterpart to before_module_with_functions and is used as a test case for adding a function. """ def another_function(): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b def my_function(): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/module_1_global_variable_added.py0000644000175100017510000000205515206113032030030 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for testing code path generation with functions. Counterpart to before_module_with_functions and is used as a test case for adding a global variable. """ GLOBAL_VARIABLE = lambda: 3 def my_function(): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/module_1_lambda_variable_added.py0000644000175100017510000000212515206113032030006 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for testing code path generation with functions. Counterpart to before_module_with_functions and is used as a test case for adding a lambda variable. """ def my_function(): a = 1 # pylint: disable=unused-variable b = lambda: 2 new_lambda_variable = lambda: 4 # pylint: disable=unused-variable return b ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/module_1_local_variable_added.py0000644000175100017510000000210515206113032027656 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for testing code path generation with functions. Counterpart to before_module_with_functions and is used as a test case for adding a variable. """ def my_function(): a = 1 # pylint: disable=unused-variable b = lambda: 2 new_local_variable = 3 # pylint: disable=unused-variable return b ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/module_1_local_variable_removed.py0000644000175100017510000000174015206113032030262 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for testing code path generation with functions. Counterpart to before_module_with_functions and is used as a test case for removing a variable. """ def my_function(): b = lambda: 2 return b ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/module_1_nested_function_2_added.py0000644000175100017510000000214515206113032030333 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for testing code path generation with functions. Counterpart to before_module_with_functions and is used as a test case for adding a nested function. """ def my_function(): a = 1 # pylint: disable=unused-variable b = lambda: 2 def nested_function(): # pylint: disable=unused-variable c = 3 return c return b ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/module_1_nested_function_added.py0000644000175100017510000000214415206113032030111 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for testing code path generation with functions. Counterpart to before_module_with_functions and is used as a test case for adding a nested function. """ def my_function(): def nested_function(): # pylint: disable=unused-variable c = 3 return c a = 1 # pylint: disable=unused-variable b = lambda: 2 return b ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/module_2.py0000644000175100017510000000403315206113032023521 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for testing code path generation with classes. Counterpart to after_module_with_classes and is used as a test case for various code changes. """ class AddLocalVariable: def my_method(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b class RemoveLocalVariable: def my_method(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b class AddLambdaVariable: def my_method(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b class RemoveLambdaVariable: def my_method(self): a = lambda: 1 # pylint: disable=unused-variable b = lambda: 2 return b class ClassWithNestedFunction: def my_method(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b class ClassWithNestedFunction2: def my_method(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b class ClassWithTwoMethods: def my_method(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b class RemoveMethod: def another_method(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b def my_method(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/module_2_modified.py0000644000175100017510000000434315206113032025365 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for testing code path generation with classes. Counterpart to before_module_with_classes and is used as a test case for various code changes. """ class AddLocalVariable: def my_method(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 new_local_variable = 3 # pylint: disable=unused-variable return b class RemoveLocalVariable: def my_method(self): b = lambda: 2 return b class AddLambdaVariable: def my_method(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 c = lambda: 3 # pylint: disable=unused-variable return b class RemoveLambdaVariable: def my_method(self): b = lambda: 2 return b class ClassWithNestedFunction: def my_method(self): def nested_function(): # pylint: disable=unused-variable c = 3 return c a = 1 # pylint: disable=unused-variable b = lambda: 2 return b class ClassWithNestedFunction2: def my_method(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 def nested_function(): # pylint: disable=unused-variable c = 3 return c return b class ClassWithTwoMethods: def another_method(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b def my_method(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b class RemoveMethod: def my_method(self): a = 1 # pylint: disable=unused-variable b = lambda: 2 return b ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/module_3.py0000644000175100017510000000174115206113032023525 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for testing code path generation with functions. Used as a test case for various code changes. """ def my_function(): a = lambda: 1 # pylint: disable=unused-variable b = lambda: 2 return b ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/module_3_modified.py0000644000175100017510000000174715206113032025373 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for testing code path generation with functions. Counterpart to before_module_with_functions and is used as a test case for removing a lambda variable. """ def my_function(): b = lambda: 2 return b ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/test_data/module_with_default_argument.py0000644000175100017510000000170415206113032027743 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for testing code path generation with functions. Used as a test case for default arguments. """ def function_with_lambda_default_argument(fn=lambda x: 1): return fn ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/util.py0000644000175100017510000001242615206113032021025 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utility functions used throughout the package. For internal use only. No backwards compatibility guarantees. """ # pytype: skip-file import logging import threading import weakref from multiprocessing.pool import ThreadPool from typing import Any from typing import Iterable from typing import TypeVar from typing import Union T = TypeVar('T') class ArgumentPlaceholder(object): """For internal use only; no backwards-compatibility guarantees. A place holder object replacing PValues in argument lists. A Fn object can take any number of "side inputs", which are PValues that will be evaluated during pipeline execution and will be provided to the function at the moment of its execution as positional or keyword arguments. This is used only internally and should never be used by user code. A custom Fn object by the time it executes will have such values replaced with real computed values. """ def __eq__(self, other): """Tests for equality of two placeholder objects. Args: other: Another placeholder object to compare to. This method is used only for test code. All placeholder objects are equal to each other. """ return isinstance(other, ArgumentPlaceholder) def __hash__(self): return hash(type(self)) def remove_objects_from_args( args: Iterable[Any], kwargs: dict[str, Any], pvalue_class: Union[type[T], tuple[type[T], ...]] ) -> tuple[list[Any], dict[str, Any], list[T]]: """For internal use only; no backwards-compatibility guarantees. Replaces all objects of a given type in args/kwargs with a placeholder. Args: args: A list of positional arguments. kwargs: A dictionary of keyword arguments. pvalue_class: A class object representing the types of arguments that must be replaced with a placeholder value (instance of ArgumentPlaceholder). Returns: A 3-tuple containing a modified list of positional arguments, a modified dictionary of keyword arguments, and a list of all objects replaced with a placeholder value. """ pvals = [] def swapper(value): pvals.append(value) return ArgumentPlaceholder() new_args = [swapper(v) if isinstance(v, pvalue_class) else v for v in args] # Make sure the order in which we process the dictionary keys is predictable # by sorting the entries first. This will be important when putting back # PValues. new_kwargs = dict((k, swapper(v)) if isinstance(v, pvalue_class) else (k, v) for k, v in sorted(kwargs.items())) return (new_args, new_kwargs, pvals) def insert_values_in_args(args, kwargs, values): """For internal use only; no backwards-compatibility guarantees. Replaces all placeholders in args/kwargs with actual values. Args: args: A list of positional arguments. kwargs: A dictionary of keyword arguments. values: A list of values that will be used to replace placeholder values. Returns: A 2-tuple containing a modified list of positional arguments, and a modified dictionary of keyword arguments. """ # Use a local iterator so that we don't modify values. v_iter = iter(values) new_args = [ next(v_iter) if isinstance(arg, ArgumentPlaceholder) else arg for arg in args ] new_kwargs = dict( (k, next(v_iter)) if isinstance(v, ArgumentPlaceholder) else (k, v) for k, v in sorted(kwargs.items())) return (new_args, new_kwargs) def run_using_threadpool(fn_to_execute, inputs, pool_size): """For internal use only; no backwards-compatibility guarantees. Runs the given function on given inputs using a thread pool. Args: fn_to_execute: Function to execute inputs: Inputs on which given function will be executed in parallel. pool_size: Size of thread pool. Returns: Results retrieved after executing the given function on given inputs. """ # ThreadPool crashes in old versions of Python (< 2.7.5) if created # from a child thread. (http://bugs.python.org/issue10015) if not hasattr(threading.current_thread(), '_children'): threading.current_thread()._children = weakref.WeakKeyDictionary() pool = ThreadPool(min(pool_size, len(inputs))) try: # We record and reset logging level here since 'apitools' library Beam # depends on updates the logging level when used with a threadpool - # https://github.com/google/apitools/issues/141 # TODO: Remove this once above issue in 'apitools' is fixed. old_level = logging.getLogger().level return pool.map(fn_to_execute, inputs) finally: pool.terminate() logging.getLogger().setLevel(old_level) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/internal/util_test.py0000644000175100017510000000435715206113032022070 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the util module.""" # pytype: skip-file import unittest from apache_beam.internal.util import ArgumentPlaceholder from apache_beam.internal.util import insert_values_in_args from apache_beam.internal.util import remove_objects_from_args class UtilTest(unittest.TestCase): def test_remove_objects_from_args(self): args, kwargs, objs = remove_objects_from_args( [1, 'a'], {'x': 1, 'y': 3.14}, (str, float)) self.assertEqual([1, ArgumentPlaceholder()], args) self.assertEqual({'x': 1, 'y': ArgumentPlaceholder()}, kwargs) self.assertEqual(['a', 3.14], objs) def test_remove_objects_from_args_nothing_to_remove(self): args, kwargs, objs = remove_objects_from_args( [1, 2], {'x': 1, 'y': 2}, (str, float)) self.assertEqual([1, 2], args) self.assertEqual({'x': 1, 'y': 2}, kwargs) self.assertEqual([], objs) def test_insert_values_in_args(self): values = ['a', 'b'] args = [1, ArgumentPlaceholder()] kwargs = {'x': 1, 'y': ArgumentPlaceholder()} args, kwargs = insert_values_in_args(args, kwargs, values) self.assertEqual([1, 'a'], args) self.assertEqual({'x': 1, 'y': 'b'}, kwargs) def test_insert_values_in_args_nothing_to_insert(self): values = [] args = [1, 'a'] kwargs = {'x': 1, 'y': 'b'} args, kwargs = insert_values_in_args(args, kwargs, values) self.assertEqual([1, 'a'], args) self.assertEqual({'x': 1, 'y': 'b'}, kwargs) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3023996 apache_beam-2.74.0/apache_beam/io/0000755000175100017510000000000015206113157016274 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/__init__.py0000644000175100017510000000330115206113032020372 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A package defining several input sources and output sinks.""" # pylint: disable=wildcard-import # isort: off from apache_beam.io.avroio import * from apache_beam.io.filebasedsink import * from apache_beam.io.iobase import Read from apache_beam.io.iobase import Sink from apache_beam.io.iobase import Write from apache_beam.io.iobase import Writer from apache_beam.io.mongodbio import * from apache_beam.io.parquetio import * from apache_beam.io.textio import * from apache_beam.io.tfrecordio import * from apache_beam.io.range_trackers import * # Protect against environments where clientslibrary is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.io.gcp.bigquery import * from apache_beam.io.gcp.pubsub import * from apache_beam.io.gcp import gcsio from apache_beam.transforms.xlang.io import * except ImportError: pass # pylint: enable=wrong-import-order, wrong-import-position ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/avroio.py0000644000175100017510000007046615206113032020152 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """``PTransforms`` for reading from and writing to Avro files. Provides two read ``PTransform``s, ``ReadFromAvro`` and ``ReadAllFromAvro``, that produces a ``PCollection`` of records. Each record of this ``PCollection`` will contain a single record read from an Avro file. Records that are of simple types will be mapped into corresponding Python types. Records that are of Avro type 'RECORD' will be mapped to Python dictionaries that comply with the schema contained in the Avro file that contains those records. In this case, keys of each dictionary will contain the corresponding field names and will be of type ``string`` while the values of the dictionary will be of the type defined in the corresponding Avro schema. For example, if schema of the Avro file is the following. {"namespace": "example.avro","type": "record","name": "User","fields": [{"name": "name", "type": "string"}, {"name": "favorite_number", "type": ["int", "null"]}, {"name": "favorite_color", "type": ["string", "null"]}]} Then records generated by read transforms will be dictionaries of the following form. {'name': 'Alyssa', 'favorite_number': 256, 'favorite_color': None}). Additionally, this module provides a write ``PTransform`` ``WriteToAvro`` that can be used to write a given ``PCollection`` of Python objects to an Avro file. """ # pytype: skip-file import os from functools import partial from typing import Any from typing import Callable from typing import Union import fastavro from fastavro.read import block_reader from fastavro.write import Writer import apache_beam as beam from apache_beam.io import filebasedsink from apache_beam.io import filebasedsource from apache_beam.io import iobase from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystems import FileSystems from apache_beam.io.iobase import Read from apache_beam.portability.api import schema_pb2 from apache_beam.transforms import PTransform from apache_beam.typehints import schemas __all__ = [ 'ReadFromAvro', 'ReadAllFromAvro', 'ReadAllFromAvroContinuously', 'WriteToAvro' ] class ReadFromAvro(PTransform): """A `PTransform` for reading records from avro files. Each record of the resulting PCollection will contain a single record read from a source. Records that are of simple types will be mapped to beam Rows with a single `record` field containing the records value. Records that are of Avro type ``RECORD`` will be mapped to Beam rows that comply with the schema contained in the Avro file that contains those records. """ def __init__( self, file_pattern=None, min_bundle_size=0, validate=True, use_fastavro=True, as_rows=False): """Initializes :class:`ReadFromAvro`. Uses source :class:`~apache_beam.io._AvroSource` to read a set of Avro files defined by a given file pattern. If ``/mypath/myavrofiles*`` is a file-pattern that points to a set of Avro files, a :class:`~apache_beam.pvalue.PCollection` for the records in these Avro files can be created in the following manner. .. testcode:: with beam.Pipeline() as p: records = p | 'Read' >> beam.io.ReadFromAvro('/mypath/myavrofiles*') .. NOTE: We're not actually interested in this error; but if we get here, it means that the way of calling this transform hasn't changed. .. testoutput:: :hide: Traceback (most recent call last): ... OSError: No files found based on the file pattern Each record of this :class:`~apache_beam.pvalue.PCollection` will contain a single record read from a source. Records that are of simple types will be mapped into corresponding Python types. Records that are of Avro type ``RECORD`` will be mapped to Python dictionaries that comply with the schema contained in the Avro file that contains those records. In this case, keys of each dictionary will contain the corresponding field names and will be of type :class:`str` while the values of the dictionary will be of the type defined in the corresponding Avro schema. For example, if schema of the Avro file is the following. :: { "namespace": "example.avro", "type": "record", "name": "User", "fields": [ {"name": "name", "type": "string"}, {"name": "favorite_number", "type": ["int", "null"]}, {"name": "favorite_color", "type": ["string", "null"]} ] } Then records generated by :class:`~apache_beam.io._AvroSource` will be dictionaries of the following form. :: {'name': 'Alyssa', 'favorite_number': 256, 'favorite_color': None}). Args: file_pattern (str): the file glob to read min_bundle_size (int): the minimum size in bytes, to be considered when splitting the input into bundles. validate (bool): flag to verify that the files exist during the pipeline creation time. use_fastavro (bool): This flag is left for API backwards compatibility and no longer has an effect. Do not use. as_rows (bool): Whether to return a schema'd PCollection of Beam rows. """ super().__init__() self._source = _FastAvroSource( file_pattern, min_bundle_size, validate=validate) if as_rows: path = FileSystems.match([file_pattern], [1])[0].metadata_list[0].path with FileSystems.open(path) as fin: avro_schema = fastavro.reader(fin).writer_schema beam_schema = avro_schema_to_beam_schema(avro_schema) self._post_process = avro_dict_to_beam_row(avro_schema, beam_schema) else: self._post_process = None def expand(self, pvalue): records = pvalue.pipeline | Read(self._source) if self._post_process: return records | beam.Map(self._post_process) else: return records def display_data(self): return {'source_dd': self._source} class ReadAllFromAvro(PTransform): """A ``PTransform`` for reading ``PCollection`` of Avro files. Uses source '_AvroSource' to read a ``PCollection`` of Avro files or file patterns and produce a ``PCollection`` of Avro records. This implementation is only tested with batch pipeline. In streaming, reading may happen with delay due to the limitation in ReShuffle involved. """ DEFAULT_DESIRED_BUNDLE_SIZE = 64 * 1024 * 1024 # 64MB def __init__( self, min_bundle_size=0, desired_bundle_size=DEFAULT_DESIRED_BUNDLE_SIZE, use_fastavro=True, with_filename=False, label='ReadAllFiles'): """Initializes ``ReadAllFromAvro``. Args: min_bundle_size: the minimum size in bytes, to be considered when splitting the input into bundles. desired_bundle_size: the desired size in bytes, to be considered when splitting the input into bundles. use_fastavro (bool): This flag is left for API backwards compatibility and no longer has an effect. Do not use. with_filename: If True, returns a Key Value with the key being the file name and the value being the actual data. If False, it only returns the data. """ source_from_file = partial(_FastAvroSource, min_bundle_size=min_bundle_size) self._read_all_files = filebasedsource.ReadAllFiles( True, CompressionTypes.AUTO, desired_bundle_size, min_bundle_size, source_from_file, with_filename) self.label = label def expand(self, pvalue): return pvalue | self.label >> self._read_all_files class ReadAllFromAvroContinuously(ReadAllFromAvro): """A ``PTransform`` for reading avro files in given file patterns. This PTransform acts as a Source and produces continuously a ``PCollection`` of Avro records. For more details, see ``ReadAllFromAvro`` for avro parsing settings; see ``apache_beam.io.fileio.MatchContinuously`` for watching settings. ReadAllFromAvroContinuously is experimental. No backwards-compatibility guarantees. Due to the limitation on Reshuffle, current implementation does not scale. """ _ARGS_FOR_MATCH = ( 'interval', 'has_deduplication', 'start_timestamp', 'stop_timestamp', 'match_updated_files', 'apply_windowing') _ARGS_FOR_READ = ( 'min_bundle_size', 'desired_bundle_size', 'use_fastavro', 'with_filename') def __init__(self, file_pattern, label='ReadAllFilesContinuously', **kwargs): """Initialize the ``ReadAllFromAvroContinuously`` transform. Accepts args for constructor args of both :class:`ReadAllFromAvro` and :class:`~apache_beam.io.fileio.MatchContinuously`. """ kwargs_for_match = { k: v for (k, v) in kwargs.items() if k in self._ARGS_FOR_MATCH } kwargs_for_read = { k: v for (k, v) in kwargs.items() if k in self._ARGS_FOR_READ } kwargs_additinal = { k: v for (k, v) in kwargs.items() if k not in self._ARGS_FOR_MATCH and k not in self._ARGS_FOR_READ } super().__init__(label=label, **kwargs_for_read, **kwargs_additinal) self._file_pattern = file_pattern self._kwargs_for_match = kwargs_for_match def expand(self, pbegin): # Importing locally to prevent circular dependency issues. from apache_beam.io.fileio import MatchContinuously # TODO(BEAM-14497) always reshuffle once gbk always trigger works. return ( pbegin | MatchContinuously(self._file_pattern, **self._kwargs_for_match) | 'ReadAllFiles' >> self._read_all_files._disable_reshuffle()) class _AvroUtils(object): @staticmethod def advance_file_past_next_sync_marker(f, sync_marker): buf_size = 10000 data = f.read(buf_size) while data: pos = data.find(sync_marker) if pos >= 0: # Adjusting the current position to the ending position of the sync # marker. backtrack = len(data) - pos - len(sync_marker) f.seek(-1 * backtrack, os.SEEK_CUR) return True else: if f.tell() >= len(sync_marker): # Backtracking in case we partially read the sync marker during the # previous read. We only have to backtrack if there are at least # len(sync_marker) bytes before current position. We only have to # backtrack (len(sync_marker) - 1) bytes. f.seek(-1 * (len(sync_marker) - 1), os.SEEK_CUR) data = f.read(buf_size) class _FastAvroSource(filebasedsource.FileBasedSource): """A source for reading Avro files using the `fastavro` library. ``_FastAvroSource`` is implemented using the file-based source framework available in module 'filebasedsource'. Hence please refer to module 'filebasedsource' to fully understand how this source implements operations common to all file-based sources such as file-pattern expansion and splitting into bundles for parallel processing. TODO: remove ``_AvroSource`` in favor of using ``_FastAvroSource`` everywhere once it has been more widely tested """ def read_records(self, file_name, range_tracker): next_block_start = -1 def split_points_unclaimed(stop_position): if next_block_start >= stop_position: # Next block starts at or after the suggested stop position. Hence # there will not be split points to be claimed for the range ending at # suggested stop position. return 0 return iobase.RangeTracker.SPLIT_POINTS_UNKNOWN range_tracker.set_split_points_unclaimed_callback(split_points_unclaimed) start_offset = range_tracker.start_position() if start_offset is None: start_offset = 0 with self.open_file(file_name) as f: blocks = block_reader(f) sync_marker = blocks._header['sync'] # We have to start at current position if previous bundle ended at the # end of a sync marker. start_offset = max(0, start_offset - len(sync_marker)) f.seek(start_offset) _AvroUtils.advance_file_past_next_sync_marker(f, sync_marker) next_block_start = f.tell() while range_tracker.try_claim(next_block_start): block = next(blocks) next_block_start = block.offset + block.size yield from block _create_avro_source = _FastAvroSource class WriteToAvro(beam.transforms.PTransform): """A ``PTransform`` for writing avro files. If the input has a schema, a corresponding avro schema will be automatically generated and used to write the output records.""" def __init__( self, file_path_prefix, schema=None, codec='deflate', file_name_suffix='', num_shards=0, shard_name_template=None, mime_type='application/x-avro', use_fastavro=True, triggering_frequency=None): """Initialize a WriteToAvro transform. Args: file_path_prefix: The file path to write to. The files written will begin with this prefix, followed by a shard identifier (see num_shards), and end in a common extension, if given by file_name_suffix. In most cases, only this argument is specified and num_shards, shard_name_template, and file_name_suffix use default values. schema: The schema to use (dict). codec: The codec to use for block-level compression. Any string supported by the Avro specification is accepted (for example 'null'). file_name_suffix: Suffix for the files written. num_shards: The number of files (shards) used for output. If not set, the service will decide on the optimal number of shards. Constraining the number of shards is likely to reduce the performance of a pipeline. Setting this value is not recommended unless you require a specific number of output files. In streaming if not set, the service will write a file per bundle. shard_name_template: A template string containing placeholders for the shard number and shard count. Currently only ``''``, ``'-SSSSS-of-NNNNN'``, ``'-W-SSSSS-of-NNNNN'`` and ``'-V-SSSSS-of-NNNNN'`` are patterns accepted by the service. When constructing a filename for a particular shard number, the upper-case letters ``S`` and ``N`` are replaced with the ``0``-padded shard number and shard count respectively. This argument can be ``''`` in which case it behaves as if num_shards was set to 1 and only one file will be generated. The default pattern used is ``'-SSSSS-of-NNNNN'`` for bounded PCollections and for ``'-W-SSSSS-of-NNNNN'`` unbounded PCollections. W is used for windowed shard naming and is replaced with ``[window.start, window.end)`` V is used for windowed shard naming and is replaced with ``[window.start.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S"), window.end.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S")`` mime_type: The MIME type to use for the produced files, if the filesystem supports specifying MIME types. use_fastavro (bool): This flag is left for API backwards compatibility and no longer has an effect. Do not use. triggering_frequency: (int) Every triggering_frequency duration, a window will be triggered and all bundles in the window will be written. If set it overrides user windowing. Mandatory for GlobalWindow. Returns: A WriteToAvro transform usable for writing. """ self._schema = schema self._sink_provider = lambda avro_schema: _create_avro_sink( file_path_prefix, avro_schema, codec, file_name_suffix, num_shards, shard_name_template, mime_type, triggering_frequency) def expand(self, pcoll): if self._schema: avro_schema = self._schema records = pcoll else: try: beam_schema = schemas.schema_from_element_type(pcoll.element_type) except TypeError as exn: raise ValueError( "An explicit schema is required to write non-schema'd PCollections." ) from exn avro_schema = beam_schema_to_avro_schema(beam_schema) records = pcoll | beam.Map( beam_row_to_avro_dict(avro_schema, beam_schema)) self._sink = self._sink_provider(avro_schema) if (not pcoll.is_bounded and self._sink.shard_name_template == filebasedsink.DEFAULT_SHARD_NAME_TEMPLATE): self._sink.shard_name_template = ( filebasedsink.DEFAULT_WINDOW_SHARD_NAME_TEMPLATE) self._sink.shard_name_format = self._sink._template_to_format( self._sink.shard_name_template) self._sink.shard_name_glob_format = self._sink._template_to_glob_format( self._sink.shard_name_template) return records | beam.io.iobase.Write(self._sink) def display_data(self): return {'sink_dd': self._sink} def _create_avro_sink( file_path_prefix, schema, codec, file_name_suffix, num_shards, shard_name_template, mime_type, triggering_frequency=60): if "class 'avro.schema" in str(type(schema)): raise ValueError( 'You are using Avro IO with fastavro (default with Beam on ' 'Python 3), but supplying a schema parsed by avro-python3. ' 'Please change the schema to a dict.') return _FastAvroSink( file_path_prefix, schema, codec, file_name_suffix, num_shards, shard_name_template, mime_type, triggering_frequency) class _BaseAvroSink(filebasedsink.FileBasedSink): """A base for a sink for avro files. """ def __init__( self, file_path_prefix, schema, codec, file_name_suffix, num_shards, shard_name_template, mime_type, triggering_frequency): super().__init__( file_path_prefix, file_name_suffix=file_name_suffix, num_shards=num_shards, shard_name_template=shard_name_template, coder=None, mime_type=mime_type, # Compression happens at the block level using the supplied codec, and # not at the file level. compression_type=CompressionTypes.UNCOMPRESSED, triggering_frequency=triggering_frequency) self._schema = schema self._codec = codec def display_data(self): res = super().display_data() res['codec'] = str(self._codec) res['schema'] = str(self._schema) return res class _FastAvroSink(_BaseAvroSink): """A sink for avro files using FastAvro. """ def __init__( self, file_path_prefix, schema, codec, file_name_suffix, num_shards, shard_name_template, mime_type, triggering_frequency): super().__init__( file_path_prefix, schema, codec, file_name_suffix, num_shards, shard_name_template, mime_type, triggering_frequency) self.file_handle = None def open(self, temp_path): self.file_handle = super().open(temp_path) return Writer(self.file_handle, self._schema, self._codec) def write_record(self, writer, value): writer.write(value) def close(self, writer): writer.flush() self.file_handle.close() AVRO_PRIMITIVES_TO_BEAM_PRIMITIVES = { 'boolean': schema_pb2.BOOLEAN, 'int': schema_pb2.INT32, 'long': schema_pb2.INT64, 'float': schema_pb2.FLOAT, 'double': schema_pb2.DOUBLE, 'bytes': schema_pb2.BYTES, 'string': schema_pb2.STRING, } BEAM_PRIMITIVES_TO_AVRO_PRIMITIVES = { v: k for k, v in AVRO_PRIMITIVES_TO_BEAM_PRIMITIVES.items() } _AvroSchemaType = Union[str, list, dict] def avro_union_type_to_beam_type(union_type: list) -> schema_pb2.FieldType: """convert an avro union type to a beam type if the union type is a nullable, and it is a nullable union of an avro type with a corresponding beam type then create a nullable beam field of the corresponding beam type, otherwise return an Any type. Args: union_type: the avro union type to convert. Returns: the beam type of the avro union. """ if len(union_type) == 2 and "null" in union_type: for avro_type in union_type: if avro_type != "null": beam_type = avro_type_to_beam_type(avro_type) beam_type.nullable = True return beam_type return schemas.typing_to_runner_api(Any) def avro_type_to_beam_type(avro_type: _AvroSchemaType) -> schema_pb2.FieldType: if isinstance(avro_type, str): return avro_type_to_beam_type({'type': avro_type}) elif isinstance(avro_type, list): # Union type return avro_union_type_to_beam_type(avro_type) type_name = avro_type['type'] if type_name in AVRO_PRIMITIVES_TO_BEAM_PRIMITIVES: return schema_pb2.FieldType( atomic_type=AVRO_PRIMITIVES_TO_BEAM_PRIMITIVES[type_name]) elif type_name in ('fixed', 'enum'): return schema_pb2.FieldType(atomic_type=schema_pb2.STRING) elif type_name == 'array': return schema_pb2.FieldType( array_type=schema_pb2.ArrayType( element_type=avro_type_to_beam_type(avro_type['items']))) elif type_name == 'map': return schema_pb2.FieldType( map_type=schema_pb2.MapType( key_type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING), value_type=avro_type_to_beam_type(avro_type['values']))) elif type_name == 'record': return schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( fields=[ schemas.schema_field( f['name'], avro_type_to_beam_type(f['type'])) for f in avro_type['fields'] ]))) else: raise ValueError(f'Unable to convert {avro_type} to a Beam schema.') def avro_schema_to_beam_schema( avro_schema: _AvroSchemaType) -> schema_pb2.Schema: beam_type = avro_type_to_beam_type(avro_schema) if isinstance(avro_schema, dict) and avro_schema['type'] == 'record': return beam_type.row_type.schema else: return schema_pb2.Schema(fields=[schemas.schema_field('record', beam_type)]) def avro_dict_to_beam_row( avro_schema: _AvroSchemaType, beam_schema: schema_pb2.Schema) -> Callable[[Any], Any]: if isinstance(avro_schema, str): return avro_dict_to_beam_row({'type': avro_schema}) if avro_schema['type'] == 'record': to_row = avro_value_to_beam_value( schema_pb2.FieldType(row_type=schema_pb2.RowType(schema=beam_schema))) else: def to_row(record): return beam.Row(record=record) return beam.typehints.with_output_types( schemas.named_tuple_from_schema(beam_schema))( to_row) def avro_value_to_beam_value( beam_type: schema_pb2.FieldType) -> Callable[[Any], Any]: type_info = beam_type.WhichOneof("type_info") if type_info == "atomic_type": return lambda value: value elif type_info == "array_type": element_converter = avro_value_to_beam_value( beam_type.array_type.element_type) return lambda value: [element_converter(e) for e in value] elif type_info == "iterable_type": element_converter = avro_value_to_beam_value( beam_type.iterable_type.element_type) return lambda value: [element_converter(e) for e in value] elif type_info == "map_type": if beam_type.map_type.key_type.atomic_type != schema_pb2.STRING: raise TypeError( f'Only strings allowed as map keys when converting from AVRO, ' f'found {beam_type}') value_converter = avro_value_to_beam_value(beam_type.map_type.value_type) return lambda value: {k: value_converter(v) for (k, v) in value.items()} elif type_info == "row_type": converters = { field.name: avro_value_to_beam_value(field.type) for field in beam_type.row_type.schema.fields } return lambda value: beam.Row( ** {name: convert(value[name]) for (name, convert) in converters.items()}) elif type_info == "logical_type": return lambda value: value else: raise ValueError(f"Unrecognized type_info: {type_info!r}") def beam_schema_to_avro_schema( beam_schema: schema_pb2.Schema) -> _AvroSchemaType: return beam_type_to_avro_type( schema_pb2.FieldType(row_type=schema_pb2.RowType(schema=beam_schema))) def unnest_primitive_type(beam_type: schema_pb2.FieldType): """unnests beam types that map to avro primitives or unions. if mapping to a avro primitive or a union, don't nest the field type for complex types, like arrays, we need to nest the type. Example: { 'type': 'string' } -> 'string' { 'type': 'array', 'items': 'string' } -> { 'type': 'array', 'items': 'string' } Args: beam_type: the beam type to map to avro. Returns: the converted avro type with the primitive or union type unnested. """ avro_type = beam_type_to_avro_type(beam_type) return avro_type['type'] if beam_type.WhichOneof( "type_info") == "atomic_type" else avro_type def beam_type_to_avro_type(beam_type: schema_pb2.FieldType) -> _AvroSchemaType: type_info = beam_type.WhichOneof("type_info") if type_info == "atomic_type": avro_primitive = BEAM_PRIMITIVES_TO_AVRO_PRIMITIVES[beam_type.atomic_type] avro_type = [ avro_primitive, 'null' ] if beam_type.nullable else avro_primitive return {'type': avro_type} elif type_info == "array_type": return { 'type': 'array', 'items': unnest_primitive_type(beam_type.array_type.element_type) } elif type_info == "iterable_type": return { 'type': 'array', 'items': unnest_primitive_type(beam_type.iterable_type.element_type) } elif type_info == "map_type": if beam_type.map_type.key_type.atomic_type != schema_pb2.STRING: raise TypeError( f'Only strings allowed as map keys when converting to AVRO, ' f'found {beam_type}') return { 'type': 'map', 'values': unnest_primitive_type(beam_type.map_type.element_type) } elif type_info == "row_type": return { 'type': 'record', 'name': beam_type.row_type.schema.id, 'fields': [{ 'name': field.name, 'type': unnest_primitive_type(field.type) } for field in beam_type.row_type.schema.fields], } else: raise ValueError(f"Unconvertable type: {beam_type}") def beam_row_to_avro_dict( avro_schema: _AvroSchemaType, beam_schema: schema_pb2.Schema): if isinstance(avro_schema, str): return beam_row_to_avro_dict({'type': avro_schema}, beam_schema) if avro_schema['type'] == 'record': return beam_value_to_avro_value( schema_pb2.FieldType(row_type=schema_pb2.RowType(schema=beam_schema))) else: convert = beam_value_to_avro_value(beam_schema) return lambda row: convert(row[0]) def beam_value_to_avro_value( beam_type: schema_pb2.FieldType) -> Callable[[Any], Any]: type_info = beam_type.WhichOneof("type_info") if type_info == "atomic_type": return lambda value: value elif type_info == "array_type": element_converter = beam_value_to_avro_value( beam_type.array_type.element_type) return lambda value: [element_converter(e) for e in value] elif type_info == "iterable_type": element_converter = beam_value_to_avro_value( beam_type.iterable_type.element_type) return lambda value: [element_converter(e) for e in value] elif type_info == "map_type": if beam_type.map_type.key_type.atomic_type != schema_pb2.STRING: raise TypeError( f'Only strings allowed as map keys when converting from AVRO, ' f'found {beam_type}') value_converter = beam_value_to_avro_value(beam_type.map_type.value_type) return lambda value: {k: value_converter(v) for (k, v) in value.items()} elif type_info == "row_type": converters = { field.name: beam_value_to_avro_value(field.type) for field in beam_type.row_type.schema.fields } return lambda value: { name: convert(getattr(value, name)) for (name, convert) in converters.items() } elif type_info == "logical_type": return lambda value: value else: raise ValueError(f"Unrecognized type_info: {type_info!r}") ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/avroio_test.py0000644000175100017510000011077415206113032021206 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import glob import json import logging import math import os import pytz import pytest import re import shutil import tempfile import unittest from typing import Any import fastavro import hamcrest as hc from fastavro.schema import parse_schema from fastavro import writer import apache_beam as beam from apache_beam import Create, schema_pb2 from apache_beam.io import avroio from apache_beam.io import filebasedsource from apache_beam.io import iobase from apache_beam.io import source_test_utils from apache_beam.io.avroio import _FastAvroSource # For testing from apache_beam.io.avroio import avro_schema_to_beam_schema # For testing from apache_beam.io.avroio import beam_schema_to_avro_schema # For testing from apache_beam.io.avroio import avro_union_type_to_beam_type # For testing from apache_beam.io.avroio import avro_dict_to_beam_row # For testing from apache_beam.io.avroio import beam_row_to_avro_dict # For testing from apache_beam.io.avroio import _create_avro_sink # For testing from apache_beam.io.filesystems import FileSystems from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.display import DisplayData from apache_beam.transforms.display_test import DisplayDataItemMatcher from apache_beam.transforms.sql import SqlTransform from apache_beam.transforms.userstate import CombiningValueStateSpec from apache_beam.transforms.util import LogElements from apache_beam.utils.timestamp import Timestamp from apache_beam.typehints import schemas from datetime import datetime # Import snappy optionally; some tests will be skipped when import fails. try: import snappy # pylint: disable=import-error except ImportError: snappy = None # pylint: disable=invalid-name logging.warning('python-snappy is not installed; some tests will be skipped.') RECORDS = [{ 'name': 'Thomas', 'favorite_number': 1, 'favorite_color': 'blue' }, { 'name': 'Henry', 'favorite_number': 3, 'favorite_color': 'green' }, { 'name': 'Toby', 'favorite_number': 7, 'favorite_color': 'brown' }, { 'name': 'Gordon', 'favorite_number': 4, 'favorite_color': 'blue' }, { 'name': 'Emily', 'favorite_number': -1, 'favorite_color': 'Red' }, { 'name': 'Percy', 'favorite_number': 6, 'favorite_color': 'Green' }] class AvroBase(object): _temp_files: list[str] = [] def __init__(self, methodName='runTest'): super().__init__(methodName) self.RECORDS = RECORDS self.SCHEMA_STRING = ''' {"namespace": "example.avro", "type": "record", "name": "User", "fields": [ {"name": "name", "type": "string"}, {"name": "favorite_number", "type": ["int", "null"]}, {"name": "favorite_color", "type": ["string", "null"]} ] } ''' def setUp(self): # Reducing the size of thread pools. Without this test execution may fail in # environments with limited amount of resources. filebasedsource.MAX_NUM_THREADS_FOR_SIZE_ESTIMATION = 2 def tearDown(self): for path in self._temp_files: if os.path.exists(path): os.remove(path) self._temp_files = [] def _write_data( self, directory=None, prefix=None, codec=None, count=None, sync_interval=None): raise NotImplementedError def _write_pattern(self, num_files, return_filenames=False): assert num_files > 0 temp_dir = tempfile.mkdtemp() file_name = None file_list = [] for _ in range(num_files): file_name = self._write_data(directory=temp_dir, prefix='mytemp') file_list.append(file_name) assert file_name file_name_prefix = file_name[:file_name.rfind(os.path.sep)] if return_filenames: return (file_name_prefix + os.path.sep + 'mytemp*', file_list) return file_name_prefix + os.path.sep + 'mytemp*' def _run_avro_test( self, pattern, desired_bundle_size, perform_splitting, expected_result): source = _FastAvroSource(pattern) if perform_splitting: assert desired_bundle_size splits = [ split for split in source.split(desired_bundle_size=desired_bundle_size) ] if len(splits) < 2: raise ValueError( 'Test is trivial. Please adjust it so that at least ' 'two splits get generated') sources_info = [(split.source, split.start_position, split.stop_position) for split in splits] source_test_utils.assert_sources_equal_reference_source( (source, None, None), sources_info) else: read_records = source_test_utils.read_from_source(source, None, None) self.assertCountEqual(expected_result, read_records) def test_schema_read_write(self): with tempfile.TemporaryDirectory() as tmp_dirname: path = os.path.join(tmp_dirname, 'tmp_filename') rows = [beam.Row(a=1, b=['x', 'y']), beam.Row(a=2, b=['t', 'u'])] stable_repr = lambda row: json.dumps(row._asdict()) with TestPipeline() as p: _ = p | Create(rows) | avroio.WriteToAvro(path) | beam.Map(print) with TestPipeline() as p: readback = ( p | avroio.ReadFromAvro(path + '*', as_rows=True) | beam.Map(stable_repr)) assert_that(readback, equal_to([stable_repr(r) for r in rows])) @pytest.mark.xlang_sql_expansion_service @unittest.skipIf( TestPipeline().get_pipeline_options().view_as(StandardOptions).runner is None, "Must be run with a runner that supports staging java artifacts.") def test_avro_schema_to_beam_schema_with_nullable_atomic_fields(self): records = [] records.extend(self.RECORDS) records.append({ 'name': 'Bruce', 'favorite_number': None, 'favorite_color': None }) avro_schema = fastavro.parse_schema(json.loads(self.SCHEMA_STRING)) beam_schema = avro_schema_to_beam_schema(avro_schema) with TestPipeline() as p: readback = ( p | Create(records) | beam.Map(avro_dict_to_beam_row(avro_schema, beam_schema)) | SqlTransform("SELECT * FROM PCOLLECTION") | beam.Map(beam_row_to_avro_dict(avro_schema, beam_schema))) assert_that(readback, equal_to(records)) def test_avro_union_type_to_beam_type_with_nullable_long(self): union_type = ['null', 'long'] beam_type = avro_union_type_to_beam_type(union_type) expected_beam_type = schema_pb2.FieldType( atomic_type=schema_pb2.INT64, nullable=True) hc.assert_that(beam_type, hc.equal_to(expected_beam_type)) def test_avro_union_type_to_beam_type_with_string_long(self): union_type = ['string', 'long'] beam_type = avro_union_type_to_beam_type(union_type) expected_beam_type = schemas.typing_to_runner_api(Any) hc.assert_that(beam_type, hc.equal_to(expected_beam_type)) def test_avro_union_type_to_beam_type_with_record_and_null(self): record_type = { 'type': 'record', 'name': 'TestRecord', 'fields': [{ 'name': 'field1', 'type': 'string' }, { 'name': 'field2', 'type': 'int' }] } union_type = [record_type, 'null'] beam_type = avro_union_type_to_beam_type(union_type) expected_beam_type = schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( fields=[ schemas.schema_field( 'field1', schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), schemas.schema_field( 'field2', schema_pb2.FieldType(atomic_type=schema_pb2.INT32)) ])), nullable=True) hc.assert_that(beam_type, hc.equal_to(expected_beam_type)) def test_avro_union_type_to_beam_type_with_nullable_annotated_string(self): annotated_string_type = {"avro.java.string": "String", "type": "string"} union_type = ['null', annotated_string_type] beam_type = avro_union_type_to_beam_type(union_type) expected_beam_type = schema_pb2.FieldType( atomic_type=schema_pb2.STRING, nullable=True) hc.assert_that(beam_type, hc.equal_to(expected_beam_type)) def test_avro_union_type_to_beam_type_with_only_null(self): union_type = ['null'] beam_type = avro_union_type_to_beam_type(union_type) expected_beam_type = schemas.typing_to_runner_api(Any) hc.assert_that(beam_type, hc.equal_to(expected_beam_type)) def test_avro_union_type_to_beam_type_with_multiple_types(self): union_type = ['null', 'string', 'int'] beam_type = avro_union_type_to_beam_type(union_type) expected_beam_type = schemas.typing_to_runner_api(Any) hc.assert_that(beam_type, hc.equal_to(expected_beam_type)) def test_avro_schema_to_beam_and_back(self): avro_schema = fastavro.parse_schema(json.loads(self.SCHEMA_STRING)) beam_schema = avro_schema_to_beam_schema(avro_schema) converted_avro_schema = beam_schema_to_avro_schema(beam_schema) expected_fields = json.loads(self.SCHEMA_STRING)["fields"] hc.assert_that( converted_avro_schema["fields"], hc.equal_to(expected_fields)) def test_read_without_splitting(self): file_name = self._write_data() expected_result = self.RECORDS self._run_avro_test(file_name, None, False, expected_result) def test_read_with_splitting(self): file_name = self._write_data() expected_result = self.RECORDS self._run_avro_test(file_name, 100, True, expected_result) def test_source_display_data(self): file_name = 'some_avro_source' source = \ _FastAvroSource( file_name, validate=False, ) dd = DisplayData.create_from(source) # No extra avro parameters for AvroSource. expected_items = [ DisplayDataItemMatcher('compression', 'auto'), DisplayDataItemMatcher('file_pattern', file_name) ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_read_display_data(self): file_name = 'some_avro_source' read = \ avroio.ReadFromAvro( file_name, validate=False) dd = DisplayData.create_from(read) # No extra avro parameters for AvroSource. expected_items = [ DisplayDataItemMatcher('compression', 'auto'), DisplayDataItemMatcher('file_pattern', file_name) ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_sink_display_data(self): file_name = 'some_avro_sink' sink = _create_avro_sink( file_name, self.SCHEMA, 'null', '.end', 0, None, 'application/x-avro') dd = DisplayData.create_from(sink) expected_items = [ DisplayDataItemMatcher('schema', str(self.SCHEMA)), DisplayDataItemMatcher( 'file_pattern', 'some_avro_sink-%(shard_num)05d-of-%(num_shards)05d.end'), DisplayDataItemMatcher('codec', 'null'), DisplayDataItemMatcher('compression', 'uncompressed') ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_write_display_data(self): file_name = 'some_avro_sink' write = avroio.WriteToAvro(file_name, self.SCHEMA) write.expand(beam.PCollection(beam.Pipeline())) dd = DisplayData.create_from(write) expected_items = [ DisplayDataItemMatcher('schema', str(self.SCHEMA)), DisplayDataItemMatcher( 'file_pattern', 'some_avro_sink-%(shard_num)05d-of-%(num_shards)05d'), DisplayDataItemMatcher('codec', 'deflate'), DisplayDataItemMatcher('compression', 'uncompressed') ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_read_reentrant_without_splitting(self): file_name = self._write_data() source = _FastAvroSource(file_name) source_test_utils.assert_reentrant_reads_succeed((source, None, None)) def test_read_reantrant_with_splitting(self): file_name = self._write_data() source = _FastAvroSource(file_name) splits = [split for split in source.split(desired_bundle_size=100000)] assert len(splits) == 1 source_test_utils.assert_reentrant_reads_succeed( (splits[0].source, splits[0].start_position, splits[0].stop_position)) def test_read_without_splitting_multiple_blocks(self): file_name = self._write_data(count=12000) expected_result = self.RECORDS * 2000 self._run_avro_test(file_name, None, False, expected_result) def test_read_with_splitting_multiple_blocks(self): file_name = self._write_data(count=12000) expected_result = self.RECORDS * 2000 self._run_avro_test(file_name, 10000, True, expected_result) def test_split_points(self): num_records = 12000 sync_interval = 16000 file_name = self._write_data(count=num_records, sync_interval=sync_interval) source = _FastAvroSource(file_name) splits = [split for split in source.split(desired_bundle_size=float('inf'))] assert len(splits) == 1 range_tracker = splits[0].source.get_range_tracker( splits[0].start_position, splits[0].stop_position) split_points_report = [] for _ in splits[0].source.read(range_tracker): split_points_report.append(range_tracker.split_points()) # There will be a total of num_blocks in the generated test file, # proportional to number of records in the file divided by syncronization # interval used by avro during write. Each block has more than 10 records. num_blocks = int(math.ceil(14.5 * num_records / sync_interval)) assert num_blocks > 1 # When reading records of the first block, range_tracker.split_points() # should return (0, iobase.RangeTracker.SPLIT_POINTS_UNKNOWN) self.assertEqual( split_points_report[:10], [(0, iobase.RangeTracker.SPLIT_POINTS_UNKNOWN)] * 10) # When reading records of last block, range_tracker.split_points() should # return (num_blocks - 1, 1) self.assertEqual(split_points_report[-10:], [(num_blocks - 1, 1)] * 10) def test_read_without_splitting_compressed_deflate(self): file_name = self._write_data(codec='deflate') expected_result = self.RECORDS self._run_avro_test(file_name, None, False, expected_result) def test_read_with_splitting_compressed_deflate(self): file_name = self._write_data(codec='deflate') expected_result = self.RECORDS self._run_avro_test(file_name, 100, True, expected_result) @unittest.skipIf(snappy is None, 'python-snappy not installed.') def test_read_without_splitting_compressed_snappy(self): file_name = self._write_data(codec='snappy') expected_result = self.RECORDS self._run_avro_test(file_name, None, False, expected_result) @unittest.skipIf(snappy is None, 'python-snappy not installed.') def test_read_with_splitting_compressed_snappy(self): file_name = self._write_data(codec='snappy') expected_result = self.RECORDS self._run_avro_test(file_name, 100, True, expected_result) def test_read_without_splitting_pattern(self): pattern = self._write_pattern(3) expected_result = self.RECORDS * 3 self._run_avro_test(pattern, None, False, expected_result) def test_read_with_splitting_pattern(self): pattern = self._write_pattern(3) expected_result = self.RECORDS * 3 self._run_avro_test(pattern, 100, True, expected_result) def test_dynamic_work_rebalancing_exhaustive(self): def compare_split_points(file_name): source = _FastAvroSource(file_name) splits = [ split for split in source.split(desired_bundle_size=float('inf')) ] assert len(splits) == 1 source_test_utils.assert_split_at_fraction_exhaustive(splits[0].source) # Adjusting block size so that we can perform a exhaustive dynamic # work rebalancing test that completes within an acceptable amount of time. file_name = self._write_data(count=5, sync_interval=2) compare_split_points(file_name) def test_corrupted_file(self): file_name = self._write_data() with open(file_name, 'rb') as f: data = f.read() # Corrupt the last character of the file which is also the last character of # the last sync_marker. # https://avro.apache.org/docs/current/spec.html#Object+Container+Files corrupted_data = bytearray(data) corrupted_data[-1] = (corrupted_data[-1] + 1) % 256 with tempfile.NamedTemporaryFile(delete=False, prefix=tempfile.template) as f: f.write(corrupted_data) corrupted_file_name = f.name source = _FastAvroSource(corrupted_file_name) with self.assertRaisesRegex(ValueError, r'expected sync marker'): source_test_utils.read_from_source(source, None, None) def test_read_from_avro(self): path = self._write_data() with TestPipeline() as p: assert_that(p | avroio.ReadFromAvro(path), equal_to(self.RECORDS)) def test_read_all_from_avro_single_file(self): path = self._write_data() with TestPipeline() as p: assert_that( p \ | Create([path]) \ | avroio.ReadAllFromAvro(), equal_to(self.RECORDS)) def test_read_all_from_avro_many_single_files(self): path1 = self._write_data() path2 = self._write_data() path3 = self._write_data() with TestPipeline() as p: assert_that( p \ | Create([path1, path2, path3]) \ | avroio.ReadAllFromAvro(), equal_to(self.RECORDS * 3)) def test_read_all_from_avro_file_pattern(self): file_pattern = self._write_pattern(5) with TestPipeline() as p: assert_that( p \ | Create([file_pattern]) \ | avroio.ReadAllFromAvro(), equal_to(self.RECORDS * 5)) def test_read_all_from_avro_many_file_patterns(self): file_pattern1 = self._write_pattern(5) file_pattern2 = self._write_pattern(2) file_pattern3 = self._write_pattern(3) with TestPipeline() as p: assert_that( p \ | Create([file_pattern1, file_pattern2, file_pattern3]) \ | avroio.ReadAllFromAvro(), equal_to(self.RECORDS * 10)) def test_read_all_from_avro_with_filename(self): file_pattern, file_paths = self._write_pattern(3, return_filenames=True) result = [(path, record) for path in file_paths for record in self.RECORDS] with TestPipeline() as p: assert_that( p \ | Create([file_pattern]) \ | avroio.ReadAllFromAvro(with_filename=True), equal_to(result)) class _WriteFilesFn(beam.DoFn): """writes a couple of files with deferral.""" COUNT_STATE = CombiningValueStateSpec('count', combine_fn=sum) def __init__(self, SCHEMA, RECORDS, tempdir): self._thread = None self.SCHEMA = SCHEMA self.RECORDS = RECORDS self.tempdir = tempdir def get_expect(self, match_updated_files): results_file1 = [('file1', x) for x in self.gen_records(1)] results_file2 = [('file2', x) for x in self.gen_records(3)] if match_updated_files: results_file1 += [('file1', x) for x in self.gen_records(2)] return results_file1 + results_file2 def gen_records(self, count): return self.RECORDS * (count // len(self.RECORDS)) + self.RECORDS[:( count % len(self.RECORDS))] def process(self, element, count_state=beam.DoFn.StateParam(COUNT_STATE)): counter = count_state.read() if counter == 0: count_state.add(1) with open(FileSystems.join(self.tempdir, 'file1'), 'wb') as f: writer(f, self.SCHEMA, self.gen_records(2)) with open(FileSystems.join(self.tempdir, 'file2'), 'wb') as f: writer(f, self.SCHEMA, self.gen_records(3)) # convert dumb key to basename in output basename = FileSystems.split(element[1][0])[1] content = element[1][1] yield basename, content def test_read_all_continuously_new(self): with TestPipeline() as pipeline: tempdir = tempfile.mkdtemp() writer_fn = self._WriteFilesFn(self.SCHEMA, self.RECORDS, tempdir) with open(FileSystems.join(tempdir, 'file1'), 'wb') as f: writer(f, writer_fn.SCHEMA, writer_fn.gen_records(1)) match_pattern = FileSystems.join(tempdir, '*') interval = 0.5 last = 2 p_read_once = ( pipeline | 'Continuously read new files' >> avroio.ReadAllFromAvroContinuously( match_pattern, with_filename=True, start_timestamp=Timestamp.now(), interval=interval, stop_timestamp=Timestamp.now() + last, match_updated_files=False) | 'add dumb key' >> beam.Map(lambda x: (0, x)) | 'Write files on-the-fly' >> beam.ParDo(writer_fn)) assert_that( p_read_once, equal_to(writer_fn.get_expect(match_updated_files=False)), label='assert read new files results') def test_read_all_continuously_update(self): with TestPipeline() as pipeline: tempdir = tempfile.mkdtemp() writer_fn = self._WriteFilesFn(self.SCHEMA, self.RECORDS, tempdir) with open(FileSystems.join(tempdir, 'file1'), 'wb') as f: writer(f, writer_fn.SCHEMA, writer_fn.gen_records(1)) match_pattern = FileSystems.join(tempdir, '*') interval = 0.5 last = 2 p_read_upd = ( pipeline | 'Continuously read updated files' >> avroio.ReadAllFromAvroContinuously( match_pattern, with_filename=True, start_timestamp=Timestamp.now(), interval=interval, stop_timestamp=Timestamp.now() + last, match_updated_files=True) | 'add dumb key' >> beam.Map(lambda x: (0, x)) | 'Write files on-the-fly' >> beam.ParDo(writer_fn)) assert_that( p_read_upd, equal_to(writer_fn.get_expect(match_updated_files=True)), label='assert read updated files results') def test_sink_transform(self): with tempfile.NamedTemporaryFile() as dst: path = dst.name with TestPipeline() as p: # pylint: disable=expression-not-assigned p \ | beam.Create(self.RECORDS) \ | avroio.WriteToAvro(path, self.SCHEMA,) with TestPipeline() as p: # json used for stable sortability readback = \ p \ | avroio.ReadFromAvro(path + '*', ) \ | beam.Map(json.dumps) assert_that(readback, equal_to([json.dumps(r) for r in self.RECORDS])) @unittest.skipIf(snappy is None, 'python-snappy not installed.') def test_sink_transform_snappy(self): with tempfile.NamedTemporaryFile() as dst: path = dst.name with TestPipeline() as p: # pylint: disable=expression-not-assigned p \ | beam.Create(self.RECORDS) \ | avroio.WriteToAvro( path, self.SCHEMA, codec='snappy') with TestPipeline() as p: # json used for stable sortability readback = \ p \ | avroio.ReadFromAvro(path + '*') \ | beam.Map(json.dumps) assert_that(readback, equal_to([json.dumps(r) for r in self.RECORDS])) def test_writer_open_and_close(self): # Create and then close a temp file so we can manually open it later dst = tempfile.NamedTemporaryFile(delete=False) dst.close() schema = parse_schema(json.loads(self.SCHEMA_STRING)) sink = _create_avro_sink( 'some_avro_sink', schema, 'null', '.end', 0, None, 'application/x-avro') w = sink.open(dst.name) sink.close(w) os.unlink(dst.name) class TestFastAvro(AvroBase, unittest.TestCase): def __init__(self, methodName='runTest'): super().__init__(methodName) self.SCHEMA = parse_schema(json.loads(self.SCHEMA_STRING)) def _write_data( self, directory=None, prefix=tempfile.template, codec='null', count=len(RECORDS), **kwargs): all_records = self.RECORDS * \ (count // len(self.RECORDS)) + self.RECORDS[:(count % len(self.RECORDS))] with tempfile.NamedTemporaryFile(delete=False, dir=directory, prefix=prefix, mode='w+b') as f: writer(f, self.SCHEMA, all_records, codec=codec, **kwargs) self._temp_files.append(f.name) return f.name class GenerateEvent(beam.PTransform): @staticmethod def sample_data(): return GenerateEvent() def expand(self, input): elemlist = [{'age': 10}, {'age': 20}, {'age': 30}] elem = elemlist return ( input | TestStream().add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 1, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 2, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 3, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 4, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 5, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 5, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 6, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 7, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 8, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 9, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 10, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 10, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 11, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 12, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 13, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 14, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 15, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 15, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 16, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 17, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 18, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 19, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 20, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 20, 0, tzinfo=pytz.UTC).timestamp()).advance_watermark_to( datetime( 2021, 3, 1, 0, 0, 25, 0, tzinfo=pytz.UTC). timestamp()).advance_watermark_to_infinity()) class WriteStreamingTest(unittest.TestCase): def setUp(self): super().setUp() self.tempdir = tempfile.mkdtemp() def tearDown(self): if os.path.exists(self.tempdir): shutil.rmtree(self.tempdir) def test_write_streaming_2_shards_default_shard_name_template( self, num_shards=2): with TestPipeline() as p: output = ( p | GenerateEvent.sample_data() | 'User windowing' >> beam.transforms.core.WindowInto( beam.transforms.window.FixedWindows(60), trigger=beam.transforms.trigger.AfterWatermark(), accumulation_mode=beam.transforms.trigger.AccumulationMode. DISCARDING, allowed_lateness=beam.utils.timestamp.Duration(seconds=0))) #AvroIO avroschema = { 'name': 'dummy', # your supposed to be file name with .avro extension 'type': 'record', # type of avro serilazation, there are more (see # above docs) 'fields': [ # this defines actual keys & their types {'name': 'age', 'type': 'int'}, ], } output2 = output | 'WriteToAvro' >> beam.io.WriteToAvro( file_path_prefix=self.tempdir + "/ouput_WriteToAvro", file_name_suffix=".avro", num_shards=num_shards, schema=avroschema) _ = output2 | 'LogElements after WriteToAvro' >> LogElements( prefix='after WriteToAvro ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToAvro-[1614556800.0, 1614556805.0)-00000-of-00002.avro # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P[\d\.]+), ' r'(?P[\d\.]+|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.avro$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToAvro*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) self.assertEqual( len(file_names), num_shards, "expected %d files, but got: %d" % (num_shards, len(file_names))) def test_write_streaming_2_shards_custom_shard_name_template( self, num_shards=2, shard_name_template='-V-SSSSS-of-NNNNN'): with TestPipeline() as p: output = (p | GenerateEvent.sample_data()) #AvroIO avroschema = { 'name': 'dummy', # your supposed to be file name with .avro extension 'type': 'record', # type of avro serilazation 'fields': [ # this defines actual keys & their types {'name': 'age', 'type': 'int'}, ], } output2 = output | 'WriteToAvro' >> beam.io.WriteToAvro( file_path_prefix=self.tempdir + "/ouput_WriteToAvro", file_name_suffix=".avro", shard_name_template=shard_name_template, num_shards=num_shards, triggering_frequency=60, schema=avroschema) _ = output2 | 'LogElements after WriteToAvro' >> LogElements( prefix='after WriteToAvro ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToAvro-[2021-03-01T00-00-00, 2021-03-01T00-01-00)- # 00000-of-00002.avro # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}), ' r'(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.avro$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToAvro*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) self.assertEqual( len(file_names), num_shards, "expected %d files, but got: %d" % (num_shards, len(file_names))) def test_write_streaming_2_shards_custom_shard_name_template_5s_window( self, num_shards=2, shard_name_template='-V-SSSSS-of-NNNNN', triggering_frequency=5): with TestPipeline() as p: output = (p | GenerateEvent.sample_data()) #AvroIO avroschema = { 'name': 'dummy', # your supposed to be file name with .avro extension 'type': 'record', # type of avro serilazation 'fields': [ # this defines actual keys & their types {'name': 'age', 'type': 'int'}, ], } output2 = output | 'WriteToAvro' >> beam.io.WriteToAvro( file_path_prefix=self.tempdir + "/ouput_WriteToAvro", file_name_suffix=".txt", shard_name_template=shard_name_template, num_shards=num_shards, triggering_frequency=triggering_frequency, schema=avroschema) _ = output2 | 'LogElements after WriteToAvro' >> LogElements( prefix='after WriteToAvro ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToAvro-[2021-03-01T00-00-00, 2021-03-01T00-01-00)- # 00000-of-00002.avro # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}), ' r'(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.txt$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToAvro*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) # for 5s window size, the input should be processed by 5 windows with # 2 shards per window self.assertEqual( len(file_names), 10, "expected %d files, but got: %d" % (num_shards, len(file_names))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3043997 apache_beam-2.74.0/apache_beam/io/aws/0000755000175100017510000000000015206113157017066 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/aws/__init__.py0000644000175100017510000000142015206113032021164 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3043997 apache_beam-2.74.0/apache_beam/io/aws/clients/0000755000175100017510000000000015206113157020527 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/aws/clients/__init__.py0000644000175100017510000000142015206113032022625 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3074002 apache_beam-2.74.0/apache_beam/io/aws/clients/s3/0000755000175100017510000000000015206113157021054 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/aws/clients/s3/__init__.py0000644000175100017510000000142015206113032023152 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/aws/clients/s3/boto3_client.py0000644000175100017510000002272715206113032024014 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file from apache_beam.io.aws.clients.s3 import messages from apache_beam.options import pipeline_options from apache_beam.utils import retry try: # pylint: disable=wrong-import-order, wrong-import-position # pylint: disable=ungrouped-imports import boto3 except ImportError: boto3 = None def get_http_error_code(exc): if hasattr(exc, 'response'): return exc.response.get('ResponseMetadata', {}).get('HTTPStatusCode') return None class Client(object): """ Wrapper for boto3 library """ def __init__(self, options): assert boto3 is not None, 'Missing boto3 requirement' if isinstance(options, pipeline_options.PipelineOptions): s3_options = options.view_as(pipeline_options.S3Options) access_key_id = s3_options.s3_access_key_id secret_access_key = s3_options.s3_secret_access_key session_token = s3_options.s3_session_token endpoint_url = s3_options.s3_endpoint_url use_ssl = not s3_options.s3_disable_ssl region_name = s3_options.s3_region_name api_version = s3_options.s3_api_version verify = s3_options.s3_verify else: access_key_id = options.get('s3_access_key_id') secret_access_key = options.get('s3_secret_access_key') session_token = options.get('s3_session_token') endpoint_url = options.get('s3_endpoint_url') use_ssl = not options.get('s3_disable_ssl', False) region_name = options.get('s3_region_name') api_version = options.get('s3_api_version') verify = options.get('s3_verify') session = boto3.session.Session() self.client = session.client( service_name='s3', region_name=region_name, api_version=api_version, use_ssl=use_ssl, verify=verify, endpoint_url=endpoint_url, aws_access_key_id=access_key_id, aws_secret_access_key=secret_access_key, aws_session_token=session_token) self._download_request = None self._download_stream = None self._download_pos = 0 def get_object_metadata(self, request): """Retrieves an object's metadata. Args: request: (GetRequest) input message Returns: (Object) The response message. """ kwargs = {'Bucket': request.bucket, 'Key': request.object} try: boto_response = self.client.head_object(**kwargs) except Exception as e: raise messages.S3ClientError(str(e), get_http_error_code(e)) item = messages.Item( boto_response['ETag'], request.object, boto_response['LastModified'], boto_response['ContentLength'], boto_response['ContentType']) return item def get_stream(self, request, start): """Opens a stream object starting at the given position. Args: request: (GetRequest) request start: (int) start offset Returns: (Stream) Boto3 stream object. """ if self._download_request and ( start != self._download_pos or request.bucket != self._download_request.bucket or request.object != self._download_request.object): self._download_stream.close() self._download_stream = None # noinspection PyProtectedMember if not self._download_stream or self._download_stream._raw_stream.closed: try: self._download_stream = self.client.get_object( Bucket=request.bucket, Key=request.object, Range='bytes={}-'.format(start))['Body'] self._download_request = request self._download_pos = start except Exception as e: raise messages.S3ClientError(str(e), get_http_error_code(e)) return self._download_stream @retry.with_exponential_backoff() def get_range(self, request, start, end): r"""Retrieves an object's contents. Args: request: (GetRequest) request start: (int) start offset end: (int) end offset (exclusive) Returns: (bytes) The response message. """ for i in range(2): try: stream = self.get_stream(request, start) data = stream.read(end - start) self._download_pos += len(data) return data except Exception as e: self._download_stream = None self._download_request = None if i == 0: # Read errors are likely with long-lived connections, retry immediately if a read fails once continue if isinstance(e, messages.S3ClientError): raise e raise messages.S3ClientError(str(e), get_http_error_code(e)) def list(self, request): r"""Retrieves a list of objects matching the criteria. Args: request: (ListRequest) input message Returns: (ListResponse) The response message. """ kwargs = {'Bucket': request.bucket, 'Prefix': request.prefix} if request.continuation_token is not None: kwargs['ContinuationToken'] = request.continuation_token try: boto_response = self.client.list_objects_v2(**kwargs) except Exception as e: raise messages.S3ClientError(str(e), get_http_error_code(e)) if boto_response['KeyCount'] == 0: message = 'Tried to list nonexistent S3 path: s3://%s/%s' % ( request.bucket, request.prefix) raise messages.S3ClientError(message, 404) items = [ messages.Item( etag=content['ETag'], key=content['Key'], last_modified=content['LastModified'], size=content['Size']) for content in boto_response['Contents'] ] try: next_token = boto_response['NextContinuationToken'] except KeyError: next_token = None response = messages.ListResponse(items, next_token) return response def create_multipart_upload(self, request): r"""Initates a multipart upload to S3 for a given object Args: request: (UploadRequest) input message Returns: (UploadResponse) The response message. """ try: boto_response = self.client.create_multipart_upload( Bucket=request.bucket, Key=request.object, ContentType=request.mime_type) response = messages.UploadResponse(boto_response['UploadId']) except Exception as e: raise messages.S3ClientError(str(e), get_http_error_code(e)) return response def upload_part(self, request): r"""Uploads part of a file to S3 during a multipart upload Args: request: (UploadPartRequest) input message Returns: (UploadPartResponse) The response message. """ try: boto_response = self.client.upload_part( Body=request.bytes, Bucket=request.bucket, Key=request.object, PartNumber=request.part_number, UploadId=request.upload_id) response = messages.UploadPartResponse( boto_response['ETag'], request.part_number) return response except Exception as e: raise messages.S3ClientError(str(e), get_http_error_code(e)) def complete_multipart_upload(self, request): r"""Completes a multipart upload to S3 Args: request: (UploadPartRequest) input message Returns: (Void) The response message. """ parts = {'Parts': request.parts} try: self.client.complete_multipart_upload( Bucket=request.bucket, Key=request.object, UploadId=request.upload_id, MultipartUpload=parts) except Exception as e: raise messages.S3ClientError(str(e), get_http_error_code(e)) def delete(self, request): r"""Deletes given object from bucket Args: request: (DeleteRequest) input message Returns: (void) Void, otherwise will raise if an error occurs """ try: self.client.delete_object(Bucket=request.bucket, Key=request.object) except Exception as e: raise messages.S3ClientError(str(e), get_http_error_code(e)) def delete_batch(self, request): aws_request = { 'Bucket': request.bucket, 'Delete': { 'Objects': [{ 'Key': object } for object in request.objects] } } try: aws_response = self.client.delete_objects(**aws_request) except Exception as e: raise messages.S3ClientError(str(e), get_http_error_code(e)) deleted = [obj['Key'] for obj in aws_response.get('Deleted', [])] failed = [obj['Key'] for obj in aws_response.get('Errors', [])] errors = [ messages.S3ClientError(obj['Message'], obj['Code']) for obj in aws_response.get('Errors', []) ] return messages.DeleteBatchResponse(deleted, failed, errors) def copy(self, request): try: copy_src = {'Bucket': request.src_bucket, 'Key': request.src_key} self.client.copy(copy_src, request.dest_bucket, request.dest_key) except Exception as e: raise messages.S3ClientError(str(e), get_http_error_code(e)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/aws/clients/s3/client_test.py0000644000175100017510000001656715206113032023752 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import os import unittest from apache_beam.io.aws import s3io from apache_beam.io.aws.clients.s3 import fake_client from apache_beam.io.aws.clients.s3 import messages from apache_beam.options import pipeline_options class ClientErrorTest(unittest.TestCase): def setUp(self): # These tests can be run locally against a mock S3 client, or as integration # tests against the real S3 client. self.USE_MOCK = True # If you're running integration tests with S3, set this variable to be an # s3 path that you have access to where test data can be written. If you're # just running tests against the mock, this can be any s3 path. It should # end with a '/'. self.TEST_DATA_PATH = 's3://random-data-sets/beam_tests/' self.test_bucket, self.test_path = s3io.parse_s3_path(self.TEST_DATA_PATH) if self.USE_MOCK: self.client = fake_client.FakeS3Client() test_data_bucket, _ = s3io.parse_s3_path(self.TEST_DATA_PATH) self.client.known_buckets.add(test_data_bucket) self.aws = s3io.S3IO(self.client) else: self.aws = s3io.S3IO(options=pipeline_options.S3Options()) def test_get_object_metadata(self): # Test nonexistent object object = self.test_path + 'nonexistent_file_doesnt_exist' request = messages.GetRequest(self.test_bucket, object) self.assertRaises( messages.S3ClientError, self.client.get_object_metadata, request) try: self.client.get_object_metadata(request) except Exception as e: self.assertIsInstance(e, messages.S3ClientError) self.assertEqual(e.code, 404) def test_get_range_nonexistent(self): # Test nonexistent object object = self.test_path + 'nonexistent_file_doesnt_exist' request = messages.GetRequest(self.test_bucket, object) self.assertRaises( messages.S3ClientError, self.client.get_range, request, 0, 10) try: self.client.get_range(request, 0, 10) except Exception as e: self.assertIsInstance(e, messages.S3ClientError) self.assertEqual(e.code, 404) def test_get_range_bad_start_end(self): file_name = self.TEST_DATA_PATH + 'get_range' contents = os.urandom(1024) with self.aws.open(file_name, 'w') as f: f.write(contents) bucket, object = s3io.parse_s3_path(file_name) response = self.client.get_range( messages.GetRequest(bucket, object), -10, 20) self.assertEqual(response, contents) response = self.client.get_range( messages.GetRequest(bucket, object), 20, 10) self.assertEqual(response, contents) # Clean up self.aws.delete(file_name) def test_upload_part_nonexistent_upload_id(self): object = self.test_path + 'upload_part' upload_id = 'not-an-id-12345' part_number = 1 contents = os.urandom(1024) request = messages.UploadPartRequest( self.test_bucket, object, upload_id, part_number, contents) self.assertRaises(messages.S3ClientError, self.client.upload_part, request) try: self.client.upload_part(request) except Exception as e: self.assertIsInstance(e, messages.S3ClientError) self.assertEqual(e.code, 404) def test_copy_nonexistent(self): src_key = self.test_path + 'not_a_real_file_does_not_exist' dest_key = self.test_path + 'destination_file_location' request = messages.CopyRequest( self.test_bucket, src_key, self.test_bucket, dest_key) with self.assertRaises(messages.S3ClientError) as e: self.client.copy(request) self.assertEqual(e.exception.code, 404) def test_upload_part_bad_number(self): object = self.test_path + 'upload_part' contents = os.urandom(1024) request = messages.UploadRequest(self.test_bucket, object, None) response = self.client.create_multipart_upload(request) upload_id = response.upload_id part_number = 0.5 request = messages.UploadPartRequest( self.test_bucket, object, upload_id, part_number, contents) self.assertRaises(messages.S3ClientError, self.client.upload_part, request) try: response = self.client.upload_part(request) except Exception as e: self.assertIsInstance(e, messages.S3ClientError) self.assertEqual(e.code, 400) def test_complete_multipart_upload_too_small(self): object = self.test_path + 'upload_part' request = messages.UploadRequest(self.test_bucket, object, None) response = self.client.create_multipart_upload(request) upload_id = response.upload_id part_number = 1 contents_1 = os.urandom(1024) request_1 = messages.UploadPartRequest( self.test_bucket, object, upload_id, part_number, contents_1) response_1 = self.client.upload_part(request_1) part_number = 2 contents_2 = os.urandom(1024) request_2 = messages.UploadPartRequest( self.test_bucket, object, upload_id, part_number, contents_2) response_2 = self.client.upload_part(request_2) parts = [{ 'PartNumber': 1, 'ETag': response_1.etag }, { 'PartNumber': 2, 'ETag': response_2.etag }] complete_request = messages.CompleteMultipartUploadRequest( self.test_bucket, object, upload_id, parts) try: self.client.complete_multipart_upload(complete_request) except Exception as e: self.assertIsInstance(e, messages.S3ClientError) self.assertEqual(e.code, 400) def test_complete_multipart_upload_too_many(self): object = self.test_path + 'upload_part' request = messages.UploadRequest(self.test_bucket, object, None) response = self.client.create_multipart_upload(request) upload_id = response.upload_id part_number = 1 contents_1 = os.urandom(5 * 1024) request_1 = messages.UploadPartRequest( self.test_bucket, object, upload_id, part_number, contents_1) response_1 = self.client.upload_part(request_1) part_number = 2 contents_2 = os.urandom(1024) request_2 = messages.UploadPartRequest( self.test_bucket, object, upload_id, part_number, contents_2) response_2 = self.client.upload_part(request_2) parts = [ { 'PartNumber': 1, 'ETag': response_1.etag }, { 'PartNumber': 2, 'ETag': response_2.etag }, { 'PartNumber': 3, 'ETag': 'fake-etag' }, ] complete_request = messages.CompleteMultipartUploadRequest( self.test_bucket, object, upload_id, parts) try: self.client.complete_multipart_upload(complete_request) except Exception as e: self.assertIsInstance(e, messages.S3ClientError) self.assertEqual(e.code, 400) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/aws/clients/s3/fake_client.py0000644000175100017510000001712215206113032023665 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import datetime import time import pytz from apache_beam.io.aws.clients.s3 import messages class FakeFile(object): def __init__(self, bucket, key, contents, etag=None): self.bucket = bucket self.key = key self.contents = contents self.last_modified = time.time() if not etag: self.etag = '"%s-1"' % ('x' * 32) else: self.etag = etag def get_metadata(self): last_modified_datetime = None if self.last_modified: last_modified_datetime = datetime.datetime.fromtimestamp( self.last_modified, pytz.utc) return messages.Item( self.etag, self.key, last_modified_datetime, len(self.contents), mime_type=None) class FakeS3Client(object): def __init__(self): self.files = {} self.list_continuation_tokens = {} self.multipart_uploads = {} # boto3 has different behavior when running some operations against a bucket # that exists vs. against one that doesn't. To emulate that behavior, the # mock client keeps a set of bucket names that it knows "exist". self.known_buckets = set() def add_file(self, f): self.files[(f.bucket, f.key)] = f if f.bucket not in self.known_buckets: self.known_buckets.add(f.bucket) def get_file(self, bucket, obj): try: return self.files[bucket, obj] except: raise messages.S3ClientError('Not Found', 404) def delete_file(self, bucket, obj): del self.files[(bucket, obj)] def get_object_metadata(self, request): r"""Retrieves an object's metadata. Args: request: (GetRequest) input message Returns: (Item) The response message. """ # TODO: Do we want to mock out a lack of credentials? file_ = self.get_file(request.bucket, request.object) return file_.get_metadata() def list(self, request): bucket = request.bucket prefix = request.prefix or '' matching_files = [] for file_bucket, file_name in sorted(iter(self.files)): if bucket == file_bucket and file_name.startswith(prefix): file_object = self.get_file(file_bucket, file_name).get_metadata() matching_files.append(file_object) if not matching_files: message = 'Tried to list nonexistent S3 path: s3://%s/%s' % ( bucket, prefix) raise messages.S3ClientError(message, 404) # Handle pagination. items_per_page = 5 if not request.continuation_token: range_start = 0 else: if request.continuation_token not in self.list_continuation_tokens: raise ValueError('Invalid page token.') range_start = self.list_continuation_tokens[request.continuation_token] del self.list_continuation_tokens[request.continuation_token] result = messages.ListResponse( items=matching_files[range_start:range_start + items_per_page]) if range_start + items_per_page < len(matching_files): next_range_start = range_start + items_per_page next_continuation_token = '_page_token_%s_%s_%d' % ( bucket, prefix, next_range_start) self.list_continuation_tokens[next_continuation_token] = next_range_start result.next_token = next_continuation_token return result def get_range(self, request, start, end): r"""Retrieves an object. Args: request: (GetRequest) request Returns: (bytes) The response message. """ file_ = self.get_file(request.bucket, request.object) # Replicates S3's behavior, per the spec here: # https://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.35 if start < 0 or end <= start: return file_.contents return file_.contents[start:end] def delete(self, request): if request.bucket not in self.known_buckets: raise messages.S3ClientError('The specified bucket does not exist', 404) if (request.bucket, request.object) in self.files: self.delete_file(request.bucket, request.object) else: # S3 doesn't raise an error if you try to delete a nonexistent file from # an extant bucket return def delete_batch(self, request): deleted, failed, errors = [], [], [] for object in request.objects: try: delete_request = messages.DeleteRequest(request.bucket, object) self.delete(delete_request) deleted.append(object) except messages.S3ClientError as e: failed.append(object) errors.append(e) return messages.DeleteBatchResponse(deleted, failed, errors) def copy(self, request): src_file = self.get_file(request.src_bucket, request.src_key) dest_file = FakeFile( request.dest_bucket, request.dest_key, src_file.contents) self.add_file(dest_file) def create_multipart_upload(self, request): # Create hash of bucket and key # Store upload_id internally upload_id = request.bucket + request.object self.multipart_uploads[upload_id] = {} return messages.UploadResponse(upload_id) def upload_part(self, request): # Save off bytes passed to internal data store upload_id, part_number = request.upload_id, request.part_number if part_number < 0 or not isinstance(part_number, int): raise messages.S3ClientError( 'Param validation failed on part number', 400) if upload_id not in self.multipart_uploads: raise messages.S3ClientError('The specified upload does not exist', 404) self.multipart_uploads[upload_id][part_number] = request.bytes etag = '"%s"' % ('x' * 32) return messages.UploadPartResponse(etag, part_number) def complete_multipart_upload(self, request): MIN_PART_SIZE = 5 * 2**10 # 5 KiB parts_received = self.multipart_uploads[request.upload_id] # Check that we got all the parts that they intended to send part_numbers_to_confirm = set(part['PartNumber'] for part in request.parts) # Make sure all the expected parts are present if part_numbers_to_confirm != set(parts_received.keys()): raise messages.S3ClientError( 'One or more of the specified parts could not be found', 400) # Sort by part number sorted_parts = sorted(parts_received.items(), key=lambda pair: pair[0]) sorted_bytes = [bytes_ for (_, bytes_) in sorted_parts] # Make sure that the parts aren't too small (except the last part) part_sizes = [len(bytes_) for bytes_ in sorted_bytes] if any(size < MIN_PART_SIZE for size in part_sizes[:-1]): e_message = """ All parts but the last must be larger than %d bytes """ % MIN_PART_SIZE raise messages.S3ClientError(e_message, 400) # String together all bytes for the given upload final_contents = b''.join(sorted_bytes) # Create FakeFile object num_parts = len(parts_received) etag = '"%s-%d"' % ('x' * 32, num_parts) file_ = FakeFile(request.bucket, request.object, final_contents, etag=etag) # Store FakeFile in self.files self.add_file(file_) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/aws/clients/s3/messages.py0000644000175100017510000000766715206113032023245 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file class GetRequest(): """ S3 request object for `Get` command """ def __init__(self, bucket, object): self.bucket = bucket self.object = object class UploadResponse(): """ S3 response object for `StartUpload` command """ def __init__(self, upload_id): self.upload_id = upload_id class UploadRequest(): """ S3 request object for `StartUpload` command """ def __init__(self, bucket, object, mime_type): self.bucket = bucket self.object = object self.mime_type = mime_type class UploadPartRequest(): """ S3 request object for `UploadPart` command """ def __init__(self, bucket, object, upload_id, part_number, bytes): self.bucket = bucket self.object = object self.upload_id = upload_id self.part_number = part_number self.bytes = bytes # self.mime_type = mime_type class UploadPartResponse(): """ S3 response object for `UploadPart` command """ def __init__(self, etag, part_number): self.etag = etag self.part_number = part_number class CompleteMultipartUploadRequest(): """ S3 request object for `UploadPart` command """ def __init__(self, bucket, object, upload_id, parts): # parts is a list of objects of the form # {'ETag': response.etag, 'PartNumber': response.part_number} self.bucket = bucket self.object = object self.upload_id = upload_id self.parts = parts # self.mime_type = mime_type class ListRequest(): """ S3 request object for `List` command """ def __init__(self, bucket, prefix, continuation_token=None): self.bucket = bucket self.prefix = prefix self.continuation_token = continuation_token class ListResponse(): """ S3 response object for `List` command """ def __init__(self, items, next_token=None): self.items = items self.next_token = next_token class Item(): """ An item in S3 """ def __init__(self, etag, key, last_modified, size, mime_type=None): self.etag = etag self.key = key self.last_modified = last_modified self.size = size self.mime_type = mime_type class DeleteRequest(): """ S3 request object for `Delete` command """ def __init__(self, bucket, object): self.bucket = bucket self.object = object class DeleteBatchRequest(): def __init__(self, bucket, objects): # `objects` is a list of strings corresponding to the keys to be deleted # in the bucket self.bucket = bucket self.objects = objects class DeleteBatchResponse(): def __init__(self, deleted, failed, errors): # `deleted` is a list of strings corresponding to the keys that were deleted # `failed` is a list of strings corresponding to the keys that caused errors # `errors` is a list of S3ClientErrors, aligned with the order of `failed` self.deleted = deleted self.failed = failed self.errors = errors class CopyRequest(): def __init__(self, src_bucket, src_key, dest_bucket, dest_key): self.src_bucket = src_bucket self.src_key = src_key self.dest_bucket = dest_bucket self.dest_key = dest_key class S3ClientError(Exception): def __init__(self, message=None, code=None): self.message = message self.code = code ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/aws/s3filesystem.py0000644000175100017510000002501115206113032022061 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """S3 file system implementation for accessing files on AWS S3.""" # pytype: skip-file import traceback from apache_beam.io.aws import s3io from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystem import CompressedFile from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystem import FileMetadata from apache_beam.io.filesystem import FileSystem __all__ = ['S3FileSystem'] class S3FileSystem(FileSystem): """An S3 `FileSystem` implementation for accessing files on AWS S3 """ CHUNK_SIZE = s3io.MAX_BATCH_OPERATION_SIZE S3_PREFIX = 's3://' def __init__(self, pipeline_options): """Initializes a connection to S3. Connection configuration is done by passing pipeline options. See :class:`~apache_beam.options.pipeline_options.S3Options`. """ super().__init__(pipeline_options) self._options = pipeline_options @classmethod def scheme(cls): """URI scheme for the FileSystem """ return 's3' def join(self, basepath, *paths): """Join two or more pathname components for the filesystem Args: basepath: string path of the first component of the path paths: path components to be added Returns: full path after combining all of the return nulled components """ if not basepath.startswith(S3FileSystem.S3_PREFIX): raise ValueError('Basepath %r must be S3 path.' % basepath) path = basepath for p in paths: path = path.rstrip('/') + '/' + p.lstrip('/') return path def split(self, path): """Splits the given path into two parts. Splits the path into a pair (head, tail) such that tail contains the last component of the path and head contains everything up to that. Head will include the S3 prefix ('s3://'). Args: path: path as a string Returns: a pair of path components as strings. """ path = path.strip() if not path.startswith(S3FileSystem.S3_PREFIX): raise ValueError('Path %r must be S3 path.' % path) prefix_len = len(S3FileSystem.S3_PREFIX) last_sep = path[prefix_len:].rfind('/') if last_sep >= 0: last_sep += prefix_len if last_sep > 0: return (path[:last_sep], path[last_sep + 1:]) elif last_sep < 0: return (path, '') else: raise ValueError('Invalid path: %s' % path) def mkdirs(self, path): """Recursively create directories for the provided path. Args: path: string path of the directory structure that should be created Raises: IOError: if leaf directory already exists. """ pass def has_dirs(self): """Whether this FileSystem supports directories.""" return False def _list(self, dir_or_prefix): """List files in a location. Listing is non-recursive, for filesystems that support directories. Args: dir_or_prefix: (string) A directory or location prefix (for filesystems that don't have directories). Returns: Generator of ``FileMetadata`` objects. Raises: ``BeamIOError``: if listing fails, but not if no files were found. """ try: for path, (size, updated) in s3io.S3IO(options=self._options).list_files( dir_or_prefix, with_metadata=True): yield FileMetadata(path, size, updated) except Exception as e: # pylint: disable=broad-except raise BeamIOError("List operation failed", {dir_or_prefix: e}) def _path_open( self, path, mode, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO): """Helper functions to open a file in the provided mode. """ compression_type = FileSystem._get_compression_type(path, compression_type) mime_type = CompressionTypes.mime_type(compression_type, mime_type) raw_file = s3io.S3IO(options=self._options).open( path, mode, mime_type=mime_type) if compression_type == CompressionTypes.UNCOMPRESSED: return raw_file return CompressedFile(raw_file, compression_type=compression_type) def create( self, path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO): """Returns a write channel for the given file path. Args: path: string path of the file object to be written to the system mime_type: MIME type to specify the type of content in the file object compression_type: Type of compression to be used for this object Returns: file handle with a close function for the user to use """ return self._path_open(path, 'wb', mime_type, compression_type) def open( self, path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO): """Returns a read channel for the given file path. Args: path: string path of the file object to be written to the system mime_type: MIME type to specify the type of content in the file object compression_type: Type of compression to be used for this object Returns: file handle with a close function for the user to use """ return self._path_open(path, 'rb', mime_type, compression_type) def copy(self, source_file_names, destination_file_names): """Recursively copy the file tree from the source to the destination Args: source_file_names: list of source file objects that needs to be copied destination_file_names: list of destination of the new object Raises: ``BeamIOError``: if any of the copy operations fail """ if not len(source_file_names) == len(destination_file_names): message = 'Unable to copy unequal number of sources and destinations' raise BeamIOError(message) src_dest_pairs = list(zip(source_file_names, destination_file_names)) return s3io.S3IO(options=self._options).copy_paths(src_dest_pairs) def rename(self, source_file_names, destination_file_names): """Rename the files at the source list to the destination list. Source and destination lists should be of the same size. Args: source_file_names: List of file paths that need to be moved destination_file_names: List of destination_file_names for the files Raises: ``BeamIOError``: if any of the rename operations fail """ if not len(source_file_names) == len(destination_file_names): message = 'Unable to rename unequal number of sources and destinations' raise BeamIOError(message) src_dest_pairs = list(zip(source_file_names, destination_file_names)) results = s3io.S3IO(options=self._options).rename_files(src_dest_pairs) exceptions = {(src, dest): error for (src, dest, error) in results if error is not None} if exceptions: raise BeamIOError("Rename operation failed", exceptions) def exists(self, path): """Check if the provided path exists on the FileSystem. Args: path: string path that needs to be checked. Returns: boolean flag indicating if path exists """ try: return s3io.S3IO(options=self._options).exists(path) except Exception as e: # pylint: disable=broad-except raise BeamIOError("exists() operation failed", {path: e}) def size(self, path): """Get size of path on the FileSystem. Args: path: string path in question. Returns: int size of path according to the FileSystem. Raises: ``BeamIOError``: if path doesn't exist. """ try: return s3io.S3IO(options=self._options).size(path) except Exception as e: # pylint: disable=broad-except raise BeamIOError("size() operation failed", {path: e}) def last_updated(self, path): """Get UNIX Epoch time in seconds on the FileSystem. Args: path: string path of file. Returns: float UNIX Epoch time Raises: ``BeamIOError``: if path doesn't exist. """ try: return s3io.S3IO(options=self._options).last_updated(path) except Exception as e: # pylint: disable=broad-except raise BeamIOError("last_updated operation failed", {path: e}) def checksum(self, path): """Fetch checksum metadata of a file on the :class:`~apache_beam.io.filesystem.FileSystem`. Args: path: string path of a file. Returns: string containing checksum Raises: ``BeamIOError``: if path isn't a file or doesn't exist. """ try: return s3io.S3IO(options=self._options).checksum(path) except Exception as e: # pylint: disable=broad-except raise BeamIOError("Checksum operation failed", {path: e}) def metadata(self, path): """Fetch metadata fields of a file on the FileSystem. Args: path: string path of a file. Returns: :class:`~apache_beam.io.filesystem.FileMetadata`. Raises: ``BeamIOError``: if path isn't a file or doesn't exist. """ try: file_metadata = s3io.S3IO(options=self._options)._status(path) return FileMetadata( path, file_metadata['size'], file_metadata['last_updated']) except Exception as e: # pylint: disable=broad-except raise BeamIOError("Metadata operation failed", {path: e}) def delete(self, paths): """Deletes files or directories at the provided paths. Directories will be deleted recursively. Args: paths: list of paths that give the file objects to be deleted """ results = s3io.S3IO(options=self._options).delete_paths(paths) exceptions = { path: error for (path, error) in results.items() if error is not None } if exceptions: raise BeamIOError("Delete operation failed", exceptions) def report_lineage(self, path, lineage): try: components = s3io.parse_s3_path(path, object_optional=True) except ValueError: # report lineage is fail-safe traceback.print_exc() return if components and not components[-1]: components = components[:-1] lineage.add('s3', *components, last_segment_sep='/') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/aws/s3filesystem_test.py0000644000175100017510000002464315206113032023132 0ustar00runnerrunner# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the S3 File System""" # pytype: skip-file import logging import unittest import mock from apache_beam.io.aws.clients.s3 import messages from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystem import FileMetadata from apache_beam.options.pipeline_options import PipelineOptions # Protect against environments where boto3 library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.io.aws import s3filesystem except ImportError: s3filesystem = None # type: ignore[assignment] # pylint: enable=wrong-import-order, wrong-import-position @unittest.skipIf(s3filesystem is None, 'AWS dependencies are not installed') class S3FileSystemTest(unittest.TestCase): def setUp(self): pipeline_options = PipelineOptions() self.fs = s3filesystem.S3FileSystem(pipeline_options=pipeline_options) def test_scheme(self): self.assertEqual(self.fs.scheme(), 's3') self.assertEqual(s3filesystem.S3FileSystem.scheme(), 's3') def test_join(self): self.assertEqual( 's3://bucket/path/to/file', self.fs.join('s3://bucket/path', 'to', 'file')) self.assertEqual( 's3://bucket/path/to/file', self.fs.join('s3://bucket/path', 'to/file')) self.assertEqual( 's3://bucket/path/to/file', self.fs.join('s3://bucket/path', '/to/file')) self.assertEqual( 's3://bucket/path/to/file', self.fs.join('s3://bucket/path/', 'to', 'file')) self.assertEqual( 's3://bucket/path/to/file', self.fs.join('s3://bucket/path/', 'to/file')) self.assertEqual( 's3://bucket/path/to/file', self.fs.join('s3://bucket/path/', '/to/file')) with self.assertRaises(ValueError): self.fs.join('/bucket/path/', '/to/file') def test_split(self): self.assertEqual(('s3://foo/bar', 'baz'), self.fs.split('s3://foo/bar/baz')) self.assertEqual(('s3://foo', ''), self.fs.split('s3://foo/')) self.assertEqual(('s3://foo', ''), self.fs.split('s3://foo')) with self.assertRaises(ValueError): self.fs.split('/no/s3/prefix') @mock.patch('apache_beam.io.aws.s3filesystem.s3io') def test_match_single(self, unused_mock_arg): # Prepare mocks. s3io_mock = mock.MagicMock() s3filesystem.s3io.S3IO = lambda options: s3io_mock # type: ignore[misc] s3io_mock._status.return_value = {'size': 1, 'last_updated': 9999999.0} expected_results = [FileMetadata('s3://bucket/file1', 1, 9999999.0)] match_result = self.fs.match(['s3://bucket/file1'])[0] self.assertEqual(match_result.metadata_list, expected_results) s3io_mock._status.assert_called_once_with('s3://bucket/file1') @mock.patch('apache_beam.io.aws.s3filesystem.s3io') def test_match_multiples(self, unused_mock_arg): # Prepare mocks. s3io_mock = mock.MagicMock() s3filesystem.s3io.S3IO = lambda options: s3io_mock # type: ignore[misc] s3io_mock.list_files.return_value = iter([ ('s3://bucket/file1', (1, 9999999.0)), ('s3://bucket/file2', (2, 8888888.0)) ]) expected_results = set([ FileMetadata('s3://bucket/file1', 1, 9999999.0), FileMetadata('s3://bucket/file2', 2, 8888888.0) ]) match_result = self.fs.match(['s3://bucket/'])[0] self.assertEqual(set(match_result.metadata_list), expected_results) s3io_mock.list_files.assert_called_once_with( 's3://bucket/', with_metadata=True) @mock.patch('apache_beam.io.aws.s3filesystem.s3io') def test_match_multiples_limit(self, unused_mock_arg): # Prepare mocks. s3io_mock = mock.MagicMock() limit = 1 s3filesystem.s3io.S3IO = lambda options: s3io_mock # type: ignore[misc] s3io_mock.list_files.return_value = iter([ ('s3://bucket/file1', (1, 99999.0)) ]) expected_results = set([FileMetadata('s3://bucket/file1', 1, 99999.0)]) match_result = self.fs.match(['s3://bucket/'], [limit])[0] self.assertEqual(set(match_result.metadata_list), expected_results) self.assertEqual(len(match_result.metadata_list), limit) s3io_mock.list_files.assert_called_once_with( 's3://bucket/', with_metadata=True) @mock.patch('apache_beam.io.aws.s3filesystem.s3io') def test_match_multiples_error(self, unused_mock_arg): # Prepare mocks. s3io_mock = mock.MagicMock() s3filesystem.s3io.S3IO = lambda options: s3io_mock # type: ignore[misc] exception = IOError('Failed') s3io_mock.list_files.side_effect = exception with self.assertRaises(BeamIOError) as error: self.fs.match(['s3://bucket/']) self.assertIn('Match operation failed', str(error.exception)) s3io_mock.list_files.assert_called_once_with( 's3://bucket/', with_metadata=True) @mock.patch('apache_beam.io.aws.s3filesystem.s3io') def test_match_multiple_patterns(self, unused_mock_arg): # Prepare mocks. s3io_mock = mock.MagicMock() s3filesystem.s3io.S3IO = lambda options: s3io_mock # type: ignore[misc] s3io_mock.list_files.side_effect = [ iter([('s3://bucket/file1', (1, 99999.0))]), iter([('s3://bucket/file2', (2, 88888.0))]), ] expected_results = [[FileMetadata('s3://bucket/file1', 1, 99999.0)], [FileMetadata('s3://bucket/file2', 2, 88888.0)]] result = self.fs.match(['s3://bucket/file1*', 's3://bucket/file2*']) self.assertEqual([mr.metadata_list for mr in result], expected_results) @mock.patch('apache_beam.io.aws.s3filesystem.s3io') def test_create(self, unused_mock_arg): # Prepare mocks. s3io_mock = mock.MagicMock() s3filesystem.s3io.S3IO = lambda options: s3io_mock # type: ignore[misc] # Issue file copy _ = self.fs.create('s3://bucket/from1', 'application/octet-stream') s3io_mock.open.assert_called_once_with( 's3://bucket/from1', 'wb', mime_type='application/octet-stream') @mock.patch('apache_beam.io.aws.s3filesystem.s3io') def test_open(self, unused_mock_arg): # Prepare mocks. s3io_mock = mock.MagicMock() s3filesystem.s3io.S3IO = lambda options: s3io_mock # type: ignore[misc] # Issue file copy _ = self.fs.open('s3://bucket/from1', 'application/octet-stream') s3io_mock.open.assert_called_once_with( 's3://bucket/from1', 'rb', mime_type='application/octet-stream') @mock.patch('apache_beam.io.aws.s3filesystem.s3io') def test_copy_file(self, unused_mock_arg): # Prepare mocks. s3io_mock = mock.MagicMock() s3filesystem.s3io.S3IO = lambda options: s3io_mock # type: ignore[misc] sources = ['s3://bucket/from1', 's3://bucket/from2'] destinations = ['s3://bucket/to1', 's3://bucket/to2'] # Issue file copy self.fs.copy(sources, destinations) src_dest_pairs = list(zip(sources, destinations)) s3io_mock.copy_paths.assert_called_once_with(src_dest_pairs) @mock.patch('apache_beam.io.aws.s3filesystem.s3io') def test_copy_file_error(self, unused_mock_arg): # Prepare mocks. s3io_mock = mock.MagicMock() s3filesystem.s3io.S3IO = lambda options: s3io_mock # type: ignore[misc] sources = ['s3://bucket/from1', 's3://bucket/from2', 's3://bucket/from3'] destinations = ['s3://bucket/to1', 's3://bucket/to2'] # Issue file copy with self.assertRaises(BeamIOError): self.fs.copy(sources, destinations) @mock.patch('apache_beam.io.aws.s3filesystem.s3io') def test_delete(self, unused_mock_arg): # Prepare mocks. s3io_mock = mock.MagicMock() s3filesystem.s3io.S3IO = lambda options: s3io_mock # type: ignore[misc] s3io_mock.size.return_value = 0 files = [ 's3://bucket/from1', 's3://bucket/from2', 's3://bucket/from3', ] # Issue batch delete. self.fs.delete(files) s3io_mock.delete_paths.assert_called_once_with(files) @mock.patch('apache_beam.io.aws.s3filesystem.s3io') def test_delete_error(self, unused_mock_arg): # Prepare mocks. s3io_mock = mock.MagicMock() s3filesystem.s3io.S3IO = lambda options: s3io_mock # type: ignore[misc] problematic_directory = 's3://nonexistent-bucket/tree/' exception = messages.S3ClientError('Not found', 404) s3io_mock.delete_paths.return_value = { problematic_directory: exception, 's3://bucket/object1': None, 's3://bucket/object2': None, } s3io_mock.size.return_value = 0 files = [ problematic_directory, 's3://bucket/object1', 's3://bucket/object2', ] expected_results = {problematic_directory: exception} # Issue batch delete. with self.assertRaises(BeamIOError) as error: self.fs.delete(files) self.assertIn('Delete operation failed', str(error.exception)) self.assertEqual(error.exception.exception_details, expected_results) s3io_mock.delete_paths.assert_called() @mock.patch('apache_beam.io.aws.s3filesystem.s3io') def test_rename(self, unused_mock_arg): # Prepare mocks. s3io_mock = mock.MagicMock() s3filesystem.s3io.S3IO = lambda options: s3io_mock # type: ignore[misc] sources = ['s3://bucket/from1', 's3://bucket/from2'] destinations = ['s3://bucket/to1', 's3://bucket/to2'] # Issue file copy self.fs.rename(sources, destinations) src_dest_pairs = list(zip(sources, destinations)) s3io_mock.rename_files.assert_called_once_with(src_dest_pairs) def test_lineage(self): self._verify_lineage("s3://bucket/", ("bucket", )) self._verify_lineage("s3://bucket/foo/bar.txt", ("bucket", "foo/bar.txt")) def _verify_lineage(self, uri, expected_segments): lineage_mock = mock.MagicMock() self.fs.report_lineage(uri, lineage_mock) lineage_mock.add.assert_called_once_with( "s3", *expected_segments, last_segment_sep='/') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/aws/s3io.py0000644000175100017510000005456315206113032020322 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """AWS S3 client """ # pytype: skip-file import errno import io import logging import re import time import traceback from apache_beam.io.aws.clients.s3 import messages from apache_beam.io.filesystemio import Downloader from apache_beam.io.filesystemio import DownloaderStream from apache_beam.io.filesystemio import Uploader from apache_beam.io.filesystemio import UploaderStream from apache_beam.utils import retry try: # pylint: disable=wrong-import-order, wrong-import-position # pylint: disable=ungrouped-imports from apache_beam.io.aws.clients.s3 import boto3_client BOTO3_INSTALLED = True except ImportError: BOTO3_INSTALLED = False MAX_BATCH_OPERATION_SIZE = 100 def parse_s3_path(s3_path, object_optional=False): """Return the bucket and object names of the given s3:// path.""" match = re.match('^s3://([^/]+)/(.*)$', s3_path) if match is None or (match.group(2) == '' and not object_optional): raise ValueError('S3 path must be in the form s3:///.') return match.group(1), match.group(2) class S3IO(object): """S3 I/O client.""" def __init__(self, client=None, options=None): if client is None and options is None: raise ValueError('Must provide one of client or options') if client is not None: self.client = client elif BOTO3_INSTALLED: self.client = boto3_client.Client(options=options) else: message = 'AWS dependencies are not installed, and no alternative ' \ 'client was provided to S3IO.' raise RuntimeError(message) def open( self, filename, mode='r', read_buffer_size=16 * 1024 * 1024, mime_type='application/octet-stream'): """Open an S3 file path for reading or writing. Args: filename (str): S3 file path in the form ``s3:///``. mode (str): ``'r'`` for reading or ``'w'`` for writing. read_buffer_size (int): Buffer size to use during read operations. mime_type (str): Mime type to set for write operations. Returns: S3 file object. Raises: ValueError: Invalid open file mode. """ if mode == 'r' or mode == 'rb': downloader = S3Downloader( self.client, filename, buffer_size=read_buffer_size) return io.BufferedReader( DownloaderStream(downloader, mode=mode), buffer_size=read_buffer_size) elif mode == 'w' or mode == 'wb': uploader = S3Uploader(self.client, filename, mime_type) return io.BufferedWriter( UploaderStream(uploader, mode=mode), buffer_size=128 * 1024) else: raise ValueError('Invalid file open mode: %s.' % mode) def list_files(self, path, with_metadata=False): """Lists files matching the prefix. Args: path: S3 file path pattern in the form s3:///[name]. with_metadata: Experimental. Specify whether returns file metadata. Returns: If ``with_metadata`` is False: generator of tuple(file name, size); if ``with_metadata`` is True: generator of tuple(file name, tuple(size, timestamp)). """ bucket, prefix = parse_s3_path(path, object_optional=True) request = messages.ListRequest(bucket=bucket, prefix=prefix) file_info = set() counter = 0 start_time = time.time() if with_metadata: logging.debug("Starting the file information of the input") else: logging.debug("Starting the size estimation of the input") while True: #The list operation will raise an exception #when trying to list a nonexistent S3 path. #This should not be an issue here. #Ignore this exception or it will break the procedure. try: response = retry.with_exponential_backoff( retry_filter=retry.retry_on_server_errors_and_timeout_filter)( self.client.list)( request) except messages.S3ClientError as e: if e.code == 404: break else: raise e for item in response.items: file_name = 's3://%s/%s' % (bucket, item.key) if file_name not in file_info: file_info.add(file_name) counter += 1 if counter % 10000 == 0: if with_metadata: logging.info( "Finished computing file information of: %s files", len(file_info)) else: logging.info( "Finished computing size of: %s files", len(file_info)) if with_metadata: yield file_name, ( item.size, self._updated_to_seconds(item.last_modified)) else: yield file_name, item.size if response.next_token: request.continuation_token = response.next_token else: break logging.log( # do not spam logs when list_files is likely used to check empty folder logging.INFO if counter > 0 else logging.DEBUG, "Finished listing %s files in %s seconds.", counter, time.time() - start_time) return file_info def checksum(self, path): """Looks up the checksum of an S3 object. Args: path: S3 file path pattern in the form s3:///. """ return self._s3_object(path).etag @retry.with_exponential_backoff( retry_filter=retry.retry_on_server_errors_and_timeout_filter) def copy(self, src, dest): """Copies a single S3 file object from src to dest. Args: src: S3 file path pattern in the form s3:///. dest: S3 file path pattern in the form s3:///. Raises: TimeoutError: on timeout. """ src_bucket, src_key = parse_s3_path(src) dest_bucket, dest_key = parse_s3_path(dest) request = messages.CopyRequest(src_bucket, src_key, dest_bucket, dest_key) self.client.copy(request) # We intentionally do not decorate this method with a retry, since the # underlying copy and delete operations are already idempotent operations # protected by retry decorators. def copy_paths(self, src_dest_pairs): """Copies the given S3 objects from src to dest. This can handle directory or file paths. Args: src_dest_pairs: list of (src, dest) tuples of s3:/// file paths to copy from src to dest Returns: List of tuples of (src, dest, exception) in the same order as the src_dest_pairs argument, where exception is None if the operation succeeded or the relevant exception if the operation failed. """ if not src_dest_pairs: return [] results = [] for src_path, dest_path in src_dest_pairs: # Copy a directory with self.copy_tree if src_path.endswith('/') and dest_path.endswith('/'): try: results += self.copy_tree(src_path, dest_path) except messages.S3ClientError as err: results.append((src_path, dest_path, err)) # Copy individual files with self.copy elif not src_path.endswith('/') and not dest_path.endswith('/'): src_bucket, src_key = parse_s3_path(src_path) dest_bucket, dest_key = parse_s3_path(dest_path) request = messages.CopyRequest( src_bucket, src_key, dest_bucket, dest_key) try: self.client.copy(request) results.append((src_path, dest_path, None)) except messages.S3ClientError as err: results.append((src_path, dest_path, err)) # Mismatched paths (one directory, one non-directory) get an error result else: e = messages.S3ClientError( "Can't copy mismatched paths (one directory, one non-directory):" + ' %s, %s' % (src_path, dest_path), 400) results.append((src_path, dest_path, e)) return results # We intentionally do not decorate this method with a retry, since the # underlying copy and delete operations are already idempotent operations # protected by retry decorators. def copy_tree(self, src, dest): """Renames the given S3 directory and it's contents recursively from src to dest. Args: src: S3 file path pattern in the form s3:////. dest: S3 file path pattern in the form s3:////. Returns: List of tuples of (src, dest, exception) where exception is None if the operation succeeded or the relevant exception if the operation failed. """ assert src.endswith('/') assert dest.endswith('/') results = [] for entry, _ in self.list_files(src): rel_path = entry[len(src):] try: self.copy(entry, dest + rel_path) results.append((entry, dest + rel_path, None)) except messages.S3ClientError as e: results.append((entry, dest + rel_path, e)) return results @retry.with_exponential_backoff( retry_filter=retry.retry_on_server_errors_and_timeout_filter) def delete(self, path): """Deletes a single S3 file object from src to dest. Args: src: S3 file path pattern in the form s3:////. dest: S3 file path pattern in the form s3:////. Returns: List of tuples of (src, dest, exception) in the same order as the src_dest_pairs argument, where exception is None if the operation succeeded or the relevant exception if the operation failed. """ bucket, object_path = parse_s3_path(path) request = messages.DeleteRequest(bucket, object_path) try: self.client.delete(request) except messages.S3ClientError as e: if e.code == 404: return # Same behavior as GCS - don't surface a 404 error else: logging.error('HTTP error while deleting file %s: %s', path, 3) raise e # We intentionally do not decorate this method with a retry, since the # underlying copy and delete operations are already idempotent operations # protected by retry decorators. def delete_paths(self, paths): """Deletes the given S3 objects from src to dest. This can handle directory or file paths. Args: src: S3 file path pattern in the form s3:////. dest: S3 file path pattern in the form s3:////. Returns: List of tuples of (src, dest, exception) in the same order as the src_dest_pairs argument, where exception is None if the operation succeeded or the relevant exception if the operation failed. """ directories, not_directories = [], [] for path in paths: if path.endswith('/'): directories.append(path) else: not_directories.append(path) results = {} for directory in directories: dir_result = dict(self.delete_tree(directory)) results.update(dir_result) not_directory_results = dict(self.delete_files(not_directories)) results.update(not_directory_results) return results # We intentionally do not decorate this method with a retry, since the # underlying copy and delete operations are already idempotent operations # protected by retry decorators. def delete_files(self, paths, max_batch_size=1000): """Deletes the given S3 file object from src to dest. Args: paths: List of S3 file paths in the form s3:/// max_batch_size: Largest number of keys to send to the client to be deleted simultaneously Returns: List of tuples of (path, exception) in the same order as the paths argument, where exception is None if the operation succeeded or the relevant exception if the operation failed. """ if not paths: return [] # Sort paths into bucket: [keys] buckets, keys = zip(*[parse_s3_path(path) for path in paths]) grouped_keys = {bucket: [] for bucket in buckets} for bucket, key in zip(buckets, keys): grouped_keys[bucket].append(key) # For each bucket, delete minibatches of keys results = {} for bucket, keys in grouped_keys.items(): for i in range(0, len(keys), max_batch_size): minibatch_keys = keys[i:i + max_batch_size] results.update(self._delete_minibatch(bucket, minibatch_keys)) # Organize final results final_results = [(path, results[parse_s3_path(path)]) for path in paths] return final_results @retry.with_exponential_backoff( retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _delete_minibatch(self, bucket, keys): """A helper method. Boto3 allows batch deletions for files within the same bucket. Args: bucket: String bucket name keys: List of keys to be deleted in the bucket Returns: dict of the form {(bucket, key): error}, where error is None if the operation succeeded """ request = messages.DeleteBatchRequest(bucket, keys) results = {} try: response = self.client.delete_batch(request) for key in response.deleted: results[(bucket, key)] = None for key, error in zip(response.failed, response.errors): results[(bucket, key)] = error except messages.S3ClientError as e: for key in keys: results[(bucket, key)] = e return results # We intentionally do not decorate this method with a retry, since the # underlying copy and delete operations are already idempotent operations # protected by retry decorators. def delete_tree(self, root): """Deletes all objects under the given S3 directory. Args: path: S3 root path in the form s3://// (ending with a "/") Returns: List of tuples of (path, exception), where each path is an object under the given root. exception is None if the operation succeeded or the relevant exception if the operation failed. """ assert root.endswith('/') paths = [p for p, _ in self.list_files(root)] return self.delete_files(paths) def size(self, path): """Returns the size of a single S3 object. This method does not perform glob expansion. Hence the given path must be for a single S3 object. Returns: size of the S3 object in bytes. """ return self._s3_object(path).size # We intentionally do not decorate this method with a retry, since the # underlying copy and delete operations are already idempotent operations # protected by retry decorators. def rename(self, src, dest): """Renames the given S3 object from src to dest. Args: src: S3 file path pattern in the form s3:///. dest: S3 file path pattern in the form s3:///. """ self.copy(src, dest) self.delete(src) def last_updated(self, path): """Returns the last updated epoch time of a single S3 object. This method does not perform glob expansion. Hence the given path must be for a single S3 object. Returns: last updated time of the S3 object in second. """ return self._updated_to_seconds(self._s3_object(path).last_modified) def exists(self, path): """Returns whether the given S3 object exists. Args: path: S3 file path pattern in the form s3:///. """ try: self._s3_object(path) return True except messages.S3ClientError as e: if e.code == 404: # HTTP 404 indicates that the file did not exist return False else: # We re-raise all other exceptions raise def _status(self, path): """For internal use only; no backwards-compatibility guarantees. Returns supported fields (checksum, last_updated, size) of a single object as a dict at once. This method does not perform glob expansion. Hence the given path must be for a single S3 object. Returns: dict of fields of the S3 object. """ s3_object = self._s3_object(path) file_status = {} if hasattr(s3_object, 'etag'): file_status['checksum'] = s3_object.etag if hasattr(s3_object, 'last_modified'): file_status['last_updated'] = self._updated_to_seconds( s3_object.last_modified) if hasattr(s3_object, 'size'): file_status['size'] = s3_object.size return file_status @retry.with_exponential_backoff( retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _s3_object(self, path): """Returns a S3 object metadata for the given path This method does not perform glob expansion. Hence the given path must be for a single S3 object. Returns: S3 object metadata. """ bucket, object = parse_s3_path(path) request = messages.GetRequest(bucket, object) return self.client.get_object_metadata(request) @staticmethod def _updated_to_seconds(updated): """Helper function transform the updated field of response to seconds""" return ( time.mktime(updated.timetuple()) - time.timezone + updated.microsecond / 1000000.0) def rename_files(self, src_dest_pairs): """Renames the given S3 objects from src to dest. Args: src_dest_pairs: list of (src, dest) tuples of s3:/// file paths to rename from src to dest Returns: List of tuples of (src, dest, exception) in the same order as the src_dest_pairs argument, where exception is None if the operation succeeded or the relevant exception if the operation failed. """ if not src_dest_pairs: return [] # TODO: Throw value error if path has directory for src, dest in src_dest_pairs: if src.endswith('/') or dest.endswith('/'): raise ValueError('Cannot rename a directory') copy_results = self.copy_paths(src_dest_pairs) paths_to_delete = [src for (src, _, err) in copy_results if err is None] delete_results = self.delete_files(paths_to_delete) delete_results_dict = {src: err for (src, err) in delete_results} rename_results = [] for src, dest, err in copy_results: if err is not None: rename_results.append((src, dest, err)) elif delete_results_dict[src] is not None: rename_results.append((src, dest, delete_results_dict[src])) else: rename_results.append((src, dest, None)) return rename_results class S3Downloader(Downloader): def __init__(self, client, path, buffer_size): self._client = client self._path = path self._bucket, self._name = parse_s3_path(path) self._buffer_size = buffer_size # Get object state. self._get_request = ( messages.GetRequest(bucket=self._bucket, object=self._name)) try: metadata = self._get_object_metadata(self._get_request) except messages.S3ClientError as e: if e.code == 404: raise IOError(errno.ENOENT, 'Not found: %s' % self._path) else: logging.error('HTTP error while requesting file %s: %s', self._path, 3) raise self._size = metadata.size @retry.with_exponential_backoff( retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _get_object_metadata(self, get_request): return self._client.get_object_metadata(get_request) @property def size(self): return self._size def get_range(self, start, end): return self._client.get_range(self._get_request, start, end) class S3Uploader(Uploader): def __init__(self, client, path, mime_type='application/octet-stream'): self._client = client self._path = path self._bucket, self._name = parse_s3_path(path) self._mime_type = mime_type self.part_number = 1 self.buffer = b'' self.last_error = None self.upload_id = None self.parts = [] self._start_upload() # There is retry logic in the underlying transfer library but we should make # it more explicit so we can control the retry parameters. @retry.no_retries # Using no_retries marks this as an integration point. def _start_upload(self): # The uploader by default transfers data in chunks of 1024 * 1024 bytes at # a time, buffering writes until that size is reached. try: request = messages.UploadRequest( self._bucket, self._name, self._mime_type) response = self._client.create_multipart_upload(request) self.upload_id = response.upload_id except Exception as e: # pylint: disable=broad-except logging.error( 'Error in _start_upload while inserting file %s: %s', self._path, traceback.format_exc()) self.last_error = e raise e def put(self, data): MIN_WRITE_SIZE = 5 * 1024 * 1024 MAX_WRITE_SIZE = 5 * 1024 * 1024 * 1024 # TODO: Byte strings might not be the most performant way to handle this self.buffer += data.tobytes() while len(self.buffer) >= MIN_WRITE_SIZE: # Take the first chunk off the buffer and write it to S3 chunk = self.buffer[:MAX_WRITE_SIZE] self._write_to_s3(chunk) # Remove the written chunk from the buffer self.buffer = self.buffer[MAX_WRITE_SIZE:] def _write_to_s3(self, data): try: request = messages.UploadPartRequest( self._bucket, self._name, self.upload_id, self.part_number, data) response = self._client.upload_part(request) self.parts.append({ 'ETag': response.etag, 'PartNumber': response.part_number }) self.part_number = self.part_number + 1 except messages.S3ClientError as e: self.last_error = e if e.code == 404: raise IOError(errno.ENOENT, 'Not found: %s' % self._path) else: logging.error('HTTP error while requesting file %s: %s', self._path, 3) raise def finish(self): if len(self.buffer) > 0: # writes with zero length or mid size files between # MIN_WRITE_SIZE = 5 * 1024 * 1024 # MAX_WRITE_SIZE = 5 * 1024 * 1024 * 1024 # as we will reach this finish() with len(self.buffer) == 0 # which will fail self._write_to_s3(self.buffer) if self.last_error is not None: raise self.last_error # pylint: disable=raising-bad-type request = messages.CompleteMultipartUploadRequest( self._bucket, self._name, self.upload_id, self.parts) self._client.complete_multipart_upload(request) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/aws/s3io_test.py0000644000175100017510000006330015206113032021346 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for S3 client.""" # pytype: skip-file import logging import os import random import time import unittest from apache_beam.io.aws import s3io from apache_beam.io.aws.clients.s3 import fake_client from apache_beam.io.aws.clients.s3 import messages from apache_beam.options import pipeline_options class TestS3PathParser(unittest.TestCase): BAD_S3_PATHS = [ 's3://', 's3://bucket', 's3:///name', 's3:///', 's3:/blah/bucket/name', ] def test_s3_path(self): self.assertEqual(s3io.parse_s3_path('s3://bucket/name'), ('bucket', 'name')) self.assertEqual( s3io.parse_s3_path('s3://bucket/name/sub'), ('bucket', 'name/sub')) def test_bad_s3_path(self): for path in self.BAD_S3_PATHS: self.assertRaises(ValueError, s3io.parse_s3_path, path) self.assertRaises(ValueError, s3io.parse_s3_path, 's3://bucket/') def test_s3_path_object_optional(self): self.assertEqual( s3io.parse_s3_path('s3://bucket/name', object_optional=True), ('bucket', 'name')) self.assertEqual( s3io.parse_s3_path('s3://bucket/', object_optional=True), ('bucket', '')) def test_bad_s3_path_object_optional(self): for path in self.BAD_S3_PATHS: self.assertRaises(ValueError, s3io.parse_s3_path, path, True) class TestS3IO(unittest.TestCase): def _insert_random_file(self, client, path, size): bucket, name = s3io.parse_s3_path(path) contents = os.urandom(size) fakeFile = fake_client.FakeFile(bucket, name, contents) if self.USE_MOCK: self.client.add_file(fakeFile) else: f = self.aws.open(path, 'w') f.write(contents) f.close() return fakeFile def setUp(self): # These tests can be run locally against a mock S3 client, or as integration # tests against the real S3 client. self.USE_MOCK = True # If you're running integration tests with S3, set this variable to be an # s3 path that you have access to where test data can be written. If you're # just running tests against the mock, this can be any s3 path. It should # end with a '/'. self.TEST_DATA_PATH = 's3://random-data-sets/beam_tests/' if self.USE_MOCK: self.client = fake_client.FakeS3Client() test_data_bucket, _ = s3io.parse_s3_path(self.TEST_DATA_PATH) self.client.known_buckets.add(test_data_bucket) self.aws = s3io.S3IO(self.client) else: self.aws = s3io.S3IO(options=pipeline_options.S3Options()) self.client = self.aws.client def test_size(self): file_name = self.TEST_DATA_PATH + 'dummy_file' file_size = 1234 self._insert_random_file(self.client, file_name, file_size) self.assertTrue(self.aws.exists(file_name)) self.assertEqual(1234, self.aws.size(file_name)) # Clean up self.aws.delete(file_name) def test_last_updated(self): file_name = self.TEST_DATA_PATH + 'dummy_file' file_size = 1234 self._insert_random_file(self.client, file_name, file_size) self.assertTrue(self.aws.exists(file_name)) # The time difference should be tiny for the mock client. # A loose tolerance is for the consideration of real s3 client. tolerance = 5 * 60 # 5 mins result = self.aws.last_updated(file_name) self.assertAlmostEqual(result, time.time(), delta=tolerance) # Clean up self.aws.delete(file_name) def test_checksum(self): file_name = self.TEST_DATA_PATH + 'checksum' file_size = 1024 file_ = self._insert_random_file(self.client, file_name, file_size) original_etag = self.aws.checksum(file_name) self.aws.delete(file_name) with self.aws.open(file_name, 'w') as f: f.write(file_.contents) rewritten_etag = self.aws.checksum(file_name) self.assertEqual(original_etag, rewritten_etag) self.assertEqual(len(original_etag), 36) self.assertTrue(original_etag.endswith('-1"')) # Clean up self.aws.delete(file_name) def test_file_status(self): file_name = self.TEST_DATA_PATH + 'metadata' file_size = 1024 self._insert_random_file(self.client, file_name, file_size) file_checksum = self.aws.checksum(file_name) file_timestamp = self.aws.last_updated(file_name) file_status = self.aws._status(file_name) self.assertEqual(file_status['size'], file_size) self.assertEqual(file_status['checksum'], file_checksum) self.assertEqual(file_status['last_updated'], file_timestamp) # Clean up self.aws.delete(file_name) def test_copy(self): src_file_name = self.TEST_DATA_PATH + 'source' dest_file_name = self.TEST_DATA_PATH + 'dest' file_size = 1024 self._insert_random_file(self.client, src_file_name, file_size) self.assertTrue( src_file_name in dict(self.aws.list_files(self.TEST_DATA_PATH))) self.assertFalse( dest_file_name in dict(self.aws.list_files(self.TEST_DATA_PATH))) self.aws.copy(src_file_name, dest_file_name) self.assertTrue( src_file_name in dict(self.aws.list_files(self.TEST_DATA_PATH))) self.assertTrue( dest_file_name in dict(self.aws.list_files(self.TEST_DATA_PATH))) # Clean up self.aws.delete_files([src_file_name, dest_file_name]) # Test copy of non-existent files. with self.assertRaises(messages.S3ClientError) as err: self.aws.copy( self.TEST_DATA_PATH + 'non-existent', self.TEST_DATA_PATH + 'non-existent-destination') self.assertTrue('Not Found' in err.exception.message) def test_copy_paths(self): from_name_pattern = self.TEST_DATA_PATH + 'copy_me_%d' to_name_pattern = self.TEST_DATA_PATH + 'destination_%d' file_size = 1024 num_files = 10 src_dest_pairs = [(from_name_pattern % i, to_name_pattern % i) for i in range(num_files)] result = self.aws.copy_paths(src_dest_pairs) self.assertTrue(result) for i, (src, dest, exception) in enumerate(result): self.assertEqual(src, from_name_pattern % i) self.assertEqual(dest, to_name_pattern % i) self.assertTrue(isinstance(exception, messages.S3ClientError)) self.assertEqual(exception.code, 404) self.assertFalse(self.aws.exists(from_name_pattern % i)) self.assertFalse(self.aws.exists(to_name_pattern % i)) # Insert some files. for i in range(num_files): self._insert_random_file(self.client, from_name_pattern % i, file_size) # Check files inserted properly. for i in range(num_files): self.assertTrue(self.aws.exists(from_name_pattern % i)) # Execute batch copy. result = self.aws.copy_paths(src_dest_pairs) # Check files copied properly. for i in range(num_files): self.assertTrue(self.aws.exists(from_name_pattern % i)) self.assertTrue(self.aws.exists(to_name_pattern % i)) # Check results for i, (src, dest, exception) in enumerate(result): self.assertEqual(src_dest_pairs[i], (src, dest)) self.assertEqual(exception, None) # Clean up all_files = set().union(*[set(pair) for pair in src_dest_pairs]) self.aws.delete_files(all_files) def test_copy_paths_error(self): n_real_files = 3 # Create some files from_path = self.TEST_DATA_PATH + 'copy_paths/' files = [from_path + '%d' % i for i in range(n_real_files)] to_path = self.TEST_DATA_PATH + 'destination/' destinations = [to_path + '%d' % i for i in range(n_real_files)] for file_ in files: self._insert_random_file(self.client, file_, 1024) # Add nonexistent files to the sources and destinations sources = files + [ from_path + 'X', from_path + 'fake_directory_1/', from_path + 'fake_directory_2/' ] destinations += [ to_path + 'X', to_path + 'fake_directory_1/', to_path + 'fake_directory_2' ] result = self.aws.copy_paths(list(zip(sources, destinations))) # The copy_paths function of class S3IO does not return one single # result when copying a directory. Instead, it returns the results # of copying every file in the source directory. self.assertEqual(len(result), len(sources) - 1) for _, _, err in result[:n_real_files]: self.assertTrue(err is None) for _, _, err in result[n_real_files:]: self.assertIsInstance(err, messages.S3ClientError) # For the same reason of copy_paths function of S3IO above # skip this assert. #self.assertEqual(result[-3][2].code, 404) self.assertEqual(result[-2][2].code, 404) self.assertEqual(result[-1][2].code, 400) # Clean up self.aws.delete_files(files) self.aws.delete_files(destinations) def test_copy_tree(self): src_dir_name = self.TEST_DATA_PATH + 'source/' dest_dir_name = self.TEST_DATA_PATH + 'dest/' file_size = 1024 paths = ['a', 'b/c', 'b/d'] for path in paths: src_file_name = src_dir_name + path dest_file_name = dest_dir_name + path self._insert_random_file(self.client, src_file_name, file_size) self.assertTrue( src_file_name in dict(self.aws.list_files(self.TEST_DATA_PATH))) self.assertFalse( dest_file_name in dict(self.aws.list_files(self.TEST_DATA_PATH))) results = self.aws.copy_tree(src_dir_name, dest_dir_name) for src_file_name, dest_file_name, err in results: self.assertTrue(src_dir_name in src_file_name) self.assertTrue(dest_dir_name in dest_file_name) self.assertIsNone(err) self.assertTrue( src_file_name in dict(self.aws.list_files(self.TEST_DATA_PATH))) self.assertTrue( dest_file_name in dict(self.aws.list_files(self.TEST_DATA_PATH))) # Clean up for path in paths: src_file_name = src_dir_name + path dest_file_name = dest_dir_name + path self.aws.delete_files([src_file_name, dest_file_name]) def test_rename(self): src_file_name = self.TEST_DATA_PATH + 'source' dest_file_name = self.TEST_DATA_PATH + 'dest' file_size = 1024 self._insert_random_file(self.client, src_file_name, file_size) self.assertTrue( src_file_name in dict(self.aws.list_files(self.TEST_DATA_PATH))) self.assertFalse( dest_file_name in dict(self.aws.list_files(self.TEST_DATA_PATH))) self.aws.rename(src_file_name, dest_file_name) self.assertFalse( src_file_name in dict(self.aws.list_files(self.TEST_DATA_PATH))) self.assertTrue( dest_file_name in dict(self.aws.list_files(self.TEST_DATA_PATH))) # Clean up self.aws.delete_files([src_file_name, dest_file_name]) def test_rename_files(self): from_name_pattern = self.TEST_DATA_PATH + 'to_rename_%d' to_name_pattern = self.TEST_DATA_PATH + 'been_renamed_%d' file_size = 1024 num_files = 10 src_dest_pairs = [(from_name_pattern % i, to_name_pattern % i) for i in range(num_files)] result = self.aws.rename_files(src_dest_pairs) self.assertTrue(result) for i, (src, dest, exception) in enumerate(result): self.assertEqual(src, from_name_pattern % i) self.assertEqual(dest, to_name_pattern % i) self.assertTrue(isinstance(exception, messages.S3ClientError)) self.assertEqual(exception.code, 404) self.assertFalse(self.aws.exists(from_name_pattern % i)) self.assertFalse(self.aws.exists(to_name_pattern % i)) # Insert some files. for i in range(num_files): self._insert_random_file(self.client, from_name_pattern % i, file_size) # Check files inserted properly. for i in range(num_files): self.assertTrue(self.aws.exists(from_name_pattern % i)) self.assertFalse(self.aws.exists(to_name_pattern % i)) # Execute batch rename. self.aws.rename_files(src_dest_pairs) # Check files were renamed properly. for i in range(num_files): self.assertFalse(self.aws.exists(from_name_pattern % i)) self.assertTrue(self.aws.exists(to_name_pattern % i)) # Clean up all_files = set().union(*[set(pair) for pair in src_dest_pairs]) self.aws.delete_files(all_files) def test_rename_files_with_errors(self): real_prefix = self.TEST_DATA_PATH + 'rename_batch_%s' fake_prefix = 's3://fake-bucket-68ae4b0ef7b9/rename_batch_%s' src_dest_pairs = [(prefix % 'src', prefix % 'dest') for prefix in (real_prefix, fake_prefix)] # Create the file in the real bucket self._insert_random_file(self.client, real_prefix % 'src', 1024) # Execute batch rename result = self.aws.rename_files(src_dest_pairs) # First is the file in the real bucket, which shouldn't throw an error self.assertEqual(result[0][0], src_dest_pairs[0][0]) self.assertEqual(result[0][1], src_dest_pairs[0][1]) self.assertIsNone(result[0][2]) # Second is the file in the fake bucket, which should throw a 404 self.assertEqual(result[1][0], src_dest_pairs[1][0]) self.assertEqual(result[1][1], src_dest_pairs[1][1]) self.assertEqual(result[1][2].code, 404) # Clean up self.aws.delete(real_prefix % 'dest') def test_rename_files_with_errors_directory(self): # Make file dir_name = self.TEST_DATA_PATH + 'rename_dir/' file_name = dir_name + 'file' self._insert_random_file(self.client, file_name, 1024) self.assertTrue(self.aws.exists(file_name)) with self.assertRaises(ValueError): self.aws.rename_files([(file_name, self.TEST_DATA_PATH + 'dir_dest/')]) # Clean up self.aws.delete(file_name) def test_delete_paths(self): # Make files prefix = self.TEST_DATA_PATH + 'delete_paths/' file_names = [prefix + 'a', prefix + 'b/c'] for file_name in file_names: self._insert_random_file(self.client, file_name, 1024) self.assertTrue(self.aws.exists(file_names[0])) self.assertTrue(self.aws.exists(file_names[1])) # Delete paths paths = [prefix + 'a', prefix + 'b/'] self.aws.delete_paths(paths) self.assertFalse(self.aws.exists(file_names[0])) self.assertFalse(self.aws.exists(file_names[1])) def test_delete(self): file_name = self.TEST_DATA_PATH + 'delete_file' file_size = 1024 # Test deletion of non-existent file (shouldn't raise any error) self.aws.delete(file_name) # Create the file and check that it was created self._insert_random_file(self.aws.client, file_name, file_size) files = dict(self.aws.list_files(self.TEST_DATA_PATH)) self.assertTrue(file_name in files) # Delete the file and check that it was deleted self.aws.delete(file_name) self.assertFalse(self.aws.exists(file_name)) def test_delete_files(self, *unused_args): file_name_pattern = self.TEST_DATA_PATH + 'delete_batch/%d' file_size = 1024 num_files = 5 # Test deletion of non-existent files. result = self.aws.delete_files( [file_name_pattern % i for i in range(num_files)]) self.assertTrue(result) for i, (file_name, exception) in enumerate(result): self.assertEqual(file_name, file_name_pattern % i) self.assertEqual(exception, None) self.assertFalse(self.aws.exists(file_name_pattern % i)) # Insert some files. for i in range(num_files): self._insert_random_file(self.client, file_name_pattern % i, file_size) # Check files inserted properly. for i in range(num_files): self.assertTrue(self.aws.exists(file_name_pattern % i)) # Execute batch delete. self.aws.delete_files([file_name_pattern % i for i in range(num_files)]) # Check files deleted properly. for i in range(num_files): self.assertFalse(self.aws.exists(file_name_pattern % i)) def test_delete_files_with_errors(self, *unused_args): real_file = self.TEST_DATA_PATH + 'delete_batch/file' fake_file = 's3://fake-bucket-68ae4b0ef7b9/delete_batch/file' filenames = [real_file, fake_file] result = self.aws.delete_files(filenames) # First is the file in the real bucket, which shouldn't throw an error self.assertEqual(result[0][0], filenames[0]) self.assertIsNone(result[0][1]) # Second is the file in the fake bucket, which should throw a 404 self.assertEqual(result[1][0], filenames[1]) self.assertEqual(result[1][1].code, 404) def test_delete_tree(self): root_path = self.TEST_DATA_PATH + 'delete_tree/' leaf_paths = ['a', 'b/c', 'b/d', 'b/d/e'] paths = [root_path + leaf for leaf in leaf_paths] # Create file tree file_size = 1024 for path in paths: self._insert_random_file(self.client, path, file_size) # Check that the files exist for path in paths: self.assertTrue(self.aws.exists(path)) # Delete the tree self.aws.delete_tree(root_path) # Check that the files have been deleted for path in paths: self.assertFalse(self.aws.exists(path)) def test_exists(self): file_name = self.TEST_DATA_PATH + 'exists' file_size = 1024 self.assertFalse(self.aws.exists(file_name)) self._insert_random_file(self.aws.client, file_name, file_size) self.assertTrue(self.aws.exists(file_name)) # Clean up self.aws.delete(file_name) self.assertFalse(self.aws.exists(file_name)) def test_file_mode(self): file_name = self.TEST_DATA_PATH + 'jerry/pigpen/bobby' with self.aws.open(file_name, 'w') as f: assert f.mode == 'w' with self.aws.open(file_name, 'r') as f: assert f.mode == 'r' # Clean up self.aws.delete(file_name) def test_full_file_read(self): file_name = self.TEST_DATA_PATH + 'jerry/pigpen/phil' file_size = 1024 f = self._insert_random_file(self.aws.client, file_name, file_size) contents = f.contents f = self.aws.open(file_name) self.assertEqual(f.mode, 'r') f.seek(0, os.SEEK_END) self.assertEqual(f.tell(), file_size) self.assertEqual(f.read(), b'') f.seek(0) self.assertEqual(f.read(), contents) # Clean up self.aws.delete(file_name) def test_file_write(self): file_name = self.TEST_DATA_PATH + 'write_file' file_size = 8 * 1024 * 1024 + 2000 contents = os.urandom(file_size) f = self.aws.open(file_name, 'w') self.assertEqual(f.mode, 'w') f.write(contents[0:1000]) f.write(contents[1000:1024 * 1024]) f.write(contents[1024 * 1024:]) f.close() new_f = self.aws.open(file_name, 'r') new_f_contents = new_f.read() self.assertEqual(new_f_contents, contents) # Clean up self.aws.delete(file_name) def test_file_mime_type(self): if self.USE_MOCK: self.skipTest("The boto3_client mock doesn't support mime_types") mime_type = 'example/example' file_name = self.TEST_DATA_PATH + 'write_file' f = self.aws.open(file_name, 'w', mime_type=mime_type) f.write(b'a string of binary text') f.close() bucket, key = s3io.parse_s3_path(file_name) metadata = self.client.get_object_metadata(messages.GetRequest(bucket, key)) self.assertEqual(mime_type, metadata.mime_type) # Clean up self.aws.delete(file_name) def test_file_random_seek(self): file_name = self.TEST_DATA_PATH + 'write_seek_file' file_size = 5 * 1024 * 1024 - 100 contents = os.urandom(file_size) with self.aws.open(file_name, 'w') as wf: wf.write(contents) f = self.aws.open(file_name) random.seed(0) for _ in range(0, 10): a = random.randint(0, file_size - 1) b = random.randint(0, file_size - 1) start, end = min(a, b), max(a, b) f.seek(start) self.assertEqual(f.tell(), start) self.assertEqual(f.read(end - start + 1), contents[start:end + 1]) self.assertEqual(f.tell(), end + 1) # Clean up self.aws.delete(file_name) def test_file_flush(self): file_name = self.TEST_DATA_PATH + 'flush_file' file_size = 5 * 1024 * 1024 + 2000 contents = os.urandom(file_size) f = self.aws.open(file_name, 'w') self.assertEqual(f.mode, 'w') f.write(contents[0:1000]) f.flush() f.write(contents[1000:1024 * 1024]) f.flush() f.flush() # Should be a NOOP. f.write(contents[1024 * 1024:]) f.close( ) # This should al`read`y call the equivalent of flush() in its body new_f = self.aws.open(file_name, 'r') new_f_contents = new_f.read() self.assertEqual(new_f_contents, contents) # Clean up self.aws.delete(file_name) def test_file_iterator(self): file_name = self.TEST_DATA_PATH + 'iterate_file' lines = [] line_count = 10 for _ in range(line_count): line_length = random.randint(100, 500) line = os.urandom(line_length).replace(b'\n', b' ') + b'\n' lines.append(line) contents = b''.join(lines) with self.aws.open(file_name, 'w') as wf: wf.write(contents) f = self.aws.open(file_name) read_lines = 0 for line in f: read_lines += 1 self.assertEqual(read_lines, line_count) # Clean up self.aws.delete(file_name) def test_file_read_line(self): file_name = self.TEST_DATA_PATH + 'read_line_file' lines = [] # Set a small buffer size to exercise refilling the buffer. # First line is carefully crafted so the newline falls as the last character # of the buffer to exercise this code path. read_buffer_size = 1099 lines.append(b'x' * 1023 + b'\n') for _ in range(1, 1000): line_length = random.randint(100, 500) line = os.urandom(line_length).replace(b'\n', b' ') + b'\n' lines.append(line) contents = b''.join(lines) file_size = len(contents) with self.aws.open(file_name, 'wb') as wf: wf.write(contents) f = self.aws.open(file_name, 'rb', read_buffer_size=read_buffer_size) # Test read of first two lines. f.seek(0) self.assertEqual(f.readline(), lines[0]) self.assertEqual(f.tell(), len(lines[0])) self.assertEqual(f.readline(), lines[1]) # Test read at line boundary. f.seek(file_size - len(lines[-1]) - 1) self.assertEqual(f.readline(), b'\n') # Test read at end of file. f.seek(file_size) self.assertEqual(f.readline(), b'') # Test reads at random positions. random.seed(0) for _ in range(0, 10): start = random.randint(0, file_size - 1) line_index = 0 # Find line corresponding to start index. chars_left = start while True: next_line_length = len(lines[line_index]) if chars_left - next_line_length < 0: break chars_left -= next_line_length line_index += 1 f.seek(start) self.assertEqual(f.readline(), lines[line_index][chars_left:]) # Clean up self.aws.delete(file_name) def test_file_close(self): file_name = self.TEST_DATA_PATH + 'close_file' file_size = 5 * 1024 * 1024 + 2000 contents = os.urandom(file_size) f = self.aws.open(file_name, 'w') self.assertEqual(f.mode, 'w') f.write(contents) f.close() f.close() # This should not crash. with self.aws.open(file_name, 'r') as f: read_contents = f.read() self.assertEqual(read_contents, contents) # Clean up self.aws.delete(file_name) def test_context_manager(self): # Test writing with a context manager. file_name = self.TEST_DATA_PATH + 'context_manager_file' file_size = 1024 contents = os.urandom(file_size) with self.aws.open(file_name, 'w') as f: f.write(contents) with self.aws.open(file_name, 'r') as f: self.assertEqual(f.read(), contents) # Clean up self.aws.delete(file_name) def test_list_files(self): objects = [ ('jerry/pigpen/phil', 5), ('jerry/pigpen/bobby', 3), ('jerry/billy/bobby', 4), ] for (object_name, size) in objects: file_name = self.TEST_DATA_PATH + object_name self._insert_random_file(self.aws.client, file_name, size) test_cases = [ ( self.TEST_DATA_PATH + 'j', [ ('jerry/pigpen/phil', 5), ('jerry/pigpen/bobby', 3), ('jerry/billy/bobby', 4), ]), ( self.TEST_DATA_PATH + 'jerry/', [ ('jerry/pigpen/phil', 5), ('jerry/pigpen/bobby', 3), ('jerry/billy/bobby', 4), ]), ( self.TEST_DATA_PATH + 'jerry/pigpen/phil', [ ('jerry/pigpen/phil', 5), ]), ] for file_pattern, expected_object_names in test_cases: expected_file_names = [(self.TEST_DATA_PATH + object_name, size) for (object_name, size) in expected_object_names] self.assertEqual( set(dict(self.aws.list_files(file_pattern)).items()), set(expected_file_names)) # Clean up for (object_name, size) in objects: self.aws.delete(self.TEST_DATA_PATH + object_name) def test_midsize_file(self): file_name = self.TEST_DATA_PATH + 'midsized' file_size = 6 * 1024 * 1024 self._insert_random_file(self.aws.client, file_name, file_size) with self.aws.open(file_name, 'r') as f: self.assertEqual(len(f.read()), file_size) self.aws.delete(file_name) def test_zerosize_file(self): file_name = self.TEST_DATA_PATH + 'zerosized' file_size = 0 self._insert_random_file(self.aws.client, file_name, file_size) with self.aws.open(file_name, 'r') as f: self.assertEqual(len(f.read()), file_size) self.aws.delete(file_name) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3094003 apache_beam-2.74.0/apache_beam/io/azure/0000755000175100017510000000000015206113157017422 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/azure/__init__.py0000644000175100017510000000142015206113032021520 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/azure/blobstoragefilesystem.py0000644000175100017510000002533515206113032024404 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Azure Blob Storage Implementation for accesing files on Azure Blob Storage. """ import traceback from apache_beam.io.azure import blobstorageio from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystem import CompressedFile from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystem import FileMetadata from apache_beam.io.filesystem import FileSystem __all__ = ['BlobStorageFileSystem'] class BlobStorageFileSystem(FileSystem): """An Azure Blob Storage ``FileSystem`` implementation for accesing files on Azure Blob Storage. """ CHUNK_SIZE = blobstorageio.MAX_BATCH_OPERATION_SIZE AZURE_FILE_SYSTEM_PREFIX = 'azfs://' def __init__(self, pipeline_options): super().__init__(pipeline_options) self._pipeline_options = pipeline_options @classmethod def scheme(cls): """URI scheme for the FileSystem """ return 'azfs' def join(self, basepath, *paths): """Join two or more pathname components for the filesystem Args: basepath: string path of the first component of the path paths: path components to be added Returns: full path after combining all the passed components """ if not basepath.startswith(BlobStorageFileSystem.AZURE_FILE_SYSTEM_PREFIX): raise ValueError( 'Basepath %r must be an Azure Blob Storage path.' % basepath) path = basepath for p in paths: path = path.rstrip('/') + '/' + p.lstrip('/') return path def split(self, path): """Splits the given path into two parts. Splits the path into a pair (head, tail) such that tail contains the last component of the path and head contains everything up to that. For file-systems other than the local file-system, head should include the prefix. Args: path: path as a string Returns: a pair of path components as strings. """ path = path.strip() if not path.startswith(BlobStorageFileSystem.AZURE_FILE_SYSTEM_PREFIX): raise ValueError('Path %r must be Azure Blob Storage path.' % path) prefix_len = len(BlobStorageFileSystem.AZURE_FILE_SYSTEM_PREFIX) last_sep = path[prefix_len:].rfind('/') if last_sep >= 0: last_sep += prefix_len if last_sep > 0: return (path[:last_sep], path[last_sep + 1:]) elif last_sep < 0: return (path, '') else: raise ValueError('Invalid path: %s' % path) def mkdirs(self, path): """Recursively create directories for the provided path. Args: path: string path of the directory structure that should be created Raises: IOError: if leaf directory already exists. """ pass def has_dirs(self): """Whether this FileSystem supports directories.""" return False def _list(self, dir_or_prefix): """List files in a location. Listing is non-recursive (for filesystems that support directories). Args: dir_or_prefix: (string) A directory or location prefix (for filesystems that don't have directories). Returns: Generator of ``FileMetadata`` objects. Raises: ``BeamIOError``: if listing fails, but not if no files were found. """ try: for path, (size, updated) in self._blobstorageIO().list_files( dir_or_prefix, with_metadata=True): yield FileMetadata(path, size, updated) except Exception as e: # pylint: disable=broad-except raise BeamIOError("List operation failed", {dir_or_prefix: e}) def _blobstorageIO(self): return blobstorageio.BlobStorageIO(pipeline_options=self._pipeline_options) def _path_open( self, path, mode, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO): """Helper functions to open a file in the provided mode. """ compression_type = FileSystem._get_compression_type(path, compression_type) mime_type = CompressionTypes.mime_type(compression_type, mime_type) raw_file = self._blobstorageIO().open(path, mode, mime_type=mime_type) if compression_type == CompressionTypes.UNCOMPRESSED: return raw_file return CompressedFile(raw_file, compression_type=compression_type) def create( self, path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO): """Returns a write channel for the given file path. Args: path: string path of the file object to be written to the system mime_type: MIME type to specify the type of content in the file object compression_type: Type of compression to be used for this object Returns: file handle with a close function for the user to use """ return self._path_open(path, 'wb', mime_type, compression_type) def open( self, path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO): """Returns a read channel for the given file path. Args: path: string path of the file object to be read mime_type: MIME type to specify the type of content in the file object compression_type: Type of compression to be used for this object Returns: file handle with a close function for the user to use """ return self._path_open(path, 'rb', mime_type, compression_type) def copy(self, source_file_names, destination_file_names): """Recursively copy the file tree from the source to the destination Args: source_file_names: list of source file objects that needs to be copied destination_file_names: list of destination of the new object Raises: ``BeamIOError``: if any of the copy operations fail """ if not len(source_file_names) == len(destination_file_names): message = 'Unable to copy unequal number of sources and destinations.' raise BeamIOError(message) src_dest_pairs = list(zip(source_file_names, destination_file_names)) return self._blobstorageIO().copy_paths(src_dest_pairs) def rename(self, source_file_names, destination_file_names): """Rename the files at the source list to the destination list. Source and destination lists should be of the same size. Args: source_file_names: List of file paths that need to be moved destination_file_names: List of destination_file_names for the files Raises: ``BeamIOError``: if any of the rename operations fail """ if not len(source_file_names) == len(destination_file_names): message = 'Unable to rename unequal number of sources and destinations.' raise BeamIOError(message) src_dest_pairs = list(zip(source_file_names, destination_file_names)) results = self._blobstorageIO().rename_files(src_dest_pairs) # Retrieve exceptions. exceptions = {(src, dest): error for (src, dest, error) in results if error is not None} if exceptions: raise BeamIOError("Rename operation failed.", exceptions) def exists(self, path): """Check if the provided path exists on the FileSystem. Args: path: string path that needs to be checked. Returns: boolean flag indicating if path exists """ try: return self._blobstorageIO().exists(path) except Exception as e: # pylint: disable=broad-except raise BeamIOError("Exists operation failed", {path: e}) def size(self, path): """Get size in bytes of a file on the FileSystem. Args: path: string filepath of file. Returns: int size of file according to the FileSystem. Raises: ``BeamIOError``: if path doesn't exist. """ try: return self._blobstorageIO().size(path) except Exception as e: # pylint: disable=broad-except raise BeamIOError("Size operation failed", {path: e}) def last_updated(self, path): """Get UNIX Epoch time in seconds on the FileSystem. Args: path: string path of file. Returns: float UNIX Epoch time Raises: ``BeamIOError``: if path doesn't exist. """ try: return self._blobstorageIO().last_updated(path) except Exception as e: # pylint: disable=broad-except raise BeamIOError("Last updated operation failed", {path: e}) def checksum(self, path): """Fetch checksum metadata of a file on the :class:`~apache_beam.io.filesystem.FileSystem`. Args: path: string path of a file. Returns: string containing checksum Raises: ``BeamIOError``: if path isn't a file or doesn't exist. """ try: return self._blobstorageIO().checksum(path) except Exception as e: # pylint: disable=broad-except raise BeamIOError("Checksum operation failed", {path, e}) def metadata(self, path): """Fetch metadata fields of a file on the FileSystem. Args: path: string path of a file. Returns: :class:`~apache_beam.io.filesystem.FileMetadata`. Raises: ``BeamIOError``: if path isn't a file or doesn't exist. """ try: file_metadata = self._blobstorageIO()._status(path) return FileMetadata( path, file_metadata['size'], file_metadata['last_updated']) except Exception as e: # pylint: disable=broad-except raise BeamIOError("Metadata operation failed", {path: e}) def delete(self, paths): """Deletes files or directories at the provided paths. Directories will be deleted recursively. Args: paths: list of paths that give the file objects to be deleted Raises: ``BeamIOError``: if any of the delete operations fail """ results = self._blobstorageIO().delete_paths(paths) # Retrieve exceptions. exceptions = { path: error for (path, error) in results.items() if error is not None } if exceptions: raise BeamIOError("Delete operation failed", exceptions) def report_lineage(self, path, lineage): try: components = blobstorageio.parse_azfs_path( path, blob_optional=True, get_account=True) except ValueError: # report lineage is fail-safe traceback.print_exc() return if components and not components[-1]: components = components[:-1] lineage.add('abs', *components, last_segment_sep='/') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/azure/blobstoragefilesystem_test.py0000644000175100017510000003245415206113032025443 0ustar00runnerrunner# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for Azure Blob Storage File System.""" # pytype: skip-file import logging import unittest import mock from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystem import FileMetadata from apache_beam.options.pipeline_options import PipelineOptions # Protect against environments where azure library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.io.azure import blobstoragefilesystem from apache_beam.io.azure import blobstorageio except ImportError: blobstoragefilesystem = None # type: ignore[assignment] # pylint: enable=wrong-import-order, wrong-import-position @unittest.skipIf( blobstoragefilesystem is None, 'Azure dependencies are not installed') class BlobStorageFileSystemTest(unittest.TestCase): def setUp(self): pipeline_options = PipelineOptions() self.fs = blobstoragefilesystem.BlobStorageFileSystem( pipeline_options=pipeline_options) def test_scheme(self): self.assertEqual(self.fs.scheme(), 'azfs') self.assertEqual( blobstoragefilesystem.BlobStorageFileSystem.scheme(), 'azfs') def test_join(self): self.assertEqual( 'azfs://account-name/container/path/to/file', self.fs.join('azfs://account-name/container/path', 'to', 'file')) self.assertEqual( 'azfs://account-name/container/path/to/file', self.fs.join('azfs://account-name/container/path', 'to/file')) self.assertEqual( 'azfs://account-name/container/path/to/file', self.fs.join('azfs://account-name/container/path', '/to/file')) self.assertEqual( 'azfs://account-name/container/path/to/file', self.fs.join('azfs://account-name/container/path', 'to', 'file')) self.assertEqual( 'azfs://account-name/container/path/to/file', self.fs.join('azfs://account-name/container/path', 'to/file')) self.assertEqual( 'azfs://account-name/container/path/to/file', self.fs.join('azfs://account-name/container/path', '/to/file')) with self.assertRaises(ValueError): self.fs.join('account-name/container/path', '/to/file') def test_split(self): self.assertEqual(('azfs://foo/bar', 'baz'), self.fs.split('azfs://foo/bar/baz')) self.assertEqual(('azfs://foo', ''), self.fs.split('azfs://foo/')) self.assertEqual(('azfs://foo', ''), self.fs.split('azfs://foo')) with self.assertRaises(ValueError): self.fs.split('/no/azfs/prefix') @mock.patch('apache_beam.io.azure.blobstoragefilesystem.blobstorageio') def test_match_single(self, unused_mock_blobstorageio): # Prepare mocks. blobstorageio_mock = mock.MagicMock() blobstoragefilesystem.blobstorageio.BlobStorageIO = \ lambda pipeline_options: blobstorageio_mock blobstorageio_mock.exists.return_value = True blobstorageio_mock._status.return_value = { 'size': 1, 'last_updated': 99999.0 } expected_results = [ FileMetadata('azfs://storageaccount/container/file1', 1, 99999.0) ] match_result = self.fs.match(['azfs://storageaccount/container/file1'])[0] self.assertEqual(match_result.metadata_list, expected_results) blobstorageio_mock._status.assert_called_once_with( 'azfs://storageaccount/container/file1') @mock.patch('apache_beam.io.azure.blobstoragefilesystem.blobstorageio') def test_match_multiples(self, unused_mock_blobstorageio): # Prepare mocks. blobstorageio_mock = mock.MagicMock() blobstoragefilesystem.blobstorageio.BlobStorageIO = \ lambda pipeline_options: blobstorageio_mock blobstorageio_mock.list_files.return_value = iter([ ('azfs://storageaccount/container/file1', (1, 99999.0)), ('azfs://storageaccount/container/file2', (2, 88888.0)) ]) expected_results = set([ FileMetadata('azfs://storageaccount/container/file1', 1, 99999.0), FileMetadata('azfs://storageaccount/container/file2', 2, 88888.0), ]) match_result = self.fs.match(['azfs://storageaccount/container/'])[0] self.assertEqual(set(match_result.metadata_list), expected_results) blobstorageio_mock.list_files.assert_called_once_with( 'azfs://storageaccount/container/', with_metadata=True) @mock.patch('apache_beam.io.azure.blobstoragefilesystem.blobstorageio') def test_match_multiples_limit(self, unused_mock_blobstorageio): # Prepare mocks. blobstorageio_mock = mock.MagicMock() limit = 1 blobstoragefilesystem.blobstorageio.BlobStorageIO = \ lambda pipeline_options: blobstorageio_mock blobstorageio_mock.list_files.return_value = iter([ ('azfs://storageaccount/container/file1', (1, 99999.0)) ]) expected_results = set( [FileMetadata('azfs://storageaccount/container/file1', 1, 99999.0)]) match_result = self.fs.match(['azfs://storageaccount/container/'], [limit])[0] self.assertEqual(set(match_result.metadata_list), expected_results) self.assertEqual(len(match_result.metadata_list), limit) blobstorageio_mock.list_files.assert_called_once_with( 'azfs://storageaccount/container/', with_metadata=True) @mock.patch('apache_beam.io.azure.blobstoragefilesystem.blobstorageio') def test_match_multiples_error(self, unused_mock_blobstorageio): # Prepare mocks. blobstorageio_mock = mock.MagicMock() blobstoragefilesystem.blobstorageio.BlobStorageIO = \ lambda pipeline_options: blobstorageio_mock exception = IOError('Failed') blobstorageio_mock.list_files.side_effect = exception with self.assertRaisesRegex(BeamIOError, r'^Match operation failed') as error: self.fs.match(['azfs://storageaccount/container/']) self.assertRegex( str(error.exception.exception_details), r'azfs://storageaccount/container/.*%s' % exception) blobstorageio_mock.list_files.assert_called_once_with( 'azfs://storageaccount/container/', with_metadata=True) @mock.patch('apache_beam.io.azure.blobstoragefilesystem.blobstorageio') def test_match_multiple_patterns(self, unused_mock_blobstorageio): # Prepare mocks. blobstorageio_mock = mock.MagicMock() blobstoragefilesystem.blobstorageio.BlobStorageIO = \ lambda pipeline_options: blobstorageio_mock blobstorageio_mock.list_files.side_effect = [ iter([('azfs://storageaccount/container/file1', (1, 99999.0))]), iter([('azfs://storageaccount/container/file2', (2, 88888.0))]), ] expected_results = [ [FileMetadata('azfs://storageaccount/container/file1', 1, 99999.0)], [FileMetadata('azfs://storageaccount/container/file2', 2, 88888.0)] ] result = self.fs.match([ 'azfs://storageaccount/container/file1*', 'azfs://storageaccount/container/file2*' ]) self.assertEqual([mr.metadata_list for mr in result], expected_results) @mock.patch('apache_beam.io.azure.blobstoragefilesystem.blobstorageio') def test_create(self, unused_mock_blobstorageio): # Prepare mocks. blobstorageio_mock = mock.MagicMock() blobstoragefilesystem.blobstorageio.BlobStorageIO = \ lambda pipeline_options: blobstorageio_mock # Issue file copy. _ = self.fs.create( 'azfs://storageaccount/container/file1', 'application/octet-stream') blobstorageio_mock.open.assert_called_once_with( 'azfs://storageaccount/container/file1', 'wb', mime_type='application/octet-stream') @mock.patch('apache_beam.io.azure.blobstoragefilesystem.blobstorageio') def test_open(self, unused_mock_blobstorageio): # Prepare mocks. blobstorageio_mock = mock.MagicMock() blobstoragefilesystem.blobstorageio.BlobStorageIO = \ lambda pipeline_options: blobstorageio_mock # Issue file copy. _ = self.fs.open( 'azfs://storageaccount/container/file1', 'application/octet-stream') blobstorageio_mock.open.assert_called_once_with( 'azfs://storageaccount/container/file1', 'rb', mime_type='application/octet-stream') @mock.patch('apache_beam.io.azure.blobstoragefilesystem.blobstorageio') def test_copy_file(self, unused_mock_blobstorageio): # Prepare mocks. blobstorageio_mock = mock.MagicMock() blobstoragefilesystem.blobstorageio.BlobStorageIO = \ lambda pipeline_options: blobstorageio_mock sources = [ 'azfs://storageaccount/container/from1', 'azfs://storageaccount/container/from2', ] destinations = [ 'azfs://storageaccount/container/to1', 'azfs://storageaccount/container/to2', ] # Issue file copy. self.fs.copy(sources, destinations) src_dest_pairs = list(zip(sources, destinations)) blobstorageio_mock.copy_paths.assert_called_once_with(src_dest_pairs) @mock.patch('apache_beam.io.azure.blobstoragefilesystem.blobstorageio') def test_copy_file_error(self, unused_mock_blobstorageio): # Prepare mocks. blobstorageio_mock = mock.MagicMock() blobstoragefilesystem.blobstorageio.BlobStorageIO = \ lambda pipeline_options: blobstorageio_mock sources = [ 'azfs://storageaccount/container/from1', 'azfs://storageaccount/container/from2', 'azfs://storageaccount/container/from3', ] destinations = [ 'azfs://storageaccount/container/to1', 'azfs://storageaccount/container/to2', ] # Issue file copy. with self.assertRaises(BeamIOError): self.fs.copy(sources, destinations) @mock.patch('apache_beam.io.azure.blobstoragefilesystem.blobstorageio') def test_delete(self, unused_mock_blobstorageio): # Prepare mocks. blobstorageio_mock = mock.MagicMock() blobstoragefilesystem.blobstorageio.BlobStorageIO = \ lambda pipeline_options: blobstorageio_mock blobstorageio_mock.size.return_value = 0 files = [ 'azfs://storageaccount/container/from1', 'azfs://storageaccount/container/from2', 'azfs://storageaccount/container/from3', ] # Issue batch delete operation. self.fs.delete(files) blobstorageio_mock.delete_paths.assert_called_once_with(files) @mock.patch('apache_beam.io.azure.blobstoragefilesystem.blobstorageio') def test_delete_error(self, unused_mock_blobstorageio): # Prepare mocks. blobstorageio_mock = mock.MagicMock() blobstoragefilesystem.blobstorageio.BlobStorageIO = \ lambda pipeline_options: blobstorageio_mock nonexistent_directory = 'azfs://storageaccount/nonexistent-container/tree/' exception = blobstorageio.BlobStorageError('Not found', 404) blobstorageio_mock.delete_paths.return_value = { nonexistent_directory: exception, 'azfs://storageaccount/container/blob1': None, 'azfs://storageaccount/container/blob2': None, } blobstorageio_mock.size.return_value = 0 files = [ nonexistent_directory, 'azfs://storageaccount/container/blob1', 'azfs://storageaccount/container/blob2', ] expected_results = {nonexistent_directory: exception} # Issue batch delete. with self.assertRaises(BeamIOError) as error: self.fs.delete(files) self.assertIn('Delete operation failed', str(error.exception)) self.assertEqual(error.exception.exception_details, expected_results) blobstorageio_mock.delete_paths.assert_called() @mock.patch('apache_beam.io.azure.blobstoragefilesystem.blobstorageio') def test_rename(self, unused_mock_blobstorageio): # Prepare mocks. blobstorageio_mock = mock.MagicMock() blobstoragefilesystem.blobstorageio.BlobStorageIO = \ lambda pipeline_options: blobstorageio_mock sources = [ 'azfs://storageaccount/container/original_blob1', 'azfs://storageaccount/container/original_blob2', ] destinations = [ 'azfs://storageaccount/container/renamed_blob1', 'azfs://storageaccount/container/renamed_blob2', ] # Issue bath rename. self.fs.rename(sources, destinations) src_dest_pairs = list(zip(sources, destinations)) blobstorageio_mock.rename_files.assert_called_once_with(src_dest_pairs) def test_lineage(self): self._verify_lineage( "azfs://storageaccount/container/", ("storageaccount", "container")) self._verify_lineage( "azfs://storageaccount/container/foo/bar.txt", ("storageaccount", "container", "foo/bar.txt")) def _verify_lineage(self, uri, expected_segments): lineage_mock = mock.MagicMock() self.fs.report_lineage(uri, lineage_mock) lineage_mock.add.assert_called_once_with( "abs", *expected_segments, last_segment_sep='/') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/azure/blobstorageio.py0000644000175100017510000005646215206113032022634 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Azure Blob Storage client. """ # pytype: skip-file import errno import io import logging import os import re import tempfile import time from apache_beam.internal.azure import auth from apache_beam.io.filesystemio import Downloader from apache_beam.io.filesystemio import DownloaderStream from apache_beam.io.filesystemio import Uploader from apache_beam.io.filesystemio import UploaderStream from apache_beam.options.pipeline_options import AzureOptions from apache_beam.utils import retry _LOGGER = logging.getLogger(__name__) try: # pylint: disable=wrong-import-order, wrong-import-position # pylint: disable=ungrouped-imports from azure.core.exceptions import ResourceNotFoundError from azure.storage.blob import BlobServiceClient from azure.storage.blob import ContentSettings AZURE_DEPS_INSTALLED = True except ImportError: AZURE_DEPS_INSTALLED = False DEFAULT_READ_BUFFER_SIZE = 16 * 1024 * 1024 MAX_BATCH_OPERATION_SIZE = 100 def parse_azfs_path(azfs_path, blob_optional=False, get_account=False): """Return the storage account, the container and blob names of the given azfs:// path. """ match = re.match( '^azfs://([a-z0-9]{3,24})/([a-z0-9](?![a-z0-9-]*--[a-z0-9-]*)' '[a-z0-9-]{1,61}[a-z0-9])/(.*)$', azfs_path) if match is None or (match.group(3) == '' and not blob_optional): raise ValueError( 'Azure Blob Storage path must be in the form ' 'azfs:////.') result = None if get_account: result = match.group(1), match.group(2), match.group(3) else: result = match.group(2), match.group(3) return result def get_azfs_url(storage_account, container, blob=''): """Returns the url in the form of https://account.blob.core.windows.net/container/blob-name """ return 'https://' + storage_account + '.blob.core.windows.net/' + \ container + '/' + blob class Blob(): """A Blob in Azure Blob Storage.""" def __init__(self, etag, name, last_updated, size, mime_type): self.etag = etag self.name = name self.last_updated = last_updated self.size = size self.mime_type = mime_type class BlobStorageIOError(IOError, retry.PermanentException): """Blob Strorage IO error that should not be retried.""" pass class BlobStorageError(Exception): """Blob Storage client error.""" def __init__(self, message=None, code=None): self.message = message self.code = code class BlobStorageIO(object): """Azure Blob Storage I/O client.""" def __init__(self, client=None, pipeline_options=None): if client is None: azure_options = pipeline_options.view_as(AzureOptions) connect_str = azure_options.azure_connection_string or \ os.getenv('AZURE_STORAGE_CONNECTION_STRING') if connect_str: self.client = BlobServiceClient.from_connection_string( conn_str=connect_str) else: credential = auth.get_service_credentials(pipeline_options) self.client = BlobServiceClient( account_url=azure_options.blob_service_endpoint, credential=credential) else: self.client = client if not AZURE_DEPS_INSTALLED: raise RuntimeError('Azure dependencies are not installed. Unable to run.') def open( self, filename, mode='r', read_buffer_size=DEFAULT_READ_BUFFER_SIZE, mime_type='application/octet-stream'): """Open an Azure Blob Storage file path for reading or writing. Args: filename (str): Azure Blob Storage file path in the form ``azfs:////``. mode (str): ``'r'`` for reading or ``'w'`` for writing. read_buffer_size (int): Buffer size to use during read operations. mime_type (str): Mime type to set for write operations. Returns: Azure Blob Storage file object. Raises: ValueError: Invalid open file mode. """ if mode == 'r' or mode == 'rb': downloader = BlobStorageDownloader( self.client, filename, buffer_size=read_buffer_size) return io.BufferedReader( DownloaderStream( downloader, read_buffer_size=read_buffer_size, mode=mode), buffer_size=read_buffer_size) elif mode == 'w' or mode == 'wb': uploader = BlobStorageUploader(self.client, filename, mime_type) return io.BufferedWriter( UploaderStream(uploader, mode=mode), buffer_size=128 * 1024) else: raise ValueError('Invalid file open mode: %s.' % mode) @retry.with_exponential_backoff( retry_filter=retry.retry_on_beam_io_error_filter) def copy(self, src, dest): """Copies a single Azure Blob Storage blob from src to dest. Args: src: Blob Storage file path pattern in the form azfs:////[name]. dest: Blob Storage file path pattern in the form azfs:////[name]. Raises: TimeoutError: on timeout. """ src_storage_account, src_container, src_blob = parse_azfs_path( src, get_account=True) dest_container, dest_blob = parse_azfs_path(dest) source_blob = get_azfs_url(src_storage_account, src_container, src_blob) copied_blob = self.client.get_blob_client(dest_container, dest_blob) try: copied_blob.start_copy_from_url(source_blob) except ResourceNotFoundError as e: message = e.reason code = e.status_code raise BlobStorageError(message, code) # We intentionally do not decorate this method with a retry, since the # underlying copy operation is already an idempotent operation protected # by retry decorators. def copy_tree(self, src, dest): """Renames the given Azure Blob storage directory and its contents recursively from src to dest. Args: src: Blob Storage file path pattern in the form azfs:////[name]. dest: Blob Storage file path pattern in the form azfs:////[name]. Returns: List of tuples of (src, dest, exception) where exception is None if the operation succeeded or the relevant exception if the operation failed. """ assert src.endswith('/') assert dest.endswith('/') results = [] for entry, _ in self.list_files(src): rel_path = entry[len(src):] try: self.copy(entry, dest + rel_path) results.append((entry, dest + rel_path, None)) except BlobStorageError as e: results.append((entry, dest + rel_path, e)) return results # We intentionally do not decorate this method with a retry, since the # underlying copy operation is already an idempotent operation protected # by retry decorators. def copy_paths(self, src_dest_pairs): """Copies the given Azure Blob Storage blobs from src to dest. This can handle directory or file paths. Args: src_dest_pairs: List of (src, dest) tuples of azfs:////[name] file paths to copy from src to dest. Returns: List of tuples of (src, dest, exception) in the same order as the src_dest_pairs argument, where exception is None if the operation succeeded or the relevant exception if the operation failed. """ if not src_dest_pairs: return [] results = [] for src_path, dest_path in src_dest_pairs: # Case 1. They are directories. if src_path.endswith('/') and dest_path.endswith('/'): try: results += self.copy_tree(src_path, dest_path) except BlobStorageError as e: results.append((src_path, dest_path, e)) # Case 2. They are individual blobs. elif not src_path.endswith('/') and not dest_path.endswith('/'): try: self.copy(src_path, dest_path) results.append((src_path, dest_path, None)) except BlobStorageError as e: results.append((src_path, dest_path, e)) # Mismatched paths (one directory, one non-directory) get an error. else: e = BlobStorageError( "Unable to copy mismatched paths" + "(directory, non-directory): %s, %s" % (src_path, dest_path), 400) results.append((src_path, dest_path, e)) return results # We intentionally do not decorate this method with a retry, since the # underlying copy and delete operations are already idempotent operations # protected by retry decorators. def rename(self, src, dest): """Renames the given Azure Blob Storage blob from src to dest. Args: src: Blob Storage file path pattern in the form azfs:////[name]. dest: Blob Storage file path pattern in the form azfs:////[name]. """ self.copy(src, dest) self.delete(src) # We intentionally do not decorate this method with a retry, since the # underlying copy and delete operations are already idempotent operations # protected by retry decorators. def rename_files(self, src_dest_pairs): """Renames the given Azure Blob Storage blobs from src to dest. Args: src_dest_pairs: List of (src, dest) tuples of azfs:////[name] file paths to rename from src to dest. Returns: List of tuples of (src, dest, exception) in the same order as the src_dest_pairs argument, where exception is None if the operation succeeded or the relevant exception if the operation failed. """ if not src_dest_pairs: return [] for src, dest in src_dest_pairs: if src.endswith('/') or dest.endswith('/'): raise ValueError('Unable to rename a directory.') # Results from copy operation. copy_results = self.copy_paths(src_dest_pairs) paths_to_delete = \ [src for (src, _, error) in copy_results if error is None] # Results from delete operation. delete_results = self.delete_files(paths_to_delete) # Get rename file results (list of tuples). results = [] # Using a dictionary will make the operation faster. delete_results_dict = {src: error for (src, error) in delete_results} for src, dest, error in copy_results: # If there was an error in the copy operation. if error is not None: results.append((src, dest, error)) # If there was an error in the delete operation. elif delete_results_dict[src] is not None: results.append((src, dest, delete_results_dict[src])) # If there was no error in the operations. else: results.append((src, dest, None)) return results def exists(self, path): """Returns whether the given Azure Blob Storage blob exists. Args: path: Azure Blob Storage file path pattern in the form azfs:////[name]. """ try: self._blob_properties(path) return True except ResourceNotFoundError as e: if e.status_code == 404: # HTTP 404 indicates that the file did not exist. return False else: # We re-raise all other exceptions. raise def size(self, path): """Returns the size of a single Blob Storage blob. This method does not perform glob expansion. Hence the given path must be for a single Blob Storage blob. Returns: size of the Blob Storage blob in bytes. """ return self._blob_properties(path).size def last_updated(self, path): """Returns the last updated epoch time of a single Azure Blob Storage blob. This method does not perform glob expansion. Hence the given path must be for a single Azure Blob Storage blob. Returns: last updated time of the Azure Blob Storage blob in seconds. """ return self._updated_to_seconds(self._blob_properties(path).last_modified) def checksum(self, path): """Looks up the checksum of an Azure Blob Storage blob. Args: path: Azure Blob Storage file path pattern in the form azfs:////[name]. """ return self._blob_properties(path).etag def _status(self, path): """For internal use only; no backwards-compatibility guarantees. Returns supported fields (checksum, last_updated, size) of a single object as a dict at once. This method does not perform glob expansion. Hence the given path must be for a single blob property. Returns: dict of fields of the blob property. """ properties = self._blob_properties(path) file_status = {} if hasattr(properties, 'etag'): file_status['checksum'] = properties.etag if hasattr(properties, 'last_modified'): file_status['last_updated'] = self._updated_to_seconds( properties.last_modified) if hasattr(properties, 'size'): file_status['size'] = properties.size return file_status @retry.with_exponential_backoff( retry_filter=retry.retry_on_beam_io_error_filter) def _blob_properties(self, path): """Returns a blob properties object for the given path This method does not perform glob expansion. Hence the given path must be for a single blob properties object. Returns: blob properties. """ container, blob = parse_azfs_path(path) blob_to_check = self.client.get_blob_client(container, blob) try: properties = blob_to_check.get_blob_properties() except ResourceNotFoundError as e: message = e.reason code = e.status_code raise BlobStorageError(message, code) return properties @staticmethod def _updated_to_seconds(updated): """Helper function transform the updated field of response to seconds""" return ( time.mktime(updated.timetuple()) - time.timezone + updated.microsecond / 1000000.0) @retry.with_exponential_backoff( retry_filter=retry.retry_on_beam_io_error_filter) def delete(self, path): """Deletes a single blob at the given Azure Blob Storage path. Args: path: Azure Blob Storage file path pattern in the form azfs:////[name]. """ container, blob = parse_azfs_path(path) blob_to_delete = self.client.get_blob_client(container, blob) try: blob_to_delete.delete_blob() except ResourceNotFoundError as e: if e.status_code == 404: # Return success when the file doesn't exist anymore for idempotency. return else: logging.error('HTTP error while deleting file %s', path) raise e # We intentionally do not decorate this method with a retry, since the # underlying copy and delete operations are already idempotent operations # protected by retry decorators. def delete_paths(self, paths): """Deletes the given Azure Blob Storage blobs from src to dest. This can handle directory or file paths. Args: paths: list of Azure Blob Storage paths in the form azfs:////[name] that give the file blobs to be deleted. Returns: List of tuples of (src, dest, exception) in the same order as the src_dest_pairs argument, where exception is None if the operation succeeded or the relevant exception if the operation failed. """ directories, blobs = [], [] # Retrieve directories and not directories. for path in paths: if path.endswith('/'): directories.append(path) else: blobs.append(path) results = {} for directory in directories: directory_result = dict(self.delete_tree(directory)) results.update(directory_result) blobs_results = dict(self.delete_files(blobs)) results.update(blobs_results) return results # We intentionally do not decorate this method with a retry, since the # underlying copy and delete operations are already idempotent operations # protected by retry decorators. def delete_tree(self, root): """Deletes all blobs under the given Azure BlobStorage virtual directory. Args: path: Azure Blob Storage file path pattern in the form azfs:////[name] (ending with a "/"). Returns: List of tuples of (path, exception), where each path is a blob under the given root. exception is None if the operation succeeded or the relevant exception if the operation failed. """ assert root.endswith('/') # Get the blob under the root directory. paths_to_delete = [p for p, _ in self.list_files(root)] return self.delete_files(paths_to_delete) # We intentionally do not decorate this method with a retry, since the # underlying copy and delete operations are already idempotent operations # protected by retry decorators. def delete_files(self, paths): """Deletes the given Azure Blob Storage blobs from src to dest. Args: paths: list of Azure Blob Storage paths in the form azfs:////[name] that give the file blobs to be deleted. Returns: List of tuples of (src, dest, exception) in the same order as the src_dest_pairs argument, where exception is None if the operation succeeded or the relevant exception if the operation failed. """ if not paths: return [] # Group blobs into containers. containers, blobs = zip(*[parse_azfs_path(path, get_account=False) \ for path in paths]) grouped_blobs = {container: [] for container in containers} # Fill dictionary. for container, blob in zip(containers, blobs): grouped_blobs[container].append(blob) results = {} # Delete minibatches of blobs for each container. for container, blobs in grouped_blobs.items(): for i in range(0, len(blobs), MAX_BATCH_OPERATION_SIZE): blobs_to_delete = blobs[i:i + MAX_BATCH_OPERATION_SIZE] results.update(self._delete_batch(container, blobs_to_delete)) final_results = \ [(path, results[parse_azfs_path(path, get_account=False)]) \ for path in paths] return final_results @retry.with_exponential_backoff( retry_filter=retry.retry_on_beam_io_error_filter) def _delete_batch(self, container, blobs): """A helper method. Azure Blob Storage Python Client allows batch deletions for blobs within the same container. Args: container: container name. blobs: list of blobs to be deleted. Returns: Dictionary of the form {(container, blob): error}, where error is None if the operation succeeded. """ container_client = self.client.get_container_client(container) results = {} for blob in blobs: try: response = container_client.delete_blob(blob) results[(container, blob)] = response except ResourceNotFoundError as e: results[(container, blob)] = e.status_code return results def list_files(self, path, with_metadata=False): """Lists files matching the prefix. Args: path: Azure Blob Storage file path pattern in the form azfs:////[name]. with_metadata: Experimental. Specify whether returns file metadata. Returns: If ``with_metadata`` is False: generator of tuple(file name, size); if ``with_metadata`` is True: generator of tuple(file name, tuple(size, timestamp)). """ storage_account, container, blob = parse_azfs_path( path, blob_optional=True, get_account=True) file_info = set() counter = 0 start_time = time.time() if with_metadata: logging.debug("Starting the file information of the input") else: logging.debug("Starting the size estimation of the input") container_client = self.client.get_container_client(container) response = retry.with_exponential_backoff( retry_filter=retry.retry_on_beam_io_error_filter)( container_client.list_blobs)( name_starts_with=blob) for item in response: file_name = "azfs://%s/%s/%s" % (storage_account, container, item.name) if file_name not in file_info: file_info.add(file_name) counter += 1 if counter % 10000 == 0: if with_metadata: logging.info( "Finished computing file information of: %s files", len(file_info)) else: logging.info("Finished computing size of: %s files", len(file_info)) if with_metadata: yield file_name, ( item.size, self._updated_to_seconds(item.last_modified)) else: yield file_name, item.size logging.log( # do not spam logs when list_files is likely used to check empty folder logging.INFO if counter > 0 else logging.DEBUG, "Finished listing %s files in %s seconds.", counter, time.time() - start_time) class BlobStorageDownloader(Downloader): def __init__(self, client, path, buffer_size): self._client = client self._path = path self._container, self._blob = parse_azfs_path(path) self._buffer_size = buffer_size self._blob_to_download = self._client.get_blob_client( self._container, self._blob) try: properties = self._get_object_properties() except ResourceNotFoundError as http_error: if http_error.status_code == 404: raise IOError(errno.ENOENT, 'Not found: %s' % self._path) else: _LOGGER.error( 'HTTP error while requesting file %s: %s', self._path, http_error) raise self._size = properties.size @retry.with_exponential_backoff( retry_filter=retry.retry_on_beam_io_error_filter) def _get_object_properties(self): return self._blob_to_download.get_blob_properties() @property def size(self): return self._size def get_range(self, start, end): # Download_blob first parameter is offset and second is length (exclusive). blob_data = self._blob_to_download.download_blob(start, end - start) # Returns the content as bytes. return blob_data.readall() class BlobStorageUploader(Uploader): def __init__(self, client, path, mime_type='application/octet-stream'): self._client = client self._path = path self._container, self._blob = parse_azfs_path(path) self._content_settings = ContentSettings(mime_type) self._blob_to_upload = self._client.get_blob_client( self._container, self._blob) # Temporary file. self._temporary_file = tempfile.NamedTemporaryFile() def put(self, data): self._temporary_file.write(data.tobytes()) def finish(self): self._temporary_file.seek(0) # The temporary file is deleted immediately after the operation. with open(self._temporary_file.name, "rb") as f: self._blob_to_upload.upload_blob( f.read(), overwrite=True, content_settings=self._content_settings) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/azure/blobstorageio_test.py0000644000175100017510000000621315206113032023660 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for Azure Blob Storage client. """ # pytype: skip-file import logging import unittest # Protect against environments where azure library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.io.azure import blobstorageio except ImportError: blobstorageio = None # type: ignore[assignment] # pylint: enable=wrong-import-order, wrong-import-position @unittest.skipIf(blobstorageio is None, 'Azure dependencies are not installed') class TestAZFSPathParser(unittest.TestCase): BAD_AZFS_PATHS = [ 'azfs://' 'azfs://storage-account/' 'azfs://storage-account/**' 'azfs://storage-account/**/*' 'azfs://container' 'azfs:///name' 'azfs:///' 'azfs:/blah/container/name' 'azfs://ab/container/name' 'azfs://accountwithmorethan24chars/container/name' 'azfs://***/container/name' 'azfs://storageaccount/my--container/name' 'azfs://storageaccount/CONTAINER/name' 'azfs://storageaccount/ct/name' ] def test_azfs_path(self): self.assertEqual( blobstorageio.parse_azfs_path( 'azfs://storageaccount/container/name', get_account=True), ('storageaccount', 'container', 'name')) self.assertEqual( blobstorageio.parse_azfs_path( 'azfs://storageaccount/container/name/sub', get_account=True), ('storageaccount', 'container', 'name/sub')) def test_bad_azfs_path(self): for path in self.BAD_AZFS_PATHS: self.assertRaises(ValueError, blobstorageio.parse_azfs_path, path) self.assertRaises( ValueError, blobstorageio.parse_azfs_path, 'azfs://storageaccount/container/') def test_azfs_path_blob_optional(self): self.assertEqual( blobstorageio.parse_azfs_path( 'azfs://storageaccount/container/name', blob_optional=True, get_account=True), ('storageaccount', 'container', 'name')) self.assertEqual( blobstorageio.parse_azfs_path( 'azfs://storageaccount/container/', blob_optional=True, get_account=True), ('storageaccount', 'container', '')) def test_bad_azfs_path_blob_optional(self): for path in self.BAD_AZFS_PATHS: self.assertRaises(ValueError, blobstorageio.parse_azfs_path, path, True) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3114007 apache_beam-2.74.0/apache_beam/io/components/0000755000175100017510000000000015206113157020461 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/components/__init__.py0000644000175100017510000000142015206113032022557 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/components/adaptive_throttler.py0000644000175100017510000001612115206113032024730 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # Utility functions & classes that are _not_ specific to the datastore client. # # For internal use only; no backwards-compatibility guarantees. # pytype: skip-file import logging import random import time from apache_beam.io.components import util from apache_beam.metrics.metric import Metrics _SECONDS_TO_MILLISECONDS = 1_000 class ThrottlingSignaler(object): """A class that handles signaling throttling of remote requests to the SDK harness. """ def __init__(self, namespace: str = ""): self.throttling_metric = Metrics.counter( namespace, "cumulativeThrottlingSeconds") def signal_throttled(self, seconds: int): """Signals to the runner that requests have been throttled for some amount of time. Args: seconds: int, duration of throttling in seconds. """ self.throttling_metric.inc(seconds) class AdaptiveThrottler(object): """Implements adaptive throttling. See https://landing.google.com/sre/book/chapters/handling-overload.html#client-side-throttling-a7sYUg for a full discussion of the use case and algorithm applied. """ # The target minimum number of requests per samplePeriodMs, even if no # requests succeed. Must be greater than 0, else we could throttle to zero. # Because every decision is probabilistic, there is no guarantee that the # request rate in any given interval will not be zero. (This is the +1 from # the formula in # https://landing.google.com/sre/book/chapters/handling-overload.html ) MIN_REQUESTS = 1 def __init__(self, window_ms, bucket_ms, overload_ratio): """Initializes AdaptiveThrottler. Args: window_ms: int, length of history to consider, in ms, to set throttling. bucket_ms: int, granularity of time buckets that we store data in, in ms. overload_ratio: float, the target ratio between requests sent and successful requests. This is "K" in the formula in https://landing.google.com/sre/book/chapters/handling-overload.html. """ self._all_requests = util.MovingSum(window_ms, bucket_ms) self._successful_requests = util.MovingSum(window_ms, bucket_ms) self._overload_ratio = float(overload_ratio) self._random = random.Random() def _throttling_probability(self, now): if not self._all_requests.has_data(now): return 0 all_requests = self._all_requests.sum(now) successful_requests = self._successful_requests.sum(now) return max( 0, (all_requests - self._overload_ratio * successful_requests) / (all_requests + AdaptiveThrottler.MIN_REQUESTS)) def throttle_request(self, now): """Determines whether one RPC attempt should be throttled. This should be called once each time the caller intends to send an RPC; if it returns true, drop or delay that request (calling this function again after the delay). Args: now: int, time in ms since the epoch Returns: bool, True if the caller should throttle or delay the request. """ throttling_probability = self._throttling_probability(now) self._all_requests.add(now, 1) return self._random.uniform(0, 1) < throttling_probability def successful_request(self, now): """Notifies the throttler of a successful request. Must be called once for each request (for which throttle_request was previously called) that succeeded. Args: now: int, time in ms since the epoch """ self._successful_requests.add(now, 1) class ReactiveThrottler(AdaptiveThrottler): """ A wrapper around the AdaptiveThrottler that also handles logging and signaling throttling to the SDK harness using the provided namespace. For usage, instantiate one instance of a ReactiveThrottler class for a PTransform. When making remote calls to a service, preface that call with the throttle() method to potentially pre-emptively throttle the request. This will throttle future calls based on the failure rate of preceding calls, with higher failure rates leading to longer periods of throttling to allow system recovery. capture the timestamp of the attempted request, then execute the request code. On a success, call successful_request(timestamp) to report the success to the throttler. This flow looks like the following: def remote_call(): throttler.throttle() try: timestamp = time.time() result = make_request() throttler.successful_request(timestamp) return result except Exception as e: # do any error handling you want to do raise """ def __init__( self, window_ms: int, bucket_ms: int, overload_ratio: float, namespace: str = '', throttle_delay_secs: int = 5): """Initializes the ReactiveThrottler. Args: window_ms: int, length of history to consider, in ms, to set throttling. bucket_ms: int, granularity of time buckets that we store data in, in ms. overload_ratio: float, the target ratio between requests sent and successful requests. This is "K" in the formula in https://landing.google.com/sre/book/chapters/handling-overload.html. namespace: str, the namespace to use for logging and signaling throttling is occurring throttle_delay_secs: int, the amount of time in seconds to wait after preemptively throttled requests """ self.throttling_signaler = ThrottlingSignaler(namespace=namespace) self.logger = logging.getLogger(namespace) self.throttle_delay_secs = throttle_delay_secs super().__init__( window_ms=window_ms, bucket_ms=bucket_ms, overload_ratio=overload_ratio) def throttle(self): """ Stops request code from advancing while the underlying AdaptiveThrottler is signaling to preemptively throttle the request. Automatically handles logging the throttling and signaling to the SDK harness that the request is being throttled. This should be called in any context where a call to a remote service is being contacted prior to the call being performed. """ while self.throttle_request(time.time() * _SECONDS_TO_MILLISECONDS): self.logger.info( "Delaying request for %d seconds due to previous failures", self.throttle_delay_secs) time.sleep(self.throttle_delay_secs) self.throttling_signaler.signal_throttled(self.throttle_delay_secs) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/components/adaptive_throttler_test.py0000644000175100017510000000762115206113032025774 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest from mock import patch from apache_beam.io.components.adaptive_throttler import AdaptiveThrottler class AdaptiveThrottlerTest(unittest.TestCase): START_TIME = 1500000000000 SAMPLE_PERIOD = 60000 BUCKET = 1000 OVERLOAD_RATIO = 2 def setUp(self): self._throttler = AdaptiveThrottler( AdaptiveThrottlerTest.SAMPLE_PERIOD, AdaptiveThrottlerTest.BUCKET, AdaptiveThrottlerTest.OVERLOAD_RATIO) # As far as practical, keep these tests aligned with # AdaptiveThrottlerTest.java. def test_no_initial_throttling(self): self.assertEqual( 0, self._throttler._throttling_probability( AdaptiveThrottlerTest.START_TIME)) def test_no_throttling_if_no_errors(self): for t in range(AdaptiveThrottlerTest.START_TIME, AdaptiveThrottlerTest.START_TIME + 20): self.assertFalse(self._throttler.throttle_request(t)) self._throttler.successful_request(t) self.assertEqual( 0, self._throttler._throttling_probability( AdaptiveThrottlerTest.START_TIME + 20)) def test_no_throttling_after_errors_expire(self): for t in range(AdaptiveThrottlerTest.START_TIME, AdaptiveThrottlerTest.START_TIME + AdaptiveThrottlerTest.SAMPLE_PERIOD, 100): self._throttler.throttle_request(t) # And no sucessful_request self.assertLess( 0, self._throttler._throttling_probability( AdaptiveThrottlerTest.START_TIME + AdaptiveThrottlerTest.SAMPLE_PERIOD)) for t in range(AdaptiveThrottlerTest.START_TIME + AdaptiveThrottlerTest.SAMPLE_PERIOD, AdaptiveThrottlerTest.START_TIME + AdaptiveThrottlerTest.SAMPLE_PERIOD * 2, 100): self._throttler.throttle_request(t) self._throttler.successful_request(t) self.assertEqual( 0, self._throttler._throttling_probability( AdaptiveThrottlerTest.START_TIME + AdaptiveThrottlerTest.SAMPLE_PERIOD * 2)) @patch('random.Random') def test_throttling_after_errors(self, mock_random): mock_random().uniform.side_effect = [x / 10.0 for x in range(0, 10)] * 2 self._throttler = AdaptiveThrottler( AdaptiveThrottlerTest.SAMPLE_PERIOD, AdaptiveThrottlerTest.BUCKET, AdaptiveThrottlerTest.OVERLOAD_RATIO) for t in range(AdaptiveThrottlerTest.START_TIME, AdaptiveThrottlerTest.START_TIME + 20): throttled = self._throttler.throttle_request(t) # 1/3rd of requests succeeding. if t % 3 == 1: self._throttler.successful_request(t) if t > AdaptiveThrottlerTest.START_TIME + 10: # Roughly 1/3rd succeeding, 1/3rd failing, 1/3rd throttled. self.assertAlmostEqual( 0.33, self._throttler._throttling_probability(t), delta=0.1) # Given the mocked random numbers, we expect 10..13 to be throttled and # 14+ to be unthrottled. self.assertEqual(t < AdaptiveThrottlerTest.START_TIME + 14, throttled) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/components/rate_limiter.py0000644000175100017510000002170315206113032023506 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Rate Limiter classes for controlling access to external resources. """ import abc import logging import math import random import threading import time import grpc from envoy_data_plane.envoy.extensions.common.ratelimit.v3 import RateLimitDescriptor from envoy_data_plane.envoy.extensions.common.ratelimit.v3 import RateLimitDescriptorEntry from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitRequest from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponse from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponseCode from apache_beam.io.components import adaptive_throttler from apache_beam.metrics import Metrics _LOGGER = logging.getLogger(__name__) _RPC_MAX_RETRIES = 5 _RPC_RETRY_DELAY_SECONDS = 10 class RateLimiter(abc.ABC): """Abstract base class for RateLimiters.""" def __init__(self, namespace: str = ""): # Metrics collected from the RateLimiter # Metric updates are thread safe self.throttling_signaler = adaptive_throttler.ThrottlingSignaler( namespace=namespace) self.requests_counter = Metrics.counter(namespace, 'RatelimitRequestsTotal') self.requests_allowed = Metrics.counter( namespace, 'RatelimitRequestsAllowed') self.requests_throttled = Metrics.counter( namespace, 'RatelimitRequestsThrottled') self.rpc_errors = Metrics.counter(namespace, 'RatelimitRpcErrors') self.rpc_retries = Metrics.counter(namespace, 'RatelimitRpcRetries') self.rpc_latency = Metrics.distribution(namespace, 'RatelimitRpcLatencyMs') @abc.abstractmethod def allow(self, **kwargs) -> bool: """Applies rate limiting to the request. This method checks if the request is permitted by the rate limiting policy. Depending on the implementation and configuration, it may block (sleep) until the request is allowed, or return false if the rate limit retry is exceeded. Args: **kwargs: Keyword arguments specific to the RateLimiter implementation. Returns: bool: True if the request is allowed, False if retries exceeded. Raises: Exception: If an underlying infrastructure error occurs (e.g. RPC failure). """ pass class EnvoyRateLimiter(RateLimiter): """Rate limiter implementation that uses an external Envoy Rate Limit Service. This limiter connects to a gRPC Envoy Rate Limit Service (RLS) to determine whether a request should be allowed. It supports defining a domain and a list of descriptors that correspond to the rate limit configuration in the RLS. """ def __init__( self, service_address: str, domain: str, descriptors: list[dict[str, str]], timeout: float = 5.0, block_until_allowed: bool = True, retries: int = 3, namespace: str = ''): """ Args: service_address: Address of the Envoy RLS (e.g., 'localhost:8081'). domain: The rate limit domain. descriptors: List of descriptors (key-value pairs). retries: Number of retries to attempt if rate limited, respected only if block_until_allowed is False. timeout: gRPC timeout in seconds. block_until_allowed: If enabled blocks until RateLimiter gets the token. namespace: the namespace to use for logging and signaling throttling is occurring. """ super().__init__(namespace=namespace) self.service_address = service_address self.domain = domain self.descriptors = descriptors self.retries = retries self.timeout = timeout self.block_until_allowed = block_until_allowed self._stub = None self._lock = threading.Lock() class RateLimitServiceStub(object): """ Wrapper for gRPC stub to be compatible with envoy_data_plane messages. The envoy-data-plane package uses 'betterproto' which generates async stubs for 'grpclib'. As Beam uses standard synchronous 'grpcio', RateLimitServiceStub is a bridge class to use the betterproto Message types (RateLimitRequest) with a standard grpcio Channel. """ def __init__(self, channel): self.ShouldRateLimit = channel.unary_unary( '/envoy.service.ratelimit.v3.RateLimitService/ShouldRateLimit', request_serializer=RateLimitRequest.SerializeToString, response_deserializer=RateLimitResponse.FromString, ) def init_connection(self): if self._stub is None: # Acquire lock to safegaurd againest multiple DoFn threads sharing the # same RateLimiter instance, which is the case when using Shared(). with self._lock: if self._stub is None: channel = grpc.insecure_channel(self.service_address) self._stub = EnvoyRateLimiter.RateLimitServiceStub(channel) def allow(self, hits_added: int = 1) -> bool: """Calls the Envoy RLS to apply rate limits. Sends a rate limit request to the configured Envoy Rate Limit Service. If 'block_until_allowed' is True, this method will sleep and retry if the limit is exceeded, effectively blocking until the request is permitted. If 'block_until_allowed' is False, it will return False after the retry limit is exceeded. Args: hits_added: Number of hits to add to the rate limit. Returns: bool: True if the request is allowed, False if retries exceeded. """ self.init_connection() # execute thread-safe gRPC call # Convert descriptors to proto format proto_descriptors = [] for d in self.descriptors: entries = [] for k, v in d.items(): entries.append(RateLimitDescriptorEntry(key=k, value=v)) proto_descriptors.append(RateLimitDescriptor(entries=entries)) request = RateLimitRequest( domain=self.domain, descriptors=proto_descriptors, hits_addend=hits_added) self.requests_counter.inc() attempt = 0 throttled = False while True: if not self.block_until_allowed and attempt > self.retries: break # retry loop for retry_attempt in range(_RPC_MAX_RETRIES): try: start_time = time.time() response = self._stub.ShouldRateLimit(request, timeout=self.timeout) self.rpc_latency.update(int((time.time() - start_time) * 1000)) break except grpc.RpcError as e: if retry_attempt == _RPC_MAX_RETRIES - 1: _LOGGER.error( "[EnvoyRateLimiter] ratelimit service call failed: %s", e) self.rpc_errors.inc() raise e self.rpc_retries.inc() _LOGGER.warning( "[EnvoyRateLimiter] ratelimit service call failed, retrying: %s", e) time.sleep(_RPC_RETRY_DELAY_SECONDS) if response.overall_code == RateLimitResponseCode.OK: self.requests_allowed.inc() throttled = True break elif response.overall_code == RateLimitResponseCode.OVER_LIMIT: self.requests_throttled.inc() # Ratelimit exceeded, sleep for duration until reset and retry # multiple rules can be set in the RLS config, so we need to find the # max duration sleep_s = 0.0 if response.statuses: for status in response.statuses: if status.code == RateLimitResponseCode.OVER_LIMIT: dur = status.duration_until_reset # duration_until_reset is converted to timedelta by betterproto val = dur.total_seconds() if val > sleep_s: sleep_s = val # Add 1% additive jitter to prevent thundering herd jitter = random.uniform(0, 0.01 * sleep_s) sleep_s += jitter _LOGGER.warning("[EnvoyRateLimiter] Throttled for %s seconds", sleep_s) # signal throttled time to backend self.throttling_signaler.signal_throttled(math.ceil(sleep_s)) time.sleep(sleep_s) attempt += 1 else: _LOGGER.error( "[EnvoyRateLimiter] Unknown code from RLS: %s", response.overall_code) break return throttled def __getstate__(self): state = self.__dict__.copy() if '_lock' in state: del state['_lock'] if '_stub' in state: del state['_stub'] return state def __setstate__(self, state): self.__dict__.update(state) self._lock = threading.Lock() self._stub = None ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/components/rate_limiter_test.py0000644000175100017510000001172515206113032024550 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest from datetime import timedelta from unittest import mock import grpc from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponse from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponseCode from envoy_data_plane.envoy.service.ratelimit.v3 import RateLimitResponseDescriptorStatus from apache_beam.io.components import rate_limiter class EnvoyRateLimiterTest(unittest.TestCase): def setUp(self): self.service_address = 'localhost:8081' self.domain = 'test_domain' self.descriptors = [{'key': 'value'}] self.limiter = rate_limiter.EnvoyRateLimiter( self.service_address, self.domain, self.descriptors, timeout=0.1, # Fast timeout for tests block_until_allowed=False, retries=2, namespace='test_namespace') @mock.patch('grpc.insecure_channel') def test_allow_success(self, mock_channel): # Mock successful OK response mock_stub = mock.Mock() mock_response = RateLimitResponse(overall_code=RateLimitResponseCode.OK) mock_stub.ShouldRateLimit.return_value = mock_response # Inject mock stub self.limiter._stub = mock_stub allowed = self.limiter.allow() self.assertTrue(allowed) mock_stub.ShouldRateLimit.assert_called_once() @mock.patch('grpc.insecure_channel') def test_allow_over_limit_retries_exceeded(self, mock_channel): # Mock OVER_LIMIT response mock_stub = mock.Mock() mock_response = RateLimitResponse( overall_code=RateLimitResponseCode.OVER_LIMIT) mock_stub.ShouldRateLimit.return_value = mock_response self.limiter._stub = mock_stub # block_until_allowed is False, so it should eventually return False # We mock time.sleep to run fast with mock.patch('time.sleep'): allowed = self.limiter.allow() self.assertFalse(allowed) # Should be called 1 (initial) + 2 (retries) + 1 (last check > retries # logic depends on loop) # Logic: attempt starts at 0. # Loop 1: attempt 0. status OVER_LIMIT. sleep. attempt becomes 1. # Loop 2: attempt 1. status OVER_LIMIT. sleep. attempt becomes 2. # Loop 3: attempt 2. status OVER_LIMIT. sleep. attempt becomes 3. # Loop 4: attempt 3 > retries(2). Break. # Total calls: 3 self.assertEqual(mock_stub.ShouldRateLimit.call_count, 3) @mock.patch('grpc.insecure_channel') def test_allow_rpc_error_retry(self, mock_channel): # Mock RpcError then Success mock_stub = mock.Mock() mock_response = RateLimitResponse(overall_code=RateLimitResponseCode.OK) # Side effect: Error, Error, Success error = grpc.RpcError() mock_stub.ShouldRateLimit.side_effect = [error, error, mock_response] self.limiter._stub = mock_stub with mock.patch('time.sleep'): allowed = self.limiter.allow() self.assertTrue(allowed) self.assertEqual(mock_stub.ShouldRateLimit.call_count, 3) @mock.patch('grpc.insecure_channel') def test_allow_rpc_error_fail(self, mock_channel): # Mock Persistent RpcError mock_stub = mock.Mock() error = grpc.RpcError() mock_stub.ShouldRateLimit.side_effect = error self.limiter._stub = mock_stub with mock.patch('time.sleep'): with self.assertRaises(grpc.RpcError): self.limiter.allow() # The inner loop tries 5 times for connection errors self.assertEqual(mock_stub.ShouldRateLimit.call_count, 5) @mock.patch('grpc.insecure_channel') @mock.patch('random.uniform', return_value=0.0) def test_extract_duration_from_response(self, mock_random, mock_channel): # Mock OVER_LIMIT with specific duration mock_stub = mock.Mock() # Valid until 5 seconds status = RateLimitResponseDescriptorStatus( code=RateLimitResponseCode.OVER_LIMIT, duration_until_reset=timedelta(seconds=5)) mock_response = RateLimitResponse( overall_code=RateLimitResponseCode.OVER_LIMIT, statuses=[status]) mock_stub.ShouldRateLimit.return_value = mock_response self.limiter._stub = mock_stub self.limiter.retries = 0 # Single attempt with mock.patch('time.sleep') as mock_sleep: self.limiter.allow() # Should sleep for 5 seconds (jitter is 0.0) mock_sleep.assert_called_with(5.0) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/components/util.py0000644000175100017510000000714415206113032022006 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # Utility functions & classes that are _not_ specific to the datastore client. # # For internal use only; no backwards-compatibility guarantees. # pytype: skip-file import math class MovingSum(object): """Class that keeps track of a rolling window sum. For use in tracking recent performance of the connector. Intended to be similar to org.apache.beam.sdk.util.MovingFunction(..., Sum.ofLongs()), but for convenience we expose the count of entries as well so this doubles as a moving average tracker. """ def __init__(self, window_ms, bucket_ms): if window_ms < bucket_ms or bucket_ms <= 0: raise ValueError("window_ms >= bucket_ms > 0 please") self._num_buckets = int(math.ceil(window_ms / bucket_ms)) self._bucket_ms = bucket_ms self._Reset(now=0) # initialize the moving window members def _Reset(self, now): self._current_index = 0 # pointer into self._buckets self._current_ms_since_epoch = math.floor( now / self._bucket_ms) * self._bucket_ms # _buckets is a list where each element is a list [sum, num_samples] # This is a circular buffer where # [_current_index] represents the time range # [_current_ms_since_epoch, _current_ms_since_epoch+_bucket_ms) # [_current_index-1] represents immediatly prior time range # [_current_ms_since_epoch-_bucket_ms, _current_ms_since_epoch) # etc, wrapping around from the start to the end of the array, so # [_current_index+1] is the element representing the oldest bucket. self._buckets = [[0, 0] for _ in range(0, self._num_buckets)] def _Flush(self, now): """ Args: now: int, milliseconds since epoch """ if now >= (self._current_ms_since_epoch + self._bucket_ms * self._num_buckets): # Time moved forward so far that all currently held data is outside of # the window. It is faster to simply reset our data. self._Reset(now) return while now > self._current_ms_since_epoch + self._bucket_ms: # Advance time by one _bucket_ms, setting the new bucket's counts to 0. self._current_ms_since_epoch += self._bucket_ms self._current_index = (self._current_index + 1) % self._num_buckets self._buckets[self._current_index] = [0, 0] # Intentional dead reckoning here; we don't care about staying precisely # aligned with multiples of _bucket_ms since the epoch, we just need our # buckets to represent the most recent _window_ms time window. def sum(self, now): self._Flush(now) return sum(bucket[0] for bucket in self._buckets) def add(self, now, inc): self._Flush(now) bucket = self._buckets[self._current_index] bucket[0] += inc bucket[1] += 1 def count(self, now): self._Flush(now) return sum(bucket[1] for bucket in self._buckets) def has_data(self, now): return self.count(now) > 0 ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/concat_source.py0000644000175100017510000002430715206113032021473 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """For internal use only; no backwards-compatibility guarantees. Concat Source, which reads the union of several other sources. """ # pytype: skip-file import bisect import threading from apache_beam.io import iobase class ConcatSource(iobase.BoundedSource): """For internal use only; no backwards-compatibility guarantees. A ``BoundedSource`` that can group a set of ``BoundedSources``. Primarily for internal use, use the ``apache_beam.Flatten`` transform to create the union of several reads. """ def __init__(self, sources): self._source_bundles = [ source if isinstance(source, iobase.SourceBundle) else iobase.SourceBundle(None, source, None, None) for source in sources ] @property def sources(self): return [s.source for s in self._source_bundles] def estimate_size(self): return sum(s.source.estimate_size() for s in self._source_bundles) def split( self, desired_bundle_size=None, start_position=None, stop_position=None): if start_position or stop_position: raise ValueError( 'Multi-level initial splitting is not supported. Expected start and ' 'stop positions to be None. Received %r and %r respectively.' % (start_position, stop_position)) for source in self._source_bundles: # We assume all sub-sources to produce bundles that specify weight using # the same unit. For example, all sub-sources may specify the size in # bytes as their weight. for bundle in source.source.split(desired_bundle_size, source.start_position, source.stop_position): yield bundle def get_range_tracker(self, start_position=None, stop_position=None): if start_position is None: start_position = (0, None) if stop_position is None: stop_position = (len(self._source_bundles), None) return ConcatRangeTracker( start_position, stop_position, self._source_bundles) def read(self, range_tracker): start_source, _ = range_tracker.start_position() stop_source, stop_pos = range_tracker.stop_position() if stop_pos is not None: stop_source += 1 for source_ix in range(start_source, stop_source): if not range_tracker.try_claim((source_ix, None)): break for record in self._source_bundles[source_ix].source.read( range_tracker.sub_range_tracker(source_ix)): yield record def default_output_coder(self): if self._source_bundles: # Getting coder from the first sub-sources. This assumes all sub-sources # to produce the same coder. return self._source_bundles[0].source.default_output_coder() else: return super().default_output_coder() class ConcatRangeTracker(iobase.RangeTracker): """For internal use only; no backwards-compatibility guarantees. Range tracker for ConcatSource""" def __init__(self, start, end, source_bundles): """Initializes ``ConcatRangeTracker`` Args: start: start position, a tuple of (source_index, source_position) end: end position, a tuple of (source_index, source_position) source_bundles: the list of source bundles in the ConcatSource """ super().__init__() self._start = start self._end = end self._source_bundles = source_bundles self._lock = threading.RLock() # Lazily-initialized list of RangeTrackers corresponding to each source. self._range_trackers = [None] * len(source_bundles) # The currently-being-iterated-over (and latest claimed) source. self._claimed_source_ix = self._start[0] # Now compute cumulative progress through the sources for converting # between global fractions and fractions within specific sources. # TODO(robertwb): Implement fraction-at-position to properly scale # partial start and end sources. # Note, however, that in practice splits are typically on source # boundaries anyways. last = end[0] if end[1] is None else end[0] + 1 self._cumulative_weights = ( [0] * start[0] + self._compute_cumulative_weights(source_bundles[start[0]:last]) + [1] * (len(source_bundles) - last - start[0])) @staticmethod def _compute_cumulative_weights(source_bundles): # Two adjacent sources must differ so that they can be uniquely # identified by a single global fraction. Let min_diff be the # smallest allowable difference between sources. min_diff = 1e-5 # For the computation below, we need weights for all sources. # Substitute average weights for those whose weights are # unspecified (or 1.0 for everything if none are known). known = [s.weight for s in source_bundles if s.weight is not None] avg = sum(known) / len(known) if known else 1.0 weights = [s.weight or avg for s in source_bundles] # Now compute running totals of the percent done upon reaching # each source, with respect to the start and end positions. # E.g. if the weights were [100, 20, 3] we would produce # [0.0, 100/123, 120/123, 1.0] total = float(sum(weights)) running_total = [0] for w in weights: running_total.append(max(min_diff, min(1, running_total[-1] + w / total))) running_total[-1] = 1 # In case of rounding error. # There are issues if, due to rouding error or greatly differing sizes, # two adjacent running total weights are equal. Normalize this things so # that this never happens. for k in range(1, len(running_total)): if running_total[k] == running_total[k - 1]: for j in range(k): running_total[j] *= (1 - min_diff) return running_total def start_position(self): return self._start def stop_position(self): return self._end def try_claim(self, pos): source_ix, source_pos = pos with self._lock: if source_ix > self._end[0]: return False elif source_ix == self._end[0] and self._end[1] is None: return False else: assert source_ix >= self._claimed_source_ix self._claimed_source_ix = source_ix if source_pos is None: return True else: return self.sub_range_tracker(source_ix).try_claim(source_pos) def try_split(self, pos): source_ix, source_pos = pos with self._lock: if source_ix < self._claimed_source_ix: # Already claimed. return None elif source_ix > self._end[0]: # After end. return None elif source_ix == self._end[0] and self._end[1] is None: # At/after end. return None else: if source_ix > self._claimed_source_ix: # Prefer to split on even boundary. split_pos = None ratio = self._cumulative_weights[source_ix] else: # Split the current subsource. split = self.sub_range_tracker(source_ix).try_split(source_pos) if not split: return None split_pos, frac = split ratio = self.local_to_global(source_ix, frac) self._end = source_ix, split_pos self._cumulative_weights = [ min(w / ratio, 1) for w in self._cumulative_weights ] return (source_ix, split_pos), ratio def set_current_position(self, pos): raise NotImplementedError('Should only be called on sub-trackers') def position_at_fraction(self, fraction): source_ix, source_frac = self.global_to_local(fraction) last = self._end[0] if self._end[1] is None else self._end[0] + 1 if source_ix == last: return (source_ix, None) else: return ( source_ix, self.sub_range_tracker(source_ix).position_at_fraction(source_frac)) def fraction_consumed(self): with self._lock: if self._claimed_source_ix == len(self._source_bundles): return 1.0 else: return self.local_to_global( self._claimed_source_ix, self.sub_range_tracker(self._claimed_source_ix).fraction_consumed()) def local_to_global(self, source_ix, source_frac): cw = self._cumulative_weights # The global fraction is the fraction to source_ix plus some portion of # the way towards the next source. return cw[source_ix] + source_frac * (cw[source_ix + 1] - cw[source_ix]) def global_to_local(self, frac): if frac == 1: last = self._end[0] if self._end[1] is None else self._end[0] + 1 return (last, None) else: cw = self._cumulative_weights # Find the last source that starts at or before frac. source_ix = bisect.bisect(cw, frac) - 1 # Return this source, converting what's left of frac after starting # this source into a value in [0.0, 1.0) representing how far we are # towards the next source. return ( source_ix, (frac - cw[source_ix]) / (cw[source_ix + 1] - cw[source_ix])) def sub_range_tracker(self, source_ix): assert self._start[0] <= source_ix <= self._end[0] if self._range_trackers[source_ix] is None: with self._lock: if self._range_trackers[source_ix] is None: source = self._source_bundles[source_ix] if source_ix == self._start[0] and self._start[1] is not None: start = self._start[1] else: start = source.start_position if source_ix == self._end[0] and self._end[1] is not None: stop = self._end[1] else: stop = source.stop_position self._range_trackers[source_ix] = source.source.get_range_tracker( start, stop) return self._range_trackers[source_ix] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/concat_source_test.py0000644000175100017510000002227015206113032022527 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the sources framework.""" # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam.io import iobase from apache_beam.io import range_trackers from apache_beam.io import source_test_utils from apache_beam.io.concat_source import ConcatSource from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to __all__ = ['RangeSource'] class RangeSource(iobase.BoundedSource): __hash__ = None # type: ignore[assignment] def __init__(self, start, end, split_freq=1): assert start <= end self._start = start self._end = end self._split_freq = split_freq def _normalize(self, start_position, end_position): return ( self._start if start_position is None else start_position, self._end if end_position is None else end_position) def _round_up(self, index): """Rounds up to the nearest mulitple of split_freq.""" return index - index % -self._split_freq def estimate_size(self): return self._end - self._start def split(self, desired_bundle_size, start_position=None, end_position=None): start, end = self._normalize(start_position, end_position) for sub_start in range(start, end, desired_bundle_size): sub_end = min(self._end, sub_start + desired_bundle_size) yield iobase.SourceBundle( sub_end - sub_start, RangeSource(sub_start, sub_end, self._split_freq), sub_start, sub_end) def get_range_tracker(self, start_position, end_position): start, end = self._normalize(start_position, end_position) return range_trackers.OffsetRangeTracker(start, end) def read(self, range_tracker): for k in range(self._round_up(range_tracker.start_position()), self._round_up(range_tracker.stop_position())): if k % self._split_freq == 0: if not range_tracker.try_claim(k): return yield k # For testing def __eq__(self, other): return ( type(self) == type(other) and self._start == other._start and self._end == other._end and self._split_freq == other._split_freq) class ConcatSourceTest(unittest.TestCase): def test_range_source(self): source_test_utils.assert_split_at_fraction_exhaustive(RangeSource(0, 10, 3)) def test_conact_source(self): source = ConcatSource([ RangeSource(0, 4), RangeSource(4, 8), RangeSource(8, 12), RangeSource(12, 16), ]) self.assertEqual( list(source.read(source.get_range_tracker())), list(range(16))) self.assertEqual( list(source.read(source.get_range_tracker((1, None), (2, 10)))), list(range(4, 10))) range_tracker = source.get_range_tracker(None, None) self.assertEqual(range_tracker.position_at_fraction(0), (0, 0)) self.assertEqual(range_tracker.position_at_fraction(.5), (2, 8)) self.assertEqual(range_tracker.position_at_fraction(.625), (2, 10)) # Simulate a read. self.assertEqual(range_tracker.try_claim((0, None)), True) self.assertEqual(range_tracker.sub_range_tracker(0).try_claim(2), True) self.assertEqual(range_tracker.fraction_consumed(), 0.125) self.assertEqual(range_tracker.try_claim((1, None)), True) self.assertEqual(range_tracker.sub_range_tracker(1).try_claim(6), True) self.assertEqual(range_tracker.fraction_consumed(), 0.375) self.assertEqual(range_tracker.try_split((0, 1)), None) self.assertEqual(range_tracker.try_split((1, 5)), None) self.assertEqual(range_tracker.try_split((3, 14)), ((3, None), 0.75)) self.assertEqual(range_tracker.try_claim((3, None)), False) self.assertEqual(range_tracker.sub_range_tracker(1).try_claim(7), True) self.assertEqual(range_tracker.try_claim((2, None)), True) self.assertEqual(range_tracker.sub_range_tracker(2).try_claim(9), True) self.assertEqual(range_tracker.try_split((2, 8)), None) self.assertEqual(range_tracker.try_split((2, 11)), ((2, 11), 11. / 12)) self.assertEqual(range_tracker.sub_range_tracker(2).try_claim(10), True) self.assertEqual(range_tracker.sub_range_tracker(2).try_claim(11), False) def test_fraction_consumed_at_end(self): source = ConcatSource([ RangeSource(0, 2), RangeSource(2, 4), ]) range_tracker = source.get_range_tracker((2, None), None) self.assertEqual(range_tracker.fraction_consumed(), 1.0) def test_estimate_size(self): source = ConcatSource([ RangeSource(0, 10), RangeSource(10, 100), RangeSource(100, 1000), ]) self.assertEqual(source.estimate_size(), 1000) def test_position_at_fration(self): ranges = [(0, 4), (4, 16), (16, 24), (24, 32)] source = ConcatSource([ iobase.SourceBundle((range[1] - range[0]) / 32., RangeSource(*range), None, None) for range in ranges ]) range_tracker = source.get_range_tracker() self.assertEqual(range_tracker.position_at_fraction(0), (0, 0)) self.assertEqual(range_tracker.position_at_fraction(.01), (0, 1)) self.assertEqual(range_tracker.position_at_fraction(.1), (0, 4)) self.assertEqual(range_tracker.position_at_fraction(.125), (1, 4)) self.assertEqual(range_tracker.position_at_fraction(.2), (1, 7)) self.assertEqual(range_tracker.position_at_fraction(.7), (2, 23)) self.assertEqual(range_tracker.position_at_fraction(.75), (3, 24)) self.assertEqual(range_tracker.position_at_fraction(.8), (3, 26)) self.assertEqual(range_tracker.position_at_fraction(1), (4, None)) range_tracker = source.get_range_tracker((1, None), (3, None)) self.assertEqual(range_tracker.position_at_fraction(0), (1, 4)) self.assertEqual(range_tracker.position_at_fraction(.01), (1, 5)) self.assertEqual(range_tracker.position_at_fraction(.5), (1, 14)) self.assertEqual(range_tracker.position_at_fraction(.599), (1, 16)) self.assertEqual(range_tracker.position_at_fraction(.601), (2, 17)) self.assertEqual(range_tracker.position_at_fraction(1), (3, None)) def test_empty_source(self): read_all = source_test_utils.read_from_source empty = RangeSource(0, 0) self.assertEqual(read_all(ConcatSource([])), []) self.assertEqual(read_all(ConcatSource([empty])), []) self.assertEqual(read_all(ConcatSource([empty, empty])), []) range10 = RangeSource(0, 10) self.assertEqual(read_all(ConcatSource([range10]), (0, None), (0, 0)), []) self.assertEqual(read_all(ConcatSource([range10]), (0, 10), (1, None)), []) self.assertEqual( read_all(ConcatSource([range10, range10]), (0, 10), (1, 0)), []) def test_single_source(self): read_all = source_test_utils.read_from_source range10 = RangeSource(0, 10) self.assertEqual(read_all(ConcatSource([range10])), list(range(10))) self.assertEqual( read_all(ConcatSource([range10]), (0, 5)), list(range(5, 10))) self.assertEqual( read_all(ConcatSource([range10]), None, (0, 5)), list(range(5))) def test_source_with_empty_ranges(self): read_all = source_test_utils.read_from_source empty = RangeSource(0, 0) self.assertEqual(read_all(empty), []) range10 = RangeSource(0, 10) self.assertEqual( read_all(ConcatSource([empty, empty, range10])), list(range(10))) self.assertEqual( read_all(ConcatSource([empty, range10, empty])), list(range(10))) self.assertEqual( read_all(ConcatSource([range10, empty, range10, empty])), list(range(10)) + list(range(10))) def test_source_with_empty_ranges_exhastive(self): empty = RangeSource(0, 0) source = ConcatSource([ empty, RangeSource(0, 10), empty, empty, RangeSource(10, 13), RangeSource(13, 17), empty, ]) source_test_utils.assert_split_at_fraction_exhaustive(source) def test_run_concat_direct(self): source = ConcatSource([ RangeSource(0, 10), RangeSource(10, 100), RangeSource(100, 1000), ]) with TestPipeline() as pipeline: pcoll = pipeline | beam.io.Read(source) assert_that(pcoll, equal_to(list(range(1000)))) def test_conact_source_exhaustive(self): source = ConcatSource([ RangeSource(0, 10), RangeSource(100, 110), RangeSource(1000, 1010), ]) source_test_utils.assert_split_at_fraction_exhaustive(source) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/debezium.py0000644000175100017510000001744615206113032020456 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Unbounded source transform for `Debezium `_. This transform is currently supported by Beam portable Flink, Spark, and Dataflow v2 runners. **Setup** Transform provided in this module is cross-language transform implemented in the Beam Java SDK. During the pipeline construction, Python SDK will connect to a Java expansion service to expand this transform. To facilitate this, a small amount of setup is needed before using this transform in a Beam Python pipeline. There are several ways to setup cross-language Debezium transform. * Option 1: use the default expansion service * Option 2: specify a custom expansion service See below for details regarding each of these options. *Option 1: Use the default expansion service* This is the recommended and easiest setup option for using Python Debezium transform. This option requires following pre-requisites before running the Beam pipeline. * Install Java runtime in the computer from where the pipeline is constructed and make sure that 'java' command is available. In this option, Python SDK will either download (for released Beam version) or build (when running from a Beam Git clone) a expansion service jar and use that to expand transforms. Currently Debezium transform use the 'beam-sdks-java-io-debezium-expansion-service' jar for this purpose. *Option 2: specify a custom expansion service* In this option, you startup your own expansion service and provide that as a parameter when using the transform provided in this module. This option requires following pre-requisites before running the Beam pipeline. * Startup your own expansion service. * Update your pipeline to provide the expansion service address when initiating Debezium transform provided in this module. Flink Users can use the built-in Expansion Service of the Flink Runner's Job Server. If you start Flink's Job Server, the expansion service will be started on port 8097. For a different address, please set the expansion_service parameter. **More information** For more information regarding cross-language transforms see: - https://beam.apache.org/roadmap/portability/ For more information specific to Flink runner see: - https://beam.apache.org/documentation/runners/flink/ """ # pytype: skip-file import json from enum import Enum from typing import NamedTuple from typing import Optional from apache_beam.transforms import Map from apache_beam.transforms import PTransform from apache_beam.transforms.external import BeamJarExpansionService from apache_beam.transforms.external import ExternalTransform from apache_beam.transforms.external import NamedTupleBasedPayloadBuilder __all__ = ['ReadFromDebezium', 'DriverClassName'] def default_io_expansion_service(): return BeamJarExpansionService( 'sdks:java:io:debezium:expansion-service:shadowJar') class DriverClassName(Enum): MYSQL = 'MySQL' POSTGRESQL = 'PostgreSQL' ORACLE = 'Oracle' DB2 = 'Db2' ReadFromDebeziumSchema = NamedTuple( 'ReadFromDebeziumSchema', [('connector_class', str), ('username', str), ('password', str), ('host', str), ('port', str), ('max_number_of_records', Optional[int]), ('connection_properties', list[str]), ('start_offset', Optional[list[str]]), ('offset_storage_path', Optional[str])]) class ReadFromDebezium(PTransform): """ An external PTransform which reads from Debezium and returns a Dictionary for each item in the specified database connection. Experimental; no backwards compatibility guarantees. """ URN = 'beam:transform:org.apache.beam:debezium_read:v1' def __init__( self, connector_class, username, password, host, port, max_number_of_records=None, connection_properties=None, start_offset=None, offset_storage_path=None, expansion_service=None): """ Initializes a read operation from Debezium. :param connector_class: name of the jdbc driver class :param username: database username :param password: database password :param host: database host :param port: database port :param max_number_of_records: maximum number of records to be fetched before stop. :param connection_properties: properties of the debezium connection passed as string with format [propertyName=property;]* :param start_offset: starting offset to resume the connector from a previously seen position. Provided as a list of "key=value" strings, where numeric values are encoded as their decimal string representation. Example for PostgreSQL:: start_offset=["lsn=28160840"] Example for MySQL:: start_offset=["file=binlog.000001", "pos=156"] Obtain the offset from the JSON output of a previous pipeline run (the "metadata" field contains connector-specific position info) or via ``SourceRecord.sourceOffset()`` in a custom Java SourceRecordMapper. :param offset_storage_path: path to a file where the connector offset is automatically saved after each checkpoint and loaded on pipeline startup, allowing the pipeline to resume from where it left off. Supports any filesystem available to the Beam runner (local, GCS, S3, etc.). Example:: offset_storage_path=( "gs://my-bucket/debezium/offset.json" ) When set, takes precedence over ``start_offset``. :param expansion_service: The address (host:port) of the ExpansionService. """ self.params = ReadFromDebeziumSchema( connector_class=connector_class.value, username=username, password=password, host=host, port=str(port), max_number_of_records=max_number_of_records, connection_properties=connection_properties, start_offset=start_offset, offset_storage_path=offset_storage_path) self.expansion_service = expansion_service or default_io_expansion_service() def expand(self, pbegin): return ( pbegin | ExternalTransform( self.URN, NamedTupleBasedPayloadBuilder(self.params), self.expansion_service, ).with_output_types(str) | 'JsonToDict' >> Map(json.loads).with_output_types(dict)) ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3204014 apache_beam-2.74.0/apache_beam/io/external/0000755000175100017510000000000015206113157020116 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/__init__.py0000644000175100017510000000142015206113032022214 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3204014 apache_beam-2.74.0/apache_beam/io/external/gcp/0000755000175100017510000000000015206113157020667 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/gcp/__init__.py0000644000175100017510000000142015206113032022765 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/gcp/pubsub.py0000644000175100017510000001454015206113032022535 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import typing import apache_beam as beam from apache_beam.io.gcp import pubsub from apache_beam.transforms import Map from apache_beam.transforms.external import ExternalTransform from apache_beam.transforms.external import NamedTupleBasedPayloadBuilder ReadFromPubsubSchema = typing.NamedTuple( 'ReadFromPubsubSchema', [ ('topic', typing.Optional[str]), ('subscription', typing.Optional[str]), ('id_label', typing.Optional[str]), ('with_attributes', bool), ('timestamp_attribute', typing.Optional[str]), ]) class ReadFromPubSub(beam.PTransform): """An external ``PTransform`` for reading from Cloud Pub/Sub. Experimental; no backwards compatibility guarantees. It requires special preparation of the Java SDK. See https://github.com/apache/beam/issues/19728. """ URN = 'beam:transform:org.apache.beam:pubsub_read:v1' def __init__( self, topic=None, subscription=None, id_label=None, with_attributes=False, timestamp_attribute=None, expansion_service=None): """Initializes ``ReadFromPubSub``. Args: topic: Cloud Pub/Sub topic in the form "projects//topics/". If provided, subscription must be None. subscription: Existing Cloud Pub/Sub subscription to use in the form "projects//subscriptions/". If not specified, a temporary subscription will be created from the specified topic. If provided, topic must be None. id_label: The attribute on incoming Pub/Sub messages to use as a unique record identifier. When specified, the value of this attribute (which can be any string that uniquely identifies the record) will be used for deduplication of messages. If not provided, we cannot guarantee that no duplicate data will be delivered on the Pub/Sub stream. In this case, deduplication of the stream will be strictly best effort. with_attributes: True - output elements will be :class:`~apache_beam.io.gcp.pubsub.PubsubMessage` objects. False - output elements will be of type ``bytes`` (message data only). timestamp_attribute: Message value to use as element timestamp. If None, uses message publishing time as the timestamp. Timestamp values should be in one of two formats: - A numerical value representing the number of milliseconds since the Unix epoch. - A string in RFC 3339 format, UTC timezone. Example: ``2015-10-29T23:41:41.123Z``. The sub-second component of the timestamp is optional, and digits beyond the first three (i.e., time units smaller than milliseconds) may be ignored. """ self.params = ReadFromPubsubSchema( topic=topic, subscription=subscription, id_label=id_label, with_attributes=with_attributes, timestamp_attribute=timestamp_attribute) self.expansion_service = expansion_service def expand(self, pbegin): pcoll = pbegin.apply( ExternalTransform( self.URN, NamedTupleBasedPayloadBuilder(self.params), self.expansion_service)) if self.params.with_attributes: pcoll = pcoll | 'FromProto' >> Map(pubsub.PubsubMessage._from_proto_str) pcoll.element_type = pubsub.PubsubMessage else: pcoll.element_type = bytes return pcoll WriteToPubsubSchema = typing.NamedTuple( 'WriteToPubsubSchema', [ ('topic', str), ('id_label', typing.Optional[str]), # this is not implemented yet on the Java side: # ('with_attributes', bool), ('timestamp_attribute', typing.Optional[str]), ]) class WriteToPubSub(beam.PTransform): """An external ``PTransform`` for writing messages to Cloud Pub/Sub. Experimental; no backwards compatibility guarantees. It requires special preparation of the Java SDK. See https://github.com/apache/beam/issues/19728. """ URN = 'beam:transform:org.apache.beam:pubsub_write:v1' def __init__( self, topic, with_attributes=False, id_label=None, timestamp_attribute=None, expansion_service=None): """Initializes ``WriteToPubSub``. Args: topic: Cloud Pub/Sub topic in the form "/topics//". with_attributes: True - input elements will be :class:`~apache_beam.io.gcp.pubsub.PubsubMessage` objects. False - input elements will be of type ``bytes`` (message data only). id_label: If set, will set an attribute for each Cloud Pub/Sub message with the given name and a unique value. This attribute can then be used in a ReadFromPubSub PTransform to deduplicate messages. timestamp_attribute: If set, will set an attribute for each Cloud Pub/Sub message with the given name and the message's publish time as the value. """ self.params = WriteToPubsubSchema( topic=topic, id_label=id_label, # with_attributes=with_attributes, timestamp_attribute=timestamp_attribute) self.expansion_service = expansion_service self.with_attributes = with_attributes def expand(self, pvalue): if self.with_attributes: pcoll = pvalue | 'ToProto' >> Map(pubsub.WriteToPubSub.to_proto_str) else: pcoll = pvalue | 'ToProto' >> Map( lambda x: pubsub.PubsubMessage(x, {})._to_proto_str()) pcoll.element_type = bytes return pcoll.apply( ExternalTransform( self.URN, NamedTupleBasedPayloadBuilder(self.params), self.expansion_service)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/generate_sequence.py0000644000175100017510000000464315206113032024151 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file from apache_beam.transforms.external import ExternalTransform from apache_beam.transforms.external import ImplicitSchemaPayloadBuilder class GenerateSequence(ExternalTransform): """ An external PTransform which provides a bounded or unbounded stream of integers. Note: To use this transform, you need to start the Java expansion service. Please refer to the portability documentation on how to do that. The expansion service address has to be provided when instantiating this transform. During pipeline translation this transform will be replaced by the Java SDK's GenerateSequence. If you start Flink's job server, the expansion service will be started on port 8097. This is also the configured default for this transform. For a different address, please set the expansion_service parameter. For more information see: - https://beam.apache.org/documentation/runners/flink/ - https://beam.apache.org/roadmap/portability/ Note: Runners need to support translating Read operations in order to use this source. At the moment only the Flink Runner supports this. Experimental; no backwards compatibility guarantees. """ URN = 'beam:external:java:generate_sequence:v1' def __init__( self, start, stop=None, elements_per_period=None, max_read_time=None, expansion_service=None): super().__init__( self.URN, ImplicitSchemaPayloadBuilder({ 'start': start, 'stop': stop, 'elements_per_period': elements_per_period, 'max_read_time': max_read_time, }), expansion_service) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/generate_sequence_test.py0000644000175100017510000000640115206113032025202 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for cross-language generate sequence.""" # pytype: skip-file import logging import os import re import unittest import pytest from apache_beam.io.external.generate_sequence import GenerateSequence from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.external import ExternalTransform from apache_beam.transforms.external import JavaClassLookupPayloadBuilder from apache_beam.transforms.external import JavaExternalTransform @pytest.mark.uses_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_PORT'), "EXPANSION_PORT environment var is not provided.") class XlangGenerateSequenceTest(unittest.TestCase): def test_generate_sequence(self): port = os.environ.get('EXPANSION_PORT') address = 'localhost:%s' % port try: with TestPipeline() as p: res = ( p | GenerateSequence(start=1, stop=10, expansion_service=address)) assert_that(res, equal_to(list(range(1, 10)))) except RuntimeError as e: if re.search(GenerateSequence.URN, str(e)): print("looks like URN not implemented in expansion service, skipping.") else: raise e def test_generate_sequence_java_class_lookup_payload_builder(self): port = os.environ.get('EXPANSION_PORT') address = 'localhost:%s' % port with TestPipeline() as p: payload_builder = JavaClassLookupPayloadBuilder( 'org.apache.beam.sdk.io.GenerateSequence') payload_builder.with_constructor_method('from', 1) payload_builder.add_builder_method('to', 10) res = ( p | ExternalTransform(None, payload_builder, expansion_service=address)) assert_that(res, equal_to(list(range(1, 10)))) def test_generate_sequence_java_external_transform(self): port = os.environ.get('EXPANSION_PORT') address = 'localhost:%s' % port with TestPipeline() as p: java_transform = JavaExternalTransform( 'org.apache.beam.sdk.io.GenerateSequence', expansion_service=address) # We have to use 'getattr' below for builder method 'from' of Java # 'GenerateSequence' class since 'from' is a reserved keyword for Python. java_transform = getattr(java_transform, 'from')(1).to(10) res = (p | java_transform) assert_that(res, equal_to(list(range(1, 10)))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/kafka.py0000644000175100017510000000222715206113032021540 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # For backwards compatibility. # pytype: skip-file import warnings # pylint: disable=unused-import from apache_beam.io.kafka import ReadFromKafka from apache_beam.io.kafka import ReadFromKafkaSchema from apache_beam.io.kafka import WriteToKafka from apache_beam.io.kafka import WriteToKafkaSchema warnings.warn( "Kafka transforms moved to apache_beam.io.kafka", DeprecationWarning) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/snowflake.py0000644000175100017510000000226715206113032022460 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # For backwards compatibility. # pytype: skip-file import warnings # pylint: disable=unused-import from apache_beam.io.snowflake import CreateDisposition from apache_beam.io.snowflake import ReadFromSnowflake from apache_beam.io.snowflake import WriteDisposition from apache_beam.io.snowflake import WriteToSnowflake warnings.warn( "Snowflake transforms moved to apache_beam.io.snowflake", DeprecationWarning) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/xlang_bigqueryio_it_test.py0000644000175100017510000005336415206113032025576 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for cross-language BigQuery sources and sinks.""" # pytype: skip-file import datetime import logging import os import secrets import time import unittest from decimal import Decimal import pytest from hamcrest.core import assert_that as hamcrest_assert from hamcrest.core.core.allof import all_of import apache_beam as beam from apache_beam.io.gcp.bigquery import StorageWriteToBigQuery from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultMatcher from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultStreamingMatcher from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.transforms.periodicsequence import PeriodicImpulse from apache_beam.utils.timestamp import Timestamp # Protect against environments where bigquery library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.io.gcp.gcsio import GcsIO except ImportError: GcsIO = None try: from apitools.base.py.exceptions import HttpError except ImportError: HttpError = None # pylint: enable=wrong-import-order, wrong-import-position _LOGGER = logging.getLogger(__name__) @pytest.mark.uses_gcp_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_JARS'), "EXPANSION_JARS environment var is not provided, " "indicating that jars have not been built") class BigQueryXlangStorageWriteIT(unittest.TestCase): BIGQUERY_DATASET = 'python_xlang_storage_write' ELEMENTS = [ # (int, float, numeric, string, bool, bytes, timestamp) { "int": 1, "float": 0.1, "numeric": Decimal("1.11"), "str": "a", "bool": True, "bytes": b'a', "timestamp": Timestamp(1000, 100) }, { "int": 2, "float": 0.2, "numeric": Decimal("2.22"), "str": "b", "bool": False, "bytes": b'b', "timestamp": Timestamp(2000, 200) }, { "int": 3, "float": 0.3, "numeric": Decimal("3.33"), "str": "c", "bool": True, "bytes": b'd', "timestamp": Timestamp(3000, 300) }, { "int": 4, "float": 0.4, "numeric": Decimal("4.44"), "str": "d", "bool": False, "bytes": b'd', "timestamp": Timestamp(4000, 400) } ] ALL_TYPES_SCHEMA = ( "int:INTEGER,float:FLOAT,numeric:NUMERIC,str:STRING," "bool:BOOLEAN,bytes:BYTES,timestamp:TIMESTAMP") def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.args = self.test_pipeline.get_full_options_as_args() self.project = self.test_pipeline.get_option('project') self._runner = PipelineOptions(self.args).get_all_options()['runner'] self.bigquery_client = BigQueryWrapper.from_pipeline_options( self.test_pipeline.options) self.dataset_id = '%s_%s_%s' % ( self.BIGQUERY_DATASET, str(int(time.time())), secrets.token_hex(3)) self.bigquery_client.get_or_create_dataset(self.project, self.dataset_id) _LOGGER.info( "Created dataset %s in project %s", self.dataset_id, self.project) def tearDown(self): try: _LOGGER.info( "Deleting dataset %s in project %s", self.dataset_id, self.project) self.bigquery_client._delete_dataset( project_id=self.project, dataset_id=self.dataset_id, delete_contents=True) except HttpError: _LOGGER.debug( 'Failed to clean up dataset %s in project %s', self.dataset_id, self.project) def parse_expected_data(self, expected_elements): if not isinstance(expected_elements, list): expected_elements = [expected_elements] data = [] for row in expected_elements: values = list(row.values()) for i, val in enumerate(values): if isinstance(val, Timestamp): # BigQuery matcher query returns a datetime.datetime object values[i] = val.to_utc_datetime().replace( tzinfo=datetime.timezone.utc) data.append(tuple(values)) return data def assert_iceberg_tables_created( self, table_prefix, storage_uri, expected_count=1): """Verify that Iceberg table directories are created in the warehouse location. Args: table_prefix: The table name prefix to look for storage_uri: The GCS storage URI (e.g., 'gs://bucket/path') expected_count: Expected number of table directories """ if GcsIO is None: _LOGGER.warning( "GcsIO not available, skipping warehouse location verification") return gcs_io = GcsIO() # Parse the storage URI to get bucket and prefix if not storage_uri.startswith('gs://'): raise ValueError(f'Storage URI must start with gs://, got: {storage_uri}') # Remove 'gs://' prefix and split bucket from path path_parts = storage_uri[5:].split('/', 1) bucket_name = path_parts[0] base_prefix = path_parts[1] if len(path_parts) > 1 else '' # Construct the full prefix to search for table directories # Following the pattern: # {base_prefix}/{project}/{dataset}/{table_prefix} search_prefix = ( f"{base_prefix}/" f"{self.project}/{self.dataset_id}/{table_prefix}") # List objects in the bucket with the constructed prefix try: objects = gcs_io.list_files(f"gs://{bucket_name}/{search_prefix}") object_count = len(list(objects)) if object_count < expected_count: raise AssertionError( f"Expected at least {expected_count} objects in warehouse " f"location gs://{bucket_name}/{search_prefix}, but found " f"{object_count}") _LOGGER.info( "Successfully verified %s objects created in " "warehouse location gs://%s/%s", object_count, bucket_name, search_prefix) except Exception as e: raise AssertionError( f"Failed to verify table creation in warehouse location " f"gs://{bucket_name}/{search_prefix}: {str(e)}") def run_storage_write_test( self, table_name, items, schema, use_at_least_once=False): table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table_name) bq_matcher = BigqueryFullResultMatcher( project=self.project, query="SELECT * FROM %s" % '{}.{}'.format(self.dataset_id, table_name), data=self.parse_expected_data(items)) with beam.Pipeline(argv=self.args) as p: _ = ( p | beam.Create(items) | beam.io.WriteToBigQuery( table=table_id, method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, schema=schema, use_at_least_once=use_at_least_once)) hamcrest_assert(p, bq_matcher) def test_all_types(self): table_name = "all_types" schema = self.ALL_TYPES_SCHEMA self.run_storage_write_test(table_name, self.ELEMENTS, schema) def test_with_at_least_once_semantics(self): table_name = "with_at_least_once_semantics" schema = self.ALL_TYPES_SCHEMA self.run_storage_write_test( table_name, self.ELEMENTS, schema, use_at_least_once=True) def test_nested_records_and_lists(self): table_name = "nested_records_and_lists" schema = { "fields": [{ "name": "repeated_int", "type": "INTEGER", "mode": "REPEATED" }, { "name": "struct", "type": "STRUCT", "fields": [{ "name": "nested_int", "type": "INTEGER" }, { "name": "nested_str", "type": "STRING" }] }, { "name": "repeated_struct", "type": "STRUCT", "mode": "REPEATED", "fields": [{ "name": "nested_numeric", "type": "NUMERIC" }, { "name": "nested_bytes", "type": "BYTES" }] }] } items = [{ "repeated_int": [1, 2, 3], "struct": { "nested_int": 1, "nested_str": "a" }, "repeated_struct": [{ "nested_numeric": Decimal("1.23"), "nested_bytes": b'a' }, { "nested_numeric": Decimal("3.21"), "nested_bytes": b'aa' }] }] self.run_storage_write_test(table_name, items, schema) def test_write_with_beam_rows(self): table = 'write_with_beam_rows' table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table) row_elements = [ beam.Row( my_int=e['int'], my_float=e['float'], my_numeric=e['numeric'], my_string=e['str'], my_bool=e['bool'], my_bytes=e['bytes'], my_timestamp=e['timestamp']) for e in self.ELEMENTS ] bq_matcher = BigqueryFullResultMatcher( project=self.project, query="SELECT * FROM {}.{}".format(self.dataset_id, table), data=self.parse_expected_data(self.ELEMENTS)) with beam.Pipeline(argv=self.args) as p: _ = ( p | beam.Create(row_elements) | StorageWriteToBigQuery(table=table_id)) hamcrest_assert(p, bq_matcher) def test_write_with_clustering(self): table = 'write_with_clustering' table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table) bq_matcher = BigqueryFullResultMatcher( project=self.project, query="SELECT * FROM {}.{}".format(self.dataset_id, table), data=self.parse_expected_data(self.ELEMENTS)) with beam.Pipeline(argv=self.args) as p: _ = ( p | "Create test data" >> beam.Create(self.ELEMENTS) | beam.io.WriteToBigQuery( table=table_id, method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, schema=self.ALL_TYPES_SCHEMA, create_disposition='CREATE_IF_NEEDED', write_disposition='WRITE_TRUNCATE', additional_bq_parameters={'clustering': { 'fields': ['int'] }})) # After pipeline finishes, verify clustering is applied table = self.bigquery_client.get_table(self.project, self.dataset_id, table) clustering_fields = table.clustering.fields if table.clustering else [] self.assertEqual(clustering_fields, ['int']) hamcrest_assert(p, bq_matcher) def test_write_with_beam_rows_cdc(self): table = 'write_with_beam_rows_cdc' table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table) expected_data_on_bq = [ # (name, value) { "name": "cdc_test", "value": 5, } ] rows_with_cdc = [ beam.Row( row_mutation_info=beam.Row( mutation_type="UPSERT", change_sequence_number="AAA/2"), record=beam.Row(name="cdc_test", value=5)), beam.Row( row_mutation_info=beam.Row( mutation_type="UPSERT", change_sequence_number="AAA/1"), record=beam.Row(name="cdc_test", value=3)) ] bq_matcher = BigqueryFullResultMatcher( project=self.project, query="SELECT * FROM {}.{}".format(self.dataset_id, table), data=self.parse_expected_data(expected_data_on_bq)) with beam.Pipeline(argv=self.args) as p: _ = ( p | beam.Create(rows_with_cdc) | beam.io.WriteToBigQuery( table=table_id, method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, use_at_least_once=True, use_cdc_writes=True, primary_key=["name"])) hamcrest_assert(p, bq_matcher) def test_write_with_dicts_cdc(self): table = 'write_with_dicts_cdc' table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table) expected_data_on_bq = [ # (name, value) { "name": "cdc_test", "value": 5, } ] data_with_cdc = [ # record: (name, value) { 'row_mutation_info': { 'mutation_type': 'UPSERT', 'change_sequence_number': 'AAA/2' }, 'record': { 'name': 'cdc_test', 'value': 5 } }, { 'row_mutation_info': { 'mutation_type': 'UPSERT', 'change_sequence_number': 'AAA/1' }, 'record': { 'name': 'cdc_test', 'value': 3 } } ] schema = { "fields": [ # include both record and mutation info fields as part of the schema { "name": "row_mutation_info", "type": "STRUCT", "fields": [ # setting both fields are required { "name": "mutation_type", "type": "STRING", "mode": "REQUIRED" }, { "name": "change_sequence_number", "type": "STRING", "mode": "REQUIRED" } ] }, { "name": "record", "type": "STRUCT", "fields": [{ "name": "name", "type": "STRING" }, { "name": "value", "type": "INTEGER" }] } ] } bq_matcher = BigqueryFullResultMatcher( project=self.project, query="SELECT * FROM {}.{}".format(self.dataset_id, table), data=self.parse_expected_data(expected_data_on_bq)) with beam.Pipeline(argv=self.args) as p: _ = ( p | beam.Create(data_with_cdc) | beam.io.WriteToBigQuery( table=table_id, method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, use_at_least_once=True, use_cdc_writes=True, schema=schema, primary_key=["name"])) hamcrest_assert(p, bq_matcher) def test_write_to_dynamic_destinations(self): base_table_spec = '{}.dynamic_dest_'.format(self.dataset_id) spec_with_project = '{}:{}'.format(self.project, base_table_spec) tables = [base_table_spec + str(record['int']) for record in self.ELEMENTS] bq_matchers = [ BigqueryFullResultMatcher( project=self.project, query="SELECT * FROM %s" % tables[i], data=self.parse_expected_data(self.ELEMENTS[i])) for i in range(len(tables)) ] with beam.Pipeline(argv=self.args) as p: _ = ( p | beam.Create(self.ELEMENTS) | beam.io.WriteToBigQuery( table=lambda record: spec_with_project + str(record['int']), method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, schema=self.ALL_TYPES_SCHEMA, use_at_least_once=False)) hamcrest_assert(p, all_of(*bq_matchers)) def test_write_to_dynamic_destinations_with_beam_rows(self): base_table_spec = '{}.dynamic_dest_'.format(self.dataset_id) spec_with_project = '{}:{}'.format(self.project, base_table_spec) tables = [base_table_spec + str(record['int']) for record in self.ELEMENTS] bq_matchers = [ BigqueryFullResultMatcher( project=self.project, query="SELECT * FROM %s" % tables[i], data=self.parse_expected_data(self.ELEMENTS[i])) for i in range(len(tables)) ] row_elements = [ beam.Row( my_int=e['int'], my_float=e['float'], my_numeric=e['numeric'], my_string=e['str'], my_bool=e['bool'], my_bytes=e['bytes'], my_timestamp=e['timestamp']) for e in self.ELEMENTS ] with beam.Pipeline(argv=self.args) as p: _ = ( p | beam.Create(row_elements) | beam.io.WriteToBigQuery( table=lambda record: spec_with_project + str(record.my_int), method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, use_at_least_once=False)) hamcrest_assert(p, all_of(*bq_matchers)) def run_streaming(self, table_name, num_streams=0, use_at_least_once=False): elements = self.ELEMENTS.copy() schema = self.ALL_TYPES_SCHEMA table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table_name) bq_matcher = BigqueryFullResultStreamingMatcher( project=self.project, query="SELECT * FROM {}.{}".format(self.dataset_id, table_name), data=self.parse_expected_data(self.ELEMENTS)) args = self.test_pipeline.get_full_options_as_args( on_success_matcher=bq_matcher, streaming=True, allow_unsafe_triggers=True) auto_sharding = (num_streams == 0) with beam.Pipeline(argv=args) as p: _ = ( p | PeriodicImpulse(0, 4, 1) | beam.Map(lambda t: elements[t]) | beam.io.WriteToBigQuery( table=table_id, method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, schema=schema, triggering_frequency=1, with_auto_sharding=auto_sharding, num_storage_api_streams=num_streams, use_at_least_once=use_at_least_once)) hamcrest_assert(p, bq_matcher) def skip_if_not_dataflow_runner(self) -> bool: # skip if dataflow runner is not specified if not self._runner or "dataflowrunner" not in self._runner.lower(): self.skipTest( "Streaming with exactly-once route has the requirement " "`beam:requirement:pardo:on_window_expiration:v1`, " "which is currently only supported by the Dataflow runner") def test_streaming_with_fixed_num_streams(self): self.skip_if_not_dataflow_runner() table = 'streaming_fixed_num_streams' self.run_streaming(table_name=table, num_streams=4) @unittest.skip( "Streaming to the Storage Write API sink with autosharding is broken " "with Dataflow Runner V2.") def test_streaming_with_auto_sharding(self): self.skip_if_not_dataflow_runner() table = 'streaming_with_auto_sharding' self.run_streaming(table_name=table) def test_streaming_with_at_least_once(self): table = 'streaming_with_at_least_once' self.run_streaming(table_name=table, use_at_least_once=True) def test_write_with_big_lake_configuration(self): """Test BigQuery Storage Write API with BigLake configuration.""" table = 'write_with_big_lake_config' table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table) # BigLake configuration with required parameters (matching Java test) big_lake_config = { 'connectionId': 'projects/apache-beam-testing/locations/us/connections/apache-beam-testing-storageapi-biglake-nodelete', # pylint: disable=line-too-long 'storageUri': 'gs://apache-beam-testing-bq-biglake/BigQueryXlangStorageWriteIT', # pylint: disable=line-too-long 'fileFormat': 'parquet', 'tableFormat': 'iceberg' } bq_matcher = BigqueryFullResultMatcher( project=self.project, query="SELECT * FROM {}.{}".format(self.dataset_id, table), data=self.parse_expected_data(self.ELEMENTS)) with beam.Pipeline(argv=self.args) as p: _ = ( p | "Create test data" >> beam.Create(self.ELEMENTS) | beam.io.WriteToBigQuery( table=table_id, method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, schema=self.ALL_TYPES_SCHEMA, create_disposition='CREATE_IF_NEEDED', write_disposition='WRITE_TRUNCATE', big_lake_configuration=big_lake_config)) hamcrest_assert(p, bq_matcher) # Verify that the table directory was created in the warehouse location self.assert_iceberg_tables_created(table, big_lake_config['storageUri']) def test_write_with_managed_transform(self): table = 'write_with_managed_transform' table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table) row_elements = [ beam.Row( my_int=e['int'], my_float=e['float'], my_string=e['str'], my_bool=e['bool'], my_bytes=e['bytes'], my_timestamp=e['timestamp']) for e in self.ELEMENTS ] expected = [] for e in self.ELEMENTS: del e["numeric"] expected.append(e) bq_matcher = BigqueryFullResultMatcher( project=self.project, query="SELECT * FROM {}.{}".format(self.dataset_id, table), data=self.parse_expected_data(expected)) with beam.Pipeline(argv=self.args) as p: _ = ( p | beam.Create(row_elements) | beam.managed.Write("bigquery", config={"table": table_id})) hamcrest_assert(p, bq_matcher) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/xlang_debeziumio_it_test.py0000644000175100017510000001145615206113032025547 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import subprocess import unittest from apache_beam.io.debezium import DriverClassName from apache_beam.io.debezium import ReadFromDebezium from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.utils import subprocess_server # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from testcontainers.postgres import PostgresContainer except ImportError: PostgresContainer = None NUM_RECORDS = 1 def _disable_debezium_test(): # disable if run on > ReadFromDebezium( username=self.username, password=self.password, host=self.host, port=self.port, max_number_of_records=NUM_RECORDS, connector_class=self.connector_class, connection_properties=self.connection_properties)) assert_that(results, equal_to(expected_response)) # Creating a container with testcontainers sometimes raises ReadTimeout # error. In java there are 2 retries set by default. def start_db_container(self, retries): for i in range(retries): try: self.db = PostgresContainer( 'quay.io/debezium/example-postgres:latest', username=self.username, password=self.password, dbname=self.database) self.db.start() break except Exception as e: # pylint: disable=bare-except if i == retries - 1: logging.error('Unable to initialize DB container.') raise e if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/xlang_jdbcio_it_test.py0000644000175100017510000003755215206113032024652 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import datetime import logging import os import time import typing import unittest from decimal import Decimal import pytest from parameterized import parameterized import apache_beam as beam from apache_beam import coders from apache_beam.io import jdbc from apache_beam.io.jdbc import ReadFromJdbc from apache_beam.io.jdbc import WriteToJdbc from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.utils.timestamp import Timestamp # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: import sqlalchemy except ImportError: sqlalchemy = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from testcontainers.mysql import MySqlContainer from testcontainers.postgres import PostgresContainer except ImportError: PostgresContainer = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports ROW_COUNT = 10 POSTGRES_BINARY_TYPE = ('BYTEA', 'BYTEA') MYSQL_BINARY_TYPE = ('BINARY(10)', 'VARBINARY(10)') JdbcTestRow = typing.NamedTuple( "JdbcTestRow", [("f_id", int), ("f_id_long", int), ("f_float", float), ("f_char", str), ("f_varchar", str), ("f_bytes", bytes), ("f_varbytes", bytes), ("f_timestamp", Timestamp), ("f_decimal", Decimal), ("f_date", datetime.date), ("f_time", datetime.time)], ) coders.registry.register_row(JdbcTestRow) CustomSchemaRow = typing.NamedTuple( "CustomSchemaRow", [ ("renamed_id", int), ("renamed_id_long", int), ("renamed_float", float), ("renamed_char", str), ("renamed_varchar", str), ("renamed_bytes", bytes), ("renamed_varbytes", bytes), ("renamed_timestamp", Timestamp), ("renamed_decimal", Decimal), ("renamed_date", datetime.date), ("renamed_time", datetime.time), ], ) # Need to put inside enforce_millis_instant_for_timestamp context to align # with the same setup in ReadFromJdbc.__init__. Remove once Beam moved to # micros instant for timestamp # Alternatively, use coders.registry.register_coder(CustomSchemaRow, RowCoder) with jdbc.enforce_millis_instant_for_timestamp(): coders.registry.register_row(CustomSchemaRow) SimpleRow = typing.NamedTuple( "SimpleRow", [("id", int), ("name", str), ("value", float)]) coders.registry.register_row(SimpleRow) @pytest.mark.uses_gcp_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_JARS'), "EXPANSION_JARS environment var is not provided, " "indicating that jars have not been built") @unittest.skipIf(sqlalchemy is None, 'sql alchemy package is not installed.') @unittest.skipIf( PostgresContainer is None, 'testcontainers package is not installed') @unittest.skipIf( TestPipeline().get_pipeline_options().view_as(StandardOptions).runner is None, 'Do not run this test on precommit suites.') @unittest.skipIf( TestPipeline().get_pipeline_options().view_as(StandardOptions).runner is not None and "dataflowrunner" in TestPipeline().get_pipeline_options().view_as( StandardOptions).runner.lower(), 'Do not run this test on dataflow runner.') class CrossLanguageJdbcIOTest(unittest.TestCase): DbData = typing.NamedTuple( 'DbData', [('container_fn', typing.Any), ('classpath', list[str]), ('db_string', str), ('connector', str)]) DB_CONTAINER_CLASSPATH_STRING = { 'postgres': DbData( lambda: PostgresContainer('postgres:12.3'), None, 'postgresql', 'org.postgresql.Driver'), 'mysql': DbData( lambda: MySqlContainer(dialect='pymysql'), ['mysql:mysql-connector-java:8.0.28'], 'mysql', 'com.mysql.cj.jdbc.Driver') } @classmethod def setUpClass(cls): cls.containers = {} cls.engines = {} cls.jdbc_configs = {} for db_type, db_data in cls.DB_CONTAINER_CLASSPATH_STRING.items(): container = cls.start_container(db_data.container_fn) cls.containers[db_type] = container cls.engines[db_type] = sqlalchemy.create_engine( container.get_connection_url()) cls.jdbc_configs[db_type] = { 'username': 'test', 'password': 'test', 'host': container.get_container_host_ip(), 'port': container.get_exposed_port(container.port), 'database_name': 'test', 'driver_class_name': db_data.connector, 'classpath': db_data.classpath, 'jdbc_url': ( f'jdbc:{db_data.db_string}://{container.get_container_host_ip()}:' f'{container.get_exposed_port(container.port)}/test'), 'binary_type': POSTGRES_BINARY_TYPE if db_type == 'postgres' else MYSQL_BINARY_TYPE } @classmethod def tearDownClass(cls): for db_type, container in cls.containers.items(): if container: # Sometimes stopping the container raises ReadTimeout. We can ignore it # here to avoid the test failure. try: container.stop() except Exception: # pylint: disable=broad-except logging.warning("Could not stop %s container", db_type) @classmethod def start_container(cls, container_init, max_retries=3): # Creating a container with testcontainers sometimes raises ReadTimeout # error. In java there are 2 retries set by default. for attempt in range(max_retries): try: container = container_init() container.start() return container except Exception: # pylint: disable=broad-except if attempt == max_retries - 1: logging.error( 'Failed to initialize container after %s attempts', max_retries) raise def create_test_table(self, connection, table_name, database): binary_type = self.jdbc_configs[database]['binary_type'] connection.execute( sqlalchemy.text( f"CREATE TABLE IF NOT EXISTS {table_name}" + "(f_id INTEGER, f_id_long BIGINT, f_float DOUBLE PRECISION, " + "f_char CHAR(10), f_varchar VARCHAR(10), " + f"f_bytes {binary_type[0]}, f_varbytes {binary_type[1]}, " + "f_timestamp TIMESTAMP(3), f_decimal DECIMAL(10, 2), " + "f_date DATE, f_time TIME(3))")) def generate_test_data(self, count): return [ JdbcTestRow( i - 3, i - 3, i + 0.1, f'Test{i}', f'Test{i}', f'Test{i}'.encode(), f'Test{i}'.encode(), Timestamp.of(seconds=round(time.time(), 3)), Decimal(f'{i-1}.23'), datetime.date(1969 + i, i % 12 + 1, i % 31 + 1), datetime.time(i % 24, i % 60, i % 60, (i * 1000) % 1_000_000)) for i in range(count) ] @parameterized.expand(['postgres', 'mysql']) def test_xlang_jdbc_write_read(self, database): table_name = f"jdbc_write_read_test_{database}" with self.engines[database].begin() as connection: self.create_test_table(connection, table_name, database) test_rows = self.generate_test_data(ROW_COUNT) expected_rows = [] for row in test_rows: f_char = row.f_char + ' ' * (10 - len(row.f_char)) f_bytes = row.f_bytes if database != 'postgres': f_bytes = row.f_bytes + b'\0' * (10 - len(row.f_bytes)) expected_rows.append( JdbcTestRow( row.f_id, row.f_id, row.f_float, f_char, row.f_varchar, f_bytes, row.f_bytes, row.f_timestamp, row.f_decimal, row.f_date, row.f_time)) config = self.jdbc_configs[database] with TestPipeline() as p: p.not_use_test_runner_api = True _ = ( p | beam.Create(test_rows).with_output_types(JdbcTestRow) | 'Write to jdbc' >> WriteToJdbc( table_name=table_name, driver_class_name=config['driver_class_name'], jdbc_url=config['jdbc_url'], username=config['username'], password=config['password'], classpath=config['classpath'], )) with TestPipeline() as p: p.not_use_test_runner_api = True result = ( p | 'Read from jdbc' >> ReadFromJdbc( table_name=table_name, driver_class_name=config['driver_class_name'], jdbc_url=config['jdbc_url'], username=config['username'], password=config['password'], classpath=config['classpath'])) assert_that(result, equal_to(expected_rows)) # Try the same read using the partitioned reader code path. # Outputs should be the same. with TestPipeline() as p: p.not_use_test_runner_api = True result = ( p | 'Partitioned read from jdbc' >> ReadFromJdbc( table_name=table_name, partition_column='f_id', partitions=3, driver_class_name=config['driver_class_name'], jdbc_url=config['jdbc_url'], username=config['username'], password=config['password'], classpath=config['classpath'])) assert_that(result, equal_to(expected_rows)) @parameterized.expand(['postgres', 'mysql']) def test_xlang_jdbc_read_with_explicit_schema(self, database): if self.containers[database] is None: self.skipTest(f"{database} container could not be initialized") table_name = f"jdbc_schema_test_{database}" with self.engines[database].begin() as connection: self.create_test_table(connection, table_name, database) test_rows = self.generate_test_data(ROW_COUNT) expected_rows = [] for row in test_rows: f_char = row.f_char f_bytes = row.f_bytes if database != 'postgres': f_bytes = row.f_bytes + b'\0' * (10 - len(row.f_bytes)) expected_rows.append( CustomSchemaRow( row.f_id, row.f_id, row.f_float, f_char, row.f_varchar, f_bytes, row.f_bytes, row.f_timestamp, row.f_decimal, row.f_date, row.f_time)) def custom_row_equals(expected, actual): return ( expected.renamed_id == actual.renamed_id and expected.renamed_id_long == actual.renamed_id_long and expected.renamed_float == actual.renamed_float and expected.renamed_char.rstrip() == actual.renamed_char.rstrip() and expected.renamed_varchar == actual.renamed_varchar and expected.renamed_bytes == actual.renamed_bytes and expected.renamed_timestamp == actual.renamed_timestamp and expected.renamed_decimal == actual.renamed_decimal and expected.renamed_date == actual.renamed_date and expected.renamed_time == actual.renamed_time) config = self.jdbc_configs[database] # Run write pipeline with TestPipeline() as p: p.not_use_test_runner_api = True _ = ( p | beam.Create(test_rows).with_output_types(JdbcTestRow) | 'Write to jdbc' >> WriteToJdbc( table_name=table_name, driver_class_name=config['driver_class_name'], jdbc_url=config['jdbc_url'], username=config['username'], password=config['password'], classpath=config['classpath'], )) # Run read pipeline with custom schema with TestPipeline() as p: p.not_use_test_runner_api = True result = ( p | 'Read from jdbc with schema' >> ReadFromJdbc( table_name=table_name, driver_class_name=config['driver_class_name'], jdbc_url=config['jdbc_url'], username=config['username'], password=config['password'], classpath=config['classpath'], schema=CustomSchemaRow)) assert_that(result, equal_to(expected_rows, equals_fn=custom_row_equals)) @parameterized.expand(['postgres', 'mysql']) def test_xlang_jdbc_custom_statements(self, database): # Skip if container wasn't initialized if self.containers[database] is None: self.skipTest(f"{database} container could not be initialized") # Create a simple table for this test table_name = f"jdbc_custom_statements_{database}" with self.engines[database].begin() as connection: connection.execute( sqlalchemy.text( f"CREATE TABLE IF NOT EXISTS {table_name}" + "(id INTEGER, name VARCHAR(50), value DOUBLE PRECISION)")) test_rows = [ SimpleRow(1, "Item1", 10.5), SimpleRow(2, "Item2", 20.75), SimpleRow(3, "Item3", 30.25), SimpleRow(4, "Item4", 40.0), SimpleRow(-5, "Item5", 50.5) ] config = self.jdbc_configs[database] write_statement = f"INSERT INTO {table_name} (id, name, value) VALUES \ (?, ?, ?)" with TestPipeline() as p: p.not_use_test_runner_api = True _ = ( p | beam.Create(test_rows).with_output_types(SimpleRow) | 'Write with custom statement' >> WriteToJdbc( table_name="", driver_class_name=config['driver_class_name'], jdbc_url=config['jdbc_url'], username=config['username'], password=config['password'], classpath=config['classpath'], statement=write_statement)) # Schema inference fails when there is a WHERE clause, so we pass explicit # schema. read_query = f"SELECT id, name, value FROM {table_name} WHERE value > 25.0" expected_filtered_rows = [row for row in test_rows if row.value > 25.0] with TestPipeline() as p: p.not_use_test_runner_api = True result = ( p | 'Read with custom query' >> ReadFromJdbc( table_name="", driver_class_name=config['driver_class_name'], jdbc_url=config['jdbc_url'], username=config['username'], password=config['password'], classpath=config['classpath'], query=read_query, schema=SimpleRow)) assert_that(result, equal_to(expected_filtered_rows)) # JdbcIO#readWithPartitions requires custom queries to be passed as a # wrapped subquery to table_name. with TestPipeline() as p: p.not_use_test_runner_api = True result = ( p | 'Read with custom query' >> ReadFromJdbc( table_name=f"({read_query}) as subq", driver_class_name=config['driver_class_name'], jdbc_url=config['jdbc_url'], username=config['username'], password=config['password'], classpath=config['classpath'], partition_column="id", schema=SimpleRow)) assert_that(result, equal_to(expected_filtered_rows)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/xlang_kafkaio_it_test.py0000644000175100017510000002015115206113032025010 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration test for Python cross-language pipelines for Java KafkaIO.""" import contextlib import logging import os import socket import subprocess import sys import time import typing import unittest import uuid import pytest import apache_beam as beam from apache_beam.coders.coders import VarIntCoder from apache_beam.io.kafka import ReadFromKafka from apache_beam.io.kafka import WriteToKafka from apache_beam.metrics import Metrics from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.userstate import BagStateSpec from apache_beam.transforms.userstate import CombiningValueStateSpec from apache_beam.utils import subprocess_server NUM_RECORDS = 1000 class CollectingFn(beam.DoFn): BUFFER_STATE = BagStateSpec('buffer', VarIntCoder()) COUNT_STATE = CombiningValueStateSpec('count', sum) def process( self, element, buffer_state=beam.DoFn.StateParam(BUFFER_STATE), count_state=beam.DoFn.StateParam(COUNT_STATE)): value = int(element[1].decode()) buffer_state.add(value) count_state.add(1) count = count_state.read() if count >= NUM_RECORDS: yield sum(buffer_state.read()) count_state.clear() buffer_state.clear() class CrossLanguageKafkaIO(object): def __init__( self, bootstrap_servers, topic, null_key, expansion_service=None): self.bootstrap_servers = bootstrap_servers self.topic = topic self.null_key = null_key self.expansion_service = expansion_service self.sum_counter = Metrics.counter('source', 'elements_sum') def build_write_pipeline(self, pipeline): _ = ( pipeline | 'Generate' >> beam.Create(range(NUM_RECORDS)) # pylint: disable=bad-option-value | 'MakeKV' >> beam.Map( lambda x: (None if self.null_key else b'key', str(x).encode())). with_output_types(tuple[typing.Optional[bytes], bytes]) | 'WriteToKafka' >> WriteToKafka( producer_config={'bootstrap.servers': self.bootstrap_servers}, topic=self.topic, expansion_service=self.expansion_service)) def build_read_pipeline(self, pipeline, max_num_records=None): kafka_records = ( pipeline | 'ReadFromKafka' >> ReadFromKafka( consumer_config={ 'bootstrap.servers': self.bootstrap_servers, 'auto.offset.reset': 'earliest' }, topics=[self.topic], max_num_records=max_num_records, expansion_service=self.expansion_service)) if max_num_records: return kafka_records return ( kafka_records | 'CalculateSum' >> beam.ParDo(CollectingFn()) | 'SetSumCounter' >> beam.Map(self.sum_counter.inc)) def run_xlang_kafkaio(self, pipeline): self.build_write_pipeline(pipeline) self.build_read_pipeline(pipeline) pipeline.run(False) class CrossLanguageKafkaIOTest(unittest.TestCase): @unittest.skipUnless( os.environ.get('LOCAL_KAFKA_JAR'), "LOCAL_KAFKA_JAR environment var is not provided.") def test_local_kafkaio_populated_key(self): kafka_topic = 'xlang_kafkaio_test_populated_key_{}'.format(uuid.uuid4()) local_kafka_jar = os.environ.get('LOCAL_KAFKA_JAR') with self.local_kafka_service(local_kafka_jar) as kafka_port: bootstrap_servers = '{}:{}'.format( self.get_platform_localhost(), kafka_port) pipeline_creator = CrossLanguageKafkaIO( bootstrap_servers, kafka_topic, False) self.run_kafka_write(pipeline_creator) self.run_kafka_read(pipeline_creator, b'key') @unittest.skipUnless( os.environ.get('LOCAL_KAFKA_JAR'), "LOCAL_KAFKA_JAR environment var is not provided.") def test_local_kafkaio_null_key(self): kafka_topic = 'xlang_kafkaio_test_null_key_{}'.format(uuid.uuid4()) local_kafka_jar = os.environ.get('LOCAL_KAFKA_JAR') with self.local_kafka_service(local_kafka_jar) as kafka_port: bootstrap_servers = '{}:{}'.format( self.get_platform_localhost(), kafka_port) pipeline_creator = CrossLanguageKafkaIO( bootstrap_servers, kafka_topic, True) self.run_kafka_write(pipeline_creator) self.run_kafka_read(pipeline_creator, None) @pytest.mark.uses_io_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_PORT'), "EXPANSION_PORT environment var is not provided.") @unittest.skipUnless( os.environ.get('KAFKA_BOOTSTRAP_SERVER'), "KAFKA_BOOTSTRAP_SERVER environment var is not provided.") def test_hosted_kafkaio_populated_key(self): kafka_topic = 'xlang_kafkaio_test_populated_key_{}'.format(uuid.uuid4()) bootstrap_servers = os.environ.get('KAFKA_BOOTSTRAP_SERVER') pipeline_creator = CrossLanguageKafkaIO( bootstrap_servers, kafka_topic, False, 'localhost:%s' % os.environ.get('EXPANSION_PORT')) self.run_kafka_write(pipeline_creator) self.run_kafka_read(pipeline_creator, b'key') @pytest.mark.uses_io_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_PORT'), "EXPANSION_PORT environment var is not provided.") @unittest.skipUnless( os.environ.get('KAFKA_BOOTSTRAP_SERVER'), "KAFKA_BOOTSTRAP_SERVER environment var is not provided.") def test_hosted_kafkaio_null_key(self): kafka_topic = 'xlang_kafkaio_test_null_key_{}'.format(uuid.uuid4()) bootstrap_servers = os.environ.get('KAFKA_BOOTSTRAP_SERVER') pipeline_creator = CrossLanguageKafkaIO( bootstrap_servers, kafka_topic, True, 'localhost:%s' % os.environ.get('EXPANSION_PORT')) self.run_kafka_write(pipeline_creator) self.run_kafka_read(pipeline_creator, None) def run_kafka_write(self, pipeline_creator): with TestPipeline() as pipeline: pipeline.not_use_test_runner_api = True pipeline_creator.build_write_pipeline(pipeline) def run_kafka_read(self, pipeline_creator, expected_key): with TestPipeline() as pipeline: pipeline.not_use_test_runner_api = True result = pipeline_creator.build_read_pipeline(pipeline, NUM_RECORDS) assert_that( result, equal_to([(expected_key, str(i).encode()) for i in range(NUM_RECORDS)])) def get_platform_localhost(self): if sys.platform == 'darwin': return 'host.docker.internal' else: return 'localhost' def get_open_port(self): s = None try: s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) except: # pylint: disable=bare-except # Above call will fail for nodes that only support IPv6. s = socket.socket(socket.AF_INET6, socket.SOCK_STREAM) s.bind(('localhost', 0)) s.listen(1) port = s.getsockname()[1] s.close() return port @contextlib.contextmanager def local_kafka_service(self, local_kafka_jar_file): kafka_port = str(self.get_open_port()) zookeeper_port = str(self.get_open_port()) kafka_server = None try: kafka_server = subprocess.Popen([ subprocess_server.JavaHelper.get_java(), '-jar', local_kafka_jar_file, kafka_port, zookeeper_port ]) time.sleep(3) yield kafka_port finally: if kafka_server: kafka_server.kill() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/xlang_kafkaio_perf_test.py0000644000175100017510000001212015206113032025325 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import sys import apache_beam as beam from apache_beam.io import iobase from apache_beam.io import kafka from apache_beam.testing.load_tests.load_test import LoadTest from apache_beam.testing.load_tests.load_test import LoadTestOptions from apache_beam.testing.load_tests.load_test_metrics_utils import CountMessages from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime from apache_beam.testing.synthetic_pipeline import SyntheticSource from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.transforms.util import Reshuffle WRITE_NAMESPACE = 'write' READ_NAMESPACE = 'read' _LOGGER = logging.getLogger(__name__) class KafkaIOTestOptions(LoadTestOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--test_class', required=True, help='Test class to run.') parser.add_argument('--kafka_topic', required=True, help='Kafka topic.') parser.add_argument( '--bootstrap_servers', help='URL TO Kafka Bootstrap service.') parser.add_argument( '--read_timeout', type=int, required=True, help='Time to wait for the events to be processed by the read pipeline' ' (in seconds)') class KafkaIOPerfTest: """Performance test for cross-language Kafka IO pipeline.""" def run(self): write_test = _KafkaIOBatchWritePerfTest() read_test = _KafkaIOSDFReadPerfTest() write_test.run() read_test.run() class _KafkaIOBatchWritePerfTest(LoadTest): def __init__(self): super().__init__(WRITE_NAMESPACE) self.test_options = self.pipeline.get_pipeline_options().view_as( KafkaIOTestOptions) self.kafka_topic = self.test_options.kafka_topic # otherwise see 'ValueError: Unexpected DoFn type: beam:dofn:javasdk:0.1' self.pipeline.not_use_test_runner_api = True def test(self): _ = ( self.pipeline | 'Generate records' >> iobase.Read( SyntheticSource(self.parse_synthetic_source_options())) \ .with_output_types(tuple[bytes, bytes]) | 'Count records' >> beam.ParDo(CountMessages(self.metrics_namespace)) | 'Avoid Fusion' >> Reshuffle() | 'Measure time' >> beam.ParDo(MeasureTime(self.metrics_namespace)) | 'WriteToKafka' >> kafka.WriteToKafka( producer_config={ 'bootstrap.servers': self.test_options.bootstrap_servers }, topic=self.kafka_topic)) def cleanup(self): pass class _KafkaIOSDFReadPerfTest(LoadTest): def __init__(self): super().__init__(READ_NAMESPACE) self.test_options = self.pipeline.get_pipeline_options().view_as( KafkaIOTestOptions) self.timeout_ms = self.test_options.read_timeout * 1000 self.kafka_topic = self.test_options.kafka_topic # otherwise see 'ValueError: Unexpected DoFn type: beam:dofn:javasdk:0.1' self.pipeline.not_use_test_runner_api = True def test(self): _ = ( self.pipeline | 'ReadFromKafka' >> kafka.ReadFromKafka( consumer_config={ 'bootstrap.servers': self.test_options.bootstrap_servers, 'auto.offset.reset': 'earliest' }, topics=[self.kafka_topic]) | 'Count records' >> beam.ParDo(CountMessages(self.metrics_namespace)) | 'Measure time' >> beam.ParDo(MeasureTime(self.metrics_namespace))) def cleanup(self): total_messages = self._metrics_monitor.get_counter_metric( self.result, CountMessages.LABEL) expected_records = self.input_options['num_records'] assert total_messages >= expected_records, ( f"Expected at least {expected_records} messages, " f"but got {total_messages}") _LOGGER.info( "Read %d messages (expected: %d)", total_messages, expected_records) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) test_options = TestPipeline().get_pipeline_options().view_as( KafkaIOTestOptions) supported_test_classes = list( filter( lambda s: s.endswith('PerfTest') and not s.startswith('_'), dir(sys.modules[__name__]))) if test_options.test_class not in supported_test_classes: raise RuntimeError( f'Test {test_options.test_class} not found. ' 'Supported tests are {supported_test_classes}') getattr(sys.modules[__name__], test_options.test_class)().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/xlang_kinesisio_it_test.py0000644000175100017510000002666315206113032025416 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Integration test for Python cross-language pipelines for Java KinesisIO. If you want to run the tests on localstack then run it just with pipeline options. To test it on a real AWS account you need to pass some additional params, e.g.: python setup.py nosetests \ --tests=apache_beam.io.external.xlang_kinesisio_it_test \ --test-pipeline-options=" --use_real_aws --aws_kinesis_stream= --aws_access_key= --aws_secret_key= --aws_region= --runner=FlinkRunner" """ # pytype: skip-file import argparse import logging import time import unittest import uuid import apache_beam as beam from apache_beam.io.kinesis import InitialPositionInStream from apache_beam.io.kinesis import ReadDataFromKinesis from apache_beam.io.kinesis import WatermarkPolicy from apache_beam.io.kinesis import WriteToKinesis from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: import boto3 except ImportError: boto3 = None try: from testcontainers.core.container import DockerContainer except ImportError: DockerContainer = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports LOCALSTACK_VERSION = '3.8.1' NUM_RECORDS = 10 MAX_READ_TIME = 5 * 60 * 1000 # 5min NOW_SECONDS = time.time() NOW_MILLIS = NOW_SECONDS * 1000 REQUEST_RECORDS_LIMIT = 1000 RECORD = b'record' + str(uuid.uuid4()).encode() @unittest.skipUnless(DockerContainer, 'testcontainers is not installed.') @unittest.skipUnless(boto3, 'boto3 is not installed.') @unittest.skipUnless( TestPipeline().get_pipeline_options().view_as(StandardOptions).runner, 'Do not run this test on precommit suites.') class CrossLanguageKinesisIOTest(unittest.TestCase): @unittest.skipUnless( TestPipeline().get_option('aws_kinesis_stream'), 'Cannot test on real aws without pipeline options provided') def test_kinesis_io_roundtrip(self): # TODO: enable this test for localstack once # https://github.com/apache/beam/issues/20416 is resolved self.run_kinesis_write() self.run_kinesis_read() @unittest.skipIf( TestPipeline().get_option('aws_kinesis_stream'), 'Do not test on localstack when pipeline options were provided') def test_kinesis_write(self): # TODO: remove this test once # https://github.com/apache/beam/issues/20416 is resolved self.run_kinesis_write() self.run_kinesis_write_with_aggregation() records = self.kinesis_helper.read_from_stream(self.aws_kinesis_stream) self.assertEqual( sorted(records), sorted([RECORD + str(i).encode() for i in range(NUM_RECORDS)])) def run_kinesis_write(self): with TestPipeline(options=PipelineOptions(self.pipeline_args)) as p: p.not_use_test_runner_api = True _ = ( p | 'Impulse' >> beam.Impulse() | 'Generate' >> beam.FlatMap(lambda x: range(NUM_RECORDS)) # pylint: disable=bad-option-value | 'Map to bytes' >> beam.Map(lambda x: RECORD + str(x).encode()).with_output_types(bytes) | 'WriteToKinesis' >> WriteToKinesis( stream_name=self.aws_kinesis_stream, aws_access_key=self.aws_access_key, aws_secret_key=self.aws_secret_key, region=self.aws_region, service_endpoint=self.aws_service_endpoint, verify_certificate=(not self.use_localstack), partition_key='1')) def run_kinesis_write_with_aggregation(self): with TestPipeline(options=PipelineOptions(self.pipeline_args)) as p: p.not_use_test_runner_api = True _ = ( p | 'Impulse' >> beam.Impulse() | 'Generate' >> beam.FlatMap(lambda x: range(NUM_RECORDS)) # pylint: disable=bad-option-value | 'Map to bytes' >> beam.Map(lambda x: RECORD + str(x).encode()).with_output_types(bytes) | 'WriteToKinesis' >> WriteToKinesis( stream_name=self.aws_kinesis_stream, aws_access_key=self.aws_access_key, aws_secret_key=self.aws_secret_key, region=self.aws_region, service_endpoint=self.aws_service_endpoint, verify_certificate=(not self.use_localstack), partition_key='1', aggregation_enabled=True)) def run_kinesis_read(self): records = [RECORD + str(i).encode() for i in range(NUM_RECORDS)] with TestPipeline(options=PipelineOptions(self.pipeline_args)) as p: result = ( p | 'ReadFromKinesis' >> ReadDataFromKinesis( stream_name=self.aws_kinesis_stream, aws_access_key=self.aws_access_key, aws_secret_key=self.aws_secret_key, region=self.aws_region, service_endpoint=self.aws_service_endpoint, verify_certificate=not self.use_localstack, max_num_records=NUM_RECORDS, max_read_time=MAX_READ_TIME, request_records_limit=REQUEST_RECORDS_LIMIT, watermark_policy=WatermarkPolicy.ARRIVAL_TIME, watermark_idle_duration_threshold=MAX_READ_TIME, initial_position_in_stream=InitialPositionInStream.AT_TIMESTAMP, initial_timestamp_in_stream=NOW_MILLIS, ).with_output_types(bytes)) assert_that(result, equal_to(records)) def set_localstack(self): self.localstack = DockerContainer('localstack/localstack:{}' .format(LOCALSTACK_VERSION))\ .with_bind_ports(4566, 4566) for i in range(4510, 4560): self.localstack = self.localstack.with_bind_ports(i, i) # Repeat if ReadTimeout is raised. for i in range(4): try: self.localstack.start() break except Exception as e: # pylint: disable=bare-except if i == 3: logging.error('Could not initialize localstack container') raise e self.aws_service_endpoint = 'https://{}:{}'.format( self.localstack.get_container_host_ip(), self.localstack.get_exposed_port('4566'), ) def setUp(self): parser = argparse.ArgumentParser() parser.add_argument( '--aws_kinesis_stream', default='beam_kinesis_xlang', help='Kinesis stream name', ) parser.add_argument( '--aws_access_key', default='accesskey', help=('Aws access key'), ) parser.add_argument( '--aws_secret_key', default='secretkey', help='Aws secret key', ) parser.add_argument( '--aws_region', default='us-east-1', help='Aws region', ) parser.add_argument( '--aws_service_endpoint', default=None, help='Url to external aws endpoint', ) parser.add_argument( '--use_real_aws', default=False, dest='use_real_aws', action='store_true', help='Flag whether to use real aws for the tests purpose', ) parser.add_argument( '--expansion_service', help='Url to externally launched expansion service.', ) pipeline = TestPipeline() argv = pipeline.get_full_options_as_args() known_args, self.pipeline_args = parser.parse_known_args(argv) self.aws_kinesis_stream = known_args.aws_kinesis_stream self.aws_access_key = known_args.aws_access_key self.aws_secret_key = known_args.aws_secret_key self.aws_region = known_args.aws_region self.aws_service_endpoint = known_args.aws_service_endpoint self.use_localstack = not known_args.use_real_aws self.expansion_service = known_args.expansion_service if self.use_localstack: self.set_localstack() self.kinesis_helper = KinesisHelper( self.aws_access_key, self.aws_secret_key, self.aws_region, self.aws_service_endpoint.replace('https', 'http') if self.aws_service_endpoint else None, ) if self.use_localstack: self.kinesis_helper.create_stream(self.aws_kinesis_stream) def tearDown(self): if self.use_localstack: self.kinesis_helper.delete_stream(self.aws_kinesis_stream) try: self.localstack.stop() except: # pylint: disable=bare-except logging.error('Could not stop the localstack container') class KinesisHelper: def __init__(self, access_key, secret_key, region, service_endpoint): self.kinesis_client = boto3.client( service_name='kinesis', region_name=region, endpoint_url=service_endpoint, aws_access_key_id=access_key, aws_secret_access_key=secret_key, ) def create_stream(self, stream_name): # localstack could not have initialized in the container yet so repeat retries = 10 for i in range(retries): try: self.kinesis_client.create_stream( StreamName=stream_name, ShardCount=1, ) time.sleep(2) break except Exception as e: if i == retries - 1: logging.error('Could not create kinesis stream') raise e # Wait for the stream to be active self.get_first_shard_id(stream_name) def delete_stream(self, stream_name): self.kinesis_client.delete_stream( StreamName=stream_name, EnforceConsumerDeletion=True, ) def get_first_shard_id(self, stream_name): retries = 10 stream = self.kinesis_client.describe_stream(StreamName=stream_name) for i in range(retries): if stream['StreamDescription']['StreamStatus'] == 'ACTIVE': break time.sleep(2) if i == retries - 1: logging.error('Could not initialize kinesis stream') raise RuntimeError( "Unable to initialize Kinesis Stream %s. Status: %s", stream['StreamDescription']['StreamName'], stream['StreamDescription']['StreamStatus']) stream = self.kinesis_client.describe_stream(StreamName=stream_name) return stream['StreamDescription']['Shards'][0]['ShardId'] def read_from_stream(self, stream_name): shard_id = self.get_first_shard_id(stream_name) shard_iterator = self.kinesis_client.get_shard_iterator( StreamName=stream_name, ShardId=shard_id, ShardIteratorType=InitialPositionInStream.AT_TIMESTAMP, Timestamp=str(NOW_SECONDS), ) result = self.kinesis_client.get_records( ShardIterator=shard_iterator['ShardIterator'], Limit=NUM_RECORDS, ) return [record['Data'] for record in result['Records']] if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/xlang_parquetio_test.py0000644000175100017510000000573615206113032024734 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for cross-language parquet io read/write.""" # pytype: skip-file import logging import os import re import unittest import apache_beam as beam from apache_beam import coders from apache_beam.coders.avro_record import AvroRecord from apache_beam.options.pipeline_options import DebugOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.transforms.external import ImplicitSchemaPayloadBuilder PARQUET_WRITE_URN = "beam:transforms:xlang:test:parquet_write" # TODO: enable test_xlang_parquetio_write after fixing BEAM-10507 # @pytest.mark.uses_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_JAR'), "EXPANSION_JAR environment variable is not set.") @unittest.skipUnless( os.environ.get('EXPANSION_PORT'), "EXPANSION_PORT environment var is not provided.") class XlangParquetIOTest(unittest.TestCase): # TODO: add verification for the file written by external transform # after fixing BEAM-7612 def test_xlang_parquetio_write(self): expansion_jar = os.environ.get('EXPANSION_JAR') port = os.environ.get('EXPANSION_PORT') address = 'localhost:%s' % port try: with TestPipeline() as p: p.get_pipeline_options().view_as(DebugOptions).experiments.append( 'jar_packages=' + expansion_jar) p.not_use_test_runner_api = True _ = p \ | beam.Create([ AvroRecord({"name": "abc"}), AvroRecord({"name": "def"}), AvroRecord({"name": "ghi"})]) \ | beam.ExternalTransform( PARQUET_WRITE_URN, ImplicitSchemaPayloadBuilder({'data': '/tmp/test.parquet'}), address) except RuntimeError as e: if re.search(PARQUET_WRITE_URN, str(e)): print("looks like URN not implemented in expansion service, skipping.") else: raise e class AvroTestCoder(coders.AvroGenericCoder): SCHEMA = """ { "type": "record", "name": "testrecord", "fields": [ {"name": "name", "type": "string"} ] } """ def __init__(self): super().__init__(self.SCHEMA) coders.registry.register_coder(AvroRecord, AvroTestCoder) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/external/xlang_snowflakeio_it_test.py0000644000175100017510000002160415206113032025730 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Integration test for cross-language snowflake io operations. Example run: python setup.py nosetests --tests=apache_beam.io.external.snowflake_test \ --test-pipeline-options=" --server_name= --username= --password= --private_key_path= --raw_private_key= --private_key_passphrase= --o_auth_token= --staging_bucket_name= --storage_integration_name= --database= --schema= --role= --warehouse= --table= --runner=FlinkRunner" """ # pytype: skip-file import argparse import binascii import logging import unittest from typing import ByteString from typing import NamedTuple import apache_beam as beam from apache_beam import coders from apache_beam.io.snowflake import CreateDisposition from apache_beam.io.snowflake import ReadFromSnowflake from apache_beam.io.snowflake import WriteDisposition from apache_beam.io.snowflake import WriteToSnowflake from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem except ImportError: GCSFileSystem = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports SCHEMA_STRING = """ {"schema":[ {"dataType":{"type":"integer","precision":38,"scale":0},"name":"number_column","nullable":false}, {"dataType":{"type":"boolean"},"name":"boolean_column","nullable":false}, {"dataType":{"type":"binary","size":100},"name":"bytes_column","nullable":true} ]} """ TestRow = NamedTuple( 'TestRow', [ ('number_column', int), ('boolean_column', bool), ('bytes_column', ByteString), ]) coders.registry.register_coder(TestRow, coders.RowCoder) NUM_RECORDS = 100 @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed') @unittest.skipIf( TestPipeline().get_option('server_name') is None, 'Snowflake IT test requires external configuration to be run.') class SnowflakeTest(unittest.TestCase): def test_snowflake_write_read(self): self.run_write() self.run_read() def run_write(self): def user_data_mapper(test_row): return [ str(test_row.number_column).encode('utf-8'), str(test_row.boolean_column).encode('utf-8'), binascii.hexlify(test_row.bytes_column), ] with TestPipeline(options=PipelineOptions(self.pipeline_args)) as p: p.not_use_test_runner_api = True _ = ( p | 'Impulse' >> beam.Impulse() | 'Generate' >> beam.FlatMap(lambda x: range(NUM_RECORDS)) # pylint: disable=bad-option-value | 'Map to TestRow' >> beam.Map( lambda num: TestRow( num, num % 2 == 0, b"test" + str(num).encode())) | WriteToSnowflake( server_name=self.server_name, username=self.username, password=self.password, o_auth_token=self.o_auth_token, private_key_path=self.private_key_path, raw_private_key=self.raw_private_key, private_key_passphrase=self.private_key_passphrase, schema=self.schema, database=self.database, role=self.role, warehouse=self.warehouse, staging_bucket_name=self.staging_bucket_name, storage_integration_name=self.storage_integration_name, create_disposition=CreateDisposition.CREATE_IF_NEEDED, write_disposition=WriteDisposition.TRUNCATE, table_schema=SCHEMA_STRING, user_data_mapper=user_data_mapper, table=self.table, query=None, expansion_service=self.expansion_service, )) def run_read(self): def csv_mapper(bytes_array): return TestRow( int(bytes_array[0]), bytes_array[1] == b'true', binascii.unhexlify(bytes_array[2])) with TestPipeline(options=PipelineOptions(self.pipeline_args)) as p: result = ( p | ReadFromSnowflake( server_name=self.server_name, username=self.username, password=self.password, o_auth_token=self.o_auth_token, private_key_path=self.private_key_path, raw_private_key=self.raw_private_key, private_key_passphrase=self.private_key_passphrase, schema=self.schema, database=self.database, role=self.role, warehouse=self.warehouse, staging_bucket_name=self.staging_bucket_name, storage_integration_name=self.storage_integration_name, csv_mapper=csv_mapper, table=self.table, query=None, expansion_service=self.expansion_service, ).with_output_types(TestRow)) assert_that( result, equal_to([ TestRow(i, i % 2 == 0, b'test' + str(i).encode()) for i in range(NUM_RECORDS) ])) @classmethod def tearDownClass(cls): GCSFileSystem(pipeline_options=PipelineOptions()) \ .delete([cls.staging_bucket_name]) @classmethod def setUpClass(cls): parser = argparse.ArgumentParser() parser.add_argument( '--server_name', required=True, help=( 'Snowflake server name of the form ' 'https://.snowflakecomputing.com'), ) parser.add_argument( '--username', help='Snowflake username', ) parser.add_argument( '--password', help='Snowflake password', ) parser.add_argument( '--private_key_path', help='Path to private key', ) parser.add_argument( '--raw_private_key', help='Raw private key', ) parser.add_argument( '--private_key_passphrase', help='Password to private key', ) parser.add_argument( '--o_auth_token', help='OAuth token', ) parser.add_argument( '--staging_bucket_name', required=True, help='GCP staging bucket name (must end with backslash)', ) parser.add_argument( '--storage_integration_name', required=True, help='Snowflake integration name', ) parser.add_argument( '--database', required=True, help='Snowflake database name', ) parser.add_argument( '--schema', required=True, help='Snowflake schema name', ) parser.add_argument( '--table', required=True, help='Snowflake table name', ) parser.add_argument( '--role', help='Snowflake role', ) parser.add_argument( '--warehouse', help='Snowflake warehouse name', ) parser.add_argument( '--expansion_service', help='Url to externally launched expansion service.', ) pipeline = TestPipeline() argv = pipeline.get_full_options_as_args() known_args, cls.pipeline_args = parser.parse_known_args(argv) cls.server_name = known_args.server_name cls.database = known_args.database cls.schema = known_args.schema cls.table = known_args.table cls.username = known_args.username cls.password = known_args.password cls.private_key_path = known_args.private_key_path cls.raw_private_key = known_args.raw_private_key cls.private_key_passphrase = known_args.private_key_passphrase cls.o_auth_token = known_args.o_auth_token cls.staging_bucket_name = known_args.staging_bucket_name cls.storage_integration_name = known_args.storage_integration_name cls.role = known_args.role cls.warehouse = known_args.warehouse cls.expansion_service = known_args.expansion_service if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/filebasedio_perf_test.py0000644000175100017510000001515615206113032023167 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Performance tests for file based io connectors.""" import logging import sys import uuid import apache_beam as beam from apache_beam import typehints from apache_beam.io.filesystems import FileSystems from apache_beam.io.iobase import Read from apache_beam.io.textio import ReadFromText from apache_beam.io.textio import WriteToText from apache_beam.testing.load_tests.load_test import LoadTest from apache_beam.testing.load_tests.load_test import LoadTestOptions from apache_beam.testing.load_tests.load_test_metrics_utils import CountMessages from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime from apache_beam.testing.synthetic_pipeline import SyntheticSource from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.util import Reshuffle WRITE_NAMESPACE = 'write' READ_NAMESPACE = 'read' _LOGGER = logging.getLogger(__name__) class FileBasedIOTestOptions(LoadTestOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--test_class', required=True, help='Test class to run.') parser.add_argument( '--filename_prefix', required=True, help='Destination prefix for files generated by the test.') parser.add_argument( '--compression_type', default='auto', help='File compression type for writing and reading test files.') parser.add_argument( '--number_of_shards', type=int, default=0, help='Number of files this test will create during the write phase.') parser.add_argument( '--dataset_size', type=int, help='Size of data saved on the target filesystem (bytes).') @typehints.with_output_types(bytes) @typehints.with_input_types(tuple[bytes, bytes]) class SyntheticRecordToStrFn(beam.DoFn): """ A DoFn that convert key-value bytes from synthetic source to string record. It uses base64 to convert random bytes emitted from the synthetic source. Therefore, every 3 bytes give 4 bytes long ascii characters. Output length = 4(ceil[len(key)/3] + ceil[len(value)/3]) + 1 """ def process(self, element): import base64 yield base64.b64encode(element[0]) + b' ' + base64.b64encode(element[1]) class CreateFolderFn(beam.DoFn): """Create folder at pipeline runtime.""" def __init__(self, folder): self.folder = folder def process(self, element): from apache_beam.io.filesystems import FileSystems # pylint: disable=reimported filesystem = FileSystems.get_filesystem(self.folder) if filesystem.has_dirs() and not filesystem.exists(self.folder): filesystem.mkdirs(self.folder) class TextIOPerfTest: def run(self): write_test = _TextIOWritePerfTest(need_cleanup=False) read_test = _TextIOReadPerfTest(input_folder=write_test.output_folder) write_test.run() read_test.run() class _TextIOWritePerfTest(LoadTest): def __init__(self, need_cleanup=True): super().__init__(WRITE_NAMESPACE) self.need_cleanup = need_cleanup self.test_options = self.pipeline.get_pipeline_options().view_as( FileBasedIOTestOptions) self.output_folder = FileSystems.join( self.test_options.filename_prefix, str(uuid.uuid4())) def test(self): # first makedir if needed _ = ( self.pipeline | beam.Impulse() | beam.ParDo(CreateFolderFn(self.output_folder))) # write to text _ = ( self.pipeline | 'Produce rows' >> Read( SyntheticSource(self.parse_synthetic_source_options())) | 'Count records' >> beam.ParDo(CountMessages(self.metrics_namespace)) | 'Format' >> beam.ParDo(SyntheticRecordToStrFn()) | 'Avoid Fusion' >> Reshuffle() | 'Measure time' >> beam.ParDo(MeasureTime(self.metrics_namespace)) | 'Write Text' >> WriteToText( file_path_prefix=FileSystems.join(self.output_folder, 'test'), compression_type=self.test_options.compression_type, num_shards=self.test_options.number_of_shards)) def cleanup(self): if not self.need_cleanup: return try: FileSystems.delete([self.output_folder]) except IOError: # may not have delete permission, just raise a warning _LOGGER.warning( 'Unable to delete file %s during cleanup.', self.output_folder) class _TextIOReadPerfTest(LoadTest): def __init__(self, input_folder): super().__init__(READ_NAMESPACE) self.test_options = self.pipeline.get_pipeline_options().view_as( FileBasedIOTestOptions) self.input_folder = input_folder def test(self): output = ( self.pipeline | 'Read from text' >> ReadFromText(file_pattern=FileSystems.join(self.input_folder, '*')) | 'Count records' >> beam.ParDo(CountMessages(self.metrics_namespace)) | 'Measure time' >> beam.ParDo(MeasureTime(self.metrics_namespace)) | 'Count' >> beam.combiners.Count.Globally()) assert_that(output, equal_to([self.input_options['num_records']])) def cleanup(self): try: #FileSystems.delete([self.input_folder]) pass except IOError: # may not have delete permission, just raise a warning _LOGGER.warning( 'Unable to delete file %s during cleanup.', self.input_folder) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) test_options = TestPipeline().get_pipeline_options().view_as( FileBasedIOTestOptions) supported_test_classes = list( filter( lambda s: s.endswith('PerfTest') and not s.startswith('_'), dir(sys.modules[__name__]))) if test_options.test_class not in supported_test_classes: raise RuntimeError( f'Test {test_options.test_class} not found. ' 'Supported tests are {supported_test_classes}') getattr(sys.modules[__name__], test_options.test_class)().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/filebasedsink.py0000644000175100017510000005103515206113032021445 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """File-based sink.""" # pytype: skip-file import logging import os import re import time import uuid from apache_beam.internal import util from apache_beam.io import iobase from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystems import FileSystems from apache_beam.options.value_provider import StaticValueProvider from apache_beam.options.value_provider import ValueProvider from apache_beam.options.value_provider import check_accessible from apache_beam.transforms import window from apache_beam.transforms.display import DisplayDataItem DEFAULT_SHARD_NAME_TEMPLATE = '-SSSSS-of-NNNNN' DEFAULT_WINDOW_SHARD_NAME_TEMPLATE = '-W-SSSSS-of-NNNNN' DEFAULT_TRIGGERING_FREQUENCY = 0 __all__ = ['FileBasedSink'] _LOGGER = logging.getLogger(__name__) class FileBasedSink(iobase.Sink): """A sink to a GCS or local files. To implement a file-based sink, extend this class and override either :meth:`.write_record()` or :meth:`.write_encoded_record()`. If needed, also overwrite :meth:`.open()` and/or :meth:`.close()` to customize the file handling or write headers and footers. The output of this write is a :class:`~apache_beam.pvalue.PCollection` of all written shards. """ # Max number of threads to be used for renaming. _MAX_RENAME_THREADS = 64 __hash__ = None # type: ignore[assignment] def __init__( self, file_path_prefix, coder, file_name_suffix='', num_shards=0, shard_name_template=None, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO, *, max_records_per_shard=None, max_bytes_per_shard=None, skip_if_empty=False, convert_fn=None, triggering_frequency=None): """ Raises: TypeError: if file path parameters are not a :class:`str` or :class:`~apache_beam.options.value_provider.ValueProvider`, or if **compression_type** is not member of :class:`~apache_beam.io.filesystem.CompressionTypes`. ValueError: if **shard_name_template** is not of expected format. """ if not isinstance(file_path_prefix, (str, ValueProvider)): raise TypeError( 'file_path_prefix must be a string or ValueProvider;' 'got %r instead' % file_path_prefix) if not isinstance(file_name_suffix, (str, ValueProvider)): raise TypeError( 'file_name_suffix must be a string or ValueProvider;' 'got %r instead' % file_name_suffix) if not CompressionTypes.is_valid_compression_type(compression_type): raise TypeError( 'compression_type must be CompressionType object but ' 'was %s' % type(compression_type)) if shard_name_template is None: shard_name_template = DEFAULT_SHARD_NAME_TEMPLATE elif shard_name_template == '': num_shards = 1 if triggering_frequency is None: triggering_frequency = DEFAULT_TRIGGERING_FREQUENCY if isinstance(file_path_prefix, str): file_path_prefix = StaticValueProvider(str, file_path_prefix) if isinstance(file_name_suffix, str): file_name_suffix = StaticValueProvider(str, file_name_suffix) self.file_path_prefix = file_path_prefix self.file_name_suffix = file_name_suffix self.num_shards = num_shards self.coder = coder self.shard_name_template = shard_name_template self.shard_name_format = self._template_to_format(shard_name_template) self.shard_name_glob_format = self._template_to_glob_format( shard_name_template) self.compression_type = compression_type self.mime_type = mime_type self.max_records_per_shard = max_records_per_shard self.max_bytes_per_shard = max_bytes_per_shard self.skip_if_empty = skip_if_empty self.convert_fn = convert_fn self.triggering_frequency = triggering_frequency def display_data(self): return { 'shards': DisplayDataItem(self.num_shards, label='Number of Shards').drop_if_default(0), 'compression': DisplayDataItem(str(self.compression_type)), 'file_pattern': DisplayDataItem( '{}{}{}'.format( self.file_path_prefix, self.shard_name_format, self.file_name_suffix), label='File Pattern') } @check_accessible(['file_path_prefix']) def open(self, temp_path): """Opens ``temp_path``, returning an opaque file handle object. The returned file handle is passed to ``write_[encoded_]record`` and ``close``. """ writer = FileSystems.create( temp_path, self.mime_type, self.compression_type) if self.max_bytes_per_shard: self.byte_counter = _ByteCountingWriter(writer) return self.byte_counter else: return writer def write_record(self, file_handle, value): """Writes a single record go the file handle returned by ``open()``. By default, calls ``write_encoded_record`` after encoding the record with this sink's Coder. """ self.write_encoded_record(file_handle, self.coder.encode(value)) def write_encoded_record(self, file_handle, encoded_value): """Writes a single encoded record to the file handle returned by ``open()``. """ raise NotImplementedError def close(self, file_handle): """Finalize and close the file handle returned from ``open()``. Called after all records are written. By default, calls ``file_handle.close()`` iff it is not None. """ if file_handle is not None: file_handle.close() @check_accessible(['file_path_prefix', 'file_name_suffix']) def initialize_write(self): file_path_prefix = self.file_path_prefix.get() tmp_dir = self._create_temp_dir(file_path_prefix) FileSystems.mkdirs(tmp_dir) return tmp_dir def _create_temp_dir(self, file_path_prefix): base_path, last_component = FileSystems.split(file_path_prefix) if not last_component: # Trying to re-split the base_path to check if it's a root. new_base_path, _ = FileSystems.split(base_path) if base_path == new_base_path: raise ValueError( 'Cannot create a temporary directory for root path ' 'prefix %s. Please specify a file path prefix with ' 'at least two components.' % file_path_prefix) path_components = [ base_path, 'beam-temp-' + last_component + '-' + uuid.uuid1().hex ] return FileSystems.join(*path_components) @check_accessible(['file_path_prefix', 'file_name_suffix']) def open_writer(self, init_result, uid): # A proper suffix is needed for AUTO compression detection. # We also ensure there will be no collisions with uid and a # (possibly unsharded) file_path_prefix and a (possibly empty) # file_name_suffix. from apache_beam.pvalue import EmptySideInput # Handle case where init_result is EmptySideInput (empty collection) # TODO: https://github.com/apache/beam/issues/36563 for Prism if isinstance(init_result, EmptySideInput): # Fall back to creating a temporary directory based on file_path_prefix _LOGGER.warning( 'Initialization result collection was empty, falling back to ' 'creating temporary directory. This may indicate an issue with ' 'the pipeline initialization phase.') file_path_prefix = self.file_path_prefix.get() init_result = self._create_temp_dir(file_path_prefix) FileSystems.mkdirs(init_result) file_path_prefix = self.file_path_prefix.get() file_name_suffix = self.file_name_suffix.get() suffix = ('.' + os.path.basename(file_path_prefix) + file_name_suffix) writer_path = FileSystems.join(init_result, uid) + suffix return FileBasedSinkWriter(self, writer_path) @check_accessible(['file_path_prefix', 'file_name_suffix']) def _get_final_name(self, shard_num, num_shards, w=None): if w is None or isinstance(w, window.GlobalWindow): window_utc = None else: window_utc = ( '[' + w.start.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S") + ', ' + w.end.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S") + ')') return ''.join([ self.file_path_prefix.get(), self.shard_name_format % dict( shard_num=shard_num, num_shards=num_shards, uuid=(uuid.uuid4()), window=w, window_utc=window_utc), self.file_name_suffix.get() ]) @check_accessible(['file_path_prefix', 'file_name_suffix']) def _get_final_name_glob(self, num_shards, w=None): if w is None or isinstance(w, window.GlobalWindow): window_utc = None else: window_utc = ( '[' + w.start.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S") + ', ' + w.end.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S") + ')') return ''.join([ self.file_path_prefix.get(), self.shard_name_glob_format % dict( num_shards=num_shards, uuid=(uuid.uuid4()), window=w, window_utc=window_utc), self.file_name_suffix.get() ]) def pre_finalize(self, init_result, writer_results): num_shards = len(list(writer_results)) dst_glob = self._get_final_name_glob(num_shards) dst_glob_files = [ file_metadata.path for mr in FileSystems.match([dst_glob]) for file_metadata in mr.metadata_list ] if dst_glob_files: _LOGGER.warning( 'Deleting %d existing files in target path matching: %s', len(dst_glob_files), self.shard_name_glob_format) FileSystems.delete(dst_glob_files) def pre_finalize_windowed(self, init_result, writer_results, window=None): num_shards = len(list(writer_results)) dst_glob = self._get_final_name_glob(num_shards, window) dst_glob_files = [ file_metadata.path for mr in FileSystems.match([dst_glob]) for file_metadata in mr.metadata_list ] if dst_glob_files: _LOGGER.warning( 'Deleting %d existing files in target path matching: %s', len(dst_glob_files), self.shard_name_glob_format) FileSystems.delete(dst_glob_files) def _check_state_for_finalize_write( self, writer_results, num_shards, window=None): """Checks writer output files' states. Returns: src_files, dst_files: Lists of files to rename. For each i, finalize_write should rename(src_files[i], dst_files[i]). delete_files: Src files to delete. These could be leftovers from an incomplete (non-atomic) rename operation. num_skipped: Tally of writer results files already renamed, such as from a previous run of finalize_write(). """ if not writer_results: return [], [], [], 0 src_glob = FileSystems.join(FileSystems.split(writer_results[0])[0], '*') dst_glob = self._get_final_name_glob(num_shards, window) src_glob_files = set( file_metadata.path for mr in FileSystems.match([src_glob]) for file_metadata in mr.metadata_list) dst_glob_files = set( file_metadata.path for mr in FileSystems.match([dst_glob]) for file_metadata in mr.metadata_list) src_files = [] dst_files = [] delete_files = [] num_skipped = 0 for shard_num, src in enumerate(writer_results): final_name = self._get_final_name(shard_num, num_shards, window) dst = final_name src_exists = src in src_glob_files dst_exists = dst in dst_glob_files if not src_exists and not dst_exists: raise BeamIOError( 'src and dst files do not exist. src: %s, dst: %s' % (src, dst)) if not src_exists and dst_exists: _LOGGER.debug('src: %s -> dst: %s already renamed, skipping', src, dst) num_skipped += 1 continue if (src_exists and dst_exists and FileSystems.checksum(src) == FileSystems.checksum(dst)): _LOGGER.debug('src: %s == dst: %s, deleting src', src, dst) delete_files.append(src) continue src_files.append(src) dst_files.append(dst) self._report_sink_lineage(dst_glob, dst_files) return src_files, dst_files, delete_files, num_skipped def _report_sink_lineage(self, dst_glob, dst_files): """ Report sink Lineage. Report every file if number of files no more than 10, otherwise only report glob. """ # There is rollup at the higher level, but this loses glob information. # Better to report multiple globs than just the parent directory. if len(dst_files) <= 10: for dst in dst_files: FileSystems.report_sink_lineage(dst) else: FileSystems.report_sink_lineage(dst_glob) @check_accessible(['file_path_prefix']) def finalize_write( self, init_result, writer_results, unused_pre_finalize_results): #Legacy finalize_write now has shares the implementation with #finalize_windowed_write when window is None. return self.finalize_windowed_write( init_result, writer_results, unused_pre_finalize_results, None) @check_accessible(['file_path_prefix']) def finalize_windowed_write( self, init_result, writer_results, unused_pre_finalize_results, w=None): writer_results = sorted(writer_results) num_shards = len(writer_results) src_files, dst_files, delete_files, num_skipped = ( self._check_state_for_finalize_write(writer_results, num_shards, w)) num_skipped += len(delete_files) FileSystems.delete(delete_files) num_shards_to_finalize = len(src_files) min_threads = min(num_shards_to_finalize, FileBasedSink._MAX_RENAME_THREADS) num_threads = max(1, min_threads) chunk_size = FileSystems.get_chunk_size(self.file_path_prefix.get()) source_file_batch = [ src_files[i:i + chunk_size] for i in range(0, len(src_files), chunk_size) ] destination_file_batch = [ dst_files[i:i + chunk_size] for i in range(0, len(dst_files), chunk_size) ] if num_shards_to_finalize: start_time = time.time() def _rename_batch(batch): """_rename_batch executes batch rename operations.""" source_files, destination_files = batch exceptions = [] try: FileSystems.rename(source_files, destination_files) return exceptions except BeamIOError as exp: if exp.exception_details is None: raise for (src, dst), exception in exp.exception_details.items(): if exception: _LOGGER.error( ('Exception in _rename_batch. src: %s, ' 'dst: %s, err: %s'), src, dst, exception) exceptions.append(exception) else: _LOGGER.debug('Rename successful: %s -> %s', src, dst) return exceptions if w is None or isinstance(w, window.GlobalWindow): # bounded input was handled by finalize_write legacy method # the implementation here should be called by finalize_write # Use a thread pool for renaming operations. exception_batches = util.run_using_threadpool( _rename_batch, list(zip(source_file_batch, destination_file_batch)), num_threads) all_exceptions = [ e for exception_batch in exception_batches for e in exception_batch ] if all_exceptions: raise Exception( 'Encountered exceptions in finalize_write: %s' % all_exceptions) yield from dst_files else: # unbounded input batch = list([src_files, dst_files]) exception_batches = _rename_batch(batch) all_exceptions = [ e for exception_batch in exception_batches for e in exception_batch ] if all_exceptions: raise Exception( 'Encountered exceptions in finalize_write: %s' % all_exceptions) yield from dst_files _LOGGER.info( 'Renamed %d shards in %.2f seconds.', num_shards_to_finalize, time.time() - start_time) else: _LOGGER.warning( 'No shards found to finalize. num_shards: %d, skipped: %d', num_shards, num_skipped) try: FileSystems.delete([init_result]) except IOError: # This error is not serious, we simply log it. _LOGGER.info('Unable to delete file: %s', init_result) @staticmethod def _template_replace_window(shard_name_template): match = re.search('W+', shard_name_template) if match: shard_name_template = shard_name_template.replace( match.group(0), '%%(window)0%ds' % len(match.group(0))) match = re.search('V+', shard_name_template) if match: shard_name_template = shard_name_template.replace( match.group(0), '%%(window_utc)0%ds' % len(match.group(0))) return shard_name_template @staticmethod def _template_replace_uuid(shard_name_template): match = re.search('U+', shard_name_template) if match: shard_name_template = shard_name_template.replace( match.group(0), '%%(uuid)0%dd' % len(match.group(0))) return shard_name_template @staticmethod def _template_replace_num_shards(shard_name_template): match = re.search('N+', shard_name_template) if match: shard_name_template = shard_name_template.replace( match.group(0), '%%(num_shards)0%dd' % len(match.group(0))) return shard_name_template @staticmethod def _template_replace_shard_num(shard_name_template): match = re.search('S+', shard_name_template) if match is None: # shard name is required in the template. raise ValueError( "Shard number pattern S+ not found in shard_name_template: %s" % shard_name_template) return shard_name_template.replace( match.group(0), '%%(shard_num)0%dd' % len(match.group(0))) @staticmethod def _template_to_format(shard_name_template): if not shard_name_template: return '' # shard_num is required in the template, while others are optional. replace_funcs = [ FileBasedSink._template_replace_shard_num, FileBasedSink._template_replace_num_shards, FileBasedSink._template_replace_uuid, FileBasedSink._template_replace_window ] for func in replace_funcs: shard_name_template = func(shard_name_template) return shard_name_template @staticmethod def _template_to_glob_format(shard_name_template): if not shard_name_template: return '' match = re.search('S+', shard_name_template) if match is None: raise ValueError( "Shard number pattern S+ not found in shard_name_template: %s" % shard_name_template) shard_name_format = shard_name_template.replace(match.group(0), '*') return FileBasedSink._template_replace_num_shards(shard_name_format) def __eq__(self, other): # TODO: Clean up workitem_test which uses this. # pylint: disable=unidiomatic-typecheck return type(self) == type(other) and self.__dict__ == other.__dict__ class FileBasedSinkWriter(iobase.Writer): """The writer for FileBasedSink. """ def __init__(self, sink, temp_shard_path): self.sink = sink self.temp_shard_path = temp_shard_path self.temp_handle = self.sink.open(temp_shard_path) self.num_records_written = 0 def write(self, value): self.num_records_written += 1 self.sink.write_record(self.temp_handle, value) def at_capacity(self): return ( self.sink.max_records_per_shard and self.num_records_written >= self.sink.max_records_per_shard) or ( self.sink.max_bytes_per_shard and self.sink.byte_counter.bytes_written >= self.sink.max_bytes_per_shard) def close(self): self.sink.close(self.temp_handle) return self.temp_shard_path class _ByteCountingWriter: def __init__(self, writer): self.writer = writer self.bytes_written = 0 def write(self, bs): self.bytes_written += len(bs) self.writer.write(bs) def flush(self): self.writer.flush() def close(self): self.writer.close() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/filebasedsink_test.py0000644000175100017510000003437015206113032022507 0ustar00runnerrunner# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for file sinks.""" # pytype: skip-file import glob import logging import os import shutil import tempfile import unittest import hamcrest as hc import mock import apache_beam as beam from apache_beam.coders import coders from apache_beam.io import filebasedsink from apache_beam.io.filesystem import BeamIOError from apache_beam.options.value_provider import StaticValueProvider from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.transforms.display import DisplayData from apache_beam.transforms.display_test import DisplayDataItemMatcher _LOGGER = logging.getLogger(__name__) # TODO: Refactor code so all io tests are using same library # TestCaseWithTempDirCleanup class. class _TestCaseWithTempDirCleanUp(unittest.TestCase): """Base class for TestCases that deals with TempDir clean-up. Inherited test cases will call self._new_tempdir() to start a temporary dir which will be deleted at the end of the tests (when tearDown() is called). """ def setUp(self): self._tempdirs = [] def tearDown(self): for path in self._tempdirs: if os.path.exists(path): shutil.rmtree(path) self._tempdirs = [] def _new_tempdir(self): result = tempfile.mkdtemp() self._tempdirs.append(result) return result def _create_temp_file(self, name='', suffix='', dir=None, content=None): if not name: name = tempfile.template if not dir: dir = self._new_tempdir() file_name = tempfile.NamedTemporaryFile( delete=False, prefix=name, dir=dir, suffix=suffix).name if content: with open(file_name, 'w') as f: f.write(content) return file_name class MyFileBasedSink(filebasedsink.FileBasedSink): def open(self, temp_path): # TODO: Fix main session pickling. # file_handle = super().open(temp_path) file_handle = filebasedsink.FileBasedSink.open(self, temp_path) file_handle.write(b'[start]') return file_handle def write_encoded_record(self, file_handle, encoded_value): file_handle.write(b'[') file_handle.write(encoded_value) file_handle.write(b']') def close(self, file_handle): file_handle.write(b'[end]') # TODO: Fix main session pickling. # file_handle = super().close(file_handle) file_handle = filebasedsink.FileBasedSink.close(self, file_handle) class TestFileBasedSink(_TestCaseWithTempDirCleanUp): def _common_init(self, sink): # Manually invoke the generic Sink API. init_token = sink.initialize_write() writer1 = sink.open_writer(init_token, '1') writer1.write('a') writer1.write('b') res1 = writer1.close() writer2 = sink.open_writer(init_token, '2') writer2.write('x') writer2.write('y') writer2.write('z') res2 = writer2.close() return init_token, [res1, res2] def test_file_sink_writing(self): temp_path = os.path.join(self._new_tempdir(), 'FileBasedSink') sink = MyFileBasedSink( temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder()) init_token, writer_results = self._common_init(sink) pre_finalize_results = sink.pre_finalize(init_token, writer_results) finalize_res1 = list( sink.finalize_write(init_token, writer_results, pre_finalize_results)) # Retry the finalize operation (as if the first attempt was lost). finalize_res2 = list( sink.finalize_write(init_token, writer_results, pre_finalize_results)) # Check the results. shard1 = temp_path + '-00000-of-00002.output' shard2 = temp_path + '-00001-of-00002.output' self.assertEqual(finalize_res1, [shard1, shard2]) self.assertEqual(finalize_res2, []) self.assertEqual(open(shard1).read(), '[start][a][b][end]') self.assertEqual(open(shard2).read(), '[start][x][y][z][end]') # Check that any temp files are deleted. self.assertCountEqual([shard1, shard2], glob.glob(temp_path + '*')) def test_file_sink_display_data(self): temp_path = os.path.join(self._new_tempdir(), 'display') sink = MyFileBasedSink( temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder()) dd = DisplayData.create_from(sink) expected_items = [ DisplayDataItemMatcher('compression', 'auto'), DisplayDataItemMatcher( 'file_pattern', '{}{}'.format( temp_path, '-%(shard_num)05d-of-%(num_shards)05d.output')) ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_empty_write(self): temp_path = tempfile.NamedTemporaryFile().name sink = MyFileBasedSink( temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder()) with TestPipeline() as p: p | beam.Create([]) | beam.io.Write(sink) # pylint: disable=expression-not-assigned self.assertEqual( open(temp_path + '-00000-of-00001.output').read(), '[start][end]') def test_static_value_provider_empty_write(self): temp_path = StaticValueProvider( value_type=str, value=tempfile.NamedTemporaryFile().name) sink = MyFileBasedSink( temp_path, file_name_suffix=StaticValueProvider(value_type=str, value='.output'), coder=coders.ToBytesCoder()) with TestPipeline() as p: p | beam.Create([]) | beam.io.Write(sink) # pylint: disable=expression-not-assigned self.assertEqual( open(temp_path.get() + '-00000-of-00001.output').read(), '[start][end]') def test_fixed_shard_write(self): temp_path = os.path.join(self._new_tempdir(), 'empty') sink = MyFileBasedSink( temp_path, file_name_suffix='.output', num_shards=3, shard_name_template='_NN_SSS_', coder=coders.ToBytesCoder()) with TestPipeline() as p: p | beam.Create(['a', 'b']) | beam.io.Write(sink) # pylint: disable=expression-not-assigned concat = ''.join( open(temp_path + '_03_%03d_.output' % shard_num).read() for shard_num in range(3)) self.assertTrue('][a][' in concat, concat) self.assertTrue('][b][' in concat, concat) # Not using 'test' in name so that 'nose' doesn't pick this as a test. def run_temp_dir_check( self, no_dir_path, dir_path, no_dir_root_path, dir_root_path, prefix, separator): def _get_temp_dir(file_path_prefix): sink = MyFileBasedSink( file_path_prefix, file_name_suffix='.output', coder=coders.ToBytesCoder()) return sink.initialize_write() temp_dir = _get_temp_dir(no_dir_path) self.assertTrue(temp_dir.startswith(prefix)) last_sep = temp_dir.rfind(separator) self.assertTrue(temp_dir[last_sep + 1:].startswith('beam-temp')) temp_dir = _get_temp_dir(dir_path) self.assertTrue(temp_dir.startswith(prefix)) last_sep = temp_dir.rfind(separator) self.assertTrue(temp_dir[last_sep + 1:].startswith('beam-temp')) with self.assertRaises(ValueError): _get_temp_dir(no_dir_root_path) with self.assertRaises(ValueError): _get_temp_dir(dir_root_path) def test_temp_dir_uniqueness(self): temp_path = os.path.join(self._new_tempdir(), 'unique') sink = MyFileBasedSink(temp_path, coder=coders.ToBytesCoder()) init_list = [''] * 1000 temp_dir_list = [sink._create_temp_dir(temp_path) for _ in init_list] temp_dir_set = set(temp_dir_list) self.assertEqual(len(temp_dir_list), len(temp_dir_set)) def test_temp_dir_gcs(self): try: self.run_temp_dir_check( 'gs://aaa/bbb', 'gs://aaa/bbb/', 'gs://aaa', 'gs://aaa/', 'gs://', '/') except ValueError: _LOGGER.debug('Ignoring test since GCP module is not installed') @mock.patch('apache_beam.io.localfilesystem.os') def test_temp_dir_local(self, filesystem_os_mock): # Here we test a unix-like mock file-system # (not really testing Unix or Windows since we mock the function of 'os' # module). def _fake_unix_split(path): sep = path.rfind('/') if sep < 0: raise ValueError('Path must contain a separator') return (path[:sep], path[sep + 1:]) def _fake_unix_join(base, path): return base + '/' + path filesystem_os_mock.path.abspath = lambda a: a filesystem_os_mock.path.split.side_effect = _fake_unix_split filesystem_os_mock.path.join.side_effect = _fake_unix_join self.run_temp_dir_check('/aaa/bbb', '/aaa/bbb/', '/', '/', '/', '/') def test_file_sink_multi_shards(self): temp_path = os.path.join(self._new_tempdir(), 'multishard') sink = MyFileBasedSink( temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder()) # Manually invoke the generic Sink API. init_token = sink.initialize_write() num_shards = 1000 writer_results = [] for i in range(num_shards): uuid = 'uuid-%05d' % i writer = sink.open_writer(init_token, uuid) writer.write('a') writer.write('b') writer.write(uuid) writer_results.append(writer.close()) pre_finalize_results = sink.pre_finalize(init_token, writer_results) res = sorted( sink.finalize_write(init_token, writer_results, pre_finalize_results)) for i in range(num_shards): shard_name = '%s-%05d-of-%05d.output' % (temp_path, i, num_shards) uuid = 'uuid-%05d' % i self.assertEqual(res[i], shard_name) self.assertEqual( open(shard_name).read(), ('[start][a][b][%s][end]' % uuid)) # Check that any temp files are deleted. self.assertCountEqual(res, glob.glob(temp_path + '*')) @mock.patch.object(filebasedsink.FileSystems, 'rename') def test_file_sink_rename_error(self, rename_mock): temp_path = os.path.join(self._new_tempdir(), 'rename_error') sink = MyFileBasedSink( temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder()) init_token, writer_results = self._common_init(sink) pre_finalize_results = sink.pre_finalize(init_token, writer_results) error_str = 'mock rename error description' rename_mock.side_effect = BeamIOError( 'mock rename error', {('src', 'dst'): error_str}) with self.assertRaisesRegex(Exception, error_str): list( sink.finalize_write(init_token, writer_results, pre_finalize_results)) def test_file_sink_src_missing(self): temp_path = os.path.join(self._new_tempdir(), 'src_missing') sink = MyFileBasedSink( temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder()) init_token, writer_results = self._common_init(sink) pre_finalize_results = sink.pre_finalize(init_token, writer_results) os.remove(writer_results[0]) with self.assertRaisesRegex(Exception, r'not exist'): list( sink.finalize_write(init_token, writer_results, pre_finalize_results)) def test_file_sink_dst_matches_src(self): temp_path = os.path.join(self._new_tempdir(), 'dst_matches_src') sink = MyFileBasedSink( temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder()) init_token, [res1, res2] = self._common_init(sink) pre_finalize_results = sink.pre_finalize(init_token, [res1, res2]) list(sink.finalize_write(init_token, [res1, res2], pre_finalize_results)) self.assertFalse(os.path.exists(res1)) self.assertFalse(os.path.exists(res2)) shard1 = temp_path + '-00000-of-00002.output' shard2 = temp_path + '-00001-of-00002.output' self.assertEqual(open(shard1).read(), '[start][a][b][end]') self.assertEqual(open(shard2).read(), '[start][x][y][z][end]') os.makedirs(os.path.dirname(res1)) shutil.copyfile(shard1, res1) shutil.copyfile(shard2, res2) list(sink.finalize_write(init_token, [res1, res2], pre_finalize_results)) def test_pre_finalize(self): temp_path = os.path.join(self._new_tempdir(), 'pre_finalize') sink = MyFileBasedSink( temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder()) init_token, [res1, res2] = self._common_init(sink) # no-op sink.pre_finalize(init_token, [res1, res2]) # Create finalized outputs from a previous run, which pre_finalize should # delete. shard1 = temp_path + '-00000-of-00002.output' shard2 = temp_path + '-00001-of-00002.output' with open(shard1, 'w') as f: f.write('foo') with open(shard2, 'w') as f: f.write('foo') self.assertTrue(os.path.exists(res1)) self.assertTrue(os.path.exists(res2)) self.assertTrue(os.path.exists(shard1)) self.assertTrue(os.path.exists(shard2)) sink.pre_finalize(init_token, [res1, res2]) self.assertTrue(os.path.exists(res1)) self.assertTrue(os.path.exists(res2)) self.assertFalse(os.path.exists(shard1)) self.assertFalse(os.path.exists(shard2)) @mock.patch.object(filebasedsink.FileSystems, 'delete') def test_pre_finalize_error(self, delete_mock): temp_path = os.path.join(self._new_tempdir(), 'pre_finalize') sink = MyFileBasedSink( temp_path, file_name_suffix='.output', coder=coders.ToBytesCoder()) init_token, [res1, res2] = self._common_init(sink) # no-op sink.pre_finalize(init_token, [res1, res2]) # Create finalized outputs from a previous run, which pre_finalize should # delete. shard1 = temp_path + '-00000-of-00002.output' shard2 = temp_path + '-00001-of-00002.output' with open(shard1, 'w') as f: f.write('foo') with open(shard2, 'w') as f: f.write('foo') error_str = 'mock rename error description' delete_mock.side_effect = BeamIOError( 'mock rename error', {shard2: error_str}) with self.assertRaisesRegex(Exception, error_str): sink.pre_finalize(init_token, [res1, res2]) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/filebasedsource.py0000644000175100017510000004446215206113032022007 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A framework for developing sources for new file types. To create a source for a new file type a sub-class of :class:`FileBasedSource` should be created. Sub-classes of :class:`FileBasedSource` must implement the method :meth:`FileBasedSource.read_records()`. Please read the documentation of that method for more details. For an example implementation of :class:`FileBasedSource` see :class:`~apache_beam.io._AvroSource`. """ # pytype: skip-file from typing import Callable from typing import Iterable from typing import Union from apache_beam.internal import pickler from apache_beam.io import concat_source from apache_beam.io import iobase from apache_beam.io import range_trackers from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystem import FileMetadata from apache_beam.io.filesystems import FileSystems from apache_beam.io.restriction_trackers import OffsetRange from apache_beam.options.value_provider import StaticValueProvider from apache_beam.options.value_provider import ValueProvider from apache_beam.options.value_provider import check_accessible from apache_beam.transforms.core import DoFn from apache_beam.transforms.core import ParDo from apache_beam.transforms.core import PTransform from apache_beam.transforms.display import DisplayDataItem from apache_beam.transforms.util import Reshuffle MAX_NUM_THREADS_FOR_SIZE_ESTIMATION = 25 __all__ = ['FileBasedSource'] class FileBasedSource(iobase.BoundedSource): """A :class:`~apache_beam.io.iobase.BoundedSource` for reading a file glob of a given type.""" MIN_NUMBER_OF_FILES_TO_STAT = 100 MIN_FRACTION_OF_FILES_TO_STAT = 0.01 def __init__( self, file_pattern, min_bundle_size=0, compression_type=CompressionTypes.AUTO, splittable=True, validate=True): """Initializes :class:`FileBasedSource`. Args: file_pattern (str): the file glob to read a string or a :class:`~apache_beam.options.value_provider.ValueProvider` (placeholder to inject a runtime value). min_bundle_size (int): minimum size of bundles that should be generated when performing initial splitting on this source. compression_type (str): Used to handle compressed output files. Typical value is :attr:`CompressionTypes.AUTO `, in which case the final file path's extension will be used to detect the compression. splittable (bool): whether :class:`FileBasedSource` should try to logically split a single file into data ranges so that different parts of the same file can be read in parallel. If set to :data:`False`, :class:`FileBasedSource` will prevent both initial and dynamic splitting of sources for single files. File patterns that represent multiple files may still get split into sources for individual files. Even if set to :data:`True` by the user, :class:`FileBasedSource` may choose to not split the file, for example, for compressed files where currently it is not possible to efficiently read a data range without decompressing the whole file. validate (bool): Boolean flag to verify that the files exist during the pipeline creation time. Raises: TypeError: when **compression_type** is not valid or if **file_pattern** is not a :class:`str` or a :class:`~apache_beam.options.value_provider.ValueProvider`. ValueError: when compression and splittable files are specified. IOError: when the file pattern specified yields an empty result. """ if not isinstance(file_pattern, (str, ValueProvider)): raise TypeError( '%s: file_pattern must be of type string' ' or ValueProvider; got %r instead' % (self.__class__.__name__, file_pattern)) if isinstance(file_pattern, str): file_pattern = StaticValueProvider(str, file_pattern) self._pattern = file_pattern self._concat_source = None self._min_bundle_size = min_bundle_size if not CompressionTypes.is_valid_compression_type(compression_type): raise TypeError( 'compression_type must be CompressionType object but ' 'was %s' % type(compression_type)) self._compression_type = compression_type self._splittable = splittable if validate and file_pattern.is_accessible(): self._validate() def display_data(self): return { 'file_pattern': DisplayDataItem( str(self._pattern), label="File Pattern"), 'compression': DisplayDataItem( str(self._compression_type), label='Compression Type') } @check_accessible(['_pattern']) def _get_concat_source(self) -> concat_source.ConcatSource: if self._concat_source is None: pattern = self._pattern.get() single_file_sources = [] match_result = FileSystems.match([pattern])[0] files_metadata = match_result.metadata_list # We create a reference for FileBasedSource that will be serialized along # with each _SingleFileSource. To prevent this FileBasedSource from having # a reference to ConcatSource (resulting in quadratic space complexity) # we clone it here. file_based_source_ref = pickler.roundtrip(self) for file_metadata in files_metadata: file_name = file_metadata.path file_size = file_metadata.size_in_bytes if file_size == 0: continue # Ignoring empty file. # We determine splittability of this specific file. splittable = ( self.splittable and _determine_splittability_from_compression_type( file_name, self._compression_type)) single_file_source = _SingleFileSource( file_based_source_ref, file_name, 0, file_size, min_bundle_size=self._min_bundle_size, splittable=splittable) single_file_sources.append(single_file_source) FileSystems.report_source_lineage(pattern) self._concat_source = concat_source.ConcatSource(single_file_sources) return self._concat_source def open_file(self, file_name): return FileSystems.open( file_name, 'application/octet-stream', compression_type=self._compression_type) @check_accessible(['_pattern']) def _validate(self): """Validate if there are actual files in the specified glob pattern """ pattern = self._pattern.get() # Limit the responses as we only want to check if something exists match_result = FileSystems.match([pattern], limits=[1])[0] if len(match_result.metadata_list) <= 0: raise IOError('No files found based on the file pattern %s' % pattern) def split( self, desired_bundle_size=None, start_position=None, stop_position=None): return self._get_concat_source().split( desired_bundle_size=desired_bundle_size, start_position=start_position, stop_position=stop_position) def estimate_size(self): return self._get_concat_source().estimate_size() def read(self, range_tracker): return self._get_concat_source().read(range_tracker) def get_range_tracker(self, start_position, stop_position): return self._get_concat_source().get_range_tracker( start_position, stop_position) def read_records(self, file_name, offset_range_tracker): """Returns a generator of records created by reading file 'file_name'. Args: file_name: a ``string`` that gives the name of the file to be read. Method ``FileBasedSource.open_file()`` must be used to open the file and create a seekable file object. offset_range_tracker: a object of type ``OffsetRangeTracker``. This defines the byte range of the file that should be read. See documentation in ``iobase.BoundedSource.read()`` for more information on reading records while complying to the range defined by a given ``RangeTracker``. Returns: an iterator that gives the records read from the given file. """ raise NotImplementedError @property def splittable(self): return self._splittable def _determine_splittability_from_compression_type(file_path, compression_type): if compression_type == CompressionTypes.AUTO: compression_type = CompressionTypes.detect_compression_type(file_path) return compression_type == CompressionTypes.UNCOMPRESSED class _SingleFileSource(iobase.BoundedSource): """Denotes a source for a specific file type.""" def __init__( self, file_based_source, file_name, start_offset, stop_offset, min_bundle_size=0, splittable=True): if not isinstance(start_offset, int): raise TypeError( 'start_offset must be a number. Received: %r' % start_offset) if stop_offset != range_trackers.OffsetRangeTracker.OFFSET_INFINITY: if not isinstance(stop_offset, int): raise TypeError( 'stop_offset must be a number. Received: %r' % stop_offset) if start_offset >= stop_offset: raise ValueError( 'start_offset must be smaller than stop_offset. Received %d and %d ' 'for start and stop offsets respectively' % (start_offset, stop_offset)) self._file_name = file_name self._is_gcs_file = file_name.startswith('gs://') if file_name else False self._start_offset = start_offset self._stop_offset = stop_offset self._min_bundle_size = min_bundle_size self._file_based_source = file_based_source self._splittable = splittable def split(self, desired_bundle_size, start_offset=None, stop_offset=None): if start_offset is None: start_offset = self._start_offset if stop_offset is None: stop_offset = self._stop_offset if self._splittable: splits = OffsetRange(start_offset, stop_offset).split( desired_bundle_size, self._min_bundle_size) for split in splits: yield iobase.SourceBundle( split.stop - split.start, _SingleFileSource( # Copying this so that each sub-source gets a fresh instance. pickler.roundtrip(self._file_based_source), self._file_name, split.start, split.stop, min_bundle_size=self._min_bundle_size, splittable=self._splittable), split.start, split.stop) else: # Returning a single sub-source with end offset set to OFFSET_INFINITY (so # that all data of the source gets read) since this source is # unsplittable. Choosing size of the file as end offset will be wrong for # certain unsplittable source, e.g., compressed sources. yield iobase.SourceBundle( stop_offset - start_offset, _SingleFileSource( self._file_based_source, self._file_name, start_offset, range_trackers.OffsetRangeTracker.OFFSET_INFINITY, min_bundle_size=self._min_bundle_size, splittable=self._splittable), start_offset, range_trackers.OffsetRangeTracker.OFFSET_INFINITY) def estimate_size(self): return self._stop_offset - self._start_offset def get_range_tracker(self, start_position, stop_position): if start_position is None: start_position = self._start_offset if stop_position is None: # If file is unsplittable we choose OFFSET_INFINITY as the default end # offset so that all data of the source gets read. Choosing size of the # file as end offset will be wrong for certain unsplittable source, for # e.g., compressed sources. stop_position = ( self._stop_offset if self._splittable else range_trackers.OffsetRangeTracker.OFFSET_INFINITY) range_tracker = range_trackers.OffsetRangeTracker( start_position, stop_position) if not self._splittable: range_tracker = range_trackers.UnsplittableRangeTracker(range_tracker) return range_tracker def read(self, range_tracker): return self._file_based_source.read_records(self._file_name, range_tracker) def default_output_coder(self): return self._file_based_source.default_output_coder() class _ExpandIntoRanges(DoFn): def __init__( self, splittable, compression_type, desired_bundle_size, min_bundle_size): self._desired_bundle_size = desired_bundle_size self._min_bundle_size = min_bundle_size self._splittable = splittable self._compression_type = compression_type self._size_track = None def process(self, element: Union[str, FileMetadata], *args, **kwargs) -> Iterable[tuple[FileMetadata, OffsetRange]]: if isinstance(element, FileMetadata): metadata_list = [element] else: match_results = FileSystems.match([element]) metadata_list = match_results[0].metadata_list for metadata in metadata_list: FileSystems.report_source_lineage(metadata.path) splittable = ( self._splittable and _determine_splittability_from_compression_type( metadata.path, self._compression_type)) if splittable: for split in OffsetRange(0, metadata.size_in_bytes).split( self._desired_bundle_size, self._min_bundle_size): yield (metadata, split) else: yield ( metadata, OffsetRange(0, range_trackers.OffsetRangeTracker.OFFSET_INFINITY)) class _ReadRange(DoFn): def __init__( self, source_from_file: Union[str, iobase.BoundedSource], with_filename: bool = False) -> None: self._source_from_file = source_from_file self._with_filename = with_filename def process(self, element, *args, **kwargs): metadata, range = element source = self._source_from_file(metadata.path) # Following split() operation has to be performed to create a proper # _SingleFileSource. Otherwise what we have is a ConcatSource that contains # a single _SingleFileSource. ConcatSource.read() expects a RangeTracker for # sub-source range and reads full sub-sources (not byte ranges). source_list = list(source.split(float('inf'))) # Handle the case of an empty source. if not source_list: return source = source_list[0].source for record in source.read(range.new_tracker()): if self._with_filename: yield (metadata.path, record) else: yield record class ReadAllFiles(PTransform): """A Read transform that reads a PCollection of files. Pipeline authors should not use this directly. This is to be used by Read PTransform authors who wishes to implement file-based Read transforms that read a PCollection of files. """ def __init__( self, splittable: bool, compression_type, desired_bundle_size: int, min_bundle_size: int, source_from_file: Callable[[str], iobase.BoundedSource], with_filename: bool = False): """ Args: splittable: If False, files won't be split into sub-ranges. If True, files may or may not be split into data ranges. compression_type: A ``CompressionType`` object that specifies the compression type of the files that will be processed. If ``CompressionType.AUTO``, system will try to automatically determine the compression type based on the extension of files. desired_bundle_size: the desired size of data ranges that should be generated when splitting a file into data ranges. min_bundle_size: minimum size of data ranges that should be generated when splitting a file into data ranges. source_from_file: a function that produces a ``BoundedSource`` given a file name. System will use this function to generate ``BoundedSource`` objects for file paths. Note that file paths passed to this will be for individual files, not for file patterns even if the ``PCollection`` of files processed by the transform consist of file patterns. with_filename: If True, returns a Key Value with the key being the file name and the value being the actual data. If False, it only returns the data. """ self._splittable = splittable self._compression_type = compression_type self._desired_bundle_size = desired_bundle_size self._min_bundle_size = min_bundle_size self._source_from_file = source_from_file self._with_filename = with_filename # TODO(BEAM-14497) always reshuffle once gbk always trigger works. self._is_reshuffle = True def _disable_reshuffle(self): # TODO(BEAM-14497) Remove this private method once gbk always trigger works. # # Currently Reshuffle() holds elements until the stage is completed. When # ReadRange is needed instantly after match (like read continuously), the # reshard is temporarily disabled. However, the read then does not scale and # is deemed experimental. self._is_reshuffle = False return self def expand(self, pvalue): pvalue = ( pvalue | 'ExpandIntoRanges' >> ParDo( _ExpandIntoRanges( self._splittable, self._compression_type, self._desired_bundle_size, self._min_bundle_size))) if self._is_reshuffle: pvalue = pvalue | 'Reshard' >> Reshuffle() return ( pvalue | 'ReadRange' >> ParDo( _ReadRange( self._source_from_file, with_filename=self._with_filename))) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/filebasedsource_test.py0000644000175100017510000006667015206113032023053 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import bz2 import gzip import io import logging import math import os import random import tempfile import unittest import hamcrest as hc import apache_beam as beam from apache_beam.io import filebasedsource from apache_beam.io import iobase from apache_beam.io import range_trackers # importing following private classes for testing from apache_beam.io.concat_source import ConcatSource from apache_beam.io.filebasedsource import FileBasedSource from apache_beam.io.filebasedsource import _SingleFileSource as SingleFileSource from apache_beam.io.filesystem import CompressionTypes from apache_beam.options.value_provider import RuntimeValueProvider from apache_beam.options.value_provider import StaticValueProvider from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.display import DisplayData from apache_beam.transforms.display_test import DisplayDataItemMatcher class LineSource(FileBasedSource): def read_records(self, file_name, range_tracker): f = self.open_file(file_name) try: start = range_tracker.start_position() if start > 0: # Any line that starts after 'start' does not belong to the current # bundle. Seeking to (start - 1) and skipping a line moves the current # position to the starting position of the first line that belongs to # the current bundle. start -= 1 f.seek(start) line = f.readline() start += len(line) current = start line = f.readline() while range_tracker.try_claim(current): # When the source is unsplittable, try_claim is not enough to determine # whether the file has reached to the end. if not line: return yield line.rstrip(b'\n') current += len(line) line = f.readline() finally: f.close() class EOL(object): LF = 1 CRLF = 2 MIXED = 3 LF_WITH_NOTHING_AT_LAST_LINE = 4 def write_data( num_lines, no_data=False, directory=None, prefix=tempfile.template, eol=EOL.LF): """Writes test data to a temporary file. Args: num_lines (int): The number of lines to write. no_data (bool): If :data:`True`, empty lines will be written, otherwise each line will contain a concatenation of b'line' and the line number. directory (str): The name of the directory to create the temporary file in. prefix (str): The prefix to use for the temporary file. eol (int): The line ending to use when writing. :class:`~apache_beam.io.filebasedsource_test.EOL` exposes attributes that can be used here to define the eol. Returns: Tuple[str, List[bytes]]: A tuple of the filename and a list of the written data. """ all_data = [] with tempfile.NamedTemporaryFile(delete=False, dir=directory, prefix=prefix) as f: sep_values = [b'\n', b'\r\n'] for i in range(num_lines): data = b'' if no_data else b'line' + str(i).encode() all_data.append(data) if eol == EOL.LF: sep = sep_values[0] elif eol == EOL.CRLF: sep = sep_values[1] elif eol == EOL.MIXED: sep = sep_values[i % len(sep_values)] elif eol == EOL.LF_WITH_NOTHING_AT_LAST_LINE: sep = b'' if i == (num_lines - 1) else sep_values[0] else: raise ValueError('Received unknown value %s for eol.' % eol) f.write(data + sep) return f.name, all_data def _write_prepared_data( data, directory=None, prefix=tempfile.template, suffix=''): with tempfile.NamedTemporaryFile(delete=False, dir=directory, prefix=prefix, suffix=suffix) as f: f.write(data) return f.name def write_prepared_pattern(data, suffixes=None): assert data, 'Data (%s) seems to be empty' % data if suffixes is None: suffixes = [''] * len(data) temp_dir = tempfile.mkdtemp() for i, d in enumerate(data): file_name = _write_prepared_data( d, temp_dir, prefix='mytemp', suffix=suffixes[i]) return file_name[:file_name.rfind(os.path.sep)] + os.path.sep + 'mytemp*' def write_pattern(lines_per_file, no_data=False): """Writes a pattern of temporary files. Args: lines_per_file (List[int]): The number of lines to write per file. no_data (bool): If :data:`True`, empty lines will be written, otherwise each line will contain a concatenation of b'line' and the line number. Returns: Tuple[str, List[bytes]]: A tuple of the filename pattern and a list of the written data. """ temp_dir = tempfile.mkdtemp() all_data = [] file_name = None start_index = 0 for i in range(len(lines_per_file)): file_name, data = write_data(lines_per_file[i], no_data=no_data, directory=temp_dir, prefix='mytemp') all_data.extend(data) start_index += lines_per_file[i] assert file_name return ( file_name[:file_name.rfind(os.path.sep)] + os.path.sep + 'mytemp*', all_data) class TestConcatSource(unittest.TestCase): class DummySource(iobase.BoundedSource): def __init__(self, values): self._values = values def split( self, desired_bundle_size, start_position=None, stop_position=None): # simply devides values into two bundles middle = len(self._values) // 2 yield iobase.SourceBundle( 0.5, TestConcatSource.DummySource(self._values[:middle]), None, None) yield iobase.SourceBundle( 0.5, TestConcatSource.DummySource(self._values[middle:]), None, None) def get_range_tracker(self, start_position, stop_position): if start_position is None: start_position = 0 if stop_position is None: stop_position = len(self._values) return range_trackers.OffsetRangeTracker(start_position, stop_position) def read(self, range_tracker): for index, value in enumerate(self._values): if not range_tracker.try_claim(index): return yield value def estimate_size(self): return len(self._values) # Assuming each value to be 1 byte. def setUp(self): # Reducing the size of thread pools. Without this test execution may fail in # environments with limited amount of resources. filebasedsource.MAX_NUM_THREADS_FOR_SIZE_ESTIMATION = 2 def test_read(self): sources = [ TestConcatSource.DummySource(range(start, start + 10)) for start in [0, 10, 20] ] concat = ConcatSource(sources) range_tracker = concat.get_range_tracker(None, None) read_data = [value for value in concat.read(range_tracker)] self.assertCountEqual(list(range(30)), read_data) def test_split(self): sources = [ TestConcatSource.DummySource(list(range(start, start + 10))) for start in [0, 10, 20] ] concat = ConcatSource(sources) splits = [split for split in concat.split()] self.assertEqual(6, len(splits)) # Reading all splits read_data = [] for split in splits: range_tracker_for_split = split.source.get_range_tracker( split.start_position, split.stop_position) read_data.extend( [value for value in split.source.read(range_tracker_for_split)]) self.assertCountEqual(list(range(30)), read_data) def test_estimate_size(self): sources = [ TestConcatSource.DummySource(range(start, start + 10)) for start in [0, 10, 20] ] concat = ConcatSource(sources) self.assertEqual(30, concat.estimate_size()) class TestFileBasedSource(unittest.TestCase): def setUp(self): # Reducing the size of thread pools. Without this test execution may fail in # environments with limited amount of resources. filebasedsource.MAX_NUM_THREADS_FOR_SIZE_ESTIMATION = 2 def test_string_or_value_provider_only(self): str_file_pattern = tempfile.NamedTemporaryFile(delete=False).name self.assertEqual( str_file_pattern, FileBasedSource(str_file_pattern)._pattern.value) static_vp_file_pattern = StaticValueProvider( value_type=str, value=str_file_pattern) self.assertEqual( static_vp_file_pattern, FileBasedSource(static_vp_file_pattern)._pattern) runtime_vp_file_pattern = RuntimeValueProvider( option_name='arg', value_type=str, default_value=str_file_pattern) self.assertEqual( runtime_vp_file_pattern, FileBasedSource(runtime_vp_file_pattern)._pattern) # Reset runtime options to avoid side-effects in other tests. RuntimeValueProvider.set_runtime_options(None) invalid_file_pattern = 123 with self.assertRaises(TypeError): FileBasedSource(invalid_file_pattern) def test_validation_file_exists(self): file_name, _ = write_data(10) LineSource(file_name) def test_validation_directory_non_empty(self): temp_dir = tempfile.mkdtemp() file_name, _ = write_data(10, directory=temp_dir) LineSource(file_name) def test_validation_failing(self): no_files_found_error = 'No files found based on the file pattern*' with self.assertRaisesRegex(IOError, no_files_found_error): LineSource('dummy_pattern') with self.assertRaisesRegex(IOError, no_files_found_error): temp_dir = tempfile.mkdtemp() LineSource(os.path.join(temp_dir, '*')) def test_validation_file_missing_verification_disabled(self): LineSource('dummy_pattern', validate=False) def test_fully_read_single_file(self): file_name, expected_data = write_data(10) assert len(expected_data) == 10 fbs = LineSource(file_name) range_tracker = fbs.get_range_tracker(None, None) read_data = [record for record in fbs.read(range_tracker)] self.assertCountEqual(expected_data, read_data) def test_single_file_display_data(self): file_name, _ = write_data(10) fbs = LineSource(file_name) dd = DisplayData.create_from(fbs) expected_items = [ DisplayDataItemMatcher('file_pattern', file_name), DisplayDataItemMatcher('compression', 'auto') ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_fully_read_file_pattern(self): pattern, expected_data = write_pattern([5, 3, 12, 8, 8, 4]) assert len(expected_data) == 40 fbs = LineSource(pattern) range_tracker = fbs.get_range_tracker(None, None) read_data = [record for record in fbs.read(range_tracker)] self.assertCountEqual(expected_data, read_data) def test_fully_read_file_pattern_with_empty_files(self): pattern, expected_data = write_pattern([5, 0, 12, 0, 8, 0]) assert len(expected_data) == 25 fbs = LineSource(pattern) range_tracker = fbs.get_range_tracker(None, None) read_data = [record for record in fbs.read(range_tracker)] self.assertCountEqual(expected_data, read_data) def test_estimate_size_of_file(self): file_name, expected_data = write_data(10) assert len(expected_data) == 10 fbs = LineSource(file_name) self.assertEqual(10 * 6, fbs.estimate_size()) def test_estimate_size_of_pattern(self): pattern, expected_data = write_pattern([5, 3, 10, 8, 8, 4]) assert len(expected_data) == 38 fbs = LineSource(pattern) self.assertEqual(38 * 6, fbs.estimate_size()) pattern, expected_data = write_pattern([5, 3, 9]) assert len(expected_data) == 17 fbs = LineSource(pattern) self.assertEqual(17 * 6, fbs.estimate_size()) def test_estimate_size_with_sampling_same_size(self): num_files = 2 * FileBasedSource.MIN_NUMBER_OF_FILES_TO_STAT pattern, _ = write_pattern([10] * num_files) # Each line will be of length 6 since write_pattern() uses # ('line' + line number + '\n') as data. self.assertEqual( 6 * 10 * num_files, FileBasedSource(pattern).estimate_size()) def test_estimate_size_with_sampling_different_sizes(self): num_files = 2 * FileBasedSource.MIN_NUMBER_OF_FILES_TO_STAT # Each line will be of length 8 since write_pattern() uses # ('line' + line number + '\n') as data. base_size = 500 variance = 5 sizes = [] for _ in range(num_files): sizes.append( int(random.uniform(base_size - variance, base_size + variance))) pattern, _ = write_pattern(sizes) tolerance = 0.05 self.assertAlmostEqual( base_size * 8 * num_files, FileBasedSource(pattern).estimate_size(), delta=base_size * 8 * num_files * tolerance) def test_splits_into_subranges(self): pattern, expected_data = write_pattern([5, 9, 6]) assert len(expected_data) == 20 fbs = LineSource(pattern) splits = [split for split in fbs.split(desired_bundle_size=15)] expected_num_splits = ( math.ceil(float(6 * 5) / 15) + math.ceil(float(6 * 9) / 15) + math.ceil(float(6 * 6) / 15)) assert len(splits) == expected_num_splits def test_read_splits_single_file(self): file_name, expected_data = write_data(100) assert len(expected_data) == 100 fbs = LineSource(file_name) splits = [split for split in fbs.split(desired_bundle_size=33)] # Reading all splits read_data = [] for split in splits: source = split.source range_tracker = source.get_range_tracker( split.start_position, split.stop_position) data_from_split = [data for data in source.read(range_tracker)] read_data.extend(data_from_split) self.assertCountEqual(expected_data, read_data) def test_read_splits_file_pattern(self): pattern, expected_data = write_pattern([34, 66, 40, 24, 24, 12]) assert len(expected_data) == 200 fbs = LineSource(pattern) splits = [split for split in fbs.split(desired_bundle_size=50)] # Reading all splits read_data = [] for split in splits: source = split.source range_tracker = source.get_range_tracker( split.start_position, split.stop_position) data_from_split = [data for data in source.read(range_tracker)] read_data.extend(data_from_split) self.assertCountEqual(expected_data, read_data) def _run_source_test(self, pattern, expected_data, splittable=True): with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> beam.io.Read( LineSource(pattern, splittable=splittable)) assert_that(pcoll, equal_to(expected_data)) def test_source_file(self): file_name, expected_data = write_data(100) assert len(expected_data) == 100 self._run_source_test(file_name, expected_data) def test_source_pattern(self): pattern, expected_data = write_pattern([34, 66, 40, 24, 24, 12]) assert len(expected_data) == 200 self._run_source_test(pattern, expected_data) def test_unsplittable_does_not_split(self): pattern, expected_data = write_pattern([5, 9, 6]) assert len(expected_data) == 20 fbs = LineSource(pattern, splittable=False) splits = [split for split in fbs.split(desired_bundle_size=15)] self.assertEqual(3, len(splits)) def test_source_file_unsplittable(self): file_name, expected_data = write_data(100) assert len(expected_data) == 100 self._run_source_test(file_name, expected_data, False) def test_source_pattern_unsplittable(self): pattern, expected_data = write_pattern([34, 66, 40, 24, 24, 12]) assert len(expected_data) == 200 self._run_source_test(pattern, expected_data, False) def test_read_file_bzip2(self): _, lines = write_data(10) filename = tempfile.NamedTemporaryFile( delete=False, prefix=tempfile.template).name with bz2.BZ2File(filename, 'wb') as f: f.write(b'\n'.join(lines)) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> beam.io.Read( LineSource( filename, splittable=False, compression_type=CompressionTypes.BZIP2)) assert_that(pcoll, equal_to(lines)) def test_read_file_gzip(self): _, lines = write_data(10) filename = tempfile.NamedTemporaryFile( delete=False, prefix=tempfile.template).name with gzip.GzipFile(filename, 'wb') as f: f.write(b'\n'.join(lines)) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> beam.io.Read( LineSource( filename, splittable=False, compression_type=CompressionTypes.GZIP)) assert_that(pcoll, equal_to(lines)) def test_read_pattern_bzip2(self): _, lines = write_data(200) splits = [0, 34, 100, 140, 164, 188, 200] chunks = [lines[splits[i - 1]:splits[i]] for i in range(1, len(splits))] compressed_chunks = [] for c in chunks: compressobj = bz2.BZ2Compressor() compressed_chunks.append( compressobj.compress(b'\n'.join(c)) + compressobj.flush()) file_pattern = write_prepared_pattern(compressed_chunks) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> beam.io.Read( LineSource( file_pattern, splittable=False, compression_type=CompressionTypes.BZIP2)) assert_that(pcoll, equal_to(lines)) def test_read_pattern_gzip(self): _, lines = write_data(200) splits = [0, 34, 100, 140, 164, 188, 200] chunks = [lines[splits[i - 1]:splits[i]] for i in range(1, len(splits))] compressed_chunks = [] for c in chunks: out = io.BytesIO() with gzip.GzipFile(fileobj=out, mode="wb") as f: f.write(b'\n'.join(c)) compressed_chunks.append(out.getvalue()) file_pattern = write_prepared_pattern(compressed_chunks) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> beam.io.Read( LineSource( file_pattern, splittable=False, compression_type=CompressionTypes.GZIP)) assert_that(pcoll, equal_to(lines)) def test_read_auto_single_file_bzip2(self): _, lines = write_data(10) filename = tempfile.NamedTemporaryFile( delete=False, prefix=tempfile.template, suffix='.bz2').name with bz2.BZ2File(filename, 'wb') as f: f.write(b'\n'.join(lines)) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> beam.io.Read( LineSource(filename, compression_type=CompressionTypes.AUTO)) assert_that(pcoll, equal_to(lines)) def test_read_auto_single_file_gzip(self): _, lines = write_data(10) filename = tempfile.NamedTemporaryFile( delete=False, prefix=tempfile.template, suffix='.gz').name with gzip.GzipFile(filename, 'wb') as f: f.write(b'\n'.join(lines)) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> beam.io.Read( LineSource(filename, compression_type=CompressionTypes.AUTO)) assert_that(pcoll, equal_to(lines)) def test_read_auto_pattern(self): _, lines = write_data(200) splits = [0, 34, 100, 140, 164, 188, 200] chunks = [lines[splits[i - 1]:splits[i]] for i in range(1, len(splits))] compressed_chunks = [] for c in chunks: out = io.BytesIO() with gzip.GzipFile(fileobj=out, mode="wb") as f: f.write(b'\n'.join(c)) compressed_chunks.append(out.getvalue()) file_pattern = write_prepared_pattern( compressed_chunks, suffixes=['.gz'] * len(chunks)) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> beam.io.Read( LineSource(file_pattern, compression_type=CompressionTypes.AUTO)) assert_that(pcoll, equal_to(lines)) def test_read_auto_pattern_compressed_and_uncompressed(self): _, lines = write_data(200) splits = [0, 34, 100, 140, 164, 188, 200] chunks = [lines[splits[i - 1]:splits[i]] for i in range(1, len(splits))] chunks_to_write = [] for i, c in enumerate(chunks): if i % 2 == 0: out = io.BytesIO() with gzip.GzipFile(fileobj=out, mode="wb") as f: f.write(b'\n'.join(c)) chunks_to_write.append(out.getvalue()) else: chunks_to_write.append(b'\n'.join(c)) file_pattern = write_prepared_pattern( chunks_to_write, suffixes=(['.gz', ''] * 3)) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> beam.io.Read( LineSource(file_pattern, compression_type=CompressionTypes.AUTO)) assert_that(pcoll, equal_to(lines)) def test_splits_get_coder_from_fbs(self): class DummyCoder(object): val = 12345 class FileBasedSourceWithCoder(LineSource): def default_output_coder(self): return DummyCoder() pattern, expected_data = write_pattern([34, 66, 40, 24, 24, 12]) self.assertEqual(200, len(expected_data)) fbs = FileBasedSourceWithCoder(pattern) splits = [split for split in fbs.split(desired_bundle_size=50)] self.assertTrue(len(splits)) for split in splits: self.assertEqual(DummyCoder.val, split.source.default_output_coder().val) class TestSingleFileSource(unittest.TestCase): def setUp(self): # Reducing the size of thread pools. Without this test execution may fail in # environments with limited amount of resources. filebasedsource.MAX_NUM_THREADS_FOR_SIZE_ESTIMATION = 2 def test_source_creation_fails_for_non_number_offsets(self): start_not_a_number_error = 'start_offset must be a number*' stop_not_a_number_error = 'stop_offset must be a number*' file_name = 'dummy_pattern' fbs = LineSource(file_name, validate=False) with self.assertRaisesRegex(TypeError, start_not_a_number_error): SingleFileSource( fbs, file_name='dummy_file', start_offset='aaa', stop_offset='bbb') with self.assertRaisesRegex(TypeError, start_not_a_number_error): SingleFileSource( fbs, file_name='dummy_file', start_offset='aaa', stop_offset=100) with self.assertRaisesRegex(TypeError, stop_not_a_number_error): SingleFileSource( fbs, file_name='dummy_file', start_offset=100, stop_offset='bbb') with self.assertRaisesRegex(TypeError, stop_not_a_number_error): SingleFileSource( fbs, file_name='dummy_file', start_offset=100, stop_offset=None) with self.assertRaisesRegex(TypeError, start_not_a_number_error): SingleFileSource( fbs, file_name='dummy_file', start_offset=None, stop_offset=100) def test_source_creation_display_data(self): file_name = 'dummy_pattern' fbs = LineSource(file_name, validate=False) dd = DisplayData.create_from(fbs) expected_items = [ DisplayDataItemMatcher('compression', 'auto'), DisplayDataItemMatcher('file_pattern', file_name) ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_source_creation_fails_if_start_lg_stop(self): start_larger_than_stop_error = ( 'start_offset must be smaller than stop_offset*') fbs = LineSource('dummy_pattern', validate=False) SingleFileSource( fbs, file_name='dummy_file', start_offset=99, stop_offset=100) with self.assertRaisesRegex(ValueError, start_larger_than_stop_error): SingleFileSource( fbs, file_name='dummy_file', start_offset=100, stop_offset=99) with self.assertRaisesRegex(ValueError, start_larger_than_stop_error): SingleFileSource( fbs, file_name='dummy_file', start_offset=100, stop_offset=100) def test_estimates_size(self): fbs = LineSource('dummy_pattern', validate=False) # Should simply return stop_offset - start_offset source = SingleFileSource( fbs, file_name='dummy_file', start_offset=0, stop_offset=100) self.assertEqual(100, source.estimate_size()) source = SingleFileSource( fbs, file_name='dummy_file', start_offset=10, stop_offset=100) self.assertEqual(90, source.estimate_size()) def test_read_range_at_beginning(self): fbs = LineSource('dummy_pattern', validate=False) file_name, expected_data = write_data(10) assert len(expected_data) == 10 source = SingleFileSource(fbs, file_name, 0, 10 * 6) range_tracker = source.get_range_tracker(0, 20) read_data = [value for value in source.read(range_tracker)] self.assertCountEqual(expected_data[:4], read_data) def test_read_range_at_end(self): fbs = LineSource('dummy_pattern', validate=False) file_name, expected_data = write_data(10) assert len(expected_data) == 10 source = SingleFileSource(fbs, file_name, 0, 10 * 6) range_tracker = source.get_range_tracker(40, 60) read_data = [value for value in source.read(range_tracker)] self.assertCountEqual(expected_data[-3:], read_data) def test_read_range_at_middle(self): fbs = LineSource('dummy_pattern', validate=False) file_name, expected_data = write_data(10) assert len(expected_data) == 10 source = SingleFileSource(fbs, file_name, 0, 10 * 6) range_tracker = source.get_range_tracker(20, 40) read_data = [value for value in source.read(range_tracker)] self.assertCountEqual(expected_data[4:7], read_data) def test_produces_splits_desiredsize_large_than_size(self): fbs = LineSource('dummy_pattern', validate=False) file_name, expected_data = write_data(10) assert len(expected_data) == 10 source = SingleFileSource(fbs, file_name, 0, 10 * 6) splits = [split for split in source.split(desired_bundle_size=100)] self.assertEqual(1, len(splits)) self.assertEqual(60, splits[0].weight) self.assertEqual(0, splits[0].start_position) self.assertEqual(60, splits[0].stop_position) range_tracker = splits[0].source.get_range_tracker(None, None) read_data = [value for value in splits[0].source.read(range_tracker)] self.assertCountEqual(expected_data, read_data) def test_produces_splits_desiredsize_smaller_than_size(self): fbs = LineSource('dummy_pattern', validate=False) file_name, expected_data = write_data(10) assert len(expected_data) == 10 source = SingleFileSource(fbs, file_name, 0, 10 * 6) splits = [split for split in source.split(desired_bundle_size=25)] self.assertEqual(3, len(splits)) read_data = [] for split in splits: source = split.source range_tracker = source.get_range_tracker( split.start_position, split.stop_position) data_from_split = [data for data in source.read(range_tracker)] read_data.extend(data_from_split) self.assertCountEqual(expected_data, read_data) def test_produce_split_with_start_and_end_positions(self): fbs = LineSource('dummy_pattern', validate=False) file_name, expected_data = write_data(10) assert len(expected_data) == 10 source = SingleFileSource(fbs, file_name, 0, 10 * 6) splits = [ split for split in source.split( desired_bundle_size=15, start_offset=10, stop_offset=50) ] self.assertEqual(3, len(splits)) read_data = [] for split in splits: source = split.source range_tracker = source.get_range_tracker( split.start_position, split.stop_position) data_from_split = [data for data in source.read(range_tracker)] read_data.extend(data_from_split) self.assertCountEqual(expected_data[2:9], read_data) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/fileio.py0000644000175100017510000007713715206113032020124 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """``PTransforms`` for manipulating files in Apache Beam. Provides reading ``PTransform``\\s, ``MatchFiles``, ``MatchAll``, that produces a ``PCollection`` of records representing a file and its metadata; and ``ReadMatches``, which takes in a ``PCollection`` of file metadata records, and produces a ``PCollection`` of ``ReadableFile`` objects. These transforms currently do not support splitting by themselves. Writing to Files ================ The transforms in this file include ``WriteToFiles``, which allows you to write a ``beam.PCollection`` to files, and gives you many options to customize how to do this. The ``WriteToFiles`` transform supports bounded and unbounded PCollections (i.e. it can be used both batch and streaming pipelines). For streaming pipelines, it currently does not have support for multiple trigger firings on the same window. File Naming ----------- One of the parameters received by ``WriteToFiles`` is a function specifying how to name the files that are written. This is a function that takes in the following parameters: - window - pane - shard_index - total_shards - compression - destination It should return a file name that is unique for a combination of these parameters. The default naming strategy is to name files in the format `$prefix-$start-$end-$pane-$shard-of-$numShards$suffix$compressionSuffix`, where: - `$prefix` is, by default, `"output"`. - `$start` and `$end` are the boundaries of the window for the data being written. These are omitted if we're using the Global window. - `$pane` is the index for the number of firing for a window. - `$shard` and `$numShards` are the current shard number, and the total number of shards for this window firing. - `$suffix` is, by default, an empty string, but it can be set by the user via ``default_file_naming``. Dynamic Destinations -------------------- If the elements in the input ``beam.PCollection`` can be partitioned into groups that should be treated differently (e.g. some events are to be stored as CSV, while some others are to be stored as Avro files), it is possible to do this by passing a `destination` parameter to ``WriteToFiles``. Something like the following:: my_pcollection | beam.io.fileio.WriteToFiles( path='/my/file/path', destination=lambda record: 'avro' if record['type'] == 'A' else 'csv', sink=lambda dest: AvroSink() if dest == 'avro' else CsvSink(), file_naming=beam.io.fileio.destination_prefix_naming()) In this transform, depending on the type of a record, it will be written down to a destination named `'avro'`, or `'csv'`. The value returned by the `destination` call is then passed to the `sink` call, to determine what sort of sink will be used for each destination. The return type of the `destination` parameter can be anything, as long as elements can be grouped by it. """ # pytype: skip-file import collections import logging import random import uuid from collections import namedtuple from functools import partial from typing import Any from typing import BinaryIO # pylint: disable=unused-import from typing import Callable from typing import Iterable from typing import Union import apache_beam as beam from apache_beam.io import filesystem from apache_beam.io import filesystems from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystem import CompressionTypes from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.value_provider import StaticValueProvider from apache_beam.options.value_provider import ValueProvider from apache_beam.transforms.periodicsequence import PeriodicImpulse from apache_beam.transforms.userstate import CombiningValueStateSpec from apache_beam.transforms.window import BoundedWindow from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import GlobalWindow from apache_beam.transforms.window import IntervalWindow from apache_beam.utils.timestamp import MAX_TIMESTAMP from apache_beam.utils.timestamp import Timestamp __all__ = [ 'EmptyMatchTreatment', 'MatchFiles', 'MatchAll', 'MatchContinuously', 'ReadableFile', 'ReadMatches', 'WriteToFiles' ] _LOGGER = logging.getLogger(__name__) FileMetadata = namedtuple("FileMetadata", "mime_type compression_type") CreateFileMetadataFn = Callable[[str, str], FileMetadata] class EmptyMatchTreatment(object): """How to treat empty matches in ``MatchAll`` and ``MatchFiles`` transforms. If empty matches are disallowed, an error will be thrown if a pattern does not match any files.""" ALLOW = 'ALLOW' DISALLOW = 'DISALLOW' ALLOW_IF_WILDCARD = 'ALLOW_IF_WILDCARD' @staticmethod def allow_empty_match(pattern, setting): if setting == EmptyMatchTreatment.ALLOW: return True elif setting == EmptyMatchTreatment.ALLOW_IF_WILDCARD and '*' in pattern: return True elif setting == EmptyMatchTreatment.DISALLOW: return False else: raise ValueError(setting) class _MatchAllFn(beam.DoFn): def __init__(self, empty_match_treatment): self._empty_match_treatment = empty_match_treatment def process(self, file_pattern: str) -> list[filesystem.FileMetadata]: # TODO: Should we batch the lookups? match_results = filesystems.FileSystems.match([file_pattern]) match_result = match_results[0] if (not match_result.metadata_list and not EmptyMatchTreatment.allow_empty_match(file_pattern, self._empty_match_treatment)): raise BeamIOError( 'Empty match for pattern %s. Disallowed.' % file_pattern) return match_result.metadata_list class MatchFiles(beam.PTransform): """Matches a file pattern using ``FileSystems.match``. This ``PTransform`` returns a ``PCollection`` of matching files in the form of ``FileMetadata`` objects.""" def __init__( self, file_pattern: str, empty_match_treatment=EmptyMatchTreatment.ALLOW_IF_WILDCARD): self._file_pattern = file_pattern self._empty_match_treatment = empty_match_treatment def expand(self, pcoll) -> beam.PCollection[filesystem.FileMetadata]: return pcoll.pipeline | beam.Create([self._file_pattern]) | MatchAll( empty_match_treatment=self._empty_match_treatment) class MatchAll(beam.PTransform): """Matches file patterns from the input PCollection via ``FileSystems.match``. This ``PTransform`` returns a ``PCollection`` of matching files in the form of ``FileMetadata`` objects.""" def __init__(self, empty_match_treatment=EmptyMatchTreatment.ALLOW): self._empty_match_treatment = empty_match_treatment def expand( self, pcoll: beam.PCollection, ) -> beam.PCollection[filesystem.FileMetadata]: return pcoll | beam.ParDo(_MatchAllFn(self._empty_match_treatment)) class ReadableFile(object): """A utility class for accessing files.""" def __init__(self, metadata, compression=None): self.metadata = metadata self._compression = compression def open(self, mime_type='text/plain', compression_type=None): compression = ( compression_type or self._compression or filesystems.CompressionTypes.AUTO) return filesystems.FileSystems.open( self.metadata.path, mime_type=mime_type, compression_type=compression) def read(self, mime_type='application/octet-stream'): return self.open(mime_type).read() def read_utf8(self): return self.open().read().decode('utf-8') class _ReadMatchesFn(beam.DoFn): def __init__(self, compression, skip_directories): self._compression = compression self._skip_directories = skip_directories def process( self, file_metadata: Union[str, filesystem.FileMetadata], ) -> Iterable[ReadableFile]: metadata = ( filesystem.FileMetadata(file_metadata, 0) if isinstance( file_metadata, str) else file_metadata) if ((metadata.path.endswith('/') or metadata.path.endswith('\\')) and self._skip_directories): return elif metadata.path.endswith('/') or metadata.path.endswith('\\'): raise BeamIOError( 'Directories are not allowed in ReadMatches transform.' 'Found %s.' % metadata.path) # TODO: Mime type? Other arguments? Maybe arguments passed in to transform? yield ReadableFile(metadata, self._compression) class MatchContinuously(beam.PTransform): """Checks for new files for a given pattern every interval. This ``PTransform`` returns a ``PCollection`` of matching files in the form of ``FileMetadata`` objects. MatchContinuously is experimental. No backwards-compatibility guarantees. Matching continuously scales poorly, as it is stateful, and requires storing file ids in memory. In addition, because it is memory-only, if a pipeline is restarted, already processed files will be reprocessed. Consider an alternate technique, such as Pub/Sub Notifications (https://cloud.google.com/storage/docs/pubsub-notifications) when using GCS if possible. """ def __init__( self, file_pattern, interval=360.0, has_deduplication=True, start_timestamp=Timestamp.now(), stop_timestamp=MAX_TIMESTAMP, match_updated_files=False, apply_windowing=False, empty_match_treatment=EmptyMatchTreatment.ALLOW): """Initializes a MatchContinuously transform. Args: file_pattern: The file path to read from. interval: Interval at which to check for files in seconds. has_deduplication: Whether files already read are discarded or not. start_timestamp: Timestamp for start file checking. stop_timestamp: Timestamp after which no more files will be checked. match_updated_files: (When has_deduplication is set to True) whether match file with timestamp changes. apply_windowing: Whether each element should be assigned to individual window. If false, all elements will reside in global window. """ self.file_pattern = file_pattern self.interval = interval self.has_deduplication = has_deduplication self.start_ts = start_timestamp self.stop_ts = stop_timestamp self.match_upd = match_updated_files self.apply_windowing = apply_windowing self.empty_match_treatment = empty_match_treatment _LOGGER.warning( 'Matching Continuously is stateful, and can scale poorly. ' 'Consider using Pub/Sub Notifications ' '(https://cloud.google.com/storage/docs/pubsub-notifications) ' 'if possible') def expand(self, pbegin) -> beam.PCollection[filesystem.FileMetadata]: # invoke periodic impulse impulse = pbegin | PeriodicImpulse( start_timestamp=self.start_ts, stop_timestamp=self.stop_ts, fire_interval=self.interval) # match file pattern periodically match_files = ( impulse | 'GetFilePattern' >> beam.Map(lambda x: self.file_pattern) | MatchAll(self.empty_match_treatment)) # apply deduplication strategy if required if self.has_deduplication: # Making a Key Value so each file has its own state. match_files = match_files | 'ToKV' >> beam.Map(lambda x: (x.path, x)) if self.match_upd: match_files = match_files | 'RemoveOldAlreadyRead' >> beam.ParDo( _RemoveOldDuplicates()) else: match_files = match_files | 'RemoveAlreadyRead' >> beam.ParDo( _RemoveDuplicates()) # apply windowing if required. Apply at last because deduplication relies on # the global window. if self.apply_windowing: match_files = match_files | beam.WindowInto(FixedWindows(self.interval)) return match_files class ReadMatches(beam.PTransform): """Converts each result of MatchFiles() or MatchAll() to a ReadableFile. This helps read in a file's contents or obtain a file descriptor.""" def __init__(self, compression=None, skip_directories=True): self._compression = compression self._skip_directories = skip_directories def expand( self, pcoll: beam.PCollection[Union[str, filesystem.FileMetadata]], ) -> beam.PCollection[ReadableFile]: return pcoll | beam.ParDo( _ReadMatchesFn(self._compression, self._skip_directories)) class FileSink(object): """Specifies how to write elements to individual files in ``WriteToFiles``. A Sink class must implement the following: - The ``open`` method, which initializes writing to a file handler (it is not responsible for opening the file handler itself). - The ``write`` method, which writes an element to the file that was passed in ``open``. - The ``flush`` method, which flushes any buffered state. This is most often called before closing a file (but not exclusively called in that situation). The sink is not responsible for closing the file handler. A Sink class can override the following: - The ``create_metadata`` method, which creates all metadata passed to Filesystems.create. """ def create_metadata( self, destination: str, full_file_name: str) -> FileMetadata: return FileMetadata( mime_type="application/octet-stream", compression_type=CompressionTypes.AUTO) def open(self, fh: BinaryIO) -> None: raise NotImplementedError def write(self, record): raise NotImplementedError def flush(self): raise NotImplementedError @beam.typehints.with_input_types(str) class TextSink(FileSink): """A sink that encodes utf8 elements, and writes to file handlers. This sink simply calls file_handler.write(record.encode('utf8') + '\n') on all records that come into it. """ def open(self, fh): self._fh = fh def write(self, record): self._fh.write(record.encode('utf8')) self._fh.write(b'\n') def flush(self): self._fh.flush() def prefix_naming(prefix): return default_file_naming(prefix) _DEFAULT_FILE_NAME_TEMPLATE = ( '{prefix}-{start}-{end}-{pane}-' '{shard:05d}-of-{total_shards:05d}' '{suffix}{compression}') def _format_shard( window, pane, shard_index, total_shards, compression, prefix, suffix): kwargs = { 'prefix': prefix, 'start': '', 'end': '', 'pane': '', 'shard': 0, 'total_shards': 0, 'suffix': '', 'compression': '' } if total_shards is not None and shard_index is not None: kwargs['shard'] = int(shard_index) kwargs['total_shards'] = int(total_shards) if window != GlobalWindow(): kwargs['start'] = window.start.to_utc_datetime().isoformat() kwargs['end'] = window.end.to_utc_datetime().isoformat() # TODO(https://github.com/apache/beam/issues/18721): Add support for PaneInfo # If the PANE is the ONLY firing in the window, we don't add it. #if pane and not (pane.is_first and pane.is_last): # kwargs['pane'] = pane.index if suffix: kwargs['suffix'] = suffix if compression: kwargs['compression'] = '.%s' % compression # Remove separators for unused template parts. format = _DEFAULT_FILE_NAME_TEMPLATE if shard_index is None: format = format.replace('-{shard:05d}', '') if total_shards is None: format = format.replace('-of-{total_shards:05d}', '') for name, value in kwargs.items(): if value in (None, ''): format = format.replace('-{%s}' % name, '') return format.format(**kwargs) FileNaming = Callable[[Any, Any, int, int, Any, str, str], str] def destination_prefix_naming(suffix=None) -> FileNaming: def _inner(window, pane, shard_index, total_shards, compression, destination): prefix = str(destination) return _format_shard( window, pane, shard_index, total_shards, compression, prefix, suffix) return _inner def default_file_naming(prefix, suffix=None) -> FileNaming: def _inner(window, pane, shard_index, total_shards, compression, destination): return _format_shard( window, pane, shard_index, total_shards, compression, prefix, suffix) return _inner def single_file_naming(prefix, suffix=None) -> FileNaming: def _inner(window, pane, shard_index, total_shards, compression, destination): assert shard_index in (0, None), shard_index assert total_shards in (1, None), total_shards return _format_shard(window, pane, None, None, compression, prefix, suffix) return _inner _FileResult = collections.namedtuple( 'FileResult', [ 'file_name', 'shard_index', 'total_shards', 'window', 'pane', 'destination' ]) # Adding a class to contain PyDoc. class FileResult(_FileResult): """A descriptor of a file that has been written.""" pass class WriteToFiles(beam.PTransform): r"""Write the incoming PCollection to a set of output files. The incoming ``PCollection`` may be bounded or unbounded. **Note:** For unbounded ``PCollection``\s, this transform does not support multiple firings per Window (due to the fact that files are named only by their destination, and window, at the moment). """ # We allow up to 20 different destinations to be written in a single bundle. # Too many files will add memory pressure to the worker, so we let it be 20. MAX_NUM_WRITERS_PER_BUNDLE = 20 DEFAULT_SHARDING = 5 def __init__( self, path, file_naming=None, destination=None, temp_directory=None, sink=None, shards=None, output_fn=None, max_writers_per_bundle=MAX_NUM_WRITERS_PER_BUNDLE): """Initializes a WriteToFiles transform. Args: path (str, ValueProvider): The directory to write files into. file_naming (callable): A callable that takes in a window, pane, shard_index, total_shards and compression; and returns a file name. destination (callable): If this argument is provided, the sink parameter must also be a callable. temp_directory (str, ValueProvider): To ensure atomicity in the transform, the output is written into temporary files, which are written to a directory that is meant to be temporary as well. Once the whole output has been written, the files are moved into their final destination, and given their final names. By default, the temporary directory will be within the temp_location of your pipeline. sink (callable, ~apache_beam.io.fileio.FileSink): The sink to use to write into a file. It should implement the methods of a ``FileSink``. Pass a class signature or an instance of FileSink to this parameter. If none is provided, a ``TextSink`` is used. shards (int): The number of shards per destination and trigger firing. output_fn (callable, optional): A callable to process the output. This parameter is currently unused and retained for backward compatibility. max_writers_per_bundle (int): The number of writers that can be open concurrently in a single worker that's processing one bundle. """ self.path = ( path if isinstance(path, ValueProvider) else StaticValueProvider( str, path)) self.file_naming_fn = file_naming or default_file_naming('output') self.destination_fn = self._get_destination_fn(destination) self._temp_directory = temp_directory self.sink_fn = self._get_sink_fn(sink) self.shards = shards or WriteToFiles.DEFAULT_SHARDING self.output_fn = output_fn or (lambda x: x) self._max_num_writers_per_bundle = max_writers_per_bundle @staticmethod def _get_sink_fn(input_sink) -> Callable[[Any], FileSink]: if isinstance(input_sink, type) and issubclass(input_sink, FileSink): return lambda x: input_sink() elif isinstance(input_sink, FileSink): kls = input_sink.__class__ return lambda x: kls() elif callable(input_sink): return input_sink else: return lambda x: TextSink() @staticmethod def _get_destination_fn(destination) -> Callable[[Any], str]: if isinstance(destination, ValueProvider): return lambda elm: destination.get() elif callable(destination): return destination else: return lambda elm: destination def expand(self, pcoll): p = pcoll.pipeline if not self._temp_directory: temp_location = ( p.options.view_as(GoogleCloudOptions).temp_location or self.path.get()) dir_uid = str(uuid.uuid4()) self._temp_directory = StaticValueProvider( str, filesystems.FileSystems.join(temp_location, '.temp%s' % dir_uid)) _LOGGER.info('Added temporary directory %s', self._temp_directory.get()) output = ( pcoll | beam.ParDo( _WriteUnshardedRecordsFn( base_path=self._temp_directory, destination_fn=self.destination_fn, sink_fn=self.sink_fn, max_writers_per_bundle=self._max_num_writers_per_bundle)). with_outputs( _WriteUnshardedRecordsFn.SPILLED_RECORDS, _WriteUnshardedRecordsFn.WRITTEN_FILES)) written_files_pc = output[_WriteUnshardedRecordsFn.WRITTEN_FILES] spilled_records_pc = output[_WriteUnshardedRecordsFn.SPILLED_RECORDS] more_written_files_pc = ( spilled_records_pc | beam.ParDo( _AppendShardedDestination(self.destination_fn, self.shards)) | "GroupRecordsByDestinationAndShard" >> beam.GroupByKey() | beam.ParDo( _WriteShardedRecordsFn( self._temp_directory, self.sink_fn, self.shards))) files_by_destination_pc = ( (written_files_pc, more_written_files_pc) | beam.Flatten() | beam.Map(lambda file_result: (file_result.destination, file_result)) | "GroupTempFilesByDestination" >> beam.GroupByKey()) # Now we should take the temporary files, and write them to the final # destination, with their proper names. file_results = ( files_by_destination_pc | beam.ParDo( _MoveTempFilesIntoFinalDestinationFn( self.path, self.file_naming_fn, self._temp_directory))) return file_results def _create_writer( base_path, writer_key: tuple[str, IntervalWindow], create_metadata_fn: CreateFileMetadataFn, ): try: filesystems.FileSystems.mkdirs(base_path) except IOError: # Directory already exists. pass destination = writer_key[0] # The file name has a prefix determined by destination+window, along with # a random string. This allows us to retrieve orphaned files later on. file_name = '%s_%s' % (abs(hash(writer_key)), uuid.uuid4()) full_file_name = filesystems.FileSystems.join(base_path, file_name) metadata = create_metadata_fn(destination, full_file_name) return full_file_name, filesystems.FileSystems.create( full_file_name, **metadata._asdict()) class _MoveTempFilesIntoFinalDestinationFn(beam.DoFn): def __init__(self, path, file_naming_fn, temp_dir): self.path = path self.file_naming_fn = file_naming_fn self.temporary_directory = temp_dir def process(self, element, w=beam.DoFn.WindowParam): destination = element[0] # list of FileResult objects for temp files temp_file_results = list(element[1]) # list of FileResult objects for final files final_file_results = [] for i, r in enumerate(temp_file_results): # TODO(pabloem): Handle compression for files. final_file_name = self.file_naming_fn( r.window, r.pane, i, len(temp_file_results), '', destination) final_file_results.append( FileResult( final_file_name, i, len(temp_file_results), r.window, r.pane, destination)) move_from = [f.file_name for f in temp_file_results] move_to = [f.file_name for f in final_file_results] _LOGGER.info( 'Moving %d temporary files to dir: %s as %s', len(move_from), self.path.get(), move_to) try: filesystems.FileSystems.mkdirs(self.path.get()) except IOError as e: cause = repr(e) if 'FileExistsError' not in cause: # Usually harmless. Especially if see FileExistsError so no need to log _LOGGER.debug('Fail to create dir for final destination: %s', cause) try: filesystems.FileSystems.rename( move_from, [filesystems.FileSystems.join(self.path.get(), f) for f in move_to]) except BeamIOError: # This error is not serious, because it may happen on a retry of the # bundle. We simply log it. _LOGGER.debug( 'Exception occurred during moving files: %s. This may be due to a' ' bundle being retried.', move_from) yield from final_file_results _LOGGER.debug( 'Checking orphaned temporary files for destination %s and window %s', destination, w) writer_key = (destination, w) self._check_orphaned_files(writer_key) def _check_orphaned_files(self, writer_key): try: prefix = filesystems.FileSystems.join( self.temporary_directory.get(), str(abs(hash(writer_key)))) match_result = filesystems.FileSystems.match(['%s*' % prefix]) orphaned_files = [m.path for m in match_result[0].metadata_list] if len(orphaned_files) > 0: _LOGGER.warning( 'Some files may be left orphaned in the temporary folder: %s. ' 'This may be a result of retried work items or insufficient' 'permissions to delete these temp files.', orphaned_files) except BeamIOError as e: _LOGGER.warning('Exceptions when checking orphaned files: %s', e) class _WriteShardedRecordsFn(beam.DoFn): def __init__( self, base_path, sink_fn: Callable[[Any], FileSink], shards: int): self.base_path = base_path self.sink_fn = sink_fn self.shards = shards def process( self, element, w=beam.DoFn.WindowParam, pane=beam.DoFn.PaneInfoParam): destination_and_shard = element[0] destination = destination_and_shard[0] shard = destination_and_shard[1] records = element[1] sink = self.sink_fn(destination) full_file_name, writer = _create_writer( base_path=self.base_path.get(), writer_key=(destination, w), create_metadata_fn=sink.create_metadata) sink.open(writer) for r in records: sink.write(r) sink.flush() writer.close() _LOGGER.info( 'Writing file %s for destination %s and shard %s', full_file_name, destination, repr(shard)) yield FileResult( full_file_name, shard_index=shard, total_shards=self.shards, window=w, pane=pane, destination=destination) class _AppendShardedDestination(beam.DoFn): def __init__(self, destination: Callable[[Any], str], shards: int): self.destination_fn = destination self.shards = shards # We start the shards for a single destination at an arbitrary point. self._shard_counter: collections.defaultdict[ str, int] = collections.defaultdict(lambda: random.randrange(self.shards)) def _next_shard_for_destination(self, destination): self._shard_counter[destination] = ((self._shard_counter[destination] + 1) % self.shards) return self._shard_counter[destination] def process(self, record): destination = self.destination_fn(record) shard = self._next_shard_for_destination(destination) yield ((destination, shard), record) class _WriteUnshardedRecordsFn(beam.DoFn): SPILLED_RECORDS = 'spilled_records' WRITTEN_FILES = 'written_files' _writers_and_sinks: dict[tuple[str, BoundedWindow], tuple[BinaryIO, FileSink]] = None _file_names: dict[tuple[str, BoundedWindow], str] = None def __init__( self, base_path, destination_fn, sink_fn, max_writers_per_bundle=WriteToFiles.MAX_NUM_WRITERS_PER_BUNDLE): self.base_path = base_path self.destination_fn = destination_fn self.sink_fn = sink_fn self.max_num_writers_per_bundle = max_writers_per_bundle def start_bundle(self): self._writers_and_sinks = {} self._file_names = {} def process( self, record, w=beam.DoFn.WindowParam, pane=beam.DoFn.PaneInfoParam): destination = self.destination_fn(record) writer, sink = self._get_or_create_writer_and_sink(destination, w) if not writer: return [beam.pvalue.TaggedOutput(self.SPILLED_RECORDS, record)] else: sink.write(record) def _get_or_create_writer_and_sink(self, destination, window): """Returns a tuple of writer, sink.""" writer_key = (destination, window) if writer_key in self._writers_and_sinks: return self._writers_and_sinks.get(writer_key) elif len(self._writers_and_sinks) >= self.max_num_writers_per_bundle: # The writer does not exist, and we have too many writers already. return None, None else: # The writer does not exist, but we can still create a new one. sink = self.sink_fn(destination) full_file_name, writer = _create_writer( base_path=self.base_path.get(), writer_key=writer_key, create_metadata_fn=sink.create_metadata) sink.open(writer) self._writers_and_sinks[writer_key] = (writer, sink) self._file_names[writer_key] = full_file_name return self._writers_and_sinks[writer_key] def finish_bundle(self): for key, (writer, sink) in self._writers_and_sinks.items(): sink.flush() writer.close() file_result = FileResult( self._file_names[key], shard_index=-1, total_shards=0, window=key[1], pane=None, # TODO(pabloem): get the pane info destination=key[0]) yield beam.pvalue.TaggedOutput( self.WRITTEN_FILES, beam.transforms.window.WindowedValue( file_result, timestamp=key[1].start, windows=[key[1]] # TODO(pabloem) HOW DO WE GET THE PANE )) class _RemoveDuplicates(beam.DoFn): """Internal DoFn that filters out filenames already seen (even though the file has updated).""" COUNT_STATE = CombiningValueStateSpec('count', combine_fn=sum) def process( self, element: tuple[str, filesystem.FileMetadata], count_state=beam.DoFn.StateParam(COUNT_STATE) ) -> Iterable[filesystem.FileMetadata]: path = element[0] file_metadata = element[1] counter = count_state.read() if counter == 0: count_state.add(1) _LOGGER.debug('Generated entry for file %s', path) yield file_metadata else: _LOGGER.debug('File %s was already read, seen %d times', path, counter) class _RemoveOldDuplicates(beam.DoFn): """Internal DoFn that filters out filenames already seen and timestamp unchanged.""" TIME_STATE = CombiningValueStateSpec( 'count', combine_fn=partial(max, default=0.0)) def process( self, element: tuple[str, filesystem.FileMetadata], time_state=beam.DoFn.StateParam(TIME_STATE) ) -> Iterable[filesystem.FileMetadata]: path = element[0] file_metadata = element[1] new_ts = file_metadata.last_updated_in_seconds old_ts = time_state.read() if old_ts < new_ts: time_state.add(new_ts) _LOGGER.debug('Generated entry for file %s', path) yield file_metadata else: _LOGGER.debug('File %s was already read', path) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/fileio_test.py0000644000175100017510000006725615206113032021164 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for transforms defined in apache_beam.io.fileio.""" # pytype: skip-file import csv import io import json import logging import os import unittest import uuid import warnings import pytest from hamcrest.library.text import stringmatches import apache_beam as beam from apache_beam.io import fileio from apache_beam.io.filebasedsink_test import _TestCaseWithTempDirCleanUp from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystems import FileSystems from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.test_utils import compute_hash from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.testing.util import matches_all from apache_beam.transforms import trigger from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import GlobalWindow from apache_beam.transforms.window import IntervalWindow from apache_beam.utils.timestamp import Timestamp warnings.filterwarnings( 'ignore', category=FutureWarning, module='apache_beam.io.fileio_test') def _get_file_reader(readable_file): return io.TextIOWrapper(readable_file.open()) class MatchTest(_TestCaseWithTempDirCleanUp): def test_basic_two_files(self): files = [] tempdir = '%s%s' % (self._new_tempdir(), os.sep) # Create a couple files to be matched files.append(self._create_temp_file(dir=tempdir)) files.append(self._create_temp_file(dir=tempdir)) with TestPipeline() as p: files_pc = ( p | fileio.MatchFiles(FileSystems.join(tempdir, '*')) | beam.Map(lambda x: x.path)) assert_that(files_pc, equal_to(files)) def test_match_all_two_directories(self): files = [] directories = [] for _ in range(2): # TODO: What about this having to append the ending slash? d = '%s%s' % (self._new_tempdir(), os.sep) directories.append(d) files.append(self._create_temp_file(dir=d)) files.append(self._create_temp_file(dir=d)) with TestPipeline() as p: files_pc = ( p | beam.Create([FileSystems.join(d, '*') for d in directories]) | fileio.MatchAll() | beam.Map(lambda x: x.path)) assert_that(files_pc, equal_to(files)) def test_match_files_one_directory_failure1(self): directories = [ '%s%s' % (self._new_tempdir(), os.sep), '%s%s' % (self._new_tempdir(), os.sep) ] files = [] files.append(self._create_temp_file(dir=directories[0])) files.append(self._create_temp_file(dir=directories[0])) with self.assertRaisesRegex(Exception, "Empty match for pattern"): with TestPipeline() as p: files_pc = ( p | beam.Create([FileSystems.join(d, '*') for d in directories]) | fileio.MatchAll(fileio.EmptyMatchTreatment.DISALLOW) | beam.Map(lambda x: x.path)) assert_that(files_pc, equal_to(files)) def test_match_files_one_directory_failure2(self): directories = [ '%s%s' % (self._new_tempdir(), os.sep), '%s%s' % (self._new_tempdir(), os.sep) ] files = [] files.append(self._create_temp_file(dir=directories[0])) files.append(self._create_temp_file(dir=directories[0])) with TestPipeline() as p: files_pc = ( p | beam.Create([FileSystems.join(d, '*') for d in directories]) | fileio.MatchAll(fileio.EmptyMatchTreatment.ALLOW_IF_WILDCARD) | beam.Map(lambda x: x.path)) assert_that(files_pc, equal_to(files)) class ReadTest(_TestCaseWithTempDirCleanUp): def test_basic_file_name_provided(self): content = 'TestingMyContent\nIn multiple lines\nhaha!' dir = '%s%s' % (self._new_tempdir(), os.sep) self._create_temp_file(dir=dir, content=content) with TestPipeline() as p: content_pc = ( p | beam.Create([FileSystems.join(dir, '*')]) | fileio.MatchAll() | fileio.ReadMatches() | beam.FlatMap(lambda f: f.read().decode('utf-8').splitlines())) assert_that(content_pc, equal_to(content.splitlines())) def test_csv_file_source(self): content = 'name,year,place\ngoogle,1999,CA\nspotify,2006,sweden' rows = [r.split(',') for r in content.split('\n')] dir = '%s%s' % (self._new_tempdir(), os.sep) self._create_temp_file(dir=dir, content=content) with TestPipeline() as p: content_pc = ( p | beam.Create([FileSystems.join(dir, '*')]) | fileio.MatchAll() | fileio.ReadMatches() | beam.FlatMap(lambda rf: csv.reader(_get_file_reader(rf)))) assert_that(content_pc, equal_to(rows)) def test_infer_compressed_file(self): dir = '%s%s' % (self._new_tempdir(), os.sep) file_contents = b'compressed_contents!' import gzip with gzip.GzipFile(os.path.join(dir, 'compressed.gz'), 'w') as f: f.write(file_contents) file_contents2 = b'compressed_contents_bz2!' import bz2 with bz2.BZ2File(os.path.join(dir, 'compressed2.bz2'), 'w') as f: f.write(file_contents2) with TestPipeline() as p: content_pc = ( p | beam.Create([FileSystems.join(dir, '*')]) | fileio.MatchAll() | fileio.ReadMatches() | beam.Map(lambda rf: rf.open().readline())) assert_that(content_pc, equal_to([file_contents, file_contents2])) def test_read_bz2_compressed_file_without_suffix(self): dir = '%s%s' % (self._new_tempdir(), os.sep) file_contents = b'compressed_contents!' import bz2 with bz2.BZ2File(os.path.join(dir, 'compressed'), 'w') as f: f.write(file_contents) with TestPipeline() as p: content_pc = ( p | beam.Create([FileSystems.join(dir, '*')]) | fileio.MatchAll() | fileio.ReadMatches() | beam.Map( lambda rf: rf.open(compression_type=CompressionTypes.BZIP2).read( len(file_contents)))) assert_that(content_pc, equal_to([file_contents])) def test_read_gzip_compressed_file_without_suffix(self): dir = '%s%s' % (self._new_tempdir(), os.sep) file_contents = b'compressed_contents!' import gzip with gzip.GzipFile(os.path.join(dir, 'compressed'), 'w') as f: f.write(file_contents) with TestPipeline() as p: content_pc = ( p | beam.Create([FileSystems.join(dir, '*')]) | fileio.MatchAll() | fileio.ReadMatches() | beam.Map( lambda rf: rf.open(compression_type=CompressionTypes.GZIP).read( len(file_contents)))) assert_that(content_pc, equal_to([file_contents])) def test_string_filenames_and_skip_directory(self): content = 'thecontent\n' files = [] tempdir = '%s%s' % (self._new_tempdir(), os.sep) # Create a couple files to be matched files.append(self._create_temp_file(dir=tempdir, content=content)) files.append(self._create_temp_file(dir=tempdir, content=content)) with TestPipeline() as p: contents_pc = ( p | beam.Create(files + ['%s/' % tempdir]) | fileio.ReadMatches() | beam.FlatMap(lambda x: x.read().decode('utf-8').splitlines())) assert_that(contents_pc, equal_to(content.splitlines() * 2)) def test_fail_on_directories(self): content = 'thecontent\n' files = [] tempdir = '%s%s' % (self._new_tempdir(), os.sep) # Create a couple files to be matched files.append(self._create_temp_file(dir=tempdir, content=content)) files.append(self._create_temp_file(dir=tempdir, content=content)) with self.assertRaisesRegex(Exception, "Directories are not allowed"): with TestPipeline() as p: _ = ( p | beam.Create(files + ['%s/' % tempdir]) | fileio.ReadMatches(skip_directories=False) | beam.Map(lambda x: x.read_utf8())) class MatchIntegrationTest(unittest.TestCase): INPUT_FILE = 'gs://dataflow-samples/shakespeare/kinglear.txt' KINGLEAR_CHECKSUM = 'f418b25f1507f5a901257026b035ac2857a7ab87' INPUT_FILE_LARGE = ( 'gs://dataflow-samples/wikipedia_edits/wiki_data-00000000000*.json') WIKI_FILES = [ 'gs://dataflow-samples/wikipedia_edits/wiki_data-000000000000.json', 'gs://dataflow-samples/wikipedia_edits/wiki_data-000000000001.json', 'gs://dataflow-samples/wikipedia_edits/wiki_data-000000000002.json', 'gs://dataflow-samples/wikipedia_edits/wiki_data-000000000003.json', 'gs://dataflow-samples/wikipedia_edits/wiki_data-000000000004.json', 'gs://dataflow-samples/wikipedia_edits/wiki_data-000000000005.json', 'gs://dataflow-samples/wikipedia_edits/wiki_data-000000000006.json', 'gs://dataflow-samples/wikipedia_edits/wiki_data-000000000007.json', 'gs://dataflow-samples/wikipedia_edits/wiki_data-000000000008.json', 'gs://dataflow-samples/wikipedia_edits/wiki_data-000000000009.json', ] def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) @pytest.mark.it_postcommit def test_transform_on_gcs(self): args = self.test_pipeline.get_full_options_as_args() with beam.Pipeline(argv=args) as p: matches_pc = ( p | beam.Create([self.INPUT_FILE, self.INPUT_FILE_LARGE]) | fileio.MatchAll() | 'GetPath' >> beam.Map(lambda metadata: metadata.path)) assert_that( matches_pc, equal_to([self.INPUT_FILE] + self.WIKI_FILES), label='Matched Files') checksum_pc = ( p | 'SingleFile' >> beam.Create([self.INPUT_FILE]) | 'MatchOneAll' >> fileio.MatchAll() | fileio.ReadMatches() | 'ReadIn' >> beam.Map(lambda x: x.read_utf8().split('\n')) | 'Checksums' >> beam.Map(compute_hash)) assert_that( checksum_pc, equal_to([self.KINGLEAR_CHECKSUM]), label='Assert Checksums') class MatchContinuouslyTest(_TestCaseWithTempDirCleanUp): def test_with_deduplication(self): files = [] tempdir = '%s%s' % (self._new_tempdir(), os.sep) # Create a file to be matched before pipeline files.append(self._create_temp_file(dir=tempdir)) # Add file name that will be created mid-pipeline files.append(FileSystems.join(tempdir, 'extra')) interval = 0.2 start = Timestamp.now() stop = start + interval + 0.1 def _create_extra_file(element): writer = FileSystems.create(FileSystems.join(tempdir, 'extra')) writer.close() return element.path with TestPipeline() as p: match_continiously = ( p | fileio.MatchContinuously( file_pattern=FileSystems.join(tempdir, '*'), interval=interval, start_timestamp=start, stop_timestamp=stop) | beam.Map(_create_extra_file)) assert_that(match_continiously, equal_to(files)) def test_without_deduplication(self): interval = 0.2 start = Timestamp.now() stop = start + interval + 0.1 files = [] tempdir = '%s%s' % (self._new_tempdir(), os.sep) # Create a file to be matched before pipeline starts file = self._create_temp_file(dir=tempdir) # Add file twice, since it will be matched for every interval files += [file, file] # Add file name that will be created mid-pipeline files.append(FileSystems.join(tempdir, 'extra')) def _create_extra_file(element): writer = FileSystems.create(FileSystems.join(tempdir, 'extra')) writer.close() return element.path with TestPipeline() as p: match_continiously = ( p | fileio.MatchContinuously( file_pattern=FileSystems.join(tempdir, '*'), interval=interval, has_deduplication=False, start_timestamp=start, stop_timestamp=stop) | beam.Map(_create_extra_file)) assert_that(match_continiously, equal_to(files)) def test_match_updated_files(self): files = [] tempdir = '%s%s' % (self._new_tempdir(), os.sep) def _create_extra_file(element): writer = FileSystems.create(FileSystems.join(tempdir, 'extra')) writer.close() return element.path # Create two files to be matched before pipeline files.append(self._create_temp_file(dir=tempdir)) writer = FileSystems.create(FileSystems.join(tempdir, 'extra')) writer.close() # Add file name that will be created mid-pipeline files.append(FileSystems.join(tempdir, 'extra')) files.append(FileSystems.join(tempdir, 'extra')) interval = 0.2 start = Timestamp.now() stop = start + interval + 0.1 with TestPipeline() as p: match_continiously = ( p | fileio.MatchContinuously( file_pattern=FileSystems.join(tempdir, '*'), interval=interval, start_timestamp=start, stop_timestamp=stop, match_updated_files=True) | beam.Map(_create_extra_file)) assert_that(match_continiously, equal_to(files)) class WriteFilesTest(_TestCaseWithTempDirCleanUp): SIMPLE_COLLECTION = [ { 'project': 'beam', 'foundation': 'apache' }, { 'project': 'prometheus', 'foundation': 'cncf' }, { 'project': 'flink', 'foundation': 'apache' }, { 'project': 'grpc', 'foundation': 'cncf' }, { 'project': 'spark', 'foundation': 'apache' }, { 'project': 'kubernetes', 'foundation': 'cncf' }, { 'project': 'spark', 'foundation': 'apache' }, { 'project': 'knative', 'foundation': 'cncf' }, { 'project': 'linux', 'foundation': 'linux' }, ] LARGER_COLLECTION = ['{:05d}'.format(i) for i in range(200)] CSV_HEADERS = ['project', 'foundation'] SIMPLE_COLLECTION_VALIDATION_SET = {(elm['project'], elm['foundation']) for elm in SIMPLE_COLLECTION} class CsvSink(fileio.TextSink): def __init__(self, headers): self.headers = headers def write(self, record): self._fh.write(','.join([record[h] for h in self.headers]).encode('utf8')) self._fh.write('\n'.encode('utf8')) class JsonSink(fileio.TextSink): def write(self, record): self._fh.write(json.dumps(record).encode('utf8')) self._fh.write('\n'.encode('utf8')) def test_write_to_single_file_batch(self): dir = self._new_tempdir() with TestPipeline() as p: _ = ( p | beam.Create(WriteFilesTest.SIMPLE_COLLECTION) | "Serialize" >> beam.Map(json.dumps) | beam.io.fileio.WriteToFiles(path=dir)) with TestPipeline() as p: result = ( p | fileio.MatchFiles(FileSystems.join(dir, '*')) | fileio.ReadMatches() | beam.FlatMap(lambda f: f.read_utf8().strip().split('\n')) | beam.Map(json.loads)) assert_that(result, equal_to([row for row in self.SIMPLE_COLLECTION])) def test_write_to_dynamic_destination(self): sink_params = [ fileio.TextSink, # pass a type signature fileio.TextSink() # pass a FileSink object ] # Test assumes that all records will be handled by same worker process, # pin to FnApiRunner to guarantee hthis runner = 'FnApiRunner' for sink in sink_params: dir = self._new_tempdir() with TestPipeline(runner) as p: _ = ( p | "Create" >> beam.Create(range(100)) | beam.Map(lambda x: str(x)) | fileio.WriteToFiles( path=dir, destination=lambda n: "odd" if int(n) % 2 else "even", sink=sink, file_naming=fileio.destination_prefix_naming("test"))) with TestPipeline(runner) as p: result = ( p | fileio.MatchFiles(FileSystems.join(dir, '*')) | fileio.ReadMatches() | beam.Map( lambda f: ( os.path.basename(f.metadata.path).split('-')[0], sorted( map(int, f.read_utf8().strip().split('\n')))))) assert_that( result, equal_to([('odd', list(range(1, 100, 2))), ('even', list(range(0, 100, 2)))])) def test_write_to_different_file_types_some_spilling(self): dir = self._new_tempdir() with TestPipeline() as p: _ = ( p | beam.Create(WriteFilesTest.SIMPLE_COLLECTION) | beam.io.fileio.WriteToFiles( path=dir, destination=lambda record: record['foundation'], sink=lambda dest: ( WriteFilesTest.CsvSink(WriteFilesTest.CSV_HEADERS) if dest == 'apache' else WriteFilesTest.JsonSink()), file_naming=fileio.destination_prefix_naming(), max_writers_per_bundle=1)) with TestPipeline() as p: cncf_res = ( p | fileio.MatchFiles(FileSystems.join(dir, 'cncf*')) | fileio.ReadMatches() | beam.FlatMap(lambda f: f.read_utf8().strip().split('\n')) | beam.Map(json.loads)) apache_res = ( p | "MatchApache" >> fileio.MatchFiles(FileSystems.join(dir, 'apache*')) | "ReadApache" >> fileio.ReadMatches() | "MapApache" >> beam.FlatMap(lambda rf: csv.reader(_get_file_reader(rf)))) assert_that( cncf_res, equal_to([ row for row in self.SIMPLE_COLLECTION if row['foundation'] == 'cncf' ]), label='verifyCNCF') assert_that( apache_res, equal_to([[row['project'], row['foundation']] for row in self.SIMPLE_COLLECTION if row['foundation'] == 'apache']), label='verifyApache') @unittest.skip('https://github.com/apache/beam/issues/21269') def test_find_orphaned_files(self): dir = self._new_tempdir() write_transform = beam.io.fileio.WriteToFiles(path=dir) def write_orphaned_file(temp_dir, writer_key): temp_dir_path = FileSystems.join(dir, temp_dir) file_prefix_dir = FileSystems.join( temp_dir_path, str(abs(hash(writer_key)))) file_name = '%s_%s' % (file_prefix_dir, uuid.uuid4()) with FileSystems.create(file_name) as f: f.write(b'Hello y\'all') return file_name with TestPipeline() as p: _ = ( p | beam.Create(WriteFilesTest.SIMPLE_COLLECTION) | "Serialize" >> beam.Map(json.dumps) | write_transform) # Pre-create the temp directory. temp_dir_path = FileSystems.mkdirs( FileSystems.join(dir, write_transform._temp_directory.get())) write_orphaned_file( write_transform._temp_directory.get(), (None, GlobalWindow())) f2 = write_orphaned_file( write_transform._temp_directory.get(), ('other-dest', GlobalWindow())) temp_dir_path = FileSystems.join(dir, write_transform._temp_directory.get()) leftovers = FileSystems.match(['%s%s*' % (temp_dir_path, os.sep)]) found_files = [m.path for m in leftovers[0].metadata_list] self.assertListEqual(found_files, [f2]) def test_write_to_different_file_types(self): dir = self._new_tempdir() with TestPipeline() as p: _ = ( p | beam.Create(WriteFilesTest.SIMPLE_COLLECTION) | beam.io.fileio.WriteToFiles( path=dir, destination=lambda record: record['foundation'], sink=lambda dest: ( WriteFilesTest.CsvSink(WriteFilesTest.CSV_HEADERS) if dest == 'apache' else WriteFilesTest.JsonSink()), file_naming=fileio.destination_prefix_naming())) with TestPipeline() as p: cncf_res = ( p | fileio.MatchFiles(FileSystems.join(dir, 'cncf*')) | fileio.ReadMatches() | beam.FlatMap(lambda f: f.read_utf8().strip().split('\n')) | beam.Map(json.loads)) apache_res = ( p | "MatchApache" >> fileio.MatchFiles(FileSystems.join(dir, 'apache*')) | "ReadApache" >> fileio.ReadMatches() | "MapApache" >> beam.FlatMap(lambda rf: csv.reader(_get_file_reader(rf)))) assert_that( cncf_res, equal_to([ row for row in self.SIMPLE_COLLECTION if row['foundation'] == 'cncf' ]), label='verifyCNCF') assert_that( apache_res, equal_to([[row['project'], row['foundation']] for row in self.SIMPLE_COLLECTION if row['foundation'] == 'apache']), label='verifyApache') def record_dofn(self): class RecordDoFn(beam.DoFn): def process(self, element): WriteFilesTest.all_records.append(element) return RecordDoFn() def test_streaming_complex_timing(self): # Use state on the TestCase class, since other references would be pickled # into a closure and not have the desired side effects. # # TODO(https://github.com/apache/beam/issues/18987): Use assert_that after # it works for the cases here in streaming mode. WriteFilesTest.all_records = [] dir = '%s%s' % (self._new_tempdir(), os.sep) # Setting up the input (TestStream) ts = TestStream().advance_watermark_to(0) for elm in WriteFilesTest.LARGER_COLLECTION: timestamp = int(elm) ts.add_elements([('key', '%s' % elm)]) if timestamp % 5 == 0 and timestamp != 0: # TODO(https://github.com/apache/beam/issues/18721): Add many firings # per window after getting PaneInfo. ts.advance_processing_time(5) ts.advance_watermark_to(timestamp) ts.advance_watermark_to_infinity() def no_colon_file_naming(*args): file_name = fileio.destination_prefix_naming()(*args) return file_name.replace(':', '_') # The pipeline that we are testing options = PipelineOptions() options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: res = ( p | ts | beam.WindowInto( FixedWindows(10), trigger=trigger.AfterWatermark(), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.GroupByKey() | beam.FlatMap(lambda x: x[1])) # Triggering after 5 processing-time seconds, and on the watermark. Also # discarding old elements. _ = ( res | beam.io.fileio.WriteToFiles( path=dir, file_naming=no_colon_file_naming, max_writers_per_bundle=0) | beam.Map(lambda fr: FileSystems.join(dir, fr.file_name)) | beam.ParDo(self.record_dofn())) # Verification pipeline with TestPipeline() as p: files = (p | beam.io.fileio.MatchFiles(FileSystems.join(dir, '*'))) file_names = (files | beam.Map(lambda fm: fm.path)) file_contents = ( files | beam.io.fileio.ReadMatches() | beam.Map( lambda rf: (rf.metadata.path, rf.read_utf8().strip().split('\n'))) ) content = ( file_contents | beam.FlatMap(lambda fc: [ln.strip() for ln in fc[1]])) assert_that( file_names, equal_to(WriteFilesTest.all_records), label='AssertFilesMatch') assert_that( content, matches_all(WriteFilesTest.LARGER_COLLECTION), label='AssertContentsMatch') def test_streaming_different_file_types(self): dir = self._new_tempdir() input = iter(WriteFilesTest.SIMPLE_COLLECTION) ts = ( TestStream().advance_watermark_to(0).add_elements( [next(input), next(input)]).advance_watermark_to(10).add_elements( [next(input), next(input)]).advance_watermark_to(20).add_elements([ next(input), next(input) ]).advance_watermark_to(30).add_elements([ next(input), next(input) ]).advance_watermark_to(40).advance_watermark_to_infinity()) def no_colon_file_naming(*args): file_name = fileio.destination_prefix_naming()(*args) return file_name.replace(':', '_') with TestPipeline() as p: _ = ( p | ts | beam.WindowInto(FixedWindows(10)) | beam.io.fileio.WriteToFiles( path=dir, destination=lambda record: record['foundation'], sink=lambda dest: ( WriteFilesTest.CsvSink(WriteFilesTest.CSV_HEADERS) if dest == 'apache' else WriteFilesTest.JsonSink()), file_naming=no_colon_file_naming, max_writers_per_bundle=0, )) with TestPipeline() as p: cncf_files = ( p | fileio.MatchFiles(FileSystems.join(dir, 'cncf*')) | "CncfFileNames" >> beam.Map(lambda fm: fm.path)) apache_files = ( p | "MatchApache" >> fileio.MatchFiles(FileSystems.join(dir, 'apache*')) | "ApacheFileNames" >> beam.Map(lambda fm: fm.path)) assert_that( cncf_files, matches_all([ stringmatches.matches_regexp( '.*cncf-1970-01-01T00_00_00-1970-01-01T00_00_10.*'), stringmatches.matches_regexp( '.*cncf-1970-01-01T00_00_10-1970-01-01T00_00_20.*'), stringmatches.matches_regexp( '.*cncf-1970-01-01T00_00_20-1970-01-01T00_00_30.*'), stringmatches.matches_regexp( '.*cncf-1970-01-01T00_00_30-1970-01-01T00_00_40.*') ]), label='verifyCNCFFiles') assert_that( apache_files, matches_all([ stringmatches.matches_regexp( '.*apache-1970-01-01T00_00_00-1970-01-01T00_00_10.*'), stringmatches.matches_regexp( '.*apache-1970-01-01T00_00_10-1970-01-01T00_00_20.*'), stringmatches.matches_regexp( '.*apache-1970-01-01T00_00_20-1970-01-01T00_00_30.*'), stringmatches.matches_regexp( '.*apache-1970-01-01T00_00_30-1970-01-01T00_00_40.*') ]), label='verifyApacheFiles') def test_shard_naming(self): namer = fileio.default_file_naming(prefix='/path/to/file', suffix='.txt') self.assertEqual( namer(GlobalWindow(), None, None, None, None, None), '/path/to/file.txt') self.assertEqual( namer(GlobalWindow(), None, 1, 5, None, None), '/path/to/file-00001-of-00005.txt') self.assertEqual( namer(GlobalWindow(), None, 1, 5, 'gz', None), '/path/to/file-00001-of-00005.txt.gz') self.assertEqual( namer(IntervalWindow(0, 100), None, 1, 5, None, None), '/path/to/file' '-1970-01-01T00:00:00-1970-01-01T00:01:40-00001-of-00005.txt') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/filesystem.py0000644000175100017510000007523615206113032021037 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """File system abstraction for file-based sources and sinks. Note to implementors: "path" arguments will be URLs in the form scheme://foo/bar. The exception is LocalFileSystem, which gets unix-style paths in the form /foo/bar. """ # pytype: skip-file import abc import bz2 import io import logging import lzma import os import posixpath import re import time import zlib from typing import BinaryIO # pylint: disable=unused-import from typing import Iterator from typing import Optional import zstandard from apache_beam.utils.plugin import BeamPlugin logger = logging.getLogger(__name__) DEFAULT_READ_BUFFER_SIZE = 16 * 1024 * 1024 __all__ = [ 'CompressionTypes', 'CompressedFile', 'FileMetadata', 'FileSystem', 'MatchResult' ] class CompressionTypes(object): """Enum-like class representing known compression types.""" # Detect compression based on filename extension. # # The following extensions are currently recognized by auto-detection: # .bz2 (implies BZIP2 as described below). # .gz (implies GZIP as described below) # .deflate (implies DEFLATE as described below) # .zst (implies ZSTD as described below) # .zst (implies ZSTD as described below) # .xz (implies LZMA as described below) # .lzma (implies LZMA as described below) # Any non-recognized extension implies UNCOMPRESSED as described below. AUTO = 'auto' # BZIP2 compression. BZIP2 = 'bzip2' # DEFLATE compression DEFLATE = 'deflate' # ZSTD compression ZSTD = 'zstd' # GZIP compression (deflate with GZIP headers). GZIP = 'gzip' # LZMA compression LZMA = 'lzma' # Uncompressed (i.e., may be split). UNCOMPRESSED = 'uncompressed' @classmethod def is_valid_compression_type(cls, compression_type): """Returns True for valid compression types, False otherwise.""" types = set([ CompressionTypes.AUTO, CompressionTypes.BZIP2, CompressionTypes.DEFLATE, CompressionTypes.GZIP, CompressionTypes.ZSTD, CompressionTypes.LZMA, CompressionTypes.UNCOMPRESSED ]) return compression_type in types @classmethod def mime_type(cls, compression_type, default='application/octet-stream'): mime_types_by_compression_type = { cls.BZIP2: 'application/x-bz2', cls.DEFLATE: 'application/x-deflate', cls.GZIP: 'application/x-gzip', cls.ZSTD: 'application/zstd', cls.LZMA: 'application/lzma' } return mime_types_by_compression_type.get(compression_type, default) @classmethod def detect_compression_type(cls, file_path): """Returns the compression type of a file (based on its suffix).""" compression_types_by_suffix = { '.bz2': cls.BZIP2, '.deflate': cls.DEFLATE, '.gz': cls.GZIP, '.zst': cls.ZSTD, '.zstd': cls.ZSTD, '.xz': cls.LZMA, '.lzma': cls.LZMA } lowercased_path = file_path.lower() for suffix, compression_type in compression_types_by_suffix.items(): if lowercased_path.endswith(suffix): return compression_type return cls.UNCOMPRESSED class CompressedFile(object): """File wrapper for easier handling of compressed files.""" # XXX: This class is not thread safe in the read path. # The bit mask to use for the wbits parameters of the zlib compressor and # decompressor objects. _gzip_mask = zlib.MAX_WBITS | 16 # Mask when using GZIP headers. def __init__( self, fileobj: BinaryIO, compression_type=CompressionTypes.GZIP, read_size=DEFAULT_READ_BUFFER_SIZE): if not fileobj: raise ValueError('File object must not be None') if not CompressionTypes.is_valid_compression_type(compression_type): raise TypeError( 'compression_type must be CompressionType object but ' 'was %s' % type(compression_type)) if compression_type in (CompressionTypes.AUTO, CompressionTypes.UNCOMPRESSED): raise ValueError( 'Cannot create object with unspecified or no compression') self._file = fileobj self._compression_type = compression_type if self._file.tell() != 0: raise ValueError( 'File object must be at position 0 but was %d' % self._file.tell()) self._uncompressed_position = 0 self._uncompressed_size: Optional[int] = None if self.readable(): self._read_size = read_size self._read_buffer = io.BytesIO() self._read_position = 0 self._read_eof = False self._initialize_decompressor() else: self._decompressor = None if self.writeable(): self._initialize_compressor() else: self._compressor = None def _initialize_decompressor(self): if self._compression_type == CompressionTypes.BZIP2: self._decompressor = bz2.BZ2Decompressor() elif self._compression_type == CompressionTypes.DEFLATE: self._decompressor = zlib.decompressobj() elif self._compression_type == CompressionTypes.ZSTD: # hardcoded max_window_size to avoid too much memory # errors when reading big files, please refer # to the following issue for further explanation: # https://github.com/indygreg/python-zstandard/issues/157 self._decompressor = zstandard.ZstdDecompressor( max_window_size=2147483648).decompressobj() elif self._compression_type == CompressionTypes.LZMA: self._decompressor = lzma.LZMADecompressor() else: assert self._compression_type == CompressionTypes.GZIP self._decompressor = zlib.decompressobj(self._gzip_mask) def _initialize_compressor(self): if self._compression_type == CompressionTypes.BZIP2: self._compressor = bz2.BZ2Compressor() elif self._compression_type == CompressionTypes.DEFLATE: self._compressor = zlib.compressobj( zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED) elif self._compression_type == CompressionTypes.ZSTD: self._compressor = zstandard.ZstdCompressor().compressobj() elif self._compression_type == CompressionTypes.LZMA: self._compressor = lzma.LZMACompressor() else: assert self._compression_type == CompressionTypes.GZIP self._compressor = zlib.compressobj( zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, self._gzip_mask) def readable(self) -> bool: mode = self._file.mode return 'r' in mode or 'a' in mode def writeable(self) -> bool: mode = self._file.mode return 'w' in mode or 'a' in mode def write(self, data: bytes) -> None: """Write data to file.""" if not self._compressor: raise ValueError('compressor not initialized') self._uncompressed_position += len(data) compressed = self._compressor.compress(data) if compressed: self._file.write(compressed) def _fetch_to_internal_buffer(self, num_bytes: int) -> None: """Fetch up to num_bytes into the internal buffer.""" if (not self._read_eof and self._read_position > 0 and (self._read_buffer.tell() - self._read_position) < num_bytes): # There aren't enough number of bytes to accommodate a read, so we # prepare for a possibly large read by clearing up all internal buffers # but without dropping any previous held data. self._read_buffer.seek(self._read_position) data = self._read_buffer.read() self._clear_read_buffer() self._read_buffer.write(data) assert self._decompressor while not self._read_eof and (self._read_buffer.tell() - self._read_position) < num_bytes: # Continue reading from the underlying file object until enough bytes are # available, or EOF is reached. if not self._decompressor.unused_data: buf = self._file.read(self._read_size) else: # Any uncompressed data at the end of the stream of a gzip or bzip2 # file that is not corrupted points to a concatenated compressed # file. We read concatenated files by recursively creating decompressor # objects for the unused compressed data. buf = self._decompressor.unused_data self._initialize_decompressor() if buf: decompressed = self._decompressor.decompress(buf) del buf # Free up some possibly large and no-longer-needed memory. self._read_buffer.write(decompressed) else: # EOF of current stream reached. if (self._compression_type == CompressionTypes.BZIP2 or self._compression_type == CompressionTypes.DEFLATE or self._compression_type == CompressionTypes.ZSTD or self._compression_type == CompressionTypes.GZIP or self._compression_type == CompressionTypes.LZMA): pass else: # Deflate, Gzip and bzip2 formats do not require flushing # remaining data in the decompressor into the read buffer when # fully decompressing files. self._read_buffer.write(self._decompressor.flush()) # Record that we have hit the end of file, so we won't unnecessarily # repeat the completeness verification step above. self._read_eof = True def _read_from_internal_buffer(self, read_fn): """Read from the internal buffer by using the supplied read_fn.""" self._read_buffer.seek(self._read_position) result = read_fn() self._read_position += len(result) self._uncompressed_position += len(result) self._read_buffer.seek(0, os.SEEK_END) # Allow future writes. return result def read(self, num_bytes: Optional[int] = DEFAULT_READ_BUFFER_SIZE) -> bytes: if not self._decompressor: raise ValueError('decompressor not initialized') self._fetch_to_internal_buffer(num_bytes) return self._read_from_internal_buffer( lambda: self._read_buffer.read(num_bytes)) def readline(self) -> bytes: """Equivalent to standard file.readline(). Same return conventions apply.""" if not self._decompressor: raise ValueError('decompressor not initialized') bytes_io = io.BytesIO() while True: # Ensure that the internal buffer has at least half the read_size. Going # with half the _read_size (as opposed to a full _read_size) to ensure # that actual fetches are more evenly spread out, as opposed to having 2 # consecutive reads at the beginning of a read. self._fetch_to_internal_buffer(self._read_size // 2) line = self._read_from_internal_buffer( lambda: self._read_buffer.readline()) bytes_io.write(line) if line.endswith(b'\n') or not line: break # Newline or EOF reached. return bytes_io.getvalue() def closed(self) -> bool: return not self._file or self._file.closed def close(self) -> None: if self.readable(): self._read_buffer.close() if self.writeable(): assert self._compressor self._file.write(self._compressor.flush()) self._file.close() def flush(self) -> None: if self.writeable(): assert self._compressor self._file.write(self._compressor.flush()) self._file.flush() @property def seekable(self) -> bool: return 'r' in self._file.mode def _clear_read_buffer(self) -> None: """Clears the read buffer by removing all the contents and resetting _read_position to 0""" self._read_position = 0 self._read_buffer.seek(0) self._read_buffer.truncate(0) def _rewind_file(self) -> None: """Seeks to the beginning of the input file. Input file's EOF marker is cleared and _uncompressed_position is reset to zero""" self._file.seek(0, os.SEEK_SET) self._read_eof = False self._uncompressed_position = 0 def _rewind(self) -> None: """Seeks to the beginning of the input file and resets the internal read buffer. The decompressor object is re-initialized to ensure that no data left in it's buffer.""" self._clear_read_buffer() self._rewind_file() # Re-initialize decompressor to clear any data buffered prior to rewind self._initialize_decompressor() def seek(self, offset: int, whence: int = os.SEEK_SET) -> None: """Set the file's current offset. Seeking behavior: * seeking from the end :data:`os.SEEK_END` the whole file is decompressed once to determine its size. Therefore it is preferred to use :data:`os.SEEK_SET` or :data:`os.SEEK_CUR` to avoid the processing overhead * seeking backwards from the current position rewinds the file to ``0`` and decompresses the chunks to the requested offset * seeking is only supported in files opened for reading * if the new offset is out of bound, it is adjusted to either ``0`` or ``EOF``. Args: offset (int): seek offset in the uncompressed content represented as number whence (int): seek mode. Supported modes are :data:`os.SEEK_SET` (absolute seek), :data:`os.SEEK_CUR` (seek relative to the current position), and :data:`os.SEEK_END` (seek relative to the end, offset should be negative). Raises: IOError: When this buffer is closed. ValueError: When whence is invalid or the file is not seekable """ if whence == os.SEEK_SET: absolute_offset = offset elif whence == os.SEEK_CUR: absolute_offset = self._uncompressed_position + offset elif whence == os.SEEK_END: # Determine and cache the uncompressed size of the file if not self._uncompressed_size: logger.warning( "Seeking relative from end of file is requested. " "Need to decompress the whole file once to determine " "its size. This might take a while...") uncompress_start_time = time.time() while self.read(self._read_size): pass uncompress_end_time = time.time() logger.warning( "Full file decompression for seek " "from end took %.2f secs", (uncompress_end_time - uncompress_start_time)) self._uncompressed_size = self._uncompressed_position absolute_offset = self._uncompressed_size + offset else: raise ValueError("Whence mode %r is invalid." % whence) # Determine how many bytes needs to be read before we reach # the requested offset. Rewind if we already passed the position. if absolute_offset < self._uncompressed_position: self._rewind() bytes_to_skip = absolute_offset - self._uncompressed_position # Read until the desired position is reached or EOF occurs. while bytes_to_skip: data = self.read(min(self._read_size, bytes_to_skip)) if not data: break bytes_to_skip -= len(data) def tell(self) -> int: """Returns current position in uncompressed file.""" return self._uncompressed_position def __enter__(self): return self def __exit__(self, exception_type, exception_value, traceback): self.close() class FileMetadata(object): """Metadata about a file path that is the output of FileSystem.match. Fields: path: [Required] file path. size_in_bytes: [Required] file size in bytes. last_updated_in_seconds: [Optional] last modified timestamp of the file, or valued 0.0 if not specified. """ def __init__( self, path: str, size_in_bytes: int, last_updated_in_seconds: float = 0.0): assert isinstance(path, str) and path, "Path should be a string" assert isinstance(size_in_bytes, int) and size_in_bytes >= 0, \ "Invalid value for size_in_bytes should %s (of type %s)" % ( size_in_bytes, type(size_in_bytes)) self.path = path self.size_in_bytes = size_in_bytes self.last_updated_in_seconds = last_updated_in_seconds def __eq__(self, other): """Note: This is only used in tests where we verify that mock objects match. """ return ( isinstance(other, FileMetadata) and self.path == other.path and self.size_in_bytes == other.size_in_bytes and self.last_updated_in_seconds == other.last_updated_in_seconds) def __hash__(self): return hash((self.path, self.size_in_bytes, self.last_updated_in_seconds)) def __repr__(self): if self.last_updated_in_seconds == 0.0: return 'FileMetadata(%s, %s)' % (self.path, self.size_in_bytes) else: return 'FileMetadata(%s, %s, %s)' % ( self.path, self.size_in_bytes, self.last_updated_in_seconds) class MatchResult(object): """Result from the ``FileSystem`` match operation which contains the list of matched ``FileMetadata``. """ def __init__(self, pattern: str, metadata_list: list[FileMetadata]) -> None: self.metadata_list = metadata_list self.pattern = pattern class BeamIOError(IOError): def __init__(self, msg, exception_details=None): """Class representing the errors thrown in the batch file operations. Args: msg: Message string for the exception thrown exception_details: Optional map of individual input to exception for failed operations in batch. This parameter is optional so if specified the user can assume that the all errors in the filesystem operation have been reported. When the details are missing then the operation may have failed anywhere so the user should use match to determine the current state of the system. """ message = "%s with exceptions %s" % (msg, exception_details) super().__init__(message) self.exception_details = exception_details class FileSystem(BeamPlugin, metaclass=abc.ABCMeta): """A class that defines the functions that can be performed on a filesystem. All methods are abstract and they are for file system providers to implement. Clients should use the FileSystems class to interact with the correct file system based on the provided file pattern scheme. """ CHUNK_SIZE = 1 # Chuck size in the batch operations def __init__(self, pipeline_options): """ Args: pipeline_options: Instance of ``PipelineOptions`` or dict of options and values (like ``RuntimeValueProvider.runtime_options``). """ @staticmethod def _get_compression_type(path, compression_type): if compression_type == CompressionTypes.AUTO: compression_type = CompressionTypes.detect_compression_type(path) elif not CompressionTypes.is_valid_compression_type(compression_type): raise TypeError( 'compression_type must be CompressionType object but ' 'was %s' % type(compression_type)) return compression_type @classmethod def scheme(cls): """URI scheme for the FileSystem """ raise NotImplementedError @abc.abstractmethod def join(self, basepath: str, *paths: str) -> str: """Join two or more pathname components for the filesystem Args: basepath: string path of the first component of the path paths: path components to be added Returns: full path after combining all the passed components """ raise NotImplementedError @abc.abstractmethod def split(self, path: str) -> tuple[str, str]: """Splits the given path into two parts. Splits the path into a pair (head, tail) such that tail contains the last component of the path and head contains everything up to that. For file-systems other than the local file-system, head should include the prefix. Args: path: path as a string Returns: a pair of path components as strings. """ raise NotImplementedError @abc.abstractmethod def mkdirs(self, path): """Recursively create directories for the provided path. Args: path: string path of the directory structure that should be created Raises: IOError: if leaf directory already exists. """ raise NotImplementedError @abc.abstractmethod def has_dirs(self): """Whether this FileSystem supports directories.""" raise NotImplementedError @abc.abstractmethod def _list(self, dir_or_prefix): """List files in a location. Listing is non-recursive (for filesystems that support directories). Args: dir_or_prefix: (string) A directory or location prefix (for filesystems that don't have directories). Returns: Generator of ``FileMetadata`` objects. Raises: ``BeamIOError``: if listing fails, but not if no files were found. """ raise NotImplementedError @staticmethod def _split_scheme(url_or_path): match = re.match(r'(^[a-z]+)://(.*)', url_or_path) if match is not None: return match.groups() return None, url_or_path @staticmethod def _combine_scheme(scheme, path): if scheme is None: return path return '{}://{}'.format(scheme, path) def _url_dirname(self, url_or_path): """Like posixpath.dirname, but preserves scheme:// prefix. Args: url_or_path: A string in the form of scheme://some/path OR /some/path. """ scheme, path = self._split_scheme(url_or_path) return self._combine_scheme(scheme, posixpath.dirname(path)) def match_files(self, file_metas: list[FileMetadata], pattern: str) -> Iterator[FileMetadata]: """Filter :class:`FileMetadata` objects by *pattern* Args: file_metas (list of :class:`FileMetadata`): Files to consider when matching pattern (str): File pattern See Also: :meth:`translate_pattern` Returns: Generator of matching :class:`FileMetadata` """ re_pattern = re.compile(self.translate_pattern(pattern)) match = re_pattern.match for file_metadata in file_metas: if match(file_metadata.path): yield file_metadata @staticmethod def translate_pattern(pattern: str) -> str: """ Translate a *pattern* to a regular expression. There is no way to quote meta-characters. Pattern syntax: The pattern syntax is based on the fnmatch_ syntax, with the following differences: - ``*`` Is equivalent to ``[^/\\]*`` rather than ``.*``. - ``**`` Is equivalent to ``.*``. See also: :meth:`match` uses this method This method is based on `Python 2.7's fnmatch.translate`_. The code in this method is licensed under PYTHON SOFTWARE FOUNDATION LICENSE VERSION 2. .. _`fnmatch`: https://docs.python.org/2/library/fnmatch.html .. _`Python 2.7's fnmatch.translate`: https://github.com/python/cpython\ /blob/170ea8ccd4235d28538ab713041502d07ad1cacd/Lib/fnmatch.py#L85-L120 """ i, n = 0, len(pattern) res = '' while i < n: c = pattern[i] i = i + 1 if c == '*': # One char lookahead for "**" if i < n and pattern[i] == "*": res = res + '.*' i = i + 1 else: res = res + r'[^/\\]*' elif c == '?': res = res + '.' elif c == '[': j = i if j < n and pattern[j] == '!': j = j + 1 if j < n and pattern[j] == ']': j = j + 1 while j < n and pattern[j] != ']': j = j + 1 if j >= n: res = res + r'\[' else: stuff = pattern[i:j].replace('\\', '\\\\') i = j + 1 if stuff[0] == '!': stuff = '^' + stuff[1:] elif stuff[0] == '^': stuff = '\\' + stuff res = '%s[%s]' % (res, stuff) else: res = res + re.escape(c) logger.debug('translate_pattern: %r -> %r', pattern, res) return r'(?ms)' + res + r'\Z' def match(self, patterns, limits=None): """Find all matching paths to the patterns provided. See Also: :meth:`translate_pattern` Patterns ending with '/' or '\\' will be appended with '*'. Args: patterns: list of string for the file path pattern to match against limits: list of maximum number of responses that need to be fetched Returns: list of ``MatchResult`` objects. Raises: ``BeamIOError``: if any of the pattern match operations fail """ if limits is None: limits = [None] * len(patterns) else: err_msg = "Patterns and limits should be equal in length" assert len(patterns) == len(limits), err_msg def _match(pattern, limit): """Find all matching paths to the pattern provided.""" if pattern.endswith('/') or pattern.endswith('\\'): pattern += '*' # Get the part of the pattern before the first globbing character. # For example scheme://path/foo* will become scheme://path/foo for # filesystems like GCS, or converted to scheme://path for filesystems with # directories. prefix_or_dir = re.match('^[^[*?]*', pattern).group(0) file_metadatas = [] if prefix_or_dir == pattern: # Short-circuit calling self.list() if there's no glob pattern to match. if self.exists(pattern): file_metadatas = [self.metadata(pattern)] else: if self.has_dirs(): prefix_dirname = self._url_dirname(prefix_or_dir) if not prefix_dirname == prefix_or_dir: logger.debug( "Changed prefix_or_dir %r -> %r", prefix_or_dir, prefix_dirname) prefix_or_dir = prefix_dirname logger.debug("Listing files in %r", prefix_or_dir) file_metadatas = self._list(prefix_or_dir) metadata_list = [] for file_metadata in self.match_files(file_metadatas, pattern): if limit is not None and len(metadata_list) >= limit: break metadata_list.append(file_metadata) return MatchResult(pattern, metadata_list) exceptions = {} result = [] for pattern, limit in zip(patterns, limits): try: result.append(_match(pattern, limit)) except Exception as e: # pylint: disable=broad-except exceptions[pattern] = e if exceptions: raise BeamIOError("Match operation failed", exceptions) return result @abc.abstractmethod def create( self, path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO) -> BinaryIO: """Returns a write channel for the given file path. Args: path: string path of the file object to be written to the system mime_type: MIME type to specify the type of content in the file object compression_type: Type of compression to be used for this object Returns: file handle with a close function for the user to use """ raise NotImplementedError @abc.abstractmethod def open( self, path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO) -> BinaryIO: """Returns a read channel for the given file path. Args: path: string path of the file object to be read mime_type: MIME type to specify the type of content in the file object compression_type: Type of compression to be used for this object Returns: file handle with a close function for the user to use """ raise NotImplementedError @abc.abstractmethod def copy(self, source_file_names, destination_file_names): """Recursively copy the file tree from the source to the destination Args: source_file_names: list of source file objects that needs to be copied destination_file_names: list of destination of the new object Raises: ``BeamIOError``: if any of the copy operations fail """ raise NotImplementedError @abc.abstractmethod def rename(self, source_file_names, destination_file_names): """Rename the files at the source list to the destination list. Source and destination lists should be of the same size. Args: source_file_names: List of file paths that need to be moved destination_file_names: List of destination_file_names for the files Raises: ``BeamIOError``: if any of the rename operations fail """ raise NotImplementedError @abc.abstractmethod def exists(self, path: str) -> bool: """Check if the provided path exists on the FileSystem. Args: path: string path that needs to be checked. Returns: boolean flag indicating if path exists """ raise NotImplementedError @abc.abstractmethod def size(self, path: str) -> int: """Get size in bytes of a file on the FileSystem. Args: path: string filepath of file. Returns: int size of file according to the FileSystem. Raises: ``BeamIOError``: if path doesn't exist. """ raise NotImplementedError @abc.abstractmethod def last_updated(self, path): """Get UNIX Epoch time in seconds on the FileSystem. Args: path: string path of file. Returns: float UNIX Epoch time Raises: ``BeamIOError``: if path doesn't exist. """ raise NotImplementedError def checksum(self, path): """Fetch checksum metadata of a file on the :class:`~apache_beam.io.filesystem.FileSystem`. This operation returns checksum metadata as stored in the underlying FileSystem. It should not need to read file data to obtain this value. Checksum type and format are FileSystem dependent and are not compatible between FileSystems. FileSystem implementations may return file size if a checksum isn't available. Args: path: string path of a file. Returns: string containing checksum Raises: ``BeamIOError``: if path isn't a file or doesn't exist. """ raise NotImplementedError @abc.abstractmethod def metadata(self, path): """Fetch metadata of a file on the :class:`~apache_beam.io.filesystem.FileSystem`. This operation returns metadata as stored in the underlying FileSystem. It should not need to read file data to obtain this value. For web based file systems, this method should also incur as few as possible requests. Args: path: string path of a file. Returns: :class:`~apache_beam.io.filesystem.FileMetadata`. Raises: ``BeamIOError``: if path isn't a file or doesn't exist. """ raise NotImplementedError @abc.abstractmethod def delete(self, paths): """Deletes files or directories at the provided paths. Directories will be deleted recursively. Args: paths: list of paths that give the file objects to be deleted Raises: ``BeamIOError``: if any of the delete operations fail """ raise NotImplementedError def report_lineage(self, path, unused_lineage): """ Report Lineage metrics for path. Unless override by FileSystem implementations, default to no-op. """ pass ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/filesystem_test.py0000644000175100017510000005153015206113032022065 0ustar00runnerrunner# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for filesystem module.""" # pytype: skip-file import bz2 import gzip import logging import lzma import ntpath import os import posixpath import tempfile import unittest import zlib from io import BytesIO import zstandard from parameterized import param from parameterized import parameterized from apache_beam.io.filesystem import CompressedFile from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystem import FileMetadata from apache_beam.io.filesystem import FileSystem class TestingFileSystem(FileSystem): def __init__(self, pipeline_options, has_dirs=False): super().__init__(pipeline_options) self._has_dirs = has_dirs self._files = {} @classmethod def scheme(cls): # Required for FileSystems.get_filesystem(). return 'test' def join(self, basepath, *paths): raise NotImplementedError def split(self, path): raise NotImplementedError def mkdirs(self, path): raise NotImplementedError def has_dirs(self): return self._has_dirs def _insert_random_file(self, path, size): self._files[path] = size def _list(self, dir_or_prefix): for path, size in self._files.items(): if path.startswith(dir_or_prefix): yield FileMetadata(path, size) def create( self, path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO): raise NotImplementedError def open( self, path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO): raise NotImplementedError def copy(self, source_file_names, destination_file_names): raise NotImplementedError def rename(self, source_file_names, destination_file_names): raise NotImplementedError def exists(self, path): raise NotImplementedError def size(self, path): raise NotImplementedError def last_updated(self, path): raise NotImplementedError def checksum(self, path): raise NotImplementedError def metadata(self, path): raise NotImplementedError def delete(self, paths): raise NotImplementedError class TestFileSystem(unittest.TestCase): def setUp(self): self.fs = TestingFileSystem(pipeline_options=None) def _flatten_match(self, match_results): return [ file_metadata for match_result in match_results for file_metadata in match_result.metadata_list ] @parameterized.expand([ ('gs://gcsio-test/**', all), # Does not match root-level files ('gs://gcsio-test/**/*', lambda n, i: n not in ['cat.png']), # Only matches root-level files ('gs://gcsio-test/*', [('cat.png', 19)]), ( 'gs://gcsio-test/cow/**', [ ('cow/cat/fish', 2), ('cow/cat/blubber', 3), ('cow/dog/blubber', 4), ]), ( 'gs://gcsio-test/cow/ca**', [ ('cow/cat/fish', 2), ('cow/cat/blubber', 3), ]), ( 'gs://gcsio-test/apple/[df]ish/ca*', [ ('apple/fish/cat', 10), ('apple/fish/cart', 11), ('apple/fish/carl', 12), ('apple/dish/cat', 14), ('apple/dish/carl', 15), ]), ( 'gs://gcsio-test/apple/?ish/?a?', [ ('apple/fish/cat', 10), ('apple/dish/bat', 13), ('apple/dish/cat', 14), ]), ( 'gs://gcsio-test/apple/fish/car?', [ ('apple/fish/cart', 11), ('apple/fish/carl', 12), ]), ( 'gs://gcsio-test/apple/fish/b*', [ ('apple/fish/blubber', 6), ('apple/fish/blowfish', 7), ('apple/fish/bambi', 8), ('apple/fish/balloon', 9), ]), ( 'gs://gcsio-test/apple/f*/b*', [ ('apple/fish/blubber', 6), ('apple/fish/blowfish', 7), ('apple/fish/bambi', 8), ('apple/fish/balloon', 9), ]), ( 'gs://gcsio-test/apple/dish/[cb]at', [ ('apple/dish/bat', 13), ('apple/dish/cat', 14), ]), ( 'gs://gcsio-test/banana/cyrano.m?', [ ('banana/cyrano.md', 17), ('banana/cyrano.mb', 18), ]), ]) def test_match_glob(self, file_pattern, expected_object_names): objects = [ ('cow/cat/fish', 2), ('cow/cat/blubber', 3), ('cow/dog/blubber', 4), ('apple/dog/blubber', 5), ('apple/fish/blubber', 6), ('apple/fish/blowfish', 7), ('apple/fish/bambi', 8), ('apple/fish/balloon', 9), ('apple/fish/cat', 10), ('apple/fish/cart', 11), ('apple/fish/carl', 12), ('apple/dish/bat', 13), ('apple/dish/cat', 14), ('apple/dish/carl', 15), ('banana/cat', 16), ('banana/cyrano.md', 17), ('banana/cyrano.mb', 18), ('cat.png', 19) ] bucket_name = 'gcsio-test' if callable(expected_object_names): # A hack around the fact that the parameters do not have access to # the "objects" list. if expected_object_names is all: # It's a placeholder for "all" objects expected_object_names = objects else: # It's a filter function of type (str, int) -> bool # that returns true for expected objects filter_func = expected_object_names expected_object_names = [(short_path, size) for short_path, size in objects if filter_func(short_path, size)] for object_name, size in objects: file_name = 'gs://%s/%s' % (bucket_name, object_name) self.fs._insert_random_file(file_name, size) expected_file_names = [('gs://%s/%s' % (bucket_name, object_name), size) for object_name, size in expected_object_names] actual_file_names = [ (file_metadata.path, file_metadata.size_in_bytes) for file_metadata in self._flatten_match(self.fs.match([file_pattern])) ] self.assertEqual(set(actual_file_names), set(expected_file_names)) # Check if limits are followed correctly limit = 3 expected_num_items = min(len(expected_object_names), limit) self.assertEqual( len(self._flatten_match(self.fs.match([file_pattern], [limit]))), expected_num_items) @parameterized.expand([ param(os_path=posixpath, sep_re='/'), param(os_path=ntpath, sep_re='\\\\'), ]) def test_translate_pattern(self, os_path, sep_re): star = r'[^/\\]*' double_star = r'.*' join = os_path.join sep = os_path.sep pattern__expected = [ (join('a', '*'), sep_re.join(['a', star])), (join('b', '*') + sep, sep_re.join(['b', star]) + sep_re), (r'*[abc\]', star + r'[abc\\]'), (join('d', '**', '*'), sep_re.join(['d', double_star, star])), ] for pattern, expected in pattern__expected: expected = r'(?ms)' + expected + r'\Z' result = self.fs.translate_pattern(pattern) self.assertEqual(expected, result) class TestFileSystemWithDirs(TestFileSystem): def setUp(self): self.fs = TestingFileSystem(pipeline_options=None, has_dirs=True) class TestCompressedFile(unittest.TestCase): """Base class for TestCases that deals with TempDir clean-up. Inherited test cases will call self._new_tempdir() to start a temporary dir which will be deleted at the end of the tests (when tearDown() is called). """ content = b"""- the BEAM - How things really are we would like to know. Does Time flow, is it elastic, or is it atomized in instants hammered around the clock's face? ... - May Swenson""" # Keep the read block size small so that we exercise the seek functionality # in compressed file and not just in the internal buffer read_block_size = 4 def setUp(self): self._tempfiles = [] def tearDown(self): for path in self._tempfiles: if os.path.exists(path): os.remove(path) def _create_temp_file(self): path = tempfile.NamedTemporaryFile(delete=False).name self._tempfiles.append(path) return path def _create_compressed_file(self, compression_type, content): file_name = self._create_temp_file() if compression_type == CompressionTypes.DEFLATE: with open(file_name, 'wb') as f: f.write(zlib.compress(content)) elif compression_type == CompressionTypes.BZIP2 or \ compression_type == CompressionTypes.GZIP: compress_open = bz2.BZ2File \ if compression_type == CompressionTypes.BZIP2 \ else gzip.open with compress_open(file_name, 'wb') as f: f.write(content) elif compression_type == CompressionTypes.ZSTD: compress_open = zstandard.open with compress_open(file_name, 'wb') as f: f.write(content) elif compression_type == CompressionTypes.LZMA: compress_open = lzma.open with compress_open(file_name, 'wb') as f: f.write(content) else: assert False, "Invalid compression type: %s" % compression_type return file_name def test_seekable_enabled_on_read(self): with open(self._create_temp_file(), 'rb') as f: readable = CompressedFile(f) self.assertTrue(readable.seekable) def test_seekable_disabled_on_write(self): with open(self._create_temp_file(), 'wb') as f: writeable = CompressedFile(f) self.assertFalse(writeable.seekable) def test_seekable_disabled_on_append(self): with open(self._create_temp_file(), 'ab') as f: writeable = CompressedFile(f) self.assertFalse(writeable.seekable) def test_seek_set(self): for compression_type in [CompressionTypes.BZIP2, CompressionTypes.DEFLATE, CompressionTypes.GZIP, CompressionTypes.ZSTD, CompressionTypes.LZMA]: file_name = self._create_compressed_file(compression_type, self.content) with open(file_name, 'rb') as f: compressed_fd = CompressedFile( f, compression_type, read_size=self.read_block_size) reference_fd = BytesIO(self.content) # Note: BytesIO's tell() reports out of bound positions (if we seek # beyond the file), therefore we need to cap it to max_position # _CompressedFile.tell() always stays within the bounds of the # uncompressed content. # Negative seek position argument is not supported for BytesIO with # whence set to SEEK_SET. for seek_position in (0, 1, len(self.content) - 1, len(self.content), len(self.content) + 1): compressed_fd.seek(seek_position, os.SEEK_SET) reference_fd.seek(seek_position, os.SEEK_SET) uncompressed_line = compressed_fd.readline() reference_line = reference_fd.readline() self.assertEqual(uncompressed_line, reference_line) uncompressed_position = compressed_fd.tell() reference_position = reference_fd.tell() max_position = len(self.content) reference_position = min(reference_position, max_position) self.assertEqual(uncompressed_position, reference_position) def test_seek_cur(self): for compression_type in [CompressionTypes.BZIP2, CompressionTypes.DEFLATE, CompressionTypes.GZIP, CompressionTypes.ZSTD, CompressionTypes.LZMA]: file_name = self._create_compressed_file(compression_type, self.content) with open(file_name, 'rb') as f: compressed_fd = CompressedFile( f, compression_type, read_size=self.read_block_size) reference_fd = BytesIO(self.content) # Test out of bound, inbound seeking in both directions # Note: BytesIO's seek() reports out of bound positions (if we seek # beyond the file), therefore we need to cap it to max_position (to # make it consistent with the old StringIO behavior for seek_position in (-1, 0, 1, len(self.content) // 2, len(self.content) // 2, -1 * len(self.content) // 2): compressed_fd.seek(seek_position, os.SEEK_CUR) reference_fd.seek(seek_position, os.SEEK_CUR) uncompressed_line = compressed_fd.readline() expected_line = reference_fd.readline() self.assertEqual(uncompressed_line, expected_line) reference_position = reference_fd.tell() uncompressed_position = compressed_fd.tell() max_position = len(self.content) reference_position = min(reference_position, max_position) reference_fd.seek(reference_position, os.SEEK_SET) self.assertEqual(uncompressed_position, reference_position) def test_read_from_end_returns_no_data(self): for compression_type in [CompressionTypes.BZIP2, CompressionTypes.DEFLATE, CompressionTypes.GZIP, CompressionTypes.ZSTD, CompressionTypes.LZMA]: file_name = self._create_compressed_file(compression_type, self.content) with open(file_name, 'rb') as f: compressed_fd = CompressedFile( f, compression_type, read_size=self.read_block_size) seek_position = 0 compressed_fd.seek(seek_position, os.SEEK_END) expected_data = b'' uncompressed_data = compressed_fd.read(10) self.assertEqual(uncompressed_data, expected_data) def test_seek_outside(self): for compression_type in [CompressionTypes.BZIP2, CompressionTypes.DEFLATE, CompressionTypes.GZIP, CompressionTypes.ZSTD, CompressionTypes.LZMA]: file_name = self._create_compressed_file(compression_type, self.content) with open(file_name, 'rb') as f: compressed_fd = CompressedFile( f, compression_type, read_size=self.read_block_size) for whence in (os.SEEK_CUR, os.SEEK_SET, os.SEEK_END): seek_position = -1 * len(self.content) - 10 compressed_fd.seek(seek_position, whence) expected_position = 0 uncompressed_position = compressed_fd.tell() self.assertEqual(uncompressed_position, expected_position) seek_position = len(self.content) + 20 compressed_fd.seek(seek_position, whence) expected_position = len(self.content) uncompressed_position = compressed_fd.tell() self.assertEqual(uncompressed_position, expected_position) def test_read_and_seek_back_to_beginning(self): for compression_type in [CompressionTypes.BZIP2, CompressionTypes.DEFLATE, CompressionTypes.GZIP, CompressionTypes.ZSTD, CompressionTypes.LZMA]: file_name = self._create_compressed_file(compression_type, self.content) with open(file_name, 'rb') as f: compressed_fd = CompressedFile( f, compression_type, read_size=self.read_block_size) first_pass = compressed_fd.readline() compressed_fd.seek(0, os.SEEK_SET) second_pass = compressed_fd.readline() self.assertEqual(first_pass, second_pass) def test_read(self): for compression_type in [CompressionTypes.BZIP2, CompressionTypes.DEFLATE, CompressionTypes.GZIP, CompressionTypes.ZSTD, CompressionTypes.LZMA]: file_name = self._create_compressed_file(compression_type, self.content) with open(file_name, 'rb') as f: compressed_fd = CompressedFile( f, compression_type, read_size=self.read_block_size) data = compressed_fd.read() self.assertEqual(data, self.content) def test_tell(self): lines = [b'line%d\n' % i for i in range(10)] tmpfile = self._create_temp_file() with open(tmpfile, 'wb') as f: writeable = CompressedFile(f) current_offset = 0 for line in lines: writeable.write(line) current_offset += len(line) self.assertEqual(current_offset, writeable.tell()) with open(tmpfile, 'rb') as f: readable = CompressedFile(f) current_offset = 0 while True: line = readable.readline() current_offset += len(line) self.assertEqual(current_offset, readable.tell()) if not line: break def test_concatenated_compressed_file(self): # The test apache_beam.io.textio_test.test_read_gzip_concat # does not encounter the problem in the Beam 2.13 and earlier # code base because the test data is too small: the data is # smaller than read_size, so it goes through logic in the code # that avoids the problem in the code. So, this test sets # read_size smaller and test data bigger, in order to # encounter the problem. It would be difficult to test in the # textio_test module, because you'd need very large test data # because default read_size is 16MiB, and the ReadFromText # interface does not allow you to modify the read_size. import random import threading from six import int2byte num_test_lines = 10 timeout = 30 read_size = (64 << 10) # set much smaller than the line size byte_table = tuple(int2byte(i) for i in range(32, 96)) def generate_random_line(): byte_list = list( b for i in range(4096) for b in random.sample(byte_table, 64)) byte_list.append(b'\n') return b''.join(byte_list) def create_test_file(compression_type, lines): filenames = [] file_name = self._create_temp_file() if compression_type == CompressionTypes.BZIP2: compress_factory = bz2.BZ2File elif compression_type == CompressionTypes.GZIP: compress_factory = gzip.open elif compression_type == CompressionTypes.ZSTD: compress_factory = zstandard.open elif compression_type == CompressionTypes.LZMA: compress_factory = lzma.open else: assert False, "Invalid compression type: %s" % compression_type for line in lines: filenames.append(self._create_temp_file()) with compress_factory(filenames[-1], 'wb') as f: f.write(line) with open(file_name, 'wb') as o: for name in filenames: with open(name, 'rb') as i: o.write(i.read()) return file_name # I remember some time ago when a job ran with a real concatenated # gzip file, I got into an endless loop in the beam filesystem module. # That's why I put this handler in to trap an endless loop. However, # this unit test doesn't encounter an endless loop, it encounters a # different error, in the Beam 2.13 and earlier implementation. # So it's not strictly necessary to have this handler in this unit test. def timeout_handler(): raise IOError('Exiting due to likley infinite loop logic in code.') timer = threading.Timer(timeout, timeout_handler) try: test_lines = tuple(generate_random_line() for i in range(num_test_lines)) for compression_type in [CompressionTypes.BZIP2, CompressionTypes.GZIP, CompressionTypes.ZSTD, CompressionTypes.LZMA]: file_name = create_test_file(compression_type, test_lines) timer.start() with open(file_name, 'rb') as f: data = CompressedFile(f, compression_type, read_size=read_size) for written_line in test_lines: read_line = data.readline() self.assertEqual(written_line, read_line) timer.cancel() # Starting a new timer for the next iteration/test. timer = threading.Timer(timeout, timeout_handler) finally: timer.cancel() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/filesystemio.py0000644000175100017510000002015115206113032021351 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utilities for ``FileSystem`` implementations.""" # pytype: skip-file import abc import io import os __all__ = [ 'Downloader', 'Uploader', 'DownloaderStream', 'UploaderStream', 'PipeStream' ] class Downloader(metaclass=abc.ABCMeta): """Download interface for a single file. Implementations should support random access reads. """ @property @abc.abstractmethod def size(self): """Size of file to download.""" @abc.abstractmethod def get_range(self, start, end): """Retrieve a given byte range [start, end) from this download. Range must be in this form: 0 <= start < end: Fetch the bytes from start to end. Args: start: (int) Initial byte offset. end: (int) Final byte offset, exclusive. Returns: (string) A buffer containing the requested data. """ class Uploader(metaclass=abc.ABCMeta): """Upload interface for a single file.""" @abc.abstractmethod def put(self, data): """Write data to file sequentially. Args: data: (memoryview) Data to write. """ @abc.abstractmethod def finish(self): """Signal to upload any remaining data and close the file. File should be fully written upon return from this method. Raises: Any error encountered during the upload. """ class DownloaderStream(io.RawIOBase): """Provides a stream interface for Downloader objects.""" def __init__( self, downloader, read_buffer_size=io.DEFAULT_BUFFER_SIZE, mode='rb'): """Initializes the stream. Args: downloader: (Downloader) Filesystem dependent implementation. read_buffer_size: (int) Buffer size to use during read operations. mode: (string) Python mode attribute for this stream. """ self._downloader = downloader self.mode = mode self._position = 0 self._reader_buffer_size = read_buffer_size def readinto(self, b): """Read up to len(b) bytes into b. Returns number of bytes read (0 for EOF). Args: b: (bytearray/memoryview) Buffer to read into. """ self._checkClosed() if self._position >= self._downloader.size: return 0 start = self._position end = min(self._position + len(b), self._downloader.size) data = self._downloader.get_range(start, end) self._position += len(data) b[:len(data)] = data return len(data) def seek(self, offset, whence=os.SEEK_SET): """Set the stream's current offset. Note if the new offset is out of bound, it is adjusted to either 0 or EOF. Args: offset: seek offset as number. whence: seek mode. Supported modes are os.SEEK_SET (absolute seek), os.SEEK_CUR (seek relative to the current position), and os.SEEK_END (seek relative to the end, offset should be negative). Raises: ``ValueError``: When this stream is closed or if whence is invalid. """ self._checkClosed() if whence == os.SEEK_SET: self._position = offset elif whence == os.SEEK_CUR: self._position += offset elif whence == os.SEEK_END: self._position = self._downloader.size + offset else: raise ValueError('Whence mode %r is invalid.' % whence) self._position = min(self._position, self._downloader.size) self._position = max(self._position, 0) return self._position def tell(self): """Tell the stream's current offset. Returns: current offset in reading this stream. Raises: ``ValueError``: When this stream is closed. """ self._checkClosed() return self._position def seekable(self): return True def readable(self): return True def readall(self): """Read until EOF, using multiple read() call.""" res = [] while True: data = self.read(self._reader_buffer_size) if not data: break res.append(data) return b''.join(res) class UploaderStream(io.RawIOBase): """Provides a stream interface for Uploader objects.""" def __init__(self, uploader, mode='wb'): """Initializes the stream. Args: uploader: (Uploader) Filesystem dependent implementation. mode: (string) Python mode attribute for this stream. """ self._uploader = uploader self.mode = mode self._position = 0 def tell(self): return self._position def write(self, b): """Write bytes from b. Returns number of bytes written (<= len(b)). Args: b: (memoryview) Buffer with data to write. """ self._checkClosed() self._uploader.put(b) bytes_written = len(b) self._position += bytes_written return bytes_written def close(self): """Complete the upload and close this stream. This method has no effect if the stream is already closed. Raises: Any error encountered by the uploader. """ if not self.closed: self._uploader.finish() super().close() def writable(self): return True class PipeStream(object): """A class that presents a pipe connection as a readable stream. Not thread-safe. Remembers the last ``size`` bytes read and allows rewinding the stream by that amount exactly. See BEAM-6380 for more. """ def __init__(self, recv_pipe): self.conn = recv_pipe self.closed = False self.position = 0 self.remaining = b'' # Data and position of last block streamed. Allows limited seeking backwards # of stream. self.last_block_position = None self.last_block = b'' def read(self, size): """Read data from the wrapped pipe connection. Args: size: Number of bytes to read. Actual number of bytes read is always equal to size unless EOF is reached. Returns: data read as str. """ data_list = [] bytes_read = 0 last_block_position = self.position while bytes_read < size: bytes_from_remaining = min(size - bytes_read, len(self.remaining)) data_list.append(self.remaining[0:bytes_from_remaining]) self.remaining = self.remaining[bytes_from_remaining:] self.position += bytes_from_remaining bytes_read += bytes_from_remaining if not self.remaining: try: self.remaining = self.conn.recv_bytes() except EOFError: break last_block = b''.join(data_list) if last_block: self.last_block_position = last_block_position self.last_block = last_block return last_block def tell(self): """Tell the file's current offset. Returns: current offset in reading this file. Raises: ``ValueError``: When this stream is closed. """ self._check_open() return self.position def seek(self, offset, whence=os.SEEK_SET): # The apitools library used by the gcsio.Uploader class insists on seeking # to the end of a stream to do a check before completing an upload, so we # must have this no-op method here in that case. if whence == os.SEEK_END and offset == 0: return elif whence == os.SEEK_SET: if offset == self.position: return elif offset == self.last_block_position and self.last_block: self.position = offset self.remaining = b''.join([self.last_block, self.remaining]) self.last_block = b'' return raise NotImplementedError( 'offset: %s, whence: %s, position: %s, last: %s' % (offset, whence, self.position, self.last_block_position)) def _check_open(self): if self.closed: raise IOError('Stream is closed.') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/filesystemio_test.py0000644000175100017510000001701015206113032022410 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for filesystemio.""" # pytype: skip-file import io import logging import multiprocessing import os import threading import unittest from apache_beam.io import filesystemio _LOGGER = logging.getLogger(__name__) class FakeDownloader(filesystemio.Downloader): def __init__(self, data): self._data = data self.last_read_size = -1 @property def size(self): return len(self._data) def get_range(self, start, end): self.last_read_size = end - start return self._data[start:end] class FakeUploader(filesystemio.Uploader): def __init__(self): self.data = b'' self.last_write_size = -1 self.finished = False def last_error(self): return None def put(self, data): assert not self.finished self.data += data.tobytes() self.last_write_size = len(data) def finish(self): self.finished = True class TestDownloaderStream(unittest.TestCase): def test_file_attributes(self): downloader = FakeDownloader(data=None) stream = filesystemio.DownloaderStream(downloader) self.assertEqual(stream.mode, 'rb') self.assertTrue(stream.readable()) self.assertFalse(stream.writable()) self.assertTrue(stream.seekable()) def test_read_empty(self): downloader = FakeDownloader(data=b'') stream = filesystemio.DownloaderStream(downloader) self.assertEqual(stream.read(), b'') def test_read(self): data = b'abcde' downloader = FakeDownloader(data) stream = filesystemio.DownloaderStream(downloader) # Read size is exactly what was passed to read() (unbuffered). self.assertEqual(stream.read(1), data[0:1]) self.assertEqual(downloader.last_read_size, 1) self.assertEqual(stream.read(), data[1:]) self.assertEqual(downloader.last_read_size, len(data) - 1) def test_read_buffered(self): data = b'abcde' downloader = FakeDownloader(data) buffer_size = 2 stream = io.BufferedReader( filesystemio.DownloaderStream(downloader), buffer_size) # Verify that buffering works and is reading ahead. self.assertEqual(stream.read(1), data[0:1]) self.assertEqual(downloader.last_read_size, buffer_size) self.assertEqual(stream.read(), data[1:]) class TestUploaderStream(unittest.TestCase): def test_file_attributes(self): uploader = FakeUploader() stream = filesystemio.UploaderStream(uploader) self.assertEqual(stream.mode, 'wb') self.assertFalse(stream.readable()) self.assertTrue(stream.writable()) self.assertFalse(stream.seekable()) def test_write_empty(self): uploader = FakeUploader() stream = filesystemio.UploaderStream(uploader) data = b'' stream.write(memoryview(data)) self.assertEqual(uploader.data, data) def test_write(self): data = b'abcde' uploader = FakeUploader() stream = filesystemio.UploaderStream(uploader) # Unbuffered writes. stream.write(memoryview(data[0:1])) self.assertEqual(uploader.data[0], data[0]) self.assertEqual(uploader.last_write_size, 1) stream.write(memoryview(data[1:])) self.assertEqual(uploader.data, data) self.assertEqual(uploader.last_write_size, len(data) - 1) def test_write_buffered(self): data = b'abcde' uploader = FakeUploader() buffer_size = 2 stream = io.BufferedWriter( filesystemio.UploaderStream(uploader), buffer_size) # Verify that buffering works: doesn't write to uploader until buffer is # filled. stream.write(data[0:1]) self.assertEqual(-1, uploader.last_write_size) stream.write(data[1:]) stream.close() self.assertEqual(data, uploader.data) class TestPipeStream(unittest.TestCase): def _read_and_verify(self, stream, expected, buffer_size, success): data_list = [] bytes_read = 0 seen_last_block = False while True: data = stream.read(buffer_size) self.assertLessEqual(len(data), buffer_size) if len(data) < buffer_size: # Test the constraint that the pipe stream returns less than the buffer # size only when at the end of the stream. if data: self.assertFalse(seen_last_block) seen_last_block = True if not data: break data_list.append(data) bytes_read += len(data) self.assertEqual(stream.tell(), bytes_read) self.assertEqual(b''.join(data_list), expected) success[0] = True def _read_and_seek(self, stream, expected, buffer_size, success): data_list = [] bytes_read = 0 while True: data = stream.read(buffer_size) # Test bad seek positions. with self.assertRaises(NotImplementedError): stream.seek(bytes_read + 1) with self.assertRaises(NotImplementedError): stream.seek(bytes_read - 1) # Rewind stream and test that it reads back the same data again. stream.seek(bytes_read) data2 = stream.read(buffer_size) self.assertEqual(data, data2) if not data: break data_list.append(data) bytes_read += len(data) self.assertEqual(stream.tell(), bytes_read) self.assertEqual(len(b''.join(data_list)), len(expected)) self.assertEqual(b''.join(data_list), expected) success[0] = True def test_pipe_stream(self): block_sizes = list(4**i for i in range(0, 12)) data_blocks = list(os.urandom(size) for size in block_sizes) expected = b''.join(data_blocks) buffer_sizes = [100001, 512 * 1024, 1024 * 1024] for buffer_size in buffer_sizes: for target in [self._read_and_verify, self._read_and_seek]: _LOGGER.info('buffer_size=%s, target=%s' % (buffer_size, target)) parent_conn, child_conn = multiprocessing.Pipe() stream = filesystemio.PipeStream(child_conn) success = [False] child_thread = threading.Thread( target=target, args=(stream, expected, buffer_size, success)) child_thread.start() for data in data_blocks: parent_conn.send_bytes(data) parent_conn.close() child_thread.join() self.assertTrue(success[0], 'error in test thread') def test_pipe_stream_rewind_buffer(self): buffer_size = 512 data = os.urandom(buffer_size) parent_conn, child_conn = multiprocessing.Pipe() parent_conn.send_bytes(data) parent_conn.close() stream = filesystemio.PipeStream(child_conn) # Regular read. read_data = stream.read(buffer_size) self.assertEqual(data, read_data) # Rewind buffer_size bytes. stream.seek(0) read_data = stream.read(buffer_size) self.assertEqual(data, read_data) # Read 0 bytes. Rewind buffer still points to offset 0. read_data = stream.read(buffer_size) self.assertFalse(read_data) stream.seek(0) read_data = stream.read(buffer_size) self.assertEqual(data, read_data) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/filesystems.py0000644000175100017510000003104415206113032021207 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """FileSystems interface class for accessing the correct filesystem""" # pytype: skip-file import logging import re from typing import BinaryIO # pylint: disable=unused-import from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystem import FileSystem from apache_beam.metrics.metric import Lineage from apache_beam.options.value_provider import RuntimeValueProvider _LOGGER = logging.getLogger(__name__) # All filesystem implements should be added here as # best effort imports. We don't want to force loading # a module if the user doesn't supply the correct # packages that these filesystems rely on. # # pylint: disable=wrong-import-position, unused-import try: from apache_beam.io.hadoopfilesystem import HadoopFileSystem except ModuleNotFoundError: # optional file system packages are not installed. pass except ImportError as e: _LOGGER.warning( 'Failed to import HadoopFileSystem; ' 'loading of this filesystem will be skipped. ' 'Error details: %s', e) try: from apache_beam.io.localfilesystem import LocalFileSystem except ModuleNotFoundError: pass except ImportError as e: # optional file system packages are installed but failed to load. _LOGGER.warning( 'Failed to import LocalFileSystem; ' 'loading of this filesystem will be skipped. ' 'Error details: %s', e) try: from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem except ModuleNotFoundError: pass except ImportError as e: _LOGGER.warning( 'Failed to import GCSFileSystem; ' 'loading of this filesystem will be skipped. ' 'Error details: %s', e) try: from apache_beam.io.aws.s3filesystem import S3FileSystem except ModuleNotFoundError: pass except ImportError as e: _LOGGER.warning( 'Failed to import S3FileSystem; ' 'loading of this filesystem will be skipped. ' 'Error details: %s', e) try: from apache_beam.io.azure.blobstoragefilesystem import BlobStorageFileSystem except ModuleNotFoundError: pass except ImportError as e: _LOGGER.warning( 'Failed to import BlobStorageFileSystem; ' 'loading of this filesystem will be skipped. ' 'Error details: %s', e) # pylint: enable=wrong-import-position, unused-import __all__ = ['FileSystems'] class FileSystems(object): """A class that defines the functions that can be performed on a filesystem. All methods are static and access the underlying registered filesystems. """ URI_SCHEMA_PATTERN = re.compile('(?P[a-zA-Z][-a-zA-Z0-9+.]*)://.*') _pipeline_options = None @classmethod def set_options(cls, pipeline_options): """Set filesystem options. Args: pipeline_options: Instance of ``PipelineOptions``. """ cls._pipeline_options = pipeline_options @staticmethod def get_scheme(path): match_result = FileSystems.URI_SCHEMA_PATTERN.match(path.strip()) if match_result is None: return None return match_result.groupdict()['scheme'] @staticmethod def get_filesystem(path): # type: (str) -> FileSystem """Get the correct filesystem for the specified path """ try: path_scheme = FileSystems.get_scheme(path) systems = [ fs for fs in FileSystem.get_all_subclasses() if fs.scheme() == path_scheme ] if len(systems) == 0: raise ValueError( 'Unable to get filesystem from specified path, please use the ' 'correct path or ensure the required dependency is installed, ' 'e.g., pip install apache-beam[gcp]. Path specified: %s' % path) elif len(systems) == 1: # Pipeline options could come either from the Pipeline itself (using # direct runner), or via RuntimeValueProvider (other runners). options = ( FileSystems._pipeline_options or RuntimeValueProvider.runtime_options) return systems[0](pipeline_options=options) else: raise ValueError('Found more than one filesystem for path %s' % path) except ValueError: raise except Exception as e: raise BeamIOError('Unable to get the Filesystem', {path: e}) @staticmethod def join(basepath, *paths): # type: (str, *str) -> str """Join two or more pathname components for the filesystem Args: basepath: string path of the first component of the path paths: path components to be added Returns: full path after combining all the passed components """ filesystem = FileSystems.get_filesystem(basepath) return filesystem.join(basepath, *paths) @staticmethod def split(path): """Splits the given path into two parts. Splits the path into a pair (head, tail) such that tail contains the last component of the path and head contains everything up to that. For file-systems other than the local file-system, head should include the prefix. Args: path: path as a string Returns: a pair of path components as strings. """ filesystem = FileSystems.get_filesystem(path) return filesystem.split(path) @staticmethod def mkdirs(path): """Recursively create directories for the provided path. Args: path: string path of the directory structure that should be created Raises: IOError: if leaf directory already exists. """ filesystem = FileSystems.get_filesystem(path) return filesystem.mkdirs(path) @staticmethod def match(patterns, limits=None): """Find all matching paths to the patterns provided. Pattern matching is done using each filesystem's ``match`` method (e.g. :meth:`.filesystem.FileSystem.match`). .. note:: - Depending on the :class:`.FileSystem` implementation, file listings (the ``.FileSystem._list`` method) may not be recursive. - If the file listing is not recursive, a pattern like ``scheme://path/*/foo`` will not be able to mach any files. See Also: :meth:`.filesystem.FileSystem.match` Pattern syntax: The pattern syntax is based on the fnmatch_ syntax, with the following differences: - ``*`` Is equivalent to ``[^/\\]*`` rather than ``.*``. - ``**`` Is equivalent to ``.*``. .. _`fnmatch`: https://docs.python.org/2/library/fnmatch.html Args: patterns: list of string for the file path pattern to match against limits: list of maximum number of responses that need to be fetched Returns: list of ``MatchResult`` objects. Raises: ``BeamIOError``: if any of the pattern match operations fail """ if len(patterns) == 0: return [] filesystem = FileSystems.get_filesystem(patterns[0]) return filesystem.match(patterns, limits) @staticmethod def create( path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO): # type: (...) -> BinaryIO """Returns a write channel for the given file path. Args: path: string path of the file object to be written to the system mime_type: MIME type to specify the type of content in the file object compression_type: Type of compression to be used for this object. See ``CompressionTypes`` for possible values. Returns: file handle with a ``close`` function for the user to use. """ filesystem = FileSystems.get_filesystem(path) return filesystem.create(path, mime_type, compression_type) @staticmethod def open( path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO): # type: (...) -> BinaryIO """Returns a read channel for the given file path. Args: path: string path of the file object to be written to the system mime_type: MIME type to specify the type of content in the file object compression_type: Type of compression to be used for this object. See ``CompressionTypes`` for possible values. Returns: file handle with a ``close`` function for the user to use. """ filesystem = FileSystems.get_filesystem(path) return filesystem.open(path, mime_type, compression_type) @staticmethod def copy(source_file_names, destination_file_names): """Recursively copy the file list from the source to the destination Args: source_file_names: list of source file objects that needs to be copied destination_file_names: list of destination of the new object Raises: ``BeamIOError``: if any of the copy operations fail """ if len(source_file_names) == 0: return filesystem = FileSystems.get_filesystem(source_file_names[0]) return filesystem.copy(source_file_names, destination_file_names) @staticmethod def rename(source_file_names, destination_file_names): """Rename the files at the source list to the destination list. Source and destination lists should be of the same size. Args: source_file_names: List of file paths that need to be moved destination_file_names: List of destination_file_names for the files Raises: ``BeamIOError``: if any of the rename operations fail """ if len(source_file_names) == 0: return filesystem = FileSystems.get_filesystem(source_file_names[0]) return filesystem.rename(source_file_names, destination_file_names) @staticmethod def exists(path): """Check if the provided path exists on the FileSystem. Args: path: string path that needs to be checked. Returns: boolean flag indicating if path exists """ filesystem = FileSystems.get_filesystem(path) return filesystem.exists(path) @staticmethod def last_updated(path): """Get UNIX Epoch time in seconds on the FileSystem. Args: path: string path of file. Returns: float UNIX Epoch time Raises: ``BeamIOError``: if path doesn't exist. """ filesystem = FileSystems.get_filesystem(path) return filesystem.last_updated(path) @staticmethod def checksum(path): """Fetch checksum metadata of a file on the :class:`~apache_beam.io.filesystem.FileSystem`. This operation returns checksum metadata as stored in the underlying FileSystem. It should not read any file data. Checksum type and format are FileSystem dependent and are not compatible between FileSystems. Args: path: string path of a file. Returns: string containing checksum Raises: ``BeamIOError``: if path isn't a file or doesn't exist. """ filesystem = FileSystems.get_filesystem(path) return filesystem.checksum(path) @staticmethod def delete(paths): """Deletes files or directories at the provided paths. Directories will be deleted recursively. Args: paths: list of paths that give the file objects to be deleted Raises: ``BeamIOError``: if any of the delete operations fail """ if isinstance(paths, str): raise BeamIOError( 'Delete passed string argument instead of list: %s' % paths) if len(paths) == 0: return filesystem = FileSystems.get_filesystem(paths[0]) return filesystem.delete(paths) @staticmethod def get_chunk_size(path): """Get the correct chunk size for the FileSystem. Args: path: string path that needs to be checked. Returns: integer size for parallelization in the FS operations. """ filesystem = FileSystems.get_filesystem(path) return filesystem.CHUNK_SIZE @staticmethod def report_source_lineage(path): """ Report source :class:`~apache_beam.metrics.metric.Lineage`. Args: path: string path to be reported. """ FileSystems.get_filesystem(path).report_lineage(path, Lineage.sources()) @staticmethod def report_sink_lineage(path): """ Report sink :class:`~apache_beam.metrics.metric.Lineage`. Args: path: string path to be reported. """ FileSystems.get_filesystem(path).report_lineage(path, Lineage.sinks()) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/filesystems_test.py0000644000175100017510000002016415206113032022247 0ustar00runnerrunner# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for LocalFileSystem.""" # pytype: skip-file import filecmp import logging import os import shutil import tempfile import unittest import mock from apache_beam.io import localfilesystem from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystems import FileSystems def _gen_fake_join(separator): """Returns a callable that joins paths with the given separator.""" def _join(first_path, *paths): return separator.join((first_path.rstrip(separator), ) + paths) return _join class FileSystemsTest(unittest.TestCase): def setUp(self): self.tmpdir = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.tmpdir) def test_get_scheme(self): self.assertIsNone(FileSystems.get_scheme('/abc/cdf')) self.assertIsNone(FileSystems.get_scheme('c:\\abc\\cdf')) self.assertEqual(FileSystems.get_scheme('gs://abc/cdf'), 'gs') def test_get_filesystem(self): self.assertTrue( isinstance( FileSystems.get_filesystem('/tmp'), localfilesystem.LocalFileSystem)) self.assertTrue( isinstance( FileSystems.get_filesystem('c:\\abc\\def'), localfilesystem.LocalFileSystem)) with self.assertRaises(ValueError): FileSystems.get_filesystem('error://abc/def') @mock.patch('apache_beam.io.localfilesystem.os') def test_unix_path_join(self, *unused_mocks): # Test joining of Unix paths. localfilesystem.os.path.join.side_effect = _gen_fake_join('/') self.assertEqual( '/tmp/path/to/file', FileSystems.join('/tmp/path', 'to', 'file')) self.assertEqual( '/tmp/path/to/file', FileSystems.join('/tmp/path', 'to/file')) self.assertEqual( '/tmp/path/to/file', FileSystems.join('/', 'tmp/path', 'to/file')) self.assertEqual( '/tmp/path/to/file', FileSystems.join('/tmp/', 'path', 'to/file')) @mock.patch('apache_beam.io.localfilesystem.os') def test_windows_path_join(self, *unused_mocks): # Test joining of Windows paths. localfilesystem.os.path.join.side_effect = _gen_fake_join('\\') self.assertEqual( r'C:\tmp\path\to\file', FileSystems.join(r'C:\tmp\path', 'to', 'file')) self.assertEqual( r'C:\tmp\path\to\file', FileSystems.join(r'C:\tmp\path', r'to\file')) self.assertEqual( r'C:\tmp\path\to\file', FileSystems.join(r'C:\tmp\path\\', 'to', 'file')) def test_mkdirs(self): path = os.path.join(self.tmpdir, 't1/t2') FileSystems.mkdirs(path) self.assertTrue(os.path.isdir(path)) def test_mkdirs_failed(self): path = os.path.join(self.tmpdir, 't1/t2') FileSystems.mkdirs(path) # Check IOError if existing directory is created with self.assertRaises(IOError): FileSystems.mkdirs(path) with self.assertRaises(IOError): FileSystems.mkdirs(os.path.join(self.tmpdir, 't1')) def test_match_file(self): path = os.path.join(self.tmpdir, 'f1') open(path, 'a').close() # Match files in the temp directory result = FileSystems.match([path])[0] files = [f.path for f in result.metadata_list] self.assertEqual(files, [path]) def test_match_file_empty(self): path = os.path.join(self.tmpdir, 'f2') # Does not exist # Match files in the temp directory result = FileSystems.match([path])[0] files = [f.path for f in result.metadata_list] self.assertEqual(files, []) def test_match_file_exception(self): # Match files with None so that it throws an exception with self.assertRaisesRegex(BeamIOError, r'^Unable to get the Filesystem') as error: FileSystems.match([None]) self.assertEqual(list(error.exception.exception_details), [None]) def test_match_directory_with_files(self): path1 = os.path.join(self.tmpdir, 'f1') path2 = os.path.join(self.tmpdir, 'f2') open(path1, 'a').close() open(path2, 'a').close() # Match both the files in the directory path = os.path.join(self.tmpdir, '*') result = FileSystems.match([path])[0] files = [f.path for f in result.metadata_list] self.assertCountEqual(files, [path1, path2]) def test_match_directory(self): result = FileSystems.match([self.tmpdir])[0] files = [f.path for f in result.metadata_list] self.assertEqual(files, [self.tmpdir]) def test_copy(self): path1 = os.path.join(self.tmpdir, 'f1') path2 = os.path.join(self.tmpdir, 'f2') with open(path1, 'a') as f: f.write('Hello') FileSystems.copy([path1], [path2]) self.assertTrue(filecmp.cmp(path1, path2)) def test_copy_error(self): path1 = os.path.join(self.tmpdir, 'f1') path2 = os.path.join(self.tmpdir, 'f2') with self.assertRaisesRegex(BeamIOError, r'^Copy operation failed') as error: FileSystems.copy([path1], [path2]) self.assertEqual( list(error.exception.exception_details.keys()), [(path1, path2)]) def test_copy_directory(self): path_t1 = os.path.join(self.tmpdir, 't1') path_t2 = os.path.join(self.tmpdir, 't2') FileSystems.mkdirs(path_t1) FileSystems.mkdirs(path_t2) path1 = os.path.join(path_t1, 'f1') path2 = os.path.join(path_t2, 'f1') with open(path1, 'a') as f: f.write('Hello') FileSystems.copy([path_t1], [path_t2]) self.assertTrue(filecmp.cmp(path1, path2)) def test_rename(self): path1 = os.path.join(self.tmpdir, 'f1') path2 = os.path.join(self.tmpdir, 'f2') with open(path1, 'a') as f: f.write('Hello') FileSystems.rename([path1], [path2]) self.assertTrue(FileSystems.exists(path2)) self.assertFalse(FileSystems.exists(path1)) def test_rename_error(self): path1 = os.path.join(self.tmpdir, 'f1') path2 = os.path.join(self.tmpdir, 'f2') with self.assertRaisesRegex(BeamIOError, r'^Rename operation failed') as error: FileSystems.rename([path1], [path2]) self.assertEqual( list(error.exception.exception_details.keys()), [(path1, path2)]) def test_rename_directory(self): path_t1 = os.path.join(self.tmpdir, 't1') path_t2 = os.path.join(self.tmpdir, 't2') FileSystems.mkdirs(path_t1) path1 = os.path.join(path_t1, 'f1') path2 = os.path.join(path_t2, 'f1') with open(path1, 'a') as f: f.write('Hello') FileSystems.rename([path_t1], [path_t2]) self.assertTrue(FileSystems.exists(path_t2)) self.assertFalse(FileSystems.exists(path_t1)) self.assertTrue(FileSystems.exists(path2)) self.assertFalse(FileSystems.exists(path1)) def test_exists(self): path1 = os.path.join(self.tmpdir, 'f1') path2 = os.path.join(self.tmpdir, 'f2') with open(path1, 'a') as f: f.write('Hello') self.assertTrue(FileSystems.exists(path1)) self.assertFalse(FileSystems.exists(path2)) def test_delete(self): path1 = os.path.join(self.tmpdir, 'f1') with open(path1, 'a') as f: f.write('Hello') self.assertTrue(FileSystems.exists(path1)) FileSystems.delete([path1]) self.assertFalse(FileSystems.exists(path1)) def test_delete_error(self): path1 = os.path.join(self.tmpdir, 'f1') with self.assertRaisesRegex(BeamIOError, r'^Delete operation failed') as error: FileSystems.delete([path1]) self.assertEqual(list(error.exception.exception_details.keys()), [path1]) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3214016 apache_beam-2.74.0/apache_beam/io/flink/0000755000175100017510000000000015206113157017377 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/flink/__init__.py0000644000175100017510000000142015206113032021475 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/flink/flink_streaming_impulse_source.py0000644000175100017510000000535715206113032026245 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A PTransform that provides an unbounded, streaming source of empty byte arrays. This can only be used with the flink runner. """ # pytype: skip-file import json from typing import Any from apache_beam import PTransform from apache_beam import Windowing from apache_beam import pvalue from apache_beam.transforms.window import GlobalWindows class FlinkStreamingImpulseSource(PTransform): URN = "flink:transform:streaming_impulse:v1" config: dict[str, Any] = {} def expand(self, pbegin): assert isinstance(pbegin, pvalue.PBegin), ( 'Input to transform must be a PBegin but found %s' % pbegin) return pvalue.PCollection(pbegin.pipeline, is_bounded=False) def get_windowing(self, unused_inputs): return Windowing(GlobalWindows()) def infer_output_type(self, unused_input_type): return bytes def to_runner_api_parameter(self, context): assert isinstance(self, FlinkStreamingImpulseSource), \ "expected instance of StreamingImpulseSource, but got %s" % self.__class__ return (self.URN, json.dumps(self.config)) def set_interval_ms(self, interval_ms): """Sets the interval (in milliseconds) between messages in the stream. """ self.config["interval_ms"] = interval_ms return self def set_message_count(self, message_count): """If non-zero, the stream will produce only this many total messages. Otherwise produces an unbounded number of messages. """ self.config["message_count"] = message_count return self @staticmethod @PTransform.register_urn(URN, None) def from_runner_api_parameter(_ptransform, spec_parameter, _context): if isinstance(spec_parameter, bytes): spec_parameter = spec_parameter.decode('utf-8') config = json.loads(spec_parameter) instance = FlinkStreamingImpulseSource() if "interval_ms" in config: instance.set_interval_ms(config["interval_ms"]) if "message_count" in config: instance.set_message_count(config["message_count"]) return instance ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/flink/flink_streaming_impulse_source_test.py0000644000175100017510000000262415206113032027276 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for flink_streaming_impulse_source.""" # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam.io.flink.flink_streaming_impulse_source import FlinkStreamingImpulseSource class FlinkStreamingImpulseSourceTest(unittest.TestCase): def test_serialization(self): p = beam.Pipeline() # pylint: disable=expression-not-assigned p | FlinkStreamingImpulseSource() # Test that roundtrip through Runner API works beam.Pipeline.from_runner_api(p.to_runner_api(), p.runner, p._options) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3664064 apache_beam-2.74.0/apache_beam/io/gcp/0000755000175100017510000000000015206113157017045 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/__init__.py0000644000175100017510000000272115206113032021150 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # Important: the MIME library in the Python 3.x standard library used by # apitools causes uploads containing '\r\n' to be corrupted, unless we # patch the BytesGenerator class to write contents verbatim. try: # pylint: disable=wrong-import-order, wrong-import-position # pylint: disable=ungrouped-imports import email.generator as email_generator from apitools.base.py import transfer class _WrapperNamespace(object): class BytesGenerator(email_generator.BytesGenerator): def _write_lines(self, lines): self.write(lines) transfer.email_generator = _WrapperNamespace except ImportError: # We may not have the GCP dependencies installed, so we pass in this case. pass ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/big_query_query_to_table_it_test.py0000644000175100017510000002354315206113032026235 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Integration test for Google Cloud BigQuery. """ # pytype: skip-file import base64 import datetime import logging import secrets import time import unittest import pytest from hamcrest.core.core.allof import all_of from tenacity import retry from tenacity import stop_after_attempt from apache_beam.io.gcp import big_query_query_to_table_pipeline from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryMatcher from apache_beam.testing import test_utils from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline # pylint: disable=wrong-import-order, wrong-import-position try: from apitools.base.py.exceptions import HttpError except ImportError: pass _LOGGER = logging.getLogger(__name__) WAIT_UNTIL_FINISH_DURATION_MS = 15 * 60 * 1000 BIG_QUERY_DATASET_ID = 'python_query_to_table_' NEW_TYPES_INPUT_TABLE = 'python_new_types_table' NEW_TYPES_OUTPUT_SCHEMA = ( '{"fields": [{"name": "bytes","type": "BYTES"},' '{"name": "date","type": "DATE"},{"name": "time","type": "TIME"}]}') NEW_TYPES_OUTPUT_VERIFY_QUERY = ('SELECT bytes, date, time FROM `%s`;') NEW_TYPES_OUTPUT_EXPECTED = [( b'xyw', datetime.date(2011, 1, 1), datetime.time(23, 59, 59, 999999), ), ( b'abc', datetime.date(2000, 1, 1), datetime.time(0, 0), ), ( b'\xe4\xbd\xa0\xe5\xa5\xbd', datetime.date(3000, 12, 31), datetime.time(23, 59, 59, 990000), ), ( b'\xab\xac\xad', datetime.date(2000, 1, 1), datetime.time(0, 0), )] LEGACY_QUERY = ( 'SELECT * FROM (SELECT "apple" as fruit), (SELECT "orange" as fruit),') STANDARD_QUERY = ( 'SELECT * FROM (SELECT "apple" as fruit) ' 'UNION ALL (SELECT "orange" as fruit)') NEW_TYPES_QUERY = ('SELECT bytes, date, time FROM [%s.%s]') DIALECT_OUTPUT_SCHEMA = ('{"fields": [{"name": "fruit","type": "STRING"}]}') DIALECT_OUTPUT_VERIFY_QUERY = ('SELECT fruit from `%s`;') DIALECT_OUTPUT_EXPECTED = [('apple', ), ('orange', )] class BigQueryQueryToTableIT(unittest.TestCase): def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.runner_name = type(self.test_pipeline.runner).__name__ self.project = self.test_pipeline.get_option('project') self.bigquery_client = BigQueryWrapper() self.dataset_id = '%s%d%s' % ( BIG_QUERY_DATASET_ID, int(time.time()), secrets.token_hex(3)) self.bigquery_client.get_or_create_dataset(self.project, self.dataset_id) self.output_table = "%s.output_table" % (self.dataset_id) def tearDown(self): request = bigquery.BigqueryDatasetsDeleteRequest( projectId=self.project, datasetId=self.dataset_id, deleteContents=True) try: self.bigquery_client.client.datasets.Delete(request) except HttpError: _LOGGER.debug('Failed to clean up dataset %s' % self.dataset_id) def _setup_new_types_env(self): table_schema = bigquery.TableSchema() table_field = bigquery.TableFieldSchema() table_field.name = 'bytes' table_field.type = 'BYTES' table_schema.fields.append(table_field) table_field = bigquery.TableFieldSchema() table_field.name = 'date' table_field.type = 'DATE' table_schema.fields.append(table_field) table_field = bigquery.TableFieldSchema() table_field.name = 'time' table_field.type = 'TIME' table_schema.fields.append(table_field) table = bigquery.Table( tableReference=bigquery.TableReference( projectId=self.project, datasetId=self.dataset_id, tableId=NEW_TYPES_INPUT_TABLE), schema=table_schema) request = bigquery.BigqueryTablesInsertRequest( projectId=self.project, datasetId=self.dataset_id, table=table) self.bigquery_client.client.tables.Insert(request) # Call get_table so that we wait until the table is visible. _ = self.bigquery_client.get_table( self.project, self.dataset_id, NEW_TYPES_INPUT_TABLE) table_data = [{ 'bytes': b'xyw', 'date': '2011-01-01', 'time': '23:59:59.999999' }, { 'bytes': b'abc', 'date': '2000-01-01', 'time': '00:00:00' }, { 'bytes': b'\xe4\xbd\xa0\xe5\xa5\xbd', 'date': '3000-12-31', 'time': '23:59:59.990000' }, { 'bytes': b'\xab\xac\xad', 'date': '2000-01-01', 'time': '00:00:00' }] # the API Tools bigquery client expects byte values to be base-64 encoded # TODO https://github.com/apache/beam/issues/19073: upgrade to # google-cloud-bigquery which does not require handling the encoding in # beam for row in table_data: row['bytes'] = base64.b64encode(row['bytes']).decode('utf-8') passed, errors = self.bigquery_client.insert_rows( self.project, self.dataset_id, NEW_TYPES_INPUT_TABLE, table_data) self.assertTrue(passed, 'Error in BQ setup: %s' % errors) @pytest.mark.it_postcommit @retry(reraise=True, stop=stop_after_attempt(3)) def test_big_query_legacy_sql(self): verify_query = DIALECT_OUTPUT_VERIFY_QUERY % self.output_table expected_checksum = test_utils.compute_hash(DIALECT_OUTPUT_EXPECTED) pipeline_verifiers = [ PipelineStateMatcher(), BigqueryMatcher( project=self.project, query=verify_query, checksum=expected_checksum) ] extra_opts = { 'query': LEGACY_QUERY, 'output': self.output_table, 'output_schema': DIALECT_OUTPUT_SCHEMA, 'use_standard_sql': False, 'wait_until_finish_duration': WAIT_UNTIL_FINISH_DURATION_MS, 'on_success_matcher': all_of(*pipeline_verifiers), } options = self.test_pipeline.get_full_options_as_args(**extra_opts) big_query_query_to_table_pipeline.run_bq_pipeline(options) @pytest.mark.it_postcommit def test_big_query_standard_sql(self): verify_query = DIALECT_OUTPUT_VERIFY_QUERY % self.output_table expected_checksum = test_utils.compute_hash(DIALECT_OUTPUT_EXPECTED) pipeline_verifiers = [ PipelineStateMatcher(), BigqueryMatcher( project=self.project, query=verify_query, checksum=expected_checksum) ] extra_opts = { 'query': STANDARD_QUERY, 'output': self.output_table, 'output_schema': DIALECT_OUTPUT_SCHEMA, 'use_standard_sql': True, 'wait_until_finish_duration': WAIT_UNTIL_FINISH_DURATION_MS, 'on_success_matcher': all_of(*pipeline_verifiers), } options = self.test_pipeline.get_full_options_as_args(**extra_opts) big_query_query_to_table_pipeline.run_bq_pipeline(options) @pytest.mark.it_postcommit def test_big_query_new_types(self): expected_checksum = test_utils.compute_hash(NEW_TYPES_OUTPUT_EXPECTED) verify_query = NEW_TYPES_OUTPUT_VERIFY_QUERY % self.output_table pipeline_verifiers = [ PipelineStateMatcher(), BigqueryMatcher( project=self.project, query=verify_query, checksum=expected_checksum) ] self._setup_new_types_env() extra_opts = { 'query': NEW_TYPES_QUERY % (self.dataset_id, NEW_TYPES_INPUT_TABLE), 'output': self.output_table, 'output_schema': NEW_TYPES_OUTPUT_SCHEMA, 'use_standard_sql': False, 'wait_until_finish_duration': WAIT_UNTIL_FINISH_DURATION_MS, 'use_json_exports': True, 'on_success_matcher': all_of(*pipeline_verifiers) } options = self.test_pipeline.get_full_options_as_args(**extra_opts) big_query_query_to_table_pipeline.run_bq_pipeline(options) @pytest.mark.it_postcommit def test_big_query_new_types_avro(self): expected_checksum = test_utils.compute_hash(NEW_TYPES_OUTPUT_EXPECTED) verify_query = NEW_TYPES_OUTPUT_VERIFY_QUERY % self.output_table pipeline_verifiers = [ PipelineStateMatcher(), BigqueryMatcher( project=self.project, query=verify_query, checksum=expected_checksum) ] self._setup_new_types_env() extra_opts = { 'query': NEW_TYPES_QUERY % (self.dataset_id, NEW_TYPES_INPUT_TABLE), 'output': self.output_table, 'output_schema': NEW_TYPES_OUTPUT_SCHEMA, 'use_standard_sql': False, 'wait_until_finish_duration': WAIT_UNTIL_FINISH_DURATION_MS, 'on_success_matcher': all_of(*pipeline_verifiers), } options = self.test_pipeline.get_full_options_as_args(**extra_opts) big_query_query_to_table_pipeline.run_bq_pipeline(options) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/big_query_query_to_table_pipeline.py0000644000175100017510000000715515206113032026370 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Job that reads and writes data to BigQuery. A Dataflow job that reads from BQ using a query and then writes to a big query table at the end of the pipeline. """ # pytype: skip-file # pylint: disable=wrong-import-order, wrong-import-position import argparse import apache_beam as beam from apache_beam.io.gcp.bigquery_tools import parse_table_schema_from_json from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline def run_bq_pipeline(argv=None): """Run the sample BigQuery pipeline. Args: argv: Arguments to the run function. """ parser = argparse.ArgumentParser() parser.add_argument( '--query', required=True, help='Query to process for the table.') parser.add_argument( '--output', required=True, help='Output BQ table to write results to.') parser.add_argument( '--output_schema', dest='output_schema', required=True, help='Schema for output BQ table.') parser.add_argument( '--use_standard_sql', action='store_true', dest='use_standard_sql', help='Output BQ table to write results to.') parser.add_argument( '--kms_key', default=None, help='Use this Cloud KMS key with BigQuery.') parser.add_argument( '--native', default=False, action='store_true', help='Use NativeSources and Sinks.') parser.add_argument( '--use_json_exports', default=False, action='store_true', help='Use JSON as the file format for exports.') known_args, pipeline_args = parser.parse_known_args(argv) table_schema = parse_table_schema_from_json(known_args.output_schema) kms_key = known_args.kms_key options = PipelineOptions(pipeline_args) p = TestPipeline(options=options) # Note to future modifiers: Keep using BigQuerySource if known_args.native is # True. if known_args.native: data = p | 'read' >> beam.io.Read( beam.io.BigQuerySource( query=known_args.query, use_standard_sql=known_args.use_standard_sql, kms_key=kms_key)) else: data = p | 'read' >> beam.io.gcp.bigquery.ReadFromBigQuery( query=known_args.query, project=options.view_as(GoogleCloudOptions).project, use_standard_sql=known_args.use_standard_sql, use_json_exports=known_args.use_json_exports, kms_key=kms_key) temp_file_format = ( 'NEWLINE_DELIMITED_JSON' if known_args.use_json_exports else 'AVRO') _ = data | 'write' >> beam.io.WriteToBigQuery( known_args.output, schema=table_schema, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=beam.io.BigQueryDisposition.WRITE_EMPTY, temp_file_format=temp_file_format, kms_key=kms_key) result = p.run() result.wait_until_finish() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery.py0000644000175100017510000041004715206113032021244 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """BigQuery sources and sinks. This module implements reading from and writing to BigQuery tables. It relies on several classes exposed by the BigQuery API: TableSchema, TableFieldSchema, TableRow, and TableCell. The default mode is to return table rows read from a BigQuery source as dictionaries. Similarly a Write transform to a BigQuerySink accepts PCollections of dictionaries. This is done for more convenient programming. If desired, the native TableRow objects can be used throughout to represent rows (use an instance of TableRowJsonCoder as a coder argument when creating the sources or sinks respectively). Also, for programming convenience, instances of TableReference and TableSchema have a string representation that can be used for the corresponding arguments: - TableReference can be a PROJECT:DATASET.TABLE or DATASET.TABLE string. - TableSchema can be a NAME:TYPE{,NAME:TYPE}* string (e.g. 'month:STRING,event_count:INTEGER'). The syntax supported is described here: https://cloud.google.com/bigquery/bq-command-line-tool-quickstart BigQuery sources can be used as main inputs or side inputs. A main input (common case) is expected to be massive and will be split into manageable chunks and processed in parallel. Side inputs are expected to be small and will be read completely every time a ParDo DoFn gets executed. In the example below the lambda function implementing the DoFn for the Map transform will get on each call *one* row of the main table and *all* rows of the side table. The runner may use some caching techniques to share the side inputs between calls in order to avoid excessive reading::: main_table = pipeline | 'VeryBig' >> beam.io.ReadFromBigQuery(...) side_table = pipeline | 'NotBig' >> beam.io.ReadFromBigQuery(...) results = ( main_table | 'ProcessData' >> beam.Map( lambda element, side_input: ..., AsList(side_table))) There is no difference in how main and side inputs are read. What makes the side_table a 'side input' is the AsList wrapper used when passing the table as a parameter to the Map transform. AsList signals to the execution framework that its input should be made available whole. The main and side inputs are implemented differently. Reading a BigQuery table as main input entails exporting the table to a set of GCS files (in AVRO or in JSON format) and then processing those files. Users may provide a query to read from rather than reading all of a BigQuery table. If specified, the result obtained by executing the specified query will be used as the data of the input transform.:: query_results = pipeline | beam.io.gcp.bigquery.ReadFromBigQuery( query='SELECT year, mean_temp FROM samples.weather_stations') When creating a BigQuery input transform, users should provide either a query or a table. Pipeline construction will fail with a validation error if neither or both are specified. When reading via `ReadFromBigQuery` using `EXPORT`, bytes are returned decoded as bytes. This is due to the fact that ReadFromBigQuery uses Avro exports by default. When reading from BigQuery using `apache_beam.io.BigQuerySource`, bytes are returned as base64-encoded bytes. To get base64-encoded bytes using `ReadFromBigQuery`, you can use the flag `use_json_exports` to export data as JSON, and receive base64-encoded bytes. ReadAllFromBigQuery ------------------- Beam 2.27.0 introduces a new transform called `ReadAllFromBigQuery` which allows you to define table and query reads from BigQuery at pipeline runtime.::: read_requests = p | beam.Create([ ReadFromBigQueryRequest(query='SELECT * FROM mydataset.mytable'), ReadFromBigQueryRequest(table='myproject.mydataset.mytable')]) results = read_requests | ReadAllFromBigQuery() A good application for this transform is in streaming pipelines to refresh a side input coming from BigQuery. This would work like so::: side_input = ( p | 'PeriodicImpulse' >> PeriodicImpulse( first_timestamp, last_timestamp, interval, True) | 'MapToReadRequest' >> beam.Map( lambda x: ReadFromBigQueryRequest(table='dataset.table')) | beam.io.ReadAllFromBigQuery()) main_input = ( p | 'MpImpulse' >> beam.Create(sample_main_input_elements) | 'MapMpToTimestamped' >> beam.Map(lambda src: TimestampedValue(src, src)) | 'WindowMpInto' >> beam.WindowInto( window.FixedWindows(main_input_windowing_interval))) result = ( main_input | 'ApplyCrossJoin' >> beam.FlatMap( cross_join, rights=beam.pvalue.AsIter(side_input))) **Note**: This transform is supported on Portable and Dataflow v2 runners. **Note**: This transform does not currently clean up temporary datasets created for its execution. (BEAM-11359) Writing Data to BigQuery ======================== The `WriteToBigQuery` transform is the recommended way of writing data to BigQuery. It supports a large set of parameters to customize how you'd like to write to BigQuery. Table References ---------------- This transform allows you to provide static `project`, `dataset` and `table` parameters which point to a specific BigQuery table to be created. The `table` parameter can also be a dynamic parameter (i.e. a callable), which receives an element to be written to BigQuery, and returns the table that that element should be sent to. You may also provide a tuple of PCollectionView elements to be passed as side inputs to your callable. For example, suppose that one wishes to send events of different types to different tables, and the table names are computed at pipeline runtime, one may do something like the following:: with Pipeline() as p: elements = (p | 'Create elements' >> beam.Create([ {'type': 'error', 'timestamp': '12:34:56', 'message': 'bad'}, {'type': 'user_log', 'timestamp': '12:34:59', 'query': 'flu symptom'}, ])) table_names = (p | 'Create table_names' >> beam.Create([ ('error', 'my_project:dataset1.error_table_for_today'), ('user_log', 'my_project:dataset1.query_table_for_today'), ])) table_names_dict = beam.pvalue.AsDict(table_names) elements | beam.io.gcp.bigquery.WriteToBigQuery( table=lambda row, table_dict: table_dict[row['type']], table_side_inputs=(table_names_dict,)) In the example above, the `table_dict` argument passed to the function in `table_dict` is the side input coming from `table_names_dict`, which is passed as part of the `table_side_inputs` argument. Schemas --------- This transform also allows you to provide a static or dynamic `schema` parameter (i.e. a callable). If providing a callable, this should take in a table reference (as returned by the `table` parameter), and return the corresponding schema for that table. This allows to provide different schemas for different tables:: def compute_table_name(row): ... errors_schema = {'fields': [ {'name': 'type', 'type': 'STRING', 'mode': 'NULLABLE'}, {'name': 'message', 'type': 'STRING', 'mode': 'NULLABLE'}]} queries_schema = {'fields': [ {'name': 'type', 'type': 'STRING', 'mode': 'NULLABLE'}, {'name': 'query', 'type': 'STRING', 'mode': 'NULLABLE'}]} with Pipeline() as p: elements = (p | beam.Create([ {'type': 'error', 'timestamp': '12:34:56', 'message': 'bad'}, {'type': 'user_log', 'timestamp': '12:34:59', 'query': 'flu symptom'}, ])) elements | beam.io.gcp.bigquery.WriteToBigQuery( table=compute_table_name, schema=lambda table: (errors_schema if 'errors' in table else queries_schema)) It may be the case that schemas are computed at pipeline runtime. In cases like these, one can also provide a `schema_side_inputs` parameter, which is a tuple of PCollectionViews to be passed to the schema callable (much like the `table_side_inputs` parameter). Additional Parameters for BigQuery Tables ----------------------------------------- This sink is able to create tables in BigQuery if they don't already exist. It also relies on creating temporary tables when performing file loads. The WriteToBigQuery transform creates tables using the BigQuery API by inserting a load job (see the API reference [1]), or by inserting a new table (see the API reference for that [2][3]). When creating a new BigQuery table, there are a number of extra parameters that one may need to specify. For example, clustering, partitioning, data encoding, etc. It is possible to provide these additional parameters by passing a Python dictionary as `additional_bq_parameters` to the transform. As an example, to create a table that has specific partitioning, and clustering properties, one would do the following:: additional_bq_parameters = { 'timePartitioning': {'type': 'DAY'}, 'clustering': {'fields': ['country']}} with Pipeline() as p: elements = (p | beam.Create([ {'country': 'mexico', 'timestamp': '12:34:56', 'query': 'acapulco'}, {'country': 'canada', 'timestamp': '12:34:59', 'query': 'influenza'}, ])) elements | beam.io.gcp.bigquery.WriteToBigQuery( table='project_name1:dataset_2.query_events_table', additional_bq_parameters=additional_bq_parameters) Much like the schema case, the parameter with `additional_bq_parameters` can also take a callable that receives a table reference. [1] https://cloud.google.com/bigquery/docs/reference/rest/v2/Job\ #jobconfigurationload [2] https://cloud.google.com/bigquery/docs/reference/rest/v2/tables/insert [3] https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#resource Chaining of operations after WriteToBigQuery -------------------------------------------- WriteToBigQuery returns an object with several PCollections that consist of metadata about the write operations. These are useful to inspect the write operation and follow with the results:: schema = {'fields': [ {'name': 'column', 'type': 'STRING', 'mode': 'NULLABLE'}]} error_schema = {'fields': [ {'name': 'destination', 'type': 'STRING', 'mode': 'NULLABLE'}, {'name': 'row', 'type': 'STRING', 'mode': 'NULLABLE'}, {'name': 'error_message', 'type': 'STRING', 'mode': 'NULLABLE'}]} with Pipeline() as p: result = (p | 'Create Columns' >> beam.Create([ {'column': 'value'}, {'bad_column': 'bad_value'} ]) | 'Write Data' >> WriteToBigQuery( method=WriteToBigQuery.Method.STREAMING_INSERTS, table=my_table, schema=schema, insert_retry_strategy=RetryStrategy.RETRY_NEVER )) _ = (result.failed_rows_with_errors | 'Get Errors' >> beam.Map(lambda e: { "destination": e[0], "row": json.dumps(e[1]), "error_message": e[2][0]['message'] }) | 'Write Errors' >> WriteToBigQuery( method=WriteToBigQuery.Method.STREAMING_INSERTS, table=error_log_table, schema=error_schema, )) Often, the simplest use case is to chain an operation after writing data to BigQuery.To do this, one can chain the operation after one of the output PCollections. A generic way in which this operation (independent of write method) could look like:: def chain_after(result): try: # This works for FILE_LOADS, where we run load and possibly copy jobs. return (result.destination_load_jobid_pairs, result.destination_copy_jobid_pairs) | beam.Flatten() except AttributeError: # Works for STREAMING_INSERTS, where we return the rows BigQuery rejected return result.failed_rows result = (pcoll | WriteToBigQuery(...)) _ = (chain_after(result) | beam.Reshuffle() # Force a 'commit' of the intermediate date | MyOperationAfterWriteToBQ()) Attributes can be accessed using dot notation or bracket notation: result.failed_rows <--> result['FailedRows'] result.failed_rows_with_errors <--> result['FailedRowsWithErrors'] result.destination_load_jobid_pairs <--> result['destination_load_jobid_pairs'] result.destination_file_pairs <--> result['destination_file_pairs'] result.destination_copy_jobid_pairs <--> result['destination_copy_jobid_pairs'] Writing with Storage Write API using Cross Language --------------------------------------------------- This sink is able to write with BigQuery's Storage Write API. To do so, specify the method `WriteToBigQuery.Method.STORAGE_WRITE_API`. This will use the StorageWriteToBigQuery() transform to discover and use the Java implementation. Using this transform directly will require the use of beam.Row() elements. Similar to streaming inserts, it returns two dead-letter queue PCollections: one containing just the failed rows and the other containing failed rows and errors. They can be accessed with `failed_rows` and `failed_rows_with_errors`, respectively. See the examples above for how to do this. *** Short introduction to BigQuery concepts *** Tables have rows (TableRow) and each row has cells (TableCell). A table has a schema (TableSchema), which in turn describes the schema of each cell (TableFieldSchema). The terms field and cell are used interchangeably. TableSchema: Describes the schema (types and order) for values in each row. Has one attribute, 'field', which is list of TableFieldSchema objects. TableFieldSchema: Describes the schema (type, name) for one field. Has several attributes, including 'name' and 'type'. Common values for the type attribute are: 'STRING', 'INTEGER', 'FLOAT', 'BOOLEAN', 'NUMERIC', 'GEOGRAPHY'. All possible values are described at: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types TableRow: Holds all values in a table row. Has one attribute, 'f', which is a list of TableCell instances. TableCell: Holds the value for one cell (or field). Has one attribute, 'v', which is a JsonValue instance. This class is defined in apitools.base.py.extra_types.py module. As of Beam 2.7.0, the NUMERIC data type is supported. This data type supports high-precision decimal numbers (precision of 38 digits, scale of 9 digits). The GEOGRAPHY data type works with Well-Known Text (See https://en.wikipedia.org/wiki/Well-known_text) format for reading and writing to BigQuery. BigQuery IO requires values of BYTES datatype to be encoded using base64 encoding when writing to BigQuery. **Updates to the I/O connector code** For any significant updates to this I/O connector, please consider involving corresponding code reviewers mentioned in https://github.com/apache/beam/blob/master/sdks/python/OWNERS """ # pytype: skip-file import collections import io import itertools import json import logging import random import secrets import time import uuid import warnings from dataclasses import dataclass from typing import Optional from typing import Union import fastavro from objsize import get_deep_size import apache_beam as beam from apache_beam import coders from apache_beam import pvalue from apache_beam.internal.gcp.json_value import from_json_value from apache_beam.internal.gcp.json_value import to_json_value from apache_beam.io import range_trackers from apache_beam.io.avroio import _create_avro_source as create_avro_source from apache_beam.io.filesystems import CompressionTypes from apache_beam.io.filesystems import FileSystems from apache_beam.io.gcp import bigquery_schema_tools from apache_beam.io.gcp import bigquery_tools from apache_beam.io.gcp.bigquery_io_metadata import create_bigquery_io_metadata from apache_beam.io.gcp.bigquery_read_internal import _BigQueryReadSplit from apache_beam.io.gcp.bigquery_read_internal import _JsonToDictCoder from apache_beam.io.gcp.bigquery_read_internal import _PassThroughThenCleanup from apache_beam.io.gcp.bigquery_read_internal import _PassThroughThenCleanupTempDatasets from apache_beam.io.gcp.bigquery_read_internal import bigquery_export_destination_uri from apache_beam.io.gcp.bigquery_tools import RetryStrategy from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.io.iobase import BoundedSource from apache_beam.io.iobase import RangeTracker from apache_beam.io.iobase import SDFBoundedSourceReader from apache_beam.io.iobase import SourceBundle from apache_beam.io.textio import _TextSource as TextSource from apache_beam.metrics import Metrics from apache_beam.metrics.metric import Lineage from apache_beam.options import value_provider as vp from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.value_provider import StaticValueProvider from apache_beam.options.value_provider import ValueProvider from apache_beam.options.value_provider import check_accessible from apache_beam.pvalue import PCollection from apache_beam.transforms import DoFn from apache_beam.transforms import ParDo from apache_beam.transforms import PTransform from apache_beam.transforms.display import DisplayDataItem from apache_beam.transforms.external import BeamJarExpansionService from apache_beam.transforms.external import SchemaAwareExternalTransform from apache_beam.transforms.sideinputs import SIDE_INPUT_PREFIX from apache_beam.transforms.sideinputs import get_sideinput_index from apache_beam.transforms.util import ReshufflePerKey from apache_beam.typehints.row_type import RowTypeConstraint from apache_beam.typehints.schemas import schema_from_element_type from apache_beam.utils import retry from apache_beam.utils.annotations import deprecated try: from apache_beam.io.gcp.internal.clients.bigquery import DatasetReference from apache_beam.io.gcp.internal.clients.bigquery import JobReference from apache_beam.io.gcp.internal.clients.bigquery import TableReference except ImportError: DatasetReference = None TableReference = None JobReference = None _LOGGER = logging.getLogger(__name__) try: import google.cloud.bigquery_storage_v1 as bq_storage except ImportError: _LOGGER.info( 'No module named google.cloud.bigquery_storage_v1. ' 'As a result, the ReadFromBigQuery transform *CANNOT* be ' 'used with `method=DIRECT_READ`.') __all__ = [ 'TableRowJsonCoder', 'BigQueryDisposition', 'BigQuerySource', 'BigQuerySink', 'BigQueryQueryPriority', 'WriteToBigQuery', 'WriteResult', 'ReadFromBigQuery', 'ReadFromBigQueryRequest', 'ReadAllFromBigQuery', 'SCHEMA_AUTODETECT', ] """ Template for BigQuery jobs created by BigQueryIO. This template is: `"beam_bq_job_{job_type}_{job_id}_{step_id}_{random}"`, where: - `job_type` represents the BigQuery job type (e.g. extract / copy / load / query). - `job_id` is the Beam job name. - `step_id` is a UUID representing the Dataflow step that created the BQ job. - `random` is a random string. NOTE: This job name template does not have backwards compatibility guarantees. """ BQ_JOB_NAME_TEMPLATE = "beam_bq_job_{job_type}_{job_id}_{step_id}{random}" """ The maximum number of times that a bundle of rows that errors out should be sent for insertion into BigQuery. The default is 10,000 with exponential backoffs, so a bundle of rows may be tried for a very long time. You may reduce this property to reduce the number of retries. """ MAX_INSERT_RETRIES = 10000 """ The maximum byte size for a BigQuery legacy streaming insert payload. Note: The actual limit is 10MB, but we set it to 9MB to make room for request overhead: https://cloud.google.com/bigquery/quotas#streaming_inserts """ MAX_INSERT_PAYLOAD_SIZE = 9 << 20 @deprecated(since='2.11.0', current="bigquery_tools.parse_table_reference") def _parse_table_reference(table, dataset=None, project=None): return bigquery_tools.parse_table_reference(table, dataset, project) @deprecated( since='2.11.0', current="bigquery_tools.parse_table_schema_from_json") def parse_table_schema_from_json(schema_string): return bigquery_tools.parse_table_schema_from_json(schema_string) @deprecated(since='2.11.0', current="bigquery_tools.default_encoder") def default_encoder(obj): return bigquery_tools.default_encoder(obj) @deprecated(since='2.11.0', current="bigquery_tools.RowAsDictJsonCoder") def RowAsDictJsonCoder(*args, **kwargs): return bigquery_tools.RowAsDictJsonCoder(*args, **kwargs) @deprecated(since='2.11.0', current="bigquery_tools.BigQueryWrapper") def BigQueryWrapper(*args, **kwargs): return bigquery_tools.BigQueryWrapper(*args, **kwargs) class TableRowJsonCoder(coders.Coder): """A coder for a TableRow instance to/from a JSON string. Note that the encoding operation (used when writing to sinks) requires the table schema in order to obtain the ordered list of field names. Reading from sources on the other hand does not need the table schema. """ def __init__(self, table_schema=None): # The table schema is needed for encoding TableRows as JSON (writing to # sinks) because the ordered list of field names is used in the JSON # representation. self.table_schema = table_schema # Precompute field names since we need them for row encoding. if self.table_schema: self.field_names = tuple(fs.name for fs in self.table_schema.fields) self.field_types = tuple(fs.type for fs in self.table_schema.fields) def encode(self, table_row): if self.table_schema is None: raise AttributeError( 'The TableRowJsonCoder requires a table schema for ' 'encoding operations. Please specify a table_schema argument.') try: return json.dumps( collections.OrderedDict( zip( self.field_names, [from_json_value(f.v) for f in table_row.f])), allow_nan=False, default=bigquery_tools.default_encoder) except ValueError as e: raise ValueError('%s. %s' % (e, bigquery_tools.JSON_COMPLIANCE_ERROR)) def decode(self, encoded_table_row): od = json.loads( encoded_table_row, object_pairs_hook=collections.OrderedDict) return bigquery.TableRow( f=[bigquery.TableCell(v=to_json_value(e)) for e in od.values()]) class BigQueryDisposition(object): """Class holding standard strings used for create and write dispositions.""" CREATE_NEVER = 'CREATE_NEVER' CREATE_IF_NEEDED = 'CREATE_IF_NEEDED' WRITE_TRUNCATE = 'WRITE_TRUNCATE' WRITE_APPEND = 'WRITE_APPEND' WRITE_EMPTY = 'WRITE_EMPTY' @staticmethod def validate_create(disposition): values = ( BigQueryDisposition.CREATE_NEVER, BigQueryDisposition.CREATE_IF_NEEDED) if disposition not in values: raise ValueError( 'Invalid create disposition %s. Expecting %s' % (disposition, values)) return disposition @staticmethod def validate_write(disposition): values = ( BigQueryDisposition.WRITE_TRUNCATE, BigQueryDisposition.WRITE_APPEND, BigQueryDisposition.WRITE_EMPTY) if disposition not in values: raise ValueError( 'Invalid write disposition %s. Expecting %s' % (disposition, values)) return disposition class BigQueryQueryPriority(object): """Class holding standard strings used for query priority.""" INTERACTIVE = 'INTERACTIVE' BATCH = 'BATCH' # ----------------------------------------------------------------------------- # BigQuerySource, BigQuerySink. @deprecated(since='2.25.0', current="ReadFromBigQuery") def BigQuerySource( table=None, dataset=None, project=None, query=None, validate=False, coder=None, use_standard_sql=False, flatten_results=True, kms_key=None, use_dataflow_native_source=False): if use_dataflow_native_source: warnings.warn( "Native sources no longer implemented; " "falling back to standard Beam source.") return ReadFromBigQuery( table=table, dataset=dataset, project=project, query=query, validate=validate, coder=coder, use_standard_sql=use_standard_sql, flatten_results=flatten_results, use_json_exports=True, kms_key=kms_key) @deprecated(since='2.25.0', current="ReadFromBigQuery") def _BigQuerySource(*args, **kwargs): """A source based on a BigQuery table.""" warnings.warn( "Native sources no longer implemented; " "falling back to standard Beam source.") return ReadFromBigQuery(*args, **kwargs) # TODO(https://github.com/apache/beam/issues/21622): remove the serialization # restriction in transform implementation once InteractiveRunner can work # without runner api roundtrips. @dataclass class _BigQueryExportResult: coder: beam.coders.Coder paths: list[str] class _CustomBigQuerySource(BoundedSource): def __init__( self, method, gcs_location=None, table=None, dataset=None, project=None, query=None, validate=False, pipeline_options=None, coder=None, use_standard_sql=False, flatten_results=True, kms_key=None, bigquery_job_labels=None, use_json_exports=False, job_name=None, step_name=None, unique_id=None, temp_dataset=None, query_priority=BigQueryQueryPriority.BATCH): if table is not None and query is not None: raise ValueError( 'Both a BigQuery table and a query were specified.' ' Please specify only one of these.') elif table is None and query is None: raise ValueError('A BigQuery table or a query must be specified') elif table is not None: self.table_reference = bigquery_tools.parse_table_reference( table, dataset, project) self.query = None self.use_legacy_sql = True else: if isinstance(query, str): query = StaticValueProvider(str, query) self.query = query # TODO(BEAM-1082): Change the internal flag to be standard_sql self.use_legacy_sql = not use_standard_sql self.table_reference = None self.method = method self.gcs_location = gcs_location self.project = project self.validate = validate self.flatten_results = flatten_results self.coder = coder or _JsonToDictCoder self.kms_key = kms_key self.export_result = None self.options = pipeline_options self.bq_io_metadata = None # Populate in setup, as it may make an RPC self.bigquery_job_labels = bigquery_job_labels or {} self.use_json_exports = use_json_exports self.temp_dataset = temp_dataset self.query_priority = query_priority self._job_name = job_name or 'BQ_EXPORT_JOB' self._step_name = step_name self._source_uuid = unique_id def _get_bq_metadata(self): if not self.bq_io_metadata: self.bq_io_metadata = create_bigquery_io_metadata(self._step_name) return self.bq_io_metadata def display_data(self): export_format = 'JSON' if self.use_json_exports else 'AVRO' return { 'method': str(self.method), 'table': str(self.table_reference), 'query': str(self.query), 'project': str(self.project), 'use_legacy_sql': self.use_legacy_sql, 'bigquery_job_labels': json.dumps(self.bigquery_job_labels), 'export_file_format': export_format, 'launchesBigQueryJobs': DisplayDataItem( True, label="This Dataflow job launches bigquery jobs."), } def estimate_size(self): bq = bigquery_tools.BigQueryWrapper.from_pipeline_options(self.options) if self.table_reference is not None: table_ref = self.table_reference if (isinstance(self.table_reference, vp.ValueProvider) and self.table_reference.is_accessible()): table_ref = bigquery_tools.parse_table_reference( self.table_reference.get(), project=self._get_project()) elif isinstance(self.table_reference, vp.ValueProvider): # Size estimation is best effort. We return None as we have # no access to the table that we're querying. return None if not table_ref.projectId: table_ref.projectId = self._get_project() table = bq.get_table( table_ref.projectId, table_ref.datasetId, table_ref.tableId) return int(table.numBytes) elif self.query is not None and self.query.is_accessible(): project = self._get_project() query_job_name = bigquery_tools.generate_bq_job_name( self._job_name, self._source_uuid, bigquery_tools.BigQueryJobTypes.QUERY, '%s_%s' % (int(time.time()), random.randint(0, 1000))) job = bq._start_query_job( project, self.query.get(), self.use_legacy_sql, self.flatten_results, job_id=query_job_name, priority=self.query_priority, dry_run=True, kms_key=self.kms_key, job_labels=self._get_bq_metadata().add_additional_bq_job_labels( self.bigquery_job_labels)) if job.statistics.totalBytesProcessed is None: # Some queries may not have access to `totalBytesProcessed` as a # result of row-level security. # > BigQuery hides sensitive statistics on all queries against # > tables with row-level security. # See cloud.google.com/bigquery/docs/managing-row-level-security # and cloud.google.com/bigquery/docs/best-practices-row-level-security return None return int(job.statistics.totalBytesProcessed) else: # Size estimation is best effort. We return None as we have # no access to the query that we're running. return None def _get_project(self): """Returns the project that queries and exports will be billed to.""" project = self.options.view_as(GoogleCloudOptions).project if isinstance(project, vp.ValueProvider): project = project.get() if self.temp_dataset: return self.temp_dataset.projectId if not project: project = self.project return project def _create_source(self, path, coder): if not self.use_json_exports: return create_avro_source(path, validate=self.validate) else: return TextSource( path, min_bundle_size=0, compression_type=CompressionTypes.UNCOMPRESSED, strip_trailing_newlines=True, coder=coder, validate=self.validate) def split(self, desired_bundle_size, start_position=None, stop_position=None): if self.export_result is None: bq = bigquery_tools.BigQueryWrapper( temp_dataset_id=( self.temp_dataset.datasetId if self.temp_dataset else None), client=bigquery_tools.BigQueryWrapper._bigquery_client(self.options)) if self.query is not None: self._setup_temporary_dataset(bq) self.table_reference = self._execute_query(bq) if isinstance(self.table_reference, vp.ValueProvider): self.table_reference = bigquery_tools.parse_table_reference( self.table_reference.get(), project=self._get_project()) elif not self.table_reference.projectId: self.table_reference.projectId = self._get_project() Lineage.sources().add( 'bigquery', self.table_reference.projectId, self.table_reference.datasetId, self.table_reference.tableId) schema, metadata_list = self._export_files(bq) self.export_result = _BigQueryExportResult( coder=self.coder(schema), paths=[metadata.path for metadata in metadata_list]) if self.query is not None: bq.clean_up_temporary_dataset(self._get_project()) for path in self.export_result.paths: source = self._create_source(path, self.export_result.coder) yield SourceBundle( weight=1.0, source=source, start_position=None, stop_position=None) def get_range_tracker(self, start_position, stop_position): class CustomBigQuerySourceRangeTracker(RangeTracker): """A RangeTracker that always returns positions as None.""" def start_position(self): return None def stop_position(self): return None return CustomBigQuerySourceRangeTracker() def read(self, range_tracker): raise NotImplementedError('BigQuery source must be split before being read') @check_accessible(['query']) def _setup_temporary_dataset(self, bq): if self.temp_dataset: # Temp dataset was provided by the user so we can just return. return location = bq.get_query_location( self._get_project(), self.query.get(), self.use_legacy_sql) bq.create_temporary_dataset( self._get_project(), location, kms_key=self.kms_key) @check_accessible(['query']) def _execute_query(self, bq): query_job_name = bigquery_tools.generate_bq_job_name( self._job_name, self._source_uuid, bigquery_tools.BigQueryJobTypes.QUERY, '%s_%s' % (int(time.time()), random.randint(0, 1000))) job = bq._start_query_job( self._get_project(), self.query.get(), self.use_legacy_sql, self.flatten_results, job_id=query_job_name, priority=self.query_priority, kms_key=self.kms_key, job_labels=self._get_bq_metadata().add_additional_bq_job_labels( self.bigquery_job_labels)) job_ref = job.jobReference bq.wait_for_bq_job(job_ref, max_retries=0) return bq._get_temp_table(self._get_project()) def _export_files(self, bq): """Runs a BigQuery export job. Returns: bigquery.TableSchema instance, a list of FileMetadata instances """ job_labels = self._get_bq_metadata().add_additional_bq_job_labels( self.bigquery_job_labels) export_job_name = bigquery_tools.generate_bq_job_name( self._job_name, self._source_uuid, bigquery_tools.BigQueryJobTypes.EXPORT, '%s_%s' % (int(time.time()), random.randint(0, 1000))) temp_location = self.options.view_as(GoogleCloudOptions).temp_location gcs_location = bigquery_export_destination_uri( self.gcs_location, temp_location, self._source_uuid) try: if self.use_json_exports: job_ref = bq.perform_extract_job([gcs_location], export_job_name, self.table_reference, bigquery_tools.FileFormat.JSON, project=self._get_project(), job_labels=job_labels, include_header=False) else: job_ref = bq.perform_extract_job([gcs_location], export_job_name, self.table_reference, bigquery_tools.FileFormat.AVRO, project=self._get_project(), include_header=False, job_labels=job_labels, use_avro_logical_types=True) bq.wait_for_bq_job(job_ref) except Exception as exn: # pylint: disable=broad-except # The error messages thrown in this case are generic and misleading, # so leave this breadcrumb in case it's the root cause. logging.warning( "Error exporting table: %s. " "Note that external tables cannot be exported: " "https://cloud.google.com/bigquery/docs/external-tables" "#external_table_limitations", exn) raise metadata_list = FileSystems.match([gcs_location])[0].metadata_list if isinstance(self.table_reference, vp.ValueProvider): table_ref = bigquery_tools.parse_table_reference( self.table_reference.get(), project=self.project) else: table_ref = self.table_reference table = bq.get_table( table_ref.projectId, table_ref.datasetId, table_ref.tableId) return table.schema, metadata_list class _CustomBigQueryStorageSource(BoundedSource): """A base class for BoundedSource implementations which read from BigQuery using the BigQuery Storage API. Args: table (str, TableReference): The ID of the table. If **dataset** argument is :data:`None` then the table argument must contain the entire table reference specified as: ``'PROJECT:DATASET.TABLE'`` or must specify a TableReference. dataset (str): Optional ID of the dataset containing this table or :data:`None` if the table argument specifies a TableReference. project (str): Optional ID of the project containing this table or :data:`None` if the table argument specifies a TableReference. selected_fields (List[str]): Optional List of names of the fields in the table that should be read. If empty, all fields will be read. If the specified field is a nested field, all the sub-fields in the field will be selected. The output field order is unrelated to the order of fields in selected_fields. row_restriction (str): Optional SQL text filtering statement, similar to a WHERE clause in a query. Aggregates are not supported. Restricted to a maximum length for 1 MB. use_native_datetime (bool): If :data:`True`, BigQuery DATETIME fields will be returned as native Python datetime objects. If :data:`False`, DATETIME fields will be returned as formatted strings (for example: 2021-01-01T12:59:59). The default is :data:`False`. """ # The maximum number of streams which will be requested when creating a read # session, regardless of the desired bundle size. MAX_SPLIT_COUNT = 10000 # The minimum number of streams which will be requested when creating a read # session, regardless of the desired bundle size. Note that the server may # still choose to return fewer than ten streams based on the layout of the # table. MIN_SPLIT_COUNT = 10 def __init__( self, method: str, query_priority: [BigQueryQueryPriority] = BigQueryQueryPriority.BATCH, table: Optional[Union[str, TableReference]] = None, dataset: Optional[str] = None, project: Optional[str] = None, query: Optional[str] = None, selected_fields: Optional[list[str]] = None, row_restriction: Optional[str] = None, pipeline_options: Optional[GoogleCloudOptions] = None, unique_id: Optional[uuid.UUID] = None, bigquery_job_labels: Optional[dict] = None, bigquery_dataset_labels: Optional[dict] = None, job_name: Optional[str] = None, step_name: Optional[str] = None, use_standard_sql: Optional[bool] = False, flatten_results: Optional[bool] = True, kms_key: Optional[str] = None, temp_dataset: Optional[DatasetReference] = None, temp_table: Optional[TableReference] = None, use_native_datetime: Optional[bool] = False, timeout: Optional[float] = None): if table is not None and query is not None: raise ValueError( 'Both a BigQuery table and a query were specified.' ' Please specify only one of these.') elif table is None and query is None: raise ValueError('A BigQuery table or a query must be specified') elif table is not None: self.table_reference = bigquery_tools.parse_table_reference( table, dataset, project) self.query = None self.use_legacy_sql = True else: if isinstance(query, str): query = StaticValueProvider(str, query) self.query = query # TODO(BEAM-1082): Change the internal flag to be standard_sql self.use_legacy_sql = not use_standard_sql self.table_reference = None self.method = method self.project = project self.selected_fields = selected_fields self.row_restriction = row_restriction self.pipeline_options = pipeline_options self.split_result = None self.bigquery_job_labels = bigquery_job_labels or {} self.bigquery_dataset_labels = bigquery_dataset_labels or {} self.bq_io_metadata = None # Populate in setup, as it may make an RPC self.flatten_results = flatten_results self.kms_key = kms_key self.temp_table = temp_table self.query_priority = query_priority self.use_native_datetime = use_native_datetime self.timeout = timeout self._job_name = job_name or 'BQ_DIRECT_READ_JOB' self._step_name = step_name self._source_uuid = unique_id def _get_project(self): """Returns the project that queries and exports will be billed to.""" if self.pipeline_options: project = self.pipeline_options.view_as(GoogleCloudOptions).project if isinstance(project, vp.ValueProvider): project = project.get() if project: return project return self.project def _get_parent_project(self): """Returns the project that will be billed.""" if self.temp_table: return self.temp_table.projectId project = self.pipeline_options.view_as(GoogleCloudOptions).project if isinstance(project, vp.ValueProvider): project = project.get() if not project: project = self.project return project def _get_table_size(self, bq, table_reference): project = ( table_reference.projectId if table_reference.projectId else self._get_parent_project()) table = bq.get_table( project, table_reference.datasetId, table_reference.tableId) return table.numBytes def _get_bq_metadata(self): if not self.bq_io_metadata: self.bq_io_metadata = create_bigquery_io_metadata(self._step_name) return self.bq_io_metadata @check_accessible(['query']) def _setup_temporary_dataset(self, bq): if self.temp_table: # Temp dataset was provided by the user so we can just return. return location = bq.get_query_location( self._get_parent_project(), self.query.get(), self.use_legacy_sql) _LOGGER.warning("### Labels: %s", str(self.bigquery_dataset_labels)) bq.create_temporary_dataset( self._get_parent_project(), location, self.bigquery_dataset_labels, kms_key=self.kms_key) @check_accessible(['query']) def _execute_query(self, bq): query_job_name = bigquery_tools.generate_bq_job_name( self._job_name, self._source_uuid, bigquery_tools.BigQueryJobTypes.QUERY, '%s_%s' % (int(time.time()), random.randint(0, 1000))) job = bq._start_query_job( self._get_parent_project(), self.query.get(), self.use_legacy_sql, self.flatten_results, job_id=query_job_name, priority=self.query_priority, kms_key=self.kms_key, job_labels=self._get_bq_metadata().add_additional_bq_job_labels( self.bigquery_job_labels)) job_ref = job.jobReference bq.wait_for_bq_job(job_ref, max_retries=0) table_reference = bq._get_temp_table(self._get_parent_project()) return table_reference def display_data(self): return { 'method': self.method, 'output_format': 'ARROW' if self.use_native_datetime else 'AVRO', 'project': str(self.project), 'table_reference': str(self.table_reference), 'query': str(self.query), 'use_legacy_sql': self.use_legacy_sql, 'use_native_datetime': self.use_native_datetime, 'selected_fields': str(self.selected_fields), 'row_restriction': str(self.row_restriction), 'launchesBigQueryJobs': DisplayDataItem( True, label="This Dataflow job launches bigquery jobs."), } def estimate_size(self): # Returns the pre-filtering size of the (temporary) table being read. bq = bigquery_tools.BigQueryWrapper.from_pipeline_options( self.pipeline_options) if self.table_reference is not None: table_ref = self.table_reference if (isinstance(self.table_reference, vp.ValueProvider) and self.table_reference.is_accessible()): table_ref = bigquery_tools.parse_table_reference( self.table_reference.get(), project=self._get_project()) elif isinstance(self.table_reference, vp.ValueProvider): # Size estimation is best effort. We return None as we have # no access to the table that we're querying. return None return self._get_table_size(bq, table_ref) elif self.query is not None and self.query.is_accessible(): query_job_name = bigquery_tools.generate_bq_job_name( self._job_name, self._source_uuid, bigquery_tools.BigQueryJobTypes.QUERY, '%s_%s' % (int(time.time()), random.randint(0, 1000))) job = bq._start_query_job( self._get_parent_project(), self.query.get(), self.use_legacy_sql, self.flatten_results, job_id=query_job_name, priority=self.query_priority, dry_run=True, kms_key=self.kms_key, job_labels=self._get_bq_metadata().add_additional_bq_job_labels( self.bigquery_job_labels)) if job.statistics.totalBytesProcessed is None: # Some queries may not have access to `totalBytesProcessed` as a # result of row-level security # > BigQuery hides sensitive statistics on all queries against # > tables with row-level security. # See cloud.google.com/bigquery/docs/managing-row-level-security # and cloud.google.com/bigquery/docs/best-practices-row-level-security return None return int(job.statistics.totalBytesProcessed) else: # Size estimation is best effort. We return None as we have # no access to the query that we're running. return None def split(self, desired_bundle_size, start_position=None, stop_position=None): if self.split_result is None: bq = bigquery_tools.BigQueryWrapper( temp_table_ref=(self.temp_table if self.temp_table else None), client=bigquery_tools.BigQueryWrapper._bigquery_client( self.pipeline_options)) if self.query is not None: self._setup_temporary_dataset(bq) self.table_reference = self._execute_query(bq) if not self.table_reference.projectId: self.table_reference.projectId = self._get_project() requested_session = bq_storage.types.ReadSession() requested_session.table = 'projects/{}/datasets/{}/tables/{}'.format( self.table_reference.projectId, self.table_reference.datasetId, self.table_reference.tableId) Lineage.sources().add( 'bigquery', self.table_reference.projectId, self.table_reference.datasetId, self.table_reference.tableId) if self.use_native_datetime: requested_session.data_format = bq_storage.types.DataFormat.ARROW requested_session.read_options\ .arrow_serialization_options.buffer_compression = \ bq_storage.types.ArrowSerializationOptions.CompressionCodec.LZ4_FRAME else: requested_session.data_format = bq_storage.types.DataFormat.AVRO if self.selected_fields is not None: requested_session.read_options.selected_fields = self.selected_fields if self.row_restriction is not None: requested_session.read_options.row_restriction = self.row_restriction storage_client = bq_storage.BigQueryReadClient() stream_count = 0 if desired_bundle_size > 0: table_size = self._get_table_size(bq, self.table_reference) stream_count = min( int(table_size / desired_bundle_size), _CustomBigQueryStorageSource.MAX_SPLIT_COUNT) stream_count = max( stream_count, _CustomBigQueryStorageSource.MIN_SPLIT_COUNT) parent = 'projects/{}'.format(self.table_reference.projectId) read_session = storage_client.create_read_session( parent=parent, read_session=requested_session, max_stream_count=stream_count) if self.use_native_datetime: display_schema = "Arrow Schema:" + str(read_session.arrow_schema) else: display_schema = "Avro Schema:" + str(read_session.avro_schema) _LOGGER.info( 'Sent BigQuery Storage API CreateReadSession request: \n %s \n' 'Received %d streams\ndata_format: %s\n' 'estimated_total_bytes_scanned: %s\n%s.', requested_session, len(read_session.streams), read_session.data_format, read_session.estimated_total_bytes_scanned, display_schema) self.split_result = [ _CustomBigQueryStorageStreamSource( stream.name, self.use_native_datetime, self.timeout) for stream in read_session.streams ] for source in self.split_result: yield SourceBundle( weight=1.0, source=source, start_position=None, stop_position=None) def get_range_tracker(self, start_position, stop_position): class NonePositionRangeTracker(RangeTracker): """A RangeTracker that always returns positions as None. Prevents the BigQuery Storage source from being read() before being split().""" def start_position(self): return None def stop_position(self): return None return NonePositionRangeTracker() def read(self, range_tracker): raise NotImplementedError( 'BigQuery storage source must be split before being read') class _CustomBigQueryStorageStreamSource(BoundedSource): """A source representing a single stream in a read session.""" # Runner will act on this counter on scaling event, if supported THROTTLE_COUNTER = Metrics.counter(__name__, 'cumulativeThrottlingSeconds') def __init__( self, read_stream_name: str, use_native_datetime: Optional[bool] = True, timeout: Optional[float] = None): self.read_stream_name = read_stream_name self.use_native_datetime = use_native_datetime self.timeout = timeout def display_data(self): return { 'output_format': 'ARROW' if self.use_native_datetime else 'AVRO', 'read_stream': str(self.read_stream_name), 'use_native_datetime': str(self.use_native_datetime) } def estimate_size(self): # The size of stream source cannot be estimate due to server-side liquid # sharding. # TODO(https://github.com/apache/beam/issues/21126): Implement progress # reporting. return None def split(self, desired_bundle_size, start_position=None, stop_position=None): # A stream source can't be split without reading from it due to # server-side liquid sharding. A split will simply return the current source # for now. return SourceBundle( weight=1.0, source=_CustomBigQueryStorageStreamSource( self.read_stream_name, self.use_native_datetime), start_position=None, stop_position=None) def get_range_tracker(self, start_position, stop_position): # TODO(https://github.com/apache/beam/issues/21127): Implement dynamic work # rebalancing. assert start_position is None # Defaulting to the start of the stream. start_position = 0 # Since the streams are unsplittable we choose OFFSET_INFINITY as the # default end offset so that all data of the source gets read. stop_position = range_trackers.OffsetRangeTracker.OFFSET_INFINITY range_tracker = range_trackers.OffsetRangeTracker( start_position, stop_position) # Ensuring that all try_split() calls will be ignored by the Rangetracker. range_tracker = range_trackers.UnsplittableRangeTracker(range_tracker) return range_tracker def read(self, range_tracker): _LOGGER.info( "Started BigQuery Storage API read from stream %s.", self.read_stream_name) if self.use_native_datetime: return self.read_arrow() else: return self.read_avro() @staticmethod def retry_delay_callback(delay): _LOGGER.info("retry delay: %f", delay) _CustomBigQueryStorageStreamSource.THROTTLE_COUNTER.inc(delay) def read_arrow(self): storage_client = bq_storage.BigQueryReadClient() read_rows_kwargs = {'retry_delay_callback': self.retry_delay_callback} if self.timeout is not None: read_rows_kwargs['timeout'] = self.timeout row_iter = iter( storage_client.read_rows(self.read_stream_name, **read_rows_kwargs).rows()) row = next(row_iter, None) # Handling the case where the user might provide very selective filters # which can result in read_rows_response being empty. if row is None: return iter([]) while row is not None: py_row = dict(map(lambda item: (item[0], item[1].as_py()), row.items())) row = next(row_iter, None) yield py_row def read_avro(self): storage_client = bq_storage.BigQueryReadClient() read_rows_kwargs = {'retry_delay_callback': self.retry_delay_callback} if self.timeout is not None: read_rows_kwargs['timeout'] = self.timeout read_rows_iterator = iter( storage_client.read_rows(self.read_stream_name, **read_rows_kwargs)) # Handling the case where the user might provide very selective filters # which can result in read_rows_response being empty. first_read_rows_response = next(read_rows_iterator, None) if first_read_rows_response is None: return iter([]) row_reader = _ReadReadRowsResponsesWithFastAvro( read_rows_iterator, first_read_rows_response) return iter(row_reader) class _ReadReadRowsResponsesWithFastAvro(): """An iterator that deserializes ReadRowsResponses using the fastavro library.""" def __init__(self, read_rows_iterator, read_rows_response): self.read_rows_iterator = read_rows_iterator self.read_rows_response = read_rows_response self.avro_schema = fastavro.parse_schema( json.loads(self.read_rows_response.avro_schema.schema)) self.bytes_reader = io.BytesIO( self.read_rows_response.avro_rows.serialized_binary_rows) def __iter__(self): return self def __next__(self): try: return fastavro.schemaless_reader(self.bytes_reader, self.avro_schema) except (StopIteration, EOFError): self.read_rows_response = next(self.read_rows_iterator, None) if self.read_rows_response is not None: self.bytes_reader = io.BytesIO( self.read_rows_response.avro_rows.serialized_binary_rows) return fastavro.schemaless_reader(self.bytes_reader, self.avro_schema) else: raise StopIteration @deprecated(since='2.11.0', current="WriteToBigQuery") def BigQuerySink(*args, validate=False, **kwargs): """A deprecated alias for WriteToBigQuery.""" warnings.warn( "Native sinks no longer implemented; " "falling back to standard Beam sink.") return WriteToBigQuery(*args, validate=validate, **kwargs) _KNOWN_TABLES = set() class BigQueryWriteFn(DoFn): """A ``DoFn`` that streams writes to BigQuery once the table is created.""" DEFAULT_MAX_BUFFERED_ROWS = 2000 DEFAULT_MAX_BATCH_SIZE = 500 FAILED_ROWS = 'FailedRows' FAILED_ROWS_WITH_ERRORS = 'FailedRowsWithErrors' STREAMING_API_LOGGING_FREQUENCY_SEC = 300 def __init__( self, batch_size, schema=None, create_disposition=None, write_disposition=None, kms_key=None, test_client=None, max_buffered_rows=None, retry_strategy=None, additional_bq_parameters=None, ignore_insert_ids=False, with_batched_input=False, ignore_unknown_columns=False, max_retries=MAX_INSERT_RETRIES, max_insert_payload_size=MAX_INSERT_PAYLOAD_SIZE): """Initialize a WriteToBigQuery transform. Args: batch_size: Number of rows to be written to BQ per streaming API insert. schema: The schema to be used if the BigQuery table to write has to be created. This can be either specified as a 'bigquery.TableSchema' object or a single string of the form 'field1:type1,field2:type2,field3:type3' that defines a comma separated list of fields. Here 'type' should specify the BigQuery type of the field. Single string based schemas do not support nested fields, repeated fields, or specifying a BigQuery mode for fields (mode will always be set to 'NULLABLE'). create_disposition: A string describing what happens if the table does not exist. Possible values are: - BigQueryDisposition.CREATE_IF_NEEDED: create if does not exist. - BigQueryDisposition.CREATE_NEVER: fail the write if does not exist. write_disposition: A string describing what happens if the table has already some data. Possible values are: - BigQueryDisposition.WRITE_TRUNCATE: delete existing rows. - BigQueryDisposition.WRITE_APPEND: add to existing rows. - BigQueryDisposition.WRITE_EMPTY: fail the write if table not empty. For streaming pipelines WriteTruncate can not be used. kms_key: Optional Cloud KMS key name for use when creating new tables. test_client: Override the default bigquery client used for testing. max_buffered_rows: The maximum number of rows that are allowed to stay buffered when running dynamic destinations. When destinations are dynamic, it is important to keep caches small even when a single batch has not been completely filled up. retry_strategy: The strategy to use when retrying streaming inserts into BigQuery. Options are shown in bigquery_tools.RetryStrategy attrs. additional_bq_parameters (dict, callable): A set of additional parameters to be passed when creating a BigQuery table. These are passed when triggering a load job for FILE_LOADS, and when creating a new table for STREAMING_INSERTS. ignore_insert_ids: When using the STREAMING_INSERTS method to write data to BigQuery, `insert_ids` are a feature of BigQuery that support deduplication of events. If your use case is not sensitive to duplication of data inserted to BigQuery, set `ignore_insert_ids` to True to increase the throughput for BQ writing. See: https://cloud.google.com/bigquery/streaming-data-into-bigquery#disabling_best_effort_de-duplication with_batched_input: Whether the input has already been batched per destination. If not, perform best-effort batching per destination within a bundle. ignore_unknown_columns: Accept rows that contain values that do not match the schema. The unknown values are ignored. Default is False, which treats unknown values as errors. See reference: https://cloud.google.com/bigquery/docs/reference/rest/v2/tabledata/insertAll max_retries: The number of times that we will retry inserting a group of rows into BigQuery. By default, we retry 10000 times with exponential backoffs (effectively retry forever). max_insert_payload_size: The maximum byte size for a BigQuery legacy streaming insert payload. """ self.schema = schema self.test_client = test_client self.create_disposition = create_disposition self.write_disposition = write_disposition if write_disposition in (BigQueryDisposition.WRITE_EMPTY, BigQueryDisposition.WRITE_TRUNCATE): raise ValueError( 'Write disposition %s is not supported for' ' streaming inserts to BigQuery' % write_disposition) self._rows_buffer = [] self._reset_rows_buffer() self._total_buffered_rows = 0 self.kms_key = kms_key self._max_batch_size = batch_size or BigQueryWriteFn.DEFAULT_MAX_BATCH_SIZE self._max_buffered_rows = ( max_buffered_rows or BigQueryWriteFn.DEFAULT_MAX_BUFFERED_ROWS) self._retry_strategy = retry_strategy or RetryStrategy.RETRY_ALWAYS self.ignore_insert_ids = ignore_insert_ids self.with_batched_input = with_batched_input self.additional_bq_parameters = additional_bq_parameters or {} # accumulate the total time spent in exponential backoff self._throttled_secs = Metrics.counter( BigQueryWriteFn, "cumulativeThrottlingSeconds") self.batch_size_metric = Metrics.distribution(self.__class__, "batch_size") self.batch_latency_metric = Metrics.distribution( self.__class__, "batch_latency_ms") self.failed_rows_metric = Metrics.distribution( self.__class__, "rows_failed_per_batch") self.bigquery_wrapper = None self.streaming_api_logging_frequency_sec = ( BigQueryWriteFn.STREAMING_API_LOGGING_FREQUENCY_SEC) self.ignore_unknown_columns = ignore_unknown_columns self._max_retries = max_retries self._max_insert_payload_size = max_insert_payload_size def display_data(self): return { 'max_batch_size': self._max_batch_size, 'max_buffered_rows': self._max_buffered_rows, 'retry_strategy': self._retry_strategy, 'create_disposition': str(self.create_disposition), 'write_disposition': str(self.write_disposition), 'additional_bq_parameters': str(self.additional_bq_parameters), 'ignore_insert_ids': str(self.ignore_insert_ids), 'ignore_unknown_columns': str(self.ignore_unknown_columns) } def _reset_rows_buffer(self): self._rows_buffer = collections.defaultdict(lambda: []) self._destination_buffer_byte_size = collections.defaultdict(lambda: 0) @staticmethod def get_table_schema(schema): """Transform the table schema into a bigquery.TableSchema instance. Args: schema: The schema to be used if the BigQuery table to write has to be created. This is a dictionary object created in the WriteToBigQuery transform. Returns: table_schema: The schema to be used if the BigQuery table to write has to be created but in the bigquery.TableSchema format. """ if schema is None: return schema elif isinstance(schema, str): return bigquery_tools.parse_table_schema_from_json(schema) elif isinstance(schema, dict): return bigquery_tools.parse_table_schema_from_json(json.dumps(schema)) else: raise TypeError('Unexpected schema argument: %s.' % schema) def start_bundle(self): self._reset_rows_buffer() if not self.bigquery_wrapper: self.bigquery_wrapper = bigquery_tools.BigQueryWrapper( client=self.test_client) ( bigquery_tools.BigQueryWrapper.HISTOGRAM_METRIC_LOGGER. minimum_logging_frequency_msec ) = self.streaming_api_logging_frequency_sec * 1000 self._backoff_calculator = iter( retry.FuzzedExponentialIntervals( initial_delay_secs=0.2, num_retries=self._max_retries, max_delay_secs=1500)) def _create_table_if_needed(self, table_reference, schema=None): str_table_reference = '%s:%s.%s' % ( table_reference.projectId, table_reference.datasetId, table_reference.tableId) if str_table_reference in _KNOWN_TABLES: return if self.create_disposition == BigQueryDisposition.CREATE_NEVER: # If we never want to create the table, we assume it already exists, # and avoid the get-or-create step. return _LOGGER.debug( 'Creating or getting table %s with schema %s.', table_reference, schema) table_schema = self.get_table_schema(schema) if table_reference.projectId is None: table_reference.projectId = vp.RuntimeValueProvider.get_value( 'project', str, '') self.bigquery_wrapper.get_or_create_table( table_reference.projectId, table_reference.datasetId, table_reference.tableId, table_schema, self.create_disposition, self.write_disposition, additional_create_parameters=self.additional_bq_parameters) _KNOWN_TABLES.add(str_table_reference) def _check_row_size(self, row_and_insert_id) -> tuple[int, Optional[str]]: """Returns error string when the row estimated size is too big""" row_byte_size = get_deep_size(row_and_insert_id) # Check if individual row exceeds size limit if row_byte_size >= self._max_insert_payload_size: row_mb_size = row_byte_size / 1_000_000 max_mb_size = self._max_insert_payload_size / 1_000_000 return ( row_byte_size, ( f"Received row with size {row_mb_size}MB that exceeds " f"the maximum insert payload size set ({max_mb_size}MB).")) return (row_byte_size, None) def process( self, element, window_value=DoFn.WindowedValueParam, *schema_side_inputs): destination = bigquery_tools.get_hashable_destination(element[0]) if callable(self.schema): schema = self.schema(destination, *schema_side_inputs) elif isinstance(self.schema, vp.ValueProvider): schema = self.schema.get() else: schema = self.schema self._create_table_if_needed( bigquery_tools.parse_table_reference(destination), schema) if not self.with_batched_input: row_and_insert_id = element[1] row_byte_size, row_too_big_error = self._check_row_size(row_and_insert_id) # send large rows that exceed BigQuery insert limits to DLQ if row_too_big_error is not None: return [ pvalue.TaggedOutput( BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS, window_value.with_value( (destination, row_and_insert_id[0], row_too_big_error))), pvalue.TaggedOutput( BigQueryWriteFn.FAILED_ROWS, window_value.with_value((destination, row_and_insert_id[0]))) ] # Flush current batch first if adding this row will exceed our limits # limits: byte size; number of rows if ((self._destination_buffer_byte_size[destination] + row_byte_size > self._max_insert_payload_size) or len(self._rows_buffer[destination]) >= self._max_batch_size): flushed_batch = self._flush_batch(destination) # After flushing our existing batch, we now buffer the current row # for the next flush self._rows_buffer[destination].append((row_and_insert_id, window_value)) self._destination_buffer_byte_size[destination] = row_byte_size return flushed_batch self._rows_buffer[destination].append((row_and_insert_id, window_value)) self._destination_buffer_byte_size[destination] += row_byte_size self._total_buffered_rows += 1 if self._total_buffered_rows >= self._max_buffered_rows: return self._flush_all_batches() else: # The input is already batched per destination # but we verify the payload size. # The batch might be split into smaller batches batched_rows = element[1] failed_outputs = [] current_batch = [] current_batch_size = 0 for row in batched_rows: row_byte_size, row_too_big_error = self._check_row_size(row) # Check if individual row exceeds size limit if row_too_big_error is not None: failed_outputs.extend([ pvalue.TaggedOutput( BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS, window_value.with_value( (destination, row[0], row_too_big_error))), pvalue.TaggedOutput( BigQueryWriteFn.FAILED_ROWS, window_value.with_value((destination, row[0]))) ]) continue # Check if adding this row would exceed batch size limit if (len(current_batch) != 0 and current_batch_size + row_byte_size > self._max_insert_payload_size): self._rows_buffer[destination].extend( ((batch_row, window_value) for batch_row in current_batch)) failed_outputs.extend(self._flush_batch(destination)) # Start new batch with current row current_batch = [row] current_batch_size = row_byte_size else: current_batch.append(row) current_batch_size += row_byte_size if current_batch: for batch_row in current_batch: self._rows_buffer[destination].append((batch_row, window_value)) failed_outputs.extend(self._flush_batch(destination)) return failed_outputs def finish_bundle(self): bigquery_tools.BigQueryWrapper.HISTOGRAM_METRIC_LOGGER.log_metrics( reset_after_logging=True) return self._flush_all_batches() def _flush_all_batches(self): _LOGGER.debug( 'Attempting to flush to all destinations. Total buffered: %s', self._total_buffered_rows) return itertools.chain( *[ self._flush_batch(destination) for destination in list(self._rows_buffer.keys()) if self._rows_buffer[destination] ]) def _flush_batch(self, destination): # Flush the current batch of rows to BigQuery. rows_and_insert_ids_with_windows = self._rows_buffer[destination] table_reference = bigquery_tools.parse_table_reference(destination) if table_reference.projectId is None: table_reference.projectId = vp.RuntimeValueProvider.get_value( 'project', str, '') _LOGGER.debug( 'Flushing data to %s. Total %s rows.', destination, len(rows_and_insert_ids_with_windows)) self.batch_size_metric.update(len(rows_and_insert_ids_with_windows)) rows_and_insert_ids, window_values = zip(*rows_and_insert_ids_with_windows) rows = [r[0] for r in rows_and_insert_ids] if self.ignore_insert_ids: insert_ids = [None for r in rows_and_insert_ids] else: insert_ids = [r[1] for r in rows_and_insert_ids] while True: start = time.time() passed, errors = self.bigquery_wrapper.insert_rows( project_id=table_reference.projectId, dataset_id=table_reference.datasetId, table_id=table_reference.tableId, rows=rows, insert_ids=insert_ids, skip_invalid_rows=True, ignore_unknown_values=self.ignore_unknown_columns) self.batch_latency_metric.update((time.time() - start) * 1000) failed_rows = [( rows[entry['index']], entry["errors"], window_values[entry['index']]) for entry in errors] failed_insert_ids = [insert_ids[entry['index']] for entry in errors] retry_backoff = next(self._backoff_calculator, None) # If retry_backoff is None, then we will not retry and must log. should_retry = any( RetryStrategy.should_retry( self._retry_strategy, entry['errors'][0]['reason']) for entry in errors) and retry_backoff is not None if not passed: self.failed_rows_metric.update(len(failed_rows)) message = ( 'There were errors inserting to BigQuery. Will{} retry. ' 'Errors were {}'.format(("" if should_retry else " not"), errors)) # The log level is: # - WARNING when should_retry is true, else ERROR. if (should_retry and self._retry_strategy in [RetryStrategy.RETRY_ON_TRANSIENT_ERROR, RetryStrategy.RETRY_ALWAYS]): log_level = logging.WARN else: log_level = logging.ERROR _LOGGER.log(log_level, message) if not should_retry: break else: _LOGGER.info( 'Sleeping %s seconds before retrying insertion.', retry_backoff) time.sleep(retry_backoff) # We can now safely discard all information about successful rows and # just focus on the failed ones rows = [fr[0] for fr in failed_rows] window_values = [fr[2] for fr in failed_rows] insert_ids = failed_insert_ids self._throttled_secs.inc(retry_backoff) self._total_buffered_rows -= len(self._rows_buffer[destination]) del self._rows_buffer[destination] if destination in self._destination_buffer_byte_size: del self._destination_buffer_byte_size[destination] return itertools.chain( [ pvalue.TaggedOutput( BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS, w.with_value((destination, row, err))) for row, err, w in failed_rows ], [ pvalue.TaggedOutput( BigQueryWriteFn.FAILED_ROWS, w.with_value((destination, row))) for row, unused_err, w in failed_rows ]) # The number of shards per destination when writing via streaming inserts. DEFAULT_SHARDS_PER_DESTINATION = 500 # The max duration a batch of elements is allowed to be buffered before being # flushed to BigQuery. DEFAULT_BATCH_BUFFERING_DURATION_LIMIT_SEC = 0.2 class _StreamToBigQuery(PTransform): def __init__( self, table_reference, table_side_inputs, schema_side_inputs, schema, batch_size, triggering_frequency, create_disposition, write_disposition, kms_key, retry_strategy, additional_bq_parameters, ignore_insert_ids, ignore_unknown_columns, with_auto_sharding, num_streaming_keys=DEFAULT_SHARDS_PER_DESTINATION, test_client=None, max_retries=MAX_INSERT_RETRIES, max_insert_payload_size=MAX_INSERT_PAYLOAD_SIZE): self.table_reference = table_reference self.table_side_inputs = table_side_inputs self.schema_side_inputs = schema_side_inputs self.schema = schema self.batch_size = batch_size self.triggering_frequency = triggering_frequency self.create_disposition = create_disposition self.write_disposition = write_disposition self.kms_key = kms_key self.retry_strategy = retry_strategy self.test_client = test_client self.additional_bq_parameters = additional_bq_parameters self.ignore_insert_ids = ignore_insert_ids self.ignore_unknown_columns = ignore_unknown_columns self.with_auto_sharding = with_auto_sharding self._num_streaming_keys = num_streaming_keys self._max_retries = max_retries self._max_insert_payload_size = max_insert_payload_size class InsertIdPrefixFn(DoFn): def start_bundle(self): self.prefix = str(uuid.uuid4()) self._row_count = 0 def process(self, element): key = element[0] value = element[1] insert_id = '%s-%s' % (self.prefix, self._row_count) self._row_count += 1 yield (key, (value, insert_id)) def expand(self, input): bigquery_write_fn = BigQueryWriteFn( schema=self.schema, batch_size=self.batch_size, create_disposition=self.create_disposition, write_disposition=self.write_disposition, kms_key=self.kms_key, retry_strategy=self.retry_strategy, test_client=self.test_client, additional_bq_parameters=self.additional_bq_parameters, ignore_insert_ids=self.ignore_insert_ids, ignore_unknown_columns=self.ignore_unknown_columns, with_batched_input=self.with_auto_sharding, max_retries=self._max_retries, max_insert_payload_size=self._max_insert_payload_size) def _add_random_shard(element): key = element[0] value = element[1] return ((key, random.randint(0, self._num_streaming_keys)), value) def _restore_table_ref(sharded_table_ref_elems_kv): sharded_table_ref = sharded_table_ref_elems_kv[0] table_ref = bigquery_tools.parse_table_reference(sharded_table_ref) return (table_ref, sharded_table_ref_elems_kv[1]) tagged_data = ( input | 'AppendDestination' >> beam.ParDo( bigquery_tools.AppendDestinationsFn(self.table_reference), *self.table_side_inputs) | 'AddInsertIds' >> beam.ParDo(_StreamToBigQuery.InsertIdPrefixFn()) | 'ToHashableTableRef' >> beam.Map(bigquery_tools.to_hashable_table_ref)) if not self.with_auto_sharding: tagged_data = ( tagged_data | 'WithFixedSharding' >> beam.Map(_add_random_shard) | 'CommitInsertIds' >> ReshufflePerKey() | 'DropShard' >> beam.Map(lambda kv: (kv[0][0], kv[1]))) else: # Auto-sharding is achieved via GroupIntoBatches.WithShardedKey # transform which shards, groups and at the same time batches the table # rows to be inserted to BigQuery. # Firstly the keys of tagged_data (table references) are converted to a # hashable format. This is needed to work with the keyed states used by # GroupIntoBatches. After grouping and batching is done, original table # references are restored. tagged_data = ( tagged_data | 'WithAutoSharding' >> beam.GroupIntoBatches.WithShardedKey( (self.batch_size or BigQueryWriteFn.DEFAULT_MAX_BUFFERED_ROWS), self.triggering_frequency or DEFAULT_BATCH_BUFFERING_DURATION_LIMIT_SEC) | 'DropShard' >> beam.Map(lambda kv: (kv[0].key, kv[1]))) return ( tagged_data | 'FromHashableTableRef' >> beam.Map(_restore_table_ref) | 'StreamInsertRows' >> ParDo( bigquery_write_fn, *self.schema_side_inputs).with_outputs( BigQueryWriteFn.FAILED_ROWS, BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS, main='main')) # Flag to be passed to WriteToBigQuery to force schema autodetection SCHEMA_AUTODETECT = 'SCHEMA_AUTODETECT' class WriteToBigQuery(PTransform): """Write data to BigQuery. This transform receives a PCollection of elements to be inserted into BigQuery tables. The elements would come in as Python dictionaries, or as `TableRow` instances. """ class Method(object): DEFAULT = 'DEFAULT' STREAMING_INSERTS = 'STREAMING_INSERTS' FILE_LOADS = 'FILE_LOADS' STORAGE_WRITE_API = 'STORAGE_WRITE_API' def __init__( self, table, dataset=None, project=None, schema=None, create_disposition=BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=BigQueryDisposition.WRITE_APPEND, kms_key=None, batch_size=None, max_file_size=None, max_partition_size=None, max_files_per_bundle=None, test_client=None, custom_gcs_temp_location=None, method=None, insert_retry_strategy=None, additional_bq_parameters=None, table_side_inputs=None, schema_side_inputs=None, triggering_frequency=None, use_at_least_once=False, validate=True, temp_file_format=None, ignore_insert_ids=False, # TODO(https://github.com/apache/beam/issues/20712): Switch the default # when the feature is mature. with_auto_sharding=False, num_storage_api_streams=0, ignore_unknown_columns=False, load_job_project_id=None, max_retries=MAX_INSERT_RETRIES, max_insert_payload_size=MAX_INSERT_PAYLOAD_SIZE, num_streaming_keys=DEFAULT_SHARDS_PER_DESTINATION, use_cdc_writes: bool = False, primary_key: list[str] = None, expansion_service=None, big_lake_configuration=None): """Initialize a WriteToBigQuery transform. Args: table (str, callable, ValueProvider): The ID of the table, or a callable that returns it. The ID must contain only letters ``a-z``, ``A-Z``, numbers ``0-9``, or connectors ``-_``. If dataset argument is :data:`None` then the table argument must contain the entire table reference specified as: ``'DATASET.TABLE'`` or ``'PROJECT:DATASET.TABLE'``. If it's a callable, it must receive one argument representing an element to be written to BigQuery, and return a TableReference, or a string table name as specified above. dataset (str): The ID of the dataset containing this table or :data:`None` if the table reference is specified entirely by the table argument. project (str): The ID of the project containing this table or :data:`None` if the table reference is specified entirely by the table argument. schema (str,dict,ValueProvider,callable): The schema to be used if the BigQuery table to write has to be created. This can be either specified as a :class:`~apache_beam.io.gcp.internal.clients.bigquery.\ bigquery_v2_messages.TableSchema`. or a `ValueProvider` that has a JSON string, or a python dictionary, or the string or dictionary itself, object or a single string of the form ``'field1:type1,field2:type2,field3:type3'`` that defines a comma separated list of fields. Here ``'type'`` should specify the BigQuery type of the field. Single string based schemas do not support nested fields, repeated fields, or specifying a BigQuery mode for fields (mode will always be set to ``'NULLABLE'``). If a callable, then it should receive a destination (in the form of a str, and return a str, dict or TableSchema). One may also pass ``SCHEMA_AUTODETECT`` here when using JSON-based file loads, and BigQuery will try to infer the schema for the files that are being loaded. create_disposition (BigQueryDisposition): A string describing what happens if the table does not exist. Possible values are: * :attr:`BigQueryDisposition.CREATE_IF_NEEDED`: create if does not exist. * :attr:`BigQueryDisposition.CREATE_NEVER`: fail the write if does not exist. write_disposition (BigQueryDisposition): A string describing what happens if the table has already some data. Possible values are: * :attr:`BigQueryDisposition.WRITE_TRUNCATE`: delete existing rows. * :attr:`BigQueryDisposition.WRITE_APPEND`: add to existing rows. * :attr:`BigQueryDisposition.WRITE_EMPTY`: fail the write if table not empty. For streaming pipelines WriteTruncate can not be used. kms_key (str): Optional Cloud KMS key name for use when creating new tables. batch_size (int): Number of rows to be written to BQ per streaming API insert. The default is 500. test_client: Override the default bigquery client used for testing. max_file_size (int): The maximum size for a file to be written and then loaded into BigQuery. The default value is 4TB, which is 80% of the limit of 5TB for BigQuery to load any file. max_partition_size (int): Maximum byte size for each load job to BigQuery. Defaults to 15TB. Applicable to FILE_LOADS only. max_files_per_bundle(int): The maximum number of files to be concurrently written by a worker. The default here is 20. Larger values will allow writing to multiple destinations without having to reshard - but they increase the memory burden on the workers. custom_gcs_temp_location (str): A GCS location to store files to be used for file loads into BigQuery. By default, this will use the pipeline's temp_location, but for pipelines whose temp_location is not appropriate for BQ File Loads, users should pass a specific one. method: The method to use to write to BigQuery. It may be STREAMING_INSERTS, FILE_LOADS, STORAGE_WRITE_API or DEFAULT. An introduction on loading data to BigQuery: https://cloud.google.com/bigquery/docs/loading-data. DEFAULT will use STREAMING_INSERTS on Streaming pipelines and FILE_LOADS on Batch pipelines. Note: FILE_LOADS currently does not support BigQuery's JSON data type: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#json_type"> insert_retry_strategy: The strategy to use when retrying streaming inserts into BigQuery. Options are shown in bigquery_tools.RetryStrategy attrs. Default is to retry always. This means that whenever there are rows that fail to be inserted to BigQuery, they will be retried indefinitely. Other retry strategy settings will produce a deadletter PCollection as output. Appropriate values are: * `RetryStrategy.RETRY_ALWAYS`: retry all rows if there are any kind of errors. Note that this will hold your pipeline back if there are errors until you cancel or update it. * `RetryStrategy.RETRY_NEVER`: rows with errors will not be retried. Instead they will be output to a dead letter queue under the `'FailedRows'` tag. * `RetryStrategy.RETRY_ON_TRANSIENT_ERROR`: retry rows with transient errors (e.g. timeouts). Rows with permanent errors will be output to dead letter queue under `'FailedRows'` tag. additional_bq_parameters (dict, callable): Additional parameters to pass to BQ when creating / loading data into a table. If a callable, it should be a function that receives a table reference indicating the destination and returns a dictionary. These can be 'timePartitioning', 'clustering', etc. They are passed directly to the job load configuration. See https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#jobconfigurationload table_side_inputs (tuple): A tuple with ``AsSideInput`` PCollections to be passed to the table callable (if one is provided). schema_side_inputs: A tuple with ``AsSideInput`` PCollections to be passed to the schema callable (if one is provided). triggering_frequency (float): When method is FILE_LOADS: Value will be converted to int. Every triggering_frequency seconds, a BigQuery load job will be triggered for all the data written since the last load job. BigQuery has limits on how many load jobs can be triggered per day, so be careful not to set this duration too low, or you may exceed daily quota. Often this is set to 5 or 10 minutes to ensure that the project stays well under the BigQuery quota. See https://cloud.google.com/bigquery/quota-policy for more information about BigQuery quotas. When method is STREAMING_INSERTS and with_auto_sharding=True: A streaming inserts batch will be submitted at least every triggering_frequency seconds when data is waiting. The batch can be sent earlier if it reaches the maximum batch size set by batch_size. Default value is 0.2 seconds. When method is STORAGE_WRITE_API: A stream of rows will be committed every triggering_frequency seconds. By default, this will be 5 seconds to ensure exactly-once semantics. use_at_least_once: Intended only for STORAGE_WRITE_API. When True, will use at-least-once semantics. This is cheaper and provides lower latency, but will potentially duplicate records. validate: Indicates whether to perform validation checks on inputs. This parameter is primarily used for testing. temp_file_format: The format to use for file loads into BigQuery. The options are NEWLINE_DELIMITED_JSON or AVRO, with NEWLINE_DELIMITED_JSON being used by default. For advantages and limitations of the two formats, see https://cloud.google.com/bigquery/docs/loading-data-cloud-storage-avro and https://cloud.google.com/bigquery/docs/loading-data-cloud-storage-json. ignore_insert_ids: When using the STREAMING_INSERTS method to write data to BigQuery, `insert_ids` are a feature of BigQuery that support deduplication of events. If your use case is not sensitive to duplication of data inserted to BigQuery, set `ignore_insert_ids` to True to increase the throughput for BQ writing. See: https://cloud.google.com/bigquery/streaming-data-into-bigquery#disabling_best_effort_de-duplication with_auto_sharding: Experimental. If true, enables using a dynamically determined number of shards to write to BigQuery. This can be used for all of FILE_LOADS, STREAMING_INSERTS, and STORAGE_WRITE_API. Only applicable to unbounded input. num_storage_api_streams: Specifies the number of write streams that the Storage API sink will use. This parameter is only applicable when writing unbounded data. ignore_unknown_columns: Accept rows that contain values that do not match the schema. The unknown values are ignored. Default is False, which treats unknown values as errors. This option is only valid for method=STREAMING_INSERTS. See reference: https://cloud.google.com/bigquery/docs/reference/rest/v2/tabledata/insertAll load_job_project_id: Specifies an alternate GCP project id to use for billingBatch File Loads. By default, the project id of the table is used. num_streaming_keys: The number of shards per destination when writing via streaming inserts. expansion_service: The address (host:port) of the expansion service. If no expansion service is provided, will attempt to run the default GCP expansion service. Used for STORAGE_WRITE_API method. max_retries: The number of times that we will retry inserting a group of rows into BigQuery. By default, we retry 10000 times with exponential backoffs (effectively retry forever). max_insert_payload_size: The maximum byte size for a BigQuery legacy streaming insert payload. use_cdc_writes: Configure the usage of CDC writes on BigQuery. The argument can be used by passing True and the Beam Rows will be sent as they are to the BigQuery sink which expects a 'record' and 'row_mutation_info' properties. Used for STORAGE_WRITE_API, working on 'at least once' mode. primary_key: When using CDC write on BigQuery and CREATE_IF_NEEDED mode for the underlying tables a list of column names is required to be configured as the primary key. Used for STORAGE_WRITE_API, working on 'at least once' mode. """ self._table = table self._dataset = dataset self._project = project self.table_reference = bigquery_tools.parse_table_reference( table, dataset, project) self.create_disposition = BigQueryDisposition.validate_create( create_disposition) self.write_disposition = BigQueryDisposition.validate_write( write_disposition) if schema == SCHEMA_AUTODETECT: self.schema = schema else: self.schema = bigquery_tools.get_dict_table_schema(schema) self.batch_size = batch_size self.kms_key = kms_key self.test_client = test_client # TODO(pabloem): Consider handling ValueProvider for this location. self.custom_gcs_temp_location = custom_gcs_temp_location self.max_file_size = max_file_size self.max_partition_size = max_partition_size self.max_files_per_bundle = max_files_per_bundle self.method = method or WriteToBigQuery.Method.DEFAULT self.triggering_frequency = triggering_frequency self.use_at_least_once = use_at_least_once self.expansion_service = expansion_service self.with_auto_sharding = with_auto_sharding self._num_storage_api_streams = num_storage_api_streams self.insert_retry_strategy = insert_retry_strategy self._validate = validate self._temp_file_format = temp_file_format or bigquery_tools.FileFormat.JSON self.additional_bq_parameters = additional_bq_parameters or {} self.table_side_inputs = table_side_inputs or () self.schema_side_inputs = schema_side_inputs or () self._ignore_insert_ids = ignore_insert_ids self._ignore_unknown_columns = ignore_unknown_columns self.load_job_project_id = load_job_project_id self._max_retries = max_retries self._max_insert_payload_size = max_insert_payload_size self._num_streaming_keys = num_streaming_keys self._use_cdc_writes = use_cdc_writes self._primary_key = primary_key self._big_lake_configuration = big_lake_configuration # Dict/schema methods were moved to bigquery_tools, but keep references # here for backward compatibility. get_table_schema_from_string = \ staticmethod(bigquery_tools.get_table_schema_from_string) table_schema_to_dict = staticmethod(bigquery_tools.table_schema_to_dict) get_dict_table_schema = staticmethod(bigquery_tools.get_dict_table_schema) def _compute_method(self, experiments, is_streaming_pipeline): # If the new BQ sink is not activated for experiment flags, then we use # streaming inserts by default (it gets overridden in dataflow_runner.py). if self.method == self.Method.DEFAULT and is_streaming_pipeline: return self.Method.STREAMING_INSERTS elif self.method == self.Method.DEFAULT and not is_streaming_pipeline: return self.Method.FILE_LOADS else: return self.method def expand(self, pcoll): p = pcoll.pipeline if (isinstance(self.table_reference, TableReference) and self.table_reference.projectId is None): self.table_reference.projectId = pcoll.pipeline.options.view_as( GoogleCloudOptions).project # TODO(pabloem): Use a different method to determine if streaming or batch. is_streaming_pipeline = p.options.view_as(StandardOptions).streaming if not is_streaming_pipeline and self.with_auto_sharding: raise ValueError( 'with_auto_sharding is not applicable to batch pipelines.') experiments = p.options.view_as(DebugOptions).experiments or [] method_to_use = self._compute_method(experiments, is_streaming_pipeline) if method_to_use == WriteToBigQuery.Method.STREAMING_INSERTS: if self.schema == SCHEMA_AUTODETECT: raise ValueError( 'Schema auto-detection is not supported for streaming ' 'inserts into BigQuery. Only for File Loads.') if self.triggering_frequency is not None and not self.with_auto_sharding: raise ValueError( 'triggering_frequency with STREAMING_INSERTS can only be used with ' 'with_auto_sharding=True.') if self._max_insert_payload_size > MAX_INSERT_PAYLOAD_SIZE: raise ValueError( 'max_insert_payload_size can only go up to ' f'{MAX_INSERT_PAYLOAD_SIZE} bytes, as per BigQuery quota limits: ' 'https://cloud.google.com/bigquery/quotas#streaming_inserts.') if self._max_retries > MAX_INSERT_RETRIES: raise ValueError( 'max_retries cannot be more than ' f'{MAX_INSERT_RETRIES}, hence please reduce the value.') outputs = pcoll | _StreamToBigQuery( table_reference=self.table_reference, table_side_inputs=self.table_side_inputs, schema_side_inputs=self.schema_side_inputs, schema=self.schema, batch_size=self.batch_size, triggering_frequency=self.triggering_frequency, create_disposition=self.create_disposition, write_disposition=self.write_disposition, kms_key=self.kms_key, retry_strategy=self.insert_retry_strategy, additional_bq_parameters=self.additional_bq_parameters, ignore_insert_ids=self._ignore_insert_ids, ignore_unknown_columns=self._ignore_unknown_columns, with_auto_sharding=self.with_auto_sharding, test_client=self.test_client, max_insert_payload_size=self._max_insert_payload_size, max_retries=self._max_retries, num_streaming_keys=self._num_streaming_keys) return WriteResult( method=WriteToBigQuery.Method.STREAMING_INSERTS, failed_rows=outputs[BigQueryWriteFn.FAILED_ROWS], failed_rows_with_errors=outputs[ BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS]) elif method_to_use == WriteToBigQuery.Method.FILE_LOADS: if self._temp_file_format == bigquery_tools.FileFormat.AVRO: if self.schema == SCHEMA_AUTODETECT: raise ValueError( 'Schema auto-detection is not supported when using Avro based ' 'file loads into BigQuery. Please specify a schema or set ' 'temp_file_format="NEWLINE_DELIMITED_JSON"') if self.schema is None: raise ValueError( 'A schema must be provided when writing to BigQuery using ' 'Avro based file loads') if self.schema and type(self.schema) is dict: def find_in_nested_dict(schema): for field in schema['fields']: if field['type'] == 'JSON': logging.warning( 'Found JSON type in TableSchema for "File_LOADS" write ' 'method. Make sure the TableSchema field is a parsed ' 'JSON to ensure the read as a JSON type. Otherwise it ' 'will read as a raw (escaped) string.') elif field['type'] == 'STRUCT': find_in_nested_dict(field) find_in_nested_dict(self.schema) from apache_beam.io.gcp.bigquery_file_loads import BigQueryBatchFileLoads # Only cast to int when a value is given. # We only use an int for BigQueryBatchFileLoads if self.triggering_frequency is not None: triggering_frequency = int(self.triggering_frequency) else: triggering_frequency = self.triggering_frequency output = pcoll | BigQueryBatchFileLoads( destination=self.table_reference, schema=self.schema, project=self._project, create_disposition=self.create_disposition, write_disposition=self.write_disposition, triggering_frequency=triggering_frequency, with_auto_sharding=self.with_auto_sharding, temp_file_format=self._temp_file_format, max_file_size=self.max_file_size, max_partition_size=self.max_partition_size, max_files_per_bundle=self.max_files_per_bundle, custom_gcs_temp_location=self.custom_gcs_temp_location, test_client=self.test_client, table_side_inputs=self.table_side_inputs, schema_side_inputs=self.schema_side_inputs, additional_bq_parameters=self.additional_bq_parameters, validate=self._validate, is_streaming_pipeline=is_streaming_pipeline, load_job_project_id=self.load_job_project_id) return WriteResult( method=WriteToBigQuery.Method.FILE_LOADS, destination_load_jobid_pairs=output[ BigQueryBatchFileLoads.DESTINATION_JOBID_PAIRS], destination_file_pairs=output[ BigQueryBatchFileLoads.DESTINATION_FILE_PAIRS], destination_copy_jobid_pairs=output[ BigQueryBatchFileLoads.DESTINATION_COPY_JOBID_PAIRS]) elif method_to_use == WriteToBigQuery.Method.STORAGE_WRITE_API: return pcoll | StorageWriteToBigQuery( table=self.table_reference, schema=self.schema, table_side_inputs=self.table_side_inputs, create_disposition=self.create_disposition, write_disposition=self.write_disposition, additional_bq_parameters=self.additional_bq_parameters, triggering_frequency=self.triggering_frequency, use_at_least_once=self.use_at_least_once, with_auto_sharding=self.with_auto_sharding, num_storage_api_streams=self._num_storage_api_streams, use_cdc_writes=self._use_cdc_writes, primary_key=self._primary_key, big_lake_configuration=self._big_lake_configuration, expansion_service=self.expansion_service) else: raise ValueError(f"Unsupported method {method_to_use}") def display_data(self): res = {} if self.table_reference is not None and isinstance(self.table_reference, TableReference): tableSpec = '{}.{}'.format( self.table_reference.datasetId, self.table_reference.tableId) if self.table_reference.projectId is not None: tableSpec = '{}:{}'.format(self.table_reference.projectId, tableSpec) res['table'] = DisplayDataItem(tableSpec, label='Table') res['validation'] = DisplayDataItem( self._validate, label="Validation Enabled") return res def to_runner_api_parameter(self, context): from apache_beam.internal import pickler # It'd be nice to name these according to their actual # names/positions in the orignal argument list, but such a # transformation is currently irreversible given how # remove_objects_from_args and insert_values_in_args # are currently implemented. def serialize(side_inputs): return {(SIDE_INPUT_PREFIX + '%s') % ix: si.to_runner_api( context).SerializeToString() for ix, si in enumerate(side_inputs)} table_side_inputs = serialize(self.table_side_inputs) schema_side_inputs = serialize(self.schema_side_inputs) config = { 'table': self._table, 'dataset': self._dataset, 'project': self._project, 'schema': self.schema, 'create_disposition': self.create_disposition, 'write_disposition': self.write_disposition, 'kms_key': self.kms_key, 'batch_size': self.batch_size, 'max_file_size': self.max_file_size, 'max_files_per_bundle': self.max_files_per_bundle, 'custom_gcs_temp_location': self.custom_gcs_temp_location, 'method': self.method, 'insert_retry_strategy': self.insert_retry_strategy, 'additional_bq_parameters': self.additional_bq_parameters, 'table_side_inputs': table_side_inputs, 'schema_side_inputs': schema_side_inputs, 'triggering_frequency': self.triggering_frequency, 'validate': self._validate, 'temp_file_format': self._temp_file_format, 'ignore_insert_ids': self._ignore_insert_ids, 'with_auto_sharding': self.with_auto_sharding, } return 'beam:transform:write_to_big_query:v0', pickler.dumps(config) @PTransform.register_urn('beam:transform:write_to_big_query:v0', bytes) def from_runner_api(unused_ptransform, payload, context): from apache_beam.internal import pickler from apache_beam.portability.api import beam_runner_api_pb2 config = pickler.loads(payload) def deserialize(side_inputs): deserialized_side_inputs = {} for k, v in side_inputs.items(): side_input = beam_runner_api_pb2.SideInput() side_input.ParseFromString(v) deserialized_side_inputs[k] = side_input # This is an ordered list stored as a dict (see the comments in # to_runner_api_parameter above). indexed_side_inputs = [( get_sideinput_index(tag), pvalue.AsSideInput.from_runner_api(si, context)) for tag, si in deserialized_side_inputs.items()] return [si for _, si in sorted(indexed_side_inputs)] config['table_side_inputs'] = deserialize(config['table_side_inputs']) config['schema_side_inputs'] = deserialize(config['schema_side_inputs']) return WriteToBigQuery(**config) class WriteResult: """The result of a WriteToBigQuery transform. """ def __init__( self, method: str = None, destination_load_jobid_pairs: PCollection[tuple[str, JobReference]] = None, destination_file_pairs: PCollection[tuple[str, tuple[str, int]]] = None, destination_copy_jobid_pairs: PCollection[tuple[str, JobReference]] = None, failed_rows: PCollection[tuple[str, dict]] = None, failed_rows_with_errors: PCollection[tuple[str, dict, list]] = None): self._method = method self._destination_load_jobid_pairs = destination_load_jobid_pairs self._destination_file_pairs = destination_file_pairs self._destination_copy_jobid_pairs = destination_copy_jobid_pairs self._failed_rows = failed_rows self._failed_rows_with_errors = failed_rows_with_errors from apache_beam.io.gcp.bigquery_file_loads import BigQueryBatchFileLoads self.attributes = { BigQueryWriteFn.FAILED_ROWS: WriteResult.failed_rows, BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS: WriteResult. failed_rows_with_errors, BigQueryBatchFileLoads.DESTINATION_JOBID_PAIRS: WriteResult. destination_load_jobid_pairs, BigQueryBatchFileLoads.DESTINATION_FILE_PAIRS: WriteResult. destination_file_pairs, BigQueryBatchFileLoads.DESTINATION_COPY_JOBID_PAIRS: WriteResult. destination_copy_jobid_pairs, } def validate(self, valid_methods, attribute): if self._method not in valid_methods: raise AttributeError( f'Cannot get {attribute} because it is not produced ' f'by the {self._method} write method. Note: only ' f'{valid_methods} produces this attribute.') @property def destination_load_jobid_pairs( self) -> PCollection[tuple[str, JobReference]]: """A ``FILE_LOADS`` method attribute Returns: A PCollection of the table destinations that were successfully loaded to using the batch load API, along with the load job IDs. Raises: AttributeError: if accessed with a write method besides ``FILE_LOADS``.""" self.validate([WriteToBigQuery.Method.FILE_LOADS], 'DESTINATION_JOBID_PAIRS') return self._destination_load_jobid_pairs @property def destination_file_pairs(self) -> PCollection[tuple[str, tuple[str, int]]]: """A ``FILE_LOADS`` method attribute Returns: A PCollection of the table destinations along with the temp files used as sources to load from. Raises: AttributeError: if accessed with a write method besides ``FILE_LOADS``.""" self.validate([WriteToBigQuery.Method.FILE_LOADS], 'DESTINATION_FILE_PAIRS') return self._destination_file_pairs @property def destination_copy_jobid_pairs( self) -> PCollection[tuple[str, JobReference]]: """A ``FILE_LOADS`` method attribute Returns: A PCollection of the table destinations that were successfully copied to, along with the copy job ID. Raises: AttributeError: if accessed with a write method besides ``FILE_LOADS``.""" self.validate([WriteToBigQuery.Method.FILE_LOADS], 'DESTINATION_COPY_JOBID_PAIRS') return self._destination_copy_jobid_pairs @property def failed_rows(self) -> PCollection[tuple[str, dict]]: """A ``[STREAMING_INSERTS, STORAGE_WRITE_API]`` method attribute Returns: A PCollection of rows that failed when inserting to BigQuery. Raises: AttributeError: if accessed with a write method besides ``[STREAMING_INSERTS, STORAGE_WRITE_API]``.""" self.validate([ WriteToBigQuery.Method.STREAMING_INSERTS, WriteToBigQuery.Method.STORAGE_WRITE_API ], 'FAILED_ROWS') return self._failed_rows @property def failed_rows_with_errors(self) -> PCollection[tuple[str, dict, list]]: """A ``[STREAMING_INSERTS, STORAGE_WRITE_API]`` method attribute Returns: A PCollection of rows that failed when inserting to BigQuery, along with their errors. Raises: AttributeError: if accessed with a write method besides ``[STREAMING_INSERTS, STORAGE_WRITE_API]``.""" self.validate([ WriteToBigQuery.Method.STREAMING_INSERTS, WriteToBigQuery.Method.STORAGE_WRITE_API ], 'FAILED_ROWS_WITH_ERRORS') return self._failed_rows_with_errors def __getitem__(self, key): if key not in self.attributes: raise AttributeError( f'Error trying to access nonexistent attribute `{key}` in write ' 'result. Please see __documentation__ for available attributes.') return self.attributes[key].__get__(self, WriteResult) class StorageWriteToBigQuery(PTransform): """Writes data to BigQuery using Storage API. Supports dynamic destinations. Dynamic schemas are not supported yet. Experimental; no backwards compatibility guarantees. """ IDENTIFIER = "beam:schematransform:org.apache.beam:bigquery_storage_write:v2" FAILED_ROWS = "FailedRows" FAILED_ROWS_WITH_ERRORS = "FailedRowsWithErrors" # fields for rows sent to Storage API with dynamic destinations DESTINATION = "destination" RECORD = "record" # field names for rows sent to Storage API for CDC functionality CDC_INFO = "row_mutation_info" CDC_MUTATION_TYPE = "mutation_type" CDC_SQN = "change_sequence_number" # magic string to tell Java that these rows are going to dynamic destinations DYNAMIC_DESTINATIONS = "DYNAMIC_DESTINATIONS" def __init__( self, table, table_side_inputs=None, schema=None, create_disposition=BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=BigQueryDisposition.WRITE_APPEND, additional_bq_parameters=None, triggering_frequency=0, use_at_least_once=False, with_auto_sharding=False, num_storage_api_streams=0, use_cdc_writes: bool = False, primary_key: list[str] = None, big_lake_configuration=None, expansion_service=None): self._table = table self._table_side_inputs = table_side_inputs self._schema = schema self._create_disposition = create_disposition self._write_disposition = write_disposition self.additional_bq_parameters = additional_bq_parameters self._triggering_frequency = triggering_frequency self._use_at_least_once = use_at_least_once self._with_auto_sharding = with_auto_sharding self._num_storage_api_streams = num_storage_api_streams self._use_cdc_writes = use_cdc_writes self._primary_key = primary_key self._big_lake_configuration = big_lake_configuration self._expansion_service = expansion_service or BeamJarExpansionService( 'sdks:java:io:google-cloud-platform:expansion-service:build') def expand(self, input): if self._schema is None: try: schema = schema_from_element_type(input.element_type) is_rows = True except TypeError as exn: raise ValueError( "A schema is required in order to prepare rows " "for writing with STORAGE_WRITE_API.") from exn elif callable(self._schema): raise NotImplementedError( "Writing with dynamic schemas is not " "supported for this write method.") elif isinstance(self._schema, vp.ValueProvider): schema = self._schema.get() is_rows = False else: schema = self._schema is_rows = False table = bigquery_tools.get_hashable_destination(self._table) # if writing to one destination, just convert to Beam rows and send over if not callable(table): if is_rows: input_beam_rows = input else: input_beam_rows = ( input | "Convert dict to Beam Row" >> self.ConvertToBeamRows( schema, False).with_output_types()) # For dynamic destinations, we first figure out where each row is going. # Then we send (destination, record) rows over to Java SchemaTransform. # We need to do this here because there are obstacles to passing the # destinations function to Java else: # call function and append destination to each row input_rows = ( input | "Append dynamic destinations" >> beam.ParDo( bigquery_tools.AppendDestinationsFn(table), *self._table_side_inputs)) # if input type is Beam Row, just wrap everything in another Row if is_rows: input_beam_rows = ( input_rows | "Wrap in Beam Row" >> beam.Map( lambda row: beam.Row( **{ StorageWriteToBigQuery.DESTINATION: row[0], StorageWriteToBigQuery.RECORD: row[1] })).with_output_types( RowTypeConstraint.from_fields([ (StorageWriteToBigQuery.DESTINATION, str), (StorageWriteToBigQuery.RECORD, input.element_type) ]))) # otherwise, convert to Beam Rows else: input_beam_rows = ( input_rows | "Convert dict to Beam Row" >> self.ConvertToBeamRows( schema, True).with_output_types()) # communicate to Java that this write should use dynamic destinations table = StorageWriteToBigQuery.DYNAMIC_DESTINATIONS clustering_fields = [] if self.additional_bq_parameters: if callable(self.additional_bq_parameters): raise NotImplementedError( "Currently, dynamic clustering and timepartitioning is not " "supported for STORAGE_WRITE_API write method.") clustering_fields = ( self.additional_bq_parameters.get("clustering", {}).get("fields", [])) output = ( input_beam_rows | SchemaAwareExternalTransform( identifier=StorageWriteToBigQuery.IDENTIFIER, expansion_service=self._expansion_service, rearrange_based_on_discovery=True, table=table, create_disposition=self._create_disposition, write_disposition=self._write_disposition, triggering_frequency_seconds=self._triggering_frequency, auto_sharding=self._with_auto_sharding, num_streams=self._num_storage_api_streams, use_at_least_once_semantics=self._use_at_least_once, use_cdc_writes=self._use_cdc_writes, primary_key=self._primary_key, clustering_fields=clustering_fields, big_lake_configuration=self._big_lake_configuration, error_handling={ 'output': StorageWriteToBigQuery.FAILED_ROWS_WITH_ERRORS })) failed_rows_with_errors = output[ StorageWriteToBigQuery.FAILED_ROWS_WITH_ERRORS] failed_rows = failed_rows_with_errors | beam.Map( lambda row_and_error: row_and_error[0]) if not is_rows: # return back from Beam Rows to Python dict elements failed_rows = failed_rows | beam.Map(lambda row: row._asdict()) failed_rows_with_errors = failed_rows_with_errors | beam.Map( lambda row: { "error_message": row.error_message, "failed_row": row.failed_row. _asdict() }) return WriteResult( method=WriteToBigQuery.Method.STORAGE_WRITE_API, failed_rows=failed_rows, failed_rows_with_errors=failed_rows_with_errors) class ConvertToBeamRowsSetupSchema: def __init__(self, schema): self._value = schema def __enter__(self): if not isinstance(self._value, (bigquery.TableSchema, bigquery.TableFieldSchema)): return bigquery_tools.get_bq_tableschema(self._value) return self._value def __exit__(self, *args): pass class ConvertToBeamRows(PTransform): def __init__(self, schema, dynamic_destinations): self.schema = schema self.dynamic_destinations = dynamic_destinations def expand(self, input_dicts): if self.dynamic_destinations: return ( input_dicts | "Convert dict to Beam Row" >> beam.Map( lambda row, schema=DoFn.SetupContextParam( StorageWriteToBigQuery.ConvertToBeamRowsSetupSchema, args= [self.schema]): beam.Row( **{ StorageWriteToBigQuery.DESTINATION: row[0], StorageWriteToBigQuery.RECORD: bigquery_tools. beam_row_from_dict(row[1], schema) }))) else: return ( input_dicts | "Convert dict to Beam Row" >> beam.Map( lambda row, schema=DoFn.SetupContextParam( StorageWriteToBigQuery.ConvertToBeamRowsSetupSchema, args=[ self.schema ]): bigquery_tools.beam_row_from_dict(row, schema))) def with_output_types(self): row_type_hints = bigquery_tools.get_beam_typehints_from_tableschema( self.schema) if self.dynamic_destinations: type_hint = RowTypeConstraint.from_fields([ (StorageWriteToBigQuery.DESTINATION, str), ( StorageWriteToBigQuery.RECORD, RowTypeConstraint.from_fields(row_type_hints)) ]) else: type_hint = RowTypeConstraint.from_fields(row_type_hints) return super().with_output_types(type_hint) class ReadFromBigQuery(PTransform): # pylint: disable=line-too-long,W1401 """Read data from BigQuery. This PTransform uses a BigQuery export job to take a snapshot of the table on GCS, and then reads from each produced file. File format is Avro by default. Args: method: The method to use to read from BigQuery. It may be EXPORT or DIRECT_READ. EXPORT invokes a BigQuery export request (https://cloud.google.com/bigquery/docs/exporting-data). DIRECT_READ reads directly from BigQuery storage using the BigQuery Read API (https://cloud.google.com/bigquery/docs/reference/storage). If unspecified, the default is currently EXPORT. timeout (float): The timeout for the read operation in seconds. This only impacts DIRECT_READ. If None, the client default will be used. use_native_datetime (bool): By default this transform exports BigQuery DATETIME fields as formatted strings (for example: 2021-01-01T12:59:59). If :data:`True`, BigQuery DATETIME fields will be returned as native Python datetime objects. This can only be used when 'method' is 'DIRECT_READ'. table (str, callable, ValueProvider): The ID of the table, or a callable that returns it. If dataset argument is :data:`None` then the table argument must contain the entire table reference specified as: ``'DATASET.TABLE'`` or ``'PROJECT:DATASET.TABLE'``. If it's a callable, it must receive one argument representing an element to be written to BigQuery, and return a TableReference, or a string table name as specified above. dataset (str): The ID of the dataset containing this table or :data:`None` if the table reference is specified entirely by the table argument. project (str): The ID of the project containing this table. query (str, ValueProvider): A query to be used instead of arguments table, dataset, and project. validate (bool): If :data:`True`, various checks will be done when source gets initialized (e.g., is table present?). This should be :data:`True` for most scenarios in order to catch errors as early as possible (pipeline construction instead of pipeline execution). It should be :data:`False` if the table is created during pipeline execution by a previous step. Set this to :data:`False` if the BigQuery export method is slow due to checking file existence. coder (~apache_beam.coders.coders.Coder): The coder for the table rows. If :data:`None`, then the default coder is _JsonToDictCoder, which will interpret every row as a JSON serialized dictionary. use_standard_sql (bool): Specifies whether to use BigQuery's standard SQL dialect for this query. The default value is :data:`False`. If set to :data:`True`, the query will use BigQuery's updated SQL dialect with improved standards compliance. This parameter is ignored for table inputs. flatten_results (bool): Flattens all nested and repeated fields in the query results. The default value is :data:`True`. kms_key (str): Optional Cloud KMS key name for use when creating new temporary tables. gcs_location (str, ValueProvider): The name of the Google Cloud Storage bucket where the extracted table should be written as a string or a :class:`~apache_beam.options.value_provider.ValueProvider`. If :data:`None`, then the temp_location parameter is used. bigquery_job_labels (dict): A dictionary with string labels to be passed to BigQuery export and query jobs created by this transform. See: https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#JobConfiguration use_json_exports (bool): By default, this transform works by exporting BigQuery data into Avro files, and reading those files. With this parameter, the transform will instead export to JSON files. JSON files are slower to read due to their larger size. When using JSON exports, the BigQuery types for DATE, DATETIME, TIME, and TIMESTAMP will be exported as strings. This behavior is consistent with BigQuerySource. When using Avro exports, these fields will be exported as native Python types (datetime.date, datetime.datetime, datetime.datetime, and datetime.datetime respectively). Avro exports are recommended. To learn more about BigQuery types, and Time-related type representations, see: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types To learn more about type conversions between BigQuery and Avro, see: https://cloud.google.com/bigquery/docs/loading-data-cloud-storage-avro#avro_conversions temp_dataset (``apache_beam.io.gcp.internal.clients.bigquery.DatasetReference``): Temporary dataset reference to use when reading from BigQuery using a query. When reading using a query, BigQuery source will create a temporary dataset and a temporary table to store the results of the query. With this option, you can set an existing dataset to create the temporary table in. BigQuery source will create a temporary table in that dataset, and will remove it once it is not needed. Job needs access to create and delete tables within the given dataset. Dataset name should *not* start with the reserved prefix `beam_temp_dataset_`. query_priority (BigQueryQueryPriority): By default, this transform runs queries with BATCH priority. Use :attr:`BigQueryQueryPriority.INTERACTIVE` to run queries with INTERACTIVE priority. This option is ignored when reading from a table rather than a query. To learn more about query priority, see: https://cloud.google.com/bigquery/docs/running-queries output_type (str): By default, this source yields Python dictionaries (`PYTHON_DICT`). There is experimental support for producing a PCollection with a schema and yielding Beam Rows via the option `BEAM_ROW`. For more information on schemas, see https://beam.apache.org/documentation/programming-guide/#what-is-a-schema) """ class Method(object): EXPORT = 'EXPORT' # This is currently the default. DIRECT_READ = 'DIRECT_READ' COUNTER = 0 def __init__( self, gcs_location=None, method=None, use_native_datetime=False, output_type=None, timeout=None, *args, **kwargs): self.method = method or ReadFromBigQuery.Method.EXPORT self.use_native_datetime = use_native_datetime self.output_type = output_type self._args = args self._kwargs = kwargs if timeout is not None: self._kwargs['timeout'] = timeout if self.method == ReadFromBigQuery.Method.EXPORT \ and self.use_native_datetime is True: raise TypeError( 'The "use_native_datetime" parameter cannot be True for EXPORT.' ' Please set the "use_native_datetime" parameter to False *OR*' ' set the "method" parameter to ReadFromBigQuery.Method.DIRECT_READ.') if gcs_location and self.method == ReadFromBigQuery.Method.EXPORT: if not isinstance(gcs_location, (str, ValueProvider)): raise TypeError( '%s: gcs_location must be of type string' ' or ValueProvider; got %r instead' % (self.__class__.__name__, type(gcs_location))) if isinstance(gcs_location, str): gcs_location = StaticValueProvider(str, gcs_location) if self.output_type == 'BEAM_ROW' and self._kwargs.get('query', None) is not None: raise ValueError( "Both a query and an output type of 'BEAM_ROW' were specified. " "'BEAM_ROW' is not currently supported with queries.") self.gcs_location = gcs_location self.bigquery_dataset_labels = { 'type': 'bq_direct_read_' + str(uuid.uuid4())[0:10] } def expand(self, pcoll): if self.method == ReadFromBigQuery.Method.EXPORT: output_pcollection = self._expand_export(pcoll) elif self.method == ReadFromBigQuery.Method.DIRECT_READ: output_pcollection = self._expand_direct_read(pcoll) else: raise ValueError( 'The method to read from BigQuery must be either EXPORT ' 'or DIRECT_READ.') return self._expand_output_type(output_pcollection) def _expand_output_type(self, output_pcollection): if self.output_type == 'PYTHON_DICT' or self.output_type is None: return output_pcollection elif self.output_type == 'BEAM_ROW': table_details = bigquery_tools.parse_table_reference( table=self._kwargs.get("table", None), dataset=self._kwargs.get("dataset", None), project=self._kwargs.get("project", None)) if isinstance(self._kwargs['table'], ValueProvider): raise TypeError( '%s: table must be of type string' '; got ValueProvider instead' % self.__class__.__name__) elif callable(self._kwargs['table']): raise TypeError( '%s: table must be of type string' '; got a callable instead' % self.__class__.__name__) return output_pcollection | bigquery_schema_tools.convert_to_usertype( bigquery_tools.BigQueryWrapper().get_table( project_id=table_details.projectId, dataset_id=table_details.datasetId, table_id=table_details.tableId).schema, self._kwargs.get('selected_fields', None)) else: raise ValueError( 'The output type from BigQuery must be either PYTHON_DICT ' 'or BEAM_ROW.') def _expand_export(self, pcoll): # TODO(https://github.com/apache/beam/issues/20683): Make ReadFromBQ rely # on ReadAllFromBQ implementation. temp_location = pcoll.pipeline.options.view_as( GoogleCloudOptions).temp_location job_name = pcoll.pipeline.options.view_as(GoogleCloudOptions).job_name gcs_location_vp = self.gcs_location unique_id = str(uuid.uuid4())[0:10] def file_path_to_remove(unused_elm): gcs_location = bigquery_export_destination_uri( gcs_location_vp, temp_location, unique_id, True) return gcs_location + '/' files_to_remove_pcoll = beam.pvalue.AsList( pcoll.pipeline | 'FilesToRemoveImpulse' >> beam.Create([None]) | 'MapFilesToRemove' >> beam.Map(file_path_to_remove)) try: step_name = self.label except AttributeError: step_name = 'ReadFromBigQuery_%d' % ReadFromBigQuery.COUNTER ReadFromBigQuery.COUNTER += 1 return ( pcoll | beam.io.Read( _CustomBigQuerySource( gcs_location=self.gcs_location, pipeline_options=pcoll.pipeline.options, method=self.method, job_name=job_name, step_name=step_name, unique_id=unique_id, *self._args, **self._kwargs)) | _PassThroughThenCleanup(files_to_remove_pcoll)) def _expand_direct_read(self, pcoll): project_id = None temp_table_ref = None if 'temp_dataset' in self._kwargs: temp_table_ref = bigquery.TableReference( projectId=self._kwargs['temp_dataset'].projectId, datasetId=self._kwargs['temp_dataset'].datasetId, tableId='beam_temp_table_' + uuid.uuid4().hex) else: project_id = pcoll.pipeline.options.view_as(GoogleCloudOptions).project pipeline_details = {} if temp_table_ref is not None: pipeline_details['temp_table_ref'] = temp_table_ref elif project_id is not None: pipeline_details['project_id'] = project_id pipeline_details['bigquery_dataset_labels'] = self.bigquery_dataset_labels def _get_pipeline_details(unused_elm): return pipeline_details project_to_cleanup_pcoll = beam.pvalue.AsList( pcoll.pipeline | 'ProjectToCleanupImpulse' >> beam.Create([None]) | 'MapProjectToCleanup' >> beam.Map(_get_pipeline_details)) return ( pcoll | beam.io.Read( _CustomBigQueryStorageSource( pipeline_options=pcoll.pipeline.options, method=self.method, use_native_datetime=self.use_native_datetime, temp_table=temp_table_ref, bigquery_dataset_labels=self.bigquery_dataset_labels, *self._args, **self._kwargs)) | _PassThroughThenCleanupTempDatasets(project_to_cleanup_pcoll)) class ReadFromBigQueryRequest: """ Class that defines data to read from BQ. """ def __init__( self, query: str = None, use_standard_sql: bool = True, table: Union[str, TableReference] = None, flatten_results: bool = False): """ Only one of query or table should be specified. :param query: SQL query to fetch data. :param use_standard_sql: Specifies whether to use BigQuery's standard SQL dialect for this query. The default value is :data:`True`. If set to :data:`False`, the query will use BigQuery's legacy SQL dialect. This parameter is ignored for table inputs. :param table: The ID of the table to read. Table should define project and dataset (ex.: ``'PROJECT:DATASET.TABLE'``). :param flatten_results: Flattens all nested and repeated fields in the query results. The default value is :data:`False`. """ self.flatten_results = flatten_results self.query = query self.use_standard_sql = use_standard_sql self.table = table self.validate() # We use this internal object ID to generate BigQuery export directories # and to create BigQuery job names self.obj_id = '%d_%s' % (int(time.time()), secrets.token_hex(3)) def validate(self): if self.table is not None and self.query is not None: raise ValueError( 'Both a BigQuery table and a query were specified.' ' Please specify only one of these.') elif self.table is None and self.query is None: raise ValueError('A BigQuery table or a query must be specified') if self.table is not None: if isinstance(self.table, str): assert self.table.find('.'), ( 'Expected a table reference ' '(PROJECT:DATASET.TABLE or DATASET.TABLE) instead of %s' % self.table) class ReadAllFromBigQuery(PTransform): """Read data from BigQuery. PTransform:ReadFromBigQueryRequest->Rows This PTransform uses a BigQuery export job to take a snapshot of the table on GCS, and then reads from each produced file. Data is exported into a new subdirectory for each export using UUIDs generated in `ReadFromBigQueryRequest` objects. It is recommended not to use this PTransform for streaming jobs on GlobalWindow, since it will not be able to cleanup snapshots. Args: gcs_location (str): The name of the Google Cloud Storage bucket where the extracted table should be written as a string. If :data:`None`, then the temp_location parameter is used. validate (bool): If :data:`True`, various checks will be done when source gets initialized (e.g., is table present?). Set this to :data:`False` if the BigQuery export method is slow due to checking file existence. kms_key (str): Experimental. Optional Cloud KMS key name for use when creating new temporary tables. """ COUNTER = 0 def __init__( self, gcs_location: Union[str, ValueProvider] = None, validate: bool = False, kms_key: str = None, temp_dataset: Union[str, DatasetReference] = None, bigquery_job_labels: dict[str, str] = None, query_priority: str = BigQueryQueryPriority.BATCH): if gcs_location: if not isinstance(gcs_location, (str, ValueProvider)): raise TypeError( '%s: gcs_location must be of type string' ' or ValueProvider; got %r instead' % (self.__class__.__name__, type(gcs_location))) self.gcs_location = gcs_location self.validate = validate self.kms_key = kms_key self.bigquery_job_labels = bigquery_job_labels self.temp_dataset = temp_dataset self.query_priority = query_priority def expand(self, pcoll): job_name = pcoll.pipeline.options.view_as(GoogleCloudOptions).job_name project = pcoll.pipeline.options.view_as(GoogleCloudOptions).project unique_id = str(uuid.uuid4())[0:10] try: step_name = self.label except AttributeError: step_name = 'ReadAllFromBigQuery_%d' % ReadAllFromBigQuery.COUNTER ReadAllFromBigQuery.COUNTER += 1 sources_to_read, cleanup_locations = ( pcoll | beam.ParDo( _BigQueryReadSplit( options=pcoll.pipeline.options, gcs_location=self.gcs_location, validate=self.validate, bigquery_job_labels=self.bigquery_job_labels, job_name=job_name, step_name=step_name, unique_id=unique_id, kms_key=self.kms_key, project=project, temp_dataset=self.temp_dataset, query_priority=self.query_priority)).with_outputs( "location_to_cleanup", main="files_to_read") ) return ( sources_to_read | SDFBoundedSourceReader(data_to_display=self.display_data()) | _PassThroughThenCleanup(beam.pvalue.AsIter(cleanup_locations))) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_avro_tools.py0000644000175100017510000001054115206113032023506 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tools used tool work with Avro files in the context of BigQuery. Classes, constants and functions in this file are experimental and have no backwards compatibility guarantees. NOTHING IN THIS FILE HAS BACKWARDS COMPATIBILITY GUARANTEES. """ from typing import Any # BigQuery types as listed in # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types # with aliases (RECORD, BOOLEAN, FLOAT, INTEGER) as defined in # https://developers.google.com/resources/api-libraries/documentation/bigquery/v2/java/latest/com/google/api/services/bigquery/model/TableFieldSchema.html#setType-java.lang.String- BIG_QUERY_TO_AVRO_TYPES = { "STRUCT": "record", "RECORD": "record", "STRING": "string", "BOOL": "boolean", "BOOLEAN": "boolean", "BYTES": "bytes", "FLOAT64": "double", "FLOAT": "double", "INT64": "long", "INTEGER": "long", "TIME": { "type": "long", "logicalType": "time-micros", }, "TIMESTAMP": { "type": "long", "logicalType": "timestamp-micros", }, "DATE": { "type": "int", "logicalType": "date", }, "DATETIME": "string", "NUMERIC": { "type": "bytes", "logicalType": "decimal", # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#numeric-type "precision": 38, "scale": 9, }, "GEOGRAPHY": "string", } def get_record_schema_from_dict_table_schema( schema_name: str, table_schema: dict[str, Any], namespace: str = "apache_beam.io.gcp.bigquery") -> dict[str, Any]: # noqa: F821 """Convert a table schema into an Avro schema. Args: schema_name (str): The name of the record. table_schema (Dict[str, Any]): A BigQuery table schema in dict form. namespace (str): The namespace of the Avro schema. Returns: Dict[str, Any]: The schema as an Avro RecordSchema. """ avro_fields = [ table_field_to_avro_field(field, ".".join((namespace, schema_name))) for field in table_schema["fields"] ] return { "type": "record", "name": schema_name, "fields": avro_fields, "doc": "Translated Avro Schema for {}".format(schema_name), "namespace": namespace, } def table_field_to_avro_field(table_field: dict[str, Any], namespace: str) -> dict[str, Any]: # noqa: F821 """Convert a BigQuery field to an avro field. Args: table_field (Dict[str, Any]): A BigQuery field in dict form. Returns: Dict[str, Any]: An equivalent Avro field in dict form. """ assert "type" in table_field, \ "Unable to get type for table field {}".format(table_field) assert table_field["type"] in BIG_QUERY_TO_AVRO_TYPES, \ "Unable to map BigQuery field type {} to avro type".format( table_field["type"]) avro_type = BIG_QUERY_TO_AVRO_TYPES[table_field["type"]] if avro_type == "record": element_type = get_record_schema_from_dict_table_schema( table_field["name"], table_field, namespace=".".join((namespace, table_field["name"]))) else: element_type = avro_type field_mode = table_field.get("mode", "NULLABLE") if field_mode in (None, "NULLABLE"): field_type = ["null", element_type] elif field_mode == "REQUIRED": field_type = element_type elif field_mode == "REPEATED": field_type = {"type": "array", "items": element_type} else: raise ValueError("Unknown BigQuery field mode: {}".format(field_mode)) avro_field = {"type": field_type, "name": table_field["name"]} doc = table_field.get("description") if doc: avro_field["doc"] = doc return avro_field ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_avro_tools_test.py0000644000175100017510000001064715206113032024554 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import logging import unittest from fastavro.schema import parse_schema from apache_beam.io.gcp import bigquery_avro_tools from apache_beam.io.gcp import bigquery_tools from apache_beam.io.gcp.bigquery_test import HttpError from apache_beam.io.gcp.internal.clients import bigquery @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestBigQueryToAvroSchema(unittest.TestCase): def test_convert_bigquery_schema_to_avro_schema(self): subfields = [ bigquery.TableFieldSchema( name="species", type="STRING", mode="NULLABLE"), ] fields = [ bigquery.TableFieldSchema( name="number", type="INTEGER", mode="REQUIRED"), bigquery.TableFieldSchema( name="species", type="STRING", mode="NULLABLE"), bigquery.TableFieldSchema(name="quality", type="FLOAT"), # default to NULLABLE bigquery.TableFieldSchema(name="grade", type="FLOAT64"), # default to NULLABLE bigquery.TableFieldSchema(name="quantity", type="INTEGER"), # default to NULLABLE bigquery.TableFieldSchema(name="dependents", type="INT64"), # default to NULLABLE bigquery.TableFieldSchema( name="birthday", type="TIMESTAMP", mode="NULLABLE"), bigquery.TableFieldSchema( name="birthdayMoney", type="NUMERIC", mode="NULLABLE"), bigquery.TableFieldSchema( name="flighted", type="BOOL", mode="NULLABLE"), bigquery.TableFieldSchema( name="flighted2", type="BOOLEAN", mode="NULLABLE"), bigquery.TableFieldSchema(name="sound", type="BYTES", mode="NULLABLE"), bigquery.TableFieldSchema( name="anniversaryDate", type="DATE", mode="NULLABLE"), bigquery.TableFieldSchema( name="anniversaryDatetime", type="DATETIME", mode="NULLABLE"), bigquery.TableFieldSchema( name="anniversaryTime", type="TIME", mode="NULLABLE"), bigquery.TableFieldSchema( name="scion", type="RECORD", mode="NULLABLE", fields=subfields), bigquery.TableFieldSchema( name="family", type="STRUCT", mode="NULLABLE", fields=subfields), bigquery.TableFieldSchema( name="associates", type="RECORD", mode="REPEATED", fields=subfields), bigquery.TableFieldSchema( name="geoPositions", type="GEOGRAPHY", mode="NULLABLE"), ] table_schema = bigquery.TableSchema(fields=fields) avro_schema = bigquery_avro_tools.get_record_schema_from_dict_table_schema( "root", bigquery_tools.get_dict_table_schema(table_schema)) parsed_schema = parse_schema(avro_schema) self.assertEqual(type(parsed_schema), dict) # names: key -> name, value -> different types allowed names = { "number": 4, "species": 2, "quality": 2, "grade": 2, "quantity": 2, "dependents": 2, "birthday": 2, "birthdayMoney": 2, "flighted": 2, "flighted2": 2, "sound": 2, "anniversaryDate": 2, "anniversaryDatetime": 2, "anniversaryTime": 2, "scion": 2, "family": 2, "associates": 2, "geoPositions": 2, } # simple test case to check if the schema is parsed right. fields = parsed_schema["fields"] for i in range(len(fields)): field_ = fields[i] assert 'name' in field_ and field_['name'] in names self.assertEqual(len(field_['type']), names[field_['name']]) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_biglake_test.py0000644000175100017510000001063315206113032023756 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for BigQuery BigLake configuration.""" import unittest from unittest import mock from apache_beam.io.gcp import bigquery @mock.patch('apache_beam.io.gcp.bigquery.BeamJarExpansionService') class BigQueryBigLakeTest(unittest.TestCase): """Test BigLake configuration support in BigQuery Storage Write API.""" def test_storage_write_to_bigquery_with_biglake_config( self, mock_expansion_service): """Test that StorageWriteToBigQuery accepts bigLakeConfiguration.""" big_lake_config = { 'connectionId': ( 'projects/test-project/locations/us/connections/test-connection'), 'storageUri': 'gs://test-bucket/test-path', 'fileFormat': 'parquet', 'tableFormat': 'iceberg' } # Test that the constructor accepts the bigLakeConfiguration parameter transform = bigquery.StorageWriteToBigQuery( table='test-project:test_dataset.test_table', big_lake_configuration=big_lake_config) # Verify the configuration is stored self.assertEqual(transform._big_lake_configuration, big_lake_config) def test_storage_write_to_bigquery_without_biglake_config( self, mock_expansion_service): """Test that StorageWriteToBigQuery works without bigLakeConfiguration.""" transform = bigquery.StorageWriteToBigQuery( table='test-project:test_dataset.test_table') # Verify the configuration is None by default self.assertIsNone(transform._big_lake_configuration) def test_biglake_config_passed_to_external_transform( self, mock_expansion_service): """Test that StorageWriteToBigQuery accepts bigLakeConfiguration.""" big_lake_config = { 'connection_id': 'projects/my-project/locations/us/connections/my-conn', 'table_format': 'ICEBERG' } # Mock the expansion service to avoid JAR dependency mock_expansion_service.return_value = mock.MagicMock() # Create the transform transform = bigquery.StorageWriteToBigQuery( table='my-project:my_dataset.my_table', big_lake_configuration=big_lake_config) # Verify the big_lake_configuration is stored correctly self.assertEqual(transform._big_lake_configuration, big_lake_config) # Verify that the transform has the expected identifier self.assertEqual( transform.IDENTIFIER, "beam:schematransform:org.apache.beam:bigquery_storage_write:v2") # Verify that the expansion service was created (mocked) mock_expansion_service.assert_called_once_with( 'sdks:java:io:google-cloud-platform:expansion-service:build') def test_biglake_config_validation(self, mock_expansion_service): """Test validation of bigLakeConfiguration parameters.""" # Test with minimal required configuration minimal_config = { 'connectionId': ( 'projects/test-project/locations/us/connections/test-connection'), 'storageUri': 'gs://test-bucket/test-path' } transform = bigquery.StorageWriteToBigQuery( table='test-project:test_dataset.test_table', big_lake_configuration=minimal_config) self.assertEqual(transform._big_lake_configuration, minimal_config) # Test with full configuration full_config = { 'connectionId': ( 'projects/test-project/locations/us/connections/test-connection'), 'storageUri': 'gs://test-bucket/test-path', 'fileFormat': 'parquet', 'tableFormat': 'iceberg' } transform = bigquery.StorageWriteToBigQuery( table='test-project:test_dataset.test_table', big_lake_configuration=full_config) self.assertEqual(transform._big_lake_configuration, full_config) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_change_history.py0000644000175100017510000014330615206113032024333 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Streaming source for BigQuery change history (APPENDS/CHANGES functions). This module provides ``ReadBigQueryChangeHistory``, a streaming PTransform that continuously polls BigQuery APPENDS() or CHANGES() functions and emits changed rows as an unbounded PCollection. **Status: Experimental**: API may change without notice. Usage:: import apache_beam as beam from apache_beam.io.gcp.bigquery_change_history import ReadBigQueryChangeHistory with beam.Pipeline(options=pipeline_options) as p: changes = ( p | ReadBigQueryChangeHistory( table='my-project:my_dataset.my_table', change_function='APPENDS', poll_interval_sec=60)) Architecture: Poll: Polling SDF emits lightweight _QueryRange instructions. Query: _ExecuteQueryFn runs the BQ query, writes to a temp table. Read: SDF reads temp table via Storage Read API with dynamic splitting. Cleanup: Stateful DoFn tracks stream completion, deletes temp tables. """ import dataclasses import datetime import logging import sys import time import uuid from typing import Any from typing import Iterable from typing import Optional import apache_beam as beam from apache_beam.io.gcp import bigquery_tools from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.io.iobase import WatermarkEstimator from apache_beam.io.restriction_trackers import OffsetRange from apache_beam.io.restriction_trackers import OffsetRestrictionTracker from apache_beam.io.watermark_estimators import ManualWatermarkEstimator from apache_beam.metrics import Metrics from apache_beam.transforms.core import WatermarkEstimatorProvider from apache_beam.transforms.window import TimestampedValue from apache_beam.utils import retry from apache_beam.utils.timestamp import MAX_TIMESTAMP from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import Timestamp try: from google.cloud import bigquery_storage_v1 as bq_storage except ImportError: bq_storage = None # type: ignore try: import pyarrow except ImportError: pyarrow = None # type: ignore _LOGGER = logging.getLogger(__name__) __all__ = ['ReadBigQueryChangeHistory'] # Max time range for CHANGES() queries: 1 day. _MAX_CHANGES_RANGE = Duration(seconds=86400) # Side output tag for cleanup signals between the Read SDF and Cleanup DoFn. _CLEANUP_TAG = 'cleanup' # Default number of Storage Read API streams to request. # Matches ReadFromBigQuery's MIN_SPLIT_COUNT to enable parallelism. # The server may return fewer streams if the table is small. _DEFAULT_MAX_STREAMS = 10 # Default table expiration for auto-created temp datasets: 24 hours in ms. # Tables created in the dataset auto-expire after this duration if not # explicitly deleted, acting as a safety net for orphaned temp tables # (e.g. pipeline crash before cleanup runs). _DEFAULT_TABLE_EXPIRATION_MS = 24 * 60 * 60 * 1000 @dataclasses.dataclass class _QueryResult: """Bridges the Query step (query execution) to the Read SDF. After _ExecuteQueryFn runs a CHANGES/APPENDS query, it emits a _QueryResult pointing to the temp table containing query results. The Read SDF reads rows from that temp table via the Storage Read API. range_start/range_end define the time window this query covers as Beam Timestamps (int microseconds internally). The Read SDF uses range_start to set an initial watermark hold so the runner doesn't advance the watermark past the data's timestamps. """ temp_table_ref: 'bigquery.TableReference' range_start: Timestamp range_end: Timestamp @dataclasses.dataclass class _PollConfig: """Input element for the polling SDF. Only contains start_time (Beam Timestamp), which _PollWatermarkEstimatorProvider uses to initialize the watermark hold. All other config is passed via _PollChangeHistoryFn.__init__. """ start_time: Timestamp @dataclasses.dataclass class _QueryRange: """Lightweight instruction emitted by the polling SDF. Contains only the time range to query as Beam Timestamps (int microseconds internally). Static config (table, project, etc.) is held by _ExecuteQueryFn which receives these after a Reshuffle commit boundary, preventing duplicate queries on SDF re-dispatch. """ chunk_start: Timestamp chunk_end: Timestamp class _StreamRestriction: """Restriction carrying BQ Storage stream names for cross-worker safety. Unlike a plain OffsetRange(0, N), this restriction is self-contained: each split carries the actual stream name strings so it can be processed on any worker. Composes an OffsetRange for offset logic. """ __slots__ = ('stream_names', 'range') def __init__( self, stream_names: tuple[str, ...], start: int, stop: int) -> None: self.stream_names = stream_names # tuple of BQ stream name strings self.range = OffsetRange(start, stop) @property def start(self) -> int: return self.range.start @property def stop(self) -> int: return self.range.stop def __eq__(self, other: object) -> bool: if not isinstance(other, _StreamRestriction): return False return ( self.stream_names == other.stream_names and self.range == other.range) def __hash__(self) -> int: return hash((type(self), self.stream_names, self.range)) def __repr__(self) -> str: return ( '_StreamRestriction(streams=%d, start=%d, stop=%d)' % (len(self.stream_names), self.start, self.stop)) def size(self) -> int: return self.range.size() class _StreamRestrictionTracker(beam.io.iobase.RestrictionTracker): """Tracker for _StreamRestriction, delegating offset logic to OffsetRestrictionTracker.""" def __init__(self, restriction: _StreamRestriction) -> None: self._stream_names = restriction.stream_names self._offset_tracker = OffsetRestrictionTracker(restriction.range) def current_restriction(self) -> _StreamRestriction: r = self._offset_tracker.current_restriction() return _StreamRestriction(self._stream_names, r.start, r.stop) def try_claim(self, position: int) -> bool: return self._offset_tracker.try_claim(position) def try_split( self, fraction_of_remainder: float ) -> Optional[tuple[_StreamRestriction, _StreamRestriction]]: result = self._offset_tracker.try_split(fraction_of_remainder) if result is not None: primary, residual = result return ( _StreamRestriction(self._stream_names, primary.start, primary.stop), _StreamRestriction(self._stream_names, residual.start, residual.stop)) return None def check_done(self) -> None: self._offset_tracker.check_done() def current_progress(self): return self._offset_tracker.current_progress() def is_bounded(self) -> bool: return True class _NonSplittableOffsetTracker(OffsetRestrictionTracker): """OffsetRestrictionTracker that allows checkpointing but prevents splitting. Checkpointing (fraction=0) is required for defer_remainder(). All other split fractions are refused, ensuring the polling SDF runs as a singleton. """ def try_split( self, fraction_of_remainder: float ) -> Optional[tuple[OffsetRange, OffsetRange]]: if fraction_of_remainder == 0: return super().try_split(fraction_of_remainder) return None class _PollWatermarkEstimator(WatermarkEstimator): """Watermark estimator that tracks both a watermark hold and poll cursor. The watermark hold (reported via current_watermark) is set to start_ts: the earliest data timestamp emitted by the current poll. This prevents downstream stages from seeing data as late. The poll cursor (last_end) tracks where the next poll should start. This is separate from the watermark so we can hold the watermark back at start_ts while still advancing the poll cursor to end_ts. All timestamps are Beam Timestamps (int microseconds internally). State is checkpointed as (watermark_hold, last_end) so both values survive SDF re-dispatch. """ def __init__(self, state: tuple[Timestamp, Timestamp]) -> None: self._watermark_hold, self._last_end = state def observe_timestamp(self, timestamp: Timestamp) -> None: pass def current_watermark(self) -> Timestamp: return self._watermark_hold def get_estimator_state(self) -> tuple[Timestamp, Timestamp]: return (self._watermark_hold, self._last_end) def set_watermark(self, timestamp: Timestamp) -> None: if not isinstance(timestamp, Timestamp): raise ValueError('set_watermark expects a Timestamp as input') if self._watermark_hold and self._watermark_hold > timestamp: raise ValueError( 'Watermark must be monotonically increasing. ' 'Provided %s < current %s' % (timestamp, self._watermark_hold)) self._watermark_hold = timestamp def advance_poll_cursor(self, end: Timestamp) -> None: """Record end so the next poll starts from here. Only advances forward: if end is earlier than the current cursor (e.g. BQ clock regression), the cursor stays put so the next poll doesn't re-query an already-covered range. """ self._last_end = max(self._last_end, end) def poll_cursor(self) -> Timestamp: """Return the start Timestamp for the next poll.""" return self._last_end class _PollWatermarkEstimatorProvider(WatermarkEstimatorProvider): """Provider for _PollWatermarkEstimator. Initializes with watermark hold at start_time and poll cursor at start_time (first poll will query from start_time). """ def initial_estimator_state( self, element: _PollConfig, restriction: OffsetRange) -> tuple[Timestamp, Timestamp]: return (element.start_time, element.start_time) def create_watermark_estimator( self, estimator_state: tuple[Timestamp, Timestamp]) -> _PollWatermarkEstimator: return _PollWatermarkEstimator(estimator_state) def build_changes_query( table: str, start: Timestamp, end: Timestamp, change_function: str, change_type_column: str = 'change_type', change_timestamp_column: str = 'change_timestamp', columns: Optional[list[str]] = None, row_filter: Optional[str] = None) -> str: """Build a CHANGES() or APPENDS() SQL query. Args: table: Table name as 'project.dataset.table' or 'project:dataset.table'. start: Start timestamp (Beam Timestamp). Inclusive. end: End timestamp (Beam Timestamp). Exclusive. change_function: 'CHANGES' or 'APPENDS'. change_type_column: Output column name for _CHANGE_TYPE pseudo-column. change_timestamp_column: Output column name for _CHANGE_TIMESTAMP pseudo-column. columns: Optional list of column names to select. If None, selects all columns. Pseudo-columns are always appended regardless. row_filter: Optional SQL WHERE clause (without the WHERE keyword). Applied after the CHANGES/APPENDS function. Returns: SQL string. """ # Normalize 'project:dataset.table' to 'project.dataset.table' table = table.replace(':', '.') start_iso = start.to_rfc3339() end_iso = end.to_rfc3339() # Pseudo-columns (_CHANGE_TYPE, _CHANGE_TIMESTAMP) can't be written to # destination tables with their original names. Rename them so they can # be persisted to the temp table for Storage Read API reading. pseudo = ( f"_CHANGE_TYPE AS {change_type_column}, " f"_CHANGE_TIMESTAMP AS {change_timestamp_column}") if columns is None: select = f"SELECT * EXCEPT(_CHANGE_TYPE, _CHANGE_TIMESTAMP), {pseudo}" else: select = f"SELECT {', '.join(columns)}, {pseudo}" from_clause = ( f"FROM {change_function}" f"(TABLE `{table}`, " f"TIMESTAMP '{start_iso}', " f"TIMESTAMP '{end_iso}')") where = f" WHERE {row_filter}" if row_filter else "" return f"{select} {from_clause}{where}" def compute_ranges(start: Timestamp, end: Timestamp, change_function: str) -> list[tuple[Timestamp, Timestamp]]: """Split [start, end) into query-safe chunks. CHANGES() has a max 1-day range. APPENDS() has no limit. Args: start: Start Timestamp. Inclusive. end: End Timestamp. Exclusive. change_function: 'CHANGES' or 'APPENDS'. Returns: List of (start, end) Timestamp tuples. Empty if end <= start. """ if end <= start: return [] if change_function != 'CHANGES': return [(start, end)] # CHANGES: chunk into <=1-day ranges ranges = [] current = start while current < end: chunk_end = min(current + _MAX_CHANGES_RANGE, end) ranges.append((current, chunk_end)) current = chunk_end return ranges def _utc(ts: Timestamp) -> str: """Format a Beam Timestamp as a concise UTC string for logging.""" return ts.to_utc_datetime(has_tz=True).strftime('%Y-%m-%dT%H:%M:%S.%f') # ============================================================================= # Poll: _PollChangeHistoryFn (Polling SDF) # ============================================================================= class _PollChangeHistoryFn(beam.DoFn, beam.transforms.core.RestrictionProvider): """SDF that periodically emits _QueryRange instructions. Uses defer_remainder() for poll timing and _PollWatermarkEstimator to control the watermark. The watermark is initially held at start_time, then advanced to start_ts of each poll. All timestamps are Beam Timestamps (int microseconds internally). Durations (buffer, poll_interval) are Beam Durations. Derives start_ts from the poll cursor. On each poll: 1. start_ts = poll cursor (last end_ts, or start_time on first poll) 2. end_ts = bq_now - buffer 3. Computes query chunks, yields _QueryRange per chunk 4. Advances poll cursor to end_ts (for next poll's start) 5. Advances watermark to start_ts (earliest data in this poll) 6. Defers to next poll interval """ def __init__( self, table: str, project: str, change_function: str, buffer: Duration, start_time: Timestamp, stop_time: Timestamp, poll_interval: Duration, location: Optional[str] = None) -> None: self._table = table self._project = project self._change_function = change_function self._buffer = buffer self._start_time = start_time self._stop_time = stop_time self._poll_interval = poll_interval self._location = location def setup(self) -> None: self._bq_wrapper = bigquery_tools.BigQueryWrapper() if self._location is None: table_ref = bigquery_tools.parse_table_reference( self._table, project=self._project) self._location = self._bq_wrapper.get_table_location( table_ref.projectId, table_ref.datasetId, table_ref.tableId) _LOGGER.info( '[Poll] Inferred location=%s from source table %s', self._location, self._table) @retry.with_exponential_backoff( num_retries=3, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _get_bq_timestamp(self) -> Timestamp: """Query BigQuery for the current server timestamp. Returns a Beam Timestamp created from integer microseconds. Uses BQ's CURRENT_TIMESTAMP instead of the local clock to avoid data loss from clock skew between the worker VM and BigQuery. """ request = bigquery.BigqueryJobsQueryRequest( projectId=self._project, queryRequest=bigquery.QueryRequest( query='SELECT UNIX_MICROS(CURRENT_TIMESTAMP()) AS ts', useLegacySql=False, location=self._location)) response = self._bq_wrapper.client.jobs.Query(request) return Timestamp(micros=int(response.rows[0].f[0].v.string_value)) def initial_restriction(self, element: _PollConfig) -> OffsetRange: return OffsetRange(0, sys.maxsize) def create_tracker( self, restriction: OffsetRange) -> _NonSplittableOffsetTracker: # Guarantee at least one poll cycle: restriction.start == 0 on the first # invocation (from initial_restriction). After the first try_claim(0) + # defer_remainder, subsequent invocations arrive with start >= 1. if restriction.start > 0 and time.time() >= float(self._stop_time): _LOGGER.info( '[Poll] create_tracker: stop_time reached, ' 'returning empty range to terminate SDF') return _NonSplittableOffsetTracker( OffsetRange(restriction.start, restriction.start)) return _NonSplittableOffsetTracker(restriction) def restriction_size( self, element: _PollConfig, restriction: OffsetRange) -> int: return 1 def split(self, element: _PollConfig, restriction: OffsetRange) -> Iterable[OffsetRange]: yield restriction def truncate(self, element: _PollConfig, restriction: OffsetRange) -> None: return None def _next_poll_time(self, start_ts: Timestamp, now: float) -> Optional[Timestamp]: """Return a Timestamp to defer to, or None if we should poll now.""" earliest = start_ts + self._buffer + self._poll_interval if now < float(earliest): return earliest return None def _emit_query_ranges( self, start_ts: Timestamp, end_ts: Timestamp, watermark_estimator: _PollWatermarkEstimator) -> Iterable[_QueryRange]: """Compute and yield _QueryRange elements, advancing estimator state.""" ranges = compute_ranges(start_ts, end_ts, self._change_function) _LOGGER.info( '[Poll] %d chunks for [%s, %s)', len(ranges), _utc(start_ts), _utc(end_ts)) Metrics.counter('BigQueryChangeHistory', 'polls').inc() watermark_estimator.advance_poll_cursor(end_ts) watermark_estimator.set_watermark(start_ts) _LOGGER.info( '[Poll] Watermark=%s (start_ts), cursor=%s (end_ts)', _utc(start_ts), _utc(end_ts)) for chunk_start, chunk_end in ranges: yield TimestampedValue( _QueryRange(chunk_start=chunk_start, chunk_end=chunk_end), start_ts) @beam.DoFn.unbounded_per_element() def process( self, _: _PollConfig, restriction_tracker=beam.DoFn.RestrictionParam(), watermark_estimator=beam.DoFn.WatermarkEstimatorParam( _PollWatermarkEstimatorProvider()) ) -> Iterable[_QueryRange]: now = time.time() start_ts = watermark_estimator.poll_cursor() defer_to = self._next_poll_time(start_ts, now) if defer_to is not None: restriction_tracker.defer_remainder(defer_to) return # Use BQ server time instead of local clock to avoid data loss # from clock skew between the worker VM and BigQuery. bq_now = self._get_bq_timestamp() end_ts = min(bq_now - self._buffer, self._stop_time) _LOGGER.info( '[Poll] Polling: start=%s, end=%s, watermark=%s, ' 'clock_skew=%.3fs', _utc(start_ts), _utc(end_ts), _utc(watermark_estimator.current_watermark()), float(bq_now) - now) current_index = restriction_tracker.current_restriction().start if not restriction_tracker.try_claim(current_index): return restriction_tracker.defer_remainder(Timestamp.of(now) + self._poll_interval) yield from self._emit_query_ranges(start_ts, end_ts, watermark_estimator) class _ExecuteQueryFn(beam.DoFn): """Executes a BQ CHANGES/APPENDS query from a _QueryRange instruction. """ def __init__( self, table: str, project: str, change_function: str, temp_dataset: str, location: Optional[str], change_type_column: str = 'change_type', change_timestamp_column: str = 'change_timestamp', columns: Optional[list[str]] = None, row_filter: Optional[str] = None) -> None: self._table = table self._project = project self._change_function = change_function self._temp_dataset = temp_dataset self._location = location self._change_type_column = change_type_column self._change_timestamp_column = change_timestamp_column self._columns = columns self._row_filter = row_filter def setup(self) -> None: self._bq_wrapper = bigquery_tools.BigQueryWrapper() if self._location is None: table_ref = bigquery_tools.parse_table_reference( self._table, project=self._project) self._location = self._bq_wrapper.get_table_location( table_ref.projectId, table_ref.datasetId, table_ref.tableId) _LOGGER.info( '[Query] Inferred location=%s from source table %s', self._location, self._table) self._bq_wrapper.get_or_create_dataset( self._project, self._temp_dataset, location=self._location, default_table_expiration_ms=_DEFAULT_TABLE_EXPIRATION_MS) def process(self, qr: _QueryRange) -> Iterable[_QueryResult]: """Execute the BQ query described by a _QueryRange and yield _QueryResult. """ sql = build_changes_query( self._table, qr.chunk_start, qr.chunk_end, self._change_function, self._change_type_column, self._change_timestamp_column, self._columns, self._row_filter) temp_table_id = f'beam_ch_temp_{uuid.uuid4().hex[:8]}' job_id = f'beam_ch_{uuid.uuid4().hex[:12]}' _LOGGER.info( '[Query] job_id=%s, temp_table=%s.%s, range=[%s, %s)', job_id, self._temp_dataset, temp_table_id, _utc(qr.chunk_start), _utc(qr.chunk_end)) temp_table_ref = bigquery.TableReference( projectId=self._project, datasetId=self._temp_dataset, tableId=temp_table_id) reference = bigquery.JobReference( jobId=job_id, projectId=self._project, location=self._location) request = bigquery.BigqueryJobsInsertRequest( projectId=self._project, job=bigquery.Job( configuration=bigquery.JobConfiguration( query=bigquery.JobConfigurationQuery( query=sql, useLegacySql=False, destinationTable=temp_table_ref, writeDisposition='WRITE_TRUNCATE', ), ), jobReference=reference)) _LOGGER.info('[Query] Submitting BQ job %s...', job_id) response = self._bq_wrapper._start_job(request) _LOGGER.info('[Query] BQ job %s submitted, waiting...', job_id) self._bq_wrapper.wait_for_bq_job( response.jobReference, sleep_duration_sec=2) _LOGGER.info( '[Query] BQ job %s DONE. Results in %s.%s', job_id, self._temp_dataset, temp_table_id) Metrics.counter('BigQueryChangeHistory', 'queries').inc() yield _QueryResult( temp_table_ref=temp_table_ref, range_start=qr.chunk_start, range_end=qr.chunk_end) class _CDCWatermarkEstimatorProvider(WatermarkEstimatorProvider): """WatermarkEstimatorProvider that initializes the hold from _QueryResult. Uses range_start from the element to set the initial watermark hold. This prevents the runner from advancing the watermark past the data's timestamps before any rows are emitted. """ def initial_estimator_state( self, element: _QueryResult, restriction: _StreamRestriction) -> Timestamp: return element.range_start def create_watermark_estimator( self, estimator_state: Timestamp) -> ManualWatermarkEstimator: return ManualWatermarkEstimator(estimator_state) # ============================================================================= # Read: _ReadStorageStreamsSDF # ============================================================================= class _ReadStorageStreamsSDF(beam.DoFn, beam.transforms.core.RestrictionProvider): """SDF that reads a temp table via BigQuery Storage Read API. Note on SDF lifecycle: the runner decomposes this SDF into three internal wrapper DoFns, each a separately deserialized copy: - Stage A (PairWithRestriction): calls initial_restriction(): no setup() - Stage B (SplitAndSizeRestrictions): calls split(), restriction_size() - Stage C (ProcessSizedElements): calls setup(), then process() Because initial_restriction() runs on a different copy than process(), _ensure_client() lazily creates a gRPC client on whichever copy needs one. The _StreamRestriction carries stream names directly so no shared state is needed between copies. Each element is a _QueryResult pointing to a temp table. Watermark: Uses ManualWatermarkEstimator so the watermark only advances as fast as the change-timestamp values we emit. Emits: Main output: TimestampedValue(row_dict, event_timestamp) Side output (_CLEANUP_TAG): (table_key, (streams_read, total_streams)) """ def __init__( self, batch_arrow_read: bool = True, change_timestamp_column: str = 'change_timestamp', max_split_rounds: int = 1, emit_raw_batches: bool = False) -> None: self._batch_arrow_read = batch_arrow_read self._change_timestamp_column = change_timestamp_column self._max_split_rounds = max_split_rounds self._emit_raw_batches = emit_raw_batches self._storage_client = None def _ensure_client(self) -> None: """Lazily initialize the Storage client. Called from both setup() and initial_restriction() because the runner may invoke initial_restriction on the RestrictionProvider instance before setup() runs (or on a separately deserialized copy). """ if self._storage_client is None: _LOGGER.info('[Read] creating BigQueryReadClient') self._storage_client = bq_storage.BigQueryReadClient() def setup(self) -> None: self._ensure_client() def _split_all_streams( self, stream_names: tuple[str, ...], max_split_rounds: int) -> tuple[str, ...]: """Split each stream at fraction=0.5 for up to max_split_rounds rounds. Each round attempts to split every stream in the current list. A successful split replaces the original stream with primary + remainder. A refused split (both fields empty) keeps the original stream intact. Stops when max_split_rounds is reached or a full round produces zero new splits. BQ's server-side granularity controls how many splits are possible. Small tables may not split at all; large tables may allow multiple rounds of doubling. """ result = list(stream_names) no_split = set() for round_num in range(1, max_split_rounds + 1): new_result = [] made_progress = False for name in result: if name in no_split: new_result.append(name) continue response = self._storage_client.split_read_stream( request=bq_storage.types.SplitReadStreamRequest( name=name, fraction=0.5)) primary = response.primary_stream.name remainder = response.remainder_stream.name if primary and remainder: new_result.extend([primary, remainder]) made_progress = True else: new_result.append(name) no_split.add(name) result = new_result _LOGGER.info( '[Read] _split_all_streams round %d/%d: %d streams ' '(progress=%s)', round_num, max_split_rounds, len(result), made_progress) if not made_progress: break return tuple(result) def initial_restriction(self, element: _QueryResult) -> _StreamRestriction: """Create ReadSession and return _StreamRestriction with stream names. When max_split_rounds > 0, uses SplitReadStream to subdivide each stream at fraction=0.5 for up to max_split_rounds rounds, maximizing parallelism beyond what CreateReadSession provides. """ self._ensure_client() table_key = bigquery_tools.get_hashable_destination(element.temp_table_ref) session = self._create_read_session(element.temp_table_ref) stream_names = tuple(s.name for s in session.streams) original_count = len(stream_names) _LOGGER.info( '[Read] initial_restriction for %s: %d streams from CreateReadSession', table_key, original_count) if self._max_split_rounds > 0: stream_names = self._split_all_streams( stream_names, self._max_split_rounds) _LOGGER.info( '[Read] initial_restriction for %s: %d -> %d streams ' 'after SplitReadStream', table_key, original_count, len(stream_names)) return _StreamRestriction(stream_names, 0, len(stream_names)) def create_tracker( self, restriction: _StreamRestriction) -> _StreamRestrictionTracker: return _StreamRestrictionTracker(restriction) def restriction_size( self, element: _QueryResult, restriction: _StreamRestriction) -> int: return restriction.size() def split(self, element: _QueryResult, restriction: _StreamRestriction) -> Iterable[_StreamRestriction]: """Yield one _StreamRestriction per stream for parallel distribution.""" if restriction.size() <= 1: yield restriction else: for i in range(restriction.start, restriction.stop): yield _StreamRestriction(restriction.stream_names, i, i + 1) def is_bounded(self) -> bool: return True def process( self, element: _QueryResult, restriction_tracker=beam.DoFn.RestrictionParam(), watermark_estimator=beam.DoFn.WatermarkEstimatorParam( _CDCWatermarkEstimatorProvider())): self._ensure_client() table_key = bigquery_tools.get_hashable_destination(element.temp_table_ref) _LOGGER.info( '[Read] Processing %s, range=[%s, %s), ' 'initial watermark=%s', table_key, _utc(element.range_start), _utc(element.range_end), _utc(watermark_estimator.current_watermark())) restriction = restriction_tracker.current_restriction() stream_names = restriction.stream_names total_streams = len(stream_names) streams_read = 0 _LOGGER.info( '[Read] Reading streams [%d, %d) of %d total for %s', restriction.start, restriction.stop, total_streams, table_key) for i in range(restriction.start, restriction.stop): if not restriction_tracker.try_claim(i): _LOGGER.info( '[Read] try_claim(%d) FAILED for %s: ' 'runner split or checkpoint, breaking', i, table_key) break stream_name = stream_names[i] _LOGGER.info( '[Read] try_claim(%d) succeeded: reading stream %s', i, stream_name) stream_rows = 0 if self._emit_raw_batches: stream_batches = 0 for raw_batch in self._read_stream_raw(stream_name): yield TimestampedValue(raw_batch, element.range_start) stream_batches += 1 Metrics.counter('BigQueryChangeHistory', 'batches_emitted').inc(stream_batches) else: for row in self._read_stream(stream_name): ts = row.get(self._change_timestamp_column) if ts is None: raise ValueError( 'Row missing %r column. Row keys: %s' % (self._change_timestamp_column, list(row.keys()))) if isinstance(ts, datetime.datetime): ts = Timestamp.from_utc_datetime(ts) yield TimestampedValue(row, ts) stream_rows += 1 Metrics.counter('BigQueryChangeHistory', 'rows_emitted').inc(stream_rows) streams_read += 1 _LOGGER.info( '[Read] Finished reading stream %d for %s: %d rows', i, table_key, stream_rows) Metrics.counter('BigQueryChangeHistory', 'streams_read').inc() # Advance watermark to range_end after reading all streams. The # initial hold was set to range_start by _CDCWatermarkEstimatorProvider. watermark_estimator.set_watermark(element.range_end) _LOGGER.info( '[Read] Watermark advanced to %s (range_end) for %s', _utc(element.range_end), table_key) # Release the storage client so the gRPC channel doesn't go stale # between process() calls. _ensure_client() will create a fresh one. self._storage_client = None # Emit cleanup signal. Every split that reads at least one stream # reports how many it read. if streams_read > 0: _LOGGER.info( '[Read] Emitting cleanup signal for %s: ' 'streams_read=%d, total_streams=%d', table_key, streams_read, total_streams) yield beam.pvalue.TaggedOutput( _CLEANUP_TAG, (table_key, (streams_read, total_streams))) def _create_read_session(self, table_ref: 'bigquery.TableReference') -> Any: """Create a BigQuery Storage ReadSession for the given table.""" table_path = ( f'projects/{table_ref.projectId}/' f'datasets/{table_ref.datasetId}/' f'tables/{table_ref.tableId}') requested_session = bq_storage.types.ReadSession() requested_session.table = table_path requested_session.data_format = bq_storage.types.DataFormat.ARROW read_options = requested_session.read_options read_options.arrow_serialization_options.buffer_compression = ( bq_storage.types.ArrowSerializationOptions.CompressionCodec.ZSTD) session = self._storage_client.create_read_session( parent=f'projects/{table_ref.projectId}', read_session=requested_session, max_stream_count=_DEFAULT_MAX_STREAMS) _LOGGER.info( '[Read] _create_read_session: table=%s, %d streams', table_path, len(session.streams)) return session def _read_stream(self, stream_name: str) -> Iterable[dict[str, Any]]: """Read all rows from a single Storage API stream as dicts. When batch_arrow_read is enabled, converts entire Arrow RecordBatches at once using to_pylist() instead of calling .as_py() on each cell individually. This is ~1.5x faster for large tables at the cost of ~2x peak memory per batch. """ if self._batch_arrow_read: yield from self._read_stream_batch(stream_name) else: yield from self._read_stream_row_by_row(stream_name) def _read_stream_row_by_row(self, stream_name: str) -> Iterable[dict[str, Any]]: """Row-by-row Arrow conversion (lower memory than batch mode).""" t0 = time.time() row_count = 0 for row in self._storage_client.read_rows(stream_name).rows(): yield dict((item[0], item[1].as_py()) for item in row.items()) row_count += 1 elapsed = time.time() - t0 _LOGGER.info( '[Read] row_by_row: %d rows in %.2fs (%.0f rows/s)', row_count, elapsed, row_count / elapsed if elapsed > 0 else 0) def _read_stream_batch(self, stream_name: str) -> Iterable[dict[str, Any]]: """Batch-convert Arrow RecordBatches for high throughput.""" schema = None row_count = 0 t0 = time.time() for response in self._storage_client.read_rows(stream_name): if schema is None and response.arrow_schema.serialized_schema: schema = pyarrow.ipc.read_schema( pyarrow.py_buffer(response.arrow_schema.serialized_schema)) batch_bytes = response.arrow_record_batch.serialized_record_batch if batch_bytes and schema is not None: batch = pyarrow.ipc.read_record_batch( pyarrow.py_buffer(batch_bytes), schema) yield from batch.to_pylist() row_count += batch.num_rows elapsed = time.time() - t0 _LOGGER.info( '[Read] batch_read: %d rows in %.2fs (%.0f rows/s)', row_count, elapsed, row_count / elapsed if elapsed > 0 else 0) def _read_stream_raw(self, stream_name: str) -> Iterable[tuple[bytes, bytes]]: """Yield raw (schema_bytes, batch_bytes) without decompression. Used when emit_raw_batches is enabled to defer decompression and Arrow-to-Python conversion to a downstream DoFn after reshuffling. Schema bytes are included in each tuple so each batch is self-contained and can be decoded independently. """ schema_bytes = b'' batch_count = 0 t0 = time.time() for response in self._storage_client.read_rows(stream_name): if not schema_bytes and response.arrow_schema.serialized_schema: schema_bytes = bytes(response.arrow_schema.serialized_schema) batch_bytes = response.arrow_record_batch.serialized_record_batch if batch_bytes and schema_bytes: yield (schema_bytes, bytes(batch_bytes)) batch_count += 1 elapsed = time.time() - t0 _LOGGER.info('[Read] raw_read: %d batches in %.2fs', batch_count, elapsed) class _DecompressArrowBatchesFn(beam.DoFn): """Decompress and convert raw Arrow batches to timestamped row dicts. Receives individual (schema_bytes, batch_bytes) tuples after Reshuffle and converts each batch to individual row dicts with event timestamps extracted from the change_timestamp column. """ def __init__(self, change_timestamp_column: str = 'change_timestamp') -> None: self._change_timestamp_column = change_timestamp_column def process(self, element: tuple[bytes, bytes]) -> Iterable[dict[str, Any]]: schema_bytes, batch_bytes = element schema = pyarrow.ipc.read_schema(pyarrow.py_buffer(schema_bytes)) batch = pyarrow.ipc.read_record_batch( pyarrow.py_buffer(batch_bytes), schema) rows = batch.to_pylist() for row in rows: ts = row.get(self._change_timestamp_column) if ts is None: raise ValueError( 'Row missing %r column. Row keys: %s' % (self._change_timestamp_column, list(row.keys()))) if isinstance(ts, datetime.datetime): ts = Timestamp.from_utc_datetime(ts) yield TimestampedValue(row, ts) Metrics.counter('BigQueryChangeHistory', 'rows_emitted').inc(len(rows)) # ============================================================================= # Cleanup: _CleanupTempTablesFn # ============================================================================= class _CleanupTempTablesFn(beam.DoFn): """Stateful DoFn that deletes temp tables after all streams are read. Receives cleanup signals from the Read SDF as: (table_key, (streams_read_count, total_streams)) Accumulates streams_read across all signals for the same table_key. When streams_read >= total_streams, deletes the temp table. The >= (rather than ==) guards against duplicate delivery in at-least-once runners. """ STREAMS_READ = beam.transforms.userstate.CombiningValueStateSpec( 'streams_read', sum) def setup(self) -> None: _LOGGER.info('[Cleanup] setup: creating BigQueryWrapper') self._bq_wrapper = bigquery_tools.BigQueryWrapper() def process( self, element: tuple[str, tuple[int, int]], streams_read=beam.DoFn.StateParam(STREAMS_READ) ) -> None: table_key = element[0] split_count = element[1][0] total_streams = element[1][1] _LOGGER.info( '[Cleanup] Received cleanup signal for %s: ' 'split_count=%d, total_streams=%d', table_key, split_count, total_streams) streams_read.add(split_count) current_read = streams_read.read() _LOGGER.info( '[Cleanup] State for %s: streams_read=%d/%d', table_key, current_read, total_streams) if current_read >= total_streams: parsed = bigquery_tools.parse_table_reference(table_key) _LOGGER.info( '[Cleanup] All streams read: DELETING temp table %s', table_key) self._bq_wrapper._delete_table( parsed.projectId, parsed.datasetId, parsed.tableId) _LOGGER.info('[Cleanup] Deleted temp table %s', table_key) Metrics.counter('BigQueryChangeHistory', 'temp_tables_deleted').inc() streams_read.clear() else: _LOGGER.info( '[Cleanup] Not yet complete for %s (%d/%d), ' 'waiting for more signals', table_key, current_read, total_streams) # ============================================================================= # Public API: ReadBigQueryChangeHistory # ============================================================================= class ReadBigQueryChangeHistory(beam.PTransform): """Streaming source for BigQuery change history. Continuously polls BigQuery APPENDS() or CHANGES() functions and emits changed rows as an unbounded PCollection of dicts. Args: table: BigQuery table to read changes from. Format: 'project:dataset.table' or 'project.dataset.table'. poll_interval_sec: Seconds between polls. Default 60. start_time: Start reading from this timestamp (float, epoch seconds). Default: current time when pipeline starts. stop_time: Stop polling at this timestamp. Default: run forever. change_function: 'CHANGES' or 'APPENDS'. Default 'APPENDS'. buffer_sec: Safety buffer in seconds behind now(). Default 10. BQ does not fail or wait if the query end_ts is less than BQ's CURRENT_TIMESTAMP. This is an extra guardrail to protect against silent data. project: GCP project ID. Default: from pipeline options. temp_dataset: Dataset for temp tables. If None (default), a per-pipeline dataset is auto-created with a 24-hour table expiration as a safety net for orphaned tables. Set this to use an existing dataset (e.g. if your service account lacks bigquery.datasets.create permission). location: BigQuery geographic location for query jobs and temp dataset (e.g. 'US', 'us-central1'). If None (default), inferred from the source table. change_type_column: Output column name for the _CHANGE_TYPE pseudo-column. Default 'change_type'. Change this if your source table already has a column named 'change_type'. change_timestamp_column: Output column name for the _CHANGE_TIMESTAMP pseudo-column. Default 'change_timestamp'. Change this if your source table already has a column named 'change_timestamp'. This column is also used internally to extract event timestamps for watermark tracking. columns: Optional list of column names to select from the source table. If None (default), all columns are selected. The pseudo-columns (change_type, change_timestamp) are always included regardless of this setting. row_filter: Optional SQL boolean expression used as a WHERE clause on the CHANGES/APPENDS query. Do not include the WHERE keyword. Example: ``'status = "active" AND region = "US"'``. batch_arrow_read: If True (default), convert Arrow RecordBatches in bulk using to_pylist() instead of per-cell .as_py() calls. This is 1.5x faster for large tables at the cost of ~2x peak memory per RecordBatch. Set to False for minimal memory usage. max_split_rounds: Maximum number of recursive SplitReadStream rounds. Each round splits every stream at fraction=0.5, potentially doubling the stream count (if BQ allows). Default 1 (one round of splitting). Set 0 to disable splitting entirely. Set higher for very large tables where more parallelism is needed. reshuffle_decompress: If True (default), the Read SDF emits raw compressed Arrow batches instead of decoded rows. The batches are reshuffled for fan-out and then decoded in a separate DoFn. This spreads decompression and Arrow-to-Python conversion CPU across more workers. Set to False to decode rows inline within the Read SDF. """ def __init__( self, table: str, poll_interval_sec: float = 60, start_time: Optional[float] = None, stop_time: Optional[float] = None, change_function: str = 'APPENDS', buffer_sec: float = 10, project: Optional[str] = None, temp_dataset: Optional[str] = None, location: Optional[str] = None, change_type_column: str = 'change_type', change_timestamp_column: str = 'change_timestamp', columns: Optional[list[str]] = None, row_filter: Optional[str] = None, batch_arrow_read: bool = True, max_split_rounds: int = 1, reshuffle_decompress: bool = True) -> None: super().__init__() if bq_storage is None: raise ImportError( 'google-cloud-bigquery-storage is required for ' 'ReadBigQueryChangeHistory. Install it with: ' 'pip install google-cloud-bigquery-storage') if pyarrow is None: raise ImportError( 'pyarrow is required for ReadBigQueryChangeHistory. ' 'Install it with: pip install pyarrow') if change_function not in ('CHANGES', 'APPENDS'): raise ValueError( f"change_function must be 'CHANGES' or 'APPENDS', " f"got '{change_function}'") if poll_interval_sec < 15: raise ValueError( f'poll_interval_sec must be >= 15, got {poll_interval_sec}') if buffer_sec < 0: raise ValueError(f'buffer_sec must be >= 0, got {buffer_sec}') self._table = table self._poll_interval_sec = poll_interval_sec self._start_time = start_time self._stop_time = stop_time self._change_function = change_function self._buffer_sec = buffer_sec self._project = project self._temp_dataset = temp_dataset self._location = location self._change_type_column = change_type_column self._change_timestamp_column = change_timestamp_column self._columns = columns self._row_filter = row_filter self._batch_arrow_read = batch_arrow_read self._max_split_rounds = max_split_rounds self._reshuffle_decompress = reshuffle_decompress def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection: project = self._project if project is None: project = pbegin.pipeline.options.view_as( beam.options.pipeline_options.GoogleCloudOptions).project if project is None: raise ValueError( 'project must be specified either in ReadBigQueryChangeHistory ' 'or in pipeline options (--project)') start_time = Timestamp(self._start_time or time.time()) stop_time = ( Timestamp(self._stop_time) if self._stop_time is not None else MAX_TIMESTAMP) buffer = Duration(seconds=self._buffer_sec) poll_interval = Duration(seconds=self._poll_interval_sec) temp_dataset = self._temp_dataset if temp_dataset is None: temp_dataset = f'beam_ch_temp_{uuid.uuid4().hex[:12]}' _LOGGER.info( '[ReadBigQueryChangeHistory] expand: table=%s, project=%s, ' 'change_function=%s, poll_interval=%d sec, buffer=%d sec, ' 'temp_dataset=%s, start_time=%s, stop_time=%s', self._table, project, self._change_function, self._poll_interval_sec, self._buffer_sec, temp_dataset, _utc(start_time), _utc(stop_time) if stop_time != MAX_TIMESTAMP else 'INF') # Custom polling SDF emits lightweight _QueryRange instructions. # The SDF uses defer_remainder() for poll timing and # _PollWatermarkEstimator to hold the watermark at data timestamps. # On the first invocation it handles the full historical range # [start_time, now - buffer) in a single poll. config = _PollConfig(start_time=start_time) query_ranges = ( pbegin | 'CreatePollConfig' >> beam.Create([config]) | 'PollChangeHistory' >> beam.ParDo( _PollChangeHistoryFn( table=self._table, project=project, change_function=self._change_function, buffer=buffer, start_time=start_time, stop_time=stop_time, poll_interval=poll_interval, location=self._location))) # CommitQueryResults: Reshuffle commits _QueryResult (temp table ref) # so that if the Read SDF retries, it re-reads the existing temp table # instead of re-running the BQ query. # Possible edge-case is that if ReadStorageStreams doesn't read the temp # table within 24 hours (table expiration) it can end up in a bad state by # trying to query a non-existing table. query_results = ( query_ranges | 'CommitQueryRanges' >> beam.Reshuffle() | 'ExecuteQueries' >> beam.ParDo( _ExecuteQueryFn( table=self._table, project=project, change_function=self._change_function, temp_dataset=temp_dataset, location=self._location, change_type_column=self._change_type_column, change_timestamp_column=self._change_timestamp_column, columns=self._columns, row_filter=self._row_filter)) | 'CommitQueryResults' >> beam.Reshuffle()) emit_raw = self._reshuffle_decompress read_sdf = beam.ParDo( _ReadStorageStreamsSDF( batch_arrow_read=self._batch_arrow_read, change_timestamp_column=self._change_timestamp_column, max_split_rounds=self._max_split_rounds, emit_raw_batches=emit_raw)) if emit_raw: read_sdf = read_sdf.with_output_types(tuple[bytes, bytes]) else: read_sdf = read_sdf.with_output_types(dict[str, Any]) read_outputs = ( query_results | 'ReadStorageStreams' >> read_sdf.with_outputs( _CLEANUP_TAG, main='rows')) _ = ( read_outputs[_CLEANUP_TAG] | 'CleanupTempTables' >> beam.ParDo(_CleanupTempTablesFn())) if emit_raw: # Reshuffle raw Arrow batches for fan-out, then decompress and # convert to timestamped row dicts in a separate DoFn. rows = ( read_outputs['rows'] | 'ReshuffleForFanout' >> beam.Reshuffle() | 'DecompressBatches' >> beam.ParDo( _DecompressArrowBatchesFn( change_timestamp_column=(self._change_timestamp_column)))) return rows else: return read_outputs['rows'] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_change_history_it_test.py0000644000175100017510000005477415206113032026100 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration tests for BigQuery change history streaming source. """ import logging import secrets import time import unittest import uuid import apache_beam as beam from apache_beam.io.gcp import bigquery_tools from apache_beam.io.gcp.bigquery_change_history import ReadBigQueryChangeHistory from apache_beam.io.gcp.bigquery_change_history import _CleanupTempTablesFn from apache_beam.io.gcp.bigquery_change_history import _ExecuteQueryFn from apache_beam.io.gcp.bigquery_change_history import _PollChangeHistoryFn from apache_beam.io.gcp.bigquery_change_history import _PollConfig from apache_beam.io.gcp.bigquery_change_history import _QueryRange from apache_beam.io.gcp.bigquery_change_history import _QueryResult from apache_beam.io.gcp.bigquery_change_history import _ReadStorageStreamsSDF from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import Timestamp _LOGGER = logging.getLogger(__name__) class BigQueryChangeHistoryIntegrationBase(unittest.TestCase): """Base class for integration tests against real BigQuery. Gets project from pipeline options (--project). Creates two unique temp datasets per test class: - dataset: for change-history-enabled source tables - temp_dataset: for pipeline temp tables (query results, etc.) Both are deleted with all contents in tearDownClass. """ @classmethod def setUpClass(cls): cls.test_pipeline = TestPipeline(is_integration_test=True) cls.project = cls.test_pipeline.get_option('project') cls.args = cls.test_pipeline.get_full_options_as_args() cls.bq_wrapper = BigQueryWrapper() suffix = secrets.token_hex(4) cls.dataset = f'beam_ch_src_{suffix}' cls.temp_dataset = f'beam_ch_tmp_{suffix}' cls.bq_wrapper.get_or_create_dataset(cls.project, cls.dataset) ds = cls.bq_wrapper.client.datasets.Get( bigquery.BigqueryDatasetsGetRequest( projectId=cls.project, datasetId=cls.dataset)) cls.location = ds.location cls.bq_wrapper.get_or_create_dataset( cls.project, cls.temp_dataset, location=cls.location) _LOGGER.info( 'Created datasets: source=%s, temp=%s (location=%s)', cls.dataset, cls.temp_dataset, cls.location) @classmethod def tearDownClass(cls): for dataset in (cls.dataset, cls.temp_dataset): try: cls.bq_wrapper.client.datasets.Delete( bigquery.BigqueryDatasetsDeleteRequest( projectId=cls.project, datasetId=dataset, deleteContents=True)) _LOGGER.info('Deleted dataset %s', dataset) except Exception as e: _LOGGER.warning('Failed to clean up dataset %s: %s', dataset, e) @classmethod def _create_temp_table_with_data(cls, table_id, rows, schema=None): """Create a table in the temp dataset and insert rows via streaming.""" if schema is None: schema = [('id', 'INTEGER'), ('name', 'STRING'), ('value', 'FLOAT')] table_schema = bigquery.TableSchema() for field_name, field_type in schema: field = bigquery.TableFieldSchema() field.name = field_name field.type = field_type table_schema.fields.append(field) table = bigquery.Table( tableReference=bigquery.TableReference( projectId=cls.project, datasetId=cls.temp_dataset, tableId=table_id), schema=table_schema) request = bigquery.BigqueryTablesInsertRequest( projectId=cls.project, datasetId=cls.temp_dataset, table=table) cls.bq_wrapper.client.tables.Insert(request) # Wait for table to be visible cls.bq_wrapper.get_table(cls.project, cls.temp_dataset, table_id) if rows: cls.bq_wrapper.insert_rows(cls.project, cls.temp_dataset, table_id, rows) # Give streaming buffer time to flush time.sleep(5) return bigquery.TableReference( projectId=cls.project, datasetId=cls.temp_dataset, tableId=table_id) @classmethod def _create_change_history_table(cls, table_id, rows=None): """Create a table with enable_change_history via DDL.""" ddl = ( f'CREATE TABLE IF NOT EXISTS ' f'`{cls.project}.{cls.dataset}.{table_id}` ' f'(id INT64, name STRING, value FLOAT64) ' f'OPTIONS (enable_change_history = true)') job_id = f'beam_ch_ddl_{uuid.uuid4().hex[:8]}' reference = bigquery.JobReference(jobId=job_id, projectId=cls.project) request = bigquery.BigqueryJobsInsertRequest( projectId=cls.project, job=bigquery.Job( configuration=bigquery.JobConfiguration( query=bigquery.JobConfigurationQuery( query=ddl, useLegacySql=False)), jobReference=reference)) response = cls.bq_wrapper._start_job(request) cls.bq_wrapper.wait_for_bq_job(response.jobReference, sleep_duration_sec=2) # Wait for table to be visible cls.bq_wrapper.get_table(cls.project, cls.dataset, table_id) if rows: cls.bq_wrapper.insert_rows(cls.project, cls.dataset, table_id, rows) time.sleep(5) return bigquery.TableReference( projectId=cls.project, datasetId=cls.dataset, tableId=table_id) @classmethod def _run_dml(cls, sql): """Run a DML statement (INSERT/UPDATE/DELETE) and wait for completion.""" job_id = f'beam_ch_dml_{uuid.uuid4().hex[:8]}' reference = bigquery.JobReference(jobId=job_id, projectId=cls.project) request = bigquery.BigqueryJobsInsertRequest( projectId=cls.project, job=bigquery.Job( configuration=bigquery.JobConfiguration( query=bigquery.JobConfigurationQuery( query=sql, useLegacySql=False)), jobReference=reference)) response = cls.bq_wrapper._start_job(request) cls.bq_wrapper.wait_for_bq_job(response.jobReference, sleep_duration_sec=2) class CleanupTempTablesFnTest(BigQueryChangeHistoryIntegrationBase): """Integration tests for _CleanupTempTablesFn against real BigQuery.""" def test_single_complete_signal_deletes_table(self): """A single signal with streams_read == total deletes the temp table.""" table_id = f'cleanup_test_{secrets.token_hex(4)}' table_ref = self._create_temp_table_with_data( table_id, [{ 'id': 1, 'name': 'a', 'value': 1.0 }]) table_key = bigquery_tools.get_hashable_destination(table_ref) # Feed cleanup signal: all 5 streams read out of 5 with beam.Pipeline(argv=self.args) as p: _ = ( p | beam.Create([(table_key, (5, 5))]) | beam.ParDo(_CleanupTempTablesFn())) # Verify table was deleted time.sleep(2) with self.assertRaises(Exception): self.bq_wrapper.get_table(self.project, self.temp_dataset, table_id) def test_partial_signals_then_complete(self): """Partial signals don't delete; final signal triggers cleanup.""" table_id = f'cleanup_partial_{secrets.token_hex(4)}' table_ref = self._create_temp_table_with_data( table_id, [{ 'id': 1, 'name': 'a', 'value': 1.0 }]) table_key = bigquery_tools.get_hashable_destination(table_ref) # Feed two partial signals: 3/10 + 7/10 = 10/10 with beam.Pipeline(argv=self.args) as p: _ = ( p | beam.Create([ (table_key, (3, 10)), (table_key, (7, 10)), ]) | beam.ParDo(_CleanupTempTablesFn())) time.sleep(2) with self.assertRaises(Exception): self.bq_wrapper.get_table(self.project, self.temp_dataset, table_id) class ReadStorageStreamsSDFTest(BigQueryChangeHistoryIntegrationBase): """Integration tests for _ReadStorageStreamsSDF against real BigQuery. Tables must include change_timestamp (TIMESTAMP) and change_type (STRING) columns to match the schema that _ExecuteQueryFn produces in the real pipeline. The Read SDF extracts event timestamps from change_timestamp. """ _READ_SCHEMA = [ ('id', 'INTEGER'), ('name', 'STRING'), ('value', 'FLOAT'), ('change_timestamp', 'TIMESTAMP'), ('change_type', 'STRING'), ] def test_reads_rows_from_temp_table(self): """SDF reads rows from a real temp table via Storage Read API.""" table_id = f'sdf_test_{secrets.token_hex(4)}' now = time.time() now_ts = Timestamp(now) rows = [ { 'id': 1, 'name': 'alice', 'value': 10.0, 'change_timestamp': now, 'change_type': 'INSERT' }, { 'id': 2, 'name': 'bob', 'value': 20.0, 'change_timestamp': now, 'change_type': 'INSERT' }, { 'id': 3, 'name': 'charlie', 'value': 30.0, 'change_timestamp': now, 'change_type': 'INSERT' }, ] table_ref = self._create_temp_table_with_data( table_id, rows, schema=self._READ_SCHEMA) query_result = _QueryResult( temp_table_ref=table_ref, range_start=now_ts - 60, range_end=now_ts + 60) with beam.Pipeline(argv=self.args) as p: outputs = ( p | beam.Create([query_result]) | beam.ParDo(_ReadStorageStreamsSDF()).with_outputs( 'cleanup', main='rows')) # Check that we get 3 rows row_count = ( outputs['rows'] | 'CountRows' >> beam.combiners.Count.Globally()) assert_that(row_count, equal_to([3]), label='CheckRowCount') def test_cleanup_signal_emitted(self): """SDF emits cleanup signal with correct counts.""" table_id = f'sdf_cleanup_{secrets.token_hex(4)}' now = time.time() now_ts = Timestamp(now) rows = [{ 'id': 1, 'name': 'a', 'value': 1.0, 'change_timestamp': now, 'change_type': 'INSERT' }] table_ref = self._create_temp_table_with_data( table_id, rows, schema=self._READ_SCHEMA) query_result = _QueryResult( temp_table_ref=table_ref, range_start=now_ts - 60, range_end=now_ts + 60) with beam.Pipeline(argv=self.args) as p: outputs = ( p | beam.Create([query_result]) | beam.ParDo(_ReadStorageStreamsSDF()).with_outputs( 'cleanup', main='rows')) # Verify cleanup signal cleanup_table_keys = ( outputs['cleanup'] | 'ExtractKey' >> beam.Map(lambda x: x[0])) assert_that( cleanup_table_keys, equal_to([bigquery_tools.get_hashable_destination(table_ref)]), label='CheckCleanupKey') def test_empty_table(self): """Empty table produces 0 rows and cleanup signal.""" table_id = f'sdf_empty_{secrets.token_hex(4)}' now_ts = Timestamp(time.time()) table_ref = self._create_temp_table_with_data( table_id, [], schema=self._READ_SCHEMA) query_result = _QueryResult( temp_table_ref=table_ref, range_start=now_ts - 60, range_end=now_ts + 60) with beam.Pipeline(argv=self.args) as p: outputs = ( p | beam.Create([query_result]) | beam.ParDo(_ReadStorageStreamsSDF()).with_outputs( 'cleanup', main='rows')) row_count = ( outputs['rows'] | 'CountRows' >> beam.combiners.Count.Globally()) assert_that(row_count, equal_to([0]), label='CheckZeroRows') class PollChangeHistoryFnTest(BigQueryChangeHistoryIntegrationBase): """Integration test for _PollChangeHistoryFn in isolation.""" def test_poll_emits_query_ranges(self): """Poll SDF emits _QueryRange elements with valid time ranges.""" table_str = f'{self.project}:{self.dataset}.nonexistent' start_time = Timestamp(time.time() - 120) config = _PollConfig(start_time=start_time) poll_sdf = _PollChangeHistoryFn( table=table_str, project=self.project, change_function='APPENDS', buffer=Duration(seconds=0), start_time=start_time, stop_time=Timestamp(time.time() + 5), poll_interval=Duration(seconds=30), location=self.location) with beam.Pipeline(argv=self.args) as p: ranges = (p | beam.Create([config]) | beam.ParDo(poll_sdf)) # assert_that works directly on unbounded PCollections (no GBK). def check_ranges(actual): assert len(actual) >= 1, f'Expected >= 1 range, got {len(actual)}' for r in actual: assert r.chunk_start < r.chunk_end, ( f'Invalid range: {r.chunk_start} >= {r.chunk_end}') assert_that(ranges, check_ranges) class ExecuteQueryFnTest(BigQueryChangeHistoryIntegrationBase): """Integration test for _ExecuteQueryFn in isolation.""" @classmethod def setUpClass(cls): super().setUpClass() cls.test_table_id = f'exec_test_{secrets.token_hex(4)}' cls.test_table_ref = cls._create_change_history_table( cls.test_table_id, rows=[ { 'id': 1, 'name': 'row1', 'value': 1.0 }, { 'id': 2, 'name': 'row2', 'value': 2.0 }, ]) cls.insert_time = time.time() time.sleep(10) def test_execute_query_produces_query_result(self): """ExecuteQueryFn creates a temp table from a _QueryRange.""" table_str = f'{self.project}:{self.dataset}.{self.test_table_id}' query_range = _QueryRange( chunk_start=Timestamp(self.insert_time - 120), chunk_end=Timestamp(time.time())) with beam.Pipeline(argv=self.args) as p: results = ( p | beam.Create([query_range]) | beam.ParDo( _ExecuteQueryFn( table=table_str, project=self.project, change_function='APPENDS', temp_dataset=self.temp_dataset, location=self.location))) result_count = results | beam.combiners.Count.Globally() assert_that(result_count, equal_to([1]), label='CheckOneResult') class ChangesEndToEndTest(BigQueryChangeHistoryIntegrationBase): """End-to-end test using CHANGES function to capture all mutation types. Creates a change-history-enabled table, performs INSERT, UPDATE, and DELETE operations via DML, then reads back via CHANGES to verify all change types appear. """ @classmethod def setUpClass(cls): super().setUpClass() cls.test_table_id = f'e2e_changes_{secrets.token_hex(4)}' fq_table = f'{cls.project}.{cls.dataset}.{cls.test_table_id}' # Create a change-history-enabled table and insert initial rows via DML. # DML inserts (not streaming inserts) are immediately visible and avoid # streaming-buffer flush delays. cls._create_change_history_table(cls.test_table_id) cls.dml_start_time = time.time() cls._run_dml( f"INSERT INTO `{fq_table}` (id, name, value) " f"VALUES (1, 'alice', 10.0), (2, 'bob', 20.0), (3, 'charlie', 30.0)") cls._run_dml(f"UPDATE `{fq_table}` SET value = 25.0 WHERE id = 2") cls._run_dml(f"DELETE FROM `{fq_table}` WHERE id = 3") _LOGGER.info('Waiting for change history propagation...') time.sleep(15) def test_changes_captures_insert_update_delete(self): """ReadBigQueryChangeHistory with CHANGES sees all mutation types.""" table_str = f'{self.project}:{self.dataset}.{self.test_table_id}' start_time = self.dml_start_time - 120 stop_time = time.time() + 15 with beam.Pipeline(argv=self.args) as p: rows = ( p | ReadBigQueryChangeHistory( table=table_str, poll_interval_sec=15, start_time=start_time, stop_time=stop_time, change_function='CHANGES', buffer_sec=10, project=self.project, temp_dataset=self.temp_dataset, location=self.location)) def check_rows(actual): by_type = {} for row in actual: ct = row['change_type'] by_type.setdefault(ct, []).append(row) # BQ CHANGES returns: # INSERT: 3 (original rows) # UPDATE: 1 (bob with new value=25.0) # DELETE: 2 (bob's pre-update row + charlie's explicit delete) inserts = sorted(by_type.get('INSERT', []), key=lambda r: r['id']) assert len(inserts) == 3, ( f'Expected 3 INSERTs, got {len(inserts)}: {inserts}') updates = by_type.get('UPDATE', []) assert len(updates) == 1, ( f'Expected 1 UPDATE, got {len(updates)}: {updates}') assert updates[0]['id'] == 2 and updates[0]['value'] == 25.0, ( f'Unexpected UPDATE row: {updates[0]}') deletes = sorted(by_type.get('DELETE', []), key=lambda r: r['id']) assert len(deletes) == 2, ( f'Expected 2 DELETEs, got {len(deletes)}: {deletes}') delete_ids = {r['id'] for r in deletes} assert delete_ids == { 2, 3 }, (f'Expected DELETE ids {{2, 3}}, got {delete_ids}') assert_that(rows, check_rows) _EXPECTED_ROWS = [ { 'id': 1, 'name': 'alice', 'value': 10.0, 'change_type': 'INSERT' }, { 'id': 2, 'name': 'bob', 'value': 20.0, 'change_type': 'INSERT' }, { 'id': 3, 'name': 'charlie', 'value': 30.0, 'change_type': 'INSERT' }, ] class EndToEndTest(BigQueryChangeHistoryIntegrationBase): """End-to-end test using the public ReadBigQueryChangeHistory API. Creates a change-history-enabled table, inserts rows, then runs the full pipeline via the public PTransform and verifies rows come through. """ @classmethod def setUpClass(cls): super().setUpClass() cls.test_table_id = f'e2e_test_{secrets.token_hex(4)}' cls.test_table_ref = cls._create_change_history_table( cls.test_table_id, rows=[ { 'id': 1, 'name': 'alice', 'value': 10.0 }, { 'id': 2, 'name': 'bob', 'value': 20.0 }, { 'id': 3, 'name': 'charlie', 'value': 30.0 }, ]) cls.insert_time = time.time() # Wait for streaming buffer + change history propagation _LOGGER.info('Waiting for streaming buffer to flush...') time.sleep(15) def test_public_api_reads_inserted_rows(self): """ReadBigQueryChangeHistory PTransform with polling SDF.""" table_str = f'{self.project}:{self.dataset}.{self.test_table_id}' start_time = self.insert_time - 120 # 2 min before insert stop_time = time.time() + 15 with beam.Pipeline(argv=self.args) as p: rows = ( p | ReadBigQueryChangeHistory( table=table_str, poll_interval_sec=15, start_time=start_time, stop_time=stop_time, change_function='APPENDS', buffer_sec=10, project=self.project, temp_dataset=self.temp_dataset, location=self.location)) def check_rows(actual): assert len(actual) == 3, f'Expected 3 rows, got {len(actual)}' got = sorted([{ k: v for k, v in row.items() if k != 'change_timestamp' } for row in actual], key=lambda r: r['id']) assert got == _EXPECTED_ROWS, ( f'Row mismatch:\n got: {got}\n expected: ' f'{_EXPECTED_ROWS}') assert_that(rows, check_rows) def test_public_api_reads_inserted_row_inline_decompress(self): """ReadBigQueryChangeHistory with inline decompression (no reshuffle).""" table_str = f'{self.project}:{self.dataset}.{self.test_table_id}' start_time = self.insert_time - 120 # 2 min before insert stop_time = time.time() + 15 with beam.Pipeline(argv=self.args) as p: rows = ( p | ReadBigQueryChangeHistory( table=table_str, poll_interval_sec=15, start_time=start_time, stop_time=stop_time, change_function='APPENDS', buffer_sec=10, project=self.project, temp_dataset=self.temp_dataset, location=self.location, reshuffle_decompress=False)) def check_rows(actual): assert len(actual) == 3, f'Expected 3 rows, got {len(actual)}' got = sorted([{ k: v for k, v in row.items() if k != 'change_timestamp' } for row in actual], key=lambda r: r['id']) assert got == _EXPECTED_ROWS, ( f'Row mismatch:\n got: {got}\n expected: ' f'{_EXPECTED_ROWS}') assert_that(rows, check_rows) def test_public_api_reads_inserted_row_with_split_streams(self): """ReadBigQueryChangeHistory with max_split_rounds=1.""" table_str = f'{self.project}:{self.dataset}.{self.test_table_id}' start_time = self.insert_time - 120 # 2 min before insert stop_time = time.time() + 15 with beam.Pipeline(argv=self.args) as p: rows = ( p | ReadBigQueryChangeHistory( table=table_str, poll_interval_sec=15, start_time=start_time, stop_time=stop_time, change_function='APPENDS', buffer_sec=10, project=self.project, temp_dataset=self.temp_dataset, location=self.location, max_split_rounds=10)) def check_rows(actual): assert len(actual) == 3, f'Expected 3 rows, got {len(actual)}' got = sorted([{ k: v for k, v in row.items() if k != 'change_timestamp' } for row in actual], key=lambda r: r['id']) assert got == _EXPECTED_ROWS, ( f'Row mismatch:\n got: {got}\n expected: ' f'{_EXPECTED_ROWS}') assert_that(rows, check_rows) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_change_history_test.py0000644000175100017510000001705515206113032025373 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for BigQuery change history streaming source (no GCP required). Tests: - build_changes_query format - compute_ranges chunking - ReadBigQueryChangeHistory validation """ import datetime import unittest from apache_beam.io.gcp.bigquery_change_history import ReadBigQueryChangeHistory from apache_beam.io.gcp.bigquery_change_history import build_changes_query from apache_beam.io.gcp.bigquery_change_history import compute_ranges from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import Timestamp # Protect against environments where apitools is not available. try: from apitools.base.py.exceptions import HttpError except ImportError: HttpError = None # type: ignore _DAY = Duration(seconds=86400) def _ts(*args, **kwargs) -> Timestamp: """Create a UTC datetime and return a Beam Timestamp.""" # pyrefly: ignore[bad-keyword-argument] dt = datetime.datetime(*args, tzinfo=datetime.timezone.utc, **kwargs) return Timestamp(dt.timestamp()) class BuildChangesQueryTest(unittest.TestCase): """Tests for build_changes_query().""" def test_appends_query_format(self): start = _ts(2025, 1, 1, 0, 0, 0) end = _ts(2025, 1, 1, 1, 0, 0) sql = build_changes_query( 'myproject.mydataset.mytable', start, end, 'APPENDS') self.assertIn('APPENDS', sql) self.assertIn('TABLE `myproject.mydataset.mytable`', sql) self.assertIn('2025-01-01T00:00:00', sql) self.assertIn('2025-01-01T01:00:00', sql) def test_changes_query_format(self): start = _ts(2025, 6, 15, 12, 0, 0) end = _ts(2025, 6, 15, 18, 0, 0) sql = build_changes_query('proj.ds.tbl', start, end, 'CHANGES') self.assertIn('CHANGES', sql) self.assertIn('TABLE `proj.ds.tbl`', sql) def test_columns_select(self): start = _ts(2025, 1, 1) end = _ts(2025, 1, 2) sql = build_changes_query( 'proj.ds.tbl', start, end, 'APPENDS', columns=['col_a', 'col_b']) self.assertIn('SELECT col_a, col_b, _CHANGE_TYPE AS', sql) self.assertNotIn('EXCEPT', sql) def test_columns_none_selects_all(self): start = _ts(2025, 1, 1) end = _ts(2025, 1, 2) sql = build_changes_query( 'proj.ds.tbl', start, end, 'APPENDS', columns=None) self.assertIn('SELECT * EXCEPT', sql) def test_row_filter(self): start = _ts(2025, 1, 1) end = _ts(2025, 1, 2) sql = build_changes_query( 'proj.ds.tbl', start, end, 'APPENDS', row_filter='status = "active"') self.assertIn('WHERE status = "active"', sql) def test_no_row_filter(self): start = _ts(2025, 1, 1) end = _ts(2025, 1, 2) sql = build_changes_query( 'proj.ds.tbl', start, end, 'APPENDS', row_filter=None) self.assertNotIn('WHERE', sql) def test_columns_and_row_filter(self): start = _ts(2025, 1, 1) end = _ts(2025, 1, 2) sql = build_changes_query( 'proj.ds.tbl', start, end, 'CHANGES', columns=['id', 'name'], row_filter='id > 100') self.assertIn('SELECT id, name, _CHANGE_TYPE AS', sql) self.assertNotIn('EXCEPT', sql) self.assertIn('WHERE id > 100', sql) def test_colon_normalized_to_dot(self): start = _ts(2025, 1, 1) end = _ts(2025, 1, 2) sql = build_changes_query( 'myproject:mydataset.mytable', start, end, 'APPENDS') self.assertIn('TABLE `myproject.mydataset.mytable`', sql) # Verify colon in table ref is normalized (timestamps contain colons) table_part = sql.split('TABLE')[1].split(',')[0] self.assertNotIn(':', table_part) def test_microsecond_precision(self): """Verify sub-second precision is preserved in ISO output.""" # 2025-01-01T00:00:00.123456Z start = Timestamp(micros=_ts(2025, 1, 1).micros + 123456) end = start + Duration(seconds=1) sql = build_changes_query('proj.ds.tbl', start, end, 'APPENDS') self.assertIn('2025-01-01T00:00:00.123456Z', sql) class ComputeRangesTest(unittest.TestCase): """Tests for compute_ranges().""" def test_appends_single_range(self): """APPENDS has no chunking — returns single range even for multi-day.""" start = Timestamp(0) end = Timestamp(0) + _DAY * Duration(seconds=5) ranges = compute_ranges(start, end, 'APPENDS') self.assertEqual(len(ranges), 1) self.assertEqual(ranges[0], (start, end)) def test_changes_single_day(self): """CHANGES within 1 day: single range.""" start = Timestamp(0) end = Timestamp(0) + _DAY ranges = compute_ranges(start, end, 'CHANGES') self.assertEqual(len(ranges), 1) self.assertEqual(ranges[0], (start, end)) def test_changes_multi_day(self): """CHANGES spanning 3 days: should chunk into 3 ranges.""" start = Timestamp(0) end = Timestamp(0) + _DAY * Duration(seconds=3) ranges = compute_ranges(start, end, 'CHANGES') self.assertEqual(len(ranges), 3) # Verify no gaps for i in range(len(ranges) - 1): self.assertEqual(ranges[i][1], ranges[i + 1][0]) self.assertEqual(ranges[0][0], start) self.assertEqual(ranges[-1][1], end) def test_changes_partial_day(self): """CHANGES spanning 1.5 days: should chunk into 2 ranges.""" start = Timestamp(0) one_day = Timestamp(0) + _DAY end = Timestamp(micros=_DAY.micros + _DAY.micros // 2) ranges = compute_ranges(start, end, 'CHANGES') self.assertEqual(len(ranges), 2) self.assertEqual(ranges[0], (start, one_day)) self.assertEqual(ranges[1], (one_day, end)) def test_zero_range(self): """end <= start: empty list.""" t100 = Timestamp(micros=100) t50 = Timestamp(micros=50) self.assertEqual(compute_ranges(t100, t100, 'CHANGES'), []) self.assertEqual(compute_ranges(t100, t50, 'CHANGES'), []) self.assertEqual(compute_ranges(t100, t100, 'APPENDS'), []) def test_exact_day_boundary(self): """Exactly 2 days: should produce 2 chunks.""" start = Timestamp(0) end = Timestamp(0) + _DAY * Duration(seconds=2) ranges = compute_ranges(start, end, 'CHANGES') self.assertEqual(len(ranges), 2) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class ValidationTest(unittest.TestCase): """Tests for ReadBigQueryChangeHistory validation.""" def test_invalid_change_function(self): with self.assertRaises(ValueError): ReadBigQueryChangeHistory(table='p:d.t', change_function='INVALID') def test_invalid_poll_interval(self): with self.assertRaises(ValueError): ReadBigQueryChangeHistory(table='p:d.t', poll_interval_sec=0) with self.assertRaises(ValueError): ReadBigQueryChangeHistory(table='p:d.t', poll_interval_sec=-1) def test_default_buffer(self): t = ReadBigQueryChangeHistory(table='p:d.t', change_function='CHANGES') self.assertEqual(t._buffer_sec, 10) t = ReadBigQueryChangeHistory(table='p:d.t', change_function='APPENDS') self.assertEqual(t._buffer_sec, 10) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_file_loads.py0000644000175100017510000014667715206113032023444 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Functionality to perform file loads into BigQuery for Batch and Streaming pipelines. This source is able to work around BigQuery load quotas and limitations. When destinations are dynamic, or when data for a single job is too large, the data will be split into multiple jobs. NOTHING IN THIS FILE HAS BACKWARDS COMPATIBILITY GUARANTEES. """ # pytype: skip-file import hashlib import io import logging import random import time import uuid import apache_beam as beam from apache_beam import pvalue from apache_beam.io import filesystems as fs from apache_beam.io.gcp import bigquery_tools from apache_beam.io.gcp.bigquery_io_metadata import create_bigquery_io_metadata from apache_beam.metrics.metric import Lineage from apache_beam.options import value_provider as vp from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.transforms import trigger from apache_beam.transforms import util from apache_beam.transforms.display import DisplayDataItem from apache_beam.transforms.window import GlobalWindows # Protect against environments where bigquery library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apitools.base.py.exceptions import HttpError except ImportError: pass _LOGGER = logging.getLogger(__name__) ONE_TERABYTE = (1 << 40) # The maximum file size for imports is 5TB. We keep our files under that. _DEFAULT_MAX_FILE_SIZE = 4 * ONE_TERABYTE _DEFAULT_MAX_WRITERS_PER_BUNDLE = 20 # The maximum size for a single load job is one terabyte _MAXIMUM_LOAD_SIZE = 15 * ONE_TERABYTE # Big query only supports up to 10 thousand URIs for a single load job. _MAXIMUM_SOURCE_URIS = 10 * 1000 # If triggering_frequency is supplied, we will trigger the file write after # this many records are written. _FILE_TRIGGERING_RECORD_COUNT = 500000 # If using auto-sharding for unbounded data, we batch the records before # triggering file write to avoid generating too many small files. _FILE_TRIGGERING_BATCHING_DURATION_SECS = 1 # How many seconds we wait before polling a pending job _SLEEP_DURATION_BETWEEN_POLLS = 10 def _generate_job_name(job_name, job_type, step_name): return bigquery_tools.generate_bq_job_name( job_name=job_name, step_id=step_name, job_type=job_type, random=_bq_uuid()) def file_prefix_generator( with_validation=True, pipeline_gcs_location=None, temp_location=None): def _generate_file_prefix(unused_elm): # If a gcs location is provided to the pipeline, then we shall use that. # Otherwise, we shall use the temp_location from pipeline options. gcs_base = pipeline_gcs_location.get() if not gcs_base: gcs_base = temp_location # This will fail at pipeline execution time, but will fail early, as this # step doesn't have any dependencies (and thus will be one of the first # stages to be run). if with_validation and (not gcs_base or not gcs_base.startswith('gs://')): raise ValueError( 'Invalid GCS location: %r.\n' 'Writing to BigQuery with FILE_LOADS method requires a' ' GCS location to be provided to write files to be loaded' ' into BigQuery. Please provide a GCS bucket through' ' custom_gcs_temp_location in the constructor of WriteToBigQuery' ' or the fallback option --temp_location, or pass' ' method="STREAMING_INSERTS" to WriteToBigQuery.' % gcs_base) prefix_uuid = _bq_uuid() return fs.FileSystems.join(gcs_base, 'bq_load', prefix_uuid) return _generate_file_prefix def _make_new_file_writer( file_prefix, destination, file_format, schema=None, schema_side_inputs=tuple()): destination = bigquery_tools.get_hashable_destination(destination) # Windows does not allow : on filenames. Replacing with underscore. # Other disallowed characters are: # https://docs.microsoft.com/en-us/windows/desktop/fileio/naming-a-file destination = destination.replace(':', '.') directory = fs.FileSystems.join(file_prefix, destination) if not fs.FileSystems.exists(directory): try: fs.FileSystems.mkdirs(directory) except IOError: # Concurrent workers may race to create the same directory. # Ignore the IOError if another worker successfully created it. if not fs.FileSystems.exists(directory): raise file_name = str(uuid.uuid4()) file_path = fs.FileSystems.join(file_prefix, destination, file_name) if file_format == bigquery_tools.FileFormat.AVRO: if callable(schema): schema = schema(destination, *schema_side_inputs) elif isinstance(schema, vp.ValueProvider): schema = schema.get() writer = bigquery_tools.AvroRowWriter( fs.FileSystems.create(file_path, "application/avro"), schema) elif file_format == bigquery_tools.FileFormat.JSON: writer = bigquery_tools.JsonRowWriter( fs.FileSystems.create(file_path, "application/text")) else: raise ValueError(( 'Only AVRO and JSON are supported as intermediate formats for ' 'BigQuery WriteRecordsToFile, got: {}.').format(file_format)) return file_path, writer def _bq_uuid(seed=None): if not seed: return str(uuid.uuid4()).replace("-", "") else: return str(hashlib.md5(seed.encode('utf8')).hexdigest()) class _ShardDestinations(beam.DoFn): """Adds a shard number to the key of the KV element. Experimental; no backwards compatibility guarantees.""" DEFAULT_SHARDING_FACTOR = 10 def __init__(self, sharding_factor=DEFAULT_SHARDING_FACTOR): self.sharding_factor = sharding_factor def start_bundle(self): self._shard_count = random.randrange(self.sharding_factor) def process(self, element): destination = element[0] row = element[1] sharded_destination = ( destination, self._shard_count % self.sharding_factor) self._shard_count += 1 yield (sharded_destination, row) class WriteRecordsToFile(beam.DoFn): """Write input records to files before triggering a load job. This transform keeps up to ``max_files_per_bundle`` files open to write to. It receives (destination, record) tuples, and it writes the records to different files for each destination. If there are more than ``max_files_per_bundle`` destinations that we need to write to, then those records are grouped by their destination, and later written to files by ``WriteGroupedRecordsToFile``. It outputs two PCollections. """ UNWRITTEN_RECORD_TAG = 'UnwrittenRecords' WRITTEN_FILE_TAG = 'WrittenFiles' def __init__( self, schema, max_files_per_bundle=_DEFAULT_MAX_WRITERS_PER_BUNDLE, max_file_size=_DEFAULT_MAX_FILE_SIZE, file_format=None): """Initialize a :class:`WriteRecordsToFile`. Args: max_files_per_bundle (int): The maximum number of files that can be kept open during execution of this step in a worker. This is to avoid over- whelming the worker memory. max_file_size (int): The maximum size in bytes for a file to be used in an export job. """ self.schema = schema self.max_files_per_bundle = max_files_per_bundle self.max_file_size = max_file_size self.file_format = file_format or bigquery_tools.FileFormat.JSON def display_data(self): return { 'max_files_per_bundle': self.max_files_per_bundle, 'max_file_size': str(self.max_file_size), 'file_format': self.file_format, } def start_bundle(self): self._destination_to_file_writer = {} def process(self, element, file_prefix, *schema_side_inputs): """Take a tuple with (destination, row) and write to file or spill out. Destination may be a ``TableReference`` or a string, and row is a Python dictionary for a row to be inserted to BigQuery.""" destination = bigquery_tools.get_hashable_destination(element[0]) row = element[1] if destination not in self._destination_to_file_writer: if len(self._destination_to_file_writer) < self.max_files_per_bundle: self._destination_to_file_writer[destination] = _make_new_file_writer( file_prefix, destination, self.file_format, self.schema, schema_side_inputs) else: yield pvalue.TaggedOutput( WriteRecordsToFile.UNWRITTEN_RECORD_TAG, element) return (file_path, writer) = self._destination_to_file_writer[destination] # TODO(pabloem): Is it possible for this to throw exception? writer.write(row) file_size = writer.tell() if file_size > self.max_file_size: writer.close() self._destination_to_file_writer.pop(destination) yield pvalue.TaggedOutput( WriteRecordsToFile.WRITTEN_FILE_TAG, (destination, (file_path, file_size))) def finish_bundle(self): for destination, file_path_writer in \ self._destination_to_file_writer.items(): (file_path, writer) = file_path_writer file_size = writer.tell() writer.close() yield pvalue.TaggedOutput( WriteRecordsToFile.WRITTEN_FILE_TAG, GlobalWindows.windowed_value((destination, (file_path, file_size)))) self._destination_to_file_writer = {} class WriteGroupedRecordsToFile(beam.DoFn): """Receives collection of dest-iterable(records), writes it to files. This is different from ``WriteRecordsToFile`` because it receives records grouped by destination. This means that it's not necessary to keep multiple file descriptors open, because we know for sure when records for a single destination have been written out. Experimental; no backwards compatibility guarantees. """ def __init__( self, schema, max_file_size=_DEFAULT_MAX_FILE_SIZE, file_format=None): self.schema = schema self.max_file_size = max_file_size self.file_format = file_format or bigquery_tools.FileFormat.JSON def process(self, element, file_prefix, *schema_side_inputs): destination = bigquery_tools.get_hashable_destination(element[0]) rows = element[1] file_path, writer = None, None for row in rows: if writer is None: (file_path, writer) = _make_new_file_writer( file_prefix, destination, self.file_format, self.schema, schema_side_inputs) writer.write(row) file_size = writer.tell() if file_size > self.max_file_size: writer.close() yield (destination, (file_path, file_size)) file_path, writer = None, None if writer is not None: writer.close() yield (destination, (file_path, file_size)) class UpdateDestinationSchema(beam.DoFn): """Update destination schema based on data that is about to be copied into it. Unlike load and query jobs, BigQuery copy jobs do not support schema field addition or relaxation on the destination table. This DoFn fills that gap by updating the destination table schemas to be compatible with the data coming from the source table so that schema field modification options are respected regardless of whether data is loaded directly to the destination table or loaded into temporary tables before being copied into the destination. This transform takes as input a (destination, job_reference) pair where the job_reference refers to a completed load job into a temporary table. This transform emits (destination, job_reference) pairs where the job_reference refers to a submitted load job for performing the schema modification in JSON format. Note that the input and output job references are not the same. Experimental; no backwards compatibility guarantees. """ def __init__( self, project=None, write_disposition=None, test_client=None, additional_bq_parameters=None, step_name=None, load_job_project_id=None): self.project = project self._test_client = test_client self._write_disposition = write_disposition self._additional_bq_parameters = additional_bq_parameters or {} self._step_name = step_name self._load_job_project_id = load_job_project_id def start_bundle(self): self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self._test_client) self._bq_io_metadata = create_bigquery_io_metadata(self._step_name) self.pending_jobs = [] def display_data(self): return { 'write_disposition': str(self._write_disposition), 'additional_bq_params': str(self._additional_bq_parameters), } def process(self, element, schema_mod_job_name_prefix): destination = element[0] temp_table_load_job_reference = element[1] if callable(self._additional_bq_parameters): additional_parameters = self._additional_bq_parameters(destination) elif isinstance(self._additional_bq_parameters, vp.ValueProvider): additional_parameters = self._additional_bq_parameters.get() else: additional_parameters = self._additional_bq_parameters # When writing to normal tables WRITE_TRUNCATE will overwrite the schema but # when writing to a partition, care needs to be taken to update the schema # even on WRITE_TRUNCATE. if (self._write_disposition not in ('WRITE_TRUNCATE', 'WRITE_APPEND') or not additional_parameters or not additional_parameters.get("schemaUpdateOptions")): # No need to modify schema of destination table return table_reference = bigquery_tools.parse_table_reference(destination) if table_reference.projectId is None: table_reference.projectId = vp.RuntimeValueProvider.get_value( 'project', str, '') or self.project try: # Check if destination table exists destination_table = self.bq_wrapper.get_table( project_id=table_reference.projectId, dataset_id=table_reference.datasetId, table_id=table_reference.tableId) except HttpError as exn: if exn.status_code == 404: # Destination table does not exist, so no need to modify its schema # ahead of the copy jobs. return else: raise temp_table_load_job = self.bq_wrapper.get_job( project=temp_table_load_job_reference.projectId, job_id=temp_table_load_job_reference.jobId, location=temp_table_load_job_reference.location) temp_table_schema = temp_table_load_job.configuration.load.schema if bigquery_tools.check_schema_equal(temp_table_schema, destination_table.schema, ignore_descriptions=True, ignore_field_order=True): # Destination table schema is already the same as the temp table schema, # so no need to run a job to update the destination table schema. return destination_hash = _bq_uuid( '%s:%s.%s' % ( table_reference.projectId, table_reference.datasetId, table_reference.tableId)) uid = _bq_uuid() job_name = '%s_%s_%s' % (schema_mod_job_name_prefix, destination_hash, uid) _LOGGER.info( 'Triggering schema modification job %s on %s', job_name, table_reference) # Trigger potential schema modification by loading zero rows into the # destination table with the temporary table schema. schema_update_job_reference = self.bq_wrapper.perform_load_job( destination=table_reference, source_stream=io.BytesIO(), # file with zero rows job_id=job_name, schema=temp_table_schema, write_disposition='WRITE_APPEND', create_disposition='CREATE_NEVER', additional_load_parameters=additional_parameters, job_labels=self._bq_io_metadata.add_additional_bq_job_labels(), # JSON format is hardcoded because zero rows load(unlike AVRO) and # a nested schema(unlike CSV, which a default one) is permitted. source_format="NEWLINE_DELIMITED_JSON", load_job_project_id=self._load_job_project_id) self.pending_jobs.append( GlobalWindows.windowed_value( (destination, schema_update_job_reference))) def finish_bundle(self): # Unlike the other steps, schema update is not always necessary. # In that case, return a None value to avoid blocking in streaming context. # Otherwise, the streaming pipeline would get stuck waiting for the # TriggerCopyJobs side-input. if not self.pending_jobs: return [GlobalWindows.windowed_value(None)] for windowed_value in self.pending_jobs: job_ref = windowed_value.value[1] self.bq_wrapper.wait_for_bq_job( job_ref, sleep_duration_sec=_SLEEP_DURATION_BETWEEN_POLLS) return self.pending_jobs class TriggerCopyJobs(beam.DoFn): """Launches jobs to copy from temporary tables into the main target table. When a job needs to write to multiple destination tables, or when a single destination table needs to have multiple load jobs to write to it, files are loaded into temporary tables, and those tables are later copied to the destination tables. This transform emits (destination, job_reference) pairs. TODO(BEAM-7822): In file loads method of writing to BigQuery, copying from temp_tables to destination_table is not atomic. See: https://issues.apache.org/jira/browse/BEAM-7822 """ TRIGGER_DELETE_TEMP_TABLES = 'TriggerDeleteTempTables' def __init__( self, project=None, create_disposition=None, write_disposition=None, test_client=None, step_name=None, load_job_project_id=None): self.project = project self.create_disposition = create_disposition self.write_disposition = write_disposition self.test_client = test_client self._observed_tables = set() self.bq_io_metadata = None self._step_name = step_name self.load_job_project_id = load_job_project_id def display_data(self): return { 'launchesBigQueryJobs': DisplayDataItem( True, label="This Dataflow job launches bigquery jobs.") } def setup(self): self._observed_tables = set() def start_bundle(self): self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self.test_client) if not self.bq_io_metadata: self.bq_io_metadata = create_bigquery_io_metadata(self._step_name) self.pending_jobs = [] def process( self, element_list, job_name_prefix=None, unused_schema_mod_jobs=None): if isinstance(element_list, tuple): # Allow this for streaming update compatibility while fixing BEAM-24535. self.process_one(element_list, job_name_prefix) else: for element in element_list: self.process_one(element, job_name_prefix) def process_one(self, element, job_name_prefix): destination, job_reference = element copy_to_reference = bigquery_tools.parse_table_reference(destination) if copy_to_reference.projectId is None: copy_to_reference.projectId = vp.RuntimeValueProvider.get_value( 'project', str, '') or self.project copy_from_reference = bigquery_tools.parse_table_reference(destination) copy_from_reference.tableId = job_reference.jobId if copy_from_reference.projectId is None: copy_from_reference.projectId = vp.RuntimeValueProvider.get_value( 'project', str, '') or self.project _LOGGER.info( "Triggering copy job from %s to %s", copy_from_reference, copy_to_reference) wait_for_job, write_disposition = ( self._determine_write_disposition(copy_to_reference)) if not self.bq_io_metadata: self.bq_io_metadata = create_bigquery_io_metadata(self._step_name) project_id = ( copy_to_reference.projectId if self.load_job_project_id is None else self.load_job_project_id) copy_job_name = '%s_%s' % ( job_name_prefix, _bq_uuid( '%s:%s.%s' % ( copy_from_reference.projectId, copy_from_reference.datasetId, copy_from_reference.tableId))) job_reference = self.bq_wrapper._insert_copy_job( project_id, copy_job_name, copy_from_reference, copy_to_reference, create_disposition=self.create_disposition, write_disposition=write_disposition, job_labels=self.bq_io_metadata.add_additional_bq_job_labels()) if wait_for_job: self.bq_wrapper.wait_for_bq_job(job_reference, sleep_duration_sec=10) self.pending_jobs.append( GlobalWindows.windowed_value((destination, job_reference))) def _determine_write_disposition(self, copy_to_reference) -> tuple[bool, str]: """ Determines the write disposition for a BigQuery copy job, based on destination. When the write_disposition for a job is WRITE_TRUNCATE, multiple copy jobs to the same destination can interfere with each other, truncate data, and write to the BigQuery table repeatedly. To prevent this, the first copy job runs with the user's specified write_disposition, but subsequent jobs must always use WRITE_APPEND. This ensures that subsequent copy jobs do not clear out data appended by previous jobs. Args: copy_to_reference: The reference to the destination table. Returns: A tuple containing a boolean indicating whether to wait for the job to complete and the write disposition to use for the job. """ full_table_ref = '%s:%s.%s' % ( copy_to_reference.projectId, copy_to_reference.datasetId, copy_to_reference.tableId) if full_table_ref not in self._observed_tables: write_disposition = self.write_disposition wait_for_job = True self._observed_tables.add(full_table_ref) Lineage.sinks().add( 'bigquery', copy_to_reference.projectId, copy_to_reference.datasetId, copy_to_reference.tableId) else: wait_for_job = False write_disposition = 'WRITE_APPEND' return wait_for_job, write_disposition def finish_bundle(self): for windowed_value in self.pending_jobs: job_ref = windowed_value.value[1] self.bq_wrapper.wait_for_bq_job( job_ref, sleep_duration_sec=_SLEEP_DURATION_BETWEEN_POLLS) yield windowed_value yield pvalue.TaggedOutput( TriggerCopyJobs.TRIGGER_DELETE_TEMP_TABLES, GlobalWindows.windowed_value(None)) class TriggerLoadJobs(beam.DoFn): """Triggers the import jobs to BQ. Experimental; no backwards compatibility guarantees. """ TEMP_TABLES = 'TemporaryTables' ONGOING_JOBS = 'OngoingJobs' def __init__( self, schema=None, project=None, create_disposition=None, write_disposition=None, test_client=None, temporary_tables=False, additional_bq_parameters=None, source_format=None, step_name=None, load_job_project_id=None): self.schema = schema self.project = project self.test_client = test_client self.temporary_tables = temporary_tables self.additional_bq_parameters = additional_bq_parameters or {} self.source_format = source_format self.bq_io_metadata = None self._step_name = step_name self.load_job_project_id = load_job_project_id if self.temporary_tables: # If we are loading into temporary tables, we rely on the default create # and write dispositions, which mean that a new table will be created. self.create_disposition = None self.write_disposition = None else: self.create_disposition = create_disposition self.write_disposition = write_disposition def display_data(self): result = { 'create_disposition': str(self.create_disposition), 'write_disposition': str(self.write_disposition), 'additional_bq_params': str(self.additional_bq_parameters), 'schema': str(self.schema), 'launchesBigQueryJobs': DisplayDataItem( True, label="This Dataflow job launches bigquery jobs."), 'source_format': str(self.source_format), } return result def start_bundle(self): self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self.test_client) if not self.bq_io_metadata: self.bq_io_metadata = create_bigquery_io_metadata(self._step_name) self.pending_jobs = [] self.schema_cache = {} def process( self, element, load_job_name_prefix, pane_info=beam.DoFn.PaneInfoParam, *schema_side_inputs): # Each load job is assumed to have files respecting these constraints: # 1. Total size of all files < 15 TB (Max size for load jobs) # 2. Total no. of files in a single load job < 10,000 # This assumption means that there will always be a single load job # triggered for each partition of files. destination = element[0] partition_key, files = element[1] if callable(self.schema): schema = self.schema(destination, *schema_side_inputs) elif isinstance(self.schema, vp.ValueProvider): schema = self.schema.get() else: schema = self.schema if callable(self.additional_bq_parameters): additional_parameters = self.additional_bq_parameters(destination) elif isinstance(self.additional_bq_parameters, vp.ValueProvider): additional_parameters = self.additional_bq_parameters.get() else: additional_parameters = self.additional_bq_parameters table_reference = bigquery_tools.parse_table_reference(destination) if table_reference.projectId is None: table_reference.projectId = vp.RuntimeValueProvider.get_value( 'project', str, '') or self.project # Load jobs for a single destination are always triggered from the same # worker. This means that we can generate a deterministic numbered job id, # and not need to worry. destination_hash = _bq_uuid( '%s:%s.%s' % ( table_reference.projectId, table_reference.datasetId, table_reference.tableId)) job_name = '%s_%s_pane%s_partition%s' % ( load_job_name_prefix, destination_hash, pane_info.index, partition_key) _LOGGER.info('Load job has %s files. Job name is %s.', len(files), job_name) create_disposition = self.create_disposition if self.temporary_tables: # we need to create temp tables, so we need a schema. # if there is no input schema, fetch the destination table's schema if schema is None: hashed_dest = bigquery_tools.get_hashable_destination(table_reference) if hashed_dest in self.schema_cache: schema = self.schema_cache[hashed_dest] else: try: schema = bigquery_tools.table_schema_to_dict( bigquery_tools.BigQueryWrapper().get_table( project_id=table_reference.projectId, dataset_id=table_reference.datasetId, table_id=table_reference.tableId).schema) self.schema_cache[hashed_dest] = schema except Exception as e: _LOGGER.warning( "Input schema is absent and could not fetch the final " "destination table's schema [%s]. Creating temp table [%s] " "will likely fail: %s", hashed_dest, job_name, e) # If we are using temporary tables, then we must always create the # temporary tables, so we replace the create_disposition. create_disposition = 'CREATE_IF_NEEDED' # For temporary tables, we create a new table with the name with JobId. table_reference.tableId = job_name yield pvalue.TaggedOutput( TriggerLoadJobs.TEMP_TABLES, bigquery_tools.get_hashable_destination(table_reference)) else: Lineage.sinks().add( 'bigquery', table_reference.projectId, table_reference.datasetId, table_reference.tableId) _LOGGER.info( 'Triggering job %s to load data to BigQuery table %s.' 'Schema: %s. Additional parameters: %s. Source format: %s', job_name, table_reference, schema, additional_parameters, self.source_format, ) if not self.bq_io_metadata: self.bq_io_metadata = create_bigquery_io_metadata(self._step_name) job_reference = self.bq_wrapper.perform_load_job( destination=table_reference, source_uris=files, job_id=job_name, schema=schema, write_disposition=self.write_disposition, create_disposition=create_disposition, additional_load_parameters=additional_parameters, source_format=self.source_format, job_labels=self.bq_io_metadata.add_additional_bq_job_labels(), load_job_project_id=self.load_job_project_id) yield pvalue.TaggedOutput( TriggerLoadJobs.ONGOING_JOBS, (destination, job_reference)) self.pending_jobs.append( GlobalWindows.windowed_value((destination, job_reference))) def finish_bundle(self): for windowed_value in self.pending_jobs: job_ref = windowed_value.value[1] self.bq_wrapper.wait_for_bq_job( job_ref, sleep_duration_sec=_SLEEP_DURATION_BETWEEN_POLLS) return self.pending_jobs class PartitionFiles(beam.DoFn): MULTIPLE_PARTITIONS_TAG = 'MULTIPLE_PARTITIONS' SINGLE_PARTITION_TAG = 'SINGLE_PARTITION' class Partition(object): def __init__(self, max_size, max_files, files=None, size=0): self.max_size = max_size self.max_files = max_files self.files = files if files is not None else [] self.size = size def can_accept(self, file_size, no_of_files=1): if (((self.size + file_size) <= self.max_size) and ((len(self.files) + no_of_files) <= self.max_files)): return True else: return False def add(self, file_path, file_size): self.files.append(file_path) self.size += file_size def __init__(self, max_partition_size, max_files_per_partition): self.max_partition_size = max_partition_size self.max_files_per_partition = max_files_per_partition def process(self, element): destination = element[0] files = element[1] partitions = [] if not files: _LOGGER.warning( 'Ignoring a BigQuery batch load partition to %s ' 'that contains no source URIs.', destination) return latest_partition = PartitionFiles.Partition( self.max_partition_size, self.max_files_per_partition) for file_path, file_size in files: if latest_partition.can_accept(file_size): latest_partition.add(file_path, file_size) else: partitions.append(latest_partition.files) latest_partition = PartitionFiles.Partition( self.max_partition_size, self.max_files_per_partition) latest_partition.add(file_path, file_size) partitions.append(latest_partition.files) if len(partitions) > 1: output_tag = PartitionFiles.MULTIPLE_PARTITIONS_TAG else: output_tag = PartitionFiles.SINGLE_PARTITION_TAG # we also pass along the index of partition as a key, which is used # to create a deterministic load job name for key, partition in enumerate(partitions): yield pvalue.TaggedOutput(output_tag, (destination, (key, partition))) class DeleteTablesFn(beam.DoFn): def __init__(self, test_client=None): self.test_client = test_client def start_bundle(self): self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self.test_client) def process(self, table_reference): _LOGGER.info("Deleting table %s", table_reference) table_reference = bigquery_tools.parse_table_reference(table_reference) self.bq_wrapper._delete_table( table_reference.projectId, table_reference.datasetId, table_reference.tableId) class BigQueryBatchFileLoads(beam.PTransform): """Takes in a set of elements, and inserts them to BigQuery via batch loads. """ DESTINATION_JOBID_PAIRS = 'destination_load_jobid_pairs' DESTINATION_FILE_PAIRS = 'destination_file_pairs' DESTINATION_COPY_JOBID_PAIRS = 'destination_copy_jobid_pairs' COUNT = 0 def __init__( self, destination, project=None, schema=None, custom_gcs_temp_location=None, create_disposition=None, write_disposition=None, triggering_frequency=None, with_auto_sharding=False, temp_file_format=None, max_file_size=None, max_files_per_bundle=None, max_partition_size=None, max_files_per_partition=None, additional_bq_parameters=None, table_side_inputs=None, schema_side_inputs=None, test_client=None, validate=True, is_streaming_pipeline=False, load_job_project_id=None): self.destination = destination self.project = project self.create_disposition = create_disposition self.write_disposition = write_disposition self.triggering_frequency = triggering_frequency self.with_auto_sharding = with_auto_sharding self.max_file_size = max_file_size or _DEFAULT_MAX_FILE_SIZE self.max_files_per_bundle = ( max_files_per_bundle or _DEFAULT_MAX_WRITERS_PER_BUNDLE) self.max_partition_size = max_partition_size or _MAXIMUM_LOAD_SIZE self.max_files_per_partition = ( max_files_per_partition or _MAXIMUM_SOURCE_URIS) if (isinstance(custom_gcs_temp_location, str) or custom_gcs_temp_location is None): self._custom_gcs_temp_location = vp.StaticValueProvider( str, custom_gcs_temp_location or '') elif isinstance(custom_gcs_temp_location, vp.ValueProvider): self._custom_gcs_temp_location = custom_gcs_temp_location else: raise ValueError('custom_gcs_temp_location must be str or ValueProvider') self.test_client = test_client self.schema = schema self._temp_file_format = temp_file_format or bigquery_tools.FileFormat.JSON # If we have multiple destinations, then we will have multiple load jobs, # thus we will need temporary tables for atomicity. self.dynamic_destinations = bool(callable(destination)) self.additional_bq_parameters = additional_bq_parameters or {} self.table_side_inputs = table_side_inputs or () self.schema_side_inputs = schema_side_inputs or () self.is_streaming_pipeline = is_streaming_pipeline self.load_job_project_id = load_job_project_id self._validate = validate if self._validate: self.verify() def verify(self): if (isinstance(self._custom_gcs_temp_location.get(), vp.StaticValueProvider) and not self._custom_gcs_temp_location.get().startswith('gs://')): # Only fail if the custom location is provided, and it is not a GCS # location. raise ValueError( 'Invalid GCS location: %r.\n' 'Writing to BigQuery with FILE_LOADS method requires a ' 'GCS location to be provided to write files to be ' 'loaded into BigQuery. Please provide a GCS bucket, or ' 'pass method="STREAMING_INSERTS" to WriteToBigQuery.' % self._custom_gcs_temp_location.get()) if self.is_streaming_pipeline and not self.triggering_frequency: raise ValueError( 'triggering_frequency must be specified to use file' 'loads in streaming') elif not self.is_streaming_pipeline and self.triggering_frequency: raise ValueError( 'triggering_frequency can only be used with file' 'loads in streaming') if not self.is_streaming_pipeline and self.with_auto_sharding: return ValueError( 'with_auto_sharding can only be used with file loads in streaming.') def _window_fn(self): """Set the correct WindowInto PTransform""" # The user-supplied triggering_frequency is often chosen to control how # many BigQuery load jobs are triggered, to prevent going over BigQuery's # daily quota for load jobs. If this is set to a large value, currently we # have to buffer all the data until the trigger fires. Instead we ensure # that the files are written if a threshold number of records are ready. # We use only the user-supplied trigger on the actual BigQuery load. # This allows us to offload the data to the filesystem. # # In the case of dynamic sharding, however, we use a default trigger since # the transform performs sharding also batches elements to avoid generating # too many tiny files. User trigger is applied right after writes to limit # the number of load jobs. if self.is_streaming_pipeline and not self.with_auto_sharding: return beam.WindowInto(beam.window.GlobalWindows(), trigger=trigger.Repeatedly( trigger.AfterAny( trigger.AfterProcessingTime( self.triggering_frequency), trigger.AfterCount( _FILE_TRIGGERING_RECORD_COUNT))), accumulation_mode=trigger.AccumulationMode\ .DISCARDING) else: return beam.WindowInto(beam.window.GlobalWindows()) def _maybe_apply_user_trigger(self, destination_file_kv_pc): if self.is_streaming_pipeline: # Apply the user's trigger back before we start triggering load jobs return ( destination_file_kv_pc | "ApplyUserTrigger" >> beam.WindowInto( beam.window.GlobalWindows(), trigger=trigger.Repeatedly( trigger.AfterAll( trigger.AfterProcessingTime(self.triggering_frequency), trigger.AfterCount(1))), accumulation_mode=trigger.AccumulationMode.DISCARDING)) else: return destination_file_kv_pc def _write_files(self, destination_data_kv_pc, file_prefix_pcv): outputs = ( destination_data_kv_pc | beam.ParDo( WriteRecordsToFile( schema=self.schema, max_files_per_bundle=self.max_files_per_bundle, max_file_size=self.max_file_size, file_format=self._temp_file_format), file_prefix_pcv, *self.schema_side_inputs).with_outputs( WriteRecordsToFile.UNWRITTEN_RECORD_TAG, WriteRecordsToFile.WRITTEN_FILE_TAG)) # A PCollection of (destination, file) tuples. It lists files with records, # and the destination each file is meant to be imported into. destination_files_kv_pc = outputs[WriteRecordsToFile.WRITTEN_FILE_TAG] # A PCollection of (destination, record) tuples. These are later sharded, # grouped, and all records for each destination-shard is written to files. # This PCollection is necessary because not all records can be written into # files in ``WriteRecordsToFile``. unwritten_records_pc = outputs[WriteRecordsToFile.UNWRITTEN_RECORD_TAG] more_destination_files_kv_pc = ( unwritten_records_pc | beam.ParDo(_ShardDestinations()) | "GroupShardedRows" >> beam.GroupByKey() | "DropShardNumber" >> beam.Map(lambda x: (x[0][0], x[1])) | "WriteGroupedRecordsToFile" >> beam.ParDo( WriteGroupedRecordsToFile( schema=self.schema, file_format=self._temp_file_format), file_prefix_pcv, *self.schema_side_inputs)) # TODO(https://github.com/apache/beam/issues/20285): Remove the identity # transform. We flatten both PCollection paths and use an identity function # to work around a flatten optimization issue where the wrong coder is # being used. all_destination_file_pairs_pc = ( (destination_files_kv_pc, more_destination_files_kv_pc) | "DestinationFilesUnion" >> beam.Flatten() | "IdentityWorkaround" >> beam.Map(lambda x: x)) return self._maybe_apply_user_trigger(all_destination_file_pairs_pc) def _write_files_with_auto_sharding( self, destination_data_kv_pc, file_prefix_pcv): clock = self.test_client.test_clock if self.test_client else time.time # Auto-sharding is achieved via GroupIntoBatches.WithShardedKey # transform which shards, groups and at the same time batches the table rows # to be inserted to BigQuery. # Firstly, the keys of tagged_data (table references) are converted to a # hashable format. This is needed to work with the keyed states used by. # GroupIntoBatches. After grouping and batching is done, table references # are restored. destination_files_kv_pc = ( destination_data_kv_pc | 'ToHashableTableRef' >> beam.Map(bigquery_tools.to_hashable_table_ref) | 'WithAutoSharding' >> util.GroupIntoBatches.WithShardedKey( batch_size=_FILE_TRIGGERING_RECORD_COUNT, max_buffering_duration_secs=_FILE_TRIGGERING_BATCHING_DURATION_SECS, clock=clock) | 'FromHashableTableRefAndDropShard' >> beam.Map( lambda kvs: (bigquery_tools.parse_table_reference(kvs[0].key), kvs[1])) | beam.ParDo( WriteGroupedRecordsToFile( schema=self.schema, file_format=self._temp_file_format), file_prefix_pcv, *self.schema_side_inputs)) return self._maybe_apply_user_trigger(destination_files_kv_pc) def _load_data( self, partitions_using_temp_tables, partitions_direct_to_destination, load_job_name_pcv, schema_mod_job_name_pcv, copy_job_name_pcv, p, step_name): """Load data to BigQuery Data is loaded into BigQuery in the following two ways: 1. Single partition: When there is a single partition of files destined to a single destination, a single load job is triggered. 2. Multiple partitions and/or Dynamic Destinations: When there are multiple partitions of files destined for a single destination or when Dynamic Destinations are used, multiple load jobs need to be triggered for each partition/destination. Load Jobs are triggered to temporary tables, and those are later copied to the actual appropriate destination table. This ensures atomicity when only some of the load jobs would fail but not other. If any of them fails, then copy jobs are not triggered. """ self.reshuffle_before_load = not p.options.is_compat_version_prior_to( "2.65.0") if self.reshuffle_before_load: # Ensure that TriggerLoadJob retry inputs are deterministic by breaking # fusion for inputs. partitions_using_temp_tables = ( partitions_using_temp_tables | "ReshuffleBeforeLoadWithTempTables" >> beam.Reshuffle()) partitions_direct_to_destination = ( partitions_direct_to_destination | "ReshuffleBeforeLoadWithoutTempTables" >> beam.Reshuffle()) # Load data using temp tables trigger_loads_outputs = ( partitions_using_temp_tables | "TriggerLoadJobsWithTempTables" >> beam.ParDo( TriggerLoadJobs( schema=self.schema, project=self.project, write_disposition=self.write_disposition, create_disposition=self.create_disposition, test_client=self.test_client, temporary_tables=True, additional_bq_parameters=self.additional_bq_parameters, source_format=self._temp_file_format, step_name=step_name, load_job_project_id=self.load_job_project_id), load_job_name_pcv, *self.schema_side_inputs).with_outputs( TriggerLoadJobs.TEMP_TABLES, TriggerLoadJobs.ONGOING_JOBS, main='main')) finished_temp_tables_load_job_ids_pc = trigger_loads_outputs['main'] temp_tables_load_job_ids_pc = trigger_loads_outputs[ TriggerLoadJobs.ONGOING_JOBS] temp_tables_pc = trigger_loads_outputs[TriggerLoadJobs.TEMP_TABLES] schema_mod_job_ids_pc = ( finished_temp_tables_load_job_ids_pc | beam.ParDo( UpdateDestinationSchema( project=self.project, write_disposition=self.write_disposition, test_client=self.test_client, additional_bq_parameters=self.additional_bq_parameters, step_name=step_name, load_job_project_id=self.load_job_project_id), schema_mod_job_name_pcv)) if self.write_disposition in ('WRITE_EMPTY', 'WRITE_TRUNCATE'): # All loads going to the same table must be processed together so that # the truncation happens only once. See # https://github.com/apache/beam/issues/24535. finished_temp_tables_load_job_ids_list_pc = ( finished_temp_tables_load_job_ids_pc | beam.MapTuple( lambda destination, job_reference: ( bigquery_tools.parse_table_reference(destination).tableId, (destination, job_reference))) | beam.GroupByKey() | beam.MapTuple(lambda tableId, batch: list(batch))) else: # Loads can happen in parallel. finished_temp_tables_load_job_ids_list_pc = ( finished_temp_tables_load_job_ids_pc # This name is to ensure update compat. | "Map()" >> beam.Map(lambda x: [x])) copy_job_outputs = ( finished_temp_tables_load_job_ids_list_pc | beam.ParDo( TriggerCopyJobs( project=self.project, create_disposition=self.create_disposition, write_disposition=self.write_disposition, test_client=self.test_client, step_name=step_name, load_job_project_id=self.load_job_project_id), copy_job_name_pcv, pvalue.AsIter(schema_mod_job_ids_pc)).with_outputs( TriggerCopyJobs.TRIGGER_DELETE_TEMP_TABLES, main='main')) destination_copy_job_ids_pc = copy_job_outputs['main'] trigger_delete = copy_job_outputs[ TriggerCopyJobs.TRIGGER_DELETE_TEMP_TABLES] _ = ( temp_tables_pc | "RemoveTempTables/AddUselessValue" >> beam.Map( lambda x, unused_trigger: (x, None), pvalue.AsList(trigger_delete)) | "RemoveTempTables/DeduplicateTables" >> beam.GroupByKey() | "RemoveTempTables/GetTableNames" >> beam.Keys() | "RemoveTempTables/Delete" >> beam.ParDo( DeleteTablesFn(self.test_client))) # Load data directly to destination table destination_load_job_ids_pc = ( partitions_direct_to_destination | "TriggerLoadJobsWithoutTempTables" >> beam.ParDo( TriggerLoadJobs( schema=self.schema, write_disposition=self.write_disposition, create_disposition=self.create_disposition, test_client=self.test_client, temporary_tables=False, additional_bq_parameters=self.additional_bq_parameters, source_format=self._temp_file_format, step_name=step_name, load_job_project_id=self.load_job_project_id), load_job_name_pcv, *self.schema_side_inputs).with_outputs( TriggerLoadJobs.ONGOING_JOBS, main='main') )[TriggerLoadJobs.ONGOING_JOBS] destination_load_job_ids_pc = ( (temp_tables_load_job_ids_pc, destination_load_job_ids_pc) | beam.Flatten()) return destination_load_job_ids_pc, destination_copy_job_ids_pc def expand(self, pcoll): p = pcoll.pipeline self.project = self.project or p.options.view_as(GoogleCloudOptions).project try: step_name = self.label except AttributeError: step_name = 'BigQueryBatchFileLoads_%d' % BigQueryBatchFileLoads.COUNT BigQueryBatchFileLoads.COUNT += 1 temp_location = p.options.view_as(GoogleCloudOptions).temp_location job_name = ( p.options.view_as(GoogleCloudOptions).job_name or 'AUTOMATIC_JOB_NAME') empty_pc = p | "ImpulseEmptyPC" >> beam.Create([]) singleton_pc = p | "ImpulseSingleElementPC" >> beam.Create([None]) load_job_name_pcv = pvalue.AsSingleton( singleton_pc | "LoadJobNamePrefix" >> beam.Map( lambda _: _generate_job_name( job_name, bigquery_tools.BigQueryJobTypes.LOAD, 'LOAD_STEP'))) schema_mod_job_name_pcv = pvalue.AsSingleton( singleton_pc | "SchemaModJobNamePrefix" >> beam.Map( lambda _: _generate_job_name( job_name, bigquery_tools.BigQueryJobTypes.LOAD, 'SCHEMA_MOD_STEP'))) copy_job_name_pcv = pvalue.AsSingleton( singleton_pc | "CopyJobNamePrefix" >> beam.Map( lambda _: _generate_job_name( job_name, bigquery_tools.BigQueryJobTypes.COPY, 'COPY_STEP'))) file_prefix_pcv = pvalue.AsSingleton( singleton_pc | "GenerateFilePrefix" >> beam.Map( file_prefix_generator( self._validate, self._custom_gcs_temp_location, temp_location))) destination_data_kv_pc = ( pcoll | "RewindowIntoGlobal" >> self._window_fn() | "AppendDestination" >> beam.ParDo( bigquery_tools.AppendDestinationsFn(self.destination), *self.table_side_inputs)) if not self.with_auto_sharding: all_destination_file_pairs_pc = self._write_files( destination_data_kv_pc, file_prefix_pcv) else: all_destination_file_pairs_pc = self._write_files_with_auto_sharding( destination_data_kv_pc, file_prefix_pcv) grouped_files_pc = ( all_destination_file_pairs_pc | "GroupFilesByTableDestinations" >> beam.GroupByKey()) partitions = ( grouped_files_pc | beam.ParDo( PartitionFiles( self.max_partition_size, self.max_files_per_partition)).with_outputs( PartitionFiles.MULTIPLE_PARTITIONS_TAG, PartitionFiles.SINGLE_PARTITION_TAG)) multiple_partitions_per_destination_pc = partitions[ PartitionFiles.MULTIPLE_PARTITIONS_TAG] single_partition_per_destination_pc = partitions[ PartitionFiles.SINGLE_PARTITION_TAG] # When using dynamic destinations, elements with both single as well as # multiple partitions are loaded into BigQuery using temporary tables to # ensure atomicity. if self.dynamic_destinations: all_partitions = (( multiple_partitions_per_destination_pc, single_partition_per_destination_pc) | "FlattenPartitions" >> beam.Flatten()) destination_load_job_ids_pc, destination_copy_job_ids_pc = ( self._load_data(all_partitions, empty_pc, load_job_name_pcv, schema_mod_job_name_pcv, copy_job_name_pcv, p, step_name)) else: destination_load_job_ids_pc, destination_copy_job_ids_pc = ( self._load_data(multiple_partitions_per_destination_pc, single_partition_per_destination_pc, load_job_name_pcv, schema_mod_job_name_pcv, copy_job_name_pcv, p, step_name)) return { self.DESTINATION_JOBID_PAIRS: destination_load_job_ids_pc, self.DESTINATION_FILE_PAIRS: all_destination_file_pairs_pc, self.DESTINATION_COPY_JOBID_PAIRS: destination_copy_job_ids_pc, } ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_file_loads_test.py0000644000175100017510000015523715206113032024473 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for BigQuery file loads utilities.""" # pytype: skip-file import logging import os import secrets import time import unittest from unittest.mock import Mock from unittest.mock import call import mock import pytest from hamcrest.core import assert_that as hamcrest_assert from hamcrest.core.core.allof import all_of from hamcrest.core.core.is_ import is_ from parameterized import param from parameterized import parameterized import apache_beam as beam from apache_beam.io.filebasedsink_test import _TestCaseWithTempDirCleanUp from apache_beam.io.gcp import bigquery from apache_beam.io.gcp import bigquery_file_loads as bqfl from apache_beam.io.gcp import bigquery_tools from apache_beam.io.gcp.bigquery import BigQueryDisposition from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.internal.clients import bigquery as bigquery_api from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultMatcher from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultStreamingMatcher from apache_beam.metrics.metric import Lineage from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.runners.dataflow.test_dataflow_runner import TestDataflowRunner from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import combiners from apache_beam.transforms.window import TimestampedValue from apache_beam.typehints.typehints import Tuple from apache_beam.utils import timestamp try: from apitools.base.py.exceptions import HttpError except ImportError: raise unittest.SkipTest('GCP dependencies are not installed') _LOGGER = logging.getLogger(__name__) _DESTINATION_ELEMENT_PAIRS = [ # DESTINATION 1 ('project1:dataset1.table1', { 'name': 'beam', 'language': 'py' }), ('project1:dataset1.table1', { 'name': 'beam', 'language': 'java' }), ('project1:dataset1.table1', { 'name': 'beam', 'language': 'go' }), ('project1:dataset1.table1', { 'name': 'flink', 'language': 'java' }), ('project1:dataset1.table1', { 'name': 'flink', 'language': 'scala' }), # DESTINATION 3 ('project1:dataset1.table3', { 'name': 'spark', 'language': 'scala' }), # DESTINATION 1 ('project1:dataset1.table1', { 'name': 'spark', 'language': 'py' }), ('project1:dataset1.table1', { 'name': 'spark', 'language': 'scala' }), # DESTINATION 2 ('project1:dataset1.table2', { 'name': 'beam', 'foundation': 'apache' }), ('project1:dataset1.table2', { 'name': 'flink', 'foundation': 'apache' }), ('project1:dataset1.table2', { 'name': 'spark', 'foundation': 'apache' }), ] _DISTINCT_DESTINATIONS = list({elm[0] for elm in _DESTINATION_ELEMENT_PAIRS}) _ELEMENTS = [elm[1] for elm in _DESTINATION_ELEMENT_PAIRS] _ELEMENTS_SCHEMA = bigquery.WriteToBigQuery.get_dict_table_schema( bigquery_api.TableSchema( fields=[ bigquery_api.TableFieldSchema( name="name", type="STRING", mode="REQUIRED"), bigquery_api.TableFieldSchema(name="language", type="STRING"), bigquery_api.TableFieldSchema(name="foundation", type="STRING"), ])) class TestWriteRecordsToFile(_TestCaseWithTempDirCleanUp): maxDiff = None def _consume_input(self, fn, checks=None): if checks is None: return with TestPipeline() as p: output_pcs = ( p | beam.Create(_DESTINATION_ELEMENT_PAIRS, reshuffle=False) | beam.ParDo(fn, self.tmpdir).with_outputs( fn.WRITTEN_FILE_TAG, fn.UNWRITTEN_RECORD_TAG)) checks(output_pcs) return output_pcs @parameterized.expand([ param(file_format=bigquery_tools.FileFormat.AVRO), param(file_format=bigquery_tools.FileFormat.JSON), param(file_format=None), ]) def test_files_created(self, file_format): """Test that the files are created and written.""" fn = bqfl.WriteRecordsToFile( schema=_ELEMENTS_SCHEMA, file_format=file_format) self.tmpdir = self._new_tempdir() def check_files_created(output_pcs): dest_file_pc = output_pcs[bqfl.WriteRecordsToFile.WRITTEN_FILE_TAG] files = dest_file_pc | "GetFiles" >> beam.Map(lambda x: x[1][0]) file_count = files | "CountFiles" >> combiners.Count.Globally() _ = files | "FilesExist" >> beam.Map( lambda x: hamcrest_assert(os.path.exists(x), is_(True))) assert_that(file_count, equal_to([3]), label='check file count') destinations = ( dest_file_pc | "GetDests" >> beam.Map(lambda x: bigquery_tools.get_hashable_destination(x[0]))) assert_that( destinations, equal_to(list(_DISTINCT_DESTINATIONS)), label='check destinations ') self._consume_input(fn, check_files_created) def test_many_files(self): """Forces records to be written to many files. For each destination multiple files are necessary. This is because the max file length is very small, so only a couple records fit in each file. """ fn = bqfl.WriteRecordsToFile(schema=_ELEMENTS_SCHEMA, max_file_size=50) self.tmpdir = self._new_tempdir() def check_many_files(output_pcs): dest_file_pc = output_pcs[bqfl.WriteRecordsToFile.WRITTEN_FILE_TAG] files_per_dest = ( dest_file_pc | beam.Map(lambda x: x).with_output_types( beam.typehints.KV[str, Tuple[str, int]]) | combiners.Count.PerKey()) files_per_dest = ( files_per_dest | "GetDests" >> beam.Map( lambda x: (bigquery_tools.get_hashable_destination(x[0]), x[1]))) assert_that( files_per_dest, equal_to([('project1:dataset1.table1', 4), ('project1:dataset1.table2', 2), ('project1:dataset1.table3', 1)])) # Check that the files exist _ = dest_file_pc | beam.Map(lambda x: x[1][0]) | beam.Map( lambda x: hamcrest_assert(os.path.exists(x), is_(True))) self._consume_input(fn, check_many_files) @parameterized.expand([ param(file_format=bigquery_tools.FileFormat.AVRO), param(file_format=bigquery_tools.FileFormat.JSON), ]) def test_records_are_spilled(self, file_format): """Forces records to be written to many files. For each destination multiple files are necessary, and at most two files can be created. This forces records to be spilled to the next stage of processing. """ fn = bqfl.WriteRecordsToFile( schema=_ELEMENTS_SCHEMA, max_files_per_bundle=2, file_format=file_format) self.tmpdir = self._new_tempdir() def check_many_files(output_pcs): dest_file_pc = output_pcs[bqfl.WriteRecordsToFile.WRITTEN_FILE_TAG] spilled_records_pc = output_pcs[ bqfl.WriteRecordsToFile.UNWRITTEN_RECORD_TAG] spilled_records_count = (spilled_records_pc | combiners.Count.Globally()) assert_that(spilled_records_count, equal_to([3]), label='spilled count') files_per_dest = ( dest_file_pc | beam.Map(lambda x: x).with_output_types( beam.typehints.KV[str, Tuple[str, int]]) | combiners.Count.PerKey()) files_per_dest = ( files_per_dest | "GetDests" >> beam.Map( lambda x: (bigquery_tools.get_hashable_destination(x[0]), x[1]))) # Only table1 and table3 get files. table2 records get spilled. assert_that( files_per_dest, equal_to([('project1:dataset1.table1', 1), ('project1:dataset1.table3', 1)]), label='file count') # Check that the files exist _ = dest_file_pc | beam.Map(lambda x: x[1][0]) | beam.Map( lambda x: hamcrest_assert(os.path.exists(x), is_(True))) self._consume_input(fn, check_many_files) class TestWriteGroupedRecordsToFile(_TestCaseWithTempDirCleanUp): def _consume_input(self, fn, input, checks): if checks is None: return with TestPipeline() as p: res = ( p | beam.Create(input) | beam.GroupByKey() | beam.ParDo(fn, self.tmpdir)) checks(res) return res @parameterized.expand([ param(file_format=bigquery_tools.FileFormat.AVRO), param(file_format=bigquery_tools.FileFormat.JSON), param(file_format=None), ]) def test_files_are_created(self, file_format): """Test that the files are created and written.""" fn = bqfl.WriteGroupedRecordsToFile( schema=_ELEMENTS_SCHEMA, file_format=file_format) self.tmpdir = self._new_tempdir() def check_files_created(output_pc): files = output_pc | "GetFiles" >> beam.Map(lambda x: x[1][0]) file_count = files | "CountFiles" >> combiners.Count.Globally() _ = files | "FilesExist" >> beam.Map( lambda x: hamcrest_assert(os.path.exists(x), is_(True))) assert_that(file_count, equal_to([3]), label='check file count') destinations = ( output_pc | "GetDests" >> beam.Map(lambda x: bigquery_tools.get_hashable_destination(x[0]))) assert_that( destinations, equal_to(list(_DISTINCT_DESTINATIONS)), label='check destinations ') self._consume_input(fn, _DESTINATION_ELEMENT_PAIRS, check_files_created) def test_multiple_files(self): """Forces records to be written to many files. For each destination multiple files are necessary. This is because the max file length is very small, so only a couple records fit in each file. """ fn = bqfl.WriteGroupedRecordsToFile( schema=_ELEMENTS_SCHEMA, max_file_size=50) self.tmpdir = self._new_tempdir() def check_multiple_files(output_pc): files_per_dest = output_pc | combiners.Count.PerKey() files_per_dest = ( files_per_dest | "GetDests" >> beam.Map( lambda x: (bigquery_tools.get_hashable_destination(x[0]), x[1]))) assert_that( files_per_dest, equal_to([ ('project1:dataset1.table1', 4), ('project1:dataset1.table2', 2), ('project1:dataset1.table3', 1), ])) # Check that the files exist _ = output_pc | beam.Map(lambda x: x[1][0]) | beam.Map(os.path.exists) self._consume_input(fn, _DESTINATION_ELEMENT_PAIRS, check_multiple_files) class TestPartitionFiles(unittest.TestCase): _ELEMENTS = [( 'destination0', [('file0', 50), ('file1', 50), ('file2', 50), ('file3', 50)]), ('destination1', [('file0', 50), ('file1', 50)])] def test_partition(self): partition = bqfl.PartitionFiles.Partition(1000, 1) self.assertEqual(partition.can_accept(50), True) self.assertEqual(partition.can_accept(2000), False) self.assertEqual(partition.can_accept(1000), True) partition.add('file1', 50) self.assertEqual(partition.files, ['file1']) self.assertEqual(partition.size, 50) self.assertEqual(partition.can_accept(50), False) self.assertEqual(partition.can_accept(0), False) def test_partition_files_dofn_file_split(self): """Force partitions to split based on max_files""" multiple_partitions_result = [('destination0', (0, ['file0', 'file1'])), ('destination0', (1, ['file2', 'file3']))] single_partition_result = [('destination1', (0, ['file0', 'file1']))] with TestPipeline() as p: destination_file_pairs = p | beam.Create(self._ELEMENTS, reshuffle=False) partitioned_files = ( destination_file_pairs | beam.ParDo(bqfl.PartitionFiles(1000, 2)).with_outputs( bqfl.PartitionFiles.MULTIPLE_PARTITIONS_TAG, bqfl.PartitionFiles.SINGLE_PARTITION_TAG)) multiple_partitions = partitioned_files[bqfl.PartitionFiles\ .MULTIPLE_PARTITIONS_TAG] single_partition = partitioned_files[bqfl.PartitionFiles\ .SINGLE_PARTITION_TAG] assert_that( multiple_partitions, equal_to(multiple_partitions_result), label='CheckMultiplePartitions') assert_that( single_partition, equal_to(single_partition_result), label='CheckSinglePartition') def test_partition_files_dofn_size_split(self): """Force partitions to split based on max_partition_size""" multiple_partitions_result = [ ('destination0', (0, ['file0', 'file1', 'file2'])), ('destination0', (1, ['file3'])) ] single_partition_result = [('destination1', (0, ['file0', 'file1']))] with TestPipeline() as p: destination_file_pairs = p | beam.Create(self._ELEMENTS, reshuffle=False) partitioned_files = ( destination_file_pairs | beam.ParDo(bqfl.PartitionFiles(150, 10)).with_outputs( bqfl.PartitionFiles.MULTIPLE_PARTITIONS_TAG, bqfl.PartitionFiles.SINGLE_PARTITION_TAG)) multiple_partitions = partitioned_files[bqfl.PartitionFiles\ .MULTIPLE_PARTITIONS_TAG] single_partition = partitioned_files[bqfl.PartitionFiles\ .SINGLE_PARTITION_TAG] assert_that( multiple_partitions, equal_to(multiple_partitions_result), label='CheckMultiplePartitions') assert_that( single_partition, equal_to(single_partition_result), label='CheckSinglePartition') class TestBigQueryFileLoads(_TestCaseWithTempDirCleanUp): def test_trigger_load_jobs_with_empty_files(self): destination = "project:dataset.table" empty_files = [] load_job_prefix = "test_prefix" with beam.Pipeline() as p: partitions = ( p | beam.Create([(destination, empty_files)]) | beam.ParDo(bqfl.PartitionFiles(1000, 10)).with_outputs( bqfl.PartitionFiles.MULTIPLE_PARTITIONS_TAG, bqfl.PartitionFiles.SINGLE_PARTITION_TAG)) _ = ( partitions[bqfl.PartitionFiles.SINGLE_PARTITION_TAG] | beam.ParDo(bqfl.TriggerLoadJobs(), load_job_prefix)) def test_records_traverse_transform_with_mocks(self): destination = 'project1:dataset1.table1' job_reference = bigquery_api.JobReference() job_reference.projectId = 'project1' job_reference.jobId = 'job_name1' result_job = bigquery_api.Job() result_job.jobReference = job_reference mock_job = mock.Mock() mock_job.status.state = 'DONE' mock_job.status.errorResult = None mock_job.jobReference = job_reference bq_client = mock.Mock() bq_client.jobs.Get.return_value = mock_job bq_client.jobs.Insert.return_value = result_job transform = bqfl.BigQueryBatchFileLoads( destination, custom_gcs_temp_location=self._new_tempdir(), test_client=bq_client, validate=False, temp_file_format=bigquery_tools.FileFormat.JSON) # Need to test this with the FnApiRunner to avoid serializing mocks with TestPipeline('FnApiRunner') as p: outputs = p | beam.Create(_ELEMENTS) | transform dest_files = outputs[bqfl.BigQueryBatchFileLoads.DESTINATION_FILE_PAIRS] dest_job = outputs[bqfl.BigQueryBatchFileLoads.DESTINATION_JOBID_PAIRS] jobs = dest_job | "GetJobs" >> beam.Map(lambda x: x[1]) files = dest_files | "GetFiles" >> beam.Map(lambda x: x[1][0]) destinations = ( dest_files | "GetDests" >> beam.Map( lambda x: (bigquery_tools.get_hashable_destination(x[0]), x[1])) | "GetUniques" >> combiners.Count.PerKey() | "GetFinalDests" >> beam.Keys()) # All files exist _ = ( files | beam.Map(lambda x: hamcrest_assert(os.path.exists(x), is_(True)))) # One file per destination assert_that( files | combiners.Count.Globally(), equal_to([1]), label='CountFiles') assert_that( destinations, equal_to([destination]), label='CheckDestinations') assert_that(jobs, equal_to([job_reference]), label='CheckJobs') @parameterized.expand([ param(compat_version=None), param(compat_version="2.64.0"), ]) @mock.patch( 'apache_beam.coders.coders._should_force_use_dill', return_value=True) def test_reshuffle_before_load(self, mock_force_dill, compat_version): destination = 'project1:dataset1.table1' job_reference = bigquery_api.JobReference() job_reference.projectId = 'project1' job_reference.jobId = 'job_name1' result_job = bigquery_api.Job() result_job.jobReference = job_reference mock_job = mock.Mock() mock_job.status.state = 'DONE' mock_job.status.errorResult = None mock_job.jobReference = job_reference bq_client = mock.Mock() bq_client.jobs.Get.return_value = mock_job bq_client.jobs.Insert.return_value = result_job transform = bqfl.BigQueryBatchFileLoads( destination, custom_gcs_temp_location=self._new_tempdir(), test_client=bq_client, validate=False, temp_file_format=bigquery_tools.FileFormat.JSON) options = PipelineOptions( update_compatibility_version=compat_version, # Disable unrelated compatibility change. force_cloudpickle_deterministic_coders=True) # Need to test this with the DirectRunner to avoid serializing mocks with TestPipeline('DirectRunner', options=options) as p: _ = p | beam.Create(_ELEMENTS) | transform reshuffle_before_load = compat_version is None assert transform.reshuffle_before_load == reshuffle_before_load def test_load_job_id_used(self): job_reference = bigquery_api.JobReference() job_reference.projectId = 'loadJobProject' job_reference.jobId = 'job_name1' result_job = bigquery_api.Job() result_job.jobReference = job_reference mock_job = mock.Mock() mock_job.status.state = 'DONE' mock_job.status.errorResult = None mock_job.jobReference = job_reference bq_client = mock.Mock() bq_client.jobs.Get.return_value = mock_job bq_client.jobs.Insert.return_value = result_job transform = bqfl.BigQueryBatchFileLoads( 'project1:dataset1.table1', custom_gcs_temp_location=self._new_tempdir(), test_client=bq_client, validate=False, load_job_project_id='loadJobProject') # TODO(https://github.com/apache/beam/issues/34549): This test relies on # lineage metrics which Prism doesn't seem to handle correctly. Defaulting # to FnApiRunner instead. with TestPipeline('FnApiRunner') as p: outputs = p | beam.Create(_ELEMENTS) | transform jobs = outputs[bqfl.BigQueryBatchFileLoads.DESTINATION_JOBID_PAIRS] \ | "GetJobs" >> beam.Map(lambda x: x[1]) assert_that(jobs, equal_to([job_reference]), label='CheckJobProjectIds') self.assertSetEqual( Lineage.query(p.result.metrics(), Lineage.SINK), set(["bigquery:project1.dataset1.table1"])) def test_load_job_id_use_for_copy_job(self): destination = 'project1:dataset1.table1' job_reference = bigquery_api.JobReference() job_reference.projectId = 'loadJobProject' job_reference.jobId = 'job_name1' result_job = mock.Mock() result_job.jobReference = job_reference mock_job = mock.Mock() mock_job.status.state = 'DONE' mock_job.status.errorResult = None mock_job.jobReference = job_reference bq_client = mock.Mock() bq_client.jobs.Get.return_value = mock_job bq_client.jobs.Insert.return_value = result_job bq_client.tables.Delete.return_value = None # TODO(https://github.com/apache/beam/issues/34549): This test relies on # lineage metrics which Prism doesn't seem to handle correctly. Defaulting # to FnApiRunner instead. with TestPipeline('FnApiRunner') as p: outputs = ( p | beam.Create(_ELEMENTS, reshuffle=False) | bqfl.BigQueryBatchFileLoads( destination, custom_gcs_temp_location=self._new_tempdir(), test_client=bq_client, validate=False, temp_file_format=bigquery_tools.FileFormat.JSON, max_file_size=45, max_partition_size=80, max_files_per_partition=2, load_job_project_id='loadJobProject')) dest_copy_jobs = outputs[ bqfl.BigQueryBatchFileLoads.DESTINATION_COPY_JOBID_PAIRS] copy_jobs = dest_copy_jobs | "GetCopyJobs" >> beam.Map(lambda x: x[1]) assert_that( copy_jobs, equal_to([ job_reference, job_reference, job_reference, job_reference, job_reference, job_reference ]), label='CheckCopyJobProjectIds') self.assertSetEqual( Lineage.query(p.result.metrics(), Lineage.SINK), set(["bigquery:project1.dataset1.table1"])) @mock.patch('time.sleep') def test_wait_for_load_job_completion(self, sleep_mock): job_1 = bigquery_api.Job() job_1.jobReference = bigquery_api.JobReference() job_1.jobReference.projectId = 'project1' job_1.jobReference.jobId = 'jobId1' job_2 = bigquery_api.Job() job_2.jobReference = bigquery_api.JobReference() job_2.jobReference.projectId = 'project1' job_2.jobReference.jobId = 'jobId2' job_1_waiting = mock.Mock() job_1_waiting.status.state = 'RUNNING' job_2_done = mock.Mock() job_2_done.status.state = 'DONE' job_2_done.status.errorResult = None job_1_done = mock.Mock() job_1_done.status.state = 'DONE' job_1_done.status.errorResult = None bq_client = mock.Mock() bq_client.jobs.Get.side_effect = [ job_1_waiting, job_2_done, job_1_done, job_2_done ] partition_1 = ('project:dataset.table0', (0, ['file0'])) partition_2 = ('project:dataset.table1', (1, ['file1'])) bq_client.jobs.Insert.side_effect = [job_1, job_2] test_job_prefix = "test_job" expected_dest_jobref_list = [(partition_1[0], job_1.jobReference), (partition_2[0], job_2.jobReference)] with TestPipeline('DirectRunner') as p: partitions = p | beam.Create([partition_1, partition_2]) outputs = ( partitions | beam.ParDo( bqfl.TriggerLoadJobs(test_client=bq_client), test_job_prefix)) assert_that(outputs, equal_to(expected_dest_jobref_list)) sleep_mock.assert_called_once() @mock.patch('time.sleep') def test_one_load_job_failed_after_waiting(self, sleep_mock): job_1 = bigquery_api.Job() job_1.jobReference = bigquery_api.JobReference() job_1.jobReference.projectId = 'project1' job_1.jobReference.jobId = 'jobId1' job_2 = bigquery_api.Job() job_2.jobReference = bigquery_api.JobReference() job_2.jobReference.projectId = 'project1' job_2.jobReference.jobId = 'jobId2' job_1_waiting = mock.Mock() job_1_waiting.status.state = 'RUNNING' job_2_done = mock.Mock() job_2_done.status.state = 'DONE' job_2_done.status.errorResult = None job_1_error = mock.Mock() job_1_error.status.state = 'DONE' job_1_error.status.errorResult = 'Some problems happened' bq_client = mock.Mock() bq_client.jobs.Get.side_effect = [ job_1_waiting, job_2_done, job_1_error, job_2_done ] partition_1 = ('project:dataset.table0', (0, ['file0'])) partition_2 = ('project:dataset.table1', (1, ['file1'])) bq_client.jobs.Insert.side_effect = [job_1, job_2] test_job_prefix = "test_job" with self.assertRaises(Exception): with TestPipeline('DirectRunner') as p: partitions = p | beam.Create([partition_1, partition_2]) _ = ( partitions | beam.ParDo( bqfl.TriggerLoadJobs(test_client=bq_client), test_job_prefix)) sleep_mock.assert_called_once() def test_multiple_partition_files(self): destination = 'project1:dataset1.table1' job_reference = bigquery_api.JobReference() job_reference.projectId = 'project1' job_reference.jobId = 'job_name1' result_job = mock.Mock() result_job.jobReference = job_reference mock_job = mock.Mock() mock_job.status.state = 'DONE' mock_job.status.errorResult = None mock_job.jobReference = job_reference bq_client = mock.Mock() bq_client.jobs.Get.return_value = mock_job bq_client.jobs.Insert.return_value = result_job bq_client.tables.Delete.return_value = None # TODO(https://github.com/apache/beam/issues/34549): This test relies on # lineage metrics which Prism doesn't seem to handle correctly. Defaulting # to FnApiRunner instead. with TestPipeline('FnApiRunner') as p: outputs = ( p | beam.Create(_ELEMENTS, reshuffle=False) | bqfl.BigQueryBatchFileLoads( destination, custom_gcs_temp_location=self._new_tempdir(), test_client=bq_client, validate=False, temp_file_format=bigquery_tools.FileFormat.JSON, max_file_size=45, max_partition_size=80, max_files_per_partition=2)) dest_files = outputs[bqfl.BigQueryBatchFileLoads.DESTINATION_FILE_PAIRS] dest_load_jobs = outputs[ bqfl.BigQueryBatchFileLoads.DESTINATION_JOBID_PAIRS] dest_copy_jobs = outputs[ bqfl.BigQueryBatchFileLoads.DESTINATION_COPY_JOBID_PAIRS] load_jobs = dest_load_jobs | "GetLoadJobs" >> beam.Map(lambda x: x[1]) copy_jobs = dest_copy_jobs | "GetCopyJobs" >> beam.Map(lambda x: x[1]) files = dest_files | "GetFiles" >> beam.Map(lambda x: x[1][0]) destinations = ( dest_files | "GetDests" >> beam.Map( lambda x: (bigquery_tools.get_hashable_destination(x[0]), x[1])) | "GetUniques" >> combiners.Count.PerKey() | "GetFinalDests" >> beam.Keys()) # All files exist _ = ( files | beam.Map(lambda x: hamcrest_assert(os.path.exists(x), is_(True)))) # One file per destination assert_that( files | "CountFiles" >> combiners.Count.Globally(), equal_to([6]), label='CheckFileCount') assert_that( destinations, equal_to([destination]), label='CheckDestinations') assert_that( load_jobs | "CountLoadJobs" >> combiners.Count.Globally(), equal_to([6]), label='CheckLoadJobCount') assert_that( copy_jobs | "CountCopyJobs" >> combiners.Count.Globally(), equal_to([6]), label='CheckCopyJobCount') self.assertSetEqual( Lineage.query(p.result.metrics(), Lineage.SINK), set(["bigquery:project1.dataset1.table1"])) @parameterized.expand([ param(write_disposition=BigQueryDisposition.WRITE_TRUNCATE), param(write_disposition=BigQueryDisposition.WRITE_EMPTY) ]) @mock.patch( 'apache_beam.io.gcp.bigquery_file_loads.TriggerCopyJobs.process', wraps=lambda *x: None) def test_multiple_partition_files_write_dispositions( self, mock_call_process, write_disposition): destination = 'project1:dataset1.table1' job_reference = bigquery_api.JobReference() job_reference.projectId = 'project1' job_reference.jobId = 'job_name1' result_job = mock.Mock() result_job.jobReference = job_reference mock_job = mock.Mock() mock_job.status.state = 'DONE' mock_job.status.errorResult = None mock_job.jobReference = job_reference bq_client = mock.Mock() bq_client.jobs.Get.return_value = mock_job bq_client.jobs.Insert.return_value = result_job bq_client.tables.Delete.return_value = None with TestPipeline('DirectRunner') as p: _ = ( p | beam.Create(_ELEMENTS, reshuffle=False) | bqfl.BigQueryBatchFileLoads( destination, custom_gcs_temp_location=self._new_tempdir(), test_client=bq_client, validate=False, temp_file_format=bigquery_tools.FileFormat.JSON, max_file_size=45, max_partition_size=80, max_files_per_partition=2, write_disposition=write_disposition)) # TriggerCopyJob only processes once self.assertEqual(mock_call_process.call_count, 1) @mock.patch( 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper.wait_for_bq_job') @mock.patch( 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._insert_copy_job') @mock.patch( 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._start_job', wraps=BigQueryWrapper._start_job) def test_multiple_identical_destinations_on_write_truncate( self, mock_perform_start_job, mock_insert_copy_job, mock_wait_for_bq_job): """ Test that multiple identical table names, but under different datasets are handled correctly. This essentially means that the `write_disposition` is set to `WRITE_TRUNCATE` for the first job and `WRITE_APPEND` for the rest. Previously this was not the case and all jobs were set to `WRITE_APPEND` from the 2nd table that was named identically with at least one previous table - but from different dataset. """ def dynamic_destination_resolver(element, *side_inputs): """A dynamic destination resolver that returns a destination strictly the same table, but different dataset.""" if element['name'] == 'beam': return 'project1:dataset1.table1' elif element['name'] == 'flink': return 'project1:dataset2.table1' return 'project1:dataset3.table1' job_reference = bigquery_api.JobReference() job_reference.projectId = 'project1' job_reference.jobId = 'job_name1' result_job = mock.Mock() result_job.jobReference = job_reference mock_job = mock.Mock() mock_job.status.state = 'DONE' mock_job.status.errorResult = None mock_job.jobReference = job_reference bq_client = mock.Mock() bq_client.jobs.Get.return_value = mock_job bq_client.jobs.Insert.return_value = result_job bq_client.tables.Delete.return_value = None m = bigquery_tools.BigQueryWrapper(bq_client) m.wait_for_bq_job = mock.Mock() m.wait_for_bq_job.return_value = None mock_jobs = [ Mock(jobReference=bigquery_api.JobReference(jobId=f'job_name{i}')) # Order matters in a sense to prove that jobs with different ids # (`2` & `3`) are run with `WRITE_APPEND` without this current fix. for i in [1, 1, 1, 1, 1] ] mock_perform_start_job.side_effect = mock_jobs # For now we don't care about the return value. mock_insert_copy_job.return_value = None # Pin to FnApiRunner for now to make mocks act appropriately. # TODO(https://github.com/apache/beam/issues/34549) with TestPipeline('FnApiRunner') as p: _ = ( p | beam.Create([ { 'name': 'beam', 'language': 'java' }, { 'name': 'flink', 'language': 'java' }, { 'name': 'beam', 'language': 'java' }, { 'name': 'spark', 'language': 'java' }, { 'name': 'beam', 'language': 'java' }, ], reshuffle=False) | bqfl.BigQueryBatchFileLoads( dynamic_destination_resolver, custom_gcs_temp_location=self._new_tempdir(), test_client=bq_client, validate=False, temp_file_format=bigquery_tools.FileFormat.JSON, max_file_size=45, max_partition_size=80, max_files_per_partition=3, write_disposition=BigQueryDisposition.WRITE_TRUNCATE)) from apache_beam.io.gcp.internal.clients.bigquery import TableReference mock_insert_copy_job.assert_has_calls( [ call( 'project1', mock.ANY, TableReference( datasetId='dataset1', projectId='project1', tableId='job_name1'), TableReference( datasetId='dataset1', projectId='project1', tableId='table1'), create_disposition=None, write_disposition='WRITE_TRUNCATE', job_labels={'step_name': 'bigquerybatchfileloads'}), call( 'project1', mock.ANY, TableReference( datasetId='dataset1', projectId='project1', tableId='job_name1'), TableReference( datasetId='dataset1', projectId='project1', tableId='table1'), create_disposition=None, write_disposition='WRITE_APPEND', job_labels={'step_name': 'bigquerybatchfileloads'}), call( 'project1', mock.ANY, TableReference( datasetId='dataset2', projectId='project1', tableId='job_name1'), TableReference( datasetId='dataset2', projectId='project1', tableId='table1'), create_disposition=None, # Previously this was `WRITE_APPEND`. write_disposition='WRITE_TRUNCATE', job_labels={'step_name': 'bigquerybatchfileloads'}), call( 'project1', mock.ANY, TableReference( datasetId='dataset3', projectId='project1', tableId='job_name1'), TableReference( datasetId='dataset3', projectId='project1', tableId='table1'), create_disposition=None, # Previously this was `WRITE_APPEND`. write_disposition='WRITE_TRUNCATE', job_labels={'step_name': 'bigquerybatchfileloads'}), ], any_order=True) self.assertEqual(4, mock_insert_copy_job.call_count) @parameterized.expand([ param(is_streaming=False, with_auto_sharding=False, compat_version=None), param(is_streaming=True, with_auto_sharding=False, compat_version=None), param(is_streaming=True, with_auto_sharding=True, compat_version=None), param( is_streaming=True, with_auto_sharding=False, compat_version="2.64.0"), param( is_streaming=True, with_auto_sharding=True, compat_version="2.64.0"), ]) @mock.patch( 'apache_beam.coders.coders._should_force_use_dill', return_value=True) def test_triggering_frequency( self, mock_force_dill, is_streaming, with_auto_sharding, compat_version): destination = 'project1:dataset1.table1' job_reference = bigquery_api.JobReference() job_reference.projectId = 'project1' job_reference.jobId = 'job_name1' result_job = bigquery_api.Job() result_job.jobReference = job_reference mock_job = mock.Mock() mock_job.status.state = 'DONE' mock_job.status.errorResult = None mock_job.jobReference = job_reference bq_client = mock.Mock() bq_client.jobs.Get.return_value = mock_job bq_client.jobs.Insert.return_value = result_job # Insert a fake clock to work with auto-sharding which needs a processing # time timer. class _FakeClock(object): def __init__(self, now=time.time()): self._now = now def __call__(self): return self._now start_time = timestamp.Timestamp(0) bq_client.test_clock = _FakeClock(now=start_time) triggering_frequency = 20 if is_streaming else None transform = bqfl.BigQueryBatchFileLoads( destination, custom_gcs_temp_location=self._new_tempdir(), test_client=bq_client, validate=False, temp_file_format=bigquery_tools.FileFormat.JSON, is_streaming_pipeline=is_streaming, triggering_frequency=triggering_frequency, with_auto_sharding=with_auto_sharding) # Need to test this with the DirectRunner to avoid serializing mocks test_options = PipelineOptions( flags=['--allow_unsafe_triggers'], update_compatibility_version=compat_version) test_options.view_as(StandardOptions).streaming = is_streaming with TestPipeline(runner='BundleBasedDirectRunner', options=test_options) as p: if is_streaming: _SIZE = len(_ELEMENTS) fisrt_batch = [ TimestampedValue(value, start_time + i + 1) for i, value in enumerate(_ELEMENTS[:_SIZE // 2]) ] second_batch = [ TimestampedValue(value, start_time + _SIZE // 2 + i + 1) for i, value in enumerate(_ELEMENTS[_SIZE // 2:]) ] # Advance processing time between batches of input elements to fire the # user triggers. Intentionally advance the processing time twice for the # auto-sharding case since we need to first fire the timer and then # fire the trigger. test_stream = ( TestStream().advance_watermark_to(start_time).add_elements( fisrt_batch).advance_processing_time( 30).advance_processing_time(30).add_elements(second_batch). advance_processing_time(30).advance_processing_time( 30).advance_watermark_to_infinity()) input = p | test_stream else: input = p | beam.Create(_ELEMENTS) outputs = input | transform dest_files = outputs[bqfl.BigQueryBatchFileLoads.DESTINATION_FILE_PAIRS] dest_job = outputs[bqfl.BigQueryBatchFileLoads.DESTINATION_JOBID_PAIRS] files = dest_files | "GetFiles" >> beam.Map(lambda x: x[1][0]) destinations = ( dest_files | "GetDests" >> beam.Map( lambda x: (bigquery_tools.get_hashable_destination(x[0]), x[1])) | "GetUniques" >> combiners.Count.PerKey() | "GetFinalDests" >> beam.Keys()) jobs = dest_job | "GetJobs" >> beam.Map(lambda x: x[1]) # Check that all files exist. _ = ( files | beam.Map(lambda x: hamcrest_assert(os.path.exists(x), is_(True)))) # Expect two load jobs are generated in the streaming case due to the # triggering frequency. Grouping is per trigger so we expect two entries # in the output as opposed to one. file_count = files | combiners.Count.Globally().without_defaults() expected_file_count = [1, 1] if is_streaming else [1] expected_destinations = [destination, destination ] if is_streaming else [destination] expected_jobs = [job_reference, job_reference ] if is_streaming else [job_reference] assert_that(file_count, equal_to(expected_file_count), label='CountFiles') assert_that( destinations, equal_to(expected_destinations), label='CheckDestinations') assert_that(jobs, equal_to(expected_jobs), label='CheckJobs') class BigQueryFileLoadsIT(unittest.TestCase): BIG_QUERY_DATASET_ID = 'python_bq_file_loads_' BIG_QUERY_SCHEMA = ( '{"fields": [{"name": "name","type": "STRING"},' '{"name": "language","type": "STRING"}]}') BIG_QUERY_SCHEMA_2 = ( '{"fields": [{"name": "name","type": "STRING"},' '{"name": "foundation","type": "STRING"}]}') BIG_QUERY_STREAMING_SCHEMA = ({ 'fields': [{ 'name': 'Integr', 'type': 'INTEGER', 'mode': 'NULLABLE' }] }) def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.runner_name = type(self.test_pipeline.runner).__name__ self.project = self.test_pipeline.get_option('project') self.dataset_id = '%s%d%s' % ( self.BIG_QUERY_DATASET_ID, int(time.time()), secrets.token_hex(3)) self.bigquery_client = bigquery_tools.BigQueryWrapper() self.bigquery_client.get_or_create_dataset(self.project, self.dataset_id) self.output_table = "%s.output_table" % (self.dataset_id) _LOGGER.info( "Created dataset %s in project %s", self.dataset_id, self.project) @pytest.mark.it_postcommit def test_batch_copy_jobs_with_no_input_schema(self): schema_1 = "col_1:INTEGER" schema_2 = "col_2:INTEGER" # create two tables with different schemas # test to make sure this works with dynamic destinations too self.bigquery_client.get_or_create_table( project_id=self.project, dataset_id=self.dataset_id, table_id="output_table_1", schema=bigquery_tools.get_table_schema_from_string(schema_1), create_disposition='CREATE_IF_NEEDED', write_disposition='WRITE_APPEND') self.bigquery_client.get_or_create_table( project_id=self.project, dataset_id=self.dataset_id, table_id="output_table_2", schema=bigquery_tools.get_table_schema_from_string(schema_2), create_disposition='CREATE_IF_NEEDED', write_disposition='WRITE_APPEND') verifiers = [ BigqueryFullResultMatcher( project=self.project, query="SELECT * FROM %s" % (self.output_table + "_1"), data=[(i, ) for i in range(5)]), BigqueryFullResultMatcher( project=self.project, query="SELECT * FROM %s" % (self.output_table + "_2"), data=[(i, ) for i in range(5, 10)]) ] output = self.output_table def callable_table(el: dict): dest = output if "col_1" in el: dest += "_1" elif "col_2" in el: dest += "_2" return dest args = self.test_pipeline.get_full_options_as_args() with beam.Pipeline(argv=args) as p: # 0...4 going to table 1 # 5...9 going to table 2 items = [{"col_1": i} for i in range(5)] items.extend([{"col_2": i} for i in range(5, 10)]) _ = ( p | beam.Create(items) | bigquery.WriteToBigQuery( table=callable_table, create_disposition="CREATE_NEVER", write_disposition="WRITE_APPEND", # reduce load job size to induce copy jobs max_file_size=10, max_partition_size=20)) hamcrest_assert(p, all_of(*verifiers)) @pytest.mark.it_postcommit def test_multiple_destinations_transform(self): output_table_1 = '%s%s' % (self.output_table, 1) output_table_2 = '%s%s' % (self.output_table, 2) output_table_3 = '%s%s' % (self.output_table, 3) output_table_4 = '%s%s' % (self.output_table, 4) schema1 = bigquery.WriteToBigQuery.get_dict_table_schema( bigquery_tools.parse_table_schema_from_json(self.BIG_QUERY_SCHEMA)) schema2 = bigquery.WriteToBigQuery.get_dict_table_schema( bigquery_tools.parse_table_schema_from_json(self.BIG_QUERY_SCHEMA_2)) schema_kv_pairs = [(output_table_1, schema1), (output_table_2, schema2), (output_table_3, schema1), (output_table_4, schema2)] pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query="SELECT name, language FROM %s" % output_table_1, data=[(d['name'], d['language']) for d in _ELEMENTS if 'language' in d]), BigqueryFullResultMatcher( project=self.project, query="SELECT name, foundation FROM %s" % output_table_2, data=[(d['name'], d['foundation']) for d in _ELEMENTS if 'foundation' in d]), BigqueryFullResultMatcher( project=self.project, query="SELECT name, language FROM %s" % output_table_3, data=[(d['name'], d['language']) for d in _ELEMENTS if 'language' in d]), BigqueryFullResultMatcher( project=self.project, query="SELECT name, foundation FROM %s" % output_table_4, data=[(d['name'], d['foundation']) for d in _ELEMENTS if 'foundation' in d]) ] args = self.test_pipeline.get_full_options_as_args() with beam.Pipeline(argv=args) as p: input = p | beam.Create(_ELEMENTS, reshuffle=False) schema_map_pcv = beam.pvalue.AsDict( p | "MakeSchemas" >> beam.Create(schema_kv_pairs)) table_record_pcv = beam.pvalue.AsDict( p | "MakeTables" >> beam.Create([('table1', output_table_1), ('table2', output_table_2)])) # Get all input in same machine input = ( input | beam.Map(lambda x: (None, x)) | beam.GroupByKey() | beam.FlatMap(lambda elm: elm[1])) _ = ( input | "WriteWithMultipleDestsFreely" >> bigquery.WriteToBigQuery( table=lambda x, tables: (tables['table1'] if 'language' in x else tables['table2']), table_side_inputs=(table_record_pcv, ), schema=lambda dest, schema_map: schema_map.get(dest, None), schema_side_inputs=(schema_map_pcv, ), create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=beam.io.BigQueryDisposition.WRITE_EMPTY)) _ = ( input | "WriteWithMultipleDests" >> bigquery.WriteToBigQuery( table=lambda x: (output_table_3 if 'language' in x else output_table_4), schema=lambda dest, schema_map: schema_map.get(dest, None), schema_side_inputs=(schema_map_pcv, ), create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=beam.io.BigQueryDisposition.WRITE_EMPTY, max_file_size=20, max_files_per_bundle=-1)) hamcrest_assert(p, all_of(*pipeline_verifiers)) @pytest.mark.it_postcommit def test_bqfl_streaming(self): if isinstance(self.test_pipeline.runner, TestDataflowRunner): self.skipTest("TestStream is not supported on TestDataflowRunner") output_table = '%s_%s' % (self.output_table, 'ints') _SIZE = 100 schema = self.BIG_QUERY_STREAMING_SCHEMA l = [{'Integr': i} for i in range(_SIZE)] bq_matcher = BigqueryFullResultStreamingMatcher( project=self.project, query="SELECT Integr FROM %s" % output_table, data=[(i, ) for i in range(100)]) args = self.test_pipeline.get_full_options_as_args( on_success_matcher=bq_matcher, streaming=True, allow_unsafe_triggers=True) with beam.Pipeline(argv=args) as p: stream_source = ( TestStream().advance_watermark_to(0).advance_processing_time( 100).add_elements(l[:_SIZE // 4]). advance_processing_time(100).advance_watermark_to(100).add_elements( l[_SIZE // 4:2 * _SIZE // 4]).advance_processing_time( 100).advance_watermark_to(200).add_elements( l[2 * _SIZE // 4:3 * _SIZE // 4]).advance_processing_time( 100).advance_watermark_to(300).add_elements( l[3 * _SIZE // 4:]).advance_processing_time( 100).advance_watermark_to_infinity()) _ = (p | stream_source | bigquery.WriteToBigQuery(output_table, schema=schema, method=bigquery.WriteToBigQuery \ .Method.FILE_LOADS, triggering_frequency=100)) hamcrest_assert(p, bq_matcher) @pytest.mark.it_postcommit def test_bqfl_streaming_with_copy_jobs(self): if isinstance(self.test_pipeline.runner, TestDataflowRunner): self.skipTest("TestStream is not supported on TestDataflowRunner") output_table = '%s_%s' % (self.output_table, 'with_copy_jobs') _SIZE = 100 schema = self.BIG_QUERY_STREAMING_SCHEMA l = [{'Integr': i} for i in range(_SIZE)] bq_matcher = BigqueryFullResultStreamingMatcher( project=self.project, query="SELECT Integr FROM %s" % output_table, data=[(i, ) for i in range(100)]) args = self.test_pipeline.get_full_options_as_args( on_success_matcher=bq_matcher, streaming=True, allow_unsafe_triggers=True) # Override these parameters to induce copy jobs bqfl._DEFAULT_MAX_FILE_SIZE = 100 bqfl._MAXIMUM_LOAD_SIZE = 200 with beam.Pipeline(argv=args) as p: stream_source = ( TestStream().advance_watermark_to(0).advance_processing_time( 100).add_elements(l[:_SIZE // 4]). advance_processing_time(100).advance_watermark_to(100).add_elements( l[_SIZE // 4:2 * _SIZE // 4]).advance_processing_time( 100).advance_watermark_to(200).add_elements( l[2 * _SIZE // 4:3 * _SIZE // 4]).advance_processing_time( 100).advance_watermark_to(300).add_elements( l[3 * _SIZE // 4:]).advance_processing_time(100). advance_watermark_to_infinity().advance_processing_time(100)) _ = (p | stream_source | bigquery.WriteToBigQuery(output_table, schema=schema, method=bigquery.WriteToBigQuery \ .Method.FILE_LOADS, triggering_frequency=100)) hamcrest_assert(p, bq_matcher) @pytest.mark.it_postcommit def test_bqfl_streaming_with_dynamic_destinations(self): if isinstance(self.test_pipeline.runner, TestDataflowRunner): self.skipTest("TestStream is not supported on TestDataflowRunner") even_table = '%s_%s' % (self.output_table, "dynamic_dest_0") odd_table = '%s_%s' % (self.output_table, "dynamic_dest_1") output_table = lambda row: even_table if ( row['Integr'] % 2 == 0) else odd_table _SIZE = 100 schema = self.BIG_QUERY_STREAMING_SCHEMA l = [{'Integr': i} for i in range(_SIZE)] pipeline_verifiers = [ BigqueryFullResultStreamingMatcher( project=self.project, query="SELECT Integr FROM %s" % even_table, data=[(i, ) for i in range(0, 100, 2)]), BigqueryFullResultStreamingMatcher( project=self.project, query="SELECT Integr FROM %s" % odd_table, data=[(i, ) for i in range(1, 100, 2)]) ] args = self.test_pipeline.get_full_options_as_args( on_success_matcher=all_of(*pipeline_verifiers), streaming=True, allow_unsafe_triggers=True) with beam.Pipeline(argv=args) as p: stream_source = ( TestStream().advance_watermark_to(0).advance_processing_time( 100).add_elements(l[:_SIZE // 4]). advance_processing_time(100).advance_watermark_to(100).add_elements( l[_SIZE // 4:2 * _SIZE // 4]).advance_processing_time( 100).advance_watermark_to(200).add_elements( l[2 * _SIZE // 4:3 * _SIZE // 4]).advance_processing_time( 100).advance_watermark_to(300).add_elements( l[3 * _SIZE // 4:]).advance_processing_time(100). advance_watermark_to_infinity().advance_processing_time(100)) _ = (p | stream_source | bigquery.WriteToBigQuery(output_table, schema=schema, method=bigquery.WriteToBigQuery \ .Method.FILE_LOADS, triggering_frequency=100)) hamcrest_assert(p, all_of(*pipeline_verifiers)) @pytest.mark.it_postcommit def test_one_job_fails_all_jobs_fail(self): # If one of the import jobs fails, then other jobs must not be performed. # This is to avoid reinsertion of some records when a pipeline fails and # is rerun. output_table_1 = '%s%s' % (self.output_table, 1) output_table_2 = '%s%s' % (self.output_table, 2) self.bigquery_client.get_or_create_table( self.project, self.dataset_id, output_table_1.split('.')[1], bigquery_tools.parse_table_schema_from_json(self.BIG_QUERY_SCHEMA), None, None) self.bigquery_client.get_or_create_table( self.project, self.dataset_id, output_table_2.split('.')[1], bigquery_tools.parse_table_schema_from_json(self.BIG_QUERY_SCHEMA_2), None, None) pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query="SELECT name, language FROM %s" % output_table_1, data=[]), BigqueryFullResultMatcher( project=self.project, query="SELECT name, foundation FROM %s" % output_table_2, data=[]) ] args = self.test_pipeline.get_full_options_as_args() with self.assertRaises(Exception): # The pipeline below fails because neither a schema nor SCHEMA_AUTODETECT # are specified. with beam.Pipeline(argv=args) as p: input = p | beam.Create(_ELEMENTS) input2 = p | "Broken record" >> beam.Create(['language_broken_record']) input = (input, input2) | beam.Flatten() _ = ( input | "WriteWithMultipleDests" >> bigquery.WriteToBigQuery( table=lambda x: (output_table_1 if 'language' in x else output_table_2), create_disposition=( beam.io.BigQueryDisposition.CREATE_IF_NEEDED), write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND, temp_file_format=bigquery_tools.FileFormat.JSON)) hamcrest_assert(p, all_of(*pipeline_verifiers)) def tearDown(self): request = bigquery_api.BigqueryDatasetsDeleteRequest( projectId=self.project, datasetId=self.dataset_id, deleteContents=True) try: _LOGGER.info( "Deleting dataset %s in project %s", self.dataset_id, self.project) self.bigquery_client.client.datasets.Delete(request) except HttpError: _LOGGER.debug( 'Failed to clean up dataset %s in project %s', self.dataset_id, self.project) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_geography_it_test.py0000644000175100017510000004377415206113032025055 0ustar00runnerrunner#!/usr/bin/env python # -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration tests for BigQuery GEOGRAPHY data type support.""" import logging import os import secrets import time import unittest import hamcrest as hc import pytest import apache_beam as beam from apache_beam.io.gcp.bigquery import ReadFromBigQuery from apache_beam.io.gcp.bigquery import WriteToBigQuery from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultMatcher from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to try: from apitools.base.py.exceptions import HttpError except ImportError: HttpError = None _LOGGER = logging.getLogger(__name__) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class BigQueryGeographyIntegrationTests(unittest.TestCase): """Integration tests for BigQuery GEOGRAPHY data type.""" BIG_QUERY_DATASET_ID = 'python_geography_it_test_' def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.runner_name = type(self.test_pipeline.runner).__name__ self.project = self.test_pipeline.get_option('project') self.bigquery_client = BigQueryWrapper() self.dataset_id = '%s%d%s' % ( self.BIG_QUERY_DATASET_ID, int(time.time()), secrets.token_hex(3)) self.bigquery_client.get_or_create_dataset(self.project, self.dataset_id) _LOGGER.info( "Created dataset %s in project %s", self.dataset_id, self.project) def tearDown(self): request = bigquery.BigqueryDatasetsDeleteRequest( projectId=self.project, datasetId=self.dataset_id, deleteContents=True) try: _LOGGER.info( "Deleting dataset %s in project %s", self.dataset_id, self.project) self.bigquery_client.client.datasets.Delete(request) except HttpError: _LOGGER.debug( 'Failed to clean up dataset %s in project %s', self.dataset_id, self.project) def create_geography_table(self, table_name, include_repeated=False): """Create a table with various GEOGRAPHY field configurations.""" table_schema = bigquery.TableSchema() # ID field id_field = bigquery.TableFieldSchema() id_field.name = 'id' id_field.type = 'INTEGER' id_field.mode = 'REQUIRED' table_schema.fields.append(id_field) # Required GEOGRAPHY field geo_required = bigquery.TableFieldSchema() geo_required.name = 'location' geo_required.type = 'GEOGRAPHY' geo_required.mode = 'REQUIRED' table_schema.fields.append(geo_required) # Nullable GEOGRAPHY field geo_nullable = bigquery.TableFieldSchema() geo_nullable.name = 'optional_location' geo_nullable.type = 'GEOGRAPHY' geo_nullable.mode = 'NULLABLE' table_schema.fields.append(geo_nullable) if include_repeated: # Repeated GEOGRAPHY field geo_repeated = bigquery.TableFieldSchema() geo_repeated.name = 'path' geo_repeated.type = 'GEOGRAPHY' geo_repeated.mode = 'REPEATED' table_schema.fields.append(geo_repeated) table = bigquery.Table( tableReference=bigquery.TableReference( projectId=self.project, datasetId=self.dataset_id, tableId=table_name), schema=table_schema) request = bigquery.BigqueryTablesInsertRequest( projectId=self.project, datasetId=self.dataset_id, table=table) self.bigquery_client.client.tables.Insert(request) # Wait for table to be available _ = self.bigquery_client.get_table( self.project, self.dataset_id, table_name) @pytest.mark.it_postcommit def test_geography_write_and_read_basic_geometries(self): """Test writing and reading basic GEOGRAPHY geometries.""" table_name = 'geography_basic_geometries' table_id = '{}.{}'.format(self.dataset_id, table_name) # Test data with various WKT geometry types input_data = [ { 'id': 1, 'location': 'POINT(30 10)', 'optional_location': ('POINT(-122.4194 37.7749)') # San Francisco }, { 'id': 2, 'location': 'LINESTRING(30 10, 10 30, 40 40)', 'optional_location': None }, { 'id': 3, 'location': ('POLYGON((30 10, 40 40, 20 40, 10 20, 30 10))'), 'optional_location': ('POLYGON((0 0, 0 1, 1 1, 1 0, 0 0))') }, { 'id': 4, 'location': ('MULTIPOINT((10 40), (40 30), (20 20), (30 10))'), 'optional_location': 'POINT(0 0)' }, { 'id': 5, 'location': ( 'MULTILINESTRING((10 10, 20 20, 10 40), ' '(40 40, 30 30, 40 20, 30 10))'), 'optional_location': None } ] table_schema = { "fields": [{ "name": "id", "type": "INTEGER", "mode": "REQUIRED" }, { "name": "location", "type": "GEOGRAPHY", "mode": "REQUIRED" }, { "name": "optional_location", "type": "GEOGRAPHY", "mode": "NULLABLE" }] } # Write data to BigQuery with TestPipeline(is_integration_test=True) as p: _ = ( p | 'CreateData' >> beam.Create(input_data) | 'WriteToBQ' >> WriteToBigQuery( table=table_id, schema=table_schema, method=WriteToBigQuery.Method.STREAMING_INSERTS, project=self.project)) # Read data back and verify with TestPipeline(is_integration_test=True) as p: result = ( p | 'ReadFromBQ' >> ReadFromBigQuery( table=table_id, project=self.project, method=ReadFromBigQuery.Method.DIRECT_READ) | 'ExtractGeography' >> beam.Map( lambda row: (row['id'], row['location'], row['optional_location']))) expected_data = [ (1, 'POINT(30 10)', 'POINT(-122.4194 37.7749)'), (2, 'LINESTRING(30 10, 10 30, 40 40)', None), ( 3, 'POLYGON((30 10, 40 40, 20 40, 10 20, 30 10))', 'POLYGON((0 0, 0 1, 1 1, 1 0, 0 0))'), (4, 'MULTIPOINT(20 20, 10 40, 40 30, 30 10)', 'POINT(0 0)'), ( 5, 'MULTILINESTRING((10 10, 20 20, 10 40), ' '(40 40, 30 30, 40 20, 30 10))', None) ] assert_that(result, equal_to(expected_data)) @pytest.mark.it_postcommit def test_geography_write_with_beam_rows(self): """Test writing GEOGRAPHY data using Beam Rows with GeographyType.""" table_name = 'geography_beam_rows' table_id = '{}.{}'.format(self.dataset_id, table_name) # Create the table first self.create_geography_table(table_name) # Create Beam Rows with GeographyType row_elements = [ beam.Row(id=1, location='POINT(1 1)', optional_location='POINT(2 2)'), beam.Row( id=2, location='LINESTRING(0 0, 1 1, 2 2)', optional_location=None), beam.Row( id=3, location='POLYGON((0 0, 0 1, 1 1, 1 0, 0 0))', optional_location='POINT(3 3)') ] # Expected data for verification expected_data = [(1, 'POINT(1 1)', 'POINT(2 2)'), (2, 'LINESTRING(0 0, 1 1, 2 2)', None), (3, 'POLYGON((0 0, 0 1, 1 1, 1 0, 0 0))', 'POINT(3 3)')] pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query=( "SELECT id, location, optional_location FROM %s ORDER BY id" % table_id), data=expected_data) ] args = self.test_pipeline.get_full_options_as_args() with beam.Pipeline(argv=args) as p: _ = ( p | 'CreateRows' >> beam.Create(row_elements) | 'ConvertToDict' >> beam.Map( lambda row: { 'id': row.id, 'location': row.location, 'optional_location': row.optional_location }) | 'WriteToBQ' >> WriteToBigQuery( table=table_id, method=WriteToBigQuery.Method.STREAMING_INSERTS, schema={ "fields": [{ "name": "id", "type": "INTEGER", "mode": "REQUIRED" }, { "name": "location", "type": "GEOGRAPHY", "mode": "REQUIRED" }, { "name": "optional_location", "type": "GEOGRAPHY", "mode": "NULLABLE" }] })) # Wait a bit for streaming inserts to complete time.sleep(5) # Verify the data was written correctly hc.assert_that(None, hc.all_of(*pipeline_verifiers)) @pytest.mark.it_postcommit def test_geography_repeated_fields(self): """Test GEOGRAPHY fields with REPEATED mode.""" table_name = 'geography_repeated' table_id = '{}.{}'.format(self.dataset_id, table_name) input_data = [ { 'id': 1, 'location': 'POINT(0 0)', 'optional_location': 'POINT(1 1)', 'path': ['POINT(0 0)', 'POINT(1 1)', 'POINT(2 2)'] }, { 'id': 2, 'location': 'POINT(10 10)', 'optional_location': None, 'path': ['LINESTRING(0 0, 5 5)', 'LINESTRING(5 5, 10 10)'] }, { 'id': 3, 'location': 'POLYGON((0 0, 0 1, 1 1, 1 0, 0 0))', 'optional_location': 'POINT(0.5 0.5)', 'path': [] # Empty array } ] table_schema = { "fields": [{ "name": "id", "type": "INTEGER", "mode": "REQUIRED" }, { "name": "location", "type": "GEOGRAPHY", "mode": "REQUIRED" }, { "name": "optional_location", "type": "GEOGRAPHY", "mode": "NULLABLE" }, { "name": "path", "type": "GEOGRAPHY", "mode": "REPEATED" }] } # Write data args = self.test_pipeline.get_full_options_as_args() with beam.Pipeline(argv=args) as p: _ = ( p | 'CreateData' >> beam.Create(input_data) | 'WriteToBQ' >> WriteToBigQuery( table=table_id, schema=table_schema, method=WriteToBigQuery.Method.STREAMING_INSERTS)) # Read and verify with beam.Pipeline(argv=args) as p: result = ( p | 'ReadFromBQ' >> ReadFromBigQuery( table=table_id, method=ReadFromBigQuery.Method.DIRECT_READ, project=self.project) | 'ExtractData' >> beam.Map( lambda row: (row['id'], len(row['path']) if row['path'] else 0))) expected_counts = [(1, 3), (2, 2), (3, 0)] assert_that(result, equal_to(expected_counts)) @pytest.mark.it_postcommit def test_geography_complex_geometries(self): """Test complex GEOGRAPHY geometries and edge cases.""" table_name = 'geography_complex' table_id = '{}.{}'.format(self.dataset_id, table_name) # Complex geometries including collections and high precision coordinates input_data = [ { 'id': 1, 'location': ( 'GEOMETRYCOLLECTION(POINT(4 6), LINESTRING(4 6, 7 10))'), 'optional_location': None }, { 'id': 2, 'location': ( 'MULTIPOLYGON(((0 0, 1 0, 1 1, 0 1, 0 0)), ' '((2 2, 3 2, 3 3, 2 3, 2 2)))'), # Fixed orientation 'optional_location': ('POINT(-122.419416 37.774929)' ) # High precision }, { 'id': 3, 'location': ('POLYGON((0 0, 0 5, 5 5, 5 0, 0 0))' ), # Simple polygon without holes 'optional_location': ('LINESTRING(-122 37, -121 38)' ) # Fixed non-antipodal coordinates } ] table_schema = { "fields": [{ "name": "id", "type": "INTEGER", "mode": "REQUIRED" }, { "name": "location", "type": "GEOGRAPHY", "mode": "REQUIRED" }, { "name": "optional_location", "type": "GEOGRAPHY", "mode": "NULLABLE" }] } expected_data = [(1, 'LINESTRING(4 6, 7 10)', None), ( 2, 'MULTIPOLYGON(((0 0, 1 0, 1 1, 0 1, 0 0)), ' '((2 2, 3 2, 3 3, 2 3, 2 2)))', 'POINT(-122.419416 37.774929)'), ( 3, 'POLYGON((0 0, 0 5, 5 5, 5 0, 0 0))', 'LINESTRING(-122 37, -121 38)')] pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query=( "SELECT id, location, optional_location FROM %s ORDER BY id" % table_id), data=expected_data) ] args = self.test_pipeline.get_full_options_as_args() with beam.Pipeline(argv=args) as p: _ = ( p | 'CreateData' >> beam.Create(input_data) | 'WriteToBQ' >> WriteToBigQuery( table=table_id, schema=table_schema, method=WriteToBigQuery.Method.STREAMING_INSERTS)) hc.assert_that(p, hc.all_of(*pipeline_verifiers)) @pytest.mark.uses_gcp_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_JARS'), "EXPANSION_JARS environment var is not provided, " "indicating that jars have not been built") def test_geography_storage_write_api(self): """Test GEOGRAPHY with Storage Write API method.""" table_name = 'geography_storage_write' table_id = '{}.{}'.format(self.dataset_id, table_name) input_data = [{ 'id': 1, 'location': 'POINT(0 0)', 'optional_location': 'POINT(1 1)' }, { 'id': 2, 'location': 'LINESTRING(0 0, 1 1)', 'optional_location': None }] table_schema = { "fields": [{ "name": "id", "type": "INTEGER", "mode": "REQUIRED" }, { "name": "location", "type": "GEOGRAPHY", "mode": "REQUIRED" }, { "name": "optional_location", "type": "GEOGRAPHY", "mode": "NULLABLE" }] } expected_data = [(1, 'POINT(0 0)', 'POINT(1 1)'), (2, 'LINESTRING(0 0, 1 1)', None)] pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query=( "SELECT id, location, optional_location FROM %s ORDER BY id" % table_id), data=expected_data) ] args = self.test_pipeline.get_full_options_as_args() with beam.Pipeline(argv=args) as p: _ = ( p | 'CreateData' >> beam.Create(input_data) | 'WriteToBQ' >> WriteToBigQuery( table=table_id, schema=table_schema, method=WriteToBigQuery.Method.STORAGE_WRITE_API)) hc.assert_that(p, hc.all_of(*pipeline_verifiers)) @pytest.mark.it_postcommit def test_geography_file_loads_method(self): """Test GEOGRAPHY with FILE_LOADS method.""" table_name = 'geography_file_loads' table_id = '{}.{}'.format(self.dataset_id, table_name) input_data = [ { 'id': i, 'location': f'POINT({i} {i})', 'optional_location': ( f'POINT({i+10} {i+10})' if i % 2 == 0 else None) } for i in range(1, 11) # 10 records ] table_schema = { "fields": [{ "name": "id", "type": "INTEGER", "mode": "REQUIRED" }, { "name": "location", "type": "GEOGRAPHY", "mode": "REQUIRED" }, { "name": "optional_location", "type": "GEOGRAPHY", "mode": "NULLABLE" }] } # Verify count and some sample data pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query="SELECT COUNT(*) as count FROM %s" % table_id, data=[(10, )]) ] args = self.test_pipeline.get_full_options_as_args() gcs_temp_location = ( f'gs://temp-storage-for-end-to-end-tests/' f'bq_it_test_{int(time.time())}') with beam.Pipeline(argv=args) as p: _ = ( p | 'CreateData' >> beam.Create(input_data) | 'WriteToBQ' >> WriteToBigQuery( table=table_id, schema=table_schema, method=WriteToBigQuery.Method.FILE_LOADS, custom_gcs_temp_location=gcs_temp_location)) hc.assert_that(p, hc.all_of(*pipeline_verifiers)) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_io_metadata.py0000644000175100017510000000657415206113032023601 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Metadata for use in BigQueryIO, i.e. a job_id to use in BQ job labels.""" # pytype: skip-file import re from apache_beam.io.gcp import gce_metadata_util _VALID_CLOUD_LABEL_PATTERN = re.compile(r'^[a-z0-9\_\-]{1,63}$') def _sanitize_value(value): """Sanitizes a value into a valid BigQuery label value.""" return re.sub(r'[^\w-]+', '', value.lower().replace('/', '-'))[0:63] def _is_valid_cloud_label_value(label_value): """Returns true if label_value is a valid cloud label string. This function can return false in cases where the label value is valid. However, it will not return true in a case where the lavel value is invalid. This is because a stricter set of allowed characters is used in this validator, because foreign language characters are not accepted. Thus, this should not be used as a generic validator for all cloud labels. See Also: https://cloud.google.com/compute/docs/labeling-resources Args: label_value: The label value to validate. Returns: True if the label value is a valid """ return _VALID_CLOUD_LABEL_PATTERN.match(label_value) def create_bigquery_io_metadata(step_name=None): """Creates a BigQueryIOMetadata. This will request metadata properly based on which runner is being used. """ dataflow_job_id = gce_metadata_util.fetch_dataflow_job_id() # If a dataflow_job id is returned on GCE metadata. Then it means # This program is running on a Dataflow GCE VM. is_dataflow_runner = bool(dataflow_job_id) kwargs = {} if is_dataflow_runner: # Only use this label if it is validated already. # As we do not want a bad label to fail the BQ job. if _is_valid_cloud_label_value(dataflow_job_id): kwargs['beam_job_id'] = dataflow_job_id if step_name: step_name = _sanitize_value(step_name) if _is_valid_cloud_label_value(step_name): kwargs['step_name'] = step_name return BigQueryIOMetadata(**kwargs) class BigQueryIOMetadata(object): """Metadata class for BigQueryIO. i.e. to use as BQ job labels. Do not construct directly, use the create_bigquery_io_metadata factory. Which will request metadata properly based on which runner is being used. """ def __init__(self, beam_job_id=None, step_name=None): self.beam_job_id = beam_job_id self.step_name = step_name def add_additional_bq_job_labels(self, job_labels=None): job_labels = job_labels or {} if self.beam_job_id and 'beam_job_id' not in job_labels: job_labels['beam_job_id'] = self.beam_job_id if self.step_name and 'step_name' not in job_labels: job_labels['step_name'] = self.step_name return job_labels ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_io_metadata_test.py0000644000175100017510000000514715206113032024633 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for bigquery_io_metadata.""" # pytype: skip-file import logging import unittest from apache_beam.io.gcp import bigquery_io_metadata class BigqueryIoMetadataTest(unittest.TestCase): def test_is_valid_cloud_label_value(self): # A dataflow job ID. # Lowercase letters, numbers, underscores and hyphens are allowed. test_str = '2020-06-29_15_26_09-12838749047888422749' self.assertTrue(bigquery_io_metadata._is_valid_cloud_label_value(test_str)) # At least one character. test_str = '0' self.assertTrue(bigquery_io_metadata._is_valid_cloud_label_value(test_str)) # Up to 63 characters. test_str = '0123456789abcdefghij0123456789abcdefghij0123456789abcdefghij012' self.assertTrue(bigquery_io_metadata._is_valid_cloud_label_value(test_str)) # Lowercase letters allowed test_str = 'abcdefghijklmnopqrstuvwxyz' for test_char in test_str: self.assertTrue( bigquery_io_metadata._is_valid_cloud_label_value(test_char)) # Empty strings not allowed. test_str = '' self.assertFalse(bigquery_io_metadata._is_valid_cloud_label_value(test_str)) # 64 or more characters not allowed. test_str = ( '0123456789abcdefghij0123456789abcdefghij0123456789abcdefghij0123') self.assertFalse(bigquery_io_metadata._is_valid_cloud_label_value(test_str)) # Uppercase letters not allowed test_str = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ' for test_char in test_str: self.assertFalse( bigquery_io_metadata._is_valid_cloud_label_value(test_char)) # Special characters besides hyphens are not allowed test_str = '!@#$%^&*()+=[{]};:\'\"\\|,<.>?/`~' for test_char in test_str: self.assertFalse( bigquery_io_metadata._is_valid_cloud_label_value(test_char)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_io_read_it_test.py0000644000175100017510000000451715206113032024462 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A Dataflow job that counts the number of rows in a BQ table. Can be configured to simulate slow reading for a given number of rows. """ # pytype: skip-file import logging import unittest import pytest from hamcrest.core.core.allof import all_of from apache_beam.io.gcp import bigquery_io_read_pipeline from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline class BigqueryIOReadIT(unittest.TestCase): DEFAULT_DATASET = "big_query_import_export" DEFAULT_TABLE_PREFIX = "export_" NUM_RECORDS = { "empty": 0, "1M": 10592, "1G": 11110839, "1T": 11110839000, } def run_bigquery_io_read_pipeline(self, input_size, beam_bq_source=False): test_pipeline = TestPipeline(is_integration_test=True) pipeline_verifiers = [ PipelineStateMatcher(), ] extra_opts = { 'input_table': self.DEFAULT_DATASET + "." + self.DEFAULT_TABLE_PREFIX + input_size, 'num_records': self.NUM_RECORDS[input_size], 'beam_bq_source': str(beam_bq_source), 'on_success_matcher': all_of(*pipeline_verifiers) } bigquery_io_read_pipeline.run( test_pipeline.get_full_options_as_args(**extra_opts)) @pytest.mark.it_postcommit def test_bigquery_read_custom_1M_python(self): self.run_bigquery_io_read_pipeline('1M', True) @pytest.mark.it_postcommit def test_bigquery_read_1M_python(self): self.run_bigquery_io_read_pipeline('1M') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_io_read_pipeline.py0000644000175100017510000000607615206113032024616 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A Dataflow job that counts the number of rows in a BQ table. Can be configured to simulate slow reading for a given number of rows. """ # pytype: skip-file import argparse import logging import random import time import apache_beam as beam from apache_beam.io.gcp.bigquery import ReadFromBigQuery from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class RowToStringWithSlowDown(beam.DoFn): def process(self, element, num_slow=0, *args, **kwargs): if num_slow == 0: yield ['row'] else: rand = random.random() * 100 if rand < num_slow: time.sleep(0.01) yield ['slow_row'] else: yield ['row'] def run(argv=None): parser = argparse.ArgumentParser() parser.add_argument( '--input_table', required=True, help='Input table to process.') parser.add_argument( '--num_records', required=True, help='The expected number of records', type=int) parser.add_argument( '--num_slow', default=0, help=( 'Percentage of rows that will be slow. ' 'Must be in the range [0, 100)')) parser.add_argument( '--beam_bq_source', default=False, type=bool, help=( 'Whether to use the new ReadFromBigQuery' ' transform, or the BigQuerySource.')) known_args, pipeline_args = parser.parse_known_args(argv) options = PipelineOptions(pipeline_args) with TestPipeline(options=options) as p: if known_args.beam_bq_source: reader = ReadFromBigQuery( table='%s:%s' % (options.view_as(GoogleCloudOptions).project, known_args.input_table)) else: reader = beam.io.Read(beam.io.BigQuerySource(known_args.input_table)) # pylint: disable=expression-not-assigned count = ( p | 'read' >> reader | 'row to string' >> beam.ParDo( RowToStringWithSlowDown(), num_slow=known_args.num_slow) | 'count' >> beam.combiners.Count.Globally()) assert_that(count, equal_to([known_args.num_records])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_json_it_test.py0000644000175100017510000003747315206113032024040 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Integration tests for BigQuery's JSON data type """ import argparse import json import logging import time import unittest from random import randint import pytest import apache_beam as beam from apache_beam.io.gcp.bigquery import ReadFromBigQuery from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # pylint: disable=wrong-import-order, wrong-import-position try: from google.api_core import exceptions as gexc from google.cloud import bigquery except ImportError: gexc = None bigquery = None # pylint: enable=wrong-import-order, wrong-import-position _LOGGER = logging.getLogger(__name__) PROJECT = 'apache-beam-testing' DATASET_ID = 'bq_jsontype_test_nodelete' JSON_TABLE_NAME = 'json_data' JSON_TABLE_DESTINATION = f"{PROJECT}:{DATASET_ID}.{JSON_TABLE_NAME}" STREAMING_TEST_TABLE = "py_streaming_test" \ f"{time.time_ns() // 1000}_{randint(0,32)}" FILE_LOAD_TABLE = "py_fileload_test" \ f"{time.time_ns() // 1000}_{randint(0,32)}" class BigQueryJsonIT(unittest.TestCase): created_tables = set() @classmethod def setUpClass(cls): cls.test_pipeline = TestPipeline(is_integration_test=True) @classmethod def tearDownClass(cls): if cls.created_tables: client = bigquery.Client(project=PROJECT) for ref in cls.created_tables: try: client.delete_table(ref[len(PROJECT) + 1:]) # need dataset:table except gexc.NotFound: pass # just skip def run_test_write(self, options): json_table_schema = self.generate_schema() rows_to_write = [] json_data = self.generate_data() for country_code, country in json_data.items(): cities_to_write = [] for city_name, city in country["cities"].items(): cities_to_write.append({'city_name': city_name, 'city': city}) rows_to_write.append({ 'country_code': country_code, 'country': country["country"], 'stats': country["stats"], 'cities': cities_to_write, 'landmarks': country["landmarks"] }) parser = argparse.ArgumentParser() parser.add_argument('--write_method') parser.add_argument('--output') parser.add_argument('--unescape', required=False) known_args, pipeline_args = parser.parse_known_args(options) self.created_tables.add(known_args.output) with beam.Pipeline(argv=pipeline_args) as p: _ = ( p | "Create rows with JSON data" >> beam.Create(rows_to_write) | "Write to BigQuery" >> beam.io.WriteToBigQuery( method=known_args.write_method, table=known_args.output, schema=json_table_schema, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, )) extra_opts = { 'read_method': "EXPORT", 'input': known_args.output, 'unescape': known_args.unescape } read_options = self.test_pipeline.get_full_options_as_args(**extra_opts) self.read_and_validate_rows(read_options) def read_and_validate_rows(self, options): json_data = self.generate_data() parser = argparse.ArgumentParser() parser.add_argument('--read_method') parser.add_argument('--query') parser.add_argument('--input') parser.add_argument('--unescape', required=False) known_args, pipeline_args = parser.parse_known_args(options) # TODO(yathu) remove this conversion when FILE_LOAD produces unescaped # JSON string def maybe_unescape(value): if known_args.unescape: value = bytes(value, "utf-8").decode("unicode_escape")[1:-1] return json.loads(value) class CompareJson(beam.DoFn, unittest.TestCase): def process(self, row): country_code = row["country_code"] expected = json_data[country_code] # Test country (JSON String) country_actual = maybe_unescape(row["country"]) country_expected = json.loads(expected["country"]) self.assertTrue(country_expected == country_actual) # Test stats (JSON String in BigQuery struct) for stat, value in row["stats"].items(): stats_actual = maybe_unescape(value) stats_expected = json.loads(expected["stats"][stat]) self.assertTrue(stats_expected == stats_actual) # Test cities (JSON String in BigQuery array of structs) for city_row in row["cities"]: city = city_row["city"] city_name = city_row["city_name"] city_actual = maybe_unescape(city) city_expected = json.loads(expected["cities"][city_name]) self.assertTrue(city_expected == city_actual) # Test landmarks (JSON String in BigQuery array) landmarks_actual = row["landmarks"] landmarks_expected = expected["landmarks"] for i in range(len(landmarks_actual)): l_actual = maybe_unescape(landmarks_actual[i]) l_expected = json.loads(landmarks_expected[i]) self.assertTrue(l_expected == l_actual) method = ReadFromBigQuery.Method.DIRECT_READ if \ known_args.read_method == "DIRECT_READ" else \ ReadFromBigQuery.Method.EXPORT if known_args.query: json_query_data = self.generate_query_data() with beam.Pipeline(argv=pipeline_args) as p: data = p | 'Read rows' >> ReadFromBigQuery( query=known_args.query, method=method, use_standard_sql=True) assert_that(data, equal_to(json_query_data)) else: with beam.Pipeline(argv=pipeline_args) as p: _ = p | 'Read rows' >> ReadFromBigQuery( table=known_args.input, method=method, ) | 'Validate rows' >> beam.ParDo(CompareJson()) @pytest.mark.it_postcommit def test_direct_read(self): extra_opts = { 'read_method': "DIRECT_READ", 'input': JSON_TABLE_DESTINATION, } options = self.test_pipeline.get_full_options_as_args(**extra_opts) self.read_and_validate_rows(options) @pytest.mark.it_postcommit def test_export_read(self): extra_opts = { 'read_method': "EXPORT", 'input': JSON_TABLE_DESTINATION, } options = self.test_pipeline.get_full_options_as_args(**extra_opts) self.read_and_validate_rows(options) @pytest.mark.it_postcommit def test_query_read(self): extra_opts = { 'query': "SELECT " "country_code, " "country.past_leaders[2] AS past_leader, " "stats.gdp_per_capita[\"gdp_per_capita\"] AS gdp, " "cities[OFFSET(1)].city.name AS city_name, " "landmarks[OFFSET(1)][\"name\"] AS landmark_name " f"FROM `{PROJECT}.{DATASET_ID}.{JSON_TABLE_NAME}`", } options = self.test_pipeline.get_full_options_as_args(**extra_opts) self.read_and_validate_rows(options) @pytest.mark.it_postcommit def test_streaming_inserts(self): extra_opts = { 'output': f"{PROJECT}:{DATASET_ID}.{STREAMING_TEST_TABLE}", 'write_method': "STREAMING_INSERTS" } options = self.test_pipeline.get_full_options_as_args(**extra_opts) self.run_test_write(options) @pytest.mark.it_postcommit def test_file_loads_write(self): extra_opts = { 'output': f"{PROJECT}:{DATASET_ID}.{FILE_LOAD_TABLE}", 'write_method': "FILE_LOADS", "unescape": "True" } options = self.test_pipeline.get_full_options_as_args(**extra_opts) self.run_test_write(options) # Schema for writing to BigQuery def generate_schema(self): from apache_beam.io.gcp.internal.clients.bigquery import TableFieldSchema from apache_beam.io.gcp.internal.clients.bigquery import TableSchema json_fields = [ TableFieldSchema(name='country_code', type='STRING', mode='NULLABLE'), TableFieldSchema(name='country', type='JSON', mode='NULLABLE'), TableFieldSchema( name='stats', type='STRUCT', mode='NULLABLE', fields=[ TableFieldSchema( name="gdp_per_capita", type='JSON', mode='NULLABLE'), TableFieldSchema( name="co2_emissions", type='JSON', mode='NULLABLE'), ]), TableFieldSchema( name='cities', type='STRUCT', mode='REPEATED', fields=[ TableFieldSchema( name="city_name", type='STRING', mode='NULLABLE'), TableFieldSchema(name="city", type='JSON', mode='NULLABLE'), ]), TableFieldSchema(name='landmarks', type='JSON', mode='REPEATED'), ] schema = TableSchema(fields=json_fields) return schema # Expected data for query test def generate_query_data(self): query_data = [{ 'country_code': 'usa', 'past_leader': '\"George W. Bush\"', 'gdp': '58559.675', 'city_name': '\"Los Angeles\"', 'landmark_name': '\"Golden Gate Bridge\"' }, { 'country_code': 'aus', 'past_leader': '\"Kevin Rudd\"', 'gdp': '58043.581', 'city_name': '\"Melbourne\"', 'landmark_name': '\"Great Barrier Reef\"' }, { 'country_code': 'special', 'past_leader': '\"!@#$%^&*()_+\"', 'gdp': '421.7', 'city_name': '\"Bikini Bottom\"', 'landmark_name': "\"Willy Wonka's Factory\"" }] return query_data def generate_data(self): # Raw country data usa = { "name": "United States of America", "population": 329484123, "cities": { "nyc": { "name": "New York City", "state": "NY", "population": 8622357 }, "la": { "name": "Los Angeles", "state": "CA", "population": 4085014 }, "chicago": { "name": "Chicago", "state": "IL", "population": 2670406 }, }, "past_leaders": [ "Donald Trump", "Barack Obama", "George W. Bush", "Bill Clinton" ], "in_northern_hemisphere": True } aus = { "name": "Australia", "population": 25687041, "cities": { "sydney": { "name": "Sydney", "state": "New South Wales", "population": 5367206 }, "melbourne": { "name": "Melbourne", "state": "Victoria", "population": 5159211 }, "brisbane": { "name": "Brisbane", "state": "Queensland", "population": 2560720 } }, "past_leaders": [ "Malcolm Turnbull", "Tony Abbot", "Kevin Rudd", ], "in_northern_hemisphere": False } special = { "name": "newline\n, form\f, tab\t, \"quotes\", " "\\backslash\\, backspace\b, \u0000_hex_\u0f0f", "population": -123456789, "cities": { "basingse": { "name": "Ba Sing Se", "state": "The Earth Kingdom", "population": 200000 }, "bikinibottom": { "name": "Bikini Bottom", "state": "The Pacific Ocean", "population": 50000 } }, "past_leaders": [ "1", "2", "!@#$%^&*()_+", ], "in_northern_hemisphere": True } landmarks = { "usa_0": { "name": "Statue of Liberty", "cool rating": None }, "usa_1": { "name": "Golden Gate Bridge", "cool rating": "very cool" }, "usa_2": { "name": "Grand Canyon", "cool rating": "very very cool" }, "aus_0": { "name": "Sydney Opera House", "cool rating": "amazing" }, "aus_1": { "name": "Great Barrier Reef", "cool rating": None }, "special_0": { "name": "Hogwarts School of WitchCraft and Wizardry", "cool rating": "magical" }, "special_1": { "name": "Willy Wonka's Factory", "cool rating": None }, "special_2": { "name": "Rivendell", "cool rating": "precious" }, } stats = { "usa_gdp_per_capita": { "gdp_per_capita": 58559.675, "currency": "constant 2015 US$" }, "usa_co2_emissions": { "co2 emissions": 15.241, "measurement": "metric tons per capita", "year": 2018 }, "aus_gdp_per_capita": { "gdp_per_capita": 58043.581, "currency": "constant 2015 US$" }, "aus_co2_emissions": { "co2 emissions": 15.476, "measurement": "metric tons per capita", "year": 2018 }, "special_gdp_per_capita": { "gdp_per_capita": 421.70, "currency": "constant 200 BC gold" }, "special_co2_emissions": { "co2 emissions": -10.79, "measurement": "metric tons per capita", "year": 2018 } } data = { "usa": { "country": json.dumps(usa), "cities": { "nyc": json.dumps(usa["cities"]["nyc"]), "la": json.dumps(usa["cities"]["la"]), "chicago": json.dumps(usa["cities"]["chicago"]) }, "landmarks": [ json.dumps(landmarks["usa_0"]), json.dumps(landmarks["usa_1"]), json.dumps(landmarks["usa_2"]) ], "stats": { "gdp_per_capita": json.dumps(stats["usa_gdp_per_capita"]), "co2_emissions": json.dumps(stats["usa_co2_emissions"]) } }, "aus": { "country": json.dumps(aus), "cities": { "sydney": json.dumps(aus["cities"]["sydney"]), "melbourne": json.dumps(aus["cities"]["melbourne"]), "brisbane": json.dumps(aus["cities"]["brisbane"]) }, "landmarks": [ json.dumps(landmarks["aus_0"]), json.dumps(landmarks["aus_1"]) ], "stats": { "gdp_per_capita": json.dumps(stats["aus_gdp_per_capita"]), "co2_emissions": json.dumps(stats["aus_co2_emissions"]) } }, "special": { "country": json.dumps(special), "cities": { "basingse": json.dumps(special["cities"]["basingse"]), "bikinibottom": json.dumps(special["cities"]["bikinibottom"]) }, "landmarks": [ json.dumps(landmarks["special_0"]), json.dumps(landmarks["special_1"]), json.dumps(landmarks["special_2"]) ], "stats": { "gdp_per_capita": json.dumps(stats["special_gdp_per_capita"]), "co2_emissions": json.dumps(stats["special_co2_emissions"]) } } } return data if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_read_internal.py0000644000175100017510000004102715206113032024131 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Internal library for reading data from BigQuery. NOTHING IN THIS FILE HAS BACKWARDS COMPATIBILITY GUARANTEES. """ import collections import decimal import json import logging import secrets import time import uuid from typing import TYPE_CHECKING from typing import Any from typing import Iterable from typing import Optional from typing import Union import apache_beam as beam from apache_beam.coders import coders from apache_beam.io.avroio import _create_avro_source from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystems import FileSystems from apache_beam.io.gcp import bigquery_tools from apache_beam.io.gcp.bigquery_io_metadata import create_bigquery_io_metadata from apache_beam.io.iobase import BoundedSource from apache_beam.io.textio import _TextSource from apache_beam.metrics.metric import Lineage from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.value_provider import ValueProvider from apache_beam.transforms import PTransform if TYPE_CHECKING: from apache_beam.io.gcp.bigquery import ReadFromBigQueryRequest try: from apache_beam.io.gcp.internal.clients.bigquery import DatasetReference from apache_beam.io.gcp.internal.clients.bigquery import TableReference except ImportError: DatasetReference = None TableReference = None _LOGGER = logging.getLogger(__name__) def bigquery_export_destination_uri( gcs_location_vp: Union[str, Optional[ValueProvider]], temp_location: Optional[str], unique_id: str, directory_only: bool = False, ) -> str: """Returns the fully qualified Google Cloud Storage URI where the extracted table should be written. """ file_pattern = 'bigquery-table-dump-*.json' gcs_location = None if gcs_location_vp is not None: if isinstance(gcs_location_vp, ValueProvider): gcs_location = gcs_location_vp.get() else: gcs_location = gcs_location_vp if gcs_location is not None: gcs_base = gcs_location elif temp_location is not None: gcs_base = temp_location _LOGGER.debug("gcs_location is empty, using temp_location instead") else: raise ValueError( 'ReadFromBigQuery requires a GCS location to be provided. Neither ' 'gcs_location in the constructor nor the fallback option ' '--temp_location is set.') if not unique_id: unique_id = uuid.uuid4().hex if directory_only: return FileSystems.join(gcs_base, unique_id) else: return FileSystems.join(gcs_base, unique_id, file_pattern) class _PassThroughThenCleanup(PTransform): """A PTransform that invokes a DoFn after the input PCollection has been processed. DoFn should have arguments (element, side_input, cleanup_signal). Utilizes readiness of PCollection to trigger DoFn. """ def __init__(self, side_input=None): self.side_input = side_input def expand(self, input): class PassThrough(beam.DoFn): def process(self, element): yield element class RemoveExtractedFiles(beam.DoFn): def process(self, unused_element, unused_signal, gcs_locations): FileSystems.delete(list(gcs_locations)) main_output, cleanup_signal = input | beam.ParDo( PassThrough()).with_outputs( 'cleanup_signal', main='main') cleanup_input = input.pipeline | beam.Create([None]) _ = cleanup_input | beam.ParDo( RemoveExtractedFiles(), beam.pvalue.AsSingleton(cleanup_signal), self.side_input, ) return main_output class _PassThroughThenCleanupTempDatasets(PTransform): """A PTransform that invokes a DoFn after the input PCollection has been processed. DoFn should have arguments (element, side_input, cleanup_signal). Utilizes readiness of PCollection to trigger DoFn. """ def __init__(self, side_input=None): self.side_input = side_input def expand(self, input): pipeline_options = input.pipeline.options class PassThrough(beam.DoFn): def process(self, element): yield element class CleanUpProjects(beam.DoFn): def process(self, unused_element, unused_signal, pipeline_details): bq = bigquery_tools.BigQueryWrapper.from_pipeline_options( pipeline_options) pipeline_details = pipeline_details[0] if 'temp_table_ref' in pipeline_details.keys(): temp_table_ref = pipeline_details['temp_table_ref'] bq._clean_up_beam_labelled_temporary_datasets( project_id=temp_table_ref.projectId, dataset_id=temp_table_ref.datasetId, table_id=temp_table_ref.tableId) elif 'project_id' in pipeline_details.keys(): bq._clean_up_beam_labelled_temporary_datasets( project_id=pipeline_details['project_id'], labels=pipeline_details['bigquery_dataset_labels']) main_output, cleanup_signal = input | beam.ParDo( PassThrough()).with_outputs( 'cleanup_signal', main='main') cleanup_input = input.pipeline | beam.Create([None]) _ = cleanup_input | beam.ParDo( CleanUpProjects(), beam.pvalue.AsSingleton(cleanup_signal), self.side_input, ) return main_output class _BigQueryReadSplit(beam.transforms.DoFn): """Starts the process of reading from BigQuery. This transform will start a BigQuery export job, and output a number of file sources that are consumed downstream. """ def __init__( self, options: PipelineOptions, gcs_location: Union[str, ValueProvider] = None, validate: bool = False, use_json_exports: bool = False, bigquery_job_labels: dict[str, str] = None, step_name: str = None, job_name: str = None, unique_id: str = None, kms_key: str = None, project: str = None, temp_dataset: Union[str, DatasetReference] = None, query_priority: Optional[str] = None): self.options = options self.validate = validate self.use_json_exports = use_json_exports self.gcs_location = gcs_location self.bigquery_job_labels = bigquery_job_labels or {} self._step_name = step_name self._job_name = job_name or 'BQ_READ_SPLIT' self._source_uuid = unique_id self.kms_key = kms_key self.project = project self.temp_dataset = temp_dataset self.query_priority = query_priority self.bq_io_metadata = None def display_data(self): return { 'use_json_exports': str(self.use_json_exports), 'gcs_location': str(self.gcs_location), 'bigquery_job_labels': json.dumps(self.bigquery_job_labels), 'kms_key': str(self.kms_key), 'project': str(self.project), 'temp_dataset': str(self.temp_dataset) } def _get_temp_dataset_id(self): if self.temp_dataset is None: return None elif isinstance(self.temp_dataset, DatasetReference): return self.temp_dataset.datasetId elif isinstance(self.temp_dataset, str): return self.temp_dataset else: raise ValueError("temp_dataset has to be either str or DatasetReference") def _get_temp_dataset_project(self): """Returns the project ID for temporary dataset operations. If temp_dataset is a DatasetReference, returns its projectId. Otherwise, returns the pipeline project for billing. """ if isinstance(self.temp_dataset, DatasetReference): return self.temp_dataset.projectId else: return self._get_project() def start_bundle(self): self.bq = bigquery_tools.BigQueryWrapper( temp_dataset_id=self._get_temp_dataset_id(), client=bigquery_tools.BigQueryWrapper._bigquery_client(self.options)) def process(self, element: 'ReadFromBigQueryRequest') -> Iterable[BoundedSource]: if element.query is not None: if not self.bq.created_temp_dataset: self._setup_temporary_dataset(self.bq, element) table_reference = self._execute_query(self.bq, element) else: assert element.table table_reference = bigquery_tools.parse_table_reference( element.table, project=self._get_project()) if not table_reference.projectId: table_reference.projectId = self._get_project() schema, metadata_list = self._export_files( self.bq, element, table_reference) for metadata in metadata_list: yield self._create_source(metadata.path, schema) Lineage.sources().add( 'bigquery', table_reference.projectId, table_reference.datasetId, table_reference.tableId) if element.query is not None: self.bq._delete_table( table_reference.projectId, table_reference.datasetId, table_reference.tableId) def finish_bundle(self): if self.bq.created_temp_dataset: # Use the same project that was used to create the temp dataset temp_dataset_project = self._get_temp_dataset_project() self.bq.clean_up_temporary_dataset(temp_dataset_project) def _get_bq_metadata(self): if not self.bq_io_metadata: self.bq_io_metadata = create_bigquery_io_metadata(self._step_name) return self.bq_io_metadata def _create_source(self, path, schema): if not self.use_json_exports: return _create_avro_source(path, validate=self.validate) else: return _TextSource( path, min_bundle_size=0, compression_type=CompressionTypes.UNCOMPRESSED, strip_trailing_newlines=True, coder=_JsonToDictCoder(schema), validate=self.validate) def _setup_temporary_dataset( self, bq: bigquery_tools.BigQueryWrapper, element: 'ReadFromBigQueryRequest'): location = bq.get_query_location( self._get_project(), element.query, not element.use_standard_sql) # Use the project from temp_dataset if it's a DatasetReference, # otherwise use the pipeline project temp_dataset_project = self._get_temp_dataset_project() bq.create_temporary_dataset( temp_dataset_project, location, kms_key=self.kms_key) def _execute_query( self, bq: bigquery_tools.BigQueryWrapper, element: 'ReadFromBigQueryRequest'): query_job_name = bigquery_tools.generate_bq_job_name( self._job_name, self._source_uuid, bigquery_tools.BigQueryJobTypes.QUERY, '%s_%s' % (int(time.time()), secrets.token_hex(3))) job = bq._start_query_job( self._get_project(), element.query, not element.use_standard_sql, element.flatten_results, job_id=query_job_name, priority=self.query_priority, kms_key=self.kms_key, job_labels=self._get_bq_metadata().add_additional_bq_job_labels( self.bigquery_job_labels)) job_ref = job.jobReference bq.wait_for_bq_job(job_ref, max_retries=0) return bq._get_temp_table(self._get_project()) def _export_files( self, bq: bigquery_tools.BigQueryWrapper, element: 'ReadFromBigQueryRequest', table_reference: TableReference): """Runs a BigQuery export job. Returns: bigquery.TableSchema instance, a list of FileMetadata instances """ job_labels = self._get_bq_metadata().add_additional_bq_job_labels( self.bigquery_job_labels) export_job_name = bigquery_tools.generate_bq_job_name( self._job_name, self._source_uuid, bigquery_tools.BigQueryJobTypes.EXPORT, element.obj_id) temp_location = self.options.view_as(GoogleCloudOptions).temp_location gcs_location = bigquery_export_destination_uri( self.gcs_location, temp_location, '%s%s' % (self._source_uuid, element.obj_id)) try: if self.use_json_exports: job_ref = bq.perform_extract_job([gcs_location], export_job_name, table_reference, bigquery_tools.FileFormat.JSON, project=self._get_project(), job_labels=job_labels, include_header=False) else: job_ref = bq.perform_extract_job([gcs_location], export_job_name, table_reference, bigquery_tools.FileFormat.AVRO, project=self._get_project(), include_header=False, job_labels=job_labels, use_avro_logical_types=True) bq.wait_for_bq_job(job_ref) except Exception as exn: # pylint: disable=broad-except # The error messages thrown in this case are generic and misleading, # so leave this breadcrumb in case it's the root cause. logging.warning( "Error exporting table: %s. " "Note that external tables cannot be exported: " "https://cloud.google.com/bigquery/docs/external-tables" "#external_table_limitations", exn) raise metadata_list = FileSystems.match([gcs_location])[0].metadata_list if isinstance(table_reference, ValueProvider): table_ref = bigquery_tools.parse_table_reference( element.table, project=self._get_project()) else: table_ref = table_reference table = bq.get_table( table_ref.projectId, table_ref.datasetId, table_ref.tableId) return table.schema, metadata_list def _get_project(self): """Returns the project that queries and exports will be billed to.""" project = self.options.view_as(GoogleCloudOptions).project if isinstance(project, ValueProvider): project = project.get() if not project: project = self.project return project FieldSchema = collections.namedtuple('FieldSchema', 'fields mode name type') class _JsonToDictCoder(coders.Coder): """A coder for a JSON string to a Python dict.""" def __init__(self, table_schema): self.fields = self._convert_to_tuple(table_schema.fields) self._converters = { 'INTEGER': int, 'INT64': int, 'FLOAT': float, 'FLOAT64': float, 'NUMERIC': self._to_decimal, 'BYTES': self._to_bytes, } @staticmethod def _to_decimal(value): return decimal.Decimal(value) @staticmethod def _to_bytes(value): """Converts value from str to bytes.""" return value.encode('utf-8') @classmethod def _convert_to_tuple(cls, table_field_schemas): """Recursively converts the list of TableFieldSchema instances to the list of tuples to prevent errors when pickling and unpickling TableFieldSchema instances. """ if not table_field_schemas: return [] return [ FieldSchema(cls._convert_to_tuple(x.fields), x.mode, x.name, x.type) for x in table_field_schemas ] def decode(self, value): value = json.loads(value.decode('utf-8')) return self._decode_row(value, self.fields) def _decode_row(self, row: dict[str, Any], schema_fields: list[FieldSchema]): for field in schema_fields: if field.name not in row: # The field exists in the schema, but it doesn't exist in this row. # It probably means its value was null, as the extract to JSON job # doesn't preserve null fields row[field.name] = None continue if field.mode == 'REPEATED': for i, elem in enumerate(row[field.name]): row[field.name][i] = self._decode_data(elem, field) else: row[field.name] = self._decode_data(row[field.name], field) return row def _decode_data(self, obj: Any, field: FieldSchema): if not field.fields: try: return self._converters[field.type](obj) except KeyError: # No need to do any conversion return obj return self._decode_row(obj, field.fields) def is_deterministic(self): return True def to_type_hint(self): return dict ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_read_internal_test.py0000644000175100017510000001556215206113032025175 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for BigQuery read internal module.""" import unittest from unittest import mock from apache_beam.io.gcp import bigquery_read_internal from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.value_provider import StaticValueProvider try: from apache_beam.io.gcp.internal.clients.bigquery import DatasetReference except ImportError: DatasetReference = None class BigQueryReadSplitTest(unittest.TestCase): """Tests for _BigQueryReadSplit DoFn.""" def setUp(self): if DatasetReference is None: self.skipTest('BigQuery dependencies are not installed') self.options = PipelineOptions() self.gcp_options = self.options.view_as(GoogleCloudOptions) self.gcp_options.project = 'test-project' def test_get_temp_dataset_project_with_string_temp_dataset(self): """Test _get_temp_dataset_project with string temp_dataset.""" split = bigquery_read_internal._BigQueryReadSplit( options=self.options, temp_dataset='temp_dataset_id') # Should return the pipeline project when temp_dataset is a string self.assertEqual(split._get_temp_dataset_project(), 'test-project') def test_get_temp_dataset_project_with_dataset_reference(self): """Test _get_temp_dataset_project with DatasetReference temp_dataset.""" dataset_ref = DatasetReference( projectId='custom-project', datasetId='temp_dataset_id') split = bigquery_read_internal._BigQueryReadSplit( options=self.options, temp_dataset=dataset_ref) # Should return the project from DatasetReference self.assertEqual(split._get_temp_dataset_project(), 'custom-project') def test_get_temp_dataset_project_with_none_temp_dataset(self): """Test _get_temp_dataset_project with None temp_dataset.""" split = bigquery_read_internal._BigQueryReadSplit( options=self.options, temp_dataset=None) # Should return the pipeline project when temp_dataset is None self.assertEqual(split._get_temp_dataset_project(), 'test-project') def test_get_temp_dataset_project_with_value_provider_project(self): """Test _get_temp_dataset_project with ValueProvider project.""" self.gcp_options.project = StaticValueProvider(str, 'vp-project') dataset_ref = DatasetReference( projectId='custom-project', datasetId='temp_dataset_id') split = bigquery_read_internal._BigQueryReadSplit( options=self.options, temp_dataset=dataset_ref) # Should still return the project from DatasetReference self.assertEqual(split._get_temp_dataset_project(), 'custom-project') @mock.patch('apache_beam.io.gcp.bigquery_tools.BigQueryWrapper') def test_setup_temporary_dataset_uses_correct_project(self, mock_bq_wrapper): """Test that _setup_temporary_dataset uses the correct project.""" dataset_ref = DatasetReference( projectId='custom-project', datasetId='temp_dataset_id') split = bigquery_read_internal._BigQueryReadSplit( options=self.options, temp_dataset=dataset_ref) # Mock the BigQueryWrapper instance mock_bq = mock.Mock() mock_bq.get_query_location.return_value = 'US' # Mock ReadFromBigQueryRequest mock_element = mock.Mock() mock_element.query = 'SELECT * FROM table' mock_element.use_standard_sql = True # Call _setup_temporary_dataset split._setup_temporary_dataset(mock_bq, mock_element) # Verify that create_temporary_dataset was called with the custom project mock_bq.create_temporary_dataset.assert_called_once_with( 'custom-project', 'US', kms_key=None) # Verify that get_query_location was called with the pipeline project mock_bq.get_query_location.assert_called_once_with( 'test-project', 'SELECT * FROM table', False) @mock.patch('apache_beam.io.gcp.bigquery_tools.BigQueryWrapper') def test_finish_bundle_uses_correct_project(self, mock_bq_wrapper): """Test that finish_bundle uses the correct project for cleanup.""" dataset_ref = DatasetReference( projectId='custom-project', datasetId='temp_dataset_id') split = bigquery_read_internal._BigQueryReadSplit( options=self.options, temp_dataset=dataset_ref) # Mock the BigQueryWrapper instance mock_bq = mock.Mock() mock_bq.created_temp_dataset = True split.bq = mock_bq # Call finish_bundle split.finish_bundle() # Verify that clean_up_temporary_dataset was called with the custom project mock_bq.clean_up_temporary_dataset.assert_called_once_with('custom-project') @mock.patch('apache_beam.io.gcp.bigquery_tools.BigQueryWrapper') def test_setup_temporary_dataset_with_string_temp_dataset( self, mock_bq_wrapper): """Test _setup_temporary_dataset with string temp_dataset uses pipeline project.""" split = bigquery_read_internal._BigQueryReadSplit( options=self.options, temp_dataset='temp_dataset_id') # Mock the BigQueryWrapper instance mock_bq = mock.Mock() mock_bq.get_query_location.return_value = 'US' # Mock ReadFromBigQueryRequest mock_element = mock.Mock() mock_element.query = 'SELECT * FROM table' mock_element.use_standard_sql = True # Call _setup_temporary_dataset split._setup_temporary_dataset(mock_bq, mock_element) # Verify that create_temporary_dataset was called with the pipeline project mock_bq.create_temporary_dataset.assert_called_once_with( 'test-project', 'US', kms_key=None) @mock.patch('apache_beam.io.gcp.bigquery_tools.BigQueryWrapper') def test_finish_bundle_with_string_temp_dataset(self, mock_bq_wrapper): """Test finish_bundle with string temp_dataset uses pipeline project.""" split = bigquery_read_internal._BigQueryReadSplit( options=self.options, temp_dataset='temp_dataset_id') # Mock the BigQueryWrapper instance mock_bq = mock.Mock() mock_bq.created_temp_dataset = True split.bq = mock_bq # Call finish_bundle split.finish_bundle() # Verify that clean_up_temporary_dataset was called with the pipeline # project mock_bq.clean_up_temporary_dataset.assert_called_once_with('test-project') if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_read_it_test.py0000644000175100017510000007277615206113032024007 0ustar00runnerrunner#!/usr/bin/env python # -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for BigQuery sources and sinks.""" # pytype: skip-file import base64 import datetime import logging import secrets import time import unittest import uuid from decimal import Decimal from functools import wraps import pytest import apache_beam as beam import apache_beam.io.gcp.bigquery from apache_beam.io.gcp import bigquery_schema_tools from apache_beam.io.gcp import bigquery_tools from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.options.value_provider import StaticValueProvider from apache_beam.runners.interactive import interactive_beam from apache_beam.runners.interactive.interactive_runner import InteractiveRunner from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.utils.timestamp import Timestamp # Protect against environments where bigquery library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apitools.base.py.exceptions import HttpError except ImportError: HttpError = None # pylint: enable=wrong-import-order, wrong-import-position _LOGGER = logging.getLogger(__name__) def skip(runners): if not isinstance(runners, list): runners = [runners] def inner(fn): @wraps(fn) def wrapped(self): if self.runner_name in runners: self.skipTest( 'This test doesn\'t work on these runners: {}'.format(runners)) else: return fn(self) return wrapped return inner def datetime_to_utc(element): for k, v in element.items(): if isinstance(v, (datetime.time, datetime.date)): element[k] = str(v) if isinstance(v, datetime.datetime) and v.tzinfo: # For datetime objects, we'll offset = v.utcoffset() utc_dt = (v - offset).strftime('%Y-%m-%d %H:%M:%S.%f UTC') element[k] = utc_dt return element class BigQueryReadIntegrationTests(unittest.TestCase): BIG_QUERY_DATASET_ID = 'python_read_table_' @classmethod def setUpClass(cls): cls.test_pipeline = TestPipeline(is_integration_test=True) cls.args = cls.test_pipeline.get_full_options_as_args() cls.runner_name = type(cls.test_pipeline.runner).__name__ cls.project = cls.test_pipeline.get_option('project') cls.bigquery_client = BigQueryWrapper() cls.dataset_id = '%s%d%s' % ( cls.BIG_QUERY_DATASET_ID, int(time.time()), secrets.token_hex(3)) cls.bigquery_client.get_or_create_dataset(cls.project, cls.dataset_id) _LOGGER.info( "Created dataset %s in project %s", cls.dataset_id, cls.project) @classmethod def tearDownClass(cls): request = bigquery.BigqueryDatasetsDeleteRequest( projectId=cls.project, datasetId=cls.dataset_id, deleteContents=True) try: _LOGGER.debug( "Deleting dataset %s in project %s", cls.dataset_id, cls.project) cls.bigquery_client.client.datasets.Delete(request) except HttpError: _LOGGER.warning( 'Failed to clean up dataset %s in project %s', cls.dataset_id, cls.project) class ReadTests(BigQueryReadIntegrationTests): TABLE_DATA = [{ 'number': 1, 'str': 'abc' }, { 'number': 2, 'str': 'def' }, { 'number': 3, 'str': '你好' }, { 'number': 4, 'str': 'привет' }] @classmethod def setUpClass(cls): super(ReadTests, cls).setUpClass() cls.table_name = 'python_read_table' cls.create_table(cls.table_name) table_id = '{}.{}'.format(cls.dataset_id, cls.table_name) cls.query = 'SELECT number, str FROM `%s`' % table_id @classmethod def create_table(cls, table_name): table_schema = bigquery.TableSchema() table_field = bigquery.TableFieldSchema() table_field.name = 'number' table_field.type = 'INTEGER' table_schema.fields.append(table_field) table_field = bigquery.TableFieldSchema() table_field.name = 'str' table_field.type = 'STRING' table_schema.fields.append(table_field) table = bigquery.Table( tableReference=bigquery.TableReference( projectId=cls.project, datasetId=cls.dataset_id, tableId=table_name), schema=table_schema) request = bigquery.BigqueryTablesInsertRequest( projectId=cls.project, datasetId=cls.dataset_id, table=table) cls.bigquery_client.client.tables.Insert(request) # Call get_table so that we wait until the table is visible. _ = cls.bigquery_client.get_table(cls.project, cls.dataset_id, table_name) cls.bigquery_client.insert_rows( cls.project, cls.dataset_id, table_name, cls.TABLE_DATA) @skip(['PortableRunner', 'FlinkRunner']) @pytest.mark.it_postcommit def test_native_source(self): with beam.Pipeline(argv=self.args) as p: result = ( p | 'read' >> beam.io.Read( beam.io.BigQuerySource(query=self.query, use_standard_sql=True))) assert_that(result, equal_to(self.TABLE_DATA)) @pytest.mark.it_postcommit def test_iobase_source(self): query = StaticValueProvider(str, self.query) with beam.Pipeline(argv=self.args) as p: result = ( p | 'read with value provider query' >> beam.io.ReadFromBigQuery( query=query, use_standard_sql=True, project=self.project)) assert_that(result, equal_to(self.TABLE_DATA)) @pytest.mark.it_postcommit def test_table_schema_retrieve(self): the_table = bigquery_tools.BigQueryWrapper().get_table( project_id="apache-beam-testing", dataset_id="beam_bigquery_io_test", table_id="table_schema_retrieve") table = the_table.schema utype = bigquery_schema_tools.\ generate_user_type_from_bq_schema(table) with beam.Pipeline(argv=self.args) as p: result = ( p | apache_beam.io.gcp.bigquery.ReadFromBigQuery( gcs_location="gs://bqio_schema_test", dataset="beam_bigquery_io_test", table="table_schema_retrieve", project="apache-beam-testing", output_type='BEAM_ROW')) assert_that( result, equal_to([ utype( id=1, name='customer1', type='test', times=Timestamp(1633262400)), utype( id=3, name='customer1', type='test', times=Timestamp(1664798400)), utype( id=2, name='customer2', type='test', times=Timestamp(1601726400)), utype( id=4, name='customer2', type='test', times=Timestamp(1570104000)) ])) @pytest.mark.it_postcommit def test_table_schema_retrieve_specifying_only_table(self): the_table = bigquery_tools.BigQueryWrapper().get_table( project_id="apache-beam-testing", dataset_id="beam_bigquery_io_test", table_id="table_schema_retrieve") table = the_table.schema utype = bigquery_schema_tools.\ generate_user_type_from_bq_schema(table) with beam.Pipeline(argv=self.args) as p: result = ( p | apache_beam.io.gcp.bigquery.ReadFromBigQuery( gcs_location="gs://bqio_schema_test", table="apache-beam-testing:" "beam_bigquery_io_test." "table_schema_retrieve", output_type='BEAM_ROW')) assert_that( result, equal_to([ utype( id=1, name='customer1', type='test', times=Timestamp(1633262400)), utype( id=3, name='customer1', type='test', times=Timestamp(1664798400)), utype( id=2, name='customer2', type='test', times=Timestamp(1601726400)), utype( id=4, name='customer2', type='test', times=Timestamp(1570104000)) ])) @pytest.mark.it_postcommit def test_table_schema_retrieve_with_direct_read(self): the_table = bigquery_tools.BigQueryWrapper().get_table( project_id="apache-beam-testing", dataset_id="beam_bigquery_io_test", table_id="table_schema_retrieve") table = the_table.schema utype = bigquery_schema_tools.\ generate_user_type_from_bq_schema(table) with beam.Pipeline(argv=self.args) as p: result = ( p | apache_beam.io.gcp.bigquery.ReadFromBigQuery( method=beam.io.ReadFromBigQuery.Method.DIRECT_READ, table="apache-beam-testing:" "beam_bigquery_io_test." "table_schema_retrieve", output_type='BEAM_ROW')) assert_that( result, equal_to([ utype( id=1, name='customer1', type='test', times=Timestamp(1633262400)), utype( id=3, name='customer1', type='test', times=Timestamp(1664798400)), utype( id=2, name='customer2', type='test', times=Timestamp(1601726400)), utype( id=4, name='customer2', type='test', times=Timestamp(1570104000)) ])) class ReadUsingStorageApiTests(BigQueryReadIntegrationTests): BIG_QUERY_DATASET_ID = 'python_read_table_' TABLE_DATA = [{ 'number': 1, 'string': '你好', 'time': '12:44:31', 'datetime': '2018-12-31 12:44:31', 'rec': None }, { 'number': 4, 'string': 'привет', 'time': '12:44:31', 'datetime': '2018-12-31 12:44:31', 'rec': { 'rec_datetime': '2018-12-31 12:44:31', 'rec_rec': { 'rec_rec_datetime': '2018-12-31 12:44:31' } }, }] @classmethod def setUpClass(cls): super(ReadUsingStorageApiTests, cls).setUpClass() cls.table_name = '%s%d%s' % ( cls.BIG_QUERY_DATASET_ID, int(time.time()), secrets.token_hex(3)) cls._create_table(cls.table_name) table_id = '{}.{}'.format(cls.dataset_id, cls.table_name) cls.query = 'SELECT * FROM `%s`' % table_id # Materializing the newly created Table to ensure the Read API can stream. cls.temp_table_reference = cls._execute_query(cls.project, cls.query) @classmethod def tearDownClass(cls): cls.bigquery_client.clean_up_temporary_dataset(cls.project) super(ReadUsingStorageApiTests, cls).tearDownClass() @classmethod def _create_table(cls, table_name): table_schema = bigquery.TableSchema() number = bigquery.TableFieldSchema() number.name = 'number' number.type = 'INTEGER' table_schema.fields.append(number) string = bigquery.TableFieldSchema() string.name = 'string' string.type = 'STRING' table_schema.fields.append(string) time = bigquery.TableFieldSchema() time.name = 'time' time.type = 'TIME' table_schema.fields.append(time) datetime = bigquery.TableFieldSchema() datetime.name = 'datetime' datetime.type = 'DATETIME' table_schema.fields.append(datetime) rec = bigquery.TableFieldSchema() rec.name = 'rec' rec.type = 'RECORD' rec_datetime = bigquery.TableFieldSchema() rec_datetime.name = 'rec_datetime' rec_datetime.type = 'DATETIME' rec.fields.append(rec_datetime) rec_rec = bigquery.TableFieldSchema() rec_rec.name = 'rec_rec' rec_rec.type = 'RECORD' rec_rec_datetime = bigquery.TableFieldSchema() rec_rec_datetime.name = 'rec_rec_datetime' rec_rec_datetime.type = 'DATETIME' rec_rec.fields.append(rec_rec_datetime) rec.fields.append(rec_rec) table_schema.fields.append(rec) table = bigquery.Table( tableReference=bigquery.TableReference( projectId=cls.project, datasetId=cls.dataset_id, tableId=table_name), schema=table_schema) request = bigquery.BigqueryTablesInsertRequest( projectId=cls.project, datasetId=cls.dataset_id, table=table) cls.bigquery_client.client.tables.Insert(request) # Call get_table so that we wait until the table is visible. _ = cls.bigquery_client.get_table(cls.project, cls.dataset_id, table_name) cls.bigquery_client.insert_rows( cls.project, cls.dataset_id, table_name, cls.TABLE_DATA) @classmethod def _setup_temporary_dataset(cls, project, query): location = cls.bigquery_client.get_query_location(project, query, False) cls.bigquery_client.create_temporary_dataset(project, location) @classmethod def _execute_query(cls, project, query): query_job_name = bigquery_tools.generate_bq_job_name( 'materializing_table_before_reading', str(uuid.uuid4())[0:10], bigquery_tools.BigQueryJobTypes.QUERY, '%d_%s' % (int(time.time()), secrets.token_hex(3))) cls._setup_temporary_dataset(cls.project, cls.query) job = cls.bigquery_client._start_query_job( project, query, use_legacy_sql=False, flatten_results=False, job_id=query_job_name, priority=beam.io.BigQueryQueryPriority.BATCH) job_ref = job.jobReference cls.bigquery_client.wait_for_bq_job(job_ref, max_retries=0) return cls.bigquery_client._get_temp_table(project) @pytest.mark.it_postcommit def test_iobase_source(self): EXPECTED_TABLE_DATA = [ { 'number': 1, 'string': '你好', 'time': datetime.time(12, 44, 31), 'datetime': '2018-12-31T12:44:31', 'rec': None, }, { 'number': 4, 'string': 'привет', 'time': datetime.time(12, 44, 31), 'datetime': '2018-12-31T12:44:31', 'rec': { 'rec_datetime': '2018-12-31T12:44:31', 'rec_rec': { 'rec_rec_datetime': '2018-12-31T12:44:31', } }, } ] with beam.Pipeline(argv=self.args) as p: result = ( p | 'Read with BigQuery Storage API' >> beam.io.ReadFromBigQuery( method=beam.io.ReadFromBigQuery.Method.DIRECT_READ, table=self.temp_table_reference)) assert_that(result, equal_to(EXPECTED_TABLE_DATA)) @pytest.mark.it_postcommit def test_iobase_source_with_native_datetime(self): EXPECTED_TABLE_DATA = [ { 'number': 1, 'string': '你好', 'time': datetime.time(12, 44, 31), 'datetime': datetime.datetime(2018, 12, 31, 12, 44, 31), 'rec': None, }, { 'number': 4, 'string': 'привет', 'time': datetime.time(12, 44, 31), 'datetime': datetime.datetime(2018, 12, 31, 12, 44, 31), 'rec': { 'rec_datetime': datetime.datetime(2018, 12, 31, 12, 44, 31), 'rec_rec': { 'rec_rec_datetime': datetime.datetime( 2018, 12, 31, 12, 44, 31) } }, } ] with beam.Pipeline(argv=self.args) as p: result = ( p | 'Read with BigQuery Storage API' >> beam.io.ReadFromBigQuery( method=beam.io.ReadFromBigQuery.Method.DIRECT_READ, table=self.temp_table_reference, use_native_datetime=True)) assert_that(result, equal_to(EXPECTED_TABLE_DATA)) @pytest.mark.it_postcommit def test_iobase_source_with_column_selection(self): EXPECTED_TABLE_DATA = [{'number': 1}, {'number': 4}] with beam.Pipeline(argv=self.args) as p: result = ( p | 'Read with BigQuery Storage API' >> beam.io.ReadFromBigQuery( method=beam.io.ReadFromBigQuery.Method.DIRECT_READ, table=self.temp_table_reference, selected_fields=['number'])) assert_that(result, equal_to(EXPECTED_TABLE_DATA)) @pytest.mark.it_postcommit def test_iobase_source_with_row_restriction(self): EXPECTED_TABLE_DATA = [{ 'number': 1, 'string': '你好', 'time': datetime.time(12, 44, 31), 'datetime': datetime.datetime(2018, 12, 31, 12, 44, 31), 'rec': None }] with beam.Pipeline(argv=self.args) as p: result = ( p | 'Read with BigQuery Storage API' >> beam.io.ReadFromBigQuery( method=beam.io.ReadFromBigQuery.Method.DIRECT_READ, table=self.temp_table_reference, row_restriction='number < 2', use_native_datetime=True)) assert_that(result, equal_to(EXPECTED_TABLE_DATA)) @pytest.mark.it_postcommit def test_iobase_source_with_column_selection_and_row_restriction(self): EXPECTED_TABLE_DATA = [{'string': 'привет'}] with beam.Pipeline(argv=self.args) as p: result = ( p | 'Read with BigQuery Storage API' >> beam.io.ReadFromBigQuery( method=beam.io.ReadFromBigQuery.Method.DIRECT_READ, table=self.temp_table_reference, row_restriction='number > 2', selected_fields=['string'])) assert_that(result, equal_to(EXPECTED_TABLE_DATA)) @pytest.mark.it_postcommit def test_iobase_source_with_column_selection_and_row_restriction_rows(self): with beam.Pipeline(argv=self.args) as p: result = ( p | 'Read with BigQuery Storage API' >> beam.io.ReadFromBigQuery( method=beam.io.ReadFromBigQuery.Method.DIRECT_READ, table=self.temp_table_reference, row_restriction='number > 2', selected_fields=['string'], output_type='BEAM_ROW')) assert_that( result | beam.Map(lambda row: row.string), equal_to(['привет'])) @pytest.mark.it_postcommit def test_iobase_source_with_very_selective_filters(self): with beam.Pipeline(argv=self.args) as p: result = ( p | 'Read with BigQuery Storage API' >> beam.io.ReadFromBigQuery( method=beam.io.ReadFromBigQuery.Method.DIRECT_READ, project=self.temp_table_reference.projectId, dataset=self.temp_table_reference.datasetId, table=self.temp_table_reference.tableId, row_restriction='number > 4', selected_fields=['string'])) assert_that(result, equal_to([])) @pytest.mark.it_postcommit def test_iobase_source_with_query(self): EXPECTED_TABLE_DATA = [ { 'number': 1, 'string': '你好', 'time': datetime.time(12, 44, 31), 'datetime': datetime.datetime(2018, 12, 31, 12, 44, 31), 'rec': None, }, { 'number': 4, 'string': 'привет', 'time': datetime.time(12, 44, 31), 'datetime': datetime.datetime(2018, 12, 31, 12, 44, 31), 'rec': { 'rec_datetime': datetime.datetime(2018, 12, 31, 12, 44, 31), 'rec_rec': { 'rec_rec_datetime': datetime.datetime( 2018, 12, 31, 12, 44, 31) } }, } ] query = StaticValueProvider(str, self.query) with beam.Pipeline(argv=self.args) as p: result = ( p | 'Direct read with query' >> beam.io.ReadFromBigQuery( method=beam.io.ReadFromBigQuery.Method.DIRECT_READ, use_native_datetime=True, use_standard_sql=True, project=self.project, query=query)) assert_that(result, equal_to(EXPECTED_TABLE_DATA)) @pytest.mark.it_postcommit def test_iobase_source_with_query_and_filters(self): EXPECTED_TABLE_DATA = [{'string': 'привет'}] query = StaticValueProvider(str, self.query) with beam.Pipeline(argv=self.args) as p: result = ( p | 'Direct read with query' >> beam.io.ReadFromBigQuery( method=beam.io.ReadFromBigQuery.Method.DIRECT_READ, row_restriction='number > 2', selected_fields=['string'], use_standard_sql=True, project=self.project, query=query)) assert_that(result, equal_to(EXPECTED_TABLE_DATA)) class ReadNewTypesTests(BigQueryReadIntegrationTests): @classmethod def setUpClass(cls): super(ReadNewTypesTests, cls).setUpClass() cls.table_name = 'python_new_types' cls.create_table(cls.table_name) table_id = '{}.{}'.format(cls.dataset_id, cls.table_name) cls.query = 'SELECT float, numeric, bytes, date, time, datetime,' \ 'timestamp, geo FROM `%s`' % table_id @classmethod def create_table(cls, table_name): table_schema = bigquery.TableSchema() table_field = bigquery.TableFieldSchema() table_field.name = 'float' table_field.type = 'FLOAT' table_schema.fields.append(table_field) table_field = bigquery.TableFieldSchema() table_field.name = 'numeric' table_field.type = 'NUMERIC' table_schema.fields.append(table_field) table_field = bigquery.TableFieldSchema() table_field.name = 'bytes' table_field.type = 'BYTES' table_schema.fields.append(table_field) table_field = bigquery.TableFieldSchema() table_field.name = 'date' table_field.type = 'DATE' table_schema.fields.append(table_field) table_field = bigquery.TableFieldSchema() table_field.name = 'time' table_field.type = 'TIME' table_schema.fields.append(table_field) table_field = bigquery.TableFieldSchema() table_field.name = 'datetime' table_field.type = 'DATETIME' table_schema.fields.append(table_field) table_field = bigquery.TableFieldSchema() table_field.name = 'timestamp' table_field.type = 'TIMESTAMP' table_schema.fields.append(table_field) table_field = bigquery.TableFieldSchema() table_field.name = 'geo' table_field.type = 'GEOGRAPHY' table_schema.fields.append(table_field) table = bigquery.Table( tableReference=bigquery.TableReference( projectId=cls.project, datasetId=cls.dataset_id, tableId=table_name), schema=table_schema) request = bigquery.BigqueryTablesInsertRequest( projectId=cls.project, datasetId=cls.dataset_id, table=table) cls.bigquery_client.client.tables.Insert(request) # Call get_table so that we wait until the table is visible. _ = cls.bigquery_client.get_table(cls.project, cls.dataset_id, table_name) row_data = { 'float': 0.33, 'numeric': Decimal('10'), 'bytes': base64.b64encode(b'\xab\xac').decode('utf-8'), 'date': '3000-12-31', 'time': '23:59:59', 'datetime': '2018-12-31T12:44:31', 'timestamp': '2018-12-31 12:44:31.744957 UTC', 'geo': 'POINT(30 10)' } table_data = [row_data] # add rows with only one key value pair and None values for all other keys for key, value in row_data.items(): table_data.append({key: value}) cls.bigquery_client.insert_rows( cls.project, cls.dataset_id, table_name, table_data) def get_expected_data(self, native=True): byts = b'\xab\xac' expected_row = { 'float': 0.33, 'numeric': Decimal('10'), 'bytes': base64.b64encode(byts) if native else byts, 'date': '3000-12-31', 'time': '23:59:59', 'datetime': '2018-12-31T12:44:31', 'timestamp': '2018-12-31 12:44:31.744957 UTC', 'geo': 'POINT(30 10)' } expected_data = [expected_row] # add rows with only one key value pair and None values for all other keys for key, value in expected_row.items(): row = {k: None for k in expected_row} row[key] = value expected_data.append(row) return expected_data @skip(['PortableRunner', 'FlinkRunner']) @pytest.mark.it_postcommit def test_native_source(self): with beam.Pipeline(argv=self.args) as p: result = ( p | 'read' >> beam.io.Read( beam.io.BigQuerySource(query=self.query, use_standard_sql=True))) assert_that(result, equal_to(self.get_expected_data())) @pytest.mark.it_postcommit def test_iobase_source(self): with beam.Pipeline(argv=self.args) as p: result = ( p | 'read' >> beam.io.ReadFromBigQuery( query=self.query, use_standard_sql=True, project=self.project, bigquery_job_labels={'launcher': 'apache_beam_tests'}) | beam.Map(datetime_to_utc)) assert_that(result, equal_to(self.get_expected_data(native=False))) class ReadAllBQTests(BigQueryReadIntegrationTests): TABLE_DATA_AVAILABILITY_WAIT_SECONDS = 30 TABLE_DATA_1 = [{ 'number': 1, 'str': 'abc' }, { 'number': 2, 'str': 'def' }, { 'number': 3, 'str': '你好' }, { 'number': 4, 'str': 'привет' }] TABLE_DATA_2 = [{ 'number': 10, 'str': 'abcd' }, { 'number': 20, 'str': 'defg' }, { 'number': 30, 'str': '你好' }, { 'number': 40, 'str': 'привет' }] TABLE_DATA_3 = [{'number': 10, 'str': 'abcde', 'extra': 3}] @classmethod def setUpClass(cls): super(ReadAllBQTests, cls).setUpClass() cls.SCHEMA_BQ = cls.create_bq_schema() cls.SCHEMA_BQ_WITH_EXTRA = cls.create_bq_schema(True) cls.table_name1 = 'python_rd_table_1' cls.table_schema1 = cls.create_table( cls.table_name1, cls.TABLE_DATA_1, cls.SCHEMA_BQ) table_id1 = '{}.{}'.format(cls.dataset_id, cls.table_name1) cls.query1 = 'SELECT number, str FROM `%s`' % table_id1 cls.table_name2 = 'python_rd_table_2' cls.table_schema2 = cls.create_table( cls.table_name2, cls.TABLE_DATA_2, cls.SCHEMA_BQ) cls.query2 = 'SELECT number, str FROM [%s:%s.%s]' % ( cls.project, cls.dataset_id, cls.table_name2) cls.table_name3 = 'python_rd_table_3' cls.table_schema3 = cls.create_table( cls.table_name3, cls.TABLE_DATA_3, cls.SCHEMA_BQ_WITH_EXTRA) table_id3 = '{}.{}'.format(cls.dataset_id, cls.table_name3) cls.query3 = 'SELECT number, str, extra FROM `%s`' % table_id3 @classmethod def create_table(cls, table_name, data, table_schema): table = bigquery.Table( tableReference=bigquery.TableReference( projectId=cls.project, datasetId=cls.dataset_id, tableId=table_name), schema=table_schema) request = bigquery.BigqueryTablesInsertRequest( projectId=cls.project, datasetId=cls.dataset_id, table=table) cls.bigquery_client.client.tables.Insert(request) # Call get_table so that we wait until the table is visible. _ = cls.bigquery_client.get_table(cls.project, cls.dataset_id, table_name) cls.bigquery_client.insert_rows( cls.project, cls.dataset_id, table_name, data) time.sleep(cls.TABLE_DATA_AVAILABILITY_WAIT_SECONDS) return table_schema @classmethod def create_bq_schema(cls, with_extra=False): table_schema = bigquery.TableSchema() table_field = bigquery.TableFieldSchema() table_field.name = 'number' table_field.type = 'INTEGER' table_field.mode = 'NULLABLE' table_schema.fields.append(table_field) table_field = bigquery.TableFieldSchema() table_field.name = 'str' table_field.type = 'STRING' table_field.mode = 'NULLABLE' table_schema.fields.append(table_field) if with_extra: table_field = bigquery.TableFieldSchema() table_field.name = 'extra' table_field.type = 'INTEGER' table_field.mode = 'NULLABLE' table_schema.fields.append(table_field) return table_schema @skip(['PortableRunner', 'FlinkRunner']) @pytest.mark.it_postcommit def test_read_queries(self): with beam.Pipeline(argv=self.args) as p: result = ( p | beam.Create([ beam.io.ReadFromBigQueryRequest(query=self.query1), beam.io.ReadFromBigQueryRequest( query=self.query2, use_standard_sql=False), beam.io.ReadFromBigQueryRequest( table='%s.%s' % (self.dataset_id, self.table_name3)) ]) | beam.io.ReadAllFromBigQuery()) assert_that( result, equal_to(self.TABLE_DATA_1 + self.TABLE_DATA_2 + self.TABLE_DATA_3)) class ReadInteractiveRunnerTests(BigQueryReadIntegrationTests): @skip(['PortableRunner', 'FlinkRunner']) @pytest.mark.it_postcommit def test_read_in_interactive_runner(self): p = beam.Pipeline(InteractiveRunner(), argv=self.args) pcoll = p | beam.io.ReadFromBigQuery(query="SELECT 1") result = interactive_beam.collect(pcoll) assert result.iloc[0, 0] == 1 if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_read_perf_test.py0000644000175100017510000001163715206113032024314 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A performance test for reading data from a BigQuery table. Besides of the standard options, there are options with special meaning: * input_dataset - BQ dataset id. * input_table - BQ table id. The table will be created and populated with data from Synthetic Source if it does not exist. * input_options - options for Synthetic Source: num_records - number of rows to be inserted, value_size - the length of a single row, key_size - required option, but its value has no meaning. Example test run on DataflowRunner: python -m apache_beam.io.gcp.bigquery_read_perf_test \ --test-pipeline-options=" --runner=TestDataflowRunner --project=... --region=... --staging_location=gs://... --temp_location=gs://... --sdk_location=.../dist/apache-beam-x.x.x.dev0.tar.gz --publish_to_big_query=true --metrics_dataset=gs://... --metrics_table=... --input_dataset=... --input_table=... --input_options='{ \"num_records\": 1024, \"key_size\": 1, \"value_size\": 1024, }'" """ # pytype: skip-file import logging from apache_beam import Map from apache_beam import ParDo from apache_beam.io import BigQueryDisposition from apache_beam.io import BigQuerySource from apache_beam.io import Read from apache_beam.io import WriteToBigQuery from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.bigquery_tools import parse_table_schema_from_json from apache_beam.testing.load_tests.load_test import LoadTest from apache_beam.testing.load_tests.load_test_metrics_utils import CountMessages from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime from apache_beam.testing.synthetic_pipeline import SyntheticSource from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.combiners import Count # pylint: disable=wrong-import-order, wrong-import-position try: from apitools.base.py.exceptions import HttpError except ImportError: HttpError = None # pylint: enable=wrong-import-order, wrong-import-position class BigQueryReadPerfTest(LoadTest): def __init__(self): super().__init__() self.input_dataset = self.pipeline.get_option('input_dataset') self.input_table = self.pipeline.get_option('input_table') self._check_for_input_data() def _check_for_input_data(self): """Checks if a BQ table with input data exists and creates it if not.""" wrapper = BigQueryWrapper() try: wrapper.get_table(self.project_id, self.input_dataset, self.input_table) except HttpError as exn: if exn.status_code == 404: self._create_input_data() def _create_input_data(self): """ Runs an additional pipeline which creates test data and waits for its completion. """ SCHEMA = parse_table_schema_from_json( '{"fields": [{"name": "data", "type": "BYTES"}]}') def format_record(record): # Since Synthetic Source returns data as a dictionary, we should skip one # of the part import base64 return {'data': base64.b64encode(record[1])} with TestPipeline() as p: ( # pylint: disable=expression-not-assigned p | 'Produce rows' >> Read( SyntheticSource(self.parse_synthetic_source_options())) | 'Format' >> Map(format_record) | 'Write to BigQuery' >> WriteToBigQuery( dataset=self.input_dataset, table=self.input_table, schema=SCHEMA, create_disposition=BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=BigQueryDisposition.WRITE_EMPTY)) def test(self): output = ( self.pipeline | 'Read from BigQuery' >> Read( BigQuerySource(dataset=self.input_dataset, table=self.input_table)) | 'Count messages' >> ParDo(CountMessages(self.metrics_namespace)) | 'Measure time' >> ParDo(MeasureTime(self.metrics_namespace)) | 'Count' >> Count.Globally()) assert_that(output, equal_to([self.input_options['num_records']])) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) BigQueryReadPerfTest().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_schema_tools.py0000644000175100017510000001141315206113032023776 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tools used tool work with Schema types in the context of BigQuery. Classes, constants and functions in this file are experimental and have no backwards compatibility guarantees. NOTHING IN THIS FILE HAS BACKWARDS COMPATIBILITY GUARANTEES. """ import datetime from typing import Optional from typing import Sequence import numpy as np import apache_beam as beam import apache_beam.io.gcp.bigquery_tools import apache_beam.typehints.schemas import apache_beam.utils.proto_utils import apache_beam.utils.timestamp from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.portability.api import schema_pb2 from apache_beam.transforms import DoFn # BigQuery types as listed in # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types # with aliases (RECORD, BOOLEAN, FLOAT, INTEGER) as defined in # https://developers.google.com/resources/api-libraries/documentation/bigquery/v2/java/latest/com/google/api/services/bigquery/model/TableFieldSchema.html#setType-java.lang.String- BIG_QUERY_TO_PYTHON_TYPES = { "STRING": str, "INTEGER": np.int64, "FLOAT64": np.float64, "FLOAT": np.float64, "BOOLEAN": bool, "BYTES": bytes, "TIMESTAMP": apache_beam.utils.timestamp.Timestamp, "GEOGRAPHY": str, #TODO(https://github.com/apache/beam/issues/20810): # Finish mappings for all BQ types } def generate_user_type_from_bq_schema( the_table_schema, selected_fields: 'bigquery.TableSchema' = None) -> type: """Convert a schema of type TableSchema into a pcollection element. Args: the_table_schema: A BQ schema of type TableSchema selected_fields: if not None, the subset of fields to consider Returns: type: type that can be used to work with pCollections. """ the_schema = beam.io.gcp.bigquery_tools.get_dict_table_schema( the_table_schema) if the_schema == {}: raise ValueError("Encountered an empty schema") field_names_and_types = [] for field in the_schema['fields']: if selected_fields is not None and field['name'] not in selected_fields: continue if field['type'] in BIG_QUERY_TO_PYTHON_TYPES: typ = bq_field_to_type(field['type'], field['mode']) else: raise ValueError( f"Encountered " f"an unsupported type: {field['type']!r}") field_names_and_types.append((field['name'], typ)) sample_schema = beam.typehints.schemas.named_fields_to_schema( field_names_and_types) usertype = beam.typehints.schemas.named_tuple_from_schema(sample_schema) return usertype def bq_field_to_type(field, mode): if mode == 'NULLABLE' or mode is None or mode == '': return Optional[BIG_QUERY_TO_PYTHON_TYPES[field]] elif mode == 'REPEATED': return Sequence[BIG_QUERY_TO_PYTHON_TYPES[field]] elif mode == 'REQUIRED': return BIG_QUERY_TO_PYTHON_TYPES[field] else: raise ValueError(f"Encountered an unsupported mode: {mode!r}") def convert_to_usertype(table_schema, selected_fields=None): usertype = generate_user_type_from_bq_schema(table_schema, selected_fields) return beam.ParDo(BeamSchemaConversionDoFn(usertype)) class BeamSchemaConversionDoFn(DoFn): def __init__(self, pcoll_val_ctor): self._pcoll_val_ctor = pcoll_val_ctor def process(self, dict_of_tuples): for k, v in dict_of_tuples.items(): if isinstance(v, datetime.datetime): dict_of_tuples[k] = beam.utils.timestamp.Timestamp.from_utc_datetime(v) yield self._pcoll_val_ctor(**dict_of_tuples) def infer_output_type(self, input_type): return self._pcoll_val_ctor @classmethod def _from_serialized_schema(cls, schema_str): return cls( apache_beam.typehints.schemas.named_tuple_from_schema( apache_beam.utils.proto_utils.parse_Bytes( schema_str, schema_pb2.Schema))) def __reduce__(self): # when pickling, use bytes representation of the schema. return ( self._from_serialized_schema, ( beam.typehints.schemas.named_tuple_to_schema( self._pcoll_val_ctor).SerializeToString(), )) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_schema_tools_test.py0000644000175100017510000003230715206113032025042 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import logging import typing import unittest.mock import mock import numpy as np import apache_beam as beam import apache_beam.io.gcp.bigquery from apache_beam.io.gcp import bigquery_schema_tools from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.options import value_provider try: from apitools.base.py.exceptions import HttpError except ImportError: HttpError = None @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestBigQueryToSchema(unittest.TestCase): def test_check_schema_conversions(self): fields = [ bigquery.TableFieldSchema(name='stn', type='STRING', mode="NULLABLE"), bigquery.TableFieldSchema(name='temp', type='FLOAT64', mode="REPEATED"), bigquery.TableFieldSchema(name='count', type='INTEGER', mode=None) ] schema = bigquery.TableSchema(fields=fields) usertype = bigquery_schema_tools.generate_user_type_from_bq_schema( the_table_schema=schema) self.assertEqual( usertype.__annotations__, { 'stn': typing.Optional[str], 'temp': typing.Sequence[np.float64], 'count': typing.Optional[np.int64] }) def test_check_conversion_with_selected_fields(self): fields = [ bigquery.TableFieldSchema(name='stn', type='STRING', mode="NULLABLE"), bigquery.TableFieldSchema(name='temp', type='FLOAT64', mode="REPEATED"), bigquery.TableFieldSchema(name='count', type='INTEGER', mode=None) ] schema = bigquery.TableSchema(fields=fields) usertype = bigquery_schema_tools.generate_user_type_from_bq_schema( the_table_schema=schema, selected_fields=['stn', 'count']) self.assertEqual( usertype.__annotations__, { 'stn': typing.Optional[str], 'count': typing.Optional[np.int64] }) def test_check_conversion_with_empty_schema(self): fields = [] schema = bigquery.TableSchema(fields=fields) usertype = bigquery_schema_tools.generate_user_type_from_bq_schema( the_table_schema=schema) self.assertEqual(usertype.__annotations__, {}) def test_check_schema_conversions_with_timestamp(self): fields = [ bigquery.TableFieldSchema(name='stn', type='STRING', mode="NULLABLE"), bigquery.TableFieldSchema(name='temp', type='FLOAT64', mode="REPEATED"), bigquery.TableFieldSchema( name='times', type='TIMESTAMP', mode="NULLABLE") ] schema = bigquery.TableSchema(fields=fields) usertype = bigquery_schema_tools.generate_user_type_from_bq_schema( the_table_schema=schema) self.assertEqual( usertype.__annotations__, { 'stn': typing.Optional[str], 'temp': typing.Sequence[np.float64], 'times': typing.Optional[apache_beam.utils.timestamp.Timestamp] }) def test_unsupported_type(self): fields = [ bigquery.TableFieldSchema( name='number', type='DOUBLE', mode="NULLABLE"), bigquery.TableFieldSchema(name='temp', type='FLOAT64', mode="REPEATED"), bigquery.TableFieldSchema(name='count', type='INTEGER', mode=None) ] schema = bigquery.TableSchema(fields=fields) with self.assertRaisesRegex(ValueError, "Encountered an unsupported type: 'DOUBLE'"): bigquery_schema_tools.generate_user_type_from_bq_schema( the_table_schema=schema) def test_unsupported_mode(self): fields = [ bigquery.TableFieldSchema(name='number', type='INTEGER', mode="NESTED"), bigquery.TableFieldSchema(name='temp', type='FLOAT64', mode="REPEATED"), bigquery.TableFieldSchema(name='count', type='INTEGER', mode=None) ] schema = bigquery.TableSchema(fields=fields) with self.assertRaisesRegex(ValueError, "Encountered an unsupported mode: 'NESTED'"): bigquery_schema_tools.generate_user_type_from_bq_schema( the_table_schema=schema) @mock.patch.object(BigQueryWrapper, 'get_table') def test_bad_schema_public_api_export(self, get_table): fields = [ bigquery.TableFieldSchema(name='stn', type='DOUBLE', mode="NULLABLE"), bigquery.TableFieldSchema(name='temp', type='FLOAT64', mode="REPEATED"), bigquery.TableFieldSchema(name='count', type='INTEGER', mode=None) ] schema = bigquery.TableSchema(fields=fields) table = apache_beam.io.gcp.internal.clients.bigquery.\ bigquery_v2_messages.Table( schema=schema) get_table.return_value = table with self.assertRaisesRegex(ValueError, "Encountered an unsupported type: 'DOUBLE'"): p = apache_beam.Pipeline() _ = p | apache_beam.io.gcp.bigquery.ReadFromBigQuery( table="dataset.sample_table", method="EXPORT", project="project", output_type='BEAM_ROW') @mock.patch.object(BigQueryWrapper, 'get_table') def test_bad_schema_public_api_direct_read(self, get_table): fields = [ bigquery.TableFieldSchema(name='stn', type='DOUBLE', mode="NULLABLE"), bigquery.TableFieldSchema(name='temp', type='FLOAT64', mode="REPEATED"), bigquery.TableFieldSchema(name='count', type='INTEGER', mode=None) ] schema = bigquery.TableSchema(fields=fields) table = apache_beam.io.gcp.internal.clients.bigquery. \ bigquery_v2_messages.Table( schema=schema) get_table.return_value = table with self.assertRaisesRegex(ValueError, "Encountered an unsupported type: 'DOUBLE'"): p = apache_beam.Pipeline() _ = p | apache_beam.io.gcp.bigquery.ReadFromBigQuery( table="dataset.sample_table", method="DIRECT_READ", project="project", output_type='BEAM_ROW') def test_unsupported_value_provider(self): with self.assertRaisesRegex(TypeError, 'ReadFromBigQuery: table must be of type string' '; got ValueProvider instead'): p = apache_beam.Pipeline() _ = p | apache_beam.io.gcp.bigquery.ReadFromBigQuery( table=value_provider.ValueProvider(), output_type='BEAM_ROW') def test_unsupported_callable(self): def filterTable(table): if table is not None: return table res = filterTable with self.assertRaisesRegex(TypeError, 'ReadFromBigQuery: table must be of type string' '; got a callable instead'): p = apache_beam.Pipeline() _ = p | apache_beam.io.gcp.bigquery.ReadFromBigQuery( table=res, output_type='BEAM_ROW') def test_unsupported_query_export(self): with self.assertRaisesRegex( ValueError, "Both a query and an output type of 'BEAM_ROW' were specified. " "'BEAM_ROW' is not currently supported with queries."): p = apache_beam.Pipeline() _ = p | apache_beam.io.gcp.bigquery.ReadFromBigQuery( table="project:dataset.sample_table", method="EXPORT", query='SELECT name FROM dataset.sample_table', output_type='BEAM_ROW') def test_unsupported_query_direct_read(self): with self.assertRaisesRegex( ValueError, "Both a query and an output type of 'BEAM_ROW' were specified. " "'BEAM_ROW' is not currently supported with queries."): p = apache_beam.Pipeline() _ = p | apache_beam.io.gcp.bigquery.ReadFromBigQuery( table="project:dataset.sample_table", method="DIRECT_READ", query='SELECT name FROM dataset.sample_table', output_type='BEAM_ROW') def test_geography_type_support(self): """Test that GEOGRAPHY type is properly supported in schema conversion.""" fields = [ bigquery.TableFieldSchema( name='location', type='GEOGRAPHY', mode="NULLABLE"), bigquery.TableFieldSchema( name='locations', type='GEOGRAPHY', mode="REPEATED"), bigquery.TableFieldSchema( name='required_location', type='GEOGRAPHY', mode="REQUIRED") ] schema = bigquery.TableSchema(fields=fields) usertype = bigquery_schema_tools.generate_user_type_from_bq_schema( the_table_schema=schema) expected_annotations = { 'location': typing.Optional[str], 'locations': typing.Sequence[str], 'required_location': str } self.assertEqual(usertype.__annotations__, expected_annotations) def test_geography_in_bq_to_python_types_mapping(self): """Test that GEOGRAPHY is included in BIG_QUERY_TO_PYTHON_TYPES mapping.""" from apache_beam.io.gcp.bigquery_schema_tools import BIG_QUERY_TO_PYTHON_TYPES self.assertIn("GEOGRAPHY", BIG_QUERY_TO_PYTHON_TYPES) self.assertEqual(BIG_QUERY_TO_PYTHON_TYPES["GEOGRAPHY"], str) def test_geography_field_type_conversion(self): """Test bq_field_to_type function with GEOGRAPHY fields.""" from apache_beam.io.gcp.bigquery_schema_tools import bq_field_to_type # Test required GEOGRAPHY field result = bq_field_to_type("GEOGRAPHY", "REQUIRED") self.assertEqual(result, str) # Test nullable GEOGRAPHY field result = bq_field_to_type("GEOGRAPHY", "NULLABLE") self.assertEqual(result, typing.Optional[str]) # Test repeated GEOGRAPHY field result = bq_field_to_type("GEOGRAPHY", "REPEATED") self.assertEqual(result, typing.Sequence[str]) # Test GEOGRAPHY field with None mode (should default to nullable) result = bq_field_to_type("GEOGRAPHY", None) self.assertEqual(result, typing.Optional[str]) # Test GEOGRAPHY field with empty mode (should default to nullable) result = bq_field_to_type("GEOGRAPHY", "") self.assertEqual(result, typing.Optional[str]) def test_convert_to_usertype_with_geography(self): """Test convert_to_usertype function with GEOGRAPHY fields.""" schema = bigquery.TableSchema( fields=[ bigquery.TableFieldSchema( name='id', type='INTEGER', mode="REQUIRED"), bigquery.TableFieldSchema( name='location', type='GEOGRAPHY', mode="NULLABLE"), bigquery.TableFieldSchema( name='name', type='STRING', mode="REQUIRED") ]) conversion_transform = bigquery_schema_tools.convert_to_usertype(schema) # Verify the transform is created successfully self.assertIsNotNone(conversion_transform) # The transform should be a ParDo with BeamSchemaConversionDoFn self.assertIsInstance(conversion_transform, beam.ParDo) def test_beam_schema_conversion_dofn_with_geography(self): """Test BeamSchemaConversionDoFn with GEOGRAPHY data.""" from apache_beam.io.gcp.bigquery_schema_tools import BeamSchemaConversionDoFn # Create a user type with GEOGRAPHY field fields = [ bigquery.TableFieldSchema(name='id', type='INTEGER', mode="REQUIRED"), bigquery.TableFieldSchema( name='location', type='GEOGRAPHY', mode="NULLABLE") ] schema = bigquery.TableSchema(fields=fields) usertype = bigquery_schema_tools.generate_user_type_from_bq_schema(schema) # Create the DoFn dofn = BeamSchemaConversionDoFn(usertype) # Test processing a dictionary with GEOGRAPHY data input_dict = {'id': 1, 'location': 'POINT(30 10)'} results = list(dofn.process(input_dict)) self.assertEqual(len(results), 1) result = results[0] self.assertEqual(result.id, 1) self.assertEqual(result.location, 'POINT(30 10)') def test_geography_with_complex_wkt(self): """Test GEOGRAPHY type with complex Well-Known Text geometries.""" fields = [ bigquery.TableFieldSchema( name='simple_point', type='GEOGRAPHY', mode="NULLABLE"), bigquery.TableFieldSchema( name='linestring', type='GEOGRAPHY', mode="NULLABLE"), bigquery.TableFieldSchema( name='polygon', type='GEOGRAPHY', mode="NULLABLE"), bigquery.TableFieldSchema( name='multigeometry', type='GEOGRAPHY', mode="NULLABLE") ] schema = bigquery.TableSchema(fields=fields) usertype = bigquery_schema_tools.generate_user_type_from_bq_schema(schema) # All GEOGRAPHY fields should map to Optional[str] expected_annotations = { 'simple_point': typing.Optional[str], 'linestring': typing.Optional[str], 'polygon': typing.Optional[str], 'multigeometry': typing.Optional[str] } self.assertEqual(usertype.__annotations__, expected_annotations) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_test.py0000644000175100017510000032115315206113032022302 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for BigQuery sources and sinks.""" # pytype: skip-file import datetime import decimal import gc import json import logging import os import pickle import re import secrets import time import unittest import uuid import hamcrest as hc import mock import pytest import pytz import requests from parameterized import param from parameterized import parameterized import apache_beam as beam from apache_beam.internal import pickler from apache_beam.internal.gcp.json_value import to_json_value from apache_beam.io.filebasedsink_test import _TestCaseWithTempDirCleanUp from apache_beam.io.filesystems import FileSystems from apache_beam.io.gcp import bigquery as beam_bq from apache_beam.io.gcp import bigquery_tools from apache_beam.io.gcp.bigquery import MAX_INSERT_RETRIES from apache_beam.io.gcp.bigquery import ReadFromBigQuery from apache_beam.io.gcp.bigquery import TableRowJsonCoder from apache_beam.io.gcp.bigquery import WriteToBigQuery from apache_beam.io.gcp.bigquery import _StreamToBigQuery from apache_beam.io.gcp.bigquery_read_internal import _BigQueryReadSplit from apache_beam.io.gcp.bigquery_read_internal import _JsonToDictCoder from apache_beam.io.gcp.bigquery_read_internal import bigquery_export_destination_uri from apache_beam.io.gcp.bigquery_tools import JSON_COMPLIANCE_ERROR from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.bigquery_tools import RetryStrategy from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.io.gcp.pubsub import ReadFromPubSub from apache_beam.io.gcp.tests import utils from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultMatcher from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultStreamingMatcher from apache_beam.io.gcp.tests.bigquery_matcher import BigQueryTableMatcher from apache_beam.metrics.metric import Lineage from apache_beam.options import value_provider from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.value_provider import RuntimeValueProvider from apache_beam.options.value_provider import StaticValueProvider from apache_beam.runners.dataflow.test_dataflow_runner import TestDataflowRunner from apache_beam.runners.runner import PipelineState from apache_beam.testing import test_utils from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.display import DisplayData from apache_beam.transforms.display_test import DisplayDataItemMatcher # Protect against environments where bigquery library is not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from apitools.base.py.exceptions import HttpError from apitools.base.py.exceptions import HttpForbiddenError from google.api_core import exceptions from google.cloud import bigquery as gcp_bigquery from google.cloud import bigquery_storage_v1 as bq_storage from apache_beam.io.gcp.internal.clients.bigquery import bigquery_v2_client except ImportError: gcp_bigquery = None bq_storage = None HttpError = None HttpForbiddenError = None exceptions = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports _LOGGER = logging.getLogger(__name__) _ELEMENTS = [ { 'name': 'beam', 'language': 'py' }, { 'name': 'beam', 'language': 'java' }, { 'name': 'beam', 'language': 'go' }, { 'name': 'flink', 'language': 'java' }, { 'name': 'flink', 'language': 'scala' }, { 'name': 'spark', 'language': 'scala' }, { 'name': 'spark', 'language': 'py' }, { 'name': 'spark', 'language': 'scala' }, { 'name': 'beam', 'foundation': 'apache' }, { 'name': 'flink', 'foundation': 'apache' }, { 'name': 'spark', 'foundation': 'apache' }, ] def _load_or_default(filename): try: with open(filename) as f: return json.load(f) except: # pylint: disable=bare-except return {} @unittest.skipIf( HttpError is None or gcp_bigquery is None, 'GCP dependencies are not installed') class TestTableRowJsonCoder(unittest.TestCase): def test_row_as_table_row(self): schema_definition = [('s', 'STRING'), ('i', 'INTEGER'), ('f', 'FLOAT'), ('b', 'BOOLEAN'), ('n', 'NUMERIC'), ('r', 'RECORD'), ('g', 'GEOGRAPHY')] data_definition = [ 'abc', 123, 123.456, True, decimal.Decimal('987654321.987654321'), { 'a': 'b' }, 'LINESTRING(1 2, 3 4, 5 6, 7 8)' ] str_def = ( '{"s": "abc", ' '"i": 123, ' '"f": 123.456, ' '"b": true, ' '"n": "987654321.987654321", ' '"r": {"a": "b"}, ' '"g": "LINESTRING(1 2, 3 4, 5 6, 7 8)"}') schema = bigquery.TableSchema( fields=[ bigquery.TableFieldSchema(name=k, type=v) for k, v in schema_definition ]) coder = TableRowJsonCoder(table_schema=schema) def value_or_decimal_to_json(val): if isinstance(val, decimal.Decimal): return to_json_value(str(val)) else: return to_json_value(val) test_row = bigquery.TableRow( f=[ bigquery.TableCell(v=value_or_decimal_to_json(e)) for e in data_definition ]) self.assertEqual(str_def, coder.encode(test_row)) self.assertEqual(test_row, coder.decode(coder.encode(test_row))) # A coder without schema can still decode. self.assertEqual( test_row, TableRowJsonCoder().decode(coder.encode(test_row))) def test_row_and_no_schema(self): coder = TableRowJsonCoder() test_row = bigquery.TableRow( f=[ bigquery.TableCell(v=to_json_value(e)) for e in ['abc', 123, 123.456, True] ]) with self.assertRaisesRegex(AttributeError, r'^The TableRowJsonCoder requires'): coder.encode(test_row) def json_compliance_exception(self, value): with self.assertRaisesRegex(ValueError, re.escape(JSON_COMPLIANCE_ERROR)): schema_definition = [('f', 'FLOAT')] schema = bigquery.TableSchema( fields=[ bigquery.TableFieldSchema(name=k, type=v) for k, v in schema_definition ]) coder = TableRowJsonCoder(table_schema=schema) test_row = bigquery.TableRow( f=[bigquery.TableCell(v=to_json_value(value))]) coder.encode(test_row) def test_invalid_json_nan(self): self.json_compliance_exception(float('nan')) def test_invalid_json_inf(self): self.json_compliance_exception(float('inf')) def test_invalid_json_neg_inf(self): self.json_compliance_exception(float('-inf')) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestJsonToDictCoder(unittest.TestCase): @staticmethod def _make_schema(fields): def _fill_schema(fields): for field in fields: table_field = bigquery.TableFieldSchema() table_field.name, table_field.type, table_field.mode, nested_fields, \ = field if nested_fields: table_field.fields = list(_fill_schema(nested_fields)) yield table_field schema = bigquery.TableSchema() schema.fields = list(_fill_schema(fields)) return schema def test_coder_is_pickable(self): try: schema = self._make_schema([ ( 'record', 'RECORD', 'NULLABLE', [ ('float', 'FLOAT', 'NULLABLE', []), ]), ('integer', 'INTEGER', 'NULLABLE', []), ]) coder = _JsonToDictCoder(schema) pickler.loads(pickler.dumps(coder)) except pickle.PicklingError: self.fail('{} is not pickable'.format(coder.__class__.__name__)) def test_values_are_converted(self): input_row = b'{"float": "10.5", "string": "abc"}' expected_row = {'float': 10.5, 'string': 'abc'} schema = self._make_schema([('float', 'FLOAT', 'NULLABLE', []), ('string', 'STRING', 'NULLABLE', [])]) coder = _JsonToDictCoder(schema) actual = coder.decode(input_row) self.assertEqual(expected_row, actual) def test_null_fields_are_preserved(self): input_row = b'{"float": "10.5"}' expected_row = {'float': 10.5, 'string': None} schema = self._make_schema([('float', 'FLOAT', 'NULLABLE', []), ('string', 'STRING', 'NULLABLE', [])]) coder = _JsonToDictCoder(schema) actual = coder.decode(input_row) self.assertEqual(expected_row, actual) def test_record_field_is_properly_converted(self): input_row = b'{"record": {"float": "55.5"}, "integer": 10}' expected_row = {'record': {'float': 55.5}, 'integer': 10} schema = self._make_schema([ ( 'record', 'RECORD', 'NULLABLE', [ ('float', 'FLOAT', 'NULLABLE', []), ]), ('integer', 'INTEGER', 'NULLABLE', []), ]) coder = _JsonToDictCoder(schema) actual = coder.decode(input_row) self.assertEqual(expected_row, actual) def test_record_and_repeatable_field_is_properly_converted(self): input_row = b'{"record": [{"float": "55.5"}, {"float": "65.5"}], ' \ b'"integer": 10}' expected_row = {'record': [{'float': 55.5}, {'float': 65.5}], 'integer': 10} schema = self._make_schema([ ( 'record', 'RECORD', 'REPEATED', [ ('float', 'FLOAT', 'NULLABLE', []), ]), ('integer', 'INTEGER', 'NULLABLE', []), ]) coder = _JsonToDictCoder(schema) actual = coder.decode(input_row) self.assertEqual(expected_row, actual) def test_repeatable_field_is_properly_converted(self): input_row = b'{"repeated": ["55.5", "65.5"], "integer": "10"}' expected_row = {'repeated': [55.5, 65.5], 'integer': 10} schema = self._make_schema([ ('repeated', 'FLOAT', 'REPEATED', []), ('integer', 'INTEGER', 'NULLABLE', []), ]) coder = _JsonToDictCoder(schema) actual = coder.decode(input_row) self.assertEqual(expected_row, actual) @unittest.skipIf( HttpError is None or HttpForbiddenError is None, 'GCP dependencies are not installed') class TestReadFromBigQuery(unittest.TestCase): @classmethod def setUpClass(cls): class UserDefinedOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument('--gcs_location') cls.UserDefinedOptions = UserDefinedOptions def tearDown(self): # Reset runtime options to avoid side-effects caused by other tests. RuntimeValueProvider.set_runtime_options(None) @classmethod def tearDownClass(cls): # Unset the option added in setupClass to avoid interfere with other tests. # Force a gc so PipelineOptions.__subclass__() no longer contains it. del cls.UserDefinedOptions gc.collect() def test_get_destination_uri_empty_runtime_vp(self): with self.assertRaisesRegex(ValueError, '^ReadFromBigQuery requires a GCS ' 'location to be provided'): # Don't provide any runtime values. RuntimeValueProvider.set_runtime_options({}) options = self.UserDefinedOptions() bigquery_export_destination_uri( options.gcs_location, None, uuid.uuid4().hex) def test_get_destination_uri_none(self): with self.assertRaisesRegex(ValueError, '^ReadFromBigQuery requires a GCS ' 'location to be provided'): bigquery_export_destination_uri(None, None, uuid.uuid4().hex) def test_get_destination_uri_runtime_vp(self): # Provide values at job-execution time. RuntimeValueProvider.set_runtime_options({'gcs_location': 'gs://bucket'}) options = self.UserDefinedOptions() unique_id = uuid.uuid4().hex uri = bigquery_export_destination_uri(options.gcs_location, None, unique_id) self.assertEqual( uri, 'gs://bucket/' + unique_id + '/bigquery-table-dump-*.json') def test_get_destination_uri_static_vp(self): unique_id = uuid.uuid4().hex uri = bigquery_export_destination_uri( StaticValueProvider(str, 'gs://bucket'), None, unique_id) self.assertEqual( uri, 'gs://bucket/' + unique_id + '/bigquery-table-dump-*.json') def test_get_destination_uri_fallback_temp_location(self): # Don't provide any runtime values. RuntimeValueProvider.set_runtime_options({}) options = self.UserDefinedOptions() with self.assertLogs('apache_beam.io.gcp.bigquery_read_internal', level='DEBUG') as context: bigquery_export_destination_uri( options.gcs_location, 'gs://bucket', uuid.uuid4().hex) self.assertEqual( context.output, [ 'DEBUG:apache_beam.io.gcp.bigquery_read_internal:gcs_location is ' 'empty, using temp_location instead' ]) @mock.patch.object(BigQueryWrapper, '_delete_table') @mock.patch.object(BigQueryWrapper, '_delete_dataset') @mock.patch('apache_beam.io.gcp.internal.clients.bigquery.BigqueryV2') def test_temp_dataset_is_configurable( self, api, delete_dataset, delete_table): temp_dataset = bigquery.DatasetReference( projectId='temp-project', datasetId='bq_dataset') bq = BigQueryWrapper(client=api, temp_dataset_id=temp_dataset.datasetId) gcs_location = 'gs://gcs_location' c = beam.io.gcp.bigquery._CustomBigQuerySource( query='select * from test_table', gcs_location=gcs_location, method=beam.io.ReadFromBigQuery.Method.EXPORT, validate=True, pipeline_options=beam.options.pipeline_options.PipelineOptions(), job_name='job_name', step_name='step_name', project='execution_project', **{'temp_dataset': temp_dataset}) c._setup_temporary_dataset(bq) api.datasets.assert_not_called() # User provided temporary dataset should not be deleted but the temporary # table created by Beam should be deleted. bq.clean_up_temporary_dataset(temp_dataset.projectId) delete_dataset.assert_not_called() delete_table.assert_called_with( temp_dataset.projectId, temp_dataset.datasetId, mock.ANY) @parameterized.expand([ param( exception_type=exceptions.Forbidden if exceptions else None, error_message='accessDenied'), param( exception_type=exceptions.ServiceUnavailable if exceptions else None, error_message='backendError'), ]) def test_create_temp_dataset_exception(self, exception_type, error_message): # Uses the FnApiRunner to ensure errors are mocked/passed through correctly with mock.patch.object(bigquery_v2_client.BigqueryV2.JobsService, 'Insert'),\ mock.patch.object(BigQueryWrapper, 'get_or_create_dataset') as mock_insert, \ mock.patch('time.sleep'), \ self.assertRaises(Exception) as exc,\ beam.Pipeline('FnApiRunner') as p: mock_insert.side_effect = exception_type(error_message) _ = p | ReadFromBigQuery( project='apache-beam-testing', query='SELECT * FROM `project.dataset.table`', gcs_location='gs://temp_location') mock_insert.assert_called() self.assertIn(error_message, str(exc.exception)) @parameterized.expand([ # read without exception param(responses=[], expected_retries=0), # first attempt returns a Http 500 blank error and retries # second attempt returns a Http 408 blank error and retries, # third attempt passes param( responses=[ HttpForbiddenError( response={'status': 500}, content="something", url="") if HttpForbiddenError else None, HttpForbiddenError( response={'status': 408}, content="blank", url="") if HttpForbiddenError else None ], expected_retries=2), # first attempts returns a 403 rateLimitExceeded error # second attempt returns a 429 blank error # third attempt returns a Http 403 rateLimitExceeded error # fourth attempt passes param( responses=[ exceptions.Forbidden( "some message", errors=({ "message": "transient", "reason": "rateLimitExceeded" }, )) if exceptions else None, exceptions.ResourceExhausted("some message") if exceptions else None, HttpForbiddenError( response={'status': 403}, content={ "error": { "errors": [{ "message": "transient", "reason": "rateLimitExceeded" }] } }, url="") if HttpForbiddenError else None, ], expected_retries=3), ]) def test_get_table_transient_exception(self, responses, expected_retries): class DummyTable: class DummySchema: fields = [] numBytes = 5 schema = DummySchema() # TODO(https://github.com/apache/beam/issues/34549): This test relies on # lineage metrics which Prism doesn't seem to handle correctly. Defaulting # to FnApiRunner instead. with mock.patch('time.sleep'), \ mock.patch.object(bigquery_v2_client.BigqueryV2.TablesService, 'Get') as mock_get_table, \ mock.patch.object(BigQueryWrapper, 'wait_for_bq_job'), \ mock.patch.object(BigQueryWrapper, 'perform_extract_job'), \ mock.patch.object(FileSystems, 'match'), \ mock.patch.object(FileSystems, 'delete'), \ beam.Pipeline('FnApiRunner') as p: call_counter = 0 def store_callback(unused_request): nonlocal call_counter if call_counter < len(responses): exception = responses[call_counter] call_counter += 1 raise exception else: call_counter += 1 return DummyTable() mock_get_table.side_effect = store_callback _ = p | beam.io.ReadFromBigQuery( table="project.dataset.table", gcs_location="gs://some_bucket") # ReadFromBigQuery export mode calls get_table() twice. Once to get # metadata (numBytes), and once to retrieve the table's schema # Any additional calls are retries self.assertEqual(expected_retries, mock_get_table.call_count - 2) self.assertSetEqual( Lineage.query(p.result.metrics(), Lineage.SOURCE), set(["bigquery:project.dataset.table"])) @parameterized.expand([ # first attempt returns a Http 429 with transient reason and retries # second attempt returns a Http 403 with non-transient reason and fails param( responses=[ HttpForbiddenError( response={'status': 429}, content={ "error": { "errors": [{ "message": "transient", "reason": "rateLimitExceeded" }] } }, url="") if HttpForbiddenError else None, HttpForbiddenError( response={'status': 403}, content={ "error": { "errors": [{ "message": "transient", "reason": "accessDenied" }] } }, url="") if HttpForbiddenError else None ], expected_retries=1), # first attempt returns a transient 403 error and retries # second attempt returns a 403 error with bad contents and fails param( responses=[ HttpForbiddenError( response={'status': 403}, content={ "error": { "errors": [{ "message": "transient", "reason": "rateLimitExceeded" }] } }, url="") if HttpForbiddenError else None, HttpError( response={'status': 403}, content="bad contents", url="") if HttpError else None ], expected_retries=1), # first attempt returns a transient 403 error and retries # second attempt returns a 429 error and retries # third attempt returns a 403 with non-transient reason and fails param( responses=[ exceptions.Forbidden( "some error", errors=({ "message": "transient", "reason": "rateLimitExceeded" }, )) if exceptions else None, exceptions.ResourceExhausted("some transient error") if exceptions else None, exceptions.Forbidden( "some error", errors=({ "message": "transient", "reason": "accessDenied" }, )) if exceptions else None, ], expected_retries=2), ]) def test_get_table_non_transient_exception(self, responses, expected_retries): class DummyTable: class DummySchema: fields = [] numBytes = 5 schema = DummySchema() with mock.patch('time.sleep'), \ mock.patch.object(bigquery_v2_client.BigqueryV2.TablesService, 'Get') as mock_get_table, \ mock.patch.object(BigQueryWrapper, 'wait_for_bq_job'), \ mock.patch.object(BigQueryWrapper, 'perform_extract_job'), \ mock.patch.object(FileSystems, 'match'), \ mock.patch.object(FileSystems, 'delete'), \ self.assertRaises(Exception), \ beam.Pipeline() as p: call_counter = 0 def store_callback(unused_request): nonlocal call_counter if call_counter < len(responses): exception = responses[call_counter] call_counter += 1 raise exception else: call_counter += 1 return DummyTable() mock_get_table.side_effect = store_callback _ = p | beam.io.ReadFromBigQuery( table="project.dataset.table", gcs_location="gs://some_bucket") # ReadFromBigQuery export mode calls get_table() twice. Once to get # metadata (numBytes), and once to retrieve the table's schema # However, the second call is never reached because this test will always # fail before it does so # After the first call, any additional calls are retries self.assertEqual(expected_retries, mock_get_table.call_count - 1) @parameterized.expand([ param( exception_type=exceptions.BadRequest if exceptions else None, error_message='invalidQuery'), param( exception_type=exceptions.NotFound if exceptions else None, error_message='notFound'), param( exception_type=exceptions.Forbidden if exceptions else None, error_message='responseTooLarge') ]) def test_query_job_exception(self, exception_type, error_message): # TODO(https://github.com/apache/beam/issues/34549): This test relies on # mocking which prism doesn't seem to fully handle correctly (mocks get # mixed between test runs). Pinning to FnApiRunner for now. with mock.patch.object(beam.io.gcp.bigquery._CustomBigQuerySource, 'estimate_size') as mock_estimate,\ mock.patch.object(BigQueryWrapper, 'get_query_location') as mock_query_location,\ mock.patch.object(bigquery_v2_client.BigqueryV2.JobsService, 'Insert') as mock_query_job,\ mock.patch.object(bigquery_v2_client.BigqueryV2.DatasetsService, 'Get'), \ mock.patch('time.sleep'), \ self.assertRaises(Exception) as exc, \ beam.Pipeline('FnApiRunner') as p: mock_estimate.return_value = None mock_query_location.return_value = None mock_query_job.side_effect = exception_type(error_message) _ = p | ReadFromBigQuery( query='SELECT * FROM `project.dataset.table`', gcs_location='gs://temp_location') mock_query_job.assert_called() self.assertIn(error_message, exc.exception.args[0]) @parameterized.expand([ param( exception_type=exceptions.BadRequest if exceptions else None, error_message='invalid'), param( exception_type=exceptions.Forbidden if exceptions else None, error_message='accessDenied') ]) def test_read_export_exception(self, exception_type, error_message): with mock.patch.object(beam.io.gcp.bigquery._CustomBigQuerySource, 'estimate_size') as mock_estimate,\ mock.patch.object(bigquery_v2_client.BigqueryV2.TablesService, 'Get'),\ mock.patch.object(bigquery_v2_client.BigqueryV2.JobsService, 'Insert') as mock_query_job, \ mock.patch('time.sleep'), \ self.assertRaises(Exception) as exc,\ beam.Pipeline() as p: mock_estimate.return_value = None mock_query_job.side_effect = exception_type(error_message) _ = p | ReadFromBigQuery( project='apache-beam-testing', method=ReadFromBigQuery.Method.EXPORT, table='project:dataset.table', gcs_location="gs://temp_location") mock_query_job.assert_called() self.assertIn(error_message, str(exc.exception)) def test_read_direct_lineage(self): # TODO(https://github.com/apache/beam/issues/34549): This test relies on # lineage metrics which Prism doesn't seem to handle correctly. Defaulting # to FnApiRunner instead. with mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client'),\ mock.patch.object(bq_storage.BigQueryReadClient, 'create_read_session'),\ beam.Pipeline('FnApiRunner') as p: _ = p | ReadFromBigQuery( method=ReadFromBigQuery.Method.DIRECT_READ, table='project:dataset.table') self.assertSetEqual( Lineage.query(p.result.metrics(), Lineage.SOURCE), set(["bigquery:project.dataset.table"])) def test_read_all_lineage(self): # TODO(https://github.com/apache/beam/issues/34549): This test relies on # lineage metrics which Prism doesn't seem to handle correctly. Defaulting # to FnApiRunner instead. with mock.patch.object(_BigQueryReadSplit, '_export_files') as export, \ beam.Pipeline('FnApiRunner') as p: export.return_value = (None, []) _ = ( p | beam.Create([ beam.io.ReadFromBigQueryRequest(table='project1:dataset1.table1'), beam.io.ReadFromBigQueryRequest(table='project2:dataset2.table2') ]) | beam.io.ReadAllFromBigQuery(gcs_location='gs://bucket/tmp')) self.assertSetEqual( Lineage.query(p.result.metrics(), Lineage.SOURCE), set([ 'bigquery:project1.dataset1.table1', 'bigquery:project2.dataset2.table2' ])) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestBigQuerySink(unittest.TestCase): def test_table_spec_display_data(self): sink = beam.io.BigQuerySink('dataset.table') dd = DisplayData.create_from(sink) expected_items = [ DisplayDataItemMatcher('table', 'dataset.table'), DisplayDataItemMatcher('validation', False) ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_parse_schema_descriptor(self): sink = beam.io.BigQuerySink('dataset.table', schema='s:STRING, n:INTEGER') self.assertEqual(sink.table_reference.datasetId, 'dataset') self.assertEqual(sink.table_reference.tableId, 'table') result_schema = { field['name']: field['type'] for field in sink.schema['fields'] } self.assertEqual({'n': 'INTEGER', 's': 'STRING'}, result_schema) def test_project_table_display_data(self): sinkq = beam.io.BigQuerySink('project:dataset.table') dd = DisplayData.create_from(sinkq) expected_items = [ DisplayDataItemMatcher('table', 'project:dataset.table'), DisplayDataItemMatcher('validation', False) ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestWriteToBigQuery(unittest.TestCase): def _cleanup_files(self): if os.path.exists('insert_calls1'): os.remove('insert_calls1') if os.path.exists('insert_calls2'): os.remove('insert_calls2') def setUp(self): self._cleanup_files() def tearDown(self): self._cleanup_files() def test_noop_schema_parsing(self): expected_table_schema = None table_schema = beam.io.gcp.bigquery.BigQueryWriteFn.get_table_schema( schema=None) self.assertEqual(expected_table_schema, table_schema) def test_dict_schema_parsing(self): schema = { 'fields': [{ 'name': 's', 'type': 'STRING', 'mode': 'NULLABLE' }, { 'name': 'n', 'type': 'INTEGER', 'mode': 'NULLABLE' }, { 'name': 'r', 'type': 'RECORD', 'mode': 'NULLABLE', 'fields': [{ 'name': 'x', 'type': 'INTEGER', 'mode': 'NULLABLE' }] }] } table_schema = beam.io.gcp.bigquery.BigQueryWriteFn.get_table_schema(schema) string_field = bigquery.TableFieldSchema( name='s', type='STRING', mode='NULLABLE') nested_field = bigquery.TableFieldSchema( name='x', type='INTEGER', mode='NULLABLE') number_field = bigquery.TableFieldSchema( name='n', type='INTEGER', mode='NULLABLE') record_field = bigquery.TableFieldSchema( name='r', type='RECORD', mode='NULLABLE', fields=[nested_field]) expected_table_schema = bigquery.TableSchema( fields=[string_field, number_field, record_field]) self.assertEqual(expected_table_schema, table_schema) def test_string_schema_parsing(self): schema = 's:STRING, n:INTEGER' expected_dict_schema = { 'fields': [{ 'name': 's', 'type': 'STRING', 'mode': 'NULLABLE' }, { 'name': 'n', 'type': 'INTEGER', 'mode': 'NULLABLE' }] } dict_schema = ( beam.io.gcp.bigquery.WriteToBigQuery.get_dict_table_schema(schema)) self.assertEqual(expected_dict_schema, dict_schema) def test_table_schema_parsing(self): string_field = bigquery.TableFieldSchema( name='s', type='STRING', mode='NULLABLE') nested_field = bigquery.TableFieldSchema( name='x', type='INTEGER', mode='NULLABLE') number_field = bigquery.TableFieldSchema( name='n', type='INTEGER', mode='NULLABLE') record_field = bigquery.TableFieldSchema( name='r', type='RECORD', mode='NULLABLE', fields=[nested_field]) schema = bigquery.TableSchema( fields=[string_field, number_field, record_field]) expected_dict_schema = { 'fields': [{ 'name': 's', 'type': 'STRING', 'mode': 'NULLABLE' }, { 'name': 'n', 'type': 'INTEGER', 'mode': 'NULLABLE' }, { 'name': 'r', 'type': 'RECORD', 'mode': 'NULLABLE', 'fields': [{ 'name': 'x', 'type': 'INTEGER', 'mode': 'NULLABLE' }] }] } dict_schema = ( beam.io.gcp.bigquery.WriteToBigQuery.get_dict_table_schema(schema)) self.assertEqual(expected_dict_schema, dict_schema) def test_table_schema_parsing_end_to_end(self): string_field = bigquery.TableFieldSchema( name='s', type='STRING', mode='NULLABLE') nested_field = bigquery.TableFieldSchema( name='x', type='INTEGER', mode='NULLABLE') number_field = bigquery.TableFieldSchema( name='n', type='INTEGER', mode='NULLABLE') record_field = bigquery.TableFieldSchema( name='r', type='RECORD', mode='NULLABLE', fields=[nested_field]) schema = bigquery.TableSchema( fields=[string_field, number_field, record_field]) table_schema = beam.io.gcp.bigquery.BigQueryWriteFn.get_table_schema( beam.io.gcp.bigquery.WriteToBigQuery.get_dict_table_schema(schema)) self.assertEqual(table_schema, schema) def test_none_schema_parsing(self): schema = None expected_dict_schema = None dict_schema = ( beam.io.gcp.bigquery.WriteToBigQuery.get_dict_table_schema(schema)) self.assertEqual(expected_dict_schema, dict_schema) def test_noop_dict_schema_parsing(self): schema = { 'fields': [{ 'name': 's', 'type': 'STRING', 'mode': 'NULLABLE' }, { 'name': 'n', 'type': 'INTEGER', 'mode': 'NULLABLE' }] } expected_dict_schema = schema dict_schema = ( beam.io.gcp.bigquery.WriteToBigQuery.get_dict_table_schema(schema)) self.assertEqual(expected_dict_schema, dict_schema) def test_schema_autodetect_not_allowed_with_avro_file_loads(self): with TestPipeline() as p: pc = p | beam.Impulse() with self.assertRaisesRegex(ValueError, '^A schema must be provided'): _ = ( pc | 'No Schema' >> beam.io.gcp.bigquery.WriteToBigQuery( "dataset.table", schema=None, temp_file_format=bigquery_tools.FileFormat.AVRO)) with self.assertRaisesRegex(ValueError, '^Schema auto-detection is not supported'): _ = ( pc | 'Schema Autodetected' >> beam.io.gcp.bigquery.WriteToBigQuery( "dataset.table", schema=beam.io.gcp.bigquery.SCHEMA_AUTODETECT, temp_file_format=bigquery_tools.FileFormat.AVRO)) def test_to_from_runner_api(self): """Tests that serialization of WriteToBigQuery is correct. This is not intended to be a change-detector test. As such, this only tests the more complicated serialization logic of parameters: ValueProviders, callables, and side inputs. """ FULL_OUTPUT_TABLE = 'test_project:output_table' p = TestPipeline() # Used for testing side input parameters. table_record_pcv = beam.pvalue.AsDict( p | "MakeTable" >> beam.Create([('table', FULL_OUTPUT_TABLE)])) # Used for testing value provider parameters. schema = value_provider.StaticValueProvider(str, '"a:str"') original = WriteToBigQuery( table=lambda _, side_input: side_input['table'], table_side_inputs=(table_record_pcv, ), schema=schema) # pylint: disable=expression-not-assigned p | beam.Create([]) | 'MyWriteToBigQuery' >> original # Run the pipeline through to generate a pipeline proto from an empty # context. This ensures that the serialization code ran. pipeline_proto, context = TestPipeline.from_runner_api( p.to_runner_api(), p.runner, p.get_pipeline_options()).to_runner_api( return_context=True) # Find the transform from the context. write_to_bq_id = [ k for k, v in pipeline_proto.components.transforms.items() if v.unique_name == 'MyWriteToBigQuery' ][0] deserialized_node = context.transforms.get_by_id(write_to_bq_id) deserialized = deserialized_node.transform self.assertIsInstance(deserialized, WriteToBigQuery) # Test that the serialization of a value provider is correct. self.assertEqual(original.schema, deserialized.schema) # Test that the serialization of a callable is correct. self.assertEqual( deserialized._table(None, {'table': FULL_OUTPUT_TABLE}), FULL_OUTPUT_TABLE) # Test that the serialization of a side input is correct. self.assertEqual( len(original.table_side_inputs), len(deserialized.table_side_inputs)) original_side_input_data = original.table_side_inputs[0]._side_input_data() deserialized_side_input_data = deserialized.table_side_inputs[ 0]._side_input_data() self.assertEqual( original_side_input_data.access_pattern, deserialized_side_input_data.access_pattern) self.assertEqual( original_side_input_data.window_mapping_fn, deserialized_side_input_data.window_mapping_fn) self.assertEqual( original_side_input_data.view_fn, deserialized_side_input_data.view_fn) def test_streaming_triggering_frequency_without_auto_sharding(self): def noop(table, **kwargs): return [] client = mock.Mock() client.insert_rows_json = mock.Mock(side_effect=noop) opt = StandardOptions() opt.streaming = True with self.assertRaises(ValueError, msg="triggering_frequency with STREAMING_INSERTS" + "can only be used with with_auto_sharding=True"): with beam.Pipeline(runner='BundleBasedDirectRunner', options=opt) as p: _ = ( p | beam.Create([{ 'columnA': 'value1' }]) | WriteToBigQuery( table='project:dataset.table', schema={ 'fields': [{ 'name': 'columnA', 'type': 'STRING', 'mode': 'NULLABLE' }] }, create_disposition='CREATE_NEVER', triggering_frequency=1, with_auto_sharding=False, test_client=client)) def test_streaming_triggering_frequency_with_auto_sharding(self): def noop(table, **kwargs): return [] client = mock.Mock() client.insert_rows_json = mock.Mock(side_effect=noop) opt = StandardOptions() opt.streaming = True with beam.Pipeline(runner='BundleBasedDirectRunner', options=opt) as p: _ = ( p | beam.Create([{ 'columnA': 'value1' }]) | WriteToBigQuery( table='project:dataset.table', schema={ 'fields': [{ 'name': 'columnA', 'type': 'STRING', 'mode': 'NULLABLE' }] }, create_disposition='CREATE_NEVER', triggering_frequency=1, with_auto_sharding=True, test_client=client)) @mock.patch('google.cloud.bigquery.Client.insert_rows_json') def test_streaming_inserts_flush_on_byte_size_limit(self, mock_insert): mock_insert.return_value = [] table = 'project:dataset.table' rows = [ { 'columnA': 'value1' }, { 'columnA': 'value2' }, # this very large row exceeds max size, so should be sent to DLQ { 'columnA': "large_string" * 100 } ] with beam.Pipeline() as p: failed_rows = ( p | beam.Create(rows) | WriteToBigQuery( table=table, method='STREAMING_INSERTS', create_disposition='CREATE_NEVER', schema='columnA:STRING', max_insert_payload_size=500)) expected_failed_rows = [(table, rows[2])] assert_that(failed_rows.failed_rows, equal_to(expected_failed_rows)) self.assertEqual(2, mock_insert.call_count) def test_max_retries_exceeds_limit(self): table = 'project:dataset.table' rows = [{'columnA': 'value1'}, {'columnA': 'value2'}] with beam.Pipeline() as p: data = p | beam.Create(rows) with self.assertRaises(ValueError) as context: _ = data | 'WriteToBQ' >> WriteToBigQuery( table=table, schema='columnA:STRING', method='STREAMING_INSERTS', max_retries=MAX_INSERT_RETRIES + 1 # Exceeds the limit of 10000 ) self.assertIn( 'max_retries cannot be more than 10000', str(context.exception)) @parameterized.expand([ param( exception_type=exceptions.Forbidden if exceptions else None, error_message='accessDenied'), param( exception_type=exceptions.ServiceUnavailable if exceptions else None, error_message='backendError') ]) @unittest.skip('Not compatible with new GCS client. See GH issue #26334.') def test_load_job_exception(self, exception_type, error_message): with mock.patch.object(bigquery_v2_client.BigqueryV2.JobsService, 'Insert') as mock_load_job,\ mock.patch('apache_beam.io.gcp.internal.clients' '.storage.storage_v1_client.StorageV1.ObjectsService'),\ mock.patch('time.sleep'),\ self.assertRaises(Exception) as exc,\ beam.Pipeline() as p: mock_load_job.side_effect = exception_type(error_message) _ = ( p | beam.Create([{ 'columnA': 'value1' }]) | WriteToBigQuery( table='project:dataset.table', schema={ 'fields': [{ 'name': 'columnA', 'type': 'STRING', 'mode': 'NULLABLE' }] }, create_disposition='CREATE_NEVER', custom_gcs_temp_location="gs://temp_location", method='FILE_LOADS')) mock_load_job.assert_called() self.assertIn(error_message, exc.exception.args[0]) @parameterized.expand([ param( exception_type=exceptions.ServiceUnavailable if exceptions else None, error_message='backendError'), param( exception_type=exceptions.InternalServerError if exceptions else None, error_message='internalError'), ]) @unittest.skip('Not compatible with new GCS client. See GH issue #26334.') def test_copy_load_job_exception(self, exception_type, error_message): from apache_beam.io.gcp import bigquery_file_loads old_max_file_size = bigquery_file_loads._DEFAULT_MAX_FILE_SIZE old_max_partition_size = bigquery_file_loads._MAXIMUM_LOAD_SIZE old_max_files_per_partition = bigquery_file_loads._MAXIMUM_SOURCE_URIS bigquery_file_loads._DEFAULT_MAX_FILE_SIZE = 15 bigquery_file_loads._MAXIMUM_LOAD_SIZE = 30 bigquery_file_loads._MAXIMUM_SOURCE_URIS = 1 with mock.patch.object(bigquery_v2_client.BigqueryV2.JobsService, 'Insert') as mock_insert_copy_job, \ mock.patch.object(BigQueryWrapper, 'perform_load_job') as mock_load_job, \ mock.patch.object(BigQueryWrapper, 'wait_for_bq_job'), \ mock.patch('apache_beam.io.gcp.internal.clients' '.storage.storage_v1_client.StorageV1.ObjectsService'),\ mock.patch('time.sleep'), \ self.assertRaises(Exception) as exc, \ beam.Pipeline() as p: mock_insert_copy_job.side_effect = exception_type(error_message) dummy_job_reference = beam.io.gcp.internal.clients.bigquery.JobReference() dummy_job_reference.jobId = 'job_id' dummy_job_reference.location = 'US' dummy_job_reference.projectId = 'apache-beam-testing' mock_load_job.return_value = dummy_job_reference _ = ( p | beam.Create([{ 'columnA': 'value1' }, { 'columnA': 'value2' }, { 'columnA': 'value3' }]) | WriteToBigQuery( table='project:dataset.table', schema={ 'fields': [{ 'name': 'columnA', 'type': 'STRING', 'mode': 'NULLABLE' }] }, create_disposition='CREATE_NEVER', custom_gcs_temp_location="gs://temp_location", method='FILE_LOADS')) bigquery_file_loads._DEFAULT_MAX_FILE_SIZE = old_max_file_size bigquery_file_loads._MAXIMUM_LOAD_SIZE = old_max_partition_size bigquery_file_loads._MAXIMUM_SOURCE_URIS = old_max_files_per_partition self.assertEqual(4, mock_insert_copy_job.call_count) self.assertIn(error_message, exc.exception.args[0]) @unittest.skipIf( HttpError is None or exceptions is None, 'GCP dependencies are not installed') class BigQueryStreamingInsertsErrorHandling(unittest.TestCase): # Running tests with a variety of exceptions from https://googleapis.dev # /python/google-api-core/latest/_modules/google/api_core/exceptions.html. # Choosing some exceptions that produce reasons included in # bigquery_tools._NON_TRANSIENT_ERRORS and some that are not @parameterized.expand([ # reason not in _NON_TRANSIENT_ERRORS for row 1 on first attempt # transient error retried and succeeds on second attempt, 0 rows sent to # failed rows param( insert_response=[ exceptions.TooManyRequests if exceptions else None, None ], error_reason='Too Many Requests', # not in _NON_TRANSIENT_ERRORS failed_rows=[]), # reason not in _NON_TRANSIENT_ERRORS for row 1 on both attempts, sent to # failed rows after hitting max_retries param( insert_response=[ exceptions.InternalServerError if exceptions else None, exceptions.InternalServerError if exceptions else None ], error_reason='Internal Server Error', # not in _NON_TRANSIENT_ERRORS failed_rows=['value1', 'value3', 'value5']), # reason in _NON_TRANSIENT_ERRORS for row 1 on both attempts, sent to # failed_rows after hitting max_retries param( insert_response=[ exceptions.Forbidden if exceptions else None, exceptions.Forbidden if exceptions else None ], error_reason='Forbidden', # in _NON_TRANSIENT_ERRORS failed_rows=['value1', 'value3', 'value5']), ]) def test_insert_rows_json_exception_retry_always( self, insert_response, error_reason, failed_rows): # In this test, a pipeline will always retry all caught exception types # since RetryStrategy is not set and defaults to RETRY_ALWAYS with mock.patch('time.sleep'): call_counter = 0 mock_response = mock.Mock() mock_response.reason = error_reason def store_callback(table, **kwargs): nonlocal call_counter # raise exception if insert_response element is an exception if insert_response[call_counter]: exception_type = insert_response[call_counter] call_counter += 1 raise exception_type('some exception', response=mock_response) # return empty list if not insert_response element, indicating # successful call to insert_rows_json else: call_counter += 1 return [] client = mock.Mock() client.insert_rows_json.side_effect = store_callback # Using the bundle based direct runner to avoid pickling problems # with mocks. with beam.Pipeline(runner='BundleBasedDirectRunner') as p: bq_write_out = ( p | beam.Create([{ 'columnA': 'value1', 'columnB': 'value2' }, { 'columnA': 'value3', 'columnB': 'value4' }, { 'columnA': 'value5', 'columnB': 'value6' }]) # Using _StreamToBigQuery in order to be able to pass max_retries # in order to limit run time of test with RETRY_ALWAYS | _StreamToBigQuery( table_reference='project:dataset.table', table_side_inputs=[], schema_side_inputs=[], schema='anyschema', batch_size=None, triggering_frequency=None, create_disposition='CREATE_NEVER', write_disposition=None, kms_key=None, retry_strategy=RetryStrategy.RETRY_ALWAYS, additional_bq_parameters=[], ignore_insert_ids=False, ignore_unknown_columns=False, with_auto_sharding=False, test_client=client, max_retries=len(insert_response) - 1, num_streaming_keys=500)) failed_values = ( bq_write_out[beam_bq.BigQueryWriteFn.FAILED_ROWS] | beam.Map(lambda x: x[1]['columnA'])) assert_that(failed_values, equal_to(failed_rows)) # Running tests with a variety of exceptions from https://googleapis.dev # /python/google-api-core/latest/_modules/google/api_core/exceptions.html. # Choosing some exceptions that produce reasons that are included in # bigquery_tools._NON_TRANSIENT_ERRORS and some that are not @parameterized.expand([ param( # not in _NON_TRANSIENT_ERRORS exception_type=exceptions.BadGateway if exceptions else None, error_reason='Bad Gateway'), param( # in _NON_TRANSIENT_ERRORS exception_type=exceptions.Unauthorized if exceptions else None, error_reason='Unauthorized'), ]) @mock.patch('time.sleep') @mock.patch('google.cloud.bigquery.Client.insert_rows_json') def test_insert_rows_json_exception_retry_never( self, mock_send, unused_mock_sleep, exception_type, error_reason): # In this test, a pipeline will never retry caught exception types # since RetryStrategy is set to RETRY_NEVER mock_response = mock.Mock() mock_response.reason = error_reason mock_send.side_effect = [ exception_type('some exception', response=mock_response) ] with beam.Pipeline(runner='BundleBasedDirectRunner') as p: bq_write_out = ( p | beam.Create([{ 'columnA': 'value1' }, { 'columnA': 'value2' }]) | WriteToBigQuery( table='project:dataset.table', schema={ 'fields': [{ 'name': 'columnA', 'type': 'STRING', 'mode': 'NULLABLE' }] }, create_disposition='CREATE_NEVER', method='STREAMING_INSERTS', insert_retry_strategy=RetryStrategy.RETRY_NEVER)) failed_values = ( bq_write_out[beam_bq.BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS] | beam.Map(lambda x: x[1]['columnA'])) assert_that(failed_values, equal_to(['value1', 'value2'])) self.assertEqual(1, mock_send.call_count) # Running tests with a variety of exceptions from https://googleapis.dev # /python/google-api-core/latest/_modules/google/api_core/exceptions.html. # Choosing some exceptions that produce reasons that are included in # bigquery_tools._NON_TRANSIENT_ERRORS and some that are not @parameterized.expand([ param( exception_type=exceptions.DeadlineExceeded if exceptions else None, error_reason='Deadline Exceeded', # not in _NON_TRANSIENT_ERRORS failed_values=[], expected_call_count=2), param( exception_type=exceptions.Conflict if exceptions else None, error_reason='Conflict', # not in _NON_TRANSIENT_ERRORS failed_values=[], expected_call_count=2), param( exception_type=exceptions.TooManyRequests if exceptions else None, error_reason='Too Many Requests', # not in _NON_TRANSIENT_ERRORS failed_values=[], expected_call_count=2), param( exception_type=exceptions.InternalServerError if exceptions else None, error_reason='Internal Server Error', # not in _NON_TRANSIENT_ERRORS failed_values=[], expected_call_count=2), param( exception_type=exceptions.BadGateway if exceptions else None, error_reason='Bad Gateway', # not in _NON_TRANSIENT_ERRORS failed_values=[], expected_call_count=2), param( exception_type=exceptions.ServiceUnavailable if exceptions else None, error_reason='Service Unavailable', # not in _NON_TRANSIENT_ERRORS failed_values=[], expected_call_count=2), param( exception_type=exceptions.GatewayTimeout if exceptions else None, error_reason='Gateway Timeout', # not in _NON_TRANSIENT_ERRORS failed_values=[], expected_call_count=2), param( exception_type=exceptions.BadRequest if exceptions else None, error_reason='Bad Request', # in _NON_TRANSIENT_ERRORS failed_values=['value1', 'value2'], expected_call_count=1), param( exception_type=exceptions.Unauthorized if exceptions else None, error_reason='Unauthorized', # in _NON_TRANSIENT_ERRORS failed_values=['value1', 'value2'], expected_call_count=1), param( exception_type=exceptions.Forbidden if exceptions else None, error_reason='Forbidden', # in _NON_TRANSIENT_ERRORS failed_values=['value1', 'value2'], expected_call_count=1), param( exception_type=exceptions.NotFound if exceptions else None, error_reason='Not Found', # in _NON_TRANSIENT_ERRORS failed_values=['value1', 'value2'], expected_call_count=1), param( exception_type=exceptions.MethodNotImplemented if exceptions else None, error_reason='Not Implemented', # in _NON_TRANSIENT_ERRORS failed_values=['value1', 'value2'], expected_call_count=1), ]) @mock.patch('time.sleep') @mock.patch('google.cloud.bigquery.Client.insert_rows_json') def test_insert_rows_json_exception_retry_on_transient_error( self, mock_send, unused_mock_sleep, exception_type, error_reason, failed_values, expected_call_count): # In this test, a pipeline will only retry caught exception types # with reasons that are not in _NON_TRANSIENT_ERRORS since RetryStrategy is # set to RETRY_ON_TRANSIENT_ERROR mock_response = mock.Mock() mock_response.reason = error_reason mock_send.side_effect = [ exception_type('some exception', response=mock_response), # Return no exception and no errors on 2nd call, if there is a 2nd call [] ] with beam.Pipeline(runner='BundleBasedDirectRunner') as p: bq_write_out = ( p | beam.Create([{ 'columnA': 'value1' }, { 'columnA': 'value2' }]) | WriteToBigQuery( table='project:dataset.table', schema={ 'fields': [{ 'name': 'columnA', 'type': 'STRING', 'mode': 'NULLABLE' }] }, create_disposition='CREATE_NEVER', method='STREAMING_INSERTS', insert_retry_strategy=RetryStrategy.RETRY_ON_TRANSIENT_ERROR)) failed_values_out = ( bq_write_out[beam_bq.BigQueryWriteFn.FAILED_ROWS] | beam.Map(lambda x: x[1]['columnA'])) assert_that(failed_values_out, equal_to(failed_values)) self.assertEqual(expected_call_count, mock_send.call_count) # Running tests with persistent exceptions with exception types not # caught in BigQueryWrapper._insert_all_rows but retriable by # retry.with_exponential_backoff @parameterized.expand([ param( exception_type=requests.exceptions.ConnectionError, error_message='some connection error'), param( exception_type=requests.exceptions.Timeout, error_message='some timeout error'), ]) @mock.patch('time.sleep') @mock.patch('google.cloud.bigquery.Client.insert_rows_json') def test_insert_rows_json_persistent_retriable_exception( self, mock_send, unused_mock_sleep, exception_type, error_message): # In this test, each insert_rows_json call will result in an exception # and be retried with retry.with_exponential_backoff until MAX_RETRIES is # reached mock_send.side_effect = exception_type(error_message) # Expecting 1 initial call plus maximum number of retries expected_call_count = 1 + bigquery_tools.MAX_RETRIES with self.assertRaises(Exception) as exc: with beam.Pipeline() as p: _ = ( p | beam.Create([{ 'columnA': 'value1' }, { 'columnA': 'value2' }]) | WriteToBigQuery( table='project:dataset.table', schema={ 'fields': [{ 'name': 'columnA', 'type': 'STRING', 'mode': 'NULLABLE' }] }, create_disposition='CREATE_NEVER', method='STREAMING_INSERTS')) self.assertEqual(expected_call_count, mock_send.call_count) self.assertIn(error_message, exc.exception.args[0]) # Running tests with intermittent exceptions with exception types not # caught in BigQueryWrapper._insert_all_rows but retriable by # retry.with_exponential_backoff @parameterized.expand([ param( exception_type=requests.exceptions.ConnectionError, error_message='some connection error'), param( exception_type=requests.exceptions.Timeout, error_message='some timeout error'), ]) @mock.patch('time.sleep') @mock.patch('google.cloud.bigquery.Client.insert_rows_json') def test_insert_rows_json_intermittent_retriable_exception( self, mock_send, unused_mock_sleep, exception_type, error_message): # In this test, the first 2 insert_rows_json calls will result in an # exception and be retried with retry.with_exponential_backoff. The last # call will not raise an exception and will succeed. mock_send.side_effect = [ exception_type(error_message), exception_type(error_message), [] ] # This relies on DirectRunner-specific mocking behavior which can be # inconsistent on Prism with beam.Pipeline('FnApiRunner') as p: _ = ( p | beam.Create([{ 'columnA': 'value1' }, { 'columnA': 'value2' }]) | WriteToBigQuery( table='project:dataset.table', schema={ 'fields': [{ 'name': 'columnA', 'type': 'STRING', 'mode': 'NULLABLE' }] }, create_disposition='CREATE_NEVER', method='STREAMING_INSERTS')) self.assertEqual(3, mock_send.call_count) # Running tests with a variety of error reasons from # https://cloud.google.com/bigquery/docs/error-messages # This covers the scenario when # the google.cloud.bigquery.Client.insert_rows_json call returns an error list # rather than raising an exception. # Choosing some error reasons that are included in # bigquery_tools._NON_TRANSIENT_ERRORS and some that are not @parameterized.expand([ # reason in _NON_TRANSIENT_ERRORS for row 1, sent to failed_rows param( insert_response=[ [{ 'index': 0, 'errors': [{ 'reason': 'invalid' }] }], [{ 'index': 0, 'errors': [{ 'reason': 'invalid' }] }], ], failed_rows=['value1']), # reason in _NON_TRANSIENT_ERRORS for row 1 # reason not in _NON_TRANSIENT_ERRORS for row 2 on 1st run # row 1 sent to failed_rows param( insert_response=[ [{ 'index': 0, 'errors': [{ 'reason': 'invalid' }] }, { 'index': 1, 'errors': [{ 'reason': 'internalError' }] }], [{ 'index': 0, 'errors': [{ 'reason': 'invalid' }] }], ], failed_rows=['value1']), # reason not in _NON_TRANSIENT_ERRORS for row 1 on first attempt # transient error succeeds on second attempt, 0 rows sent to failed rows param( insert_response=[ [{ 'index': 0, 'errors': [{ 'reason': 'internalError' }] }], [], ], failed_rows=[]), ]) def test_insert_rows_json_errors_retry_always( self, insert_response, failed_rows, unused_sleep_mock=None): # In this test, a pipeline will always retry all errors # since RetryStrategy is not set and defaults to RETRY_ALWAYS with mock.patch('time.sleep'): call_counter = 0 def store_callback(table, **kwargs): nonlocal call_counter response = insert_response[call_counter] call_counter += 1 return response client = mock.Mock() client.insert_rows_json = mock.Mock(side_effect=store_callback) # Using the bundle based direct runner to avoid pickling problems # with mocks. with beam.Pipeline(runner='BundleBasedDirectRunner') as p: bq_write_out = ( p | beam.Create([{ 'columnA': 'value1', 'columnB': 'value2' }, { 'columnA': 'value3', 'columnB': 'value4' }, { 'columnA': 'value5', 'columnB': 'value6' }]) # Using _StreamToBigQuery in order to be able to pass max_retries # in order to limit run time of test with RETRY_ALWAYS | _StreamToBigQuery( table_reference='project:dataset.table', table_side_inputs=[], schema_side_inputs=[], schema='anyschema', batch_size=None, triggering_frequency=None, create_disposition='CREATE_NEVER', write_disposition=None, kms_key=None, retry_strategy=RetryStrategy.RETRY_ALWAYS, additional_bq_parameters=[], ignore_insert_ids=False, ignore_unknown_columns=False, with_auto_sharding=False, test_client=client, max_retries=len(insert_response) - 1, num_streaming_keys=500)) failed_values = ( bq_write_out[beam_bq.BigQueryWriteFn.FAILED_ROWS] | beam.Map(lambda x: x[1]['columnA'])) assert_that(failed_values, equal_to(failed_rows)) # Running tests with a variety of error reasons from # https://cloud.google.com/bigquery/docs/error-messages # This covers the scenario when # the google.cloud.bigquery.Client.insert_rows_json call returns an error list # rather than raising an exception. # Choosing some error reasons that are included in # bigquery_tools._NON_TRANSIENT_ERRORS and some that are not @parameterized.expand([ # reason in _NON_TRANSIENT_ERRORS for row 1, sent to failed_rows param( insert_response=[ [{ 'index': 0, 'errors': [{ 'reason': 'invalidQuery' }] }], ], streaming=False), # reason not in _NON_TRANSIENT_ERRORS for row 1, sent to failed_rows param( insert_response=[ [{ 'index': 0, 'errors': [{ 'reason': 'internalError' }] }], ], streaming=False), param( insert_response=[ [{ 'index': 0, 'errors': [{ 'reason': 'invalid' }] }], ], streaming=True), # reason not in _NON_TRANSIENT_ERRORS for row 1, sent to failed_rows param( insert_response=[ [{ 'index': 0, 'errors': [{ 'reason': 'internalError' }] }], ], streaming=True), ]) @mock.patch('time.sleep') @mock.patch('google.cloud.bigquery.Client.insert_rows_json') def test_insert_rows_json_errors_retry_never( self, mock_send, unused_mock_sleep, insert_response, streaming): # In this test, a pipeline will never retry errors since RetryStrategy is # set to RETRY_NEVER mock_send.side_effect = insert_response opt = StandardOptions() opt.streaming = streaming with beam.Pipeline(runner='BundleBasedDirectRunner', options=opt) as p: bq_write_out = ( p | beam.Create([{ 'columnA': 'value1' }, { 'columnA': 'value2' }]) | WriteToBigQuery( table='project:dataset.table', schema={ 'fields': [{ 'name': 'columnA', 'type': 'STRING', 'mode': 'NULLABLE' }] }, create_disposition='CREATE_NEVER', method='STREAMING_INSERTS', insert_retry_strategy=RetryStrategy.RETRY_NEVER)) failed_values = ( bq_write_out[beam_bq.BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS] | beam.Map(lambda x: x[1]['columnA'])) assert_that(failed_values, equal_to(['value1'])) self.assertEqual(1, mock_send.call_count) # Running tests with a variety of error reasons from # https://cloud.google.com/bigquery/docs/error-messages # This covers the scenario when # the google.cloud.bigquery.Client.insert_rows_json call returns an error list # rather than raising an exception. # Choosing some error reasons that are included in # bigquery_tools._NON_TRANSIENT_ERRORS and some that are not @parameterized.expand([ # reason in _NON_TRANSIENT_ERRORS for row 1, sent to failed_rows param( insert_response=[ [{ 'index': 0, 'errors': [{ 'reason': 'invalid' }] }], ], failed_rows=['value1'], streaming=False), # reason not in _NON_TRANSIENT_ERRORS for row 1 on 1st attempt # transient error succeeds on 2nd attempt, 0 rows sent to failed rows param( insert_response=[ [{ 'index': 0, 'errors': [{ 'reason': 'internalError' }] }], [], ], failed_rows=[], streaming=False), # reason in _NON_TRANSIENT_ERRORS for row 1 # reason not in _NON_TRANSIENT_ERRORS for row 2 on 1st and 2nd attempt # all rows with errors are retried when any row has a retriable error # row 1 sent to failed_rows after final attempt param( insert_response=[ [{ 'index': 0, 'errors': [{ 'reason': 'invalid' }] }, { 'index': 1, 'errors': [{ 'reason': 'internalError' }] }], [ { 'index': 0, 'errors': [{ 'reason': 'invalid' }] }, ], ], failed_rows=['value1'], streaming=False), # reason in _NON_TRANSIENT_ERRORS for row 1, sent to failed_rows param( insert_response=[ [{ 'index': 0, 'errors': [{ 'reason': 'invalid' }] }], ], failed_rows=['value1'], streaming=True), # reason not in _NON_TRANSIENT_ERRORS for row 1 on 1st attempt # transient error succeeds on 2nd attempt, 0 rows sent to failed rows param( insert_response=[ [{ 'index': 0, 'errors': [{ 'reason': 'internalError' }] }], [], ], failed_rows=[], streaming=True), # reason in _NON_TRANSIENT_ERRORS for row 1 # reason not in _NON_TRANSIENT_ERRORS for row 2 on 1st and 2nd attempt # all rows with errors are retried when any row has a retriable error # row 1 sent to failed_rows after final attempt param( insert_response=[ [{ 'index': 0, 'errors': [{ 'reason': 'invalid' }] }, { 'index': 1, 'errors': [{ 'reason': 'internalError' }] }], [ { 'index': 0, 'errors': [{ 'reason': 'invalid' }] }, ], ], failed_rows=['value1'], streaming=True), ]) @mock.patch('time.sleep') @mock.patch('google.cloud.bigquery.Client.insert_rows_json') def test_insert_rows_json_errors_retry_on_transient_error( self, mock_send, unused_mock_sleep, insert_response, failed_rows, streaming=False): # In this test, a pipeline will only retry errors with reasons that are not # in _NON_TRANSIENT_ERRORS since RetryStrategy is set to # RETRY_ON_TRANSIENT_ERROR call_counter = 0 def store_callback(table, **kwargs): nonlocal call_counter response = insert_response[call_counter] call_counter += 1 return response mock_send.side_effect = store_callback opt = StandardOptions() opt.streaming = streaming # Using the bundle based direct runner to avoid pickling problems # with mocks. with beam.Pipeline(runner='BundleBasedDirectRunner', options=opt) as p: bq_write_out = ( p | beam.Create([{ 'columnA': 'value1' }, { 'columnA': 'value2' }, { 'columnA': 'value3' }]) | WriteToBigQuery( table='project:dataset.table', schema={ 'fields': [{ 'name': 'columnA', 'type': 'STRING', 'mode': 'NULLABLE' }] }, create_disposition='CREATE_NEVER', method='STREAMING_INSERTS', insert_retry_strategy=RetryStrategy.RETRY_ON_TRANSIENT_ERROR)) failed_values = ( bq_write_out[beam_bq.BigQueryWriteFn.FAILED_ROWS] | beam.Map(lambda x: x[1]['columnA'])) assert_that(failed_values, equal_to(failed_rows)) def test_with_batched_input_exceeds_size_limit(self): from apache_beam.transforms import window from apache_beam.utils.windowed_value import WindowedValue client = mock.Mock() client.tables.Get.return_value = bigquery.Table( tableReference=bigquery.TableReference( projectId='project-id', datasetId='dataset_id', tableId='table_id')) client.insert_rows_json.return_value = [] fn = beam.io.gcp.bigquery.BigQueryWriteFn( batch_size=10, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND, kms_key=None, with_batched_input=True, max_insert_payload_size=500, # Small size to trigger the check test_client=client) fn.start_bundle() # Create rows where the third one exceeds the size limit rows = [ ({ 'month': 1 }, 'insertid1'), ({ 'month': 2 }, 'insertid2'), ({ 'columnA': 'large_string' * 100 }, 'insertid3') # This exceeds 500 bytes ] # Create a windowed value for testing test_window = window.GlobalWindow() test_timestamp = window.MIN_TIMESTAMP windowed_value = WindowedValue(None, test_timestamp, [test_window]) # Process the batched input result = fn.process(('project-id:dataset_id.table_id', rows), windowed_value) # Convert generator to list to check results failed_rows = list(result) # Should have 2 failed outputs (FAILED_ROWS_WITH_ERRORS and FAILED_ROWS) self.assertEqual(len(failed_rows), 2) # Check that the large row was sent to DLQ failed_with_errors = [ r for r in failed_rows if r.tag == fn.FAILED_ROWS_WITH_ERRORS ] failed_without_errors = [r for r in failed_rows if r.tag == fn.FAILED_ROWS] self.assertEqual(len(failed_with_errors), 1) self.assertEqual(len(failed_without_errors), 1) # Verify the error message destination, row, error = failed_with_errors[0].value.value self.assertEqual(destination, 'project-id:dataset_id.table_id') self.assertEqual(row, {'columnA': 'large_string' * 100}) self.assertIn('exceeds the maximum insert payload size', error) # Verify that only the valid rows were sent to BigQuery self.assertTrue(client.insert_rows_json.called) # Check that the rows were inserted (might be in multiple batches) total_rows = [] for call in client.insert_rows_json.call_args_list: total_rows.extend(call[1]['json_rows']) self.assertEqual(len(total_rows), 2) self.assertEqual(total_rows[0], {'month': 1}) self.assertEqual(total_rows[1], {'month': 2}) def test_with_batched_input_splits_large_batch(self): from apache_beam.transforms import window from apache_beam.utils.windowed_value import WindowedValue client = mock.Mock() client.tables.Get.return_value = bigquery.Table( tableReference=bigquery.TableReference( projectId='project-id', datasetId='dataset_id', tableId='table_id')) client.insert_rows_json.return_value = [] create_disposition = beam.io.BigQueryDisposition.CREATE_IF_NEEDED write_disposition = beam.io.BigQueryDisposition.WRITE_APPEND fn = beam.io.gcp.bigquery.BigQueryWriteFn( batch_size=10, create_disposition=create_disposition, write_disposition=write_disposition, kms_key=None, with_batched_input=True, max_insert_payload_size=800, # Small size to force batch splitting test_client=client) fn.start_bundle() # Create rows that together exceed the batch size limit. # Each row with 'data' * 10 is about 200 bytes # So 2 rows should fit, 3rd should cause flush. rows = [ ({ 'data': 'x' * 10 }, 'insertid1'), ({ 'data': 'y' * 10 }, 'insertid2'), ({ 'data': 'z' * 10 }, 'insertid3'), # This should go in second batch ] # Create a windowed value for testing test_window = window.GlobalWindow() test_timestamp = window.MIN_TIMESTAMP windowed_value = WindowedValue(None, test_timestamp, [test_window]) # Process the batched input result = fn.process(('project-id:dataset_id.table_id', rows), windowed_value) # Convert generator to list (should be empty as no failures) failed_rows = list(result) self.assertEqual(len(failed_rows), 0) # With 800 byte limit and ~341 byte rows # we should be able to fit 2 rows per batch. # So we expect 2 calls: [row1, row2] and [row3] self.assertEqual(client.insert_rows_json.call_count, 2) # Check first batch (2 rows) first_call = client.insert_rows_json.call_args_list[0][1] self.assertEqual(len(first_call['json_rows']), 2) self.assertEqual(first_call['json_rows'][0], {'data': 'x' * 10}) self.assertEqual(first_call['json_rows'][1], {'data': 'y' * 10}) # Check second batch (1 row) second_call = client.insert_rows_json.call_args_list[1][1] self.assertEqual(len(second_call['json_rows']), 1) self.assertEqual(second_call['json_rows'][0], {'data': 'z' * 10}) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class BigQueryStreamingInsertTransformTests(unittest.TestCase): def test_dofn_client_process_performs_batching(self): client = mock.Mock() client.tables.Get.return_value = bigquery.Table( tableReference=bigquery.TableReference( projectId='project-id', datasetId='dataset_id', tableId='table_id')) client.insert_rows_json.return_value = [] create_disposition = beam.io.BigQueryDisposition.CREATE_NEVER write_disposition = beam.io.BigQueryDisposition.WRITE_APPEND fn = beam.io.gcp.bigquery.BigQueryWriteFn( batch_size=2, create_disposition=create_disposition, write_disposition=write_disposition, kms_key=None, test_client=client) fn.process(('project-id:dataset_id.table_id', {'month': 1})) # InsertRows not called as batch size is not hit yet self.assertFalse(client.insert_rows_json.called) def test_dofn_client_process_flush_called(self): client = mock.Mock() client.tables.Get.return_value = bigquery.Table( tableReference=bigquery.TableReference( projectId='project-id', datasetId='dataset_id', tableId='table_id')) client.insert_rows_json.return_value = [] create_disposition = beam.io.BigQueryDisposition.CREATE_NEVER write_disposition = beam.io.BigQueryDisposition.WRITE_APPEND fn = beam.io.gcp.bigquery.BigQueryWriteFn( batch_size=2, create_disposition=create_disposition, write_disposition=write_disposition, kms_key=None, test_client=client) fn.start_bundle() fn.process(('project-id:dataset_id.table_id', ({'month': 1}, 'insertid1'))) fn.process(('project-id:dataset_id.table_id', ({'month': 2}, 'insertid2'))) fn.finish_bundle() # InsertRows called as batch size is hit self.assertTrue(client.insert_rows_json.called) def test_dofn_client_finish_bundle_flush_called(self): client = mock.Mock() client.tables.Get.return_value = bigquery.Table( tableReference=bigquery.TableReference( projectId='project-id', datasetId='dataset_id', tableId='table_id')) client.insert_rows_json.return_value = [] create_disposition = beam.io.BigQueryDisposition.CREATE_IF_NEEDED write_disposition = beam.io.BigQueryDisposition.WRITE_APPEND fn = beam.io.gcp.bigquery.BigQueryWriteFn( batch_size=2, create_disposition=create_disposition, write_disposition=write_disposition, kms_key=None, test_client=client) fn.start_bundle() # Destination is a tuple of (destination, schema) to ensure the table is # created. fn.process(('project-id:dataset_id.table_id', ({'month': 1}, 'insertid3'))) self.assertTrue(client.tables.Get.called) # InsertRows not called as batch size is not hit self.assertFalse(client.insert_rows_json.called) fn.finish_bundle() # InsertRows called in finish bundle self.assertTrue(client.insert_rows_json.called) def test_dofn_client_no_records(self): client = mock.Mock() client.tables.Get.return_value = bigquery.Table( tableReference=bigquery.TableReference( projectId='project-id', datasetId='dataset_id', tableId='table_id')) client.tabledata.InsertAll.return_value = \ bigquery.TableDataInsertAllResponse(insertErrors=[]) create_disposition = beam.io.BigQueryDisposition.CREATE_NEVER write_disposition = beam.io.BigQueryDisposition.WRITE_APPEND fn = beam.io.gcp.bigquery.BigQueryWriteFn( batch_size=2, create_disposition=create_disposition, write_disposition=write_disposition, kms_key=None, test_client=client) fn.start_bundle() # InsertRows not called as batch size is not hit self.assertFalse(client.tabledata.InsertAll.called) fn.finish_bundle() # InsertRows not called in finish bundle as no records self.assertFalse(client.tabledata.InsertAll.called) def test_with_batched_input(self): client = mock.Mock() client.tables.Get.return_value = bigquery.Table( tableReference=bigquery.TableReference( projectId='project-id', datasetId='dataset_id', tableId='table_id')) client.insert_rows_json.return_value = [] create_disposition = beam.io.BigQueryDisposition.CREATE_IF_NEEDED write_disposition = beam.io.BigQueryDisposition.WRITE_APPEND fn = beam.io.gcp.bigquery.BigQueryWriteFn( batch_size=10, create_disposition=create_disposition, write_disposition=write_disposition, kms_key=None, with_batched_input=True, test_client=client) fn.start_bundle() # Destination is a tuple of (destination, schema) to ensure the table is # created. fn.process(( 'project-id:dataset_id.table_id', [({ 'month': 1 }, 'insertid3'), ({ 'month': 2 }, 'insertid2'), ({ 'month': 3 }, 'insertid1')])) # InsertRows called since the input is already batched. self.assertTrue(client.insert_rows_json.called) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class PipelineBasedStreamingInsertTest(_TestCaseWithTempDirCleanUp): @mock.patch('time.sleep') def test_failure_has_same_insert_ids(self, unused_mock_sleep): tempdir = '%s%s' % (self._new_tempdir(), os.sep) file_name_1 = os.path.join(tempdir, 'file1') file_name_2 = os.path.join(tempdir, 'file2') def store_callback(table, **kwargs): insert_ids = [r for r in kwargs['row_ids']] colA_values = [r['columnA'] for r in kwargs['json_rows']] json_output = {'insertIds': insert_ids, 'colA_values': colA_values} # The first time we try to insert, we save those insertions in # file insert_calls1. if not os.path.exists(file_name_1): with open(file_name_1, 'w') as f: json.dump(json_output, f) raise RuntimeError() else: with open(file_name_2, 'w') as f: json.dump(json_output, f) return [] client = mock.Mock() client.insert_rows_json = mock.Mock(side_effect=store_callback) # Using the bundle based direct runner to avoid pickling problems # with mocks. with beam.Pipeline(runner='BundleBasedDirectRunner') as p: _ = ( p | beam.Create([{ 'columnA': 'value1', 'columnB': 'value2' }, { 'columnA': 'value3', 'columnB': 'value4' }, { 'columnA': 'value5', 'columnB': 'value6' }]) | _StreamToBigQuery( table_reference='project:dataset.table', table_side_inputs=[], schema_side_inputs=[], schema='anyschema', batch_size=None, triggering_frequency=None, create_disposition='CREATE_NEVER', write_disposition=None, kms_key=None, retry_strategy=None, additional_bq_parameters=[], ignore_insert_ids=False, ignore_unknown_columns=False, with_auto_sharding=False, test_client=client, num_streaming_keys=500)) with open(file_name_1) as f1, open(file_name_2) as f2: self.assertEqual(json.load(f1), json.load(f2)) @parameterized.expand([ param(retry_strategy=RetryStrategy.RETRY_ALWAYS), param(retry_strategy=RetryStrategy.RETRY_NEVER), param(retry_strategy=RetryStrategy.RETRY_ON_TRANSIENT_ERROR), ]) def test_failure_in_some_rows_does_not_duplicate(self, retry_strategy=None): with mock.patch('time.sleep'): # In this test we simulate a failure to write out two out of three rows. # Row 0 and row 2 fail to be written on the first attempt, and then # succeed on the next attempt (if there is one). tempdir = '%s%s' % (self._new_tempdir(), os.sep) file_name_1 = os.path.join(tempdir, 'file1_partial') file_name_2 = os.path.join(tempdir, 'file2_partial') def store_callback(table, **kwargs): insert_ids = [r for r in kwargs['row_ids']] colA_values = [r['columnA'] for r in kwargs['json_rows']] # The first time this function is called, all rows are included # so we need to filter out 'failed' rows. json_output_1 = { 'insertIds': [insert_ids[1]], 'colA_values': [colA_values[1]] } # The second time this function is called, only rows 0 and 2 are incl # so we don't need to filter any of them. We just write them all out. json_output_2 = {'insertIds': insert_ids, 'colA_values': colA_values} # The first time we try to insert, we save those insertions in # file insert_calls1. if not os.path.exists(file_name_1): with open(file_name_1, 'w') as f: json.dump(json_output_1, f) return [ { 'index': 0, 'errors': [{ 'reason': 'i dont like this row' }, { 'reason': 'its bad' }] }, { 'index': 2, 'errors': [{ 'reason': 'i het this row' }, { 'reason': 'its no gud' }] }, ] else: with open(file_name_2, 'w') as f: json.dump(json_output_2, f) return [] client = mock.Mock() client.insert_rows_json = mock.Mock(side_effect=store_callback) # The expected rows to be inserted according to the insert strategy if retry_strategy == RetryStrategy.RETRY_NEVER: result = ['value3'] else: # RETRY_ALWAYS and RETRY_ON_TRANSIENT_ERRORS should insert all rows result = ['value1', 'value3', 'value5'] # Using the bundle based direct runner to avoid pickling problems # with mocks. with beam.Pipeline(runner='BundleBasedDirectRunner') as p: bq_write_out = ( p | beam.Create([{ 'columnA': 'value1', 'columnB': 'value2' }, { 'columnA': 'value3', 'columnB': 'value4' }, { 'columnA': 'value5', 'columnB': 'value6' }]) | _StreamToBigQuery( table_reference='project:dataset.table', table_side_inputs=[], schema_side_inputs=[], schema='anyschema', batch_size=None, triggering_frequency=None, create_disposition='CREATE_NEVER', write_disposition=None, kms_key=None, retry_strategy=retry_strategy, additional_bq_parameters=[], ignore_insert_ids=False, ignore_unknown_columns=False, with_auto_sharding=False, test_client=client, num_streaming_keys=500)) failed_values = ( bq_write_out[beam_bq.BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS] | beam.Map(lambda x: x[1]['columnA'])) assert_that( failed_values, equal_to(list({'value1', 'value3', 'value5'}.difference(result)))) data1 = _load_or_default(file_name_1) data2 = _load_or_default(file_name_2) self.assertListEqual( sorted(data1.get('colA_values', []) + data2.get('colA_values', [])), result) self.assertEqual(len(data1['colA_values']), 1) @parameterized.expand([ param(with_auto_sharding=False), param(with_auto_sharding=True), ]) def test_batch_size_with_auto_sharding(self, with_auto_sharding): tempdir = '%s%s' % (self._new_tempdir(), os.sep) file_name_1 = os.path.join(tempdir, 'file1') file_name_2 = os.path.join(tempdir, 'file2') def store_callback(table, **kwargs): insert_ids = [r for r in kwargs['row_ids']] colA_values = [r['columnA'] for r in kwargs['json_rows']] json_output = {'insertIds': insert_ids, 'colA_values': colA_values} # Expect two batches of rows will be inserted. Store them separately. if not os.path.exists(file_name_1): with open(file_name_1, 'w') as f: json.dump(json_output, f) else: with open(file_name_2, 'w') as f: json.dump(json_output, f) return [] client = mock.Mock() client.insert_rows_json = mock.Mock(side_effect=store_callback) # Using the bundle based direct runner to avoid pickling problems # with mocks. with beam.Pipeline(runner='BundleBasedDirectRunner') as p: _ = ( p | beam.Create([{ 'columnA': 'value1', 'columnB': 'value2' }, { 'columnA': 'value3', 'columnB': 'value4' }, { 'columnA': 'value5', 'columnB': 'value6' }]) | _StreamToBigQuery( table_reference='project:dataset.table', table_side_inputs=[], schema_side_inputs=[], schema='anyschema', # Set a batch size such that the input elements will be inserted # in 2 batches. batch_size=2, triggering_frequency=None, create_disposition='CREATE_NEVER', write_disposition=None, kms_key=None, retry_strategy=None, additional_bq_parameters=[], ignore_insert_ids=False, ignore_unknown_columns=False, with_auto_sharding=with_auto_sharding, test_client=client, num_streaming_keys=500)) with open(file_name_1) as f1, open(file_name_2) as f2: out1 = json.load(f1)['colA_values'] out2 = json.load(f2)['colA_values'] out_all = out1 + out2 out_all.sort() self.assertEqual(out_all, ['value1', 'value3', 'value5']) self.assertEqual(len(out1), 2) self.assertEqual(len(out2), 1) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class BigQueryStreamingInsertTransformIntegrationTests(unittest.TestCase): BIG_QUERY_DATASET_ID = 'python_bq_streaming_inserts_' def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.runner_name = type(self.test_pipeline.runner).__name__ self.project = self.test_pipeline.get_option('project') self.dataset_id = '%s%d%s' % ( self.BIG_QUERY_DATASET_ID, int(time.time()), secrets.token_hex(3)) self.bigquery_client = bigquery_tools.BigQueryWrapper() self.bigquery_client.get_or_create_dataset(self.project, self.dataset_id) self.output_table = "%s.output_table" % (self.dataset_id) _LOGGER.info( "Created dataset %s in project %s", self.dataset_id, self.project) @pytest.mark.it_postcommit def test_value_provider_transform(self): output_table_1 = '%s%s' % (self.output_table, 1) output_table_2 = '%s%s' % (self.output_table, 2) schema = { 'fields': [{ 'name': 'name', 'type': 'STRING', 'mode': 'NULLABLE' }, { 'name': 'language', 'type': 'STRING', 'mode': 'NULLABLE' }] } additional_bq_parameters = { 'timePartitioning': { 'type': 'DAY' }, 'clustering': { 'fields': ['language'] } } table_ref = bigquery_tools.parse_table_reference(output_table_1) table_ref2 = bigquery_tools.parse_table_reference(output_table_2) pipeline_verifiers = [ BigQueryTableMatcher( project=self.project, dataset=table_ref.datasetId, table=table_ref.tableId, expected_properties=additional_bq_parameters), BigQueryTableMatcher( project=self.project, dataset=table_ref2.datasetId, table=table_ref2.tableId, expected_properties=additional_bq_parameters), BigqueryFullResultMatcher( project=self.project, query="SELECT name, language FROM %s" % output_table_1, data=[(d['name'], d['language']) for d in _ELEMENTS if 'language' in d]), BigqueryFullResultMatcher( project=self.project, query="SELECT name, language FROM %s" % output_table_2, data=[(d['name'], d['language']) for d in _ELEMENTS if 'language' in d]) ] args = self.test_pipeline.get_full_options_as_args( on_success_matcher=hc.all_of(*pipeline_verifiers)) with beam.Pipeline(argv=args) as p: input = p | beam.Create([row for row in _ELEMENTS if 'language' in row]) _ = ( input | "WriteWithMultipleDests" >> beam.io.gcp.bigquery.WriteToBigQuery( table=value_provider.StaticValueProvider( str, '%s:%s' % (self.project, output_table_1)), schema=value_provider.StaticValueProvider(dict, schema), additional_bq_parameters=additional_bq_parameters, method='STREAMING_INSERTS')) _ = ( input | "WriteWithMultipleDests2" >> beam.io.gcp.bigquery.WriteToBigQuery( table=value_provider.StaticValueProvider( str, '%s:%s' % (self.project, output_table_2)), schema=beam.io.gcp.bigquery.SCHEMA_AUTODETECT, additional_bq_parameters=lambda _: additional_bq_parameters, method='FILE_LOADS')) @pytest.mark.it_postcommit def test_multiple_destinations_transform(self): streaming = self.test_pipeline.options.view_as(StandardOptions).streaming if streaming and isinstance(self.test_pipeline.runner, TestDataflowRunner): self.skipTest("TestStream is not supported on TestDataflowRunner") output_table_1 = '%s%s' % (self.output_table, 1) output_table_2 = '%s%s' % (self.output_table, 2) full_output_table_1 = '%s:%s' % (self.project, output_table_1) full_output_table_2 = '%s:%s' % (self.project, output_table_2) schema1 = { 'fields': [{ 'name': 'name', 'type': 'STRING', 'mode': 'NULLABLE' }, { 'name': 'language', 'type': 'STRING', 'mode': 'NULLABLE' }] } schema2 = { 'fields': [{ 'name': 'name', 'type': 'STRING', 'mode': 'NULLABLE' }, { 'name': 'foundation', 'type': 'STRING', 'mode': 'NULLABLE' }] } bad_record = {'language': 1, 'manguage': 2} if streaming: pipeline_verifiers = [ PipelineStateMatcher(PipelineState.RUNNING), BigqueryFullResultStreamingMatcher( project=self.project, query="SELECT name, language FROM %s" % output_table_1, data=[(d['name'], d['language']) for d in _ELEMENTS if 'language' in d]), BigqueryFullResultStreamingMatcher( project=self.project, query="SELECT name, foundation FROM %s" % output_table_2, data=[(d['name'], d['foundation']) for d in _ELEMENTS if 'foundation' in d]) ] else: pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query="SELECT name, language FROM %s" % output_table_1, data=[(d['name'], d['language']) for d in _ELEMENTS if 'language' in d]), BigqueryFullResultMatcher( project=self.project, query="SELECT name, foundation FROM %s" % output_table_2, data=[(d['name'], d['foundation']) for d in _ELEMENTS if 'foundation' in d]) ] args = self.test_pipeline.get_full_options_as_args( on_success_matcher=hc.all_of(*pipeline_verifiers)) with beam.Pipeline(argv=args) as p: if streaming: _SIZE = len(_ELEMENTS) test_stream = ( TestStream().advance_watermark_to(0).add_elements( _ELEMENTS[:_SIZE // 2]).advance_watermark_to(100).add_elements( _ELEMENTS[_SIZE // 2:]).advance_watermark_to_infinity()) input = p | test_stream else: input = p | beam.Create(_ELEMENTS) schema_table_pcv = beam.pvalue.AsDict( p | "MakeSchemas" >> beam.Create([(full_output_table_1, schema1), (full_output_table_2, schema2)])) table_record_pcv = beam.pvalue.AsDict( p | "MakeTables" >> beam.Create([('table1', full_output_table_1), ('table2', full_output_table_2)])) input2 = p | "Broken record" >> beam.Create([bad_record]) input = (input, input2) | beam.Flatten() r = ( input | "WriteWithMultipleDests" >> beam.io.gcp.bigquery.WriteToBigQuery( table=lambda x, tables: (tables['table1'] if 'language' in x else tables['table2']), table_side_inputs=(table_record_pcv, ), schema=lambda dest, table_map: table_map.get(dest, None), schema_side_inputs=(schema_table_pcv, ), insert_retry_strategy=RetryStrategy.RETRY_ON_TRANSIENT_ERROR, method='STREAMING_INSERTS')) assert_that( r[beam.io.gcp.bigquery.BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS] | beam.Map(lambda elm: (elm[0], elm[1])), equal_to([(full_output_table_1, bad_record)])) assert_that( r[beam.io.gcp.bigquery.BigQueryWriteFn.FAILED_ROWS], equal_to([(full_output_table_1, bad_record)]), label='FailedRowsMatch') def tearDown(self): request = bigquery.BigqueryDatasetsDeleteRequest( projectId=self.project, datasetId=self.dataset_id, deleteContents=True) try: _LOGGER.info( "Deleting dataset %s in project %s", self.dataset_id, self.project) self.bigquery_client.client.datasets.Delete(request) except HttpError: _LOGGER.debug( 'Failed to clean up dataset %s in project %s', self.dataset_id, self.project) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class PubSubBigQueryIT(unittest.TestCase): INPUT_TOPIC = 'psit_topic_output' INPUT_SUB = 'psit_subscription_input' BIG_QUERY_DATASET_ID = 'python_pubsub_bq_' SCHEMA = { 'fields': [{ 'name': 'number', 'type': 'INTEGER', 'mode': 'NULLABLE' }] } _SIZE = 4 WAIT_UNTIL_FINISH_DURATION = 15 * 60 * 1000 def setUp(self): # Set up PubSub self.test_pipeline = TestPipeline(is_integration_test=True) self.runner_name = type(self.test_pipeline.runner).__name__ self.project = self.test_pipeline.get_option('project') self.uuid = str(uuid.uuid4()) from google.cloud import pubsub self.pub_client = pubsub.PublisherClient() self.input_topic = self.pub_client.create_topic( name=self.pub_client.topic_path( self.project, self.INPUT_TOPIC + self.uuid)) self.sub_client = pubsub.SubscriberClient() self.input_sub = self.sub_client.create_subscription( name=self.sub_client.subscription_path( self.project, self.INPUT_SUB + self.uuid), topic=self.input_topic.name) # Set up BQ self.dataset_ref = utils.create_bq_dataset( self.project, self.BIG_QUERY_DATASET_ID) self.output_table = "%s.output_table" % (self.dataset_ref.dataset_id) def tearDown(self): # Tear down PubSub test_utils.cleanup_topics(self.pub_client, [self.input_topic]) test_utils.cleanup_subscriptions(self.sub_client, [self.input_sub]) # Tear down BigQuery utils.delete_bq_dataset(self.project, self.dataset_ref) def _run_pubsub_bq_pipeline(self, method, triggering_frequency=None): l = [i for i in range(self._SIZE)] matchers = [ PipelineStateMatcher(PipelineState.RUNNING), BigqueryFullResultStreamingMatcher( project=self.project, query="SELECT number FROM %s" % self.output_table, data=[(i, ) for i in l]) ] args = self.test_pipeline.get_full_options_as_args( on_success_matcher=hc.all_of(*matchers), wait_until_finish_duration=self.WAIT_UNTIL_FINISH_DURATION, streaming=True, allow_unsafe_triggers=True) def add_schema_info(element): yield {'number': element} messages = [str(i).encode('utf-8') for i in l] for message in messages: self.pub_client.publish(self.input_topic.name, message) with beam.Pipeline(argv=args) as p: mesages = ( p | ReadFromPubSub(subscription=self.input_sub.name) | beam.ParDo(add_schema_info)) _ = mesages | WriteToBigQuery( self.output_table, schema=self.SCHEMA, method=method, triggering_frequency=triggering_frequency) @pytest.mark.it_postcommit def test_streaming_inserts(self): self._run_pubsub_bq_pipeline(WriteToBigQuery.Method.STREAMING_INSERTS) @pytest.mark.it_postcommit def test_file_loads(self): self._run_pubsub_bq_pipeline( WriteToBigQuery.Method.FILE_LOADS, triggering_frequency=20) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class BigQueryFileLoadsIntegrationTests(unittest.TestCase): BIG_QUERY_DATASET_ID = 'python_bq_file_loads_' def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.runner_name = type(self.test_pipeline.runner).__name__ self.project = self.test_pipeline.get_option('project') self.dataset_id = '%s%d%s' % ( self.BIG_QUERY_DATASET_ID, int(time.time()), secrets.token_hex(3)) self.bigquery_client = bigquery_tools.BigQueryWrapper() self.bigquery_client.get_or_create_dataset(self.project, self.dataset_id) self.output_table = '%s.output_table' % (self.dataset_id) self.table_ref = bigquery_tools.parse_table_reference(self.output_table) _LOGGER.info( 'Created dataset %s in project %s', self.dataset_id, self.project) @pytest.mark.it_postcommit def test_avro_file_load(self): # Construct elements such that they can be written via Avro but not via # JSON. See BEAM-8841. from apache_beam.io.gcp import bigquery_file_loads old_max_files = bigquery_file_loads._MAXIMUM_SOURCE_URIS old_max_file_size = bigquery_file_loads._DEFAULT_MAX_FILE_SIZE bigquery_file_loads._MAXIMUM_SOURCE_URIS = 1 bigquery_file_loads._DEFAULT_MAX_FILE_SIZE = 100 elements = [ { 'name': 'Negative infinity', 'value': -float('inf'), 'timestamp': datetime.datetime(1970, 1, 1, tzinfo=pytz.utc), }, { 'name': 'Not a number', 'value': float('nan'), 'timestamp': datetime.datetime(2930, 12, 9, tzinfo=pytz.utc), }, ] schema = beam.io.gcp.bigquery.WriteToBigQuery.get_dict_table_schema( bigquery.TableSchema( fields=[ bigquery.TableFieldSchema( name='name', type='STRING', mode='REQUIRED'), bigquery.TableFieldSchema( name='value', type='FLOAT', mode='REQUIRED'), bigquery.TableFieldSchema( name='timestamp', type='TIMESTAMP', mode='REQUIRED'), ])) pipeline_verifiers = [ # Some gymnastics here to avoid comparing NaN since NaN is not equal to # anything, including itself. BigqueryFullResultMatcher( project=self.project, query="SELECT name, value, timestamp FROM {} WHERE value<0".format( self.output_table), data=[(d['name'], d['value'], d['timestamp']) for d in elements[:1]], ), BigqueryFullResultMatcher( project=self.project, query="SELECT name, timestamp FROM {}".format(self.output_table), data=[(d['name'], d['timestamp']) for d in elements], ), ] args = self.test_pipeline.get_full_options_as_args( on_success_matcher=hc.all_of(*pipeline_verifiers), ) with beam.Pipeline(argv=args) as p: input = p | 'CreateInput' >> beam.Create(elements) schema_pc = p | 'CreateSchema' >> beam.Create([schema]) _ = ( input | 'WriteToBigQuery' >> beam.io.gcp.bigquery.WriteToBigQuery( table='%s:%s' % (self.project, self.output_table), schema=lambda _, schema: schema, schema_side_inputs=(beam.pvalue.AsSingleton(schema_pc), ), method='FILE_LOADS', temp_file_format=bigquery_tools.FileFormat.AVRO, )) bigquery_file_loads._MAXIMUM_SOURCE_URIS = old_max_files bigquery_file_loads._DEFAULT_MAX_FILE_SIZE = old_max_file_size def tearDown(self): request = bigquery.BigqueryDatasetsDeleteRequest( projectId=self.project, datasetId=self.dataset_id, deleteContents=True) try: _LOGGER.info( "Deleting dataset %s in project %s", self.dataset_id, self.project) self.bigquery_client.client.datasets.Delete(request) except HttpError: _LOGGER.debug( 'Failed to clean up dataset %s in project %s', self.dataset_id, self.project) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_tools.py0000644000175100017510000020267015206113032022465 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tools used by BigQuery sources and sinks. Classes, constants and functions in this file are experimental and have no backwards compatibility guarantees. These tools include wrappers and clients to interact with BigQuery APIs. NOTHING IN THIS FILE HAS BACKWARDS COMPATIBILITY GUARANTEES. """ # pytype: skip-file # pylint: disable=wrong-import-order, wrong-import-position import datetime import decimal import io import json import logging import re import sys import time import uuid from json.decoder import JSONDecodeError from typing import Optional from typing import Sequence from typing import TypeVar from typing import Union import fastavro import numpy as np import apache_beam from apache_beam import coders from apache_beam.internal.gcp import auth from apache_beam.internal.gcp.json_value import from_json_value from apache_beam.internal.http_client import get_new_http from apache_beam.internal.metrics.metric import MetricLogger from apache_beam.internal.metrics.metric import ServiceCallMetric from apache_beam.io.gcp import bigquery_avro_tools from apache_beam.io.gcp import resource_identifiers from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.metrics import monitoring_infos from apache_beam.metrics.metric import Metrics from apache_beam.options import value_provider from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.transforms import DoFn from apache_beam.typehints.row_type import RowTypeConstraint from apache_beam.typehints.typehints import Any from apache_beam.utils import retry from apache_beam.utils.histogram import LinearBucket # Protect against environments where bigquery library is not available. try: import regex from apitools.base.py.exceptions import HttpError from apitools.base.py.exceptions import HttpForbiddenError from apitools.base.py.transfer import Upload from google.api_core.client_info import ClientInfo from google.api_core.exceptions import ClientError from google.api_core.exceptions import GoogleAPICallError from google.cloud import bigquery as gcp_bigquery except Exception: gcp_bigquery = None pass try: from orjson import dumps as fast_json_dumps from orjson import loads as fast_json_loads except ImportError: fast_json_dumps = json.dumps fast_json_loads = json.loads # pylint: enable=wrong-import-order, wrong-import-position # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from apache_beam.io.gcp.internal.clients.bigquery import TableReference except ImportError: TableReference = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports _LOGGER = logging.getLogger(__name__) JSON_COMPLIANCE_ERROR = 'NAN, INF and -INF values are not JSON compliant.' MAX_RETRIES = 3 UNKNOWN_MIME_TYPE = 'application/octet-stream' # Timeout for a BQ streaming insert RPC. Set to a maximum of 2 minutes. BQ_STREAMING_INSERT_TIMEOUT_SEC = 120 _PROJECT_PATTERN = r'([a-z0-9.-]+:)?[a-z][a-z0-9-]*[a-z0-9]' _DATASET_PATTERN = r'\w{1,1024}' _TABLE_PATTERN = r'[\p{L}\p{M}\p{N}\p{Pc}\p{Pd}\p{Zs}$]{1,1024}' # TODO(https://github.com/apache/beam/issues/25946): Add support for # more Beam portable schema types as Python types BIGQUERY_TYPE_TO_PYTHON_TYPE = { "STRING": str, "BOOL": bool, "BOOLEAN": bool, "BYTES": bytes, "INT64": np.int64, "INTEGER": np.int64, "FLOAT64": np.float64, "FLOAT": np.float64, "NUMERIC": decimal.Decimal, "TIMESTAMP": apache_beam.utils.timestamp.Timestamp, "GEOGRAPHY": str, } class FileFormat(object): CSV = 'CSV' JSON = 'NEWLINE_DELIMITED_JSON' AVRO = 'AVRO' class ExportCompression(object): GZIP = 'GZIP' DEFLATE = 'DEFLATE' SNAPPY = 'SNAPPY' NONE = 'NONE' def default_encoder(obj): if isinstance(obj, decimal.Decimal): return str(obj) elif isinstance(obj, bytes): # on python 3 base64-encoded bytes are decoded to strings # before being sent to BigQuery return obj.decode('utf-8') elif isinstance(obj, (datetime.date, datetime.time)): return str(obj) elif isinstance(obj, datetime.datetime): return obj.isoformat() _LOGGER.error("Unable to serialize %r to JSON", obj) raise TypeError( "Object of type '%s' is not JSON serializable" % type(obj).__name__) def get_hashable_destination(destination): """Parses a table reference into a (project, dataset, table) tuple. Args: destination: Either a TableReference object from the bigquery API. The object has the following attributes: projectId, datasetId, and tableId. Or a string representing the destination containing 'PROJECT:DATASET.TABLE'. Returns: A string representing the destination containing 'PROJECT:DATASET.TABLE'. """ if isinstance(destination, TableReference): return '%s:%s.%s' % ( destination.projectId, destination.datasetId, destination.tableId) else: return destination V = TypeVar('V') def to_hashable_table_ref( table_ref_elem_kv: tuple[Union[str, TableReference], V]) -> tuple[str, V]: """Turns the key of the input tuple to its string representation. The key should be either a string or a TableReference. Args: table_ref_elem_kv: A tuple of table reference and element. Returns: A tuple of string representation of input table and input element. """ table_ref = table_ref_elem_kv[0] hashable_table_ref = get_hashable_destination(table_ref) return (hashable_table_ref, table_ref_elem_kv[1]) def parse_table_schema_from_json(schema_string): """Parse the Table Schema provided as string. Args: schema_string: String serialized table schema, should be a valid JSON. Returns: A TableSchema of the BigQuery export from either the Query or the Table. """ try: json_schema = json.loads(schema_string) except JSONDecodeError as e: raise ValueError( 'Unable to parse JSON schema: %s - %r' % (schema_string, e)) def _parse_schema_field(field): """Parse a single schema field from dictionary. Args: field: Dictionary object containing serialized schema. Returns: A TableFieldSchema for a single column in BigQuery. """ schema = bigquery.TableFieldSchema() schema.name = field['name'] schema.type = field['type'] if 'mode' in field: schema.mode = field['mode'] else: schema.mode = 'NULLABLE' if 'description' in field: schema.description = field['description'] if 'fields' in field: schema.fields = [_parse_schema_field(x) for x in field['fields']] return schema fields = [_parse_schema_field(f) for f in json_schema['fields']] return bigquery.TableSchema(fields=fields) def parse_table_reference(table, dataset=None, project=None): """Parses a table reference into a (project, dataset, table) tuple. Args: table: The ID of the table. The ID must contain only letters (a-z, A-Z), numbers (0-9), connectors (-_). If dataset argument is None then the table argument must contain the entire table reference: 'DATASET.TABLE' or 'PROJECT:DATASET.TABLE'. This argument can be a TableReference instance in which case dataset and project are ignored and the reference is returned as a result. Additionally, for date partitioned tables, appending '$YYYYmmdd' to the table name is supported, e.g. 'DATASET.TABLE$YYYYmmdd'. dataset: The ID of the dataset containing this table or null if the table reference is specified entirely by the table argument. project: The ID of the project containing this table or null if the table reference is specified entirely by the table (and possibly dataset) argument. Returns: A TableReference object from the bigquery API. The object has the following attributes: projectId, datasetId, and tableId. If the input is a TableReference object, a new object will be returned. Raises: ValueError: if the table reference as a string does not match the expected format. """ if isinstance(table, TableReference): return TableReference( projectId=table.projectId, datasetId=table.datasetId, tableId=table.tableId) elif callable(table): return table elif isinstance(table, value_provider.ValueProvider): return table table_reference = TableReference() # If dataset argument is not specified, the expectation is that the # table argument will contain a full table reference instead of just a # table name. if dataset is None: pattern = ( f'((?P{_PROJECT_PATTERN})[:\\.])?' f'(?P{_DATASET_PATTERN})\\.(?P{_TABLE_PATTERN})') match = regex.fullmatch(pattern, table) if not match: raise ValueError( 'Expected a table reference (PROJECT:DATASET.TABLE or ' 'DATASET.TABLE) instead of %s.' % table) table_reference.projectId = match.group('project') table_reference.datasetId = match.group('dataset') table_reference.tableId = match.group('table') else: table_reference.projectId = project table_reference.datasetId = dataset table_reference.tableId = table return table_reference # ----------------------------------------------------------------------------- # BigQueryWrapper. def _build_job_labels(input_labels): """Builds job label protobuf structure.""" input_labels = input_labels or {} result = bigquery.JobConfiguration.LabelsValue() for k, v in input_labels.items(): result.additionalProperties.append( bigquery.JobConfiguration.LabelsValue.AdditionalProperty( key=k, value=v, )) return result def _build_dataset_labels(input_labels): """Builds dataset label protobuf structure.""" input_labels = input_labels or {} result = bigquery.Dataset.LabelsValue() for k, v in input_labels.items(): result.additionalProperties.append( bigquery.Dataset.LabelsValue.AdditionalProperty( key=k, value=v, )) return result def _build_filter_from_labels(labels): filter_str = '' for key, value in labels.items(): filter_str += 'labels.' + key + ':' + value + ' ' return filter_str def _build_dataset_encryption_config(kms_key): return bigquery.EncryptionConfiguration(kmsKeyName=kms_key) class BigQueryWrapper(object): """BigQuery client wrapper with utilities for querying. The wrapper is used to organize all the BigQuery integration points and offer a common place where retry logic for failures can be controlled. In addition, it offers various functions used both in sources and sinks (e.g., find and create tables, query a table, etc.). Note that client parameter in constructor is only for testing purposes and should not be used in production code. """ # If updating following names, also update the corresponding pydocs in # bigquery.py. TEMP_TABLE = 'beam_temp_table_' TEMP_DATASET = 'beam_temp_dataset_' HISTOGRAM_METRIC_LOGGER = MetricLogger() def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None): self.client = client or BigQueryWrapper._bigquery_client(PipelineOptions()) self.gcp_bq_client = client or gcp_bigquery.Client( client_info=ClientInfo( user_agent="apache-beam-%s" % apache_beam.__version__)) self._unique_row_id = 0 # For testing scenarios where we pass in a client we do not want a # randomized prefix for row IDs. self._row_id_prefix = '' if client else uuid.uuid4() self._latency_histogram_metric = Metrics.histogram( self.__class__, 'latency_histogram_ms', LinearBucket(0, 20, 3000), BigQueryWrapper.HISTOGRAM_METRIC_LOGGER) if temp_dataset_id is not None and temp_table_ref is not None: raise ValueError( 'Both a BigQuery temp_dataset_id and a temp_table_ref were specified.' ' Please specify only one of these.') if temp_dataset_id and temp_dataset_id.startswith(self.TEMP_DATASET): raise ValueError( 'User provided temp dataset ID cannot start with %r' % self.TEMP_DATASET) if temp_table_ref is not None: self.temp_table_ref = temp_table_ref self.temp_dataset_id = temp_table_ref.datasetId else: self.temp_table_ref = None self._temporary_table_suffix = uuid.uuid4().hex self.temp_dataset_id = temp_dataset_id or self._get_temp_dataset() self.created_temp_dataset = False @property def unique_row_id(self): """Returns a unique row ID (str) used to avoid multiple insertions. If the row ID is provided, BigQuery will make a best effort to not insert the same row multiple times for fail and retry scenarios in which the insert request may be issued several times. This comes into play for sinks executed in a local runner. Returns: a unique row ID string """ self._unique_row_id += 1 return '%s_%d' % (self._row_id_prefix, self._unique_row_id) def _get_temp_table(self, project_id): if self.temp_table_ref: return self.temp_table_ref return parse_table_reference( table=BigQueryWrapper.TEMP_TABLE + self._temporary_table_suffix, dataset=self.temp_dataset_id, project=project_id) def _get_temp_table_project(self, fallback_project_id): """Returns the project ID for temporary table operations. If temp_table_ref exists, returns its projectId. Otherwise, returns the fallback_project_id. """ if self.temp_table_ref: return self.temp_table_ref.projectId else: return fallback_project_id def _get_temp_dataset(self): if self.temp_table_ref: return self.temp_table_ref.datasetId return BigQueryWrapper.TEMP_DATASET + self._temporary_table_suffix @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def get_query_location(self, project_id, query, use_legacy_sql): """ Get the location of tables referenced in a query. This method returns the location of the first available referenced table for user in the query and depends on the BigQuery service to provide error handling for queries that reference tables in multiple locations. """ reference = bigquery.JobReference( jobId=uuid.uuid4().hex, projectId=project_id) request = bigquery.BigqueryJobsInsertRequest( projectId=project_id, job=bigquery.Job( configuration=bigquery.JobConfiguration( dryRun=True, query=bigquery.JobConfigurationQuery( query=query, useLegacySql=use_legacy_sql, )), jobReference=reference)) response = self.client.jobs.Insert(request) if response.statistics is None: # This behavior is only expected in tests _LOGGER.warning( "Unable to get location, missing response.statistics. Query: %s", query) return None referenced_tables = response.statistics.query.referencedTables if referenced_tables: # Guards against both non-empty and non-None for table in referenced_tables: try: location = self.get_table_location( table.projectId, table.datasetId, table.tableId) except HttpForbiddenError: # Permission access for table (i.e. from authorized_view), # try next one continue _LOGGER.info( "Using location %r from table %r referenced by query %s", location, table, query) return location _LOGGER.debug( "Query %s does not reference any tables or " "you don't have permission to inspect them.", query) return None @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _insert_copy_job( self, project_id, job_id, from_table_reference, to_table_reference, create_disposition=None, write_disposition=None, job_labels=None): reference = bigquery.JobReference() reference.jobId = job_id reference.projectId = project_id request = bigquery.BigqueryJobsInsertRequest( projectId=project_id, job=bigquery.Job( configuration=bigquery.JobConfiguration( copy=bigquery.JobConfigurationTableCopy( destinationTable=to_table_reference, sourceTable=from_table_reference, createDisposition=create_disposition, writeDisposition=write_disposition, ), labels=_build_job_labels(job_labels), ), jobReference=reference, )) return self._start_job(request).jobReference @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _insert_load_job( self, project_id, job_id, table_reference, source_uris=None, source_stream=None, schema=None, write_disposition=None, create_disposition=None, additional_load_parameters=None, source_format=None, job_labels=None): if not source_uris and not source_stream: _LOGGER.warning( 'Both source URIs and source stream are not provided. BigQuery load ' 'job will not load any data.') if source_uris and source_stream: raise ValueError( 'Only one of source_uris and source_stream may be specified. ' 'Got both.') if source_uris is None: source_uris = [] additional_load_parameters = additional_load_parameters or {} job_schema = None if schema == 'SCHEMA_AUTODETECT' else schema reference = bigquery.JobReference(jobId=job_id, projectId=project_id) request = bigquery.BigqueryJobsInsertRequest( projectId=project_id, job=bigquery.Job( configuration=bigquery.JobConfiguration( load=bigquery.JobConfigurationLoad( sourceUris=source_uris, destinationTable=table_reference, schema=job_schema, writeDisposition=write_disposition, createDisposition=create_disposition, sourceFormat=source_format, useAvroLogicalTypes=True, autodetect=schema == 'SCHEMA_AUTODETECT', **additional_load_parameters), labels=_build_job_labels(job_labels), ), jobReference=reference, )) return self._start_job(request, stream=source_stream).jobReference @staticmethod def _parse_location_from_exc(content, job_id): """Parse job location from Exception content.""" if isinstance(content, bytes): content = content.decode('ascii', 'replace') # search for "Already Exists: Job :." m = re.search(r"Already Exists: Job \S+\:(\S+)\." + job_id, content) if not m: _LOGGER.warning( "Not able to parse BigQuery load job location for %s", job_id) return None return m.group(1) def _start_job( self, request: 'bigquery.BigqueryJobsInsertRequest', stream=None, ): """Inserts a BigQuery job. If the job exists already, it returns it. Args: request (bigquery.BigqueryJobsInsertRequest): An insert job request. stream (IO[bytes]): A bytes IO object open for reading. """ try: upload = None if stream: upload = Upload.FromStream(stream, mime_type=UNKNOWN_MIME_TYPE) response = self.client.jobs.Insert(request, upload=upload) _LOGGER.info( "Started BigQuery job: %s\n " "bq show -j --format=prettyjson --project_id=%s %s", response.jobReference, response.jobReference.projectId, response.jobReference.jobId) return response except HttpError as exn: if exn.status_code == 409: jobId = request.job.jobReference.jobId _LOGGER.info( "BigQuery job %s already exists, will not retry inserting it: %s", request.job.jobReference, exn) job_location = self._parse_location_from_exc(exn.content, jobId) response = request.job if not response.jobReference.location and job_location: # Request not constructed with location response.jobReference.location = job_location return response else: _LOGGER.info( "Failed to insert job %s: %s", request.job.jobReference, exn) raise @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _start_query_job( self, project_id, query, use_legacy_sql, flatten_results, job_id, priority, dry_run=False, kms_key=None, job_labels=None): reference = bigquery.JobReference(jobId=job_id, projectId=project_id) request = bigquery.BigqueryJobsInsertRequest( projectId=project_id, job=bigquery.Job( configuration=bigquery.JobConfiguration( dryRun=dry_run, query=bigquery.JobConfigurationQuery( query=query, useLegacySql=use_legacy_sql, allowLargeResults=not dry_run, destinationTable=self._get_temp_table( self._get_temp_table_project(project_id)) if not dry_run else None, flattenResults=flatten_results, priority=priority, destinationEncryptionConfiguration=bigquery. EncryptionConfiguration(kmsKeyName=kms_key)), labels=_build_job_labels(job_labels), ), jobReference=reference)) return self._start_job(request) def wait_for_bq_job(self, job_reference, sleep_duration_sec=5, max_retries=0): """Poll job until it is DONE. Args: job_reference: bigquery.JobReference instance. sleep_duration_sec: Specifies the delay in seconds between retries. max_retries: The total number of times to retry. If equals to 0, the function waits forever. Raises: `RuntimeError`: If the job is FAILED or the number of retries has been reached. """ retry = 0 while True: retry += 1 job = self.get_job( job_reference.projectId, job_reference.jobId, job_reference.location) _LOGGER.info('Job %s status: %s', job.id, job.status.state) if job.status.state == 'DONE' and job.status.errorResult: raise RuntimeError( 'BigQuery job {} failed. Error Result: {}'.format( job_reference.jobId, job.status.errorResult)) elif job.status.state == 'DONE': return True else: time.sleep(sleep_duration_sec) if max_retries != 0 and retry >= max_retries: raise RuntimeError('The maximum number of retries has been reached') @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _get_query_results( self, project_id, job_id, page_token=None, max_results=10000, location=None): request = bigquery.BigqueryJobsGetQueryResultsRequest( jobId=job_id, pageToken=page_token, projectId=project_id, maxResults=max_results, location=location) response = self.client.jobs.GetQueryResults(request) return response @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_timeout_or_quota_issues_filter) def _insert_all_rows( self, project_id, dataset_id, table_id, rows, insert_ids, skip_invalid_rows=False, ignore_unknown_values=False): """Calls the insertAll BigQuery API endpoint. Docs for this BQ call: https://cloud.google.com/bigquery/docs/reference\ /rest/v2/tabledata/insertAll.""" # The rows argument is a list of # bigquery.TableDataInsertAllRequest.RowsValueListEntry instances as # required by the InsertAll() method. resource = resource_identifiers.BigQueryTable( project_id, dataset_id, table_id) labels = { # TODO(ajamato): Add Ptransform label. monitoring_infos.SERVICE_LABEL: 'BigQuery', # Refer to any method which writes elements to BigQuery in batches # as "BigQueryBatchWrite". I.e. storage API's insertAll, or future # APIs introduced. monitoring_infos.METHOD_LABEL: 'BigQueryBatchWrite', monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.BIGQUERY_PROJECT_ID_LABEL: project_id, monitoring_infos.BIGQUERY_DATASET_LABEL: dataset_id, monitoring_infos.BIGQUERY_TABLE_LABEL: table_id, } service_call_metric = ServiceCallMetric( request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN, base_labels=labels) started_millis = int(time.time() * 1000) try: table_ref_str = '%s.%s.%s' % (project_id, dataset_id, table_id) errors = self.gcp_bq_client.insert_rows_json( table_ref_str, json_rows=rows, row_ids=insert_ids, skip_invalid_rows=skip_invalid_rows, ignore_unknown_values=ignore_unknown_values, timeout=BQ_STREAMING_INSERT_TIMEOUT_SEC) if not errors: service_call_metric.call('ok') else: for insert_error in errors: service_call_metric.call(insert_error['errors'][0]) except (ClientError, GoogleAPICallError) as e: # e.code contains the numeric http status code. service_call_metric.call(e.code) # Package exception with required fields error = {'message': e.message, 'reason': e.response.reason} # Add all rows to the errors list along with the error errors = [{"index": i, "errors": [error]} for i, _ in enumerate(rows)] except HttpError as e: service_call_metric.call(e) # Re-raise the exception so that we re-try appropriately. raise finally: self._latency_histogram_metric.update( int(time.time() * 1000) - started_millis) return not errors, errors @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_timeout_or_quota_issues_filter) def get_table(self, project_id, dataset_id, table_id): """Lookup a table's metadata object. Args: client: bigquery.BigqueryV2 instance project_id: table lookup parameter dataset_id: table lookup parameter table_id: table lookup parameter Returns: bigquery.Table instance Raises: HttpError: if lookup failed. """ request = bigquery.BigqueryTablesGetRequest( projectId=project_id, datasetId=dataset_id, tableId=table_id) response = self.client.tables.Get(request) return response def _create_table( self, project_id, dataset_id, table_id, schema, additional_parameters=None): valid_tablename = regex.fullmatch(_TABLE_PATTERN, table_id, regex.ASCII) if not valid_tablename: raise ValueError( 'Invalid BigQuery table name: %s \n' 'See https://cloud.google.com/bigquery/docs/tables#table_naming' % table_id) additional_parameters = additional_parameters or {} table = bigquery.Table( tableReference=TableReference( projectId=project_id, datasetId=dataset_id, tableId=table_id), schema=schema, **additional_parameters) request = bigquery.BigqueryTablesInsertRequest( projectId=project_id, datasetId=dataset_id, table=table) response = self.client.tables.Insert(request) _LOGGER.debug("Created the table with id %s", table_id) # The response is a bigquery.Table instance. return response @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def get_or_create_dataset( self, project_id, dataset_id, location=None, labels=None, kms_key=None, default_table_expiration_ms=None): # Check if dataset already exists otherwise create it try: dataset = self.client.datasets.Get( bigquery.BigqueryDatasetsGetRequest( projectId=project_id, datasetId=dataset_id)) self.created_temp_dataset = False return dataset except HttpError as exn: if exn.status_code == 404: _LOGGER.info( 'Dataset %s:%s does not exist so we will create it as temporary ' 'with location=%s', project_id, dataset_id, location) dataset_reference = bigquery.DatasetReference( projectId=project_id, datasetId=dataset_id) dataset = bigquery.Dataset(datasetReference=dataset_reference) if location is not None: dataset.location = location if labels is not None: dataset.labels = _build_dataset_labels(labels) if kms_key is not None: dataset.defaultEncryptionConfiguration = ( _build_dataset_encryption_config(kms_key)) if default_table_expiration_ms is not None: dataset.defaultTableExpirationMs = default_table_expiration_ms request = bigquery.BigqueryDatasetsInsertRequest( projectId=project_id, dataset=dataset) response = self.client.datasets.Insert(request) self.created_temp_dataset = True # The response is a bigquery.Dataset instance. return response else: raise @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _is_table_empty(self, project_id, dataset_id, table_id): request = bigquery.BigqueryTabledataListRequest( projectId=project_id, datasetId=dataset_id, tableId=table_id, maxResults=1) response = self.client.tabledata.List(request) # The response is a bigquery.TableDataList instance. return response.totalRows == 0 @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _delete_table(self, project_id, dataset_id, table_id): request = bigquery.BigqueryTablesDeleteRequest( projectId=project_id, datasetId=dataset_id, tableId=table_id) try: self.client.tables.Delete(request) except HttpError as exn: if exn.status_code == 404: _LOGGER.warning( 'Table %s:%s.%s does not exist', project_id, dataset_id, table_id) return else: raise @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _delete_dataset(self, project_id, dataset_id, delete_contents=True): request = bigquery.BigqueryDatasetsDeleteRequest( projectId=project_id, datasetId=dataset_id, deleteContents=delete_contents) try: self.client.datasets.Delete(request) except HttpError as exn: if exn.status_code == 404: _LOGGER.warning('Dataset %s:%s does not exist', project_id, dataset_id) return else: raise @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def get_table_location(self, project_id, dataset_id, table_id): table = self.get_table(project_id, dataset_id, table_id) return table.location # Returns true if the temporary dataset was provided by the user. def is_user_configured_dataset(self): return ( self.temp_dataset_id and not self.temp_dataset_id.startswith(self.TEMP_DATASET)) @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def create_temporary_dataset( self, project_id, location, labels=None, kms_key=None): self.get_or_create_dataset( project_id, self.temp_dataset_id, location=location, labels=labels, kms_key=kms_key) if (project_id is not None and not self.is_user_configured_dataset() and not self.created_temp_dataset): # Unittests don't pass projectIds so they can be run without error # User configured datasets are allowed to pre-exist. raise RuntimeError( 'Dataset %s:%s already exists so cannot be used as temporary.' % (project_id, self.temp_dataset_id)) @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def clean_up_temporary_dataset(self, project_id): temp_table = self._get_temp_table(project_id) try: self.client.datasets.Get( bigquery.BigqueryDatasetsGetRequest( projectId=project_id, datasetId=temp_table.datasetId)) except HttpError as exn: if exn.status_code == 404: _LOGGER.warning( 'Dataset %s:%s does not exist', project_id, temp_table.datasetId) return else: raise try: # We do not want to delete temporary datasets configured by the user hence # we just delete the temporary table in that case. if not self.is_user_configured_dataset(): self._delete_dataset(temp_table.projectId, temp_table.datasetId, True) else: self._delete_table( temp_table.projectId, temp_table.datasetId, temp_table.tableId) self.created_temp_dataset = False except HttpError as exn: if exn.status_code == 403: _LOGGER.warning( 'Permission denied to delete temporary dataset %s:%s for clean up', temp_table.projectId, temp_table.datasetId) return else: raise @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _clean_up_beam_labelled_temporary_datasets( self, project_id, dataset_id=None, table_id=None, labels=None): if isinstance(labels, dict): filter_str = _build_filter_from_labels(labels) if not self.is_user_configured_dataset() and labels is not None: response = ( self.client.datasets.List( bigquery.BigqueryDatasetsListRequest( projectId=project_id, filter=filter_str))) for dataset in response.datasets: try: dataset_id = dataset.datasetReference.datasetId self._delete_dataset(project_id, dataset_id, True) except HttpError as exn: if exn.status_code == 403: _LOGGER.warning( 'Permission denied to delete temporary dataset %s:%s for ' 'clean up.', project_id, dataset_id) return else: raise else: try: self._delete_table(project_id, dataset_id, table_id) except HttpError as exn: if exn.status_code == 403: _LOGGER.warning( 'Permission denied to delete temporary table %s:%s.%s for ' 'clean up.', project_id, dataset_id, table_id) return else: raise @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def get_job(self, project, job_id, location=None): request = bigquery.BigqueryJobsGetRequest() request.jobId = job_id request.projectId = project request.location = location return self.client.jobs.Get(request) def perform_load_job( self, destination, job_id, source_uris=None, source_stream=None, schema=None, write_disposition=None, create_disposition=None, additional_load_parameters=None, source_format=None, job_labels=None, load_job_project_id=None): """Starts a job to load data into BigQuery. Returns: bigquery.JobReference with the information about the job that was started. """ project_id = ( destination.projectId if load_job_project_id is None else load_job_project_id) return self._insert_load_job( project_id, job_id, destination, source_uris=source_uris, source_stream=source_stream, schema=schema, create_disposition=create_disposition, write_disposition=write_disposition, additional_load_parameters=additional_load_parameters, source_format=source_format, job_labels=job_labels) @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def perform_extract_job( self, destination, job_id, table_reference, destination_format, project=None, include_header=True, compression=ExportCompression.NONE, use_avro_logical_types=False, job_labels=None): """Starts a job to export data from BigQuery. Returns: bigquery.JobReference with the information about the job that was started. """ job_project = project or table_reference.projectId job_reference = bigquery.JobReference(jobId=job_id, projectId=job_project) request = bigquery.BigqueryJobsInsertRequest( projectId=job_project, job=bigquery.Job( configuration=bigquery.JobConfiguration( extract=bigquery.JobConfigurationExtract( destinationUris=destination, sourceTable=table_reference, printHeader=include_header, destinationFormat=destination_format, compression=compression, useAvroLogicalTypes=use_avro_logical_types, ), labels=_build_job_labels(job_labels), ), jobReference=job_reference, )) return self._start_job(request).jobReference @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry. retry_if_valid_input_but_server_error_and_timeout_filter) def get_or_create_table( self, project_id, dataset_id, table_id, schema, create_disposition, write_disposition, additional_create_parameters=None): """Gets or creates a table based on create and write dispositions. The function mimics the behavior of BigQuery import jobs when using the same create and write dispositions. Args: project_id: The project id owning the table. dataset_id: The dataset id owning the table. table_id: The table id. schema: A bigquery.TableSchema instance or None. create_disposition: CREATE_NEVER or CREATE_IF_NEEDED. write_disposition: WRITE_APPEND, WRITE_EMPTY or WRITE_TRUNCATE. Returns: A bigquery.Table instance if table was found or created. Raises: `RuntimeError`: For various mismatches between the state of the table and the create/write dispositions passed in. For example if the table is not empty and WRITE_EMPTY was specified then an error will be raised since the table was expected to be empty. """ from apache_beam.io.gcp.bigquery import BigQueryDisposition found_table = None try: found_table = self.get_table(project_id, dataset_id, table_id) except HttpError as exn: if exn.status_code == 404: if create_disposition == BigQueryDisposition.CREATE_NEVER: raise RuntimeError( 'Table %s:%s.%s not found but create disposition is CREATE_NEVER.' % (project_id, dataset_id, table_id)) else: raise # If table exists already then handle the semantics for WRITE_EMPTY and # WRITE_TRUNCATE write dispositions. if found_table and write_disposition in ( BigQueryDisposition.WRITE_EMPTY, BigQueryDisposition.WRITE_TRUNCATE): # Delete the table and recreate it (later) if WRITE_TRUNCATE was # specified. if write_disposition == BigQueryDisposition.WRITE_TRUNCATE: self._delete_table(project_id, dataset_id, table_id) elif (write_disposition == BigQueryDisposition.WRITE_EMPTY and not self._is_table_empty(project_id, dataset_id, table_id)): raise RuntimeError( 'Table %s:%s.%s is not empty but write disposition is WRITE_EMPTY.' % (project_id, dataset_id, table_id)) # Create a new table potentially reusing the schema from a previously # found table in case the schema was not specified. if schema is None and found_table is None: raise RuntimeError( 'Table %s:%s.%s requires a schema. None can be inferred because the ' 'table does not exist.' % (project_id, dataset_id, table_id)) if found_table and write_disposition != BigQueryDisposition.WRITE_TRUNCATE: return found_table else: created_table = None try: created_table = self._create_table( project_id=project_id, dataset_id=dataset_id, table_id=table_id, schema=schema or found_table.schema, additional_parameters=additional_create_parameters) except HttpError as exn: if exn.status_code == 409: _LOGGER.debug( 'Skipping Creation. Table %s:%s.%s already exists.' % (project_id, dataset_id, table_id)) created_table = self.get_table(project_id, dataset_id, table_id) else: raise _LOGGER.info( 'Created table %s.%s.%s with schema %s. ' 'Result: %s.', project_id, dataset_id, table_id, schema or found_table.schema, created_table) # if write_disposition == BigQueryDisposition.WRITE_TRUNCATE we delete # the table before this point. if write_disposition == BigQueryDisposition.WRITE_TRUNCATE: # BigQuery can route data to the old table for 2 mins max so wait # that much time before creating the table and writing it _LOGGER.warning( 'Sleeping for 150 seconds before the write as ' + 'BigQuery inserts can be routed to deleted table ' + 'for 2 mins after the delete and create.') # TODO(BEAM-2673): Remove this sleep by migrating to load api time.sleep(150) return created_table else: return created_table def run_query( self, project_id, query, use_legacy_sql, flatten_results, priority, dry_run=False, job_labels=None): job = self._start_query_job( project_id, query, use_legacy_sql, flatten_results, job_id=uuid.uuid4().hex, priority=priority, dry_run=dry_run, job_labels=job_labels) job_id = job.jobReference.jobId location = job.jobReference.location if dry_run: # If this was a dry run then the fact that we get here means the # query has no errors. The start_query_job would raise an error otherwise. return page_token = None while True: response = self._get_query_results( project_id, job_id, page_token, location=location) if not response.jobComplete: # The jobComplete field can be False if the query request times out # (default is 10 seconds). Note that this is a timeout for the query # request not for the actual execution of the query in the service. If # the request times out we keep trying. This situation is quite possible # if the query will return a large number of rows. _LOGGER.info('Waiting on response from query: %s ...', query) time.sleep(1.0) continue # We got some results. The last page is signalled by a missing pageToken. yield response.rows, response.schema if not response.pageToken: break page_token = response.pageToken def insert_rows( self, project_id, dataset_id, table_id, rows, insert_ids=None, skip_invalid_rows=False, ignore_unknown_values=False): """Inserts rows into the specified table. Args: project_id: The project id owning the table. dataset_id: The dataset id owning the table. table_id: The table id. rows: A list of plain Python dictionaries. Each dictionary is a row and each key in it is the name of a field. skip_invalid_rows: If there are rows with insertion errors, whether they should be skipped, and all others should be inserted successfully. ignore_unknown_values: Set this option to true to ignore unknown column names. If the input rows contain columns that are not part of the existing table's schema, those columns are ignored, and the rows are successfully inserted. Returns: A tuple (bool, errors). If first element is False then the second element will be a bigquery.InsertErrorsValueListEntry instance containing specific errors. """ # Prepare rows for insertion. Of special note is the row ID that we add to # each row in order to help BigQuery avoid inserting a row multiple times. # BigQuery will do a best-effort if unique IDs are provided. This situation # can happen during retries on failures. # TODO(silviuc): Must add support to writing TableRow's instead of dicts. insert_ids = [ str(self.unique_row_id) if not insert_ids else insert_ids[i] for i, _ in enumerate(rows) ] rows = [ fast_json_loads(fast_json_dumps(r, default=default_encoder)) for r in rows ] result, errors = self._insert_all_rows( project_id, dataset_id, table_id, rows, insert_ids, skip_invalid_rows=skip_invalid_rows, ignore_unknown_values=ignore_unknown_values) return result, errors def _convert_cell_value_to_dict(self, value, field): if field.type == 'STRING': # Input: "XYZ" --> Output: "XYZ" return value elif field.type == 'BOOLEAN': # Input: "true" --> Output: True return value == 'true' elif field.type == 'INTEGER': # Input: "123" --> Output: 123 return int(value) elif field.type == 'FLOAT': # Input: "1.23" --> Output: 1.23 return float(value) elif field.type == 'TIMESTAMP': # The UTC should come from the timezone library but this is a known # issue in python 2.7 so we'll just hardcode it as we're reading using # utcfromtimestamp. # Input: 1478134176.985864 --> Output: "2016-11-03 00:49:36.985864 UTC" dt = datetime.datetime.utcfromtimestamp(float(value)) return dt.strftime('%Y-%m-%d %H:%M:%S.%f UTC') elif field.type == 'BYTES': # Input: "YmJi" --> Output: "YmJi" return value elif field.type == 'DATE': # Input: "2016-11-03" --> Output: "2016-11-03" return value elif field.type == 'DATETIME': # Input: "2016-11-03T00:49:36" --> Output: "2016-11-03T00:49:36" return value elif field.type == 'TIME': # Input: "00:49:36" --> Output: "00:49:36" return value elif field.type == 'RECORD': # Note that a schema field object supports also a RECORD type. However # when querying, the repeated and/or record fields are flattened # unless we pass the flatten_results flag as False to the source return self.convert_row_to_dict(value, field) elif field.type == 'NUMERIC': return decimal.Decimal(value) elif field.type == 'GEOGRAPHY': return value else: raise RuntimeError('Unexpected field type: %s' % field.type) def convert_row_to_dict(self, row, schema): """Converts a TableRow instance using the schema to a Python dict.""" result = {} for index, field in enumerate(schema.fields): value = None if isinstance(schema, bigquery.TableSchema): cell = row.f[index] value = from_json_value(cell.v) if cell.v is not None else None elif isinstance(schema, bigquery.TableFieldSchema): cell = row['f'][index] value = cell['v'] if 'v' in cell else None if field.mode == 'REPEATED': if value is None: # Ideally this should never happen as repeated fields default to # returning an empty list result[field.name] = [] else: result[field.name] = [ self._convert_cell_value_to_dict(x['v'], field) for x in value ] elif value is None: if not field.mode == 'NULLABLE': raise ValueError( 'Received \'None\' as the value for the field %s ' 'but the field is not NULLABLE.' % field.name) result[field.name] = None else: result[field.name] = self._convert_cell_value_to_dict(value, field) return result @staticmethod def from_pipeline_options(pipeline_options: PipelineOptions): return BigQueryWrapper( client=BigQueryWrapper._bigquery_client(pipeline_options)) @staticmethod def _bigquery_client(pipeline_options: PipelineOptions): return bigquery.BigqueryV2( http=get_new_http(), credentials=auth.get_service_credentials(pipeline_options), response_encoding='utf8', additional_http_headers={ "user-agent": "apache-beam-%s" % apache_beam.__version__ }) class RowAsDictJsonCoder(coders.Coder): """A coder for a table row (represented as a dict) to/from a JSON string. This is the default coder for sources and sinks if the coder argument is not specified. """ def encode(self, table_row): # The normal error when dumping NAN/INF values is: # ValueError: Out of range float values are not JSON compliant # This code will catch this error to emit an error that explains # to the programmer that they have used NAN/INF values. try: return json.dumps( table_row, allow_nan=False, ensure_ascii=False, default=default_encoder).encode('utf-8') except ValueError as e: raise ValueError( '%s. %s. Row: %r' % (e, JSON_COMPLIANCE_ERROR, table_row)) def decode(self, encoded_table_row): return json.loads(encoded_table_row.decode('utf-8')) def to_type_hint(self): return Any class JsonRowWriter(io.IOBase): """ A writer which provides an IOBase-like interface for writing table rows (represented as dicts) as newline-delimited JSON strings. """ def __init__(self, file_handle): """Initialize an JsonRowWriter. Args: file_handle (io.IOBase): Output stream to write to. """ if not file_handle.writable(): raise ValueError("Output stream must be writable") self._file_handle = file_handle self._coder = RowAsDictJsonCoder() def close(self): self._file_handle.close() @property def closed(self): return self._file_handle.closed def flush(self): self._file_handle.flush() def read(self, size=-1): raise io.UnsupportedOperation("JsonRowWriter is not readable") def tell(self): return self._file_handle.tell() def writable(self): return self._file_handle.writable() def write(self, row): return self._file_handle.write(self._coder.encode(row) + b'\n') class AvroRowWriter(io.IOBase): """ A writer which provides an IOBase-like interface for writing table rows (represented as dicts) as Avro records. """ def __init__(self, file_handle, schema): """Initialize an AvroRowWriter. Args: file_handle (io.IOBase): Output stream to write Avro records to. schema (Dict[Text, Any]): BigQuery table schema. """ if not file_handle.writable(): raise ValueError("Output stream must be writable") self._file_handle = file_handle avro_schema = fastavro.parse_schema( get_avro_schema_from_table_schema(schema)) self._avro_writer = fastavro.write.Writer(self._file_handle, avro_schema) def close(self): if not self._file_handle.closed: self.flush() self._file_handle.close() @property def closed(self): return self._file_handle.closed def flush(self): if self._file_handle.closed: raise ValueError("flush on closed file") self._avro_writer.flush() self._file_handle.flush() def read(self, size=-1): raise io.UnsupportedOperation("AvroRowWriter is not readable") def tell(self): # Flush the fastavro Writer to the underlying stream, otherwise there isn't # a reliable way to determine how many bytes have been written. self._avro_writer.flush() return self._file_handle.tell() def writable(self): return self._file_handle.writable() def write(self, row): try: self._avro_writer.write(row) except (TypeError, ValueError) as ex: _, _, tb = sys.exc_info() raise ex.__class__( "Error writing row to Avro: {}\nSchema: {}\nRow: {}".format( ex, self._avro_writer.schema, row)).with_traceback(tb) class RetryStrategy(object): RETRY_ALWAYS = 'RETRY_ALWAYS' RETRY_NEVER = 'RETRY_NEVER' RETRY_ON_TRANSIENT_ERROR = 'RETRY_ON_TRANSIENT_ERROR' # Values below may be found in reasons provided either in an # error returned by a client method or by an http response as # defined in google.api_core.exceptions _NON_TRANSIENT_ERRORS = { 'invalid', 'invalidQuery', 'notImplemented', 'Bad Request', 'Unauthorized', 'Forbidden', 'Not Found', 'Not Implemented', } @staticmethod def should_retry(strategy, error_message): if strategy == RetryStrategy.RETRY_ALWAYS: return True elif strategy == RetryStrategy.RETRY_NEVER: return False elif (strategy == RetryStrategy.RETRY_ON_TRANSIENT_ERROR and error_message not in RetryStrategy._NON_TRANSIENT_ERRORS): return True else: return False class AppendDestinationsFn(DoFn): """Adds the destination to an element, making it a KV pair. Outputs a PCollection of KV-pairs where the key is a TableReference for the destination, and the value is the record itself. Experimental; no backwards compatibility guarantees. """ def __init__(self, destination): self._display_destination = destination self.destination = AppendDestinationsFn._get_table_fn(destination) def display_data(self): return {'destination': str(self._display_destination)} @staticmethod def _value_provider_or_static_val(elm): if isinstance(elm, value_provider.ValueProvider): return elm else: # The type argument is a NoOp, because we assume the argument already has # the proper formatting. return value_provider.StaticValueProvider(lambda x: x, value=elm) @staticmethod def _get_table_fn(destination): if callable(destination): return destination else: return lambda x: AppendDestinationsFn._value_provider_or_static_val( destination).get() def process(self, element, *side_inputs): yield (self.destination(element, *side_inputs), element) def beam_row_from_dict(row: dict, schema): """Converts a dictionary row to a Beam Row. Nested records and lists are supported. Args: row (dict): The row to convert. schema (str, dict, ~apache_beam.io.gcp.internal.clients.bigquery.\ bigquery_v2_messages.TableSchema): The table schema. Will be used to help convert the row. Returns: ~apache_beam.pvalue.Row: The converted row. """ if not isinstance(schema, (bigquery.TableSchema, bigquery.TableFieldSchema)): schema = get_bq_tableschema(schema) beam_row = {} for field in schema.fields: name = field.name mode = field.mode.upper() type = field.type.upper() # When writing with Storage Write API via xlang, we give the Beam Row # PCollection a hint on the schema using `with_output_types`. # This requires that each row has all the fields in the schema. # However, it's possible that some nullable fields don't appear in the row. # For this case, we create the field with a `None` value # None is also set when a repeated field is missing as BigQuery # converts Null Repeated fields to empty lists if name not in row and mode != "REQUIRED": row[name] = None value = row[name] if type in ["RECORD", "STRUCT"] and value: # if this is a list of records, we create a list of Beam Rows if mode == "REPEATED": list_of_beam_rows = [] for record in value: list_of_beam_rows.append(beam_row_from_dict(record, field)) beam_row[name] = list_of_beam_rows # otherwise, create a Beam Row from this record else: beam_row[name] = beam_row_from_dict(value, field) else: beam_row[name] = value return apache_beam.pvalue.Row(**beam_row) def get_table_schema_from_string(schema): """Transform the string table schema into a :class:`~apache_beam.io.gcp.internal.clients.bigquery.\ bigquery_v2_messages.TableSchema` instance. Args: schema (str): The string schema to be used if the BigQuery table to write has to be created. Returns: ~apache_beam.io.gcp.internal.clients.bigquery.\ bigquery_v2_messages.TableSchema: The schema to be used if the BigQuery table to write has to be created but in the :class:`~apache_beam.io.gcp.internal.clients.bigquery.\ bigquery_v2_messages.TableSchema` format. """ table_schema = bigquery.TableSchema() schema_list = [s.strip() for s in schema.split(',')] for field_and_type in schema_list: field_name, field_type = field_and_type.split(':') field_schema = bigquery.TableFieldSchema() field_schema.name = field_name field_schema.type = field_type field_schema.mode = 'NULLABLE' table_schema.fields.append(field_schema) return table_schema def table_schema_to_dict(table_schema): """Create a dictionary representation of table schema for serialization """ def get_table_field(field): """Create a dictionary representation of a table field """ result = {} result['name'] = field.name result['type'] = field.type result['mode'] = getattr(field, 'mode', 'NULLABLE') if hasattr(field, 'description') and field.description is not None: result['description'] = field.description if hasattr(field, 'fields') and field.fields: result['fields'] = [get_table_field(f) for f in field.fields] return result if not isinstance(table_schema, bigquery.TableSchema): raise ValueError("Table schema must be of the type bigquery.TableSchema") schema = {'fields': []} for field in table_schema.fields: schema['fields'].append(get_table_field(field)) return schema def get_dict_table_schema(schema): """Transform the table schema into a dictionary instance. Args: schema (str, dict, ~apache_beam.io.gcp.internal.clients.bigquery.\ bigquery_v2_messages.TableSchema): The schema to be used if the BigQuery table to write has to be created. This can either be a dict or string or in the TableSchema format. Returns: Dict[str, Any]: The schema to be used if the BigQuery table to write has to be created but in the dictionary format. """ if (isinstance(schema, (dict, value_provider.ValueProvider)) or callable(schema) or schema is None): return schema elif isinstance(schema, str): table_schema = get_table_schema_from_string(schema) return table_schema_to_dict(table_schema) elif isinstance(schema, bigquery.TableSchema): return table_schema_to_dict(schema) else: raise TypeError('Unexpected schema argument: %s.' % schema) def get_bq_tableschema(schema): """Convert the table schema to a TableSchema object. Args: schema (str, dict, ~apache_beam.io.gcp.internal.clients.bigquery.\ bigquery_v2_messages.TableSchema): The schema to be used if the BigQuery table to write has to be created. This can either be a dict or string or in the TableSchema format. Returns: ~apache_beam.io.gcp.internal.clients.bigquery.\ bigquery_v2_messages.TableSchema: The schema as a TableSchema object. """ if (isinstance(schema, (bigquery.TableSchema, value_provider.ValueProvider)) or callable(schema) or schema is None): return schema elif isinstance(schema, str): return get_table_schema_from_string(schema) elif isinstance(schema, dict): schema_string = json.dumps(schema) return parse_table_schema_from_json(schema_string) else: raise TypeError('Unexpected schema argument: %s.' % schema) def get_avro_schema_from_table_schema(schema): """Transform the table schema into an Avro schema. Args: schema (str, dict, ~apache_beam.io.gcp.internal.clients.bigquery.\ bigquery_v2_messages.TableSchema): The TableSchema to convert to Avro schema. This can either be a dict or string or in the TableSchema format. Returns: Dict[str, Any]: An Avro schema, which can be used by fastavro. """ dict_table_schema = get_dict_table_schema(schema) return bigquery_avro_tools.get_record_schema_from_dict_table_schema( "root", dict_table_schema) def get_beam_typehints_from_tableschema(schema): """Extracts Beam Python type hints from the schema. Args: schema (~apache_beam.io.gcp.internal.clients.bigquery.\ bigquery_v2_messages.TableSchema): The TableSchema to extract type hints from. Returns: List[Tuple[str, Any]]: A list of type hints that describe the input schema. Nested and repeated fields are supported. """ if not isinstance(schema, (bigquery.TableSchema, bigquery.TableFieldSchema)): schema = get_bq_tableschema(schema) typehints = [] for field in schema.fields: name, field_type, mode = field.name, field.type.upper(), field.mode.upper() if field_type in ["STRUCT", "RECORD"]: # Structs can be represented as Beam Rows. typehint = RowTypeConstraint.from_fields( get_beam_typehints_from_tableschema(field)) elif field_type in BIGQUERY_TYPE_TO_PYTHON_TYPE: typehint = BIGQUERY_TYPE_TO_PYTHON_TYPE[field_type] else: raise ValueError( f"Converting BigQuery type [{field_type}] to " "Python Beam type is not supported.") if mode == "REPEATED": typehint = Sequence[typehint] elif mode != "REQUIRED": typehint = Optional[typehint] typehints.append((name, typehint)) return typehints class BigQueryJobTypes: EXPORT = 'EXPORT' COPY = 'COPY' LOAD = 'LOAD' QUERY = 'QUERY' def generate_bq_job_name(job_name, step_id, job_type, random=None): from apache_beam.io.gcp.bigquery import BQ_JOB_NAME_TEMPLATE random = ("_%s" % random) if random else "" return str.format( BQ_JOB_NAME_TEMPLATE, job_type=job_type, job_id=job_name.replace("-", ""), step_id=step_id, random=random) def check_schema_equal( left: Union['bigquery.TableSchema', 'bigquery.TableFieldSchema'], right: Union['bigquery.TableSchema', 'bigquery.TableFieldSchema'], *, ignore_descriptions: bool = False, ignore_field_order: bool = False) -> bool: """Check whether schemas are equivalent. This comparison function differs from using == to compare TableSchema because it ignores categories, policy tags, descriptions (optionally), and field ordering (optionally). Args: left (~apache_beam.io.gcp.internal.clients.bigquery.\ bigquery_v2_messages.TableSchema, ~apache_beam.io.gcp.internal.clients.\ bigquery.bigquery_v2_messages.TableFieldSchema): One schema to compare. right (~apache_beam.io.gcp.internal.clients.bigquery.\ bigquery_v2_messages.TableSchema, ~apache_beam.io.gcp.internal.clients.\ bigquery.bigquery_v2_messages.TableFieldSchema): The other schema to compare. ignore_descriptions (bool): (optional) Whether or not to ignore field descriptions when comparing. Defaults to False. ignore_field_order (bool): (optional) Whether or not to ignore struct field order when comparing. Defaults to False. Returns: bool: True if the schemas are equivalent, False otherwise. """ if type(left) != type(right) or not isinstance( left, (bigquery.TableSchema, bigquery.TableFieldSchema)): return False if isinstance(left, bigquery.TableFieldSchema): if left.name != right.name: return False if left.type != right.type: # Check for type aliases if sorted( (left.type, right.type)) not in (["BOOL", "BOOLEAN"], ["FLOAT", "FLOAT64"], ["INT64", "INTEGER"], ["RECORD", "STRUCT"]): return False if left.mode != right.mode: return False if not ignore_descriptions and left.description != right.description: return False if isinstance(left, bigquery.TableSchema) or left.type in ("RECORD", "STRUCT"): if len(left.fields) != len(right.fields): return False if ignore_field_order: left_fields = sorted(left.fields, key=lambda field: field.name) right_fields = sorted(right.fields, key=lambda field: field.name) else: left_fields = left.fields right_fields = right.fields for left_field, right_field in zip(left_fields, right_fields): if not check_schema_equal(left_field, right_field, ignore_descriptions=ignore_descriptions, ignore_field_order=ignore_field_order): return False return True ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_tools_test.py0000644000175100017510000013120515206113032023517 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import datetime import decimal import io import json import logging import math import re import unittest from typing import Optional from typing import Sequence import fastavro import mock import numpy as np import pytz from parameterized import parameterized import apache_beam as beam from apache_beam.io.gcp import resource_identifiers from apache_beam.io.gcp.bigquery_tools import JSON_COMPLIANCE_ERROR from apache_beam.io.gcp.bigquery_tools import AvroRowWriter from apache_beam.io.gcp.bigquery_tools import BigQueryJobTypes from apache_beam.io.gcp.bigquery_tools import JsonRowWriter from apache_beam.io.gcp.bigquery_tools import RowAsDictJsonCoder from apache_beam.io.gcp.bigquery_tools import beam_row_from_dict from apache_beam.io.gcp.bigquery_tools import check_schema_equal from apache_beam.io.gcp.bigquery_tools import generate_bq_job_name from apache_beam.io.gcp.bigquery_tools import get_beam_typehints_from_tableschema from apache_beam.io.gcp.bigquery_tools import parse_table_reference from apache_beam.io.gcp.bigquery_tools import parse_table_schema_from_json from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.metrics import monitoring_infos from apache_beam.metrics.execution import MetricsEnvironment from apache_beam.options.value_provider import StaticValueProvider from apache_beam.typehints.row_type import RowTypeConstraint from apache_beam.utils.timestamp import Timestamp # Protect against environments where bigquery library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apitools.base.py.exceptions import HttpError from apitools.base.py.exceptions import HttpForbiddenError from google.api_core.exceptions import ClientError from google.api_core.exceptions import DeadlineExceeded from google.api_core.exceptions import InternalServerError except ImportError: ClientError = None DeadlineExceeded = None HttpError = None HttpForbiddenError = None InternalServerError = None google = None # pylint: enable=wrong-import-order, wrong-import-position @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestTableSchemaParser(unittest.TestCase): def test_parse_table_schema_from_json(self): string_field = bigquery.TableFieldSchema( name='s', type='STRING', mode='NULLABLE', description='s description') number_field = bigquery.TableFieldSchema( name='n', type='INTEGER', mode='REQUIRED', description='n description') record_field = bigquery.TableFieldSchema( name='r', type='RECORD', mode='REQUIRED', description='r description', fields=[string_field, number_field]) expected_schema = bigquery.TableSchema(fields=[record_field]) json_str = json.dumps({ 'fields': [{ 'name': 'r', 'type': 'RECORD', 'mode': 'REQUIRED', 'description': 'r description', 'fields': [{ 'name': 's', 'type': 'STRING', 'mode': 'NULLABLE', 'description': 's description' }, { 'name': 'n', 'type': 'INTEGER', 'mode': 'REQUIRED', 'description': 'n description' }] }] }) self.assertEqual(parse_table_schema_from_json(json_str), expected_schema) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestTableReferenceParser(unittest.TestCase): def test_calling_with_table_reference(self): table_ref = bigquery.TableReference() table_ref.projectId = 'test_project' table_ref.datasetId = 'test_dataset' table_ref.tableId = 'test_table' parsed_ref = parse_table_reference(table_ref) self.assertEqual(table_ref, parsed_ref) self.assertIsNot(table_ref, parsed_ref) def test_calling_with_callable(self): callable_ref = lambda: 'foo' parsed_ref = parse_table_reference(callable_ref) self.assertIs(callable_ref, parsed_ref) def test_calling_with_value_provider(self): value_provider_ref = StaticValueProvider(str, 'test_dataset.test_table') parsed_ref = parse_table_reference(value_provider_ref) self.assertIs(value_provider_ref, parsed_ref) @parameterized.expand([ ('project:dataset.test_table', 'project', 'dataset', 'test_table'), ('project:dataset.test-table', 'project', 'dataset', 'test-table'), ('project:dataset.test- table', 'project', 'dataset', 'test- table'), ('project.dataset. test_table', 'project', 'dataset', ' test_table'), ('project.dataset.test$table', 'project', 'dataset', 'test$table'), ]) def test_calling_with_fully_qualified_table_ref( self, fully_qualified_table: str, project_id: str, dataset_id: str, table_id: str, ): parsed_ref = parse_table_reference(fully_qualified_table) self.assertIsInstance(parsed_ref, bigquery.TableReference) self.assertEqual(parsed_ref.projectId, project_id) self.assertEqual(parsed_ref.datasetId, dataset_id) self.assertEqual(parsed_ref.tableId, table_id) def test_calling_with_partially_qualified_table_ref(self): datasetId = 'test_dataset' tableId = 'test_table' partially_qualified_table = '{}.{}'.format(datasetId, tableId) parsed_ref = parse_table_reference(partially_qualified_table) self.assertIsInstance(parsed_ref, bigquery.TableReference) self.assertEqual(parsed_ref.datasetId, datasetId) self.assertEqual(parsed_ref.tableId, tableId) def test_calling_with_insufficient_table_ref(self): table = 'test_table' self.assertRaises(ValueError, parse_table_reference, table) def test_calling_with_all_arguments(self): projectId = 'test_project' datasetId = 'test_dataset' tableId = 'test_table' parsed_ref = parse_table_reference( tableId, dataset=datasetId, project=projectId) self.assertIsInstance(parsed_ref, bigquery.TableReference) self.assertEqual(parsed_ref.projectId, projectId) self.assertEqual(parsed_ref.datasetId, datasetId) self.assertEqual(parsed_ref.tableId, tableId) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestBigQueryWrapper(unittest.TestCase): def test_delete_non_existing_dataset(self): client = mock.Mock() client.datasets.Delete.side_effect = HttpError( response={'status': '404'}, url='', content='') wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) wrapper._delete_dataset('', '') self.assertTrue(client.datasets.Delete.called) @mock.patch('time.sleep', return_value=None) def test_delete_dataset_retries_fail(self, patched_time_sleep): client = mock.Mock() client.datasets.Delete.side_effect = ValueError("Cannot delete") wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) with self.assertRaises(ValueError): wrapper._delete_dataset('', '') self.assertEqual( beam.io.gcp.bigquery_tools.MAX_RETRIES + 1, client.datasets.Delete.call_count) self.assertTrue(client.datasets.Delete.called) def test_delete_non_existing_table(self): client = mock.Mock() client.tables.Delete.side_effect = HttpError( response={'status': '404'}, url='', content='') wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) wrapper._delete_table('', '', '') self.assertTrue(client.tables.Delete.called) @mock.patch('time.sleep', return_value=None) def test_delete_table_retries_fail(self, patched_time_sleep): client = mock.Mock() client.tables.Delete.side_effect = ValueError("Cannot delete") wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) with self.assertRaises(ValueError): wrapper._delete_table('', '', '') self.assertTrue(client.tables.Delete.called) @mock.patch('time.sleep', return_value=None) def test_delete_dataset_retries_for_timeouts(self, patched_time_sleep): client = mock.Mock() client.datasets.Delete.side_effect = [ HttpError(response={'status': '408'}, url='', content=''), bigquery.BigqueryDatasetsDeleteResponse() ] wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) wrapper._delete_dataset('', '') self.assertTrue(client.datasets.Delete.called) # the function _insert_all_rows() in the wrapper calls google.cloud.bigquery, # so we have to skip that when this library is not accessible @unittest.skipIf( beam.io.gcp.bigquery_tools.gcp_bigquery is None, "bigquery library not available in this env") @mock.patch('time.sleep', return_value=None) @mock.patch( 'apitools.base.py.base_api._SkipGetCredentials', return_value=True) @mock.patch('google.cloud._http.JSONConnection.http') def test_user_agent_insert_all( self, http_mock, patched_skip_get_credentials, patched_sleep): wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper() try: wrapper._insert_all_rows('p', 'd', 't', [{'name': 'any'}], None) except: # pylint: disable=bare-except # Ignore errors. The errors come from the fact that we did not mock # the response from the API, so the overall insert_all_rows call fails # soon after the BQ API is called. pass call = http_mock.request.mock_calls[-2] self.assertIn('apache-beam-', call[2]['headers']['User-Agent']) # the function create_temporary_dataset() in the wrapper does not call # google.cloud.bigquery, so it is fine to just mock it @mock.patch( 'apache_beam.io.gcp.bigquery_tools.gcp_bigquery', return_value=mock.Mock()) @mock.patch( 'apitools.base.py.base_api._SkipGetCredentials', return_value=True) @mock.patch('time.sleep', return_value=None) def test_user_agent_create_temporary_dataset( self, sleep_mock, skip_get_credentials_mock, gcp_bigquery_mock): wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper() request_mock = mock.Mock() wrapper.client._http.request = request_mock try: wrapper.create_temporary_dataset('project-id', 'location') except: # pylint: disable=bare-except # Ignore errors. The errors come from the fact that we did not mock # the response from the API, so the overall create_dataset call fails # soon after the BQ API is called. pass call = request_mock.mock_calls[-1] self.assertIn('apache-beam-', call[2]['headers']['user-agent']) @mock.patch('time.sleep', return_value=None) def test_delete_table_retries_for_timeouts(self, patched_time_sleep): client = mock.Mock() client.tables.Delete.side_effect = [ HttpError(response={'status': '408'}, url='', content=''), bigquery.BigqueryTablesDeleteResponse() ] wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) wrapper._delete_table('', '', '') self.assertTrue(client.tables.Delete.called) @mock.patch('time.sleep', return_value=None) def test_temporary_dataset_is_unique(self, patched_time_sleep): client = mock.Mock() client.datasets.Get.return_value = bigquery.Dataset( datasetReference=bigquery.DatasetReference( projectId='project-id', datasetId='dataset_id')) wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) with self.assertRaises(RuntimeError): wrapper.create_temporary_dataset('project-id', 'location') self.assertTrue(client.datasets.Get.called) def test_get_or_create_dataset_created(self): client = mock.Mock() client.datasets.Get.side_effect = HttpError( response={'status': '404'}, url='', content='') client.datasets.Insert.return_value = bigquery.Dataset( datasetReference=bigquery.DatasetReference( projectId='project-id', datasetId='dataset_id')) wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) new_dataset = wrapper.get_or_create_dataset('project-id', 'dataset_id') self.assertEqual(new_dataset.datasetReference.datasetId, 'dataset_id') def test_create_temporary_dataset_with_kms_key(self): kms_key = ( 'projects/my-project/locations/global/keyRings/my-kr/' 'cryptoKeys/my-key') client = mock.Mock() client.datasets.Get.side_effect = HttpError( response={'status': '404'}, url='', content='') client.datasets.Insert.return_value = bigquery.Dataset( datasetReference=bigquery.DatasetReference( projectId='project-id', datasetId='temp_dataset')) wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) try: wrapper.create_temporary_dataset( 'project-id', 'location', kms_key=kms_key) except Exception: pass args, _ = client.datasets.Insert.call_args insert_request = args[0] # BigqueryDatasetsInsertRequest inserted_dataset = insert_request.dataset # Actual Dataset object # Assertions self.assertIsNotNone(inserted_dataset.defaultEncryptionConfiguration) self.assertEqual( inserted_dataset.defaultEncryptionConfiguration.kmsKeyName, kms_key) def test_get_or_create_dataset_fetched(self): client = mock.Mock() client.datasets.Get.return_value = bigquery.Dataset( datasetReference=bigquery.DatasetReference( projectId='project-id', datasetId='dataset_id')) wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) new_dataset = wrapper.get_or_create_dataset('project-id', 'dataset_id') self.assertEqual(new_dataset.datasetReference.datasetId, 'dataset_id') def test_get_or_create_table(self): client = mock.Mock() client.tables.Insert.return_value = 'table_id' client.tables.Get.side_effect = [None, 'table_id'] wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) new_table = wrapper.get_or_create_table( 'project-id', 'dataset_id', 'table_id', bigquery.TableSchema( fields=[ bigquery.TableFieldSchema( name='b', type='BOOLEAN', mode='REQUIRED') ]), False, False) self.assertEqual(new_table, 'table_id') def test_get_or_create_table_race_condition(self): client = mock.Mock() client.tables.Insert.side_effect = HttpError( response={'status': '409'}, url='', content='') client.tables.Get.side_effect = [None, 'table_id'] wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) new_table = wrapper.get_or_create_table( 'project-id', 'dataset_id', 'table_id', bigquery.TableSchema( fields=[ bigquery.TableFieldSchema( name='b', type='BOOLEAN', mode='REQUIRED') ]), False, False) self.assertEqual(new_table, 'table_id') def test_get_or_create_table_intermittent_exception(self): client = mock.Mock() client.tables.Insert.side_effect = [ HttpError(response={'status': '408'}, url='', content=''), 'table_id' ] client.tables.Get.side_effect = [None, 'table_id'] wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) new_table = wrapper.get_or_create_table( 'project-id', 'dataset_id', 'table_id', bigquery.TableSchema( fields=[ bigquery.TableFieldSchema( name='b', type='BOOLEAN', mode='REQUIRED') ]), False, False) self.assertEqual(new_table, 'table_id') @parameterized.expand(['', 'a' * 1025]) def test_get_or_create_table_invalid_tablename(self, table_id): client = mock.Mock() client.tables.Get.side_effect = [None] wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) self.assertRaises( ValueError, wrapper.get_or_create_table, 'project-id', 'dataset_id', table_id, bigquery.TableSchema( fields=[ bigquery.TableFieldSchema( name='b', type='BOOLEAN', mode='REQUIRED') ]), False, False) def test_wait_for_job_returns_true_when_job_is_done(self): def make_response(state): m = mock.Mock() m.status.errorResult = None m.status.state = state return m client, job_ref = mock.Mock(), mock.Mock() wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) # Return 'DONE' the second time get_job is called. wrapper.get_job = mock.Mock( side_effect=[make_response('RUNNING'), make_response('DONE')]) result = wrapper.wait_for_bq_job( job_ref, sleep_duration_sec=0, max_retries=5) self.assertTrue(result) def test_wait_for_job_retries_fail(self): client, response, job_ref = mock.Mock(), mock.Mock(), mock.Mock() response.status.state = 'RUNNING' wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) # Return 'RUNNING' response forever. wrapper.get_job = lambda *args: response with self.assertRaises(RuntimeError) as context: wrapper.wait_for_bq_job(job_ref, sleep_duration_sec=0, max_retries=5) self.assertEqual( 'The maximum number of retries has been reached', str(context.exception)) def test_get_query_location(self): client = mock.Mock() query = """ SELECT av.column1, table.column1 FROM `dataset.authorized_view` as av JOIN `dataset.table` as table ON av.column2 = table.column2 """ job = mock.MagicMock(spec=bigquery.Job) job.statistics.query.referencedTables = [ bigquery.TableReference( projectId="first_project_id", datasetId="first_dataset", tableId="table_used_by_authorized_view"), bigquery.TableReference( projectId="second_project_id", datasetId="second_dataset", tableId="table"), ] client.jobs.Insert.return_value = job wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) wrapper.get_table_location = mock.Mock( side_effect=[ HttpForbiddenError(response={'status': '404'}, url='', content=''), "US" ]) location = wrapper.get_query_location( project_id="second_project_id", query=query, use_legacy_sql=False) self.assertEqual("US", location) def test_perform_load_job_source_mutual_exclusivity(self): client = mock.Mock() wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) # Both source_uri and source_stream specified. with self.assertRaises(ValueError): wrapper.perform_load_job( destination=parse_table_reference('project:dataset.table'), job_id='job_id', source_uris=['gs://example.com/*'], source_stream=io.BytesIO()) # Neither source_uri nor source_stream specified. wrapper.perform_load_job( destination=parse_table_reference('project:dataset.table'), job_id='J') def test_perform_load_job_with_source_stream(self): client = mock.Mock() wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) wrapper.perform_load_job( destination=parse_table_reference('project:dataset.table'), job_id='job_id', source_stream=io.BytesIO(b'some,data')) client.jobs.Insert.assert_called_once() upload = client.jobs.Insert.call_args[1]["upload"] self.assertEqual(b'some,data', upload.stream.read()) def test_perform_load_job_with_load_job_id(self): client = mock.Mock() wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) wrapper.perform_load_job( destination=parse_table_reference('project:dataset.table'), job_id='job_id', source_uris=['gs://example.com/*'], load_job_project_id='loadId') call_args = client.jobs.Insert.call_args self.assertEqual('loadId', call_args[0][0].projectId) def verify_write_call_metric( self, project_id, dataset_id, table_id, status, count): """Check if an metric was recorded for the BQ IO write API call.""" process_wide_monitoring_infos = list( MetricsEnvironment.process_wide_container(). to_runner_api_monitoring_infos(None).values()) resource = resource_identifiers.BigQueryTable( project_id, dataset_id, table_id) labels = { # TODO(ajamato): Add Ptransform label. monitoring_infos.SERVICE_LABEL: 'BigQuery', # Refer to any method which writes elements to BigQuery in batches # as "BigQueryBatchWrite". I.e. storage API's insertAll, or future # APIs introduced. monitoring_infos.METHOD_LABEL: 'BigQueryBatchWrite', monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.BIGQUERY_PROJECT_ID_LABEL: project_id, monitoring_infos.BIGQUERY_DATASET_LABEL: dataset_id, monitoring_infos.BIGQUERY_TABLE_LABEL: table_id, monitoring_infos.STATUS_LABEL: status, } expected_mi = monitoring_infos.int64_counter( monitoring_infos.API_REQUEST_COUNT_URN, count, labels=labels) expected_mi.ClearField("start_time") found = False for actual_mi in process_wide_monitoring_infos: actual_mi.ClearField("start_time") if expected_mi == actual_mi: found = True break self.assertTrue( found, "Did not find write call metric with status: %s" % status) @unittest.skipIf(ClientError is None, 'GCP dependencies are not installed') def test_insert_rows_sets_metric_on_failure(self): MetricsEnvironment.process_wide_container().reset() client = mock.Mock() client.insert_rows_json = mock.Mock( # Fail a few times, then succeed. side_effect=[ DeadlineExceeded("Deadline Exceeded"), InternalServerError("Internal Error"), [], ]) wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) wrapper.insert_rows("my_project", "my_dataset", "my_table", []) # Expect two failing calls, then a success (i.e. two retries). self.verify_write_call_metric( "my_project", "my_dataset", "my_table", "deadline_exceeded", 1) self.verify_write_call_metric( "my_project", "my_dataset", "my_table", "internal", 1) self.verify_write_call_metric( "my_project", "my_dataset", "my_table", "ok", 1) @unittest.skipIf(ClientError is None, 'GCP dependencies are not installed') def test_start_query_job_priority_configuration(self): client = mock.Mock() wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) query_result = mock.Mock() query_result.pageToken = None wrapper._get_query_results = mock.Mock(return_value=query_result) wrapper._start_query_job( "my_project", "my_query", use_legacy_sql=False, flatten_results=False, job_id="my_job_id", priority=beam.io.BigQueryQueryPriority.BATCH) self.assertEqual( client.jobs.Insert.call_args[0][0].job.configuration.query.priority, 'BATCH') wrapper._start_query_job( "my_project", "my_query", use_legacy_sql=False, flatten_results=False, job_id="my_job_id", priority=beam.io.BigQueryQueryPriority.INTERACTIVE) self.assertEqual( client.jobs.Insert.call_args[0][0].job.configuration.query.priority, 'INTERACTIVE') def test_get_temp_table_project_with_temp_table_ref(self): """Test _get_temp_table_project returns project from temp_table_ref.""" client = mock.Mock() temp_table_ref = bigquery.TableReference( projectId='temp-project', datasetId='temp_dataset', tableId='temp_table') wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper( client, temp_table_ref=temp_table_ref) result = wrapper._get_temp_table_project('fallback-project') self.assertEqual(result, 'temp-project') def test_get_temp_table_project_without_temp_table_ref(self): """Test _get_temp_table_project returns fallback when no temp_table_ref.""" client = mock.Mock() wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) result = wrapper._get_temp_table_project('fallback-project') self.assertEqual(result, 'fallback-project') @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestRowAsDictJsonCoder(unittest.TestCase): def test_row_as_dict(self): coder = RowAsDictJsonCoder() test_value = {'s': 'abc', 'i': 123, 'f': 123.456, 'b': True} self.assertEqual(test_value, coder.decode(coder.encode(test_value))) def test_decimal_in_row_as_dict(self): decimal_value = decimal.Decimal('123456789.987654321') coder = RowAsDictJsonCoder() # Bigquery IO uses decimals to represent NUMERIC types. # To export to BQ, it's necessary to convert to strings, due to the # lower precision of JSON numbers. This means that we can't recognize # a NUMERIC when we decode from JSON, thus we match the string here. test_value = {'f': 123.456, 'b': True, 'numerico': decimal_value} output_value = {'f': 123.456, 'b': True, 'numerico': str(decimal_value)} self.assertEqual(output_value, coder.decode(coder.encode(test_value))) def json_compliance_exception(self, value): with self.assertRaisesRegex(ValueError, re.escape(JSON_COMPLIANCE_ERROR)): coder = RowAsDictJsonCoder() test_value = {'s': value} coder.decode(coder.encode(test_value)) def test_invalid_json_nan(self): self.json_compliance_exception(float('nan')) def test_invalid_json_inf(self): self.json_compliance_exception(float('inf')) def test_invalid_json_neg_inf(self): self.json_compliance_exception(float('-inf')) def test_ensure_ascii(self): coder = RowAsDictJsonCoder() test_value = {'s': '🎉'} output_value = b'{"s": "\xf0\x9f\x8e\x89"}' self.assertEqual(output_value, coder.encode(test_value)) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestJsonRowWriter(unittest.TestCase): def test_write_row(self): rows = [ { 'name': 'beam', 'game': 'dream' }, { 'name': 'team', 'game': 'cream' }, ] with io.BytesIO() as buf: # Mock close() so we can access the buffer contents # after JsonRowWriter is closed. with mock.patch.object(buf, 'close') as mock_close: writer = JsonRowWriter(buf) for row in rows: writer.write(row) writer.close() mock_close.assert_called_once() buf.seek(0) read_rows = [ json.loads(row) for row in buf.getvalue().strip().decode('utf-8').split('\n') ] self.assertEqual(read_rows, rows) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestAvroRowWriter(unittest.TestCase): def test_write_row(self): schema = bigquery.TableSchema( fields=[ bigquery.TableFieldSchema(name='stamp', type='TIMESTAMP'), bigquery.TableFieldSchema( name='number', type='FLOAT', mode='REQUIRED'), ]) stamp = datetime.datetime(2020, 2, 25, 12, 0, 0, tzinfo=pytz.utc) with io.BytesIO() as buf: # Mock close() so we can access the buffer contents # after AvroRowWriter is closed. with mock.patch.object(buf, 'close') as mock_close: writer = AvroRowWriter(buf, schema) writer.write({'stamp': stamp, 'number': float('NaN')}) writer.close() mock_close.assert_called_once() buf.seek(0) records = [r for r in fastavro.reader(buf)] self.assertEqual(len(records), 1) self.assertTrue(math.isnan(records[0]['number'])) self.assertEqual(records[0]['stamp'], stamp) class TestBQJobNames(unittest.TestCase): def test_simple_names(self): self.assertEqual( "beam_bq_job_EXPORT_beamappjobtest_abcd", generate_bq_job_name( "beamapp-job-test", "abcd", BigQueryJobTypes.EXPORT)) self.assertEqual( "beam_bq_job_LOAD_beamappjobtest_abcd", generate_bq_job_name("beamapp-job-test", "abcd", BigQueryJobTypes.LOAD)) self.assertEqual( "beam_bq_job_QUERY_beamappjobtest_abcd", generate_bq_job_name( "beamapp-job-test", "abcd", BigQueryJobTypes.QUERY)) self.assertEqual( "beam_bq_job_COPY_beamappjobtest_abcd", generate_bq_job_name("beamapp-job-test", "abcd", BigQueryJobTypes.COPY)) def test_random_in_name(self): self.assertEqual( "beam_bq_job_COPY_beamappjobtest_abcd_randome", generate_bq_job_name( "beamapp-job-test", "abcd", BigQueryJobTypes.COPY, "randome")) def test_matches_template(self): base_pattern = "beam_bq_job_[A-Z]+_[a-z0-9-]+_[a-z0-9-]+(_[a-z0-9-]+)?" job_name = generate_bq_job_name( "beamapp-job-test", "abcd", BigQueryJobTypes.COPY, "randome") self.assertRegex(job_name, base_pattern) job_name = generate_bq_job_name( "beamapp-job-test", "abcd", BigQueryJobTypes.COPY) self.assertRegex(job_name, base_pattern) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestCheckSchemaEqual(unittest.TestCase): def test_simple_schemas(self): schema1 = bigquery.TableSchema(fields=[]) self.assertTrue(check_schema_equal(schema1, schema1)) schema2 = bigquery.TableSchema( fields=[ bigquery.TableFieldSchema(name="a", mode="NULLABLE", type="INT64") ]) self.assertTrue(check_schema_equal(schema2, schema2)) self.assertFalse(check_schema_equal(schema1, schema2)) schema3 = bigquery.TableSchema( fields=[ bigquery.TableFieldSchema( name="b", mode="REPEATED", type="RECORD", fields=[ bigquery.TableFieldSchema( name="c", mode="REQUIRED", type="BOOL") ]) ]) self.assertTrue(check_schema_equal(schema3, schema3)) self.assertFalse(check_schema_equal(schema2, schema3)) def test_field_order(self): """Test that field order is ignored when ignore_field_order=True.""" schema1 = bigquery.TableSchema( fields=[ bigquery.TableFieldSchema( name="a", mode="REQUIRED", type="FLOAT64"), bigquery.TableFieldSchema(name="b", mode="REQUIRED", type="INT64"), ]) schema2 = bigquery.TableSchema(fields=list(reversed(schema1.fields))) self.assertFalse(check_schema_equal(schema1, schema2)) self.assertTrue( check_schema_equal(schema1, schema2, ignore_field_order=True)) def test_descriptions(self): """ Test that differences in description are ignored when ignore_descriptions=True. """ schema1 = bigquery.TableSchema( fields=[ bigquery.TableFieldSchema( name="a", mode="REQUIRED", type="FLOAT64", description="Field A", ), bigquery.TableFieldSchema( name="b", mode="REQUIRED", type="INT64", ), ]) schema2 = bigquery.TableSchema( fields=[ bigquery.TableFieldSchema( name="a", mode="REQUIRED", type="FLOAT64", description="Field A is for Apple"), bigquery.TableFieldSchema( name="b", mode="REQUIRED", type="INT64", description="Field B", ), ]) self.assertFalse(check_schema_equal(schema1, schema2)) self.assertTrue( check_schema_equal(schema1, schema2, ignore_descriptions=True)) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestBeamRowFromDict(unittest.TestCase): DICT_ROW = { "str": "a", "bool": True, "bytes": b'a', "int": 1, "float": 0.1, "numeric": decimal.Decimal("1.11"), "timestamp": Timestamp(1000, 100) } def get_schema_fields_with_mode(self, mode): return [{ "name": "str", "type": "STRING", "mode": mode }, { "name": "bool", "type": "boolean", "mode": mode }, { "name": "bytes", "type": "BYTES", "mode": mode }, { "name": "int", "type": "INTEGER", "mode": mode }, { "name": "float", "type": "Float", "mode": mode }, { "name": "numeric", "type": "NUMERIC", "mode": mode }, { "name": "timestamp", "type": "TIMESTAMP", "mode": mode }] def test_dict_to_beam_row_all_types_required(self): schema = {"fields": self.get_schema_fields_with_mode("REQUIRED")} expected_beam_row = beam.Row( str="a", bool=True, bytes=b'a', int=1, float=0.1, numeric=decimal.Decimal("1.11"), timestamp=Timestamp(1000, 100)) self.assertEqual( expected_beam_row, beam_row_from_dict(self.DICT_ROW, schema)) def test_dict_to_beam_row_all_types_repeated(self): schema = {"fields": self.get_schema_fields_with_mode("REPEATED")} dict_row = { "str": ["a", "b"], "bool": [True, False], "bytes": [b'a', b'b'], "int": [1, 2], "float": [0.1, 0.2], "numeric": [decimal.Decimal("1.11"), decimal.Decimal("2.22")], "timestamp": [Timestamp(1000, 100), Timestamp(2000, 200)] } expected_beam_row = beam.Row( str=["a", "b"], bool=[True, False], bytes=[b'a', b'b'], int=[1, 2], float=[0.1, 0.2], numeric=[decimal.Decimal("1.11"), decimal.Decimal("2.22")], timestamp=[Timestamp(1000, 100), Timestamp(2000, 200)]) self.assertEqual(expected_beam_row, beam_row_from_dict(dict_row, schema)) def test_dict_to_beam_row_all_types_nullable(self): schema_fields_with_nested = [{ "name": "nested_record", "type": "record", "mode": "repeated", "fields": self.get_schema_fields_with_mode("nullable") }] schema_fields_with_nested.extend( self.get_schema_fields_with_mode("nullable")) schema = {"fields": schema_fields_with_nested} dict_row = {k: None for k in self.DICT_ROW} # input dict row with missing nullable fields should still yield a full # Beam Row del dict_row['str'] del dict_row['bool'] expected_beam_row = beam.Row( nested_record=None, str=None, bool=None, bytes=None, int=None, float=None, numeric=None, timestamp=None) self.assertEqual(expected_beam_row, beam_row_from_dict(dict_row, schema)) def test_dict_to_beam_row_nested_record(self): schema_fields_with_nested = [{ "name": "nested_record", "type": "record", "fields": self.get_schema_fields_with_mode("required") }] schema_fields_with_nested.extend( self.get_schema_fields_with_mode("required")) schema = {"fields": schema_fields_with_nested} dict_row = { "nested_record": self.DICT_ROW, "str": "a", "bool": True, "bytes": b'a', "int": 1, "float": 0.1, "numeric": decimal.Decimal("1.11"), "timestamp": Timestamp(1000, 100) } expected_beam_row = beam.Row( nested_record=beam.Row( str="a", bool=True, bytes=b'a', int=1, float=0.1, numeric=decimal.Decimal("1.11"), timestamp=Timestamp(1000, 100)), str="a", bool=True, bytes=b'a', int=1, float=0.1, numeric=decimal.Decimal("1.11"), timestamp=Timestamp(1000, 100)) self.assertEqual(expected_beam_row, beam_row_from_dict(dict_row, schema)) def test_dict_to_beam_row_repeated_nested_record(self): schema_fields_with_repeated_nested_record = [{ "name": "nested_repeated_record", "type": "record", "mode": "repeated", "fields": self.get_schema_fields_with_mode("required") }] schema = {"fields": schema_fields_with_repeated_nested_record} dict_row = { "nested_repeated_record": [self.DICT_ROW, self.DICT_ROW, self.DICT_ROW], } beam_row = beam.Row( str="a", bool=True, bytes=b'a', int=1, float=0.1, numeric=decimal.Decimal("1.11"), timestamp=Timestamp(1000, 100)) expected_beam_row = beam.Row( nested_repeated_record=[beam_row, beam_row, beam_row]) self.assertEqual(expected_beam_row, beam_row_from_dict(dict_row, schema)) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestBeamTypehintFromSchema(unittest.TestCase): EXPECTED_TYPEHINTS = [("str", str), ("bool", bool), ("bytes", bytes), ("int", np.int64), ("float", np.float64), ("numeric", decimal.Decimal), ("timestamp", Timestamp)] def get_schema_fields_with_mode(self, mode): return [{ "name": "str", "type": "STRING", "mode": mode }, { "name": "bool", "type": "boolean", "mode": mode }, { "name": "bytes", "type": "BYTES", "mode": mode }, { "name": "int", "type": "INTEGER", "mode": mode }, { "name": "float", "type": "Float", "mode": mode }, { "name": "numeric", "type": "NUMERIC", "mode": mode }, { "name": "timestamp", "type": "TIMESTAMP", "mode": mode }] def test_typehints_from_required_schema(self): schema = {"fields": self.get_schema_fields_with_mode("required")} typehints = get_beam_typehints_from_tableschema(schema) self.assertEqual(typehints, self.EXPECTED_TYPEHINTS) def test_typehints_from_repeated_schema(self): schema = {"fields": self.get_schema_fields_with_mode("repeated")} typehints = get_beam_typehints_from_tableschema(schema) expected_repeated_typehints = [(name, Sequence[type]) for name, type in self.EXPECTED_TYPEHINTS] self.assertEqual(typehints, expected_repeated_typehints) def test_typehints_from_nullable_schema(self): schema = {"fields": self.get_schema_fields_with_mode("nullable")} typehints = get_beam_typehints_from_tableschema(schema) expected_nullable_typehints = [(name, Optional[type]) for name, type in self.EXPECTED_TYPEHINTS] self.assertEqual(typehints, expected_nullable_typehints) def test_typehints_from_schema_with_struct(self): schema = { "fields": [{ "name": "record", "type": "record", "mode": "required", "fields": self.get_schema_fields_with_mode("required") }] } typehints = get_beam_typehints_from_tableschema(schema) expected_typehints = [ ("record", RowTypeConstraint.from_fields(self.EXPECTED_TYPEHINTS)) ] self.assertEqual(typehints, expected_typehints) def test_typehints_from_schema_with_repeated_struct(self): schema = { "fields": [{ "name": "record", "type": "record", "mode": "repeated", "fields": self.get_schema_fields_with_mode("required") }] } typehints = get_beam_typehints_from_tableschema(schema) expected_typehints = [( "record", Sequence[RowTypeConstraint.from_fields(self.EXPECTED_TYPEHINTS)])] self.assertEqual(typehints, expected_typehints) @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestGeographyTypeSupport(unittest.TestCase): """Tests for GEOGRAPHY data type support in BigQuery.""" def test_geography_in_bigquery_type_mapping(self): """Test that GEOGRAPHY is properly mapped in type mapping.""" from apache_beam.io.gcp.bigquery_tools import BIGQUERY_TYPE_TO_PYTHON_TYPE self.assertIn("GEOGRAPHY", BIGQUERY_TYPE_TO_PYTHON_TYPE) self.assertEqual(BIGQUERY_TYPE_TO_PYTHON_TYPE["GEOGRAPHY"], str) def test_geography_field_conversion(self): """Test that GEOGRAPHY fields are converted correctly.""" from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper # Create a mock field with GEOGRAPHY type field = bigquery.TableFieldSchema() field.type = 'GEOGRAPHY' field.name = 'location' field.mode = 'NULLABLE' wrapper = BigQueryWrapper(client=mock.Mock()) # Test various WKT formats test_cases = [ "POINT(30 10)", "LINESTRING(30 10, 10 30, 40 40)", "POLYGON((30 10, 40 40, 20 40, 10 20, 30 10))", "MULTIPOINT((10 40), (40 30), (20 20), (30 10))", "GEOMETRYCOLLECTION(POINT(4 6),LINESTRING(4 6,7 10))" ] for wkt_value in test_cases: result = wrapper._convert_cell_value_to_dict(wkt_value, field) self.assertEqual(result, wkt_value) self.assertIsInstance(result, str) def test_geography_typehints_from_schema(self): """Test that GEOGRAPHY fields generate correct type hints.""" schema = { "fields": [{ "name": "location", "type": "GEOGRAPHY", "mode": "REQUIRED" }, { "name": "optional_location", "type": "GEOGRAPHY", "mode": "NULLABLE" }, { "name": "locations", "type": "GEOGRAPHY", "mode": "REPEATED" }] } typehints = get_beam_typehints_from_tableschema(schema) expected_typehints = [("location", str), ("optional_location", Optional[str]), ("locations", Sequence[str])] self.assertEqual(typehints, expected_typehints) def test_geography_beam_row_conversion(self): """Test converting dictionary with GEOGRAPHY to Beam Row.""" schema = { "fields": [{ "name": "id", "type": "INTEGER", "mode": "REQUIRED" }, { "name": "location", "type": "GEOGRAPHY", "mode": "NULLABLE" }, { "name": "name", "type": "STRING", "mode": "REQUIRED" }] } row_dict = {"id": 1, "location": "POINT(30 10)", "name": "Test Location"} beam_row = beam_row_from_dict(row_dict, schema) self.assertEqual(beam_row.id, 1) self.assertEqual(beam_row.location, "POINT(30 10)") self.assertEqual(beam_row.name, "Test Location") def test_geography_beam_row_conversion_with_null(self): """Test converting dictionary with null GEOGRAPHY to Beam Row.""" schema = { "fields": [{ "name": "id", "type": "INTEGER", "mode": "REQUIRED" }, { "name": "location", "type": "GEOGRAPHY", "mode": "NULLABLE" }] } row_dict = {"id": 1, "location": None} beam_row = beam_row_from_dict(row_dict, schema) self.assertEqual(beam_row.id, 1) self.assertIsNone(beam_row.location) def test_geography_beam_row_conversion_repeated(self): """Test converting dictionary with repeated GEOGRAPHY to Beam Row.""" schema = { "fields": [{ "name": "id", "type": "INTEGER", "mode": "REQUIRED" }, { "name": "locations", "type": "GEOGRAPHY", "mode": "REPEATED" }] } row_dict = { "id": 1, "locations": ["POINT(30 10)", "POINT(40 20)", "LINESTRING(0 0, 1 1)"] } beam_row = beam_row_from_dict(row_dict, schema) self.assertEqual(beam_row.id, 1) self.assertEqual(len(beam_row.locations), 3) self.assertEqual(beam_row.locations[0], "POINT(30 10)") self.assertEqual(beam_row.locations[1], "POINT(40 20)") self.assertEqual(beam_row.locations[2], "LINESTRING(0 0, 1 1)") def test_geography_json_encoding(self): """Test that GEOGRAPHY values are properly JSON encoded.""" coder = RowAsDictJsonCoder() row_with_geography = {"id": 1, "location": "POINT(30 10)", "name": "Test"} encoded = coder.encode(row_with_geography) decoded = coder.decode(encoded) self.assertEqual(decoded["location"], "POINT(30 10)") self.assertIsInstance(decoded["location"], str) def test_geography_with_special_characters(self): """Test GEOGRAPHY values with special characters and geometries.""" from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper field = bigquery.TableFieldSchema() field.type = 'GEOGRAPHY' field.name = 'complex_geo' field.mode = 'NULLABLE' wrapper = BigQueryWrapper(client=mock.Mock()) # Test complex WKT with various coordinate systems and precision complex_wkt = ( "POLYGON((-122.4194 37.7749, -122.4094 37.7849, " "-122.3994 37.7749, -122.4194 37.7749))") result = wrapper._convert_cell_value_to_dict(complex_wkt, field) self.assertEqual(result, complex_wkt) self.assertIsInstance(result, str) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_write_it_test.py0000644000175100017510000004733415206113032024216 0ustar00runnerrunner#!/usr/bin/env python # -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for BigQuery sources and sinks.""" # pytype: skip-file import base64 import datetime import logging import secrets import time import unittest from decimal import Decimal import hamcrest as hc import mock import pytest import pytz from parameterized import param from parameterized import parameterized from tenacity import retry from tenacity import stop_after_attempt import apache_beam as beam from apache_beam.io.gcp.bigquery import BigQueryWriteFn from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.bigquery_tools import FileFormat from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultMatcher from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # Protect against environments where bigquery library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apitools.base.py.exceptions import HttpError except ImportError: HttpError = None # pylint: enable=wrong-import-order, wrong-import-position _LOGGER = logging.getLogger(__name__) class BigQueryWriteIntegrationTests(unittest.TestCase): BIG_QUERY_DATASET_ID = 'python_write_to_table_' def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.runner_name = type(self.test_pipeline.runner).__name__ self.project = self.test_pipeline.get_option('project') self.bigquery_client = BigQueryWrapper() self.dataset_id = '%s%d%s' % ( self.BIG_QUERY_DATASET_ID, int(time.time()), secrets.token_hex(3)) self.bigquery_client.get_or_create_dataset(self.project, self.dataset_id) _LOGGER.info( "Created dataset %s in project %s", self.dataset_id, self.project) def tearDown(self): request = bigquery.BigqueryDatasetsDeleteRequest( projectId=self.project, datasetId=self.dataset_id, deleteContents=True) try: _LOGGER.info( "Deleting dataset %s in project %s", self.dataset_id, self.project) self.bigquery_client.client.datasets.Delete(request) except HttpError: _LOGGER.debug( 'Failed to clean up dataset %s in project %s', self.dataset_id, self.project) def create_table(self, table_name): table_schema = bigquery.TableSchema() table_field = bigquery.TableFieldSchema() table_field.name = 'int64' table_field.type = 'INT64' table_field.mode = 'REQUIRED' table_schema.fields.append(table_field) table_field = bigquery.TableFieldSchema() table_field.name = 'bytes' table_field.type = 'BYTES' table_schema.fields.append(table_field) table_field = bigquery.TableFieldSchema() table_field.name = 'date' table_field.type = 'DATE' table_schema.fields.append(table_field) table_field = bigquery.TableFieldSchema() table_field.name = 'time' table_field.type = 'TIME' table_schema.fields.append(table_field) table = bigquery.Table( tableReference=bigquery.TableReference( projectId=self.project, datasetId=self.dataset_id, tableId=table_name), schema=table_schema) request = bigquery.BigqueryTablesInsertRequest( projectId=self.project, datasetId=self.dataset_id, table=table) self.bigquery_client.client.tables.Insert(request) @pytest.mark.it_postcommit def test_big_query_write(self): table_name = 'python_write_table' table_id = '{}.{}'.format(self.dataset_id, table_name) input_data = [ { 'number': 1, 'str': 'abc' }, { 'number': 2, 'str': 'def' }, { 'number': 3, 'str': '你好' }, { 'number': 4, 'str': 'привет' }, ] table_schema = { "fields": [{ "name": "number", "type": "INTEGER" }, { "name": "str", "type": "STRING" }] } pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query="SELECT number, str FROM %s" % table_id, data=[( 1, 'abc', ), ( 2, 'def', ), ( 3, '你好', ), ( 4, 'привет', )]) ] args = self.test_pipeline.get_full_options_as_args( on_success_matcher=hc.all_of(*pipeline_verifiers)) with beam.Pipeline(argv=args) as p: # pylint: disable=expression-not-assigned ( p | 'create' >> beam.Create(input_data) | 'write' >> beam.io.WriteToBigQuery( table_id, schema=table_schema, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=beam.io.BigQueryDisposition.WRITE_EMPTY)) @pytest.mark.it_postcommit def test_big_query_write_schema_autodetect(self): if self.runner_name == 'TestDataflowRunner': self.skipTest('DataflowRunner does not support schema autodetection') table_name = 'python_write_table_schema_autodetect' table_id = '{}.{}'.format(self.dataset_id, table_name) input_data = [ { 'number': 1, 'str': 'abc' }, { 'number': 2, 'str': 'def' }, ] pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query="SELECT number, str FROM %s" % table_id, data=[( 1, 'abc', ), ( 2, 'def', )]) ] args = self.test_pipeline.get_full_options_as_args( on_success_matcher=hc.all_of(*pipeline_verifiers)) with beam.Pipeline(argv=args) as p: # pylint: disable=expression-not-assigned ( p | 'create' >> beam.Create(input_data) | 'write' >> beam.io.WriteToBigQuery( table_id, method=beam.io.WriteToBigQuery.Method.FILE_LOADS, schema=beam.io.gcp.bigquery.SCHEMA_AUTODETECT, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=beam.io.BigQueryDisposition.WRITE_EMPTY, temp_file_format=FileFormat.JSON)) @pytest.mark.it_postcommit def test_big_query_write_new_types(self): table_name = 'python_new_types_table' table_id = '{}.{}'.format(self.dataset_id, table_name) row_data = { 'float': 0.33, 'numeric': Decimal('10'), 'bytes': base64.b64encode(b'\xab\xac').decode('utf-8'), 'date': '3000-12-31', 'time': '23:59:59', 'datetime': '2018-12-31T12:44:31', 'timestamp': '2018-12-31 12:44:31.744957 UTC', 'geo': 'POINT(30 10)' } input_data = [row_data] # add rows with only one key value pair and None values for all other keys for key, value in row_data.items(): input_data.append({key: value}) table_schema = { "fields": [{ "name": "float", "type": "FLOAT" }, { "name": "numeric", "type": "NUMERIC" }, { "name": "bytes", "type": "BYTES" }, { "name": "date", "type": "DATE" }, { "name": "time", "type": "TIME" }, { "name": "datetime", "type": "DATETIME" }, { "name": "timestamp", "type": "TIMESTAMP" }, { "name": "geo", "type": "GEOGRAPHY" }] } expected_row = ( 0.33, Decimal('10'), b'\xab\xac', datetime.date(3000, 12, 31), datetime.time(23, 59, 59), datetime.datetime(2018, 12, 31, 12, 44, 31), datetime.datetime(2018, 12, 31, 12, 44, 31, 744957, tzinfo=pytz.utc), 'POINT(30 10)', ) expected_data = [expected_row] # add rows with only one key value pair and None values for all other keys for i, value in enumerate(expected_row): row = [None] * len(expected_row) row[i] = value expected_data.append(tuple(row)) pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query='SELECT float, numeric, bytes, date, time, datetime,' 'timestamp, geo FROM %s' % table_id, data=expected_data) ] args = self.test_pipeline.get_full_options_as_args( on_success_matcher=hc.all_of(*pipeline_verifiers)) with beam.Pipeline(argv=args) as p: # pylint: disable=expression-not-assigned ( p | 'create' >> beam.Create(input_data) | 'write' >> beam.io.WriteToBigQuery( table_id, schema=table_schema, create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=beam.io.BigQueryDisposition.WRITE_EMPTY)) @pytest.mark.it_postcommit def test_big_query_write_without_schema(self): table_name = 'python_no_schema_table' self.create_table(table_name) table_id = '{}.{}'.format(self.dataset_id, table_name) input_data = [{ 'int64': 1, 'bytes': b'xyw', 'date': '2011-01-01', 'time': '23:59:59.999999' }, { 'int64': 2, 'bytes': b'abc', 'date': '2000-01-01', 'time': '00:00:00' }, { 'int64': 3, 'bytes': b'\xe4\xbd\xa0\xe5\xa5\xbd', 'date': '3000-12-31', 'time': '23:59:59' }, { 'int64': 4, 'bytes': b'\xab\xac\xad', 'date': '2000-01-01', 'time': '00:00:00' }] # bigquery io expects bytes to be base64 encoded values for row in input_data: row['bytes'] = base64.b64encode(row['bytes']) pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query="SELECT int64, bytes, date, time FROM %s" % table_id, data=[( 1, b'xyw', datetime.date(2011, 1, 1), datetime.time(23, 59, 59, 999999), ), ( 2, b'abc', datetime.date(2000, 1, 1), datetime.time(0, 0, 0), ), ( 3, b'\xe4\xbd\xa0\xe5\xa5\xbd', datetime.date(3000, 12, 31), datetime.time(23, 59, 59), ), ( 4, b'\xab\xac\xad', datetime.date(2000, 1, 1), datetime.time(0, 0, 0), )]) ] args = self.test_pipeline.get_full_options_as_args( on_success_matcher=hc.all_of(*pipeline_verifiers)) with beam.Pipeline(argv=args) as p: # pylint: disable=expression-not-assigned ( p | 'create' >> beam.Create(input_data) | 'write' >> beam.io.WriteToBigQuery( table_id, write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND, temp_file_format=FileFormat.JSON)) @pytest.mark.it_postcommit def test_big_query_write_insert_errors_reporting(self): """ Test that errors returned by beam.io.WriteToBigQuery contain both the failed rows and the reason for it failing. """ table_name = 'python_write_table_insert_errors' table_id = '{}.{}'.format(self.dataset_id, table_name) input_data = [{ 'number': 1, 'str': 'some_string', }, { 'number': 2 }, { 'number': 3, 'str': 'some_string', 'additional_field_str': 'some_string', }] table_schema = { "fields": [{ "name": "number", "type": "INTEGER", 'mode': 'REQUIRED' }, { "name": "str", "type": "STRING", 'mode': 'REQUIRED' }] } bq_result_errors = [( { "number": 2 }, [{ "reason": "invalid", "location": "", "debugInfo": "", "message": "Missing required field: Msg_0_CLOUD_QUERY_TABLE.str." }], ), ({ "number": 3, "str": "some_string", "additional_field_str": "some_string" }, [{ "reason": "invalid", "location": "additional_field_str", "debugInfo": "", "message": "no such field: additional_field_str." }])] pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query="SELECT number, str FROM %s" % table_id, data=[(1, 'some_string')]), ] args = self.test_pipeline.get_full_options_as_args( on_success_matcher=hc.all_of(*pipeline_verifiers)) with beam.Pipeline(argv=args) as p: # pylint: disable=expression-not-assigned errors = ( p | 'create' >> beam.Create(input_data) | 'write' >> beam.io.WriteToBigQuery( table_id, schema=table_schema, method='STREAMING_INSERTS', insert_retry_strategy='RETRY_NEVER', create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND)) assert_that( errors[BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS] | 'ParseErrors' >> beam.Map(lambda err: (err[1], err[2])) | 'ToList' >> beam.combiners.ToList() | 'SortErrors' >> beam.Map( lambda errs: sorted(errs, key=lambda x: x[0].get("number", 0))), equal_to( [sorted(bq_result_errors, key=lambda x: x[0].get("number", 0))])) @pytest.mark.it_postcommit def test_big_query_write_insert_non_transient_api_call_error(self): """ Test that non-transient GoogleAPICallError errors returned by beam.io.WriteToBigQuery are not retried and result in FAILED_ROWS containing both the failed rows and the reason for failure. """ table_name = 'this_table_does_not_exist' table_id = '{}.{}'.format(self.dataset_id, table_name) input_data = [{ 'number': 1, 'str': 'some_string', }] table_schema = { "fields": [{ "name": "number", "type": "INTEGER", 'mode': 'NULLABLE' }, { "name": "str", "type": "STRING", 'mode': 'NULLABLE' }] } bq_result_errors = [({ 'number': 1, 'str': 'some_string', }, "Not Found")] args = self.test_pipeline.get_full_options_as_args() with beam.Pipeline(argv=args) as p: # pylint: disable=expression-not-assigned errors = ( p | 'create' >> beam.Create(input_data) | beam.WindowInto(beam.transforms.window.FixedWindows(10)) | 'write' >> beam.io.WriteToBigQuery( table_id, schema=table_schema, method='STREAMING_INSERTS', insert_retry_strategy='RETRY_ON_TRANSIENT_ERROR', create_disposition=beam.io.BigQueryDisposition.CREATE_NEVER, write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND)) assert_that( errors[BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS] | 'ParseErrors' >> beam.Map(lambda err: (err[1], err[2][0]["reason"])), equal_to(bq_result_errors)) @pytest.mark.it_postcommit @parameterized.expand([ param(file_format=FileFormat.AVRO), param(file_format=FileFormat.JSON), param(file_format=None), ]) @mock.patch( "apache_beam.io.gcp.bigquery_file_loads._MAXIMUM_SOURCE_URIS", new=1) @retry(reraise=True, stop=stop_after_attempt(3)) def test_big_query_write_temp_table_append_schema_update(self, file_format): """ Test that nested schema update options and schema relaxation are respected when appending to an existing table via temporary tables. _MAXIMUM_SOURCE_URIS and max_file_size are both set to 1 to force multiple load jobs and usage of temporary tables. """ table_name = 'python_append_schema_update' self.create_table(table_name) table_id = '{}.{}'.format(self.dataset_id, table_name) # bytes, date, time fields are optional and omitted in the test # only required and new columns are specified table_schema = { "fields": [{ "name": "int64", "type": "INT64", "mode": "NULLABLE", }, { "name": "bool", "type": "BOOL", }, { "name": "nested_field", "type": "RECORD", "mode": "REPEATED", "fields": [ { "name": "fruit", "type": "STRING", "mode": "NULLABLE" }, ] }] } input_data = [{ "int64": 1, "bool": True, "nested_field": [{ "fruit": "Apple" }] }, { "bool": False, "nested_field": [{ "fruit": "Mango" }] }, { "int64": None, "bool": True, "nested_field": [{ "fruit": "Banana" }] }] args = self.test_pipeline.get_full_options_as_args( on_success_matcher=BigqueryFullResultMatcher( project=self.project, query=""" SELECT bytes, date, time, int64, bool, fruit FROM {}, UNNEST(nested_field) as nested_field ORDER BY fruit """.format(table_id), data=[(None, None, None, 1, True, "Apple"), (None, None, None, None, True, "Banana"), ( None, None, None, None, False, "Mango")])) with beam.Pipeline(argv=args) as p: # pylint: disable=expression-not-assigned ( p | 'create' >> beam.Create(input_data) | 'write' >> beam.io.WriteToBigQuery( table_id, schema=table_schema, write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND, max_file_size=1, # bytes method=beam.io.WriteToBigQuery.Method.FILE_LOADS, additional_bq_parameters={ 'schemaUpdateOptions': [ 'ALLOW_FIELD_ADDITION', 'ALLOW_FIELD_RELAXATION' ] }, temp_file_format=file_format)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigquery_write_perf_test.py0000644000175100017510000000752015206113032024527 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A pipeline that writes data from Synthetic Source to a BigQuery table. Besides of the standard options, there are options with special meaning: * output_dataset - BQ dataset name. * output_table - BQ table name. The table will be removed after test completion, * input_options - options for Synthetic Source: num_records - number of rows to be inserted, value_size - the length of a single row, key_size - required option, but its value has no meaning. Example test run on DataflowRunner: python -m apache_beam.io.gcp.bigquery_write_perf_test \ --test-pipeline-options=" --runner=TestDataflowRunner --project=... --region=... --staging_location=gs://... --temp_location=gs://... --sdk_location=.../dist/apache-beam-x.x.x.dev0.tar.gz --publish_to_big_query=true --metrics_dataset=gs://... --metrics_table=... --output_dataset=... --output_table=... --input_options='{ \"num_records\": 1024, \"key_size\": 1, \"value_size\": 1024, }'" This setup will result in a table of 1MB size. """ # pytype: skip-file import logging from apache_beam import Map from apache_beam import ParDo from apache_beam.io import BigQueryDisposition from apache_beam.io import Read from apache_beam.io import WriteToBigQuery from apache_beam.io.gcp.bigquery_tools import parse_table_schema_from_json from apache_beam.io.gcp.tests import utils from apache_beam.testing.load_tests.load_test import LoadTest from apache_beam.testing.load_tests.load_test_metrics_utils import CountMessages from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime from apache_beam.testing.synthetic_pipeline import SyntheticSource class BigQueryWritePerfTest(LoadTest): def __init__(self): super().__init__() self.output_dataset = self.pipeline.get_option('output_dataset') self.output_table = self.pipeline.get_option('output_table') def test(self): SCHEMA = parse_table_schema_from_json( '{"fields": [{"name": "data", "type": "BYTES"}]}') def format_record(record): # Since Synthetic Source returns data as a dictionary, we should skip one # of the part import base64 return {'data': base64.b64encode(record[1])} ( # pylint: disable=expression-not-assigned self.pipeline | 'Produce rows' >> Read( SyntheticSource(self.parse_synthetic_source_options())) | 'Count messages' >> ParDo(CountMessages(self.metrics_namespace)) | 'Format' >> Map(format_record) | 'Measure time' >> ParDo(MeasureTime(self.metrics_namespace)) | 'Write to BigQuery' >> WriteToBigQuery( dataset=self.output_dataset, table=self.output_table, schema=SCHEMA, create_disposition=BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=BigQueryDisposition.WRITE_TRUNCATE)) def cleanup(self): """Removes an output BQ table.""" utils.delete_bq_table( self.project_id, self.output_dataset, self.output_table) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) BigQueryWritePerfTest().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigtableio.py0000644000175100017510000003536515206113032021524 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """BigTable connector This module implements writing to BigTable tables. The default mode is to set row data to write to BigTable tables. The syntax supported is described here: https://cloud.google.com/bigtable/docs/quickstart-cbt BigTable connector can be used as main outputs. A main output (common case) is expected to be massive and will be split into manageable chunks and processed in parallel. In the example below we created a list of rows then passed to the GeneratedDirectRows DoFn to set the Cells and then we call the BigTableWriteFn to insert those generated rows in the table. main_table = (p | beam.Create(self._generate()) | WriteToBigTable(project_id, instance_id, table_id)) """ # pytype: skip-file import logging import struct import apache_beam as beam from apache_beam.internal.metrics.metric import ServiceCallMetric from apache_beam.io.gcp import resource_identifiers from apache_beam.metrics import Metrics from apache_beam.metrics import monitoring_infos from apache_beam.metrics.metric import Lineage from apache_beam.transforms import PTransform from apache_beam.transforms.display import DisplayDataItem from apache_beam.transforms.external import BeamJarExpansionService from apache_beam.transforms.external import SchemaAwareExternalTransform from apache_beam.typehints.row_type import RowTypeConstraint _LOGGER = logging.getLogger(__name__) FLUSH_COUNT = 1000 MAX_ROW_BYTES = 5242880 # 5MB try: from google.cloud.bigtable import Client from google.cloud.bigtable.batcher import MutationsBatcher from google.cloud.bigtable.row import Cell from google.cloud.bigtable.row import PartialRowData except ImportError: _LOGGER.warning( 'ImportError: from google.cloud.bigtable import Client', exc_info=True) __all__ = ['WriteToBigTable', 'ReadFromBigtable'] class _BigTableWriteFn(beam.DoFn): """ Creates the connector can call and add_row to the batcher using each row in beam pipe line Args: project_id(str): GCP Project ID instance_id(str): GCP Instance ID table_id(str): GCP Table ID flush_count(int): Max number of rows to flush max_row_bytes(int) Max number of row mutations size to flush """ def __init__( self, project_id, instance_id, table_id, flush_count, max_row_bytes): """ Constructor of the Write connector of Bigtable Args: project_id(str): GCP Project of to write the Rows instance_id(str): GCP Instance to write the Rows table_id(str): GCP Table to write the `DirectRows` flush_count(int): Max number of rows to flush max_row_bytes(int) Max number of row mutations size to flush """ super().__init__() self.beam_options = { 'project_id': project_id, 'instance_id': instance_id, 'table_id': table_id, 'flush_count': flush_count, 'max_row_bytes': max_row_bytes, } self.table = None self.batcher = None self.service_call_metric = None self.written = Metrics.counter(self.__class__, 'Written Row') def __getstate__(self): return self.beam_options def __setstate__(self, options): self.beam_options = options self.table = None self.batcher = None self.service_call_metric = None self.written = Metrics.counter(self.__class__, 'Written Row') def write_mutate_metrics(self, status_list): for status in status_list: code = status.code if status else None grpc_status_string = ( ServiceCallMetric.bigtable_error_code_to_grpc_status_string(code)) self.service_call_metric.call(grpc_status_string) def start_service_call_metrics(self, project_id, instance_id, table_id): resource = resource_identifiers.BigtableTable( project_id, instance_id, table_id) labels = { monitoring_infos.SERVICE_LABEL: 'BigTable', # TODO(JIRA-11985): Add Ptransform label. monitoring_infos.METHOD_LABEL: 'google.bigtable.v2.MutateRows', monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.BIGTABLE_PROJECT_ID_LABEL: ( self.beam_options['project_id']), monitoring_infos.INSTANCE_ID_LABEL: self.beam_options['instance_id'], monitoring_infos.TABLE_ID_LABEL: self.beam_options['table_id'] } return ServiceCallMetric( request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN, base_labels=labels) def start_bundle(self): if self.table is None: client = Client(project=self.beam_options['project_id']) instance = client.instance(self.beam_options['instance_id']) self.table = instance.table(self.beam_options['table_id']) self.service_call_metric = self.start_service_call_metrics( self.beam_options['project_id'], self.beam_options['instance_id'], self.beam_options['table_id']) self.batcher = MutationsBatcher( self.table, batch_completed_callback=self.write_mutate_metrics, flush_count=self.beam_options['flush_count'], max_row_bytes=self.beam_options['max_row_bytes']) def process(self, row): self.written.inc() # You need to set the timestamp in the cells in this row object, # when we do a retry we will mutating the same object, but, with this # we are going to set our cell with new values. # Example: # direct_row.set_cell('cf1', # 'field1', # 'value1', # timestamp=datetime.now()) self.batcher.mutate(row) def finish_bundle(self): if self.batcher: self.batcher.close() self.batcher = None # Report Lineage metrics on write Lineage.sinks().add( 'bigtable', self.beam_options['project_id'], self.beam_options['instance_id'], self.beam_options['table_id']) def display_data(self): return { 'projectId': DisplayDataItem( self.beam_options['project_id'], label='Bigtable Project Id'), 'instanceId': DisplayDataItem( self.beam_options['instance_id'], label='Bigtable Instance Id'), 'tableId': DisplayDataItem( self.beam_options['table_id'], label='Bigtable Table Id') } class WriteToBigTable(beam.PTransform): """A transform that writes rows to a Bigtable table. Takes an input PCollection of `DirectRow` objects containing un-committed mutations. For more information about this row object, visit https://cloud.google.com/python/docs/reference/bigtable/latest/row#class-googlecloudbigtablerowdirectrowrowkey-tablenone If flag `use_cross_language` is set to true, this transform will use the multi-language transforms framework to inject the Java native write transform into the pipeline. """ URN = "beam:schematransform:org.apache.beam:bigtable_write:v1" def __init__( self, project_id, instance_id, table_id, use_cross_language=False, expansion_service=None, flush_count=FLUSH_COUNT, max_row_bytes=MAX_ROW_BYTES, ): """Initialize an WriteToBigTable transform. :param table_id: The ID of the table to write to. :param instance_id: The ID of the instance where the table resides. :param project_id: The GCP project ID. :param use_cross_language: If set to True, will use the Java native transform via cross-language. :param expansion_service: The address of the expansion service in the case of using cross-language. If no expansion service is provided, will attempt to run the default GCP expansion service. :type flush_count: int :param flush_count: (Optional) Max number of rows to flush. Default is FLUSH_COUNT (1000 rows). :type max_row_bytes: int :param max_row_bytes: (Optional) Max number of row mutations size to flush. Default is MAX_ROW_BYTES (5 MB). """ super().__init__() self._table_id = table_id self._instance_id = instance_id self._project_id = project_id self._use_cross_language = use_cross_language if use_cross_language: self._expansion_service = ( expansion_service or BeamJarExpansionService( 'sdks:java:io:google-cloud-platform:expansion-service:build')) self.schematransform_config = ( SchemaAwareExternalTransform.discover_config( self._expansion_service, self.URN)) self._flush_count = flush_count self._max_row_bytes = max_row_bytes def expand(self, input): if self._use_cross_language: external_write = SchemaAwareExternalTransform( identifier=self.schematransform_config.identifier, expansion_service=self._expansion_service, rearrange_based_on_discovery=True, table_id=self._table_id, instance_id=self._instance_id, project_id=self._project_id) return ( input | beam.ParDo(self._DirectRowMutationsToBeamRow()).with_output_types( RowTypeConstraint.from_fields( [("key", bytes), ("mutations", list[dict[str, bytes]])])) | external_write) else: return ( input | beam.ParDo( _BigTableWriteFn( self._project_id, self._instance_id, self._table_id, flush_count=self._flush_count, max_row_bytes=self._max_row_bytes))) class _DirectRowMutationsToBeamRow(beam.DoFn): def process(self, direct_row): args = {"key": direct_row.row_key, "mutations": []} # start accumulating mutations in a list for mutation in direct_row._get_mutations(): if mutation.__contains__("set_cell"): mutation_dict = { "type": b'SetCell', "family_name": mutation.set_cell.family_name.encode('utf-8'), "column_qualifier": mutation.set_cell.column_qualifier, "value": mutation.set_cell.value, "timestamp_micros": struct.pack( '>q', mutation.set_cell.timestamp_micros) } elif mutation.__contains__("delete_from_column"): mutation_dict = { "type": b'DeleteFromColumn', "family_name": mutation.delete_from_column.family_name.encode( 'utf-8'), "column_qualifier": mutation.delete_from_column.column_qualifier } time_range = mutation.delete_from_column.time_range if time_range.start_timestamp_micros: mutation_dict['start_timestamp_micros'] = struct.pack( '>q', time_range.start_timestamp_micros) if time_range.end_timestamp_micros: mutation_dict['end_timestamp_micros'] = struct.pack( '>q', time_range.end_timestamp_micros) elif mutation.__contains__("delete_from_family"): mutation_dict = { "type": b'DeleteFromFamily', "family_name": mutation.delete_from_family.family_name.encode( 'utf-8') } elif mutation.__contains__("delete_from_row"): mutation_dict = {"type": b'DeleteFromRow'} else: raise ValueError("Unexpected mutation") args["mutations"].append(mutation_dict) yield beam.Row(**args) class ReadFromBigtable(PTransform): """Reads rows from Bigtable. Returns a PCollection of PartialRowData objects, each representing a Bigtable row. For more information about this row object, visit https://cloud.google.com/python/docs/reference/bigtable/latest/row#class-googlecloudbigtablerowpartialrowdatarowkey """ URN = "beam:schematransform:org.apache.beam:bigtable_read:v1" def __init__(self, project_id, instance_id, table_id, expansion_service=None): """Initialize a ReadFromBigtable transform. :param table_id: The ID of the table to read from. :param instance_id: The ID of the instance where the table resides. :param project_id: The GCP project ID. :param expansion_service: The address of the expansion service. If no expansion service is provided, will attempt to run the default GCP expansion service. """ super().__init__() self._table_id = table_id self._instance_id = instance_id self._project_id = project_id self._expansion_service = ( expansion_service or BeamJarExpansionService( 'sdks:java:io:google-cloud-platform:expansion-service:build')) self.schematransform_config = SchemaAwareExternalTransform.discover_config( self._expansion_service, self.URN) def expand(self, input): external_read = SchemaAwareExternalTransform( identifier=self.schematransform_config.identifier, expansion_service=self._expansion_service, rearrange_based_on_discovery=True, table_id=self._table_id, instance_id=self._instance_id, project_id=self._project_id, flatten=False) return ( input.pipeline | external_read | beam.ParDo(self._BeamRowToPartialRowData())) # PartialRowData has some useful methods for querying data within a row. # To make use of those methods and to give Python users a more familiar # object, we process each Beam Row and return a PartialRowData equivalent. class _BeamRowToPartialRowData(beam.DoFn): def process(self, row): key = row.key families = row.column_families # initialize PartialRowData object partial_row: PartialRowData = PartialRowData(key) for fam_name, col_fam in families.items(): if fam_name not in partial_row.cells: partial_row.cells[fam_name] = {} for col_qualifier, cells in col_fam.items(): # store column qualifier as bytes to follow PartialRowData behavior col_qualifier_bytes = col_qualifier.encode() if col_qualifier not in partial_row.cells[fam_name]: partial_row.cells[fam_name][col_qualifier_bytes] = [] for cell in cells: value = cell.value timestamp_micros = cell.timestamp_micros partial_row.cells[fam_name][col_qualifier_bytes].append( Cell(value, timestamp_micros)) yield partial_row ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigtableio_it_test.py0000644000175100017510000003662115206113032023253 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration tests for BigTable service.""" import logging import os import secrets import time import unittest # pytype: skip-file from datetime import datetime from datetime import timezone import pytest import apache_beam as beam from apache_beam.io.gcp import bigtableio from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to _LOGGER = logging.getLogger(__name__) # Protect against environments where bigtable library is not available. try: from apitools.base.py.exceptions import HttpError from google.cloud.bigtable import client from google.cloud.bigtable.row import Cell from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.row import PartialRowData from google.cloud.bigtable.row_filters import TimestampRange from google.cloud.bigtable.table import Table from google.cloud.bigtable_admin_v2.types import instance except ImportError as e: client = None HttpError = None def instance_prefix(instance): datestr = "".join(filter(str.isdigit, str(datetime.now(timezone.utc).date()))) instance_id = '%s-%s-%s' % (instance, datestr, secrets.token_hex(4)) assert len(instance_id) < 34, "instance id length needs to be within [6, 33]" return instance_id @pytest.mark.uses_gcp_java_expansion_service @pytest.mark.uses_transform_service @unittest.skipIf(client is None, 'Bigtable dependencies are not installed') @unittest.skipUnless( os.environ.get('EXPANSION_JARS') or os.environ.get('TRANSFORM_SERVICE_PORT'), "A valid expansion service is not available for executing the " "cross-language test.") class TestReadFromBigTableIT(unittest.TestCase): INSTANCE = "bt-read-tests" TABLE_ID = "test-table" def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.args = self.test_pipeline.get_full_options_as_args() self.project = self.test_pipeline.get_option('project') instance_id = instance_prefix(self.INSTANCE) self.client = client.Client(admin=True, project=self.project) # create cluster and instance self.instance = self.client.instance( instance_id, display_name=self.INSTANCE, instance_type=instance.Instance.Type.DEVELOPMENT) cluster = self.instance.cluster("test-cluster", "us-central1-a") operation = self.instance.create(clusters=[cluster]) operation.result(timeout=500) _LOGGER.info( "Created instance [%s] in project [%s]", self.instance.instance_id, self.project) # create table inside instance self.table = self.instance.table(self.TABLE_ID) self.table.create() _LOGGER.info("Created table [%s]", self.table.table_id) if (os.environ.get('TRANSFORM_SERVICE_PORT')): self._transform_service_address = ( 'localhost:' + os.environ.get('TRANSFORM_SERVICE_PORT')) else: self._transform_service_address = None def tearDown(self): try: _LOGGER.info( "Deleting table [%s] and instance [%s]", self.table.table_id, self.instance.instance_id) self.table.delete() self.instance.delete() except HttpError: _LOGGER.warning( "Failed to clean up table [%s] and instance [%s]", self.table.table_id, self.instance.instance_id) def add_rows(self, num_rows, num_families, num_columns_per_family): cells = [] for i in range(1, num_rows + 1): key = 'key-' + str(i) row = DirectRow(key, self.table) for j in range(num_families): fam_name = 'test_col_fam_' + str(j) # create the table's column families one time only if i == 1: col_fam = self.table.column_family(fam_name) col_fam.create() for k in range(1, num_columns_per_family + 1): row.set_cell(fam_name, f'col-{j}-{k}', f'value-{i}-{j}-{k}') # after all mutations on the row are done, commit to Bigtable row.commit() # read the same row back from Bigtable to get the expected data # accumulate rows in `cells` read_row: PartialRowData = self.table.read_row(key) cells.append(read_row.cells) return cells def test_read_xlang(self): # create rows and retrieve expected cells expected_cells = self.add_rows( num_rows=5, num_families=3, num_columns_per_family=4) with beam.Pipeline(argv=self.args) as p: cells = ( p | bigtableio.ReadFromBigtable( project_id=self.project, instance_id=self.instance.instance_id, table_id=self.table.table_id, expansion_service=self._transform_service_address) | "Extract cells" >> beam.Map(lambda row: row._cells)) assert_that(cells, equal_to(expected_cells)) @pytest.mark.uses_gcp_java_expansion_service @pytest.mark.uses_transform_service @unittest.skipIf(client is None, 'Bigtable dependencies are not installed') @unittest.skipUnless( os.environ.get('EXPANSION_JARS') or os.environ.get('TRANSFORM_SERVICE_PORT'), "A valid expansion service is not available for executing the " "cross-language test.") class TestWriteToBigtableXlangIT(unittest.TestCase): # These are integration tests for the cross-language write transform. INSTANCE = "bt-write-xlang" TABLE_ID = "test-table" @classmethod def setUpClass(cls): cls.test_pipeline = TestPipeline(is_integration_test=True) cls.project = cls.test_pipeline.get_option('project') cls.args = cls.test_pipeline.get_full_options_as_args() instance_id = instance_prefix(cls.INSTANCE) cls.client = client.Client(admin=True, project=cls.project) # create cluster and instance cls.instance = cls.client.instance( instance_id, display_name=cls.INSTANCE, instance_type=instance.Instance.Type.DEVELOPMENT) cluster = cls.instance.cluster("test-cluster", "us-central1-a") operation = cls.instance.create(clusters=[cluster]) operation.result(timeout=500) _LOGGER.warning( "Created instance [%s] in project [%s]", cls.instance.instance_id, cls.project) def setUp(self): # create table inside instance self.table: Table = self.instance.table( '%s-%s-%s' % (self.TABLE_ID, str(int(time.time())), secrets.token_hex(3))) self.table.create() _LOGGER.info("Created table [%s]", self.table.table_id) if (os.environ.get('TRANSFORM_SERVICE_PORT')): self._transform_service_address = ( 'localhost:' + os.environ.get('TRANSFORM_SERVICE_PORT')) else: self._transform_service_address = None def tearDown(self): try: _LOGGER.info("Deleting table [%s]", self.table.table_id) self.table.delete() except HttpError: _LOGGER.warning("Failed to clean up table [%s]", self.table.table_id) @classmethod def tearDownClass(cls): try: _LOGGER.info("Deleting instance [%s]", cls.instance.instance_id) cls.instance.delete() except HttpError: _LOGGER.warning( "Failed to clean up instance [%s]", cls.instance.instance_id) def run_pipeline(self, rows): with beam.Pipeline(argv=self.args) as p: _ = ( p | beam.Create(rows) | bigtableio.WriteToBigTable( project_id=self.project, instance_id=self.instance.instance_id, table_id=self.table.table_id, use_cross_language=True, expansion_service=self._transform_service_address)) def test_set_mutation(self): row1: DirectRow = DirectRow('key-1') row2: DirectRow = DirectRow('key-2') col_fam = self.table.column_family('col_fam') col_fam.create() # expected cells row1_col1_cell = Cell(b'val1-1', 100_000_000) row1_col2_cell = Cell(b'val1-2', 200_000_000) row2_col1_cell = Cell(b'val2-1', 100_000_000) row2_col2_cell = Cell(b'val2-2', 200_000_000) # When setting this cell, we won't set a timestamp. We expect the timestamp # to default to -1, and Bigtable will set it to system time at insertion. row2_col1_no_timestamp = Cell(b'val2-2-notimestamp', time.time()) # rows sent to write transform row1.set_cell( 'col_fam', b'col-1', row1_col1_cell.value, row1_col1_cell.timestamp) row1.set_cell( 'col_fam', b'col-2', row1_col2_cell.value, row1_col2_cell.timestamp) row2.set_cell( 'col_fam', b'col-1', row2_col1_cell.value, row2_col1_cell.timestamp) row2.set_cell( 'col_fam', b'col-2', row2_col2_cell.value, row2_col2_cell.timestamp) # don't set a timestamp here. it should default to -1 row2.set_cell('col_fam', b'col-no-timestamp', row2_col1_no_timestamp.value) self.run_pipeline([row1, row2]) # after write transform executes, get actual rows from table actual_row1: PartialRowData = self.table.read_row('key-1') actual_row2: PartialRowData = self.table.read_row('key-2') # check actual rows match with expected rows (value and timestamp) self.assertEqual( row1_col1_cell, actual_row1.find_cells('col_fam', b'col-1')[0]) self.assertEqual( row1_col2_cell, actual_row1.find_cells('col_fam', b'col-2')[0]) self.assertEqual( row2_col1_cell, actual_row2.find_cells('col_fam', b'col-1')[0]) self.assertEqual( row2_col2_cell, actual_row2.find_cells('col_fam', b'col-2')[0]) # check mutation that doesn't have a timestamp set is handled properly: self.assertEqual( row2_col1_no_timestamp.value, actual_row2.find_cells('col_fam', b'col-no-timestamp')[0].value) # Bigtable sets timestamp as insertion time, which is later than the # time.time() we set when creating this test case cell_timestamp = actual_row2.find_cells('col_fam', b'col-no-timestamp')[0].timestamp self.assertTrue( row2_col1_no_timestamp.timestamp < cell_timestamp, msg="Expected cell with unset timestamp to have ingestion time " f"attached, but was {cell_timestamp}") def test_delete_cells_mutation(self): col_fam = self.table.column_family('col_fam') col_fam.create() # write a row with two columns to the table beforehand. write_row: DirectRow = DirectRow('key-1', self.table) write_row.set_cell('col_fam', b'col-1', b'val-1') write_row.set_cell('col_fam', b'col-2', b'val-2') write_row.commit() # prepare a row that will delete cells in one of the columns. delete_row: DirectRow = DirectRow('key-1') delete_row.delete_cell('col_fam', b'col-1') self.run_pipeline([delete_row]) # after transform executes, get actual row from table actual_row: PartialRowData = self.table.read_row('key-1') # we deleted all the cells in 'col-1', so this check should throw an error with self.assertRaises(KeyError): actual_row.find_cells('col_fam', b'col-1') # check the cell in col-2 is still there col2_cells = actual_row.find_cells('col_fam', b'col-2') self.assertEqual(1, len(col2_cells)) self.assertEqual(b'val-2', col2_cells[0].value) def test_delete_cells_with_timerange_mutation(self): col_fam = self.table.column_family('col_fam') col_fam.create() # write two cells in a column to the table beforehand. write_row: DirectRow = DirectRow('key-1', self.table) write_row.set_cell( 'col_fam', b'col', b'val', datetime.fromtimestamp(100_000_000, tz=timezone.utc)) write_row.commit() write_row.set_cell( 'col_fam', b'col', b'new-val', datetime.fromtimestamp(200_000_000, tz=timezone.utc)) write_row.commit() # prepare a row that will delete cells within a timestamp range. delete_row: DirectRow = DirectRow('key-1') delete_row.delete_cell( 'col_fam', b'col', time_range=TimestampRange( start=datetime.fromtimestamp(99_999_999, tz=timezone.utc), end=datetime.fromtimestamp(100_000_001, tz=timezone.utc))) self.run_pipeline([delete_row]) # after transform executes, get actual row from table actual_row: PartialRowData = self.table.read_row('key-1') # we deleted one cell within the timestamp range. # check the other (newer) cell still exists. cells = actual_row.find_cells('col_fam', b'col') self.assertEqual(1, len(cells)) self.assertEqual(b'new-val', cells[0].value) self.assertEqual( datetime.fromtimestamp(200_000_000, tz=timezone.utc), cells[0].timestamp) def test_delete_column_family_mutation(self): # create two column families col_fam = self.table.column_family('col_fam-1') col_fam.create() col_fam = self.table.column_family('col_fam-2') col_fam.create() # write a row with values in both column families to the table beforehand. write_row: DirectRow = DirectRow('key-1', self.table) write_row.set_cell('col_fam-1', b'col', b'val') write_row.set_cell('col_fam-2', b'col', b'val') write_row.commit() # prepare a row that will delete a column family from the row delete_row: DirectRow = DirectRow('key-1') delete_row.delete_cells('col_fam-1', delete_row.ALL_COLUMNS) self.run_pipeline([delete_row]) # after transform executes, get actual row from table actual_row: PartialRowData = self.table.read_row('key-1') # we deleted column family 'col_fam-1', so this check should throw an error with self.assertRaises(KeyError): actual_row.find_cells('col_fam-1', b'col') # check we have one column family left with the correct cell value self.assertEqual(1, len(actual_row.cells)) self.assertEqual(b'val', actual_row.cell_value('col_fam-2', b'col')) def test_delete_row_mutation(self): write_row1: DirectRow = DirectRow('key-1', self.table) write_row2: DirectRow = DirectRow('key-2', self.table) col_fam = self.table.column_family('col_fam') col_fam.create() # write a couple of rows to the table beforehand write_row1.set_cell('col_fam', b'col', b'val-1') write_row1.commit() write_row2.set_cell('col_fam', b'col', b'val-2') write_row2.commit() # prepare a row that will delete itself delete_row: DirectRow = DirectRow('key-1') delete_row.delete() self.run_pipeline([delete_row]) # after write transform executes, get actual rows from table actual_row1: PartialRowData = self.table.read_row('key-1') actual_row2: PartialRowData = self.table.read_row('key-2') # we deleted row with key 'key-1', check it doesn't exist anymore # the Bigtable API doesn't throw an error here, just returns a None value self.assertEqual(None, actual_row1) # check row 2 exists with the correct cell value in col self.assertEqual(b'val-2', actual_row2.cell_value('col_fam', b'col')) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/bigtableio_test.py0000644000175100017510000003602315206113032022553 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for BigTable service.""" import logging import string import unittest import uuid # pytype: skip-file from datetime import datetime from datetime import timezone from random import choice from mock import MagicMock from mock import patch import apache_beam as beam from apache_beam.internal.metrics.metric import ServiceCallMetric from apache_beam.io.gcp import bigtableio from apache_beam.io.gcp import resource_identifiers from apache_beam.metrics import monitoring_infos from apache_beam.metrics.execution import MetricsEnvironment from apache_beam.metrics.metric import Lineage from apache_beam.testing.test_pipeline import TestPipeline _LOGGER = logging.getLogger(__name__) # Protect against environments where bigtable library is not available. try: from google.cloud.bigtable import client from google.cloud.bigtable.batcher import MutationsBatcher from google.cloud.bigtable.instance import Instance from google.cloud.bigtable.row import Cell from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.row import PartialRowData from google.cloud.bigtable.row_filters import TimestampRange from google.cloud.bigtable.table import Table from google.rpc.code_pb2 import ALREADY_EXISTS from google.rpc.code_pb2 import OK from google.rpc.status_pb2 import Status except ImportError as e: client = None @unittest.skipIf(client is None, 'Bigtable dependencies are not installed') class TestBeamRowToPartialRowData(unittest.TestCase): # Beam Row schema: # - key: bytes # - column_families: dict[str, dict[str, list[beam.Row(cell_schema)]]] # cell_schema: # - value: bytes # - timestamp_micros: int def test_beam_row_to_bigtable_row(self): # create test beam row families = { 'family_1': { 'col_1': [ beam.Row(value=b'a-1', timestamp_micros=100_000_000), beam.Row(value=b'b-1', timestamp_micros=200_000_000), beam.Row(value=b'c-1', timestamp_micros=300_000_000) ], 'col_2': [ beam.Row(value=b'a-2', timestamp_micros=400_000_000), beam.Row(value=b'b-2', timestamp_micros=500_000_000), beam.Row(value=b'c-2', timestamp_micros=600_000_000) ], }, 'family_2': { 'column_qualifier': [ beam.Row(value=b'val-1', timestamp_micros=700_000_000), beam.Row(value=b'val-2', timestamp_micros=800_000_000), beam.Row(value=b'val-3', timestamp_micros=900_000_000) ] } } beam_row = beam.Row(key=b'key', column_families=families) # get equivalent bigtable row doFn = bigtableio.ReadFromBigtable._BeamRowToPartialRowData() bigtable_row: PartialRowData = next(doFn.process(beam_row)) # using bigtable utils (PartialRowData methods), check that beam row data # landed in the right cells self.assertEqual(beam_row.key, bigtable_row.row_key) self.assertEqual([ Cell(c.value, c.timestamp_micros) for c in beam_row.column_families['family_1']['col_1'] ], bigtable_row.find_cells('family_1', b'col_1')) self.assertEqual([ Cell(c.value, c.timestamp_micros) for c in beam_row.column_families['family_1']['col_2'] ], bigtable_row.find_cells('family_1', b'col_2')) self.assertEqual([ Cell(c.value, c.timestamp_micros) for c in beam_row.column_families['family_2']['column_qualifier'] ], bigtable_row.find_cells('family_2', b'column_qualifier')) @unittest.skipIf(client is None, 'Bigtable dependencies are not installed') class TestBigtableDirectRowToBeamRow(unittest.TestCase): doFn = bigtableio.WriteToBigTable._DirectRowMutationsToBeamRow() def test_set_cell(self): # create some set cell mutations direct_row: DirectRow = DirectRow('key-1') direct_row.set_cell( 'col_fam', b'col', b'a', datetime.fromtimestamp(100_000).replace(tzinfo=timezone.utc)) direct_row.set_cell( 'col_fam', b'other-col', b'b', datetime.fromtimestamp(200_000).replace(tzinfo=timezone.utc)) direct_row.set_cell( 'other_col_fam', b'col', b'c', datetime.fromtimestamp(300_000).replace(tzinfo=timezone.utc)) # get equivalent beam row beam_row = next(self.doFn.process(direct_row)) # sort both lists of mutations for convenience beam_row_mutations = sorted(beam_row.mutations, key=lambda m: m['value']) bt_row_mutations = sorted( direct_row._get_mutations(), key=lambda m: m.set_cell.value) self.assertEqual(beam_row.key, direct_row.row_key) self.assertEqual(len(beam_row_mutations), len(bt_row_mutations)) # check that the values in each beam mutation is equal to the original # Bigtable direct row mutations for i in range(len(beam_row_mutations)): beam_mutation = beam_row_mutations[i] bt_mutation = bt_row_mutations[i].set_cell self.assertEqual(beam_mutation['type'], b'SetCell') self.assertEqual( beam_mutation['family_name'].decode(), bt_mutation.family_name) self.assertEqual( beam_mutation['column_qualifier'], bt_mutation.column_qualifier) self.assertEqual(beam_mutation['value'], bt_mutation.value) self.assertEqual( int.from_bytes(beam_mutation['timestamp_micros'], 'big'), bt_mutation.timestamp_micros) def test_delete_cells(self): # create some delete cell mutations. one with a timestamp range direct_row: DirectRow = DirectRow('key-1') direct_row.delete_cell('col_fam', b'col-1') direct_row.delete_cell( 'other_col_fam', b'col-2', time_range=TimestampRange( start=datetime.fromtimestamp(10_000_000, tz=timezone.utc))) direct_row.delete_cells( 'another_col_fam', [b'col-3', b'col-4', b'col-5'], time_range=TimestampRange( start=datetime.fromtimestamp(50_000_000, tz=timezone.utc), end=datetime.fromtimestamp(100_000_000, tz=timezone.utc))) # get equivalent beam row beam_row = next(self.doFn.process(direct_row)) # sort both lists of mutations for convenience beam_row_mutations = sorted( beam_row.mutations, key=lambda m: m['column_qualifier']) bt_row_mutations = sorted( direct_row._get_mutations(), key=lambda m: m.delete_from_column.column_qualifier) self.assertEqual(beam_row.key, direct_row.row_key) self.assertEqual(len(beam_row_mutations), len(bt_row_mutations)) # check that the values in each beam mutation is equal to the original # Bigtable direct row mutations for i in range(len(beam_row_mutations)): beam_mutation = beam_row_mutations[i] bt_mutation = bt_row_mutations[i].delete_from_column print(bt_mutation) self.assertEqual(beam_mutation['type'], b'DeleteFromColumn') self.assertEqual( beam_mutation['family_name'].decode(), bt_mutation.family_name) self.assertEqual( beam_mutation['column_qualifier'], bt_mutation.column_qualifier) # check we set a timestamp range only when appropriate if bt_mutation.time_range.start_timestamp_micros: self.assertEqual( int.from_bytes(beam_mutation['start_timestamp_micros'], 'big'), bt_mutation.time_range.start_timestamp_micros) else: self.assertTrue('start_timestamp_micros' not in beam_mutation) if bt_mutation.time_range.end_timestamp_micros: self.assertEqual( int.from_bytes(beam_mutation['end_timestamp_micros'], 'big'), bt_mutation.time_range.end_timestamp_micros) else: self.assertTrue('end_timestamp_micros' not in beam_mutation) def test_delete_column_family(self): # create mutation to delete column family direct_row: DirectRow = DirectRow('key-1') direct_row.delete_cells('col_fam-1', direct_row.ALL_COLUMNS) direct_row.delete_cells('col_fam-2', direct_row.ALL_COLUMNS) # get equivalent beam row beam_row = next(self.doFn.process(direct_row)) # sort both lists of mutations for convenience beam_row_mutations = sorted( beam_row.mutations, key=lambda m: m['family_name']) bt_row_mutations = sorted( direct_row._get_mutations(), key=lambda m: m.delete_from_column.family_name) self.assertEqual(beam_row.key, direct_row.row_key) self.assertEqual(len(beam_row_mutations), len(bt_row_mutations)) # check that the values in each beam mutation is equal to the original # Bigtable direct row mutations for i in range(len(beam_row_mutations)): beam_mutation = beam_row_mutations[i] bt_mutation = bt_row_mutations[i].delete_from_family self.assertEqual(beam_mutation['type'], b'DeleteFromFamily') self.assertEqual( beam_mutation['family_name'].decode(), bt_mutation.family_name) def test_delete_row(self): # create mutation to delete the Bigtable row direct_row: DirectRow = DirectRow('key-1') direct_row.delete() # get equivalent beam row beam_row = next(self.doFn.process(direct_row)) self.assertEqual(beam_row.key, direct_row.row_key) beam_mutation = beam_row.mutations[0] self.assertEqual(beam_mutation['type'], b'DeleteFromRow') @unittest.skipIf(client is None, 'Bigtable dependencies are not installed') class TestWriteBigTable(unittest.TestCase): TABLE_PREFIX = "python-test" _PROJECT_ID = TABLE_PREFIX + "-" + str(uuid.uuid4())[:8] _INSTANCE_ID = TABLE_PREFIX + "-" + str(uuid.uuid4())[:8] _TABLE_ID = TABLE_PREFIX + "-" + str(uuid.uuid4())[:8] def setUp(self): client = MagicMock() instance = Instance(self._INSTANCE_ID, client) self.table = Table(self._TABLE_ID, instance) def test_write(self): direct_rows = [self.generate_row(i) for i in range(5)] # TODO(https://github.com/apache/beam/issues/34549): This test relies on # lineage metrics which Prism doesn't seem to handle correctly. Defaulting # to FnApiRunner instead. runner = 'FnApiRunner' with patch.object(MutationsBatcher, 'mutate'), \ patch.object(MutationsBatcher, 'close'), TestPipeline(runner) as p: _ = p | beam.Create(direct_rows) | bigtableio.WriteToBigTable( self._PROJECT_ID, self._INSTANCE_ID, self._TABLE_ID) self.assertSetEqual( Lineage.query(p.result.metrics(), Lineage.SINK), set([ f"bigtable:{self._PROJECT_ID}.{self._INSTANCE_ID}.{self._TABLE_ID}" ])) def test_write_metrics(self): MetricsEnvironment.process_wide_container().reset() write_fn = bigtableio._BigTableWriteFn( self._PROJECT_ID, self._INSTANCE_ID, self._TABLE_ID, flush_count=1000, max_row_bytes=5242880) write_fn.table = self.table write_fn.start_bundle() number_of_rows = 2 error = Status() error.message = 'Entity already exists.' error.code = ALREADY_EXISTS success = Status() success.message = 'Success' success.code = OK rows_response = [error, success] * number_of_rows with patch.object(Table, 'mutate_rows', return_value=rows_response): direct_rows = [self.generate_row(i) for i in range(number_of_rows * 2)] for direct_row in direct_rows: write_fn.process(direct_row) try: write_fn.finish_bundle() except: # pylint: disable=bare-except # Currently we fail the bundle when there are any failures. # TODO(https://github.com/apache/beam/issues/21396): remove after # bigtableio can selectively retry. pass self.verify_write_call_metric( self._PROJECT_ID, self._INSTANCE_ID, self._TABLE_ID, ServiceCallMetric.bigtable_error_code_to_grpc_status_string( ALREADY_EXISTS), 2) self.verify_write_call_metric( self._PROJECT_ID, self._INSTANCE_ID, self._TABLE_ID, ServiceCallMetric.bigtable_error_code_to_grpc_status_string(OK), 2) def generate_row(self, index=0): rand = choice(string.ascii_letters + string.digits) value = ''.join(rand for i in range(100)) column_family_id = 'cf1' key = "beam_key%s" % ('{0:07}'.format(index)) direct_row = DirectRow(row_key=key) for column_id in range(10): direct_row.set_cell( column_family_id, ('field%s' % column_id).encode('utf-8'), value, datetime.now()) return direct_row def verify_write_call_metric( self, project_id, instance_id, table_id, status, count): """Check if a metric was recorded for the Datastore IO write API call.""" process_wide_monitoring_infos = list( MetricsEnvironment.process_wide_container(). to_runner_api_monitoring_infos(None).values()) resource = resource_identifiers.BigtableTable( project_id, instance_id, table_id) labels = { monitoring_infos.SERVICE_LABEL: 'BigTable', monitoring_infos.METHOD_LABEL: 'google.bigtable.v2.MutateRows', monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.BIGTABLE_PROJECT_ID_LABEL: project_id, monitoring_infos.INSTANCE_ID_LABEL: instance_id, monitoring_infos.TABLE_ID_LABEL: table_id, monitoring_infos.STATUS_LABEL: status } expected_mi = monitoring_infos.int64_counter( monitoring_infos.API_REQUEST_COUNT_URN, count, labels=labels) expected_mi.ClearField("start_time") found = False for actual_mi in process_wide_monitoring_infos: actual_mi.ClearField("start_time") if expected_mi == actual_mi: found = True break self.assertTrue( found, "Did not find write call metric with status: %s" % status) def test_custom_flush_config(self): direct_rows = [self.generate_row(0)] with patch.object( MutationsBatcher, '__init__', return_value=None) as mock_init, \ patch.object(MutationsBatcher, 'mutate'), \ patch.object(MutationsBatcher, 'close'), TestPipeline() as p: _ = p | beam.Create(direct_rows) | bigtableio.WriteToBigTable( self._PROJECT_ID, self._INSTANCE_ID, self._TABLE_ID, flush_count=1001, max_row_bytes=5000001) mock_init.assert_called_once() call_args = mock_init.call_args.kwargs assert call_args['flush_count'] == 1001 assert call_args['max_row_bytes'] == 5000001 if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3674064 apache_beam-2.74.0/apache_beam/io/gcp/datastore/0000755000175100017510000000000015206113157021033 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/__init__.py0000644000175100017510000000142015206113032023131 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3734071 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/0000755000175100017510000000000015206113157022073 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/__init__.py0000644000175100017510000000142015206113032024171 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/datastore_write_it_pipeline.py0000644000175100017510000001373215206113032030224 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A job that write Entries into Datastore. The pipelines behave in the steps below. 1. Create and write Entities to Datastore 2. (Optional) If read limit was provided, read it and confirm that the expected Entities were read. 3. Query the written Entities and verify result. 4. Delete Entries. 5. Query the written Entities, verify no results. """ # pytype: skip-file import argparse import hashlib import logging import uuid import apache_beam as beam from apache_beam.io.gcp.datastore.v1new.datastoreio import DeleteFromDatastore from apache_beam.io.gcp.datastore.v1new.datastoreio import ReadFromDatastore from apache_beam.io.gcp.datastore.v1new.datastoreio import WriteToDatastore from apache_beam.io.gcp.datastore.v1new.types import Entity from apache_beam.io.gcp.datastore.v1new.types import Key from apache_beam.io.gcp.datastore.v1new.types import Query from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to _LOGGER = logging.getLogger(__name__) def new_pipeline_with_job_name(pipeline_options, job_name, suffix): """Create a pipeline with the given job_name and a suffix.""" gcp_options = pipeline_options.view_as(GoogleCloudOptions) # DirectRunner doesn't have a job name. if job_name: gcp_options.job_name = job_name + suffix return TestPipeline(options=pipeline_options) class EntityWrapper(object): """ Create a Cloud Datastore entity from the given string. Namespace and project are taken from the parent key. """ def __init__(self, kind, parent_key): self._kind = kind self._parent_key = parent_key def make_entity(self, content): """Create entity from given string.""" key = Key([self._kind, hashlib.sha1(content.encode('utf-8')).hexdigest()], parent=self._parent_key) entity = Entity(key) entity.set_properties({'content': str(content)}) return entity def run(argv=None): """Main entry point.""" parser = argparse.ArgumentParser() parser.add_argument( '--kind', dest='kind', default='writereadtest', help='Datastore Kind') parser.add_argument( '--num_entities', dest='num_entities', type=int, required=True, help='Number of entities to write') parser.add_argument( '--limit', dest='limit', type=int, help='Limit of number of entities to write') known_args, pipeline_args = parser.parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) gcloud_options = pipeline_options.view_as(GoogleCloudOptions) job_name = gcloud_options.job_name kind = known_args.kind num_entities = known_args.num_entities project = gcloud_options.project # Pipeline 1: Create and write the specified number of Entities to the # Cloud Datastore. ancestor_key = Key([kind, str(uuid.uuid4())], project=project) _LOGGER.info('Writing %s entities to %s', num_entities, project) p = new_pipeline_with_job_name(pipeline_options, job_name, '-write') _ = ( p | 'Input' >> beam.Create(list(range(num_entities))) | 'To String' >> beam.Map(str) | 'To Entity' >> beam.Map(EntityWrapper(kind, ancestor_key).make_entity) | 'Write to Datastore' >> WriteToDatastore(project, hint_num_workers=1)) p.run() query = Query(kind=kind, project=project, ancestor=ancestor_key) # Optional Pipeline 2: If a read limit was provided, read it and confirm # that the expected entities were read. if known_args.limit is not None: _LOGGER.info( 'Querying a limited set of %s entities and verifying count.', known_args.limit) p = new_pipeline_with_job_name(pipeline_options, job_name, '-verify-limit') query.limit = known_args.limit entities = p | 'read from datastore' >> ReadFromDatastore(query) assert_that( entities | beam.combiners.Count.Globally(), equal_to([known_args.limit])) p.run() query.limit = None # Pipeline 3: Query the written Entities and verify result. _LOGGER.info('Querying entities, asserting they match.') p = new_pipeline_with_job_name(pipeline_options, job_name, '-verify') entities = p | 'read from datastore' >> ReadFromDatastore(query) assert_that( entities | beam.combiners.Count.Globally(), equal_to([num_entities])) p.run() # Pipeline 4: Delete Entities. _LOGGER.info('Deleting entities.') p = new_pipeline_with_job_name(pipeline_options, job_name, '-delete') entities = p | 'read from datastore' >> ReadFromDatastore(query) _ = ( entities | 'To Keys' >> beam.Map(lambda entity: entity.key) | 'delete entities' >> DeleteFromDatastore(project, hint_num_workers=1)) p.run() # Pipeline 5: Query the written Entities, verify no results. _LOGGER.info('Querying for the entities to make sure there are none present.') p = new_pipeline_with_job_name(pipeline_options, job_name, '-verify-deleted') entities = p | 'read from datastore' >> ReadFromDatastore(query) assert_that(entities | beam.combiners.Count.Globally(), equal_to([0])) p.run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/datastore_write_it_test.py0000644000175100017510000000501215206113032027366 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """An integration test for datastore_write_it_pipeline This test creates entities and writes them to Cloud Datastore. Subsequently, these entities are read from Cloud Datastore, compared to the expected value for the entity, and deleted. There is no output; instead, we use `assert_that` transform to verify the results in the pipeline. """ # pytype: skip-file import logging import random import unittest from datetime import datetime import pytest from hamcrest.core.core.allof import all_of from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline try: from apache_beam.io.gcp.datastore.v1new import datastore_write_it_pipeline except ImportError: datastore_write_it_pipeline = None # type: ignore class DatastoreWriteIT(unittest.TestCase): NUM_ENTITIES = 1001 LIMIT = 500 def run_datastore_write(self, limit=None): test_pipeline = TestPipeline(is_integration_test=True) current_time = datetime.now().strftime("%m%d%H%M%S") seed = random.randint(0, 100000) kind = 'testkind%s%d' % (current_time, seed) pipeline_verifiers = [PipelineStateMatcher()] extra_opts = { 'kind': kind, 'num_entities': self.NUM_ENTITIES, 'on_success_matcher': all_of(*pipeline_verifiers) } if limit is not None: extra_opts['limit'] = limit datastore_write_it_pipeline.run( test_pipeline.get_full_options_as_args(**extra_opts)) @pytest.mark.it_postcommit @unittest.skipIf( datastore_write_it_pipeline is None, 'GCP dependencies are not installed') def test_datastore_write_limit(self): self.run_datastore_write(limit=self.LIMIT) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/datastoreio.py0000644000175100017510000005606615206113032024770 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A connector for reading from and writing to Google Cloud Datastore. This module uses the newer google-cloud-datastore client package. Its API was different enough to require extensive changes to this and associated modules. **Updates to the I/O connector code** For any significant updates to this I/O connector, please consider involving corresponding code reviewers mentioned in https://github.com/apache/beam/blob/master/sdks/python/OWNERS """ # pytype: skip-file import logging import time from apache_beam import typehints from apache_beam.internal.metrics.metric import ServiceCallMetric from apache_beam.io.components.adaptive_throttler import AdaptiveThrottler from apache_beam.io.gcp import resource_identifiers from apache_beam.io.gcp.datastore.v1new import helper from apache_beam.io.gcp.datastore.v1new import query_splitter from apache_beam.io.gcp.datastore.v1new import types from apache_beam.io.gcp.datastore.v1new import util from apache_beam.io.gcp.datastore.v1new.rampup_throttling_fn import RampupThrottlingFn from apache_beam.metrics import monitoring_infos from apache_beam.metrics.metric import Metrics from apache_beam.transforms import Create from apache_beam.transforms import DoFn from apache_beam.transforms import ParDo from apache_beam.transforms import PTransform from apache_beam.transforms import Reshuffle from apache_beam.utils import retry # Protect against environments where datastore library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from google.api_core.exceptions import ClientError from google.api_core.exceptions import GoogleAPICallError except ImportError: pass __all__ = ['ReadFromDatastore', 'WriteToDatastore', 'DeleteFromDatastore'] _LOGGER = logging.getLogger(__name__) @typehints.with_output_types(types.Entity) class ReadFromDatastore(PTransform): """A ``PTransform`` for querying Google Cloud Datastore. To read a ``PCollection[Entity]`` from a Cloud Datastore ``Query``, use the ``ReadFromDatastore`` transform by providing a `query` to read from. The project and optional namespace are set in the query. The query will be split into multiple queries to allow for parallelism. The degree of parallelism is automatically determined, but can be overridden by setting `num_splits` to a value of 1 or greater. Note: Normally, a runner will read from Cloud Datastore in parallel across many workers. However, when the `query` is configured with a `limit` or if the query contains inequality filters like `GREATER_THAN, LESS_THAN` etc., then all the returned results will be read by a single worker in order to ensure correct data. Since data is read from a single worker, this could have significant impact on the performance of the job. Using a :class:`~apache_beam.transforms.util.Reshuffle` transform after the read in this case might be beneficial for parallelizing work across workers. The semantics for query splitting is defined below: 1. If `num_splits` is equal to 0, then the number of splits will be chosen dynamically at runtime based on the query data size. 2. Any value of `num_splits` greater than `ReadFromDatastore._NUM_QUERY_SPLITS_MAX` will be capped at that value. 3. If the `query` has a user limit set, or contains inequality filters, then `num_splits` will be ignored and no split will be performed. 4. Under certain cases Cloud Datastore is unable to split query to the requested number of splits. In such cases we just use whatever Cloud Datastore returns. See https://developers.google.com/datastore/ for more details on Google Cloud Datastore. """ # An upper bound on the number of splits for a query. _NUM_QUERY_SPLITS_MAX = 50000 # A lower bound on the number of splits for a query. This is to ensure that # we parallelize the query even when Datastore statistics are not available. _NUM_QUERY_SPLITS_MIN = 12 # Default bundle size of 64MB. _DEFAULT_BUNDLE_SIZE_BYTES = 64 * 1024 * 1024 def __init__(self, query, num_splits=0): """Initialize the `ReadFromDatastore` transform. This transform outputs elements of type :class:`~apache_beam.io.gcp.datastore.v1new.types.Entity`. Args: query: (:class:`~apache_beam.io.gcp.datastore.v1new.types.Query`) query used to fetch entities. num_splits: (:class:`int`) (optional) Number of splits for the query. """ super().__init__() if not query.project: raise ValueError("query.project cannot be empty") if not query: raise ValueError("query cannot be empty") if num_splits < 0: raise ValueError("num_splits must be greater than or equal 0") self._project = query.project # using _namespace conflicts with DisplayData._namespace self._datastore_namespace = query.namespace self._query = query self._num_splits = num_splits def expand(self, pcoll): # This is a composite transform involves the following: # 1. Create a singleton of the user provided `query` and apply a ``ParDo`` # that splits the query into `num_splits` queries if possible. # # If the value of `num_splits` is 0, the number of splits will be # computed dynamically based on the size of the data for the `query`. # # 2. The resulting ``PCollection`` is sharded across workers using a # ``Reshuffle`` operation. # # 3. In the third step, a ``ParDo`` reads entities for each query and # outputs a ``PCollection[Entity]``. return ( pcoll.pipeline | 'UserQuery' >> Create([self._query]) | 'SplitQuery' >> ParDo( ReadFromDatastore._SplitQueryFn(self._num_splits)) | Reshuffle() | 'Read' >> ParDo(ReadFromDatastore._QueryFn())) def display_data(self): disp_data = { 'project': self._query.project, 'query': str(self._query), 'num_splits': self._num_splits } if self._datastore_namespace is not None: disp_data['namespace'] = self._datastore_namespace return disp_data @typehints.with_input_types(types.Query) @typehints.with_output_types(types.Query) class _SplitQueryFn(DoFn): """A `DoFn` that splits a given query into multiple sub-queries.""" def __init__(self, num_splits): super().__init__() self._num_splits = num_splits def process(self, query, *args, **kwargs): client = helper.get_client(query.project, query.namespace) try: # Short circuit estimating num_splits if split is not possible. query_splitter.validate_split(query) if self._num_splits == 0: estimated_num_splits = self.get_estimated_num_splits(client, query) else: estimated_num_splits = self._num_splits _LOGGER.info("Splitting the query into %d splits", estimated_num_splits) query_splits = query_splitter.get_splits( client, query, estimated_num_splits) except query_splitter.QuerySplitterError: _LOGGER.info( "Unable to parallelize the given query: %s", query, exc_info=True) query_splits = [query] return query_splits def display_data(self): disp_data = {'num_splits': self._num_splits} return disp_data @staticmethod def query_latest_statistics_timestamp(client): """Fetches the latest timestamp of statistics from Cloud Datastore. Cloud Datastore system tables with statistics are periodically updated. This method fetches the latest timestamp (in microseconds) of statistics update using the `__Stat_Total__` table. """ if client.namespace is None: kind = '__Stat_Total__' else: kind = '__Stat_Ns_Total__' query = client.query( kind=kind, order=[ "-timestamp", ]) entities = list(query.fetch(limit=1)) if not entities: raise RuntimeError("Datastore total statistics unavailable.") return entities[0]['timestamp'] @staticmethod def get_estimated_size_bytes(client, query): """Get the estimated size of the data returned by this instance's query. Cloud Datastore provides no way to get a good estimate of how large the result of a query is going to be. Hence we use the __Stat_Kind__ system table to get size of the entire kind as an approximate estimate, assuming exactly 1 kind is specified in the query. See https://cloud.google.com/datastore/docs/concepts/stats. """ kind_name = query.kind latest_timestamp = ( ReadFromDatastore._SplitQueryFn.query_latest_statistics_timestamp( client)) _LOGGER.info( 'Latest stats timestamp for kind %s is %s', kind_name, latest_timestamp) if client.namespace is None: kind = '__Stat_Kind__' else: kind = '__Stat_Ns_Kind__' query = client.query(kind=kind) query.add_filter('kind_name', '=', kind_name) query.add_filter('timestamp', '=', latest_timestamp) entities = list(query.fetch(limit=1)) if not entities: raise RuntimeError( 'Datastore statistics for kind %s unavailable' % kind_name) return entities[0]['entity_bytes'] @staticmethod def get_estimated_num_splits(client, query): """Computes the number of splits to be performed on the query.""" try: estimated_size_bytes = ( ReadFromDatastore._SplitQueryFn.get_estimated_size_bytes( client, query)) _LOGGER.info('Estimated size bytes for query: %s', estimated_size_bytes) num_splits = int( min( ReadFromDatastore._NUM_QUERY_SPLITS_MAX, round(( float(estimated_size_bytes) / ReadFromDatastore._DEFAULT_BUNDLE_SIZE_BYTES)))) except Exception as e: _LOGGER.warning('Failed to fetch estimated size bytes: %s', e) # Fallback in case estimated size is unavailable. num_splits = ReadFromDatastore._NUM_QUERY_SPLITS_MIN return max(num_splits, ReadFromDatastore._NUM_QUERY_SPLITS_MIN) @typehints.with_input_types(types.Query) @typehints.with_output_types(types.Entity) class _QueryFn(DoFn): """A DoFn that fetches entities from Cloud Datastore, for a given query.""" def process(self, query, *unused_args, **unused_kwargs): if query.namespace is None: query.namespace = '' _client = helper.get_client(query.project, query.namespace) client_query = query._to_client_query(_client) # Create request count metric resource = resource_identifiers.DatastoreNamespace( query.project, query.namespace) labels = { monitoring_infos.SERVICE_LABEL: 'Datastore', monitoring_infos.METHOD_LABEL: 'BatchDatastoreRead', monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.DATASTORE_NAMESPACE_LABEL: query.namespace, monitoring_infos.DATASTORE_PROJECT_ID_LABEL: query.project, monitoring_infos.STATUS_LABEL: 'ok' } service_call_metric = ServiceCallMetric( request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN, base_labels=labels) try: for client_entity in client_query.fetch(query.limit): yield types.Entity.from_client_entity(client_entity) service_call_metric.call('ok') except (ClientError, GoogleAPICallError) as e: # e.code.value contains the numeric http status code. service_call_metric.call(e.code.value) raise class _Mutate(PTransform): """A ``PTransform`` that writes mutations to Cloud Datastore. Only idempotent Datastore mutation operations (upsert and delete) are supported, as the commits are retried when failures occur. """ # Default hint for the expected number of workers in the ramp-up throttling # step for write or delete operations. _DEFAULT_HINT_NUM_WORKERS = 500 def __init__( self, mutate_fn, throttle_rampup=True, hint_num_workers=_DEFAULT_HINT_NUM_WORKERS): """Initializes a Mutate transform. Args: mutate_fn: Instance of `DatastoreMutateFn` to use. throttle_rampup: Whether to enforce a gradual ramp-up. hint_num_workers: A hint for the expected number of workers, used to estimate appropriate limits during ramp-up throttling. """ self._mutate_fn = mutate_fn self._throttle_rampup = throttle_rampup self._hint_num_workers = hint_num_workers def expand(self, pcoll): if self._throttle_rampup: throttling_fn = RampupThrottlingFn(self._hint_num_workers) pcoll = ( pcoll | 'Enforce throttling during ramp-up' >> ParDo(throttling_fn)) return pcoll | 'Write Batch to Datastore' >> ParDo(self._mutate_fn) class DatastoreMutateFn(DoFn): """A ``DoFn`` that write mutations to Datastore. Mutations are written in batches, where the maximum batch size is `util.WRITE_BATCH_SIZE`. Commits are non-transactional. If a commit fails because of a conflict over an entity group, the commit will be retried. This means that the mutation should be idempotent (`upsert` and `delete` mutations) to prevent duplicate data or errors. """ def __init__(self, project): """ Args: project: (str) cloud project id """ self._project = project self._client = None self._rpc_successes = Metrics.counter( _Mutate.DatastoreMutateFn, "datastoreRpcSuccesses") self._rpc_errors = Metrics.counter( _Mutate.DatastoreMutateFn, "datastoreRpcErrors") self._throttled_secs = Metrics.counter( _Mutate.DatastoreMutateFn, "cumulativeThrottlingSeconds") self._throttler = AdaptiveThrottler( window_ms=120000, bucket_ms=1000, overload_ratio=1.25) def _update_rpc_stats(self, successes=0, errors=0, throttled_secs=0): self._rpc_successes.inc(successes) self._rpc_errors.inc(errors) self._throttled_secs.inc(throttled_secs) def start_bundle(self): self._client = helper.get_client(self._project, namespace=None) self._init_batch() self._batch_sizer = util.DynamicBatchSizer() self._target_batch_size = self._batch_sizer.get_batch_size( time.time() * 1000) def element_to_client_batch_item(self, element): raise NotImplementedError def add_to_batch(self, client_batch_item): raise NotImplementedError @retry.with_exponential_backoff( num_retries=5, retry_filter=helper.retry_on_rpc_error) def write_mutations(self, throttler, rpc_stats_callback, throttle_delay=1): """Writes a batch of mutations to Cloud Datastore. If a commit fails, it will be retried up to 5 times. All mutations in the batch will be committed again, even if the commit was partially successful. If the retry limit is exceeded, the last exception from Cloud Datastore will be raised. Assumes that the Datastore client library does not perform any retries on commits. It has not been determined how such retries would interact with the retries and throttler used here. See ``google.cloud.datastore_v1.gapic.datastore_client_config`` for retry config. Args: rpc_stats_callback: a function to call with arguments `successes` and `failures` and `throttled_secs`; this is called to record successful and failed RPCs to Datastore and time spent waiting for throttling. throttler: (``apache_beam.io.gcp.datastore.v1new.adaptive_throttler. AdaptiveThrottler``) Throttler instance used to select requests to be throttled. throttle_delay: (:class:`float`) time in seconds to sleep when throttled. Returns: (int) The latency of the successful RPC in milliseconds. """ # Client-side throttling. while throttler.throttle_request(time.time() * 1000): _LOGGER.info( "Delaying request for %ds due to previous failures", throttle_delay) time.sleep(throttle_delay) rpc_stats_callback(throttled_secs=throttle_delay) if self._batch is None: # this will only happen when we re-try previously failed batch self._batch = self._client.batch() self._batch.begin() for element in self._batch_elements: self.add_to_batch(element) # Create request count metric resource = resource_identifiers.DatastoreNamespace(self._project, "") labels = { monitoring_infos.SERVICE_LABEL: 'Datastore', monitoring_infos.METHOD_LABEL: 'BatchDatastoreWrite', monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.DATASTORE_NAMESPACE_LABEL: "", monitoring_infos.DATASTORE_PROJECT_ID_LABEL: self._project, monitoring_infos.STATUS_LABEL: 'ok' } service_call_metric = ServiceCallMetric( request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN, base_labels=labels) try: start_time = time.time() self._batch.commit() end_time = time.time() service_call_metric.call('ok') rpc_stats_callback(successes=1) throttler.successful_request(start_time * 1000) commit_time_ms = int((end_time - start_time) * 1000) return commit_time_ms except (ClientError, GoogleAPICallError) as e: self._batch = None # e.code.value contains the numeric http status code. service_call_metric.call(e.code.value) rpc_stats_callback(errors=1) raise def process(self, element): client_element = self.element_to_client_batch_item(element) self._batch_elements.append(client_element) self.add_to_batch(client_element) self._batch_bytes_size += self._batch.mutations[-1]._pb.ByteSize() if (len(self._batch.mutations) >= self._target_batch_size or self._batch_bytes_size > util.WRITE_BATCH_MAX_BYTES_SIZE): self._flush_batch() def finish_bundle(self): if self._batch_elements: self._flush_batch() def _init_batch(self): self._batch_bytes_size = 0 self._batch = self._client.batch() self._batch.begin() self._batch_elements = [] def _flush_batch(self): # Flush the current batch of mutations to Cloud Datastore. latency_ms = self.write_mutations( self._throttler, rpc_stats_callback=self._update_rpc_stats, throttle_delay=util.WRITE_BATCH_TARGET_LATENCY_MS // 1000) _LOGGER.debug( "Successfully wrote %d mutations in %dms.", len(self._batch.mutations), latency_ms) now = time.time() * 1000 self._batch_sizer.report_latency( now, latency_ms, len(self._batch.mutations)) self._target_batch_size = self._batch_sizer.get_batch_size(now) self._init_batch() @typehints.with_input_types(types.Entity) class WriteToDatastore(_Mutate): """ Writes elements of type :class:`~apache_beam.io.gcp.datastore.v1new.types.Entity` to Cloud Datastore. Entity keys must be complete. The ``project`` field in each key must match the project ID passed to this transform. If ``project`` field in entity or property key is empty then it is filled with the project ID passed to this transform. """ def __init__( self, project, throttle_rampup=True, hint_num_workers=_Mutate._DEFAULT_HINT_NUM_WORKERS): """Initialize the `WriteToDatastore` transform. Args: project: (:class:`str`) The ID of the project to write entities to. throttle_rampup: Whether to enforce a gradual ramp-up. hint_num_workers: A hint for the expected number of workers, used to estimate appropriate limits during ramp-up throttling. """ mutate_fn = WriteToDatastore._DatastoreWriteFn(project) super().__init__(mutate_fn, throttle_rampup, hint_num_workers) class _DatastoreWriteFn(_Mutate.DatastoreMutateFn): def element_to_client_batch_item(self, element): if not isinstance(element, types.Entity): raise ValueError( 'apache_beam.io.gcp.datastore.v1new.datastoreio.Entity' ' expected, got: %s' % type(element)) if not element.key.project: element.key.project = self._project client_entity = element.to_client_entity() if client_entity.key.is_partial: raise ValueError( 'Entities to be written to Cloud Datastore must ' 'have complete keys:\n%s' % client_entity) return client_entity def add_to_batch(self, client_entity): self._batch.put(client_entity) def display_data(self): return { 'mutation': 'Write (upsert)', 'project': self._project, } @typehints.with_input_types(types.Key) class DeleteFromDatastore(_Mutate): """ Deletes elements matching input :class:`~apache_beam.io.gcp.datastore.v1new.types.Key` elements from Cloud Datastore. Keys must be complete. The ``project`` field in each key must match the project ID passed to this transform. If ``project`` field in key is empty then it is filled with the project ID passed to this transform. """ def __init__( self, project, throttle_rampup=True, hint_num_workers=_Mutate._DEFAULT_HINT_NUM_WORKERS): """Initialize the `DeleteFromDatastore` transform. Args: project: (:class:`str`) The ID of the project from which the entities will be deleted. throttle_rampup: Whether to enforce a gradual ramp-up. hint_num_workers: A hint for the expected number of workers, used to estimate appropriate limits during ramp-up throttling. """ mutate_fn = DeleteFromDatastore._DatastoreDeleteFn(project) super().__init__(mutate_fn, throttle_rampup, hint_num_workers) class _DatastoreDeleteFn(_Mutate.DatastoreMutateFn): def element_to_client_batch_item(self, element): if not isinstance(element, types.Key): raise ValueError( 'apache_beam.io.gcp.datastore.v1new.datastoreio.Key' ' expected, got: %s' % type(element)) if not element.project: element.project = self._project client_key = element.to_client_key() if client_key.is_partial: raise ValueError( 'Keys to be deleted from Cloud Datastore must be ' 'complete:\n%s' % client_key) return client_key def add_to_batch(self, client_key): self._batch.delete(client_key) def display_data(self): return { 'mutation': 'Delete', 'project': self._project, } ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/datastoreio_test.py0000644000175100017510000004707215206113032026024 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for datastoreio.""" # pytype: skip-file import datetime import math import unittest from mock import MagicMock from mock import call from mock import patch from mock import ANY # Protect against environments where datastore library is not available. try: from apache_beam.io.gcp import resource_identifiers from apache_beam.io.gcp.datastore.v1new import helper, util from apache_beam.io.gcp.datastore.v1new import query_splitter from apache_beam.io.gcp.datastore.v1new import datastoreio from apache_beam.io.gcp.datastore.v1new.datastoreio import DeleteFromDatastore from apache_beam.io.gcp.datastore.v1new.datastoreio import ReadFromDatastore from apache_beam.io.gcp.datastore.v1new.datastoreio import WriteToDatastore from apache_beam.io.gcp.datastore.v1new.types import Key from apache_beam.metrics import monitoring_infos from apache_beam.metrics.execution import MetricsEnvironment from google.cloud.datastore import client from google.cloud.datastore import entity from google.cloud.datastore import helpers from google.cloud.datastore import key from google.api_core import exceptions except ImportError: client = None # used for internal testing only class FakeMessage: def __init__(self, entity, key): self.entity = entity self.key = key def ByteSize(self): if self.entity is not None: return helpers.entity_to_protobuf(self.entity)._pb.ByteSize() else: return self.key.to_protobuf()._pb.ByteSize() # used for internal testing only class FakeMutation(object): def __init__(self, entity=None, key=None): """Fake mutation request object. Requires exactly one of entity or key to be set. Args: entity: (``google.cloud.datastore.entity.Entity``) entity representing this upsert mutation key: (``google.cloud.datastore.key.Key``) key representing this delete mutation """ self.entity = entity self.key = key self._pb = FakeMessage(entity, key) class FakeBatch(object): def __init__(self, all_batch_items=None, commit_count=None): """Fake ``google.cloud.datastore.batch.Batch`` object. Args: all_batch_items: (list) If set, will append all entities/keys added to this batch. commit_count: (list of int) If set, will increment commit_count[0] on each ``commit``. """ self._all_batch_items = all_batch_items self._commit_count = commit_count self.mutations = [] def put(self, _entity): assert isinstance(_entity, entity.Entity) self.mutations.append(FakeMutation(entity=_entity)) if self._all_batch_items is not None: self._all_batch_items.append(_entity) def delete(self, _key): assert isinstance(_key, key.Key) self.mutations.append(FakeMutation(key=_key)) if self._all_batch_items is not None: self._all_batch_items.append(_key) def begin(self): pass def commit(self): if self._commit_count: self._commit_count[0] += 1 @unittest.skipIf(client is None, 'Datastore dependencies are not installed') class MutateTest(unittest.TestCase): def test_write_mutations_no_errors(self): mock_batch = MagicMock() mock_throttler = MagicMock() rpc_stats_callback = MagicMock() mock_throttler.throttle_request.return_value = [] mutate = datastoreio._Mutate.DatastoreMutateFn("") mutate._batch = mock_batch mutate.write_mutations(mock_throttler, rpc_stats_callback) rpc_stats_callback.assert_has_calls([ call(successes=1), ]) @patch('time.sleep', return_value=None) def test_write_mutations_reconstruct_on_error(self, unused_sleep): mock_batch = MagicMock() mock_batch.begin.side_effect = [None, ValueError] mock_batch.commit.side_effect = [ exceptions.DeadlineExceeded('retryable'), None ] mock_throttler = MagicMock() rpc_stats_callback = MagicMock() mock_throttler.throttle_request.return_value = [] mutate = datastoreio._Mutate.DatastoreMutateFn("") mutate._batch = mock_batch mutate._client = MagicMock() mutate._batch_elements = [None] mock_add_to_batch = MagicMock() mutate.add_to_batch = mock_add_to_batch mutate.write_mutations(mock_throttler, rpc_stats_callback) rpc_stats_callback.assert_has_calls([ call(successes=1), ]) self.assertEqual(1, mock_add_to_batch.call_count) @patch('time.sleep', return_value=None) def test_write_mutations_throttle_delay_retryable_error(self, unused_sleep): mock_batch = MagicMock() mock_batch.commit.side_effect = [ exceptions.DeadlineExceeded('retryable'), None ] mock_throttler = MagicMock() rpc_stats_callback = MagicMock() # First try: throttle once [True, False] # Second try: no throttle [False] mock_throttler.throttle_request.side_effect = [True, False, False] mutate = datastoreio._Mutate.DatastoreMutateFn("") mutate._batch = mock_batch mutate._batch_elements = [] mutate._client = MagicMock() mutate.write_mutations(mock_throttler, rpc_stats_callback, throttle_delay=0) rpc_stats_callback.assert_has_calls([ call(successes=1), call(throttled_secs=ANY), call(errors=1), ], any_order=True) self.assertEqual(3, rpc_stats_callback.call_count) def test_write_mutations_non_retryable_error(self): mock_batch = MagicMock() mock_batch.commit.side_effect = [ exceptions.InvalidArgument('non-retryable'), ] mock_throttler = MagicMock() rpc_stats_callback = MagicMock() mock_throttler.throttle_request.return_value = False mutate = datastoreio._Mutate.DatastoreMutateFn("") mutate._batch = mock_batch with self.assertRaises(exceptions.InvalidArgument): mutate.write_mutations( mock_throttler, rpc_stats_callback, throttle_delay=0) rpc_stats_callback.assert_called_once_with(errors=1) def test_write_mutations_metric_on_failure(self): MetricsEnvironment.process_wide_container().reset() mock_batch = MagicMock() mock_batch.commit.side_effect = [ exceptions.DeadlineExceeded("Deadline Exceeded"), [] ] mock_throttler = MagicMock() rpc_stats_callback = MagicMock() mock_throttler.throttle_request.return_value = False mutate = datastoreio._Mutate.DatastoreMutateFn("my_project") mutate._batch = mock_batch mutate._batch_elements = [] mutate._client = MagicMock() mutate.write_mutations(mock_throttler, rpc_stats_callback, throttle_delay=0) self.verify_write_call_metric("my_project", "", "deadline_exceeded", 1) self.verify_write_call_metric("my_project", "", "ok", 1) def verify_write_call_metric(self, project_id, namespace, status, count): """Check if a metric was recorded for the Datastore IO write API call.""" process_wide_monitoring_infos = list( MetricsEnvironment.process_wide_container(). to_runner_api_monitoring_infos(None).values()) resource = resource_identifiers.DatastoreNamespace(project_id, namespace) labels = { monitoring_infos.SERVICE_LABEL: 'Datastore', monitoring_infos.METHOD_LABEL: 'BatchDatastoreWrite', monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.DATASTORE_NAMESPACE_LABEL: namespace, monitoring_infos.DATASTORE_PROJECT_ID_LABEL: project_id, monitoring_infos.STATUS_LABEL: status } expected_mi = monitoring_infos.int64_counter( monitoring_infos.API_REQUEST_COUNT_URN, count, labels=labels) expected_mi.ClearField("start_time") found = False for actual_mi in process_wide_monitoring_infos: actual_mi.ClearField("start_time") if expected_mi == actual_mi: found = True break self.assertTrue( found, "Did not find write call metric with status: %s" % status) @unittest.skipIf(client is None, 'Datastore dependencies are not installed') class DatastoreioTest(unittest.TestCase): _PROJECT = 'project' _KIND = 'kind' _NAMESPACE = 'namespace' def setUp(self): self._WRITE_BATCH_INITIAL_SIZE = util.WRITE_BATCH_INITIAL_SIZE self._mock_client = MagicMock() self._mock_client.project = self._PROJECT self._mock_client.namespace = self._NAMESPACE self._mock_query = MagicMock() self._mock_query.limit = None self._mock_query.order = None self._real_client = client.Client( project=self._PROJECT, namespace=self._NAMESPACE, # Don't do any network requests. _http=MagicMock()) def test_SplitQueryFn_with_num_splits(self): with patch.object(helper, 'get_client', return_value=self._mock_client): num_splits = 23 expected_num_splits = 23 def fake_get_splits(unused_client, query, num_splits): return [query] * num_splits with patch.object(query_splitter, 'get_splits', side_effect=fake_get_splits): split_query_fn = ReadFromDatastore._SplitQueryFn(num_splits) split_queries = split_query_fn.process(self._mock_query) self.assertEqual(expected_num_splits, len(split_queries)) def test_SplitQueryFn_without_num_splits(self): with patch.object(helper, 'get_client', return_value=self._mock_client): # Force _SplitQueryFn to compute the number of query splits num_splits = 0 expected_num_splits = 23 entity_bytes = ( expected_num_splits * ReadFromDatastore._DEFAULT_BUNDLE_SIZE_BYTES) with patch.object(ReadFromDatastore._SplitQueryFn, 'get_estimated_size_bytes', return_value=entity_bytes): def fake_get_splits(unused_client, query, num_splits): return [query] * num_splits with patch.object(query_splitter, 'get_splits', side_effect=fake_get_splits): split_query_fn = ReadFromDatastore._SplitQueryFn(num_splits) split_queries = split_query_fn.process(self._mock_query) self.assertEqual(expected_num_splits, len(split_queries)) def test_SplitQueryFn_with_query_limit(self): """A test that verifies no split is performed when the query has a limit.""" with patch.object(helper, 'get_client', return_value=self._mock_client): num_splits = 4 expected_num_splits = 1 self._mock_query.limit = 3 split_query_fn = ReadFromDatastore._SplitQueryFn(num_splits) split_queries = split_query_fn.process(self._mock_query) self.assertEqual(expected_num_splits, len(split_queries)) def test_SplitQueryFn_with_exception(self): """A test that verifies that no split is performed when failures occur.""" with patch.object(helper, 'get_client', return_value=self._mock_client): # Force _SplitQueryFn to compute the number of query splits num_splits = 0 expected_num_splits = 1 entity_bytes = ( expected_num_splits * ReadFromDatastore._DEFAULT_BUNDLE_SIZE_BYTES) with patch.object(ReadFromDatastore._SplitQueryFn, 'get_estimated_size_bytes', return_value=entity_bytes): with patch.object(query_splitter, 'get_splits', side_effect=query_splitter.QuerySplitterError( "Testing query split error")): split_query_fn = ReadFromDatastore._SplitQueryFn(num_splits) split_queries = split_query_fn.process(self._mock_query) self.assertEqual(expected_num_splits, len(split_queries)) self.assertEqual(self._mock_query, split_queries[0]) def test_QueryFn_metric_on_failure(self): MetricsEnvironment.process_wide_container().reset() with patch.object(helper, 'get_client', return_value=self._mock_client): self._mock_query.project = self._PROJECT self._mock_query.namespace = self._NAMESPACE _query_fn = ReadFromDatastore._QueryFn() client_query = self._mock_query._to_client_query() # Test with exception client_query.fetch.side_effect = [ exceptions.DeadlineExceeded("Deadline exceed") ] try: list(_query_fn.process(self._mock_query)) except Exception: self.verify_read_call_metric( self._PROJECT, self._NAMESPACE, "deadline_exceeded", 1) # Test success client_query.fetch.side_effect = [[]] list(_query_fn.process(self._mock_query)) self.verify_read_call_metric(self._PROJECT, self._NAMESPACE, "ok", 1) else: raise Exception('Excepted _query_fn.process call to raise an error') def verify_read_call_metric(self, project_id, namespace, status, count): """Check if a metric was recorded for the Datastore IO read API call.""" process_wide_monitoring_infos = list( MetricsEnvironment.process_wide_container(). to_runner_api_monitoring_infos(None).values()) resource = resource_identifiers.DatastoreNamespace(project_id, namespace) labels = { monitoring_infos.SERVICE_LABEL: 'Datastore', monitoring_infos.METHOD_LABEL: 'BatchDatastoreRead', monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.DATASTORE_NAMESPACE_LABEL: namespace, monitoring_infos.DATASTORE_PROJECT_ID_LABEL: project_id, monitoring_infos.STATUS_LABEL: status } expected_mi = monitoring_infos.int64_counter( monitoring_infos.API_REQUEST_COUNT_URN, count, labels=labels) expected_mi.ClearField("start_time") found = False for actual_mi in process_wide_monitoring_infos: actual_mi.ClearField("start_time") if expected_mi == actual_mi: found = True break self.assertTrue( found, "Did not find read call metric with status: %s" % status) def check_DatastoreWriteFn(self, num_entities): """A helper function to test _DatastoreWriteFn.""" with patch.object(helper, 'get_client', return_value=self._mock_client): entities = helper.create_entities(num_entities) expected_entities = [entity.to_client_entity() for entity in entities] # Infer project from write fn project arg. if num_entities: key = Key(['k1', 1234], project=self._PROJECT) expected_key = key.to_client_key() key.project = None entities[0].key = key expected_entities[0].key = expected_key all_batch_entities = [] commit_count = [0] self._mock_client.batch.side_effect = ( lambda: FakeBatch( all_batch_items=all_batch_entities, commit_count=commit_count)) datastore_write_fn = WriteToDatastore._DatastoreWriteFn(self._PROJECT) datastore_write_fn.start_bundle() for entity in entities: datastore_write_fn.process(entity) datastore_write_fn.finish_bundle() self.assertListEqual([e.key for e in all_batch_entities], [e.key for e in expected_entities]) batch_count = math.ceil(num_entities / util.WRITE_BATCH_MAX_SIZE) self.assertLessEqual(batch_count, commit_count[0]) def test_DatastoreWriteFn_with_empty_batch(self): self.check_DatastoreWriteFn(0) def test_DatastoreWriteFn_with_one_batch(self): num_entities_to_write = self._WRITE_BATCH_INITIAL_SIZE * 1 - 50 self.check_DatastoreWriteFn(num_entities_to_write) def test_DatastoreWriteFn_with_multiple_batches(self): num_entities_to_write = self._WRITE_BATCH_INITIAL_SIZE * 3 + 50 self.check_DatastoreWriteFn(num_entities_to_write) def test_DatastoreWriteFn_with_batch_size_exact_multiple(self): num_entities_to_write = self._WRITE_BATCH_INITIAL_SIZE * 2 self.check_DatastoreWriteFn(num_entities_to_write) def test_DatastoreWriteFn_with_dynamic_batch_sizes(self): num_entities_to_write = self._WRITE_BATCH_INITIAL_SIZE * 3 + 50 self.check_DatastoreWriteFn(num_entities_to_write) def test_DatastoreWriteLargeEntities(self): """100*100kB entities gets split over two Commit RPCs.""" with patch.object(helper, 'get_client', return_value=self._mock_client): entities = helper.create_entities(100) commit_count = [0] self._mock_client.batch.side_effect = ( lambda: FakeBatch(commit_count=commit_count)) datastore_write_fn = WriteToDatastore._DatastoreWriteFn(self._PROJECT) datastore_write_fn.start_bundle() for entity in entities: entity.set_properties({'large': 'A' * 100000}) datastore_write_fn.process(entity) datastore_write_fn.finish_bundle() self.assertEqual(2, commit_count[0]) def check_estimated_size_bytes(self, entity_bytes, timestamp, namespace=None): """A helper method to test get_estimated_size_bytes""" self._mock_client.namespace = namespace self._mock_client.query.return_value = self._mock_query self._mock_query.project = self._PROJECT self._mock_query.namespace = namespace self._mock_query.fetch.side_effect = [ [{ 'timestamp': timestamp }], [{ 'entity_bytes': entity_bytes }], ] self._mock_query.kind = self._KIND split_query_fn = ReadFromDatastore._SplitQueryFn(num_splits=0) self.assertEqual( entity_bytes, split_query_fn.get_estimated_size_bytes( self._mock_client, self._mock_query)) if namespace is None: ns_keyword = '_' else: ns_keyword = '_Ns_' self._mock_client.query.assert_has_calls([ call(kind='__Stat%sTotal__' % ns_keyword, order=['-timestamp']), call().fetch(limit=1), call(kind='__Stat%sKind__' % ns_keyword), call().add_filter('kind_name', '=', self._KIND), call().add_filter('timestamp', '=', timestamp), call().fetch(limit=1), ]) def get_timestamp(self): return datetime.datetime(2019, 3, 14, 15, 9, 26, 535897) def test_get_estimated_size_bytes_without_namespace(self): entity_bytes = 100 timestamp = self.get_timestamp() self.check_estimated_size_bytes(entity_bytes, timestamp) def test_get_estimated_size_bytes_with_namespace(self): entity_bytes = 100 timestamp = self.get_timestamp() self.check_estimated_size_bytes(entity_bytes, timestamp, self._NAMESPACE) def test_DatastoreDeleteFn(self): with patch.object(helper, 'get_client', return_value=self._mock_client): keys = [entity.key for entity in helper.create_entities(10)] expected_keys = [key.to_client_key() for key in keys] # Infer project from delete fn project arg. key = Key(['k1', 1234], project=self._PROJECT) expected_key = key.to_client_key() key.project = None keys.append(key) expected_keys.append(expected_key) all_batch_keys = [] self._mock_client.batch.side_effect = ( lambda: FakeBatch(all_batch_items=all_batch_keys)) datastore_delete_fn = DeleteFromDatastore._DatastoreDeleteFn( self._PROJECT) datastore_delete_fn.start_bundle() for key in keys: datastore_delete_fn.process(key) datastore_delete_fn.finish_bundle() self.assertListEqual(all_batch_keys, expected_keys) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/helper.py0000644000175100017510000000562415206113032023723 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Cloud Datastore client and test functions. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file import os import uuid from typing import Union from cachetools.func import ttl_cache from google.api_core import exceptions from google.api_core.gapic_v1 import client_info from google.cloud import environment_vars from google.cloud.datastore import __version__ from google.cloud.datastore import client from apache_beam.io.gcp.datastore.v1new import types from apache_beam.version import __version__ as beam_version # https://cloud.google.com/datastore/docs/concepts/errors#error_codes _RETRYABLE_DATASTORE_ERRORS = ( exceptions.Aborted, exceptions.DeadlineExceeded, exceptions.InternalServerError, exceptions.ServiceUnavailable, ) @ttl_cache(maxsize=128, ttl=3600) def get_client(project, namespace): """Returns a Cloud Datastore client.""" _client_info = client_info.ClientInfo( client_library_version=__version__, gapic_version=__version__, user_agent=f'beam-python-sdk/{beam_version}') _client = client.Client( project=project, namespace=namespace, client_info=_client_info) # Avoid overwriting user setting. BEAM-7608 if not os.environ.get(environment_vars.GCD_HOST, None): _client.base_url = 'https://batch-datastore.googleapis.com' # BEAM-1387 return _client def retry_on_rpc_error(exception): """A retry filter for Cloud Datastore RPCErrors.""" return isinstance(exception, _RETRYABLE_DATASTORE_ERRORS) def create_entities(count, id_or_name=False): """Creates a list of entities with random keys.""" if id_or_name: ids_or_names: list[Union[str, int]] = [ uuid.uuid4().int & ((1 << 63) - 1) for _ in range(count) ] else: ids_or_names = [str(uuid.uuid4()) for _ in range(count)] keys = [types.Key(['EntityKind', x], project='project') for x in ids_or_names] return [types.Entity(key) for key in keys] def create_client_entities(count, id_or_name=False): """Creates a list of client-style entities with random keys.""" return [ entity.to_client_entity() for entity in create_entities(count, id_or_name=id_or_name) ] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/query_splitter.py0000644000175100017510000002260215206113032025532 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Implements a Cloud Datastore query splitter. For internal use only. No backwards compatibility guarantees. """ # pytype: skip-file from apache_beam.io.gcp.datastore.v1new import types from apache_beam.options.value_provider import ValueProvider __all__ = ['QuerySplitterError', 'SplitNotPossibleError', 'get_splits'] SCATTER_PROPERTY_NAME = '__scatter__' KEY_PROPERTY_NAME = '__key__' # The number of keys to sample for each split. KEYS_PER_SPLIT = 32 class QuerySplitterError(Exception): """Top-level error type.""" class SplitNotPossibleError(QuerySplitterError): """Raised when some parameter of the query does not allow splitting.""" def get_splits(client, query, num_splits): """Returns a list of sharded queries for the given Cloud Datastore query. This will create up to the desired number of splits, however it may return less splits if the desired number of splits is unavailable. This will happen if the number of split points provided by the underlying Datastore is less than the desired number, which will occur if the number of results for the query is too small. This implementation of the QuerySplitter uses the __scatter__ property to gather random split points for a query. Note: This implementation is derived from the java query splitter in https://github.com/GoogleCloudPlatform/google-cloud-datastore/blob/master/java/datastore/src/main/java/com/google/datastore/v1/client/QuerySplitterImpl.java Args: client: the datastore client. query: the query to split. num_splits: the desired number of splits. Returns: A list of split queries, of a max length of `num_splits` Raises: QuerySplitterError: if split could not be performed owing to query or split parameters. """ if num_splits <= 1: raise SplitNotPossibleError('num_splits must be > 1, got: %d' % num_splits) validate_split(query) splits = [] client_scatter_keys = _get_scatter_keys(client, query, num_splits) last_client_key = None for next_client_key in _get_split_key(client_scatter_keys, num_splits): splits.append(_create_split(last_client_key, next_client_key, query)) last_client_key = next_client_key splits.append(_create_split(last_client_key, None, query)) return splits def validate_split(query): """ Verifies that the given query can be properly scattered. Note that equality and ancestor filters are allowed, however they may result in inefficient sharding. Raises: QuerySplitterError if split could not be performed owing to query parameters. """ if query.order: raise SplitNotPossibleError('Query cannot have any sort orders.') if query.limit is not None: raise SplitNotPossibleError('Query cannot have a limit set.') for filter in query.filters: if isinstance(filter[1], ValueProvider): filter_operator = filter[1].get() else: filter_operator = filter[1] if filter_operator in ['<', '<=', '>', '>=']: raise SplitNotPossibleError('Query cannot have any inequality filters.') def _create_scatter_query(query, num_splits): """Creates a scatter query from the given user query.""" # There is a split containing entities before and after each scatter entity: # ||---*------*------*------*------*------*------*---|| * = scatter entity # If we represent each split as a region before a scatter entity, there is an # extra region following the last scatter point. Thus, we do not need the # scatter entity for the last region. limit = (num_splits - 1) * KEYS_PER_SPLIT scatter_query = types.Query( kind=query.kind, project=query.project, namespace=query.namespace, order=[SCATTER_PROPERTY_NAME], projection=[KEY_PROPERTY_NAME], limit=limit) return scatter_query class IdOrName(object): """Represents an ID or name of a Datastore key, Implements sort ordering: by ID, then by name, keys with IDs before those with names. """ def __init__(self, id_or_name): self.id_or_name = id_or_name if isinstance(id_or_name, str): self.id = None self.name = id_or_name elif isinstance(id_or_name, int): self.id = id_or_name self.name = None else: raise TypeError('Unexpected type of id_or_name: %s' % id_or_name) def __lt__(self, other): if not isinstance(other, IdOrName): return super().__lt__(other) if self.id is not None: if other.id is None: return True else: return self.id < other.id if other.id is not None: return False return self.name < other.name def __eq__(self, other): if not isinstance(other, IdOrName): return super().__eq__(other) return self.id == other.id and self.name == other.name def __hash__(self): return hash((self.id, self.other)) def client_key_sort_key(client_key): """Key function for sorting lists of ``google.cloud.datastore.key.Key``.""" sort_key = [client_key.project, client_key.namespace or ''] # A key path is made up of (kind, id_or_name) pairs. The last pair might be # missing an id_or_name. flat_path = list(client_key.flat_path) while flat_path: sort_key.append(flat_path.pop(0)) # kind if flat_path: sort_key.append(IdOrName(flat_path.pop(0))) return sort_key def _get_scatter_keys(client, query, num_splits): """Gets a list of split keys given a desired number of splits. This list will contain multiple split keys for each split. Only a single split key will be chosen as the split point, however providing multiple keys allows for more uniform sharding. Args: client: the client to datastore containing the data. query: the user query. num_splits: the number of desired splits. Returns: A list of scatter keys returned by Datastore. """ scatter_point_query = _create_scatter_query(query, num_splits) client_query = scatter_point_query._to_client_query(client) client_key_splits = [ client_entity.key for client_entity in client_query.fetch( client=client, limit=scatter_point_query.limit) ] client_key_splits.sort(key=client_key_sort_key) return client_key_splits def _get_split_key(client_keys, num_splits): """Given a list of keys and a number of splits find the keys to split on. Args: client_keys: the list of keys. num_splits: the number of splits. Returns: A list of keys to split on. """ # If the number of keys is less than the number of splits, we are limited # in the number of splits we can make. if not client_keys or (len(client_keys) < (num_splits - 1)): return client_keys # Calculate the number of keys per split. This should be KEYS_PER_SPLIT, # but may be less if there are not KEYS_PER_SPLIT * (numSplits - 1) scatter # entities. # # Consider the following dataset, where - represents an entity and # * represents an entity that is returned as a scatter entity: # ||---*-----*----*-----*-----*------*----*----|| # If we want 4 splits in this data, the optimal split would look like: # ||---*-----*----*-----*-----*------*----*----|| # | | | # The scatter keys in the last region are not useful to us, so we never # request them: # ||---*-----*----*-----*-----*------*---------|| # | | | # With 6 scatter keys we want to set scatter points at indexes: 1, 3, 5. # # We keep this as a float so that any "fractional" keys per split get # distributed throughout the splits and don't make the last split # significantly larger than the rest. num_keys_per_split = max(1.0, float(len(client_keys)) / (num_splits - 1)) split_client_keys = [] # Grab the last sample for each split, otherwise the first split will be too # small. for i in range(1, num_splits): split_index = int(round(i * num_keys_per_split) - 1) split_client_keys.append(client_keys[split_index]) return split_client_keys def _create_split(last_client_key, next_client_key, query): """Create a new {@link Query} given the query and range. Args: last_client_key: the previous key. If null then assumed to be the beginning. next_client_key: the next key. If null then assumed to be the end. query: query to base the split query on. Returns: A split query with fetches entities in the range [last_key, next_client_key) """ if not (last_client_key or next_client_key): return query split_query = query.clone() # Copy filters and possible convert the default empty tuple to empty list. filters = list(split_query.filters) if last_client_key: filters.append((KEY_PROPERTY_NAME, '>=', last_client_key)) if next_client_key: filters.append((KEY_PROPERTY_NAME, '<', next_client_key)) split_query.filters = filters return split_query ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/query_splitter_test.py0000644000175100017510000002263015206113032026572 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Cloud Datastore query splitter test.""" # pytype: skip-file import unittest import mock # Protect against environments where datastore library is not available. try: from google.cloud.datastore import key from apache_beam.io.gcp.datastore.v1new import helper from apache_beam.io.gcp.datastore.v1new import query_splitter from apache_beam.io.gcp.datastore.v1new import types from apache_beam.io.gcp.datastore.v1new.query_splitter import SplitNotPossibleError except ImportError: query_splitter = None # type: ignore @unittest.skipIf(query_splitter is None, 'GCP dependencies are not installed') class QuerySplitterTest(unittest.TestCase): _PROJECT = 'project' _NAMESPACE = 'namespace' def create_query( self, kinds=(), order=False, limit=None, offset=None, inequality_filter=False): kind = None filters = [] if kinds: kind = kinds[0] if order: order = ['prop1'] if inequality_filter: filters = [('prop1', '>', 'value1')] return types.Query(kind=kind, filters=filters, order=order, limit=limit) def test_get_splits_query_with_order(self): query = self.create_query(kinds=['a'], order=True) with self.assertRaisesRegex(SplitNotPossibleError, r'sort orders'): query_splitter.get_splits(None, query, 3) def test_get_splits_query_with_unsupported_filter(self): query = self.create_query(kinds=['a'], inequality_filter=True) with self.assertRaisesRegex(SplitNotPossibleError, r'inequality filters'): query_splitter.get_splits(None, query, 2) def test_get_splits_query_with_limit(self): query = self.create_query(kinds=['a'], limit=10) with self.assertRaisesRegex(SplitNotPossibleError, r'limit set'): query_splitter.get_splits(None, query, 2) def test_get_splits_query_with_num_splits_of_one(self): query = self.create_query() with self.assertRaisesRegex(SplitNotPossibleError, r'num_splits'): query_splitter.get_splits(None, query, 1) def test_create_scatter_query(self): query = types.Query(kind='shakespeare-demo') num_splits = 10 scatter_query = query_splitter._create_scatter_query(query, num_splits) self.assertEqual(scatter_query.kind, query.kind) self.assertEqual( scatter_query.limit, (num_splits - 1) * query_splitter.KEYS_PER_SPLIT) self.assertEqual( scatter_query.order, [query_splitter.SCATTER_PROPERTY_NAME]) self.assertEqual( scatter_query.projection, [query_splitter.KEY_PROPERTY_NAME]) def check_get_splits(self, query, num_splits, num_entities): """A helper method to test the query_splitter get_splits method. Args: query: the query to be split num_splits: number of splits num_entities: number of scatter entities returned to the splitter. """ # Test for random long ids, string ids, and a mix of both. for id_or_name in [True, False, None]: if id_or_name is None: client_entities = helper.create_client_entities(num_entities, False) client_entities.extend( helper.create_client_entities(num_entities, True)) num_entities *= 2 else: client_entities = helper.create_client_entities( num_entities, id_or_name) mock_client = mock.MagicMock() mock_client_query = mock.MagicMock() mock_client_query.fetch.return_value = client_entities with mock.patch.object(types.Query, '_to_client_query', return_value=mock_client_query): split_queries = query_splitter.get_splits( mock_client, query, num_splits) mock_client_query.fetch.assert_called_once() # if request num_splits is greater than num_entities, the best it can # do is one entity per split. expected_num_splits = min(num_splits, num_entities + 1) self.assertEqual(len(split_queries), expected_num_splits) # Verify no gaps in key ranges. Filters should look like: # query1: (__key__ < key1) # query2: (__key__ >= key1), (__key__ < key2) # ... # queryN: (__key__ >=keyN-1) prev_client_key = None last_query_seen = False for split_query in split_queries: self.assertFalse(last_query_seen) lt_key = None gte_key = None for _filter in split_query.filters: self.assertEqual(query_splitter.KEY_PROPERTY_NAME, _filter[0]) if _filter[1] == '<': lt_key = _filter[2] elif _filter[1] == '>=': gte_key = _filter[2] # Case where the scatter query has no results. if lt_key is None and gte_key is None: self.assertEqual(1, len(split_queries)) break if prev_client_key is None: self.assertIsNone(gte_key) self.assertIsNotNone(lt_key) prev_client_key = lt_key else: self.assertEqual(prev_client_key, gte_key) prev_client_key = lt_key if lt_key is None: last_query_seen = True def test_get_splits_with_two_splits(self): query = self.create_query(kinds=['shakespeare-demo']) num_splits = 2 num_entities = 97 self.check_get_splits(query, num_splits, num_entities) def test_get_splits_with_multiple_splits(self): query = self.create_query(kinds=['shakespeare-demo']) num_splits = 4 num_entities = 369 self.check_get_splits(query, num_splits, num_entities) def test_get_splits_with_large_num_splits(self): query = self.create_query(kinds=['shakespeare-demo']) num_splits = 10 num_entities = 4 self.check_get_splits(query, num_splits, num_entities) def test_get_splits_with_small_num_entities(self): query = self.create_query(kinds=['shakespeare-demo']) num_splits = 4 num_entities = 50 self.check_get_splits(query, num_splits, num_entities) def test_get_splits_with_batch_size_exact_multiple(self): """Test get_splits when num scatter keys is a multiple of batch size.""" query = self.create_query(kinds=['shakespeare-demo']) num_splits = 4 num_entities = 400 self.check_get_splits(query, num_splits, num_entities) def test_get_splits_with_large_batch_size(self): """Test get_splits when all scatter keys are retured in a single req.""" query = self.create_query(kinds=['shakespeare-demo']) num_splits = 4 num_entities = 400 self.check_get_splits(query, num_splits, num_entities) def test_get_splits_with_num_splits_gt_entities(self): query = self.create_query(kinds=['shakespeare-demo']) num_splits = 10 num_entities = 4 self.check_get_splits(query, num_splits, num_entities) def test_id_or_name(self): id_ = query_splitter.IdOrName(1) self.assertEqual(1, id_.id) self.assertIsNone(id_.name) name = query_splitter.IdOrName('1') self.assertIsNone(name.id) self.assertEqual('1', name.name) self.assertEqual(query_splitter.IdOrName(1), query_splitter.IdOrName(1)) self.assertEqual(query_splitter.IdOrName('1'), query_splitter.IdOrName('1')) self.assertLess(query_splitter.IdOrName(2), query_splitter.IdOrName('1')) self.assertLess(query_splitter.IdOrName(1), query_splitter.IdOrName(2)) self.assertLess(query_splitter.IdOrName('1'), query_splitter.IdOrName('2')) def test_client_key_sort_key(self): k = key.Key('kind1', 1, project=self._PROJECT, namespace=self._NAMESPACE) k2 = key.Key('kind2', 'a', parent=k) k3 = key.Key('kind2', 'b', parent=k) k4 = key.Key('kind1', 'a', project=self._PROJECT, namespace=self._NAMESPACE) k5 = key.Key('kind1', 'a', project=self._PROJECT) keys = [k5, k, k4, k3, k2, k2, k] expected_sort = [k5, k, k, k2, k2, k3, k4] keys.sort(key=query_splitter.client_key_sort_key) self.assertEqual(expected_sort, keys) def test_client_key_sort_key_ids(self): k1 = key.Key('kind', 2, project=self._PROJECT) k2 = key.Key('kind', 1, project=self._PROJECT) keys = [k1, k2] expected_sort = [k2, k1] keys.sort(key=query_splitter.client_key_sort_key) self.assertEqual(expected_sort, keys) def test_client_key_sort_key_names(self): k1 = key.Key('kind', '2', project=self._PROJECT) k2 = key.Key('kind', '1', project=self._PROJECT) keys = [k1, k2] expected_sort = [k2, k1] keys.sort(key=query_splitter.client_key_sort_key) self.assertEqual(expected_sort, keys) def test_client_key_sort_key_ids_vs_names(self): # Keys with IDs always come before keys with names. k1 = key.Key('kind', '1', project=self._PROJECT) k2 = key.Key('kind', 2, project=self._PROJECT) keys = [k1, k2] expected_sort = [k2, k1] keys.sort(key=query_splitter.client_key_sort_key) self.assertEqual(expected_sort, keys) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/rampup_throttling_fn.py0000644000175100017510000000714215206113032026706 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import datetime import logging import time from typing import TypeVar from apache_beam import typehints from apache_beam.io.gcp.datastore.v1new import util from apache_beam.metrics.metric import Metrics from apache_beam.transforms import DoFn from apache_beam.utils.retry import FuzzedExponentialIntervals T = TypeVar('T') _LOG = logging.getLogger(__name__) @typehints.with_input_types(T) @typehints.with_output_types(T) class RampupThrottlingFn(DoFn): """A ``DoFn`` that throttles ramp-up following an exponential function. An implementation of a client-side throttler that enforces a gradual ramp-up, broadly in line with Datastore best practices. See also https://cloud.google.com/datastore/docs/best-practices#ramping_up_traffic. """ _BASE_BUDGET = 500 _RAMP_UP_INTERVAL = datetime.timedelta(minutes=5) def __init__(self, num_workers, *unused_args, **unused_kwargs): """Initializes a ramp-up throttler transform. Args: num_workers: A hint for the expected number of workers, used to derive the local rate limit. """ super().__init__(*unused_args, **unused_kwargs) self._num_workers = num_workers self._successful_ops = util.MovingSum(window_ms=1000, bucket_ms=1000) self._first_instant = datetime.datetime.now() self._throttled_secs = Metrics.counter( RampupThrottlingFn, "cumulativeThrottlingSeconds") def _calc_max_ops_budget( self, first_instant: datetime.datetime, current_instant: datetime.datetime): """Function that returns per-second budget according to best practices. The exact function is `500 / num_workers * 1.5^max(0, (x-5)/5)`, where x is the number of minutes since start time. """ timedelta_since_first = current_instant - first_instant growth = max( 0.0, (timedelta_since_first - self._RAMP_UP_INTERVAL) / self._RAMP_UP_INTERVAL) try: max_ops_budget = int( self._BASE_BUDGET / self._num_workers * (1.5**growth)) except OverflowError: max_ops_budget = float('inf') return max(1, max_ops_budget) def process(self, element, **kwargs): backoff = iter( FuzzedExponentialIntervals(initial_delay_secs=1, num_retries=10000)) while True: instant = datetime.datetime.now() max_ops_budget = self._calc_max_ops_budget(self._first_instant, instant) current_op_count = self._successful_ops.sum(instant.timestamp() * 1000) available_ops = max_ops_budget - current_op_count if available_ops > 0: self._successful_ops.add(instant.timestamp() * 1000, 1) yield element break else: backoff_secs = next(backoff) _LOG.info( 'Delaying by %sms to conform to gradual ramp-up.', int(1000 * backoff_secs)) time.sleep(backoff_secs) self._throttled_secs.inc(int(backoff_secs)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/rampup_throttling_fn_test.py0000644000175100017510000000567215206113032027753 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import datetime import math import sys import unittest from mock import patch from apache_beam.io.gcp.datastore.v1new.rampup_throttling_fn import RampupThrottlingFn DATE_ZERO = datetime.datetime( year=1970, month=1, day=1, tzinfo=datetime.timezone.utc) class _RampupDelayException(Exception): pass class RampupThrottlerTransformTest(unittest.TestCase): @patch('datetime.datetime') @patch('time.sleep') def test_rampup_throttling(self, mock_sleep, mock_datetime): mock_datetime.now.return_value = DATE_ZERO throttling_fn = RampupThrottlingFn(num_workers=1) rampup_schedule = [ (DATE_ZERO + datetime.timedelta(seconds=0), 500), (DATE_ZERO + datetime.timedelta(milliseconds=1), 0), (DATE_ZERO + datetime.timedelta(seconds=1), 500), (DATE_ZERO + datetime.timedelta(seconds=1, milliseconds=1), 0), (DATE_ZERO + datetime.timedelta(minutes=5), 500), (DATE_ZERO + datetime.timedelta(minutes=10), 750), (DATE_ZERO + datetime.timedelta(minutes=15), 1125), (DATE_ZERO + datetime.timedelta(minutes=30), 3796), (DATE_ZERO + datetime.timedelta(minutes=60), 43248), ] mock_sleep.side_effect = _RampupDelayException() for date, expected_budget in rampup_schedule: mock_datetime.now.return_value = date for _ in range(expected_budget): next(throttling_fn.process(None)) # Delay after budget is exhausted with self.assertRaises(_RampupDelayException): next(throttling_fn.process(None)) def test_budget_overflow(self): throttling_fn = RampupThrottlingFn(num_workers=1) normal_date = DATE_ZERO + datetime.timedelta(minutes=2000) normal_budget = throttling_fn._calc_max_ops_budget(DATE_ZERO, normal_date) self.assertNotEqual(normal_budget, float('inf')) # This tests that a previously thrown OverflowError is caught. overflow_minutes = math.log(sys.float_info.max) / math.log(1.5) * 5 overflow_date = DATE_ZERO + datetime.timedelta(minutes=overflow_minutes) overflow_budget = throttling_fn._calc_max_ops_budget( DATE_ZERO, overflow_date) self.assertEqual(overflow_budget, float('inf')) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/types.py0000644000175100017510000002341615206113032023607 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Beam Datastore types. """ # pytype: skip-file import copy from typing import Iterable from typing import Optional from typing import Union from google.cloud.datastore import entity from google.cloud.datastore import key from google.cloud.datastore import query from apache_beam.options.value_provider import ValueProvider __all__ = ['Query', 'Key', 'Entity'] class Query(object): def __init__( self, kind=None, project=None, namespace=None, ancestor=None, filters=(), projection=(), order=(), distinct_on=(), limit=None): """Represents a Datastore query. Args: kind: (str) The kind to query. project: (str) Required. Project associated with query. namespace: (str, ValueProvider(str)) (Optional) Namespace to restrict results to. ancestor: (:class:`~apache_beam.io.gcp.datastore.v1new.types.Key`) (Optional) key of the ancestor to which this query's results are restricted. filters: (sequence of tuple[str, str, str], sequence of tuple[ValueProvider(str), ValueProvider(str), ValueProvider(str)]) Property filters applied by this query. The sequence is ``(property_name, operator, value)``. projection: (sequence of string) fields returned as part of query results. order: (sequence of string) field names used to order query results. Prepend ``-`` to a field name to sort it in descending order. distinct_on: (sequence of string) field names used to group query results. limit: (int) Maximum amount of results to return. """ self.kind = kind self.project = project self.namespace = namespace self.ancestor = ancestor self.filters = filters or () self.projection = projection self.order = order self.distinct_on = distinct_on self.limit = limit def _to_client_query(self, client): """ Returns a ``google.cloud.datastore.query.Query`` instance that represents this query. Args: client: (``google.cloud.datastore.client.Client``) Datastore client instance to use. """ ancestor_client_key = None if self.ancestor is not None: ancestor_client_key = self.ancestor.to_client_key() # Resolve ValueProvider arguments. self.filters = self._set_runtime_filters() if isinstance(self.namespace, ValueProvider): self.namespace = self.namespace.get() return query.Query( client, kind=self.kind, project=self.project, namespace=self.namespace, ancestor=ancestor_client_key, filters=self.filters, projection=self.projection, order=self.order, distinct_on=self.distinct_on) def _set_runtime_filters(self): """ Extracts values from ValueProviders in `self.filters` if available :param filters: sequence of tuple[str, str, str] or sequence of tuple[ValueProvider, ValueProvider, ValueProvider] :return: tuple[str, str, str] """ runtime_filters = [] if not all(len(filter_tuple) == 3 for filter_tuple in self.filters): raise TypeError( '%s: filters must be a sequence of tuple with length=3' ' got %r instead' % (self.__class__.__name__, self.filters)) for filter_type, filter_operator, filter_value in self.filters: if isinstance(filter_type, ValueProvider): filter_type = filter_type.get() if isinstance(filter_operator, ValueProvider): filter_operator = filter_operator.get() if isinstance(filter_value, ValueProvider): filter_value = filter_value.get() runtime_filters.append((filter_type, filter_operator, filter_value)) return runtime_filters or () def clone(self): return copy.copy(self) def __repr__(self): return ( '' % ( self.kind, self.project, self.namespace, self.ancestor, self.filters, self.projection, self.order, self.distinct_on, self.limit)) class Key(object): def __init__( self, path_elements: list[Union[str, int]], parent: Optional['Key'] = None, project: Optional[str] = None, namespace: Optional[str] = None): """ Represents a Datastore key. The partition ID is represented by its components: namespace and project. If key has a parent, project and namespace should either be unset or match the parent's. Args: path_elements: (list of str and int) Key path: an alternating sequence of kind and identifier. The kind must be of type ``str`` and identifier may be a ``str`` or an ``int``. If the last identifier is omitted this is an incomplete key, which is unsupported in ``WriteToDatastore`` and ``DeleteFromDatastore``. See :class:`google.cloud.datastore.key.Key` for more details. parent: (:class:`~apache_beam.io.gcp.datastore.v1new.types.Key`) (optional) Parent for this key. project: (str) Project ID. Required unless set by parent. namespace: (str) (optional) Namespace ID """ # Verification or arguments is delegated to to_client_key(). self.path_elements = tuple(path_elements) self.parent = parent self.namespace = namespace self.project = project @staticmethod def from_client_key(client_key): return Key( client_key.flat_path, project=client_key.project, namespace=client_key.namespace) def to_client_key(self): """ Returns a :class:`google.cloud.datastore.key.Key` instance that represents this key. """ parent = self.parent if parent is not None: parent = parent.to_client_key() return key.Key( *self.path_elements, parent=parent, namespace=self.namespace, project=self.project) def __eq__(self, other): if not isinstance(other, Key): return False if self.path_elements != other.path_elements: return False if self.project != other.project: return False if self.parent is not None and other.parent is not None: return self.parent == other.parent return self.parent is None and other.parent is None __hash__ = None # type: ignore[assignment] def __repr__(self): return '<%s(%s, parent=%s, project=%s, namespace=%s)>' % ( self.__class__.__name__, str(self.path_elements), str(self.parent), self.project, self.namespace) class Entity(object): def __init__(self, key: Key, exclude_from_indexes: Iterable[str] = ()): """ Represents a Datastore entity. Does not support the property value "meaning" field. Args: key: (Key) A complete Key representing this Entity. exclude_from_indexes: (iterable of str) List of property keys whose values should not be indexed for this entity. """ self.key = key self.exclude_from_indexes = set(exclude_from_indexes) self.properties = {} def set_properties(self, property_dict): """Sets a dictionary of properties on this entity. Args: property_dict: A map from property name to value. See :class:`google.cloud.datastore.entity.Entity` documentation for allowed values. """ self.properties.update(property_dict) @staticmethod def from_client_entity(client_entity): res = Entity( Key.from_client_key(client_entity.key), exclude_from_indexes=set(client_entity.exclude_from_indexes)) for name, value in client_entity.items(): if isinstance(value, key.Key): value = Key.from_client_key(value) if isinstance(value, entity.Entity): if value.key: value = Entity.from_client_entity(value) else: value = {k: v for k, v in value.items()} res.properties[name] = value return res def to_client_entity(self): """ Returns a :class:`google.cloud.datastore.entity.Entity` instance that represents this entity. """ res = entity.Entity( key=self.key.to_client_key(), exclude_from_indexes=tuple(self.exclude_from_indexes)) for name, value in self.properties.items(): if isinstance(value, Key): if not value.project: value.project = self.key.project value = value.to_client_key() if isinstance(value, Entity): if not value.key.project: value.key.project = self.key.project value = value.to_client_entity() res[name] = value return res def __eq__(self, other): if not isinstance(other, Entity): return False return ( self.key == other.key and self.exclude_from_indexes == other.exclude_from_indexes and self.properties == other.properties) __hash__ = None # type: ignore[assignment] def __repr__(self): return "<%s(key=%s, exclude_from_indexes=%s) properties=%s>" % ( self.__class__.__name__, str(self.key), str(self.exclude_from_indexes), str(self.properties)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/types_test.py0000644000175100017510000002100015206113032024631 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for types module.""" # pytype: skip-file import datetime import logging import unittest import mock # Protect against environments where datastore library is not available. try: from google.cloud.datastore import client from google.cloud.datastore import entity from google.cloud.datastore import key from google.cloud.datastore.helpers import GeoPoint from apache_beam.io.gcp.datastore.v1new.types import Entity from apache_beam.io.gcp.datastore.v1new.types import Key from apache_beam.io.gcp.datastore.v1new.types import Query from apache_beam.options.value_provider import StaticValueProvider except ImportError: client = None _LOGGER = logging.getLogger(__name__) @unittest.skipIf(client is None, 'Datastore dependencies are not installed') class TypesTest(unittest.TestCase): _PROJECT = 'project' _NAMESPACE = 'namespace' def setUp(self): self._test_client = client.Client( project=self._PROJECT, namespace=self._NAMESPACE, # Don't do any network requests. _http=mock.MagicMock()) def _assert_keys_equal(self, beam_type, client_type, expected_project): self.assertEqual(beam_type.path_elements[0], client_type.kind) self.assertEqual(beam_type.path_elements[1], client_type.id) self.assertEqual(expected_project, client_type.project) def testEntityToClientEntity(self): # Test conversion from Beam type to client type. k = Key(['kind', 1234], project=self._PROJECT) kc = k.to_client_key() exclude_from_indexes = ('datetime', 'key') e = Entity(k, exclude_from_indexes=exclude_from_indexes) properties = { 'datetime': datetime.datetime.utcnow(), 'key_ref': Key(['kind2', 1235]), 'bool': True, 'float': 1.21, 'int': 1337, 'unicode': 'text', 'bytes': b'bytes', 'geopoint': GeoPoint(0.123, 0.456), 'none': None, 'list': [1, 2, 3], 'entity': Entity(Key(['kind', 111])), 'dict': { 'property': 5 }, } e.set_properties(properties) ec = e.to_client_entity() self.assertEqual(kc, ec.key) self.assertSetEqual(set(exclude_from_indexes), ec.exclude_from_indexes) self.assertEqual('kind', ec.kind) self.assertEqual(1234, ec.id) for name, unconverted in properties.items(): converted = ec[name] if name == 'key_ref': self.assertNotIsInstance(converted, Key) self._assert_keys_equal(unconverted, converted, self._PROJECT) elif name == 'entity': self.assertNotIsInstance(converted, Entity) self.assertNotIsInstance(converted.key, Key) self._assert_keys_equal(unconverted.key, converted.key, self._PROJECT) else: self.assertEqual(unconverted, converted) # Test reverse conversion. entity_from_client_entity = Entity.from_client_entity(ec) self.assertEqual(e, entity_from_client_entity) def testEmbeddedClientEntityWithoutKey(self): client_entity = entity.Entity(key.Key('foo', project='bar')) entity_without_key = entity.Entity() entity_without_key['test'] = True client_entity['embedded'] = entity_without_key e = Entity.from_client_entity(client_entity) self.assertIsInstance(e.properties['embedded'], dict) def testKeyToClientKey(self): k = Key(['kind1', 'parent'], project=self._PROJECT, namespace=self._NAMESPACE) ck = k.to_client_key() self.assertEqual(self._PROJECT, ck.project) self.assertEqual(self._NAMESPACE, ck.namespace) self.assertEqual(('kind1', 'parent'), ck.flat_path) self.assertEqual('kind1', ck.kind) self.assertEqual('parent', ck.id_or_name) self.assertEqual(None, ck.parent) k2 = Key(['kind2', 1234], parent=k) ck2 = k2.to_client_key() self.assertEqual(self._PROJECT, ck2.project) self.assertEqual(self._NAMESPACE, ck2.namespace) self.assertEqual(('kind1', 'parent', 'kind2', 1234), ck2.flat_path) self.assertEqual('kind2', ck2.kind) self.assertEqual(1234, ck2.id_or_name) self.assertEqual(ck, ck2.parent) def testKeyFromClientKey(self): k = Key(['k1', 1234], project=self._PROJECT, namespace=self._NAMESPACE) kfc = Key.from_client_key(k.to_client_key()) self.assertEqual(k, kfc) k2 = Key(['k2', 'adsf'], parent=k) kfc2 = Key.from_client_key(k2.to_client_key()) # Converting a key with a parent to a client_key and back loses the parent: self.assertNotEqual(k2, kfc2) self.assertTupleEqual(('k1', 1234, 'k2', 'adsf'), kfc2.path_elements) self.assertIsNone(kfc2.parent) kfc3 = Key.from_client_key(kfc2.to_client_key()) self.assertEqual(kfc2, kfc3) kfc4 = Key.from_client_key(kfc2.to_client_key()) kfc4.project = 'other' self.assertNotEqual(kfc2, kfc4) def testKeyFromClientKeyNoNamespace(self): k = Key(['k1', 1234], project=self._PROJECT) ck = k.to_client_key() self.assertEqual(None, ck.namespace) # Test that getter doesn't croak. kfc = Key.from_client_key(ck) self.assertEqual(k, kfc) def testKeyToClientKeyMissingProject(self): k = Key(['k1', 1234], namespace=self._NAMESPACE) with self.assertRaisesRegex(ValueError, r'project'): _ = Key.from_client_key(k.to_client_key()) def testQuery(self): filters = [('property_name', '=', 'value')] projection = ['f1', 'f2'] order = projection distinct_on = projection ancestor_key = Key(['kind', 'id'], project=self._PROJECT) q = Query( kind='kind', project=self._PROJECT, namespace=self._NAMESPACE, ancestor=ancestor_key, filters=filters, projection=projection, order=order, distinct_on=distinct_on) cq = q._to_client_query(self._test_client) self.assertEqual(self._PROJECT, cq.project) self.assertEqual(self._NAMESPACE, cq.namespace) self.assertEqual('kind', cq.kind) self.assertEqual(ancestor_key.to_client_key(), cq.ancestor) self.assertEqual(filters, cq.filters) self.assertEqual(projection, cq.projection) self.assertEqual(order, cq.order) self.assertEqual(distinct_on, cq.distinct_on) _LOGGER.info('query: %s', q) # Test __repr__() def testValueProviderFilters(self): self.vp_filters = [ [( StaticValueProvider(str, 'property_name'), StaticValueProvider(str, '='), StaticValueProvider(str, 'value'))], [( StaticValueProvider(str, 'property_name'), StaticValueProvider(str, '='), StaticValueProvider(str, 'value')), ('property_name', '=', 'value')], ] self.expected_filters = [ [('property_name', '=', 'value')], [('property_name', '=', 'value'), ('property_name', '=', 'value')], ] for vp_filter, exp_filter in zip(self.vp_filters, self.expected_filters): q = Query( kind='kind', project=self._PROJECT, namespace=self._NAMESPACE, filters=vp_filter) cq = q._to_client_query(self._test_client) self.assertEqual(exp_filter, cq.filters) _LOGGER.info('query: %s', q) # Test __repr__() def testValueProviderNamespace(self): self.vp_namespace = StaticValueProvider(str, 'vp_namespace') self.expected_namespace = 'vp_namespace' q = Query(kind='kind', project=self._PROJECT, namespace=self.vp_namespace) cq = q._to_client_query(self._test_client) self.assertEqual(self.expected_namespace, cq.namespace) _LOGGER.info('query: %s', q) # Test __repr__() def testQueryEmptyNamespace(self): # Test that we can pass a namespace of None. self._test_client.namespace = None q = Query(project=self._PROJECT, namespace=None) cq = q._to_client_query(self._test_client) self.assertEqual(self._test_client.project, cq.project) self.assertEqual(None, cq.namespace) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/util.py0000644000175100017510000001231015206113032023407 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # Utility functions & classes that are _not_ specific to the datastore client. # # For internal use only; no backwards-compatibility guarantees. # pytype: skip-file import math # Constants used in batched mutation RPCs: WRITE_BATCH_INITIAL_SIZE = 50 # Max allowed Datastore writes per batch, and max bytes per batch. # Note that the max bytes per batch set here is lower than the 10MB limit # actually enforced by the API, to leave space for the CommitRequest wrapper # around the mutations. # https://cloud.google.com/datastore/docs/concepts/limits WRITE_BATCH_MAX_SIZE = 500 WRITE_BATCH_MAX_BYTES_SIZE = 9000000 WRITE_BATCH_MIN_SIZE = 5 WRITE_BATCH_TARGET_LATENCY_MS = 6000 class MovingSum(object): """Class that keeps track of a rolling window sum. For use in tracking recent performance of the connector. Intended to be similar to org.apache.beam.sdk.util.MovingFunction(..., Sum.ofLongs()), but for convenience we expose the count of entries as well so this doubles as a moving average tracker. """ def __init__(self, window_ms, bucket_ms): if window_ms < bucket_ms or bucket_ms <= 0: raise ValueError("window_ms >= bucket_ms > 0 please") self._num_buckets = int(math.ceil(window_ms / bucket_ms)) self._bucket_ms = bucket_ms self._Reset(now=0) # initialize the moving window members def _Reset(self, now): self._current_index = 0 # pointer into self._buckets self._current_ms_since_epoch = math.floor( now / self._bucket_ms) * self._bucket_ms # _buckets is a list where each element is a list [sum, num_samples] # This is a circular buffer where # [_current_index] represents the time range # [_current_ms_since_epoch, _current_ms_since_epoch+_bucket_ms) # [_current_index-1] represents immediatly prior time range # [_current_ms_since_epoch-_bucket_ms, _current_ms_since_epoch) # etc, wrapping around from the start to the end of the array, so # [_current_index+1] is the element representing the oldest bucket. self._buckets = [[0, 0] for _ in range(0, self._num_buckets)] def _Flush(self, now): """ Args: now: int, milliseconds since epoch """ if now >= (self._current_ms_since_epoch + self._bucket_ms * self._num_buckets): # Time moved forward so far that all currently held data is outside of # the window. It is faster to simply reset our data. self._Reset(now) return while now > self._current_ms_since_epoch + self._bucket_ms: # Advance time by one _bucket_ms, setting the new bucket's counts to 0. self._current_ms_since_epoch += self._bucket_ms self._current_index = (self._current_index + 1) % self._num_buckets self._buckets[self._current_index] = [0, 0] # Intentional dead reckoning here; we don't care about staying precisely # aligned with multiples of _bucket_ms since the epoch, we just need our # buckets to represent the most recent _window_ms time window. def sum(self, now): self._Flush(now) return sum(bucket[0] for bucket in self._buckets) def add(self, now, inc): self._Flush(now) bucket = self._buckets[self._current_index] bucket[0] += inc bucket[1] += 1 def count(self, now): self._Flush(now) return sum(bucket[1] for bucket in self._buckets) def has_data(self, now): return self.count(now) > 0 class DynamicBatchSizer(object): """Determines request sizes for future Datastore RPCs.""" def __init__(self): self._commit_time_per_entity_ms = MovingSum( window_ms=120000, bucket_ms=10000) def get_batch_size(self, now): """Returns the recommended size for datastore RPCs at this time.""" if not self._commit_time_per_entity_ms.has_data(now): return WRITE_BATCH_INITIAL_SIZE recent_mean_latency_ms = ( self._commit_time_per_entity_ms.sum(now) // self._commit_time_per_entity_ms.count(now)) return max( WRITE_BATCH_MIN_SIZE, min( WRITE_BATCH_MAX_SIZE, WRITE_BATCH_TARGET_LATENCY_MS // max(recent_mean_latency_ms, 1))) def report_latency(self, now, latency_ms, num_mutations): """Report the latency of a Datastore RPC. Args: now: double, completion time of the RPC as seconds since the epoch. latency_ms: double, the observed latency in milliseconds for this RPC. num_mutations: int, number of mutations contained in the RPC. """ self._commit_time_per_entity_ms.add(now, latency_ms / num_mutations) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/datastore/v1new/util_test.py0000644000175100017510000000700615206113032024454 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for util.py.""" # pytype: skip-file import unittest from apache_beam.io.gcp.datastore.v1new import util class MovingSumTest(unittest.TestCase): TIMESTAMP = 1500000000 def test_bad_bucket_size(self): with self.assertRaises(ValueError): _ = util.MovingSum(1, 0) def test_bad_window_size(self): with self.assertRaises(ValueError): _ = util.MovingSum(1, 2) def test_no_data(self): ms = util.MovingSum(10, 1) self.assertEqual(0, ms.sum(MovingSumTest.TIMESTAMP)) self.assertEqual(0, ms.count(MovingSumTest.TIMESTAMP)) self.assertFalse(ms.has_data(MovingSumTest.TIMESTAMP)) def test_one_data_point(self): ms = util.MovingSum(10, 1) ms.add(MovingSumTest.TIMESTAMP, 5) self.assertEqual(5, ms.sum(MovingSumTest.TIMESTAMP)) self.assertEqual(1, ms.count(MovingSumTest.TIMESTAMP)) self.assertTrue(ms.has_data(MovingSumTest.TIMESTAMP)) def test_aggregates_within_window(self): ms = util.MovingSum(10, 1) ms.add(MovingSumTest.TIMESTAMP, 5) ms.add(MovingSumTest.TIMESTAMP + 1, 3) ms.add(MovingSumTest.TIMESTAMP + 2, 7) self.assertEqual(15, ms.sum(MovingSumTest.TIMESTAMP + 3)) self.assertEqual(3, ms.count(MovingSumTest.TIMESTAMP + 3)) def test_data_expires_from_moving_window(self): ms = util.MovingSum(5, 1) ms.add(MovingSumTest.TIMESTAMP, 5) ms.add(MovingSumTest.TIMESTAMP + 3, 3) ms.add(MovingSumTest.TIMESTAMP + 6, 7) self.assertEqual(10, ms.sum(MovingSumTest.TIMESTAMP + 7)) self.assertEqual(2, ms.count(MovingSumTest.TIMESTAMP + 7)) class DynamicWriteBatcherTest(unittest.TestCase): def setUp(self): self._batcher = util.DynamicBatchSizer() # If possible, keep these test cases aligned with the Java test cases in # DatastoreV1Test.java def test_no_data(self): self.assertEqual( util.WRITE_BATCH_INITIAL_SIZE, self._batcher.get_batch_size(0)) def test_fast_queries(self): self._batcher.report_latency(0, 1000, 200) self._batcher.report_latency(0, 1000, 200) self.assertEqual(util.WRITE_BATCH_MAX_SIZE, self._batcher.get_batch_size(0)) def test_slow_queries(self): self._batcher.report_latency(0, 10000, 200) self._batcher.report_latency(0, 10000, 200) self.assertEqual(120, self._batcher.get_batch_size(0)) def test_size_not_below_minimum(self): self._batcher.report_latency(0, 75000, 50) self._batcher.report_latency(0, 75000, 50) self.assertEqual(util.WRITE_BATCH_MIN_SIZE, self._batcher.get_batch_size(0)) def test_sliding_window(self): self._batcher.report_latency(0, 30000, 50) self._batcher.report_latency(50000, 5000, 200) self._batcher.report_latency(100000, 5000, 200) self.assertEqual(240, self._batcher.get_batch_size(150000)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/dicomclient.py0000644000175100017510000000211615206113032021701 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # For backwards compatibility. # pytype: skip-file import warnings # pylint: disable=unused-import from apache_beam.io.gcp.healthcare import dicomclient from apache_beam.io.gcp.healthcare import dicomio warnings.warn( "Dicom Client moved to apache_beam.io.gcp.healthcare.dicomclient", DeprecationWarning) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/dicomio.py0000644000175100017510000000210715206113032021032 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # For backwards compatibility. # pytype: skip-file import warnings # pylint: disable=unused-import from apache_beam.io.gcp.healthcare import dicomclient from apache_beam.io.gcp.healthcare import dicomio warnings.warn( "Dicom IO moved to apache_beam.io.gcp.healthcare.dicomcio", DeprecationWarning) ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3774076 apache_beam-2.74.0/apache_beam/io/gcp/experimental/0000755000175100017510000000000015206113157021542 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/experimental/__init__.py0000644000175100017510000000142015206113032023640 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/experimental/spannerio.py0000644000175100017510000013666515206113032024123 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Google Cloud Spanner IO Deprecated; use apache_beam.io.gcp.spanner module instead. This is an experimental module for reading and writing data from Google Cloud Spanner. Visit: https://cloud.google.com/spanner for more details. Reading Data from Cloud Spanner. To read from Cloud Spanner apply ReadFromSpanner transformation. It will return a PCollection, where each element represents an individual row returned from the read operation. Both Query and Read APIs are supported. ReadFromSpanner relies on the ReadOperation objects which is exposed by the SpannerIO API. ReadOperation holds the immutable data which is responsible to execute batch and naive reads on Cloud Spanner. This is done for more convenient programming. ReadFromSpanner reads from Cloud Spanner by providing either an 'sql' param in the constructor or 'table' name with 'columns' as list. For example::: records = (pipeline | ReadFromSpanner(PROJECT_ID, INSTANCE_ID, DB_NAME, sql='Select * from users')) records = (pipeline | ReadFromSpanner(PROJECT_ID, INSTANCE_ID, DB_NAME, table='users', columns=['id', 'name', 'email'])) You can also perform multiple reads by providing a list of ReadOperations to the ReadFromSpanner transform constructor. ReadOperation exposes two static methods. Use 'query' to perform sql based reads, 'table' to perform read from table name. For example::: read_operations = [ ReadOperation.table(table='customers', columns=['name', 'email']), ReadOperation.table(table='vendors', columns=['name', 'email']), ] all_users = pipeline | ReadFromSpanner(PROJECT_ID, INSTANCE_ID, DB_NAME, read_operations=read_operations) ...OR... read_operations = [ ReadOperation.query(sql='Select name, email from customers'), ReadOperation.query( sql='Select * from users where id <= @user_id', params={'user_id': 100}, params_type={'user_id': param_types.INT64} ), ] # `params_types` are instance of `google.cloud.spanner.param_types` all_users = pipeline | ReadFromSpanner(PROJECT_ID, INSTANCE_ID, DB_NAME, read_operations=read_operations) For more information, please review the docs on class ReadOperation. User can also able to provide the ReadOperation in form of PCollection via pipeline. For example::: users = (pipeline | beam.Create([ReadOperation...]) | ReadFromSpanner(PROJECT_ID, INSTANCE_ID, DB_NAME)) User may also create cloud spanner transaction from the transform called `create_transaction` which is available in the SpannerIO API. The transform is guaranteed to be executed on a consistent snapshot of data, utilizing the power of read only transactions. Staleness of data can be controlled by providing the `read_timestamp` or `exact_staleness` param values in the constructor. This transform requires root of the pipeline (PBegin) and returns PTransform which is passed later to the `ReadFromSpanner` constructor. `ReadFromSpanner` pass this transaction PTransform as a singleton side input to the `_NaiveSpannerReadDoFn` containing 'session_id' and 'transaction_id'. For example::: transaction = (pipeline | create_transaction(TEST_PROJECT_ID, TEST_INSTANCE_ID, DB_NAME)) users = pipeline | ReadFromSpanner(PROJECT_ID, INSTANCE_ID, DB_NAME, sql='Select * from users', transaction=transaction) tweets = pipeline | ReadFromSpanner(PROJECT_ID, INSTANCE_ID, DB_NAME, sql='Select * from tweets', transaction=transaction) For further details of this transform, please review the docs on the :meth:`create_transaction` method available in the SpannerIO API. ReadFromSpanner takes this transform in the constructor and pass this to the read pipeline as the singleton side input. Writing Data to Cloud Spanner. The WriteToSpanner transform writes to Cloud Spanner by executing a collection a input rows (WriteMutation). The mutations are grouped into batches for efficiency. WriteToSpanner transform relies on the WriteMutation objects which is exposed by the SpannerIO API. WriteMutation have five static methods (insert, update, insert_or_update, replace, delete). These methods returns the instance of the _Mutator object which contains the mutation type and the Spanner Mutation object. For more details, review the docs of the class SpannerIO.WriteMutation. For example::: mutations = [ WriteMutation.insert(table='user', columns=('name', 'email'), values=[('sara', 'sara@dev.com')]) ] _ = (p | beam.Create(mutations) | WriteToSpanner( project_id=SPANNER_PROJECT_ID, instance_id=SPANNER_INSTANCE_ID, database_id=SPANNER_DATABASE_NAME) ) You can also create WriteMutation via calling its constructor. For example::: mutations = [ WriteMutation(insert='users', columns=('name', 'email'), values=[('sara", 'sara@example.com')]) ] For more information, review the docs available on WriteMutation class. WriteToSpanner transform also takes three batching parameters (max_number_rows, max_number_cells and max_batch_size_bytes). By default, max_number_rows is set to 50 rows, max_number_cells is set to 500 cells and max_batch_size_bytes is set to 1MB (1048576 bytes). These parameter used to reduce the number of transactions sent to spanner by grouping the mutation into batches. Setting these param values either to smaller value or zero to disable batching. Unlike the Java connector, this connector does not create batches of transactions sorted by table and primary key. WriteToSpanner transforms starts with the grouping into batches. The first step in this process is to make the mutation groups of the WriteMutation objects and then filtering them into batchable and unbatchable mutation groups. There are three batching parameters (max_number_cells, max_number_rows & max_batch_size_bytes). We calculated th mutation byte size from the method available in the `google.cloud.spanner_v1.proto.mutation_pb2.Mutation.ByteSize`. if the mutation rows, cells or byte size are larger than value of the any batching parameters param, it will be tagged as "unbatchable" mutation. After this all the batchable mutation are merged into a single mutation group whos size is not larger than the "max_batch_size_bytes", after this process, all the mutation groups together to process. If the Mutation references a table or column does not exits, it will cause a exception and fails the entire pipeline. """ import typing from collections import deque from collections import namedtuple from apache_beam import Create from apache_beam import DoFn from apache_beam import Flatten from apache_beam import ParDo from apache_beam import Reshuffle from apache_beam.internal.metrics.metric import ServiceCallMetric from apache_beam.io.gcp import resource_identifiers from apache_beam.metrics import Metrics from apache_beam.metrics import monitoring_infos from apache_beam.pvalue import AsSingleton from apache_beam.pvalue import PBegin from apache_beam.pvalue import TaggedOutput from apache_beam.transforms import PTransform from apache_beam.transforms import ptransform_fn from apache_beam.transforms import window from apache_beam.transforms.display import DisplayDataItem from apache_beam.typehints import with_input_types from apache_beam.typehints import with_output_types from apache_beam.utils.annotations import deprecated # Protect against environments where spanner library is not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports # pylint: disable=unused-import try: from google.api_core.exceptions import ClientError from google.api_core.exceptions import GoogleAPICallError from google.cloud.spanner import Client from google.cloud.spanner import KeySet from google.cloud.spanner_v1 import batch from google.cloud.spanner_v1.database import BatchSnapshot except ImportError: Client = None KeySet = None BatchSnapshot = None try: from google.cloud.spanner_v1 import Mutation except ImportError: try: # Remove this and the try clause when we upgrade to google-cloud-spanner # 3.x.x. from google.cloud.spanner_v1.proto.mutation_pb2 import Mutation except ImportError: # Ignoring for environments where the Spanner library is not available. pass __all__ = [ 'create_transaction', 'ReadFromSpanner', 'ReadOperation', 'WriteToSpanner', 'WriteMutation', 'MutationGroup' ] class _SPANNER_TRANSACTION(namedtuple("SPANNER_TRANSACTION", ["transaction"])): """ Holds the spanner transaction details. """ __slots__ = () class ReadOperation(namedtuple( "ReadOperation", ["is_sql", "is_table", "read_operation", "kwargs"])): """ Encapsulates a spanner read operation. """ __slots__ = () @classmethod def query(cls, sql, params=None, param_types=None): """ A convenient method to construct ReadOperation from sql query. Args: sql: SQL query statement params: (optional) values for parameter replacement. Keys must match the names used in sql param_types: (optional) maps explicit types for one or more param values; required if parameters are passed. """ if params: assert param_types is not None return cls( is_sql=True, is_table=False, read_operation="process_query_batch", kwargs={ 'sql': sql, 'params': params, 'param_types': param_types }) @classmethod def table(cls, table, columns, index="", keyset=None): """ A convenient method to construct ReadOperation from table. Args: table: name of the table from which to fetch data. columns: names of columns to be retrieved. index: (optional) name of index to use, rather than the table's primary key. keyset: (optional) `KeySet` keys / ranges identifying rows to be retrieved. """ keyset = keyset or KeySet(all_=True) if not isinstance(keyset, KeySet): raise ValueError( "keyset must be an instance of class " "google.cloud.spanner.KeySet") return cls( is_sql=False, is_table=True, read_operation="process_read_batch", kwargs={ 'table': table, 'columns': columns, 'index': index, 'keyset': keyset }) class _BeamSpannerConfiguration(namedtuple("_BeamSpannerConfiguration", ["project", "instance", "database", "table", "query_name", "credentials", "pool", "snapshot_read_timestamp", "snapshot_exact_staleness"])): """ A namedtuple holds the immutable data of the connection string to the cloud spanner. """ @property def snapshot_options(self): snapshot_options = {} if self.snapshot_exact_staleness: snapshot_options['exact_staleness'] = self.snapshot_exact_staleness if self.snapshot_read_timestamp: snapshot_options['read_timestamp'] = self.snapshot_read_timestamp return snapshot_options @with_input_types(ReadOperation, _SPANNER_TRANSACTION) @with_output_types(list[typing.Any]) class _NaiveSpannerReadDoFn(DoFn): def __init__(self, spanner_configuration): """ A naive version of Spanner read which uses the transaction API of the cloud spanner. https://googleapis.dev/python/spanner/latest/transaction-api.html In Naive reads, this transform performs single reads, where as the Batch reads use the spanner partitioning query to create batches. Args: spanner_configuration: (_BeamSpannerConfiguration) Connection details to connect with cloud spanner. """ self._spanner_configuration = spanner_configuration self._snapshot = None self._session = None self.base_labels = { monitoring_infos.SERVICE_LABEL: 'Spanner', monitoring_infos.METHOD_LABEL: 'Read', monitoring_infos.SPANNER_PROJECT_ID: ( self._spanner_configuration.project), monitoring_infos.SPANNER_DATABASE_ID: ( self._spanner_configuration.database), } def _table_metric(self, table_id, status): database_id = self._spanner_configuration.database project_id = self._spanner_configuration.project resource = resource_identifiers.SpannerTable( project_id, database_id, table_id) labels = { **self.base_labels, monitoring_infos.RESOURCE_LABEL: resource, } if table_id: labels[monitoring_infos.SPANNER_TABLE_ID] = table_id service_call_metric = ServiceCallMetric( request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN, base_labels=labels) service_call_metric.call(str(status)) def _query_metric(self, query_name, status): project_id = self._spanner_configuration.project resource = resource_identifiers.SpannerSqlQuery(project_id, query_name) labels = { **self.base_labels, monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.SPANNER_QUERY_NAME: query_name } service_call_metric = ServiceCallMetric( request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN, base_labels=labels) service_call_metric.call(str(status)) def _get_session(self): if self._session is None: session = self._session = self._database.session() session.create() return self._session def _close_session(self): if self._session is not None: self._session.delete() def setup(self): # setting up client to connect with cloud spanner spanner_client = Client(self._spanner_configuration.project) instance = spanner_client.instance(self._spanner_configuration.instance) self._database = instance.database( self._spanner_configuration.database, pool=self._spanner_configuration.pool) def process(self, element, spanner_transaction): # `spanner_transaction` should be the instance of the _SPANNER_TRANSACTION # object. if not isinstance(spanner_transaction, _SPANNER_TRANSACTION): raise ValueError( "Invalid transaction object: %s. It should be instance " "of SPANNER_TRANSACTION object created by " "spannerio.create_transaction transform." % type(spanner_transaction)) transaction_info = spanner_transaction.transaction # We used batch snapshot to reuse the same transaction passed through the # side input self._snapshot = BatchSnapshot.from_dict(self._database, transaction_info) # getting the transaction from the snapshot's session to run read operation. # with self._snapshot.session().transaction() as transaction: with self._get_session().transaction() as transaction: table_id = self._spanner_configuration.table query_name = self._spanner_configuration.query_name or '' if element.is_sql is True: transaction_read = transaction.execute_sql metric_action = self._query_metric metric_id = query_name elif element.is_table is True: transaction_read = transaction.read metric_action = self._table_metric metric_id = table_id else: raise ValueError( "ReadOperation is improperly configure: %s" % str(element)) try: for row in transaction_read(**element.kwargs): yield row metric_action(metric_id, 'ok') except (ClientError, GoogleAPICallError) as e: metric_action(metric_id, e.code.value) raise @with_input_types(ReadOperation) @with_output_types(dict[typing.Any, typing.Any]) class _CreateReadPartitions(DoFn): """ A DoFn to create partitions. Uses the Partitioning API (PartitionRead / PartitionQuery) request to start a partitioned query operation. Returns a list of batch information needed to perform the actual queries. If the element is the instance of :class:`ReadOperation` is to perform sql query, `PartitionQuery` API is used the create partitions and returns mappings of information used perform actual partitioned reads via :meth:`process_query_batch`. If the element is the instance of :class:`ReadOperation` is to perform read from table, `PartitionRead` API is used the create partitions and returns mappings of information used perform actual partitioned reads via :meth:`process_read_batch`. """ def __init__(self, spanner_configuration): self._spanner_configuration = spanner_configuration def setup(self): spanner_client = Client( project=self._spanner_configuration.project, credentials=self._spanner_configuration.credentials) instance = spanner_client.instance(self._spanner_configuration.instance) self._database = instance.database( self._spanner_configuration.database, pool=self._spanner_configuration.pool) self._snapshot = self._database.batch_snapshot( **self._spanner_configuration.snapshot_options) self._snapshot_dict = self._snapshot.to_dict() def process(self, element): if element.is_sql is True: partitioning_action = self._snapshot.generate_query_batches elif element.is_table is True: partitioning_action = self._snapshot.generate_read_batches else: raise ValueError( "ReadOperation is improperly configure: %s" % str(element)) for p in partitioning_action(**element.kwargs): yield { "is_sql": element.is_sql, "is_table": element.is_table, "read_operation": element.read_operation, "partitions": p, "transaction_info": self._snapshot_dict } @with_input_types(int) @with_output_types(_SPANNER_TRANSACTION) class _CreateTransactionFn(DoFn): """ A DoFn to create the transaction of cloud spanner. It connects to the database and and returns the transaction_id and session_id by using the batch_snapshot.to_dict() method available in the google cloud spanner sdk. https://googleapis.dev/python/spanner/latest/database-api.html?highlight= batch_snapshot#google.cloud.spanner_v1.database.BatchSnapshot.to_dict """ def __init__( self, project_id, instance_id, database_id, credentials, pool, read_timestamp, exact_staleness): self._project_id = project_id self._instance_id = instance_id self._database_id = database_id self._credentials = credentials self._pool = pool self._snapshot_options = {} if read_timestamp: self._snapshot_options['read_timestamp'] = read_timestamp if exact_staleness: self._snapshot_options['exact_staleness'] = exact_staleness self._snapshot = None def setup(self): self._spanner_client = Client( project=self._project_id, credentials=self._credentials) self._instance = self._spanner_client.instance(self._instance_id) self._database = self._instance.database(self._database_id, pool=self._pool) def process(self, element, *args, **kwargs): self._snapshot = self._database.batch_snapshot(**self._snapshot_options) return [_SPANNER_TRANSACTION(self._snapshot.to_dict())] @ptransform_fn def create_transaction( pbegin, project_id, instance_id, database_id, credentials=None, pool=None, read_timestamp=None, exact_staleness=None): """ A PTransform method to create a batch transaction. Args: pbegin: Root of the pipeline project_id: Cloud spanner project id. Be sure to use the Project ID, not the Project Number. instance_id: Cloud spanner instance id. database_id: Cloud spanner database id. credentials: (optional) The authorization credentials to attach to requests. These credentials identify this application to the service. If none are specified, the client will attempt to ascertain the credentials from the environment. pool: (optional) session pool to be used by database. If not passed, Spanner Cloud SDK uses the BurstyPool by default. `google.cloud.spanner.BurstyPool`. Ref: https://googleapis.dev/python/spanner/latest/database-api.html?#google. cloud.spanner_v1.database.Database read_timestamp: (optional) An instance of the `datetime.datetime` object to execute all reads at the given timestamp. exact_staleness: (optional) An instance of the `datetime.timedelta` object. These timestamp bounds execute reads at a user-specified timestamp. """ assert isinstance(pbegin, PBegin) return ( pbegin | Create([1]) | ParDo( _CreateTransactionFn( project_id, instance_id, database_id, credentials, pool, read_timestamp, exact_staleness))) @with_input_types(dict[typing.Any, typing.Any]) @with_output_types(list[typing.Any]) class _ReadFromPartitionFn(DoFn): """ A DoFn to perform reads from the partition. """ def __init__(self, spanner_configuration): self._spanner_configuration = spanner_configuration self.base_labels = { monitoring_infos.SERVICE_LABEL: 'Spanner', monitoring_infos.METHOD_LABEL: 'Read', monitoring_infos.SPANNER_PROJECT_ID: ( self._spanner_configuration.project), monitoring_infos.SPANNER_DATABASE_ID: ( self._spanner_configuration.database), } self.service_metric = None def _table_metric(self, table_id): database_id = self._spanner_configuration.database project_id = self._spanner_configuration.project resource = resource_identifiers.SpannerTable( project_id, database_id, table_id) labels = { **self.base_labels, monitoring_infos.RESOURCE_LABEL: resource, } if table_id: labels[monitoring_infos.SPANNER_TABLE_ID] = table_id service_call_metric = ServiceCallMetric( request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN, base_labels=labels) return service_call_metric def _query_metric(self, query_name): project_id = self._spanner_configuration.project resource = resource_identifiers.SpannerSqlQuery(project_id, query_name) labels = { **self.base_labels, monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.SPANNER_QUERY_NAME: query_name } service_call_metric = ServiceCallMetric( request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN, base_labels=labels) return service_call_metric def setup(self): spanner_client = Client(self._spanner_configuration.project) instance = spanner_client.instance(self._spanner_configuration.instance) self._database = instance.database( self._spanner_configuration.database, pool=self._spanner_configuration.pool) self._snapshot = self._database.batch_snapshot( **self._spanner_configuration.snapshot_options) def process(self, element): self._snapshot = BatchSnapshot.from_dict( self._database, element['transaction_info']) table_id = self._spanner_configuration.table query_name = self._spanner_configuration.query_name or '' if element['is_sql'] is True: read_action = self._snapshot.process_query_batch self.service_metric = self._query_metric(query_name) elif element['is_table'] is True: read_action = self._snapshot.process_read_batch self.service_metric = self._table_metric(table_id) else: raise ValueError( "ReadOperation is improperly configure: %s" % str(element)) try: for row in read_action(element['partitions']): yield row self.service_metric.call('ok') except (ClientError, GoogleAPICallError) as e: self.service_metric(str(e.code.value)) raise def teardown(self): if self._snapshot: self._snapshot.close() @deprecated(since='2.68', current='apache_beam.io.gcp.spanner.ReadFromSpanner') class ReadFromSpanner(PTransform): """ A PTransform to perform reads from cloud spanner. ReadFromSpanner uses BatchAPI to perform all read operations. """ def __init__( self, project_id, instance_id, database_id, pool=None, read_timestamp=None, exact_staleness=None, credentials=None, sql=None, params=None, param_types=None, # with_query table=None, query_name=None, columns=None, index="", keyset=None, # with_table read_operations=None, # for read all transaction=None): """ A PTransform that uses Spanner Batch API to perform reads. Args: project_id: Cloud spanner project id. Be sure to use the Project ID, not the Project Number. instance_id: Cloud spanner instance id. database_id: Cloud spanner database id. pool: (optional) session pool to be used by database. If not passed, Spanner Cloud SDK uses the BurstyPool by default. `google.cloud.spanner.BurstyPool`. Ref: https://googleapis.dev/python/spanner/latest/database-api.html?#google. cloud.spanner_v1.database.Database read_timestamp: (optional) An instance of the `datetime.datetime` object to execute all reads at the given timestamp. By default, set to `None`. exact_staleness: (optional) An instance of the `datetime.timedelta` object. These timestamp bounds execute reads at a user-specified timestamp. By default, set to `None`. credentials: (optional) The authorization credentials to attach to requests. These credentials identify this application to the service. If none are specified, the client will attempt to ascertain the credentials from the environment. By default, set to `None`. sql: (optional) SQL query statement. params: (optional) Values for parameter replacement. Keys must match the names used in sql. By default, set to `None`. param_types: (optional) maps explicit types for one or more param values; required if params are passed. By default, set to `None`. table: (optional) Name of the table from which to fetch data. By default, set to `None`. columns: (optional) List of names of columns to be retrieved; required if the table is passed. By default, set to `None`. index: (optional) name of index to use, rather than the table's primary key. By default, set to `None`. keyset: (optional) keys / ranges identifying rows to be retrieved. By default, set to `None`. read_operations: (optional) List of the objects of :class:`ReadOperation` to perform read all. By default, set to `None`. transaction: (optional) PTransform of the :meth:`create_transaction` to perform naive read on cloud spanner. By default, set to `None`. """ self._configuration = _BeamSpannerConfiguration( project=project_id, instance=instance_id, database=database_id, table=table, query_name=query_name, credentials=credentials, pool=pool, snapshot_read_timestamp=read_timestamp, snapshot_exact_staleness=exact_staleness) self._read_operations = read_operations self._transaction = transaction if self._read_operations is None: if table is not None: if columns is None: raise ValueError("Columns are required with the table name.") self._read_operations = [ ReadOperation.table( table=table, columns=columns, index=index, keyset=keyset) ] elif sql is not None: self._read_operations = [ ReadOperation.query( sql=sql, params=params, param_types=param_types) ] def expand(self, pbegin): if self._read_operations is not None and isinstance(pbegin, PBegin): pcoll = pbegin.pipeline | Create(self._read_operations) elif not isinstance(pbegin, PBegin): if self._read_operations is not None: raise ValueError( "Read operation in the constructor only works with " "the root of the pipeline.") pcoll = pbegin else: raise ValueError( "Spanner required read operation, sql or table " "with columns.") if self._transaction is None: # reading as batch read using the spanner partitioning query to create # batches. p = ( pcoll | 'Generate Partitions' >> ParDo( _CreateReadPartitions(spanner_configuration=self._configuration)) | 'Reshuffle' >> Reshuffle() | 'Read From Partitions' >> ParDo( _ReadFromPartitionFn(spanner_configuration=self._configuration))) else: # reading as naive read, in which we don't make batches and execute the # queries as a single read. p = ( pcoll | 'Reshuffle' >> Reshuffle().with_input_types(ReadOperation) | 'Perform Read' >> ParDo( _NaiveSpannerReadDoFn(spanner_configuration=self._configuration), AsSingleton(self._transaction))) return p def display_data(self): res = {} sql = [] table = [] if self._read_operations is not None: for ro in self._read_operations: if ro.is_sql is True: sql.append(ro.kwargs) elif ro.is_table is True: table.append(ro.kwargs) if sql: res['sql'] = DisplayDataItem(str(sql), label='Sql') if table: res['table'] = DisplayDataItem(str(table), label='Table') if self._transaction: res['transaction'] = DisplayDataItem( str(self._transaction), label='transaction') return res @deprecated( since='2.68', current='apache_beam.io.gcp.spanner.WriteToSpannerSchema') class WriteToSpanner(PTransform): def __init__( self, project_id, instance_id, database_id, pool=None, credentials=None, max_batch_size_bytes=1048576, max_number_rows=50, max_number_cells=500): """ A PTransform to write onto Google Cloud Spanner. Args: project_id: Cloud spanner project id. Be sure to use the Project ID, not the Project Number. instance_id: Cloud spanner instance id. database_id: Cloud spanner database id. max_batch_size_bytes: (optional) Split the mutations into batches to reduce the number of transaction sent to Spanner. By default it is set to 1 MB (1048576 Bytes). max_number_rows: (optional) Split the mutations into batches to reduce the number of transaction sent to Spanner. By default it is set to 50 rows per batch. max_number_cells: (optional) Split the mutations into batches to reduce the number of transaction sent to Spanner. By default it is set to 500 cells per batch. """ self._configuration = _BeamSpannerConfiguration( project=project_id, instance=instance_id, database=database_id, table=None, query_name=None, credentials=credentials, pool=pool, snapshot_read_timestamp=None, snapshot_exact_staleness=None) self._max_batch_size_bytes = max_batch_size_bytes self._max_number_rows = max_number_rows self._max_number_cells = max_number_cells self._database_id = database_id self._project_id = project_id self._instance_id = instance_id self._pool = pool def display_data(self): res = { 'project_id': DisplayDataItem(self._project_id, label='Project Id'), 'instance_id': DisplayDataItem(self._instance_id, label='Instance Id'), 'pool': DisplayDataItem(str(self._pool), label='Pool'), 'database': DisplayDataItem(self._database_id, label='Database'), 'batch_size': DisplayDataItem( self._max_batch_size_bytes, label="Batch Size"), 'max_number_rows': DisplayDataItem( self._max_number_rows, label="Max Rows"), 'max_number_cells': DisplayDataItem( self._max_number_cells, label="Max Cells"), } return res def expand(self, pcoll): return ( pcoll | "make batches" >> _WriteGroup( max_batch_size_bytes=self._max_batch_size_bytes, max_number_rows=self._max_number_rows, max_number_cells=self._max_number_cells) | 'Writing to spanner' >> ParDo(_WriteToSpannerDoFn(self._configuration))) class _Mutator(namedtuple('_Mutator', ["mutation", "operation", "kwargs", "rows", "cells"]) ): __slots__ = () @property def byte_size(self): if hasattr(self.mutation, '_pb'): # google-cloud-spanner 3.x return self.mutation._pb.ByteSize() else: # google-cloud-spanner 1.x return self.mutation.ByteSize() class MutationGroup(deque): """ A Bundle of Spanner Mutations (_Mutator). """ @property def info(self): cells = 0 rows = 0 bytes = 0 for m in self.__iter__(): bytes += m.byte_size rows += m.rows cells += m.cells return {"rows": rows, "cells": cells, "byte_size": bytes} def primary(self): return next(self.__iter__()) class WriteMutation(object): _OPERATION_DELETE = "delete" _OPERATION_INSERT = "insert" _OPERATION_INSERT_OR_UPDATE = "insert_or_update" _OPERATION_REPLACE = "replace" _OPERATION_UPDATE = "update" def __init__( self, insert=None, update=None, insert_or_update=None, replace=None, delete=None, columns=None, values=None, keyset=None): """ A convenient class to create Spanner Mutations for Write. User can provide the operation via constructor or via static methods. Note: If a user passing the operation via construction, make sure that it will only accept one operation at a time. For example, if a user passing a table name in the `insert` parameter, and he also passes the `update` parameter value, this will cause an error. Args: insert: (Optional) Name of the table in which rows will be inserted. update: (Optional) Name of the table in which existing rows will be updated. insert_or_update: (Optional) Table name in which rows will be written. Like insert, except that if the row already exists, then its column values are overwritten with the ones provided. Any column values not explicitly written are preserved. replace: (Optional) Table name in which rows will be replaced. Like insert, except that if the row already exists, it is deleted, and the column values provided are inserted instead. Unlike `insert_or_update`, this means any values not explicitly written become `NULL`. delete: (Optional) Table name from which rows will be deleted. Succeeds whether or not the named rows were present. columns: The names of the columns in table to be written. The list of columns must contain enough columns to allow Cloud Spanner to derive values for all primary key columns in the row(s) to be modified. values: The values to be written. `values` can contain more than one list of values. If it does, then multiple rows are written, one for each entry in `values`. Each list in `values` must have exactly as many entries as there are entries in columns above. Sending multiple lists is equivalent to sending multiple Mutations, each containing one `values` entry and repeating table and columns. keyset: (Optional) The primary keys of the rows within table to delete. Delete is idempotent. The transaction will succeed even if some or all rows do not exist. """ self._columns = columns self._values = values self._keyset = keyset self._insert = insert self._update = update self._insert_or_update = insert_or_update self._replace = replace self._delete = delete if sum([1 for x in [self._insert, self._update, self._insert_or_update, self._replace, self._delete] if x is not None]) != 1: raise ValueError( "No or more than one write mutation operation " "provided: <%s: %s>" % (self.__class__.__name__, str(self.__dict__))) def __call__(self, *args, **kwargs): if self._insert is not None: return WriteMutation.insert( table=self._insert, columns=self._columns, values=self._values) elif self._update is not None: return WriteMutation.update( table=self._update, columns=self._columns, values=self._values) elif self._insert_or_update is not None: return WriteMutation.insert_or_update( table=self._insert_or_update, columns=self._columns, values=self._values) elif self._replace is not None: return WriteMutation.replace( table=self._replace, columns=self._columns, values=self._values) elif self._delete is not None: return WriteMutation.delete(table=self._delete, keyset=self._keyset) @staticmethod def insert(table, columns, values): """Insert one or more new table rows. Args: table: Name of the table to be modified. columns: Name of the table columns to be modified. values: Values to be modified. """ rows = len(values) cells = len(columns) * len(values) return _Mutator( mutation=Mutation(insert=batch._make_write_pb(table, columns, values)), operation=WriteMutation._OPERATION_INSERT, rows=rows, cells=cells, kwargs={ "table": table, "columns": columns, "values": values }) @staticmethod def update(table, columns, values): """Update one or more existing table rows. Args: table: Name of the table to be modified. columns: Name of the table columns to be modified. values: Values to be modified. """ rows = len(values) cells = len(columns) * len(values) return _Mutator( mutation=Mutation(update=batch._make_write_pb(table, columns, values)), operation=WriteMutation._OPERATION_UPDATE, rows=rows, cells=cells, kwargs={ "table": table, "columns": columns, "values": values }) @staticmethod def insert_or_update(table, columns, values): """Insert/update one or more table rows. Args: table: Name of the table to be modified. columns: Name of the table columns to be modified. values: Values to be modified. """ rows = len(values) cells = len(columns) * len(values) return _Mutator( mutation=Mutation( insert_or_update=batch._make_write_pb(table, columns, values)), operation=WriteMutation._OPERATION_INSERT_OR_UPDATE, rows=rows, cells=cells, kwargs={ "table": table, "columns": columns, "values": values }) @staticmethod def replace(table, columns, values): """Replace one or more table rows. Args: table: Name of the table to be modified. columns: Name of the table columns to be modified. values: Values to be modified. """ rows = len(values) cells = len(columns) * len(values) return _Mutator( mutation=Mutation(replace=batch._make_write_pb(table, columns, values)), operation=WriteMutation._OPERATION_REPLACE, rows=rows, cells=cells, kwargs={ "table": table, "columns": columns, "values": values }) @staticmethod def delete(table, keyset): """Delete one or more table rows. Args: table: Name of the table to be modified. keyset: Keys/ranges identifying rows to delete. """ delete = Mutation.Delete(table=table, key_set=keyset._to_pb()) return _Mutator( mutation=Mutation(delete=delete), rows=0, cells=0, operation=WriteMutation._OPERATION_DELETE, kwargs={ "table": table, "keyset": keyset }) @with_input_types(typing.Union[MutationGroup, TaggedOutput]) @with_output_types(MutationGroup) class _BatchFn(DoFn): """ Batches mutations together. """ def __init__(self, max_batch_size_bytes, max_number_rows, max_number_cells): self._max_batch_size_bytes = max_batch_size_bytes self._max_number_rows = max_number_rows self._max_number_cells = max_number_cells def start_bundle(self): self._batch = MutationGroup() self._size_in_bytes = 0 self._rows = 0 self._cells = 0 def _reset_count(self): self._batch = MutationGroup() self._size_in_bytes = 0 self._rows = 0 self._cells = 0 def process(self, element): mg_info = element.info if mg_info['byte_size'] + self._size_in_bytes > self._max_batch_size_bytes \ or mg_info['cells'] + self._cells > self._max_number_cells \ or mg_info['rows'] + self._rows > self._max_number_rows: # Batch is full, output the batch and resetting the count. if self._batch: yield self._batch self._reset_count() self._batch.extend(element) # total byte size of the mutation group. self._size_in_bytes += mg_info['byte_size'] # total rows in the mutation group. self._rows += mg_info['rows'] # total cells in the mutation group. self._cells += mg_info['cells'] def finish_bundle(self): if self._batch is not None: yield window.GlobalWindows.windowed_value(self._batch) self._batch = None @with_input_types(MutationGroup) @with_output_types(MutationGroup) class _BatchableFilterFn(DoFn): """ Filters MutationGroups larger than the batch size to the output tagged with OUTPUT_TAG_UNBATCHABLE. """ OUTPUT_TAG_UNBATCHABLE = 'unbatchable' def __init__(self, max_batch_size_bytes, max_number_rows, max_number_cells): self._max_batch_size_bytes = max_batch_size_bytes self._max_number_rows = max_number_rows self._max_number_cells = max_number_cells self._batchable = None self._unbatchable = None def process(self, element): if element.primary().operation == WriteMutation._OPERATION_DELETE: # As delete mutations are not batchable. yield TaggedOutput(_BatchableFilterFn.OUTPUT_TAG_UNBATCHABLE, element) else: mg_info = element.info if mg_info['byte_size'] > self._max_batch_size_bytes \ or mg_info['cells'] > self._max_number_cells \ or mg_info['rows'] > self._max_number_rows: yield TaggedOutput(_BatchableFilterFn.OUTPUT_TAG_UNBATCHABLE, element) else: yield element class _WriteToSpannerDoFn(DoFn): def __init__(self, spanner_configuration): self._spanner_configuration = spanner_configuration self._db_instance = None self.batches = Metrics.counter(self.__class__, 'SpannerBatches') self.base_labels = { monitoring_infos.SERVICE_LABEL: 'Spanner', monitoring_infos.METHOD_LABEL: 'Write', monitoring_infos.SPANNER_PROJECT_ID: spanner_configuration.project, monitoring_infos.SPANNER_DATABASE_ID: spanner_configuration.database, } # table_id to metrics self.service_metrics = {} def _register_table_metric(self, table_id): if table_id in self.service_metrics: return database_id = self._spanner_configuration.database project_id = self._spanner_configuration.project resource = resource_identifiers.SpannerTable( project_id, database_id, table_id) labels = { **self.base_labels, monitoring_infos.RESOURCE_LABEL: resource, } if table_id: labels[monitoring_infos.SPANNER_TABLE_ID] = table_id service_call_metric = ServiceCallMetric( request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN, base_labels=labels) self.service_metrics[table_id] = service_call_metric def setup(self): spanner_client = Client(self._spanner_configuration.project) instance = spanner_client.instance(self._spanner_configuration.instance) self._db_instance = instance.database( self._spanner_configuration.database, pool=self._spanner_configuration.pool) def start_bundle(self): self.service_metrics = {} def process(self, element): self.batches.inc() try: with self._db_instance.batch() as b: for m in element: table_id = m.kwargs['table'] self._register_table_metric(table_id) if m.operation == WriteMutation._OPERATION_DELETE: batch_func = b.delete elif m.operation == WriteMutation._OPERATION_REPLACE: batch_func = b.replace elif m.operation == WriteMutation._OPERATION_INSERT_OR_UPDATE: batch_func = b.insert_or_update elif m.operation == WriteMutation._OPERATION_INSERT: batch_func = b.insert elif m.operation == WriteMutation._OPERATION_UPDATE: batch_func = b.update else: raise ValueError("Unknown operation action: %s" % m.operation) batch_func(**m.kwargs) except (ClientError, GoogleAPICallError) as e: for service_metric in self.service_metrics.values(): service_metric.call(str(e.code.value)) raise else: for service_metric in self.service_metrics.values(): service_metric.call('ok') @with_input_types(typing.Union[MutationGroup, _Mutator]) @with_output_types(MutationGroup) class _MakeMutationGroupsFn(DoFn): """ Make Mutation group object if the element is the instance of _Mutator. """ def process(self, element): if isinstance(element, MutationGroup): yield element elif isinstance(element, _Mutator): yield MutationGroup([element]) else: raise ValueError( "Invalid object type: %s. Object must be an instance of " "MutationGroup or WriteMutations" % str(element)) class _WriteGroup(PTransform): def __init__(self, max_batch_size_bytes, max_number_rows, max_number_cells): self._max_batch_size_bytes = max_batch_size_bytes self._max_number_rows = max_number_rows self._max_number_cells = max_number_cells def expand(self, pcoll): filter_batchable_mutations = ( pcoll | 'Making mutation groups' >> ParDo(_MakeMutationGroupsFn()) | 'Filtering Batchable Mutations' >> ParDo( _BatchableFilterFn( max_batch_size_bytes=self._max_batch_size_bytes, max_number_rows=self._max_number_rows, max_number_cells=self._max_number_cells)).with_outputs( _BatchableFilterFn.OUTPUT_TAG_UNBATCHABLE, main='batchable') ) batching_batchables = ( filter_batchable_mutations['batchable'] | ParDo( _BatchFn( max_batch_size_bytes=self._max_batch_size_bytes, max_number_rows=self._max_number_rows, max_number_cells=self._max_number_cells))) return (( batching_batchables, filter_batchable_mutations[_BatchableFilterFn.OUTPUT_TAG_UNBATCHABLE]) | 'Merging batchable and unbatchable' >> Flatten()) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/experimental/spannerio_read_it_test.py0000644000175100017510000003014715206113032026635 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import random import sys import unittest import uuid import pytest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # Protect against environments where spanner library is not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports # pylint: disable=unused-import try: from google.cloud import spanner from apache_beam.io.gcp import resource_identifiers from apache_beam.io.gcp.experimental.spannerio import ReadFromSpanner from apache_beam.io.gcp.experimental.spannerio import ReadOperation from apache_beam.io.gcp.experimental.spannerio import create_transaction from apache_beam.metrics import monitoring_infos from apache_beam.metrics.execution import MetricsEnvironment from apache_beam.metrics.metricbase import MetricName except ImportError: spanner = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports # pylint: enable=unused-import _LOGGER = logging.getLogger(__name__) _TEST_INSTANCE_ID = 'beam-test' @unittest.skipIf(spanner is None, 'GCP dependencies are not installed.') class SpannerReadIntegrationTest(unittest.TestCase): TEST_DATABASE = None _database_prefix = "pybeam-read-{}" _data = None _SPANNER_CLIENT = None _SPANNER_INSTANCE = None @classmethod def _generate_table_name(cls): cls.TEST_DATABASE = cls._database_prefix.format( ''.join(random.sample(uuid.uuid4().hex, 15))) return cls.TEST_DATABASE @classmethod def _create_database(cls): _LOGGER.info("Creating test database: %s" % cls.TEST_DATABASE) instance = cls._SPANNER_INSTANCE database = instance.database( cls.TEST_DATABASE, ddl_statements=[ """CREATE TABLE Users ( UserId INT64 NOT NULL, Key STRING(1024) ) PRIMARY KEY (UserId)""", ]) operation = database.create() _LOGGER.info("Creating database: Done! %s" % str(operation.result())) @classmethod def _add_dummy_entries(cls): _LOGGER.info("Dummy Data: Adding dummy data...") instance = cls._SPANNER_INSTANCE database = instance.database(cls.TEST_DATABASE) data = cls._data = [[x + 1, uuid.uuid4().hex] for x in range(200)] with database.batch() as batch: batch.insert(table='Users', columns=('UserId', 'Key'), values=data) @classmethod def setUpClass(cls): _LOGGER.info(".... PyVersion ---> %s" % str(sys.version)) _LOGGER.info(".... Setting up!") cls.test_pipeline = TestPipeline(is_integration_test=True) cls.args = cls.test_pipeline.get_full_options_as_args() cls.runner_name = type(cls.test_pipeline.runner).__name__ cls.project = cls.test_pipeline.get_option('project') cls.instance = ( cls.test_pipeline.get_option('instance') or _TEST_INSTANCE_ID) _ = cls._generate_table_name() spanner_client = cls._SPANNER_CLIENT = spanner.Client() _LOGGER.info(".... Spanner Client created!") cls._SPANNER_INSTANCE = spanner_client.instance(cls.instance) cls._create_database() cls._add_dummy_entries() _LOGGER.info("Spanner Read IT Setup Complete...") @pytest.mark.spannerio_it def test_read_via_table(self): _LOGGER.info("Spanner Read via table") with beam.Pipeline(argv=self.args) as p: r = p | ReadFromSpanner( self.project, self.instance, self.TEST_DATABASE, table="Users", columns=["UserId", "Key"]) assert_that(r, equal_to(self._data)) @pytest.mark.spannerio_it def test_read_via_sql(self): _LOGGER.info("Running Spanner via sql") with beam.Pipeline(argv=self.args) as p: r = p | ReadFromSpanner( self.project, self.instance, self.TEST_DATABASE, sql="select * from Users") assert_that(r, equal_to(self._data)) @pytest.mark.spannerio_it def test_transaction_table_metrics_ok_call(self): if 'DirectRunner' not in self.runner_name: raise unittest.SkipTest('This test only runs with DirectRunner.') MetricsEnvironment.process_wide_container().reset() with beam.Pipeline(argv=self.args) as p: transaction = ( p | create_transaction(self.project, self.instance, self.TEST_DATABASE)) r = p | ReadFromSpanner( self.project, self.instance, self.TEST_DATABASE, table="Users", columns=["UserId", "Key"], transaction=transaction) assert_that(r, equal_to(self._data)) self.verify_table_read_call_metric( self.project, self.TEST_DATABASE, 'Users', 'ok', 1) @pytest.mark.spannerio_it def test_transaction_table_metrics_error_call(self): if 'DirectRunner' not in self.runner_name: raise unittest.SkipTest('This test only runs with DirectRunner.') MetricsEnvironment.process_wide_container().reset() with self.assertRaises(Exception): p = beam.Pipeline(argv=self.args) transaction = ( p | create_transaction(self.project, self.instance, self.TEST_DATABASE)) _ = p | ReadFromSpanner( self.project, self.instance, self.TEST_DATABASE, table="INVALID_TABLE", columns=["UserId", "Key"], transaction=transaction) res = p.run() res.wait_until_finish() self.verify_table_read_call_metric( self.project, self.TEST_DATABASE, 'INVALID_TABLE', '404', 1) @pytest.mark.spannerio_it def test_transaction_sql_metrics_ok_call(self): if 'DirectRunner' not in self.runner_name: raise unittest.SkipTest('This test only runs with DirectRunner.') MetricsEnvironment.process_wide_container().reset() with beam.Pipeline(argv=self.args) as p: transaction = ( p | create_transaction(self.project, self.instance, self.TEST_DATABASE)) r = p | ReadFromSpanner( self.project, self.instance, self.TEST_DATABASE, sql="select * from Users", query_name='query-1', transaction=transaction) assert_that(r, equal_to(self._data)) self.verify_sql_read_call_metric( self.project, self.TEST_DATABASE, 'query-1', 'ok', 1) @pytest.mark.spannerio_it def test_transaction_sql_metrics_error_call(self): if 'DirectRunner' not in self.runner_name: raise unittest.SkipTest('This test only runs with DirectRunner.') MetricsEnvironment.process_wide_container().reset() with self.assertRaises(Exception): p = beam.Pipeline(argv=self.args) transaction = ( p | create_transaction(self.project, self.instance, self.TEST_DATABASE)) _ = p | ReadFromSpanner( self.project, self.instance, self.TEST_DATABASE, sql="select * from NonExistent", query_name="query-2", transaction=transaction) res = p.run() res.wait_until_finish() self.verify_sql_read_call_metric( self.project, self.TEST_DATABASE, 'query-2', '400', 1) @pytest.mark.spannerio_it def test_table_metrics_ok_call(self): if 'DirectRunner' not in self.runner_name: raise unittest.SkipTest('This test only runs with DirectRunner.') MetricsEnvironment.process_wide_container().reset() with beam.Pipeline(argv=self.args) as p: r = p | ReadFromSpanner( self.project, self.instance, self.TEST_DATABASE, table="Users", columns=["UserId", "Key"]) assert_that(r, equal_to(self._data)) self.verify_table_read_call_metric( self.project, self.TEST_DATABASE, 'Users', 'ok', 1) @pytest.mark.spannerio_it def test_table_metrics_error_call(self): if 'DirectRunner' not in self.runner_name: raise unittest.SkipTest('This test only runs with DirectRunner.') MetricsEnvironment.process_wide_container().reset() with self.assertRaises(Exception): p = beam.Pipeline(argv=self.args) _ = p | ReadFromSpanner( self.project, self.instance, self.TEST_DATABASE, table="INVALID_TABLE", columns=["UserId", "Key"]) res = p.run() res.wait_until_finish() self.verify_table_read_call_metric( self.project, self.TEST_DATABASE, 'INVALID_TABLE', '404', 1) @pytest.mark.spannerio_it def test_sql_metrics_ok_call(self): if 'DirectRunner' not in self.runner_name: raise unittest.SkipTest('This test only runs with DirectRunner.') MetricsEnvironment.process_wide_container().reset() with beam.Pipeline(argv=self.args) as p: r = p | ReadFromSpanner( self.project, self.instance, self.TEST_DATABASE, sql="select * from Users", query_name='query-1') assert_that(r, equal_to(self._data)) self.verify_sql_read_call_metric( self.project, self.TEST_DATABASE, 'query-1', 'ok', 1) @pytest.mark.spannerio_it def test_sql_metrics_error_call(self): if 'DirectRunner' not in self.runner_name: raise unittest.SkipTest('This test only runs with DirectRunner.') MetricsEnvironment.process_wide_container().reset() with self.assertRaises(Exception): p = beam.Pipeline(argv=self.args) _ = p | ReadFromSpanner( self.project, self.instance, self.TEST_DATABASE, sql="select * from NonExistent", query_name='query-2') res = p.run() res.wait_until_finish() self.verify_sql_read_call_metric( self.project, self.TEST_DATABASE, 'query-2', '400', 1) def verify_table_read_call_metric( self, project, database, table, status, count): resource = resource_identifiers.SpannerTable(project, database, table) labels = { monitoring_infos.SERVICE_LABEL: 'Spanner', monitoring_infos.METHOD_LABEL: 'Read', monitoring_infos.SPANNER_PROJECT_ID: project, monitoring_infos.SPANNER_DATABASE_ID: database, monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.SPANNER_TABLE_ID: table, monitoring_infos.STATUS_LABEL: status } metric_name = MetricName( None, None, urn=monitoring_infos.API_REQUEST_COUNT_URN, labels=labels) metric_value = MetricsEnvironment.process_wide_container().get_counter( metric_name).get_cumulative() self.assertEqual(metric_value, count) def verify_sql_read_call_metric( self, project, database, query_name, status, count): resource = resource_identifiers.SpannerSqlQuery(project, query_name) labels = { monitoring_infos.SERVICE_LABEL: 'Spanner', monitoring_infos.METHOD_LABEL: 'Read', monitoring_infos.SPANNER_PROJECT_ID: project, monitoring_infos.SPANNER_DATABASE_ID: database, monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.SPANNER_QUERY_NAME: query_name, monitoring_infos.STATUS_LABEL: status } metric_name = MetricName( None, None, urn=monitoring_infos.API_REQUEST_COUNT_URN, labels=labels) metric_value = MetricsEnvironment.process_wide_container().get_counter( metric_name).get_cumulative() self.assertEqual(metric_value, count) @classmethod def tearDownClass(cls): # drop the testing database after the tests database = cls._SPANNER_INSTANCE.database(cls.TEST_DATABASE) database.drop() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/experimental/spannerio_read_perf_test.py0000644000175100017510000001272015206113032027152 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A performance test for reading data from a Spanner database table. Besides of the standard options, there are options with special meaning: * spanner_instance - Spanner Instance ID. * spanner_database - Spanner Database ID. The table will be created and populated with data from Synthetic Source if it does not exist. * input_options - options for Synthetic Source: num_records - number of rows to be inserted, value_size - the length of a single row, key_size - required option, but its value has no meaning. Example test run on DataflowRunner: python -m apache_beam.io.gcp.experimental.spannerio_read_perf_test \ --test-pipeline-options=" --runner=TestDataflowRunner --project='...' --region='...' --temp_location='gs://...' --sdk_location=build/apache-beam.tar.gz --publish_to_big_query=true --metrics_dataset='...' --metrics_table='...' --spanner_instance='...' --spanner_database='...' --input_options='{ \"num_records\": 10, \"key_size\": 1, \"value_size\": 1024 }'" This setup will result in a table of 1MB size. """ from __future__ import absolute_import import logging from apache_beam import FlatMap from apache_beam import Map from apache_beam import ParDo from apache_beam.io import Read from apache_beam.io.gcp.experimental.spannerio import ReadFromSpanner from apache_beam.io.gcp.experimental.spannerio import WriteToSpanner from apache_beam.testing.load_tests.load_test import LoadTest from apache_beam.testing.load_tests.load_test_metrics_utils import CountMessages from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime from apache_beam.testing.synthetic_pipeline import SyntheticSource from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.combiners import Count # pylint: disable=wrong-import-order, wrong-import-position try: from google.api_core.exceptions import AlreadyExists from google.cloud import spanner except ImportError: spanner = None AlreadyExists = None # pylint: enable=wrong-import-order, wrong-import-position class SpannerReadPerfTest(LoadTest): def __init__(self): super().__init__() self.project = self.pipeline.get_option('project') self.spanner_instance = self.pipeline.get_option('spanner_instance') self.spanner_database = self.pipeline.get_option('spanner_database') self._init_setup() def _create_database(self): spanner_client = spanner.Client() instance = spanner_client.instance(self.spanner_instance) database = instance.database( self.spanner_database, ddl_statements=[ """CREATE TABLE test_data ( id STRING(99) NOT NULL, data BYTES(MAX) NOT NULL ) PRIMARY KEY (id)""", ]) database.create() def _init_setup(self): """Checks if a spanner database exists and creates it if not.""" try: self._create_database() self._create_input_data() except AlreadyExists: # pass if the database already exists pass def _create_input_data(self): """ Runs an additional pipeline which creates test data and waits for its completion. """ def format_record(record): import base64 return base64.b64encode(record[1]) def make_insert_mutations(element): import uuid from apache_beam.io.gcp.experimental.spannerio import WriteMutation ins_mutation = WriteMutation.insert( table='test_data', columns=('id', 'data'), values=[(str(uuid.uuid1()), element)]) return [ins_mutation] with TestPipeline() as p: ( # pylint: disable=expression-not-assigned p | 'Produce rows' >> Read( SyntheticSource(self.parse_synthetic_source_options())) | 'Format' >> Map(format_record) | 'Make mutations' >> FlatMap(make_insert_mutations) | 'Write to Spanner' >> WriteToSpanner( project_id=self.project, instance_id=self.spanner_instance, database_id=self.spanner_database, max_batch_size_bytes=5120)) def test(self): output = ( self.pipeline | 'Read from Spanner' >> ReadFromSpanner( self.project, self.spanner_instance, self.spanner_database, sql="select data from test_data") | 'Count messages' >> ParDo(CountMessages(self.metrics_namespace)) | 'Measure time' >> ParDo(MeasureTime(self.metrics_namespace)) | 'Count' >> Count.Globally()) assert_that(output, equal_to([self.input_options['num_records']])) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) SpannerReadPerfTest().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/experimental/spannerio_test.py0000644000175100017510000006532515206113032025154 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import datetime import logging import random import string import typing import unittest import mock import apache_beam as beam from apache_beam.metrics.metric import MetricsFilter from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # Protect against environments where spanner library is not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports # pylint: disable=unused-import try: from google.cloud import spanner from apache_beam.io.gcp import resource_identifiers from apache_beam.io.gcp.experimental.spannerio import MutationGroup from apache_beam.io.gcp.experimental.spannerio import ReadFromSpanner from apache_beam.io.gcp.experimental.spannerio import ReadOperation from apache_beam.io.gcp.experimental.spannerio import WriteMutation from apache_beam.io.gcp.experimental.spannerio import WriteToSpanner from apache_beam.io.gcp.experimental.spannerio import _BatchFn from apache_beam.io.gcp.experimental.spannerio import create_transaction from apache_beam.metrics import monitoring_infos from apache_beam.metrics.execution import MetricsEnvironment from apache_beam.metrics.metricbase import MetricName except ImportError: spanner = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports # pylint: enable=unused-import MAX_DB_NAME_LENGTH = 30 TEST_PROJECT_ID = 'apache-beam-testing' TEST_INSTANCE_ID = 'beam-test' TEST_DATABASE_PREFIX = 'spanner-testdb-' FAKE_TRANSACTION_INFO = {"session_id": "qwerty", "transaction_id": "qwerty"} FAKE_ROWS = [[1, 'Alice'], [2, 'Bob'], [3, 'Carl'], [4, 'Dan'], [5, 'Evan'], [6, 'Floyd']] def _generate_database_name(): mask = string.ascii_lowercase + string.digits length = MAX_DB_NAME_LENGTH - 1 - len(TEST_DATABASE_PREFIX) return TEST_DATABASE_PREFIX + ''.join( random.choice(mask) for i in range(length)) def _generate_test_data(): mask = string.ascii_lowercase + string.digits length = 100 return [( 'users', ['Key', 'Value'], [(x, ''.join(random.choice(mask) for _ in range(length))) for x in range(1, 5)])] @unittest.skipIf(spanner is None, 'GCP dependencies are not installed.') @mock.patch('apache_beam.io.gcp.experimental.spannerio.Client') @mock.patch('apache_beam.io.gcp.experimental.spannerio.BatchSnapshot') class SpannerReadTest(unittest.TestCase): def test_read_with_query_batch( self, mock_batch_snapshot_class, mock_client_class): mock_snapshot_instance = mock.MagicMock() mock_snapshot_instance.generate_query_batches.return_value = [{ 'query': { 'sql': 'SELECT * FROM users' }, 'partition': 'test_partition' } for _ in range(3)] mock_snapshot_instance.to_dict.return_value = {} mock_batch_snapshot_instance = mock.MagicMock() # Prepare process_query_batch return results for three pipelines mock_batch_snapshot_instance.process_query_batch.side_effect = [ FAKE_ROWS[0:2], FAKE_ROWS[2:4], FAKE_ROWS[4:] ] * 3 mock_client_class.return_value.instance.return_value.database.return_value \ .batch_snapshot.return_value = mock_snapshot_instance mock_batch_snapshot_class.from_dict.return_value \ = mock_batch_snapshot_instance ro = [ReadOperation.query("Select * from users")] with TestPipeline() as pipeline: read = ( pipeline | 'read' >> ReadFromSpanner( TEST_PROJECT_ID, TEST_INSTANCE_ID, _generate_database_name(), sql="SELECT * FROM users")) assert_that(read, equal_to(FAKE_ROWS), label='checkRead') with TestPipeline() as pipeline: readall = ( pipeline | 'read all' >> ReadFromSpanner( TEST_PROJECT_ID, TEST_INSTANCE_ID, _generate_database_name(), read_operations=ro)) assert_that(readall, equal_to(FAKE_ROWS), label='checkReadAll') with TestPipeline() as pipeline: readpipeline = ( pipeline | 'create reads' >> beam.Create(ro) | 'reads' >> ReadFromSpanner( TEST_PROJECT_ID, TEST_INSTANCE_ID, _generate_database_name())) assert_that(readpipeline, equal_to(FAKE_ROWS), label='checkReadPipeline') # three pipelines self.assertEqual( mock_snapshot_instance.generate_query_batches.call_count, 3) # three pipelines, each called three times self.assertEqual( mock_batch_snapshot_instance.process_query_batch.call_count, 3 * 3) def test_read_with_table_batch( self, mock_batch_snapshot_class, mock_client_class): mock_snapshot_instance = mock.MagicMock() mock_snapshot_instance.generate_read_batches.return_value = [{ 'read': { 'table': 'users', 'keyset': { 'all': True }, 'columns': ['Key', 'Value'], 'index': '' }, 'partition': 'test_partition' } for _ in range(3)] mock_snapshot_instance.to_dict.return_value = {} mock_batch_snapshot_instance = mock.MagicMock() # Prepare process_read_batch return results for three pipelines mock_batch_snapshot_instance.process_read_batch.side_effect = [ FAKE_ROWS[0:2], FAKE_ROWS[2:4], FAKE_ROWS[4:] ] * 3 mock_client_class.return_value.instance.return_value.database.return_value \ .batch_snapshot.return_value = mock_snapshot_instance mock_batch_snapshot_class.from_dict.return_value \ = mock_batch_snapshot_instance ro = [ReadOperation.table("users", ["Key", "Value"])] with TestPipeline() as pipeline: read = ( pipeline | 'read' >> ReadFromSpanner( TEST_PROJECT_ID, TEST_INSTANCE_ID, _generate_database_name(), table="users", columns=["Key", "Value"])) assert_that(read, equal_to(FAKE_ROWS), label='checkRead') with TestPipeline() as pipeline: readall = ( pipeline | 'read all' >> ReadFromSpanner( TEST_PROJECT_ID, TEST_INSTANCE_ID, _generate_database_name(), read_operations=ro)) assert_that(readall, equal_to(FAKE_ROWS), label='checkReadAll') with TestPipeline() as pipeline: readpipeline = ( pipeline | 'create reads' >> beam.Create(ro) | 'reads' >> ReadFromSpanner( TEST_PROJECT_ID, TEST_INSTANCE_ID, _generate_database_name())) assert_that(readpipeline, equal_to(FAKE_ROWS), label='checkReadPipeline') # three pipelines self.assertEqual(mock_snapshot_instance.generate_read_batches.call_count, 3) # three pipelines, each called three times self.assertEqual( mock_batch_snapshot_instance.process_read_batch.call_count, 3 * 3) with TestPipeline() as pipeline, self.assertRaises(ValueError): # Test the exception raised at pipeline construction time, when user # passes the read operations in the constructor and also in the pipeline _ = ( pipeline | 'reads error' >> ReadFromSpanner( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), table="users")) def test_read_with_index(self, mock_batch_snapshot_class, mock_client_class): mock_snapshot_instance = mock.MagicMock() mock_snapshot_instance.generate_read_batches.return_value = [{ 'read': { 'table': 'users', 'keyset': { 'all': True }, 'columns': ['Key', 'Value'], 'index': '' }, 'partition': 'test_partition' } for _ in range(3)] mock_batch_snapshot_instance = mock.MagicMock() # Prepare process_read_batch return results for three pipelines mock_batch_snapshot_instance.process_read_batch.side_effect = [ FAKE_ROWS[0:2], FAKE_ROWS[2:4], FAKE_ROWS[4:] ] * 3 mock_snapshot_instance.to_dict.return_value = {} mock_client_class.return_value.instance.return_value.database.return_value \ .batch_snapshot.return_value = mock_snapshot_instance mock_batch_snapshot_class.from_dict.return_value \ = mock_batch_snapshot_instance ro = [ReadOperation.table("users", ["Key", "Value"], index="Key")] with TestPipeline() as pipeline: read = ( pipeline | 'read' >> ReadFromSpanner( TEST_PROJECT_ID, TEST_INSTANCE_ID, _generate_database_name(), table="users", columns=["Key", "Value"])) assert_that(read, equal_to(FAKE_ROWS), label='checkRead') with TestPipeline() as pipeline: readall = ( pipeline | 'read all' >> ReadFromSpanner( TEST_PROJECT_ID, TEST_INSTANCE_ID, _generate_database_name(), read_operations=ro)) assert_that(readall, equal_to(FAKE_ROWS), label='checkReadAll') with TestPipeline() as pipeline: readpipeline = ( pipeline | 'create reads' >> beam.Create(ro) | 'reads' >> ReadFromSpanner( TEST_PROJECT_ID, TEST_INSTANCE_ID, _generate_database_name())) assert_that(readpipeline, equal_to(FAKE_ROWS), label='checkReadPipeline') # three pipelines self.assertEqual(mock_snapshot_instance.generate_read_batches.call_count, 3) # three pipelines, each called three times self.assertEqual( mock_batch_snapshot_instance.process_read_batch.call_count, 3 * 3) with TestPipeline() as pipeline, self.assertRaises(ValueError): # Test the exception raised at pipeline construction time, when user # passes the read operations in the constructor and also in the pipeline. _ = ( pipeline | 'reads error' >> ReadFromSpanner( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), table="users")) def test_read_with_transaction( self, mock_batch_snapshot_class, mock_client_class): mock_snapshot_instance = mock.MagicMock() mock_snapshot_instance.to_dict.return_value = FAKE_TRANSACTION_INFO mock_transaction_instance = mock.MagicMock() mock_transaction_instance.execute_sql.return_value = FAKE_ROWS mock_transaction_instance.read.return_value = FAKE_ROWS mock_client_class.return_value.instance.return_value.database.return_value \ .batch_snapshot.return_value = mock_snapshot_instance mock_client_class.return_value.instance.return_value.database.return_value \ .session.return_value.transaction.return_value.__enter__.return_value \ = mock_transaction_instance ro = [ReadOperation.query("Select * from users")] with TestPipeline() as p: transaction = ( p | create_transaction( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), exact_staleness=datetime.timedelta(seconds=10))) read_query = ( p | 'with query' >> ReadFromSpanner( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), transaction=transaction, sql="Select * from users")) assert_that(read_query, equal_to(FAKE_ROWS), label='checkQuery') read_table = ( p | 'with table' >> ReadFromSpanner( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), transaction=transaction, table="users", columns=["Key", "Value"])) assert_that(read_table, equal_to(FAKE_ROWS), label='checkTable') read_indexed_table = ( p | 'with index' >> ReadFromSpanner( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), transaction=transaction, table="users", index="Key", columns=["Key", "Value"])) assert_that( read_indexed_table, equal_to(FAKE_ROWS), label='checkTableIndex') read = ( p | 'read all' >> ReadFromSpanner( TEST_PROJECT_ID, TEST_INSTANCE_ID, _generate_database_name(), transaction=transaction, read_operations=ro)) assert_that(read, equal_to(FAKE_ROWS), label='checkReadAll') read_pipeline = ( p | 'create read operations' >> beam.Create(ro) | 'reads' >> ReadFromSpanner( TEST_PROJECT_ID, TEST_INSTANCE_ID, _generate_database_name(), transaction=transaction)) assert_that(read_pipeline, equal_to(FAKE_ROWS), label='checkReadPipeline') # transaction setup once self.assertEqual(mock_snapshot_instance.to_dict.call_count, 1) # three pipelines called execute_sql self.assertEqual(mock_transaction_instance.execute_sql.call_count, 3) # two pipelines called read self.assertEqual(mock_transaction_instance.read.call_count, 2) with TestPipeline() as p, self.assertRaises(ValueError): # Test the exception raised at pipeline construction time, when user # passes the read operations in the constructor and also in the pipeline. transaction = ( p | create_transaction( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), exact_staleness=datetime.timedelta(seconds=10))) _ = ( p | 'create read operations2' >> beam.Create(ro) | 'reads with error' >> ReadFromSpanner( TEST_PROJECT_ID, TEST_INSTANCE_ID, _generate_database_name(), transaction=transaction, read_operations=ro)) def test_invalid_transaction( self, mock_batch_snapshot_class, mock_client_class): # test exception raises at pipeline execution time error_string = "Invalid transaction object" with self.assertRaisesRegex(Exception, error_string), TestPipeline() as p: transaction = ( p | beam.Create([{ "invalid": "transaction" }]).with_output_types(typing.Any)) _ = ( p | 'with query' >> ReadFromSpanner( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), transaction=transaction, sql="Select * from users")) def test_display_data(self, *args): dd_sql = ReadFromSpanner( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), sql="Select * from users").display_data() dd_table = ReadFromSpanner( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), table="users", columns=['id', 'name']).display_data() dd_transaction = ReadFromSpanner( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), table="users", columns=['id', 'name'], transaction={ "transaction_id": "test123", "session_id": "test456" }).display_data() self.assertTrue("sql" in dd_sql) self.assertTrue("table" in dd_table) self.assertTrue("table" in dd_transaction) self.assertTrue("transaction" in dd_transaction) @unittest.skipIf(spanner is None, 'GCP dependencies are not installed.') @mock.patch('apache_beam.io.gcp.experimental.spannerio.Client') @mock.patch('google.cloud.spanner_v1.database.BatchCheckout') class SpannerWriteTest(unittest.TestCase): def test_spanner_write(self, mock_batch_snapshot_class, mock_batch_checkout): ks = spanner.KeySet(keys=[[1233], [1234]]) mutations = [ WriteMutation.delete("roles", ks), WriteMutation.insert( "roles", ("key", "rolename"), [('1233', "mutations-inset-1233")]), WriteMutation.insert( "roles", ("key", "rolename"), [('1234', "mutations-inset-1234")]), WriteMutation.update( "roles", ("key", "rolename"), [('1234', "mutations-inset-1233-updated")]), ] # TODO(https://github.com/apache/beam/issues/34549): This test relies on # metrics filtering which doesn't work on Prism yet because Prism renames # steps (e.g. "Do" becomes "ref_AppliedPTransform_Do_7"). # https://github.com/apache/beam/blob/5f9cd73b7c9a2f37f83971ace3a399d633201dd1/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py#L1590 p = TestPipeline('FnApiRunner') _ = ( p | beam.Create(mutations) | WriteToSpanner( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), max_batch_size_bytes=1024)) res = p.run() res.wait_until_finish() metric_results = res.metrics().query( MetricsFilter().with_name("SpannerBatches")) batches_counter = metric_results['counters'][0] self.assertEqual(batches_counter.committed, 2) self.assertEqual(batches_counter.attempted, 2) def test_spanner_bundles_size( self, mock_batch_snapshot_class, mock_batch_checkout): ks = spanner.KeySet(keys=[[1233], [1234]]) mutations = [ WriteMutation.delete("roles", ks), WriteMutation.insert( "roles", ("key", "rolename"), [('1234', "mutations-inset-1234")]) ] * 50 # TODO(https://github.com/apache/beam/issues/34549): This test relies on # metrics filtering which doesn't work on Prism yet because Prism renames # steps (e.g. "Do" becomes "ref_AppliedPTransform_Do_7"). # https://github.com/apache/beam/blob/5f9cd73b7c9a2f37f83971ace3a399d633201dd1/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py#L1590 p = TestPipeline('FnApiRunner') _ = ( p | beam.Create(mutations) | WriteToSpanner( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), max_batch_size_bytes=1024)) res = p.run() res.wait_until_finish() metric_results = res.metrics().query( MetricsFilter().with_name('SpannerBatches')) batches_counter = metric_results['counters'][0] self.assertEqual(batches_counter.committed, 53) self.assertEqual(batches_counter.attempted, 53) def test_spanner_write_mutation_groups( self, mock_batch_snapshot_class, mock_batch_checkout): ks = spanner.KeySet(keys=[[1233], [1234]]) mutation_groups = [ MutationGroup([ WriteMutation.insert( "roles", ("key", "rolename"), [('9001233', "mutations-inset-1233")]), WriteMutation.insert( "roles", ("key", "rolename"), [('9001234', "mutations-inset-1234")]) ]), MutationGroup([ WriteMutation.update( "roles", ("key", "rolename"), [('9001234', "mutations-inset-9001233-updated")]) ]), MutationGroup([WriteMutation.delete("roles", ks)]) ] # TODO(https://github.com/apache/beam/issues/34549): This test relies on # metrics filtering which doesn't work on Prism yet because Prism renames # steps (e.g. "Do" becomes "ref_AppliedPTransform_Do_7"). # https://github.com/apache/beam/blob/5f9cd73b7c9a2f37f83971ace3a399d633201dd1/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py#L1590 p = TestPipeline('FnApiRunner') _ = ( p | beam.Create(mutation_groups) | WriteToSpanner( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), max_batch_size_bytes=100)) res = p.run() res.wait_until_finish() metric_results = res.metrics().query( MetricsFilter().with_name('SpannerBatches')) batches_counter = metric_results['counters'][0] self.assertEqual(batches_counter.committed, 3) self.assertEqual(batches_counter.attempted, 3) def test_batch_byte_size( self, mock_batch_snapshot_class, mock_batch_checkout): # each mutation group byte size is 58 bytes. mutation_group = [ MutationGroup([ WriteMutation.insert( "roles", ("key", "rolename"), [('1234', "mutations-inset-1234")]) ]) ] * 50 with TestPipeline() as p: # the total 50 mutation group size will be 2900 (58 * 50) # if we want to make two batches, so batch size should be 1450 (2900 / 2) # and each bach should contains 25 mutations. res = ( p | beam.Create(mutation_group) | beam.ParDo( _BatchFn( max_batch_size_bytes=1450, max_number_rows=50, max_number_cells=500)) | beam.Map(lambda x: len(x))) assert_that(res, equal_to([25] * 2)) def test_batch_disable(self, mock_batch_snapshot_class, mock_batch_checkout): mutation_group = [ MutationGroup([ WriteMutation.insert( "roles", ("key", "rolename"), [('1234', "mutations-inset-1234")]) ]) ] * 4 with TestPipeline() as p: # to disable to batching, we need to set any of the batching parameters # either to lower value or zero res = ( p | beam.Create(mutation_group) | beam.ParDo( _BatchFn( max_batch_size_bytes=1450, max_number_rows=0, max_number_cells=500)) | beam.Map(lambda x: len(x))) assert_that(res, equal_to([1] * 4)) def test_batch_max_rows(self, mock_batch_snapshot_class, mock_batch_checkout): mutation_group = [ MutationGroup([ WriteMutation.insert( "roles", ("key", "rolename"), [ ('1234', "mutations-inset-1234"), ('1235', "mutations-inset-1235"), ]) ]) ] * 50 with TestPipeline() as p: # There are total 50 mutation groups, each contains two rows. # The total number of rows will be 100 (50 * 2). # If each batch contains 10 rows max then batch count should be 10 # (contains 5 mutation groups each). res = ( p | beam.Create(mutation_group) | beam.ParDo( _BatchFn( max_batch_size_bytes=1048576, max_number_rows=10, max_number_cells=500)) | beam.Map(lambda x: len(x))) assert_that(res, equal_to([5] * 10)) def test_batch_max_cells( self, mock_batch_snapshot_class, mock_batch_checkout): mutation_group = [ MutationGroup([ WriteMutation.insert( "roles", ("key", "rolename"), [ ('1234', "mutations-inset-1234"), ('1235', "mutations-inset-1235"), ]) ]) ] * 50 with TestPipeline() as p: # There are total 50 mutation groups, each contains two rows (or 4 cells). # The total number of cells will be 200 (50 groups * 4 cells). # If each batch contains 50 cells max then batch count should be 5. # 4 batches contains 12 mutations groups and the fifth batch should be # consists of 2 mutation group element. # No. of mutations groups per batch = Max Cells / Cells per mutation group # total_batches = Total Number of Cells / Max Cells res = ( p | beam.Create(mutation_group) | beam.ParDo( _BatchFn( max_batch_size_bytes=1048576, max_number_rows=500, max_number_cells=50)) | beam.Map(lambda x: len(x))) # Accept both optimal and suboptimal batching patterns due to Beam's # non-deterministic execution # Optimal: [12, 12, 12, 12, 2] - ideal batching without bundle # fragmentation # Suboptimal: [12, 12, 1, 12, 1, 12] - caused by bundle boundaries # interrupting batching optimal_batch_sizes = [12, 12, 12, 12, 2] suboptimal_batch_sizes = [12, 12, 1, 12, 1, 12] def validate_batching(actual_batch_sizes): actual_sorted = sorted(actual_batch_sizes) optimal_sorted = sorted(optimal_batch_sizes) suboptimal_sorted = sorted(suboptimal_batch_sizes) # Verify total element count first total_elements = sum(actual_batch_sizes) if total_elements != 50: raise AssertionError( f"Expected total of 50 elements, got {total_elements}") # Accept either optimal or known suboptimal pattern if actual_sorted == optimal_sorted: # Optimal batching achieved return True elif actual_sorted == suboptimal_sorted: # Known suboptimal pattern due to bundle fragmentation - acceptable return True else: raise AssertionError( f"Expected batch sizes {optimal_sorted} (optimal) or " f"{suboptimal_sorted} (suboptimal), got {actual_sorted}") assert_that(res, validate_batching) def test_write_mutation_error(self, *args): with self.assertRaises(ValueError): # since `WriteMutation` only accept one operation. WriteMutation(insert="table-name", update="table-name") def test_display_data(self, *args): data = WriteToSpanner( project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), max_batch_size_bytes=1024).display_data() self.assertTrue("project_id" in data) self.assertTrue("instance_id" in data) self.assertTrue("pool" in data) self.assertTrue("database" in data) self.assertTrue("batch_size" in data) self.assertTrue("max_number_rows" in data) self.assertTrue("max_number_cells" in data) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/experimental/spannerio_write_it_test.py0000644000175100017510000002300515206113032027047 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the 'License'); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an 'AS IS' BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import random import unittest import uuid import pytest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline # Protect against environments where spanner library is not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports # pylint: disable=unused-import try: from google.api_core.exceptions import NotFound from google.cloud import spanner from apache_beam.io.gcp import resource_identifiers from apache_beam.io.gcp.experimental.spannerio import MutationGroup from apache_beam.io.gcp.experimental.spannerio import WriteMutation from apache_beam.io.gcp.experimental.spannerio import WriteToSpanner from apache_beam.metrics import monitoring_infos from apache_beam.metrics.execution import MetricsEnvironment from apache_beam.metrics.metricbase import MetricName except ImportError: spanner = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports # pylint: enable=unused-import _LOGGER = logging.getLogger(__name__) _TEST_INSTANCE_ID = 'beam-test' @unittest.skipIf(spanner is None, 'GCP dependencies are not installed.') class SpannerWriteIntegrationTest(unittest.TestCase): TEST_DATABASE = None _database_prefix = 'pybeam-write-{}' _SPANNER_CLIENT = None _SPANNER_INSTANCE = None @classmethod def _generate_table_name(cls): cls.TEST_DATABASE = cls._database_prefix.format( ''.join(random.sample(uuid.uuid4().hex, 15))) return cls.TEST_DATABASE @classmethod def _create_database(cls): _LOGGER.info('Creating test database: %s' % cls.TEST_DATABASE) instance = cls._SPANNER_INSTANCE database = instance.database( cls.TEST_DATABASE, ddl_statements=[ '''CREATE TABLE Users ( UserId STRING(256) NOT NULL, Key STRING(1024) ) PRIMARY KEY (UserId)''', '''CREATE TABLE Albums ( AlbumId STRING(256) NOT NULL, Name STRING(1024) ) PRIMARY KEY (AlbumId) ''' ]) operation = database.create() _LOGGER.info('Creating database: Done! %s' % str(operation.result())) @classmethod def _count_data(cls, prefix): instance = cls._SPANNER_INSTANCE database = instance.database(cls.TEST_DATABASE) count = None with database.snapshot() as snapshot: results = snapshot.execute_sql( 'SELECT COUNT(*) FROM Users WHERE UserId ' 'LIKE "{}%"'.format(prefix)) try: count = list(results)[0][0] except IndexError: raise ValueError( "Spanner Count rows results not found for %s." % prefix) return count @classmethod def setUpClass(cls): _LOGGER.info('.... Setting up!') cls.test_pipeline = TestPipeline(is_integration_test=True) cls.args = cls.test_pipeline.get_full_options_as_args() cls.runner_name = type(cls.test_pipeline.runner).__name__ cls.project = cls.test_pipeline.get_option('project') cls.instance = ( cls.test_pipeline.get_option('instance') or _TEST_INSTANCE_ID) _ = cls._generate_table_name() spanner_client = cls._SPANNER_CLIENT = spanner.Client() _LOGGER.info('.... Spanner Client created!') cls._SPANNER_INSTANCE = spanner_client.instance(cls.instance) cls._create_database() _LOGGER.info('Spanner Write IT Setup Complete...') @pytest.mark.spannerio_it def test_write_batches(self): _prefex = 'test_write_batches' mutations = [ WriteMutation.insert( 'Users', ('UserId', 'Key'), [(_prefex + '1', _prefex + 'inset-1')]), WriteMutation.insert( 'Users', ('UserId', 'Key'), [(_prefex + '2', _prefex + 'inset-2')]), WriteMutation.insert( 'Users', ('UserId', 'Key'), [(_prefex + '3', _prefex + 'inset-3')]), WriteMutation.insert( 'Users', ('UserId', 'Key'), [(_prefex + '4', _prefex + 'inset-4')]) ] p = beam.Pipeline(argv=self.args) _ = ( p | beam.Create(mutations) | WriteToSpanner( project_id=self.project, instance_id=self.instance, database_id=self.TEST_DATABASE, max_batch_size_bytes=250)) res = p.run() res.wait_until_finish() self.assertEqual(self._count_data(_prefex), len(mutations)) @pytest.mark.spannerio_it def test_spanner_update(self): _prefex = 'test_update' # adding data to perform test instance = self._SPANNER_INSTANCE database = instance.database(self.TEST_DATABASE) data = [ (_prefex + '1', _prefex + 'inset-1'), (_prefex + '2', _prefex + 'inset-2'), (_prefex + '3', _prefex + 'inset-3'), ] with database.batch() as batch: batch.insert(table='Users', columns=('UserId', 'Key'), values=data) mutations_update = [ WriteMutation.update( 'Users', ('UserId', 'Key'), [(_prefex + '1', _prefex + 'update-1')]), WriteMutation.update( 'Users', ('UserId', 'Key'), [(_prefex + '2', _prefex + 'update-2')]), WriteMutation.delete('Users', spanner.KeySet(keys=[[_prefex + '3']])) ] p = beam.Pipeline(argv=self.args) _ = ( p | beam.Create(mutations_update) | WriteToSpanner( project_id=self.project, instance_id=self.instance, database_id=self.TEST_DATABASE)) res = p.run() res.wait_until_finish() self.assertEqual(self._count_data(_prefex), 2) @pytest.mark.spannerio_it def test_spanner_error(self): mutations_update = [ WriteMutation.update( 'Users', ('UserId', 'Key'), [('INVALD_ID', 'Error-error')]), ] with self.assertRaises(Exception): p = beam.Pipeline(argv=self.args) _ = ( p | beam.Create(mutations_update) | WriteToSpanner( project_id=self.project, instance_id=self.instance, database_id=self.TEST_DATABASE)) p.run() @pytest.mark.spannerio_it def test_metrics_ok_call(self): if 'DirectRunner' not in self.runner_name: raise unittest.SkipTest('This test only runs with DirectRunner.') MetricsEnvironment.process_wide_container().reset() _prefix = 'test_write_batches' mutations = [ WriteMutation.insert( 'Albums', ('AlbumId', 'Name'), [(_prefix + '1', _prefix + 'inset-1')]), WriteMutation.insert( 'Albums', ('AlbumId', 'Name'), [(_prefix + '2', _prefix + 'inset-2')]), ] p = beam.Pipeline(argv=self.args) _ = ( p | beam.Create(mutations) | WriteToSpanner( project_id=self.project, instance_id=self.instance, database_id=self.TEST_DATABASE)) res = p.run() res.wait_until_finish() self.verify_write_call_metric( self.project, self.TEST_DATABASE, 'Albums', 'ok', 1) @pytest.mark.spannerio_it def test_metrics_error_call(self): if 'DirectRunner' not in self.runner_name: raise unittest.SkipTest('This test only runs with DirectRunner.') MetricsEnvironment.process_wide_container().reset() _prefix = 'test_write_batches' mutations = [ WriteMutation.insert( 'Albums', ('AlbumId', 'Name'), [(_prefix + '3', _prefix + 'inset-3')]), WriteMutation.insert( 'Albums', ('AlbumId', 'Name'), [(_prefix + '3', _prefix + 'inset-3')]), ] with self.assertRaises(Exception): p = beam.Pipeline(argv=self.args) _ = ( p | beam.Create(mutations) | WriteToSpanner( project_id=self.project, instance_id=self.instance, database_id=self.TEST_DATABASE)) res = p.run() res.wait_until_finish() self.verify_write_call_metric( self.project, self.TEST_DATABASE, 'Albums', '400', 1) def verify_write_call_metric(self, project, database, table, status, count): resource = resource_identifiers.SpannerTable(project, database, table) labels = { monitoring_infos.SERVICE_LABEL: 'Spanner', monitoring_infos.METHOD_LABEL: 'Write', monitoring_infos.SPANNER_PROJECT_ID: project, monitoring_infos.SPANNER_DATABASE_ID: database, monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.SPANNER_TABLE_ID: table, monitoring_infos.STATUS_LABEL: status } metric_name = MetricName( None, None, urn=monitoring_infos.API_REQUEST_COUNT_URN, labels=labels) metric_value = MetricsEnvironment.process_wide_container().get_counter( metric_name).get_cumulative() self.assertEqual(metric_value, count) @classmethod def tearDownClass(cls): # drop the testing database after the tests database = cls._SPANNER_INSTANCE.database(cls.TEST_DATABASE) database.drop() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/experimental/spannerio_write_perf_test.py0000644000175100017510000001146115206113032027372 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A pipeline that writes data from Synthetic Source to a Spanner. Besides of the standard options, there are options with special meaning: * spanner_instance - Spanner Instance ID. * spanner_database - Spanner Database ID. * input_options - options for Synthetic Source: num_records - number of rows to be inserted, value_size - the length of a single row, key_size - required option, but its value has no meaning. Example test run on DataflowRunner: python -m apache_beam.io.gcp.experimental.spannerio_write_perf_test \ --test-pipeline-options=" --runner=TestDataflowRunner --project='...' --region='...' --temp_location='gs://...' --sdk_location=build/apache-beam.tar.gz --publish_to_big_query=true --metrics_dataset='...' --metrics_table='...' --spanner_instance='...' --spanner_database='...' --input_options='{ \"num_records\": 10, \"key_size\": 1, \"value_size\": 1024 }'" This setup will result in a table of 1MB size. """ from __future__ import absolute_import import logging import random import uuid from apache_beam import FlatMap from apache_beam import Map from apache_beam import ParDo from apache_beam.io import Read from apache_beam.io.gcp.experimental.spannerio import WriteToSpanner from apache_beam.testing.load_tests.load_test import LoadTest from apache_beam.testing.load_tests.load_test_metrics_utils import CountMessages from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime from apache_beam.testing.synthetic_pipeline import SyntheticSource # pylint: disable=wrong-import-order, wrong-import-position try: from google.cloud import spanner except ImportError: spanner = None # pylint: enable=wrong-import-order, wrong-import-position class SpannerWritePerfTest(LoadTest): TEST_DATABASE = None def __init__(self): super().__init__() self.project = self.pipeline.get_option('project') self.spanner_instance = self.pipeline.get_option('spanner_instance') self.spanner_database = self.pipeline.get_option('spanner_database') self._init_setup() def _generate_table_name(self): self.TEST_DATABASE = "{}_{}".format( self.spanner_database, ''.join(random.sample(uuid.uuid4().hex, 4))) return self.TEST_DATABASE def _create_database(self): spanner_client = spanner.Client() instance = self._SPANNER_INSTANCE = spanner_client.instance( self.spanner_instance) database = instance.database( self.TEST_DATABASE, ddl_statements=[ """CREATE TABLE test ( id STRING(99) NOT NULL, data BYTES(MAX) NOT NULL ) PRIMARY KEY (id)""", ]) database.create() def _init_setup(self): """Create database.""" self._generate_table_name() self._create_database() def test(self): def format_record(record): import base64 return base64.b64encode(record[1]) def make_insert_mutations(element): import uuid # pylint: disable=reimported from apache_beam.io.gcp.experimental.spannerio import WriteMutation ins_mutation = WriteMutation.insert( table='test', columns=('id', 'data'), values=[(str(uuid.uuid1()), element)]) return [ins_mutation] ( # pylint: disable=expression-not-assigned self.pipeline | 'Produce rows' >> Read( SyntheticSource(self.parse_synthetic_source_options())) | 'Count messages' >> ParDo(CountMessages(self.metrics_namespace)) | 'Format' >> Map(format_record) | 'Make mutations' >> FlatMap(make_insert_mutations) | 'Measure time' >> ParDo(MeasureTime(self.metrics_namespace)) | 'Write to Spanner' >> WriteToSpanner( project_id=self.project, instance_id=self.spanner_instance, database_id=self.TEST_DATABASE, max_batch_size_bytes=5120) ) def cleanup(self): """Removes test database.""" database = self._SPANNER_INSTANCE.database(self.TEST_DATABASE) database.drop() if __name__ == '__main__': logging.basicConfig(level=logging.INFO) SpannerWritePerfTest().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/gce_metadata_util.py0000644000175100017510000000304415206113032023043 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Fetches GCE metadata if the calling process is running on a GCE VM.""" # pytype: skip-file import requests BASE_METADATA_URL = "http://metadata/computeMetadata/v1/" def _fetch_metadata(key): try: headers = {"Metadata-Flavor": "Google"} uri = BASE_METADATA_URL + key resp = requests.get(uri, headers=headers, timeout=5) # 5 seconds. if resp.status_code == 200: return resp.text except requests.exceptions.RequestException: # Silently fail, may mean its running on a non DataflowRunner, # in which case it's prefectly normal. pass return "" def _fetch_custom_gce_metadata(customMetadataKey): return _fetch_metadata("instance/attributes/" + customMetadataKey) def fetch_dataflow_job_id(): return _fetch_custom_gce_metadata("job_id") ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/gcsfilesystem.py0000644000175100017510000002763215206113032022302 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """GCS file system implementation for accessing files on GCS. **Updates to the I/O connector code** For any significant updates to this I/O connector, please consider involving corresponding code reviewers mentioned in https://github.com/apache/beam/blob/master/sdks/python/OWNERS """ # pytype: skip-file import traceback from typing import BinaryIO # pylint: disable=unused-import from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystem import CompressedFile from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystem import FileMetadata from apache_beam.io.filesystem import FileSystem from apache_beam.io.gcp import gcsio __all__ = ['GCSFileSystem'] class GCSFileSystem(FileSystem): """A GCS ``FileSystem`` implementation for accessing files on GCS. """ CHUNK_SIZE = gcsio.MAX_BATCH_OPERATION_SIZE # Chuck size in batch operations GCS_PREFIX = 'gs://' def __init__(self, pipeline_options): super().__init__(pipeline_options) self._pipeline_options = pipeline_options @classmethod def scheme(cls): """URI scheme for the FileSystem """ return 'gs' def join(self, basepath, *paths): """Join two or more pathname components for the filesystem Args: basepath: string path of the first component of the path paths: path components to be added Returns: full path after combining all the passed components """ if not basepath.startswith(GCSFileSystem.GCS_PREFIX): raise ValueError('Basepath %r must be GCS path.' % basepath) path = basepath for p in paths: path = path.rstrip('/') + '/' + p.lstrip('/') return path def split(self, path): """Splits the given path into two parts. Splits the path into a pair (head, tail) such that tail contains the last component of the path and head contains everything up to that. Head will include the GCS prefix ('gs://'). Args: path: path as a string Returns: a pair of path components as strings. """ path = path.strip() if not path.startswith(GCSFileSystem.GCS_PREFIX): raise ValueError('Path %r must be GCS path.' % path) prefix_len = len(GCSFileSystem.GCS_PREFIX) last_sep = path[prefix_len:].rfind('/') if last_sep >= 0: last_sep += prefix_len if last_sep > 0: return (path[:last_sep], path[last_sep + 1:]) elif last_sep < 0: return (path, '') else: raise ValueError('Invalid path: %s' % path) def mkdirs(self, path): """Recursively create directories for the provided path. Args: path: string path of the directory structure that should be created Raises: IOError: if leaf directory already exists. """ pass def has_dirs(self): """Whether this FileSystem supports directories.""" return False def _list(self, dir_or_prefix): """List files in a location. Listing is non-recursive, for filesystems that support directories. Args: dir_or_prefix: (string) A directory or location prefix (for filesystems that don't have directories). Returns: Generator of ``FileMetadata`` objects. Raises: ``BeamIOError``: if listing fails, but not if no files were found. """ try: for path, (size, updated) in self._gcsIO().list_files(dir_or_prefix, with_metadata=True): yield FileMetadata(path, size, updated) except Exception as e: # pylint: disable=broad-except raise BeamIOError("List operation failed", {dir_or_prefix: e}) def _gcsIO(self): return gcsio.GcsIO(pipeline_options=self._pipeline_options) def _path_open( self, path, mode, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO): """Helper functions to open a file in the provided mode. """ compression_type = FileSystem._get_compression_type(path, compression_type) mime_type = CompressionTypes.mime_type(compression_type, mime_type) raw_file = self._gcsIO().open(path, mode, mime_type=mime_type) if compression_type == CompressionTypes.UNCOMPRESSED: return raw_file return CompressedFile(raw_file, compression_type=compression_type) def create( self, path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO) -> BinaryIO: """Returns a write channel for the given file path. Args: path: string path of the file object to be written to the system mime_type: MIME type to specify the type of content in the file object compression_type: Type of compression to be used for this object Returns: file handle with a close function for the user to use """ return self._path_open(path, 'wb', mime_type, compression_type) def open( self, path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO) -> BinaryIO: """Returns a read channel for the given file path. Args: path: string path of the file object to be written to the system mime_type: MIME type to specify the type of content in the file object compression_type: Type of compression to be used for this object Returns: file handle with a close function for the user to use """ return self._path_open(path, 'rb', mime_type, compression_type) def copy(self, source_file_names, destination_file_names): """Recursively copy the file tree from the source to the destination Args: source_file_names: list of source file objects that needs to be copied destination_file_names: list of destination of the new object Raises: ``BeamIOError``: if any of the copy operations fail """ err_msg = ( "source_file_names and destination_file_names should " "be equal in length") assert len(source_file_names) == len(destination_file_names), err_msg def _copy_path(source, destination): """Recursively copy the file tree from the source to the destination """ if not destination.startswith(GCSFileSystem.GCS_PREFIX): raise ValueError('Destination %r must be GCS path.' % destination) # Use copy_tree if the path ends with / as it is a directory if source.endswith('/'): self._gcsIO().copytree(source, destination) else: self._gcsIO().copy(source, destination) exceptions = {} for source, destination in zip(source_file_names, destination_file_names): try: _copy_path(source, destination) except Exception as e: # pylint: disable=broad-except exceptions[(source, destination)] = e if exceptions: raise BeamIOError("Copy operation failed", exceptions) def rename(self, source_file_names, destination_file_names): """Rename the files at the source list to the destination list. Source and destination lists should be of the same size. Args: source_file_names: List of file paths that need to be moved destination_file_names: List of destination_file_names for the files Raises: ``BeamIOError``: if any of the rename operations fail """ err_msg = ( "source_file_names and destination_file_names should " "be equal in length") assert len(source_file_names) == len(destination_file_names), err_msg gcs_batches = [] gcs_current_batch = [] for src, dest in zip(source_file_names, destination_file_names): gcs_current_batch.append((src, dest)) if len(gcs_current_batch) == self.CHUNK_SIZE: gcs_batches.append(gcs_current_batch) gcs_current_batch = [] if gcs_current_batch: gcs_batches.append(gcs_current_batch) # Execute GCS renames if any and return exceptions. exceptions = {} for batch in gcs_batches: copy_statuses = self._gcsIO().copy_batch(batch) copy_succeeded = {} delete_targets = [] for src, dest, exception in copy_statuses: if exception: exceptions[(src, dest)] = exception else: copy_succeeded[src] = dest delete_targets.append(src) delete_statuses = self._gcsIO().delete_batch(delete_targets) for src, exception in delete_statuses: if exception: dest = copy_succeeded[src] exceptions[(src, dest)] = exception if exceptions: raise BeamIOError("Rename operation failed", exceptions) def exists(self, path): """Check if the provided path exists on the FileSystem. Args: path: string path that needs to be checked. Returns: boolean flag indicating if path exists """ return self._gcsIO().exists(path) def size(self, path): """Get size of path on the FileSystem. Args: path: string path in question. Returns: int size of path according to the FileSystem. Raises: ``BeamIOError``: if path doesn't exist. """ return self._gcsIO().size(path) def last_updated(self, path): """Get UNIX Epoch time in seconds on the FileSystem. Args: path: string path of file. Returns: float UNIX Epoch time Raises: ``BeamIOError``: if path doesn't exist. """ return self._gcsIO().last_updated(path) def checksum(self, path): """Fetch checksum metadata of a file on the :class:`~apache_beam.io.filesystem.FileSystem`. Args: path: string path of a file. Returns: string containing checksum Raises: ``BeamIOError``: if path isn't a file or doesn't exist. """ try: return self._gcsIO().checksum(path) except Exception as e: # pylint: disable=broad-except raise BeamIOError("Checksum operation failed", {path: e}) def metadata(self, path): """Fetch metadata fields of a file on the FileSystem. Args: path: string path of a file. Returns: :class:`~apache_beam.io.filesystem.FileMetadata`. Raises: ``BeamIOError``: if path isn't a file or doesn't exist. """ try: file_metadata = self._gcsIO()._status(path) return FileMetadata(path, file_metadata['size'], file_metadata['updated']) except Exception as e: # pylint: disable=broad-except raise BeamIOError("Metadata operation failed", {path: e}) def delete(self, paths): """Deletes files or directories at the provided paths. Directories will be deleted recursively. Args: paths: list of paths that give the file objects to be deleted """ exceptions = {} for path in paths: if path.endswith('/'): self._gcsIO().delete(path, recursive=True) continue else: path_to_use = path match_result = self.match([path_to_use])[0] statuses = self._gcsIO().delete_batch( [m.path for m in match_result.metadata_list]) for target, exception in statuses: if exception: exceptions[target] = exception if exceptions: raise BeamIOError("Delete operation failed", exceptions) def report_lineage(self, path, lineage): try: components = gcsio.parse_gcs_path(path, object_optional=True) except ValueError: # report lineage is fail-safe traceback.print_exc() return if components and not components[-1]: components = components[:-1] lineage.add('gcs', *components, last_segment_sep='/') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/gcsfilesystem_integration_test.py0000644000175100017510000001143715206113032025740 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration tests for gcsfilesystem module. Runs tests against Google Cloud Storage service. Instantiates a TestPipeline to get options such as GCP project name, but doesn't actually start a Beam pipeline or test any specific runner. Run the following in 'sdks/python' directory to run these tests manually: scripts/run_integration_test.sh \ --test_opts apache_beam/io/gcp/gcsfilesystem_integration_test.py """ # pytype: skip-file import logging import unittest import uuid import pytest from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystems import FileSystems from apache_beam.testing.test_pipeline import TestPipeline try: from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem fs_not_available = False except ImportError: fs_not_available = True # type: ignore @unittest.skipIf(fs_not_available, 'GCP dependencies are not installed') class GcsFileSystemIntegrationTest(unittest.TestCase): INPUT_FILE = 'gs://dataflow-samples/shakespeare/kinglear.txt' def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.runner_name = type(self.test_pipeline.runner).__name__ if self.runner_name != 'TestDataflowRunner': # This test doesn't run a pipeline, so it doesn't make sense to try it on # different runners. Running with TestDataflowRunner makes sense since # it uses GoogleCloudOptions such as 'project'. raise unittest.SkipTest('This test only runs with TestDataflowRunner.') self.project = self.test_pipeline.get_option('project') self.gcs_tempdir = ( self.test_pipeline.get_option('temp_location') + '/gcs_it-' + str(uuid.uuid4())) self.fs = GCSFileSystem(self.test_pipeline.get_pipeline_options()) def tearDown(self): FileSystems.delete([self.gcs_tempdir + '/']) def _verify_copy(self, src, dest): self.assertTrue(FileSystems.exists(src), 'src does not exist: %s' % src) self.assertTrue(FileSystems.exists(dest), 'dest does not exist: %s' % dest) src_checksum = self.fs.checksum(src) dest_checksum = self.fs.checksum(dest) self.assertEqual(src_checksum, dest_checksum) def _verify_rename(self, src, dest): self.assertFalse(FileSystems.exists(src), 'file %s not renamed' % src) self.assertTrue(FileSystems.exists(dest), 'file not renamed to %s' % dest) def _test_copy(self, name, max_bytes_rewritten_per_call=None, src=None): src = src or self.INPUT_FILE dest = self.gcs_tempdir + '/%s' % name extra_kwargs = {} if max_bytes_rewritten_per_call is not None: extra_kwargs['max_bytes_rewritten_per_call'] = ( max_bytes_rewritten_per_call) self.fs.copy([src], [dest]) self._verify_copy(src, dest) @pytest.mark.it_postcommit def test_copy(self): self._test_copy("test_copy") @pytest.mark.it_postcommit def test_rename(self): num_copies = 10 srcs = [self.INPUT_FILE] * num_copies dests = [ self.gcs_tempdir + '/%s_%d' % (self.INPUT_FILE, i) for i in range(num_copies) ] self.fs.copy(srcs, dests) new_names = [ self.gcs_tempdir + '/%s_%d' % ("renamed", i) for i in range(num_copies) ] self.fs.rename(dests, new_names) for _old, _new in list(zip(dests, new_names)): self._verify_rename(_old, _new) @pytest.mark.it_postcommit def test_rename_error(self): num_copies = 10 srcs = [self.INPUT_FILE] * num_copies dests = [ self.gcs_tempdir + '/%s_%d' % (self.INPUT_FILE, i) for i in range(num_copies) ] self.fs.copy(srcs, dests) new_names = [ self.gcs_tempdir + '/%s_%d' % ("renamed", i) for i in range(num_copies) ] # corrupt one of the destination names bad_name = self.gcs_tempdir + '/errorbadwrong' dests[int(num_copies / 2)] = bad_name with self.assertRaises(BeamIOError): self.fs.rename(dests, new_names) for _old, _new in list(zip(dests, new_names)): if _old != bad_name: self._verify_rename(_old, _new) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/gcsfilesystem_test.py0000644000175100017510000003323715206113032023337 0ustar00runnerrunner# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for GCS File System.""" # pytype: skip-file import logging import unittest import mock from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystem import FileMetadata from apache_beam.options.pipeline_options import PipelineOptions # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.io.gcp import gcsfilesystem except ImportError: gcsfilesystem = None # type: ignore # pylint: enable=wrong-import-order, wrong-import-position @unittest.skipIf(gcsfilesystem is None, 'GCP dependencies are not installed') class GCSFileSystemTest(unittest.TestCase): def setUp(self): pipeline_options = PipelineOptions() self.fs = gcsfilesystem.GCSFileSystem(pipeline_options=pipeline_options) def test_scheme(self): self.assertEqual(self.fs.scheme(), 'gs') self.assertEqual(gcsfilesystem.GCSFileSystem.scheme(), 'gs') def test_join(self): self.assertEqual( 'gs://bucket/path/to/file', self.fs.join('gs://bucket/path', 'to', 'file')) self.assertEqual( 'gs://bucket/path/to/file', self.fs.join('gs://bucket/path', 'to/file')) self.assertEqual( 'gs://bucket/path/to/file', self.fs.join('gs://bucket/path', '/to/file')) self.assertEqual( 'gs://bucket/path/to/file', self.fs.join('gs://bucket/path/', 'to', 'file')) self.assertEqual( 'gs://bucket/path/to/file', self.fs.join('gs://bucket/path/', 'to/file')) self.assertEqual( 'gs://bucket/path/to/file', self.fs.join('gs://bucket/path/', '/to/file')) with self.assertRaises(ValueError): self.fs.join('/bucket/path/', '/to/file') def test_split(self): self.assertEqual(('gs://foo/bar', 'baz'), self.fs.split('gs://foo/bar/baz')) self.assertEqual(('gs://foo', ''), self.fs.split('gs://foo/')) self.assertEqual(('gs://foo', ''), self.fs.split('gs://foo')) with self.assertRaises(ValueError): self.fs.split('/no/gcs/prefix') @mock.patch('apache_beam.io.gcp.gcsfilesystem.gcsio') def test_match_multiples(self, mock_gcsio): # Prepare mocks. gcsio_mock = mock.MagicMock() gcsfilesystem.gcsio.GcsIO = lambda pipeline_options=None: gcsio_mock gcsio_mock.list_files.return_value = iter([ ('gs://bucket/file1', (1, 99999.0)), ('gs://bucket/file2', (2, 88888.0)) ]) expected_results = set([ FileMetadata('gs://bucket/file1', 1, 99999.0), FileMetadata('gs://bucket/file2', 2, 88888.0) ]) match_result = self.fs.match(['gs://bucket/'])[0] self.assertEqual(set(match_result.metadata_list), expected_results) gcsio_mock.list_files.assert_called_once_with( 'gs://bucket/', with_metadata=True) @mock.patch('apache_beam.io.gcp.gcsfilesystem.gcsio') def test_match_multiples_limit(self, mock_gcsio): # Prepare mocks. gcsio_mock = mock.MagicMock() limit = 1 gcsfilesystem.gcsio.GcsIO = lambda pipeline_options=None: gcsio_mock gcsio_mock.list_files.return_value = iter([ ('gs://bucket/file1', (1, 99999.0)) ]) expected_results = set([FileMetadata('gs://bucket/file1', 1, 99999.0)]) match_result = self.fs.match(['gs://bucket/'], [limit])[0] self.assertEqual(set(match_result.metadata_list), expected_results) self.assertEqual(len(match_result.metadata_list), limit) gcsio_mock.list_files.assert_called_once_with( 'gs://bucket/', with_metadata=True) @mock.patch('apache_beam.io.gcp.gcsfilesystem.gcsio') def test_match_multiples_error(self, mock_gcsio): # Prepare mocks. gcsio_mock = mock.MagicMock() gcsfilesystem.gcsio.GcsIO = lambda pipeline_options=None: gcsio_mock exception = IOError('Failed') gcsio_mock.list_files.side_effect = exception with self.assertRaisesRegex(BeamIOError, r'^Match operation failed') as error: self.fs.match(['gs://bucket/']) self.assertRegex( str(error.exception.exception_details), r'gs://bucket/.*%s' % exception) gcsio_mock.list_files.assert_called_once_with( 'gs://bucket/', with_metadata=True) @mock.patch('apache_beam.io.gcp.gcsfilesystem.gcsio') def test_match_multiple_patterns(self, mock_gcsio): # Prepare mocks. gcsio_mock = mock.MagicMock() gcsfilesystem.gcsio.GcsIO = lambda pipeline_options=None: gcsio_mock gcsio_mock.list_files.side_effect = [ iter([('gs://bucket/file1', (1, 99999.0))]), iter([('gs://bucket/file2', (2, 88888.0))]), ] expected_results = [[FileMetadata('gs://bucket/file1', 1, 99999.0)], [FileMetadata('gs://bucket/file2', 2, 88888.0)]] result = self.fs.match(['gs://bucket/file1*', 'gs://bucket/file2*']) self.assertEqual([mr.metadata_list for mr in result], expected_results) @mock.patch('apache_beam.io.gcp.gcsfilesystem.gcsio') def test_create(self, mock_gcsio): # Prepare mocks. gcsio_mock = mock.MagicMock() gcsfilesystem.gcsio.GcsIO = lambda pipeline_options=None: gcsio_mock # Issue file copy _ = self.fs.create('gs://bucket/from1', 'application/octet-stream') gcsio_mock.open.assert_called_once_with( 'gs://bucket/from1', 'wb', mime_type='application/octet-stream') @mock.patch('apache_beam.io.gcp.gcsfilesystem.gcsio') def test_open(self, mock_gcsio): # Prepare mocks. gcsio_mock = mock.MagicMock() gcsfilesystem.gcsio.GcsIO = lambda pipeline_options=None: gcsio_mock # Issue file copy _ = self.fs.open('gs://bucket/from1', 'application/octet-stream') gcsio_mock.open.assert_called_once_with( 'gs://bucket/from1', 'rb', mime_type='application/octet-stream') @mock.patch('apache_beam.io.gcp.gcsfilesystem.gcsio') def test_copy_file(self, mock_gcsio): # Prepare mocks. gcsio_mock = mock.MagicMock() gcsfilesystem.gcsio.GcsIO = lambda pipeline_options=None: gcsio_mock sources = ['gs://bucket/from1'] destinations = ['gs://bucket/to1'] # Issue file copy self.fs.copy(sources, destinations) gcsio_mock.copy.assert_called_once_with( 'gs://bucket/from1', 'gs://bucket/to1') @mock.patch('apache_beam.io.gcp.gcsfilesystem.gcsio') def test_copy_file_error(self, mock_gcsio): # Prepare mocks. gcsio_mock = mock.MagicMock() gcsfilesystem.gcsio.GcsIO = lambda pipeline_options=None: gcsio_mock sources = ['gs://bucket/from1'] destinations = ['gs://bucket/to1'] exception = IOError('Failed') gcsio_mock.copy.side_effect = exception # Issue batch rename. expected_results = {(s, d): exception for s, d in zip(sources, destinations)} # Issue batch copy. with self.assertRaisesRegex(BeamIOError, r'^Copy operation failed') as error: self.fs.copy(sources, destinations) self.assertEqual(error.exception.exception_details, expected_results) gcsio_mock.copy.assert_called_once_with( 'gs://bucket/from1', 'gs://bucket/to1') @mock.patch('apache_beam.io.gcp.gcsfilesystem.gcsio') def test_copy_tree(self, mock_gcsio): # Prepare mocks. gcsio_mock = mock.MagicMock() gcsfilesystem.gcsio.GcsIO = lambda pipeline_options=None: gcsio_mock sources = ['gs://bucket1/'] destinations = ['gs://bucket2/'] # Issue directory copy self.fs.copy(sources, destinations) gcsio_mock.copytree.assert_called_once_with( 'gs://bucket1/', 'gs://bucket2/') @mock.patch('apache_beam.io.gcp.gcsfilesystem.gcsio') def test_rename(self, mock_gcsio): # Prepare mocks. gcsio_mock = mock.MagicMock() gcsfilesystem.gcsio.GcsIO = lambda pipeline_options=None: gcsio_mock sources = [ 'gs://bucket/from1', 'gs://bucket/from2', 'gs://bucket/from3', ] destinations = [ 'gs://bucket/to1', 'gs://bucket/to2', 'gs://bucket/to3', ] gcsio_mock.copy_batch.side_effect = [[ ('gs://bucket/from1', 'gs://bucket/to1', None), ('gs://bucket/from2', 'gs://bucket/to2', None), ('gs://bucket/from3', 'gs://bucket/to3', None), ]] gcsio_mock.delete_batch.side_effect = [[ ('gs://bucket/from1', None), ('gs://bucket/from2', None), ('gs://bucket/from3', None), ]] # Issue batch rename. self.fs.rename(sources, destinations) gcsio_mock.copy_batch.assert_called_once_with([ ('gs://bucket/from1', 'gs://bucket/to1'), ('gs://bucket/from2', 'gs://bucket/to2'), ('gs://bucket/from3', 'gs://bucket/to3'), ]) gcsio_mock.delete_batch.assert_called_once_with([ 'gs://bucket/from1', 'gs://bucket/from2', 'gs://bucket/from3', ]) @mock.patch('apache_beam.io.gcp.gcsfilesystem.gcsio') def test_rename_delete_error(self, mock_gcsio): # Prepare mocks. gcsio_mock = mock.MagicMock() gcsfilesystem.gcsio.GcsIO = lambda pipeline_options=None: gcsio_mock sources = [ 'gs://bucket/from1', 'gs://bucket/from2', 'gs://bucket/from3', ] destinations = [ 'gs://bucket/to1', 'gs://bucket/to2', 'gs://bucket/to3', ] gcsio_mock.copy_batch.side_effect = [[ ('gs://bucket/from1', 'gs://bucket/to1', None), ('gs://bucket/from2', 'gs://bucket/to2', None), ('gs://bucket/from3', 'gs://bucket/to3', None), ]] gcsio_mock.delete_batch.side_effect = [[ ('gs://bucket/from1', None), ('gs://bucket/from2', Exception("BadThings")), ('gs://bucket/from3', None), ]] # Issue batch rename. with self.assertRaisesRegex(BeamIOError, r'^Rename operation failed'): self.fs.rename(sources, destinations) gcsio_mock.copy_batch.assert_called_once_with([ ('gs://bucket/from1', 'gs://bucket/to1'), ('gs://bucket/from2', 'gs://bucket/to2'), ('gs://bucket/from3', 'gs://bucket/to3'), ]) gcsio_mock.delete_batch.assert_called_once_with([ 'gs://bucket/from1', 'gs://bucket/from2', 'gs://bucket/from3', ]) @mock.patch('apache_beam.io.gcp.gcsfilesystem.gcsio') def test_rename_copy_error(self, mock_gcsio): # Prepare mocks. gcsio_mock = mock.MagicMock() gcsfilesystem.gcsio.GcsIO = lambda pipeline_options=None: gcsio_mock sources = [ 'gs://bucket/from1', 'gs://bucket/from2', 'gs://bucket/from3', ] destinations = [ 'gs://bucket/to1', 'gs://bucket/to2', 'gs://bucket/to3', ] gcsio_mock.copy_batch.side_effect = [[ ('gs://bucket/from1', 'gs://bucket/to1', None), ('gs://bucket/from2', 'gs://bucket/to2', Exception("BadThings")), ('gs://bucket/from3', 'gs://bucket/to3', None), ]] gcsio_mock.delete_batch.side_effect = [[ ('gs://bucket/from1', None), ('gs://bucket/from3', None), ]] # Issue batch rename. with self.assertRaisesRegex(BeamIOError, r'^Rename operation failed'): self.fs.rename(sources, destinations) gcsio_mock.copy_batch.assert_called_once_with([ ('gs://bucket/from1', 'gs://bucket/to1'), ('gs://bucket/from2', 'gs://bucket/to2'), ('gs://bucket/from3', 'gs://bucket/to3'), ]) gcsio_mock.delete_batch.assert_called_once_with([ 'gs://bucket/from1', 'gs://bucket/from3', ]) @mock.patch('apache_beam.io.gcp.gcsfilesystem.gcsio') def test_delete(self, mock_gcsio): # Prepare mocks. gcsio_mock = mock.MagicMock() gcsfilesystem.gcsio.GcsIO = lambda pipeline_options=None: gcsio_mock gcsio_mock._status.return_value = {'size': 0, 'updated': 99999.0} files = [ 'gs://bucket/from1', 'gs://bucket/from2', 'gs://bucket/from3', ] # Issue batch delete. self.fs.delete(files) gcsio_mock.delete_batch.assert_called() @mock.patch('apache_beam.io.gcp.gcsfilesystem.gcsio') def test_delete_error(self, mock_gcsio): # Prepare mocks. gcsio_mock = mock.MagicMock() gcsfilesystem.gcsio.GcsIO = lambda pipeline_options=None: gcsio_mock gcsio_mock._status.return_value = {'size': 0, 'updated': 99999.0} files = [ 'gs://bucket/from1', 'gs://bucket/from2', 'gs://bucket/from3', ] gcsio_mock.delete_batch.side_effect = [ [('gs://bucket/from1', None)], [('gs://bucket/from2', Exception("BadThings"))], [('gs://bucket/from3', None)], ] # Issue batch delete. with self.assertRaisesRegex(BeamIOError, r'^Delete operation failed'): self.fs.delete(files) gcsio_mock.delete_batch.assert_called() def test_lineage(self): self._verify_lineage("gs://bucket/", ("bucket", )) self._verify_lineage("gs://bucket/foo/bar.txt", ("bucket", "foo/bar.txt")) def _verify_lineage(self, uri, expected_segments): lineage_mock = mock.MagicMock() self.fs.report_lineage(uri, lineage_mock) lineage_mock.add.assert_called_once_with( "gcs", *expected_segments, last_segment_sep='/') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/gcsio.py0000644000175100017510000006004415206113032020517 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Google Cloud Storage client. This library evolved from the Google App Engine GCS client available at https://github.com/GoogleCloudPlatform/appengine-gcs-client. **Updates to the I/O connector code** For any significant updates to this I/O connector, please consider involving corresponding code reviewers mentioned in https://github.com/apache/beam/blob/master/sdks/python/OWNERS """ # pytype: skip-file import logging import re import time from typing import Optional from typing import Union from google.api_core.exceptions import RetryError from google.cloud import storage from google.cloud.exceptions import NotFound from google.cloud.exceptions import from_http_response from google.cloud.storage.fileio import BlobReader from google.cloud.storage.fileio import BlobWriter from google.cloud.storage.retry import DEFAULT_RETRY from apache_beam import version as beam_version from apache_beam.internal.gcp import auth from apache_beam.io.gcp import gcsio_retry from apache_beam.metrics.metric import Metrics from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions __all__ = ['GcsIO', 'create_storage_client'] _LOGGER = logging.getLogger(__name__) DEFAULT_READ_BUFFER_SIZE = 16 * 1024 * 1024 # Maximum number of operations permitted in GcsIO.copy_batch() and # GcsIO.delete_batch(). MAX_BATCH_OPERATION_SIZE = 100 def parse_gcs_path(gcs_path, object_optional=False): """Return the bucket and object names of the given gs:// path.""" match = re.match('^gs://([^/]+)/(.*)$', gcs_path) if match is None or (match.group(2) == '' and not object_optional): raise ValueError( 'GCS path must be in the form gs:///. ' f'Encountered {gcs_path!r}') return match.group(1), match.group(2) def default_gcs_bucket_name(project, region): from hashlib import md5 return 'dataflow-staging-%s-%s' % ( region, md5(project.encode('utf8')).hexdigest()) def get_or_create_default_gcs_bucket(options): """Create a default GCS bucket for this project.""" if getattr(options, 'dataflow_kms_key', None): _LOGGER.warning( 'Cannot create a default bucket when --dataflow_kms_key is set.') return None project = getattr(options, 'project', None) region = getattr(options, 'region', None) if not project or not region: return None bucket_name = default_gcs_bucket_name(project, region) bucket = GcsIO(pipeline_options=options).get_bucket(bucket_name) if bucket: return bucket else: _LOGGER.warning( 'Creating default GCS bucket for project %s: gs://%s', project, bucket_name) return GcsIO(pipeline_options=options).create_bucket( bucket_name, project, location=region) def create_storage_client(pipeline_options, use_credentials=True): """Create a GCS client for Beam via GCS Client Library. Args: pipeline_options(apache_beam.options.pipeline_options.PipelineOptions): the options of the pipeline. use_credentials(bool): whether to create an authenticated client based on pipeline options or an anonymous client. Returns: A google.cloud.storage.client.Client instance. """ if use_credentials: credentials = auth.get_service_credentials(pipeline_options) else: credentials = None if credentials: google_cloud_options = pipeline_options.view_as(GoogleCloudOptions) from google.api_core import client_info beam_client_info = client_info.ClientInfo( user_agent="apache-beam/%s (GPN:Beam)" % beam_version.__version__) extra_headers = { "x-goog-custom-audit-job": google_cloud_options.job_name if google_cloud_options.job_name else "UNKNOWN" } # Note: Custom audit entries with "job" key will overwrite the default above if google_cloud_options.gcs_custom_audit_entries is not None: extra_headers.update(google_cloud_options.gcs_custom_audit_entries) return storage.Client( credentials=credentials.get_google_auth_credentials(), project=google_cloud_options.project, client_info=beam_client_info, extra_headers=extra_headers) else: return storage.Client.create_anonymous_client() class GcsIO(object): """Google Cloud Storage I/O client.""" def __init__( self, storage_client: Optional[storage.Client] = None, pipeline_options: Optional[Union[dict, PipelineOptions]] = None) -> None: if pipeline_options is None: pipeline_options = PipelineOptions() elif isinstance(pipeline_options, dict): pipeline_options = PipelineOptions.from_dictionary(pipeline_options) if storage_client is None: storage_client = create_storage_client(pipeline_options) google_cloud_options = pipeline_options.view_as(GoogleCloudOptions) self.enable_read_bucket_metric = getattr( google_cloud_options, 'enable_bucket_read_metric_counter', False) self.enable_write_bucket_metric = getattr( google_cloud_options, 'enable_bucket_write_metric_counter', False) self.client = storage_client self._rewrite_cb = None self.bucket_to_project_number = {} self._storage_client_retry = gcsio_retry.get_retry(pipeline_options) self._use_blob_generation = getattr( google_cloud_options, 'enable_gcsio_blob_generation', False) def get_project_number(self, bucket): if bucket not in self.bucket_to_project_number: bucket_metadata = self.get_bucket(bucket_name=bucket) if bucket_metadata: self.bucket_to_project_number[bucket] = bucket_metadata.projectNumber return self.bucket_to_project_number.get(bucket, None) def get_bucket(self, bucket_name, **kwargs): """Returns an object bucket from its name, or None if it does not exist.""" try: return self.client.lookup_bucket( bucket_name, retry=self._storage_client_retry, **kwargs) except NotFound: return None def create_bucket( self, bucket_name, project, kms_key=None, location=None, soft_delete_retention_duration_seconds=0): """Create and return a GCS bucket in a specific project.""" try: bucket = self.client.bucket(bucket_name) bucket.soft_delete_policy.retention_duration_seconds = ( soft_delete_retention_duration_seconds) bucket = self.client.create_bucket( bucket_or_name=bucket, project=project, location=location, retry=self._storage_client_retry) if kms_key: bucket.default_kms_key_name(kms_key) bucket.patch() return bucket except NotFound: return None def open( self, filename, mode='r', read_buffer_size=DEFAULT_READ_BUFFER_SIZE, mime_type='application/octet-stream'): """Open a GCS file path for reading or writing. Args: filename (str): GCS file path in the form ``gs:///``. mode (str): ``'r'`` for reading or ``'w'`` for writing. read_buffer_size (int): Buffer size to use during read operations. mime_type (str): Mime type to set for write operations. Returns: GCS file object. Raises: ValueError: Invalid open file mode. """ bucket_name, blob_name = parse_gcs_path(filename) bucket = self.client.bucket(bucket_name) if mode == 'r' or mode == 'rb': blob = bucket.blob(blob_name) return BeamBlobReader( blob, chunk_size=read_buffer_size, enable_read_bucket_metric=self.enable_read_bucket_metric, retry=self._storage_client_retry) elif mode == 'w' or mode == 'wb': blob = bucket.blob(blob_name) return BeamBlobWriter( blob, mime_type, enable_write_bucket_metric=self.enable_write_bucket_metric, retry=self._storage_client_retry) else: raise ValueError('Invalid file open mode: %s.' % mode) def delete(self, path, recursive=False): """Deletes the object at the given GCS path. If the path is a directory (prefix), it deletes all blobs under that prefix when recursive=True. Args: path: GCS file path pattern in the form gs:///. recursive (bool, optional): If True, deletes all objects under the prefix when the path is a directory (default: False). """ bucket_name, blob_name = parse_gcs_path(path) bucket = self.client.bucket(bucket_name) if recursive: # List all blobs under the prefix. blobs_to_delete = bucket.list_blobs( prefix=blob_name, retry=self._storage_client_retry) # Collect full paths for batch deletion. paths_to_delete = [ f'gs://{bucket_name}/{blob.name}' for blob in blobs_to_delete ] if paths_to_delete: # Delete them in batches. results = self.delete_batch(paths_to_delete) # Log any errors encountered during batch deletion. errors = [f'{path}: {err}' for path, err in results if err is not None] if errors: _LOGGER.warning( 'Failed to delete some objects during recursive delete of %s: %s', path, ', '.join(errors)) else: # Delete only the specific blob. self._delete_blob(bucket, blob_name) def _delete_blob(self, bucket, blob_name): """Helper method to delete a single blob from GCS. Args: bucket: The GCS bucket object. blob_name: The name of the blob to delete under the bucket. """ if self._use_blob_generation: # Fetch blob generation if required. blob = bucket.get_blob(blob_name, retry=self._storage_client_retry) generation = getattr(blob, "generation", None) else: generation = None try: bucket.delete_blob( blob_name, if_generation_match=generation, retry=self._storage_client_retry) except NotFound: return def _batch_with_retry(self, requests, fn): current_requests = [*enumerate(requests)] responses = [None for _ in current_requests] @self._storage_client_retry def run_with_retry(): current_batch = self.client.batch(raise_exception=False) with current_batch: for _, request in current_requests: fn(request) last_retryable_exception = None for (i, current_pair), response in zip( [*current_requests], current_batch._responses ): responses[i] = response should_retry = ( response.status_code >= 400 and self._storage_client_retry._predicate(from_http_response(response))) if should_retry: last_retryable_exception = from_http_response(response) else: current_requests.remove((i, current_pair)) if last_retryable_exception: raise last_retryable_exception try: run_with_retry() except RetryError: pass return responses def _delete_batch_request(self, path): bucket_name, blob_name = parse_gcs_path(path) bucket = self.client.bucket(bucket_name) bucket.delete_blob(blob_name) def delete_batch(self, paths): """Deletes the objects at the given GCS paths. Args: paths: List of GCS file path patterns or Dict with GCS file path patterns as keys. The patterns are in the form gs:///, but not to exceed MAX_BATCH_OPERATION_SIZE in length. Returns: List of tuples of (path, exception) in the same order as the paths argument, where exception is None if the operation succeeded or the relevant exception if the operation failed. """ final_results = [] s = 0 if not isinstance(paths, list): paths = list(iter(paths)) while s < len(paths): if (s + MAX_BATCH_OPERATION_SIZE) < len(paths): current_paths = paths[s:s + MAX_BATCH_OPERATION_SIZE] else: current_paths = paths[s:] responses = self._batch_with_retry( current_paths, self._delete_batch_request) for i, path in enumerate(current_paths): error_code = None resp = responses[i] if resp.status_code >= 400 and resp.status_code != 404: error_code = resp.status_code final_results.append((path, error_code)) s += MAX_BATCH_OPERATION_SIZE return final_results def copy(self, src, dest): """Copies the given GCS object from src to dest. Args: src: GCS file path pattern in the form gs:///. dest: GCS file path pattern in the form gs:///. Raises: Any exceptions during copying """ src_bucket_name, src_blob_name = parse_gcs_path(src) dest_bucket_name, dest_blob_name= parse_gcs_path(dest, object_optional=True) src_bucket = self.client.bucket(src_bucket_name) if self._use_blob_generation: src_blob = src_bucket.get_blob(src_blob_name) if src_blob is None: raise NotFound("source blob %s not found during copying" % src) src_generation = src_blob.generation else: src_blob = src_bucket.blob(src_blob_name) src_generation = None dest_bucket = self.client.bucket(dest_bucket_name) if not dest_blob_name: dest_blob_name = None src_bucket.copy_blob( src_blob, dest_bucket, new_name=dest_blob_name, source_generation=src_generation, retry=self._storage_client_retry) def _copy_batch_request(self, pair): src_bucket_name, src_blob_name = parse_gcs_path(pair[0]) dest_bucket_name, dest_blob_name = parse_gcs_path(pair[1]) src_bucket = self.client.bucket(src_bucket_name) src_blob = src_bucket.blob(src_blob_name) dest_bucket = self.client.bucket(dest_bucket_name) src_bucket.copy_blob(src_blob, dest_bucket, dest_blob_name) def copy_batch(self, src_dest_pairs): """Copies the given GCS objects from src to dest. Args: src_dest_pairs: list of (src, dest) tuples of gs:/// files paths to copy from src to dest, not to exceed MAX_BATCH_OPERATION_SIZE in length. Returns: List of tuples of (src, dest, exception) in the same order as the src_dest_pairs argument, where exception is None if the operation succeeded or the relevant exception if the operation failed. """ final_results = [] s = 0 while s < len(src_dest_pairs): if (s + MAX_BATCH_OPERATION_SIZE) < len(src_dest_pairs): current_pairs = src_dest_pairs[s:s + MAX_BATCH_OPERATION_SIZE] else: current_pairs = src_dest_pairs[s:] responses = self._batch_with_retry( current_pairs, self._copy_batch_request) for i, pair in enumerate(current_pairs): error_code = None resp = responses[i] if resp.status_code >= 400: error_code = resp.status_code final_results.append((pair[0], pair[1], error_code)) s += MAX_BATCH_OPERATION_SIZE return final_results # We intentionally do not decorate this method with a retry, since the # underlying copy and delete operations are already idempotent operations. def copytree(self, src, dest): """Renames the given GCS "directory" recursively from src to dest. Args: src: GCS file path pattern in the form gs:////. dest: GCS file path pattern in the form gs:////. """ assert src.endswith('/') assert dest.endswith('/') for entry, _ in self.list_files(src): rel_path = entry[len(src):] self.copy(entry, dest + rel_path) # We intentionally do not decorate this method with a retry, since the # underlying copy and delete operations are already idempotent operations. def rename(self, src, dest): """Renames the given GCS object from src to dest. Args: src: GCS file path pattern in the form gs:///. dest: GCS file path pattern in the form gs:///. """ self.copy(src, dest) self.delete(src) def exists(self, path): """Returns whether the given GCS object exists. Args: path: GCS file path pattern in the form gs:///. """ bucket_name, blob_name = parse_gcs_path(path) bucket = self.client.bucket(bucket_name) blob = bucket.blob(blob_name) return blob.exists(retry=self._storage_client_retry) def checksum(self, path): """Looks up the checksum of a GCS object. Args: path: GCS file path pattern in the form gs:///. """ return self._gcs_object(path).crc32c def size(self, path): """Returns the size of a single GCS object. This method does not perform glob expansion. Hence the given path must be for a single GCS object. Returns: size of the GCS object in bytes. """ return self._gcs_object(path).size def kms_key(self, path): """Returns the KMS key of a single GCS object. This method does not perform glob expansion. Hence the given path must be for a single GCS object. Returns: KMS key name of the GCS object as a string, or None if it doesn't have one. """ return self._gcs_object(path).kms_key_name def last_updated(self, path): """Returns the last updated epoch time of a single GCS object. This method does not perform glob expansion. Hence the given path must be for a single GCS object. Returns: last updated time of the GCS object in second. """ return self._updated_to_seconds(self._gcs_object(path).updated) def _status(self, path): """For internal use only; no backwards-compatibility guarantees. Returns supported fields (checksum, kms_key, last_updated, size) of a single object as a dict at once. This method does not perform glob expansion. Hence the given path must be for a single GCS object. Returns: dict of fields of the GCS object. """ gcs_object = self._gcs_object(path) file_status = {} if hasattr(gcs_object, 'crc32c'): file_status['checksum'] = gcs_object.crc32c if hasattr(gcs_object, 'kms_key_name'): file_status['kms_key'] = gcs_object.kms_key_name if hasattr(gcs_object, 'updated'): file_status['updated'] = self._updated_to_seconds(gcs_object.updated) if hasattr(gcs_object, 'size'): file_status['size'] = gcs_object.size return file_status def _gcs_object(self, path): """Returns a gcs object for the given path This method does not perform glob expansion. Hence the given path must be for a single GCS object. The method will make HTTP requests. Returns: GCS object. """ bucket_name, blob_name = parse_gcs_path(path) bucket = self.client.bucket(bucket_name) blob = bucket.get_blob(blob_name, retry=self._storage_client_retry) if blob: return blob else: raise NotFound('Object %s not found', path) def list_files(self, path, with_metadata=False): """Lists files matching the prefix. Args: path: GCS file path pattern in the form gs:///[name]. with_metadata: Experimental. Specify whether returns file metadata. Returns: If ``with_metadata`` is False: generator of tuple(file name, size); if ``with_metadata`` is True: generator of tuple(file name, tuple(size, timestamp)). """ bucket_name, prefix = parse_gcs_path(path, object_optional=True) file_info = set() counter = 0 start_time = time.time() if with_metadata: _LOGGER.debug("Starting the file information of the input") else: _LOGGER.debug("Starting the size estimation of the input") bucket = self.client.bucket(bucket_name) response = self.client.list_blobs( bucket, prefix=prefix, retry=self._storage_client_retry) for item in response: file_name = 'gs://%s/%s' % (item.bucket.name, item.name) if file_name not in file_info: file_info.add(file_name) counter += 1 if counter % 10000 == 0: if with_metadata: _LOGGER.info( "Finished computing file information of: %s files", len(file_info)) else: _LOGGER.info("Finished computing size of: %s files", len(file_info)) if with_metadata: yield file_name, (item.size, self._updated_to_seconds(item.updated)) else: yield file_name, item.size _LOGGER.log( # do not spam logs when list_files is likely used to check empty folder logging.INFO if counter > 0 else logging.DEBUG, "Finished listing %s files in %s seconds.", counter, time.time() - start_time) @staticmethod def _updated_to_seconds(updated): """Helper function transform the updated field of response to seconds""" return ( time.mktime(updated.timetuple()) - time.timezone + updated.microsecond / 1000000.0) def is_soft_delete_enabled(self, gcs_path): try: bucket_name, _ = parse_gcs_path(gcs_path, object_optional=True) bucket = self.get_bucket(bucket_name) if (bucket.soft_delete_policy is not None and bucket.soft_delete_policy.retention_duration_seconds > 0): return True except Exception: _LOGGER.warning( "Unexpected error occurred when checking soft delete policy for %s" % gcs_path) return False class BeamBlobReader(BlobReader): def __init__( self, blob, chunk_size=DEFAULT_READ_BUFFER_SIZE, enable_read_bucket_metric=False, retry=DEFAULT_RETRY, raw_download=True): # By default, we always request to retrieve raw data from GCS even if the # object meets the criteria of decompressive transcoding # (https://cloud.google.com/storage/docs/transcoding). super().__init__( blob, chunk_size=chunk_size, retry=retry, raw_download=raw_download) # TODO: Remove this after # https://github.com/googleapis/python-storage/issues/1406 is fixed. # As a workaround, we manually trigger a reload here. Otherwise, an internal # call of reader.seek() will cause an exception if raw_download is set # when initializing BlobReader(), blob.reload() # TODO: Currently there is a bug in GCS server side when a client requests # a file with "content-encoding=gzip" and "content-type=application/gzip" or # "content-type=application/x-gzip", which will lead to infinite loop. # We skip the support of this type of files until the GCS bug is fixed. # Internal bug id: 203845981. if (blob.content_encoding == "gzip" and blob.content_type in ["application/gzip", "application/x-gzip"]): raise NotImplementedError("Doubly compressed files not supported.") self.enable_read_bucket_metric = enable_read_bucket_metric self.mode = "r" def read(self, size=-1): bytesRead = super().read(size) if self.enable_read_bucket_metric: Metrics.counter( self.__class__, "GCS_read_bytes_counter_" + self._blob.bucket.name).inc( len(bytesRead)) return bytesRead class BeamBlobWriter(BlobWriter): def __init__( self, blob, content_type, chunk_size=16 * 1024 * 1024, ignore_flush=True, enable_write_bucket_metric=False, retry=DEFAULT_RETRY): super().__init__( blob, content_type=content_type, chunk_size=chunk_size, ignore_flush=ignore_flush, retry=retry) self.mode = "w" self.enable_write_bucket_metric = enable_write_bucket_metric def write(self, b): bytesWritten = super().write(b) if self.enable_write_bucket_metric: Metrics.counter( self.__class__, "GCS_write_bytes_counter_" + self._blob.bucket.name).inc(bytesWritten) return bytesWritten ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/gcsio_integration_test.py0000644000175100017510000003106515206113032024162 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration tests for gcsio module. Runs tests against Google Cloud Storage service. Instantiates a TestPipeline to get options such as GCP project name, but doesn't actually start a Beam pipeline or test any specific runner. Run the following in 'sdks/python' directory to run these tests manually: scripts/run_integration_test.sh \ --test_opts apache_beam/io/gcp/gcsio_integration_test.py """ # pytype: skip-file import logging import time import unittest import uuid import zlib import mock import pytest from parameterized import parameterized from parameterized import parameterized_class from apache_beam import Create from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystems import FileSystems from apache_beam.io.textio import ReadAllFromText from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.combiners import Count try: from apache_beam.io.gcp import gcsio except ImportError: gcsio = None # type: ignore try: from google.api_core.exceptions import NotFound except ImportError: NotFound = None # Number of seconds to wait for bucket deletion or creation to propagate. WAIT_BUCKET_PROPAGATION_SECONDS = 60 @unittest.skipIf(gcsio is None, 'GCP dependencies are not installed') @parameterized_class( ('no_gcsio_throttling_counter', 'enable_gcsio_blob_generation'), [(False, False), (False, True), (True, False), (True, True)]) class GcsIOIntegrationTest(unittest.TestCase): INPUT_FILE = 'gs://dataflow-samples/shakespeare/kinglear.txt' def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.runner_name = type(self.test_pipeline.runner).__name__ if self.runner_name != 'TestDataflowRunner': # This test doesn't run a pipeline, so it doesn't make sense to try it on # different runners. Running with TestDataflowRunner makes sense since # it uses GoogleCloudOptions such as 'project'. raise unittest.SkipTest('This test only runs with TestDataflowRunner.') self.project = self.test_pipeline.get_option('project') self.gcs_tempdir = ( self.test_pipeline.get_option('temp_location') + '/gcs_it-' + str(uuid.uuid4())) def tearDown(self): FileSystems.delete([self.gcs_tempdir + '/']) def _verify_copy(self, src, dest, dest_kms_key_name=None): self.assertTrue( FileSystems.exists(src), 'source file does not exist: %s' % src) self.assertTrue( FileSystems.exists(dest), 'copied file not present in destination: %s' % dest) src_checksum = self.gcsio.checksum(src) dest_checksum = self.gcsio.checksum(dest) self.assertEqual(src_checksum, dest_checksum) actual_dest_kms_key = self.gcsio.kms_key(dest) if actual_dest_kms_key is None: self.assertEqual(actual_dest_kms_key, dest_kms_key_name) else: self.assertTrue( actual_dest_kms_key.startswith(dest_kms_key_name), "got: %s, wanted startswith: %s" % (actual_dest_kms_key, dest_kms_key_name)) @pytest.mark.it_postcommit def test_copy(self): self.gcsio = gcsio.GcsIO( pipeline_options={ "no_gcsio_throttling_counter": self.no_gcsio_throttling_counter, "enable_gcsio_blob_generation": self.enable_gcsio_blob_generation }) src = self.INPUT_FILE dest = self.gcs_tempdir + '/test_copy' self.gcsio.copy(src, dest) self._verify_copy(src, dest) unknown_src = self.test_pipeline.get_option('temp_location') + \ '/gcs_it-' + str(uuid.uuid4()) with self.assertRaises(NotFound): self.gcsio.copy(unknown_src, dest) @pytest.mark.it_postcommit def test_copy_and_delete(self): self.gcsio = gcsio.GcsIO( pipeline_options={ "no_gcsio_throttling_counter": self.no_gcsio_throttling_counter, "enable_gcsio_blob_generation": self.enable_gcsio_blob_generation }) src = self.INPUT_FILE dest = self.gcs_tempdir + '/test_copy' self.gcsio.copy(src, dest) self._verify_copy(src, dest) self.gcsio.delete(dest) # no exception if we delete an nonexistent file. self.gcsio.delete(dest) @pytest.mark.it_postcommit def test_batch_copy_and_delete(self): self.gcsio = gcsio.GcsIO( pipeline_options={ "no_gcsio_throttling_counter": self.no_gcsio_throttling_counter, "enable_gcsio_blob_generation": self.enable_gcsio_blob_generation }) num_copies = 10 srcs = [self.INPUT_FILE] * num_copies dests = [ self.gcs_tempdir + '/test_copy_batch_%d' % i for i in range(num_copies) ] src_dest_pairs = list(zip(srcs, dests)) copy_results = self.gcsio.copy_batch(src_dest_pairs) self.assertEqual(len(copy_results), len(src_dest_pairs)) for pair, result in list(zip(src_dest_pairs, copy_results)): self._verify_copy(pair[0], pair[1]) self.assertEqual( pair[0], result[0], 'copy source %s does not match %s' % (pair[0], str(result))) self.assertEqual( pair[1], result[1], 'copy destination %s does not match %s' % (pair[1], result[1])) self.assertFalse( result[2], 'response code %s indicates that copy operation did not succeed' % result[2]) delete_results = self.gcsio.delete_batch(dests) self.assertEqual(len(delete_results), len(dests)) for dest, result in list(zip(dests, delete_results)): self.assertFalse( FileSystems.exists(dest), 'deleted file still exists: %s' % dest) self.assertEqual( dest, result[0], 'delete path %s does not match %s' % (dest, result[0])) self.assertFalse( result[1], 'response code %s indicates that delete operation did not succeed' % result[1]) redelete_results = self.gcsio.delete_batch(dests) for dest, result in list(zip(dests, redelete_results)): self.assertFalse( result[1], 're-delete should not throw error: %s' % result[1]) @pytest.mark.it_postcommit @mock.patch('apache_beam.io.gcp.gcsio.default_gcs_bucket_name') @unittest.skipIf(NotFound is None, 'GCP dependencies are not installed') def test_create_default_bucket(self, mock_default_gcs_bucket_name): self.gcsio = gcsio.GcsIO() google_cloud_options = self.test_pipeline.options.view_as( GoogleCloudOptions) # overwrite kms option here, because get_or_create_default_gcs_bucket() # requires this option unset. google_cloud_options.dataflow_kms_key = None import random from hashlib import blake2b # Add a random number to avoid collision if multiple test instances # are run at the same time. To avoid too many dangling buckets if bucket # removal fails, we limit the max number of possible bucket names in this # test to 1000. overridden_bucket_name = 'gcsio-it-%d-%s-%s-%d' % ( random.randint(0, 999), google_cloud_options.region, blake2b(google_cloud_options.project.encode('utf8'), digest_size=4).hexdigest(), int(time.time())) mock_default_gcs_bucket_name.return_value = overridden_bucket_name # remove the existing bucket with the same name as the default bucket existing_bucket = self.gcsio.get_bucket(overridden_bucket_name) if existing_bucket: try: existing_bucket.delete() time.sleep(WAIT_BUCKET_PROPAGATION_SECONDS) except NotFound: # Bucket existence check from get_bucket may be inaccurate due to gcs # cache or delay pass bucket = gcsio.get_or_create_default_gcs_bucket(google_cloud_options) time.sleep(WAIT_BUCKET_PROPAGATION_SECONDS) self.assertIsNotNone(bucket) self.assertEqual(bucket.name, overridden_bucket_name) # verify soft delete policy is disabled by default in the default bucket # after creation self.assertEqual(bucket.soft_delete_policy.retention_duration_seconds, 0) max_retries = 5 retry_delay = 1 existing_bucket = None for attempt in range(max_retries): try: existing_bucket = self.gcsio.get_bucket(overridden_bucket_name) break except NotFound: if attempt < max_retries - 1: time.sleep(retry_delay) retry_delay *= 2 else: existing_bucket = None if existing_bucket: try: existing_bucket.delete() except NotFound: pass time.sleep(WAIT_BUCKET_PROPAGATION_SECONDS) bucket_after_delete = self.gcsio.get_bucket(overridden_bucket_name) self.assertIsNone(bucket_after_delete) class GcsIOReadGzipTest(unittest.TestCase): gcs_path_prefix = "gs://apache-beam-samples/textio/" gzip_test_files = [ "textio-test-data.content-type-gzip-content-encoding-gzip.1k.txt.gz", "textio-test-data.content-type-gzip-content-encoding-none.1k.txt.gz", "textio-test-data.content-type-none-content-encoding-gzip.1k.txt.gz", "textio-test-data.content-type-none-content-encoding-none.1k.txt.gz", "textio-test-data.content-type-text-content-encoding-gzip.1k.txt.gz", "textio-test-data.content-type-text-content-encoding-none.1k.txt.gz", "textio-test-data.default.1k.txt", "textio-test-data.default.1k.txt.gz", "textio-test-data.gzip-local.1k.txt.gz", ] @parameterized.expand([ (gzip_test_files[0], CompressionTypes.UNCOMPRESSED, NotImplementedError), (gzip_test_files[0], CompressionTypes.GZIP, NotImplementedError), (gzip_test_files[0], CompressionTypes.AUTO, NotImplementedError), (gzip_test_files[1], CompressionTypes.UNCOMPRESSED, UnicodeDecodeError), (gzip_test_files[1], CompressionTypes.GZIP, None), (gzip_test_files[1], CompressionTypes.AUTO, None), (gzip_test_files[2], CompressionTypes.UNCOMPRESSED, UnicodeDecodeError), (gzip_test_files[2], CompressionTypes.GZIP, None), (gzip_test_files[2], CompressionTypes.AUTO, None), (gzip_test_files[3], CompressionTypes.UNCOMPRESSED, UnicodeDecodeError), (gzip_test_files[3], CompressionTypes.GZIP, None), (gzip_test_files[3], CompressionTypes.AUTO, None), (gzip_test_files[4], CompressionTypes.UNCOMPRESSED, UnicodeDecodeError), (gzip_test_files[4], CompressionTypes.GZIP, None), (gzip_test_files[4], CompressionTypes.AUTO, None), (gzip_test_files[5], CompressionTypes.UNCOMPRESSED, UnicodeDecodeError), (gzip_test_files[5], CompressionTypes.GZIP, None), (gzip_test_files[5], CompressionTypes.AUTO, None), (gzip_test_files[6], CompressionTypes.UNCOMPRESSED, None), (gzip_test_files[6], CompressionTypes.GZIP, zlib.error), (gzip_test_files[6], CompressionTypes.AUTO, None), (gzip_test_files[7], CompressionTypes.UNCOMPRESSED, UnicodeDecodeError), (gzip_test_files[7], CompressionTypes.GZIP, None), (gzip_test_files[7], CompressionTypes.AUTO, None), (gzip_test_files[8], CompressionTypes.UNCOMPRESSED, UnicodeDecodeError), (gzip_test_files[8], CompressionTypes.GZIP, None), (gzip_test_files[8], CompressionTypes.AUTO, None), ]) @unittest.skipIf(NotFound is None, 'GCP dependencies are not installed') def test_read_gzip_file(self, file_name, compression_type, exception): p = TestPipeline(runner="Direct", is_integration_test=True) r = ( p | Create([f"{GcsIOReadGzipTest.gcs_path_prefix}{file_name}"]) | "Read File from GCS" >> ReadAllFromText(compression_type=compression_type) | Count.Globally()) assert_that(r, equal_to([1000])) if exception is None: result = p.run() result.wait_until_finish() else: with self.assertRaises(exception): result = p.run() result.wait_until_finish() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/gcsio_retry.py0000644000175100017510000000676715206113032021760 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Throttling Handler for GCSIO """ import inspect import logging import math from itertools import tee from google.api_core import __version__ as _api_core_version from google.api_core import exceptions as api_exceptions from google.api_core import retry from google.cloud.storage.retry import DEFAULT_RETRY from google.cloud.storage.retry import _should_retry # pylint: disable=protected-access from packaging import version from apache_beam.metrics.metric import Metrics from apache_beam.options.pipeline_options import GoogleCloudOptions _LOGGER = logging.getLogger(__name__) __all__ = ['DEFAULT_RETRY_WITH_THROTTLING_COUNTER'] _MIN_SLEEP_ARG_SWITCH_VERSION = version.parse("2.25.0rc0") _LEGACY_SLEEP_ARG_NAME = "next_sleep" _CURRENT_SLEEP_ARG_NAME = "sleep_iterator" class ThrottlingHandler(object): _THROTTLED_SECS = Metrics.counter('gcsio', "cumulativeThrottlingSeconds") def __init__(self): # decide which arg name google-api-core uses try: core_ver = version.parse(_api_core_version) except Exception: core_ver = version.parse("0") if core_ver < _MIN_SLEEP_ARG_SWITCH_VERSION: self._sleep_arg = _LEGACY_SLEEP_ARG_NAME else: self._sleep_arg = _CURRENT_SLEEP_ARG_NAME def __call__(self, exc): if isinstance(exc, api_exceptions.TooManyRequests): _LOGGER.debug('Caught GCS quota error (%s), retrying.', exc.reason) # TODO: revisit the logic here when gcs client library supports error # callbacks frame = inspect.currentframe() if frame is None: _LOGGER.warning('cannot inspect the current stack frame') return prev_frame = frame.f_back if prev_frame is None: _LOGGER.warning('cannot inspect the caller stack frame') return # Determine which retry helper argument to inspect in # google/api_core/retry/retry_base.py’s _retry_error_helper(): # - versions < 2.25.0rc0 use “next_sleep” # - versions ≥ 2.25.0rc0 use “sleep_iterator” if self._sleep_arg == _LEGACY_SLEEP_ARG_NAME: sleep_seconds = prev_frame.f_locals.get(self._sleep_arg, 0) else: sleep_iterator = prev_frame.f_locals.get(self._sleep_arg, iter([])) sleep_iterator, sleep_iterator_copy = tee(sleep_iterator) try: sleep_seconds = next(sleep_iterator_copy) except StopIteration: sleep_seconds = 0 ThrottlingHandler._THROTTLED_SECS.inc(math.ceil(sleep_seconds)) DEFAULT_RETRY_WITH_THROTTLING_COUNTER = retry.Retry( predicate=_should_retry, on_error=ThrottlingHandler()) def get_retry(pipeline_options): if pipeline_options.view_as(GoogleCloudOptions).no_gcsio_throttling_counter: return DEFAULT_RETRY else: return DEFAULT_RETRY_WITH_THROTTLING_COUNTER ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/gcsio_retry_test.py0000644000175100017510000000571415206113032023006 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for Throttling Handler of GCSIO.""" import unittest from unittest.mock import Mock from apache_beam.metrics.execution import MetricsContainer from apache_beam.metrics.execution import MetricsEnvironment from apache_beam.metrics.metricbase import MetricName from apache_beam.runners.worker import statesampler from apache_beam.utils import counters # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from google.api_core import exceptions as api_exceptions from apache_beam.io.gcp import gcsio_retry except ImportError: gcsio_retry = None api_exceptions = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports @unittest.skipIf((gcsio_retry is None or api_exceptions is None), 'GCP dependencies are not installed') class TestGCSIORetry(unittest.TestCase): def test_retry_on_non_retriable(self): mock = Mock(side_effect=[ Exception('Something wrong!'), ]) retry = gcsio_retry.DEFAULT_RETRY_WITH_THROTTLING_COUNTER with self.assertRaises(Exception): retry(mock)() def test_retry_on_throttling(self): mock = Mock( side_effect=[ api_exceptions.TooManyRequests("Slow down!"), api_exceptions.TooManyRequests("Slow down again!"), 12345 ]) retry = gcsio_retry.DEFAULT_RETRY_WITH_THROTTLING_COUNTER sampler = statesampler.StateSampler('', counters.CounterFactory()) statesampler.set_current_tracker(sampler) state = sampler.scoped_state( 'my_step', 'my_state', metrics_container=MetricsContainer('my_step')) try: sampler.start() with state: container = MetricsEnvironment.current_container() self.assertEqual( container.get_counter( MetricName('gcsio', "cumulativeThrottlingSeconds")).get_cumulative(), 0) self.assertEqual(12345, retry(mock)()) self.assertGreater( container.get_counter( MetricName('gcsio', "cumulativeThrottlingSeconds")).get_cumulative(), 1) finally: sampler.stop() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/gcsio_test.py0000644000175100017510000007052015206113032021556 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for Google Cloud Storage client.""" # pytype: skip-file import logging import os import random import unittest from datetime import datetime import mock from apache_beam import version as beam_version from apache_beam.metrics.execution import MetricsContainer from apache_beam.metrics.execution import MetricsEnvironment from apache_beam.metrics.metricbase import MetricName from apache_beam.pipeline import PipelineOptions from apache_beam.runners.worker import statesampler from apache_beam.utils import counters # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from google.cloud.exceptions import BadRequest from google.cloud.exceptions import NotFound from apache_beam.io.gcp import gcsio from apache_beam.io.gcp.gcsio_retry import DEFAULT_RETRY_WITH_THROTTLING_COUNTER except ImportError: NotFound = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports DEFAULT_GCP_PROJECT = 'apache-beam-testing' class FakeGcsClient(object): # Fake storage client. def __init__(self): self.buckets = {} def _add_bucket(self, bucket): self.buckets[bucket.name] = bucket return self.buckets[bucket.name] def bucket(self, name): return FakeBucket(self, name) def create_bucket(self, name): return self._add_bucket(self.bucket(name)) def get_bucket(self, name): if name in self.buckets: return self.buckets[name] else: raise NotFound("Bucket not found") def lookup_bucket(self, name): if name in self.buckets: return self.buckets[name] else: return self.create_bucket(name) def batch(self, raise_exception=True): # Return a mock object configured to act as a context manager # and provide the necessary _responses attribute after __exit__. # test_delete performs 3 deletions. num_expected_responses = 3 mock_batch = mock.Mock() # Configure the mock responses (assuming success for test_delete) # These need to be available *after* the 'with' block finishes. # We'll store them temporarily and assign in __exit__. successful_responses = [ mock.Mock(status_code=204) for _ in range(num_expected_responses) ] # Define the exit logic def mock_exit_logic(exc_type, exc_val, exc_tb): # Assign responses to the mock instance itself # so they are available after the 'with' block. mock_batch._responses = successful_responses # Configure the mock to behave like a context manager mock_batch.configure_mock( __enter__=mock.Mock(return_value=mock_batch), __exit__=mock.Mock(side_effect=mock_exit_logic)) # The loop inside _batch_with_retry calls fn(request) for each item. # The real batch object might have methods like add() or similar, # but the core logic in gcsio.py calls the passed function `fn` directly # within the `with` block. So, no specific action methods seem needed # on the mock_batch itself for this test case. return mock_batch def add_file(self, bucket, blob, contents): folder = self.lookup_bucket(bucket) holder = folder.lookup_blob(blob) holder.contents = contents def get_file(self, bucket, blob): folder = self.get_bucket(bucket.name) holder = folder.get_blob(blob.name) return holder def list_blobs(self, bucket_or_path, prefix=None, **unused_kwargs): bucket = self.get_bucket(bucket_or_path.name) if not prefix: return list(bucket.blobs.values()) else: output = [] for name in list(bucket.blobs): if name[0:len(prefix)] == prefix: output.append(bucket.blobs[name]) return output class FakeBucket(object): #Fake bucket for storing test blobs locally. def __init__(self, client, name): self.client = client self.name = name self.blobs = {} self.default_kms_key_name = None def _get_canonical_bucket(self): return self.client.get_bucket(self.name) def _create_blob(self, name): return FakeBlob(name, self) def add_blob(self, blob): bucket = self._get_canonical_bucket() bucket.blobs[blob.name] = blob return bucket.blobs[blob.name] def blob(self, name): return self._create_blob(name) def copy_blob(self, blob, dest, new_name=None, **kwargs): if self.get_blob(blob.name) is None: raise NotFound("source blob not found") if not new_name: new_name = blob.name new_blob = FakeBlob(new_name, dest) dest.add_blob(new_blob) return new_blob def get_blob(self, blob_name, **unused_kwargs): bucket = self._get_canonical_bucket() if blob_name in bucket.blobs: return bucket.blobs[blob_name] else: return None def lookup_blob(self, name): bucket = self._get_canonical_bucket() if name in bucket.blobs: return bucket.blobs[name] else: new_blob = bucket._create_blob(name) return bucket.add_blob(new_blob) def set_default_kms_key_name(self, name): self.default_kms_key_name = name def delete_blob(self, name, **kwargs): bucket = self._get_canonical_bucket() if name in bucket.blobs: del bucket.blobs[name] def list_blobs(self, prefix=None, **kwargs): bucket = self._get_canonical_bucket() return self.client.list_blobs(bucket, prefix, **kwargs) class FakeBlob(object): def __init__( self, name, bucket, size=0, contents=None, generation=1, crc32c=None, kms_key_name=None, updated=None, fail_when_getting_metadata=False, fail_when_reading=False): self.name = name self.bucket = bucket self.size = size self.contents = contents self._generation = generation self.crc32c = crc32c self.kms_key_name = kms_key_name self.updated = updated self._fail_when_getting_metadata = fail_when_getting_metadata self._fail_when_reading = fail_when_reading self.generation = random.randint(0, (1 << 63) - 1) self.content_encoding = None self.content_type = None def reload(self): pass def delete(self): self.bucket.delete_blob(self.name) def download_as_bytes(self, **kwargs): blob = self.bucket.get_blob(self.name) if blob is None: raise NotFound("blob not found") return blob.contents def __eq__(self, other): return self.bucket.get_blob(self.name) is other.bucket.get_blob(other.name) def exists(self, **kwargs): return self.bucket.get_blob(self.name) is not None @unittest.skipIf(NotFound is None, 'GCP dependencies are not installed') class TestGCSPathParser(unittest.TestCase): BAD_GCS_PATHS = [ 'gs://', 'gs://bucket', 'gs:///name', 'gs:///', 'gs:/blah/bucket/name', ] def test_gcs_path(self): self.assertEqual( gcsio.parse_gcs_path('gs://bucket/name'), ('bucket', 'name')) self.assertEqual( gcsio.parse_gcs_path('gs://bucket/name/sub'), ('bucket', 'name/sub')) def test_bad_gcs_path(self): for path in self.BAD_GCS_PATHS: self.assertRaises(ValueError, gcsio.parse_gcs_path, path) self.assertRaises(ValueError, gcsio.parse_gcs_path, 'gs://bucket/') def test_gcs_path_object_optional(self): self.assertEqual( gcsio.parse_gcs_path('gs://bucket/name', object_optional=True), ('bucket', 'name')) self.assertEqual( gcsio.parse_gcs_path('gs://bucket/', object_optional=True), ('bucket', '')) def test_bad_gcs_path_object_optional(self): for path in self.BAD_GCS_PATHS: self.assertRaises(ValueError, gcsio.parse_gcs_path, path, True) class SampleOptions(object): def __init__(self, project, region, kms_key=None): self.project = DEFAULT_GCP_PROJECT self.region = region self.dataflow_kms_key = kms_key _DEFAULT_UNIVERSE_DOMAIN = "googleapis.com" def _make_credentials(project=None, universe_domain=_DEFAULT_UNIVERSE_DOMAIN): import google.auth.credentials if project is not None: return mock.Mock( spec=google.auth.credentials.Credentials, project_id=project, universe_domain=universe_domain, ) return mock.Mock( spec=google.auth.credentials.Credentials, universe_domain=universe_domain) @unittest.skipIf(NotFound is None, 'GCP dependencies are not installed') class TestGCSIO(unittest.TestCase): def _insert_random_file( self, client, path, size=0, crc32c=None, kms_key_name=None, updated=None, fail_when_getting_metadata=False, fail_when_reading=False): bucket_name, blob_name = gcsio.parse_gcs_path(path) bucket = client.lookup_bucket(bucket_name) blob = FakeBlob( blob_name, bucket, size=size, contents=os.urandom(size), crc32c=crc32c, kms_key_name=kms_key_name, updated=updated, fail_when_getting_metadata=fail_when_getting_metadata, fail_when_reading=fail_when_reading) bucket.add_blob(blob) return blob def setUp(self): self.client = FakeGcsClient() self.gcs = gcsio.GcsIO(self.client) self.client.create_bucket("gcsio-test") def test_read_bucket_metric(self): sampler = statesampler.StateSampler('', counters.CounterFactory()) statesampler.set_current_tracker(sampler) state1 = sampler.scoped_state( 'mystep', 'myState', metrics_container=MetricsContainer('mystep')) try: sampler.start() with state1: client = FakeGcsClient() gcs = gcsio.GcsIO(client, {"enable_bucket_read_metric_counter": True}) client.create_bucket("gcsio-test") file_name = 'gs://gcsio-test/dummy_file' file_size = 1234 self._insert_random_file(client, file_name, file_size) reader = gcs.open(file_name, 'r') reader.read() container = MetricsEnvironment.current_container() self.assertEqual( container.get_counter( MetricName( "apache_beam.io.gcp.gcsio.BeamBlobReader", "GCS_read_bytes_counter_gcsio-test")).get_cumulative(), file_size) finally: sampler.stop() def test_write_bucket_metric(self): sampler = statesampler.StateSampler('', counters.CounterFactory()) statesampler.set_current_tracker(sampler) state1 = sampler.scoped_state( 'mystep', 'myState', metrics_container=MetricsContainer('mystep')) try: sampler.start() with state1: client = FakeGcsClient() gcs = gcsio.GcsIO(client, {"enable_bucket_write_metric_counter": True}) client.create_bucket("gcsio-test") file_name = 'gs://gcsio-test/dummy_file' gcsFile = gcs.open(file_name, 'w') gcsFile.write(str.encode("some text")) container = MetricsEnvironment.current_container() self.assertEqual( container.get_counter( MetricName( "apache_beam.io.gcp.gcsio.BeamBlobWriter", "GCS_write_bytes_counter_gcsio-test")).get_cumulative(), 9) finally: sampler.stop() def test_default_bucket_name(self): self.assertEqual( gcsio.default_gcs_bucket_name(DEFAULT_GCP_PROJECT, "us-central1"), 'dataflow-staging-us-central1-77b801c0838aee13391c0d1885860494') def test_default_bucket_name_failure(self): self.assertEqual( gcsio.get_or_create_default_gcs_bucket( SampleOptions( DEFAULT_GCP_PROJECT, "us-central1", kms_key="kmskey!")), None) def test_exists(self): file_name = 'gs://gcsio-test/dummy_file' file_size = 1234 self._insert_random_file(self.client, file_name, file_size) self.assertFalse(self.gcs.exists(file_name + 'xyz')) self.assertTrue(self.gcs.exists(file_name)) @mock.patch.object(FakeBlob, 'exists') def test_exists_failure(self, mock_get): # Raising an error other than 404. Raising 404 is a valid failure for # exists() call. mock_get.side_effect = BadRequest("Try again") file_name = 'gs://gcsio-test/dummy_file' file_size = 1234 self._insert_random_file(self.client, file_name, file_size) with self.assertRaises(BadRequest): self.gcs.exists(file_name) def test_checksum(self): file_name = 'gs://gcsio-test/dummy_file' file_size = 1234 checksum = 'deadbeef' self._insert_random_file(self.client, file_name, file_size, crc32c=checksum) self.assertTrue(self.gcs.exists(file_name)) self.assertEqual(checksum, self.gcs.checksum(file_name)) def test_size(self): file_name = 'gs://gcsio-test/dummy_file' file_size = 1234 self._insert_random_file(self.client, file_name, file_size) self.assertTrue(self.gcs.exists(file_name)) self.assertEqual(1234, self.gcs.size(file_name)) def test_kms_key(self): file_name = 'gs://gcsio-test/dummy_file' file_size = 1234 kms_key_name = "dummy" self._insert_random_file( self.client, file_name, file_size, kms_key_name=kms_key_name) self.assertTrue(self.gcs.exists(file_name)) self.assertEqual(kms_key_name, self.gcs.kms_key(file_name)) def test_last_updated(self): file_name = 'gs://gcsio-test/dummy_file' file_size = 1234 updated = datetime.fromtimestamp(123456.78) self._insert_random_file(self.client, file_name, file_size, updated=updated) self.assertTrue(self.gcs.exists(file_name)) self.assertEqual( gcsio.GcsIO._updated_to_seconds(updated), self.gcs.last_updated(file_name)) def test_file_status(self): file_name = 'gs://gcsio-test/dummy_file' file_size = 1234 updated = datetime.fromtimestamp(123456.78) checksum = 'deadbeef' self._insert_random_file( self.client, file_name, file_size, updated=updated, crc32c=checksum) file_checksum = self.gcs.checksum(file_name) file_status = self.gcs._status(file_name) self.assertEqual(file_status['size'], file_size) self.assertEqual(file_status['checksum'], file_checksum) self.assertEqual( file_status['updated'], gcsio.GcsIO._updated_to_seconds(updated)) def test_file_mode_calls(self): file_name = 'gs://gcsio-test/dummy_mode_file' self._insert_random_file(self.client, file_name) with mock.patch('apache_beam.io.gcp.gcsio.BeamBlobWriter') as writer: self.gcs.open(file_name, 'wb') writer.assert_called() with mock.patch('apache_beam.io.gcp.gcsio.BeamBlobReader') as reader: self.gcs.open(file_name, 'rb') reader.assert_called() def test_bad_file_modes(self): file_name = 'gs://gcsio-test/dummy_mode_file' self._insert_random_file(self.client, file_name) with self.assertRaises(ValueError): self.gcs.open(file_name, 'w+') with self.assertRaises(ValueError): self.gcs.open(file_name, 'r+b') def test_delete(self): # File path. file_name = 'gs://gcsio-test/delete_me' file_size = 1024 bucket_name, blob_name = gcsio.parse_gcs_path(file_name) # Test deletion of non-existent file. bucket = self.client.get_bucket(bucket_name) self.gcs.delete(file_name) # Insert a random file for testing. self._insert_random_file(self.client, file_name, file_size) self.assertTrue(blob_name in bucket.blobs) # Deleting the file. self.gcs.delete(file_name) self.assertFalse(blob_name in bucket.blobs) # Now test deleting a directory (prefix) with multiple files. prefix = 'gs://gcsio-test/directory_to_delete/' file_names = [f"{prefix}file1", f"{prefix}file2", f"{prefix}file3"] blobs = [gcsio.parse_gcs_path(file_name) for file_name in file_names] # Insert random files under the prefix. for file_name in file_names: self._insert_random_file(self.client, file_name, file_size) # Verify the files exist before deletion for blob in blobs: self.assertTrue(blob[1] in bucket.blobs) # Deleting the directory (all files under the prefix). self.gcs.delete(prefix, recursive=True) # Verify that the files are deleted. for blob in blobs: self.assertFalse(blob[1] in bucket.blobs) def test_copy(self): src_file_name = 'gs://gcsio-test/source' dest_file_name = 'gs://gcsio-test/dest' src_bucket_name, src_blob_name = gcsio.parse_gcs_path(src_file_name) dest_bucket_name, dest_blob_name = gcsio.parse_gcs_path(dest_file_name) src_bucket = self.client.lookup_bucket(src_bucket_name) dest_bucket = self.client.lookup_bucket(dest_bucket_name) file_size = 1024 self._insert_random_file(self.client, src_file_name, file_size) self.assertTrue(src_blob_name in src_bucket.blobs) self.assertFalse(dest_blob_name in dest_bucket.blobs) self.gcs.copy(src_file_name, dest_file_name) self.assertTrue(src_blob_name in src_bucket.blobs) self.assertTrue(dest_blob_name in dest_bucket.blobs) # Test copy of non-existent files. with self.assertRaises(NotFound): self.gcs.copy( 'gs://gcsio-test/non-existent', 'gs://gcsio-test/non-existent-destination') @staticmethod def _fake_batch_responses(status_codes): return mock.Mock( __enter__=mock.Mock(), __exit__=mock.Mock(), _responses=[ mock.Mock( **{ 'json.return_value': { 'error': { 'message': 'error' } }, 'request.method': 'BATCH', 'request.url': 'contentid://None', }, status_code=code, ) for code in status_codes ], ) @mock.patch('apache_beam.io.gcp.gcsio.MAX_BATCH_OPERATION_SIZE', 3) @mock.patch('time.sleep', mock.Mock()) def test_copy_batch(self): src_dest_pairs = [ (f'gs://source_bucket/file{i}.txt', f'gs://dest_bucket/file{i}.txt') for i in range(7) ] gcs_io = gcsio.GcsIO( storage_client=mock.Mock( batch=mock.Mock( side_effect=[ self._fake_batch_responses([200, 404, 429]), self._fake_batch_responses([429]), self._fake_batch_responses([429]), self._fake_batch_responses([200]), self._fake_batch_responses([200, 429, 200]), self._fake_batch_responses([200]), self._fake_batch_responses([200]), ]), )) results = gcs_io.copy_batch(src_dest_pairs) expected = [ ('gs://source_bucket/file0.txt', 'gs://dest_bucket/file0.txt', None), ('gs://source_bucket/file1.txt', 'gs://dest_bucket/file1.txt', 404), ('gs://source_bucket/file2.txt', 'gs://dest_bucket/file2.txt', None), ('gs://source_bucket/file3.txt', 'gs://dest_bucket/file3.txt', None), ('gs://source_bucket/file4.txt', 'gs://dest_bucket/file4.txt', None), ('gs://source_bucket/file5.txt', 'gs://dest_bucket/file5.txt', None), ('gs://source_bucket/file6.txt', 'gs://dest_bucket/file6.txt', None), ] self.assertEqual(results, expected) @mock.patch('time.sleep', mock.Mock()) @mock.patch('time.monotonic', mock.Mock(side_effect=[0, 120])) def test_copy_batch_timeout_exceeded(self): src_dest_pairs = [ ('gs://source_bucket/file0.txt', 'gs://dest_bucket/file0.txt') ] gcs_io = gcsio.GcsIO( storage_client=mock.Mock( batch=mock.Mock(side_effect=[self._fake_batch_responses([429])]))) results = gcs_io.copy_batch(src_dest_pairs) expected = [ ('gs://source_bucket/file0.txt', 'gs://dest_bucket/file0.txt', 429), ] self.assertEqual(results, expected) def test_copytree(self): src_dir_name = 'gs://gcsio-test/source/' dest_dir_name = 'gs://gcsio-test/dest/' file_size = 1024 paths = ['a', 'b/c', 'b/d'] for path in paths: src_file_name = src_dir_name + path dest_file_name = dest_dir_name + path src_bucket_name, src_blob_name = gcsio.parse_gcs_path(src_file_name) dest_bucket_name, dest_blob_name = gcsio.parse_gcs_path(dest_file_name) src_bucket = self.client.lookup_bucket(src_bucket_name) dest_bucket = self.client.lookup_bucket(dest_bucket_name) file_size = 1024 self._insert_random_file(self.client, src_file_name, file_size) self.assertTrue(src_blob_name in src_bucket.blobs) self.assertFalse(dest_blob_name in dest_bucket.blobs) self.gcs.copytree(src_dir_name, dest_dir_name) for path in paths: src_file_name = src_dir_name + path dest_file_name = dest_dir_name + path src_bucket_name, src_blob_name = gcsio.parse_gcs_path(src_file_name) dest_bucket_name, dest_blob_name = gcsio.parse_gcs_path(dest_file_name) src_bucket = self.client.lookup_bucket(src_bucket_name) dest_bucket = self.client.lookup_bucket(dest_bucket_name) self.assertTrue(src_blob_name in src_bucket.blobs) self.assertTrue(dest_blob_name in dest_bucket.blobs) def test_rename(self): src_file_name = 'gs://gcsio-test/source' dest_file_name = 'gs://gcsio-test/dest' src_bucket_name, src_blob_name = gcsio.parse_gcs_path(src_file_name) dest_bucket_name, dest_blob_name = gcsio.parse_gcs_path(dest_file_name) file_size = 1024 self._insert_random_file(self.client, src_file_name, file_size) src_bucket = self.client.lookup_bucket(src_bucket_name) dest_bucket = self.client.lookup_bucket(dest_bucket_name) self.assertTrue(src_blob_name in src_bucket.blobs) self.assertFalse(dest_blob_name in dest_bucket.blobs) self.gcs.rename(src_file_name, dest_file_name) self.assertFalse(src_blob_name in src_bucket.blobs) self.assertTrue(dest_blob_name in dest_bucket.blobs) def test_file_buffered_read_call(self): file_name = 'gs://gcsio-test/read_line_file' read_buffer_size = 1024 self._insert_random_file(self.client, file_name, 10240) bucket_name, blob_name = gcsio.parse_gcs_path(file_name) bucket = self.client.get_bucket(bucket_name) blob = bucket.get_blob(blob_name) with mock.patch('apache_beam.io.gcp.gcsio.BeamBlobReader') as reader: self.gcs.open(file_name, read_buffer_size=read_buffer_size) reader.assert_called_with( blob, chunk_size=read_buffer_size, enable_read_bucket_metric=False, retry=DEFAULT_RETRY_WITH_THROTTLING_COUNTER) def test_file_write_call(self): file_name = 'gs://gcsio-test/write_file' with mock.patch('apache_beam.io.gcp.gcsio.BeamBlobWriter') as writer: self.gcs.open(file_name, 'w') writer.assert_called() def test_list_files(self): bucket_name = 'gcsio-test' objects = [ ('cow/cat/fish', 2), ('cow/cat/blubber', 3), ('cow/dog/blubber', 4), ] for (object_name, size) in objects: file_name = 'gs://%s/%s' % (bucket_name, object_name) self._insert_random_file(self.client, file_name, size) test_cases = [ ( 'gs://gcsio-test/c', [ ('cow/cat/fish', 2), ('cow/cat/blubber', 3), ('cow/dog/blubber', 4), ]), ( 'gs://gcsio-test/cow/', [ ('cow/cat/fish', 2), ('cow/cat/blubber', 3), ('cow/dog/blubber', 4), ]), ('gs://gcsio-test/cow/cat/fish', [ ('cow/cat/fish', 2), ]), ] for file_pattern, expected_object_names in test_cases: expected_file_names = [('gs://%s/%s' % (bucket_name, object_name), size) for (object_name, size) in expected_object_names] self.assertEqual( set(self.gcs.list_files(file_pattern)), set(expected_file_names)) def test_downloader_fail_non_existent_object(self): file_name = 'gs://gcsio-metrics-test/dummy_mode_file' with self.assertRaises(NotFound): with self.gcs.open(file_name, 'r') as f: f.read(1) def test_blob_delete(self): file_name = 'gs://gcsio-test/delete_me' file_size = 1024 bucket_name, blob_name = gcsio.parse_gcs_path(file_name) # Test deletion of non-existent file. bucket = self.client.get_bucket(bucket_name) self.gcs.delete(file_name) self._insert_random_file(self.client, file_name, file_size) self.assertTrue(blob_name in bucket.blobs) blob = bucket.get_blob(blob_name) self.assertIsNotNone(blob) blob.delete() self.assertFalse(blob_name in bucket.blobs) @mock.patch('google.cloud._http.JSONConnection._do_request') @mock.patch('apache_beam.internal.gcp.auth.get_service_credentials') def test_headers(self, mock_get_service_credentials, mock_do_request): from apache_beam.internal.gcp.auth import _ApitoolsCredentialsAdapter mock_get_service_credentials.return_value = _ApitoolsCredentialsAdapter( _make_credentials("test-project")) options = PipelineOptions([ "--job_name=test-job-name", "--gcs_custom_audit_entry=user=test-user-id", "--gcs_custom_audit_entries={\"id\": \"1234\", \"status\": \"ok\"}" ]) gcs = gcsio.GcsIO(pipeline_options=options) # no HTTP request when initializing GcsIO mock_do_request.assert_not_called() import requests response = requests.Response() response.status_code = 200 mock_do_request.return_value = response # The function of get_bucket() is supposed to send only one HTTP request gcs.get_bucket("test-bucket") mock_do_request.assert_called_once() call_args = mock_do_request.call_args[0] # Headers are specified as the third argument of # google.cloud._http.JSONConnection._do_request actual_headers = call_args[2] beam_user_agent = "apache-beam/%s (GPN:Beam)" % beam_version.__version__ self.assertIn(beam_user_agent, actual_headers['User-Agent']) self.assertEqual(actual_headers['x-goog-custom-audit-job'], 'test-job-name') self.assertEqual(actual_headers['x-goog-custom-audit-user'], 'test-user-id') self.assertEqual(actual_headers['x-goog-custom-audit-id'], '1234') self.assertEqual(actual_headers['x-goog-custom-audit-status'], 'ok') @mock.patch('google.cloud._http.JSONConnection._do_request') @mock.patch('apache_beam.internal.gcp.auth.get_service_credentials') def test_create_default_bucket( self, mock_get_service_credentials, mock_do_request): from apache_beam.internal.gcp.auth import _ApitoolsCredentialsAdapter mock_get_service_credentials.return_value = _ApitoolsCredentialsAdapter( _make_credentials("test-project")) gcs = gcsio.GcsIO(pipeline_options={"job_name": "test-job-name"}) # no HTTP request when initializing GcsIO mock_do_request.assert_not_called() import requests response = requests.Response() response.status_code = 200 mock_do_request.return_value = response # The function of create_bucket() is supposed to send only one HTTP request gcs.create_bucket("test-bucket", "test-project") mock_do_request.assert_called_once() call_args = mock_do_request.call_args[0] # Request data is specified as the fourth argument of # google.cloud._http.JSONConnection._do_request actual_request_data = call_args[3] import json request_data_json = json.loads(actual_request_data) # verify soft delete policy is disabled by default in the bucket creation # request self.assertEqual( request_data_json['softDeletePolicy']['retentionDurationSeconds'], 0) @mock.patch("apache_beam.io.gcp.gcsio.GcsIO.get_bucket") def test_is_soft_delete_enabled(self, mock_get_bucket): bucket = mock.MagicMock() mock_get_bucket.return_value = bucket # soft delete policy enabled bucket.soft_delete_policy.retention_duration_seconds = 1024 self.assertTrue( self.gcs.is_soft_delete_enabled("gs://beam_with_soft_delete/tmp")) # soft delete policy disabled bucket.soft_delete_policy.retention_duration_seconds = 0 self.assertFalse( self.gcs.is_soft_delete_enabled("gs://beam_without_soft_delete/tmp")) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3864086 apache_beam-2.74.0/apache_beam/io/gcp/healthcare/0000755000175100017510000000000015206113157021145 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/healthcare/__init__.py0000644000175100017510000000142015206113032023243 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/healthcare/dicomclient.py0000644000175100017510000000766115206113032024013 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file from google.auth import default from google.auth.transport import requests class DicomApiHttpClient: """Creates a client that communicates with the GCP Healthcare API (https://cloud.google.com/healthcare-api), over HTTP. Note that this client strictly uses v1 of this API.""" healthcare_base_url = "https://healthcare.googleapis.com/v1" session = None def get_session(self, credential): if self.session: return self.session # if the credential is not provided, use the default credential. if not credential: credential, _ = default() new_seesion = requests.AuthorizedSession(credential) self.session = new_seesion return new_seesion def qido_search( self, project_id, region, dataset_id, dicom_store_id, search_type, params=None, credential=None): """Function for searching a DICOM store""" # Sends a request to the GCP HCLS API v1. api_endpoint = "{}/projects/{}/locations/{}".format( self.healthcare_base_url, project_id, region) # base of dicomweb path. dicomweb_path = "{}/datasets/{}/dicomStores/{}/dicomWeb/{}".format( api_endpoint, dataset_id, dicom_store_id, search_type) # Make an authenticated API request session = self.get_session(credential) headers = {"Content-Type": "application/dicom+json; charset=utf-8"} page_size = 500 if params and 'limit' in params: page_size = params['limit'] elif params: params['limit'] = page_size else: params = {'limit': page_size} offset = 0 output = [] # iterate to get all the results while True: params['offset'] = offset response = session.get(dicomweb_path, headers=headers, params=params) status = response.status_code if status != 200: if offset == 0: return [], status params['offset'] = offset - 1 params['limit'] = 1 response = session.get(dicomweb_path, headers=headers, params=params) check_status = response.status_code if check_status == 200: # if the number of results equals to page size return output, check_status else: # something wrong with the request or server return [], status results = response.json() output += results if len(results) < page_size: # got all the results, return break offset += len(results) return output, status def dicomweb_store_instance( self, project_id, region, dataset_id, dicom_store_id, dcm_file, credential=None): """Function for storing a DICOM instance.""" api_endpoint = "{}/projects/{}/locations/{}".format( self.healthcare_base_url, project_id, region) dicomweb_path = "{}/datasets/{}/dicomStores/{}/dicomWeb/studies".format( api_endpoint, dataset_id, dicom_store_id) # Make an authenticated API request session = self.get_session(credential) content_type = "application/dicom" headers = {"Content-Type": content_type} response = session.post(dicomweb_path, data=dcm_file, headers=headers) return None, response.status_code ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/healthcare/dicomio.py0000644000175100017510000005044115206113032023136 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """DICOM IO connector This module implements several tools to facilitate the interaction between a Google Cloud Healthcare DICOM store and a Beam pipeline. For more details on DICOM store and API: https://cloud.google.com/healthcare/docs/how-tos/dicom The DICOM IO connector can be used to search metadata or write DICOM files to DICOM store. When used together with Google Pubsub message connector, the `FormatToQido` PTransform implemented in this module can be used to convert Pubsub messages to search requests. Since Traceability is crucial for healthcare API users, every input or error message will be recorded in the output of the DICOM IO connector. As a result, every PTransform in this module will return a PCollection of dict that encodes results and detailed error messages. Search instance's metadata (QIDO request) =================================================== DicomSearch() wraps the QIDO request client and supports 3 levels of search. Users should specify the level by setting the 'search_type' entry in the input dict. They can also refine the search by adding tags to filter the results using the 'params' entry. Here is a sample usage: with Pipeline() as p: input_dict = p | beam.Create( [{'project_id': 'abc123', 'type': 'instances',...}, {'project_id': 'dicom_go', 'type': 'series',...}]) results = input_dict | io.gcp.DicomSearch() results | 'print successful search' >> beam.Map( lambda x: print(x['result'] if x['success'] else None)) results | 'print failed search' >> beam.Map( lambda x: print(x['result'] if not x['success'] else None)) In the example above, successful qido search results and error messages for failed requests are printed. When used in real life, user can choose to filter those data and output them to wherever they want. Convert DICOM Pubsub message to Qido search request =================================================== Healthcare API users might read messages from Pubsub to monitor the store operations (e.g. new file) in a DICOM storage. Pubsub message encode DICOM as a web store path as well as instance ids. If users are interested in getting new instance's metadata, they can use the `FormatToQido` transform to convert the message into Qido Search dict then use the `DicomSearch` transform. Here is a sample usage: pipeline_options = PipelineOptions() pipeline_options.view_as(StandardOptions).streaming = True p = beam.Pipeline(options=pipeline_options) pubsub = p | beam.io.ReadStringFromPubsub(subscription='a_dicom_store') results = pubsub | FormatToQido() success = results | 'filter message' >> beam.Filter(lambda x: x['success']) qido_dict = success | 'get qido request' >> beam.Map(lambda x: x['result']) metadata = qido_dict | DicomSearch() In the example above, the pipeline is listening to a pubsub topic and waiting for messages from DICOM API. When a new DICOM file comes into the storage, the pipeline will receive a pubsub message, convert it to a Qido request dict, and feed it to DicomSearch() PTransform. As a result, users can get the metadata for every new DICOM file. Note that not every pubsub message received is from DICOM API, so we to filter the results first. Store a DICOM file in a DICOM storage =================================================== UploadToDicomStore() wraps store request API and users can use it to send a DICOM file to a DICOM store. It supports two types of input: 1. fileio object 2. file data in byte[]. Users should set the 'input_type' when initialzing this PTransform. Here are the examples: input_dict = {'project_id': 'abc123', 'type': 'instances',...} str_input = json.dumps(dict_input) temp_dir = '%s%s' % (self._new_tempdir(), os.sep) self._create_temp_file(dir=temp_dir, content=str_input) with Pipeline() as p: results = ( p | beam.Create([FileSystems.join(temp_dir, '*')]) | fileio.MatchAll() | fileio.ReadMatches() | UploadToDicomStore(input_dict, 'fileio')) input_dict = {'project_id': 'abc123', 'type': 'instances',...} str_input = json.dumps(dict_input) bytes_input = bytes(str_input.encode("utf-8")) with Pipeline() as p: results = ( p | beam.Create([bytes_input]) | UploadToDicomStore(input_dict, 'bytes')) The first example uses a PCollection of fileio objects as input. UploadToDicomStore will read DICOM files from the objects and send them to a DICOM storage. The second example uses a PCollection of byte[] as input. UploadToDicomStore will directly send those DICOM files to a DICOM storage. Users can also get the operation results in the output PCollection if they want to handle the failed store requests. """ # pytype: skip-file from concurrent.futures import ThreadPoolExecutor from concurrent.futures import as_completed import apache_beam as beam from apache_beam.io.gcp.healthcare.dicomclient import DicomApiHttpClient from apache_beam.transforms import PTransform class DicomSearch(PTransform): """A PTransform used for retrieving DICOM instance metadata from Google Cloud DICOM store. It takes a PCollection of dicts as input and return a PCollection of dict as results: INPUT: The input dict represents DICOM web path parameters, which has the following string keys and values: { 'project_id': str, 'region': str, 'dataset_id': str, 'dicom_store_id': str, 'search_type': str, 'params': dict(str,str) (Optional), } Key-value pairs: project_id: Id of the project in which the DICOM store is located. (Required) region: Region where the DICOM store resides. (Required) dataset_id: Id of the dataset where DICOM store belongs to. (Required) dicom_store_id: Id of the dicom store. (Required) search_type: Which type of search it is, could only be one of the three values: 'instances', 'series', or 'studies'. (Required) params: A dict of str:str pairs used to refine QIDO search. (Optional) Supported tags in three categories: 1.Studies: * StudyInstanceUID, * PatientName, * PatientID, * AccessionNumber, * ReferringPhysicianName, * StudyDate, 2.Series: all study level search terms and * SeriesInstanceUID, * Modality, 3.Instances: all study/series level search terms and * SOPInstanceUID, e.g. {"StudyInstanceUID":"1","SeriesInstanceUID":"2"} OUTPUT: The output dict wraps results as well as error messages: { 'result': a list of dicts in JSON style. 'success': boolean value telling whether the operation is successful. 'input': detail ids and dicomweb path for this retrieval. 'status': status code from the server, used as error message. } """ def __init__( self, buffer_size=8, max_workers=5, client=None, credential=None): """Initializes DicomSearch. Args: buffer_size: # type: Int. Size of the request buffer. max_workers: # type: Int. Maximum number of threads a worker can create. If it is set to one, all the request will be processed sequentially in a worker. client: # type: object. If it is specified, all the Api calls will made by this client instead of the default one (DicomApiHttpClient). credential: # type: Google credential object, if it is specified, the Http client will use it to create sessions instead of the default. """ self.buffer_size = buffer_size self.max_workers = max_workers self.client = client or DicomApiHttpClient() self.credential = credential def expand(self, pcoll): return pcoll | beam.ParDo( _QidoReadFn( self.buffer_size, self.max_workers, self.client, self.credential)) class _QidoReadFn(beam.DoFn): """A DoFn for executing every qido query request.""" def __init__(self, buffer_size, max_workers, client, credential=None): self.buffer_size = buffer_size self.max_workers = max_workers self.client = client self.credential = credential def start_bundle(self): self.buffer = [] def finish_bundle(self): for item in self._flush(): yield item def validate_element(self, element): # Check if all required keys present. required_keys = [ 'project_id', 'region', 'dataset_id', 'dicom_store_id', 'search_type' ] for key in required_keys: if key not in element: error_message = 'Must have %s in the dict.' % (key) return False, error_message # Check if return type is correct. if element['search_type'] in ['instances', "studies", "series"]: return True, None else: error_message = ( 'Search type can only be "studies", ' '"instances" or "series"') return False, error_message def process( self, element, window=beam.DoFn.WindowParam, timestamp=beam.DoFn.TimestampParam): # Check if the element is valid valid, error_message = self.validate_element(element) if valid: self.buffer.append((element, window, timestamp)) if len(self.buffer) >= self.buffer_size: for item in self._flush(): yield item else: # Return this when the input dict dose not meet the requirements out = {} out['result'] = [] out['status'] = error_message out['input'] = element out['success'] = False yield out def make_request(self, element): # Sending Qido request to DICOM Api project_id = element['project_id'] region = element['region'] dataset_id = element['dataset_id'] dicom_store_id = element['dicom_store_id'] search_type = element['search_type'] params = element['params'] if 'params' in element else None # Call qido search http client result, status_code = self.client.qido_search( project_id, region, dataset_id, dicom_store_id, search_type, params, self.credential ) out = {} out['result'] = result out['status'] = status_code out['input'] = element out['success'] = (status_code == 200) return out def process_buffer_element(self, buffer_element): # Thread job runner - each thread makes a Qido search request value = self.make_request(buffer_element[0]) windows = [buffer_element[1]] timestamp = buffer_element[2] return beam.utils.windowed_value.WindowedValue( value=value, timestamp=timestamp, windows=windows) def _flush(self): # Create thread pool executor and process the buffered elements in paralllel executor = ThreadPoolExecutor(max_workers=self.max_workers) futures = [ executor.submit(self.process_buffer_element, ele) for ele in self.buffer ] self.buffer = [] for f in as_completed(futures): yield f.result() class FormatToQido(PTransform): """A PTransform for converting pubsub messages into search input dict. Takes PCollection of string as input and returns a PCollection of dict as results. Note that some pubsub messages may not be from DICOM API, which will be recorded as failed conversions. INPUT: The input are normally strings from Pubsub topic: "projects/PROJECT_ID/locations/LOCATION/datasets/DATASET_ID/ dicomStores/DICOM_STORE_ID/dicomWeb/studies/STUDY_UID/ series/SERIES_UID/instances/INSTANCE_UID" OUTPUT: The output dict encodes results as well as error messages: { 'result': a dict representing instance level qido search request. 'success': boolean value telling whether the conversion is successful. 'input': input pubsub message string. } """ def __init__(self, credential=None): """Initializes FormatToQido. Args: credential: # type: Google credential object, if it is specified, the Http client will use it instead of the default one. """ self.credential = credential def expand(self, pcoll): return pcoll | beam.ParDo(_ConvertStringToQido()) class _ConvertStringToQido(beam.DoFn): """A DoFn for converting pubsub string to qido search parameters.""" def process(self, element): # Some constants for DICOM pubsub message NUM_PUBSUB_STR_ENTRIES = 15 NUM_DICOM_WEBPATH_PARAMETERS = 5 NUM_TOTAL_PARAMETERS = 8 INDEX_PROJECT_ID = 1 INDEX_REGION = 3 INDEX_DATASET_ID = 5 INDEX_DICOMSTORE_ID = 7 INDEX_STUDY_ID = 10 INDEX_SERIE_ID = 12 INDEX_INSTANCE_ID = 14 entries = element.split('/') # Output dict with error message, used when # receiving invalid pubsub string. error_dict = {} error_dict['result'] = {} error_dict['input'] = element error_dict['success'] = False if len(entries) != NUM_PUBSUB_STR_ENTRIES: return [error_dict] required_keys = [ 'projects', 'locations', 'datasets', 'dicomStores', 'dicomWeb', 'studies', 'series', 'instances' ] # Check if the required keys present and # the positions of those keys are correct for i in range(NUM_DICOM_WEBPATH_PARAMETERS): if required_keys[i] != entries[i * 2]: return [error_dict] for i in range(NUM_DICOM_WEBPATH_PARAMETERS, NUM_TOTAL_PARAMETERS): if required_keys[i] != entries[i * 2 - 1]: return [error_dict] # Compose dicom webpath parameters for qido search qido_dict = {} qido_dict['project_id'] = entries[INDEX_PROJECT_ID] qido_dict['region'] = entries[INDEX_REGION] qido_dict['dataset_id'] = entries[INDEX_DATASET_ID] qido_dict['dicom_store_id'] = entries[INDEX_DICOMSTORE_ID] qido_dict['search_type'] = 'instances' # Compose instance level params for qido search params = {} params['StudyInstanceUID'] = entries[INDEX_STUDY_ID] params['SeriesInstanceUID'] = entries[INDEX_SERIE_ID] params['SOPInstanceUID'] = entries[INDEX_INSTANCE_ID] qido_dict['params'] = params out = {} out['result'] = qido_dict out['input'] = element out['success'] = True return [out] class UploadToDicomStore(PTransform): """A PTransform for storing instances to a DICOM store. Takes PCollection of byte[] as input and return a PCollection of dict as results. The inputs are normally DICOM file in bytes or str filename. INPUT: This PTransform supports two types of input: 1. Byte[]: representing dicom file. 2. Fileio object: stream file object. OUTPUT: The output dict encodes status as well as error messages: { 'success': boolean value telling whether the store is successful. 'input': undeliverable data. Exactly the same as the input, only set if the operation is failed. 'status': status code from the server, used as error messages. } """ def __init__( self, destination_dict, input_type, buffer_size=8, max_workers=5, client=None, credential=None): """Initializes UploadToDicomStore. Args: destination_dict: # type: python dict, encodes DICOM endpoint information: { 'project_id': str, 'region': str, 'dataset_id': str, 'dicom_store_id': str, } Key-value pairs: * project_id: Id of the project in which DICOM store locates. (Required) * region: Region where the DICOM store resides. (Required) * dataset_id: Id of the dataset where DICOM store belongs to. (Required) * dicom_store_id: Id of the dicom store. (Required) input_type: # type: string, could only be 'bytes' or 'fileio' buffer_size: # type: Int. Size of the request buffer. max_workers: # type: Int. Maximum number of threads a worker can create. If it is set to one, all the request will be processed sequentially in a worker. client: # type: object. If it is specified, all the Api calls will made by this client instead of the default one (DicomApiHttpClient). credential: # type: Google credential object, if it is specified, the Http client will use it instead of the default one. """ self.destination_dict = destination_dict # input_type pre-check if input_type not in ['bytes', 'fileio']: raise ValueError("input_type could only be 'bytes' or 'fileio'") self.input_type = input_type self.buffer_size = buffer_size self.max_workers = max_workers self.client = client self.credential = credential def expand(self, pcoll): return pcoll | beam.ParDo( _StoreInstance( self.destination_dict, self.input_type, self.buffer_size, self.max_workers, self.client, self.credential)) class _StoreInstance(beam.DoFn): """A DoFn read or fetch dicom files then push it to a dicom store.""" def __init__( self, destination_dict, input_type, buffer_size, max_workers, client, credential=None): # pre-check destination dict required_keys = ['project_id', 'region', 'dataset_id', 'dicom_store_id'] for key in required_keys: if key not in destination_dict: raise ValueError('Must have %s in the dict.' % (key)) self.destination_dict = destination_dict self.input_type = input_type self.buffer_size = buffer_size self.max_workers = max_workers self.client = client self.credential = credential def start_bundle(self): self.buffer = [] def finish_bundle(self): for item in self._flush(): yield item def process( self, element, window=beam.DoFn.WindowParam, timestamp=beam.DoFn.TimestampParam): self.buffer.append((element, window, timestamp)) if len(self.buffer) >= self.buffer_size: for item in self._flush(): yield item def make_request(self, dicom_file): # Send file to DICOM store and records the results. project_id = self.destination_dict['project_id'] region = self.destination_dict['region'] dataset_id = self.destination_dict['dataset_id'] dicom_store_id = self.destination_dict['dicom_store_id'] # Feed the dicom file into store client if self.client: _, status_code = self.client.dicomweb_store_instance( project_id, region, dataset_id, dicom_store_id, dicom_file, self.credential ) else: _, status_code = DicomApiHttpClient().dicomweb_store_instance( project_id, region, dataset_id, dicom_store_id, dicom_file, self.credential ) out = {} out['status'] = status_code out['success'] = (status_code == 200) return out def read_dicom_file(self, buffer_element): # Read the file based on different input. If the read fails ,return # an error dict which records input and error messages. try: if self.input_type == 'fileio': f = buffer_element.open() data = f.read() f.close() return True, data else: return True, buffer_element except Exception as error_message: error_out = {} error_out['status'] = error_message error_out['success'] = False return False, error_out def process_buffer_element(self, buffer_element): # Thread job runner - each thread stores a DICOM file success, read_result = self.read_dicom_file(buffer_element[0]) windows = [buffer_element[1]] timestamp = buffer_element[2] value = None if success: value = self.make_request(read_result) else: value = read_result # save the undeliverable data if not value['success']: value['input'] = buffer_element[0] return beam.utils.windowed_value.WindowedValue( value=value, timestamp=timestamp, windows=windows) def _flush(self): # Create thread pool executor and process the buffered elements in paralllel executor = ThreadPoolExecutor(max_workers=self.max_workers) futures = [ executor.submit(self.process_buffer_element, ele) for ele in self.buffer ] self.buffer = [] for f in as_completed(futures): yield f.result() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/healthcare/dicomio_integration_test.py0000644000175100017510000002200415206113032026572 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Integration test for Google Cloud DICOM IO connector. This e2e test will first create a temporary empty DICOM storage and send 18 DICOM files from gs://apache-beam-samples/healthcare/dicom/io_test_files to it. The test will compare the metadata of a persistent DICOM storage, which reprensets ground truths and has 18 files stored, to the temporary storage in order to check if the connectors are functioning correctly. """ # pytype: skip-file import datetime import random import string import unittest import pytest import apache_beam as beam from apache_beam.io import fileio from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from google.auth import default from google.auth.transport import requests from apache_beam.io.gcp.healthcare.dicomclient import DicomApiHttpClient from apache_beam.io.gcp.healthcare.dicomio import DicomSearch from apache_beam.io.gcp.healthcare.dicomio import UploadToDicomStore except ImportError: DicomSearch = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports REGION = 'us-central1' DATA_SET_ID = 'apache-beam-integration-testing' HEALTHCARE_BASE_URL = 'https://healthcare.googleapis.com/v1' GCS_BASE_URL = 'https://storage.googleapis.com/storage/v1' PERSISTENT_DICOM_STORE_NAME = "dicom_it_persistent_store" BUCKET_NAME = 'apache-beam-samples' DICOM_DIR_PATH = 'healthcare/dicom' DICOM_FILES_PATH = 'gs://' + BUCKET_NAME + '/' + DICOM_DIR_PATH METADATA_DIR_PATH = DICOM_DIR_PATH + '/io_test_metadata/' META_DATA_ALL_NAME = 'Dicom_io_it_test_data.json' META_DATA_REFINED_NAME = 'Dicom_io_it_test_refined_data.json' NUM_INSTANCE = 18 RAND_LEN = 15 SCOPES = ['https://www.googleapis.com/auth/cloud-platform'] # Tag 00081190 contains temp store name which contains currentDate VOLATILE_TAGS = {"00081190"} def normalize_outer(elem: dict) -> dict: elem = dict(elem) # shallow copy elem["result"] = [normalize_instance(d) for d in elem.get("result", [])] return elem def normalize_instance(instance: dict) -> dict: return {k: v for k, v in instance.items() if k not in VOLATILE_TAGS} def random_string_generator(length): letters_and_digits = string.ascii_letters + string.digits result = ''.join((random.choice(letters_and_digits) for i in range(length))) return result def create_dicom_store(project_id, dataset_id, region, dicom_store_id): # Create a an empty DICOM store credential, _ = default(SCOPES) session = requests.AuthorizedSession(credential) api_endpoint = "{}/projects/{}/locations/{}".format( HEALTHCARE_BASE_URL, project_id, region) # base of dicomweb path. dicomweb_path = "{}/datasets/{}/dicomStores".format(api_endpoint, dataset_id) response = session.post( dicomweb_path, params={"dicomStoreId": dicom_store_id}) response.raise_for_status() return response.status_code def delete_dicom_store(project_id, dataset_id, region, dicom_store_id): # Delete an existing DICOM store credential, _ = default(SCOPES) session = requests.AuthorizedSession(credential) api_endpoint = "{}/projects/{}/locations/{}".format( HEALTHCARE_BASE_URL, project_id, region) # base of dicomweb path. dicomweb_path = "{}/datasets/{}/dicomStores/{}".format( api_endpoint, dataset_id, dicom_store_id) response = session.delete(dicomweb_path) response.raise_for_status() return response.status_code def get_gcs_file_http(file_name): # Get gcs file from REST Api file_name = file_name.replace('/', '%2F') api_endpoint = "{}/b/{}/o/{}?alt=media".format( GCS_BASE_URL, BUCKET_NAME, file_name) credential, _ = default(SCOPES) session = requests.AuthorizedSession(credential) response = session.get(api_endpoint) response.raise_for_status() return response.json() @unittest.skipIf(DicomSearch is None, 'GCP dependencies are not installed') class DICOMIoIntegrationTest(unittest.TestCase): def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.project = self.test_pipeline.get_option('project') self.expected_output_all_metadata = get_gcs_file_http( METADATA_DIR_PATH + META_DATA_ALL_NAME) self.expected_output_refined_metadata = get_gcs_file_http( METADATA_DIR_PATH + META_DATA_REFINED_NAME) # create a temp Dicom store based on the time stamp self.temp_dicom_store = "DICOM_store_" + datetime.datetime.now().strftime( '%Y-%m-%d_%H%M%S.%f_') + random_string_generator(RAND_LEN) create_dicom_store(self.project, DATA_SET_ID, REGION, self.temp_dicom_store) def tearDown(self): # clean up the temp Dicom store delete_dicom_store(self.project, DATA_SET_ID, REGION, self.temp_dicom_store) @pytest.mark.it_postcommit def test_dicom_search_instances(self): # Search and compare the metadata of a persistent DICOM store. # Both refine and comprehensive search will be tested. input_dict_all = {} input_dict_all['project_id'] = self.project input_dict_all['region'] = REGION input_dict_all['dataset_id'] = DATA_SET_ID input_dict_all['dicom_store_id'] = PERSISTENT_DICOM_STORE_NAME input_dict_all['search_type'] = "instances" input_dict_refine = {} input_dict_refine['project_id'] = self.project input_dict_refine['region'] = REGION input_dict_refine['dataset_id'] = DATA_SET_ID input_dict_refine['dicom_store_id'] = PERSISTENT_DICOM_STORE_NAME input_dict_refine['search_type'] = "instances" input_dict_refine['params'] = { 'StudyInstanceUID': 'study_000000001', 'limit': 500, 'offset': 0 } expected_dict_all = {} expected_dict_all['result'] = self.expected_output_all_metadata expected_dict_all['status'] = 200 expected_dict_all['input'] = input_dict_all expected_dict_all['success'] = True expected_dict_refine = {} expected_dict_refine['result'] = self.expected_output_refined_metadata expected_dict_refine['status'] = 200 expected_dict_refine['input'] = input_dict_refine expected_dict_refine['success'] = True with self.test_pipeline as p: results_all = ( p | 'create all dict' >> beam.Create([input_dict_all]) | 'search all' >> DicomSearch() | 'normalize all' >> beam.Map(normalize_outer)) results_refine = ( p | 'create refine dict' >> beam.Create([input_dict_refine]) | 'search refine' >> DicomSearch() | 'normalize refine' >> beam.Map(normalize_outer)) expected_all_norm = normalize_outer(expected_dict_all) expected_refine_norm = normalize_outer(expected_dict_refine) assert_that( results_all, equal_to([expected_all_norm]), label='all search assert') assert_that( results_refine, equal_to([expected_refine_norm]), label='refine search assert') @pytest.mark.it_postcommit def test_dicom_store_instance_from_gcs(self): # Store DICOM files to a empty DICOM store from a GCS bucket, # then check if the store metadata match. input_dict_store = {} input_dict_store['project_id'] = self.project input_dict_store['region'] = REGION input_dict_store['dataset_id'] = DATA_SET_ID input_dict_store['dicom_store_id'] = self.temp_dicom_store expected_output = [True] * NUM_INSTANCE with self.test_pipeline as p: gcs_path = DICOM_FILES_PATH + "/io_test_files/*" results = ( p | fileio.MatchFiles(gcs_path) | fileio.ReadMatches() | UploadToDicomStore(input_dict_store, 'fileio') | beam.Map(lambda x: x['success'])) assert_that( results, equal_to(expected_output), label='store first assert') # Check the metadata using client credential, _ = default(SCOPES) result, status_code = DicomApiHttpClient().qido_search( self.project, REGION, DATA_SET_ID, self.temp_dicom_store, 'instances', credential=credential ) self.assertEqual(status_code, 200) actual_norm = [normalize_instance(r) for r in result] expected_norm = [ normalize_instance(r) for r in self.expected_output_all_metadata ] # Order-insensitive deep equality self.assertCountEqual(actual_norm, expected_norm) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/healthcare/dicomio_test.py0000644000175100017510000004372715206113032024206 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for Dicom IO connectors.""" # pytype: skip-file import json import os import unittest from mock import patch import apache_beam as beam from apache_beam.io import fileio from apache_beam.io.filebasedsink_test import _TestCaseWithTempDirCleanUp from apache_beam.io.filesystems import FileSystems from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # Protect against environments where gcp library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.io.gcp.healthcare.dicomio import DicomSearch from apache_beam.io.gcp.healthcare.dicomio import FormatToQido from apache_beam.io.gcp.healthcare.dicomio import UploadToDicomStore except ImportError: DicomSearch = None # type: ignore # pylint: enable=wrong-import-order, wrong-import-position class MockHttpClient(): # A mock HTTP client that talks directly to an in-memory Dicom store. def __init__(self): # set 5 mock dicom instances dicom_metadata = [] dicom_metadata.append({ 'PatientName': 'Albert', 'Age': 21, 'TestResult': 'Positive' }) dicom_metadata.append({ 'PatientName': 'Albert', 'Age': 21, 'TestResult': 'Negative' }) dicom_metadata.append({ 'PatientName': 'Brian', 'Age': 20, 'TestResult': 'Positive' }) dicom_metadata.append({ 'PatientName': 'Colin', 'Age': 25, 'TestResult': 'Negative' }) dicom_metadata.append({ 'PatientName': 'Daniel', 'Age': 22, 'TestResult': 'Negative' }) dicom_metadata.append({ 'PatientName': 'Eric', 'Age': 50, 'TestResult': 'Negative' }) self.dicom_metadata = dicom_metadata # ids for this dicom store self.project_id = 'test_project' self.region = 'test_region' self.dataset_id = 'test_dataset_id' self.dicom_store_id = 'test_dicom_store_id' def qido_search( self, project_id, region, dataset_id, dicom_store_id, search_type, params=None, credential=None): # qido search function for this mock client if project_id != self.project_id or region != self.region or \ dataset_id != self.dataset_id or dicom_store_id != self.dicom_store_id: return [], 204 # only supports all instance search in this client if not params: return self.dicom_metadata, 200 # only supports patient name filter in this client patient_name = params['PatientName'] out = [] for meta in self.dicom_metadata: if meta['PatientName'] == patient_name: out.append(meta) return out, 200 def dicomweb_store_instance( self, project_id, region, dataset_id, dicom_store_id, dcm_file, credential=None): if project_id != self.project_id or region != self.region or \ dataset_id != self.dataset_id or dicom_store_id != self.dicom_store_id: return [], 204 # convert the bytes file back to dict string_array = dcm_file.decode('utf-8') metadata_dict = json.loads(string_array) self.dicom_metadata.append(metadata_dict) return None, 200 @unittest.skipIf(DicomSearch is None, 'GCP dependencies are not installed') class TestFormatToQido(unittest.TestCase): valid_pubsub_string = ( "projects/PROJECT_ID/locations/LOCATION/datasets" "/DATASET_ID/dicomStores/DICOM_STORE_ID/dicomWeb/" "studies/STUDY_UID/series/SERIES_UID/instances/INSTANCE_UID") expected_valid_pubsub_dict = { 'result': { 'project_id': 'PROJECT_ID', 'region': 'LOCATION', 'dataset_id': 'DATASET_ID', 'dicom_store_id': 'DICOM_STORE_ID', 'search_type': 'instances', 'params': { 'StudyInstanceUID': 'STUDY_UID', 'SeriesInstanceUID': 'SERIES_UID', 'SOPInstanceUID': 'INSTANCE_UID' } }, 'input': valid_pubsub_string, 'success': True } invalid_pubsub_string = "this is not a valid pubsub message" expected_invalid_pubsub_dict = { 'result': {}, 'input': 'this is not a valid pubsub message', 'success': False } def test_normal_convert(self): with TestPipeline() as p: convert_result = ( p | beam.Create([self.valid_pubsub_string]) | FormatToQido()) assert_that(convert_result, equal_to([self.expected_valid_pubsub_dict])) def test_failed_convert(self): with TestPipeline() as p: convert_result = ( p | beam.Create([self.invalid_pubsub_string]) | FormatToQido()) assert_that(convert_result, equal_to([self.expected_invalid_pubsub_dict])) @unittest.skipIf(DicomSearch is None, 'GCP dependencies are not installed') class TestDicomSearch(unittest.TestCase): @patch("apache_beam.io.gcp.healthcare.dicomio.DicomApiHttpClient") def test_successful_search(self, MockClient): input_dict = {} input_dict['project_id'] = "test_project" input_dict['region'] = "test_region" input_dict['dataset_id'] = "test_dataset_id" input_dict['dicom_store_id'] = "test_dicom_store_id" input_dict['search_type'] = "instances" mc = MockHttpClient() MockClient.return_value = mc expected_dict = {} expected_dict['result'] = mc.dicom_metadata expected_dict['status'] = 200 expected_dict['input'] = input_dict expected_dict['success'] = True with TestPipeline() as p: results = (p | beam.Create([input_dict]) | DicomSearch()) assert_that(results, equal_to([expected_dict])) @patch("apache_beam.io.gcp.healthcare.dicomio.DicomApiHttpClient") def test_Qido_search_small_buffer_flush(self, MockClient): input_dict = {} input_dict['project_id'] = "test_project" input_dict['region'] = "test_region" input_dict['dataset_id'] = "test_dataset_id" input_dict['dicom_store_id'] = "test_dicom_store_id" input_dict['search_type'] = "instances" mc = MockHttpClient() MockClient.return_value = mc expected_dict = {} expected_dict['result'] = mc.dicom_metadata expected_dict['status'] = 200 expected_dict['input'] = input_dict expected_dict['success'] = True with TestPipeline() as p: results = (p | beam.Create([input_dict] * 5) | DicomSearch(buffer_size=1)) assert_that(results, equal_to([expected_dict] * 5)) @patch("apache_beam.io.gcp.healthcare.dicomio.DicomApiHttpClient") def test_param_dict_passing(self, MockClient): input_dict = {} input_dict['project_id'] = "test_project" input_dict['region'] = "test_region" input_dict['dataset_id'] = "test_dataset_id" input_dict['dicom_store_id'] = "test_dicom_store_id" input_dict['search_type'] = "instances" input_dict['params'] = {'PatientName': 'Brian'} expected_dict = {} expected_dict['result'] = [{ 'PatientName': 'Brian', 'Age': 20, 'TestResult': 'Positive' }] expected_dict['status'] = 200 expected_dict['input'] = input_dict expected_dict['success'] = True mc = MockHttpClient() MockClient.return_value = mc with TestPipeline() as p: results = (p | beam.Create([input_dict]) | DicomSearch()) assert_that(results, equal_to([expected_dict])) @patch("apache_beam.io.gcp.healthcare.dicomio.DicomApiHttpClient") def test_wrong_input_type(self, MockClient): input_dict = {} input_dict['project_id'] = "test_project" input_dict['region'] = "test_region" input_dict['dataset_id'] = "test_dataset_id" input_dict['dicom_store_id'] = "test_dicom_store_id" input_dict['search_type'] = "not exist type" expected_invalid_dict = {} expected_invalid_dict['result'] = [] expected_invalid_dict[ 'status'] = 'Search type can only be "studies", "instances" or "series"' expected_invalid_dict['input'] = input_dict expected_invalid_dict['success'] = False mc = MockHttpClient() MockClient.return_value = mc with TestPipeline() as p: results = (p | beam.Create([input_dict]) | DicomSearch()) assert_that(results, equal_to([expected_invalid_dict])) @patch("apache_beam.io.gcp.healthcare.dicomio.DicomApiHttpClient") def test_missing_project_id(self, MockClient): input_dict = {} input_dict['dataset_id'] = "test_dataset" input_dict['region'] = "test_region" expected_invalid_dict = {} expected_invalid_dict['result'] = [] expected_invalid_dict['status'] = 'Must have project_id in the dict.' expected_invalid_dict['input'] = input_dict expected_invalid_dict['success'] = False mc = MockHttpClient() MockClient.return_value = mc with TestPipeline() as p: results = (p | beam.Create([input_dict]) | DicomSearch()) assert_that(results, equal_to([expected_invalid_dict])) @patch("apache_beam.io.gcp.healthcare.dicomio.DicomApiHttpClient") def test_missing_dataset_id(self, MockClient): input_dict = {} input_dict['project_id'] = "test_project" input_dict['region'] = "test_region" expected_invalid_dict = {} expected_invalid_dict['result'] = [] expected_invalid_dict['status'] = 'Must have dataset_id in the dict.' expected_invalid_dict['input'] = input_dict expected_invalid_dict['success'] = False mc = MockHttpClient() MockClient.return_value = mc with TestPipeline() as p: results = (p | beam.Create([input_dict]) | DicomSearch()) assert_that(results, equal_to([expected_invalid_dict])) @patch("apache_beam.io.gcp.healthcare.dicomio.DicomApiHttpClient") def test_missing_region(self, MockClient): input_dict = {} input_dict['project_id'] = "test_project" input_dict['dataset_id'] = "test_dataset" expected_invalid_dict = {} expected_invalid_dict['result'] = [] expected_invalid_dict['status'] = 'Must have region in the dict.' expected_invalid_dict['input'] = input_dict expected_invalid_dict['success'] = False mc = MockHttpClient() MockClient.return_value = mc with TestPipeline() as p: results = (p | beam.Create([input_dict]) | DicomSearch()) assert_that(results, equal_to([expected_invalid_dict])) @patch("apache_beam.io.gcp.healthcare.dicomio.DicomApiHttpClient") def test_client_search_notfound(self, MockClient): input_dict = {} # search instances in a not exist store input_dict['project_id'] = "wrong_project" input_dict['region'] = "wrong_region" input_dict['dataset_id'] = "wrong_dataset_id" input_dict['dicom_store_id'] = "wrong_dicom_store_id" input_dict['search_type'] = "instances" expected_invalid_dict = {} expected_invalid_dict['result'] = [] expected_invalid_dict['status'] = 204 expected_invalid_dict['input'] = input_dict expected_invalid_dict['success'] = False mc = MockHttpClient() MockClient.return_value = mc with TestPipeline() as p: results = (p | beam.Create([input_dict]) | DicomSearch()) assert_that(results, equal_to([expected_invalid_dict])) @unittest.skipIf(DicomSearch is None, 'GCP dependencies are not installed') class TestDicomStoreInstance(_TestCaseWithTempDirCleanUp): @patch("apache_beam.io.gcp.healthcare.dicomio.DicomApiHttpClient") def test_store_byte_file(self, MockClient): input_dict = {} input_dict['project_id'] = "test_project" input_dict['region'] = "test_region" input_dict['dataset_id'] = "test_dataset_id" input_dict['dicom_store_id'] = "test_dicom_store_id" mc = MockHttpClient() MockClient.return_value = mc dict_input = {'PatientName': 'George', 'Age': 23, 'TestResult': 'Negative'} str_input = json.dumps(dict_input) bytes_input = bytes(str_input.encode("utf-8")) with TestPipeline() as p: results = ( p | beam.Create([bytes_input]) | UploadToDicomStore(input_dict, 'bytes') | beam.Map(lambda x: x['success'])) assert_that(results, equal_to([True])) self.assertTrue(dict_input in mc.dicom_metadata) @patch("apache_beam.io.gcp.healthcare.dicomio.DicomApiHttpClient") def test_store_byte_file_small_buffer_flush(self, MockClient): input_dict = {} input_dict['project_id'] = "test_project" input_dict['region'] = "test_region" input_dict['dataset_id'] = "test_dataset_id" input_dict['dicom_store_id'] = "test_dicom_store_id" mc = MockHttpClient() MockClient.return_value = mc dict_input_1 = { 'PatientName': 'George', 'Age': 23, 'TestResult': 'Negative' } str_input_1 = json.dumps(dict_input_1) bytes_input_1 = bytes(str_input_1.encode("utf-8")) dict_input_2 = {'PatientName': 'Peter', 'Age': 54, 'TestResult': 'Positive'} str_input_2 = json.dumps(dict_input_2) bytes_input_2 = bytes(str_input_2.encode("utf-8")) dict_input_3 = {'PatientName': 'Zen', 'Age': 27, 'TestResult': 'Negative'} str_input_3 = json.dumps(dict_input_3) bytes_input_3 = bytes(str_input_3.encode("utf-8")) with TestPipeline() as p: results = ( p | beam.Create([bytes_input_1, bytes_input_2, bytes_input_3]) | UploadToDicomStore(input_dict, 'bytes', buffer_size=1) | beam.Map(lambda x: x['success'])) assert_that(results, equal_to([True] * 3)) self.assertTrue(dict_input_1 in mc.dicom_metadata) self.assertTrue(dict_input_2 in mc.dicom_metadata) self.assertTrue(dict_input_3 in mc.dicom_metadata) @patch("apache_beam.io.gcp.healthcare.dicomio.DicomApiHttpClient") def test_store_fileio_file(self, MockClient): input_dict = {} input_dict['project_id'] = "test_project" input_dict['region'] = "test_region" input_dict['dataset_id'] = "test_dataset_id" input_dict['dicom_store_id'] = "test_dicom_store_id" mc = MockHttpClient() MockClient.return_value = mc dict_input = {'PatientName': 'George', 'Age': 23, 'TestResult': 'Negative'} str_input = json.dumps(dict_input) temp_dir = '%s%s' % (self._new_tempdir(), os.sep) self._create_temp_file(dir=temp_dir, content=str_input) with TestPipeline() as p: results = ( p | beam.Create([FileSystems.join(temp_dir, '*')]) | fileio.MatchAll() | fileio.ReadMatches() | UploadToDicomStore(input_dict, 'fileio') | beam.Map(lambda x: x['success'])) assert_that(results, equal_to([True])) self.assertTrue(dict_input in mc.dicom_metadata) @patch("apache_beam.io.gcp.healthcare.dicomio.DicomApiHttpClient") def test_store_fileio_file_small_buffer_flush(self, MockClient): input_dict = {} input_dict['project_id'] = "test_project" input_dict['region'] = "test_region" input_dict['dataset_id'] = "test_dataset_id" input_dict['dicom_store_id'] = "test_dicom_store_id" mc = MockHttpClient() MockClient.return_value = mc temp_dir = '%s%s' % (self._new_tempdir(), os.sep) dict_input_1 = { 'PatientName': 'George', 'Age': 23, 'TestResult': 'Negative' } str_input_1 = json.dumps(dict_input_1) self._create_temp_file(dir=temp_dir, content=str_input_1) dict_input_2 = {'PatientName': 'Peter', 'Age': 54, 'TestResult': 'Positive'} str_input_2 = json.dumps(dict_input_2) self._create_temp_file(dir=temp_dir, content=str_input_2) dict_input_3 = {'PatientName': 'Zen', 'Age': 27, 'TestResult': 'Negative'} str_input_3 = json.dumps(dict_input_3) self._create_temp_file(dir=temp_dir, content=str_input_3) with TestPipeline() as p: results = ( p | beam.Create([FileSystems.join(temp_dir, '*')]) | fileio.MatchAll() | fileio.ReadMatches() | UploadToDicomStore(input_dict, 'fileio', buffer_size=1) | beam.Map(lambda x: x['success'])) assert_that(results, equal_to([True] * 3)) self.assertTrue(dict_input_1 in mc.dicom_metadata) self.assertTrue(dict_input_2 in mc.dicom_metadata) self.assertTrue(dict_input_3 in mc.dicom_metadata) @patch("apache_beam.io.gcp.healthcare.dicomio.DicomApiHttpClient") def test_destination_notfound(self, MockClient): input_dict = {} # search instances in a not exist store input_dict['project_id'] = "wrong_project" input_dict['region'] = "wrong_region" input_dict['dataset_id'] = "wrong_dataset_id" input_dict['dicom_store_id'] = "wrong_dicom_store_id" expected_invalid_dict = {} expected_invalid_dict['status'] = 204 expected_invalid_dict['input'] = '' expected_invalid_dict['success'] = False mc = MockHttpClient() MockClient.return_value = mc with TestPipeline() as p: results = ( p | beam.Create(['']) | UploadToDicomStore(input_dict, 'bytes')) assert_that(results, equal_to([expected_invalid_dict])) @patch("apache_beam.io.gcp.healthcare.dicomio.DicomApiHttpClient") def test_missing_parameters(self, MockClient): input_dict = {} input_dict['project_id'] = "test_project" input_dict['region'] = "test_region" expected_invalid_dict = {} expected_invalid_dict['result'] = [] expected_invalid_dict['status'] = 'Must have dataset_id in the dict.' expected_invalid_dict['input'] = input_dict expected_invalid_dict['success'] = False mc = MockHttpClient() MockClient.return_value = mc with self.assertRaisesRegex(ValueError, "Must have dataset_id in the dict."): p = TestPipeline() _ = (p | beam.Create(['']) | UploadToDicomStore(input_dict, 'bytes')) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3874085 apache_beam-2.74.0/apache_beam/io/gcp/internal/0000755000175100017510000000000015206113157020661 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/internal/__init__.py0000644000175100017510000000142015206113032022757 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3874085 apache_beam-2.74.0/apache_beam/io/gcp/internal/clients/0000755000175100017510000000000015206113157022322 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/internal/clients/__init__.py0000644000175100017510000000142015206113032024420 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3884087 apache_beam-2.74.0/apache_beam/io/gcp/internal/clients/bigquery/0000755000175100017510000000000015206113157024151 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/internal/clients/bigquery/__init__.py0000644000175100017510000000254615206113032026261 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Common imports for generated bigquery client library.""" # pylint:disable=wildcard-import import pkgutil # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apitools.base.py import * from apache_beam.io.gcp.internal.clients.bigquery.bigquery_v2_client import * from apache_beam.io.gcp.internal.clients.bigquery.bigquery_v2_messages import * except ImportError: pass # pylint: enable=wrong-import-order, wrong-import-position __path__ = pkgutil.extend_path(__path__, __name__) # type: ignore ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_client.py0000644000175100017510000012505415206113032030316 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Generated client library for bigquery version v2.""" # NOTE: This file is autogenerated and should not be edited by hand. from apitools.base.py import base_api from apache_beam.io.gcp.internal.clients.bigquery import \ bigquery_v2_messages as messages class BigqueryV2(base_api.BaseApiClient): """Generated client library for service bigquery version v2.""" MESSAGES_MODULE = messages BASE_URL = 'https://bigquery.googleapis.com/bigquery/v2/' MTLS_BASE_URL = 'https://www.mtls.googleapis.com/bigquery/v2/' _PACKAGE = 'bigquery' _SCOPES = [ 'https://www.googleapis.com/auth/bigquery', 'https://www.googleapis.com/auth/bigquery.insertdata', 'https://www.googleapis.com/auth/bigquery.readonly', 'https://www.googleapis.com/auth/cloud-platform', 'https://www.googleapis.com/auth/cloud-platform.read-only', 'https://www.googleapis.com/auth/devstorage.full_control', 'https://www.googleapis.com/auth/devstorage.read_only', 'https://www.googleapis.com/auth/devstorage.read_write' ] _VERSION = 'v2' _CLIENT_ID = '1042881264118.apps.googleusercontent.com' _CLIENT_SECRET = 'x_Tw5K8nnjoRAqULM9PFAC2b' _USER_AGENT = 'x_Tw5K8nnjoRAqULM9PFAC2b' _CLIENT_CLASS_NAME = 'BigqueryV2' _URL_VERSION = 'v2' _API_KEY = None def __init__( self, url='', credentials=None, get_credentials=True, http=None, model=None, log_request=False, log_response=False, credentials_args=None, default_global_params=None, additional_http_headers=None, response_encoding=None): """Create a new bigquery handle.""" url = url or self.BASE_URL super().__init__( url, credentials=credentials, get_credentials=get_credentials, http=http, model=model, log_request=log_request, log_response=log_response, credentials_args=credentials_args, default_global_params=default_global_params, additional_http_headers=additional_http_headers, response_encoding=response_encoding) self.datasets = self.DatasetsService(self) self.jobs = self.JobsService(self) self.models = self.ModelsService(self) self.projects = self.ProjectsService(self) self.routines = self.RoutinesService(self) self.rowAccessPolicies = self.RowAccessPoliciesService(self) self.tabledata = self.TabledataService(self) self.tables = self.TablesService(self) class DatasetsService(base_api.BaseApiService): """Service class for the datasets resource.""" _NAME = 'datasets' def __init__(self, client): super().__init__(client) self._upload_configs = {} def Delete(self, request, global_params=None): r"""Deletes the dataset specified by the datasetId value. Before you can delete a dataset, you must delete all its tables, either manually or by specifying deleteContents. Immediately after deletion, you can create another dataset with the same name. Args: request: (BigqueryDatasetsDeleteRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (BigqueryDatasetsDeleteResponse) The response message. """ config = self.GetMethodConfig('Delete') return self._RunMethod(config, request, global_params=global_params) Delete.method_config = lambda: base_api.ApiMethodInfo( http_method='DELETE', method_id='bigquery.datasets.delete', ordered_params=['projectId', 'datasetId'], path_params=['datasetId', 'projectId'], query_params=['deleteContents'], relative_path='projects/{projectId}/datasets/{datasetId}', request_field='', request_type_name='BigqueryDatasetsDeleteRequest', response_type_name='BigqueryDatasetsDeleteResponse', supports_download=False, ) def Get(self, request, global_params=None): r"""Returns the dataset specified by datasetID. Args: request: (BigqueryDatasetsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Dataset) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod(config, request, global_params=global_params) Get.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='bigquery.datasets.get', ordered_params=['projectId', 'datasetId'], path_params=['datasetId', 'projectId'], query_params=[], relative_path='projects/{projectId}/datasets/{datasetId}', request_field='', request_type_name='BigqueryDatasetsGetRequest', response_type_name='Dataset', supports_download=False, ) def Insert(self, request, global_params=None): r"""Creates a new empty dataset. Args: request: (BigqueryDatasetsInsertRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Dataset) The response message. """ config = self.GetMethodConfig('Insert') return self._RunMethod(config, request, global_params=global_params) Insert.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='bigquery.datasets.insert', ordered_params=['projectId'], path_params=['projectId'], query_params=[], relative_path='projects/{projectId}/datasets', request_field='dataset', request_type_name='BigqueryDatasetsInsertRequest', response_type_name='Dataset', supports_download=False, ) def List(self, request, global_params=None): r"""Lists all datasets in the specified project to which you have been granted the READER dataset role. Args: request: (BigqueryDatasetsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (DatasetList) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod(config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='bigquery.datasets.list', ordered_params=['projectId'], path_params=['projectId'], query_params=['all', 'filter', 'maxResults', 'pageToken'], relative_path='projects/{projectId}/datasets', request_field='', request_type_name='BigqueryDatasetsListRequest', response_type_name='DatasetList', supports_download=False, ) def Patch(self, request, global_params=None): r"""Updates information in an existing dataset. The update method replaces the entire dataset resource, whereas the patch method only replaces fields that are provided in the submitted dataset resource. This method supports patch semantics. Args: request: (BigqueryDatasetsPatchRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Dataset) The response message. """ config = self.GetMethodConfig('Patch') return self._RunMethod(config, request, global_params=global_params) Patch.method_config = lambda: base_api.ApiMethodInfo( http_method='PATCH', method_id='bigquery.datasets.patch', ordered_params=['projectId', 'datasetId'], path_params=['datasetId', 'projectId'], query_params=[], relative_path='projects/{projectId}/datasets/{datasetId}', request_field='dataset', request_type_name='BigqueryDatasetsPatchRequest', response_type_name='Dataset', supports_download=False, ) def Update(self, request, global_params=None): r"""Updates information in an existing dataset. The update method replaces the entire dataset resource, whereas the patch method only replaces fields that are provided in the submitted dataset resource. Args: request: (BigqueryDatasetsUpdateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Dataset) The response message. """ config = self.GetMethodConfig('Update') return self._RunMethod(config, request, global_params=global_params) Update.method_config = lambda: base_api.ApiMethodInfo( http_method='PUT', method_id='bigquery.datasets.update', ordered_params=['projectId', 'datasetId'], path_params=['datasetId', 'projectId'], query_params=[], relative_path='projects/{projectId}/datasets/{datasetId}', request_field='dataset', request_type_name='BigqueryDatasetsUpdateRequest', response_type_name='Dataset', supports_download=False, ) class JobsService(base_api.BaseApiService): """Service class for the jobs resource.""" _NAME = 'jobs' def __init__(self, client): super().__init__(client) self._upload_configs = { 'Insert': base_api.ApiUploadInfo( accept=['*/*'], max_size=None, resumable_multipart=True, resumable_path= '/resumable/upload/bigquery/v2/projects/{projectId}/jobs', simple_multipart=True, simple_path='/upload/bigquery/v2/projects/{projectId}/jobs', ), } def Cancel(self, request, global_params=None): r"""Requests that a job be cancelled. This call will return immediately, and the client will need to poll for the job status to see if the cancel completed successfully. Cancelled jobs may still incur costs. Args: request: (BigqueryJobsCancelRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (JobCancelResponse) The response message. """ config = self.GetMethodConfig('Cancel') return self._RunMethod(config, request, global_params=global_params) Cancel.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='bigquery.jobs.cancel', ordered_params=['projectId', 'jobId'], path_params=['jobId', 'projectId'], query_params=['location'], relative_path='projects/{projectId}/jobs/{jobId}/cancel', request_field='', request_type_name='BigqueryJobsCancelRequest', response_type_name='JobCancelResponse', supports_download=False, ) def Get(self, request, global_params=None): r"""Returns information about a specific job. Job information is available for a six month period after creation. Requires that you're the person who ran the job, or have the Is Owner project role. Args: request: (BigqueryJobsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Job) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod(config, request, global_params=global_params) Get.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='bigquery.jobs.get', ordered_params=['projectId', 'jobId'], path_params=['jobId', 'projectId'], query_params=['location'], relative_path='projects/{projectId}/jobs/{jobId}', request_field='', request_type_name='BigqueryJobsGetRequest', response_type_name='Job', supports_download=False, ) def GetQueryResults(self, request, global_params=None): r"""Retrieves the results of a query job. Args: request: (BigqueryJobsGetQueryResultsRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (GetQueryResultsResponse) The response message. """ config = self.GetMethodConfig('GetQueryResults') return self._RunMethod(config, request, global_params=global_params) GetQueryResults.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='bigquery.jobs.getQueryResults', ordered_params=['projectId', 'jobId'], path_params=['jobId', 'projectId'], query_params= ['location', 'maxResults', 'pageToken', 'startIndex', 'timeoutMs'], relative_path='projects/{projectId}/queries/{jobId}', request_field='', request_type_name='BigqueryJobsGetQueryResultsRequest', response_type_name='GetQueryResultsResponse', supports_download=False, ) def Insert(self, request, global_params=None, upload=None): r"""Starts a new asynchronous job. Requires the Can View project role. Args: request: (BigqueryJobsInsertRequest) input message global_params: (StandardQueryParameters, default: None) global arguments upload: (Upload, default: None) If present, upload this stream with the request. Returns: (Job) The response message. """ config = self.GetMethodConfig('Insert') upload_config = self.GetUploadConfig('Insert') return self._RunMethod( config, request, global_params=global_params, upload=upload, upload_config=upload_config) Insert.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='bigquery.jobs.insert', ordered_params=['projectId'], path_params=['projectId'], query_params=[], relative_path='projects/{projectId}/jobs', request_field='job', request_type_name='BigqueryJobsInsertRequest', response_type_name='Job', supports_download=False, ) def List(self, request, global_params=None): r"""Lists all jobs that you started in the specified project. Job information is available for a six month period after creation. The job list is sorted in reverse chronological order, by job creation time. Requires the Can View project role, or the Is Owner project role if you set the allUsers property. Args: request: (BigqueryJobsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (JobList) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod(config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='bigquery.jobs.list', ordered_params=['projectId'], path_params=['projectId'], query_params=[ 'allUsers', 'maxCreationTime', 'maxResults', 'minCreationTime', 'pageToken', 'parentJobId', 'projection', 'stateFilter' ], relative_path='projects/{projectId}/jobs', request_field='', request_type_name='BigqueryJobsListRequest', response_type_name='JobList', supports_download=False, ) def Query(self, request, global_params=None): r"""Runs a BigQuery SQL query synchronously and returns query results if the query completes within a specified timeout. Args: request: (BigqueryJobsQueryRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (QueryResponse) The response message. """ config = self.GetMethodConfig('Query') return self._RunMethod(config, request, global_params=global_params) Query.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='bigquery.jobs.query', ordered_params=['projectId'], path_params=['projectId'], query_params=[], relative_path='projects/{projectId}/queries', request_field='queryRequest', request_type_name='BigqueryJobsQueryRequest', response_type_name='QueryResponse', supports_download=False, ) class ModelsService(base_api.BaseApiService): """Service class for the models resource.""" _NAME = 'models' def __init__(self, client): super().__init__(client) self._upload_configs = {} def Delete(self, request, global_params=None): r"""Deletes the model specified by modelId from the dataset. Args: request: (BigqueryModelsDeleteRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (BigqueryModelsDeleteResponse) The response message. """ config = self.GetMethodConfig('Delete') return self._RunMethod(config, request, global_params=global_params) Delete.method_config = lambda: base_api.ApiMethodInfo( flat_path= 'projects/{projectsId}/datasets/{datasetsId}/models/{modelsId}', http_method='DELETE', method_id='bigquery.models.delete', ordered_params=['projectId', 'datasetId', 'modelId'], path_params=['datasetId', 'modelId', 'projectId'], query_params=[], relative_path= 'projects/{+projectId}/datasets/{+datasetId}/models/{+modelId}', request_field='', request_type_name='BigqueryModelsDeleteRequest', response_type_name='BigqueryModelsDeleteResponse', supports_download=False, ) def Get(self, request, global_params=None): r"""Gets the specified model resource by model ID. Args: request: (BigqueryModelsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Model) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod(config, request, global_params=global_params) Get.method_config = lambda: base_api.ApiMethodInfo( flat_path= 'projects/{projectsId}/datasets/{datasetsId}/models/{modelsId}', http_method='GET', method_id='bigquery.models.get', ordered_params=['projectId', 'datasetId', 'modelId'], path_params=['datasetId', 'modelId', 'projectId'], query_params=[], relative_path= 'projects/{+projectId}/datasets/{+datasetId}/models/{+modelId}', request_field='', request_type_name='BigqueryModelsGetRequest', response_type_name='Model', supports_download=False, ) def List(self, request, global_params=None): r"""Lists all models in the specified dataset. Requires the READER dataset role. Args: request: (BigqueryModelsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ListModelsResponse) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod(config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( flat_path='projects/{projectsId}/datasets/{datasetsId}/models', http_method='GET', method_id='bigquery.models.list', ordered_params=['projectId', 'datasetId'], path_params=['datasetId', 'projectId'], query_params=['maxResults', 'pageToken'], relative_path='projects/{+projectId}/datasets/{+datasetId}/models', request_field='', request_type_name='BigqueryModelsListRequest', response_type_name='ListModelsResponse', supports_download=False, ) def Patch(self, request, global_params=None): r"""Patch specific fields in the specified model. Args: request: (BigqueryModelsPatchRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Model) The response message. """ config = self.GetMethodConfig('Patch') return self._RunMethod(config, request, global_params=global_params) Patch.method_config = lambda: base_api.ApiMethodInfo( flat_path= 'projects/{projectsId}/datasets/{datasetsId}/models/{modelsId}', http_method='PATCH', method_id='bigquery.models.patch', ordered_params=['projectId', 'datasetId', 'modelId'], path_params=['datasetId', 'modelId', 'projectId'], query_params=[], relative_path= 'projects/{+projectId}/datasets/{+datasetId}/models/{+modelId}', request_field='model', request_type_name='BigqueryModelsPatchRequest', response_type_name='Model', supports_download=False, ) class ProjectsService(base_api.BaseApiService): """Service class for the projects resource.""" _NAME = 'projects' def __init__(self, client): super().__init__(client) self._upload_configs = {} def GetServiceAccount(self, request, global_params=None): r"""Returns the email address of the service account for your project used for interactions with Google Cloud KMS. Args: request: (BigqueryProjectsGetServiceAccountRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (GetServiceAccountResponse) The response message. """ config = self.GetMethodConfig('GetServiceAccount') return self._RunMethod(config, request, global_params=global_params) GetServiceAccount.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='bigquery.projects.getServiceAccount', ordered_params=['projectId'], path_params=['projectId'], query_params=[], relative_path='projects/{projectId}/serviceAccount', request_field='', request_type_name='BigqueryProjectsGetServiceAccountRequest', response_type_name='GetServiceAccountResponse', supports_download=False, ) def List(self, request, global_params=None): r"""Lists all projects to which you have been granted any project role. Args: request: (BigqueryProjectsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ProjectList) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod(config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='bigquery.projects.list', ordered_params=[], path_params=[], query_params=['maxResults', 'pageToken'], relative_path='projects', request_field='', request_type_name='BigqueryProjectsListRequest', response_type_name='ProjectList', supports_download=False, ) class RoutinesService(base_api.BaseApiService): """Service class for the routines resource.""" _NAME = 'routines' def __init__(self, client): super().__init__(client) self._upload_configs = {} def Delete(self, request, global_params=None): r"""Deletes the routine specified by routineId from the dataset. Args: request: (BigqueryRoutinesDeleteRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (BigqueryRoutinesDeleteResponse) The response message. """ config = self.GetMethodConfig('Delete') return self._RunMethod(config, request, global_params=global_params) Delete.method_config = lambda: base_api.ApiMethodInfo( flat_path= 'projects/{projectsId}/datasets/{datasetsId}/routines/{routinesId}', http_method='DELETE', method_id='bigquery.routines.delete', ordered_params=['projectId', 'datasetId', 'routineId'], path_params=['datasetId', 'projectId', 'routineId'], query_params=[], relative_path= 'projects/{+projectId}/datasets/{+datasetId}/routines/{+routineId}', request_field='', request_type_name='BigqueryRoutinesDeleteRequest', response_type_name='BigqueryRoutinesDeleteResponse', supports_download=False, ) def Get(self, request, global_params=None): r"""Gets the specified routine resource by routine ID. Args: request: (BigqueryRoutinesGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Routine) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod(config, request, global_params=global_params) Get.method_config = lambda: base_api.ApiMethodInfo( flat_path= 'projects/{projectsId}/datasets/{datasetsId}/routines/{routinesId}', http_method='GET', method_id='bigquery.routines.get', ordered_params=['projectId', 'datasetId', 'routineId'], path_params=['datasetId', 'projectId', 'routineId'], query_params=['readMask'], relative_path= 'projects/{+projectId}/datasets/{+datasetId}/routines/{+routineId}', request_field='', request_type_name='BigqueryRoutinesGetRequest', response_type_name='Routine', supports_download=False, ) def Insert(self, request, global_params=None): r"""Creates a new routine in the dataset. Args: request: (BigqueryRoutinesInsertRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Routine) The response message. """ config = self.GetMethodConfig('Insert') return self._RunMethod(config, request, global_params=global_params) Insert.method_config = lambda: base_api.ApiMethodInfo( flat_path='projects/{projectsId}/datasets/{datasetsId}/routines', http_method='POST', method_id='bigquery.routines.insert', ordered_params=['projectId', 'datasetId'], path_params=['datasetId', 'projectId'], query_params=[], relative_path='projects/{+projectId}/datasets/{+datasetId}/routines', request_field='routine', request_type_name='BigqueryRoutinesInsertRequest', response_type_name='Routine', supports_download=False, ) def List(self, request, global_params=None): r"""Lists all routines in the specified dataset. Requires the READER dataset role. Args: request: (BigqueryRoutinesListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ListRoutinesResponse) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod(config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( flat_path='projects/{projectsId}/datasets/{datasetsId}/routines', http_method='GET', method_id='bigquery.routines.list', ordered_params=['projectId', 'datasetId'], path_params=['datasetId', 'projectId'], query_params=['filter', 'maxResults', 'pageToken', 'readMask'], relative_path='projects/{+projectId}/datasets/{+datasetId}/routines', request_field='', request_type_name='BigqueryRoutinesListRequest', response_type_name='ListRoutinesResponse', supports_download=False, ) def Update(self, request, global_params=None): r"""Updates information in an existing routine. The update method replaces the entire Routine resource. Args: request: (BigqueryRoutinesUpdateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Routine) The response message. """ config = self.GetMethodConfig('Update') return self._RunMethod(config, request, global_params=global_params) Update.method_config = lambda: base_api.ApiMethodInfo( flat_path= 'projects/{projectsId}/datasets/{datasetsId}/routines/{routinesId}', http_method='PUT', method_id='bigquery.routines.update', ordered_params=['projectId', 'datasetId', 'routineId'], path_params=['datasetId', 'projectId', 'routineId'], query_params=[], relative_path= 'projects/{+projectId}/datasets/{+datasetId}/routines/{+routineId}', request_field='routine', request_type_name='BigqueryRoutinesUpdateRequest', response_type_name='Routine', supports_download=False, ) class RowAccessPoliciesService(base_api.BaseApiService): """Service class for the rowAccessPolicies resource.""" _NAME = 'rowAccessPolicies' def __init__(self, client): super().__init__(client) self._upload_configs = {} def List(self, request, global_params=None): r"""Lists all row access policies on the specified table. Args: request: (BigqueryRowAccessPoliciesListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ListRowAccessPoliciesResponse) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod(config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( flat_path= 'projects/{projectsId}/datasets/{datasetsId}/tables/{tablesId}/rowAccessPolicies', http_method='GET', method_id='bigquery.rowAccessPolicies.list', ordered_params=['projectId', 'datasetId', 'tableId'], path_params=['datasetId', 'projectId', 'tableId'], query_params=['pageSize', 'pageToken'], relative_path= 'projects/{+projectId}/datasets/{+datasetId}/tables/{+tableId}/rowAccessPolicies', request_field='', request_type_name='BigqueryRowAccessPoliciesListRequest', response_type_name='ListRowAccessPoliciesResponse', supports_download=False, ) class TabledataService(base_api.BaseApiService): """Service class for the tabledata resource.""" _NAME = 'tabledata' def __init__(self, client): super().__init__(client) self._upload_configs = {} def InsertAll(self, request, global_params=None): r"""Streams data into BigQuery one record at a time without needing to run a load job. Requires the WRITER dataset role. Args: request: (BigqueryTabledataInsertAllRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (TableDataInsertAllResponse) The response message. """ config = self.GetMethodConfig('InsertAll') return self._RunMethod(config, request, global_params=global_params) InsertAll.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='bigquery.tabledata.insertAll', ordered_params=['projectId', 'datasetId', 'tableId'], path_params=['datasetId', 'projectId', 'tableId'], query_params=[], relative_path= 'projects/{projectId}/datasets/{datasetId}/tables/{tableId}/insertAll', request_field='tableDataInsertAllRequest', request_type_name='BigqueryTabledataInsertAllRequest', response_type_name='TableDataInsertAllResponse', supports_download=False, ) def List(self, request, global_params=None): r"""Retrieves table data from a specified set of rows. Requires the READER dataset role. Args: request: (BigqueryTabledataListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (TableDataList) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod(config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='bigquery.tabledata.list', ordered_params=['projectId', 'datasetId', 'tableId'], path_params=['datasetId', 'projectId', 'tableId'], query_params= ['maxResults', 'pageToken', 'selectedFields', 'startIndex'], relative_path= 'projects/{projectId}/datasets/{datasetId}/tables/{tableId}/data', request_field='', request_type_name='BigqueryTabledataListRequest', response_type_name='TableDataList', supports_download=False, ) class TablesService(base_api.BaseApiService): """Service class for the tables resource.""" _NAME = 'tables' def __init__(self, client): super().__init__(client) self._upload_configs = {} def Delete(self, request, global_params=None): r"""Deletes the table specified by tableId from the dataset. If the table contains data, all the data will be deleted. Args: request: (BigqueryTablesDeleteRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (BigqueryTablesDeleteResponse) The response message. """ config = self.GetMethodConfig('Delete') return self._RunMethod(config, request, global_params=global_params) Delete.method_config = lambda: base_api.ApiMethodInfo( http_method='DELETE', method_id='bigquery.tables.delete', ordered_params=['projectId', 'datasetId', 'tableId'], path_params=['datasetId', 'projectId', 'tableId'], query_params=[], relative_path= 'projects/{projectId}/datasets/{datasetId}/tables/{tableId}', request_field='', request_type_name='BigqueryTablesDeleteRequest', response_type_name='BigqueryTablesDeleteResponse', supports_download=False, ) def Get(self, request, global_params=None): r"""Gets the specified table resource by table ID. This method does not return the data in the table, it only returns the table resource, which describes the structure of this table. Args: request: (BigqueryTablesGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Table) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod(config, request, global_params=global_params) Get.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='bigquery.tables.get', ordered_params=['projectId', 'datasetId', 'tableId'], path_params=['datasetId', 'projectId', 'tableId'], query_params=['selectedFields'], relative_path= 'projects/{projectId}/datasets/{datasetId}/tables/{tableId}', request_field='', request_type_name='BigqueryTablesGetRequest', response_type_name='Table', supports_download=False, ) def GetIamPolicy(self, request, global_params=None): r"""Gets the access control policy for a resource. Returns an empty policy if the resource exists and does not have a policy set. Args: request: (BigqueryTablesGetIamPolicyRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Policy) The response message. """ config = self.GetMethodConfig('GetIamPolicy') return self._RunMethod(config, request, global_params=global_params) GetIamPolicy.method_config = lambda: base_api.ApiMethodInfo( flat_path= 'projects/{projectsId}/datasets/{datasetsId}/tables/{tablesId}:getIamPolicy', http_method='POST', method_id='bigquery.tables.getIamPolicy', ordered_params=['resource'], path_params=['resource'], query_params=[], relative_path='{+resource}:getIamPolicy', request_field='getIamPolicyRequest', request_type_name='BigqueryTablesGetIamPolicyRequest', response_type_name='Policy', supports_download=False, ) def Insert(self, request, global_params=None): r"""Creates a new, empty table in the dataset. Args: request: (BigqueryTablesInsertRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Table) The response message. """ config = self.GetMethodConfig('Insert') return self._RunMethod(config, request, global_params=global_params) Insert.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='bigquery.tables.insert', ordered_params=['projectId', 'datasetId'], path_params=['datasetId', 'projectId'], query_params=[], relative_path='projects/{projectId}/datasets/{datasetId}/tables', request_field='table', request_type_name='BigqueryTablesInsertRequest', response_type_name='Table', supports_download=False, ) def List(self, request, global_params=None): r"""Lists all tables in the specified dataset. Requires the READER dataset role. Args: request: (BigqueryTablesListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (TableList) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod(config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='bigquery.tables.list', ordered_params=['projectId', 'datasetId'], path_params=['datasetId', 'projectId'], query_params=['maxResults', 'pageToken'], relative_path='projects/{projectId}/datasets/{datasetId}/tables', request_field='', request_type_name='BigqueryTablesListRequest', response_type_name='TableList', supports_download=False, ) def Patch(self, request, global_params=None): r"""Updates information in an existing table. The update method replaces the entire table resource, whereas the patch method only replaces fields that are provided in the submitted table resource. This method supports patch semantics. Args: request: (BigqueryTablesPatchRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Table) The response message. """ config = self.GetMethodConfig('Patch') return self._RunMethod(config, request, global_params=global_params) Patch.method_config = lambda: base_api.ApiMethodInfo( http_method='PATCH', method_id='bigquery.tables.patch', ordered_params=['projectId', 'datasetId', 'tableId'], path_params=['datasetId', 'projectId', 'tableId'], query_params=[], relative_path= 'projects/{projectId}/datasets/{datasetId}/tables/{tableId}', request_field='table', request_type_name='BigqueryTablesPatchRequest', response_type_name='Table', supports_download=False, ) def SetIamPolicy(self, request, global_params=None): r"""Sets the access control policy on the specified resource. Replaces any existing policy. Can return `NOT_FOUND`, `INVALID_ARGUMENT`, and `PERMISSION_DENIED` errors. Args: request: (BigqueryTablesSetIamPolicyRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Policy) The response message. """ config = self.GetMethodConfig('SetIamPolicy') return self._RunMethod(config, request, global_params=global_params) SetIamPolicy.method_config = lambda: base_api.ApiMethodInfo( flat_path= 'projects/{projectsId}/datasets/{datasetsId}/tables/{tablesId}:setIamPolicy', http_method='POST', method_id='bigquery.tables.setIamPolicy', ordered_params=['resource'], path_params=['resource'], query_params=[], relative_path='{+resource}:setIamPolicy', request_field='setIamPolicyRequest', request_type_name='BigqueryTablesSetIamPolicyRequest', response_type_name='Policy', supports_download=False, ) def TestIamPermissions(self, request, global_params=None): r"""Returns permissions that a caller has on the specified resource. If the resource does not exist, this will return an empty set of permissions, not a `NOT_FOUND` error. Note: This operation is designed to be used for building permission-aware UIs and command-line tools, not for authorization checking. This operation may "fail open" without warning. Args: request: (BigqueryTablesTestIamPermissionsRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (TestIamPermissionsResponse) The response message. """ config = self.GetMethodConfig('TestIamPermissions') return self._RunMethod(config, request, global_params=global_params) TestIamPermissions.method_config = lambda: base_api.ApiMethodInfo( flat_path= 'projects/{projectsId}/datasets/{datasetsId}/tables/{tablesId}:testIamPermissions', http_method='POST', method_id='bigquery.tables.testIamPermissions', ordered_params=['resource'], path_params=['resource'], query_params=[], relative_path='{+resource}:testIamPermissions', request_field='testIamPermissionsRequest', request_type_name='BigqueryTablesTestIamPermissionsRequest', response_type_name='TestIamPermissionsResponse', supports_download=False, ) def Update(self, request, global_params=None): r"""Updates information in an existing table. The update method replaces the entire table resource, whereas the patch method only replaces fields that are provided in the submitted table resource. Args: request: (BigqueryTablesUpdateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Table) The response message. """ config = self.GetMethodConfig('Update') return self._RunMethod(config, request, global_params=global_params) Update.method_config = lambda: base_api.ApiMethodInfo( http_method='PUT', method_id='bigquery.tables.update', ordered_params=['projectId', 'datasetId', 'tableId'], path_params=['datasetId', 'projectId', 'tableId'], query_params=[], relative_path= 'projects/{projectId}/datasets/{datasetId}/tables/{tableId}', request_field='table', request_type_name='BigqueryTablesUpdateRequest', response_type_name='Table', supports_download=False, ) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py0000644000175100017510000070101615206113032030645 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Generated message classes for bigquery version v2. A data platform for customers to create, manage, share and query data. """ # NOTE: This file is autogenerated and should not be edited by hand. from apitools.base.protorpclite import message_types as _message_types from apitools.base.protorpclite import messages as _messages from apitools.base.py import encoding from apitools.base.py import extra_types package = 'bigquery' class AggregateClassificationMetrics(_messages.Message): r"""Aggregate metrics for classification/classifier models. For multi-class models, the metrics are either macro-averaged or micro-averaged. When macro- averaged, the metrics are calculated for each label and then an unweighted average is taken of those values. When micro-averaged, the metric is calculated globally by counting the total number of correctly predicted rows. Fields: accuracy: Accuracy is the fraction of predictions given the correct label. For multiclass this is a micro-averaged metric. f1Score: The F1 score is an average of recall and precision. For multiclass this is a macro-averaged metric. logLoss: Logarithmic Loss. For multiclass this is a macro-averaged metric. precision: Precision is the fraction of actual positive predictions that had positive actual labels. For multiclass this is a macro-averaged metric treating each class as a binary classifier. recall: Recall is the fraction of actual positive labels that were given a positive prediction. For multiclass this is a macro-averaged metric. rocAuc: Area Under a ROC Curve. For multiclass this is a macro-averaged metric. threshold: Threshold at which the metrics are computed. For binary classification models this is the positive class threshold. For multi- class classfication models this is the confidence threshold. """ accuracy = _messages.FloatField(1) f1Score = _messages.FloatField(2) logLoss = _messages.FloatField(3) precision = _messages.FloatField(4) recall = _messages.FloatField(5) rocAuc = _messages.FloatField(6) threshold = _messages.FloatField(7) class Argument(_messages.Message): r"""Input/output argument of a function or a stored procedure. Enums: ArgumentKindValueValuesEnum: Optional. Defaults to FIXED_TYPE. ModeValueValuesEnum: Optional. Specifies whether the argument is input or output. Can be set for procedures only. Fields: argumentKind: Optional. Defaults to FIXED_TYPE. dataType: Required unless argument_kind = ANY_TYPE. mode: Optional. Specifies whether the argument is input or output. Can be set for procedures only. name: Optional. The name of this argument. Can be absent for function return argument. """ class ArgumentKindValueValuesEnum(_messages.Enum): r"""Optional. Defaults to FIXED_TYPE. Values: ARGUMENT_KIND_UNSPECIFIED: FIXED_TYPE: The argument is a variable with fully specified type, which can be a struct or an array, but not a table. ANY_TYPE: The argument is any type, including struct or array, but not a table. To be added: FIXED_TABLE, ANY_TABLE """ ARGUMENT_KIND_UNSPECIFIED = 0 FIXED_TYPE = 1 ANY_TYPE = 2 class ModeValueValuesEnum(_messages.Enum): r"""Optional. Specifies whether the argument is input or output. Can be set for procedures only. Values: MODE_UNSPECIFIED: IN: The argument is input-only. OUT: The argument is output-only. INOUT: The argument is both an input and an output. """ MODE_UNSPECIFIED = 0 IN = 1 OUT = 2 INOUT = 3 argumentKind = _messages.EnumField('ArgumentKindValueValuesEnum', 1) dataType = _messages.MessageField('StandardSqlDataType', 2) mode = _messages.EnumField('ModeValueValuesEnum', 3) name = _messages.StringField(4) class ArimaCoefficients(_messages.Message): r"""Arima coefficients. Fields: autoRegressiveCoefficients: Auto-regressive coefficients, an array of double. interceptCoefficient: Intercept coefficient, just a double not an array. movingAverageCoefficients: Moving-average coefficients, an array of double. """ autoRegressiveCoefficients = _messages.FloatField(1, repeated=True) interceptCoefficient = _messages.FloatField(2) movingAverageCoefficients = _messages.FloatField(3, repeated=True) class ArimaFittingMetrics(_messages.Message): r"""ARIMA model fitting metrics. Fields: aic: AIC. logLikelihood: Log-likelihood. variance: Variance. """ aic = _messages.FloatField(1) logLikelihood = _messages.FloatField(2) variance = _messages.FloatField(3) class ArimaForecastingMetrics(_messages.Message): r"""Model evaluation metrics for ARIMA forecasting models. Enums: SeasonalPeriodsValueListEntryValuesEnum: Fields: arimaFittingMetrics: Arima model fitting metrics. arimaSingleModelForecastingMetrics: Repeated as there can be many metric sets (one for each model) in auto-arima and the large-scale case. hasDrift: Whether Arima model fitted with drift or not. It is always false when d is not 1. nonSeasonalOrder: Non-seasonal order. seasonalPeriods: Seasonal periods. Repeated because multiple periods are supported for one time series. timeSeriesId: Id to differentiate different time series for the large- scale case. """ class SeasonalPeriodsValueListEntryValuesEnum(_messages.Enum): r"""SeasonalPeriodsValueListEntryValuesEnum enum type. Values: SEASONAL_PERIOD_TYPE_UNSPECIFIED: NO_SEASONALITY: No seasonality DAILY: Daily period, 24 hours. WEEKLY: Weekly period, 7 days. MONTHLY: Monthly period, 30 days or irregular. QUARTERLY: Quarterly period, 90 days or irregular. YEARLY: Yearly period, 365 days or irregular. """ SEASONAL_PERIOD_TYPE_UNSPECIFIED = 0 NO_SEASONALITY = 1 DAILY = 2 WEEKLY = 3 MONTHLY = 4 QUARTERLY = 5 YEARLY = 6 arimaFittingMetrics = _messages.MessageField( 'ArimaFittingMetrics', 1, repeated=True) arimaSingleModelForecastingMetrics = _messages.MessageField( 'ArimaSingleModelForecastingMetrics', 2, repeated=True) hasDrift = _messages.BooleanField(3, repeated=True) nonSeasonalOrder = _messages.MessageField('ArimaOrder', 4, repeated=True) seasonalPeriods = _messages.EnumField( 'SeasonalPeriodsValueListEntryValuesEnum', 5, repeated=True) timeSeriesId = _messages.StringField(6, repeated=True) class ArimaModelInfo(_messages.Message): r"""Arima model information. Enums: SeasonalPeriodsValueListEntryValuesEnum: Fields: arimaCoefficients: Arima coefficients. arimaFittingMetrics: Arima fitting metrics. hasDrift: Whether Arima model fitted with drift or not. It is always false when d is not 1. nonSeasonalOrder: Non-seasonal order. seasonalPeriods: Seasonal periods. Repeated because multiple periods are supported for one time series. timeSeriesId: The id to indicate different time series. """ class SeasonalPeriodsValueListEntryValuesEnum(_messages.Enum): r"""SeasonalPeriodsValueListEntryValuesEnum enum type. Values: SEASONAL_PERIOD_TYPE_UNSPECIFIED: NO_SEASONALITY: No seasonality DAILY: Daily period, 24 hours. WEEKLY: Weekly period, 7 days. MONTHLY: Monthly period, 30 days or irregular. QUARTERLY: Quarterly period, 90 days or irregular. YEARLY: Yearly period, 365 days or irregular. """ SEASONAL_PERIOD_TYPE_UNSPECIFIED = 0 NO_SEASONALITY = 1 DAILY = 2 WEEKLY = 3 MONTHLY = 4 QUARTERLY = 5 YEARLY = 6 arimaCoefficients = _messages.MessageField('ArimaCoefficients', 1) arimaFittingMetrics = _messages.MessageField('ArimaFittingMetrics', 2) hasDrift = _messages.BooleanField(3) nonSeasonalOrder = _messages.MessageField('ArimaOrder', 4) seasonalPeriods = _messages.EnumField( 'SeasonalPeriodsValueListEntryValuesEnum', 5, repeated=True) timeSeriesId = _messages.StringField(6) class ArimaOrder(_messages.Message): r"""Arima order, can be used for both non-seasonal and seasonal parts. Fields: d: Order of the differencing part. p: Order of the autoregressive part. q: Order of the moving-average part. """ d = _messages.IntegerField(1) p = _messages.IntegerField(2) q = _messages.IntegerField(3) class ArimaResult(_messages.Message): r"""(Auto-)arima fitting result. Wrap everything in ArimaResult for easier refactoring if we want to use model-specific iteration results. Enums: SeasonalPeriodsValueListEntryValuesEnum: Fields: arimaModelInfo: This message is repeated because there are multiple arima models fitted in auto-arima. For non-auto-arima model, its size is one. seasonalPeriods: Seasonal periods. Repeated because multiple periods are supported for one time series. """ class SeasonalPeriodsValueListEntryValuesEnum(_messages.Enum): r"""SeasonalPeriodsValueListEntryValuesEnum enum type. Values: SEASONAL_PERIOD_TYPE_UNSPECIFIED: NO_SEASONALITY: No seasonality DAILY: Daily period, 24 hours. WEEKLY: Weekly period, 7 days. MONTHLY: Monthly period, 30 days or irregular. QUARTERLY: Quarterly period, 90 days or irregular. YEARLY: Yearly period, 365 days or irregular. """ SEASONAL_PERIOD_TYPE_UNSPECIFIED = 0 NO_SEASONALITY = 1 DAILY = 2 WEEKLY = 3 MONTHLY = 4 QUARTERLY = 5 YEARLY = 6 arimaModelInfo = _messages.MessageField('ArimaModelInfo', 1, repeated=True) seasonalPeriods = _messages.EnumField( 'SeasonalPeriodsValueListEntryValuesEnum', 2, repeated=True) class ArimaSingleModelForecastingMetrics(_messages.Message): r"""Model evaluation metrics for a single ARIMA forecasting model. Enums: SeasonalPeriodsValueListEntryValuesEnum: Fields: arimaFittingMetrics: Arima fitting metrics. hasDrift: Is arima model fitted with drift or not. It is always false when d is not 1. nonSeasonalOrder: Non-seasonal order. seasonalPeriods: Seasonal periods. Repeated because multiple periods are supported for one time series. timeSeriesId: The id to indicate different time series. """ class SeasonalPeriodsValueListEntryValuesEnum(_messages.Enum): r"""SeasonalPeriodsValueListEntryValuesEnum enum type. Values: SEASONAL_PERIOD_TYPE_UNSPECIFIED: NO_SEASONALITY: No seasonality DAILY: Daily period, 24 hours. WEEKLY: Weekly period, 7 days. MONTHLY: Monthly period, 30 days or irregular. QUARTERLY: Quarterly period, 90 days or irregular. YEARLY: Yearly period, 365 days or irregular. """ SEASONAL_PERIOD_TYPE_UNSPECIFIED = 0 NO_SEASONALITY = 1 DAILY = 2 WEEKLY = 3 MONTHLY = 4 QUARTERLY = 5 YEARLY = 6 arimaFittingMetrics = _messages.MessageField('ArimaFittingMetrics', 1) hasDrift = _messages.BooleanField(2) nonSeasonalOrder = _messages.MessageField('ArimaOrder', 3) seasonalPeriods = _messages.EnumField( 'SeasonalPeriodsValueListEntryValuesEnum', 4, repeated=True) timeSeriesId = _messages.StringField(5) class AuditConfig(_messages.Message): r"""Specifies the audit configuration for a service. The configuration determines which permission types are logged, and what identities, if any, are exempted from logging. An AuditConfig must have one or more AuditLogConfigs. If there are AuditConfigs for both `allServices` and a specific service, the union of the two AuditConfigs is used for that service: the log_types specified in each AuditConfig are enabled, and the exempted_members in each AuditLogConfig are exempted. Example Policy with multiple AuditConfigs: { "audit_configs": [ { "service": "allServices", "audit_log_configs": [ { "log_type": "DATA_READ", "exempted_members": [ "user:jose@example.com" ] }, { "log_type": "DATA_WRITE" }, { "log_type": "ADMIN_READ" } ] }, { "service": "sampleservice.googleapis.com", "audit_log_configs": [ { "log_type": "DATA_READ" }, { "log_type": "DATA_WRITE", "exempted_members": [ "user:aliya@example.com" ] } ] } ] } For sampleservice, this policy enables DATA_READ, DATA_WRITE and ADMIN_READ logging. It also exempts jose@example.com from DATA_READ logging, and aliya@example.com from DATA_WRITE logging. Fields: auditLogConfigs: The configuration for logging of each type of permission. service: Specifies a service that will be enabled for audit logging. For example, `storage.googleapis.com`, `cloudsql.googleapis.com`. `allServices` is a special value that covers all services. """ auditLogConfigs = _messages.MessageField('AuditLogConfig', 1, repeated=True) service = _messages.StringField(2) class AuditLogConfig(_messages.Message): r"""Provides the configuration for logging a type of permissions. Example: { "audit_log_configs": [ { "log_type": "DATA_READ", "exempted_members": [ "user:jose@example.com" ] }, { "log_type": "DATA_WRITE" } ] } This enables 'DATA_READ' and 'DATA_WRITE' logging, while exempting jose@example.com from DATA_READ logging. Enums: LogTypeValueValuesEnum: The log type that this config enables. Fields: exemptedMembers: Specifies the identities that do not cause logging for this type of permission. Follows the same format of Binding.members. logType: The log type that this config enables. """ class LogTypeValueValuesEnum(_messages.Enum): r"""The log type that this config enables. Values: LOG_TYPE_UNSPECIFIED: Default case. Should never be this. ADMIN_READ: Admin reads. Example: CloudIAM getIamPolicy DATA_WRITE: Data writes. Example: CloudSQL Users create DATA_READ: Data reads. Example: CloudSQL Users list """ LOG_TYPE_UNSPECIFIED = 0 ADMIN_READ = 1 DATA_WRITE = 2 DATA_READ = 3 exemptedMembers = _messages.StringField(1, repeated=True) logType = _messages.EnumField('LogTypeValueValuesEnum', 2) class BigQueryModelTraining(_messages.Message): r"""A BigQueryModelTraining object. Fields: currentIteration: [Output-only, Beta] Index of current ML training iteration. Updated during create model query job to show job progress. expectedTotalIterations: [Output-only, Beta] Expected number of iterations for the create model query job specified as num_iterations in the input query. The actual total number of iterations may be less than this number due to early stop. """ currentIteration = _messages.IntegerField(1, variant=_messages.Variant.INT32) expectedTotalIterations = _messages.IntegerField(2) class BigqueryDatasetsDeleteRequest(_messages.Message): r"""A BigqueryDatasetsDeleteRequest object. Fields: datasetId: Dataset ID of dataset being deleted deleteContents: If True, delete all the tables in the dataset. If False and the dataset contains tables, the request will fail. Default is False projectId: Project ID of the dataset being deleted """ datasetId = _messages.StringField(1, required=True) deleteContents = _messages.BooleanField(2) projectId = _messages.StringField(3, required=True) class BigqueryDatasetsDeleteResponse(_messages.Message): r"""An empty BigqueryDatasetsDelete response.""" class BigqueryDatasetsGetRequest(_messages.Message): r"""A BigqueryDatasetsGetRequest object. Fields: datasetId: Dataset ID of the requested dataset projectId: Project ID of the requested dataset """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) class BigqueryDatasetsInsertRequest(_messages.Message): r"""A BigqueryDatasetsInsertRequest object. Fields: dataset: A Dataset resource to be passed as the request body. projectId: Project ID of the new dataset """ dataset = _messages.MessageField('Dataset', 1) projectId = _messages.StringField(2, required=True) class BigqueryDatasetsListRequest(_messages.Message): r"""A BigqueryDatasetsListRequest object. Fields: all: Whether to list all datasets, including hidden ones filter: An expression for filtering the results of the request by label. The syntax is "labels.[:]". Multiple filters can be ANDed together by connecting with a space. Example: "labels.department:receiving labels.active". See Filtering datasets using labels for details. maxResults: The maximum number of results to return pageToken: Page token, returned by a previous call, to request the next page of results projectId: Project ID of the datasets to be listed """ all = _messages.BooleanField(1) filter = _messages.StringField(2) maxResults = _messages.IntegerField(3, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(4) projectId = _messages.StringField(5, required=True) class BigqueryDatasetsPatchRequest(_messages.Message): r"""A BigqueryDatasetsPatchRequest object. Fields: dataset: A Dataset resource to be passed as the request body. datasetId: Dataset ID of the dataset being updated projectId: Project ID of the dataset being updated """ dataset = _messages.MessageField('Dataset', 1) datasetId = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) class BigqueryDatasetsUpdateRequest(_messages.Message): r"""A BigqueryDatasetsUpdateRequest object. Fields: dataset: A Dataset resource to be passed as the request body. datasetId: Dataset ID of the dataset being updated projectId: Project ID of the dataset being updated """ dataset = _messages.MessageField('Dataset', 1) datasetId = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) class BigqueryJobsCancelRequest(_messages.Message): r"""A BigqueryJobsCancelRequest object. Fields: jobId: [Required] Job ID of the job to cancel location: The geographic location of the job. Required except for US and EU. See details at https://cloud.google.com/bigquery/docs/locations#spec ifying_your_location. projectId: [Required] Project ID of the job to cancel """ jobId = _messages.StringField(1, required=True) location = _messages.StringField(2) projectId = _messages.StringField(3, required=True) class BigqueryJobsGetQueryResultsRequest(_messages.Message): r"""A BigqueryJobsGetQueryResultsRequest object. Fields: jobId: [Required] Job ID of the query job location: The geographic location where the job should run. Required except for US and EU. See details at https://cloud.google.com/bigquery/d ocs/locations#specifying_your_location. maxResults: Maximum number of results to read pageToken: Page token, returned by a previous call, to request the next page of results projectId: [Required] Project ID of the query job startIndex: Zero-based index of the starting row timeoutMs: How long to wait for the query to complete, in milliseconds, before returning. Default is 10 seconds. If the timeout passes before the job completes, the 'jobComplete' field in the response will be false """ jobId = _messages.StringField(1, required=True) location = _messages.StringField(2) maxResults = _messages.IntegerField(3, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(4) projectId = _messages.StringField(5, required=True) startIndex = _messages.IntegerField(6, variant=_messages.Variant.UINT64) timeoutMs = _messages.IntegerField(7, variant=_messages.Variant.UINT32) class BigqueryJobsGetRequest(_messages.Message): r"""A BigqueryJobsGetRequest object. Fields: jobId: [Required] Job ID of the requested job location: The geographic location of the job. Required except for US and EU. See details at https://cloud.google.com/bigquery/docs/locations#spec ifying_your_location. projectId: [Required] Project ID of the requested job """ jobId = _messages.StringField(1, required=True) location = _messages.StringField(2) projectId = _messages.StringField(3, required=True) class BigqueryJobsInsertRequest(_messages.Message): r"""A BigqueryJobsInsertRequest object. Fields: job: A Job resource to be passed as the request body. projectId: Project ID of the project that will be billed for the job """ job = _messages.MessageField('Job', 1) projectId = _messages.StringField(2, required=True) class BigqueryJobsListRequest(_messages.Message): r"""A BigqueryJobsListRequest object. Enums: ProjectionValueValuesEnum: Restrict information returned to a set of selected fields StateFilterValueValuesEnum: Filter for job state Fields: allUsers: Whether to display jobs owned by all users in the project. Default false maxCreationTime: Max value for job creation time, in milliseconds since the POSIX epoch. If set, only jobs created before or at this timestamp are returned maxResults: Maximum number of results to return minCreationTime: Min value for job creation time, in milliseconds since the POSIX epoch. If set, only jobs created after or at this timestamp are returned pageToken: Page token, returned by a previous call, to request the next page of results parentJobId: If set, retrieves only jobs whose parent is this job. Otherwise, retrieves only jobs which have no parent projectId: Project ID of the jobs to list projection: Restrict information returned to a set of selected fields stateFilter: Filter for job state """ class ProjectionValueValuesEnum(_messages.Enum): r"""Restrict information returned to a set of selected fields Values: full: Includes all job data minimal: Does not include the job configuration """ full = 0 minimal = 1 class StateFilterValueValuesEnum(_messages.Enum): r"""Filter for job state Values: done: Finished jobs pending: Pending jobs running: Running jobs """ done = 0 pending = 1 running = 2 allUsers = _messages.BooleanField(1) maxCreationTime = _messages.IntegerField(2, variant=_messages.Variant.UINT64) maxResults = _messages.IntegerField(3, variant=_messages.Variant.UINT32) minCreationTime = _messages.IntegerField(4, variant=_messages.Variant.UINT64) pageToken = _messages.StringField(5) parentJobId = _messages.StringField(6) projectId = _messages.StringField(7, required=True) projection = _messages.EnumField('ProjectionValueValuesEnum', 8) stateFilter = _messages.EnumField( 'StateFilterValueValuesEnum', 9, repeated=True) class BigqueryJobsQueryRequest(_messages.Message): r"""A BigqueryJobsQueryRequest object. Fields: projectId: Project ID of the project billed for the query queryRequest: A QueryRequest resource to be passed as the request body. """ projectId = _messages.StringField(1, required=True) queryRequest = _messages.MessageField('QueryRequest', 2) class BigqueryModelsDeleteRequest(_messages.Message): r"""A BigqueryModelsDeleteRequest object. Fields: datasetId: Required. Dataset ID of the model to delete. modelId: Required. Model ID of the model to delete. projectId: Required. Project ID of the model to delete. """ datasetId = _messages.StringField(1, required=True) modelId = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) class BigqueryModelsDeleteResponse(_messages.Message): r"""An empty BigqueryModelsDelete response.""" class BigqueryModelsGetRequest(_messages.Message): r"""A BigqueryModelsGetRequest object. Fields: datasetId: Required. Dataset ID of the requested model. modelId: Required. Model ID of the requested model. projectId: Required. Project ID of the requested model. """ datasetId = _messages.StringField(1, required=True) modelId = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) class BigqueryModelsListRequest(_messages.Message): r"""A BigqueryModelsListRequest object. Fields: datasetId: Required. Dataset ID of the models to list. maxResults: The maximum number of results to return in a single response page. Leverage the page tokens to iterate through the entire collection. pageToken: Page token, returned by a previous call to request the next page of results projectId: Required. Project ID of the models to list. """ datasetId = _messages.StringField(1, required=True) maxResults = _messages.IntegerField(2, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(3) projectId = _messages.StringField(4, required=True) class BigqueryModelsPatchRequest(_messages.Message): r"""A BigqueryModelsPatchRequest object. Fields: datasetId: Required. Dataset ID of the model to patch. model: A Model resource to be passed as the request body. modelId: Required. Model ID of the model to patch. projectId: Required. Project ID of the model to patch. """ datasetId = _messages.StringField(1, required=True) model = _messages.MessageField('Model', 2) modelId = _messages.StringField(3, required=True) projectId = _messages.StringField(4, required=True) class BigqueryProjectsGetServiceAccountRequest(_messages.Message): r"""A BigqueryProjectsGetServiceAccountRequest object. Fields: projectId: Project ID for which the service account is requested. """ projectId = _messages.StringField(1, required=True) class BigqueryProjectsListRequest(_messages.Message): r"""A BigqueryProjectsListRequest object. Fields: maxResults: Maximum number of results to return pageToken: Page token, returned by a previous call, to request the next page of results """ maxResults = _messages.IntegerField(1, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(2) class BigqueryRoutinesDeleteRequest(_messages.Message): r"""A BigqueryRoutinesDeleteRequest object. Fields: datasetId: Required. Dataset ID of the routine to delete projectId: Required. Project ID of the routine to delete routineId: Required. Routine ID of the routine to delete """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) routineId = _messages.StringField(3, required=True) class BigqueryRoutinesDeleteResponse(_messages.Message): r"""An empty BigqueryRoutinesDelete response.""" class BigqueryRoutinesGetRequest(_messages.Message): r"""A BigqueryRoutinesGetRequest object. Fields: datasetId: Required. Dataset ID of the requested routine projectId: Required. Project ID of the requested routine readMask: If set, only the Routine fields in the field mask are returned in the response. If unset, all Routine fields are returned. routineId: Required. Routine ID of the requested routine """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) readMask = _messages.StringField(3) routineId = _messages.StringField(4, required=True) class BigqueryRoutinesInsertRequest(_messages.Message): r"""A BigqueryRoutinesInsertRequest object. Fields: datasetId: Required. Dataset ID of the new routine projectId: Required. Project ID of the new routine routine: A Routine resource to be passed as the request body. """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) routine = _messages.MessageField('Routine', 3) class BigqueryRoutinesListRequest(_messages.Message): r"""A BigqueryRoutinesListRequest object. Fields: datasetId: Required. Dataset ID of the routines to list filter: If set, then only the Routines matching this filter are returned. The current supported form is either "routine_type:" or "routineType:", where is a RoutineType enum. Example: "routineType:SCALAR_FUNCTION". maxResults: The maximum number of results to return in a single response page. Leverage the page tokens to iterate through the entire collection. pageToken: Page token, returned by a previous call, to request the next page of results projectId: Required. Project ID of the routines to list readMask: If set, then only the Routine fields in the field mask, as well as project_id, dataset_id and routine_id, are returned in the response. If unset, then the following Routine fields are returned: etag, project_id, dataset_id, routine_id, routine_type, creation_time, last_modified_time, and language. """ datasetId = _messages.StringField(1, required=True) filter = _messages.StringField(2) maxResults = _messages.IntegerField(3, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(4) projectId = _messages.StringField(5, required=True) readMask = _messages.StringField(6) class BigqueryRoutinesUpdateRequest(_messages.Message): r"""A BigqueryRoutinesUpdateRequest object. Fields: datasetId: Required. Dataset ID of the routine to update projectId: Required. Project ID of the routine to update routine: A Routine resource to be passed as the request body. routineId: Required. Routine ID of the routine to update """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) routine = _messages.MessageField('Routine', 3) routineId = _messages.StringField(4, required=True) class BigqueryRowAccessPoliciesListRequest(_messages.Message): r"""A BigqueryRowAccessPoliciesListRequest object. Fields: datasetId: Required. Dataset ID of row access policies to list. pageSize: The maximum number of results to return in a single response page. Leverage the page tokens to iterate through the entire collection. pageToken: Page token, returned by a previous call, to request the next page of results. projectId: Required. Project ID of the row access policies to list. tableId: Required. Table ID of the table to list row access policies. """ datasetId = _messages.StringField(1, required=True) pageSize = _messages.IntegerField(2, variant=_messages.Variant.INT32) pageToken = _messages.StringField(3) projectId = _messages.StringField(4, required=True) tableId = _messages.StringField(5, required=True) class BigqueryTabledataInsertAllRequest(_messages.Message): r"""A BigqueryTabledataInsertAllRequest object. Fields: datasetId: Dataset ID of the destination table. projectId: Project ID of the destination table. tableDataInsertAllRequest: A TableDataInsertAllRequest resource to be passed as the request body. tableId: Table ID of the destination table. """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) tableDataInsertAllRequest = _messages.MessageField( 'TableDataInsertAllRequest', 3) tableId = _messages.StringField(4, required=True) class BigqueryTabledataListRequest(_messages.Message): r"""A BigqueryTabledataListRequest object. Fields: datasetId: Dataset ID of the table to read maxResults: Maximum number of results to return pageToken: Page token, returned by a previous call, identifying the result set projectId: Project ID of the table to read selectedFields: List of fields to return (comma-separated). If unspecified, all fields are returned startIndex: Zero-based index of the starting row to read tableId: Table ID of the table to read """ datasetId = _messages.StringField(1, required=True) maxResults = _messages.IntegerField(2, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(3) projectId = _messages.StringField(4, required=True) selectedFields = _messages.StringField(5) startIndex = _messages.IntegerField(6, variant=_messages.Variant.UINT64) tableId = _messages.StringField(7, required=True) class BigqueryTablesDeleteRequest(_messages.Message): r"""A BigqueryTablesDeleteRequest object. Fields: datasetId: Dataset ID of the table to delete projectId: Project ID of the table to delete tableId: Table ID of the table to delete """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) tableId = _messages.StringField(3, required=True) class BigqueryTablesDeleteResponse(_messages.Message): r"""An empty BigqueryTablesDelete response.""" class BigqueryTablesGetIamPolicyRequest(_messages.Message): r"""A BigqueryTablesGetIamPolicyRequest object. Fields: getIamPolicyRequest: A GetIamPolicyRequest resource to be passed as the request body. resource: REQUIRED: The resource for which the policy is being requested. See the operation documentation for the appropriate value for this field. """ getIamPolicyRequest = _messages.MessageField('GetIamPolicyRequest', 1) resource = _messages.StringField(2, required=True) class BigqueryTablesGetRequest(_messages.Message): r"""A BigqueryTablesGetRequest object. Fields: datasetId: Dataset ID of the requested table projectId: Project ID of the requested table selectedFields: List of fields to return (comma-separated). If unspecified, all fields are returned tableId: Table ID of the requested table """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) selectedFields = _messages.StringField(3) tableId = _messages.StringField(4, required=True) class BigqueryTablesInsertRequest(_messages.Message): r"""A BigqueryTablesInsertRequest object. Fields: datasetId: Dataset ID of the new table projectId: Project ID of the new table table: A Table resource to be passed as the request body. """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) table = _messages.MessageField('Table', 3) class BigqueryTablesListRequest(_messages.Message): r"""A BigqueryTablesListRequest object. Fields: datasetId: Dataset ID of the tables to list maxResults: Maximum number of results to return pageToken: Page token, returned by a previous call, to request the next page of results projectId: Project ID of the tables to list """ datasetId = _messages.StringField(1, required=True) maxResults = _messages.IntegerField(2, variant=_messages.Variant.UINT32) pageToken = _messages.StringField(3) projectId = _messages.StringField(4, required=True) class BigqueryTablesPatchRequest(_messages.Message): r"""A BigqueryTablesPatchRequest object. Fields: datasetId: Dataset ID of the table to update projectId: Project ID of the table to update table: A Table resource to be passed as the request body. tableId: Table ID of the table to update """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) table = _messages.MessageField('Table', 3) tableId = _messages.StringField(4, required=True) class BigqueryTablesSetIamPolicyRequest(_messages.Message): r"""A BigqueryTablesSetIamPolicyRequest object. Fields: resource: REQUIRED: The resource for which the policy is being specified. See the operation documentation for the appropriate value for this field. setIamPolicyRequest: A SetIamPolicyRequest resource to be passed as the request body. """ resource = _messages.StringField(1, required=True) setIamPolicyRequest = _messages.MessageField('SetIamPolicyRequest', 2) class BigqueryTablesTestIamPermissionsRequest(_messages.Message): r"""A BigqueryTablesTestIamPermissionsRequest object. Fields: resource: REQUIRED: The resource for which the policy detail is being requested. See the operation documentation for the appropriate value for this field. testIamPermissionsRequest: A TestIamPermissionsRequest resource to be passed as the request body. """ resource = _messages.StringField(1, required=True) testIamPermissionsRequest = _messages.MessageField( 'TestIamPermissionsRequest', 2) class BigqueryTablesUpdateRequest(_messages.Message): r"""A BigqueryTablesUpdateRequest object. Fields: datasetId: Dataset ID of the table to update projectId: Project ID of the table to update table: A Table resource to be passed as the request body. tableId: Table ID of the table to update """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) table = _messages.MessageField('Table', 3) tableId = _messages.StringField(4, required=True) class BigtableColumn(_messages.Message): r"""A BigtableColumn object. Fields: encoding: [Optional] The encoding of the values when the type is not STRING. Acceptable encoding values are: TEXT - indicates values are alphanumeric text strings. BINARY - indicates values are encoded using HBase Bytes.toBytes family of functions. 'encoding' can also be set at the column family level. However, the setting at this level takes precedence if 'encoding' is set at both levels. fieldName: [Optional] If the qualifier is not a valid BigQuery field identifier i.e. does not match [a-zA-Z][a-zA-Z0-9_]*, a valid identifier must be provided as the column field name and is used as field name in queries. onlyReadLatest: [Optional] If this is set, only the latest version of value in this column are exposed. 'onlyReadLatest' can also be set at the column family level. However, the setting at this level takes precedence if 'onlyReadLatest' is set at both levels. qualifierEncoded: [Required] Qualifier of the column. Columns in the parent column family that has this exact qualifier are exposed as . field. If the qualifier is valid UTF-8 string, it can be specified in the qualifier_string field. Otherwise, a base-64 encoded value must be set to qualifier_encoded. The column field name is the same as the column qualifier. However, if the qualifier is not a valid BigQuery field identifier i.e. does not match [a-zA-Z][a-zA-Z0-9_]*, a valid identifier must be provided as field_name. qualifierString: A string attribute. type: [Optional] The type to convert the value in cells of this column. The values are expected to be encoded using HBase Bytes.toBytes function when using the BINARY encoding value. Following BigQuery types are allowed (case-sensitive) - BYTES STRING INTEGER FLOAT BOOLEAN Default type is BYTES. 'type' can also be set at the column family level. However, the setting at this level takes precedence if 'type' is set at both levels. """ encoding = _messages.StringField(1) fieldName = _messages.StringField(2) onlyReadLatest = _messages.BooleanField(3) qualifierEncoded = _messages.BytesField(4) qualifierString = _messages.StringField(5) type = _messages.StringField(6) class BigtableColumnFamily(_messages.Message): r"""A BigtableColumnFamily object. Fields: columns: [Optional] Lists of columns that should be exposed as individual fields as opposed to a list of (column name, value) pairs. All columns whose qualifier matches a qualifier in this list can be accessed as .. Other columns can be accessed as a list through .Column field. encoding: [Optional] The encoding of the values when the type is not STRING. Acceptable encoding values are: TEXT - indicates values are alphanumeric text strings. BINARY - indicates values are encoded using HBase Bytes.toBytes family of functions. This can be overridden for a specific column by listing that column in 'columns' and specifying an encoding for it. familyId: Identifier of the column family. onlyReadLatest: [Optional] If this is set only the latest version of value are exposed for all columns in this column family. This can be overridden for a specific column by listing that column in 'columns' and specifying a different setting for that column. type: [Optional] The type to convert the value in cells of this column family. The values are expected to be encoded using HBase Bytes.toBytes function when using the BINARY encoding value. Following BigQuery types are allowed (case-sensitive) - BYTES STRING INTEGER FLOAT BOOLEAN Default type is BYTES. This can be overridden for a specific column by listing that column in 'columns' and specifying a type for it. """ columns = _messages.MessageField('BigtableColumn', 1, repeated=True) encoding = _messages.StringField(2) familyId = _messages.StringField(3) onlyReadLatest = _messages.BooleanField(4) type = _messages.StringField(5) class BigtableOptions(_messages.Message): r"""A BigtableOptions object. Fields: columnFamilies: [Optional] List of column families to expose in the table schema along with their types. This list restricts the column families that can be referenced in queries and specifies their value types. You can use this list to do type conversions - see the 'type' field for more details. If you leave this list empty, all column families are present in the table schema and their values are read as BYTES. During a query only the column families referenced in that query are read from Bigtable. ignoreUnspecifiedColumnFamilies: [Optional] If field is true, then the column families that are not specified in columnFamilies list are not exposed in the table schema. Otherwise, they are read with BYTES type values. The default value is false. readRowkeyAsString: [Optional] If field is true, then the rowkey column families will be read and converted to string. Otherwise they are read with BYTES type values and users need to manually cast them with CAST if necessary. The default value is false. """ columnFamilies = _messages.MessageField( 'BigtableColumnFamily', 1, repeated=True) ignoreUnspecifiedColumnFamilies = _messages.BooleanField(2) readRowkeyAsString = _messages.BooleanField(3) class BinaryClassificationMetrics(_messages.Message): r"""Evaluation metrics for binary classification/classifier models. Fields: aggregateClassificationMetrics: Aggregate classification metrics. binaryConfusionMatrixList: Binary confusion matrix at multiple thresholds. negativeLabel: Label representing the negative class. positiveLabel: Label representing the positive class. """ aggregateClassificationMetrics = _messages.MessageField( 'AggregateClassificationMetrics', 1) binaryConfusionMatrixList = _messages.MessageField( 'BinaryConfusionMatrix', 2, repeated=True) negativeLabel = _messages.StringField(3) positiveLabel = _messages.StringField(4) class BinaryConfusionMatrix(_messages.Message): r"""Confusion matrix for binary classification models. Fields: accuracy: The fraction of predictions given the correct label. f1Score: The equally weighted average of recall and precision. falseNegatives: Number of false samples predicted as false. falsePositives: Number of false samples predicted as true. positiveClassThreshold: Threshold value used when computing each of the following metric. precision: The fraction of actual positive predictions that had positive actual labels. recall: The fraction of actual positive labels that were given a positive prediction. trueNegatives: Number of true samples predicted as false. truePositives: Number of true samples predicted as true. """ accuracy = _messages.FloatField(1) f1Score = _messages.FloatField(2) falseNegatives = _messages.IntegerField(3) falsePositives = _messages.IntegerField(4) positiveClassThreshold = _messages.FloatField(5) precision = _messages.FloatField(6) recall = _messages.FloatField(7) trueNegatives = _messages.IntegerField(8) truePositives = _messages.IntegerField(9) class Binding(_messages.Message): r"""Associates `members` with a `role`. Fields: condition: The condition that is associated with this binding. If the condition evaluates to `true`, then this binding applies to the current request. If the condition evaluates to `false`, then this binding does not apply to the current request. However, a different role binding might grant the same role to one or more of the members in this binding. To learn which resources support conditions in their IAM policies, see the [IAM documentation](https://cloud.google.com/iam/help/conditions/resource- policies). members: Specifies the identities requesting access for a Cloud Platform resource. `members` can have the following values: * `allUsers`: A special identifier that represents anyone who is on the internet; with or without a Google account. * `allAuthenticatedUsers`: A special identifier that represents anyone who is authenticated with a Google account or a service account. * `user:{emailid}`: An email address that represents a specific Google account. For example, `alice@example.com` . * `serviceAccount:{emailid}`: An email address that represents a service account. For example, `my-other-app@appspot.gserviceaccount.com`. * `group:{emailid}`: An email address that represents a Google group. For example, `admins@example.com`. * `deleted:user:{emailid}?uid={uniqueid}`: An email address (plus unique identifier) representing a user that has been recently deleted. For example, `alice@example.com?uid=123456789012345678901`. If the user is recovered, this value reverts to `user:{emailid}` and the recovered user retains the role in the binding. * `deleted:serviceAccount:{emailid}?uid={uniqueid}`: An email address (plus unique identifier) representing a service account that has been recently deleted. For example, `my-other- app@appspot.gserviceaccount.com?uid=123456789012345678901`. If the service account is undeleted, this value reverts to `serviceAccount:{emailid}` and the undeleted service account retains the role in the binding. * `deleted:group:{emailid}?uid={uniqueid}`: An email address (plus unique identifier) representing a Google group that has been recently deleted. For example, `admins@example.com?uid=123456789012345678901`. If the group is recovered, this value reverts to `group:{emailid}` and the recovered group retains the role in the binding. * `domain:{domain}`: The G Suite domain (primary) that represents all the users of that domain. For example, `google.com` or `example.com`. role: Role that is assigned to `members`. For example, `roles/viewer`, `roles/editor`, or `roles/owner`. """ condition = _messages.MessageField('Expr', 1) members = _messages.StringField(2, repeated=True) role = _messages.StringField(3) class BqmlIterationResult(_messages.Message): r"""A BqmlIterationResult object. Fields: durationMs: [Output-only, Beta] Time taken to run the training iteration in milliseconds. evalLoss: [Output-only, Beta] Eval loss computed on the eval data at the end of the iteration. The eval loss is used for early stopping to avoid overfitting. No eval loss if eval_split_method option is specified as no_split or auto_split with input data size less than 500 rows. index: [Output-only, Beta] Index of the ML training iteration, starting from zero for each training run. learnRate: [Output-only, Beta] Learning rate used for this iteration, it varies for different training iterations if learn_rate_strategy option is not constant. trainingLoss: [Output-only, Beta] Training loss computed on the training data at the end of the iteration. The training loss function is defined by model type. """ durationMs = _messages.IntegerField(1) evalLoss = _messages.FloatField(2) index = _messages.IntegerField(3, variant=_messages.Variant.INT32) learnRate = _messages.FloatField(4) trainingLoss = _messages.FloatField(5) class BqmlTrainingRun(_messages.Message): r"""A BqmlTrainingRun object. Messages: TrainingOptionsValue: [Output-only, Beta] Training options used by this training run. These options are mutable for subsequent training runs. Default values are explicitly stored for options not specified in the input query of the first training run. For subsequent training runs, any option not explicitly specified in the input query will be copied from the previous training run. Fields: iterationResults: [Output-only, Beta] List of each iteration results. startTime: [Output-only, Beta] Training run start time in milliseconds since the epoch. state: [Output-only, Beta] Different state applicable for a training run. IN PROGRESS: Training run is in progress. FAILED: Training run ended due to a non-retryable failure. SUCCEEDED: Training run successfully completed. CANCELLED: Training run cancelled by the user. trainingOptions: [Output-only, Beta] Training options used by this training run. These options are mutable for subsequent training runs. Default values are explicitly stored for options not specified in the input query of the first training run. For subsequent training runs, any option not explicitly specified in the input query will be copied from the previous training run. """ class TrainingOptionsValue(_messages.Message): r"""[Output-only, Beta] Training options used by this training run. These options are mutable for subsequent training runs. Default values are explicitly stored for options not specified in the input query of the first training run. For subsequent training runs, any option not explicitly specified in the input query will be copied from the previous training run. Fields: earlyStop: A boolean attribute. l1Reg: A number attribute. l2Reg: A number attribute. learnRate: A number attribute. learnRateStrategy: A string attribute. lineSearchInitLearnRate: A number attribute. maxIteration: A string attribute. minRelProgress: A number attribute. warmStart: A boolean attribute. """ earlyStop = _messages.BooleanField(1) l1Reg = _messages.FloatField(2) l2Reg = _messages.FloatField(3) learnRate = _messages.FloatField(4) learnRateStrategy = _messages.StringField(5) lineSearchInitLearnRate = _messages.FloatField(6) maxIteration = _messages.IntegerField(7) minRelProgress = _messages.FloatField(8) warmStart = _messages.BooleanField(9) iterationResults = _messages.MessageField( 'BqmlIterationResult', 1, repeated=True) startTime = _message_types.DateTimeField(2) state = _messages.StringField(3) trainingOptions = _messages.MessageField('TrainingOptionsValue', 4) class CategoricalValue(_messages.Message): r"""Representative value of a categorical feature. Fields: categoryCounts: Counts of all categories for the categorical feature. If there are more than ten categories, we return top ten (by count) and return one more CategoryCount with category "_OTHER_" and count as aggregate counts of remaining categories. """ categoryCounts = _messages.MessageField('CategoryCount', 1, repeated=True) class CategoryCount(_messages.Message): r"""Represents the count of a single category within the cluster. Fields: category: The name of category. count: The count of training samples matching the category within the cluster. """ category = _messages.StringField(1) count = _messages.IntegerField(2) class Cluster(_messages.Message): r"""Message containing the information about one cluster. Fields: centroidId: Centroid id. count: Count of training data rows that were assigned to this cluster. featureValues: Values of highly variant features for this cluster. """ centroidId = _messages.IntegerField(1) count = _messages.IntegerField(2) featureValues = _messages.MessageField('FeatureValue', 3, repeated=True) class ClusterInfo(_messages.Message): r"""Information about a single cluster for clustering model. Fields: centroidId: Centroid id. clusterRadius: Cluster radius, the average distance from centroid to each point assigned to the cluster. clusterSize: Cluster size, the total number of points assigned to the cluster. """ centroidId = _messages.IntegerField(1) clusterRadius = _messages.FloatField(2) clusterSize = _messages.IntegerField(3) class Clustering(_messages.Message): r"""A Clustering object. Fields: fields: [Repeated] One or more fields on which data should be clustered. Only top-level, non-repeated, simple-type fields are supported. When you cluster a table using multiple columns, the order of columns you specify is important. The order of the specified columns determines the sort order of the data. """ fields = _messages.StringField(1, repeated=True) class ClusteringMetrics(_messages.Message): r"""Evaluation metrics for clustering models. Fields: clusters: [Beta] Information for all clusters. daviesBouldinIndex: Davies-Bouldin index. meanSquaredDistance: Mean of squared distances between each sample to its cluster centroid. """ clusters = _messages.MessageField('Cluster', 1, repeated=True) daviesBouldinIndex = _messages.FloatField(2) meanSquaredDistance = _messages.FloatField(3) class ConfusionMatrix(_messages.Message): r"""Confusion matrix for multi-class classification models. Fields: confidenceThreshold: Confidence threshold used when computing the entries of the confusion matrix. rows: One row per actual label. """ confidenceThreshold = _messages.FloatField(1) rows = _messages.MessageField('Row', 2, repeated=True) class ConnectionProperty(_messages.Message): r"""A ConnectionProperty object. Fields: key: [Required] Name of the connection property to set. value: [Required] Value of the connection property. """ key = _messages.StringField(1) value = _messages.StringField(2) class CsvOptions(_messages.Message): r"""A CsvOptions object. Fields: allowJaggedRows: [Optional] Indicates if BigQuery should accept rows that are missing trailing optional columns. If true, BigQuery treats missing trailing columns as null values. If false, records with missing trailing columns are treated as bad records, and if there are too many bad records, an invalid error is returned in the job result. The default value is false. allowQuotedNewlines: [Optional] Indicates if BigQuery should allow quoted data sections that contain newline characters in a CSV file. The default value is false. encoding: [Optional] The character encoding of the data. The supported values are UTF-8 or ISO-8859-1. The default value is UTF-8. BigQuery decodes the data after the raw, binary data has been split using the values of the quote and fieldDelimiter properties. fieldDelimiter: [Optional] The separator for fields in a CSV file. BigQuery converts the string to ISO-8859-1 encoding, and then uses the first byte of the encoded string to split the data in its raw, binary state. BigQuery also supports the escape sequence "\t" to specify a tab separator. The default value is a comma (','). quote: [Optional] The value that is used to quote data sections in a CSV file. BigQuery converts the string to ISO-8859-1 encoding, and then uses the first byte of the encoded string to split the data in its raw, binary state. The default value is a double-quote ('"'). If your data does not contain quoted sections, set the property value to an empty string. If your data contains quoted newline characters, you must also set the allowQuotedNewlines property to true. skipLeadingRows: [Optional] The number of rows at the top of a CSV file that BigQuery will skip when reading the data. The default value is 0. This property is useful if you have header rows in the file that should be skipped. When autodetect is on, the behavior is the following: * skipLeadingRows unspecified - Autodetect tries to detect headers in the first row. If they are not detected, the row is read as data. Otherwise data is read starting from the second row. * skipLeadingRows is 0 - Instructs autodetect that there are no headers and data should be read starting from the first row. * skipLeadingRows = N > 0 - Autodetect skips N-1 rows and tries to detect headers in row N. If headers are not detected, row N is just skipped. Otherwise row N is used to extract column names for the detected schema. """ allowJaggedRows = _messages.BooleanField(1) allowQuotedNewlines = _messages.BooleanField(2) encoding = _messages.StringField(3) fieldDelimiter = _messages.StringField(4) quote = _messages.StringField(5, default='"') skipLeadingRows = _messages.IntegerField(6) class DataSplitResult(_messages.Message): r"""Data split result. This contains references to the training and evaluation data tables that were used to train the model. Fields: evaluationTable: Table reference of the evaluation data after split. trainingTable: Table reference of the training data after split. """ evaluationTable = _messages.MessageField('TableReference', 1) trainingTable = _messages.MessageField('TableReference', 2) class Dataset(_messages.Message): r"""A Dataset object. Messages: AccessValueListEntry: A AccessValueListEntry object. LabelsValue: The labels associated with this dataset. You can use these to organize and group your datasets. You can set this property when inserting or updating a dataset. See Creating and Updating Dataset Labels for more information. Fields: access: [Optional] An array of objects that define dataset access for one or more entities. You can set this property when inserting or updating a dataset in order to control who is allowed to access the data. If unspecified at dataset creation time, BigQuery adds default dataset access for the following entities: access.specialGroup: projectReaders; access.role: READER; access.specialGroup: projectWriters; access.role: WRITER; access.specialGroup: projectOwners; access.role: OWNER; access.userByEmail: [dataset creator email]; access.role: OWNER; creationTime: [Output-only] The time when this dataset was created, in milliseconds since the epoch. datasetReference: [Required] A reference that identifies the dataset. defaultEncryptionConfiguration: A EncryptionConfiguration attribute. defaultPartitionExpirationMs: [Optional] The default partition expiration for all partitioned tables in the dataset, in milliseconds. Once this property is set, all newly-created partitioned tables in the dataset will have an expirationMs property in the timePartitioning settings set to this value, and changing the value will only affect new tables, not existing ones. The storage in a partition will have an expiration time of its partition time plus this value. Setting this property overrides the use of defaultTableExpirationMs for partitioned tables: only one of defaultTableExpirationMs and defaultPartitionExpirationMs will be used for any new partitioned table. If you provide an explicit timePartitioning.expirationMs when creating or updating a partitioned table, that value takes precedence over the default partition expiration time indicated by this property. defaultTableExpirationMs: [Optional] The default lifetime of all tables in the dataset, in milliseconds. The minimum value is 3600000 milliseconds (one hour). Once this property is set, all newly-created tables in the dataset will have an expirationTime property set to the creation time plus the value in this property, and changing the value will only affect new tables, not existing ones. When the expirationTime for a given table is reached, that table will be deleted automatically. If a table's expirationTime is modified or removed before the table expires, or if you provide an explicit expirationTime when creating a table, that value takes precedence over the default expiration time indicated by this property. description: [Optional] A user-friendly description of the dataset. etag: [Output-only] A hash of the resource. friendlyName: [Optional] A descriptive name for the dataset. id: [Output-only] The fully-qualified unique name of the dataset in the format projectId:datasetId. The dataset name without the project name is given in the datasetId field. When creating a new dataset, leave this field blank, and instead specify the datasetId field. kind: [Output-only] The resource type. labels: The labels associated with this dataset. You can use these to organize and group your datasets. You can set this property when inserting or updating a dataset. See Creating and Updating Dataset Labels for more information. lastModifiedTime: [Output-only] The date when this dataset or any of its tables was last modified, in milliseconds since the epoch. location: The geographic location where the dataset should reside. The default value is US. See details at https://cloud.google.com/bigquery/docs/locations. satisfiesPZS: [Output-only] Reserved for future use. selfLink: [Output-only] A URL that can be used to access the resource again. You can use this URL in Get or Update requests to the resource. """ class AccessValueListEntry(_messages.Message): r"""A AccessValueListEntry object. Fields: domain: [Pick one] A domain to grant access to. Any users signed in with the domain specified will be granted the specified access. Example: "example.com". Maps to IAM policy member "domain:DOMAIN". groupByEmail: [Pick one] An email address of a Google Group to grant access to. Maps to IAM policy member "group:GROUP". iamMember: [Pick one] Some other type of member that appears in the IAM Policy but isn't a user, group, domain, or special group. role: [Required] An IAM role ID that should be granted to the user, group, or domain specified in this access entry. The following legacy mappings will be applied: OWNER roles/bigquery.dataOwner WRITER roles/bigquery.dataEditor READER roles/bigquery.dataViewer This field will accept any of the above formats, but will return only the legacy format. For example, if you set this field to "roles/bigquery.dataOwner", it will be returned back as "OWNER". routine: [Pick one] A routine from a different dataset to grant access to. Queries executed against that routine will have read access to views/tables/routines in this dataset. Only UDF is supported for now. The role field is not required when this field is set. If that routine is updated by any user, access to the routine needs to be granted again via an update operation. specialGroup: [Pick one] A special group to grant access to. Possible values include: projectOwners: Owners of the enclosing project. projectReaders: Readers of the enclosing project. projectWriters: Writers of the enclosing project. allAuthenticatedUsers: All authenticated BigQuery users. Maps to similarly-named IAM members. userByEmail: [Pick one] An email address of a user to grant access to. For example: fred@example.com. Maps to IAM policy member "user:EMAIL" or "serviceAccount:EMAIL". view: [Pick one] A view from a different dataset to grant access to. Queries executed against that view will have read access to tables in this dataset. The role field is not required when this field is set. If that view is updated by any user, access to the view needs to be granted again via an update operation. """ domain = _messages.StringField(1) groupByEmail = _messages.StringField(2) iamMember = _messages.StringField(3) role = _messages.StringField(4) routine = _messages.MessageField('RoutineReference', 5) specialGroup = _messages.StringField(6) userByEmail = _messages.StringField(7) view = _messages.MessageField('TableReference', 8) @encoding.MapUnrecognizedFields('additionalProperties') class LabelsValue(_messages.Message): r"""The labels associated with this dataset. You can use these to organize and group your datasets. You can set this property when inserting or updating a dataset. See Creating and Updating Dataset Labels for more information. Messages: AdditionalProperty: An additional property for a LabelsValue object. Fields: additionalProperties: Additional properties of type LabelsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a LabelsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField( 'AdditionalProperty', 1, repeated=True) access = _messages.MessageField('AccessValueListEntry', 1, repeated=True) creationTime = _messages.IntegerField(2) datasetReference = _messages.MessageField('DatasetReference', 3) defaultEncryptionConfiguration = _messages.MessageField( 'EncryptionConfiguration', 4) defaultPartitionExpirationMs = _messages.IntegerField(5) defaultTableExpirationMs = _messages.IntegerField(6) description = _messages.StringField(7) etag = _messages.StringField(8) friendlyName = _messages.StringField(9) id = _messages.StringField(10) kind = _messages.StringField(11, default='bigquery#dataset') labels = _messages.MessageField('LabelsValue', 12) lastModifiedTime = _messages.IntegerField(13) location = _messages.StringField(14) satisfiesPZS = _messages.BooleanField(15) selfLink = _messages.StringField(16) class DatasetList(_messages.Message): r"""A DatasetList object. Messages: DatasetsValueListEntry: A DatasetsValueListEntry object. Fields: datasets: An array of the dataset resources in the project. Each resource contains basic information. For full information about a particular dataset resource, use the Datasets: get method. This property is omitted when there are no datasets in the project. etag: A hash value of the results page. You can use this property to determine if the page has changed since the last request. kind: The list type. This property always returns the value "bigquery#datasetList". nextPageToken: A token that can be used to request the next results page. This property is omitted on the final results page. """ class DatasetsValueListEntry(_messages.Message): r"""A DatasetsValueListEntry object. Messages: LabelsValue: The labels associated with this dataset. You can use these to organize and group your datasets. Fields: datasetReference: The dataset reference. Use this property to access specific parts of the dataset's ID, such as project ID or dataset ID. friendlyName: A descriptive name for the dataset, if one exists. id: The fully-qualified, unique, opaque ID of the dataset. kind: The resource type. This property always returns the value "bigquery#dataset". labels: The labels associated with this dataset. You can use these to organize and group your datasets. location: The geographic location where the data resides. """ @encoding.MapUnrecognizedFields('additionalProperties') class LabelsValue(_messages.Message): r"""The labels associated with this dataset. You can use these to organize and group your datasets. Messages: AdditionalProperty: An additional property for a LabelsValue object. Fields: additionalProperties: Additional properties of type LabelsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a LabelsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField( 'AdditionalProperty', 1, repeated=True) datasetReference = _messages.MessageField('DatasetReference', 1) friendlyName = _messages.StringField(2) id = _messages.StringField(3) kind = _messages.StringField(4, default='bigquery#dataset') labels = _messages.MessageField('LabelsValue', 5) location = _messages.StringField(6) datasets = _messages.MessageField('DatasetsValueListEntry', 1, repeated=True) etag = _messages.StringField(2) kind = _messages.StringField(3, default='bigquery#datasetList') nextPageToken = _messages.StringField(4) class DatasetReference(_messages.Message): r"""A DatasetReference object. Fields: datasetId: [Required] A unique ID for this dataset, without the project name. The ID must contain only letters (a-z, A-Z), numbers (0-9), or underscores (_). The maximum length is 1,024 characters. projectId: [Optional] The ID of the project containing this dataset. """ datasetId = _messages.StringField(1) projectId = _messages.StringField(2) class DestinationTableProperties(_messages.Message): r"""A DestinationTableProperties object. Messages: LabelsValue: [Optional] The labels associated with this table. You can use these to organize and group your tables. This will only be used if the destination table is newly created. If the table already exists and labels are different than the current labels are provided, the job will fail. Fields: description: [Optional] The description for the destination table. This will only be used if the destination table is newly created. If the table already exists and a value different than the current description is provided, the job will fail. friendlyName: [Optional] The friendly name for the destination table. This will only be used if the destination table is newly created. If the table already exists and a value different than the current friendly name is provided, the job will fail. labels: [Optional] The labels associated with this table. You can use these to organize and group your tables. This will only be used if the destination table is newly created. If the table already exists and labels are different than the current labels are provided, the job will fail. """ @encoding.MapUnrecognizedFields('additionalProperties') class LabelsValue(_messages.Message): r"""[Optional] The labels associated with this table. You can use these to organize and group your tables. This will only be used if the destination table is newly created. If the table already exists and labels are different than the current labels are provided, the job will fail. Messages: AdditionalProperty: An additional property for a LabelsValue object. Fields: additionalProperties: Additional properties of type LabelsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a LabelsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField( 'AdditionalProperty', 1, repeated=True) description = _messages.StringField(1) friendlyName = _messages.StringField(2) labels = _messages.MessageField('LabelsValue', 3) class EncryptionConfiguration(_messages.Message): r"""A EncryptionConfiguration object. Fields: kmsKeyName: [Optional] Describes the Cloud KMS encryption key that will be used to protect destination BigQuery table. The BigQuery Service Account associated with your project requires access to this encryption key. """ kmsKeyName = _messages.StringField(1) class Entry(_messages.Message): r"""A single entry in the confusion matrix. Fields: itemCount: Number of items being predicted as this label. predictedLabel: The predicted label. For confidence_threshold > 0, we will also add an entry indicating the number of items under the confidence threshold. """ itemCount = _messages.IntegerField(1) predictedLabel = _messages.StringField(2) class ErrorProto(_messages.Message): r"""A ErrorProto object. Fields: debugInfo: Debugging information. This property is internal to Google and should not be used. location: Specifies where the error occurred, if present. message: A human-readable description of the error. reason: A short error code that summarizes the error. """ debugInfo = _messages.StringField(1) location = _messages.StringField(2) message = _messages.StringField(3) reason = _messages.StringField(4) class EvaluationMetrics(_messages.Message): r"""Evaluation metrics of a model. These are either computed on all training data or just the eval data based on whether eval data was used during training. These are not present for imported models. Fields: arimaForecastingMetrics: Populated for ARIMA models. binaryClassificationMetrics: Populated for binary classification/classifier models. clusteringMetrics: Populated for clustering models. multiClassClassificationMetrics: Populated for multi-class classification/classifier models. rankingMetrics: Populated for implicit feedback type matrix factorization models. regressionMetrics: Populated for regression models and explicit feedback type matrix factorization models. """ arimaForecastingMetrics = _messages.MessageField('ArimaForecastingMetrics', 1) binaryClassificationMetrics = _messages.MessageField( 'BinaryClassificationMetrics', 2) clusteringMetrics = _messages.MessageField('ClusteringMetrics', 3) multiClassClassificationMetrics = _messages.MessageField( 'MultiClassClassificationMetrics', 4) rankingMetrics = _messages.MessageField('RankingMetrics', 5) regressionMetrics = _messages.MessageField('RegressionMetrics', 6) class ExplainQueryStage(_messages.Message): r"""A ExplainQueryStage object. Fields: completedParallelInputs: Number of parallel input segments completed. computeMsAvg: Milliseconds the average shard spent on CPU-bound tasks. computeMsMax: Milliseconds the slowest shard spent on CPU-bound tasks. computeRatioAvg: Relative amount of time the average shard spent on CPU- bound tasks. computeRatioMax: Relative amount of time the slowest shard spent on CPU- bound tasks. endMs: Stage end time represented as milliseconds since epoch. id: Unique ID for stage within plan. inputStages: IDs for stages that are inputs to this stage. name: Human-readable name for stage. parallelInputs: Number of parallel input segments to be processed. readMsAvg: Milliseconds the average shard spent reading input. readMsMax: Milliseconds the slowest shard spent reading input. readRatioAvg: Relative amount of time the average shard spent reading input. readRatioMax: Relative amount of time the slowest shard spent reading input. recordsRead: Number of records read into the stage. recordsWritten: Number of records written by the stage. shuffleOutputBytes: Total number of bytes written to shuffle. shuffleOutputBytesSpilled: Total number of bytes written to shuffle and spilled to disk. slotMs: Slot-milliseconds used by the stage. startMs: Stage start time represented as milliseconds since epoch. status: Current status for the stage. steps: List of operations within the stage in dependency order (approximately chronological). waitMsAvg: Milliseconds the average shard spent waiting to be scheduled. waitMsMax: Milliseconds the slowest shard spent waiting to be scheduled. waitRatioAvg: Relative amount of time the average shard spent waiting to be scheduled. waitRatioMax: Relative amount of time the slowest shard spent waiting to be scheduled. writeMsAvg: Milliseconds the average shard spent on writing output. writeMsMax: Milliseconds the slowest shard spent on writing output. writeRatioAvg: Relative amount of time the average shard spent on writing output. writeRatioMax: Relative amount of time the slowest shard spent on writing output. """ completedParallelInputs = _messages.IntegerField(1) computeMsAvg = _messages.IntegerField(2) computeMsMax = _messages.IntegerField(3) computeRatioAvg = _messages.FloatField(4) computeRatioMax = _messages.FloatField(5) endMs = _messages.IntegerField(6) id = _messages.IntegerField(7) inputStages = _messages.IntegerField(8, repeated=True) name = _messages.StringField(9) parallelInputs = _messages.IntegerField(10) readMsAvg = _messages.IntegerField(11) readMsMax = _messages.IntegerField(12) readRatioAvg = _messages.FloatField(13) readRatioMax = _messages.FloatField(14) recordsRead = _messages.IntegerField(15) recordsWritten = _messages.IntegerField(16) shuffleOutputBytes = _messages.IntegerField(17) shuffleOutputBytesSpilled = _messages.IntegerField(18) slotMs = _messages.IntegerField(19) startMs = _messages.IntegerField(20) status = _messages.StringField(21) steps = _messages.MessageField('ExplainQueryStep', 22, repeated=True) waitMsAvg = _messages.IntegerField(23) waitMsMax = _messages.IntegerField(24) waitRatioAvg = _messages.FloatField(25) waitRatioMax = _messages.FloatField(26) writeMsAvg = _messages.IntegerField(27) writeMsMax = _messages.IntegerField(28) writeRatioAvg = _messages.FloatField(29) writeRatioMax = _messages.FloatField(30) class ExplainQueryStep(_messages.Message): r"""A ExplainQueryStep object. Fields: kind: Machine-readable operation type. substeps: Human-readable stage descriptions. """ kind = _messages.StringField(1) substeps = _messages.StringField(2, repeated=True) class Explanation(_messages.Message): r"""Explanation for a single feature. Fields: attribution: Attribution of feature. featureName: Full name of the feature. For non-numerical features, will be formatted like .. Overall size of feature name will always be truncated to first 120 characters. """ attribution = _messages.FloatField(1) featureName = _messages.StringField(2) class Expr(_messages.Message): r"""Represents a textual expression in the Common Expression Language (CEL) syntax. CEL is a C-like expression language. The syntax and semantics of CEL are documented at https://github.com/google/cel-spec. Example (Comparison): title: "Summary size limit" description: "Determines if a summary is less than 100 chars" expression: "document.summary.size() < 100" Example (Equality): title: "Requestor is owner" description: "Determines if requestor is the document owner" expression: "document.owner == request.auth.claims.email" Example (Logic): title: "Public documents" description: "Determine whether the document should be publicly visible" expression: "document.type != 'private' && document.type != 'internal'" Example (Data Manipulation): title: "Notification string" description: "Create a notification string with a timestamp." expression: "'New message received at ' + string(document.create_time)" The exact variables and functions that may be referenced within an expression are determined by the service that evaluates it. See the service documentation for additional information. Fields: description: Optional. Description of the expression. This is a longer text which describes the expression, e.g. when hovered over it in a UI. expression: Textual representation of an expression in Common Expression Language syntax. location: Optional. String indicating the location of the expression for error reporting, e.g. a file name and a position in the file. title: Optional. Title for the expression, i.e. a short string describing its purpose. This can be used e.g. in UIs which allow to enter the expression. """ description = _messages.StringField(1) expression = _messages.StringField(2) location = _messages.StringField(3) title = _messages.StringField(4) class ExternalDataConfiguration(_messages.Message): r"""A ExternalDataConfiguration object. Fields: autodetect: Try to detect schema and format options automatically. Any option specified explicitly will be honored. bigtableOptions: [Optional] Additional options if sourceFormat is set to BIGTABLE. compression: [Optional] The compression type of the data source. Possible values include GZIP and NONE. The default value is NONE. This setting is ignored for Google Cloud Bigtable, Google Cloud Datastore backups and Avro formats. connectionId: [Optional, Trusted Tester] Connection for external data source. csvOptions: Additional properties to set if sourceFormat is set to CSV. googleSheetsOptions: [Optional] Additional options if sourceFormat is set to GOOGLE_SHEETS. hivePartitioningOptions: [Optional, Trusted Tester] Options to configure hive partitioning support. ignoreUnknownValues: [Optional] Indicates if BigQuery should allow extra values that are not represented in the table schema. If true, the extra values are ignored. If false, records with extra columns are treated as bad records, and if there are too many bad records, an invalid error is returned in the job result. The default value is false. The sourceFormat property determines what BigQuery treats as an extra value: CSV: Trailing columns JSON: Named values that don't match any column names Google Cloud Bigtable: This setting is ignored. Google Cloud Datastore backups: This setting is ignored. Avro: This setting is ignored. maxBadRecords: [Optional] The maximum number of bad records that BigQuery can ignore when reading data. If the number of bad records exceeds this value, an invalid error is returned in the job result. This is only valid for CSV, JSON, and Google Sheets. The default value is 0, which requires that all records are valid. This setting is ignored for Google Cloud Bigtable, Google Cloud Datastore backups and Avro formats. schema: [Optional] The schema for the data. Schema is required for CSV and JSON formats. Schema is disallowed for Google Cloud Bigtable, Cloud Datastore backups, and Avro formats. sourceFormat: [Required] The data format. For CSV files, specify "CSV". For Google sheets, specify "GOOGLE_SHEETS". For newline-delimited JSON, specify "NEWLINE_DELIMITED_JSON". For Avro files, specify "AVRO". For Google Cloud Datastore backups, specify "DATASTORE_BACKUP". [Beta] For Google Cloud Bigtable, specify "BIGTABLE". sourceUris: [Required] The fully-qualified URIs that point to your data in Google Cloud. For Google Cloud Storage URIs: Each URI can contain one '*' wildcard character and it must come after the 'bucket' name. Size limits related to load jobs apply to external data sources. For Google Cloud Bigtable URIs: Exactly one URI can be specified and it has be a fully specified and valid HTTPS URL for a Google Cloud Bigtable table. For Google Cloud Datastore backups, exactly one URI can be specified. Also, the '*' wildcard character is not allowed. """ autodetect = _messages.BooleanField(1) bigtableOptions = _messages.MessageField('BigtableOptions', 2) compression = _messages.StringField(3) connectionId = _messages.StringField(4) csvOptions = _messages.MessageField('CsvOptions', 5) googleSheetsOptions = _messages.MessageField('GoogleSheetsOptions', 6) hivePartitioningOptions = _messages.MessageField('HivePartitioningOptions', 7) ignoreUnknownValues = _messages.BooleanField(8) maxBadRecords = _messages.IntegerField(9, variant=_messages.Variant.INT32) schema = _messages.MessageField('TableSchema', 10) sourceFormat = _messages.StringField(11) sourceUris = _messages.StringField(12, repeated=True) class FeatureValue(_messages.Message): r"""Representative value of a single feature within the cluster. Fields: categoricalValue: The categorical feature value. featureColumn: The feature column name. numericalValue: The numerical feature value. This is the centroid value for this feature. """ categoricalValue = _messages.MessageField('CategoricalValue', 1) featureColumn = _messages.StringField(2) numericalValue = _messages.FloatField(3) class GetIamPolicyRequest(_messages.Message): r"""Request message for `GetIamPolicy` method. Fields: options: OPTIONAL: A `GetPolicyOptions` object for specifying options to `GetIamPolicy`. """ options = _messages.MessageField('GetPolicyOptions', 1) class GetPolicyOptions(_messages.Message): r"""Encapsulates settings provided to GetIamPolicy. Fields: requestedPolicyVersion: Optional. The policy format version to be returned. Valid values are 0, 1, and 3. Requests specifying an invalid value will be rejected. Requests for policies with any conditional bindings must specify version 3. Policies without any conditional bindings may specify any valid value or leave the field unset. To learn which resources support conditions in their IAM policies, see the [IAM documentation](https://cloud.google.com/iam/help/conditions/resource- policies). """ requestedPolicyVersion = _messages.IntegerField( 1, variant=_messages.Variant.INT32) class GetQueryResultsResponse(_messages.Message): r"""A GetQueryResultsResponse object. Fields: cacheHit: Whether the query result was fetched from the query cache. errors: [Output-only] The first errors or warnings encountered during the running of the job. The final message includes the number of errors that caused the process to stop. Errors here do not necessarily mean that the job has completed or was unsuccessful. etag: A hash of this response. jobComplete: Whether the query has completed or not. If rows or totalRows are present, this will always be true. If this is false, totalRows will not be available. jobReference: Reference to the BigQuery Job that was created to run the query. This field will be present even if the original request timed out, in which case GetQueryResults can be used to read the results once the query has completed. Since this API only returns the first page of results, subsequent pages can be fetched via the same mechanism (GetQueryResults). kind: The resource type of the response. numDmlAffectedRows: [Output-only] The number of rows affected by a DML statement. Present only for DML statements INSERT, UPDATE or DELETE. pageToken: A token used for paging results. rows: An object with as many results as can be contained within the maximum permitted reply size. To get any additional rows, you can call GetQueryResults and specify the jobReference returned above. Present only when the query completes successfully. schema: The schema of the results. Present only when the query completes successfully. totalBytesProcessed: The total number of bytes processed for this query. totalRows: The total number of rows in the complete query result set, which can be more than the number of rows in this single page of results. Present only when the query completes successfully. """ cacheHit = _messages.BooleanField(1) errors = _messages.MessageField('ErrorProto', 2, repeated=True) etag = _messages.StringField(3) jobComplete = _messages.BooleanField(4) jobReference = _messages.MessageField('JobReference', 5) kind = _messages.StringField(6, default='bigquery#getQueryResultsResponse') numDmlAffectedRows = _messages.IntegerField(7) pageToken = _messages.StringField(8) rows = _messages.MessageField('TableRow', 9, repeated=True) schema = _messages.MessageField('TableSchema', 10) totalBytesProcessed = _messages.IntegerField(11) totalRows = _messages.IntegerField(12, variant=_messages.Variant.UINT64) class GetServiceAccountResponse(_messages.Message): r"""A GetServiceAccountResponse object. Fields: email: The service account email address. kind: The resource type of the response. """ email = _messages.StringField(1) kind = _messages.StringField(2, default='bigquery#getServiceAccountResponse') class GlobalExplanation(_messages.Message): r"""Global explanations containing the top most important features after training. Fields: classLabel: Class label for this set of global explanations. Will be empty/null for binary logistic and linear regression models. Sorted alphabetically in descending order. explanations: A list of the top global explanations. Sorted by absolute value of attribution in descending order. """ classLabel = _messages.StringField(1) explanations = _messages.MessageField('Explanation', 2, repeated=True) class GoogleSheetsOptions(_messages.Message): r"""A GoogleSheetsOptions object. Fields: range: [Optional] Range of a sheet to query from. Only used when non- empty. Typical format: sheet_name!top_left_cell_id:bottom_right_cell_id For example: sheet1!A1:B20 skipLeadingRows: [Optional] The number of rows at the top of a sheet that BigQuery will skip when reading the data. The default value is 0. This property is useful if you have header rows that should be skipped. When autodetect is on, behavior is the following: * skipLeadingRows unspecified - Autodetect tries to detect headers in the first row. If they are not detected, the row is read as data. Otherwise data is read starting from the second row. * skipLeadingRows is 0 - Instructs autodetect that there are no headers and data should be read starting from the first row. * skipLeadingRows = N > 0 - Autodetect skips N-1 rows and tries to detect headers in row N. If headers are not detected, row N is just skipped. Otherwise row N is used to extract column names for the detected schema. """ range = _messages.StringField(1) skipLeadingRows = _messages.IntegerField(2) class HivePartitioningOptions(_messages.Message): r"""A HivePartitioningOptions object. Fields: mode: [Optional] When set, what mode of hive partitioning to use when reading data. The following modes are supported. (1) AUTO: automatically infer partition key name(s) and type(s). (2) STRINGS: automatically infer partition key name(s). All types are interpreted as strings. (3) CUSTOM: partition key schema is encoded in the source URI prefix. Not all storage formats support hive partitioning. Requesting hive partitioning on an unsupported format will lead to an error. Currently supported types include: AVRO, CSV, JSON, ORC and Parquet. requirePartitionFilter: [Optional] If set to true, queries over this table require a partition filter that can be used for partition elimination to be specified. Note that this field should only be true when creating a permanent external table or querying a temporary external table. Hive- partitioned loads with requirePartitionFilter explicitly set to true will fail. sourceUriPrefix: [Optional] When hive partition detection is requested, a common prefix for all source uris should be supplied. The prefix must end immediately before the partition key encoding begins. For example, consider files following this data layout. gs://bucket/path_to_table/dt=2019-01-01/country=BR/id=7/file.avro gs://bucket/path_to_table/dt=2018-12-31/country=CA/id=3/file.avro When hive partitioning is requested with either AUTO or STRINGS detection, the common prefix can be either of gs://bucket/path_to_table or gs://bucket/path_to_table/ (trailing slash does not matter). """ mode = _messages.StringField(1) requirePartitionFilter = _messages.BooleanField(2) sourceUriPrefix = _messages.StringField(3) class IterationResult(_messages.Message): r"""Information about a single iteration of the training run. Fields: arimaResult: A ArimaResult attribute. clusterInfos: Information about top clusters for clustering models. durationMs: Time taken to run the iteration in milliseconds. evalLoss: Loss computed on the eval data at the end of iteration. index: Index of the iteration, 0 based. learnRate: Learn rate used for this iteration. trainingLoss: Loss computed on the training data at the end of iteration. """ arimaResult = _messages.MessageField('ArimaResult', 1) clusterInfos = _messages.MessageField('ClusterInfo', 2, repeated=True) durationMs = _messages.IntegerField(3) evalLoss = _messages.FloatField(4) index = _messages.IntegerField(5, variant=_messages.Variant.INT32) learnRate = _messages.FloatField(6) trainingLoss = _messages.FloatField(7) class Job(_messages.Message): r"""A Job object. Fields: configuration: [Required] Describes the job configuration. etag: [Output-only] A hash of this resource. id: [Output-only] Opaque ID field of the job jobReference: [Optional] Reference describing the unique-per-user name of the job. kind: [Output-only] The type of the resource. selfLink: [Output-only] A URL that can be used to access this resource again. statistics: [Output-only] Information about the job, including starting time and ending time of the job. status: [Output-only] The status of this job. Examine this value when polling an asynchronous job to see if the job is complete. user_email: [Output-only] Email address of the user who ran the job. """ configuration = _messages.MessageField('JobConfiguration', 1) etag = _messages.StringField(2) id = _messages.StringField(3) jobReference = _messages.MessageField('JobReference', 4) kind = _messages.StringField(5, default='bigquery#job') selfLink = _messages.StringField(6) statistics = _messages.MessageField('JobStatistics', 7) status = _messages.MessageField('JobStatus', 8) user_email = _messages.StringField(9) class JobCancelResponse(_messages.Message): r"""A JobCancelResponse object. Fields: job: The final state of the job. kind: The resource type of the response. """ job = _messages.MessageField('Job', 1) kind = _messages.StringField(2, default='bigquery#jobCancelResponse') class JobConfiguration(_messages.Message): r"""A JobConfiguration object. Messages: LabelsValue: The labels associated with this job. You can use these to organize and group your jobs. Label keys and values can be no longer than 63 characters, can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. Label values are optional. Label keys must start with a letter and each label in the list must have a different key. Fields: copy: [Pick one] Copies a table. dryRun: [Optional] If set, don't actually run this job. A valid query will return a mostly empty response with some processing statistics, while an invalid query will return the same error it would if it wasn't a dry run. Behavior of non-query jobs is undefined. extract: [Pick one] Configures an extract job. jobTimeoutMs: [Optional] Job timeout in milliseconds. If this time limit is exceeded, BigQuery may attempt to terminate the job. jobType: [Output-only] The type of the job. Can be QUERY, LOAD, EXTRACT, COPY or UNKNOWN. labels: The labels associated with this job. You can use these to organize and group your jobs. Label keys and values can be no longer than 63 characters, can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. Label values are optional. Label keys must start with a letter and each label in the list must have a different key. load: [Pick one] Configures a load job. query: [Pick one] Configures a query job. """ @encoding.MapUnrecognizedFields('additionalProperties') class LabelsValue(_messages.Message): r"""The labels associated with this job. You can use these to organize and group your jobs. Label keys and values can be no longer than 63 characters, can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. Label values are optional. Label keys must start with a letter and each label in the list must have a different key. Messages: AdditionalProperty: An additional property for a LabelsValue object. Fields: additionalProperties: Additional properties of type LabelsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a LabelsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField( 'AdditionalProperty', 1, repeated=True) copy = _messages.MessageField('JobConfigurationTableCopy', 1) dryRun = _messages.BooleanField(2) extract = _messages.MessageField('JobConfigurationExtract', 3) jobTimeoutMs = _messages.IntegerField(4) jobType = _messages.StringField(5) labels = _messages.MessageField('LabelsValue', 6) load = _messages.MessageField('JobConfigurationLoad', 7) query = _messages.MessageField('JobConfigurationQuery', 8) class JobConfigurationExtract(_messages.Message): r"""A JobConfigurationExtract object. Fields: compression: [Optional] The compression type to use for exported files. Possible values include GZIP, DEFLATE, SNAPPY, and NONE. The default value is NONE. DEFLATE and SNAPPY are only supported for Avro. Not applicable when extracting models. destinationFormat: [Optional] The exported file format. Possible values include CSV, NEWLINE_DELIMITED_JSON, PARQUET or AVRO for tables and ML_TF_SAVED_MODEL or ML_XGBOOST_BOOSTER for models. The default value for tables is CSV. Tables with nested or repeated fields cannot be exported as CSV. The default value for models is ML_TF_SAVED_MODEL. destinationUri: [Pick one] DEPRECATED: Use destinationUris instead, passing only one URI as necessary. The fully-qualified Google Cloud Storage URI where the extracted table should be written. destinationUris: [Pick one] A list of fully-qualified Google Cloud Storage URIs where the extracted table should be written. fieldDelimiter: [Optional] Delimiter to use between fields in the exported data. Default is ','. Not applicable when extracting models. printHeader: [Optional] Whether to print out a header row in the results. Default is true. Not applicable when extracting models. sourceModel: A reference to the model being exported. sourceTable: A reference to the table being exported. useAvroLogicalTypes: [Optional] If destinationFormat is set to "AVRO", this flag indicates whether to enable extracting applicable column types (such as TIMESTAMP) to their corresponding AVRO logical types (timestamp-micros), instead of only using their raw types (avro-long). Not applicable when extracting models. """ compression = _messages.StringField(1) destinationFormat = _messages.StringField(2) destinationUri = _messages.StringField(3) destinationUris = _messages.StringField(4, repeated=True) fieldDelimiter = _messages.StringField(5) printHeader = _messages.BooleanField(6, default=True) sourceModel = _messages.MessageField('ModelReference', 7) sourceTable = _messages.MessageField('TableReference', 8) useAvroLogicalTypes = _messages.BooleanField(9) class JobConfigurationLoad(_messages.Message): r"""A JobConfigurationLoad object. Fields: allowJaggedRows: [Optional] Accept rows that are missing trailing optional columns. The missing values are treated as nulls. If false, records with missing trailing columns are treated as bad records, and if there are too many bad records, an invalid error is returned in the job result. The default value is false. Only applicable to CSV, ignored for other formats. allowQuotedNewlines: Indicates if BigQuery should allow quoted data sections that contain newline characters in a CSV file. The default value is false. autodetect: [Optional] Indicates if we should automatically infer the options and schema for CSV and JSON sources. clustering: [Beta] Clustering specification for the destination table. Must be specified with time-based partitioning, data in the table will be first partitioned and subsequently clustered. createDisposition: [Optional] Specifies whether the job is allowed to create new tables. The following values are supported: CREATE_IF_NEEDED: If the table does not exist, BigQuery creates the table. CREATE_NEVER: The table must already exist. If it does not, a 'notFound' error is returned in the job result. The default value is CREATE_IF_NEEDED. Creation, truncation and append actions occur as one atomic update upon job completion. decimalTargetTypes: [Trusted Tester] Defines the list of possible SQL data types to which the source decimal values are converted. This list and the precision and the scale parameters of the decimal field determine the target type. In the order of NUMERIC, BIGNUMERIC, and STRING, a type is picked if it is in the specified list and if it supports the precision and the scale. STRING supports all precision and scale values. If none of the listed types supports the precision and the scale, the type supporting the widest range in the specified list is picked, and if a value exceeds the supported range when reading the data, an error will be thrown. For example: suppose decimal_target_type = ["NUMERIC", "BIGNUMERIC"]. Then if (precision,scale) is: * (38,9) -> NUMERIC; * (39,9) -> BIGNUMERIC (NUMERIC cannot hold 30 integer digits); * (38,10) -> BIGNUMERIC (NUMERIC cannot hold 10 fractional digits); * (76,38) -> BIGNUMERIC; * (77,38) -> BIGNUMERIC (error if value exeeds supported range). For duplicated types in this field, only one will be considered and the rest will be ignored. The order of the types in this field is ignored. For example, ["BIGNUMERIC", "NUMERIC"] is the same as ["NUMERIC", "BIGNUMERIC"] and NUMERIC always takes precedence over BIGNUMERIC. destinationEncryptionConfiguration: Custom encryption configuration (e.g., Cloud KMS keys). destinationTable: [Required] The destination table to load the data into. destinationTableProperties: [Beta] [Optional] Properties with which to create the destination table if it is new. encoding: [Optional] The character encoding of the data. The supported values are UTF-8 or ISO-8859-1. The default value is UTF-8. BigQuery decodes the data after the raw, binary data has been split using the values of the quote and fieldDelimiter properties. fieldDelimiter: [Optional] The separator for fields in a CSV file. The separator can be any ISO-8859-1 single-byte character. To use a character in the range 128-255, you must encode the character as UTF8. BigQuery converts the string to ISO-8859-1 encoding, and then uses the first byte of the encoded string to split the data in its raw, binary state. BigQuery also supports the escape sequence "\t" to specify a tab separator. The default value is a comma (','). hivePartitioningOptions: [Optional, Trusted Tester] Options to configure hive partitioning support. ignoreUnknownValues: [Optional] Indicates if BigQuery should allow extra values that are not represented in the table schema. If true, the extra values are ignored. If false, records with extra columns are treated as bad records, and if there are too many bad records, an invalid error is returned in the job result. The default value is false. The sourceFormat property determines what BigQuery treats as an extra value: CSV: Trailing columns JSON: Named values that don't match any column names maxBadRecords: [Optional] The maximum number of bad records that BigQuery can ignore when running the job. If the number of bad records exceeds this value, an invalid error is returned in the job result. This is only valid for CSV and JSON. The default value is 0, which requires that all records are valid. nullMarker: [Optional] Specifies a string that represents a null value in a CSV file. For example, if you specify "\\N", BigQuery interprets "\\N" as a null value when loading a CSV file. The default value is the empty string. If you set this property to a custom value, BigQuery throws an error if an empty string is present for all data types except for STRING and BYTE. For STRING and BYTE columns, BigQuery interprets the empty string as an empty value. projectionFields: If sourceFormat is set to "DATASTORE_BACKUP", indicates which entity properties to load into BigQuery from a Cloud Datastore backup. Property names are case sensitive and must be top-level properties. If no properties are specified, BigQuery loads all properties. If any named property isn't found in the Cloud Datastore backup, an invalid error is returned in the job result. quote: [Optional] The value that is used to quote data sections in a CSV file. BigQuery converts the string to ISO-8859-1 encoding, and then uses the first byte of the encoded string to split the data in its raw, binary state. The default value is a double-quote ('"'). If your data does not contain quoted sections, set the property value to an empty string. If your data contains quoted newline characters, you must also set the allowQuotedNewlines property to true. rangePartitioning: [TrustedTester] Range partitioning specification for this table. Only one of timePartitioning and rangePartitioning should be specified. schema: [Optional] The schema for the destination table. The schema can be omitted if the destination table already exists, or if you're loading data from Google Cloud Datastore. schemaInline: [Deprecated] The inline schema. For CSV schemas, specify as "Field1:Type1[,Field2:Type2]*". For example, "foo:STRING, bar:INTEGER, baz:FLOAT". schemaInlineFormat: [Deprecated] The format of the schemaInline property. schemaUpdateOptions: Allows the schema of the destination table to be updated as a side effect of the load job if a schema is autodetected or supplied in the job configuration. Schema update options are supported in two cases: when writeDisposition is WRITE_APPEND; when writeDisposition is WRITE_TRUNCATE and the destination table is a partition of a table, specified by partition decorators. For normal tables, WRITE_TRUNCATE will always overwrite the schema. One or more of the following values are specified: ALLOW_FIELD_ADDITION: allow adding a nullable field to the schema. ALLOW_FIELD_RELAXATION: allow relaxing a required field in the original schema to nullable. skipLeadingRows: [Optional] The number of rows at the top of a CSV file that BigQuery will skip when loading the data. The default value is 0. This property is useful if you have header rows in the file that should be skipped. sourceFormat: [Optional] The format of the data files. For CSV files, specify "CSV". For datastore backups, specify "DATASTORE_BACKUP". For newline-delimited JSON, specify "NEWLINE_DELIMITED_JSON". For Avro, specify "AVRO". For parquet, specify "PARQUET". For orc, specify "ORC". The default value is CSV. sourceUris: [Required] The fully-qualified URIs that point to your data in Google Cloud. For Google Cloud Storage URIs: Each URI can contain one '*' wildcard character and it must come after the 'bucket' name. Size limits related to load jobs apply to external data sources. For Google Cloud Bigtable URIs: Exactly one URI can be specified and it has be a fully specified and valid HTTPS URL for a Google Cloud Bigtable table. For Google Cloud Datastore backups: Exactly one URI can be specified. Also, the '*' wildcard character is not allowed. timePartitioning: Time-based partitioning specification for the destination table. Only one of timePartitioning and rangePartitioning should be specified. useAvroLogicalTypes: [Optional] If sourceFormat is set to "AVRO", indicates whether to enable interpreting logical types into their corresponding types (ie. TIMESTAMP), instead of only using their raw types (ie. INTEGER). writeDisposition: [Optional] Specifies the action that occurs if the destination table already exists. The following values are supported: WRITE_TRUNCATE: If the table already exists, BigQuery overwrites the table data. WRITE_APPEND: If the table already exists, BigQuery appends the data to the table. WRITE_EMPTY: If the table already exists and contains data, a 'duplicate' error is returned in the job result. The default value is WRITE_APPEND. Each action is atomic and only occurs if BigQuery is able to complete the job successfully. Creation, truncation and append actions occur as one atomic update upon job completion. """ allowJaggedRows = _messages.BooleanField(1) allowQuotedNewlines = _messages.BooleanField(2) autodetect = _messages.BooleanField(3) clustering = _messages.MessageField('Clustering', 4) createDisposition = _messages.StringField(5) decimalTargetTypes = _messages.StringField(6, repeated=True) destinationEncryptionConfiguration = _messages.MessageField( 'EncryptionConfiguration', 7) destinationTable = _messages.MessageField('TableReference', 8) destinationTableProperties = _messages.MessageField( 'DestinationTableProperties', 9) encoding = _messages.StringField(10) fieldDelimiter = _messages.StringField(11) hivePartitioningOptions = _messages.MessageField( 'HivePartitioningOptions', 12) ignoreUnknownValues = _messages.BooleanField(13) maxBadRecords = _messages.IntegerField(14, variant=_messages.Variant.INT32) nullMarker = _messages.StringField(15) projectionFields = _messages.StringField(16, repeated=True) quote = _messages.StringField(17, default='"') rangePartitioning = _messages.MessageField('RangePartitioning', 18) schema = _messages.MessageField('TableSchema', 19) schemaInline = _messages.StringField(20) schemaInlineFormat = _messages.StringField(21) schemaUpdateOptions = _messages.StringField(22, repeated=True) skipLeadingRows = _messages.IntegerField(23, variant=_messages.Variant.INT32) sourceFormat = _messages.StringField(24) sourceUris = _messages.StringField(25, repeated=True) timePartitioning = _messages.MessageField('TimePartitioning', 26) useAvroLogicalTypes = _messages.BooleanField(27) writeDisposition = _messages.StringField(28) class JobConfigurationQuery(_messages.Message): r"""A JobConfigurationQuery object. Messages: TableDefinitionsValue: [Optional] If querying an external data source outside of BigQuery, describes the data format, location and other properties of the data source. By defining these properties, the data source can then be queried as if it were a standard BigQuery table. Fields: allowLargeResults: [Optional] If true and query uses legacy SQL dialect, allows the query to produce arbitrarily large result tables at a slight cost in performance. Requires destinationTable to be set. For standard SQL queries, this flag is ignored and large results are always allowed. However, you must still set destinationTable when result size exceeds the allowed maximum response size. clustering: [Beta] Clustering specification for the destination table. Must be specified with time-based partitioning, data in the table will be first partitioned and subsequently clustered. connectionProperties: Connection properties. createDisposition: [Optional] Specifies whether the job is allowed to create new tables. The following values are supported: CREATE_IF_NEEDED: If the table does not exist, BigQuery creates the table. CREATE_NEVER: The table must already exist. If it does not, a 'notFound' error is returned in the job result. The default value is CREATE_IF_NEEDED. Creation, truncation and append actions occur as one atomic update upon job completion. defaultDataset: [Optional] Specifies the default dataset to use for unqualified table names in the query. Note that this does not alter behavior of unqualified dataset names. destinationEncryptionConfiguration: Custom encryption configuration (e.g., Cloud KMS keys). destinationTable: [Optional] Describes the table where the query results should be stored. If not present, a new table will be created to store the results. This property must be set for large results that exceed the maximum response size. flattenResults: [Optional] If true and query uses legacy SQL dialect, flattens all nested and repeated fields in the query results. allowLargeResults must be true if this is set to false. For standard SQL queries, this flag is ignored and results are never flattened. maximumBillingTier: [Optional] Limits the billing tier for this job. Queries that have resource usage beyond this tier will fail (without incurring a charge). If unspecified, this will be set to your project default. maximumBytesBilled: [Optional] Limits the bytes billed for this job. Queries that will have bytes billed beyond this limit will fail (without incurring a charge). If unspecified, this will be set to your project default. parameterMode: Standard SQL only. Set to POSITIONAL to use positional (?) query parameters or to NAMED to use named (@myparam) query parameters in this query. preserveNulls: [Deprecated] This property is deprecated. priority: [Optional] Specifies a priority for the query. Possible values include INTERACTIVE and BATCH. The default value is INTERACTIVE. query: [Required] SQL query text to execute. The useLegacySql field can be used to indicate whether the query uses legacy SQL or standard SQL. queryParameters: Query parameters for standard SQL queries. rangePartitioning: [TrustedTester] Range partitioning specification for this table. Only one of timePartitioning and rangePartitioning should be specified. schemaUpdateOptions: Allows the schema of the destination table to be updated as a side effect of the query job. Schema update options are supported in two cases: when writeDisposition is WRITE_APPEND; when writeDisposition is WRITE_TRUNCATE and the destination table is a partition of a table, specified by partition decorators. For normal tables, WRITE_TRUNCATE will always overwrite the schema. One or more of the following values are specified: ALLOW_FIELD_ADDITION: allow adding a nullable field to the schema. ALLOW_FIELD_RELAXATION: allow relaxing a required field in the original schema to nullable. tableDefinitions: [Optional] If querying an external data source outside of BigQuery, describes the data format, location and other properties of the data source. By defining these properties, the data source can then be queried as if it were a standard BigQuery table. timePartitioning: Time-based partitioning specification for the destination table. Only one of timePartitioning and rangePartitioning should be specified. useLegacySql: Specifies whether to use BigQuery's legacy SQL dialect for this query. The default value is true. If set to false, the query will use BigQuery's standard SQL: https://cloud.google.com/bigquery/sql- reference/ When useLegacySql is set to false, the value of flattenResults is ignored; query will be run as if flattenResults is false. useQueryCache: [Optional] Whether to look for the result in the query cache. The query cache is a best-effort cache that will be flushed whenever tables in the query are modified. Moreover, the query cache is only available when a query does not have a destination table specified. The default value is true. userDefinedFunctionResources: Describes user-defined function resources used in the query. writeDisposition: [Optional] Specifies the action that occurs if the destination table already exists. The following values are supported: WRITE_TRUNCATE: If the table already exists, BigQuery overwrites the table data and uses the schema from the query result. WRITE_APPEND: If the table already exists, BigQuery appends the data to the table. WRITE_EMPTY: If the table already exists and contains data, a 'duplicate' error is returned in the job result. The default value is WRITE_EMPTY. Each action is atomic and only occurs if BigQuery is able to complete the job successfully. Creation, truncation and append actions occur as one atomic update upon job completion. """ @encoding.MapUnrecognizedFields('additionalProperties') class TableDefinitionsValue(_messages.Message): r"""[Optional] If querying an external data source outside of BigQuery, describes the data format, location and other properties of the data source. By defining these properties, the data source can then be queried as if it were a standard BigQuery table. Messages: AdditionalProperty: An additional property for a TableDefinitionsValue object. Fields: additionalProperties: Additional properties of type TableDefinitionsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a TableDefinitionsValue object. Fields: key: Name of the additional property. value: A ExternalDataConfiguration attribute. """ key = _messages.StringField(1) value = _messages.MessageField('ExternalDataConfiguration', 2) additionalProperties = _messages.MessageField( 'AdditionalProperty', 1, repeated=True) allowLargeResults = _messages.BooleanField(1, default=False) clustering = _messages.MessageField('Clustering', 2) connectionProperties = _messages.MessageField( 'ConnectionProperty', 3, repeated=True) createDisposition = _messages.StringField(4) defaultDataset = _messages.MessageField('DatasetReference', 5) destinationEncryptionConfiguration = _messages.MessageField( 'EncryptionConfiguration', 6) destinationTable = _messages.MessageField('TableReference', 7) flattenResults = _messages.BooleanField(8, default=True) maximumBillingTier = _messages.IntegerField( 9, variant=_messages.Variant.INT32, default=1) maximumBytesBilled = _messages.IntegerField(10) parameterMode = _messages.StringField(11) preserveNulls = _messages.BooleanField(12) priority = _messages.StringField(13) query = _messages.StringField(14) queryParameters = _messages.MessageField('QueryParameter', 15, repeated=True) rangePartitioning = _messages.MessageField('RangePartitioning', 16) schemaUpdateOptions = _messages.StringField(17, repeated=True) tableDefinitions = _messages.MessageField('TableDefinitionsValue', 18) timePartitioning = _messages.MessageField('TimePartitioning', 19) useLegacySql = _messages.BooleanField(20, default=True) useQueryCache = _messages.BooleanField(21, default=True) userDefinedFunctionResources = _messages.MessageField( 'UserDefinedFunctionResource', 22, repeated=True) writeDisposition = _messages.StringField(23) class JobConfigurationTableCopy(_messages.Message): r"""A JobConfigurationTableCopy object. Fields: createDisposition: [Optional] Specifies whether the job is allowed to create new tables. The following values are supported: CREATE_IF_NEEDED: If the table does not exist, BigQuery creates the table. CREATE_NEVER: The table must already exist. If it does not, a 'notFound' error is returned in the job result. The default value is CREATE_IF_NEEDED. Creation, truncation and append actions occur as one atomic update upon job completion. destinationEncryptionConfiguration: Custom encryption configuration (e.g., Cloud KMS keys). destinationExpirationTime: [Optional] The time when the destination table expires. Expired tables will be deleted and their storage reclaimed. destinationTable: [Required] The destination table operationType: [Optional] Supported operation types in table copy job. sourceTable: [Pick one] Source table to copy. sourceTables: [Pick one] Source tables to copy. writeDisposition: [Optional] Specifies the action that occurs if the destination table already exists. The following values are supported: WRITE_TRUNCATE: If the table already exists, BigQuery overwrites the table data. WRITE_APPEND: If the table already exists, BigQuery appends the data to the table. WRITE_EMPTY: If the table already exists and contains data, a 'duplicate' error is returned in the job result. The default value is WRITE_EMPTY. Each action is atomic and only occurs if BigQuery is able to complete the job successfully. Creation, truncation and append actions occur as one atomic update upon job completion. """ createDisposition = _messages.StringField(1) destinationEncryptionConfiguration = _messages.MessageField( 'EncryptionConfiguration', 2) destinationExpirationTime = _messages.MessageField('extra_types.JsonValue', 3) destinationTable = _messages.MessageField('TableReference', 4) operationType = _messages.StringField(5) sourceTable = _messages.MessageField('TableReference', 6) sourceTables = _messages.MessageField('TableReference', 7, repeated=True) writeDisposition = _messages.StringField(8) class JobList(_messages.Message): r"""A JobList object. Messages: JobsValueListEntry: A JobsValueListEntry object. Fields: etag: A hash of this page of results. jobs: List of jobs that were requested. kind: The resource type of the response. nextPageToken: A token to request the next page of results. """ class JobsValueListEntry(_messages.Message): r"""A JobsValueListEntry object. Fields: configuration: [Full-projection-only] Specifies the job configuration. errorResult: A result object that will be present only if the job has failed. id: Unique opaque ID of the job. jobReference: Job reference uniquely identifying the job. kind: The resource type. state: Running state of the job. When the state is DONE, errorResult can be checked to determine whether the job succeeded or failed. statistics: [Output-only] Information about the job, including starting time and ending time of the job. status: [Full-projection-only] Describes the state of the job. user_email: [Full-projection-only] Email address of the user who ran the job. """ configuration = _messages.MessageField('JobConfiguration', 1) errorResult = _messages.MessageField('ErrorProto', 2) id = _messages.StringField(3) jobReference = _messages.MessageField('JobReference', 4) kind = _messages.StringField(5, default='bigquery#job') state = _messages.StringField(6) statistics = _messages.MessageField('JobStatistics', 7) status = _messages.MessageField('JobStatus', 8) user_email = _messages.StringField(9) etag = _messages.StringField(1) jobs = _messages.MessageField('JobsValueListEntry', 2, repeated=True) kind = _messages.StringField(3, default='bigquery#jobList') nextPageToken = _messages.StringField(4) class JobReference(_messages.Message): r"""A JobReference object. Fields: jobId: [Required] The ID of the job. The ID must contain only letters (a-z, A-Z), numbers (0-9), underscores (_), or dashes (-). The maximum length is 1,024 characters. location: The geographic location of the job. See details at https://cloud.google.com/bigquery/docs/locations#specifying_your_locatio n. projectId: [Required] The ID of the project containing this job. """ jobId = _messages.StringField(1) location = _messages.StringField(2) projectId = _messages.StringField(3) class JobStatistics(_messages.Message): r"""A JobStatistics object. Messages: ReservationUsageValueListEntry: A ReservationUsageValueListEntry object. Fields: completionRatio: [TrustedTester] [Output-only] Job progress (0.0 -> 1.0) for LOAD and EXTRACT jobs. creationTime: [Output-only] Creation time of this job, in milliseconds since the epoch. This field will be present on all jobs. endTime: [Output-only] End time of this job, in milliseconds since the epoch. This field will be present whenever a job is in the DONE state. extract: [Output-only] Statistics for an extract job. load: [Output-only] Statistics for a load job. numChildJobs: [Output-only] Number of child jobs executed. parentJobId: [Output-only] If this is a child job, the id of the parent. query: [Output-only] Statistics for a query job. quotaDeferments: [Output-only] Quotas which delayed this job's start time. reservationUsage: [Output-only] Job resource usage breakdown by reservation. reservation_id: [Output-only] Name of the primary reservation assigned to this job. Note that this could be different than reservations reported in the reservation usage field if parent reservations were used to execute this job. rowLevelSecurityStatistics: [Output-only] [Preview] Statistics for row- level security. Present only for query and extract jobs. scriptStatistics: [Output-only] Statistics for a child job of a script. startTime: [Output-only] Start time of this job, in milliseconds since the epoch. This field will be present when the job transitions from the PENDING state to either RUNNING or DONE. totalBytesProcessed: [Output-only] [Deprecated] Use the bytes processed in the query statistics instead. totalSlotMs: [Output-only] Slot-milliseconds for the job. transactionInfoTemplate: [Output-only] [Alpha] Information of the multi- statement transaction if this job is part of one. """ class ReservationUsageValueListEntry(_messages.Message): r"""A ReservationUsageValueListEntry object. Fields: name: [Output-only] Reservation name or "unreserved" for on-demand resources usage. slotMs: [Output-only] Slot-milliseconds the job spent in the given reservation. """ name = _messages.StringField(1) slotMs = _messages.IntegerField(2) completionRatio = _messages.FloatField(1) creationTime = _messages.IntegerField(2) endTime = _messages.IntegerField(3) extract = _messages.MessageField('JobStatistics4', 4) load = _messages.MessageField('JobStatistics3', 5) numChildJobs = _messages.IntegerField(6) parentJobId = _messages.StringField(7) query = _messages.MessageField('JobStatistics2', 8) quotaDeferments = _messages.StringField(9, repeated=True) reservationUsage = _messages.MessageField( 'ReservationUsageValueListEntry', 10, repeated=True) reservation_id = _messages.StringField(11) rowLevelSecurityStatistics = _messages.MessageField( 'RowLevelSecurityStatistics', 12) scriptStatistics = _messages.MessageField('ScriptStatistics', 13) startTime = _messages.IntegerField(14) totalBytesProcessed = _messages.IntegerField(15) totalSlotMs = _messages.IntegerField(16) transactionInfoTemplate = _messages.MessageField('TransactionInfo', 17) class JobStatistics2(_messages.Message): r"""A JobStatistics2 object. Messages: ReservationUsageValueListEntry: A ReservationUsageValueListEntry object. Fields: billingTier: [Output-only] Billing tier for the job. cacheHit: [Output-only] Whether the query result was fetched from the query cache. ddlAffectedRowAccessPolicyCount: [Output-only] [Preview] The number of row access policies affected by a DDL statement. Present only for DROP ALL ROW ACCESS POLICIES queries. ddlOperationPerformed: The DDL operation performed, possibly dependent on the pre-existence of the DDL target. Possible values (new values might be added in the future): "CREATE": The query created the DDL target. "SKIP": No-op. Example cases: the query is CREATE TABLE IF NOT EXISTS while the table already exists, or the query is DROP TABLE IF EXISTS while the table does not exist. "REPLACE": The query replaced the DDL target. Example case: the query is CREATE OR REPLACE TABLE, and the table already exists. "DROP": The query deleted the DDL target. ddlTargetRoutine: The DDL target routine. Present only for CREATE/DROP FUNCTION/PROCEDURE queries. ddlTargetRowAccessPolicy: [Output-only] [Preview] The DDL target row access policy. Present only for CREATE/DROP ROW ACCESS POLICY queries. ddlTargetTable: [Output-only] The DDL target table. Present only for CREATE/DROP TABLE/VIEW and DROP ALL ROW ACCESS POLICIES queries. estimatedBytesProcessed: [Output-only] The original estimate of bytes processed for the job. modelTraining: [Output-only, Beta] Information about create model query job progress. modelTrainingCurrentIteration: [Output-only, Beta] Deprecated; do not use. modelTrainingExpectedTotalIteration: [Output-only, Beta] Deprecated; do not use. numDmlAffectedRows: [Output-only] The number of rows affected by a DML statement. Present only for DML statements INSERT, UPDATE or DELETE. queryPlan: [Output-only] Describes execution plan for the query. referencedRoutines: [Output-only] Referenced routines (persistent user- defined functions and stored procedures) for the job. referencedTables: [Output-only] Referenced tables for the job. Queries that reference more than 50 tables will not have a complete list. reservationUsage: [Output-only] Job resource usage breakdown by reservation. schema: [Output-only] The schema of the results. Present only for successful dry run of non-legacy SQL queries. statementType: The type of query statement, if valid. Possible values (new values might be added in the future): "SELECT": SELECT query. "INSERT": INSERT query; see https://cloud.google.com/bigquery/docs/reference/standard-sql/data- manipulation-language. "UPDATE": UPDATE query; see https://cloud.google.com/bigquery/docs/reference/standard-sql/data- manipulation-language. "DELETE": DELETE query; see https://cloud.google.com/bigquery/docs/reference/standard-sql/data- manipulation-language. "MERGE": MERGE query; see https://cloud.google.com/bigquery/docs/reference/standard-sql/data- manipulation-language. "ALTER_TABLE": ALTER TABLE query. "ALTER_VIEW": ALTER VIEW query. "ASSERT": ASSERT condition AS 'description'. "CREATE_FUNCTION": CREATE FUNCTION query. "CREATE_MODEL": CREATE [OR REPLACE] MODEL ... AS SELECT ... . "CREATE_PROCEDURE": CREATE PROCEDURE query. "CREATE_TABLE": CREATE [OR REPLACE] TABLE without AS SELECT. "CREATE_TABLE_AS_SELECT": CREATE [OR REPLACE] TABLE ... AS SELECT ... . "CREATE_VIEW": CREATE [OR REPLACE] VIEW ... AS SELECT ... . "DROP_FUNCTION" : DROP FUNCTION query. "DROP_PROCEDURE": DROP PROCEDURE query. "DROP_TABLE": DROP TABLE query. "DROP_VIEW": DROP VIEW query. timeline: [Output-only] [Beta] Describes a timeline of job execution. totalBytesBilled: [Output-only] Total bytes billed for the job. totalBytesProcessed: [Output-only] Total bytes processed for the job. totalBytesProcessedAccuracy: [Output-only] For dry-run jobs, totalBytesProcessed is an estimate and this field specifies the accuracy of the estimate. Possible values can be: UNKNOWN: accuracy of the estimate is unknown. PRECISE: estimate is precise. LOWER_BOUND: estimate is lower bound of what the query would cost. UPPER_BOUND: estimate is upper bound of what the query would cost. totalPartitionsProcessed: [Output-only] Total number of partitions processed from all partitioned tables referenced in the job. totalSlotMs: [Output-only] Slot-milliseconds for the job. undeclaredQueryParameters: Standard SQL only: list of undeclared query parameters detected during a dry run validation. """ class ReservationUsageValueListEntry(_messages.Message): r"""A ReservationUsageValueListEntry object. Fields: name: [Output-only] Reservation name or "unreserved" for on-demand resources usage. slotMs: [Output-only] Slot-milliseconds the job spent in the given reservation. """ name = _messages.StringField(1) slotMs = _messages.IntegerField(2) billingTier = _messages.IntegerField(1, variant=_messages.Variant.INT32) cacheHit = _messages.BooleanField(2) ddlAffectedRowAccessPolicyCount = _messages.IntegerField(3) ddlOperationPerformed = _messages.StringField(4) ddlTargetRoutine = _messages.MessageField('RoutineReference', 5) ddlTargetRowAccessPolicy = _messages.MessageField( 'RowAccessPolicyReference', 6) ddlTargetTable = _messages.MessageField('TableReference', 7) estimatedBytesProcessed = _messages.IntegerField(8) modelTraining = _messages.MessageField('BigQueryModelTraining', 9) modelTrainingCurrentIteration = _messages.IntegerField( 10, variant=_messages.Variant.INT32) modelTrainingExpectedTotalIteration = _messages.IntegerField(11) numDmlAffectedRows = _messages.IntegerField(12) queryPlan = _messages.MessageField('ExplainQueryStage', 13, repeated=True) referencedRoutines = _messages.MessageField( 'RoutineReference', 14, repeated=True) referencedTables = _messages.MessageField('TableReference', 15, repeated=True) reservationUsage = _messages.MessageField( 'ReservationUsageValueListEntry', 16, repeated=True) schema = _messages.MessageField('TableSchema', 17) statementType = _messages.StringField(18) timeline = _messages.MessageField('QueryTimelineSample', 19, repeated=True) totalBytesBilled = _messages.IntegerField(20) totalBytesProcessed = _messages.IntegerField(21) totalBytesProcessedAccuracy = _messages.StringField(22) totalPartitionsProcessed = _messages.IntegerField(23) totalSlotMs = _messages.IntegerField(24) undeclaredQueryParameters = _messages.MessageField( 'QueryParameter', 25, repeated=True) class JobStatistics3(_messages.Message): r"""A JobStatistics3 object. Fields: badRecords: [Output-only] The number of bad records encountered. Note that if the job has failed because of more bad records encountered than the maximum allowed in the load job configuration, then this number can be less than the total number of bad records present in the input data. inputFileBytes: [Output-only] Number of bytes of source data in a load job. inputFiles: [Output-only] Number of source files in a load job. outputBytes: [Output-only] Size of the loaded data in bytes. Note that while a load job is in the running state, this value may change. outputRows: [Output-only] Number of rows imported in a load job. Note that while an import job is in the running state, this value may change. """ badRecords = _messages.IntegerField(1) inputFileBytes = _messages.IntegerField(2) inputFiles = _messages.IntegerField(3) outputBytes = _messages.IntegerField(4) outputRows = _messages.IntegerField(5) class JobStatistics4(_messages.Message): r"""A JobStatistics4 object. Fields: destinationUriFileCounts: [Output-only] Number of files per destination URI or URI pattern specified in the extract configuration. These values will be in the same order as the URIs specified in the 'destinationUris' field. inputBytes: [Output-only] Number of user bytes extracted into the result. This is the byte count as computed by BigQuery for billing purposes. """ destinationUriFileCounts = _messages.IntegerField(1, repeated=True) inputBytes = _messages.IntegerField(2) class JobStatus(_messages.Message): r"""A JobStatus object. Fields: errorResult: [Output-only] Final error result of the job. If present, indicates that the job has completed and was unsuccessful. errors: [Output-only] The first errors encountered during the running of the job. The final message includes the number of errors that caused the process to stop. Errors here do not necessarily mean that the job has completed or was unsuccessful. state: [Output-only] Running state of the job. """ errorResult = _messages.MessageField('ErrorProto', 1) errors = _messages.MessageField('ErrorProto', 2, repeated=True) state = _messages.StringField(3) @encoding.MapUnrecognizedFields('additionalProperties') class JsonObject(_messages.Message): r"""Represents a single JSON object. Messages: AdditionalProperty: An additional property for a JsonObject object. Fields: additionalProperties: Additional properties of type JsonObject """ class AdditionalProperty(_messages.Message): r"""An additional property for a JsonObject object. Fields: key: Name of the additional property. value: A JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('JsonValue', 2) additionalProperties = _messages.MessageField( 'AdditionalProperty', 1, repeated=True) JsonValue = extra_types.JsonValue class ListModelsResponse(_messages.Message): r"""A ListModelsResponse object. Fields: models: Models in the requested dataset. Only the following fields are populated: model_reference, model_type, creation_time, last_modified_time and labels. nextPageToken: A token to request the next page of results. """ models = _messages.MessageField('Model', 1, repeated=True) nextPageToken = _messages.StringField(2) class ListRoutinesResponse(_messages.Message): r"""A ListRoutinesResponse object. Fields: nextPageToken: A token to request the next page of results. routines: Routines in the requested dataset. Unless read_mask is set in the request, only the following fields are populated: etag, project_id, dataset_id, routine_id, routine_type, creation_time, last_modified_time, and language. """ nextPageToken = _messages.StringField(1) routines = _messages.MessageField('Routine', 2, repeated=True) class ListRowAccessPoliciesResponse(_messages.Message): r"""Response message for the ListRowAccessPolicies method. Fields: nextPageToken: A token to request the next page of results. rowAccessPolicies: Row access policies on the requested table. """ nextPageToken = _messages.StringField(1) rowAccessPolicies = _messages.MessageField( 'RowAccessPolicy', 2, repeated=True) class LocationMetadata(_messages.Message): r"""BigQuery-specific metadata about a location. This will be set on google.cloud.location.Location.metadata in Cloud Location API responses. Fields: legacyLocationId: The legacy BigQuery location ID, e.g. "EU" for the "europe" location. This is for any API consumers that need the legacy "US" and "EU" locations. """ legacyLocationId = _messages.StringField(1) class MaterializedViewDefinition(_messages.Message): r"""A MaterializedViewDefinition object. Fields: enableRefresh: [Optional] [TrustedTester] Enable automatic refresh of the materialized view when the base table is updated. The default value is "true". lastRefreshTime: [Output-only] [TrustedTester] The time when this materialized view was last modified, in milliseconds since the epoch. query: [Required] A query whose result is persisted. refreshIntervalMs: [Optional] [TrustedTester] The maximum frequency at which this materialized view will be refreshed. The default value is "1800000" (30 minutes). """ enableRefresh = _messages.BooleanField(1) lastRefreshTime = _messages.IntegerField(2) query = _messages.StringField(3) refreshIntervalMs = _messages.IntegerField(4) class Model(_messages.Message): r"""A Model object. Enums: ModelTypeValueValuesEnum: Output only. Type of the model resource. Messages: LabelsValue: The labels associated with this model. You can use these to organize and group your models. Label keys and values can be no longer than 63 characters, can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. Label values are optional. Label keys must start with a letter and each label in the list must have a different key. Fields: creationTime: Output only. The time when this model was created, in millisecs since the epoch. description: Optional. A user-friendly description of this model. encryptionConfiguration: Custom encryption configuration (e.g., Cloud KMS keys). This shows the encryption configuration of the model data while stored in BigQuery storage. This field can be used with PatchModel to update encryption key for an already encrypted model. etag: Output only. A hash of this resource. expirationTime: Optional. The time when this model expires, in milliseconds since the epoch. If not present, the model will persist indefinitely. Expired models will be deleted and their storage reclaimed. The defaultTableExpirationMs property of the encapsulating dataset can be used to set a default expirationTime on newly created models. featureColumns: Output only. Input feature columns that were used to train this model. friendlyName: Optional. A descriptive name for this model. labelColumns: Output only. Label columns that were used to train this model. The output of the model will have a "predicted_" prefix to these columns. labels: The labels associated with this model. You can use these to organize and group your models. Label keys and values can be no longer than 63 characters, can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. Label values are optional. Label keys must start with a letter and each label in the list must have a different key. lastModifiedTime: Output only. The time when this model was last modified, in millisecs since the epoch. location: Output only. The geographic location where the model resides. This value is inherited from the dataset. modelReference: Required. Unique identifier for this model. modelType: Output only. Type of the model resource. trainingRuns: Output only. Information for all training runs in increasing order of start_time. """ class ModelTypeValueValuesEnum(_messages.Enum): r"""Output only. Type of the model resource. Values: MODEL_TYPE_UNSPECIFIED: LINEAR_REGRESSION: Linear regression model. LOGISTIC_REGRESSION: Logistic regression based classification model. KMEANS: K-means clustering model. MATRIX_FACTORIZATION: Matrix factorization model. DNN_CLASSIFIER: [Beta] DNN classifier model. TENSORFLOW: [Beta] An imported TensorFlow model. DNN_REGRESSOR: [Beta] DNN regressor model. BOOSTED_TREE_REGRESSOR: [Beta] Boosted tree regressor model. BOOSTED_TREE_CLASSIFIER: [Beta] Boosted tree classifier model. ARIMA: [Beta] ARIMA model. AUTOML_REGRESSOR: [Beta] AutoML Tables regression model. AUTOML_CLASSIFIER: [Beta] AutoML Tables classification model. """ MODEL_TYPE_UNSPECIFIED = 0 LINEAR_REGRESSION = 1 LOGISTIC_REGRESSION = 2 KMEANS = 3 MATRIX_FACTORIZATION = 4 DNN_CLASSIFIER = 5 TENSORFLOW = 6 DNN_REGRESSOR = 7 BOOSTED_TREE_REGRESSOR = 8 BOOSTED_TREE_CLASSIFIER = 9 ARIMA = 10 AUTOML_REGRESSOR = 11 AUTOML_CLASSIFIER = 12 @encoding.MapUnrecognizedFields('additionalProperties') class LabelsValue(_messages.Message): r"""The labels associated with this model. You can use these to organize and group your models. Label keys and values can be no longer than 63 characters, can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. Label values are optional. Label keys must start with a letter and each label in the list must have a different key. Messages: AdditionalProperty: An additional property for a LabelsValue object. Fields: additionalProperties: Additional properties of type LabelsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a LabelsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField( 'AdditionalProperty', 1, repeated=True) creationTime = _messages.IntegerField(1) description = _messages.StringField(2) encryptionConfiguration = _messages.MessageField('EncryptionConfiguration', 3) etag = _messages.StringField(4) expirationTime = _messages.IntegerField(5) featureColumns = _messages.MessageField('StandardSqlField', 6, repeated=True) friendlyName = _messages.StringField(7) labelColumns = _messages.MessageField('StandardSqlField', 8, repeated=True) labels = _messages.MessageField('LabelsValue', 9) lastModifiedTime = _messages.IntegerField(10) location = _messages.StringField(11) modelReference = _messages.MessageField('ModelReference', 12) modelType = _messages.EnumField('ModelTypeValueValuesEnum', 13) trainingRuns = _messages.MessageField('TrainingRun', 14, repeated=True) class ModelDefinition(_messages.Message): r"""A ModelDefinition object. Messages: ModelOptionsValue: [Output-only, Beta] Model options used for the first training run. These options are immutable for subsequent training runs. Default values are used for any options not specified in the input query. Fields: modelOptions: [Output-only, Beta] Model options used for the first training run. These options are immutable for subsequent training runs. Default values are used for any options not specified in the input query. trainingRuns: [Output-only, Beta] Information about ml training runs, each training run comprises of multiple iterations and there may be multiple training runs for the model if warm start is used or if a user decides to continue a previously cancelled query. """ class ModelOptionsValue(_messages.Message): r"""[Output-only, Beta] Model options used for the first training run. These options are immutable for subsequent training runs. Default values are used for any options not specified in the input query. Fields: labels: A string attribute. lossType: A string attribute. modelType: A string attribute. """ labels = _messages.StringField(1, repeated=True) lossType = _messages.StringField(2) modelType = _messages.StringField(3) modelOptions = _messages.MessageField('ModelOptionsValue', 1) trainingRuns = _messages.MessageField('BqmlTrainingRun', 2, repeated=True) class ModelReference(_messages.Message): r"""A ModelReference object. Fields: datasetId: [Required] The ID of the dataset containing this model. modelId: [Required] The ID of the model. The ID must contain only letters (a-z, A-Z), numbers (0-9), or underscores (_). The maximum length is 1,024 characters. projectId: [Required] The ID of the project containing this model. """ datasetId = _messages.StringField(1) modelId = _messages.StringField(2) projectId = _messages.StringField(3) class MultiClassClassificationMetrics(_messages.Message): r"""Evaluation metrics for multi-class classification/classifier models. Fields: aggregateClassificationMetrics: Aggregate classification metrics. confusionMatrixList: Confusion matrix at different thresholds. """ aggregateClassificationMetrics = _messages.MessageField( 'AggregateClassificationMetrics', 1) confusionMatrixList = _messages.MessageField( 'ConfusionMatrix', 2, repeated=True) class Policy(_messages.Message): r"""An Identity and Access Management (IAM) policy, which specifies access controls for Google Cloud resources. A `Policy` is a collection of `bindings`. A `binding` binds one or more `members` to a single `role`. Members can be user accounts, service accounts, Google groups, and domains (such as G Suite). A `role` is a named list of permissions; each `role` can be an IAM predefined role or a user-created custom role. For some types of Google Cloud resources, a `binding` can also specify a `condition`, which is a logical expression that allows access to a resource only if the expression evaluates to `true`. A condition can add constraints based on attributes of the request, the resource, or both. To learn which resources support conditions in their IAM policies, see the [IAM documentation](https://cloud.google.com/iam/help/conditions/resource- policies). **JSON example:** { "bindings": [ { "role": "roles/resourcemanager.organizationAdmin", "members": [ "user:mike@example.com", "group:admins@example.com", "domain:google.com", "serviceAccount:my-project-id@appspot.gserviceaccount.com" ] }, { "role": "roles/resourcemanager.organizationViewer", "members": [ "user:eve@example.com" ], "condition": { "title": "expirable access", "description": "Does not grant access after Sep 2020", "expression": "request.time < timestamp('2020-10-01T00:00:00.000Z')", } } ], "etag": "BwWWja0YfJA=", "version": 3 } **YAML example:** bindings: - members: - user:mike@example.com - group:admins@example.com - domain:google.com - serviceAccount:my-project-id@appspot.gserviceaccount.com role: roles/resourcemanager.organizationAdmin - members: - user:eve@example.com role: roles/resourcemanager.organizationViewer condition: title: expirable access description: Does not grant access after Sep 2020 expression: request.time < timestamp('2020-10-01T00:00:00.000Z') - etag: BwWWja0YfJA= - version: 3 For a description of IAM and its features, see the [IAM documentation](https://cloud.google.com/iam/docs/). Fields: auditConfigs: Specifies cloud audit logging configuration for this policy. bindings: Associates a list of `members` to a `role`. Optionally, may specify a `condition` that determines how and when the `bindings` are applied. Each of the `bindings` must contain at least one member. etag: `etag` is used for optimistic concurrency control as a way to help prevent simultaneous updates of a policy from overwriting each other. It is strongly suggested that systems make use of the `etag` in the read- modify-write cycle to perform policy updates in order to avoid race conditions: An `etag` is returned in the response to `getIamPolicy`, and systems are expected to put that etag in the request to `setIamPolicy` to ensure that their change will be applied to the same version of the policy. **Important:** If you use IAM Conditions, you must include the `etag` field whenever you call `setIamPolicy`. If you omit this field, then IAM allows you to overwrite a version `3` policy with a version `1` policy, and all of the conditions in the version `3` policy are lost. version: Specifies the format of the policy. Valid values are `0`, `1`, and `3`. Requests that specify an invalid value are rejected. Any operation that affects conditional role bindings must specify version `3`. This requirement applies to the following operations: * Getting a policy that includes a conditional role binding * Adding a conditional role binding to a policy * Changing a conditional role binding in a policy * Removing any role binding, with or without a condition, from a policy that includes conditions **Important:** If you use IAM Conditions, you must include the `etag` field whenever you call `setIamPolicy`. If you omit this field, then IAM allows you to overwrite a version `3` policy with a version `1` policy, and all of the conditions in the version `3` policy are lost. If a policy does not include any conditions, operations on that policy may specify any valid version or leave the field unset. To learn which resources support conditions in their IAM policies, see the [IAM documentation](https://cloud.google.com/iam/help/conditions/resource- policies). """ auditConfigs = _messages.MessageField('AuditConfig', 1, repeated=True) bindings = _messages.MessageField('Binding', 2, repeated=True) etag = _messages.BytesField(3) version = _messages.IntegerField(4, variant=_messages.Variant.INT32) class ProjectList(_messages.Message): r"""A ProjectList object. Messages: ProjectsValueListEntry: A ProjectsValueListEntry object. Fields: etag: A hash of the page of results kind: The type of list. nextPageToken: A token to request the next page of results. projects: Projects to which you have at least READ access. totalItems: The total number of projects in the list. """ class ProjectsValueListEntry(_messages.Message): r"""A ProjectsValueListEntry object. Fields: friendlyName: A descriptive name for this project. id: An opaque ID of this project. kind: The resource type. numericId: The numeric ID of this project. projectReference: A unique reference to this project. """ friendlyName = _messages.StringField(1) id = _messages.StringField(2) kind = _messages.StringField(3, default='bigquery#project') numericId = _messages.IntegerField(4, variant=_messages.Variant.UINT64) projectReference = _messages.MessageField('ProjectReference', 5) etag = _messages.StringField(1) kind = _messages.StringField(2, default='bigquery#projectList') nextPageToken = _messages.StringField(3) projects = _messages.MessageField('ProjectsValueListEntry', 4, repeated=True) totalItems = _messages.IntegerField(5, variant=_messages.Variant.INT32) class ProjectReference(_messages.Message): r"""A ProjectReference object. Fields: projectId: [Required] ID of the project. Can be either the numeric ID or the assigned ID of the project. """ projectId = _messages.StringField(1) class QueryParameter(_messages.Message): r"""A QueryParameter object. Fields: name: [Optional] If unset, this is a positional parameter. Otherwise, should be unique within a query. parameterType: [Required] The type of this parameter. parameterValue: [Required] The value of this parameter. """ name = _messages.StringField(1) parameterType = _messages.MessageField('QueryParameterType', 2) parameterValue = _messages.MessageField('QueryParameterValue', 3) class QueryParameterType(_messages.Message): r"""A QueryParameterType object. Messages: StructTypesValueListEntry: A StructTypesValueListEntry object. Fields: arrayType: [Optional] The type of the array's elements, if this is an array. structTypes: [Optional] The types of the fields of this struct, in order, if this is a struct. type: [Required] The top level type of this field. """ class StructTypesValueListEntry(_messages.Message): r"""A StructTypesValueListEntry object. Fields: description: [Optional] Human-oriented description of the field. name: [Optional] The name of this field. type: [Required] The type of this field. """ description = _messages.StringField(1) name = _messages.StringField(2) type = _messages.MessageField('QueryParameterType', 3) arrayType = _messages.MessageField('QueryParameterType', 1) structTypes = _messages.MessageField( 'StructTypesValueListEntry', 2, repeated=True) type = _messages.StringField(3) class QueryParameterValue(_messages.Message): r"""A QueryParameterValue object. Messages: StructValuesValue: [Optional] The struct field values, in order of the struct type's declaration. Fields: arrayValues: [Optional] The array values, if this is an array type. structValues: [Optional] The struct field values, in order of the struct type's declaration. value: [Optional] The value of this value, if a simple scalar type. """ @encoding.MapUnrecognizedFields('additionalProperties') class StructValuesValue(_messages.Message): r"""[Optional] The struct field values, in order of the struct type's declaration. Messages: AdditionalProperty: An additional property for a StructValuesValue object. Fields: additionalProperties: Additional properties of type StructValuesValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a StructValuesValue object. Fields: key: Name of the additional property. value: A QueryParameterValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('QueryParameterValue', 2) additionalProperties = _messages.MessageField( 'AdditionalProperty', 1, repeated=True) arrayValues = _messages.MessageField('QueryParameterValue', 1, repeated=True) structValues = _messages.MessageField('StructValuesValue', 2) value = _messages.StringField(3) class QueryRequest(_messages.Message): r"""A QueryRequest object. Messages: LabelsValue: The labels associated with this job. You can use these to organize and group your jobs. Label keys and values can be no longer than 63 characters, can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. Label values are optional. Label keys must start with a letter and each label in the list must have a different key. Fields: connectionProperties: Connection properties. defaultDataset: [Optional] Specifies the default datasetId and projectId to assume for any unqualified table names in the query. If not set, all table names in the query string must be qualified in the format 'datasetId.tableId'. dryRun: [Optional] If set to true, BigQuery doesn't run the job. Instead, if the query is valid, BigQuery returns statistics about the job such as how many bytes would be processed. If the query is invalid, an error returns. The default value is false. kind: The resource type of the request. labels: The labels associated with this job. You can use these to organize and group your jobs. Label keys and values can be no longer than 63 characters, can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. Label values are optional. Label keys must start with a letter and each label in the list must have a different key. location: The geographic location where the job should run. See details at https://cloud.google.com/bigquery/docs/locations#specifying_your_locatio n. maxResults: [Optional] The maximum number of rows of data to return per page of results. Setting this flag to a small value such as 1000 and then paging through results might improve reliability when the query result set is large. In addition to this limit, responses are also limited to 10 MB. By default, there is no maximum row count, and only the byte limit applies. maximumBytesBilled: [Optional] Limits the bytes billed for this job. Queries that will have bytes billed beyond this limit will fail (without incurring a charge). If unspecified, this will be set to your project default. parameterMode: Standard SQL only. Set to POSITIONAL to use positional (?) query parameters or to NAMED to use named (@myparam) query parameters in this query. preserveNulls: [Deprecated] This property is deprecated. query: [Required] A query string, following the BigQuery query syntax, of the query to execute. Example: "SELECT count(f1) FROM [myProjectId:myDatasetId.myTableId]". queryParameters: Query parameters for Standard SQL queries. requestId: A unique user provided identifier to ensure idempotent behavior for queries. Note that this is different from the job_id. It has the following properties: 1. It is case-sensitive, limited to up to 36 ASCII characters. A UUID is recommended. 2. Read only queries can ignore this token since they are nullipotent by definition. 3. For the purposes of idempotency ensured by the request_id, a request is considered duplicate of another only if they have the same request_id and are actually duplicates. When determining whether a request is a duplicate of the previous request, all parameters in the request that may affect the behavior are considered. For example, query, connection_properties, query_parameters, use_legacy_sql are parameters that affect the result and are considered when determining whether a request is a duplicate, but properties like timeout_ms don't affect the result and are thus not considered. Dry run query requests are never considered duplicate of another request. 4. When a duplicate mutating query request is detected, it returns: a. the results of the mutation if it completes successfully within the timeout. b. the running operation if it is still in progress at the end of the timeout. 5. Its lifetime is limited to 15 minutes. In other words, if two requests are sent with the same request_id, but more than 15 minutes apart, idempotency is not guaranteed. timeoutMs: [Optional] How long to wait for the query to complete, in milliseconds, before the request times out and returns. Note that this is only a timeout for the request, not the query. If the query takes longer to run than the timeout value, the call returns without any results and with the 'jobComplete' flag set to false. You can call GetQueryResults() to wait for the query to complete and read the results. The default value is 10000 milliseconds (10 seconds). useLegacySql: Specifies whether to use BigQuery's legacy SQL dialect for this query. The default value is true. If set to false, the query will use BigQuery's standard SQL: https://cloud.google.com/bigquery/sql- reference/ When useLegacySql is set to false, the value of flattenResults is ignored; query will be run as if flattenResults is false. useQueryCache: [Optional] Whether to look for the result in the query cache. The query cache is a best-effort cache that will be flushed whenever tables in the query are modified. The default value is true. """ @encoding.MapUnrecognizedFields('additionalProperties') class LabelsValue(_messages.Message): r"""The labels associated with this job. You can use these to organize and group your jobs. Label keys and values can be no longer than 63 characters, can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. Label values are optional. Label keys must start with a letter and each label in the list must have a different key. Messages: AdditionalProperty: An additional property for a LabelsValue object. Fields: additionalProperties: Additional properties of type LabelsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a LabelsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField( 'AdditionalProperty', 1, repeated=True) connectionProperties = _messages.MessageField( 'ConnectionProperty', 1, repeated=True) defaultDataset = _messages.MessageField('DatasetReference', 2) dryRun = _messages.BooleanField(3) kind = _messages.StringField(4, default='bigquery#queryRequest') labels = _messages.MessageField('LabelsValue', 5) location = _messages.StringField(6) maxResults = _messages.IntegerField(7, variant=_messages.Variant.UINT32) maximumBytesBilled = _messages.IntegerField(8) parameterMode = _messages.StringField(9) preserveNulls = _messages.BooleanField(10) query = _messages.StringField(11) queryParameters = _messages.MessageField('QueryParameter', 12, repeated=True) requestId = _messages.StringField(13) timeoutMs = _messages.IntegerField(14, variant=_messages.Variant.UINT32) useLegacySql = _messages.BooleanField(15, default=True) useQueryCache = _messages.BooleanField(16, default=True) class QueryResponse(_messages.Message): r"""A QueryResponse object. Fields: cacheHit: Whether the query result was fetched from the query cache. errors: [Output-only] The first errors or warnings encountered during the running of the job. The final message includes the number of errors that caused the process to stop. Errors here do not necessarily mean that the job has completed or was unsuccessful. jobComplete: Whether the query has completed or not. If rows or totalRows are present, this will always be true. If this is false, totalRows will not be available. jobReference: Reference to the Job that was created to run the query. This field will be present even if the original request timed out, in which case GetQueryResults can be used to read the results once the query has completed. Since this API only returns the first page of results, subsequent pages can be fetched via the same mechanism (GetQueryResults). kind: The resource type. numDmlAffectedRows: [Output-only] The number of rows affected by a DML statement. Present only for DML statements INSERT, UPDATE or DELETE. pageToken: A token used for paging results. rows: An object with as many results as can be contained within the maximum permitted reply size. To get any additional rows, you can call GetQueryResults and specify the jobReference returned above. schema: The schema of the results. Present only when the query completes successfully. totalBytesProcessed: The total number of bytes processed for this query. If this query was a dry run, this is the number of bytes that would be processed if the query were run. totalRows: The total number of rows in the complete query result set, which can be more than the number of rows in this single page of results. """ cacheHit = _messages.BooleanField(1) errors = _messages.MessageField('ErrorProto', 2, repeated=True) jobComplete = _messages.BooleanField(3) jobReference = _messages.MessageField('JobReference', 4) kind = _messages.StringField(5, default='bigquery#queryResponse') numDmlAffectedRows = _messages.IntegerField(6) pageToken = _messages.StringField(7) rows = _messages.MessageField('TableRow', 8, repeated=True) schema = _messages.MessageField('TableSchema', 9) totalBytesProcessed = _messages.IntegerField(10) totalRows = _messages.IntegerField(11, variant=_messages.Variant.UINT64) class QueryTimelineSample(_messages.Message): r"""A QueryTimelineSample object. Fields: activeUnits: Total number of units currently being processed by workers. This does not correspond directly to slot usage. This is the largest value observed since the last sample. completedUnits: Total parallel units of work completed by this query. elapsedMs: Milliseconds elapsed since the start of query execution. pendingUnits: Total parallel units of work remaining for the active stages. totalSlotMs: Cumulative slot-ms consumed by the query. """ activeUnits = _messages.IntegerField(1) completedUnits = _messages.IntegerField(2) elapsedMs = _messages.IntegerField(3) pendingUnits = _messages.IntegerField(4) totalSlotMs = _messages.IntegerField(5) class RangePartitioning(_messages.Message): r"""A RangePartitioning object. Messages: RangeValue: [TrustedTester] [Required] Defines the ranges for range partitioning. Fields: field: [TrustedTester] [Required] The table is partitioned by this field. The field must be a top-level NULLABLE/REQUIRED field. The only supported type is INTEGER/INT64. range: [TrustedTester] [Required] Defines the ranges for range partitioning. """ class RangeValue(_messages.Message): r"""[TrustedTester] [Required] Defines the ranges for range partitioning. Fields: end: [TrustedTester] [Required] The end of range partitioning, exclusive. interval: [TrustedTester] [Required] The width of each interval. start: [TrustedTester] [Required] The start of range partitioning, inclusive. """ end = _messages.IntegerField(1) interval = _messages.IntegerField(2) start = _messages.IntegerField(3) field = _messages.StringField(1) range = _messages.MessageField('RangeValue', 2) class RankingMetrics(_messages.Message): r"""Evaluation metrics used by weighted-ALS models specified by feedback_type=implicit. Fields: averageRank: Determines the goodness of a ranking by computing the percentile rank from the predicted confidence and dividing it by the original rank. meanAveragePrecision: Calculates a precision per user for all the items by ranking them and then averages all the precisions across all the users. meanSquaredError: Similar to the mean squared error computed in regression and explicit recommendation models except instead of computing the rating directly, the output from evaluate is computed against a preference which is 1 or 0 depending on if the rating exists or not. normalizedDiscountedCumulativeGain: A metric to determine the goodness of a ranking calculated from the predicted confidence by comparing it to an ideal rank measured by the original ratings. """ averageRank = _messages.FloatField(1) meanAveragePrecision = _messages.FloatField(2) meanSquaredError = _messages.FloatField(3) normalizedDiscountedCumulativeGain = _messages.FloatField(4) class RegressionMetrics(_messages.Message): r"""Evaluation metrics for regression and explicit feedback type matrix factorization models. Fields: meanAbsoluteError: Mean absolute error. meanSquaredError: Mean squared error. meanSquaredLogError: Mean squared log error. medianAbsoluteError: Median absolute error. rSquared: R^2 score. """ meanAbsoluteError = _messages.FloatField(1) meanSquaredError = _messages.FloatField(2) meanSquaredLogError = _messages.FloatField(3) medianAbsoluteError = _messages.FloatField(4) rSquared = _messages.FloatField(5) class Routine(_messages.Message): r"""A user-defined function or a stored procedure. Enums: DeterminismLevelValueValuesEnum: Optional. [Experimental] The determinism level of the JavaScript UDF if defined. LanguageValueValuesEnum: Optional. Defaults to "SQL". RoutineTypeValueValuesEnum: Required. The type of routine. Fields: arguments: Optional. creationTime: Output only. The time when this routine was created, in milliseconds since the epoch. definitionBody: Required. The body of the routine. For functions, this is the expression in the AS clause. If language=SQL, it is the substring inside (but excluding) the parentheses. For example, for the function created with the following statement: `CREATE FUNCTION JoinLines(x string, y string) as (concat(x, "\n", y))` The definition_body is `concat(x, "\n", y)` (\n is not replaced with linebreak). If language=JAVASCRIPT, it is the evaluated string in the AS clause. For example, for the function created with the following statement: `CREATE FUNCTION f() RETURNS STRING LANGUAGE js AS 'return "\n";\n'` The definition_body is `return "\n";\n` Note that both \n are replaced with linebreaks. description: Optional. [Experimental] The description of the routine if defined. determinismLevel: Optional. [Experimental] The determinism level of the JavaScript UDF if defined. etag: Output only. A hash of this resource. importedLibraries: Optional. If language = "JAVASCRIPT", this field stores the path of the imported JAVASCRIPT libraries. language: Optional. Defaults to "SQL". lastModifiedTime: Output only. The time when this routine was last modified, in milliseconds since the epoch. returnType: Optional if language = "SQL"; required otherwise. If absent, the return type is inferred from definition_body at query time in each query that references this routine. If present, then the evaluated result will be cast to the specified returned type at query time. For example, for the functions created with the following statements: * `CREATE FUNCTION Add(x FLOAT64, y FLOAT64) RETURNS FLOAT64 AS (x + y);` * `CREATE FUNCTION Increment(x FLOAT64) AS (Add(x, 1));` * `CREATE FUNCTION Decrement(x FLOAT64) RETURNS FLOAT64 AS (Add(x, -1));` The return_type is `{type_kind: "FLOAT64"}` for `Add` and `Decrement`, and is absent for `Increment` (inferred as FLOAT64 at query time). Suppose the function `Add` is replaced by `CREATE OR REPLACE FUNCTION Add(x INT64, y INT64) AS (x + y);` Then the inferred return type of `Increment` is automatically changed to INT64 at query time, while the return type of `Decrement` remains FLOAT64. routineReference: Required. Reference describing the ID of this routine. routineType: Required. The type of routine. """ class DeterminismLevelValueValuesEnum(_messages.Enum): r"""Optional. [Experimental] The determinism level of the JavaScript UDF if defined. Values: DETERMINISM_LEVEL_UNSPECIFIED: The determinism of the UDF is unspecified. DETERMINISTIC: The UDF is deterministic, meaning that 2 function calls with the same inputs always produce the same result, even across 2 query runs. NOT_DETERMINISTIC: The UDF is not deterministic. """ DETERMINISM_LEVEL_UNSPECIFIED = 0 DETERMINISTIC = 1 NOT_DETERMINISTIC = 2 class LanguageValueValuesEnum(_messages.Enum): r"""Optional. Defaults to "SQL". Values: LANGUAGE_UNSPECIFIED: SQL: SQL language. JAVASCRIPT: JavaScript language. """ LANGUAGE_UNSPECIFIED = 0 SQL = 1 JAVASCRIPT = 2 class RoutineTypeValueValuesEnum(_messages.Enum): r"""Required. The type of routine. Values: ROUTINE_TYPE_UNSPECIFIED: SCALAR_FUNCTION: Non-builtin permanent scalar function. PROCEDURE: Stored procedure. """ ROUTINE_TYPE_UNSPECIFIED = 0 SCALAR_FUNCTION = 1 PROCEDURE = 2 arguments = _messages.MessageField('Argument', 1, repeated=True) creationTime = _messages.IntegerField(2) definitionBody = _messages.StringField(3) description = _messages.StringField(4) determinismLevel = _messages.EnumField('DeterminismLevelValueValuesEnum', 5) etag = _messages.StringField(6) importedLibraries = _messages.StringField(7, repeated=True) language = _messages.EnumField('LanguageValueValuesEnum', 8) lastModifiedTime = _messages.IntegerField(9) returnType = _messages.MessageField('StandardSqlDataType', 10) routineReference = _messages.MessageField('RoutineReference', 11) routineType = _messages.EnumField('RoutineTypeValueValuesEnum', 12) class RoutineReference(_messages.Message): r"""A RoutineReference object. Fields: datasetId: [Required] The ID of the dataset containing this routine. projectId: [Required] The ID of the project containing this routine. routineId: [Required] The ID of the routine. The ID must contain only letters (a-z, A-Z), numbers (0-9), or underscores (_). The maximum length is 256 characters. """ datasetId = _messages.StringField(1) projectId = _messages.StringField(2) routineId = _messages.StringField(3) class Row(_messages.Message): r"""A single row in the confusion matrix. Fields: actualLabel: The original label of this row. entries: Info describing predicted label distribution. """ actualLabel = _messages.StringField(1) entries = _messages.MessageField('Entry', 2, repeated=True) class RowAccessPolicy(_messages.Message): r"""Represents access on a subset of rows on the specified table, defined by its filter predicate. Access to the subset of rows is controlled by its IAM policy. Fields: creationTime: Output only. The time when this row access policy was created, in milliseconds since the epoch. etag: Output only. A hash of this resource. filterPredicate: Required. A SQL boolean expression that represents the rows defined by this row access policy, similar to the boolean expression in a WHERE clause of a SELECT query on a table. References to other tables, routines, and temporary functions are not supported. Examples: region="EU" date_field = CAST('2019-9-27' as DATE) nullable_field is not NULL numeric_field BETWEEN 1.0 AND 5.0 lastModifiedTime: Output only. The time when this row access policy was last modified, in milliseconds since the epoch. rowAccessPolicyReference: Required. Reference describing the ID of this row access policy. """ creationTime = _messages.StringField(1) etag = _messages.StringField(2) filterPredicate = _messages.StringField(3) lastModifiedTime = _messages.StringField(4) rowAccessPolicyReference = _messages.MessageField( 'RowAccessPolicyReference', 5) class RowAccessPolicyReference(_messages.Message): r"""A RowAccessPolicyReference object. Fields: datasetId: [Required] The ID of the dataset containing this row access policy. policyId: [Required] The ID of the row access policy. The ID must contain only letters (a-z, A-Z), numbers (0-9), or underscores (_). The maximum length is 256 characters. projectId: [Required] The ID of the project containing this row access policy. tableId: [Required] The ID of the table containing this row access policy. """ datasetId = _messages.StringField(1) policyId = _messages.StringField(2) projectId = _messages.StringField(3) tableId = _messages.StringField(4) class RowLevelSecurityStatistics(_messages.Message): r"""A RowLevelSecurityStatistics object. Fields: rowLevelSecurityApplied: [Output-only] [Preview] Whether any accessed data was protected by row access policies. """ rowLevelSecurityApplied = _messages.BooleanField(1) class ScriptStackFrame(_messages.Message): r"""A ScriptStackFrame object. Fields: endColumn: [Output-only] One-based end column. endLine: [Output-only] One-based end line. procedureId: [Output-only] Name of the active procedure, empty if in a top-level script. startColumn: [Output-only] One-based start column. startLine: [Output-only] One-based start line. text: [Output-only] Text of the current statement/expression. """ endColumn = _messages.IntegerField(1, variant=_messages.Variant.INT32) endLine = _messages.IntegerField(2, variant=_messages.Variant.INT32) procedureId = _messages.StringField(3) startColumn = _messages.IntegerField(4, variant=_messages.Variant.INT32) startLine = _messages.IntegerField(5, variant=_messages.Variant.INT32) text = _messages.StringField(6) class ScriptStatistics(_messages.Message): r"""A ScriptStatistics object. Fields: evaluationKind: [Output-only] Whether this child job was a statement or expression. stackFrames: Stack trace showing the line/column/procedure name of each frame on the stack at the point where the current evaluation happened. The leaf frame is first, the primary script is last. Never empty. """ evaluationKind = _messages.StringField(1) stackFrames = _messages.MessageField('ScriptStackFrame', 2, repeated=True) class SetIamPolicyRequest(_messages.Message): r"""Request message for `SetIamPolicy` method. Fields: policy: REQUIRED: The complete policy to be applied to the `resource`. The size of the policy is limited to a few 10s of KB. An empty policy is a valid policy but certain Cloud Platform services (such as Projects) might reject them. updateMask: OPTIONAL: A FieldMask specifying which fields of the policy to modify. Only the fields in the mask will be modified. If no mask is provided, the following default mask is used: `paths: "bindings, etag"` """ policy = _messages.MessageField('Policy', 1) updateMask = _messages.StringField(2) class SnapshotDefinition(_messages.Message): r"""A SnapshotDefinition object. Fields: baseTableReference: [Required] Reference describing the ID of the table that is snapshotted. snapshotTime: [Required] The time at which the base table was snapshot. """ baseTableReference = _messages.MessageField('TableReference', 1) snapshotTime = _message_types.DateTimeField(2) class StandardQueryParameters(_messages.Message): r"""Query parameters accepted by all methods. Enums: AltValueValuesEnum: Data format for the response. Fields: alt: Data format for the response. fields: Selector specifying which fields to include in a partial response. key: API key. Your API key identifies your project and provides you with API access, quota, and reports. Required unless you provide an OAuth 2.0 token. oauth_token: OAuth 2.0 token for the current user. prettyPrint: Returns response with indentations and line breaks. quotaUser: An opaque string that represents a user for quota purposes. Must not exceed 40 characters. trace: A tracing token of the form "token:" to include in api requests. userIp: Deprecated. Please use quotaUser instead. """ class AltValueValuesEnum(_messages.Enum): r"""Data format for the response. Values: json: Responses with Content-Type of application/json """ json = 0 alt = _messages.EnumField('AltValueValuesEnum', 1, default='json') fields = _messages.StringField(2) key = _messages.StringField(3) oauth_token = _messages.StringField(4) prettyPrint = _messages.BooleanField(5, default=True) quotaUser = _messages.StringField(6) trace = _messages.StringField(7) userIp = _messages.StringField(8) class StandardSqlDataType(_messages.Message): r"""The type of a variable, e.g., a function argument. Examples: INT64: {type_kind="INT64"} ARRAY: {type_kind="ARRAY", array_element_type="STRING"} STRUCT>: {type_kind="STRUCT", struct_type={fields=[ {name="x", type={type_kind="STRING"}}, {name="y", type={type_kind="ARRAY", array_element_type="DATE"}} ]}} Enums: TypeKindValueValuesEnum: Required. The top level type of this field. Can be any standard SQL data type (e.g., "INT64", "DATE", "ARRAY"). Fields: arrayElementType: The type of the array's elements, if type_kind = "ARRAY". structType: The fields of this struct, in order, if type_kind = "STRUCT". typeKind: Required. The top level type of this field. Can be any standard SQL data type (e.g., "INT64", "DATE", "ARRAY"). """ class TypeKindValueValuesEnum(_messages.Enum): r"""Required. The top level type of this field. Can be any standard SQL data type (e.g., "INT64", "DATE", "ARRAY"). Values: TYPE_KIND_UNSPECIFIED: Invalid type. INT64: Encoded as a string in decimal format. BOOL: Encoded as a boolean "false" or "true". FLOAT64: Encoded as a number, or string "NaN", "Infinity" or "-Infinity". STRING: Encoded as a string value. BYTES: Encoded as a base64 string per RFC 4648, section 4. TIMESTAMP: Encoded as an RFC 3339 timestamp with mandatory "Z" time zone string: 1985-04-12T23:20:50.52Z DATE: Encoded as RFC 3339 full-date format string: 1985-04-12 TIME: Encoded as RFC 3339 partial-time format string: 23:20:50.52 DATETIME: Encoded as RFC 3339 full-date "T" partial-time: 1985-04-12T23:20:50.52 GEOGRAPHY: Encoded as WKT NUMERIC: Encoded as a decimal string. BIGNUMERIC: Encoded as a decimal string. ARRAY: Encoded as a list with types matching Type.array_type. STRUCT: Encoded as a list with fields of type Type.struct_type[i]. List is used because a JSON object cannot have duplicate field names. """ TYPE_KIND_UNSPECIFIED = 0 INT64 = 1 BOOL = 2 FLOAT64 = 3 STRING = 4 BYTES = 5 TIMESTAMP = 6 DATE = 7 TIME = 8 DATETIME = 9 GEOGRAPHY = 10 NUMERIC = 11 BIGNUMERIC = 12 ARRAY = 13 STRUCT = 14 arrayElementType = _messages.MessageField('StandardSqlDataType', 1) structType = _messages.MessageField('StandardSqlStructType', 2) typeKind = _messages.EnumField('TypeKindValueValuesEnum', 3) class StandardSqlField(_messages.Message): r"""A field or a column. Fields: name: Optional. The name of this field. Can be absent for struct fields. type: Optional. The type of this parameter. Absent if not explicitly specified (e.g., CREATE FUNCTION statement can omit the return type; in this case the output parameter does not have this "type" field). """ name = _messages.StringField(1) type = _messages.MessageField('StandardSqlDataType', 2) class StandardSqlStructType(_messages.Message): r"""A StandardSqlStructType object. Fields: fields: A StandardSqlField attribute. """ fields = _messages.MessageField('StandardSqlField', 1, repeated=True) class Streamingbuffer(_messages.Message): r"""A Streamingbuffer object. Fields: estimatedBytes: [Output-only] A lower-bound estimate of the number of bytes currently in the streaming buffer. estimatedRows: [Output-only] A lower-bound estimate of the number of rows currently in the streaming buffer. oldestEntryTime: [Output-only] Contains the timestamp of the oldest entry in the streaming buffer, in milliseconds since the epoch, if the streaming buffer is available. """ estimatedBytes = _messages.IntegerField(1, variant=_messages.Variant.UINT64) estimatedRows = _messages.IntegerField(2, variant=_messages.Variant.UINT64) oldestEntryTime = _messages.IntegerField(3, variant=_messages.Variant.UINT64) class Table(_messages.Message): r"""A Table object. Messages: LabelsValue: The labels associated with this table. You can use these to organize and group your tables. Label keys and values can be no longer than 63 characters, can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. Label values are optional. Label keys must start with a letter and each label in the list must have a different key. Fields: clustering: [Beta] Clustering specification for the table. Must be specified with partitioning, data in the table will be first partitioned and subsequently clustered. creationTime: [Output-only] The time when this table was created, in milliseconds since the epoch. description: [Optional] A user-friendly description of this table. encryptionConfiguration: Custom encryption configuration (e.g., Cloud KMS keys). etag: [Output-only] A hash of the table metadata. Used to ensure there were no concurrent modifications to the resource when attempting an update. Not guaranteed to change when the table contents or the fields numRows, numBytes, numLongTermBytes or lastModifiedTime change. expirationTime: [Optional] The time when this table expires, in milliseconds since the epoch. If not present, the table will persist indefinitely. Expired tables will be deleted and their storage reclaimed. The defaultTableExpirationMs property of the encapsulating dataset can be used to set a default expirationTime on newly created tables. externalDataConfiguration: [Optional] Describes the data format, location, and other properties of a table stored outside of BigQuery. By defining these properties, the data source can then be queried as if it were a standard BigQuery table. friendlyName: [Optional] A descriptive name for this table. id: [Output-only] An opaque ID uniquely identifying the table. kind: [Output-only] The type of the resource. labels: The labels associated with this table. You can use these to organize and group your tables. Label keys and values can be no longer than 63 characters, can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. Label values are optional. Label keys must start with a letter and each label in the list must have a different key. lastModifiedTime: [Output-only] The time when this table was last modified, in milliseconds since the epoch. location: [Output-only] The geographic location where the table resides. This value is inherited from the dataset. materializedView: [Optional] Materialized view definition. model: [Output-only, Beta] Present iff this table represents a ML model. Describes the training information for the model, and it is required to run 'PREDICT' queries. numBytes: [Output-only] The size of this table in bytes, excluding any data in the streaming buffer. numLongTermBytes: [Output-only] The number of bytes in the table that are considered "long-term storage". numPhysicalBytes: [Output-only] [TrustedTester] The physical size of this table in bytes, excluding any data in the streaming buffer. This includes compression and storage used for time travel. numRows: [Output-only] The number of rows of data in this table, excluding any data in the streaming buffer. rangePartitioning: [TrustedTester] Range partitioning specification for this table. Only one of timePartitioning and rangePartitioning should be specified. requirePartitionFilter: [Optional] If set to true, queries over this table require a partition filter that can be used for partition elimination to be specified. schema: [Optional] Describes the schema of this table. selfLink: [Output-only] A URL that can be used to access this resource again. snapshotDefinition: [Output-only] Snapshot definition. streamingBuffer: [Output-only] Contains information regarding this table's streaming buffer, if one is present. This field will be absent if the table is not being streamed to or if there is no data in the streaming buffer. tableReference: [Required] Reference describing the ID of this table. timePartitioning: Time-based partitioning specification for this table. Only one of timePartitioning and rangePartitioning should be specified. type: [Output-only] Describes the table type. The following values are supported: TABLE: A normal BigQuery table. VIEW: A virtual table defined by a SQL query. SNAPSHOT: An immutable, read-only table that is a copy of another table. [TrustedTester] MATERIALIZED_VIEW: SQL query whose result is persisted. EXTERNAL: A table that references data stored in an external storage system, such as Google Cloud Storage. The default value is TABLE. view: [Optional] The view definition. """ @encoding.MapUnrecognizedFields('additionalProperties') class LabelsValue(_messages.Message): r"""The labels associated with this table. You can use these to organize and group your tables. Label keys and values can be no longer than 63 characters, can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. Label values are optional. Label keys must start with a letter and each label in the list must have a different key. Messages: AdditionalProperty: An additional property for a LabelsValue object. Fields: additionalProperties: Additional properties of type LabelsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a LabelsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField( 'AdditionalProperty', 1, repeated=True) clustering = _messages.MessageField('Clustering', 1) creationTime = _messages.IntegerField(2) description = _messages.StringField(3) encryptionConfiguration = _messages.MessageField('EncryptionConfiguration', 4) etag = _messages.StringField(5) expirationTime = _messages.IntegerField(6) externalDataConfiguration = _messages.MessageField( 'ExternalDataConfiguration', 7) friendlyName = _messages.StringField(8) id = _messages.StringField(9) kind = _messages.StringField(10, default='bigquery#table') labels = _messages.MessageField('LabelsValue', 11) lastModifiedTime = _messages.IntegerField( 12, variant=_messages.Variant.UINT64) location = _messages.StringField(13) materializedView = _messages.MessageField('MaterializedViewDefinition', 14) model = _messages.MessageField('ModelDefinition', 15) numBytes = _messages.IntegerField(16) numLongTermBytes = _messages.IntegerField(17) numPhysicalBytes = _messages.IntegerField(18) numRows = _messages.IntegerField(19, variant=_messages.Variant.UINT64) rangePartitioning = _messages.MessageField('RangePartitioning', 20) requirePartitionFilter = _messages.BooleanField(21, default=False) schema = _messages.MessageField('TableSchema', 22) selfLink = _messages.StringField(23) snapshotDefinition = _messages.MessageField('SnapshotDefinition', 24) streamingBuffer = _messages.MessageField('Streamingbuffer', 25) tableReference = _messages.MessageField('TableReference', 26) timePartitioning = _messages.MessageField('TimePartitioning', 27) type = _messages.StringField(28) view = _messages.MessageField('ViewDefinition', 29) class TableCell(_messages.Message): r"""A TableCell object. Fields: v: A extra_types.JsonValue attribute. """ v = _messages.MessageField('extra_types.JsonValue', 1) class TableDataInsertAllRequest(_messages.Message): r"""A TableDataInsertAllRequest object. Messages: RowsValueListEntry: A RowsValueListEntry object. Fields: ignoreUnknownValues: [Optional] Accept rows that contain values that do not match the schema. The unknown values are ignored. Default is false, which treats unknown values as errors. kind: The resource type of the response. rows: The rows to insert. skipInvalidRows: [Optional] Insert all valid rows of a request, even if invalid rows exist. The default value is false, which causes the entire request to fail if any invalid rows exist. templateSuffix: If specified, treats the destination table as a base template, and inserts the rows into an instance table named "{destination}{templateSuffix}". BigQuery will manage creation of the instance table, using the schema of the base template table. See https://cloud.google.com/bigquery/streaming-data-into-bigquery#template- tables for considerations when working with templates tables. """ class RowsValueListEntry(_messages.Message): r"""A RowsValueListEntry object. Fields: insertId: [Optional] A unique ID for each row. BigQuery uses this property to detect duplicate insertion requests on a best-effort basis. json: [Required] A JSON object that contains a row of data. The object's properties and values must match the destination table's schema. """ insertId = _messages.StringField(1) json = _messages.MessageField('JsonObject', 2) ignoreUnknownValues = _messages.BooleanField(1) kind = _messages.StringField(2, default='bigquery#tableDataInsertAllRequest') rows = _messages.MessageField('RowsValueListEntry', 3, repeated=True) skipInvalidRows = _messages.BooleanField(4) templateSuffix = _messages.StringField(5) class TableDataInsertAllResponse(_messages.Message): r"""A TableDataInsertAllResponse object. Messages: InsertErrorsValueListEntry: A InsertErrorsValueListEntry object. Fields: insertErrors: An array of errors for rows that were not inserted. kind: The resource type of the response. """ class InsertErrorsValueListEntry(_messages.Message): r"""A InsertErrorsValueListEntry object. Fields: errors: Error information for the row indicated by the index property. index: The index of the row that error applies to. """ errors = _messages.MessageField('ErrorProto', 1, repeated=True) index = _messages.IntegerField(2, variant=_messages.Variant.UINT32) insertErrors = _messages.MessageField( 'InsertErrorsValueListEntry', 1, repeated=True) kind = _messages.StringField(2, default='bigquery#tableDataInsertAllResponse') class TableDataList(_messages.Message): r"""A TableDataList object. Fields: etag: A hash of this page of results. kind: The resource type of the response. pageToken: A token used for paging results. Providing this token instead of the startIndex parameter can help you retrieve stable results when an underlying table is changing. rows: Rows of results. totalRows: The total number of rows in the complete table. """ etag = _messages.StringField(1) kind = _messages.StringField(2, default='bigquery#tableDataList') pageToken = _messages.StringField(3) rows = _messages.MessageField('TableRow', 4, repeated=True) totalRows = _messages.IntegerField(5) class TableFieldSchema(_messages.Message): r"""A TableFieldSchema object. Messages: CategoriesValue: [Optional] The categories attached to this field, used for field-level access control. PolicyTagsValue: A PolicyTagsValue object. Fields: categories: [Optional] The categories attached to this field, used for field-level access control. description: [Optional] The field description. The maximum length is 1,024 characters. fields: [Optional] Describes the nested schema fields if the type property is set to RECORD. mode: [Optional] The field mode. Possible values include NULLABLE, REQUIRED and REPEATED. The default value is NULLABLE. name: [Required] The field name. The name must contain only letters (a-z, A-Z), numbers (0-9), or underscores (_), and must start with a letter or underscore. The maximum length is 128 characters. policyTags: A PolicyTagsValue attribute. type: [Required] The field data type. Possible values include STRING, BYTES, INTEGER, INT64 (same as INTEGER), FLOAT, FLOAT64 (same as FLOAT), BOOLEAN, BOOL (same as BOOLEAN), TIMESTAMP, DATE, TIME, DATETIME, RECORD (where RECORD indicates that the field contains a nested schema) or STRUCT (same as RECORD). """ class CategoriesValue(_messages.Message): r"""[Optional] The categories attached to this field, used for field-level access control. Fields: names: A list of category resource names. For example, "projects/1/taxonomies/2/categories/3". At most 5 categories are allowed. """ names = _messages.StringField(1, repeated=True) class PolicyTagsValue(_messages.Message): r"""A PolicyTagsValue object. Fields: names: A list of category resource names. For example, "projects/1/location/eu/taxonomies/2/policyTags/3". At most 1 policy tag is allowed. """ names = _messages.StringField(1, repeated=True) categories = _messages.MessageField('CategoriesValue', 1) description = _messages.StringField(2) fields = _messages.MessageField('TableFieldSchema', 3, repeated=True) mode = _messages.StringField(4) name = _messages.StringField(5) policyTags = _messages.MessageField('PolicyTagsValue', 6) type = _messages.StringField(7) class TableList(_messages.Message): r"""A TableList object. Messages: TablesValueListEntry: A TablesValueListEntry object. Fields: etag: A hash of this page of results. kind: The type of list. nextPageToken: A token to request the next page of results. tables: Tables in the requested dataset. totalItems: The total number of tables in the dataset. """ class TablesValueListEntry(_messages.Message): r"""A TablesValueListEntry object. Messages: LabelsValue: The labels associated with this table. You can use these to organize and group your tables. ViewValue: Additional details for a view. Fields: clustering: [Beta] Clustering specification for this table, if configured. creationTime: The time when this table was created, in milliseconds since the epoch. expirationTime: [Optional] The time when this table expires, in milliseconds since the epoch. If not present, the table will persist indefinitely. Expired tables will be deleted and their storage reclaimed. friendlyName: The user-friendly name for this table. id: An opaque ID of the table kind: The resource type. labels: The labels associated with this table. You can use these to organize and group your tables. rangePartitioning: The range partitioning specification for this table, if configured. tableReference: A reference uniquely identifying the table. timePartitioning: The time-based partitioning specification for this table, if configured. type: The type of table. Possible values are: TABLE, VIEW. view: Additional details for a view. """ @encoding.MapUnrecognizedFields('additionalProperties') class LabelsValue(_messages.Message): r"""The labels associated with this table. You can use these to organize and group your tables. Messages: AdditionalProperty: An additional property for a LabelsValue object. Fields: additionalProperties: Additional properties of type LabelsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a LabelsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField( 'AdditionalProperty', 1, repeated=True) class ViewValue(_messages.Message): r"""Additional details for a view. Fields: useLegacySql: True if view is defined in legacy SQL dialect, false if in standard SQL. """ useLegacySql = _messages.BooleanField(1) clustering = _messages.MessageField('Clustering', 1) creationTime = _messages.IntegerField(2) expirationTime = _messages.IntegerField(3) friendlyName = _messages.StringField(4) id = _messages.StringField(5) kind = _messages.StringField(6, default='bigquery#table') labels = _messages.MessageField('LabelsValue', 7) rangePartitioning = _messages.MessageField('RangePartitioning', 8) tableReference = _messages.MessageField('TableReference', 9) timePartitioning = _messages.MessageField('TimePartitioning', 10) type = _messages.StringField(11) view = _messages.MessageField('ViewValue', 12) etag = _messages.StringField(1) kind = _messages.StringField(2, default='bigquery#tableList') nextPageToken = _messages.StringField(3) tables = _messages.MessageField('TablesValueListEntry', 4, repeated=True) totalItems = _messages.IntegerField(5, variant=_messages.Variant.INT32) class TableReference(_messages.Message): r"""A TableReference object. Fields: datasetId: [Required] The ID of the dataset containing this table. projectId: [Required] The ID of the project containing this table. tableId: [Required] The ID of the table. The ID must contain only letters (a-z, A-Z), numbers (0-9), or underscores (_). The maximum length is 1,024 characters. """ datasetId = _messages.StringField(1) projectId = _messages.StringField(2) tableId = _messages.StringField(3) class TableRow(_messages.Message): r"""A TableRow object. Fields: f: Represents a single row in the result set, consisting of one or more fields. """ f = _messages.MessageField('TableCell', 1, repeated=True) class TableSchema(_messages.Message): r"""A TableSchema object. Fields: fields: Describes the fields in a table. """ fields = _messages.MessageField('TableFieldSchema', 1, repeated=True) class TestIamPermissionsRequest(_messages.Message): r"""Request message for `TestIamPermissions` method. Fields: permissions: The set of permissions to check for the `resource`. Permissions with wildcards (such as '*' or 'storage.*') are not allowed. For more information see [IAM Overview](https://cloud.google.com/iam/docs/overview#permissions). """ permissions = _messages.StringField(1, repeated=True) class TestIamPermissionsResponse(_messages.Message): r"""Response message for `TestIamPermissions` method. Fields: permissions: A subset of `TestPermissionsRequest.permissions` that the caller is allowed. """ permissions = _messages.StringField(1, repeated=True) class TimePartitioning(_messages.Message): r"""A TimePartitioning object. Fields: expirationMs: [Optional] Number of milliseconds for which to keep the storage for partitions in the table. The storage in a partition will have an expiration time of its partition time plus this value. field: [Beta] [Optional] If not set, the table is partitioned by pseudo column, referenced via either '_PARTITIONTIME' as TIMESTAMP type, or '_PARTITIONDATE' as DATE type. If field is specified, the table is instead partitioned by this field. The field must be a top-level TIMESTAMP or DATE field. Its mode must be NULLABLE or REQUIRED. requirePartitionFilter: A boolean attribute. type: [Required] The supported types are DAY, HOUR, MONTH, and YEAR, which will generate one partition per day, hour, month, and year, respectively. When the type is not specified, the default behavior is DAY. """ expirationMs = _messages.IntegerField(1) field = _messages.StringField(2) requirePartitionFilter = _messages.BooleanField(3) type = _messages.StringField(4) class TrainingOptions(_messages.Message): r"""A TrainingOptions object. Enums: DataFrequencyValueValuesEnum: The data frequency of a time series. DataSplitMethodValueValuesEnum: The data split type for training and evaluation, e.g. RANDOM. DistanceTypeValueValuesEnum: Distance type for clustering models. FeedbackTypeValueValuesEnum: Feedback type that specifies which algorithm to run for matrix factorization. HolidayRegionValueValuesEnum: The geographical region based on which the holidays are considered in time series modeling. If a valid value is specified, then holiday effects modeling is enabled. KmeansInitializationMethodValueValuesEnum: The method used to initialize the centroids for kmeans algorithm. LearnRateStrategyValueValuesEnum: The strategy to determine learn rate for the current iteration. LossTypeValueValuesEnum: Type of loss function used during training run. OptimizationStrategyValueValuesEnum: Optimization strategy for training linear regression models. Messages: LabelClassWeightsValue: Weights associated with each label class, for rebalancing the training data. Only applicable for classification models. Fields: autoArima: Whether to enable auto ARIMA or not. autoArimaMaxOrder: The max value of non-seasonal p and q. batchSize: Batch size for dnn models. dataFrequency: The data frequency of a time series. dataSplitColumn: The column to split data with. This column won't be used as a feature. 1. When data_split_method is CUSTOM, the corresponding column should be boolean. The rows with true value tag are eval data, and the false are training data. 2. When data_split_method is SEQ, the first DATA_SPLIT_EVAL_FRACTION rows (from smallest to largest) in the corresponding column are used as training data, and the rest are eval data. It respects the order in Orderable data types: https://cloud.google.com/bigquery/docs/reference/standard-sql/data- types#data-type-properties dataSplitEvalFraction: The fraction of evaluation data over the whole input data. The rest of data will be used as training data. The format should be double. Accurate to two decimal places. Default value is 0.2. dataSplitMethod: The data split type for training and evaluation, e.g. RANDOM. distanceType: Distance type for clustering models. dropout: Dropout probability for dnn models. earlyStop: Whether to stop early when the loss doesn't improve significantly any more (compared to min_relative_progress). Used only for iterative training algorithms. feedbackType: Feedback type that specifies which algorithm to run for matrix factorization. hiddenUnits: Hidden units for dnn models. holidayRegion: The geographical region based on which the holidays are considered in time series modeling. If a valid value is specified, then holiday effects modeling is enabled. horizon: The number of periods ahead that need to be forecasted. includeDrift: Include drift when fitting an ARIMA model. initialLearnRate: Specifies the initial learning rate for the line search learn rate strategy. inputLabelColumns: Name of input label columns in training data. itemColumn: Item column specified for matrix factorization models. kmeansInitializationColumn: The column used to provide the initial centroids for kmeans algorithm when kmeans_initialization_method is CUSTOM. kmeansInitializationMethod: The method used to initialize the centroids for kmeans algorithm. l1Regularization: L1 regularization coefficient. l2Regularization: L2 regularization coefficient. labelClassWeights: Weights associated with each label class, for rebalancing the training data. Only applicable for classification models. learnRate: Learning rate in training. Used only for iterative training algorithms. learnRateStrategy: The strategy to determine learn rate for the current iteration. lossType: Type of loss function used during training run. maxIterations: The maximum number of iterations in training. Used only for iterative training algorithms. maxTreeDepth: Maximum depth of a tree for boosted tree models. minRelativeProgress: When early_stop is true, stops training when accuracy improvement is less than 'min_relative_progress'. Used only for iterative training algorithms. minSplitLoss: Minimum split loss for boosted tree models. modelUri: [Beta] Google Cloud Storage URI from which the model was imported. Only applicable for imported models. nonSeasonalOrder: A specification of the non-seasonal part of the ARIMA model: the three components (p, d, q) are the AR order, the degree of differencing, and the MA order. numClusters: Number of clusters for clustering models. numFactors: Num factors specified for matrix factorization models. optimizationStrategy: Optimization strategy for training linear regression models. preserveInputStructs: Whether to preserve the input structs in output feature names. Suppose there is a struct A with field b. When false (default), the output feature name is A_b. When true, the output feature name is A.b. subsample: Subsample fraction of the training data to grow tree to prevent overfitting for boosted tree models. timeSeriesDataColumn: Column to be designated as time series data for ARIMA model. timeSeriesIdColumn: The id column that will be used to indicate different time series to forecast in parallel. timeSeriesTimestampColumn: Column to be designated as time series timestamp for ARIMA model. userColumn: User column specified for matrix factorization models. walsAlpha: Hyperparameter for matrix factoration when implicit feedback type is specified. warmStart: Whether to train a model from the last checkpoint. """ class DataFrequencyValueValuesEnum(_messages.Enum): r"""The data frequency of a time series. Values: DATA_FREQUENCY_UNSPECIFIED: AUTO_FREQUENCY: Automatically inferred from timestamps. YEARLY: Yearly data. QUARTERLY: Quarterly data. MONTHLY: Monthly data. WEEKLY: Weekly data. DAILY: Daily data. HOURLY: Hourly data. """ DATA_FREQUENCY_UNSPECIFIED = 0 AUTO_FREQUENCY = 1 YEARLY = 2 QUARTERLY = 3 MONTHLY = 4 WEEKLY = 5 DAILY = 6 HOURLY = 7 class DataSplitMethodValueValuesEnum(_messages.Enum): r"""The data split type for training and evaluation, e.g. RANDOM. Values: DATA_SPLIT_METHOD_UNSPECIFIED: RANDOM: Splits data randomly. CUSTOM: Splits data with the user provided tags. SEQUENTIAL: Splits data sequentially. NO_SPLIT: Data split will be skipped. AUTO_SPLIT: Splits data automatically: Uses NO_SPLIT if the data size is small. Otherwise uses RANDOM. """ DATA_SPLIT_METHOD_UNSPECIFIED = 0 RANDOM = 1 CUSTOM = 2 SEQUENTIAL = 3 NO_SPLIT = 4 AUTO_SPLIT = 5 class DistanceTypeValueValuesEnum(_messages.Enum): r"""Distance type for clustering models. Values: DISTANCE_TYPE_UNSPECIFIED: EUCLIDEAN: Eculidean distance. COSINE: Cosine distance. """ DISTANCE_TYPE_UNSPECIFIED = 0 EUCLIDEAN = 1 COSINE = 2 class FeedbackTypeValueValuesEnum(_messages.Enum): r"""Feedback type that specifies which algorithm to run for matrix factorization. Values: FEEDBACK_TYPE_UNSPECIFIED: IMPLICIT: Use weighted-als for implicit feedback problems. EXPLICIT: Use nonweighted-als for explicit feedback problems. """ FEEDBACK_TYPE_UNSPECIFIED = 0 IMPLICIT = 1 EXPLICIT = 2 class HolidayRegionValueValuesEnum(_messages.Enum): r"""The geographical region based on which the holidays are considered in time series modeling. If a valid value is specified, then holiday effects modeling is enabled. Values: HOLIDAY_REGION_UNSPECIFIED: Holiday region unspecified. GLOBAL: Global. NA: North America. JAPAC: Japan and Asia Pacific: Korea, Greater China, India, Australia, and New Zealand. EMEA: Europe, the Middle East and Africa. LAC: Latin America and the Caribbean. AE: United Arab Emirates AR: Argentina AT: Austria AU: Australia BE: Belgium BR: Brazil CA: Canada CH: Switzerland CL: Chile CN: China CO: Colombia CS: Czechoslovakia CZ: Czech Republic DE: Germany DK: Denmark DZ: Algeria EC: Ecuador EE: Estonia EG: Egypt ES: Spain FI: Finland FR: France GB: Great Britain (United Kingdom) GR: Greece HK: Hong Kong HU: Hungary ID: Indonesia IE: Ireland IL: Israel IN: India IR: Iran IT: Italy JP: Japan KR: Korea (South) LV: Latvia MA: Morocco MX: Mexico MY: Malaysia NG: Nigeria NL: Netherlands NO: Norway NZ: New Zealand PE: Peru PH: Philippines PK: Pakistan PL: Poland PT: Portugal RO: Romania RS: Serbia RU: Russian Federation SA: Saudi Arabia SE: Sweden SG: Singapore SI: Slovenia SK: Slovakia TH: Thailand TR: Turkey TW: Taiwan UA: Ukraine US: United States VE: Venezuela VN: Viet Nam ZA: South Africa """ HOLIDAY_REGION_UNSPECIFIED = 0 GLOBAL = 1 NA = 2 JAPAC = 3 EMEA = 4 LAC = 5 AE = 6 AR = 7 AT = 8 AU = 9 BE = 10 BR = 11 CA = 12 CH = 13 CL = 14 CN = 15 CO = 16 CS = 17 CZ = 18 DE = 19 DK = 20 DZ = 21 EC = 22 EE = 23 EG = 24 ES = 25 FI = 26 FR = 27 GB = 28 GR = 29 HK = 30 HU = 31 ID = 32 IE = 33 IL = 34 IN = 35 IR = 36 IT = 37 JP = 38 KR = 39 LV = 40 MA = 41 MX = 42 MY = 43 NG = 44 NL = 45 NO = 46 NZ = 47 PE = 48 PH = 49 PK = 50 PL = 51 PT = 52 RO = 53 RS = 54 RU = 55 SA = 56 SE = 57 SG = 58 SI = 59 SK = 60 TH = 61 TR = 62 TW = 63 UA = 64 US = 65 VE = 66 VN = 67 ZA = 68 class KmeansInitializationMethodValueValuesEnum(_messages.Enum): r"""The method used to initialize the centroids for kmeans algorithm. Values: KMEANS_INITIALIZATION_METHOD_UNSPECIFIED: RANDOM: Initializes the centroids randomly. CUSTOM: Initializes the centroids using data specified in kmeans_initialization_column. KMEANS_PLUS_PLUS: Initializes with kmeans++. """ KMEANS_INITIALIZATION_METHOD_UNSPECIFIED = 0 RANDOM = 1 CUSTOM = 2 KMEANS_PLUS_PLUS = 3 class LearnRateStrategyValueValuesEnum(_messages.Enum): r"""The strategy to determine learn rate for the current iteration. Values: LEARN_RATE_STRATEGY_UNSPECIFIED: LINE_SEARCH: Use line search to determine learning rate. CONSTANT: Use a constant learning rate. """ LEARN_RATE_STRATEGY_UNSPECIFIED = 0 LINE_SEARCH = 1 CONSTANT = 2 class LossTypeValueValuesEnum(_messages.Enum): r"""Type of loss function used during training run. Values: LOSS_TYPE_UNSPECIFIED: MEAN_SQUARED_LOSS: Mean squared loss, used for linear regression. MEAN_LOG_LOSS: Mean log loss, used for logistic regression. """ LOSS_TYPE_UNSPECIFIED = 0 MEAN_SQUARED_LOSS = 1 MEAN_LOG_LOSS = 2 class OptimizationStrategyValueValuesEnum(_messages.Enum): r"""Optimization strategy for training linear regression models. Values: OPTIMIZATION_STRATEGY_UNSPECIFIED: BATCH_GRADIENT_DESCENT: Uses an iterative batch gradient descent algorithm. NORMAL_EQUATION: Uses a normal equation to solve linear regression problem. """ OPTIMIZATION_STRATEGY_UNSPECIFIED = 0 BATCH_GRADIENT_DESCENT = 1 NORMAL_EQUATION = 2 @encoding.MapUnrecognizedFields('additionalProperties') class LabelClassWeightsValue(_messages.Message): r"""Weights associated with each label class, for rebalancing the training data. Only applicable for classification models. Messages: AdditionalProperty: An additional property for a LabelClassWeightsValue object. Fields: additionalProperties: Additional properties of type LabelClassWeightsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a LabelClassWeightsValue object. Fields: key: Name of the additional property. value: A number attribute. """ key = _messages.StringField(1) value = _messages.FloatField(2) additionalProperties = _messages.MessageField( 'AdditionalProperty', 1, repeated=True) autoArima = _messages.BooleanField(1) autoArimaMaxOrder = _messages.IntegerField(2) batchSize = _messages.IntegerField(3) dataFrequency = _messages.EnumField('DataFrequencyValueValuesEnum', 4) dataSplitColumn = _messages.StringField(5) dataSplitEvalFraction = _messages.FloatField(6) dataSplitMethod = _messages.EnumField('DataSplitMethodValueValuesEnum', 7) distanceType = _messages.EnumField('DistanceTypeValueValuesEnum', 8) dropout = _messages.FloatField(9) earlyStop = _messages.BooleanField(10) feedbackType = _messages.EnumField('FeedbackTypeValueValuesEnum', 11) hiddenUnits = _messages.IntegerField(12, repeated=True) holidayRegion = _messages.EnumField('HolidayRegionValueValuesEnum', 13) horizon = _messages.IntegerField(14) includeDrift = _messages.BooleanField(15) initialLearnRate = _messages.FloatField(16) inputLabelColumns = _messages.StringField(17, repeated=True) itemColumn = _messages.StringField(18) kmeansInitializationColumn = _messages.StringField(19) kmeansInitializationMethod = _messages.EnumField( 'KmeansInitializationMethodValueValuesEnum', 20) l1Regularization = _messages.FloatField(21) l2Regularization = _messages.FloatField(22) labelClassWeights = _messages.MessageField('LabelClassWeightsValue', 23) learnRate = _messages.FloatField(24) learnRateStrategy = _messages.EnumField( 'LearnRateStrategyValueValuesEnum', 25) lossType = _messages.EnumField('LossTypeValueValuesEnum', 26) maxIterations = _messages.IntegerField(27) maxTreeDepth = _messages.IntegerField(28) minRelativeProgress = _messages.FloatField(29) minSplitLoss = _messages.FloatField(30) modelUri = _messages.StringField(31) nonSeasonalOrder = _messages.MessageField('ArimaOrder', 32) numClusters = _messages.IntegerField(33) numFactors = _messages.IntegerField(34) optimizationStrategy = _messages.EnumField( 'OptimizationStrategyValueValuesEnum', 35) preserveInputStructs = _messages.BooleanField(36) subsample = _messages.FloatField(37) timeSeriesDataColumn = _messages.StringField(38) timeSeriesIdColumn = _messages.StringField(39) timeSeriesTimestampColumn = _messages.StringField(40) userColumn = _messages.StringField(41) walsAlpha = _messages.FloatField(42) warmStart = _messages.BooleanField(43) class TrainingRun(_messages.Message): r"""Information about a single training query run for the model. Fields: dataSplitResult: Data split result of the training run. Only set when the input data is actually split. evaluationMetrics: The evaluation metrics over training/eval data that were computed at the end of training. globalExplanations: Global explanations for important features of the model. For multi-class models, there is one entry for each label class. For other models, there is only one entry in the list. results: Output of each iteration run, results.size() <= max_iterations. startTime: The start time of this training run. trainingOptions: Options that were used for this training run, includes user specified and default options that were used. """ dataSplitResult = _messages.MessageField('DataSplitResult', 1) evaluationMetrics = _messages.MessageField('EvaluationMetrics', 2) globalExplanations = _messages.MessageField( 'GlobalExplanation', 3, repeated=True) results = _messages.MessageField('IterationResult', 4, repeated=True) startTime = _messages.StringField(5) trainingOptions = _messages.MessageField('TrainingOptions', 6) class TransactionInfo(_messages.Message): r"""A TransactionInfo object. Fields: transactionId: [Output-only] // [Alpha] Id of the transaction. """ transactionId = _messages.StringField(1) class UserDefinedFunctionResource(_messages.Message): r"""This is used for defining User Defined Function (UDF) resources only when using legacy SQL. Users of Standard SQL should leverage either DDL (e.g. CREATE [TEMPORARY] FUNCTION ... ) or the Routines API to define UDF resources. For additional information on migrating, see: https://cloud.google.com/bigquery/docs/reference/standard-sql/migrating- from-legacy-sql#differences_in_user-defined_javascript_functions Fields: inlineCode: [Pick one] An inline resource that contains code for a user- defined function (UDF). Providing a inline code resource is equivalent to providing a URI for a file containing the same code. resourceUri: [Pick one] A code resource to load from a Google Cloud Storage URI (gs://bucket/path). """ inlineCode = _messages.StringField(1) resourceUri = _messages.StringField(2) class ViewDefinition(_messages.Message): r"""A ViewDefinition object. Fields: query: [Required] A query that BigQuery executes when the view is referenced. useLegacySql: Specifies whether to use BigQuery's legacy SQL for this view. The default value is true. If set to false, the view will use BigQuery's standard SQL: https://cloud.google.com/bigquery/sql- reference/ Queries and views that reference this view must use the same flag value. userDefinedFunctionResources: Describes user-defined function resources used in the query. """ query = _messages.StringField(1) useLegacySql = _messages.BooleanField(2) userDefinedFunctionResources = _messages.MessageField( 'UserDefinedFunctionResource', 3, repeated=True) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/pubsub.py0000644000175100017510000007134315206113032020717 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Google Cloud PubSub sources and sinks. Cloud Pub/Sub sources are currently supported only in streaming pipelines, during remote execution. Cloud Pub/Sub sinks (WriteToPubSub) support both streaming and batch pipelines. This API is currently under development and is subject to change. **Updates to the I/O connector code** For any significant updates to this I/O connector, please consider involving corresponding code reviewers mentioned in https://github.com/apache/beam/blob/master/sdks/python/OWNERS """ # pytype: skip-file import re from typing import Any from typing import NamedTuple from typing import Optional from typing import Union from apache_beam import coders from apache_beam.io import iobase from apache_beam.io.iobase import Read from apache_beam.metrics.metric import Lineage from apache_beam.transforms import DoFn from apache_beam.transforms import Flatten from apache_beam.transforms import Map from apache_beam.transforms import ParDo from apache_beam.transforms import PTransform from apache_beam.transforms.display import DisplayDataItem from apache_beam.utils.annotations import deprecated try: from google.cloud import pubsub except ImportError: pubsub = None __all__ = [ 'MultipleReadFromPubSub', 'PubsubMessage', 'PubSubSourceDescriptor', 'ReadFromPubSub', 'ReadStringsFromPubSub', 'WriteStringsToPubSub', 'WriteToPubSub' ] class PubsubMessage(object): """Represents a Cloud Pub/Sub message. Message payload includes the data and attributes fields. For the payload to be valid, at least one of its fields must be non-empty. Attributes: data: (bytes) Message data. May be None. attributes: (dict) Key-value map of str to str, containing both user-defined and service generated attributes (such as id_label and timestamp_attribute). May be None. message_id: (str) ID of the message, assigned by the pubsub service when the message is published. Guaranteed to be unique within the topic. Will be reset to None if the message is being written to pubsub. publish_time: (datetime) Time at which the message was published. Will be reset to None if the Message is being written to pubsub. ordering_key: (str) If non-empty, identifies related messages for which publish order is respected by the PubSub subscription. """ def __init__( self, data, attributes, message_id=None, publish_time=None, ordering_key=""): if data is None and not attributes: raise ValueError( 'Either data (%r) or attributes (%r) must be set.', data, attributes) self.data = data self.attributes = attributes self.message_id = message_id self.publish_time = publish_time self.ordering_key = ordering_key def __hash__(self): return hash((self.data, frozenset(self.attributes.items()))) def __eq__(self, other): return isinstance(other, PubsubMessage) and ( self.data == other.data and self.attributes == other.attributes) def __repr__(self): return 'PubsubMessage(%s, %s)' % (self.data, self.attributes) @staticmethod def _from_proto_str(proto_msg: bytes) -> 'PubsubMessage': """Construct from serialized form of ``PubsubMessage``. Args: proto_msg: String containing a serialized protobuf of type https://cloud.google.com/pubsub/docs/reference/rpc/google.pubsub.v1#google.pubsub.v1.PubsubMessage Returns: A new PubsubMessage object. """ msg = pubsub.types.PubsubMessage.deserialize(proto_msg) # Convert ScalarMapContainer to dict. attributes = dict(msg.attributes) return PubsubMessage( msg.data, attributes, msg.message_id, msg.publish_time, msg.ordering_key) def _to_proto_str(self, for_publish=False): """Get serialized form of ``PubsubMessage``. The serialized message is validated against pubsub message limits specified at https://cloud.google.com/pubsub/quotas#resource_limits Args: proto_msg: str containing a serialized protobuf. for_publish: bool, if True strip out message fields which cannot be published (currently message_id and publish_time) per https://cloud.google.com/pubsub/docs/reference/rpc/google.pubsub.v1#pubsubmessage Returns: A str containing a serialized protobuf of type https://cloud.google.com/pubsub/docs/reference/rpc/google.pubsub.v1#google.pubsub.v1.PubsubMessage containing the payload of this object. """ if len(self.data) > (10_000_000): raise ValueError('A pubsub message data field must not exceed 10MB') if self.attributes: if len(self.attributes) > 100: raise ValueError( 'A pubsub message must not have more than 100 attributes.') for key, value in self.attributes.items(): if len(key) > 256: raise ValueError( 'A pubsub message attribute key must not exceed 256 bytes.') if len(value) > 1024: raise ValueError( 'A pubsub message attribute value must not exceed 1024 bytes') message_id = None publish_time = None if not for_publish: if self.message_id: message_id = self.message_id if self.publish_time: publish_time = self.publish_time if len(self.ordering_key) > 1024: raise ValueError( 'A pubsub message ordering key must not exceed 1024 bytes.') msg = pubsub.types.PubsubMessage( data=self.data, attributes=self.attributes, message_id=message_id, publish_time=publish_time, ordering_key=self.ordering_key) serialized = pubsub.types.PubsubMessage.serialize(msg) if len(serialized) > (10_000_000): raise ValueError( 'Serialized pubsub message exceeds the publish request limit of 10MB') return serialized @staticmethod def _from_message(msg: Any) -> 'PubsubMessage': """Construct from ``google.cloud.pubsub_v1.subscriber.message.Message``. https://googleapis.github.io/google-cloud-python/latest/pubsub/subscriber/api/message.html """ # Convert ScalarMapContainer to dict. attributes = dict(msg.attributes) pubsubmessage = PubsubMessage(msg.data, attributes) if msg.message_id: pubsubmessage.message_id = msg.message_id if msg.publish_time: pubsubmessage.publish_time = msg.publish_time if msg.ordering_key: pubsubmessage.ordering_key = msg.ordering_key return pubsubmessage class ReadFromPubSub(PTransform): """A ``PTransform`` for reading from Cloud Pub/Sub.""" # Implementation note: This ``PTransform`` is overridden by Directrunner. def __init__( self, topic: Optional[str] = None, subscription: Optional[str] = None, id_label: Optional[str] = None, with_attributes: bool = False, timestamp_attribute: Optional[str] = None) -> None: """Initializes ``ReadFromPubSub``. Args: topic: Cloud Pub/Sub topic in the form "projects//topics/". If provided, subscription must be None. subscription: Existing Cloud Pub/Sub subscription to use in the form "projects//subscriptions/". If not specified, a temporary subscription will be created from the specified topic. If provided, topic must be None. id_label: The attribute on incoming Pub/Sub messages to use as a unique record identifier. When specified, the value of this attribute (which can be any string that uniquely identifies the record) will be used for deduplication of messages. If not provided, we cannot guarantee that no duplicate data will be delivered on the Pub/Sub stream. In this case, deduplication of the stream will be strictly best effort. with_attributes: True - output elements will be :class:`~PubsubMessage` objects. False - output elements will be of type ``bytes`` (message data only). timestamp_attribute: Message value to use as element timestamp. If None, uses message publishing time as the timestamp. Timestamp values should be in one of two formats: - A numerical value representing the number of milliseconds since the Unix epoch. - A string in RFC 3339 format, UTC timezone. Example: ``2015-10-29T23:41:41.123Z``. The sub-second component of the timestamp is optional, and digits beyond the first three (i.e., time units smaller than milliseconds) may be ignored. """ super().__init__() self.with_attributes = with_attributes self._source = _PubSubSource( topic=topic, subscription=subscription, id_label=id_label, with_attributes=self.with_attributes, timestamp_attribute=timestamp_attribute) def expand(self, pvalue): # TODO(BEAM-27443): Apply a proper transform rather than Read. pcoll = pvalue.pipeline | Read(self._source) # explicit element_type required after native read, otherwise coder error pcoll.element_type = bytes return self.expand_continued(pcoll) def expand_continued(self, pcoll): pcoll = pcoll | ParDo( _AddMetricsPassThrough( project=self._source.project, topic=self._source.topic_name, sub=self._source.subscription_name)).with_output_types(bytes) if self.with_attributes: pcoll = pcoll | Map(PubsubMessage._from_proto_str) pcoll.element_type = PubsubMessage return pcoll def to_runner_api_parameter(self, context): # Required as this is identified by type in PTransformOverrides. # TODO(https://github.com/apache/beam/issues/18713): Use an actual URN here. return self.to_runner_api_pickled(context) class _AddMetricsPassThrough(DoFn): def __init__(self, project, topic=None, sub=None): self.project = project self.topic = topic self.sub = sub self.reported_lineage = False def setup(self): self.reported_lineage = False def process(self, element: bytes): self.report_lineage_once() yield element def report_lineage_once(self): if not self.reported_lineage: self.reported_lineage = True if self.topic is not None: Lineage.sources().add( 'pubsub', self.project, self.topic, subtype='topic') elif self.sub is not None: Lineage.sources().add( 'pubsub', self.project, self.sub, subtype='subscription') @deprecated(since='2.7.0', extra_message='Use ReadFromPubSub instead.') def ReadStringsFromPubSub(topic=None, subscription=None, id_label=None): return _ReadStringsFromPubSub(topic, subscription, id_label) class _ReadStringsFromPubSub(PTransform): """This class is deprecated. Use ``ReadFromPubSub`` instead.""" def __init__(self, topic=None, subscription=None, id_label=None): super().__init__() self.topic = topic self.subscription = subscription self.id_label = id_label def expand(self, pvalue): p = ( pvalue.pipeline | ReadFromPubSub( self.topic, self.subscription, self.id_label, with_attributes=False) | 'DecodeString' >> Map(lambda b: b.decode('utf-8'))) p.element_type = str return p @deprecated(since='2.7.0', extra_message='Use WriteToPubSub instead.') def WriteStringsToPubSub(topic): return _WriteStringsToPubSub(topic) class _WriteStringsToPubSub(PTransform): """This class is deprecated. Use ``WriteToPubSub`` instead.""" def __init__(self, topic): """Initializes ``_WriteStringsToPubSub``. Attributes: topic: Cloud Pub/Sub topic in the form "/topics//". """ super().__init__() self.topic = topic def expand(self, pcoll): pcoll = pcoll | 'EncodeString' >> Map(lambda s: s.encode('utf-8')) pcoll.element_type = bytes return pcoll | WriteToPubSub(self.topic) class _AddMetricsAndMap(DoFn): def __init__(self, fn, project, topic=None): self.project = project self.topic = topic self.fn = fn self.reported_lineage = False def setup(self): self.reported_lineage = False def process(self, element): self.report_lineage_once() yield self.fn(element) def report_lineage_once(self): if not self.reported_lineage: self.reported_lineage = True Lineage.sinks().add('pubsub', self.project, self.topic, subtype='topic') class WriteToPubSub(PTransform): """A ``PTransform`` for writing messages to Cloud Pub/Sub. This transform supports both streaming and batch pipelines. In streaming mode, messages are written continuously as they arrive. In batch mode, all messages are written when the pipeline completes. """ # Implementation note: This ``PTransform`` is overridden by Directrunner. def __init__( self, topic: str, with_attributes: bool = False, id_label: Optional[str] = None, timestamp_attribute: Optional[str] = None) -> None: """Initializes ``WriteToPubSub``. Args: topic: Cloud Pub/Sub topic in the form "projects//topics/". with_attributes: True - input elements will be :class:`~PubsubMessage` objects. False - input elements will be of type ``bytes`` (message data only). id_label: If set, will set an attribute for each Cloud Pub/Sub message with the given name and a unique value. This attribute can then be used in a ReadFromPubSub PTransform to deduplicate messages. timestamp_attribute: If set, will set an attribute for each Cloud Pub/Sub message with the given name and the message's publish time as the value. """ super().__init__() self.with_attributes = with_attributes self.id_label = id_label self.timestamp_attribute = timestamp_attribute self.project, self.topic_name = parse_topic(topic) self.full_topic = topic self._sink = _PubSubSink(topic, id_label, timestamp_attribute) self.pipeline_options = None # Will be set during expand() @staticmethod def message_to_proto_str(element: PubsubMessage) -> bytes: if not isinstance(element, PubsubMessage): raise TypeError( 'Unexpected element. Type: %s (expected: PubsubMessage), ' 'value: %r' % (type(element), element)) return element._to_proto_str(for_publish=True) @staticmethod def bytes_to_proto_str(element: Union[bytes, str]) -> bytes: msg = PubsubMessage(element, {}) return msg._to_proto_str(for_publish=True) def expand(self, pcoll): # Store pipeline options for use in DoFn self.pipeline_options = pcoll.pipeline.options if pcoll.pipeline else None if self.with_attributes: pcoll = pcoll | 'ToProtobufX' >> ParDo( _AddMetricsAndMap( self.message_to_proto_str, self.project, self.topic_name)).with_input_types(PubsubMessage) else: pcoll = pcoll | 'ToProtobufY' >> ParDo( _AddMetricsAndMap( self.bytes_to_proto_str, self.project, self.topic_name)).with_input_types(Union[bytes, str]) pcoll.element_type = bytes return pcoll | ParDo(_PubSubWriteDoFn(self)) def to_runner_api_parameter(self, context): # Required as this is identified by type in PTransformOverrides. # TODO(https://github.com/apache/beam/issues/18713): Use an actual URN here. return self.to_runner_api_pickled(context) def display_data(self): return { 'topic': DisplayDataItem(self.full_topic, label='Pubsub Topic'), 'id_label': DisplayDataItem(self.id_label, label='ID Label Attribute'), 'with_attributes': DisplayDataItem( True, label='With Attributes').drop_if_none(), 'timestamp_attribute': DisplayDataItem( self.timestamp_attribute, label='Timestamp Attribute'), } PROJECT_ID_REGEXP = '[a-z][-a-z0-9:.]{4,61}[a-z0-9]' SUBSCRIPTION_REGEXP = 'projects/([^/]+)/subscriptions/(.+)' TOPIC_REGEXP = 'projects/([^/]+)/topics/(.+)' def parse_topic(full_topic: str) -> tuple[str, str]: match = re.match(TOPIC_REGEXP, full_topic) if not match: raise ValueError( 'PubSub topic must be in the form "projects//topics' '/" (got %r).' % full_topic) project, topic_name = match.group(1), match.group(2) if not re.match(PROJECT_ID_REGEXP, project): raise ValueError('Invalid PubSub project name: %r.' % project) return project, topic_name def parse_subscription(full_subscription): match = re.match(SUBSCRIPTION_REGEXP, full_subscription) if not match: raise ValueError( 'PubSub subscription must be in the form "projects/' '/subscriptions/" (got %r).' % full_subscription) project, subscription_name = match.group(1), match.group(2) if not re.match(PROJECT_ID_REGEXP, project): raise ValueError('Invalid PubSub project name: %r.' % project) return project, subscription_name # TODO(BEAM-27443): Remove (or repurpose as a proper PTransform). class _PubSubSource(iobase.SourceBase): """Source for a Cloud Pub/Sub topic or subscription. This ``NativeSource`` is overridden by a native Pubsub implementation. Attributes: with_attributes: If False, will fetch just message data. Otherwise, fetches ``PubsubMessage`` protobufs. """ def __init__( self, topic: Optional[str] = None, subscription: Optional[str] = None, id_label: Optional[str] = None, with_attributes: bool = False, timestamp_attribute: Optional[str] = None): self.coder = coders.BytesCoder() self.full_topic = topic self.full_subscription = subscription self.topic_name = None self.subscription_name = None self.id_label = id_label self.with_attributes = with_attributes self.timestamp_attribute = timestamp_attribute # Perform some validation on the topic and subscription. if not (topic or subscription): raise ValueError('Either a topic or subscription must be provided.') if topic and subscription: raise ValueError('Only one of topic or subscription should be provided.') if topic: self.project, self.topic_name = parse_topic(topic) if subscription: self.project, self.subscription_name = parse_subscription(subscription) def display_data(self): return { 'id_label': DisplayDataItem(self.id_label, label='ID Label Attribute').drop_if_none(), 'topic': DisplayDataItem(self.full_topic, label='Pubsub Topic').drop_if_none(), 'subscription': DisplayDataItem( self.full_subscription, label='Pubsub Subscription').drop_if_none(), 'with_attributes': DisplayDataItem( self.with_attributes, label='With Attributes').drop_if_none(), 'timestamp_attribute': DisplayDataItem( self.timestamp_attribute, label='Timestamp Attribute').drop_if_none(), } def default_output_coder(self): return self.coder def is_bounded(self): return False class _PubSubWriteDoFn(DoFn): """DoFn for writing messages to Cloud Pub/Sub. This DoFn handles both streaming and batch modes by buffering messages and publishing them in batches to optimize performance. """ BUFFER_SIZE_ELEMENTS = 100 FLUSH_TIMEOUT_SECS = 5 * 60 # 5 minutes def __init__(self, transform): self.project = transform.project self.short_topic_name = transform.topic_name self.id_label = transform.id_label self.timestamp_attribute = transform.timestamp_attribute self.with_attributes = transform.with_attributes # TODO(https://github.com/apache/beam/issues/18939): Add support for # id_label and timestamp_attribute. # Only raise errors for DirectRunner or batch pipelines pipeline_options = transform.pipeline_options output_labels_supported = True if pipeline_options: from apache_beam.options.pipeline_options import StandardOptions # Check if using DirectRunner try: # Get runner from pipeline options all_options = pipeline_options.get_all_options() runner_name = all_options.get('runner', StandardOptions.DEFAULT_RUNNER) # Check if it's a DirectRunner variant if (runner_name is None or (runner_name in StandardOptions.LOCAL_RUNNERS or 'DirectRunner' in str(runner_name) or 'TestDirectRunner' in str(runner_name))): output_labels_supported = False except Exception: # If we can't determine runner, assume DirectRunner for safety output_labels_supported = False # Check if in batch mode (not streaming) standard_options = pipeline_options.view_as(StandardOptions) if not standard_options.streaming: output_labels_supported = False else: # If no pipeline options available, fall back to original behavior output_labels_supported = False # Log debug information for troubleshooting import logging runner_info = getattr( pipeline_options, 'runner', 'None') if pipeline_options else 'No options' streaming_info = 'Unknown' if pipeline_options: try: standard_options = pipeline_options.view_as(StandardOptions) streaming_info = 'streaming=%s' % standard_options.streaming except Exception: streaming_info = 'streaming=unknown' logging.debug( 'PubSub unsupported feature check: runner=%s, %s', runner_info, streaming_info) if not output_labels_supported: if transform.id_label: raise NotImplementedError( f'id_label is not supported for PubSub writes with DirectRunner ' f'or in batch mode (runner={runner_info}, {streaming_info})') if transform.timestamp_attribute: raise NotImplementedError( f'timestamp_attribute is not supported for PubSub writes with ' f'DirectRunner or in batch mode ' f'(runner={runner_info}, {streaming_info})') def setup(self): from google.cloud import pubsub self._pub_client = pubsub.PublisherClient() self._topic = self._pub_client.topic_path( self.project, self.short_topic_name) def start_bundle(self): self._buffer = [] def process(self, elem): self._buffer.append(elem) if len(self._buffer) >= self.BUFFER_SIZE_ELEMENTS: self._flush() def finish_bundle(self): self._flush() def _flush(self): if not self._buffer: return import time # The elements in buffer are serialized protobuf bytes from the previous # transforms. We need to deserialize them to extract data and attributes. futures = [] for elem in self._buffer: # Deserialize the protobuf to get the original PubsubMessage pubsub_msg = PubsubMessage._from_proto_str(elem) # Publish with the correct data and attributes if self.with_attributes and pubsub_msg.attributes: future = self._pub_client.publish( self._topic, pubsub_msg.data, **pubsub_msg.attributes) else: future = self._pub_client.publish(self._topic, pubsub_msg.data) futures.append(future) timer_start = time.time() for future in futures: remaining = self.FLUSH_TIMEOUT_SECS - (time.time() - timer_start) if remaining <= 0: raise TimeoutError( f"PubSub publish timeout exceeded {self.FLUSH_TIMEOUT_SECS} seconds" ) future.result(remaining) self._buffer = [] class _PubSubSink(object): """Sink for a Cloud Pub/Sub topic. This sink works for both streaming and batch pipelines by using a DoFn that buffers and batches messages for efficient publishing. """ def __init__( self, topic: str, id_label: Optional[str], timestamp_attribute: Optional[str], ): self.coder = coders.BytesCoder() self.full_topic = topic self.id_label = id_label self.timestamp_attribute = timestamp_attribute self.project, self.topic_name = parse_topic(topic) class PubSubSourceDescriptor(NamedTuple): """A PubSub source descriptor for ``MultipleReadFromPubSub``` Attributes: source: Existing Cloud Pub/Sub topic or subscription to use in the form "projects//topics/" or "projects//subscriptions/" id_label: The attribute on incoming Pub/Sub messages to use as a unique record identifier. When specified, the value of this attribute (which can be any string that uniquely identifies the record) will be used for deduplication of messages. If not provided, we cannot guarantee that no duplicate data will be delivered on the Pub/Sub stream. In this case, deduplication of the stream will be strictly best effort. timestamp_attribute: Message value to use as element timestamp. If None, uses message publishing time as the timestamp. Timestamp values should be in one of two formats: - A numerical value representing the number of milliseconds since the Unix epoch. - A string in RFC 3339 format, UTC timezone. Example: ``2015-10-29T23:41:41.123Z``. The sub-second component of the timestamp is optional, and digits beyond the first three (i.e., time units smaller than milliseconds) may be ignored. """ source: str id_label: str = None timestamp_attribute: str = None PUBSUB_DESCRIPTOR_REGEXP = 'projects/([^/]+)/(topics|subscriptions)/(.+)' class MultipleReadFromPubSub(PTransform): """A ``PTransform`` that expands ``ReadFromPubSub`` to read from multiple ``PubSubSourceDescriptor``. The `MultipleReadFromPubSub` transform allows you to read multiple topics and/or subscriptions using just one transform. It is the recommended transform to read multiple Pub/Sub sources when the output `PCollection` are going to be flattened. The transform takes a list of `PubSubSourceDescriptor` and organize them by type (topic / subscription) and project::: topic_1 = PubSubSourceDescriptor('projects/myproject/topics/a_topic') topic_2 = PubSubSourceDescriptor( 'projects/myproject2/topics/b_topic', 'my_label', 'my_timestamp_attribute') subscription_1 = PubSubSourceDescriptor( 'projects/myproject/subscriptions/a_subscription') results = pipeline | MultipleReadFromPubSub( [topic_1, topic_2, subscription_1]) """ def __init__( self, pubsub_source_descriptors: list[PubSubSourceDescriptor], with_attributes: bool = False, ): """Initializes ``PubSubMultipleReader``. Args: pubsub_source_descriptors: List of Cloud Pub/Sub topics or subscriptions of type `~PubSubSourceDescriptor`. with_attributes: True - input elements will be :class:`~PubsubMessage` objects. False - input elements will be of type ``bytes`` (message data only). """ self.pubsub_source_descriptors = pubsub_source_descriptors self.with_attributes = with_attributes for descriptor in self.pubsub_source_descriptors: match_descriptor = re.match(PUBSUB_DESCRIPTOR_REGEXP, descriptor.source) if not match_descriptor: raise ValueError( 'PubSub source descriptor must be in the form "projects/' '/topics/" or "projects//subscription' '/" (got %r).' % descriptor.source) def expand(self, pcol): sources_pcol = [] for descriptor in self.pubsub_source_descriptors: source_match = re.match(PUBSUB_DESCRIPTOR_REGEXP, descriptor.source) source_project = source_match.group(1) source_type = source_match.group(2) source_name = source_match.group(3) read_step_name = 'PubSub %s/project:%s/Read %s' % ( source_type, source_project, source_name) if source_type == 'topics': current_source = pcol | read_step_name >> ReadFromPubSub( topic=descriptor.source, id_label=descriptor.id_label, with_attributes=self.with_attributes, timestamp_attribute=descriptor.timestamp_attribute) else: current_source = pcol | read_step_name >> ReadFromPubSub( subscription=descriptor.source, id_label=descriptor.id_label, with_attributes=self.with_attributes, timestamp_attribute=descriptor.timestamp_attribute) sources_pcol.append(current_source) return tuple(sources_pcol) | Flatten() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/pubsub_integration_test.py0000644000175100017510000003003015206113032024345 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Integration test for Google Cloud Pub/Sub. """ # pytype: skip-file import logging import time import unittest import uuid import pytest from hamcrest.core.core.allof import all_of from apache_beam.io.gcp import pubsub_it_pipeline from apache_beam.io.gcp.pubsub import PubsubMessage from apache_beam.io.gcp.pubsub import WriteToPubSub from apache_beam.io.gcp.tests.pubsub_matcher import PubSubMessageMatcher from apache_beam.runners.runner import PipelineState from apache_beam.testing import test_utils from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline INPUT_TOPIC = 'psit_topic_input' OUTPUT_TOPIC = 'psit_topic_output' INPUT_SUB = 'psit_subscription_input' OUTPUT_SUB = 'psit_subscription_output' # How long TestXXXRunner will wait for pubsub_it_pipeline to run before # cancelling it. TEST_PIPELINE_DURATION_MS = 10 * 60 * 1000 # How long PubSubMessageMatcher will wait for the correct set of messages to # appear. MESSAGE_MATCHER_TIMEOUT_S = 10 * 60 class PubSubIntegrationTest(unittest.TestCase): ID_LABEL = 'id' TIMESTAMP_ATTRIBUTE = 'timestamp' INPUT_MESSAGES = { # TODO(https://github.com/apache/beam/issues/18939): DirectRunner doesn't # support reading or writing label_ids, nor writing timestamp attributes. # Once these features exist, TestDirectRunner and TestDataflowRunner # should behave identically. 'TestDirectRunner': [ PubsubMessage(b'data001', {}), # For those elements that have the TIMESTAMP_ATTRIBUTE attribute, the # IT pipeline writes back the timestamp of each element (as reported # by Beam), as a TIMESTAMP_ATTRIBUTE + '_out' attribute. PubsubMessage( b'data002', { TIMESTAMP_ATTRIBUTE: '2018-07-11T02:02:50.149000Z', }), PubsubMessage(b'data003\xab\xac', {}), PubsubMessage( b'data004\xab\xac', { TIMESTAMP_ATTRIBUTE: '2018-07-11T02:02:50.149000Z', }) ], 'TestDataflowRunner': [ # Use ID_LABEL attribute to deduplicate messages with the same ID. PubsubMessage(b'data001', {ID_LABEL: 'foo'}), PubsubMessage(b'data001', {ID_LABEL: 'foo'}), PubsubMessage(b'data001', {ID_LABEL: 'foo'}), # For those elements that have the TIMESTAMP_ATTRIBUTE attribute, the # IT pipeline writes back the timestamp of each element (as reported # by Beam), as a TIMESTAMP_ATTRIBUTE + '_out' attribute. PubsubMessage( b'data002', { TIMESTAMP_ATTRIBUTE: '2018-07-11T02:02:50.149000Z', }), PubsubMessage(b'data003\xab\xac', {ID_LABEL: 'foo2'}), PubsubMessage(b'data003\xab\xac', {ID_LABEL: 'foo2'}), PubsubMessage(b'data003\xab\xac', {ID_LABEL: 'foo2'}), PubsubMessage( b'data004\xab\xac', { TIMESTAMP_ATTRIBUTE: '2018-07-11T02:02:50.149000Z', }) ], } EXPECTED_OUTPUT_MESSAGES = { 'TestDirectRunner': [ PubsubMessage(b'data001-seen', {'processed': 'IT'}), PubsubMessage( b'data002-seen', { TIMESTAMP_ATTRIBUTE: '2018-07-11T02:02:50.149000Z', TIMESTAMP_ATTRIBUTE + '_out': '2018-07-11T02:02:50.149000Z', 'processed': 'IT', }), PubsubMessage(b'data003\xab\xac-seen', {'processed': 'IT'}), PubsubMessage( b'data004\xab\xac-seen', { TIMESTAMP_ATTRIBUTE: '2018-07-11T02:02:50.149000Z', TIMESTAMP_ATTRIBUTE + '_out': '2018-07-11T02:02:50.149000Z', 'processed': 'IT', }) ], 'TestDataflowRunner': [ PubsubMessage(b'data001-seen', {'processed': 'IT'}), PubsubMessage( b'data002-seen', { TIMESTAMP_ATTRIBUTE + '_out': '2018-07-11T02:02:50.149000Z', 'processed': 'IT', }), PubsubMessage(b'data003\xab\xac-seen', {'processed': 'IT'}), PubsubMessage( b'data004\xab\xac-seen', { TIMESTAMP_ATTRIBUTE + '_out': '2018-07-11T02:02:50.149000Z', 'processed': 'IT', }) ], } def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.runner_name = type(self.test_pipeline.runner).__name__ self.project = self.test_pipeline.get_option('project') self.uuid = str(uuid.uuid4()) # Set up PubSub environment. from google.cloud import pubsub self.pub_client = pubsub.PublisherClient() self.input_topic = self.pub_client.create_topic( name=self.pub_client.topic_path(self.project, INPUT_TOPIC + self.uuid)) self.output_topic = self.pub_client.create_topic( name=self.pub_client.topic_path(self.project, OUTPUT_TOPIC + self.uuid)) self.sub_client = pubsub.SubscriberClient() self.input_sub = self.sub_client.create_subscription( name=self.sub_client.subscription_path( self.project, INPUT_SUB + self.uuid), topic=self.input_topic.name) self.output_sub = self.sub_client.create_subscription( name=self.sub_client.subscription_path( self.project, OUTPUT_SUB + self.uuid), topic=self.output_topic.name) # Add a 30 second sleep after resource creation to ensure subscriptions will # receive messages. time.sleep(30) def tearDown(self): test_utils.cleanup_subscriptions( self.sub_client, [self.input_sub, self.output_sub]) test_utils.cleanup_topics( self.pub_client, [self.input_topic, self.output_topic]) def _test_streaming(self, with_attributes): """Runs IT pipeline with message verifier. Args: with_attributes: False - Reads and writes message data only. True - Reads and writes message data and attributes. Also verifies id_label and timestamp_attribute features. """ # Set on_success_matcher to verify pipeline state and pubsub output. These # verifications run on a (remote) worker. # Expect the state to be RUNNING since a streaming pipeline is usually # never DONE. The test runner will cancel the pipeline after verification. state_verifier = PipelineStateMatcher(PipelineState.RUNNING) expected_messages = self.EXPECTED_OUTPUT_MESSAGES[self.runner_name] if not with_attributes: expected_messages = [pubsub_msg.data for pubsub_msg in expected_messages] if self.runner_name == 'TestDirectRunner': strip_attributes = None else: strip_attributes = [self.ID_LABEL, self.TIMESTAMP_ATTRIBUTE] pubsub_msg_verifier = PubSubMessageMatcher( self.project, self.output_sub.name, expected_messages, timeout=MESSAGE_MATCHER_TIMEOUT_S, with_attributes=with_attributes, strip_attributes=strip_attributes) extra_opts = { 'input_subscription': self.input_sub.name, 'output_topic': self.output_topic.name, 'wait_until_finish_duration': TEST_PIPELINE_DURATION_MS, 'on_success_matcher': all_of(state_verifier, pubsub_msg_verifier) } # Generate input data and inject to PubSub. for msg in self.INPUT_MESSAGES[self.runner_name]: self.pub_client.publish( self.input_topic.name, msg.data, **msg.attributes).result() # Get pipeline options from command argument: --test-pipeline-options, # and start pipeline job by calling pipeline main function. pubsub_it_pipeline.run_pipeline( argv=self.test_pipeline.get_full_options_as_args(**extra_opts), with_attributes=with_attributes, id_label=self.ID_LABEL, timestamp_attribute=self.TIMESTAMP_ATTRIBUTE) @pytest.mark.it_postcommit def test_streaming_data_only(self): self._test_streaming(with_attributes=False) @pytest.mark.it_postcommit def test_streaming_with_attributes(self): self._test_streaming(with_attributes=True) def _test_batch_write(self, with_attributes): """Tests batch mode WriteToPubSub functionality. Args: with_attributes: False - Writes message data only. True - Writes message data and attributes. """ from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.transforms import Create # Create test messages for batch mode test_messages = [ PubsubMessage(b'batch_data001', {'batch_attr': 'value1'}), PubsubMessage(b'batch_data002', {'batch_attr': 'value2'}), PubsubMessage(b'batch_data003', {'batch_attr': 'value3'}) ] pipeline_options = PipelineOptions() # Explicitly set streaming to False for batch mode pipeline_options.view_as(StandardOptions).streaming = False with TestPipeline(options=pipeline_options) as p: if with_attributes: messages = p | 'CreateMessages' >> Create(test_messages) _ = messages | 'WriteToPubSub' >> WriteToPubSub( self.output_topic.name, with_attributes=True) else: # For data-only mode, extract just the data message_data = [msg.data for msg in test_messages] messages = p | 'CreateData' >> Create(message_data) _ = messages | 'WriteToPubSub' >> WriteToPubSub( self.output_topic.name, with_attributes=False) # Verify messages were published by reading from the subscription time.sleep(10) # Allow time for messages to be published and received # Pull messages from the output subscription to verify they were written response = self.sub_client.pull( request={ "subscription": self.output_sub.name, "max_messages": 10, }) received_messages = [] for received_message in response.received_messages: if with_attributes: # Parse attributes attrs = dict(received_message.message.attributes) received_messages.append( PubsubMessage(received_message.message.data, attrs)) else: received_messages.append(received_message.message.data) # Acknowledge the message self.sub_client.acknowledge( request={ "subscription": self.output_sub.name, "ack_ids": [received_message.ack_id], }) # Verify we received the expected number of messages self.assertEqual(len(received_messages), len(test_messages)) if with_attributes: # Verify message content and attributes received_data = [msg.data for msg in received_messages] expected_data = [msg.data for msg in test_messages] self.assertEqual(sorted(received_data), sorted(expected_data)) else: # Verify message data only expected_data = [msg.data for msg in test_messages] self.assertEqual(sorted(received_messages), sorted(expected_data)) @pytest.mark.it_postcommit def test_batch_write_data_only(self): """Test WriteToPubSub in batch mode with data only.""" self._test_batch_write(with_attributes=False) @pytest.mark.it_postcommit def test_batch_write_with_attributes(self): """Test WriteToPubSub in batch mode with attributes.""" self._test_batch_write(with_attributes=True) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/pubsub_io_perf_test.py0000644000175100017510000001663715206113032023466 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Performance PubsubIO streaming test for Write/Read operations. Caution: only test runners (e.g. TestDataflowRunner) support matchers Example for TestDataflowRunner: python -m apache_beam.io.gcp.pubsub_io_perf_test \ --test-pipeline-options=" --runner=TestDataflowRunner --sdk_location=.../dist/apache-beam-x.x.x.dev0.tar.gz --project= --temp_location=gs:///tmp --staging_location=gs:///staging --wait_until_finish_duration= --pubsub_namespace_prefix= --publish_to_big_query= --metrics_dataset= --metrics_table= --input_options='{ \"num_records\": \"key_size\": 1 \"value_size\": }'" """ # pytype: skip-file import logging import sys from hamcrest import all_of import apache_beam as beam from apache_beam.io import Read from apache_beam.io import ReadFromPubSub from apache_beam.io.gcp.tests.pubsub_matcher import PubSubMessageMatcher from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.load_tests.load_test import LoadTest from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime from apache_beam.testing.synthetic_pipeline import SyntheticSource from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.transforms import trigger from apache_beam.transforms import window # pylint: disable=wrong-import-order, wrong-import-position try: from google.cloud import pubsub except ImportError: pubsub = None # pylint: enable=wrong-import-order, wrong-import-position WRITE_METRICS_NAMESPACE = 'pubsub_io_perf_write' READ_METRICS_NAMESPACE = 'pubsub_io_perf_read' MATCHER_TIMEOUT = 60 * 15 MATCHER_PULL_TIMEOUT = 60 * 5 class PubsubIOPerfTest(LoadTest): def _setup_env(self): if not self.pipeline.get_option('pubsub_namespace_prefix'): logging.error('--pubsub_namespace_prefix argument is required.') sys.exit(1) if not self.pipeline.get_option('wait_until_finish_duration'): logging.error('--wait_until_finish_duration argument is required.') sys.exit(1) self.num_of_messages = int(self.input_options.get('num_records')) pubsub_namespace_prefix = self.pipeline.get_option( 'pubsub_namespace_prefix') self.pubsub_namespace = pubsub_namespace_prefix + unique_id def _setup_pubsub(self): self.pub_client = pubsub.PublisherClient() self.topic_name = self.pub_client.topic_path( self.project_id, self.pubsub_namespace) self.matcher_topic_name = self.pub_client.topic_path( self.project_id, self.pubsub_namespace + '_matcher') self.sub_client = pubsub.SubscriberClient() self.read_sub_name = self.sub_client.subscription_path( self.project_id, self.pubsub_namespace + '_read', ) self.read_matcher_sub_name = self.sub_client.subscription_path( self.project_id, self.pubsub_namespace + '_read_matcher', ) class PubsubWritePerfTest(PubsubIOPerfTest): def __init__(self): super().__init__(WRITE_METRICS_NAMESPACE) self._setup_env() self._setup_pubsub() self._setup_pipeline() def test(self): def to_pubsub_message(element): import uuid from apache_beam.io import PubsubMessage return PubsubMessage( data=element[1], attributes={'id': str(uuid.uuid1()).encode('utf-8')}, ) _ = ( self.pipeline | 'Create input' >> Read( SyntheticSource(self.parse_synthetic_source_options())) | 'Format to pubsub message in bytes' >> beam.Map(to_pubsub_message) | 'Measure time' >> beam.ParDo(MeasureTime(self.metrics_namespace)) | 'Write to Pubsub' >> beam.io.WriteToPubSub( self.topic_name, with_attributes=True, id_label='id', )) def _setup_pipeline(self): options = PipelineOptions(self.pipeline.get_full_options_as_args()) options.view_as(StandardOptions).streaming = True self.pipeline = TestPipeline(options=options) def _setup_pubsub(self): super()._setup_pubsub() _ = self.pub_client.create_topic(name=self.topic_name) _ = self.sub_client.create_subscription( name=self.read_sub_name, topic=self.topic_name, ) class PubsubReadPerfTest(PubsubIOPerfTest): def __init__(self): super().__init__(READ_METRICS_NAMESPACE) self._setup_env() self._setup_pubsub() self._setup_pipeline() def test(self): _ = ( self.pipeline | 'Read from pubsub' >> ReadFromPubSub( subscription=self.read_sub_name, with_attributes=True, id_label='id', ) | beam.Map(lambda x: bytes(1)).with_output_types(bytes) | 'Measure time' >> beam.ParDo(MeasureTime(self.metrics_namespace)) | 'Window' >> beam.WindowInto( window.GlobalWindows(), trigger=trigger.Repeatedly( trigger.AfterCount(self.num_of_messages)), accumulation_mode=trigger.AccumulationMode.DISCARDING) | 'Count messages' >> beam.CombineGlobally( beam.combiners.CountCombineFn()).without_defaults(). with_output_types(int) | 'Convert to bytes' >> beam.Map(lambda count: str(count).encode('utf-8')) | 'Write to Pubsub' >> beam.io.WriteToPubSub(self.matcher_topic_name)) def _setup_pubsub(self): super()._setup_pubsub() _ = self.pub_client.create_topic(name=self.matcher_topic_name) _ = self.sub_client.create_subscription( name=self.read_matcher_sub_name, topic=self.matcher_topic_name, ) def _setup_pipeline(self): pubsub_msg_verifier = PubSubMessageMatcher( self.project_id, self.read_matcher_sub_name, expected_msg=[str(self.num_of_messages).encode('utf-8')], timeout=MATCHER_TIMEOUT, pull_timeout=MATCHER_PULL_TIMEOUT, ) extra_opts = { 'on_success_matcher': all_of(pubsub_msg_verifier), 'streaming': True, } args = self.pipeline.get_full_options_as_args(**extra_opts) self.pipeline = TestPipeline(options=PipelineOptions(args)) def cleanup(self): self.sub_client.delete_subscription(subscription=self.read_sub_name) self.sub_client.delete_subscription(subscription=self.read_matcher_sub_name) self.pub_client.delete_topic(topic=self.topic_name) self.pub_client.delete_topic(topic=self.matcher_topic_name) if __name__ == '__main__': import uuid unique_id = str(uuid.uuid4()) logging.basicConfig(level=logging.INFO) PubsubWritePerfTest().run() PubsubReadPerfTest().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/pubsub_it_pipeline.py0000644000175100017510000000620615206113032023274 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Test pipeline for use by pubsub_integration_test. """ # pytype: skip-file import argparse import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions def run_pipeline(argv, with_attributes, id_label, timestamp_attribute): """Build and run the pipeline.""" parser = argparse.ArgumentParser() parser.add_argument( '--output_topic', required=True, help=( 'Output PubSub topic of the form ' '"projects//topic/".')) parser.add_argument( '--input_subscription', required=True, help=( 'Input PubSub subscription of the form ' '"projects//subscriptions/."')) known_args, pipeline_args = parser.parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(StandardOptions).streaming = True p = beam.Pipeline(options=pipeline_options) runner_name = type(p.runner).__name__ # Read from PubSub into a PCollection. if runner_name == 'TestDirectRunner': messages = p | beam.io.ReadFromPubSub( subscription=known_args.input_subscription, with_attributes=with_attributes, timestamp_attribute=timestamp_attribute) else: messages = p | beam.io.ReadFromPubSub( subscription=known_args.input_subscription, id_label=id_label, with_attributes=with_attributes, timestamp_attribute=timestamp_attribute) def add_attribute(msg, timestamp=beam.DoFn.TimestampParam): msg.data += b'-seen' msg.attributes['processed'] = 'IT' if timestamp_attribute in msg.attributes: msg.attributes[timestamp_attribute + '_out'] = timestamp.to_rfc3339() return msg def modify_data(data): return data + b'-seen' if with_attributes: output = messages | 'add_attribute' >> beam.Map(add_attribute) else: output = messages | 'modify_data' >> beam.Map(modify_data) # Write to PubSub. if runner_name == 'TestDirectRunner': _ = output | beam.io.WriteToPubSub( known_args.output_topic, with_attributes=with_attributes) else: _ = output | beam.io.WriteToPubSub( known_args.output_topic, id_label=id_label, with_attributes=with_attributes, timestamp_attribute=timestamp_attribute) result = p.run() result.wait_until_finish() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/pubsub_test.py0000644000175100017510000012401715206113032021753 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for PubSub sources and sinks.""" # pytype: skip-file import logging import unittest import hamcrest as hc import mock import apache_beam as beam from apache_beam import Pipeline from apache_beam.io import Read from apache_beam.io import Write from apache_beam.io.gcp.pubsub import MultipleReadFromPubSub from apache_beam.io.gcp.pubsub import PubsubMessage from apache_beam.io.gcp.pubsub import PubSubSourceDescriptor from apache_beam.io.gcp.pubsub import ReadFromPubSub from apache_beam.io.gcp.pubsub import ReadStringsFromPubSub from apache_beam.io.gcp.pubsub import WriteStringsToPubSub from apache_beam.io.gcp.pubsub import WriteToPubSub from apache_beam.io.gcp.pubsub import _PubSubSink from apache_beam.io.gcp.pubsub import _PubSubSource from apache_beam.metrics.metric import Lineage from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners import pipeline_context from apache_beam.runners.direct import transform_evaluator from apache_beam.runners.direct.direct_runner import _DirectReadFromPubSub from apache_beam.runners.direct.direct_runner import _get_transform_overrides from apache_beam.runners.direct.transform_evaluator import _PubSubReadEvaluator from apache_beam.testing import test_utils from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import TestWindowedValue from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import window from apache_beam.transforms.core import Create from apache_beam.transforms.display import DisplayData from apache_beam.transforms.display_test import DisplayDataItemMatcher from apache_beam.utils import proto_utils from apache_beam.utils import timestamp # Protect against environments where the PubSub library is not available. try: from google.cloud import pubsub except ImportError: pubsub = None class TestPubsubMessage(unittest.TestCase): def test_payload_valid(self): _ = PubsubMessage('', None) _ = PubsubMessage('data', None) _ = PubsubMessage(None, {'k': 'v'}) def test_payload_invalid(self): with self.assertRaisesRegex(ValueError, r'data.*attributes.*must be set'): _ = PubsubMessage(None, None) with self.assertRaisesRegex(ValueError, r'data.*attributes.*must be set'): _ = PubsubMessage(None, {}) @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') def test_proto_conversion(self): data = b'data' attributes = {'k1': 'v1', 'k2': 'v2'} m = PubsubMessage(data, attributes) m_converted = PubsubMessage._from_proto_str(m._to_proto_str()) self.assertEqual(m_converted.data, data) self.assertEqual(m_converted.attributes, attributes) @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') def test_payload_publish_invalid(self): with self.assertRaisesRegex(ValueError, r'data field.*10MB'): msg = PubsubMessage(b'0' * 1024 * 1024 * 11, None) msg._to_proto_str(for_publish=True) with self.assertRaisesRegex(ValueError, 'attribute key'): msg = PubsubMessage(b'0', {'0' * 257: '0'}) msg._to_proto_str(for_publish=True) with self.assertRaisesRegex(ValueError, 'attribute value'): msg = PubsubMessage(b'0', {'0' * 100: '0' * 1025}) msg._to_proto_str(for_publish=True) with self.assertRaisesRegex(ValueError, '100 attributes'): attributes = {} for i in range(0, 101): attributes[str(i)] = str(i) msg = PubsubMessage(b'0', attributes) msg._to_proto_str(for_publish=True) with self.assertRaisesRegex(ValueError, 'ordering key'): msg = PubsubMessage(b'0', None, ordering_key='0' * 1301) msg._to_proto_str(for_publish=True) def test_eq(self): a = PubsubMessage(b'abc', {1: 2, 3: 4}) b = PubsubMessage(b'abc', {1: 2, 3: 4}) c = PubsubMessage(b'abc', {1: 2}) self.assertTrue(a == b) self.assertTrue(a != c) self.assertTrue(b != c) def test_hash(self): a = PubsubMessage(b'abc', {1: 2, 3: 4}) b = PubsubMessage(b'abc', {1: 2, 3: 4}) c = PubsubMessage(b'abc', {1: 2}) self.assertTrue(hash(a) == hash(b)) self.assertTrue(hash(a) != hash(c)) self.assertTrue(hash(b) != hash(c)) def test_repr(self): a = PubsubMessage(b'abc', {1: 2, 3: 4}) b = PubsubMessage(b'abc', {1: 2, 3: 4}) c = PubsubMessage(b'abc', {1: 2}) self.assertTrue(repr(a) == repr(b)) self.assertTrue(repr(a) != repr(c)) self.assertTrue(repr(b) != repr(c)) @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') class TestReadFromPubSubOverride(unittest.TestCase): def test_expand_with_topic(self): options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True p = TestPipeline(options=options) pcoll = ( p | ReadFromPubSub( 'projects/fakeprj/topics/a_topic', None, 'a_label', with_attributes=False, timestamp_attribute=None) | beam.Map(lambda x: x)) self.assertEqual(bytes, pcoll.element_type) # Apply the necessary PTransformOverrides. overrides = _get_transform_overrides(options) p.replace_all(overrides) # Note that the direct output of ReadFromPubSub will be replaced # by a PTransformOverride, so we use a no-op Map. read_transform = pcoll.producer.inputs[0].producer.transform # Ensure that the properties passed through correctly source = read_transform._source self.assertEqual('a_topic', source.topic_name) self.assertEqual('a_label', source.id_label) def test_expand_with_subscription(self): options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True p = TestPipeline(options=options) pcoll = ( p | ReadFromPubSub( None, 'projects/fakeprj/subscriptions/a_subscription', 'a_label', with_attributes=False, timestamp_attribute=None) | beam.Map(lambda x: x)) self.assertEqual(bytes, pcoll.element_type) # Apply the necessary PTransformOverrides. overrides = _get_transform_overrides(options) p.replace_all(overrides) # Note that the direct output of ReadFromPubSub will be replaced # by a PTransformOverride, so we use a no-op Map. read_transform = pcoll.producer.inputs[0].producer.transform # Ensure that the properties passed through correctly source = read_transform._source self.assertEqual('a_subscription', source.subscription_name) self.assertEqual('a_label', source.id_label) def test_expand_with_no_topic_or_subscription(self): with self.assertRaisesRegex( ValueError, "Either a topic or subscription must be provided."): ReadFromPubSub( None, None, 'a_label', with_attributes=False, timestamp_attribute=None) def test_expand_with_both_topic_and_subscription(self): with self.assertRaisesRegex( ValueError, "Only one of topic or subscription should be provided."): ReadFromPubSub( 'a_topic', 'a_subscription', 'a_label', with_attributes=False, timestamp_attribute=None) def test_expand_with_other_options(self): options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True p = TestPipeline(options=options) pcoll = ( p | ReadFromPubSub( 'projects/fakeprj/topics/a_topic', None, 'a_label', with_attributes=True, timestamp_attribute='time') | beam.Map(lambda x: x)) self.assertEqual(PubsubMessage, pcoll.element_type) # Apply the necessary PTransformOverrides. overrides = _get_transform_overrides(options) p.replace_all(overrides) # Note that the direct output of ReadFromPubSub will be replaced # by a PTransformOverride, so we use a no-op Map. read_transform = pcoll.producer.inputs[0].producer.transform # Ensure that the properties passed through correctly source = read_transform._source self.assertTrue(source.with_attributes) self.assertEqual('time', source.timestamp_attribute) @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') class TestMultiReadFromPubSubOverride(unittest.TestCase): def test_expand_with_multiple_sources(self): options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True p = TestPipeline(options=options) topics = [ 'projects/fakeprj/topics/a_topic', 'projects/fakeprj2/topics/b_topic' ] subscriptions = ['projects/fakeprj/subscriptions/a_subscription'] pubsub_sources = [ PubSubSourceDescriptor(descriptor) for descriptor in topics + subscriptions ] pcoll = (p | MultipleReadFromPubSub(pubsub_sources) | beam.Map(lambda x: x)) # Apply the necessary PTransformOverrides. overrides = _get_transform_overrides(options) p.replace_all(overrides) self.assertEqual(bytes, pcoll.element_type) # Ensure that the sources are passed through correctly read_transforms = pcoll.producer.inputs[0].producer.inputs topics_list = [] subscription_list = [] for read_transform in read_transforms: source = read_transform.producer.transform._source if source.full_topic: topics_list.append(source.full_topic) else: subscription_list.append(source.full_subscription) self.assertEqual(topics_list, topics) self.assertEqual(subscription_list, subscriptions) def test_expand_with_multiple_sources_and_attributes(self): options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True p = TestPipeline(options=options) topics = [ 'projects/fakeprj/topics/a_topic', 'projects/fakeprj2/topics/b_topic' ] subscriptions = ['projects/fakeprj/subscriptions/a_subscription'] pubsub_sources = [ PubSubSourceDescriptor(descriptor) for descriptor in topics + subscriptions ] pcoll = ( p | MultipleReadFromPubSub(pubsub_sources, with_attributes=True) | beam.Map(lambda x: x)) # Apply the necessary PTransformOverrides. overrides = _get_transform_overrides(options) p.replace_all(overrides) self.assertEqual(PubsubMessage, pcoll.element_type) # Ensure that the sources are passed through correctly read_transforms = pcoll.producer.inputs[0].producer.inputs topics_list = [] subscription_list = [] for read_transform in read_transforms: source = read_transform.producer.transform._source if source.full_topic: topics_list.append(source.full_topic) else: subscription_list.append(source.full_subscription) self.assertEqual(topics_list, topics) self.assertEqual(subscription_list, subscriptions) def test_expand_with_multiple_sources_and_other_options(self): options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True p = TestPipeline(options=options) sources = [ 'projects/fakeprj/topics/a_topic', 'projects/fakeprj2/topics/b_topic', 'projects/fakeprj/subscriptions/a_subscription' ] id_labels = ['a_label_topic', 'b_label_topic', 'a_label_subscription'] timestamp_attributes = ['a_ta_topic', 'b_ta_topic', 'a_ta_subscription'] pubsub_sources = [ PubSubSourceDescriptor( source=source, id_label=id_label, timestamp_attribute=timestamp_attribute) for source, id_label, timestamp_attribute in zip( sources, id_labels, timestamp_attributes) ] pcoll = (p | MultipleReadFromPubSub(pubsub_sources) | beam.Map(lambda x: x)) # Apply the necessary PTransformOverrides. overrides = _get_transform_overrides(options) p.replace_all(overrides) self.assertEqual(bytes, pcoll.element_type) # Ensure that the sources are passed through correctly read_transforms = pcoll.producer.inputs[0].producer.inputs for i, read_transform in enumerate(read_transforms): id_label = id_labels[i] timestamp_attribute = timestamp_attributes[i] source = read_transform.producer.transform._source self.assertEqual(source.id_label, id_label) self.assertEqual(source.with_attributes, False) self.assertEqual(source.timestamp_attribute, timestamp_attribute) def test_expand_with_wrong_source(self): with self.assertRaisesRegex( ValueError, r'PubSub source descriptor must be in the form ' r'"projects//topics/"' ' or "projects//subscription/".*'): MultipleReadFromPubSub([PubSubSourceDescriptor('not_a_proper_source')]) @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') class TestWriteStringsToPubSubOverride(unittest.TestCase): @mock.patch.object(Pipeline, '_assert_not_applying_PDone', mock.Mock()) def test_expand_deprecated(self): options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True p = TestPipeline(options=options) pcoll = ( p | ReadFromPubSub('projects/fakeprj/topics/baz') | WriteStringsToPubSub('projects/fakeprj/topics/a_topic') | beam.Map(lambda x: x)) # Apply the necessary PTransformOverrides. overrides = _get_transform_overrides(options) p.replace_all(overrides) # Note that the direct output of ReadFromPubSub will be replaced # by a PTransformOverride, so we use a no-op Map. write_transform = pcoll.producer.inputs[0].producer.transform # Ensure that the properties passed through correctly self.assertEqual('a_topic', write_transform.dofn.short_topic_name) @mock.patch.object(Pipeline, '_assert_not_applying_PDone', mock.Mock()) def test_expand(self): options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True p = TestPipeline(options=options) pcoll = ( p | ReadFromPubSub('projects/fakeprj/topics/baz') | beam.Map(lambda x: PubsubMessage(x)) | WriteToPubSub( 'projects/fakeprj/topics/a_topic', with_attributes=True) | beam.Map(lambda x: x)) # Apply the necessary PTransformOverrides. overrides = _get_transform_overrides(options) p.replace_all(overrides) # Note that the direct output of ReadFromPubSub will be replaced # by a PTransformOverride, so we use a no-op Map. write_transform = pcoll.producer.inputs[0].producer.transform # Ensure that the properties passed through correctly self.assertEqual('a_topic', write_transform.dofn.short_topic_name) self.assertEqual(True, write_transform.dofn.with_attributes) # TODO(https://github.com/apache/beam/issues/18939): These properties # aren't supported yet in direct runner. self.assertEqual(None, write_transform.dofn.id_label) self.assertEqual(None, write_transform.dofn.timestamp_attribute) @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') class TestPubSubSource(unittest.TestCase): def test_display_data_topic(self): source = _PubSubSource('projects/fakeprj/topics/a_topic', None, 'a_label') dd = DisplayData.create_from(source) expected_items = [ DisplayDataItemMatcher('topic', 'projects/fakeprj/topics/a_topic'), DisplayDataItemMatcher('id_label', 'a_label'), DisplayDataItemMatcher('with_attributes', False), ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_display_data_subscription(self): source = _PubSubSource( None, 'projects/fakeprj/subscriptions/a_subscription', 'a_label') dd = DisplayData.create_from(source) expected_items = [ DisplayDataItemMatcher( 'subscription', 'projects/fakeprj/subscriptions/a_subscription'), DisplayDataItemMatcher('id_label', 'a_label'), DisplayDataItemMatcher('with_attributes', False), ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_display_data_no_subscription(self): source = _PubSubSource('projects/fakeprj/topics/a_topic') dd = DisplayData.create_from(source) expected_items = [ DisplayDataItemMatcher('topic', 'projects/fakeprj/topics/a_topic'), DisplayDataItemMatcher('with_attributes', False), ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') class TestPubSubSink(unittest.TestCase): def test_display_data(self): sink = WriteToPubSub( 'projects/fakeprj/topics/a_topic', id_label='id', timestamp_attribute='time') dd = DisplayData.create_from(sink) expected_items = [ DisplayDataItemMatcher('topic', 'projects/fakeprj/topics/a_topic'), DisplayDataItemMatcher('id_label', 'id'), DisplayDataItemMatcher('with_attributes', True), DisplayDataItemMatcher('timestamp_attribute', 'time'), ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) class TestPubSubReadEvaluator(object): """Wrapper of _PubSubReadEvaluator that makes it bounded.""" _pubsub_read_evaluator = _PubSubReadEvaluator def __init__(self, *args, **kwargs): self._evaluator = self._pubsub_read_evaluator(*args, **kwargs) def start_bundle(self): return self._evaluator.start_bundle() def process_element(self, element): return self._evaluator.process_element(element) def finish_bundle(self): result = self._evaluator.finish_bundle() result.unprocessed_bundles = [] result.keyed_watermark_holds = {None: None} return result transform_evaluator.TransformEvaluatorRegistry._test_evaluators_overrides = { _DirectReadFromPubSub: TestPubSubReadEvaluator, # type: ignore[dict-item] } @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') @mock.patch('google.cloud.pubsub.SubscriberClient') class TestReadFromPubSub(unittest.TestCase): def test_read_messages_success(self, mock_pubsub): data = b'data' publish_time_secs = 1520861821 publish_time_nanos = 234567000 attributes = {'key': 'value'} ack_id = 'ack_id' pull_response = test_utils.create_pull_response([ test_utils.PullResponseMessage( data, attributes, publish_time_secs, publish_time_nanos, ack_id) ]) expected_elements = [ TestWindowedValue( PubsubMessage(data, attributes), timestamp.Timestamp(1520861821.234567), [window.GlobalWindow()]) ] mock_pubsub.return_value.pull.return_value = pull_response options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: pcoll = ( p | ReadFromPubSub( 'projects/fakeprj/topics/a_topic', None, None, with_attributes=True)) assert_that(pcoll, equal_to(expected_elements), reify_windows=True) mock_pubsub.return_value.acknowledge.assert_has_calls( [mock.call(subscription=mock.ANY, ack_ids=[ack_id])]) mock_pubsub.return_value.close.assert_has_calls([mock.call()]) def test_read_strings_success(self, mock_pubsub): data = '🤷 ¯\\_(ツ)_/¯' data_encoded = data.encode('utf-8') ack_id = 'ack_id' pull_response = test_utils.create_pull_response( [test_utils.PullResponseMessage(data_encoded, ack_id=ack_id)]) expected_elements = [data] mock_pubsub.return_value.pull.return_value = pull_response options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: pcoll = ( p | ReadStringsFromPubSub( 'projects/fakeprj/topics/a_topic', None, None)) assert_that(pcoll, equal_to(expected_elements)) mock_pubsub.return_value.acknowledge.assert_has_calls( [mock.call(subscription=mock.ANY, ack_ids=[ack_id])]) mock_pubsub.return_value.close.assert_has_calls([mock.call()]) def test_read_data_success(self, mock_pubsub): data_encoded = '🤷 ¯\\_(ツ)_/¯'.encode('utf-8') ack_id = 'ack_id' pull_response = test_utils.create_pull_response( [test_utils.PullResponseMessage(data_encoded, ack_id=ack_id)]) expected_elements = [data_encoded] mock_pubsub.return_value.pull.return_value = pull_response options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: pcoll = ( p | ReadFromPubSub('projects/fakeprj/topics/a_topic', None, None)) assert_that(pcoll, equal_to(expected_elements)) mock_pubsub.return_value.acknowledge.assert_has_calls( [mock.call(subscription=mock.ANY, ack_ids=[ack_id])]) mock_pubsub.return_value.close.assert_has_calls([mock.call()]) def test_read_messages_timestamp_attribute_milli_success(self, mock_pubsub): data = b'data' attributes = {'time': '1337'} publish_time_secs = 1520861821 publish_time_nanos = 234567000 ack_id = 'ack_id' pull_response = test_utils.create_pull_response([ test_utils.PullResponseMessage( data, attributes, publish_time_secs, publish_time_nanos, ack_id) ]) expected_elements = [ TestWindowedValue( PubsubMessage(data, attributes), timestamp.Timestamp(micros=int(attributes['time']) * 1000), [window.GlobalWindow()]), ] mock_pubsub.return_value.pull.return_value = pull_response options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: pcoll = ( p | ReadFromPubSub( 'projects/fakeprj/topics/a_topic', None, None, with_attributes=True, timestamp_attribute='time')) assert_that(pcoll, equal_to(expected_elements), reify_windows=True) mock_pubsub.return_value.acknowledge.assert_has_calls( [mock.call(subscription=mock.ANY, ack_ids=[ack_id])]) mock_pubsub.return_value.close.assert_has_calls([mock.call()]) def test_read_messages_timestamp_attribute_rfc3339_success(self, mock_pubsub): data = b'data' attributes = {'time': '2018-03-12T13:37:01.234567Z'} publish_time_secs = 1337000000 publish_time_nanos = 133700000 ack_id = 'ack_id' pull_response = test_utils.create_pull_response([ test_utils.PullResponseMessage( data, attributes, publish_time_secs, publish_time_nanos, ack_id) ]) expected_elements = [ TestWindowedValue( PubsubMessage(data, attributes), timestamp.Timestamp.from_rfc3339(attributes['time']), [window.GlobalWindow()]), ] mock_pubsub.return_value.pull.return_value = pull_response options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: pcoll = ( p | ReadFromPubSub( 'projects/fakeprj/topics/a_topic', None, None, with_attributes=True, timestamp_attribute='time')) assert_that(pcoll, equal_to(expected_elements), reify_windows=True) mock_pubsub.return_value.acknowledge.assert_has_calls( [mock.call(subscription=mock.ANY, ack_ids=[ack_id])]) mock_pubsub.return_value.close.assert_has_calls([mock.call()]) def test_read_messages_timestamp_attribute_missing(self, mock_pubsub): data = b'data' attributes = {} publish_time_secs = 1520861821 publish_time_nanos = 234567000 publish_time = '2018-03-12T13:37:01.234567Z' ack_id = 'ack_id' pull_response = test_utils.create_pull_response([ test_utils.PullResponseMessage( data, attributes, publish_time_secs, publish_time_nanos, ack_id) ]) expected_elements = [ TestWindowedValue( PubsubMessage(data, attributes), timestamp.Timestamp.from_rfc3339(publish_time), [window.GlobalWindow()]), ] mock_pubsub.return_value.pull.return_value = pull_response options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: pcoll = ( p | ReadFromPubSub( 'projects/fakeprj/topics/a_topic', None, None, with_attributes=True, timestamp_attribute='nonexistent')) assert_that(pcoll, equal_to(expected_elements), reify_windows=True) mock_pubsub.return_value.acknowledge.assert_has_calls( [mock.call(subscription=mock.ANY, ack_ids=[ack_id])]) mock_pubsub.return_value.close.assert_has_calls([mock.call()]) def test_read_messages_timestamp_attribute_fail_parse(self, mock_pubsub): data = b'data' attributes = {'time': '1337 unparseable'} publish_time_secs = 1520861821 publish_time_nanos = 234567000 ack_id = 'ack_id' pull_response = test_utils.create_pull_response([ test_utils.PullResponseMessage( data, attributes, publish_time_secs, publish_time_nanos, ack_id) ]) mock_pubsub.return_value.pull.return_value = pull_response options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True p = TestPipeline(options=options) _ = ( p | ReadFromPubSub( 'projects/fakeprj/topics/a_topic', None, None, with_attributes=True, timestamp_attribute='time')) with self.assertRaisesRegex(ValueError, r'parse'): p.run() mock_pubsub.return_value.acknowledge.assert_not_called() mock_pubsub.return_value.close.assert_has_calls([mock.call()]) def test_read_message_id_label_unsupported(self, unused_mock_pubsub): # id_label is unsupported in DirectRunner. options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True with self.assertRaisesRegex(NotImplementedError, r'id_label is not supported'): with TestPipeline(options=options) as p: _ = ( p | ReadFromPubSub( 'projects/fakeprj/topics/a_topic', None, 'a_label')) def test_runner_api_transformation_with_topic(self, unused_mock_pubsub): source = _PubSubSource( topic='projects/fakeprj/topics/a_topic', subscription=None, id_label='a_label', timestamp_attribute='b_label', with_attributes=True) transform = Read(source) context = pipeline_context.PipelineContext() proto_transform_spec = transform.to_runner_api(context) self.assertEqual( common_urns.composites.PUBSUB_READ.urn, proto_transform_spec.urn) pubsub_read_payload = ( proto_utils.parse_Bytes( proto_transform_spec.payload, beam_runner_api_pb2.PubSubReadPayload)) self.assertEqual( 'projects/fakeprj/topics/a_topic', pubsub_read_payload.topic) self.assertEqual('a_label', pubsub_read_payload.id_attribute) self.assertEqual('b_label', pubsub_read_payload.timestamp_attribute) self.assertEqual('', pubsub_read_payload.subscription) self.assertTrue(pubsub_read_payload.with_attributes) proto_transform = beam_runner_api_pb2.PTransform( unique_name="dummy_label", spec=proto_transform_spec) transform_from_proto = Read.from_runner_api_parameter( proto_transform, pubsub_read_payload, None) self.assertTrue(isinstance(transform_from_proto, Read)) self.assertTrue(isinstance(transform_from_proto.source, _PubSubSource)) self.assertEqual( 'projects/fakeprj/topics/a_topic', transform_from_proto.source.full_topic) self.assertTrue(transform_from_proto.source.with_attributes) def test_runner_api_transformation_properties_none(self, unused_mock_pubsub): # Confirming that properties stay None after a runner API transformation. source = _PubSubSource( topic='projects/fakeprj/topics/a_topic', with_attributes=True) transform = Read(source) context = pipeline_context.PipelineContext() proto_transform_spec = transform.to_runner_api(context) self.assertEqual( common_urns.composites.PUBSUB_READ.urn, proto_transform_spec.urn) pubsub_read_payload = ( proto_utils.parse_Bytes( proto_transform_spec.payload, beam_runner_api_pb2.PubSubReadPayload)) proto_transform = beam_runner_api_pb2.PTransform( unique_name="dummy_label", spec=proto_transform_spec) transform_from_proto = Read.from_runner_api_parameter( proto_transform, pubsub_read_payload, None) self.assertIsNone(transform_from_proto.source.full_subscription) self.assertIsNone(transform_from_proto.source.id_label) self.assertIsNone(transform_from_proto.source.timestamp_attribute) def test_runner_api_transformation_with_subscription( self, unused_mock_pubsub): source = _PubSubSource( topic=None, subscription='projects/fakeprj/subscriptions/a_subscription', id_label='a_label', timestamp_attribute='b_label', with_attributes=True) transform = Read(source) context = pipeline_context.PipelineContext() proto_transform_spec = transform.to_runner_api(context) self.assertEqual( common_urns.composites.PUBSUB_READ.urn, proto_transform_spec.urn) pubsub_read_payload = ( proto_utils.parse_Bytes( proto_transform_spec.payload, beam_runner_api_pb2.PubSubReadPayload)) self.assertEqual( 'projects/fakeprj/subscriptions/a_subscription', pubsub_read_payload.subscription) self.assertEqual('a_label', pubsub_read_payload.id_attribute) self.assertEqual('b_label', pubsub_read_payload.timestamp_attribute) self.assertEqual('', pubsub_read_payload.topic) self.assertTrue(pubsub_read_payload.with_attributes) proto_transform = beam_runner_api_pb2.PTransform( unique_name="dummy_label", spec=proto_transform_spec) transform_from_proto = Read.from_runner_api_parameter( proto_transform, pubsub_read_payload, None) self.assertTrue(isinstance(transform_from_proto, Read)) self.assertTrue(isinstance(transform_from_proto.source, _PubSubSource)) self.assertTrue(transform_from_proto.source.with_attributes) self.assertEqual( 'projects/fakeprj/subscriptions/a_subscription', transform_from_proto.source.full_subscription) def test_read_from_pubsub_no_overwrite(self, unused_mock): expected_elements = [ TestWindowedValue( b'apache', timestamp.Timestamp(1520861826.234567), [window.GlobalWindow()]), TestWindowedValue( b'beam', timestamp.Timestamp(1520861824.234567), [window.GlobalWindow()]) ] options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True # TODO(https://github.com/apache/beam/issues/34549): This test relies on # lineage metrics which Prism doesn't seem to handle correctly. Defaulting # to FnApiRunner instead. options.view_as(StandardOptions).runner = 'FnApiRunner' for test_case in ('topic', 'subscription'): with TestPipeline(options=options) as p: # Direct runner currently overwrites the whole ReadFromPubSub transform. # This test part of composite transform without overwrite. pcoll = p | beam.Create([b'apache', b'beam']) | beam.Map( lambda x: window.TimestampedValue(x, 1520861820.234567 + len(x))) args = {test_case: f'projects/fakeprj/{test_case}s/topic_or_sub'} pcoll = ReadFromPubSub(**args).expand_continued(pcoll) assert_that(pcoll, equal_to(expected_elements), reify_windows=True) self.assertSetEqual( Lineage.query(p.result.metrics(), Lineage.SOURCE), set([f"pubsub:{test_case}:fakeprj.topic_or_sub"])) @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') @mock.patch('google.cloud.pubsub.PublisherClient') class TestWriteToPubSub(unittest.TestCase): def test_write_messages_success(self, mock_pubsub): data = 'data' payloads = [data] options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: _ = ( p | Create(payloads) | WriteToPubSub( 'projects/fakeprj/topics/a_topic', with_attributes=False)) # Verify that publish was called (data will be protobuf serialized) mock_pubsub.return_value.publish.assert_called() # Check that the call was made with the topic and some data call_args = mock_pubsub.return_value.publish.call_args self.assertEqual(len(call_args[0]), 2) # topic and data def test_write_messages_deprecated(self, mock_pubsub): data = 'data' payloads = [data] options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: _ = ( p | Create(payloads) | WriteStringsToPubSub('projects/fakeprj/topics/a_topic')) # Verify that publish was called (data will be protobuf serialized) mock_pubsub.return_value.publish.assert_called() # Check that the call was made with the topic and some data call_args = mock_pubsub.return_value.publish.call_args self.assertEqual(len(call_args[0]), 2) # topic and data def test_write_messages_with_attributes_success(self, mock_pubsub): data = b'data' attributes = {'key': 'value'} payloads = [PubsubMessage(data, attributes)] options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: _ = ( p | Create(payloads) | WriteToPubSub( 'projects/fakeprj/topics/a_topic', with_attributes=True)) # Verify that publish was called (data will be protobuf serialized) mock_pubsub.return_value.publish.assert_called() # Check that the call was made with the topic and some data call_args = mock_pubsub.return_value.publish.call_args self.assertEqual(len(call_args[0]), 2) # topic and data def test_write_messages_batch_mode_success(self, mock_pubsub): """Test WriteToPubSub works in batch mode (non-streaming).""" data = 'data' payloads = [data] options = PipelineOptions([]) # Explicitly set streaming to False for batch mode options.view_as(StandardOptions).streaming = False with TestPipeline(options=options) as p: _ = ( p | Create(payloads) | WriteToPubSub( 'projects/fakeprj/topics/a_topic', with_attributes=False)) # Verify that publish was called (data will be protobuf serialized) mock_pubsub.return_value.publish.assert_called() # Check that the call was made with the topic and some data call_args = mock_pubsub.return_value.publish.call_args self.assertEqual(len(call_args[0]), 2) # topic and data def test_write_messages_with_attributes_batch_mode_success(self, mock_pubsub): """Test WriteToPubSub with attributes works in batch mode.""" data = b'data' attributes = {'key': 'value'} payloads = [PubsubMessage(data, attributes)] options = PipelineOptions([]) # Explicitly set streaming to False for batch mode options.view_as(StandardOptions).streaming = False with TestPipeline(options=options) as p: _ = ( p | Create(payloads) | WriteToPubSub( 'projects/fakeprj/topics/a_topic', with_attributes=True)) # Verify that publish was called (data will be protobuf serialized) mock_pubsub.return_value.publish.assert_called() # Check that the call was made with the topic and some data call_args = mock_pubsub.return_value.publish.call_args self.assertEqual(len(call_args[0]), 2) # topic and data def test_write_messages_with_attributes_error(self, mock_pubsub): data = 'data' # Sending raw data when WriteToPubSub expects a PubsubMessage object. payloads = [data] options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True with self.assertRaisesRegex(Exception, r'requires.*PubsubMessage.*applied.*str'): with TestPipeline(options=options) as p: _ = ( p | Create(payloads) | WriteToPubSub( 'projects/fakeprj/topics/a_topic', with_attributes=True)) def test_write_messages_unsupported_features(self, mock_pubsub): data = b'data' attributes = {'key': 'value'} payloads = [PubsubMessage(data, attributes)] options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True with self.assertRaisesRegex(NotImplementedError, r'id_label is not supported'): with TestPipeline(options=options) as p: _ = ( p | Create(payloads) | WriteToPubSub( 'projects/fakeprj/topics/a_topic', id_label='a_label', with_attributes=True)) options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True with self.assertRaisesRegex(NotImplementedError, r'timestamp_attribute is not supported'): with TestPipeline(options=options) as p: _ = ( p | Create(payloads) | WriteToPubSub( 'projects/fakeprj/topics/a_topic', timestamp_attribute='timestamp', with_attributes=True)) def test_runner_api_transformation(self, unused_mock_pubsub): sink = _PubSubSink( topic='projects/fakeprj/topics/a_topic', id_label=None, # We expect encoded PubSub write transform to always return attributes. timestamp_attribute=None) transform = Write(sink) context = pipeline_context.PipelineContext() proto_transform_spec = transform.to_runner_api(context) self.assertEqual( common_urns.composites.PUBSUB_WRITE.urn, proto_transform_spec.urn) pubsub_write_payload = ( proto_utils.parse_Bytes( proto_transform_spec.payload, beam_runner_api_pb2.PubSubWritePayload)) self.assertEqual( 'projects/fakeprj/topics/a_topic', pubsub_write_payload.topic) proto_transform = beam_runner_api_pb2.PTransform( unique_name="dummy_label", spec=proto_transform_spec) transform_from_proto = Write.from_runner_api_parameter( proto_transform, pubsub_write_payload, None) self.assertTrue(isinstance(transform_from_proto, Write)) self.assertTrue(isinstance(transform_from_proto.sink, _PubSubSink)) self.assertEqual( 'projects/fakeprj/topics/a_topic', transform_from_proto.sink.full_topic) def test_runner_api_transformation_properties_none(self, unused_mock_pubsub): # Confirming that properties stay None after a runner API transformation. sink = _PubSubSink( topic='projects/fakeprj/topics/a_topic', id_label=None, # We expect encoded PubSub write transform to always return attributes. timestamp_attribute=None) transform = Write(sink) context = pipeline_context.PipelineContext() proto_transform_spec = transform.to_runner_api(context) self.assertEqual( common_urns.composites.PUBSUB_WRITE.urn, proto_transform_spec.urn) pubsub_write_payload = ( proto_utils.parse_Bytes( proto_transform_spec.payload, beam_runner_api_pb2.PubSubWritePayload)) proto_transform = beam_runner_api_pb2.PTransform( unique_name="dummy_label", spec=proto_transform_spec) transform_from_proto = Write.from_runner_api_parameter( proto_transform, pubsub_write_payload, None) self.assertTrue(isinstance(transform_from_proto, Write)) self.assertTrue(isinstance(transform_from_proto.sink, _PubSubSink)) self.assertIsNone(transform_from_proto.sink.id_label) self.assertIsNone(transform_from_proto.sink.timestamp_attribute) def test_write_to_pubsub_no_overwrite(self, unused_mock): data = 'data' payloads = [data] options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True # TODO(https://github.com/apache/beam/issues/34549): This test relies on # lineage metrics which Prism doesn't seem to handle correctly. Defaulting # to FnApiRunner instead. options.view_as(StandardOptions).runner = 'FnApiRunner' with TestPipeline(options=options) as p: pcoll = p | Create(payloads) WriteToPubSub( 'projects/fakeprj/topics/a_topic', with_attributes=False).expand(pcoll) self.assertSetEqual( Lineage.query(p.result.metrics(), Lineage.SINK), set(["pubsub:topic:fakeprj.a_topic"])) def test_write_to_pubsub_with_attributes_no_overwrite(self, unused_mock): data = b'data' attributes = {'key': 'value'} payloads = [PubsubMessage(data, attributes)] options = PipelineOptions([]) options.view_as(StandardOptions).streaming = True # TODO(https://github.com/apache/beam/issues/34549): This test relies on # lineage metrics which Prism doesn't seem to handle correctly. Defaulting # to FnApiRunner instead. options.view_as(StandardOptions).runner = 'FnApiRunner' with TestPipeline(options=options) as p: pcoll = p | Create(payloads) # Avoid direct runner overwrites WriteToPubSub WriteToPubSub( 'projects/fakeprj/topics/a_topic', with_attributes=True).expand(pcoll) self.assertSetEqual( Lineage.query(p.result.metrics(), Lineage.SINK), set(["pubsub:topic:fakeprj.a_topic"])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/resource_identifiers.py0000644000175100017510000000414115206113032023623 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Helper functions to generate resource labels strings for GCP entitites These can be used on MonitoringInfo 'resource' labels. See example entities: https://s.apache.org/beam-gcp-debuggability For GCP entities, populate the RESOURCE label with the aip.dev/122 format: https://google.aip.dev/122 If an official GCP format does not exist, try to use the following format. //whatever.googleapis.com/parents/{parentId}/whatevers/{whateverId} """ def BigQueryTable(project_id, dataset_id, table_id): return '//bigquery.googleapis.com/projects/%s/datasets/%s/tables/%s' % ( project_id, dataset_id, table_id) def GoogleCloudStorageBucket(bucket_id): return '//storage.googleapis.com/buckets/%s' % bucket_id def DatastoreNamespace(project_id, namespace_id): return '//bigtable.googleapis.com/projects/%s/namespaces/%s' % ( project_id, namespace_id) def SpannerTable(project_id, database_id, table_id): return '//spanner.googleapis.com/projects/%s/topics/%s/tables/%s' % ( project_id, database_id, table_id) def SpannerSqlQuery(project_id, query_name): return '//spanner.googleapis.com/projects/%s/queries/%s' % ( project_id, query_name) def BigtableTable(project_id, instance_id, table_id): return '//bigtable.googleapis.com/projects/%s/instances/%s/tables/%s' % ( project_id, instance_id, table_id) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/spanner.py0000644000175100017510000006300215206113032021056 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """PTransforms for supporting Spanner in Python pipelines. These transforms are currently supported by Beam portable Flink and Spark runners. **Setup** Transforms provided in this module are cross-language transforms implemented in the Beam Java SDK. During the pipeline construction, Python SDK will connect to a Java expansion service to expand these transforms. To facilitate this, a small amount of setup is needed before using these transforms in a Beam Python pipeline. There are several ways to setup cross-language Spanner transforms. * Option 1: use the default expansion service * Option 2: specify a custom expansion service See below for details regarding each of these options. *Option 1: Use the default expansion service* This is the recommended and easiest setup option for using Python Spanner transforms. This option is only available for Beam 2.26.0 and later. This option requires following pre-requisites before running the Beam pipeline. * Install Java runtime in the computer from where the pipeline is constructed and make sure that 'java' command is available. In this option, Python SDK will either download (for released Beam version) or build (when running from a Beam Git clone) a expansion service jar and use that to expand transforms. Currently Spanner transforms use the 'beam-sdks-java-io-google-cloud-platform-expansion-service' jar for this purpose. *Option 2: specify a custom expansion service* In this option, you startup your own expansion service and provide that as a parameter when using the transforms provided in this module. This option requires following pre-requisites before running the Beam pipeline. * Startup your own expansion service. * Update your pipeline to provide the expansion service address when initiating Spanner transforms provided in this module. Flink Users can use the built-in Expansion Service of the Flink Runner's Job Server. If you start Flink's Job Server, the expansion service will be started on port 8097. For a different address, please set the expansion_service parameter. **More information** For more information regarding cross-language transforms see: - https://beam.apache.org/roadmap/portability/ For more information specific to Flink runner see: - https://beam.apache.org/documentation/runners/flink/ """ # pytype: skip-file from enum import Enum from enum import auto from typing import NamedTuple from typing import Optional from apache_beam.transforms.external import BeamJarExpansionService from apache_beam.transforms.external import ExternalTransform from apache_beam.transforms.external import NamedTupleBasedPayloadBuilder from apache_beam.typehints.schemas import named_tuple_to_schema __all__ = [ 'ReadFromSpanner', 'SpannerDelete', 'SpannerInsert', 'SpannerInsertOrUpdate', 'SpannerReplace', 'SpannerUpdate', 'TimestampBoundMode', 'TimeUnit', 'ReadChangeStreamFromSpanner', ] def default_io_expansion_service(): return BeamJarExpansionService( 'sdks:java:io:google-cloud-platform:expansion-service:shadowJar') class TimeUnit(Enum): NANOSECONDS = auto() MICROSECONDS = auto() MILLISECONDS = auto() SECONDS = auto() HOURS = auto() DAYS = auto() class TimestampBoundMode(Enum): MAX_STALENESS = auto() EXACT_STALENESS = auto() READ_TIMESTAMP = auto() MIN_READ_TIMESTAMP = auto() STRONG = auto() class FailureMode(Enum): FAIL_FAST = auto() REPORT_FAILURES = auto() class ReadFromSpannerSchema(NamedTuple): instance_id: str database_id: str schema: bytes sql: Optional[str] table: Optional[str] project_id: Optional[str] host: Optional[str] emulator_host: Optional[str] batching: Optional[bool] timestamp_bound_mode: Optional[str] read_timestamp: Optional[str] staleness: Optional[int] time_unit: Optional[str] class ReadChangeStreamFromSpannerSchema(NamedTuple): instance_id: str database_id: str project_id: str changeStreamName: str inclusiveStartAt: str inclusiveEndAt: Optional[str] metadataDatabase: str metadataInstance: str metadataTable: Optional[str] rpcPriority: Optional[str] watermarkRefreshRate: Optional[str] class ReadFromSpanner(ExternalTransform): """ A PTransform which reads from the specified Spanner instance's database. This transform required type of the row it has to return to provide the schema. Example:: from typing import NamedTuple from apache_beam import coders class ExampleRow(NamedTuple): id: int name: unicode coders.registry.register_coder(ExampleRow, coders.RowCoder) with Pipeline() as p: result = ( p | ReadFromSpanner( instance_id='your_instance_id', database_id='your_database_id', project_id='your_project_id', row_type=ExampleRow, sql='SELECT * FROM some_table', timestamp_bound_mode=TimestampBoundMode.MAX_STALENESS, staleness=3, time_unit=TimeUnit.HOURS, ).with_output_types(ExampleRow)) Experimental; no backwards compatibility guarantees. """ URN = 'beam:transform:org.apache.beam:spanner_read:v1' def __init__( self, project_id, instance_id, database_id, row_type=None, sql=None, table=None, host=None, emulator_host=None, batching=None, timestamp_bound_mode=None, read_timestamp=None, staleness=None, time_unit=None, expansion_service=None, ): """ Initializes a read operation from Spanner. :param project_id: Specifies the Cloud Spanner project. :param instance_id: Specifies the Cloud Spanner instance. :param database_id: Specifies the Cloud Spanner database. :param row_type: Row type that fits the given query or table. Passed as NamedTuple, e.g. NamedTuple('name', [('row_name', unicode)]) :param sql: An sql query to execute. It's results must fit the provided row_type. Don't use when table is set. :param table: A spanner table. When provided all columns from row_type will be selected to query. Don't use when query is set. :param batching: By default Batch API is used to read data from Cloud Spanner. It is useful to disable batching when the underlying query is not root-partitionable. :param host: Specifies the Cloud Spanner host. :param emulator_host: Specifies Spanner emulator host. :param timestamp_bound_mode: Defines how Cloud Spanner will choose a timestamp for a read-only transaction or a single read/query. Passed as TimestampBoundMode enum. Possible values: STRONG: A timestamp bound that will perform reads and queries at a timestamp where all previously committed transactions are visible. READ_TIMESTAMP: Returns a timestamp bound that will perform reads and queries at the given timestamp. MIN_READ_TIMESTAMP: Returns a timestamp bound that will perform reads and queries at a timestamp chosen to be at least given timestamp value. EXACT_STALENESS: Returns a timestamp bound that will perform reads and queries at an exact staleness. The timestamp is chosen soon after the read is started. MAX_STALENESS: Returns a timestamp bound that will perform reads and queries at a timestamp chosen to be at most time_unit stale. :param read_timestamp: Timestamp in string. Use only when timestamp_bound_mode is set to READ_TIMESTAMP or MIN_READ_TIMESTAMP. :param staleness: Staleness value as int. Use only when timestamp_bound_mode is set to EXACT_STALENESS or MAX_STALENESS. time_unit has to be set along with this param. :param time_unit: Time unit for staleness_value passed as TimeUnit enum. Possible values: NANOSECONDS, MICROSECONDS, MILLISECONDS, SECONDS, HOURS, DAYS. :param expansion_service: The address (host:port) of the ExpansionService. """ assert row_type assert sql or table and not (sql and table) staleness_value = int(staleness) if staleness else None if staleness_value or time_unit: assert staleness_value and time_unit and \ timestamp_bound_mode is TimestampBoundMode.MAX_STALENESS or \ timestamp_bound_mode is TimestampBoundMode.EXACT_STALENESS if read_timestamp: assert timestamp_bound_mode is TimestampBoundMode.MIN_READ_TIMESTAMP\ or timestamp_bound_mode is TimestampBoundMode.READ_TIMESTAMP super().__init__( self.URN, NamedTupleBasedPayloadBuilder( ReadFromSpannerSchema( instance_id=instance_id, database_id=database_id, sql=sql, table=table, schema=named_tuple_to_schema(row_type).SerializeToString(), project_id=project_id, host=host, emulator_host=emulator_host, batching=batching, timestamp_bound_mode=_get_enum_name(timestamp_bound_mode), read_timestamp=read_timestamp, staleness=staleness, time_unit=_get_enum_name(time_unit), ), ), expansion_service or default_io_expansion_service(), ) class WriteToSpannerSchema(NamedTuple): project_id: str instance_id: str database_id: str table: str max_batch_size_bytes: Optional[int] max_number_mutations: Optional[int] max_number_rows: Optional[int] grouping_factor: Optional[int] host: Optional[str] emulator_host: Optional[str] commit_deadline: Optional[int] max_cumulative_backoff: Optional[int] failure_mode: Optional[str] high_priority: bool _CLASS_DOC = \ """ A PTransform which writes {operation} mutations to the specified Spanner table. This transform receives rows defined as NamedTuple. Example:: from typing import NamedTuple from apache_beam import coders class {row_type}(NamedTuple): id: int name: unicode coders.registry.register_coder({row_type}, coders.RowCoder) with Pipeline() as p: _ = ( p | 'Impulse' >> beam.Impulse() | 'Generate' >> beam.FlatMap(lambda x: range(num_rows)) | 'To row' >> beam.Map(lambda n: {row_type}(n, str(n)) .with_output_types({row_type}) | 'Write to Spanner' >> Spanner{operation_suffix}( instance_id='your_instance', database_id='existing_database', project_id='your_project_id', table='your_table')) Experimental; no backwards compatibility guarantees. """ _INIT_DOC = \ """ Initializes {operation} operation to a Spanner table. :param project_id: Specifies the Cloud Spanner project. :param instance_id: Specifies the Cloud Spanner instance. :param database_id: Specifies the Cloud Spanner database. :param table: Specifies the Cloud Spanner table. :param max_batch_size_bytes: Specifies the batch size limit (max number of bytes mutated per batch). Default value is 1048576 bytes = 1MB. :param max_number_mutations: Specifies the cell mutation limit (maximum number of mutated cells per batch). Default value is 5000. :param max_number_rows: Specifies the row mutation limit (maximum number of mutated rows per batch). Default value is 500. :param grouping_factor: Specifies the multiple of max mutation (in terms of both bytes per batch and cells per batch) that is used to select a set of mutations to sort by key for batching. This sort uses local memory on the workers, so using large values can cause out of memory errors. Default value is 1000. :param host: Specifies the Cloud Spanner host. :param emulator_host: Specifies Spanner emulator host. :param commit_deadline: Specifies the deadline for the Commit API call. Default is 15 secs. DEADLINE_EXCEEDED errors will prompt a backoff/retry until the value of commit_deadline is reached. DEADLINE_EXCEEDED errors are ar reported with logging and counters. Pass seconds as value. :param max_cumulative_backoff: Specifies the maximum cumulative backoff time when retrying after DEADLINE_EXCEEDED errors. Default is 900s (15min). If the mutations still have not been written after this time, they are treated as a failure, and handled according to the setting of failure_mode. Pass seconds as value. :param failure_mode: Specifies the behavior for mutations that fail to be written to Spanner. Default is FAIL_FAST. When FAIL_FAST is set, an exception will be thrown for any failed mutation. When REPORT_FAILURES is set, processing will continue instead of throwing an exception. Note that REPORT_FAILURES can cause data loss if used incorrectly. :param expansion_service: The address (host:port) of the ExpansionService. """ def _add_doc( value, operation=None, row_type=None, operation_suffix=None, ): def _doc(obj): obj.__doc__ = value.format( operation=operation, row_type=row_type, operation_suffix=operation_suffix, ) return obj return _doc @_add_doc( _CLASS_DOC, operation='delete', row_type='ExampleKey', operation_suffix='Delete', ) class SpannerDelete(ExternalTransform): URN = 'beam:transform:org.apache.beam:spanner_delete:v1' @_add_doc(_INIT_DOC, operation='a delete') def __init__( self, project_id, instance_id, database_id, table, max_batch_size_bytes=None, max_number_mutations=None, max_number_rows=None, grouping_factor=None, host=None, emulator_host=None, commit_deadline=None, max_cumulative_backoff=None, failure_mode=None, expansion_service=None, high_priority=False, ): max_cumulative_backoff = int( max_cumulative_backoff) if max_cumulative_backoff else None commit_deadline = int(commit_deadline) if commit_deadline else None super().__init__( self.URN, NamedTupleBasedPayloadBuilder( WriteToSpannerSchema( project_id=project_id, instance_id=instance_id, database_id=database_id, table=table, max_batch_size_bytes=max_batch_size_bytes, max_number_mutations=max_number_mutations, max_number_rows=max_number_rows, grouping_factor=grouping_factor, host=host, emulator_host=emulator_host, commit_deadline=commit_deadline, max_cumulative_backoff=max_cumulative_backoff, failure_mode=_get_enum_name(failure_mode), high_priority=high_priority, ), ), expansion_service=expansion_service or default_io_expansion_service(), ) @_add_doc( _CLASS_DOC, operation='insert', row_type='ExampleRow', operation_suffix='Insert', ) class SpannerInsert(ExternalTransform): URN = 'beam:transform:org.apache.beam:spanner_insert:v1' @_add_doc(_INIT_DOC, operation='an insert') def __init__( self, project_id, instance_id, database_id, table, max_batch_size_bytes=None, max_number_mutations=None, max_number_rows=None, grouping_factor=None, host=None, emulator_host=None, commit_deadline=None, max_cumulative_backoff=None, expansion_service=None, failure_mode=None, high_priority=False, ): max_cumulative_backoff = int( max_cumulative_backoff) if max_cumulative_backoff else None commit_deadline = int(commit_deadline) if commit_deadline else None super().__init__( self.URN, NamedTupleBasedPayloadBuilder( WriteToSpannerSchema( project_id=project_id, instance_id=instance_id, database_id=database_id, table=table, max_batch_size_bytes=max_batch_size_bytes, max_number_mutations=max_number_mutations, max_number_rows=max_number_rows, grouping_factor=grouping_factor, host=host, emulator_host=emulator_host, commit_deadline=commit_deadline, max_cumulative_backoff=max_cumulative_backoff, failure_mode=_get_enum_name(failure_mode), high_priority=high_priority, ), ), expansion_service=expansion_service or default_io_expansion_service(), ) @_add_doc( _CLASS_DOC, operation='replace', row_type='ExampleRow', operation_suffix='Replace', ) class SpannerReplace(ExternalTransform): URN = 'beam:transform:org.apache.beam:spanner_replace:v1' @_add_doc(_INIT_DOC, operation='a replace') def __init__( self, project_id, instance_id, database_id, table, max_batch_size_bytes=None, max_number_mutations=None, max_number_rows=None, grouping_factor=None, host=None, emulator_host=None, commit_deadline=None, max_cumulative_backoff=None, expansion_service=None, failure_mode=None, high_priority=False, ): max_cumulative_backoff = int( max_cumulative_backoff) if max_cumulative_backoff else None commit_deadline = int(commit_deadline) if commit_deadline else None super().__init__( self.URN, NamedTupleBasedPayloadBuilder( WriteToSpannerSchema( project_id=project_id, instance_id=instance_id, database_id=database_id, table=table, max_batch_size_bytes=max_batch_size_bytes, max_number_mutations=max_number_mutations, max_number_rows=max_number_rows, grouping_factor=grouping_factor, host=host, emulator_host=emulator_host, commit_deadline=commit_deadline, max_cumulative_backoff=max_cumulative_backoff, failure_mode=_get_enum_name(failure_mode), high_priority=high_priority, ), ), expansion_service=expansion_service or default_io_expansion_service(), ) @_add_doc( _CLASS_DOC, operation='insert-or-update', row_type='ExampleRow', operation_suffix='InsertOrUpdate', ) class SpannerInsertOrUpdate(ExternalTransform): URN = 'beam:transform:org.apache.beam:spanner_insert_or_update:v1' @_add_doc(_INIT_DOC, operation='an insert-or-update') def __init__( self, project_id, instance_id, database_id, table, max_batch_size_bytes=None, max_number_mutations=None, max_number_rows=None, grouping_factor=None, host=None, emulator_host=None, commit_deadline=None, max_cumulative_backoff=None, failure_mode=None, expansion_service=None, high_priority=False, ): max_cumulative_backoff = int( max_cumulative_backoff) if max_cumulative_backoff else None commit_deadline = int(commit_deadline) if commit_deadline else None super().__init__( self.URN, NamedTupleBasedPayloadBuilder( WriteToSpannerSchema( project_id=project_id, instance_id=instance_id, database_id=database_id, table=table, max_batch_size_bytes=max_batch_size_bytes, max_number_mutations=max_number_mutations, max_number_rows=max_number_rows, grouping_factor=grouping_factor, host=host, emulator_host=emulator_host, commit_deadline=commit_deadline, max_cumulative_backoff=max_cumulative_backoff, failure_mode=_get_enum_name(failure_mode), high_priority=high_priority, ), ), expansion_service=expansion_service or default_io_expansion_service(), ) @_add_doc( _CLASS_DOC, operation='update', row_type='ExampleRow', operation_suffix='Update', ) class SpannerUpdate(ExternalTransform): URN = 'beam:transform:org.apache.beam:spanner_update:v1' @_add_doc(_INIT_DOC, operation='an update') def __init__( self, project_id, instance_id, database_id, table, max_batch_size_bytes=None, max_number_mutations=None, max_number_rows=None, grouping_factor=None, host=None, emulator_host=None, commit_deadline=None, max_cumulative_backoff=None, failure_mode=None, expansion_service=None, high_priority=False, ): max_cumulative_backoff = int( max_cumulative_backoff) if max_cumulative_backoff else None commit_deadline = int(commit_deadline) if commit_deadline else None super().__init__( self.URN, NamedTupleBasedPayloadBuilder( WriteToSpannerSchema( project_id=project_id, instance_id=instance_id, database_id=database_id, table=table, max_batch_size_bytes=max_batch_size_bytes, max_number_mutations=max_number_mutations, max_number_rows=max_number_rows, grouping_factor=grouping_factor, host=host, emulator_host=emulator_host, commit_deadline=commit_deadline, max_cumulative_backoff=max_cumulative_backoff, failure_mode=_get_enum_name(failure_mode), high_priority=high_priority, ), ), expansion_service=expansion_service or default_io_expansion_service(), ) class ReadChangeStreamFromSpanner(ExternalTransform): """ A PTransform to read Change Streams from Google Cloud Spanner. The output of this transform is a PCollection of JSON strings, where each string represents a com.google.cloud.spanner.DataChangeRecord. Example: with beam.Pipeline(options=pipeline_options) as p: p | "ReadSpannerChangeStream" >> beam_spanner.ReadChangeStreamFromSpanner( project_id="spanner-project-id", instance_id="spanner-instance-id", database_id="spanner-database-id", changeStreamName="spanner-change-stream", inclusiveStartAt="2025-05-20T10:00:00Z", metadataDatabase="spanner-metadata-database", metadataInstance="spanner-metadata-instance") Experimental; no backwards compatibility guarantees. """ URN = 'beam:transform:org.apache.beam:spanner_change_stream_reader:v1' def __init__( self, project_id, instance_id, database_id, changeStreamName, metadataDatabase, metadataInstance, inclusiveStartAt, inclusiveEndAt=None, metadataTable=None, rpcPriority=None, watermarkRefreshRate=None, expansion_service=None, ): """ Reads Change Streams from Google Cloud Spanner. :param project_id: (Required) Specifies the Cloud Spanner project. :param instance_id: (Required) Specifies the Cloud Spanner instance. :param database_id: (Required) Specifies the Cloud Spanner database. :param changeStreamName: (Required) The name of the Spanner change stream to read. :param metadataDatabase: (Required) The database where the change stream metadata is stored. :param metadataInstance: (Required) The instance where the change stream metadata database resides. :param inclusiveStartAt: (Required) An inclusive start timestamp for reading the change stream. :param inclusiveEndAt: (Optional) An inclusive end timestamp for reading the change stream. If not specified, the stream will be read indefinitely. :param metadataTable: (Optional) The name of the metadata table used by the change stream connector. If not specified, a default table name will be used. :param rpcPriority: (Optional) The RPC priority for Spanner operations. Can be 'HIGH', 'MEDIUM', or 'LOW'. :param watermarkRefreshRate: (Optional) The duration at which the watermark is refreshed. """ super().__init__( self.URN, NamedTupleBasedPayloadBuilder( ReadChangeStreamFromSpannerSchema( instance_id=instance_id, database_id=database_id, project_id=project_id, changeStreamName=changeStreamName, inclusiveStartAt=inclusiveStartAt, inclusiveEndAt=inclusiveEndAt, metadataDatabase=metadataDatabase, metadataInstance=metadataInstance, metadataTable=metadataTable, rpcPriority=rpcPriority, watermarkRefreshRate=watermarkRefreshRate, ), ), expansion_service=expansion_service or default_io_expansion_service(), ) def _get_enum_name(enum): return None if enum is None else enum.name ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/spanner_wrapper.py0000644000175100017510000000476315206113032022627 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the 'License'); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an 'AS IS' BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import uuid from apache_beam.utils import retry try: from google.cloud import spanner except ImportError: spanner = None _LOGGER = logging.getLogger(__name__) MAX_RETRIES = 3 class SpannerWrapper(object): TEMP_DATABASE_PREFIX = 'temp-' def __init__(self, project_id, temp_database_id=None): self._spanner_client = spanner.Client(project=project_id) self._spanner_instance = self._spanner_client.instance("beam-test") self._test_database = None if temp_database_id and temp_database_id.startswith( self.TEMP_DATABASE_PREFIX): raise ValueError( 'User provided temp database ID cannot start with %r' % self.TEMP_DATABASE_PREFIX) if temp_database_id is not None: self._test_database = temp_database_id else: self._test_database = self._get_temp_database() def _get_temp_database(self): uniq_id = uuid.uuid4().hex[:10] return f'{self.TEMP_DATABASE_PREFIX}{uniq_id}' @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _create_database(self): _LOGGER.info('Creating test database: %s' % self._test_database) instance = self._spanner_instance database = instance.database( self._test_database, ddl_statements=[ '''CREATE TABLE tmp_table ( UserId STRING(256) NOT NULL, Key STRING(1024) ) PRIMARY KEY (UserId)''' ]) operation = database.create() _LOGGER.info('Creating database: Done! %s' % str(operation.result())) def _delete_database(self): if (self._spanner_instance): database = self._spanner_instance.database(self._test_database) database.drop() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.3984096 apache_beam-2.74.0/apache_beam/io/gcp/tests/0000755000175100017510000000000015206113157020207 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/tests/__init__.py0000644000175100017510000000142015206113032022305 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/tests/bigquery_matcher.py0000644000175100017510000002174115206113032024110 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Bigquery data verifier for end-to-end test.""" # pytype: skip-file import concurrent import logging import time from hamcrest.core.base_matcher import BaseMatcher from apache_beam.io.gcp import bigquery_tools from apache_beam.testing.test_utils import compute_hash from apache_beam.testing.util import BeamAssertException from apache_beam.testing.util import equal_to from apache_beam.utils import retry __all__ = ['BigqueryMatcher', 'BigQueryTableMatcher'] # Protect against environments where bigquery library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from google.cloud import bigquery from google.cloud.exceptions import GoogleCloudError except ImportError: bigquery = None # pylint: enable=wrong-import-order, wrong-import-position MAX_RETRIES = 5 _LOGGER = logging.getLogger(__name__) def retry_on_http_timeout_and_value_error(exception): """Filter allowing retries on Bigquery errors and value error.""" return isinstance( exception, (GoogleCloudError, ValueError, concurrent.futures.TimeoutError)) class BigqueryMatcher(BaseMatcher): """Matcher that verifies the checksum of Bigquery data with given query. Fetch Bigquery data with given query, compute a hash string and compare with expected checksum. """ def __init__(self, project, query, checksum, timeout_secs=0): """Initialize BigQueryMatcher object. Args: project: The name (string) of the project. query: The query (string) to perform. checksum: SHA-1 hash generated from a sorted list of lines read from expected output. timeout_secs: Duration to retry query until checksum matches. This is useful for DF streaming pipelines or BQ streaming inserts. The default (0) never retries. """ if bigquery is None: raise ImportError('Bigquery dependencies are not installed.') if not query or not isinstance(query, str): raise ValueError('Invalid argument: query. Please use non-empty string') if not checksum or not isinstance(checksum, str): raise ValueError( 'Invalid argument: checksum. Please use non-empty string') self.project = project self.query = query self.expected_checksum = checksum self.checksum = None self.timeout_secs = timeout_secs def _matches(self, _): @retry.with_exponential_backoff( num_retries=1000, initial_delay_secs=0.5, max_delay_secs=30, stop_after_secs=self.timeout_secs, ) def get_checksum(): response = self._query_with_retry() _LOGGER.info( 'Read from given query (%s), total rows %d', self.query, len(response)) self.checksum = compute_hash(response) _LOGGER.info('Generate checksum: %s', self.checksum) if self.checksum != self.expected_checksum: # This exception is never raised beyond the enclosing method. raise ValueError( 'Checksums do not match. Expected: %s, got: %s' % (self.expected_checksum, self.checksum)) if self.checksum is None: try: get_checksum() except ValueError: pass return self.checksum == self.expected_checksum @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry_on_http_timeout_and_value_error) def _query_with_retry(self): """Run Bigquery query with retry if got error http response""" _LOGGER.info('Attempting to perform query %s to BQ', self.query) # Create client here since it throws an exception if pickled. bigquery_client = bigquery.Client(self.project) query_job = bigquery_client.query(self.query) rows = query_job.result(timeout=60) return [row.values() for row in rows] def describe_to(self, description): description \ .append_text("Expected checksum is ") \ .append_text(self.expected_checksum) def describe_mismatch(self, pipeline_result, mismatch_description): mismatch_description \ .append_text("Actual checksum is ") \ .append_text(self.checksum) class BigqueryFullResultMatcher(BigqueryMatcher): """Matcher that verifies Bigquery data with given query. Fetch Bigquery data with given query, compare to the expected data. """ def __init__(self, project, query, data): """Initialize BigQueryMatcher object. Args: project: The name (string) of the project. query: The query (string) to perform. data: List of tuples with the expected data. """ super().__init__(project, query, 'unused_checksum') self.expected_data = data self.actual_data = None def _matches(self, _): if self.actual_data is None: self.actual_data = self._get_query_result() _LOGGER.info('Result of query is: %r', self.actual_data) try: equal_to(self.expected_data)(self.actual_data) return True except BeamAssertException: return False def _get_query_result(self): return self._query_with_retry() def describe_to(self, description): description \ .append_text("Expected data is ") \ .append_text(self.expected_data) def describe_mismatch(self, pipeline_result, mismatch_description): mismatch_description \ .append_text("Actual data is ") \ .append_text(self.actual_data) class BigqueryFullResultStreamingMatcher(BigqueryFullResultMatcher): """ Matcher that verifies Bigquery data with given query. Fetch Bigquery data with given query, compare to the expected data. This matcher polls BigQuery until the no. of records in BigQuery is equal to the no. of records in expected data. A timeout can be specified. """ DEFAULT_TIMEOUT = 5 * 60 def __init__(self, project, query, data, timeout=DEFAULT_TIMEOUT): super().__init__(project, query, data) self.timeout = timeout def _get_query_result(self): start_time = time.time() while time.time() - start_time <= self.timeout: response = self._query_with_retry() if len(response) >= len(self.expected_data): return response _LOGGER.debug('Query result contains %d rows' % len(response)) time.sleep(1) raise TimeoutError('Timeout exceeded for matcher.') # noqa: F821 class BigQueryTableMatcher(BaseMatcher): """Matcher that verifies the properties of a Table in BigQuery.""" def __init__(self, project, dataset, table, expected_properties): if bigquery is None: raise ImportError('Bigquery dependencies are not installed.') self.project = project self.dataset = dataset self.table = table self.expected_properties = expected_properties @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry_on_http_timeout_and_value_error) def _get_table_with_retry(self, bigquery_wrapper): return bigquery_wrapper.get_table(self.project, self.dataset, self.table) def _matches(self, _): _LOGGER.info('Start verify Bigquery table properties.') # Run query bigquery_wrapper = bigquery_tools.BigQueryWrapper() self.actual_table = self._get_table_with_retry(bigquery_wrapper) _LOGGER.info('Table proto is %s', self.actual_table) return all( self._match_property(v, self._get_or_none(self.actual_table, k)) for k, v in self.expected_properties.items()) @staticmethod def _get_or_none(obj, attr): try: return obj.__getattribute__(attr) except AttributeError: try: return obj.get(attr, None) except TypeError: return None @staticmethod def _match_property(expected, actual): _LOGGER.info("Matching %s to %s", expected, actual) if isinstance(expected, dict): return all( BigQueryTableMatcher._match_property( v, BigQueryTableMatcher._get_or_none(actual, k)) for k, v in expected.items()) else: return expected == actual def describe_to(self, description): description \ .append_text("Expected table attributes are ") \ .append_text(sorted((k, v) for k, v in self.expected_properties.items())) def describe_mismatch(self, pipeline_result, mismatch_description): mismatch_description \ .append_text("Actual table attributes are ") \ .append_text(sorted((k, self._get_or_none(self.actual_table, k)) for k in self.expected_properties)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/tests/bigquery_matcher_test.py0000644000175100017510000001432215206113032025144 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit test for Bigquery verifier""" # pytype: skip-file import logging import unittest import mock from hamcrest import assert_that as hc_assert_that from apache_beam.io.gcp import bigquery_tools from apache_beam.io.gcp.tests import bigquery_matcher as bq_verifier from apache_beam.testing.test_utils import patch_retry # Protect against environments where bigquery library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from google.cloud import bigquery from google.cloud.exceptions import NotFound except ImportError: bigquery = None NotFound = None # pylint: enable=wrong-import-order, wrong-import-position @unittest.skipIf(bigquery is None, 'Bigquery dependencies are not installed.') @mock.patch.object(bigquery, 'Client') class BigqueryMatcherTest(unittest.TestCase): def setUp(self): self._mock_result = mock.Mock() patch_retry(self, bq_verifier) def test_bigquery_matcher_success(self, mock_bigquery): mock_query_result = [mock.Mock(), mock.Mock(), mock.Mock()] mock_query_result[0].values.return_value = [] mock_query_result[1].values.return_value = None mock_query_result[2].values.return_value = None mock_query = mock_bigquery.return_value.query mock_query.return_value.result.return_value = mock_query_result matcher = bq_verifier.BigqueryMatcher( 'mock_project', 'mock_query', '59f9d6bdee30d67ea73b8aded121c3a0280f9cd8') hc_assert_that(self._mock_result, matcher) self.assertEqual(1, mock_query.call_count) def test_bigquery_matcher_success_streaming_retry(self, mock_bigquery): # Simulate case where a streaming insert takes time to process, such that # the first query result is incomplete (empty). empty_query_result = [] mock_query_result = [mock.Mock(), mock.Mock(), mock.Mock()] mock_query_result[0].values.return_value = [] mock_query_result[1].values.return_value = None mock_query_result[2].values.return_value = None mock_query = mock_bigquery.return_value.query mock_query.return_value.result.side_effect = [ empty_query_result, mock_query_result ] matcher = bq_verifier.BigqueryMatcher( 'mock_project', 'mock_query', '59f9d6bdee30d67ea73b8aded121c3a0280f9cd8', timeout_secs=5, ) hc_assert_that(self._mock_result, matcher) self.assertEqual(2, mock_query.call_count) def test_bigquery_matcher_query_error_retry(self, mock_bigquery): mock_query = mock_bigquery.return_value.query mock_query.side_effect = NotFound('table not found') matcher = bq_verifier.BigqueryMatcher( 'mock_project', 'mock_query', 'mock_checksum') with self.assertRaises(NotFound): hc_assert_that(self._mock_result, matcher) self.assertEqual(bq_verifier.MAX_RETRIES + 1, mock_query.call_count) def test_bigquery_matcher_query_error_checksum(self, mock_bigquery): empty_query_result = [] mock_query = mock_bigquery.return_value.query mock_query.return_value.result.return_value = empty_query_result matcher = bq_verifier.BigqueryMatcher( 'mock_project', 'mock_query', '59f9d6bdee30d67ea73b8aded121c3a0280f9cd8', ) with self.assertRaisesRegex(AssertionError, r'Expected checksum'): hc_assert_that(self._mock_result, matcher) self.assertEqual(1, mock_query.call_count) @unittest.skipIf(bigquery is None, 'Bigquery dependencies are not installed.') @mock.patch.object(bigquery_tools, 'BigQueryWrapper') class BigqueryTableMatcherTest(unittest.TestCase): def setUp(self): self._mock_result = mock.Mock() patch_retry(self, bq_verifier) def test_bigquery_table_matcher_success(self, mock_bigquery): mock_query_result = mock.Mock( partitioning='a lot of partitioning', clustering={'column': 'FRIENDS'}) mock_bigquery.return_value.get_table.return_value = mock_query_result matcher = bq_verifier.BigQueryTableMatcher( 'mock_project', 'mock_dataset', 'mock_table', { 'partitioning': 'a lot of partitioning', 'clustering': { 'column': 'FRIENDS' } }) hc_assert_that(self._mock_result, matcher) def test_bigquery_table_matcher_query_error_retry(self, mock_bigquery): mock_query = mock_bigquery.return_value.get_table mock_query.side_effect = ValueError('table not found') matcher = bq_verifier.BigQueryTableMatcher( 'mock_project', 'mock_dataset', 'mock_table', { 'partitioning': 'a lot of partitioning', 'clustering': { 'column': 'FRIENDS' } }) with self.assertRaises(ValueError): hc_assert_that(self._mock_result, matcher) self.assertEqual(bq_verifier.MAX_RETRIES + 1, mock_query.call_count) @unittest.skipIf(bigquery is None, 'Bigquery dependencies are not installed.') @mock.patch.object(bigquery, 'Client') class BigqueryFullResultStreamingMatcherTest(unittest.TestCase): def setUp(self): self.timeout = 0.01 def test__get_query_result_timeout(self, mock_bigquery): mock_query = mock_bigquery.return_value.query mock_query.return_value.result.return_value = [] matcher = bq_verifier.BigqueryFullResultStreamingMatcher( 'some-project', 'some-query', [1, 2, 3], timeout=self.timeout) with self.assertRaises(TimeoutError): # noqa: F821 matcher._get_query_result() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/tests/pubsub_matcher.py0000644000175100017510000001712615206113032023563 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """PubSub verifier used for end-to-end test.""" # pytype: skip-file import logging import time from collections import Counter from hamcrest.core.base_matcher import BaseMatcher from apache_beam.io.gcp.pubsub import PubsubMessage __all__ = ['PubSubMessageMatcher'] # Protect against environments where pubsub library is not available. try: from google.cloud import pubsub except ImportError: pubsub = None DEFAULT_TIMEOUT = 5 * 60 DEFAULT_SLEEP_TIME = 1 DEFAULT_MAX_MESSAGES_IN_ONE_PULL = 50 DEFAULT_PULL_TIMEOUT = 30.0 _LOGGER = logging.getLogger(__name__) class PubSubMessageMatcher(BaseMatcher): """Matcher that verifies messages from given subscription. This matcher can block the test and keep pulling messages from given subscription until all expected messages are shown or timeout. """ def __init__( self, project, sub_name, expected_msg=None, expected_msg_len=None, timeout=DEFAULT_TIMEOUT, with_attributes=False, strip_attributes=None, sleep_time=DEFAULT_SLEEP_TIME, max_messages_in_one_pull=DEFAULT_MAX_MESSAGES_IN_ONE_PULL, pull_timeout=DEFAULT_PULL_TIMEOUT): """Initialize PubSubMessageMatcher object. Args: project: A name string of project. sub_name: A name string of subscription which is attached to output. expected_msg: A string list that contains expected message data pulled from the subscription. See also: with_attributes. expected_msg_len: Number of expected messages pulled from the subscription. timeout: Timeout in seconds to wait for all expected messages appears. with_attributes: If True, will match against both message data and attributes. If True, expected_msg should be a list of ``PubsubMessage`` objects. Otherwise, it should be a list of ``bytes``. strip_attributes: List of strings. If with_attributes==True, strip the attributes keyed by these values from incoming messages. If a key is missing, will add an attribute with an error message as value to prevent a successful match. sleep_time: Time in seconds between which the pulls from pubsub are done. max_messages_in_one_pull: Maximum number of messages pulled from pubsub at once. pull_timeout: Time in seconds after which the pull from pubsub is repeated """ if pubsub is None: raise ImportError('PubSub dependencies are not installed.') if not project: raise ValueError('Invalid project %s.' % project) if not sub_name: raise ValueError('Invalid subscription %s.' % sub_name) if not expected_msg_len and not expected_msg: raise ValueError( 'Required expected_msg: {} or expected_msg_len: {}.'.format( expected_msg, expected_msg_len)) if expected_msg and not isinstance(expected_msg, list): raise ValueError('Invalid expected messages %s.' % expected_msg) if expected_msg_len and not isinstance(expected_msg_len, int): raise ValueError('Invalid expected messages %s.' % expected_msg_len) self.project = project self.sub_name = sub_name self.expected_msg = expected_msg self.expected_msg_len = expected_msg_len or len(self.expected_msg) self.timeout = timeout self.messages = None self.messages_all_details = None self.with_attributes = with_attributes self.strip_attributes = strip_attributes self.sleep_time = sleep_time self.max_messages_in_one_pull = max_messages_in_one_pull self.pull_timeout = pull_timeout def _matches(self, _): if self.messages is None: self.messages, self.messages_all_details = self._wait_for_messages( self.expected_msg_len, self.timeout) if self.expected_msg: return Counter(self.messages) == Counter(self.expected_msg) else: return len(self.messages) == self.expected_msg_len def _wait_for_messages(self, expected_num, timeout): """Wait for messages from given subscription.""" total_messages = [] total_messages_all_details = [] sub_client = pubsub.SubscriberClient() start_time = time.time() while time.time() - start_time <= timeout: response = sub_client.pull( subscription=self.sub_name, max_messages=self.max_messages_in_one_pull, timeout=self.pull_timeout) for rm in response.received_messages: msg = PubsubMessage._from_message(rm.message) full_message = ( msg.data, msg.attributes, msg.attributes, msg.publish_time, msg.ordering_key) if not self.with_attributes: total_messages.append(msg.data) total_messages_all_details.append(full_message) continue if self.strip_attributes: for attr in self.strip_attributes: try: del msg.attributes[attr] except KeyError: msg.attributes[attr] = ( 'PubSubMessageMatcher error: ' 'expected attribute not found.') total_messages.append(msg) total_messages_all_details.append(full_message) ack_ids = [rm.ack_id for rm in response.received_messages] if ack_ids: sub_client.acknowledge(subscription=self.sub_name, ack_ids=ack_ids) if len(total_messages) >= expected_num: break time.sleep(self.sleep_time) if time.time() - start_time > timeout: _LOGGER.error( 'Timeout after %d sec. Received %d messages from %s.', timeout, len(total_messages), self.sub_name) return total_messages, total_messages_all_details def describe_to(self, description): description.append_text('Expected %d messages.' % self.expected_msg_len) def describe_mismatch(self, _, mismatch_description): c_expected = Counter(self.expected_msg) c_actual = Counter(self.messages) mismatch_description.append_text("Got %d messages. " % (len(self.messages))) if self.expected_msg: expected = (c_expected - c_actual).items() unexpected = (c_actual - c_expected).items() unexpected_keys = [repr(item[0]) for item in unexpected] if self.with_attributes: unexpected_all_details = [ x for x in self.messages_all_details if 'PubsubMessage(%s, %s)' % (repr(x[0]), x[1]) in unexpected_keys ] else: unexpected_all_details = [ x for x in self.messages_all_details if repr(x[0]) in unexpected_keys ] mismatch_description.append_text( "Diffs (item, count):\n" " Expected but not in actual: %s\n" " Unexpected: %s\n" " Unexpected (with all details): %s" % (expected, unexpected, unexpected_all_details)) if self.with_attributes and self.strip_attributes: mismatch_description.append_text( '\n Stripped attributes: %r' % self.strip_attributes) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/tests/pubsub_matcher_test.py0000644000175100017510000001727615206113032024630 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit test for PubSub verifier.""" # pytype: skip-file import logging import unittest import mock from hamcrest import assert_that as hc_assert_that from apache_beam.io.gcp.pubsub import PubsubMessage from apache_beam.io.gcp.tests.pubsub_matcher import PubSubMessageMatcher from apache_beam.testing.test_utils import PullResponseMessage from apache_beam.testing.test_utils import create_pull_response try: from google.cloud import pubsub except ImportError: pubsub = None @unittest.skipIf(pubsub is None, 'PubSub dependencies are not installed.') @mock.patch('time.sleep', return_value=None) @mock.patch('google.cloud.pubsub.SubscriberClient') class PubSubMatcherTest(unittest.TestCase): def setUp(self): self.mock_presult = mock.MagicMock() def init_matcher( self, expected_msg=None, with_attributes=False, strip_attributes=None): self.pubsub_matcher = PubSubMessageMatcher( 'mock_project', 'mock_sub_name', expected_msg, with_attributes=with_attributes, strip_attributes=strip_attributes) def init_counter_matcher(self, expected_msg_len=1): self.pubsub_matcher = PubSubMessageMatcher( 'mock_project', 'mock_sub_name', expected_msg_len=expected_msg_len) def test_message_matcher_success(self, mock_get_sub, unsued_mock): self.init_matcher(expected_msg=[b'a', b'b']) mock_sub = mock_get_sub.return_value mock_sub.pull.side_effect = [ create_pull_response([PullResponseMessage(b'a', {})]), create_pull_response([PullResponseMessage(b'b', {})]), ] hc_assert_that(self.mock_presult, self.pubsub_matcher) self.assertEqual(mock_sub.pull.call_count, 2) self.assertEqual(mock_sub.acknowledge.call_count, 2) def test_message_matcher_attributes_success(self, mock_get_sub, unsued_mock): self.init_matcher( expected_msg=[PubsubMessage(b'a', {'k': 'v'})], with_attributes=True) mock_sub = mock_get_sub.return_value mock_sub.pull.side_effect = [ create_pull_response([PullResponseMessage(b'a', {'k': 'v'})]) ] hc_assert_that(self.mock_presult, self.pubsub_matcher) self.assertEqual(mock_sub.pull.call_count, 1) self.assertEqual(mock_sub.acknowledge.call_count, 1) def test_message_matcher_attributes_fail(self, mock_get_sub, unsued_mock): self.init_matcher( expected_msg=[PubsubMessage(b'a', {})], with_attributes=True) mock_sub = mock_get_sub.return_value # Unexpected attribute 'k'. mock_sub.pull.side_effect = [ create_pull_response([PullResponseMessage(b'a', {'k': 'v'})]) ] with self.assertRaisesRegex(AssertionError, r'Unexpected'): hc_assert_that(self.mock_presult, self.pubsub_matcher) self.assertEqual(mock_sub.pull.call_count, 1) self.assertEqual(mock_sub.acknowledge.call_count, 1) def test_message_matcher_strip_success(self, mock_get_sub, unsued_mock): self.init_matcher( expected_msg=[PubsubMessage(b'a', {'k': 'v'})], with_attributes=True, strip_attributes=['id', 'timestamp']) mock_sub = mock_get_sub.return_value mock_sub.pull.side_effect = [ create_pull_response([ PullResponseMessage( b'a', { 'id': 'foo', 'timestamp': 'bar', 'k': 'v' }) ]) ] hc_assert_that(self.mock_presult, self.pubsub_matcher) self.assertEqual(mock_sub.pull.call_count, 1) self.assertEqual(mock_sub.acknowledge.call_count, 1) def test_message_matcher_strip_fail(self, mock_get_sub, unsued_mock): self.init_matcher( expected_msg=[PubsubMessage(b'a', {'k': 'v'})], with_attributes=True, strip_attributes=['id', 'timestamp']) mock_sub = mock_get_sub.return_value # Message is missing attribute 'timestamp'. mock_sub.pull.side_effect = [ create_pull_response( [PullResponseMessage(b'a', { 'id': 'foo', 'k': 'v' })]) ] with self.assertRaisesRegex(AssertionError, r'Stripped attributes'): hc_assert_that(self.mock_presult, self.pubsub_matcher) self.assertEqual(mock_sub.pull.call_count, 1) self.assertEqual(mock_sub.acknowledge.call_count, 1) def test_message_matcher_mismatch(self, mock_get_sub, unused_mock): self.init_matcher(expected_msg=[b'a']) mock_sub = mock_get_sub.return_value mock_sub.pull.side_effect = [ create_pull_response( [PullResponseMessage(b'c', {}), PullResponseMessage(b'd', {})]), ] with self.assertRaises(AssertionError) as error: hc_assert_that(self.mock_presult, self.pubsub_matcher) self.assertEqual(mock_sub.pull.call_count, 1) self.assertCountEqual([b'c', b'd'], self.pubsub_matcher.messages) self.assertIn( '\nExpected: Expected 1 messages.\n but: Got 2 messages.', str(error.exception.args[0])) self.assertEqual(mock_sub.pull.call_count, 1) self.assertEqual(mock_sub.acknowledge.call_count, 1) def test_message_matcher_timeout(self, mock_get_sub, unused_mock): self.init_matcher(expected_msg=[b'a']) mock_sub = mock_get_sub.return_value mock_sub.return_value.full_name.return_value = 'mock_sub' self.pubsub_matcher.timeout = 0.1 with self.assertRaisesRegex(AssertionError, r'Expected 1.*\n.*Got 0'): hc_assert_that(self.mock_presult, self.pubsub_matcher) self.assertTrue(mock_sub.pull.called) self.assertEqual(mock_sub.acknowledge.call_count, 0) def test_message_count_matcher_below_fail(self, mock_get_sub, unused_mock): self.init_counter_matcher(expected_msg_len=1) mock_sub = mock_get_sub.return_value mock_sub.pull.side_effect = [ create_pull_response( [PullResponseMessage(b'c', {}), PullResponseMessage(b'd', {})]), ] with self.assertRaises(AssertionError) as error: hc_assert_that(self.mock_presult, self.pubsub_matcher) self.assertEqual(mock_sub.pull.call_count, 1) self.assertIn( '\nExpected: Expected 1 messages.\n but: Got 2 messages.', str(error.exception.args[0])) def test_message_count_matcher_above_fail(self, mock_get_sub, unused_mock): self.init_counter_matcher(expected_msg_len=1) mock_sub = mock_get_sub.return_value self.pubsub_matcher.timeout = 0.1 with self.assertRaisesRegex(AssertionError, r'Expected 1.*\n.*Got 0'): hc_assert_that(self.mock_presult, self.pubsub_matcher) self.assertTrue(mock_sub.pull.called) self.assertEqual(mock_sub.acknowledge.call_count, 0) def test_message_count_matcher_success(self, mock_get_sub, unused_mock): self.init_counter_matcher(expected_msg_len=15) mock_sub = mock_get_sub.return_value mock_sub.pull.side_effect = [ create_pull_response( [PullResponseMessage(b'a', {'foo': 'bar'}) for _ in range(15)]) ] hc_assert_that(self.mock_presult, self.pubsub_matcher) self.assertEqual(mock_sub.pull.call_count, 1) self.assertEqual(mock_sub.acknowledge.call_count, 1) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/tests/utils.py0000644000175100017510000001230215206113032021707 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utility methods for testing on GCP.""" # pytype: skip-file import logging import secrets import time from apache_beam.io import filesystems from apache_beam.io.gcp.pubsub import PubsubMessage from apache_beam.utils import retry # Protect against environments where bigquery library is not available. try: from google.api_core import exceptions as gexc from google.cloud import bigquery except ImportError: gexc = None bigquery = None _LOGGER = logging.getLogger(__name__) class GcpTestIOError(retry.PermanentException): """Basic GCP IO error for testing. Function that raises this error should not be retried.""" pass @retry.with_exponential_backoff( num_retries=3, retry_filter=retry.retry_on_server_errors_filter) def create_bq_dataset(project, dataset_base_name): """Creates an empty BigQuery dataset. Args: project: Project to work in. dataset_base_name: Prefix for dataset id. Returns: A ``google.cloud.bigquery.dataset.DatasetReference`` object pointing to the new dataset. """ client = bigquery.Client(project=project) unique_dataset_name = '%s%d%s' % ( dataset_base_name, int(time.time()), secrets.token_hex(3)) dataset_ref = client.dataset(unique_dataset_name, project=project) dataset = bigquery.Dataset(dataset_ref) client.create_dataset(dataset) return dataset_ref @retry.with_exponential_backoff( num_retries=3, retry_filter=retry.retry_on_server_errors_filter) def delete_bq_dataset(project, dataset_ref): """Deletes a BigQuery dataset and its contents. Args: project: Project to work in. dataset_ref: A ``google.cloud.bigquery.dataset.DatasetReference`` object pointing to the dataset to delete. """ client = bigquery.Client(project=project) client.delete_dataset(dataset_ref, delete_contents=True) @retry.with_exponential_backoff( num_retries=3, retry_filter=retry.retry_on_server_errors_filter) def delete_bq_table(project, dataset_id, table_id): """Delete a BiqQuery table. Args: project: Name of the project. dataset_id: Name of the dataset where table is. table_id: Name of the table. """ _LOGGER.info( 'Clean up a BigQuery table with project: %s, dataset: %s, ' 'table: %s.', project, dataset_id, table_id) client = bigquery.Client(project=project) table_ref = client.dataset(dataset_id).table(table_id) try: client.delete_table(table_ref) except gexc.NotFound: raise GcpTestIOError('BigQuery table does not exist: %s' % table_ref) @retry.with_exponential_backoff( num_retries=3, retry_filter=retry.retry_on_server_errors_filter) def delete_directory(directory): """Delete a directory in a filesystem. Args: directory: Full path to a directory supported by Beam filesystems (e.g. "gs://mybucket/mydir/", "s3://...", ...) """ filesystems.FileSystems.delete([directory]) def write_to_pubsub( pub_client, topic_path, messages, with_attributes=False, chunk_size=100, delay_between_chunks=0.1): for start in range(0, len(messages), chunk_size): message_chunk = messages[start:start + chunk_size] if with_attributes: futures = [ pub_client.publish(topic_path, message.data, **message.attributes) for message in message_chunk ] else: futures = [ pub_client.publish(topic_path, message) for message in message_chunk ] for future in futures: future.result() time.sleep(delay_between_chunks) def read_from_pubsub( sub_client, subscription_path, with_attributes=False, number_of_elements=None, timeout=None): if number_of_elements is None and timeout is None: raise ValueError("Either number_of_elements or timeout must be specified.") messages = [] start_time = time.time() while ((number_of_elements is None or len(messages) < number_of_elements) and (timeout is None or (time.time() - start_time) < timeout)): try: response = sub_client.pull( subscription_path, max_messages=1000, retry=None, timeout=10) except (gexc.RetryError, gexc.DeadlineExceeded): continue ack_ids = [msg.ack_id for msg in response.received_messages] sub_client.acknowledge(subscription=subscription_path, ack_ids=ack_ids) for msg in response.received_messages: message = PubsubMessage._from_message(msg.message) if with_attributes: messages.append(message) else: messages.append(message.data) return messages ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/tests/utils_test.py0000644000175100017510000002316215206113032022754 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unittest for GCP testing utils.""" # pytype: skip-file import logging import unittest import mock from apache_beam.io.gcp.pubsub import PubsubMessage from apache_beam.io.gcp.tests import utils from apache_beam.testing import test_utils # Protect against environments where bigquery library is not available. try: from google.api_core import exceptions as gexc from google.cloud import bigquery from google.cloud import pubsub except ImportError: gexc = None bigquery = None pubsub = None @unittest.skipIf(bigquery is None, 'Bigquery dependencies are not installed.') @mock.patch.object(bigquery, 'Client') class UtilsTest(unittest.TestCase): def setUp(self): test_utils.patch_retry(self, utils) @mock.patch.object(bigquery, 'Dataset') def test_create_bq_dataset(self, mock_dataset, mock_client): mock_client.dataset.return_value = 'dataset_ref' mock_dataset.return_value = 'dataset_obj' utils.create_bq_dataset('project', 'dataset_base_name') mock_client.return_value.create_dataset.assert_called_with('dataset_obj') def test_delete_bq_dataset(self, mock_client): utils.delete_bq_dataset('project', 'dataset_ref') mock_client.return_value.delete_dataset.assert_called_with( 'dataset_ref', delete_contents=mock.ANY) def test_delete_table_succeeds(self, mock_client): mock_client.return_value.dataset.return_value.table.return_value = ( 'table_ref') utils.delete_bq_table('unused_project', 'unused_dataset', 'unused_table') mock_client.return_value.delete_table.assert_called_with('table_ref') def test_delete_table_fails_not_found(self, mock_client): mock_client.return_value.dataset.return_value.table.return_value = ( 'table_ref') mock_client.return_value.delete_table.side_effect = gexc.NotFound('test') with self.assertRaisesRegex(Exception, r'does not exist:.*table_ref'): utils.delete_bq_table('unused_project', 'unused_dataset', 'unused_table') @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed') class PubSubUtilTest(unittest.TestCase): def test_write_to_pubsub(self): mock_pubsub = mock.Mock() topic_path = "project/fakeproj/topics/faketopic" data = b'data' utils.write_to_pubsub(mock_pubsub, topic_path, [data]) mock_pubsub.publish.assert_has_calls( [mock.call(topic_path, data), mock.call().result()]) def test_write_to_pubsub_with_attributes(self): mock_pubsub = mock.Mock() topic_path = "project/fakeproj/topics/faketopic" data = b'data' attributes = {'key': 'value'} message = PubsubMessage(data, attributes) utils.write_to_pubsub( mock_pubsub, topic_path, [message], with_attributes=True) mock_pubsub.publish.assert_has_calls( [mock.call(topic_path, data, **attributes), mock.call().result()]) def test_write_to_pubsub_delay(self): number_of_elements = 2 chunk_size = 1 mock_pubsub = mock.Mock() topic_path = "project/fakeproj/topics/faketopic" data = b'data' with mock.patch('apache_beam.io.gcp.tests.utils.time') as mock_time: utils.write_to_pubsub( mock_pubsub, topic_path, [data] * number_of_elements, chunk_size=chunk_size, delay_between_chunks=123) mock_time.sleep.assert_called_with(123) mock_pubsub.publish.assert_has_calls( [mock.call(topic_path, data), mock.call().result()] * number_of_elements) def test_write_to_pubsub_many_chunks(self): number_of_elements = 83 chunk_size = 11 mock_pubsub = mock.Mock() topic_path = "project/fakeproj/topics/faketopic" data_list = [ 'data {}'.format(i).encode("utf-8") for i in range(number_of_elements) ] utils.write_to_pubsub( mock_pubsub, topic_path, data_list, chunk_size=chunk_size) call_list = [] for start in range(0, number_of_elements, chunk_size): # Publish a batch of messages call_list += [ mock.call(topic_path, data) for data in data_list[start:start + chunk_size] ] # Wait for those messages to be received call_list += [ mock.call().result() for _ in data_list[start:start + chunk_size] ] mock_pubsub.publish.assert_has_calls(call_list) def test_read_from_pubsub(self): mock_pubsub = mock.Mock() subscription_path = "project/fakeproj/subscriptions/fakesub" data = b'data' ack_id = 'ack_id' pull_response = test_utils.create_pull_response( [test_utils.PullResponseMessage(data, ack_id=ack_id)]) mock_pubsub.pull.return_value = pull_response output = utils.read_from_pubsub( mock_pubsub, subscription_path, number_of_elements=1) self.assertEqual([data], output) mock_pubsub.acknowledge.assert_called_once_with( subscription=subscription_path, ack_ids=[ack_id]) def test_read_from_pubsub_with_attributes(self): mock_pubsub = mock.Mock() subscription_path = "project/fakeproj/subscriptions/fakesub" data = b'data' ack_id = 'ack_id' attributes = {'key': 'value'} message = PubsubMessage(data, attributes) pull_response = test_utils.create_pull_response( [test_utils.PullResponseMessage(data, attributes, ack_id=ack_id)]) mock_pubsub.pull.return_value = pull_response output = utils.read_from_pubsub( mock_pubsub, subscription_path, with_attributes=True, number_of_elements=1) self.assertEqual([message], output) mock_pubsub.acknowledge.assert_called_once_with( subscription=subscription_path, ack_ids=[ack_id]) def test_read_from_pubsub_flaky(self): number_of_elements = 10 mock_pubsub = mock.Mock() subscription_path = "project/fakeproj/subscriptions/fakesub" data = b'data' ack_id = 'ack_id' pull_response = test_utils.create_pull_response( [test_utils.PullResponseMessage(data, ack_id=ack_id)]) class FlakyPullResponse(object): def __init__(self, pull_response): self.pull_response = pull_response self._state = -1 def __call__(self, *args, **kwargs): self._state += 1 if self._state % 3 == 0: raise gexc.RetryError("", "") if self._state % 3 == 1: raise gexc.DeadlineExceeded("") if self._state % 3 == 2: return self.pull_response mock_pubsub.pull.side_effect = FlakyPullResponse(pull_response) output = utils.read_from_pubsub( mock_pubsub, subscription_path, number_of_elements=number_of_elements) self.assertEqual([data] * number_of_elements, output) self._assert_ack_ids_equal(mock_pubsub, [ack_id] * number_of_elements) def test_read_from_pubsub_many(self): response_size = 33 number_of_elements = 100 mock_pubsub = mock.Mock() subscription_path = "project/fakeproj/subscriptions/fakesub" data_list = [ 'data {}'.format(i).encode("utf-8") for i in range(number_of_elements) ] attributes_list = [{ 'key': 'value {}'.format(i) } for i in range(number_of_elements)] ack_ids = ['ack_id_{}'.format(i) for i in range(number_of_elements)] messages = [ PubsubMessage(data, attributes) for data, attributes in zip(data_list, attributes_list) ] response_messages = [ test_utils.PullResponseMessage(data, attributes, ack_id=ack_id) for data, attributes, ack_id in zip( data_list, attributes_list, ack_ids) ] class SequentialPullResponse(object): def __init__(self, response_messages, response_size): self.response_messages = response_messages self.response_size = response_size self._index = 0 def __call__(self, *args, **kwargs): start = self._index self._index += self.response_size response = test_utils.create_pull_response( self.response_messages[start:start + self.response_size]) return response mock_pubsub.pull.side_effect = SequentialPullResponse( response_messages, response_size) output = utils.read_from_pubsub( mock_pubsub, subscription_path, with_attributes=True, number_of_elements=number_of_elements) self.assertEqual(messages, output) self._assert_ack_ids_equal(mock_pubsub, ack_ids) def test_read_from_pubsub_invalid_arg(self): sub_client = mock.Mock() subscription_path = "project/fakeproj/subscriptions/fakesub" with self.assertRaisesRegex(ValueError, "number_of_elements"): utils.read_from_pubsub(sub_client, subscription_path) with self.assertRaisesRegex(ValueError, "number_of_elements"): utils.read_from_pubsub( sub_client, subscription_path, with_attributes=True) def _assert_ack_ids_equal(self, mock_pubsub, ack_ids): actual_ack_ids = [ ack_id for args_list in mock_pubsub.acknowledge.call_args_list for ack_id in args_list[1]["ack_ids"] ] self.assertEqual(actual_ack_ids, ack_ids) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/gcp/tests/xlang_spannerio_it_test.py0000644000175100017510000003235315206113032025501 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import argparse import logging import os import time import unittest import uuid from typing import NamedTuple from typing import Optional import pytest import apache_beam as beam from apache_beam import coders from apache_beam.io.gcp.spanner import ReadFromSpanner from apache_beam.io.gcp.spanner import SpannerDelete from apache_beam.io.gcp.spanner import SpannerInsert from apache_beam.io.gcp.spanner import SpannerInsertOrUpdate from apache_beam.io.gcp.spanner import SpannerReplace from apache_beam.io.gcp.spanner import SpannerUpdate from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.utils.timestamp import Timestamp # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from google.cloud import spanner except ImportError: spanner = None try: from testcontainers.core.container import DockerContainer except ImportError: DockerContainer = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports TIMESTAMPS = [Timestamp.of(1234567890.0 + i) for i in range(1000)] class SpannerTestKey(NamedTuple): f_string: str class SpannerTestRow(NamedTuple): f_string: str f_int64: Optional[int] f_boolean: Optional[bool] f_timestamp: Optional[Timestamp] class SpannerPartTestRow(NamedTuple): f_string: str f_int64: Optional[int] f_timestamp: Optional[Timestamp] @pytest.mark.uses_gcp_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_JARS'), "EXPANSION_JARS environment var is not provided, " "indicating that jars have not been built") @unittest.skipIf(spanner is None, 'GCP dependencies are not installed.') @unittest.skipIf( DockerContainer is None, 'testcontainers package is not installed.') class CrossLanguageSpannerIOTest(unittest.TestCase): @classmethod def setUpClass(cls): parser = argparse.ArgumentParser() parser.add_argument( '--spanner_instance_id', default='beam-test', help='Spanner instance id', ) parser.add_argument( '--spanner_project_id', default='beam-testing', help='GCP project with spanner instance', ) parser.add_argument( '--use_real_spanner', action='store_true', default=False, help='Whether to use emulator or real spanner instance', ) pipeline = TestPipeline(is_integration_test=True) runner_name = type(pipeline.runner).__name__ if 'DataflowRunner' in runner_name: pytest.skip("Spanner emulator not compatible with dataflow runner.") argv = pipeline.get_full_options_as_args() known_args, _ = parser.parse_known_args(argv) cls.project_id = known_args.spanner_project_id cls.instance_id = known_args.spanner_instance_id use_spanner_emulator = not known_args.use_real_spanner cls.table = 'xlang_beam_spanner' cls.spanner_helper = SpannerHelper( cls.project_id, cls.instance_id, cls.table, use_spanner_emulator) coders.registry.register_coder(SpannerTestRow, coders.RowCoder) coders.registry.register_coder(SpannerPartTestRow, coders.RowCoder) coders.registry.register_coder(SpannerTestKey, coders.RowCoder) @classmethod def tearDownClass(cls): cls.spanner_helper.shutdown() def setUp(self): self.database_id = f'xlang_beam{uuid.uuid4()}'.replace('-', '')[:30] self.spanner_helper.create_database(self.database_id) def tearDown(self): self.spanner_helper.drop_database(self.database_id) def test_spanner_insert_or_update(self): self.spanner_helper.insert_values( self.database_id, [('or_update0', 5, False, TIMESTAMPS[1].to_rfc3339()), ('or_update1', 9, False, TIMESTAMPS[0].to_rfc3339())]) def to_row_fn(i): return SpannerTestRow( f_int64=i, f_string=f'or_update{i}', f_boolean=i % 2 == 0, f_timestamp=TIMESTAMPS[i]) self.run_write_pipeline(3, to_row_fn, SpannerTestRow, SpannerInsertOrUpdate) results = self.spanner_helper.read_data( self.database_id, prefix='or_update') self.assertEqual(len(results), 3) for i, row in enumerate(results): self.assertEqual(row[0], f'or_update{i}') self.assertEqual(row[1], i) self.assertEqual(row[2], i % 2 == 0) self.assertEqual(row[3].timestamp_pb(), TIMESTAMPS[i].to_proto()) def test_spanner_insert(self): def to_row_fn(num): return SpannerTestRow( f_string=f'insert{num}', f_int64=num, f_boolean=None, f_timestamp=TIMESTAMPS[num]) self.run_write_pipeline(1000, to_row_fn, SpannerTestRow, SpannerInsert) def compare_row(row): return row[1] results = sorted( self.spanner_helper.read_data(self.database_id, 'insert'), key=compare_row) self.assertEqual(len(results), 1000) for i, row in enumerate(results): self.assertEqual(row[0], f'insert{i}') self.assertEqual(row[1], i) self.assertIsNone(row[2]) self.assertEqual(row[3].timestamp_pb(), TIMESTAMPS[i].to_proto()) def test_spanner_replace(self): self.spanner_helper.insert_values( self.database_id, [('replace0', 0, True, TIMESTAMPS[10].to_rfc3339()), ('replace1', 1, False, TIMESTAMPS[11].to_rfc3339())]) def to_row_fn(num): return SpannerPartTestRow( f_string=f'replace{num}', f_int64=num + 10, f_timestamp=TIMESTAMPS[num]) self.run_write_pipeline(2, to_row_fn, SpannerPartTestRow, SpannerReplace) results = self.spanner_helper.read_data(self.database_id, prefix='replace') for i in range(len(results)): results[i][3] = results[i][3].timestamp_pb() self.assertEqual( results, [['replace0', 10, None, TIMESTAMPS[0].to_proto()], ['replace1', 11, None, TIMESTAMPS[1].to_proto()]]) def test_spanner_update(self): self.spanner_helper.insert_values( self.database_id, [('update0', 5, False, TIMESTAMPS[10].to_rfc3339()), ('update1', 9, False, TIMESTAMPS[100].to_rfc3339())]) def to_row_fn(num): return SpannerPartTestRow( f_string=f'update{num}', f_int64=num + 10, f_timestamp=TIMESTAMPS[num]) self.run_write_pipeline(2, to_row_fn, SpannerPartTestRow, SpannerUpdate) results = self.spanner_helper.read_data(self.database_id, 'update') for i in range(len(results)): results[i][3] = results[i][3].timestamp_pb() self.assertEqual( results, [['update0', 10, False, TIMESTAMPS[0].to_proto()], ['update1', 11, False, TIMESTAMPS[1].to_proto()]]) def test_spanner_delete(self): self.spanner_helper.insert_values( self.database_id, values=[ ('delete0', 0, None, TIMESTAMPS[0].to_rfc3339()), ('delete6', 6, False, TIMESTAMPS[0].to_rfc3339()), ('delete20', 20, True, TIMESTAMPS[0].to_rfc3339()), ]) def to_row_fn(num): return SpannerTestKey(f_string=f'delete{num}') self.run_write_pipeline(10, to_row_fn, SpannerTestKey, SpannerDelete) results = self.spanner_helper.read_data(self.database_id, prefix='delete') for i in range(len(results)): results[i][3] = results[i][3].timestamp_pb() self.assertEqual( results, [['delete20', 20, True, TIMESTAMPS[0].to_proto()]]) def test_spanner_read_query(self): self.insert_read_values('query_read') self.run_read_pipeline('query_read', query=f'SELECT * FROM {self.table}') def test_spanner_read_table(self): self.insert_read_values('table_read') self.run_read_pipeline('table_read', table=self.table) def run_read_pipeline(self, prefix, table=None, query=None): with TestPipeline(is_integration_test=True) as p: p.not_use_test_runner_api = True result = ( p | ReadFromSpanner( instance_id=self.instance_id, database_id=self.database_id, project_id=self.project_id, row_type=SpannerTestRow, sql=query, table=table, emulator_host=self.spanner_helper.get_emulator_host(), )) assert_that( result, equal_to([ SpannerTestRow( f_int64=0, f_string=f'{prefix}0', f_boolean=None, f_timestamp=TIMESTAMPS[0]), SpannerTestRow( f_int64=1, f_string=f'{prefix}1', f_boolean=True, f_timestamp=TIMESTAMPS[1]), SpannerTestRow( f_int64=2, f_string=f'{prefix}2', f_boolean=False, f_timestamp=TIMESTAMPS[2]), ])) def run_write_pipeline( self, num_rows, to_row_fn, row_type, spanner_transform=None): with TestPipeline(is_integration_test=True) as p: p.not_use_test_runner_api = True _ = ( p | 'Impulse' >> beam.Impulse() | 'Generate' >> beam.FlatMap(lambda x: range(num_rows)) # pylint: disable=bad-option-value | 'Map to row' >> beam.Map(to_row_fn).with_output_types(row_type) | 'Write to Spanner' >> spanner_transform( instance_id=self.instance_id, database_id=self.database_id, project_id=self.project_id, table=self.table, failure_mode=beam.io.gcp.spanner.FailureMode.REPORT_FAILURES, emulator_host=self.spanner_helper.get_emulator_host(), )) def insert_read_values(self, prefix): self.spanner_helper.insert_values( self.database_id, values=[ (f'{prefix}0', 0, None, TIMESTAMPS[0].to_rfc3339()), (f'{prefix}1', 1, True, TIMESTAMPS[1].to_rfc3339()), (f'{prefix}2', 2, False, TIMESTAMPS[2].to_rfc3339()), ]) def retry(fn, retries, err_msg, *args, **kwargs): for _ in range(retries): try: return fn(*args, **kwargs) except: # pylint: disable=bare-except pass logging.error(err_msg) raise RuntimeError(err_msg) class SpannerHelper(object): def __init__(self, project_id, instance_id, table, use_emulator): self.use_emulator = use_emulator self.table = table self.host = None if use_emulator: self.emulator = DockerContainer( 'gcr.io/cloud-spanner-emulator/emulator:latest').with_exposed_ports( 9010, 9020) retry(self.emulator.start, 3, 'Could not start spanner emulator.') time.sleep(3) self.host = f'{self.emulator.get_container_host_ip()}:' \ f'{self.emulator.get_exposed_port(9010)}' os.environ['SPANNER_EMULATOR_HOST'] = self.host self.client = spanner.Client(project_id) self.instance = self.client.instance(instance_id) if use_emulator: self.create_instance() def create_instance(self): self.instance.create().result(120) def create_database(self, database_id): database = self.instance.database( database_id, ddl_statements=[ f''' CREATE TABLE {self.table} ( f_string STRING(1024) NOT NULL, f_int64 INT64, f_boolean BOOL, f_timestamp TIMESTAMP ) PRIMARY KEY (f_string)''' ]) database.create().result(120) def insert_values(self, database_id, values, columns=None): values = values or [] columns = columns or ('f_string', 'f_int64', 'f_boolean', 'f_timestamp') with self.instance.database(database_id).batch() as batch: batch.insert( table=self.table, columns=columns, values=values, ) def get_emulator_host(self): return f'http://{self.host}' def read_data(self, database_id, prefix): database = self.instance.database(database_id) with database.snapshot() as snapshot: results = snapshot.execute_sql( f'''SELECT * FROM {self.table} WHERE f_string LIKE "{prefix}%" ORDER BY f_int64''') try: rows = list(results) if results else None except IndexError: raise ValueError(f"Spanner results not found for {prefix}.") return rows def drop_database(self, database_id): database = self.instance.database(database_id) database.drop() def shutdown(self): if self.use_emulator: try: self.emulator.stop() except: # pylint: disable=bare-except logging.error('Could not stop Spanner Cloud emulator.') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/hadoopfilesystem.py0000644000175100017510000003420015206113032022214 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """:class:`~apache_beam.io.filesystem.FileSystem` implementation for accessing Hadoop Distributed File System files.""" # pytype: skip-file import io import logging import posixpath import re from typing import BinaryIO # pylint: disable=unused-import from apache_beam.io import filesystemio from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystem import CompressedFile from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystem import FileMetadata from apache_beam.io.filesystem import FileSystem from apache_beam.options.pipeline_options import HadoopFileSystemOptions from apache_beam.options.pipeline_options import PipelineOptions try: import hdfs except ImportError: hdfs = None __all__ = ['HadoopFileSystem'] _HDFS_PREFIX = 'hdfs:/' _URL_RE = re.compile(r'^' + _HDFS_PREFIX + r'(/.*)') _FULL_URL_RE = re.compile(r'^' + _HDFS_PREFIX + r'/([^/]+)(/.*)*') _COPY_BUFFER_SIZE = 2**16 _DEFAULT_BUFFER_SIZE = 20 * 1024 * 1024 # WebHDFS FileChecksum property constants. _FILE_CHECKSUM_ALGORITHM = 'algorithm' _FILE_CHECKSUM_BYTES = 'bytes' _FILE_CHECKSUM_LENGTH = 'length' # WebHDFS FileStatus property constants. _FILE_STATUS_LENGTH = 'length' _FILE_STATUS_UPDATED = 'modificationTime' _FILE_STATUS_PATH_SUFFIX = 'pathSuffix' _FILE_STATUS_TYPE = 'type' _FILE_STATUS_TYPE_DIRECTORY = 'DIRECTORY' _FILE_STATUS_TYPE_FILE = 'FILE' _LOGGER = logging.getLogger(__name__) class HdfsDownloader(filesystemio.Downloader): def __init__(self, hdfs_client, path): self._hdfs_client = hdfs_client self._path = path self._size = self._hdfs_client.status(path)[_FILE_STATUS_LENGTH] @property def size(self): return self._size def get_range(self, start, end): with self._hdfs_client.read(self._path, offset=start, length=end - start) as reader: return reader.read() class HdfsUploader(filesystemio.Uploader): def __init__(self, hdfs_client, path): self._hdfs_client = hdfs_client if self._hdfs_client.status(path, strict=False) is not None: raise BeamIOError('Path already exists: %s' % path) self._handle_context = self._hdfs_client.write(path) self._handle = self._handle_context.__enter__() def put(self, data): # hdfs uses an async writer which first add data to a queue. To avoid buffer # gets reused upstream a deepcopy is required here. self._handle.write(bytes(data)) def finish(self): self._handle.__exit__(None, None, None) self._handle = None self._handle_context = None class HadoopFileSystem(FileSystem): """``FileSystem`` implementation that supports HDFS. URL arguments to methods expect strings starting with ``hdfs://``. """ def __init__(self, pipeline_options): """Initializes a connection to HDFS. Connection configuration is done by passing pipeline options. See :class:`~apache_beam.options.pipeline_options.HadoopFileSystemOptions`. """ super().__init__(pipeline_options) if hdfs is None: raise ImportError( 'Failed to import hdfs. You can ensure it is ' 'installed by installing the hadoop beam extra') logging.getLogger('hdfs.client').setLevel(logging.WARN) if pipeline_options is None: raise ValueError('pipeline_options is not set') if isinstance(pipeline_options, PipelineOptions): hdfs_options = pipeline_options.view_as(HadoopFileSystemOptions) hdfs_host = hdfs_options.hdfs_host hdfs_port = hdfs_options.hdfs_port hdfs_user = hdfs_options.hdfs_user self._full_urls = hdfs_options.hdfs_full_urls else: hdfs_host = pipeline_options.get('hdfs_host') hdfs_port = pipeline_options.get('hdfs_port') hdfs_user = pipeline_options.get('hdfs_user') self._full_urls = pipeline_options.get('hdfs_full_urls', False) if hdfs_host is None: raise ValueError('hdfs_host is not set') if hdfs_port is None: raise ValueError('hdfs_port is not set') if hdfs_user is None: raise ValueError('hdfs_user is not set') if not isinstance(self._full_urls, bool): raise ValueError( 'hdfs_full_urls should be bool, got: %s', self._full_urls) self._hdfs_client = hdfs.InsecureClient( 'http://%s:%s' % (hdfs_host, str(hdfs_port)), user=hdfs_user) @classmethod def scheme(cls): return 'hdfs' def _parse_url(self, url): """Verifies that url begins with hdfs:// prefix, strips it and adds a leading /. Parsing behavior is determined by HadoopFileSystemOptions.hdfs_full_urls. Args: url: (str) A URL in the form hdfs://path/... or in the form hdfs://server/path/... Raises: ValueError if the URL doesn't match the expect format. Returns: (str, str) If using hdfs_full_urls, for an input of 'hdfs://server/path/...' will return (server, '/path/...'). Otherwise, for an input of 'hdfs://path/...', will return ('', '/path/...'). """ if not self._full_urls: m = _URL_RE.match(url) if m is None: raise ValueError('Could not parse url: %s' % url) return '', m.group(1) else: m = _FULL_URL_RE.match(url) if m is None: raise ValueError('Could not parse url: %s' % url) return m.group(1), m.group(2) or '/' def join(self, base_url, *paths): """Join two or more pathname components. Args: base_url: string path of the first component of the path. Must start with hdfs://. paths: path components to be added Returns: Full url after combining all the passed components. """ server, basepath = self._parse_url(base_url) return _HDFS_PREFIX + self._join(server, basepath, *paths) def _join(self, server, basepath, *paths): res = posixpath.join(basepath, *paths) if server: server = '/' + server return server + res def split(self, url): server, rel_path = self._parse_url(url) if server: server = '/' + server head, tail = posixpath.split(rel_path) return _HDFS_PREFIX + server + head, tail def mkdirs(self, url): _, path = self._parse_url(url) if self._exists(path): raise BeamIOError('Path already exists: %s' % path) return self._mkdirs(path) def _mkdirs(self, path): self._hdfs_client.makedirs(path) def has_dirs(self): return True def _list(self, url): try: server, path = self._parse_url(url) for res in self._hdfs_client.list(path, status=True): yield FileMetadata( _HDFS_PREFIX + self._join(server, path, res[0]), res[1][_FILE_STATUS_LENGTH], res[1][_FILE_STATUS_UPDATED] / 1000.0) except Exception as e: # pylint: disable=broad-except raise BeamIOError('List operation failed', {url: e}) @staticmethod def _add_compression(stream, path, mime_type, compression_type): if mime_type != 'application/octet-stream': _LOGGER.warning( 'Mime types are not supported. Got non-default mime_type:' ' %s', mime_type) if compression_type == CompressionTypes.AUTO: compression_type = CompressionTypes.detect_compression_type(path) if compression_type != CompressionTypes.UNCOMPRESSED: return CompressedFile(stream) return stream def create( self, url, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO) -> BinaryIO: """ Returns: A Python File-like object. """ _, path = self._parse_url(url) return self._create(path, mime_type, compression_type) def _create( self, path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO): stream = io.BufferedWriter( filesystemio.UploaderStream(HdfsUploader(self._hdfs_client, path)), buffer_size=_DEFAULT_BUFFER_SIZE) return self._add_compression(stream, path, mime_type, compression_type) def open( self, url, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO) -> BinaryIO: """ Returns: A Python File-like object. """ _, path = self._parse_url(url) return self._open(path, mime_type, compression_type) def _open( self, path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO): stream = io.BufferedReader( filesystemio.DownloaderStream(HdfsDownloader(self._hdfs_client, path)), buffer_size=_DEFAULT_BUFFER_SIZE) return self._add_compression(stream, path, mime_type, compression_type) def copy(self, source_file_names, destination_file_names): """ It is an error if any file to copy already exists at the destination. Raises ``BeamIOError`` if any error occurred. Args: source_file_names: iterable of URLs. destination_file_names: iterable of URLs. """ if len(source_file_names) != len(destination_file_names): raise BeamIOError( 'source_file_names and destination_file_names should ' 'be equal in length: %d != %d' % (len(source_file_names), len(destination_file_names))) def _copy_file(source, destination): with self._open(source) as f1: with self._create(destination) as f2: while True: buf = f1.read(_COPY_BUFFER_SIZE) if not buf: break f2.write(buf) def _copy_path(source, destination): """Recursively copy the file tree from the source to the destination.""" if self._hdfs_client.status( source)[_FILE_STATUS_TYPE] != _FILE_STATUS_TYPE_DIRECTORY: _copy_file(source, destination) return for path, dirs, files in self._hdfs_client.walk(source): for dir in dirs: new_dir = self._join('', destination, dir) if not self._exists(new_dir): self._mkdirs(new_dir) rel_path = posixpath.relpath(path, source) if rel_path == '.': rel_path = '' for file in files: _copy_file( self._join('', path, file), self._join('', destination, rel_path, file)) exceptions = {} for source, destination in zip(source_file_names, destination_file_names): try: _, rel_source = self._parse_url(source) _, rel_destination = self._parse_url(destination) _copy_path(rel_source, rel_destination) except Exception as e: # pylint: disable=broad-except exceptions[(source, destination)] = e if exceptions: raise BeamIOError('Copy operation failed', exceptions) def rename(self, source_file_names, destination_file_names): exceptions = {} for source, destination in zip(source_file_names, destination_file_names): try: _, rel_source = self._parse_url(source) _, rel_destination = self._parse_url(destination) try: self._hdfs_client.rename(rel_source, rel_destination) except hdfs.HdfsError as e: raise BeamIOError( 'libhdfs error in renaming %s to %s' % (source, destination), e) except Exception as e: # pylint: disable=broad-except exceptions[(source, destination)] = e if exceptions: raise BeamIOError('Rename operation failed', exceptions) def exists(self, url: str) -> bool: """Checks existence of url in HDFS. Args: url: String in the form hdfs://... Returns: True if url exists as a file or directory in HDFS. """ _, path = self._parse_url(url) return self._exists(path) def _exists(self, path): """Returns True if path exists as a file or directory in HDFS. Args: path: String in the form /... """ return self._hdfs_client.status(path, strict=False) is not None def size(self, url): """Fetches file size for a URL. Returns: int size of path according to the FileSystem. Raises: ``BeamIOError``: if url doesn't exist. """ return self.metadata(url).size_in_bytes def last_updated(self, url): """Fetches last updated time for a URL. Args: url: string url of file. Returns: float UNIX Epoch time Raises: ``BeamIOError``: if path doesn't exist. """ return self.metadata(url).last_updated_in_seconds def checksum(self, url): """Fetches a checksum description for a URL. Returns: String describing the checksum. Raises: ``BeamIOError``: if url doesn't exist. """ _, path = self._parse_url(url) file_checksum = self._hdfs_client.checksum(path) return '%s-%d-%s' % ( file_checksum[_FILE_CHECKSUM_ALGORITHM], file_checksum[_FILE_CHECKSUM_LENGTH], file_checksum[_FILE_CHECKSUM_BYTES], ) def metadata(self, url): """Fetch metadata fields of a file on the FileSystem. Args: url: string url of a file. Returns: :class:`~apache_beam.io.filesystem.FileMetadata`. Raises: ``BeamIOError``: if url doesn't exist. """ _, path = self._parse_url(url) status = self._hdfs_client.status(path, strict=False) if status is None: raise BeamIOError('File not found: %s' % url) return FileMetadata( url, status[_FILE_STATUS_LENGTH], status[_FILE_STATUS_UPDATED] / 1000.0) def delete(self, urls): exceptions = {} for url in urls: try: _, path = self._parse_url(url) self._hdfs_client.delete(path, recursive=True) except Exception as e: # pylint: disable=broad-except exceptions[url] = e if exceptions: raise BeamIOError("Delete operation failed", exceptions) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/hadoopfilesystem_test.py0000644000175100017510000005421415206113032023262 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for :class:`HadoopFileSystem`.""" # pytype: skip-file import io import logging import posixpath import time import unittest from parameterized import parameterized_class from apache_beam.io import hadoopfilesystem as hdfs from apache_beam.io.filesystem import BeamIOError from apache_beam.options.pipeline_options import HadoopFileSystemOptions from apache_beam.options.pipeline_options import PipelineOptions try: import hdfs as actual_hdfs except ImportError: actual_hdfs = None class FakeFile(io.BytesIO): """File object for FakeHdfs""" __hash__ = None # type: ignore[assignment] def __init__(self, path, mode='', type='FILE', time_ms=None): io.BytesIO.__init__(self) if time_ms is None: time_ms = int(time.time() * 1000) self.time_ms = time_ms self.stat = {'path': path, 'mode': mode, 'type': type} self.saved_data = None def __eq__(self, other): """Equality of two files. Timestamp not included in comparison""" return self.stat == other.stat and self.getvalue() == self.getvalue() def close(self): self.saved_data = self.getvalue() io.BytesIO.close(self) def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): self.close() @property def size(self): if self.closed: # pylint: disable=using-constant-test if self.saved_data is None: return 0 return len(self.saved_data) return len(self.getvalue()) def get_file_status(self): """Returns a partial WebHDFS FileStatus object.""" return { hdfs._FILE_STATUS_PATH_SUFFIX: posixpath.basename(self.stat['path']), hdfs._FILE_STATUS_LENGTH: self.size, hdfs._FILE_STATUS_TYPE: self.stat['type'], hdfs._FILE_STATUS_UPDATED: self.time_ms } def get_file_checksum(self): """Returns a WebHDFS FileChecksum object.""" return { hdfs._FILE_CHECKSUM_ALGORITHM: 'fake_algo', hdfs._FILE_CHECKSUM_BYTES: 'checksum_byte_sequence', hdfs._FILE_CHECKSUM_LENGTH: 5, } class FakeHdfsError(Exception): """Generic error for FakeHdfs methods.""" class FakeHdfs(object): """Fake implementation of ``hdfs.Client``.""" def __init__(self): self.files = {} def write(self, path): if self.status(path, strict=False) is not None: raise FakeHdfsError('Path already exists: %s' % path) new_file = FakeFile(path, 'wb') self.files[path] = new_file return new_file def read(self, path, offset=0, length=None): old_file = self.files.get(path, None) if old_file is None: raise FakeHdfsError('Path not found: %s' % path) if old_file.stat['type'] == 'DIRECTORY': raise FakeHdfsError('Cannot open a directory: %s' % path) if not old_file.closed: raise FakeHdfsError('File already opened: %s' % path) # old_file is closed and can't be operated upon. Return a copy instead. new_file = FakeFile(path, 'rb') if old_file.saved_data: if length is None: new_file.write(old_file.saved_data) else: new_file.write(old_file.saved_data[:offset + length]) new_file.seek(offset) return new_file def list(self, path, status=False): if not status: raise ValueError('status must be True') fs = self.status(path, strict=False) if (fs is not None and fs[hdfs._FILE_STATUS_TYPE] == hdfs._FILE_STATUS_TYPE_FILE): raise ValueError( 'list must be called on a directory, got file: %s' % path) result = [] for file in self.files.values(): if file.stat['path'].startswith(path): fs = file.get_file_status() result.append((fs[hdfs._FILE_STATUS_PATH_SUFFIX], fs)) return result def makedirs(self, path): self.files[path] = FakeFile(path, type='DIRECTORY') def status(self, path, strict=True): f = self.files.get(path) if f is None: if strict: raise FakeHdfsError('Path not found: %s' % path) else: return f return f.get_file_status() def delete(self, path, recursive=True): if not recursive: raise FakeHdfsError('Non-recursive mode not implemented') _ = self.status(path) for filepath in list(self.files): if filepath.startswith(path): del self.files[filepath] def walk(self, path): paths = [path] while paths: path = paths.pop() files = [] dirs = [] for full_path in self.files: if not full_path.startswith(path): continue short_path = posixpath.relpath(full_path, path) if '/' not in short_path: if self.status(full_path)[hdfs._FILE_STATUS_TYPE] == 'DIRECTORY': if short_path != '.': dirs.append(short_path) else: files.append(short_path) yield path, dirs, files paths = [posixpath.join(path, dir) for dir in dirs] def rename(self, path1, path2): if self.status(path1, strict=False) is None: raise FakeHdfsError('Path1 not found: %s' % path1) files_to_rename = [ path for path in self.files if path == path1 or path.startswith(path1 + '/') ] for fullpath in files_to_rename: f = self.files.pop(fullpath) newpath = path2 + fullpath[len(path1):] f.stat['path'] = newpath self.files[newpath] = f def checksum(self, path): f = self.files.get(path, None) if f is None: raise FakeHdfsError('Path not found: %s' % path) return f.get_file_checksum() @parameterized_class(('full_urls', ), [(False, ), (True, )]) @unittest.skipIf(actual_hdfs is None, "hdfs extra not installed") class HadoopFileSystemTest(unittest.TestCase): def setUp(self): self._fake_hdfs = FakeHdfs() hdfs.hdfs.InsecureClient = (lambda *args, **kwargs: self._fake_hdfs) pipeline_options = PipelineOptions() hdfs_options = pipeline_options.view_as(HadoopFileSystemOptions) hdfs_options.hdfs_host = '' hdfs_options.hdfs_port = 0 hdfs_options.hdfs_user = '' self.fs = hdfs.HadoopFileSystem(pipeline_options) self.fs._full_urls = self.full_urls if self.full_urls: self.tmpdir = 'hdfs://test_dir' else: self.tmpdir = 'hdfs://server/test_dir' for filename in ['old_file1', 'old_file2']: url = self.fs.join(self.tmpdir, filename) self.fs.create(url).close() def test_scheme(self): self.assertEqual(self.fs.scheme(), 'hdfs') self.assertEqual(hdfs.HadoopFileSystem.scheme(), 'hdfs') def test_parse_url(self): cases = [ ('hdfs://', ('', '/'), False), ('hdfs://', None, True), ('hdfs://a', ('', '/a'), False), ('hdfs://a', ('a', '/'), True), ('hdfs://a/', ('', '/a/'), False), ('hdfs://a/', ('a', '/'), True), ('hdfs://a/b', ('', '/a/b'), False), ('hdfs://a/b', ('a', '/b'), True), ('hdfs://a/b/', ('', '/a/b/'), False), ('hdfs://a/b/', ('a', '/b/'), True), ('hdfs:/a/b', None, False), ('hdfs:/a/b', None, True), ('invalid', None, False), ('invalid', None, True), ] for url, expected, full_urls in cases: if self.full_urls != full_urls: continue try: result = self.fs._parse_url(url) except ValueError: self.assertIsNone(expected, msg=(url, expected, full_urls)) continue self.assertEqual(expected, result, msg=(url, expected, full_urls)) def test_url_join(self): self.assertEqual( 'hdfs://tmp/path/to/file', self.fs.join('hdfs://tmp/path', 'to', 'file')) self.assertEqual( 'hdfs://tmp/path/to/file', self.fs.join('hdfs://tmp/path', 'to/file')) self.assertEqual('hdfs://tmp/path/', self.fs.join('hdfs://tmp/path/', '')) if not self.full_urls: self.assertEqual('hdfs://bar', self.fs.join('hdfs://foo', '/bar')) self.assertEqual('hdfs://bar', self.fs.join('hdfs://foo/', '/bar')) with self.assertRaises(ValueError): self.fs.join('/no/scheme', 'file') else: self.assertEqual('hdfs://foo/bar', self.fs.join('hdfs://foo', '/bar')) self.assertEqual('hdfs://foo/bar', self.fs.join('hdfs://foo/', '/bar')) def test_url_split(self): self.assertEqual(('hdfs://tmp/path/to', 'file'), self.fs.split('hdfs://tmp/path/to/file')) if not self.full_urls: self.assertEqual(('hdfs://', 'tmp'), self.fs.split('hdfs://tmp')) self.assertEqual(('hdfs://tmp', ''), self.fs.split('hdfs://tmp/')) self.assertEqual(('hdfs://tmp', 'a'), self.fs.split('hdfs://tmp/a')) else: self.assertEqual(('hdfs://tmp/', ''), self.fs.split('hdfs://tmp')) self.assertEqual(('hdfs://tmp/', ''), self.fs.split('hdfs://tmp/')) self.assertEqual(('hdfs://tmp/', 'a'), self.fs.split('hdfs://tmp/a')) self.assertEqual(('hdfs://tmp/a', ''), self.fs.split('hdfs://tmp/a/')) with self.assertRaisesRegex(ValueError, r'parse'): self.fs.split('tmp') def test_mkdirs(self): url = self.fs.join(self.tmpdir, 't1/t2') self.fs.mkdirs(url) self.assertTrue(self.fs.exists(url)) def test_mkdirs_failed(self): url = self.fs.join(self.tmpdir, 't1/t2') self.fs.mkdirs(url) with self.assertRaises(IOError): self.fs.mkdirs(url) def test_match_file(self): expected_files = [ self.fs.join(self.tmpdir, filename) for filename in ['old_file1', 'old_file2'] ] match_patterns = expected_files result = self.fs.match(match_patterns) returned_files = [ f.path for match_result in result for f in match_result.metadata_list ] self.assertCountEqual(expected_files, returned_files) def test_match_file_with_limits(self): expected_files = [ self.fs.join(self.tmpdir, filename) for filename in ['old_file1', 'old_file2'] ] result = self.fs.match([self.tmpdir + '/'], [1])[0] files = [f.path for f in result.metadata_list] self.assertEqual(len(files), 1) self.assertIn(files[0], expected_files) def test_match_file_with_zero_limit(self): result = self.fs.match([self.tmpdir + '/'], [0])[0] self.assertEqual(len(result.metadata_list), 0) def test_match_file_empty(self): url = self.fs.join(self.tmpdir, 'nonexistent_file') result = self.fs.match([url])[0] files = [f.path for f in result.metadata_list] self.assertEqual(files, []) def test_match_file_error(self): url = self.fs.join(self.tmpdir, 'old_file1') bad_url = 'bad_url' with self.assertRaisesRegex(BeamIOError, r'^Match operation failed .* %s' % bad_url): result = self.fs.match([bad_url, url])[0] files = [f.path for f in result.metadata_list] self.assertEqual(files, [self.fs._parse_url(url)]) def test_match_directory(self): expected_files = [ self.fs.join(self.tmpdir, filename) for filename in ['old_file1', 'old_file2'] ] # Listing without a trailing '/' should return the directory itself and not # its contents. The fake HDFS client here has a "sparse" directory # structure, so listing without a '/' will return no results. result = self.fs.match([self.tmpdir + '/'])[0] files = [f.path for f in result.metadata_list] self.assertCountEqual(files, expected_files) def test_match_directory_trailing_slash(self): expected_files = [ self.fs.join(self.tmpdir, filename) for filename in ['old_file1', 'old_file2'] ] result = self.fs.match([self.tmpdir + '/'])[0] files = [f.path for f in result.metadata_list] self.assertCountEqual(files, expected_files) def test_create_success(self): url = self.fs.join(self.tmpdir, 'new_file') handle = self.fs.create(url) self.assertIsNotNone(handle) _, url = self.fs._parse_url(url) expected_file = FakeFile(url, 'wb') self.assertEqual(self._fake_hdfs.files[url], expected_file) def test_create_write_read_compressed(self): url = self.fs.join(self.tmpdir, 'new_file.gz') handle = self.fs.create(url) self.assertIsNotNone(handle) _, path = self.fs._parse_url(url) expected_file = FakeFile(path, 'wb') self.assertEqual(self._fake_hdfs.files[path], expected_file) data = b'abc' * 10 handle.write(data) # Compressed data != original data self.assertNotEqual(data, self._fake_hdfs.files[path].getvalue()) handle.close() handle = self.fs.open(url) read_data = handle.read(len(data)) self.assertEqual(data, read_data) handle.close() def test_random_read_large_file(self): # this tests HdfsDownloader.get_range() works properly with # filesystemio.readinto when reading a file of size larger than the buffer. url = self.fs.join(self.tmpdir, 'read_length') handle = self.fs.create(url) data = b'test' * 10_000_000 handle.write(data) handle.close() handle = self.fs.open(url) handle.seek(100) # read 3 bytes read_data = handle.read(3) self.assertEqual(data[100:103], read_data) # read 4 bytes read_data = handle.read(4) self.assertEqual(data[103:107], read_data) def test_open(self): url = self.fs.join(self.tmpdir, 'old_file1') handle = self.fs.open(url) expected_data = b'' data = handle.read() self.assertEqual(data, expected_data) def test_open_bad_path(self): with self.assertRaises(FakeHdfsError): self.fs.open(self.fs.join(self.tmpdir, 'nonexistent/path')) def _cmpfiles(self, url1, url2): with self.fs.open(url1) as f1: with self.fs.open(url2) as f2: data1 = f1.read() data2 = f2.read() return data1 == data2 def test_copy_file(self): url1 = self.fs.join(self.tmpdir, 'new_file1') url2 = self.fs.join(self.tmpdir, 'new_file2') url3 = self.fs.join(self.tmpdir, 'new_file3') with self.fs.create(url1) as f1: f1.write(b'Hello') self.fs.copy([url1, url1], [url2, url3]) self.assertTrue(self._cmpfiles(url1, url2)) self.assertTrue(self._cmpfiles(url1, url3)) def test_copy_file_overwrite_error(self): url1 = self.fs.join(self.tmpdir, 'new_file1') url2 = self.fs.join(self.tmpdir, 'new_file2') with self.fs.create(url1) as f1: f1.write(b'Hello') with self.fs.create(url2) as f2: f2.write(b'nope') with self.assertRaisesRegex(BeamIOError, r'already exists.*%s' % posixpath.basename(url2)): self.fs.copy([url1], [url2]) def test_copy_file_error(self): url1 = self.fs.join(self.tmpdir, 'new_file1') url2 = self.fs.join(self.tmpdir, 'new_file2') url3 = self.fs.join(self.tmpdir, 'new_file3') url4 = self.fs.join(self.tmpdir, 'new_file4') with self.fs.create(url3) as f: f.write(b'Hello') with self.assertRaisesRegex(BeamIOError, r'^Copy operation failed .*%s.*%s.* not found' % (url1, url2)): self.fs.copy([url1, url3], [url2, url4]) self.assertTrue(self._cmpfiles(url3, url4)) def test_copy_directory(self): url_t1 = self.fs.join(self.tmpdir, 't1') url_t1_inner = self.fs.join(self.tmpdir, 't1/inner') url_t2 = self.fs.join(self.tmpdir, 't2') url_t2_inner = self.fs.join(self.tmpdir, 't2/inner') self.fs.mkdirs(url_t1) self.fs.mkdirs(url_t1_inner) self.fs.mkdirs(url_t2) url1 = self.fs.join(url_t1_inner, 'f1') url2 = self.fs.join(url_t2_inner, 'f1') with self.fs.create(url1) as f: f.write(b'Hello') self.fs.copy([url_t1], [url_t2]) self.assertTrue(self._cmpfiles(url1, url2)) def test_copy_directory_overwrite_error(self): url_t1 = self.fs.join(self.tmpdir, 't1') url_t1_inner = self.fs.join(self.tmpdir, 't1/inner') url_t2 = self.fs.join(self.tmpdir, 't2') url_t2_inner = self.fs.join(self.tmpdir, 't2/inner') self.fs.mkdirs(url_t1) self.fs.mkdirs(url_t1_inner) self.fs.mkdirs(url_t2) self.fs.mkdirs(url_t2_inner) url1 = self.fs.join(url_t1, 'f1') url1_inner = self.fs.join(url_t1_inner, 'f2') url2 = self.fs.join(url_t2, 'f1') unused_url2_inner = self.fs.join(url_t2_inner, 'f2') url3_inner = self.fs.join(url_t2_inner, 'f3') for url in [url1, url1_inner, url3_inner]: with self.fs.create(url) as f: f.write(b'Hello') with self.fs.create(url2) as f: f.write(b'nope') with self.assertRaisesRegex(BeamIOError, r'already exists'): self.fs.copy([url_t1], [url_t2]) def test_rename_file(self): url1 = self.fs.join(self.tmpdir, 'f1') url2 = self.fs.join(self.tmpdir, 'f2') with self.fs.create(url1) as f: f.write(b'Hello') self.fs.rename([url1], [url2]) self.assertFalse(self.fs.exists(url1)) self.assertTrue(self.fs.exists(url2)) def test_rename_file_error(self): url1 = self.fs.join(self.tmpdir, 'f1') url2 = self.fs.join(self.tmpdir, 'f2') url3 = self.fs.join(self.tmpdir, 'f3') url4 = self.fs.join(self.tmpdir, 'f4') with self.fs.create(url3) as f: f.write(b'Hello') with self.assertRaisesRegex(BeamIOError, r'^Rename operation failed .*%s.*%s' % (url1, url2)): self.fs.rename([url1, url3], [url2, url4]) self.assertFalse(self.fs.exists(url3)) self.assertTrue(self.fs.exists(url4)) def test_rename_directory(self): url_t1 = self.fs.join(self.tmpdir, 't1') url_t2 = self.fs.join(self.tmpdir, 't2') self.fs.mkdirs(url_t1) url1 = self.fs.join(url_t1, 'f1') url2 = self.fs.join(url_t2, 'f1') with self.fs.create(url1) as f: f.write(b'Hello') self.fs.rename([url_t1], [url_t2]) self.assertFalse(self.fs.exists(url_t1)) self.assertTrue(self.fs.exists(url_t2)) self.assertFalse(self.fs.exists(url1)) self.assertTrue(self.fs.exists(url2)) def test_exists(self): url1 = self.fs.join(self.tmpdir, 'old_file1') url2 = self.fs.join(self.tmpdir, 'nonexistent') self.assertTrue(self.fs.exists(url1)) self.assertFalse(self.fs.exists(url2)) def test_size(self): url = self.fs.join(self.tmpdir, 'f1') with self.fs.create(url) as f: f.write(b'Hello') self.assertEqual(5, self.fs.size(url)) def test_checksum(self): url = self.fs.join(self.tmpdir, 'f1') with self.fs.create(url) as f: f.write(b'Hello') self.assertEqual( 'fake_algo-5-checksum_byte_sequence', self.fs.checksum(url)) def test_last_updated(self): url = self.fs.join(self.tmpdir, 'f1') with self.fs.create(url) as f: f.write(b'Hello') # The time difference should be tiny for the mock hdfs. # A loose tolerance is for the consideration of real web hdfs. tolerance = 5 * 60 # 5 mins result = self.fs.last_updated(url) self.assertAlmostEqual(result, time.time(), delta=tolerance) def test_delete_file(self): url = self.fs.join(self.tmpdir, 'old_file1') self.assertTrue(self.fs.exists(url)) self.fs.delete([url]) self.assertFalse(self.fs.exists(url)) def test_delete_dir(self): url_t1 = self.fs.join(self.tmpdir, 'new_dir1') url_t2 = self.fs.join(url_t1, 'new_dir2') url1 = self.fs.join(url_t2, 'new_file1') url2 = self.fs.join(url_t2, 'new_file2') self.fs.mkdirs(url_t1) self.fs.mkdirs(url_t2) self.fs.create(url1).close() self.fs.create(url2).close() self.assertTrue(self.fs.exists(url1)) self.assertTrue(self.fs.exists(url2)) self.fs.delete([url_t1]) self.assertFalse(self.fs.exists(url_t1)) self.assertFalse(self.fs.exists(url_t2)) self.assertFalse(self.fs.exists(url2)) self.assertFalse(self.fs.exists(url1)) def test_delete_error(self): url1 = self.fs.join(self.tmpdir, 'nonexistent') url2 = self.fs.join(self.tmpdir, 'old_file1') self.assertTrue(self.fs.exists(url2)) _, path1 = self.fs._parse_url(url1) with self.assertRaisesRegex(BeamIOError, r'^Delete operation failed .* %s' % path1): self.fs.delete([url1, url2]) self.assertFalse(self.fs.exists(url2)) @unittest.skipIf(actual_hdfs is None, "hdfs extra not installed") class HadoopFileSystemRuntimeValueProviderTest(unittest.TestCase): """Tests pipeline_options, in the form of a RuntimeValueProvider.runtime_options object.""" def setUp(self): self._fake_hdfs = FakeHdfs() hdfs.hdfs.InsecureClient = (lambda *args, **kwargs: self._fake_hdfs) def test_dict_options(self): pipeline_options = { 'hdfs_host': '', 'hdfs_port': 0, 'hdfs_user': '', } self.fs = hdfs.HadoopFileSystem(pipeline_options=pipeline_options) self.assertFalse(self.fs._full_urls) def test_dict_options_missing(self): with self.assertRaisesRegex(ValueError, r'hdfs_host'): self.fs = hdfs.HadoopFileSystem( pipeline_options={ 'hdfs_port': 0, 'hdfs_user': '', }) with self.assertRaisesRegex(ValueError, r'hdfs_port'): self.fs = hdfs.HadoopFileSystem( pipeline_options={ 'hdfs_host': '', 'hdfs_user': '', }) with self.assertRaisesRegex(ValueError, r'hdfs_user'): self.fs = hdfs.HadoopFileSystem( pipeline_options={ 'hdfs_host': '', 'hdfs_port': 0, }) def test_dict_options_full_urls(self): pipeline_options = { 'hdfs_host': '', 'hdfs_port': 0, 'hdfs_user': '', 'hdfs_full_urls': 'invalid', } with self.assertRaisesRegex(ValueError, r'hdfs_full_urls'): self.fs = hdfs.HadoopFileSystem(pipeline_options=pipeline_options) pipeline_options['hdfs_full_urls'] = True self.fs = hdfs.HadoopFileSystem(pipeline_options=pipeline_options) self.assertTrue(self.fs._full_urls) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/iobase.py0000644000175100017510000022427315206113032020112 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Sources and sinks. A Source manages record-oriented data input from a particular kind of source (e.g. a set of files, a database table, etc.). The reader() method of a source returns a reader object supporting the iterator protocol; iteration yields raw records of unprocessed, serialized data. A Sink manages record-oriented data output to a particular kind of sink (e.g. a set of files, a database table, etc.). The writer() method of a sink returns a writer object supporting writing records of serialized data to the sink. """ # pytype: skip-file import logging import math import random import uuid from collections import namedtuple from typing import Any from typing import Iterator from typing import Optional from typing import Union import apache_beam as beam from apache_beam import coders from apache_beam import pvalue from apache_beam.coders.coders import _MemoizingPickleCoder from apache_beam.internal import pickler from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.pvalue import AsIter from apache_beam.pvalue import AsSingleton from apache_beam.transforms import Impulse from apache_beam.transforms import PTransform from apache_beam.transforms import core from apache_beam.transforms import ptransform from apache_beam.transforms import window from apache_beam.transforms.core import DoFn from apache_beam.transforms.display import DisplayDataItem from apache_beam.transforms.display import HasDisplayData from apache_beam.utils import timestamp from apache_beam.utils import urns from apache_beam.utils.windowed_value import WindowedValue __all__ = [ 'BoundedSource', 'RangeTracker', 'Read', 'RestrictionProgress', 'RestrictionTracker', 'WatermarkEstimator', 'Sink', 'Write', 'Writer' ] _LOGGER = logging.getLogger(__name__) # Encapsulates information about a bundle of a source generated when method # BoundedSource.split() is invoked. # This is a named 4-tuple that has following fields. # * weight - a number that represents the size of the bundle. This value will # be used to compare the relative sizes of bundles generated by the # current source. # The weight returned here could be specified using a unit of your # choice (for example, bundles of sizes 100MB, 200MB, and 700MB may # specify weights 100, 200, 700 or 1, 2, 7) but all bundles of a # source should specify the weight using the same unit. # * source - a BoundedSource object for the bundle. # * start_position - starting position of the bundle # * stop_position - ending position of the bundle. # # Type for start and stop positions are specific to the bounded source and must # be consistent throughout. SourceBundle = namedtuple( 'SourceBundle', 'weight source start_position stop_position') class SourceBase(HasDisplayData, urns.RunnerApiFn): """Base class for all sources that can be passed to beam.io.Read(...). """ urns.RunnerApiFn.register_pickle_urn(python_urns.PICKLED_SOURCE) def default_output_coder(self): raise NotImplementedError def is_bounded(self) -> bool: raise NotImplementedError class BoundedSource(SourceBase): """A source that reads a finite amount of input records. This class defines following operations which can be used to read the source efficiently. * Size estimation - method ``estimate_size()`` may return an accurate estimation in bytes for the size of the source. * Splitting into bundles of a given size - method ``split()`` can be used to split the source into a set of sub-sources (bundles) based on a desired bundle size. * Getting a RangeTracker - method ``get_range_tracker()`` should return a ``RangeTracker`` object for a given position range for the position type of the records returned by the source. * Reading the data - method ``read()`` can be used to read data from the source while respecting the boundaries defined by a given ``RangeTracker``. A runner will perform reading the source in two steps. (1) Method ``get_range_tracker()`` will be invoked with start and end positions to obtain a ``RangeTracker`` for the range of positions the runner intends to read. Source must define a default initial start and end position range. These positions must be used if the start and/or end positions passed to the method ``get_range_tracker()`` are ``None`` (2) Method read() will be invoked with the ``RangeTracker`` obtained in the previous step. **Mutability** A ``BoundedSource`` object should not be mutated while its methods (for example, ``read()``) are being invoked by a runner. Runner implementations may invoke methods of ``BoundedSource`` objects through multi-threaded and/or reentrant execution modes. """ def estimate_size(self) -> Optional[int]: """Estimates the size of source in bytes. An estimate of the total size (in bytes) of the data that would be read from this source. This estimate is in terms of external storage size, before performing decompression or other processing. Returns: estimated size of the source if the size can be determined, ``None`` otherwise. """ raise NotImplementedError def split( self, desired_bundle_size: int, start_position: Optional[Any] = None, stop_position: Optional[Any] = None, ) -> Iterator[SourceBundle]: """Splits the source into a set of bundles. Bundles should be approximately of size ``desired_bundle_size`` bytes. Args: desired_bundle_size: the desired size (in bytes) of the bundles returned. start_position: if specified the given position must be used as the starting position of the first bundle. stop_position: if specified the given position must be used as the ending position of the last bundle. Returns: an iterator of objects of type 'SourceBundle' that gives information about the generated bundles. """ raise NotImplementedError def get_range_tracker( self, start_position: Optional[Any], stop_position: Optional[Any], ) -> 'RangeTracker': """Returns a RangeTracker for a given position range. Framework may invoke ``read()`` method with the RangeTracker object returned here to read data from the source. Args: start_position: starting position of the range. If 'None' default start position of the source must be used. stop_position: ending position of the range. If 'None' default stop position of the source must be used. Returns: a ``RangeTracker`` for the given position range. """ raise NotImplementedError def read(self, range_tracker): """Returns an iterator that reads data from the source. The returned set of data must respect the boundaries defined by the given ``RangeTracker`` object. For example: * Returned set of data must be for the range ``[range_tracker.start_position, range_tracker.stop_position)``. Note that a source may decide to return records that start after ``range_tracker.stop_position``. See documentation in class ``RangeTracker`` for more details. Also, note that framework might invoke ``range_tracker.try_split()`` to perform dynamic split operations. range_tracker.stop_position may be updated dynamically due to successful dynamic split operations. * Method ``range_tracker.try_split()`` must be invoked for every record that starts at a split point. * Method ``range_tracker.record_current_position()`` may be invoked for records that do not start at split points. Args: range_tracker: a ``RangeTracker`` whose boundaries must be respected when reading data from the source. A runner that reads this source muss pass a ``RangeTracker`` object that is not ``None``. Returns: an iterator of data read by the source. """ raise NotImplementedError def default_output_coder(self): """Coder that should be used for the records returned by the source. Should be overridden by sources that produce objects that can be encoded more efficiently than pickling. """ return coders.registry.get_coder(object) def is_bounded(self): return True class RangeTracker(object): """A thread safe object used by Dataflow source framework. A Dataflow source is defined using a ''BoundedSource'' and a ''RangeTracker'' pair. A ''RangeTracker'' is used by Dataflow source framework to perform dynamic work rebalancing of position-based sources. **Position-based sources** A position-based source is one where the source can be described by a range of positions of an ordered type and the records returned by the reader can be described by positions of the same type. In case a record occupies a range of positions in the source, the most important thing about the record is the position where it starts. Defining the semantics of positions for a source is entirely up to the source class, however the chosen definitions have to obey certain properties in order to make it possible to correctly split the source into parts, including dynamic splitting. Two main aspects need to be defined: 1. How to assign starting positions to records. 2. Which records should be read by a source with a range '[A, B)'. Moreover, reading a range must be *efficient*, i.e., the performance of reading a range should not significantly depend on the location of the range. For example, reading the range [A, B) should not require reading all data before 'A'. The sections below explain exactly what properties these definitions must satisfy, and how to use a ``RangeTracker`` with a properly defined source. **Properties of position-based sources** The main requirement for position-based sources is *associativity*: reading records from '[A, B)' and records from '[B, C)' should give the same records as reading from '[A, C)', where 'A <= B <= C'. This property ensures that no matter how a range of positions is split into arbitrarily many sub-ranges, the total set of records described by them stays the same. The other important property is how the source's range relates to positions of records in the source. In many sources each record can be identified by a unique starting position. In this case: * All records returned by a source '[A, B)' must have starting positions in this range. * All but the last record should end within this range. The last record may or may not extend past the end of the range. * Records should not overlap. Such sources should define "read '[A, B)'" as "read from the first record starting at or after 'A', up to but not including the first record starting at or after 'B'". Some examples of such sources include reading lines or CSV from a text file, reading keys and values from a BigTable, etc. The concept of *split points* allows to extend the definitions for dealing with sources where some records cannot be identified by a unique starting position. In all cases, all records returned by a source '[A, B)' must *start* at or after 'A'. **Split points** Some sources may have records that are not directly addressable. For example, imagine a file format consisting of a sequence of compressed blocks. Each block can be assigned an offset, but records within the block cannot be directly addressed without decompressing the block. Let us refer to this hypothetical format as CBF (Compressed Blocks Format). Many such formats can still satisfy the associativity property. For example, in CBF, reading '[A, B)' can mean "read all the records in all blocks whose starting offset is in '[A, B)'". To support such complex formats, we introduce the notion of *split points*. We say that a record is a split point if there exists a position 'A' such that the record is the first one to be returned when reading the range '[A, infinity)'. In CBF, the only split points would be the first records in each block. Split points allow us to define the meaning of a record's position and a source's range in all cases: * For a record that is at a split point, its position is defined to be the largest 'A' such that reading a source with the range '[A, infinity)' returns this record. * Positions of other records are only required to be non-decreasing. * Reading the source '[A, B)' must return records starting from the first split point at or after 'A', up to but not including the first split point at or after 'B'. In particular, this means that the first record returned by a source MUST always be a split point. * Positions of split points must be unique. As a result, for any decomposition of the full range of the source into position ranges, the total set of records will be the full set of records in the source, and each record will be read exactly once. **Consumed positions** As the source is being read, and records read from it are being passed to the downstream transforms in the pipeline, we say that positions in the source are being *consumed*. When a reader has read a record (or promised to a caller that a record will be returned), positions up to and including the record's start position are considered *consumed*. Dynamic splitting can happen only at *unconsumed* positions. If the reader just returned a record at offset 42 in a file, dynamic splitting can happen only at offset 43 or beyond, as otherwise that record could be read twice (by the current reader and by a reader of the task starting at 43). """ SPLIT_POINTS_UNKNOWN = object() def start_position(self): """Returns the starting position of the current range, inclusive.""" raise NotImplementedError(type(self)) def stop_position(self): """Returns the ending position of the current range, exclusive.""" raise NotImplementedError(type(self)) def try_claim(self, position): # pylint: disable=unused-argument """Atomically determines if a record at a split point is within the range. This method should be called **if and only if** the record is at a split point. This method may modify the internal state of the ``RangeTracker`` by updating the last-consumed position to ``position``. ** Thread safety ** Methods of the class ``RangeTracker`` including this method may get invoked by different threads, hence must be made thread-safe, e.g. by using a single lock object. Args: position: starting position of a record being read by a source. Returns: ``True``, if the given position falls within the current range, returns ``False`` otherwise. """ raise NotImplementedError def set_current_position(self, position): """Updates the last-consumed position to the given position. A source may invoke this method for records that do not start at split points. This may modify the internal state of the ``RangeTracker``. If the record starts at a split point, method ``try_claim()`` **must** be invoked instead of this method. Args: position: starting position of a record being read by a source. """ raise NotImplementedError def position_at_fraction(self, fraction): """Returns the position at the given fraction. Given a fraction within the range [0.0, 1.0) this method will return the position at the given fraction compared to the position range [self.start_position, self.stop_position). ** Thread safety ** Methods of the class ``RangeTracker`` including this method may get invoked by different threads, hence must be made thread-safe, e.g. by using a single lock object. Args: fraction: a float value within the range [0.0, 1.0). Returns: a position within the range [self.start_position, self.stop_position). """ raise NotImplementedError def try_split(self, position): """Atomically splits the current range. Determines a position to split the current range, split_position, based on the given position. In most cases split_position and position will be the same. Splits the current range '[self.start_position, self.stop_position)' into a "primary" part '[self.start_position, split_position)' and a "residual" part '[split_position, self.stop_position)', assuming the current last-consumed position is within '[self.start_position, split_position)' (i.e., split_position has not been consumed yet). If successful, updates the current range to be the primary and returns a tuple (split_position, split_fraction). split_fraction should be the fraction of size of range '[self.start_position, split_position)' compared to the original (before split) range '[self.start_position, self.stop_position)'. If the split_position has already been consumed, returns ``None``. ** Thread safety ** Methods of the class ``RangeTracker`` including this method may get invoked by different threads, hence must be made thread-safe, e.g. by using a single lock object. Args: position: suggested position where the current range should try to be split at. Returns: a tuple containing the split position and split fraction if split is successful. Returns ``None`` otherwise. """ raise NotImplementedError def fraction_consumed(self): """Returns the approximate fraction of consumed positions in the source. ** Thread safety ** Methods of the class ``RangeTracker`` including this method may get invoked by different threads, hence must be made thread-safe, e.g. by using a single lock object. Returns: the approximate fraction of positions that have been consumed by successful 'try_split()' and 'try_claim()' calls, or 0.0 if no such calls have happened. """ raise NotImplementedError def split_points(self): """Gives the number of split points consumed and remaining. For a ``RangeTracker`` used by a ``BoundedSource`` (within a ``BoundedSource.read()`` invocation) this method produces a 2-tuple that gives the number of split points consumed by the ``BoundedSource`` and the number of split points remaining within the range of the ``RangeTracker`` that has not been consumed by the ``BoundedSource``. More specifically, given that the position of the current record being read by ``BoundedSource`` is current_position this method produces a tuple that consists of (1) number of split points in the range [self.start_position(), current_position) without including the split point that is currently being consumed. This represents the total amount of parallelism in the consumed part of the source. (2) number of split points within the range [current_position, self.stop_position()) including the split point that is currently being consumed. This represents the total amount of parallelism in the unconsumed part of the source. Methods of the class ``RangeTracker`` including this method may get invoked by different threads, hence must be made thread-safe, e.g. by using a single lock object. ** General information about consumed and remaining number of split points returned by this method. ** * Before a source read (``BoundedSource.read()`` invocation) claims the first split point, number of consumed split points is 0. This condition holds independent of whether the input is "splittable". A splittable source is a source that has more than one split point. * Any source read that has only claimed one split point has 0 consumed split points since the first split point is the current split point and is still being processed. This condition holds independent of whether the input is splittable. * For an empty source read which never invokes ``RangeTracker.try_claim()``, the consumed number of split points is 0. This condition holds independent of whether the input is splittable. * For a source read which has invoked ``RangeTracker.try_claim()`` n times, the consumed number of split points is n -1. * If a ``BoundedSource`` sets a callback through function ``set_split_points_unclaimed_callback()``, ``RangeTracker`` can use that callback when determining remaining number of split points. * Remaining split points should include the split point that is currently being consumed by the source read. Hence if the above callback returns an integer value n, remaining number of split points should be (n + 1). * After last split point is claimed remaining split points becomes 1, because this unfinished read itself represents an unfinished split point. * After all records of the source has been consumed, remaining number of split points becomes 0 and consumed number of split points becomes equal to the total number of split points within the range being read by the source. This method does not address this condition and will continue to report number of consumed split points as ("total number of split points" - 1) and number of remaining split points as 1. A runner that performs the reading of the source can detect when all records have been consumed and adjust remaining and consumed number of split points accordingly. ** Examples ** (1) A "perfectly splittable" input which can be read in parallel down to the individual records. Consider a perfectly splittable input that consists of 50 split points. * Before a source read (``BoundedSource.read()`` invocation) claims the first split point, number of consumed split points is 0 number of remaining split points is 50. * After claiming first split point, consumed number of split points is 0 and remaining number of split is 50. * After claiming split point #30, consumed number of split points is 29 and remaining number of split points is 21. * After claiming all 50 split points, consumed number of split points is 49 and remaining number of split points is 1. (2) a "block-compressed" file format such as ``avroio``, in which a block of records has to be read as a whole, but different blocks can be read in parallel. Consider a block compressed input that consists of 5 blocks. * Before a source read (``BoundedSource.read()`` invocation) claims the first split point (first block), number of consumed split points is 0 number of remaining split points is 5. * After claiming first split point, consumed number of split points is 0 and remaining number of split is 5. * After claiming split point #3, consumed number of split points is 2 and remaining number of split points is 3. * After claiming all 5 split points, consumed number of split points is 4 and remaining number of split points is 1. (3) an "unsplittable" input such as a cursor in a database or a gzip compressed file. Such an input is considered to have only a single split point. Number of consumed split points is always 0 and number of remaining split points is always 1. By default ``RangeTracker` returns ``RangeTracker.SPLIT_POINTS_UNKNOWN`` for both consumed and remaining number of split points, which indicates that the number of split points consumed and remaining is unknown. Returns: A pair that gives consumed and remaining number of split points. Consumed number of split points should be an integer larger than or equal to zero or ``RangeTracker.SPLIT_POINTS_UNKNOWN``. Remaining number of split points should be an integer larger than zero or ``RangeTracker.SPLIT_POINTS_UNKNOWN``. """ return ( RangeTracker.SPLIT_POINTS_UNKNOWN, RangeTracker.SPLIT_POINTS_UNKNOWN) def set_split_points_unclaimed_callback(self, callback): """Sets a callback for determining the unclaimed number of split points. By invoking this function, a ``BoundedSource`` can set a callback function that may get invoked by the ``RangeTracker`` to determine the number of unclaimed split points. A split point is unclaimed if ``RangeTracker.try_claim()`` method has not been successfully invoked for that particular split point. The callback function accepts a single parameter, a stop position for the BoundedSource (stop_position). If the record currently being consumed by the ``BoundedSource`` is at position current_position, callback should return the number of split points within the range (current_position, stop_position). Note that, this should not include the split point that is currently being consumed by the source. This function must be implemented by subclasses before being used. Args: callback: a function that takes a single parameter, a stop position, and returns unclaimed number of split points for the source read operation that is calling this function. Value returned from callback should be either an integer larger than or equal to zero or ``RangeTracker.SPLIT_POINTS_UNKNOWN``. """ raise NotImplementedError class Sink(HasDisplayData): """This class is deprecated, no backwards-compatibility guarantees. A resource that can be written to using the ``beam.io.Write`` transform. Here ``beam`` stands for Apache Beam Python code imported in following manner. ``import apache_beam as beam``. A parallel write to an ``iobase.Sink`` consists of three phases: 1. A sequential *initialization* phase (e.g., creating a temporary output directory, etc.) 2. A parallel write phase where workers write *bundles* of records 3. A sequential *finalization* phase (e.g., committing the writes, merging output files, etc.) Implementing a new sink requires extending two classes. 1. iobase.Sink ``iobase.Sink`` is an immutable logical description of the location/resource to write to. Depending on the type of sink, it may contain fields such as the path to an output directory on a filesystem, a database table name, etc. ``iobase.Sink`` provides methods for performing a write operation to the sink described by it. To this end, implementors of an extension of ``iobase.Sink`` must implement three methods: ``initialize_write()``, ``open_writer()``, and ``finalize_write()``. 2. iobase.Writer ``iobase.Writer`` is used to write a single bundle of records. An ``iobase.Writer`` defines two methods: ``write()`` which writes a single record from the bundle and ``close()`` which is called once at the end of writing a bundle. See also ``apache_beam.io.filebasedsink.FileBasedSink`` which provides a simpler API for writing sinks that produce files. **Execution of the Write transform** ``initialize_write()``, ``pre_finalize()``, and ``finalize_write()`` are conceptually called once. However, implementors must ensure that these methods are *idempotent*, as they may be called multiple times on different machines in the case of failure/retry. A method may be called more than once concurrently, in which case it's okay to have a transient failure (such as due to a race condition). This failure should not prevent subsequent retries from succeeding. ``initialize_write()`` should perform any initialization that needs to be done prior to writing to the sink. ``initialize_write()`` may return a result (let's call this ``init_result``) that contains any parameters it wants to pass on to its writers about the sink. For example, a sink that writes to a file system may return an ``init_result`` that contains a dynamically generated unique directory to which data should be written. To perform writing of a bundle of elements, Dataflow execution engine will create an ``iobase.Writer`` using the implementation of ``iobase.Sink.open_writer()``. When invoking ``open_writer()`` execution engine will provide the ``init_result`` returned by ``initialize_write()`` invocation as well as a *bundle id* (let's call this ``bundle_id``) that is unique for each invocation of ``open_writer()``. Execution engine will then invoke ``iobase.Writer.write()`` implementation for each element that has to be written. Once all elements of a bundle are written, execution engine will invoke ``iobase.Writer.close()`` implementation which should return a result (let's call this ``write_result``) that contains information that encodes the result of the write and, in most cases, some encoding of the unique bundle id. For example, if each bundle is written to a unique temporary file, ``close()`` method may return an object that contains the temporary file name. After writing of all bundles is complete, execution engine will invoke ``pre_finalize()`` and then ``finalize_write()`` implementation. The execution of a write transform can be illustrated using following pseudo code (assume that the outer for loop happens in parallel across many machines):: init_result = sink.initialize_write() write_results = [] for bundle in partition(pcoll): writer = sink.open_writer(init_result, generate_bundle_id()) for elem in bundle: writer.write(elem) write_results.append(writer.close()) pre_finalize_result = sink.pre_finalize(init_result, write_results) sink.finalize_write(init_result, write_results, pre_finalize_result) **init_result** Methods of 'iobase.Sink' should agree on the 'init_result' type that will be returned when initializing the sink. This type can be a client-defined object or an existing type. The returned type must be picklable using Dataflow coder ``coders.PickleCoder``. Returning an init_result is optional. **bundle_id** In order to ensure fault-tolerance, a bundle may be executed multiple times (e.g., in the event of failure/retry or for redundancy). However, exactly one of these executions will have its result passed to the ``iobase.Sink.finalize_write()`` method. Each call to ``iobase.Sink.open_writer()`` is passed a unique bundle id when it is called by the ``WriteImpl`` transform, so even redundant or retried bundles will have a unique way of identifying their output. The bundle id should be used to guarantee that a bundle's output is unique. This uniqueness guarantee is important; if a bundle is to be output to a file, for example, the name of the file must be unique to avoid conflicts with other writers. The bundle id should be encoded in the writer result returned by the writer and subsequently used by the ``finalize_write()`` method to identify the results of successful writes. For example, consider the scenario where a Writer writes files containing serialized records and the ``finalize_write()`` is to merge or rename these output files. In this case, a writer may use its unique id to name its output file (to avoid conflicts) and return the name of the file it wrote as its writer result. The ``finalize_write()`` will then receive an ``Iterable`` of output file names that it can then merge or rename using some bundle naming scheme. **write_result** ``iobase.Writer.close()`` and ``finalize_write()`` implementations must agree on type of the ``write_result`` object returned when invoking ``iobase.Writer.close()``. This type can be a client-defined object or an existing type. The returned type must be picklable using Dataflow coder ``coders.PickleCoder``. Returning a ``write_result`` when ``iobase.Writer.close()`` is invoked is optional but if unique ``write_result`` objects are not returned, sink should, guarantee idempotency when same bundle is written multiple times due to failure/retry or redundancy. **More information** For more information on creating new sinks please refer to the official documentation at ``https://beam.apache.org/documentation/sdks/python-custom-io#creating-sinks`` """ # Whether Beam should skip writing any shards if all are empty. skip_if_empty = False # the Write triggering frequency in streaming with GlobalWindow triggering_frequency = 0 def initialize_write(self): """Initializes the sink before writing begins. Invoked before any data is written to the sink. Please see documentation in ``iobase.Sink`` for an example. Returns: An object that contains any sink specific state generated by initialization. This object will be passed to open_writer() and finalize_write() methods. """ raise NotImplementedError def open_writer(self, init_result, uid): """Opens a writer for writing a bundle of elements to the sink. Args: init_result: the result of initialize_write() invocation. uid: a unique identifier generated by the system. Returns: an ``iobase.Writer`` that can be used to write a bundle of records to the current sink. """ raise NotImplementedError def pre_finalize(self, init_result, writer_results): """Pre-finalization stage for sink for bounded PCollections. Called after all bundle writes are complete and before finalize_write. Used to setup and verify filesystem and sink states. Args: init_result: the result of ``initialize_write()`` invocation. writer_results: an iterable containing results of ``Writer.close()`` invocations. This will only contain results of successful writes, and will only contain the result of a single successful write for a given bundle. Returns: An object that contains any sink specific state generated. This object will be passed to finalize_write(). """ raise NotImplementedError def pre_finalize_windowed(self, init_result, writer_results, window=None): """Pre-finalization stage for sink for unbounded PCollections. Called after all bundle writes are complete and before finalize_write. Used to setup and verify filesystem and sink states. Args: init_result: the result of ``initialize_write()`` invocation. writer_results: an iterable containing results of ``Writer.close()`` invocations. This will only contain results of successful writes, and will only contain the result of a single successful write for a given bundle. window: DoFn window Returns: An object that contains any sink specific state generated. This object will be passed to finalize_windowed_write(). """ raise NotImplementedError def finalize_write(self, init_result, writer_results, pre_finalize_result): """Finalizes the sink after all data is written to it. (batch) Given the result of initialization and an iterable of results from bundle writes, performs finalization after writing and closes the sink. Called after all bundle writes are complete. The bundle write results that are passed to finalize are those returned by bundles that completed successfully. Although bundles may have been run multiple times (for fault-tolerance), only one writer result will be passed to finalize for each bundle. An implementation of finalize should perform clean up of any failed and successfully retried bundles. Note that these failed bundles will not have their writer result passed to finalize, so finalize should be capable of locating any temporary/partial output written by failed bundles. If all retries of a bundle fails, the whole pipeline will fail *without* finalize_write() being invoked. A best practice is to make finalize atomic. If this is impossible given the semantics of the sink, finalize should be idempotent, as it may be called multiple times in the case of failure/retry or for redundancy. Note that the iteration order of the writer results is not guaranteed to be consistent if finalize is called multiple times. Args: init_result: the result of ``initialize_write()`` invocation. writer_results: an iterable containing results of ``Writer.close()`` invocations. This will only contain results of successful writes, and will only contain the result of a single successful write for a given bundle. pre_finalize_result: the result of ``pre_finalize()`` invocation. """ raise NotImplementedError def finalize_windowed_write( self, init_result, writer_results, pre_finalize_result, window=None): """Finalizes the sink after all data is written to it for a window. Similar to ``finalize_write()``. Args: init_result: the result of ``initialize_write()`` invocation. writer_results: an iterable containing results of ``Writer.close()`` invocations. This will only contain results of successful writes, and will only contain the result of a single successful write for a given bundle. pre_finalize_result: the result of ``pre_finalize()`` invocation. window: DoFn window """ raise NotImplementedError class Writer(object): """This class is deprecated, no backwards-compatibility guarantees. Writes a bundle of elements from a ``PCollection`` to a sink. A Writer ``iobase.Writer.write()`` writes and elements to the sink while ``iobase.Writer.close()`` is called after all elements in the bundle have been written. See ``iobase.Sink`` for more detailed documentation about the process of writing to a sink. """ def write(self, value): """Writes a value to the sink using the current writer. """ raise NotImplementedError def close(self): """Closes the current writer. Please see documentation in ``iobase.Sink`` for an example. Returns: An object representing the writes that were performed by the current writer. """ raise NotImplementedError def at_capacity(self) -> bool: """Returns whether this writer should be considered at capacity and a new one should be created. """ return False class Read(ptransform.PTransform): """A transform that reads a PCollection.""" # Import runners here to prevent circular imports from apache_beam.runners.pipeline_context import PipelineContext def __init__(self, source: SourceBase) -> None: """Initializes a Read transform. Args: source: Data source to read from. """ super().__init__() self.source = source @staticmethod def get_desired_chunk_size(total_size): if total_size: # 1MB = 1 shard, 1GB = 32 shards, 1TB = 1000 shards, 1PB = 32k shards chunk_size = max(1 << 20, 1000 * int(math.sqrt(total_size))) else: chunk_size = 64 << 20 # 64mb return chunk_size def expand(self, pbegin): if isinstance(self.source, BoundedSource): coders.registry.register_coder(BoundedSource, _MemoizingPickleCoder) display_data = self.source.display_data() or {} display_data['source'] = self.source.__class__ return ( pbegin | Impulse() | 'EmitSource' >> core.Map(lambda _: self.source).with_output_types(BoundedSource) | SDFBoundedSourceReader(display_data)) elif isinstance(self.source, ptransform.PTransform): # The Read transform can also admit a full PTransform as an input # rather than an anctual source. If the input is a PTransform, then # just apply it directly. return pbegin.pipeline | self.source else: # Treat Read itself as a primitive. return pvalue.PCollection( pbegin.pipeline, is_bounded=self.source.is_bounded()) def get_windowing(self, unused_inputs) -> core.Windowing: return core.Windowing(window.GlobalWindows()) def _infer_output_coder(self, input_type=None, input_coder=None) -> Optional[coders.Coder]: if isinstance(self.source, SourceBase): return self.source.default_output_coder() else: return None def display_data(self): return { 'source': DisplayDataItem(self.source.__class__, label='Read Source'), 'source_dd': self.source } def to_runner_api_parameter( self, context: PipelineContext, ) -> tuple[str, Any]: from apache_beam.io.gcp.pubsub import _PubSubSource if isinstance(self.source, _PubSubSource): return ( common_urns.composites.PUBSUB_READ.urn, beam_runner_api_pb2.PubSubReadPayload( topic=self.source.full_topic, subscription=self.source.full_subscription, timestamp_attribute=self.source.timestamp_attribute, with_attributes=self.source.with_attributes, id_attribute=self.source.id_label)) if isinstance(self.source, BoundedSource): return ( common_urns.deprecated_primitives.READ.urn, beam_runner_api_pb2.ReadPayload( source=self.source.to_runner_api(context), is_bounded=beam_runner_api_pb2.IsBounded.BOUNDED if self.source.is_bounded() else beam_runner_api_pb2.IsBounded.UNBOUNDED)) elif isinstance(self.source, ptransform.PTransform): return self.source.to_runner_api_parameter(context) raise NotImplementedError( "to_runner_api_parameter not " "implemented for type") @staticmethod def from_runner_api_parameter( transform: beam_runner_api_pb2.PTransform, payload: Union[beam_runner_api_pb2.ReadPayload, beam_runner_api_pb2.PubSubReadPayload], context: PipelineContext, ) -> "Read": if transform.spec.urn == common_urns.composites.PUBSUB_READ.urn: assert isinstance(payload, beam_runner_api_pb2.PubSubReadPayload) # Importing locally to prevent circular dependencies. # TODO(BEAM-27443): Remove the need for this. from apache_beam.io.gcp.pubsub import _PubSubSource source = _PubSubSource( topic=payload.topic or None, subscription=payload.subscription or None, id_label=payload.id_attribute or None, with_attributes=payload.with_attributes, timestamp_attribute=payload.timestamp_attribute or None) return Read(source) else: assert isinstance(payload, beam_runner_api_pb2.ReadPayload) return Read(SourceBase.from_runner_api(payload.source, context)) @staticmethod def _from_runner_api_parameter_read( transform: beam_runner_api_pb2.PTransform, payload: beam_runner_api_pb2.ReadPayload, context: PipelineContext, ) -> "Read": """Method for type proxying when calling register_urn due to limitations in type exprs in Python""" return Read.from_runner_api_parameter(transform, payload, context) @staticmethod def _from_runner_api_parameter_pubsub_read( transform: beam_runner_api_pb2.PTransform, payload: beam_runner_api_pb2.PubSubReadPayload, context: PipelineContext, ) -> "Read": """Method for type proxying when calling register_urn due to limitations in type exprs in Python""" return Read.from_runner_api_parameter(transform, payload, context) ptransform.PTransform.register_urn( common_urns.deprecated_primitives.READ.urn, beam_runner_api_pb2.ReadPayload, Read._from_runner_api_parameter_read, ) # TODO(BEAM-27443): Remove. ptransform.PTransform.register_urn( common_urns.composites.PUBSUB_READ.urn, beam_runner_api_pb2.PubSubReadPayload, Read._from_runner_api_parameter_pubsub_read, ) class Write(ptransform.PTransform): """A ``PTransform`` that writes to a sink. A sink should inherit ``iobase.Sink``. Such implementations are handled using a composite transform that consists of three ``ParDo``s - (1) a ``ParDo`` performing a global initialization (2) a ``ParDo`` performing a parallel write and (3) a ``ParDo`` performing a global finalization. In the case of an empty ``PCollection``, only the global initialization and finalization will be performed. Currently only batch workflows support custom sinks. Example usage:: pcollection | beam.io.Write(MySink()) This returns a ``pvalue.PValue`` object that represents the end of the Pipeline. The sink argument may also be a full PTransform, in which case it will be applied directly. This allows composite sink-like transforms (e.g. a sink with some pre-processing DoFns) to be used the same as all other sinks. This transform also supports sinks that inherit ``iobase.NativeSink``. These are sinks that are implemented natively by the Dataflow service and hence should not be updated by users. These sinks are processed using a Dataflow native write transform. """ # Import runners here to prevent circular imports from apache_beam.runners.pipeline_context import PipelineContext def __init__(self, sink): """Initializes a Write transform. Args: sink: Data sink to write to. """ super().__init__() self.sink = sink def display_data(self): return {'sink': self.sink.__class__, 'sink_dd': self.sink} def expand(self, pcoll): # Importing locally to prevent circular dependencies. from apache_beam.io.gcp.pubsub import _PubSubSink if isinstance(self.sink, _PubSubSink): # TODO(BEAM-27443): Remove the need for special casing here. return pvalue.PDone(pcoll.pipeline) elif isinstance(self.sink, Sink): # A custom sink return pcoll | WriteImpl(self.sink) elif isinstance(self.sink, ptransform.PTransform): # This allows "composite" sinks to be used like non-composite ones. return pcoll | self.sink else: raise ValueError( 'A sink must inherit iobase.Sink, iobase.NativeSink, ' 'or be a PTransform. Received : %r' % self.sink) def to_runner_api_parameter( self, context: PipelineContext, ) -> tuple[str, Any]: # TODO(BEAM-27443): Remove the need for special casing here. # Importing locally to prevent circular dependencies. from apache_beam.io.gcp.pubsub import _PubSubSink if isinstance(self.sink, _PubSubSink): payload = beam_runner_api_pb2.PubSubWritePayload( topic=self.sink.full_topic, id_attribute=self.sink.id_label, timestamp_attribute=self.sink.timestamp_attribute) return (common_urns.composites.PUBSUB_WRITE.urn, payload) else: return super().to_runner_api_parameter(context) @staticmethod @ptransform.PTransform.register_urn( common_urns.composites.PUBSUB_WRITE.urn, beam_runner_api_pb2.PubSubWritePayload) def from_runner_api_parameter( ptransform: Any, payload: beam_runner_api_pb2.PubSubWritePayload, unused_context: PipelineContext, ) -> "Write": if ptransform.spec.urn != common_urns.composites.PUBSUB_WRITE.urn: raise ValueError( 'Write transform cannot be constructed for the given proto %r', ptransform) if not payload.topic: raise NotImplementedError( "from_runner_api_parameter does not " "handle empty or None topic") # Importing locally to prevent circular dependencies. from apache_beam.io.gcp.pubsub import _PubSubSink sink = _PubSubSink( topic=payload.topic, id_label=payload.id_attribute or None, timestamp_attribute=payload.timestamp_attribute or None) return Write(sink) class WriteImpl(ptransform.PTransform): """Implements the writing of custom sinks.""" def __init__(self, sink: Sink) -> None: super().__init__() self.sink = sink def _expand_bounded(self, pcoll, min_shards): """Handles the expansion logic for a bounded PCollection.""" do_once = pcoll.pipeline | 'DoOnce' >> core.Create([None]) init_result_coll = do_once | 'InitializeWrite' >> core.Map( lambda _, sink: sink.initialize_write(), self.sink) if (min_shards >= 1): if min_shards == 1: keyed_pcoll = pcoll | core.Map(lambda x: (None, x)) else: keyed_pcoll = pcoll | core.ParDo(_RoundRobinKeyFn(), count=min_shards) write_result_coll = ( keyed_pcoll | core.WindowInto(window.GlobalWindows()) | core.GroupByKey() | 'WriteBundles' >> core.ParDo( _WriteKeyedBundleDoFn(self.sink), AsSingleton(init_result_coll))) else: # min_shards is 0, num_shards was undef _LOGGER.info( "*** WriteImpl min_shards undef so it's 1, and we write per Bundle") min_shards = 1 write_result_coll = ( pcoll | core.WindowInto(window.GlobalWindows()) | 'WriteBundles' >> core.ParDo( _WriteBundleDoFn(self.sink), AsSingleton(init_result_coll)) | 'Pair' >> core.Map(lambda x: (None, x)) | core.GroupByKey() | 'Extract' >> core.FlatMap(lambda x: x[1])) pre_finalize_coll = ( do_once | 'PreFinalize' >> core.FlatMap( _pre_finalize, self.sink, AsSingleton(init_result_coll), AsIter(write_result_coll))) return ( do_once | 'FinalizeWrite' >> core.FlatMap( _finalize_write, self.sink, AsSingleton(init_result_coll), AsIter(write_result_coll), min_shards, AsSingleton(pre_finalize_coll)).with_output_types(str)) def _apply_windowing(self, pcoll): """ Applies windowing to an unbounded PCollection based on the sink's triggering frequency. """ use_fixed_windows = ( self.sink.triggering_frequency is not None and self.sink.triggering_frequency > 0) if isinstance(pcoll.windowing.windowfn, window.GlobalWindows): if not use_fixed_windows: raise ValueError( 'To write a GlobalWindow unbounded PCollection, ' 'triggering_frequency must be set and be greater than 0') return pcoll | 'ApplyFixedWindows' >> core.WindowInto( window.FixedWindows(self.sink.triggering_frequency), trigger=beam.transforms.trigger.AfterWatermark(), accumulation_mode=beam.transforms.trigger.AccumulationMode.DISCARDING, allowed_lateness=beam.utils.timestamp.Duration(seconds=0)) # Keep user-defined windowing unless triggering_frequency is specified. if use_fixed_windows: return pcoll | 'ApplyFixedWindows' >> core.WindowInto( window.FixedWindows(self.sink.triggering_frequency), trigger=beam.transforms.trigger.AfterWatermark(), accumulation_mode=beam.transforms.trigger.AccumulationMode.DISCARDING, allowed_lateness=beam.utils.timestamp.Duration(seconds=0)) return pcoll # Keep original windowing def _expand_unbounded(self, pcoll, min_shards): """Handles the expansion logic for an unbounded PCollection.""" windowed_pcoll = self._apply_windowing(pcoll) if self.sink.convert_fn is not None: windowed_pcoll = windowed_pcoll | core.ParDo(self.sink.convert_fn) init_result_window_coll = ( windowed_pcoll | 'PairInit' >> core.Map(lambda x: (None, x)) | 'GroupInit' >> core.GroupByKey() | 'InitializeWindowedWrite' >> core.Map( lambda _, sink: sink.initialize_write(), self.sink)) if min_shards > 1: keyed_pcoll = windowed_pcoll | core.ParDo( _RoundRobinKeyFn(), count=min_shards) grouped_pcoll = keyed_pcoll | 'Group by random key' >> core.GroupByKey() else: # min_shards is 1 min_shards = 1 _LOGGER.info("Writing unbounded data with a single shard per window.") keyed_pcoll = windowed_pcoll | core.Map(lambda x: (None, x)) grouped_pcoll = keyed_pcoll | 'Group by window' >> core.GroupByKey() write_result_coll = ( grouped_pcoll | 'WriteWindowedBundles' >> core.ParDo( _WriteWindowedBundleDoFn(sink=self.sink, per_key=True), AsSingleton(init_result_window_coll)) | 'PairFinal' >> core.Map(lambda x: (None, x)) | 'GroupFinal' >> core.GroupByKey() | 'ExtractFinal' >> core.Map(lambda x: x[1])) pre_finalized_write_result_coll = ( write_result_coll | 'PreFinalize' >> core.ParDo( _PreFinalizeWindowedBundleDoFn(self.sink), AsSingleton(init_result_window_coll))) return ( pre_finalized_write_result_coll | 'FinalizeWrite' >> core.FlatMap( _finalize_windowed_write, self.sink, AsSingleton(init_result_window_coll), AsSingleton(write_result_coll), min_shards, AsIter(pre_finalized_write_result_coll)).with_output_types(str)) def expand(self, pcoll): min_shards = getattr(self.sink, 'num_shards', 0) if (pcoll.is_bounded): return self._expand_bounded(pcoll, min_shards) else: return self._expand_unbounded(pcoll, min_shards) class _WriteBundleDoFn(core.DoFn): """A DoFn for writing elements to an iobase.Writer. Opens a writer at the first element and closes the writer at finish_bundle(). """ def __init__(self, sink): self.sink = sink def display_data(self): return {'sink_dd': self.sink} def start_bundle(self): self.writer = None def process(self, element, init_result): if self.writer is None: # We ignore UUID collisions here since they are extremely rare. self.writer = self.sink.open_writer(init_result, str(uuid.uuid4())) self.writer.write(element) if self.writer.at_capacity(): yield self.writer.close() self.writer = None def finish_bundle(self): if self.writer is not None: yield WindowedValue( self.writer.close(), window.GlobalWindow().max_timestamp(), [window.GlobalWindow()]) class _PreFinalizeWindowedBundleDoFn(core.DoFn): """A DoFn for writing elements to an iobase.Writer. Opens a writer at the first element and closes the writer at finish_bundle(). """ def __init__( self, sink, destination_fn=None, temp_directory=None, ): self.sink = sink self._temp_directory = temp_directory self.destination_fn = destination_fn def display_data(self): return {'sink_dd': self.sink} def process( self, element, init_result, w=core.DoFn.WindowParam, pane=core.DoFn.PaneInfoParam): self.sink.pre_finalize_windowed( init_result=init_result, writer_results=element, window=w) yield element class _WriteWindowedBundleDoFn(core.DoFn): """A DoFn for writing elements to an iobase.Writer. Opens a writer at the first element and closes the writer at finish_bundle(). """ def __init__(self, sink, per_key=False): self.sink = sink self.per_key = per_key def display_data(self): return {'sink_dd': self.sink} def start_bundle(self): self.writer = {} self.window = {} self.init_result = {} def process( self, element, init_result, w=core.DoFn.WindowParam, pane=core.DoFn.PaneInfoParam): if self.per_key: w_key = "%s_%s" % (w, element[0]) # key else: w_key = w if not w in self.writer: # We ignore UUID collisions here since they are extremely rare. self.window[w_key] = w self.writer[w_key] = self.sink.open_writer( init_result, '%s_%s' % (w_key, uuid.uuid4())) self.init_result[w_key] = init_result if self.per_key: for e in element[1]: # values self.writer[w_key].write(e) # value else: self.writer[w_key].write(element) if self.writer[w_key].at_capacity(): yield self.writer[w_key].close() self.writer[w_key] = None def finish_bundle(self): for w_key, writer in self.writer.items(): w = self.window[w_key] if writer is not None: closed = writer.temp_shard_path try: closed = writer.close() # TODO : improve sink closing for streaming except ValueError as exp: _LOGGER.info( "*** _WriteWindowedBundleDoFn finish_bundle closed ERROR %s", exp) yield WindowedValue( closed, timestamp=w.start, windows=[w] # TODO(pabloem) HOW DO WE GET THE PANE ) class _WriteKeyedBundleDoFn(core.DoFn): def __init__(self, sink): self.sink = sink def display_data(self): return {'sink_dd': self.sink} def process(self, element, init_result): bundle = element writer = self.sink.open_writer(init_result, str(uuid.uuid4())) for e in bundle[1]: # values writer.write(e) return [window.TimestampedValue(writer.close(), timestamp.MAX_TIMESTAMP)] def _pre_finalize(unused_element, sink, init_result, write_results): return sink.pre_finalize(init_result, write_results) def _finalize_write( unused_element, sink, init_result, write_results, min_shards, pre_finalize_results): write_results = list(write_results) extra_shards = [] if len(write_results) < min_shards: if write_results or not sink.skip_if_empty: _LOGGER.debug( 'Creating %s empty shard(s).', min_shards - len(write_results)) for _ in range(min_shards - len(write_results)): writer = sink.open_writer(init_result, str(uuid.uuid4())) extra_shards.append(writer.close()) outputs = sink.finalize_write( init_result, write_results + extra_shards, pre_finalize_results) if outputs: return ( window.TimestampedValue(v, timestamp.MAX_TIMESTAMP) for v in outputs) def _finalize_windowed_write( unused_element, sink, init_result, write_results, min_shards, pre_finalize_results, w=DoFn.WindowParam): write_results = list(write_results) extra_shards = [] if len(write_results) < min_shards: if write_results or not sink.skip_if_empty: _LOGGER.debug( 'Creating %s empty shard(s).', min_shards - len(write_results)) for _ in range(min_shards - len(write_results)): writer = sink.open_writer(init_result, str(uuid.uuid4())) extra_shards.append(writer.close()) outputs = sink.finalize_windowed_write( init_result, write_results + extra_shards, pre_finalize_results, w) if outputs: return (window.TimestampedValue(v, w.end) for v in outputs) class _RoundRobinKeyFn(core.DoFn): def start_bundle(self): self.counter = None def process(self, element, count): if self.counter is None: self.counter = random.randrange(0, count) self.counter = (1 + self.counter) % count yield self.counter, element class RestrictionTracker(object): """Manages access to a restriction. Keeps track of the restrictions claimed part for a Splittable DoFn. The restriction may be modified by different threads, however the system will ensure sufficient locking such that no methods on the restriction tracker will be called concurrently. See following documents for more details. * https://s.apache.org/splittable-do-fn * https://s.apache.org/splittable-do-fn-python-sdk """ def current_restriction(self): """Returns the current restriction. Returns a restriction accurately describing the full range of work the current ``DoFn.process()`` call will do, including already completed work. The current restriction returned by method may be updated dynamically due to due to concurrent invocation of other methods of the ``RestrictionTracker``, For example, ``split()``. This API is required to be implemented. Returns: a restriction object. """ raise NotImplementedError def current_progress(self) -> 'RestrictionProgress': """Returns a RestrictionProgress object representing the current progress. This API is recommended to be implemented. The runner can do a better job at parallel processing with better progress signals. """ raise NotImplementedError def check_done(self): """Checks whether the restriction has been fully processed. Called by the SDK harness after iterator returned by ``DoFn.process()`` has been fully read. This method must raise a `ValueError` if there is still any unclaimed work remaining in the restriction when this method is invoked. Exception raised must have an informative error message. This API is required to be implemented in order to make sure no data loss during SDK processing. Returns: ``True`` if current restriction has been fully processed. Raises: ValueError: if there is still any unclaimed work remaining. """ raise NotImplementedError def try_split(self, fraction_of_remainder): """Splits current restriction based on fraction_of_remainder. If splitting the current restriction is possible, the current restriction is split into a primary and residual restriction pair. This invocation updates the ``current_restriction()`` to be the primary restriction effectively having the current ``DoFn.process()`` execution responsible for performing the work that the primary restriction represents. The residual restriction will be executed in a separate ``DoFn.process()`` invocation (likely in a different process). The work performed by executing the primary and residual restrictions as separate ``DoFn.process()`` invocations MUST be equivalent to the work performed as if this split never occurred. The ``fraction_of_remainder`` should be used in a best effort manner to choose a primary and residual restriction based upon the fraction of the remaining work that the current ``DoFn.process()`` invocation is responsible for. For example, if a ``DoFn.process()`` was reading a file with a restriction representing the offset range [100, 200) and has processed up to offset 130 with a fraction_of_remainder of 0.7, the primary and residual restrictions returned would be [100, 179), [179, 200) (note: current_offset + fraction_of_remainder * remaining_work = 130 + 0.7 * 70 = 179). ``fraction_of_remainder`` = 0 means a checkpoint is required. The API is recommended to be implemented for batch pipeline given that it is very important for pipeline scaling and end to end pipeline execution. The API is required to be implemented for a streaming pipeline. Args: fraction_of_remainder: A hint as to the fraction of work the primary restriction should represent based upon the current known remaining amount of work. Returns: (primary_restriction, residual_restriction) if a split was possible, otherwise returns ``None``. """ raise NotImplementedError def try_claim(self, position): """Attempts to claim the block of work in the current restriction identified by the given position. Each claimed position MUST be a valid split point. If this succeeds, the DoFn MUST execute the entire block of work. If it fails, the ``DoFn.process()`` MUST return ``None`` without performing any additional work or emitting output (note that emitting output or performing work from ``DoFn.process()`` is also not allowed before the first call of this method). The API is required to be implemented. Args: position: current position that wants to be claimed. Returns: ``True`` if the position can be claimed as current_position. Otherwise, returns ``False``. """ raise NotImplementedError def is_bounded(self): """Returns whether the amount of work represented by the current restriction is bounded. The boundedness of the restriction is used to determine the default behavior of how to truncate restrictions when a pipeline is being `drained `_. # pylint: disable=line-too-long If the restriction is bounded, then the entire restriction will be processed otherwise the restriction will be processed till a checkpoint is possible. The API is required to be implemented. Returns: ``True`` if the restriction represents a finite amount of work. Otherwise, returns ``False``. """ raise NotImplementedError class WatermarkEstimator(object): """A WatermarkEstimator which is used for estimating output_watermark based on the timestamp of output records or manual modifications. Please refer to ``watermark_estiamtors`` for commonly used watermark estimators. The base class provides common APIs that are called by the framework, which are also accessible inside a DoFn.process() body. Derived watermark estimator should implement all APIs listed below. Additional methods can be implemented and will be available when invoked within a DoFn. Internal state must not be updated asynchronously. """ def get_estimator_state(self): """Get current state of the WatermarkEstimator instance, which can be used to recreate the WatermarkEstimator when processing the restriction. See WatermarkEstimatorProvider.create_watermark_estimator. """ raise NotImplementedError(type(self)) def current_watermark(self) -> timestamp.Timestamp: """Return estimated output_watermark. This function must return monotonically increasing watermarks.""" raise NotImplementedError(type(self)) def observe_timestamp(self, timestamp: timestamp.Timestamp) -> None: """Update tracking watermark with latest output timestamp. Args: timestamp: the `timestamp.Timestamp` of current output element. This is called with the timestamp of every element output from the DoFn. """ raise NotImplementedError(type(self)) class RestrictionProgress(object): """Used to record the progress of a restriction.""" def __init__(self, **kwargs): # Only accept keyword arguments. self._fraction = kwargs.pop('fraction', None) self._completed = kwargs.pop('completed', None) self._remaining = kwargs.pop('remaining', None) assert not kwargs def __repr__(self): return 'RestrictionProgress(fraction=%s, completed=%s, remaining=%s)' % ( self._fraction, self._completed, self._remaining) @property def completed_work(self) -> float: if self._completed is not None: return self._completed elif self._remaining is not None and self._fraction is not None: return self._remaining * self._fraction / (1 - self._fraction) else: return self._fraction @property def remaining_work(self) -> float: if self._remaining is not None: return self._remaining elif self._completed is not None and self._fraction: return self._completed * (1 - self._fraction) / self._fraction else: return 1 - self._fraction @property def total_work(self) -> float: return self.completed_work + self.remaining_work @property def fraction_completed(self) -> float: if self._fraction is not None: return self._fraction else: return float(self._completed) / self.total_work @property def fraction_remaining(self) -> float: if self._fraction is not None: return 1 - self._fraction else: return float(self._remaining) / self.total_work def with_completed(self, completed: int) -> 'RestrictionProgress': return RestrictionProgress( fraction=self._fraction, remaining=self._remaining, completed=completed) class _SDFBoundedSourceRestriction(object): """ A restriction wraps SourceBundle and RangeTracker. """ def __init__(self, source_bundle, range_tracker=None): self._source_bundle = source_bundle self._range_tracker = range_tracker def __reduce__(self): # The instance of RangeTracker shouldn't be serialized. return (self.__class__, (self._source_bundle, )) def range_tracker(self): if not self._range_tracker: self._range_tracker = self._source_bundle.source.get_range_tracker( self._source_bundle.start_position, self._source_bundle.stop_position) return self._range_tracker def weight(self): return self._source_bundle.weight def source(self): return self._source_bundle.source def try_split(self, fraction_of_remainder): try: consumed_fraction = self.range_tracker().fraction_consumed() fraction = ( consumed_fraction + (1 - consumed_fraction) * fraction_of_remainder) position = self.range_tracker().position_at_fraction(fraction) # Need to stash current stop_pos before splitting since # range_tracker.split will update its stop_pos if splits # successfully. stop_pos = self._source_bundle.stop_position split_result = self.range_tracker().try_split(position) if split_result: split_pos, split_fraction = split_result primary_weight = self._source_bundle.weight * split_fraction residual_weight = self._source_bundle.weight - primary_weight # Update self to primary weight and end position. self._source_bundle = SourceBundle( primary_weight, self._source_bundle.source, self._source_bundle.start_position, split_pos) return ( self, _SDFBoundedSourceRestriction( SourceBundle( residual_weight, self._source_bundle.source, split_pos, stop_pos))) except Exception: # For any exceptions from underlying trySplit calls, the wrapper will # think that the source refuses to split at this point. In this case, # no split happens at the wrapper level. return None class _SDFBoundedSourceRestrictionTracker(RestrictionTracker): """An `iobase.RestrictionTracker` implementations for wrapping BoundedSource with SDF. The tracked restriction is a _SDFBoundedSourceRestriction, which wraps SourceBundle and RangeTracker. Delegated RangeTracker guarantees synchronization safety. """ def __init__(self, restriction): if not isinstance(restriction, _SDFBoundedSourceRestriction): raise ValueError( 'Initializing SDFBoundedSourceRestrictionTracker' ' requires a _SDFBoundedSourceRestriction. Got %s instead.' % restriction) self.restriction = restriction def current_progress(self) -> RestrictionProgress: return RestrictionProgress( fraction=self.restriction.range_tracker().fraction_consumed()) def current_restriction(self): self.restriction.range_tracker() return self.restriction def start_pos(self): return self.restriction.range_tracker().start_position() def stop_pos(self): return self.restriction.range_tracker().stop_position() def try_claim(self, position): return self.restriction.range_tracker().try_claim(position) def try_split(self, fraction_of_remainder): return self.restriction.try_split(fraction_of_remainder) def check_done(self): return self.restriction.range_tracker().fraction_consumed() >= 1.0 def is_bounded(self): return True class _SDFBoundedSourceWrapperRestrictionCoder(coders.Coder): def decode(self, value): return _SDFBoundedSourceRestriction(SourceBundle(*pickler.loads(value))) def encode(self, restriction): return pickler.dumps(( restriction._source_bundle.weight, restriction._source_bundle.source, restriction._source_bundle.start_position, restriction._source_bundle.stop_position)) class _SDFBoundedSourceRestrictionProvider(core.RestrictionProvider): """ A `RestrictionProvider` that is used by SDF for `BoundedSource`. This restriction provider initializes restriction based on input element that is expected to be of BoundedSource type. """ def __init__(self, desired_chunk_size=None, restriction_coder=None): self._desired_chunk_size = desired_chunk_size self._restriction_coder = ( restriction_coder or _SDFBoundedSourceWrapperRestrictionCoder()) def _check_source(self, src): if not isinstance(src, BoundedSource): raise RuntimeError( 'SDFBoundedSourceRestrictionProvider can only utilize BoundedSource') def initial_restriction(self, element_source: BoundedSource): self._check_source(element_source) range_tracker = element_source.get_range_tracker(None, None) return _SDFBoundedSourceRestriction( SourceBundle( None, element_source, range_tracker.start_position(), range_tracker.stop_position())) def create_tracker(self, restriction): return _SDFBoundedSourceRestrictionTracker(restriction) def split(self, element, restriction): if self._desired_chunk_size is None: try: estimated_size = restriction.source().estimate_size() except NotImplementedError: estimated_size = None self._desired_chunk_size = Read.get_desired_chunk_size(estimated_size) # Invoke source.split to get initial splitting results. source_bundles = restriction.source().split(self._desired_chunk_size) for source_bundle in source_bundles: yield _SDFBoundedSourceRestriction(source_bundle) def restriction_size(self, element, restriction): return restriction.weight() def restriction_coder(self): return self._restriction_coder class SDFBoundedSourceReader(PTransform): """A ``PTransform`` that uses SDF to read from each ``BoundedSource`` in a PCollection. NOTE: This transform can only be used with beam_fn_api enabled. """ def __init__(self, data_to_display=None): self._data_to_display = data_to_display or {} super().__init__() def _create_sdf_bounded_source_dofn(self): class SDFBoundedSourceDoFn(core.DoFn): def __init__(self, dd): self._dd = dd def display_data(self): return self._dd def process( self, unused_element, restriction_tracker=core.DoFn.RestrictionParam( _SDFBoundedSourceRestrictionProvider())): current_restriction = restriction_tracker.current_restriction() assert isinstance(current_restriction, _SDFBoundedSourceRestriction) return current_restriction.source().read( current_restriction.range_tracker()) return SDFBoundedSourceDoFn(self._data_to_display) def expand(self, pvalue): return pvalue | core.ParDo(self._create_sdf_bounded_source_dofn()) def get_windowing(self, unused_inputs): return core.Windowing(window.GlobalWindows()) def display_data(self): return self._data_to_display ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/iobase_it_test.py0000644000175100017510000000505315206113032021636 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import unittest import uuid import apache_beam as beam from apache_beam.io.textio import WriteToText from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.transforms.window import FixedWindows # End-to-End tests for iobase # Usage: # cd sdks/python # pip install build && python -m build --sdist # DataflowRunner: # python -m pytest -o log_cli=True -o log_level=Info \ # apache_beam/io/iobase_it_test.py::IOBaseITTest \ # --test-pipeline-options="--runner=TestDataflowRunner \ # --project=apache-beam-testing --region=us-central1 \ # --temp_location=gs://apache-beam-testing-temp/temp \ # --sdk_location=dist/apache_beam-2.65.0.dev0.tar.gz" class IOBaseITTest(unittest.TestCase): def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.runner_name = type(self.test_pipeline.runner).__name__ def test_unbounded_pcoll_without_global_window(self): # https://github.com/apache/beam/issues/25598 args = self.test_pipeline.get_full_options_as_args(streaming=True) topic = 'projects/pubsub-public-data/topics/taxirides-realtime' unique_id = str(uuid.uuid4()) output_file = f'gs://apache-beam-testing-integration-testing/iobase/test-{unique_id}' # pylint: disable=line-too-long p = beam.Pipeline(argv=args) # Read from Pub/Sub with fixed windowing lines = ( p | "ReadFromPubSub" >> beam.io.ReadFromPubSub(topic=topic) | "WindowInto" >> beam.WindowInto(FixedWindows(10))) # Write to text file _ = lines | 'WriteToText' >> WriteToText(output_file) result = p.run() result.wait_until_finish(duration=60 * 1000) if __name__ == "__main__": logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/iobase_test.py0000644000175100017510000002150315206113032021140 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for classes in iobase.py.""" # pytype: skip-file import unittest import mock import apache_beam as beam from apache_beam.io.concat_source import ConcatSource from apache_beam.io.concat_source_test import RangeSource from apache_beam.io import iobase from apache_beam.io import range_trackers from apache_beam.io.iobase import SourceBundle from apache_beam.options.pipeline_options import DebugOptions from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class SDFBoundedSourceRestrictionProviderTest(unittest.TestCase): def setUp(self): self.initial_range_start = 0 self.initial_range_stop = 4 self.initial_range_source = RangeSource( self.initial_range_start, self.initial_range_stop) self.sdf_restriction_provider = ( iobase._SDFBoundedSourceRestrictionProvider(desired_chunk_size=2)) def test_initial_restriction(self): element = self.initial_range_source restriction = (self.sdf_restriction_provider.initial_restriction(element)) self.assertTrue( isinstance(restriction, iobase._SDFBoundedSourceRestriction)) self.assertTrue(isinstance(restriction._source_bundle, SourceBundle)) self.assertEqual( self.initial_range_start, restriction._source_bundle.start_position) self.assertEqual( self.initial_range_stop, restriction._source_bundle.stop_position) self.assertTrue(isinstance(restriction._source_bundle.source, RangeSource)) self.assertEqual(restriction._range_tracker, None) def test_create_tracker(self): expected_start = 1 expected_stop = 3 source_bundle = SourceBundle( expected_stop - expected_start, RangeSource(1, 3), expected_start, expected_stop) restriction_tracker = ( self.sdf_restriction_provider.create_tracker( iobase._SDFBoundedSourceRestriction(source_bundle))) self.assertTrue( isinstance( restriction_tracker, iobase._SDFBoundedSourceRestrictionTracker)) self.assertEqual(expected_start, restriction_tracker.start_pos()) self.assertEqual(expected_stop, restriction_tracker.stop_pos()) def test_simple_source_split(self): element = self.initial_range_source restriction = (self.sdf_restriction_provider.initial_restriction(element)) expect_splits = [(0, 2), (2, 4)] split_bundles = list( self.sdf_restriction_provider.split(element, restriction)) self.assertTrue( all( isinstance(bundle._source_bundle, SourceBundle) for bundle in split_bundles)) splits = ([( bundle._source_bundle.start_position, bundle._source_bundle.stop_position) for bundle in split_bundles]) self.assertEqual(expect_splits, list(splits)) def test_concat_source_split(self): element = self.initial_range_source initial_concat_source = ConcatSource([self.initial_range_source]) sdf_concat_restriction_provider = ( iobase._SDFBoundedSourceRestrictionProvider(desired_chunk_size=2)) restriction = (self.sdf_restriction_provider.initial_restriction(element)) expect_splits = [(0, 2), (2, 4)] split_bundles = list( sdf_concat_restriction_provider.split( initial_concat_source, restriction)) self.assertTrue( all( isinstance(bundle._source_bundle, SourceBundle) for bundle in split_bundles)) splits = ([( bundle._source_bundle.start_position, bundle._source_bundle.stop_position) for bundle in split_bundles]) self.assertEqual(expect_splits, list(splits)) def test_restriction_size(self): element = self.initial_range_source restriction = (self.sdf_restriction_provider.initial_restriction(element)) split_1, split_2 = self.sdf_restriction_provider.split(element, restriction) split_1_size = self.sdf_restriction_provider.restriction_size( element, split_1) split_2_size = self.sdf_restriction_provider.restriction_size( element, split_2) self.assertEqual(2, split_1_size) self.assertEqual(2, split_2_size) class SDFBoundedSourceRestrictionTrackerTest(unittest.TestCase): def setUp(self): self.initial_start_pos = 0 self.initial_stop_pos = 4 source_bundle = SourceBundle( self.initial_stop_pos - self.initial_start_pos, RangeSource(self.initial_start_pos, self.initial_stop_pos), self.initial_start_pos, self.initial_stop_pos) self.sdf_restriction_tracker = ( iobase._SDFBoundedSourceRestrictionTracker( iobase._SDFBoundedSourceRestriction(source_bundle))) def test_current_restriction_before_split(self): current_restriction = (self.sdf_restriction_tracker.current_restriction()) self.assertEqual( self.initial_start_pos, current_restriction._source_bundle.start_position) self.assertEqual( self.initial_stop_pos, current_restriction._source_bundle.stop_position) self.assertEqual( self.initial_start_pos, current_restriction._range_tracker.start_position()) self.assertEqual( self.initial_stop_pos, current_restriction._range_tracker.stop_position()) def test_current_restriction_after_split(self): fraction_of_remainder = 0.5 self.sdf_restriction_tracker.try_claim(1) expected_restriction, _ = ( self.sdf_restriction_tracker.try_split(fraction_of_remainder)) current_restriction = self.sdf_restriction_tracker.current_restriction() self.assertEqual( expected_restriction._source_bundle, current_restriction._source_bundle) self.assertTrue(current_restriction._range_tracker) def test_try_split_at_remainder(self): fraction_of_remainder = 0.4 expected_primary = (0, 2, 2.0) expected_residual = (2, 4, 2.0) self.sdf_restriction_tracker.try_claim(0) actual_primary, actual_residual = ( self.sdf_restriction_tracker.try_split(fraction_of_remainder)) self.assertEqual( expected_primary, ( actual_primary._source_bundle.start_position, actual_primary._source_bundle.stop_position, actual_primary._source_bundle.weight)) self.assertEqual( expected_residual, ( actual_residual._source_bundle.start_position, actual_residual._source_bundle.stop_position, actual_residual._source_bundle.weight)) self.assertEqual( actual_primary._source_bundle.weight, self.sdf_restriction_tracker.current_restriction().weight()) def test_try_split_with_any_exception(self): source_bundle = SourceBundle( range_trackers.OffsetRangeTracker.OFFSET_INFINITY, RangeSource(0, range_trackers.OffsetRangeTracker.OFFSET_INFINITY), 0, range_trackers.OffsetRangeTracker.OFFSET_INFINITY) self.sdf_restriction_tracker = ( iobase._SDFBoundedSourceRestrictionTracker( iobase._SDFBoundedSourceRestriction(source_bundle))) self.sdf_restriction_tracker.try_claim(0) self.assertIsNone(self.sdf_restriction_tracker.try_split(0.5)) class UseSdfBoundedSourcesTests(unittest.TestCase): def _run_sdf_wrapper_pipeline(self, source, expected_values): with beam.Pipeline() as p: experiments = (p._options.view_as(DebugOptions).experiments or []) # Setup experiment option to enable using SDFBoundedSourceWrapper if 'beam_fn_api' not in experiments: # Required so mocking below doesn't mock Create used in assert_that. experiments.append('beam_fn_api') p._options.view_as(DebugOptions).experiments = experiments actual = p | beam.io.Read(source) assert_that(actual, equal_to(expected_values)) @mock.patch('apache_beam.io.iobase.SDFBoundedSourceReader.expand') def test_sdf_wrapper_overrides_read(self, sdf_wrapper_mock_expand): def _fake_wrapper_expand(pbegin): return pbegin | beam.Map(lambda x: 'fake') sdf_wrapper_mock_expand.side_effect = _fake_wrapper_expand self._run_sdf_wrapper_pipeline(RangeSource(0, 4), ['fake']) def test_sdf_wrap_range_source(self): self._run_sdf_wrapper_pipeline(RangeSource(0, 4), [0, 1, 2, 3]) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/jdbc.py0000644000175100017510000004023115206113032017540 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """PTransforms for supporting Jdbc in Python pipelines. These transforms are currently supported by Beam portable Flink, Spark, and Dataflow v2 runners. **Setup** Transforms provided in this module are cross-language transforms implemented in the Beam Java SDK. During the pipeline construction, Python SDK will connect to a Java expansion service to expand these transforms. To facilitate this, a small amount of setup is needed before using these transforms in a Beam Python pipeline. There are several ways to setup cross-language Jdbc transforms. * Option 1: use the default expansion service * Option 2: specify a custom expansion service See below for details regarding each of these options. *Option 1: Use the default expansion service* This is the recommended and easiest setup option for using Python Jdbc transforms. This option is only available for Beam 2.24.0 and later. This option requires following pre-requisites before running the Beam pipeline. * Install Java runtime in the computer from where the pipeline is constructed and make sure that 'java' command is available. In this option, Python SDK will either download (for released Beam version) or build (when running from a Beam Git clone) a expansion service jar and use that to expand transforms. Currently Jdbc transforms use the 'beam-sdks-java-io-expansion-service' jar for this purpose. The transforms in this file support an extra `classpath` argument, where one can specify any extra JARs to be included in the classpath for the expansion service. Users will need to specify this option to include the JDBC driver for the database that you're trying to use. **By default, a Postgres JDBC driver** is included (i.e. the Java package `"org.postgresql:postgresql:42.2.16"`). *Option 2: specify a custom expansion service* In this option, you startup your own expansion service and provide that as a parameter when using the transforms provided in this module. This option requires following pre-requisites before running the Beam pipeline. * Startup your own expansion service. * Update your pipeline to provide the expansion service address when initiating Jdbc transforms provided in this module. Flink Users can use the built-in Expansion Service of the Flink Runner's Job Server. If you start Flink's Job Server, the expansion service will be started on port 8097. For a different address, please set the expansion_service parameter. **More information** For more information regarding cross-language transforms see: - https://beam.apache.org/roadmap/portability/ For more information specific to Flink runner see: - https://beam.apache.org/documentation/runners/flink/ """ # pytype: skip-file import contextlib import typing import numpy as np from apache_beam.coders import RowCoder from apache_beam.transforms.external import BeamJarExpansionService from apache_beam.transforms.external import ExternalTransform from apache_beam.transforms.external import NamedTupleBasedPayloadBuilder from apache_beam.typehints.schemas import JdbcDateType # pylint: disable=unused-import from apache_beam.typehints.schemas import JdbcTimeType # pylint: disable=unused-import from apache_beam.typehints.schemas import LogicalType from apache_beam.typehints.schemas import MillisInstant from apache_beam.typehints.schemas import typing_to_runner_api __all__ = [ 'WriteToJdbc', 'ReadFromJdbc', ] def default_io_expansion_service(classpath=None): return BeamJarExpansionService( ':sdks:java:extensions:schemaio-expansion-service:shadowJar', classpath=classpath) JdbcConfigSchema = typing.NamedTuple( 'JdbcConfigSchema', [('location', str), ('config', bytes), ('dataSchema', typing.Optional[bytes])], ) Config = typing.NamedTuple( 'Config', [('driver_class_name', str), ('jdbc_url', str), ('username', str), ('password', str), ('connection_properties', typing.Optional[str]), ('connection_init_sqls', typing.Optional[list[str]]), ('read_query', typing.Optional[str]), ('write_statement', typing.Optional[str]), ('fetch_size', typing.Optional[np.int16]), ('disable_autocommit', typing.Optional[bool]), ('output_parallelization', typing.Optional[bool]), ('autosharding', typing.Optional[bool]), ('partition_column', typing.Optional[str]), ('partitions', typing.Optional[np.int16]), ('max_connections', typing.Optional[np.int16]), ('driver_jars', typing.Optional[str]), ('write_batch_size', typing.Optional[np.int64])], ) DEFAULT_JDBC_CLASSPATH = ['org.postgresql:postgresql:42.2.16'] class WriteToJdbc(ExternalTransform): """A PTransform which writes Rows to the specified database via JDBC. This transform receives Rows defined as NamedTuple type and registered in the coders registry, e.g.:: ExampleRow = typing.NamedTuple('ExampleRow', [('id', int), ('name', unicode)]) coders.registry.register_coder(ExampleRow, coders.RowCoder) with TestPipeline() as p: _ = ( p | beam.Create([ExampleRow(1, 'abc')]) .with_output_types(ExampleRow) | 'Write to jdbc' >> WriteToJdbc( table_name='jdbc_external_test_write' driver_class_name='org.postgresql.Driver', jdbc_url='jdbc:postgresql://localhost:5432/example', username='postgres', password='postgres', )) table_name is a required paramater, and by default, the write_statement is generated from it. The generated write_statement can be overridden by passing in a write_statment. Experimental; no backwards compatibility guarantees. """ URN = 'beam:transform:org.apache.beam:schemaio_jdbc_write:v1' def __init__( self, table_name, driver_class_name, jdbc_url, username, password, statement=None, connection_properties=None, connection_init_sqls=None, autosharding=False, max_connections=None, driver_jars=None, expansion_service=None, classpath=None, write_batch_size=None, ): """ Initializes a write operation to Jdbc. :param driver_class_name: name of the jdbc driver class :param jdbc_url: full jdbc url to the database. :param username: database username :param password: database password :param statement: sql statement to be executed :param connection_properties: properties of the jdbc connection passed as string with format [propertyName=property;]* :param connection_init_sqls: required only for MySql and MariaDB. passed as list of strings :param autosharding: enable automatic re-sharding of bundles to scale the number of shards with the number of workers. :param max_connections: sets the maximum total number of connections. use a negative value for no limit. :param driver_jars: comma separated paths for JDBC drivers. if not specified, the default classloader is used to load the driver jars. :param expansion_service: The address (host:port) of the ExpansionService. :param classpath: A list of JARs or Java packages to include in the classpath for the expansion service. This option is usually needed for `jdbc` to include extra JDBC driver packages. The packages can be in these three formats: (1) A local file, (2) A URL, (3) A gradle-style identifier of a Maven package (e.g. "org.postgresql:postgresql:42.3.1"). By default, this argument includes a Postgres SQL JDBC driver. :param write_batch_size: sets the maximum size in number of SQL statement for the batch. default is {@link JdbcIO.DEFAULT_BATCH_SIZE} """ classpath = classpath or DEFAULT_JDBC_CLASSPATH super().__init__( self.URN, NamedTupleBasedPayloadBuilder( JdbcConfigSchema( location=table_name, config=RowCoder( typing_to_runner_api(Config).row_type.schema).encode( Config( driver_class_name=driver_class_name, jdbc_url=jdbc_url, username=username, password=password, connection_properties=connection_properties, connection_init_sqls=connection_init_sqls, write_statement=statement, write_batch_size=write_batch_size, read_query=None, fetch_size=None, disable_autocommit=None, output_parallelization=None, autosharding=autosharding, max_connections=max_connections, driver_jars=driver_jars, partitions=None, partition_column=None)), dataSchema=None), ), expansion_service or default_io_expansion_service(classpath), ) @contextlib.contextmanager def enforce_millis_instant_for_timestamp(): old_registry = LogicalType._known_logical_types LogicalType._known_logical_types = old_registry.copy() try: LogicalType.register_logical_type(MillisInstant) LogicalType.register_logical_type(JdbcDateType) yield finally: LogicalType._known_logical_types = old_registry class ReadFromJdbc(ExternalTransform): """A PTransform which reads Rows from the specified database via JDBC. This transform delivers Rows defined as NamedTuple registered in the coders registry, e.g.:: ExampleRow = typing.NamedTuple('ExampleRow', [('id', int), ('name', unicode)]) coders.registry.register_coder(ExampleRow, coders.RowCoder) with TestPipeline() as p: result = ( p | 'Read from jdbc' >> ReadFromJdbc( table_name='jdbc_external_test_read' driver_class_name='org.postgresql.Driver', jdbc_url='jdbc:postgresql://localhost:5432/example', username='postgres', password='postgres', )) table_name is a required paramater, and by default, the read_query is generated from it. The generated read_query can be overridden by passing in a read_query. Experimental; no backwards compatibility guarantees. """ URN = 'beam:transform:org.apache.beam:schemaio_jdbc_read:v1' def __init__( self, table_name, driver_class_name, jdbc_url, username, password, query=None, disable_autocommit=None, output_parallelization=None, fetch_size=None, partition_column=None, partitions=None, connection_properties=None, connection_init_sqls=None, max_connections=None, driver_jars=None, expansion_service=None, classpath=None, schema=None): """ Initializes a read operation from Jdbc. :param driver_class_name: name of the jdbc driver class :param jdbc_url: full jdbc url to the database. :param username: database username :param password: database password :param query: sql query to be executed :param disable_autocommit: disable autocommit on read :param output_parallelization: is output parallelization on :param fetch_size: how many rows to fetch :param partition_column: enable partitioned reads by splitting on this column :param partitions: override the default number of splits when using partition_column :param connection_properties: properties of the jdbc connection passed as string with format [propertyName=property;]* :param connection_init_sqls: required only for MySql and MariaDB. passed as list of strings :param max_connections: sets the maximum total number of connections. use a negative value for no limit. :param driver_jars: comma separated paths for JDBC drivers. if not specified, the default classloader is used to load the driver jars. :param expansion_service: The address (host:port) of the ExpansionService. :param classpath: A list of JARs or Java packages to include in the classpath for the expansion service. This option is usually needed for `jdbc` to include extra JDBC driver packages. The packages can be in these three formats: (1) A local file, (2) A URL, (3) A gradle-style identifier of a Maven package (e.g. "org.postgresql:postgresql:42.3.1"). By default, this argument includes a Postgres SQL JDBC driver. :param schema: Optional custom schema for the returned rows. If provided, this should be a NamedTuple type that defines the structure of the output PCollection elements. This bypasses automatic schema inference during pipeline construction. """ # override new portable Date type with the current Jdbc type # TODO(https://github.com/apache/beam/issues/28359): # switch JdbcIO to return portable Date type LogicalType.register_logical_type(JdbcDateType) classpath = classpath or DEFAULT_JDBC_CLASSPATH dataSchema = None if schema is not None: with enforce_millis_instant_for_timestamp(): # Convert Python schema to Beam Schema proto schema_proto = typing_to_runner_api(schema).row_type.schema # Serialize the proto to bytes for transmission dataSchema = schema_proto.SerializeToString() super().__init__( self.URN, NamedTupleBasedPayloadBuilder( JdbcConfigSchema( location=table_name, config=RowCoder( typing_to_runner_api(Config).row_type.schema).encode( Config( driver_class_name=driver_class_name, jdbc_url=jdbc_url, username=username, password=password, connection_properties=connection_properties, connection_init_sqls=connection_init_sqls, write_statement=None, write_batch_size=None, read_query=query, fetch_size=fetch_size, disable_autocommit=disable_autocommit, output_parallelization=output_parallelization, autosharding=None, max_connections=max_connections, driver_jars=driver_jars, partition_column=partition_column, partitions=partitions)), dataSchema=dataSchema), ), expansion_service or default_io_expansion_service(classpath), ) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/kafka.py0000644000175100017510000003437315206113032017725 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unbounded source and sink transforms for `Kafka `_. These transforms are currently supported by Beam portable runners (for example, portable Flink and Spark) as well as Dataflow runner. **Setup** Transforms provided in this module are cross-language transforms implemented in the Beam Java SDK. During the pipeline construction, Python SDK will connect to a Java expansion service to expand these transforms. To facilitate this, a small amount of setup is needed before using these transforms in a Beam Python pipeline. There are several ways to setup cross-language Kafka transforms. * Option 1: use the default expansion service * Option 2: specify a custom expansion service See below for details regarding each of these options. *Option 1: Use the default expansion service* This is the recommended and easiest setup option for using Python Kafka transforms. This option is only available for Beam 2.22.0 and later. This option requires following pre-requisites before running the Beam pipeline. * Install Java runtime in the computer from where the pipeline is constructed and make sure that 'java' command is available or set JAVA_HOME environment variable. In this option, Python SDK will either download (for released Beam version) or build (when running from a Beam Git clone) an expansion service jar and use that to expand transforms. Currently Kafka transforms use the 'beam-sdks-java-io-expansion-service' jar for this purpose. Note that the KafkaIO read transform can be compiled in two modes * `ReadFromKafkaViaUnbounded` (legacy) * `ReadFromKafkaViaSDF` (default) To use the legacy mode, the `use_deprecated_read` flag should be specified within the IO expansion service. For example, kafka.default_io_expansion_service( append_args=["--experiments=use_deprecated_read"] ) To use the legacy mode using Dataflow runner, the `use_unbounded_sdf_wrapper` flag should be specified within the IO expansion service. For example, kafka.default_io_expansion_service( append_args=["--experiments=use_unbounded_sdf_wrapper"] ) *Option 2: specify a custom expansion service* In this option, you startup your own expansion service and provide that as a parameter when using the transforms provided in this module. This option requires following pre-requisites before running the Beam pipeline. * Startup your own expansion service. * Update your pipeline to provide the expansion service address when initiating Kafka transforms provided in this module. Flink Users can use the built-in Expansion Service of the Flink Runner's Job Server. If you start Flink's Job Server, the expansion service will be started on port 8097. For a different address, please set the expansion_service parameter. **More information** For more information regarding cross-language transforms see: - https://beam.apache.org/roadmap/portability/ For more information specific to Flink runner see: - https://beam.apache.org/documentation/runners/flink/ """ # pytype: skip-file import collections import typing import numpy as np from apache_beam.transforms.external import BeamJarExpansionService from apache_beam.transforms.external import ExternalTransform from apache_beam.transforms.external import NamedTupleBasedPayloadBuilder ReadFromKafkaSchema = typing.NamedTuple( 'ReadFromKafkaSchema', [('consumer_config', typing.Mapping[str, str]), ('topics', list[str]), ('key_deserializer', str), ('value_deserializer', str), ('start_read_time', typing.Optional[int]), ('max_num_records', typing.Optional[int]), ('max_read_time', typing.Optional[int]), ('commit_offset_in_finalize', bool), ('timestamp_policy', str), ('consumer_polling_timeout', typing.Optional[int]), ('redistribute', typing.Optional[bool]), ('redistribute_num_keys', typing.Optional[np.int32]), ('allow_duplicates', typing.Optional[bool]), ('dynamic_read_poll_interval_seconds', typing.Optional[int]), ('consumer_factory_fn_class', typing.Optional[str]), ( 'consumer_factory_fn_params', typing.Optional[collections.abc.Mapping[str, str]])]) def default_io_expansion_service(append_args=None): return BeamJarExpansionService( 'sdks:java:io:expansion-service:shadowJar', append_args=append_args) class ReadFromKafka(ExternalTransform): """ An external PTransform which reads from Kafka and returns a KV pair for each item in the specified Kafka topics. If no Kafka Deserializer for key/value is provided, then the data will be returned as a raw byte array. Experimental; no backwards compatibility guarantees. """ # Returns the key/value data as raw byte arrays byte_array_deserializer = ( 'org.apache.kafka.common.serialization.ByteArrayDeserializer') processing_time_policy = 'ProcessingTime' create_time_policy = 'CreateTime' log_append_time = 'LogAppendTime' URN_WITH_METADATA = ( 'beam:transform:org.apache.beam:kafka_read_with_metadata:v1') URN_WITHOUT_METADATA = ( 'beam:transform:org.apache.beam:kafka_read_without_metadata:v1') def __init__( self, consumer_config, topics, key_deserializer=byte_array_deserializer, value_deserializer=byte_array_deserializer, start_read_time=None, max_num_records=None, max_read_time=None, commit_offset_in_finalize=False, timestamp_policy=processing_time_policy, consumer_polling_timeout=2, with_metadata=False, expansion_service=None, redistribute=False, redistribute_num_keys=np.int32(0), allow_duplicates=False, dynamic_read_poll_interval_seconds: typing.Optional[int] = None, consumer_factory_fn_class: typing.Optional[str] = None, consumer_factory_fn_params: typing.Optional[ collections.abc.Mapping] = None): """ Initializes a read operation from Kafka. :param consumer_config: A dictionary containing the consumer configuration. :param topics: A list of topic strings. :param key_deserializer: A fully-qualified Java class name of a Kafka Deserializer for the topic's key, e.g. 'org.apache.kafka.common.serialization.LongDeserializer'. Default: 'org.apache.kafka.common.serialization.ByteArrayDeserializer'. :param value_deserializer: A fully-qualified Java class name of a Kafka Deserializer for the topic's value, e.g. 'org.apache.kafka.common.serialization.LongDeserializer'. Default: 'org.apache.kafka.common.serialization.ByteArrayDeserializer'. :param start_read_time: Use timestamp to set up start offset in milliseconds epoch. :param max_num_records: Maximum amount of records to be read. Mainly used for tests and demo applications. :param max_read_time: Maximum amount of time in seconds the transform executes. Mainly used for tests and demo applications. :param commit_offset_in_finalize: Whether to commit offsets when finalizing. :param timestamp_policy: The built-in timestamp policy which is used for extracting timestamp from KafkaRecord. :param consumer_polling_timeout: Kafka client polling request timeout time in seconds. A lower timeout optimizes for latency. Increase the timeout if the consumer is not fetching any records. Default is 2 seconds. :param with_metadata: whether the returned PCollection should contain Kafka related metadata or not. If False (default), elements of the returned PCollection will be of type 'bytes' if True, elements of the returned PCollection will be of the type 'Row'. Note that, currently this only works when using default key and value deserializers where Java Kafka Reader reads keys and values as 'byte[]'. :param expansion_service: The address (host:port) of the ExpansionService. :param redistribute: whether a Redistribute transform should be applied immediately after the read. :param redistribute_num_keys: Configures how many keys the Redistribute spreads the data across. :param allow_duplicates: whether the Redistribute transform allows for duplicates (this serves solely as a hint to the underlying runner). :param dynamic_read_poll_interval_seconds: The interval in seconds at which to check for new partitions. If not None, dynamic partition discovery is enabled. :param consumer_factory_fn_class: A fully qualified classpath to an existing provided consumerFactoryFn. If not None, this will construct Kafka consumers with a custom configuration. :param consumer_factory_fn_params: A map which specifies the parameters for the provided consumer_factory_fn_class. If not None, the values in this map will be used when constructing the consumer_factory_fn_class object. This cannot be null if the consumer_factory_fn_class is not null. """ if timestamp_policy not in [ReadFromKafka.processing_time_policy, ReadFromKafka.create_time_policy, ReadFromKafka.log_append_time]: raise ValueError( 'timestamp_policy should be one of ' '[ProcessingTime, CreateTime, LogAppendTime]') super().__init__( self.URN_WITH_METADATA if with_metadata else self.URN_WITHOUT_METADATA, NamedTupleBasedPayloadBuilder( ReadFromKafkaSchema( consumer_config=consumer_config, topics=topics, key_deserializer=key_deserializer, value_deserializer=value_deserializer, max_num_records=max_num_records, max_read_time=max_read_time, start_read_time=start_read_time, commit_offset_in_finalize=commit_offset_in_finalize, timestamp_policy=timestamp_policy, consumer_polling_timeout=consumer_polling_timeout, redistribute=redistribute, redistribute_num_keys=redistribute_num_keys, allow_duplicates=allow_duplicates, dynamic_read_poll_interval_seconds= dynamic_read_poll_interval_seconds, consumer_factory_fn_class=consumer_factory_fn_class, consumer_factory_fn_params=consumer_factory_fn_params)), expansion_service or default_io_expansion_service()) WriteToKafkaSchema = typing.NamedTuple( 'WriteToKafkaSchema', [ ('producer_config', typing.Mapping[str, str]), ('topic', str), ('key_serializer', str), ('value_serializer', str), ]) class WriteToKafka(ExternalTransform): """ An external PTransform which writes KV data to a specified Kafka topic. If no Kafka Serializer for key/value is provided, then key/value are assumed to be byte arrays. Experimental; no backwards compatibility guarantees. When with_headers=True, the input PCollection elements must be beam.Row objects with the following schema: - key: bytes (required) - The key of the record. - value: bytes (required) - The value of the record. - headers: List[Row(key=str, value=bytes)] (optional) - Record headers. - topic: str (optional) - Per-record topic override. - partition: int (optional) - Per-record partition. - timestamp: int (optional) - Per-record timestamp in milliseconds. """ # Default serializer which passes raw bytes to Kafka byte_array_serializer = ( 'org.apache.kafka.common.serialization.ByteArraySerializer') URN = 'beam:transform:org.apache.beam:kafka_write:v1' URN_WITH_HEADERS = ( 'beam:transform:org.apache.beam:kafka_write_with_headers:v1') def __init__( self, producer_config, topic, key_serializer=byte_array_serializer, value_serializer=byte_array_serializer, with_headers=False, expansion_service=None): """ Initializes a write operation to Kafka. :param producer_config: A dictionary containing the producer configuration. :param topic: A Kafka topic name. :param key_serializer: A fully-qualified Java class name of a Kafka Serializer for the topic's key, e.g. 'org.apache.kafka.common.serialization.LongSerializer'. Default: 'org.apache.kafka.common.serialization.ByteArraySerializer'. :param value_serializer: A fully-qualified Java class name of a Kafka Serializer for the topic's value, e.g. 'org.apache.kafka.common.serialization.LongSerializer'. Default: 'org.apache.kafka.common.serialization.ByteArraySerializer'. :param with_headers: If True, input elements must be beam.Row objects containing 'key', 'value', and optional 'headers' fields. Only ByteArraySerializer is supported when with_headers=True. :param expansion_service: The address (host:port) of the ExpansionService. """ if with_headers and (key_serializer != self.byte_array_serializer or value_serializer != self.byte_array_serializer): raise ValueError( 'WriteToKafka(with_headers=True) only supports ' 'ByteArraySerializer for key and value.') urn = self.URN_WITH_HEADERS if with_headers else self.URN super().__init__( urn, NamedTupleBasedPayloadBuilder( WriteToKafkaSchema( producer_config=producer_config, topic=topic, key_serializer=key_serializer, value_serializer=value_serializer, )), expansion_service or default_io_expansion_service()) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/kinesis.py0000644000175100017510000003204415206113032020306 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """PTransforms for supporting Kinesis streaming in Python pipelines. These transforms are currently supported by Beam Flink and Spark portable runners. **Setup** Transforms provided in this module are cross-language transforms implemented in the Beam Java SDK. During the pipeline construction, Python SDK will connect to a Java expansion service to expand these transforms. To facilitate this, a small amount of setup is needed before using these transforms in a Beam Python pipeline. There are several ways to setup cross-language Kinesis transforms. * Option 1: use the default expansion service * Option 2: specify a custom expansion service See below for details regarding each of these options. *Option 1: Use the default expansion service* This is the recommended and easiest setup option for using Python Kinesis transforms. This option is only available for Beam 2.25.0 and later. This option requires following pre-requisites before running the Beam pipeline. * Install Java runtime in the computer from where the pipeline is constructed and make sure that 'java' command is available or set JAVA_HOME environment variable. In this option, Python SDK will either download (for released Beam version) or build (when running from a Beam Git clone) a expansion service jar and use that to expand transforms. Currently Kinesis transforms use the 'beam-sdks-java-io-amazon-web-services2-expansion-service' jar for this purpose. *Option 2: specify a custom expansion service* In this option, you startup your own expansion service and provide that as a parameter when using the transforms provided in this module. This option requires following pre-requisites before running the Beam pipeline. * Startup your own expansion service. * Update your pipeline to provide the expansion service address when initiating Kinesis transforms provided in this module. Flink Users can use the built-in Expansion Service of the Flink Runner's Job Server. If you start Flink's Job Server, the expansion service will be started on port 8097. For a different address, please set the expansion_service parameter. **More information** For more information regarding cross-language transforms see: - https://beam.apache.org/roadmap/portability/ For more information specific to Flink runner see: - https://beam.apache.org/documentation/runners/flink/ """ # pytype: skip-file import logging import time from typing import NamedTuple from typing import Optional from apache_beam import BeamJarExpansionService from apache_beam import ExternalTransform from apache_beam import NamedTupleBasedPayloadBuilder __all__ = [ 'WriteToKinesis', 'ReadDataFromKinesis', 'InitialPositionInStream', 'WatermarkPolicy', ] def default_io_expansion_service(): return BeamJarExpansionService( 'sdks:java:io:amazon-web-services2:expansion-service:shadowJar') WriteToKinesisSchema = NamedTuple( 'WriteToKinesisSchema', [ ('stream_name', str), ('aws_access_key', str), ('aws_secret_key', str), ('region', str), ('partition_key', str), ('service_endpoint', Optional[str]), ('verify_certificate', Optional[bool]), ('aggregation_enabled', Optional[bool]), ('aggregation_max_bytes', Optional[int]), ('aggregation_max_buffered_time', Optional[int]), ('aggregation_shard_refresh_interval', Optional[int]), ], ) class WriteToKinesis(ExternalTransform): """ An external PTransform which writes byte array stream to Amazon Kinesis. Experimental; no backwards compatibility guarantees. """ URN = 'beam:transform:org.apache.beam:kinesis_write:v2' def __init__( self, stream_name, aws_access_key, aws_secret_key, region, partition_key, service_endpoint=None, verify_certificate=None, producer_properties=None, expansion_service=None, aggregation_enabled=None, aggregation_max_bytes=51200, aggregation_max_buffered_time=100, aggregation_shard_refresh_interval=2, ): """ Initializes a write operation to Kinesis. :param stream_name: Kinesis stream name. :param aws_access_key: Kinesis access key. :param aws_secret_key: Kinesis access key secret. :param region: AWS region. Example: 'us-east-1'. :param service_endpoint: Kinesis service endpoint :param verify_certificate: Enable or disable certificate verification. Never set to False on production. True by default. :param partition_key: Specify default partition key. :param producer_properties: (Deprecated) This option no longer is available since the AWS IOs upgraded to v2. Trying to set it will lead to an error. For more info, see https://github.com/apache/beam/issues/33430. :param expansion_service: The address (host:port) of the ExpansionService. :param aggregation_enabled: Enable or disable aggregation. :param aggregation_max_bytes: Maximum number of bytes to buffer before sending a batch of records. Defaults to 51200. :param aggregation_max_buffered_time: Maximum time(millisecond) to buffer records before sending a batch of records. Defaults to 100. :param aggregation_shard_refresh_interval: Interval in minutes to refresh the shard map. Defaults to 2. """ if producer_properties is not None: raise ValueError( 'producer_properties is no longer supported and will be removed ' + 'in a future release.') super().__init__( self.URN, NamedTupleBasedPayloadBuilder( WriteToKinesisSchema( stream_name=stream_name, aws_access_key=aws_access_key, aws_secret_key=aws_secret_key, region=region, partition_key=partition_key, service_endpoint=service_endpoint, verify_certificate=verify_certificate, aggregation_enabled=aggregation_enabled, aggregation_max_bytes=aggregation_max_bytes, aggregation_max_buffered_time=aggregation_max_buffered_time, aggregation_shard_refresh_interval= aggregation_shard_refresh_interval, )), expansion_service or default_io_expansion_service(), ) ReadFromKinesisSchema = NamedTuple( 'ReadFromKinesisSchema', [ ('stream_name', str), ('aws_access_key', str), ('aws_secret_key', str), ('region', str), ('service_endpoint', Optional[str]), ('verify_certificate', Optional[bool]), ('max_num_records', Optional[int]), ('max_read_time', Optional[int]), ('initial_position_in_stream', Optional[str]), ('initial_timestamp_in_stream', Optional[int]), ('request_records_limit', Optional[int]), ('up_to_date_threshold', Optional[int]), ('max_capacity_per_shard', Optional[int]), ('watermark_policy', Optional[str]), ('watermark_idle_duration_threshold', Optional[int]), ('rate_limit', Optional[int]), ], ) class ReadDataFromKinesis(ExternalTransform): """ An external PTransform which reads byte array stream from Amazon Kinesis. Experimental; no backwards compatibility guarantees. """ URN = 'beam:transform:org.apache.beam:kinesis_read_data:v2' def __init__( self, stream_name, aws_access_key, aws_secret_key, region, service_endpoint=None, verify_certificate=None, max_num_records=None, max_read_time=None, initial_position_in_stream=None, initial_timestamp_in_stream=None, request_records_limit=None, up_to_date_threshold=None, max_capacity_per_shard=None, watermark_policy=None, watermark_idle_duration_threshold=None, rate_limit=None, expansion_service=None, ): """ Initializes a read operation from Kinesis. :param stream_name: Kinesis stream name. :param aws_access_key: Kinesis access key. :param aws_secret_key: Kinesis access key secret. :param region: AWS region. Example: 'us-east-1'. :param service_endpoint: Kinesis service endpoint :param verify_certificate: Enable or disable certificate verification. Never set to False on production. True by default. :param max_num_records: Specifies to read at most a given number of records. Must be greater than 0. :param max_read_time: Specifies to read records during x milliseconds. :param initial_timestamp_in_stream: Specify reading beginning at the given timestamp in milliseconds. Must be in the past. :param initial_position_in_stream: Specify reading from some initial position in stream. Possible values: LATEST - Start after the most recent data record (fetch new data). TRIM_HORIZON - Start from the oldest available data record. AT_TIMESTAMP - Start from the record at or after the specified server-side timestamp. :param request_records_limit: Specifies the maximum number of records in GetRecordsResult returned by GetRecords call which is limited by 10K records. If should be adjusted according to average size of data record to prevent shard overloading. More at: docs.aws.amazon.com/kinesis/latest/APIReference/API_GetRecords.html :param up_to_date_threshold: Specifies how late in milliseconds records consumed by this source can be to still be considered on time. Defaults to zero. :param max_capacity_per_shard: Specifies the maximum number of messages per one shard. Defaults to 10'000. :param watermark_policy: Specifies the watermark policy. Possible values: PROCESSING_TYPE, ARRIVAL_TIME. Defaults to ARRIVAL_TIME. :param watermark_idle_duration_threshold: Use only when watermark policy is ARRIVAL_TIME. Denotes the duration for which the watermark can be idle. Passed in milliseconds. :param rate_limit: Sets fixed rate policy for given milliseconds value. By default there is no rate limit. :param expansion_service: The address (host:port) of the ExpansionService. """ WatermarkPolicy.validate_param(watermark_policy) InitialPositionInStream.validate_param(initial_position_in_stream) if watermark_idle_duration_threshold: assert WatermarkPolicy.ARRIVAL_TIME == watermark_policy if request_records_limit: assert 0 < request_records_limit <= 10000 initial_timestamp_in_stream = int( initial_timestamp_in_stream) if initial_timestamp_in_stream else None if initial_timestamp_in_stream and initial_timestamp_in_stream < time.time( ): logging.warning('Provided timestamp emplaced not in the past.') super().__init__( self.URN, NamedTupleBasedPayloadBuilder( ReadFromKinesisSchema( stream_name=stream_name, aws_access_key=aws_access_key, aws_secret_key=aws_secret_key, region=region, service_endpoint=service_endpoint, verify_certificate=verify_certificate, max_num_records=max_num_records, max_read_time=max_read_time, initial_position_in_stream=initial_position_in_stream, initial_timestamp_in_stream=initial_timestamp_in_stream, request_records_limit=request_records_limit, up_to_date_threshold=up_to_date_threshold, max_capacity_per_shard=max_capacity_per_shard, watermark_policy=watermark_policy, watermark_idle_duration_threshold= watermark_idle_duration_threshold, rate_limit=rate_limit, )), expansion_service or default_io_expansion_service(), ) class InitialPositionInStream: LATEST = 'LATEST' TRIM_HORIZON = 'TRIM_HORIZON' AT_TIMESTAMP = 'AT_TIMESTAMP' @staticmethod def validate_param(param): if param and not hasattr(InitialPositionInStream, param): raise RuntimeError('Invalid initial position in stream: {}'.format(param)) class WatermarkPolicy: PROCESSING_TYPE = 'PROCESSING_TYPE' ARRIVAL_TIME = 'ARRIVAL_TIME' @staticmethod def validate_param(param): if param and not hasattr(WatermarkPolicy, param): raise RuntimeError('Invalid watermark policy: {}'.format(param)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/localfilesystem.py0000644000175100017510000002605215206113032022042 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Local File system implementation for accessing files on disk.""" # pytype: skip-file import io import os import shutil from typing import BinaryIO # pylint: disable=unused-import from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystem import CompressedFile from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystem import FileMetadata from apache_beam.io.filesystem import FileSystem __all__ = ['LocalFileSystem'] class LocalFileSystem(FileSystem): """A Local ``FileSystem`` implementation for accessing files on disk. """ @classmethod def scheme(cls): """URI scheme for the FileSystem """ return None def join(self, basepath, *paths): """Join two or more pathname components for the filesystem Args: basepath: string path of the first component of the path paths: path components to be added Returns: full path after combining all the passed components """ return os.path.join(basepath, *paths) def split(self, path): """Splits the given path into two parts. Splits the path into a pair (head, tail) such that tail contains the last component of the path and head contains everything up to that. Args: path: path as a string Returns: a pair of path components as strings. """ return os.path.split(os.path.abspath(path)) def mkdirs(self, path): """Recursively create directories for the provided path. Args: path: string path of the directory structure that should be created Raises: IOError: if leaf directory already exists. """ try: os.makedirs(path) except OSError as err: raise IOError(err) def has_dirs(self): """Whether this FileSystem supports directories.""" return True def _url_dirname(self, url_or_path): """Pass through to os.path.dirname. This version uses os.path instead of posixpath to be compatible with the host OS. Args: url_or_path: A string in the form of /some/path. """ return os.path.dirname(url_or_path) def _list(self, dir_or_prefix): """List files in a location. Listing is non-recursive, for filesystems that support directories. Args: dir_or_prefix: (string) A directory or location prefix (for filesystems that don't have directories). Returns: Generator of ``FileMetadata`` objects. Raises: ``BeamIOError``: if listing fails, but not if no files were found. """ if not self.exists(dir_or_prefix): return def list_files(root): for dirpath, _, files in os.walk(root): for filename in files: yield self.join(dirpath, filename) try: for f in list_files(dir_or_prefix): try: yield FileMetadata(f, os.path.getsize(f), os.path.getmtime(f)) except OSError: # Files may disappear, such as when listing /tmp. pass except Exception as e: # pylint: disable=broad-except raise BeamIOError("List operation failed", {dir_or_prefix: e}) def _path_open( self, path, mode, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO): """Helper functions to open a file in the provided mode. """ compression_type = FileSystem._get_compression_type(path, compression_type) raw_file = io.open(path, mode) if compression_type == CompressionTypes.UNCOMPRESSED: return raw_file else: return CompressedFile(raw_file, compression_type=compression_type) def create( self, path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO) -> BinaryIO: """Returns a write channel for the given file path. Args: path: string path of the file object to be written to the system mime_type: MIME type to specify the type of content in the file object compression_type: Type of compression to be used for this object Returns: file handle with a close function for the user to use """ dirname = os.path.dirname(path) if dirname: os.makedirs(os.path.dirname(path), exist_ok=True) return self._path_open(path, 'wb', mime_type, compression_type) def open( self, path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO) -> BinaryIO: """Returns a read channel for the given file path. Args: path: string path of the file object to be written to the system mime_type: MIME type to specify the type of content in the file object compression_type: Type of compression to be used for this object Returns: file handle with a close function for the user to use """ return self._path_open(path, 'rb', mime_type, compression_type) def copy(self, source_file_names, destination_file_names): """Recursively copy the file tree from the source to the destination Args: source_file_names: list of source file objects that needs to be copied destination_file_names: list of destination of the new object Raises: ``BeamIOError``: if any of the copy operations fail """ err_msg = ( "source_file_names and destination_file_names should " "be equal in length") assert len(source_file_names) == len(destination_file_names), err_msg def _copy_path(source, destination): """Recursively copy the file tree from the source to the destination """ try: if os.path.exists(destination): if os.path.isdir(destination): shutil.rmtree(destination) else: os.remove(destination) if os.path.isdir(source): shutil.copytree(source, destination) else: shutil.copy2(source, destination) except OSError as err: raise IOError(err) exceptions = {} for source, destination in zip(source_file_names, destination_file_names): try: _copy_path(source, destination) except Exception as e: # pylint: disable=broad-except exceptions[(source, destination)] = e if exceptions: raise BeamIOError("Copy operation failed", exceptions) def rename(self, source_file_names, destination_file_names): """Rename the files at the source list to the destination list. Source and destination lists should be of the same size. Args: source_file_names: List of file paths that need to be moved destination_file_names: List of destination_file_names for the files Raises: ``BeamIOError``: if any of the rename operations fail """ err_msg = ( "source_file_names and destination_file_names should " "be equal in length") assert len(source_file_names) == len(destination_file_names), err_msg def _rename_file(source, destination): """Rename a single file object""" try: os.rename(source, destination) except OSError as err: raise IOError(err) exceptions = {} for source, destination in zip(source_file_names, destination_file_names): try: _rename_file(source, destination) except Exception as e: # pylint: disable=broad-except exceptions[(source, destination)] = e if exceptions: raise BeamIOError("Rename operation failed", exceptions) def exists(self, path): """Check if the provided path exists on the FileSystem. Args: path: string path that needs to be checked. Returns: boolean flag indicating if path exists """ return os.path.exists(path) def size(self, path): """Get size of path on the FileSystem. Args: path: string path in question. Returns: int size of path according to the FileSystem. Raises: ``BeamIOError``: if path doesn't exist. """ try: return os.path.getsize(path) except Exception as e: # pylint: disable=broad-except raise BeamIOError("Size operation failed", {path: e}) def last_updated(self, path): """Get UNIX Epoch time in seconds on the FileSystem. Args: path: string path of file. Returns: float UNIX Epoch time Raises: ``BeamIOError``: if path doesn't exist. """ if not self.exists(path): raise BeamIOError('Path does not exist: %s' % path) return os.path.getmtime(path) def checksum(self, path): """Fetch checksum metadata of a file on the :class:`~apache_beam.io.filesystem.FileSystem`. Args: path: string path of a file. Returns: string containing file size. Raises: ``BeamIOError``: if path isn't a file or doesn't exist. """ if not self.exists(path): raise BeamIOError('Path does not exist: %s' % path) return str(os.path.getsize(path)) def metadata(self, path): """Fetch metadata fields of a file on the FileSystem. Args: path: string path of a file. Returns: :class:`~apache_beam.io.filesystem.FileMetadata`. Raises: ``BeamIOError``: if path isn't a file or doesn't exist. """ if not self.exists(path): raise BeamIOError('Path does not exist: %s' % path) return FileMetadata(path, os.path.getsize(path), os.path.getmtime(path)) def delete(self, paths): """Deletes files or directories at the provided paths. Directories will be deleted recursively. Args: paths: list of paths that give the file objects to be deleted Raises: ``BeamIOError``: if any of the delete operations fail """ def _delete_path(path): """Recursively delete the file or directory at the provided path. """ try: if os.path.isdir(path): shutil.rmtree(path) else: os.remove(path) except OSError as err: raise IOError(err) exceptions = {} def try_delete(path): try: _delete_path(path) except Exception as e: # pylint: disable=broad-except exceptions[path] = e for match_result in self.match(paths): metadata_list = match_result.metadata_list if not metadata_list: exceptions[match_result.pattern] = \ IOError('No files found to delete under: %s' % match_result.pattern) for metadata in match_result.metadata_list: try_delete(metadata.path) if exceptions: raise BeamIOError("Delete operation failed", exceptions) def report_lineage(self, path, lineage): lineage.add('filesystem', 'localhost', path, last_segment_sep='/') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/localfilesystem_test.py0000644000175100017510000003704715206113032023107 0ustar00runnerrunner# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for LocalFileSystem.""" import contextlib import filecmp import logging import os import shutil import tempfile import unittest import mock from parameterized import param from parameterized import parameterized from apache_beam.io import localfilesystem from apache_beam.io.filesystem import BeamIOError from apache_beam.options.pipeline_options import PipelineOptions # pytype: skip-file def _gen_fake_join(separator): """Returns a callable that joins paths with the given separator.""" def _join(first_path, *paths): return separator.join((first_path.rstrip(separator), ) + paths) return _join def _gen_fake_split(separator): """Returns a callable that splits a with the given separator.""" def _split(path): sep_index = path.rfind(separator) if sep_index >= 0: return (path[:sep_index], path[sep_index + 1:]) else: return (path, '') return _split class LocalFileSystemTest(unittest.TestCase): def setUp(self): self.tmpdir = tempfile.mkdtemp() pipeline_options = PipelineOptions() self.fs = localfilesystem.LocalFileSystem(pipeline_options) def tearDown(self): shutil.rmtree(self.tmpdir) @contextlib.contextmanager def tmpdir_as_cwd(self): """Context manager that sets the current working directory to a temp dir.""" old_cwd = os.getcwd() os.chdir(self.tmpdir) try: yield finally: os.chdir(old_cwd) def test_scheme(self): self.assertIsNone(self.fs.scheme()) self.assertIsNone(localfilesystem.LocalFileSystem.scheme()) def test_create_cwd_file(self): with self.tmpdir_as_cwd(): with self.fs.create("blah.txt") as f: f.write(b"blah") with open("blah.txt", "rb") as f: assert f.read() == b"blah" @mock.patch('apache_beam.io.localfilesystem.os') def test_unix_path_join(self, *unused_mocks): # Test joining of Unix paths. localfilesystem.os.path.join.side_effect = _gen_fake_join('/') self.assertEqual( '/tmp/path/to/file', self.fs.join('/tmp/path', 'to', 'file')) self.assertEqual('/tmp/path/to/file', self.fs.join('/tmp/path', 'to/file')) @mock.patch('apache_beam.io.localfilesystem.os') def test_windows_path_join(self, *unused_mocks): # Test joining of Windows paths. localfilesystem.os.path.join.side_effect = _gen_fake_join('\\') self.assertEqual( r'C:\tmp\path\to\file', self.fs.join(r'C:\tmp\path', 'to', 'file')) self.assertEqual( r'C:\tmp\path\to\file', self.fs.join(r'C:\tmp\path', r'to\file')) @mock.patch('apache_beam.io.localfilesystem.os') def test_unix_path_split(self, os_mock): os_mock.path.abspath.side_effect = lambda a: a os_mock.path.split.side_effect = _gen_fake_split('/') self.assertEqual(('/tmp/path/to', 'file'), self.fs.split('/tmp/path/to/file')) # Actual os.path.split will split following to '/' and 'tmp' when run in # Unix. self.assertEqual(('', 'tmp'), self.fs.split('/tmp')) @mock.patch('apache_beam.io.localfilesystem.os') def test_windows_path_split(self, os_mock): os_mock.path.abspath = lambda a: a os_mock.path.split.side_effect = _gen_fake_split('\\') self.assertEqual((r'C:\tmp\path\to', 'file'), self.fs.split(r'C:\tmp\path\to\file')) # Actual os.path.split will split following to 'C:\' and 'tmp' when run in # Windows. self.assertEqual((r'C:', 'tmp'), self.fs.split(r'C:\tmp')) def test_mkdirs(self): path = os.path.join(self.tmpdir, 't1/t2') self.fs.mkdirs(path) self.assertTrue(os.path.isdir(path)) def test_mkdirs_failed(self): path = os.path.join(self.tmpdir, 't1/t2') self.fs.mkdirs(path) # Check IOError if existing directory is created with self.assertRaises(IOError): self.fs.mkdirs(path) with self.assertRaises(IOError): self.fs.mkdirs(os.path.join(self.tmpdir, 't1')) def test_match_file(self): path = os.path.join(self.tmpdir, 'f1') open(path, 'a').close() # Match files in the temp directory result = self.fs.match([path])[0] files = [f.path for f in result.metadata_list] self.assertEqual(files, [path]) def test_match_file_empty(self): path = os.path.join(self.tmpdir, 'f2') # Does not exist # Match files in the temp directory result = self.fs.match([path])[0] files = [f.path for f in result.metadata_list] self.assertEqual(files, []) def test_match_file_exception(self): # Match files with None so that it throws an exception with self.assertRaisesRegex(BeamIOError, r'^Match operation failed') as error: self.fs.match([None]) self.assertEqual(list(error.exception.exception_details.keys()), [None]) @parameterized.expand([ param('*', files=['a', 'b', os.path.join('c', 'x')], expected=['a', 'b']), param( '**', files=['a', os.path.join('b', 'x'), os.path.join('c', 'x')], expected=['a', os.path.join('b', 'x'), os.path.join('c', 'x')]), param( os.path.join('*', '*'), files=[ 'a', os.path.join('b', 'x'), os.path.join('c', 'x'), os.path.join('d', 'x', 'y') ], expected=[os.path.join('b', 'x'), os.path.join('c', 'x')]), param( os.path.join('**', '*'), files=[ 'a', os.path.join('b', 'x'), os.path.join('c', 'x'), os.path.join('d', 'x', 'y') ], expected=[ os.path.join('b', 'x'), os.path.join('c', 'x'), os.path.join('d', 'x', 'y') ]), ]) def test_match_glob(self, pattern, files, expected): for filename in files: full_path = os.path.join(self.tmpdir, filename) dirname = os.path.dirname(full_path) if not dirname == full_path: # Make sure we don't go outside the tmpdir assert os.path.commonprefix([self.tmpdir, full_path]) == self.tmpdir try: self.fs.mkdirs(dirname) except IOError: # Directory exists pass open(full_path, 'a').close() # create empty file # Match both the files in the directory full_pattern = os.path.join(self.tmpdir, pattern) result = self.fs.match([full_pattern])[0] files = [os.path.relpath(f.path, self.tmpdir) for f in result.metadata_list] self.assertCountEqual(files, expected) def test_match_directory(self): result = self.fs.match([self.tmpdir])[0] files = [f.path for f in result.metadata_list] self.assertEqual(files, [self.tmpdir]) def test_match_directory_contents(self): path1 = os.path.join(self.tmpdir, 'f1') path2 = os.path.join(self.tmpdir, 'f2') open(path1, 'a').close() open(path2, 'a').close() result = self.fs.match([os.path.join(self.tmpdir, '*')])[0] files = [f.path for f in result.metadata_list] self.assertCountEqual(files, [path1, path2]) def test_copy(self): path1 = os.path.join(self.tmpdir, 'f1') path2 = os.path.join(self.tmpdir, 'f2') with open(path1, 'a') as f: f.write('Hello') self.fs.copy([path1], [path2]) self.assertTrue(filecmp.cmp(path1, path2)) def test_copy_error(self): path1 = os.path.join(self.tmpdir, 'f1') path2 = os.path.join(self.tmpdir, 'f2') with self.assertRaisesRegex(BeamIOError, r'^Copy operation failed') as error: self.fs.copy([path1], [path2]) self.assertEqual( list(error.exception.exception_details.keys()), [(path1, path2)]) def test_copy_directory(self): path_t1 = os.path.join(self.tmpdir, 't1') path_t2 = os.path.join(self.tmpdir, 't2') self.fs.mkdirs(path_t1) self.fs.mkdirs(path_t2) path1 = os.path.join(path_t1, 'f1') path2 = os.path.join(path_t2, 'f1') with open(path1, 'a') as f: f.write('Hello') self.fs.copy([path_t1], [path_t2]) self.assertTrue(filecmp.cmp(path1, path2)) def test_rename(self): path1 = os.path.join(self.tmpdir, 'f1') path2 = os.path.join(self.tmpdir, 'f2') with open(path1, 'a') as f: f.write('Hello') self.fs.rename([path1], [path2]) self.assertTrue(self.fs.exists(path2)) self.assertFalse(self.fs.exists(path1)) def test_rename_error(self): path1 = os.path.join(self.tmpdir, 'f1') path2 = os.path.join(self.tmpdir, 'f2') with self.assertRaisesRegex(BeamIOError, r'^Rename operation failed') as error: self.fs.rename([path1], [path2]) self.assertEqual( list(error.exception.exception_details.keys()), [(path1, path2)]) def test_rename_directory(self): path_t1 = os.path.join(self.tmpdir, 't1') path_t2 = os.path.join(self.tmpdir, 't2') self.fs.mkdirs(path_t1) path1 = os.path.join(path_t1, 'f1') path2 = os.path.join(path_t2, 'f1') with open(path1, 'a') as f: f.write('Hello') self.fs.rename([path_t1], [path_t2]) self.assertTrue(self.fs.exists(path_t2)) self.assertFalse(self.fs.exists(path_t1)) self.assertTrue(self.fs.exists(path2)) self.assertFalse(self.fs.exists(path1)) def test_exists(self): path1 = os.path.join(self.tmpdir, 'f1') path2 = os.path.join(self.tmpdir, 'f2') with open(path1, 'a') as f: f.write('Hello') self.assertTrue(self.fs.exists(path1)) self.assertFalse(self.fs.exists(path2)) def test_checksum(self): path1 = os.path.join(self.tmpdir, 'f1') path2 = os.path.join(self.tmpdir, 'f2') with open(path1, 'a') as f: f.write('Hello') with open(path2, 'a') as f: f.write('foo') # tests that localfilesystem checksum returns file size checksum1 = self.fs.checksum(path1) checksum2 = self.fs.checksum(path2) self.assertEqual(checksum1, str(5)) self.assertEqual(checksum2, str(3)) # tests that fs.checksum and str(fs.size) are consistent self.assertEqual(checksum1, str(self.fs.size(path1))) self.assertEqual(checksum2, str(self.fs.size(path2))) def make_tree(self, path, value, expected_leaf_count=None): """Create a file+directory structure from a simple dict-based DSL :param path: root path to create directories+files under :param value: a specification of what ``path`` should contain: ``None`` to make it an empty directory, a string literal to make it a file with those contents, and a ``dict`` to make it a non-empty directory and recurse :param expected_leaf_count: only be set at the top of a recursive call stack; after the whole tree has been created, verify the presence and number of all files+directories, as a sanity check """ if value is None: # empty directory os.makedirs(path) elif isinstance(value, str): # file with string-literal contents dir = os.path.dirname(path) if not os.path.exists(dir): os.makedirs(dir) with open(path, 'a') as f: f.write(value) elif isinstance(value, dict): # recurse to create a subdirectory tree for basename, v in value.items(): self.make_tree(os.path.join(path, basename), v) else: raise Exception('Unexpected value in tempdir tree: %s' % value) if expected_leaf_count is not None: self.assertEqual(self.check_tree(path, value), expected_leaf_count) def check_tree(self, path, value, expected_leaf_count=None): """Verify a directory+file structure according to the rules described in ``make_tree`` :param path: path to check under :param value: DSL-representation of expected files+directories under ``path`` :return: number of leaf files/directories that were verified """ actual_leaf_count = None if value is None: # empty directory self.assertTrue(os.path.exists(path), msg=path) self.assertEqual(os.listdir(path), []) actual_leaf_count = 1 elif isinstance(value, str): # file with string-literal contents with open(path, 'r') as f: self.assertEqual(f.read(), value, msg=path) actual_leaf_count = 1 elif isinstance(value, dict): # recurse to check subdirectory tree actual_leaf_count = sum([ self.check_tree(os.path.join(path, basename), v) for basename, v in value.items() ]) else: raise Exception('Unexpected value in tempdir tree: %s' % value) if expected_leaf_count is not None: self.assertEqual(actual_leaf_count, expected_leaf_count) return actual_leaf_count _test_tree = { 'path1': '111', 'path2': { '2': '222', 'emptydir': None }, 'aaa': { 'b1': 'b1', 'b2': None, 'bbb': { 'ccc': { 'ddd': 'DDD' } }, 'c': None } } def test_delete_globs(self): dir = os.path.join(self.tmpdir, 'dir') self.make_tree(dir, self._test_tree, expected_leaf_count=7) self.fs.delete([os.path.join(dir, 'path*'), os.path.join(dir, 'aaa', 'b*')]) # One empty nested directory is left self.check_tree(dir, {'aaa': {'c': None}}, expected_leaf_count=1) def test_recursive_delete(self): dir = os.path.join(self.tmpdir, 'dir') self.make_tree(dir, self._test_tree, expected_leaf_count=7) self.fs.delete([dir]) self.check_tree(self.tmpdir, {'': None}, expected_leaf_count=1) def test_delete_glob_errors(self): dir = os.path.join(self.tmpdir, 'dir') self.make_tree(dir, self._test_tree, expected_leaf_count=7) with self.assertRaisesRegex(BeamIOError, r'^Delete operation failed') as error: self.fs.delete([ os.path.join(dir, 'path*'), os.path.join(dir, 'aaa', 'b*'), os.path.join(dir, 'aaa', 'd*') # doesn't match anything, will raise ]) self.check_tree(dir, {'aaa': {'c': None}}, expected_leaf_count=1) self.assertEqual( list(error.exception.exception_details.keys()), [os.path.join(dir, 'aaa', 'd*')]) with self.assertRaisesRegex(BeamIOError, r'^Delete operation failed') as error: self.fs.delete([ os.path.join(dir, 'path*') # doesn't match anything, will raise ]) self.check_tree(dir, {'aaa': {'c': None}}, expected_leaf_count=1) self.assertEqual( list(error.exception.exception_details.keys()), [os.path.join(dir, 'path*')]) def test_delete(self): path1 = os.path.join(self.tmpdir, 'f1') with open(path1, 'a') as f: f.write('Hello') self.assertTrue(self.fs.exists(path1)) self.fs.delete([path1]) self.assertFalse(self.fs.exists(path1)) def test_delete_error(self): path1 = os.path.join(self.tmpdir, 'f1') with self.assertRaisesRegex(BeamIOError, r'^Delete operation failed') as error: self.fs.delete([path1]) self.assertEqual(list(error.exception.exception_details.keys()), [path1]) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/mongodbio.py0000644000175100017510000007035315206113032020623 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """This module implements IO classes to read and write data on MongoDB. Read from MongoDB ----------------- :class:`ReadFromMongoDB` is a ``PTransform`` that reads from a configured MongoDB source and returns a ``PCollection`` of dict representing MongoDB documents. To configure MongoDB source, the URI to connect to MongoDB server, database name, collection name needs to be provided. Example usage:: pipeline | ReadFromMongoDB(uri='mongodb://localhost:27017', db='testdb', coll='input') To read from MongoDB Atlas, use ``bucket_auto`` option to enable ``@bucketAuto`` MongoDB aggregation instead of ``splitVector`` command which is a high-privilege function that cannot be assigned to any user in Atlas. Example usage:: pipeline | ReadFromMongoDB(uri='mongodb+srv://user:pwd@cluster0.mongodb.net', db='testdb', coll='input', bucket_auto=True) Write to MongoDB: ----------------- :class:`WriteToMongoDB` is a ``PTransform`` that writes MongoDB documents to configured sink, and the write is conducted through a mongodb bulk_write of ``ReplaceOne`` operations. If the document's _id field already existed in the MongoDB collection, it results in an overwrite, otherwise, a new document will be inserted. Example usage:: pipeline | WriteToMongoDB(uri='mongodb://localhost:27017', db='testdb', coll='output', batch_size=10) No backward compatibility guarantees. Everything in this module is experimental. """ # pytype: skip-file import itertools import json import logging import math import struct from typing import Union import apache_beam as beam from apache_beam.io import iobase from apache_beam.io.range_trackers import LexicographicKeyRangeTracker from apache_beam.io.range_trackers import OffsetRangeTracker from apache_beam.io.range_trackers import OrderedPositionRangeTracker from apache_beam.transforms import DoFn from apache_beam.transforms import PTransform from apache_beam.transforms import Reshuffle _LOGGER = logging.getLogger(__name__) try: # Mongodb has its own bundled bson, which is not compatible with bson package. # (https://github.com/py-bson/bson/issues/82). Try to import objectid and if # it fails because bson package is installed, MongoDB IO will not work but at # least rest of the SDK will work. from bson import json_util from bson import objectid from bson.objectid import ObjectId # pymongo also internally depends on bson. from pymongo import ASCENDING from pymongo import DESCENDING from pymongo import MongoClient from pymongo import ReplaceOne from pymongo.driver_info import DriverInfo except ImportError: objectid = None json_util = None ObjectId = None ASCENDING = 1 DESCENDING = -1 MongoClient = None ReplaceOne = None DriverInfo = None _LOGGER.warning("Could not find a compatible bson package.") __all__ = ["ReadFromMongoDB", "WriteToMongoDB"] class ReadFromMongoDB(PTransform): """A ``PTransform`` to read MongoDB documents into a ``PCollection``.""" def __init__( self, uri="mongodb://localhost:27017", db=None, coll=None, filter=None, projection=None, extra_client_params=None, bucket_auto=False, ): """Initialize a :class:`ReadFromMongoDB` Args: uri (str): The MongoDB connection string following the URI format. db (str): The MongoDB database name. coll (str): The MongoDB collection name. filter: A `bson.SON `_ object specifying elements which must be present for a document to be included in the result set. projection: A list of field names that should be returned in the result set or a dict specifying the fields to include or exclude. extra_client_params(dict): Optional `MongoClient `_ parameters. bucket_auto (bool): If :data:`True`, use MongoDB `$bucketAuto` aggregation to split collection into bundles instead of `splitVector` command, which does not work with MongoDB Atlas. If :data:`False` (the default), use `splitVector` command for bundling. Returns: :class:`~apache_beam.transforms.ptransform.PTransform` """ if extra_client_params is None: extra_client_params = {} if not isinstance(db, str): raise ValueError("ReadFromMongDB db param must be specified as a string") if not isinstance(coll, str): raise ValueError( "ReadFromMongDB coll param must be specified as a string") self._mongo_source = _BoundedMongoSource( uri=uri, db=db, coll=coll, filter=filter, projection=projection, extra_client_params=extra_client_params, bucket_auto=bucket_auto, ) def expand(self, pcoll): return pcoll | iobase.Read(self._mongo_source) class _ObjectIdRangeTracker(OrderedPositionRangeTracker): """RangeTracker for tracking mongodb _id of bson ObjectId type.""" def position_to_fraction( self, pos: ObjectId, start: ObjectId, end: ObjectId, ): """Returns the fraction of keys in the range [start, end) that are less than the given key. """ pos_number = _ObjectIdHelper.id_to_int(pos) start_number = _ObjectIdHelper.id_to_int(start) end_number = _ObjectIdHelper.id_to_int(end) return (pos_number - start_number) / (end_number - start_number) def fraction_to_position( self, fraction: float, start: ObjectId, end: ObjectId, ): """Converts a fraction between 0 and 1 to a position between start and end. """ start_number = _ObjectIdHelper.id_to_int(start) end_number = _ObjectIdHelper.id_to_int(end) total = end_number - start_number pos = int(total * fraction + start_number) # make sure split position is larger than start position and smaller than # end position. if pos <= start_number: return _ObjectIdHelper.increment_id(start, 1) if pos >= end_number: return _ObjectIdHelper.increment_id(end, -1) return _ObjectIdHelper.int_to_id(pos) class _BoundedMongoSource(iobase.BoundedSource): """A MongoDB source that reads a finite amount of input records. This class defines following operations which can be used to read MongoDB source efficiently. * Size estimation - method ``estimate_size()`` may return an accurate estimation in bytes for the size of the source. * Splitting into bundles of a given size - method ``split()`` can be used to split the source into a set of sub-sources (bundles) based on a desired bundle size. * Getting a RangeTracker - method ``get_range_tracker()`` should return a ``RangeTracker`` object for a given position range for the position type of the records returned by the source. * Reading the data - method ``read()`` can be used to read data from the source while respecting the boundaries defined by a given ``RangeTracker``. A runner will perform reading the source in two steps. (1) Method ``get_range_tracker()`` will be invoked with start and end positions to obtain a ``RangeTracker`` for the range of positions the runner intends to read. Source must define a default initial start and end position range. These positions must be used if the start and/or end positions passed to the method ``get_range_tracker()`` are ``None`` (2) Method read() will be invoked with the ``RangeTracker`` obtained in the previous step. **Mutability** A ``_BoundedMongoSource`` object should not be mutated while its methods (for example, ``read()``) are being invoked by a runner. Runner implementations may invoke methods of ``_BoundedMongoSource`` objects through multi-threaded and/or reentrant execution modes. """ def __init__( self, uri=None, db=None, coll=None, filter=None, projection=None, extra_client_params=None, bucket_auto=False, ): if extra_client_params is None: extra_client_params = {} if filter is None: filter = {} self.uri = uri self.db = db self.coll = coll self.filter = filter self.projection = projection self.spec = extra_client_params self.bucket_auto = bucket_auto if "driver" not in self.spec: self.spec["driver"] = DriverInfo( name="Apache Beam", version=beam.__version__, ) def estimate_size(self): with MongoClient(self.uri, **self.spec) as client: return client[self.db].command("collstats", self.coll).get("size") def _estimate_average_document_size(self): with MongoClient(self.uri, **self.spec) as client: return client[self.db].command("collstats", self.coll).get("avgObjSize") def split( self, desired_bundle_size: int, start_position: Union[int, str, bytes, ObjectId] = None, stop_position: Union[int, str, bytes, ObjectId] = None, ): """Splits the source into a set of bundles. Bundles should be approximately of size ``desired_bundle_size`` bytes. Args: desired_bundle_size: the desired size (in bytes) of the bundles returned. start_position: if specified the given position must be used as the starting position of the first bundle. stop_position: if specified the given position must be used as the ending position of the last bundle. Returns: an iterator of objects of type 'SourceBundle' that gives information about the generated bundles. """ desired_bundle_size_in_mb = desired_bundle_size // 1024 // 1024 # for desired bundle size, if desired chunk size smaller than 1mb, use # MongoDB default split size of 1mb. desired_bundle_size_in_mb = max(desired_bundle_size_in_mb, 1) is_initial_split = start_position is None and stop_position is None start_position, stop_position = self._replace_none_positions( start_position, stop_position ) if self.bucket_auto: # Use $bucketAuto for bundling split_keys = [] weights = [] for bucket in self._get_auto_buckets( desired_bundle_size_in_mb, start_position, stop_position, is_initial_split, ): split_keys.append({"_id": bucket["_id"]["max"]}) weights.append(bucket["count"]) else: # Use splitVector for bundling split_keys = self._get_split_keys( desired_bundle_size_in_mb, start_position, stop_position) weights = itertools.cycle((desired_bundle_size_in_mb, )) bundle_start = start_position for split_key_id, weight in zip(split_keys, weights): if bundle_start >= stop_position: break bundle_end = min(stop_position, split_key_id["_id"]) yield iobase.SourceBundle( weight=weight, source=self, start_position=bundle_start, stop_position=bundle_end, ) bundle_start = bundle_end # add range of last split_key to stop_position if bundle_start < stop_position: # bucket_auto mode can come here if not split due to single document weight = 1 if self.bucket_auto else desired_bundle_size_in_mb yield iobase.SourceBundle( weight=weight, source=self, start_position=bundle_start, stop_position=stop_position, ) def get_range_tracker( self, start_position: Union[int, str, ObjectId] = None, stop_position: Union[int, str, ObjectId] = None, ) -> Union[ _ObjectIdRangeTracker, OffsetRangeTracker, LexicographicKeyRangeTracker]: """Returns a RangeTracker for a given position range depending on type. Args: start_position: starting position of the range. If 'None' default start position of the source must be used. stop_position: ending position of the range. If 'None' default stop position of the source must be used. Returns: a ``_ObjectIdRangeTracker``, ``OffsetRangeTracker`` or ``LexicographicKeyRangeTracker`` depending on the given position range. """ start_position, stop_position = self._replace_none_positions( start_position, stop_position ) if isinstance(start_position, ObjectId): return _ObjectIdRangeTracker(start_position, stop_position) if isinstance(start_position, int): return OffsetRangeTracker(start_position, stop_position) if isinstance(start_position, str): return LexicographicKeyRangeTracker(start_position, stop_position) raise NotImplementedError( f"RangeTracker for {type(start_position)} not implemented!") def read(self, range_tracker): """Returns an iterator that reads data from the source. The returned set of data must respect the boundaries defined by the given ``RangeTracker`` object. For example: * Returned set of data must be for the range ``[range_tracker.start_position, range_tracker.stop_position)``. Note that a source may decide to return records that start after ``range_tracker.stop_position``. See documentation in class ``RangeTracker`` for more details. Also, note that framework might invoke ``range_tracker.try_split()`` to perform dynamic split operations. range_tracker.stop_position may be updated dynamically due to successful dynamic split operations. * Method ``range_tracker.try_split()`` must be invoked for every record that starts at a split point. * Method ``range_tracker.record_current_position()`` may be invoked for records that do not start at split points. Args: range_tracker: a ``RangeTracker`` whose boundaries must be respected when reading data from the source. A runner that reads this source muss pass a ``RangeTracker`` object that is not ``None``. Returns: an iterator of data read by the source. """ with MongoClient(self.uri, **self.spec) as client: all_filters = self._merge_id_filter( range_tracker.start_position(), range_tracker.stop_position()) docs_cursor = ( client[self.db][self.coll].find( filter=all_filters, projection=self.projection).sort([("_id", ASCENDING)])) for doc in docs_cursor: if not range_tracker.try_claim(doc["_id"]): return yield doc def display_data(self): """Returns the display data associated to a pipeline component.""" res = super().display_data() res["database"] = self.db res["collection"] = self.coll res["filter"] = json.dumps(self.filter, default=json_util.default) res["projection"] = str(self.projection) res["bucket_auto"] = self.bucket_auto return res @staticmethod def _range_is_not_splittable( start_pos: Union[int, str, ObjectId], end_pos: Union[int, str, ObjectId], ): """Return `True` if splitting range doesn't make sense (single document is not splittable), Return `False` otherwise. """ return (( isinstance(start_pos, ObjectId) and start_pos >= _ObjectIdHelper.increment_id(end_pos, -1)) or (isinstance(start_pos, int) and start_pos >= end_pos - 1) or (isinstance(start_pos, str) and start_pos >= end_pos)) def _get_split_keys( self, desired_chunk_size_in_mb: int, start_pos: Union[int, str, ObjectId], end_pos: Union[int, str, ObjectId], ): """Calls MongoDB `splitVector` command to get document ids at split position. """ # single document not splittable if self._range_is_not_splittable(start_pos, end_pos): return [] with MongoClient(self.uri, **self.spec) as client: name_space = "%s.%s" % (self.db, self.coll) return client[self.db].command( "splitVector", name_space, keyPattern={"_id": 1}, # Ascending index min={"_id": start_pos}, max={"_id": end_pos}, maxChunkSize=desired_chunk_size_in_mb, )["splitKeys"] def _get_auto_buckets( self, desired_chunk_size_in_mb: int, start_pos: Union[int, str, ObjectId], end_pos: Union[int, str, ObjectId], is_initial_split: bool, ) -> list: """Use MongoDB `$bucketAuto` aggregation to split collection into bundles instead of `splitVector` command, which does not work with MongoDB Atlas. """ # single document not splittable if self._range_is_not_splittable(start_pos, end_pos): return [] if is_initial_split and not self.filter: # total collection size in MB size_in_mb = self.estimate_size() / float(1 << 20) else: # size of documents within start/end id range and possibly filtered documents_count = self._count_id_range(start_pos, end_pos) avg_document_size = self._estimate_average_document_size() size_in_mb = documents_count * avg_document_size / float(1 << 20) if size_in_mb == 0: # no documents not splittable (maybe a result of filtering) return [] bucket_count = math.ceil(size_in_mb / desired_chunk_size_in_mb) with beam.io.mongodbio.MongoClient(self.uri, **self.spec) as client: pipeline = [ { # filter by positions and by the custom filter if any "$match": self._merge_id_filter(start_pos, end_pos) }, { "$bucketAuto": { "groupBy": "$_id", "buckets": bucket_count } }, ] buckets = list( # Use `allowDiskUse` option to avoid aggregation limit of 100 Mb RAM client[self.db][self.coll].aggregate(pipeline, allowDiskUse=True)) if buckets: buckets[-1]["_id"]["max"] = end_pos return buckets def _merge_id_filter( self, start_position: Union[int, str, bytes, ObjectId], stop_position: Union[int, str, bytes, ObjectId] = None, ) -> dict: """Merge the default filter (if any) with refined _id field range of range_tracker. $gte specifies start position (inclusive) and $lt specifies the end position (exclusive), see more at https://docs.mongodb.com/manual/reference/operator/query/gte/ and https://docs.mongodb.com/manual/reference/operator/query/lt/ """ if stop_position is None: id_filter = {"_id": {"$gte": start_position}} else: id_filter = {"_id": {"$gte": start_position, "$lt": stop_position}} if self.filter: all_filters = { # see more at # https://docs.mongodb.com/manual/reference/operator/query/and/ "$and": [self.filter.copy(), id_filter] } else: all_filters = id_filter return all_filters def _get_head_document_id(self, sort_order): with MongoClient(self.uri, **self.spec) as client: cursor = ( client[self.db][self.coll].find(filter={}, projection=[]).sort([ ("_id", sort_order) ]).limit(1)) try: return cursor[0]["_id"] except IndexError: raise ValueError("Empty Mongodb collection") def _replace_none_positions(self, start_position, stop_position): if start_position is None: start_position = self._get_head_document_id(ASCENDING) if stop_position is None: last_doc_id = self._get_head_document_id(DESCENDING) # increment last doc id binary value by 1 to make sure the last document # is not excluded if isinstance(last_doc_id, ObjectId): stop_position = _ObjectIdHelper.increment_id(last_doc_id, 1) elif isinstance(last_doc_id, int): stop_position = last_doc_id + 1 elif isinstance(last_doc_id, str): stop_position = last_doc_id + '\x00' return start_position, stop_position def _count_id_range(self, start_position, stop_position): """Number of documents between start_position (inclusive) and stop_position (exclusive), respecting the custom filter if any. """ with MongoClient(self.uri, **self.spec) as client: return client[self.db][self.coll].count_documents( filter=self._merge_id_filter(start_position, stop_position)) class _ObjectIdHelper: """A Utility class to manipulate bson object ids.""" @classmethod def id_to_int(cls, _id: Union[int, ObjectId]) -> int: """ Args: _id: ObjectId required for each MongoDB document _id field. Returns: Converted integer value of ObjectId's 12 bytes binary value. """ if isinstance(_id, int): return _id # converts object id binary to integer # id object is bytes type with size of 12 ints = struct.unpack(">III", _id.binary) return (ints[0] << 64) + (ints[1] << 32) + ints[2] @classmethod def int_to_id(cls, number): """ Args: number(int): The integer value to be used to convert to ObjectId. Returns: The ObjectId that has the 12 bytes binary converted from the integer value. """ # converts integer value to object id. Int value should be less than # (2 ^ 96) so it can be convert to 12 bytes required by object id. if number < 0 or number >= (1 << 96): raise ValueError("number value must be within [0, %s)" % (1 << 96)) ints = [ (number & 0xFFFFFFFF0000000000000000) >> 64, (number & 0x00000000FFFFFFFF00000000) >> 32, number & 0x0000000000000000FFFFFFFF, ] number_bytes = struct.pack(">III", *ints) return ObjectId(number_bytes) @classmethod def increment_id( cls, _id: ObjectId, inc: int, ) -> ObjectId: """ Increment object_id binary value by inc value and return new object id. Args: _id: The `_id` to change. inc(int): The incremental int value to be added to `_id`. Returns: `_id` incremented by `inc` value """ id_number = _ObjectIdHelper.id_to_int(_id) new_number = id_number + inc if new_number < 0 or new_number >= (1 << 96): raise ValueError( "invalid incremental, inc value must be within [" "%s, %s)" % (0 - id_number, 1 << 96 - id_number)) return _ObjectIdHelper.int_to_id(new_number) class WriteToMongoDB(PTransform): """WriteToMongoDB is a ``PTransform`` that writes a ``PCollection`` of mongodb document to the configured MongoDB server. In order to make the document writes idempotent so that the bundles are retry-able without creating duplicates, the PTransform added 2 transformations before final write stage: a ``GenerateId`` transform and a ``Reshuffle`` transform.:: ----------------------------------------------- Pipeline --> |GenerateId --> Reshuffle --> WriteToMongoSink| ----------------------------------------------- (WriteToMongoDB) The ``GenerateId`` transform adds a random and unique*_id* field to the documents if they don't already have one, it uses the same format as MongoDB default. The ``Reshuffle`` transform makes sure that no fusion happens between ``GenerateId`` and the final write stage transform,so that the set of documents and their unique IDs are not regenerated if final write step is retried due to a failure. This prevents duplicate writes of the same document with different unique IDs. """ def __init__( self, uri="mongodb://localhost:27017", db=None, coll=None, batch_size=100, extra_client_params=None, ): """ Args: uri (str): The MongoDB connection string following the URI format db (str): The MongoDB database name coll (str): The MongoDB collection name batch_size(int): Number of documents per bulk_write to MongoDB, default to 100 extra_client_params(dict): Optional `MongoClient `_ parameters as keyword arguments Returns: :class:`~apache_beam.transforms.ptransform.PTransform` """ if extra_client_params is None: extra_client_params = {} if not isinstance(db, str): raise ValueError("WriteToMongoDB db param must be specified as a string") if not isinstance(coll, str): raise ValueError( "WriteToMongoDB coll param must be specified as a string") self._uri = uri self._db = db self._coll = coll self._batch_size = batch_size self._spec = extra_client_params def expand(self, pcoll): return ( pcoll | beam.ParDo(_GenerateObjectIdFn()) | Reshuffle() | beam.ParDo( _WriteMongoFn( self._uri, self._db, self._coll, self._batch_size, self._spec))) class _GenerateObjectIdFn(DoFn): def process(self, element, *args, **kwargs): # if _id field already exist we keep it as it is, otherwise the ptransform # generates a new _id field to achieve idempotent write to mongodb. if "_id" not in element: # object.ObjectId() generates a unique identifier that follows mongodb # default format, if _id is not present in document, mongodb server # generates it with this same function upon write. However the # uniqueness of generated id may not be guaranteed if the work load are # distributed across too many processes. See more on the ObjectId format # https://docs.mongodb.com/manual/reference/bson-types/#objectid. element["_id"] = objectid.ObjectId() yield element class _WriteMongoFn(DoFn): def __init__( self, uri=None, db=None, coll=None, batch_size=100, extra_params=None): if extra_params is None: extra_params = {} self.uri = uri self.db = db self.coll = coll self.spec = extra_params self.batch_size = batch_size self.batch = [] def finish_bundle(self): self._flush() def process(self, element, *args, **kwargs): self.batch.append(element) if len(self.batch) >= self.batch_size: self._flush() def _flush(self): if len(self.batch) == 0: return with _MongoSink(self.uri, self.db, self.coll, self.spec) as sink: sink.write(self.batch) self.batch = [] def display_data(self): res = super().display_data() res["database"] = self.db res["collection"] = self.coll res["batch_size"] = self.batch_size return res class _MongoSink: def __init__(self, uri=None, db=None, coll=None, extra_params=None): if extra_params is None: extra_params = {} self.uri = uri self.db = db self.coll = coll self.spec = extra_params self.client = None if "driver" not in self.spec: self.spec["driver"] = DriverInfo( name="Apache Beam", version=beam.__version__, ) def write(self, documents): if self.client is None: self.client = MongoClient(host=self.uri, **self.spec) requests = [] for doc in documents: # match document based on _id field, if not found in current collection, # insert new one, otherwise overwrite it. requests.append( ReplaceOne( filter={"_id": doc.get("_id", None)}, replacement=doc, upsert=True)) resp = self.client[self.db][self.coll].bulk_write(requests) _LOGGER.debug( "BulkWrite to MongoDB result in nModified:%d, nUpserted:%d, " "nMatched:%d, Errors:%s" % ( resp.modified_count, resp.upserted_count, resp.matched_count, resp.bulk_api_result.get("writeErrors"), )) def __enter__(self): if self.client is None: self.client = MongoClient(host=self.uri, **self.spec) return self def __exit__(self, exc_type, exc_val, exc_tb): if self.client is not None: self.client.close() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/mongodbio_it_test.py0000644000175100017510000001253515206113032022354 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import argparse import logging import time from pymongo import MongoClient import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to _LOGGER = logging.getLogger(__name__) class GenerateDocs(beam.DoFn): def process(self, num_docs, *args, **kwargs): for i in range(num_docs): yield {'number': i, 'number_mod_2': i % 2, 'number_mod_3': i % 3} def run(argv=None): default_db = 'beam_mongodbio_it_db' default_coll = 'integration_test_%d' % time.time() parser = argparse.ArgumentParser() parser.add_argument( '--mongo_uri', default='mongodb://localhost:27017', help='mongo uri string for connection') parser.add_argument( '--mongo_db', default=default_db, help='mongo uri string for connection') parser.add_argument( '--mongo_coll', default=default_coll, help='mongo uri string for connection') parser.add_argument( '--num_documents', default=100000, help='The expected number of documents to be generated ' 'for write or read', type=int) parser.add_argument( '--batch_size', default=10000, type=int, help=('batch size for writing to mongodb')) known_args, pipeline_args = parser.parse_known_args(argv) # Test Write to MongoDB with TestPipeline(options=PipelineOptions(pipeline_args)) as p: start_time = time.time() _LOGGER.info('Writing %d documents to mongodb', known_args.num_documents) _ = ( p | beam.Create([known_args.num_documents]) | 'Create documents' >> beam.ParDo(GenerateDocs()) | 'WriteToMongoDB' >> beam.io.WriteToMongoDB( known_args.mongo_uri, known_args.mongo_db, known_args.mongo_coll, known_args.batch_size)) elapsed = time.time() - start_time _LOGGER.info( 'Writing %d documents to mongodb finished in %.3f seconds' % (known_args.num_documents, elapsed)) # Test Read from MongoDB total_sum = sum(range(known_args.num_documents)) mod_3_sum = sum( num for num in range(known_args.num_documents) if num % 3 == 0) mod_3_count = sum( 1 for num in range(known_args.num_documents) if num % 3 == 0) # yapf: disable read_cases = [ # (reader_params, expected) ( { 'projection': ['number'] }, { 'number_sum': total_sum, 'docs_count': known_args.num_documents } ), ( { 'filter': {'number_mod_3': 0}, 'projection': ['number'] }, { 'number_sum': mod_3_sum, 'docs_count': mod_3_count } ), ( { 'projection': ['number'], 'bucket_auto': True }, { 'number_sum': total_sum, 'docs_count': known_args.num_documents } ), ( { 'filter': {'number_mod_3': 0}, 'projection': ['number'], 'bucket_auto': True }, { 'number_sum': mod_3_sum, 'docs_count': mod_3_count } ), ] # yapf: enable for reader_params, expected in read_cases: with TestPipeline(options=PipelineOptions(pipeline_args)) as p: start_time = time.time() _LOGGER.info('=' * 80) _LOGGER.info( 'Reading from mongodb %s:%s', known_args.mongo_db, known_args.mongo_coll) _LOGGER.info('reader params : %s', reader_params) _LOGGER.info('expected results: %s', expected) docs = ( p | 'ReadFromMongoDB' >> beam.io.ReadFromMongoDB( known_args.mongo_uri, known_args.mongo_db, known_args.mongo_coll, **reader_params) | 'Map' >> beam.Map(lambda doc: doc['number'])) number_sum = (docs | 'Combine' >> beam.CombineGlobally(sum)) docs_count = (docs | 'Count' >> beam.combiners.Count.Globally()) r = ([number_sum, docs_count] | 'Flatten' >> beam.Flatten()) assert_that(r, equal_to([expected['number_sum'], expected['docs_count']])) elapsed = time.time() - start_time _LOGGER.info( 'Reading documents from mongodb finished in %.3f seconds', elapsed) # Clean-up with MongoClient(host=known_args.mongo_uri) as client: client.drop_database(known_args.mongo_db) if __name__ == "__main__": logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/mongodbio_test.py0000644000175100017510000006267115206113032021666 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import datetime import logging import random import unittest from typing import Union from unittest import TestCase import mock from bson import ObjectId from bson import objectid from parameterized import parameterized_class from pymongo import ASCENDING from pymongo import ReplaceOne import apache_beam as beam from apache_beam.io import ReadFromMongoDB from apache_beam.io import WriteToMongoDB from apache_beam.io import source_test_utils from apache_beam.io.mongodbio import _BoundedMongoSource from apache_beam.io.mongodbio import _GenerateObjectIdFn from apache_beam.io.mongodbio import _MongoSink from apache_beam.io.mongodbio import _ObjectIdHelper from apache_beam.io.mongodbio import _ObjectIdRangeTracker from apache_beam.io.mongodbio import _WriteMongoFn from apache_beam.io.range_trackers import LexicographicKeyRangeTracker from apache_beam.io.range_trackers import OffsetRangeTracker from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class _MockMongoColl(object): """Fake mongodb collection cursor.""" def __init__(self, docs): self.docs = docs def __getitem__(self, index): return self.docs[index] def __len__(self): return len(self.docs) @staticmethod def _make_filter(conditions): assert isinstance(conditions, dict) checks = [] for field, value in conditions.items(): if isinstance(value, dict): for op, val in value.items(): if op == '$gte': op = '__ge__' elif op == '$lt': op = '__lt__' else: raise Exception('Operator "{0}" not supported.'.format(op)) checks.append((field, op, val)) else: checks.append((field, '__eq__', value)) def func(doc): for field, op, value in checks: if not getattr(doc[field], op)(value): return False return True return func def _filter(self, filter): match = [] if not filter: return self all_filters = [] if '$and' in filter: for item in filter['$and']: all_filters.append(self._make_filter(item)) else: all_filters.append(self._make_filter(filter)) for doc in self.docs: if not all(check(doc) for check in all_filters): continue match.append(doc) return match @staticmethod def _projection(docs, projection=None): if projection: return [{ k: v for k, v in doc.items() if k in projection or k == '_id' } for doc in docs] return docs def find(self, filter=None, projection=None, **kwargs): return _MockMongoColl(self._projection(self._filter(filter), projection)) def sort(self, sort_items): key, order = sort_items[0] self.docs = sorted( self.docs, key=lambda x: x[key], reverse=(order != ASCENDING)) return self def limit(self, num): return _MockMongoColl(self.docs[0:num]) def count_documents(self, filter): return len(self._filter(filter)) def aggregate(self, pipeline, **kwargs): # Simulate $bucketAuto aggregate pipeline. # Example splits doc count for the total of 5 docs: # - 1 bucket: [5] # - 2 buckets: [3, 2] # - 3 buckets: [2, 2, 1] # - 4 buckets: [2, 1, 1, 1] # - 5 buckets: [1, 1, 1, 1, 1] match_step = next((step for step in pipeline if '$match' in step), None) bucket_auto_step = next(step for step in pipeline if '$bucketAuto' in step) if match_step is None: docs = self.docs else: docs = self.find(filter=match_step['$match']) doc_count = len(docs) bucket_count = min(bucket_auto_step['$bucketAuto']['buckets'], doc_count) # bucket_count ≠ 0 bucket_len, remainder = divmod(doc_count, bucket_count) bucket_sizes = ( remainder * [bucket_len + 1] + (bucket_count - remainder) * [bucket_len]) buckets = [] start = 0 for bucket_size in bucket_sizes: stop = start + bucket_size if stop >= doc_count: # MongoDB: the last bucket's 'max' is inclusive stop = doc_count - 1 count = stop - start + 1 else: # non-last bucket's 'max' is exclusive and == next bucket's 'min' count = stop - start buckets.append({ '_id': { 'min': docs[start]['_id'], 'max': docs[stop]['_id'], }, 'count': count }) start = stop return buckets class _MockMongoDb(object): """Fake Mongo Db.""" def __init__(self, docs): self.docs = docs def __getitem__(self, coll_name): return _MockMongoColl(self.docs) def command(self, command, *args, **kwargs): if command == 'collstats': return {'size': 5 * 1024 * 1024, 'avgObjSize': 1 * 1024 * 1024} if command == 'splitVector': return self.get_split_keys(command, *args, **kwargs) def get_split_keys(self, command, ns, min, max, maxChunkSize, **kwargs): # simulate mongo db splitVector command, return split keys base on chunk # size, assuming every doc is of size 1mb start_id = min['_id'] end_id = max['_id'] if start_id >= end_id: return [] start_index = 0 end_index = 0 # get split range of [min, max] for doc in self.docs: if doc['_id'] < start_id: start_index += 1 if doc['_id'] <= end_id: end_index += 1 else: break # Return ids of elements in the range with chunk size skip and exclude # head element. For simplicity of tests every document is considered 1Mb # by default. return { 'splitKeys': [{ '_id': x['_id'] } for x in self.docs[start_index:end_index:maxChunkSize]][1:] } class _MockMongoClient: def __init__(self, docs): self.docs = docs def __getitem__(self, db_name): return _MockMongoDb(self.docs) def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): pass # Generate test data for MongoDB collections of different types OBJECT_IDS = [ objectid.ObjectId.from_datetime( datetime.datetime(year=2020, month=i + 1, day=i + 1)) for i in range(5) ] INT_IDS = [n for n in range(5)] # [0, 1, 2, 3, 4] STR_IDS_1 = [str(n) for n in range(5)] # ['0', '1', '2', '3', '4'] # ['aaaaa', 'bbbbb', 'ccccc', 'ddddd', 'eeeee'] STR_IDS_2 = [chr(97 + n) * 5 for n in range(5)] # ['AAAAAAAAAAAAAAAAAAAA', 'BBBBBBBBBBBBBBBBBBBB', ..., 'EEEEEEEEEEEEEEEEEEEE'] STR_IDS_3 = [chr(65 + n) * 20 for n in range(5)] @parameterized_class(('bucket_auto', '_ids', 'min_id', 'max_id'), [ ( None, OBJECT_IDS, _ObjectIdHelper.int_to_id(0), _ObjectIdHelper.int_to_id(2**96 - 1)), ( True, OBJECT_IDS, _ObjectIdHelper.int_to_id(0), _ObjectIdHelper.int_to_id(2**96 - 1)), ( None, INT_IDS, 0, 2**96 - 1, ), ( True, INT_IDS, 0, 2**96 - 1, ), ( None, STR_IDS_1, chr(0), chr(0x10ffff), ), ( True, STR_IDS_1, chr(0), chr(0x10ffff), ), ( None, STR_IDS_2, chr(0), chr(0x10ffff), ), ( True, STR_IDS_2, chr(0), chr(0x10ffff), ), ( None, STR_IDS_3, chr(0), chr(0x10ffff), ), ( True, STR_IDS_3, chr(0), chr(0x10ffff), ), ]) class MongoSourceTest(unittest.TestCase): @mock.patch('apache_beam.io.mongodbio.MongoClient') def setUp(self, mock_client): self._docs = [{'_id': self._ids[i], 'x': i} for i in range(len(self._ids))] mock_client.return_value = _MockMongoClient(self._docs) self.mongo_source = self._create_source(bucket_auto=self.bucket_auto) @staticmethod def _create_source(filter=None, bucket_auto=None): kwargs = {} if filter is not None: kwargs['filter'] = filter if bucket_auto is not None: kwargs['bucket_auto'] = bucket_auto return _BoundedMongoSource('mongodb://test', 'testdb', 'testcoll', **kwargs) def _increment_id( self, _id: Union[ObjectId, int, str], inc: int, ) -> Union[ObjectId, int, str]: """Helper method to increment `_id` of different types.""" if isinstance(_id, ObjectId): return _ObjectIdHelper.increment_id(_id, inc) if isinstance(_id, int): return _id + inc if isinstance(_id, str): index = self._ids.index(_id) + inc if index <= 0: return self._ids[0] if index >= len(self._ids): return self._ids[-1] return self._ids[index] @mock.patch('apache_beam.io.mongodbio.MongoClient') def test_estimate_size(self, mock_client): mock_client.return_value = _MockMongoClient(self._docs) self.assertEqual(self.mongo_source.estimate_size(), 5 * 1024 * 1024) @mock.patch('apache_beam.io.mongodbio.MongoClient') def test_estimate_average_document_size(self, mock_client): mock_client.return_value = _MockMongoClient(self._docs) self.assertEqual( self.mongo_source._estimate_average_document_size(), 1 * 1024 * 1024) @mock.patch('apache_beam.io.mongodbio.MongoClient') def test_split(self, mock_client): mock_client.return_value = _MockMongoClient(self._docs) for size_mb, expected_split_count in [(0.5, 5), (1, 5), (2, 3), (10, 1)]: size = size_mb * 1024 * 1024 splits = list( self.mongo_source.split( start_position=None, stop_position=None, desired_bundle_size=size)) self.assertEqual(len(splits), expected_split_count) reference_info = (self.mongo_source, None, None) sources_info = ([ (split.source, split.start_position, split.stop_position) for split in splits ]) source_test_utils.assert_sources_equal_reference_source( reference_info, sources_info) @mock.patch('apache_beam.io.mongodbio.MongoClient') def test_split_single_document(self, mock_client): mock_client.return_value = _MockMongoClient(self._docs[0:1]) for size_mb in [1, 5]: size = size_mb * 1024 * 1024 splits = list( self.mongo_source.split( start_position=None, stop_position=None, desired_bundle_size=size)) self.assertEqual(len(splits), 1) _id = self._docs[0]['_id'] assert _id == splits[0].start_position assert _id <= splits[0].stop_position if isinstance(_id, (ObjectId, int)): # We can unambiguously determine next `_id` assert self._increment_id(_id, 1) == splits[0].stop_position @mock.patch('apache_beam.io.mongodbio.MongoClient') def test_split_no_documents(self, mock_client): mock_client.return_value = _MockMongoClient([]) with self.assertRaises(ValueError) as cm: list( self.mongo_source.split( start_position=None, stop_position=None, desired_bundle_size=1024 * 1024)) self.assertEqual(str(cm.exception), 'Empty Mongodb collection') @mock.patch('apache_beam.io.mongodbio.MongoClient') def test_split_filtered(self, mock_client): # filtering 2 documents: 2 <= 'x' < 4 filtered_mongo_source = self._create_source( filter={'x': { '$gte': 2, '$lt': 4 }}, bucket_auto=self.bucket_auto) mock_client.return_value = _MockMongoClient(self._docs) for size_mb, (bucket_auto_count, split_vector_count) in [(1, (2, 5)), (2, (1, 3)), (10, (1, 1))]: size = size_mb * 1024 * 1024 splits = list( filtered_mongo_source.split( start_position=None, stop_position=None, desired_bundle_size=size)) if self.bucket_auto: self.assertEqual(len(splits), bucket_auto_count) else: # Note: splitVector mode does not respect filter self.assertEqual(len(splits), split_vector_count) reference_info = ( filtered_mongo_source, self._docs[2]['_id'], self._docs[4]['_id']) sources_info = ([ (split.source, split.start_position, split.stop_position) for split in splits ]) source_test_utils.assert_sources_equal_reference_source( reference_info, sources_info) @mock.patch('apache_beam.io.mongodbio.MongoClient') def test_split_filtered_empty(self, mock_client): # filtering doesn't match any documents filtered_mongo_source = self._create_source( filter={'x': { '$lt': 0 }}, bucket_auto=self.bucket_auto) mock_client.return_value = _MockMongoClient(self._docs) for size_mb, (bucket_auto_count, split_vector_count) in [(1, (1, 5)), (2, (1, 3)), (10, (1, 1))]: size = size_mb * 1024 * 1024 splits = list( filtered_mongo_source.split( start_position=None, stop_position=None, desired_bundle_size=size)) if self.bucket_auto: # Note: if filter matches no docs - one split covers entire range self.assertEqual(len(splits), bucket_auto_count) else: # Note: splitVector mode does not respect filter self.assertEqual(len(splits), split_vector_count) reference_info = ( filtered_mongo_source, # range to match no documents: self._increment_id(self._docs[-1]['_id'], 1), self._increment_id(self._docs[-1]['_id'], 2), ) sources_info = ([ (split.source, split.start_position, split.stop_position) for split in splits ]) source_test_utils.assert_sources_equal_reference_source( reference_info, sources_info) @mock.patch('apache_beam.io.mongodbio.MongoClient') def test_dynamic_work_rebalancing(self, mock_client): mock_client.return_value = _MockMongoClient(self._docs) splits = list( self.mongo_source.split(desired_bundle_size=3000 * 1024 * 1024)) assert len(splits) == 1 source_test_utils.assert_split_at_fraction_exhaustive( splits[0].source, splits[0].start_position, splits[0].stop_position) @mock.patch('apache_beam.io.mongodbio.MongoClient') def test_get_range_tracker(self, mock_client): mock_client.return_value = _MockMongoClient(self._docs) if self._ids == OBJECT_IDS: self.assertIsInstance( self.mongo_source.get_range_tracker(None, None), _ObjectIdRangeTracker, ) elif self._ids == INT_IDS: self.assertIsInstance( self.mongo_source.get_range_tracker(None, None), OffsetRangeTracker, ) elif self._ids == STR_IDS_1: self.assertIsInstance( self.mongo_source.get_range_tracker(None, None), LexicographicKeyRangeTracker, ) @mock.patch('apache_beam.io.mongodbio.MongoClient') def test_read(self, mock_client): mock_tracker = mock.MagicMock() test_cases = [ { # range covers the first(inclusive) to third(exclusive) documents 'start': self._ids[0], 'stop': self._ids[2], 'expected': self._docs[0:2] }, { # range covers from the first to the third documents 'start': self.min_id, # smallest possible id 'stop': self._ids[2], 'expected': self._docs[0:2] }, { # range covers from the third to last documents 'start': self._ids[2], 'stop': self.max_id, # largest possible id 'expected': self._docs[2:] }, { # range covers all documents 'start': self.min_id, 'stop': self.max_id, 'expected': self._docs }, { # range doesn't include any document 'start': self._increment_id(self._ids[2], 1), 'stop': self._increment_id(self._ids[3], -1), 'expected': [] }, ] mock_client.return_value = _MockMongoClient(self._docs) for case in test_cases: mock_tracker.start_position.return_value = case['start'] mock_tracker.stop_position.return_value = case['stop'] result = list(self.mongo_source.read(mock_tracker)) self.assertListEqual(case['expected'], result) def test_display_data(self): data = self.mongo_source.display_data() self.assertTrue('database' in data) self.assertTrue('collection' in data) def test_range_is_not_splittable(self): self.assertTrue( self.mongo_source._range_is_not_splittable( _ObjectIdHelper.int_to_id(1), _ObjectIdHelper.int_to_id(1), )) self.assertTrue( self.mongo_source._range_is_not_splittable( _ObjectIdHelper.int_to_id(1), _ObjectIdHelper.int_to_id(2), )) self.assertFalse( self.mongo_source._range_is_not_splittable( _ObjectIdHelper.int_to_id(1), _ObjectIdHelper.int_to_id(3), )) self.assertTrue(self.mongo_source._range_is_not_splittable(0, 0)) self.assertTrue(self.mongo_source._range_is_not_splittable(0, 1)) self.assertFalse(self.mongo_source._range_is_not_splittable(0, 2)) self.assertTrue(self.mongo_source._range_is_not_splittable("AAA", "AAA")) self.assertFalse( self.mongo_source._range_is_not_splittable("AAA", "AAA\x00")) self.assertFalse(self.mongo_source._range_is_not_splittable("AAA", "AAB")) @parameterized_class(('bucket_auto', ), [(False, ), (True, )]) class ReadFromMongoDBTest(unittest.TestCase): @mock.patch('apache_beam.io.mongodbio.MongoClient') def test_read_from_mongodb(self, mock_client): documents = [{ '_id': objectid.ObjectId(), 'x': i, 'selected': 1, 'unselected': 2 } for i in range(3)] mock_client.return_value = _MockMongoClient(documents) projection = ['x', 'selected'] projected_documents = [{ k: v for k, v in e.items() if k in projection or k == '_id' } for e in documents] with TestPipeline() as p: docs = p | 'ReadFromMongoDB' >> ReadFromMongoDB( uri='mongodb://test', db='db', coll='collection', projection=projection, bucket_auto=self.bucket_auto) assert_that(docs, equal_to(projected_documents)) class GenerateObjectIdFnTest(unittest.TestCase): def test_process(self): with TestPipeline() as p: output = ( p | "Create" >> beam.Create([{ 'x': 1 }, { 'x': 2, '_id': 123 }]) | "Generate ID" >> beam.ParDo(_GenerateObjectIdFn()) | "Check" >> beam.Map(lambda x: '_id' in x)) assert_that(output, equal_to([True] * 2)) class WriteMongoFnTest(unittest.TestCase): @mock.patch('apache_beam.io.mongodbio._MongoSink') def test_process(self, mock_sink): docs = [{'x': 1}, {'x': 2}, {'x': 3}] with TestPipeline() as p: _ = ( p | "Create" >> beam.Create(docs) | "Write" >> beam.ParDo(_WriteMongoFn(batch_size=2))) p.run() self.assertEqual( 2, mock_sink.return_value.__enter__.return_value.write.call_count) def test_display_data(self): data = _WriteMongoFn(batch_size=10).display_data() self.assertEqual(10, data['batch_size']) class MongoSinkTest(unittest.TestCase): @mock.patch('apache_beam.io.mongodbio.MongoClient') def test_write(self, mock_client): docs = [{'x': 1}, {'x': 2}, {'x': 3}] _MongoSink(uri='test', db='test', coll='test').write(docs) self.assertTrue( mock_client.return_value.__getitem__.return_value.__getitem__. return_value.bulk_write.called) class WriteToMongoDBTest(unittest.TestCase): @mock.patch('apache_beam.io.mongodbio.MongoClient') def test_write_to_mongodb_with_existing_id(self, mock_client): _id = objectid.ObjectId() docs = [{'x': 1, '_id': _id}] expected_update = [ ReplaceOne({'_id': _id}, { 'x': 1, '_id': _id }, True, None) ] with TestPipeline() as p: _ = ( p | "Create" >> beam.Create(docs) | "Write" >> WriteToMongoDB(db='test', coll='test')) p.run() mock_client.return_value.__getitem__.return_value.__getitem__. \ return_value.bulk_write.assert_called_with(expected_update) @mock.patch('apache_beam.io.mongodbio.MongoClient') def test_write_to_mongodb_with_generated_id(self, mock_client): docs = [{'x': 1}] expected_update = [ ReplaceOne({'_id': mock.ANY}, { 'x': 1, '_id': mock.ANY }, True, None) ] with TestPipeline() as p: _ = ( p | "Create" >> beam.Create(docs) | "Write" >> WriteToMongoDB(db='test', coll='test')) p.run() mock_client.return_value.__getitem__.return_value.__getitem__. \ return_value.bulk_write.assert_called_with(expected_update) class ObjectIdHelperTest(TestCase): def test_conversion(self): test_cases = [ (objectid.ObjectId('000000000000000000000000'), 0), (objectid.ObjectId('000000000000000100000000'), 2**32), (objectid.ObjectId('0000000000000000ffffffff'), 2**32 - 1), (objectid.ObjectId('000000010000000000000000'), 2**64), (objectid.ObjectId('00000000ffffffffffffffff'), 2**64 - 1), (objectid.ObjectId('ffffffffffffffffffffffff'), 2**96 - 1), ] for (_id, number) in test_cases: self.assertEqual(_id, _ObjectIdHelper.int_to_id(number)) self.assertEqual(number, _ObjectIdHelper.id_to_int(_id)) # random tests for _ in range(100): _id = objectid.ObjectId() number = int(_id.binary.hex(), 16) self.assertEqual(_id, _ObjectIdHelper.int_to_id(number)) self.assertEqual(number, _ObjectIdHelper.id_to_int(_id)) def test_increment_id(self): test_cases = [ ( objectid.ObjectId("000000000000000100000000"), objectid.ObjectId("0000000000000000ffffffff"), ), ( objectid.ObjectId("000000010000000000000000"), objectid.ObjectId("00000000ffffffffffffffff"), ), ] for first, second in test_cases: self.assertEqual(second, _ObjectIdHelper.increment_id(first, -1)) self.assertEqual(first, _ObjectIdHelper.increment_id(second, 1)) for _ in range(100): _id = objectid.ObjectId() self.assertLess(_id, _ObjectIdHelper.increment_id(_id, 1)) self.assertGreater(_id, _ObjectIdHelper.increment_id(_id, -1)) class ObjectRangeTrackerTest(TestCase): def test_fraction_position_conversion(self): start_int = 0 stop_int = 2**96 - 1 start = _ObjectIdHelper.int_to_id(start_int) stop = _ObjectIdHelper.int_to_id(stop_int) test_cases = ([start_int, stop_int, 2**32, 2**32 - 1, 2**64, 2**64 - 1] + [random.randint(start_int, stop_int) for _ in range(100)]) tracker = _ObjectIdRangeTracker() for pos in test_cases: _id = _ObjectIdHelper.int_to_id(pos - start_int) desired_fraction = (pos - start_int) / (stop_int - start_int) self.assertAlmostEqual( tracker.position_to_fraction(_id, start, stop), desired_fraction, places=20) convert_id = tracker.fraction_to_position( (pos - start_int) / (stop_int - start_int), start, stop) # due to precision loss, the convert fraction is only gonna be close to # original fraction. convert_fraction = tracker.position_to_fraction(convert_id, start, stop) self.assertGreater(convert_id, start) self.assertLess(convert_id, stop) self.assertAlmostEqual(convert_fraction, desired_fraction, places=20) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/parquetio.py0000644000175100017510000007564715206113032020672 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """``PTransforms`` for reading from and writing to Parquet files. Provides two read ``PTransform``\\s, ``ReadFromParquet`` and ``ReadAllFromParquet``, that produces a ``PCollection`` of records. Each record of this ``PCollection`` will contain a single record read from a Parquet file. Records that are of simple types will be mapped into corresponding Python types. The actual parquet file operations are done by pyarrow. Source splitting is supported at row group granularity. Additionally, this module provides a write ``PTransform`` ``WriteToParquet`` that can be used to write a given ``PCollection`` of Python objects to a Parquet file. """ # pytype: skip-file from functools import partial from typing import Iterator from packaging import version from apache_beam.io import filebasedsink from apache_beam.io import filebasedsource from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystems import FileSystems from apache_beam.io.iobase import RangeTracker from apache_beam.io.iobase import Read from apache_beam.io.iobase import Write from apache_beam.portability.api import schema_pb2 from apache_beam.transforms import DoFn from apache_beam.transforms import ParDo from apache_beam.transforms import PTransform from apache_beam.transforms import window from apache_beam.typehints import schemas from apache_beam.utils.windowed_value import WindowedValue try: import pyarrow as pa paTable = pa.Table import pyarrow.parquet as pq # pylint: disable=ungrouped-imports from apache_beam.typehints import arrow_type_compatibility except ImportError: pa = None paTable = None pq = None ARROW_MAJOR_VERSION = None arrow_type_compatibility = None else: base_pa_version = version.parse(pa.__version__).base_version ARROW_MAJOR_VERSION, _, _ = map(int, base_pa_version.split('.')) __all__ = [ 'ReadFromParquet', 'ReadAllFromParquet', 'ReadFromParquetBatched', 'ReadAllFromParquetBatched', 'WriteToParquet', 'WriteToParquetBatched' ] class _ArrowTableToRowDictionaries(DoFn): """ A DoFn that consumes an Arrow table and yields a python dictionary for each row in the table.""" def process(self, table, with_filename=False): if with_filename: file_name = table[0] table = table[1] num_rows = table.num_rows data_items = table.to_pydict().items() for n in range(num_rows): row = {} for column, values in data_items: row[column] = values[n] if with_filename: yield (file_name, row) else: yield row class _RowDictionariesToArrowTable(DoFn): """ A DoFn that consumes python dictionarys and yields a pyarrow table.""" def __init__( self, schema, row_group_buffer_size=64 * 1024 * 1024, record_batch_size=1000): self._schema = schema self._row_group_buffer_size = row_group_buffer_size self._buffer = [[] for _ in range(len(schema.names))] self._buffer_size = record_batch_size self._record_batches = [] self._record_batches_byte_size = 0 self._window = None def process(self, row, w=DoFn.WindowParam, pane=DoFn.PaneInfoParam): self._window = w if len(self._buffer[0]) >= self._buffer_size: self._flush_buffer() if self._record_batches_byte_size >= self._row_group_buffer_size: table = self._create_table() yield table # reorder the data in columnar format. for i, n in enumerate(self._schema.names): # Handle missing nullable fields by using None as default value field = self._schema.field(i) if field.nullable and n not in row: self._buffer[i].append(None) else: self._buffer[i].append(row[n]) def finish_bundle(self): if len(self._buffer[0]) > 0: self._flush_buffer() if self._record_batches_byte_size > 0: table = self._create_table() if self._window is None or isinstance(self._window, window.GlobalWindow): # bounded input yield window.GlobalWindows.windowed_value_at_end_of_window(table) else: # unbounded input yield WindowedValue( table, timestamp=self._window. end, #or it could be max of timestamp of the rows processed windows=[self._window] # TODO(pabloem) HOW DO WE GET THE PANE ) def display_data(self): res = super().display_data() res['row_group_buffer_size'] = str(self._row_group_buffer_size) res['buffer_size'] = str(self._buffer_size) return res def _create_table(self): table = pa.Table.from_batches(self._record_batches, schema=self._schema) self._record_batches = [] self._record_batches_byte_size = 0 return table def _flush_buffer(self): arrays = [[] for _ in range(len(self._schema.names))] for x, y in enumerate(self._buffer): arrays[x] = pa.array(y, type=self._schema.types[x]) self._buffer[x] = [] rb = pa.RecordBatch.from_arrays(arrays, schema=self._schema) self._record_batches.append(rb) size = 0 for x in arrays: for b in x.buffers(): if b is not None: size = size + b.size self._record_batches_byte_size = self._record_batches_byte_size + size class _ArrowTableToBeamRows(DoFn): def __init__(self, beam_type): self._beam_type = beam_type @DoFn.yields_batches def process(self, element) -> Iterator[paTable]: yield element def infer_output_type(self, input_type): return self._beam_type class _BeamRowsToArrowTable(DoFn): @DoFn.yields_elements def process_batch(self, element: paTable) -> Iterator[paTable]: yield element class ReadFromParquetBatched(PTransform): """A :class:`~apache_beam.transforms.ptransform.PTransform` for reading Parquet files as a `PCollection` of `pyarrow.Table`. This `PTransform` is currently experimental. No backward-compatibility guarantees.""" def __init__( self, file_pattern=None, min_bundle_size=0, validate=True, columns=None): """ Initializes :class:`~ReadFromParquetBatched` An alternative to :class:`~ReadFromParquet` that yields each row group from the Parquet file as a `pyarrow.Table`. These Table instances can be processed directly, or converted to a pandas DataFrame for processing. For more information on supported types and schema, please see the pyarrow documentation. .. testcode:: with beam.Pipeline() as p: dataframes = p \\ | 'Read' >> beam.io.ReadFromParquetBatched('/mypath/mypqfiles*') \\ | 'Convert to pandas' >> beam.Map(lambda table: table.to_pandas()) .. NOTE: We're not actually interested in this error; but if we get here, it means that the way of calling this transform hasn't changed. .. testoutput:: :hide: Traceback (most recent call last): ... OSError: No files found based on the file pattern See also: :class:`~ReadFromParquet`. Args: file_pattern (str): the file glob to read min_bundle_size (int): the minimum size in bytes, to be considered when splitting the input into bundles. validate (bool): flag to verify that the files exist during the pipeline creation time. columns (List[str]): list of columns that will be read from files. A column name may be a prefix of a nested field, e.g. 'a' will select 'a.b', 'a.c', and 'a.d.e' """ super().__init__() self._source = _ParquetSource( file_pattern, min_bundle_size, validate=validate, columns=columns, ) def expand(self, pvalue): return pvalue.pipeline | Read(self._source) def display_data(self): return {'source_dd': self._source} class ReadFromParquet(PTransform): """A `PTransform` for reading Parquet files.""" def __init__( self, file_pattern=None, min_bundle_size=0, validate=True, columns=None, as_rows=False): """Initializes :class:`ReadFromParquet`. Uses source ``_ParquetSource`` to read a set of Parquet files defined by a given file pattern. If ``/mypath/myparquetfiles*`` is a file-pattern that points to a set of Parquet files, a :class:`~apache_beam.pvalue.PCollection` for the records in these Parquet files can be created in the following manner. .. testcode:: with beam.Pipeline() as p: records = p | 'Read' >> beam.io.ReadFromParquet('/mypath/mypqfiles*') .. NOTE: We're not actually interested in this error; but if we get here, it means that the way of calling this transform hasn't changed. .. testoutput:: :hide: Traceback (most recent call last): ... OSError: No files found based on the file pattern Each element of this :class:`~apache_beam.pvalue.PCollection` will contain a Python dictionary representing a single record. The keys will be of type :class:`str` and named after their corresponding column names. The values will be of the type defined in the corresponding Parquet schema. Records that are of simple types will be mapped into corresponding Python types. Records that are of complex types like list and struct will be mapped to Python list and dictionary respectively. For more information on supported types and schema, please see the pyarrow documentation. See also: :class:`~ReadFromParquetBatched`. Args: file_pattern (str): the file glob to read min_bundle_size (int): the minimum size in bytes, to be considered when splitting the input into bundles. validate (bool): flag to verify that the files exist during the pipeline creation time. columns (List[str]): list of columns that will be read from files. A column name may be a prefix of a nested field, e.g. 'a' will select 'a.b', 'a.c', and 'a.d.e' as_rows (bool): whether to output a schema'd PCollection of Beam rows rather than Python dictionaries. """ super().__init__() self._source = _ParquetSource( file_pattern, min_bundle_size, validate=validate, columns=columns, ) if as_rows: if columns is None: filter_schema = lambda schema: schema else: top_level_columns = set(c.split('.')[0] for c in columns) filter_schema = lambda schema: schema_pb2.Schema( fields=[f for f in schema.fields if f.name in top_level_columns]) path = FileSystems.match([file_pattern], [1])[0].metadata_list[0].path with FileSystems.open(path) as fin: self._schema = filter_schema( arrow_type_compatibility.beam_schema_from_arrow_schema( pq.read_schema(fin))) else: self._schema = None def expand(self, pvalue): arrow_batches = pvalue | Read(self._source) if self._schema is None: return arrow_batches | ParDo(_ArrowTableToRowDictionaries()) else: return arrow_batches | ParDo( _ArrowTableToBeamRows(schemas.named_tuple_from_schema(self._schema))) def display_data(self): return {'source_dd': self._source} class ReadAllFromParquetBatched(PTransform): """A ``PTransform`` for reading ``PCollection`` of Parquet files. Uses source ``_ParquetSource`` to read a ``PCollection`` of Parquet files or file patterns and produce a ``PCollection`` of ``pyarrow.Table``, one for each Parquet file row group. This ``PTransform`` is currently experimental. No backward-compatibility guarantees. """ DEFAULT_DESIRED_BUNDLE_SIZE = 64 * 1024 * 1024 # 64MB def __init__( self, min_bundle_size=0, desired_bundle_size=DEFAULT_DESIRED_BUNDLE_SIZE, columns=None, with_filename=False, label='ReadAllFiles'): """Initializes ``ReadAllFromParquet``. Args: min_bundle_size: the minimum size in bytes, to be considered when splitting the input into bundles. desired_bundle_size: the desired size in bytes, to be considered when splitting the input into bundles. columns: list of columns that will be read from files. A column name may be a prefix of a nested field, e.g. 'a' will select 'a.b', 'a.c', and 'a.d.e' with_filename: If True, returns a Key Value with the key being the file name and the value being the actual data. If False, it only returns the data. """ super().__init__() source_from_file = partial( _ParquetSource, min_bundle_size=min_bundle_size, columns=columns) self._read_all_files = filebasedsource.ReadAllFiles( True, CompressionTypes.UNCOMPRESSED, desired_bundle_size, min_bundle_size, source_from_file, with_filename) self.label = label def expand(self, pvalue): return pvalue | self.label >> self._read_all_files class ReadAllFromParquet(PTransform): def __init__(self, with_filename=False, **kwargs): self._with_filename = with_filename self._read_batches = ReadAllFromParquetBatched( with_filename=self._with_filename, **kwargs) def expand(self, pvalue): return pvalue | self._read_batches | ParDo( _ArrowTableToRowDictionaries(), with_filename=self._with_filename) class _ParquetUtils(object): @staticmethod def find_first_row_group_index(pf, start_offset): for i in range(_ParquetUtils.get_number_of_row_groups(pf)): row_group_start_offset = _ParquetUtils.get_offset(pf, i) if row_group_start_offset >= start_offset: return i return -1 @staticmethod def get_offset(pf, row_group_index): first_column_metadata =\ pf.metadata.row_group(row_group_index).column(0) if first_column_metadata.has_dictionary_page: return first_column_metadata.dictionary_page_offset else: return first_column_metadata.data_page_offset @staticmethod def get_number_of_row_groups(pf): return pf.metadata.num_row_groups class _ParquetSource(filebasedsource.FileBasedSource): """A source for reading Parquet files. """ def __init__( self, file_pattern, min_bundle_size=0, validate=False, columns=None): super().__init__( file_pattern=file_pattern, min_bundle_size=min_bundle_size, validate=validate) self._columns = columns def read_records(self, file_name, range_tracker): next_block_start = -1 def split_points_unclaimed(stop_position): if next_block_start >= stop_position: # Next block starts at or after the suggested stop position. Hence # there will not be split points to be claimed for the range ending at # suggested stop position. return 0 return RangeTracker.SPLIT_POINTS_UNKNOWN range_tracker.set_split_points_unclaimed_callback(split_points_unclaimed) start_offset = range_tracker.start_position() if start_offset is None: start_offset = 0 with self.open_file(file_name) as f: pf = pq.ParquetFile(f) # find the first dictionary page (or data page if there's no dictionary # page available) offset after the given start_offset. This offset is also # the starting offset of any row group since the Parquet specification # describes that the data pages always come first before the meta data in # each row group. index = _ParquetUtils.find_first_row_group_index(pf, start_offset) if index != -1: next_block_start = _ParquetUtils.get_offset(pf, index) else: next_block_start = range_tracker.stop_position() number_of_row_groups = _ParquetUtils.get_number_of_row_groups(pf) while range_tracker.try_claim(next_block_start): table = pf.read_row_group(index, self._columns) if index + 1 < number_of_row_groups: index = index + 1 next_block_start = _ParquetUtils.get_offset(pf, index) else: next_block_start = range_tracker.stop_position() yield table _create_parquet_source = _ParquetSource class WriteToParquet(PTransform): """A ``PTransform`` for writing parquet files. """ def __init__( self, file_path_prefix, schema=None, row_group_buffer_size=64 * 1024 * 1024, record_batch_size=1000, codec='none', use_deprecated_int96_timestamps=False, use_compliant_nested_type=False, file_name_suffix='', num_shards=0, shard_name_template=None, mime_type='application/x-parquet', triggering_frequency=None, ): """Initialize a WriteToParquet transform. Writes parquet files from a :class:`~apache_beam.pvalue.PCollection` of records. Each record is a dictionary with keys of a string type that represent column names. Schema must be specified like the example below. .. testsetup:: from tempfile import NamedTemporaryFile import glob import os import pyarrow filename = NamedTemporaryFile(delete=False).name .. testcode:: with beam.Pipeline() as p: records = p | 'Read' >> beam.Create( [{'name': 'foo', 'age': 10}, {'name': 'bar', 'age': 20}] ) _ = records | 'Write' >> beam.io.WriteToParquet(filename, pyarrow.schema( [('name', pyarrow.binary()), ('age', pyarrow.int64())] ) ) .. testcleanup:: for output in glob.glob('{}*'.format(filename)): os.remove(output) For more information on supported types and schema, please see the pyarrow document. Args: file_path_prefix: The file path to write to. The files written will begin with this prefix, followed by a shard identifier (see num_shards), and end in a common extension, if given by file_name_suffix. In most cases, only this argument is specified and num_shards, shard_name_template, and file_name_suffix use default values. schema: The schema to use, as type of ``pyarrow.Schema``. row_group_buffer_size: The byte size of the row group buffer. Note that this size is for uncompressed data on the memory and normally much bigger than the actual row group size written to a file. record_batch_size: The number of records in each record batch. Record batch is a basic unit used for storing data in the row group buffer. A higher record batch size implies low granularity on a row group buffer size. For configuring a row group size based on the number of records, set ``row_group_buffer_size`` to 1 and use ``record_batch_size`` to adjust the value. codec: The codec to use for block-level compression. Any string supported by the pyarrow specification is accepted. use_deprecated_int96_timestamps: Write nanosecond resolution timestamps to INT96 Parquet format. Defaults to False. use_compliant_nested_type: Write compliant Parquet nested type (lists). file_name_suffix: Suffix for the files written. num_shards: The number of files (shards) used for output. If not set, the service will decide on the optimal number of shards. Constraining the number of shards is likely to reduce the performance of a pipeline. Setting this value is not recommended unless you require a specific number of output files. shard_name_template: A template string containing placeholders for the shard number and shard count. Currently only ``''``, ``'-SSSSS-of-NNNNN'``, ``'-W-SSSSS-of-NNNNN'`` and ``'-V-SSSSS-of-NNNNN'`` are patterns accepted by the service. When constructing a filename for a particular shard number, the upper-case letters ``S`` and ``N`` are replaced with the ``0``-padded shard number and shard count respectively. This argument can be ``''`` in which case it behaves as if num_shards was set to 1 and only one file will be generated. The default pattern used is ``'-SSSSS-of-NNNNN'`` for bounded PCollections and for ``'-W-SSSSS-of-NNNNN'`` unbounded PCollections. W is used for windowed shard naming and is replaced with ``[window.start, window.end)`` V is used for windowed shard naming and is replaced with ``[window.start.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S"), window.end.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S")`` mime_type: The MIME type to use for the produced files, if the filesystem supports specifying MIME types. triggering_frequency: (int) Every triggering_frequency duration, a window will be triggered and all bundles in the window will be written. If set it overrides user windowing. Mandatory for GlobalWindow. Returns: A WriteToParquet transform usable for writing. """ super().__init__() self._schema = schema self._row_group_buffer_size = row_group_buffer_size self._record_batch_size = record_batch_size self._sink = \ _create_parquet_sink( file_path_prefix, schema, codec, use_deprecated_int96_timestamps, use_compliant_nested_type, file_name_suffix, num_shards, shard_name_template, mime_type, triggering_frequency ) def expand(self, pcoll): if (not pcoll.is_bounded and self._sink.shard_name_template == filebasedsink.DEFAULT_SHARD_NAME_TEMPLATE): self._sink.shard_name_template = ( filebasedsink.DEFAULT_WINDOW_SHARD_NAME_TEMPLATE) self._sink.shard_name_format = self._sink._template_to_format( self._sink.shard_name_template) self._sink.shard_name_glob_format = self._sink._template_to_glob_format( self._sink.shard_name_template) if self._schema is None: try: beam_schema = schemas.schema_from_element_type(pcoll.element_type) except TypeError as exn: raise ValueError( "A schema is required to write non-schema'd data.") from exn self._sink._schema = ( arrow_type_compatibility.arrow_schema_from_beam_schema(beam_schema)) convert_fn = _BeamRowsToArrowTable() else: convert_fn = _RowDictionariesToArrowTable( self._schema, self._row_group_buffer_size, self._record_batch_size) if pcoll.is_bounded: return pcoll | ParDo(convert_fn) | Write(self._sink) else: self._sink.convert_fn = convert_fn return pcoll | Write(self._sink) def display_data(self): return { 'sink_dd': self._sink, 'row_group_buffer_size': str(self._row_group_buffer_size) } class WriteToParquetBatched(PTransform): """A ``PTransform`` for writing parquet files from a `PCollection` of `pyarrow.Table`. This ``PTransform`` is currently experimental. No backward-compatibility guarantees. """ def __init__( self, file_path_prefix, schema=None, codec='none', use_deprecated_int96_timestamps=False, use_compliant_nested_type=False, file_name_suffix='', num_shards=0, shard_name_template=None, mime_type='application/x-parquet', triggering_frequency=None): """Initialize a WriteToParquetBatched transform. Writes parquet files from a :class:`~apache_beam.pvalue.PCollection` of records. Each record is a pa.Table Schema must be specified like the example below. .. testsetup:: batched from tempfile import NamedTemporaryFile import glob import os import pyarrow filename = NamedTemporaryFile(delete=False).name .. testcode:: batched table = pyarrow.Table.from_pylist([{'name': 'foo', 'age': 10}, {'name': 'bar', 'age': 20}]) with beam.Pipeline() as p: records = p | 'Read' >> beam.Create([table]) _ = records | 'Write' >> beam.io.WriteToParquetBatched(filename, pyarrow.schema( [('name', pyarrow.string()), ('age', pyarrow.int64())] ) ) .. testcleanup:: batched for output in glob.glob('{}*'.format(filename)): os.remove(output) For more information on supported types and schema, please see the pyarrow document. Args: file_path_prefix: The file path to write to. The files written will begin with this prefix, followed by a shard identifier (see num_shards), and end in a common extension, if given by file_name_suffix. In most cases, only this argument is specified and num_shards, shard_name_template, and file_name_suffix use default values. schema: The schema to use, as type of ``pyarrow.Schema``. codec: The codec to use for block-level compression. Any string supported by the pyarrow specification is accepted. use_deprecated_int96_timestamps: Write nanosecond resolution timestamps to INT96 Parquet format. Defaults to False. use_compliant_nested_type: Write compliant Parquet nested type (lists). file_name_suffix: Suffix for the files written. num_shards: The number of files (shards) used for output. If not set, the service will decide on the optimal number of shards. Constraining the number of shards is likely to reduce the performance of a pipeline. Setting this value is not recommended unless you require a specific number of output files. shard_name_template: A template string containing placeholders for the shard number and shard count. When constructing a filename for a particular shard number, the upper-case letters 'S' and 'N' are replaced with the 0-padded shard number and shard count respectively. W is used for windowed shard naming and is replaced with ``[window.start, window.end)`` V is used for windowed shard naming and is replaced with ``[window.start.to_utc_datetime().isoformat(), window.end.to_utc_datetime().isoformat()`` This argument can be '' in which case it behaves as if num_shards was set to 1 and only one file will be generated. The default pattern used is '-SSSSS-of-NNNNN' if None is passed as the shard_name_template and the PCollection is bounded. The default pattern used is '-W-SSSSS-of-NNNNN' if None is passed as the shard_name_template and the PCollection is unbounded. mime_type: The MIME type to use for the produced files, if the filesystem supports specifying MIME types. triggering_frequency: (int) Every triggering_frequency duration, a window will be triggered and all bundles in the window will be written. Returns: A WriteToParquetBatched transform usable for writing. """ super().__init__() self._sink = \ _create_parquet_sink( file_path_prefix, schema, codec, use_deprecated_int96_timestamps, use_compliant_nested_type, file_name_suffix, num_shards, shard_name_template, mime_type, triggering_frequency ) def expand(self, pcoll): if (not pcoll.is_bounded and self._sink.shard_name_template == filebasedsink.DEFAULT_SHARD_NAME_TEMPLATE): self._sink.shard_name_template = ( filebasedsink.DEFAULT_WINDOW_SHARD_NAME_TEMPLATE) self._sink.shard_name_format = self._sink._template_to_format( self._sink.shard_name_template) self._sink.shard_name_glob_format = self._sink._template_to_glob_format( self._sink.shard_name_template) return pcoll | Write(self._sink) def display_data(self): return {'sink_dd': self._sink} def _create_parquet_sink( file_path_prefix, schema, codec, use_deprecated_int96_timestamps, use_compliant_nested_type, file_name_suffix, num_shards, shard_name_template, mime_type, triggering_frequency=60): return \ _ParquetSink( file_path_prefix, schema, codec, use_deprecated_int96_timestamps, use_compliant_nested_type, file_name_suffix, num_shards, shard_name_template, mime_type, triggering_frequency ) class _ParquetSink(filebasedsink.FileBasedSink): """A sink for parquet files from batches.""" def __init__( self, file_path_prefix, schema, codec, use_deprecated_int96_timestamps, use_compliant_nested_type, file_name_suffix, num_shards, shard_name_template, mime_type, triggering_frequency): super().__init__( file_path_prefix, file_name_suffix=file_name_suffix, num_shards=num_shards, shard_name_template=shard_name_template, coder=None, mime_type=mime_type, # Compression happens at the block level using the supplied codec, and # not at the file level. compression_type=CompressionTypes.UNCOMPRESSED, triggering_frequency=triggering_frequency) self._schema = schema self._codec = codec if ARROW_MAJOR_VERSION == 1 and self._codec.lower() == "lz4": raise ValueError( "Due to ARROW-9424, writing with LZ4 compression is not supported in " "pyarrow 1.x, please use a different pyarrow version or a different " f"codec. Your pyarrow version: {pa.__version__}") self._use_deprecated_int96_timestamps = use_deprecated_int96_timestamps if use_compliant_nested_type and ARROW_MAJOR_VERSION < 4: raise ValueError( "With ARROW-11497, use_compliant_nested_type is only supported in " "pyarrow version >= 4.x, please use a different pyarrow version. " f"Your pyarrow version: {pa.__version__}") self._use_compliant_nested_type = use_compliant_nested_type self._file_handle = None def open(self, temp_path): self._file_handle = super().open(temp_path) if ARROW_MAJOR_VERSION < 4: return pq.ParquetWriter( self._file_handle, self._schema, compression=self._codec, use_deprecated_int96_timestamps=self._use_deprecated_int96_timestamps) return pq.ParquetWriter( self._file_handle, self._schema, compression=self._codec, use_deprecated_int96_timestamps=self._use_deprecated_int96_timestamps, use_compliant_nested_type=self._use_compliant_nested_type) def write_record(self, writer, table: paTable): writer.write_table(table) def close(self, writer): writer.close() if self._file_handle: self._file_handle.close() self._file_handle = None def display_data(self): res = super().display_data() res['codec'] = str(self._codec) res['schema'] = str(self._schema) return res ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/parquetio_it_test.py0000644000175100017510000001412115206113032022401 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import string import unittest import uuid from collections import Counter from datetime import datetime import pytest import pytz import apache_beam as beam from apache_beam import Create from apache_beam import DoFn from apache_beam import FlatMap from apache_beam import Flatten from apache_beam import Map from apache_beam import ParDo from apache_beam import Reshuffle from apache_beam.io.filesystems import FileSystems from apache_beam.io.parquetio import ReadAllFromParquet from apache_beam.io.parquetio import WriteToParquet from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import BeamAssertException from apache_beam.transforms import CombineGlobally from apache_beam.transforms.combiners import Count from apache_beam.transforms.periodicsequence import PeriodicImpulse try: import pyarrow as pa except ImportError: pa = None @unittest.skipIf(pa is None, "PyArrow is not installed.") class TestParquetIT(unittest.TestCase): def setUp(self): pass def tearDown(self): pass @pytest.mark.it_postcommit def test_parquetio_it(self): file_prefix = "parquet_it_test" init_size = 10 data_size = 20000 with TestPipeline(is_integration_test=True) as p: pcol = self._generate_data(p, file_prefix, init_size, data_size) self._verify_data(pcol, init_size, data_size) @staticmethod def _sum_verifier(init_size, data_size, x): expected = sum(range(data_size)) * init_size if x != expected: raise BeamAssertException( "incorrect sum: expected(%d) actual(%d)" % (expected, x)) return [] @staticmethod def _count_verifier(init_size, data_size, x): name, count = x[0].decode('utf-8'), x[1] counter = Counter( [string.ascii_uppercase[x % 26] for x in range(0, data_size * 4, 4)]) expected_count = counter[name[0]] * init_size if count != expected_count: raise BeamAssertException( "incorrect count(%s): expected(%d) actual(%d)" % (name, expected_count, count)) return [] def _verify_data(self, pcol, init_size, data_size): read = pcol | 'read' >> ReadAllFromParquet() v1 = ( read | 'get_number' >> Map(lambda x: x['number']) | 'sum_globally' >> CombineGlobally(sum) | 'validate_number' >> FlatMap(lambda x: TestParquetIT._sum_verifier(init_size, data_size, x))) v2 = ( read | 'make_pair' >> Map(lambda x: (x['name'], x['number'])) | 'count_per_key' >> Count.PerKey() | 'validate_name' >> FlatMap( lambda x: TestParquetIT._count_verifier(init_size, data_size, x))) _ = ((v1, v2, pcol) | 'flatten' >> Flatten() | 'reshuffle' >> Reshuffle() | 'cleanup' >> Map(lambda x: FileSystems.delete([x]))) def _generate_data(self, p, output_prefix, init_size, data_size): init_data = [x for x in range(init_size)] lines = ( p | 'create' >> Create(init_data) | 'produce' >> ParDo(ProducerFn(data_size))) schema = pa.schema([('name', pa.binary()), ('number', pa.int64())]) files = lines | 'write' >> WriteToParquet( output_prefix, schema, codec='snappy', file_name_suffix='.parquet') return files class ProducerFn(DoFn): def __init__(self, number): super().__init__() self._number = number self._string_index = 0 self._number_index = 0 def process(self, element): self._string_index = 0 self._number_index = 0 for _ in range(self._number): yield {'name': self.get_string(4), 'number': self.get_int()} def get_string(self, length): s = [] for _ in range(length): s.append(string.ascii_uppercase[self._string_index]) self._string_index = (self._string_index + 1) % 26 return ''.join(s) def get_int(self): i = self._number_index self._number_index = self._number_index + 1 return i @unittest.skipIf(pa is None, "PyArrow is not installed.") class WriteStreamingIT(unittest.TestCase): def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.runner_name = type(self.test_pipeline.runner).__name__ super().setUp() def test_write_streaming_2_shards_default_shard_name_template( self, num_shards=2): args = self.test_pipeline.get_full_options_as_args(streaming=True) unique_id = str(uuid.uuid4()) output_file = f'gs://apache-beam-testing-integration-testing/iobase/test-{unique_id}' # pylint: disable=line-too-long p = beam.Pipeline(argv=args) pyschema = pa.schema([('age', pa.int64())]) _ = ( p | "generate impulse" >> PeriodicImpulse( start_timestamp=datetime(2021, 3, 1, 0, 0, 1, 0, tzinfo=pytz.UTC).timestamp(), stop_timestamp=datetime(2021, 3, 1, 0, 0, 20, 0, tzinfo=pytz.UTC).timestamp(), fire_interval=1) | "generate data" >> beam.Map(lambda t: {'age': t * 10}) | 'WriteToParquet' >> beam.io.WriteToParquet( file_path_prefix=output_file, file_name_suffix=".parquet", num_shards=num_shards, triggering_frequency=60, schema=pyschema)) result = p.run() result.wait_until_finish(duration=600 * 1000) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/parquetio_test.py0000644000175100017510000011342515206113032021714 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import glob import json import logging import os import re import shutil import tempfile import unittest from datetime import datetime from tempfile import TemporaryDirectory import hamcrest as hc import pandas import pytest import pytz from parameterized import param from parameterized import parameterized import apache_beam as beam from apache_beam import Create from apache_beam import Map from apache_beam.io import filebasedsource from apache_beam.io import source_test_utils from apache_beam.io.iobase import RangeTracker from apache_beam.io.parquetio import ReadAllFromParquet from apache_beam.io.parquetio import ReadAllFromParquetBatched from apache_beam.io.parquetio import ReadFromParquet from apache_beam.io.parquetio import ReadFromParquetBatched from apache_beam.io.parquetio import WriteToParquet from apache_beam.io.parquetio import WriteToParquetBatched from apache_beam.io.parquetio import _create_parquet_sink from apache_beam.io.parquetio import _create_parquet_source from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.display import DisplayData from apache_beam.transforms.display_test import DisplayDataItemMatcher from apache_beam.transforms.util import LogElements try: import pyarrow as pa import pyarrow.parquet as pq ARROW_MAJOR_VERSION, _, _ = map(int, pa.__version__.split('.')) except ImportError: pa = None pq = None ARROW_MAJOR_VERSION = 0 @unittest.skipIf(pa is None, "PyArrow is not installed.") @pytest.mark.uses_pyarrow class TestParquet(unittest.TestCase): def setUp(self): # Reducing the size of thread pools. Without this test execution may fail in # environments with limited amount of resources. filebasedsource.MAX_NUM_THREADS_FOR_SIZE_ESTIMATION = 2 self.temp_dir = tempfile.mkdtemp() self.RECORDS = [{ 'name': 'Thomas', 'favorite_number': 1, 'favorite_color': 'blue' }, { 'name': 'Henry', 'favorite_number': 3, 'favorite_color': 'green' }, { 'name': 'Toby', 'favorite_number': 7, 'favorite_color': 'brown' }, { 'name': 'Gordon', 'favorite_number': 4, 'favorite_color': 'blue' }, { 'name': 'Emily', 'favorite_number': -1, 'favorite_color': 'Red' }, { 'name': 'Percy', 'favorite_number': 6, 'favorite_color': 'Green' }, { 'name': 'Peter', 'favorite_number': 3, 'favorite_color': None }] self.SCHEMA = pa.schema([('name', pa.string(), False), ('favorite_number', pa.int64(), False), ('favorite_color', pa.string())]) self.SCHEMA96 = pa.schema([('name', pa.string(), False), ('favorite_number', pa.timestamp('ns'), False), ('favorite_color', pa.string())]) self.RECORDS_NESTED = [{ 'items': [ { 'name': 'Thomas', 'favorite_number': 1, 'favorite_color': 'blue' }, { 'name': 'Henry', 'favorite_number': 3, 'favorite_color': 'green' }, ] }, { 'items': [ { 'name': 'Toby', 'favorite_number': 7, 'favorite_color': 'brown' }, ] }] self.SCHEMA_NESTED = pa.schema([( 'items', pa.list_( pa.struct([('name', pa.string(), False), ('favorite_number', pa.int64(), False), ('favorite_color', pa.string())])))]) def tearDown(self): shutil.rmtree(self.temp_dir) def _record_to_columns(self, records, schema): col_list = [] for n in schema.names: column = [] for r in records: column.append(r[n]) col_list.append(column) return col_list def _records_as_arrow(self, schema=None, count=None): if schema is None: schema = self.SCHEMA if count is None: count = len(self.RECORDS) len_records = len(self.RECORDS) data = [] for i in range(count): data.append(self.RECORDS[i % len_records]) col_data = self._record_to_columns(data, schema) col_array = [pa.array(c, schema.types[cn]) for cn, c in enumerate(col_data)] return pa.Table.from_arrays(col_array, schema=schema) def _write_data( self, directory=None, schema=None, prefix=tempfile.template, row_group_size=1000, codec='none', count=None): if directory is None: directory = self.temp_dir with tempfile.NamedTemporaryFile(delete=False, dir=directory, prefix=prefix) as f: table = self._records_as_arrow(schema, count) pq.write_table( table, f, row_group_size=row_group_size, compression=codec, use_deprecated_int96_timestamps=True) return f.name def _write_pattern(self, num_files, with_filename=False): assert num_files > 0 temp_dir = tempfile.mkdtemp(dir=self.temp_dir) file_list = [] for _ in range(num_files): file_list.append(self._write_data(directory=temp_dir, prefix='mytemp')) if with_filename: return (temp_dir + os.path.sep + 'mytemp*', file_list) return temp_dir + os.path.sep + 'mytemp*' def _run_parquet_test( self, pattern, columns, desired_bundle_size, perform_splitting, expected_result): source = _create_parquet_source(pattern, columns=columns) if perform_splitting: assert desired_bundle_size sources_info = [ (split.source, split.start_position, split.stop_position) for split in source.split(desired_bundle_size=desired_bundle_size) ] if len(sources_info) < 2: raise ValueError( 'Test is trivial. Please adjust it so that at least ' 'two splits get generated') source_test_utils.assert_sources_equal_reference_source( (source, None, None), sources_info) else: read_records = source_test_utils.read_from_source(source, None, None) self.assertCountEqual(expected_result, read_records) def test_read_without_splitting(self): file_name = self._write_data() expected_result = [self._records_as_arrow()] self._run_parquet_test(file_name, None, None, False, expected_result) def test_read_with_splitting(self): file_name = self._write_data() expected_result = [self._records_as_arrow()] self._run_parquet_test(file_name, None, 100, True, expected_result) def test_source_display_data(self): file_name = 'some_parquet_source' source = \ _create_parquet_source( file_name, validate=False ) dd = DisplayData.create_from(source) expected_items = [ DisplayDataItemMatcher('compression', 'auto'), DisplayDataItemMatcher('file_pattern', file_name) ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_read_display_data(self): file_name = 'some_parquet_source' read = \ ReadFromParquet( file_name, validate=False) read_batched = \ ReadFromParquetBatched( file_name, validate=False) expected_items = [ DisplayDataItemMatcher('compression', 'auto'), DisplayDataItemMatcher('file_pattern', file_name) ] hc.assert_that( DisplayData.create_from(read).items, hc.contains_inanyorder(*expected_items)) hc.assert_that( DisplayData.create_from(read_batched).items, hc.contains_inanyorder(*expected_items)) def test_sink_display_data(self): file_name = 'some_parquet_sink' sink = _create_parquet_sink( file_name, self.SCHEMA, 'none', False, False, '.end', 0, None, 'application/x-parquet') dd = DisplayData.create_from(sink) expected_items = [ DisplayDataItemMatcher('schema', str(self.SCHEMA)), DisplayDataItemMatcher( 'file_pattern', 'some_parquet_sink-%(shard_num)05d-of-%(num_shards)05d.end'), DisplayDataItemMatcher('codec', 'none'), DisplayDataItemMatcher('compression', 'uncompressed') ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_write_display_data(self): file_name = 'some_parquet_sink' write = WriteToParquet(file_name, self.SCHEMA) dd = DisplayData.create_from(write) expected_items = [ DisplayDataItemMatcher('codec', 'none'), DisplayDataItemMatcher('schema', str(self.SCHEMA)), DisplayDataItemMatcher('row_group_buffer_size', str(64 * 1024 * 1024)), DisplayDataItemMatcher( 'file_pattern', 'some_parquet_sink-%(shard_num)05d-of-%(num_shards)05d'), DisplayDataItemMatcher('compression', 'uncompressed') ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_write_batched_display_data(self): file_name = 'some_parquet_sink' write = WriteToParquetBatched(file_name, self.SCHEMA) dd = DisplayData.create_from(write) expected_items = [ DisplayDataItemMatcher('codec', 'none'), DisplayDataItemMatcher('schema', str(self.SCHEMA)), DisplayDataItemMatcher( 'file_pattern', 'some_parquet_sink-%(shard_num)05d-of-%(num_shards)05d'), DisplayDataItemMatcher('compression', 'uncompressed') ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) @unittest.skipIf( ARROW_MAJOR_VERSION >= 13, 'pyarrow 13.x and above does not throw ArrowInvalid error') def test_sink_transform_int96(self): with self.assertRaisesRegex(Exception, 'would lose data'): # Should throw an error "ArrowInvalid: Casting from timestamp[ns] to # timestamp[us] would lose data" dst = tempfile.NamedTemporaryFile() path = dst.name with TestPipeline() as p: _ = p \ | Create(self.RECORDS) \ | WriteToParquet( path, self.SCHEMA96, num_shards=1, shard_name_template='') def test_sink_transform(self): with TemporaryDirectory() as tmp_dirname: path = os.path.join(tmp_dirname + "tmp_filename") with TestPipeline() as p: _ = p \ | Create(self.RECORDS) \ | WriteToParquet( path, self.SCHEMA, num_shards=1, shard_name_template='') with TestPipeline() as p: # json used for stable sortability readback = \ p \ | ReadFromParquet(path) \ | Map(json.dumps) assert_that(readback, equal_to([json.dumps(r) for r in self.RECORDS])) def test_sink_transform_batched(self): with TemporaryDirectory() as tmp_dirname: path = os.path.join(tmp_dirname + "tmp_filename") with TestPipeline() as p: _ = p \ | Create([self._records_as_arrow()]) \ | WriteToParquetBatched( path, self.SCHEMA, num_shards=1, shard_name_template='') with TestPipeline() as p: # json used for stable sortability readback = \ p \ | ReadFromParquet(path) \ | Map(json.dumps) assert_that(readback, equal_to([json.dumps(r) for r in self.RECORDS])) def test_sink_transform_compliant_nested_type(self): if ARROW_MAJOR_VERSION < 4: return unittest.skip( 'Writing with compliant nested type is only ' 'supported in pyarrow 4.x and above') with TemporaryDirectory() as tmp_dirname: path = os.path.join(tmp_dirname + 'tmp_filename') with TestPipeline() as p: _ = p \ | Create(self.RECORDS_NESTED) \ | WriteToParquet( path, self.SCHEMA_NESTED, num_shards=1, shard_name_template='', use_compliant_nested_type=True) with TestPipeline() as p: # json used for stable sortability readback = \ p \ | ReadFromParquet(path) \ | Map(json.dumps) assert_that( readback, equal_to([json.dumps(r) for r in self.RECORDS_NESTED])) def test_schema_read_write(self): with TemporaryDirectory() as tmp_dirname: path = os.path.join(tmp_dirname, 'tmp_filename') rows = [beam.Row(a=1, b='x'), beam.Row(a=2, b='y')] stable_repr = lambda row: json.dumps(row._asdict()) with TestPipeline() as p: _ = p | Create(rows) | WriteToParquet(path) with TestPipeline() as p: readback = ( p | ReadFromParquet(path + '*', as_rows=True) | Map(stable_repr)) assert_that(readback, equal_to([stable_repr(r) for r in rows])) def test_write_with_nullable_fields_missing_data(self): """Test WriteToParquet with nullable fields where some fields are missing. This test addresses the bug reported in: https://github.com/apache/beam/issues/35791 where WriteToParquet fails with a KeyError if any nullable field is missing in the data. """ # Define PyArrow schema with all fields nullable schema = pa.schema([ pa.field("id", pa.int64(), nullable=True), pa.field("name", pa.string(), nullable=True), pa.field("age", pa.int64(), nullable=True), pa.field("email", pa.string(), nullable=True), ]) # Sample data with missing nullable fields data = [ { 'id': 1, 'name': 'Alice', 'age': 30 }, # missing 'email' { 'id': 2, 'name': 'Bob', 'age': 25, 'email': 'bob@example.com' }, # all fields present { 'id': 3, 'name': 'Charlie', 'age': None, 'email': None }, # explicit None values { 'id': 4, 'name': 'David' }, # missing 'age' and 'email' ] with TemporaryDirectory() as tmp_dirname: path = os.path.join(tmp_dirname, 'nullable_test') # Write data with missing nullable fields - this should not raise KeyError with TestPipeline() as p: _ = ( p | Create(data) | WriteToParquet( path, schema, num_shards=1, shard_name_template='')) # Read back and verify the data with TestPipeline() as p: readback = ( p | ReadFromParquet(path + '*') | Map(json.dumps, sort_keys=True)) # Expected data should have None for missing nullable fields expected_data = [ { 'id': 1, 'name': 'Alice', 'age': 30, 'email': None }, { 'id': 2, 'name': 'Bob', 'age': 25, 'email': 'bob@example.com' }, { 'id': 3, 'name': 'Charlie', 'age': None, 'email': None }, { 'id': 4, 'name': 'David', 'age': None, 'email': None }, ] assert_that( readback, equal_to([json.dumps(r, sort_keys=True) for r in expected_data])) def test_batched_read(self): with TemporaryDirectory() as tmp_dirname: path = os.path.join(tmp_dirname + "tmp_filename") with TestPipeline() as p: _ = p \ | Create(self.RECORDS, reshuffle=False) \ | WriteToParquet( path, self.SCHEMA, num_shards=1, shard_name_template='') with TestPipeline() as p: # json used for stable sortability readback = \ p \ | ReadFromParquetBatched(path) assert_that(readback, equal_to([self._records_as_arrow()])) @parameterized.expand([ param(compression_type='snappy'), param(compression_type='gzip'), param(compression_type='brotli'), param(compression_type='lz4'), param(compression_type='zstd') ]) def test_sink_transform_compressed(self, compression_type): if compression_type == 'lz4' and ARROW_MAJOR_VERSION == 1: return unittest.skip( "Writing with LZ4 compression is not supported in " "pyarrow 1.x") with TemporaryDirectory() as tmp_dirname: path = os.path.join(tmp_dirname + "tmp_filename") with TestPipeline() as p: _ = p \ | Create(self.RECORDS) \ | WriteToParquet( path, self.SCHEMA, codec=compression_type, num_shards=1, shard_name_template='') with TestPipeline() as p: # json used for stable sortability readback = \ p \ | ReadFromParquet(path + '*') \ | Map(json.dumps) assert_that(readback, equal_to([json.dumps(r) for r in self.RECORDS])) def test_read_reentrant(self): file_name = self._write_data(count=6, row_group_size=3) source = _create_parquet_source(file_name) source_test_utils.assert_reentrant_reads_succeed((source, None, None)) def test_read_without_splitting_multiple_row_group(self): file_name = self._write_data(count=12000, row_group_size=1000) # We expect 12000 elements, split into batches of 1000 elements. Create # a list of pa.Table instances to model this expecation expected_result = [ pa.Table.from_batches([batch]) for batch in self._records_as_arrow( count=12000).to_batches(max_chunksize=1000) ] self._run_parquet_test(file_name, None, None, False, expected_result) def test_read_with_splitting_multiple_row_group(self): file_name = self._write_data(count=12000, row_group_size=1000) # We expect 12000 elements, split into batches of 1000 elements. Create # a list of pa.Table instances to model this expecation expected_result = [ pa.Table.from_batches([batch]) for batch in self._records_as_arrow( count=12000).to_batches(max_chunksize=1000) ] self._run_parquet_test(file_name, None, 10000, True, expected_result) def test_dynamic_work_rebalancing(self): # This test depends on count being sufficiently large + the ratio of # count to row_group_size also being sufficiently large (but the required # ratio to pass varies for values of row_group_size and, somehow, the # version of pyarrow being tested against.) file_name = self._write_data(count=320, row_group_size=20) source = _create_parquet_source(file_name) splits = [split for split in source.split(desired_bundle_size=float('inf'))] assert len(splits) == 1 source_test_utils.assert_split_at_fraction_exhaustive( splits[0].source, splits[0].start_position, splits[0].stop_position) def test_min_bundle_size(self): file_name = self._write_data(count=120, row_group_size=20) source = _create_parquet_source( file_name, min_bundle_size=100 * 1024 * 1024) splits = [split for split in source.split(desired_bundle_size=1)] self.assertEqual(len(splits), 1) source = _create_parquet_source(file_name, min_bundle_size=0) splits = [split for split in source.split(desired_bundle_size=1)] self.assertNotEqual(len(splits), 1) def _convert_to_timestamped_record(self, record): timestamped_record = record.copy() timestamped_record['favorite_number'] =\ pandas.Timestamp(timestamped_record['favorite_number']) return timestamped_record def test_int96_type_conversion(self): file_name = self._write_data( count=120, row_group_size=20, schema=self.SCHEMA96) orig = self._records_as_arrow(count=120, schema=self.SCHEMA96) expected_result = [ pa.Table.from_batches([batch], schema=self.SCHEMA96) for batch in orig.to_batches(max_chunksize=20) ] self._run_parquet_test(file_name, None, None, False, expected_result) def test_split_points(self): file_name = self._write_data(count=12000, row_group_size=3000) source = _create_parquet_source(file_name) splits = [split for split in source.split(desired_bundle_size=float('inf'))] assert len(splits) == 1 range_tracker = splits[0].source.get_range_tracker( splits[0].start_position, splits[0].stop_position) split_points_report = [] for _ in splits[0].source.read(range_tracker): split_points_report.append(range_tracker.split_points()) # There are a total of four row groups. Each row group has 3000 records. # When reading records of the first group, range_tracker.split_points() # should return (0, iobase.RangeTracker.SPLIT_POINTS_UNKNOWN) self.assertEqual( split_points_report, [ (0, RangeTracker.SPLIT_POINTS_UNKNOWN), (1, RangeTracker.SPLIT_POINTS_UNKNOWN), (2, RangeTracker.SPLIT_POINTS_UNKNOWN), (3, 1), ]) def test_selective_columns(self): file_name = self._write_data() orig = self._records_as_arrow() name_column = self.SCHEMA.field('name') expected_result = [ pa.Table.from_arrays( [orig.column('name')], schema=pa.schema([('name', name_column.type, name_column.nullable) ])) ] self._run_parquet_test(file_name, ['name'], None, False, expected_result) def test_sink_transform_multiple_row_group(self): with TemporaryDirectory() as tmp_dirname: path = os.path.join(tmp_dirname + "tmp_filename") # Pin to FnApiRunner since test assumes fixed bundle size with TestPipeline('FnApiRunner') as p: # writing 623200 bytes of data _ = p \ | Create(self.RECORDS * 4000) \ | WriteToParquet( path, self.SCHEMA, num_shards=1, codec='none', shard_name_template='', row_group_buffer_size=250000) self.assertEqual(pq.read_metadata(path).num_row_groups, 3) def test_read_all_from_parquet_single_file(self): path = self._write_data() with TestPipeline() as p: assert_that( p \ | Create([path]) \ | ReadAllFromParquet(), equal_to(self.RECORDS)) with TestPipeline() as p: assert_that( p \ | Create([path]) \ | ReadAllFromParquetBatched(), equal_to([self._records_as_arrow()])) def test_read_all_from_parquet_many_single_files(self): path1 = self._write_data() path2 = self._write_data() path3 = self._write_data() with TestPipeline() as p: assert_that( p \ | Create([path1, path2, path3]) \ | ReadAllFromParquet(), equal_to(self.RECORDS * 3)) with TestPipeline() as p: assert_that( p \ | Create([path1, path2, path3]) \ | ReadAllFromParquetBatched(), equal_to([self._records_as_arrow()] * 3)) def test_read_all_from_parquet_file_pattern(self): file_pattern = self._write_pattern(5) with TestPipeline() as p: assert_that( p \ | Create([file_pattern]) \ | ReadAllFromParquet(), equal_to(self.RECORDS * 5)) with TestPipeline() as p: assert_that( p \ | Create([file_pattern]) \ | ReadAllFromParquetBatched(), equal_to([self._records_as_arrow()] * 5)) def test_read_all_from_parquet_many_file_patterns(self): file_pattern1 = self._write_pattern(5) file_pattern2 = self._write_pattern(2) file_pattern3 = self._write_pattern(3) with TestPipeline() as p: assert_that( p \ | Create([file_pattern1, file_pattern2, file_pattern3]) \ | ReadAllFromParquet(), equal_to(self.RECORDS * 10)) with TestPipeline() as p: assert_that( p \ | Create([file_pattern1, file_pattern2, file_pattern3]) \ | ReadAllFromParquetBatched(), equal_to([self._records_as_arrow()] * 10)) def test_read_all_from_parquet_with_filename(self): file_pattern, file_paths = self._write_pattern(3, with_filename=True) result = [(path, record) for path in file_paths for record in self.RECORDS] with TestPipeline() as p: assert_that( p \ | Create([file_pattern]) \ | ReadAllFromParquet(with_filename=True), equal_to(result)) class GenerateEvent(beam.PTransform): @staticmethod def sample_data(): return GenerateEvent() def expand(self, input): elemlist = [{'age': 10}, {'age': 20}, {'age': 30}] elem = elemlist return ( input | TestStream().add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 1, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 2, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 3, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 4, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 5, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 5, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 6, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 7, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 8, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 9, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 10, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 10, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 11, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 12, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 13, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 14, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 15, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 15, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 16, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 17, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 18, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 19, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 20, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 20, 0, tzinfo=pytz.UTC).timestamp()).advance_watermark_to( datetime( 2021, 3, 1, 0, 0, 25, 0, tzinfo=pytz.UTC). timestamp()).advance_watermark_to_infinity()) class WriteStreamingTest(unittest.TestCase): def setUp(self): super().setUp() self.tempdir = tempfile.mkdtemp() def tearDown(self): if os.path.exists(self.tempdir): shutil.rmtree(self.tempdir) def test_write_streaming_2_shards_default_shard_name_template( self, num_shards=2): with TestPipeline() as p: output = (p | GenerateEvent.sample_data()) #ParquetIO pyschema = pa.schema([('age', pa.int64())]) output2 = output | 'WriteToParquet' >> beam.io.WriteToParquet( file_path_prefix=self.tempdir + "/ouput_WriteToParquet", file_name_suffix=".parquet", num_shards=num_shards, triggering_frequency=60, schema=pyschema) _ = output2 | 'LogElements after WriteToParquet' >> LogElements( prefix='after WriteToParquet ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToParquet-[1614556800.0, 1614556805.0)-00000-of-00002.parquet # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P[\d\.]+), ' r'(?P[\d\.]+|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.parquet$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToParquet*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) self.assertEqual( len(file_names), num_shards, "expected %d files, but got: %d" % (num_shards, len(file_names))) def test_write_streaming_2_shards_custom_shard_name_template( self, num_shards=2, shard_name_template='-V-SSSSS-of-NNNNN'): with TestPipeline() as p: output = (p | GenerateEvent.sample_data()) #ParquetIO pyschema = pa.schema([('age', pa.int64())]) output2 = output | 'WriteToParquet' >> beam.io.WriteToParquet( file_path_prefix=self.tempdir + "/ouput_WriteToParquet", file_name_suffix=".parquet", shard_name_template=shard_name_template, num_shards=num_shards, triggering_frequency=60, schema=pyschema) _ = output2 | 'LogElements after WriteToParquet' >> LogElements( prefix='after WriteToParquet ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToParquet-[2021-03-01T00-00-00, 2021-03-01T00-01-00)- # 00000-of-00002.parquet # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}), ' r'(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.parquet$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToParquet*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) self.assertEqual( len(file_names), num_shards, "expected %d files, but got: %d" % (num_shards, len(file_names))) def test_write_streaming_2_shards_custom_shard_name_template_5s_window( self, num_shards=2, shard_name_template='-V-SSSSS-of-NNNNN', triggering_frequency=5): with TestPipeline() as p: output = (p | GenerateEvent.sample_data()) #ParquetIO pyschema = pa.schema([('age', pa.int64())]) output2 = output | 'WriteToParquet' >> beam.io.WriteToParquet( file_path_prefix=self.tempdir + "/ouput_WriteToParquet", file_name_suffix=".parquet", shard_name_template=shard_name_template, num_shards=num_shards, triggering_frequency=triggering_frequency, schema=pyschema) _ = output2 | 'LogElements after WriteToParquet' >> LogElements( prefix='after WriteToParquet ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToParquet-[2021-03-01T00-00-00, 2021-03-01T00-01-00)- # 00000-of-00002.parquet # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}), ' r'(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.parquet$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToParquet*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) # for 5s window size, the input should be processed by 5 windows with # 2 shards per window self.assertEqual( len(file_names), 10, "expected %d files, but got: %d" % (num_shards, len(file_names))) def test_write_streaming_undef_shards_default_shard_name_template_windowed_pcoll( # pylint: disable=line-too-long self): with TestPipeline() as p: output = ( p | GenerateEvent.sample_data() | 'User windowing' >> beam.transforms.core.WindowInto( beam.transforms.window.FixedWindows(10), trigger=beam.transforms.trigger.AfterWatermark(), accumulation_mode=beam.transforms.trigger.AccumulationMode. DISCARDING, allowed_lateness=beam.utils.timestamp.Duration(seconds=0))) #ParquetIO pyschema = pa.schema([('age', pa.int64())]) output2 = output | 'WriteToParquet' >> beam.io.WriteToParquet( file_path_prefix=self.tempdir + "/ouput_WriteToParquet", file_name_suffix=".parquet", num_shards=0, schema=pyschema) _ = output2 | 'LogElements after WriteToParquet' >> LogElements( prefix='after WriteToParquet ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToParquet-[1614556800.0, 1614556805.0)-00000-of-00002.parquet # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P[\d\.]+), ' r'(?P[\d\.]+|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.parquet$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToParquet*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) self.assertGreaterEqual( len(file_names), 1 * 3, #25s of data covered by 3 10s windows "expected %d files, but got: %d" % (1 * 3, len(file_names))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/range_trackers.py0000644000175100017510000003613115206113032021634 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """iobase.RangeTracker implementations provided with Apache Beam. """ # pytype: skip-file import codecs import logging import math import threading from typing import Union from apache_beam.io import iobase __all__ = [ 'OffsetRangeTracker', 'LexicographicKeyRangeTracker', 'OrderedPositionRangeTracker', 'UnsplittableRangeTracker' ] _LOGGER = logging.getLogger(__name__) class OffsetRangeTracker(iobase.RangeTracker): """A 'RangeTracker' for non-negative positions of type 'long'.""" # Offset corresponding to infinity. This can only be used as the upper-bound # of a range, and indicates reading all of the records until the end without # specifying exactly what the end is. # Infinite ranges cannot be split because it is impossible to estimate # progress within them. OFFSET_INFINITY = float('inf') def __init__(self, start, end): super().__init__() if start is None: raise ValueError('Start offset must not be \'None\'') if end is None: raise ValueError('End offset must not be \'None\'') assert isinstance(start, int) if end != self.OFFSET_INFINITY: assert isinstance(end, int) assert start <= end self._start_offset = start self._stop_offset = end self._last_record_start = -1 self._last_attempted_record_start = -1 self._offset_of_last_split_point = -1 self._lock = threading.Lock() self._split_points_seen = 0 self._split_points_unclaimed_callback = None def start_position(self): return self._start_offset def stop_position(self): return self._stop_offset @property def last_record_start(self): return self._last_record_start @property def last_attempted_record_start(self): """Return current value of last_attempted_record_start. last_attempted_record_start records a valid position that tried to be claimed by calling try_claim(). This value is only updated by `try_claim()` no matter `try_claim()` returns `True` or `False`. """ return self._last_attempted_record_start def _validate_record_start(self, record_start, split_point): # This function must only be called under the lock self.lock. if not self._lock.locked(): raise ValueError( 'This function must only be called under the lock self.lock.') if record_start < self._last_record_start: raise ValueError( 'Trying to return a record [starting at %d] which is before the ' 'last-returned record [starting at %d]' % (record_start, self._last_record_start)) if (split_point and self._offset_of_last_split_point != -1 and record_start == self._offset_of_last_split_point): raise ValueError( 'Record at a split point has same offset as the previous split ' 'point: %d' % record_start) if not split_point and self._last_record_start == -1: raise ValueError( 'The first record [starting at %d] must be at a split point' % record_start) def try_claim(self, record_start): with self._lock: # Attempted claim should be monotonous. if record_start <= self._last_attempted_record_start: raise ValueError( 'Trying to return a record [starting at %d] which is not greater' 'than the last-attempted record [starting at %d]' % (record_start, self._last_attempted_record_start)) self._validate_record_start(record_start, True) self._last_attempted_record_start = record_start if record_start >= self.stop_position(): return False self._offset_of_last_split_point = record_start self._last_record_start = record_start self._split_points_seen += 1 return True def set_current_position(self, record_start): with self._lock: self._validate_record_start(record_start, False) self._last_record_start = record_start def try_split(self, split_offset): assert isinstance(split_offset, int) with self._lock: if self._stop_offset == OffsetRangeTracker.OFFSET_INFINITY: _LOGGER.debug( 'refusing to split %r at %d: stop position unspecified', self, split_offset) return if self._last_record_start == -1: _LOGGER.debug( 'Refusing to split %r at %d: unstarted', self, split_offset) return if split_offset <= self._last_record_start: _LOGGER.debug( 'Refusing to split %r at %d: already past proposed stop offset', self, split_offset) return if (split_offset < self.start_position() or split_offset >= self.stop_position()): _LOGGER.debug( 'Refusing to split %r at %d: proposed split position out of range', self, split_offset) return _LOGGER.debug('Agreeing to split %r at %d', self, split_offset) split_fraction = ( float(split_offset - self._start_offset) / (self._stop_offset - self._start_offset)) self._stop_offset = split_offset return self._stop_offset, split_fraction def fraction_consumed(self): with self._lock: # self.last_record_start may become larger than self.end_offset when # reading the records since any record that starts before the first 'split # point' at or after the defined 'stop offset' is considered to be within # the range of the OffsetRangeTracker. Hence fraction could be > 1. # self.last_record_start is initialized to -1, hence fraction may be < 0. # Bounding the to range [0, 1]. return self.position_to_fraction( self._last_record_start, self.start_position(), self.stop_position()) def position_to_fraction(self, pos, start, stop): fraction = 1.0 * (pos - start) / (stop - start) if start != stop else 0.0 return max(0.0, min(1.0, fraction)) def position_at_fraction(self, fraction): if self.stop_position() == OffsetRangeTracker.OFFSET_INFINITY: raise Exception( 'get_position_for_fraction_consumed is not applicable for an ' 'unbounded range') return int( math.ceil( self.start_position() + fraction * (self.stop_position() - self.start_position()))) def split_points(self): with self._lock: split_points_consumed = ( 0 if self._split_points_seen == 0 else self._split_points_seen - 1) split_points_unclaimed = ( self._split_points_unclaimed_callback(self.stop_position()) if self._split_points_unclaimed_callback else iobase.RangeTracker.SPLIT_POINTS_UNKNOWN) split_points_remaining = ( iobase.RangeTracker.SPLIT_POINTS_UNKNOWN if split_points_unclaimed == iobase.RangeTracker.SPLIT_POINTS_UNKNOWN else (split_points_unclaimed + 1)) return (split_points_consumed, split_points_remaining) def set_split_points_unclaimed_callback(self, callback): self._split_points_unclaimed_callback = callback class OrderedPositionRangeTracker(iobase.RangeTracker): """ An abstract base class for range trackers whose positions are comparable. Subclasses only need to implement the mapping from position ranges to and from the closed interval [0, 1]. """ UNSTARTED = object() def __init__(self, start_position=None, stop_position=None): self._start_position = start_position self._stop_position = stop_position self._lock = threading.Lock() self._last_claim = self.UNSTARTED def start_position(self): return self._start_position def stop_position(self): with self._lock: return self._stop_position def try_claim(self, position): with self._lock: if self._last_claim is not self.UNSTARTED and position < self._last_claim: raise ValueError( "Positions must be claimed in order: " "claim '%s' attempted after claim '%s'" % (position, self._last_claim)) elif self._start_position is not None and position < self._start_position: raise ValueError( "Claim '%s' is before start '%s'" % (position, self._start_position)) if self._stop_position is None or position < self._stop_position: self._last_claim = position return True else: return False def position_at_fraction(self, fraction): return self.fraction_to_position( fraction, self._start_position, self._stop_position) def try_split(self, position): with self._lock: if ((self._stop_position is not None and position >= self._stop_position) or (self._start_position is not None and position <= self._start_position)): _LOGGER.debug( 'Refusing to split %r at %d: proposed split position out of range', self, position) return if self._last_claim is self.UNSTARTED or self._last_claim < position: fraction = self.position_to_fraction( position, start=self._start_position, end=self._stop_position) self._stop_position = position return position, fraction def fraction_consumed(self): if self._last_claim is self.UNSTARTED: return 0 else: return self.position_to_fraction( self._last_claim, self._start_position, self._stop_position) def fraction_to_position(self, fraction, start, end): """ Converts a fraction between 0 and 1 to a position between start and end. """ raise NotImplementedError def position_to_fraction(self, position, start, end): """Returns the fraction of keys in the range [start, end) that are less than the given key. """ raise NotImplementedError class UnsplittableRangeTracker(iobase.RangeTracker): """A RangeTracker that always ignores split requests. This can be used to make a given :class:`~apache_beam.io.iobase.RangeTracker` object unsplittable by ignoring all calls to :meth:`.try_split()`. All other calls will be delegated to the given :class:`~apache_beam.io.iobase.RangeTracker`. """ def __init__(self, range_tracker): """Initializes UnsplittableRangeTracker. Args: range_tracker (~apache_beam.io.iobase.RangeTracker): a :class:`~apache_beam.io.iobase.RangeTracker` to which all method calls except calls to :meth:`.try_split()` will be delegated. """ assert isinstance(range_tracker, iobase.RangeTracker) self._range_tracker = range_tracker def start_position(self): return self._range_tracker.start_position() def stop_position(self): return self._range_tracker.stop_position() def position_at_fraction(self, fraction): return self._range_tracker.position_at_fraction(fraction) def try_claim(self, position): return self._range_tracker.try_claim(position) def try_split(self, position): return None def set_current_position(self, position): self._range_tracker.set_current_position(position) def fraction_consumed(self): return self._range_tracker.fraction_consumed() def split_points(self): # An unsplittable range only contains a single split point. return (0, 1) def set_split_points_unclaimed_callback(self, callback): self._range_tracker.set_split_points_unclaimed_callback(callback) class LexicographicKeyRangeTracker(OrderedPositionRangeTracker): """A range tracker that tracks progress through a lexicographically ordered keyspace of strings. """ @classmethod def fraction_to_position( cls, fraction: float, start: Union[bytes, str] = None, end: Union[bytes, str] = None, ) -> Union[bytes, str]: """Linearly interpolates a key that is lexicographically fraction of the way between start and end. """ assert 0 <= fraction <= 1, fraction if start is None: start = b'' if fraction == 0: return start if fraction == 1: return end if not end: common_prefix_len = len(start) - len(start.lstrip(b'\xFF')) else: for ix, (s, e) in enumerate(zip(start, end)): if s != e: common_prefix_len = ix break else: common_prefix_len = min(len(start), len(end)) # Convert the relative precision of fraction (~53 bits) to an absolute # precision needed to represent values between start and end distinctly. prec = common_prefix_len + int(-math.log(fraction, 256)) + 7 istart = cls._bytestring_to_int(start, prec) iend = cls._bytestring_to_int(end, prec) if end else 1 << (prec * 8) ikey = istart + int((iend - istart) * fraction) # Could be equal due to rounding. # Adjust to ensure we never return the actual start and end # unless fraction is exatly 0 or 1. if ikey == istart: ikey += 1 elif ikey == iend: ikey -= 1 position: bytes = cls._bytestring_from_int(ikey, prec).rstrip(b'\0') if isinstance(start, bytes): return position return position.decode(encoding='unicode_escape', errors='replace') @classmethod def position_to_fraction( cls, key: Union[bytes, str] = None, start: Union[bytes, str] = None, end: Union[bytes, str] = None, ) -> float: """Returns the fraction of keys in the range [start, end) that are less than the given key. """ if not key: return 0 if start is None: start = '' if isinstance(key, str) else b'' prec = len(start) + 7 if key.startswith(start): # Higher absolute precision needed for very small values of fixed # relative position. trailing_symbol = '\0' if isinstance(key, str) else b'\0' prec = max( prec, len(key) - len(key[len(start):].strip(trailing_symbol)) + 7) istart = cls._bytestring_to_int(start, prec) ikey = cls._bytestring_to_int(key, prec) iend = cls._bytestring_to_int(end, prec) if end else 1 << (prec * 8) return float(ikey - istart) / (iend - istart) @staticmethod def _bytestring_to_int(s: Union[bytes, str], prec: int) -> int: """Returns int(256**prec * f) where f is the fraction represented by interpreting '.' + s as a base-256 floating point number. """ if not s: return 0 if isinstance(s, str): s = s.encode() # str -> bytes if len(s) < prec: s += b'\0' * (prec - len(s)) else: s = s[:prec] h = codecs.encode(s, encoding='hex') return int(h, base=16) @staticmethod def _bytestring_from_int(i: int, prec: int) -> bytes: """Inverse of _bytestring_to_int.""" h = '%x' % i return codecs.decode('0' * (2 * prec - len(h)) + h, encoding='hex') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/range_trackers_test.py0000644000175100017510000004635315206113032022702 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the range_trackers module.""" # pytype: skip-file import copy import logging import math import unittest from typing import Optional from typing import Union from apache_beam.io import range_trackers class OffsetRangeTrackerTest(unittest.TestCase): def test_try_return_record_simple_sparse(self): tracker = range_trackers.OffsetRangeTracker(100, 200) self.assertTrue(tracker.try_claim(110)) self.assertTrue(tracker.try_claim(140)) self.assertTrue(tracker.try_claim(183)) self.assertFalse(tracker.try_claim(210)) def test_try_return_record_simple_dense(self): tracker = range_trackers.OffsetRangeTracker(3, 6) self.assertTrue(tracker.try_claim(3)) self.assertTrue(tracker.try_claim(4)) self.assertTrue(tracker.try_claim(5)) self.assertFalse(tracker.try_claim(6)) def test_try_claim_update_last_attempt(self): tracker = range_trackers.OffsetRangeTracker(1, 2) self.assertTrue(tracker.try_claim(1)) self.assertEqual(1, tracker.last_attempted_record_start) self.assertFalse(tracker.try_claim(3)) self.assertEqual(3, tracker.last_attempted_record_start) self.assertFalse(tracker.try_claim(6)) self.assertEqual(6, tracker.last_attempted_record_start) with self.assertRaises(Exception): tracker.try_claim(6) def test_set_current_position(self): tracker = range_trackers.OffsetRangeTracker(0, 6) self.assertTrue(tracker.try_claim(2)) # Cannot set current position before successful claimed pos. with self.assertRaises(Exception): tracker.set_current_position(1) self.assertFalse(tracker.try_claim(10)) tracker.set_current_position(11) self.assertEqual(10, tracker.last_attempted_record_start) self.assertEqual(11, tracker.last_record_start) def test_try_return_record_continuous_until_split_point(self): tracker = range_trackers.OffsetRangeTracker(9, 18) # Return records with gaps of 2; every 3rd record is a split point. self.assertTrue(tracker.try_claim(10)) tracker.set_current_position(12) tracker.set_current_position(14) self.assertTrue(tracker.try_claim(16)) # Out of range, but not a split point... tracker.set_current_position(18) tracker.set_current_position(20) # Out of range AND a split point. self.assertFalse(tracker.try_claim(22)) def test_split_at_offset_fails_if_unstarted(self): tracker = range_trackers.OffsetRangeTracker(100, 200) self.assertFalse(tracker.try_split(150)) def test_split_at_offset(self): tracker = range_trackers.OffsetRangeTracker(100, 200) self.assertTrue(tracker.try_claim(110)) # Example positions we shouldn't split at, when last record starts at 110: self.assertFalse(tracker.try_split(109)) self.assertFalse(tracker.try_split(110)) self.assertFalse(tracker.try_split(200)) self.assertFalse(tracker.try_split(210)) # Example positions we *should* split at: self.assertTrue(copy.copy(tracker).try_split(111)) self.assertTrue(copy.copy(tracker).try_split(129)) self.assertTrue(copy.copy(tracker).try_split(130)) self.assertTrue(copy.copy(tracker).try_split(131)) self.assertTrue(copy.copy(tracker).try_split(150)) self.assertTrue(copy.copy(tracker).try_split(199)) # If we split at 170 and then at 150: self.assertTrue(tracker.try_split(170)) self.assertTrue(tracker.try_split(150)) # Should be able to return a record starting before the new stop offset. # Returning records starting at the same offset is ok. self.assertTrue(copy.copy(tracker).try_claim(135)) self.assertTrue(copy.copy(tracker).try_claim(135)) # Should be able to return a record starting right before the new stop # offset. self.assertTrue(copy.copy(tracker).try_claim(149)) # Should not be able to return a record starting at or after the new stop # offset. self.assertFalse(tracker.try_claim(150)) self.assertFalse(tracker.try_claim(151)) # Should accept non-splitpoint records starting after stop offset. tracker.set_current_position(152) tracker.set_current_position(160) tracker.set_current_position(171) def test_get_position_for_fraction_dense(self): # Represents positions 3, 4, 5. tracker = range_trackers.OffsetRangeTracker(3, 6) # Position must be an integer type. self.assertTrue(isinstance(tracker.position_at_fraction(0.0), int)) # [3, 3) represents 0.0 of [3, 6) self.assertEqual(3, tracker.position_at_fraction(0.0)) # [3, 4) represents up to 1/3 of [3, 6) self.assertEqual(4, tracker.position_at_fraction(1.0 / 6)) self.assertEqual(4, tracker.position_at_fraction(0.333)) # [3, 5) represents up to 2/3 of [3, 6) self.assertEqual(5, tracker.position_at_fraction(0.334)) self.assertEqual(5, tracker.position_at_fraction(0.666)) # Any fraction consumed over 2/3 means the whole [3, 6) has been consumed. self.assertEqual(6, tracker.position_at_fraction(0.667)) def test_get_fraction_consumed_dense(self): tracker = range_trackers.OffsetRangeTracker(3, 6) self.assertEqual(0, tracker.fraction_consumed()) self.assertTrue(tracker.try_claim(3)) self.assertEqual(0.0, tracker.fraction_consumed()) self.assertTrue(tracker.try_claim(4)) self.assertEqual(1.0 / 3, tracker.fraction_consumed()) self.assertTrue(tracker.try_claim(5)) self.assertEqual(2.0 / 3, tracker.fraction_consumed()) tracker.set_current_position(6) self.assertEqual(1.0, tracker.fraction_consumed()) tracker.set_current_position(7) self.assertFalse(tracker.try_claim(7)) def test_get_fraction_consumed_sparse(self): tracker = range_trackers.OffsetRangeTracker(100, 200) self.assertEqual(0, tracker.fraction_consumed()) self.assertTrue(tracker.try_claim(110)) # Consumed positions through 110 = total 10 positions of 100 done. self.assertEqual(0.10, tracker.fraction_consumed()) self.assertTrue(tracker.try_claim(150)) self.assertEqual(0.50, tracker.fraction_consumed()) self.assertTrue(tracker.try_claim(195)) self.assertEqual(0.95, tracker.fraction_consumed()) def test_everything_with_unbounded_range(self): tracker = range_trackers.OffsetRangeTracker( 100, range_trackers.OffsetRangeTracker.OFFSET_INFINITY) self.assertTrue(tracker.try_claim(150)) self.assertTrue(tracker.try_claim(250)) # get_position_for_fraction_consumed should fail for an unbounded range with self.assertRaises(Exception): tracker.position_at_fraction(0.5) def test_try_return_first_record_not_split_point(self): with self.assertRaises(Exception): range_trackers.OffsetRangeTracker(100, 200).set_current_position(120) def test_try_return_record_non_monotonic(self): tracker = range_trackers.OffsetRangeTracker(100, 200) self.assertTrue(tracker.try_claim(120)) with self.assertRaises(Exception): tracker.try_claim(110) def test_try_split_points(self): tracker = range_trackers.OffsetRangeTracker(100, 400) def dummy_callback(stop_position): return int(stop_position // 5) tracker.set_split_points_unclaimed_callback(dummy_callback) self.assertEqual(tracker.split_points(), (0, 81)) self.assertTrue(tracker.try_claim(120)) self.assertEqual(tracker.split_points(), (0, 81)) self.assertTrue(tracker.try_claim(140)) self.assertEqual(tracker.split_points(), (1, 81)) tracker.try_split(200) self.assertEqual(tracker.split_points(), (1, 41)) self.assertTrue(tracker.try_claim(150)) self.assertEqual(tracker.split_points(), (2, 41)) self.assertTrue(tracker.try_claim(180)) self.assertEqual(tracker.split_points(), (3, 41)) self.assertFalse(tracker.try_claim(210)) self.assertEqual(tracker.split_points(), (3, 41)) class OrderedPositionRangeTrackerTest(unittest.TestCase): class DoubleRangeTracker(range_trackers.OrderedPositionRangeTracker): @staticmethod def fraction_to_position(fraction, start, end): return start + (end - start) * fraction @staticmethod def position_to_fraction(pos, start, end): return float(pos - start) / (end - start) def test_try_claim(self): tracker = self.DoubleRangeTracker(10, 20) self.assertTrue(tracker.try_claim(10)) self.assertTrue(tracker.try_claim(15)) self.assertFalse(tracker.try_claim(20)) self.assertFalse(tracker.try_claim(25)) def test_fraction_consumed(self): tracker = self.DoubleRangeTracker(10, 20) self.assertEqual(0, tracker.fraction_consumed()) tracker.try_claim(10) self.assertEqual(0, tracker.fraction_consumed()) tracker.try_claim(15) self.assertEqual(.5, tracker.fraction_consumed()) tracker.try_claim(17) self.assertEqual(.7, tracker.fraction_consumed()) tracker.try_claim(25) self.assertEqual(.7, tracker.fraction_consumed()) def test_try_split(self): tracker = self.DoubleRangeTracker(10, 20) tracker.try_claim(15) self.assertEqual(.5, tracker.fraction_consumed()) # Split at 18. self.assertEqual((18, 0.8), tracker.try_split(18)) # Fraction consumed reflects smaller range. self.assertEqual(.625, tracker.fraction_consumed()) # We can claim anything less than 18, self.assertTrue(tracker.try_claim(17)) # but can't split before claimed 17, self.assertIsNone(tracker.try_split(16)) # nor claim anything at or after 18. self.assertFalse(tracker.try_claim(18)) self.assertFalse(tracker.try_claim(19)) def test_claim_order(self): tracker = self.DoubleRangeTracker(10, 20) tracker.try_claim(12) tracker.try_claim(15) with self.assertRaises(ValueError): tracker.try_claim(13) def test_out_of_range(self): tracker = self.DoubleRangeTracker(10, 20) # Can't claim before range. with self.assertRaises(ValueError): tracker.try_claim(-5) # Can't split before range. self.assertFalse(tracker.try_split(-5)) # Reject useless split at start position. self.assertFalse(tracker.try_split(10)) # Can't split after range. self.assertFalse(tracker.try_split(25)) tracker.try_split(15) # Can't split after modified range. self.assertFalse(tracker.try_split(17)) # Reject useless split at end position. self.assertFalse(tracker.try_split(15)) self.assertTrue(tracker.try_split(14)) class UnsplittableRangeTrackerTest(unittest.TestCase): def test_try_claim(self): tracker = range_trackers.UnsplittableRangeTracker( range_trackers.OffsetRangeTracker(100, 200)) self.assertTrue(tracker.try_claim(110)) self.assertTrue(tracker.try_claim(140)) self.assertTrue(tracker.try_claim(183)) self.assertFalse(tracker.try_claim(210)) def test_try_split_fails(self): tracker = range_trackers.UnsplittableRangeTracker( range_trackers.OffsetRangeTracker(100, 200)) self.assertTrue(tracker.try_claim(110)) # Out of range self.assertFalse(tracker.try_split(109)) self.assertFalse(tracker.try_split(210)) # Within range. But splitting is still unsuccessful. self.assertFalse(copy.copy(tracker).try_split(111)) self.assertFalse(copy.copy(tracker).try_split(130)) self.assertFalse(copy.copy(tracker).try_split(199)) class LexicographicKeyRangeTrackerTest(unittest.TestCase): """Tests of LexicographicKeyRangeTracker.""" key_to_fraction = ( range_trackers.LexicographicKeyRangeTracker.position_to_fraction) fraction_to_key = ( range_trackers.LexicographicKeyRangeTracker.fraction_to_position) def _check( self, fraction: Optional[float] = None, key: Union[bytes, str] = None, start: Union[bytes, str] = None, end: Union[bytes, str] = None, delta: float = 0.0): assert key is not None or fraction is not None if fraction is None: fraction = self.key_to_fraction(key, start, end) elif key is None: key = self.fraction_to_key(fraction, start, end) if key is None and end is None and fraction == 1: # No way to distinguish from fraction == 0. computed_fraction = 1 else: computed_fraction = self.key_to_fraction(key, start, end) computed_key = self.fraction_to_key(fraction, start, end) if delta: self.assertAlmostEqual( computed_fraction, fraction, delta=delta, places=None, msg=str(locals())) else: self.assertEqual(computed_fraction, fraction, str(locals())) self.assertEqual(computed_key, key, str(locals())) def test_key_to_fraction_no_endpoints(self): self._check(key=b'\x07', fraction=7 / 256.) self._check(key=b'\xFF', fraction=255 / 256.) self._check(key=b'\x01\x02\x03', fraction=(2**16 + 2**9 + 3) / (2.0**24)) self._check(key=b'UUUUUUT', fraction=1 / 3) self._check(key=b'3333334', fraction=1 / 5) self._check(key=b'$\x92I$\x92I$', fraction=1 / 7, delta=1e-3) self._check(key=b'\x01\x02\x03', fraction=(2**16 + 2**9 + 3) / (2.0**24)) def test_key_to_fraction(self): # test no key, no start self._check(end=b'eeeeee', fraction=0.0) self._check(end='eeeeee', fraction=0.0) # test no fraction self._check(key=b'bbbbbb', start=b'aaaaaa', end=b'eeeeee') self._check(key='bbbbbb', start='aaaaaa', end='eeeeee') # test no start self._check(key=b'eeeeee', end=b'eeeeee', fraction=1.0) self._check(key='eeeeee', end='eeeeee', fraction=1.0) self._check(key=b'\x19YYYYY@', end=b'eeeeee', fraction=0.25) self._check(key=b'2\xb2\xb2\xb2\xb2\xb2\x80', end='eeeeee', fraction=0.5) self._check(key=b'L\x0c\x0c\x0c\x0c\x0b\xc0', end=b'eeeeee', fraction=0.75) # test bytes keys self._check(key=b'\x87', start=b'\x80', fraction=7 / 128.) self._check(key=b'\x07', end=b'\x10', fraction=7 / 16.) self._check(key=b'\x47', start=b'\x40', end=b'\x80', fraction=7 / 64.) self._check(key=b'\x47\x80', start=b'\x40', end=b'\x80', fraction=15 / 128.) # test string keys self._check(key='aaaaaa', start='aaaaaa', end='eeeeee', fraction=0.0) self._check(key='bbbbbb', start='aaaaaa', end='eeeeee', fraction=0.25) self._check(key='cccccc', start='aaaaaa', end='eeeeee', fraction=0.5) self._check(key='dddddd', start='aaaaaa', end='eeeeee', fraction=0.75) self._check(key='eeeeee', start='aaaaaa', end='eeeeee', fraction=1.0) def test_key_to_fraction_common_prefix(self): # test bytes keys self._check( key=b'a' * 100 + b'b', start=b'a' * 100 + b'a', end=b'a' * 100 + b'c', fraction=0.5) self._check( key=b'a' * 100 + b'b', start=b'a' * 100 + b'a', end=b'a' * 100 + b'e', fraction=0.25) self._check( key=b'\xFF' * 100 + b'\x40', start=b'\xFF' * 100, end=None, fraction=0.25) self._check( key=b'foob', start=b'fooa\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFE', end=b'foob\x00\x00\x00\x00\x00\x00\x00\x00\x02', fraction=0.5) # test string keys self._check( key='a' * 100 + 'a', start='a' * 100 + 'a', end='a' * 100 + 'e', fraction=0.0) self._check( key='a' * 100 + 'b', start='a' * 100 + 'a', end='a' * 100 + 'e', fraction=0.25) self._check( key='a' * 100 + 'c', start='a' * 100 + 'a', end='a' * 100 + 'e', fraction=0.5) self._check( key='a' * 100 + 'd', start='a' * 100 + 'a', end='a' * 100 + 'e', fraction=0.75) self._check( key='a' * 100 + 'e', start='a' * 100 + 'a', end='a' * 100 + 'e', fraction=1.0) def test_tiny(self): # test bytes keys self._check(fraction=.5**20, key=b'\0\0\x10') self._check(fraction=.5**20, start=b'a', end=b'b', key=b'a\0\0\x10') self._check(fraction=.5**20, start=b'a', end=b'c', key=b'a\0\0\x20') self._check( fraction=.5**20, start=b'xy_a', end=b'xy_c', key=b'xy_a\0\0\x20') self._check( fraction=.5**20, start=b'\xFF\xFF\x80', key=b'\xFF\xFF\x80\x00\x08') self._check( fraction=.5**20 / 3, start=b'xy_a', end=b'xy_c', key=b'xy_a\x00\x00\n\xaa\xaa\xaa\xaa\xaa', delta=1e-15) self._check(fraction=.5**100, key=b'\0' * 12 + b'\x10') # test string keys self._check(fraction=.5**20, start='a', end='b', key='a\0\0\x10') self._check(fraction=.5**20, start='a', end='c', key='a\0\0\x20') self._check(fraction=.5**20, start='xy_a', end='xy_c', key='xy_a\0\0\x20') def test_lots(self): for fraction in (0, 1, .5, .75, 7. / 512, 1 - 7. / 4096): self._check(fraction) self._check(fraction, start=b'\x01') self._check(fraction, end=b'\xF0') self._check(fraction, start=b'0x75', end=b'\x76') self._check(fraction, start=b'0x75', end=b'\x77') self._check(fraction, start=b'0x75', end=b'\x78') self._check( fraction, start=b'a' * 100 + b'\x80', end=b'a' * 100 + b'\x81') self._check( fraction, start=b'a' * 101 + b'\x80', end=b'a' * 101 + b'\x81') self._check( fraction, start=b'a' * 102 + b'\x80', end=b'a' * 102 + b'\x81') for fraction in (.3, 1 / 3., 1 / math.e, .001, 1e-30, .99, .999999): self._check(fraction, delta=1e-14) self._check(fraction, start=b'\x01', delta=1e-14) self._check(fraction, end=b'\xF0', delta=1e-14) self._check(fraction, start=b'0x75', end=b'\x76', delta=1e-14) self._check(fraction, start=b'0x75', end=b'\x77', delta=1e-14) self._check(fraction, start=b'0x75', end=b'\x78', delta=1e-14) self._check( fraction, start=b'a' * 100 + b'\x80', end=b'a' * 100 + b'\x81', delta=1e-14) def test_good_prec(self): # There should be about 7 characters (~53 bits) of precision # (beyond the common prefix of start and end). self._check( 1 / math.e, start='AAAAAAA', end='zzzzzzz', key='VNg/ot\x82', delta=1e-14) self._check( 1 / math.e, start=b'abc_abc', end=b'abc_xyz', key=b'abc_i\xe0\xf4\x84\x86\x99\x96', delta=1e-15) # This remains true even if the start and end keys are given to # high precision. self._check( 1 / math.e, start=b'abcd_abc\0\0\0\0\0_______________abc', end=b'abcd_xyz\0\0\0\0\0\0_______________abc', key=b'abcd_i\xe0\xf4\x84\x86\x99\x96', delta=1e-15) # For very small fractions, however, higher precision is used to # accurately represent small increments in the keyspace. self._check( 1e-20 / math.e, start=b'abcd_abc', end=b'abcd_xyz', key=b'abcd_abc\x00\x00\x00\x00\x00\x01\x91#\x172N\xbb', delta=1e-35) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/requestresponse.py0000644000175100017510000007077215206113032022122 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """``PTransform`` for reading from and writing to Web APIs.""" import abc import concurrent.futures import contextlib import enum import json import logging import sys import time from datetime import timedelta from typing import Any from typing import Generic from typing import Mapping from typing import Optional from typing import TypeVar from typing import Union from google.api_core.exceptions import TooManyRequests import apache_beam as beam from apache_beam import pvalue from apache_beam.coders import coders from apache_beam.io.components.adaptive_throttler import AdaptiveThrottler from apache_beam.metrics import Metrics from apache_beam.transforms.util import BatchElements from apache_beam.utils import retry try: import redis except ImportError: redis = None RequestT = TypeVar('RequestT') ResponseT = TypeVar('ResponseT') # DEFAULT_TIMEOUT_SECS represents the time interval for completing the request # with external source. DEFAULT_TIMEOUT_SECS = 30 # DEFAULT_CACHE_ENTRY_TTL_SEC represents the total time-to-live # for cache record. DEFAULT_CACHE_ENTRY_TTL_SEC = 24 * 60 * 60 MSEC_TO_SEC = 1000 _LOGGER = logging.getLogger(__name__) __all__ = [ 'RequestResponseIO', 'ExponentialBackOffRepeater', 'DefaultThrottler', 'NoOpsRepeater', 'RedisCache', ] class UserCodeExecutionException(Exception): """Base class for errors related to calling Web APIs.""" class UserCodeQuotaException(UserCodeExecutionException): """Extends ``UserCodeExecutionException`` to signal specifically that the Web API client encountered a Quota or API overuse related error. """ class UserCodeTimeoutException(UserCodeExecutionException): """Extends ``UserCodeExecutionException`` to signal a user code timeout.""" def retry_on_exception(exception: Exception): """retry on exceptions caused by unavailability of the remote server.""" return isinstance( exception, (TooManyRequests, UserCodeTimeoutException, UserCodeQuotaException)) class _MetricsCollector: """A metrics collector that tracks RequestResponseIO related usage.""" def __init__(self, namespace: str): """ Args: namespace: Namespace for the metrics. """ self.requests = Metrics.counter(namespace, 'requests') self.responses = Metrics.counter(namespace, 'responses') self.failures = Metrics.counter(namespace, 'failures') self.throttled_requests = Metrics.counter(namespace, 'throttled_requests') self.throttled_secs = Metrics.counter( namespace, 'cumulativeThrottlingSeconds') self.timeout_requests = Metrics.counter(namespace, 'requests_timed_out') self.call_counter = Metrics.counter(namespace, 'call_invocations') self.setup_counter = Metrics.counter(namespace, 'setup_counter') self.teardown_counter = Metrics.counter(namespace, 'teardown_counter') self.backoff_counter = Metrics.counter(namespace, 'backoff_counter') self.sleeper_counter = Metrics.counter(namespace, 'sleeper_counter') self.should_backoff_counter = Metrics.counter( namespace, 'should_backoff_counter') class Caller(contextlib.AbstractContextManager, abc.ABC, Generic[RequestT, ResponseT]): """Interface for user custom code intended for API calls. For setup and teardown of clients when applicable, implement the ``__enter__`` and ``__exit__`` methods respectively.""" @abc.abstractmethod def __call__(self, request: RequestT, *args, **kwargs) -> ResponseT: """Calls a Web API with the ``RequestT`` and returns a ``ResponseT``. ``RequestResponseIO`` expects implementations of the ``__call__`` method to throw either a ``UserCodeExecutionException``, ``UserCodeQuotaException``, or ``UserCodeTimeoutException``. """ pass def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): return None def get_cache_key(self, request: RequestT) -> str: """Returns the request to be cached. This is how the response will be looked up in the cache as well. By default, entire request is cached as the key for the cache. Implement this method to override the key for the cache. For example, in `BigTableEnrichmentHandler`, the row key for the element is returned here. """ return "" def batch_elements_kwargs(self) -> Mapping[str, Any]: """Returns a kwargs suitable for `beam.BatchElements`.""" return {} class ShouldBackOff(abc.ABC): """ Provides mechanism to apply adaptive throttling. """ pass class Repeater(abc.ABC): """Provides mechanism to repeat requests for a configurable condition.""" @abc.abstractmethod def repeat( self, caller: Caller[RequestT, ResponseT], request: RequestT, timeout: float, metrics_collector: Optional[_MetricsCollector]) -> ResponseT: """Implements a repeater strategy for RequestResponseIO when a repeater is enabled. Args: caller: a `~apache_beam.io.requestresponse.Caller` object that calls the API. request: input request to repeat. timeout: time to wait for the request to complete. metrics_collector: (Optional) a `~apache_beam.io.requestresponse._MetricsCollector` object to collect the metrics for RequestResponseIO. """ pass def _execute_request( caller: Caller[RequestT, ResponseT], request: RequestT, timeout: float, metrics_collector: Optional[_MetricsCollector] = None) -> ResponseT: with concurrent.futures.ThreadPoolExecutor() as executor: future = executor.submit(caller, request) try: return future.result(timeout=timeout) except TooManyRequests as e: _LOGGER.info( 'request could not be completed. got code %i from the service.', e.code) raise e except concurrent.futures.TimeoutError: if metrics_collector: metrics_collector.timeout_requests.inc(1) raise UserCodeTimeoutException( f'Timeout {timeout} exceeded ' f'while completing request: {request}') except RuntimeError: if metrics_collector: metrics_collector.failures.inc(1) raise UserCodeExecutionException('could not complete request') class ExponentialBackOffRepeater(Repeater): """Configure exponential backoff retry strategy. It retries for exceptions due to the remote service such as TooManyRequests (HTTP 429), UserCodeTimeoutException, UserCodeQuotaException. It utilizes the decorator :func:`apache_beam.utils.retry.with_exponential_backoff`. """ def __init__(self): pass @retry.with_exponential_backoff( num_retries=2, retry_filter=retry_on_exception) def repeat( self, caller: Caller[RequestT, ResponseT], request: RequestT, timeout: float, metrics_collector: Optional[_MetricsCollector] = None) -> ResponseT: """repeat method is called from the RequestResponseIO when a repeater is enabled. Args: caller: a `~apache_beam.io.requestresponse.Caller` object that calls the API. request: input request to repeat. timeout: time to wait for the request to complete. metrics_collector: (Optional) a `~apache_beam.io.requestresponse._MetricsCollector` object to collect the metrics for RequestResponseIO. """ return _execute_request(caller, request, timeout, metrics_collector) class NoOpsRepeater(Repeater): """Executes a request just once irrespective of any exception. """ def repeat( self, caller: Caller[RequestT, ResponseT], request: RequestT, timeout: float, metrics_collector: Optional[_MetricsCollector]) -> ResponseT: return _execute_request(caller, request, timeout, metrics_collector) class PreCallThrottler(abc.ABC): """Provides a throttle mechanism before sending request.""" pass class DefaultThrottler(PreCallThrottler): """Default throttler that uses :class:`apache_beam.io.components.adaptive_throttler.AdaptiveThrottler` Args: window_ms (int): length of history to consider, in ms, to set throttling. bucket_ms (int): granularity of time buckets that we store data in, in ms. overload_ratio (float): the target ratio between requests sent and successful requests. This is "K" in the formula in https://landing.google.com/sre/book/chapters/handling-overload.html. delay_secs (int): minimum number of seconds to throttle a request. """ def __init__( self, window_ms: int = 1, bucket_ms: int = 1, overload_ratio: float = 2, delay_secs: int = 5): self.throttler = AdaptiveThrottler( window_ms=window_ms, bucket_ms=bucket_ms, overload_ratio=overload_ratio) self.delay_secs = delay_secs class _FilterCacheReadFn(beam.DoFn): """A `DoFn` that partitions cache reads. It emits to main output for successful cache read requests or to the tagged output - `cache_misses` - otherwise.""" def process(self, element: tuple[RequestT, ResponseT], *args, **kwargs): if not element[1]: yield pvalue.TaggedOutput('cache_misses', element[0]) else: yield element class _Call(beam.PTransform[beam.PCollection[RequestT], beam.PCollection[ResponseT]]): """(Internal-only) PTransform that invokes a remote function on each element of the input PCollection. This PTransform uses a `Caller` object to invoke the actual API calls, and uses ``__enter__`` and ``__exit__`` to manage setup and teardown of clients when applicable. Additionally, a timeout value is specified to regulate the duration of each call, defaults to 30 seconds. Args: caller: a `Caller` object that invokes API call. timeout (float): timeout value in seconds to wait for response from API. should_backoff: (Optional) provides methods for backoff. repeater: (Optional) provides methods to repeat requests to API. throttler: (Optional) provides methods to pre-throttle a request. """ def __init__( self, caller: Caller[RequestT, ResponseT], timeout: Optional[float] = DEFAULT_TIMEOUT_SECS, should_backoff: Optional[ShouldBackOff] = None, repeater: Repeater = None, throttler: PreCallThrottler = None, ): self._caller = caller self._timeout = timeout self._should_backoff = should_backoff self._repeater = repeater self._throttler = throttler def expand( self, requests: beam.PCollection[RequestT]) -> beam.PCollection[ResponseT]: return requests | beam.ParDo( _CallDoFn(self._caller, self._timeout, self._repeater, self._throttler)) class _CallDoFn(beam.DoFn): def setup(self): self._caller.__enter__() self._metrics_collector = _MetricsCollector(self._caller.__str__()) self._metrics_collector.setup_counter.inc(1) def __init__( self, caller: Caller[RequestT, ResponseT], timeout: float, repeater: Repeater, throttler: PreCallThrottler): self._metrics_collector = None self._caller = caller self._timeout = timeout self._repeater = repeater self._throttler = throttler def process(self, request: RequestT, *args, **kwargs): self._metrics_collector.requests.inc(1) is_throttled_request = False if self._throttler: while self._throttler.throttler.throttle_request(time.time() * MSEC_TO_SEC): _LOGGER.info( "Delaying request for %d seconds" % self._throttler.delay_secs) time.sleep(self._throttler.delay_secs) self._metrics_collector.throttled_secs.inc(self._throttler.delay_secs) is_throttled_request = True if is_throttled_request: self._metrics_collector.throttled_requests.inc(1) try: req_time = time.time() response = self._repeater.repeat( self._caller, request, self._timeout, self._metrics_collector) self._metrics_collector.responses.inc(1) self._throttler.throttler.successful_request(req_time * MSEC_TO_SEC) yield response except Exception as e: raise e def teardown(self): self._metrics_collector.teardown_counter.inc(1) self._caller.__exit__(*sys.exc_info()) class Cache(abc.ABC): """Base Cache class for :class:`apache_beam.io.requestresponse.RequestResponseIO`. For adding cache support to RequestResponseIO, implement this class. """ @abc.abstractmethod def get_read(self): """returns a PTransform that reads from the cache.""" pass @abc.abstractmethod def get_write(self): """returns a PTransform that writes to the cache.""" pass @property @abc.abstractmethod def request_coder(self): """request coder to use with Cache.""" pass @request_coder.setter @abc.abstractmethod def request_coder(self, request_coder: coders.Coder): """sets the request coder to use with Cache.""" pass @property @abc.abstractmethod def source_caller(self): """Actual caller that is using the cache.""" pass @source_caller.setter @abc.abstractmethod def source_caller(self, caller: Caller): """Sets the source caller for :class:`apache_beam.io.requestresponse.RequestResponseIO` to pull cache request key from respective callers.""" pass class _RedisMode(enum.Enum): """ Mode of operation for redis cache when using `~apache_beam.io.requestresponse._RedisCaller`. """ READ = 0 WRITE = 1 class _RedisCaller(Caller): """An implementation of `~apache_beam.io.requestresponse.Caller` for Redis client. It provides the functionality for making requests to Redis server using :class:`apache_beam.io.requestresponse.RequestResponseIO`. """ def __init__( self, host: str, port: int, time_to_live: Union[int, timedelta], *, request_coder: Optional[coders.Coder], response_coder: Optional[coders.Coder], kwargs: Optional[dict[str, Any]] = None, source_caller: Optional[Caller] = None, mode: _RedisMode, ): """ Args: host (str): The hostname or IP address of the Redis server. port (int): The port number of the Redis server. time_to_live: `(Union[int, timedelta])` The time-to-live (TTL) for records stored in Redis. Provide an integer (in seconds) or a `datetime.timedelta` object. request_coder: (Optional[`coders.Coder`]) coder for requests stored in Redis. response_coder: (Optional[`coders.Coder`]) coder for decoding responses received from Redis. kwargs: Optional(Dict[str, Any]) additional keyword arguments that are required to connect to your redis server. Same as `redis.Redis()`. source_caller: (Optional[`Caller`]): The source caller using this Redis cache in case of fetching the cache request to store in Redis. mode: `_RedisMode` An enum type specifying the operational mode of the `_RedisCaller`. """ self.host, self.port = host, port self.time_to_live = time_to_live self.request_coder = request_coder self.response_coder = response_coder self.kwargs = kwargs self.source_caller = source_caller self.mode = mode def __enter__(self): self.client = redis.Redis(self.host, self.port, **self.kwargs) def _read_cache(self, element): cache_request = self.source_caller.get_cache_key(element) # check if the caller is a enrichment handler. EnrichmentHandler # provides the request format for cache. if cache_request: encoded_request = self.request_coder.encode(cache_request) else: encoded_request = self.request_coder.encode(element) encoded_response = self.client.get(encoded_request) if not encoded_response: # no cache entry present for this request. return element, None if self.response_coder is None: try: response_dict = json.loads(encoded_response.decode('utf-8')) response = beam.Row(**response_dict) except Exception: _LOGGER.warning( 'cannot decode response from redis cache for %s.' % element) return element, None else: response = self.response_coder.decode(encoded_response) return element, response def _write_cache(self, element): cache_request = self.source_caller.get_cache_key(element[0]) if cache_request: encoded_request = self.request_coder.encode(cache_request) else: encoded_request = self.request_coder.encode(element[0]) if self.response_coder is None: try: encoded_response = json.dumps(element[1]._asdict()).encode('utf-8') except Exception: _LOGGER.warning( 'cannot encode response %s for %s to store in ' 'redis cache.' % (element[1], element[0])) return element else: encoded_response = self.response_coder.encode(element[1]) # Write to cache with TTL. Set nx to True to prevent overwriting for the # same key. self.client.set( encoded_request, encoded_response, self.time_to_live, nx=True) return element def __call__(self, element, *args, **kwargs): if self.mode == _RedisMode.READ: if isinstance(element, list): responses = [self._read_cache(e) for e in element] return responses else: return self._read_cache(element) else: if isinstance(element, list): responses = [self._write_cache(e) for e in element] return responses else: return self._write_cache(element) def __exit__(self, exc_type, exc_val, exc_tb): self.client.close() class _ReadFromRedis(beam.PTransform[beam.PCollection[RequestT], beam.PCollection[ResponseT]]): """A `PTransform` that performs Redis cache read.""" def __init__( self, host: str, port: int, time_to_live: Union[int, timedelta], *, kwargs: Optional[dict[str, Any]] = None, request_coder: Optional[coders.Coder], response_coder: Optional[coders.Coder], source_caller: Optional[Caller[RequestT, ResponseT]] = None, ): """ Args: host (str): The hostname or IP address of the Redis server. port (int): The port number of the Redis server. time_to_live: `(Union[int, timedelta])` The time-to-live (TTL) for records stored in Redis. Provide an integer (in seconds) or a `datetime.timedelta` object. kwargs: Optional(Dict[str, Any]) additional keyword arguments that are required to connect to your redis server. Same as `redis.Redis()`. request_coder: (Optional[`coders.Coder`]) coder for requests stored in Redis. response_coder: (Optional[`coders.Coder`]) coder for decoding responses received from Redis. source_caller: (Optional[`Caller`]): The source caller using this Redis cache in case of fetching the cache request to store in Redis. """ self.request_coder = request_coder self.response_coder = response_coder self.redis_caller = _RedisCaller( host, port, time_to_live, request_coder=self.request_coder, response_coder=self.response_coder, kwargs=kwargs, source_caller=source_caller, mode=_RedisMode.READ) def expand( self, requests: beam.PCollection[RequestT]) -> beam.PCollection[ResponseT]: return requests | RequestResponseIO(self.redis_caller) class _WriteToRedis(beam.PTransform[beam.PCollection[tuple[RequestT, ResponseT]], beam.PCollection[ResponseT]]): """A `PTransfrom` that performs write to Redis cache.""" def __init__( self, host: str, port: int, time_to_live: Union[int, timedelta], *, kwargs: Optional[dict[str, Any]] = None, request_coder: Optional[coders.Coder], response_coder: Optional[coders.Coder], source_caller: Optional[Caller[RequestT, ResponseT]] = None, ): """ Args: host (str): The hostname or IP address of the Redis server. port (int): The port number of the Redis server. time_to_live: `(Union[int, timedelta])` The time-to-live (TTL) for records stored in Redis. Provide an integer (in seconds) or a `datetime.timedelta` object. kwargs: Optional(Dict[str, Any]) additional keyword arguments that are required to connect to your redis server. Same as `redis.Redis()`. request_coder: (Optional[`coders.Coder`]) coder for requests stored in Redis. response_coder: (Optional[`coders.Coder`]) coder for decoding responses received from Redis. source_caller: (Optional[`Caller`]): The source caller using this Redis cache in case of fetching the cache request to store in Redis. """ self.request_coder = request_coder self.response_coder = response_coder self.redis_caller = _RedisCaller( host, port, time_to_live, request_coder=self.request_coder, response_coder=self.response_coder, kwargs=kwargs, source_caller=source_caller, mode=_RedisMode.WRITE) def expand( self, elements: beam.PCollection[tuple[RequestT, ResponseT]] ) -> beam.PCollection[ResponseT]: return elements | RequestResponseIO(self.redis_caller) def ensure_coders_exist(request_coder): """checks if the coder exists to encode the request for caching.""" if not request_coder: raise ValueError( 'need request coder to be able to use ' 'Cache with RequestResponseIO.') class RedisCache(Cache): """Configure cache using Redis for :class:`apache_beam.io.requestresponse.RequestResponseIO`.""" def __init__( self, host: str, port: int, time_to_live: Union[int, timedelta] = DEFAULT_CACHE_ENTRY_TTL_SEC, *, request_coder: Optional[coders.Coder] = None, response_coder: Optional[coders.Coder] = None, **kwargs, ): """ Args: host (str): The hostname or IP address of the Redis server. port (int): The port number of the Redis server. time_to_live: `(Union[int, timedelta])` The time-to-live (TTL) for records stored in Redis. Provide an integer (in seconds) or a `datetime.timedelta` object. request_coder: (Optional[`coders.Coder`]) coder for encoding requests. response_coder: (Optional[`coders.Coder`]) coder for decoding responses received from Redis. kwargs: Optional additional keyword arguments that are required to connect to your redis server. Same as `redis.Redis()`. """ self._host = host self._port = port self._time_to_live = time_to_live self._request_coder = request_coder self._response_coder = response_coder self._kwargs = kwargs if kwargs else {} self._source_caller = None if redis is None: raise ImportError( 'Failed to import redis. You can ensure it is ' 'installed by installing the redis beam extra') def get_read(self): """get_read returns a PTransform for reading from the cache.""" ensure_coders_exist(self._request_coder) return _ReadFromRedis( self._host, self._port, time_to_live=self._time_to_live, kwargs=self._kwargs, request_coder=self._request_coder, response_coder=self._response_coder, source_caller=self._source_caller) def get_write(self): """returns a PTransform for writing to the cache.""" ensure_coders_exist(self._request_coder) return _WriteToRedis( self._host, self._port, time_to_live=self._time_to_live, kwargs=self._kwargs, request_coder=self._request_coder, response_coder=self._response_coder, source_caller=self._source_caller) @property def source_caller(self): return self._source_caller @source_caller.setter def source_caller(self, source_caller: Caller): self._source_caller = source_caller @property def request_coder(self): return self._request_coder @request_coder.setter def request_coder(self, request_coder: coders.Coder): self._request_coder = request_coder class FlattenBatch(beam.DoFn): """Flatten a batched PCollection.""" def process(self, elements, *args, **kwargs): for element in elements: yield element class RequestResponseIO(beam.PTransform[beam.PCollection[RequestT], beam.PCollection[ResponseT]]): """A :class:`RequestResponseIO` transform to read and write to APIs. Processes an input :class:`~apache_beam.pvalue.PCollection` of requests by making a call to the API as defined in `Caller`'s `__call__` method and returns a :class:`~apache_beam.pvalue.PCollection` of responses. """ def __init__( self, caller: Caller[RequestT, ResponseT], timeout: Optional[float] = DEFAULT_TIMEOUT_SECS, should_backoff: Optional[ShouldBackOff] = None, repeater: Repeater = ExponentialBackOffRepeater(), cache: Optional[Cache] = None, throttler: PreCallThrottler = DefaultThrottler(), ): """ Instantiates a RequestResponseIO transform. Args: caller: an implementation of `Caller` object that makes call to the API. timeout (float): timeout value in seconds to wait for response from API. should_backoff: (Optional) provides methods for backoff. repeater: provides method to repeat failed requests to API due to service errors. Defaults to :class:`apache_beam.io.requestresponse.ExponentialBackOffRepeater` to repeat requests with exponential backoff. cache: (Optional) a `~apache_beam.io.requestresponse.Cache` object to use the appropriate cache. throttler: provides methods to pre-throttle a request. Defaults to :class:`apache_beam.io.requestresponse.DefaultThrottler` for client-side adaptive throttling using :class:`apache_beam.io.components.adaptive_throttler.AdaptiveThrottler` """ self._caller = caller self._timeout = timeout self._should_backoff = should_backoff if repeater: self._repeater = repeater else: self._repeater = NoOpsRepeater() self._cache = cache self._throttler = throttler self._batching_kwargs = self._caller.batch_elements_kwargs() def expand( self, requests: beam.PCollection[RequestT]) -> beam.PCollection[ResponseT]: # TODO(riteshghorse): handle Throttle PTransforms when available. if self._cache: self._cache.source_caller = self._caller inputs = requests if self._cache: # read from cache. outputs = inputs | self._cache.get_read() # filter responses that are None and send them to the Call transform # to fetch a value from external service. cached_responses, inputs = (outputs | beam.ParDo(_FilterCacheReadFn() ).with_outputs( 'cache_misses', main='cached_responses')) # Batch elements if batching is enabled. if self._batching_kwargs: inputs = inputs | BatchElements(**self._batching_kwargs) if isinstance(self._throttler, DefaultThrottler): # DefaultThrottler applies throttling in the DoFn of # Call PTransform. responses = ( inputs | _Call( caller=self._caller, timeout=self._timeout, should_backoff=self._should_backoff, repeater=self._repeater, throttler=self._throttler)) else: # No throttling mechanism. The requests are made to the external source # as they come. responses = ( inputs | _Call( caller=self._caller, timeout=self._timeout, should_backoff=self._should_backoff, repeater=self._repeater)) # if batching is enabled then handle accordingly. if self._batching_kwargs: responses = responses | "FlattenBatch" >> beam.ParDo(FlattenBatch()) if self._cache: # write to cache. _ = responses | self._cache.get_write() return (cached_responses, responses) | beam.Flatten() return responses ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/requestresponse_it_test.py0000644000175100017510000002460615206113032023650 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import base64 import logging import sys import time import typing import unittest from dataclasses import dataclass from typing import Union import pytest import urllib3 import apache_beam as beam from apache_beam.coders import coders from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline # pylint: disable=ungrouped-imports try: from testcontainers.redis import RedisContainer from apache_beam.io.requestresponse import Caller from apache_beam.io.requestresponse import RedisCache from apache_beam.io.requestresponse import RequestResponseIO from apache_beam.io.requestresponse import UserCodeExecutionException from apache_beam.io.requestresponse import UserCodeQuotaException except ImportError: raise unittest.SkipTest('RequestResponseIO dependencies are not installed.') _HTTP_PATH = '/v1/echo' _PAYLOAD = base64.b64encode(bytes('payload', 'utf-8')) _HTTP_ENDPOINT_ADDRESS_FLAG = '--httpEndpointAddress' _LOGGER = logging.getLogger(__name__) class EchoITOptions(PipelineOptions): """Shared options for running integration tests on a deployed ``EchoServiceGrpc`` See https://github.com/apache/beam/tree/master/.test -infra/mock-apis#integration for details on how to acquire values required by ``EchoITOptions``. """ @classmethod def _add_argparse_args(cls, parser) -> None: parser.add_argument( _HTTP_ENDPOINT_ADDRESS_FLAG, default='http://10.138.0.32:8080', dest='http_endpoint_address', help='The HTTP address of the Echo API endpoint; must being with ' 'http(s)://') parser.add_argument( '--neverExceedQuotaId', default='echo-should-never-exceed-quota', dest='never_exceed_quota_id', help='The ID for an allocated quota that should never exceed.') parser.add_argument( '--shouldExceedQuotaId', default='echo-should-exceed-quota', dest='should_exceed_quota_id', help='The ID for an allocated quota that should exceed.') @dataclass class EchoResponse: id: str payload: bytes # TODO(riteshghorse,damondouglas) replace Echo(Request|Response) with proto # generated classes from .test-infra/mock-apis: class Request(typing.NamedTuple): id: str payload: bytes class EchoHTTPCaller(Caller[Request, EchoResponse]): """Implements ``Caller`` to call the ``EchoServiceGrpc``'s HTTP handler. The purpose of ``EchoHTTPCaller`` is to support integration tests. """ def __init__(self, url: str): self.url = url + _HTTP_PATH def __call__(self, request: Request, *args, **kwargs) -> EchoResponse: """Overrides ``Caller``'s call method invoking the ``EchoServiceGrpc``'s HTTP handler with an ``EchoRequest``, returning either a successful ``EchoResponse`` or throwing either a ``UserCodeExecutionException``, ``UserCodeTimeoutException``, or a ``UserCodeQuotaException``. """ try: http = urllib3.PoolManager() resp = http.request( "POST", self.url, json={ "id": request.id, "payload": str(request.payload, 'utf-8') }, retries=False) if resp.status < 300: resp_body = resp.json() resp_id = resp_body['id'] payload = resp_body['payload'] return EchoResponse(id=resp_id, payload=bytes(payload, 'utf-8')) if resp.status == 429: # Too Many Requests raise UserCodeQuotaException(resp.reason) else: raise UserCodeExecutionException(resp.status, resp.reason, request) except urllib3.exceptions.HTTPError as e: raise UserCodeExecutionException(e) class ValidateResponse(beam.DoFn): """Validates response received from Mock API server.""" def process(self, element, *args, **kwargs): if (element.id != 'echo-should-never-exceed-quota' or element.payload != _PAYLOAD): raise ValueError( 'got EchoResponse(id: %s, payload: %s), want ' 'EchoResponse(id: echo-should-never-exceed-quota, ' 'payload: %s' % (element.id, element.payload, _PAYLOAD)) @pytest.mark.uses_mock_api class EchoHTTPCallerTestIT(unittest.TestCase): options: Union[EchoITOptions, None] = None client: Union[EchoHTTPCaller, None] = None @classmethod def setUpClass(cls) -> None: cls.options = EchoITOptions() http_endpoint_address = cls.options.http_endpoint_address if not http_endpoint_address or http_endpoint_address == '': raise unittest.SkipTest(f'{_HTTP_ENDPOINT_ADDRESS_FLAG} is required.') cls.client = EchoHTTPCaller(http_endpoint_address) @classmethod def _get_client_and_options(cls) -> tuple[EchoHTTPCaller, EchoITOptions]: assert cls.options is not None assert cls.client is not None return cls.client, cls.options def test_request_response_io(self): client, options = EchoHTTPCallerTestIT._get_client_and_options() req = Request(id=options.never_exceed_quota_id, payload=_PAYLOAD) with TestPipeline(is_integration_test=True) as test_pipeline: output = ( test_pipeline | 'Create PCollection' >> beam.Create([req]) | 'RRIO Transform' >> RequestResponseIO(client) | 'Validate' >> beam.ParDo(ValidateResponse())) self.assertIsNotNone(output) class ValidateCacheResponses(beam.DoFn): """Validates that the responses are fetched from the cache.""" def process(self, element, *args, **kwargs): if not element[1] or 'cached-' not in element[1]: raise ValueError( 'responses not fetched from cache even though cache ' 'entries are present.') class ValidateCallerResponses(beam.DoFn): """Validates that the responses are fetched from the caller.""" def process(self, element, *args, **kwargs): if not element[1] or 'ACK-' not in element[1]: raise ValueError('responses not fetched from caller when they should.') class FakeCallerForCache(Caller[str, str]): def __init__(self, use_cache: bool = False): self.use_cache = use_cache def __enter__(self): pass def __call__(self, element, *args, **kwargs): if self.use_cache: return None, None return element, 'ACK-{element}' def __exit__(self, exc_type, exc_val, exc_tb): pass @pytest.mark.uses_testcontainer class TestRedisCache(unittest.TestCase): def setUp(self) -> None: self.retries = 5 self._start_container() def test_rrio_cache_all_miss(self): """Cache is empty so all responses are fetched from caller.""" caller = FakeCallerForCache() req = ['redis', 'cachetools', 'memcache'] cache = RedisCache( self.host, self.port, time_to_live=30, request_coder=coders.StrUtf8Coder(), response_coder=coders.StrUtf8Coder()) with TestPipeline(is_integration_test=True) as p: _ = ( p | beam.Create(req) | RequestResponseIO(caller, cache=cache) | beam.ParDo(ValidateCallerResponses())) def test_rrio_cache_all_hit(self): """Validate that records are fetched from cache.""" caller = FakeCallerForCache() requests = ['foo', 'bar'] responses = ['cached-foo', 'cached-bar'] coder = coders.StrUtf8Coder() for i in range(len(requests)): enc_req = coder.encode(requests[i]) enc_resp = coder.encode(responses[i]) self.client.setex(enc_req, 120, enc_resp) cache = RedisCache( self.host, self.port, time_to_live=30, request_coder=coders.StrUtf8Coder(), response_coder=coders.StrUtf8Coder()) with TestPipeline(is_integration_test=True) as p: _ = ( p | beam.Create(requests) | RequestResponseIO(caller, cache=cache) | beam.ParDo(ValidateCacheResponses())) def test_rrio_cache_miss_and_hit(self): """Run two back-to-back pipelines, one with pulling the data from caller and other from the cache.""" caller = FakeCallerForCache() requests = ['beam', 'flink', 'spark'] cache = RedisCache( self.host, self.port, request_coder=coders.StrUtf8Coder(), response_coder=coders.StrUtf8Coder()) with TestPipeline(is_integration_test=True) as p: _ = ( p | beam.Create(requests) | RequestResponseIO(caller, cache=cache) | beam.ParDo(ValidateCallerResponses())) caller = FakeCallerForCache(use_cache=True) with TestPipeline(is_integration_test=True) as p: _ = ( p | beam.Create(requests) | RequestResponseIO(caller, cache=cache) | beam.ParDo(ValidateCallerResponses())) def test_rrio_no_coder_exception(self): caller = FakeCallerForCache() requests = ['beam', 'flink', 'spark'] cache = RedisCache(self.host, self.port) with self.assertRaises(ValueError): test_pipeline = beam.Pipeline() _ = ( test_pipeline | beam.Create(requests) | RequestResponseIO(caller, cache=cache)) res = test_pipeline.run() res.wait_until_finish() def tearDown(self) -> None: self.container.stop() def _start_container(self): for i in range(self.retries): try: self.container = RedisContainer(image='redis:7.2.4') self.container.start() self.host = self.container.get_container_host_ip() self.port = self.container.get_exposed_port(6379) self.client = self.container.get_client() break except Exception as e: if i == self.retries - 1: _LOGGER.error('Unable to start redis container for RRIO tests.') raise e # Add a small delay between retries to avoid rapid successive failures time.sleep(2) if __name__ == '__main__': unittest.main(argv=sys.argv[:1]) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/requestresponse_test.py0000644000175100017510000001365415206113032023155 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import time import unittest from tenacity import retry from tenacity import retry_if_exception_type from tenacity import stop_after_attempt import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline # pylint: disable=ungrouped-imports try: from google.api_core.exceptions import TooManyRequests from apache_beam.io.requestresponse import Caller from apache_beam.io.requestresponse import DefaultThrottler from apache_beam.io.requestresponse import RequestResponseIO from apache_beam.io.requestresponse import retry_on_exception except ImportError: raise unittest.SkipTest('RequestResponseIO dependencies are not installed.') _LOGGER = logging.getLogger() MAX_TEST_RETRIES = 3 class AckCaller(Caller[str, str]): """AckCaller acknowledges the incoming request by returning a request with ACK.""" def __enter__(self): pass def __call__(self, request: str): return f"ACK: {request}" def __exit__(self, exc_type, exc_val, exc_tb): return None class CallerWithTimeout(AckCaller): """CallerWithTimeout sleeps for 2 seconds before responding. Used to test timeout in RequestResponseIO.""" def __call__(self, request: str, *args, **kwargs): time.sleep(2) return f"ACK: {request}" class CallerWithRuntimeError(AckCaller): """CallerWithRuntimeError raises a `RuntimeError` for RequestResponseIO to raise a UserCodeExecutionException.""" def __call__(self, request: str, *args, **kwargs): if not request: raise RuntimeError("Exception expected, not an error.") class CallerThatRetries(AckCaller): def __init__(self): self.count = -1 def __call__(self, request: str, *args, **kwargs): try: pass except Exception as e: raise e finally: self.count += 1 raise TooManyRequests('retries = %d' % self.count) class TestCaller(unittest.TestCase): def test_valid_call(self): caller = AckCaller() with TestPipeline() as test_pipeline: output = ( test_pipeline | beam.Create(["sample_request"]) | RequestResponseIO(caller=caller)) self.assertIsNotNone(output) def test_call_timeout(self): caller = CallerWithTimeout() with self.assertRaisesRegex(Exception, "Timeout"): with TestPipeline() as test_pipeline: _ = ( test_pipeline | beam.Create(["timeout_request"]) | RequestResponseIO(caller=caller, timeout=1)) def test_call_runtime_error(self): caller = CallerWithRuntimeError() with self.assertRaisesRegex(Exception, "could not complete request"): with TestPipeline() as test_pipeline: _ = ( test_pipeline | beam.Create([""]) | RequestResponseIO(caller=caller)) def test_retry_on_exception(self): self.assertFalse(retry_on_exception(RuntimeError())) self.assertTrue(retry_on_exception(TooManyRequests("HTTP 429"))) def test_caller_backoff_retry_strategy(self): caller = CallerThatRetries() with self.assertRaises(Exception) as cm: with TestPipeline() as test_pipeline: _ = ( test_pipeline | beam.Create(["sample_request"]) | RequestResponseIO(caller=caller)) self.assertRegex(str(cm.exception), 'retries = 2') def test_caller_no_retry_strategy(self): caller = CallerThatRetries() with self.assertRaises(Exception) as cm: with TestPipeline() as test_pipeline: _ = ( test_pipeline | beam.Create(["sample_request"]) | RequestResponseIO(caller=caller, repeater=None)) self.assertRegex(str(cm.exception), 'retries = 0') @retry( retry=retry_if_exception_type(IndexError), reraise=True, stop=stop_after_attempt(MAX_TEST_RETRIES)) def test_default_throttler(self): caller = CallerWithTimeout() throttler = DefaultThrottler( window_ms=10000, bucket_ms=5000, overload_ratio=1) # manually override the number of received requests for testing. throttler.throttler._all_requests.add(time.time() * 1000, 100) # TODO(https://github.com/apache/beam/issues/34549): This test relies on # metrics filtering which doesn't work on Prism yet because Prism renames # steps (e.g. "Do" becomes "ref_AppliedPTransform_Do_7"). # https://github.com/apache/beam/blob/5f9cd73b7c9a2f37f83971ace3a399d633201dd1/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py#L1590 test_pipeline = TestPipeline('FnApiRunner') _ = ( test_pipeline | beam.Create(['sample_request']) | RequestResponseIO(caller=caller, throttler=throttler)) result = test_pipeline.run() result.wait_until_finish() metrics = result.metrics().query( beam.metrics.MetricsFilter().with_name('throttled_requests')) self.assertEqual(metrics['counters'][0].committed, 1) metrics = result.metrics().query( beam.metrics.MetricsFilter().with_name('cumulativeThrottlingSeconds')) self.assertGreater(metrics['counters'][0].committed, 0) metrics = result.metrics().query( beam.metrics.MetricsFilter().with_name('responses')) self.assertEqual(metrics['counters'][0].committed, 1) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/restriction_trackers.py0000644000175100017510000001372315206113032023107 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """`iobase.RestrictionTracker` implementations provided with Apache Beam.""" # pytype: skip-file from apache_beam.io.iobase import RestrictionProgress from apache_beam.io.iobase import RestrictionTracker from apache_beam.io.range_trackers import OffsetRangeTracker class OffsetRange(object): def __init__(self, start, stop): if start > stop: raise ValueError( 'Start offset must be not be larger than the stop offset. ' 'Received %d and %d respectively.' % (start, stop)) self.start = start self.stop = stop def __eq__(self, other): if not isinstance(other, OffsetRange): return False return self.start == other.start and self.stop == other.stop def __hash__(self): return hash((type(self), self.start, self.stop)) def __repr__(self): return 'OffsetRange(start=%s, stop=%s)' % (self.start, self.stop) def split(self, desired_num_offsets_per_split, min_num_offsets_per_split=1): current_split_start = self.start max_split_size = max( desired_num_offsets_per_split, min_num_offsets_per_split) while current_split_start < self.stop: current_split_stop = min(current_split_start + max_split_size, self.stop) remaining = self.stop - current_split_stop # Avoiding a small split at the end. if (remaining < desired_num_offsets_per_split // 4 or remaining < min_num_offsets_per_split): current_split_stop = self.stop yield OffsetRange(current_split_start, current_split_stop) current_split_start = current_split_stop def split_at(self, split_pos) -> tuple['OffsetRange', 'OffsetRange']: return OffsetRange(self.start, split_pos), OffsetRange(split_pos, self.stop) def new_tracker(self): return OffsetRangeTracker(self.start, self.stop) def size(self): return self.stop - self.start class OffsetRestrictionTracker(RestrictionTracker): """An `iobase.RestrictionTracker` implementations for an offset range. Offset range is represented as OffsetRange. """ def __init__(self, offset_range: OffsetRange) -> None: assert isinstance(offset_range, OffsetRange), offset_range self._range = offset_range self._current_position = None self._last_claim_attempt = None self._checkpointed = False def check_done(self): if (self._range.start != self._range.stop and (self._last_claim_attempt is None or self._last_claim_attempt < self._range.stop - 1)): raise ValueError( 'OffsetRestrictionTracker is not done since work in range [%s, %s) ' 'has not been claimed.' % ( self._last_claim_attempt if self._last_claim_attempt is not None else self._range.start, self._range.stop)) def current_restriction(self): return self._range def current_progress(self) -> RestrictionProgress: if self._current_position is None: fraction = 0.0 elif self._range.stop == self._range.start: # If self._current_position is not None, we must be done. fraction = 1.0 else: fraction = ( float(self._current_position - self._range.start) / (self._range.stop - self._range.start)) return RestrictionProgress(fraction=fraction) def start_position(self): return self._range.start def stop_position(self): return self._range.stop def try_claim(self, position): if (self._last_claim_attempt is not None and position <= self._last_claim_attempt): raise ValueError( 'Positions claimed should strictly increase. Trying to claim ' 'position %d while last claim attempt was %d.' % (position, self._last_claim_attempt)) self._last_claim_attempt = position if position < self._range.start: raise ValueError( 'Position to be claimed cannot be smaller than the start position ' 'of the range. Tried to claim position %r for the range [%r, %r)' % (position, self._range.start, self._range.stop)) if self._range.start <= position < self._range.stop: self._current_position = position return True return False def try_split(self, fraction_of_remainder): if not self._checkpointed: if self._last_claim_attempt is None: cur = self._range.start - 1 else: cur = self._last_claim_attempt split_point = ( cur + int(max(1, (self._range.stop - cur) * fraction_of_remainder))) if split_point < self._range.stop: if fraction_of_remainder == 0: self._checkpointed = True self._range, residual_range = self._range.split_at(split_point) return self._range, residual_range def is_bounded(self): return True class UnsplittableRestrictionTracker(RestrictionTracker): """An `iobase.RestrictionTracker` that wraps another but does not split.""" def __init__(self, underling_tracker): self._underling_tracker = underling_tracker def try_split(self, fraction_of_remainder): return False # __getattribute__ is used rather than __getattr__ to override the # stubs in the baseclass. def __getattribute__(self, name): if name.startswith('_') or name in ('try_split', ): return super().__getattribute__(name) else: return getattr(self._underling_tracker, name) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/restriction_trackers_test.py0000644000175100017510000001617315206113032024150 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the range_trackers module.""" # pytype: skip-file import logging import unittest from apache_beam.io.restriction_trackers import OffsetRange from apache_beam.io.restriction_trackers import OffsetRestrictionTracker class OffsetRangeTest(unittest.TestCase): def test_create(self): OffsetRange(0, 10) OffsetRange(10, 10) OffsetRange(10, 100) with self.assertRaises(ValueError): OffsetRange(10, 9) def test_split_respects_desired_num_splits(self): range = OffsetRange(10, 100) splits = list(range.split(desired_num_offsets_per_split=25)) self.assertEqual(4, len(splits)) self.assertIn(OffsetRange(10, 35), splits) self.assertIn(OffsetRange(35, 60), splits) self.assertIn(OffsetRange(60, 85), splits) self.assertIn(OffsetRange(85, 100), splits) def test_split_respects_min_num_splits(self): range = OffsetRange(10, 100) splits = list( range.split( desired_num_offsets_per_split=5, min_num_offsets_per_split=25)) self.assertEqual(3, len(splits)) self.assertIn(OffsetRange(10, 35), splits) self.assertIn(OffsetRange(35, 60), splits) self.assertIn(OffsetRange(60, 100), splits) def test_split_no_small_split_at_end(self): range = OffsetRange(10, 90) splits = list(range.split(desired_num_offsets_per_split=25)) self.assertEqual(3, len(splits)) self.assertIn(OffsetRange(10, 35), splits) self.assertIn(OffsetRange(35, 60), splits) self.assertIn(OffsetRange(60, 90), splits) def test_split_at(self): range = OffsetRange(0, 10) cur, residual = range.split_at(5) self.assertEqual(cur, OffsetRange(0, 5)) self.assertEqual(residual, OffsetRange(5, 10)) class OffsetRestrictionTrackerTest(unittest.TestCase): def test_try_claim(self): tracker = OffsetRestrictionTracker(OffsetRange(100, 200)) self.assertEqual(OffsetRange(100, 200), tracker.current_restriction()) self.assertTrue(tracker.try_claim(100)) self.assertTrue(tracker.try_claim(150)) self.assertTrue(tracker.try_claim(199)) self.assertFalse(tracker.try_claim(200)) def test_checkpoint_unstarted(self): tracker = OffsetRestrictionTracker(OffsetRange(100, 200)) _, checkpoint = tracker.try_split(0) self.assertEqual(OffsetRange(100, 100), tracker.current_restriction()) self.assertEqual(OffsetRange(100, 200), checkpoint) def test_checkpoint_just_started(self): tracker = OffsetRestrictionTracker(OffsetRange(100, 200)) self.assertTrue(tracker.try_claim(100)) _, checkpoint = tracker.try_split(0) self.assertEqual(OffsetRange(100, 101), tracker.current_restriction()) self.assertEqual(OffsetRange(101, 200), checkpoint) def test_checkpoint_regular(self): tracker = OffsetRestrictionTracker(OffsetRange(100, 200)) self.assertTrue(tracker.try_claim(105)) self.assertTrue(tracker.try_claim(110)) _, checkpoint = tracker.try_split(0) self.assertEqual(OffsetRange(100, 111), tracker.current_restriction()) self.assertEqual(OffsetRange(111, 200), checkpoint) def test_checkpoint_claimed_last(self): tracker = OffsetRestrictionTracker(OffsetRange(100, 200)) self.assertTrue(tracker.try_claim(105)) self.assertTrue(tracker.try_claim(110)) self.assertTrue(tracker.try_claim(199)) checkpoint = tracker.try_split(0) self.assertEqual(OffsetRange(100, 200), tracker.current_restriction()) self.assertEqual(None, checkpoint) def test_checkpoint_after_failed_claim(self): tracker = OffsetRestrictionTracker(OffsetRange(100, 200)) self.assertTrue(tracker.try_claim(105)) self.assertTrue(tracker.try_claim(110)) self.assertTrue(tracker.try_claim(160)) self.assertFalse(tracker.try_claim(240)) self.assertIsNone(tracker.try_split(0)) self.assertTrue(OffsetRange(100, 200), tracker.current_restriction()) def test_non_monotonic_claim(self): with self.assertRaises(ValueError): tracker = OffsetRestrictionTracker(OffsetRange(100, 200)) self.assertTrue(tracker.try_claim(105)) self.assertTrue(tracker.try_claim(110)) self.assertTrue(tracker.try_claim(103)) def test_claim_before_starting_range(self): with self.assertRaises(ValueError): tracker = OffsetRestrictionTracker(OffsetRange(100, 200)) tracker.try_claim(90) def test_check_done_after_try_claim_past_end_of_range(self): tracker = OffsetRestrictionTracker(OffsetRange(100, 200)) self.assertTrue(tracker.try_claim(150)) self.assertTrue(tracker.try_claim(175)) self.assertFalse(tracker.try_claim(220)) tracker.check_done() def test_check_done_after_try_claim_right_before_end_of_range(self): tracker = OffsetRestrictionTracker(OffsetRange(100, 200)) self.assertTrue(tracker.try_claim(150)) self.assertTrue(tracker.try_claim(175)) self.assertTrue(tracker.try_claim(199)) tracker.check_done() def test_check_done_when_not_done(self): tracker = OffsetRestrictionTracker(OffsetRange(100, 200)) self.assertTrue(tracker.try_claim(150)) self.assertTrue(tracker.try_claim(175)) with self.assertRaises(ValueError): tracker.check_done() def test_check_done_with_no_claims(self): tracker = OffsetRestrictionTracker(OffsetRange(100, 200)) with self.assertRaises(ValueError): tracker.check_done() def test_try_split(self): tracker = OffsetRestrictionTracker(OffsetRange(100, 200)) tracker.try_claim(100) cur, residual = tracker.try_split(0.5) self.assertEqual(OffsetRange(100, 150), cur) self.assertEqual(OffsetRange(150, 200), residual) self.assertEqual(cur, tracker.current_restriction()) def test_try_split_when_restriction_is_done(self): tracker = OffsetRestrictionTracker(OffsetRange(100, 200)) tracker.try_claim(199) self.assertIsNone(tracker.try_split(0.5)) tracker.try_claim(200) self.assertIsNone(tracker.try_split(0.5)) def test_check_done_empty_range(self): tracker = OffsetRestrictionTracker(OffsetRange(0, 0)) tracker.check_done() def test_try_claim_empty_range(self): tracker = OffsetRestrictionTracker(OffsetRange(0, 0)) self.assertFalse(tracker.try_claim(0)) def test_checkpoint_empty_range(self): tracker = OffsetRestrictionTracker(OffsetRange(0, 0)) self.assertIsNone(tracker.try_split(0)) self.assertFalse(tracker.try_claim(0)) self.assertIsNone(tracker.try_split(0)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/snowflake.py0000644000175100017510000004420315206113032020632 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Snowflake transforms tested against Flink portable runner. **Setup** Transforms provided in this module are cross-language transforms implemented in the Beam Java SDK. During the pipeline construction, Python SDK will connect to a Java expansion service to expand these transforms. To facilitate this, a small amount of setup is needed before using these transforms in a Beam Python pipeline. There are several ways to setup cross-language Snowflake transforms. * Option 1: use the default expansion service * Option 2: specify a custom expansion service See below for details regarding each of these options. *Option 1: Use the default expansion service* This is the recommended and easiest setup option for using Python Snowflake transforms.This option requires following pre-requisites before running the Beam pipeline. * Install Java runtime in the computer from where the pipeline is constructed and make sure that 'java' command is available or set JAVA_HOME environment variable. In this option, Python SDK will either download (for released Beam version) or build (when running from a Beam Git clone) a expansion service jar and use that to expand transforms. Currently Snowflake transforms use the 'beam-sdks-java-io-snowflake-expansion-service' jar for this purpose. *Option 2: specify a custom expansion service* In this option, you startup your own expansion service and provide that as a parameter when using the transforms provided in this module. This option requires following pre-requisites before running the Beam pipeline. * Startup your own expansion service. * Update your pipeline to provide the expansion service address when initiating Snowflake transforms provided in this module. Flink Users can use the built-in Expansion Service of the Flink Runner's Job Server. If you start Flink's Job Server, the expansion service will be started on port 8097. For a different address, please set the expansion_service parameter. **More information** For more information regarding cross-language transforms see: - https://beam.apache.org/roadmap/portability/ For more information specific to Flink runner see: - https://beam.apache.org/documentation/runners/flink/ """ # pytype: skip-file from typing import Iterable from typing import NamedTuple from typing import Optional import apache_beam as beam from apache_beam.transforms.external import BeamJarExpansionService from apache_beam.transforms.external import ExternalTransform from apache_beam.transforms.external import NamedTupleBasedPayloadBuilder __all__ = [ 'ReadFromSnowflake', 'WriteToSnowflake', 'WriteDisposition', 'CreateDisposition', ] def default_io_expansion_service(): return BeamJarExpansionService( 'sdks:java:io:snowflake:expansion-service:shadowJar') ReadFromSnowflakeSchema = NamedTuple( 'ReadFromSnowflakeSchema', [ ('server_name', str), ('schema', str), ('database', str), ('staging_bucket_name', str), ('storage_integration_name', str), ('username', Optional[str]), ('password', Optional[str]), ('private_key_path', Optional[str]), ('raw_private_key', Optional[str]), ('private_key_passphrase', Optional[str]), ('o_auth_token', Optional[str]), ('table', Optional[str]), ('query', Optional[str]), ('role', Optional[str]), ('warehouse', Optional[str]), ]) class ReadFromSnowflake(beam.PTransform): """ An external PTransform which reads from Snowflake. """ URN = 'beam:transform:org.apache.beam:snowflake_read:v1' def __init__( self, server_name, schema, database, staging_bucket_name, storage_integration_name, csv_mapper, username=None, password=None, private_key_path=None, raw_private_key=None, private_key_passphrase=None, o_auth_token=None, table=None, query=None, role=None, warehouse=None, expansion_service=None): """ Initializes a read operation from Snowflake. Required parameters: :param server_name: full Snowflake server name with the following format https://account.region.gcp.snowflakecomputing.com. :param schema: name of the Snowflake schema in the database to use. :param database: name of the Snowflake database to use. :param staging_bucket_name: name of the Google Cloud Storage bucket. Bucket will be used as a temporary location for storing CSV files. Those temporary directories will be named 'sf_copy_csv_DATE_TIME_RANDOMSUFFIX' and they will be removed automatically once Read operation finishes. :param storage_integration_name: is the name of storage integration object created according to Snowflake documentation. :param csv_mapper: specifies a function which must translate user-defined object to array of strings. SnowflakeIO uses a COPY INTO statement to move data from a Snowflake table to Google Cloud Storage as CSV files.These files are then downloaded via FileIO and processed line by line. Each line is split into an array of Strings using the OpenCSV The csv_mapper function job is to give the user the possibility to convert the array of Strings to a user-defined type, ie. GenericRecord for Avro or Parquet files, or custom objects. Example: def csv_mapper(strings_array) return User(strings_array[0], int(strings_array[1]))) :param table: specifies a Snowflake table name. :param query: specifies a Snowflake custom SQL query. :param role: specifies a Snowflake role. :param warehouse: specifies a Snowflake warehouse name. :param expansion_service: specifies URL of expansion service. Authentication parameters: :param username: specifies username for username/password authentication method. :param password: specifies password for username/password authentication method. :param private_key_path: specifies a private key file for key/ pair authentication method. :param raw_private_key: specifies a private key for key/ pair authentication method. :param private_key_passphrase: specifies password for key/ pair authentication method. :param o_auth_token: specifies access token for OAuth authentication method. """ verify_credentials( username=username, password=password, private_key_path=private_key_path, raw_private_key=raw_private_key, o_auth_token=o_auth_token, ) self.params = ReadFromSnowflakeSchema( server_name=server_name, schema=schema, database=database, staging_bucket_name=staging_bucket_name, storage_integration_name=storage_integration_name, username=username, password=password, private_key_path=private_key_path, raw_private_key=raw_private_key, private_key_passphrase=private_key_passphrase, o_auth_token=o_auth_token, table=table, query=query, role=role, warehouse=warehouse, ) self.csv_mapper = csv_mapper self.expansion_service = expansion_service or default_io_expansion_service() def expand(self, pbegin): return ( pbegin | ExternalTransform( self.URN, NamedTupleBasedPayloadBuilder(self.params), self.expansion_service, ) | 'CSV to array mapper' >> beam.Map(lambda csv: csv.split(b',')) | 'CSV mapper' >> beam.Map(self.csv_mapper)) WriteToSnowflakeSchema = NamedTuple( 'WriteToSnowflakeSchema', [ ('server_name', str), ('schema', str), ('database', str), ('staging_bucket_name', str), ('storage_integration_name', str), ('create_disposition', str), ('write_disposition', str), ('table_schema', str), ('username', Optional[str]), ('password', Optional[str]), ('private_key_path', Optional[str]), ('raw_private_key', Optional[str]), ('private_key_passphrase', Optional[str]), ('o_auth_token', Optional[str]), ('table', Optional[str]), ('query', Optional[str]), ('role', Optional[str]), ('warehouse', Optional[str]), ], ) class WriteToSnowflake(beam.PTransform): """ An external PTransform which writes to Snowflake. """ URN = 'beam:transform:org.apache.beam:snowflake_write:v1' def __init__( self, server_name, schema, database, staging_bucket_name, storage_integration_name, create_disposition, write_disposition, table_schema, user_data_mapper, username=None, password=None, private_key_path=None, raw_private_key=None, private_key_passphrase=None, o_auth_token=None, table=None, query=None, role=None, warehouse=None, expansion_service=None, ): # pylint: disable=line-too-long """ Initializes a write operation to Snowflake. Required parameters: :param server_name: full Snowflake server name with the following format https://account.region.gcp.snowflakecomputing.com. :param schema: name of the Snowflake schema in the database to use. :param database: name of the Snowflake database to use. :param staging_bucket_name: name of the Google Cloud Storage bucket. Bucket will be used as a temporary location for storing CSV files. :param storage_integration_name: is the name of a Snowflake storage integration object created according to Snowflake documentation for the GCS bucket. :param user_data_mapper: specifies a function which maps data from a PCollection to an array of String values before the write operation saves the data to temporary .csv files. Example: def user_data_mapper(user): return [user.name, str(user.age)] :param table: specifies a Snowflake table name :param query: specifies a custom SQL query :param role: specifies a Snowflake role. :param warehouse: specifies a Snowflake warehouse name. :param expansion_service: specifies URL of expansion service. Authentication parameters: :param username: specifies username for username/password authentication method. :param password: specifies password for username/password authentication method. :param private_key_path: specifies a private key file for key/ pair authentication method. :param raw_private_key: specifies a private key for key/ pair authentication method. :param private_key_passphrase: specifies password for key/ pair authentication method. :param o_auth_token: specifies access token for OAuth authentication method. Additional parameters: :param create_disposition: Defines the behaviour of the write operation if the target table does not exist. The following values are supported: CREATE_IF_NEEDED - default behaviour. The write operation checks whether the specified target table exists; if it does not, the write operation attempts to create the table Specify the schema for the target table using the table_schema parameter. CREATE_NEVER - The write operation fails if the target table does not exist. :param write_disposition: Defines the write behaviour based on the table where data will be written to. The following values are supported: APPEND - Default behaviour. Written data is added to the existing rows in the table, EMPTY - The target table must be empty; otherwise, the write operation fails, TRUNCATE - The write operation deletes all rows from the target table before writing to it. :param table_schema: When the create_disposition parameter is set to CREATE_IF_NEEDED, the table_schema parameter enables specifying the schema for the created target table. A table schema is as JSON with the following structure: .. code-block:: none {"schema":[ { "dataType":{"type":""}, "name":"", "nullable": }, ]} All supported data types: .. code-block:: json {"schema":[ {"dataType":{"type":"date"},"name":"","nullable":false}, {"dataType":{"type":"datetime"},"name":"","nullable":false}, {"dataType":{"type":"time"},"name":"","nullable":false}, {"dataType":{"type":"timestamp"},"name":"","nullable":false}, {"dataType":{"type":"timestamp_ltz"},"name":"","nullable":false}, {"dataType":{"type":"timestamp_ntz"},"name":"","nullable":false}, {"dataType":{"type":"timestamp_tz"},"name":"","nullable":false}, {"dataType":{"type":"boolean"},"name":"","nullable":false}, {"dataType":{"type":"decimal","precision":38,"scale":1},"name":"","nullable":true}, {"dataType":{"type":"double"},"name":"","nullable":false}, {"dataType":{"type":"float"},"name":"","nullable":false}, {"dataType":{"type":"integer","precision":38,"scale":0},"name":"","nullable":false}, {"dataType":{"type":"number","precision":38,"scale":1},"name":"","nullable":false}, {"dataType":{"type":"numeric","precision":38,"scale":2},"name":"","nullable":false}, {"dataType":{"type":"real"},"name":"","nullable":false}, {"dataType":{"type":"array"},"name":"","nullable":false}, {"dataType":{"type":"object"},"name":"","nullable":false}, {"dataType":{"type":"variant"},"name":"","nullable":true}, {"dataType":{"type":"binary","size":null},"name":"","nullable":false}, {"dataType":{"type":"char","length":1},"name":"","nullable":false}, {"dataType":{"type":"string","length":null},"name":"","nullable":false}, {"dataType":{"type":"text","length":null},"name":"","nullable":false}, {"dataType":{"type":"varbinary","size":null},"name":"","nullable":false}, {"dataType":{"type":"varchar","length":100},"name":"","nullable":false}, {"dataType":{"type":"geography"},"name":"","nullable":true}] } """ verify_credentials( username=username, password=password, private_key_path=private_key_path, raw_private_key=raw_private_key, o_auth_token=o_auth_token, ) WriteDisposition.VerifyParam(write_disposition) CreateDisposition.VerifyParam(create_disposition) self.params = WriteToSnowflakeSchema( server_name=server_name, schema=schema, database=database, staging_bucket_name=staging_bucket_name, storage_integration_name=storage_integration_name, create_disposition=create_disposition, write_disposition=write_disposition, table_schema=table_schema, username=username, password=password, private_key_path=private_key_path, raw_private_key=raw_private_key, private_key_passphrase=private_key_passphrase, o_auth_token=o_auth_token, table=table, query=query, role=role, warehouse=warehouse, ) self.user_data_mapper = user_data_mapper self.expansion_service = expansion_service or default_io_expansion_service() def expand(self, pbegin): return ( pbegin | 'User data mapper' >> beam.Map( self.user_data_mapper).with_output_types(Iterable[bytes]) | ExternalTransform( self.URN, NamedTupleBasedPayloadBuilder(self.params), self.expansion_service)) class CreateDisposition: """ Enum class for possible values of create dispositions: CREATE_IF_NEEDED: default behaviour. The write operation checks whether the specified target table exists; if it does not, the write operation attempts to create the table Specify the schema for the target table using the table_schema parameter. CREATE_NEVER: The write operation fails if the target table does not exist. """ CREATE_IF_NEEDED = 'CREATE_IF_NEEDED' CREATE_NEVER = 'CREATE_NEVER' @staticmethod def VerifyParam(field): if field and not hasattr(CreateDisposition, field): raise RuntimeError( 'Create disposition has to be one of the following values:' 'CREATE_IF_NEEDED, CREATE_NEVER. Got: {}'.format(field)) class WriteDisposition: """ Enum class for possible values of write dispositions: APPEND: Default behaviour. Written data is added to the existing rows in the table, EMPTY: The target table must be empty; otherwise, the write operation fails, TRUNCATE: The write operation deletes all rows from the target table before writing to it. """ APPEND = 'APPEND' EMPTY = 'EMPTY' TRUNCATE = 'TRUNCATE' @staticmethod def VerifyParam(field): if field and not hasattr(WriteDisposition, field): raise RuntimeError( 'Write disposition has to be one of the following values:' 'APPEND, EMPTY, TRUNCATE. Got: {}'.format(field)) def verify_credentials( username, password, private_key_path, raw_private_key, o_auth_token): if not (o_auth_token or (username and password) or (username and (private_key_path or raw_private_key))): raise RuntimeError('Snowflake credentials are not set correctly.') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/source_test_utils.py0000644000175100017510000006304415206113032022424 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Helper functions and test harnesses for source implementations. This module contains helper functions and test harnesses for checking correctness of source (a subclass of ``iobase.BoundedSource``) and range tracker (a subclass of``iobase.RangeTracker``) implementations. Contains a few lightweight utilities (e.g. reading items from a source such as ``readFromSource()``, as well as heavyweight property testing and stress testing harnesses that help getting a large amount of test coverage with few code. Most notable ones are: * ``assertSourcesEqualReferenceSource()`` helps testing that the data read by the union of sources produced by ``BoundedSource.split()`` is the same as data read by the original source. * If your source implements dynamic work rebalancing, use the ``assertSplitAtFraction()`` family of functions - they test behavior of ``RangeTracker.try_split()``, in particular, that various consistency properties are respected and the total set of data read by the source is preserved when splits happen. Use ``assertSplitAtFractionBehavior()`` to test individual cases of ``RangeTracker.try_split()`` and use ``assertSplitAtFractionExhaustive()`` as a heavy-weight stress test including concurrency. We strongly recommend to use both. For example usages, see the unit tests of modules such as * apache_beam.io.source_test_utils_test.py * apache_beam.io.avroio_test.py """ # pytype: skip-file import logging import threading import weakref from collections import namedtuple from multiprocessing.pool import ThreadPool from apache_beam.io import iobase from apache_beam.testing.util import equal_to __all__ = [ 'read_from_source', 'assert_sources_equal_reference_source', 'assert_reentrant_reads_succeed', 'assert_split_at_fraction_behavior', 'assert_split_at_fraction_binary', 'assert_split_at_fraction_exhaustive', 'assert_split_at_fraction_fails', 'assert_split_at_fraction_succeeds_and_consistent' ] _LOGGER = logging.getLogger(__name__) class ExpectedSplitOutcome(object): MUST_SUCCEED_AND_BE_CONSISTENT = 1 MUST_FAIL = 2 MUST_BE_CONSISTENT_IF_SUCCEEDS = 3 SplitAtFractionResult = namedtuple( 'SplitAtFractionResult', 'num_primary_items num_residual_items') SplitFractionStatistics = namedtuple( 'SplitFractionStatistics', 'successful_fractions non_trivial_fractions') def read_from_source(source, start_position=None, stop_position=None): """Reads elements from the given ```BoundedSource```. Only reads elements within the given position range. Args: source (~apache_beam.io.iobase.BoundedSource): :class:`~apache_beam.io.iobase.BoundedSource` implementation. start_position (int): start position for reading. stop_position (int): stop position for reading. Returns: List[str]: the set of values read from the sources. """ values = [] range_tracker = source.get_range_tracker(start_position, stop_position) assert isinstance(range_tracker, iobase.RangeTracker) reader = source.read(range_tracker) for value in reader: values.append(value) return values def _ThreadPool(threads): # ThreadPool crashes in old versions of Python (< 2.7.5) if created from a # child thread. (http://bugs.python.org/issue10015) if not hasattr(threading.current_thread(), '_children'): threading.current_thread()._children = weakref.WeakKeyDictionary() return ThreadPool(threads) def assert_sources_equal_reference_source(reference_source_info, sources_info): """Tests if a reference source is equal to a given set of sources. Given a reference source (a :class:`~apache_beam.io.iobase.BoundedSource` and a position range) and a list of sources, assert that the union of the records read from the list of sources is equal to the records read from the reference source. Args: reference_source_info\ (Tuple[~apache_beam.io.iobase.BoundedSource, int, int]): a three-tuple that gives the reference :class:`~apache_beam.io.iobase.BoundedSource`, position to start reading at, and position to stop reading at. sources_info\ (Iterable[Tuple[~apache_beam.io.iobase.BoundedSource, int, int]]): a set of sources. Each source is a three-tuple that is of the same format described above. Raises: ValueError: if the set of data produced by the reference source and the given set of sources are not equivalent. """ if not (isinstance(reference_source_info, tuple) and len(reference_source_info) == 3 and isinstance(reference_source_info[0], iobase.BoundedSource)): raise ValueError( 'reference_source_info must a three-tuple where first' 'item of the tuple gives a ' 'iobase.BoundedSource. Received: %r' % reference_source_info) reference_records = read_from_source(*reference_source_info) source_records = [] for source_info in sources_info: assert isinstance(source_info, tuple) assert len(source_info) == 3 if not (isinstance(source_info, tuple) and len(source_info) == 3 and isinstance(source_info[0], iobase.BoundedSource)): raise ValueError( 'source_info must a three tuple where first' 'item of the tuple gives a ' 'iobase.BoundedSource. Received: %r' % source_info) if (type(reference_source_info[0].default_output_coder()) != type(source_info[0].default_output_coder())): raise ValueError( 'Reference source %r and the source %r must use the same coder. ' 'They are using %r and %r respectively instead.' % ( reference_source_info[0], source_info[0], type(reference_source_info[0].default_output_coder()), type(source_info[0].default_output_coder()))) source_records.extend(read_from_source(*source_info)) if len(reference_records) != len(source_records): raise ValueError( 'Reference source must produce the same number of records as the ' 'list of sources. Number of records were %d and %d instead.' % (len(reference_records), len(source_records))) if equal_to(reference_records)(source_records): raise ValueError( 'Reference source and provided list of sources must produce the ' 'same set of records.') def assert_reentrant_reads_succeed(source_info): """Tests if a given source can be read in a reentrant manner. Assume that given source produces the set of values ``{v1, v2, v3, ... vn}``. For ``i`` in range ``[1, n-1]`` this method performs a reentrant read after reading ``i`` elements and verifies that both the original and reentrant read produce the expected set of values. Args: source_info (Tuple[~apache_beam.io.iobase.BoundedSource, int, int]): a three-tuple that gives the reference :class:`~apache_beam.io.iobase.BoundedSource`, position to start reading at, and a position to stop reading at. Raises: ValueError: if source is too trivial or reentrant read result in an incorrect read. """ source, start_position, stop_position = source_info assert isinstance(source, iobase.BoundedSource) expected_values = [ val for val in source.read( source.get_range_tracker(start_position, stop_position)) ] if len(expected_values) < 2: raise ValueError( 'Source is too trivial since it produces only %d ' 'values. Please give a source that reads at least 2 ' 'values.' % len(expected_values)) for i in range(1, len(expected_values) - 1): read_iter = source.read( source.get_range_tracker(start_position, stop_position)) original_read = [] for _ in range(i): original_read.append(next(read_iter)) # Reentrant read reentrant_read = [ val for val in source.read( source.get_range_tracker(start_position, stop_position)) ] # Continuing original read. for val in read_iter: original_read.append(val) if equal_to(original_read)(expected_values): raise ValueError( 'Source did not produce expected values when ' 'performing a reentrant read after reading %d values. ' 'Expected %r received %r.' % (i, expected_values, original_read)) if equal_to(reentrant_read)(expected_values): raise ValueError( 'A reentrant read of source after reading %d values ' 'did not produce expected values. Expected %r ' 'received %r.' % (i, expected_values, reentrant_read)) def assert_split_at_fraction_behavior( source, num_items_to_read_before_split, split_fraction, expected_outcome): """Verifies the behaviour of splitting a source at a given fraction. Asserts that splitting a :class:`~apache_beam.io.iobase.BoundedSource` either fails after reading **num_items_to_read_before_split** items, or succeeds in a way that is consistent according to :func:`assert_split_at_fraction_succeeds_and_consistent()`. Args: source (~apache_beam.io.iobase.BoundedSource): the source to perform dynamic splitting on. num_items_to_read_before_split (int): number of items to read before splitting. split_fraction (float): fraction to split at. expected_outcome (int): a value from :class:`~apache_beam.io.source_test_utils.ExpectedSplitOutcome`. Returns: Tuple[int, int]: a tuple that gives the number of items produced by reading the two ranges produced after dynamic splitting. If splitting did not occur, the first value of the tuple will represent the full set of records read by the source while the second value of the tuple will be ``-1``. """ assert isinstance(source, iobase.BoundedSource) expected_items = read_from_source(source, None, None) return _assert_split_at_fraction_behavior( source, expected_items, num_items_to_read_before_split, split_fraction, expected_outcome) def _assert_split_at_fraction_behavior( source, expected_items, num_items_to_read_before_split, split_fraction, expected_outcome, start_position=None, stop_position=None): range_tracker = source.get_range_tracker(start_position, stop_position) assert isinstance(range_tracker, iobase.RangeTracker) current_items = [] reader = source.read(range_tracker) # Reading 'num_items_to_read_before_split' items. reader_iter = iter(reader) for _ in range(num_items_to_read_before_split): current_items.append(next(reader_iter)) suggested_split_position = range_tracker.position_at_fraction(split_fraction) stop_position_before_split = range_tracker.stop_position() split_result = range_tracker.try_split(suggested_split_position) if split_result is not None: if len(split_result) != 2: raise ValueError( 'Split result must be a tuple that contains split ' 'position and split fraction. Received: %r' % (split_result, )) if range_tracker.stop_position() != split_result[0]: raise ValueError( 'After a successful split, the stop position of the ' 'RangeTracker must be the same as the returned split ' 'position. Observed %r and %r which are different.' % (range_tracker.stop_position() % (split_result[0], ))) if split_fraction < 0 or split_fraction > 1: raise ValueError( 'Split fraction must be within the range [0,1]', 'Observed split fraction was %r.' % (split_result[1], )) stop_position_after_split = range_tracker.stop_position() if split_result and stop_position_after_split == stop_position_before_split: raise ValueError( 'Stop position %r did not change after a successful ' 'split of source %r at fraction %r.' % (stop_position_before_split, source, split_fraction)) if expected_outcome == ExpectedSplitOutcome.MUST_SUCCEED_AND_BE_CONSISTENT: if not split_result: raise ValueError( 'Expected split of source %r at fraction %r to be ' 'successful after reading %d elements. But ' 'the split failed.' % (source, split_fraction, num_items_to_read_before_split)) elif expected_outcome == ExpectedSplitOutcome.MUST_FAIL: if split_result: raise ValueError( 'Expected split of source %r at fraction %r after ' 'reading %d elements to fail. But splitting ' 'succeeded with result %r.' % ( source, split_fraction, num_items_to_read_before_split, split_result)) elif (expected_outcome != ExpectedSplitOutcome.MUST_BE_CONSISTENT_IF_SUCCEEDS): raise ValueError('Unknown type of expected outcome: %r' % expected_outcome) current_items.extend([value for value in reader_iter]) residual_range = ( split_result[0], stop_position_before_split) if split_result else None return _verify_single_split_fraction_result( source, expected_items, current_items, split_result, (range_tracker.start_position(), range_tracker.stop_position()), residual_range, split_fraction) def _range_to_str(start, stop): return '[' + (str(start) + ',' + str(stop) + ')') def _verify_single_split_fraction_result( source, expected_items, current_items, split_successful, primary_range, residual_range, split_fraction): assert primary_range primary_items = read_from_source(source, *primary_range) if not split_successful: # For unsuccessful splits, residual_range should be None. assert not residual_range residual_items = ( read_from_source(source, *residual_range) if split_successful else []) total_items = primary_items + residual_items if current_items != primary_items: raise ValueError( 'Current source %r and a source created using the ' 'range of the primary source %r determined ' 'by performing dynamic work rebalancing at fraction ' '%r produced different values. Expected ' 'these sources to produce the same list of values.' % (source, _range_to_str(*primary_range), split_fraction)) if expected_items != total_items: raise ValueError( 'Items obtained by reading the source %r for primary ' 'and residual ranges %s and %s did not produce the ' 'expected list of values.' % (source, _range_to_str(*primary_range), _range_to_str(*residual_range))) result = (len(primary_items), len(residual_items) if split_successful else -1) return result def assert_split_at_fraction_succeeds_and_consistent( source, num_items_to_read_before_split, split_fraction): """Verifies some consistency properties of dynamic work rebalancing. Equivalent to the following pseudocode::: original_range_tracker = source.getRangeTracker(None, None) original_reader = source.read(original_range_tracker) items_before_split = read N items from original_reader suggested_split_position = original_range_tracker.position_for_fraction( split_fraction) original_stop_position - original_range_tracker.stop_position() split_result = range_tracker.try_split() split_position, split_fraction = split_result primary_range_tracker = source.get_range_tracker( original_range_tracker.start_position(), split_position) residual_range_tracker = source.get_range_tracker(split_position, original_stop_position) assert that: items when reading source.read(primary_range_tracker) == items_before_split + items from continuing to read 'original_reader' assert that: items when reading source.read(original_range_tracker) = items when reading source.read(primary_range_tracker) + items when reading source.read(residual_range_tracker) Args: source: source to perform dynamic work rebalancing on. num_items_to_read_before_split: number of items to read before splitting. split_fraction: fraction to split at. """ assert_split_at_fraction_behavior( source, num_items_to_read_before_split, split_fraction, ExpectedSplitOutcome.MUST_SUCCEED_AND_BE_CONSISTENT) def assert_split_at_fraction_fails( source, num_items_to_read_before_split, split_fraction): """Asserts that dynamic work rebalancing at a given fraction fails. Asserts that trying to perform dynamic splitting after reading 'num_items_to_read_before_split' items from the source fails. Args: source: source to perform dynamic splitting on. num_items_to_read_before_split: number of items to read before splitting. split_fraction: fraction to split at. """ assert_split_at_fraction_behavior( source, num_items_to_read_before_split, split_fraction, ExpectedSplitOutcome.MUST_FAIL) def assert_split_at_fraction_binary( source, expected_items, num_items_to_read_before_split, left_fraction, left_result, right_fraction, right_result, stats, start_position=None, stop_position=None): """Performs dynamic work rebalancing for fractions within a given range. Asserts that given a start position, a source can be split at every interesting fraction (halfway between two fractions that differ by at least one item) and the results are consistent if a split succeeds. Args: source: source to perform dynamic splitting on. expected_items: total set of items expected when reading the source. num_items_to_read_before_split: number of items to read before splitting. left_fraction: left fraction for binary splitting. left_result: result received by splitting at left fraction. right_fraction: right fraction for binary splitting. right_result: result received by splitting at right fraction. stats: a ``SplitFractionStatistics`` for storing results. """ assert right_fraction > left_fraction if right_fraction - left_fraction < 0.001: # This prevents infinite recursion. return middle_fraction = (left_fraction + right_fraction) / 2 if left_result is None: left_result = _assert_split_at_fraction_behavior( source, expected_items, num_items_to_read_before_split, left_fraction, ExpectedSplitOutcome.MUST_BE_CONSISTENT_IF_SUCCEEDS) if right_result is None: right_result = _assert_split_at_fraction_behavior( source, expected_items, num_items_to_read_before_split, right_fraction, ExpectedSplitOutcome.MUST_BE_CONSISTENT_IF_SUCCEEDS) middle_result = _assert_split_at_fraction_behavior( source, expected_items, num_items_to_read_before_split, middle_fraction, ExpectedSplitOutcome.MUST_BE_CONSISTENT_IF_SUCCEEDS) if middle_result[1] != -1: stats.successful_fractions.append(middle_fraction) if middle_result[1] > 0: stats.non_trivial_fractions.append(middle_fraction) # Two split results are equivalent if primary and residual ranges of them # produce the same number of records (simply checking the size of primary # enough since the total number of records is constant). if left_result[0] != middle_result[0]: assert_split_at_fraction_binary( source, expected_items, num_items_to_read_before_split, left_fraction, left_result, middle_fraction, middle_result, stats) # We special case right_fraction=1.0 since that could fail due to being out # of range. (even if a dynamic split fails at 'middle_fraction' and at # fraction 1.0, there might be fractions in range ('middle_fraction', 1.0) # where dynamic splitting succeeds). if right_fraction == 1.0 or middle_result[0] != right_result[0]: assert_split_at_fraction_binary( source, expected_items, num_items_to_read_before_split, middle_fraction, middle_result, right_fraction, right_result, stats) MAX_CONCURRENT_SPLITTING_TRIALS_PER_ITEM = 100 MAX_CONCURRENT_SPLITTING_TRIALS_TOTAL = 1000 def assert_split_at_fraction_exhaustive( source, start_position=None, stop_position=None, perform_multi_threaded_test=True): """Performs and tests dynamic work rebalancing exhaustively. Asserts that for each possible start position, a source can be split at every interesting fraction (halfway between two fractions that differ by at least one item) and the results are consistent if a split succeeds. Verifies multi threaded splitting as well. Args: source (~apache_beam.io.iobase.BoundedSource): the source to perform dynamic splitting on. perform_multi_threaded_test (bool): if :data:`True` performs a multi-threaded test, otherwise this test is skipped. Raises: ValueError: if the exhaustive splitting test fails. """ expected_items = read_from_source(source, start_position, stop_position) if not expected_items: raise ValueError('Source %r is empty.' % source) if len(expected_items) == 1: raise ValueError('Source %r only reads a single item.' % source) all_non_trivial_fractions = [] any_successful_fractions = False any_non_trivial_fractions = False for i in range(len(expected_items)): stats = SplitFractionStatistics([], []) assert_split_at_fraction_binary( source, expected_items, i, 0.0, None, 1.0, None, stats) if stats.successful_fractions: any_successful_fractions = True if stats.non_trivial_fractions: any_non_trivial_fractions = True all_non_trivial_fractions.append(stats.non_trivial_fractions) if not any_successful_fractions: raise ValueError( 'SplitAtFraction test completed vacuously: no ' 'successful split fractions found') if not any_non_trivial_fractions: raise ValueError( 'SplitAtFraction test completed vacuously: no non-trivial split ' 'fractions found') if not perform_multi_threaded_test: return num_total_trials = 0 for i in range(len(expected_items)): non_trivial_fractions = [2.0] # 2.0 is larger than any valid fraction. non_trivial_fractions.extend(all_non_trivial_fractions[i]) min_non_trivial_fraction = min(non_trivial_fractions) if min_non_trivial_fraction == 2.0: # This will not happen all the time. Otherwise previous test will fail # due to vacuousness. continue num_trials = 0 have_success = False have_failure = False thread_pool = _ThreadPool(2) try: while True: num_trials += 1 if num_trials > MAX_CONCURRENT_SPLITTING_TRIALS_PER_ITEM: _LOGGER.warning( 'After %d concurrent splitting trials at item #%d, observed ' 'only %s, giving up on this item', num_trials, i, 'success' if have_success else 'failure') break if _assert_split_at_fraction_concurrent(source, expected_items, i, min_non_trivial_fraction, thread_pool): have_success = True else: have_failure = True if have_success and have_failure: _LOGGER.info( '%d trials to observe both success and failure of ' 'concurrent splitting at item #%d', num_trials, i) break finally: thread_pool.close() num_total_trials += num_trials if num_total_trials > MAX_CONCURRENT_SPLITTING_TRIALS_TOTAL: _LOGGER.warning( 'After %d total concurrent splitting trials, considered ' 'only %d items, giving up.', num_total_trials, i) break _LOGGER.info( '%d total concurrent splitting trials for %d items', num_total_trials, len(expected_items)) def _assert_split_at_fraction_concurrent( source, expected_items, num_items_to_read_before_splitting, split_fraction, thread_pool=None): range_tracker = source.get_range_tracker(None, None) stop_position_before_split = range_tracker.stop_position() reader = source.read(range_tracker) reader_iter = iter(reader) current_items = [] for _ in range(num_items_to_read_before_splitting): current_items.append(next(reader_iter)) def read_or_split(test_params): if test_params[0]: return [val for val in test_params[1]] else: position = test_params[1].position_at_fraction(test_params[2]) result = test_params[1].try_split(position) return result inputs = [] pool = thread_pool if thread_pool else _ThreadPool(2) try: inputs.append([True, reader_iter]) inputs.append([False, range_tracker, split_fraction]) results = pool.map(read_or_split, inputs) finally: if not thread_pool: pool.close() current_items.extend(results[0]) primary_range = ( range_tracker.start_position(), range_tracker.stop_position()) split_result = results[1] residual_range = ( split_result[0], stop_position_before_split) if split_result else None res = _verify_single_split_fraction_result( source, expected_items, current_items, split_result, primary_range, residual_range, split_fraction) return res[1] > 0 ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/source_test_utils_test.py0000644000175100017510000001147515206113032023464 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import tempfile import unittest from apache_beam.io import source_test_utils from apache_beam.io.filebasedsource_test import LineSource class SourceTestUtilsTest(unittest.TestCase): def _create_file_with_data(self, lines): assert isinstance(lines, list) with tempfile.NamedTemporaryFile(delete=False) as f: for line in lines: f.write(line + b'\n') return f.name def _create_data(self, num_lines): return [b'line ' + str(i).encode('latin1') for i in range(num_lines)] def _create_source(self, data): source = LineSource(self._create_file_with_data(data)) # By performing initial splitting, we can get a source for a single file. # This source, that uses OffsetRangeTracker, is better for testing purposes, # than using the original source for a file-pattern. for bundle in source.split(float('inf')): return bundle.source def test_read_from_source(self): data = self._create_data(100) source = self._create_source(data) self.assertCountEqual( data, source_test_utils.read_from_source(source, None, None)) def test_source_equals_reference_source(self): data = self._create_data(100) reference_source = self._create_source(data) sources_info = [(split.source, split.start_position, split.stop_position) for split in reference_source.split(desired_bundle_size=50)] if len(sources_info) < 2: raise ValueError( 'Test is too trivial since splitting only generated %d' 'bundles. Please adjust the test so that at least ' 'two splits get generated.' % len(sources_info)) source_test_utils.assert_sources_equal_reference_source( (reference_source, None, None), sources_info) def test_split_at_fraction_successful(self): data = self._create_data(100) source = self._create_source(data) result1 = source_test_utils.assert_split_at_fraction_behavior( source, 10, 0.5, source_test_utils.ExpectedSplitOutcome.MUST_SUCCEED_AND_BE_CONSISTENT) result2 = source_test_utils.assert_split_at_fraction_behavior( source, 20, 0.5, source_test_utils.ExpectedSplitOutcome.MUST_SUCCEED_AND_BE_CONSISTENT) self.assertEqual(result1, result2) self.assertEqual(100, result1[0] + result1[1]) result3 = source_test_utils.assert_split_at_fraction_behavior( source, 30, 0.8, source_test_utils.ExpectedSplitOutcome.MUST_SUCCEED_AND_BE_CONSISTENT) result4 = source_test_utils.assert_split_at_fraction_behavior( source, 50, 0.8, source_test_utils.ExpectedSplitOutcome.MUST_SUCCEED_AND_BE_CONSISTENT) self.assertEqual(result3, result4) self.assertEqual(100, result3[0] + result4[1]) self.assertTrue(result1[0] < result3[0]) self.assertTrue(result1[1] > result3[1]) def test_split_at_fraction_fails(self): data = self._create_data(100) source = self._create_source(data) result = source_test_utils.assert_split_at_fraction_behavior( source, 90, 0.1, source_test_utils.ExpectedSplitOutcome.MUST_FAIL) self.assertEqual(result[0], 100) self.assertEqual(result[1], -1) with self.assertRaises(ValueError): source_test_utils.assert_split_at_fraction_behavior( source, 10, 0.5, source_test_utils.ExpectedSplitOutcome.MUST_FAIL) def test_split_at_fraction_binary(self): data = self._create_data(100) source = self._create_source(data) stats = source_test_utils.SplitFractionStatistics([], []) source_test_utils.assert_split_at_fraction_binary( source, data, 10, 0.5, None, 0.8, None, stats) # These lists should not be empty now. self.assertTrue(stats.successful_fractions) self.assertTrue(stats.non_trivial_fractions) def test_split_at_fraction_exhaustive(self): data = self._create_data(10) source = self._create_source(data) source_test_utils.assert_split_at_fraction_exhaustive(source) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/sources_test.py0000644000175100017510000000761315206113032021367 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the sources framework.""" # pytype: skip-file import logging import os import tempfile import unittest import apache_beam as beam from apache_beam import coders from apache_beam.io import iobase from apache_beam.io import range_trackers from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class LineSource(iobase.BoundedSource): """A simple source that reads lines from a given file.""" TEST_BUNDLE_SIZE = 10 def __init__(self, file_name): self._file_name = file_name def read(self, range_tracker): with open(self._file_name, 'rb') as f: start = range_tracker.start_position() f.seek(start) if start > 0: f.seek(-1, os.SEEK_CUR) start -= 1 start += len(f.readline()) current = start line = f.readline() while range_tracker.try_claim(current): if not line: return yield line.rstrip(b'\n') current += len(line) line = f.readline() def split(self, desired_bundle_size, start_position=None, stop_position=None): assert start_position is None assert stop_position is None size = self.estimate_size() bundle_start = 0 while bundle_start < size: bundle_stop = min(bundle_start + LineSource.TEST_BUNDLE_SIZE, size) yield iobase.SourceBundle( bundle_stop - bundle_start, self, bundle_start, bundle_stop) bundle_start = bundle_stop def get_range_tracker(self, start_position, stop_position): if start_position is None: start_position = 0 if stop_position is None: stop_position = self._get_file_size() return range_trackers.OffsetRangeTracker(start_position, stop_position) def default_output_coder(self): return coders.BytesCoder() def estimate_size(self): return self._get_file_size() def _get_file_size(self): with open(self._file_name, 'rb') as f: f.seek(0, os.SEEK_END) return f.tell() class SourcesTest(unittest.TestCase): def _create_temp_file(self, contents): with tempfile.NamedTemporaryFile(delete=False) as f: f.write(contents) return f.name def test_read_from_source(self): file_name = self._create_temp_file(b'aaaa\nbbbb\ncccc\ndddd') source = LineSource(file_name) range_tracker = source.get_range_tracker(None, None) result = [line for line in source.read(range_tracker)] self.assertCountEqual([b'aaaa', b'bbbb', b'cccc', b'dddd'], result) self.assertTrue( range_tracker.last_attempted_record_start >= range_tracker.stop_position()) def test_source_estimated_size(self): file_name = self._create_temp_file(b'aaaa\n') source = LineSource(file_name) self.assertEqual(5, source.estimate_size()) def test_run_direct(self): file_name = self._create_temp_file(b'aaaa\nbbbb\ncccc\ndddd') with TestPipeline() as pipeline: pcoll = pipeline | beam.io.Read(LineSource(file_name)) assert_that(pcoll, equal_to([b'aaaa', b'bbbb', b'cccc', b'dddd'])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/textio.py0000644000175100017510000012766715206113032020175 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A source and a sink for reading from and writing to text files.""" # pytype: skip-file import logging import os from functools import partial from typing import TYPE_CHECKING from typing import Any from typing import Optional from typing import Union from apache_beam import typehints from apache_beam.coders import coders from apache_beam.io import filebasedsink from apache_beam.io import filebasedsource from apache_beam.io import iobase from apache_beam.io.filebasedsource import ReadAllFiles from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.iobase import Read from apache_beam.io.iobase import Write from apache_beam.transforms import PTransform from apache_beam.transforms.display import DisplayDataItem if TYPE_CHECKING: from apache_beam.io import fileio __all__ = [ 'ReadFromText', 'ReadFromTextWithFilename', 'ReadAllFromText', 'ReadAllFromTextContinuously', 'WriteToText', 'ReadFromCsv', 'WriteToCsv', 'ReadFromJson', 'WriteToJson', ] _LOGGER = logging.getLogger(__name__) class _TextSource(filebasedsource.FileBasedSource): r"""A source for reading text files. Parses a text file as newline-delimited elements. Supports newline delimiters '\n' and '\r\n. This implementation reads encoded text and uses the input coder's encoding to decode from bytes to str. This does not support ``UTF-16`` or ``UTF-32`` encodings. """ DEFAULT_READ_BUFFER_SIZE = 8192 class ReadBuffer(object): # A buffer that gives the buffered data and next position in the # buffer that should be read. def __init__(self, data, position): self._data = data self._position = position @property def data(self): return self._data @data.setter def data(self, value): assert isinstance(value, bytes) self._data = value @property def position(self): return self._position @position.setter def position(self, value): assert isinstance(value, int) if value > len(self._data): raise ValueError( 'Cannot set position to %d since it\'s larger than ' 'size of data %d.' % (value, len(self._data))) self._position = value def reset(self): self.data = b'' self.position = 0 def __init__( self, file_pattern, min_bundle_size, compression_type, strip_trailing_newlines, coder: coders.Coder, buffer_size=DEFAULT_READ_BUFFER_SIZE, validate=True, skip_header_lines=0, header_processor_fns=(None, None), delimiter=None, escapechar=None): """Initialize a _TextSource Args: header_processor_fns (tuple): a tuple of a `header_matcher` function and a `header_processor` function. The `header_matcher` should return `True` for all lines at the start of the file that are part of the file header and `False` otherwise. These header lines will not be yielded when reading records and instead passed into `header_processor` to be handled. If `skip_header_lines` and a `header_matcher` are both provided, the value of `skip_header_lines` lines will be skipped and the header will be processed from there. delimiter (bytes) Optional: delimiter to split records. Must not self-overlap, because self-overlapping delimiters cause ambiguous parsing. escapechar (bytes) Optional: a single byte to escape the records delimiter, can also escape itself. Raises: ValueError: if skip_lines is negative. Please refer to documentation in class `ReadFromText` for the rest of the arguments. """ super().__init__( file_pattern, min_bundle_size, compression_type=compression_type, validate=validate) self._strip_trailing_newlines = strip_trailing_newlines self._compression_type = compression_type self._coder = coder self._buffer_size = buffer_size if skip_header_lines < 0: raise ValueError( 'Cannot skip negative number of header lines: %d' % skip_header_lines) elif skip_header_lines > 10: _LOGGER.warning( 'Skipping %d header lines. Skipping large number of header ' 'lines might significantly slow down processing.') self._skip_header_lines = skip_header_lines self._header_matcher, self._header_processor = header_processor_fns if delimiter is not None: if not isinstance(delimiter, bytes) or len(delimiter) == 0: raise ValueError('Delimiter must be a non-empty bytes sequence.') if self._is_self_overlapping(delimiter): raise ValueError('Delimiter must not self-overlap.') self._delimiter = delimiter if escapechar is not None: if not (isinstance(escapechar, bytes) and len(escapechar) == 1): raise ValueError( "escapechar must be bytes of size 1: '%s'" % escapechar) self._escapechar = escapechar def display_data(self): parent_dd = super().display_data() parent_dd['strip_newline'] = DisplayDataItem( self._strip_trailing_newlines, label='Strip Trailing New Lines') parent_dd['buffer_size'] = DisplayDataItem( self._buffer_size, label='Buffer Size') parent_dd['coder'] = DisplayDataItem(self._coder.__class__, label='Coder') return parent_dd def read_records(self, file_name, range_tracker): start_offset = range_tracker.start_position() read_buffer = _TextSource.ReadBuffer(b'', 0) next_record_start_position = -1 def split_points_unclaimed(stop_position): return ( 0 if stop_position <= next_record_start_position else iobase.RangeTracker.SPLIT_POINTS_UNKNOWN) range_tracker.set_split_points_unclaimed_callback(split_points_unclaimed) with self.open_file(file_name) as file_to_read: position_after_processing_header_lines = ( self._process_header(file_to_read, read_buffer)) start_offset = max(start_offset, position_after_processing_header_lines) if start_offset > position_after_processing_header_lines: # Seeking to one delimiter length before the start index and ignoring # the current line. If start_position is at beginning of the line, that # line belongs to the current bundle, hence ignoring that is incorrect. # Seeking to one delimiter before prevents that. if self._delimiter is not None and start_offset >= len(self._delimiter): required_position = start_offset - len(self._delimiter) else: required_position = start_offset - 1 if self._escapechar is not None: # Need more bytes to check if the delimiter is escaped. # Seek until the first escapechar if any. while required_position > 0: file_to_read.seek(required_position - 1) if file_to_read.read(1) == self._escapechar: required_position -= 1 else: break file_to_read.seek(required_position) read_buffer.reset() sep_bounds = self._find_separator_bounds(file_to_read, read_buffer) if not sep_bounds: # Could not find a delimiter after required_position. This means that # none of the records within the file belongs to the current source. return _, sep_end = sep_bounds read_buffer.data = read_buffer.data[sep_end:] next_record_start_position = required_position + sep_end else: next_record_start_position = position_after_processing_header_lines while range_tracker.try_claim(next_record_start_position): record, num_bytes_to_next_record = self._read_record(file_to_read, read_buffer) # For compressed text files that use an unsplittable OffsetRangeTracker # with infinity as the end position, above 'try_claim()' invocation # would pass for an empty record at the end of file that is not # followed by a new line character. Since such a record is at the last # position of a file, it should not be a part of the considered range. # We do this check to ignore such records. if len(record) == 0 and num_bytes_to_next_record < 0: # pylint: disable=len-as-condition break # Record delimiter must be larger than zero bytes. assert num_bytes_to_next_record != 0 if num_bytes_to_next_record > 0: next_record_start_position += num_bytes_to_next_record yield self._coder.decode(record) if num_bytes_to_next_record < 0: break def _process_header(self, file_to_read, read_buffer): # Returns a tuple containing the position in file after processing header # records and a list of decoded header lines that match # 'header_matcher'. header_lines = [] position = self._skip_lines( file_to_read, read_buffer, self._skip_header_lines) if self._skip_header_lines else 0 if self._header_matcher: while True: record, num_bytes_to_next_record = self._read_record(file_to_read, read_buffer) decoded_line = self._coder.decode(record) if not self._header_matcher(decoded_line): # We've read past the header section at this point, so go back a line. file_to_read.seek(position) read_buffer.reset() break header_lines.append(decoded_line) if num_bytes_to_next_record < 0: break position += num_bytes_to_next_record if self._header_processor: self._header_processor(header_lines) return position def _find_separator_bounds(self, file_to_read, read_buffer): # Determines the start and end positions within 'read_buffer.data' of the # next delimiter starting from position 'read_buffer.position'. # Use the custom delimiter to be used in place of # the default ones ('\n' or '\r\n')' # This method may increase the size of buffer but it will not decrease the # size of it. current_pos = read_buffer.position # b'\n' use as default delimiter = self._delimiter or b'\n' delimiter_len = len(delimiter) while True: if current_pos >= len(read_buffer.data) - delimiter_len + 1: # Ensuring that there are enough bytes to determine # at current_pos. if not self._try_to_ensure_num_bytes_in_buffer( file_to_read, read_buffer, current_pos + delimiter_len): return # Using find() here is more efficient than a linear scan # of the byte array. next_delim = read_buffer.data.find(delimiter, current_pos) if next_delim >= 0: if (self._delimiter is None and read_buffer.data[next_delim - 1:next_delim] == b'\r'): if self._escapechar is not None and self._is_escaped(read_buffer, next_delim - 1): # Accept '\n' as a default delimiter, because '\r' is escaped. return (next_delim, next_delim + 1) else: # Accept both '\r\n' and '\n' as a default delimiter. return (next_delim - 1, next_delim + 1) else: if self._escapechar is not None and self._is_escaped(read_buffer, next_delim): # Skip an escaped delimiter. current_pos = next_delim + delimiter_len + 1 continue else: # Found a delimiter. Accepting that as the next delimiter. return (next_delim, next_delim + delimiter_len) elif self._delimiter is not None: # Corner case: custom delimiter is truncated at the end of the buffer. next_delim = read_buffer.data.find( delimiter[0], len(read_buffer.data) - delimiter_len + 1) if next_delim >= 0: # Delimiters longer than 1 byte may cross the buffer boundary. # Defer full matching till the next iteration. current_pos = next_delim continue current_pos = len(read_buffer.data) def _try_to_ensure_num_bytes_in_buffer( self, file_to_read, read_buffer, num_bytes): # Tries to ensure that there are at least num_bytes bytes in the buffer. # Returns True if this can be fulfilled, returned False if this cannot be # fulfilled due to reaching EOF. while len(read_buffer.data) < num_bytes: read_data = file_to_read.read(self._buffer_size) if not read_data: return False read_buffer.data += read_data return True def _skip_lines(self, file_to_read, read_buffer, num_lines): """Skip num_lines from file_to_read, return num_lines+1 start position.""" if file_to_read.tell() > 0: file_to_read.seek(0) position = 0 for _ in range(num_lines): _, num_bytes_to_next_record = self._read_record(file_to_read, read_buffer) if num_bytes_to_next_record < 0: # We reached end of file. It is OK to just break here # because subsequent _read_record will return same result. break position += num_bytes_to_next_record return position def _read_record(self, file_to_read, read_buffer): # Returns a tuple containing the current_record and number of bytes to the # next record starting from 'read_buffer.position'. If EOF is # reached, returns a tuple containing the current record and -1. if read_buffer.position > self._buffer_size: # read_buffer is too large. Truncating and adjusting it. read_buffer.data = read_buffer.data[read_buffer.position:] read_buffer.position = 0 record_start_position_in_buffer = read_buffer.position sep_bounds = self._find_separator_bounds(file_to_read, read_buffer) read_buffer.position = sep_bounds[1] if sep_bounds else len( read_buffer.data) if not sep_bounds: # Reached EOF. Bytes up to the EOF is the next record. Returning '-1' for # the starting position of the next record. return (read_buffer.data[record_start_position_in_buffer:], -1) if self._strip_trailing_newlines: # Current record should not contain the delimiter. return ( read_buffer.data[record_start_position_in_buffer:sep_bounds[0]], sep_bounds[1] - record_start_position_in_buffer) else: # Current record should contain the delimiter. return ( read_buffer.data[record_start_position_in_buffer:sep_bounds[1]], sep_bounds[1] - record_start_position_in_buffer) @staticmethod def _is_self_overlapping(delimiter): # A delimiter self-overlaps if it has a prefix that is also its suffix. for i in range(1, len(delimiter)): if delimiter[0:i] == delimiter[len(delimiter) - i:]: return True return False def _is_escaped(self, read_buffer, position): # Returns True if byte at position is preceded with an odd number # of escapechar bytes or False if preceded by 0 or even escapes # (the even number means that all the escapes are escaped themselves). escape_count = 0 for current_pos in reversed(range(0, position)): if read_buffer.data[current_pos:current_pos + 1] != self._escapechar: break escape_count += 1 return escape_count % 2 == 1 def output_type_hint(self): try: return self._coder.to_type_hint() except NotImplementedError: return Any class _TextSourceWithFilename(_TextSource): def read_records(self, file_name, range_tracker): records = super().read_records(file_name, range_tracker) for record in records: yield (file_name, record) def output_type_hint(self): return typehints.KV[str, super().output_type_hint()] class _TextSink(filebasedsink.FileBasedSink): """A sink to a GCS or local text file or files.""" def __init__( self, file_path_prefix, file_name_suffix='', append_trailing_newlines=True, num_shards=0, shard_name_template=None, coder: coders.Coder = coders.ToBytesCoder(), compression_type=CompressionTypes.AUTO, header=None, footer=None, *, max_records_per_shard=None, max_bytes_per_shard=None, skip_if_empty=False, triggering_frequency=None): """Initialize a _TextSink. Args: file_path_prefix: The file path to write to. The files written will begin with this prefix, followed by a shard identifier (see num_shards), and end in a common extension, if given by file_name_suffix. In most cases, only this argument is specified and num_shards, shard_name_template, and file_name_suffix use default values. file_name_suffix: Suffix for the files written. append_trailing_newlines: indicate whether this sink should write an additional newline char after writing each element. num_shards: The number of files (shards) used for output. If not set, the service will decide on the optimal number of shards. Constraining the number of shards is likely to reduce the performance of a pipeline. Setting this value is not recommended unless you require a specific number of output files. In streaming if not set, the service will write a file per bundle. shard_name_template: A template string containing placeholders for the shard number and shard count. Currently only ``''``, ``'-SSSSS-of-NNNNN'``, ``'-W-SSSSS-of-NNNNN'`` and ``'-V-SSSSS-of-NNNNN'`` are patterns accepted by the service. When constructing a filename for a particular shard number, the upper-case letters ``S`` and ``N`` are replaced with the ``0``-padded shard number and shard count respectively. This argument can be ``''`` in which case it behaves as if num_shards was set to 1 and only one file will be generated. The default pattern used is ``'-SSSSS-of-NNNNN'`` for bounded PCollections and for ``'-W-SSSSS-of-NNNNN'`` unbounded PCollections. W is used for windowed shard naming and is replaced with ``[window.start, window.end)`` V is used for windowed shard naming and is replaced with ``[window.start.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S"), window.end.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S")`` coder: Coder used to encode each line. compression_type: Used to handle compressed output files. Typical value is CompressionTypes.AUTO, in which case the final file path's extension (as determined by file_path_prefix, file_name_suffix, num_shards and shard_name_template) will be used to detect the compression. header: String to write at beginning of file as a header. If not None and append_trailing_newlines is set, '\n' will be added. footer: String to write at the end of file as a footer. If not None and append_trailing_newlines is set, '\n' will be added. max_records_per_shard: Maximum number of records to write to any individual shard. max_bytes_per_shard: Target maximum number of bytes to write to any individual shard. This may be exceeded slightly, as a new shard is created once this limit is hit, but the remainder of a given record, a subsequent newline, and a footer may cause the actual shard size to exceed this value. This also tracks the uncompressed, not compressed, size of the shard. skip_if_empty: Don't write any shards if the PCollection is empty. triggering_frequency: (int) Every triggering_frequency duration, a window will be triggered and all bundles in the window will be written. If set it overrides user windowing. Mandatory for GlobalWindow. Returns: A _TextSink object usable for writing. """ super().__init__( file_path_prefix, file_name_suffix=file_name_suffix, num_shards=num_shards, shard_name_template=shard_name_template, coder=coder, mime_type='text/plain', compression_type=compression_type, max_records_per_shard=max_records_per_shard, max_bytes_per_shard=max_bytes_per_shard, skip_if_empty=skip_if_empty, triggering_frequency=triggering_frequency) self._append_trailing_newlines = append_trailing_newlines self._header = header self._footer = footer def open(self, temp_path): file_handle = super().open(temp_path) if self._header is not None: file_handle.write(coders.ToBytesCoder().encode(self._header)) if self._append_trailing_newlines: file_handle.write(b'\n') return file_handle def close(self, file_handle): if self._footer is not None: file_handle.write(coders.ToBytesCoder().encode(self._footer)) if self._append_trailing_newlines: file_handle.write(b'\n') super().close(file_handle) def display_data(self): dd_parent = super().display_data() dd_parent['append_newline'] = DisplayDataItem( self._append_trailing_newlines, label='Append Trailing New Lines') return dd_parent def write_encoded_record(self, file_handle, encoded_value): """Writes a single encoded record.""" file_handle.write(encoded_value) if self._append_trailing_newlines: file_handle.write(b'\n') def _create_text_source( file_pattern=None, min_bundle_size=None, compression_type=None, strip_trailing_newlines=None, coder=None, validate=False, skip_header_lines=None, delimiter=None, escapechar=None): return _TextSource( file_pattern=file_pattern, min_bundle_size=min_bundle_size, compression_type=compression_type, strip_trailing_newlines=strip_trailing_newlines, coder=coder, validate=validate, skip_header_lines=skip_header_lines, delimiter=delimiter, escapechar=escapechar) class ReadAllFromText(PTransform): """A ``PTransform`` for reading a ``PCollection`` of text files. Reads a ``PCollection`` of text files or file patterns and produces a ``PCollection`` of strings. Parses a text file as newline-delimited elements, by default assuming UTF-8 encoding. Supports newline delimiters '\\n' and '\\r\\n'. If `with_filename` is ``True`` the output will include the file name. This is similar to ``ReadFromTextWithFilename`` but this ``PTransform`` can be placed anywhere in the pipeline. If reading from a text file that that requires a different encoding, you may provide a custom :class:`~apache_beam.coders.coders.Coder` that encodes and decodes with the appropriate codec. For example, see the implementation of :class:`~apache_beam.coders.coders.StrUtf8Coder`. This does not support ``UTF-16`` or ``UTF-32`` encodings. This implementation is only tested with batch pipeline. In streaming, reading may happen with delay due to the limitation in ReShuffle involved. """ DEFAULT_DESIRED_BUNDLE_SIZE = 64 * 1024 * 1024 # 64MB def __init__( self, min_bundle_size=0, desired_bundle_size=DEFAULT_DESIRED_BUNDLE_SIZE, compression_type=CompressionTypes.AUTO, strip_trailing_newlines=True, validate=False, coder: coders.Coder = coders.StrUtf8Coder(), skip_header_lines=0, with_filename=False, delimiter=None, escapechar=None, **kwargs): """Initialize the ``ReadAllFromText`` transform. Args: min_bundle_size: Minimum size of bundles that should be generated when splitting this source into bundles. See ``FileBasedSource`` for more details. desired_bundle_size: Desired size of bundles that should be generated when splitting this source into bundles. See ``FileBasedSource`` for more details. compression_type: Used to handle compressed input files. Typical value is ``CompressionTypes.AUTO``, in which case the underlying file_path's extension will be used to detect the compression. strip_trailing_newlines: Indicates whether this source should remove the newline char in each line it reads before decoding that line. validate: flag to verify that the files exist during the pipeline creation time. skip_header_lines: Number of header lines to skip. Same number is skipped from each source file. Must be 0 or higher. Large number of skipped lines might impact performance. coder: Coder used to decode each line. with_filename: If True, returns a Key Value with the key being the file name and the value being the actual data. If False, it only returns the data. delimiter (bytes) Optional: delimiter to split records. Must not self-overlap, because self-overlapping delimiters cause ambiguous parsing. escapechar (bytes) Optional: a single byte to escape the records delimiter, can also escape itself. """ super().__init__(**kwargs) self._source_from_file = partial( _create_text_source, min_bundle_size=min_bundle_size, compression_type=compression_type, strip_trailing_newlines=strip_trailing_newlines, validate=validate, coder=coder, skip_header_lines=skip_header_lines, delimiter=delimiter, escapechar=escapechar) self._desired_bundle_size = desired_bundle_size self._min_bundle_size = min_bundle_size self._compression_type = compression_type self._with_filename = with_filename self._read_all_files = ReadAllFiles( True, self._compression_type, self._desired_bundle_size, self._min_bundle_size, self._source_from_file, self._with_filename) def expand(self, pvalue): return pvalue | 'ReadAllFiles' >> self._read_all_files class ReadAllFromTextContinuously(ReadAllFromText): """A ``PTransform`` for reading text files in given file patterns. This PTransform acts as a Source and produces continuously a ``PCollection`` of strings. For more details, see ``ReadAllFromText`` for text parsing settings; see ``apache_beam.io.fileio.MatchContinuously`` for watching settings. ReadAllFromTextContinuously is experimental. No backwards-compatibility guarantees. Due to the limitation on Reshuffle, current implementation does not scale. """ _ARGS_FOR_MATCH = ( 'interval', 'has_deduplication', 'start_timestamp', 'stop_timestamp', 'match_updated_files', 'apply_windowing') _ARGS_FOR_READ = ( 'min_bundle_size', 'desired_bundle_size', 'compression_type', 'strip_trailing_newlines', 'validate', 'coder', 'skip_header_lines', 'with_filename', 'delimiter', 'escapechar') def __init__(self, file_pattern, **kwargs): """Initialize the ``ReadAllFromTextContinuously`` transform. Accepts args for constructor args of both :class:`ReadAllFromText` and :class:`~apache_beam.io.fileio.MatchContinuously`. """ kwargs_for_match = { k: v for (k, v) in kwargs.items() if k in self._ARGS_FOR_MATCH } kwargs_for_read = { k: v for (k, v) in kwargs.items() if k in self._ARGS_FOR_READ } kwargs_additinal = { k: v for (k, v) in kwargs.items() if k not in self._ARGS_FOR_MATCH and k not in self._ARGS_FOR_READ } super().__init__(**kwargs_for_read, **kwargs_additinal) self._file_pattern = file_pattern self._kwargs_for_match = kwargs_for_match def expand(self, pbegin): # Importing locally to prevent circular dependency issues. from apache_beam.io.fileio import MatchContinuously # TODO(BEAM-14497) always reshuffle once gbk always trigger works. return ( pbegin | MatchContinuously(self._file_pattern, **self._kwargs_for_match) | 'ReadAllFiles' >> self._read_all_files._disable_reshuffle()) class ReadFromText(PTransform): r"""A :class:`~apache_beam.transforms.ptransform.PTransform` for reading text files. Parses a text file as newline-delimited elements, by default assuming ``UTF-8`` encoding. Supports newline delimiters ``\n`` and ``\r\n`` or specified delimiter. If reading from a text file that that requires a different encoding, you may provide a custom :class:`~apache_beam.coders.coders.Coder` that encodes and decodes with the appropriate codec. For example, see the implementation of :class:`~apache_beam.coders.coders.StrUtf8Coder`. This does not support ``UTF-16`` or ``UTF-32`` encodings. """ _source_class = _TextSource def __init__( self, file_pattern=None, min_bundle_size=0, compression_type=CompressionTypes.AUTO, strip_trailing_newlines=True, coder: coders.Coder = coders.StrUtf8Coder(), validate=True, skip_header_lines=0, delimiter=None, escapechar=None, **kwargs): """Initialize the :class:`ReadFromText` transform. Args: file_pattern (str): The file path to read from as a local file path or a GCS ``gs://`` path. The path can contain glob characters (``*``, ``?``, and ``[...]`` sets). min_bundle_size (int): Minimum size of bundles that should be generated when splitting this source into bundles. See :class:`~apache_beam.io.filebasedsource.FileBasedSource` for more details. compression_type (str): Used to handle compressed input files. Typical value is :attr:`CompressionTypes.AUTO `, in which case the underlying file_path's extension will be used to detect the compression. strip_trailing_newlines (bool): Indicates whether this source should remove the newline char in each line it reads before decoding that line. validate (bool): flag to verify that the files exist during the pipeline creation time. skip_header_lines (int): Number of header lines to skip. Same number is skipped from each source file. Must be 0 or higher. Large number of skipped lines might impact performance. coder (~apache_beam.coders.coders.Coder): Coder used to decode each line. delimiter (bytes) Optional: delimiter to split records. Must not self-overlap, because self-overlapping delimiters cause ambiguous parsing. escapechar (bytes) Optional: a single byte to escape the records delimiter, can also escape itself. """ super().__init__(**kwargs) if file_pattern: try: if not os.path.dirname(file_pattern): file_pattern = os.path.join('.', file_pattern) except TypeError: pass self._source = self._source_class( file_pattern, min_bundle_size, compression_type, strip_trailing_newlines, coder, validate=validate, skip_header_lines=skip_header_lines, delimiter=delimiter, escapechar=escapechar) def expand(self, pvalue): return pvalue.pipeline | Read(self._source).with_output_types( self._source.output_type_hint()) class ReadFromTextWithFilename(ReadFromText): r"""A :class:`~apache_beam.io.textio.ReadFromText` for reading text files returning the name of the file and the content of the file. This class extend ReadFromText class just setting a different _source_class attribute. """ _source_class = _TextSourceWithFilename class WriteToText(PTransform): """A :class:`~apache_beam.transforms.ptransform.PTransform` for writing to text files.""" def __init__( self, file_path_prefix: str, file_name_suffix='', append_trailing_newlines=True, num_shards=0, shard_name_template: Optional[str] = None, coder: coders.Coder = coders.ToBytesCoder(), compression_type=CompressionTypes.AUTO, header=None, footer=None, *, max_records_per_shard=None, max_bytes_per_shard=None, skip_if_empty=False, triggering_frequency=None): r"""Initialize a :class:`WriteToText` transform. Args: file_path_prefix (str): The file path to write to. The files written will begin with this prefix, followed by a shard identifier (see **num_shards**), and end in a common extension, if given by **file_name_suffix**. In most cases, only this argument is specified and **num_shards**, **shard_name_template**, and **file_name_suffix** use default values. file_name_suffix (str): Suffix for the files written. append_trailing_newlines (bool): indicate whether this sink should write an additional newline char after writing each element. num_shards (int): The number of files (shards) used for output. If not set, the service will decide on the optimal number of shards. Constraining the number of shards is likely to reduce the performance of a pipeline. Setting this value is not recommended unless you require a specific number of output files. shard_name_template (str): A template string containing placeholders for the shard number and shard count. Currently only ``''``, ``'-SSSSS-of-NNNNN'``, ``'-W-SSSSS-of-NNNNN'`` and ``'-V-SSSSS-of-NNNNN'`` are patterns accepted by the service. When constructing a filename for a particular shard number, the upper-case letters ``S`` and ``N`` are replaced with the ``0``-padded shard number and shard count respectively. This argument can be ``''`` in which case it behaves as if num_shards was set to 1 and only one file will be generated. The default pattern used is ``'-SSSSS-of-NNNNN'`` for bounded PCollections and for ``'-W-SSSSS-of-NNNNN'`` unbounded PCollections. W is used for windowed shard naming and is replaced with ``[window.start, window.end)`` V is used for windowed shard naming and is replaced with ``[window.start.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S"), window.end.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S")`` coder (~apache_beam.coders.coders.Coder): Coder used to encode each line. compression_type (str): Used to handle compressed output files. Typical value is :class:`CompressionTypes.AUTO `, in which case the final file path's extension (as determined by **file_path_prefix**, **file_name_suffix**, **num_shards** and **shard_name_template**) will be used to detect the compression. header (str): String to write at beginning of file as a header. If not :data:`None` and **append_trailing_newlines** is set, ``\n`` will be added. footer (str): String to write at the end of file as a footer. If not :data:`None` and **append_trailing_newlines** is set, ``\n`` will be added. max_records_per_shard: Maximum number of records to write to any individual shard. max_bytes_per_shard: Target maximum number of bytes to write to any individual shard. This may be exceeded slightly, as a new shard is created once this limit is hit, but the remainder of a given record, a subsequent newline, and a footer may cause the actual shard size to exceed this value. This also tracks the uncompressed, not compressed, size of the shard. skip_if_empty: Don't write any shards if the PCollection is empty. In case of an empty PCollection, this will still delete existing files having same file path and not create new ones. triggering_frequency: (int) Every triggering_frequency duration, a window will be triggered and all bundles in the window will be written. """ self._sink = _TextSink( file_path_prefix, file_name_suffix, append_trailing_newlines, num_shards, shard_name_template, coder, compression_type, header, footer, max_records_per_shard=max_records_per_shard, max_bytes_per_shard=max_bytes_per_shard, skip_if_empty=skip_if_empty, triggering_frequency=triggering_frequency) def expand(self, pcoll): if (not pcoll.is_bounded and self._sink.shard_name_template == filebasedsink.DEFAULT_SHARD_NAME_TEMPLATE): self._sink.shard_name_template = ( filebasedsink.DEFAULT_WINDOW_SHARD_NAME_TEMPLATE) self._sink.shard_name_format = self._sink._template_to_format( self._sink.shard_name_template) self._sink.shard_name_glob_format = self._sink._template_to_glob_format( self._sink.shard_name_template) return pcoll | Write(self._sink) try: import pandas def append_pandas_args(src, exclude): def append(dest): state = None skip = False extra_lines = [] for line in src.__doc__.split('\n'): if line.strip() == 'Parameters': indent = len(line) - len(line.lstrip()) extra_lines = ['\n\nPandas Parameters'] state = 'append' continue elif line.strip().startswith('Returns'): break if state == 'append': if skip: if line and not line[indent:].startswith(' '): skip = False if any(line.strip().startswith(arg + ' : ') for arg in exclude): skip = True if not skip: extra_lines.append(line[indent:]) # Expand title underline due to Parameters -> Pandas Parameters. extra_lines[1] += '-------' dest.__doc__ += '\n'.join(extra_lines) return dest return append @append_pandas_args( pandas.read_csv, exclude=['filepath_or_buffer', 'iterator']) def ReadFromCsv( path: str, *, splittable: bool = True, filename_column: Optional[str] = None, **kwargs): """A PTransform for reading comma-separated values (csv) files into a PCollection. Args: path (str): The file path to read from. The path can contain glob characters such as ``*`` and ``?``. splittable (bool): Whether the csv files are splittable at line boundaries, i.e. each line of this file represents a complete record. This should be set to False if single records span multiple lines (e.g. a quoted field has a newline inside of it). Setting this to false may disable liquid sharding. filename_column (str): If not None, the name of the column to add to each record, containing the filename of the source file. **kwargs: Extra arguments passed to `pandas.read_csv` (see below). """ from apache_beam.dataframe.io import ReadViaPandas return 'ReadFromCsv' >> ReadViaPandas( 'csv', path, splittable=splittable, filename_column=filename_column, **kwargs) @append_pandas_args( pandas.DataFrame.to_csv, exclude=['path_or_buf', 'index', 'index_label']) def WriteToCsv( path: str, num_shards: Optional[int] = None, file_naming: Optional['fileio.FileNaming'] = None, **kwargs): # pylint: disable=line-too-long """A PTransform for writing a schema'd PCollection as a (set of) comma-separated values (csv) files. Args: path (str): The file path to write to. The files written will begin with this prefix, followed by a shard identifier (see `num_shards`) according to the `file_naming` parameter. num_shards (optional int): The number of shards to use in the distributed write. Defaults to None, letting the system choose an optimal value. file_naming (optional callable): A file-naming strategy, determining the actual shard names given their shard number, etc. See the section on `file naming `_ Defaults to `fileio.default_file_naming`, which names files as `path-XXXXX-of-NNNNN`. **kwargs: Extra arguments passed to `pandas.Dataframe.to_csv` (see below). """ from apache_beam.dataframe.io import WriteViaPandas if num_shards is not None: kwargs['num_shards'] = num_shards if file_naming is not None: kwargs['file_naming'] = file_naming return 'WriteToCsv' >> WriteViaPandas('csv', path, index=False, **kwargs) @append_pandas_args(pandas.read_json, exclude=['path_or_buf']) def ReadFromJson( path: str, *, orient: str = 'records', lines: bool = True, dtype: Union[bool, dict[str, Any]] = False, **kwargs): """A PTransform for reading json values from files into a PCollection. Args: path (str): The file path to read from. The path can contain glob characters such as ``*`` and ``?``. orient (str): Format of the json elements in the file. Default to 'records', meaning the file is expected to contain a list of json objects like `{field1: value1, field2: value2, ...}`. lines (bool): Whether each line should be considered a separate record, as opposed to the entire file being a valid JSON object or list. Defaults to True (unlike Pandas). dtype (bool): If True, infer dtypes; if a dict of column to dtype, then use those; if False, then don’t infer dtypes at all. Defaults to False (unlike Pandas). **kwargs: Extra arguments passed to `pandas.read_json` (see below). """ from apache_beam.dataframe.io import ReadViaPandas return 'ReadFromJson' >> ReadViaPandas( 'json', path, orient=orient, lines=lines, dtype=dtype, **kwargs) @append_pandas_args( pandas.DataFrame.to_json, exclude=['path_or_buf', 'index']) def WriteToJson( path: str, *, num_shards: Optional[int] = None, file_naming: Optional['fileio.FileNaming'] = None, orient: str = 'records', lines: Optional[bool] = None, **kwargs): # pylint: disable=line-too-long """A PTransform for writing a PCollection as json values to files. Args: path (str): The file path to write to. The files written will begin with this prefix, followed by a shard identifier (see `num_shards`) according to the `file_naming` parameter. num_shards (optional int): The number of shards to use in the distributed write. Defaults to None, letting the system choose an optimal value. file_naming (optional callable): A file-naming strategy, determining the actual shard names given their shard number, etc. See the section on `file naming `_ Defaults to `fileio.default_file_naming`, which names files as `path-XXXXX-of-NNNNN`. orient (str): Format of the json elements in the file. Default to 'records', meaning the file will to contain a list of json objects like `{field1: value1, field2: value2, ...}`. lines (bool): Whether each line should be considered a separate record, as opposed to the entire file being a valid JSON object or list. Defaults to True if orient is 'records' (unlike Pandas). **kwargs: Extra arguments passed to `pandas.Dataframe.to_json` (see below). """ from apache_beam.dataframe.io import WriteViaPandas if num_shards is not None: kwargs['num_shards'] = num_shards if file_naming is not None: kwargs['file_naming'] = file_naming if lines is None: lines = orient == 'records' return 'WriteToJson' >> WriteViaPandas( 'json', path, orient=orient, lines=lines, **kwargs) except ImportError: def no_pandas(*args, **kwargs): raise ImportError('Please install apache_beam[dataframe]') for transform in ('ReadFromCsv', 'WriteToCsv', 'ReadFromJson', 'WriteToJson'): globals()[transform] = no_pandas ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/textio_test.py0000644000175100017510000024754515206113032021232 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for textio module.""" # pytype: skip-file import bz2 import glob import gzip import logging import os import platform import re import shutil import tempfile import unittest import zlib from datetime import datetime import pytz import apache_beam as beam from apache_beam import coders from apache_beam.io import iobase from apache_beam.io import source_test_utils from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystems import FileSystems # Importing following private classes for testing. from apache_beam.io.textio import ReadAllFromText from apache_beam.io.textio import ReadAllFromTextContinuously from apache_beam.io.textio import ReadFromText from apache_beam.io.textio import ReadFromTextWithFilename from apache_beam.io.textio import WriteToText from apache_beam.io.textio import _TextSink as TextSink from apache_beam.io.textio import _TextSource as TextSource from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.test_utils import TempDir from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.core import Create from apache_beam.transforms.userstate import CombiningValueStateSpec from apache_beam.transforms.util import LogElements from apache_beam.utils.timestamp import Timestamp class DummyCoder(coders.Coder): def encode(self, x): raise ValueError def decode(self, x): return (x * 2).decode('utf-8') def to_type_hint(self): return str class EOL(object): LF = 1 CRLF = 2 MIXED = 3 LF_WITH_NOTHING_AT_LAST_LINE = 4 CUSTOM_DELIMITER = 5 def write_data( num_lines, no_data=False, directory=None, prefix=tempfile.template, eol=EOL.LF, custom_delimiter=None, line_value=b'line'): """Writes test data to a temporary file. Args: num_lines (int): The number of lines to write. no_data (bool): If :data:`True`, empty lines will be written, otherwise each line will contain a concatenation of b'line' and the line number. directory (str): The name of the directory to create the temporary file in. prefix (str): The prefix to use for the temporary file. eol (int): The line ending to use when writing. :class:`~apache_beam.io.textio_test.EOL` exposes attributes that can be used here to define the eol. custom_delimiter (bytes): The custom delimiter. line_value (bytes): Default value for test data, default b'line' Returns: Tuple[str, List[str]]: A tuple of the filename and a list of the utf-8 decoded written data. """ all_data = [] with tempfile.NamedTemporaryFile(delete=False, dir=directory, prefix=prefix) as f: sep_values = [b'\n', b'\r\n'] for i in range(num_lines): data = b'' if no_data else line_value + str(i).encode() all_data.append(data) if eol == EOL.LF: sep = sep_values[0] elif eol == EOL.CRLF: sep = sep_values[1] elif eol == EOL.MIXED: sep = sep_values[i % len(sep_values)] elif eol == EOL.LF_WITH_NOTHING_AT_LAST_LINE: sep = b'' if i == (num_lines - 1) else sep_values[0] elif eol == EOL.CUSTOM_DELIMITER: if custom_delimiter is None or len(custom_delimiter) == 0: raise ValueError('delimiter can not be null or empty') else: sep = custom_delimiter else: raise ValueError('Received unknown value %s for eol.' % eol) f.write(data + sep) return f.name, [line.decode('utf-8') for line in all_data] def write_pattern(lines_per_file, no_data=False, return_filenames=False): """Writes a pattern of temporary files. Args: lines_per_file (List[int]): The number of lines to write per file. no_data (bool): If :data:`True`, empty lines will be written, otherwise each line will contain a concatenation of b'line' and the line number. return_filenames (bool): If True, returned list will contain (filename, data) pairs. Returns: Tuple[str, List[Union[str, (str, str)]]]: A tuple of the filename pattern and a list of the utf-8 decoded written data or (filename, data) pairs. """ temp_dir = tempfile.mkdtemp() all_data = [] file_name = None start_index = 0 for i in range(len(lines_per_file)): file_name, data = write_data(lines_per_file[i], no_data=no_data, directory=temp_dir, prefix='mytemp') if return_filenames: all_data.extend(zip([file_name] * len(data), data)) else: all_data.extend(data) start_index += lines_per_file[i] assert file_name return ( file_name[:file_name.rfind(os.path.sep)] + os.path.sep + 'mytemp*', all_data) class TextSourceTest(unittest.TestCase): # Number of records that will be written by most tests. DEFAULT_NUM_RECORDS = 100 def _run_read_test( self, file_or_pattern, expected_data, buffer_size=DEFAULT_NUM_RECORDS, compression=CompressionTypes.UNCOMPRESSED, delimiter=None, escapechar=None): # Since each record usually takes more than 1 byte, default buffer size is # smaller than the total size of the file. This is done to # increase test coverage for cases that hit the buffer boundary. kwargs = {} if delimiter: kwargs['delimiter'] = delimiter if escapechar: kwargs['escapechar'] = escapechar source = TextSource( file_or_pattern, 0, compression, True, coders.StrUtf8Coder(), buffer_size, **kwargs) range_tracker = source.get_range_tracker(None, None) read_data = list(source.read(range_tracker)) self.assertCountEqual(expected_data, read_data) @unittest.skipIf(platform.system() == 'Windows', 'Skipping on Windows') def test_read_from_text_file_pattern_with_dot_slash(self): cwd = os.getcwd() expected = ['abc', 'de'] with TempDir() as temp_dir: temp_dir.create_temp_file(suffix='.txt', lines=[b'a', b'b', b'c']) temp_dir.create_temp_file(suffix='.txt', lines=[b'd', b'e']) os.chdir(temp_dir.get_path()) with TestPipeline() as p: dot_slash = p | 'ReadDotSlash' >> ReadFromText('./*.txt') no_dot_slash = p | 'ReadNoSlash' >> ReadFromText('*.txt') assert_that(dot_slash, equal_to(expected)) assert_that(no_dot_slash, equal_to(expected)) os.chdir(cwd) def test_read_from_text_with_value_provider(self): class UserDefinedOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument( '--file_pattern', help='This keyword argument is a value provider', default='some value') options = UserDefinedOptions(['--file_pattern', 'abc']) with self.assertRaises(OSError): with TestPipeline(options=options) as pipeline: _ = pipeline | 'Read' >> ReadFromText(options.file_pattern) def test_read_single_file(self): file_name, expected_data = write_data(TextSourceTest.DEFAULT_NUM_RECORDS) assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS self._run_read_test(file_name, expected_data) def test_read_single_file_smaller_than_default_buffer(self): file_name, expected_data = write_data(TextSourceTest.DEFAULT_NUM_RECORDS) self._run_read_test( file_name, expected_data, buffer_size=TextSource.DEFAULT_READ_BUFFER_SIZE) def test_read_single_file_larger_than_default_buffer(self): file_name, expected_data = write_data(TextSource.DEFAULT_READ_BUFFER_SIZE) self._run_read_test( file_name, expected_data, buffer_size=TextSource.DEFAULT_READ_BUFFER_SIZE) def test_read_file_pattern(self): pattern, expected_data = write_pattern( [TextSourceTest.DEFAULT_NUM_RECORDS * 5, TextSourceTest.DEFAULT_NUM_RECORDS * 3, TextSourceTest.DEFAULT_NUM_RECORDS * 12, TextSourceTest.DEFAULT_NUM_RECORDS * 8, TextSourceTest.DEFAULT_NUM_RECORDS * 8, TextSourceTest.DEFAULT_NUM_RECORDS * 4]) assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS * 40 self._run_read_test(pattern, expected_data) def test_read_single_file_windows_eol(self): file_name, expected_data = write_data(TextSourceTest.DEFAULT_NUM_RECORDS, eol=EOL.CRLF) assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS self._run_read_test(file_name, expected_data) def test_read_single_file_mixed_eol(self): file_name, expected_data = write_data(TextSourceTest.DEFAULT_NUM_RECORDS, eol=EOL.MIXED) assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS self._run_read_test(file_name, expected_data) def test_read_single_file_last_line_no_eol(self): file_name, expected_data = write_data( TextSourceTest.DEFAULT_NUM_RECORDS, eol=EOL.LF_WITH_NOTHING_AT_LAST_LINE) assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS self._run_read_test(file_name, expected_data) def test_read_single_file_single_line_no_eol(self): file_name, expected_data = write_data( 1, eol=EOL.LF_WITH_NOTHING_AT_LAST_LINE) assert len(expected_data) == 1 self._run_read_test(file_name, expected_data) def test_read_empty_single_file(self): file_name, written_data = write_data( 1, no_data=True, eol=EOL.LF_WITH_NOTHING_AT_LAST_LINE) assert len(written_data) == 1 # written data has a single entry with an empty string. Reading the source # should not produce anything since we only wrote a single empty string # without an end of line character. self._run_read_test(file_name, []) def test_read_single_file_last_line_no_eol_gzip(self): file_name, expected_data = write_data( TextSourceTest.DEFAULT_NUM_RECORDS, eol=EOL.LF_WITH_NOTHING_AT_LAST_LINE) gzip_file_name = file_name + '.gz' with open(file_name, 'rb') as src, gzip.open(gzip_file_name, 'wb') as dst: dst.writelines(src) assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS self._run_read_test( gzip_file_name, expected_data, compression=CompressionTypes.GZIP) def test_read_single_file_single_line_no_eol_gzip(self): file_name, expected_data = write_data( 1, eol=EOL.LF_WITH_NOTHING_AT_LAST_LINE) gzip_file_name = file_name + '.gz' with open(file_name, 'rb') as src, gzip.open(gzip_file_name, 'wb') as dst: dst.writelines(src) assert len(expected_data) == 1 self._run_read_test( gzip_file_name, expected_data, compression=CompressionTypes.GZIP) def test_read_empty_single_file_no_eol_gzip(self): file_name, written_data = write_data( 1, no_data=True, eol=EOL.LF_WITH_NOTHING_AT_LAST_LINE) gzip_file_name = file_name + '.gz' with open(file_name, 'rb') as src, gzip.open(gzip_file_name, 'wb') as dst: dst.writelines(src) assert len(written_data) == 1 # written data has a single entry with an empty string. Reading the source # should not produce anything since we only wrote a single empty string # without an end of line character. self._run_read_test(gzip_file_name, [], compression=CompressionTypes.GZIP) def test_read_single_file_with_empty_lines(self): file_name, expected_data = write_data( TextSourceTest.DEFAULT_NUM_RECORDS, no_data=True, eol=EOL.LF) assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS assert not expected_data[0] self._run_read_test(file_name, expected_data) def test_read_single_file_without_striping_eol_lf(self): file_name, written_data = write_data(TextSourceTest.DEFAULT_NUM_RECORDS, eol=EOL.LF) assert len(written_data) == TextSourceTest.DEFAULT_NUM_RECORDS source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, False, coders.StrUtf8Coder()) range_tracker = source.get_range_tracker(None, None) read_data = list(source.read(range_tracker)) self.assertCountEqual([line + '\n' for line in written_data], read_data) def test_read_single_file_without_striping_eol_crlf(self): file_name, written_data = write_data(TextSourceTest.DEFAULT_NUM_RECORDS, eol=EOL.CRLF) assert len(written_data) == TextSourceTest.DEFAULT_NUM_RECORDS source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, False, coders.StrUtf8Coder()) range_tracker = source.get_range_tracker(None, None) read_data = list(source.read(range_tracker)) self.assertCountEqual([line + '\r\n' for line in written_data], read_data) def test_read_file_pattern_with_empty_files(self): pattern, expected_data = write_pattern( [5 * TextSourceTest.DEFAULT_NUM_RECORDS, 3 * TextSourceTest.DEFAULT_NUM_RECORDS, 12 * TextSourceTest.DEFAULT_NUM_RECORDS, 8 * TextSourceTest.DEFAULT_NUM_RECORDS, 8 * TextSourceTest.DEFAULT_NUM_RECORDS, 4 * TextSourceTest.DEFAULT_NUM_RECORDS], no_data=True) assert len(expected_data) == 40 * TextSourceTest.DEFAULT_NUM_RECORDS assert not expected_data[0] self._run_read_test(pattern, expected_data) def test_read_after_splitting(self): file_name, expected_data = write_data(10) assert len(expected_data) == 10 source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder()) splits = list(source.split(desired_bundle_size=33)) reference_source_info = (source, None, None) sources_info = ([(split.source, split.start_position, split.stop_position) for split in splits]) source_test_utils.assert_sources_equal_reference_source( reference_source_info, sources_info) def test_header_processing(self): file_name, expected_data = write_data(10) assert len(expected_data) == 10 def header_matcher(line): return line in expected_data[:5] header_lines = [] def store_header(lines): for line in lines: header_lines.append(line) source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder(), header_processor_fns=(header_matcher, store_header)) splits = list(source.split(desired_bundle_size=100000)) assert len(splits) == 1 range_tracker = splits[0].source.get_range_tracker( splits[0].start_position, splits[0].stop_position) read_data = list(source.read_records(file_name, range_tracker)) self.assertCountEqual(expected_data[:5], header_lines) self.assertCountEqual(expected_data[5:], read_data) def test_progress(self): file_name, expected_data = write_data(10) assert len(expected_data) == 10 source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder()) splits = list(source.split(desired_bundle_size=100000)) assert len(splits) == 1 fraction_consumed_report = [] split_points_report = [] range_tracker = splits[0].source.get_range_tracker( splits[0].start_position, splits[0].stop_position) for _ in splits[0].source.read(range_tracker): fraction_consumed_report.append(range_tracker.fraction_consumed()) split_points_report.append(range_tracker.split_points()) self.assertEqual([float(i) / 10 for i in range(0, 10)], fraction_consumed_report) expected_split_points_report = [((i - 1), iobase.RangeTracker.SPLIT_POINTS_UNKNOWN) for i in range(1, 10)] # At last split point, the remaining split points callback returns 1 since # the expected position of next record becomes equal to the stop position. expected_split_points_report.append((9, 1)) self.assertEqual(expected_split_points_report, split_points_report) def test_read_reentrant_without_splitting(self): file_name, expected_data = write_data(10) assert len(expected_data) == 10 source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder()) source_test_utils.assert_reentrant_reads_succeed((source, None, None)) def test_read_reentrant_after_splitting(self): file_name, expected_data = write_data(10) assert len(expected_data) == 10 source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder()) splits = list(source.split(desired_bundle_size=100000)) assert len(splits) == 1 source_test_utils.assert_reentrant_reads_succeed( (splits[0].source, splits[0].start_position, splits[0].stop_position)) def test_dynamic_work_rebalancing(self): file_name, expected_data = write_data(5) assert len(expected_data) == 5 source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder()) splits = list(source.split(desired_bundle_size=100000)) assert len(splits) == 1 source_test_utils.assert_split_at_fraction_exhaustive( splits[0].source, splits[0].start_position, splits[0].stop_position) def test_dynamic_work_rebalancing_windows_eol(self): file_name, expected_data = write_data(15, eol=EOL.CRLF) assert len(expected_data) == 15 source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder()) splits = list(source.split(desired_bundle_size=100000)) assert len(splits) == 1 source_test_utils.assert_split_at_fraction_exhaustive( splits[0].source, splits[0].start_position, splits[0].stop_position, perform_multi_threaded_test=False) def test_dynamic_work_rebalancing_mixed_eol(self): file_name, expected_data = write_data(5, eol=EOL.MIXED) assert len(expected_data) == 5 source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder()) splits = list(source.split(desired_bundle_size=100000)) assert len(splits) == 1 source_test_utils.assert_split_at_fraction_exhaustive( splits[0].source, splits[0].start_position, splits[0].stop_position, perform_multi_threaded_test=False) def test_read_from_text_single_file(self): file_name, expected_data = write_data(5) assert len(expected_data) == 5 with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText(file_name) assert_that(pcoll, equal_to(expected_data)) def test_read_from_text_with_file_name_single_file(self): file_name, data = write_data(5) expected_data = [(file_name, el) for el in data] assert len(expected_data) == 5 with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromTextWithFilename(file_name) assert_that(pcoll, equal_to(expected_data)) def test_read_all_single_file(self): file_name, expected_data = write_data(5) assert len(expected_data) == 5 with TestPipeline() as pipeline: pcoll = pipeline | 'Create' >> Create( [file_name]) | 'ReadAll' >> ReadAllFromText() assert_that(pcoll, equal_to(expected_data)) def test_read_all_many_single_files(self): file_name1, expected_data1 = write_data(5) assert len(expected_data1) == 5 file_name2, expected_data2 = write_data(10) assert len(expected_data2) == 10 file_name3, expected_data3 = write_data(15) assert len(expected_data3) == 15 expected_data = [] expected_data.extend(expected_data1) expected_data.extend(expected_data2) expected_data.extend(expected_data3) with TestPipeline() as pipeline: pcoll = pipeline | 'Create' >> Create([ file_name1, file_name2, file_name3 ]) | 'ReadAll' >> ReadAllFromText() assert_that(pcoll, equal_to(expected_data)) def test_read_all_unavailable_files_ignored(self): file_name1, expected_data1 = write_data(5) assert len(expected_data1) == 5 file_name2, expected_data2 = write_data(10) assert len(expected_data2) == 10 file_name3, expected_data3 = write_data(15) assert len(expected_data3) == 15 file_name4 = "/unavailable_file" expected_data = [] expected_data.extend(expected_data1) expected_data.extend(expected_data2) expected_data.extend(expected_data3) with TestPipeline() as pipeline: pcoll = ( pipeline | 'Create' >> Create([file_name1, file_name2, file_name3, file_name4]) | 'ReadAll' >> ReadAllFromText()) assert_that(pcoll, equal_to(expected_data)) class _WriteFilesFn(beam.DoFn): """writes a couple of files with deferral.""" COUNT_STATE = CombiningValueStateSpec('count', combine_fn=sum) def __init__(self, temp_path): self.temp_path = temp_path def process(self, element, count_state=beam.DoFn.StateParam(COUNT_STATE)): counter = count_state.read() if counter == 0: count_state.add(1) with open(FileSystems.join(self.temp_path, 'file1'), 'w') as f: f.write('second A\nsecond B') with open(FileSystems.join(self.temp_path, 'file2'), 'w') as f: f.write('first') # convert dumb key to basename in output basename = FileSystems.split(element[1][0])[1] content = element[1][1] yield basename, content def test_read_all_continuously_new(self): with TempDir() as tempdir, TestPipeline() as pipeline: temp_path = tempdir.get_path() # create a temp file at the beginning with open(FileSystems.join(temp_path, 'file1'), 'w') as f: f.write('first') match_pattern = FileSystems.join(temp_path, '*') interval = 0.5 last = 2 p_read_once = ( pipeline | 'Continuously read new files' >> ReadAllFromTextContinuously( match_pattern, with_filename=True, start_timestamp=Timestamp.now(), interval=interval, stop_timestamp=Timestamp.now() + last, match_updated_files=False) | 'add dumb key' >> beam.Map(lambda x: (0, x)) | 'Write files on-the-fly' >> beam.ParDo(self._WriteFilesFn(temp_path))) assert_that( p_read_once, equal_to([('file1', 'first'), ('file2', 'first')]), label='assert read new files results') def test_read_all_continuously_update(self): with TempDir() as tempdir, TestPipeline() as pipeline: temp_path = tempdir.get_path() # create a temp file at the beginning with open(FileSystems.join(temp_path, 'file1'), 'w') as f: f.write('first') match_pattern = FileSystems.join(temp_path, '*') interval = 0.5 last = 2 p_read_upd = ( pipeline | 'Continuously read updated files' >> ReadAllFromTextContinuously( match_pattern, with_filename=True, start_timestamp=Timestamp.now(), interval=interval, stop_timestamp=Timestamp.now() + last, match_updated_files=True) | 'add dumb key' >> beam.Map(lambda x: (0, x)) | 'Write files on-the-fly' >> beam.ParDo(self._WriteFilesFn(temp_path))) assert_that( p_read_upd, equal_to([('file1', 'first'), ('file1', 'second A'), ('file1', 'second B'), ('file2', 'first')]), label='assert read updated files results') def test_read_from_text_single_file_with_coder(self): file_name, expected_data = write_data(5) assert len(expected_data) == 5 with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText(file_name, coder=DummyCoder()) assert_that(pcoll, equal_to([record * 2 for record in expected_data])) def test_read_from_text_file_pattern(self): pattern, expected_data = write_pattern([5, 3, 12, 8, 8, 4]) assert len(expected_data) == 40 with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText(pattern) assert_that(pcoll, equal_to(expected_data)) def test_read_from_text_with_file_name_file_pattern(self): pattern, expected_data = write_pattern( lines_per_file=[5, 5], return_filenames=True) assert len(expected_data) == 10 with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromTextWithFilename(pattern) assert_that(pcoll, equal_to(expected_data)) def test_read_all_file_pattern(self): pattern, expected_data = write_pattern([5, 3, 12, 8, 8, 4]) assert len(expected_data) == 40 with TestPipeline() as pipeline: pcoll = ( pipeline | 'Create' >> Create([pattern]) | 'ReadAll' >> ReadAllFromText()) assert_that(pcoll, equal_to(expected_data)) def test_read_all_many_file_patterns(self): pattern1, expected_data1 = write_pattern([5, 3, 12, 8, 8, 4]) assert len(expected_data1) == 40 pattern2, expected_data2 = write_pattern([3, 7, 9]) assert len(expected_data2) == 19 pattern3, expected_data3 = write_pattern([11, 20, 5, 5]) assert len(expected_data3) == 41 expected_data = [] expected_data.extend(expected_data1) expected_data.extend(expected_data2) expected_data.extend(expected_data3) with TestPipeline() as pipeline: pcoll = pipeline | 'Create' >> Create( [pattern1, pattern2, pattern3]) | 'ReadAll' >> ReadAllFromText() assert_that(pcoll, equal_to(expected_data)) def test_read_all_with_filename(self): pattern, expected_data = write_pattern([5, 3], return_filenames=True) assert len(expected_data) == 8 with TestPipeline() as pipeline: pcoll = ( pipeline | 'Create' >> Create([pattern]) | 'ReadAll' >> ReadAllFromText(with_filename=True)) assert_that(pcoll, equal_to(expected_data)) def test_read_auto_bzip2(self): _, lines = write_data(15) with TempDir() as tempdir: file_name = tempdir.create_temp_file(suffix='.bz2') with bz2.BZ2File(file_name, 'wb') as f: f.write('\n'.join(lines).encode('utf-8')) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText(file_name) assert_that(pcoll, equal_to(lines)) def test_read_auto_deflate(self): _, lines = write_data(15) with TempDir() as tempdir: file_name = tempdir.create_temp_file(suffix='.deflate') with open(file_name, 'wb') as f: f.write(zlib.compress('\n'.join(lines).encode('utf-8'))) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText(file_name) assert_that(pcoll, equal_to(lines)) def test_read_auto_gzip(self): _, lines = write_data(15) with TempDir() as tempdir: file_name = tempdir.create_temp_file(suffix='.gz') with gzip.GzipFile(file_name, 'wb') as f: f.write('\n'.join(lines).encode('utf-8')) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText(file_name) assert_that(pcoll, equal_to(lines)) def test_read_bzip2(self): _, lines = write_data(15) with TempDir() as tempdir: file_name = tempdir.create_temp_file() with bz2.BZ2File(file_name, 'wb') as f: f.write('\n'.join(lines).encode('utf-8')) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText( file_name, compression_type=CompressionTypes.BZIP2) assert_that(pcoll, equal_to(lines)) def test_read_corrupted_bzip2_fails(self): _, lines = write_data(15) with TempDir() as tempdir: file_name = tempdir.create_temp_file() with bz2.BZ2File(file_name, 'wb') as f: f.write('\n'.join(lines).encode('utf-8')) with open(file_name, 'wb') as f: f.write(b'corrupt') with self.assertRaises(Exception): with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText( file_name, compression_type=CompressionTypes.BZIP2) assert_that(pcoll, equal_to(lines)) def test_read_bzip2_concat(self): with TempDir() as tempdir: bzip2_file_name1 = tempdir.create_temp_file() lines = ['a', 'b', 'c'] with bz2.BZ2File(bzip2_file_name1, 'wb') as dst: data = '\n'.join(lines) + '\n' dst.write(data.encode('utf-8')) bzip2_file_name2 = tempdir.create_temp_file() lines = ['p', 'q', 'r'] with bz2.BZ2File(bzip2_file_name2, 'wb') as dst: data = '\n'.join(lines) + '\n' dst.write(data.encode('utf-8')) bzip2_file_name3 = tempdir.create_temp_file() lines = ['x', 'y', 'z'] with bz2.BZ2File(bzip2_file_name3, 'wb') as dst: data = '\n'.join(lines) + '\n' dst.write(data.encode('utf-8')) final_bzip2_file = tempdir.create_temp_file() with open(bzip2_file_name1, 'rb') as src, open( final_bzip2_file, 'wb') as dst: dst.writelines(src.readlines()) with open(bzip2_file_name2, 'rb') as src, open( final_bzip2_file, 'ab') as dst: dst.writelines(src.readlines()) with open(bzip2_file_name3, 'rb') as src, open( final_bzip2_file, 'ab') as dst: dst.writelines(src.readlines()) with TestPipeline() as pipeline: lines = pipeline | 'ReadFromText' >> beam.io.ReadFromText( final_bzip2_file, compression_type=beam.io.filesystem.CompressionTypes.BZIP2) expected = ['a', 'b', 'c', 'p', 'q', 'r', 'x', 'y', 'z'] assert_that(lines, equal_to(expected)) def test_read_deflate(self): _, lines = write_data(15) with TempDir() as tempdir: file_name = tempdir.create_temp_file() with open(file_name, 'wb') as f: f.write(zlib.compress('\n'.join(lines).encode('utf-8'))) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText( file_name, 0, CompressionTypes.DEFLATE, True, coders.StrUtf8Coder()) assert_that(pcoll, equal_to(lines)) def test_read_corrupted_deflate_fails(self): _, lines = write_data(15) with TempDir() as tempdir: file_name = tempdir.create_temp_file() with open(file_name, 'wb') as f: f.write(zlib.compress('\n'.join(lines).encode('utf-8'))) with open(file_name, 'wb') as f: f.write(b'corrupt') with self.assertRaises(Exception): with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText( file_name, 0, CompressionTypes.DEFLATE, True, coders.StrUtf8Coder()) assert_that(pcoll, equal_to(lines)) def test_read_deflate_concat(self): with TempDir() as tempdir: deflate_file_name1 = tempdir.create_temp_file() lines = ['a', 'b', 'c'] with open(deflate_file_name1, 'wb') as dst: data = '\n'.join(lines) + '\n' dst.write(zlib.compress(data.encode('utf-8'))) deflate_file_name2 = tempdir.create_temp_file() lines = ['p', 'q', 'r'] with open(deflate_file_name2, 'wb') as dst: data = '\n'.join(lines) + '\n' dst.write(zlib.compress(data.encode('utf-8'))) deflate_file_name3 = tempdir.create_temp_file() lines = ['x', 'y', 'z'] with open(deflate_file_name3, 'wb') as dst: data = '\n'.join(lines) + '\n' dst.write(zlib.compress(data.encode('utf-8'))) final_deflate_file = tempdir.create_temp_file() with open(deflate_file_name1, 'rb') as src, \ open(final_deflate_file, 'wb') as dst: dst.writelines(src.readlines()) with open(deflate_file_name2, 'rb') as src, \ open(final_deflate_file, 'ab') as dst: dst.writelines(src.readlines()) with open(deflate_file_name3, 'rb') as src, \ open(final_deflate_file, 'ab') as dst: dst.writelines(src.readlines()) with TestPipeline() as pipeline: lines = pipeline | 'ReadFromText' >> beam.io.ReadFromText( final_deflate_file, compression_type=beam.io.filesystem.CompressionTypes.DEFLATE) expected = ['a', 'b', 'c', 'p', 'q', 'r', 'x', 'y', 'z'] assert_that(lines, equal_to(expected)) def test_read_gzip(self): _, lines = write_data(15) with TempDir() as tempdir: file_name = tempdir.create_temp_file() with gzip.GzipFile(file_name, 'wb') as f: f.write('\n'.join(lines).encode('utf-8')) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText( file_name, 0, CompressionTypes.GZIP, True, coders.StrUtf8Coder()) assert_that(pcoll, equal_to(lines)) def test_read_corrupted_gzip_fails(self): _, lines = write_data(15) with TempDir() as tempdir: file_name = tempdir.create_temp_file() with gzip.GzipFile(file_name, 'wb') as f: f.write('\n'.join(lines).encode('utf-8')) with open(file_name, 'wb') as f: f.write(b'corrupt') with self.assertRaises(Exception): with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText( file_name, 0, CompressionTypes.GZIP, True, coders.StrUtf8Coder()) assert_that(pcoll, equal_to(lines)) def test_read_gzip_concat(self): with TempDir() as tempdir: gzip_file_name1 = tempdir.create_temp_file() lines = ['a', 'b', 'c'] with gzip.open(gzip_file_name1, 'wb') as dst: data = '\n'.join(lines) + '\n' dst.write(data.encode('utf-8')) gzip_file_name2 = tempdir.create_temp_file() lines = ['p', 'q', 'r'] with gzip.open(gzip_file_name2, 'wb') as dst: data = '\n'.join(lines) + '\n' dst.write(data.encode('utf-8')) gzip_file_name3 = tempdir.create_temp_file() lines = ['x', 'y', 'z'] with gzip.open(gzip_file_name3, 'wb') as dst: data = '\n'.join(lines) + '\n' dst.write(data.encode('utf-8')) final_gzip_file = tempdir.create_temp_file() with open(gzip_file_name1, 'rb') as src, \ open(final_gzip_file, 'wb') as dst: dst.writelines(src.readlines()) with open(gzip_file_name2, 'rb') as src, \ open(final_gzip_file, 'ab') as dst: dst.writelines(src.readlines()) with open(gzip_file_name3, 'rb') as src, \ open(final_gzip_file, 'ab') as dst: dst.writelines(src.readlines()) with TestPipeline() as pipeline: lines = pipeline | 'ReadFromText' >> beam.io.ReadFromText( final_gzip_file, compression_type=beam.io.filesystem.CompressionTypes.GZIP) expected = ['a', 'b', 'c', 'p', 'q', 'r', 'x', 'y', 'z'] assert_that(lines, equal_to(expected)) def test_read_all_gzip(self): _, lines = write_data(100) with TempDir() as tempdir: file_name = tempdir.create_temp_file() with gzip.GzipFile(file_name, 'wb') as f: f.write('\n'.join(lines).encode('utf-8')) with TestPipeline() as pipeline: pcoll = ( pipeline | Create([file_name]) | 'ReadAll' >> ReadAllFromText(compression_type=CompressionTypes.GZIP)) assert_that(pcoll, equal_to(lines)) def test_read_gzip_large(self): _, lines = write_data(10000) with TempDir() as tempdir: file_name = tempdir.create_temp_file() with gzip.GzipFile(file_name, 'wb') as f: f.write('\n'.join(lines).encode('utf-8')) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText( file_name, 0, CompressionTypes.GZIP, True, coders.StrUtf8Coder()) assert_that(pcoll, equal_to(lines)) def test_read_gzip_large_after_splitting(self): _, lines = write_data(10000) with TempDir() as tempdir: file_name = tempdir.create_temp_file() with gzip.GzipFile(file_name, 'wb') as f: f.write('\n'.join(lines).encode('utf-8')) source = TextSource( file_name, 0, CompressionTypes.GZIP, True, coders.StrUtf8Coder()) splits = list(source.split(desired_bundle_size=1000)) if len(splits) > 1: raise ValueError( 'FileBasedSource generated more than one initial ' 'split for a compressed file.') reference_source_info = (source, None, None) sources_info = ([ (split.source, split.start_position, split.stop_position) for split in splits ]) source_test_utils.assert_sources_equal_reference_source( reference_source_info, sources_info) def test_read_gzip_empty_file(self): with TempDir() as tempdir: file_name = tempdir.create_temp_file() with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText( file_name, 0, CompressionTypes.GZIP, True, coders.StrUtf8Coder()) assert_that(pcoll, equal_to([])) def _remove_lines(self, lines, sublist_lengths, num_to_remove): """Utility function to remove num_to_remove lines from each sublist. Args: lines: list of items. sublist_lengths: list of integers representing length of sublist corresponding to each source file. num_to_remove: number of lines to remove from each sublist. Returns: remaining lines. """ curr = 0 result = [] for offset in sublist_lengths: end = curr + offset start = min(curr + num_to_remove, end) result += lines[start:end] curr += offset return result def _read_skip_header_lines(self, file_or_pattern, skip_header_lines): """Simple wrapper function for instantiating TextSource.""" source = TextSource( file_or_pattern, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder(), skip_header_lines=skip_header_lines) range_tracker = source.get_range_tracker(None, None) return list(source.read(range_tracker)) def test_read_skip_header_single(self): file_name, expected_data = write_data(TextSourceTest.DEFAULT_NUM_RECORDS) assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS skip_header_lines = 1 expected_data = self._remove_lines( expected_data, [TextSourceTest.DEFAULT_NUM_RECORDS], skip_header_lines) read_data = self._read_skip_header_lines(file_name, skip_header_lines) self.assertEqual(len(expected_data), len(read_data)) self.assertCountEqual(expected_data, read_data) def test_read_skip_header_pattern(self): line_counts = [ TextSourceTest.DEFAULT_NUM_RECORDS * 5, TextSourceTest.DEFAULT_NUM_RECORDS * 3, TextSourceTest.DEFAULT_NUM_RECORDS * 12, TextSourceTest.DEFAULT_NUM_RECORDS * 8, TextSourceTest.DEFAULT_NUM_RECORDS * 8, TextSourceTest.DEFAULT_NUM_RECORDS * 4 ] skip_header_lines = 2 pattern, data = write_pattern(line_counts) expected_data = self._remove_lines(data, line_counts, skip_header_lines) read_data = self._read_skip_header_lines(pattern, skip_header_lines) self.assertEqual(len(expected_data), len(read_data)) self.assertCountEqual(expected_data, read_data) def test_read_skip_header_pattern_insufficient_lines(self): line_counts = [ 5, 3, # Fewer lines in file than we want to skip 12, 8, 8, 4 ] skip_header_lines = 4 pattern, data = write_pattern(line_counts) data = self._remove_lines(data, line_counts, skip_header_lines) read_data = self._read_skip_header_lines(pattern, skip_header_lines) self.assertEqual(len(data), len(read_data)) self.assertCountEqual(data, read_data) def test_read_gzip_with_skip_lines(self): _, lines = write_data(15) with TempDir() as tempdir: file_name = tempdir.create_temp_file() with gzip.GzipFile(file_name, 'wb') as f: f.write('\n'.join(lines).encode('utf-8')) with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText( file_name, 0, CompressionTypes.GZIP, True, coders.StrUtf8Coder(), skip_header_lines=2) assert_that(pcoll, equal_to(lines[2:])) def test_read_after_splitting_skip_header(self): file_name, expected_data = write_data(100) assert len(expected_data) == 100 source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder(), skip_header_lines=2) splits = list(source.split(desired_bundle_size=33)) reference_source_info = (source, None, None) sources_info = ([(split.source, split.start_position, split.stop_position) for split in splits]) self.assertGreater(len(sources_info), 1) reference_lines = source_test_utils.read_from_source(*reference_source_info) split_lines = [] for source_info in sources_info: split_lines.extend(source_test_utils.read_from_source(*source_info)) self.assertEqual(expected_data[2:], reference_lines) self.assertEqual(reference_lines, split_lines) def test_custom_delimiter_read_from_text(self): file_name, expected_data = write_data( 5, eol=EOL.CUSTOM_DELIMITER, custom_delimiter=b'@#') assert len(expected_data) == 5 with TestPipeline() as pipeline: pcoll = pipeline | 'Read' >> ReadFromText(file_name, delimiter=b'@#') assert_that(pcoll, equal_to(expected_data)) def test_custom_delimiter_read_all_single_file(self): file_name, expected_data = write_data( 5, eol=EOL.CUSTOM_DELIMITER, custom_delimiter=b'@#') assert len(expected_data) == 5 with TestPipeline() as pipeline: pcoll = pipeline | 'Create' >> Create( [file_name]) | 'ReadAll' >> ReadAllFromText(delimiter=b'@#') assert_that(pcoll, equal_to(expected_data)) def test_invalid_delimiters_are_rejected(self): file_name, _ = write_data(1) for delimiter in (b'', '', '\r\n', 'a', 1): with self.assertRaises( ValueError, msg='Delimiter must be a non-empty bytes sequence.'): _ = TextSource( file_pattern=file_name, min_bundle_size=0, buffer_size=6, compression_type=CompressionTypes.UNCOMPRESSED, strip_trailing_newlines=True, coder=coders.StrUtf8Coder(), delimiter=delimiter, ) def test_non_self_overlapping_delimiter_is_accepted(self): file_name, _ = write_data(1) for delimiter in (b'\n', b'\r\n', b'*', b'abc', b'cabdab', b'abcabd'): _ = TextSource( file_pattern=file_name, min_bundle_size=0, buffer_size=6, compression_type=CompressionTypes.UNCOMPRESSED, strip_trailing_newlines=True, coder=coders.StrUtf8Coder(), delimiter=delimiter, ) def test_self_overlapping_delimiter_is_rejected(self): file_name, _ = write_data(1) for delimiter in (b'||', b'***', b'aba', b'abcab'): with self.assertRaises(ValueError, msg='Delimiter must not self-overlap.'): _ = TextSource( file_pattern=file_name, min_bundle_size=0, buffer_size=6, compression_type=CompressionTypes.UNCOMPRESSED, strip_trailing_newlines=True, coder=coders.StrUtf8Coder(), delimiter=delimiter, ) def test_read_with_customer_delimiter(self): delimiters = [ b'\n', b'\r\n', b'*|', b'*', b'*=-', ] for delimiter in delimiters: file_name, expected_data = write_data( 10, eol=EOL.CUSTOM_DELIMITER, custom_delimiter=delimiter) assert len(expected_data) == 10 source = TextSource( file_pattern=file_name, min_bundle_size=0, compression_type=CompressionTypes.UNCOMPRESSED, strip_trailing_newlines=True, coder=coders.StrUtf8Coder(), delimiter=delimiter) range_tracker = source.get_range_tracker(None, None) read_data = list(source.read(range_tracker)) self.assertEqual(read_data, expected_data) def test_read_with_custom_delimiter_around_split_point(self): for delimiter in (b'\n', b'\r\n', b'@#', b'abc'): file_name, expected_data = write_data( 20, eol=EOL.CUSTOM_DELIMITER, custom_delimiter=delimiter) assert len(expected_data) == 20 for desired_bundle_size in (4, 5, 6, 7): source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder(), delimiter=delimiter) splits = list(source.split(desired_bundle_size=desired_bundle_size)) reference_source_info = (source, None, None) sources_info = ([ (split.source, split.start_position, split.stop_position) for split in splits ]) source_test_utils.assert_sources_equal_reference_source( reference_source_info, sources_info) def test_read_with_customer_delimiter_truncated(self): """ Corner case: delimiter truncated at the end of the file Use delimiter with length = 3, buffer_size = 6 and line_value with length = 4 to split the delimiter """ delimiter = b'@$*' file_name, expected_data = write_data( 10, eol=EOL.CUSTOM_DELIMITER, line_value=b'a' * 4, custom_delimiter=delimiter) assert len(expected_data) == 10 source = TextSource( file_pattern=file_name, min_bundle_size=0, buffer_size=6, compression_type=CompressionTypes.UNCOMPRESSED, strip_trailing_newlines=True, coder=coders.StrUtf8Coder(), delimiter=delimiter, ) range_tracker = source.get_range_tracker(None, None) read_data = list(source.read(range_tracker)) self.assertEqual(read_data, expected_data) def test_read_with_customer_delimiter_over_buffer_size(self): """ Corner case: delimiter is on border of size of buffer """ file_name, expected_data = write_data(3, eol=EOL.CRLF, line_value=b'\rline') assert len(expected_data) == 3 self._run_read_test( file_name, expected_data, buffer_size=7, delimiter=b'\r\n') def test_read_with_customer_delimiter_truncated_and_not_equal(self): """ Corner case: delimiter truncated at the end of the file and only part of delimiter equal end of buffer Use delimiter with length = 3, buffer_size = 6 and line_value with length = 4 to split the delimiter """ write_delimiter = b'@$' read_delimiter = b'@$*' file_name, expected_data = write_data( 10, eol=EOL.CUSTOM_DELIMITER, line_value=b'a' * 4, custom_delimiter=write_delimiter) # In this case check, that the line won't be splitted write_delimiter_encode = write_delimiter.decode('utf-8') expected_data_str = [ write_delimiter_encode.join(expected_data) + write_delimiter_encode ] source = TextSource( file_pattern=file_name, min_bundle_size=0, buffer_size=6, compression_type=CompressionTypes.UNCOMPRESSED, strip_trailing_newlines=True, coder=coders.StrUtf8Coder(), delimiter=read_delimiter, ) range_tracker = source.get_range_tracker(None, None) read_data = list(source.read(range_tracker)) self.assertEqual(read_data, expected_data_str) def test_read_crlf_split_by_buffer(self): file_name, expected_data = write_data(3, eol=EOL.CRLF) assert len(expected_data) == 3 self._run_read_test(file_name, expected_data, buffer_size=6) def test_read_escaped_lf(self): file_name, expected_data = write_data( self.DEFAULT_NUM_RECORDS, eol=EOL.LF, line_value=b'li\\\nne') assert len(expected_data) == self.DEFAULT_NUM_RECORDS self._run_read_test(file_name, expected_data, escapechar=b'\\') def test_read_escaped_crlf(self): file_name, expected_data = write_data( TextSource.DEFAULT_READ_BUFFER_SIZE, eol=EOL.CRLF, line_value=b'li\\\r\\\nne') assert len(expected_data) == TextSource.DEFAULT_READ_BUFFER_SIZE self._run_read_test(file_name, expected_data, escapechar=b'\\') def test_read_escaped_cr_before_not_escaped_lf(self): file_name, expected_data_temp = write_data( self.DEFAULT_NUM_RECORDS, eol=EOL.CRLF, line_value=b'li\\\r\nne') expected_data = [] for line in expected_data_temp: expected_data += line.split("\n") assert len(expected_data) == self.DEFAULT_NUM_RECORDS * 2 self._run_read_test(file_name, expected_data, escapechar=b'\\') def test_read_escaped_custom_delimiter_crlf(self): file_name, expected_data = write_data( self.DEFAULT_NUM_RECORDS, eol=EOL.CRLF, line_value=b'li\\\r\nne') assert len(expected_data) == self.DEFAULT_NUM_RECORDS self._run_read_test( file_name, expected_data, delimiter=b'\r\n', escapechar=b'\\') def test_read_escaped_custom_delimiter(self): file_name, expected_data = write_data( TextSource.DEFAULT_READ_BUFFER_SIZE, eol=EOL.CUSTOM_DELIMITER, custom_delimiter=b'*|', line_value=b'li\\*|ne') assert len(expected_data) == TextSource.DEFAULT_READ_BUFFER_SIZE self._run_read_test( file_name, expected_data, delimiter=b'*|', escapechar=b'\\') def test_read_escaped_lf_at_buffer_edge(self): file_name, expected_data = write_data(3, eol=EOL.LF, line_value=b'line\\\n') assert len(expected_data) == 3 self._run_read_test( file_name, expected_data, buffer_size=5, escapechar=b'\\') def test_read_escaped_crlf_split_by_buffer(self): file_name, expected_data = write_data( 3, eol=EOL.CRLF, line_value=b'line\\\r\n') assert len(expected_data) == 3 self._run_read_test( file_name, expected_data, buffer_size=6, delimiter=b'\r\n', escapechar=b'\\') def test_read_escaped_lf_after_splitting(self): file_name, expected_data = write_data(3, line_value=b'line\\\n') assert len(expected_data) == 3 source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder(), escapechar=b'\\') splits = list(source.split(desired_bundle_size=6)) reference_source_info = (source, None, None) sources_info = ([(split.source, split.start_position, split.stop_position) for split in splits]) source_test_utils.assert_sources_equal_reference_source( reference_source_info, sources_info) def test_read_escaped_lf_after_splitting_many(self): file_name, expected_data = write_data( 3, line_value=b'\\\\\\\\\\\n') # 5 escapes assert len(expected_data) == 3 source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder(), escapechar=b'\\') splits = list(source.split(desired_bundle_size=6)) reference_source_info = (source, None, None) sources_info = ([(split.source, split.start_position, split.stop_position) for split in splits]) source_test_utils.assert_sources_equal_reference_source( reference_source_info, sources_info) def test_read_escaped_escapechar_after_splitting(self): file_name, expected_data = write_data(3, line_value=b'line\\\\*|') assert len(expected_data) == 3 source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder(), delimiter=b'*|', escapechar=b'\\') splits = list(source.split(desired_bundle_size=8)) reference_source_info = (source, None, None) sources_info = ([(split.source, split.start_position, split.stop_position) for split in splits]) source_test_utils.assert_sources_equal_reference_source( reference_source_info, sources_info) def test_read_escaped_escapechar_after_splitting_many(self): file_name, expected_data = write_data( 3, line_value=b'\\\\\\\\\\\\*|') # 6 escapes assert len(expected_data) == 3 source = TextSource( file_name, 0, CompressionTypes.UNCOMPRESSED, True, coders.StrUtf8Coder(), delimiter=b'*|', escapechar=b'\\') splits = list(source.split(desired_bundle_size=8)) reference_source_info = (source, None, None) sources_info = ([(split.source, split.start_position, split.stop_position) for split in splits]) source_test_utils.assert_sources_equal_reference_source( reference_source_info, sources_info) class TextSinkTest(unittest.TestCase): def setUp(self): super().setUp() self.lines = [b'Line %d' % d for d in range(100)] self.tempdir = tempfile.mkdtemp() self.path = self._create_temp_file() def tearDown(self): if os.path.exists(self.tempdir): shutil.rmtree(self.tempdir) def _create_temp_file(self, name='', suffix=''): if not name: name = tempfile.template file_name = tempfile.NamedTemporaryFile( delete=True, prefix=name, dir=self.tempdir, suffix=suffix).name return file_name def _write_lines(self, sink, lines): f = sink.open(self.path) for line in lines: sink.write_record(f, line) sink.close(f) def test_write_text_file(self): sink = TextSink(self.path) self._write_lines(sink, self.lines) with open(self.path, 'rb') as f: self.assertEqual(f.read().splitlines(), self.lines) def test_write_text_file_empty(self): sink = TextSink(self.path) self._write_lines(sink, []) with open(self.path, 'rb') as f: self.assertEqual(f.read().splitlines(), []) def test_write_bzip2_file(self): sink = TextSink(self.path, compression_type=CompressionTypes.BZIP2) self._write_lines(sink, self.lines) with bz2.BZ2File(self.path, 'rb') as f: self.assertEqual(f.read().splitlines(), self.lines) def test_write_bzip2_file_auto(self): self.path = self._create_temp_file(suffix='.bz2') sink = TextSink(self.path) self._write_lines(sink, self.lines) with bz2.BZ2File(self.path, 'rb') as f: self.assertEqual(f.read().splitlines(), self.lines) def test_write_gzip_file(self): sink = TextSink(self.path, compression_type=CompressionTypes.GZIP) self._write_lines(sink, self.lines) with gzip.GzipFile(self.path, 'rb') as f: self.assertEqual(f.read().splitlines(), self.lines) def test_write_gzip_file_auto(self): self.path = self._create_temp_file(suffix='.gz') sink = TextSink(self.path) self._write_lines(sink, self.lines) with gzip.GzipFile(self.path, 'rb') as f: self.assertEqual(f.read().splitlines(), self.lines) def test_write_gzip_file_empty(self): sink = TextSink(self.path, compression_type=CompressionTypes.GZIP) self._write_lines(sink, []) with gzip.GzipFile(self.path, 'rb') as f: self.assertEqual(f.read().splitlines(), []) def test_write_deflate_file(self): sink = TextSink(self.path, compression_type=CompressionTypes.DEFLATE) self._write_lines(sink, self.lines) with open(self.path, 'rb') as f: self.assertEqual(zlib.decompress(f.read()).splitlines(), self.lines) def test_write_deflate_file_auto(self): self.path = self._create_temp_file(suffix='.deflate') sink = TextSink(self.path) self._write_lines(sink, self.lines) with open(self.path, 'rb') as f: self.assertEqual(zlib.decompress(f.read()).splitlines(), self.lines) def test_write_deflate_file_empty(self): sink = TextSink(self.path, compression_type=CompressionTypes.DEFLATE) self._write_lines(sink, []) with open(self.path, 'rb') as f: self.assertEqual(zlib.decompress(f.read()).splitlines(), []) def test_write_text_file_with_header(self): header = b'header1\nheader2' sink = TextSink(self.path, header=header) self._write_lines(sink, self.lines) with open(self.path, 'rb') as f: self.assertEqual(f.read().splitlines(), header.splitlines() + self.lines) def test_write_text_file_with_footer(self): footer = b'footer1\nfooter2' sink = TextSink(self.path, footer=footer) self._write_lines(sink, self.lines) with open(self.path, 'rb') as f: self.assertEqual(f.read().splitlines(), self.lines + footer.splitlines()) def test_write_text_file_empty_with_header(self): header = b'header1\nheader2' sink = TextSink(self.path, header=header) self._write_lines(sink, []) with open(self.path, 'rb') as f: self.assertEqual(f.read().splitlines(), header.splitlines()) def test_write_pipeline(self): with TestPipeline() as pipeline: pcoll = pipeline | beam.core.Create(self.lines) pcoll | 'Write' >> WriteToText(self.path) # pylint: disable=expression-not-assigned read_result = [] for file_name in glob.glob(self.path + '*'): with open(file_name, 'rb') as f: read_result.extend(f.read().splitlines()) self.assertEqual(sorted(read_result), sorted(self.lines)) def test_write_pipeline_non_globalwindow_input(self): with TestPipeline() as p: _ = ( p | beam.core.Create(self.lines) | beam.WindowInto(beam.transforms.window.FixedWindows(1)) | 'Write' >> WriteToText(self.path)) read_result = [] for file_name in glob.glob(self.path + '*'): with open(file_name, 'rb') as f: read_result.extend(f.read().splitlines()) self.assertEqual(sorted(read_result), sorted(self.lines)) def test_write_pipeline_auto_compression(self): with TestPipeline() as pipeline: pcoll = pipeline | beam.core.Create(self.lines) pcoll | 'Write' >> WriteToText(self.path, file_name_suffix='.gz') # pylint: disable=expression-not-assigned read_result = [] for file_name in glob.glob(self.path + '*'): with gzip.GzipFile(file_name, 'rb') as f: read_result.extend(f.read().splitlines()) self.assertEqual(sorted(read_result), sorted(self.lines)) def test_write_pipeline_auto_compression_unsharded(self): with TestPipeline() as pipeline: pcoll = pipeline | 'Create' >> beam.core.Create(self.lines) pcoll | 'Write' >> WriteToText( # pylint: disable=expression-not-assigned self.path + '.gz', shard_name_template='') read_result = [] for file_name in glob.glob(self.path + '*'): with gzip.GzipFile(file_name, 'rb') as f: read_result.extend(f.read().splitlines()) self.assertEqual(sorted(read_result), sorted(self.lines)) def test_write_pipeline_header(self): with TestPipeline() as pipeline: pcoll = pipeline | 'Create' >> beam.core.Create(self.lines) header_text = 'foo' pcoll | 'Write' >> WriteToText( # pylint: disable=expression-not-assigned self.path + '.gz', shard_name_template='', header=header_text) read_result = [] for file_name in glob.glob(self.path + '*'): with gzip.GzipFile(file_name, 'rb') as f: read_result.extend(f.read().splitlines()) # header_text is automatically encoded in WriteToText self.assertEqual(read_result[0], header_text.encode('utf-8')) self.assertEqual(sorted(read_result[1:]), sorted(self.lines)) def test_write_pipeline_footer(self): with TestPipeline() as pipeline: footer_text = 'footer' pcoll = pipeline | beam.core.Create(self.lines) pcoll | 'Write' >> WriteToText( # pylint: disable=expression-not-assigned self.path, footer=footer_text) read_result = [] for file_name in glob.glob(self.path + '*'): with open(file_name, 'rb') as f: read_result.extend(f.read().splitlines()) self.assertEqual(sorted(read_result[:-1]), sorted(self.lines)) self.assertEqual(read_result[-1], footer_text.encode('utf-8')) def test_write_empty(self): with TestPipeline() as p: # pylint: disable=expression-not-assigned p | beam.core.Create([]) | WriteToText(self.path) outputs = glob.glob(self.path + '*') self.assertEqual(len(outputs), 1) with open(outputs[0], 'rb') as f: self.assertEqual(list(f.read().splitlines()), []) def test_write_empty_skipped(self): with TestPipeline() as p: # pylint: disable=expression-not-assigned p | beam.core.Create([]) | WriteToText(self.path, skip_if_empty=True) outputs = list(glob.glob(self.path + '*')) self.assertEqual(outputs, []) def test_write_max_records_per_shard(self): records_per_shard = 13 lines = [str(i).encode('utf-8') for i in range(100)] with TestPipeline() as p: # pylint: disable=expression-not-assigned p | beam.core.Create(lines) | WriteToText( self.path, max_records_per_shard=records_per_shard) read_result = [] for file_name in glob.glob(self.path + '*'): with open(file_name, 'rb') as f: shard_lines = list(f.read().splitlines()) self.assertLessEqual(len(shard_lines), records_per_shard) read_result.extend(shard_lines) self.assertEqual(sorted(read_result), sorted(lines)) def test_write_max_bytes_per_shard(self): bytes_per_shard = 300 max_len = 100 lines = [b'x' * i for i in range(max_len)] header = b'a' * 20 footer = b'b' * 30 with TestPipeline() as p: # pylint: disable=expression-not-assigned p | beam.core.Create(lines) | WriteToText( self.path, header=header, footer=footer, max_bytes_per_shard=bytes_per_shard) read_result = [] for file_name in glob.glob(self.path + '*'): with open(file_name, 'rb') as f: contents = f.read() self.assertLessEqual( len(contents), bytes_per_shard + max_len + len(footer) + 2) shard_lines = list(contents.splitlines()) self.assertEqual(shard_lines[0], header) self.assertEqual(shard_lines[-1], footer) read_result.extend(shard_lines[1:-1]) self.assertEqual(sorted(read_result), sorted(lines)) class CsvTest(unittest.TestCase): def test_csv_read_write(self): records = [beam.Row(a='str', b=ix) for ix in range(3)] with tempfile.TemporaryDirectory() as dest: with TestPipeline() as p: # pylint: disable=expression-not-assigned p | beam.Create(records) | beam.io.WriteToCsv(os.path.join(dest, 'out')) with TestPipeline() as p: pcoll = ( p | beam.io.ReadFromCsv(os.path.join(dest, 'out*')) | beam.Map(lambda t: beam.Row(**dict(zip(type(t)._fields, t))))) assert_that(pcoll, equal_to(records)) def test_csv_read_with_filename(self): records = [beam.Row(a='str', b=ix) for ix in range(3)] with tempfile.TemporaryDirectory() as dest: file_path = os.path.join(dest, 'out.csv') with TestPipeline() as p: # pylint: disable=expression-not-assigned p | beam.Create(records) | beam.io.WriteToCsv(file_path) with TestPipeline() as p: pcoll = ( p | beam.io.ReadFromCsv( file_path + '*', filename_column='source_filename') | beam.Map(lambda t: beam.Row(**dict(zip(type(t)._fields, t))))) # Get the sharded file name files = glob.glob(file_path + '*') self.assertEqual(len(files), 1) sharded_file_path = files[0] expected = [ beam.Row(a=r.a, b=r.b, source_filename=sharded_file_path) for r in records ] assert_that(pcoll, equal_to(expected)) def test_non_utf8_csv_read_write(self): content = b"\xe0,\xe1,\xe2\n0,1,2\n1,2,3\n" with tempfile.TemporaryDirectory() as dest: input_fn = os.path.join(dest, 'input.csv') with open(input_fn, 'wb') as f: f.write(content) with TestPipeline() as p: r1 = ( p | 'Read' >> beam.io.ReadFromCsv(input_fn, encoding="latin1") | 'ToDict' >> beam.Map(lambda x: x._asdict())) assert_that( r1, equal_to([{ "\u00e0": 0, "\u00e1": 1, "\u00e2": 2 }, { "\u00e0": 1, "\u00e1": 2, "\u00e2": 3 }])) with TestPipeline() as p: _ = ( p | 'Read' >> beam.io.ReadFromCsv(input_fn, encoding="latin1") | 'Write' >> beam.io.WriteToCsv( os.path.join(dest, 'out'), encoding="latin1")) with TestPipeline() as p: r2 = ( p | 'Read' >> beam.io.ReadFromCsv( os.path.join(dest, 'out*'), encoding="latin1") | 'ToDict' >> beam.Map(lambda x: x._asdict())) assert_that( r2, equal_to([{ "\u00e0": 0, "\u00e1": 1, "\u00e2": 2 }, { "\u00e0": 1, "\u00e1": 2, "\u00e2": 3 }])) class JsonTest(unittest.TestCase): def test_json_read_write(self): records = [beam.Row(a='str', b=ix) for ix in range(3)] with tempfile.TemporaryDirectory() as dest: with TestPipeline() as p: # pylint: disable=expression-not-assigned p | beam.Create(records) | beam.io.WriteToJson( os.path.join(dest, 'out')) with TestPipeline() as p: pcoll = ( p | beam.io.ReadFromJson(os.path.join(dest, 'out*')) | beam.Map(lambda t: beam.Row(**dict(zip(type(t)._fields, t))))) assert_that(pcoll, equal_to(records)) def test_numeric_strings_preserved(self): records = [ beam.Row( as_string=str(ix), as_float_string=str(float(ix)), as_int=ix, as_float=float(ix)) for ix in range(3) ] with tempfile.TemporaryDirectory() as dest: with TestPipeline() as p: # pylint: disable=expression-not-assigned p | beam.Create(records) | beam.io.WriteToJson( os.path.join(dest, 'out')) with TestPipeline() as p: pcoll = ( p | beam.io.ReadFromJson(os.path.join(dest, 'out*')) | beam.Map(lambda t: beam.Row(**dict(zip(type(t)._fields, t))))) assert_that(pcoll, equal_to(records)) # This test should be redundant as Python equality does not equate # numeric values with their string representations, but this is much # more explicit about what we're asserting here. def check_types(element): for a, b in zip(element, records[0]): assert type(a) == type(b), (a, b, type(a), type(b)) _ = pcoll | beam.Map(check_types) class GenerateEvent(beam.PTransform): @staticmethod def sample_data(): return GenerateEvent() def expand(self, input): elemlist = [{'age': 10}, {'age': 20}, {'age': 30}] elem = elemlist return ( input | TestStream().add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 1, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 2, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 3, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 4, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 5, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 5, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 6, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 7, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 8, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 9, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 10, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 10, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 11, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 12, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 13, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 14, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 15, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 15, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 16, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 17, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 18, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 19, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 20, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 20, 0, tzinfo=pytz.UTC).timestamp()).advance_watermark_to( datetime( 2021, 3, 1, 0, 0, 25, 0, tzinfo=pytz.UTC). timestamp()).advance_watermark_to_infinity()) class WriteStreamingTest(unittest.TestCase): def setUp(self): super().setUp() self.tempdir = tempfile.mkdtemp() def tearDown(self): if os.path.exists(self.tempdir): shutil.rmtree(self.tempdir) def test_write_streaming_2_shards_default_shard_name_template( self, num_shards=2): with TestPipeline() as p: output = (p | GenerateEvent.sample_data()) #TextIO output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( file_path_prefix=self.tempdir + "/ouput_WriteToText", file_name_suffix=".txt", num_shards=num_shards, triggering_frequency=60) _ = output2 | 'LogElements after WriteToText' >> LogElements( prefix='after WriteToText ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToText-[1614556800.0, 1614556805.0)-00000-of-00002.txt # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P[\d\.]+), ' r'(?P[\d\.]+|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.txt$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToText*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) self.assertEqual( len(file_names), num_shards, "expected %d files, but got: %d" % (num_shards, len(file_names))) def test_write_streaming_2_shards_default_shard_name_template_windowed_pcoll( self, num_shards=2): with TestPipeline() as p: output = ( p | GenerateEvent.sample_data() | 'User windowing' >> beam.transforms.core.WindowInto( beam.transforms.window.FixedWindows(10), trigger=beam.transforms.trigger.AfterWatermark(), accumulation_mode=beam.transforms.trigger.AccumulationMode. DISCARDING, allowed_lateness=beam.utils.timestamp.Duration(seconds=0))) #TextIO output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( file_path_prefix=self.tempdir + "/ouput_WriteToText", file_name_suffix=".txt", num_shards=num_shards, ) _ = output2 | 'LogElements after WriteToText' >> LogElements( prefix='after WriteToText ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToText-[1614556800.0, 1614556805.0)-00000-of-00002.txt # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P[\d\.]+), ' r'(?P[\d\.]+|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.txt$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToText*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) self.assertEqual( len(file_names), num_shards * 3, #25s of data covered by 3 10s windows "expected %d files, but got: %d" % (num_shards * 3, len(file_names))) def test_write_streaming_undef_shards_default_shard_name_template_windowed_pcoll( # pylint: disable=line-too-long self): with TestPipeline() as p: output = ( p | GenerateEvent.sample_data() | 'User windowing' >> beam.transforms.core.WindowInto( beam.transforms.window.FixedWindows(10), trigger=beam.transforms.trigger.AfterWatermark(), accumulation_mode=beam.transforms.trigger.AccumulationMode. DISCARDING, allowed_lateness=beam.utils.timestamp.Duration(seconds=0))) #TextIO output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( file_path_prefix=self.tempdir + "/ouput_WriteToText", file_name_suffix=".txt", num_shards=0, ) _ = output2 | 'LogElements after WriteToText' >> LogElements( prefix='after WriteToText ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToText-[1614556800.0, 1614556805.0)-00000-of-00002.txt # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P[\d\.]+), ' r'(?P[\d\.]+|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.txt$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToText*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) self.assertGreaterEqual( len(file_names), 1 * 3, #25s of data covered by 3 10s windows "expected %d files, but got: %d" % (1 * 3, len(file_names))) def test_write_streaming_undef_shards_default_shard_name_template_windowed_pcoll_and_trig_freq( # pylint: disable=line-too-long self): with TestPipeline() as p: output = ( p | GenerateEvent.sample_data() | 'User windowing' >> beam.transforms.core.WindowInto( beam.transforms.window.FixedWindows(60), trigger=beam.transforms.trigger.AfterWatermark(), accumulation_mode=beam.transforms.trigger.AccumulationMode. DISCARDING, allowed_lateness=beam.utils.timestamp.Duration(seconds=0))) #TextIO output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( file_path_prefix=self.tempdir + "/ouput_WriteToText", file_name_suffix=".txt", num_shards=0, triggering_frequency=10, ) _ = output2 | 'LogElements after WriteToText' >> LogElements( prefix='after WriteToText ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToText-[1614556800.0, 1614556805.0)-00000-of-00002.txt # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P[\d\.]+), ' r'(?P[\d\.]+|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.txt$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToText*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) self.assertGreaterEqual( len(file_names), 1 * 3, #25s of data covered by 3 10s windows "expected %d files, but got: %d" % (1 * 3, len(file_names))) def test_write_streaming_undef_shards_default_shard_name_template_global_window_pcoll( # pylint: disable=line-too-long self): with TestPipeline() as p: output = (p | GenerateEvent.sample_data()) #TextIO output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( file_path_prefix=self.tempdir + "/ouput_WriteToText", file_name_suffix=".txt", num_shards=0, #0 means undef nb of shards, same as omitted/default triggering_frequency=60, ) _ = output2 | 'LogElements after WriteToText' >> LogElements( prefix='after WriteToText ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToText-[1614556800.0, 1614556805.0)-00000-of-00002.txt # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P[\d\.]+), ' r'(?P[\d\.]+|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.txt$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToText*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) self.assertGreaterEqual( len(file_names), 1, #25s of data covered by 60s windows "expected %d files, but got: %d" % (1, len(file_names))) def test_write_streaming_2_shards_custom_shard_name_template( self, num_shards=2, shard_name_template='-V-SSSSS-of-NNNNN'): with TestPipeline() as p: output = (p | GenerateEvent.sample_data()) #TextIO output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( file_path_prefix=self.tempdir + "/ouput_WriteToText", file_name_suffix=".txt", shard_name_template=shard_name_template, num_shards=num_shards, triggering_frequency=60, ) _ = output2 | 'LogElements after WriteToText' >> LogElements( prefix='after WriteToText ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToText-[2021-03-01T00-00-00, 2021-03-01T00-01-00)- # 00000-of-00002.txt # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}), ' r'(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.txt$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToText*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) self.assertEqual( len(file_names), num_shards, "expected %d files, but got: %d" % (num_shards, len(file_names))) def test_write_streaming_2_shards_custom_shard_name_template_5s_window( self, num_shards=2, shard_name_template='-V-SSSSS-of-NNNNN', triggering_frequency=5): with TestPipeline() as p: output = (p | GenerateEvent.sample_data()) #TextIO output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( file_path_prefix=self.tempdir + "/ouput_WriteToText", file_name_suffix=".txt", shard_name_template=shard_name_template, num_shards=num_shards, triggering_frequency=triggering_frequency, ) _ = output2 | 'LogElements after WriteToText' >> LogElements( prefix='after WriteToText ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToText-[2021-03-01T00-00-00, 2021-03-01T00-01-00)- # 00000-of-00002.txt # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}), ' r'(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.txt$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToText*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) # for 5s window size, the input should be processed by 5 windows with # 2 shards per window self.assertEqual( len(file_names), 10, "expected %d files, but got: %d" % (num_shards, len(file_names))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/tfrecordio.py0000644000175100017510000003367115206113032021010 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """TFRecord sources and sinks.""" # pytype: skip-file import codecs import logging import struct from functools import partial from apache_beam import coders from apache_beam.io import filebasedsink from apache_beam.io.filebasedsource import FileBasedSource from apache_beam.io.filebasedsource import ReadAllFiles from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.iobase import Read from apache_beam.io.iobase import Write from apache_beam.transforms import PTransform try: import crcmod except ImportError: crcmod = None __all__ = ['ReadFromTFRecord', 'ReadAllFromTFRecord', 'WriteToTFRecord'] _LOGGER = logging.getLogger(__name__) def _default_crc32c_fn(value): """Calculates crc32c of a bytes object using either snappy or google-crc32c or crcmod.""" if not _default_crc32c_fn.fn: try: import snappy # pylint: disable=import-error # Support multiple versions of python-snappy: # https://github.com/andrix/python-snappy/pull/53 if getattr(snappy, '_crc32c', None): _default_crc32c_fn.fn = snappy._crc32c # pylint: disable=protected-access elif getattr(snappy, '_snappy', None): _default_crc32c_fn.fn = snappy._snappy._crc32c # pylint: disable=protected-access except ImportError: pass if not _default_crc32c_fn.fn: try: import google_crc32c # pylint: disable=import-error if getattr(google_crc32c, 'value', None): _default_crc32c_fn.fn = google_crc32c.value # pylint: disable=protected-access except ImportError: pass if not _default_crc32c_fn.fn: if crcmod is None: raise RuntimeError( 'Could not find python-snappy, google-crc32c, or crcmod. To allow ' 'execution to succeed, make sure that one of these packages is ' 'installed or pip install apache-beam[tfrecord]') _LOGGER.warning( 'Couldn\'t find python-snappy or google-crc32c so the ' 'implementation of _TFRecordUtil._masked_crc32c is not as fast ' 'as it could be.') _default_crc32c_fn.fn = crcmod.predefined.mkPredefinedCrcFun('crc-32c') return _default_crc32c_fn.fn(value) _default_crc32c_fn.fn = None # type: ignore class _TFRecordUtil(object): """Provides basic TFRecord encoding/decoding with consistency checks. For detailed TFRecord format description see: https://www.tensorflow.org/versions/r1.11/api_guides/python/python_io#TFRecords_Format_Details Note that masks and length are represented in LittleEndian order. """ @classmethod def _masked_crc32c(cls, value, crc32c_fn=_default_crc32c_fn): """Compute a masked crc32c checksum for a value. Args: value: A bytes object for which we compute the crc. crc32c_fn: A function that can compute a crc32c. This is a performance hook that also helps with testing. Callers are not expected to make use of it directly. Returns: Masked crc32c checksum. """ crc = crc32c_fn(value) return (((crc >> 15) | (crc << 17)) + 0xa282ead8) & 0xffffffff @staticmethod def encoded_num_bytes(record): """Return the number of bytes consumed by a record in its encoded form.""" # 16 = 8 (Length) + 4 (crc of length) + 4 (crc of data) return len(record) + 16 @classmethod def write_record(cls, file_handle, value): """Encode a value as a TFRecord. Args: file_handle: The file to write to. value: A bytes object representing content of the record. """ encoded_length = struct.pack(b'> self._read_all_files class ReadFromTFRecord(PTransform): """Transform for reading TFRecord sources.""" def __init__( self, file_pattern, coder=coders.BytesCoder(), compression_type=CompressionTypes.AUTO, validate=True): """Initialize a ReadFromTFRecord transform. Args: file_pattern: A file glob pattern to read TFRecords from. coder: Coder used to decode each record. compression_type: Used to handle compressed input files. Default value is CompressionTypes.AUTO, in which case the file_path's extension will be used to detect the compression. validate: Boolean flag to verify that the files exist during the pipeline creation time. Returns: A ReadFromTFRecord transform object. """ super().__init__() self._source = _TFRecordSource( file_pattern, coder, compression_type, validate) def expand(self, pvalue): return pvalue.pipeline | Read(self._source) class _TFRecordSink(filebasedsink.FileBasedSink): """Sink for writing TFRecords files. For detailed TFRecord format description see: https://www.tensorflow.org/versions/r1.11/api_guides/python/python_io#TFRecords_Format_Details """ def __init__( self, file_path_prefix, coder, file_name_suffix, num_shards, shard_name_template, compression_type, triggering_frequency=60): """Initialize a TFRecordSink. See WriteToTFRecord for details.""" super().__init__( file_path_prefix=file_path_prefix, coder=coder, file_name_suffix=file_name_suffix, num_shards=num_shards, shard_name_template=shard_name_template, mime_type='application/octet-stream', compression_type=compression_type, triggering_frequency=triggering_frequency) def write_encoded_record(self, file_handle, value): _TFRecordUtil.write_record(file_handle, value) class WriteToTFRecord(PTransform): """Transform for writing to TFRecord sinks.""" def __init__( self, file_path_prefix, coder=coders.BytesCoder(), file_name_suffix='', num_shards=0, shard_name_template=None, compression_type=CompressionTypes.AUTO, triggering_frequency=None): """Initialize WriteToTFRecord transform. Args: file_path_prefix: The file path to write to. The files written will begin with this prefix, followed by a shard identifier (see num_shards), and end in a common extension, if given by file_name_suffix. coder: Coder used to encode each record. file_name_suffix: Suffix for the files written. num_shards: The number of files (shards) used for output. If not set, the default value will be used. In streaming if not set, the service will write a file per bundle. shard_name_template: A template string containing placeholders for the shard number and shard count. Currently only ``''``, ``'-SSSSS-of-NNNNN'``, ``'-W-SSSSS-of-NNNNN'`` and ``'-V-SSSSS-of-NNNNN'`` are patterns accepted by the service. When constructing a filename for a particular shard number, the upper-case letters ``S`` and ``N`` are replaced with the ``0``-padded shard number and shard count respectively. This argument can be ``''`` in which case it behaves as if num_shards was set to 1 and only one file will be generated. The default pattern used is ``'-SSSSS-of-NNNNN'`` for bounded PCollections and for ``'-W-SSSSS-of-NNNNN'`` unbounded PCollections. W is used for windowed shard naming and is replaced with ``[window.start, window.end)`` V is used for windowed shard naming and is replaced with ``[window.start.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S"), window.end.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S")`` compression_type: Used to handle compressed output files. Typical value is CompressionTypes.AUTO, in which case the file_path's extension will be used to detect the compression. triggering_frequency: (int) Every triggering_frequency duration, a window will be triggered and all bundles in the window will be written. If set it overrides user windowing. Mandatory for GlobalWindow. Returns: A WriteToTFRecord transform object. """ super().__init__() self._sink = _TFRecordSink( file_path_prefix, coder, file_name_suffix, num_shards, shard_name_template, compression_type, triggering_frequency) def expand(self, pcoll): if (not pcoll.is_bounded and self._sink.shard_name_template == filebasedsink.DEFAULT_SHARD_NAME_TEMPLATE): self._sink.shard_name_template = ( filebasedsink.DEFAULT_WINDOW_SHARD_NAME_TEMPLATE) self._sink.shard_name_format = self._sink._template_to_format( self._sink.shard_name_template) self._sink.shard_name_glob_format = self._sink._template_to_glob_format( self._sink.shard_name_template) return pcoll | Write(self._sink) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/tfrecordio_test.py0000644000175100017510000007502015206113032022041 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import binascii import glob import gzip import io import json import logging import os import pickle import random import re import shutil import tempfile import unittest import zlib from datetime import datetime import pytz import apache_beam as beam from apache_beam import Create from apache_beam import coders from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.tfrecordio import ReadAllFromTFRecord from apache_beam.io.tfrecordio import ReadFromTFRecord from apache_beam.io.tfrecordio import WriteToTFRecord from apache_beam.io.tfrecordio import _TFRecordSink from apache_beam.io.tfrecordio import _TFRecordUtil from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.test_utils import TempDir from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.util import LogElements try: import tensorflow.compat.v1 as tf # pylint: disable=import-error except ImportError: try: import tensorflow as tf # pylint: disable=import-error except ImportError: tf = None # pylint: disable=invalid-name logging.warning('Tensorflow is not installed, so skipping some tests.') try: import crcmod except ImportError: crcmod = None # Created by running following code in python: # >>> import tensorflow as tf # >>> import base64 # >>> writer = tf.python_io.TFRecordWriter('/tmp/python_foo.tfrecord') # >>> writer.write(b'foo') # >>> writer.close() # >>> with open('/tmp/python_foo.tfrecord', 'rb') as f: # ... data = base64.b64encode(f.read()) # ... print(data) FOO_RECORD_BASE64 = b'AwAAAAAAAACwmUkOZm9vYYq+/g==' # Same as above but containing two records [b'foo', b'bar'] FOO_BAR_RECORD_BASE64 = b'AwAAAAAAAACwmUkOZm9vYYq+/gMAAAAAAAAAsJlJDmJhckYA5cg=' def _write_file(path, base64_records): record = binascii.a2b_base64(base64_records) with open(path, 'wb') as f: f.write(record) def _write_file_deflate(path, base64_records): record = binascii.a2b_base64(base64_records) with open(path, 'wb') as f: f.write(zlib.compress(record)) def _write_file_gzip(path, base64_records): record = binascii.a2b_base64(base64_records) with gzip.GzipFile(path, 'wb') as f: f.write(record) class TestTFRecordUtil(unittest.TestCase): def setUp(self): self.record = binascii.a2b_base64(FOO_RECORD_BASE64) def _as_file_handle(self, contents): result = io.BytesIO() result.write(contents) result.seek(0) return result def _increment_value_at_index(self, value, index): l = list(value) l[index] = l[index] + 1 return bytes(l) def _test_error(self, record, error_text): with self.assertRaisesRegex(ValueError, re.escape(error_text)): _TFRecordUtil.read_record(self._as_file_handle(record)) def test_masked_crc32c(self): self.assertEqual(0xfd7fffa, _TFRecordUtil._masked_crc32c(b'\x00' * 32)) self.assertEqual(0xf909b029, _TFRecordUtil._masked_crc32c(b'\xff' * 32)) self.assertEqual(0xfebe8a61, _TFRecordUtil._masked_crc32c(b'foo')) self.assertEqual( 0xe4999b0, _TFRecordUtil._masked_crc32c(b'\x03\x00\x00\x00\x00\x00\x00\x00')) @unittest.skipIf(crcmod is None, 'crcmod not installed.') def test_masked_crc32c_crcmod(self): crc32c_fn = crcmod.predefined.mkPredefinedCrcFun('crc-32c') self.assertEqual( 0xfd7fffa, _TFRecordUtil._masked_crc32c(b'\x00' * 32, crc32c_fn=crc32c_fn)) self.assertEqual( 0xf909b029, _TFRecordUtil._masked_crc32c(b'\xff' * 32, crc32c_fn=crc32c_fn)) self.assertEqual( 0xfebe8a61, _TFRecordUtil._masked_crc32c(b'foo', crc32c_fn=crc32c_fn)) self.assertEqual( 0xe4999b0, _TFRecordUtil._masked_crc32c( b'\x03\x00\x00\x00\x00\x00\x00\x00', crc32c_fn=crc32c_fn)) def test_write_record(self): file_handle = io.BytesIO() _TFRecordUtil.write_record(file_handle, b'foo') self.assertEqual(self.record, file_handle.getvalue()) def test_read_record(self): actual = _TFRecordUtil.read_record(self._as_file_handle(self.record)) self.assertEqual(b'foo', actual) def test_read_record_invalid_record(self): self._test_error(b'bar', 'Not a valid TFRecord. Fewer than 12 bytes') def test_read_record_invalid_length_mask(self): record = self._increment_value_at_index(self.record, 9) self._test_error(record, 'Mismatch of length mask') def test_read_record_invalid_data_mask(self): record = self._increment_value_at_index(self.record, 16) self._test_error(record, 'Mismatch of data mask') def test_compatibility_read_write(self): for record in [b'', b'blah', b'another blah']: file_handle = io.BytesIO() _TFRecordUtil.write_record(file_handle, record) file_handle.seek(0) actual = _TFRecordUtil.read_record(file_handle) self.assertEqual(record, actual) class TestTFRecordSink(unittest.TestCase): def _write_lines(self, sink, path, lines): f = sink.open(path) for l in lines: sink.write_record(f, l) sink.close(f) def test_write_record_single(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result') record = binascii.a2b_base64(FOO_RECORD_BASE64) sink = _TFRecordSink( path, coder=coders.BytesCoder(), file_name_suffix='', num_shards=0, shard_name_template=None, compression_type=CompressionTypes.UNCOMPRESSED) self._write_lines(sink, path, [b'foo']) with open(path, 'rb') as f: self.assertEqual(f.read(), record) def test_write_record_multiple(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result') record = binascii.a2b_base64(FOO_BAR_RECORD_BASE64) sink = _TFRecordSink( path, coder=coders.BytesCoder(), file_name_suffix='', num_shards=0, shard_name_template=None, compression_type=CompressionTypes.UNCOMPRESSED) self._write_lines(sink, path, [b'foo', b'bar']) with open(path, 'rb') as f: self.assertEqual(f.read(), record) @unittest.skipIf(tf is None, 'tensorflow not installed.') class TestWriteToTFRecord(TestTFRecordSink): def test_write_record_gzip(self): with TempDir() as temp_dir: file_path_prefix = temp_dir.create_temp_file('result') with TestPipeline() as p: input_data = [b'foo', b'bar'] _ = p | beam.Create(input_data) | WriteToTFRecord( file_path_prefix, compression_type=CompressionTypes.GZIP) actual = [] file_name = glob.glob(file_path_prefix + '-*')[0] for r in tf.python_io.tf_record_iterator( file_name, options=tf.python_io.TFRecordOptions( tf.python_io.TFRecordCompressionType.GZIP)): actual.append(r) self.assertEqual(sorted(actual), sorted(input_data)) def test_write_record_auto(self): with TempDir() as temp_dir: file_path_prefix = temp_dir.create_temp_file('result') with TestPipeline() as p: input_data = [b'foo', b'bar'] _ = p | beam.Create(input_data) | WriteToTFRecord( file_path_prefix, file_name_suffix='.gz') actual = [] file_name = glob.glob(file_path_prefix + '-*.gz')[0] for r in tf.python_io.tf_record_iterator( file_name, options=tf.python_io.TFRecordOptions( tf.python_io.TFRecordCompressionType.GZIP)): actual.append(r) self.assertEqual(sorted(actual), sorted(input_data)) class TestReadFromTFRecord(unittest.TestCase): def test_process_single(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result') _write_file(path, FOO_RECORD_BASE64) with TestPipeline() as p: result = ( p | ReadFromTFRecord( path, coder=coders.BytesCoder(), compression_type=CompressionTypes.AUTO, validate=True)) assert_that(result, equal_to([b'foo'])) def test_process_multiple(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result') _write_file(path, FOO_BAR_RECORD_BASE64) with TestPipeline() as p: result = ( p | ReadFromTFRecord( path, coder=coders.BytesCoder(), compression_type=CompressionTypes.AUTO, validate=True)) assert_that(result, equal_to([b'foo', b'bar'])) def test_process_deflate(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result') _write_file_deflate(path, FOO_BAR_RECORD_BASE64) with TestPipeline() as p: result = ( p | ReadFromTFRecord( path, coder=coders.BytesCoder(), compression_type=CompressionTypes.DEFLATE, validate=True)) assert_that(result, equal_to([b'foo', b'bar'])) def test_process_gzip_with_coder(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result') _write_file_gzip(path, FOO_BAR_RECORD_BASE64) with TestPipeline() as p: result = ( p | ReadFromTFRecord( path, coder=coders.BytesCoder(), compression_type=CompressionTypes.GZIP, validate=True)) assert_that(result, equal_to([b'foo', b'bar'])) def test_process_gzip_without_coder(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result') _write_file_gzip(path, FOO_BAR_RECORD_BASE64) with TestPipeline() as p: result = ( p | ReadFromTFRecord(path, compression_type=CompressionTypes.GZIP)) assert_that(result, equal_to([b'foo', b'bar'])) def test_process_auto(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result.gz') _write_file_gzip(path, FOO_BAR_RECORD_BASE64) with TestPipeline() as p: result = ( p | ReadFromTFRecord( path, coder=coders.BytesCoder(), compression_type=CompressionTypes.AUTO, validate=True)) assert_that(result, equal_to([b'foo', b'bar'])) def test_process_gzip_auto(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result.gz') _write_file_gzip(path, FOO_BAR_RECORD_BASE64) with TestPipeline() as p: result = ( p | ReadFromTFRecord(path, compression_type=CompressionTypes.AUTO)) assert_that(result, equal_to([b'foo', b'bar'])) class TestReadAllFromTFRecord(unittest.TestCase): def _write_glob(self, temp_dir, suffix, include_empty=False): for _ in range(3): path = temp_dir.create_temp_file(suffix) _write_file(path, FOO_BAR_RECORD_BASE64) if include_empty: path = temp_dir.create_temp_file(suffix) _write_file(path, '') def test_process_single(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result') _write_file(path, FOO_RECORD_BASE64) with TestPipeline() as p: result = ( p | Create([path]) | ReadAllFromTFRecord( coder=coders.BytesCoder(), compression_type=CompressionTypes.AUTO)) assert_that(result, equal_to([b'foo'])) def test_process_multiple(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result') _write_file(path, FOO_BAR_RECORD_BASE64) with TestPipeline() as p: result = ( p | Create([path]) | ReadAllFromTFRecord( coder=coders.BytesCoder(), compression_type=CompressionTypes.AUTO)) assert_that(result, equal_to([b'foo', b'bar'])) def test_process_with_filename(self): with TempDir() as temp_dir: num_files = 3 files = [] for i in range(num_files): path = temp_dir.create_temp_file('result%s' % i) _write_file(path, FOO_BAR_RECORD_BASE64) files.append(path) content = [b'foo', b'bar'] expected = [(file, line) for file in files for line in content] pattern = temp_dir.get_path() + os.path.sep + '*' with TestPipeline() as p: result = ( p | Create([pattern]) | ReadAllFromTFRecord( coder=coders.BytesCoder(), compression_type=CompressionTypes.AUTO, with_filename=True)) assert_that(result, equal_to(expected)) def test_process_glob(self): with TempDir() as temp_dir: self._write_glob(temp_dir, 'result') glob = temp_dir.get_path() + os.path.sep + '*result' with TestPipeline() as p: result = ( p | Create([glob]) | ReadAllFromTFRecord( coder=coders.BytesCoder(), compression_type=CompressionTypes.AUTO)) assert_that(result, equal_to([b'foo', b'bar'] * 3)) def test_process_glob_with_empty_file(self): with TempDir() as temp_dir: self._write_glob(temp_dir, 'result', include_empty=True) glob = temp_dir.get_path() + os.path.sep + '*result' with TestPipeline() as p: result = ( p | Create([glob]) | ReadAllFromTFRecord( coder=coders.BytesCoder(), compression_type=CompressionTypes.AUTO)) assert_that(result, equal_to([b'foo', b'bar'] * 3)) def test_process_multiple_globs(self): with TempDir() as temp_dir: globs = [] for i in range(3): suffix = 'result' + str(i) self._write_glob(temp_dir, suffix) globs.append(temp_dir.get_path() + os.path.sep + '*' + suffix) with TestPipeline() as p: result = ( p | Create(globs) | ReadAllFromTFRecord( coder=coders.BytesCoder(), compression_type=CompressionTypes.AUTO)) assert_that(result, equal_to([b'foo', b'bar'] * 9)) def test_process_deflate(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result') _write_file_deflate(path, FOO_BAR_RECORD_BASE64) with TestPipeline() as p: result = ( p | Create([path]) | ReadAllFromTFRecord( coder=coders.BytesCoder(), compression_type=CompressionTypes.DEFLATE)) assert_that(result, equal_to([b'foo', b'bar'])) def test_process_gzip(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result') _write_file_gzip(path, FOO_BAR_RECORD_BASE64) with TestPipeline() as p: result = ( p | Create([path]) | ReadAllFromTFRecord( coder=coders.BytesCoder(), compression_type=CompressionTypes.GZIP)) assert_that(result, equal_to([b'foo', b'bar'])) def test_process_auto(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result.gz') _write_file_gzip(path, FOO_BAR_RECORD_BASE64) with TestPipeline() as p: result = ( p | Create([path]) | ReadAllFromTFRecord( coder=coders.BytesCoder(), compression_type=CompressionTypes.AUTO)) assert_that(result, equal_to([b'foo', b'bar'])) class TestEnd2EndWriteAndRead(unittest.TestCase): def create_inputs(self): input_array = [[random.random() - 0.5 for _ in range(15)] for _ in range(12)] memfile = io.BytesIO() pickle.dump(input_array, memfile) return memfile.getvalue() def test_end2end(self): with TempDir() as temp_dir: file_path_prefix = temp_dir.create_temp_file('result') # Generate a TFRecord file. with TestPipeline() as p: expected_data = [self.create_inputs() for _ in range(0, 10)] _ = p | beam.Create(expected_data) | WriteToTFRecord(file_path_prefix) # Read the file back and compare. with TestPipeline() as p: actual_data = p | ReadFromTFRecord(file_path_prefix + '-*') assert_that(actual_data, equal_to(expected_data)) def test_end2end_auto_compression(self): with TempDir() as temp_dir: file_path_prefix = temp_dir.create_temp_file('result') # Generate a TFRecord file. with TestPipeline() as p: expected_data = [self.create_inputs() for _ in range(0, 10)] _ = p | beam.Create(expected_data) | WriteToTFRecord( file_path_prefix, file_name_suffix='.gz') # Read the file back and compare. with TestPipeline() as p: actual_data = p | ReadFromTFRecord(file_path_prefix + '-*') assert_that(actual_data, equal_to(expected_data)) def test_end2end_auto_compression_unsharded(self): with TempDir() as temp_dir: file_path_prefix = temp_dir.create_temp_file('result') # Generate a TFRecord file. with TestPipeline() as p: expected_data = [self.create_inputs() for _ in range(0, 10)] _ = p | beam.Create(expected_data) | WriteToTFRecord( file_path_prefix + '.gz', shard_name_template='') # Read the file back and compare. with TestPipeline() as p: actual_data = p | ReadFromTFRecord(file_path_prefix + '.gz') assert_that(actual_data, equal_to(expected_data)) @unittest.skipIf(tf is None, 'tensorflow not installed.') def test_end2end_example_proto(self): with TempDir() as temp_dir: file_path_prefix = temp_dir.create_temp_file('result') example = tf.train.Example() example.features.feature['int'].int64_list.value.extend(list(range(3))) example.features.feature['bytes'].bytes_list.value.extend( [b'foo', b'bar']) with TestPipeline() as p: _ = p | beam.Create([example]) | WriteToTFRecord( file_path_prefix, coder=beam.coders.ProtoCoder(example.__class__)) # Read the file back and compare. with TestPipeline() as p: actual_data = ( p | ReadFromTFRecord( file_path_prefix + '-*', coder=beam.coders.ProtoCoder(example.__class__))) assert_that(actual_data, equal_to([example])) def test_end2end_read_write_read(self): with TempDir() as temp_dir: path = temp_dir.create_temp_file('result') with TestPipeline() as p: # Initial read to validate the pipeline doesn't fail before the file is # created. _ = p | ReadFromTFRecord(path + '-*', validate=False) expected_data = [self.create_inputs() for _ in range(0, 10)] _ = p | beam.Create(expected_data) | WriteToTFRecord( path, file_name_suffix='.gz') # Read the file back and compare. with TestPipeline() as p: actual_data = p | ReadFromTFRecord(path + '-*', validate=True) assert_that(actual_data, equal_to(expected_data)) class GenerateEvent(beam.PTransform): @staticmethod def sample_data(): return GenerateEvent() def expand(self, input): elemlist = [{'age': 10}, {'age': 20}, {'age': 30}] elem = elemlist return ( input | TestStream().add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 1, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 2, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 3, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 4, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 5, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 5, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 6, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 7, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 8, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 9, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 10, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 10, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 11, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 12, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 13, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 14, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 15, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 15, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 16, 0, tzinfo=pytz.UTC).timestamp()). add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 17, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 18, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 19, 0, tzinfo=pytz.UTC).timestamp()). advance_watermark_to( datetime(2021, 3, 1, 0, 0, 20, 0, tzinfo=pytz.UTC).timestamp()).add_elements( elements=elem, event_timestamp=datetime( 2021, 3, 1, 0, 0, 20, 0, tzinfo=pytz.UTC).timestamp()).advance_watermark_to( datetime( 2021, 3, 1, 0, 0, 25, 0, tzinfo=pytz.UTC). timestamp()).advance_watermark_to_infinity()) class WriteStreamingTest(unittest.TestCase): def setUp(self): super().setUp() self.tempdir = tempfile.mkdtemp() def tearDown(self): if os.path.exists(self.tempdir): shutil.rmtree(self.tempdir) def test_write_streaming_2_shards_default_shard_name_template( self, num_shards=2): with TestPipeline() as p: output = ( p | GenerateEvent.sample_data() | 'User windowing' >> beam.transforms.core.WindowInto( beam.transforms.window.FixedWindows(60), trigger=beam.transforms.trigger.AfterWatermark(), accumulation_mode=beam.transforms.trigger.AccumulationMode. DISCARDING, allowed_lateness=beam.utils.timestamp.Duration(seconds=0)) | "encode" >> beam.Map(lambda s: json.dumps(s).encode('utf-8'))) #TFrecordIO output2 = output | 'WriteToTFRecord' >> beam.io.WriteToTFRecord( file_path_prefix=self.tempdir + "/ouput_WriteToTFRecord", file_name_suffix=".tfrecord", num_shards=num_shards, ) _ = output2 | 'LogElements after WriteToTFRecord' >> LogElements( prefix='after WriteToTFRecord ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToTFRecord-[1614556800.0, 1614556805.0)-00000-of-00002.tfrecord # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P[\d\.]+), ' r'(?P[\d\.]+|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.tfrecord$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToTFRecord*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) self.assertEqual( len(file_names), num_shards, "expected %d files, but got: %d" % (num_shards, len(file_names))) def test_write_streaming_2_shards_custom_shard_name_template( self, num_shards=2, shard_name_template='-V-SSSSS-of-NNNNN'): with TestPipeline() as p: output = ( p | GenerateEvent.sample_data() | "encode" >> beam.Map(lambda s: json.dumps(s).encode('utf-8'))) #TFrecordIO output2 = output | 'WriteToTFRecord' >> beam.io.WriteToTFRecord( file_path_prefix=self.tempdir + "/ouput_WriteToTFRecord", file_name_suffix=".tfrecord", shard_name_template=shard_name_template, num_shards=num_shards, triggering_frequency=60, ) _ = output2 | 'LogElements after WriteToTFRecord' >> LogElements( prefix='after WriteToTFRecord ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToTFRecord-[2021-03-01T00-00-00, 2021-03-01T00-01-00)- # 00000-of-00002.tfrecord # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}), ' r'(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.tfrecord$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToTFRecord*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) self.assertEqual( len(file_names), num_shards, "expected %d files, but got: %d" % (num_shards, len(file_names))) def test_write_streaming_2_shards_custom_shard_name_template_5s_window( self, num_shards=2, shard_name_template='-V-SSSSS-of-NNNNN', triggering_frequency=5): with TestPipeline() as p: output = ( p | GenerateEvent.sample_data() | "encode" >> beam.Map(lambda s: json.dumps(s).encode('utf-8'))) #TFrecordIO output2 = output | 'WriteToTFRecord' >> beam.io.WriteToTFRecord( file_path_prefix=self.tempdir + "/ouput_WriteToTFRecord", file_name_suffix=".tfrecord", shard_name_template=shard_name_template, num_shards=num_shards, triggering_frequency=triggering_frequency, ) _ = output2 | 'LogElements after WriteToTFRecord' >> LogElements( prefix='after WriteToTFRecord ', with_window=True, level=logging.INFO) # Regex to match the expected windowed file pattern # Example: # ouput_WriteToTFRecord-[2021-03-01T00-00-00, 2021-03-01T00-01-00)- # 00000-of-00002.tfrecord # It captures: window_interval, shard_num, total_shards pattern_string = ( r'.*-\[(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}), ' r'(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}|Infinity)\)-' r'(?P\d{5})-of-(?P\d{5})\.tfrecord$') pattern = re.compile(pattern_string) file_names = [] for file_name in glob.glob(self.tempdir + '/ouput_WriteToTFRecord*'): match = pattern.match(file_name) self.assertIsNotNone( match, f"File name {file_name} did not match expected pattern.") if match: file_names.append(file_name) print("Found files matching expected pattern:", file_names) # for 5s window size, the input should be processed by 5 windows with # 2 shards per window self.assertEqual( len(file_names), 10, "expected %d files, but got: %d" % (num_shards, len(file_names))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/utils.py0000644000175100017510000000445315206113032020004 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utils for the io library. * CountingSource: Subclass of iobase.BoundedSource. Used on transforms.ptransform_test.test_read_metrics. """ # pytype: skip-file from apache_beam.io import iobase from apache_beam.io.range_trackers import OffsetRangeTracker from apache_beam.metrics import Metrics class CountingSource(iobase.BoundedSource): def __init__(self, count): self.records_read = Metrics.counter(self.__class__, 'recordsRead') self._count = count def estimate_size(self): return self._count def get_range_tracker(self, start_position, stop_position): if start_position is None: start_position = 0 if stop_position is None: stop_position = self._count return OffsetRangeTracker(start_position, stop_position) def read(self, range_tracker): for i in range(range_tracker.start_position(), range_tracker.stop_position()): if not range_tracker.try_claim(i): return self.records_read.inc() yield i def split(self, desired_bundle_size, start_position=None, stop_position=None): if start_position is None: start_position = 0 if stop_position is None: stop_position = self._count bundle_start = start_position while bundle_start < stop_position: bundle_stop = min(stop_position, bundle_start + desired_bundle_size) yield iobase.SourceBundle( weight=(bundle_stop - bundle_start), source=self, start_position=bundle_start, stop_position=bundle_stop) bundle_start = bundle_stop ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/utils_test.py0000644000175100017510000000441015206113032021034 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import mock from apache_beam.io import OffsetRangeTracker from apache_beam.io import source_test_utils from apache_beam.io.utils import CountingSource class CountingSourceTest(unittest.TestCase): def setUp(self): self.source = CountingSource(10) def test_estimate_size(self): self.assertEqual(10, self.source.estimate_size()) @mock.patch('apache_beam.io.utils.OffsetRangeTracker') def test_get_range_tracker(self, mock_tracker): _ = self.source.get_range_tracker(None, None) mock_tracker.assert_called_with(0, 10) _ = self.source.get_range_tracker(3, 7) mock_tracker.assert_called_with(3, 7) def test_read(self): tracker = OffsetRangeTracker(3, 6) res = list(self.source.read(tracker)) self.assertEqual([3, 4, 5], res) def test_split(self): for size in [1, 3, 10]: splits = list(self.source.split(desired_bundle_size=size)) reference_info = (self.source, None, None) sources_info = ([ (split.source, split.start_position, split.stop_position) for split in splits ]) source_test_utils.assert_sources_equal_reference_source( reference_info, sources_info) def test_dynamic_work_rebalancing(self): splits = list(self.source.split(desired_bundle_size=20)) assert len(splits) == 1 source_test_utils.assert_split_at_fraction_exhaustive( splits[0].source, splits[0].start_position, splits[0].stop_position) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/watermark_estimators.py0000644000175100017510000001222115206113032023103 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A collection of WatermarkEstimator implementations that SplittableDoFns can use.""" # pytype: skip-file from apache_beam.io.iobase import WatermarkEstimator from apache_beam.transforms.core import WatermarkEstimatorProvider from apache_beam.utils.timestamp import Timestamp class MonotonicWatermarkEstimator(WatermarkEstimator): """A WatermarkEstimator which assumes that timestamps of all ouput records are increasing monotonically. """ def __init__(self, timestamp): """For a new pair, the initial value is None. When resuming processing, the initial timestamp will be the last reported watermark. """ self._watermark = timestamp self._last_observed_timestamp = timestamp def observe_timestamp(self, timestamp): self._last_observed_timestamp = timestamp def current_watermark(self): if self._last_observed_timestamp is not None \ and self._last_observed_timestamp >= self._watermark: self._watermark = self._last_observed_timestamp return self._watermark def get_estimator_state(self): return self._watermark @staticmethod def default_provider(): """Provide a default WatermarkEstimatorProvider for MonotonicWatermarkEstimator. """ class DefaultMonotonicWatermarkEstimator(WatermarkEstimatorProvider): def initial_estimator_state(self, element, restriction): return None def create_watermark_estimator(self, estimator_state): return MonotonicWatermarkEstimator(estimator_state) return DefaultMonotonicWatermarkEstimator() class WalltimeWatermarkEstimator(WatermarkEstimator): """A WatermarkEstimator which uses processing time as the estimated watermark. """ def __init__(self, timestamp=None): self._timestamp = timestamp or Timestamp.now() def observe_timestamp(self, timestamp): pass def current_watermark(self): self._timestamp = max(self._timestamp, Timestamp.now()) return self._timestamp def get_estimator_state(self): return self._timestamp @staticmethod def default_provider(): """Provide a default WatermarkEstimatorProvider for WalltimeWatermarkEstimator. """ class DefaultWalltimeWatermarkEstimator(WatermarkEstimatorProvider): def initial_estimator_state(self, element, restriction): return None def create_watermark_estimator(self, estimator_state): return WalltimeWatermarkEstimator(estimator_state) return DefaultWalltimeWatermarkEstimator() class ManualWatermarkEstimator(WatermarkEstimator): """A WatermarkEstimator which is controlled manually from within a DoFn. The DoFn must invoke set_watermark to advance the watermark. """ def __init__(self, watermark): self._watermark = watermark def observe_timestamp(self, timestamp): pass def current_watermark(self): return self._watermark def get_estimator_state(self): return self._watermark def set_watermark(self, timestamp): # pylint: disable=line-too-long """Sets a timestamp before or at the timestamps of all future elements produced by the associated DoFn. This can be approximate. If records are output that violate this guarantee, they will be considered late, which will affect how they will be processed. See https://beam.apache.org/documentation/programming-guide/#watermarks-and-late-data for more information on late data and how to handle it. However, this value should be as late as possible. Downstream windows may not be able to close until this watermark passes their end. """ if not isinstance(timestamp, Timestamp): raise ValueError('set_watermark expects a Timestamp as input') if self._watermark and self._watermark > timestamp: raise ValueError( 'Watermark must be monotonically increasing.' 'Provided watermark %s is less than ' 'current watermark %s', timestamp, self._watermark) self._watermark = timestamp @staticmethod def default_provider(): """Provide a default WatermarkEstimatorProvider for WalltimeWatermarkEstimator. """ class DefaultManualWatermarkEstimatorProvider(WatermarkEstimatorProvider): def initial_estimator_state(self, element, restriction): return None def create_watermark_estimator(self, estimator_state): return ManualWatermarkEstimator(estimator_state) return DefaultManualWatermarkEstimatorProvider() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/io/watermark_estimators_test.py0000644000175100017510000001164615206113032024154 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for built-in WatermarkEstimators""" # pytype: skip-file import unittest import mock from apache_beam.io.iobase import WatermarkEstimator from apache_beam.io.watermark_estimators import ManualWatermarkEstimator from apache_beam.io.watermark_estimators import MonotonicWatermarkEstimator from apache_beam.io.watermark_estimators import WalltimeWatermarkEstimator from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import Timestamp class MonotonicWatermarkEstimatorTest(unittest.TestCase): def test_initialize_from_state(self): timestamp = Timestamp(10) watermark_estimator = MonotonicWatermarkEstimator(timestamp) self.assertIsInstance(watermark_estimator, WatermarkEstimator) self.assertEqual(watermark_estimator.get_estimator_state(), timestamp) def test_observe_timestamp(self): watermark_estimator = MonotonicWatermarkEstimator(Timestamp(10)) watermark_estimator.observe_timestamp(Timestamp(15)) self.assertEqual(watermark_estimator.current_watermark(), Timestamp(15)) watermark_estimator.observe_timestamp(Timestamp(20)) self.assertEqual(watermark_estimator.current_watermark(), Timestamp(20)) watermark_estimator.observe_timestamp(Timestamp(20)) self.assertEqual(watermark_estimator.current_watermark(), Timestamp(20)) watermark_estimator.observe_timestamp(Timestamp(10)) self.assertEqual(watermark_estimator.current_watermark(), Timestamp(20)) def test_get_estimator_state(self): watermark_estimator = MonotonicWatermarkEstimator(Timestamp(10)) self.assertEqual(watermark_estimator.get_estimator_state(), Timestamp(10)) watermark_estimator.observe_timestamp(Timestamp(15)) # State only progresses when we request a new watermark self.assertEqual(watermark_estimator.get_estimator_state(), Timestamp(10)) self.assertEqual(watermark_estimator.current_watermark(), Timestamp(15)) self.assertEqual(watermark_estimator.get_estimator_state(), Timestamp(15)) class WalltimeWatermarkEstimatorTest(unittest.TestCase): @mock.patch('apache_beam.utils.timestamp.Timestamp.now') def test_initialization(self, mock_timestamp): now_time = Timestamp.now() - Duration(10) mock_timestamp.side_effect = lambda: now_time watermark_estimator = WalltimeWatermarkEstimator() self.assertIsInstance(watermark_estimator, WatermarkEstimator) self.assertEqual(watermark_estimator.get_estimator_state(), now_time) def test_observe_timestamp(self): now_time = Timestamp.now() + Duration(10) watermark_estimator = WalltimeWatermarkEstimator(now_time) watermark_estimator.observe_timestamp(Timestamp(10)) watermark_estimator.observe_timestamp(Timestamp(10)) self.assertEqual(watermark_estimator.current_watermark(), now_time) def test_advance_watermark_with_incorrect_sys_clock(self): initial_timestamp = Timestamp.now() + Duration(100) watermark_estimator = WalltimeWatermarkEstimator(initial_timestamp) self.assertEqual(watermark_estimator.current_watermark(), initial_timestamp) self.assertEqual( watermark_estimator.get_estimator_state(), initial_timestamp) class ManualWatermarkEstimatorTest(unittest.TestCase): def test_initialization(self): watermark_estimator = ManualWatermarkEstimator(None) self.assertIsNone(watermark_estimator.get_estimator_state()) self.assertIsNone(watermark_estimator.current_watermark()) watermark_estimator = ManualWatermarkEstimator(Timestamp(10)) self.assertEqual(watermark_estimator.get_estimator_state(), Timestamp(10)) def test_set_watermark(self): watermark_estimator = ManualWatermarkEstimator(None) self.assertIsNone(watermark_estimator.current_watermark()) watermark_estimator.observe_timestamp(Timestamp(10)) self.assertIsNone(watermark_estimator.current_watermark()) watermark_estimator.set_watermark(Timestamp(20)) self.assertEqual(watermark_estimator.current_watermark(), Timestamp(20)) watermark_estimator.set_watermark(Timestamp(30)) self.assertEqual(watermark_estimator.current_watermark(), Timestamp(30)) with self.assertRaises(ValueError): watermark_estimator.set_watermark(Timestamp(25)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.4034102 apache_beam-2.74.0/apache_beam/metrics/0000755000175100017510000000000015206113157017333 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/metrics/__init__.py0000644000175100017510000000156415206113032021442 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from apache_beam.metrics.metric import Metrics from apache_beam.metrics.metric import MetricsFilter ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/metrics/cells.pxd0000644000175100017510000000423315206113032021144 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # cimport cython cimport libc.stdint from cpython.datetime cimport datetime cdef class MetricCell(object): cdef object _lock cpdef bint update(self, value) except -1 cdef datetime _start_time cdef class CounterCell(MetricCell): cdef readonly libc.stdint.int64_t value @cython.locals(ivalue=libc.stdint.int64_t) cpdef bint update(self, value) except -1 # Not using AbstractMetricCell so that data can be typed. cdef class DistributionCell(MetricCell): cdef readonly DistributionData data @cython.locals(ivalue=libc.stdint.int64_t) cdef inline bint _update(self, value) except -1 cdef class AbstractMetricCell(MetricCell): cdef readonly object data_class cdef public object data cdef bint _update_locked(self, value) except -1 cdef class GaugeCell(AbstractMetricCell): pass cdef class StringSetCell(AbstractMetricCell): pass cdef class BoundedTrieCell(AbstractMetricCell): pass cdef class DistributionData(object): cdef readonly libc.stdint.int64_t sum cdef readonly libc.stdint.int64_t count cdef readonly libc.stdint.int64_t min cdef readonly libc.stdint.int64_t max cdef class _BoundedTrieNode(object): cdef readonly libc.stdint.int64_t _size cdef readonly dict _children cdef readonly bint _truncated cdef class BoundedTrieData(object): cdef readonly libc.stdint.int64_t _bound cdef readonly object _singleton cdef readonly _BoundedTrieNode _root ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/metrics/cells.py0000644000175100017510000007073515206113032021013 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ This file contains metric cell classes. A metric cell is used to accumulate in-memory changes to a metric. It represents a specific metric in a single context. """ # pytype: skip-file # ruff: noqa: UP006 import copy import logging import threading import time from datetime import datetime from datetime import timezone from typing import Iterable from typing import Optional from typing import Set from apache_beam.portability.api import metrics_pb2 from apache_beam.utils.histogram import Histogram try: import cython except ImportError: class fake_cython: compiled = False globals()['cython'] = fake_cython __all__ = [ 'MetricCell', 'MetricCellFactory', 'DistributionResult', 'GaugeResult' ] _LOGGER = logging.getLogger(__name__) class MetricCell(object): """For internal use only; no backwards-compatibility guarantees. Accumulates in-memory changes to a metric. A MetricCell represents a specific metric in a single context and bundle. All subclasses must be thread safe, as these are used in the pipeline runners, and may be subject to parallel/concurrent updates. Cells should only be used directly within a runner. """ def __init__(self): self._lock = threading.Lock() self._start_time = None def update(self, value): raise NotImplementedError def get_cumulative(self): raise NotImplementedError def to_runner_api_monitoring_info(self, name, transform_id): if not self._start_time: self._start_time = datetime.now(timezone.utc) mi = self.to_runner_api_monitoring_info_impl(name, transform_id) mi.start_time.FromDatetime(self._start_time) return mi def to_runner_api_monitoring_info_impl(self, name, transform_id): raise NotImplementedError def reset(self): # type: () -> None raise NotImplementedError def __reduce__(self): raise NotImplementedError class MetricCellFactory(object): def __call__(self): # type: () -> MetricCell raise NotImplementedError class CounterCell(MetricCell): """For internal use only; no backwards-compatibility guarantees. Tracks the current value and delta of a counter metric. Each cell tracks the state of an integer metric independently per context per bundle. Therefore, each metric has a different cell in each bundle, cells are aggregated by the runner. This class is thread safe. """ def __init__(self, *args): super().__init__(*args) self.value = 0 def reset(self): # type: () -> None self.value = 0 def combine(self, other): # type: (CounterCell) -> CounterCell result = CounterCell() result.inc(self.value + other.value) return result def inc(self, n=1): self.update(n) def dec(self, n=1): self.update(-n) def update(self, value): # type: (int) -> None if cython.compiled: ivalue = value # Since We hold the GIL, no need for another lock. # And because the C threads won't preempt and interleave # each other. # Assuming there is no code trying to access the counters # directly by circumventing the GIL. self.value += ivalue else: with self._lock: self.value += value def get_cumulative(self): # type: () -> int with self._lock: return self.value def to_runner_api_monitoring_info_impl(self, name, transform_id): from apache_beam.metrics import monitoring_infos if not name.urn: # User counter case. return monitoring_infos.int64_user_counter( name.namespace, name.name, self.get_cumulative(), ptransform=transform_id) else: # Arbitrary URN case. return monitoring_infos.int64_counter( name.urn, self.get_cumulative(), labels=name.labels) class DistributionCell(MetricCell): """For internal use only; no backwards-compatibility guarantees. Tracks the current value and delta for a distribution metric. Each cell tracks the state of a metric independently per context per bundle. Therefore, each metric has a different cell in each bundle, that is later aggregated. This class is thread safe. """ def __init__(self, *args): super().__init__(*args) self.data = DistributionData.identity_element() def reset(self): # type: () -> None self.data = DistributionData.identity_element() def combine(self, other): # type: (DistributionCell) -> DistributionCell result = DistributionCell() result.data = self.data.combine(other.data) return result def update(self, value): if cython.compiled: # We will hold the GIL throughout the entire _update. self._update(value) else: with self._lock: self._update(value) def _update(self, value): if cython.compiled: ivalue = value else: ivalue = int(value) self.data.count = self.data.count + 1 self.data.sum = self.data.sum + ivalue if ivalue < self.data.min: self.data.min = ivalue if ivalue > self.data.max: self.data.max = ivalue def get_cumulative(self): # type: () -> DistributionData with self._lock: return self.data.get_cumulative() def to_runner_api_monitoring_info_impl(self, name, transform_id): from apache_beam.metrics import monitoring_infos return monitoring_infos.int64_user_distribution( name.namespace, name.name, self.get_cumulative(), ptransform=transform_id) class AbstractMetricCell(MetricCell): """For internal use only; no backwards-compatibility guarantees. Tracks the current value and delta for a metric with a data class. This class is thread safe. """ def __init__(self, data_class): super().__init__() self.data_class = data_class self.data = self.data_class.identity_element() def reset(self): self.data = self.data_class.identity_element() def combine(self, other: 'AbstractMetricCell') -> 'AbstractMetricCell': result = type(self)() # type: ignore[call-arg] result.data = self.data.combine(other.data) return result def set(self, value): with self._lock: self._update_locked(value) def update(self, value): with self._lock: self._update_locked(value) def _update_locked(self, value): raise NotImplementedError(type(self)) def get_cumulative(self): with self._lock: return self.data.get_cumulative() def to_runner_api_monitoring_info_impl(self, name, transform_id): raise NotImplementedError(type(self)) class GaugeCell(AbstractMetricCell): """For internal use only; no backwards-compatibility guarantees. Tracks the current value and delta for a gauge metric. Each cell tracks the state of a metric independently per context per bundle. Therefore, each metric has a different cell in each bundle, that is later aggregated. This class is thread safe. """ def __init__(self): super().__init__(GaugeData) def _update_locked(self, value): # Set the value directly without checking timestamp, because # this value is naturally the latest value. self.data.value = int(value) self.data.timestamp = time.time() def to_runner_api_monitoring_info_impl(self, name, transform_id): from apache_beam.metrics import monitoring_infos return monitoring_infos.int64_user_gauge( name.namespace, name.name, self.get_cumulative(), ptransform=transform_id) class StringSetCell(AbstractMetricCell): """For internal use only; no backwards-compatibility guarantees. Tracks the current value for a StringSet metric. Each cell tracks the state of a metric independently per context per bundle. Therefore, each metric has a different cell in each bundle, that is later aggregated. This class is thread safe. """ def __init__(self): super().__init__(StringSetData) def add(self, value): self.update(value) def _update_locked(self, value): self.data.add(value) def to_runner_api_monitoring_info_impl(self, name, transform_id): from apache_beam.metrics import monitoring_infos return monitoring_infos.user_set_string( name.namespace, name.name, self.get_cumulative(), ptransform=transform_id) class BoundedTrieCell(AbstractMetricCell): """For internal use only; no backwards-compatibility guarantees. Tracks the current value for a BoundedTrie metric. Each cell tracks the state of a metric independently per context per bundle. Therefore, each metric has a different cell in each bundle, that is later aggregated. This class is thread safe. """ def __init__(self): super().__init__(BoundedTrieData) def add(self, value): self.update(value) def _update_locked(self, value): self.data.add(value) def to_runner_api_monitoring_info_impl(self, name, transform_id): from apache_beam.metrics import monitoring_infos return monitoring_infos.user_bounded_trie( name.namespace, name.name, self.get_cumulative(), ptransform=transform_id) class DistributionResult(object): """The result of a Distribution metric.""" def __init__(self, data): # type: (DistributionData) -> None self.data = data def __eq__(self, other): # type: (object) -> bool if isinstance(other, DistributionResult): return self.data == other.data else: return False def __hash__(self): # type: () -> int return hash(self.data) def __repr__(self): # type: () -> str return ( 'DistributionResult(sum={}, count={}, min={}, max={}, ' 'mean={})'.format(self.sum, self.count, self.min, self.max, self.mean)) @property def max(self): # type: () -> Optional[int] return self.data.max if self.data.count else None @property def min(self): # type: () -> Optional[int] return self.data.min if self.data.count else None @property def count(self): # type: () -> Optional[int] return self.data.count @property def sum(self): # type: () -> Optional[int] return self.data.sum @property def mean(self): # type: () -> Optional[float] """Returns the float mean of the distribution. If the distribution contains no elements, it returns None. """ if self.data.count == 0: return None return self.data.sum / self.data.count class GaugeResult(object): def __init__(self, data): # type: (GaugeData) -> None self.data = data def __eq__(self, other): # type: (object) -> bool if isinstance(other, GaugeResult): return self.data == other.data else: return False def __hash__(self): # type: () -> int return hash(self.data) def __repr__(self): return ''.format( self.value, self.timestamp) @property def value(self): # type: () -> Optional[int] return self.data.value @property def timestamp(self): # type: () -> Optional[int] return self.data.timestamp class GaugeData(object): """For internal use only; no backwards-compatibility guarantees. The data structure that holds data about a gauge metric. Gauge metrics are restricted to integers only. This object is not thread safe, so it's not supposed to be modified by other than the GaugeCell that contains it. """ def __init__(self, value, timestamp=None): # type: (Optional[int], Optional[int]) -> None self.value = value self.timestamp = timestamp if timestamp is not None else 0 def __eq__(self, other): # type: (object) -> bool if isinstance(other, GaugeData): return self.value == other.value and self.timestamp == other.timestamp else: return False def __hash__(self): # type: () -> int return hash((self.value, self.timestamp)) def __repr__(self): # type: () -> str return ''.format( self.value, self.timestamp) def get_cumulative(self): # type: () -> GaugeData return GaugeData(self.value, timestamp=self.timestamp) def get_result(self): # type: () -> GaugeResult return GaugeResult(self.get_cumulative()) def combine(self, other): # type: (Optional[GaugeData]) -> GaugeData if other is None: return self if other.timestamp > self.timestamp: return other else: return self @staticmethod def singleton(value, timestamp=None): # type: (Optional[int], Optional[int]) -> GaugeData return GaugeData(value, timestamp=timestamp) @staticmethod def identity_element(): # type: () -> GaugeData return GaugeData(0, timestamp=0) class DistributionData(object): """For internal use only; no backwards-compatibility guarantees. The data structure that holds data about a distribution metric. Distribution metrics are restricted to distributions of integers only. This object is not thread safe, so it's not supposed to be modified by other than the DistributionCell that contains it. """ def __init__(self, sum, count, min, max): # type: (int, int, int, int) -> None if count: self.sum = sum self.count = count self.min = min self.max = max else: self.sum = self.count = 0 self.min = 2**63 - 1 # Avoid Wimplicitly-unsigned-literal caused by -2**63. self.max = -self.min - 1 def __eq__(self, other): # type: (object) -> bool if isinstance(other, DistributionData): return ( self.sum == other.sum and self.count == other.count and self.min == other.min and self.max == other.max) else: return False def __hash__(self): # type: () -> int return hash((self.sum, self.count, self.min, self.max)) def __repr__(self): # type: () -> str return 'DistributionData(sum={}, count={}, min={}, max={})'.format( self.sum, self.count, self.min, self.max) def get_cumulative(self): # type: () -> DistributionData return DistributionData(self.sum, self.count, self.min, self.max) def get_result(self) -> DistributionResult: return DistributionResult(self.get_cumulative()) def combine(self, other): # type: (Optional[DistributionData]) -> DistributionData if other is None: return self return DistributionData( self.sum + other.sum, self.count + other.count, self.min if self.min < other.min else other.min, self.max if self.max > other.max else other.max) @staticmethod def singleton(value): # type: (int) -> DistributionData return DistributionData(value, 1, value, value) @staticmethod def identity_element(): # type: () -> DistributionData return DistributionData(0, 0, 2**63 - 1, -2**63) class StringSetData(object): """For internal use only; no backwards-compatibility guarantees. The data structure that holds data about a StringSet metric. StringSet metrics are restricted to set of strings only. This object is not thread safe, so it's not supposed to be modified by other than the StringSetCell that contains it. The summation of all string length for a StringSetData cannot exceed 1 MB. Further addition of elements are dropped. """ _STRING_SET_SIZE_LIMIT = 1_000_000 def __init__(self, string_set: Optional[Set] = None, string_size: int = 0): self.string_set = string_set or set() if not string_size: string_size = 0 for s in self.string_set: string_size += len(s) self.string_size = string_size def __eq__(self, other: object) -> bool: if isinstance(other, StringSetData): return ( self.string_size == other.string_size and self.string_set == other.string_set) else: return False def __hash__(self) -> int: return hash(self.string_set) def __repr__(self) -> str: return 'StringSetData{}:{}'.format(self.string_set, self.string_size) def get_cumulative(self) -> "StringSetData": return StringSetData(set(self.string_set), self.string_size) def get_result(self) -> Set[str]: return set(self.string_set) def add(self, *strings): """ Add strings into this StringSetData and return the result StringSetData. Reuse the original StringSetData's set. """ self.string_size = self.add_until_capacity( self.string_set, self.string_size, strings) return self def combine(self, other: "StringSetData") -> "StringSetData": """ Combines this StringSetData with other, both original StringSetData are left intact. """ if other is None: return self if not other.string_set: return self elif not self.string_set: return other combined = set(self.string_set) string_size = self.add_until_capacity( combined, self.string_size, other.string_set) return StringSetData(combined, string_size) @classmethod def add_until_capacity( cls, combined: set, current_size: int, others: Iterable[str]): """ Add strings into set until reach capacity. Return the all string size of added set. """ if current_size > cls._STRING_SET_SIZE_LIMIT: return current_size for string in others: if string not in combined: combined.add(string) current_size += len(string) if current_size > cls._STRING_SET_SIZE_LIMIT: _LOGGER.warning( "StringSet metrics reaches capacity. Further incoming elements " "won't be recorded. Current size: %d, last element size: %d.", current_size, len(string)) break return current_size @staticmethod def singleton(value: str) -> "StringSetData": return StringSetData({value}) @staticmethod def identity_element() -> "StringSetData": return StringSetData() class _BoundedTrieNode(object): def __init__(self): # invariant: size = len(self.flattened()) = min(1, sum(size of children)) self._size = 1 self._children: Optional[dict[str, '_BoundedTrieNode']] = {} self._truncated = False def to_proto(self) -> metrics_pb2.BoundedTrieNode: return metrics_pb2.BoundedTrieNode( truncated=self._truncated, children={ name: child.to_proto() for name, child in self._children.items() } if self._children else None) @staticmethod def from_proto(proto: metrics_pb2.BoundedTrieNode) -> '_BoundedTrieNode': node = _BoundedTrieNode() if proto.truncated: node._truncated = True node._children = None else: node._children = { name: _BoundedTrieNode.from_proto(child) for name, child in proto.children.items() } node._size = max(1, sum(child._size for child in node._children.values())) return node def size(self): return self._size def contains(self, segments): if self._truncated or not segments: return True head, *tail = segments return head in self._children and self._children[head].contains(tail) def add(self, segments) -> int: if self._truncated or not segments: return 0 head, *tail = segments was_empty = not self._children child = self._children.get(head, None) # type: ignore[union-attr] if child is None: child = self._children[head] = _BoundedTrieNode() # type: ignore[index] delta = 0 if was_empty else 1 else: delta = 0 if tail: delta += child.add(tail) self._size += delta return delta def add_all(self, segments_iter): return sum(self.add(segments) for segments in segments_iter) def trim(self) -> int: if not self._children: return 0 max_child = max(self._children.values(), key=lambda child: child._size) if max_child._size == 1: delta = 1 - self._size self._truncated = True self._children = None else: delta = max_child.trim() self._size += delta return delta def merge(self, other: '_BoundedTrieNode') -> int: if self._truncated: delta = 0 elif other._truncated: delta = 1 - self._size self._truncated = True self._children = None elif not other._children: delta = 0 elif not self._children: self._children = other._children delta = other._size - self._size else: delta = 0 other_child: '_BoundedTrieNode' self_child: Optional['_BoundedTrieNode'] for prefix, other_child in other._children.items(): self_child = self._children.get(prefix, None) if self_child is None: self._children[prefix] = other_child delta += other_child._size else: delta += self_child.merge(other_child) self._size += delta return delta def flattened(self): if self._truncated: yield (True, ) elif not self._children: yield (False, ) else: for prefix, child in sorted(self._children.items()): for flattened in child.flattened(): yield (prefix, ) + flattened def __hash__(self): return self._truncated or hash(sorted(self._children.items())) def __eq__(self, other): if isinstance(other, _BoundedTrieNode): return ( self._truncated == other._truncated and self._children == other._children) else: return False def __repr__(self): return repr(set(''.join(str(s) for s in t) for t in self.flattened())) class BoundedTrieData(object): _DEFAULT_BOUND = 100 def __init__(self, *, root=None, singleton=None, bound=_DEFAULT_BOUND): self._singleton = singleton self._root = root self._bound = bound assert singleton is None or root is None def size(self): if self._singleton is not None: return 1 elif self._root is not None: return self._root.size() else: return 0 def contains(self, value): if self._singleton is not None: return tuple(value) == self._singleton elif self._root is not None: return self._root.contains(value) else: return False def flattened(self): return self.as_trie().flattened() def to_proto(self) -> metrics_pb2.BoundedTrie: return metrics_pb2.BoundedTrie( bound=self._bound, singleton=self._singleton if self._singleton else None, root=self._root.to_proto() if self._root else None) @staticmethod def from_proto(proto: metrics_pb2.BoundedTrie) -> 'BoundedTrieData': return BoundedTrieData( bound=proto.bound, singleton=tuple(proto.singleton) if proto.singleton else None, root=( _BoundedTrieNode.from_proto(proto.root) if proto.HasField('root') else None)) def as_trie(self): if self._root is not None: return self._root else: root = _BoundedTrieNode() if self._singleton is not None: root.add(self._singleton) return root def __eq__(self, other: object) -> bool: if isinstance(other, BoundedTrieData): return self.as_trie() == other.as_trie() else: return False def __hash__(self) -> int: return hash(self.as_trie()) def __repr__(self) -> str: return 'BoundedTrieData({})'.format(self.as_trie()) def get_cumulative(self) -> "BoundedTrieData": return copy.deepcopy(self) def get_result(self) -> Set[tuple]: if self._root is None: if self._singleton is None: return set() else: return set([self._singleton + (False, )]) else: return set(self._root.flattened()) def add(self, segments): if self._root is None and self._singleton is None: self._singleton = segments elif self._singleton is not None and self._singleton == segments: # Optimize for the common case of re-adding the same value. return else: if self._root is None: self._root = self.as_trie() self._singleton = None self._root.add(segments) if self._root._size > self._bound: self._root.trim() def combine(self, other: "BoundedTrieData") -> "BoundedTrieData": if self._root is None and self._singleton is None: return other elif other._root is None and other._singleton is None: return self else: if self._root is None and other._root is not None: self, other = other, self combined = copy.deepcopy(self.as_trie()) if other._root is not None: combined.merge(other._root) else: combined.add(other._singleton) self._bound = min(self._bound, other._bound) while combined._size > self._bound: combined.trim() return BoundedTrieData(root=combined) @staticmethod def singleton(value: str) -> "BoundedTrieData": s = BoundedTrieData() s.add(value) return s @staticmethod def identity_element() -> "BoundedTrieData": return BoundedTrieData() class HistogramCell(MetricCell): """For internal use only; no backwards-compatibility guarantees. Tracks the current value and delta for a histogram metric. Each cell tracks the state of a metric independently per context per bundle. Therefore, each metric has a different cell in each bundle, that is later aggregated. This class is thread safe since underlying histogram object is thread safe. """ def __init__(self, bucket_type): self._bucket_type = bucket_type self.data = HistogramData.identity_element(bucket_type) def reset(self): self.data = HistogramData.identity_element(self._bucket_type) def combine(self, other: 'HistogramCell') -> 'HistogramCell': result = HistogramCell(self._bucket_type) result.data = self.data.combine(other.data) return result def update(self, value): self.data.histogram.record(value) def get_cumulative(self) -> 'HistogramData': return self.data.get_cumulative() def to_runner_api_monitoring_info(self, name, transform_id): # Histogram metric is currently worker-local and internal # use only. This method should be implemented when runners # support Histogram metric reporting. return None class HistogramCellFactory(MetricCellFactory): def __init__(self, bucket_type): self._bucket_type = bucket_type def __call__(self): return HistogramCell(self._bucket_type) def __eq__(self, other): if not isinstance(other, HistogramCellFactory): return False return self._bucket_type == other._bucket_type def __hash__(self): return hash(self._bucket_type) class HistogramResult(object): def __init__(self, data: 'HistogramData') -> None: self.data = data def __eq__(self, other): if isinstance(other, HistogramResult): return self.data == other.data else: return False def __hash__(self): return hash(self.data) def __repr__(self): return ''.format( self.data.histogram.get_percentile_info()) @property def p99(self): return self.data.histogram.p99() @property def p95(self): return self.data.histogram.p95() @property def p90(self): return self.data.histogram.p90() @property def histogram(self): return self.data.histogram class HistogramData(object): """For internal use only; no backwards-compatibility guarantees. The data structure that holds data about a histogram metric. This object is not thread safe, so it's not supposed to be modified outside the HistogramCell. """ def __init__(self, histogram): self.histogram = histogram def __eq__(self, other): return self.histogram == other.histogram def __hash__(self): return hash(self.histogram) def __repr__(self): return 'HistogramData({})'.format(self.histogram.get_percentile_info()) def get_cumulative(self) -> 'HistogramData': return HistogramData(self.histogram) def combine(self, other: Optional['HistogramData']) -> 'HistogramData': if other is None: return self return HistogramData(self.histogram.combine(other.histogram)) @staticmethod def identity_element(bucket_type) -> 'HistogramData': return HistogramData(Histogram(bucket_type)) def get_result(self) -> 'HistogramResult': return HistogramResult(self.get_cumulative()) def to_proto(self) -> metrics_pb2.HistogramValue: return self.histogram.to_runner_api() @classmethod def from_proto(cls, proto: metrics_pb2.HistogramValue): return cls(Histogram.from_runner_api(proto)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/metrics/cells_test.py0000644000175100017510000003756615206113032022057 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import copy import itertools import random import threading import unittest from apache_beam.metrics.cells import BoundedTrieData from apache_beam.metrics.cells import CounterCell from apache_beam.metrics.cells import DistributionCell from apache_beam.metrics.cells import DistributionData from apache_beam.metrics.cells import GaugeCell from apache_beam.metrics.cells import GaugeData from apache_beam.metrics.cells import HistogramCell from apache_beam.metrics.cells import HistogramCellFactory from apache_beam.metrics.cells import HistogramData from apache_beam.metrics.cells import StringSetCell from apache_beam.metrics.cells import StringSetData from apache_beam.metrics.cells import _BoundedTrieNode from apache_beam.metrics.metricbase import MetricName from apache_beam.utils.histogram import Histogram from apache_beam.utils.histogram import LinearBucket class TestCounterCell(unittest.TestCase): @classmethod def _modify_counter(cls, d): for i in range(cls.NUM_ITERATIONS): d.inc(i) NUM_THREADS = 5 NUM_ITERATIONS = 100 def test_parallel_access(self): # We create NUM_THREADS threads that concurrently modify the counter. threads = [] c = CounterCell() for _ in range(TestCounterCell.NUM_THREADS): t = threading.Thread( target=TestCounterCell._modify_counter, args=(c, )) threads.append(t) t.start() for t in threads: t.join() total = ( self.NUM_ITERATIONS * (self.NUM_ITERATIONS - 1) // 2 * self.NUM_THREADS) self.assertEqual(c.get_cumulative(), total) def test_basic_operations(self): c = CounterCell() c.inc(2) self.assertEqual(c.get_cumulative(), 2) c.dec(10) self.assertEqual(c.get_cumulative(), -8) c.dec() self.assertEqual(c.get_cumulative(), -9) c.inc() self.assertEqual(c.get_cumulative(), -8) def test_start_time_set(self): c = CounterCell() c.inc(2) name = MetricName('namespace', 'name1') mi = c.to_runner_api_monitoring_info(name, 'transform_id') self.assertGreater(mi.start_time.seconds, 0) class TestDistributionCell(unittest.TestCase): @classmethod def _modify_distribution(cls, d): for i in range(cls.NUM_ITERATIONS): d.update(i) NUM_THREADS = 5 NUM_ITERATIONS = 100 def test_parallel_access(self): # We create NUM_THREADS threads that concurrently modify the distribution. threads = [] d = DistributionCell() for _ in range(TestDistributionCell.NUM_THREADS): t = threading.Thread( target=TestDistributionCell._modify_distribution, args=(d, )) threads.append(t) t.start() for t in threads: t.join() total = ( self.NUM_ITERATIONS * (self.NUM_ITERATIONS - 1) // 2 * self.NUM_THREADS) count = (self.NUM_ITERATIONS * self.NUM_THREADS) self.assertEqual( d.get_cumulative(), DistributionData(total, count, 0, self.NUM_ITERATIONS - 1)) def test_basic_operations(self): d = DistributionCell() d.update(10) self.assertEqual(d.get_cumulative(), DistributionData(10, 1, 10, 10)) d.update(2) self.assertEqual(d.get_cumulative(), DistributionData(12, 2, 2, 10)) d.update(900) self.assertEqual(d.get_cumulative(), DistributionData(912, 3, 2, 900)) def test_integer_only(self): d = DistributionCell() d.update(3.1) d.update(3.2) d.update(3.3) self.assertEqual(d.get_cumulative(), DistributionData(9, 3, 3, 3)) def test_start_time_set(self): d = DistributionCell() d.update(3.1) name = MetricName('namespace', 'name1') mi = d.to_runner_api_monitoring_info(name, 'transform_id') self.assertGreater(mi.start_time.seconds, 0) class TestGaugeCell(unittest.TestCase): def test_basic_operations(self): g = GaugeCell() g.set(10) self.assertEqual(g.get_cumulative().value, GaugeData(10).value) g.set(2) self.assertEqual(g.get_cumulative().value, 2) def test_integer_only(self): g = GaugeCell() g.set(3.3) self.assertEqual(g.get_cumulative().value, 3) def test_combine_appropriately(self): g1 = GaugeCell() g1.set(3) g2 = GaugeCell() g2.set(1) # THe second Gauge, with value 1 was the most recent, so it should be # the final result. result = g2.combine(g1) self.assertEqual(result.data.value, 1) def test_start_time_set(self): g1 = GaugeCell() g1.set(3) name = MetricName('namespace', 'name1') mi = g1.to_runner_api_monitoring_info(name, 'transform_id') self.assertGreater(mi.start_time.seconds, 0) class TestStringSetCell(unittest.TestCase): def test_not_leak_mutable_set(self): c = StringSetCell() c.add('test') c.add('another') s = c.get_cumulative() self.assertEqual(s, StringSetData({'test', 'another'}, 11)) s.add('yet another') self.assertEqual(c.get_cumulative(), StringSetData({'test', 'another'}, 11)) def test_combine_appropriately(self): s1 = StringSetCell() s1.add('1') s1.add('2') s2 = StringSetCell() s2.add('1') s2.add('3') result = s2.combine(s1) self.assertEqual(result.data, StringSetData({'1', '2', '3'})) def test_add_size_tracked_correctly(self): s = StringSetCell() s.add('1') s.add('2') self.assertEqual(s.data.string_size, 2) s.add('2') s.add('3') self.assertEqual(s.data.string_size, 3) class TestBoundedTrieNode(unittest.TestCase): @classmethod def random_segments_fixed_depth(cls, n, depth, overlap, rand): if depth == 0: yield from ((), ) * n else: seen = [] to_string = lambda ix: chr(ord('a') + ix) if ix < 26 else f'z{ix}' for suffix in cls.random_segments_fixed_depth(n, depth - 1, overlap, rand): if not seen or rand.random() > overlap: prefix = to_string(len(seen)) seen.append(prefix) else: prefix = rand.choice(seen) yield (prefix, ) + suffix @classmethod def random_segments(cls, n, min_depth, max_depth, overlap, rand): for depth, segments in zip( itertools.cycle(range(min_depth, max_depth + 1)), cls.random_segments_fixed_depth(n, max_depth, overlap, rand)): yield segments[:depth] def assert_covers(self, node, expected, max_truncated=0): self.assert_covers_flattened(node.flattened(), expected, max_truncated) def assert_covers_flattened(self, flattened, expected, max_truncated=0): expected = set(expected) # Split node into the exact and truncated segments. partitioned = {True: set(), False: set()} for segments in flattened: partitioned[segments[-1]].add(segments[:-1]) exact, truncated = partitioned[False], partitioned[True] # Check we cover both parts. self.assertLessEqual(len(truncated), max_truncated, truncated) self.assertTrue(exact.issubset(expected), exact - expected) seen_truncated = set() for segments in expected - exact: found = 0 for ix in range(len(segments)): if segments[:ix] in truncated: seen_truncated.add(segments[:ix]) found += 1 if found != 1: self.fail( f"Expected exactly one prefix of {segments} " f"to occur in {truncated}, found {found}") self.assertEqual(seen_truncated, truncated, truncated - seen_truncated) def run_covers_test(self, flattened, expected, max_truncated): def parse(s): return tuple(s.strip('*')) + (s.endswith('*'), ) self.assert_covers_flattened([parse(s) for s in flattened], [tuple(s) for s in expected], max_truncated) def test_covers_exact(self): self.run_covers_test(['ab', 'ac', 'cd'], ['ab', 'ac', 'cd'], 0) with self.assertRaises(AssertionError): self.run_covers_test(['ab', 'ac', 'cd'], ['ac', 'cd'], 0) with self.assertRaises(AssertionError): self.run_covers_test(['ab', 'ac'], ['ab', 'ac', 'cd'], 0) with self.assertRaises(AssertionError): self.run_covers_test(['a*', 'cd'], ['ab', 'ac', 'cd'], 0) def test_covers_trunacted(self): self.run_covers_test(['a*', 'cd'], ['ab', 'ac', 'cd'], 1) self.run_covers_test(['a*', 'cd'], ['ab', 'ac', 'abcde', 'cd'], 1) with self.assertRaises(AssertionError): self.run_covers_test(['ab', 'ac', 'cd'], ['ac', 'cd'], 1) with self.assertRaises(AssertionError): self.run_covers_test(['ab', 'ac'], ['ab', 'ac', 'cd'], 1) with self.assertRaises(AssertionError): self.run_covers_test(['a*', 'c*'], ['ab', 'ac', 'cd'], 1) with self.assertRaises(AssertionError): self.run_covers_test(['a*', 'c*'], ['ab', 'ac'], 1) def run_test(self, to_add): everything = list(set(to_add)) all_prefixees = set( segments[:ix] for segments in everything for ix in range(len(segments))) everything_deduped = set(everything) - all_prefixees # Check basic addition. node = _BoundedTrieNode() total_size = node.size() self.assertEqual(total_size, 1) for segments in everything: total_size += node.add(segments) self.assertEqual(node.size(), len(everything_deduped), node) self.assertEqual(node.size(), total_size, node) self.assert_covers(node, everything_deduped) # Check merging node0 = _BoundedTrieNode() node0.add_all(everything[0::2]) node1 = _BoundedTrieNode() node1.add_all(everything[1::2]) pre_merge_size = node0.size() merge_delta = node0.merge(node1) self.assertEqual(node0.size(), pre_merge_size + merge_delta) self.assertEqual(node0, node) # Check trimming. if node.size() > 1: trim_delta = node.trim() self.assertLess(trim_delta, 0, node) self.assertEqual(node.size(), total_size + trim_delta) self.assert_covers(node, everything_deduped, max_truncated=1) if node.size() > 1: trim2_delta = node.trim() self.assertLess(trim2_delta, 0) self.assertEqual(node.size(), total_size + trim_delta + trim2_delta) self.assert_covers(node, everything_deduped, max_truncated=2) # Adding after trimming should be a no-op. node_copy = copy.deepcopy(node) for segments in everything: self.assertEqual(node.add(segments), 0) self.assertEqual(node, node_copy) # Merging after trimming should be a no-op. self.assertEqual(node.merge(node0), 0) self.assertEqual(node.merge(node1), 0) self.assertEqual(node, node_copy) if node._truncated: expected_delta = 0 else: expected_delta = 2 # Adding something new is not. new_values = [('new1', ), ('new2', 'new2.1')] self.assertEqual(node.add_all(new_values), expected_delta) self.assert_covers( node, list(everything_deduped) + new_values, max_truncated=2) # Nor is merging something new. new_values_node = _BoundedTrieNode() new_values_node.add_all(new_values) self.assertEqual(node_copy.merge(new_values_node), expected_delta) self.assert_covers( node_copy, list(everything_deduped) + new_values, max_truncated=2) def run_fuzz(self, iterations=10, **params): for _ in range(iterations): seed = random.getrandbits(64) segments = self.random_segments(**params, rand=random.Random(seed)) try: self.run_test(segments) except: print("SEED", seed) raise def test_trivial(self): self.run_test([('a', 'b'), ('a', 'c')]) def test_flat(self): self.run_test([('a', 'a'), ('b', 'b'), ('c', 'c')]) def test_deep(self): self.run_test([('a', ) * 10, ('b', ) * 12]) def test_small(self): self.run_fuzz(n=5, min_depth=2, max_depth=3, overlap=0.5) def test_medium(self): self.run_fuzz(n=20, min_depth=2, max_depth=4, overlap=0.5) def test_large_sparse(self): self.run_fuzz(n=120, min_depth=2, max_depth=4, overlap=0.2) def test_large_dense(self): self.run_fuzz(n=120, min_depth=2, max_depth=4, overlap=0.8) def test_bounded_trie_data_combine(self): empty = BoundedTrieData() # The merging here isn't complicated we're just ensuring that # BoundedTrieData invokes _BoundedTrieNode correctly. singletonA = BoundedTrieData(singleton=('a', 'a')) singletonB = BoundedTrieData(singleton=('b', 'b')) lots_root = _BoundedTrieNode() lots_root.add_all([('c', 'c'), ('d', 'd')]) lots = BoundedTrieData(root=lots_root) self.assertEqual(empty.get_result(), set()) self.assertEqual( empty.combine(singletonA).get_result(), set([('a', 'a', False)])) self.assertEqual( singletonA.combine(empty).get_result(), set([('a', 'a', False)])) self.assertEqual( singletonA.combine(singletonB).get_result(), set([('a', 'a', False), ('b', 'b', False)])) self.assertEqual( singletonA.combine(lots).get_result(), set([('a', 'a', False), ('c', 'c', False), ('d', 'd', False)])) self.assertEqual( lots.combine(singletonA).get_result(), set([('a', 'a', False), ('c', 'c', False), ('d', 'd', False)])) def test_bounded_trie_data_combine_trim(self): left = _BoundedTrieNode() left.add_all([('a', 'x'), ('b', 'd')]) right = _BoundedTrieNode() right.add_all([('a', 'y'), ('c', 'd')]) self.assertEqual( BoundedTrieData(root=left).combine( BoundedTrieData(root=right, bound=3)).get_result(), set([('a', True), ('b', 'd', False), ('c', 'd', False)])) def test_merge_on_empty_node(self): root1 = _BoundedTrieNode() root2 = _BoundedTrieNode() root2.add_all([["a", "b", "c"], ["a", "b", "d"], ["a", "e"]]) self.assertEqual(2, root1.merge(root2)) self.assertEqual(3, root1.size()) self.assertFalse(root1._truncated) def test_merge_with_empty_node(self): root1 = _BoundedTrieNode() root1.add_all([["a", "b", "c"], ["a", "b", "d"], ["a", "e"]]) root2 = _BoundedTrieNode() self.assertEqual(0, root1.merge(root2)) self.assertEqual(3, root1.size()) self.assertFalse(root1._truncated) class TestHistogramCell(unittest.TestCase): @classmethod def _modify_histogram(cls, d): for i in range(cls.NUM_ITERATIONS): d.update(i) NUM_THREADS = 5 NUM_ITERATIONS = 100 def test_parallel_access(self): # We create NUM_THREADS threads that concurrently modify the distribution. threads = [] bucket_type = LinearBucket(0, 1, 100) d = HistogramCell(bucket_type) for _ in range(TestHistogramCell.NUM_THREADS): t = threading.Thread( target=TestHistogramCell._modify_histogram, args=(d, )) threads.append(t) t.start() for t in threads: t.join() histogram = Histogram(bucket_type) for _ in range(self.NUM_THREADS): for i in range(self.NUM_ITERATIONS): histogram.record(i) self.assertEqual(d.get_cumulative(), HistogramData(histogram)) def test_basic_operations(self): d = HistogramCellFactory(LinearBucket(0, 1, 10))() d.update(10) self.assertEqual( str(d.get_cumulative()), 'HistogramData(Total count: 1, P99: >=10, P90: >=10, P50: >=10)') d.update(0) self.assertEqual( str(d.get_cumulative()), 'HistogramData(Total count: 2, P99: >=10, P90: >=10, P50: 1)') d.update(5) self.assertEqual( str(d.get_cumulative()), 'HistogramData(Total count: 3, P99: >=10, P90: >=10, P50: 6)') if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/metrics/execution.pxd0000644000175100017510000000263015206113032022044 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # cimport cython cimport libc.stdint from apache_beam.metrics.cells cimport MetricCell cdef object get_current_tracker cdef class _TypedMetricName(object): cdef readonly object cell_type cdef readonly object metric_name cdef readonly object fast_name cdef libc.stdint.int64_t _hash cdef object _DEFAULT cdef class MetricUpdater(object): cdef _TypedMetricName typed_metric_name cdef object default_value cdef bint process_wide # bint is used to represent C++ bool. cdef class MetricsContainer(object): cdef object step_name cdef object lock cdef public dict metrics cpdef MetricCell get_metric_cell(self, metric_key) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/metrics/execution.py0000644000175100017510000003225115206113032021703 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ This module is for internal use only; no backwards-compatibility guarantees. The classes in this file keep shared state, and organize metrics information. Available classes: - MetricKey - Internal key for a metric. - MetricResult - Current status of a metric's updates/commits. - _MetricsEnvironment - Keeps track of MetricsContainer and other metrics information for every single execution working thread. - MetricsContainer - Holds the metrics of a single step and a single unit-of-commit (bundle). """ # pytype: skip-file # ruff: noqa: UP006 import threading from typing import TYPE_CHECKING from typing import Any from typing import Dict from typing import FrozenSet from typing import Optional from typing import Type from typing import Union from typing import cast from apache_beam.metrics import monitoring_infos from apache_beam.metrics.cells import BoundedTrieCell from apache_beam.metrics.cells import CounterCell from apache_beam.metrics.cells import DistributionCell from apache_beam.metrics.cells import GaugeCell from apache_beam.metrics.cells import HistogramCellFactory from apache_beam.metrics.cells import HistogramData from apache_beam.metrics.cells import StringSetCell from apache_beam.metrics.cells import StringSetData from apache_beam.runners.worker import statesampler from apache_beam.runners.worker.statesampler import get_current_tracker if TYPE_CHECKING: from apache_beam.metrics.cells import BoundedTrieData from apache_beam.metrics.cells import DistributionData from apache_beam.metrics.cells import GaugeData from apache_beam.metrics.cells import MetricCell from apache_beam.metrics.cells import MetricCellFactory from apache_beam.metrics.metricbase import MetricName from apache_beam.portability.api import metrics_pb2 class MetricKey(object): """Key used to identify instance of metric cell. Metrics are internally keyed by the name of the step they're associated with, the name and namespace (if it is a user defined metric) of the metric, and any extra label metadata added by the runner specific metric collection service. """ def __init__(self, step, metric, labels=None): """Initializes ``MetricKey``. Args: step: A string with the step this metric cell is part of. metric: A ``MetricName`` namespace+name that identifies a metric. labels: An arbitrary set of labels that also identifies the metric. """ self.step = step self.metric = metric self.labels = labels if labels else {} def __eq__(self, other): return ( self.step == other.step and self.metric == other.metric and self.labels == other.labels) def __hash__(self): return hash((self.step, self.metric, frozenset(self.labels))) def __repr__(self): return 'MetricKey(step={}, metric={}, labels={})'.format( self.step, self.metric, self.labels) class MetricResult(object): """Keeps track of the status of a metric within a single bundle. It contains the physical and logical updates to the metric. Physical updates are updates that have not necessarily been committed, but that have been made during pipeline execution. Logical updates are updates that have been committed. Attributes: key: A ``MetricKey`` that identifies the metric and bundle of this result. committed: The committed updates of the metric. This attribute's type is of metric type result (e.g. int, DistributionResult, GaugeResult). attempted: The logical updates of the metric. This attribute's type is that of metric type result (e.g. int, DistributionResult, GaugeResult). """ def __init__(self, key, committed, attempted): """Initializes ``MetricResult``. Args: key: A ``MetricKey`` object. committed: Metric data that has been committed (e.g. logical updates) attempted: Metric data that has been attempted (e.g. physical updates) """ self.key = key self.committed = committed self.attempted = attempted def __eq__(self, other): return ( self.key == other.key and self.committed == other.committed and self.attempted == other.attempted) def __hash__(self): return hash((self.key, self.committed, self.attempted)) def __repr__(self): return 'MetricResult(key={}, committed={}, attempted={})'.format( self.key, str(self.committed), str(self.attempted)) def __str__(self): return repr(self) @property def result(self): """Short-hand for falling back to attempted metrics if it seems that committed was not populated (e.g. due to not being supported on a given runner""" return self.committed if self.committed else self.attempted class _MetricsEnvironment(object): """Holds the MetricsContainer for every thread and other metric information. This class is not meant to be instantiated, instead being used to keep track of global state. """ def current_container(self): """Returns the current MetricsContainer.""" sampler = statesampler.get_current_tracker() if sampler is None: return None return sampler.current_state().metrics_container def process_wide_container(self): """Returns the MetricsContainer for process wide metrics, e.g. memory.""" return PROCESS_WIDE_METRICS_CONTAINER MetricsEnvironment = _MetricsEnvironment() class _TypedMetricName(object): """Like MetricName, but also stores the cell type of the metric.""" def __init__( self, cell_type, # type: Union[Type[MetricCell], MetricCellFactory] metric_name # type: Union[str, MetricName] ): # type: (...) -> None self.cell_type = cell_type self.metric_name = metric_name if isinstance(metric_name, str): self.fast_name = metric_name else: self.fast_name = metric_name.fast_name() # Cached for speed, as this is used as a key for every counter update. self._hash = hash((cell_type, self.fast_name)) def __eq__(self, other): return self is other or ( self.cell_type == other.cell_type and self.fast_name == other.fast_name) def __hash__(self): return self._hash def __str__(self): return '%s %s' % (self.cell_type, self.metric_name) def __reduce__(self): return _TypedMetricName, (self.cell_type, self.metric_name) _DEFAULT = None # type: Any class MetricUpdater(object): """A callable that updates the metric as quickly as possible.""" def __init__( self, cell_type, # type: Union[Type[MetricCell], MetricCellFactory] metric_name, # type: Union[str, MetricName] default_value=None, process_wide=False): self.process_wide = process_wide self.typed_metric_name = _TypedMetricName(cell_type, metric_name) self.default_value = default_value def __call__(self, value=_DEFAULT): # type: (Any) -> None if value is _DEFAULT: if self.default_value is _DEFAULT: raise ValueError( 'Missing value for update of %s' % self.typed_metric_name.fast_name) value = self.default_value if self.process_wide: MetricsEnvironment.process_wide_container().get_metric_cell( self.typed_metric_name).update(value) else: tracker = get_current_tracker() if tracker is not None: tracker.update_metric(self.typed_metric_name, value) def __reduce__(self): return MetricUpdater, ( self.typed_metric_name.cell_type, self.typed_metric_name.metric_name, self.default_value) class MetricsContainer(object): """Holds the metrics of a single step and a single bundle. Or the metrics associated with the process/SDK harness. I.e. memory usage. """ def __init__(self, step_name): self.step_name = step_name self.lock = threading.Lock() self.metrics = {} # type: Dict[_TypedMetricName, MetricCell] def get_counter(self, metric_name): # type: (MetricName) -> CounterCell return cast( CounterCell, self.get_metric_cell(_TypedMetricName(CounterCell, metric_name))) def get_distribution(self, metric_name): # type: (MetricName) -> DistributionCell return cast( DistributionCell, self.get_metric_cell(_TypedMetricName(DistributionCell, metric_name))) def get_gauge(self, metric_name): # type: (MetricName) -> GaugeCell return cast( GaugeCell, self.get_metric_cell(_TypedMetricName(GaugeCell, metric_name))) def get_string_set(self, metric_name): # type: (MetricName) -> StringSetCell return cast( StringSetCell, self.get_metric_cell(_TypedMetricName(StringSetCell, metric_name))) def get_bounded_trie(self, metric_name): return self.get_metric_cell(_TypedMetricName(BoundedTrieCell, metric_name)) def get_metric_cell(self, typed_metric_name): # type: (_TypedMetricName) -> MetricCell cell = self.metrics.get(typed_metric_name, None) if cell is None: with self.lock: cell = self.metrics[typed_metric_name] = typed_metric_name.cell_type() return cell def get_cumulative(self): # type: () -> MetricUpdates """Return MetricUpdates with cumulative values of all metrics in container. This returns all the cumulative values for all metrics. """ counters = { MetricKey(self.step_name, k.metric_name): v.get_cumulative() for k, v in self.metrics.items() if k.cell_type == CounterCell } distributions = { MetricKey(self.step_name, k.metric_name): v.get_cumulative() for k, v in self.metrics.items() if k.cell_type == DistributionCell } gauges = { MetricKey(self.step_name, k.metric_name): v.get_cumulative() for k, v in self.metrics.items() if k.cell_type == GaugeCell } string_sets = { MetricKey(self.step_name, k.metric_name): v.get_cumulative() for k, v in self.metrics.items() if k.cell_type == StringSetCell } bounded_tries = { MetricKey(self.step_name, k.metric_name): v.get_cumulative() for k, v in self.metrics.items() if k.cell_type == BoundedTrieCell } histograms = { MetricKey(self.step_name, k.metric_name): v.get_cumulative() for k, v in self.metrics.items() if isinstance(k.cell_type, HistogramCellFactory) } return MetricUpdates( counters, distributions, gauges, string_sets, bounded_tries, histograms) def to_runner_api(self): return [ cell.to_runner_api_user_metric(key.metric_name) for key, cell in self.metrics.items() ] def to_runner_api_monitoring_infos(self, transform_id): # type: (str) -> Dict[FrozenSet, metrics_pb2.MonitoringInfo] """Returns a list of MonitoringInfos for the metrics in this container.""" with self.lock: items = list(self.metrics.items()) all_metrics = [ cell.to_runner_api_monitoring_info(key.metric_name, transform_id) for key, cell in items ] return { monitoring_infos.to_key(mi): mi for mi in all_metrics if mi is not None } def reset(self): # type: () -> None """Reset all metrics in the MetricsContainer. This does not delete added metrics. """ for metric in self.metrics.values(): metric.reset() def __reduce__(self): raise NotImplementedError PROCESS_WIDE_METRICS_CONTAINER = MetricsContainer(None) class MetricUpdates(object): """Contains updates for several metrics. A metric update is an object containing information to update a metric. For Distribution metrics, it is DistributionData, and for Counter metrics, it's an int. """ def __init__( self, counters=None, # type: Optional[Dict[MetricKey, int]] distributions=None, # type: Optional[Dict[MetricKey, DistributionData]] gauges=None, # type: Optional[Dict[MetricKey, GaugeData]] string_sets=None, # type: Optional[Dict[MetricKey, StringSetData]] bounded_tries=None, # type: Optional[Dict[MetricKey, BoundedTrieData]] histograms=None, # type: Optional[Dict[MetricKey, HistogramData]] ): # type: (...) -> None """Create a MetricUpdates object. Args: counters: Dictionary of MetricKey:MetricUpdate updates. distributions: Dictionary of MetricKey:MetricUpdate objects. gauges: Dictionary of MetricKey:MetricUpdate objects. string_sets: Dictionary of MetricKey:MetricUpdate objects. bounded_tries: Dictionary of MetricKey:MetricUpdate objects. """ self.counters = counters or {} self.distributions = distributions or {} self.gauges = gauges or {} self.string_sets = string_sets or {} self.bounded_tries = bounded_tries or {} self.histograms = histograms or {} ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/metrics/execution_test.py0000644000175100017510000001154615206113032022746 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import functools import unittest from apache_beam.metrics.execution import MetricKey from apache_beam.metrics.execution import MetricsContainer from apache_beam.metrics.metricbase import MetricName class TestMetricKey(unittest.TestCase): def test_equality_for_key_with_labels(self): test_labels = {'label1', 'value1'} test_object = MetricKey( 'step', MetricName('namespace', 'name'), labels=test_labels) same_labels = MetricKey( 'step', MetricName('namespace', 'name'), labels={'label1', 'value1'}) same_label_reference = MetricKey( 'step', MetricName('namespace', 'name'), labels=test_labels) self.assertEqual(test_object, same_labels) self.assertEqual(test_object, same_label_reference) self.assertEqual(hash(test_object), hash(same_labels)) self.assertEqual(hash(test_object), hash(same_label_reference)) def test_inequality_for_key_with_labels(self): test_labels = {'label1', 'value1'} test_object = MetricKey( 'step', MetricName('namespace', 'name'), labels=test_labels) no_labels = MetricKey('step', MetricName('namespace', 'name')) diff_label_key = MetricKey( 'step', MetricName('namespace', 'name'), labels={'l1_diff', 'value1'}) diff_label_value = MetricKey( 'step', MetricName('namespace', 'name'), labels={'label1', 'v1_diff'}) self.assertNotEqual(test_object, no_labels) self.assertNotEqual(test_object, diff_label_key) self.assertNotEqual(test_object, diff_label_value) self.assertNotEqual(hash(test_object), hash(no_labels)) self.assertNotEqual(hash(test_object), hash(diff_label_key)) self.assertNotEqual(hash(test_object), hash(diff_label_value)) def test_equality_for_key_with_no_labels(self): test_object = MetricKey('step', MetricName('namespace', 'name')) same = MetricKey('step', MetricName('namespace', 'name')) self.assertEqual(test_object, same) self.assertEqual(hash(test_object), hash(same)) diff_step = MetricKey('step_diff', MetricName('namespace', 'name')) diff_namespace = MetricKey('step', MetricName('namespace_diff', 'name')) diff_name = MetricKey('step', MetricName('namespace', 'name_diff')) self.assertNotEqual(test_object, diff_step) self.assertNotEqual(test_object, diff_namespace) self.assertNotEqual(test_object, diff_name) self.assertNotEqual(hash(test_object), hash(diff_step)) self.assertNotEqual(hash(test_object), hash(diff_namespace)) self.assertNotEqual(hash(test_object), hash(diff_name)) class TestMetricsContainer(unittest.TestCase): def test_add_to_counter(self): mc = MetricsContainer('astep') counter = mc.get_counter(MetricName('namespace', 'name')) counter.inc() counter = mc.get_counter(MetricName('namespace', 'name')) self.assertEqual(counter.value, 1) def test_get_cumulative_or_updates(self): mc = MetricsContainer('astep') all_values = [] for i in range(1, 11): counter = mc.get_counter(MetricName('namespace', 'name{}'.format(i))) distribution = mc.get_distribution( MetricName('namespace', 'name{}'.format(i))) gauge = mc.get_gauge(MetricName('namespace', 'name{}'.format(i))) str_set = mc.get_string_set(MetricName('namespace', 'name{}'.format(i))) counter.inc(i) distribution.update(i) gauge.set(i) str_set.add(str(i % 7)) all_values.append(i) # Retrieve ALL updates. cumulative = mc.get_cumulative() self.assertEqual(len(cumulative.counters), 10) self.assertEqual(len(cumulative.distributions), 10) self.assertEqual(len(cumulative.gauges), 10) self.assertEqual(len(cumulative.string_sets), 10) self.assertEqual( set(all_values), {v for _, v in cumulative.counters.items()}) self.assertEqual( set(all_values), {v.value for _, v in cumulative.gauges.items()}) self.assertEqual( {str(i % 7) for i in all_values}, functools.reduce( set.union, (v.string_set for _, v in cumulative.string_sets.items()))) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/metrics/metric.py0000644000175100017510000004250315206113032021164 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ User-facing classes for Metrics API. The classes in this file allow users to define and use metrics to be collected and displayed as part of their pipeline execution. - Metrics - This class lets pipeline and transform writers create and access metric objects such as counters, distributions, etc. """ # pytype: skip-file # mypy: disallow-untyped-defs import logging import re from typing import TYPE_CHECKING from typing import Iterable from typing import Iterator from typing import Optional from typing import Union from apache_beam.metrics import cells from apache_beam.metrics.cells import HistogramCellFactory from apache_beam.metrics.execution import MetricResult from apache_beam.metrics.execution import MetricUpdater from apache_beam.metrics.metricbase import BoundedTrie from apache_beam.metrics.metricbase import Counter from apache_beam.metrics.metricbase import Distribution from apache_beam.metrics.metricbase import Gauge from apache_beam.metrics.metricbase import Histogram from apache_beam.metrics.metricbase import MetricName from apache_beam.metrics.metricbase import StringSet if TYPE_CHECKING: from apache_beam.internal.metrics.metric import MetricLogger from apache_beam.metrics.execution import MetricKey from apache_beam.metrics.metricbase import Metric from apache_beam.utils.histogram import BucketType __all__ = ['Metrics', 'MetricsFilter', 'Lineage'] _LOGGER = logging.getLogger(__name__) class Metrics(object): """Lets users create/access metric objects during pipeline execution.""" @staticmethod def get_namespace(namespace: Union[type, str]) -> str: if isinstance(namespace, type): return '{}.{}'.format(namespace.__module__, namespace.__name__) elif isinstance(namespace, str): return namespace else: raise ValueError('Unknown namespace type') @staticmethod def counter( namespace: Union[type, str], name: str) -> 'Metrics.DelegatingCounter': """Obtains or creates a Counter metric. Args: namespace: A class or string that gives the namespace to a metric name: A string that gives a unique name to a metric Returns: A Counter object. """ namespace = Metrics.get_namespace(namespace) return Metrics.DelegatingCounter(MetricName(namespace, name)) @staticmethod def distribution( namespace: Union[type, str], name: str, process_wide: bool = False) -> 'Metrics.DelegatingDistribution': """Obtains or creates a Distribution metric. Distribution metrics are restricted to integer-only distributions. Args: namespace: A class or string that gives the namespace to a metric name: A string that gives a unique name to a metric Returns: A Distribution object. """ namespace = Metrics.get_namespace(namespace) return Metrics.DelegatingDistribution( MetricName(namespace, name), process_wide=process_wide) @staticmethod def gauge( namespace: Union[type, str], name: str, process_wide: bool = False) -> 'Metrics.DelegatingGauge': """Obtains or creates a Gauge metric. Gauge metrics are restricted to integer-only values. Args: namespace: A class or string that gives the namespace to a metric name: A string that gives a unique name to a metric process_wide: Whether or not the metric is specific to the current bundle or should be calculated for the entire process. Returns: A Gauge object. """ namespace = Metrics.get_namespace(namespace) return Metrics.DelegatingGauge( MetricName(namespace, name), process_wide=process_wide) @staticmethod def string_set( namespace: Union[type, str], name: str) -> 'Metrics.DelegatingStringSet': """Obtains or creates a String set metric. String set metrics are restricted to string values. Args: namespace: A class or string that gives the namespace to a metric name: A string that gives a unique name to a metric Returns: A StringSet object. """ namespace = Metrics.get_namespace(namespace) return Metrics.DelegatingStringSet(MetricName(namespace, name)) @staticmethod def bounded_trie( namespace: Union[type, str], name: str) -> 'Metrics.DelegatingBoundedTrie': """Obtains or creates a Bounded Trie metric. Args: namespace: A class or string that gives the namespace to a metric name: A string that gives a unique name to a metric Returns: A BoundedTrie object. """ namespace = Metrics.get_namespace(namespace) return Metrics.DelegatingBoundedTrie(MetricName(namespace, name)) @staticmethod def histogram( namespace: Union[type, str], name: str, bucket_type: 'BucketType', logger: Optional['MetricLogger'] = None) -> 'Metrics.DelegatingHistogram': """Obtains or creates a Histogram metric. Args: namespace: A class or string that gives the namespace to a metric name: A string that gives a unique name to a metric bucket_type: A type of bucket used in a histogram. A subclass of apache_beam.utils.histogram.BucketType logger: MetricLogger for logging locally aggregated metric Returns: A Histogram object. """ namespace = Metrics.get_namespace(namespace) return Metrics.DelegatingHistogram( MetricName(namespace, name), bucket_type, logger) class DelegatingHistogram(Histogram): """Metrics Histogram that Delegates functionality to MetricsEnvironment.""" def __init__( self, metric_name: MetricName, bucket_type: 'BucketType', logger: Optional['MetricLogger']) -> None: super().__init__(metric_name) self.metric_name = metric_name self.cell_type = HistogramCellFactory(bucket_type) self.logger = logger self.updater = MetricUpdater(self.cell_type, self.metric_name) def update(self, value: object) -> None: self.updater(value) if self.logger: self.logger.update(self.cell_type, self.metric_name, value) class DelegatingCounter(Counter): """Metrics Counter that Delegates functionality to MetricsEnvironment.""" def __init__( self, metric_name: MetricName, process_wide: bool = False) -> None: super().__init__(metric_name) self.inc = MetricUpdater( # type: ignore[method-assign] cells.CounterCell, metric_name, default_value=1, process_wide=process_wide) class DelegatingDistribution(Distribution): """Metrics Distribution Delegates functionality to MetricsEnvironment.""" def __init__( self, metric_name: MetricName, process_wide: bool = False) -> None: super().__init__(metric_name) self.update = MetricUpdater(cells.DistributionCell, metric_name, process_wide=process_wide) # type: ignore[method-assign] class DelegatingGauge(Gauge): """Metrics Gauge that Delegates functionality to MetricsEnvironment.""" def __init__( self, metric_name: MetricName, process_wide: bool = False) -> None: super().__init__(metric_name) self.set = MetricUpdater( # type: ignore[method-assign] cells.GaugeCell, metric_name, process_wide=process_wide) class DelegatingStringSet(StringSet): """Metrics StringSet that Delegates functionality to MetricsEnvironment.""" def __init__(self, metric_name: MetricName) -> None: super().__init__(metric_name) self.add = MetricUpdater(cells.StringSetCell, metric_name) # type: ignore[method-assign] class DelegatingBoundedTrie(BoundedTrie): """Metrics StringSet that Delegates functionality to MetricsEnvironment.""" def __init__(self, metric_name: MetricName) -> None: super().__init__(metric_name) self.add = MetricUpdater(cells.BoundedTrieCell, metric_name) # type: ignore[method-assign] class MetricResults(object): COUNTERS = "counters" DISTRIBUTIONS = "distributions" GAUGES = "gauges" STRINGSETS = "string_sets" BOUNDED_TRIES = "bounded_tries" HISTOGRAMS = "histograms" @staticmethod def _matches_name(filter: 'MetricsFilter', metric_key: 'MetricKey') -> bool: if ((filter.namespaces and metric_key.metric.namespace not in filter.namespaces) or (filter.names and metric_key.metric.name not in filter.names)): return False else: return True @staticmethod def _is_sub_list(needle: list[str], haystack: list[str]) -> bool: """True iff `needle` is a sub-list of `haystack` (i.e. a contiguous slice of `haystack` exactly matches `needle`""" needle_len = len(needle) haystack_len = len(haystack) for i in range(0, haystack_len - needle_len + 1): if haystack[i:i + needle_len] == needle: return True return False @staticmethod def _matches_sub_path(actual_scope: str, filter_scope: str) -> bool: """True iff the '/'-delimited pieces of filter_scope exist as a sub-list of the '/'-delimited pieces of actual_scope""" return bool( actual_scope and MetricResults._is_sub_list( filter_scope.split('/'), actual_scope.split('/'))) @staticmethod def _matches_scope(filter: 'MetricsFilter', metric_key: 'MetricKey') -> bool: if not filter.steps: return True for step in filter.steps: if MetricResults._matches_sub_path(metric_key.step, step): return True return False @staticmethod def matches( filter: Optional['MetricsFilter'], metric_key: 'MetricKey') -> bool: if filter is None: return True if (MetricResults._matches_name(filter, metric_key) and MetricResults._matches_scope(filter, metric_key)): return True return False def query( self, filter: Optional['MetricsFilter'] = None ) -> dict[str, list['MetricResult']]: """Queries the runner for existing user metrics that match the filter. It should return a dictionary, with lists of each kind of metric, and each list contains the corresponding kind of MetricResult. Like so: { "counters": [MetricResult(counter_key, committed, attempted), ...], "distributions": [MetricResult(dist_key, committed, attempted), ...], "gauges": [], // Empty list if nothing matched the filter. "string_sets": [] [MetricResult(string_set_key, committed, attempted), ...] } The committed / attempted values are DistributionResult / GaugeResult / int / set objects. """ raise NotImplementedError class MetricsFilter(object): """Simple object to filter metrics results. If filters by matching a result's step-namespace-name with three internal sets. No execution/matching logic is added to this object, so that it may be used to construct arguments as an RPC request. It is left for runners to implement matching logic by themselves. Note: This class only supports user defined metrics. """ def __init__(self) -> None: self._names: set[str] = set() self._namespaces: set[str] = set() self._steps: set[str] = set() @property def steps(self) -> frozenset[str]: return frozenset(self._steps) @property def names(self) -> frozenset[str]: return frozenset(self._names) @property def namespaces(self) -> frozenset[str]: return frozenset(self._namespaces) def with_metric(self, metric: 'Metric') -> 'MetricsFilter': name = metric.metric_name.name or '' namespace = metric.metric_name.namespace or '' return self.with_name(name).with_namespace(namespace) def with_name(self, name: str) -> 'MetricsFilter': return self.with_names([name]) def with_names(self, names: Iterable[str]) -> 'MetricsFilter': if isinstance(names, str): raise ValueError('Names must be a collection, not a string') self._names.update(names) return self def with_namespace(self, namespace: Union[type, str]) -> 'MetricsFilter': return self.with_namespaces([namespace]) def with_namespaces( self, namespaces: Iterable[Union[type, str]]) -> 'MetricsFilter': if isinstance(namespaces, str): raise ValueError('Namespaces must be an iterable, not a string') self._namespaces.update([Metrics.get_namespace(ns) for ns in namespaces]) return self def with_step(self, step: str) -> 'MetricsFilter': return self.with_steps([step]) def with_steps(self, steps: Iterable[str]) -> 'MetricsFilter': if isinstance(steps, str): raise ValueError('Steps must be an iterable, not a string') self._steps.update(steps) return self class Lineage: """Standard collection of metrics used to record source and sinks information for lineage tracking.""" LINEAGE_NAMESPACE = "lineage" SOURCE = "sources_v2" SINK = "sinks_v2" _METRICS = { SOURCE: Metrics.bounded_trie(LINEAGE_NAMESPACE, SOURCE), SINK: Metrics.bounded_trie(LINEAGE_NAMESPACE, SINK) } def __init__(self, label: str) -> None: """Create a Lineage with valid label (:data:`~Lineage.SOURCE` or :data:`~Lineage.SINK`) """ self.metric = Lineage._METRICS[label] @classmethod def sources(cls) -> 'Lineage': return cls(Lineage.SOURCE) @classmethod def sinks(cls) -> 'Lineage': return cls(Lineage.SINK) _RESERVED_CHARS = re.compile(r'[:\s.]') @staticmethod def wrap_segment(segment: str) -> str: """Wrap segment to valid segment name. Specifically, If there are reserved chars (colon, whitespace, dot), escape with backtick. If the segment is already wrapped, return the original. """ if segment.startswith("`") and segment.endswith("`"): return segment if Lineage._RESERVED_CHARS.search(segment): return "`" + segment + "`" return segment @staticmethod def get_fq_name( system: str, *segments: str, subtype: Optional[str] = None) -> str: """Assemble fully qualified name (`FQN `_). Format: - `system:segment1.segment2` - `system:subtype:segment1.segment2` - `system:`segment1.with.dots:colons`.segment2` This helper method is for internal and testing usage only. """ segs = '.'.join(map(Lineage.wrap_segment, segments)) if subtype: return ':'.join((system, subtype, segs)) return ':'.join((system, segs)) @staticmethod def _get_fqn_parts( system: str, *segments: str, subtype: Optional[str] = None, last_segment_sep: Optional[str] = None) -> Iterator[str]: yield system + ':' if subtype: yield subtype + ':' if segments: for segment in segments[:-1]: yield segment + '.' if last_segment_sep: sub_segments = segments[-1].split(last_segment_sep) for sub_segment in sub_segments[:-1]: yield sub_segment + last_segment_sep yield sub_segments[-1] else: yield segments[-1] def add( self, system: str, *segments: str, subtype: Optional[str] = None, last_segment_sep: Optional[str] = None) -> None: """ Adds the given details as Lineage. For asset level lineage the resource location should be specified as Dataplex FQN, see https://cloud.google.com/data-catalog/docs/fully-qualified-names Example of adding FQN components: - `add("system", "segment1", "segment2")` - `add("system", "segment1", "segment2", subtype="subtype")` Example of adding a FQN: - `add("system:segment1.segment2")` - `add("system:subtype:segment1.segment2")` The first positional argument serves as system, if full segments are provided, or the full FQN if it is provided as a single argument. """ self.add_raw( *self._get_fqn_parts( system, *segments, subtype=subtype, last_segment_sep=last_segment_sep)) def add_raw(self, *rollup_segments: str) -> None: """Adds the given fqn as lineage. `rollup_segments` should be an iterable of strings whose concatenation is a valid Dataplex FQN. In particular, this means they will often have trailing delimiters. """ self.metric.add(rollup_segments) @staticmethod def query(results: MetricResults, label: str, truncated_marker: str = '*') -> set[str]: if not label in Lineage._METRICS: raise ValueError("Label {} does not exist for Lineage", label) response = results.query( MetricsFilter().with_namespace(Lineage.LINEAGE_NAMESPACE).with_name( label))[MetricResults.BOUNDED_TRIES] result = set() for metric in response: for fqn in metric.committed.flattened(): result.add(''.join(fqn[:-1]) + (truncated_marker if fqn[-1] else '')) for fqn in metric.attempted.flattened(): result.add(''.join(fqn[:-1]) + (truncated_marker if fqn[-1] else '')) return result ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/metrics/metric_test.py0000644000175100017510000002712115206113032022222 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import re import unittest import hamcrest as hc import pytest import apache_beam as beam from apache_beam import metrics from apache_beam.metrics.cells import DistributionData from apache_beam.metrics.execution import MetricKey from apache_beam.metrics.execution import MetricsContainer from apache_beam.metrics.execution import MetricsEnvironment from apache_beam.metrics.metric import Lineage from apache_beam.metrics.metric import MetricResults from apache_beam.metrics.metric import Metrics from apache_beam.metrics.metric import MetricsFilter from apache_beam.metrics.metricbase import MetricName from apache_beam.runners.direct.direct_runner import BundleBasedDirectRunner from apache_beam.runners.worker import statesampler from apache_beam.testing.metric_result_matchers import DistributionMatcher from apache_beam.testing.metric_result_matchers import MetricResultMatcher from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.utils import counters from apache_beam.utils.histogram import LinearBucket class NameTest(unittest.TestCase): def test_basic_metric_name(self): name = MetricName('namespace1', 'name1') self.assertEqual(name.namespace, 'namespace1') self.assertEqual(name.name, 'name1') self.assertEqual(name, MetricName('namespace1', 'name1')) key = MetricKey('step1', name) self.assertEqual(key.step, 'step1') self.assertEqual(key.metric.namespace, 'namespace1') self.assertEqual(key.metric.name, 'name1') self.assertEqual(key, MetricKey('step1', MetricName('namespace1', 'name1'))) class MetricResultsTest(unittest.TestCase): def test_metric_filter_namespace_matching(self): filter = MetricsFilter().with_namespace('ns1') name = MetricName('ns1', 'name1') key = MetricKey('step1', name) self.assertTrue(MetricResults.matches(filter, key)) def test_metric_filter_name_matching(self): filter = MetricsFilter().with_name('name1').with_namespace('ns1') name = MetricName('ns1', 'name1') key = MetricKey('step1', name) self.assertTrue(MetricResults.matches(filter, key)) filter = MetricsFilter().with_name('name1') name = MetricName('ns1', 'name1') key = MetricKey('step1', name) self.assertTrue(MetricResults.matches(filter, key)) def test_metric_filter_step_matching(self): name = MetricName('ns1', 'name1') filter = MetricsFilter().with_step('Step1') key = MetricKey('Step1', name) self.assertTrue(MetricResults.matches(filter, key)) key = MetricKey('Step10', name) self.assertFalse(MetricResults.matches(filter, key)) key = MetricKey('Step10/Step1', name) self.assertTrue(MetricResults.matches(filter, key)) key = MetricKey('Top1/Outer1/Inner1', name) filter = MetricsFilter().with_step('Top1/Outer1/Inner1') self.assertTrue(MetricResults.matches(filter, key)) filter = MetricsFilter().with_step('Top1/Outer1') self.assertTrue(MetricResults.matches(filter, key)) filter = MetricsFilter().with_step('Outer1/Inner1') self.assertTrue(MetricResults.matches(filter, key)) filter = MetricsFilter().with_step('Top1/Inner1') self.assertFalse(MetricResults.matches(filter, key)) class MetricsTest(unittest.TestCase): def test_get_namespace_class(self): class MyClass(object): pass self.assertEqual( '{}.{}'.format(MyClass.__module__, MyClass.__name__), Metrics.get_namespace(MyClass)) def test_get_namespace_string(self): namespace = 'MyNamespace' self.assertEqual(namespace, Metrics.get_namespace(namespace)) def test_get_namespace_error(self): with self.assertRaises(ValueError): Metrics.get_namespace(object()) def test_counter_empty_name(self): with self.assertRaises(ValueError): Metrics.counter("namespace", "") def test_counter_empty_namespace(self): with self.assertRaises(ValueError): Metrics.counter("", "names") def test_distribution_empty_name(self): with self.assertRaises(ValueError): Metrics.distribution("namespace", "") def test_distribution_empty_namespace(self): with self.assertRaises(ValueError): Metrics.distribution("", "names") # Do not change the behaviour of str(), do tno update/delete this test case # if the behaviour of str() is changed. Doing so will # break end user beam code which depends on the str() behaviour. def test_user_metric_name_str(self): mn = MetricName("my_namespace", "my_name") expected_str = 'MetricName(namespace=my_namespace, name=my_name)' self.assertEqual(str(mn), expected_str) def test_general_urn_metric_name_str(self): mn = MetricName( "my_namespace", "my_name", urn='my_urn', labels={'key': 'value'}) expected_str = ( "MetricName(namespace=my_namespace, name=my_name, " "urn=my_urn, labels={'key': 'value'})") self.assertEqual(str(mn), expected_str) @pytest.mark.it_validatesrunner def test_user_counter_using_pardo(self): class SomeDoFn(beam.DoFn): """A custom dummy DoFn using yield.""" static_counter_elements = metrics.Metrics.counter( "SomeDoFn", 'metrics_static_counter_element') def __init__(self): self.user_counter_elements = metrics.Metrics.counter( self.__class__, 'metrics_user_counter_element') def process(self, element): self.static_counter_elements.inc(2) self.user_counter_elements.inc() distro = Metrics.distribution(self.__class__, 'element_dist') distro.update(element) yield element pipeline = TestPipeline() nums = pipeline | 'Input' >> beam.Create([1, 2, 3, 4]) results = nums | 'ApplyPardo' >> beam.ParDo(SomeDoFn()) assert_that(results, equal_to([1, 2, 3, 4])) res = pipeline.run() res.wait_until_finish() # Verify static counter. metric_results = ( res.metrics().query( MetricsFilter().with_metric(SomeDoFn.static_counter_elements))) outputs_static_counter = metric_results['counters'][0] self.assertEqual( outputs_static_counter.key.metric.name, 'metrics_static_counter_element') self.assertEqual(outputs_static_counter.committed, 8) # Verify user counter. metric_results = ( res.metrics().query( MetricsFilter().with_name('metrics_user_counter_element'))) outputs_user_counter = metric_results['counters'][0] self.assertEqual( outputs_user_counter.key.metric.name, 'metrics_user_counter_element') self.assertEqual(outputs_user_counter.committed, 4) # Verify user distribution counter. metric_results = res.metrics().query() matcher = MetricResultMatcher( step=hc.contains_string('ApplyPardo'), namespace=hc.contains_string('SomeDoFn'), name='element_dist', committed=DistributionMatcher( sum_value=hc.greater_than_or_equal_to(0), count_value=hc.greater_than_or_equal_to(0), min_value=hc.greater_than_or_equal_to(0), max_value=hc.greater_than_or_equal_to(0))) hc.assert_that(metric_results['distributions'], hc.has_item(matcher)) def test_create_counter_distribution(self): sampler = statesampler.StateSampler('', counters.CounterFactory()) statesampler.set_current_tracker(sampler) state1 = sampler.scoped_state( 'mystep', 'myState', metrics_container=MetricsContainer('mystep')) try: sampler.start() with state1: counter_ns = 'aCounterNamespace' distro_ns = 'aDistributionNamespace' name = 'a_name' counter = Metrics.counter(counter_ns, name) distro = Metrics.distribution(distro_ns, name) counter.inc(10) counter.dec(3) distro.update(10) distro.update(2) self.assertTrue(isinstance(counter, Metrics.DelegatingCounter)) self.assertTrue(isinstance(distro, Metrics.DelegatingDistribution)) del distro del counter container = MetricsEnvironment.current_container() self.assertEqual( container.get_counter(MetricName(counter_ns, name)).get_cumulative(), 7) self.assertEqual( container.get_distribution(MetricName(distro_ns, name)).get_cumulative(), DistributionData(12, 2, 2, 10)) finally: sampler.stop() class LineageTest(unittest.TestCase): def test_fq_name(self): test_cases = { "apache-beam": "apache-beam", "`apache-beam`": "`apache-beam`", "apache.beam": "`apache.beam`", "apache:beam": "`apache:beam`", "apache beam": "`apache beam`", "`apache beam`": "`apache beam`", "apache\tbeam": "`apache\tbeam`", "apache\nbeam": "`apache\nbeam`" } for k, v in test_cases.items(): self.assertEqual("apache:" + v, Lineage.get_fq_name("apache", k)) self.assertEqual( "apache:beam:" + v, Lineage.get_fq_name("apache", k, subtype="beam")) self.assertEqual( "apache:beam:" + v + '.' + v, Lineage.get_fq_name("apache", k, k, subtype="beam")) def test_add(self): lineage = Lineage(Lineage.SOURCE) added = set() # override lineage.metric = added lineage.add("s", "1", "2") lineage.add("s:3.4") lineage.add("s", "5", "6.7") lineage.add("s", "1", "2", subtype="t") lineage.add("sys", "seg1", "seg2", "seg3/part2/part3", last_segment_sep='/') self.assertSetEqual( added, {('s:', '1.', '2'), ('s:3.4:', ), ('s:', '5.', '6.7'), ('s:', 't:', '1.', '2'), ('sys:', 'seg1.', 'seg2.', 'seg3/', 'part2/', 'part3')}) class HistogramTest(unittest.TestCase): def test_histogram(self): class WordExtractingDoFn(beam.DoFn): def __init__(self): self.word_lengths_dist = Metrics.histogram( self.__class__, 'latency_histogram_ms', LinearBucket(0, 1, num_buckets=10)) def process(self, element): text_line = element.strip() words = re.findall(r'[\w\']+', text_line, re.UNICODE) for w in words: self.word_lengths_dist.update(len(w)) return words with beam.Pipeline(runner=BundleBasedDirectRunner()) as p: lines = p | 'read' >> beam.Create(["x x x yyyyyy yyyyyy yyyyyy"]) _ = ( lines | 'split' >> (beam.ParDo(WordExtractingDoFn()).with_output_types(str))) result = p.result filter = MetricsFilter().with_name('latency_histogram_ms') query_result = result.metrics().query(filter) histogram = query_result['histograms'][0].committed.histogram assert histogram._buckets == {1: 3, 6: 3} assert histogram.total_count() == 6 assert 1 < histogram.get_linear_interpolation(0.50) < 3 assert histogram.get_linear_interpolation(0.99) > 3 if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/metrics/metricbase.py0000644000175100017510000001207415206113032022017 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ The classes in this file are interfaces for metrics. They are not intended to be subclassed or created directly by users. To work with and access metrics, users should use the classes and methods exposed in metric.py. Available classes: - Metric - Base interface of a metrics object. - Counter - Counter metric interface. Allows a count to be incremented or decremented during pipeline execution. - Distribution - Distribution Metric interface. Allows statistics about the distribution of a variable to be collected during pipeline execution. - Gauge - Gauge Metric interface. Allows to track the latest value of a variable during pipeline execution. - MetricName - Namespace and name used to refer to a Metric. """ # pytype: skip-file from typing import Optional __all__ = [ 'Metric', 'Counter', 'Distribution', 'Gauge', 'StringSet', 'BoundedTrie', 'Histogram', 'MetricName' ] class MetricName(object): """The name of a metric. The name of a metric consists of a namespace and a name. The namespace allows grouping related metrics together and also prevents collisions between multiple metrics of the same name. """ def __init__( self, namespace: Optional[str], name: Optional[str], urn: Optional[str] = None, labels: Optional[dict[str, str]] = None) -> None: """Initializes ``MetricName``. Note: namespace and name should be set for user metrics, urn and labels should be set for an arbitrary metric to package into a MonitoringInfo. Args: namespace: A string with the namespace of a metric. name: A string with the name of a metric. urn: URN to populate on a MonitoringInfo, when sending to RunnerHarness. labels: Labels to populate on a MonitoringInfo """ if not urn: if not namespace: raise ValueError('Metric namespace must be non-empty') if not name: raise ValueError('Metric name must be non-empty') self.namespace = namespace self.name = name self.urn = urn self.labels = labels if labels else {} def __eq__(self, other): return ( self.namespace == other.namespace and self.name == other.name and self.urn == other.urn and self.labels == other.labels) def __str__(self): if self.urn: return 'MetricName(namespace={}, name={}, urn={}, labels={})'.format( self.namespace, self.name, self.urn, self.labels) else: # User counter case. return 'MetricName(namespace={}, name={})'.format( self.namespace, self.name) def __hash__(self): return hash((self.namespace, self.name, self.urn) + tuple(self.labels.items())) def fast_name(self): name = self.name or '' namespace = self.namespace or '' urn = self.urn or '' labels = '' if self.labels: labels = '_'.join(['%s=%s' % (k, v) for (k, v) in self.labels.items()]) return '%d_%s%s%s%s' % (len(name), name, namespace, urn, labels) class Metric(object): """Base interface of a metric object.""" def __init__(self, metric_name: MetricName) -> None: self.metric_name = metric_name class Counter(Metric): """Counter metric interface. Allows a count to be incremented/decremented during pipeline execution.""" def inc(self, n=1): raise NotImplementedError def dec(self, n=1): self.inc(-n) class Distribution(Metric): """Distribution Metric interface. Allows statistics about the distribution of a variable to be collected during pipeline execution.""" def update(self, value): raise NotImplementedError class Gauge(Metric): """Gauge Metric interface. Allows tracking of the latest value of a variable during pipeline execution.""" def set(self, value): raise NotImplementedError class StringSet(Metric): """StringSet Metric interface. Reports set of unique string values during pipeline execution..""" def add(self, value): raise NotImplementedError class BoundedTrie(Metric): """BoundedTrie Metric interface. Reports set of unique string values during pipeline execution..""" def add(self, value): raise NotImplementedError class Histogram(Metric): """Histogram Metric interface. Allows statistics about the percentile of a variable to be collected during pipeline execution.""" def update(self, value): raise NotImplementedError ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/metrics/monitoring_infos.py0000644000175100017510000005312615206113032023267 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import collections import time from functools import reduce from typing import Hashable from typing import Union from apache_beam.coders import coder_impl from apache_beam.coders import coders from apache_beam.metrics.cells import BoundedTrieData from apache_beam.metrics.cells import DistributionData from apache_beam.metrics.cells import DistributionResult from apache_beam.metrics.cells import GaugeData from apache_beam.metrics.cells import GaugeResult from apache_beam.metrics.cells import HistogramData from apache_beam.metrics.cells import StringSetData from apache_beam.portability import common_urns from apache_beam.portability.api import metrics_pb2 SAMPLED_BYTE_SIZE_URN = ( common_urns.monitoring_info_specs.SAMPLED_BYTE_SIZE.spec.urn) ELEMENT_COUNT_URN = common_urns.monitoring_info_specs.ELEMENT_COUNT.spec.urn START_BUNDLE_MSECS_URN = ( common_urns.monitoring_info_specs.START_BUNDLE_MSECS.spec.urn) PROCESS_BUNDLE_MSECS_URN = ( common_urns.monitoring_info_specs.PROCESS_BUNDLE_MSECS.spec.urn) FINISH_BUNDLE_MSECS_URN = ( common_urns.monitoring_info_specs.FINISH_BUNDLE_MSECS.spec.urn) TOTAL_MSECS_URN = common_urns.monitoring_info_specs.TOTAL_MSECS.spec.urn USER_COUNTER_URN = common_urns.monitoring_info_specs.USER_SUM_INT64.spec.urn USER_HISTOGRAM_URN = common_urns.monitoring_info_specs.USER_HISTOGRAM.spec.urn USER_DISTRIBUTION_URN = ( common_urns.monitoring_info_specs.USER_DISTRIBUTION_INT64.spec.urn) USER_GAUGE_URN = common_urns.monitoring_info_specs.USER_LATEST_INT64.spec.urn USER_STRING_SET_URN = common_urns.monitoring_info_specs.USER_SET_STRING.spec.urn USER_BOUNDED_TRIE_URN = ( common_urns.monitoring_info_specs.USER_BOUNDED_TRIE.spec.urn) USER_METRIC_URNS = set([ USER_COUNTER_URN, USER_DISTRIBUTION_URN, USER_GAUGE_URN, USER_STRING_SET_URN, USER_BOUNDED_TRIE_URN, USER_HISTOGRAM_URN ]) WORK_REMAINING_URN = common_urns.monitoring_info_specs.WORK_REMAINING.spec.urn WORK_COMPLETED_URN = common_urns.monitoring_info_specs.WORK_COMPLETED.spec.urn DATA_CHANNEL_READ_INDEX = ( common_urns.monitoring_info_specs.DATA_CHANNEL_READ_INDEX.spec.urn) API_REQUEST_COUNT_URN = ( common_urns.monitoring_info_specs.API_REQUEST_COUNT.spec.urn) # TODO(ajamato): Implement the remaining types, i.e. Double types # Extrema types, etc. See: # https://s.apache.org/beam-fn-api-metrics SUM_INT64_TYPE = common_urns.monitoring_info_types.SUM_INT64_TYPE.urn DISTRIBUTION_INT64_TYPE = ( common_urns.monitoring_info_types.DISTRIBUTION_INT64_TYPE.urn) LATEST_INT64_TYPE = common_urns.monitoring_info_types.LATEST_INT64_TYPE.urn PROGRESS_TYPE = common_urns.monitoring_info_types.PROGRESS_TYPE.urn STRING_SET_TYPE = common_urns.monitoring_info_types.SET_STRING_TYPE.urn BOUNDED_TRIE_TYPE = common_urns.monitoring_info_types.BOUNDED_TRIE_TYPE.urn HISTOGRAM_TYPE = common_urns.monitoring_info_types.HISTOGRAM.urn COUNTER_TYPES = set([SUM_INT64_TYPE]) DISTRIBUTION_TYPES = set([DISTRIBUTION_INT64_TYPE]) GAUGE_TYPES = set([LATEST_INT64_TYPE]) STRING_SET_TYPES = set([STRING_SET_TYPE]) BOUNDED_TRIE_TYPES = set([BOUNDED_TRIE_TYPE]) HISTOGRAM_TYPES = set([HISTOGRAM_TYPE]) # TODO(migryz) extract values from beam_fn_api.proto::MonitoringInfoLabels PCOLLECTION_LABEL = ( common_urns.monitoring_info_labels.PCOLLECTION.label_props.name) PTRANSFORM_LABEL = ( common_urns.monitoring_info_labels.TRANSFORM.label_props.name) NAMESPACE_LABEL = ( common_urns.monitoring_info_labels.NAMESPACE.label_props.name) NAME_LABEL = (common_urns.monitoring_info_labels.NAME.label_props.name) SERVICE_LABEL = (common_urns.monitoring_info_labels.SERVICE.label_props.name) METHOD_LABEL = (common_urns.monitoring_info_labels.METHOD.label_props.name) RESOURCE_LABEL = (common_urns.monitoring_info_labels.RESOURCE.label_props.name) STATUS_LABEL = (common_urns.monitoring_info_labels.STATUS.label_props.name) BIGQUERY_PROJECT_ID_LABEL = ( common_urns.monitoring_info_labels.BIGQUERY_PROJECT_ID.label_props.name) BIGQUERY_DATASET_LABEL = ( common_urns.monitoring_info_labels.BIGQUERY_DATASET.label_props.name) BIGQUERY_TABLE_LABEL = ( common_urns.monitoring_info_labels.BIGQUERY_TABLE.label_props.name) BIGQUERY_VIEW_LABEL = ( common_urns.monitoring_info_labels.BIGQUERY_VIEW.label_props.name) BIGQUERY_QUERY_NAME_LABEL = ( common_urns.monitoring_info_labels.BIGQUERY_QUERY_NAME.label_props.name) GCS_PROJECT_ID_LABEL = ( common_urns.monitoring_info_labels.GCS_PROJECT_ID.label_props.name) GCS_BUCKET_LABEL = ( common_urns.monitoring_info_labels.GCS_BUCKET.label_props.name) DATASTORE_PROJECT_ID_LABEL = ( common_urns.monitoring_info_labels.DATASTORE_PROJECT.label_props.name) DATASTORE_NAMESPACE_LABEL = ( common_urns.monitoring_info_labels.DATASTORE_NAMESPACE.label_props.name) SPANNER_PROJECT_ID = ( common_urns.monitoring_info_labels.SPANNER_PROJECT_ID.label_props.name) SPANNER_DATABASE_ID = ( common_urns.monitoring_info_labels.SPANNER_DATABASE_ID.label_props.name) SPANNER_TABLE_ID = ( common_urns.monitoring_info_labels.SPANNER_TABLE_ID.label_props.name) SPANNER_QUERY_NAME = ( common_urns.monitoring_info_labels.SPANNER_QUERY_NAME.label_props.name) BIGTABLE_PROJECT_ID_LABEL = ( common_urns.monitoring_info_labels.BIGTABLE_PROJECT_ID.label_props.name) INSTANCE_ID_LABEL = ( common_urns.monitoring_info_labels.INSTANCE_ID.label_props.name) TABLE_ID_LABEL = common_urns.monitoring_info_labels.TABLE_ID.label_props.name def extract_counter_value(monitoring_info_proto): """Returns the counter value of the monitoring info.""" if not is_counter(monitoring_info_proto): raise ValueError('Unsupported type %s' % monitoring_info_proto.type) # Only SUM_INT64_TYPE is currently supported. return coders.VarIntCoder().decode(monitoring_info_proto.payload) def extract_gauge_value(monitoring_info_proto): """Returns a tuple containing (timestamp, value)""" if not is_gauge(monitoring_info_proto): raise ValueError('Unsupported type %s' % monitoring_info_proto.type) # Only LATEST_INT64_TYPE is currently supported. return _decode_gauge(coders.VarIntCoder(), monitoring_info_proto.payload) def extract_distribution(monitoring_info_proto): """Returns a tuple of (count, sum, min, max). Args: proto: The monitoring info for the distribution. """ if not is_distribution(monitoring_info_proto): raise ValueError('Unsupported type %s' % monitoring_info_proto.type) # Only DISTRIBUTION_INT64_TYPE is currently supported. return _decode_distribution( coders.VarIntCoder(), monitoring_info_proto.payload) def extract_string_set_value(monitoring_info_proto): if not is_string_set(monitoring_info_proto): raise ValueError('Unsupported type %s' % monitoring_info_proto.type) coder = coders.IterableCoder(coders.StrUtf8Coder()) return set(coder.decode(monitoring_info_proto.payload)) def extract_bounded_trie_value(monitoring_info_proto): if not is_bounded_trie(monitoring_info_proto): raise ValueError('Unsupported type %s' % monitoring_info_proto.type) return BoundedTrieData.from_proto( metrics_pb2.BoundedTrie.FromString(monitoring_info_proto.payload)) def extract_histogram_value(monitoring_info_proto): if not is_histogram(monitoring_info_proto): raise ValueError('Unsupported type %s' % monitoring_info_proto.type) return HistogramData.from_proto( metrics_pb2.HistogramValue.FromString(monitoring_info_proto.payload)) def create_labels(ptransform=None, namespace=None, name=None, pcollection=None): """Create the label dictionary based on the provided values. Args: ptransform: The ptransform id used as a label. pcollection: The pcollection id used as a label. """ labels = {} if ptransform: labels[PTRANSFORM_LABEL] = ptransform if namespace: labels[NAMESPACE_LABEL] = namespace if name: labels[NAME_LABEL] = name if pcollection: labels[PCOLLECTION_LABEL] = pcollection return labels def int64_user_counter( namespace, name, metric, ptransform=None) -> metrics_pb2.MonitoringInfo: """Return the counter monitoring info for the specifed URN, metric and labels. Args: urn: The URN of the monitoring info/metric. metric: The payload field to use in the monitoring info or an int value. ptransform: The ptransform id used as a label. """ labels = create_labels(ptransform=ptransform, namespace=namespace, name=name) if isinstance(metric, int): metric = coders.VarIntCoder().encode(metric) return create_monitoring_info( USER_COUNTER_URN, SUM_INT64_TYPE, metric, labels) def int64_counter( urn, metric, ptransform=None, pcollection=None, labels=None) -> metrics_pb2.MonitoringInfo: """Return the counter monitoring info for the specifed URN, metric and labels. Args: urn: The URN of the monitoring info/metric. metric: The payload field to use in the monitoring info or an int value. ptransform: The ptransform id used as a label. pcollection: The pcollection id used as a label. """ labels = labels or {} labels.update(create_labels(ptransform=ptransform, pcollection=pcollection)) if isinstance(metric, int): metric = coders.VarIntCoder().encode(metric) return create_monitoring_info(urn, SUM_INT64_TYPE, metric, labels) def int64_user_distribution( namespace, name, metric, ptransform=None) -> metrics_pb2.MonitoringInfo: """Return the distribution monitoring info for the URN, metric and labels. Args: urn: The URN of the monitoring info/metric. metric: The DistributionData for the metric. ptransform: The ptransform id used as a label. """ labels = create_labels(ptransform=ptransform, namespace=namespace, name=name) payload = _encode_distribution( coders.VarIntCoder(), metric.count, metric.sum, metric.min, metric.max) return create_monitoring_info( USER_DISTRIBUTION_URN, DISTRIBUTION_INT64_TYPE, payload, labels) def int64_distribution( urn, metric, ptransform=None, pcollection=None) -> metrics_pb2.MonitoringInfo: """Return a distribution monitoring info for the URN, metric and labels. Args: urn: The URN of the monitoring info/metric. metric: The DistributionData for the metric. ptransform: The ptransform id used as a label. pcollection: The pcollection id used as a label. """ labels = create_labels(ptransform=ptransform, pcollection=pcollection) payload = _encode_distribution( coders.VarIntCoder(), metric.count, metric.sum, metric.min, metric.max) return create_monitoring_info(urn, DISTRIBUTION_INT64_TYPE, payload, labels) def int64_user_gauge( namespace, name, metric, ptransform=None) -> metrics_pb2.MonitoringInfo: """Return the gauge monitoring info for the URN, metric and labels. Args: namespace: User-defined namespace of gauge metric. name: Name of gauge metric. metric: The GaugeData containing the metrics. ptransform: The ptransform id used as a label. """ labels = create_labels(ptransform=ptransform, namespace=namespace, name=name) if isinstance(metric, GaugeData): coder = coders.VarIntCoder() value = metric.value timestamp = metric.timestamp else: raise TypeError( 'Expected GaugeData metric type but received %s with value %s' % (type(metric), metric)) payload = _encode_gauge(coder, timestamp, value) return create_monitoring_info( USER_GAUGE_URN, LATEST_INT64_TYPE, payload, labels) def int64_gauge(urn, metric, ptransform=None) -> metrics_pb2.MonitoringInfo: """Return the gauge monitoring info for the URN, metric and labels. Args: urn: The URN of the monitoring info/metric. metric: An int representing the value. The current time will be used for the timestamp. ptransform: The ptransform id used as a label. """ labels = create_labels(ptransform=ptransform) if isinstance(metric, int): value = metric time_ms = int(time.time()) * 1000 else: raise TypeError( 'Expected int metric type but received %s with value %s' % (type(metric), metric)) coder = coders.VarIntCoder() payload = coder.encode(time_ms) + coder.encode(value) return create_monitoring_info(urn, LATEST_INT64_TYPE, payload, labels) def user_set_string(namespace, name, metric, ptransform=None): """Return the string set monitoring info for the URN, metric and labels. Args: namespace: User-defined namespace of StringSet. name: Name of StringSet. metric: The StringSetData representing the metrics. ptransform: The ptransform id used as a label. """ labels = create_labels(ptransform=ptransform, namespace=namespace, name=name) if isinstance(metric, StringSetData): metric = metric.string_set if isinstance(metric, set): metric = list(metric) if isinstance(metric, list): metric = coders.IterableCoder(coders.StrUtf8Coder()).encode(metric) return create_monitoring_info( USER_STRING_SET_URN, STRING_SET_TYPE, metric, labels) def user_histogram(namespace, name, metric: HistogramData, ptransform=None): """Return the histogram monitoring info for the URN, metric and labels. Args: namespace: User-defined namespace of Histogram. name: Name of Histogram. metric: The Histogram representing the metrics. ptransform: The ptransform id used as a label. """ labels = create_labels(ptransform=ptransform, namespace=namespace, name=name) metric_proto = metric.to_proto() return create_monitoring_info( USER_HISTOGRAM_URN, HISTOGRAM_TYPE, metric_proto.SerializeToString(), labels) def user_bounded_trie(namespace, name, metric, ptransform=None): """Return the string set monitoring info for the URN, metric and labels. Args: namespace: User-defined namespace of BoundedTrie. name: Name of BoundedTrie. metric: The BoundedTrieData representing the metrics. ptransform: The ptransform id used as a label. """ labels = create_labels(ptransform=ptransform, namespace=namespace, name=name) return create_monitoring_info( USER_BOUNDED_TRIE_URN, BOUNDED_TRIE_TYPE, metric.to_proto().SerializeToString(), labels) def create_monitoring_info( urn, type_urn, payload, labels=None) -> metrics_pb2.MonitoringInfo: """Return the monitoring info for the URN, type, metric and labels. Args: urn: The URN of the monitoring info/metric. type_urn: The URN of the type of the monitoring info/metric. i.e. beam:metrics:sum_int_64, beam:metrics:latest_int_64. payload: The payload field to use in the monitoring info. labels: The label dictionary to use in the MonitoringInfo. """ try: return metrics_pb2.MonitoringInfo( urn=urn, type=type_urn, labels=labels or {}, payload=payload) except TypeError as e: raise RuntimeError( f'Failed to create MonitoringInfo for urn {urn} type {type_urn} ' f'labels {labels} and payload {payload}') from e def is_counter(monitoring_info_proto): """Returns true if the monitoring info is a coutner metric.""" return monitoring_info_proto.type in COUNTER_TYPES def is_gauge(monitoring_info_proto): """Returns true if the monitoring info is a gauge metric.""" return monitoring_info_proto.type in GAUGE_TYPES def is_distribution(monitoring_info_proto): """Returns true if the monitoring info is a distrbution metric.""" return monitoring_info_proto.type in DISTRIBUTION_TYPES def is_histogram(monitoring_info_proto): """Returns true if the monitoring info is a distrbution metric.""" return monitoring_info_proto.type in HISTOGRAM_TYPES def is_string_set(monitoring_info_proto): """Returns true if the monitoring info is a StringSet metric.""" return monitoring_info_proto.type in STRING_SET_TYPES def is_bounded_trie(monitoring_info_proto): """Returns true if the monitoring info is a BoundedTrie metric.""" return monitoring_info_proto.type in BOUNDED_TRIE_TYPES def is_user_monitoring_info(monitoring_info_proto): """Returns true if the monitoring info is a user metric.""" return monitoring_info_proto.urn in USER_METRIC_URNS def extract_metric_result_map_value( monitoring_info_proto ) -> Union[None, int, DistributionResult, GaugeResult, set, BoundedTrieData]: """Returns the relevant GaugeResult, DistributionResult or int value for counter metric, set for StringSet metric. These are the proper format for use in the MetricResult.query() result. """ # Returns a metric result (AKA the legacy format). # from the MonitoringInfo if is_counter(monitoring_info_proto): return extract_counter_value(monitoring_info_proto) if is_distribution(monitoring_info_proto): (count, sum, min, max) = extract_distribution(monitoring_info_proto) return DistributionResult(DistributionData(sum, count, min, max)) if is_gauge(monitoring_info_proto): (timestamp, value) = extract_gauge_value(monitoring_info_proto) return GaugeResult(GaugeData(value, timestamp)) if is_string_set(monitoring_info_proto): return extract_string_set_value(monitoring_info_proto) if is_bounded_trie(monitoring_info_proto): return extract_bounded_trie_value(monitoring_info_proto) return None def parse_namespace_and_name(monitoring_info_proto): """Returns the (namespace, name) tuple of the URN in the monitoring info.""" # Remove the URN prefix which indicates that it is a user counter. if is_user_monitoring_info(monitoring_info_proto): labels = monitoring_info_proto.labels return labels[NAMESPACE_LABEL], labels[NAME_LABEL] # If it is not a user counter, just use the first part of the URN, i.e. 'beam' split = monitoring_info_proto.urn.split(':', 1) return split[0], split[1] def get_step_name(monitoring_info_proto): """Returns a step name for the given monitoring info or None if step name cannot be specified.""" # Right now only metrics that have a PTRANSFORM are taken into account return monitoring_info_proto.labels.get(PTRANSFORM_LABEL) def to_key( monitoring_info_proto: metrics_pb2.MonitoringInfo) -> frozenset[Hashable]: """Returns a key based on the URN and labels. This is useful in maps to prevent reporting the same MonitoringInfo twice. """ key_items: list[Hashable] = list(monitoring_info_proto.labels.items()) key_items.append(monitoring_info_proto.urn) return frozenset(key_items) def sum_payload_combiner(payload_a, payload_b): coder = coders.VarIntCoder() return coder.encode(coder.decode(payload_a) + coder.decode(payload_b)) def distribution_payload_combiner(payload_a, payload_b): coder = coders.VarIntCoder() (count_a, sum_a, min_a, max_a) = _decode_distribution(coder, payload_a) (count_b, sum_b, min_b, max_b) = _decode_distribution(coder, payload_b) return _encode_distribution( coder, count_a + count_b, sum_a + sum_b, min(min_a, min_b), max(max_a, max_b)) _KNOWN_COMBINERS = { SUM_INT64_TYPE: sum_payload_combiner, DISTRIBUTION_INT64_TYPE: distribution_payload_combiner, } def consolidate(metrics, key=to_key): grouped = collections.defaultdict(list) for metric in metrics: grouped[key(metric)].append(metric) for values in grouped.values(): if len(values) == 1: yield values[0] else: combiner = _KNOWN_COMBINERS.get(values[0].type) if combiner: def merge(a, b): # pylint: disable=cell-var-from-loop return metrics_pb2.MonitoringInfo( urn=a.urn, type=a.type, labels=dict((label, value) for label, value in a.labels.items() if b.labels.get(label) == value), payload=combiner(a.payload, b.payload)) yield reduce(merge, values) else: for value in values: yield value def _decode_gauge(coder, payload): """Returns a tuple of (timestamp, value).""" timestamp_coder = coders.VarIntCoder().get_impl() stream = coder_impl.create_InputStream(payload) time_ms = timestamp_coder.decode_from_stream(stream, True) return (time_ms / 1000.0, coder.get_impl().decode_from_stream(stream, True)) def _encode_gauge(coder, timestamp, value): timestamp_coder = coders.VarIntCoder().get_impl() stream = coder_impl.create_OutputStream() timestamp_coder.encode_to_stream(int(timestamp * 1000), stream, True) coder.get_impl().encode_to_stream(value, stream, True) return stream.get() def _decode_distribution(value_coder, payload): """Returns a tuple of (count, sum, min, max).""" count_coder = coders.VarIntCoder().get_impl() value_coder = value_coder.get_impl() stream = coder_impl.create_InputStream(payload) return ( count_coder.decode_from_stream(stream, True), value_coder.decode_from_stream(stream, True), value_coder.decode_from_stream(stream, True), value_coder.decode_from_stream(stream, True)) def _encode_distribution(value_coder, count, sum, min, max): count_coder = coders.VarIntCoder().get_impl() value_coder = value_coder.get_impl() stream = coder_impl.create_OutputStream() count_coder.encode_to_stream(count, stream, True) value_coder.encode_to_stream(sum, stream, True) value_coder.encode_to_stream(min, stream, True) value_coder.encode_to_stream(max, stream, True) return stream.get() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/metrics/monitoring_infos_test.py0000644000175100017510000001561115206113032024323 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest from apache_beam.metrics import monitoring_infos from apache_beam.metrics.cells import CounterCell from apache_beam.metrics.cells import GaugeCell from apache_beam.metrics.cells import HistogramCell from apache_beam.metrics.cells import HistogramData from apache_beam.metrics.cells import StringSetCell from apache_beam.utils.histogram import Histogram from apache_beam.utils.histogram import LinearBucket class MonitoringInfosTest(unittest.TestCase): def test_parse_namespace_and_name_for_nonuser_metric(self): input = monitoring_infos.create_monitoring_info( "beam:dummy:metric", "typeurn", None) namespace, name = monitoring_infos.parse_namespace_and_name(input) self.assertEqual(namespace, "beam") self.assertEqual(name, "dummy:metric") def test_parse_namespace_and_name_for_user_counter_metric(self): urn = monitoring_infos.USER_COUNTER_URN labels = {} labels[monitoring_infos.NAMESPACE_LABEL] = "counternamespace" labels[monitoring_infos.NAME_LABEL] = "countername" input = monitoring_infos.create_monitoring_info( urn, "typeurn", None, labels) namespace, name = monitoring_infos.parse_namespace_and_name(input) self.assertEqual(namespace, "counternamespace") self.assertEqual(name, "countername") def test_parse_namespace_and_name_for_user_distribution_metric(self): urn = monitoring_infos.USER_DISTRIBUTION_URN labels = {} labels[monitoring_infos.NAMESPACE_LABEL] = "counternamespace" labels[monitoring_infos.NAME_LABEL] = "countername" input = monitoring_infos.create_monitoring_info( urn, "typeurn", None, labels) namespace, name = monitoring_infos.parse_namespace_and_name(input) self.assertEqual(namespace, "counternamespace") self.assertEqual(name, "countername") def test_parse_namespace_and_name_for_user_gauge_metric(self): urn = monitoring_infos.USER_GAUGE_URN labels = {} labels[monitoring_infos.NAMESPACE_LABEL] = "counternamespace" labels[monitoring_infos.NAME_LABEL] = "countername" input = monitoring_infos.create_monitoring_info( urn, "typeurn", None, labels) namespace, name = monitoring_infos.parse_namespace_and_name(input) self.assertEqual(namespace, "counternamespace") self.assertEqual(name, "countername") def test_parse_namespace_and_name_for_user_string_set_metric(self): urn = monitoring_infos.USER_STRING_SET_URN labels = {} labels[monitoring_infos.NAMESPACE_LABEL] = "stringsetnamespace" labels[monitoring_infos.NAME_LABEL] = "stringsetname" input = monitoring_infos.create_monitoring_info( urn, "typeurn", None, labels) namespace, name = monitoring_infos.parse_namespace_and_name(input) self.assertEqual(namespace, "stringsetnamespace") self.assertEqual(name, "stringsetname") def test_parse_namespace_and_name_for_user_histogram_metric(self): urn = monitoring_infos.USER_HISTOGRAM_URN labels = {} labels[monitoring_infos.NAMESPACE_LABEL] = "histogramnamespace" labels[monitoring_infos.NAME_LABEL] = "histogramname" input = monitoring_infos.create_monitoring_info( urn, "typeurn", None, labels) namespace, name = monitoring_infos.parse_namespace_and_name(input) self.assertEqual(name, "histogramname") self.assertEqual(namespace, "histogramnamespace") def test_int64_user_gauge(self): metric = GaugeCell().get_cumulative() result = monitoring_infos.int64_user_gauge( 'gaugenamespace', 'gaugename', metric) _, gauge_value = monitoring_infos.extract_gauge_value(result) self.assertEqual(0, gauge_value) def test_int64_user_counter(self): expected_labels = {} expected_labels[monitoring_infos.NAMESPACE_LABEL] = "counternamespace" expected_labels[monitoring_infos.NAME_LABEL] = "countername" metric = CounterCell().get_cumulative() result = monitoring_infos.int64_user_counter( 'counternamespace', 'countername', metric) counter_value = monitoring_infos.extract_counter_value(result) self.assertEqual(0, counter_value) self.assertEqual(result.labels, expected_labels) def test_int64_counter(self): expected_labels = {} expected_labels[monitoring_infos.PCOLLECTION_LABEL] = "collectionname" expected_labels[monitoring_infos.PTRANSFORM_LABEL] = "ptransformname" expected_labels[monitoring_infos.SERVICE_LABEL] = "BigQuery" labels = { monitoring_infos.SERVICE_LABEL: "BigQuery", } metric = CounterCell().get_cumulative() result = monitoring_infos.int64_counter( monitoring_infos.API_REQUEST_COUNT_URN, metric, ptransform="ptransformname", pcollection="collectionname", labels=labels) counter_value = monitoring_infos.extract_counter_value(result) self.assertEqual(0, counter_value) self.assertEqual(result.labels, expected_labels) def test_user_set_string(self): expected_labels = {} expected_labels[monitoring_infos.NAMESPACE_LABEL] = "stringsetnamespace" expected_labels[monitoring_infos.NAME_LABEL] = "stringsetname" metric = StringSetCell().get_cumulative() result = monitoring_infos.user_set_string( 'stringsetnamespace', 'stringsetname', metric) string_set_value = monitoring_infos.extract_string_set_value(result) self.assertEqual(set(), string_set_value) self.assertEqual(result.labels, expected_labels) def test_user_histogram(self): datapoints = [5, 50, 90] expected_labels = {} expected_labels[monitoring_infos.NAMESPACE_LABEL] = "histogramnamespace" expected_labels[monitoring_infos.NAME_LABEL] = "histogramname" cell = HistogramCell(LinearBucket(0, 1, 100)) for point in datapoints: cell.update(point) metric = cell.get_cumulative() result = monitoring_infos.user_histogram( 'histogramnamespace', 'histogramname', metric) histogramvalue = monitoring_infos.extract_histogram_value(result) self.assertEqual(result.labels, expected_labels) exp_histogram = Histogram(LinearBucket(0, 1, 100)) for point in datapoints: exp_histogram.record(point) self.assertEqual(HistogramData(exp_histogram), histogramvalue) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.4044104 apache_beam-2.74.0/apache_beam/ml/0000755000175100017510000000000015206113157016275 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/__init__.py0000644000175100017510000000161415206113032020400 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Contains packages for supported machine learning transforms.""" from apache_beam.ml.transforms.base import MLTransform ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.4084108 apache_beam-2.74.0/apache_beam/ml/anomaly/0000755000175100017510000000000015206113157017735 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/__init__.py0000644000175100017510000000142015206113032022033 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/aggregations.py0000644000175100017510000002771015206113032022760 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import collections import math import statistics from collections.abc import Callable from collections.abc import Iterable from typing import Any from typing import Optional from apache_beam.ml.anomaly.base import DEFAULT_MISSING_LABEL from apache_beam.ml.anomaly.base import DEFAULT_NORMAL_LABEL from apache_beam.ml.anomaly.base import DEFAULT_OUTLIER_LABEL from apache_beam.ml.anomaly.base import AggregationFn from apache_beam.ml.anomaly.base import AnomalyPrediction from apache_beam.ml.anomaly.specifiable import specifiable class _AggModelIdMixin: def __init__(self, agg_model_id: Optional[str] = None): self._agg_model_id = agg_model_id def _set_agg_model_id_if_unset(self, agg_model_id: str) -> None: if self._agg_model_id is None: self._agg_model_id = agg_model_id def add_model_id(self, result_dict): result_dict["model_id"] = self._agg_model_id class _SourcePredictionMixin: def __init__(self, include_source_predictions): self._include_source_predictions = include_source_predictions def add_source_predictions(self, result_dict, source_predictions): if self._include_source_predictions: result_dict["source_predictions"] = list(source_predictions) class LabelAggregation(AggregationFn, _AggModelIdMixin, _SourcePredictionMixin): """Aggregates anomaly predictions based on their labels. This is an abstract base class for `AggregationFn`s that combine multiple `AnomalyPrediction` objects into a single `AnomalyPrediction` based on the labels of the input predictions. Args: agg_func (Callable[[Iterable[int]], int]): A function that aggregates a collection of anomaly labels (integers) into a single label. agg_model_id (Optional[str]): The model id used in aggregated predictions. Defaults to None. include_source_predictions (bool): If True, include the input predictions in the `source_predictions` of the output. Defaults to False. """ def __init__( self, agg_func: Callable[[Iterable[int]], int], agg_model_id: Optional[str] = None, include_source_predictions: bool = False, normal_label: int = DEFAULT_NORMAL_LABEL, outlier_label: int = DEFAULT_OUTLIER_LABEL, missing_label: int = DEFAULT_MISSING_LABEL, ): self._agg = agg_func self._normal_label = normal_label self._outlier_label = outlier_label self._missing_label = missing_label _AggModelIdMixin.__init__(self, agg_model_id) _SourcePredictionMixin.__init__(self, include_source_predictions) def apply( self, predictions: Iterable[AnomalyPrediction]) -> AnomalyPrediction: """Applies the label aggregation function to a list of predictions. Args: predictions (Iterable[AnomalyPrediction]): A collection of `AnomalyPrediction` objects to be aggregated. Returns: AnomalyPrediction: A single `AnomalyPrediction` object with the aggregated label. The aggregated label is determined as follows: - If there are any non-missing and non-error labels, the `agg_func` is applied to aggregate them. - If all labels are error labels (`None`), the aggregated label is also `None`. - If there are a mix of missing and error labels, the aggregated label is the `missing_label`. """ result_dict: dict[str, Any] = {} _AggModelIdMixin.add_model_id(self, result_dict) _SourcePredictionMixin.add_source_predictions( self, result_dict, predictions) labels = [ prediction.label for prediction in predictions if prediction.label is not None and prediction.label != self._missing_label ] if len(labels) > 0: # apply aggregation_fn if there is any non-None and non-missing label result_dict["label"] = self._agg(labels) elif all(map(lambda x: x.label is None, predictions)): # all are error labels (None) -- all scores are error result_dict["label"] = None else: # some missing labels with some error labels (None) result_dict["label"] = self._missing_label return AnomalyPrediction(**result_dict) class ScoreAggregation(AggregationFn, _AggModelIdMixin, _SourcePredictionMixin): """Aggregates anomaly predictions based on their scores. This is an abstract base class for `AggregationFn`s that combine multiple `AnomalyPrediction` objects into a single `AnomalyPrediction` based on the scores of the input predictions. Args: agg_func (Callable[[Iterable[float]], float]): A function that aggregates a collection of anomaly scores (floats) into a single score. agg_model_id (Optional[str]): The model id used in aggregated predictions. Defaults to None. include_source_predictions (bool): If True, include the input predictions in the `source_predictions` of the output. Defaults to False. """ def __init__( self, agg_func: Callable[[Iterable[float]], float], agg_model_id: Optional[str] = None, include_source_predictions: bool = False): self._agg = agg_func _AggModelIdMixin.__init__(self, agg_model_id) _SourcePredictionMixin.__init__(self, include_source_predictions) def apply( self, predictions: Iterable[AnomalyPrediction]) -> AnomalyPrediction: """Applies the score aggregation function to a list of predictions. Args: predictions (Iterable[AnomalyPrediction]): A collection of `AnomalyPrediction` objects to be aggregated. Returns: AnomalyPrediction: A single `AnomalyPrediction` object with the aggregated score. The aggregated score is determined as follows: - If there are any non-missing and non-error scores, the `agg_func` is applied to aggregate them. - If all scores are error scores (`None`), the aggregated score is also `None`. - If there are a mix of missing (`NaN`) and error scores (`None`), the aggregated score is `NaN`. """ result_dict: dict[str, Any] = {} _AggModelIdMixin.add_model_id(self, result_dict) _SourcePredictionMixin.add_source_predictions( self, result_dict, predictions) scores = [ prediction.score for prediction in predictions if prediction.score is not None and not math.isnan(prediction.score) ] if len(scores) > 0: # apply aggregation_fn if there is any non-None and non-NaN score result_dict["score"] = self._agg(scores) elif all(map(lambda x: x.score is None, predictions)): # all are error scores (None) result_dict["score"] = None else: # some missing scores (NaN) with some error scores (None) result_dict["score"] = float("NaN") return AnomalyPrediction(**result_dict) @specifiable class MajorityVote(LabelAggregation): """Aggregates anomaly labels using majority voting. This `AggregationFn` implements a majority voting strategy to combine anomaly labels from multiple `AnomalyPrediction` objects. It counts the occurrences of normal and outlier labels and selects the label with the higher count as the aggregated label. In case of a tie, a tie-breaker label is used. Example: If input labels are [normal, outlier, outlier, normal, outlier], and normal_label=0, outlier_label=1, then the aggregated label will be outlier (1) because outliers have a majority (3 vs 2). Args: normal_label (int): The integer label for normal predictions. Defaults to 0. outlier_label (int): The integer label for outlier predictions. Defaults to 1. tie_breaker (int): The label to return if there is a tie in votes. Defaults to 0 (normal_label). **kwargs: Additional keyword arguments to pass to the base `LabelAggregation` class. """ def __init__(self, tie_breaker=DEFAULT_NORMAL_LABEL, **kwargs): self._tie_breaker = tie_breaker def inner(predictions: Iterable[int]) -> int: counters = collections.Counter(predictions) if counters[self._normal_label] < counters[self._outlier_label]: vote = self._outlier_label elif counters[self._normal_label] > counters[self._outlier_label]: vote = self._normal_label else: vote = self._tie_breaker return vote super().__init__(agg_func=inner, **kwargs) # And scheme @specifiable class AllVote(LabelAggregation): """Aggregates anomaly labels using an "all vote" (AND) scheme. This `AggregationFn` implements an "all vote" strategy. It aggregates anomaly labels such that the result is considered an outlier only if all input `AnomalyPrediction` objects are labeled as outliers. Example: If input labels are [outlier, outlier, outlier], and outlier_label=1, then the aggregated label will be outlier (1). If input labels are [outlier, normal, outlier], and outlier_label=1, then the aggregated label will be normal (0). Args: normal_label (int): The integer label for normal predictions. Defaults to 0. outlier_label (int): The integer label for outlier predictions. Defaults to 1. **kwargs: Additional keyword arguments to pass to the base `LabelAggregation` class. """ def __init__(self, **kwargs): def inner(predictions: Iterable[int]) -> int: return self._outlier_label if all( map(lambda p: p == self._outlier_label, predictions)) else self._normal_label super().__init__(agg_func=inner, **kwargs) # Or scheme @specifiable class AnyVote(LabelAggregation): """Aggregates anomaly labels using an "any vote" (OR) scheme. This `AggregationFn` implements an "any vote" strategy. It aggregates anomaly labels such that the result is considered an outlier if at least one of the input `AnomalyPrediction` objects is labeled as an outlier. Example: If input labels are [normal, normal, outlier], and outlier_label=1, then the aggregated label will be outlier (1). If input labels are [normal, normal, normal], and outlier_label=1, then the aggregated label will be normal (0). Args: normal_label (int): The integer label for normal predictions. Defaults to 0. outlier_label (int): The integer label for outlier predictions. Defaults to 1. **kwargs: Additional keyword arguments to pass to the base `LabelAggregation` class. """ def __init__(self, **kwargs): def inner(predictions: Iterable[int]) -> int: return self._outlier_label if any( map(lambda p: p == self._outlier_label, predictions)) else self._normal_label super().__init__(agg_func=inner, **kwargs) @specifiable class AverageScore(ScoreAggregation): """Aggregates anomaly scores by calculating their average. This `AggregationFn` computes the average of the anomaly scores from a collection of `AnomalyPrediction` objects. Args: **kwargs: Additional keyword arguments to pass to the base `ScoreAggregation` class. """ def __init__(self, **kwargs): super().__init__(agg_func=statistics.mean, **kwargs) @specifiable class MaxScore(ScoreAggregation): """Aggregates anomaly scores by selecting the maximum score. This `AggregationFn` selects the highest anomaly score from a collection of `AnomalyPrediction` objects as the aggregated score. Args: **kwargs: Additional keyword arguments to pass to the base `ScoreAggregation` class. """ def __init__(self, **kwargs): super().__init__(agg_func=max, **kwargs) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/aggregations_test.py0000644000175100017510000001222715206113032024014 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest from apache_beam.ml.anomaly import aggregations from apache_beam.ml.anomaly.base import AnomalyPrediction class LabelAggTestWithMissingOrError(unittest.TestCase): def test_default(self): normal = AnomalyPrediction(label=0) outlier = AnomalyPrediction(label=1) missing = AnomalyPrediction(label=-2) error = AnomalyPrediction(label=None) vote = aggregations.MajorityVote().apply # missing and error labels are ignored if there is any normal/outlier self.assertEqual(vote([normal, missing, error]), normal) self.assertEqual(vote([outlier, missing, error]), outlier) # if there is any missing among errors, return missing self.assertEqual(vote([error, missing, error]), missing) # return error only when all are errors self.assertEqual(vote([error, error, error]), error) class MajorityVoteTest(unittest.TestCase): def test_default(self): normal = AnomalyPrediction(label=0) outlier = AnomalyPrediction(label=1) vote = aggregations.MajorityVote().apply self.assertEqual(vote([]), AnomalyPrediction()) self.assertEqual(vote([normal]), normal) self.assertEqual(vote([outlier]), outlier) self.assertEqual(vote([outlier, normal, normal]), normal) self.assertEqual(vote([outlier, normal, outlier]), outlier) # use normal to break ties by default self.assertEqual(vote([outlier, normal]), normal) def test_tie_breaker(self): normal = AnomalyPrediction(label=0) outlier = AnomalyPrediction(label=1) vote = aggregations.MajorityVote(tie_breaker=1).apply self.assertEqual(vote([outlier, normal]), outlier) class AllVoteTest(unittest.TestCase): def test_default(self): normal = AnomalyPrediction(label=0) outlier = AnomalyPrediction(label=1) vote = aggregations.AllVote().apply self.assertEqual(vote([]), AnomalyPrediction()) self.assertEqual(vote([normal]), normal) self.assertEqual(vote([outlier]), outlier) # outlier is only labeled when everyone is outlier self.assertEqual(vote([normal, normal, normal]), normal) self.assertEqual(vote([outlier, normal, normal]), normal) self.assertEqual(vote([outlier, normal, outlier]), normal) self.assertEqual(vote([outlier, outlier, outlier]), outlier) class AnyVoteTest(unittest.TestCase): def test_default(self): normal = AnomalyPrediction(label=0) outlier = AnomalyPrediction(label=1) vote = aggregations.AnyVote().apply self.assertEqual(vote([]), AnomalyPrediction()) self.assertEqual(vote([normal]), normal) self.assertEqual(vote([outlier]), outlier) # outlier is labeled when at least one is outlier self.assertEqual(vote([normal, normal, normal]), normal) self.assertEqual(vote([outlier, normal, normal]), outlier) self.assertEqual(vote([outlier, normal, outlier]), outlier) self.assertEqual(vote([outlier, outlier, outlier]), outlier) class ScoreAggTestWithMissingOrError(unittest.TestCase): def test_default(self): normal = AnomalyPrediction(score=1.0) missing = AnomalyPrediction(score=float("NaN")) error = AnomalyPrediction(score=None) avg = aggregations.AverageScore().apply # missing and error scores are ignored if there is any normal/outlier self.assertEqual(avg([normal, missing, error]), normal) # if there is any missing among errors, return missing. # note that NaN != NaN, so we cannot use `assertEqual` here. self.assertTrue(avg([error, missing, error]).score) # return error only when all are errors self.assertEqual(avg([error, error, error]), error) class AverageScoreTest(unittest.TestCase): def test_default(self): avg = aggregations.AverageScore().apply self.assertEqual(avg([]), AnomalyPrediction()) self.assertEqual( avg([AnomalyPrediction(score=1)]), AnomalyPrediction(score=1)) self.assertEqual( avg([AnomalyPrediction(score=1), AnomalyPrediction(score=2)]), AnomalyPrediction(score=1.5)) class MaxScoreTest(unittest.TestCase): def test_default(self): avg = aggregations.MaxScore().apply self.assertEqual(avg([]), AnomalyPrediction()) self.assertEqual( avg([AnomalyPrediction(score=1)]), AnomalyPrediction(score=1)) self.assertEqual( avg([AnomalyPrediction(score=1), AnomalyPrediction(score=2)]), AnomalyPrediction(score=2)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/base.py0000644000175100017510000001647715206113032021230 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Base classes for anomaly detection """ from __future__ import annotations import abc from collections.abc import Iterable from dataclasses import dataclass from typing import Optional import apache_beam as beam __all__ = [ "AnomalyPrediction", "AnomalyResult", "ThresholdFn", "AggregationFn", "AnomalyDetector", "EnsembleAnomalyDetector" ] DEFAULT_NORMAL_LABEL = 0 DEFAULT_OUTLIER_LABEL = 1 DEFAULT_MISSING_LABEL = -2 @dataclass(frozen=True) class AnomalyPrediction(): """A dataclass for anomaly detection predictions.""" #: The ID of detector (model) that generates the prediction. model_id: Optional[str] = None #: The outlier score resulting from applying the detector to the input data. score: Optional[float] = None #: The outlier label (normal or outlier) derived from the outlier score. label: Optional[int] = None #: The threshold used to determine the label. threshold: Optional[float] = None #: Additional information about the prediction. info: str = "" #: If enabled, a list of `AnomalyPrediction` objects used to derive the #: aggregated prediction. source_predictions: Optional[Iterable[AnomalyPrediction]] = None @dataclass(frozen=True) class AnomalyResult(): """A dataclass for the anomaly detection results""" #: The original input data. example: beam.Row #: The iterable of `AnomalyPrediction` objects containing the predictions. #: Expect length 1 if it is a result for a non-ensemble detector or an #: ensemble detector with an aggregation strategy applied. predictions: Iterable[AnomalyPrediction] class ThresholdFn(abc.ABC): """An abstract base class for threshold functions. Args: normal_label: The integer label used to identify normal data. Defaults to 0. outlier_label: The integer label used to identify outlier data. Defaults to 1. missing_label: The integer label used when a score is missing because the model is not ready to score. """ def __init__( self, normal_label: int = DEFAULT_NORMAL_LABEL, outlier_label: int = DEFAULT_OUTLIER_LABEL, missing_label: int = DEFAULT_MISSING_LABEL): self._normal_label = normal_label self._outlier_label = outlier_label self._missing_label = missing_label @property @abc.abstractmethod def is_stateful(self) -> bool: """Indicates whether the threshold function is stateful or not.""" raise NotImplementedError @property @abc.abstractmethod def threshold(self) -> Optional[float]: """Retrieves the current threshold value, or None if not set.""" raise NotImplementedError @abc.abstractmethod def apply(self, score: Optional[float]) -> Optional[int]: """Applies the threshold function to a given score to classify it as normal or outlier. Args: score: The outlier score generated from the detector (model). Returns: The label assigned to the score, either `self._normal_label` or `self._outlier_label` """ raise NotImplementedError class AggregationFn(abc.ABC): """An abstract base class for aggregation functions.""" @abc.abstractmethod def apply( self, predictions: Iterable[AnomalyPrediction]) -> AnomalyPrediction: """Applies the aggregation function to an iterable of predictions, either on their outlier scores or labels. Args: predictions: An Iterable of `AnomalyPrediction` objects to aggregate. Returns: An `AnomalyPrediction` object containing the aggregated result. """ raise NotImplementedError class AnomalyDetector(abc.ABC): """An abstract base class for anomaly detectors. Args: model_id: The ID of detector (model). Defaults to the value of the `spec_type` attribute, or 'unknown' if not set. features: An Iterable of strings representing the names of the input features in the `beam.Row` target: The name of the target field in the `beam.Row`. threshold_criterion: An optional `ThresholdFn` to apply to the outlier score and yield a label. """ def __init__( self, model_id: Optional[str] = None, features: Optional[Iterable[str]] = None, target: Optional[str] = None, threshold_criterion: Optional[ThresholdFn] = None, **kwargs): self._model_id = model_id if model_id is not None else getattr( self, 'spec_type', lambda: "unknown")() self._features = features self._target = target self._threshold_criterion = threshold_criterion @abc.abstractmethod def learn_one(self, x: beam.Row) -> None: """Trains the detector on a single data instance. Args: x: A `beam.Row` representing the data instance. """ raise NotImplementedError @abc.abstractmethod def score_one(self, x: beam.Row) -> Optional[float]: """Scores a single data instance for anomalies. Args: x: A `beam.Row` representing the data instance. Returns: The outlier score as a float. None if an exception occurs during scoring, and NaN if the model is not ready. """ raise NotImplementedError class EnsembleAnomalyDetector(AnomalyDetector): """An abstract base class for an ensemble of anomaly (sub-)detectors. Args: sub_detectors: A list of `AnomalyDetector` used in this ensemble model. aggregation_strategy: An optional `AggregationFn` to apply to the predictions from all sub-detectors and yield an aggregated result. model_id: Inherited from `AnomalyDetector`. features: Inherited from `AnomalyDetector`. target: Inherited from `AnomalyDetector`. threshold_criterion: Inherited from `AnomalyDetector`. """ def __init__( self, sub_detectors: Optional[list[AnomalyDetector]] = None, aggregation_strategy: Optional[AggregationFn] = None, **kwargs): if "model_id" not in kwargs or kwargs["model_id"] is None: kwargs["model_id"] = getattr(self, 'spec_type', lambda: 'custom')() super().__init__(**kwargs) self._aggregation_strategy = aggregation_strategy self._sub_detectors = sub_detectors def learn_one(self, x: beam.Row) -> None: """Inherited from `AnomalyDetector.learn_one`. This method is never called during ensemble detector training. The training process is done on each sub-detector independently and in parallel. """ raise NotImplementedError def score_one(self, x: beam.Row) -> float: """Inherited from `AnomalyDetector.score_one`. This method is never called during ensemble detector scoring. The scoring process is done on sub-detector independently and in parallel, and then the results are aggregated in the pipeline. """ raise NotImplementedError ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/base_test.py0000644000175100017510000001734115206113032022256 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from __future__ import annotations import copy import logging import unittest from parameterized import parameterized from apache_beam.ml.anomaly.base import AggregationFn from apache_beam.ml.anomaly.base import AnomalyDetector from apache_beam.ml.anomaly.base import EnsembleAnomalyDetector from apache_beam.ml.anomaly.base import ThresholdFn from apache_beam.ml.anomaly.specifiable import _KNOWN_SPECIFIABLE from apache_beam.ml.anomaly.specifiable import Spec from apache_beam.ml.anomaly.specifiable import Specifiable from apache_beam.ml.anomaly.specifiable import specifiable class TestAnomalyDetector(unittest.TestCase): def setUp(self) -> None: self.saved_specifiable = copy.deepcopy(_KNOWN_SPECIFIABLE) def tearDown(self) -> None: _KNOWN_SPECIFIABLE.clear() _KNOWN_SPECIFIABLE.update(self.saved_specifiable) @parameterized.expand([(False, False), (True, False), (False, True), (True, True)]) def test_model_id_and_spec(self, on_demand_init, just_in_time_init): @specifiable( on_demand_init=on_demand_init, just_in_time_init=just_in_time_init) class DummyThreshold(ThresholdFn): def __init__(self, my_threshold_arg=None): ... def is_stateful(self): return False def threshold(self): ... def apply(self, x): ... @specifiable( on_demand_init=on_demand_init, just_in_time_init=just_in_time_init) class Dummy(AnomalyDetector): def __init__(self, my_arg=None, **kwargs): self._my_arg = my_arg super().__init__(**kwargs) def learn_one(self): ... def score_one(self): ... def __eq__(self, value) -> bool: return isinstance(value, Dummy) and \ self._my_arg == value._my_arg a = Dummy( my_arg="abc", target="ABC", threshold_criterion=(t1 := DummyThreshold(2))) # The class attributes can only be accessed when # (1) on_demand_init == False, just_in_time_init == False # In this case, the true __init__ is called immediately during object # initialization # (2) just_in_time_init == True # In this case, regardless what on_demand_init is, the true __init__ # is called when we first access any class attribute (delay init). if just_in_time_init or not on_demand_init: self.assertEqual(a._model_id, "Dummy") self.assertEqual(a._target, "ABC") self.assertEqual(a._my_arg, "abc") assert isinstance(a, Specifiable) self.assertEqual( a.init_kwargs, { "my_arg": "abc", "target": "ABC", "threshold_criterion": t1, }) b = Dummy( my_arg="efg", model_id="my_dummy", target="EFG", threshold_criterion=(t2 := DummyThreshold(3))) # See the comment above for more details. if just_in_time_init or not on_demand_init: self.assertEqual(b._model_id, "my_dummy") self.assertEqual(b._target, "EFG") self.assertEqual(b._my_arg, "efg") assert isinstance(b, Specifiable) self.assertEqual( b.init_kwargs, { "model_id": "my_dummy", "my_arg": "efg", "target": "EFG", "threshold_criterion": t2, }) spec = b.to_spec() expected_spec = Spec( type="Dummy", config={ "my_arg": "efg", "model_id": "my_dummy", "target": "EFG", "threshold_criterion": Spec( type="DummyThreshold", config={"my_threshold_arg": 3}), }) self.assertEqual(spec, expected_spec) b_dup = Specifiable.from_spec(spec) # We need to manually call the original __init__ function in one scenario. # See the comment above for more details. if on_demand_init and not just_in_time_init: b.run_original_init() self.assertEqual(b, b_dup) class TestEnsembleAnomalyDetector(unittest.TestCase): def setUp(self) -> None: self.saved_specifiable = copy.deepcopy(_KNOWN_SPECIFIABLE) def tearDown(self) -> None: _KNOWN_SPECIFIABLE.clear() _KNOWN_SPECIFIABLE.update(self.saved_specifiable) @parameterized.expand([(False, False), (True, False), (False, True), (True, True)]) def test_model_id_and_spec(self, on_demand_init, just_in_time_init): @specifiable( on_demand_init=on_demand_init, just_in_time_init=just_in_time_init) class DummyAggregation(AggregationFn): def apply(self, x): ... @specifiable( on_demand_init=on_demand_init, just_in_time_init=just_in_time_init) class DummyEnsemble(EnsembleAnomalyDetector): def __init__(self, my_ensemble_arg=None, **kwargs): super().__init__(**kwargs) self._my_ensemble_arg = my_ensemble_arg def learn_one(self): ... def score_one(self): ... def __eq__(self, value) -> bool: return isinstance(value, DummyEnsemble) and \ self._my_ensemble_arg == value._my_ensemble_arg @specifiable( on_demand_init=on_demand_init, just_in_time_init=just_in_time_init) class DummyWeakLearner(AnomalyDetector): def __init__(self, my_arg=None, **kwargs): super().__init__(**kwargs) self._my_arg = my_arg def learn_one(self): ... def score_one(self): ... def __eq__(self, value) -> bool: return isinstance(value, DummyWeakLearner) \ and self._my_arg == value._my_arg # See the comment in TestAnomalyDetector for more details. if just_in_time_init or not on_demand_init: a = DummyEnsemble() self.assertEqual(a._model_id, "DummyEnsemble") b = DummyEnsemble(model_id="my_dummy_ensemble") self.assertEqual(b._model_id, "my_dummy_ensemble") c = EnsembleAnomalyDetector() self.assertEqual(c._model_id, "custom") d = EnsembleAnomalyDetector(model_id="my_dummy_ensemble_2") self.assertEqual(d._model_id, "my_dummy_ensemble_2") d1 = DummyWeakLearner(my_arg=1) d2 = DummyWeakLearner(my_arg=2) ensemble = DummyEnsemble( my_ensemble_arg=123, learners=[d1, d2], aggregation_strategy=DummyAggregation()) expected_spec = Spec( type="DummyEnsemble", config={ "my_ensemble_arg": 123, "learners": [ Spec(type="DummyWeakLearner", config={"my_arg": 1}), Spec(type="DummyWeakLearner", config={"my_arg": 2}) ], "aggregation_strategy": Spec( type="DummyAggregation", config={}, ), }) assert isinstance(ensemble, Specifiable) spec = ensemble.to_spec() self.assertEqual(spec, expected_spec) ensemble_dup = Specifiable.from_spec(spec) # See the comment in TestAnomalyDetector for more details. if on_demand_init and not just_in_time_init: ensemble.run_original_init() self.assertEqual(ensemble, ensemble_dup) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.4664168 apache_beam-2.74.0/apache_beam/ml/anomaly/detectors/0000755000175100017510000000000015206113157021731 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/detectors/__init__.py0000644000175100017510000000171115206113032024032 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from apache_beam.ml.anomaly.detectors.iqr import IQR from apache_beam.ml.anomaly.detectors.robust_zscore import RobustZScore from apache_beam.ml.anomaly.detectors.zscore import ZScore ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/detectors/iqr.py0000644000175100017510000001106715206113032023073 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import math from typing import Optional import apache_beam as beam from apache_beam.ml.anomaly.base import AnomalyDetector from apache_beam.ml.anomaly.specifiable import specifiable from apache_beam.ml.anomaly.thresholds import FixedThreshold from apache_beam.ml.anomaly.univariate.base import EPSILON from apache_beam.ml.anomaly.univariate.quantile import BufferedSlidingQuantileTracker # pylint: disable=line-too-long from apache_beam.ml.anomaly.univariate.quantile import QuantileTracker from apache_beam.ml.anomaly.univariate.quantile import SecondaryBufferedQuantileTracker # pylint: disable=line-too-long DEFAULT_WINDOW_SIZE = 1000 @specifiable class IQR(AnomalyDetector): """Interquartile Range (IQR) anomaly detector. This class implements an anomaly detection algorithm based on the Interquartile Range (IQR) [#]_ . It calculates the IQR using quantile trackers for Q1 (25th percentile) and Q3 (75th percentile) and scores data points based on their deviation from these quartiles. The score is calculated as follows: * If a data point is above Q3, the score is (value - Q3) / IQR. * If a data point is below Q1, the score is (Q1 - value) / IQR. * If a data point is within the IQR (Q1 <= value <= Q3), the score is 0. Initializes the IQR anomaly detector. Args: q1_tracker: Optional QuantileTracker for Q1 (25th percentile). If None, a BufferedSlidingQuantileTracker with a default window size is used. q3_tracker: Optional QuantileTracker for Q3 (75th percentile). If None, a SecondaryBufferedQuantileTracker based on q1_tracker is used. threshold_criterion: Optional ThresholdFn to apply on the score. Defaults to `FixedThreshold(1.5)` since outliers are commonly defined as data points that fall below Q1 - 1.5 IQR or above Q3 + 1.5 IQR. **kwargs: Additional keyword arguments. .. [#] https://en.wikipedia.org/wiki/Interquartile_range """ def __init__( self, q1_tracker: Optional[QuantileTracker] = None, q3_tracker: Optional[QuantileTracker] = None, **kwargs): if "threshold_criterion" not in kwargs: kwargs["threshold_criterion"] = FixedThreshold(1.5) super().__init__(**kwargs) self._q1_tracker = q1_tracker or \ BufferedSlidingQuantileTracker(DEFAULT_WINDOW_SIZE, 0.25) assert self._q1_tracker._q == 0.25, \ "q1_tracker must be initialized with q = 0.25" self._q3_tracker = q3_tracker or \ SecondaryBufferedQuantileTracker(self._q1_tracker, 0.75) assert self._q3_tracker._q == 0.75, \ "q3_tracker must be initialized with q = 0.75" def learn_one(self, x: beam.Row) -> None: """Updates the quantile trackers with a new data point. Args: x: A `beam.Row` containing a single numerical value. """ if len(x.__dict__) != 1: raise ValueError( "IQR.learn_one expected univariate input, but got %s", str(x)) v = next(iter(x)) self._q1_tracker.push(v) self._q3_tracker.push(v) def score_one(self, x: beam.Row) -> Optional[float]: """Scores a data point based on its deviation from the IQR. Args: x: A `beam.Row` containing a single numerical value. Returns: float | None: The anomaly score. """ if len(x.__dict__) != 1: raise ValueError( "IQR.score_one expected univariate input, but got %s", str(x)) v = next(iter(x)) if v is None or math.isnan(v): return None q1 = self._q1_tracker.get() q3 = self._q3_tracker.get() # not enough data points to compute median or median absolute deviation if math.isnan(q1) or math.isnan(q3): return float('NaN') iqr = q3 - q1 if abs(iqr) < EPSILON: return 0.0 if v > q3: return (v - q3) / iqr if v < q1: return (q1 - v) / iqr # q1 <= v <= q3, normal points return 0 ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/detectors/iqr_test.py0000644000175100017510000000274215206113032024132 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import math import unittest import apache_beam as beam from apache_beam.ml.anomaly.detectors.iqr import IQR class IQRTest(unittest.TestCase): input = [ beam.Row(x=1), beam.Row(x=1), beam.Row(x=5), beam.Row(x=9), beam.Row(x=20), beam.Row(x=10), beam.Row(x=1) ] def test_with_default_trackers(self): iqr = IQR() scores = [] for row in IQRTest.input: scores.append(iqr.score_one(row)) iqr.learn_one(row) self.assertTrue(math.isnan(scores[0])) self.assertEqual( scores[1:], [0.0, 0.0, 3.0, 2.8, 0.125, 0.12903225806451613]) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/detectors/offline.py0000644000175100017510000001061015206113032023713 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from typing import Any from typing import Optional from typing import SupportsFloat from typing import SupportsInt from typing import TypeVar import apache_beam as beam from apache_beam.ml.anomaly.base import AnomalyDetector from apache_beam.ml.anomaly.base import AnomalyPrediction from apache_beam.ml.anomaly.specifiable import specifiable from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import PredictionT KeyT = TypeVar('KeyT') @specifiable class OfflineDetector(AnomalyDetector): """A offline anomaly detector that uses a provided model handler for scoring. Args: keyed_model_handler: The model handler to use for inference. Requires a `KeyModelHandler[Any, Row, PredictionT, Any]` instance. run_inference_args: Optional arguments to pass to RunInference **kwargs: Additional keyword arguments to pass to the base AnomalyDetector class. """ @staticmethod def score_prediction_adapter( keyed_prediction: tuple[KeyT, PredictionResult] ) -> tuple[KeyT, AnomalyPrediction]: """Extracts a float score from `PredictionResult.inference` and wraps it. Takes a keyed `PredictionResult` from common ModelHandler output, assumes its `inference` attribute is a float-convertible score, and returns the key paired with an `AnomalyPrediction` containing that float score. Args: keyed_prediction: tuple of `(key, PredictionResult)`. `PredictionResult` must have an `inference` attribute supporting float conversion. Returns: tuple of `(key, AnomalyPrediction)` with the extracted score. Raises: AssertionError: If `PredictionResult.inference` doesn't support float(). """ key, prediction = keyed_prediction score = prediction.inference assert isinstance(score, SupportsFloat) return key, AnomalyPrediction(score=float(score)) @staticmethod def label_prediction_adapter( keyed_prediction: tuple[KeyT, PredictionResult] ) -> tuple[KeyT, AnomalyPrediction]: """Extracts an integer label from `PredictionResult.inference` and wraps it. Takes a keyed `PredictionResult`, assumes its `inference` attribute is an integer-convertible label, and returns the key paired with an `AnomalyPrediction` containing that integer label. Args: keyed_prediction: tuple of `(key, PredictionResult)`. `PredictionResult` must have an `inference` attribute supporting int conversion. Returns: tuple of `(key, AnomalyPrediction)` with the extracted label. Raises: AssertionError: If `PredictionResult.inference` doesn't support int(). """ key, prediction = keyed_prediction label = prediction.inference assert isinstance(label, SupportsInt) return key, AnomalyPrediction(label=int(label)) def __init__( self, keyed_model_handler: KeyedModelHandler[Any, beam.Row, PredictionT, Any], run_inference_args: Optional[dict[str, Any]] = None, **kwargs): super().__init__(**kwargs) # TODO: validate the model handler type self._keyed_model_handler = keyed_model_handler self._run_inference_args = run_inference_args or {} # always override model_identifier with model_id from the detector self._run_inference_args["model_identifier"] = self._model_id def learn_one(self, x: beam.Row) -> None: """Not implemented since OfflineDetector invokes RunInference directly.""" raise NotImplementedError def score_one(self, x: beam.Row) -> Optional[float]: """Not implemented since OfflineDetector invokes RunInference directly.""" raise NotImplementedError ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/detectors/pyod_adapter.py0000644000175100017510000001002415206113032024743 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import pickle from collections.abc import Iterable from collections.abc import Sequence from typing import Any from typing import Optional import numpy as np from pyod.models.base import BaseDetector as PyODBaseDetector import apache_beam as beam from apache_beam.io.filesystems import FileSystems from apache_beam.ml.anomaly.detectors.offline import OfflineDetector from apache_beam.ml.anomaly.specifiable import specifiable from apache_beam.ml.anomaly.thresholds import FixedThreshold from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import _PostProcessingModelHandler from apache_beam.ml.inference.utils import _convert_to_result # Turn the used ModelHandler into specifiable, but without lazy init. KeyedModelHandler = specifiable( # type: ignore[misc] KeyedModelHandler, on_demand_init=False, just_in_time_init=False) _PostProcessingModelHandler = specifiable( # type: ignore[misc] _PostProcessingModelHandler, on_demand_init=False, just_in_time_init=False) @specifiable class PyODModelHandler(ModelHandler[beam.Row, PredictionResult, PyODBaseDetector]): """Implementation of the ModelHandler interface for PyOD [#]_ Models. The ModelHandler processes input data as `beam.Row` objects. **NOTE:** This API and its implementation are currently under active development and may not be backward compatible. Args: model_uri: The URI specifying the location of the pickled PyOD model. .. [#] https://github.com/yzhao062/pyod """ def __init__(self, model_uri: str): self._model_uri = model_uri def load_model(self) -> PyODBaseDetector: file = FileSystems.open(self._model_uri, 'rb') return pickle.load(file) def run_inference( self, batch: Sequence[beam.Row], model: PyODBaseDetector, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: np_batch = [] for row in batch: np_batch.append(np.fromiter(row, dtype=np.float64)) # stack a batch of samples into a 2-D array for better performance vectorized_batch = np.stack(np_batch, axis=0) predictions = model.decision_function(vectorized_batch) return _convert_to_result(batch, predictions, model_id=self._model_uri) class PyODFactory(): @staticmethod def create_detector(model_uri: str, **kwargs) -> OfflineDetector: """A utility function to create OfflineDetector for a PyOD model. **NOTE:** This API and its implementation are currently under active development and may not be backward compatible. Args: model_uri: The URI specifying the location of the pickled PyOD model. **kwargs: Additional keyword arguments. """ model_handler = KeyedModelHandler( PyODModelHandler(model_uri=model_uri)).with_postprocess_fn( OfflineDetector.score_prediction_adapter) m = model_handler.load_model() assert (isinstance(m, PyODBaseDetector)) threshold = float(m.threshold_) detector = OfflineDetector( model_handler, threshold_criterion=FixedThreshold(threshold), **kwargs) # type: ignore[arg-type] return detector ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/detectors/pyod_adapter_test.py0000644000175100017510000001372015206113032026010 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import os.path import pickle import shutil import tempfile import unittest import numpy as np from parameterized import parameterized import apache_beam as beam from apache_beam.ml.anomaly.base import AnomalyPrediction from apache_beam.ml.anomaly.base import AnomalyResult from apache_beam.ml.anomaly.transforms import AnomalyDetection from apache_beam.ml.anomaly.transforms_test import _keyed_result_is_equal_to from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # Protect against environments where onnx and pytorch library is not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from pyod.models.iforest import IForest from apache_beam.ml.anomaly.detectors.pyod_adapter import PyODFactory except ImportError: raise unittest.SkipTest('PyOD dependencies are not installed') class PyODIForestTest(unittest.TestCase): def setUp(self) -> None: self.tmp_dir = tempfile.mkdtemp() seed = 1234 model = IForest(random_state=seed) model.fit(self.get_train_data()) self.pickled_model_uri = os.path.join(self.tmp_dir, 'iforest_pickled') with open(self.pickled_model_uri, 'wb') as fp: pickle.dump(model, fp) def tearDown(self) -> None: shutil.rmtree(self.tmp_dir) def get_train_data(self): return [ np.array([1, 5], dtype="float32"), np.array([2, 6], dtype="float32"), np.array([3, 4], dtype="float32"), np.array([2, 6], dtype="float32"), np.array([10, 10], dtype="float32"), # need an outlier in training data np.array([3, 4], dtype="float32"), np.array([2, 6], dtype="float32"), np.array([2, 6], dtype="float32"), np.array([2, 5], dtype="float32"), ] def get_test_data(self): return [ np.array([2, 6], dtype="float32"), np.array([100, 100], dtype="float32"), ] def get_test_data_with_target(self): return [ np.array([2, 6, 0], dtype="float32"), np.array([100, 100, 1], dtype="float32"), ] @parameterized.expand([True, False]) def test_scoring_with_matched_features(self, with_target): if with_target: rows = [beam.Row(a=2, b=6, target=0), beam.Row(a=100, b=100, target=1)] field_names = ["a", "b", "target"] # The selected features should match the features used for training detector = PyODFactory.create_detector( self.pickled_model_uri, features=["a", "b"]) input_data = self.get_test_data_with_target() else: rows = [beam.Row(a=2, b=6), beam.Row(a=100, b=100)] field_names = ["a", "b"] detector = PyODFactory.create_detector(self.pickled_model_uri) input_data = self.get_test_data() expected_out = [( 0, AnomalyResult( example=rows[0], predictions=[ AnomalyPrediction( model_id='OfflineDetector', score=-0.20316164744828075, label=0, threshold=8.326672684688674e-17, info='', source_predictions=None) ])), ( 0, AnomalyResult( example=rows[1], predictions=[ AnomalyPrediction( model_id='OfflineDetector', score=0.179516865091218, label=1, threshold=8.326672684688674e-17, info='', source_predictions=None) ]))] options = PipelineOptions([]) with beam.Pipeline(options=options) as p: out = ( p | beam.Create(input_data) | beam.Map(lambda x: beam.Row(**dict(zip(field_names, map(int, x))))) | beam.WithKeys(0) | AnomalyDetection(detector=detector)) assert_that(out, equal_to(expected_out, _keyed_result_is_equal_to)) def test_scoring_with_unmatched_features(self): # The model is trained with two features: a, b, but the input features of # scoring has one more feature (target). # In this case, we should either get rid of the extra feature(s) from # the scoring input or set `features` when creating the offline detector # (see the `test_scoring_with_matched_features`) detector = PyODFactory.create_detector(self.pickled_model_uri) options = PipelineOptions([]) # This should raise a ValueError with message # "X has 3 features, but IsolationForest is expecting 2 features as input." with self.assertRaisesRegex(Exception, "is expecting 2 features"): with beam.Pipeline(options=options) as p: _ = ( p | beam.Create(self.get_test_data_with_target()) | beam.Map( lambda x: beam.Row( **dict(zip(["a", "b", "target"], map(int, x))))) | beam.WithKeys(0) | AnomalyDetection(detector=detector)) if __name__ == '__main__': logging.getLogger().setLevel(logging.WARNING) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/detectors/robust_zscore.py0000644000175100017510000001033115206113032025174 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import math from typing import Optional import apache_beam as beam from apache_beam.ml.anomaly.base import AnomalyDetector from apache_beam.ml.anomaly.specifiable import specifiable from apache_beam.ml.anomaly.thresholds import FixedThreshold from apache_beam.ml.anomaly.univariate.base import EPSILON from apache_beam.ml.anomaly.univariate.mad import MadTracker # pylint: disable=line-too-long @specifiable class RobustZScore(AnomalyDetector): """Robust Z-Score anomaly detector. This class implements an detection algorithm based on Robust Z-Score (also known as Modified Z-Score), which is a robust alternative to the traditional Z-score [#]_. It uses the median and Median Absolute Deviation (MAD) to compute a score that is less sensitive to outliers. The score is calculated as: `|0.6745 * (value - median) / MAD|` Important: In the streaming setting, we use the online version of median and MAD in the calculation. Therefore, the score computed here does not exactly match its batch counterpart. This implementation is adapted from the implementation within PySAD [#]_: https://github.com/selimfirat/pysad/blob/master/pysad/models/median_absolute_deviation.py The batch version can be seen at PyOD [#]_: https://github.com/yzhao062/pyod/blob/master/pyod/models/mad.py Args: mad_tracker: Optional `MadTracker` instance. If None, a default `MadTracker` is created. threshold_criterion: threshold_criterion: Optional `ThresholdFn` to apply on the score. Defaults to `FixedThreshold(3)` due to the commonly used 3-sigma rule. **kwargs: Additional keyword arguments. .. [#] Hoaglin, David C.. (2013). Volume 16: How to Detect and Handle Outliers. .. [#] Yilmaz, Selim & Kozat, Suleyman. (2020). PySAD: A Streaming Anomaly Detection Framework in Python. 10.48550/arXiv.2009.02572. .. [#] Zhao, Y., Nasrullah, Z. and Li, Z.. (2019). PyOD: A Python Toolbox for Scalable Outlier Detection. Journal of machine learning research (JMLR), 20(96), pp.1-7. """ # pylint: enable=line-too-long SCALE_FACTOR = 0.6745 def __init__(self, mad_tracker: Optional[MadTracker] = None, **kwargs): if "threshold_criterion" not in kwargs: kwargs["threshold_criterion"] = FixedThreshold(3) super().__init__(**kwargs) self._mad_tracker = mad_tracker or MadTracker() def learn_one(self, x: beam.Row) -> None: """Updates the `MadTracker` with a new data point. Args: x: A `beam.Row` containing a single numerical value. """ if len(x.__dict__) != 1: raise ValueError( "RobustZScore.learn_one expected univariate input, but got %s", str(x)) v = next(iter(x)) self._mad_tracker.push(v) def score_one(self, x: beam.Row) -> Optional[float]: """Scores a data point using the Robust Z-Score. Args: x: A `beam.Row` containing a single numerical value. Returns: float | None: The Robust Z-Score. """ if len(x.__dict__) != 1: raise ValueError( "RobustZScore.score_one expected univariate input, but got %s", str(x)) v = next(iter(x)) if v is None or math.isnan(v): return None median = self._mad_tracker.get_median() mad = self._mad_tracker.get() # not enough data points to compute median or median absolute deviation if math.isnan(mad) or math.isnan(median): return float('NaN') if abs(mad) < EPSILON: return 0.0 return abs(RobustZScore.SCALE_FACTOR * (v - median) / mad) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/detectors/robust_zscore_test.py0000644000175100017510000000301015206113032026227 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import math import unittest import apache_beam as beam from apache_beam.ml.anomaly.detectors.robust_zscore import RobustZScore class RobustZScoreTest(unittest.TestCase): input = [ beam.Row(x=1), beam.Row(x=1), beam.Row(x=5), beam.Row(x=7), beam.Row(x=20), beam.Row(x=6), beam.Row(x=1) ] def test_with_default_trackers(self): zscore = RobustZScore() scores = [] for row in RobustZScoreTest.input: scores.append(zscore.score_one(row)) zscore.learn_one(row) self.assertTrue(math.isnan(scores[0])) self.assertEqual(scores[1:], [0.0, 0.0, 0.0, 5.73325, 0.168625, 1.349]) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/detectors/zscore.py0000644000175100017510000001122515206113032023601 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import math from typing import Optional import apache_beam as beam from apache_beam.ml.anomaly.base import AnomalyDetector from apache_beam.ml.anomaly.specifiable import specifiable from apache_beam.ml.anomaly.thresholds import FixedThreshold from apache_beam.ml.anomaly.univariate.base import EPSILON from apache_beam.ml.anomaly.univariate.mean import IncSlidingMeanTracker from apache_beam.ml.anomaly.univariate.mean import MeanTracker from apache_beam.ml.anomaly.univariate.stdev import IncSlidingStdevTracker from apache_beam.ml.anomaly.univariate.stdev import StdevTracker DEFAULT_WINDOW_SIZE = 1000 # pylint: disable=line-too-long @specifiable class ZScore(AnomalyDetector): """Z-Score anomaly detector. This class implements an anomaly detection algorithm based on Z-Score (also known as Standard Score [#]_ ), which measures how many standard deviations a data point is from the mean. The score is calculated as: `| (value - mean) / stdev |` Important: In the streaming setting, we use the online version of mean and standard deviation in the calculation. This implementation is adapted from the implementations within PySAD [#]_ and River [#]_: * https://github.com/selimfirat/pysad/blob/master/pysad/models/standard_absolute_deviation.py * https://github.com/online-ml/river/blob/main/river/anomaly/sad.py Args: sub_stat_tracker: Optional `MeanTracker` instance. If None, an `IncSlidingMeanTracker` with a default window size 1000 is created. stdev_tracker: Optional `StdevTracker` instance. If None, an `IncSlidingStdevTracker` with a default window size 1000 is created. threshold_criterion: threshold_criterion: Optional `ThresholdFn` to apply on the score. Defaults to `FixedThreshold(3)` due to the commonly used 3-sigma rule. **kwargs: Additional keyword arguments. .. [#] https://en.wikipedia.org/wiki/Standard_score .. [#] Yilmaz, Selim & Kozat, Suleyman. (2020). PySAD: A Streaming Anomaly Detection Framework in Python. 10.48550/arXiv.2009.02572. .. [#] Jacob Montiel, Max Halford, Saulo Martiello Mastelini, Geoffrey Bolmier, Raphaël Sourty, et al.. (2021). River: machine learning for streaming data in Python. Journal of Machine Learning Research, 2021, 22, pp.1-8. """ # pylint: enable=line-too-long def __init__( self, sub_stat_tracker: Optional[MeanTracker] = None, stdev_tracker: Optional[StdevTracker] = None, **kwargs): if "threshold_criterion" not in kwargs: kwargs["threshold_criterion"] = FixedThreshold(3) super().__init__(**kwargs) self._sub_stat_tracker = sub_stat_tracker or IncSlidingMeanTracker( DEFAULT_WINDOW_SIZE) self._stdev_tracker = stdev_tracker or IncSlidingStdevTracker( DEFAULT_WINDOW_SIZE) def learn_one(self, x: beam.Row) -> None: """Updates the mean and standard deviation trackers with a new data point. Args: x: A `beam.Row` containing a single numerical value. """ if len(x.__dict__) != 1: raise ValueError( "ZScore.learn_one expected univariate input, but got %s", str(x)) v = next(iter(x)) self._stdev_tracker.push(v) self._sub_stat_tracker.push(v) def score_one(self, x: beam.Row) -> Optional[float]: """Scores a data point using the Z-Score. Args: x: A `beam.Row` containing a single numerical value. Returns: float | None: The Z-Score. """ if len(x.__dict__) != 1: raise ValueError( "ZScore.score_one expected univariate input, but got %s", str(x)) v = next(iter(x)) if v is None or math.isnan(v): return None sub_stat = self._sub_stat_tracker.get() stdev = self._stdev_tracker.get() # not enough data points to compute sub_stat or standard deviation if math.isnan(stdev) or math.isnan(sub_stat): return float('NaN') if abs(stdev) < EPSILON: return 0.0 return abs((v - sub_stat) / stdev) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/detectors/zscore_test.py0000644000175100017510000000453215206113032024643 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import math import unittest import apache_beam as beam from apache_beam.ml.anomaly.detectors.zscore import ZScore from apache_beam.ml.anomaly.univariate.mean import IncSlidingMeanTracker from apache_beam.ml.anomaly.univariate.stdev import IncSlidingStdevTracker class ZScoreTest(unittest.TestCase): input = [ beam.Row(x=1), beam.Row(x=1), beam.Row(x=5), beam.Row(x=9), beam.Row(x=20), beam.Row(x=10), beam.Row(x=1) ] def test_with_default_trackers(self): zscore = ZScore() scores = [] for row in ZScoreTest.input: scores.append(zscore.score_one(row)) zscore.learn_one(row) self.assertTrue(math.isnan(scores[0])) self.assertTrue(math.isnan(scores[1])) self.assertEqual( scores[2:], [ 0.0, 2.8867513459481287, 4.177863742936748, 0.35502819053868157, 0.932910509720565 ]) def test_with_custom_mean_tracker(self): zscore = ZScore( sub_stat_tracker=IncSlidingMeanTracker(3), stdev_tracker=IncSlidingStdevTracker(3)) scores = [] for row in ZScoreTest.input: scores.append(zscore.score_one(row)) zscore.learn_one(row) self.assertTrue(math.isnan(scores[0])) self.assertTrue(math.isnan(scores[1])) self.assertEqual( scores[2:], [ 0.0, 2.8867513459481287, 3.75, 0.17165643016914964, 1.9727878476642864 ]) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/specifiable.py0000644000175100017510000003447215206113032022557 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A module that provides utilities to turn a class into a Specifiable subclass. """ from __future__ import annotations import collections import dataclasses import inspect import logging import os from collections.abc import Callable from typing import Any from typing import Optional from typing import Protocol from typing import TypeVar from typing import Union from typing import overload from typing import runtime_checkable from typing_extensions import Self __all__ = ["Spec", "Specifiable", "specifiable"] _FALLBACK_SUBSPACE = "*" _ACCEPTED_SUBSPACES = [ "EnsembleAnomalyDetector", "AnomalyDetector", "BaseTracker", "ThresholdFn", "AggregationFn", _FALLBACK_SUBSPACE, ] #: A nested dictionary for efficient lookup of Specifiable subclasses. #: Structure: `_KNOWN_SPECIFIABLE[subspace][spec_type]`, where `subspace` is one #: of the accepted subspaces that the class belongs to and `spec_type` is the #: class name by default. Users can also specify a different value for #: `spec_type` when applying the `specifiable` decorator to an existing class. _KNOWN_SPECIFIABLE = collections.defaultdict(dict) T = TypeVar('T', bound=type) BUILTIN_TYPES_IN_SPEC = (int, float, complex, str, bytes, bytearray) def _class_to_subspace(cls: type) -> str: """ Search the class hierarchy to find the subspace: the closest ancestor class in the class's method resolution order (MRO) whose name is found in the accepted subspace list. This is usually called when registering a new specifiable class. """ if hasattr(cls, "mro"): # some classes do not have "mro", such as functions. for c in cls.mro(): if c.__name__ in _ACCEPTED_SUBSPACES: return c.__name__ return _FALLBACK_SUBSPACE def _spec_type_to_subspace(spec_type: str) -> str: """ Look for the subspace for a spec type. This is usually called to retrieve the subspace of a registered specifiable class. """ for subspace in _ACCEPTED_SUBSPACES: if spec_type in _KNOWN_SPECIFIABLE[subspace]: return subspace raise ValueError(f"subspace for {spec_type} not found.") @dataclasses.dataclass(frozen=True) class Spec(): """ Dataclass for storing specifications of specifiable objects. Objects can be initialized using the data in their corresponding spec. """ #: A string indicating the concrete `Specifiable` class type: str #: An optional dictionary of keyword arguments for the `__init__` method of #: the class. If None, when we materialize this Spec, we only return the #: class without instantiate any objects from it. config: Optional[dict[str, Any]] = dataclasses.field(default_factory=dict) def _specifiable_from_spec_helper(v, _run_init): if isinstance(v, Spec): return Specifiable.from_spec(v, _run_init) if isinstance(v, list): return [_specifiable_from_spec_helper(e, _run_init) for e in v] # TODO: support spec treatment for more types if not isinstance(v, BUILTIN_TYPES_IN_SPEC): logging.warning( "Type %s is not a recognized supported type for the " "specification. It will be included without conversion.", str(type(v))) return v def _specifiable_to_spec_helper(v): if isinstance(v, Specifiable): return v.to_spec() if isinstance(v, list): return [_specifiable_to_spec_helper(e) for e in v] if inspect.isfunction(v): if not hasattr(v, "spec_type"): _register(v, inject_spec_type=False) return Spec(type=_get_default_spec_type(v), config=None) if inspect.isclass(v): if not hasattr(v, "spec_type"): _register(v, inject_spec_type=False) return Spec(type=_get_default_spec_type(v), config=None) # TODO: support spec treatment for more types if not isinstance(v, BUILTIN_TYPES_IN_SPEC): logging.warning( "Type %s is not a recognized supported type for the " "specification. It will be included without conversion.", str(type(v))) return v @runtime_checkable class Specifiable(Protocol): """Protocol that a specifiable class needs to implement.""" @classmethod def spec_type(cls) -> str: pass @classmethod def from_spec(cls, spec: Spec, _run_init: bool = True) -> Union[Self, type[Self]]: """Generate a `Specifiable` subclass object based on a spec. Args: spec: the specification of a `Specifiable` subclass object _run_init: whether to call `__init__` or not for the initial instantiation Returns: Self: the `Specifiable` subclass object """ if spec.type is None: raise ValueError(f"Spec type not found in {spec}") subspace = _spec_type_to_subspace(spec.type) subclass: type[Self] = _KNOWN_SPECIFIABLE[subspace].get(spec.type, None) if subclass is None: raise ValueError(f"Unknown spec type '{spec.type}' in {spec}") if spec.config is None: # when functions or classes are used as arguments, we won't try to # create an instance. return subclass kwargs = { k: _specifiable_from_spec_helper(v, _run_init) for k, v in spec.config.items() } if _run_init: kwargs["_run_init"] = True return subclass(**kwargs) def to_spec(self) -> Spec: """Generate a spec from a `Specifiable` subclass object. Returns: Spec: The specification of the instance. """ if getattr(type(self), 'spec_type', None) is None: raise ValueError( f"'{type(self).__name__}' not registered as Specifiable. " f"Decorate ({type(self).__name__}) with @specifiable") args = { k: _specifiable_to_spec_helper(v) for k, v in self.init_kwargs.items() } return Spec(type=self.spec_type(), config=args) def run_original_init(self) -> None: """Invoke the original __init__ method with original keyword arguments""" pass @classmethod def unspecifiable(cls) -> None: """Resume the class structure prior to specifiable""" pass def _get_default_spec_type(cls): spec_type = cls.__name__ if inspect.isfunction(cls) and cls.__name__ == "": # for lambda functions, we need to include more information to distinguish # among them spec_type = '' % ( os.path.basename(cls.__code__.co_filename), cls.__code__.co_firstlineno) return spec_type # Register a `Specifiable` subclass in `KNOWN_SPECIFIABLE` def _register(cls: type, spec_type=None, inject_spec_type=True) -> None: assert spec_type is None or inject_spec_type, \ "need to inject spec_type to class if spec_type is not None" if spec_type is None: # Use default spec_type for a class if users do not specify one. spec_type = _get_default_spec_type(cls) subspace = _class_to_subspace(cls) if spec_type in _KNOWN_SPECIFIABLE[subspace]: if cls is not _KNOWN_SPECIFIABLE[subspace][spec_type]: # only raise exception if we register the same spec type with a different # class raise ValueError( f"{spec_type} is already registered for " f"specifiable class {_KNOWN_SPECIFIABLE[subspace][spec_type]}. " "Please specify a different spec_type by @specifiable(spec_type=...)." ) else: _KNOWN_SPECIFIABLE[subspace][spec_type] = cls if inject_spec_type: setattr(cls, cls.__name__ + '__spec_type', spec_type) # cls.__spec_type = spec_type # Keep a copy of arguments that are used to call the `__init__` method when the # object is initialized. def _get_init_kwargs(inst, init_method, *args, **kwargs): params = dict( zip(inspect.signature(init_method).parameters.keys(), (None, ) + args)) del params['self'] params.update(**kwargs) return params @overload def specifiable( my_cls: None = None, /, *, spec_type: Optional[str] = None, on_demand_init: bool = True, just_in_time_init: bool = True) -> Callable[[T], T]: pass @overload def specifiable( my_cls: T, /, *, spec_type: Optional[str] = None, on_demand_init: bool = True, just_in_time_init: bool = True) -> T: pass def specifiable( my_cls: Optional[T] = None, /, *, spec_type: Optional[str] = None, on_demand_init: bool = True, just_in_time_init: bool = True) -> Union[T, Callable[[T], T]]: """A decorator that turns a class into a `Specifiable` subclass by implementing the `Specifiable` protocol. To use the decorator, simply place `@specifiable` before the class definition:: @specifiable class Foo(): ... For finer control, the decorator can accept arguments:: @specifiable(spec_type="My Class", on_demand_init=False) class Bar(): ... Args: spec_type: The value of the `type` field in the Spec of a `Specifiable` subclass. If not provided, the class name is used. This argument is useful when registering multiple classes with the same base name; in such cases, one can specify `spec_type` to different values to resolve conflict. on_demand_init: If True, allow on-demand object initialization. The original `__init__` method will be called when `_run_init=True` is passed to the object's initialization function. just_in_time_init: If True, allow just-in-time object initialization. The original `__init__` method will be called when the first time an attribute is accessed. """ def _wrapper(cls: T) -> T: def new_init(self, *args, **kwargs): self._initialized = False self._in_init = False run_init_request = False if "_run_init" in kwargs: run_init_request = kwargs["_run_init"] del kwargs["_run_init"] if 'init_kwargs' not in self.__dict__: # If it is a child specifiable (i.e.g init_kwargs not set), we determine # whether to skip the original __init__ call based on options: # on_demand_init, just_in_time_init and _run_init. # Otherwise (i.e. init_kwargs is set), we always call the original # __init__ method for ancestor specifiable. self.init_kwargs = _get_init_kwargs( self, original_init, *args, **kwargs) logging.debug("Record init params in %s.new_init", class_name) if (on_demand_init and not run_init_request) or \ (not on_demand_init and just_in_time_init): logging.debug("Skip original %s.__init__", class_name) return logging.debug("Call original %s.__init__ in new_init", class_name) original_init(self, *args, **kwargs) self._initialized = True def run_original_init(self) -> None: """Execute the original `__init__` method with its saved arguments. For instances of the `Specifiable` class, initialization is deferred (lazy initialization). This function forces the execution of the original `__init__` method using the arguments captured during the object's initial instantiation. """ self._in_init = True original_init(self, **self.init_kwargs) self._in_init = False self._initialized = True # __getattr__ is only called when an attribute is not found in the object def new_getattr(self, name): logging.debug( "Trying to access %s.%s, but it is not found.", class_name, name) # Fix the infinite loop issue when pickling a Specifiable if name in ["_in_init", "__getstate__"] and name not in self.__dict__: raise AttributeError( f"'{type(self).__name__}' object has no attribute '{name}'") # If the attribute is not found during or after initialization, then # it is a missing attribute. if self._in_init or self._initialized: raise AttributeError( f"'{type(self).__name__}' object has no attribute '{name}'") # Here, we know the object is not initialized, then we will call original # init method. logging.debug("Call original %s.__init__ in new_getattr", class_name) run_original_init(self) # __getattribute__ is call for every attribute regardless whether it is # present in the object. In this case, we don't cause an infinite loop # if the attribute does not exist. logging.debug( "Call original %s.__getattribute__(%s) in new_getattr", class_name, name) return self.__getattribute__(name) def spec_type_func(cls): return getattr(cls, spec_type_attr_name) def unspecifiable(cls): delattr(cls, spec_type_attr_name) cls.__init__ = original_init if just_in_time_init: delattr(cls, '__getattr__') delattr(cls, 'spec_type') delattr(cls, 'run_original_init') delattr(cls, 'to_spec') delattr(cls, 'from_spec') delattr(cls, 'unspecifiable') spec_type_attr_name = cls.__name__ + "__spec_type" # the class is registered if hasattr(cls, spec_type_attr_name): return cls # start of the function body of _wrapper _register(cls, spec_type) class_name = cls.__name__ original_init = cls.__init__ # type: ignore[misc] cls.__init__ = new_init # type: ignore[misc] if just_in_time_init: cls.__getattr__ = new_getattr cls.spec_type = classmethod(spec_type_func) cls.run_original_init = run_original_init cls.to_spec = Specifiable.to_spec cls.from_spec = Specifiable.from_spec cls.unspecifiable = classmethod(unspecifiable) return cls # end of the function body of _wrapper # When this decorator is called with arguments, i.e.. # "@specifiable(arg1=...,arg2=...)", it is equivalent to assigning # specifiable(arg1=..., arg2=...) to a variable, say decor_func, and then # calling "@decor_func". if my_cls is None: return _wrapper # When this decorator is called without an argument, i.e. "@specifiable", # we return the augmented class. return _wrapper(my_cls) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/specifiable_test.py0000644000175100017510000004650115206113032023612 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import copy import dataclasses import logging import os import unittest from typing import Optional import pytest from parameterized import parameterized from apache_beam.internal.cloudpickle import cloudpickle from apache_beam.ml.anomaly.specifiable import _FALLBACK_SUBSPACE from apache_beam.ml.anomaly.specifiable import _KNOWN_SPECIFIABLE from apache_beam.ml.anomaly.specifiable import Spec from apache_beam.ml.anomaly.specifiable import Specifiable from apache_beam.ml.anomaly.specifiable import specifiable class TestSpecifiable(unittest.TestCase): def setUp(self) -> None: self.saved_specifiable = copy.deepcopy(_KNOWN_SPECIFIABLE) def tearDown(self) -> None: _KNOWN_SPECIFIABLE.clear() _KNOWN_SPECIFIABLE.update(self.saved_specifiable) def test_decorator_in_function_form(self): class A(): pass class B(): pass # class is not decorated and thus not registered self.assertNotIn("A", _KNOWN_SPECIFIABLE[_FALLBACK_SUBSPACE]) # apply the decorator function to an existing class A = specifiable(A) self.assertEqual(A.spec_type(), "A") self.assertTrue(isinstance(A(), Specifiable)) self.assertIn("A", _KNOWN_SPECIFIABLE[_FALLBACK_SUBSPACE]) self.assertEqual(_KNOWN_SPECIFIABLE[_FALLBACK_SUBSPACE]["A"], A) # Re-registering spec_type with the same class is allowed A = specifiable(A) # Raise an error when re-registering spec_type with a different class self.assertRaises(ValueError, specifiable(spec_type='A'), B) # Applying the decorator function to an existing class with a different # spec_type will have no effect. A = specifiable(spec_type="A_DUP")(A) self.assertEqual(A.spec_type(), "A") def test_decorator_in_syntactic_sugar_form(self): # call decorator without parameters @specifiable class B(): pass self.assertTrue(isinstance(B(), Specifiable)) self.assertIn("B", _KNOWN_SPECIFIABLE[_FALLBACK_SUBSPACE]) self.assertEqual(_KNOWN_SPECIFIABLE[_FALLBACK_SUBSPACE]["B"], B) # call decorator with parameters @specifiable(spec_type="C_TYPE") class C(): pass self.assertTrue(isinstance(C(), Specifiable)) self.assertIn("C_TYPE", _KNOWN_SPECIFIABLE[_FALLBACK_SUBSPACE]) self.assertEqual(_KNOWN_SPECIFIABLE[_FALLBACK_SUBSPACE]["C_TYPE"], C) def test_init_params_in_specifiable(self): @specifiable class ParentWithInitParams(): def __init__(self, arg_1, arg_2=2, arg_3="3", **kwargs): pass parent = ParentWithInitParams(10, arg_3="30", arg_4=40) assert isinstance(parent, Specifiable) self.assertEqual( parent.init_kwargs, { 'arg_1': 10, 'arg_3': '30', 'arg_4': 40 }) # inheritance of a Specifiable subclass @specifiable class ChildWithInitParams(ParentWithInitParams): def __init__(self, new_arg_1, new_arg_2=200, new_arg_3="300", **kwargs): super().__init__(**kwargs) child = ChildWithInitParams( 1000, arg_1=11, arg_2=20, new_arg_2=2000, arg_4=4000) assert isinstance(child, Specifiable) self.assertEqual( child.init_kwargs, { 'new_arg_1': 1000, 'arg_1': 11, 'arg_2': 20, 'new_arg_2': 2000, 'arg_4': 4000 }) # composite of Specifiable subclasses @specifiable class CompositeWithInitParams(): def __init__( self, my_parent: Optional[ParentWithInitParams] = None, my_child: Optional[ChildWithInitParams] = None): pass composite = CompositeWithInitParams(parent, child) assert isinstance(composite, Specifiable) self.assertEqual( composite.init_kwargs, { 'my_parent': parent, 'my_child': child }) def test_from_spec_on_unknown_spec_type(self): self.assertRaises(ValueError, Specifiable.from_spec, Spec(type="unknown")) # To test from_spec and to_spec with/without just_in_time_init. @parameterized.expand([(False, False), (True, False), (False, True), (True, True)]) def test_from_spec_and_to_spec(self, on_demand_init, just_in_time_init): @specifiable( on_demand_init=on_demand_init, just_in_time_init=just_in_time_init) @dataclasses.dataclass class Product(): name: str price: float @specifiable( on_demand_init=on_demand_init, just_in_time_init=just_in_time_init) class Entry(): def __init__(self, product: Product, quantity: int = 1): self._product = product self._quantity = quantity def __eq__(self, value) -> bool: return isinstance(value, Entry) and \ self._product == value._product and \ self._quantity == value._quantity @specifiable( on_demand_init=on_demand_init, just_in_time_init=just_in_time_init) @dataclasses.dataclass class ShoppingCart(): user_id: str entries: list[Entry] orange = Product("orange", 1.0) expected_orange_spec = Spec( "Product", config={ 'name': 'orange', 'price': 1.0 }) assert isinstance(orange, Specifiable) self.assertEqual(orange.to_spec(), expected_orange_spec) entry_1 = Entry(product=orange) expected_entry_spec_1 = Spec( "Entry", config={ 'product': expected_orange_spec, }) assert isinstance(entry_1, Specifiable) self.assertEqual(entry_1.to_spec(), expected_entry_spec_1) banana = Product("banana", 0.5) expected_banana_spec = Spec( "Product", config={ 'name': 'banana', 'price': 0.5 }) entry_2 = Entry(product=banana, quantity=5) expected_entry_spec_2 = Spec( "Entry", config={ 'product': expected_banana_spec, 'quantity': 5 }) shopping_cart = ShoppingCart(user_id="test", entries=[entry_1, entry_2]) expected_shopping_cart_spec = Spec( "ShoppingCart", config={ "user_id": "test", "entries": [expected_entry_spec_1, expected_entry_spec_2] }) assert isinstance(shopping_cart, Specifiable) self.assertEqual(shopping_cart.to_spec(), expected_shopping_cart_spec) if on_demand_init and not just_in_time_init: orange.run_original_init() banana.run_original_init() entry_1.run_original_init() entry_2.run_original_init() shopping_cart.run_original_init() self.assertEqual(Specifiable.from_spec(expected_orange_spec), orange) self.assertEqual(Specifiable.from_spec(expected_entry_spec_1), entry_1) self.assertEqual( Specifiable.from_spec(expected_shopping_cart_spec), shopping_cart) class TestInitCallCount(unittest.TestCase): def setUp(self) -> None: self.saved_specifiable = copy.deepcopy(_KNOWN_SPECIFIABLE) def tearDown(self) -> None: _KNOWN_SPECIFIABLE.clear() _KNOWN_SPECIFIABLE.update(self.saved_specifiable) def test_on_demand_init(self): @specifiable(on_demand_init=True, just_in_time_init=False) class FooOnDemand(): counter = 0 def __init__(self, arg): self.my_arg = arg * 10 FooOnDemand.counter += 1 # increment it when __init__ is called foo = FooOnDemand(123) self.assertEqual(FooOnDemand.counter, 0) self.assertIn("init_kwargs", foo.__dict__) self.assertEqual(foo.__dict__["init_kwargs"], {"arg": 123}) self.assertNotIn("my_arg", foo.__dict__) self.assertRaises(AttributeError, getattr, foo, "my_arg") self.assertRaises(AttributeError, lambda: foo.my_arg) self.assertRaises(AttributeError, getattr, foo, "unknown_arg") self.assertRaises(AttributeError, lambda: foo.unknown_arg) self.assertEqual(FooOnDemand.counter, 0) # __init__ is called when _run_init=True is used foo_2 = FooOnDemand(456, _run_init=True) self.assertEqual(FooOnDemand.counter, 1) self.assertIn("init_kwargs", foo_2.__dict__) self.assertEqual(foo_2.__dict__["init_kwargs"], {"arg": 456}) self.assertIn("my_arg", foo_2.__dict__) self.assertEqual(foo_2.my_arg, 4560) self.assertEqual(FooOnDemand.counter, 1) def test_just_in_time_init(self): @specifiable(on_demand_init=False, just_in_time_init=True) class FooJustInTime(): counter = 0 def __init__(self, arg): self.my_arg = arg * 10 FooJustInTime.counter += 1 # increment it when __init__ is called foo = FooJustInTime(321) self.assertEqual(FooJustInTime.counter, 0) self.assertIn("init_kwargs", foo.__dict__) self.assertEqual(foo.__dict__["init_kwargs"], {"arg": 321}) # __init__ hasn't been called yet self.assertNotIn("my_arg", foo.__dict__) self.assertEqual(FooJustInTime.counter, 0) # __init__ is called when trying to access a class attribute self.assertEqual(foo.my_arg, 3210) self.assertEqual(FooJustInTime.counter, 1) self.assertRaises(AttributeError, lambda: foo.unknown_arg) self.assertEqual(FooJustInTime.counter, 1) def test_on_demand_and_just_in_time_init(self): @specifiable(on_demand_init=True, just_in_time_init=True) class FooOnDemandAndJustInTime(): counter = 0 def __init__(self, arg): self.my_arg = arg * 10 FooOnDemandAndJustInTime.counter += 1 foo = FooOnDemandAndJustInTime(987) self.assertEqual(FooOnDemandAndJustInTime.counter, 0) self.assertIn("init_kwargs", foo.__dict__) self.assertEqual(foo.__dict__["init_kwargs"], {"arg": 987}) self.assertNotIn("my_arg", foo.__dict__) self.assertEqual(FooOnDemandAndJustInTime.counter, 0) # __init__ is called when trying to access a class attribute self.assertEqual(foo.my_arg, 9870) self.assertEqual(FooOnDemandAndJustInTime.counter, 1) # __init__ is called when _run_init=True is used foo_2 = FooOnDemandAndJustInTime(789, _run_init=True) self.assertEqual(FooOnDemandAndJustInTime.counter, 2) self.assertIn("init_kwargs", foo_2.__dict__) self.assertEqual(foo_2.__dict__["init_kwargs"], {"arg": 789}) self.assertEqual(FooOnDemandAndJustInTime.counter, 2) # __init__ is NOT called after it is initialized self.assertEqual(foo_2.my_arg, 7890) self.assertEqual(FooOnDemandAndJustInTime.counter, 2) @specifiable(on_demand_init=True, just_in_time_init=True) class FooForPickle(): counter = 0 def __init__(self, arg): self.my_arg = arg * 10 type(self).counter += 1 @pytest.mark.uses_dill def test_on_dill_pickle(self): pytest.importorskip("dill") FooForPickle = TestInitCallCount.FooForPickle import dill FooForPickle.counter = 0 foo = FooForPickle(456) self.assertEqual(FooForPickle.counter, 0) new_foo = dill.loads(dill.dumps(foo)) self.assertEqual(FooForPickle.counter, 0) self.assertEqual(new_foo.__dict__, foo.__dict__) self.assertEqual(foo.my_arg, 4560) self.assertEqual(FooForPickle.counter, 1) new_foo_2 = dill.loads(dill.dumps(foo)) self.assertEqual(FooForPickle.counter, 1) self.assertEqual(new_foo_2.__dict__, foo.__dict__) def test_on_pickle(self): FooForPickle = TestInitCallCount.FooForPickle # Note that pickle does not support classes/functions nested in a function. import pickle FooForPickle.counter = 0 foo = FooForPickle(456) self.assertEqual(FooForPickle.counter, 0) new_foo = pickle.loads(pickle.dumps(foo)) self.assertEqual(FooForPickle.counter, 0) self.assertEqual(new_foo.__dict__, foo.__dict__) self.assertEqual(foo.my_arg, 4560) self.assertEqual(FooForPickle.counter, 1) new_foo_2 = pickle.loads(pickle.dumps(foo)) self.assertEqual(FooForPickle.counter, 1) self.assertEqual(new_foo_2.__dict__, foo.__dict__) FooForPickle.counter = 0 foo = FooForPickle(456) self.assertEqual(FooForPickle.counter, 0) new_foo = cloudpickle.loads(cloudpickle.dumps(foo)) self.assertEqual(FooForPickle.counter, 0) self.assertEqual(new_foo.__dict__, foo.__dict__) self.assertEqual(foo.my_arg, 4560) self.assertEqual(FooForPickle.counter, 1) new_foo_2 = cloudpickle.loads(cloudpickle.dumps(foo)) self.assertEqual(FooForPickle.counter, 1) self.assertEqual(new_foo_2.__dict__, foo.__dict__) @specifiable class Parent(): counter = 0 parent_class_var = 1000 def __init__(self, p): self.parent_inst_var = p * 10 Parent.counter += 1 @specifiable class Child_1(Parent): counter = 0 child_class_var = 2001 def __init__(self, c): super().__init__(c) self.child_inst_var = c + 1 Child_1.counter += 1 @specifiable class Child_2(Parent): counter = 0 child_class_var = 2001 def __init__(self, c): self.child_inst_var = c + 1 super().__init__(c) Child_2.counter += 1 @specifiable class Child_Error_1(Parent): counter = 0 child_class_var = 2001 def __init__(self, c): # read an instance var in child that doesn't exist self.child_inst_var += 1 super().__init__(c) Child_2.counter += 1 @specifiable class Child_Error_2(Parent): counter = 0 child_class_var = 2001 def __init__(self, c): # read an instance var in parent without calling parent's __init__. self.parent_inst_var += 1 Child_2.counter += 1 class TestNestedSpecifiable(unittest.TestCase): def setUp(self) -> None: self.saved_specifiable = copy.deepcopy(_KNOWN_SPECIFIABLE) def tearDown(self) -> None: _KNOWN_SPECIFIABLE.clear() _KNOWN_SPECIFIABLE.update(self.saved_specifiable) @parameterized.expand([[Child_1, 0], [Child_2, 0], [Child_1, 1], [Child_2, 1], [Child_1, 2], [Child_2, 2]]) def test_nested_specifiable(self, Child, mode): Parent.counter = 0 Child.counter = 0 child = Child(5) self.assertEqual(Parent.counter, 0) self.assertEqual(Child.counter, 0) # accessing class vars won't trigger __init__ self.assertEqual(child.parent_class_var, 1000) self.assertEqual(child.child_class_var, 2001) self.assertEqual(Parent.counter, 0) self.assertEqual(Child.counter, 0) # accessing instance var will trigger __init__ if mode == 0: self.assertEqual(child.parent_inst_var, 50) elif mode == 1: self.assertEqual(child.child_inst_var, 6) else: self.assertRaises(AttributeError, lambda: child.unknown_var) self.assertEqual(Parent.counter, 1) self.assertEqual(Child.counter, 1) # after initialization, it won't trigger __init__ again self.assertEqual(child.parent_inst_var, 50) self.assertEqual(child.child_inst_var, 6) self.assertRaises(AttributeError, lambda: child.unknown_var) self.assertEqual(Parent.counter, 1) self.assertEqual(Child.counter, 1) def test_error_in_child(self): Parent.counter = 0 child_1 = Child_Error_1(5) self.assertEqual(child_1.child_class_var, 2001) # error during child initialization self.assertRaises(AttributeError, lambda: child_1.child_inst_var) self.assertEqual(Parent.counter, 0) self.assertEqual(Child_1.counter, 0) child_2 = Child_Error_2(5) self.assertEqual(child_2.child_class_var, 2001) # error during child initialization self.assertRaises(AttributeError, lambda: child_2.parent_inst_var) self.assertEqual(Parent.counter, 0) self.assertEqual(Child_2.counter, 0) def my_normal_func(x, y): return x + y @specifiable class Wrapper(): def __init__(self, func=None, cls=None, **kwargs): self._func = func if cls is not None: self._cls = cls(**kwargs) def run_func(self, x, y): return self._func(x, y) def run_func_in_class(self, x, y): return self._cls.apply(x, y) class TestFunctionAsArgument(unittest.TestCase): def setUp(self) -> None: self.saved_specifiable = copy.deepcopy(_KNOWN_SPECIFIABLE) def tearDown(self) -> None: _KNOWN_SPECIFIABLE.clear() _KNOWN_SPECIFIABLE.update(self.saved_specifiable) def test_normal_function(self): w = Wrapper(my_normal_func) self.assertEqual(w.run_func(1, 2), 3) w_spec = w.to_spec() self.assertEqual( w_spec, Spec( type='Wrapper', config={'func': Spec(type="my_normal_func", config=None)})) w_2 = Specifiable.from_spec(w_spec) self.assertEqual(w_2.run_func(2, 3), 5) def test_lambda_function(self): my_lambda_func = lambda x, y: x - y w = Wrapper(my_lambda_func) self.assertEqual(w.run_func(3, 2), 1) w_spec = w.to_spec() self.assertEqual( w_spec, Spec( type='Wrapper', config={ 'func': Spec( type= f"", # pylint: disable=line-too-long config=None) })) w_2 = Specifiable.from_spec(w_spec) self.assertEqual(w_2.run_func(5, 3), 2) class TestClassAsArgument(unittest.TestCase): def setUp(self) -> None: self.saved_specifiable = copy.deepcopy(_KNOWN_SPECIFIABLE) def tearDown(self) -> None: _KNOWN_SPECIFIABLE.clear() _KNOWN_SPECIFIABLE.update(self.saved_specifiable) def test_normal_class(self): class InnerClass(): def __init__(self, multiplier): self._multiplier = multiplier def apply(self, x, y): return x * y * self._multiplier w = Wrapper(cls=InnerClass, multiplier=10) self.assertEqual(w.run_func_in_class(2, 3), 60) w_spec = w.to_spec() self.assertEqual( w_spec, Spec( type='Wrapper', config={ 'cls': Spec(type='InnerClass', config=None), 'multiplier': 10 })) w_2 = Specifiable.from_spec(w_spec) self.assertEqual(w_2.run_func_in_class(5, 3), 150) class TestUncommonUsages(unittest.TestCase): def test_double_specifiable(self): @specifiable @specifiable class ZZ(): def __init__(self, a): self.a = a assert issubclass(ZZ, Specifiable) c = ZZ("b") c.run_original_init() self.assertEqual(c.a, "b") def test_unspecifiable(self): class YY(): def __init__(self, x): self.x = x assert False YY = specifiable(YY) assert issubclass(YY, Specifiable) y = YY(1) # __init__ is called (with assertion error raised) when attribute is first # accessed self.assertRaises(AssertionError, lambda: y.x) # unspecifiable YY YY.unspecifiable() # __init__ is called immediately self.assertRaises(AssertionError, YY, 1) self.assertFalse(hasattr(YY, 'run_original_init')) self.assertFalse(hasattr(YY, 'spec_type')) self.assertFalse(hasattr(YY, 'to_spec')) self.assertFalse(hasattr(YY, 'from_spec')) self.assertFalse(hasattr(YY, 'unspecifiable')) # make YY specifiable again YY = specifiable(YY) assert issubclass(YY, Specifiable) y = YY(1) self.assertRaises(AssertionError, lambda: y.x) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/thresholds.py0000644000175100017510000001436015206113032022462 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from __future__ import annotations import math from typing import Optional from apache_beam.ml.anomaly.base import ThresholdFn from apache_beam.ml.anomaly.specifiable import specifiable from apache_beam.ml.anomaly.univariate.quantile import BufferedSlidingQuantileTracker # pylint: disable=line-too-long from apache_beam.ml.anomaly.univariate.quantile import QuantileTracker @specifiable class FixedThreshold(ThresholdFn): """Applies a fixed cutoff value to anomaly scores. This `ThresholdFn` is stateless and uses a pre-defined cutoff value to classify anomaly scores. Scores below the cutoff are considered normal, while scores at or above the cutoff are classified as outliers. Args: cutoff (float): The fixed threshold value. Anomaly scores at or above this value will be labeled as outliers. **kwargs: Additional keyword arguments to be passed to the base `ThresholdFn` constructor. """ def __init__(self, cutoff: float, **kwargs): super().__init__(**kwargs) self._cutoff = cutoff @property def is_stateful(self) -> bool: """Indicates whether this ThresholdFn is stateful. Returns: bool: Always False for `FixedThreshold` as it is stateless. """ return False @property def threshold(self) -> float: """Returns the fixed cutoff threshold value. Returns: float: The fixed threshold value. """ return self._cutoff def apply(self, score: Optional[float]) -> Optional[int]: """Applies the fixed threshold to an anomaly score. Classifies the given anomaly score as normal or outlier based on the predefined cutoff. Args: score (Optional[float]): The input anomaly score. Returns: Optional[int]: The anomaly label: - `normal_label` if the score is less than the threshold. - `outlier_label` if the score is at or above the threshold. - `missing_label` if the score is `NaN` (detector not ready). - `None` if the score is `None` (detector ready, but unable to produce score). """ # score error: detector is ready but is unable to produce the score due to # errors such as ill-formatted input. if score is None: return None # score missing: detector is not yet ready if math.isnan(score): return self._missing_label if score < self.threshold: return self._normal_label return self._outlier_label @specifiable class QuantileThreshold(ThresholdFn): """Applies a quantile-based dynamic threshold to anomaly scores. This `ThresholdFn` is stateful and uses a quantile tracker to dynamically determine the threshold for anomaly detection. It estimates the specified quantile of the incoming anomaly scores and uses this quantile value as the threshold. The threshold adapts over time as more data is processed, making it suitable for scenarios where the distribution of anomaly scores may change. Args: quantile (Optional[float]): The quantile to be tracked (e.g., 0.95 for the 95th percentile). This value determines the dynamic threshold. Defaults to 0.95. quantile_tracker (Optional[BufferedQuantileTracker]): An optional pre-initialized quantile tracker. If provided, this tracker will be used; otherwise, a `BufferedSlidingQuantileTracker` will be created with a default window size of 100. **kwargs: Additional keyword arguments to be passed to the base `ThresholdFn` constructor. """ def __init__( self, quantile: Optional[float] = 0.95, quantile_tracker: Optional[QuantileTracker] = None, **kwargs): super().__init__(**kwargs) if quantile_tracker is not None: self._tracker = quantile_tracker else: self._tracker = BufferedSlidingQuantileTracker( window_size=100, q=quantile) @property def is_stateful(self) -> bool: """Indicates whether this ThresholdFn is stateful. Returns: bool: Always True for `QuantileThreshold` as it is stateful. """ return True @property def threshold(self) -> float: """Returns the current quantile-based threshold value. Returns: float: The dynamically calculated threshold value based on the quantile tracker. """ return self._tracker.get() def apply(self, score: Optional[float]) -> Optional[int]: """Applies the quantile-based threshold to an anomaly score. Updates the quantile tracker with the given score and classifies the score as normal or outlier based on the current quantile threshold. Args: score (Optional[float]): The input anomaly score. Returns: Optional[int]: The anomaly label: - `normal_label` if the score is less than the threshold. - `outlier_label` if the score is at or above the threshold. - `missing_label` if the score is `NaN` (detector not ready). - `None` if the score is `None` (detector ready, but unable to produce score). """ # score error: detector is ready but is unable to produce the score due to # errors such as ill-formatted input. if score is None: # store NaN instead of None in quantile tracker to simplify tracker logic. # After all, it is only a place holder and will not be used in calculation self._tracker.push(float("NaN")) return None self._tracker.push(score) # score missing: detector is not yet ready if math.isnan(score): return self._missing_label if score < self.threshold: return self._normal_label return self._outlier_label ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/thresholds_test.py0000644000175100017510000000616215206113032023522 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest from apache_beam.ml.anomaly import thresholds from apache_beam.ml.anomaly.specifiable import Spec from apache_beam.ml.anomaly.univariate.quantile import BufferedSlidingQuantileTracker # pylint: disable=line-too-long from apache_beam.ml.anomaly.univariate.quantile import SimpleSlidingQuantileTracker # pylint: disable=line-too-long class TestFixedThreshold(unittest.TestCase): def test_apply_only(self): threshold_fn = thresholds.FixedThreshold(2) self.assertEqual(threshold_fn.apply(1.0), 0) self.assertEqual(threshold_fn.apply(2.0), 1) self.assertEqual(threshold_fn.apply(None), None) self.assertEqual(threshold_fn.apply(float('NaN')), -2) class TestQuantileThreshold(unittest.TestCase): def test_apply_only(self): threshold_fn = thresholds.QuantileThreshold(0.9) self.assertEqual(threshold_fn.apply(1.0), 1) self.assertEqual(threshold_fn.apply(2.0), 1) self.assertEqual(threshold_fn.apply(1.2), 0) self.assertEqual(threshold_fn.apply(None), None) self.assertEqual(threshold_fn.apply(float('NaN')), -2) def test_quantile_tracker(self): t1 = thresholds.QuantileThreshold() self.assertTrue(isinstance(t1._tracker, BufferedSlidingQuantileTracker)) self.assertEqual(t1._tracker._q, 0.95) self.assertEqual(t1.to_spec(), Spec("QuantileThreshold", config={})) t2 = thresholds.QuantileThreshold(quantile=0.99) self.assertTrue(isinstance(t2._tracker, BufferedSlidingQuantileTracker)) self.assertEqual(t2._tracker._q, 0.99) self.assertEqual( t2.to_spec(), Spec("QuantileThreshold", config={"quantile": 0.99})) # argument quantile=0.9 is not used because quantile_tracker is set t3 = thresholds.QuantileThreshold( quantile=0.9, quantile_tracker=SimpleSlidingQuantileTracker(50, 0.975)) self.assertTrue(isinstance(t3._tracker, SimpleSlidingQuantileTracker)) self.assertEqual(t3._tracker._q, 0.975) self.assertEqual( t3.to_spec(), Spec( "QuantileThreshold", config={ 'quantile': 0.9, 'quantile_tracker': Spec( type='SimpleSlidingQuantileTracker', config={ 'window_size': 50, 'q': 0.975 }) })) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/transforms.py0000644000175100017510000006016015206113032022500 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import dataclasses import uuid from collections.abc import Callable from collections.abc import Iterable from typing import Any from typing import Optional from typing import TypeVar from typing import Union import apache_beam as beam from apache_beam.coders import CloudpickleCoder from apache_beam.ml.anomaly import aggregations from apache_beam.ml.anomaly.base import AggregationFn from apache_beam.ml.anomaly.base import AnomalyDetector from apache_beam.ml.anomaly.base import AnomalyPrediction from apache_beam.ml.anomaly.base import AnomalyResult from apache_beam.ml.anomaly.base import EnsembleAnomalyDetector from apache_beam.ml.anomaly.base import ThresholdFn from apache_beam.ml.anomaly.detectors.offline import OfflineDetector from apache_beam.ml.anomaly.specifiable import Spec from apache_beam.ml.anomaly.specifiable import Specifiable from apache_beam.ml.inference.base import RunInference from apache_beam.transforms.userstate import ReadModifyWriteStateSpec from apache_beam.typehints.typehints import TupleConstraint KeyT = TypeVar('KeyT') TempKeyT = TypeVar('TempKeyT', bound=str) InputT = beam.Row OutputT = AnomalyResult KeyedInputT = tuple[KeyT, beam.Row] KeyedOutputT = tuple[KeyT, AnomalyResult] NestedKeyedInputT = tuple[KeyT, tuple[TempKeyT, beam.Row]] NestedKeyedOutputT = tuple[KeyT, tuple[TempKeyT, AnomalyResult]] class _ScoreAndLearnDoFn(beam.DoFn): """Scores and learns from incoming data using an anomaly detection model. This DoFn applies an anomaly detection model to score incoming data and then updates the model with the same data. It maintains the model state using Beam's state management. """ MODEL_STATE_INDEX = ReadModifyWriteStateSpec( 'saved_model', CloudpickleCoder()) def __init__(self, detector_spec: Spec): self._detector_spec = detector_spec assert isinstance(self._detector_spec.config, dict) self._detector_spec.config["_run_init"] = True def score_and_learn(self, data): """Scores and learns from a single data point. Args: data: A `beam.Row` representing the input data point. Returns: float: The anomaly score predicted by the model. """ assert self._underlying if self._underlying._features is not None: x = beam.Row(**{f: getattr(data, f) for f in self._underlying._features}) else: x = beam.Row(**data._asdict()) # score the incoming data using the existing model y_pred = self._underlying.score_one(x) # then update the model with the same data self._underlying.learn_one(x) return y_pred def process( self, element: NestedKeyedInputT, model_state=beam.DoFn.StateParam(MODEL_STATE_INDEX), **kwargs) -> Iterable[NestedKeyedOutputT]: k1, (k2, data) = element self._underlying: AnomalyDetector = model_state.read() if self._underlying is None: self._underlying = Specifiable.from_spec(self._detector_spec) yield k1, (k2, AnomalyResult( example=data, predictions=[ AnomalyPrediction( model_id=self._underlying._model_id, score=self.score_and_learn(data)) ])) model_state.write(self._underlying) class RunScoreAndLearn(beam.PTransform[beam.PCollection[NestedKeyedInputT], beam.PCollection[NestedKeyedOutputT]]): """Applies the _ScoreAndLearnDoFn to a PCollection of data. This PTransform scores and learns from data points using an anomaly detection model. Args: detector: The anomaly detection model to use. """ def __init__(self, detector: AnomalyDetector): self._detector = detector def expand( self, input: beam.PCollection[NestedKeyedInputT] ) -> beam.PCollection[NestedKeyedOutputT]: return input | beam.ParDo(_ScoreAndLearnDoFn(self._detector.to_spec())) class _BaseThresholdDoFn(beam.DoFn): """Applies a ThresholdFn to anomaly detection results. This abstract base class defines the structure for DoFns that use a `ThresholdFn` to convert anomaly scores into anomaly labels (e.g., normal or outlier). It handles the core logic of applying the threshold function and updating the prediction labels within `AnomalyResult` objects. Args: threshold_fn_spec (Spec): Specification defining the `ThresholdFn` to be used. """ def __init__(self, threshold_fn_spec: Spec): self._threshold_fn_spec = threshold_fn_spec def _apply_threshold_to_predictions( self, result: AnomalyResult) -> AnomalyResult: """Updates the prediction labels in an AnomalyResult using the ThresholdFn. Args: result (AnomalyResult): The input `AnomalyResult` containing anomaly scores. Returns: AnomalyResult: A new `AnomalyResult` with updated prediction labels and threshold values. """ predictions = [ dataclasses.replace( p, label=self._threshold_fn.apply(p.score), threshold=self._threshold_fn.threshold) for p in result.predictions ] return dataclasses.replace(result, predictions=predictions) class _StatelessThresholdDoFn(_BaseThresholdDoFn): """Applies a stateless ThresholdFn to anomaly detection results. This DoFn is designed for stateless `ThresholdFn` implementations. It initializes the `ThresholdFn` once during setup and applies it to each incoming element without maintaining any state across elements. Args: threshold_fn_spec (Spec): Specification defining the `ThresholdFn` to be used. Raises: AssertionError: If the provided `threshold_fn_spec` leads to the creation of a stateful `ThresholdFn`. """ def __init__(self, threshold_fn_spec: Spec): assert isinstance(threshold_fn_spec.config, dict) threshold_fn_spec.config["_run_init"] = True self._threshold_fn = Specifiable.from_spec(threshold_fn_spec) assert isinstance(self._threshold_fn, ThresholdFn) assert not self._threshold_fn.is_stateful, \ "This DoFn can only take stateless function as threshold_fn" def process(self, element: NestedKeyedOutputT, **kwargs) -> Iterable[NestedKeyedOutputT]: """Processes a batch of anomaly results using a stateless ThresholdFn. Args: element (tuple[Any, tuple[Any, AnomalyResult]]): A tuple representing an element in the Beam pipeline. It is expected to be in the format `(key1, (key2, AnomalyResult))`, where key1 is the original input key, and key2 is a disambiguating key for distinct data points. **kwargs: Additional keyword arguments passed to the `process` method in Beam DoFns. Yields: Iterable[tuple[Any, tuple[Any, AnomalyResult]]]: An iterable containing a single output element with the same structure as the input, but with the `AnomalyResult` having updated prediction labels based on the stateless `ThresholdFn`. """ k1, (k2, result) = element yield k1, (k2, self._apply_threshold_to_predictions(result)) class _StatefulThresholdDoFn(_BaseThresholdDoFn): """Applies a stateful ThresholdFn to anomaly detection results. This DoFn is designed for stateful `ThresholdFn` implementations. It leverages Beam's state management to persist and update the state of the `ThresholdFn` across multiple elements. This is necessary for `ThresholdFn`s that need to accumulate information or adapt over time, such as quantile-based thresholds. Args: threshold_fn_spec (Spec): Specification defining the `ThresholdFn` to be used. Raises: AssertionError: If the provided `threshold_fn_spec` leads to the creation of a stateless `ThresholdFn`. """ THRESHOLD_STATE_INDEX = ReadModifyWriteStateSpec( 'saved_tracker', CloudpickleCoder()) def __init__(self, threshold_fn_spec: Spec): assert isinstance(threshold_fn_spec.config, dict) threshold_fn_spec.config["_run_init"] = True threshold_fn = Specifiable.from_spec(threshold_fn_spec) assert isinstance(threshold_fn, ThresholdFn) assert threshold_fn.is_stateful, \ "This DoFn can only take stateful function as threshold_fn" self._threshold_fn_spec = threshold_fn_spec def process( self, element: NestedKeyedOutputT, threshold_state=beam.DoFn.StateParam(THRESHOLD_STATE_INDEX), **kwargs) -> Iterable[NestedKeyedOutputT]: """Processes a batch of anomaly results using a stateful ThresholdFn. For each input element, this DoFn retrieves the stateful `ThresholdFn` from Beam state, initializes it if it's the first time, applies it to update the prediction labels in the `AnomalyResult`, and then updates the state in Beam for future elements. Args: element (tuple[Any, tuple[Any, AnomalyResult]]): A tuple representing an element in the Beam pipeline. It is expected to be in the format `(key1, (key2, AnomalyResult))`, where key1 is the original input key, and key2 is a disambiguating key for distinct data points. threshold_state: A Beam state parameter that provides access to the persisted state of the `ThresholdFn`. It is automatically managed by Beam. **kwargs: Additional keyword arguments passed to the `process` method in Beam DoFns. Yields: Iterable[tuple[Any, tuple[Any, AnomalyResult]]]: An iterable containing a single output element with the same structure as the input, but with the `AnomalyResult` having updated prediction labels based on the stateful `ThresholdFn`. """ k1, (k2, result) = element self._threshold_fn = threshold_state.read() if self._threshold_fn is None: self._threshold_fn = Specifiable.from_spec(self._threshold_fn_spec) yield k1, (k2, self._apply_threshold_to_predictions(result)) threshold_state.write(self._threshold_fn) class RunThresholdCriterion( beam.PTransform[beam.PCollection[NestedKeyedOutputT], beam.PCollection[NestedKeyedOutputT]]): """Applies a threshold criterion to anomaly detection results. This PTransform applies a `ThresholdFn` to the anomaly scores in `AnomalyResult` objects, updating the prediction labels. It handles both stateful and stateless `ThresholdFn` implementations. Args: threshold_criterion: The `ThresholdFn` to apply. """ def __init__(self, threshold_criterion: ThresholdFn): self._threshold_fn = threshold_criterion def expand( self, input: beam.PCollection[NestedKeyedOutputT] ) -> beam.PCollection[NestedKeyedOutputT]: if self._threshold_fn.is_stateful: return ( input | beam.ParDo(_StatefulThresholdDoFn(self._threshold_fn.to_spec()))) else: return ( input | beam.ParDo(_StatelessThresholdDoFn(self._threshold_fn.to_spec()))) class RunAggregationStrategy( beam.PTransform[beam.PCollection[NestedKeyedOutputT], beam.PCollection[NestedKeyedOutputT]]): """Applies an aggregation strategy to grouped anomaly detection results. This PTransform aggregates anomaly predictions from multiple models or data points using an `AggregationFn`. It handles both custom and simple aggregation strategies. Args: aggregation_strategy: The `AggregationFn` to use. agg_model_id: The model ID for aggregation. """ def __init__( self, aggregation_strategy: Optional[AggregationFn], agg_model_id: str): self._aggregation_fn = aggregation_strategy self._agg_model_id = agg_model_id def expand( self, input: beam.PCollection[NestedKeyedOutputT] ) -> beam.PCollection[NestedKeyedOutputT]: post_gbk = ( input | beam.MapTuple(lambda k, v: ((k, v[0]), v[1])) | beam.GroupByKey()) if self._aggregation_fn is None: # simply put predictions into an iterable (list) ret = ( post_gbk | beam.MapTuple( lambda k, v: ( k[0], ( k[1], AnomalyResult( example=v[0].example, predictions=[ prediction for result in v for prediction in result.predictions ]))))) return ret # create a new aggregation_fn from spec and make sure it is initialized aggregation_fn_spec = self._aggregation_fn.to_spec() aggregation_fn_spec.config["_run_init"] = True aggregation_fn = Specifiable.from_spec(aggregation_fn_spec) # if no _agg_model_id is set in the aggregation function, use # model id from the ensemble instance if (isinstance(aggregation_fn, aggregations._AggModelIdMixin)): aggregation_fn._set_agg_model_id_if_unset(self._agg_model_id) # post_gbk is a PCollection of ((original_key, temp_key), AnomalyResult). # We use (original_key, temp_key) as the key for GroupByKey() so that # scores from multiple detectors per data point are grouped. ret = ( post_gbk | beam.MapTuple( lambda k, v, agg=aggregation_fn: ( k[0], ( k[1], AnomalyResult( example=v[0].example, predictions=[ agg.apply([ prediction for result in v for prediction in result.predictions ]) ]))))) return ret class RunOneDetector(beam.PTransform[beam.PCollection[NestedKeyedInputT], beam.PCollection[NestedKeyedOutputT]]): """Runs a single anomaly detector on a PCollection of data. This PTransform applies a single `AnomalyDetector` to the input data, including scoring, learning, and thresholding. Args: detector: The `AnomalyDetector` to run. """ def __init__(self, detector): self._detector = detector def expand( self, input: beam.PCollection[NestedKeyedInputT] ) -> beam.PCollection[NestedKeyedOutputT]: model_id = getattr( self._detector, "_model_id", getattr(self._detector, "_key", "unknown_model")) model_uuid = f"{model_id}:{uuid.uuid4().hex[:6]}" ret = ( input | beam.Reshuffle() | f"Score and Learn ({model_uuid})" >> RunScoreAndLearn(self._detector)) if self._detector._threshold_criterion: ret = ( ret | f"Run Threshold Criterion ({model_uuid})" >> RunThresholdCriterion(self._detector._threshold_criterion)) return ret class RunOfflineDetector(beam.PTransform[beam.PCollection[NestedKeyedInputT], beam.PCollection[NestedKeyedOutputT]]): """Runs a offline anomaly detector on a PCollection of data. This PTransform applies a `OfflineDetector` to the input data, handling custom input/output conversion and inference. Args: offline_detector: The `OfflineDetector` to run. """ def __init__(self, offline_detector: OfflineDetector): self._offline_detector = offline_detector def _restore_and_convert( self, elem: tuple[tuple[Any, Any, beam.Row], Any]) -> NestedKeyedOutputT: """Converts the model output to AnomalyResult. Args: elem: A tuple containing the combined key (original key, temp key, row) and the output from RunInference. Returns: A tuple containing the keyed AnomalyResult. """ (orig_key, temp_key, row), prediction = elem assert isinstance(prediction, AnomalyPrediction), ( "Wrong model handler output type." + f"Expected: 'AnomalyPrediction', but got '{type(prediction).__name__}'. " + # pylint: disable=line-too-long "Consider adding a post-processing function via `with_postprocess_fn` " + f"to convert from '{type(prediction).__name__}' to 'AnomalyPrediction', " + # pylint: disable=line-too-long "or use `score_prediction_adapter` or `label_prediction_adapter` to " + "perform the conversion.") result = AnomalyResult( example=row, predictions=[ dataclasses.replace( prediction, model_id=self._offline_detector._model_id) ]) return orig_key, (temp_key, result) def _select_features(self, elem: tuple[Any, beam.Row]) -> tuple[Any, beam.Row]: assert self._offline_detector._features is not None k, v = elem row_dict = v._asdict() return ( k, beam.Row(**{k: row_dict[k] for k in self._offline_detector._features})) def expand( self, input: beam.PCollection[NestedKeyedInputT] ) -> beam.PCollection[NestedKeyedOutputT]: model_uuid = f"{self._offline_detector._model_id}:{uuid.uuid4().hex[:6]}" # Call RunInference Transform with the keyed model handler run_inference = RunInference( self._offline_detector._keyed_model_handler, **self._offline_detector._run_inference_args) # ((orig_key, temp_key, beam.Row), beam.Row) rekeyed_model_input = input | "Rekey" >> beam.Map( lambda x: ((x[0], x[1][0], x[1][1]), x[1][1])) if self._offline_detector._features is not None: rekeyed_model_input = rekeyed_model_input | "Select Features" >> beam.Map( self._select_features) # ((orig_key, temp_key, beam.Row), AnomalyPrediction) rekeyed_model_output = ( rekeyed_model_input | f"Call RunInference ({model_uuid})" >> run_inference) ret = ( rekeyed_model_output | "Restore keys and convert model output" >> beam.Map(self._restore_and_convert)) if self._offline_detector._threshold_criterion: ret = ( ret | f"Run Threshold Criterion ({model_uuid})" >> RunThresholdCriterion(self._offline_detector._threshold_criterion)) return ret class RunEnsembleDetector(beam.PTransform[beam.PCollection[NestedKeyedInputT], beam.PCollection[NestedKeyedOutputT]] ): """Runs an ensemble of anomaly detectors on a PCollection of data. This PTransform applies an `EnsembleAnomalyDetector` to the input data, running each sub-detector and aggregating the results. Args: ensemble_detector: The `EnsembleAnomalyDetector` to run. """ def __init__(self, ensemble_detector: EnsembleAnomalyDetector): self._ensemble_detector = ensemble_detector def expand( self, input: beam.PCollection[NestedKeyedInputT] ) -> beam.PCollection[NestedKeyedOutputT]: model_uuid = f"{self._ensemble_detector._model_id}:{uuid.uuid4().hex[:6]}" assert self._ensemble_detector._sub_detectors is not None if not self._ensemble_detector._sub_detectors: raise ValueError(f"No detectors found at {model_uuid}") results = [] for idx, detector in enumerate(self._ensemble_detector._sub_detectors): if isinstance(detector, EnsembleAnomalyDetector): results.append( input | f"Run Ensemble Detector at index {idx} ({model_uuid})" >> RunEnsembleDetector(detector)) elif isinstance(detector, OfflineDetector): results.append( input | f"Run Offline Detector at index {idx} ({model_uuid})" >> RunOfflineDetector(detector)) else: results.append( input | f"Run One Detector at index {idx} ({model_uuid})" >> RunOneDetector(detector)) if self._ensemble_detector._aggregation_strategy is None: aggregation_type = "Simple" else: aggregation_type = "Custom" ret = ( results | beam.Flatten() | f"Run {aggregation_type} Aggregation Strategy ({model_uuid})" >> RunAggregationStrategy( self._ensemble_detector._aggregation_strategy, self._ensemble_detector._model_id)) if self._ensemble_detector._threshold_criterion: ret = ( ret | f"Run Threshold Criterion ({model_uuid})" >> RunThresholdCriterion(self._ensemble_detector._threshold_criterion)) return ret class AnomalyDetection(beam.PTransform[beam.PCollection[Union[InputT, KeyedInputT]], beam.PCollection[Union[OutputT, KeyedOutputT]]]): """Performs anomaly detection on a PCollection of data. This PTransform applies an `AnomalyDetector` or `EnsembleAnomalyDetector` to the input data and returns a PCollection of `AnomalyResult` objects. Examples:: # Run a single anomaly detector p | AnomalyDetection(ZScore(features=["x1"])) # Run an ensemble anomaly detector sub_detectors = [ZScore(features=["x1"]), IQR(features=["x2"])] p | AnomalyDetection(EnsembleAnomalyDetector( sub_detectors, aggregation_strategy=AnyVote())) Args: detector: The `AnomalyDetector` or `EnsembleAnomalyDetector` to use. """ def __init__( self, detector: AnomalyDetector, ) -> None: self._root_detector = detector def expand( self, input: beam.PCollection[Union[InputT, KeyedInputT]], ) -> beam.PCollection[Union[OutputT, KeyedOutputT]]: # Add a temporary unique key per data point to facilitate grouping the # outputs from multiple anomaly detectors for the same data point. # # Unique key generation options: # (1) Timestamp-based methods: https://docs.python.org/3/library/time.html # (2) UUID module: https://docs.python.org/3/library/uuid.html # # Timestamp precision on Windows can lead to key collisions (see PEP 564: # https://peps.python.org/pep-0564/#windows). Only time.perf_counter_ns() # provides sufficient precision for our needs. # # Performance note: # $ python -m timeit -n 100000 "import uuid; uuid.uuid1()" # 100000 loops, best of 5: 806 nsec per loop # $ python -m timeit -n 100000 "import uuid; uuid.uuid4()" # 100000 loops, best of 5: 1.53 usec per loop # $ python -m timeit -n 100000 "import time; time.perf_counter_ns()" # 100000 loops, best of 5: 82.3 nsec per loop # # We select uuid.uuid1() for its inclusion of node information, making it # more suitable for parallel execution environments. if isinstance(input.element_type, TupleConstraint): keyed_input = input else: # Add a default key 0 if the input is unkeyed. keyed_input = input | beam.WithKeys(0) add_temp_key_fn: Callable[[KeyedInputT], NestedKeyedInputT] run_detector: beam.PTransform if isinstance(self._root_detector, EnsembleAnomalyDetector): add_temp_key_fn = lambda e: (e[0], (str(uuid.uuid1()), e[1])) run_detector = RunEnsembleDetector(self._root_detector) else: # If there is only one non-ensemble detector, temp key can be the same # because we don't need it to identify each input during result # aggregation. add_temp_key_fn = lambda e: (e[0], ("", e[1])) if isinstance(self._root_detector, OfflineDetector): run_detector = RunOfflineDetector(self._root_detector) else: run_detector = RunOneDetector(self._root_detector) nested_keyed_input = ( keyed_input | "Add temp key" >> beam.Map(add_temp_key_fn)) nested_keyed_output = nested_keyed_input | run_detector # Remove the temporary key and simplify the output. remove_temp_key_fn: Callable[[NestedKeyedOutputT], KeyedOutputT] \ = lambda e: (e[0], e[1][1]) keyed_output = nested_keyed_output | "Remove temp key" >> beam.Map( remove_temp_key_fn) if isinstance(input.element_type, TupleConstraint): ret = keyed_output else: # Remove the default key if the input is unkeyed. ret = keyed_output | beam.Values() return ret ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/transforms_test.py0000644000175100017510000005146015206113032023542 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import itertools import logging import math import os import pickle import shutil import tempfile import unittest from collections.abc import Iterable from typing import Any import mock import numpy from parameterized import parameterized import apache_beam as beam from apache_beam.ml.anomaly.aggregations import AnyVote from apache_beam.ml.anomaly.base import AnomalyPrediction from apache_beam.ml.anomaly.base import AnomalyResult from apache_beam.ml.anomaly.base import EnsembleAnomalyDetector from apache_beam.ml.anomaly.detectors.offline import OfflineDetector from apache_beam.ml.anomaly.detectors.zscore import ZScore from apache_beam.ml.anomaly.specifiable import Spec from apache_beam.ml.anomaly.specifiable import Specifiable from apache_beam.ml.anomaly.specifiable import _spec_type_to_subspace from apache_beam.ml.anomaly.specifiable import specifiable from apache_beam.ml.anomaly.thresholds import FixedThreshold from apache_beam.ml.anomaly.thresholds import QuantileThreshold from apache_beam.ml.anomaly.transforms import AnomalyDetection from apache_beam.ml.anomaly.transforms import _StatefulThresholdDoFn from apache_beam.ml.anomaly.transforms import _StatelessThresholdDoFn from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.base import _PostProcessingModelHandler from apache_beam.ml.inference.base import _PreProcessingModelHandler from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerNumpy from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to def _prediction_iterable_is_equal_to( a: Iterable[AnomalyPrediction], b: Iterable[AnomalyPrediction]): a_list = list(a) b_list = list(b) if len(a_list) != len(b_list): return False a_dict = {} b_dict = {} for i in a_list: a_dict[i.model_id] = i for i in b_list: b_dict[i.model_id] = i for k, a_val in a_dict.items(): if k not in b_dict or not _prediction_is_equal_to(a_val, b_dict[k]): return False return True def _prediction_is_equal_to(a: AnomalyPrediction, b: AnomalyPrediction): if a.model_id != b.model_id: return False if a.threshold != b.threshold: return False if a.score != b.score: if not (a.score is not None and b.score is not None and math.isnan(a.score) and math.isnan(b.score)): return False if a.label != b.label: return False if a.info != b.info: return False if a.source_predictions is None and b.source_predictions is None: return True if a.source_predictions is not None and b.source_predictions is not None: return _prediction_iterable_is_equal_to( a.source_predictions, b.source_predictions) return False def _unkeyed_result_is_equal_to(a: AnomalyResult, b: AnomalyResult): return a.example._asdict() == b.example._asdict() and \ _prediction_iterable_is_equal_to(a.predictions, b.predictions) def _keyed_result_is_equal_to( a: tuple[int, AnomalyResult], b: tuple[int, AnomalyResult]): return a[0] == b[0] and _unkeyed_result_is_equal_to(a[1], b[1]) class TestAnomalyDetection(unittest.TestCase): class TestData: unkeyed_input = [ beam.Row(x1=1, x2=4), beam.Row(x1=2, x2=4), beam.Row(x1=3, x2=5), beam.Row(x1=10, x2=4), # outlier in key=1, with respect to x1 beam.Row(x1=2, x2=10), # outlier in key=1, with respect to x2 beam.Row(x1=3, x2=4), ] keyed_input = list(zip(itertools.repeat(1), unkeyed_input)) + [(2, beam.Row(x1=100, x2=5))] zscore_x1_expected_predictions = [ AnomalyPrediction( model_id='zscore_x1', score=float('NaN'), label=-2, threshold=3), AnomalyPrediction( model_id='zscore_x1', score=float('NaN'), label=-2, threshold=3), AnomalyPrediction( model_id='zscore_x1', score=2.1213203435596424, label=0, threshold=3), AnomalyPrediction( model_id='zscore_x1', score=8.0, label=1, threshold=3), AnomalyPrediction( model_id='zscore_x1', score=0.4898979485566356, label=0, threshold=3), AnomalyPrediction( model_id='zscore_x1', score=0.16452254913212455, label=0, threshold=3), AnomalyPrediction( model_id='zscore_x1', score=float('NaN'), label=-2, threshold=3), ] zscore_x2_expected_predictions = [ AnomalyPrediction( model_id='zscore_x2', score=float('NaN'), label=-2, threshold=2), AnomalyPrediction( model_id='zscore_x2', score=float('NaN'), label=-2, threshold=2), AnomalyPrediction(model_id='zscore_x2', score=0, label=0, threshold=2), AnomalyPrediction( model_id='zscore_x2', score=0.5773502691896252, label=0, threshold=2), AnomalyPrediction( model_id='zscore_x2', score=11.5, label=1, threshold=2), AnomalyPrediction( model_id='zscore_x2', score=0.5368754921931594, label=0, threshold=2), AnomalyPrediction( model_id='zscore_x2', score=float('NaN'), label=-2, threshold=2), ] aggregated_expected_predictions = [ AnomalyPrediction(model_id="custom", label=-2), AnomalyPrediction(model_id="custom", label=-2), AnomalyPrediction(model_id="custom", label=0), AnomalyPrediction(model_id="custom", label=1), AnomalyPrediction(model_id="custom", label=1), AnomalyPrediction(model_id="custom", label=0), AnomalyPrediction(model_id="custom", label=-2), ] keyed_zscore_x1_expected = [ (input[0], AnomalyResult(example=input[1], predictions=[decision])) for input, decision in zip(keyed_input, zscore_x1_expected_predictions) ] unkeyed_zscore_x1_expected = [ AnomalyResult(example=input, predictions=[decision]) for input, decision in zip( unkeyed_input, zscore_x1_expected_predictions) ] keyed_ensemble_expected = [( input[0], AnomalyResult(example=input[1], predictions=[decision1, decision2])) for input, decision1, decision2 in zip( keyed_input, zscore_x1_expected_predictions, zscore_x2_expected_predictions)] unkeyed_ensemble_expected = [ AnomalyResult(example=input, predictions=[decision1, decision2]) for input, decision1, decision2 in zip( unkeyed_input, zscore_x1_expected_predictions, zscore_x2_expected_predictions) ] keyed_ensemble_agg_expected = [ (input[0], AnomalyResult(example=input[1], predictions=[prediction])) for input, prediction in zip( keyed_input, aggregated_expected_predictions) ] unkeyed_ensemble_agg_expected = [ AnomalyResult(example=input, predictions=[prediction]) for input, prediction in zip( unkeyed_input, aggregated_expected_predictions) ] @parameterized.expand([ (TestData.keyed_input, TestData.keyed_zscore_x1_expected), (TestData.unkeyed_input, TestData.unkeyed_zscore_x1_expected), ]) def test_one_detector(self, input, expected): detector = ZScore(features=["x1"], model_id="zscore_x1") with TestPipeline() as p: result = (p | beam.Create(input) | AnomalyDetection(detector)) if isinstance(input[0], tuple): assert_that(result, equal_to(expected, _keyed_result_is_equal_to)) else: assert_that(result, equal_to(expected, _unkeyed_result_is_equal_to)) @parameterized.expand([ (TestData.keyed_input, TestData.keyed_ensemble_expected), (TestData.unkeyed_input, TestData.unkeyed_ensemble_expected), ]) def test_multiple_detectors_without_aggregation(self, input, expected): sub_detectors = [] sub_detectors.append(ZScore(features=["x1"], model_id="zscore_x1")) sub_detectors.append( ZScore( features=["x2"], threshold_criterion=FixedThreshold(2), model_id="zscore_x2")) with beam.Pipeline() as p: result = ( p | beam.Create(input) | AnomalyDetection(EnsembleAnomalyDetector(sub_detectors))) if isinstance(input[0], tuple): assert_that(result, equal_to(expected, _keyed_result_is_equal_to)) else: assert_that(result, equal_to(expected, _unkeyed_result_is_equal_to)) @parameterized.expand([ (TestData.keyed_input, TestData.keyed_ensemble_agg_expected), (TestData.unkeyed_input, TestData.unkeyed_ensemble_agg_expected), ]) def test_multiple_sub_detectors_with_aggregation(self, input, expected): sub_detectors = [] sub_detectors.append(ZScore(features=["x1"], model_id="zscore_x1")) sub_detectors.append( ZScore( features=["x2"], threshold_criterion=FixedThreshold(2), model_id="zscore_x2")) with beam.Pipeline() as p: result = ( p | beam.Create(input) | AnomalyDetection( EnsembleAnomalyDetector( sub_detectors, aggregation_strategy=AnyVote()))) if isinstance(input[0], tuple): assert_that(result, equal_to(expected, _keyed_result_is_equal_to)) else: assert_that(result, equal_to(expected, _unkeyed_result_is_equal_to)) class FakeNumpyModel(): def __init__(self): self.total_predict_calls = 0 def predict(self, input_vector: numpy.ndarray): self.total_predict_calls += 1 return [input_vector[0][0] * 10 - input_vector[0][1]] def _to_keyed_numpy_array(t: tuple[Any, beam.Row]): """Converts an Apache Beam Row to a NumPy array.""" return t[0], numpy.array(list(t[1])) class TestOfflineDetector(unittest.TestCase): class TestData: unkeyed_input = [ beam.Row(x=1, y=2), beam.Row(x=2, y=4), beam.Row(x=3, y=6), ] keyed_input = list(zip(itertools.repeat(1), unkeyed_input)) expected_predictions = [ AnomalyPrediction( model_id='OfflineDetector', score=8.0, label=None, threshold=None, info='', source_predictions=None), AnomalyPrediction( model_id='OfflineDetector', score=16.0, label=None, threshold=None, info='', source_predictions=None), AnomalyPrediction( model_id='OfflineDetector', score=24.0, label=None, threshold=None, info='', source_predictions=None), ] keyed_expected = [ (input[0], AnomalyResult(example=input[1], predictions=[prediction])) for input, prediction in zip(keyed_input, expected_predictions) ] unkeyed_expected = [ AnomalyResult(example=input, predictions=[prediction]) for input, prediction in zip(unkeyed_input, expected_predictions) ] def setUp(self): global SklearnModelHandlerNumpy, KeyedModelHandler global _PreProcessingModelHandler, _PostProcessingModelHandler # Make model handlers into Specifiable SklearnModelHandlerNumpy = specifiable(SklearnModelHandlerNumpy) KeyedModelHandler = specifiable(KeyedModelHandler) _PreProcessingModelHandler = specifiable(_PreProcessingModelHandler) _PostProcessingModelHandler = specifiable(_PostProcessingModelHandler) self.tmpdir = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.tmpdir) # Make the model handlers back to normal SklearnModelHandlerNumpy.unspecifiable() KeyedModelHandler.unspecifiable() _PreProcessingModelHandler.unspecifiable() _PostProcessingModelHandler.unspecifiable() @parameterized.expand([ (TestData.keyed_input, TestData.keyed_expected), (TestData.unkeyed_input, TestData.unkeyed_expected), ]) def test_default_inference_fn(self, input, expected): temp_file_name = self.tmpdir + os.sep + 'pickled_file' with open(temp_file_name, 'wb') as file: pickle.dump(FakeNumpyModel(), file) keyed_model_handler = KeyedModelHandler( SklearnModelHandlerNumpy(model_uri=temp_file_name)).with_preprocess_fn( _to_keyed_numpy_array).with_postprocess_fn( OfflineDetector.score_prediction_adapter) detector = OfflineDetector(keyed_model_handler=keyed_model_handler) detector_spec = detector.to_spec() expected_spec = Spec( type='OfflineDetector', config={ 'keyed_model_handler': Spec( type='_PostProcessingModelHandler', config={ 'base': Spec( type='_PreProcessingModelHandler', config={ 'base': Spec( type='KeyedModelHandler', config={ 'unkeyed': Spec( type='SklearnModelHandlerNumpy', config={'model_uri': temp_file_name}) }), 'preprocess_fn': Spec( type='_to_keyed_numpy_array', config=None) }), 'postprocess_fn': Spec( type='score_prediction_adapter', config=None) }) }) self.assertEqual(detector_spec, expected_spec) self.assertEqual(_spec_type_to_subspace('SklearnModelHandlerNumpy'), '*') self.assertEqual(_spec_type_to_subspace('_PreProcessingModelHandler'), '*') self.assertEqual(_spec_type_to_subspace('_PostProcessingModelHandler'), '*') self.assertEqual(_spec_type_to_subspace('_to_keyed_numpy_array'), '*') self.assertEqual(_spec_type_to_subspace('score_prediction_adapter'), '*') # Make sure the spec from the detector can be used to reconstruct the same # detector detector_new = Specifiable.from_spec(detector_spec) with TestPipeline() as p: result = (p | beam.Create(input) | AnomalyDetection(detector_new)) if isinstance(input[0], tuple): assert_that(result, equal_to(expected, _keyed_result_is_equal_to)) else: assert_that(result, equal_to(expected, _unkeyed_result_is_equal_to)) def test_run_inference_args(self): model_handler = SklearnModelHandlerNumpy(model_uri="unused") detector = OfflineDetector( keyed_model_handler=model_handler, run_inference_args={"inference_args": { "multiplier": 10 }}) p = TestPipeline() # patch the RunInference in "apache_beam.ml.anomaly.transforms" where # it is imported and call with mock.patch('apache_beam.ml.anomaly.transforms.RunInference') as mock_run_inference: # pylint: disable=line-too-long # make the actual RunInference as the sideeffect, so we record the call # information but also create the true RunInference instance. mock_run_inference.side_effect = RunInference try: p = TestPipeline() _ = ( p | beam.Create(TestOfflineDetector.TestData.unkeyed_input) | AnomalyDetection(detector)) except: # pylint: disable=bare-except pass call_args = mock_run_inference.call_args[1] self.assertEqual( call_args, { 'inference_args': { 'multiplier': 10 }, 'model_identifier': 'OfflineDetector' }) R = beam.Row(x=10, y=20) class TestStatelessThresholdDoFn(unittest.TestCase): def test_dofn_on_single_prediction(self): input = [ (1, (2, AnomalyResult(R, [AnomalyPrediction(score=1)]))), (1, (3, AnomalyResult(R, [AnomalyPrediction(score=2)]))), (1, (4, AnomalyResult(R, [AnomalyPrediction(score=3)]))), ] expected = [ ( 1, ( 2, AnomalyResult( R, [AnomalyPrediction(score=1, label=0, threshold=2)]))), ( 1, ( 3, AnomalyResult( R, [AnomalyPrediction(score=2, label=1, threshold=2)]))), ( 1, ( 4, AnomalyResult( R, [AnomalyPrediction(score=3, label=1, threshold=2)]))), ] with TestPipeline() as p: result = ( p | beam.Create(input) | beam.ParDo( _StatelessThresholdDoFn( FixedThreshold(2, normal_label=0, outlier_label=1).to_spec()))) assert_that(result, equal_to(expected)) def test_dofn_on_multiple_predictions(self): input = [ ( 1, ( 2, AnomalyResult( R, [AnomalyPrediction(score=1), AnomalyPrediction(score=4)]))), ( 1, ( 3, AnomalyResult( R, [AnomalyPrediction(score=2), AnomalyPrediction(score=0.5) ]))), ] expected = [ ( 1, ( 2, AnomalyResult( R, [ AnomalyPrediction(score=1, label=0, threshold=2), AnomalyPrediction(score=4, label=1, threshold=2) ]))), ( 1, ( 3, AnomalyResult( R, [ AnomalyPrediction(score=2, label=1, threshold=2), AnomalyPrediction(score=0.5, label=0, threshold=2) ]))), ] with TestPipeline() as p: result = ( p | beam.Create(input) | beam.ParDo( _StatelessThresholdDoFn( FixedThreshold(2, normal_label=0, outlier_label=1).to_spec()))) assert_that(result, equal_to(expected)) class TestStatefulThresholdDoFn(unittest.TestCase): def test_dofn_on_single_prediction(self): # use the input data with two keys to test stateful threshold function input = [ (1, (2, AnomalyResult(R, [AnomalyPrediction(score=1)]))), (1, (3, AnomalyResult(R, [AnomalyPrediction(score=2)]))), (1, (4, AnomalyResult(R, [AnomalyPrediction(score=3)]))), (2, (2, AnomalyResult(R, [AnomalyPrediction(score=10)]))), (2, (3, AnomalyResult(R, [AnomalyPrediction(score=20)]))), (2, (4, AnomalyResult(R, [AnomalyPrediction(score=30)]))), ] expected = [ ( 1, ( 2, AnomalyResult( R, [AnomalyPrediction(score=1, label=1, threshold=1)]))), ( 1, ( 3, AnomalyResult( R, [AnomalyPrediction(score=2, label=1, threshold=1.5)]))), ( 2, ( 2, AnomalyResult( R, [AnomalyPrediction(score=10, label=1, threshold=10)]))), ( 2, ( 3, AnomalyResult( R, [AnomalyPrediction(score=20, label=1, threshold=15)]))), ( 1, ( 4, AnomalyResult( R, [AnomalyPrediction(score=3, label=1, threshold=2)]))), ( 2, ( 4, AnomalyResult( R, [AnomalyPrediction(score=30, label=1, threshold=20)]))), ] with TestPipeline() as p: result = ( p | beam.Create(input) # use median just for test convenience | beam.ParDo( _StatefulThresholdDoFn( QuantileThreshold( quantile=0.5, normal_label=0, outlier_label=1).to_spec()))) assert_that(result, equal_to(expected)) if __name__ == '__main__': logging.getLogger().setLevel(logging.WARNING) unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.4724176 apache_beam-2.74.0/apache_beam/ml/anomaly/univariate/0000755000175100017510000000000015206113157022104 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/univariate/__init__.py0000644000175100017510000000142015206113032024202 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/univariate/base.py0000644000175100017510000000530515206113032023363 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import abc from collections import deque from enum import Enum EPSILON = 1e-9 class BaseTracker(abc.ABC): """Abstract base class for all univariate trackers.""" @abc.abstractmethod def push(self, x): """Push a new value to the tracker. Args: x: The value to be pushed. """ raise NotImplementedError() @abc.abstractmethod def get(self): """Get the current tracking value. Returns: The current tracked value, the type of which depends on the specific tracker implementation. """ raise NotImplementedError() class WindowMode(Enum): """Enum representing the window mode for windowed trackers.""" #: operating on all data points from the beginning. LANDMARK = 1 #: operating on a fixed-size sliding window of recent data points. SLIDING = 2 class WindowedTracker(BaseTracker): """Abstract base class for trackers that operate on a data window. This class provides a foundation for trackers that maintain a window of data, either as a landmark window or a sliding window. It provides basic push and pop operations. Args: window_mode: A `WindowMode` enum specifying whether the window is `LANDMARK` or `SLIDING`. **kwargs: Keyword arguments. For `SLIDING` window mode, `window_size` can be specified to set the maximum size of the sliding window. Defaults to 100. """ def __init__(self, window_mode, **kwargs): if window_mode == WindowMode.SLIDING: self._window_size = kwargs.get("window_size", 100) self._queue = deque(maxlen=self._window_size) self._n = 0 self._window_mode = window_mode def push(self, x): """Adds a new value to the data window. Args: x: The value to be added to the window. """ self._queue.append(x) def pop(self): """Removes and returns the oldest value from the data window (FIFO). Returns: The oldest value from the window. """ return self._queue.popleft() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/univariate/mad.py0000644000175100017510000000642515206113032023216 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Trackers for calculating median absolute deviation in windowed fashion.""" from typing import Optional from apache_beam.ml.anomaly.specifiable import specifiable from apache_beam.ml.anomaly.univariate.base import BaseTracker from apache_beam.ml.anomaly.univariate.median import MedianTracker @specifiable class MadTracker(BaseTracker): """Tracks the Median Absolute Deviation (MAD) of a stream of values. This class calculates the MAD, a robust measure of statistical dispersion, in an online setting. Similar functionality is available in the River library: https://github.com/online-ml/river/blob/main/river/stats/mad.py Important: This online version of MAD that does not exactly match its batch counterpart. In a streaming data context, where the true median is initially unknown, we employ an iterative estimation process. For each incoming data point, we first update the estimated median, and then calculate the absolute difference between the data point and this updated median. To maintain computational efficiency, previously calculated absolute differences are not recalculated with each subsequent median update. Args: median_tracker: An optional `MedianTracker` instance for tracking the median of the input values. If None, a default `MedianTracker` is created. diff_median_tracker: An optional `MedianTracker` instance for tracking the median of the absolute deviations from the median. If None, a default `MedianTracker` is created. """ def __init__( self, median_tracker: Optional[MedianTracker] = None, diff_median_tracker: Optional[MedianTracker] = None): self._median_tracker = median_tracker or MedianTracker() self._diff_median_tracker = diff_median_tracker or MedianTracker() def push(self, x): """Adds a new value to the tracker and updates the MAD. Args: x: The value to be added to the tracked stream. """ self._median_tracker.push(x) median = self._median_tracker.get() self._diff_median_tracker.push(abs(x - median)) def get(self): """Retrieves the current MAD value. Returns: float: The MAD of the values within the defined window. Returns `NaN` if the window is empty. """ return self._diff_median_tracker.get() def get_median(self): """Retrieves the current median value. Returns: float: The median of the values within the defined window. Returns `NaN` if the window is empty. """ return self._median_tracker.get() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/univariate/mad_test.py0000644000175100017510000000313615206113032024251 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import math import unittest from apache_beam.ml.anomaly.univariate.mad import MadTracker class MadTest(unittest.TestCase): def test_default_tracker(self): t = MadTracker() self.assertTrue(math.isnan(t.get())) t.push(4.0) # median=4, abs_dev=[0] self.assertEqual(t.get(), 0) t.push(2.0) # median=3, abs_dev=[0, 1] self.assertEqual(t.get(), 0.5) t.push(5.0) # median=4, abs_dev=[0, 1, 1] self.assertEqual(t.get(), 1.0) t.push(3.0) # median=3.5 abs_dev=[0, 0.5, 1, 1] self.assertEqual(t.get(), 0.75) t.push(0.0) # median=3 abs_dev=[0, 0.5, 1, 1, 3] self.assertEqual(t.get(), 1.0) t.push(4.0) # median=3.5 abs_dev[0, 0.5, 0.5, 1, 1, 3] self.assertEqual(t.get(), 0.75) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/univariate/mean.py0000644000175100017510000001071215206113032023367 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Trackers for calculating mean in windowed fashion. This module defines different types of mean trackers that operate on windows of data. It includes: * `SimpleSlidingMeanTracker`: Calculates mean using numpy in a sliding window. * `IncLandmarkMeanTracker`: Incremental mean tracker in landmark window mode. * `IncSlidingMeanTracker`: Incremental mean tracker in sliding window mode. """ import math import warnings import numpy as np from apache_beam.ml.anomaly.specifiable import specifiable from apache_beam.ml.anomaly.univariate.base import BaseTracker from apache_beam.ml.anomaly.univariate.base import WindowedTracker from apache_beam.ml.anomaly.univariate.base import WindowMode class MeanTracker(BaseTracker): """Abstract base class for mean trackers. Currently, it does not add any specific functionality but provides a type hierarchy for mean trackers. """ pass @specifiable class SimpleSlidingMeanTracker(WindowedTracker, MeanTracker): """Sliding window mean tracker that calculates mean using NumPy. This tracker uses NumPy's `nanmean` function to calculate the mean of the values currently in the sliding window. It's a simple, non-incremental approach. Args: window_size: The size of the sliding window. """ def __init__(self, window_size): super().__init__(window_mode=WindowMode.SLIDING, window_size=window_size) def get(self): """Calculates and returns the mean of the current sliding window. Returns: float: The mean of the values in the current sliding window. Returns NaN if the window is empty. """ if len(self._queue) == 0: return float('nan') with warnings.catch_warnings(record=False): warnings.simplefilter("ignore") return np.nanmean(self._queue) class IncMeanTracker(WindowedTracker, MeanTracker): """Base class for incremental mean trackers. This class implements incremental calculation of the mean, which is more efficient for streaming data as it updates the mean with each new data point instead of recalculating from scratch. Args: window_mode: A `WindowMode` enum specifying whether the window is `LANDMARK` or `SLIDING`. **kwargs: Keyword arguments passed to the parent class constructor. """ def __init__(self, window_mode, **kwargs): super().__init__(window_mode=window_mode, **kwargs) self._mean = 0 def push(self, x): """Pushes a new value and updates the incremental mean. Args: x: The new value to be pushed. """ if not math.isnan(x): self._n += 1 delta = x - self._mean else: delta = 0 if self._window_mode == WindowMode.SLIDING: if len(self._queue) >= self._window_size and \ not math.isnan(old_x := self.pop()): self._n -= 1 delta += (self._mean - old_x) super().push(x) if self._n > 0: self._mean += delta / self._n else: self._mean = 0 def get(self): """Returns the current incremental mean. Returns: float: The current incremental mean value. Returns NaN if no valid (non-NaN) values have been pushed. """ if self._n < 1: # keep it consistent with numpy return float("nan") return self._mean @specifiable class IncLandmarkMeanTracker(IncMeanTracker): """Landmark window mean tracker using incremental calculation.""" def __init__(self): super().__init__(window_mode=WindowMode.LANDMARK) @specifiable class IncSlidingMeanTracker(IncMeanTracker): """Sliding window mean tracker using incremental calculation. Args: window_size: The size of the sliding window. """ def __init__(self, window_size): super().__init__(window_mode=WindowMode.SLIDING, window_size=window_size) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/univariate/mean_test.py0000644000175100017510000001142015206113032024423 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import math import random import time import unittest import warnings from parameterized import parameterized from apache_beam.ml.anomaly.univariate.mean import IncLandmarkMeanTracker from apache_beam.ml.anomaly.univariate.mean import IncSlidingMeanTracker from apache_beam.ml.anomaly.univariate.mean import SimpleSlidingMeanTracker FLOAT64_MAX = 1.79769313486231570814527423731704356798070e+308 class LandmarkMeanTest(unittest.TestCase): def test_without_nan(self): t = IncLandmarkMeanTracker() self.assertTrue(math.isnan(t.get())) # Returns NaN if tracker is empty t.push(1) self.assertEqual(t.get(), 1.0) t.push(3) self.assertEqual(t.get(), 2.0) t.push(8) self.assertEqual(t.get(), 4.0) t.push(16) self.assertEqual(t.get(), 7.0) t.push(-3) self.assertEqual(t.get(), 5.0) def test_with_nan(self): t = IncLandmarkMeanTracker() t.push(float('nan')) self.assertTrue(math.isnan(t.get())) # NaN is ignored t.push(1) self.assertEqual(t.get(), 1.0) t.push(float('nan')) self.assertEqual(t.get(), 1.0) t.push(float('nan')) self.assertEqual(t.get(), 1.0) t.push(float('nan')) self.assertEqual(t.get(), 1.0) def test_with_float64_max(self): t = IncLandmarkMeanTracker() t.push(FLOAT64_MAX) self.assertEqual(t.get(), FLOAT64_MAX) t.push(FLOAT64_MAX) self.assertEqual(t.get(), FLOAT64_MAX) def test_accuracy_fuzz(self): seed = int(time.time()) random.seed(seed) print("Random seed: %d" % seed) for _ in range(10): numbers = [] for _ in range(5000): numbers.append(random.randint(0, 1000)) with warnings.catch_warnings(record=False): warnings.simplefilter("ignore") t1 = IncLandmarkMeanTracker() t2 = SimpleSlidingMeanTracker(len(numbers)) for v in numbers: t1.push(v) t2.push(v) self.assertTrue(abs(t1.get() - t2.get()) < 1e-9) class SlidingMeanTest(unittest.TestCase): @parameterized.expand([SimpleSlidingMeanTracker, IncSlidingMeanTracker]) def test_without_nan(self, tracker): t = tracker(3) self.assertTrue(math.isnan(t.get())) # Returns NaN if tracker is empty t.push(1) self.assertEqual(t.get(), 1.0) t.push(3) self.assertEqual(t.get(), 2.0) t.push(8) self.assertEqual(t.get(), 4.0) t.push(16) self.assertEqual(t.get(), 9.0) t.push(-3) self.assertEqual(t.get(), 7.0) @parameterized.expand([SimpleSlidingMeanTracker, IncSlidingMeanTracker]) def test_with_nan(self, tracker): t = tracker(3) t.push(float('nan')) self.assertTrue(math.isnan(t.get())) # NaN is ignored t.push(1) self.assertEqual(t.get(), 1.0) # flush the only number out t.push(float('nan')) self.assertEqual(t.get(), 1.0) t.push(float('nan')) self.assertEqual(t.get(), 1.0) t.push(float('nan')) self.assertTrue(math.isnan(t.get())) # All values in the tracker are NaN t.push(4) self.assertEqual(t.get(), 4.0) @parameterized.expand([SimpleSlidingMeanTracker, IncSlidingMeanTracker]) def test_with_float64_max(self, tracker): t = tracker(2) t.push(FLOAT64_MAX) self.assertEqual(t.get(), FLOAT64_MAX) t.push(FLOAT64_MAX) if tracker is IncSlidingMeanTracker: self.assertEqual(t.get(), FLOAT64_MAX) self.assertFalse(math.isinf(t.get())) else: # SimpleSlidingMean (using Numpy) returns inf when it computes the # average of [float64_max, float64_max]. self.assertTrue(math.isinf(t.get())) def test_accuracy_fuzz(self): seed = int(time.time()) random.seed(seed) print("Random seed: %d" % seed) for _ in range(10): numbers = [] for _ in range(5000): numbers.append(random.randint(0, 1000)) t1 = IncSlidingMeanTracker(100) t2 = SimpleSlidingMeanTracker(100) for v in numbers: t1.push(v) t2.push(v) self.assertTrue(abs(t1.get() - t2.get()) < 1e-9) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/univariate/median.py0000644000175100017510000000473315206113032023712 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Trackers for calculating median in windowed fashion.""" from typing import Optional from apache_beam.ml.anomaly.specifiable import specifiable from apache_beam.ml.anomaly.univariate.base import BaseTracker from apache_beam.ml.anomaly.univariate.quantile import BufferedSlidingQuantileTracker # pylint: disable=line-too-long from apache_beam.ml.anomaly.univariate.quantile import QuantileTracker DEFAULT_WINDOW_SIZE = 1000 @specifiable class MedianTracker(BaseTracker): """Tracks the median of a stream of values using a quantile tracker. This wrapper class encapsulates a `QuantileTracker` configured specifically for the 0.5 quantile (median). Args: quantile_tracker: An optional `QuantileTracker` instance. If not provided, a `BufferedSlidingQuantileTracker` with a default window size 1000 and q=0.5 is created. Raises: AssertionError: If the provided quantile_tracker is not initialized with q=0.5. """ def __init__(self, quantile_tracker: Optional[QuantileTracker] = None): self._quantile_tracker = quantile_tracker or BufferedSlidingQuantileTracker( DEFAULT_WINDOW_SIZE, 0.5) assert self._quantile_tracker._q == 0.5, \ "quantile_tracker must be initialized with q = 0.5" def push(self, x): """Pushes a new value and updates the internal quantile tracker. Args: x: The new value to be pushed. """ self._quantile_tracker.push(x) def get(self): """Calculates and returns the median (q = 0.5). Returns: float: The median of the values in the window setting specified in the internal quantile tracker. Returns NaN if the window is empty. """ return self._quantile_tracker.get() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/univariate/median_test.py0000644000175100017510000000402215206113032024740 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import math import unittest from apache_beam.ml.anomaly.univariate.median import MedianTracker from apache_beam.ml.anomaly.univariate.quantile import SimpleSlidingQuantileTracker # pylint: disable=line-too-long class MedianTest(unittest.TestCase): def test_default_tracker(self): t = MedianTracker() self.assertTrue(math.isnan(t.get())) t.push(1.0) self.assertEqual(t.get(), 1.0) t.push(2.0) self.assertEqual(t.get(), 1.5) t.push(1.0) self.assertEqual(t.get(), 1.0) t.push(20.0) self.assertEqual(t.get(), 1.5) t.push(10.0) self.assertEqual(t.get(), 2.0) def test_custom_tracker(self): t = MedianTracker(SimpleSlidingQuantileTracker(3, 0.5)) self.assertTrue(math.isnan(t.get())) t.push(1.0) self.assertEqual(t.get(), 1.0) t.push(2.0) self.assertEqual(t.get(), 1.5) t.push(1.0) self.assertEqual(t.get(), 1.0) t.push(20.0) self.assertEqual(t.get(), 2.0) t.push(10.0) self.assertEqual(t.get(), 10.0) def test_wrong_tracker(self): t = MedianTracker(SimpleSlidingQuantileTracker(50, 0.1)) self.assertRaises(AssertionError, lambda: t.run_original_init()) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/univariate/perf_test.py0000644000175100017510000000540215206113032024442 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import random import statistics import time import timeit import unittest import warnings from apache_beam.ml.anomaly.univariate.mean import * from apache_beam.ml.anomaly.univariate.quantile import * from apache_beam.ml.anomaly.univariate.stdev import * seed_value_time = int(time.time()) random.seed(seed_value_time) print(f"{'Seed value':32s}{seed_value_time}") numbers = [] for _ in range(50000): numbers.append(random.randint(0, 1000)) def run_tracker(tracker, numbers): for i in range(len(numbers)): tracker.push(numbers[i]) _ = tracker.get() def print_result(tracker, number=10, repeat=5): runtimes = timeit.repeat( lambda: run_tracker(tracker, numbers), number=number, repeat=repeat) mean = statistics.mean(runtimes) sd = statistics.stdev(runtimes) print(f"{tracker.__class__.__name__:32s}{mean:.6f} ± {sd:.6f}") class PerfTest(unittest.TestCase): def test_mean_perf(self): print() print_result(IncLandmarkMeanTracker()) print_result(IncSlidingMeanTracker(100)) # SimpleSlidingMeanTracker (numpy-based batch approach) is an order of # magnitude slower than other methods. To prevent excessively long test # runs, we reduce the number of repetitions. print_result(SimpleSlidingMeanTracker(100), number=1) def test_stdev_perf(self): print() print_result(IncLandmarkStdevTracker()) print_result(IncSlidingStdevTracker(100)) # Same as test_mean_perf, we reduce the number of repetitions here. print_result(SimpleSlidingStdevTracker(100), number=1) def test_quantile_perf(self): print() with warnings.catch_warnings(record=False): warnings.simplefilter("ignore") print_result(BufferedLandmarkQuantileTracker(0.5)) print_result(BufferedSlidingQuantileTracker(100, 0.5)) # Same as test_mean_perf, we reduce the number of repetitions here. print_result(SimpleSlidingQuantileTracker(100, 0.5), number=1) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/univariate/quantile.py0000644000175100017510000001624415206113032024277 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Trackers for calculating quantiles in windowed fashion. This module defines different types of quantile trackers that operate on windows of data. It includes: * `SimpleSlidingQuantileTracker`: Calculates quantile using numpy in a sliding window. * `BufferedLandmarkQuantileTracker`: Sortedlist based quantile tracker in landmark window mode. * `BufferedSlidingQuantileTracker`: Sortedlist based quantile tracker in sliding window mode. """ import math import typing import warnings import numpy as np from sortedcontainers import SortedList from apache_beam.ml.anomaly.specifiable import specifiable from apache_beam.ml.anomaly.univariate.base import BaseTracker from apache_beam.ml.anomaly.univariate.base import WindowedTracker from apache_beam.ml.anomaly.univariate.base import WindowMode class QuantileTracker(BaseTracker): """Abstract base class for quantile trackers. Currently, it does not add any specific functionality but provides a type hierarchy for quantile trackers. """ def __init__(self, q): assert 0 <= q <= 1, "quantile argument should be between 0 and 1" self._q = q @specifiable class SimpleSlidingQuantileTracker(WindowedTracker, QuantileTracker): """Sliding window quantile tracker using NumPy. This tracker uses NumPy's `nanquantile` function to calculate the specified quantile of the values currently in the sliding window. It's a simple, non-incremental approach. Args: window_size: The size of the sliding window. q: The quantile to calculate, a float between 0 and 1 (inclusive). """ def __init__(self, window_size, q): super().__init__(window_mode=WindowMode.SLIDING, window_size=window_size) QuantileTracker.__init__(self, q) def get(self): """Calculates and returns the specified quantile of the current sliding window. Returns: float: The specified quantile of the values in the current sliding window. Returns NaN if the window is empty. """ with warnings.catch_warnings(record=False): warnings.simplefilter("ignore") return np.nanquantile(self._queue, self._q) class BufferedQuantileTracker(WindowedTracker, QuantileTracker): """Abstract base class for buffered quantile trackers. Warning: Buffered quantile trackers are NOT truly incremental in the sense that they don't update the quantile in constant time per new data point. They maintain a sorted list of all values in the window. Args: window_mode: A `WindowMode` enum specifying whether the window is `LANDMARK` or `SLIDING`. q: The quantile to calculate, a float between 0 and 1 (inclusive). **kwargs: Keyword arguments passed to the parent class constructor. """ def __init__(self, window_mode, q, **kwargs): super().__init__(window_mode, **kwargs) QuantileTracker.__init__(self, q) self._sorted_items = SortedList() def push(self, x): """Pushes a new value, maintains the sorted list, and manages the window. Args: x: The new value to be pushed. """ if not math.isnan(x): self._sorted_items.add(x) if self._window_mode == WindowMode.SLIDING: if (len(self._queue) >= self._window_size and not math.isnan(old_x := self.pop())): self._sorted_items.discard(old_x) super().push(x) @staticmethod def _get_helper(sorted_items, q): n = len(sorted_items) if n < 1: return float("nan") pos = q * (n - 1) lo = math.floor(pos) lo_value = typing.cast(float, sorted_items[lo]) # Use linear interpolation to yield the requested quantile hi = min(lo + 1, n - 1) hi_value: float = typing.cast(float, sorted_items[hi]) return lo_value + (hi_value - lo_value) * (pos - lo) def get(self): """Returns the current quantile value using the sorted list. Calculates the quantile using linear interpolation on the sorted values. Returns: float: The calculated quantile value. Returns NaN if the window is empty. """ return self._get_helper(self._sorted_items, self._q) @specifiable class SecondaryBufferedQuantileTracker(WindowedTracker, QuantileTracker): """A secondary quantile tracker that shares its data with a master tracker. This tracker acts as a read-only view of the master tracker's data, providing quantile calculations without maintaining its own independent buffer. It relies on the master's sorted items for quantile estimations. Args: master: The BufferedQuantileTracker instance to share data with. q: A list of quantiles to track. """ def __init__(self, master: QuantileTracker, q): assert isinstance(master, BufferedQuantileTracker), \ "Cannot create secondary tracker from non-BufferedQuantileTracker" self._master = master super().__init__(self._master._window_mode) QuantileTracker.__init__(self, q) self._sorted_items = self._master._sorted_items def push(self, x): """Does nothing, as this is a secondary tracker. """ pass def get(self): """Returns the calculated quantiles based on the master tracker's buffer. Returns: A list of calculated quantiles. """ return self._master._get_helper(self._master._sorted_items, self._q) @specifiable class BufferedLandmarkQuantileTracker(BufferedQuantileTracker): """Landmark quantile tracker using a sorted list for quantile calculation. Warning: Landmark quantile trackers have unbounded memory consumption as they store all pushed values in a sorted list. Avoid using in production for long-running streams. Args: q: The quantile to calculate, a float between 0 and 1 (inclusive). """ def __init__(self, q): warnings.warn( "Quantile trackers should not be used in production due to " "the unbounded memory consumption.") super().__init__(window_mode=WindowMode.LANDMARK, q=q) @specifiable class BufferedSlidingQuantileTracker(BufferedQuantileTracker): """Sliding window quantile tracker using a sorted list for quantile calculation. Warning: Maintains a sorted list of values within the sliding window to calculate the specified quantile. Memory consumption is bounded by the window size but can still be significant for large windows. Args: window_size: The size of the sliding window. q: The quantile to calculate, a float between 0 and 1 (inclusive). """ def __init__(self, window_size, q): super().__init__( window_mode=WindowMode.SLIDING, q=q, window_size=window_size) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/univariate/quantile_test.py0000644000175100017510000001124215206113032025327 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import math import random import time import unittest import warnings from parameterized import parameterized from apache_beam.ml.anomaly.univariate.quantile import BufferedLandmarkQuantileTracker # pylint: disable=line-too-long from apache_beam.ml.anomaly.univariate.quantile import BufferedSlidingQuantileTracker from apache_beam.ml.anomaly.univariate.quantile import SimpleSlidingQuantileTracker # pylint: disable=line-too-long class LandmarkQuantileTest(unittest.TestCase): def test_without_nan(self): with warnings.catch_warnings(record=False): warnings.simplefilter("ignore") t = BufferedLandmarkQuantileTracker(0.5) self.assertTrue(math.isnan(t.get())) t.push(1) self.assertEqual(t.get(), 1.0) t.push(2) self.assertEqual(t.get(), 1.5) t.push(2) self.assertEqual(t.get(), 2.0) t.push(0) self.assertEqual(t.get(), 1.5) t.push(3) self.assertEqual(t.get(), 2.0) t.push(1) self.assertEqual(t.get(), 1.5) def test_with_nan(self): with warnings.catch_warnings(record=False): warnings.simplefilter("ignore") t = BufferedLandmarkQuantileTracker(0.2) self.assertTrue(math.isnan(t.get())) t.push(1) self.assertEqual(t.get(), 1.0) t.push(float('nan')) self.assertEqual(t.get(), 1.0) t.push(2) self.assertEqual(t.get(), 1.2) t.push(float('nan')) self.assertEqual(t.get(), 1.2) t.push(0) self.assertEqual(t.get(), 0.4) def test_accuracy_fuzz(self): seed = int(time.time()) random.seed(seed) print("Random seed: %d" % seed) def _accuracy_helper(): numbers = [] for _ in range(5000): numbers.append(random.randint(0, 1000)) with warnings.catch_warnings(record=False): warnings.simplefilter("ignore") t1 = BufferedLandmarkQuantileTracker(0.5) t2 = SimpleSlidingQuantileTracker(len(numbers), 0.5) for v in numbers: t1.push(v) t2.push(v) self.assertTrue(abs(t1.get() - t2.get()) < 1e-9) for _ in range(10): _accuracy_helper() class SlidingQuantileTest(unittest.TestCase): @parameterized.expand( [ #SimpleSlidingQuantileTracker, BufferedSlidingQuantileTracker ]) def test_without_nan(self, tracker): t = tracker(3, 0.5) self.assertTrue(math.isnan(t.get())) t.push(1) self.assertEqual(t.get(), 1.0) t.push(2) self.assertEqual(t.get(), 1.5) t.push(2) self.assertEqual(t.get(), 2.0) t.push(0) self.assertEqual(t.get(), 2.0) t.push(3) self.assertEqual(t.get(), 2.0) t.push(1) self.assertEqual(t.get(), 1.0) @parameterized.expand( [SimpleSlidingQuantileTracker, BufferedSlidingQuantileTracker]) def test_with_nan(self, tracker): t = tracker(3, 0.8) self.assertTrue(math.isnan(t.get())) t.push(1) self.assertEqual(t.get(), 1.0) t.push(2) self.assertEqual(t.get(), 1.8) t.push(float('nan')) self.assertEqual(t.get(), 1.8) t.push(2) self.assertEqual(t.get(), 2.0) t.push(0) self.assertEqual(t.get(), 1.6) t.push(float('nan')) self.assertEqual(t.get(), 1.6) t.push(float('nan')) self.assertEqual(t.get(), 0.0) t.push(float('nan')) self.assertTrue(math.isnan(t.get())) t.push(3) self.assertEqual(t.get(), 3.0) t.push(1) self.assertEqual(t.get(), 2.6) def test_accuracy_fuzz(self): seed = int(time.time()) random.seed(seed) print("Random seed: %d" % seed) def _accuracy_helper(): numbers = [] for _ in range(5000): numbers.append(random.randint(0, 1000)) t1 = BufferedSlidingQuantileTracker(100, 0.1) t2 = SimpleSlidingQuantileTracker(100, 0.1) for v in numbers: t1.push(v) t2.push(v) self.assertTrue(abs(t1.get() - t2.get()) < 1e-9) for _ in range(10): _accuracy_helper() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/univariate/stdev.py0000644000175100017510000001164415206113032023601 0ustar00runnerrunner# this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Trackers for calculating standard deviation in windowed fashion. This module defines different types of standard deviation trackers that operate on windows of data. It includes: * `SimpleSlidingStdevTracker`: Calculates stdev using numpy in a sliding window. * `IncLandmarkStdevTracker`: Incremental stdev tracker in landmark window mode. * `IncSlidingStdevTracker`: Incremental stdev tracker in sliding window mode. """ import math import warnings import numpy as np from apache_beam.ml.anomaly.specifiable import specifiable from apache_beam.ml.anomaly.univariate.base import BaseTracker from apache_beam.ml.anomaly.univariate.base import WindowedTracker from apache_beam.ml.anomaly.univariate.base import WindowMode class StdevTracker(BaseTracker): """Abstract base class for standard deviation trackers. Currently, it does not add any specific functionality but provides a type hierarchy for stdev trackers. """ pass @specifiable class SimpleSlidingStdevTracker(WindowedTracker, StdevTracker): """Sliding window standard deviation tracker using NumPy. This tracker uses NumPy's `nanvar` function to calculate the variance of the values currently in the sliding window and then takes the square root to get the standard deviation. It's a simple, non-incremental approach. """ def __init__(self, window_size): super().__init__(window_mode=WindowMode.SLIDING, window_size=window_size) def get(self): """Calculates and returns the stdev of the current sliding window. Returns: float: The standard deviation of the values in the current sliding window. Returns NaN if the window contains fewer than 2 elements. """ with warnings.catch_warnings(record=False): warnings.simplefilter("ignore") # We do not use nanstd, since nanstd([]) gives 0, which is incorrect. # Use nanvar instead. return math.sqrt(np.nanvar(self._queue, ddof=1)) class IncStdevTracker(WindowedTracker, StdevTracker): """Abstract base class for incremental standard deviation trackers. This class implements an online algorithm for calculating standard deviation, updating the standard deviation incrementally as new data points arrive. Args: window_mode: A `WindowMode` enum specifying whether the window is `LANDMARK` or `SLIDING`. **kwargs: Keyword arguments passed to the parent class constructor. """ def __init__(self, window_mode, **kwargs): super().__init__(window_mode, **kwargs) self._mean = 0 self._m2 = 0 def push(self, x): """Pushes a new value and updates the incremental standard deviation calculation. Implements Welford's online algorithm for variance, then derives stdev. Args: x: The new value to be pushed. """ if not math.isnan(x): self._n += 1 delta1 = x - self._mean else: delta1 = 0 if self._window_mode == WindowMode.SLIDING: if (len(self._queue) >= self._window_size and not math.isnan(old_x := self.pop())): self._n -= 1 delta2 = self._mean - old_x else: delta2 = 0 super().push(x) else: delta2 = 0 if self._n > 0: self._mean += (delta1 + delta2) / self._n if delta1 != 0: self._m2 += delta1 * (x - self._mean) if delta2 != 0: self._m2 += delta2 * (old_x - self._mean) else: self._mean = 0 self._m2 = 0 def get(self): """Returns the current incremental standard deviation. Returns: float: The current incremental standard deviation value. Returns NaN if fewer than 2 valid (non-NaN) values have been pushed. """ if self._n < 2: # keep it consistent with numpy return float("nan") dof = self._n - 1 return math.sqrt(self._m2 / dof) @specifiable class IncLandmarkStdevTracker(IncStdevTracker): """Landmark window standard deviation tracker using incremental calculation.""" # pylint: disable=line-too-long def __init__(self): super().__init__(window_mode=WindowMode.LANDMARK) @specifiable class IncSlidingStdevTracker(IncStdevTracker): """Sliding window standard deviation tracker using incremental calculation. Args: window_size: The size of the sliding window. """ def __init__(self, window_size): super().__init__(window_mode=WindowMode.SLIDING, window_size=window_size) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/anomaly/univariate/stdev_test.py0000644000175100017510000001076615206113032024644 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import math import random import time import unittest from parameterized import parameterized from apache_beam.ml.anomaly.univariate.stdev import IncLandmarkStdevTracker from apache_beam.ml.anomaly.univariate.stdev import IncSlidingStdevTracker from apache_beam.ml.anomaly.univariate.stdev import SimpleSlidingStdevTracker class LandmarkStdevTest(unittest.TestCase): def test_without_nan(self): t = IncLandmarkStdevTracker() self.assertTrue(math.isnan(t.get())) t.push(1) self.assertTrue(math.isnan(t.get())) t.push(2) self.assertEqual(t.get(), 0.7071067811865476) t.push(3) self.assertEqual(t.get(), 1.0) t.push(10) self.assertEqual(t.get(), 4.08248290463863) def test_with_nan(self): t = IncLandmarkStdevTracker() t.push(float('nan')) self.assertTrue(math.isnan(t.get())) # NaN is ignored t.push(float('nan')) self.assertTrue(math.isnan(t.get())) # NaN is ignored t.push(2) self.assertTrue(math.isnan(t.get())) t.push(1) self.assertEqual(t.get(), 0.7071067811865476) t.push(3) self.assertEqual(t.get(), 1.0) # flush the only number out t.push(float('nan')) self.assertEqual(t.get(), 1.0) t.push(10) self.assertEqual(t.get(), 4.08248290463863) def test_accuracy_fuzz(self): seed = int(time.time()) random.seed(seed) print("Random seed: %d" % seed) for _ in range(10): numbers = [] for _ in range(5000): numbers.append(random.randint(0, 1000)) t1 = IncLandmarkStdevTracker() t2 = SimpleSlidingStdevTracker(len(numbers)) for v in numbers: t1.push(v) t2.push(v) self.assertTrue((math.isnan(t1.get()) and math.isnan(t2.get())) or abs(t1.get() - t2.get()) < 1e-9) class SlidingStdevTest(unittest.TestCase): @parameterized.expand([SimpleSlidingStdevTracker, IncSlidingStdevTracker]) def test_without_nan(self, tracker): t = tracker(3) self.assertTrue(math.isnan(t.get())) t.push(2) self.assertTrue(math.isnan(t.get())) t.push(1) self.assertEqual(t.get(), 0.7071067811865476) t.push(3) self.assertEqual(t.get(), 1.0) t.push(10) self.assertEqual(t.get(), 4.725815626252609) @parameterized.expand([SimpleSlidingStdevTracker, IncSlidingStdevTracker]) def test_stdev_with_nan(self, tracker): t = tracker(3) t.push(float('nan')) self.assertTrue(math.isnan(t.get())) # NaN is ignored t.push(float('nan')) self.assertTrue(math.isnan(t.get())) # NaN is ignored t.push(1) self.assertTrue(math.isnan(t.get())) t.push(2) self.assertEqual(t.get(), 0.7071067811865476) t.push(3) self.assertEqual(t.get(), 1.0) # flush the only number out t.push(float('nan')) self.assertEqual(t.get(), 0.7071067811865476) t.push(float('nan')) self.assertTrue(math.isnan(t.get())) t.push(float('nan')) self.assertTrue(math.isnan(t.get())) if tracker is IncSlidingStdevTracker: self.assertEqual(t._m2, 0) self.assertEqual(t._mean, 0) t.push(4) self.assertTrue(math.isnan(t.get())) t.push(5) self.assertEqual(t.get(), 0.7071067811865476) def test_accuracy_fuzz(self): seed = int(time.time()) random.seed(seed) print("Random seed: %d" % seed) for _ in range(10): numbers = [] for _ in range(5000): numbers.append(random.randint(0, 1000)) t1 = IncSlidingStdevTracker(100) t2 = SimpleSlidingStdevTracker(100) for v in numbers: t1.push(v) t2.push(v) self.assertTrue((math.isnan(t1.get()) and math.isnan(t2.get())) or abs(t1.get() - t2.get()) < 1e-9) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.477418 apache_beam-2.74.0/apache_beam/ml/gcp/0000755000175100017510000000000015206113157017046 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/__init__.py0000644000175100017510000000206115206113032021146 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # TODO(https://github.com/apache/beam/issues/25625) # update code. For example, # use class something(beam.PTransform) instead of decorator @beam.ptransform_fn # on a function. # update type annotations to follow pep 484 https://peps.python.org/pep-0484/ # and also update doc strings. ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/cloud_dlp.py0000644000175100017510000002027015206113032021356 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """``PTransforms`` that implement Google Cloud Data Loss Prevention functionality. """ import logging from google.cloud import dlp_v2 from apache_beam import typehints from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.transforms import DoFn from apache_beam.transforms import ParDo from apache_beam.transforms import PTransform __all__ = ['MaskDetectedDetails', 'InspectForDetails'] _LOGGER = logging.getLogger(__name__) @typehints.with_input_types(str) @typehints.with_output_types(str) class MaskDetectedDetails(PTransform): """Scrubs sensitive information detected in text. The ``PTransform`` returns a ``PCollection`` of ``str`` Example usage:: pipeline | MaskDetectedDetails(project='example-gcp-project', deidentification_config={ 'info_type_transformations: { 'transformations': [{ 'primitive_transformation': { 'character_mask_config': { 'masking_character': '#' } } }] } }, inspection_config={'info_types': [{'name': 'EMAIL_ADDRESS'}]}) """ def __init__( self, project=None, deidentification_template_name=None, deidentification_config=None, inspection_template_name=None, inspection_config=None, timeout=None): """Initializes a :class:`MaskDetectedDetails` transform. Args: project: Optional. GCP project name in which inspection will be performed deidentification_template_name (str): Either this or `deidentification_config` required. Name of deidentification template to be used on detected sensitive information instances in text. deidentification_config (``Union[dict, google.cloud.dlp_v2.types.DeidentifyConfig]``): Configuration for the de-identification of the content item. If both template name and config are supplied, config is more important. inspection_template_name (str): This or `inspection_config` required. Name of inspection template to be used to detect sensitive data in text. inspection_config (``Union[dict, google.cloud.dlp_v2.types.InspectConfig]``): Configuration for the inspector used to detect sensitive data in text. If both template name and config are supplied, config takes precedence. timeout (float): Optional. The amount of time, in seconds, to wait for the request to complete. """ self.config = {} self.project = project self.timeout = timeout if deidentification_template_name is not None \ and deidentification_config is not None: raise ValueError( 'Both deidentification_template_name and ' 'deidentification_config were specified.' ' Please specify only one of these.') elif deidentification_template_name is None \ and deidentification_config is None: raise ValueError( 'deidentification_template_name or ' 'deidentification_config must be specified.') elif deidentification_template_name is not None: self.config['deidentify_template_name'] = deidentification_template_name else: self.config['deidentify_config'] = deidentification_config if inspection_config is None and inspection_template_name is None: raise ValueError( 'inspection_template_name or inspection_config must be specified') if inspection_template_name is not None: self.config['inspect_template_name'] = inspection_template_name if inspection_config is not None: self.config['inspect_config'] = inspection_config def expand(self, pcoll): if self.project is None: self.project = pcoll.pipeline.options.view_as(GoogleCloudOptions).project if self.project is None: raise ValueError( 'GCP project name needs to be specified in "project" pipeline option') return ( pcoll | ParDo(_DeidentifyFn(self.config, self.timeout, self.project))) @typehints.with_input_types(str) @typehints.with_output_types(list[dlp_v2.types.dlp.Finding]) class InspectForDetails(PTransform): """Inspects input text for sensitive information. the ``PTransform`` returns a ``PCollection`` of ``List[google.cloud.dlp_v2.proto.dlp_pb2.Finding]`` Example usage:: pipeline | InspectForDetails(project='example-gcp-project', inspection_config={'info_types': [{'name': 'EMAIL_ADDRESS'}]}) """ def __init__( self, project=None, inspection_template_name=None, inspection_config=None, timeout=None): """Initializes a :class:`InspectForDetails` transform. Args: project: Optional. GCP project name in which inspection will be performed inspection_template_name (str): This or `inspection_config` required. Name of inspection template to be used to detect sensitive data in text. inspection_config (``Union[dict, google.cloud.dlp_v2.types.InspectConfig]``): Configuration for the inspector used to detect sensitive data in text. If both template name and config are supplied, config takes precedence. timeout (float): Optional. The amount of time, in seconds, to wait for the request to complete. """ self.timeout = timeout self.config = {} self.project = project if inspection_config is None and inspection_template_name is None: raise ValueError( 'inspection_template_name or inspection_config must be specified') if inspection_template_name is not None: self.config['inspect_template_name'] = inspection_template_name if inspection_config is not None: self.config['inspect_config'] = inspection_config def expand(self, pcoll): if self.project is None: self.project = pcoll.pipeline.options.view_as(GoogleCloudOptions).project if self.project is None: raise ValueError( 'GCP project name needs to be specified in "project" pipeline option') return pcoll | ParDo(_InspectFn(self.config, self.timeout, self.project)) class _DeidentifyFn(DoFn): def __init__(self, config=None, timeout=None, project=None, client=None): self.config = config self.timeout = timeout self.client = client self.project = project self.params = {} def setup(self): if self.client is None: self.client = dlp_v2.DlpServiceClient() self.params = { 'timeout': self.timeout, } self.parent = self.client.common_project_path(self.project) def process(self, element, **kwargs): request = {'item': {'value': element}, 'parent': self.parent} request.update(self.config) operation = self.client.deidentify_content(request=request, **self.params) yield operation.item.value class _InspectFn(DoFn): def __init__(self, config=None, timeout=None, project=None): self.config = config self.timeout = timeout self.client = None self.project = project self.params = {} def setup(self): if self.client is None: self.client = dlp_v2.DlpServiceClient() self.params = { 'timeout': self.timeout, } self.parent = self.client.common_project_path(self.project) def process(self, element, **kwargs): request = {'item': {'value': element}, 'parent': self.parent} request.update(self.config) operation = self.client.inspect_content(request=request, **self.params) hits = [x for x in operation.result.findings] yield hits ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/cloud_dlp_it_test.py0000644000175100017510000000635315206113032023117 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration tests for Google Cloud Video Intelligence API transforms.""" import logging import unittest import pytest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # Protect against environments with google-cloud-dlp unavailable. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from google.cloud import dlp_v2 except ImportError: dlp_v2 = None else: from apache_beam.ml.gcp.cloud_dlp import InspectForDetails from apache_beam.ml.gcp.cloud_dlp import MaskDetectedDetails # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports _LOGGER = logging.getLogger(__name__) INSPECT_CONFIG = {"info_types": [{"name": "EMAIL_ADDRESS"}]} DEIDENTIFY_CONFIG = { "info_type_transformations": { "transformations": [{ "primitive_transformation": { "character_mask_config": { "masking_character": '#' } } }] } } def extract_inspection_results(response): yield beam.pvalue.TaggedOutput('info_type', response[0].info_type.name) @unittest.skipIf(dlp_v2 is None, 'GCP dependencies are not installed') class CloudDLPIT(unittest.TestCase): def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.runner_name = type(self.test_pipeline.runner).__name__ self.project = self.test_pipeline.get_option('project') @pytest.mark.it_postcommit def test_deidentification(self): with TestPipeline(is_integration_test=True) as p: output = ( p | beam.Create(["mary.sue@example.com"]) | MaskDetectedDetails( project=self.project, deidentification_config=DEIDENTIFY_CONFIG, inspection_config=INSPECT_CONFIG)) assert_that(output, equal_to(['####################'])) @pytest.mark.it_postcommit def test_inspection(self): with TestPipeline(is_integration_test=True) as p: output = ( p | beam.Create(["mary.sue@example.com"]) | InspectForDetails( project=self.project, inspection_config=INSPECT_CONFIG) | beam.ParDo(extract_inspection_results).with_outputs( 'quote', 'info_type')) assert_that(output.info_type, equal_to(['EMAIL_ADDRESS']), 'Type matches') if __name__ == '__main__': logging.getLogger().setLevel(logging.WARN) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/cloud_dlp_test.py0000644000175100017510000001331515206113032022417 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for Google Cloud Video Intelligence API transforms.""" import logging import unittest import mock import apache_beam as beam from apache_beam.metrics import Metrics from apache_beam.testing.test_pipeline import TestPipeline # Protect against environments with google-cloud-dlp unavailable. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from google.cloud import dlp_v2 except ImportError: dlp_v2 = None else: from google.cloud.dlp_v2.types import dlp from apache_beam.ml.gcp.cloud_dlp import InspectForDetails from apache_beam.ml.gcp.cloud_dlp import MaskDetectedDetails from apache_beam.ml.gcp.cloud_dlp import _DeidentifyFn from apache_beam.ml.gcp.cloud_dlp import _InspectFn # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports _LOGGER = logging.getLogger(__name__) @unittest.skipIf(dlp_v2 is None, 'GCP dependencies are not installed') class TestDeidentifyText(unittest.TestCase): def test_exception_raised_when_no_config_is_provided(self): with self.assertRaises(ValueError): with TestPipeline() as p: # pylint: disable=expression-not-assigned p | MaskDetectedDetails() @unittest.skipIf(dlp_v2 is None, 'GCP dependencies are not installed') class TestDeidentifyFn(unittest.TestCase): def test_deidentify_called(self): class ClientMock(object): def deidentify_content(self, *args, **kwargs): # Check that we can marshal a valid request. dlp.DeidentifyContentRequest(kwargs['request']) called = Metrics.counter('test_deidentify_text', 'called') called.inc() operation = mock.Mock() item = mock.Mock() item.value = [None] operation.item = item return operation def common_project_path(self, *args): return 'test' with mock.patch('google.cloud.dlp_v2.DlpServiceClient', ClientMock): # TODO(https://github.com/apache/beam/issues/34549): This test relies on # metrics filtering which doesn't work on Prism yet because Prism renames # steps (e.g. "Do" becomes "ref_AppliedPTransform_Do_7"). # https://github.com/apache/beam/blob/5f9cd73b7c9a2f37f83971ace3a399d633201dd1/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py#L1590 p = TestPipeline('FnApiRunner') config = { "deidentify_config": { "info_type_transformations": { "transformations": [{ "primitive_transformation": { "character_mask_config": { "masking_character": '#' } } }] } } } # pylint: disable=expression-not-assigned ( p | beam.Create(['mary.sue@example.com', 'john.doe@example.com']) | beam.ParDo(_DeidentifyFn(config=config))) result = p.run() result.wait_until_finish() called = result.metrics().query()['counters'][0] self.assertEqual(called.result, 2) @unittest.skipIf(dlp_v2 is None, 'GCP dependencies are not installed') class TestInspectText(unittest.TestCase): def test_exception_raised_then_no_config_provided(self): with self.assertRaises(ValueError): with TestPipeline() as p: #pylint: disable=expression-not-assigned p | InspectForDetails() @unittest.skipIf(dlp_v2 is None, 'GCP dependencies are not installed') class TestInspectFn(unittest.TestCase): def test_inspect_called(self): class ClientMock(object): def inspect_content(self, *args, **kwargs): # Check that we can marshal a valid request. dlp.InspectContentRequest(kwargs['request']) called = Metrics.counter('test_inspect_text', 'called') called.inc() operation = mock.Mock() operation.result = mock.Mock() operation.result.findings = [None] return operation def common_project_path(self, *args): return 'test' with mock.patch('google.cloud.dlp_v2.DlpServiceClient', ClientMock): # TODO(https://github.com/apache/beam/issues/34549): This test relies on # metrics filtering which doesn't work on Prism yet because Prism renames # steps (e.g. "Do" becomes "ref_AppliedPTransform_Do_7"). # https://github.com/apache/beam/blob/5f9cd73b7c9a2f37f83971ace3a399d633201dd1/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py#L1590 p = TestPipeline('FnApiRunner') config = {"inspect_config": {"info_types": [{"name": "EMAIL_ADDRESS"}]}} # pylint: disable=expression-not-assigned ( p | beam.Create(['mary.sue@example.com', 'john.doe@example.com']) | beam.ParDo(_InspectFn(config=config))) result = p.run() result.wait_until_finish() called = result.metrics().query()['counters'][0] self.assertEqual(called.result, 2) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/naturallanguageml.py0000644000175100017510000001163715206113032023123 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from typing import Mapping from typing import Optional from typing import Sequence from typing import Union import apache_beam as beam from apache_beam.metrics import Metrics try: from google.cloud import language from google.cloud import language_v1 except ImportError: raise ImportError( 'Google Cloud Natural Language API not supported for this execution ' 'environment (could not import Natural Language API client).') __all__ = ['Document', 'AnnotateText'] class Document(object): """Represents the input to :class:`AnnotateText` transform. Args: content (str): The content of the input or the Google Cloud Storage URI where the file is stored. type (`Union[str, google.cloud.language_v1.Document.Type]`): Text type. Possible values are `HTML`, `PLAIN_TEXT`. The default value is `PLAIN_TEXT`. language_hint (`Optional[str]`): The language of the text. If not specified, language will be automatically detected. Values should conform to ISO-639-1 standard. encoding (`Optional[str]`): Text encoding. Possible values are: `NONE`, `UTF8`, `UTF16`, `UTF32`. The default value is `UTF8`. from_gcs (bool): Whether the content should be interpret as a Google Cloud Storage URI. The default value is :data:`False`. """ def __init__( self, content: str, type: Union[str, language_v1.Document.Type] = 'PLAIN_TEXT', language_hint: Optional[str] = None, encoding: Optional[str] = 'UTF8', from_gcs: bool = False): self.content = content self.type = type self.encoding = encoding self.language_hint = language_hint self.from_gcs = from_gcs @staticmethod def to_dict(document: 'Document') -> Mapping[str, Optional[str]]: if document.from_gcs: dict_repr = {'gcs_content_uri': document.content} else: dict_repr = {'content': document.content} dict_repr.update({ 'type': document.type, 'language': document.language_hint }) return dict_repr @beam.ptransform_fn def AnnotateText( pcoll: beam.pvalue.PCollection, features: Union[Mapping[str, bool], language_v1.AnnotateTextRequest.Features], timeout: Optional[float] = None, metadata: Optional[Sequence[tuple[str, str]]] = None): """A :class:`~apache_beam.transforms.ptransform.PTransform` for annotating text using the Google Cloud Natural Language API: https://cloud.google.com/natural-language/docs. Args: pcoll (:class:`~apache_beam.pvalue.PCollection`): An input PCollection of :class:`Document` objects. features (`Union[Mapping[str, bool], types.AnnotateTextRequest.Features]`): A dictionary of natural language operations to be performed on given text in the following format:: {'extact_syntax'=True, 'extract_entities'=True} timeout (`Optional[float]`): The amount of time, in seconds, to wait for the request to complete. The timeout applies to each individual retry attempt. metadata (`Optional[Sequence[Tuple[str, str]]]`): Additional metadata that is provided to the method. """ return pcoll | beam.ParDo(_AnnotateTextFn(features, timeout, metadata)) @beam.typehints.with_input_types(Document) @beam.typehints.with_output_types(language_v1.AnnotateTextResponse) class _AnnotateTextFn(beam.DoFn): def __init__( self, features: Union[Mapping[str, bool], language_v1.AnnotateTextRequest.Features], timeout: Optional[float], metadata: Optional[Sequence[tuple[str, str]]] = None): self.features = features self.timeout = timeout self.metadata = metadata self.api_calls = Metrics.counter(self.__class__.__name__, 'api_calls') self.client = None def setup(self): self.client = self._get_api_client() @staticmethod def _get_api_client() -> language.LanguageServiceClient: return language.LanguageServiceClient() def process(self, element): response = self.client.annotate_text( document=Document.to_dict(element), features=self.features, encoding_type=element.encoding, timeout=self.timeout, metadata=self.metadata) self.api_calls.inc() yield response ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/naturallanguageml_test.py0000644000175100017510000000443015206113032024153 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file """Unit tests for Google Cloud Natural Language API transform.""" import unittest from apache_beam.metrics import MetricsFilter # Protect against environments where Google Cloud Natural Language client # is not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from google.cloud import language except ImportError: language = None else: from apache_beam.ml.gcp import naturallanguageml # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports @unittest.skipIf(language is None, 'GCP dependencies are not installed') class NaturalLanguageMlTest(unittest.TestCase): def assertCounterEqual(self, pipeline_result, counter_name, expected): metrics = pipeline_result.metrics().query( MetricsFilter().with_name(counter_name)) try: counter = metrics['counters'][0] self.assertEqual(expected, counter.result) except IndexError: raise AssertionError('Counter "{}" was not found'.format(counter_name)) def test_document_source(self): document = naturallanguageml.Document('Hello, world!') dict_ = naturallanguageml.Document.to_dict(document) self.assertTrue('content' in dict_) self.assertFalse('gcs_content_uri' in dict_) document = naturallanguageml.Document('gs://sample/location', from_gcs=True) dict_ = naturallanguageml.Document.to_dict(document) self.assertFalse('content' in dict_) self.assertTrue('gcs_content_uri' in dict_) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/naturallanguageml_test_it.py0000644000175100017510000000462115206113032024651 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import pytest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # Protect against environments where Google Cloud Natural Language client is # not available. try: from apache_beam.ml.gcp.naturallanguageml import AnnotateText from apache_beam.ml.gcp.naturallanguageml import Document from apache_beam.ml.gcp.naturallanguageml import enums from apache_beam.ml.gcp.naturallanguageml import types except ImportError: AnnotateText = None def extract(response): yield beam.pvalue.TaggedOutput('language', response.language) yield beam.pvalue.TaggedOutput( 'parts_of_speech', [ enums.PartOfSpeech.Tag(x.part_of_speech.tag).name for x in response.tokens ]) @pytest.mark.it_postcommit @unittest.skipIf(AnnotateText is None, 'GCP dependencies are not installed') class NaturalLanguageMlTestIT(unittest.TestCase): def test_analyzing_syntax(self): with TestPipeline(is_integration_test=True) as p: output = ( p | beam.Create([Document('Unified programming model.')]) | AnnotateText( types.AnnotateTextRequest.Features(extract_syntax=True)) | beam.ParDo(extract).with_outputs('language', 'parts_of_speech')) assert_that(output.language, equal_to(['en']), label='verify_language') assert_that( output.parts_of_speech, equal_to([['ADJ', 'NOUN', 'NOUN', 'PUNCT']]), label='verify_parts_of_speech') if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/recommendations_ai.py0000644000175100017510000004704015206113032023255 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A connector for sending API requests to the GCP Recommendations AI API (https://cloud.google.com/recommendations). """ from __future__ import absolute_import from typing import Sequence from cachetools.func import ttl_cache from google.api_core.retry import Retry from apache_beam import pvalue from apache_beam.metrics import Metrics from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.transforms import DoFn from apache_beam.transforms import ParDo from apache_beam.transforms import PTransform from apache_beam.transforms.util import GroupIntoBatches from apache_beam.utils import retry # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from google.cloud import recommendationengine except ImportError: raise ImportError( 'Google Cloud Recommendation AI not supported for this execution ' 'environment (could not import google.cloud.recommendationengine).') # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports __all__ = [ 'CreateCatalogItem', 'WriteUserEvent', 'ImportCatalogItems', 'ImportUserEvents', 'PredictUserEvent' ] FAILED_CATALOG_ITEMS = "failed_catalog_items" MAX_RETRIES = 5 @ttl_cache(maxsize=128, ttl=3600) def get_recommendation_prediction_client(): """Returns a Recommendation AI - Prediction Service client.""" _client = recommendationengine.PredictionServiceClient() return _client @ttl_cache(maxsize=128, ttl=3600) def get_recommendation_catalog_client(): """Returns a Recommendation AI - Catalog Service client.""" _client = recommendationengine.CatalogServiceClient() return _client @ttl_cache(maxsize=128, ttl=3600) def get_recommendation_user_event_client(): """Returns a Recommendation AI - UserEvent Service client.""" _client = recommendationengine.UserEventServiceClient() return _client class CreateCatalogItem(PTransform): """Creates catalog item records. The ``PTransform`` returns a ``PCollectionTuple`` of successfully created catalog items (``created_catalog_items``) and failures (``failed_catalog_items``). Example usage:: result = ( pipeline | CreateCatalogItem( project='example-gcp-project', catalog_name='my-catalog')) created = result.created_catalog_items """ def __init__( self, project: str = None, retry: Retry = None, timeout: float = 120, metadata: Sequence[tuple[str, str]] = (), catalog_name: str = "default_catalog"): """Initializes a :class:`CreateCatalogItem` transform. Args: project (str): Optional. GCP project name in which the catalog data will be imported. retry: Optional. Designation of what errors, if any, should be retried. timeout (float): Optional. The amount of time, in seconds, to wait for the request to complete. metadata: Optional. Strings which should be sent along with the request as metadata. catalog_name (str): Optional. Name of the catalog. Default: 'default_catalog' """ self.project = project self.retry = retry self.timeout = timeout self.metadata = metadata self.catalog_name = catalog_name def expand(self, pcoll): if self.project is None: self.project = pcoll.pipeline.options.view_as(GoogleCloudOptions).project if self.project is None: raise ValueError( """GCP project name needs to be specified in "project" pipeline option""") pardo = ParDo( _CreateCatalogItemFn( self.project, self.retry, self.timeout, self.metadata, self.catalog_name)) return pcoll | pardo.with_outputs( FAILED_CATALOG_ITEMS, main='created_catalog_items') class _CreateCatalogItemFn(DoFn): def __init__( self, project: str = None, retry: Retry = None, timeout: float = 120, metadata: Sequence[tuple[str, str]] = (), catalog_name: str = None): self._client = None self.retry = retry self.timeout = timeout self.metadata = metadata self.parent = f"projects/{project}/locations/global/catalogs/{catalog_name}" self.counter = Metrics.counter(self.__class__, "api_calls") def setup(self): if self._client is None: self._client = get_recommendation_catalog_client() def process(self, element): catalog_item = recommendationengine.CatalogItem(element) request = recommendationengine.CreateCatalogItemRequest( parent=self.parent, catalog_item=catalog_item) @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_timeout_or_quota_issues_filter ) def create_item(): return self._client.create_catalog_item( request=request, retry=self.retry, timeout=self.timeout, metadata=self.metadata) try: created_catalog_item = create_item() self.counter.inc() yield recommendationengine.CatalogItem.to_dict(created_catalog_item) except Exception: yield pvalue.TaggedOutput( FAILED_CATALOG_ITEMS, recommendationengine.CatalogItem.to_dict(catalog_item)) class ImportCatalogItems(PTransform): """Imports catalogitems in bulk. The `PTransform` returns a PCollectionTuple with PCollections of successfully and failed imported CatalogItems. Example usage:: pipeline | ImportCatalogItems( project='example-gcp-project', catalog_name='my-catalog') """ def __init__( self, max_batch_size: int = 5000, project: str = None, retry: Retry = None, timeout: float = 120, metadata: Sequence[tuple[str, str]] = (), catalog_name: str = "default_catalog"): """Initializes a :class:`ImportCatalogItems` transform Args: batch_size (int): Required. Maximum number of catalogitems per request. project (str): Optional. GCP project name in which the catalog data will be imported. retry: Optional. Designation of what errors, if any, should be retried. timeout (float): Optional. The amount of time, in seconds, to wait for the request to complete. metadata: Optional. Strings which should be sent along with the request as metadata. catalog_name (str): Optional. Name of the catalog. Default: 'default_catalog' """ self.max_batch_size = max_batch_size self.project = project self.retry = retry self.timeout = timeout self.metadata = metadata self.catalog_name = catalog_name def expand(self, pcoll): if self.project is None: self.project = pcoll.pipeline.options.view_as(GoogleCloudOptions).project if self.project is None: raise ValueError( 'GCP project name needs to be specified in "project" pipeline option') return ( pcoll | GroupIntoBatches.WithShardedKey(self.max_batch_size) | ParDo( _ImportCatalogItemsFn( self.project, self.retry, self.timeout, self.metadata, self.catalog_name))) class _ImportCatalogItemsFn(DoFn): def __init__( self, project=None, retry=None, timeout=120, metadata=None, catalog_name=None): self._client = None self.retry = retry self.timeout = timeout self.metadata = metadata self.parent = f"projects/{project}/locations/global/catalogs/{catalog_name}" self.counter = Metrics.counter(self.__class__, "api_calls") def setup(self): if self._client is None: self.client = get_recommendation_catalog_client() def process(self, element): catalog_items = [recommendationengine.CatalogItem(e) for e in element[1]] catalog_inline_source = recommendationengine.CatalogInlineSource( {"catalog_items": catalog_items}) input_config = recommendationengine.InputConfig( catalog_inline_source=catalog_inline_source) request = recommendationengine.ImportCatalogItemsRequest( parent=self.parent, input_config=input_config) try: operation = self._client.import_catalog_items( request=request, retry=self.retry, timeout=self.timeout, metadata=self.metadata) self.counter.inc(len(catalog_items)) yield operation.result() except Exception: yield pvalue.TaggedOutput(FAILED_CATALOG_ITEMS, catalog_items) class WriteUserEvent(PTransform): """Write user event information. The `PTransform` returns a PCollectionTuple with PCollections of successfully and failed written UserEvents. Example usage:: pipeline | WriteUserEvent( project='example-gcp-project', catalog_name='my-catalog', event_store='my_event_store') """ def __init__( self, project: str = None, retry: Retry = None, timeout: float = 120, metadata: Sequence[tuple[str, str]] = (), catalog_name: str = "default_catalog", event_store: str = "default_event_store"): """Initializes a :class:`WriteUserEvent` transform. Args: project (str): Optional. GCP project name in which the catalog data will be imported. retry: Optional. Designation of what errors, if any, should be retried. timeout (float): Optional. The amount of time, in seconds, to wait for the request to complete. metadata: Optional. Strings which should be sent along with the request as metadata. catalog_name (str): Optional. Name of the catalog. Default: 'default_catalog' event_store (str): Optional. Name of the event store. Default: 'default_event_store' """ self.project = project self.retry = retry self.timeout = timeout self.metadata = metadata self.catalog_name = catalog_name self.event_store = event_store def expand(self, pcoll): if self.project is None: self.project = pcoll.pipeline.options.view_as(GoogleCloudOptions).project if self.project is None: raise ValueError( 'GCP project name needs to be specified in "project" pipeline option') return pcoll | ParDo( _WriteUserEventFn( self.project, self.retry, self.timeout, self.metadata, self.catalog_name, self.event_store)) class _WriteUserEventFn(DoFn): FAILED_USER_EVENTS = "failed_user_events" def __init__( self, project=None, retry=None, timeout=120, metadata=None, catalog_name=None, event_store=None): self._client = None self.retry = retry self.timeout = timeout self.metadata = metadata self.parent = f"projects/{project}/locations/global/catalogs/"\ f"{catalog_name}/eventStores/{event_store}" self.counter = Metrics.counter(self.__class__, "api_calls") def setup(self): if self._client is None: self._client = get_recommendation_user_event_client() def process(self, element): user_event = recommendationengine.UserEvent(element) request = recommendationengine.WriteUserEventRequest( parent=self.parent, user_event=user_event) try: created_user_event = self._client.write_user_event(request) self.counter.inc() yield recommendationengine.UserEvent.to_dict(created_user_event) except Exception: yield pvalue.TaggedOutput( self.FAILED_USER_EVENTS, recommendationengine.UserEvent.to_dict(user_event)) class ImportUserEvents(PTransform): """Imports userevents in bulk. The `PTransform` returns a PCollectionTuple with PCollections of successfully and failed imported UserEvents. Example usage:: pipeline | ImportUserEvents( project='example-gcp-project', catalog_name='my-catalog', event_store='my_event_store') """ def __init__( self, max_batch_size: int = 5000, project: str = None, retry: Retry = None, timeout: float = 120, metadata: Sequence[tuple[str, str]] = (), catalog_name: str = "default_catalog", event_store: str = "default_event_store"): """Initializes a :class:`WriteUserEvent` transform. Args: batch_size (int): Required. Maximum number of catalogitems per request. project (str): Optional. GCP project name in which the catalog data will be imported. retry: Optional. Designation of what errors, if any, should be retried. timeout (float): Optional. The amount of time, in seconds, to wait for the request to complete. metadata: Optional. Strings which should be sent along with the request as metadata. catalog_name (str): Optional. Name of the catalog. Default: 'default_catalog' event_store (str): Optional. Name of the event store. Default: 'default_event_store' """ self.max_batch_size = max_batch_size self.project = project self.retry = retry self.timeout = timeout self.metadata = metadata self.catalog_name = catalog_name self.event_store = event_store def expand(self, pcoll): if self.project is None: self.project = pcoll.pipeline.options.view_as(GoogleCloudOptions).project if self.project is None: raise ValueError( 'GCP project name needs to be specified in "project" pipeline option') return ( pcoll | GroupIntoBatches.WithShardedKey(self.max_batch_size) | ParDo( _ImportUserEventsFn( self.project, self.retry, self.timeout, self.metadata, self.catalog_name, self.event_store))) class _ImportUserEventsFn(DoFn): FAILED_USER_EVENTS = "failed_user_events" def __init__( self, project=None, retry=None, timeout=120, metadata=None, catalog_name=None, event_store=None): self._client = None self.retry = retry self.timeout = timeout self.metadata = metadata self.parent = f"projects/{project}/locations/global/catalogs/"\ f"{catalog_name}/eventStores/{event_store}" self.counter = Metrics.counter(self.__class__, "api_calls") def setup(self): if self._client is None: self.client = get_recommendation_user_event_client() def process(self, element): user_events = [recommendationengine.UserEvent(e) for e in element[1]] user_event_inline_source = recommendationengine.UserEventInlineSource( {"user_events": user_events}) input_config = recommendationengine.InputConfig( user_event_inline_source=user_event_inline_source) request = recommendationengine.ImportUserEventsRequest( parent=self.parent, input_config=input_config) try: operation = self._client.write_user_event(request) self.counter.inc(len(user_events)) yield recommendationengine.PredictResponse.to_dict(operation.result()) except Exception: yield pvalue.TaggedOutput(self.FAILED_USER_EVENTS, user_events) class PredictUserEvent(PTransform): """Make a recommendation prediction. The `PTransform` returns a PCollection Example usage:: pipeline | PredictUserEvent( project='example-gcp-project', catalog_name='my-catalog', event_store='my_event_store', placement_id='recently_viewed_default') """ def __init__( self, project: str = None, retry: Retry = None, timeout: float = 120, metadata: Sequence[tuple[str, str]] = (), catalog_name: str = "default_catalog", event_store: str = "default_event_store", placement_id: str = None): """Initializes a :class:`PredictUserEvent` transform. Args: project (str): Optional. GCP project name in which the catalog data will be imported. retry: Optional. Designation of what errors, if any, should be retried. timeout (float): Optional. The amount of time, in seconds, to wait for the request to complete. metadata: Optional. Strings which should be sent along with the request as metadata. catalog_name (str): Optional. Name of the catalog. Default: 'default_catalog' event_store (str): Optional. Name of the event store. Default: 'default_event_store' placement_id (str): Required. ID of the recommendation engine placement. This id is used to identify the set of models that will be used to make the prediction. """ self.project = project self.retry = retry self.timeout = timeout self.metadata = metadata self.placement_id = placement_id self.catalog_name = catalog_name self.event_store = event_store if placement_id is None: raise ValueError('placement_id must be specified') else: self.placement_id = placement_id def expand(self, pcoll): if self.project is None: self.project = pcoll.pipeline.options.view_as(GoogleCloudOptions).project if self.project is None: raise ValueError( 'GCP project name needs to be specified in "project" pipeline option') return pcoll | ParDo( _PredictUserEventFn( self.project, self.retry, self.timeout, self.metadata, self.catalog_name, self.event_store, self.placement_id)) class _PredictUserEventFn(DoFn): FAILED_PREDICTIONS = "failed_predictions" def __init__( self, project=None, retry=None, timeout=120, metadata=None, catalog_name=None, event_store=None, placement_id=None): self._client = None self.retry = retry self.timeout = timeout self.metadata = metadata self.name = f"projects/{project}/locations/global/catalogs/"\ f"{catalog_name}/eventStores/{event_store}/placements/"\ f"{placement_id}" self.counter = Metrics.counter(self.__class__, "api_calls") def setup(self): if self._client is None: self._client = get_recommendation_prediction_client() def process(self, element): user_event = recommendationengine.UserEvent(element) request = recommendationengine.PredictRequest( name=self.name, user_event=user_event) try: prediction = self._client.predict(request) self.counter.inc() yield [ recommendationengine.PredictResponse.to_dict(p) for p in prediction.pages ] except Exception: yield pvalue.TaggedOutput(self.FAILED_PREDICTIONS, user_event) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/recommendations_ai_test.py0000644000175100017510000001644015206113032024314 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for Recommendations AI transforms.""" from __future__ import absolute_import import unittest import mock import apache_beam as beam from apache_beam.metrics import MetricsFilter # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from google.cloud import recommendationengine from apache_beam.ml.gcp import recommendations_ai except ImportError: recommendationengine = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports @unittest.skipIf( recommendationengine is None, "Recommendations AI dependencies not installed.") class RecommendationsAICatalogItemTest(unittest.TestCase): def setUp(self): self._mock_client = mock.Mock() self._mock_client.create_catalog_item.return_value = ( recommendationengine.CatalogItem()) self.m2 = mock.Mock() self.m2.result.return_value = None self._mock_client.import_catalog_items.return_value = self.m2 self._catalog_item = { "id": "12345", "title": "Sample laptop", "description": "Indisputably the most fantastic laptop ever created.", "language_code": "en", "category_hierarchies": [{ "categories": ["Electronic", "Computers"] }] } def test_CreateCatalogItem(self): expected_counter = 1 with mock.patch.object(recommendations_ai, 'get_recommendation_catalog_client', return_value=self._mock_client): p = beam.Pipeline() create_outputs = ( p | "Create data" >> beam.Create([self._catalog_item]) | "Create CatalogItem" >> recommendations_ai.CreateCatalogItem(project="test")) _ = ( create_outputs.created_catalog_items | 'CountCreated' >> beam.combiners.Count.Globally()) _ = ( create_outputs.failed_catalog_items | 'CountFailed' >> beam.combiners.Count.Globally()) result = p.run() result.wait_until_finish() read_filter = MetricsFilter().with_name('api_calls') query_result = result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] self.assertTrue(read_counter.result == expected_counter) def test_ImportCatalogItems(self): expected_counter = 1 with mock.patch.object(recommendations_ai, 'get_recommendation_catalog_client', return_value=self._mock_client): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create([ (self._catalog_item["id"], self._catalog_item), (self._catalog_item["id"], self._catalog_item) ]) | "Create CatalogItems" >> recommendations_ai.ImportCatalogItems(project="test")) result = p.run() result.wait_until_finish() read_filter = MetricsFilter().with_name('api_calls') query_result = result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] self.assertTrue(read_counter.result == expected_counter) @unittest.skipIf( recommendationengine is None, "Recommendations AI dependencies not installed.") class RecommendationsAIUserEventTest(unittest.TestCase): def setUp(self): self._mock_client = mock.Mock() self._mock_client.write_user_event.return_value = ( recommendationengine.UserEvent()) self.m2 = mock.Mock() self.m2.result.return_value = None self._mock_client.import_user_events.return_value = self.m2 self._user_event = { "event_type": "page-visit", "user_info": { "visitor_id": "1" } } def test_CreateUserEvent(self): expected_counter = 1 with mock.patch.object(recommendations_ai, 'get_recommendation_user_event_client', return_value=self._mock_client): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create([self._user_event]) | "Create UserEvent" >> recommendations_ai.WriteUserEvent(project="test")) result = p.run() result.wait_until_finish() read_filter = MetricsFilter().with_name('api_calls') query_result = result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] self.assertTrue(read_counter.result == expected_counter) def test_ImportUserEvents(self): expected_counter = 1 with mock.patch.object(recommendations_ai, 'get_recommendation_user_event_client', return_value=self._mock_client): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create([ (self._user_event["user_info"]["visitor_id"], self._user_event), (self._user_event["user_info"]["visitor_id"], self._user_event) ]) | "Create UserEvents" >> recommendations_ai.ImportUserEvents(project="test")) result = p.run() result.wait_until_finish() read_filter = MetricsFilter().with_name('api_calls') query_result = result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] self.assertTrue(read_counter.result == expected_counter) @unittest.skipIf( recommendationengine is None, "Recommendations AI dependencies not installed.") class RecommendationsAIPredictTest(unittest.TestCase): def setUp(self): self._mock_client = mock.Mock() self._mock_client.predict.return_value = [ recommendationengine.PredictResponse() ] self._user_event = { "event_type": "page-visit", "user_info": { "visitor_id": "1" } } def test_Predict(self): expected_counter = 1 with mock.patch.object(recommendations_ai, 'get_recommendation_prediction_client', return_value=self._mock_client): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create([self._user_event]) | "Prediction UserEvents" >> recommendations_ai.PredictUserEvent( project="test", placement_id="recently_viewed_default")) result = p.run() result.wait_until_finish() read_filter = MetricsFilter().with_name('api_calls') query_result = result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] self.assertTrue(read_counter.result == expected_counter) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/recommendations_ai_test_it.py0000644000175100017510000001054215206113032025005 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration tests for Recommendations AI transforms.""" from __future__ import absolute_import import random import unittest from datetime import datetime import pytest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.testing.util import is_not_empty # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from google.cloud import recommendationengine from apache_beam.ml.gcp import recommendations_ai except ImportError: recommendationengine = None # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports GCP_TEST_PROJECT = 'apache-beam-testing' CATALOG_ITEM = { "id": f"aitest-{int(datetime.now().timestamp())}-" f"{int(random.randint(1,10000))}", "title": "Sample laptop", "description": "Indisputably the most fantastic laptop ever created.", "language_code": "en", "category_hierarchies": [{ "categories": ["Electronic", "Computers"] }] } def extract_id(response): yield response["id"] def extract_event_type(response): yield response["event_type"] def extract_prediction(response): yield response[0]["results"] @pytest.mark.it_postcommit @unittest.skipIf( recommendationengine is None, "Recommendations AI dependencies not installed.") class RecommendationAIIT(unittest.TestCase): test_ran = False def test_create_catalog_item(self): with TestPipeline(is_integration_test=True) as p: RecommendationAIIT.test_ran = True create_outputs = ( p | 'Create data' >> beam.Create([CATALOG_ITEM]) | 'Create CatalogItem' >> recommendations_ai.CreateCatalogItem(project=GCP_TEST_PROJECT)) output = ( create_outputs.created_catalog_items | beam.ParDo(extract_id) | beam.combiners.ToList()) _ = create_outputs.failed_catalog_items | beam.combiners.Count.Globally() assert_that(output, equal_to([[CATALOG_ITEM["id"]]])) def test_create_user_event(self): USER_EVENT = {"event_type": "page-visit", "user_info": {"visitor_id": "1"}} with TestPipeline(is_integration_test=True) as p: RecommendationAIIT.test_ran = True output = ( p | 'Create data' >> beam.Create([USER_EVENT]) | 'Create UserEvent' >> recommendations_ai.WriteUserEvent(project=GCP_TEST_PROJECT) | beam.ParDo(extract_event_type) | beam.combiners.ToList()) assert_that(output, equal_to([[USER_EVENT["event_type"]]])) def test_predict(self): USER_EVENT = {"event_type": "page-visit", "user_info": {"visitor_id": "1"}} with TestPipeline(is_integration_test=True) as p: RecommendationAIIT.test_ran = True output = ( p | 'Create data' >> beam.Create([USER_EVENT]) | 'Predict UserEvent' >> recommendations_ai.PredictUserEvent( project=GCP_TEST_PROJECT, placement_id="recently_viewed_default") | beam.ParDo(extract_prediction)) assert_that(output, is_not_empty()) @classmethod def tearDownClass(cls): if not cls.test_ran: raise unittest.SkipTest('all test skipped') client = recommendationengine.CatalogServiceClient() parent = ( f'projects/{GCP_TEST_PROJECT}/locations/' 'global/catalogs/default_catalog') for item in list(client.list_catalog_items(parent=parent)): client.delete_catalog_item( name=f"projects/{GCP_TEST_PROJECT}/locations/global/catalogs/" f"default_catalog/catalogItems/{item.id}") if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/videointelligenceml.py0000644000175100017510000001772015206113032023441 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A connector for sending API requests to the GCP Video Intelligence API.""" from typing import Optional from typing import Union from cachetools.func import ttl_cache from apache_beam import typehints from apache_beam.metrics import Metrics from apache_beam.transforms import DoFn from apache_beam.transforms import ParDo from apache_beam.transforms import PTransform try: from google.cloud import videointelligence except ImportError: raise ImportError( 'Google Cloud Video Intelligence not supported for this execution ' 'environment (could not import google.cloud.videointelligence).') __all__ = ['AnnotateVideo', 'AnnotateVideoWithContext'] @ttl_cache(maxsize=128, ttl=3600) def get_videointelligence_client(): """Returns a Cloud Video Intelligence client.""" _client = videointelligence.VideoIntelligenceServiceClient() return _client class AnnotateVideo(PTransform): """A ``PTransform`` for annotating video using the GCP Video Intelligence API ref: https://cloud.google.com/video-intelligence/docs Sends each element to the GCP Video Intelligence API. Element is a Union[str, bytes] of either an URI (e.g. a GCS URI) or bytes base64-encoded video data. Accepts an `AsDict` side input that maps each video to a video context. """ def __init__( self, features, location_id=None, metadata=None, timeout=120, context_side_input=None): """ Args: features: (List[``videointelligence_v1.Feature``]) Required. The Video Intelligence API features to detect location_id: (str) Optional. Cloud region where annotation should take place. If no region is specified, a region will be determined based on video file location. metadata: (Sequence[Tuple[str, str]]) Optional. Additional metadata that is provided to the method. timeout: (int) Optional. The time in seconds to wait for the response from the Video Intelligence API context_side_input: (beam.pvalue.AsDict) Optional. An ``AsDict`` of a PCollection to be passed to the _VideoAnnotateFn as the video context mapping containing additional video context and/or feature-specific parameters. Example usage:: video_contexts = [('gs://cloud-samples-data/video/cat.mp4', Union[dict, ``videointelligence_v1.VideoContext``]), ('gs://some-other-video/sample.mp4', Union[dict, ``videointelligence_v1.VideoContext``]),] context_side_input = ( p | "Video contexts" >> beam.Create(video_contexts) ) videointelligenceml.AnnotateVideo(features, context_side_input=beam.pvalue.AsDict(context_side_input))) """ super().__init__() self.features = features self.location_id = location_id self.metadata = metadata self.timeout = timeout self.context_side_input = context_side_input def expand(self, pvalue): return pvalue | ParDo( _VideoAnnotateFn( features=self.features, location_id=self.location_id, metadata=self.metadata, timeout=self.timeout), context_side_input=self.context_side_input) @typehints.with_input_types( Union[str, bytes], Optional[videointelligence.VideoContext]) class _VideoAnnotateFn(DoFn): """A DoFn that sends each input element to the GCP Video Intelligence API service and outputs an element with the return result of the API (``google.cloud.videointelligence_v1.AnnotateVideoResponse``). """ def __init__(self, features, location_id, metadata, timeout): super().__init__() self._client = None self.features = features self.location_id = location_id self.metadata = metadata self.timeout = timeout self.counter = Metrics.counter(self.__class__, "API Calls") def start_bundle(self): self._client = get_videointelligence_client() def _annotate_video(self, element, video_context): if isinstance(element, str): # Is element an URI to a GCS bucket response = self._client.annotate_video( input_uri=element, features=self.features, video_context=video_context, location_id=self.location_id, metadata=self.metadata) else: # Is element raw bytes response = self._client.annotate_video( input_content=element, features=self.features, video_context=video_context, location_id=self.location_id, metadata=self.metadata) return response def process(self, element, context_side_input=None, *args, **kwargs): if context_side_input: # If we have a side input video context, use that video_context = context_side_input.get(element) else: video_context = None response = self._annotate_video(element, video_context) self.counter.inc() yield response.result(timeout=self.timeout) class AnnotateVideoWithContext(AnnotateVideo): """A ``PTransform`` for annotating video using the GCP Video Intelligence API ref: https://cloud.google.com/video-intelligence/docs Sends each element to the GCP Video Intelligence API. Element is a tuple of (Union[str, bytes], Optional[videointelligence.VideoContext]) where the former is either an URI (e.g. a GCS URI) or bytes base64-encoded video data """ def __init__(self, features, location_id=None, metadata=None, timeout=120): """ Args: features: (List[``videointelligence_v1.Feature``]) Required. the Video Intelligence API features to detect location_id: (str) Optional. Cloud region where annotation should take place. If no region is specified, a region will be determined based on video file location. metadata: (Sequence[Tuple[str, str]]) Optional. Additional metadata that is provided to the method. timeout: (int) Optional. The time in seconds to wait for the response from the Video Intelligence API """ super().__init__( features=features, location_id=location_id, metadata=metadata, timeout=timeout) def expand(self, pvalue): return pvalue | ParDo( _VideoAnnotateFnWithContext( features=self.features, location_id=self.location_id, metadata=self.metadata, timeout=self.timeout)) @typehints.with_input_types( tuple[Union[str, bytes], Optional[videointelligence.VideoContext]]) class _VideoAnnotateFnWithContext(_VideoAnnotateFn): """A DoFn that unpacks each input tuple to element, video_context variables and sends these to the GCP Video Intelligence API service and outputs an element with the return result of the API (``google.cloud.videointelligence_v1.AnnotateVideoResponse``). """ def __init__(self, features, location_id, metadata, timeout): super().__init__( features=features, location_id=location_id, metadata=metadata, timeout=timeout) def process(self, element, *args, **kwargs): element, video_context = element # Unpack (video, video_context) tuple response = self._annotate_video(element, video_context) self.counter.inc() yield response.result(timeout=self.timeout) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/videointelligenceml_test.py0000644000175100017510000001720515206113032024476 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for videointelligenceml.""" # pytype: skip-file import logging import unittest import mock import apache_beam as beam from apache_beam.metrics import MetricsFilter from apache_beam.typehints.decorators import TypeCheckError # Protect against environments where video intelligence lib is not available. # pylint: disable=ungrouped-imports try: from google.cloud import videointelligence from google.cloud.videointelligence import VideoIntelligenceServiceClient from apache_beam.ml.gcp import videointelligenceml except ImportError: VideoIntelligenceServiceClient = None @unittest.skipIf( VideoIntelligenceServiceClient is None, 'Video intelligence dependencies are not installed') class VideoIntelligenceTest(unittest.TestCase): def setUp(self): self._mock_client = mock.Mock() self.m2 = mock.Mock() self.m2.result.return_value = None self._mock_client.annotate_video.return_value = self.m2 self.features = [videointelligence.Feature.LABEL_DETECTION] self.location_id = 'us-west1' config = videointelligence.SpeechTranscriptionConfig( language_code='en-US', enable_automatic_punctuation=True) self.video_ctx = videointelligence.VideoContext( speech_transcription_config=config) def test_AnnotateVideo_with_side_input_context(self): videos_to_annotate = [ 'gs://cloud-samples-data/video/cat.mp4', 'gs://some-other-video/sample.mp4', 'gs://some-other-video/sample_2.mp4' ] video_contexts = [ ('gs://cloud-samples-data/video/cat.mp4', self.video_ctx), ('gs://some-other-video/sample.mp4', self.video_ctx), ] expected_counter = len(videos_to_annotate) with mock.patch.object(videointelligenceml, 'get_videointelligence_client', return_value=self._mock_client): p = beam.Pipeline() context_side_input = (p | "Video contexts" >> beam.Create(video_contexts)) _ = ( p | "Create data" >> beam.Create(videos_to_annotate) | "Annotate video" >> videointelligenceml.AnnotateVideo( self.features, context_side_input=beam.pvalue.AsDict(context_side_input))) result = p.run() result.wait_until_finish() read_filter = MetricsFilter().with_name('API Calls') query_result = result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] self.assertTrue(read_counter.committed == expected_counter) def test_AnnotateVideo_URIs(self): videos_to_annotate = [ 'gs://cloud-samples-data/video/cat.mp4', 'gs://cloud-samples-data/video/cat.mp4' ] expected_counter = len(videos_to_annotate) with mock.patch.object(videointelligenceml, 'get_videointelligence_client', return_value=self._mock_client): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create(videos_to_annotate) | "Annotate video" >> videointelligenceml.AnnotateVideo(self.features)) result = p.run() result.wait_until_finish() read_filter = MetricsFilter().with_name('API Calls') query_result = result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] self.assertTrue(read_counter.committed == expected_counter) def test_AnnotateVideoWithContext_b64_content(self): base_64_encoded_video = \ b'YmVnaW4gNjQ0IGNhdC12aWRlby5tcDRNICAgICgmOVQ+NyFNPCMwUi4uZmFrZV92aWRlb' videos_to_annotate = [ (base_64_encoded_video, self.video_ctx), (base_64_encoded_video, None), (base_64_encoded_video, self.video_ctx), ] expected_counter = len(videos_to_annotate) with mock.patch.object(videointelligenceml, 'get_videointelligence_client', return_value=self._mock_client): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create(videos_to_annotate) | "Annotate video" >> videointelligenceml.AnnotateVideoWithContext( self.features)) result = p.run() result.wait_until_finish() read_filter = MetricsFilter().with_name('API Calls') query_result = result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] self.assertTrue(read_counter.committed == expected_counter) def test_AnnotateVideo_b64_content(self): base_64_encoded_video = \ b'YmVnaW4gNjQ0IGNhdC12aWRlby5tcDRNICAgICgmOVQ+NyFNPCMwUi4uZmFrZV92aWRlb' videos_to_annotate = [ base_64_encoded_video, base_64_encoded_video, base_64_encoded_video, ] expected_counter = len(videos_to_annotate) with mock.patch.object(videointelligenceml, 'get_videointelligence_client', return_value=self._mock_client): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create(videos_to_annotate) | "Annotate video" >> videointelligenceml.AnnotateVideo(self.features)) result = p.run() result.wait_until_finish() read_filter = MetricsFilter().with_name('API Calls') query_result = result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] self.assertTrue(read_counter.committed == expected_counter) def test_AnnotateVideoWithContext_bad_input(self): """AnnotateVideoWithContext should not accept videos without context""" videos_to_annotate = [ 'gs://cloud-samples-data/video/cat.mp4', 'gs://cloud-samples-data/video/cat.mp4' ] with mock.patch.object(videointelligenceml, 'get_videointelligence_client', return_value=self._mock_client): with self.assertRaises(TypeCheckError): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create(videos_to_annotate) | "Annotate video" >> videointelligenceml.AnnotateVideoWithContext( self.features)) result = p.run() result.wait_until_finish() def test_AnnotateVideo_bad_input(self): videos_to_annotate = [123456789, 123456789, 123456789] with mock.patch.object(videointelligenceml, 'get_videointelligence_client', return_value=self._mock_client): with self.assertRaises(TypeCheckError): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create(videos_to_annotate) | "Annotate video" >> videointelligenceml.AnnotateVideo( self.features)) result = p.run() result.wait_until_finish() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/videointelligenceml_test_it.py0000644000175100017510000000541315206113032025170 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file """An integration test that labels entities appearing in a video and checks if some expected entities were properly recognized.""" import unittest import hamcrest as hc import pytest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import matches_all # Protect against environments where Google Cloud VideoIntelligence client is # not available. # pylint: disable=ungrouped-imports try: from google.cloud.videointelligence import enums from google.cloud.videointelligence import types from apache_beam.ml.gcp.videointelligenceml import AnnotateVideoWithContext except ImportError: AnnotateVideoWithContext = None def extract_entities_descriptions(response): for result in response.annotation_results: for segment in result.segment_presence_label_annotations: yield segment.entity.description @pytest.mark.it_postcommit @unittest.skipIf( AnnotateVideoWithContext is None, 'GCP dependencies are not installed') class VideoIntelligenceMlTestIT(unittest.TestCase): VIDEO_PATH = 'gs://apache-beam-samples/advanced_analytics/video/' \ 'gbikes_dinosaur.mp4' def test_label_detection_with_video_context(self): with TestPipeline(is_integration_test=True) as p: output = ( p | beam.Create([( self.VIDEO_PATH, types.VideoContext( label_detection_config=types.LabelDetectionConfig( label_detection_mode=enums.LabelDetectionMode.SHOT_MODE, model='builtin/latest')))]) | AnnotateVideoWithContext(features=[enums.Feature.LABEL_DETECTION]) | beam.ParDo(extract_entities_descriptions) | beam.combiners.ToList()) # Search for at least one entity that contains 'bicycle'. assert_that( output, matches_all([hc.has_item(hc.contains_string('bicycle'))])) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/visionml.py0000644000175100017510000002443015206113032021253 0ustar00runnerrunner# pylint: skip-file # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A connector for sending API requests to the GCP Vision API. """ from typing import Optional from typing import Union from cachetools.func import ttl_cache from apache_beam import typehints from apache_beam.metrics import Metrics from apache_beam.transforms import DoFn from apache_beam.transforms import FlatMap from apache_beam.transforms import ParDo from apache_beam.transforms import PTransform from apache_beam.transforms import util try: from google.cloud import vision except ImportError: raise ImportError( 'Google Cloud Vision not supported for this execution environment ' '(could not import google.cloud.vision).') __all__ = [ 'AnnotateImage', 'AnnotateImageWithContext', ] @ttl_cache(maxsize=128, ttl=3600) def get_vision_client(client_options=None): """Returns a Cloud Vision API client.""" _client = vision.ImageAnnotatorClient(client_options=client_options) return _client class AnnotateImage(PTransform): """A ``PTransform`` for annotating images using the GCP Vision API. ref: https://cloud.google.com/vision/docs/ Batches elements together using ``util.BatchElements`` PTransform and sends each batch of elements to the GCP Vision API. Element is a Union[str, bytes] of either an URI (e.g. a GCS URI) or bytes base64-encoded image data. Accepts an `AsDict` side input that maps each image to an image context. """ MAX_BATCH_SIZE = 5 MIN_BATCH_SIZE = 1 def __init__( self, features, retry=None, timeout=120, max_batch_size=None, min_batch_size=None, client_options=None, context_side_input=None, metadata=None): """ Args: features: (List[``vision.Feature``]) Required. The Vision API features to detect retry: (google.api_core.retry.Retry) Optional. A retry object used to retry requests. If None is specified (default), requests will not be retried. timeout: (float) Optional. The time in seconds to wait for the response from the Vision API. Default is 120. max_batch_size: (int) Optional. Maximum number of images to batch in the same request to the Vision API. Default is 5 (which is also the Vision API max). This parameter is primarily intended for testing. min_batch_size: (int) Optional. Minimum number of images to batch in the same request to the Vision API. Default is None. This parameter is primarily intended for testing. client_options: (Union[dict, google.api_core.client_options.ClientOptions]) Optional. Client options used to set user options on the client. API Endpoint should be set through client_options. context_side_input: (beam.pvalue.AsDict) Optional. An ``AsDict`` of a PCollection to be passed to the _ImageAnnotateFn as the image context mapping containing additional image context and/or feature-specific parameters. Example usage:: image_contexts = [(''gs://cloud-samples-data/vision/ocr/sign.jpg'', Union[dict, ``vision.ImageContext()``]), (''gs://cloud-samples-data/vision/ocr/sign.jpg'', Union[dict, ``vision.ImageContext()``]),] context_side_input = ( p | "Image contexts" >> beam.Create(image_contexts) ) visionml.AnnotateImage(features, context_side_input=beam.pvalue.AsDict(context_side_input))) metadata: (Optional[Sequence[Tuple[str, str]]]): Optional. Additional metadata that is provided to the method. """ super().__init__() self.features = features self.retry = retry self.timeout = timeout self.max_batch_size = max_batch_size or AnnotateImage.MAX_BATCH_SIZE if self.max_batch_size > AnnotateImage.MAX_BATCH_SIZE: raise ValueError( 'Max batch_size exceeded. ' 'Batch size needs to be smaller than {}'.format( AnnotateImage.MAX_BATCH_SIZE)) self.min_batch_size = min_batch_size or AnnotateImage.MIN_BATCH_SIZE self.client_options = client_options self.context_side_input = context_side_input self.metadata = metadata def expand(self, pvalue): return ( pvalue | FlatMap(self._create_image_annotation_pairs, self.context_side_input) | util.BatchElements( min_batch_size=self.min_batch_size, max_batch_size=self.max_batch_size) | ParDo( _ImageAnnotateFn( features=self.features, retry=self.retry, timeout=self.timeout, client_options=self.client_options, metadata=self.metadata))) @typehints.with_input_types(Union[str, bytes], Optional[vision.ImageContext]) @typehints.with_output_types(list[vision.AnnotateImageRequest]) def _create_image_annotation_pairs(self, element, context_side_input): if context_side_input: # If we have a side input image context, use that image_context = context_side_input.get(element) else: image_context = None if isinstance(element, str): image = vision.Image( {'source': vision.ImageSource({'image_uri': element})}) else: # Typehint checks only allows str or bytes image = vision.Image(content=element) request = vision.AnnotateImageRequest({ 'image': image, 'features': self.features, 'image_context': image_context }) yield request class AnnotateImageWithContext(AnnotateImage): """A ``PTransform`` for annotating images using the GCP Vision API. ref: https://cloud.google.com/vision/docs/ Batches elements together using ``util.BatchElements`` PTransform and sends each batch of elements to the GCP Vision API. Element is a tuple of:: (Union[str, bytes], Optional[``vision.ImageContext``]) where the former is either an URI (e.g. a GCS URI) or bytes base64-encoded image data. """ def __init__( self, features, retry=None, timeout=120, max_batch_size=None, min_batch_size=None, client_options=None, metadata=None): """ Args: features: (List[``vision.Feature``]) Required. The Vision API features to detect retry: (google.api_core.retry.Retry) Optional. A retry object used to retry requests. If None is specified (default), requests will not be retried. timeout: (float) Optional. The time in seconds to wait for the response from the Vision API. Default is 120. max_batch_size: (int) Optional. Maximum number of images to batch in the same request to the Vision API. Default is 5 (which is also the Vision API max). This parameter is primarily intended for testing. min_batch_size: (int) Optional. Minimum number of images to batch in the same request to the Vision API. Default is None. This parameter is primarily intended for testing. client_options: (Union[dict, google.api_core.client_options.ClientOptions]) Optional. Client options used to set user options on the client. API Endpoint should be set through client_options. metadata: (Optional[Sequence[Tuple[str, str]]]): Optional. Additional metadata that is provided to the method. """ super().__init__( features=features, retry=retry, timeout=timeout, max_batch_size=max_batch_size, min_batch_size=min_batch_size, client_options=client_options, metadata=metadata) def expand(self, pvalue): return ( pvalue | FlatMap(self._create_image_annotation_pairs) | util.BatchElements( min_batch_size=self.min_batch_size, max_batch_size=self.max_batch_size) | ParDo( _ImageAnnotateFn( features=self.features, retry=self.retry, timeout=self.timeout, client_options=self.client_options, metadata=self.metadata))) @typehints.with_input_types( tuple[Union[str, bytes], Optional[vision.ImageContext]]) @typehints.with_output_types(list[vision.AnnotateImageRequest]) def _create_image_annotation_pairs(self, element, **kwargs): element, image_context = element # Unpack (image, image_context) tuple if isinstance(element, str): image = vision.Image( {'source': vision.ImageSource({'image_uri': element})}) else: # Typehint checks only allows str or bytes image = vision.Image({"content": element}) request = vision.AnnotateImageRequest({ 'image': image, 'features': self.features, 'image_context': image_context }) yield request @typehints.with_input_types(list[vision.AnnotateImageRequest]) class _ImageAnnotateFn(DoFn): """A DoFn that sends each input element to the GCP Vision API. Returns ``google.cloud.vision.BatchAnnotateImagesResponse``. """ def __init__(self, features, retry, timeout, client_options, metadata): super().__init__() self._client = None self.features = features self.retry = retry self.timeout = timeout self.client_options = client_options self.metadata = metadata self.counter = Metrics.counter(self.__class__, "API Calls") def setup(self): self._client = get_vision_client(self.client_options) def process(self, element, *args, **kwargs): response = self._client.batch_annotate_images( requests=element, retry=self.retry, timeout=self.timeout, metadata=self.metadata) self.counter.inc() yield response ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/visionml_test.py0000644000175100017510000002212415206113032022310 0ustar00runnerrunner# pylint: skip-file # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for visionml.""" # pytype: skip-file import logging import unittest import mock import apache_beam as beam from apache_beam.metrics import MetricsFilter from apache_beam.typehints.decorators import TypeCheckError # Protect against environments where vision lib is not available. try: from google.cloud import vision from google.cloud.vision import ImageAnnotatorClient from apache_beam.ml.gcp import visionml except ImportError: ImageAnnotatorClient = None @unittest.skipIf( ImageAnnotatorClient is None, 'Vision dependencies are not installed') class VisionTest(unittest.TestCase): def setUp(self): self._mock_client = mock.Mock() self._mock_client.batch_annotate_images.return_value = None feature_type = vision.Feature.Type.TEXT_DETECTION self.features = [ vision.Feature({ 'type': feature_type, 'max_results': 3, 'model': "builtin/stable" }) ] self.img_ctx = vision.ImageContext() self.min_batch_size = 1 self.max_batch_size = 1 def test_AnnotateImage_URIs(self): images_to_annotate = [ 'gs://cloud-samples-data/vision/ocr/sign.jpg', 'gs://cloud-samples-data/vision/ocr/sign.jpg' ] expected_counter = len(images_to_annotate) with mock.patch.object(visionml, 'get_vision_client', return_value=self._mock_client): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create(images_to_annotate) | "Annotate image" >> visionml.AnnotateImage( self.features, min_batch_size=self.min_batch_size, max_batch_size=self.max_batch_size)) result = p.run() result.wait_until_finish() read_filter = MetricsFilter().with_name('API Calls') query_result = result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] self.assertTrue(read_counter.result == expected_counter) def test_AnnotateImage_URI_with_side_input_context(self): images_to_annotate = [ 'gs://cloud-samples-data/vision/ocr/sign.jpg', 'gs://cloud-samples-data/vision/ocr/sign.jpg' ] image_contexts = [ ('gs://cloud-samples-data/vision/ocr/sign.jpg', self.img_ctx), ('gs://cloud-samples-data/vision/ocr/sign.jpg', self.img_ctx), ] expected_counter = len(images_to_annotate) with mock.patch.object(visionml, 'get_vision_client', return_value=self._mock_client): p = beam.Pipeline() context_side_input = (p | "Image contexts" >> beam.Create(image_contexts)) _ = ( p | "Create data" >> beam.Create(images_to_annotate) | "Annotate image" >> visionml.AnnotateImage( self.features, min_batch_size=self.min_batch_size, max_batch_size=self.max_batch_size, context_side_input=beam.pvalue.AsDict(context_side_input))) result = p.run() result.wait_until_finish() read_filter = MetricsFilter().with_name('API Calls') query_result = result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] self.assertTrue(read_counter.result == expected_counter) def test_AnnotateImage_b64_content(self): base_64_encoded_image = \ b'YmVnaW4gNjQ0IGNhdC12aWRlby5tcDRNICAgICgmOVQ+NyFNPCMwUi4uZmFrZV92aWRlb' images_to_annotate = [ base_64_encoded_image, base_64_encoded_image, base_64_encoded_image, ] expected_counter = len(images_to_annotate) with mock.patch.object(visionml, 'get_vision_client', return_value=self._mock_client): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create(images_to_annotate) | "Annotate image" >> visionml.AnnotateImage( self.features, min_batch_size=self.min_batch_size, max_batch_size=self.max_batch_size)) result = p.run() result.wait_until_finish() read_filter = MetricsFilter().with_name('API Calls') query_result = result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] self.assertTrue(read_counter.result == expected_counter) def test_AnnotateImageWithContext_URIs(self): images_to_annotate = [ ('gs://cloud-samples-data/vision/ocr/sign.jpg', self.img_ctx), ('gs://cloud-samples-data/vision/ocr/sign.jpg', None), ('gs://cloud-samples-data/vision/ocr/sign.jpg', self.img_ctx), ] batch_size = 5 expected_counter = 1 # All images should fit in the same batch with mock.patch.object(visionml, 'get_vision_client', return_value=self._mock_client): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create(images_to_annotate) | "Annotate image" >> visionml.AnnotateImageWithContext( self.features, min_batch_size=batch_size, max_batch_size=batch_size)) result = p.run() result.wait_until_finish() read_filter = MetricsFilter().with_name('API Calls') query_result = result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] self.assertTrue(read_counter.result == expected_counter) def test_AnnotateImageWithContext_bad_input(self): """AnnotateImageWithContext should not accept images without context""" images_to_annotate = [ 'gs://cloud-samples-data/vision/ocr/sign.jpg', 'gs://cloud-samples-data/vision/ocr/sign.jpg' ] with mock.patch.object(visionml, 'get_vision_client', return_value=self._mock_client): with self.assertRaises(TypeCheckError): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create(images_to_annotate) | "Annotate image" >> visionml.AnnotateImageWithContext( self.features)) result = p.run() result.wait_until_finish() def test_AnnotateImage_bad_input(self): images_to_annotate = [123456789, 123456789, 123456789] with mock.patch.object(visionml, 'get_vision_client', return_value=self._mock_client): with self.assertRaises(TypeCheckError): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create(images_to_annotate) | "Annotate image" >> visionml.AnnotateImage(self.features)) result = p.run() result.wait_until_finish() def test_AnnotateImage_URIs_large_batch(self): images_to_annotate = [ 'gs://cloud-samples-data/vision/ocr/sign.jpg', 'gs://cloud-samples-data/vision/ocr/sign.jpg', 'gs://cloud-samples-data/vision/ocr/sign.jpg', 'gs://cloud-samples-data/vision/ocr/sign.jpg', 'gs://cloud-samples-data/vision/ocr/sign.jpg', 'gs://cloud-samples-data/vision/ocr/sign.jpg', 'gs://cloud-samples-data/vision/ocr/sign.jpg', 'gs://cloud-samples-data/vision/ocr/sign.jpg', 'gs://cloud-samples-data/vision/ocr/sign.jpg', 'gs://cloud-samples-data/vision/ocr/sign.jpg', 'gs://cloud-samples-data/vision/ocr/sign.jpg', ] batch_size = 5 expected_counter = 3 # All 11 images should fit in 3 batches with mock.patch.object(visionml, 'get_vision_client', return_value=self._mock_client): p = beam.Pipeline() _ = ( p | "Create data" >> beam.Create(images_to_annotate) | "Annotate image" >> visionml.AnnotateImage( self.features, max_batch_size=batch_size, min_batch_size=batch_size)) result = p.run() result.wait_until_finish() read_filter = MetricsFilter().with_name('API Calls') query_result = result.metrics().query(read_filter) if query_result['counters']: read_counter = query_result['counters'][0] self.assertTrue(read_counter.result == expected_counter) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/gcp/visionml_test_it.py0000644000175100017510000000504715206113032023011 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import pytest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # Protect against environments where Google Cloud Vision client is not # available. # pylint: disable=ungrouped-imports try: from google.cloud import vision from apache_beam.ml.gcp.visionml import AnnotateImage except ImportError: vision = None def extract(response): for r in response.responses: for text_annotation in r.text_annotations: yield text_annotation.description @pytest.mark.it_postcommit @unittest.skipIf(vision is None, 'GCP dependencies are not installed') class VisionMlTestIT(unittest.TestCase): def test_text_detection_with_language_hint(self): IMAGES_TO_ANNOTATE = [ 'gs://apache-beam-samples/advanced_analytics/vision/sign.jpg' ] IMAGE_CONTEXT = [vision.ImageContext({'language_hints': ['en']})] with TestPipeline(is_integration_test=True) as p: contexts = p | 'Create context' >> beam.Create( dict(zip(IMAGES_TO_ANNOTATE, IMAGE_CONTEXT))) output = ( p | beam.Create(IMAGES_TO_ANNOTATE) | AnnotateImage( features=[ vision.Feature({'type_': vision.Feature.Type.TEXT_DETECTION}) ], context_side_input=beam.pvalue.AsDict(contexts)) | beam.ParDo(extract)) assert_that( output, equal_to([ 'WAITING?\nPLEASE\nTURN OFF\nYOUR\nENGINE', 'WAITING', '?', 'PLEASE', 'TURN', 'OFF', 'YOUR', 'ENGINE' ])) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003200000000000010210 xustar0026 mtime=1779996271.49742 apache_beam-2.74.0/apache_beam/ml/inference/0000755000175100017510000000000015206113157020233 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/__init__.py0000644000175100017510000000242015206113032022332 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A package with various modules for running inferences and predictions on models. This package contains support for popular frameworks as well as an interface for adding unsupported frameworks. Note: on top of the frameworks captured in submodules below, Beam also has a supported TensorFlow model handler via the tfx-bsl library. See https://beam.apache.org/documentation/ml/about-ml/#tensorflow for more information on using TensorFlow in Beam. """ from apache_beam.ml.inference.base import RunInference ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/agent_development_kit.py0000644000175100017510000002621415206113032025151 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ModelHandler for running agents built with the Google Agent Development Kit. This module provides :class:`ADKAgentModelHandler`, a Beam :class:`~apache_beam.ml.inference.base.ModelHandler` that wraps an ADK :class:`google.adk.agents.llm_agent.LlmAgent` so it can be used with the :class:`~apache_beam.ml.inference.base.RunInference` transform. Typical usage:: import apache_beam as beam from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.agent_development_kit import ADKAgentModelHandler from google.adk.agents import LlmAgent agent = LlmAgent( name="my_agent", model="gemini-2.0-flash", instruction="You are a helpful assistant.", ) with beam.Pipeline() as p: results = ( p | beam.Create(["What is the capital of France?"]) | RunInference(ADKAgentModelHandler(agent=agent)) ) If your agent contains state that is not picklable (e.g. tool closures that capture unpicklable objects), pass a zero-arg factory callable instead:: handler = ADKAgentModelHandler(agent=lambda: LlmAgent(...)) """ import asyncio import logging import uuid from collections.abc import Callable from collections.abc import Iterable from collections.abc import Sequence from typing import Any from typing import Optional from typing import Union from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import PredictionResult try: from google.adk import sessions from google.adk.agents import Agent from google.adk.runners import Runner from google.adk.sessions import BaseSessionService from google.adk.sessions import InMemorySessionService from google.genai.types import Content as genai_Content from google.genai.types import Part as genai_Part ADK_AVAILABLE = True except ImportError: ADK_AVAILABLE = False genai_Content = Any # type: ignore[assignment, misc] genai_Part = Any # type: ignore[assignment, misc] LOGGER = logging.getLogger("ADKAgentModelHandler") # Type alias for an agent or factory that produces one _AgentOrFactory = Union["Agent", Callable[[], "Agent"]] class ADKAgentModelHandler(ModelHandler[str | genai_Content, PredictionResult, "Runner"]): """ModelHandler for running ADK agents with the Beam RunInference transform. Accepts either a fully constructed :class:`google.adk.agents.Agent` or a zero-arg factory callable that produces one. The factory form is useful when the agent contains state that is not picklable and therefore cannot be serialized alongside the pipeline graph. Each call to :meth:`run_inference` invokes the agent once per element in the batch. By default every invocation uses a fresh, isolated session (stateless). Stateful multi-turn conversations can be achieved by passing a ``session_id`` key inside ``inference_args``; elements sharing the same ``session_id`` will continue the same conversation history. When using stateful conversations, it is recommended to use a custom session_service_factory to provide a session service implementation which can be managed across multiple workers (e.g. :class:`~google.adk.sessions.DatabaseSessionService`). The default :class:`~google.adk.sessions.InMemorySessionService` will not correctly track the same session across multiple workers. Args: agent: A pre-constructed :class:`~google.adk.agents.Agent` instance, or a zero-arg callable that returns one. The callable form defers agent construction to worker ``load_model`` time, which is useful when the agent cannot be serialized. app_name: The ADK application name used to namespace sessions. Defaults to ``"beam_inference"``. session_service_factory: Optional zero-arg callable returning a :class:`~google.adk.sessions.BaseSessionService`. When ``None``, an :class:`~google.adk.sessions.InMemorySessionService` is created automatically. min_batch_size: Optional minimum batch size. max_batch_size: Optional maximum batch size. max_batch_duration_secs: Optional maximum time to buffer a batch before emitting; used in streaming contexts. max_batch_weight: Optional maximum total weight of a batch. element_size_fn: Optional function that returns the size (weight) of an element. """ def __init__( self, agent: _AgentOrFactory, app_name: str = "beam_inference", session_service_factory: Optional[Callable[[], "BaseSessionService"]] = None, *, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, **kwargs): if not ADK_AVAILABLE: raise ImportError( "google-adk is required to use ADKAgentModelHandler. " "Install it with: pip install google-adk") if agent is None: raise ValueError("'agent' must be an Agent instance or a callable.") self._agent_or_factory = agent self._app_name = app_name self._session_service_factory = session_service_factory super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, **kwargs) def load_model(self) -> "Runner": """Instantiates the ADK Runner on the worker. Resolves the agent (calling the factory if a callable was provided), then creates a :class:`~google.adk.runners.Runner` backed by the configured session service. Returns: A fully initialised :class:`~google.adk.runners.Runner`. """ if callable(self._agent_or_factory) and not isinstance( self._agent_or_factory, Agent): agent = self._agent_or_factory() else: agent = self._agent_or_factory if self._session_service_factory is not None: session_service = self._session_service_factory() else: session_service = InMemorySessionService() runner = Runner( agent=agent, app_name=self._app_name, session_service=session_service, ) LOGGER.info( "Loaded ADK Runner for agent '%s' (app_name='%s')", agent.name, self._app_name, ) return runner def run_inference( self, batch: Sequence[str | genai_Content], model: "Runner", inference_args: Optional[dict[str, Any]] = None, ) -> Iterable[PredictionResult]: """Runs the ADK agent on each element in the batch. Each element is sent to the agent as a new user turn. The final response text from the agent is returned as the ``inference`` field of a :class:`~apache_beam.ml.inference.base.PredictionResult`. Args: batch: A sequence of inputs, each of which is either a ``str`` (the user message text) or a :class:`google.genai.types.Content` object (for richer multi-part messages). model: The :class:`~google.adk.runners.Runner` returned by :meth:`load_model`. inference_args: Optional dict of extra arguments. Supported keys: - ``"session_id"`` (:class:`str`): If supplied, all elements in this batch share this session ID, enabling stateful multi-turn conversations. If omitted, each element receives a unique auto- generated session ID. - ``"user_id"`` (:class:`str`): The user identifier to pass to the runner. Defaults to ``"beam_user"``. Returns: An iterable of :class:`~apache_beam.ml.inference.base.PredictionResult`, one per input element. """ if inference_args is None: inference_args = {} user_id: str = inference_args.get("user_id", "beam_user") agent_invocations = [] elements_with_sessions = [] for element in batch: session_id: str = inference_args.get("session_id", str(uuid.uuid4())) # Wrap plain strings in a Content object if isinstance(element, str): # pyrefly: ignore[bad-instantiation] message = genai_Content(role="user", parts=[genai_Part(text=element)]) else: # Assume the caller has already constructed a types.Content object message = element agent_invocations.append( self._invoke_agent( model, user_id, session_id, self._app_name, message)) elements_with_sessions.append(element) # Run all agent invocations concurrently async def _run_concurrently(): return await asyncio.gather(*agent_invocations) response_texts = asyncio.run(_run_concurrently()) results = [] for i, element in enumerate(elements_with_sessions): results.append( PredictionResult( example=element, inference=response_texts[i], model_id=model.agent.name, )) return results @staticmethod async def _invoke_agent( runner: "Runner", user_id: str, session_id: str, app_name: str, message: genai_Content, ) -> Optional[str]: """Drives the ADK event loop and returns the final response text. Args: runner: The ADK Runner to invoke. user_id: The user ID for this invocation. session_id: The session ID for this invocation. message: The :class:`google.genai.types.Content` to send. Returns: The text of the agent's final response, or ``None`` if the agent produced no final text response. """ # Check for your specific session ID try: # Attempt to get the specific session await runner.session_service.get_session(session_id) except Exception as e: await runner.session_service.create_session( app_name=app_name, user_id=user_id, session_id=session_id, ) async for event in runner.run_async( user_id=user_id, session_id=session_id, new_message=message, ): if event.is_final_response(): if event.content and event.content.parts: return "".join([p.text for p in event.content.parts]) raise ValueError( f"Agent {runner.agent.name} did not return a response, " f"final event: {event}") raise ValueError(f"Agent {runner.agent.name} did not return a response") def get_metrics_namespace(self) -> str: return "ADKAgentModelHandler" ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/agent_development_kit_test.py0000644000175100017510000002762115206113032026213 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import asyncio import unittest from unittest import mock try: from google.adk.agents import Agent from apache_beam.ml.inference.agent_development_kit import ADKAgentModelHandler from apache_beam.ml.inference.base import PredictionResult except ImportError: raise unittest.SkipTest('google-adk dependencies are not installed') def _make_mock_agent(name: str = "test_agent") -> mock.MagicMock: """Returns a mock that quacks like a google.adk.agents.Agent.""" agent = mock.MagicMock() agent.name = name return agent def _make_mock_runner( agent: mock.MagicMock, final_text: str = "Hello from agent", ) -> mock.MagicMock: """Returns a mock Runner whose run_async yields one final-response event.""" part = mock.MagicMock() part.text = final_text content = mock.MagicMock() content.parts = [part] event = mock.MagicMock() event.is_final_response.return_value = True event.content = content async def _async_gen(*args, **kwargs): yield event runner = mock.MagicMock() runner.agent = agent runner.run_async = mock.MagicMock(side_effect=_async_gen) runner.session_service = mock.MagicMock() runner.session_service.get_session = mock.AsyncMock( side_effect=Exception("Session not found")) runner.session_service.create_session = mock.AsyncMock() return runner # --------------------------------------------------------------------------- # Helper: patch ADK imports inside the module under test so tests work even # when google-adk is installed (avoids constructing real ADK objects). # --------------------------------------------------------------------------- _MODULE = "apache_beam.ml.inference.agent_development_kit" class TestADKAgentModelHandlerInit(unittest.TestCase): """Tests for __init__ argument validation.""" def test_raises_if_agent_is_none(self): with self.assertRaises((ValueError, TypeError)): ADKAgentModelHandler(agent=None) def test_accepts_agent_object(self): agent = _make_mock_agent() handler = ADKAgentModelHandler(agent=agent) self.assertEqual(handler._agent_or_factory, agent) def test_accepts_agent_factory_callable(self): agent = _make_mock_agent() factory = lambda: agent handler = ADKAgentModelHandler(agent=factory) self.assertTrue(callable(handler._agent_or_factory)) def test_default_app_name(self): agent = _make_mock_agent() handler = ADKAgentModelHandler(agent=agent) self.assertEqual(handler._app_name, "beam_inference") def test_custom_app_name(self): agent = _make_mock_agent() handler = ADKAgentModelHandler(agent=agent, app_name="my_app") self.assertEqual(handler._app_name, "my_app") def test_metrics_namespace(self): agent = _make_mock_agent() handler = ADKAgentModelHandler(agent=agent) self.assertEqual(handler.get_metrics_namespace(), "ADKAgentModelHandler") class TestLoadModel(unittest.TestCase): """Tests for load_model / Runner construction.""" def test_load_model_with_agent_object(self): def get_current_time(city: str) -> dict: """Returns the current time in a specified city.""" return {"status": "success", "city": city, "time": "10:30 AM"} agent = Agent( model='gemini-3-flash-preview', name='root_agent', description="Tells the current time in a specified city.", instruction= "You are a helpful assistant that tells the current time in cities. " "Use the 'get_current_time' tool for this purpose.", tools=[get_current_time], ) handler = ADKAgentModelHandler(agent=agent, app_name="test_app") runner = handler.load_model() self.assertEqual(agent, runner.agent) @mock.patch(f"{_MODULE}.Runner") @mock.patch(f"{_MODULE}.InMemorySessionService") def test_load_model_calls_factory(self, mock_session_cls, mock_runner_cls): agent = _make_mock_agent() factory = mock.MagicMock(return_value=agent) handler = ADKAgentModelHandler(agent=factory) handler.load_model() factory.assert_called_once() mock_runner_cls.assert_called_once_with( agent=agent, app_name="beam_inference", session_service=mock_session_cls.return_value, ) class TestRunInference(unittest.TestCase): """Tests for run_inference output and batching.""" def test_string_input_returns_prediction_result(self): agent = _make_mock_agent() runner = _make_mock_runner(agent, final_text="Paris") handler = ADKAgentModelHandler(agent=agent) results = list( handler.run_inference( batch=["What is the capital of France?"], model=runner)) self.assertEqual(len(results), 1) pr = results[0] self.assertIsInstance(pr, PredictionResult) self.assertEqual(pr.example, "What is the capital of France?") self.assertEqual(pr.inference, "Paris") self.assertEqual(pr.model_id, "test_agent") def test_batch_of_strings(self): agent = _make_mock_agent() runner = _make_mock_runner(agent, final_text="answer") handler = ADKAgentModelHandler(agent=agent) results = list( handler.run_inference(batch=["q1", "q2", "q3"], model=runner)) self.assertEqual(len(results), 3) self.assertEqual([r.example for r in results], ["q1", "q2", "q3"]) def test_content_object_input(self): """Non-string inputs (types.Content) are passed through unchanged.""" agent = _make_mock_agent() runner = _make_mock_runner(agent, final_text="Berlin") content_input = mock.MagicMock() # simulates types.Content handler = ADKAgentModelHandler(agent=agent) results = list(handler.run_inference(batch=[content_input], model=runner)) self.assertEqual(len(results), 1) self.assertEqual(results[0].example, content_input) self.assertEqual(results[0].inference, "Berlin") def test_none_inference_args_uses_defaults(self): agent = _make_mock_agent() runner = _make_mock_runner(agent) handler = ADKAgentModelHandler(agent=agent) results = list( handler.run_inference( batch=["hello"], model=runner, inference_args=None)) self.assertEqual(len(results), 1) def test_custom_user_id_passed_to_runner(self): agent = _make_mock_agent() runner = _make_mock_runner(agent) handler = ADKAgentModelHandler(agent=agent) handler.run_inference( batch=["hi"], model=runner, inference_args={"user_id": "custom_user"}, ) call_kwargs = runner.run_async.call_args[1] self.assertEqual(call_kwargs["user_id"], "custom_user") class TestSessionManagement(unittest.TestCase): """Tests for session creation and session_id handling.""" def test_each_element_gets_unique_session_by_default(self): agent = _make_mock_agent() runner = _make_mock_runner(agent) handler = ADKAgentModelHandler(agent=agent) handler.run_inference(batch=["a", "b", "c"], model=runner) # create_session should have been called 3 times with distinct session IDs calls = runner.session_service.create_session.call_args_list self.assertEqual(len(calls), 3) session_ids = [c[1]["session_id"] for c in calls] self.assertEqual(len(set(session_ids)), 3, "Expected unique session IDs") def test_shared_session_id_from_inference_args(self): agent = _make_mock_agent() runner = _make_mock_runner(agent) handler = ADKAgentModelHandler(agent=agent) handler.run_inference( batch=["turn1", "turn2"], model=runner, inference_args={"session_id": "my-session"}, ) calls = runner.session_service.create_session.call_args_list session_ids = [c[1]["session_id"] for c in calls] self.assertTrue( all(sid == "my-session" for sid in session_ids), "All elements should share the provided session_id", ) def test_session_created_with_correct_app_name(self): agent = _make_mock_agent() runner = _make_mock_runner(agent) handler = ADKAgentModelHandler(agent=agent, app_name="my_app") handler.run_inference(batch=["hello"], model=runner) call_kwargs = runner.session_service.create_session.call_args[1] self.assertEqual(call_kwargs["app_name"], "my_app") class TestResponseExtraction(unittest.TestCase): """Tests for extraction of the final response from the event stream.""" def test_raises_when_no_final_response(self): """Agent emits only non-final events; should raise ValueError.""" agent = _make_mock_agent() # Build a runner that yields only non-final events non_final_event = mock.MagicMock() non_final_event.is_final_response.return_value = False async def _async_gen(*args, **kwargs): yield non_final_event runner = mock.MagicMock() runner.agent = agent runner.run_async = mock.MagicMock(side_effect=_async_gen) runner.session_service = mock.MagicMock() runner.session_service.get_session = mock.AsyncMock(side_effect=Exception()) runner.session_service.create_session = mock.AsyncMock() handler = ADKAgentModelHandler(agent=agent) with self.assertRaisesRegex(ValueError, "did not return a response"): list(handler.run_inference(batch=["hello"], model=runner)) def test_raises_when_final_event_has_no_content(self): agent = _make_mock_agent() event = mock.MagicMock() event.is_final_response.return_value = True event.content = None async def _async_gen(*args, **kwargs): yield event runner = mock.MagicMock() runner.agent = agent runner.run_async = mock.MagicMock(side_effect=_async_gen) runner.session_service = mock.MagicMock() runner.session_service.get_session = mock.AsyncMock(side_effect=Exception()) runner.session_service.create_session = mock.AsyncMock() handler = ADKAgentModelHandler(agent=agent) with self.assertRaisesRegex(ValueError, "did not return a response"): list(handler.run_inference(batch=["hello"], model=runner)) def test_stops_after_first_final_response(self): """Multiple final events: only the first one's text should be used.""" agent = _make_mock_agent() def _make_event(text: str): part = mock.MagicMock() part.text = text content = mock.MagicMock() content.parts = [part] event = mock.MagicMock() event.is_final_response.return_value = True event.content = content return event async def _async_gen(*args, **kwargs): yield _make_event("first") yield _make_event("second") runner = mock.MagicMock() runner.agent = agent runner.run_async = mock.MagicMock(side_effect=_async_gen) runner.session_service = mock.MagicMock() runner.session_service.get_session = mock.AsyncMock(side_effect=Exception()) runner.session_service.create_session = mock.AsyncMock() handler = ADKAgentModelHandler(agent=agent) results = list(handler.run_inference(batch=["hi"], model=runner)) self.assertEqual(results[0].inference, "first") def test_invoke_agent_static_method_directly(self): """Unit test the async _invoke_agent helper directly.""" agent = _make_mock_agent() runner = _make_mock_runner(agent, final_text="direct result") result = asyncio.run( ADKAgentModelHandler._invoke_agent( runner, "user", "session-1", "test_app", mock.MagicMock())) self.assertEqual(result, "direct result") if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/anthropic_inference.py0000644000175100017510000002507015206113032024606 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A ModelHandler for Anthropic Claude models using the Messages API. This module provides an integration between Apache Beam's RunInference transform and Anthropic's Claude models, enabling batch inference in Beam pipelines. Example usage:: from apache_beam.ml.inference.anthropic_inference import ( AnthropicModelHandler, message_from_string, ) from apache_beam.ml.inference.base import RunInference # Basic usage model_handler = AnthropicModelHandler( model_name='claude-haiku-4-5', api_key='your-api-key', request_fn=message_from_string, ) # With system prompt and structured output model_handler = AnthropicModelHandler( model_name='claude-haiku-4-5', api_key='your-api-key', request_fn=message_from_string, system='You are a helpful assistant that responds concisely.', output_config={ 'format': { 'type': 'json_schema', 'schema': { 'type': 'object', 'properties': { 'answer': {'type': 'string'}, 'confidence': {'type': 'number'}, }, 'required': ['answer', 'confidence'], 'additionalProperties': False, }, }, }, ) with beam.Pipeline() as p: results = ( p | beam.Create(['What is Apache Beam?', 'Explain MapReduce.']) | RunInference(model_handler) ) """ import logging from collections.abc import Callable from collections.abc import Iterable from collections.abc import Sequence from typing import Any from typing import Optional from typing import Union from anthropic import Anthropic from anthropic import APIStatusError from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RemoteModelHandler LOGGER = logging.getLogger("AnthropicModelHandler") def _retry_on_appropriate_error(exception: Exception) -> bool: """Retry filter that returns True for 5xx errors or 429 (rate limiting). Args: exception: the exception encountered during the request/response loop. Returns: True if the exception is retriable (429 or 5xx), False otherwise. """ if not isinstance(exception, APIStatusError): return False return exception.status_code == 429 or exception.status_code >= 500 def message_from_string( model_name: str, batch: Sequence[str], client: Anthropic, inference_args: dict[str, Any]): """Request function that sends string prompts to Claude as user messages. Each string in the batch is sent as a separate request. The results are returned as a list of response objects. Args: model_name: the Claude model to use (e.g. 'claude-haiku-4-5'). batch: the string inputs to send to Claude. client: the Anthropic client instance. inference_args: additional arguments passed to the messages.create call. Common args include 'max_tokens', 'system', 'temperature', 'top_p'. """ max_tokens = inference_args.pop('max_tokens', 1024) responses = [] for prompt in batch: response = client.messages.create( model=model_name, max_tokens=max_tokens, messages=[{ "role": "user", "content": prompt }], **inference_args) responses.append(response) return responses def message_from_conversation( model_name: str, batch: Sequence[list[dict[str, str]]], client: Anthropic, inference_args: dict[str, Any]): """Request function that sends multi-turn conversations to Claude. Each element in the batch is a list of message dicts with 'role' and 'content' keys, representing a multi-turn conversation. Args: model_name: the Claude model to use. batch: a sequence of conversations (each a list of message dicts). client: the Anthropic client instance. inference_args: additional arguments passed to the messages.create call. """ max_tokens = inference_args.pop('max_tokens', 1024) responses = [] for messages in batch: response = client.messages.create( model=model_name, max_tokens=max_tokens, messages=messages, **inference_args) responses.append(response) return responses class AnthropicModelHandler(RemoteModelHandler[Any, PredictionResult, Anthropic]): def __init__( self, model_name: str, request_fn: Callable[[str, Sequence[Any], Anthropic, dict[str, Any]], Any], api_key: Optional[str] = None, *, system: Optional[Union[str, list[dict[str, str]]]] = None, output_config: Optional[dict[str, Any]] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, **kwargs): """Implementation of the ModelHandler interface for Anthropic Claude. **NOTE:** This API and its implementation are under development and do not provide backward compatibility guarantees. This handler connects to the Anthropic Messages API to run inference using Claude models. It supports text generation from string prompts or multi-turn conversations, with optional system prompts and structured output schemas. Args: model_name: the Claude model to send requests to (e.g. 'claude-sonnet-4-6', 'claude-haiku-4-5'). request_fn: the function to use to send requests. Should take the model name, batch, client, and inference_args and return the responses from Claude. Built-in options are message_from_string and message_from_conversation. api_key: the Anthropic API key. If not provided, the client will look for the ANTHROPIC_API_KEY environment variable. system: optional system prompt to set the model's behavior for all requests. Can be a string or a list of content blocks (dicts with 'type' and 'text' keys). This is applied to every request in the pipeline. Per-request overrides can be passed via inference_args. output_config: optional output configuration to constrain responses to a structured schema. The value is passed directly to the Anthropic API as the 'output_config' parameter. This uses the GA API shape with a nested 'format' key. Example:: output_config={ 'format': { 'type': 'json_schema', 'schema': { 'type': 'object', 'properties': { 'answer': {'type': 'string'}, }, 'required': ['answer'], 'additionalProperties': False, }, }, } min_batch_size: optional. the minimum batch size to use when batching inputs. max_batch_size: optional. the maximum batch size to use when batching inputs. max_batch_duration_secs: optional. the maximum amount of time to buffer a batch before emitting; used in streaming contexts. max_batch_weight: optional. the maximum total weight of a batch. element_size_fn: optional. a function that returns the size (weight) of an element. """ self._batching_kwargs = {} self._env_vars = kwargs.get('env_vars', {}) if min_batch_size is not None: self._batching_kwargs["min_batch_size"] = min_batch_size if max_batch_size is not None: self._batching_kwargs["max_batch_size"] = max_batch_size if max_batch_duration_secs is not None: self._batching_kwargs["max_batch_duration_secs"] = max_batch_duration_secs if max_batch_weight is not None: self._batching_kwargs["max_batch_weight"] = max_batch_weight if element_size_fn is not None: self._batching_kwargs['element_size_fn'] = element_size_fn self.model_name = model_name self.request_fn = request_fn self.api_key = api_key self.system = system self.output_config = output_config super().__init__( namespace='AnthropicModelHandler', retry_filter=_retry_on_appropriate_error, **kwargs) def batch_elements_kwargs(self): return self._batching_kwargs def create_client(self) -> Anthropic: """Creates the Anthropic client used to send requests. If api_key was provided at construction time, it is used directly. Otherwise, the client will fall back to the ANTHROPIC_API_KEY environment variable. """ if self.api_key: return Anthropic(api_key=self.api_key) return Anthropic() def request( self, batch: Sequence[Any], model: Anthropic, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """Sends a prediction request to the Anthropic API. Handler-level system and output_config are injected into inference_args before calling the request function. Per-request values in inference_args take precedence over handler-level values. Args: batch: a sequence of inputs to be passed to the request function. model: an Anthropic client instance. inference_args: additional arguments to send as part of the prediction request (e.g. max_tokens, temperature, system). Returns: An iterable of PredictionResults. """ if inference_args is None: inference_args = {} if self.system is not None and 'system' not in inference_args: inference_args['system'] = self.system if self.output_config is not None and 'output_config' not in inference_args: inference_args['output_config'] = self.output_config responses = self.request_fn(self.model_name, batch, model, inference_args) return utils._convert_to_result(batch, responses, self.model_name) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/anthropic_inference_it_test.py0000644000175100017510000001563615206113032026350 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-End test for Anthropic Claude Remote Inference""" import logging import os import unittest import pytest try: from apache_beam.ml.inference.anthropic_inference import AnthropicModelHandler from apache_beam.ml.inference.anthropic_inference import message_from_conversation from apache_beam.ml.inference.anthropic_inference import message_from_string except ImportError: raise unittest.SkipTest("Anthropic dependencies are not installed") import apache_beam as beam from apache_beam.ml.inference.base import RunInference from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import is_not_empty _ANTHROPIC_API_KEY = os.environ.get('ANTHROPIC_API_KEY', None) _TEST_MODEL = 'claude-haiku-4-5' def _extract_text(prediction_result): return prediction_result.inference.content[0].text @pytest.mark.anthropic_postcommit class AnthropicInferenceIT(unittest.TestCase): @unittest.skipIf( _ANTHROPIC_API_KEY is None, 'ANTHROPIC_API_KEY environment variable is not set') def test_anthropic_text_generation(self): handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=message_from_string, api_key=_ANTHROPIC_API_KEY, max_batch_size=1, ) prompts = [ 'What is Apache Beam in one sentence?', 'Name three distributed computing frameworks.', ] with TestPipeline() as p: results = ( p | beam.Create(prompts) | RunInference(handler) | beam.Map(_extract_text)) assert_that(results, is_not_empty()) @unittest.skipIf( _ANTHROPIC_API_KEY is None, 'ANTHROPIC_API_KEY environment variable is not set') def test_anthropic_conversation(self): handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=message_from_conversation, api_key=_ANTHROPIC_API_KEY, max_batch_size=1, ) conversations = [ [ { "role": "user", "content": "What is 2 + 2?" }, { "role": "assistant", "content": "4" }, { "role": "user", "content": "Add 3 to that." }, ], ] with TestPipeline() as p: results = ( p | beam.Create(conversations) | RunInference(handler) | beam.Map(_extract_text)) assert_that(results, is_not_empty()) @unittest.skipIf( _ANTHROPIC_API_KEY is None, 'ANTHROPIC_API_KEY environment variable is not set') def test_anthropic_with_system_prompt(self): handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=message_from_string, api_key=_ANTHROPIC_API_KEY, system='You are a pirate. Respond only in pirate speak.', max_batch_size=1, ) prompts = ['What is your name?'] with TestPipeline() as p: results = ( p | beam.Create(prompts) | RunInference(handler) | beam.Map(_extract_text)) assert_that(results, is_not_empty()) @unittest.skipIf( _ANTHROPIC_API_KEY is None, 'ANTHROPIC_API_KEY environment variable is not set') def test_anthropic_system_prompt_with_structured_output(self): handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=message_from_string, api_key=_ANTHROPIC_API_KEY, system=( "You are a counting bot. When asked to count objects, convert " "responses such that numbers that are multiples of 3 are written " "as 'Fizz' instead of the number."), output_config={ 'format': { 'type': 'json_schema', 'schema': { 'type': 'object', 'properties': { 'items': { 'type': 'array', 'items': { 'type': 'object', 'properties': { 'name': { 'type': 'string' }, 'count': { 'type': 'string' }, }, 'required': ['name', 'count'], 'additionalProperties': False, } } }, 'required': ['items'], 'additionalProperties': False, }, }, }, max_batch_size=1, ) prompts = ['I have 2 apples, 3 bananas, and 6 oranges. Count them.'] with TestPipeline() as p: results = ( p | beam.Create(prompts) | RunInference(handler) | beam.Map(_extract_text)) def verify_fizz(response_text): import json data = json.loads(response_text) items = data.get('items', []) if not items: raise ValueError('Expected items in response') found_banana = False found_orange = False for item in items: name = item['name'].lower() count = str(item['count']) if 'banana' in name: found_banana = True if count != 'Fizz': raise ValueError('Expected banana count Fizz, ' 'got %s' % count) elif 'orange' in name: found_orange = True if count != 'Fizz': raise ValueError('Expected orange count Fizz, ' 'got %s' % count) elif 'apple' in name: if count != '2': raise ValueError('Expected apple count 2, ' 'got %s' % count) if not found_banana or not found_orange: raise ValueError('Missing expected items: %s' % response_text) return response_text _ = results | beam.Map(verify_fizz) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/anthropic_inference_test.py0000644000175100017510000002743415206113032025653 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest from dataclasses import dataclass from unittest import mock try: from anthropic import APIStatusError from apache_beam.ml.inference.anthropic_inference import AnthropicModelHandler from apache_beam.ml.inference.anthropic_inference import _retry_on_appropriate_error from apache_beam.ml.inference.anthropic_inference import message_from_conversation from apache_beam.ml.inference.anthropic_inference import message_from_string except ImportError: raise unittest.SkipTest('Anthropic dependencies are not installed') import apache_beam as beam from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to _TEST_MODEL = 'claude-haiku-4-5' @dataclass class FakeContentBlock: text: str type: str = 'text' @dataclass class FakeMessage: """Picklable stand-in for anthropic.types.Message.""" content: list model: str = _TEST_MODEL stop_reason: str = 'end_turn' def _make_fake_response(text): return FakeMessage(content=[FakeContentBlock(text=text)]) class RetryOnErrorTest(unittest.TestCase): def test_retry_on_rate_limit(self): e = APIStatusError( message="Rate limited", response=mock.MagicMock(status_code=429, headers={}), body=None) self.assertTrue(_retry_on_appropriate_error(e)) def test_retry_on_server_error(self): e = APIStatusError( message="Internal server error", response=mock.MagicMock(status_code=500, headers={}), body=None) self.assertTrue(_retry_on_appropriate_error(e)) def test_retry_on_503(self): e = APIStatusError( message="Service unavailable", response=mock.MagicMock(status_code=503, headers={}), body=None) self.assertTrue(_retry_on_appropriate_error(e)) def test_no_retry_on_400(self): e = APIStatusError( message="Bad request", response=mock.MagicMock(status_code=400, headers={}), body=None) self.assertFalse(_retry_on_appropriate_error(e)) def test_no_retry_on_401(self): e = APIStatusError( message="Unauthorized", response=mock.MagicMock(status_code=401, headers={}), body=None) self.assertFalse(_retry_on_appropriate_error(e)) def test_no_retry_on_non_api_error(self): self.assertFalse(_retry_on_appropriate_error(ValueError("oops"))) self.assertFalse(_retry_on_appropriate_error(RuntimeError("fail"))) class MessageFromStringTest(unittest.TestCase): def test_sends_each_prompt(self): client = mock.MagicMock() client.messages.create.side_effect = [ _make_fake_response("answer 1"), _make_fake_response("answer 2"), ] results = message_from_string(_TEST_MODEL, ['hello', 'world'], client, {}) self.assertEqual(len(results), 2) self.assertEqual(client.messages.create.call_count, 2) call_args = client.messages.create.call_args_list[0] self.assertEqual(call_args.kwargs['model'], _TEST_MODEL) self.assertEqual( call_args.kwargs['messages'], [{ "role": "user", "content": "hello" }]) def test_passes_inference_args(self): client = mock.MagicMock() client.messages.create.return_value = _make_fake_response("ok") message_from_string( _TEST_MODEL, ['test'], client, { 'max_tokens': 2048, 'temperature': 0.5 }) call_args = client.messages.create.call_args self.assertEqual(call_args.kwargs['max_tokens'], 2048) self.assertEqual(call_args.kwargs['temperature'], 0.5) def test_default_max_tokens(self): client = mock.MagicMock() client.messages.create.return_value = _make_fake_response("ok") message_from_string(_TEST_MODEL, ['test'], client, {}) call_args = client.messages.create.call_args self.assertEqual(call_args.kwargs['max_tokens'], 1024) class MessageFromConversationTest(unittest.TestCase): def test_sends_conversation(self): client = mock.MagicMock() client.messages.create.return_value = _make_fake_response("Paris!") convo = [ { "role": "user", "content": "What is the capital of France?" }, ] results = message_from_conversation(_TEST_MODEL, [convo], client, {}) self.assertEqual(len(results), 1) call_args = client.messages.create.call_args self.assertEqual(call_args.kwargs['messages'], convo) class AnthropicModelHandlerTest(unittest.TestCase): @mock.patch('apache_beam.ml.inference.anthropic_inference.Anthropic') def test_create_client_with_api_key(self, mock_anthropic): handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=message_from_string, api_key='test-key-123') handler.create_client() mock_anthropic.assert_called_once_with(api_key='test-key-123') @mock.patch('apache_beam.ml.inference.anthropic_inference.Anthropic') def test_create_client_from_env(self, mock_anthropic): handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=message_from_string) handler.create_client() mock_anthropic.assert_called_once_with() def test_request_returns_prediction_results(self): handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=message_from_string, api_key='fake') mock_client = mock.MagicMock() resp1 = _make_fake_response("answer 1") resp2 = _make_fake_response("answer 2") mock_client.messages.create.side_effect = [resp1, resp2] results = list(handler.request(['q1', 'q2'], mock_client, {})) self.assertEqual(len(results), 2) self.assertIsInstance(results[0], PredictionResult) self.assertEqual(results[0].example, 'q1') self.assertEqual(results[0].inference, resp1) self.assertEqual(results[0].model_id, _TEST_MODEL) self.assertEqual(results[1].example, 'q2') self.assertEqual(results[1].inference, resp2) def test_batch_elements_kwargs(self): handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=message_from_string, api_key='fake', min_batch_size=2, max_batch_size=10) kwargs = handler.batch_elements_kwargs() self.assertEqual(kwargs['min_batch_size'], 2) self.assertEqual(kwargs['max_batch_size'], 10) def _fake_request_fn(model_name, batch, client, inference_args): """A picklable request function that returns fake responses.""" return [ FakeMessage(content=[FakeContentBlock(text=f'answer for: {p}')]) for p in batch ] class SystemPromptTest(unittest.TestCase): def test_system_prompt_injected(self): handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=message_from_string, api_key='fake', system='Be concise.') mock_client = mock.MagicMock() mock_client.messages.create.return_value = _make_fake_response("ok") handler.request(['test'], mock_client, {}) call_args = mock_client.messages.create.call_args self.assertEqual(call_args.kwargs['system'], 'Be concise.') def test_system_prompt_not_overridden_by_handler(self): handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=message_from_string, api_key='fake', system='Handler system prompt.') mock_client = mock.MagicMock() mock_client.messages.create.return_value = _make_fake_response("ok") handler.request(['test'], mock_client, {'system': 'Per-request override.'}) call_args = mock_client.messages.create.call_args self.assertEqual(call_args.kwargs['system'], 'Per-request override.') def test_no_system_prompt_when_none(self): handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=message_from_string, api_key='fake') mock_client = mock.MagicMock() mock_client.messages.create.return_value = _make_fake_response("ok") handler.request(['test'], mock_client, {}) call_args = mock_client.messages.create.call_args self.assertNotIn('system', call_args.kwargs) class OutputConfigTest(unittest.TestCase): _SCHEMA = { 'format': { 'type': 'json_schema', 'schema': { 'type': 'object', 'properties': { 'answer': { 'type': 'string' } }, 'required': ['answer'], 'additionalProperties': False, }, }, } def test_output_config_injected(self): handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=message_from_string, api_key='fake', output_config=self._SCHEMA) mock_client = mock.MagicMock() mock_client.messages.create.return_value = ( _make_fake_response('{"answer":"ok"}')) handler.request(['test'], mock_client, {}) call_args = mock_client.messages.create.call_args self.assertEqual(call_args.kwargs['output_config'], self._SCHEMA) def test_output_config_not_overridden_by_handler(self): handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=message_from_string, api_key='fake', output_config=self._SCHEMA) mock_client = mock.MagicMock() mock_client.messages.create.return_value = _make_fake_response('{}') override = {'format': {'type': 'text'}} handler.request(['test'], mock_client, {'output_config': override}) call_args = mock_client.messages.create.call_args self.assertEqual(call_args.kwargs['output_config'], override) def test_no_output_config_when_none(self): handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=message_from_string, api_key='fake') mock_client = mock.MagicMock() mock_client.messages.create.return_value = _make_fake_response("ok") handler.request(['test'], mock_client, {}) call_args = mock_client.messages.create.call_args self.assertNotIn('output_config', call_args.kwargs) class AnthropicRunInferencePipelineTest(unittest.TestCase): def test_pipeline_e2e(self): """Full pipeline test with a fake request function.""" handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=_fake_request_fn, api_key='fake-key', max_batch_size=5, ) prompts = ['What is Beam?', 'What is MapReduce?'] with TestPipeline() as p: results = ( p | beam.Create(prompts) | RunInference(handler) | beam.Map(lambda r: r.example)) assert_that(results, equal_to(prompts)) def test_pipeline_with_system_prompt(self): """Pipeline test that verifies system prompt flows through.""" handler = AnthropicModelHandler( model_name=_TEST_MODEL, request_fn=_fake_request_fn, api_key='fake-key', system='You respond in haiku form.', max_batch_size=5, ) prompts = ['Tell me about Beam.'] with TestPipeline() as p: results = ( p | beam.Create(prompts) | RunInference(handler) | beam.Map(lambda r: r.example)) assert_that(results, equal_to(prompts)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/base.py0000644000175100017510000025605015206113032021517 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # TODO: https://github.com/apache/beam/issues/21822 # mypy: ignore-errors """An extensible run inference transform. Users of this module can extend the ModelHandler class for any machine learning framework. A ModelHandler implementation is a required parameter of RunInference. The transform handles standard inference functionality, like metric collection, sharing model between threads, and batching elements. """ import functools import logging import os import pickle import sys import threading import time import uuid from abc import ABC from abc import abstractmethod from collections import OrderedDict from collections import defaultdict from collections.abc import Callable from collections.abc import Iterable from collections.abc import Mapping from collections.abc import Sequence from copy import deepcopy from dataclasses import dataclass from datetime import datetime from datetime import timedelta from typing import Any from typing import Generic from typing import NamedTuple from typing import Optional from typing import TypeVar from typing import Union import apache_beam as beam from apache_beam.io.components.adaptive_throttler import ReactiveThrottler from apache_beam.utils import multi_process_shared from apache_beam.utils import retry from apache_beam.utils import shared try: from apache_beam.io.components.rate_limiter import RateLimiter except ImportError: RateLimiter = None try: # pylint: disable=wrong-import-order, wrong-import-position import resource from apache_beam.ml.inference.model_manager import ModelManager except ImportError: resource = None # type: ignore[assignment] ModelManager = None # type: ignore[assignment] _NANOSECOND_TO_MILLISECOND = 1_000_000 _NANOSECOND_TO_MICROSECOND = 1_000 _MILLISECOND_TO_SECOND = 1_000 ModelT = TypeVar('ModelT') ExampleT = TypeVar('ExampleT') PreProcessT = TypeVar('PreProcessT') PredictionT = TypeVar('PredictionT') PostProcessT = TypeVar('PostProcessT') _INPUT_TYPE = TypeVar('_INPUT_TYPE') _OUTPUT_TYPE = TypeVar('_OUTPUT_TYPE') KeyT = TypeVar('KeyT') # We use NamedTuple to define the structure of the PredictionResult, # however, as support for generic NamedTuples is not available in Python # versions prior to 3.11, we use the __new__ method to provide default # values for the fields while maintaining backwards compatibility. class PredictionResult(NamedTuple('PredictionResult', [('example', _INPUT_TYPE), ('inference', _OUTPUT_TYPE), ('model_id', Optional[str])])): __slots__ = () def __new__(cls, example, inference, model_id=None): return super().__new__(cls, example, inference, model_id) PredictionResult.__doc__ = """A NamedTuple containing both input and output from the inference.""" PredictionResult.example.__doc__ = """The input example.""" PredictionResult.inference.__doc__ = """Results for the inference on the model for the given example.""" PredictionResult.model_id.__doc__ = """Model ID used to run the prediction.""" class RateLimitExceeded(RuntimeError): """RateLimit Exceeded to process a batch of requests.""" pass class ModelMetadata(NamedTuple): model_id: str model_name: str class RunInferenceDLQ(NamedTuple): failed_inferences: beam.PCollection failed_preprocessing: Sequence[beam.PCollection] failed_postprocessing: Sequence[beam.PCollection] class _ModelLoadStats(NamedTuple): model_tag: str load_latency: Optional[int] byte_size: Optional[int] ModelMetadata.model_id.__doc__ = """Unique identifier for the model. This can be a file path or a URL where the model can be accessed. It is used to load the model for inference.""" ModelMetadata.model_name.__doc__ = """Human-readable name for the model. This can be used to identify the model in the metrics generated by the RunInference transform.""" def _to_milliseconds(time_ns: int) -> int: return int(time_ns / _NANOSECOND_TO_MILLISECOND) def _to_microseconds(time_ns: int) -> int: return int(time_ns / _NANOSECOND_TO_MICROSECOND) @dataclass(frozen=True) class KeyModelPathMapping(Generic[KeyT]): """ Dataclass for mapping 1 or more keys to 1 model path. This is used in conjunction with a KeyedModelHandler with many model handlers to update a set of keys' model handlers with the new path. Given `KeyModelPathMapping(keys: ['key1', 'key2'], update_path: 'updated/path', model_id: 'id1')`, all examples with keys `key1` or `key2` will have their corresponding model handler's update_model function called with 'updated/path' and their metrics will correspond with 'id1'. For more information see the KeyedModelHandler documentation https://beam.apache.org/releases/pydoc/current/apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.KeyedModelHandler documentation and the website section on model updates https://beam.apache.org/documentation/ml/about-ml/#automatic-model-refresh """ keys: list[KeyT] update_path: str model_id: str = '' class ModelHandler(Generic[ExampleT, PredictionT, ModelT]): """Has the ability to load and apply an ML model.""" def __init__( self, *, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, large_model: bool = False, model_copies: Optional[int] = None, **kwargs): """Initializes the ModelHandler. Args: min_batch_size: the minimum batch size to use when batching inputs. max_batch_size: the maximum batch size to use when batching inputs. max_batch_duration_secs: the maximum amount of time to buffer a batch before emitting; used in streaming contexts. max_batch_weight: the maximum weight of a batch. Requires element_size_fn. element_size_fn: a function that returns the size (weight) of an element. batch_length_fn: a callable mapping an element to its length (int). When set together with max_batch_duration_secs, enables length-aware bucketed keying so that elements of similar length are batched together, reducing padding waste for variable-length inputs. Bucket assignment uses bisect_right so boundaries are lower-inclusive: e.g., for boundaries [10, 50], buckets are (-inf, 10), [10, 50), [50, inf). batch_bucket_boundaries: a sorted list of positive boundary values for length bucketing. Boundaries are lower-inclusive (bisect_right semantics): bucket i covers lengths in [boundaries[i-1], boundaries[i]). Requires batch_length_fn. Defaults to [16, 32, 64, 128, 256, 512] when batch_length_fn is set. large_model: set to true if your model is large enough to run into memory pressure if you load multiple copies. model_copies: The exact number of models that you would like loaded onto your machine. kwargs: 'env_vars' can be used to set environment variables before loading the model. """ self._env_vars = kwargs.get('env_vars', {}) self._batching_kwargs: dict[str, Any] = {} if min_batch_size is not None: self._batching_kwargs['min_batch_size'] = min_batch_size if max_batch_size is not None: self._batching_kwargs['max_batch_size'] = max_batch_size if max_batch_duration_secs is not None: self._batching_kwargs['max_batch_duration_secs'] = max_batch_duration_secs if max_batch_weight is not None: self._batching_kwargs['max_batch_weight'] = max_batch_weight if element_size_fn is not None: self._batching_kwargs['element_size_fn'] = element_size_fn if batch_length_fn is not None: self._batching_kwargs['length_fn'] = batch_length_fn if batch_bucket_boundaries is not None: self._batching_kwargs['bucket_boundaries'] = batch_bucket_boundaries self._large_model = large_model self._model_copies = model_copies self._share_across_processes = large_model or (model_copies is not None) def load_model(self) -> ModelT: """Loads and initializes a model for processing.""" raise NotImplementedError(type(self)) def run_inference( self, batch: Sequence[ExampleT], model: ModelT, inference_args: Optional[dict[str, Any]] = None) -> Iterable[PredictionT]: """Runs inferences on a batch of examples. Args: batch: A sequence of examples or features. model: The model used to make inferences. inference_args: Extra arguments for models whose inference call requires extra parameters. Returns: An Iterable of Predictions. """ raise NotImplementedError(type(self)) def get_num_bytes(self, batch: Sequence[ExampleT]) -> int: """ Returns: The number of bytes of data for a batch. """ return len(pickle.dumps(batch)) def get_metrics_namespace(self) -> str: """ Returns: A namespace for metrics collected by the RunInference transform. """ return 'RunInference' def get_resource_hints(self) -> dict: """ Returns: Resource hints for the transform. """ return {} def batch_elements_kwargs(self) -> Mapping[str, Any]: """ Returns: kwargs suitable for beam.BatchElements. """ return getattr(self, '_batching_kwargs', {}) def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): """ Allows model handlers to provide some validation to make sure passed in inference args are valid. Some ModelHandlers throw here to disallow inference args altogether. """ pass def update_model_path(self, model_path: Optional[str] = None): """ Update the model path produced by side inputs. update_model_path should be used when a ModelHandler represents a single model, not multiple models. This will be true in most cases. For more information see the website section on model updates https://beam.apache.org/documentation/ml/about-ml/#automatic-model-refresh """ pass def update_model_paths( self, model: ModelT, model_paths: Optional[Union[str, list[KeyModelPathMapping]]] = None): """ Update the model paths produced by side inputs. update_model_paths should be used when updating multiple models at once (e.g. when using a KeyedModelHandler that holds multiple models). For more information see the KeyedModelHandler documentation https://beam.apache.org/releases/pydoc/current/apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.KeyedModelHandler documentation and the website section on model updates https://beam.apache.org/documentation/ml/about-ml/#automatic-model-refresh """ pass def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: """Gets all preprocessing functions to be run before batching/inference. Functions are in order that they should be applied.""" return [] def get_postprocess_fns(self) -> Iterable[Callable[[Any], Any]]: """Gets all postprocessing functions to be run after inference. Functions are in order that they should be applied.""" return [] def should_skip_batching(self) -> bool: """Whether RunInference's batching should be skipped. Can be flipped to True by using `with_no_batching`""" return False def set_environment_vars(self): """Sets environment variables using a dictionary provided via kwargs. Keys are the env variable name, and values are the env variable value. Child ModelHandler classes should set _env_vars via kwargs in __init__, or else call super().__init__().""" env_vars = getattr(self, '_env_vars', {}) for env_variable, env_value in env_vars.items(): os.environ[env_variable] = env_value def with_preprocess_fn( self, fn: Callable[[PreProcessT], ExampleT] ) -> 'ModelHandler[PreProcessT, PredictionT, ModelT]': """Returns a new ModelHandler with a preprocessing function associated with it. The preprocessing function will be run before batching/inference and should map your input PCollection to the base ModelHandler's input type. If you apply multiple preprocessing functions, they will be run on your original PCollection in order from last applied to first applied.""" return _PreProcessingModelHandler(self, fn) def with_postprocess_fn( self, fn: Callable[[PredictionT], PostProcessT] ) -> 'ModelHandler[ExampleT, PostProcessT, ModelT]': """Returns a new ModelHandler with a postprocessing function associated with it. The postprocessing function will be run after inference and should map the base ModelHandler's output type to your desired output type. If you apply multiple postprocessing functions, they will be run on your original inference result in order from first applied to last applied.""" return _PostProcessingModelHandler(self, fn) def with_no_batching( self ) -> """ModelHandler[Union[ ExampleT, Iterable[ExampleT]], PostProcessT, ModelT, PostProcessT]""": """Returns a new ModelHandler which does not require batching of inputs so that RunInference will skip this step. RunInference will expect the input to be pre-batched and passed in as an Iterable of records. If you skip batching, any preprocessing functions should accept a batch of data, not just a single record. This option is only recommended if you want to do custom batching yourself. If you just want to pass in records without a batching dimension, it is recommended to (1) add `max_batch_size=1` to `batch_elements_kwargs` and (2) remove the batching dimension as part of your inference call (by calling `record=batch[0]`)""" return _PrebatchedModelHandler(self) def share_model_across_processes(self) -> bool: """Returns a boolean representing whether or not a model should be shared across multiple processes instead of being loaded per process. This is primary useful for large models that can't fit multiple copies in memory. Multi-process support may vary by runner, but this will fallback to loading per process as necessary. See https://beam.apache.org/releases/pydoc/current/apache_beam.utils.multi_process_shared.html""" return getattr(self, '_share_across_processes', False) def model_copies(self) -> int: """Returns the maximum number of model copies that should be loaded at one time. This only impacts model handlers that are using share_model_across_processes to share their model across processes instead of being loaded per process.""" return getattr(self, '_model_copies', None) or 1 def override_metrics(self, metrics_namespace: str = '') -> bool: """Returns a boolean representing whether or not a model handler will override metrics reporting. If True, RunInference will not report any metrics.""" return False def should_garbage_collect_on_timeout(self) -> bool: """Whether the model should be garbage collected if model loading or inference timeout, or if it should be left for future calls. Usually should not be overriden unless the model handler implements other mechanisms for garbage collection.""" return self.share_model_across_processes() class RemoteModelHandler(ABC, ModelHandler[ExampleT, PredictionT, ModelT]): """Has the ability to call a model at a remote endpoint.""" def __init__( self, namespace: str = '', num_retries: int = 5, throttle_delay_secs: int = 5, retry_filter: Callable[[Exception], bool] = lambda x: True, *, window_ms: int = 1 * _MILLISECOND_TO_SECOND, bucket_ms: int = 1 * _MILLISECOND_TO_SECOND, overload_ratio: float = 2, rate_limiter: Optional[RateLimiter] = None): """Initializes a ReactiveThrottler class for enabling client-side throttling for remote calls to an inference service. Also wraps provided calls to the service with retry logic. See https://s.apache.org/beam-client-side-throttling for more details on the configuration of the throttling and retry mechanics. Args: namespace: the metrics and logging namespace num_retries: the maximum number of times to retry a request on retriable errors before failing throttle_delay_secs: the amount of time to throttle when the client-side elects to throttle retry_filter: a function accepting an exception as an argument and returning a boolean. On a true return, the run_inference call will be retried. Defaults to always retrying. window_ms: length of history to consider, in ms, to set throttling. bucket_ms: granularity of time buckets that we store data in, in ms. overload_ratio: the target ratio between requests sent and successful requests. This is "K" in the formula in https://landing.google.com/sre/book/chapters/handling-overload.html. rate_limiter: A RateLimiter object for setting a global rate limit. """ # Configure ReactiveThrottler for client-side throttling behavior. self.throttler = ReactiveThrottler( window_ms=window_ms, bucket_ms=bucket_ms, overload_ratio=overload_ratio, namespace=namespace, throttle_delay_secs=throttle_delay_secs) self.logger = logging.getLogger(namespace) self.num_retries = num_retries self.retry_filter = retry_filter self._rate_limiter = rate_limiter self._shared_rate_limiter = None self._shared_handle = shared.Shared() def __init_subclass__(cls): if cls.load_model is not RemoteModelHandler.load_model: raise Exception( "Cannot override RemoteModelHandler.load_model, ", "implement create_client instead.") if cls.run_inference is not RemoteModelHandler.run_inference: raise Exception( "Cannot override RemoteModelHandler.run_inference, ", "implement request instead.") @abstractmethod def create_client(self) -> ModelT: """Creates the client that is used to make the remote inference request in request(). All relevant arguments should be passed to __init__(). """ raise NotImplementedError(type(self)) def load_model(self) -> ModelT: return self.create_client() def retry_on_exception(func): @functools.wraps(func) def wrapper(self, *args, **kwargs): return retry.with_exponential_backoff( num_retries=self.num_retries, retry_filter=self.retry_filter)(func)(self, *args, **kwargs) return wrapper @retry_on_exception def run_inference( self, batch: Sequence[ExampleT], model: ModelT, inference_args: Optional[dict[str, Any]] = None) -> Iterable[PredictionT]: """Runs inferences on a batch of examples. Calls a remote model for predictions and will retry if a retryable exception is raised. Args: batch: A sequence of examples or features. model: The model used to make inferences. inference_args: Extra arguments for models whose inference call requires extra parameters. Returns: An Iterable of Predictions. """ if self._rate_limiter: if self._shared_rate_limiter is None: def init_limiter(): return self._rate_limiter self._shared_rate_limiter = self._shared_handle.acquire(init_limiter) if not self._shared_rate_limiter.allow(hits_added=len(batch)): raise RateLimitExceeded( "Rate Limit Exceeded, " "Could not process this batch.") self.throttler.throttle() try: req_time = time.time() predictions = self.request(batch, model, inference_args) self.throttler.successful_request(req_time * _MILLISECOND_TO_SECOND) return predictions except Exception as e: self.logger.error("exception raised as part of request, got %s", e) raise @abstractmethod def request( self, batch: Sequence[ExampleT], model: ModelT, inference_args: Optional[dict[str, Any]] = None) -> Iterable[PredictionT]: """Makes a request to a remote inference service and returns the response. Should raise an exception of some kind if there is an error to enable the retry and client-side throttling logic to work. Returns an iterable of the desired prediction type. This method should return the values directly, as handling return values as a generator can prevent the retry logic from functioning correctly. Args: batch: A sequence of examples or features. model: The model used to make inferences. inference_args: Extra arguments for models whose inference call requires extra parameters. Returns: An Iterable of Predictions. """ raise NotImplementedError(type(self)) class _ModelHandlerManager: """ A class for efficiently managing copies of multiple model handlers. Will load a single copy of each model from the model handler into a multi_process_shared object and then return a lookup key for that object. Used for KeyedModelHandler only. """ def __init__(self, mh_map: dict[str, ModelHandler]): """ Args: mh_map: A map from keys to model handlers which can be used to load a model. """ self._max_models = None # Map keys to model handlers self._mh_map: dict[str, ModelHandler] = mh_map # Map keys to the last updated model path for that key self._key_to_last_update: dict[str, str] = defaultdict(str) # Map key for a model to a unique tag that will persist for the life of # that model in memory. A new tag will be generated if a model is swapped # out of memory and reloaded. self._tag_map: dict[str, str] = OrderedDict() # Map a tag to a multiprocessshared model object for that tag. Each entry # of this map should last as long as the corresponding entry in _tag_map. self._proxy_map: dict[str, multi_process_shared.MultiProcessShared] = {} def load(self, key: str) -> _ModelLoadStats: """ Loads the appropriate model for the given key into memory. Args: key: the key associated with the model we'd like to load. Returns: _ModelLoadStats with tag, byte size, and latency to load the model. If the model was already loaded, byte size/latency will be None. """ # Map the key for a model to a unique tag that will persist until the model # is released. This needs to be unique between releasing/reacquiring th # model because otherwise the ProxyManager will try to reuse the model that # has been released and deleted. if key in self._tag_map: self._tag_map.move_to_end(key) return _ModelLoadStats(self._tag_map[key], None, None) else: self._tag_map[key] = uuid.uuid4().hex tag = self._tag_map[key] mh = self._mh_map[key] if self._max_models is not None and self._max_models < len(self._tag_map): # If we're about to exceed our LRU size, release the last used model. tag_to_remove = self._tag_map.popitem(last=False)[1] shared_handle, model_to_remove = self._proxy_map[tag_to_remove] shared_handle.release(model_to_remove) del self._proxy_map[tag_to_remove] # Load the new model memory_before = _get_current_process_memory_in_bytes() start_time = _to_milliseconds(time.time_ns()) shared_handle = multi_process_shared.MultiProcessShared( mh.load_model, tag=tag) model_reference = shared_handle.acquire() self._proxy_map[tag] = (shared_handle, model_reference) memory_after = _get_current_process_memory_in_bytes() end_time = _to_milliseconds(time.time_ns()) return _ModelLoadStats( tag, end_time - start_time, memory_after - memory_before) def increment_max_models(self, increment: int): """ Increments the number of models that this instance of a _ModelHandlerManager is able to hold. If it is never called, no limit is imposed. Args: increment: the amount by which we are incrementing the number of models. """ if self._max_models is None: self._max_models = 0 self._max_models += increment def update_model_handler(self, key: str, model_path: str, previous_key: str): """ Updates the model path of this model handler and removes it from memory so that it can be reloaded with the updated path. No-ops if no model update needs to be applied. Args: key: the key associated with the model we'd like to update. model_path: the new path to the model we'd like to load. previous_key: the key that is associated with the old version of this model. This will often be the same as the current key, but sometimes we will want to keep both the old and new models to serve different cohorts. In that case, the keys should be different. """ if self._key_to_last_update[key] == model_path: return self._key_to_last_update[key] = model_path if key not in self._mh_map: self._mh_map[key] = deepcopy(self._mh_map[previous_key]) self._mh_map[key].update_model_path(model_path) if key in self._tag_map: tag_to_remove = self._tag_map[key] shared_handle, model_to_remove = self._proxy_map[tag_to_remove] shared_handle.release(model_to_remove) del self._tag_map[key] del self._proxy_map[tag_to_remove] # Use a dataclass instead of named tuple because NamedTuples and generics don't # mix well across the board for all versions: # https://github.com/python/typing/issues/653 class KeyModelMapping(Generic[KeyT, ExampleT, PredictionT, ModelT]): """ Dataclass for mapping 1 or more keys to 1 model handler. Given `KeyModelMapping(['key1', 'key2'], myMh)`, all examples with keys `key1` or `key2` will be run against the model defined by the `myMh` ModelHandler. """ def __init__( self, keys: list[KeyT], mh: ModelHandler[ExampleT, PredictionT, ModelT]): self.keys = keys self.mh = mh class KeyedModelHandler(Generic[KeyT, ExampleT, PredictionT, ModelT], ModelHandler[tuple[KeyT, ExampleT], tuple[KeyT, PredictionT], Union[ModelT, _ModelHandlerManager]]): def __init__( self, unkeyed: Union[ModelHandler[ExampleT, PredictionT, ModelT], list[KeyModelMapping[KeyT, ExampleT, PredictionT, ModelT]]], max_models_per_worker_hint: Optional[int] = None): """A ModelHandler that takes keyed examples and returns keyed predictions. For example, if the original model is used with RunInference to take a PCollection[E] to a PCollection[P], this ModelHandler would take a PCollection[tuple[K, E]] to a PCollection[tuple[K, P]], making it possible to use the key to associate the outputs with the inputs. KeyedModelHandler is able to accept either a single unkeyed ModelHandler or many different model handlers corresponding to the keys for which that ModelHandler should be used. For example, the following configuration could be used to map keys 1-3 to ModelHandler1 and keys 4-5 to ModelHandler2: k1 = ['k1', 'k2', 'k3'] k2 = ['k4', 'k5'] KeyedModelHandler([KeyModelMapping(k1, mh1), KeyModelMapping(k2, mh2)]) Note that a single copy of each of these models may all be held in memory at the same time; be careful not to load too many large models or your pipeline may cause Out of Memory exceptions. KeyedModelHandlers support Automatic Model Refresh to update your model to a newer version without stopping your streaming pipeline. For an overview of this feature, see https://beam.apache.org/documentation/ml/about-ml/#automatic-model-refresh To use this feature with a KeyedModelHandler that has many models per key, you can pass in a list of KeyModelPathMapping objects to define your new model paths. For example, passing in the side input of [KeyModelPathMapping(keys=['k1', 'k2'], update_path='update/path/1'), KeyModelPathMapping(keys=['k3'], update_path='update/path/2')] will update the model corresponding to keys 'k1' and 'k2' with path 'update/path/1' and the model corresponding to 'k3' with 'update/path/2'. In order to do a side input update: (1) all restrictions mentioned in https://beam.apache.org/documentation/ml/about-ml/#automatic-model-refresh must be met, (2) all update_paths must be non-empty, even if they are not being updated from their original values, and (3) the set of keys originally defined cannot change. This means that if originally you have defined model handlers for 'key1', 'key2', and 'key3', all 3 of those keys must appear in your list of KeyModelPathMappings exactly once. No additional keys can be added. When using many models defined per key, metrics about inference and model loading will be gathered on an aggregate basis for all keys. These will be reported with no prefix. Metrics will also be gathered on a per key basis. Since some keys can share the same model, only one set of metrics will be reported per key 'cohort'. These will be reported in the form: `-`, where `` can be any key selected from the cohort. When model updates occur, the metrics will be reported in the form `--`. Loading multiple models at the same time can increase the risk of an out of memory (OOM) exception. To avoid this issue, use the parameter `max_models_per_worker_hint` to limit the number of models that are loaded at the same time. For more information about memory management, see `Use a keyed `ModelHandler _`. # pylint: disable=line-too-long Args: unkeyed: Either (a) an implementation of ModelHandler that does not require keys or (b) a list of KeyModelMappings mapping lists of keys to unkeyed ModelHandlers. max_models_per_worker_hint: A hint to the runner indicating how many models can be held in memory at one time per worker process. For example, if your worker has 8 GB of memory provisioned and your workers take up 1 GB each, you should set this to 7 to allow all models to sit in memory with some buffer. For more information about memory management, see `Use a keyed `ModelHandler _`. # pylint: disable=line-too-long """ self._metrics_collectors: dict[str, _MetricsCollector] = {} self._default_metrics_collector: _MetricsCollector = None self._metrics_namespace = '' self._single_model = not isinstance(unkeyed, list) if self._single_model: if len(unkeyed.get_preprocess_fns()) or len( unkeyed.get_postprocess_fns()): raise Exception( 'Cannot make make an unkeyed model handler with pre or ' 'postprocessing functions defined into a keyed model handler. All ' 'pre/postprocessing functions must be defined on the outer model' 'handler.') self._env_vars = getattr(unkeyed, '_env_vars', {}) self._unkeyed = unkeyed return self._max_models_per_worker_hint = max_models_per_worker_hint # To maintain an efficient representation, we will map all keys in a given # KeyModelMapping to a single id (the first key in the KeyModelMapping # list). We will then map that key to a ModelHandler. This will allow us to # quickly look up the appropriate ModelHandler for any given key. self._id_to_mh_map: dict[str, ModelHandler[ExampleT, PredictionT, ModelT]] = {} self._key_to_id_map: dict[str, str] = {} for mh_tuple in unkeyed: mh = mh_tuple.mh keys = mh_tuple.keys if len(mh.get_preprocess_fns()) or len(mh.get_postprocess_fns()): raise ValueError( 'Cannot use an unkeyed model handler with pre or ' 'postprocessing functions defined in a keyed model handler. All ' 'pre/postprocessing functions must be defined on the outer model' 'handler.') hints = mh.get_resource_hints() if len(hints) > 0: logging.warning( 'mh %s defines the following resource hints, which will be' 'ignored: %s. Resource hints are not respected when more than one ' 'model handler is used in a KeyedModelHandler. If you would like ' 'to specify resource hints, you can do so by overriding the ' 'KeyedModelHandler.get_resource_hints() method.', mh, hints) batch_kwargs = mh.batch_elements_kwargs() if len(batch_kwargs) > 0: logging.warning( 'mh %s defines the following batching kwargs which will be ' 'ignored %s. Batching kwargs are not respected when ' 'more than one model handler is used in a KeyedModelHandler. If ' 'you would like to specify resource hints, you can do so by ' 'overriding the KeyedModelHandler.batch_elements_kwargs() method.', hints, batch_kwargs) env_vars = getattr(mh, '_env_vars', {}) if len(env_vars) > 0: logging.warning( 'mh %s defines the following _env_vars which will be ignored %s. ' '_env_vars are not respected when more than one model handler is ' 'used in a KeyedModelHandler. If you need env vars set at ' 'inference time, you can do so with ' 'a custom inference function.', mh, env_vars) if len(keys) == 0: raise ValueError( f'Empty list maps to model handler {mh}. All model handlers must ' 'have one or more associated keys.') self._id_to_mh_map[keys[0]] = mh for key in keys: if key in self._key_to_id_map: raise ValueError( f'key {key} maps to multiple model handlers. All keys must map ' 'to exactly one model handler.') self._key_to_id_map[key] = keys[0] def load_model(self) -> Union[ModelT, _ModelHandlerManager]: if self._single_model: return self._unkeyed.load_model() return _ModelHandlerManager(self._id_to_mh_map) def run_inference( self, batch: Sequence[tuple[KeyT, ExampleT]], model: Union[ModelT, _ModelHandlerManager], inference_args: Optional[dict[str, Any]] = None ) -> Iterable[tuple[KeyT, PredictionT]]: if self._single_model: keys, unkeyed_batch = zip(*batch) return zip( keys, self._unkeyed.run_inference(unkeyed_batch, model, inference_args)) # The first time a MultiProcessShared ModelManager is used for inference # from this process, we should increment its max model count if self._max_models_per_worker_hint is not None: lock = threading.Lock() if lock.acquire(blocking=False): model.increment_max_models(self._max_models_per_worker_hint) self._max_models_per_worker_hint = None batch_by_key = defaultdict(list) key_by_id = defaultdict(set) for key, example in batch: batch_by_key[key].append(example) key_by_id[self._key_to_id_map[key]].add(key) predictions = [] for id, keys in key_by_id.items(): mh = self._id_to_mh_map[id] loaded_model = model.load(id) keyed_model_tag = loaded_model.model_tag if loaded_model.byte_size is not None: self._metrics_collectors[id].update_load_model_metrics( loaded_model.load_latency, loaded_model.byte_size) self._default_metrics_collector.update_load_model_metrics( loaded_model.load_latency, loaded_model.byte_size) keyed_model_shared_handle = multi_process_shared.MultiProcessShared( mh.load_model, tag=keyed_model_tag) keyed_model = keyed_model_shared_handle.acquire() start_time = _to_microseconds(time.time_ns()) num_bytes = 0 num_elements = 0 try: for key in keys: unkeyed_batches = batch_by_key[key] try: for inf in mh.run_inference(unkeyed_batches, keyed_model, inference_args): predictions.append((key, inf)) except BaseException as e: self._metrics_collectors[id].failed_batches_counter.inc() self._default_metrics_collector.failed_batches_counter.inc() raise e num_bytes += mh.get_num_bytes(unkeyed_batches) num_elements += len(unkeyed_batches) finally: keyed_model_shared_handle.release(keyed_model) end_time = _to_microseconds(time.time_ns()) inference_latency = end_time - start_time self._metrics_collectors[id].update( num_elements, num_bytes, inference_latency) self._default_metrics_collector.update( num_elements, num_bytes, inference_latency) return predictions def get_num_bytes(self, batch: Sequence[tuple[KeyT, ExampleT]]) -> int: keys, unkeyed_batch = zip(*batch) batch_bytes = len(pickle.dumps(keys)) if self._single_model: return batch_bytes + self._unkeyed.get_num_bytes(unkeyed_batch) batch_by_key = defaultdict(list) for key, examples in batch: batch_by_key[key].append(examples) for key, examples in batch_by_key.items(): mh_id = self._key_to_id_map[key] batch_bytes += self._id_to_mh_map[mh_id].get_num_bytes(examples) return batch_bytes def get_metrics_namespace(self) -> str: if self._single_model: return self._unkeyed.get_metrics_namespace() return 'BeamML_KeyedModels' def get_resource_hints(self): if self._single_model: return self._unkeyed.get_resource_hints() return {} def batch_elements_kwargs(self): if self._single_model: return self._unkeyed.batch_elements_kwargs() return {} def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): if self._single_model: return self._unkeyed.validate_inference_args(inference_args) for mh in self._id_to_mh_map.values(): mh.validate_inference_args(inference_args) def update_model_paths( self, model: Union[ModelT, _ModelHandlerManager], model_paths: list[KeyModelPathMapping[KeyT]] = None): # When there are many models, the keyed model handler is responsible for # reorganizing the model handlers into cohorts and telling the model # manager to update every cohort's associated model handler. The model # manager is responsible for performing the updates and tracking which # updates have already been applied. if model_paths is None or len(model_paths) == 0 or model is None: return if self._single_model: raise RuntimeError( 'Invalid model update: sent many model paths to ' 'update, but KeyedModelHandler is wrapping a single ' 'model.') # Map cohort ids to a dictionary mapping new model paths to the keys that # were originally in that cohort. We will use this to construct our new # cohorts. # cohort_path_mapping will be structured as follows: # { # original_cohort_id: { # 'update/path/1': ['key1FromOriginalCohort', key2FromOriginalCohort'], # 'update/path/2': ['key3FromOriginalCohort', key4FromOriginalCohort'], # } # } cohort_path_mapping: dict[KeyT, dict[str, list[KeyT]]] = {} key_modelid_mapping: dict[KeyT, str] = {} seen_keys = set() for mp in model_paths: keys = mp.keys update_path = mp.update_path model_id = mp.model_id if len(update_path) == 0: raise ValueError(f'Invalid model update, path for {keys} is empty') for key in keys: if key in seen_keys: raise ValueError( f'Invalid model update: {key} appears in multiple ' 'update lists. A single model update must provide exactly one ' 'updated path per key.') seen_keys.add(key) if key not in self._key_to_id_map: raise ValueError( f'Invalid model update: {key} appears in ' 'update, but not in the original configuration.') key_modelid_mapping[key] = model_id cohort_id = self._key_to_id_map[key] if cohort_id not in cohort_path_mapping: cohort_path_mapping[cohort_id] = defaultdict(list) cohort_path_mapping[cohort_id][update_path].append(key) for key in self._key_to_id_map: if key not in seen_keys: raise ValueError( f'Invalid model update: {key} appears in the ' 'original configuration, but not the update.') # We now have our new set of cohorts. For each one, update our local model # handler configuration and send the results to the ModelManager for old_cohort_id, path_key_mapping in cohort_path_mapping.items(): for updated_path, keys in path_key_mapping.items(): cohort_id = old_cohort_id if old_cohort_id not in keys: # Create new cohort cohort_id = keys[0] for key in keys: self._key_to_id_map[key] = cohort_id mh = self._id_to_mh_map[old_cohort_id] self._id_to_mh_map[cohort_id] = deepcopy(mh) self._id_to_mh_map[cohort_id].update_model_path(updated_path) model.update_model_handler(cohort_id, updated_path, old_cohort_id) model_id = key_modelid_mapping[cohort_id] self._metrics_collectors[cohort_id] = _MetricsCollector( self._metrics_namespace, f'{cohort_id}-{model_id}-') def update_model_path(self, model_path: Optional[str] = None): if self._single_model: return self._unkeyed.update_model_path(model_path=model_path) if model_path is not None: raise RuntimeError( 'Model updates are currently not supported for ' + 'KeyedModelHandlers with multiple different per-key ' + 'ModelHandlers.') def share_model_across_processes(self) -> bool: if self._single_model: return self._unkeyed.share_model_across_processes() return True def model_copies(self) -> int: if self._single_model: return self._unkeyed.model_copies() for mh in self._id_to_mh_map.values(): if mh.model_copies() != 1: raise ValueError( 'KeyedModelHandler cannot map records to multiple ' 'models if one or more of its ModelHandlers ' 'require multiple model copies (set via ' 'model_copies). To fix, verify that each ' 'ModelHandler is not set to load multiple copies of ' 'its model.') return 1 def override_metrics(self, metrics_namespace: str = '') -> bool: if self._single_model: return self._unkeyed.override_metrics(metrics_namespace) self._metrics_namespace = metrics_namespace self._default_metrics_collector = _MetricsCollector(metrics_namespace) for cohort_id in self._id_to_mh_map: self._metrics_collectors[cohort_id] = _MetricsCollector( metrics_namespace, f'{cohort_id}-') return True def should_garbage_collect_on_timeout(self) -> bool: return self._single_model and self.share_model_across_processes() class MaybeKeyedModelHandler(Generic[KeyT, ExampleT, PredictionT, ModelT], ModelHandler[Union[ExampleT, tuple[KeyT, ExampleT]], Union[PredictionT, tuple[KeyT, PredictionT]], ModelT]): def __init__(self, unkeyed: ModelHandler[ExampleT, PredictionT, ModelT]): """A ModelHandler that takes examples that might have keys and returns predictions that might have keys. For example, if the original model is used with RunInference to take a PCollection[E] to a PCollection[P], this ModelHandler would take either PCollection[E] to a PCollection[P] or PCollection[tuple[K, E]] to a PCollection[tuple[K, P]], depending on the whether the elements are tuples. This pattern makes it possible to associate the outputs with the inputs based on the key. Note that you cannot use this ModelHandler if E is a tuple type. In addition, either all examples should be keyed, or none of them. Args: unkeyed: An implementation of ModelHandler that does not require keys. """ if len(unkeyed.get_preprocess_fns()) or len(unkeyed.get_postprocess_fns()): raise Exception( 'Cannot make make an unkeyed model handler with pre or ' 'postprocessing functions defined into a keyed model handler. All ' 'pre/postprocessing functions must be defined on the outer model' 'handler.') self._unkeyed = unkeyed self._env_vars = getattr(unkeyed, '_env_vars', {}) def load_model(self) -> ModelT: return self._unkeyed.load_model() def run_inference( self, batch: Sequence[Union[ExampleT, tuple[KeyT, ExampleT]]], model: ModelT, inference_args: Optional[dict[str, Any]] = None ) -> Union[Iterable[PredictionT], Iterable[tuple[KeyT, PredictionT]]]: # Really the input should be # Union[Sequence[ExampleT], Sequence[tuple[KeyT, ExampleT]]] # but there's not a good way to express (or check) that. if isinstance(batch[0], tuple): is_keyed = True keys, unkeyed_batch = zip(*batch) # type: ignore[arg-type] else: is_keyed = False unkeyed_batch = batch # type: ignore[assignment] unkeyed_results = self._unkeyed.run_inference( unkeyed_batch, model, inference_args) if is_keyed: return zip(keys, unkeyed_results) else: return unkeyed_results def get_num_bytes( self, batch: Sequence[Union[ExampleT, tuple[KeyT, ExampleT]]]) -> int: # MyPy can't follow the branching logic. if isinstance(batch[0], tuple): keys, unkeyed_batch = zip(*batch) # type: ignore[arg-type] return len( pickle.dumps(keys)) + self._unkeyed.get_num_bytes(unkeyed_batch) else: return self._unkeyed.get_num_bytes(batch) # type: ignore[arg-type] def get_metrics_namespace(self) -> str: return self._unkeyed.get_metrics_namespace() def get_resource_hints(self): return self._unkeyed.get_resource_hints() def batch_elements_kwargs(self): return self._unkeyed.batch_elements_kwargs() def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): return self._unkeyed.validate_inference_args(inference_args) def update_model_path(self, model_path: Optional[str] = None): return self._unkeyed.update_model_path(model_path=model_path) def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._unkeyed.get_preprocess_fns() def get_postprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._unkeyed.get_postprocess_fns() def should_skip_batching(self) -> bool: return self._unkeyed.should_skip_batching() def share_model_across_processes(self) -> bool: return self._unkeyed.share_model_across_processes() def model_copies(self) -> int: return self._unkeyed.model_copies() class _PrebatchedModelHandler(Generic[ExampleT, PredictionT, ModelT], ModelHandler[Sequence[ExampleT], PredictionT, ModelT]): def __init__(self, base: ModelHandler[ExampleT, PredictionT, ModelT]): """A ModelHandler that skips batching in RunInference. Args: base: An implementation of the underlying model handler. """ self._base = base self._env_vars = getattr(base, '_env_vars', {}) def load_model(self) -> ModelT: return self._base.load_model() def run_inference( self, batch: Sequence[Union[ExampleT, tuple[KeyT, ExampleT]]], model: ModelT, inference_args: Optional[dict[str, Any]] = None ) -> Union[Iterable[PredictionT], Iterable[tuple[KeyT, PredictionT]]]: return self._base.run_inference(batch, model, inference_args) def get_num_bytes( self, batch: Sequence[Union[ExampleT, tuple[KeyT, ExampleT]]]) -> int: return self._base.get_num_bytes(batch) def get_metrics_namespace(self) -> str: return self._base.get_metrics_namespace() def get_resource_hints(self): return self._base.get_resource_hints() def batch_elements_kwargs(self): return self._base.batch_elements_kwargs() def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): return self._base.validate_inference_args(inference_args) def update_model_path(self, model_path: Optional[str] = None): return self._base.update_model_path(model_path=model_path) def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._base.get_preprocess_fns() def should_skip_batching(self) -> bool: return True def share_model_across_processes(self) -> bool: return self._base.share_model_across_processes() def model_copies(self) -> int: return self._base.model_copies() def get_postprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._base.get_postprocess_fns() class _PreProcessingModelHandler(Generic[ExampleT, PredictionT, ModelT, PreProcessT], ModelHandler[PreProcessT, PredictionT, ModelT]): def __init__( self, base: ModelHandler[ExampleT, PredictionT, ModelT], preprocess_fn: Callable[[PreProcessT], ExampleT]): """A ModelHandler that has a preprocessing function associated with it. Args: base: An implementation of the underlying model handler. preprocess_fn: the preprocessing function to use. """ self._base = base self._env_vars = getattr(base, '_env_vars', {}) self._preprocess_fn = preprocess_fn def load_model(self) -> ModelT: return self._base.load_model() def run_inference( self, batch: Sequence[Union[ExampleT, tuple[KeyT, ExampleT]]], model: ModelT, inference_args: Optional[dict[str, Any]] = None ) -> Union[Iterable[PredictionT], Iterable[tuple[KeyT, PredictionT]]]: return self._base.run_inference(batch, model, inference_args) def get_num_bytes( self, batch: Sequence[Union[ExampleT, tuple[KeyT, ExampleT]]]) -> int: return self._base.get_num_bytes(batch) def get_metrics_namespace(self) -> str: return self._base.get_metrics_namespace() def get_resource_hints(self): return self._base.get_resource_hints() def batch_elements_kwargs(self): return self._base.batch_elements_kwargs() def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): return self._base.validate_inference_args(inference_args) def update_model_path(self, model_path: Optional[str] = None): return self._base.update_model_path(model_path=model_path) def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return [self._preprocess_fn] + self._base.get_preprocess_fns() def should_skip_batching(self) -> bool: return self._base.should_skip_batching() def share_model_across_processes(self) -> bool: return self._base.share_model_across_processes() def model_copies(self) -> int: return self._base.model_copies() def get_postprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._base.get_postprocess_fns() class _PostProcessingModelHandler(Generic[ExampleT, PredictionT, ModelT, PostProcessT], ModelHandler[ExampleT, PostProcessT, ModelT]): def __init__( self, base: ModelHandler[ExampleT, PredictionT, ModelT], postprocess_fn: Callable[[PredictionT], PostProcessT]): """A ModelHandler that has a preprocessing function associated with it. Args: base: An implementation of the underlying model handler. postprocess_fn: the preprocessing function to use. """ self._base = base self._env_vars = getattr(base, '_env_vars', {}) self._postprocess_fn = postprocess_fn def load_model(self) -> ModelT: return self._base.load_model() def run_inference( self, batch: Sequence[Union[ExampleT, tuple[KeyT, ExampleT]]], model: ModelT, inference_args: Optional[dict[str, Any]] = None ) -> Union[Iterable[PredictionT], Iterable[tuple[KeyT, PredictionT]]]: return self._base.run_inference(batch, model, inference_args) def get_num_bytes( self, batch: Sequence[Union[ExampleT, tuple[KeyT, ExampleT]]]) -> int: return self._base.get_num_bytes(batch) def get_metrics_namespace(self) -> str: return self._base.get_metrics_namespace() def get_resource_hints(self): return self._base.get_resource_hints() def batch_elements_kwargs(self): return self._base.batch_elements_kwargs() def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): return self._base.validate_inference_args(inference_args) def update_model_path(self, model_path: Optional[str] = None): return self._base.update_model_path(model_path=model_path) def get_preprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._base.get_preprocess_fns() def should_skip_batching(self) -> bool: return self._base.should_skip_batching() def share_model_across_processes(self) -> bool: return self._base.share_model_across_processes() def model_copies(self) -> int: return self._base.model_copies() def get_postprocess_fns(self) -> Iterable[Callable[[Any], Any]]: return self._base.get_postprocess_fns() + [self._postprocess_fn] class OOMProtectedFn: def __init__(self, func): self.func = func def __call__(self, *args, **kwargs): try: return self.func(*args, **kwargs) except Exception as e: # Check string to avoid hard import dependency if 'out of memory' in str(e) and 'CUDA' in str(e): logging.warning("Caught CUDA OOM during operation. Cleaning memory.") try: import gc import torch gc.collect() torch.cuda.empty_cache() except ImportError: pass except Exception as cleanup_error: logging.error("Failed to clean up CUDA memory: %s", cleanup_error) raise e class RunInference(beam.PTransform[beam.PCollection[Union[ExampleT, Iterable[ExampleT]]], beam.PCollection[PredictionT]]): def __init__( self, model_handler: ModelHandler[ExampleT, PredictionT, Any], clock=time, inference_args: Optional[dict[str, Any]] = None, metrics_namespace: Optional[str] = None, *, model_metadata_pcoll: beam.PCollection[ModelMetadata] = None, watch_model_pattern: Optional[str] = None, model_identifier: Optional[str] = None, use_model_manager: bool = False, model_manager_args: Optional[dict[str, Any]] = None, **kwargs): """ A transform that takes a PCollection of examples (or features) for use on an ML model. The transform then outputs inferences (or predictions) for those examples in a PCollection of PredictionResults that contains the input examples and the output inferences. Models for supported frameworks can be loaded using a URI. Supported services can also be used. This transform attempts to batch examples using the beam.BatchElements transform. Batching can be configured using the ModelHandler. Args: model_handler: An implementation of ModelHandler. clock: A clock implementing time_ns. *Used for unit testing.* inference_args: Extra arguments for models whose inference call requires extra parameters. metrics_namespace: Namespace of the transform to collect metrics. model_metadata_pcoll: PCollection that emits Singleton ModelMetadata containing model path and model name, that is used as a side input to the _RunInferenceDoFn. watch_model_pattern: A glob pattern used to watch a directory for automatic model refresh. model_identifier: A string used to identify the model being loaded. You can set this if you want to reuse the same model across multiple RunInference steps and don't want to reload it twice. Note that using the same tag for different models will lead to non-deterministic results, so exercise caution when using this parameter. This only impacts models which are already being shared across processes. """ self._model_handler = model_handler self._inference_args = inference_args self._clock = clock self._metrics_namespace = metrics_namespace self._model_metadata_pcoll = model_metadata_pcoll self._with_exception_handling = False self._exception_handling_timeout = None self._timeout = None self._watch_model_pattern = watch_model_pattern self._use_model_manager = use_model_manager self._model_manager_args = model_manager_args self._kwargs = kwargs # Generate a random tag to use for shared.py and multi_process_shared.py to # allow us to effectively disambiguate in multi-model settings. Only use # the same tag if the model being loaded across multiple steps is actually # the same. self._model_tag = model_identifier if model_identifier is None: self._model_tag = uuid.uuid4().hex def annotations(self): return { 'model_handler': str(self._model_handler), 'model_handler_type': ( f'{self._model_handler.__class__.__module__}' f'.{self._model_handler.__class__.__qualname__}'), **super().annotations() } def _get_model_metadata_pcoll(self, pipeline): # avoid circular imports. # pylint: disable=wrong-import-position from apache_beam.ml.inference.utils import WatchFilePattern extra_params = {} if 'interval' in self._kwargs: extra_params['interval'] = self._kwargs['interval'] if 'stop_timestamp' in self._kwargs: extra_params['stop_timestamp'] = self._kwargs['stop_timestamp'] return ( pipeline | WatchFilePattern( file_pattern=self._watch_model_pattern, **extra_params)) # TODO(BEAM-14046): Add and link to help documentation. @classmethod def from_callable(cls, model_handler_provider, **kwargs): """Multi-language friendly constructor. Use this constructor with fully_qualified_named_transform to initialize the RunInference transform from PythonCallableSource provided by foreign SDKs. Args: model_handler_provider: A callable object that returns ModelHandler. kwargs: Keyword arguments for model_handler_provider. """ return cls(model_handler_provider(**kwargs)) def _apply_fns( self, pcoll: beam.PCollection, fns: Iterable[Callable[[Any], Any]], step_prefix: str) -> tuple[beam.PCollection, Iterable[beam.PCollection]]: bad_preprocessed = [] for idx in range(len(fns)): fn = fns[idx] if self._with_exception_handling: pcoll, bad = (pcoll | f"{step_prefix}-{idx}" >> beam.Map( fn).with_exception_handling( exc_class=self._exc_class, use_subprocess=self._use_subprocess, threshold=self._threshold, timeout = self._timeout)) bad_preprocessed.append(bad) else: pcoll = pcoll | f"{step_prefix}-{idx}" >> beam.Map(fn) return pcoll, bad_preprocessed # TODO(https://github.com/apache/beam/issues/21447): Add batch_size back off # in the case there are functional reasons large batch sizes cannot be # handled. def expand( self, pcoll: beam.PCollection[ExampleT]) -> beam.PCollection[PredictionT]: self._model_handler.validate_inference_args(self._inference_args) # DLQ pcollections bad_preprocessed = [] bad_inference = None bad_postprocessed = [] preprocess_fns = self._model_handler.get_preprocess_fns() postprocess_fns = self._model_handler.get_postprocess_fns() pcoll, bad_preprocessed = self._apply_fns( pcoll, preprocess_fns, 'BeamML_RunInference_Preprocess') resource_hints = self._model_handler.get_resource_hints() # check for the side input if self._watch_model_pattern: self._model_metadata_pcoll = self._get_model_metadata_pcoll( pcoll.pipeline) if self._model_handler.should_skip_batching(): batched_elements_pcoll = pcoll else: batched_elements_pcoll = ( pcoll # TODO(https://github.com/apache/beam/issues/21440): Hook into the # batching DoFn APIs. | beam.BatchElements(**self._model_handler.batch_elements_kwargs())) # Skip loading in setup if we are dependent on side inputs or we want to # enforce a timeout since neither of these are available in a helpful way # in setup. load_model_at_runtime = ( self._model_metadata_pcoll is not None or self._timeout is not None) run_inference_pardo = beam.ParDo( _RunInferenceDoFn( self._model_handler, self._clock, self._metrics_namespace, load_model_at_runtime, self._model_tag, self._use_model_manager, self._model_manager_args), self._inference_args, beam.pvalue.AsSingleton( self._model_metadata_pcoll, ) if self._model_metadata_pcoll else None).with_resource_hints( **resource_hints) if self._with_exception_handling: # On timeouts, report back to the central model metadata # that the model is invalid model_tag = self._model_tag share_across_processes = self._model_handler.share_model_across_processes( ) timeout = self._timeout def failure_callback(exception: Exception, element: Any): if type(exception) is not TimeoutError: return model_metadata = load_model_status(model_tag, share_across_processes) model_metadata.try_mark_current_model_invalid(timeout) logging.warning( "Inference failed with a timeout, marking the current " + "model for garbage collection") callback = None if (self._timeout is not None and self._model_handler.should_garbage_collect_on_timeout()): callback = failure_callback results, bad_inference = ( batched_elements_pcoll | 'BeamML_RunInference' >> run_inference_pardo.with_exception_handling( exc_class=self._exc_class, use_subprocess=self._use_subprocess, threshold=self._threshold, timeout = self._timeout, on_failure_callback=callback)) else: results = ( batched_elements_pcoll | 'BeamML_RunInference' >> run_inference_pardo) results, bad_postprocessed = self._apply_fns( results, postprocess_fns, 'BeamML_RunInference_Postprocess') if self._with_exception_handling: dlq = RunInferenceDLQ(bad_inference, bad_preprocessed, bad_postprocessed) return results, dlq return results def with_exception_handling( self, *, exc_class=Exception, use_subprocess=False, threshold=1, timeout: Optional[int] = None): """Automatically provides a dead letter output for skipping bad records. This can allow a pipeline to continue successfully rather than fail or continuously throw errors on retry when bad elements are encountered. This returns a tagged output with two PCollections, the first being the results of successfully processing the input PCollection, and the second being the set of bad batches of records (those which threw exceptions during processing) along with information about the errors raised. For example, one would write:: main, other = RunInference( maybe_error_raising_model_handler ).with_exception_handling() and `main` will be a PCollection of PredictionResults and `other` will contain a `RunInferenceDLQ` object with PCollections containing failed records for each failed inference, preprocess operation, or postprocess operation. To access each collection of failed records, one would write: failed_inferences = other.failed_inferences failed_preprocessing = other.failed_preprocessing failed_postprocessing = other.failed_postprocessing failed_inferences is in the form PCollection[tuple[failed batch, exception]]. failed_preprocessing is in the form list[PCollection[tuple[failed record, exception]]]], where each element of the list corresponds to a preprocess function. These PCollections are in the same order that the preprocess functions are applied. failed_postprocessing is in the form list[PCollection[tuple[failed record, exception]]]], where each element of the list corresponds to a postprocess function. These PCollections are in the same order that the postprocess functions are applied. Args: exc_class: An exception class, or tuple of exception classes, to catch. Optional, defaults to 'Exception'. use_subprocess: Whether to execute the DoFn logic in a subprocess. This allows one to recover from errors that can crash the calling process (e.g. from an underlying library causing a segfault), but is slower as elements and results must cross a process boundary. Note that this starts up a long-running process that is used to handle all the elements (until hard failure, which should be rare) rather than a new process per element, so the overhead should be minimal (and can be amortized if there's any per-process or per-bundle initialization that needs to be done). Optional, defaults to False. threshold: An upper bound on the ratio of records that can be bad before aborting the entire pipeline. Optional, defaults to 1.0 (meaning up to 100% of records can be bad and the pipeline will still succeed). timeout: The maximum amount of time in seconds given to load a model, run inference on a batch of elements and perform and pre/postprocessing operations. Since the timeout applies in multiple places, it should be equal to the maximum possible timeout for any of these operations. Note in particular that model load and inference on a single batch count to the same timeout value. When an inference fails, all related resources, including the model, will be deleted and reloaded. As a result, it is recommended to leave significant buffer and set the timeout to at least `2 * (time to load model + time to run inference on a batch of data)`. """ self._with_exception_handling = True self._exc_class = exc_class self._use_subprocess = use_subprocess self._threshold = threshold self._timeout = timeout return self class _MetricsCollector: """ A metrics collector that tracks ML related performance and memory usage. """ def __init__(self, namespace: str, prefix: str = ''): """ Args: namespace: Namespace for the metrics. prefix: Unique identifier for metrics, used when models are updated using side input. """ # Metrics if prefix: prefix = f'{prefix}_' self._inference_counter = beam.metrics.Metrics.counter( namespace, prefix + 'num_inferences') self.failed_batches_counter = beam.metrics.Metrics.counter( namespace, prefix + 'failed_batches_counter') self._inference_request_batch_size = beam.metrics.Metrics.distribution( namespace, prefix + 'inference_request_batch_size') self._inference_request_batch_byte_size = ( beam.metrics.Metrics.distribution( namespace, prefix + 'inference_request_batch_byte_size')) # Batch inference latency in microseconds. self._inference_batch_latency_micro_secs = ( beam.metrics.Metrics.distribution( namespace, prefix + 'inference_batch_latency_micro_secs')) self._model_byte_size = beam.metrics.Metrics.distribution( namespace, prefix + 'model_byte_size') # Model load latency in milliseconds. self._load_model_latency_milli_secs = beam.metrics.Metrics.distribution( namespace, prefix + 'load_model_latency_milli_secs') # Metrics cache self._load_model_latency_milli_secs_cache = None self._model_byte_size_cache = None def update_metrics_with_cache(self): if self._load_model_latency_milli_secs_cache is not None: self._load_model_latency_milli_secs.update( self._load_model_latency_milli_secs_cache) self._load_model_latency_milli_secs_cache = None if self._model_byte_size_cache is not None: self._model_byte_size.update(self._model_byte_size_cache) self._model_byte_size_cache = None def cache_load_model_metrics(self, load_model_latency_ms, model_byte_size): self._load_model_latency_milli_secs_cache = load_model_latency_ms self._model_byte_size_cache = model_byte_size def update_load_model_metrics(self, load_model_latency_ms, model_byte_size): self._load_model_latency_milli_secs.update(load_model_latency_ms) self._model_byte_size.update(model_byte_size) def update( self, examples_count: int, examples_byte_size: int, latency_micro_secs: int): self._inference_batch_latency_micro_secs.update(latency_micro_secs) self._inference_counter.inc(examples_count) self._inference_request_batch_size.update(examples_count) self._inference_request_batch_byte_size.update(examples_byte_size) class _ModelRoutingStrategy(): """A class meant to sit in a shared location for mapping incoming batches to different models. Currently only supports round-robin, but can be extended to support other protocols if needed. """ def __init__(self): self._cur_index = 0 def next_model_index(self, num_models): self._cur_index = (self._cur_index + 1) % num_models return self._cur_index class _ModelStatus(): """A class holding any metadata about a model required by RunInference. Currently, this only includes whether or not the model is valid. Uses the model tag to map models to metadata. """ def __init__(self, share_model_across_processes: bool): self._active_tags = set() self._invalid_tags = set() self._tag_mapping = {} self._model_first_seen = {} self._pending_hard_delete = [] self._share_model_across_process = share_model_across_processes def try_mark_current_model_invalid(self, min_model_life_seconds): """Mark the current model invalid. Since we don't have sufficient information to say which model is being marked invalid, but there may be multiple active models, we will mark all models currently in use as inactive so that they all get reloaded. To avoid thrashing, however, we will only mark models as invalid if they've been active at least min_model_life_seconds seconds. """ cutoff_time = datetime.now() - timedelta(seconds=min_model_life_seconds) for tag in list(self._active_tags): if cutoff_time >= self._model_first_seen[tag]: self._invalid_tags.add(tag) # Delete old models after a grace period of 2 * the model life. # This already happens automatically for shared.Shared models, so # cleanup is only necessary for multi_process_shared models. if self._share_model_across_process: self._pending_hard_delete.append(( tag, datetime.now() + 2 * timedelta(seconds=min_model_life_seconds))) self._active_tags.remove(tag) def get_valid_tag(self, tag: str) -> str: """Takes in a proposed valid tag and returns a valid one. Will always return a valid tag. If the passed in tag is valid, this function will simply return it, otherwise it will deterministically generate a new tag to use instead. The new tag will be the original tag with an incrementing suffix (e.g. `my_tag_reload_1`, `my_tag_reload_2`) for each reload """ if tag not in self._invalid_tags: if tag not in self._model_first_seen: self._model_first_seen[tag] = datetime.now() self._active_tags.add(tag) return tag if (tag in self._tag_mapping and self._tag_mapping[tag] not in self._invalid_tags): return self._tag_mapping[tag] i = 1 new_tag = f'{tag}_reload_{i}' while new_tag in self._invalid_tags: i += 1 new_tag = f'{tag}_reload_{i}' self._tag_mapping[tag] = new_tag self._model_first_seen[new_tag] = datetime.now() self._active_tags.add(new_tag) return new_tag def is_valid_tag(self, tag: str) -> bool: return tag == self.get_valid_tag(tag) def get_tags_for_garbage_collection(self) -> list[str]: # Since this function may be in multi_process_shared space, delegate model # deletion to the calling process which is not to avoid having a # multi_process_shared reference in multi_process_shared space, which # can create issues with python's multiprocessing module. # We will rely on the calling process to report back deleted models so that # we can confirm deletion. to_delete = [] cur_time = datetime.now() for i in range(len(self._pending_hard_delete)): delete_time = self._pending_hard_delete[i][1] tag = self._pending_hard_delete[i][0] if delete_time <= cur_time: to_delete.append(tag) else: # early return once we hit a model which was added later since models # are added in order. return to_delete return to_delete def mark_tags_deleted(self, deleted_tags: set[str]): while len(self._pending_hard_delete) > 0: tag = self._pending_hard_delete[0][0] if tag in deleted_tags: self._pending_hard_delete.pop(0) else: return def load_model_status( model_tag: str, share_across_processes: bool) -> _ModelStatus: tag = f'{model_tag}_model_status' if share_across_processes: return multi_process_shared.MultiProcessShared( lambda: _ModelStatus(True), tag=tag, always_proxy=True).acquire() return shared.Shared().acquire(lambda: _ModelStatus(False), tag=tag) class _ProxyLoader: """ A helper callable to wrap the loader for MultiProcessShared. """ def __init__(self, loader_func, model_tag): self.loader_func = loader_func self.model_tag = model_tag def __call__(self): # Generate a unique tag for the model being loaded so that # we will have unique instances of the model in multi_process_shared # space instead of reusing the same instance over. The instance will # be initialized and left running as a separate process, which then # can be grabbed again using the unique tag if needed during inference. unique_tag = self.model_tag + '_' + uuid.uuid4().hex # Ensure that each model loaded in a different process for parallelism multi_process_shared.MultiProcessShared( OOMProtectedFn(self.loader_func), tag=unique_tag, always_proxy=True, spawn_process=True).acquire() # Only return the tag to avoid pickling issues with the model itself. return unique_tag class _SharedModelWrapper(): """A router class to map incoming calls to the correct model. This allows us to round robin calls to models sitting in different processes so that we can more efficiently use resources (e.g. GPUs). """ def __init__( self, models: Union[list[Any], ModelManager], model_tag: str, loader_func: Optional[Callable[[], Any]] = None): self.models = models self.use_model_manager = not isinstance(models, list) self.model_tag = model_tag self.loader_func = loader_func if not self.use_model_manager and len(models) > 1: self.model_router = multi_process_shared.MultiProcessShared( lambda: _ModelRoutingStrategy(), tag=f'{model_tag}_counter', always_proxy=True).acquire() def next_model(self): if self.use_model_manager: loader_wrapper = _ProxyLoader(self.loader_func, self.model_tag) return self.models.acquire_model(self.model_tag, loader_wrapper) if len(self.models) == 1: # Short circuit if there's no routing strategy needed in order to # avoid the cross-process call return self.models[0] return self.models[self.model_router.next_model_index(len(self.models))] def release_model(self, model_tag: str, model: Any): if self.use_model_manager: self.models.release_model(model_tag, model) def all_models(self): if self.use_model_manager: return self.models.all_models()[self.model_tag] return self.models def force_reset(self): if self.use_model_manager: self.models.force_reset() class _RunInferenceDoFn(beam.DoFn, Generic[ExampleT, PredictionT]): def __init__( self, model_handler: ModelHandler[ExampleT, PredictionT, Any], clock, metrics_namespace, load_model_at_runtime: bool = False, model_tag: str = "RunInference", use_model_manager: bool = False, model_manager_args: Optional[dict[str, Any]] = None): """A DoFn implementation generic to frameworks. Args: model_handler: An implementation of ModelHandler. clock: A clock implementing time_ns. *Used for unit testing.* metrics_namespace: Namespace of the transform to collect metrics. load_model_at_runtime: Bool to indicate if model loading should be deferred to runtime - for example if we are depending on side inputs to get the model path or we want to enforce a timeout on model loading. model_tag: Tag to use to disambiguate models in multi-model settings. """ self._model_handler = model_handler self._shared_model_handle = shared.Shared() self._clock = clock self._model = None self._metrics_namespace = metrics_namespace self._load_model_at_runtime = load_model_at_runtime self._side_input_path = None # _model_tag is the original tag passed in. # _cur_tag is the tag of the actually loaded model self._model_tag = model_tag self._cur_tag = model_tag self.use_model_manager = use_model_manager self._model_manager_args = model_manager_args or {} def _load_model( self, side_input_model_path: Optional[Union[str, list[KeyModelPathMapping]]] = None ) -> _SharedModelWrapper: def load(): """Function for constructing shared LoadedModel.""" memory_before = _get_current_process_memory_in_bytes() start_time = _to_milliseconds(self._clock.time_ns()) if isinstance(side_input_model_path, str): self._model_handler.update_model_path(side_input_model_path) else: if self._model is not None: models = self._model.all_models() for m in models: self._model_handler.update_model_paths(m, side_input_model_path) model = self._model_handler.load_model() end_time = _to_milliseconds(self._clock.time_ns()) memory_after = _get_current_process_memory_in_bytes() load_model_latency_ms = end_time - start_time model_byte_size = memory_after - memory_before if self._metrics_collector: self._metrics_collector.cache_load_model_metrics( load_model_latency_ms, model_byte_size) return model # TODO(https://github.com/apache/beam/issues/21443): Investigate releasing # model. model_tag = self._model_tag if isinstance(side_input_model_path, str) and side_input_model_path != '': model_tag = side_input_model_path # Ensure the tag we're loading is valid, if not replace it with a valid tag self._cur_tag = self._model_metadata.get_valid_tag(model_tag) if self.use_model_manager: logging.info("Using Model Manager to manage models automatically.") model_manager = multi_process_shared.MultiProcessShared( lambda: ModelManager(**self._model_manager_args), tag='model_manager', always_proxy=True).acquire() model_wrapper = _SharedModelWrapper( model_manager, self._cur_tag, self._model_handler.load_model) elif self._model_handler.share_model_across_processes(): models = [] for copy_tag in _get_tags_for_copies(self._cur_tag, self._model_handler.model_copies()): models.append( multi_process_shared.MultiProcessShared( load, tag=copy_tag, always_proxy=True).acquire()) model_wrapper = _SharedModelWrapper(models, self._cur_tag) else: model = self._shared_model_handle.acquire(load, tag=self._cur_tag) model_wrapper = _SharedModelWrapper([model], self._cur_tag) # since shared_model_handle is shared across threads, the model path # might not get updated in the model handler # because we directly get cached weak ref model from shared cache, instead # of calling load(). For sanity check, call update_model_path again. if isinstance(side_input_model_path, str): self._model_handler.update_model_path(side_input_model_path) else: if self._model is not None: models = self._model.all_models() for m in models: self._model_handler.update_model_paths(m, side_input_model_path) return model_wrapper def get_metrics_collector(self, prefix: str = ''): """ Args: prefix: Unique identifier for metrics, used when models are updated using side input. """ metrics_namespace = ( self._metrics_namespace) if self._metrics_namespace else ( self._model_handler.get_metrics_namespace()) if self._model_handler.override_metrics(metrics_namespace): return None return _MetricsCollector(metrics_namespace, prefix=prefix) def setup(self): self._metrics_collector = self.get_metrics_collector() self._model_handler.set_environment_vars() self._model_metadata = load_model_status( self._model_tag, self._model_handler.share_model_across_processes()) if not self._load_model_at_runtime: self._model = self._load_model() def update_model( self, side_input_model_path: Optional[Union[str, list[KeyModelPathMapping]]] = None): self._model = self._load_model(side_input_model_path=side_input_model_path) def _run_inference(self, batch, inference_args): start_time = _to_microseconds(self._clock.time_ns()) try: model = self._model.next_model() if isinstance(model, str): # ModelManager with MultiProcessShared returns the model tag unique_tag = model model = multi_process_shared.MultiProcessShared( lambda: None, tag=model, always_proxy=True).acquire() try: result_generator = (OOMProtectedFn(self._model_handler.run_inference))( batch, model, inference_args) finally: # Always release the model so that it can be reloaded. if self.use_model_manager: self._model.release_model(self._model_tag, unique_tag) except BaseException as e: if self._metrics_collector: self._metrics_collector.failed_batches_counter.inc() if (e is pickle.PickleError and self._model_handler.share_model_across_processes()): raise TypeError( 'Pickling error encountered while running inference. ' 'This may be caused by trying to send unpickleable ' 'data to a model which is shared across processes. ' 'For more information, see ' 'https://beam.apache.org/documentation/ml/large-language-modeling/#pickling-errors' # pylint: disable=line-too-long ) from e raise e predictions = list(result_generator) end_time = _to_microseconds(self._clock.time_ns()) inference_latency = end_time - start_time num_bytes = self._model_handler.get_num_bytes(batch) num_elements = len(batch) if self._metrics_collector: self._metrics_collector.update(num_elements, num_bytes, inference_latency) return predictions def process( self, batch, inference_args, si_model_metadata: Optional[Union[ModelMetadata, list[ModelMetadata], list[KeyModelPathMapping]]]): """ When side input is enabled: The method checks if the side input model has been updated, and if so, updates the model and runs inference on the batch of data. If the side input is empty or the model has not been updated, the method simply runs inference on the batch of data. """ if not si_model_metadata: if (not self._model_metadata.is_valid_tag(self._cur_tag) or self._model is None): self.update_model(side_input_model_path=None) return self._run_inference(batch, inference_args) if isinstance(si_model_metadata, beam.pvalue.EmptySideInput): self.update_model(side_input_model_path=None) elif isinstance(si_model_metadata, list) and hasattr(si_model_metadata[0], 'keys'): # TODO(https://github.com/apache/beam/issues/27628): Update metrics here self.update_model(si_model_metadata) elif self._side_input_path != si_model_metadata.model_id: self._side_input_path = si_model_metadata.model_id self._metrics_collector = self.get_metrics_collector( prefix=si_model_metadata.model_name) lock = threading.Lock() with lock: self.update_model(si_model_metadata.model_id) return self._run_inference(batch, inference_args) return self._run_inference(batch, inference_args) def finish_bundle(self): # TODO(https://github.com/apache/beam/issues/21435): Figure out why there # is a cache. if self._metrics_collector: self._metrics_collector.update_metrics_with_cache() # Do garbage collection of old models tags_to_gc = self._model_metadata.get_tags_for_garbage_collection() if len(tags_to_gc) > 0: for unprefixed_tag in tags_to_gc: for tag in _get_tags_for_copies(unprefixed_tag, self._model_handler.model_copies()): multi_process_shared.MultiProcessShared(lambda: None, tag).unsafe_hard_delete() self._model_metadata.mark_tags_deleted(tags_to_gc) def _is_darwin() -> bool: return sys.platform == 'darwin' def _get_current_process_memory_in_bytes(): """ Returns: memory usage in bytes. """ if resource is not None: usage = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss if _is_darwin(): return usage return usage * 1024 else: logging.warning( 'Resource module is not available for current platform, ' 'memory usage cannot be fetched.') return 0 def _get_tags_for_copies(base_tag, num_copies): tags = [] for i in range(num_copies): tags.append(f'{base_tag}{i}') return tags ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/base_test.py0000644000175100017510000026323515206113032022561 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.ml.base.""" import math import multiprocessing import os import pickle import random import sys import tempfile import time import unittest import unittest.mock from collections.abc import Iterable from collections.abc import Mapping from collections.abc import Sequence from typing import Any from typing import Optional from typing import Union import pytest import apache_beam as beam from apache_beam.examples.inference import run_inference_side_inputs from apache_beam.metrics.metric import MetricsFilter from apache_beam.ml.inference import base from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.transforms.periodicsequence import TimestampedValue from apache_beam.utils import multi_process_shared class FakeModel: def predict(self, example: int) -> int: return example + 1 class FakeStatefulModel: def __init__(self, state: int): if state == 100: raise Exception('Oh no') self._state = state def predict(self, example: int) -> int: return self._state def increment_state(self, amount: int): self._state += amount class FakeSlowModel: def __init__(self, sleep_on_load_seconds=0, file_path_write_on_del=None): self._file_path_write_on_del = file_path_write_on_del time.sleep(sleep_on_load_seconds) def predict(self, example: int) -> int: time.sleep(example) return example def __del__(self): if self._file_path_write_on_del is not None: with open(self._file_path_write_on_del, 'a') as myfile: myfile.write('Deleted FakeSlowModel') class FakeIncrementingModel: def __init__(self): self._state = 0 def predict(self, example: int) -> int: self._state += 1 return self._state class FakeSlowModelHandler(base.ModelHandler[int, int, FakeModel]): def __init__( self, sleep_on_load: int, multi_process_shared=False, file_path_write_on_del=None): self._sleep_on_load = sleep_on_load self._multi_process_shared = multi_process_shared self._file_path_write_on_del = file_path_write_on_del def load_model(self): return FakeSlowModel(self._sleep_on_load, self._file_path_write_on_del) def run_inference( self, batch: Sequence[int], model: FakeModel, inference_args=None) -> Iterable[int]: for example in batch: yield model.predict(example) def share_model_across_processes(self) -> bool: return self._multi_process_shared def batch_elements_kwargs(self): return {'min_batch_size': 1, 'max_batch_size': 1} class FakeModelHandler(base.ModelHandler[int, int, FakeModel]): def __init__( self, clock=None, min_batch_size=1, max_batch_size=9999, multi_process_shared=False, state=None, incrementing=False, max_copies=1, num_bytes_per_element=None, **kwargs): self._fake_clock = clock self._min_batch_size = min_batch_size self._max_batch_size = max_batch_size self._env_vars = kwargs.get('env_vars', {}) self._multi_process_shared = multi_process_shared self._state = state self._incrementing = incrementing self._max_copies = max_copies self._num_bytes_per_element = num_bytes_per_element def load_model(self): assert (not self._incrementing or self._state is None) if self._fake_clock: self._fake_clock.current_time_ns += 500_000_000 # 500ms if self._incrementing: return FakeIncrementingModel() if self._state is not None: return FakeStatefulModel(self._state) return FakeModel() def run_inference( self, batch: Sequence[int], model: FakeModel, inference_args=None) -> Iterable[int]: multi_process_shared_loaded = "multi_process_shared" in str(type(model)) if self._multi_process_shared != multi_process_shared_loaded: raise Exception( f'Loaded model of type {type(model)}, was' + f'{"" if self._multi_process_shared else " not"} ' + 'expecting multi_process_shared_model') if self._fake_clock: self._fake_clock.current_time_ns += 3_000_000 # 3 milliseconds for example in batch: yield model.predict(example) def update_model_path(self, model_path: Optional[str] = None): pass def batch_elements_kwargs(self): return { 'min_batch_size': self._min_batch_size, 'max_batch_size': self._max_batch_size } def share_model_across_processes(self): return self._multi_process_shared def model_copies(self): return self._max_copies def get_num_bytes(self, batch: Sequence[int]) -> int: if self._num_bytes_per_element: return self._num_bytes_per_element * len(batch) return super().get_num_bytes(batch) class FakeModelHandlerReturnsPredictionResult( base.ModelHandler[int, base.PredictionResult, FakeModel]): def __init__( self, clock=None, model_id='fake_model_id_default', multi_process_shared=False, state=None): self.model_id = model_id self._fake_clock = clock self._env_vars = {} self._multi_process_shared = multi_process_shared self._state = state def load_model(self): if self._state is not None: return FakeStatefulModel(0) return FakeModel() def run_inference( self, batch: Sequence[int], model: Union[FakeModel, FakeStatefulModel], inference_args=None) -> Iterable[base.PredictionResult]: multi_process_shared_loaded = "multi_process_shared" in str(type(model)) if self._multi_process_shared != multi_process_shared_loaded: raise Exception( f'Loaded model of type {type(model)}, was' + f'{"" if self._multi_process_shared else " not"} ' + 'expecting multi_process_shared_model') for example in batch: yield base.PredictionResult( model_id=self.model_id, example=example, inference=model.predict(example)) if self._state is not None: model.increment_state(1) # type: ignore[union-attr] def update_model_path(self, model_path: Optional[str] = None): self.model_id = model_path if model_path else self.model_id def share_model_across_processes(self): return self._multi_process_shared class FakeModelHandlerNoEnvVars(base.ModelHandler[int, int, FakeModel]): def __init__( self, clock=None, min_batch_size=1, max_batch_size=9999, **kwargs): self._fake_clock = clock self._min_batch_size = min_batch_size self._max_batch_size = max_batch_size def load_model(self): if self._fake_clock: self._fake_clock.current_time_ns += 500_000_000 # 500ms return FakeModel() def run_inference( self, batch: Sequence[int], model: FakeModel, inference_args=None) -> Iterable[int]: if self._fake_clock: self._fake_clock.current_time_ns += 3_000_000 # 3 milliseconds for example in batch: yield model.predict(example) def update_model_path(self, model_path: Optional[str] = None): pass def batch_elements_kwargs(self): return { 'min_batch_size': self._min_batch_size, 'max_batch_size': self._max_batch_size } class FakeClock: def __init__(self): # Start at 10 seconds. self.current_time_ns = 10_000_000_000 def time_ns(self) -> int: return self.current_time_ns class ExtractInferences(beam.DoFn): def process(self, prediction_result): yield prediction_result.inference class FakeModelHandlerNeedsBigBatch(FakeModelHandler): def run_inference(self, batch, unused_model, inference_args=None): if len(batch) < 100: raise ValueError('Unexpectedly small batch') return batch def batch_elements_kwargs(self): return {'min_batch_size': 9999} class FakeModelHandlerFailsOnInferenceArgs(FakeModelHandler): def run_inference(self, batch, unused_model, inference_args=None): raise ValueError( 'run_inference should not be called because error should already be ' 'thrown from the validate_inference_args check.') def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): if inference_args: raise ValueError( 'inference_args were provided, but should be None because this ' 'framework does not expect extra arguments on inferences.') class FakeModelHandlerExpectedInferenceArgs(FakeModelHandler): def run_inference(self, batch, unused_model, inference_args=None): if not inference_args: raise ValueError('inference_args should exist') return batch def validate_inference_args(self, inference_args): pass class RunInferenceBaseTest(unittest.TestCase): def test_run_inference_impl_simple_examples(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] expected = [example + 1 for example in examples] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference(FakeModelHandler()) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_impl_simple_examples_multi_process_shared(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] expected = [example + 1 for example in examples] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference( FakeModelHandler(multi_process_shared=True)) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_impl_simple_examples_multi_process_shared_multi_copy( self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] expected = [example + 1 for example in examples] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference( FakeModelHandler(multi_process_shared=True, max_copies=4)) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_impl_multi_process_shared_incrementing_multi_copy( self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10, 1, 5, 3, 10, 1, 5, 3, 10, 1, 5, 3, 10] expected = [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference( FakeModelHandler( multi_process_shared=True, max_copies=4, incrementing=True, max_batch_size=1)) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_impl_mps_nobatch_incrementing_multi_copy(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10, 1, 5, 3, 10, 1, 5, 3, 10, 1, 5, 3, 10] expected = [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4] batched_examples = [[example] for example in examples] pcoll = pipeline | 'start' >> beam.Create(batched_examples) actual = pcoll | base.RunInference( FakeModelHandler( multi_process_shared=True, max_copies=4, incrementing=True).with_no_batching()) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_impl_keyed_mps_incrementing_multi_copy(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10, 1, 5, 3, 10, 1, 5, 3, 10, 1, 5, 3, 10] keyed_examples = [('abc', example) for example in examples] expected = [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4] keyed_expected = [('abc', val) for val in expected] pcoll = pipeline | 'start' >> beam.Create(keyed_examples) actual = pcoll | base.RunInference( base.KeyedModelHandler( FakeModelHandler( multi_process_shared=True, max_copies=4, incrementing=True, max_batch_size=1))) assert_that(actual, equal_to(keyed_expected), label='assert:inferences') def test_run_inference_impl_with_keyed_examples(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] keyed_examples = [(i, example) for i, example in enumerate(examples)] expected = [(i, example + 1) for i, example in enumerate(examples)] pcoll = pipeline | 'start' >> beam.Create(keyed_examples) actual = pcoll | base.RunInference( base.KeyedModelHandler(FakeModelHandler())) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_impl_with_keyed_examples_many_model_handlers(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] keyed_examples = [(i, example) for i, example in enumerate(examples)] expected = [(i, example + 1) for i, example in enumerate(examples)] expected[0] = (0, 200) pcoll = pipeline | 'start' >> beam.Create(keyed_examples) mhs = [ base.KeyModelMapping([0], FakeModelHandler( state=200, multi_process_shared=True)), base.KeyModelMapping([1, 2, 3], FakeModelHandler(multi_process_shared=True)) ] actual = pcoll | base.RunInference(base.KeyedModelHandler(mhs)) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_impl_with_keyed_examples_many_model_handlers_metrics( self): pipeline = TestPipeline() examples = [1, 5, 3, 10] metrics_namespace = 'test_namespace' keyed_examples = [(i, example) for i, example in enumerate(examples)] pcoll = pipeline | 'start' >> beam.Create(keyed_examples) mhs = [ base.KeyModelMapping([0], FakeModelHandler( state=200, multi_process_shared=True)), base.KeyModelMapping([1, 2, 3], FakeModelHandler(multi_process_shared=True)) ] _ = pcoll | base.RunInference( base.KeyedModelHandler(mhs), metrics_namespace=metrics_namespace) result = pipeline.run() result.wait_until_finish() metrics_filter = MetricsFilter().with_namespace(namespace=metrics_namespace) metrics = result.metrics().query(metrics_filter) assert len(metrics['counters']) != 0 assert len(metrics['distributions']) != 0 metrics_filter = MetricsFilter().with_name('0-_num_inferences') metrics = result.metrics().query(metrics_filter) num_inferences_counter_key_0 = metrics['counters'][0] self.assertEqual(num_inferences_counter_key_0.committed, 1) metrics_filter = MetricsFilter().with_name('1-_num_inferences') metrics = result.metrics().query(metrics_filter) num_inferences_counter_key_1 = metrics['counters'][0] self.assertEqual(num_inferences_counter_key_1.committed, 3) metrics_filter = MetricsFilter().with_name('num_inferences') metrics = result.metrics().query(metrics_filter) num_inferences_counter_aggregate = metrics['counters'][0] self.assertEqual(num_inferences_counter_aggregate.committed, 4) metrics_filter = MetricsFilter().with_name('0-_failed_batches_counter') metrics = result.metrics().query(metrics_filter) failed_batches_counter_key_0 = metrics['counters'] self.assertEqual(len(failed_batches_counter_key_0), 0) metrics_filter = MetricsFilter().with_name('failed_batches_counter') metrics = result.metrics().query(metrics_filter) failed_batches_counter_aggregate = metrics['counters'] self.assertEqual(len(failed_batches_counter_aggregate), 0) metrics_filter = MetricsFilter().with_name( '0-_load_model_latency_milli_secs') metrics = result.metrics().query(metrics_filter) load_latency_dist_key_0 = metrics['distributions'][0] self.assertEqual(load_latency_dist_key_0.committed.count, 1) metrics_filter = MetricsFilter().with_name('load_model_latency_milli_secs') metrics = result.metrics().query(metrics_filter) load_latency_dist_aggregate = metrics['distributions'][0] self.assertEqual(load_latency_dist_aggregate.committed.count, 2) def test_run_inference_impl_with_keyed_examples_many_mhs_max_models_hint( self): pipeline = TestPipeline() examples = [1, 5, 3, 10, 2, 4, 6, 8, 9, 7, 1, 5, 3, 10, 2, 4, 6, 8, 9, 7] metrics_namespace = 'test_namespace' keyed_examples = [(i, example) for i, example in enumerate(examples)] pcoll = pipeline | 'start' >> beam.Create(keyed_examples) mhs = [ base.KeyModelMapping([0, 2, 4, 6, 8], FakeModelHandler( state=200, multi_process_shared=True)), base.KeyModelMapping( [1, 3, 5, 7, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19], FakeModelHandler(multi_process_shared=True)) ] _ = pcoll | base.RunInference( base.KeyedModelHandler(mhs, max_models_per_worker_hint=1), metrics_namespace=metrics_namespace) result = pipeline.run() result.wait_until_finish() metrics_filter = MetricsFilter().with_namespace(namespace=metrics_namespace) metrics = result.metrics().query(metrics_filter) assert len(metrics['counters']) != 0 assert len(metrics['distributions']) != 0 metrics_filter = MetricsFilter().with_name('load_model_latency_milli_secs') metrics = result.metrics().query(metrics_filter) load_latency_dist_aggregate = metrics['distributions'][0] # We should flip back and forth between models a bit since # max_models_per_worker_hint=1, but we shouldn't thrash forever # since most examples belong to the second ModelMapping self.assertGreater(load_latency_dist_aggregate.committed.count, 2) self.assertLess(load_latency_dist_aggregate.committed.count, 12) def test_keyed_many_model_handlers_validation(self): def mult_two(example: str) -> int: return int(example) * 2 mhs = [ base.KeyModelMapping( [0], FakeModelHandler( state=200, multi_process_shared=True).with_preprocess_fn(mult_two)), base.KeyModelMapping([1, 2, 3], FakeModelHandler(multi_process_shared=True)) ] with self.assertRaises(ValueError): base.KeyedModelHandler(mhs) mhs = [ base.KeyModelMapping( [0], FakeModelHandler( state=200, multi_process_shared=True).with_postprocess_fn(mult_two)), base.KeyModelMapping([1, 2, 3], FakeModelHandler(multi_process_shared=True)) ] with self.assertRaises(ValueError): base.KeyedModelHandler(mhs) mhs = [ base.KeyModelMapping([0], FakeModelHandler( state=200, multi_process_shared=True)), base.KeyModelMapping([0, 1, 2, 3], FakeModelHandler(multi_process_shared=True)) ] with self.assertRaises(ValueError): base.KeyedModelHandler(mhs) mhs = [ base.KeyModelMapping([], FakeModelHandler( state=200, multi_process_shared=True)), base.KeyModelMapping([0, 1, 2, 3], FakeModelHandler(multi_process_shared=True)) ] with self.assertRaises(ValueError): base.KeyedModelHandler(mhs) def test_keyed_model_handler_get_num_bytes(self): mh = base.KeyedModelHandler(FakeModelHandler(num_bytes_per_element=10)) batch = [('key1', 1), ('key2', 2), ('key1', 3)] expected = len(pickle.dumps(('key1', 'key2', 'key1'))) + 30 actual = mh.get_num_bytes(batch) self.assertEqual(expected, actual) def test_keyed_model_handler_multiple_models_get_num_bytes(self): mhs = [ base.KeyModelMapping(['key1'], FakeModelHandler(num_bytes_per_element=10)), base.KeyModelMapping(['key2'], FakeModelHandler(num_bytes_per_element=20)) ] mh = base.KeyedModelHandler(mhs) batch = [('key1', 1), ('key2', 2), ('key1', 3)] expected = len(pickle.dumps(('key1', 'key2', 'key1'))) + 40 actual = mh.get_num_bytes(batch) self.assertEqual(expected, actual) def test_run_inference_impl_with_maybe_keyed_examples(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] keyed_examples = [(i, example) for i, example in enumerate(examples)] expected = [example + 1 for example in examples] keyed_expected = [(i, example + 1) for i, example in enumerate(examples)] model_handler = base.MaybeKeyedModelHandler(FakeModelHandler()) pcoll = pipeline | 'Unkeyed' >> beam.Create(examples) actual = pcoll | 'RunUnkeyed' >> base.RunInference(model_handler) assert_that(actual, equal_to(expected), label='CheckUnkeyed') keyed_pcoll = pipeline | 'Keyed' >> beam.Create(keyed_examples) keyed_actual = keyed_pcoll | 'RunKeyed' >> base.RunInference( model_handler) assert_that(keyed_actual, equal_to(keyed_expected), label='CheckKeyed') def test_run_inference_impl_with_keyed_examples_multi_process_shared(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] keyed_examples = [(i, example) for i, example in enumerate(examples)] expected = [(i, example + 1) for i, example in enumerate(examples)] pcoll = pipeline | 'start' >> beam.Create(keyed_examples) actual = pcoll | base.RunInference( base.KeyedModelHandler(FakeModelHandler(multi_process_shared=True))) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_impl_with_maybe_keyed_examples_multi_process_shared( self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] keyed_examples = [(i, example) for i, example in enumerate(examples)] expected = [example + 1 for example in examples] keyed_expected = [(i, example + 1) for i, example in enumerate(examples)] model_handler = base.MaybeKeyedModelHandler( FakeModelHandler(multi_process_shared=True)) pcoll = pipeline | 'Unkeyed' >> beam.Create(examples) actual = pcoll | 'RunUnkeyed' >> base.RunInference(model_handler) assert_that(actual, equal_to(expected), label='CheckUnkeyed') keyed_pcoll = pipeline | 'Keyed' >> beam.Create(keyed_examples) keyed_actual = keyed_pcoll | 'RunKeyed' >> base.RunInference( model_handler) assert_that(keyed_actual, equal_to(keyed_expected), label='CheckKeyed') def test_run_inference_preprocessing(self): def mult_two(example: str) -> int: return int(example) * 2 with TestPipeline() as pipeline: examples = ["1", "5", "3", "10"] expected = [int(example) * 2 + 1 for example in examples] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference( FakeModelHandler().with_preprocess_fn(mult_two)) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_prebatched(self): with TestPipeline() as pipeline: examples = [[1, 5], [3, 10]] expected = [int(example) + 1 for batch in examples for example in batch] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference(FakeModelHandler().with_no_batching()) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_preprocessing_multiple_fns(self): def add_one(example: str) -> int: return int(example) + 1 def mult_two(example: int) -> int: return example * 2 with TestPipeline() as pipeline: examples = ["1", "5", "3", "10"] expected = [(int(example) + 1) * 2 + 1 for example in examples] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference( FakeModelHandler().with_preprocess_fn(mult_two).with_preprocess_fn( add_one)) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_postprocessing(self): def mult_two(example: int) -> str: return str(example * 2) with TestPipeline() as pipeline: examples = [1, 5, 3, 10] expected = [str((example + 1) * 2) for example in examples] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference( FakeModelHandler().with_postprocess_fn(mult_two)) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_postprocessing_multiple_fns(self): def add_one(example: int) -> str: return str(int(example) + 1) def mult_two(example: int) -> int: return example * 2 with TestPipeline() as pipeline: examples = [1, 5, 3, 10] expected = [str(((example + 1) * 2) + 1) for example in examples] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference( FakeModelHandler().with_postprocess_fn(mult_two).with_postprocess_fn( add_one)) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_preprocessing_dlq(self): def mult_two(example: str) -> int: if example == "5": raise Exception("TEST") return int(example) * 2 with TestPipeline() as pipeline: examples = ["1", "5", "3", "10"] expected = [3, 7, 21] expected_bad = ["5"] pcoll = pipeline | 'start' >> beam.Create(examples) main, other = pcoll | base.RunInference( FakeModelHandler().with_preprocess_fn(mult_two) ).with_exception_handling() assert_that(main, equal_to(expected), label='assert:inferences') assert_that( other.failed_inferences, equal_to([]), label='assert:bad_infer') # bad will be in form [element, error]. Just pull out bad element. bad_without_error = other.failed_preprocessing[0] | beam.Map( lambda x: x[0]) assert_that( bad_without_error, equal_to(expected_bad), label='assert:failures') def test_run_inference_postprocessing_dlq(self): def mult_two(example: int) -> str: if example == 6: raise Exception("TEST") return str(example * 2) with TestPipeline() as pipeline: examples = [1, 5, 3, 10] expected = ["4", "8", "22"] expected_bad = [6] pcoll = pipeline | 'start' >> beam.Create(examples) main, other = pcoll | base.RunInference( FakeModelHandler().with_postprocess_fn(mult_two) ).with_exception_handling() assert_that(main, equal_to(expected), label='assert:inferences') assert_that( other.failed_inferences, equal_to([]), label='assert:bad_infer') # bad will be in form [element, error]. Just pull out bad element. bad_without_error = other.failed_postprocessing[0] | beam.Map( lambda x: x[0]) assert_that( bad_without_error, equal_to(expected_bad), label='assert:failures') def test_run_inference_pre_and_post_processing_dlq(self): def mult_two_pre(example: str) -> int: if example == "5": raise Exception("TEST") return int(example) * 2 def mult_two_post(example: int) -> str: if example == 7: raise Exception("TEST") return str(example * 2) with TestPipeline() as pipeline: examples = ["1", "5", "3", "10"] expected = ["6", "42"] expected_bad_pre = ["5"] expected_bad_post = [7] pcoll = pipeline | 'start' >> beam.Create(examples) main, other = pcoll | base.RunInference( FakeModelHandler().with_preprocess_fn( mult_two_pre ).with_postprocess_fn( mult_two_post )).with_exception_handling() assert_that(main, equal_to(expected), label='assert:inferences') assert_that( other.failed_inferences, equal_to([]), label='assert:bad_infer') # bad will be in form [elements, error]. Just pull out bad element. bad_without_error_pre = other.failed_preprocessing[0] | beam.Map( lambda x: x[0]) assert_that( bad_without_error_pre, equal_to(expected_bad_pre), label='assert:failures_pre') # bad will be in form [elements, error]. Just pull out bad element. bad_without_error_post = other.failed_postprocessing[0] | beam.Map( lambda x: x[0]) assert_that( bad_without_error_post, equal_to(expected_bad_post), label='assert:failures_post') def test_run_inference_keyed_pre_and_post_processing(self): def mult_two(element): return (element[0], element[1] * 2) with TestPipeline() as pipeline: examples = [1, 5, 3, 10] keyed_examples = [(i, example) for i, example in enumerate(examples)] expected = [(i, ((example * 2) + 1) * 2) for i, example in enumerate(examples)] pcoll = pipeline | 'start' >> beam.Create(keyed_examples) actual = pcoll | base.RunInference( base.KeyedModelHandler(FakeModelHandler()).with_preprocess_fn( mult_two).with_postprocess_fn(mult_two)) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_maybe_keyed_pre_and_post_processing(self): def mult_two(element): return element * 2 def mult_two_keyed(element): return (element[0], element[1] * 2) with TestPipeline() as pipeline: examples = [1, 5, 3, 10] keyed_examples = [(i, example) for i, example in enumerate(examples)] expected = [((2 * example) + 1) * 2 for example in examples] keyed_expected = [(i, ((2 * example) + 1) * 2) for i, example in enumerate(examples)] model_handler = base.MaybeKeyedModelHandler(FakeModelHandler()) pcoll = pipeline | 'Unkeyed' >> beam.Create(examples) actual = pcoll | 'RunUnkeyed' >> base.RunInference( model_handler.with_preprocess_fn(mult_two).with_postprocess_fn( mult_two)) assert_that(actual, equal_to(expected), label='CheckUnkeyed') keyed_pcoll = pipeline | 'Keyed' >> beam.Create(keyed_examples) keyed_actual = keyed_pcoll | 'RunKeyed' >> base.RunInference( model_handler.with_preprocess_fn(mult_two_keyed).with_postprocess_fn( mult_two_keyed)) assert_that(keyed_actual, equal_to(keyed_expected), label='CheckKeyed') def test_run_inference_impl_dlq(self): with TestPipeline() as pipeline: examples = [1, 'TEST', 3, 10, 'TEST2'] expected_good = [2, 4, 11] expected_bad = ['TEST', 'TEST2'] pcoll = pipeline | 'start' >> beam.Create(examples) main, other = pcoll | base.RunInference( FakeModelHandler( min_batch_size=1, max_batch_size=1 )).with_exception_handling() assert_that(main, equal_to(expected_good), label='assert:inferences') # bad.failed_inferences will be in form [batch[elements], error]. # Just pull out bad element. bad_without_error = other.failed_inferences | beam.Map(lambda x: x[0][0]) assert_that( bad_without_error, equal_to(expected_bad), label='assert:failures') def test_run_inference_timeout_on_load_dlq(self): with TestPipeline() as pipeline: examples = [1, 2] expected_good = [] expected_bad = [1, 2] pcoll = pipeline | 'start' >> beam.Create(examples) main, other = pcoll | base.RunInference( FakeSlowModelHandler(10)).with_exception_handling(timeout=2) assert_that(main, equal_to(expected_good), label='assert:inferences') # bad.failed_inferences will be in form [batch[elements], error]. # Just pull out bad element. bad_without_error = other.failed_inferences | beam.Map(lambda x: x[0][0]) assert_that( bad_without_error, equal_to(expected_bad), label='assert:failures') def test_run_inference_timeout_on_inference_dlq(self): with TestPipeline() as pipeline: examples = [10, 11] expected_good = [] expected_bad = [10, 11] pcoll = pipeline | 'start' >> beam.Create(examples) main, other = pcoll | base.RunInference( FakeSlowModelHandler(0)).with_exception_handling(timeout=5) assert_that(main, equal_to(expected_good), label='assert:inferences') # bad.failed_inferences will be in form [batch[elements], error]. # Just pull out bad element. bad_without_error = other.failed_inferences | beam.Map(lambda x: x[0][0]) assert_that( bad_without_error, equal_to(expected_bad), label='assert:failures') def test_run_inference_timeout_not_hit(self): with TestPipeline() as pipeline: examples = [1, 2] expected_good = [1, 2] expected_bad = [] pcoll = pipeline | 'start' >> beam.Create(examples) main, other = pcoll | base.RunInference( FakeSlowModelHandler(3)).with_exception_handling(timeout=500) assert_that(main, equal_to(expected_good), label='assert:inferences') # bad.failed_inferences will be in form [batch[elements], error]. # Just pull out bad element. bad_without_error = other.failed_inferences | beam.Map(lambda x: x[0][0]) assert_that( bad_without_error, equal_to(expected_bad), label='assert:failures') @unittest.skipIf( sys.platform == "win32" or sys.version_info < (3, 11), "This test relies on the __del__ lifecycle method, but __del__ does " + "not get invoked in the same way on older versions of Python or on " + "windows, breaking this test. See " + "github.com/python/cpython/issues/87950#issuecomment-1807570983 " + "for example.") def test_run_inference_timeout_does_garbage_collection(self): with tempfile.TemporaryDirectory() as tmp_dirname: tmp_path = os.path.join(tmp_dirname, 'tmp_filename') expected_file_contents = 'Deleted FakeSlowModel' with TestPipeline() as pipeline: # Start with bad example which gets timed out. # Then provide plenty of time for GC to happen. examples = [20] + [1] * 15 expected_good = [1] * 15 expected_bad = [20] pcoll = pipeline | 'start' >> beam.Create(examples) main, other = pcoll | base.RunInference( FakeSlowModelHandler( 0, True, tmp_path)).with_exception_handling(timeout=5) assert_that(main, equal_to(expected_good), label='assert:inferences') # # bad.failed_inferences will be in form [batch[elements], error]. # # Just pull out bad element. bad_without_error = other.failed_inferences | beam.Map( lambda x: x[0][0]) assert_that( bad_without_error, equal_to(expected_bad), label='assert:failures') with open(tmp_path) as f: s = f.read() self.assertEqual(s, expected_file_contents) def test_run_inference_impl_inference_args(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] pcoll = pipeline | 'start' >> beam.Create(examples) inference_args = {'key': True} actual = pcoll | base.RunInference( FakeModelHandlerExpectedInferenceArgs(), inference_args=inference_args) assert_that(actual, equal_to(examples), label='assert:inferences') def test_run_inference_metrics_with_custom_namespace(self): metrics_namespace = 'my_custom_namespace' pipeline = TestPipeline() examples = [1, 5, 3, 10] pcoll = pipeline | 'start' >> beam.Create(examples) _ = pcoll | base.RunInference( FakeModelHandler(), metrics_namespace=metrics_namespace) result = pipeline.run() result.wait_until_finish() metrics_filter = MetricsFilter().with_namespace(namespace=metrics_namespace) metrics = result.metrics().query(metrics_filter) assert len(metrics['counters']) != 0 assert len(metrics['distributions']) != 0 metrics_filter = MetricsFilter().with_namespace(namespace='fake_namespace') metrics = result.metrics().query(metrics_filter) assert len(metrics['counters']) == len(metrics['distributions']) == 0 def test_unexpected_inference_args_passed(self): with self.assertRaisesRegex(ValueError, r'inference_args were provided'): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] pcoll = pipeline | 'start' >> beam.Create(examples) inference_args = {'key': True} _ = pcoll | base.RunInference( FakeModelHandlerFailsOnInferenceArgs(), inference_args=inference_args) def test_increment_failed_batches_counter(self): with self.assertRaises(ValueError): # TODO(https://github.com/apache/beam/issues/34549): This test relies on # metrics filtering which doesn't work on Prism yet because Prism renames # steps (e.g. "Do" becomes "ref_AppliedPTransform_Do_7"). # https://github.com/apache/beam/blob/5f9cd73b7c9a2f37f83971ace3a399d633201dd1/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py#L1590 with TestPipeline('FnApiRunner') as pipeline: examples = [7] pcoll = pipeline | 'start' >> beam.Create(examples) _ = pcoll | base.RunInference(FakeModelHandlerExpectedInferenceArgs()) run_result = pipeline.run() run_result.wait_until_finish() metric_results = ( run_result.metrics().query( MetricsFilter().with_name('failed_batches_counter'))) num_failed_batches_counter = metric_results['counters'][0] self.assertEqual(num_failed_batches_counter.committed, 3) # !!!: The above will need to be updated if retry behavior changes def test_failed_batches_counter_no_failures(self): pipeline = TestPipeline() examples = [7] pcoll = pipeline | 'start' >> beam.Create(examples) inference_args = {'key': True} _ = pcoll | base.RunInference( FakeModelHandlerExpectedInferenceArgs(), inference_args=inference_args) run_result = pipeline.run() run_result.wait_until_finish() metric_results = ( run_result.metrics().query( MetricsFilter().with_name('failed_batches_counter'))) self.assertEqual(len(metric_results['counters']), 0) def test_counted_metrics(self): pipeline = TestPipeline() examples = [1, 5, 3, 10] pcoll = pipeline | 'start' >> beam.Create(examples) _ = pcoll | base.RunInference(FakeModelHandler()) run_result = pipeline.run() run_result.wait_until_finish() metric_results = ( run_result.metrics().query(MetricsFilter().with_name('num_inferences'))) num_inferences_counter = metric_results['counters'][0] self.assertEqual(num_inferences_counter.committed, 4) inference_request_batch_size = run_result.metrics().query( MetricsFilter().with_name('inference_request_batch_size')) self.assertTrue(inference_request_batch_size['distributions']) self.assertEqual( inference_request_batch_size['distributions'][0].result.sum, 4) inference_request_batch_byte_size = run_result.metrics().query( MetricsFilter().with_name('inference_request_batch_byte_size')) self.assertTrue(inference_request_batch_byte_size['distributions']) self.assertGreaterEqual( inference_request_batch_byte_size['distributions'][0].result.sum, len(pickle.dumps(examples))) inference_request_batch_byte_size = run_result.metrics().query( MetricsFilter().with_name('model_byte_size')) self.assertTrue(inference_request_batch_byte_size['distributions']) def test_timing_metrics(self): pipeline = TestPipeline() examples = [1, 5, 3, 10] pcoll = pipeline | 'start' >> beam.Create(examples) fake_clock = FakeClock() _ = pcoll | base.RunInference( FakeModelHandler(clock=fake_clock), clock=fake_clock) res = pipeline.run() res.wait_until_finish() metric_results = ( res.metrics().query( MetricsFilter().with_name('inference_batch_latency_micro_secs'))) batch_latency = metric_results['distributions'][0] self.assertEqual(batch_latency.result.count, 3) self.assertEqual(batch_latency.result.mean, 3000) metric_results = ( res.metrics().query( MetricsFilter().with_name('load_model_latency_milli_secs'))) load_model_latency = metric_results['distributions'][0] self.assertEqual(load_model_latency.result.count, 1) self.assertEqual(load_model_latency.result.mean, 500) def test_forwards_batch_args(self): examples = list(range(100)) with TestPipeline('FnApiRunner') as pipeline: pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference(FakeModelHandlerNeedsBigBatch()) assert_that(actual, equal_to(examples), label='assert:inferences') def test_run_inference_unkeyed_examples_with_keyed_model_handler(self): pipeline = TestPipeline() with self.assertRaisesRegex(Exception, "object is not iterable"): examples = [1, 3, 5] model_handler = base.KeyedModelHandler(FakeModelHandler()) _ = ( pipeline | 'Unkeyed' >> beam.Create(examples) | 'RunUnkeyed' >> base.RunInference(model_handler)) pipeline.run().wait_until_finish() def test_run_inference_keyed_examples_with_unkeyed_model_handler(self): pipeline = TestPipeline() examples = [1, 3, 5] keyed_examples = [(i, example) for i, example in enumerate(examples)] model_handler = FakeModelHandler() with self.assertRaisesRegex(Exception, "can only concatenate tuple"): _ = ( pipeline | 'keyed' >> beam.Create(keyed_examples) | 'RunKeyed' >> base.RunInference(model_handler)) pipeline.run() def test_model_handler_compatibility(self): # ** IMPORTANT ** Do not change this test to make your PR pass without # first reading below. # Be certain that the modification will not break third party # implementations of ModelHandler. # See issue https://github.com/apache/beam/issues/23484 # If this test fails, likely third party implementations of # ModelHandler will break. class ThirdPartyHandler(base.ModelHandler[int, int, FakeModel]): def __init__(self, custom_parameter=None): pass def load_model(self) -> FakeModel: return FakeModel() def run_inference( self, batch: Sequence[int], model: FakeModel, inference_args: Optional[dict[str, Any]] = None) -> Iterable[int]: yield 0 def get_num_bytes(self, batch: Sequence[int]) -> int: return 1 def get_metrics_namespace(self) -> str: return 'ThirdParty' def get_resource_hints(self) -> dict: return {} def batch_elements_kwargs(self) -> Mapping[str, Any]: return {} def validate_inference_args( self, inference_args: Optional[dict[str, Any]]): pass # This test passes if calling these methods does not cause # any runtime exceptions. third_party_model_handler = ThirdPartyHandler(custom_parameter=0) fake_model = third_party_model_handler.load_model() third_party_model_handler.run_inference([], fake_model) fake_inference_args = {'some_arg': 1} third_party_model_handler.run_inference([], fake_model, inference_args=fake_inference_args) third_party_model_handler.get_num_bytes([1, 2, 3]) third_party_model_handler.get_metrics_namespace() third_party_model_handler.get_resource_hints() third_party_model_handler.batch_elements_kwargs() third_party_model_handler.validate_inference_args({}) def test_run_inference_prediction_result_with_model_id(self): examples = [1, 5, 3, 10] expected = [ base.PredictionResult( example=example, inference=example + 1, model_id='fake_model_id_default') for example in examples ] with TestPipeline() as pipeline: pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference( FakeModelHandlerReturnsPredictionResult()) assert_that(actual, equal_to(expected), label='assert:inferences') def test_run_inference_with_iterable_side_input(self): test_pipeline = TestPipeline() side_input = ( test_pipeline | "CreateDummySideInput" >> beam.Create( [base.ModelMetadata(1, 1), base.ModelMetadata(2, 2)]) | "ApplySideInputWindow" >> beam.WindowInto( window.GlobalWindows(), trigger=trigger.Repeatedly(trigger.AfterProcessingTime(1)), accumulation_mode=trigger.AccumulationMode.DISCARDING)) test_pipeline.options.view_as(StandardOptions).streaming = True with self.assertRaises(Exception) as e: _ = ( test_pipeline | beam.Create([1, 2, 3, 4]) | base.RunInference( FakeModelHandler(), model_metadata_pcoll=side_input)) test_pipeline.run() msg = str(e.exception) self.assertIn('singleton', msg, msg='Expected singleton view error') self.assertIn('more than one', msg, msg='Expected multiple-elements error') def test_run_inference_with_iterable_side_input_multi_process_shared(self): test_pipeline = TestPipeline() side_input = ( test_pipeline | "CreateDummySideInput" >> beam.Create( [base.ModelMetadata(1, 1), base.ModelMetadata(2, 2)]) | "ApplySideInputWindow" >> beam.WindowInto( window.GlobalWindows(), trigger=trigger.Repeatedly(trigger.AfterProcessingTime(1)), accumulation_mode=trigger.AccumulationMode.DISCARDING)) test_pipeline.options.view_as(StandardOptions).streaming = True with self.assertRaises(Exception) as e: _ = ( test_pipeline | beam.Create([1, 2, 3, 4]) | base.RunInference( FakeModelHandler(multi_process_shared=True), model_metadata_pcoll=side_input)) test_pipeline.run() msg = str(e.exception) self.assertIn('singleton', msg, msg='Expected singleton view error') self.assertIn('more than one', msg, msg='Expected multiple-elements error') def test_run_inference_empty_side_input(self): model_handler = FakeModelHandlerReturnsPredictionResult() main_input_elements = [1, 2] with TestPipeline(is_integration_test=False) as pipeline: side_pcoll = pipeline | "side" >> beam.Create([]) result_pcoll = ( pipeline | beam.Create(main_input_elements) | base.RunInference(model_handler, model_metadata_pcoll=side_pcoll)) expected = [ base.PredictionResult(ele, ele + 1, 'fake_model_id_default') for ele in main_input_elements ] assert_that(result_pcoll, equal_to(expected)) def test_run_inference_side_input_in_batch(self): first_ts = math.floor(time.time()) - 30 interval = 7 sample_main_input_elements = ([ first_ts - 2, first_ts + 1, first_ts + 8, first_ts + 15, first_ts + 22, ]) sample_side_input_elements = [ (first_ts + 1, base.ModelMetadata(model_id='', model_name='')), # if model_id is empty string, we use the default model # handler model URI. ( first_ts + 8, base.ModelMetadata( model_id='fake_model_id_1', model_name='fake_model_id_1')), ( first_ts + 15, base.ModelMetadata( model_id='fake_model_id_2', model_name='fake_model_id_2')) ] model_handler = FakeModelHandlerReturnsPredictionResult() # applying GroupByKey to utilize windowing according to # https://beam.apache.org/documentation/programming-guide/#windowing-bounded-collections class _EmitElement(beam.DoFn): def process(self, element): for e in element: yield e # This test relies on poorly defined side input semantics which vary # across runners (including prism). Pinning to FnApiRunner which # consistently guarantees output. with TestPipeline('FnApiRunner') as pipeline: side_input = ( pipeline | "CreateSideInputElements" >> beam.Create(sample_side_input_elements) | beam.Map(lambda x: TimestampedValue(x[1], x[0])) | beam.WindowInto( window.FixedWindows(interval), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.Map(lambda x: ('key', x)) | beam.GroupByKey() | beam.Map(lambda x: x[1]) | "EmitSideInput" >> beam.ParDo(_EmitElement())) result_pcoll = ( pipeline | beam.Create(sample_main_input_elements) | "MapTimeStamp" >> beam.Map(lambda x: TimestampedValue(x, x)) | "ApplyWindow" >> beam.WindowInto(window.FixedWindows(interval)) | beam.Map(lambda x: ('key', x)) | "MainInputGBK" >> beam.GroupByKey() | beam.Map(lambda x: x[1]) | beam.ParDo(_EmitElement()) | "RunInference" >> base.RunInference( model_handler, model_metadata_pcoll=side_input)) expected_model_id_order = [ 'fake_model_id_default', 'fake_model_id_default', 'fake_model_id_1', 'fake_model_id_2', 'fake_model_id_2' ] expected_result = [ base.PredictionResult( example=sample_main_input_elements[i], inference=sample_main_input_elements[i] + 1, model_id=expected_model_id_order[i]) for i in range(5) ] assert_that(result_pcoll, equal_to(expected_result)) def test_run_inference_side_input_in_batch_per_key_models(self): first_ts = math.floor(time.time()) - 30 interval = 7 sample_main_input_elements = ([ ('key1', first_ts - 2), ('key2', first_ts + 1), ('key2', first_ts + 8), ('key1', first_ts + 15), ('key2', first_ts + 22), ('key1', first_ts + 29), ]) sample_side_input_elements = [ ( first_ts + 1, [ base.KeyModelPathMapping( keys=['key1'], update_path='fake_model_id_default'), base.KeyModelPathMapping( keys=['key2'], update_path='fake_model_id_default') ]), # if model_id is empty string, we use the default model # handler model URI. ( first_ts + 8, [ base.KeyModelPathMapping( keys=['key1'], update_path='fake_model_id_1'), base.KeyModelPathMapping( keys=['key2'], update_path='fake_model_id_default') ]), ( first_ts + 15, [ base.KeyModelPathMapping( keys=['key1'], update_path='fake_model_id_1'), base.KeyModelPathMapping( keys=['key2'], update_path='fake_model_id_2') ]), ] model_handler = base.KeyedModelHandler([ base.KeyModelMapping(['key1'], FakeModelHandlerReturnsPredictionResult( multi_process_shared=True, state=True)), base.KeyModelMapping(['key2'], FakeModelHandlerReturnsPredictionResult( multi_process_shared=True, state=True)) ]) class _EmitElement(beam.DoFn): def process(self, element): for e in element: yield e # This test relies on poorly defined side input semantics which vary # across runners (including prism). Pinning to FnApiRunner which # consistently guarantees output. with TestPipeline('FnApiRunner') as pipeline: side_input = ( pipeline | "CreateSideInputElements" >> beam.Create(sample_side_input_elements) | beam.Map(lambda x: TimestampedValue(x[1], x[0])) | beam.WindowInto( window.FixedWindows(interval), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.Map(lambda x: ('key', x)) | beam.GroupByKey() | beam.Map(lambda x: x[1]) | "EmitSideInput" >> beam.ParDo(_EmitElement())) result_pcoll = ( pipeline | beam.Create(sample_main_input_elements) | "MapTimeStamp" >> beam.Map(lambda x: TimestampedValue(x, x[1])) | "ApplyWindow" >> beam.WindowInto(window.FixedWindows(interval)) | "RunInference" >> base.RunInference( model_handler, model_metadata_pcoll=side_input) | beam.Map(lambda x: x[1])) expected_model_id_order = [ 'fake_model_id_default', 'fake_model_id_default', 'fake_model_id_default', 'fake_model_id_1', 'fake_model_id_2', 'fake_model_id_1', ] expected_inferences = [ 0, 0, 1, 0, 0, 1, ] expected_result = [ base.PredictionResult( example=sample_main_input_elements[i][1], inference=expected_inferences[i], model_id=expected_model_id_order[i]) for i in range(len(expected_inferences)) ] assert_that(result_pcoll, equal_to(expected_result)) def test_run_inference_side_input_in_batch_per_key_models_split_cohort(self): first_ts = math.floor(time.time()) - 30 interval = 7 sample_main_input_elements = ([ ('key1', first_ts - 2), ('key2', first_ts + 1), ('key1', first_ts + 8), ('key2', first_ts + 15), ('key1', first_ts + 22), ]) sample_side_input_elements = [ ( first_ts + 1, [ base.KeyModelPathMapping( keys=['key1', 'key2'], update_path='fake_model_id_default') ]), # if model_id is empty string, we use the default model # handler model URI. ( first_ts + 8, [ base.KeyModelPathMapping( keys=['key1'], update_path='fake_model_id_1'), base.KeyModelPathMapping( keys=['key2'], update_path='fake_model_id_default') ]), ( first_ts + 15, [ base.KeyModelPathMapping( keys=['key1'], update_path='fake_model_id_1'), base.KeyModelPathMapping( keys=['key2'], update_path='fake_model_id_2') ]), ] model_handler = base.KeyedModelHandler([ base.KeyModelMapping( ['key1', 'key2'], FakeModelHandlerReturnsPredictionResult(multi_process_shared=True)) ]) class _EmitElement(beam.DoFn): def process(self, element): for e in element: yield e # This test relies on poorly defined side input semantics which vary # across runners (including prism). Pinning to FnApiRunner which # consistently guarantees output. with TestPipeline('FnApiRunner') as pipeline: side_input = ( pipeline | "CreateSideInputElements" >> beam.Create(sample_side_input_elements) | beam.Map(lambda x: TimestampedValue(x[1], x[0])) | beam.WindowInto( window.FixedWindows(interval), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.Map(lambda x: ('key', x)) | beam.GroupByKey() | beam.Map(lambda x: x[1]) | "EmitSideInput" >> beam.ParDo(_EmitElement())) result_pcoll = ( pipeline | beam.Create(sample_main_input_elements) | "MapTimeStamp" >> beam.Map(lambda x: TimestampedValue(x, x[1])) | "ApplyWindow" >> beam.WindowInto(window.FixedWindows(interval)) | "RunInference" >> base.RunInference( model_handler, model_metadata_pcoll=side_input) | beam.Map(lambda x: x[1])) expected_model_id_order = [ 'fake_model_id_default', 'fake_model_id_default', 'fake_model_id_1', 'fake_model_id_2', 'fake_model_id_1' ] expected_result = [ base.PredictionResult( example=sample_main_input_elements[i][1], inference=sample_main_input_elements[i][1] + 1, model_id=expected_model_id_order[i]) for i in range(5) ] assert_that(result_pcoll, equal_to(expected_result)) def test_run_inference_side_input_in_batch_multi_process_shared(self): first_ts = math.floor(time.time()) - 30 interval = 7 sample_main_input_elements = ([ first_ts - 2, first_ts + 1, first_ts + 8, first_ts + 15, first_ts + 22, ]) sample_side_input_elements = [ (first_ts + 1, base.ModelMetadata(model_id='', model_name='')), # if model_id is empty string, we use the default model # handler model URI. ( first_ts + 8, base.ModelMetadata( model_id='fake_model_id_1', model_name='fake_model_id_1')), ( first_ts + 15, base.ModelMetadata( model_id='fake_model_id_2', model_name='fake_model_id_2')) ] model_handler = FakeModelHandlerReturnsPredictionResult( multi_process_shared=True) # applying GroupByKey to utilize windowing according to # https://beam.apache.org/documentation/programming-guide/#windowing-bounded-collections class _EmitElement(beam.DoFn): def process(self, element): for e in element: yield e # This test relies on poorly defined side input semantics which vary # across runners (including prism). Pinning to FnApiRunner which # consistently guarantees output. with TestPipeline('FnApiRunner') as pipeline: side_input = ( pipeline | "CreateSideInputElements" >> beam.Create(sample_side_input_elements) | beam.Map(lambda x: TimestampedValue(x[1], x[0])) | beam.WindowInto( window.FixedWindows(interval), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.Map(lambda x: ('key', x)) | beam.GroupByKey() | beam.Map(lambda x: x[1]) | "EmitSideInput" >> beam.ParDo(_EmitElement())) result_pcoll = ( pipeline | beam.Create(sample_main_input_elements) | "MapTimeStamp" >> beam.Map(lambda x: TimestampedValue(x, x)) | "ApplyWindow" >> beam.WindowInto(window.FixedWindows(interval)) | beam.Map(lambda x: ('key', x)) | "MainInputGBK" >> beam.GroupByKey() | beam.Map(lambda x: x[1]) | beam.ParDo(_EmitElement()) | "RunInference" >> base.RunInference( model_handler, model_metadata_pcoll=side_input)) expected_model_id_order = [ 'fake_model_id_default', 'fake_model_id_default', 'fake_model_id_1', 'fake_model_id_2', 'fake_model_id_2' ] expected_result = [ base.PredictionResult( example=sample_main_input_elements[i], inference=sample_main_input_elements[i] + 1, model_id=expected_model_id_order[i]) for i in range(5) ] assert_that(result_pcoll, equal_to(expected_result)) @unittest.skipIf( not TestPipeline().get_pipeline_options().view_as( StandardOptions).streaming, "SideInputs to RunInference are only supported in streaming mode.") @pytest.mark.it_postcommit @pytest.mark.sickbay_direct @pytest.mark.it_validatesrunner def test_run_inference_with_side_inputin_streaming(self): test_pipeline = TestPipeline(is_integration_test=True) test_pipeline.options.view_as(StandardOptions).streaming = True run_inference_side_inputs.run( test_pipeline.get_full_options_as_args(), save_main_session=False) def test_env_vars_set_correctly(self): handler_with_vars = FakeModelHandler(env_vars={'FOO': 'bar'}) os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) with TestPipeline() as pipeline: examples = [1, 2, 3] _ = ( pipeline | 'start' >> beam.Create(examples) | base.RunInference(handler_with_vars)) pipeline.run() self.assertTrue('FOO' in os.environ) self.assertTrue((os.environ['FOO']) == 'bar') def test_child_class_without_env_vars(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] expected = [example + 1 for example in examples] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference(FakeModelHandlerNoEnvVars()) assert_that(actual, equal_to(expected), label='assert:inferences') def test_model_handler_manager_loads_shared_model(self): mhs = { 'key1': FakeModelHandler(state=1), 'key2': FakeModelHandler(state=2), 'key3': FakeModelHandler(state=3) } mm = base._ModelHandlerManager(mh_map=mhs) tag1 = mm.load('key1').model_tag # Use bad_mh's load function to make sure we're actually loading the # version already stored bad_mh = FakeModelHandler(state=100) model1 = multi_process_shared.MultiProcessShared( bad_mh.load_model, tag=tag1).acquire() self.assertEqual(1, model1.predict(10)) tag2 = mm.load('key2').model_tag tag3 = mm.load('key3').model_tag model2 = multi_process_shared.MultiProcessShared( bad_mh.load_model, tag=tag2).acquire() model3 = multi_process_shared.MultiProcessShared( bad_mh.load_model, tag=tag3).acquire() self.assertEqual(2, model2.predict(10)) self.assertEqual(3, model3.predict(10)) def test_model_handler_manager_evicts_models(self): mh1 = FakeModelHandler(state=1) mh2 = FakeModelHandler(state=2) mh3 = FakeModelHandler(state=3) mhs = {'key1': mh1, 'key2': mh2, 'key3': mh3} mm = base._ModelHandlerManager(mh_map=mhs) mm.increment_max_models(2) tag1 = mm.load('key1').model_tag sh1 = multi_process_shared.MultiProcessShared(mh1.load_model, tag=tag1) model1 = sh1.acquire() self.assertEqual(1, model1.predict(10)) model1.increment_state(5) tag2 = mm.load('key2').model_tag tag3 = mm.load('key3').model_tag sh2 = multi_process_shared.MultiProcessShared(mh2.load_model, tag=tag2) model2 = sh2.acquire() sh3 = multi_process_shared.MultiProcessShared(mh3.load_model, tag=tag3) model3 = sh3.acquire() model2.increment_state(5) model3.increment_state(5) self.assertEqual(7, model2.predict(10)) self.assertEqual(8, model3.predict(10)) sh2.release(model2) sh3.release(model3) # model1 should have retained a valid reference to the model until released self.assertEqual(6, model1.predict(10)) sh1.release(model1) # This should now have been garbage collected, so it now it should get # recreated and shouldn't have the state updates model1 = multi_process_shared.MultiProcessShared( mh1.load_model, tag=tag1).acquire() self.assertEqual(1, model1.predict(10)) # These should not get recreated, so they should have the state updates model2 = multi_process_shared.MultiProcessShared( mh2.load_model, tag=tag2).acquire() self.assertEqual(7, model2.predict(10)) model3 = multi_process_shared.MultiProcessShared( mh3.load_model, tag=tag3).acquire() self.assertEqual(8, model3.predict(10)) def test_model_handler_manager_evicts_models_after_update(self): mh1 = FakeModelHandler(state=1) mhs = {'key1': mh1} mm = base._ModelHandlerManager(mh_map=mhs) tag1 = mm.load('key1').model_tag sh1 = multi_process_shared.MultiProcessShared(mh1.load_model, tag=tag1) model1 = sh1.acquire() self.assertEqual(1, model1.predict(10)) model1.increment_state(5) self.assertEqual(6, model1.predict(10)) mm.update_model_handler('key1', 'fake/path', 'key1') self.assertEqual(6, model1.predict(10)) sh1.release(model1) tag1 = mm.load('key1').model_tag sh1 = multi_process_shared.MultiProcessShared(mh1.load_model, tag=tag1) model1 = sh1.acquire() self.assertEqual(1, model1.predict(10)) model1.increment_state(5) self.assertEqual(6, model1.predict(10)) sh1.release(model1) # Shouldn't evict if path is the same as last update mm.update_model_handler('key1', 'fake/path', 'key1') tag1 = mm.load('key1').model_tag sh1 = multi_process_shared.MultiProcessShared(mh1.load_model, tag=tag1) model1 = sh1.acquire() self.assertEqual(6, model1.predict(10)) sh1.release(model1) def test_model_handler_manager_evicts_models_after_being_incremented(self): mh1 = FakeModelHandler(state=1) mh2 = FakeModelHandler(state=2) mh3 = FakeModelHandler(state=3) mhs = {'key1': mh1, 'key2': mh2, 'key3': mh3} mm = base._ModelHandlerManager(mh_map=mhs) mm.increment_max_models(1) mm.increment_max_models(1) tag1 = mm.load('key1').model_tag sh1 = multi_process_shared.MultiProcessShared(mh1.load_model, tag=tag1) model1 = sh1.acquire() self.assertEqual(1, model1.predict(10)) model1.increment_state(5) self.assertEqual(6, model1.predict(10)) sh1.release(model1) tag2 = mm.load('key2').model_tag tag3 = mm.load('key3').model_tag sh2 = multi_process_shared.MultiProcessShared(mh2.load_model, tag=tag2) model2 = sh2.acquire() sh3 = multi_process_shared.MultiProcessShared(mh3.load_model, tag=tag3) model3 = sh3.acquire() model2.increment_state(5) model3.increment_state(5) self.assertEqual(7, model2.predict(10)) self.assertEqual(8, model3.predict(10)) sh2.release(model2) sh3.release(model3) # This should get recreated, so it shouldn't have the state updates model1 = multi_process_shared.MultiProcessShared( mh1.load_model, tag=tag1).acquire() self.assertEqual(1, model1.predict(10)) # These should not get recreated, so they should have the state updates model2 = multi_process_shared.MultiProcessShared( mh2.load_model, tag=tag2).acquire() self.assertEqual(7, model2.predict(10)) model3 = multi_process_shared.MultiProcessShared( mh3.load_model, tag=tag3).acquire() self.assertEqual(8, model3.predict(10)) def test_run_inference_loads_different_models(self): mh1 = FakeModelHandler(incrementing=True, min_batch_size=3) with TestPipeline() as pipeline: pcoll = pipeline | 'start' >> beam.Create([1, 2, 3]) actual = ( pcoll | 'ri1' >> base.RunInference(mh1) | 'ri2' >> base.RunInference(mh1)) assert_that(actual, equal_to([1, 2, 3]), label='assert:inferences') def test_run_inference_loads_different_models_multi_process_shared(self): mh1 = FakeModelHandler( incrementing=True, min_batch_size=3, multi_process_shared=True) with TestPipeline() as pipeline: pcoll = pipeline | 'start' >> beam.Create([1, 2, 3]) actual = ( pcoll | 'ri1' >> base.RunInference(mh1) | 'ri2' >> base.RunInference(mh1)) assert_that(actual, equal_to([1, 2, 3]), label='assert:inferences') def test_runinference_loads_same_model_with_identifier_multi_process_shared( self): mh1 = FakeModelHandler( incrementing=True, min_batch_size=3, multi_process_shared=True) with TestPipeline() as pipeline: pcoll = pipeline | 'start' >> beam.Create([1, 2, 3]) actual = ( pcoll | 'ri1' >> base.RunInference( mh1, model_identifier='same_model_with_identifier_multi_process_shared' ) | 'ri2' >> base.RunInference( mh1, model_identifier='same_model_with_identifier_multi_process_shared' )) assert_that(actual, equal_to([4, 5, 6]), label='assert:inferences') def test_run_inference_watch_file_pattern_side_input_label(self): pipeline = TestPipeline() # label of the WatchPattern transform. side_input_str = 'WatchFilePattern/ApplyGlobalWindow' from apache_beam.ml.inference.utils import WatchFilePattern file_pattern_side_input = ( pipeline | 'WatchFilePattern' >> WatchFilePattern(file_pattern='fake/path/*')) pcoll = pipeline | 'start' >> beam.Create([1, 2, 3]) result_pcoll = pcoll | base.RunInference( FakeModelHandler(), model_metadata_pcoll=file_pattern_side_input) assert side_input_str in str(result_pcoll.producer.side_inputs[0]) def test_run_inference_watch_file_pattern_keyword_arg_side_input_label(self): # label of the WatchPattern transform. side_input_str = 'WatchFilePattern/ApplyGlobalWindow' pipeline = TestPipeline() pcoll = pipeline | 'start' >> beam.Create([1, 2, 3]) result_pcoll = pcoll | base.RunInference( FakeModelHandler(), watch_model_pattern='fake/path/*') assert side_input_str in str(result_pcoll.producer.side_inputs[0]) def test_model_status_provides_valid_tags(self): ms = base._ModelStatus(False) self.assertTrue(ms.is_valid_tag('tag1')) self.assertTrue(ms.is_valid_tag('tag2')) self.assertEqual('tag1', ms.get_valid_tag('tag1')) self.assertEqual('tag2', ms.get_valid_tag('tag2')) self.assertTrue(ms.is_valid_tag('tag1')) self.assertTrue(ms.is_valid_tag('tag2')) ms.try_mark_current_model_invalid(0) self.assertFalse(ms.is_valid_tag('tag1')) self.assertFalse(ms.is_valid_tag('tag2')) self.assertEqual('tag1_reload_1', ms.get_valid_tag('tag1')) self.assertEqual('tag2_reload_1', ms.get_valid_tag('tag2')) self.assertTrue(ms.is_valid_tag('tag1_reload_1')) self.assertTrue(ms.is_valid_tag('tag2_reload_1')) ms.try_mark_current_model_invalid(0) self.assertFalse(ms.is_valid_tag('tag1')) self.assertFalse(ms.is_valid_tag('tag2')) self.assertFalse(ms.is_valid_tag('tag1_reload_1')) self.assertFalse(ms.is_valid_tag('tag2_reload_1')) self.assertEqual('tag1_reload_2', ms.get_valid_tag('tag1')) self.assertEqual('tag2_reload_2', ms.get_valid_tag('tag2')) self.assertTrue(ms.is_valid_tag('tag1_reload_2')) self.assertTrue(ms.is_valid_tag('tag2_reload_2')) def test_model_status_provides_valid_garbage_collection(self): ms = base._ModelStatus(True) self.assertTrue(ms.is_valid_tag('mspvgc_tag1')) self.assertTrue(ms.is_valid_tag('mspvgc_tag2')) ms.try_mark_current_model_invalid(1) tags = ms.get_tags_for_garbage_collection() self.assertEqual(0, len(tags)) time.sleep(4) ms.try_mark_current_model_invalid(1) time.sleep(4) self.assertTrue(ms.is_valid_tag('mspvgc_tag3')) tags = ms.get_tags_for_garbage_collection() self.assertEqual(2, len(tags)) self.assertTrue('mspvgc_tag1' in tags) self.assertTrue('mspvgc_tag2' in tags) ms.try_mark_current_model_invalid(0) tags = ms.get_tags_for_garbage_collection() self.assertEqual(3, len(tags)) self.assertTrue('mspvgc_tag1' in tags) self.assertTrue('mspvgc_tag2' in tags) self.assertTrue('mspvgc_tag3' in tags) ms.mark_tags_deleted(set(['mspvgc_tag1', 'mspvgc_tag2'])) tags = ms.get_tags_for_garbage_collection() self.assertEqual(['mspvgc_tag3'], tags) ms.mark_tags_deleted(set(['mspvgc_tag1'])) tags = ms.get_tags_for_garbage_collection() self.assertEqual(1, len(tags)) self.assertEqual('mspvgc_tag3', tags[0]) ms.mark_tags_deleted(set(['mspvgc_tag3'])) tags = ms.get_tags_for_garbage_collection() self.assertEqual(0, len(tags)) def _always_retry(e: Exception) -> bool: return True class FakeRemoteModelHandler(base.RemoteModelHandler[int, int, FakeModel]): def __init__( self, clock=None, min_batch_size=1, max_batch_size=9999, retry_filter=_always_retry, **kwargs): self._fake_clock = clock self._min_batch_size = min_batch_size self._max_batch_size = max_batch_size self._env_vars = kwargs.get('env_vars', {}) self._multi_process_shared = multi_process_shared super().__init__( namespace='FakeRemoteModelHandler', retry_filter=retry_filter) def create_client(self): return FakeModel() def request(self, batch, model, inference_args=None) -> Iterable[int]: responses = [] for example in batch: responses.append(model.predict(example)) return responses def batch_elements_kwargs(self): return { 'min_batch_size': self._min_batch_size, 'max_batch_size': self._max_batch_size } class FakeAlwaysFailsRemoteModelHandler(base.RemoteModelHandler[int, int, FakeModel]): def __init__( self, clock=None, min_batch_size=1, max_batch_size=9999, retry_filter=_always_retry, **kwargs): self._fake_clock = clock self._min_batch_size = min_batch_size self._max_batch_size = max_batch_size self._env_vars = kwargs.get('env_vars', {}) super().__init__( namespace='FakeRemoteModelHandler', retry_filter=retry_filter, num_retries=2, throttle_delay_secs=1) def create_client(self): return FakeModel() def request(self, batch, model, inference_args=None) -> Iterable[int]: raise Exception def batch_elements_kwargs(self): return { 'min_batch_size': self._min_batch_size, 'max_batch_size': self._max_batch_size } class FakeFailsOnceRemoteModelHandler(base.RemoteModelHandler[int, int, FakeModel]): def __init__( self, clock=None, min_batch_size=1, max_batch_size=9999, retry_filter=_always_retry, **kwargs): self._fake_clock = clock self._min_batch_size = min_batch_size self._max_batch_size = max_batch_size self._env_vars = kwargs.get('env_vars', {}) self._should_fail = True super().__init__( namespace='FakeRemoteModelHandler', retry_filter=retry_filter, num_retries=2, throttle_delay_secs=1) def create_client(self): return FakeModel() def request(self, batch, model, inference_args=None) -> Iterable[int]: if self._should_fail: self._should_fail = False raise Exception else: self._should_fail = True responses = [] for example in batch: responses.append(model.predict(example)) return responses def batch_elements_kwargs(self): return { 'min_batch_size': self._min_batch_size, 'max_batch_size': self._max_batch_size } class RunInferenceRemoteTest(unittest.TestCase): def test_normal_model_execution(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] expected = [example + 1 for example in examples] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference(FakeRemoteModelHandler()) assert_that(actual, equal_to(expected), label='assert:inferences') def test_repeated_requests_fail(self): test_pipeline = TestPipeline() with self.assertRaises(Exception): _ = ( test_pipeline | beam.Create([1, 2, 3, 4]) | base.RunInference(FakeAlwaysFailsRemoteModelHandler())) test_pipeline.run() def test_works_on_retry(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] expected = [example + 1 for example in examples] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference(FakeFailsOnceRemoteModelHandler()) assert_that(actual, equal_to(expected), label='assert:inferences') def test_exception_on_load_model_override(self): with self.assertRaises(Exception): class _(base.RemoteModelHandler[int, int, FakeModel]): def __init__(self, clock=None, retry_filter=_always_retry, **kwargs): self._fake_clock = clock self._min_batch_size = 1 self._max_batch_size = 1 self._env_vars = kwargs.get('env_vars', {}) super().__init__( namespace='FakeRemoteModelHandler', retry_filter=retry_filter) def load_model(self): return FakeModel() def request(self, batch, model, inference_args=None) -> Iterable[int]: responses = [] for example in batch: responses.append(model.predict(example)) return responses def test_exception_on_run_inference_override(self): with self.assertRaises(Exception): class _(base.RemoteModelHandler[int, int, FakeModel]): def __init__(self, clock=None, retry_filter=_always_retry, **kwargs): self._fake_clock = clock self._min_batch_size = 1 self._max_batch_size = 1 self._env_vars = kwargs.get('env_vars', {}) super().__init__( namespace='FakeRemoteModelHandler', retry_filter=retry_filter) def create_client(self): return FakeModel() def run_inference(self, batch, model, inference_args=None) -> Iterable[int]: responses = [] for example in batch: responses.append(model.predict(example)) return responses def test_run_inference_with_rate_limiter(self): class FakeRateLimiter(base.RateLimiter): def __init__(self): super().__init__(namespace='test_namespace') def allow(self, hits_added=1): self.requests_counter.inc() return True limiter = FakeRateLimiter() with TestPipeline() as pipeline: examples = [1, 5] class ConcreteRemoteModelHandler(base.RemoteModelHandler): def create_client(self): return FakeModel() def request(self, batch, model, inference_args=None): return [model.predict(example) for example in batch] model_handler = ConcreteRemoteModelHandler( rate_limiter=limiter, namespace='test_namespace') pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference(model_handler) expected = [2, 6] assert_that(actual, equal_to(expected)) result = pipeline.run() result.wait_until_finish() metrics_filter = MetricsFilter().with_name( 'RatelimitRequestsTotal').with_namespace('test_namespace') metrics = result.metrics().query(metrics_filter) self.assertGreaterEqual(metrics['counters'][0].committed, 0) def test_run_inference_with_rate_limiter_exceeded(self): class FakeRateLimiter(base.RateLimiter): def __init__(self): super().__init__(namespace='test_namespace') def allow(self, hits_added=1): return False class ConcreteRemoteModelHandler(base.RemoteModelHandler): def create_client(self): return FakeModel() def request(self, batch, model, inference_args=None): return [model.predict(example) for example in batch] model_handler = ConcreteRemoteModelHandler( rate_limiter=FakeRateLimiter(), namespace='test_namespace', num_retries=0) with self.assertRaises(base.RateLimitExceeded): model_handler.run_inference([1], FakeModel()) class FakeModelHandlerForSizing(base.ModelHandler[int, int, FakeModel]): """A ModelHandler used to test element sizing behavior.""" def __init__( self, max_batch_size: int = 10, max_batch_weight: Optional[int] = None, element_size_fn=None): super().__init__( max_batch_size=max_batch_size, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn) def load_model(self) -> FakeModel: return FakeModel() def run_inference(self, batch, model, inference_args=None): return [model.predict(x) for x in batch] class RunInferenceSizeTest(unittest.TestCase): """Tests for ModelHandler.batch_elements_kwargs with element_size_fn.""" def test_kwargs_are_passed_correctly(self): """Adds element_size_fn without clobbering existing kwargs.""" def size_fn(x): return 10 sized_handler = FakeModelHandlerForSizing( max_batch_size=20, max_batch_weight=100, element_size_fn=size_fn) kwargs = sized_handler.batch_elements_kwargs() self.assertEqual(kwargs['max_batch_size'], 20) self.assertEqual(kwargs['max_batch_weight'], 100) self.assertIn('element_size_fn', kwargs) self.assertEqual(kwargs['element_size_fn'](1), 10) def test_sizing_with_edge_cases(self): """Allows extreme values from element_size_fn.""" zero_size_fn = lambda x: 0 sized_handler = FakeModelHandlerForSizing( max_batch_size=1, element_size_fn=zero_size_fn) kwargs = sized_handler.batch_elements_kwargs() self.assertEqual(kwargs['element_size_fn'](999), 0) large_size_fn = lambda x: 1000000 sized_handler = FakeModelHandlerForSizing( max_batch_size=1, element_size_fn=large_size_fn) kwargs = sized_handler.batch_elements_kwargs() self.assertEqual(kwargs['element_size_fn'](1), 1000000) class FakeModelHandlerForBatching(base.ModelHandler[int, int, FakeModel]): """A ModelHandler used to test batching behavior via base class __init__.""" def __init__(self, **kwargs): super().__init__(**kwargs) def load_model(self) -> FakeModel: return FakeModel() def run_inference(self, batch, model, inference_args=None): return [model.predict(x) for x in batch] class ModelHandlerBatchingArgsTest(unittest.TestCase): """Tests for ModelHandler.__init__ batching parameters.""" def test_batch_elements_kwargs_all_args(self): """All batching args passed to __init__ are in batch_elements_kwargs.""" def size_fn(x): return 10 handler = FakeModelHandlerForBatching( min_batch_size=5, max_batch_size=20, max_batch_duration_secs=30, max_batch_weight=100, element_size_fn=size_fn) kwargs = handler.batch_elements_kwargs() self.assertEqual(kwargs['min_batch_size'], 5) self.assertEqual(kwargs['max_batch_size'], 20) self.assertEqual(kwargs['max_batch_duration_secs'], 30) self.assertEqual(kwargs['max_batch_weight'], 100) self.assertIn('element_size_fn', kwargs) self.assertEqual(kwargs['element_size_fn'](1), 10) def test_batch_elements_kwargs_partial_args(self): """Only provided batching args are included in kwargs.""" handler = FakeModelHandlerForBatching(max_batch_size=50) kwargs = handler.batch_elements_kwargs() self.assertEqual(kwargs, {'max_batch_size': 50}) def test_batch_elements_kwargs_empty_when_no_args(self): """No batching kwargs when none are provided.""" handler = FakeModelHandlerForBatching() kwargs = handler.batch_elements_kwargs() self.assertEqual(kwargs, {}) def test_large_model_sets_share_across_processes(self): """Setting large_model=True enables share_model_across_processes.""" handler = FakeModelHandlerForBatching(large_model=True) self.assertTrue(handler.share_model_across_processes()) def test_model_copies_sets_share_across_processes(self): """Setting model_copies enables share_model_across_processes.""" handler = FakeModelHandlerForBatching(model_copies=2) self.assertTrue(handler.share_model_across_processes()) self.assertEqual(handler.model_copies(), 2) def test_default_share_across_processes_is_false(self): """Default share_model_across_processes is False.""" handler = FakeModelHandlerForBatching() self.assertFalse(handler.share_model_across_processes()) def test_default_model_copies_is_one(self): """Default model_copies is 1.""" handler = FakeModelHandlerForBatching() self.assertEqual(handler.model_copies(), 1) def test_env_vars_from_kwargs(self): """Environment variables can be passed via kwargs.""" handler = FakeModelHandlerForBatching(env_vars={'MY_VAR': 'value'}) self.assertEqual(handler._env_vars, {'MY_VAR': 'value'}) def test_min_batch_size_only(self): """min_batch_size can be passed alone.""" handler = FakeModelHandlerForBatching(min_batch_size=10) kwargs = handler.batch_elements_kwargs() self.assertEqual(kwargs, {'min_batch_size': 10}) def test_max_batch_duration_secs_only(self): """max_batch_duration_secs can be passed alone.""" handler = FakeModelHandlerForBatching(max_batch_duration_secs=60) kwargs = handler.batch_elements_kwargs() self.assertEqual(kwargs, {'max_batch_duration_secs': 60}) def test_batch_length_fn_and_batch_bucket_boundaries(self): """batch_length_fn and batch_bucket_boundaries passed through to kwargs.""" handler = FakeModelHandlerForBatching( batch_length_fn=len, batch_bucket_boundaries=[16, 32, 64]) kwargs = handler.batch_elements_kwargs() self.assertIs(kwargs['length_fn'], len) self.assertEqual(kwargs['bucket_boundaries'], [16, 32, 64]) def test_batch_length_fn_only(self): """batch_length_fn alone is passed through without bucket_boundaries.""" handler = FakeModelHandlerForBatching(batch_length_fn=len) kwargs = handler.batch_elements_kwargs() self.assertIs(kwargs['length_fn'], len) self.assertNotIn('bucket_boundaries', kwargs) def test_batch_bucket_boundaries_without_batch_length_fn(self): """Passing batch_bucket_boundaries without batch_length_fn should fail in BatchElements. Note: ModelHandler.__init__ doesn't validate this; the error is raised by BatchElements when batch_elements_kwargs are used.""" handler = FakeModelHandlerForBatching(batch_bucket_boundaries=[10, 20]) kwargs = handler.batch_elements_kwargs() # The kwargs are stored, but BatchElements will reject them self.assertEqual(kwargs['bucket_boundaries'], [10, 20]) self.assertNotIn('length_fn', kwargs) def test_batching_kwargs_none_values_omitted(self): """None values for batch_length_fn and batch_bucket_boundaries are not in kwargs.""" handler = FakeModelHandlerForBatching( min_batch_size=5, batch_length_fn=None, batch_bucket_boundaries=None) kwargs = handler.batch_elements_kwargs() self.assertNotIn('length_fn', kwargs) self.assertNotIn('bucket_boundaries', kwargs) self.assertEqual(kwargs['min_batch_size'], 5) class PaddingReportingStringModelHandler(base.ModelHandler[str, str, FakeModel]): """Reports each element with the max length of the batch it ran in.""" def load_model(self): return FakeModel() def run_inference(self, batch, model, inference_args=None): max_len = max(len(s) for s in batch) return [f'{s}:{max_len}' for s in batch] class RunInferenceLengthAwareBatchingTest(unittest.TestCase): """End-to-end tests for PR2 length-aware batching in RunInference.""" def test_run_inference_with_length_aware_batch_elements(self): handler = PaddingReportingStringModelHandler( min_batch_size=2, max_batch_size=2, max_batch_duration_secs=60, batch_length_fn=len, batch_bucket_boundaries=[5]) examples = ['a', 'cccccc', 'bb', 'ddddddd'] with TestPipeline('FnApiRunner') as p: results = ( p | beam.Create(examples, reshuffle=False) | base.RunInference(handler)) assert_that(results, equal_to(['a:2', 'bb:2', 'cccccc:7', 'ddddddd:7'])) class SimpleFakeModelHandler(base.ModelHandler[int, int, FakeModel]): def load_model(self): return FakeModel() def run_inference( self, batch: Sequence[int], model: FakeModel, inference_args=None) -> Iterable[int]: for example in batch: yield model.predict(example) def try_import_model_manager(): try: # pylint: disable=unused-import from apache_beam.ml.inference.model_manager import ModelManager return True except ImportError: return False class ModelManagerTest(unittest.TestCase): """Tests for RunInference with Model Manager integration.""" def tearDown(self): for p in multiprocessing.active_children(): p.terminate() p.join() @unittest.skipIf( not try_import_model_manager(), 'Model Manager not available') def test_run_inference_impl_with_model_manager(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] expected = [example + 1 for example in examples] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference( SimpleFakeModelHandler(), use_model_manager=True) assert_that(actual, equal_to(expected), label='assert:inferences') @unittest.skipIf( not try_import_model_manager(), 'Model Manager not available') def test_run_inference_impl_with_model_manager_args(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] expected = [example + 1 for example in examples] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference( SimpleFakeModelHandler(), use_model_manager=True, model_manager_args={ 'slack_percentage': 0.2, 'poll_interval': 1.0, 'peak_window_seconds': 10.0, 'min_data_points': 10, 'smoothing_factor': 0.5 }) assert_that(actual, equal_to(expected), label='assert:inferences') @unittest.skipIf( not try_import_model_manager(), 'Model Manager not available') def test_run_inference_impl_with_model_manager_oom(self): class OOMFakeModelHandler(SimpleFakeModelHandler): def run_inference( self, batch: Sequence[int], model: FakeModel, inference_args=None) -> Iterable[int]: if random.random() < 0.8: raise MemoryError("Simulated OOM") for example in batch: yield model.predict(example) def batch_elements_kwargs(self): return {'min_batch_size': 1, 'max_batch_size': 1} with self.assertRaises(Exception): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference( OOMFakeModelHandler(), use_model_manager=True) assert_that(actual, equal_to([2, 6, 4, 11]), label='assert:inferences') if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/gemini_inference.py0000644000175100017510000002473115206113032024072 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging from collections.abc import Callable from collections.abc import Iterable from collections.abc import Sequence from typing import Any from typing import Optional from typing import Union from typing import cast from google import genai from google.genai import errors from google.genai.types import HttpOptions from google.genai.types import Part from PIL.Image import Image from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RemoteModelHandler LOGGER = logging.getLogger("GeminiModelHandler") def _retry_on_appropriate_service_error(exception: Exception) -> bool: """ Retry filter that returns True if a returned HTTP error code is 5xx or 429. This is used to retry remote requests that fail, most notably 429 (throttling by the service) Args: exception: the returned exception encountered during the request/response loop. Returns: boolean indication whether or not the exception is a ServerError (5xx) or a 429 error. """ if not isinstance(exception, errors.APIError): return False return exception.code == 429 or exception.code >= 500 def generate_from_string( model_name: str, batch: Sequence[str], model: genai.Client, inference_args: dict[str, Any]): """ Request function that expects inputs to be composed of strings, then sends requests to Gemini to generate text responses based on the text prompts. Args: model_name: the Gemini model to use for the request. This model should be a text generation model. batch: the string inputs to be send to Gemini for text generation. model: the genai Client inference_args: any additional arguments passed to the generate_content call. """ return model.models.generate_content( model=model_name, contents=cast(Any, batch), **inference_args) def generate_image_from_strings_and_images( model_name: str, batch: Sequence[list[Union[str, Image, Part]]], model: genai.Client, inference_args: dict[str, Any]): """ Request function that expects inputs to be composed of lists of strings and PIL Image instances, then sends requests to Gemini to generate images based on the text prompts and contextual images. This is currently intended to be used with the gemini-2.5-flash-image model (AKA Nano Banana.) Args: model_name: the Gemini model to use for the request. This model should be an image generation model such as gemini-2.5-flash-image. batch: the inputs to be send to Gemini for image generation as prompts. Composed of text prompts and contextual pillow Images. model: the genai Client inference_args: any additional arguments passed to the generate_content call. """ return model.models.generate_content( model=model_name, contents=cast(Any, batch), **inference_args) class GeminiModelHandler(RemoteModelHandler[Any, PredictionResult, genai.Client]): def __init__( self, model_name: str, request_fn: Callable[[str, Sequence[Any], genai.Client, dict[str, Any]], Any], api_key: Optional[str] = None, project: Optional[str] = None, location: Optional[str] = None, use_vertex_flex_api: Optional[bool] = False, *, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, **kwargs): """Implementation of the ModelHandler interface for Google Gemini. **NOTE:** This API and its implementation are under development and do not provide backward compatibility guarantees. Gemini can be accessed through either the Vertex AI API or the Gemini Developer API, and this handler chooses which to connect to based upon the arguments provided. As written, this model handler operates solely on string input. Args: model_name: the Gemini model to send the request to request_fn: the function to use to send the request. Should take the model name and the parameters from request() and return the responses from Gemini. The class will handle bundling the inputs and responses together. api_key: the Gemini Developer API key to use for the requests. Setting this parameter sends requests for this job to the Gemini Developer API. If this paramter is provided, do not set the project or location parameters. project: the GCP project to use for Vertex AI requests. Setting this parameter routes requests to Vertex AI. If this paramter is provided, location must also be provided and api_key should not be set. location: the GCP project to use for Vertex AI requests. Setting this parameter routes requests to Vertex AI. If this paramter is provided, project must also be provided and api_key should not be set. use_vertex_flex_api: if true, use the Vertex Flex API. This is a cost-effective option for accessing Gemini models if you can tolerate longer response times and throttling. This is often beneficial for data processing workloads which usually have higher latency tolerance than live serving paths. See https://docs.cloud.google.com/vertex-ai/generative-ai/docs/flex-paygo for more details. min_batch_size: optional. the minimum batch size to use when batching inputs. max_batch_size: optional. the maximum batch size to use when batching inputs. max_batch_duration_secs: optional. the maximum amount of time to buffer a batch before emitting; used in streaming contexts. max_batch_weight: optional. the maximum total weight of a batch. element_size_fn: optional. a function that returns the size (weight) of an element. batch_length_fn: optional. a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: optional. a sorted list of positive boundary values for length-aware batching buckets. """ self._batching_kwargs = {} self._env_vars = kwargs.get('env_vars', {}) if min_batch_size is not None: self._batching_kwargs["min_batch_size"] = min_batch_size if max_batch_size is not None: self._batching_kwargs["max_batch_size"] = max_batch_size if max_batch_duration_secs is not None: self._batching_kwargs["max_batch_duration_secs"] = max_batch_duration_secs if max_batch_weight is not None: self._batching_kwargs["max_batch_weight"] = max_batch_weight if element_size_fn is not None: self._batching_kwargs['element_size_fn'] = element_size_fn if batch_length_fn is not None: self._batching_kwargs['length_fn'] = batch_length_fn if batch_bucket_boundaries is not None: self._batching_kwargs['bucket_boundaries'] = batch_bucket_boundaries self.model_name = model_name self.request_fn = request_fn if api_key: if project or location: raise ValueError("project and location must be None if api_key is set") self.api_key = api_key self.use_vertex = False else: if project is None or location is None: raise ValueError( "project and location must both be provided if api_key is None") self.project = project self.location = location self.use_vertex = True self.use_vertex_flex_api = use_vertex_flex_api super().__init__( namespace='GeminiModelHandler', retry_filter=_retry_on_appropriate_service_error, **kwargs) def batch_elements_kwargs(self): return self._batching_kwargs def create_client(self) -> genai.Client: """Creates the GenAI client used to send requests. Creates a version for the Vertex AI API or the Gemini Developer API based on the arguments provided when the GeminiModelHandler class is instantiated. """ if self.use_vertex: if self.use_vertex_flex_api: return genai.Client( vertexai=True, project=self.project, location=self.location, http_options=HttpOptions( api_version="v1", headers={"X-Vertex-AI-LLM-Request-Type": "flex"}, # Set timeout in the unit of millisecond. timeout=600000)) else: return genai.Client( vertexai=True, project=self.project, location=self.location) return genai.Client(api_key=self.api_key) def request( self, batch: Sequence[Any], model: genai.Client, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """ Sends a prediction request to a Gemini service containing a batch of inputs and matches that input with the prediction response from the endpoint as an iterable of PredictionResults. Args: batch: a sequence of any values to be passed to the Gemini service. Should be inputs accepted by the provided inference function. model: a genai.Client object configured to access the desired service. inference_args: any additional arguments to send as part of the prediction request. Returns: An iterable of Predictions. """ if inference_args is None: inference_args = {} # Wrap the responses in a list to prevent zip() call from treating the # response itself as an iterable of individual responses. responses = [self.request_fn(self.model_name, batch, model, inference_args)] return utils._convert_to_result(batch, responses, self.model_name) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/gemini_inference_it_test.py0000644000175100017510000000557115206113032025626 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-End test for Gemini Remote Inference""" import logging import unittest import uuid import pytest from apache_beam.io.filesystem import MatchResult from apache_beam.io.filesystems import FileSystems from apache_beam.testing.test_pipeline import TestPipeline # pylint: disable=ungrouped-imports try: from apache_beam.examples.inference import gemini_image_generation from apache_beam.examples.inference import gemini_text_classification except ImportError as e: raise unittest.SkipTest("Gemini model handler dependencies are not installed") _OUTPUT_DIR = "gs://apache-beam-ml/testing/outputs/gemini" _TEST_PROJECT = "apache-beam-testing" _TEST_REGION = "us-central1" class GeminiInference(unittest.TestCase): @pytest.mark.gemini_postcommit def test_gemini_text_classification(self): output_file = '/'.join([_OUTPUT_DIR, str(uuid.uuid4()), 'output.txt']) test_pipeline = TestPipeline(is_integration_test=True) extra_opts = { 'output': output_file, 'cloud_project': _TEST_PROJECT, 'cloud_region': _TEST_REGION } gemini_text_classification.run( test_pipeline.get_full_options_as_args(**extra_opts)) self.assertEqual(FileSystems().exists(output_file), True) def _flatten_match(self, match_results): return [ file_metadata for match_result in match_results for file_metadata in match_result.metadata_list ] @pytest.mark.gemini_postcommit def test_gemini_image_generation(self): output_dir = '/'.join([_OUTPUT_DIR, str(uuid.uuid4())]) test_pipeline = TestPipeline(is_integration_test=False) extra_opts = { 'output': output_dir, 'cloud_project': _TEST_PROJECT, 'cloud_region': _TEST_REGION } gemini_image_generation.run( test_pipeline.get_full_options_as_args(**extra_opts)) matches: MatchResult = FileSystems().match([output_dir + '/*']) self.assertGreater(len(matches), 0) for match in matches: for file in match.metadata_list: self.assertTrue(file.path.endswith(".png")) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/gemini_inference_test.py0000644000175100017510000000710415206113032025124 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest from unittest import mock try: from google.genai import errors from apache_beam.ml.inference.gemini_inference import GeminiModelHandler from apache_beam.ml.inference.gemini_inference import _retry_on_appropriate_service_error from apache_beam.ml.inference.gemini_inference import generate_from_string except ImportError: raise unittest.SkipTest('Gemini dependencies are not installed') class RetryOnClientErrorTest(unittest.TestCase): def test_retry_on_client_error_positive(self): e = errors.APIError(code=429, response_json={}) self.assertTrue(_retry_on_appropriate_service_error(e)) def test_retry_on_client_error_negative(self): e = errors.APIError(code=404, response_json={}) self.assertFalse(_retry_on_appropriate_service_error(e)) def test_retry_on_server_error(self): e = errors.APIError(code=501, response_json={}) self.assertTrue(_retry_on_appropriate_service_error(e)) class ModelHandlerArgConditions(unittest.TestCase): def test_all_params_set(self): self.assertRaises( ValueError, GeminiModelHandler, model_name="gemini-model-123", request_fn=generate_from_string, api_key="123456789", project="testproject", location="us-central1", ) def test_missing_vertex_location_param(self): self.assertRaises( ValueError, GeminiModelHandler, model_name="gemini-model-123", request_fn=generate_from_string, project="testproject", ) def test_missing_vertex_project_param(self): self.assertRaises( ValueError, GeminiModelHandler, model_name="gemini-model-123", request_fn=generate_from_string, location="us-central1", ) def test_missing_all_params(self): self.assertRaises( ValueError, GeminiModelHandler, model_name="gemini-model-123", request_fn=generate_from_string, ) @mock.patch('apache_beam.ml.inference.gemini_inference.genai.Client') @mock.patch('apache_beam.ml.inference.gemini_inference.HttpOptions') class TestGeminiModelHandler(unittest.TestCase): def test_create_client_with_flex_api( self, mock_http_options, mock_genai_client): handler = GeminiModelHandler( model_name="gemini-pro", request_fn=generate_from_string, project="test-project", location="us-central1", use_vertex_flex_api=True) handler.create_client() mock_http_options.assert_called_with( api_version="v1", headers={"X-Vertex-AI-LLM-Request-Type": "flex"}, timeout=600000, ) mock_genai_client.assert_called_with( vertexai=True, project="test-project", location="us-central1", http_options=mock_http_options.return_value) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/huggingface_inference.py0000644000175100017510000007217015206113032025071 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import sys from collections import defaultdict from collections.abc import Callable from collections.abc import Iterable from collections.abc import Sequence from enum import Enum from typing import Any from typing import Optional from typing import Union import tensorflow as tf import torch from transformers import AutoModel from transformers import Pipeline from transformers import TFAutoModel from transformers import pipeline from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.pytorch_inference import _convert_to_device _LOGGER = logging.getLogger(__name__) __all__ = [ "HuggingFaceModelHandlerTensor", "HuggingFaceModelHandlerKeyedTensor", "HuggingFacePipelineModelHandler", ] TensorInferenceFn = Callable[ [ Sequence[Union[torch.Tensor, tf.Tensor]], Union[AutoModel, TFAutoModel], str, Optional[dict[str, Any]], Optional[str], ], Iterable[PredictionResult], ] KeyedTensorInferenceFn = Callable[[ Sequence[dict[str, Union[torch.Tensor, tf.Tensor]]], Union[AutoModel, TFAutoModel], str, Optional[dict[str, Any]], Optional[str] ], Iterable[PredictionResult]] PipelineInferenceFn = Callable[ [Sequence[str], Pipeline, Optional[dict[str, Any]]], Iterable[PredictionResult]] class PipelineTask(str, Enum): """ PipelineTask defines all the tasks supported by the Hugging Face Pipelines listed at https://huggingface.co/docs/transformers/main_classes/pipelines. Only these tasks can be passed to HuggingFacePipelineModelHandler. """ AudioClassification = 'audio-classification' AutomaticSpeechRecognition = 'automatic-speech-recognition' Conversational = 'conversational' DepthEstimation = 'depth-estimation' DocumentQuestionAnswering = 'document-question-answering' FeatureExtraction = 'feature-extraction' FillMask = 'fill-mask' ImageClassification = 'image-classification' ImageSegmentation = 'image-segmentation' ImageToText = 'image-to-text' MaskGeneration = 'mask-generation' NER = 'ner' ObjectDetection = 'object-detection' QuestionAnswering = 'question-answering' SentimentAnalysis = 'sentiment-analysis' Summarization = 'summarization' TableQuestionAnswering = 'table-question-answering' TextClassification = 'text-classification' TextGeneration = 'text-generation' Text2TextGeneration = 'text2text-generation' TextToAudio = 'text-to-audio' TokenClassification = 'token-classification' Translation = 'translation' VideoClassification = 'video-classification' VisualQuestionAnswering = 'visual-question-answering' VQA = 'vqa' ZeroShotAudioClassification = 'zero-shot-audio-classification' ZeroShotClassification = 'zero-shot-classification' ZeroShotImageClassification = 'zero-shot-image-classification' ZeroShotObjectDetection = 'zero-shot-object-detection' Translation_XX_to_YY = 'translation_XX_to_YY' def _validate_constructor_args(model_uri, model_class): message = ( "Please provide both model class and model uri to load the model." "Got params as model_uri={model_uri} and " "model_class={model_class}.") if not model_uri and not model_class: raise RuntimeError( message.format(model_uri=model_uri, model_class=model_class)) elif not model_uri: raise RuntimeError( message.format(model_uri=model_uri, model_class=model_class)) elif not model_class: raise RuntimeError( message.format(model_uri=model_uri, model_class=model_class)) def no_gpu_available_warning(): _LOGGER.warning( "HuggingFaceModelHandler specified a 'GPU' device, " "but GPUs are not available. Switching to CPU.") def is_gpu_available_torch(): if torch.cuda.is_available(): return True else: no_gpu_available_warning() return False def get_device_torch(device): if device == "GPU" and is_gpu_available_torch(): return torch.device("cuda") return torch.device("cpu") def is_gpu_available_tensorflow(device): gpu_devices = tf.config.list_physical_devices(device) if len(gpu_devices) == 0: no_gpu_available_warning() return False return True def _validate_constructor_args_hf_pipeline(task, model): if not task and not model: raise RuntimeError( 'Please provide either task or model to the ' 'HuggingFacePipelineModelHandler. If the model already defines the ' 'task, no need to specify the task.') def _run_inference_torch_keyed_tensor( batch: Sequence[dict[str, torch.Tensor]], model: AutoModel, device, inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: device = get_device_torch(device) key_to_tensor_list = defaultdict(list) # torch.no_grad() mitigates GPU memory issues # https://github.com/apache/beam/issues/22811 with torch.no_grad(): for example in batch: for key, tensor in example.items(): key_to_tensor_list[key].append(tensor) key_to_batched_tensors = {} for key in key_to_tensor_list: batched_tensors = torch.stack(key_to_tensor_list[key]) batched_tensors = _convert_to_device(batched_tensors, device) key_to_batched_tensors[key] = batched_tensors predictions = model(**key_to_batched_tensors, **inference_args) return utils._convert_to_result(batch, predictions, model_id) def _run_inference_tensorflow_keyed_tensor( batch: Sequence[dict[str, tf.Tensor]], model: TFAutoModel, device, inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: if device == "GPU": is_gpu_available_tensorflow(device) key_to_tensor_list = defaultdict(list) for example in batch: for key, tensor in example.items(): key_to_tensor_list[key].append(tensor) key_to_batched_tensors = {} for key in key_to_tensor_list: batched_tensors = tf.stack(key_to_tensor_list[key], axis=0) key_to_batched_tensors[key] = batched_tensors predictions = model(**key_to_batched_tensors, **inference_args) return utils._convert_to_result(batch, predictions, model_id) class HuggingFaceModelHandlerKeyedTensor(ModelHandler[dict[str, Union[tf.Tensor, torch.Tensor]], PredictionResult, Union[AutoModel, TFAutoModel]]): def __init__( self, model_uri: str, model_class: Union[AutoModel, TFAutoModel], framework: str, device: str = "CPU", *, inference_fn: Optional[Callable[..., Iterable[PredictionResult]]] = None, load_model_args: Optional[dict[str, Any]] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, large_model: bool = False, model_copies: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, **kwargs): """ Implementation of the ModelHandler interface for HuggingFace with Keyed Tensors for PyTorch/Tensorflow backend. Example Usage model:: pcoll | RunInference(HuggingFaceModelHandlerKeyedTensor( model_uri="bert-base-uncased", model_class=AutoModelForMaskedLM, framework='pt')) Args: model_uri (str): path to the pretrained model on the hugging face models hub. model_class: model class to load the repository from model_uri. framework (str): Framework to use for the model. 'tf' for TensorFlow and 'pt' for PyTorch. device: For torch tensors, specify device on which you wish to run the model. Defaults to CPU. inference_fn: the inference function to use during RunInference. Default is _run_inference_torch_keyed_tensor or _run_inference_tensorflow_keyed_tensor depending on the input type. load_model_args (dict[str, Any]): (Optional) Keyword arguments to provide load options while loading models from Hugging Face Hub. Defaults to None. min_batch_size: the minimum batch size to use when batching inputs. max_batch_size: the maximum batch size to use when batching inputs. max_batch_duration_secs: the maximum amount of time to buffer a batch before emitting; used in streaming contexts. large_model: set to true if your model is large enough to run into memory pressure if you load multiple copies. Given a model that consumes N memory and a machine with W cores and M memory, you should set this to True if N*W > M. model_copies: The exact number of models that you would like loaded onto your machine. This can be useful if you exactly know your CPU or GPU capacity and want to maximize resource utilization. max_batch_weight: the maximum total weight of a batch. element_size_fn: a function that returns the size (weight) of an element. batch_length_fn: a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: a sorted list of positive boundary values for length-aware batching buckets. kwargs: 'env_vars' can be used to set environment variables before loading the model. **Supported Versions:** HuggingFaceModelHandler supports transformers>=4.18.0. """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries, large_model=large_model, model_copies=model_copies, **kwargs) self._model_uri = model_uri self._model_class = model_class self._device = device self._inference_fn = inference_fn self._model_config_args = load_model_args if load_model_args else {} self._framework = framework _validate_constructor_args( model_uri=self._model_uri, model_class=self._model_class) def load_model(self): """Loads and initializes the model for processing.""" model = self._model_class.from_pretrained( self._model_uri, **self._model_config_args) if self._framework == 'pt': if self._device == "GPU" and is_gpu_available_torch(): model.to(torch.device("cuda")) if callable(getattr(model, 'requires_grad_', None)): model.requires_grad_(False) return model def run_inference( self, batch: Sequence[dict[str, Union[tf.Tensor, torch.Tensor]]], model: Union[AutoModel, TFAutoModel], inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """ Runs inferences on a batch of Keyed Tensors and returns an Iterable of Tensors Predictions. This method stacks the list of Tensors in a vectorized format to optimize the inference call. Args: batch: A sequence of Keyed Tensors. These Tensors should be batchable, as this method will call `tf.stack()`/`torch.stack()` and pass in batched Tensors with dimensions (batch_size, n_features, etc.) into the model's predict() function. model: A Tensorflow/PyTorch model. inference_args: Non-batchable arguments required as inputs to the model's inference function. Unlike Tensors in `batch`, these parameters will not be dynamically batched. Returns: An Iterable of type PredictionResult. """ inference_args = {} if not inference_args else inference_args if self._inference_fn: return self._inference_fn( batch, model, self._device, inference_args, self._model_uri) if self._framework == "tf": return _run_inference_tensorflow_keyed_tensor( batch, model, self._device, inference_args, self._model_uri) else: return _run_inference_torch_keyed_tensor( batch, model, self._device, inference_args, self._model_uri) def update_model_path(self, model_path: Optional[str] = None): self._model_uri = model_path if model_path else self._model_uri def get_num_bytes( self, batch: Sequence[Union[tf.Tensor, torch.Tensor]]) -> int: """ Returns: The number of bytes of data for the Tensors batch. """ if self._framework == "tf": return sum(sys.getsizeof(element) for element in batch) else: return sum( (el.element_size() for tensor in batch for el in tensor.values())) def get_metrics_namespace(self) -> str: """ Returns: A namespace for metrics collected by the RunInference transform. """ return "BeamML_HuggingFaceModelHandler_KeyedTensor" def _default_inference_fn_torch( batch: Sequence[Union[tf.Tensor, torch.Tensor]], model: Union[AutoModel, TFAutoModel], device, inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: device = get_device_torch(device) # torch.no_grad() mitigates GPU memory issues # https://github.com/apache/beam/issues/22811 with torch.no_grad(): batched_tensors = torch.stack(batch) batched_tensors = _convert_to_device(batched_tensors, device) predictions = model(batched_tensors, **inference_args) return utils._convert_to_result(batch, predictions, model_id) def _default_inference_fn_tensorflow( batch: Sequence[Union[tf.Tensor, torch.Tensor]], model: Union[AutoModel, TFAutoModel], device, inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: if device == "GPU": is_gpu_available_tensorflow(device) batched_tensors = tf.stack(batch, axis=0) predictions = model(batched_tensors, **inference_args) return utils._convert_to_result(batch, predictions, model_id) class HuggingFaceModelHandlerTensor(ModelHandler[Union[tf.Tensor, torch.Tensor], PredictionResult, Union[AutoModel, TFAutoModel]]): def __init__( self, model_uri: str, model_class: Union[AutoModel, TFAutoModel], device: str = "CPU", *, inference_fn: Optional[Callable[..., Iterable[PredictionResult]]] = None, load_model_args: Optional[dict[str, Any]] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, large_model: bool = False, model_copies: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, **kwargs): """ Implementation of the ModelHandler interface for HuggingFace with Tensors for PyTorch/Tensorflow backend. Depending on the type of tensors, the model framework is determined automatically. Example Usage model: pcoll | RunInference(HuggingFaceModelHandlerTensor( model_uri="bert-base-uncased", model_class=AutoModelForMaskedLM)) Args: model_uri (str): path to the pretrained model on the hugging face models hub. model_class: model class to load the repository from model_uri. device: For torch tensors, specify device on which you wish to run the model. Defaults to CPU. inference_fn: the inference function to use during RunInference. Default is _run_inference_torch_keyed_tensor or _run_inference_tensorflow_keyed_tensor depending on the input type. load_model_args (dict[str, Any]): (Optional) keyword arguments to provide load options while loading models from Hugging Face Hub. Defaults to None. min_batch_size: the minimum batch size to use when batching inputs. max_batch_size: the maximum batch size to use when batching inputs. max_batch_duration_secs: the maximum amount of time to buffer a batch before emitting; used in streaming contexts. large_model: set to true if your model is large enough to run into memory pressure if you load multiple copies. Given a model that consumes N memory and a machine with W cores and M memory, you should set this to True if N*W > M. model_copies: The exact number of models that you would like loaded onto your machine. This can be useful if you exactly know your CPU or GPU capacity and want to maximize resource utilization. max_batch_weight: the maximum total weight of a batch. element_size_fn: a function that returns the size (weight) of an element. batch_length_fn: a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: a sorted list of positive boundary values for length-aware batching buckets. kwargs: 'env_vars' can be used to set environment variables before loading the model. **Supported Versions:** HuggingFaceModelHandler supports transformers>=4.18.0. """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries, large_model=large_model, model_copies=model_copies, **kwargs) self._model_uri = model_uri self._model_class = model_class self._device = device self._inference_fn = inference_fn self._model_config_args = load_model_args if load_model_args else {} self._framework = "" _validate_constructor_args( model_uri=self._model_uri, model_class=self._model_class) def load_model(self): """Loads and initializes the model for processing.""" model = self._model_class.from_pretrained( self._model_uri, **self._model_config_args) if callable(getattr(model, 'requires_grad_', None)): model.requires_grad_(False) return model def run_inference( self, batch: Sequence[Union[tf.Tensor, torch.Tensor]], model: Union[AutoModel, TFAutoModel], inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """ Runs inferences on a batch of Tensors and returns an Iterable of Tensors Predictions. This method stacks the list of Tensors in a vectorized format to optimize the inference call. Args: batch: A sequence of Tensors. These Tensors should be batchable, as this method will call `tf.stack()`/`torch.stack()` and pass in batched Tensors with dimensions (batch_size, n_features, etc.) into the model's predict() function. model: A Tensorflow/PyTorch model. inference_args (dict[str, Any]): Non-batchable arguments required as inputs to the model's inference function. Unlike Tensors in `batch`, these parameters will not be dynamically batched. Returns: An Iterable of type PredictionResult. """ inference_args = {} if not inference_args else inference_args if not self._framework: if isinstance(batch[0], tf.Tensor): self._framework = "tf" else: self._framework = "pt" if (self._framework == "pt" and self._device == "GPU" and is_gpu_available_torch()): model.to(torch.device("cuda")) if self._inference_fn: return self._inference_fn( batch, model, inference_args, inference_args, self._model_uri) if self._framework == "tf": return _default_inference_fn_tensorflow( batch, model, self._device, inference_args, self._model_uri) else: return _default_inference_fn_torch( batch, model, self._device, inference_args, self._model_uri) def update_model_path(self, model_path: Optional[str] = None): self._model_uri = model_path if model_path else self._model_uri def get_num_bytes( self, batch: Sequence[Union[tf.Tensor, torch.Tensor]]) -> int: """ Returns: The number of bytes of data for the Tensors batch. """ if self._framework == "tf": return sum(sys.getsizeof(element) for element in batch) else: return sum( (el.element_size() for tensor in batch for el in tensor.values())) def get_metrics_namespace(self) -> str: """ Returns: A namespace for metrics collected by the RunInference transform. """ return 'BeamML_HuggingFaceModelHandler_Tensor' def _default_pipeline_inference_fn( batch, pipeline, inference_args) -> Iterable[PredictionResult]: predicitons = pipeline(batch, **inference_args) return predicitons class HuggingFacePipelineModelHandler(ModelHandler[str, PredictionResult, Pipeline]): def __init__( self, task: Union[str, PipelineTask] = "", model: str = "", *, device: Optional[str] = None, inference_fn: PipelineInferenceFn = _default_pipeline_inference_fn, load_pipeline_args: Optional[dict[str, Any]] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, large_model: bool = False, model_copies: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, **kwargs): """ Implementation of the ModelHandler interface for Hugging Face Pipelines. Example Usage model:: pcoll | RunInference(HuggingFacePipelineModelHandler( task="fill-mask")) Args: task (str or enum.Enum): task supported by HuggingFace Pipelines. Accepts a string task or an enum.Enum from PipelineTask. model (str): path to the pretrained *model-id* on Hugging Face Models Hub to use custom model for the chosen task. If the `model` already defines the task then no need to specify the `task` parameter. Use the *model-id* string instead of an actual model here. Model-specific kwargs for `from_pretrained(..., **model_kwargs)` can be specified with `model_kwargs` using `load_pipeline_args`. Example Usage:: model_handler = HuggingFacePipelineModelHandler( task="text-generation", model="meta-llama/Llama-2-7b-hf", load_pipeline_args={'model_kwargs':{'quantization_map':config}}) device (str): the device (`"CPU"` or `"GPU"`) on which you wish to run the pipeline. Defaults to GPU. If GPU is not available then it falls back to CPU. You can also use advanced option like `device_map` with key-value pair as you would do in the usual Hugging Face pipeline using `load_pipeline_args`. Ex: load_pipeline_args={'device_map':auto}). inference_fn: the inference function to use during RunInference. Default is _default_pipeline_inference_fn. load_pipeline_args (dict[str, Any]): keyword arguments to provide load options while loading pipelines from Hugging Face. Defaults to None. min_batch_size: the minimum batch size to use when batching inputs. max_batch_size: the maximum batch size to use when batching inputs. max_batch_duration_secs: the maximum amount of time to buffer a batch before emitting; used in streaming contexts. large_model: set to true if your model is large enough to run into memory pressure if you load multiple copies. Given a model that consumes N memory and a machine with W cores and M memory, you should set this to True if N*W > M. model_copies: The exact number of models that you would like loaded onto your machine. This can be useful if you exactly know your CPU or GPU capacity and want to maximize resource utilization. max_batch_weight: the maximum total weight of a batch. element_size_fn: a function that returns the size (weight) of an element. batch_length_fn: a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: a sorted list of positive boundary values for length-aware batching buckets. kwargs: 'env_vars' can be used to set environment variables before loading the model. **Supported Versions:** HuggingFacePipelineModelHandler supports transformers>=4.18.0. """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries, large_model=large_model, model_copies=model_copies, **kwargs) self._task = task self._model = model self._inference_fn = inference_fn self._load_pipeline_args = load_pipeline_args if load_pipeline_args else {} self._framework = "pt" # Check if the device is specified twice. If true then the device parameter # of model handler is overridden. self._deduplicate_device_value(device) _validate_constructor_args_hf_pipeline(self._task, self._model) def _deduplicate_device_value(self, device: Optional[str]): current_device = device.upper() if device else None if (current_device and current_device != 'CPU' and current_device != 'GPU'): raise ValueError( f"Invalid device value: {device}. Please specify " "either CPU or GPU. Defaults to GPU if no value " "is provided.") if 'device' not in self._load_pipeline_args: if current_device == 'CPU': self._load_pipeline_args['device'] = 'cpu' else: if is_gpu_available_torch(): self._load_pipeline_args['device'] = 'cuda:0' else: _LOGGER.warning( "HuggingFaceModelHandler specified a 'GPU' device, " "but GPUs are not available. Switching to CPU.") self._load_pipeline_args['device'] = 'cpu' else: if current_device: raise ValueError( '`device` specified in `load_pipeline_args`. `device` ' 'parameter for HuggingFacePipelineModelHandler will be ignored.') def load_model(self): """Loads and initializes the pipeline for processing.""" return pipeline( task=self._task, model=self._model, **self._load_pipeline_args) def run_inference( self, batch: Sequence[str], pipeline: Pipeline, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """ Runs inferences on a batch of examples passed as a string resource. These can either be string sentences, or string path to images or audio files. Args: batch: A sequence of strings resources. pipeline: A Hugging Face Pipeline. inference_args: Non-batchable arguments required as inputs to the model's inference function. Returns: An Iterable of type PredictionResult. """ inference_args = {} if not inference_args else inference_args predictions = self._inference_fn(batch, pipeline, inference_args) return utils._convert_to_result(batch, predictions) def update_model_path(self, model_path: Optional[str] = None): """ Updates the pretrained model used by the Hugging Face Pipeline task. Make sure that the new model does the same task as initial model. Args: model_path (str): (Optional) Path to the new trained model from Hugging Face. Defaults to None. """ self._model = model_path if model_path else self._model def get_num_bytes(self, batch: Sequence[str]) -> int: """ Returns: The number of bytes of input batch elements. """ return sum(sys.getsizeof(element) for element in batch) def get_metrics_namespace(self) -> str: """ Returns: A namespace for metrics collected by the RunInference transform. """ return 'BeamML_HuggingFacePipelineModelHandler' ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/huggingface_inference_it_test.py0000644000175100017510000001350015206113032026614 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-End test for Hugging Face Inference""" import logging import unittest import uuid import pytest from apache_beam.io.filesystems import FileSystems from apache_beam.testing.test_pipeline import TestPipeline try: from apache_beam.examples.inference import huggingface_language_modeling from apache_beam.examples.inference import huggingface_question_answering from apache_beam.ml.inference import pytorch_inference_it_test except ImportError: raise unittest.SkipTest( "transformers dependencies are not installed. " "Check if transformers, torch, and tensorflow " "is installed.") @pytest.mark.uses_transformers @pytest.mark.it_postcommit @pytest.mark.timeout(1800) class HuggingFaceInference(unittest.TestCase): def test_hf_language_modeling(self): test_pipeline = TestPipeline(is_integration_test=True) # Path to text file containing some sentences file_of_sentences = 'gs://apache-beam-ml/datasets/custom/hf_sentences.txt' output_file_dir = 'gs://apache-beam-ml/testing/predictions' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_name = 'stevhliu/my_awesome_eli5_mlm_model' extra_opts = { 'input': file_of_sentences, 'output': output_file, 'model_name': model_name, } huggingface_language_modeling.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) predictions = pytorch_inference_it_test.process_outputs( filepath=output_file) actuals_file = 'gs://apache-beam-ml/testing/expected_outputs/test_hf_run_inference_for_masked_lm_actuals.txt' # pylint: disable=line-too-long actuals = pytorch_inference_it_test.process_outputs(filepath=actuals_file) predictions_dict = {} for prediction in predictions: text, predicted_text = prediction.split(';') predictions_dict[text] = predicted_text.strip().lower() for actual in actuals: text, actual_predicted_text = actual.split(';') predicted_predicted_text = predictions_dict[text] self.assertEqual(actual_predicted_text, predicted_predicted_text) def test_hf_language_modeling_large_model(self): test_pipeline = TestPipeline(is_integration_test=True) # Path to text file containing some sentences file_of_sentences = 'gs://apache-beam-ml/datasets/custom/hf_sentences.txt' output_file_dir = 'gs://apache-beam-ml/testing/predictions' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_name = 'stevhliu/my_awesome_eli5_mlm_model' extra_opts = { 'input': file_of_sentences, 'output': output_file, 'model_name': model_name, 'large_model': True, } huggingface_language_modeling.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) predictions = pytorch_inference_it_test.process_outputs( filepath=output_file) actuals_file = 'gs://apache-beam-ml/testing/expected_outputs/test_hf_run_inference_for_masked_lm_actuals.txt' # pylint: disable=line-too-long actuals = pytorch_inference_it_test.process_outputs(filepath=actuals_file) predictions_dict = {} for prediction in predictions: text, predicted_text = prediction.split(';') predictions_dict[text] = predicted_text.strip().lower() for actual in actuals: text, actual_predicted_text = actual.split(';') predicted_predicted_text = predictions_dict[text] self.assertEqual(actual_predicted_text, predicted_predicted_text) def test_hf_pipeline(self): test_pipeline = TestPipeline(is_integration_test=True) # Path to text file containing some questions and context input_file = 'gs://apache-beam-ml/datasets/custom/questions.txt' output_file_dir = 'gs://apache-beam-ml/hf/testing/predictions' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) extra_opts = { 'input': input_file, 'output': output_file, 'revision': 'deedc3e42208524e0df3d9149d1f26aa6934f05f', } huggingface_question_answering.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) predictions = pytorch_inference_it_test.process_outputs( filepath=output_file) actuals_file = ( 'gs://apache-beam-ml/testing/expected_outputs/' 'test_hf_pipeline_answers.txt') actuals = pytorch_inference_it_test.process_outputs(filepath=actuals_file) predictions_dict = {} for prediction in predictions: text, predicted_text = prediction.split(';') predictions_dict[text] = predicted_text.strip() for actual in actuals: text, actual_predicted_text = actual.split(';') predicted_predicted_text = predictions_dict[text] self.assertEqual(actual_predicted_text, predicted_predicted_text) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/huggingface_inference_test.py0000644000175100017510000001307315206113032026125 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import shutil import tempfile import unittest from collections.abc import Iterable from collections.abc import Sequence from typing import Any from typing import Optional from typing import Union import pytest from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.tensorflow_inference_test import FakeTFTensorModel from apache_beam.ml.inference.tensorflow_inference_test import _compare_tensor_prediction_result # pylint: disable=ungrouped-imports try: import tensorflow as tf import torch from transformers import AutoModel from transformers import TFAutoModel from apache_beam.ml.inference.huggingface_inference import HuggingFaceModelHandlerTensor except ImportError: raise unittest.SkipTest('Transformers dependencies are not installed.') def fake_inference_fn_tensor( batch: Sequence[Union[tf.Tensor, torch.Tensor]], model: Union[AutoModel, TFAutoModel], device, inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: predictions = model.predict(batch, **inference_args) return utils._convert_to_result(batch, predictions, model_id) class FakeTorchModel: def predict(self, input: torch.Tensor): return input @pytest.mark.uses_transformers class HuggingFaceInferenceTest(unittest.TestCase): def setUp(self) -> None: self.tmpdir = tempfile.mkdtemp() def tearDown(self) -> None: shutil.rmtree(self.tmpdir) def test_predict_tensor(self): fake_model = FakeTFTensorModel() inference_runner = HuggingFaceModelHandlerTensor( model_uri='unused', model_class=TFAutoModel, inference_fn=fake_inference_fn_tensor) batched_examples = [tf.constant([1]), tf.constant([10]), tf.constant([100])] expected_predictions = [ PredictionResult(ex, pred) for ex, pred in zip( batched_examples, [tf.math.multiply(n, 10) for n in batched_examples]) ] inferences = inference_runner.run_inference(batched_examples, fake_model) for actual, expected in zip(inferences, expected_predictions): self.assertTrue(_compare_tensor_prediction_result(actual, expected)) def test_predict_tensor_with_inference_args(self): fake_model = FakeTFTensorModel() inference_runner = HuggingFaceModelHandlerTensor( model_uri='unused', model_class=TFAutoModel, inference_fn=fake_inference_fn_tensor, inference_args={"add": True}) batched_examples = [tf.constant([1]), tf.constant([10]), tf.constant([100])] expected_predictions = [ PredictionResult(ex, pred) for ex, pred in zip( batched_examples, [ tf.math.add(tf.math.multiply(n, 10), 10) for n in batched_examples ]) ] inferences = inference_runner.run_inference( batched_examples, fake_model, inference_args={"add": True}) for actual, expected in zip(inferences, expected_predictions): self.assertTrue(_compare_tensor_prediction_result(actual, expected)) def test_framework_detection_torch(self): fake_model = FakeTorchModel() inference_runner = HuggingFaceModelHandlerTensor( model_uri='unused', model_class=TFAutoModel, inference_fn=fake_inference_fn_tensor) batched_examples = [torch.tensor(1), torch.tensor(10), torch.tensor(100)] inference_runner.run_inference(batched_examples, fake_model) self.assertEqual(inference_runner._framework, "pt") def test_framework_detection_tensorflow(self): fake_model = FakeTFTensorModel() inference_runner = HuggingFaceModelHandlerTensor( model_uri='unused', model_class=TFAutoModel, inference_fn=fake_inference_fn_tensor) batched_examples = [tf.constant(1), tf.constant(10), tf.constant(100)] inference_runner.run_inference(batched_examples, fake_model) self.assertEqual(inference_runner._framework, 'tf') def test_convert_to_result_batch_processing(self): """Test that utils._convert_to_result correctly handles batches with multiple elements.""" # Test case that reproduces the bug: batch size > 1 batch = ["input1", "input2"] predictions = [{ "translation_text": "output1" }, { "translation_text": "output2" }] results = list(utils._convert_to_result(batch, predictions)) # Should return 2 results, not 1 self.assertEqual( len(results), 2, "Should return one result per batch element") # Check that each result has the correct input and output self.assertEqual(results[0].example, "input1") self.assertEqual(results[0].inference, {"translation_text": "output1"}) self.assertEqual(results[1].example, "input2") self.assertEqual(results[1].inference, {"translation_text": "output2"}) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/model_manager.py0000644000175100017510000006622215206113032023377 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for managing ML models in Apache Beam pipelines. This module provides classes and functions to efficiently manage multiple machine learning models within Apache Beam pipelines. It includes functionality for loading, caching, and updating models using multi-process shared memory, ensuring that models are reused across different workers to optimize resource usage and performance. """ import gc import heapq import itertools import logging import subprocess import threading import time from collections import Counter from collections import OrderedDict from collections import defaultdict from collections import deque from typing import Any from typing import Callable from typing import Optional import numpy as np import torch from scipy.optimize import nnls from apache_beam.metrics.metric import Metrics from apache_beam.utils import multi_process_shared logger = logging.getLogger(__name__) _MODEL_MANAGER_METRICS_NAMESPACE = "BeamML_ModelManager" class GPUMonitor: """Monitors GPU memory usage in a separate thread using nvidia-smi. This class continuously polls GPU memory statistics to track current usage and peak usage over a sliding time window. It serves as the source of truth for the ModelManager's resource decisions. Attributes: fallback_memory_mb: Default total memory if hardware detection fails. poll_interval: Seconds between memory checks. peak_window_seconds: Duration to track peak memory usage. """ def __init__( self, fallback_memory_mb: float = 16000.0, poll_interval: float = 0.5, peak_window_seconds: float = 30.0): self._current_usage = 0.0 self._peak_usage = 0.0 self._total_memory = fallback_memory_mb self._poll_interval = poll_interval self._peak_window_seconds = peak_window_seconds self._memory_history = deque() self._running = False self._thread = None self._lock = threading.Lock() def _detect_hardware(self): try: cmd = [ "nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits" ] output = subprocess.check_output(cmd, text=True).strip() self._total_memory = float(output) return True except (FileNotFoundError, subprocess.CalledProcessError): logger.warning( "nvidia-smi not found or failed. Defaulting total memory to %s MB", self._total_memory) return False except Exception as e: logger.warning( "Error parsing nvidia-smi output: %s. " "Defaulting total memory to %s MB", e, self._total_memory) return False def start(self): self._gpu_available = self._detect_hardware() if self._running or not self._gpu_available: return self._running = True self._thread = threading.Thread(target=self._poll_loop, daemon=True) self._thread.start() def stop(self): self._running = False if self._thread: self._thread.join() def reset_peak(self): with self._lock: now = time.time() self._memory_history.clear() self._memory_history.append((now, self._current_usage)) self._peak_usage = self._current_usage def get_stats(self) -> tuple[float, float, float]: with self._lock: return self._current_usage, self._peak_usage, self._total_memory def refresh(self): """Forces an immediate poll of the GPU.""" usage = self._get_nvidia_smi_used() now = time.time() with self._lock: self._current_usage = usage self._memory_history.append((now, usage)) # Recalculate peak immediately while self._memory_history and (now - self._memory_history[0][0] > self._peak_window_seconds): self._memory_history.popleft() self._peak_usage = ( max(m for _, m in self._memory_history) if self._memory_history else usage) def _get_nvidia_smi_used(self) -> float: try: cmd = [ "nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits" ] output = subprocess.check_output(cmd, text=True).strip() free_memory = float(output) return self._total_memory - free_memory except Exception as e: logger.warning('Failed to get GPU memory usage: %s', e) return 0.0 def _poll_loop(self): while self._running: usage = self._get_nvidia_smi_used() now = time.time() with self._lock: self._current_usage = usage self._memory_history.append((now, usage)) while self._memory_history and (now - self._memory_history[0][0] > self._peak_window_seconds): self._memory_history.popleft() self._peak_usage = ( max(m for _, m in self._memory_history) if self._memory_history else usage) time.sleep(self._poll_interval) class ResourceEstimator: """Estimates individual model memory usage using statistical observation. Uses Non-Negative Least Squares (NNLS) to deduce the memory footprint of individual models based on aggregate system memory readings and the configuration of active models at that time. """ def __init__( self, smoothing_factor: float = 0.2, min_data_points: int = 5, verbose_logging: bool = False): self.smoothing_factor = smoothing_factor self.min_data_points = min_data_points self.verbose_logging = verbose_logging self.estimates: dict[str, float] = {} self.history = defaultdict(lambda: deque(maxlen=20)) self.known_models = set() self._lock = threading.Lock() def logging_info(self, message: str, *args): if self.verbose_logging: logger.info(message, *args) def is_unknown(self, model_tag: str) -> bool: with self._lock: return model_tag not in self.estimates def get_estimate(self, model_tag: str, default_mb: float = 4000.0) -> float: with self._lock: return self.estimates.get(model_tag, default_mb) def set_initial_estimate(self, model_tag: str, cost: float): with self._lock: self.estimates[model_tag] = cost self.known_models.add(model_tag) Metrics.distribution( _MODEL_MANAGER_METRICS_NAMESPACE, f"memory_estimate_mb_{model_tag}", process_wide=True).update(int(cost)) self.logging_info("Initial Profile for %s: %s MB", model_tag, cost) def add_observation( self, active_snapshot: dict[str, int], peak_memory: float): if active_snapshot: model_list = "\n".join( f"\t- {model}: {count}" for model, count in sorted(active_snapshot.items())) else: model_list = "\t- None" self.logging_info( "Adding Observation:\n PeakMemory: %.1f MB\n Instances:\n%s", peak_memory, model_list) if not active_snapshot: return with self._lock: config_key = tuple(sorted(active_snapshot.items())) self.history[config_key].append(peak_memory) for tag in active_snapshot: self.known_models.add(tag) self._solve() def _solve(self): """ Solves Ax=b using raw readings (no pre-averaging) and NNLS. This creates a 'tall' matrix A where every memory reading is a separate equation. """ unique = sorted(list(self.known_models)) # We need to build the matrix first to know if we have enough data points A, b = [], [] for config_key, mem_values in self.history.items(): if not mem_values: continue # 1. Create the feature row for this configuration ONCE # (It represents the model counts + bias) counts = dict(config_key) feature_row = [counts.get(model, 0) for model in unique] feature_row.append(1) # Bias column # 2. Add a separate row to the matrix for EVERY individual reading # Instead of averaging, we flatten the history into the matrix for reading in mem_values: A.append(feature_row) # The inputs (models) stay the same b.append(reading) # The output (memory) varies due to noise # Convert to numpy for SciPy A = np.array(A) b = np.array(b) if len( self.history.keys()) < len(unique) + 1 or len(A) < self.min_data_points: # Not enough data to solve yet return self.logging_info( "Solving with %s total observations for %s models.", len(A), len(unique)) try: # Solve using Non-Negative Least Squares # x will be >= 0 x, _ = nnls(A, b) weights = x[:-1] bias = x[-1] for i, model in enumerate(unique): calculated_cost = weights[i] if model in self.estimates: old = self.estimates[model] new = (old * (1 - self.smoothing_factor)) + ( calculated_cost * self.smoothing_factor) self.estimates[model] = new else: self.estimates[model] = calculated_cost self.logging_info( "Updated Estimate for %s: %.1f MB", model, self.estimates[model]) Metrics.distribution( _MODEL_MANAGER_METRICS_NAMESPACE, f"memory_estimate_mb_{model}", process_wide=True).update(int(self.estimates[model])) self.logging_info("System Bias: %s MB", bias) except Exception as e: logger.error("Solver failed: %s", e) class QueueTicket: def __init__(self, priority, ticket_num, tag): self.priority = priority self.ticket_num = ticket_num self.tag = tag self.wake_event = threading.Event() def __lt__(self, other): return (self.priority, self.ticket_num) < (other.priority, other.ticket_num) class ModelManager: """Manages model lifecycles, caching, and resource arbitration. This class acts as the central controller for acquiring model instances. 1. LRU Caching of idle models. 2. Resource estimation and admission control (preventing OOM). 3. Dynamic eviction of low-priority models, determined by count of pending requests, when space is needed. 4. 'Isolation Mode' for safely profiling unknown models. """ def __init__( self, monitor: Optional['GPUMonitor'] = None, slack_percentage: float = 0.10, poll_interval: float = 0.5, peak_window_seconds: float = 30.0, min_data_points: int = 5, smoothing_factor: float = 0.2, eviction_cooldown_seconds: float = 10.0, min_model_copies: int = 1, wait_timeout_seconds: float = 300.0, lock_timeout_seconds: float = 60.0, verbose_logging: bool = False): self._estimator = ResourceEstimator( min_data_points=min_data_points, smoothing_factor=smoothing_factor, verbose_logging=verbose_logging) self._monitor = monitor if monitor else GPUMonitor( poll_interval=poll_interval, peak_window_seconds=peak_window_seconds) self._slack_percentage = slack_percentage self._eviction_cooldown = eviction_cooldown_seconds self._min_model_copies = min_model_copies self._wait_timeout_seconds = wait_timeout_seconds self._lock_timeout_seconds = lock_timeout_seconds self._verbose_logging = verbose_logging # Resource State self._models = defaultdict(list) # Idle LRU used to track released models that # can be freed or reused upon request. self._idle_lru = OrderedDict() self._active_counts = Counter() self._total_active_jobs = 0 self._pending_reservations = 0.0 # Isolation state used to profile unknown models, # ensuring they run alone to get accurate readings. # isolation_baseline represents the GPU usage before # loading the unknown model. self._isolation_mode = False self._isolation_baseline = 0.0 # Waiting Queue and Ticketing to make sure we have fair ordering # and also priority for unknown models. self._wait_queue = [] self._ticket_counter = itertools.count() self._cancelled_tickets = set() # TODO: Consider making the wait to be smarter, i.e. # splitting read/write etc. to avoid potential contention. self._cv = threading.Condition() self._monitor.start() def _update_model_count_metric(self): for tag, instances in self._models.items(): Metrics.distribution( _MODEL_MANAGER_METRICS_NAMESPACE, f"num_loaded_models_{tag}", process_wide=True).update(len(instances)) def _clear_all_model_metrics(self): for tag in self._models: Metrics.distribution( _MODEL_MANAGER_METRICS_NAMESPACE, f"num_loaded_models_{tag}", process_wide=True).update(0) def logging_info(self, message: str, *args): if self._verbose_logging: logger.info(message, *args) def all_models(self, tag) -> list[Any]: return self._models[tag] # Should hold _cv lock when calling def try_enter_isolation_mode(self, tag: str, ticket_num: int) -> bool: if self._total_active_jobs > 0: self.logging_info( "Waiting to enter isolation: tag=%s ticket num=%s", tag, ticket_num) self._cv.wait(timeout=self._lock_timeout_seconds) # return False since we have waited and need to re-evaluate # in caller to make sure our priority is still valid. return False self.logging_info("Unknown model %s detected. Flushing GPU.", tag) self._delete_all_models() self._isolation_mode = True self._total_active_jobs += 1 self._isolation_baseline, _, _ = self._monitor.get_stats() self._monitor.reset_peak() return True # Should hold _cv lock when calling def should_spawn_model(self, tag: str, ticket_num: int) -> bool: curr, _, total = self._monitor.get_stats() est_cost = self._estimator.get_estimate(tag) limit = total * (1 - self._slack_percentage) # Use current usage for capacity check (ignore old spikes) if (curr + self._pending_reservations + est_cost) <= limit: self._pending_reservations += est_cost self._total_active_jobs += 1 self._active_counts[tag] += 1 return True # Evict to make space (passing tag to check demand/existence) if self._evict_to_make_space(limit, est_cost, requesting_tag=tag): return True # Manually log status for debugging if we are going to wait idle_count = 0 other_idle_count = 0 for item in self._idle_lru.items(): if item[1][0] == tag: idle_count += 1 else: other_idle_count += 1 total_model_count = 0 for _, instances in self._models.items(): total_model_count += len(instances) curr, _, _ = self._monitor.get_stats() self.logging_info( "Waiting for resources to free up: " "tag=%s ticket num%s model count=%s " "idle count=%s resource usage=%.1f MB " "total models count=%s other idle=%s", tag, ticket_num, len(self._models[tag]), idle_count, curr, total_model_count, other_idle_count) # Wait since we couldn't make space and # added timeout to avoid missed notify call. self._cv.wait(timeout=self._lock_timeout_seconds) return False def _wake_next_in_queue(self): if self._wait_queue: # Clean up cancelled tickets at head of queue while self._wait_queue and self._wait_queue[ 0].ticket_num in self._cancelled_tickets: self._cancelled_tickets.remove(self._wait_queue[0].ticket_num) heapq.heappop(self._wait_queue) next_inline = self._wait_queue[0] next_inline.wake_event.set() def _wait_in_queue(self, ticket: QueueTicket): self._cv.release() try: ticket.wake_event.wait(timeout=self._lock_timeout_seconds) ticket.wake_event.clear() finally: self._cv.acquire() def acquire_model(self, tag: str, loader_func: Callable[[], Any]) -> Any: current_priority = 0 if self._estimator.is_unknown(tag) else 1 ticket_num = next(self._ticket_counter) my_ticket = QueueTicket(current_priority, ticket_num, tag) with self._cv: # FAST PATH: Grab from idle LRU if available if not self._isolation_mode: cached_instance = self._try_grab_from_lru(tag) if cached_instance: return cached_instance # SLOW PATH: Enqueue and wait for turn to acquire model, # with unknown models having priority and order enforced # by ticket number as FIFO. self.logging_info( "Acquire Queued: tag=%s, priority=%d " "total models count=%s ticket num=%s", tag, current_priority, len(self._models[tag]), ticket_num) heapq.heappush(self._wait_queue, my_ticket) est_cost = 0.0 is_unknown = False wait_time_start = time.time() try: while True: wait_time_elapsed = time.time() - wait_time_start if wait_time_elapsed > self._wait_timeout_seconds: raise RuntimeError( f"Timeout waiting to acquire model: {tag} " f"after {wait_time_elapsed:.1f} seconds.") if not self._wait_queue or self._wait_queue[ 0].ticket_num != ticket_num: self.logging_info( "Waiting for its turn: tag=%s ticket num=%s", tag, ticket_num) self._wait_in_queue(my_ticket) continue # Re-evaluate priority in case model became known during wait is_unknown = self._estimator.is_unknown(tag) real_priority = 0 if is_unknown else 1 # If priority changed, reinsert into queue and wait if current_priority != real_priority: heapq.heappop(self._wait_queue) current_priority = real_priority my_ticket = QueueTicket(current_priority, ticket_num, tag) heapq.heappush(self._wait_queue, my_ticket) self._wake_next_in_queue() continue # Try grab from LRU again in case model was released during wait cached_instance = self._try_grab_from_lru(tag) if cached_instance: return cached_instance # Path A: Isolation if is_unknown: if self.try_enter_isolation_mode(tag, ticket_num): # We got isolation, can proceed to spawn break else: # We waited, need to re-evaluate our turn # because priority may have changed during the wait continue # Path B: Concurrent else: if self._isolation_mode: self.logging_info( "Waiting due to isolation in progress: tag=%s ticket num%s", tag, ticket_num) self._wait_in_queue(my_ticket) continue if self.should_spawn_model(tag, ticket_num): est_cost = self._estimator.get_estimate(tag) # We can proceed to spawn since we have resources break else: # We waited, need to re-evaluate our turn # because priority may have changed during the wait continue finally: # Remove self from wait queue once done if self._wait_queue and self._wait_queue[0].ticket_num == ticket_num: heapq.heappop(self._wait_queue) else: # Marked as cancelled so that we skip when we reach head later self._cancelled_tickets.add(ticket_num) self._wake_next_in_queue() return self._spawn_new_model(tag, loader_func, is_unknown, est_cost) def release_model(self, tag: str, instance: Any): with self._cv: try: self._total_active_jobs -= 1 if self._active_counts[tag] > 0: self._active_counts[tag] -= 1 self._idle_lru[id(instance)] = (tag, instance, time.time()) # Update estimator with latest stats _, peak_during_job, _ = self._monitor.get_stats() if self._isolation_mode and self._active_counts[tag] == 0: # For isolation mode, we directly set the initial estimate # so that we can quickly learn the model cost. cost = max(0, peak_during_job - self._isolation_baseline) self._estimator.set_initial_estimate(tag, cost) self._isolation_mode = False self._isolation_baseline = 0.0 else: # Regular update for known models snapshot = { t: len(instances) for t, instances in self._models.items() if len(instances) > 0 } if snapshot: self._estimator.add_observation(snapshot, peak_during_job) finally: self._wake_next_in_queue() self._cv.notify_all() def _try_grab_from_lru(self, tag: str) -> Any: target_key = None target_instance = None for key, (t, instance, _) in reversed(self._idle_lru.items()): if t == tag: target_key = key target_instance = instance break if target_instance: # Found an idle model, remove from LRU and return del self._idle_lru[target_key] self._active_counts[tag] += 1 self._total_active_jobs += 1 return target_instance self.logging_info("No idle model found for tag: %s", tag) return None def _evict_to_make_space( self, limit: float, est_cost: float, requesting_tag: str) -> bool: """ Evicts models based on Demand Magnitude + Tiers. Crucially: If we have 0 active copies of 'requesting_tag', we FORCE eviction of the lowest-demand candidate to avoid starvation. Returns True if space was made, False otherwise. """ curr, _, _ = self._monitor.get_stats() projected_usage = curr + self._pending_reservations + est_cost if projected_usage <= limit: # Memory usage changed and we are already under limit return True now = time.time() # Calculate the demand from the wait queue # TODO: Also factor in the active counts to avoid thrashing demand_map = Counter() for item in self._wait_queue: demand_map[item.tag] += 1 my_demand = demand_map[requesting_tag] am_i_starving = len(self._models[requesting_tag]) == 0 candidates = [] for key, (tag, instance, release_time) in self._idle_lru.items(): candidate_demand = demand_map[tag] # TODO: Try to avoid churn if demand is similar if not am_i_starving and candidate_demand >= my_demand: continue # Attempts to score candidates based on hotness and manually # specified minimum copies. Demand is weighted heavily to # ensure we evict low-demand models first. age = now - release_time is_cold = age >= self._eviction_cooldown total_copies = len(self._models[tag]) is_surplus = total_copies > self._min_model_copies if is_cold and is_surplus: tier = 0 elif not is_cold and is_surplus: tier = 1 elif is_cold and not is_surplus: tier = 2 else: tier = 3 score = (candidate_demand * 10) + tier candidates.append((score, release_time, key, tag, instance)) candidates.sort(key=lambda x: (x[0], x[1])) # Evict candidates until we are under limit for score, _, key, tag, instance in candidates: if projected_usage <= limit: break if key not in self._idle_lru: continue self._perform_eviction(key, tag, instance, score) curr, _, _ = self._monitor.get_stats() projected_usage = curr + self._pending_reservations + est_cost return projected_usage <= limit def _delete_instance(self, instance: Any): if isinstance(instance, str): # If the instance is a string, it's a uuid used # to retrieve the model from MultiProcessShared try: multi_process_shared.MultiProcessShared( lambda: "N/A", tag=instance).unsafe_hard_delete() except (EOFError, OSError, BrokenPipeError): # This can happen even in normal operation. pass if hasattr(instance, 'mock_model_unsafe_hard_delete'): # Call the mock unsafe hard delete method for testing instance.mock_model_unsafe_hard_delete() del instance def _perform_eviction(self, key: str, tag: str, instance: Any, score: int): self.logging_info("Evicting Model: %s (Score %d)", tag, score) curr, _, _ = self._monitor.get_stats() self.logging_info("Resource Usage Before Eviction: %.1f MB", curr) if key in self._idle_lru: del self._idle_lru[key] for i, inst in enumerate(self._models[tag]): if instance == inst: del self._models[tag][i] break self._delete_instance(instance) gc.collect() torch.cuda.empty_cache() self._monitor.refresh() self._monitor.reset_peak() curr, _, _ = self._monitor.get_stats() self.logging_info("Resource Usage After Eviction: %.1f MB", curr) self._update_model_count_metric() def _spawn_new_model( self, tag: str, loader_func: Callable[[], Any], is_unknown: bool, est_cost: float) -> Any: try: with self._cv: self.logging_info("Loading Model: %s (Unknown: %s)", tag, is_unknown) baseline_snap, _, _ = self._monitor.get_stats() instance = loader_func() _, peak_during_load, _ = self._monitor.get_stats() snapshot = {tag: 1} self._estimator.add_observation( snapshot, peak_during_load - baseline_snap) if not is_unknown: self._pending_reservations = max( 0.0, self._pending_reservations - est_cost) self._models[tag].append(instance) self._update_model_count_metric() return instance except Exception as e: logger.error("Load Failed: %s. Error: %s", tag, e) with self._cv: self._total_active_jobs -= 1 if is_unknown: self._isolation_mode = False self._isolation_baseline = 0.0 else: self._pending_reservations = max( 0.0, self._pending_reservations - est_cost) self._active_counts[tag] -= 1 self._cv.notify_all() raise e def _delete_all_models(self): self._clear_all_model_metrics() self._idle_lru.clear() for _, instances in self._models.items(): for instance in instances: self._delete_instance(instance) self._models.clear() self._active_counts.clear() gc.collect() torch.cuda.empty_cache() self._monitor.refresh() self._monitor.reset_peak() def _force_reset(self): logger.warning("Force Reset Triggered") self._delete_all_models() self._models = defaultdict(list) self._idle_lru = OrderedDict() self._active_counts = Counter() self._wait_queue = [] self._total_active_jobs = 0 self._pending_reservations = 0.0 self._isolation_mode = False self._isolation_baseline = 0.0 def shutdown(self): self._delete_all_models() self._monitor.stop() def __del__(self): self.shutdown() def __exit__(self, exc_type, exc_value, traceback): self.shutdown() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/model_manager_it_test.py0000644000175100017510000001555715206113032025137 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest import apache_beam as beam from apache_beam.ml.inference.base import RunInference from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # pylint: disable=ungrouped-imports try: import torch from apache_beam.ml.inference.huggingface_inference import HuggingFacePipelineModelHandler except ImportError as e: raise unittest.SkipTest( "HuggingFace model handler dependencies are not installed") class HuggingFaceGpuTest(unittest.TestCase): # Skips the test if you run it on a machine without a GPU @unittest.skipIf( not torch.cuda.is_available(), "No GPU detected, skipping GPU test") def test_sentiment_analysis_on_gpu_large_input(self): """ Runs inference on a GPU (device=0) with a larger set of inputs. """ model_handler = HuggingFacePipelineModelHandler( task="sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english", device=0, inference_args={"batch_size": 4}) DUPLICATE_FACTOR = 2 with TestPipeline() as pipeline: examples = [ "I absolutely love this product, it's a game changer!", "This is the worst experience I have ever had.", "The weather is okay, but I wish it were sunnier.", "Apache Beam makes parallel processing incredibly efficient.", "I am extremely disappointed with the service.", "Logic and reason are the pillars of good debugging.", "I'm so happy today!", "This error message is confusing and unhelpful.", "The movie was fantastic and the acting was superb.", "I hate waiting in line for so long." ] * DUPLICATE_FACTOR pcoll = pipeline | 'CreateInputs' >> beam.Create(examples) predictions = pcoll | 'RunInference' >> RunInference( model_handler, use_model_manager=True) actual_labels = predictions | beam.Map(lambda x: x.inference['label']) expected_labels = [ 'POSITIVE', # "love this product" 'NEGATIVE', # "worst experience" 'NEGATIVE', # "weather is okay, but..." 'POSITIVE', # "incredibly efficient" 'NEGATIVE', # "disappointed" 'POSITIVE', # "pillars of good debugging" 'POSITIVE', # "so happy" 'NEGATIVE', # "confusing and unhelpful" 'POSITIVE', # "fantastic" 'NEGATIVE' # "hate waiting" ] * DUPLICATE_FACTOR assert_that( actual_labels, equal_to(expected_labels), label='CheckPredictions') @unittest.skipIf(not torch.cuda.is_available(), "No GPU detected") def test_sentiment_analysis_large_roberta_gpu(self): """ Runs inference using a Large architecture (RoBERTa-Large, ~355M params). This tests if the GPU can handle larger weights and requires more VRAM. """ model_handler = HuggingFacePipelineModelHandler( task="sentiment-analysis", model="Siebert/sentiment-roberta-large-english", device=0, inference_args={"batch_size": 2}) DUPLICATE_FACTOR = 2 with TestPipeline() as pipeline: examples = [ "I absolutely love this product, it's a game changer!", "This is the worst experience I have ever had.", "Apache Beam scales effortlessly to massive datasets.", "I am somewhat annoyed by the delay.", "The nuanced performance of this large model is impressive.", "I regret buying this immediately.", "The sunset looks beautiful tonight.", "This documentation is sparse and misleading.", "Winning the championship felt surreal.", "I'm feeling very neutral about this whole situation." ] * DUPLICATE_FACTOR pcoll = pipeline | 'CreateInputs' >> beam.Create(examples) predictions = pcoll | 'RunInference' >> RunInference( model_handler, use_model_manager=True) actual_labels = predictions | beam.Map(lambda x: x.inference['label']) expected_labels = [ 'POSITIVE', # love 'NEGATIVE', # worst 'POSITIVE', # scales effortlessly 'NEGATIVE', # annoyed 'POSITIVE', # impressive 'NEGATIVE', # regret 'POSITIVE', # beautiful 'NEGATIVE', # misleading 'POSITIVE', # surreal 'NEGATIVE' # "neutral" ] * DUPLICATE_FACTOR assert_that( actual_labels, equal_to(expected_labels), label='CheckPredictionsLarge') @unittest.skipIf(not torch.cuda.is_available(), "No GPU detected") def test_parallel_inference_branches(self): """ Tests a branching pipeline where one input source feeds two RunInference transforms running in parallel. Topology: [ Input Data ] | +--------+--------+ | | [ Translation ] [ Sentiment ] """ translator_handler = HuggingFacePipelineModelHandler( task="translation_en_to_es", model="Helsinki-NLP/opus-mt-en-es", device=0, inference_args={"batch_size": 8}) sentiment_handler = HuggingFacePipelineModelHandler( task="sentiment-analysis", model="nlptown/bert-base-multilingual-uncased-sentiment", device=0, inference_args={"batch_size": 8}) base_examples = [ "I love this product.", "This is terrible.", "Hello world.", "The service was okay.", "I am extremely angry." ] MULTIPLIER = 10 examples = base_examples * MULTIPLIER with TestPipeline() as pipeline: inputs = pipeline | 'CreateInputs' >> beam.Create(examples) _ = ( inputs | 'RunTranslation' >> RunInference( translator_handler, use_model_manager=True) | 'ExtractSpanish' >> beam.Map(lambda x: x.inference['translation_text'])) _ = ( inputs | 'RunSentiment' >> RunInference( sentiment_handler, use_model_manager=True) | 'ExtractLabel' >> beam.Map(lambda x: x.inference['label'])) if __name__ == "__main__": unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/model_manager_test.py0000644000175100017510000005645615206113032024446 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import random import threading import time import unittest from concurrent.futures import ThreadPoolExecutor from concurrent.futures import TimeoutError from unittest.mock import patch from apache_beam.utils import multi_process_shared try: from apache_beam.metrics.execution import MetricsEnvironment from apache_beam.metrics.metricbase import MetricName from apache_beam.ml.inference.model_manager import GPUMonitor from apache_beam.ml.inference.model_manager import ModelManager from apache_beam.ml.inference.model_manager import ResourceEstimator except ImportError as e: raise unittest.SkipTest("Model Manager dependencies are not installed") class MockGPUMonitor: """ Simulates GPU hardware with cumulative memory tracking. Allows simulating specific allocation spikes and baseline usage. """ def __init__(self, total_memory=12000.0, peak_window: int = 5): self._current = 0.0 self._peak = 0.0 self._total = total_memory self._lock = threading.Lock() self.running = False self.history = [] self.peak_window = peak_window def start(self): self.running = True def stop(self): self.running = False def get_stats(self): with self._lock: return self._current, self._peak, self._total def reset_peak(self): with self._lock: self._peak = self._current self.history = [self._current] def set_usage(self, current_mb): """Sets absolute usage (legacy helper).""" with self._lock: self._current = current_mb self._peak = max(self._peak, current_mb) def allocate(self, amount_mb): """Simulates memory allocation (e.g., tensors loaded to VRAM).""" with self._lock: self._current += amount_mb self.history.append(self._current) if len(self.history) > self.peak_window: self.history.pop(0) self._peak = max(self.history) def free(self, amount_mb): """Simulates memory freeing (not used often if pooling is active).""" with self._lock: self._current = max(0.0, self._current - amount_mb) self.history.append(self._current) if len(self.history) > self.peak_window: self.history.pop(0) self._peak = max(self.history) def refresh(self): """Simulates a refresh of the monitor stats (no-op for mock).""" pass class MockModel: def __init__(self, name, size, monitor): self.name = name self.size = size self.monitor = monitor self.deleted = False self.monitor.allocate(size) def mock_model_unsafe_hard_delete(self): if not self.deleted: self.monitor.free(self.size) self.deleted = True class Counter(object): def __init__(self, start=0): self.running = start self.lock = threading.Lock() def get(self): return self.running def increment(self, value=1): with self.lock: self.running += value return self.running class TestModelManager(unittest.TestCase): def setUp(self): """Force reset the Singleton ModelManager before every test.""" self.mock_monitor = MockGPUMonitor() self.manager = ModelManager(monitor=self.mock_monitor) def tearDown(self): self.manager.shutdown() def test_model_manager_deletes_multiprocessshared_instances(self): """Test that MultiProcessShared instances are deleted properly.""" model_name = "test_model_shared" tag = f"model_manager_test_{model_name}" def loader(): multi_process_shared.MultiProcessShared( lambda: Counter, tag=tag, always_proxy=True) return tag instance = self.manager.acquire_model(model_name, loader) instance_before = multi_process_shared.MultiProcessShared( Counter, tag=tag, always_proxy=True).acquire() instance_before.increment() self.assertEqual(instance_before.get(), 1) self.manager.release_model(model_name, instance) # Force delete all models self.manager._force_reset() # Verify that the MultiProcessShared instance is deleted # and the counter is reseted with self.assertRaises(Exception): instance_before.get() instance_after = multi_process_shared.MultiProcessShared( Counter, tag=tag, always_proxy=True).acquire() self.assertEqual(instance_after.get(), 0) def test_model_manager_timeout_on_acquire(self): """Test that acquiring a model times out properly.""" model_name = "timeout_model" self.manager = ModelManager( monitor=self.mock_monitor, wait_timeout_seconds=1.0, lock_timeout_seconds=1.0) def loader(): self.mock_monitor.allocate(self.mock_monitor._total) return model_name # Acquire the model in one thread to block others _ = self.manager.acquire_model(model_name, loader) def acquire_model_with_timeout(): return self.manager.acquire_model(model_name, loader) with ThreadPoolExecutor(max_workers=1000) as executor: futures = [ executor.submit(acquire_model_with_timeout) for i in range(1000) ] with self.assertRaises(RuntimeError) as context: for future in futures: future.result() self.assertIn("Timeout waiting to acquire model", str(context.exception)) # Release the initially acquired model and try to acquire again # to make sure the manager is still functional self.manager.release_model(model_name, model_name) _ = self.manager.acquire_model(model_name, loader) def test_model_manager_capacity_check(self): """ Test that the manager blocks when spawning models exceeds the limit, and unblocks when resources become available (via reuse). """ model_name = "known_model" model_cost = 3000.0 self.manager._estimator.set_initial_estimate(model_name, model_cost) acquired_refs = [] def loader(): self.mock_monitor.allocate(model_cost) return model_name # 1. Saturate GPU with 3 models (9000 MB usage) for _ in range(3): inst = self.manager.acquire_model(model_name, loader) acquired_refs.append(inst) # 2. Spawn one more (Should Block because 9000 + 3000 > Limit) def run_inference(): return self.manager.acquire_model(model_name, loader) with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(run_inference) try: future.result(timeout=0.5) self.fail("Should have blocked due to capacity") except TimeoutError: pass # 3. Release resources to unblock item_to_release = acquired_refs.pop() self.manager.release_model(model_name, item_to_release) result = future.result(timeout=2.0) self.assertIsNotNone(result) self.assertEqual(result, item_to_release) def test_model_manager_unknown_model_runs_isolated(self): """Test that a model with no history runs in isolation.""" model_name = "unknown_model_v1" self.assertTrue(self.manager._estimator.is_unknown(model_name)) def dummy_loader(): time.sleep(0.05) return "model_instance" instance = self.manager.acquire_model(model_name, dummy_loader) self.assertTrue(self.manager._isolation_mode) self.assertEqual(self.manager._total_active_jobs, 1) self.manager.release_model(model_name, instance) self.assertFalse(self.manager._isolation_mode) self.assertFalse(self.manager._estimator.is_unknown(model_name)) def test_model_manager_concurrent_execution(self): """Test that multiple small known models can run together.""" model_a = "small_model_a" model_b = "small_model_b" self.manager._estimator.set_initial_estimate(model_a, 1000.0) self.manager._estimator.set_initial_estimate(model_b, 1000.0) self.mock_monitor.set_usage(1000.0) inst_a = self.manager.acquire_model(model_a, lambda: "A") inst_b = self.manager.acquire_model(model_b, lambda: "B") self.assertEqual(self.manager._total_active_jobs, 2) self.manager.release_model(model_a, inst_a) self.manager.release_model(model_b, inst_b) self.assertEqual(self.manager._total_active_jobs, 0) def test_model_manager_concurrent_mixed_workload_convergence(self): """ Simulates a production environment with multiple model types running concurrently. Verifies that the estimator converges. """ TRUE_COSTS = {"model_small": 1500.0, "model_medium": 3000.0} def run_job(model_name): cost = TRUE_COSTS[model_name] def loader(): model = MockModel(model_name, cost, self.mock_monitor) return model instance = self.manager.acquire_model(model_name, loader) time.sleep(random.uniform(0.01, 0.05)) self.manager.release_model(model_name, instance) # Create a workload stream workload = ["model_small"] * 15 + ["model_medium"] * 15 random.shuffle(workload) with ThreadPoolExecutor(max_workers=8) as executor: futures = [executor.submit(run_job, name) for name in workload] for f in futures: f.result() est_small = self.manager._estimator.get_estimate("model_small") est_med = self.manager._estimator.get_estimate("model_medium") self.assertAlmostEqual(est_small, TRUE_COSTS["model_small"], delta=100.0) self.assertAlmostEqual(est_med, TRUE_COSTS["model_medium"], delta=100.0) def test_model_manager_oom_recovery(self): """Test that the manager recovers state if a loader crashes.""" model_name = "crasher_model" self.manager._estimator.set_initial_estimate(model_name, 1000.0) def crashing_loader(): raise RuntimeError("CUDA OOM or similar") with self.assertRaises(RuntimeError): self.manager.acquire_model(model_name, crashing_loader) self.assertEqual(self.manager._total_active_jobs, 0) self.assertEqual(self.manager._pending_reservations, 0.0) self.assertFalse(self.manager._cv._is_owned()) def test_model_manager_force_reset_on_exception(self): """Test that force_reset clears all models from the manager.""" model_name = "test_model" def dummy_loader(): self.mock_monitor.allocate(1000.0) raise RuntimeError("Simulated loader exception") try: instance = self.manager.acquire_model( model_name, lambda: "model_instance") self.manager.release_model(model_name, instance) instance = self.manager.acquire_model(model_name, dummy_loader) except RuntimeError: self.manager._force_reset() self.assertTrue(len(self.manager._models[model_name]) == 0) self.assertEqual(self.manager._total_active_jobs, 0) self.assertEqual(self.manager._pending_reservations, 0.0) self.assertFalse(self.manager._isolation_mode) pass instance = self.manager.acquire_model(model_name, lambda: "model_instance") self.manager.release_model(model_name, instance) def test_model_manager_metrics(self): """Test that distribution metrics are updated correctly.""" tag1 = "model1" tag2 = "model2" def _get_count_dist_max(tag): dist = MetricsEnvironment.process_wide_container().get_distribution( MetricName('BeamML_ModelManager', f'num_loaded_models_{tag}')) return dist.get_cumulative().max def _get_est_dist_mean(tag): dist = MetricsEnvironment.process_wide_container().get_distribution( MetricName('BeamML_ModelManager', f'memory_estimate_mb_{tag}')) val = dist.get_cumulative() return int(val.sum / val.count) if val.count > 0 else 0 # Verify that initial estimates correctly export int metrics self.manager._estimator.set_initial_estimate(tag1, 1000.5) self.assertEqual(_get_est_dist_mean(tag1), 1000) self.manager._estimator.set_initial_estimate(tag2, 2000.9) self.assertEqual(_get_est_dist_mean(tag2), 2000) # 1. Acquire a model self.manager.acquire_model( tag1, lambda: MockModel(tag1, 1000.0, self.mock_monitor)) self.assertEqual(_get_count_dist_max(tag1), 1) self.assertEqual(_get_est_dist_mean(tag1), 1000) # 2. Acquire another instance of same model self.manager.acquire_model( tag1, lambda: MockModel(tag1, 1000.0, self.mock_monitor)) self.assertEqual(_get_count_dist_max(tag1), 2) self.assertEqual(_get_est_dist_mean(tag1), 1000) # 3. Acquire a different model self.manager.acquire_model( tag2, lambda: MockModel(tag2, 2000.0, self.mock_monitor)) self.assertEqual(_get_count_dist_max(tag2), 1) self.assertEqual(_get_est_dist_mean(tag2), 2000) # tag1 max count should remain 2 self.assertEqual(_get_count_dist_max(tag1), 2) self.assertEqual(_get_est_dist_mean(tag1), 1000) # 4. Delete all models self.manager._delete_all_models() # It retains the highest count it ever saw. self.assertEqual(_get_count_dist_max(tag1), 2) self.assertEqual(_get_count_dist_max(tag2), 1) self.assertEqual(_get_est_dist_mean(tag1), 1000) self.assertEqual(_get_est_dist_mean(tag2), 2000) # 5. Repopulate and force reset self.manager.acquire_model( tag1, lambda: MockModel(tag1, 1000.0, self.mock_monitor)) # Max is still 2 from earlier in the test run self.assertEqual(_get_count_dist_max(tag1), 2) self.manager._force_reset() self.assertEqual(_get_count_dist_max(tag1), 2) def test_single_model_convergence_with_fluctuations(self): """ Tests that the estimator converges to the true usage with fluctuations. """ model_name = "fluctuating_model" model_cost = 3000.0 # Fix random seed for reproducibility random.seed(42) def loader(): self.mock_monitor.allocate(model_cost) return model_name model = self.manager.acquire_model(model_name, loader) self.manager.release_model(model_name, model) initial_est = self.manager._estimator.get_estimate(model_name) self.assertEqual(initial_est, model_cost) def run_inference(): model = self.manager.acquire_model(model_name, loader) noise = random.uniform(-300.0, 300.0) self.mock_monitor.allocate(noise) time.sleep(0.1) self.mock_monitor.free(noise) self.manager.release_model(model_name, model) return with ThreadPoolExecutor(max_workers=8) as executor: futures = [executor.submit(run_inference) for _ in range(100)] for f in futures: f.result() est_cost = self.manager._estimator.get_estimate(model_name) self.assertAlmostEqual(est_cost, model_cost, delta=100.0) class TestModelManagerEviction(unittest.TestCase): def setUp(self): self.mock_monitor = MockGPUMonitor(total_memory=12000.0) self.manager = ModelManager( monitor=self.mock_monitor, slack_percentage=0.0, min_data_points=1, eviction_cooldown_seconds=10.0, min_model_copies=1) def tearDown(self): self.manager.shutdown() def create_loader(self, name, size): return lambda: MockModel(name, size, self.mock_monitor) def test_basic_lru_eviction(self): self.manager._estimator.set_initial_estimate("A", 4000) self.manager._estimator.set_initial_estimate("B", 4000) self.manager._estimator.set_initial_estimate("C", 5000) model_a = self.manager.acquire_model("A", self.create_loader("A", 4000)) self.manager.release_model("A", model_a) model_b = self.manager.acquire_model("B", self.create_loader("B", 4000)) self.manager.release_model("B", model_b) key_a = list(self.manager._idle_lru.keys())[0] self.manager._idle_lru[key_a] = ("A", model_a, time.time() - 20.0) key_b = list(self.manager._idle_lru.keys())[1] self.manager._idle_lru[key_b] = ("B", model_b, time.time() - 20.0) model_a_again = self.manager.acquire_model( "A", self.create_loader("A", 4000)) self.manager.release_model("A", model_a_again) self.manager.acquire_model("C", self.create_loader("C", 5000)) self.assertEqual(len(self.manager.all_models("B")), 0) self.assertEqual(len(self.manager.all_models("A")), 1) def test_chained_eviction(self): self.manager._estimator.set_initial_estimate("big_guy", 8000) models = [] for i in range(4): name = f"small_{i}" m = self.manager.acquire_model(name, self.create_loader(name, 3000)) self.manager.release_model(name, m) models.append(m) self.manager.acquire_model("big_guy", self.create_loader("big_guy", 8000)) self.assertTrue(models[0].deleted) self.assertTrue(models[1].deleted) self.assertTrue(models[2].deleted) self.assertFalse(models[3].deleted) def test_active_models_are_protected(self): self.manager._estimator.set_initial_estimate("A", 6000) self.manager._estimator.set_initial_estimate("B", 4000) self.manager._estimator.set_initial_estimate("C", 4000) model_a = self.manager.acquire_model("A", self.create_loader("A", 6000)) model_b = self.manager.acquire_model("B", self.create_loader("B", 4000)) self.manager.release_model("B", model_b) key_b = list(self.manager._idle_lru.keys())[0] self.manager._idle_lru[key_b] = ("B", model_b, time.time() - 20.0) def acquire_c(): return self.manager.acquire_model("C", self.create_loader("C", 4000)) with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(acquire_c) model_c = future.result(timeout=2.0) self.assertTrue(model_b.deleted) self.assertFalse(model_a.deleted) self.manager.release_model("A", model_a) self.manager.release_model("C", model_c) def test_unknown_model_clears_memory(self): self.manager._estimator.set_initial_estimate("A", 2000) model_a = self.manager.acquire_model("A", self.create_loader("A", 2000)) self.manager.release_model("A", model_a) self.assertFalse(model_a.deleted) self.assertTrue(self.manager._estimator.is_unknown("X")) model_x = self.manager.acquire_model("X", self.create_loader("X", 10000)) self.assertTrue(model_a.deleted, "Model A should be deleted for isolation") self.assertEqual(len(self.manager.all_models("A")), 0) self.assertTrue(self.manager._isolation_mode) self.manager.release_model("X", model_x) def test_concurrent_eviction_pressure(self): def worker(idx): name = f"model_{idx % 5}" try: m = self.manager.acquire_model(name, self.create_loader(name, 4000)) time.sleep(0.001) self.manager.release_model(name, m) except Exception: pass with ThreadPoolExecutor(max_workers=8) as executor: futures = [executor.submit(worker, i) for i in range(50)] for f in futures: f.result() curr, _, _ = self.mock_monitor.get_stats() expected_usage = 0 for _, instances in self.manager._models.items(): expected_usage += len(instances) * 4000 self.assertAlmostEqual(curr, expected_usage) def test_starvation_prevention_overrides_demand(self): self.manager._estimator.set_initial_estimate("A", 12000) m_a = self.manager.acquire_model("A", self.create_loader("A", 12000)) self.manager.release_model("A", m_a) def cycle_a(): try: m = self.manager.acquire_model("A", self.create_loader("A", 12000)) time.sleep(0.3) self.manager.release_model("A", m) except Exception: pass executor = ThreadPoolExecutor(max_workers=5) for _ in range(5): executor.submit(cycle_a) def acquire_b(): return self.manager.acquire_model("B", self.create_loader("B", 4000)) b_future = executor.submit(acquire_b) model_b = b_future.result() self.assertTrue(m_a.deleted) self.manager.release_model("B", model_b) executor.shutdown(wait=True) class TestGPUMonitor(unittest.TestCase): def setUp(self): self.subprocess_patcher = patch('subprocess.check_output') self.mock_subprocess = self.subprocess_patcher.start() def tearDown(self): self.subprocess_patcher.stop() def test_init_hardware_detected(self): """Test that init correctly reads total memory when nvidia-smi exists.""" self.mock_subprocess.return_value = "24576" monitor = GPUMonitor() monitor.start() self.assertTrue(monitor._gpu_available) self.assertEqual(monitor._total_memory, 24576.0) def test_init_hardware_missing(self): """Test fallback behavior when nvidia-smi is missing.""" self.mock_subprocess.side_effect = FileNotFoundError() monitor = GPUMonitor(fallback_memory_mb=12000.0) monitor.start() self.assertFalse(monitor._gpu_available) self.assertEqual(monitor._total_memory, 12000.0) @patch('time.sleep') def test_polling_updates_stats(self, mock_sleep): """Test that the polling loop updates current and peak usage.""" def subprocess_side_effect(*args, **kwargs): if isinstance(args[0], list) and "memory.total" in args[0][1]: return "16000" if isinstance(args[0], list) and any("memory.free" in part for part in args[0]): return "12000" raise Exception("Unexpected command") self.mock_subprocess.side_effect = subprocess_side_effect self.mock_subprocess.return_value = None monitor = GPUMonitor() monitor.start() time.sleep(0.1) curr, peak, total = monitor.get_stats() monitor.stop() self.assertEqual(curr, 4000.0) self.assertEqual(peak, 4000.0) self.assertEqual(total, 16000.0) def test_reset_peak(self): """Test that resetting peak usage works.""" monitor = GPUMonitor() monitor._gpu_available = True with monitor._lock: monitor._current_usage = 2000.0 monitor._peak_usage = 8000.0 monitor._memory_history.append((time.time(), 8000.0)) monitor._memory_history.append((time.time(), 2000.0)) monitor.reset_peak() _, peak, _ = monitor.get_stats() self.assertEqual(peak, 2000.0) class TestResourceEstimatorSolver(unittest.TestCase): def setUp(self): self.estimator = ResourceEstimator() @patch('apache_beam.ml.inference.model_manager.nnls') def test_solver_respects_min_data_points(self, mock_nnls): mock_nnls.return_value = ([100.0, 50.0], 0.0) self.estimator.add_observation({'model_A': 1}, 500) self.estimator.add_observation({'model_B': 1}, 500) self.assertFalse(mock_nnls.called) self.estimator.add_observation({'model_A': 1, 'model_B': 1}, 1000) self.assertFalse(mock_nnls.called) self.estimator.add_observation({'model_A': 1}, 500) self.assertFalse(mock_nnls.called) self.estimator.add_observation({'model_B': 1}, 500) self.assertTrue(mock_nnls.called) @patch('apache_beam.ml.inference.model_manager.nnls') def test_solver_respects_unique_model_constraint(self, mock_nnls): mock_nnls.return_value = ([100.0, 100.0, 50.0], 0.0) for _ in range(5): self.estimator.add_observation({'model_A': 1, 'model_B': 1}, 800) for _ in range(5): self.estimator.add_observation({'model_C': 1}, 400) self.assertFalse(mock_nnls.called) self.estimator.add_observation({'model_A': 1}, 300) self.estimator.add_observation({'model_B': 1}, 300) self.assertTrue(mock_nnls.called) if __name__ == "__main__": unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/onnx_inference.py0000644000175100017510000001562615206113032023607 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from collections.abc import Callable from collections.abc import Iterable from collections.abc import Sequence from typing import Any from typing import Optional import numpy import onnx import onnxruntime as ort from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import PredictionResult __all__ = ['OnnxModelHandlerNumpy'] NumpyInferenceFn = Callable[ [Sequence[numpy.ndarray], ort.InferenceSession, Optional[dict[str, Any]]], Iterable[PredictionResult]] def default_numpy_inference_fn( inference_session: ort.InferenceSession, batch: Sequence[numpy.ndarray], inference_args: Optional[dict[str, Any]] = None) -> Any: ort_inputs = { inference_session.get_inputs()[0].name: numpy.stack(batch, axis=0) } if inference_args: ort_inputs = {**ort_inputs, **inference_args} ort_outs = inference_session.run(None, ort_inputs)[0] return ort_outs class OnnxModelHandlerNumpy(ModelHandler[numpy.ndarray, PredictionResult, ort.InferenceSession]): def __init__( #pylint: disable=dangerous-default-value self, model_uri: str, session_options=None, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'], provider_options=None, *, inference_fn: NumpyInferenceFn = default_numpy_inference_fn, large_model: bool = False, model_copies: Optional[int] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, **kwargs): """ Implementation of the ModelHandler interface for onnx using numpy arrays as input. Note that inputs to ONNXModelHandler should be of the same sizes Example Usage:: pcoll | RunInference(OnnxModelHandler(model_uri="my_uri")) Args: model_uri: The URI to where the model is saved. inference_fn: The inference function to use on RunInference calls. default=default_numpy_inference_fn large_model: set to true if your model is large enough to run into memory pressure if you load multiple copies. Given a model that consumes N memory and a machine with W cores and M memory, you should set this to True if N*W > M. model_copies: The exact number of models that you would like loaded onto your machine. This can be useful if you exactly know your CPU or GPU capacity and want to maximize resource utilization. min_batch_size: the minimum batch size to use when batching inputs. max_batch_size: the maximum batch size to use when batching inputs. max_batch_duration_secs: the maximum amount of time to buffer a batch before emitting; used in streaming contexts. max_batch_weight: the maximum total weight of a batch. element_size_fn: a function that returns the size (weight) of an element. batch_length_fn: a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: a sorted list of positive boundary values for length-aware batching buckets. kwargs: 'env_vars' can be used to set environment variables before loading the model. """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries, large_model=large_model, model_copies=model_copies, **kwargs) self._model_uri = model_uri self._session_options = session_options self._providers = providers self._provider_options = provider_options self._model_inference_fn = inference_fn def load_model(self) -> ort.InferenceSession: """Loads and initializes an onnx inference session for processing.""" # when path is remote, we should first load into memory then deserialize f = FileSystems.open(self._model_uri, "rb") model_proto = onnx.load(f) model_proto_bytes = model_proto if not isinstance(model_proto, bytes): if (hasattr(model_proto, "SerializeToString") and callable(model_proto.SerializeToString)): model_proto_bytes = model_proto.SerializeToString() else: raise TypeError( "No SerializeToString method is detected on loaded model. " + f"Type of model: {type(model_proto)}") ort_session = ort.InferenceSession( model_proto_bytes, sess_options=self._session_options, providers=self._providers, provider_options=self._provider_options) return ort_session def run_inference( self, batch: Sequence[numpy.ndarray], inference_session: ort.InferenceSession, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """Runs inferences on a batch of numpy arrays. Args: batch: A sequence of examples as numpy arrays. They should be single examples. inference_session: An onnx inference session. Must be runnable with input x where x is sequence of numpy array inference_args: Any additional arguments for an inference. Returns: An Iterable of type PredictionResult. """ predictions = self._model_inference_fn( inference_session, batch, inference_args) return utils._convert_to_result(batch, predictions) def get_num_bytes(self, batch: Sequence[numpy.ndarray]) -> int: """ Returns: The number of bytes of data for a batch. """ return sum((np_array.itemsize for np_array in batch)) def get_metrics_namespace(self) -> str: """ Returns: A namespace for metrics collected by the RunInference transform. """ return 'BeamML_Onnx' ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/onnx_inference_it_test.py0000644000175100017510000000620415206113032025332 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-End test for Onnx Inference""" import logging import os import unittest import uuid import pytest from apache_beam.io.filesystems import FileSystems from apache_beam.testing.test_pipeline import TestPipeline # pylint: disable=ungrouped-imports try: import onnx from apache_beam.examples.inference import onnx_sentiment_classification except ImportError as e: onnx = None def process_outputs(filepath): with FileSystems().open(filepath) as f: lines = f.readlines() lines = [l.decode('utf-8').strip('\n') for l in lines] return lines @unittest.skipIf( os.getenv('FORCE_ONNX_IT') is None and onnx is None, 'Missing dependencies. ' 'Test depends on onnx and transformers') class OnnxInference(unittest.TestCase): @pytest.mark.uses_onnx @pytest.mark.it_postcommit def test_onnx_run_inference_roberta_sentiment_classification(self): test_pipeline = TestPipeline(is_integration_test=True) # Path to text file containing some sentences file_of_sentences = ( 'gs://apache-beam-ml/testing/inputs/onnx/' 'sentiment_classification_input.txt') output_file_dir = 'local/sentiment_classification/output' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_uri = ( 'gs://apache-beam-ml/models/' 'roberta_sentiment_classification.onnx') extra_opts = { 'input': file_of_sentences, 'output': output_file, 'model_uri': model_uri, } onnx_sentiment_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) predictions = process_outputs(filepath=output_file) actuals_file = ( 'gs://apache-beam-ml/testing/expected_outputs/' 'test_onnx_run_inference_roberta_sentiment' '_classification_actuals.txt') actuals = process_outputs(filepath=actuals_file) predictions_dict = {} for prediction in predictions: text, predicted_text = prediction.split(';') predictions_dict[text] = predicted_text for actual in actuals: text, actual_predicted_text = actual.split(';') predicted_predicted_text = predictions_dict[text] self.assertEqual(actual_predicted_text, predicted_predicted_text) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/onnx_inference_test.py0000644000175100017510000004666615206113032024656 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import os import shutil import tempfile import unittest from collections import OrderedDict import numpy as np import pytest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to if bool(1): # lint doesn't like an unconditional `raise`. raise unittest.SkipTest( 'TODO: fix https://github.com/apache/beam/issues/31254') # Protect against environments where onnx and pytorch library is not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: import onnxruntime as ort import tensorflow as tf import tf2onnx import torch from skl2onnx import convert_sklearn from skl2onnx.common.data_types import FloatTensorType from sklearn import linear_model from tensorflow.keras import layers from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.onnx_inference import OnnxModelHandlerNumpy from apache_beam.ml.inference.onnx_inference import default_numpy_inference_fn except ImportError: raise unittest.SkipTest('Onnx dependencies are not installed') try: from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem except ImportError: GCSFileSystem = None # type: ignore class PytorchLinearRegression(torch.nn.Module): def __init__(self, input_dim, output_dim): super().__init__() self.linear = torch.nn.Linear(input_dim, output_dim) def forward(self, x): out = self.linear(x) return out def generate(self, x): out = self.linear(x) + 0.5 return out class TestDataAndModel(): def get_one_feature_samples(self): return [ np.array([1], dtype="float32"), np.array([5], dtype="float32"), np.array([-3], dtype="float32"), np.array([10.0], dtype="float32"), ] def get_one_feature_predictions(self): return [ PredictionResult(ex, pred) for ex, pred in zip( self.get_one_feature_samples(), [example * 2.0 + 0.5 for example in self.get_one_feature_samples()]) ] def get_two_feature_examples(self): return [ np.array([1, 5], dtype="float32"), np.array([3, 10], dtype="float32"), np.array([-14, 0], dtype="float32"), np.array([0.5, 0.5], dtype="float32") ] def get_two_feature_predictions(self): return [ PredictionResult(ex, pred) for ex, pred in zip( self.get_two_feature_examples(), [ f1 * 2.0 + f2 * 3 + 0.5 for f1, f2 in self.get_two_feature_examples() ]) ] def get_torch_one_feature_model(self): model = PytorchLinearRegression(input_dim=1, output_dim=1) model.load_state_dict( OrderedDict([('linear.weight', torch.Tensor([[2.0]])), ('linear.bias', torch.Tensor([0.5]))])) return model def get_tf_one_feature_model(self): params = [ np.array([[2.0]], dtype="float32"), np.array([0.5], dtype="float32") ] linear_layer = layers.Dense(units=1, weights=params) linear_model = tf.keras.Sequential([linear_layer]) return linear_model def get_sklearn_one_feature_model(self): x = [[0], [1]] y = [0.5, 2.5] model = linear_model.LinearRegression() model.fit(x, y) return model def get_torch_two_feature_model(self): model = PytorchLinearRegression(input_dim=2, output_dim=1) model.load_state_dict( OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])), ('linear.bias', torch.Tensor([0.5]))])) return model def get_tf_two_feature_model(self): params = [np.array([[2.0], [3]]), np.array([0.5], dtype="float32")] linear_layer = layers.Dense(units=1, weights=params) linear_model = tf.keras.Sequential([linear_layer]) return linear_model def get_sklearn_two_feature_model(self): x = [[1, 5], [3, 2], [1, 0]] y = [17.5, 12.5, 2.5] model = linear_model.LinearRegression() model.fit(x, y) return model def _compare_prediction_result(a, b): example_equal = np.array_equal(a.example, b.example) if isinstance(a.inference, dict): return all( x == y for x, y in zip(a.inference.values(), b.inference.values())) and example_equal return a.inference == b.inference and example_equal def _to_numpy(tensor): return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu( ).numpy() class TestOnnxModelHandler(OnnxModelHandlerNumpy): def __init__( #pylint: disable=dangerous-default-value self, model_uri: str, session_options=None, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'], provider_options=None, *, inference_fn=default_numpy_inference_fn, large_model=False, **kwargs): self._model_uri = model_uri self._session_options = session_options self._providers = providers self._provider_options = provider_options self._model_inference_fn = inference_fn self._env_vars = kwargs.get('env_vars', {}) self._large_model = large_model class OnnxTestBase(unittest.TestCase): def setUp(self): self.tmpdir = tempfile.mkdtemp() self.test_data_and_model = TestDataAndModel() def tearDown(self): shutil.rmtree(self.tmpdir) @pytest.mark.uses_onnx class OnnxPytorchRunInferenceTest(OnnxTestBase): def test_onnx_pytorch_run_inference(self): examples = self.test_data_and_model.get_one_feature_samples() expected_predictions = self.test_data_and_model.get_one_feature_predictions( ) model = self.test_data_and_model.get_torch_one_feature_model() path = os.path.join(self.tmpdir, 'my_onnx_pytorch_path') dummy_input = torch.randn(4, 1, requires_grad=True) torch.onnx.export( model, dummy_input, # model input path, # where to save the model export_params=True, # store the trained parameter weights opset_version=10, # the ONNX version do_constant_folding=True, # whether to execute constant- # folding for optimization input_names=['input'], # model's input names output_names=['output'], # model's output names dynamic_axes={ 'input': { 0: 'batch_size' }, 'output': { 0: 'batch_size' } }) inference_runner = TestOnnxModelHandler(path) inference_session = ort.InferenceSession( path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] ) # this list specifies priority - prioritize gpu if cuda kernel exists predictions = inference_runner.run_inference(examples, inference_session) for actual, expected in zip(predictions, expected_predictions): self.assertEqual(actual, expected) def test_num_bytes(self): inference_runner = TestOnnxModelHandler("dummy") batched_examples_int = [ np.array([1, 2, 3]), np.array([4, 5, 6]), np.array([7, 8, 9]) ] self.assertEqual( batched_examples_int[0].itemsize * 3, inference_runner.get_num_bytes(batched_examples_int)) batched_examples_float = [ np.array([1, 5], dtype=np.float32), np.array([3, 10], dtype=np.float32), np.array([-14, 0], dtype=np.float32), np.array([0.5, 0.5], dtype=np.float32) ] self.assertEqual( batched_examples_float[0].itemsize * 4, inference_runner.get_num_bytes(batched_examples_float)) def test_namespace(self): inference_runner = TestOnnxModelHandler("dummy") self.assertEqual('BeamML_Onnx', inference_runner.get_metrics_namespace()) @pytest.mark.uses_onnx class OnnxTensorflowRunInferenceTest(OnnxTestBase): def test_onnx_tensorflow_run_inference(self): examples = self.test_data_and_model.get_one_feature_samples() expected_predictions = self.test_data_and_model.get_one_feature_predictions( ) linear_model = self.test_data_and_model.get_tf_one_feature_model() path = os.path.join(self.tmpdir, 'my_onnx_tf_path') spec = (tf.TensorSpec((None, 1), tf.float32, name="input"), ) _, _ = tf2onnx.convert.from_keras(linear_model, input_signature=spec, opset=13, output_path=path) inference_runner = TestOnnxModelHandler(path) inference_session = ort.InferenceSession( path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] ) # this list specifies priority - prioritize gpu if cuda kernel exists predictions = inference_runner.run_inference(examples, inference_session) for actual, expected in zip(predictions, expected_predictions): self.assertEqual(actual, expected) @pytest.mark.uses_onnx class OnnxSklearnRunInferenceTest(OnnxTestBase): def save_model(self, model, input_dim, path): # assume float input initial_type = [('float_input', FloatTensorType([None, input_dim]))] onx = convert_sklearn(model, initial_types=initial_type) with open(path, "wb") as f: f.write(onx.SerializeToString()) def test_onnx_sklearn_run_inference(self): examples = self.test_data_and_model.get_one_feature_samples() expected_predictions = self.test_data_and_model.get_one_feature_predictions( ) linear_model = self.test_data_and_model.get_sklearn_one_feature_model() path = os.path.join(self.tmpdir, 'my_onnx_sklearn_path') self.save_model(linear_model, 1, path) inference_runner = TestOnnxModelHandler(path) inference_session = ort.InferenceSession( path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] ) # this list specifies priority - prioritize gpu if cuda kernel exists predictions = inference_runner.run_inference(examples, inference_session) for actual, expected in \ zip(predictions, expected_predictions): self.assertEqual(actual, expected) @pytest.mark.uses_onnx class OnnxPytorchRunInferencePipelineTest(OnnxTestBase): def exportModelToOnnx(self, model, path): dummy_input = torch.randn(4, 2, requires_grad=True) torch.onnx.export( model, dummy_input, # model input path, # where to save the model export_params=True, # store the trained parameter weights opset_version=10, # the ONNX version do_constant_folding=True, # whether to execute constant # folding for optimization input_names=['input'], # odel's input names output_names=['output'], # model's output names dynamic_axes={ 'input': { 0: 'batch_size' }, 'output': { 0: 'batch_size' } }) def test_pipeline_local_model_simple(self): with TestPipeline() as pipeline: path = os.path.join(self.tmpdir, 'my_onnx_pytorch_path') model = self.test_data_and_model.get_torch_two_feature_model() self.exportModelToOnnx(model, path) model_handler = TestOnnxModelHandler(path) pcoll = pipeline | 'start' >> beam.Create( self.test_data_and_model.get_two_feature_examples()) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to( self.test_data_and_model.get_two_feature_predictions(), equals_fn=_compare_prediction_result)) def test_model_handler_sets_env_vars(self): with TestPipeline() as pipeline: path = os.path.join(self.tmpdir, 'my_onnx_pytorch_path') model = self.test_data_and_model.get_torch_two_feature_model() self.exportModelToOnnx(model, path) model_handler = OnnxModelHandlerNumpy( model_uri=path, env_vars={'FOO': 'bar'}) self.assertFalse('FOO' in os.environ) _ = ( pipeline | 'start' >> beam.Create( self.test_data_and_model.get_two_feature_examples()) | RunInference(model_handler)) pipeline.run() self.assertTrue('FOO' in os.environ) self.assertTrue('bar'.equals(os.environ['FOO'])) def test_model_handler_large_model(self): with TestPipeline() as pipeline: def onxx_numpy_inference_fn( inference_session: ort.InferenceSession, batch, inference_args=None): multi_process_shared_loaded = "multi_process_shared" in str( type(inference_session)) if not multi_process_shared_loaded: raise Exception( f'Loaded model of type {type(model)}, was ' + 'expecting multi_process_shared_model') return default_numpy_inference_fn( inference_session, batch, inference_args) path = os.path.join(self.tmpdir, 'my_onnx_pytorch_path') model = self.test_data_and_model.get_torch_two_feature_model() self.exportModelToOnnx(model, path) model_handler = OnnxModelHandlerNumpy( model_uri=path, env_vars={'FOO': 'bar'}, inference_fn=onxx_numpy_inference_fn, large_model=True) self.assertFalse('FOO' in os.environ) _ = ( pipeline | 'start' >> beam.Create( self.test_data_and_model.get_two_feature_examples()) | RunInference(model_handler)) pipeline.run() @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed') def test_pipeline_gcs_model(self): with TestPipeline() as pipeline: examples = self.test_data_and_model.get_one_feature_samples() expected_predictions = ( self.test_data_and_model.get_one_feature_predictions()) gs_path = 'gs://apache-beam-ml/models/torch_2xplus5_onnx' # first need to download model from remote model_handler = TestOnnxModelHandler(gs_path) pcoll = pipeline | 'start' >> beam.Create(examples) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to(expected_predictions, equals_fn=_compare_prediction_result)) def test_invalid_input_type(self): with self.assertRaisesRegex(Exception, "Got invalid dimensions for input"): with TestPipeline() as pipeline: examples = [np.array([1], dtype="float32")] path = os.path.join(self.tmpdir, 'my_onnx_pytorch_path') model = self.test_data_and_model.get_torch_two_feature_model() self.exportModelToOnnx(model, path) model_handler = TestOnnxModelHandler(path) pcoll = pipeline | 'start' >> beam.Create(examples) # pylint: disable=expression-not-assigned pcoll | RunInference(model_handler) @pytest.mark.uses_onnx class OnnxTensorflowRunInferencePipelineTest(OnnxTestBase): def exportModelToOnnx(self, model, path): spec = (tf.TensorSpec((None, 2), tf.float32, name="input"), ) _, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=13, output_path=path) def test_pipeline_local_model_simple(self): with TestPipeline() as pipeline: path = os.path.join(self.tmpdir, 'my_onnx_tensorflow_path') model = self.test_data_and_model.get_tf_two_feature_model() self.exportModelToOnnx(model, path) model_handler = TestOnnxModelHandler(path) pcoll = pipeline | 'start' >> beam.Create( self.test_data_and_model.get_two_feature_examples()) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to( self.test_data_and_model.get_two_feature_predictions(), equals_fn=_compare_prediction_result)) @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed') def test_pipeline_gcs_model(self): with TestPipeline() as pipeline: examples = self.test_data_and_model.get_one_feature_samples() expected_predictions = ( self.test_data_and_model.get_one_feature_predictions()) gs_path = 'gs://apache-beam-ml/models/tf_2xplus5_onnx' model_handler = TestOnnxModelHandler(gs_path) pcoll = pipeline | 'start' >> beam.Create(examples) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to(expected_predictions, equals_fn=_compare_prediction_result)) def test_invalid_input_type(self): with self.assertRaisesRegex(Exception, "Got invalid dimensions for input"): with TestPipeline() as pipeline: examples = [np.array([1], dtype="float32")] path = os.path.join(self.tmpdir, 'my_onnx_tensorflow_path') model = self.test_data_and_model.get_tf_two_feature_model() self.exportModelToOnnx(model, path) model_handler = TestOnnxModelHandler(path) pcoll = pipeline | 'start' >> beam.Create(examples) # pylint: disable=expression-not-assigned pcoll | RunInference(model_handler) @pytest.mark.uses_onnx class OnnxSklearnRunInferencePipelineTest(OnnxTestBase): def save_model(self, model, input_dim, path): # assume float input initial_type = [('float_input', FloatTensorType([None, input_dim]))] onx = convert_sklearn(model, initial_types=initial_type) with open(path, "wb") as f: f.write(onx.SerializeToString()) def test_pipeline_local_model_simple(self): with TestPipeline() as pipeline: path = os.path.join(self.tmpdir, 'my_onnx_sklearn_path') model = self.test_data_and_model.get_sklearn_two_feature_model() self.save_model(model, 2, path) model_handler = TestOnnxModelHandler(path) pcoll = pipeline | 'start' >> beam.Create( self.test_data_and_model.get_two_feature_examples()) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to( self.test_data_and_model.get_two_feature_predictions(), equals_fn=_compare_prediction_result)) @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed') def test_pipeline_gcs_model(self): with TestPipeline() as pipeline: examples = (self.test_data_and_model.get_one_feature_samples()) expected_predictions = ( self.test_data_and_model.get_one_feature_predictions()) gs_path = 'gs://apache-beam-ml/models/skl_2xplus5_onnx' model_handler = TestOnnxModelHandler(gs_path) pcoll = pipeline | 'start' >> beam.Create(examples) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to(expected_predictions, equals_fn=_compare_prediction_result)) def test_invalid_input_type(self): with self.assertRaisesRegex(Exception, "InvalidArgument"): with TestPipeline() as pipeline: examples = [np.array([1], dtype="float32")] path = os.path.join(self.tmpdir, 'my_onnx_sklearn_path') model = self.test_data_and_model.get_sklearn_two_feature_model() self.save_model(model, 2, path) model_handler = TestOnnxModelHandler(path) pcoll = pipeline | 'start' >> beam.Create(examples) # pylint: disable=expression-not-assigned pcoll | RunInference(model_handler) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/pytorch_inference.py0000644000175100017510000005747015206113032024320 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging from collections import defaultdict from collections.abc import Callable from collections.abc import Iterable from collections.abc import Sequence from typing import Any from typing import Optional import torch from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import PredictionResult __all__ = [ 'PytorchModelHandlerTensor', 'PytorchModelHandlerKeyedTensor', ] TensorInferenceFn = Callable[[ Sequence[torch.Tensor], torch.nn.Module, torch.device, Optional[dict[str, Any]], Optional[str] ], Iterable[PredictionResult]] KeyedTensorInferenceFn = Callable[[ Sequence[dict[str, torch.Tensor]], torch.nn.Module, torch.device, Optional[dict[str, Any]], Optional[str] ], Iterable[PredictionResult]] def _cuda_device_is_usable() -> bool: """Returns True only when CUDA can actually allocate tensors.""" if not torch.cuda.is_available(): return False try: # Some environments report CUDA available but fail at first real use # because a driver is missing or inaccessible. torch.empty(1, device='cuda') return True except Exception: # pylint: disable=broad-except logging.warning("CUDA probe failed", exc_info=True) return False def _validate_constructor_args( state_dict_path, model_class, torch_script_model_path): message = ( "A {param1} has been supplied to the model " "handler, but the required {param2} is missing. " "Please provide the {param2} in order to " "successfully load the {param1}.") # state_dict_path and model_class are coupled with each other # raise RuntimeError if user forgets to pass any one of them. if state_dict_path and not model_class: raise RuntimeError( message.format(param1="state_dict_path", param2="model_class")) if not state_dict_path and model_class: raise RuntimeError( message.format(param1="model_class", param2="state_dict_path")) if torch_script_model_path and state_dict_path: raise RuntimeError( "Please specify either torch_script_model_path or " "(state_dict_path, model_class) to successfully load the model.") def _load_model( model_class: Optional[Callable[..., torch.nn.Module]], state_dict_path: Optional[str], device: torch.device, model_params: Optional[dict[str, Any]], torch_script_model_path: Optional[str], load_model_args: Optional[dict[str, Any]]): if device == torch.device('cuda') and not _cuda_device_is_usable(): logging.warning( "Model handler specified a 'GPU' device, but GPUs are not available. " "Switching to CPU.") device = torch.device('cpu') try: logging.info( "Loading state_dict_path %s onto a %s device", state_dict_path, device) if not torch_script_model_path: file = FileSystems.open(state_dict_path, 'rb') model = model_class(**model_params) # type: ignore[arg-type,misc] state_dict = torch.load(file, map_location=device, **load_model_args) model.load_state_dict(state_dict) model.requires_grad_(False) else: file = FileSystems.open(torch_script_model_path, 'rb') model = torch.jit.load(file, map_location=device, **load_model_args) except RuntimeError as e: if device == torch.device('cuda'): message = "Loading the model onto a GPU device failed due to an " \ f"exception:\n{e}\nAttempting to load onto a CPU device instead." logging.warning(message) return _load_model( model_class, state_dict_path, torch.device('cpu'), model_params, torch_script_model_path, load_model_args) else: raise e model.to(device) model.eval() logging.info("Finished loading PyTorch model.") return model, device def _convert_to_device(examples: torch.Tensor, device) -> torch.Tensor: """ Converts samples to a style matching given device. **NOTE:** A user may pass in device='GPU' but if GPU is not detected in the environment it must be converted back to CPU. """ if examples.device != device: examples = examples.to(device) return examples def default_tensor_inference_fn( batch: Sequence[torch.Tensor], model: torch.nn.Module, device: str, inference_args: Optional[dict[str, Any]] = None, model_id: Optional[str] = None, ) -> Iterable[PredictionResult]: # torch.no_grad() mitigates GPU memory issues # https://github.com/apache/beam/issues/22811 with torch.no_grad(): batched_tensors = torch.stack(batch) batched_tensors = _convert_to_device(batched_tensors, device) predictions = model(batched_tensors, **inference_args) return utils._convert_to_result(batch, predictions, model_id) def make_tensor_model_fn(model_fn: str) -> TensorInferenceFn: """ Produces a TensorInferenceFn that uses a method of the model other that the forward() method. Args: model_fn: A string name of the method to be used. This is accessed through getattr(model, model_fn) """ def attr_fn( batch: Sequence[torch.Tensor], model: torch.nn.Module, device: str, inference_args: Optional[dict[str, Any]] = None, model_id: Optional[str] = None, ) -> Iterable[PredictionResult]: with torch.no_grad(): batched_tensors = torch.stack(batch) batched_tensors = _convert_to_device(batched_tensors, device) pred_fn = getattr(model, model_fn) predictions = pred_fn(batched_tensors, **inference_args) return utils._convert_to_result(batch, predictions, model_id) return attr_fn class PytorchModelHandlerTensor(ModelHandler[torch.Tensor, PredictionResult, torch.nn.Module]): def __init__( self, state_dict_path: Optional[str] = None, model_class: Optional[Callable[..., torch.nn.Module]] = None, model_params: Optional[dict[str, Any]] = None, device: str = 'CPU', *, inference_fn: TensorInferenceFn = default_tensor_inference_fn, torch_script_model_path: Optional[str] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, large_model: bool = False, model_copies: Optional[int] = None, load_model_args: Optional[dict[str, Any]] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, **kwargs): """Implementation of the ModelHandler interface for PyTorch. Example Usage for torch model:: pcoll | RunInference(PytorchModelHandlerTensor(state_dict_path="my_uri", model_class="my_class")) Example Usage for torchscript model:: pcoll | RunInference(PytorchModelHandlerTensor( torch_script_model_path="my_uri")) See https://pytorch.org/tutorials/beginner/saving_loading_models.html for details Args: state_dict_path: path to the saved dictionary of the model state. model_class: class of the Pytorch model that defines the model structure. model_params: A dictionary of arguments required to instantiate the model class. device: the device on which you wish to run the model. If ``device = GPU`` then a GPU device will be used if it is available. Otherwise, it will be CPU. inference_fn: the inference function to use during RunInference. default=_default_tensor_inference_fn torch_script_model_path: Path to the torch script model. the model will be loaded using `torch.jit.load()`. `state_dict_path`, `model_class` and `model_params` arguments will be disregarded. min_batch_size: the minimum batch size to use when batching inputs. This batch will be fed into the inference_fn as a Sequence of Tensors. max_batch_size: the maximum batch size to use when batching inputs. This batch will be fed into the inference_fn as a Sequence of Tensors. max_batch_duration_secs: the maximum amount of time to buffer a batch before emitting; used in streaming contexts. large_model: set to true if your model is large enough to run into memory pressure if you load multiple copies. Given a model that consumes N memory and a machine with W cores and M memory, you should set this to True if N*W > M. model_copies: The exact number of models that you would like loaded onto your machine. This can be useful if you exactly know your CPU or GPU capacity and want to maximize resource utilization. load_model_args: a dictionary of parameters passed to the torch.load function to specify custom config for loading models. max_batch_weight: the maximum total weight of a batch. element_size_fn: a function that returns the size (weight) of an element. batch_length_fn: a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: a sorted list of positive boundary values for length-aware batching buckets. kwargs: 'env_vars' can be used to set environment variables before loading the model. **Supported Versions:** RunInference APIs in Apache Beam have been tested with PyTorch 1.9 and 1.10. """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries, large_model=large_model, model_copies=model_copies, **kwargs) self._state_dict_path = state_dict_path if device == 'GPU': logging.info("Device is set to CUDA") self._device = torch.device('cuda') else: logging.info("Device is set to CPU") self._device = torch.device('cpu') self._model_class = model_class self._model_params = model_params if model_params else {} self._inference_fn = inference_fn self._torch_script_model_path = torch_script_model_path self._load_model_args = load_model_args if load_model_args else {} _validate_constructor_args( state_dict_path=self._state_dict_path, model_class=self._model_class, torch_script_model_path=self._torch_script_model_path) def load_model(self) -> torch.nn.Module: """Loads and initializes a Pytorch model for processing.""" model, device = _load_model( model_class=self._model_class, state_dict_path=self._state_dict_path, device=self._device, model_params=self._model_params, torch_script_model_path=self._torch_script_model_path, load_model_args=self._load_model_args ) self._device = device return model def update_model_path(self, model_path: Optional[str] = None): if self._torch_script_model_path: self._torch_script_model_path = ( model_path if model_path else self._torch_script_model_path) else: self._state_dict_path = ( model_path if model_path else self._state_dict_path) def run_inference( self, batch: Sequence[torch.Tensor], model: torch.nn.Module, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """ Runs inferences on a batch of Tensors and returns an Iterable of Tensor Predictions. This method stacks the list of Tensors in a vectorized format to optimize the inference call. Args: batch: A sequence of Tensors. These Tensors should be batchable, as this method will call `torch.stack()` and pass in batched Tensors with dimensions (batch_size, n_features, etc.) into the model's forward() function. model: A PyTorch model. inference_args: Non-batchable arguments required as inputs to the model's forward() function. Unlike Tensors in `batch`, these parameters will not be dynamically batched Returns: An Iterable of type PredictionResult. """ inference_args = {} if not inference_args else inference_args model_id = ( self._state_dict_path if not self._torch_script_model_path else self._torch_script_model_path) return self._inference_fn( batch, model, self._device, inference_args, model_id) def get_num_bytes(self, batch: Sequence[torch.Tensor]) -> int: """ Returns: The number of bytes of data for a batch of Tensors. """ return sum((el.element_size() for tensor in batch for el in tensor)) def get_metrics_namespace(self) -> str: """ Returns: A namespace for metrics collected by the RunInference transform. """ return 'BeamML_PyTorch' def default_keyed_tensor_inference_fn( batch: Sequence[dict[str, torch.Tensor]], model: torch.nn.Module, device: str, inference_args: Optional[dict[str, Any]] = None, model_id: Optional[str] = None, ) -> Iterable[PredictionResult]: # If elements in `batch` are provided as a dictionaries from key to Tensors, # then iterate through the batch list, and group Tensors to the same key key_to_tensor_list = defaultdict(list) # torch.no_grad() mitigates GPU memory issues # https://github.com/apache/beam/issues/22811 with torch.no_grad(): for example in batch: for key, tensor in example.items(): key_to_tensor_list[key].append(tensor) key_to_batched_tensors = {} for key in key_to_tensor_list: batched_tensors = torch.stack(key_to_tensor_list[key]) batched_tensors = _convert_to_device(batched_tensors, device) key_to_batched_tensors[key] = batched_tensors predictions = model(**key_to_batched_tensors, **inference_args) return utils._convert_to_result(batch, predictions, model_id) def make_keyed_tensor_model_fn(model_fn: str) -> KeyedTensorInferenceFn: """ Produces a KeyedTensorInferenceFn that uses a method of the model other that the forward() method. Args: model_fn: A string name of the method to be used. This is accessed through getattr(model, model_fn) """ def attr_fn( batch: Sequence[dict[str, torch.Tensor]], model: torch.nn.Module, device: str, inference_args: Optional[dict[str, Any]] = None, model_id: Optional[str] = None, ) -> Iterable[PredictionResult]: # If elements in `batch` are provided as a dictionaries from key to Tensors, # then iterate through the batch list, and group Tensors to the same key key_to_tensor_list = defaultdict(list) # torch.no_grad() mitigates GPU memory issues # https://github.com/apache/beam/issues/22811 with torch.no_grad(): for example in batch: for key, tensor in example.items(): key_to_tensor_list[key].append(tensor) key_to_batched_tensors = {} for key in key_to_tensor_list: batched_tensors = torch.stack(key_to_tensor_list[key]) batched_tensors = _convert_to_device(batched_tensors, device) key_to_batched_tensors[key] = batched_tensors pred_fn = getattr(model, model_fn) predictions = pred_fn(**key_to_batched_tensors, **inference_args) return utils._convert_to_result(batch, predictions, model_id) return attr_fn class PytorchModelHandlerKeyedTensor(ModelHandler[dict[str, torch.Tensor], PredictionResult, torch.nn.Module]): def __init__( self, state_dict_path: Optional[str] = None, model_class: Optional[Callable[..., torch.nn.Module]] = None, model_params: Optional[dict[str, Any]] = None, device: str = 'CPU', *, inference_fn: KeyedTensorInferenceFn = default_keyed_tensor_inference_fn, torch_script_model_path: Optional[str] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, large_model: bool = False, model_copies: Optional[int] = None, load_model_args: Optional[dict[str, Any]] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, **kwargs): """Implementation of the ModelHandler interface for PyTorch. Example Usage for torch model:: pcoll | RunInference(PytorchModelHandlerKeyedTensor( state_dict_path="my_uri", model_class="my_class")) Example Usage for torchscript model:: pcoll | RunInference(PytorchModelHandlerKeyedTensor( torch_script_model_path="my_uri")) **NOTE:** This API and its implementation are under development and do not provide backward compatibility guarantees. See https://pytorch.org/tutorials/beginner/saving_loading_models.html for details Args: state_dict_path: path to the saved dictionary of the model state. model_class: class of the Pytorch model that defines the model structure. model_params: A dictionary of arguments required to instantiate the model class. device: the device on which you wish to run the model. If ``device = GPU`` then a GPU device will be used if it is available. Otherwise, it will be CPU. inference_fn: the function to invoke on run_inference. default = default_keyed_tensor_inference_fn torch_script_model_path: Path to the torch script model. the model will be loaded using `torch.jit.load()`. `state_dict_path`, `model_class` and `model_params` arguments will be disregarded. min_batch_size: the minimum batch size to use when batching inputs. This batch will be fed into the inference_fn as a Sequence of Keyed Tensors. max_batch_size: the maximum batch size to use when batching inputs. This batch will be fed into the inference_fn as a Sequence of Keyed Tensors. max_batch_duration_secs: the maximum amount of time to buffer a batch before emitting; used in streaming contexts. large_model: set to true if your model is large enough to run into memory pressure if you load multiple copies. Given a model that consumes N memory and a machine with W cores and M memory, you should set this to True if N*W > M. model_copies: The exact number of models that you would like loaded onto your machine. This can be useful if you exactly know your CPU or GPU capacity and want to maximize resource utilization. load_model_args: a dictionary of parameters passed to the torch.load function to specify custom config for loading models. max_batch_weight: the maximum total weight of a batch. element_size_fn: a function that returns the size (weight) of an element. batch_length_fn: a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: a sorted list of positive boundary values for length-aware batching buckets. kwargs: 'env_vars' can be used to set environment variables before loading the model. **Supported Versions:** RunInference APIs in Apache Beam have been tested on torch>=1.9.0,<1.14.0. """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries, large_model=large_model, model_copies=model_copies, **kwargs) self._state_dict_path = state_dict_path if device == 'GPU': logging.info("Device is set to CUDA") self._device = torch.device('cuda') else: logging.info("Device is set to CPU") self._device = torch.device('cpu') self._model_class = model_class self._model_params = model_params if model_params else {} self._inference_fn = inference_fn self._torch_script_model_path = torch_script_model_path self._load_model_args = load_model_args if load_model_args else {} _validate_constructor_args( state_dict_path=self._state_dict_path, model_class=self._model_class, torch_script_model_path=self._torch_script_model_path) def load_model(self) -> torch.nn.Module: """Loads and initializes a Pytorch model for processing.""" model, device = _load_model( model_class=self._model_class, state_dict_path=self._state_dict_path, device=self._device, model_params=self._model_params, torch_script_model_path=self._torch_script_model_path, load_model_args=self._load_model_args ) self._device = device return model def update_model_path(self, model_path: Optional[str] = None): if self._torch_script_model_path: self._torch_script_model_path = ( model_path if model_path else self._torch_script_model_path) else: self._state_dict_path = ( model_path if model_path else self._state_dict_path) def run_inference( self, batch: Sequence[dict[str, torch.Tensor]], model: torch.nn.Module, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """ Runs inferences on a batch of Keyed Tensors and returns an Iterable of Tensor Predictions. For the same key across all examples, this will stack all Tensors values in a vectorized format to optimize the inference call. Args: batch: A sequence of keyed Tensors. These Tensors should be batchable, as this method will call `torch.stack()` and pass in batched Tensors with dimensions (batch_size, n_features, etc.) into the model's forward() function. model: A PyTorch model. inference_args: Non-batchable arguments required as inputs to the model's forward() function. Unlike Tensors in `batch`, these parameters will not be dynamically batched Returns: An Iterable of type PredictionResult. """ inference_args = {} if not inference_args else inference_args model_id = ( self._state_dict_path if not self._torch_script_model_path else self._torch_script_model_path) return self._inference_fn( batch, model, self._device, inference_args, model_id) def get_num_bytes(self, batch: Sequence[torch.Tensor]) -> int: """ Returns: The number of bytes of data for a batch of dict of Tensors. """ # If elements in `batch` are provided as a dictionaries from key to Tensors return sum( (el.element_size() for tensor in batch for el in tensor.values())) def get_metrics_namespace(self) -> str: """ Returns: A namespace for metrics collected by the RunInference transform. """ return 'BeamML_PyTorch' ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/pytorch_inference_it_test.py0000644000175100017510000002641315206113032026044 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-End test for Pytorch Inference""" import logging import os import unittest import uuid import pytest from apache_beam.io.filesystems import FileSystems from apache_beam.testing.test_pipeline import TestPipeline # pylint: disable=ungrouped-imports try: import torch from apache_beam.examples.inference import pytorch_image_classification from apache_beam.examples.inference import pytorch_image_segmentation from apache_beam.examples.inference import pytorch_language_modeling from apache_beam.examples.inference import pytorch_model_per_key_image_segmentation except ImportError as e: torch = None # pylint: disable=line-too-long _EXPECTED_OUTPUTS = { 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005001.JPEG': '681', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005002.JPEG': '333', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005003.JPEG': '711', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005004.JPEG': '286', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005005.JPEG': '433', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005006.JPEG': '290', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005007.JPEG': '890', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005008.JPEG': '592', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005009.JPEG': '406', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005010.JPEG': '996', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005011.JPEG': '327', 'gs://apache-beam-ml/datasets/imagenet/raw-data/validation/ILSVRC2012_val_00005012.JPEG': '573' } def process_outputs(filepath): with FileSystems().open(filepath) as f: lines = f.readlines() lines = [l.decode('utf-8').strip('\n') for l in lines] return lines @unittest.skipIf( os.getenv('FORCE_TORCH_IT') is None and torch is None, 'Missing dependencies. ' 'Test depends on torch, torchvision, pillow, and transformers') class PyTorchInference(unittest.TestCase): @pytest.mark.uses_pytorch @pytest.mark.it_postcommit def test_torch_run_inference_imagenet_mobilenetv2(self): test_pipeline = TestPipeline(is_integration_test=True) # text files containing absolute path to the imagenet validation data on GCS file_of_image_names = 'gs://apache-beam-ml/testing/inputs/it_mobilenetv2_imagenet_validation_inputs.txt' # pylint: disable=line-too-long output_file_dir = 'gs://apache-beam-ml/testing/predictions' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_state_dict_path = 'gs://apache-beam-ml/models/imagenet_classification_mobilenet_v2.pt' # pylint: disable=line-too-long extra_opts = { 'input': file_of_image_names, 'output': output_file, 'model_state_dict_path': model_state_dict_path, } pytorch_image_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) predictions = process_outputs(filepath=output_file) for prediction in predictions: filename, prediction = prediction.split(',') self.assertEqual(_EXPECTED_OUTPUTS[filename], prediction) @pytest.mark.uses_pytorch @pytest.mark.it_postcommit def test_torch_run_inference_coco_maskrcnn_resnet50_fpn(self): test_pipeline = TestPipeline(is_integration_test=True) # text files containing absolute path to the coco validation data on GCS file_of_image_names = 'gs://apache-beam-ml/testing/inputs/it_coco_validation_inputs.txt' # pylint: disable=line-too-long output_file_dir = 'gs://apache-beam-ml/testing/predictions' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_state_dict_path = 'gs://apache-beam-ml/models/torchvision.models.detection.maskrcnn_resnet50_fpn.pth' # pylint: disable=line-too-long images_dir = 'gs://apache-beam-ml/datasets/coco/raw-data/val2017' extra_opts = { 'input': file_of_image_names, 'output': output_file, 'model_state_dict_path': model_state_dict_path, 'images_dir': images_dir, } pytorch_image_segmentation.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) predictions = process_outputs(filepath=output_file) actuals_file = 'gs://apache-beam-ml/testing/expected_outputs/test_torch_run_inference_coco_maskrcnn_resnet50_fpn_actuals.txt' # pylint: disable=line-too-long actuals = process_outputs(filepath=actuals_file) predictions_dict = {} for prediction in predictions: filename, prediction_labels = prediction.split(';') predictions_dict[filename] = prediction_labels for actual in actuals: filename, actual_labels = actual.split(';') prediction_labels = predictions_dict[filename] self.assertEqual(actual_labels, prediction_labels) @pytest.mark.uses_pytorch @pytest.mark.it_postcommit @pytest.mark.timeout(1800) def test_torch_run_inference_coco_maskrcnn_resnet50_fpn_v1_and_v2(self): test_pipeline = TestPipeline(is_integration_test=True) # text files containing absolute path to the coco validation data on GCS file_of_image_names = 'gs://apache-beam-ml/testing/inputs/it_coco_validation_inputs.txt' # pylint: disable=line-too-long output_file_dir = 'gs://apache-beam-ml/testing/predictions' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_state_dict_paths = [ 'gs://apache-beam-ml/models/torchvision.models.detection.maskrcnn_resnet50_fpn.pth', # pylint: disable=line-too-long 'gs://apache-beam-ml/models/torchvision.models.detection.maskrcnn_resnet50_fpn_v2.pth' # pylint: disable=line-too-long ] images_dir = 'gs://apache-beam-ml/datasets/coco/raw-data/val2017' extra_opts = { 'input': file_of_image_names, 'output': output_file, 'model_state_dict_paths': ','.join(model_state_dict_paths), 'images_dir': images_dir, } pytorch_model_per_key_image_segmentation.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) predictions = process_outputs(filepath=output_file) actuals_file = 'gs://apache-beam-ml/testing/expected_outputs/test_torch_run_inference_coco_maskrcnn_resnet50_fpn_v1_and_v2_actuals.txt' # pylint: disable=line-too-long actuals = process_outputs(filepath=actuals_file) predictions_dict = {} for prediction in predictions: p = prediction.split('---') filename = p[0] v1predictions = p[1] v2predictions = p[2] predictions_dict[filename] = (v1predictions, v2predictions) for actual in actuals: a = actual.split('---') filename = a[0] v1actuals = a[1] v2actuals = a[2] v1prediction_labels, v2prediction_labels = predictions_dict[filename] self.assertEqual(v1actuals, v1prediction_labels) self.assertEqual(v2actuals, v2prediction_labels) @pytest.mark.uses_pytorch @pytest.mark.it_postcommit def test_torch_run_inference_bert_for_masked_lm(self): test_pipeline = TestPipeline(is_integration_test=True) # Path to text file containing some sentences file_of_sentences = 'gs://apache-beam-ml/datasets/custom/sentences.txt' # pylint: disable=line-too-long output_file_dir = 'gs://apache-beam-ml/testing/predictions' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_state_dict_path = 'gs://apache-beam-ml/models/huggingface.BertForMaskedLM.bert-base-uncased.pth' # pylint: disable=line-too-long extra_opts = { 'input': file_of_sentences, 'output': output_file, 'model_state_dict_path': model_state_dict_path, } pytorch_language_modeling.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) predictions = process_outputs(filepath=output_file) actuals_file = 'gs://apache-beam-ml/testing/expected_outputs/test_torch_run_inference_bert_for_masked_lm_actuals.txt' # pylint: disable=line-too-long actuals = process_outputs(filepath=actuals_file) predictions_dict = {} for prediction in predictions: text, predicted_text = prediction.split(';') predictions_dict[text] = predicted_text for actual in actuals: text, actual_predicted_text = actual.split(';') predicted_predicted_text = predictions_dict[text] self.assertEqual(actual_predicted_text, predicted_predicted_text) @pytest.mark.uses_pytorch @pytest.mark.it_postcommit def test_torch_run_inference_bert_for_masked_lm_large_model(self): test_pipeline = TestPipeline(is_integration_test=True) # Path to text file containing some sentences file_of_sentences = 'gs://apache-beam-ml/datasets/custom/sentences.txt' # pylint: disable=line-too-long output_file_dir = 'gs://apache-beam-ml/testing/predictions' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_state_dict_path = 'gs://apache-beam-ml/models/huggingface.BertForMaskedLM.bert-base-uncased.pth' # pylint: disable=line-too-long extra_opts = { 'input': file_of_sentences, 'output': output_file, 'model_state_dict_path': model_state_dict_path, 'large_model': True, } pytorch_language_modeling.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) predictions = process_outputs(filepath=output_file) actuals_file = 'gs://apache-beam-ml/testing/expected_outputs/test_torch_run_inference_bert_for_masked_lm_actuals.txt' # pylint: disable=line-too-long actuals = process_outputs(filepath=actuals_file) predictions_dict = {} for prediction in predictions: text, predicted_text = prediction.split(';') predictions_dict[text] = predicted_text for actual in actuals: text, actual_predicted_text = actual.split(';') predicted_predicted_text = predictions_dict[text] self.assertEqual(actual_predicted_text, predicted_predicted_text) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/pytorch_inference_test.py0000644000175100017510000011352415206113032025350 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import os import shutil import tempfile import unittest from collections import OrderedDict import numpy as np import pytest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # Protect against environments where pytorch library is not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: import torch from apache_beam.ml.inference import pytorch_inference from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerKeyedTensor from apache_beam.ml.inference.pytorch_inference import PytorchModelHandlerTensor from apache_beam.ml.inference.pytorch_inference import default_keyed_tensor_inference_fn from apache_beam.ml.inference.pytorch_inference import default_tensor_inference_fn from apache_beam.ml.inference.pytorch_inference import make_keyed_tensor_model_fn from apache_beam.ml.inference.pytorch_inference import make_tensor_model_fn except ImportError: raise unittest.SkipTest('PyTorch dependencies are not installed') try: from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem except ImportError: GCSFileSystem = None # type: ignore TWO_FEATURES_EXAMPLES = [ torch.from_numpy(np.array([1, 5], dtype="float32")), torch.from_numpy(np.array([3, 10], dtype="float32")), torch.from_numpy(np.array([-14, 0], dtype="float32")), torch.from_numpy(np.array([0.5, 0.5], dtype="float32")), ] TWO_FEATURES_PREDICTIONS = [ PredictionResult(ex, pred) for ex, pred in zip( TWO_FEATURES_EXAMPLES, torch.Tensor( [f1 * 2.0 + f2 * 3 + 0.5 for f1, f2 in TWO_FEATURES_EXAMPLES]).reshape(-1, 1)) ] TWO_FEATURES_DICT_OUT_PREDICTIONS = [ PredictionResult( p.example, { "output1": p.inference, "output2": p.inference }) for p in TWO_FEATURES_PREDICTIONS ] KEYED_TORCH_EXAMPLES = [ { 'k1': torch.from_numpy(np.array([1], dtype="float32")), 'k2': torch.from_numpy(np.array([1.5], dtype="float32")) }, { 'k1': torch.from_numpy(np.array([5], dtype="float32")), 'k2': torch.from_numpy(np.array([5.5], dtype="float32")) }, { 'k1': torch.from_numpy(np.array([-3], dtype="float32")), 'k2': torch.from_numpy(np.array([-3.5], dtype="float32")) }, { 'k1': torch.from_numpy(np.array([10.0], dtype="float32")), 'k2': torch.from_numpy(np.array([10.5], dtype="float32")) }, ] KEYED_TORCH_PREDICTIONS = [ PredictionResult(ex, pred) for ex, pred in zip( KEYED_TORCH_EXAMPLES, torch.Tensor( [(example['k1'] * 2.0 + 0.5) + (example['k2'] * 2.0 + 0.5) for example in KEYED_TORCH_EXAMPLES]).reshape(-1, 1)) ] KEYED_TORCH_HELPER_PREDICTIONS = [ PredictionResult(ex, pred) for ex, pred in zip( KEYED_TORCH_EXAMPLES, torch.Tensor( [(example['k1'] * 2.0 + 0.5) + (example['k2'] * 2.0 + 0.5) + 0.5 for example in KEYED_TORCH_EXAMPLES]).reshape(-1, 1)) ] KEYED_TORCH_DICT_OUT_PREDICTIONS = [ PredictionResult( p.example, { "output1": p.inference, "output2": p.inference }) for p in KEYED_TORCH_PREDICTIONS ] class TestPytorchModelHandlerForInferenceOnly(PytorchModelHandlerTensor): def __init__(self, device, *, inference_fn=default_tensor_inference_fn): self._device = device self._inference_fn = inference_fn self._state_dict_path = None self._torch_script_model_path = None class TestPytorchModelHandlerKeyedTensorForInferenceOnly( PytorchModelHandlerKeyedTensor): def __init__(self, device, *, inference_fn=default_keyed_tensor_inference_fn): self._device = device self._inference_fn = inference_fn self._state_dict_path = None self._torch_script_model_path = None def _compare_prediction_result(x, y): if isinstance(x.example, dict): example_equals = all( torch.equal(x, y) for x, y in zip(x.example.values(), y.example.values())) else: example_equals = torch.equal(x.example, y.example) if not example_equals: return False if isinstance(x.inference, dict): return all( torch.equal(x, y) for x, y in zip(x.inference.values(), y.inference.values())) return torch.equal(x.inference, y.inference) def custom_tensor_inference_fn( batch, model, device, inference_args, model_id=None): predictions = [ PredictionResult(ex, pred) for ex, pred in zip( batch, torch.Tensor([item * 2.0 + 1.5 for item in batch]).reshape(-1, 1)) ] return predictions class PytorchLinearRegression(torch.nn.Module): def __init__(self, input_dim, output_dim): super().__init__() self.linear = torch.nn.Linear(input_dim, output_dim) def forward(self, x): out = self.linear(x) return out def generate(self, x): out = self.linear(x) + 0.5 return out class PytorchLinearRegressionDict(torch.nn.Module): def __init__(self, input_dim, output_dim): super().__init__() self.linear = torch.nn.Linear(input_dim, output_dim) def forward(self, x): out = self.linear(x) return {'output1': out, 'output2': out} class PytorchLinearRegressionKeyedBatchAndExtraInferenceArgs(torch.nn.Module): """ A linear model with batched keyed inputs and non-batchable extra args. Note: k1 and k2 are batchable examples passed in as a dict from str to tensor. prediction_param_array, prediction_param_bool are non-batchable extra args (typically model-related info) used to configure the model before its predict call is invoked """ def __init__(self, input_dim, output_dim): super().__init__() self.linear = torch.nn.Linear(input_dim, output_dim) def forward(self, k1, k2, prediction_param_array, prediction_param_bool): if not prediction_param_bool: raise ValueError("Expected prediction_param_bool to be True") if not torch.all(prediction_param_array): raise ValueError("Expected prediction_param_array to be all True") out = self.linear(k1) + self.linear(k2) return out @pytest.mark.uses_pytorch class PytorchRunInferenceTest(unittest.TestCase): def test_run_inference_single_tensor_feature(self): examples = [ torch.from_numpy(np.array([1], dtype="float32")), torch.from_numpy(np.array([5], dtype="float32")), torch.from_numpy(np.array([-3], dtype="float32")), torch.from_numpy(np.array([10.0], dtype="float32")), ] expected_predictions = [ PredictionResult(ex, pred) for ex, pred in zip( examples, torch.Tensor( [example * 2.0 + 0.5 for example in examples]).reshape(-1, 1)) ] model = PytorchLinearRegression(input_dim=1, output_dim=1) model.load_state_dict( OrderedDict([('linear.weight', torch.Tensor([[2.0]])), ('linear.bias', torch.Tensor([0.5]))])) model.eval() inference_runner = TestPytorchModelHandlerForInferenceOnly( torch.device('cpu')) predictions = inference_runner.run_inference(examples, model) for actual, expected in zip(predictions, expected_predictions): self.assertEqual(actual, expected) def test_run_inference_multiple_tensor_features(self): model = PytorchLinearRegression(input_dim=2, output_dim=1) model.load_state_dict( OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])), ('linear.bias', torch.Tensor([0.5]))])) model.eval() inference_runner = TestPytorchModelHandlerForInferenceOnly( torch.device('cpu')) predictions = inference_runner.run_inference(TWO_FEATURES_EXAMPLES, model) for actual, expected in zip(predictions, TWO_FEATURES_PREDICTIONS): self.assertEqual(actual, expected) def test_run_inference_multiple_tensor_features_dict_output(self): model = PytorchLinearRegressionDict(input_dim=2, output_dim=1) model.load_state_dict( OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])), ('linear.bias', torch.Tensor([0.5]))])) model.eval() inference_runner = TestPytorchModelHandlerForInferenceOnly( torch.device('cpu')) predictions = inference_runner.run_inference(TWO_FEATURES_EXAMPLES, model) for actual, expected in zip(predictions, TWO_FEATURES_DICT_OUT_PREDICTIONS): self.assertEqual(actual, expected) def test_run_inference_custom(self): examples = [ torch.from_numpy(np.array([1], dtype="float32")), torch.from_numpy(np.array([5], dtype="float32")), torch.from_numpy(np.array([-3], dtype="float32")), torch.from_numpy(np.array([10.0], dtype="float32")), ] expected_predictions = [ PredictionResult(ex, pred) for ex, pred in zip( examples, torch.Tensor( [example * 2.0 + 1.5 for example in examples]).reshape(-1, 1)) ] model = PytorchLinearRegression(input_dim=1, output_dim=1) model.load_state_dict( OrderedDict([('linear.weight', torch.Tensor([[2.0]])), ('linear.bias', torch.Tensor([0.5]))])) model.eval() inference_runner = TestPytorchModelHandlerForInferenceOnly( torch.device('cpu'), inference_fn=custom_tensor_inference_fn) predictions = inference_runner.run_inference(examples, model) for actual, expected in zip(predictions, expected_predictions): self.assertEqual(actual, expected) def test_run_inference_keyed(self): """ This tests for inputs that are passed as a dictionary from key to tensor instead of a standard non-keyed tensor example. Example: Typical input format is input = torch.tensor([1, 2, 3]) But Pytorch syntax allows inputs to have the form input = { 'k1' : torch.tensor([1, 2, 3]), 'k2' : torch.tensor([4, 5, 6]) } """ class PytorchLinearRegressionMultipleArgs(torch.nn.Module): def __init__(self, input_dim, output_dim): super().__init__() self.linear = torch.nn.Linear(input_dim, output_dim) def forward(self, k1, k2): out = self.linear(k1) + self.linear(k2) return out model = PytorchLinearRegressionMultipleArgs(input_dim=1, output_dim=1) model.load_state_dict( OrderedDict([('linear.weight', torch.Tensor([[2.0]])), ('linear.bias', torch.Tensor([0.5]))])) model.eval() inference_runner = TestPytorchModelHandlerKeyedTensorForInferenceOnly( torch.device('cpu')) predictions = inference_runner.run_inference(KEYED_TORCH_EXAMPLES, model) for actual, expected in zip(predictions, KEYED_TORCH_PREDICTIONS): self.assertTrue(_compare_prediction_result(actual, expected)) def test_run_inference_keyed_dict_output(self): class PytorchLinearRegressionMultipleArgsDict(torch.nn.Module): def __init__(self, input_dim, output_dim): super().__init__() self.linear = torch.nn.Linear(input_dim, output_dim) def forward(self, k1, k2): out = self.linear(k1) + self.linear(k2) return {'output1': out, 'output2': out} model = PytorchLinearRegressionMultipleArgsDict(input_dim=1, output_dim=1) model.load_state_dict( OrderedDict([('linear.weight', torch.Tensor([[2.0]])), ('linear.bias', torch.Tensor([0.5]))])) model.eval() inference_runner = TestPytorchModelHandlerKeyedTensorForInferenceOnly( torch.device('cpu')) predictions = inference_runner.run_inference(KEYED_TORCH_EXAMPLES, model) for actual, expected in zip(predictions, KEYED_TORCH_DICT_OUT_PREDICTIONS): self.assertTrue(_compare_prediction_result(actual, expected)) def test_inference_runner_inference_args(self): """ This tests for non-batchable input arguments. Since we do the batching for the user, we have to distinguish between the inputs that should be batched and the ones that should not be batched. """ inference_args = { 'prediction_param_array': torch.from_numpy( np.array([1, 2], dtype="float32")), 'prediction_param_bool': True } model = PytorchLinearRegressionKeyedBatchAndExtraInferenceArgs( input_dim=1, output_dim=1) model.load_state_dict( OrderedDict([('linear.weight', torch.Tensor([[2.0]])), ('linear.bias', torch.Tensor([0.5]))])) model.eval() inference_runner = TestPytorchModelHandlerKeyedTensorForInferenceOnly( torch.device('cpu')) predictions = inference_runner.run_inference( batch=KEYED_TORCH_EXAMPLES, model=model, inference_args=inference_args) for actual, expected in zip(predictions, KEYED_TORCH_PREDICTIONS): self.assertEqual(actual, expected) def test_run_inference_helper(self): examples = [ torch.from_numpy(np.array([1], dtype="float32")), torch.from_numpy(np.array([5], dtype="float32")), torch.from_numpy(np.array([-3], dtype="float32")), torch.from_numpy(np.array([10.0], dtype="float32")), ] expected_predictions = [ PredictionResult(ex, pred) for ex, pred in zip( examples, torch.Tensor( [example * 2.0 + 1.0 for example in examples]).reshape(-1, 1)) ] gen_fn = make_tensor_model_fn('generate') model = PytorchLinearRegression(input_dim=1, output_dim=1) model.load_state_dict( OrderedDict([('linear.weight', torch.Tensor([[2.0]])), ('linear.bias', torch.Tensor([0.5]))])) model.eval() inference_runner = TestPytorchModelHandlerForInferenceOnly( torch.device('cpu'), inference_fn=gen_fn) predictions = inference_runner.run_inference(examples, model) for actual, expected in zip(predictions, expected_predictions): self.assertEqual(actual, expected) def test_run_inference_keyed_helper(self): """ This tests for inputs that are passed as a dictionary from key to tensor instead of a standard non-keyed tensor example. Example: Typical input format is input = torch.tensor([1, 2, 3]) But Pytorch syntax allows inputs to have the form input = { 'k1' : torch.tensor([1, 2, 3]), 'k2' : torch.tensor([4, 5, 6]) } """ class PytorchLinearRegressionMultipleArgs(torch.nn.Module): def __init__(self, input_dim, output_dim): super().__init__() self.linear = torch.nn.Linear(input_dim, output_dim) def forward(self, k1, k2): out = self.linear(k1) + self.linear(k2) return out def generate(self, k1, k2): out = self.linear(k1) + self.linear(k2) + 0.5 return out model = PytorchLinearRegressionMultipleArgs(input_dim=1, output_dim=1) model.load_state_dict( OrderedDict([('linear.weight', torch.Tensor([[2.0]])), ('linear.bias', torch.Tensor([0.5]))])) model.eval() gen_fn = make_keyed_tensor_model_fn('generate') inference_runner = TestPytorchModelHandlerKeyedTensorForInferenceOnly( torch.device('cpu'), inference_fn=gen_fn) predictions = inference_runner.run_inference(KEYED_TORCH_EXAMPLES, model) for actual, expected in zip(predictions, KEYED_TORCH_HELPER_PREDICTIONS): self.assertTrue(_compare_prediction_result(actual, expected)) def test_num_bytes(self): inference_runner = TestPytorchModelHandlerForInferenceOnly( torch.device('cpu')) examples = torch.from_numpy( np.array([1, 5, 3, 10, -14, 0, 0.5, 0.5], dtype="float32")).reshape(-1, 2) self.assertEqual((examples[0].element_size()) * 8, inference_runner.get_num_bytes(examples)) def test_namespace(self): inference_runner = TestPytorchModelHandlerForInferenceOnly( torch.device('cpu')) self.assertEqual('BeamML_PyTorch', inference_runner.get_metrics_namespace()) @pytest.mark.uses_pytorch class PytorchRunInferencePipelineTest(unittest.TestCase): def setUp(self): self.tmpdir = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.tmpdir) def test_pipeline_local_model_simple(self): with TestPipeline() as pipeline: state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])), ('linear.bias', torch.Tensor([0.5]))]) path = os.path.join(self.tmpdir, 'my_state_dict_path') torch.save(state_dict, path) model_handler = PytorchModelHandlerTensor( state_dict_path=path, model_class=PytorchLinearRegression, model_params={ 'input_dim': 2, 'output_dim': 1 }) pcoll = pipeline | 'start' >> beam.Create(TWO_FEATURES_EXAMPLES) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to( TWO_FEATURES_PREDICTIONS, equals_fn=_compare_prediction_result)) def test_pipeline_local_model_large(self): with TestPipeline() as pipeline: def batch_validator_tensor_inference_fn( batch, model, device, inference_args, model_id, ): multi_process_shared_loaded = "multi_process_shared" in str(type(model)) if not multi_process_shared_loaded: raise Exception( f'Loaded model of type {type(model)}, was ' + 'expecting multi_process_shared_model') return default_tensor_inference_fn( batch, model, device, inference_args, model_id) state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])), ('linear.bias', torch.Tensor([0.5]))]) path = os.path.join(self.tmpdir, 'my_state_dict_path') torch.save(state_dict, path) model_handler = PytorchModelHandlerTensor( state_dict_path=path, model_class=PytorchLinearRegression, model_params={ 'input_dim': 2, 'output_dim': 1 }, inference_fn=batch_validator_tensor_inference_fn, large_model=True) pcoll = pipeline | 'start' >> beam.Create(TWO_FEATURES_EXAMPLES) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to( TWO_FEATURES_PREDICTIONS, equals_fn=_compare_prediction_result)) def test_pipeline_local_model_extra_inference_args(self): with TestPipeline() as pipeline: inference_args = { 'prediction_param_array': torch.from_numpy( np.array([1, 2], dtype="float32")), 'prediction_param_bool': True } state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0]])), ('linear.bias', torch.Tensor([0.5]))]) path = os.path.join(self.tmpdir, 'my_state_dict_path') torch.save(state_dict, path) model_handler = PytorchModelHandlerKeyedTensor( state_dict_path=path, model_class=PytorchLinearRegressionKeyedBatchAndExtraInferenceArgs, model_params={ 'input_dim': 1, 'output_dim': 1 }) pcoll = pipeline | 'start' >> beam.Create(KEYED_TORCH_EXAMPLES) inference_args_side_input = ( pipeline | 'create side' >> beam.Create(inference_args)) predictions = pcoll | RunInference( model_handler=model_handler, inference_args=beam.pvalue.AsDict(inference_args_side_input)) assert_that( predictions, equal_to( KEYED_TORCH_PREDICTIONS, equals_fn=_compare_prediction_result)) def test_pipeline_local_model_extra_inference_args_large(self): with TestPipeline() as pipeline: inference_args = { 'prediction_param_array': torch.from_numpy( np.array([1, 2], dtype="float32")), 'prediction_param_bool': True } state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0]])), ('linear.bias', torch.Tensor([0.5]))]) path = os.path.join(self.tmpdir, 'my_state_dict_path') torch.save(state_dict, path) def batch_validator_keyed_tensor_inference_fn( batch, model, device, inference_args, model_id, ): multi_process_shared_loaded = "multi_process_shared" in str(type(model)) if not multi_process_shared_loaded: raise Exception( f'Loaded model of type {type(model)}, was ' + 'expecting multi_process_shared_model') return default_keyed_tensor_inference_fn( batch, model, device, inference_args, model_id) model_handler = PytorchModelHandlerKeyedTensor( state_dict_path=path, model_class=PytorchLinearRegressionKeyedBatchAndExtraInferenceArgs, model_params={ 'input_dim': 1, 'output_dim': 1 }, inference_fn=batch_validator_keyed_tensor_inference_fn, large_model=True) pcoll = pipeline | 'start' >> beam.Create(KEYED_TORCH_EXAMPLES) inference_args_side_input = ( pipeline | 'create side' >> beam.Create(inference_args)) predictions = pcoll | RunInference( model_handler=model_handler, inference_args=beam.pvalue.AsDict(inference_args_side_input)) assert_that( predictions, equal_to( KEYED_TORCH_PREDICTIONS, equals_fn=_compare_prediction_result)) def test_pipeline_local_model_extra_inference_args_batching_args(self): with TestPipeline() as pipeline: inference_args = { 'prediction_param_array': torch.from_numpy( np.array([1, 2], dtype="float32")), 'prediction_param_bool': True } state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0]])), ('linear.bias', torch.Tensor([0.5]))]) path = os.path.join(self.tmpdir, 'my_state_dict_path') torch.save(state_dict, path) def batch_validator_keyed_tensor_inference_fn( batch, model, device, inference_args, model_id, ): if len(batch) != 2: raise Exception( f'Expected batch of size 2, received batch of size {len(batch)}') return default_keyed_tensor_inference_fn( batch, model, device, inference_args, model_id) model_handler = PytorchModelHandlerKeyedTensor( state_dict_path=path, model_class=PytorchLinearRegressionKeyedBatchAndExtraInferenceArgs, model_params={ 'input_dim': 1, 'output_dim': 1 }, inference_fn=batch_validator_keyed_tensor_inference_fn, min_batch_size=2, max_batch_size=2) pcoll = pipeline | 'start' >> beam.Create(KEYED_TORCH_EXAMPLES) inference_args_side_input = ( pipeline | 'create side' >> beam.Create(inference_args)) predictions = pcoll | RunInference( model_handler=model_handler, inference_args=beam.pvalue.AsDict(inference_args_side_input)) assert_that( predictions, equal_to( KEYED_TORCH_PREDICTIONS, equals_fn=_compare_prediction_result)) @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed') def test_pipeline_gcs_model(self): with TestPipeline() as pipeline: examples = torch.from_numpy( np.array([1, 5, 3, 10], dtype="float32").reshape(-1, 1)) expected_predictions = [ PredictionResult(ex, pred) for ex, pred in zip( examples, torch.Tensor( [example * 2.0 + 0.5 for example in examples]).reshape(-1, 1)) ] gs_pth = 'gs://apache-beam-ml/models/' \ 'pytorch_lin_reg_model_2x+0.5_state_dict.pth' model_handler = PytorchModelHandlerTensor( state_dict_path=gs_pth, model_class=PytorchLinearRegression, model_params={ 'input_dim': 1, 'output_dim': 1 }) pcoll = pipeline | 'start' >> beam.Create(examples) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to(expected_predictions, equals_fn=_compare_prediction_result)) @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed') def test_pipeline_gcs_model_control_batching(self): with TestPipeline() as pipeline: examples = torch.from_numpy( np.array([1, 5, 3, 10], dtype="float32").reshape(-1, 1)) expected_predictions = [ PredictionResult(ex, pred) for ex, pred in zip( examples, torch.Tensor( [example * 2.0 + 0.5 for example in examples]).reshape(-1, 1)) ] def batch_validator_tensor_inference_fn( batch, model, device, inference_args, model_id, ): if len(batch) != 2: raise Exception( f'Expected batch of size 2, received batch of size {len(batch)}') return default_tensor_inference_fn( batch, model, device, inference_args, model_id) gs_pth = 'gs://apache-beam-ml/models/' \ 'pytorch_lin_reg_model_2x+0.5_state_dict.pth' model_handler = PytorchModelHandlerTensor( state_dict_path=gs_pth, model_class=PytorchLinearRegression, model_params={ 'input_dim': 1, 'output_dim': 1 }, inference_fn=batch_validator_tensor_inference_fn, min_batch_size=2, max_batch_size=2) pcoll = pipeline | 'start' >> beam.Create(examples) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to(expected_predictions, equals_fn=_compare_prediction_result)) def test_invalid_input_type(self): with self.assertRaisesRegex(Exception, "expected Tensor as element"): with TestPipeline() as pipeline: examples = np.array([1, 5, 3, 10], dtype="float32").reshape(-1, 1) state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0]])), ('linear.bias', torch.Tensor([0.5]))]) path = os.path.join(self.tmpdir, 'my_state_dict_path') torch.save(state_dict, path) model_handler = PytorchModelHandlerTensor( state_dict_path=path, model_class=PytorchLinearRegression, model_params={ 'input_dim': 1, 'output_dim': 1 }) pcoll = pipeline | 'start' >> beam.Create(examples) # pylint: disable=expression-not-assigned pcoll | RunInference(model_handler) def test_gpu_auto_convert_to_cpu(self): """ This tests the scenario in which the user defines `device='GPU'` for the PytorchModelHandlerX, but runs the pipeline on a machine without GPU, we automatically detect this discrepancy and do automatic conversion to CPU. A warning is also logged to inform the user. """ with self.assertLogs() as log: with TestPipeline() as pipeline: examples = torch.from_numpy( np.array([1, 5, 3, 10], dtype="float32").reshape(-1, 1)) state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0]])), ('linear.bias', torch.Tensor([0.5]))]) path = os.path.join(self.tmpdir, 'my_state_dict_path') torch.save(state_dict, path) model_handler = PytorchModelHandlerTensor( state_dict_path=path, model_class=PytorchLinearRegression, model_params={ 'input_dim': 1, 'output_dim': 1 }, device='GPU') # Upon initialization, device is cuda self.assertEqual(model_handler._device, torch.device('cuda')) pcoll = pipeline | 'start' >> beam.Create(examples) # pylint: disable=expression-not-assigned pcoll | RunInference(model_handler) # During model loading, device converted to cuda self.assertEqual(model_handler._device, torch.device('cuda')) self.assertIn("INFO:root:Device is set to CUDA", log.output) self.assertIn( "WARNING:root:Model handler specified a 'GPU' device, but GPUs " \ "are not available. Switching to CPU.", log.output) def test_load_torch_script_model(self): torch_model = PytorchLinearRegression(2, 1) torch_script_model = torch.jit.script(torch_model) torch_script_path = os.path.join(self.tmpdir, 'torch_script_model.pt') torch.jit.save(torch_script_model, torch_script_path) model_handler = PytorchModelHandlerTensor( torch_script_model_path=torch_script_path) torch_script_model = model_handler.load_model() self.assertTrue(isinstance(torch_script_model, torch.jit.ScriptModule)) def test_inference_torch_script_model(self): torch_model = PytorchLinearRegression(2, 1) torch_model.load_state_dict( OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])), ('linear.bias', torch.Tensor([0.5]))])) torch_script_model = torch.jit.script(torch_model) torch_script_path = os.path.join(self.tmpdir, 'torch_script_model.pt') torch.jit.save(torch_script_model, torch_script_path) model_handler = PytorchModelHandlerTensor( torch_script_model_path=torch_script_path) with TestPipeline() as pipeline: pcoll = pipeline | 'start' >> beam.Create(TWO_FEATURES_EXAMPLES) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to( TWO_FEATURES_PREDICTIONS, equals_fn=_compare_prediction_result)) def test_torch_model_class_none(self): torch_model = PytorchLinearRegression(2, 1) torch_path = os.path.join(self.tmpdir, 'torch_model.pt') torch.save(torch_model, torch_path) with self.assertRaisesRegex( RuntimeError, "A state_dict_path has been supplied to the model " "handler, but the required model_class is missing. " "Please provide the model_class in order to"): _ = PytorchModelHandlerTensor(state_dict_path=torch_path) with self.assertRaisesRegex( RuntimeError, "A state_dict_path has been supplied to the model " "handler, but the required model_class is missing. " "Please provide the model_class in order to"): _ = (PytorchModelHandlerKeyedTensor(state_dict_path=torch_path)) def test_torch_model_state_dict_none(self): with self.assertRaisesRegex( RuntimeError, "A model_class has been supplied to the model " "handler, but the required state_dict_path is missing. " "Please provide the state_dict_path in order to"): _ = PytorchModelHandlerTensor(model_class=PytorchLinearRegression) with self.assertRaisesRegex( RuntimeError, "A model_class has been supplied to the model " "handler, but the required state_dict_path is missing. " "Please provide the state_dict_path in order to"): _ = PytorchModelHandlerKeyedTensor(model_class=PytorchLinearRegression) def test_specify_torch_script_path_and_state_dict_path(self): torch_model = PytorchLinearRegression(2, 1) torch_path = os.path.join(self.tmpdir, 'torch_model.pt') torch.save(torch_model, torch_path) torch_script_model = torch.jit.script(torch_model) torch_script_path = os.path.join(self.tmpdir, 'torch_script_model.pt') torch.jit.save(torch_script_model, torch_script_path) with self.assertRaisesRegex( RuntimeError, "Please specify either torch_script_model_path or "): _ = PytorchModelHandlerTensor( state_dict_path=torch_path, model_class=PytorchLinearRegression, torch_script_model_path=torch_script_path) def test_prediction_result_model_id_with_torch_script_model(self): torch_model = PytorchLinearRegression(2, 1) torch_script_model = torch.jit.script(torch_model) torch_script_path = os.path.join(self.tmpdir, 'torch_script_model.pt') torch.jit.save(torch_script_model, torch_script_path) model_handler = PytorchModelHandlerTensor( torch_script_model_path=torch_script_path) def check_torch_script_model_id(element): assert ('torch_script_model.pt' in element.model_id) is True with TestPipeline() as pipeline: pcoll = pipeline | 'start' >> beam.Create(TWO_FEATURES_EXAMPLES) predictions = pcoll | RunInference(model_handler) _ = predictions | beam.Map(check_torch_script_model_id) def test_prediction_result_model_id_with_torch_model(self): # weights associated with PytorchLinearRegression class state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])), ('linear.bias', torch.Tensor([0.5]))]) torch_path = os.path.join(self.tmpdir, 'torch_model.pt') torch.save(state_dict, torch_path) model_handler = PytorchModelHandlerTensor( state_dict_path=torch_path, model_class=PytorchLinearRegression, model_params={ 'input_dim': 2, 'output_dim': 1 }) def check_torch_script_model_id(element): assert ('torch_model.pt' in element.model_id) is True with TestPipeline() as pipeline: pcoll = pipeline | 'start' >> beam.Create(TWO_FEATURES_EXAMPLES) predictions = pcoll | RunInference(model_handler) _ = predictions | beam.Map(check_torch_script_model_id) def test_env_vars_set_correctly_tensor_handler(self): torch_model = PytorchLinearRegression(2, 1) torch_model.load_state_dict( OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])), ('linear.bias', torch.Tensor([0.5]))])) torch_script_model = torch.jit.script(torch_model) torch_script_path = os.path.join(self.tmpdir, 'torch_script_model.pt') torch.jit.save(torch_script_model, torch_script_path) handler_with_vars = PytorchModelHandlerTensor( torch_script_model_path=torch_script_path, env_vars={'FOO': 'bar'}) os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) with TestPipeline() as pipeline: _ = ( pipeline | 'start' >> beam.Create(TWO_FEATURES_EXAMPLES) | RunInference(handler_with_vars)) pipeline.run() self.assertTrue('FOO' in os.environ) self.assertTrue((os.environ['FOO']) == 'bar') def test_env_vars_set_correctly_keyed_tensor_handler(self): os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) with TestPipeline() as pipeline: inference_args = { 'prediction_param_array': torch.from_numpy( np.array([1, 2], dtype="float32")), 'prediction_param_bool': True } state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0]])), ('linear.bias', torch.Tensor([0.5]))]) path = os.path.join(self.tmpdir, 'my_state_dict_path') torch.save(state_dict, path) handler_with_vars = PytorchModelHandlerKeyedTensor( env_vars={'FOO': 'bar'}, state_dict_path=path, model_class=PytorchLinearRegressionKeyedBatchAndExtraInferenceArgs, model_params={ 'input_dim': 1, 'output_dim': 1 }) inference_args_side_input = ( pipeline | 'create side' >> beam.Create(inference_args)) _ = ( pipeline | 'start' >> beam.Create(KEYED_TORCH_EXAMPLES) | RunInference( model_handler=handler_with_vars, inference_args=beam.pvalue.AsDict(inference_args_side_input))) pipeline.run() self.assertTrue('FOO' in os.environ) self.assertTrue((os.environ['FOO']) == 'bar') @pytest.mark.uses_pytorch class PytorchInferenceTestWithMocks(unittest.TestCase): def setUp(self): self._load_model = pytorch_inference._load_model pytorch_inference._load_model = unittest.mock.MagicMock( return_value=("model", "device")) self.tmpdir = tempfile.mkdtemp() self.state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])), ('linear.bias', torch.Tensor([0.5]))]) self.torch_path = os.path.join(self.tmpdir, 'torch_model.pt') torch.save(self.state_dict, self.torch_path) self.model_params = {'input_dim': 2, 'output_dim': 1} def tearDown(self): pytorch_inference._load_model = self._load_model shutil.rmtree(self.tmpdir) def test_load_model_args_tensor(self): load_model_args = {'weights_only': True} model_handler = PytorchModelHandlerTensor( state_dict_path=self.torch_path, model_class=PytorchLinearRegression, model_params=self.model_params, load_model_args=load_model_args) model_handler.load_model() pytorch_inference._load_model.assert_called_with( model_class=PytorchLinearRegression, state_dict_path=self.torch_path, device=torch.device('cpu'), model_params=self.model_params, torch_script_model_path=None, load_model_args=load_model_args) def test_load_model_args_keyed_tensor(self): load_model_args = {'weights_only': True} model_handler = PytorchModelHandlerKeyedTensor( state_dict_path=self.torch_path, model_class=PytorchLinearRegression, model_params=self.model_params, load_model_args=load_model_args) model_handler.load_model() pytorch_inference._load_model.assert_called_with( model_class=PytorchLinearRegression, state_dict_path=self.torch_path, device=torch.device('cpu'), model_params=self.model_params, torch_script_model_path=None, load_model_args=load_model_args) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/sklearn_inference.py0000644000175100017510000003121315206113032024252 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import enum import pickle import sys from collections.abc import Callable from collections.abc import Iterable from collections.abc import Sequence from typing import Any from typing import Optional import numpy import pandas from sklearn.base import BaseEstimator from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import PredictionResult try: import joblib except ImportError: # joblib is an optional dependency. pass __all__ = [ 'SklearnModelHandlerNumpy', 'SklearnModelHandlerPandas', ] NumpyInferenceFn = Callable[ [BaseEstimator, Sequence[numpy.ndarray], Optional[dict[str, Any]]], Any] class ModelFileType(enum.Enum): """Defines how a model file is serialized. Options are pickle or joblib.""" PICKLE = 1 JOBLIB = 2 def _load_model(model_uri, file_type): file = FileSystems.open(model_uri, 'rb') if file_type == ModelFileType.PICKLE: return pickle.load(file) elif file_type == ModelFileType.JOBLIB: if not joblib: raise ImportError( 'Could not import joblib in this execution environment. ' 'For help with managing dependencies on Python workers.' 'see https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/' # pylint: disable=line-too-long ) return joblib.load(file) raise AssertionError('Unsupported serialization type.') def _default_numpy_inference_fn( model: BaseEstimator, batch: Sequence[numpy.ndarray], inference_args: Optional[dict[str, Any]] = None) -> Any: inference_args = {} if not inference_args else inference_args # vectorize data for better performance vectorized_batch = numpy.stack(batch, axis=0) return model.predict(vectorized_batch, **inference_args) class SklearnModelHandlerNumpy(ModelHandler[numpy.ndarray, PredictionResult, BaseEstimator]): def __init__( self, model_uri: str, model_file_type: ModelFileType = ModelFileType.PICKLE, *, inference_fn: NumpyInferenceFn = _default_numpy_inference_fn, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, large_model: bool = False, model_copies: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, **kwargs): """ Implementation of the ModelHandler interface for scikit-learn using numpy arrays as input. Example Usage:: pcoll | RunInference(SklearnModelHandlerNumpy(model_uri="my_uri")) Args: model_uri: The URI to where the model is saved. model_file_type: The method of serialization of the argument. default=pickle inference_fn: The inference function to use. default=_default_numpy_inference_fn min_batch_size: the minimum batch size to use when batching inputs. This batch will be fed into the inference_fn as a Sequence of Numpy ndarrays. max_batch_size: the maximum batch size to use when batching inputs. This batch will be fed into the inference_fn as a Sequence of Numpy ndarrays. max_batch_duration_secs: the maximum amount of time to buffer a batch before emitting; used in streaming contexts. large_model: set to true if your model is large enough to run into memory pressure if you load multiple copies. Given a model that consumes N memory and a machine with W cores and M memory, you should set this to True if N*W > M. model_copies: The exact number of models that you would like loaded onto your machine. This can be useful if you exactly know your CPU or GPU capacity and want to maximize resource utilization. max_batch_weight: the maximum total weight of a batch. element_size_fn: a function that returns the size (weight) of an element. batch_length_fn: a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: a sorted list of positive boundary values for length-aware batching buckets. kwargs: 'env_vars' can be used to set environment variables before loading the model. """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries, large_model=large_model, model_copies=model_copies, **kwargs) self._model_uri = model_uri self._model_file_type = model_file_type self._model_inference_fn = inference_fn def load_model(self) -> BaseEstimator: """Loads and initializes a model for processing.""" return _load_model(self._model_uri, self._model_file_type) def update_model_path(self, model_path: Optional[str] = None): self._model_uri = model_path if model_path else self._model_uri def run_inference( self, batch: Sequence[numpy.ndarray], model: BaseEstimator, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """Runs inferences on a batch of numpy arrays. Args: batch: A sequence of examples as numpy arrays. They should be single examples. model: A numpy model or pipeline. Must implement predict(X). Where the parameter X is a numpy array. inference_args: Any additional arguments for an inference. Returns: An Iterable of type PredictionResult. """ predictions = self._model_inference_fn( model, batch, inference_args, ) return utils._convert_to_result( batch, predictions, model_id=self._model_uri) def get_num_bytes(self, batch: Sequence[numpy.ndarray]) -> int: """ Returns: The number of bytes of data for a batch. """ return sum(sys.getsizeof(element) for element in batch) def get_metrics_namespace(self) -> str: """ Returns: A namespace for metrics collected by the RunInference transform. """ return 'BeamML_Sklearn' PandasInferenceFn = Callable[ [BaseEstimator, Sequence[pandas.DataFrame], Optional[dict[str, Any]]], Any] def _default_pandas_inference_fn( model: BaseEstimator, batch: Sequence[pandas.DataFrame], inference_args: Optional[dict[str, Any]] = None) -> Any: # vectorize data for better performance vectorized_batch = pandas.concat(batch, axis=0) predictions = model.predict(vectorized_batch) splits = [ vectorized_batch.iloc[[i]] for i in range(vectorized_batch.shape[0]) ] return predictions, splits class SklearnModelHandlerPandas(ModelHandler[pandas.DataFrame, PredictionResult, BaseEstimator]): def __init__( self, model_uri: str, model_file_type: ModelFileType = ModelFileType.PICKLE, *, inference_fn: PandasInferenceFn = _default_pandas_inference_fn, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, large_model: bool = False, model_copies: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, **kwargs): """Implementation of the ModelHandler interface for scikit-learn that supports pandas dataframes. Example Usage:: pcoll | RunInference(SklearnModelHandlerPandas(model_uri="my_uri")) **NOTE:** This API and its implementation are under development and do not provide backward compatibility guarantees. Args: model_uri: The URI to where the model is saved. model_file_type: The method of serialization of the argument. default=pickle inference_fn: The inference function to use. default=_default_pandas_inference_fn min_batch_size: the minimum batch size to use when batching inputs. This batch will be fed into the inference_fn as a Sequence of Pandas Dataframes. max_batch_size: the maximum batch size to use when batching inputs. This batch will be fed into the inference_fn as a Sequence of Pandas Dataframes. max_batch_duration_secs: the maximum amount of time to buffer a batch before emitting; used in streaming contexts. large_model: set to true if your model is large enough to run into memory pressure if you load multiple copies. Given a model that consumes N memory and a machine with W cores and M memory, you should set this to True if N*W > M. model_copies: The exact number of models that you would like loaded onto your machine. This can be useful if you exactly know your CPU or GPU capacity and want to maximize resource utilization. max_batch_weight: the maximum total weight of a batch. element_size_fn: a function that returns the size (weight) of an element. batch_length_fn: a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: a sorted list of positive boundary values for length-aware batching buckets. kwargs: 'env_vars' can be used to set environment variables before loading the model. """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries, large_model=large_model, model_copies=model_copies, **kwargs) self._model_uri = model_uri self._model_file_type = model_file_type self._model_inference_fn = inference_fn def load_model(self) -> BaseEstimator: """Loads and initializes a model for processing.""" return _load_model(self._model_uri, self._model_file_type) def update_model_path(self, model_path: Optional[str] = None): self._model_uri = model_path if model_path else self._model_uri def run_inference( self, batch: Sequence[pandas.DataFrame], model: BaseEstimator, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """ Runs inferences on a batch of pandas dataframes. Args: batch: A sequence of examples as numpy arrays. They should be single examples. model: A dataframe model or pipeline. Must implement predict(X). Where the parameter X is a pandas dataframe. inference_args: Any additional arguments for an inference. Returns: An Iterable of type PredictionResult. """ # sklearn_inference currently only supports single rowed dataframes. for dataframe in iter(batch): if dataframe.shape[0] != 1: raise ValueError('Only dataframes with single rows are supported.') predictions, splits = self._model_inference_fn(model, batch, inference_args) return utils._convert_to_result( splits, predictions, model_id=self._model_uri) def get_num_bytes(self, batch: Sequence[pandas.DataFrame]) -> int: """ Returns: The number of bytes of data for a batch. """ return sum(df.memory_usage(deep=True).sum() for df in batch) def get_metrics_namespace(self) -> str: """ Returns: A namespace for metrics collected by the RunInference transform. """ return 'BeamML_Sklearn' ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/sklearn_inference_it_test.py0000644000175100017510000001570715206113032026017 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-End test for Sklearn Inference""" import logging import re import sys import unittest import uuid import pytest from apache_beam.examples.inference import sklearn_japanese_housing_regression from apache_beam.examples.inference import sklearn_mnist_classification from apache_beam.io.filesystems import FileSystems from apache_beam.testing.test_pipeline import TestPipeline # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports, unused-import try: from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem except ImportError: raise unittest.SkipTest('GCP dependencies are not installed') def process_outputs(filepath): with FileSystems().open(filepath) as f: lines = f.readlines() lines = [l.decode('utf-8').strip('\n') for l in lines] return lines def file_lines_sorted(filepath): with FileSystems().open(filepath) as f: lines = f.readlines() lines = [l.decode('utf-8').strip('\n') for l in lines] return sorted(lines) @pytest.mark.uses_sklearn @pytest.mark.it_postcommit class SklearnInference(unittest.TestCase): # TODO(https://github.com/apache/beam/issues/33796) use older numpy @unittest.skipIf((3, 9, 0) <= sys.version_info < (3, 11, 0), "Beam#33796") def test_sklearn_mnist_classification(self): test_pipeline = TestPipeline(is_integration_test=True) input_file = 'gs://apache-beam-ml/testing/inputs/it_mnist_data.csv' output_file_dir = 'gs://temp-storage-for-end-to-end-tests' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_path = 'gs://apache-beam-ml/models/mnist_model_svm.pickle' extra_opts = { 'input': input_file, 'output': output_file, 'model_path': model_path, } sklearn_mnist_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_output_filepath = 'gs://apache-beam-ml/testing/expected_outputs/test_sklearn_mnist_classification_actuals.txt' # pylint: disable=line-too-long expected_outputs = process_outputs(expected_output_filepath) predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) predictions_dict = {} for i in range(len(predicted_outputs)): true_label, prediction = predicted_outputs[i].split(',') predictions_dict[true_label] = prediction for i in range(len(expected_outputs)): true_label, expected_prediction = expected_outputs[i].split(',') self.assertEqual(predictions_dict[true_label], expected_prediction) # TODO(https://github.com/apache/beam/issues/33796) use older numpy @unittest.skipIf((3, 9, 0) <= sys.version_info < (3, 11, 0), "Beam#33796") def test_sklearn_mnist_classification_large_model(self): test_pipeline = TestPipeline(is_integration_test=True) input_file = 'gs://apache-beam-ml/testing/inputs/it_mnist_data.csv' output_file_dir = 'gs://temp-storage-for-end-to-end-tests' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_path = 'gs://apache-beam-ml/models/mnist_model_svm.pickle' extra_opts = { 'input': input_file, 'output': output_file, 'model_path': model_path, 'large_model': True } sklearn_mnist_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_output_filepath = 'gs://apache-beam-ml/testing/expected_outputs/test_sklearn_mnist_classification_actuals.txt' # pylint: disable=line-too-long expected_outputs = process_outputs(expected_output_filepath) predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) predictions_dict = {} for i in range(len(predicted_outputs)): true_label, prediction = predicted_outputs[i].split(',') predictions_dict[true_label] = prediction for i in range(len(expected_outputs)): true_label, expected_prediction = expected_outputs[i].split(',') self.assertEqual(predictions_dict[true_label], expected_prediction) # TODO(https://github.com/apache/beam/issues/27151) use model with sklearn 1.2 # TODO(https://github.com/apache/beam/issues/33796) use older numpy @unittest.skipIf(sys.version_info >= (3, 9, 0), "Beam#27151") def test_sklearn_regression(self): test_pipeline = TestPipeline(is_integration_test=True) input_file = 'gs://apache-beam-ml/testing/inputs/japanese_housing_test_data.csv' # pylint: disable=line-too-long output_file_dir = 'gs://temp-storage-for-end-to-end-tests' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_path = 'gs://apache-beam-ml/models/japanese_housing/' extra_opts = { 'input': input_file, 'output': output_file, 'model_path': model_path, } sklearn_japanese_housing_regression.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_output_filepath = 'gs://apache-beam-ml/testing/expected_outputs/japanese_housing_subset.txt' # pylint: disable=line-too-long expected_outputs = file_lines_sorted(expected_output_filepath) actual_outputs = file_lines_sorted(output_file) self.assertEqual(len(expected_outputs), len(actual_outputs)) for expected, actual in zip(expected_outputs, actual_outputs): expected_true, expected_predict = re.findall(r'\d+', expected) actual_true, actual_predict = re.findall(r'\d+', actual) # actual_true is the y value from the input csv file. # Therefore it should be an exact match to expected_true. self.assertEqual(actual_true, expected_true) # predictions might not be exactly equal due to differences between # environments. This code validates they are within 10 percent. percent_diff = abs(float(expected_predict) - float(actual_predict) ) / float(expected_predict) * 100.0 self.assertLess(percent_diff, 10) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/sklearn_inference_test.py0000644000175100017510000005510615206113032025320 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import io import math import os import pickle import platform import shutil import sys import tempfile import unittest from collections.abc import Sequence from typing import Any from typing import Optional import joblib import numpy import pandas from sklearn import linear_model from sklearn import svm from sklearn.base import BaseEstimator from sklearn.compose import ColumnTransformer from sklearn.pipeline import Pipeline from sklearn.preprocessing import OneHotEncoder from sklearn.preprocessing import StandardScaler import apache_beam as beam from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.sklearn_inference import ModelFileType from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerNumpy from apache_beam.ml.inference.sklearn_inference import SklearnModelHandlerPandas from apache_beam.ml.inference.sklearn_inference import _default_numpy_inference_fn from apache_beam.ml.inference.sklearn_inference import _default_pandas_inference_fn from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to def _compare_prediction_result(a, b): example_equal = numpy.array_equal(a.example, b.example) if isinstance(a.inference, dict): return all( x == y for x, y in zip(a.inference.values(), b.inference.values())) and example_equal return a.inference == b.inference and example_equal def _compare_dataframe_predictions(a_in, b_in): keys_equal = True if isinstance(a_in, tuple) and not isinstance(a_in, PredictionResult): a_key, a = a_in b_key, b = b_in keys_equal = a_key == b_key else: a = a_in b = b_in example_equal = pandas.DataFrame.equals(a.example, b.example) if isinstance(a.inference, dict): return all( math.floor(a) == math.floor(b) for a, b in zip( a.inference.values(), b.inference.values())) and example_equal inference_equal = math.floor(a.inference) == math.floor(b.inference) return inference_equal and example_equal and keys_equal class FakeModel: def __init__(self): self.total_predict_calls = 0 def predict(self, input_vector: numpy.ndarray): self.total_predict_calls += 1 return numpy.sum(input_vector, axis=1) class FakeNumpyModelDictOut: def __init__(self): self.total_predict_calls = 0 def predict(self, input_vector: numpy.ndarray): self.total_predict_calls += 1 out = numpy.sum(input_vector, axis=1) return {"out1": out, "out2": out} class FakePandasModelDictOut: def __init__(self): self.total_predict_calls = 0 def predict(self, df: pandas.DataFrame): self.total_predict_calls += 1 out = df.loc[:, 'number_2'] return {"out1": out, "out2": out} def build_model(): x = [[0, 0], [1, 1]] y = [0, 1] model = svm.SVC() model.fit(x, y) return model def pandas_dataframe(): csv_string = ( 'category_1,number_1,category_2,number_2,label,number_3\n' 'red,4,frog,5,6,7\n' 'blue,3,horse,8,9,10\n' 'red,0,cow,1,2,3\n' 'blue,4,frog,1,1,1\n' 'red,1,horse,4,2,3') csv_string_io = io.StringIO(csv_string) return pandas.read_csv(csv_string_io) def build_pandas_pipeline(): """Builds a common type of pandas pipeline with preprocessing.""" categorical_columns = ['category_1', 'category_2'] numerical_columns = ['number_1', 'number_2', 'number_3'] categorical_transformer = OneHotEncoder(handle_unknown='ignore') numerical_transformer = StandardScaler() preprocessor = ColumnTransformer( transformers=[ ("numerical", numerical_transformer, numerical_columns), ("categorical", categorical_transformer, categorical_columns), ]) pipeline = Pipeline( steps=[("preprocessor", preprocessor), ("classifier", linear_model.SGDRegressor())]) data = pandas_dataframe() labels = data['label'] pipeline.fit(data, labels) return pipeline def convert_inference_to_floor(prediction_result): return math.floor(prediction_result.inference) def alternate_numpy_inference_fn( model: BaseEstimator, batch: Sequence[numpy.ndarray], inference_args: Optional[dict[str, Any]] = None) -> Any: return [0] def alternate_pandas_inference_fn( model: BaseEstimator, batch: Sequence[pandas.DataFrame], inference_args: Optional[dict[str, Any]] = None) -> Any: # vectorize data for better performance vectorized_batch = pandas.concat(batch, axis=0) predictions = model.predict(vectorized_batch) splits = [ vectorized_batch.iloc[[i]] for i in range(vectorized_batch.shape[0]) ] predictions = predictions - 1 return predictions, splits class SkLearnRunInferenceTest(unittest.TestCase): def setUp(self): self.tmpdir = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.tmpdir) def test_predict_output(self): fake_model = FakeModel() inference_runner = SklearnModelHandlerNumpy(model_uri='unused') batched_examples = [ numpy.array([1, 2, 3]), numpy.array([4, 5, 6]), numpy.array([7, 8, 9]) ] expected_predictions = [ PredictionResult(numpy.array([1, 2, 3]), 6), PredictionResult(numpy.array([4, 5, 6]), 15), PredictionResult(numpy.array([7, 8, 9]), 24) ] inferences = inference_runner.run_inference(batched_examples, fake_model) for actual, expected in zip(inferences, expected_predictions): self.assertTrue(_compare_prediction_result(actual, expected)) def test_custom_inference_fn(self): fake_model = FakeModel() inference_runner = SklearnModelHandlerNumpy( model_uri='unused', inference_fn=alternate_numpy_inference_fn) batched_examples = [ numpy.array([1, 2, 3]), numpy.array([4, 5, 6]), numpy.array([7, 8, 9]) ] expected_predictions = [ PredictionResult(numpy.array([1, 2, 3]), 0), PredictionResult(numpy.array([4, 5, 6]), 0), PredictionResult(numpy.array([7, 8, 9]), 0) ] inferences = inference_runner.run_inference(batched_examples, fake_model) for actual, expected in zip(inferences, expected_predictions): self.assertTrue(_compare_prediction_result(actual, expected)) def test_predict_output_dict(self): fake_model = FakeNumpyModelDictOut() inference_runner = SklearnModelHandlerNumpy(model_uri='unused') batched_examples = [ numpy.array([1, 2, 3]), numpy.array([4, 5, 6]), numpy.array([7, 8, 9]) ] expected_predictions = [ PredictionResult(numpy.array([1, 2, 3]), { "out1": 6, "out2": 6 }), PredictionResult(numpy.array([4, 5, 6]), { "out1": 15, "out2": 15 }), PredictionResult(numpy.array([7, 8, 9]), { "out1": 24, "out2": 24 }) ] inferences = inference_runner.run_inference(batched_examples, fake_model) for actual, expected in zip(inferences, expected_predictions): self.assertTrue(_compare_prediction_result(actual, expected)) def test_data_vectorized(self): fake_model = FakeModel() inference_runner = SklearnModelHandlerNumpy(model_uri='unused') batched_examples = [ numpy.array([1, 2, 3]), numpy.array([4, 5, 6]), numpy.array([7, 8, 9]) ] # even though there are 3 examples, the data should # be vectorized and only 1 call should happen. inference_runner.run_inference(batched_examples, fake_model) self.assertEqual(1, fake_model.total_predict_calls) def test_num_bytes_numpy(self): inference_runner = SklearnModelHandlerNumpy(model_uri='unused') batched_examples_int = [ numpy.array([1, 2, 3]), numpy.array([4, 5, 6]), numpy.array([7, 8, 9]) ] self.assertEqual( sys.getsizeof(batched_examples_int[0]) * 3, inference_runner.get_num_bytes(batched_examples_int)) batched_examples_float = [ numpy.array([1.0, 2.0, 3.0]), numpy.array([4.1, 5.2, 6.3]), numpy.array([7.7, 8.8, 9.9]) ] self.assertEqual( sys.getsizeof(batched_examples_float[0]) * 3, inference_runner.get_num_bytes(batched_examples_float)) def test_pipeline_pickled(self): temp_file_name = self.tmpdir + os.sep + 'pickled_file' with open(temp_file_name, 'wb') as file: pickle.dump(build_model(), file) with TestPipeline() as pipeline: examples = [numpy.array([0, 0]), numpy.array([1, 1])] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | RunInference( SklearnModelHandlerNumpy(model_uri=temp_file_name)) expected = [ PredictionResult(numpy.array([0, 0]), 0), PredictionResult(numpy.array([1, 1]), 1) ] assert_that( actual, equal_to(expected, equals_fn=_compare_prediction_result)) def test_pipeline_pickled_custom_batching(self): temp_file_name = self.tmpdir + os.sep + 'pickled_file' with open(temp_file_name, 'wb') as file: pickle.dump(build_model(), file) def batch_validator_numpy_inference_fn( model: BaseEstimator, batch: Sequence[numpy.ndarray], inference_args: Optional[dict[str, Any]] = None) -> Any: if len(batch) != 2: raise Exception( f'Expected batch of size 2, received batch of size {len(batch)}') return _default_numpy_inference_fn(model, batch, inference_args) with TestPipeline() as pipeline: examples = [numpy.array([0, 0]), numpy.array([1, 1])] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | RunInference( SklearnModelHandlerNumpy( model_uri=temp_file_name, inference_fn=batch_validator_numpy_inference_fn, min_batch_size=2, max_batch_size=2)) expected = [ PredictionResult(numpy.array([0, 0]), 0), PredictionResult(numpy.array([1, 1]), 1) ] assert_that( actual, equal_to(expected, equals_fn=_compare_prediction_result)) def test_pipeline_pickled_large_model(self): temp_file_name = self.tmpdir + os.sep + 'pickled_file' with open(temp_file_name, 'wb') as file: pickle.dump(build_model(), file) def large_model_validator_numpy_inference_fn( model: BaseEstimator, batch: Sequence[numpy.ndarray], inference_args: Optional[dict[str, Any]] = None) -> Any: multi_process_shared_loaded = "multi_process_shared" in str(type(model)) if not multi_process_shared_loaded: raise Exception( f'Loaded model of type {type(model)}, was ' + 'expecting multi_process_shared_model') return _default_numpy_inference_fn(model, batch, inference_args) with TestPipeline() as pipeline: examples = [numpy.array([0, 0]), numpy.array([1, 1])] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | RunInference( SklearnModelHandlerNumpy( model_uri=temp_file_name, inference_fn=large_model_validator_numpy_inference_fn, large_model=True)) expected = [ PredictionResult(numpy.array([0, 0]), 0), PredictionResult(numpy.array([1, 1]), 1) ] assert_that( actual, equal_to(expected, equals_fn=_compare_prediction_result)) def test_pipeline_joblib(self): temp_file_name = self.tmpdir + os.sep + 'joblib_file' with open(temp_file_name, 'wb') as file: joblib.dump(build_model(), file) with TestPipeline() as pipeline: examples = [numpy.array([0, 0]), numpy.array([1, 1])] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | RunInference( SklearnModelHandlerNumpy( model_uri=temp_file_name, model_file_type=ModelFileType.JOBLIB)) expected = [ PredictionResult(numpy.array([0, 0]), 0), PredictionResult(numpy.array([1, 1]), 1) ] assert_that( actual, equal_to(expected, equals_fn=_compare_prediction_result)) def test_bad_file_raises(self): with self.assertRaises(RuntimeError): with TestPipeline() as pipeline: examples = [numpy.array([0, 0])] pcoll = pipeline | 'start' >> beam.Create(examples) _ = pcoll | RunInference( SklearnModelHandlerNumpy(model_uri='/var/bad_file_name')) pipeline.run() def test_bad_input_type_raises(self): with self.assertRaisesRegex(AssertionError, 'Unsupported serialization type'): with tempfile.NamedTemporaryFile(delete=False) as file: model_handler = SklearnModelHandlerNumpy( model_uri=file.name, model_file_type=None) model_handler.load_model() def test_env_vars_set_correctly_numpy(self): temp_file_name = self.tmpdir + os.sep + 'pickled_file' with open(temp_file_name, 'wb') as file: pickle.dump(build_model(), file) handler_with_vars = SklearnModelHandlerNumpy( env_vars={'FOO': 'bar'}, model_uri=temp_file_name) os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) examples = [numpy.array([0, 0]), numpy.array([1, 1])] with TestPipeline() as pipeline: _ = ( pipeline | 'start' >> beam.Create(examples) | RunInference(handler_with_vars)) pipeline.run() self.assertTrue('FOO' in os.environ) self.assertTrue((os.environ['FOO']) == 'bar') def test_pipeline_pandas(self): temp_file_name = self.tmpdir + os.sep + 'pickled_file' with open(temp_file_name, 'wb') as file: pickle.dump(build_pandas_pipeline(), file) with TestPipeline() as pipeline: dataframe = pandas_dataframe() splits = [dataframe.loc[[i]] for i in dataframe.index] pcoll = pipeline | 'start' >> beam.Create(splits) actual = pcoll | RunInference( SklearnModelHandlerPandas(model_uri=temp_file_name)) expected = [ PredictionResult(splits[0], 5), PredictionResult(splits[1], 8), PredictionResult(splits[2], 1), PredictionResult(splits[3], 1), PredictionResult(splits[4], 2), ] assert_that( actual, equal_to(expected, equals_fn=_compare_dataframe_predictions)) def test_pipeline_pandas_env_vars_set_correctly(self): temp_file_name = self.tmpdir + os.sep + 'pickled_file' with open(temp_file_name, 'wb') as file: pickle.dump(build_pandas_pipeline(), file) handler_with_vars = SklearnModelHandlerPandas( env_vars={'FOO': 'bar'}, model_uri=temp_file_name) os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) with TestPipeline() as pipeline: dataframe = pandas_dataframe() splits = [dataframe.loc[[i]] for i in dataframe.index] _ = ( pipeline | 'start' >> beam.Create(splits) | RunInference(handler_with_vars)) pipeline.run() self.assertTrue('FOO' in os.environ) self.assertTrue((os.environ['FOO']) == 'bar') def test_pipeline_pandas_custom_batching(self): temp_file_name = self.tmpdir + os.sep + 'pickled_file' with open(temp_file_name, 'wb') as file: pickle.dump(build_pandas_pipeline(), file) def batch_validator_pandas_inference_fn( model: BaseEstimator, batch: Sequence[numpy.ndarray], inference_args: Optional[dict[str, Any]] = None) -> Any: if len(batch) != 5: raise Exception( f'Expected batch of size 5, received batch of size {len(batch)}') return _default_pandas_inference_fn(model, batch, inference_args) with TestPipeline() as pipeline: dataframe = pandas_dataframe() splits = [dataframe.loc[[i]] for i in dataframe.index] pcoll = pipeline | 'start' >> beam.Create(splits) actual = pcoll | RunInference( SklearnModelHandlerPandas( model_uri=temp_file_name, inference_fn=batch_validator_pandas_inference_fn, min_batch_size=5, max_batch_size=5)) expected = [ PredictionResult(splits[0], 5), PredictionResult(splits[1], 8), PredictionResult(splits[2], 1), PredictionResult(splits[3], 1), PredictionResult(splits[4], 2), ] assert_that( actual, equal_to(expected, equals_fn=_compare_dataframe_predictions)) def test_pipeline_pandas_large_model(self): temp_file_name = self.tmpdir + os.sep + 'pickled_file' with open(temp_file_name, 'wb') as file: pickle.dump(build_pandas_pipeline(), file) def large_model_validator_pandas_inference_fn( model: BaseEstimator, batch: Sequence[numpy.ndarray], inference_args: Optional[dict[str, Any]] = None) -> Any: multi_process_shared_loaded = "multi_process_shared" in str(type(model)) if not multi_process_shared_loaded: raise Exception( f'Loaded model of type {type(model)}, was ' + 'expecting multi_process_shared_model') return _default_pandas_inference_fn(model, batch, inference_args) with TestPipeline() as pipeline: dataframe = pandas_dataframe() splits = [dataframe.loc[[i]] for i in dataframe.index] pcoll = pipeline | 'start' >> beam.Create(splits) actual = pcoll | RunInference( SklearnModelHandlerPandas( model_uri=temp_file_name, inference_fn=large_model_validator_pandas_inference_fn, large_model=True)) expected = [ PredictionResult(splits[0], 5), PredictionResult(splits[1], 8), PredictionResult(splits[2], 1), PredictionResult(splits[3], 1), PredictionResult(splits[4], 2), ] assert_that( actual, equal_to(expected, equals_fn=_compare_dataframe_predictions)) def test_pipeline_pandas_custom_inference(self): temp_file_name = self.tmpdir + os.sep + 'pickled_file' with open(temp_file_name, 'wb') as file: pickle.dump(build_pandas_pipeline(), file) with TestPipeline() as pipeline: dataframe = pandas_dataframe() splits = [dataframe.loc[[i]] for i in dataframe.index] pcoll = pipeline | 'start' >> beam.Create(splits) actual = pcoll | RunInference( SklearnModelHandlerPandas( model_uri=temp_file_name, inference_fn=alternate_pandas_inference_fn)) expected = [ PredictionResult(splits[0], 4), PredictionResult(splits[1], 7), PredictionResult(splits[2], 0), PredictionResult(splits[3], 0), PredictionResult(splits[4], 1), ] assert_that( actual, equal_to(expected, equals_fn=_compare_dataframe_predictions)) def test_pipeline_pandas_dict_out(self): temp_file_name = self.tmpdir + os.sep + 'pickled_file' with open(temp_file_name, 'wb') as file: pickle.dump(FakePandasModelDictOut(), file) with TestPipeline() as pipeline: dataframe = pandas_dataframe() splits = [dataframe.loc[[i]] for i in dataframe.index] pcoll = pipeline | 'start' >> beam.Create(splits) actual = pcoll | RunInference( SklearnModelHandlerPandas(model_uri=temp_file_name)) expected = [ PredictionResult(splits[0], { 'out1': 5, 'out2': 5 }), PredictionResult(splits[1], { 'out1': 8, 'out2': 8 }), PredictionResult(splits[2], { 'out1': 1, 'out2': 1 }), PredictionResult(splits[3], { 'out1': 1, 'out2': 1 }), PredictionResult(splits[4], { 'out1': 4, 'out2': 4 }), ] assert_that( actual, equal_to(expected, equals_fn=_compare_dataframe_predictions)) @unittest.skipIf(platform.system() == 'Windows', 'BEAM-14359') def test_pipeline_pandas_joblib(self): temp_file_name = self.tmpdir + os.sep + 'pickled_file' with open(temp_file_name, 'wb') as file: joblib.dump(build_pandas_pipeline(), file) with TestPipeline() as pipeline: dataframe = pandas_dataframe() splits = [dataframe.loc[[i]] for i in dataframe.index] pcoll = pipeline | 'start' >> beam.Create(splits) actual = pcoll | RunInference( SklearnModelHandlerPandas( model_uri=temp_file_name, model_file_type=ModelFileType.JOBLIB)) expected = [ PredictionResult(splits[0], 5), PredictionResult(splits[1], 8), PredictionResult(splits[2], 1), PredictionResult(splits[3], 1), PredictionResult(splits[4], 2), ] assert_that( actual, equal_to(expected, equals_fn=_compare_dataframe_predictions)) def test_pipeline_pandas_with_keys(self): temp_file_name = self.tmpdir + os.sep + 'pickled_file' with open(temp_file_name, 'wb') as file: pickle.dump(build_pandas_pipeline(), file) with TestPipeline() as pipeline: data_frame = pandas_dataframe() keys = [str(i) for i in range(5)] splits = [data_frame.loc[[i]] for i in data_frame.index] keyed_rows = [(key, value) for key, value in zip(keys, splits)] pcoll = pipeline | 'start' >> beam.Create(keyed_rows) actual = pcoll | RunInference( KeyedModelHandler( SklearnModelHandlerPandas(model_uri=temp_file_name))) expected = [ ('0', PredictionResult(splits[0], 5)), ('1', PredictionResult(splits[1], 8)), ('2', PredictionResult(splits[2], 1)), ('3', PredictionResult(splits[3], 1)), ('4', PredictionResult(splits[4], 2)), ] assert_that( actual, equal_to(expected, equals_fn=_compare_dataframe_predictions)) def test_infer_too_many_rows_in_dataframe(self): with self.assertRaisesRegex( ValueError, r'Only dataframes with single rows are supported'): data_frame_too_many_rows = pandas_dataframe() fake_model = FakeModel() inference_runner = SklearnModelHandlerPandas(model_uri='unused') inference_runner.run_inference([data_frame_too_many_rows], fake_model) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/tensorflow_inference.py0000644000175100017510000003456715206113032025034 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import enum import sys from collections.abc import Callable from collections.abc import Iterable from collections.abc import Sequence from typing import Any from typing import Optional from typing import Union import numpy import tensorflow as tf import tensorflow_hub as hub from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import PredictionResult __all__ = [ 'TFModelHandlerNumpy', 'TFModelHandlerTensor', ] TensorInferenceFn = Callable[[ tf.Module, Sequence[Union[numpy.ndarray, tf.Tensor]], dict[str, Any], Optional[str] ], Iterable[PredictionResult]] class ModelType(enum.Enum): """Defines how a model file should be loaded.""" SAVED_MODEL = 1 SAVED_WEIGHTS = 2 def _load_model(model_uri, custom_weights, load_model_args): try: model = tf.keras.models.load_model( hub.resolve(model_uri), **load_model_args) except Exception as e: raise ValueError( "Unable to load the TensorFlow model: {exception}. Make sure you've \ saved the model with TF2 format. Check out the list of TF2 Models on \ TensorFlow Hub - https://tfhub.dev/s?subtype=module,placeholder&tf-version=tf2." # pylint: disable=line-too-long .format(exception=e)) if custom_weights: model.load_weights(custom_weights) return model def _load_model_from_weights(create_model_fn, weights_path): model = create_model_fn() model.load_weights(weights_path) return model def default_numpy_inference_fn( model: tf.Module, batch: Sequence[numpy.ndarray], inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: vectorized_batch = numpy.stack(batch, axis=0) predictions = model(vectorized_batch, **inference_args) return utils._convert_to_result(batch, predictions, model_id) def default_tensor_inference_fn( model: tf.Module, batch: Sequence[tf.Tensor], inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: vectorized_batch = tf.stack(batch, axis=0) predictions = model(vectorized_batch, **inference_args) return utils._convert_to_result(batch, predictions, model_id) class TFModelHandlerNumpy(ModelHandler[numpy.ndarray, PredictionResult, tf.Module]): def __init__( self, model_uri: str, model_type: ModelType = ModelType.SAVED_MODEL, create_model_fn: Optional[Callable] = None, *, load_model_args: Optional[dict[str, Any]] = None, custom_weights: str = "", inference_fn: TensorInferenceFn = default_numpy_inference_fn, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, large_model: bool = False, model_copies: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, **kwargs): """Implementation of the ModelHandler interface for Tensorflow. Example Usage:: pcoll | RunInference(TFModelHandlerNumpy(model_uri="my_uri")) See https://www.tensorflow.org/tutorials/keras/save_and_load for details. Args: model_uri (str): path to the trained model. model_type: type of model to be loaded. Defaults to SAVED_MODEL. create_model_fn: a function that creates and returns a new tensorflow model to load the saved weights. It should be used with ModelType.SAVED_WEIGHTS. load_model_args: a dictionary of parameters to pass to the load_model function of TensorFlow to specify custom config. custom_weights (str): path to the custom weights to be applied once the model is loaded. inference_fn: inference function to use during RunInference. Defaults to default_numpy_inference_fn. large_model: set to true if your model is large enough to run into memory pressure if you load multiple copies. Given a model that consumes N memory and a machine with W cores and M memory, you should set this to True if N*W > M. model_copies: The exact number of models that you would like loaded onto your machine. This can be useful if you exactly know your CPU or GPU capacity and want to maximize resource utilization. max_batch_weight: the maximum total weight of a batch. element_size_fn: a function that returns the size (weight) of an element. batch_length_fn: a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: a sorted list of positive boundary values for length-aware batching buckets. kwargs: 'env_vars' can be used to set environment variables before loading the model. **Supported Versions:** RunInference APIs in Apache Beam have been tested with Tensorflow 2.9, 2.10, 2.11. """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries, large_model=large_model, model_copies=model_copies, **kwargs) self._model_uri = model_uri self._model_type = model_type self._inference_fn = inference_fn self._create_model_fn = create_model_fn self._load_model_args = {} if not load_model_args else load_model_args self._custom_weights = custom_weights def load_model(self) -> tf.Module: """Loads and initializes a Tensorflow model for processing.""" if self._model_type == ModelType.SAVED_WEIGHTS: if not self._create_model_fn: raise ValueError( "Callable create_model_fn must be passed" "with ModelType.SAVED_WEIGHTS") return _load_model_from_weights(self._create_model_fn, self._model_uri) return _load_model( self._model_uri, self._custom_weights, self._load_model_args) def update_model_path(self, model_path: Optional[str] = None): self._model_uri = model_path if model_path else self._model_uri def run_inference( self, batch: Sequence[numpy.ndarray], model: tf.Module, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """ Runs inferences on a batch of numpy array and returns an Iterable of numpy array Predictions. This method stacks the n-dimensional numpy array in a vectorized format to optimize the inference call. Args: batch: A sequence of numpy nd-array. These should be batchable, as this method will call `numpy.stack()` and pass in batched numpy nd-array with dimensions (batch_size, n_features, etc.) into the model's predict() function. model: A Tensorflow model. inference_args: any additional arguments for an inference. Returns: An Iterable of type PredictionResult. """ inference_args = {} if not inference_args else inference_args return self._inference_fn(model, batch, inference_args, self._model_uri) def get_num_bytes(self, batch: Sequence[numpy.ndarray]) -> int: """ Returns: The number of bytes of data for a batch of numpy arrays. """ return sum(sys.getsizeof(element) for element in batch) def get_metrics_namespace(self) -> str: """ Returns: A namespace for metrics collected by the RunInference transform. """ return 'BeamML_TF_Numpy' class TFModelHandlerTensor(ModelHandler[tf.Tensor, PredictionResult, tf.Module]): def __init__( self, model_uri: str, model_type: ModelType = ModelType.SAVED_MODEL, create_model_fn: Optional[Callable] = None, *, load_model_args: Optional[dict[str, Any]] = None, custom_weights: str = "", inference_fn: TensorInferenceFn = default_tensor_inference_fn, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, large_model: bool = False, model_copies: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, **kwargs): """Implementation of the ModelHandler interface for Tensorflow. Example Usage:: pcoll | RunInference(TFModelHandlerTensor(model_uri="my_uri")) See https://www.tensorflow.org/tutorials/keras/save_and_load for details. Args: model_uri (str): path to the trained model. model_type: type of model to be loaded. Defaults to SAVED_MODEL. create_model_fn: a function that creates and returns a new tensorflow model to load the saved weights. It should be used with ModelType.SAVED_WEIGHTS. load_model_args: a dictionary of parameters to pass to the load_model function of TensorFlow to specify custom config. custom_weights (str): path to the custom weights to be applied once the model is loaded. inference_fn: inference function to use during RunInference. Defaults to default_numpy_inference_fn. min_batch_size: the minimum batch size to use when batching inputs. max_batch_size: the maximum batch size to use when batching inputs. max_batch_duration_secs: the maximum amount of time to buffer a batch before emitting; used in streaming contexts. large_model: set to true if your model is large enough to run into memory pressure if you load multiple copies. Given a model that consumes N memory and a machine with W cores and M memory, you should set this to True if N*W > M. model_copies: The exact number of models that you would like loaded onto your machine. This can be useful if you exactly know your CPU or GPU capacity and want to maximize resource utilization. max_batch_weight: the maximum total weight of a batch. element_size_fn: a function that returns the size (weight) of an element. batch_length_fn: a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: a sorted list of positive boundary values for length-aware batching buckets. kwargs: 'env_vars' can be used to set environment variables before loading the model. **Supported Versions:** RunInference APIs in Apache Beam have been tested with Tensorflow 2.11. """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries, large_model=large_model, model_copies=model_copies, **kwargs) self._model_uri = model_uri self._model_type = model_type self._inference_fn = inference_fn self._create_model_fn = create_model_fn self._load_model_args = {} if not load_model_args else load_model_args self._custom_weights = custom_weights def load_model(self) -> tf.Module: """Loads and initializes a tensorflow model for processing.""" if self._model_type == ModelType.SAVED_WEIGHTS: if not self._create_model_fn: raise ValueError( "Callable create_model_fn must be passed" "with ModelType.SAVED_WEIGHTS") return _load_model_from_weights(self._create_model_fn, self._model_uri) return _load_model( self._model_uri, self._custom_weights, self._load_model_args) def update_model_path(self, model_path: Optional[str] = None): self._model_uri = model_path if model_path else self._model_uri def run_inference( self, batch: Sequence[tf.Tensor], model: tf.Module, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """ Runs inferences on a batch of tf.Tensor and returns an Iterable of Tensor Predictions. This method stacks the list of Tensors in a vectorized format to optimize the inference call. Args: batch: A sequence of Tensors. These Tensors should be batchable, as this method will call `tf.stack()` and pass in batched Tensors with dimensions (batch_size, n_features, etc.) into the model's predict() function. model: A Tensorflow model. inference_args: Non-batchable arguments required as inputs to the model's forward() function. Unlike Tensors in `batch`, these parameters will not be dynamically batched Returns: An Iterable of type PredictionResult. """ inference_args = {} if not inference_args else inference_args return self._inference_fn(model, batch, inference_args, self._model_uri) def get_num_bytes(self, batch: Sequence[tf.Tensor]) -> int: """ Returns: The number of bytes of data for a batch of Tensors. """ return sum(sys.getsizeof(element) for element in batch) def get_metrics_namespace(self) -> str: """ Returns: A namespace for metrics collected by the RunInference transform. """ return 'BeamML_TF_Tensor' ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/tensorflow_inference_it_test.py0000644000175100017510000002004715206113032026553 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-End test for Tensorflow Inference""" import logging import unittest import uuid from pathlib import Path import pytest from apache_beam.io.filesystems import FileSystems from apache_beam.testing.test_pipeline import TestPipeline # pylint: disable=ungrouped-imports try: import tensorflow as tf import tensorflow_hub as hub from apache_beam.examples.inference import tensorflow_imagenet_segmentation from apache_beam.examples.inference import tensorflow_mnist_classification from apache_beam.examples.inference import tensorflow_mnist_with_weights except ImportError as e: tf = None def process_outputs(filepath): with FileSystems().open(filepath) as f: lines = f.readlines() lines = [l.decode('utf-8').strip('\n') for l in lines] return lines def rmdir(directory): directory = Path(directory) for item in directory.iterdir(): if item.is_dir(): rmdir(item) else: item.unlink() directory.rmdir() def clear_tf_hub_temp_dir(model_path): # When loading from tensorflow hub using tfhub.resolve, the model is saved in # a temporary directory. That file can be persisted between test runs, in # which case tfhub.resolve will no-op. If the model is deleted and the file # isn't, tfhub.resolve will still no-op and tf.keras.models.load_model will # throw. To avoid this (and test more robustly) we delete the temporary # directory entirely between runs. local_path = hub.resolve(model_path) rmdir(local_path) @unittest.skipIf( tf is None, 'Missing dependencies. ' 'Test depends on tensorflow') @pytest.mark.uses_tf @pytest.mark.it_postcommit class TensorflowInference(unittest.TestCase): def test_tf_mnist_classification(self): test_pipeline = TestPipeline(is_integration_test=True) input_file = 'gs://apache-beam-ml/testing/inputs/it_mnist_data.csv' output_file_dir = 'gs://apache-beam-ml/testing/outputs' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_path = 'gs://apache-beam-ml/models/tensorflow/mnist/' extra_opts = { 'input': input_file, 'output': output_file, 'model_path': model_path, } tensorflow_mnist_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_output_filepath = 'gs://apache-beam-ml/testing/expected_outputs/test_sklearn_mnist_classification_actuals.txt' # pylint: disable=line-too-long expected_outputs = process_outputs(expected_output_filepath) predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) predictions_dict = {} for i in range(len(predicted_outputs)): true_label, prediction = predicted_outputs[i].split(',') predictions_dict[true_label] = prediction for i in range(len(expected_outputs)): true_label, expected_prediction = expected_outputs[i].split(',') self.assertEqual(predictions_dict[true_label], expected_prediction) def test_tf_mnist_classification_large_model(self): test_pipeline = TestPipeline(is_integration_test=True) input_file = 'gs://apache-beam-ml/testing/inputs/it_mnist_data.csv' output_file_dir = 'gs://apache-beam-ml/testing/outputs' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_path = 'gs://apache-beam-ml/models/tensorflow/mnist/' extra_opts = { 'input': input_file, 'output': output_file, 'model_path': model_path, 'large_model': True, } tensorflow_mnist_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_output_filepath = 'gs://apache-beam-ml/testing/expected_outputs/test_sklearn_mnist_classification_actuals.txt' # pylint: disable=line-too-long expected_outputs = process_outputs(expected_output_filepath) predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) predictions_dict = {} for i in range(len(predicted_outputs)): true_label, prediction = predicted_outputs[i].split(',') predictions_dict[true_label] = prediction for i in range(len(expected_outputs)): true_label, expected_prediction = expected_outputs[i].split(',') self.assertEqual(predictions_dict[true_label], expected_prediction) def test_tf_imagenet_image_segmentation(self): test_pipeline = TestPipeline(is_integration_test=True) input_file = ( 'gs://apache-beam-ml/testing/inputs/it_imagenet_input_labels.txt') image_dir = ( 'https://storage.googleapis.com/download.tensorflow.org/example_images/' ) output_file_dir = 'gs://apache-beam-ml/testing/outputs' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_path = ( 'https://tfhub.dev/google/tf2-preview/mobilenet_v2/classification/4') clear_tf_hub_temp_dir(model_path) extra_opts = { 'input': input_file, 'output': output_file, 'model_path': model_path, 'image_dir': image_dir } tensorflow_imagenet_segmentation.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_output_filepath = 'gs://apache-beam-ml/testing/expected_outputs/test_tf_imagenet_image_segmentation.txt' # pylint: disable=line-too-long expected_outputs = process_outputs(expected_output_filepath) predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) for true_label, predicted_label in zip(expected_outputs, predicted_outputs): self.assertEqual(true_label, predicted_label) def test_tf_mnist_with_weights_classification(self): test_pipeline = TestPipeline(is_integration_test=True) input_file = 'gs://apache-beam-ml/testing/inputs/it_mnist_data.csv' output_file_dir = 'gs://apache-beam-ml/testing/outputs' output_file = '/'.join([output_file_dir, str(uuid.uuid4()), 'result.txt']) model_path = 'gs://apache-beam-ml/models/tensorflow/mnist' extra_opts = { 'input': input_file, 'output': output_file, 'model_path': model_path, } tensorflow_mnist_with_weights.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_output_filepath = 'gs://apache-beam-ml/testing/expected_outputs/test_sklearn_mnist_classification_actuals.txt' # pylint: disable=line-too-long expected_outputs = process_outputs(expected_output_filepath) predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) predictions_dict = {} for i in range(len(predicted_outputs)): true_label, prediction = predicted_outputs[i].split(',') predictions_dict[true_label] = prediction for i in range(len(expected_outputs)): true_label, expected_prediction = expected_outputs[i].split(',') self.assertEqual(predictions_dict[true_label], expected_prediction) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/tensorflow_inference_test.py0000644000175100017510000004213715206113032026063 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import os import shutil import tempfile import unittest import uuid from collections.abc import Iterable from collections.abc import Sequence from typing import Any from typing import Optional from typing import Union import numpy import pytest import apache_beam as beam from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # pylint: disable=ungrouped-imports try: import tensorflow as tf from apache_beam.ml.inference import tensorflow_inference from apache_beam.ml.inference.sklearn_inference_test import _compare_prediction_result from apache_beam.ml.inference.tensorflow_inference import TFModelHandlerNumpy from apache_beam.ml.inference.tensorflow_inference import TFModelHandlerTensor except ImportError: raise unittest.SkipTest( 'Tensorflow dependencies are not installed. ' + 'Make sure you have both tensorflow and tensorflow_hub installed.') class FakeTFNumpyModel: def predict(self, input: numpy.ndarray): return numpy.multiply(input, 10) class FakeTFTensorModel: def predict(self, input: tf.Tensor, add=False): if add: return tf.math.add(tf.math.multiply(input, 10), 10) return tf.math.multiply(input, 10) def _create_mult2_model(): inputs = tf.keras.Input(shape=(3, )) outputs = tf.keras.layers.Lambda(lambda x: x * 2, dtype='float32')(inputs) return tf.keras.Model(inputs=inputs, outputs=outputs) def _compare_tensor_prediction_result(x, y): return tf.reduce_all(tf.math.equal(x.inference, y.inference)) def fake_inference_fn( model: tf.Module, batch: Union[Sequence[numpy.ndarray], Sequence[tf.Tensor]], inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: predictions = model.predict(batch, **inference_args) return utils._convert_to_result(batch, predictions, model_id) @pytest.mark.uses_tf class TFRunInferenceTest(unittest.TestCase): def setUp(self): self.tmpdir = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.tmpdir) def test_predict_numpy(self): fake_model = FakeTFNumpyModel() inference_runner = TFModelHandlerNumpy( model_uri='unused', inference_fn=fake_inference_fn) batched_examples = [numpy.array([1]), numpy.array([10]), numpy.array([100])] expected_predictions = [ PredictionResult(numpy.array([1]), 10), PredictionResult(numpy.array([10]), 100), PredictionResult(numpy.array([100]), 1000) ] inferences = inference_runner.run_inference(batched_examples, fake_model) for actual, expected in zip(inferences, expected_predictions): self.assertTrue(_compare_prediction_result(actual, expected)) def test_predict_tensor(self): fake_model = FakeTFTensorModel() inference_runner = TFModelHandlerTensor( model_uri='unused', inference_fn=fake_inference_fn) batched_examples = [ tf.convert_to_tensor(numpy.array([1])), tf.convert_to_tensor(numpy.array([10])), tf.convert_to_tensor(numpy.array([100])), ] expected_predictions = [ PredictionResult(ex, pred) for ex, pred in zip( batched_examples, [tf.math.multiply(n, 10) for n in batched_examples]) ] inferences = inference_runner.run_inference(batched_examples, fake_model) for actual, expected in zip(inferences, expected_predictions): self.assertTrue(_compare_tensor_prediction_result(actual, expected)) def test_predict_tensor_with_batch_size(self): model = _create_mult2_model() model_path = os.path.join(self.tmpdir, f'mult2_{uuid.uuid4()}.keras') tf.keras.models.save_model(model, model_path) # FnApiRunner guarantees large batches, which this pipeline assumes with TestPipeline('FnApiRunner') as pipeline: def fake_batching_inference_fn( model: tf.Module, batch: Union[Sequence[numpy.ndarray], Sequence[tf.Tensor]], inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: if len(batch) != 2: raise Exception( f'Expected batch of size 2, received batch of size {len(batch)}') batch = tf.stack(batch, axis=0) predictions = model(batch) return utils._convert_to_result(batch, predictions, model_id) model_handler = TFModelHandlerTensor( model_uri=model_path, inference_fn=fake_batching_inference_fn, load_model_args={'safe_mode': False}, min_batch_size=2, max_batch_size=2) examples = [ tf.convert_to_tensor(numpy.array([1.1, 2.2, 3.3], dtype='float32')), tf.convert_to_tensor( numpy.array([10.1, 20.2, 30.3], dtype='float32')), tf.convert_to_tensor( numpy.array([100.1, 200.2, 300.3], dtype='float32')), tf.convert_to_tensor( numpy.array([200.1, 300.2, 400.3], dtype='float32')), ] expected_predictions = [ PredictionResult(ex, pred) for ex, pred in zip( examples, [tf.math.multiply(n, 2) for n in examples]) ] pcoll = pipeline | 'start' >> beam.Create(examples) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to( expected_predictions, equals_fn=_compare_tensor_prediction_result)) def test_predict_tensor_with_large_model(self): model = _create_mult2_model() model_path = os.path.join(self.tmpdir, f'mult2_{uuid.uuid4()}.keras') tf.keras.models.save_model(model, model_path) with TestPipeline() as pipeline: def fake_batching_inference_fn( model: tf.Module, batch: Union[Sequence[numpy.ndarray], Sequence[tf.Tensor]], inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: multi_process_shared_loaded = "multi_process_shared" in str(type(model)) if not multi_process_shared_loaded: raise Exception( f'Loaded model of type {type(model)}, was ' + 'expecting multi_process_shared_model') batch = tf.stack(batch, axis=0) predictions = model(batch) return utils._convert_to_result(batch, predictions, model_id) model_handler = TFModelHandlerTensor( model_uri=model_path, inference_fn=fake_batching_inference_fn, load_model_args={'safe_mode': False}, large_model=True) examples = [ tf.convert_to_tensor(numpy.array([1.1, 2.2, 3.3], dtype='float32')), tf.convert_to_tensor( numpy.array([10.1, 20.2, 30.3], dtype='float32')), tf.convert_to_tensor( numpy.array([100.1, 200.2, 300.3], dtype='float32')), tf.convert_to_tensor( numpy.array([200.1, 300.2, 400.3], dtype='float32')), ] expected_predictions = [ PredictionResult(ex, pred) for ex, pred in zip( examples, [tf.math.multiply(n, 2) for n in examples]) ] pcoll = pipeline | 'start' >> beam.Create(examples) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to( expected_predictions, equals_fn=_compare_tensor_prediction_result)) def test_predict_numpy_with_batch_size(self): model = _create_mult2_model() model_path = os.path.join(self.tmpdir, f'mult2_{uuid.uuid4()}.keras') tf.keras.models.save_model(model, model_path) # TODO(https://github.com/apache/beam/issues/34549): This test relies on a # runner producing a single bundle or bundles of even size, neither of # which prism seems to do here with TestPipeline('FnApiRunner') as pipeline: def fake_batching_inference_fn( model: tf.Module, batch: Sequence[numpy.ndarray], inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: if len(batch) != 2: raise Exception( f'Expected batch of size 2, received batch of size {len(batch)}') vectorized_batch = numpy.stack(batch, axis=0) predictions = model.predict(vectorized_batch, **inference_args) return utils._convert_to_result(batch, predictions, model_id) model_handler = TFModelHandlerNumpy( model_uri=model_path, inference_fn=fake_batching_inference_fn, load_model_args={'safe_mode': False}, min_batch_size=2, max_batch_size=2) examples = [ numpy.array([1.1, 2.2, 3.3], dtype='float32'), numpy.array([10.1, 20.2, 30.3], dtype='float32'), numpy.array([100.1, 200.2, 300.3], dtype='float32'), numpy.array([200.1, 300.2, 400.3], dtype='float32'), ] expected_predictions = [ PredictionResult(ex, pred) for ex, pred in zip( examples, [numpy.multiply(n, 2) for n in examples]) ] pcoll = pipeline | 'start' >> beam.Create(examples) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to( expected_predictions, equals_fn=_compare_tensor_prediction_result)) def test_predict_numpy_with_large_model(self): model = _create_mult2_model() model_path = os.path.join(self.tmpdir, f'mult2_{uuid.uuid4()}.keras') tf.keras.models.save_model(model, model_path) with TestPipeline() as pipeline: def fake_inference_fn( model: tf.Module, batch: Sequence[numpy.ndarray], inference_args: dict[str, Any], model_id: Optional[str] = None) -> Iterable[PredictionResult]: multi_process_shared_loaded = "multi_process_shared" in str(type(model)) if not multi_process_shared_loaded: raise Exception( f'Loaded model of type {type(model)}, was ' + 'expecting multi_process_shared_model') vectorized_batch = numpy.stack(batch, axis=0) predictions = model.predict(vectorized_batch, **inference_args) return utils._convert_to_result(batch, predictions, model_id) model_handler = TFModelHandlerNumpy( model_uri=model_path, load_model_args={'safe_mode': False}, inference_fn=fake_inference_fn, large_model=True) examples = [ numpy.array([1.1, 2.2, 3.3], dtype='float32'), numpy.array([10.1, 20.2, 30.3], dtype='float32'), numpy.array([100.1, 200.2, 300.3], dtype='float32'), numpy.array([200.1, 300.2, 400.3], dtype='float32'), ] expected_predictions = [ PredictionResult(ex, pred) for ex, pred in zip( examples, [numpy.multiply(n, 2) for n in examples]) ] pcoll = pipeline | 'start' >> beam.Create(examples) predictions = pcoll | RunInference(model_handler) assert_that( predictions, equal_to( expected_predictions, equals_fn=_compare_tensor_prediction_result)) def test_predict_tensor_with_args(self): fake_model = FakeTFTensorModel() inference_runner = TFModelHandlerTensor( model_uri='unused', inference_fn=fake_inference_fn) batched_examples = [ tf.convert_to_tensor(numpy.array([1])), tf.convert_to_tensor(numpy.array([10])), tf.convert_to_tensor(numpy.array([100])), ] expected_predictions = [ PredictionResult(ex, pred) for ex, pred in zip( batched_examples, [ tf.math.add(tf.math.multiply(n, 10), 10) for n in batched_examples ]) ] inferences = inference_runner.run_inference( batched_examples, fake_model, inference_args={"add": True}) for actual, expected in zip(inferences, expected_predictions): self.assertTrue(_compare_tensor_prediction_result(actual, expected)) def test_predict_keyed_numpy(self): fake_model = FakeTFNumpyModel() inference_runner = KeyedModelHandler( TFModelHandlerNumpy(model_uri='unused', inference_fn=fake_inference_fn)) batched_examples = [ ('k1', numpy.array([1], dtype=numpy.int64)), ('k2', numpy.array([10], dtype=numpy.int64)), ('k3', numpy.array([100], dtype=numpy.int64)), ] expected_predictions = [ (ex[0], PredictionResult(ex[1], pred)) for ex, pred in zip( batched_examples, [numpy.multiply(n[1], 10) for n in batched_examples]) ] inferences = inference_runner.run_inference(batched_examples, fake_model) for actual, expected in zip(inferences, expected_predictions): self.assertTrue(_compare_prediction_result(actual[1], expected[1])) def test_predict_keyed_tensor(self): fake_model = FakeTFTensorModel() inference_runner = KeyedModelHandler( TFModelHandlerTensor( model_uri='unused', inference_fn=fake_inference_fn)) batched_examples = [ ('k1', tf.convert_to_tensor(numpy.array([1]))), ('k2', tf.convert_to_tensor(numpy.array([10]))), ('k3', tf.convert_to_tensor(numpy.array([100]))), ] expected_predictions = [ (ex[0], PredictionResult(ex[1], pred)) for ex, pred in zip( batched_examples, [tf.math.multiply(n[1], 10) for n in batched_examples]) ] inferences = inference_runner.run_inference(batched_examples, fake_model) for actual, expected in zip(inferences, expected_predictions): self.assertTrue(_compare_tensor_prediction_result(actual[1], expected[1])) def test_load_model_exception(self): with self.assertRaises(ValueError): tensorflow_inference._load_model( "https://tfhub.dev/google/imagenet/mobilenet_v1_075_192/quantops/classification/3", # pylint: disable=line-too-long None, {}) @pytest.mark.uses_tf class TFRunInferenceTestWithMocks(unittest.TestCase): def setUp(self): self._load_model = tensorflow_inference._load_model tensorflow_inference._load_model = unittest.mock.MagicMock() def tearDown(self): tensorflow_inference._load_model = self._load_model def test_load_model_args(self): load_model_args = {compile: False, 'custom_objects': {'optimizer': 1}} model_handler = TFModelHandlerNumpy( "dummy_model", load_model_args=load_model_args) model_handler.load_model() tensorflow_inference._load_model.assert_called_with( "dummy_model", "", load_model_args) def test_load_model_with_args_and_custom_weights(self): load_model_args = {compile: False, 'custom_objects': {'optimizer': 1}} model_handler = TFModelHandlerNumpy( "dummy_model", custom_weights="dummy_weights", load_model_args=load_model_args) model_handler.load_model() tensorflow_inference._load_model.assert_called_with( "dummy_model", "dummy_weights", load_model_args) def test_env_vars_set_correctly_tensor(self): handler_with_vars = TFModelHandlerTensor( env_vars={'FOO': 'bar'}, model_uri='unused', inference_fn=fake_inference_fn) os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) batched_examples = [ tf.convert_to_tensor(numpy.array([1])), tf.convert_to_tensor(numpy.array([10])), tf.convert_to_tensor(numpy.array([100])), ] with TestPipeline() as pipeline: _ = ( pipeline | 'start' >> beam.Create(batched_examples) | RunInference(handler_with_vars)) pipeline.run() self.assertTrue('FOO' in os.environ) self.assertTrue((os.environ['FOO']) == 'bar') def test_env_vars_set_correctly_numpy(self): handler_with_vars = TFModelHandlerNumpy( env_vars={'FOO': 'bar'}, model_uri="unused", inference_fn=fake_inference_fn) os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) batched_examples = [numpy.array([1]), numpy.array([10]), numpy.array([100])] tensorflow_inference._load_model = unittest.mock.MagicMock() with TestPipeline() as pipeline: _ = ( pipeline | 'start' >> beam.Create(batched_examples) | RunInference(handler_with_vars)) pipeline.run() self.assertTrue('FOO' in os.environ) self.assertTrue((os.environ['FOO']) == 'bar') if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/tensorrt_inference.py0000644000175100017510000002752315206113032024504 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file from __future__ import annotations import logging import threading from collections.abc import Callable from collections.abc import Iterable from collections.abc import Sequence from typing import Any from typing import Optional import numpy as np from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import PredictionResult LOGGER = logging.getLogger("TensorRTEngineHandlerNumPy") # This try/catch block allows users to submit jobs from a machine without # GPU and other dependencies (tensorrt, cuda, etc.) at job submission time. try: import tensorrt as trt TRT_LOGGER = trt.Logger(trt.Logger.INFO) trt.init_libnvinfer_plugins(TRT_LOGGER, namespace="") LOGGER.info('tensorrt module successfully imported.') except ModuleNotFoundError: TRT_LOGGER = None msg = 'tensorrt module was not found. This is ok as long as the specified ' \ 'runner has tensorrt dependencies installed.' LOGGER.warning(msg) def _load_engine(engine_path): import tensorrt as trt file = FileSystems.open(engine_path, 'rb') runtime = trt.Runtime(TRT_LOGGER) engine = runtime.deserialize_cuda_engine(file.read()) assert engine return engine def _load_onnx(onnx_path): import tensorrt as trt builder = trt.Builder(TRT_LOGGER) network = builder.create_network( flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser = trt.OnnxParser(network, TRT_LOGGER) with FileSystems.open(onnx_path) as f: if not parser.parse(f.read()): LOGGER.error("Failed to load ONNX file: %s", onnx_path) for error in range(parser.num_errors): LOGGER.error(parser.get_error(error)) raise ValueError(f"Failed to load ONNX file: {onnx_path}") return network, builder def _build_engine(network, builder): import tensorrt as trt config = builder.create_builder_config() runtime = trt.Runtime(TRT_LOGGER) plan = builder.build_serialized_network(network, config) engine = runtime.deserialize_cuda_engine(plan) builder.reset() return engine def _assign_or_fail(args): """CUDA error checking.""" from cuda import cuda err, ret = args[0], args[1:] if isinstance(err, cuda.CUresult): if err != cuda.CUresult.CUDA_SUCCESS: raise RuntimeError("Cuda Error: {}".format(err)) else: raise RuntimeError("Unknown error type: {}".format(err)) # Special case so that no unpacking is needed at call-site. if len(ret) == 1: return ret[0] return ret class TensorRTEngine: def __init__(self, engine: trt.ICudaEngine): """Implementation of the TensorRTEngine class which handles allocations associated with TensorRT engine. Example Usage:: TensorRTEngine(engine) Args: engine: trt.ICudaEngine object that contains TensorRT engine """ import tensorrt as trt from cuda import cuda self.engine = engine self.context = engine.create_execution_context() self.context_lock = threading.RLock() self.inputs = [] self.outputs = [] self.gpu_allocations = [] self.cpu_allocations = [] # TODO(https://github.com/NVIDIA/TensorRT/issues/2557): # Clean up when fixed upstream. try: _ = np.bool except AttributeError: # numpy >= 1.24.0 np.bool = np.bool_ # type: ignore # Setup I/O bindings. for i in range(self.engine.num_bindings): name = self.engine.get_binding_name(i) dtype = self.engine.get_binding_dtype(i) shape = self.engine.get_binding_shape(i) size = trt.volume(shape) * dtype.itemsize allocation = _assign_or_fail(cuda.cuMemAlloc(size)) binding = { 'index': i, 'name': name, 'dtype': np.dtype(trt.nptype(dtype)), 'shape': list(shape), 'allocation': allocation, 'size': size } self.gpu_allocations.append(allocation) if self.engine.binding_is_input(i): self.inputs.append(binding) else: self.outputs.append(binding) assert self.context assert len(self.inputs) > 0 assert len(self.outputs) > 0 assert len(self.gpu_allocations) > 0 for output in self.outputs: self.cpu_allocations.append(np.zeros(output['shape'], output['dtype'])) # Create CUDA Stream. self.stream = _assign_or_fail(cuda.cuStreamCreate(0)) def get_engine_attrs(self): """Returns TensorRT engine attributes.""" return ( self.engine, self.context, self.context_lock, self.inputs, self.outputs, self.gpu_allocations, self.cpu_allocations, self.stream) TensorRTInferenceFn = Callable[ [Sequence[np.ndarray], TensorRTEngine, Optional[dict[str, Any]]], Iterable[PredictionResult]] def _default_tensorRT_inference_fn( batch: Sequence[np.ndarray], engine: TensorRTEngine, inference_args: Optional[dict[str, Any]] = None) -> Iterable[PredictionResult]: from cuda import cuda ( engine, context, context_lock, inputs, outputs, gpu_allocations, cpu_allocations, stream) = engine.get_engine_attrs() # Process I/O and execute the network with context_lock: _assign_or_fail( cuda.cuMemcpyHtoDAsync( inputs[0]['allocation'], np.ascontiguousarray(batch), inputs[0]['size'], stream)) context.execute_async_v2(gpu_allocations, stream) for output in range(len(cpu_allocations)): _assign_or_fail( cuda.cuMemcpyDtoHAsync( cpu_allocations[output], outputs[output]['allocation'], outputs[output]['size'], stream)) _assign_or_fail(cuda.cuStreamSynchronize(stream)) predictions = [] for idx in range(len(batch)): predictions.append([prediction[idx] for prediction in cpu_allocations]) return utils._convert_to_result(batch, predictions) class TensorRTEngineHandlerNumPy(ModelHandler[np.ndarray, PredictionResult, TensorRTEngine]): def __init__( self, min_batch_size: int, max_batch_size: int, *, inference_fn: TensorRTInferenceFn = _default_tensorRT_inference_fn, large_model: bool = False, model_copies: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, **kwargs): """Implementation of the ModelHandler interface for TensorRT. Example Usage:: pcoll | RunInference( TensorRTEngineHandlerNumPy( min_batch_size=1, max_batch_size=1, engine_path="my_uri")) **NOTE:** This API and its implementation are under development and do not provide backward compatibility guarantees. Args: min_batch_size: minimum accepted batch size. max_batch_size: maximum accepted batch size. inference_fn: the inference function to use on RunInference calls. default: _default_tensorRT_inference_fn large_model: set to true if your model is large enough to run into memory pressure if you load multiple copies. Given a model that consumes N memory and a machine with W cores and M memory, you should set this to True if N*W > M. model_copies: The exact number of models that you would like loaded onto your machine. This can be useful if you exactly know your CPU or GPU capacity and want to maximize resource utilization. max_batch_duration_secs: the maximum amount of time to buffer a batch before emitting; used in streaming contexts. max_batch_weight: the maximum total weight of a batch. element_size_fn: a function that returns the size (weight) of an element. batch_length_fn: a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: a sorted list of positive boundary values for length-aware batching buckets. kwargs: Additional arguments like 'engine_path' and 'onnx_path' are currently supported. 'env_vars' can be used to set environment variables before loading the model. See https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/ for details """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries, large_model=large_model, model_copies=model_copies, **kwargs) self.inference_fn = inference_fn if 'engine_path' in kwargs: self.engine_path = kwargs.get('engine_path') elif 'onnx_path' in kwargs: self.onnx_path = kwargs.get('onnx_path') def load_model(self) -> TensorRTEngine: """Loads and initializes a TensorRT engine for processing.""" engine = _load_engine(self.engine_path) return TensorRTEngine(engine) def load_onnx(self) -> tuple[trt.INetworkDefinition, trt.Builder]: """Loads and parses an onnx model for processing.""" return _load_onnx(self.onnx_path) def build_engine( self, network: trt.INetworkDefinition, builder: trt.Builder) -> TensorRTEngine: """Build an engine according to parsed/created network.""" engine = _build_engine(network, builder) return TensorRTEngine(engine) def run_inference( self, batch: Sequence[np.ndarray], engine: TensorRTEngine, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """ Runs inferences on a batch of Tensors and returns an Iterable of TensorRT Predictions. Args: batch: A np.ndarray or a np.ndarray that represents a concatenation of multiple arrays as a batch. engine: A TensorRT engine. inference_args: Any additional arguments for an inference that are not applicable to TensorRT. Returns: An Iterable of type PredictionResult. """ return self.inference_fn(batch, engine, inference_args) def get_num_bytes(self, batch: Sequence[np.ndarray]) -> int: """ Returns: The number of bytes of data for a batch of Tensors. """ return sum((np_array.itemsize for np_array in batch)) def get_metrics_namespace(self) -> str: """ Returns a namespace for metrics collected by the RunInference transform. """ return 'BeamML_TensorRT' def validate_inference_args(self, inference_args: Optional[dict[str, Any]]): """ Currently, this model handler does not support inference args. Given that, we will throw if any are passed in. """ if inference_args: raise ValueError( 'inference_args were provided, but should be None because this ' 'framework does not expect extra arguments on inferences.') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/tensorrt_inference_test.py0000644000175100017510000004432715206113032025544 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import os import unittest import numpy as np import pytest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # Protect against environments where TensorRT python library is not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: import tensorrt as trt from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.tensorrt_inference import TensorRTEngineHandlerNumPy except ImportError: raise unittest.SkipTest('TensorRT dependencies are not installed') try: from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem except ImportError: GCSFileSystem = None # type: ignore LOGGER = trt.Logger(trt.Logger.INFO) SINGLE_FEATURE_EXAMPLES = [ np.array(1, dtype=np.float32), np.array(5, dtype=np.float32), np.array(-3, dtype=np.float32), np.array(10.0, dtype=np.float32) ] SINGLE_FEATURE_PREDICTIONS = [ PredictionResult(ex, pred) for ex, pred in zip( SINGLE_FEATURE_EXAMPLES, [[np.array([example * 2.0 + 0.5], dtype=np.float32)] for example in SINGLE_FEATURE_EXAMPLES]) ] SINGLE_FEATURE_CUSTOM_PREDICTIONS = [ PredictionResult(ex, pred) for ex, pred in zip( SINGLE_FEATURE_EXAMPLES, [[np.array([(example * 2.0 + 0.5) * 2], dtype=np.float32)] for example in SINGLE_FEATURE_EXAMPLES]) ] TWO_FEATURES_EXAMPLES = [ np.array([1, 5], dtype=np.float32), np.array([3, 10], dtype=np.float32), np.array([-14, 0], dtype=np.float32), np.array([0.5, 0.5], dtype=np.float32) ] TWO_FEATURES_PREDICTIONS = [ PredictionResult(ex, pred) for ex, pred in zip( TWO_FEATURES_EXAMPLES, [[ np.array([example[0] * 2.0 + example[1] * 3 + 0.5], dtype=np.float32) ] for example in TWO_FEATURES_EXAMPLES]) ] def _compare_prediction_result(a, b): return ((a.example == b.example).all() and all( np.array_equal(actual, expected) for actual, expected in zip(a.inference, b.inference))) def _assign_or_fail(args): """CUDA error checking.""" from cuda import cuda err, ret = args[0], args[1:] if isinstance(err, cuda.CUresult): if err != cuda.CUresult.CUDA_SUCCESS: raise RuntimeError("Cuda Error: {}".format(err)) else: raise RuntimeError("Unknown error type: {}".format(err)) # Special case so that no unpacking is needed at call-site. if len(ret) == 1: return ret[0] return ret def _custom_tensorRT_inference_fn(batch, engine, inference_args): from cuda import cuda ( engine, context, context_lock, inputs, outputs, gpu_allocations, cpu_allocations, stream) = engine.get_engine_attrs() # Process I/O and execute the network with context_lock: _assign_or_fail( cuda.cuMemcpyHtoDAsync( inputs[0]['allocation'], np.ascontiguousarray(batch), inputs[0]['size'], stream)) context.execute_async_v2(gpu_allocations, stream) for output in range(len(cpu_allocations)): _assign_or_fail( cuda.cuMemcpyDtoHAsync( cpu_allocations[output], outputs[output]['allocation'], outputs[output]['size'], stream)) _assign_or_fail(cuda.cuStreamSynchronize(stream)) return [ PredictionResult( x, [prediction[idx] * 2 for prediction in cpu_allocations]) for idx, x in enumerate(batch) ] @pytest.mark.uses_tensorrt class TensorRTRunInferenceTest(unittest.TestCase): @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed') def test_inference_single_tensor_feature_onnx(self): """ This tests ONNX parser and TensorRT engine creation from parsed ONNX network. Single feature tensors batched into size of 4 are used as input. """ inference_runner = TensorRTEngineHandlerNumPy( min_batch_size=4, max_batch_size=4, onnx_path="gs://apache-beam-ml/models/single_tensor_features_model.onnx" ) network, builder = inference_runner.load_onnx() engine = inference_runner.build_engine(network, builder) predictions = inference_runner.run_inference( SINGLE_FEATURE_EXAMPLES, engine) for actual, expected in zip(predictions, SINGLE_FEATURE_PREDICTIONS): self.assertEqual(actual, expected) @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed') def test_inference_multiple_tensor_features_onnx(self): """ This tests ONNX parser and TensorRT engine creation from parsed ONNX network. Two feature tensors batched into size of 4 are used as input. """ inference_runner = TensorRTEngineHandlerNumPy( min_batch_size=4, max_batch_size=4, onnx_path= 'gs://apache-beam-ml/models/multiple_tensor_features_model.onnx') network, builder = inference_runner.load_onnx() engine = inference_runner.build_engine(network, builder) predictions = inference_runner.run_inference(TWO_FEATURES_EXAMPLES, engine) for actual, expected in zip(predictions, TWO_FEATURES_PREDICTIONS): self.assertTrue(_compare_prediction_result(actual, expected)) def test_inference_single_tensor_feature(self): """ This tests creating TensorRT network from scratch. Test replicates the same ONNX network above but natively in TensorRT. After network creation, network is used to build a TensorRT engine. Single feature tensors batched into size of 4 are used as input. """ inference_runner = TensorRTEngineHandlerNumPy( min_batch_size=4, max_batch_size=4) builder = trt.Builder(LOGGER) network = builder.create_network( flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) input_tensor = network.add_input( name="input", dtype=trt.float32, shape=(4, 1)) weight_const = network.add_constant( (1, 1), trt.Weights((np.ascontiguousarray([2.0], dtype=np.float32)))) mm = network.add_matrix_multiply( input_tensor, trt.MatrixOperation.NONE, weight_const.get_output(0), trt.MatrixOperation.NONE) bias_const = network.add_constant( (1, 1), trt.Weights((np.ascontiguousarray([0.5], dtype=np.float32)))) bias_add = network.add_elementwise( mm.get_output(0), bias_const.get_output(0), trt.ElementWiseOperation.SUM) bias_add.get_output(0).name = "output" network.mark_output(tensor=bias_add.get_output(0)) engine = inference_runner.build_engine(network, builder) predictions = inference_runner.run_inference( SINGLE_FEATURE_EXAMPLES, engine) for actual, expected in zip(predictions, SINGLE_FEATURE_PREDICTIONS): self.assertEqual(actual, expected) def test_inference_custom_single_tensor_feature(self): """ This tests creating TensorRT network from scratch. Test replicates the same ONNX network above but natively in TensorRT. After network creation, network is used to build a TensorRT engine. Single feature tensors batched into size of 4 are used as input. This routes through a custom inference function. """ inference_runner = TensorRTEngineHandlerNumPy( min_batch_size=4, max_batch_size=4, inference_fn=_custom_tensorRT_inference_fn) builder = trt.Builder(LOGGER) network = builder.create_network( flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) input_tensor = network.add_input( name="input", dtype=trt.float32, shape=(4, 1)) weight_const = network.add_constant( (1, 1), trt.Weights((np.ascontiguousarray([2.0], dtype=np.float32)))) mm = network.add_matrix_multiply( input_tensor, trt.MatrixOperation.NONE, weight_const.get_output(0), trt.MatrixOperation.NONE) bias_const = network.add_constant( (1, 1), trt.Weights((np.ascontiguousarray([0.5], dtype=np.float32)))) bias_add = network.add_elementwise( mm.get_output(0), bias_const.get_output(0), trt.ElementWiseOperation.SUM) bias_add.get_output(0).name = "output" network.mark_output(tensor=bias_add.get_output(0)) engine = inference_runner.build_engine(network, builder) predictions = inference_runner.run_inference( SINGLE_FEATURE_EXAMPLES, engine) for actual, expected in zip(predictions, SINGLE_FEATURE_CUSTOM_PREDICTIONS): self.assertEqual(actual, expected) def test_inference_multiple_tensor_features(self): """ This tests creating TensorRT network from scratch. Test replicates the same ONNX network above but natively in TensorRT. After network creation, network is used to build a TensorRT engine. Two feature tensors batched into size of 4 are used as input. """ inference_runner = TensorRTEngineHandlerNumPy( min_batch_size=4, max_batch_size=4) builder = trt.Builder(LOGGER) network = builder.create_network( flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) input_tensor = network.add_input( name="input", dtype=trt.float32, shape=(4, 2)) weight_const = network.add_constant( (1, 2), trt.Weights((np.ascontiguousarray([2.0, 3], dtype=np.float32)))) mm = network.add_matrix_multiply( input_tensor, trt.MatrixOperation.NONE, weight_const.get_output(0), trt.MatrixOperation.TRANSPOSE) bias_const = network.add_constant( (1, 1), trt.Weights((np.ascontiguousarray([0.5], dtype=np.float32)))) bias_add = network.add_elementwise( mm.get_output(0), bias_const.get_output(0), trt.ElementWiseOperation.SUM) bias_add.get_output(0).name = "output" network.mark_output(tensor=bias_add.get_output(0)) engine = inference_runner.build_engine(network, builder) predictions = inference_runner.run_inference(TWO_FEATURES_EXAMPLES, engine) for actual, expected in zip(predictions, TWO_FEATURES_PREDICTIONS): self.assertTrue(_compare_prediction_result(actual, expected)) @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed') def test_inference_single_tensor_feature_built_engine(self): """ This tests already pre-built TensorRT engine from ONNX network. To execute this test succesfully, TensorRT engine that is used here, must have been built in the same environment with the same GPU that will be used when running a test. In other words, using the same environment and same GPU we must pre-build the engine and after we run this test. Otherwise behavior might be unpredictable, read more: https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#compatibility-serialized-engines Single feature tensors batched into size of 4 are used as input. """ inference_runner = TensorRTEngineHandlerNumPy( min_batch_size=4, max_batch_size=4, engine_path= 'gs://apache-beam-ml/models/single_tensor_features_engine.trt') engine = inference_runner.load_model() predictions = inference_runner.run_inference( SINGLE_FEATURE_EXAMPLES, engine) for actual, expected in zip(predictions, SINGLE_FEATURE_PREDICTIONS): self.assertEqual(actual, expected) @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed') def test_inference_multiple_tensor_feature_built_engine(self): """ This tests already pre-built TensorRT engine from ONNX network. To execute this test succesfully, TensorRT engine that is used here, must have been built in the same environment with the same GPU that will be used when running a test. In other words, using the same environment and same GPU we must pre-build the engine and after we run this test. Otherwise behavior might be unpredictable, read more: https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#compatibility-serialized-engines Two feature tensors batched into size of 4 are used as input. """ inference_runner = TensorRTEngineHandlerNumPy( min_batch_size=4, max_batch_size=4, engine_path= 'gs://apache-beam-ml/models/multiple_tensor_features_engine.trt') engine = inference_runner.load_model() predictions = inference_runner.run_inference(TWO_FEATURES_EXAMPLES, engine) for actual, expected in zip(predictions, TWO_FEATURES_PREDICTIONS): self.assertTrue(_compare_prediction_result(actual, expected)) def test_num_bytes(self): inference_runner = TensorRTEngineHandlerNumPy( min_batch_size=1, max_batch_size=1) examples = [ np.array([1, 5], dtype=np.float32), np.array([3, 10], dtype=np.float32), np.array([-14, 0], dtype=np.float32), np.array([0.5, 0.5], dtype=np.float32) ] self.assertEqual((examples[0].itemsize) * 4, inference_runner.get_num_bytes(examples)) def test_namespace(self): inference_runner = TensorRTEngineHandlerNumPy( min_batch_size=4, max_batch_size=4) self.assertEqual( 'RunInferenceTensorRT', inference_runner.get_metrics_namespace()) @pytest.mark.uses_tensorrt class TensorRTRunInferencePipelineTest(unittest.TestCase): @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed') def test_pipeline_single_tensor_feature_built_engine(self): with TestPipeline() as pipeline: engine_handler = TensorRTEngineHandlerNumPy( min_batch_size=4, max_batch_size=4, engine_path= 'gs://apache-beam-ml/models/single_tensor_features_engine.trt') pcoll = pipeline | 'start' >> beam.Create(SINGLE_FEATURE_EXAMPLES) predictions = pcoll | RunInference(engine_handler) assert_that( predictions, equal_to( SINGLE_FEATURE_PREDICTIONS, equals_fn=_compare_prediction_result)) @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed') def test_pipeline_single_tensor_feature_built_engine_large_engine(self): with TestPipeline() as pipeline: def fake_inference_fn(batch, engine, inference_args=None): multi_process_shared_loaded = "multi_process_shared" in str( type(engine)) if not multi_process_shared_loaded: raise Exception( f'Loaded engine of type {type(engine)}, was ' + 'expecting multi_process_shared engine') from cuda import cuda ( engine, context, context_lock, inputs, outputs, gpu_allocations, cpu_allocations, stream) = engine.get_engine_attrs() # Process I/O and execute the network with context_lock: _assign_or_fail( cuda.cuMemcpyHtoDAsync( inputs[0]['allocation'], np.ascontiguousarray(batch), inputs[0]['size'], stream)) context.execute_async_v2(gpu_allocations, stream) for output in range(len(cpu_allocations)): _assign_or_fail( cuda.cuMemcpyDtoHAsync( cpu_allocations[output], outputs[output]['allocation'], outputs[output]['size'], stream)) _assign_or_fail(cuda.cuStreamSynchronize(stream)) predictions = [] for idx in range(len(batch)): predictions.append( [prediction[idx] for prediction in cpu_allocations]) return utils._convert_to_result(batch, predictions) engine_handler = TensorRTEngineHandlerNumPy( min_batch_size=4, max_batch_size=4, engine_path= 'gs://apache-beam-ml/models/single_tensor_features_engine.trt', inference_fn=fake_inference_fn, large_model=True) pcoll = pipeline | 'start' >> beam.Create(SINGLE_FEATURE_EXAMPLES) predictions = pcoll | RunInference(engine_handler) assert_that( predictions, equal_to( SINGLE_FEATURE_PREDICTIONS, equals_fn=_compare_prediction_result)) @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed') def test_pipeline_sets_env_vars_correctly(self): with TestPipeline() as pipeline: engine_handler = TensorRTEngineHandlerNumPy( env_vars={'FOO': 'bar'}, min_batch_size=4, max_batch_size=4, engine_path= 'gs://apache-beam-ml/models/single_tensor_features_engine.trt') os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) _ = ( pipeline | 'start' >> beam.Create(SINGLE_FEATURE_EXAMPLES) | RunInference(engine_handler)) pipeline.run() self.assertTrue('FOO' in os.environ) self.assertTrue((os.environ['FOO']) == 'bar') @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed') def test_pipeline_multiple_tensor_feature_built_engine(self): with TestPipeline() as pipeline: engine_handler = TensorRTEngineHandlerNumPy( min_batch_size=4, max_batch_size=4, engine_path= 'gs://apache-beam-ml/models/multiple_tensor_features_engine.trt') pcoll = pipeline | 'start' >> beam.Create(TWO_FEATURES_EXAMPLES) predictions = pcoll | RunInference(engine_handler) assert_that( predictions, equal_to( TWO_FEATURES_PREDICTIONS, equals_fn=_compare_prediction_result)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/utils.py0000644000175100017510000001416415206113032021743 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file """ Util/helper functions used in apache_beam.ml.inference. """ import os from collections.abc import Iterable from functools import partial from typing import Any from typing import Optional from typing import Union import apache_beam as beam from apache_beam.io.fileio import EmptyMatchTreatment from apache_beam.io.fileio import MatchContinuously from apache_beam.ml.inference.base import ModelMetadata from apache_beam.ml.inference.base import PredictionResult from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.transforms.userstate import CombiningValueStateSpec from apache_beam.utils.timestamp import MAX_TIMESTAMP from apache_beam.utils.timestamp import Timestamp _START_TIME_STAMP = Timestamp.now() def _convert_to_result( batch: Iterable, predictions: Union[Iterable, dict[Any, Iterable]], model_id: Optional[str] = None, ) -> Iterable[PredictionResult]: if isinstance(predictions, dict): # Go from one dictionary of type: {key_type1: Iterable, # key_type2: Iterable, ...} where each Iterable is of # length batch_size, to a list of dictionaries: # [{key_type1: value_type1, key_type2: value_type2}] predictions_per_tensor = [ dict(zip(predictions.keys(), v)) for v in zip(*predictions.values()) ] return [ PredictionResult(x, y, model_id) for x, y in zip(batch, predictions_per_tensor) ] return [PredictionResult(x, y, model_id) for x, y in zip(batch, predictions)] class _ConvertIterToSingleton(beam.DoFn): """ Internal only; No backwards compatibility. The MatchContinuously transform examines all files present in a given directory and returns those that have timestamps older than the pipeline's start time. This can produce an Iterable rather than a Singleton. This class only returns the file path when it is first encountered, and it is cached as part of the side input caching mechanism. If the path is seen again, it will not return anything. By doing this, we can ensure that the output of this transform can be wrapped with beam.pvalue.AsSingleton(). """ COUNT_STATE = CombiningValueStateSpec('count', combine_fn=sum) def process(self, element, count_state=beam.DoFn.StateParam(COUNT_STATE)): counter = count_state.read() if counter == 0: count_state.add(1) yield element[1] class _GetLatestFileByTimeStamp(beam.DoFn): """ Internal only; No backwards compatibility. This DoFn checks the timestamps of files against the time that the pipeline began running. It returns the files that were modified after the pipeline started. If no such files are found, it returns a default file as fallback. """ TIME_STATE = CombiningValueStateSpec( 'max', combine_fn=partial(max, default=_START_TIME_STAMP)) def process(self, element, time_state=beam.DoFn.StateParam(TIME_STATE)): _, file_metadata = element new_ts = file_metadata.last_updated_in_seconds old_ts = time_state.read() if new_ts > old_ts: time_state.clear() time_state.add(new_ts) model_path = file_metadata.path else: model_path = '' model_name = os.path.splitext(os.path.basename(model_path))[0] return [ (model_path, ModelMetadata(model_id=model_path, model_name=model_name)) ] class WatchFilePattern(beam.PTransform): def __init__( self, file_pattern, interval=360, stop_timestamp=MAX_TIMESTAMP, ): """ Watches a directory for updates to files matching a given file pattern. Args: file_pattern: The file path to read from as a local file path or a GCS ``gs://`` path. The path can contain glob characters (``*``, ``?``, and ``[...]`` sets). interval: Interval at which to check for files matching file_pattern in seconds. stop_timestamp: Timestamp after which no more files will be checked. **Note**: 1. Any previously used filenames cannot be reused. If a file is added or updated to a previously used filename, this transform will ignore that update. To trigger a model update, always upload a file with unique name. 2. Initially, before the pipeline startup time, WatchFilePattern expects at least one file present that matches the file_pattern. 3. This transform is supported in streaming mode since MatchContinuously produces an unbounded source. Running in batch mode can lead to undesired results or result in pipeline being stuck. """ self.file_pattern = file_pattern self.interval = interval self.stop_timestamp = stop_timestamp def expand(self, pcoll) -> beam.PCollection[ModelMetadata]: return ( pcoll | 'MatchContinuously' >> MatchContinuously( file_pattern=self.file_pattern, interval=self.interval, stop_timestamp=self.stop_timestamp, empty_match_treatment=EmptyMatchTreatment.DISALLOW) | "AttachKey" >> beam.Map(lambda x: (x.path, x)) | "GetLatestFileMetaData" >> beam.ParDo(_GetLatestFileByTimeStamp()) | "AcceptNewSideInputOnly" >> beam.ParDo(_ConvertIterToSingleton()) | 'ApplyGlobalWindow' >> beam.transforms.WindowInto( window.GlobalWindows(), trigger=trigger.Repeatedly(trigger.AfterProcessingTime(1)), accumulation_mode=trigger.AccumulationMode.DISCARDING)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/utils_test.py0000644000175100017510000000672015206113032023001 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import apache_beam as beam from apache_beam.io.filesystem import FileMetadata from apache_beam.ml.inference import utils from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class WatchFilePatternTest(unittest.TestCase): def test_latest_file_by_timestamp_default_value(self): # match continuously returns the files in sorted timestamp order. main_input_pcoll = [ FileMetadata( 'path1.py', 10, last_updated_in_seconds=utils._START_TIME_STAMP - 20), FileMetadata( 'path2.py', 10, last_updated_in_seconds=utils._START_TIME_STAMP - 10) ] with TestPipeline() as p: files_pc = ( p | beam.Create(main_input_pcoll) | beam.Map(lambda x: (x.path, x)) | beam.ParDo(utils._GetLatestFileByTimeStamp()) | beam.Map(lambda x: x[0])) assert_that(files_pc, equal_to(['', ''])) def test_latest_file_with_timestamp_after_pipeline_construction_time(self): main_input_pcoll = [ FileMetadata( 'path1.py', 10, last_updated_in_seconds=utils._START_TIME_STAMP + 10) ] with TestPipeline() as p: files_pc = ( p | beam.Create(main_input_pcoll) | beam.Map(lambda x: (x.path, x)) | beam.ParDo(utils._GetLatestFileByTimeStamp()) | beam.Map(lambda x: x[0])) assert_that(files_pc, equal_to(['path1.py'])) def test_emitting_singleton_output(self): # match continuously returns the files in sorted timestamp order. main_input_pcoll = [ FileMetadata( 'path1.py', 10, last_updated_in_seconds=utils._START_TIME_STAMP - 20), # returns default FileMetadata( 'path2.py', 10, last_updated_in_seconds=utils._START_TIME_STAMP - 10), # returns default FileMetadata( 'path3.py', 10, last_updated_in_seconds=utils._START_TIME_STAMP + 10), FileMetadata( 'path4.py', 10, last_updated_in_seconds=utils._START_TIME_STAMP + 20) ] # returns path3.py with TestPipeline() as p: files_pc = ( p | beam.Create(main_input_pcoll) | beam.Map(lambda x: (x.path, x)) | beam.ParDo(utils._GetLatestFileByTimeStamp()) | beam.ParDo(utils._ConvertIterToSingleton()) | beam.Map(lambda x: x[0])) assert_that(files_pc, equal_to(['', 'path3.py', 'path4.py'])) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/vertex_ai_inference.py0000644000175100017510000002736415206113032024615 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import json import logging from collections.abc import Callable from collections.abc import Iterable from collections.abc import Mapping from collections.abc import Sequence from typing import Any from typing import Optional from google.api_core.exceptions import ServerError from google.api_core.exceptions import TooManyRequests from google.cloud import aiplatform from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RemoteModelHandler LOGGER = logging.getLogger("VertexAIModelHandlerJSON") # pylint: disable=line-too-long def _retry_on_appropriate_gcp_error(exception): """ Retry filter that returns True if a returned HTTP error code is 5xx or 429. This is used to retry remote requests that fail, most notably 429 (TooManyRequests.) Args: exception: the returned exception encountered during the request/response loop. Returns: boolean indication whether or not the exception is a Server Error (5xx) or a TooManyRequests (429) error. """ return isinstance(exception, (TooManyRequests, ServerError)) class VertexAIModelHandlerJSON(RemoteModelHandler[Any, PredictionResult, aiplatform.Endpoint]): def __init__( self, endpoint_id: str, project: str, location: str, experiment: Optional[str] = None, network: Optional[str] = None, private: bool = False, invoke_route: Optional[str] = None, *, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, **kwargs): """Implementation of the ModelHandler interface for Vertex AI. **NOTE:** This API and its implementation are under development and do not provide backward compatibility guarantees. Unlike other ModelHandler implementations, this does not load the model being used onto the worker and instead makes remote queries to a Vertex AI endpoint. In that way it functions more like a mid-pipeline IO. Public Vertex AI endpoints have a maximum request size of 1.5 MB. If you wish to make larger requests and use a private endpoint, provide the Compute Engine network you wish to use and set `private=True` Args: endpoint_id: the numerical ID of the Vertex AI endpoint to query project: the GCP project name where the endpoint is deployed location: the GCP location where the endpoint is deployed experiment: optional. experiment label to apply to the queries. See https://cloud.google.com/vertex-ai/docs/experiments/intro-vertex-ai-experiments for more information. network: optional. the full name of the Compute Engine network the endpoint is deployed on; used for private endpoints. The network or subnetwork Dataflow pipeline option must be set and match this network for pipeline execution. Ex: "projects/12345/global/networks/myVPC" private: optional. if the deployed Vertex AI endpoint is private, set to true. Requires a network to be provided as well. invoke_route: optional. the custom route path to use when invoking endpoints with arbitrary prediction routes. When specified, uses `Endpoint.invoke()` instead of `Endpoint.predict()`. The route should start with a forward slash, e.g., "/predict/v1". See https://cloud.google.com/vertex-ai/docs/predictions/use-arbitrary-custom-routes for more information. min_batch_size: optional. the minimum batch size to use when batching inputs. max_batch_size: optional. the maximum batch size to use when batching inputs. max_batch_duration_secs: optional. the maximum amount of time to buffer a batch before emitting; used in streaming contexts. max_batch_weight: optional. the maximum total weight of a batch. element_size_fn: optional. a function that returns the size (weight) of an element. batch_length_fn: optional. a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: optional. a sorted list of positive boundary values for length-aware batching buckets. """ self._batching_kwargs = {} self._env_vars = kwargs.get('env_vars', {}) self._invoke_route = invoke_route if min_batch_size is not None: self._batching_kwargs["min_batch_size"] = min_batch_size if max_batch_size is not None: self._batching_kwargs["max_batch_size"] = max_batch_size if max_batch_duration_secs is not None: self._batching_kwargs["max_batch_duration_secs"] = max_batch_duration_secs if max_batch_weight is not None: self._batching_kwargs["max_batch_weight"] = max_batch_weight if element_size_fn is not None: self._batching_kwargs['element_size_fn'] = element_size_fn if batch_length_fn is not None: self._batching_kwargs['length_fn'] = batch_length_fn if batch_bucket_boundaries is not None: self._batching_kwargs['bucket_boundaries'] = batch_bucket_boundaries if private and network is None: raise ValueError( "A VPC network must be provided to use a private endpoint.") # TODO: support the full list of options for aiplatform.init() # See https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform#google_cloud_aiplatform_init aiplatform.init( project=project, location=location, experiment=experiment, network=network) # Check for liveness here but don't try to actually store the endpoint # in the class yet self.endpoint_name = endpoint_id self.location = location self.is_private = private _ = self._retrieve_endpoint( self.endpoint_name, self.location, self.is_private) super().__init__( namespace='VertexAIModelHandlerJSON', retry_filter=_retry_on_appropriate_gcp_error, **kwargs) def _retrieve_endpoint( self, endpoint_id: str, location: str, is_private: bool) -> aiplatform.Endpoint: """Retrieves an AI Platform endpoint and queries it for liveness/deployed models. Args: endpoint_id: the numerical ID of the Vertex AI endpoint to retrieve. is_private: a boolean indicating if the Vertex AI endpoint is a private endpoint Returns: An aiplatform.Endpoint object Raises: ValueError: if endpoint is inactive or has no models deployed to it. """ if is_private: endpoint: aiplatform.Endpoint = aiplatform.PrivateEndpoint( endpoint_name=endpoint_id, location=location) LOGGER.debug("Treating endpoint %s as private", endpoint_id) else: endpoint = aiplatform.Endpoint( endpoint_name=endpoint_id, location=location) LOGGER.debug("Treating endpoint %s as public", endpoint_id) try: mod_list = endpoint.list_models() except Exception as e: raise ValueError( "Failed to contact endpoint %s, got exception: %s", endpoint_id, e) if len(mod_list) == 0: raise ValueError("Endpoint %s has no models deployed to it.", endpoint_id) return endpoint def create_client(self) -> aiplatform.Endpoint: """Loads the Endpoint object used to build and send prediction request to Vertex AI. """ # Check to make sure the endpoint is still active since pipeline # construction time ep = self._retrieve_endpoint( self.endpoint_name, self.location, self.is_private) return ep def request( self, batch: Sequence[Any], model: aiplatform.Endpoint, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """ Sends a prediction request to a Vertex AI endpoint containing batch of inputs and matches that input with the prediction response from the endpoint as an iterable of PredictionResults. Args: batch: a sequence of any values to be passed to the Vertex AI endpoint. Should be encoded as the model expects. model: an aiplatform.Endpoint object configured to access the desired model. inference_args: any additional arguments to send as part of the prediction request. Returns: An iterable of Predictions. """ if self._invoke_route: # Use invoke() for endpoints with custom prediction routes request_body: dict[str, Any] = {"instances": list(batch)} if inference_args: request_body["parameters"] = inference_args response = model.invoke( request_path=self._invoke_route, body=json.dumps(request_body).encode("utf-8"), headers={"Content-Type": "application/json"}) if hasattr(response, "content"): return self._parse_invoke_response(batch, response.content) return self._parse_invoke_response(batch, bytes(response)) else: prediction = model.predict( instances=list(batch), parameters=inference_args) return utils._convert_to_result( batch, prediction.predictions, prediction.deployed_model_id) def _parse_invoke_response(self, batch: Sequence[Any], response: bytes) -> Iterable[PredictionResult]: """Parses the response from Endpoint.invoke() into PredictionResults. Args: batch: the original batch of inputs. response: the raw bytes response from invoke(). Returns: An iterable of PredictionResults. """ try: response_json = json.loads(response.decode("utf-8")) except (json.JSONDecodeError, UnicodeDecodeError) as e: LOGGER.warning( "Failed to decode invoke response as JSON, returning raw bytes: %s", e) # Return raw response for each batch item return [ PredictionResult(example=example, inference=response) for example in batch ] # Handle standard Vertex AI response format with "predictions" key if isinstance(response_json, dict) and "predictions" in response_json: predictions = response_json["predictions"] model_id = response_json.get("deployedModelId") return utils._convert_to_result(batch, predictions, model_id) # Handle response as a list of predictions (one per input) if isinstance(response_json, list) and len(response_json) == len(batch): return utils._convert_to_result(batch, response_json, None) # Handle single prediction response if len(batch) == 1: return [PredictionResult(example=batch[0], inference=response_json)] # Fallback: return the full response for each batch item return [ PredictionResult(example=example, inference=response_json) for example in batch ] def batch_elements_kwargs(self) -> Mapping[str, Any]: return self._batching_kwargs ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/vertex_ai_inference_it_test.py0000644000175100017510000001065715206113032026345 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-End test for Vertex AI Remote Inference""" import logging import unittest import uuid import pytest import apache_beam as beam from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.base import RunInference from apache_beam.testing.test_pipeline import TestPipeline pytest.importorskip("vertexai", reason="Vertex AI dependencies not available") # pylint: disable=ungrouped-imports try: from apache_beam.examples.inference import vertex_ai_image_classification from apache_beam.ml.inference.vertex_ai_inference import VertexAIModelHandlerJSON except ImportError as e: raise unittest.SkipTest( "Vertex AI model handler dependencies are not installed") _INPUT = "gs://apache-beam-ml/testing/inputs/vertex_images/*/*.jpg" _OUTPUT_DIR = "gs://apache-beam-ml/testing/outputs/vertex_images" _FLOWER_ENDPOINT_ID = "5384055553544683520" _ENDPOINT_PROJECT = "apache-beam-testing" _ENDPOINT_REGION = "us-central1" _ENDPOINT_NETWORK = "projects/844138762903/global/networks/beam-test-vpc" # pylint: disable=line-too-long _SUBNETWORK = "https://www.googleapis.com/compute/v1/projects/apache-beam-testing/regions/us-central1/subnetworks/beam-test-vpc" # Constants for custom prediction routes (invoke) test # Follow beam/sdks/python/apache_beam/ml/inference/test_resources/vertex_ai_custom_prediction/README.md # to get endpoint ID after deploying invoke-enabled model _INVOKE_ENDPOINT_ID = "6890840581900075008" _INVOKE_ROUTE = "/predict" _INVOKE_OUTPUT_DIR = "gs://apache-beam-ml/testing/outputs/vertex_invoke" @pytest.mark.vertex_ai_postcommit class VertexAIInference(unittest.TestCase): def test_vertex_ai_run_flower_image_classification(self): output_file = '/'.join([_OUTPUT_DIR, str(uuid.uuid4()), 'output.txt']) test_pipeline = TestPipeline(is_integration_test=True) extra_opts = { 'input': _INPUT, 'output': output_file, 'endpoint_id': _FLOWER_ENDPOINT_ID, 'endpoint_project': _ENDPOINT_PROJECT, 'endpoint_region': _ENDPOINT_REGION, 'endpoint_network': _ENDPOINT_NETWORK, 'private': "True", 'subnetwork': _SUBNETWORK, } vertex_ai_image_classification.run( test_pipeline.get_full_options_as_args(**extra_opts)) self.assertEqual(FileSystems().exists(output_file), True) @unittest.skipIf( not _INVOKE_ENDPOINT_ID, "Invoke endpoint not configured. Set _INVOKE_ENDPOINT_ID.") def test_vertex_ai_custom_prediction_route(self): """Test custom prediction routes using invoke_route parameter. This test verifies that VertexAIModelHandlerJSON correctly uses Endpoint.invoke() instead of Endpoint.predict() when invoke_route is specified, enabling custom prediction routes. """ output_file = '/'.join( [_INVOKE_OUTPUT_DIR, str(uuid.uuid4()), 'output.txt']) test_pipeline = TestPipeline(is_integration_test=True) model_handler = VertexAIModelHandlerJSON( endpoint_id=_INVOKE_ENDPOINT_ID, project=_ENDPOINT_PROJECT, location=_ENDPOINT_REGION, invoke_route=_INVOKE_ROUTE) # Test inputs - simple data to echo back test_inputs = [{"value": 1}, {"value": 2}, {"value": 3}] with test_pipeline as p: results = ( p | "CreateInputs" >> beam.Create(test_inputs) | "RunInference" >> RunInference(model_handler) | "ExtractResults" >> beam.Map(lambda result: f"{result.example}:{result.inference}")) _ = results | "WriteOutput" >> beam.io.WriteToText( output_file, shard_name_template='') self.assertTrue(FileSystems().exists(output_file)) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/vertex_ai_inference_test.py0000644000175100017510000001076115206113032025645 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest try: from google.api_core.exceptions import TooManyRequests from apache_beam.ml.inference.vertex_ai_inference import VertexAIModelHandlerJSON from apache_beam.ml.inference.vertex_ai_inference import _retry_on_appropriate_gcp_error except ImportError: raise unittest.SkipTest('VertexAI dependencies are not installed') class RetryOnClientErrorTest(unittest.TestCase): def test_retry_on_client_error_positive(self): e = TooManyRequests(message="fake service rate limiting") self.assertTrue(_retry_on_appropriate_gcp_error(e)) def test_retry_on_client_error_negative(self): e = ValueError() self.assertFalse(_retry_on_appropriate_gcp_error(e)) class ModelHandlerArgConditions(unittest.TestCase): def test_exception_on_private_without_network(self): self.assertRaises( ValueError, VertexAIModelHandlerJSON, endpoint_id="1", project="testproject", location="us-central1", private=True) class ParseInvokeResponseTest(unittest.TestCase): """Tests for _parse_invoke_response method.""" def _create_handler_with_invoke_route(self, invoke_route="/test"): """Creates a mock handler with invoke_route for testing.""" import unittest.mock as mock # Mock both _retrieve_endpoint and aiplatform.init to prevent test # pollution of global aiplatform state with mock.patch.object(VertexAIModelHandlerJSON, '_retrieve_endpoint', return_value=None): with mock.patch('google.cloud.aiplatform.init'): handler = VertexAIModelHandlerJSON( endpoint_id="1", project="testproject", location="us-central1", invoke_route=invoke_route) return handler def test_parse_invoke_response_with_predictions_key(self): """Test parsing response with standard 'predictions' key.""" handler = self._create_handler_with_invoke_route() batch = [{"input": "test1"}, {"input": "test2"}] response = ( b'{"predictions": ["result1", "result2"], ' b'"deployedModelId": "model123"}') results = list(handler._parse_invoke_response(batch, response)) self.assertEqual(len(results), 2) self.assertEqual(results[0].example, {"input": "test1"}) self.assertEqual(results[0].inference, "result1") self.assertEqual(results[1].example, {"input": "test2"}) self.assertEqual(results[1].inference, "result2") def test_parse_invoke_response_list_format(self): """Test parsing response as a list of predictions.""" handler = self._create_handler_with_invoke_route() batch = [{"input": "test1"}, {"input": "test2"}] response = b'["result1", "result2"]' results = list(handler._parse_invoke_response(batch, response)) self.assertEqual(len(results), 2) self.assertEqual(results[0].inference, "result1") self.assertEqual(results[1].inference, "result2") def test_parse_invoke_response_single_prediction(self): """Test parsing response with a single prediction.""" handler = self._create_handler_with_invoke_route() batch = [{"input": "test1"}] response = b'{"output": "single result"}' results = list(handler._parse_invoke_response(batch, response)) self.assertEqual(len(results), 1) self.assertEqual(results[0].inference, {"output": "single result"}) def test_parse_invoke_response_non_json(self): """Test handling non-JSON response.""" handler = self._create_handler_with_invoke_route() batch = [{"input": "test1"}] response = b'not valid json' results = list(handler._parse_invoke_response(batch, response)) self.assertEqual(len(results), 1) self.assertEqual(results[0].inference, response) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/vllm_inference.py0000644000175100017510000003542415206113032023575 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import asyncio import logging import os import subprocess import sys import threading import time import uuid from collections.abc import Callable from collections.abc import Iterable from collections.abc import Sequence from dataclasses import dataclass from typing import Any from typing import Optional from openai import AsyncOpenAI from openai import OpenAI from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.utils import subprocess_server try: # VLLM logging config breaks beam logging. os.environ["VLLM_CONFIGURE_LOGGING"] = "0" import vllm # pylint: disable=unused-import logging.info('vllm module successfully imported.') os.environ["VLLM_CONFIGURE_LOGGING"] = "1" except ModuleNotFoundError: msg = 'vllm module was not found. This is ok as long as the specified ' \ 'runner has vllm dependencies installed.' logging.warning(msg) __all__ = [ 'OpenAIChatMessage', 'VLLMCompletionsModelHandler', 'VLLMChatModelHandler', ] @dataclass(frozen=True) class OpenAIChatMessage(): """" Dataclass containing previous chat messages in conversation. Role is the entity that sent the message (either 'user' or 'system'). Content is the contents of the message. """ role: str content: str def start_process(cmd) -> tuple[subprocess.Popen, int]: port, = subprocess_server.pick_port(None) cmd = [arg.replace('{{PORT}}', str(port)) for arg in cmd] # pylint: disable=not-an-iterable logging.info("Starting service with %s", str(cmd).replace("',", "'")) process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) # Emit the output of this command as info level logging. def log_stdout(): line = process.stdout.readline() while line: # The log obtained from stdout is bytes, decode it into string. # Remove newline via rstrip() to not print an empty line. logging.info(line.decode(errors='backslashreplace').rstrip()) line = process.stdout.readline() t = threading.Thread(target=log_stdout) t.daemon = True t.start() return process, port def getVLLMClient(port) -> OpenAI: openai_api_key = "EMPTY" openai_api_base = f"http://localhost:{port}/v1" return OpenAI( api_key=openai_api_key, base_url=openai_api_base, ) def getAsyncVLLMClient(port) -> AsyncOpenAI: openai_api_key = "EMPTY" openai_api_base = f"http://localhost:{port}/v1" return AsyncOpenAI( api_key=openai_api_key, base_url=openai_api_base, ) class _VLLMModelServer(): def __init__(self, model_name: str, vllm_server_kwargs: dict[str, str]): self._model_name = model_name self._vllm_server_kwargs = vllm_server_kwargs self._server_started = False self._server_process = None self._server_port: int = -1 self._server_process_lock = threading.RLock() self.start_server() def start_server(self, retries=3): with self._server_process_lock: if not self._server_started: server_cmd = [ sys.executable, '-m', 'vllm.entrypoints.openai.api_server', '--model', self._model_name, '--port', '{{PORT}}', ] for k, v in self._vllm_server_kwargs.items(): server_cmd.append(f'--{k}') # Only add values for commands with value part. if v is not None: server_cmd.append(v) self._server_process, self._server_port = start_process(server_cmd) self.check_connectivity(retries) def get_server_port(self) -> int: if not self._server_started: self.start_server() return self._server_port def check_connectivity(self, retries=3): with getVLLMClient(self._server_port) as client: while self._server_process.poll() is None: try: models = client.models.list().data logging.info('models: %s' % models) if len(models) > 0: self._server_started = True return except: # pylint: disable=bare-except pass # Sleep while bringing up the process time.sleep(5) if retries == 0: self._server_started = False raise Exception( "Failed to start vLLM server, polling process exited with code " + "%s. Next time a request is tried, the server will be restarted" % self._server_process.poll()) else: self.start_server(retries - 1) class VLLMCompletionsModelHandler(ModelHandler[str, PredictionResult, _VLLMModelServer]): def __init__( self, model_name: str, vllm_server_kwargs: Optional[dict[str, str]] = None, *, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None): """Implementation of the ModelHandler interface for vLLM using text as input. Example Usage:: pcoll | RunInference(VLLMModelHandler(model_name='facebook/opt-125m')) Args: model_name: The vLLM model. See https://docs.vllm.ai/en/latest/models/supported_models.html for supported models. vllm_server_kwargs: Any additional kwargs to be passed into your vllm server when it is being created. Will be invoked using `python -m vllm.entrypoints.openai.api_serverv `. For example, you could pass `{'echo': 'true'}` to prepend new messages with the previous message. On ~16GB GPUs, pass lower ``max-num-seqs`` and ``gpu-memory-utilization`` values (see ``apache_beam.examples.inference.vllm_text_completion``). For a list of possible kwargs, see https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#extra-parameters-for-completions-api min_batch_size: optional. the minimum batch size to use when batching inputs. max_batch_size: optional. the maximum batch size to use when batching inputs. max_batch_duration_secs: optional. the maximum amount of time to buffer a batch before emitting; used in streaming contexts. max_batch_weight: optional. the maximum total weight of a batch. element_size_fn: optional. a function that returns the size (weight) of an element. batch_length_fn: optional. a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: optional. a sorted list of positive boundary values for length-aware batching buckets. """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries) self._model_name = model_name self._vllm_server_kwargs: dict[str, str] = vllm_server_kwargs or {} def load_model(self) -> _VLLMModelServer: return _VLLMModelServer(self._model_name, self._vllm_server_kwargs) async def _async_run_inference( self, batch: Sequence[str], model: _VLLMModelServer, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: inference_args = inference_args or {} async with getAsyncVLLMClient(model.get_server_port()) as client: try: async_predictions = [ client.completions.create( model=self._model_name, prompt=prompt, **inference_args) for prompt in batch ] responses = await asyncio.gather(*async_predictions) except Exception as e: model.check_connectivity() raise e return [PredictionResult(x, y) for x, y in zip(batch, responses)] def run_inference( self, batch: Sequence[str], model: _VLLMModelServer, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """Runs inferences on a batch of text strings. Args: batch: A sequence of examples as text strings. model: A _VLLMModelServer containing info for connecting to the server. inference_args: Any additional arguments for an inference. Returns: An Iterable of type PredictionResult. """ return asyncio.run(self._async_run_inference(batch, model, inference_args)) def share_model_across_processes(self) -> bool: return True class VLLMChatModelHandler(ModelHandler[Sequence[OpenAIChatMessage], PredictionResult, _VLLMModelServer]): def __init__( self, model_name: str, chat_template_path: Optional[str] = None, vllm_server_kwargs: Optional[dict[str, str]] = None, *, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None): """ Implementation of the ModelHandler interface for vLLM using previous messages as input. Example Usage:: pcoll | RunInference(VLLMModelHandler(model_name='facebook/opt-125m')) Args: model_name: The vLLM model. See https://docs.vllm.ai/en/latest/models/supported_models.html for supported models. chat_template_path: Path to a chat template. This file must be accessible from your runner's execution environment, so it is recommended to use a cloud based file storage system (e.g. Google Cloud Storage). For info on chat templates, see: https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#chat-template vllm_server_kwargs: Any additional kwargs to be passed into your vllm server when it is being created. Will be invoked using `python -m vllm.entrypoints.openai.api_serverv `. For example, you could pass `{'echo': 'true'}` to prepend new messages with the previous message. For a list of possible kwargs, see https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#extra-parameters-for-chat-api min_batch_size: optional. the minimum batch size to use when batching inputs. max_batch_size: optional. the maximum batch size to use when batching inputs. max_batch_duration_secs: optional. the maximum amount of time to buffer a batch before emitting; used in streaming contexts. max_batch_weight: optional. the maximum total weight of a batch. element_size_fn: optional. a function that returns the size (weight) of an element. batch_length_fn: optional. a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: optional. a sorted list of positive boundary values for length-aware batching buckets. """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries) self._model_name = model_name self._vllm_server_kwargs: dict[str, str] = vllm_server_kwargs or {} self._chat_template_path = chat_template_path self._chat_file = f'template-{uuid.uuid4().hex}.jinja' def load_model(self) -> _VLLMModelServer: chat_template_contents = '' if self._chat_template_path is not None: local_chat_template_path = os.path.join(os.getcwd(), self._chat_file) if not os.path.exists(local_chat_template_path): with FileSystems.open(self._chat_template_path) as fin: chat_template_contents = fin.read().decode() with open(local_chat_template_path, 'a') as f: f.write(chat_template_contents) self._vllm_server_kwargs['chat_template'] = local_chat_template_path return _VLLMModelServer(self._model_name, self._vllm_server_kwargs) async def _async_run_inference( self, batch: Sequence[Sequence[OpenAIChatMessage]], model: _VLLMModelServer, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: inference_args = inference_args or {} async with getAsyncVLLMClient(model.get_server_port()) as client: try: async_predictions = [ client.chat.completions.create( model=self._model_name, messages=[{ "role": message.role, "content": message.content } for message in messages], **inference_args) for messages in batch ] predictions = await asyncio.gather(*async_predictions) except Exception as e: model.check_connectivity() raise e return [PredictionResult(x, y) for x, y in zip(batch, predictions)] def run_inference( self, batch: Sequence[Sequence[OpenAIChatMessage]], model: _VLLMModelServer, inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """Runs inferences on a batch of text strings. Args: batch: A sequence of examples as OpenAI messages. model: A _VLLMModelServer for connecting to the spun up server. inference_args: Any additional arguments for an inference. Returns: An Iterable of type PredictionResult. """ return asyncio.run(self._async_run_inference(batch, model, inference_args)) def share_model_across_processes(self) -> bool: return True ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/xgboost_inference.py0000644000175100017510000003310215206113032024277 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import sys from abc import ABC from collections.abc import Callable from collections.abc import Iterable from collections.abc import Sequence from typing import Any from typing import Optional from typing import Union import datatable import numpy import pandas import scipy import xgboost from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.base import ExampleT from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import ModelT from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import PredictionT __all__ = [ 'XGBoostModelHandler', 'XGBoostModelHandlerNumpy', 'XGBoostModelHandlerPandas', 'XGBoostModelHandlerSciPy', 'XGBoostModelHandlerDatatable' ] XGBoostInferenceFn = Callable[[ Sequence[object], Union[xgboost.Booster, xgboost.XGBModel], Optional[dict[str, Any]] ], Iterable[PredictionResult]] def default_xgboost_inference_fn( batch: Sequence[object], model: Union[xgboost.Booster, xgboost.XGBModel], inference_args: Optional[dict[str, Any]] = None) -> Iterable[PredictionResult]: inference_args = {} if not inference_args else inference_args if type(model) == xgboost.Booster: batch = [xgboost.DMatrix(array) for array in batch] predictions = [model.predict(el, **inference_args) for el in batch] return [PredictionResult(x, y) for x, y in zip(batch, predictions)] class XGBoostModelHandler(ModelHandler[ExampleT, PredictionT, ModelT], ABC): def __init__( self, model_class: Union[Callable[..., xgboost.Booster], Callable[..., xgboost.XGBModel]], model_state: str, inference_fn: XGBoostInferenceFn = default_xgboost_inference_fn, *, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None, max_batch_weight: Optional[int] = None, element_size_fn: Optional[Callable[[Any], int]] = None, batch_length_fn: Optional[Callable[[Any], int]] = None, batch_bucket_boundaries: Optional[list[int]] = None, **kwargs): """Implementation of the ModelHandler interface for XGBoost. Example Usage:: pcoll | RunInference( XGBoostModelHandler( model_class="XGBoost Model Class", model_state="my_model_state.json"))) See https://xgboost.readthedocs.io/en/stable/tutorials/saving_model.html for details Args: model_class: class of the XGBoost model that defines the model structure. model_state: path to a json file that contains the model's configuration. inference_fn: the inference function to use during RunInference. default=default_xgboost_inference_fn min_batch_size: optional. the minimum batch size to use when batching inputs. max_batch_size: optional. the maximum batch size to use when batching inputs. max_batch_duration_secs: optional. the maximum amount of time to buffer a batch before emitting; used in streaming contexts. max_batch_weight: optional. the maximum total weight of a batch. element_size_fn: optional. a function that returns the size (weight) of an element. batch_length_fn: optional. a callable that returns the length of an element for length-aware batching. batch_bucket_boundaries: optional. a sorted list of positive boundary values for length-aware batching buckets. kwargs: 'env_vars' can be used to set environment variables before loading the model. **Supported Versions:** RunInference APIs in Apache Beam have been tested with XGBoost 1.6.0 and 1.7.0 XGBoost 1.0.0 introduced support for using JSON to save and load XGBoost models. XGBoost 1.6.0, additional support for Universal Binary JSON. It is recommended to use a model trained in XGBoost 1.6.0 or higher. While you should be able to load models created in older versions, there are no guarantees this will work as expected. This class is the superclass of all the various XGBoostModelhandlers and should not be instantiated directly. (See instead XGBoostModelHandlerNumpy, XGBoostModelHandlerPandas, etc.) """ super().__init__( min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs, max_batch_weight=max_batch_weight, element_size_fn=element_size_fn, batch_length_fn=batch_length_fn, batch_bucket_boundaries=batch_bucket_boundaries, **kwargs) self._model_class = model_class self._model_state = model_state self._inference_fn = inference_fn def load_model(self) -> Union[xgboost.Booster, xgboost.XGBModel]: model = self._model_class() model_state_file_handler = FileSystems.open(self._model_state, 'rb') model_state_bytes = model_state_file_handler.read() # Convert into a bytearray so that the # model state can be loaded in XGBoost model_state_bytearray = bytearray(model_state_bytes) model.load_model(model_state_bytearray) return model def get_metrics_namespace(self) -> str: return 'BeamML_XGBoost' class XGBoostModelHandlerNumpy(XGBoostModelHandler[numpy.ndarray, PredictionResult, Union[xgboost.Booster, xgboost.XGBModel]]): """Implementation of the ModelHandler interface for XGBoost using numpy arrays as input. Example Usage:: pcoll | RunInference( XGBoostModelHandlerNumpy( model_class="XGBoost Model Class", model_state="my_model_state.json"))) Args: model_class: class of the XGBoost model that defines the model structure. model_state: path to a json file that contains the model's configuration. inference_fn: the inference function to use during RunInference. default=default_xgboost_inference_fn """ def run_inference( self, batch: Sequence[numpy.ndarray], model: Union[xgboost.Booster, xgboost.XGBModel], inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """Runs inferences on a batch of 2d numpy arrays. Args: batch: A sequence of examples as 2d numpy arrays. Each row in an array is a single example. The dimensions must match the dimensions of the data used to train the model. model: XGBoost booster or XBGModel (sklearn interface). Must implement predict(X). Where the parameter X is a 2d numpy array. inference_args: Any additional arguments for an inference. Returns: An Iterable of type PredictionResult. """ return self._inference_fn(batch, model, inference_args) def get_num_bytes(self, batch: Sequence[numpy.ndarray]) -> int: """ Returns: The number of bytes of data for a batch. """ return sum(sys.getsizeof(element) for element in batch) class XGBoostModelHandlerPandas(XGBoostModelHandler[pandas.DataFrame, PredictionResult, Union[xgboost.Booster, xgboost.XGBModel]]): """Implementation of the ModelHandler interface for XGBoost using pandas dataframes as input. Example Usage:: pcoll | RunInference( XGBoostModelHandlerPandas( model_class="XGBoost Model Class", model_state="my_model_state.json"))) Args: model_class: class of the XGBoost model that defines the model structure. model_state: path to a json file that contains the model's configuration. inference_fn: the inference function to use during RunInference. default=default_xgboost_inference_fn """ def run_inference( self, batch: Sequence[pandas.DataFrame], model: Union[xgboost.Booster, xgboost.XGBModel], inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """Runs inferences on a batch of pandas dataframes. Args: batch: A sequence of examples as pandas dataframes. Each row in a dataframe is a single example. The dimensions must match the dimensions of the data used to train the model. model: XGBoost booster or XBGModel (sklearn interface). Must implement predict(X). Where the parameter X is a pandas dataframe. inference_args: Any additional arguments for an inference. Returns: An Iterable of type PredictionResult. """ return self._inference_fn(batch, model, inference_args) def get_num_bytes(self, batch: Sequence[pandas.DataFrame]) -> int: """ Returns: The number of bytes of data for a batch of Numpy arrays. """ return sum(df.memory_usage(deep=True).sum() for df in batch) class XGBoostModelHandlerSciPy(XGBoostModelHandler[scipy.sparse.csr_matrix, PredictionResult, Union[xgboost.Booster, xgboost.XGBModel]]): """ Implementation of the ModelHandler interface for XGBoost using scipy matrices as input. Example Usage:: pcoll | RunInference( XGBoostModelHandlerSciPy( model_class="XGBoost Model Class", model_state="my_model_state.json"))) Args: model_class: class of the XGBoost model that defines the model structure. model_state: path to a json file that contains the model's configuration. inference_fn: the inference function to use during RunInference. default=default_xgboost_inference_fn """ def run_inference( self, batch: Sequence[scipy.sparse.csr_matrix], model: Union[xgboost.Booster, xgboost.XGBModel], inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """Runs inferences on a batch of SciPy sparse matrices. Args: batch: A sequence of examples as Scipy sparse matrices. The dimensions must match the dimensions of the data used to train the model. model: XGBoost booster or XBGModel (sklearn interface). Must implement predict(X). Where the parameter X is a SciPy sparse matrix. inference_args: Any additional arguments for an inference. Returns: An Iterable of type PredictionResult. """ return self._inference_fn(batch, model, inference_args) def get_num_bytes(self, batch: Sequence[scipy.sparse.csr_matrix]) -> int: """ Returns: The number of bytes of data for a batch. """ return sum(sys.getsizeof(element) for element in batch) class XGBoostModelHandlerDatatable(XGBoostModelHandler[datatable.Frame, PredictionResult, Union[xgboost.Booster, xgboost.XGBModel]] ): """Implementation of the ModelHandler interface for XGBoost using datatable dataframes as input. Example Usage:: pcoll | RunInference( XGBoostModelHandlerDatatable( model_class="XGBoost Model Class", model_state="my_model_state.json"))) Args: model_class: class of the XGBoost model that defines the model structure. model_state: path to a json file that contains the model's configuration. inference_fn: the inference function to use during RunInference. default=default_xgboost_inference_fn """ def run_inference( self, batch: Sequence[datatable.Frame], model: Union[xgboost.Booster, xgboost.XGBModel], inference_args: Optional[dict[str, Any]] = None ) -> Iterable[PredictionResult]: """Runs inferences on a batch of datatable dataframe. Args: batch: A sequence of examples as datatable dataframes. Each row in a dataframe is a single example. The dimensions must match the dimensions of the data used to train the model. model: XGBoost booster or XBGModel (sklearn interface). Must implement predict(X). Where the parameter X is a datatable dataframe. inference_args: Any additional arguments for an inference. Returns: An Iterable of type PredictionResult. """ return self._inference_fn(batch, model, inference_args) def get_num_bytes(self, batch: Sequence[datatable.Frame]) -> int: """ Returns: The number of bytes of data for a batch. """ return sum(sys.getsizeof(element) for element in batch) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/xgboost_inference_it_test.py0000644000175100017510000003224415206113032026040 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import os import uuid try: import unittest import pytest import xgboost from apache_beam.examples.inference import xgboost_iris_classification from apache_beam.io.filesystems import FileSystems from apache_beam.testing.test_pipeline import TestPipeline except ImportError as e: xgboost = None EXPECTED_OUTPUT_SINGLE_BATCHES = [ "0,[1 1 1 0 0 0 0 1 2 0 0 2 0 2 1 2 2 2 2 0 0 0 0 2 2 0 2 2 2 1]" ] EXPECTED_OUTPUT_MULTIPLE_BATCHES = [ "0,[1]", "1,[1]", "2,[1]", "3,[0]", "4,[0]", "5,[0]", "6,[0]", "7,[1]", "8,[2]", "9,[0]", "10,[0]", "11,[2]", "12,[0]", "13,[2]", "14,[1]", "15,[2]", "16,[2]", "17,[2]", "18,[2]", "19,[0]", "20,[0]", "21,[0]", "22,[0]", "23,[2]", "24,[2]", "25,[0]", "26,[2]", "27,[2]", "28,[2]", "29,[1]", ] def process_outputs(filepath): with FileSystems().open(filepath) as f: lines = f.readlines() lines = [l.decode('utf-8').strip('\n') for l in lines] return lines @unittest.skipIf( os.getenv('FORCE_XGBOOST_IT') is None and xgboost is None, 'Missing dependencies. ' 'Test depends on xgboost and datatable') @pytest.mark.uses_xgboost @pytest.mark.it_postcommit class XGBoostInference(unittest.TestCase): def test_iris_classification_numpy_single_batch(self): test_pipeline = TestPipeline(is_integration_test=True) input_type = 'numpy' output_file_dir = '/tmp' output_file = '/'.join( [output_file_dir, str(uuid.uuid4()), 'numpy_single_batch.txt']) model_state_path = 'gs://apache-beam-ml/models/xgboost.iris_classifier.json' extra_opts = { 'input_type': input_type, 'output': output_file, 'model_state': model_state_path, 'no_split': True } xgboost_iris_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_outputs = EXPECTED_OUTPUT_SINGLE_BATCHES predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) predictions_dict = {} for predicted_output in predicted_outputs: true_label, prediction = predicted_output.split(',') predictions_dict[true_label] = prediction for expected_output in expected_outputs: true_label, expected_prediction = expected_output.split(',') self.assertEqual(predictions_dict[true_label], expected_prediction) def test_iris_classification_numpy_single_batch_large_model(self): test_pipeline = TestPipeline(is_integration_test=True) input_type = 'numpy' output_file_dir = '/tmp' output_file = '/'.join( [output_file_dir, str(uuid.uuid4()), 'numpy_single_batch.txt']) model_state_path = 'gs://apache-beam-ml/models/xgboost.iris_classifier.json' extra_opts = { 'input_type': input_type, 'output': output_file, 'model_state': model_state_path, 'no_split': True, 'large_model': True, } xgboost_iris_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_outputs = EXPECTED_OUTPUT_SINGLE_BATCHES predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) predictions_dict = {} for predicted_output in predicted_outputs: true_label, prediction = predicted_output.split(',') predictions_dict[true_label] = prediction for expected_output in expected_outputs: true_label, expected_prediction = expected_output.split(',') self.assertEqual(predictions_dict[true_label], expected_prediction) def test_iris_classification_pandas_single_batch(self): test_pipeline = TestPipeline(is_integration_test=True) input_type = 'pandas' output_file_dir = '/tmp' output_file = '/'.join( [output_file_dir, str(uuid.uuid4()), 'pandas_single_batch.txt']) model_state_path = 'gs://apache-beam-ml/models/xgboost.iris_classifier.json' extra_opts = { 'input_type': input_type, 'output': output_file, 'model_state': model_state_path, 'no_split': True } xgboost_iris_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_outputs = EXPECTED_OUTPUT_SINGLE_BATCHES predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) predictions_dict = {} for predicted_output in predicted_outputs: true_label, prediction = predicted_output.split(',') predictions_dict[true_label] = prediction for expected_output in expected_outputs: true_label, expected_prediction = expected_output.split(',') self.assertEqual(predictions_dict[true_label], expected_prediction) def test_iris_classification_scipy_single_batch(self): test_pipeline = TestPipeline(is_integration_test=True) input_type = 'scipy' output_file_dir = '/tmp' output_file = '/'.join( [output_file_dir, str(uuid.uuid4()), 'scipy_single_batch.txt']) model_state_path = 'gs://apache-beam-ml/models/xgboost.iris_classifier.json' extra_opts = { 'input_type': input_type, 'output': output_file, 'model_state': model_state_path, 'no_split': True } xgboost_iris_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_outputs = EXPECTED_OUTPUT_SINGLE_BATCHES predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) predictions_dict = {} for predicted_output in predicted_outputs: true_label, prediction = predicted_output.split(',') predictions_dict[true_label] = prediction for expected_output in expected_outputs: true_label, expected_prediction = expected_output.split(',') self.assertEqual(predictions_dict[true_label], expected_prediction) def test_iris_classification_datatable_single_batch(self): test_pipeline = TestPipeline(is_integration_test=True) input_type = 'datatable' output_file_dir = '/tmp' output_file = '/'.join( [output_file_dir, str(uuid.uuid4()), 'datatable_single_batch.txt']) model_state_path = 'gs://apache-beam-ml/models/xgboost.iris_classifier.json' extra_opts = { 'input_type': input_type, 'output': output_file, 'model_state': model_state_path, 'no_split': True } xgboost_iris_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_outputs = EXPECTED_OUTPUT_SINGLE_BATCHES predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) predictions_dict = {} for predicted_output in predicted_outputs: true_label, prediction = predicted_output.split(',') predictions_dict[true_label] = prediction for expected_output in expected_outputs: true_label, expected_prediction = expected_output.split(',') self.assertEqual(predictions_dict[true_label], expected_prediction) def test_iris_classification_numpy_multi_batch(self): test_pipeline = TestPipeline(is_integration_test=True) input_type = 'numpy' output_file_dir = '/tmp' output_file = '/'.join( [output_file_dir, str(uuid.uuid4()), 'numpy_multi_batch.txt']) model_state_path = 'gs://apache-beam-ml/models/xgboost.iris_classifier.json' extra_opts = { 'input_type': input_type, 'output': output_file, 'model_state': model_state_path, 'split': True } xgboost_iris_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_outputs = EXPECTED_OUTPUT_MULTIPLE_BATCHES predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) predictions_dict = {} for predicted_output in predicted_outputs: true_label, prediction = predicted_output.split(',') predictions_dict[true_label] = prediction for expected_output in expected_outputs: true_label, expected_prediction = expected_output.split(',') self.assertEqual(predictions_dict[true_label], expected_prediction) def test_iris_classification_pandas_multi_batch(self): test_pipeline = TestPipeline(is_integration_test=True) input_type = 'pandas' output_file_dir = '/tmp' output_file = '/'.join( [output_file_dir, str(uuid.uuid4()), 'pandas_multi_batch.txt']) model_state_path = 'gs://apache-beam-ml/models/xgboost.iris_classifier.json' extra_opts = { 'input_type': input_type, 'output': output_file, 'model_state': model_state_path, 'split': True } xgboost_iris_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_outputs = EXPECTED_OUTPUT_MULTIPLE_BATCHES predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) predictions_dict = {} for predicted_output in predicted_outputs: true_label, prediction = predicted_output.split(',') predictions_dict[true_label] = prediction for expected_output in expected_outputs: true_label, expected_prediction = expected_output.split(',') self.assertEqual(predictions_dict[true_label], expected_prediction) def test_iris_classification_scipy_multi_batch(self): test_pipeline = TestPipeline(is_integration_test=True) input_type = 'scipy' output_file_dir = '/tmp' output_file = '/'.join( [output_file_dir, str(uuid.uuid4()), 'scipy_multi_batch.txt']) model_state_path = 'gs://apache-beam-ml/models/xgboost.iris_classifier.json' extra_opts = { 'input_type': input_type, 'output': output_file, 'model_state': model_state_path, 'split': True } xgboost_iris_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_outputs = EXPECTED_OUTPUT_MULTIPLE_BATCHES predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) predictions_dict = {} for predicted_output in predicted_outputs: true_label, prediction = predicted_output.split(',') predictions_dict[true_label] = prediction for expected_output in expected_outputs: true_label, expected_prediction = expected_output.split(',') self.assertEqual(predictions_dict[true_label], expected_prediction) def test_iris_classification_datatable_multi_batch(self): test_pipeline = TestPipeline(is_integration_test=True) input_type = 'datatable' output_file_dir = '/tmp' output_file = '/'.join( [output_file_dir, str(uuid.uuid4()), 'datatable_multi_batch.txt']) model_state_path = 'gs://apache-beam-ml/models/xgboost.iris_classifier.json' extra_opts = { 'input_type': input_type, 'output': output_file, 'model_state': model_state_path, 'split': True } xgboost_iris_classification.run( test_pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) self.assertEqual(FileSystems().exists(output_file), True) expected_outputs = EXPECTED_OUTPUT_MULTIPLE_BATCHES predicted_outputs = process_outputs(output_file) self.assertEqual(len(expected_outputs), len(predicted_outputs)) predictions_dict = {} for predicted_output in predicted_outputs: true_label, prediction = predicted_output.split(',') predictions_dict[true_label] = prediction for expected_output in expected_outputs: true_label, expected_prediction = expected_output.split(',') self.assertEqual(predictions_dict[true_label], expected_prediction) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/inference/xgboost_inference_test.py0000644000175100017510000004653415206113032025353 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import os import shutil import sys import tempfile import unittest import zipfile from typing import Any try: import datatable import numpy import pandas import pytest import scipy import xgboost import apache_beam as beam from apache_beam.ml.inference import RunInference from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.xgboost_inference import XGBoostModelHandlerDatatable from apache_beam.ml.inference.xgboost_inference import XGBoostModelHandlerNumpy from apache_beam.ml.inference.xgboost_inference import XGBoostModelHandlerPandas from apache_beam.ml.inference.xgboost_inference import XGBoostModelHandlerSciPy from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to except ImportError: raise unittest.SkipTest('XGBoost dependencies are not installed') def _compare_prediction_result(a: PredictionResult, b: PredictionResult): if isinstance(a.example, scipy.sparse.csr_matrix) and isinstance( b.example, scipy.sparse.csr_matrix): example_equal = numpy.array_equal(a.example.todense(), b.example.todense()) else: example_equal = numpy.array_equal(a.example, b.example) # type: ignore[arg-type] if isinstance(a.inference, dict): return all( x == y for x, y in zip(a.inference.values(), b.inference.values())) and example_equal return a.inference == b.inference and example_equal def _compare_keyed_prediction_result( a: tuple[Any, PredictionResult], b: tuple[Any, PredictionResult]): a_key, a_val = a b_key, b_val = b keys_equal = a_key == b_key return _compare_prediction_result(a_val, b_val) and keys_equal def predict_fn(self, data): self.inference_calls += 1 if isinstance(data, pandas.DataFrame): data = data.to_numpy() if isinstance(data, datatable.Frame): data = data.to_numpy() if isinstance(data, scipy.sparse.csr_matrix): data = data.toarray() return sum(sum(array) for array in data) @pytest.fixture(autouse=True) def predict_patched(monkeypatch): monkeypatch.setattr(xgboost.XGBClassifier, 'predict', predict_fn) def build_monkeypatched_xgboost_classifier() -> xgboost.XGBClassifier: model = xgboost.XGBClassifier() model.inference_calls = 0 model.fit([[0, 0], [0, 1], [1, 0], [1, 1]], [0, 1, 0, 1]) return model @pytest.mark.uses_xgboost class XGBoostRunInferenceTest(unittest.TestCase): def setUp(self): self.tmpdir = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.tmpdir) def test_predict_output(self): model = build_monkeypatched_xgboost_classifier() inference_runner = XGBoostModelHandlerNumpy(xgboost.XGBClassifier, 'unused') batched_examples = [ numpy.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]), numpy.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]]) ] expected_predictions = [ PredictionResult(numpy.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]), 45), PredictionResult(numpy.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]]), 9), ] inferences = inference_runner.run_inference(batched_examples, model) for actual, expected in zip(inferences, expected_predictions): self.assertTrue(_compare_prediction_result(actual, expected)) def test_single_inference_call(self): model = build_monkeypatched_xgboost_classifier() inference_runner = XGBoostModelHandlerNumpy(xgboost.XGBClassifier, 'unused') self.assertEqual(model.inference_calls, 0) batched_examples = [numpy.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])] _ = inference_runner.run_inference(batched_examples, model) self.assertEqual(model.inference_calls, 1) def test_multiple_inference_calls(self): model = build_monkeypatched_xgboost_classifier() inference_runner = XGBoostModelHandlerNumpy(xgboost.XGBClassifier, 'unused') self.assertEqual(model.inference_calls, 0) batched_examples = [ numpy.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]), numpy.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]), numpy.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]), numpy.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]), numpy.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) ] _ = inference_runner.run_inference(batched_examples, model) self.assertEqual(model.inference_calls, 5) def test_num_bytes_numpy(self): inference_runner = XGBoostModelHandlerNumpy( model_class=xgboost.XGBClassifier, model_state='unused') batched_examples_int = [ numpy.array([[1, 1], [2, 2]]), numpy.array([[2, 4], [6, 8]]), ] self.assertEqual( sys.getsizeof(batched_examples_int[0]) + sys.getsizeof(batched_examples_int[1]), inference_runner.get_num_bytes(batched_examples_int)) batched_examples_float = [ numpy.array([[1.0, 1.0], [2.0, 2.0]]), numpy.array([[2.0, 4.0], [6.0, 8.0]]), ] self.assertEqual( sys.getsizeof(batched_examples_float[0]) + sys.getsizeof(batched_examples_float[1]), inference_runner.get_num_bytes(batched_examples_float)) def test_num_bytes_pandas(self): inference_runner = XGBoostModelHandlerPandas( model_class=xgboost.XGBClassifier, model_state='unused') batched_examples_int = [ pandas.DataFrame([[1, 1], [2, 2]]), pandas.DataFrame([[2, 4], [6, 8]]), ] self.assertEqual( batched_examples_int[0].memory_usage(deep=True).sum() + batched_examples_int[1].memory_usage(deep=True).sum(), inference_runner.get_num_bytes(batched_examples_int)) batched_examples_float = [ pandas.DataFrame([[1.0, 1.0], [2.0, 2.0]]), pandas.DataFrame([[2.0, 4.0], [6.0, 8.0]]), ] self.assertEqual( batched_examples_float[0].memory_usage(deep=True).sum() + batched_examples_float[1].memory_usage(deep=True).sum(), inference_runner.get_num_bytes(batched_examples_float)) def test_num_bytes_datatable(self): inference_runner = XGBoostModelHandlerDatatable( model_class=xgboost.XGBClassifier, model_state='unused') batched_examples_int = [ datatable.Frame([[1, 1], [2, 2]]), datatable.Frame([[2, 4], [6, 8]]), ] self.assertEqual( sys.getsizeof(batched_examples_int[0]) + sys.getsizeof(batched_examples_int[1]), inference_runner.get_num_bytes(batched_examples_int)) batched_examples_float = [ datatable.Frame([[1.0, 1.0], [2.0, 2.0]]), datatable.Frame([[2.0, 4.0], [6.0, 8.0]]), ] self.assertEqual( sys.getsizeof(batched_examples_float[0]) + sys.getsizeof(batched_examples_float[1]), inference_runner.get_num_bytes(batched_examples_float)) def test_num_bytes_scipy(self): inference_runner = XGBoostModelHandlerSciPy( model_class=xgboost.XGBClassifier, model_state='unused') batched_examples_int = [ scipy.sparse.csr_matrix([[1, 1], [2, 2]]), scipy.sparse.csr_matrix([[2, 4], [6, 8]]), ] self.assertEqual( sys.getsizeof(batched_examples_int[0]) + sys.getsizeof(batched_examples_int[1]), inference_runner.get_num_bytes(batched_examples_int)) batched_examples_float = [ scipy.sparse.csr_matrix([[1.0, 1.0], [2.0, 2.0]]), scipy.sparse.csr_matrix([[2.0, 4.0], [6.0, 8.0]]), ] self.assertEqual( sys.getsizeof(batched_examples_float[0]) + sys.getsizeof(batched_examples_float[1]), inference_runner.get_num_bytes(batched_examples_float)) def test_pipeline_numpy(self): model = build_monkeypatched_xgboost_classifier() model_state = self.tmpdir + os.sep + 'model.json' model.save_model(model_state) with TestPipeline() as pipeline: examples = [ numpy.array([[1, 1], [2, 2]]), numpy.array([[2, 4], [6, 8]]), ] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | RunInference( XGBoostModelHandlerNumpy( model_class=xgboost.XGBClassifier, model_state=model_state)) expected = [ PredictionResult(numpy.array([[1, 1], [2, 2]]), 6), PredictionResult(numpy.array([[2, 4], [6, 8]]), 20) ] assert_that( actual, equal_to(expected, equals_fn=_compare_prediction_result)) def test_pipeline_numpy_sets_env_vars_correctly(self): model = build_monkeypatched_xgboost_classifier() model_state = self.tmpdir + os.sep + 'model.json' model.save_model(model_state) os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) with TestPipeline() as pipeline: examples = [ numpy.array([[1, 1], [2, 2]]), numpy.array([[2, 4], [6, 8]]), ] handler_with_vars = XGBoostModelHandlerNumpy( env_vars={'FOO': 'bar'}, model_class=xgboost.XGBClassifier, model_state=model_state) _ = ( pipeline | 'start' >> beam.Create(examples) | RunInference(handler_with_vars)) pipeline.run() self.assertTrue('FOO' in os.environ) self.assertTrue((os.environ['FOO']) == 'bar') def test_pipeline_pandas(self): model = build_monkeypatched_xgboost_classifier() model_state = self.tmpdir + os.sep + 'model.json' model.save_model(model_state) with TestPipeline() as pipeline: examples = [ pandas.DataFrame([[1, 1], [2, 2]]), pandas.DataFrame([[2, 4], [6, 8]]), ] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | RunInference( XGBoostModelHandlerPandas( model_class=xgboost.XGBClassifier, model_state=model_state)) expected = [ PredictionResult(pandas.DataFrame([[1, 1], [2, 2]]), 6), PredictionResult(pandas.DataFrame([[2, 4], [6, 8]]), 20) ] assert_that( actual, equal_to(expected, equals_fn=_compare_prediction_result)) pipeline.run() def test_pipeline_pandas_sets_env_vars_correctly(self): model = build_monkeypatched_xgboost_classifier() model_state = self.tmpdir + os.sep + 'model.json' model.save_model(model_state) os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) with TestPipeline() as pipeline: examples = [ pandas.DataFrame([[1, 1], [2, 2]]), pandas.DataFrame([[2, 4], [6, 8]]), ] handler_with_vars = XGBoostModelHandlerPandas( env_vars={'FOO': 'bar'}, model_class=xgboost.XGBClassifier, model_state=model_state) _ = ( pipeline | 'start' >> beam.Create(examples) | RunInference(handler_with_vars)) pipeline.run() self.assertTrue('FOO' in os.environ) self.assertTrue((os.environ['FOO']) == 'bar') def test_pipeline_datatable(self): model = build_monkeypatched_xgboost_classifier() model_state = self.tmpdir + os.sep + 'model.json' model.save_model(model_state) with TestPipeline() as pipeline: examples = [ datatable.Frame([[1, 1], [2, 2]]), datatable.Frame([[2, 4], [6, 8]]), ] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | RunInference( XGBoostModelHandlerDatatable( model_class=xgboost.XGBClassifier, model_state=model_state)) expected = [ PredictionResult(datatable.Frame([[1, 1], [2, 2]]), 6), PredictionResult(datatable.Frame([[2, 4], [6, 8]]), 20) ] assert_that( actual, equal_to(expected, equals_fn=_compare_prediction_result)) def test_pipeline_datatable_sets_env_vars_correctly(self): model = build_monkeypatched_xgboost_classifier() model_state = self.tmpdir + os.sep + 'model.json' model.save_model(model_state) os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) with TestPipeline() as pipeline: examples = [ datatable.Frame([[1, 1], [2, 2]]), datatable.Frame([[2, 4], [6, 8]]), ] handler_with_vars = XGBoostModelHandlerDatatable( env_vars={'FOO': 'bar'}, model_class=xgboost.XGBClassifier, model_state=model_state) _ = ( pipeline | 'start' >> beam.Create(examples) | RunInference(handler_with_vars)) pipeline.run() self.assertTrue('FOO' in os.environ) self.assertTrue((os.environ['FOO']) == 'bar') def test_pipeline_scipy(self): model = build_monkeypatched_xgboost_classifier() model_state = self.tmpdir + os.sep + 'model.json' model.save_model(model_state) with TestPipeline() as pipeline: examples = [ scipy.sparse.csr_matrix(numpy.array([[1, 1], [2, 2]])), scipy.sparse.csr_matrix(numpy.array([[2, 4], [6, 8]])), ] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | RunInference( XGBoostModelHandlerSciPy( model_class=xgboost.XGBClassifier, model_state=model_state)) expected = [ PredictionResult( scipy.sparse.csr_matrix(numpy.array([[1, 1], [2, 2]])), 6), PredictionResult( scipy.sparse.csr_matrix(numpy.array([[2, 4], [6, 8]])), 20) ] assert_that( actual, equal_to(expected, equals_fn=_compare_prediction_result)) def test_pipeline_scipy_sets_env_vars_correctly(self): model = build_monkeypatched_xgboost_classifier() model_state = self.tmpdir + os.sep + 'model.json' model.save_model(model_state) os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) with TestPipeline() as pipeline: examples = [ scipy.sparse.csr_matrix(numpy.array([[1, 1], [2, 2]])), scipy.sparse.csr_matrix(numpy.array([[2, 4], [6, 8]])), ] handler_with_vars = XGBoostModelHandlerSciPy( env_vars={'FOO': 'bar'}, model_class=xgboost.XGBClassifier, model_state=model_state) _ = ( pipeline | 'start' >> beam.Create(examples) | RunInference(handler_with_vars)) pipeline.run() self.assertTrue('FOO' in os.environ) self.assertTrue((os.environ['FOO']) == 'bar') def test_bad_model_file_raises(self): model_state = self.tmpdir + os.sep + 'bad_file_name.json' with self.assertRaises(RuntimeError): with TestPipeline() as pipeline: examples = [ datatable.Frame([[1, 1], [2, 2]]), datatable.Frame([[2, 4], [6, 8]]), ] pcoll = pipeline | 'start' >> beam.Create(examples) _ = pcoll | RunInference( XGBoostModelHandlerNumpy(xgboost.XGBClassifier, model_state)) pipeline.run() def test_bad_input_type_raises(self): model = build_monkeypatched_xgboost_classifier() model_state = self.tmpdir + os.sep + 'model.json' model.save_model(model_state) archived_model_state = self.tmpdir + os.sep + 'model.zip' zip_file = zipfile.ZipFile(archived_model_state, "w", zipfile.ZIP_DEFLATED) zip_file.write(model_state) zip_file.close() with self.assertRaises(xgboost.core.XGBoostError): model_handler = XGBoostModelHandlerNumpy( xgboost.XGBClassifier, model_state=archived_model_state) model_handler.load_model() def test_pipeline_scipy_with_keys(self): model = build_monkeypatched_xgboost_classifier() model_state = self.tmpdir + os.sep + 'model.json' model.save_model(model_state) with TestPipeline() as pipeline: examples = [ ('0', scipy.sparse.csr_matrix([[1, 1], [2, 2]])), ('1', scipy.sparse.csr_matrix([[2, 4], [6, 8]])), ] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | RunInference( KeyedModelHandler( XGBoostModelHandlerSciPy( model_class=xgboost.XGBClassifier, model_state=model_state))) expected = [ ('0', PredictionResult(scipy.sparse.csr_matrix([[1, 1], [2, 2]]), 6)), ( '1', PredictionResult(scipy.sparse.csr_matrix([[2, 4], [6, 8]]), 20)) ] assert_that( actual, equal_to(expected, equals_fn=_compare_keyed_prediction_result)) def test_pipeline_numpy_with_keys(self): model = build_monkeypatched_xgboost_classifier() model_state = self.tmpdir + os.sep + 'model.json' model.save_model(model_state) with TestPipeline() as pipeline: examples = [ ('0', numpy.array([[1, 1], [2, 2]])), ('1', numpy.array([[2, 4], [6, 8]])), ] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | RunInference( KeyedModelHandler( XGBoostModelHandlerNumpy( model_class=xgboost.XGBClassifier, model_state=model_state))) expected = [('0', PredictionResult(numpy.array([[1, 1], [2, 2]]), 6)), ('1', PredictionResult(numpy.array([[2, 4], [6, 8]]), 20))] assert_that( actual, equal_to(expected, equals_fn=_compare_keyed_prediction_result)) def test_pipeline_pandas_with_keys(self): model = build_monkeypatched_xgboost_classifier() model_state = self.tmpdir + os.sep + 'model.json' model.save_model(model_state) with TestPipeline() as pipeline: examples = [ ('0', pandas.DataFrame([[1, 1], [2, 2]])), ('1', pandas.DataFrame([[2, 4], [6, 8]])), ] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | RunInference( KeyedModelHandler( XGBoostModelHandlerPandas( model_class=xgboost.XGBClassifier, model_state=model_state))) expected = [ ('0', PredictionResult(pandas.DataFrame([[1, 1], [2, 2]]), 6)), ('1', PredictionResult(pandas.DataFrame([[2, 4], [6, 8]]), 20)) ] assert_that( actual, equal_to(expected, equals_fn=_compare_keyed_prediction_result)) def test_pipeline_datatable_with_keys(self): model = build_monkeypatched_xgboost_classifier() model_state = self.tmpdir + os.sep + 'model.json' model.save_model(model_state) with TestPipeline() as pipeline: examples = [ ('0', datatable.Frame([[1, 1], [2, 2]])), ('1', datatable.Frame([[2, 4], [6, 8]])), ] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | RunInference( KeyedModelHandler( XGBoostModelHandlerDatatable( model_class=xgboost.XGBClassifier, model_state=model_state))) expected = [ ('0', PredictionResult(datatable.Frame([[1, 1], [2, 2]]), 6)), ('1', PredictionResult(datatable.Frame([[2, 4], [6, 8]]), 20)) ] assert_that( actual, equal_to(expected, equals_fn=_compare_keyed_prediction_result)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.4984202 apache_beam-2.74.0/apache_beam/ml/rag/0000755000175100017510000000000015206113157017046 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/__init__.py0000644000175100017510000000177215206113032021156 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Apache Beam RAG (Retrieval Augmented Generation) components. This package provides components for building RAG pipelines in Apache Beam, including: - Chunking - Embedding generation - Vector storage - Vector search enrichment """ ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5004206 apache_beam-2.74.0/apache_beam/ml/rag/chunking/0000755000175100017510000000000015206113160020646 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/chunking/__init__.py0000644000175100017510000000162115206113032022755 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Chunking components for RAG pipelines. This module provides components for splitting text into chunks for RAG pipelines. """ ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/chunking/base.py0000644000175100017510000000666515206113032022145 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import abc import functools from collections.abc import Callable from typing import Any from typing import Optional import apache_beam as beam from apache_beam.ml.rag.types import Chunk from apache_beam.ml.transforms.base import MLTransformProvider ChunkIdFn = Callable[[Chunk], str] def _assign_chunk_id(chunk_id_fn: ChunkIdFn, chunk: Chunk): chunk.id = chunk_id_fn(chunk) return chunk class ChunkingTransformProvider(MLTransformProvider): def __init__(self, chunk_id_fn: Optional[ChunkIdFn] = None): """Base class for chunking transforms in RAG pipelines. ChunkingTransformProvider defines the interface for splitting documents into chunks for embedding and retrieval. Implementations should define how to split content while preserving metadata and managing chunk IDs. The transform flow: - Takes input documents with content and metadata - Splits content into chunks using implementation-specific logic - Preserves document metadata in resulting chunks - Optionally assigns unique IDs to chunks (configurable via chunk_id_fn Example usage: >>> class MyChunker(ChunkingTransformProvider): ... def get_splitter_transform(self): ... return beam.ParDo(MySplitterDoFn()) ... >>> chunker = MyChunker(chunk_id_fn=my_id_function) >>> >>> with beam.Pipeline() as p: ... chunks = ( ... p ... | beam.Create([{'text': 'document...', 'source': 'doc.txt'}]) ... | MLTransform(...).with_transform(chunker)) Args: chunk_id_fn: Optional function to generate chunk IDs. If not provided, random UUIDs will be used. Function should take a Chunk and return str. """ self.assign_chunk_id_fn = functools.partial( _assign_chunk_id, chunk_id_fn) if chunk_id_fn is not None else None @abc.abstractmethod def get_splitter_transform( self ) -> beam.PTransform[beam.PCollection[dict[str, Any]], beam.PCollection[Chunk]]: """Creates transforms that emits splits for given content.""" raise NotImplementedError( "Subclasses must implement get_splitter_transform") def get_ptransform_for_processing( self, **kwargs ) -> beam.PTransform[beam.PCollection[dict[str, Any]], beam.PCollection[Chunk]]: """Creates transform for processing documents into chunks.""" ptransform = ( "Split document" >> self.get_splitter_transform().with_output_types(Chunk)) if self.assign_chunk_id_fn: ptransform = ( ptransform | "Assign chunk id" >> beam.Map( self.assign_chunk_id_fn).with_output_types(Chunk)) return ptransform ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/chunking/base_test.py0000644000175100017510000001073115206113032023171 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Tests for apache_beam.ml.rag.chunking.base.""" import unittest from typing import Any from typing import Optional import pytest import apache_beam as beam from apache_beam.ml.rag.chunking.base import ChunkIdFn from apache_beam.ml.rag.chunking.base import ChunkingTransformProvider from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class WordSplitter(beam.DoFn): def process(self, element): words = element['text'].split() for i, word in enumerate(words): yield Chunk( content=Content(text=word), index=i, metadata={'source': element['source']}) class InvalidChunkingProvider(ChunkingTransformProvider): def __init__(self, chunk_id_fn: Optional[ChunkIdFn] = None): super().__init__(chunk_id_fn=chunk_id_fn) class MockChunkingProvider(ChunkingTransformProvider): def __init__(self, chunk_id_fn: Optional[ChunkIdFn] = None): super().__init__(chunk_id_fn=chunk_id_fn) def get_splitter_transform( self ) -> beam.PTransform[beam.PCollection[dict[str, Any]], beam.PCollection[Chunk]]: return beam.ParDo(WordSplitter()) def chunk_equals(expected, actual): """Custom equality function for Chunk objects.""" if not isinstance(expected, Chunk) or not isinstance(actual, Chunk): return False # Don't compare IDs since they're randomly generated return ( expected.index == actual.index and expected.content == actual.content and expected.metadata == actual.metadata) def id_equals(expected, actual): """Custom equality function for Chunk object id's.""" if not isinstance(expected, Chunk) or not isinstance(actual, Chunk): return False return (expected.id == actual.id) @pytest.mark.uses_transformers class ChunkingTransformProviderTest(unittest.TestCase): def setUp(self): self.test_doc = {'text': 'hello world test', 'source': 'test.txt'} def test_doesnt_override_get_text_splitter_transform(self): provider = InvalidChunkingProvider() with self.assertRaises(NotImplementedError): provider.get_splitter_transform() def test_chunking_transform(self): """Test the complete chunking transform.""" provider = MockChunkingProvider() with TestPipeline() as p: chunks = ( p | beam.Create([self.test_doc]) | provider.get_ptransform_for_processing()) expected = [ Chunk( content=Content(text="hello"), index=0, metadata={'source': 'test.txt'}), Chunk( content=Content(text="world"), index=1, metadata={'source': 'test.txt'}), Chunk( content=Content(text="test"), index=2, metadata={'source': 'test.txt'}) ] assert_that(chunks, equal_to(expected, equals_fn=chunk_equals)) def test_custom_chunk_id_fn(self): """Test the a custom chink id function.""" def source_index_id_fn(chunk: Chunk): return f"{chunk.metadata['source']}_{chunk.index}" provider = MockChunkingProvider(chunk_id_fn=source_index_id_fn) with TestPipeline() as p: chunks = ( p | beam.Create([self.test_doc]) | provider.get_ptransform_for_processing()) expected = [ Chunk(content=Content(text="hello"), id="test.txt_0"), Chunk(content=Content(text="world"), id="test.txt_1"), Chunk(content=Content(text="test"), id="test.txt_2") ] assert_that(chunks, equal_to(expected, equals_fn=id_equals)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/chunking/langchain.py0000644000175100017510000001056515206113032023151 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from typing import Any from typing import Optional import apache_beam as beam from apache_beam.ml.rag.chunking.base import ChunkIdFn from apache_beam.ml.rag.chunking.base import ChunkingTransformProvider from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content try: from langchain.text_splitter import TextSplitter except ImportError: TextSplitter = None class LangChainChunker(ChunkingTransformProvider): def __init__( self, text_splitter: TextSplitter, document_field: str, metadata_fields: list[str], chunk_id_fn: Optional[ChunkIdFn] = None): """A ChunkingTransformProvider that uses LangChain text splitters. This provider integrates LangChain's text splitting capabilities into Beam's MLTransform framework. It supports various text splitting strategies through LangChain's TextSplitter interface, including recursive character splitting and other methods. The provider: - Takes documents with text content and metadata - Splits text using configured LangChain splitter - Preserves document metadata in resulting chunks - Assigns unique IDs to chunks (configurable via chunk_id_fn) Example usage: ```python from langchain.text_splitter import RecursiveCharacterTextSplitter splitter = RecursiveCharacterTextSplitter( chunk_size=100, chunk_overlap=20 ) chunker = LangChainChunker(text_splitter=splitter) with beam.Pipeline() as p: chunks = ( p | beam.Create([{'text': 'long document...', 'source': 'doc.txt'}]) | MLTransform(...).with_transform(chunker)) ``` Args: text_splitter: A LangChain TextSplitter instance that defines how documents are split into chunks. metadata_fields: List of field names to copy from input documents to chunk metadata. These fields will be preserved in each chunk created from the document. chunk_id_fn: Optional function that take a Chunk and return str to generate chunk IDs. If not provided, random UUIDs will be used. """ if not TextSplitter: raise ImportError( "langchain is required to use LangChainChunker" "Please install it with using `pip install langchain`.") if not isinstance(text_splitter, TextSplitter): raise TypeError("text_splitter must be a LangChain TextSplitter") if not document_field: raise ValueError("document_field cannot be empty") super().__init__(chunk_id_fn) self.text_splitter = text_splitter self.document_field = document_field self.metadata_fields = metadata_fields def get_splitter_transform( self ) -> beam.PTransform[beam.PCollection[dict[str, Any]], beam.PCollection[Chunk]]: return "Langchain text split" >> beam.ParDo( _LangChainTextSplitter( text_splitter=self.text_splitter, document_field=self.document_field, metadata_fields=self.metadata_fields)) class _LangChainTextSplitter(beam.DoFn): def __init__( self, text_splitter: TextSplitter, document_field: str, metadata_fields: list[str]): self.text_splitter = text_splitter self.document_field = document_field self.metadata_fields = metadata_fields def process(self, element): text_chunks = self.text_splitter.split_text(element[self.document_field]) metadata = {field: element[field] for field in self.metadata_fields} for i, text_chunk in enumerate(text_chunks): yield Chunk(content=Content(text=text_chunk), index=i, metadata=metadata) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/chunking/langchain_test.py0000644000175100017510000001773715206113032024220 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Tests for apache_beam.ml.rag.chunking.langchain.""" import functools import unittest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import BeamAssertException from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.testing.util import is_not_empty # pylint: disable=ungrouped-imports try: from langchain.text_splitter import CharacterTextSplitter from langchain.text_splitter import RecursiveCharacterTextSplitter from apache_beam.ml.rag.chunking.langchain import LangChainChunker LANGCHAIN_AVAILABLE = True except ImportError: LANGCHAIN_AVAILABLE = False # Import optional dependencies try: from transformers import AutoTokenizer TRANSFORMERS_AVAILABLE = True except ImportError: TRANSFORMERS_AVAILABLE = False def assert_true(elements, assert_fn, error_message_fn): if not assert_fn(elements): raise BeamAssertException(error_message_fn(elements)) return True @unittest.skipIf(not LANGCHAIN_AVAILABLE, 'langchain is not installed.') class LangChainChunkingTest(unittest.TestCase): def setUp(self): self.simple_text = { 'content': 'This is a simple test document. It has multiple sentences. ' 'We will use it to test basic splitting.', 'source': 'simple.txt', 'language': 'en' } self.complex_text = { 'content': ( 'The patient arrived at 2 p.m. yesterday. ' 'Initial assessment was completed. ' 'Lab results showed normal ranges. ' 'Follow-up scheduled for next week.'), 'source': 'medical.txt', 'language': 'en' } def test_no_metadata_fields(self): """Test chunking with no metadata fields specified.""" splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=20) provider = LangChainChunker( document_field='content', metadata_fields=[], text_splitter=splitter) with TestPipeline() as p: chunks = ( p | beam.Create([self.simple_text]) | provider.get_ptransform_for_processing()) chunks_count = chunks | beam.combiners.Count.Globally() assert_that(chunks_count, is_not_empty(), 'Has chunks') assert_that( chunks, functools.partial( assert_true, assert_fn=lambda x: (all(c.metadata == {} for c in x)), error_message_fn=lambda x: f"Expected empty metadata, actual {x}") ) def test_multiple_metadata_fields(self): """Test chunking with multiple metadata fields.""" splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=20) provider = LangChainChunker( document_field='content', metadata_fields=['source', 'language'], text_splitter=splitter) expected_metadata = {'source': 'simple.txt', 'language': 'en'} with TestPipeline() as p: chunks = ( p | beam.Create([self.simple_text]) | provider.get_ptransform_for_processing()) chunks_count = chunks | beam.combiners.Count.Globally() assert_that(chunks_count, is_not_empty(), 'Has chunks') assert_that( chunks, functools.partial( assert_true, assert_fn=lambda x: all( c.metadata == expected_metadata for c in x), error_message_fn=lambda x: f"Expected metadata {expected_metadata}, actual {x}")) def test_recursive_splitter_no_overlap(self): """Test RecursiveCharacterTextSplitter with no overlap.""" splitter = RecursiveCharacterTextSplitter( chunk_size=30, chunk_overlap=0, separators=[".", " "]) provider = LangChainChunker( document_field='content', metadata_fields=['source'], text_splitter=splitter) with TestPipeline() as p: chunks = ( p | beam.Create([self.simple_text]) | provider.get_ptransform_for_processing()) chunks_count = chunks | beam.combiners.Count.Globally() assert_that(chunks_count, is_not_empty(), 'Has chunks') assert_that( chunks, functools.partial( assert_true, assert_fn=lambda x: all(len(c.content.text) <= 30 for c in x), error_message_fn=lambda x: f"Expected len(chunk) <= 30, \ actual {[len(c.content.text) for c in x]}")) @unittest.skipIf(not TRANSFORMERS_AVAILABLE, "transformers not available") def test_huggingface_tokenizer_splitter(self): """Test text splitter created from HuggingFace tokenizer.""" tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer( tokenizer, chunk_size=10, # tokens chunk_overlap=2 # tokens ) provider = LangChainChunker( document_field='content', metadata_fields=['source'], text_splitter=splitter) with TestPipeline() as p: chunks = ( p | beam.Create([self.simple_text]) | provider.get_ptransform_for_processing()) def check_token_lengths(chunks): for chunk in chunks: # Verify each chunk's token length is within limits num_tokens = len(tokenizer.tokenize(chunk.content.text)) if not num_tokens <= 10: raise BeamAssertException( f"Chunk has {num_tokens} tokens, expected <= 10") return True chunks_count = chunks | beam.combiners.Count.Globally() assert_that(chunks_count, is_not_empty(), 'Has chunks') assert_that(chunks, check_token_lengths) def test_invalid_document_field(self): """Test that using an invalid document field raises KeyError.""" splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=20) provider = LangChainChunker( document_field='nonexistent', metadata_fields={}, text_splitter=splitter) with self.assertRaisesRegex(Exception, "nonexistent"): with TestPipeline() as p: _ = ( p | beam.Create([self.simple_text]) | provider.get_ptransform_for_processing()) def test_empty_document_field(self): """Test that using an invalid document field raises KeyError.""" splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=20) with self.assertRaises(ValueError): _ = LangChainChunker( document_field='', metadata_fields={}, text_splitter=splitter) def test_invalid_text_splitter(self): """Test that using an invalid document field raises KeyError.""" with self.assertRaises(TypeError): _ = LangChainChunker( document_field='nonexistent', text_splitter="Not a text splitter!") def test_empty_text(self): """Test that empty text produces no chunks.""" empty_doc = {'content': '', 'source': 'empty.txt'} splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=20) provider = LangChainChunker( document_field='content', metadata_fields=['source'], text_splitter=splitter) with TestPipeline() as p: chunks = ( p | beam.Create([empty_doc]) | provider.get_ptransform_for_processing()) assert_that(chunks, equal_to([])) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5034208 apache_beam-2.74.0/apache_beam/ml/rag/embeddings/0000755000175100017510000000000015206113160021141 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/embeddings/__init__.py0000644000175100017510000000161415206113032023252 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Embedding components for RAG pipelines. This module provides components for generating embeddings in RAG pipelines. """ ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/embeddings/base.py0000644000175100017510000000453015206113032022425 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """RAG-specific embedding adapters. This module provides adapters for extracting content from EmbeddableItem instances and mapping embeddings back. Adapters are used by EmbeddingsManager to support various input types. """ from collections.abc import Sequence from apache_beam.ml.rag.types import EmbeddableItem from apache_beam.ml.rag.types import Embedding from apache_beam.ml.transforms.base import EmbeddingTypeAdapter def create_text_adapter( ) -> EmbeddingTypeAdapter[EmbeddableItem, EmbeddableItem]: """Creates adapter for text content embedding. Works with any EmbeddableItem that has text content (content.text). Extracts text for embedding and maps results back as Embedding objects. Returns: EmbeddingTypeAdapter configured for text embedding """ return EmbeddingTypeAdapter( input_fn=_extract_text, output_fn=_add_embedding_fn) # Backward compatibility alias. create_rag_adapter = create_text_adapter def _extract_text(items: Sequence[EmbeddableItem]) -> list[str]: """Extract text from items for embedding.""" texts = [] for item in items: if not item.content.text: raise ValueError( f"Expected text content in {type(item).__name__} {item.id}, " "got None") texts.append(item.content.text) return texts def _add_embedding_fn( items: Sequence[EmbeddableItem], embeddings: Sequence[list[float]]) -> list[EmbeddableItem]: """Create Embeddings from items and embedding vectors.""" for item, embedding in zip(items, embeddings): item.embedding = Embedding(dense_embedding=embedding) return list(items) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/embeddings/base_test.py0000644000175100017510000001077015206113032023467 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import unittest from apache_beam.ml.rag.embeddings.base import create_text_adapter from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import EmbeddableItem from apache_beam.ml.rag.types import Embedding class RAGBaseEmbeddingsTest(unittest.TestCase): def setUp(self): self.test_chunks = [ Chunk( content=Content(text="This is a test sentence."), id="1", metadata={ "source": "test.txt", "language": "en" }), Chunk( content=Content(text="Another example."), id="2", metadata={ "source": "test2.txt", "language": "en" }) ] def test_adapter_input_conversion(self): """Test the RAG type adapter converts correctly.""" adapter = create_text_adapter() # Test input conversion texts = adapter.input_fn(self.test_chunks) self.assertEqual(texts, ["This is a test sentence.", "Another example."]) def test_adapter_input_conversion_missing_text_content(self): """Test the RAG type adapter converts correctly.""" adapter = create_text_adapter() # Test input conversion with self.assertRaisesRegex(ValueError, "Expected text content"): adapter.input_fn([ Chunk( content=Content(), id="1", metadata={ "source": "test.txt", "language": "en" }) ]) def test_adapter_output_conversion(self): """Test the RAG type adapter converts correctly.""" # Test output conversion mock_embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] # Expected outputs expected = [ Chunk( id="1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), metadata={ 'source': 'test.txt', 'language': 'en' }, content=Content(text='This is a test sentence.')), Chunk( id="2", embedding=Embedding(dense_embedding=[0.4, 0.5, 0.6]), metadata={ 'source': 'test2.txt', 'language': 'en' }, content=Content(text='Another example.')), ] adapter = create_text_adapter() embeddings = adapter.output_fn(self.test_chunks, mock_embeddings) self.assertListEqual(embeddings, expected) class ImageEmbeddableItemTest(unittest.TestCase): def test_from_image_str(self): item = EmbeddableItem.from_image('gs://bucket/img.jpg', id='img1') self.assertEqual(item.content.image, 'gs://bucket/img.jpg') self.assertIsNone(item.content.text) self.assertEqual(item.id, 'img1') def test_from_image_bytes(self): data = b'\x89PNG\r\n' item = EmbeddableItem.from_image(data, id='img2') self.assertEqual(item.content.image, data) self.assertIsNone(item.content.text) def test_from_image_with_metadata(self): item = EmbeddableItem.from_image( 'path/to/img.jpg', id='img3', metadata={'source': 'camera'}) self.assertEqual(item.metadata, {'source': 'camera'}) self.assertEqual(item.content.image, 'path/to/img.jpg') class ContentStringTest(unittest.TestCase): def test_text_content(self): item = EmbeddableItem(content=Content(text="hello"), id="1") self.assertEqual(item.content_string, "hello") def test_image_uri_content(self): item = EmbeddableItem.from_image('gs://bucket/img.jpg', id='img1') self.assertEqual(item.content_string, 'gs://bucket/img.jpg') def test_image_bytes_raises(self): item = EmbeddableItem.from_image(b'\x89PNG\r\n', id='img2') with self.assertRaisesRegex(ValueError, "EmbeddableItem does not contain.*"): item.content_string if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/embeddings/huggingface.py0000644000175100017510000001717515206113032023773 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """RAG-specific embedding implementations using HuggingFace models.""" import io from collections.abc import Sequence from typing import Optional import apache_beam as beam from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference.base import RunInference from apache_beam.ml.rag.embeddings.base import _add_embedding_fn from apache_beam.ml.rag.embeddings.base import create_text_adapter from apache_beam.ml.rag.types import EmbeddableItem from apache_beam.ml.transforms.base import EmbeddingsManager from apache_beam.ml.transforms.base import EmbeddingTypeAdapter from apache_beam.ml.transforms.base import _ImageEmbeddingHandler from apache_beam.ml.transforms.base import _TextEmbeddingHandler from apache_beam.ml.transforms.embeddings.huggingface import _SentenceTransformerModelHandler try: from sentence_transformers import SentenceTransformer except ImportError: SentenceTransformer = None try: from PIL import Image as PILImage except ImportError: PILImage = None # type: ignore[assignment] class HuggingfaceTextEmbeddings(EmbeddingsManager): def __init__( self, model_name: str, *, max_seq_length: Optional[int] = None, **kwargs): """HuggingFace text embeddings for RAG pipelines. Args: model_name: Name of the sentence-transformers model to use. max_seq_length: Maximum sequence length for the model. **kwargs: Additional arguments passed to :class:`~apache_beam.ml.transforms.base.EmbeddingsManager`, including: - ``load_model_args``: dict passed to ``SentenceTransformer()`` constructor (e.g. ``device``, ``cache_folder``). - ``min_batch_size`` / ``max_batch_size``: Control batching for inference. - ``large_model``: If True, share the model across processes to reduce memory usage. - ``inference_args``: dict passed to ``model.encode()`` (e.g. ``normalize_embeddings``). """ if not SentenceTransformer: raise ImportError( "sentence-transformers is required to use " "HuggingfaceTextEmbeddings." "Please install it with using `pip install sentence-transformers`.") super().__init__(type_adapter=create_text_adapter(), **kwargs) self.model_name = model_name self.max_seq_length = max_seq_length self.model_class = SentenceTransformer def get_model_handler(self): """Returns model handler configured with RAG adapter.""" return _SentenceTransformerModelHandler( model_class=self.model_class, max_seq_length=self.max_seq_length, model_name=self.model_name, load_model_args=self.load_model_args, min_batch_size=self.min_batch_size, max_batch_size=self.max_batch_size, large_model=self.large_model) def get_ptransform_for_processing( self, **kwargs ) -> beam.PTransform[beam.PCollection[EmbeddableItem], beam.PCollection[EmbeddableItem]]: """Returns PTransform that uses the RAG adapter.""" return RunInference( model_handler=_TextEmbeddingHandler(self), inference_args=self.inference_args).with_output_types(EmbeddableItem) def _extract_images(items: Sequence[EmbeddableItem]) -> list: """Extract images from items and convert to PIL.Image objects. Supports raw bytes, local file paths, and remote URIs (e.g. gs://, s3://) via Beam's FileSystems. """ images = [] for item in items: if not item.content.image: raise ValueError( "Expected image content in " f"{type(item).__name__} {item.id}, " "got None") img_data = item.content.image if isinstance(img_data, bytes): img = PILImage.open(io.BytesIO(img_data)) else: with FileSystems.open(img_data, 'rb') as f: img = PILImage.open(f) img.load() images.append(img.convert('RGB')) return images def _create_hf_image_adapter( ) -> EmbeddingTypeAdapter[EmbeddableItem, EmbeddableItem]: """Creates adapter for HuggingFace image embedding. Extracts content.image from EmbeddableItems and converts to PIL.Image objects. Supports both raw bytes and file paths. Returns: EmbeddingTypeAdapter for HuggingFace image embedding. """ return EmbeddingTypeAdapter( input_fn=_extract_images, output_fn=_add_embedding_fn) class HuggingfaceImageEmbeddings(EmbeddingsManager): def __init__( self, model_name: str, *, max_seq_length: Optional[int] = None, **kwargs): """HuggingFace image embeddings for RAG pipelines. Generates embeddings for images using sentence-transformers models that support image input (e.g. clip-ViT-B-32). Args: model_name: Name of the sentence-transformers model. Must be an image-text model. See https://www.sbert.net/docs/sentence_transformer/pretrained_models.html#image-text-models max_seq_length: Maximum sequence length for the model if applicable. **kwargs: Additional arguments passed to :class:`~apache_beam.ml.transforms.base.EmbeddingsManager`, including: - ``load_model_args``: dict passed to ``SentenceTransformer()`` constructor (e.g. ``device``, ``cache_folder``, ``trust_remote_code``). - ``min_batch_size`` / ``max_batch_size``: Control batching for inference. - ``large_model``: If True, share the model across processes to reduce memory usage. - ``inference_args``: dict passed to ``model.encode()`` (e.g. ``normalize_embeddings``). """ if not SentenceTransformer: raise ImportError( "sentence-transformers is required to use " "HuggingfaceImageEmbeddings. " "Please install it with `pip install sentence-transformers`.") if not PILImage: raise ImportError( "Pillow is required to use HuggingfaceImageEmbeddings. " "Please install it with `pip install pillow`.") super().__init__(type_adapter=_create_hf_image_adapter(), **kwargs) self.model_name = model_name self.max_seq_length = max_seq_length self.model_class = SentenceTransformer def get_model_handler(self): """Returns model handler configured with RAG adapter.""" return _SentenceTransformerModelHandler( model_class=self.model_class, max_seq_length=self.max_seq_length, model_name=self.model_name, load_model_args=self.load_model_args, min_batch_size=self.min_batch_size, max_batch_size=self.max_batch_size, large_model=self.large_model) def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: """Returns PTransform for image embedding.""" return RunInference( model_handler=_ImageEmbeddingHandler(self), inference_args=self.inference_args).with_output_types(EmbeddableItem) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/embeddings/huggingface_test.py0000644000175100017510000002012615206113032025020 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Tests for apache_beam.ml.rag.embeddings.huggingface.""" import io import os import shutil import tempfile import unittest import pytest import apache_beam as beam from apache_beam.ml.rag.embeddings.huggingface import HuggingfaceImageEmbeddings from apache_beam.ml.rag.embeddings.huggingface import HuggingfaceTextEmbeddings from apache_beam.ml.rag.embeddings.huggingface import _create_hf_image_adapter from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import EmbeddableItem from apache_beam.ml.rag.types import Embedding from apache_beam.ml.transforms.base import MLTransform from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # pylint: disable=unused-import try: from sentence_transformers import SentenceTransformer SENTENCE_TRANSFORMERS_AVAILABLE = True except ImportError: SENTENCE_TRANSFORMERS_AVAILABLE = False try: from PIL import Image PIL_AVAILABLE = True except ImportError: PIL_AVAILABLE = False def chunk_approximately_equals(expected, actual): """Compare embeddings allowing for numerical differences.""" if not isinstance(expected, Chunk) or not isinstance(actual, Chunk): return False return ( expected.id == actual.id and expected.metadata == actual.metadata and expected.content == actual.content and len(expected.embedding.dense_embedding) == len( actual.embedding.dense_embedding) and all(isinstance(x, float) for x in actual.embedding.dense_embedding)) @pytest.mark.uses_transformers @unittest.skipIf( not SENTENCE_TRANSFORMERS_AVAILABLE, "sentence-transformers not available") class HuggingfaceTextEmbeddingsTest(unittest.TestCase): def setUp(self): self.artifact_location = tempfile.mkdtemp(prefix='sentence_transformers_') self.test_chunks = [ Chunk( content=Content(text="This is a test sentence."), id="1", metadata={ "source": "test.txt", "language": "en" }), Chunk( content=Content(text="Another example."), id="2", metadata={ "source": "test.txt", "language": "en" }) ] def tearDown(self) -> None: shutil.rmtree(self.artifact_location) def test_embedding_pipeline(self): expected = [ Chunk( id="1", embedding=Embedding(dense_embedding=[0.0] * 384), metadata={ "source": "test.txt", "language": "en" }, content=Content(text="This is a test sentence.")), Chunk( id="2", embedding=Embedding(dense_embedding=[0.0] * 384), metadata={ "source": "test.txt", "language": "en" }, content=Content(text="Another example.")) ] embedder = HuggingfaceTextEmbeddings( model_name="sentence-transformers/all-MiniLM-L6-v2") with TestPipeline() as p: embeddings = ( p | beam.Create(self.test_chunks) | MLTransform(write_artifact_location=self.artifact_location). with_transform(embedder)) assert_that( embeddings, equal_to(expected, equals_fn=chunk_approximately_equals)) @pytest.mark.uses_transformers @unittest.skipIf( not SENTENCE_TRANSFORMERS_AVAILABLE or not PIL_AVAILABLE, "sentence-transformers or PIL not available") class HuggingfaceImageAdapterTest(unittest.TestCase): def test_image_adapter_missing_content(self): adapter = _create_hf_image_adapter() item = EmbeddableItem(content=Content(), id="no_img") with self.assertRaisesRegex(ValueError, "Expected image content"): adapter.input_fn([item]) def test_image_adapter_bytes_input(self): adapter = _create_hf_image_adapter() png_bytes = _create_png_bytes() item = EmbeddableItem.from_image(png_bytes, id='img1') images = adapter.input_fn([item]) self.assertEqual(len(images), 1) def test_image_adapter_path_input(self): adapter = _create_hf_image_adapter() png_bytes = _create_png_bytes() tmpdir = tempfile.mkdtemp() try: img_path = os.path.join(tmpdir, 'test.png') with open(img_path, 'wb') as f: f.write(png_bytes) item = EmbeddableItem.from_image(img_path, id='img2') images = adapter.input_fn([item]) self.assertEqual(len(images), 1) finally: shutil.rmtree(tmpdir) def test_image_adapter_output(self): adapter = _create_hf_image_adapter() items = [EmbeddableItem.from_image(b'\x89PNG', id='img1')] mock_embeddings = [[0.1, 0.2, 0.3]] result = adapter.output_fn(items, mock_embeddings) self.assertEqual(len(result), 1) self.assertEqual(result[0].embedding.dense_embedding, [0.1, 0.2, 0.3]) @pytest.mark.uses_transformers @unittest.skipIf( not SENTENCE_TRANSFORMERS_AVAILABLE or not PIL_AVAILABLE, "sentence-transformers or PIL not available") class HuggingfaceImageEmbeddingsTest(unittest.TestCase): def setUp(self): self.artifact_location = tempfile.mkdtemp(prefix='hf_image_') def tearDown(self) -> None: shutil.rmtree(self.artifact_location) def test_image_embedding_pipeline(self): png_bytes = _create_png_bytes() test_items = [ EmbeddableItem.from_image( png_bytes, id="img1", metadata={"source": "test"}), ] # clip-ViT-B-32 produces 512-dimensional embeddings expected = [ EmbeddableItem( id="img1", embedding=Embedding(dense_embedding=[0.0] * 512), metadata={"source": "test"}, content=Content(image=png_bytes)), ] embedder = HuggingfaceImageEmbeddings(model_name="clip-ViT-B-32") with TestPipeline() as p: embeddings = ( p | beam.Create(test_items) | MLTransform(write_artifact_location=( self.artifact_location)).with_transform(embedder)) assert_that( embeddings, equal_to(expected, equals_fn=chunk_approximately_equals)) def test_image_embedding_pipeline_from_path(self): png_bytes = _create_png_bytes() img_path = os.path.join(self.artifact_location, 'test_img.png') with open(img_path, 'wb') as f: f.write(png_bytes) test_items = [ EmbeddableItem.from_image( img_path, id="img1", metadata={"source": "test"}), ] expected = [ EmbeddableItem( id="img1", embedding=Embedding(dense_embedding=[0.0] * 512), metadata={"source": "test"}, content=Content(image=img_path)), ] artifact_location = tempfile.mkdtemp(prefix='hf_image_path_') embedder = HuggingfaceImageEmbeddings(model_name="clip-ViT-B-32") try: with TestPipeline() as p: embeddings = ( p | beam.Create(test_items) | MLTransform(write_artifact_location=( artifact_location)).with_transform(embedder)) assert_that( embeddings, equal_to(expected, equals_fn=chunk_approximately_equals)) finally: shutil.rmtree(artifact_location) def _create_png_bytes(): """Create a small valid RGB PNG image.""" from PIL import Image img = Image.new('RGB', (10, 10), color=(128, 64, 32)) buf = io.BytesIO() img.save(buf, format='PNG') return buf.getvalue() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/embeddings/vertex_ai.py0000644000175100017510000001566615206113032023515 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # Vertex AI Python SDK is required for this module. # Follow https://cloud.google.com/vertex-ai/docs/python-sdk/use-vertex-ai-python-sdk # pylint: disable=line-too-long # to install Vertex AI Python SDK. """RAG-specific embedding implementations using Vertex AI models.""" from collections.abc import Sequence from typing import Optional from google.auth.credentials import Credentials import apache_beam as beam from apache_beam.ml.inference.base import RunInference from apache_beam.ml.rag.embeddings.base import _add_embedding_fn from apache_beam.ml.rag.embeddings.base import create_text_adapter from apache_beam.ml.rag.types import EmbeddableItem from apache_beam.ml.transforms.base import EmbeddingsManager from apache_beam.ml.transforms.base import EmbeddingTypeAdapter from apache_beam.ml.transforms.base import _ImageEmbeddingHandler from apache_beam.ml.transforms.base import _TextEmbeddingHandler from apache_beam.ml.transforms.embeddings.vertex_ai import DEFAULT_TASK_TYPE from apache_beam.ml.transforms.embeddings.vertex_ai import _VertexAIImageEmbeddingHandler from apache_beam.ml.transforms.embeddings.vertex_ai import _VertexAITextEmbeddingHandler try: import vertexai from vertexai.vision_models import Image except ImportError: vertexai = None # type: ignore[assignment] class VertexAITextEmbeddings(EmbeddingsManager): def __init__( self, model_name: str, *, title: Optional[str] = None, task_type: str = DEFAULT_TASK_TYPE, project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[Credentials] = None, **kwargs): """Utilizes Vertex AI text embeddings for semantic search and RAG pipelines. Args: model_name: Name of the Vertex AI text embedding model title: Optional title for the text content task_type: Task type for embeddings (default: RETRIEVAL_DOCUMENT) project: GCP project ID location: GCP location credentials: Optional GCP credentials **kwargs: Additional arguments passed to :class:`~apache_beam.ml.transforms.base.EmbeddingsManager`. """ if not vertexai: raise ImportError( "vertexai is required to use VertexAITextEmbeddings. " "Please install it with `pip install google-cloud-aiplatform`") super().__init__(type_adapter=create_text_adapter(), **kwargs) self.model_name = model_name self.title = title self.task_type = task_type self.project = project self.location = location self.credentials = credentials def get_model_handler(self): """Returns model handler configured with RAG adapter.""" return _VertexAITextEmbeddingHandler( model_name=self.model_name, title=self.title, task_type=self.task_type, project=self.project, location=self.location, credentials=self.credentials, ) def get_ptransform_for_processing( self, **kwargs ) -> beam.PTransform[beam.PCollection[EmbeddableItem], beam.PCollection[EmbeddableItem]]: """Returns PTransform that uses the RAG adapter.""" return RunInference( model_handler=_TextEmbeddingHandler(self), inference_args=self.inference_args).with_output_types(EmbeddableItem) def _extract_images(items: Sequence[EmbeddableItem]) -> list: """Extract images from items and convert to vertexai Image objects.""" images = [] for item in items: if not item.content.image: raise ValueError( "Expected image content in " f"{type(item).__name__} {item.id}, " "got None") img_data = item.content.image if isinstance(img_data, bytes): images.append(Image(image_bytes=img_data)) else: images.append(Image.load_from_file(img_data)) return images def _create_image_adapter( ) -> EmbeddingTypeAdapter[EmbeddableItem, EmbeddableItem]: """Creates adapter for Vertex AI image embedding. Extracts content.image from EmbeddableItems and converts to vertexai.vision_models.Image objects. Supports both raw bytes and file paths/URIs. Returns: EmbeddingTypeAdapter for Vertex AI image embedding. """ return EmbeddingTypeAdapter( input_fn=_extract_images, output_fn=_add_embedding_fn) class VertexAIImageEmbeddings(EmbeddingsManager): def __init__( self, model_name: str, *, dimension: Optional[int] = None, project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[Credentials] = None, **kwargs): """Vertex AI image embeddings for RAG pipelines. Generates embeddings for images using Vertex AI multimodal embedding models. Args: model_name: Name of the Vertex AI model. dimension: Embedding dimension. Must be one of 128, 256, 512, or 1408. project: GCP project ID. location: GCP location. credentials: Optional GCP credentials. **kwargs: Additional arguments passed to :class:`~apache_beam.ml.transforms.base.EmbeddingsManager`. """ if not vertexai: raise ImportError( "vertexai is required to use " "VertexAIImageEmbeddings. " "Please install it with " "`pip install google-cloud-aiplatform`") if dimension is not None and dimension not in (128, 256, 512, 1408): raise ValueError("dimension must be one of " "128, 256, 512, or 1408") super().__init__(type_adapter=_create_image_adapter(), **kwargs) self.model_name = model_name self.dimension = dimension self.project = project self.location = location self.credentials = credentials def get_model_handler(self): """Returns model handler for image embedding.""" return _VertexAIImageEmbeddingHandler( model_name=self.model_name, dimension=self.dimension, project=self.project, location=self.location, credentials=self.credentials, ) def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: """Returns PTransform for image embedding.""" return RunInference( model_handler=_ImageEmbeddingHandler(self), inference_args=self.inference_args).with_output_types(EmbeddableItem) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/embeddings/vertex_ai_test.py0000644000175100017510000001752615206113032024551 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Tests for apache_beam.ml.rag.embeddings.vertex_ai.""" import shutil import struct import tempfile import unittest import zlib import pytest import apache_beam as beam from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import EmbeddableItem from apache_beam.ml.rag.types import Embedding from apache_beam.ml.transforms.base import MLTransform from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to pytest.importorskip("vertexai", reason="Vertex AI dependencies not available") # pylint: disable=ungrouped-imports try: import vertexai # pylint: disable=unused-import from apache_beam.ml.rag.embeddings.vertex_ai import VertexAIImageEmbeddings from apache_beam.ml.rag.embeddings.vertex_ai import VertexAITextEmbeddings from apache_beam.ml.rag.embeddings.vertex_ai import _create_image_adapter except ImportError: VertexAIImageEmbeddings = None # type: ignore VertexAITextEmbeddings = None # type: ignore _create_image_adapter = None # type: ignore def chunk_approximately_equals(expected, actual): """Compare embeddings allowing for numerical differences.""" if not isinstance(expected, Chunk) or not isinstance(actual, Chunk): return False return ( expected.id == actual.id and expected.metadata == actual.metadata and expected.content == actual.content and len(expected.embedding.dense_embedding) == len( actual.embedding.dense_embedding) and all(isinstance(x, float) for x in actual.embedding.dense_embedding)) @pytest.mark.vertex_ai_postcommit class VertexAITextEmbeddingsTest(unittest.TestCase): def setUp(self): self.artifact_location = tempfile.mkdtemp(prefix='vertex_ai_') self.test_chunks = [ Chunk( content=Content(text="This is a test sentence."), id="1", metadata={ "source": "test.txt", "language": "en" }), Chunk( content=Content(text="Another example."), id="2", metadata={ "source": "test.txt", "language": "en" }) ] def tearDown(self) -> None: shutil.rmtree(self.artifact_location) def test_embedding_pipeline(self): # gecko@002 produces 768-dimensional embeddings expected = [ Chunk( id="1", embedding=Embedding(dense_embedding=[0.0] * 768), metadata={ "source": "test.txt", "language": "en" }, content=Content(text="This is a test sentence.")), Chunk( id="2", embedding=Embedding(dense_embedding=[0.0] * 768), metadata={ "source": "test.txt", "language": "en" }, content=Content(text="Another example.")) ] embedder = VertexAITextEmbeddings(model_name="text-embedding-005") with TestPipeline() as p: embeddings = ( p | beam.Create(self.test_chunks) | MLTransform(write_artifact_location=self.artifact_location). with_transform(embedder)) assert_that( embeddings, equal_to(expected, equals_fn=chunk_approximately_equals)) class VertexAIImageAdapterTest(unittest.TestCase): def test_image_adapter_missing_content(self): adapter = _create_image_adapter() item = EmbeddableItem(content=Content(), id="no_img") with self.assertRaisesRegex(ValueError, "Expected image content"): adapter.input_fn([item]) def test_image_adapter_str_input(self): adapter = _create_image_adapter() item = EmbeddableItem.from_image('gs://bucket/img.jpg', id='img1') images = adapter.input_fn([item]) self.assertEqual(len(images), 1) def test_image_adapter_bytes_input(self): adapter = _create_image_adapter() png_bytes = _create_png_bytes() item = EmbeddableItem.from_image(png_bytes, id='img2') images = adapter.input_fn([item]) self.assertEqual(len(images), 1) def test_image_adapter_output(self): adapter = _create_image_adapter() items = [ EmbeddableItem.from_image('gs://bucket/img.jpg', id='img1'), ] mock_embeddings = [[0.1, 0.2, 0.3]] result = adapter.output_fn(items, mock_embeddings) self.assertEqual(len(result), 1) self.assertEqual(result[0].embedding.dense_embedding, [0.1, 0.2, 0.3]) @pytest.mark.vertex_ai_postcommit class VertexAIImageEmbeddingsTest(unittest.TestCase): def setUp(self): self.artifact_location = tempfile.mkdtemp(prefix='vertex_ai_img_') def tearDown(self) -> None: shutil.rmtree(self.artifact_location) def test_image_embedding_pipeline(self): png_bytes = _create_png_bytes() test_items = [ EmbeddableItem.from_image( png_bytes, id="img1", metadata={"source": "test"}), ] # Default dimension for multimodal is 1408 expected = [ EmbeddableItem( id="img1", embedding=Embedding(dense_embedding=[0.0] * 1408), metadata={"source": "test"}, content=Content(image=png_bytes)), ] embedder = VertexAIImageEmbeddings(model_name="multimodalembedding@001") with TestPipeline() as p: embeddings = ( p | beam.Create(test_items) | MLTransform(write_artifact_location=( self.artifact_location)).with_transform(embedder)) assert_that( embeddings, equal_to(expected, equals_fn=chunk_approximately_equals)) def test_image_embedding_pipeline_from_path(self): png_bytes = _create_png_bytes() test_items = [ EmbeddableItem.from_image( png_bytes, id="img1", metadata={"source": "test"}), ] expected = [ EmbeddableItem( id="img1", embedding=Embedding(dense_embedding=[0.0] * 1408), metadata={"source": "test"}, content=Content(image=png_bytes)), ] artifact_location = tempfile.mkdtemp(prefix='vertex_ai_img_path_') embedder = VertexAIImageEmbeddings(model_name="multimodalembedding@001") try: with TestPipeline() as p: embeddings = ( p | beam.Create(test_items) | MLTransform(write_artifact_location=( artifact_location)).with_transform(embedder)) assert_that( embeddings, equal_to(expected, equals_fn=chunk_approximately_equals)) finally: shutil.rmtree(artifact_location) def _create_png_bytes(): """Create a minimal valid 1x1 RGB PNG image (no PIL dependency).""" png_sig = b'\x89PNG\r\n\x1a\n' ihdr_data = struct.pack('>IIBBBBB', 1, 1, 8, 2, 0, 0, 0) ihdr_crc = zlib.crc32(b'IHDR' + ihdr_data) ihdr = struct.pack('>I', 13) + b'IHDR' + ihdr_data + struct.pack( '>I', ihdr_crc) raw = zlib.compress(b'\x00\x00\x00\x00') idat_crc = zlib.crc32(b'IDAT' + raw) idat = struct.pack('>I', len(raw)) + b'IDAT' + raw + struct.pack( '>I', idat_crc) iend_crc = zlib.crc32(b'IEND') iend = struct.pack('>I', 0) + b'IEND' + struct.pack('>I', iend_crc) return png_sig + ihdr + idat + iend if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.506421 apache_beam-2.74.0/apache_beam/ml/rag/enrichment/0000755000175100017510000000000015206113160021174 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/enrichment/__init__.py0000644000175100017510000000162015206113032023302 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Enrichment components for RAG pipelines. This module provides components for vector search enrichment in RAG pipelines. """ ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/enrichment/bigquery_vector_search.py0000644000175100017510000003565015206113032026313 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging from collections import defaultdict from dataclasses import dataclass from typing import Any from typing import Callable from typing import Optional from typing import Union from google.cloud import bigquery from apache_beam.ml.rag.types import EmbeddableItem from apache_beam.ml.rag.types import Embedding from apache_beam.transforms.enrichment import EnrichmentSourceHandler _LOGGER = logging.getLogger(__name__) @dataclass class BigQueryVectorSearchParameters: """Parameters for configuring BigQuery vector similarity search. This class is used by BigQueryVectorSearchEnrichmentHandler to perform vector similarity search using BigQuery's VECTOR_SEARCH function. It processes :class:`~apache_beam.ml.rag.types.EmbeddableItem` objects that contain :class:`~apache_beam.ml.rag.types.Embedding` and returns similar vectors from a BigQuery table. BigQueryVectorSearchEnrichmentHandler is used with :class:`~apache_beam.transforms.enrichment.Enrichment` transform to enrich EmbeddableItems with similar content from a vector database. For example: >>> # Create search parameters >>> params = BigQueryVectorSearchParameters( ... table_name='project.dataset.embeddings', ... embedding_column='embedding', ... columns=['content'], ... neighbor_count=5 ... ) >>> # Use in pipeline >>> enriched = ( ... embeddable_items ... | "Generate Embeddings" >> MLTransform(...) ... | "Find Similar" >> Enrichment( ... BigQueryVectorSearchEnrichmentHandler( ... project='my-project', ... vector_search_parameters=params ... ) ... ) ... ) BigQueryVectorSearchParameters encapsulates the configuration needed to perform vector similarity search using BigQuery's VECTOR_SEARCH function. It handles formatting the query with proper embedding vectors and metadata restrictions. Example with flattened metadata column: Table schema:: embedding: ARRAY # Vector embedding content: STRING # Document content language: STRING # Direct metadata column Code:: >>> params = BigQueryVectorSearchParameters( ... table_name='project.dataset.embeddings', ... embedding_column='embedding', ... columns=['content', 'language'], ... neighbor_count=5, ... # For column 'language', value comes from ... # embeddable_item.metadata['language'] ... metadata_restriction_template="language = '{language}'" ... ) >>> # When processing a embeddable_item with >>> # metadata={'language': 'en'}, generates: WHERE language = 'en' Example with nested repeated metadata: Table schema:: embedding: ARRAY # Vector embedding content: STRING # Document content metadata: ARRAY # Nested repeated metadata key: STRING, value: STRING >> Code:: >>> params = BigQueryVectorSearchParameters( ... table_name='project.dataset.embeddings', ... embedding_column='embedding', ... columns=['content', 'metadata'], ... neighbor_count=5, ... # check_metadata(field_name, key_to_search, ... # value_from_embeddable_item) ... metadata_restriction_template=( ... "check_metadata(metadata, 'language', '{language}')" ... ) ... ) >>> # When processing a embeddable_item with >>> # metadata={'language': 'en'}, >>> # generates: WHERE check_metadata(metadata, 'language', 'en') >>> # Searches for {key: 'language', value: 'en'} in metadata array Args: project: GCP project ID containing the BigQuery dataset table_name: Fully qualified BigQuery table name containing vectors. embedding_column: Column name containing the embedding vectors. columns: List of columns to retrieve from matched vectors. neighbor_count: Number of similar vectors to return (top-k). metadata_restriction_template: Template string or callable for filtering vectors. Template string supports two formats: 1. For flattened metadata columns: ``column_name = '{metadata_key}'`` where column_name is the BigQuery column and metadata_key is used to get the value from embeddable_item.metadata[metadata_key]. 2. For nested repeated metadata (ARRAY>): ``check_metadata(field_name, 'key_to_match', '{metadata_key}')`` where field_name is the ARRAY column in BigQuery, key_to_match is the literal key to search for in the array, and metadata_key is used to get value from embeddable_item.metadata[metadata_key]. Multiple conditions can be combined using AND/OR operators. For example:: >>> # Combine metadata check with column filter >>> template = ( ... "check_metadata(metadata, 'language', '{language}') " ... "AND source = '{source}'" ... ) >>> # When embeddable_item.metadata = {'language': 'en', >>> # 'source': 'web'} >>> # Generates: WHERE >>> # check_metadata(metadata, 'language', 'en') >>> # AND source = 'web' distance_type: Optional distance metric to use. Supported values: COSINE (default), EUCLIDEAN, or DOT_PRODUCT. options: Optional dictionary of additional VECTOR_SEARCH options. include_distance: Reurns the vector search similarity score if True. """ project: str table_name: str embedding_column: str columns: list[str] neighbor_count: int metadata_restriction_template: Optional[Union[str, Callable[[EmbeddableItem], str]]] = None distance_type: Optional[str] = None options: Optional[dict[str, Any]] = None include_distance: bool = False def _format_restrict(self, item: EmbeddableItem) -> str: assert self.metadata_restriction_template is not None, ( "metadata_restriction_template cannot be None when formatting. " "This indicates a logical error in the code." ) if callable(self.metadata_restriction_template): return self.metadata_restriction_template(item) return self.metadata_restriction_template.format(**item.metadata) def format_query(self, items: list[EmbeddableItem]) -> str: """Format the vector search query template.""" base_columns_str = ", ".join(f"base.{col}" for col in self.columns) columns_str = ", ".join(self.columns) distance_clause = ( f", distance_type => '{self.distance_type}'" if self.distance_type else "") options_clause = (f", options => {self.options}" if self.options else "") # Create metadata check function only if needed metadata_fn = """ CREATE TEMP FUNCTION check_metadata( metadata ARRAY>, search_key STRING, search_value STRING ) AS (( SELECT COUNT(*) > 0 FROM UNNEST(metadata) WHERE key = search_key AND value = search_value )); """ if self.metadata_restriction_template else "" # Group items by their metadata conditions condition_groups = defaultdict(list) if self.metadata_restriction_template: for item in items: condition = self._format_restrict(item) condition_groups[condition].append(item) else: # No metadata filtering - all items in one group condition_groups[""] = items # Generate VECTOR_SEARCH subqueries for each condition group vector_searches = [] for condition, group_items in condition_groups.items(): # Create embeddings subquery for this group embedding_unions = [] for item in group_items: if not item.dense_embedding: raise ValueError(f"Item {item.id} missing embedding") embedding_str = ( f"SELECT '{item.id}' as id, " f"{[float(x) for x in item.dense_embedding]} " f"as embedding") embedding_unions.append(embedding_str) group_embeddings = " UNION ALL ".join(embedding_unions) where_clause = f"WHERE {condition}" if condition else "" # Create VECTOR_SEARCH for this condition group vector_search = f""" SELECT query.id, ARRAY_AGG( STRUCT({"distance, " if self.include_distance else ""}\ {base_columns_str}) ) as chunks FROM VECTOR_SEARCH( (SELECT {columns_str}, {self.embedding_column} FROM `{self.table_name}` {where_clause}), '{self.embedding_column}', (SELECT * FROM ({group_embeddings})), top_k => {self.neighbor_count} {distance_clause} {options_clause} ) GROUP BY query.id """ vector_searches.append(vector_search) # Combine all vector searches combined_searches = " UNION ALL ".join(vector_searches) return f""" {metadata_fn} {combined_searches} """ class BigQueryVectorSearchEnrichmentHandler( EnrichmentSourceHandler[Union[EmbeddableItem, list[EmbeddableItem]], list[tuple[EmbeddableItem, dict[str, Any]]]]): """Enrichment handler that performs vector similarity search using BigQuery. This handler enriches EmbeddableItems by finding similar vectors in a BigQuery table using the VECTOR_SEARCH function. It supports batching requests for efficiency and preserves the original metadata while adding the search results. Example: >>> from apache_beam.ml.rag.types import EmbeddableItem >>> from apache_beam.ml.rag.types import Content, Embedding >>> >>> # Configure vector search >>> params = BigQueryVectorSearchParameters( ... table_name='project.dataset.embeddings', ... embedding_column='embedding', ... columns=['content', 'metadata'], ... neighbor_count=2, ... metadata_restriction_template="language = '{language}'" ... ) >>> >>> # Create handler >>> handler = BigQueryVectorSearchEnrichmentHandler( ... project='my-project', ... vector_search_parameters=params, ... min_batch_size=100, ... max_batch_size=1000 ... ) >>> >>> # Use in pipeline >>> with beam.Pipeline() as p: ... enriched = ( ... p ... | beam.Create([ ... EmbeddableItem( ... id='query1', ... embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), ... content=Content(text='test query'), ... metadata={'language': 'en'} ... ) ... ]) ... | Enrichment(handler) ... ) Args: vector_search_parameters: Configuration for the vector search query min_batch_size: Minimum number of items to process in one batch max_batch_size: Maximum number of items to process in one batch log_query: Debug option to log the BigQuery query **kwargs: Additional arguments passed to bigquery.Client The handler will: 1. Batch incoming embeddable_items according to batch size parameters 2. Format and execute vector search query for each batch 3. Join results back to original embeddable_items 4. Return tuples of (original_embeddable_item, search_results) """ def __init__( self, vector_search_parameters: BigQueryVectorSearchParameters, *, min_batch_size: int = 1, max_batch_size: int = 1000, log_query=False, **kwargs): self.project = vector_search_parameters.project self.vector_search_parameters = vector_search_parameters self.kwargs = kwargs self._batching_kwargs = { 'min_batch_size': min_batch_size, 'max_batch_size': max_batch_size } self.join_fn = join_fn self.use_custom_types = True self.log_query = log_query def __enter__(self): self.client = bigquery.Client(project=self.project, **self.kwargs) def __call__( self, request: Union[EmbeddableItem, list[EmbeddableItem]], *args, **kwargs) -> list[tuple[EmbeddableItem, dict[str, Any]]]: """Process request(s) using BigQuery vector search. Args: request: Single EmbeddableItem with embedding or list of EmbeddableItems with embeddings to process Returns: EmbeddableItem(s) where metadata['enrichment_output'] contains the data retrieved via BigQuery VECTOR_SEARCH. """ # Convert single request to list for uniform processing requests = request if isinstance(request, list) else [request] # Generate and execute query query = self.vector_search_parameters.format_query(requests) if self.log_query: _LOGGER.info("Executing query %s", query) query_job = self.client.query(query) results = query_job.result() # Create results dict with empty embeddable_items list as default results_by_id = {} for result_row in results: result_dict = dict(result_row.items()) results_by_id[result_row.id] = result_dict # Return all embeddable_items in original order, with empty results if # no matches response = [] for embeddable_item in requests: result_dict = results_by_id.get(embeddable_item.id, {}) response.append((embeddable_item, result_dict)) return response def __exit__(self, exc_type, exc_val, exc_tb): self.client.close() def batch_elements_kwargs(self) -> dict[str, int]: """Returns kwargs for beam.BatchElements.""" return self._batching_kwargs def join_fn(left: Embedding, right: dict[str, Any]) -> Embedding: left.metadata['enrichment_data'] = right return left ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/enrichment/bigquery_vector_search_it_test.py0000644000175100017510000010002115206113032030027 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import secrets import time import unittest import apache_beam as beam from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import Embedding from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # pylint: disable=ungrouped-imports try: from apache_beam.ml.rag.enrichment.bigquery_vector_search import BigQueryVectorSearchEnrichmentHandler from apache_beam.ml.rag.enrichment.bigquery_vector_search import BigQueryVectorSearchParameters from apache_beam.transforms.enrichment import Enrichment except ImportError: raise unittest.SkipTest('BigQuery dependencies not installed') _LOGGER = logging.getLogger(__name__) DEFAULT_TABLE_NAME = "default_table" DEFAULT_TABLE_SCHEMA = [('id', 'STRING'), ('content', 'STRING'), ('domain', 'STRING'), ('embedding', 'FLOAT64', 'REPEATED'), ( 'metadata', 'RECORD', 'REPEATED', [('key', 'STRING'), ('value', 'STRING')])] PRICE_TABLE_SCHEMA = [ ('id', 'STRING'), ('content', 'STRING'), ('price', 'FLOAT64'), # Unnested price column ('embedding', 'FLOAT64', 'REPEATED') ] # Define test data as class attributes DEFAULT_TABLE_DATA = [{ "id": "doc1", "content": "This is a test document", "domain": "medical", "embedding": [0.1, 0.2, 0.3], "metadata": [{ "key": "language", "value": "en" }] }, { "id": "doc2", "content": "Another test document", "domain": "legal", "embedding": [0.2, 0.3, 0.4], "metadata": [{ "key": "language", "value": "en" }] }, { "id": "doc3", "content": "Un document de test", "domain": "financial", "embedding": [0.3, 0.4, 0.5], "metadata": [{ "key": "language", "value": "fr" }] }] PRICE_TABLE_NAME = "price_table" PRICE_TABLE_DATA = [{ "id": "item1", "content": "Budget item", "price": 25.99, "embedding": [0.1, 0.2, 0.3] }, { "id": "item2", "content": "Affordable product", "price": 49.99, "embedding": [0.2, 0.3, 0.4] }, { "id": "item3", "content": "Mid-range option", "price": 75.50, "embedding": [0.3, 0.4, 0.5] }, { "id": "item4", "content": "Premium product", "price": 149.99, "embedding": [0.4, 0.5, 0.6] }, { "id": "item5", "content": "Luxury item", "price": 199.99, "embedding": [0.5, 0.6, 0.7] }] class BigQueryVectorSearchIT(unittest.TestCase): bigquery_dataset_id = 'python_vector_search_test_' project = "apache-beam-testing" @classmethod def setUpClass(cls): cls.bigquery_client = BigQueryWrapper() cls.dataset_id = '%s%d%s' % ( cls.bigquery_dataset_id, int(time.time()), secrets.token_hex(3)) cls.bigquery_client.get_or_create_dataset(cls.project, cls.dataset_id) _LOGGER.info( "Created dataset %s in project %s", cls.dataset_id, cls.project) @classmethod def tearDownClass(cls): request = bigquery.BigqueryDatasetsDeleteRequest( projectId=cls.project, datasetId=cls.dataset_id, deleteContents=True) try: cls.bigquery_client.client.datasets.Delete(request) except Exception: _LOGGER.warning( 'Failed to clean up dataset %s in project %s', cls.dataset_id, cls.project) class TestBigQueryVectorSearchIT(BigQueryVectorSearchIT): # Test data with embeddings @classmethod def create_table(cls, table_name, schema_fields, table_data): """Create a BigQuery table with the specified schema and data. Args: table_name: Name of the table to create schema_fields: List of field definitions in the format: (name, type, [mode, [subfields]]) table_data: List of dictionaries containing the data to insert Returns: Fully qualified table name (project.dataset.table) """ table_schema = bigquery.TableSchema() for field_def in schema_fields: field = bigquery.TableFieldSchema() field.name = field_def[0] field.type = field_def[1] if len(field_def) > 2: field.mode = field_def[2] if len(field_def) > 3: for subfield_def in field_def[3]: subfield = bigquery.TableFieldSchema() subfield.name = subfield_def[0] subfield.type = subfield_def[1] field.fields.append(subfield) table_schema.fields.append(field) table = bigquery.Table( tableReference=bigquery.TableReference( projectId=cls.project, datasetId=cls.dataset_id, tableId=table_name), schema=table_schema) request = bigquery.BigqueryTablesInsertRequest( projectId=cls.project, datasetId=cls.dataset_id, table=table) cls.bigquery_client.client.tables.Insert(request) cls.bigquery_client.insert_rows( cls.project, cls.dataset_id, table_name, table_data) # Return the fully qualified table name return f"{cls.project}.{cls.dataset_id}.{table_name}" @classmethod def setUpClass(cls): super().setUpClass() cls.default_table_fqn = cls.create_table( table_name=DEFAULT_TABLE_NAME, schema_fields=DEFAULT_TABLE_SCHEMA, table_data=DEFAULT_TABLE_DATA) cls.price_table_fqn = cls.create_table( table_name=PRICE_TABLE_NAME, schema_fields=PRICE_TABLE_SCHEMA, table_data=PRICE_TABLE_DATA) def test_basic_vector_search(self): """Test basic vector similarity search.""" test_chunks = [ Chunk( id="query1", index=0, embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="test query"), metadata={"language": "en"}) ] # Expected chunk will have enrichment_data in metadata expected_chunks = [ Chunk( id="query1", index=0, embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="test query"), metadata={ "language": "en", "enrichment_data": { "id": "query1", "chunks": [{ "content": "This is a test document", "metadata": [{ "key": "language", "value": "en" }] }, { "content": "Another test document", "metadata": [{ "key": "language", "value": "en" }] }] } }) ] params = BigQueryVectorSearchParameters( project=self.project, table_name=self.default_table_fqn, embedding_column='embedding', columns=['content', 'metadata'], neighbor_count=2, metadata_restriction_template=( "check_metadata(metadata, 'language', '{language}')")) handler = BigQueryVectorSearchEnrichmentHandler( vector_search_parameters=params) with TestPipeline(is_integration_test=True) as p: result = (p | beam.Create(test_chunks) | Enrichment(handler)) assert_that(result, equal_to(expected_chunks)) def test_include_distance(self): """Test basic vector similarity search.""" test_chunks = [ Chunk( id="query1", index=0, embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="test query"), metadata={"language": "en"}) ] # Expected chunk will have enrichment_data in metadata expected_chunks = [ Chunk( id="query1", index=0, embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="test query"), metadata={ "language": "en", "enrichment_data": { "id": "query1", "chunks": [{ "distance": 0.0, "content": "This is a test document", "metadata": [{ "key": "language", "value": "en" }] }] } }) ] params = BigQueryVectorSearchParameters( project=self.project, table_name=self.default_table_fqn, embedding_column='embedding', columns=['content', 'metadata'], neighbor_count=1, include_distance=True, metadata_restriction_template=( "check_metadata(metadata, 'language', '{language}')")) handler = BigQueryVectorSearchEnrichmentHandler( vector_search_parameters=params) with TestPipeline(is_integration_test=True) as p: result = (p | beam.Create(test_chunks) | Enrichment(handler)) assert_that(result, equal_to(expected_chunks)) def test_batched_metadata_filter_vector_search(self): """Test vector similarity search with batching.""" test_chunks = [ Chunk( id="query1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="test query 1"), metadata={"language": "en"}, index=0), Chunk( id="query2", embedding=Embedding(dense_embedding=[0.2, 0.3, 0.4]), content=Content(text="test query 2"), metadata={"language": "en"}, index=1), Chunk( id="query3", embedding=Embedding(dense_embedding=[0.3, 0.4, 0.5]), content=Content(text="test query 3"), metadata={"language": "fr"}, index=2) ] params = BigQueryVectorSearchParameters( project=self.project, table_name=self.default_table_fqn, embedding_column='embedding', columns=['content', 'metadata'], neighbor_count=2, metadata_restriction_template=( "check_metadata(metadata, 'language', '{language}')")) handler = BigQueryVectorSearchEnrichmentHandler( vector_search_parameters=params, min_batch_size=2, # Force batching max_batch_size=2 # Process 2 chunks at a time ) expected_chunks = [ Chunk( id="query1", embedding=Embedding( dense_embedding=[0.1, 0.2, 0.3], sparse_embedding=None), content=Content(text="test query 1"), metadata={ "language": "en", "enrichment_data": { "id": "query1", "chunks": [{ "content": "This is a test document", "metadata": [{ "key": "language", "value": "en" }] }, { "content": "Another test document", "metadata": [{ "key": "language", "value": "en" }] }] } }, index=0), Chunk( id="query2", embedding=Embedding( dense_embedding=[0.2, 0.3, 0.4], sparse_embedding=None), content=Content(text="test query 2"), metadata={ "language": "en", "enrichment_data": { "id": "query2", "chunks": [{ "content": "Another test document", "metadata": [{ "key": "language", "value": "en" }] }, { "content": "This is a test document", "metadata": [{ "key": "language", "value": "en" }] }] } }, index=1), Chunk( id="query3", embedding=Embedding( dense_embedding=[0.3, 0.4, 0.5], sparse_embedding=None), content=Content(text="test query 3"), metadata={ "language": "fr", "enrichment_data": { "id": "query3", "chunks": [{ "content": "Un document de test", "metadata": [{ "key": "language", "value": "fr" }] }] } }, index=2) ] with TestPipeline(is_integration_test=True) as p: result = (p | beam.Create(test_chunks) | Enrichment(handler)) assert_that(result, equal_to(expected_chunks)) def test_euclidean_distance_search(self): """Test vector similarity search using Euclidean distance.""" test_chunks = [ Chunk( id="query1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="test query 1"), metadata={"language": "en"}, index=0), Chunk( id="query2", embedding=Embedding(dense_embedding=[0.2, 0.3, 0.4]), content=Content(text="test query 2"), metadata={"language": "en"}, index=1) ] params = BigQueryVectorSearchParameters( project=self.project, table_name=self.default_table_fqn, embedding_column='embedding', columns=['content', 'metadata'], neighbor_count=2, metadata_restriction_template=( "check_metadata(metadata, 'language', '{language}')"), distance_type='EUCLIDEAN' # Use Euclidean distance ) handler = BigQueryVectorSearchEnrichmentHandler( vector_search_parameters=params, min_batch_size=2, max_batch_size=2) expected_chunks = [ Chunk( id="query1", embedding=Embedding( dense_embedding=[0.1, 0.2, 0.3], sparse_embedding=None), content=Content(text="test query 1"), metadata={ "language": "en", "enrichment_data": { "id": "query1", "chunks": [{ "content": "This is a test document", "metadata": [{ "key": "language", "value": "en" }] }, { "content": "Another test document", "metadata": [{ "key": "language", "value": "en" }] }] } }, index=0), Chunk( id="query2", embedding=Embedding( dense_embedding=[0.2, 0.3, 0.4], sparse_embedding=None), content=Content(text="test query 2"), metadata={ "language": "en", "enrichment_data": { "id": "query2", "chunks": [{ "content": "Another test document", "metadata": [{ "key": "language", "value": "en" }] }, { "content": "This is a test document", "metadata": [{ "key": "language", "value": "en" }] }] } }, index=1) ] with TestPipeline(is_integration_test=True) as p: result = (p | beam.Create(test_chunks) | Enrichment(handler)) assert_that(result, equal_to(expected_chunks)) def test_no_metadata_restriction(self): """Test vector search without metadata filtering.""" test_chunks = [ Chunk( id="query1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="test query"), metadata={'language': 'fr'}, index=0) ] params = BigQueryVectorSearchParameters( project=self.project, table_name=self.default_table_fqn, embedding_column='embedding', columns=['content', 'metadata'], neighbor_count=2, # Get top 2 matches metadata_restriction_template=None # No filtering ) handler = BigQueryVectorSearchEnrichmentHandler( vector_search_parameters=params) with TestPipeline(is_integration_test=True) as p: result = (p | beam.Create(test_chunks) | Enrichment(handler)) # Should get matches regardless of metadata assert_that( result, equal_to([ Chunk( id="query1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="test query"), metadata={ "language": "fr", "enrichment_data": { "id": "query1", "chunks": [{ "content": "This is a test document", "metadata": [{ "key": "language", "value": "en" }] }, { "content": "Another test document", "metadata": [{ "key": "language", "value": "en" }] }] } }, index=0) ])) def test_metadata_filter_leakage(self): """Test that metadata filters don't leak between batched chunks.""" test_chunks = [ Chunk( id="query1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="medical query"), metadata={ "domain": "medical", "language": "en" }, index=0), Chunk( id="query2", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="unmatched query"), metadata={ "domain": "doesntexist", "language": "en" }, index=1) ] params = BigQueryVectorSearchParameters( project=self.project, table_name=self.default_table_fqn, embedding_column='embedding', columns=['content', 'metadata', 'domain'], neighbor_count=1, metadata_restriction_template=( "domain = '{domain}' AND " "check_metadata(metadata, 'language', '{language}')")) handler = BigQueryVectorSearchEnrichmentHandler( vector_search_parameters=params, min_batch_size=2, # Force batching max_batch_size=2) with TestPipeline(is_integration_test=True) as p: result = (p | beam.Create(test_chunks) | Enrichment(handler)) assert_that( result, equal_to([ Chunk( id="query1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="medical query"), metadata={ "domain": "medical", "language": "en", "enrichment_data": { "id": "query1", "chunks": [{ "content": "This is a test document", "domain": "medical", "metadata": [{ "key": "language", "value": "en" }] }] } }, index=0), Chunk( id="query2", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="unmatched query"), metadata={ "domain": "doesntexist", "language": "en", "enrichment_data": {} }, index=1) ])) def test_price_range_filter_with_callable(self): """Test vector search with a callable that filters by price range.""" test_chunks = [ Chunk( id="low_price_query", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="cheap product"), metadata={ "min_price": 10, "max_price": 50 }, index=0), Chunk( id="high_price_query", embedding=Embedding(dense_embedding=[0.4, 0.5, 0.6]), content=Content(text="expensive product"), metadata={ "min_price": 100, "max_price": 200 }, index=1) ] def price_range_filter(chunk): min_price = chunk.metadata.get("min_price", 0) max_price = chunk.metadata.get("max_price", float('inf')) return f"price >= {min_price} AND price <= {max_price}" params = BigQueryVectorSearchParameters( project=self.project, table_name=self.price_table_fqn, embedding_column='embedding', columns=['content', 'price'], neighbor_count=5, metadata_restriction_template=price_range_filter) handler = BigQueryVectorSearchEnrichmentHandler( vector_search_parameters=params) with TestPipeline(is_integration_test=True) as p: result = (p | beam.Create(test_chunks) | Enrichment(handler)) assert_that( result, equal_to([ Chunk( id="low_price_query", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="cheap product"), metadata={ "min_price": 10, "max_price": 50, "enrichment_data": { "id": "low_price_query", "chunks": [{ "content": "Budget item", "price": 25.99 }, { "content": "Affordable product", "price": 49.99 }] } }, index=0), Chunk( id="high_price_query", embedding=Embedding(dense_embedding=[0.4, 0.5, 0.6]), content=Content(text="expensive product"), metadata={ "min_price": 100, "max_price": 200, "enrichment_data": { "id": "high_price_query", "chunks": [{ "content": "Premium product", "price": 149.99 }, { "content": "Luxury item", "price": 199.99 }] } }, index=1) ])) def test_condition_batching(self): """Test that queries with same metadata conditions are batched together.""" # Create three queries with same conditions but different embeddings test_chunks = [ Chunk( id="query1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="english query 1"), metadata={"language": "en"}, index=0), Chunk( id="query2", embedding=Embedding(dense_embedding=[0.4, 0.5, 0.6]), content=Content(text="english query 2"), metadata={"language": "en"}, index=1), Chunk( id="query3", embedding=Embedding(dense_embedding=[0.7, 0.8, 0.9]), content=Content(text="french query 3"), metadata={"language": "fr"}, index=2) ] params = BigQueryVectorSearchParameters( project=self.project, table_name=self.default_table_fqn, embedding_column='embedding', columns=['content', 'domain', 'metadata'], neighbor_count=3, metadata_restriction_template=( "check_metadata(metadata, 'language', '{language}')")) handler = BigQueryVectorSearchEnrichmentHandler( vector_search_parameters=params, min_batch_size=10, # Force batching max_batch_size=100) with TestPipeline(is_integration_test=True) as p: result = (p | beam.Create(test_chunks) | Enrichment(handler)) # All queries should be handled in a single VECTOR_SEARCH # Each should get its closest match assert_that( result, equal_to([ Chunk( id="query1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="english query 1"), metadata={ "language": "en", "enrichment_data": { "id": "query1", "chunks": [{ "content": "This is a test document", "domain": "medical", "metadata": [{ "key": "language", "value": "en" }] }, { "content": "Another test document", "domain": "legal", "metadata": [{ "key": "language", "value": "en" }] }] } }, index=0), Chunk( id="query2", embedding=Embedding(dense_embedding=[0.4, 0.5, 0.6]), content=Content(text="english query 2"), metadata={ "language": "en", "enrichment_data": { "id": "query2", "chunks": [{ "content": "Another test document", "domain": "legal", "metadata": [{ "key": "language", "value": "en" }] }, { "content": "This is a test document", "domain": "medical", "metadata": [{ "key": "language", "value": "en" }] }] } }, index=1), Chunk( id="query3", embedding=Embedding(dense_embedding=[0.7, 0.8, 0.9]), content=Content(text="french query 3"), metadata={ "language": "fr", "enrichment_data": { "id": "query3", "chunks": [{ "content": "Un document de test", "domain": "financial", "metadata": [{ "key": "language", "value": "fr" }] }] } }, index=2) ])) def test_invalid_query(self): """Test error handling for invalid queries.""" test_chunks = [ Chunk( id="query1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="test query"), metadata={"language": "en"}) ] params = BigQueryVectorSearchParameters( project=self.project, table_name=self.default_table_fqn, embedding_column='nonexistent_column', # Invalid column columns=['content'], neighbor_count=1, metadata_restriction_template=( "language = '{language}'" # Invalid template )) handler = BigQueryVectorSearchEnrichmentHandler( vector_search_parameters=params) with self.assertRaisesRegex(Exception, "Unrecognized name"): with TestPipeline() as p: _ = (p | beam.Create(test_chunks) | Enrichment(handler)) def test_empty_input(self): """Test handling of empty input.""" params = BigQueryVectorSearchParameters( project=self.project, table_name=self.default_table_fqn, embedding_column='embedding', columns=['content'], neighbor_count=1) handler = BigQueryVectorSearchEnrichmentHandler( vector_search_parameters=params) with TestPipeline(is_integration_test=True) as p: result = (p | beam.Create([]) | Enrichment(handler)) assert_that(result, equal_to([])) def test_missing_embedding(self): """Test handling of chunks with missing embeddings.""" test_chunks = [ Chunk( id="query1", embedding=None, # Missing embedding content=Content(text="test query"), metadata={"language": "en"}, index=0) ] params = BigQueryVectorSearchParameters( project=self.project, table_name=self.default_table_fqn, embedding_column='embedding', columns=['content'], neighbor_count=1) handler = BigQueryVectorSearchEnrichmentHandler( vector_search_parameters=params) with self.assertRaisesRegex(Exception, "missing embedding"): with TestPipeline() as p: _ = (p | beam.Create(test_chunks) | Enrichment(handler)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/enrichment/milvus_search.py0000644000175100017510000005550115206113032024416 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from collections.abc import Sequence from dataclasses import dataclass from dataclasses import field from enum import Enum from typing import Any from typing import Optional from typing import Union from google.protobuf.json_format import MessageToDict from pymilvus import AnnSearchRequest from pymilvus import Hit from pymilvus import Hits from pymilvus import MilvusClient from pymilvus import SearchResult from pymilvus.exceptions import MilvusException from apache_beam.ml.rag.types import EmbeddableItem from apache_beam.ml.rag.types import Embedding from apache_beam.ml.rag.utils import MilvusConnectionParameters from apache_beam.ml.rag.utils import MilvusHelpers from apache_beam.ml.rag.utils import retry_with_backoff from apache_beam.ml.rag.utils import unpack_dataclass_with_kwargs from apache_beam.transforms.enrichment import EnrichmentSourceHandler class SearchStrategy(Enum): """Search strategies for information retrieval. Args: HYBRID: Combines vector and keyword search approaches. Leverages both semantic understanding and exact matching. Typically provides the most comprehensive results. Useful for queries with both conceptual and specific keyword components. VECTOR: Vector similarity search only. Based on semantic similarity between query and documents. Effective for conceptual searches and finding related content. Less sensitive to exact terminology than keyword search. KEYWORD: Keyword/text search only. Based on exact or fuzzy matching of specific terms. Effective for precise queries where exact wording matters. Less effective for conceptual or semantic searches. """ HYBRID = "hybrid" VECTOR = "vector" KEYWORD = "keyword" class KeywordSearchMetrics(Enum): """Metrics for keyword search. Args: BM25: Range [0 to ∞), Best Match 25 ranking algorithm for text relevance. Combines term frequency, inverse document frequency, and document length. Higher scores indicate greater relevance. Higher scores indicate greater relevance. Takes into account diminishing returns of term frequency. Balances between exact matching and semantic relevance. """ BM25 = "BM25" class VectorSearchMetrics(Enum): """Metrics for vector search. Args: COSINE: Range [-1 to 1], higher values indicate greater similarity. Value 1 means vectors point in identical direction. Value 0 means vectors are perpendicular to each other (no relationship). Value -1 means vectors point in exactly opposite directions. EUCLIDEAN_DISTANCE (L2): Range [0 to ∞), lower values indicate greater similarity. Value 0 means vectors are identical. Larger values mean more dissimilarity between vectors. INNER_PRODUCT (IP): Range varies based on vector magnitudes, higher values indicate greater similarity. Value 0 means vectors are perpendicular to each other. Positive values mean vectors share some directional component. Negative values mean vectors point in opposing directions. """ COSINE = "COSINE" EUCLIDEAN_DISTANCE = "L2" INNER_PRODUCT = "IP" class MilvusBaseRanker: """Base class for ranking algorithms in Milvus hybrid search strategy.""" def __int__(self): return def dict(self): return {} def __str__(self): return self.dict().__str__() @dataclass class BaseSearchParameters: """Base parameters for both vector and keyword search operations. Args: anns_field: Approximate nearest neighbor search field indicates field name containing the embedding to search. Required for both vector and keyword search. limit: Maximum number of results to return per query. Must be positive. Defaults to 3 search results. filter: Boolean expression string for filtering search results. Example: 'price <= 1000 AND category == "electronics"'. search_params: Additional search parameters specific to the search type. Example: {"metric_type": VectorSearchMetrics.EUCLIDEAN_DISTANCE}. consistency_level: Consistency level for read operations. Options: "Strong", "Session", "Bounded", "Eventually". Defaults to "Bounded" if not specified when creating the collection. """ anns_field: str limit: int = 3 filter: str = field(default_factory=str) search_params: dict[str, Any] = field(default_factory=dict) consistency_level: Optional[str] = None def __post_init__(self): if not self.anns_field: raise ValueError( "Approximate Nearest Neighbor Search (ANNS) field must be provided") if self.limit <= 0: raise ValueError(f"Search limit must be positive, got {self.limit}") @dataclass class VectorSearchParameters(BaseSearchParameters): """Parameters for vector similarity search operations. Inherits all parameters from BaseSearchParameters with the same semantics. The anns_field should contain dense vector embeddings for this search type. Args: kwargs: Optional keyword arguments for additional vector search parameters. Enables forward compatibility. Note: For inherited parameters documentation, see BaseSearchParameters. """ kwargs: dict[str, Any] = field(default_factory=dict) @dataclass class KeywordSearchParameters(BaseSearchParameters): """Parameters for keyword/text search operations. This class inherits all parameters from BaseSearchParameters with the same semantics. The anns_field should contain sparse vector embeddings content for this search type. Args: kwargs: Optional keyword arguments for additional keyword search parameters. Enables forward compatibility. Note: For inherited parameters documentation, see BaseSearchParameters. """ kwargs: dict[str, Any] = field(default_factory=dict) @dataclass class HybridSearchParameters: """Parameters for hybrid (vector + keyword) search operations. Args: vector: Parameters for the vector search component. keyword: Parameters for the keyword search component. ranker: Ranker for combining vector and keyword search results. Example: RRFRanker(k=100). limit: Maximum number of results to return per query. Defaults to 3 search results. kwargs: Optional keyword arguments for additional hybrid search parameters. Enables forward compatibility. """ vector: VectorSearchParameters keyword: KeywordSearchParameters ranker: MilvusBaseRanker limit: int = 3 kwargs: dict[str, Any] = field(default_factory=dict) def __post_init__(self): if not self.vector or not self.keyword: raise ValueError( "Vector and keyword search parameters must be provided for " "hybrid search") if not self.ranker: raise ValueError("Ranker must be provided for hybrid search") if self.limit <= 0: raise ValueError(f"Search limit must be positive, got {self.limit}") SearchStrategyType = Union[VectorSearchParameters, KeywordSearchParameters, HybridSearchParameters] @dataclass class MilvusSearchParameters: """Parameters configuring Milvus search operations. This class encapsulates all parameters needed to execute searches against Milvus collections, supporting vector, keyword, and hybrid search strategies. Args: collection_name: Name of the collection to search in. search_strategy: Type of search to perform (VECTOR, KEYWORD, or HYBRID). partition_names: List of partition names to restrict the search to. If empty, all partitions will be searched. output_fields: List of field names to include in search results. If empty, only primary fields including distances will be returned. timeout: Search operation timeout in seconds. If not specified, the client's default timeout is used. round_decimal: Number of decimal places for distance/similarity scores. Defaults to -1 means no rounding. """ collection_name: str search_strategy: SearchStrategyType partition_names: list[str] = field(default_factory=list) output_fields: list[str] = field(default_factory=list) timeout: Optional[float] = None round_decimal: int = -1 def __post_init__(self): if not self.collection_name: raise ValueError("Collection name must be provided") if not self.search_strategy: raise ValueError("Search strategy must be provided") @dataclass class MilvusCollectionLoadParameters: """Parameters that control how Milvus loads a collection into memory. This class provides fine-grained control over collection loading, which is particularly important in resource-constrained environments. Proper configuration can significantly reduce memory usage and improve query performance by loading only necessary data. Args: refresh: If True, forces a reload of the collection even if already loaded. Ensures the most up-to-date data is in memory. resource_groups: List of resource groups to load the collection into. Can be used for load balancing across multiple query nodes. load_fields: Specify which fields to load into memory. Loading only necessary fields reduces memory usage. If empty, all fields loaded. skip_load_dynamic_field: If True, dynamic/growing fields will not be loaded into memory. Saves memory when dynamic fields aren't needed. kwargs: Optional keyword arguments for additional collection load parameters. Enables forward compatibility. """ refresh: bool = field(default_factory=bool) resource_groups: list[str] = field(default_factory=list) load_fields: list[str] = field(default_factory=list) skip_load_dynamic_field: bool = field(default_factory=bool) kwargs: dict[str, Any] = field(default_factory=dict) @dataclass class MilvusSearchResult: """Search result from Milvus per embeddable_item. Args: id: List of entity IDs returned from the search. Can be either string or integer IDs. distance: List of distances/similarity scores for each returned entity. fields: List of dictionaries containing additional field values for each entity. Each dictionary corresponds to one returned entity. """ id: list[Union[str, int]] = field(default_factory=list) distance: list[float] = field(default_factory=list) fields: list[dict[str, Any]] = field(default_factory=list) InputT, OutputT = (Union[EmbeddableItem, list[EmbeddableItem]], list[tuple[EmbeddableItem, dict[str, Any]]]) class MilvusSearchEnrichmentHandler(EnrichmentSourceHandler[InputT, OutputT]): """Enrichment handler for Milvus vector database searches. This handler is designed to work with the :class:`apache_beam.transforms.enrichment.EnrichmentSourceHandler` transform. It enables enriching data through vector similarity, keyword, or hybrid searches against Milvus collections. The handler supports different search strategies: * Vector search - For finding similar embeddings based on vector similarity * Keyword search - For text-based retrieval using BM25 or other text metrics * Hybrid search - For combining vector and keyword search results This handler queries the Milvus database per element by default. To enable batching for improved performance, set the `min_batch_size` and `max_batch_size` parameters. These control the batching behavior in the :class:`apache_beam.transforms.utils.BatchElements` transform. For memory-intensive operations, the handler allows fine-grained control over collection loading through the `collection_load_parameters`. """ def __init__( self, connection_parameters: MilvusConnectionParameters, search_parameters: MilvusSearchParameters, *, collection_load_parameters: Optional[MilvusCollectionLoadParameters], min_batch_size: int = 1, max_batch_size: int = 1000, **kwargs): """ Example Usage: connection_parameters = MilvusConnectionParameters( uri="http://localhost:19530") search_parameters = MilvusSearchParameters( collection_name="my_collection", search_strategy=VectorSearchParameters(anns_field="embedding")) collection_load_parameters = MilvusCollectionLoadParameters( load_fields=["embedding", "metadata"]), milvus_handler = MilvusSearchEnrichmentHandler( connection_parameters, search_parameters, collection_load_parameters=collection_load_parameters, min_batch_size=10, max_batch_size=100) Args: connection_parameters (MilvusConnectionParameters): Configuration for connecting to the Milvus server, including URI, credentials, and connection options. search_parameters (MilvusSearchParameters): Configuration for search operations, including collection name, search strategy, and output fields. collection_load_parameters (Optional[MilvusCollectionLoadParameters]): Parameters controlling how collections are loaded into memory, which can significantly impact resource usage and performance. min_batch_size (int): Minimum number of elements to batch together when querying Milvus. Default is 1 (no batching when max_batch_size is 1). max_batch_size (int): Maximum number of elements to batch together.Default is 1000. Higher values may improve throughput but increase memory usage. **kwargs: Additional keyword arguments for Milvus Enrichment Handler. Note: * For large collections, consider setting appropriate values in collection_load_parameters to reduce memory usage. * The search_strategy in search_parameters determines the type of search (vector, keyword, or hybrid) and associated parameters. * Batching can significantly improve performance but requires more memory. """ self._connection_parameters = connection_parameters self._search_parameters = search_parameters self._collection_load_parameters = collection_load_parameters if not self._collection_load_parameters: self._collection_load_parameters = MilvusCollectionLoadParameters() self._batching_kwargs = { 'min_batch_size': min_batch_size, 'max_batch_size': max_batch_size } self.kwargs = kwargs self._client = None self.join_fn = join_fn self.use_custom_types = True def __enter__(self): """Enters the context manager and establishes Milvus connection. Returns: Self, enabling use in 'with' statements. """ if not self._client: connection_params = unpack_dataclass_with_kwargs( self._connection_parameters) collection_load_params = unpack_dataclass_with_kwargs( self._collection_load_parameters) # Extract retry parameters from connection_params. max_retries = connection_params.pop('max_retries', 3) retry_delay = connection_params.pop('retry_delay', 1.0) retry_backoff_factor = connection_params.pop('retry_backoff_factor', 2.0) def connect_and_load(): client = MilvusClient(**connection_params) client.load_collection( collection_name=self.collection_name, partition_names=self.partition_names, **collection_load_params) return client self._client = retry_with_backoff( connect_and_load, max_retries=max_retries, retry_delay=retry_delay, retry_backoff_factor=retry_backoff_factor, operation_name="Milvus connection and collection load", exception_types=(MilvusException, )) return self def __call__( self, request: Union[EmbeddableItem, list[EmbeddableItem]], *args, **kwargs) -> list[tuple[EmbeddableItem, dict[str, Any]]]: reqs = request if isinstance(request, list) else [request] # Early return for empty requests to avoid unnecessary connection attempts if not reqs: return [] search_result = self._search_documents(reqs) return self._get_call_response(reqs, search_result) def _search_documents(self, embeddable_items: list[EmbeddableItem]): if isinstance(self.search_strategy, HybridSearchParameters): data = self._get_hybrid_search_data(embeddable_items) return self._client.hybrid_search( collection_name=self.collection_name, partition_names=self.partition_names, output_fields=self.output_fields, timeout=self.timeout, round_decimal=self.round_decimal, reqs=data, ranker=self.search_strategy.ranker, limit=self.search_strategy.limit, **self.search_strategy.kwargs) elif isinstance(self.search_strategy, VectorSearchParameters): data = list(map(self._get_vector_search_data, embeddable_items)) vector_search_params = unpack_dataclass_with_kwargs(self.search_strategy) return self._client.search( collection_name=self.collection_name, partition_names=self.partition_names, output_fields=self.output_fields, timeout=self.timeout, round_decimal=self.round_decimal, data=data, **vector_search_params) elif isinstance(self.search_strategy, KeywordSearchParameters): data = list(map(self._get_keyword_search_data, embeddable_items)) keyword_search_params = unpack_dataclass_with_kwargs(self.search_strategy) return self._client.search( collection_name=self.collection_name, partition_names=self.partition_names, output_fields=self.output_fields, timeout=self.timeout, round_decimal=self.round_decimal, data=data, **keyword_search_params) else: raise ValueError( f"Not supported search strategy yet: {self.search_strategy}") def _get_hybrid_search_data(self, embeddable_items: list[EmbeddableItem]): vector_search_data = list( map(self._get_vector_search_data, embeddable_items)) keyword_search_data = list( map(self._get_keyword_search_data, embeddable_items)) vector_search_req = AnnSearchRequest( data=vector_search_data, anns_field=self.search_strategy.vector.anns_field, param=self.search_strategy.vector.search_params, limit=self.search_strategy.vector.limit, expr=self.search_strategy.vector.filter) keyword_search_req = AnnSearchRequest( data=keyword_search_data, anns_field=self.search_strategy.keyword.anns_field, param=self.search_strategy.keyword.search_params, limit=self.search_strategy.keyword.limit, expr=self.search_strategy.keyword.filter) reqs = [vector_search_req, keyword_search_req] return reqs def _get_vector_search_data(self, embeddable_item: EmbeddableItem): if not embeddable_item.dense_embedding: raise ValueError( f"Item {embeddable_item.id} missing dense embedding required for" " vector search") return embeddable_item.dense_embedding def _get_keyword_search_data(self, embeddable_item: EmbeddableItem): has_no_text = not embeddable_item.content.text has_no_sparse = not embeddable_item.sparse_embedding if has_no_text and has_no_sparse: raise ValueError( f"Item {embeddable_item.id} missing both text content and sparse " "embedding required for keyword search") sparse_embedding = MilvusHelpers.sparse_embedding( embeddable_item.sparse_embedding) return embeddable_item.content.text or sparse_embedding def _get_call_response( self, embeddable_items: list[EmbeddableItem], search_result: SearchResult[Hits]): response = [] for i in range(len(embeddable_items)): embeddable_item = embeddable_items[i] hits: Hits = search_result[i] result = MilvusSearchResult() for j in range(len(hits)): hit: Hit = hits[j] normalized_fields = self._normalize_milvus_fields(hit.fields) result.id.append(hit.id) result.distance.append(hit.distance) result.fields.append(normalized_fields) response.append((embeddable_item, result.__dict__)) return response def _normalize_milvus_fields(self, fields: dict[str, Any]): normalized_fields = {} for field, value in fields.items(): value = self._normalize_milvus_value(value) normalized_fields[field] = value return normalized_fields def _normalize_milvus_value(self, value: Any): # Convert Milvus-specific types to Python native types. neither_str_nor_dict_nor_bytes = not isinstance(value, (str, dict, bytes)) if isinstance(value, Sequence) and neither_str_nor_dict_nor_bytes: return list(value) elif hasattr(value, 'DESCRIPTOR'): # Handle protobuf messages. return MessageToDict(value) else: # Keep other types as they are. return value def convert_sparse_embedding_to_milvus_format( self, sparse_vector: tuple[list[int], list[float]]) -> dict[int, float]: if not sparse_vector: return None # Converts sparse embedding from (indices, values) tuple format to # Milvus-compatible values dict format {dimension_index: value, ...}. indices, values = sparse_vector return {int(idx): float(val) for idx, val in zip(indices, values)} @property def collection_name(self): """Getter method for collection_name property""" return self._search_parameters.collection_name @property def search_strategy(self): """Getter method for search_strategy property""" return self._search_parameters.search_strategy @property def partition_names(self): """Getter method for partition_names property""" return self._search_parameters.partition_names @property def output_fields(self): """Getter method for output_fields property""" return self._search_parameters.output_fields @property def timeout(self): """Getter method for search timeout property""" return self._search_parameters.timeout @property def round_decimal(self): """Getter method for search round_decimal property""" return self._search_parameters.round_decimal def __exit__(self, exc_type, exc_val, exc_tb): self._client.release_collection(self.collection_name) self._client.close() self._client = None def batch_elements_kwargs(self) -> dict[str, int]: """Returns kwargs for beam.BatchElements.""" return self._batching_kwargs def join_fn(left: Embedding, right: dict[str, Any]) -> Embedding: left.metadata['enrichment_data'] = right return left ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/enrichment/milvus_search_it_test.py0000644000175100017510000010533115206113032026146 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import platform import unittest from dataclasses import dataclass from dataclasses import field import pytest import apache_beam as beam from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import Embedding from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that # pylint: disable=ungrouped-imports try: from pymilvus import DataType from pymilvus import FieldSchema from pymilvus import Function from pymilvus import FunctionType from pymilvus import RRFRanker from pymilvus.milvus_client import IndexParams from apache_beam.ml.rag.enrichment.milvus_search import HybridSearchParameters from apache_beam.ml.rag.enrichment.milvus_search import KeywordSearchMetrics from apache_beam.ml.rag.enrichment.milvus_search import KeywordSearchParameters from apache_beam.ml.rag.enrichment.milvus_search import MilvusCollectionLoadParameters from apache_beam.ml.rag.enrichment.milvus_search import MilvusConnectionParameters from apache_beam.ml.rag.enrichment.milvus_search import MilvusSearchEnrichmentHandler from apache_beam.ml.rag.enrichment.milvus_search import MilvusSearchParameters from apache_beam.ml.rag.enrichment.milvus_search import VectorSearchMetrics from apache_beam.ml.rag.enrichment.milvus_search import VectorSearchParameters from apache_beam.ml.rag.test_utils import MilvusTestHelpers from apache_beam.ml.rag.test_utils import VectorDBContainerInfo from apache_beam.transforms.enrichment import Enrichment except ImportError as e: raise unittest.SkipTest(f'Milvus dependencies not installed: {str(e)}') def _construct_index_params(): index_params = IndexParams() # Milvus doesn't support multiple indexes on the same field. This is a # limitation of Milvus - someone can only create one index per field as yet. # Cosine similarity index on first dense embedding field index_params.add_index( field_name="dense_embedding_cosine", index_name="dense_embedding_cosine_ivf_flat", index_type="IVF_FLAT", metric_type=VectorSearchMetrics.COSINE.value, params={"nlist": 1}) # Euclidean distance index on second dense embedding field index_params.add_index( field_name="dense_embedding_euclidean", index_name="dense_embedding_euclidean_ivf_flat", index_type="IVF_FLAT", metric_type=VectorSearchMetrics.EUCLIDEAN_DISTANCE.value, params={"nlist": 1}) # Inner product index on third dense embedding field index_params.add_index( field_name="dense_embedding_inner_product", index_name="dense_embedding_inner_product_ivf_flat", index_type="IVF_FLAT", metric_type=VectorSearchMetrics.INNER_PRODUCT.value, params={"nlist": 1}) index_params.add_index( field_name="sparse_embedding_inner_product", index_name="sparse_embedding_inner_product_inverted_index", index_type="SPARSE_INVERTED_INDEX", metric_type=VectorSearchMetrics.INNER_PRODUCT.value, params={ "inverted_index_algo": "TAAT_NAIVE", }) # BM25 index on sparse_embedding field. # # For deterministic testing results # 1. Using TAAT_NAIVE: Most predictable algorithm that processes each term # completely before moving to the next. # 2. Using k1=1: Moderate term frequency weighting – repeated terms matter # but with diminishing returns. # 3. Using b=0: No document length normalization – longer documents not # penalized. # This combination provides maximum transparency and predictability for # test assertions. index_params.add_index( field_name="sparse_embedding_bm25", index_name="sparse_embedding_bm25_inverted_index", index_type="SPARSE_INVERTED_INDEX", metric_type=KeywordSearchMetrics.BM25.value, params={ "inverted_index_algo": "TAAT_NAIVE", "bm25_k1": 1, "bm25_b": 0, }) return index_params @dataclass class MilvusITDataConstruct: id: int content: str domain: str cost: int metadata: dict tags: list[str] dense_embedding: list[float] sparse_embedding: dict vocabulary: dict[str, int] = field(default_factory=dict) def __getitem__(self, key): return getattr(self, key) MILVUS_IT_CONFIG = { "collection_name": "docs_catalog", "fields": [ FieldSchema( name="id", dtype=DataType.INT64, is_primary=True, auto_id=False), FieldSchema( name="content", dtype=DataType.VARCHAR, max_length=512, enable_analyzer=True), FieldSchema(name="domain", dtype=DataType.VARCHAR, max_length=128), FieldSchema(name="cost", dtype=DataType.INT32), FieldSchema(name="metadata", dtype=DataType.JSON), FieldSchema( name="tags", dtype=DataType.ARRAY, element_type=DataType.VARCHAR, max_length=64, max_capacity=64), FieldSchema( name="dense_embedding_cosine", dtype=DataType.FLOAT_VECTOR, dim=3), FieldSchema( name="dense_embedding_euclidean", dtype=DataType.FLOAT_VECTOR, dim=3), FieldSchema( name="dense_embedding_inner_product", dtype=DataType.FLOAT_VECTOR, dim=3), FieldSchema( name="sparse_embedding_bm25", dtype=DataType.SPARSE_FLOAT_VECTOR), FieldSchema( name="sparse_embedding_inner_product", dtype=DataType.SPARSE_FLOAT_VECTOR) ], "functions": [ Function( name="content_bm25_emb", input_field_names=["content"], output_field_names=["sparse_embedding_bm25"], function_type=FunctionType.BM25) ], "index": _construct_index_params, "corpus": [ MilvusITDataConstruct( id=1, content="This is a test document", domain="medical", cost=49, metadata={"language": "en"}, tags=["healthcare", "patient", "clinical"], dense_embedding=[0.1, 0.2, 0.3], sparse_embedding={ 1: 0.05, 2: 0.41, 3: 0.05, 4: 0.41 }), MilvusITDataConstruct( id=2, content="Another test document", domain="legal", cost=75, metadata={"language": "en"}, tags=["contract", "law", "regulation"], dense_embedding=[0.2, 0.3, 0.4], sparse_embedding={ 1: 0.07, 3: 3.07, 0: 0.53 }), MilvusITDataConstruct( id=3, content="وثيقة اختبار", domain="financial", cost=149, metadata={"language": "ar"}, tags=["banking", "investment", "arabic"], dense_embedding=[0.3, 0.4, 0.5], sparse_embedding={ 6: 0.62, 5: 0.62 }) ], "vocabulary": { "this": 4, "is": 2, "test": 3, "document": 1, "another": 0, "وثيقة": 6, "اختبار": 5 } } @pytest.mark.require_docker_in_docker @unittest.skipUnless( platform.system() == "Linux", "Test runs only on Linux due to lack of support, as yet, for nested " "virtualization in CI environments on Windows/macOS. Many CI providers run " "tests in virtualized environments, and nested virtualization " "(Docker inside a VM) is either unavailable or has several issues on " "non-Linux platforms.") class TestMilvusSearchEnrichment(unittest.TestCase): """Tests for search functionality across all search strategies""" _db: VectorDBContainerInfo @classmethod def setUpClass(cls): cls._db = MilvusTestHelpers.start_db_container() cls._connection_params = MilvusConnectionParameters( uri=cls._db.uri, user=cls._db.user, password=cls._db.password, db_name=cls._db.id, token=cls._db.token) cls._collection_load_params = MilvusCollectionLoadParameters() cls._collection_name = MilvusTestHelpers.initialize_db_with_data( cls._connection_params, MILVUS_IT_CONFIG) @classmethod def tearDownClass(cls): MilvusTestHelpers.stop_db_container(cls._db) cls._db = None def test_invalid_query_on_non_existent_collection(self): non_existent_collection = "nonexistent_collection" existent_field = "dense_embedding_cosine" test_chunks = [ Chunk( embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content()) ] search_parameters = MilvusSearchParameters( collection_name=non_existent_collection, search_strategy=VectorSearchParameters(anns_field=existent_field)) collection_load_parameters = MilvusCollectionLoadParameters() handler = MilvusSearchEnrichmentHandler( self._connection_params, search_parameters, collection_load_parameters=collection_load_parameters) with self.assertRaises(Exception) as context: with TestPipeline() as p: p.not_use_test_runner_api = True _ = (p | beam.Create(test_chunks) | Enrichment(handler)) expect_err_msg_contains = "collection not found" self.assertIn(expect_err_msg_contains, str(context.exception)) def test_invalid_query_on_non_existent_field(self): non_existent_field = "nonexistent_column" existent_collection = MILVUS_IT_CONFIG["collection_name"] test_chunks = [ Chunk( embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content()) ] search_parameters = MilvusSearchParameters( collection_name=existent_collection, search_strategy=VectorSearchParameters(anns_field=non_existent_field)) collection_load_parameters = MilvusCollectionLoadParameters() handler = MilvusSearchEnrichmentHandler( self._connection_params, search_parameters, collection_load_parameters=collection_load_parameters) with self.assertRaises(Exception) as context: with TestPipeline() as p: p.not_use_test_runner_api = True _ = (p | beam.Create(test_chunks) | Enrichment(handler)) expect_err_msg_contains = f"fieldName({non_existent_field}) not found" self.assertIn(expect_err_msg_contains, str(context.exception)) def test_empty_input_chunks(self): test_chunks = [] anns_field = "dense_embedding_cosine" search_parameters = MilvusSearchParameters( collection_name=MILVUS_IT_CONFIG["collection_name"], search_strategy=VectorSearchParameters(anns_field=anns_field)) collection_load_parameters = MilvusCollectionLoadParameters() handler = MilvusSearchEnrichmentHandler( self._connection_params, search_parameters, collection_load_parameters=collection_load_parameters) expected_chunks = [] with TestPipeline() as p: p.not_use_test_runner_api = True result = (p | beam.Create(test_chunks) | Enrichment(handler)) assert_that( result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent( actual, expected_chunks)) def test_filtered_search_with_cosine_similarity_and_batching(self): test_chunks = [ Chunk( id="query1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content()), Chunk( id="query2", embedding=Embedding(dense_embedding=[0.2, 0.3, 0.4]), content=Content()), Chunk( id="query3", embedding=Embedding(dense_embedding=[0.3, 0.4, 0.5]), content=Content()) ] filter_condition = 'metadata["language"] == "en"' anns_field = "dense_embedding_cosine" addition_search_params = { "metric_type": VectorSearchMetrics.COSINE.value, "nprobe": 1 } vector_search_parameters = VectorSearchParameters( anns_field=anns_field, limit=10, filter=filter_condition, search_params=addition_search_params) search_parameters = MilvusSearchParameters( collection_name=MILVUS_IT_CONFIG["collection_name"], search_strategy=vector_search_parameters, output_fields=["id", "content", "metadata"], round_decimal=1) collection_load_parameters = MilvusCollectionLoadParameters() # Force batching. min_batch_size, max_batch_size = 2, 2 handler = MilvusSearchEnrichmentHandler( connection_parameters=self._connection_params, search_parameters=search_parameters, collection_load_parameters=collection_load_parameters, min_batch_size=min_batch_size, max_batch_size=max_batch_size) expected_chunks = [ Chunk( id='query1', content=Content(), metadata={ 'enrichment_data': { 'id': [1, 2], 'distance': [1.0, 1.0], 'fields': [{ 'content': 'This is a test document', 'metadata': { 'language': 'en' }, 'id': 1 }, { 'content': 'Another test document', 'metadata': { 'language': 'en' }, 'id': 2 }] } }, embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3])), Chunk( id='query2', content=Content(), metadata={ 'enrichment_data': { 'id': [2, 1], 'distance': [1.0, 1.0], 'fields': [{ 'content': 'Another test document', 'metadata': { 'language': 'en' }, 'id': 2 }, { 'content': 'This is a test document', 'metadata': { 'language': 'en' }, 'id': 1 }] } }, embedding=Embedding(dense_embedding=[0.2, 0.3, 0.4])), Chunk( id='query3', content=Content(), metadata={ 'enrichment_data': { 'id': [2, 1], 'distance': [1.0, 1.0], 'fields': [{ 'content': 'Another test document', 'metadata': { 'language': 'en' }, 'id': 2 }, { 'content': 'This is a test document', 'metadata': { 'language': 'en' }, 'id': 1 }] } }, embedding=Embedding(dense_embedding=[0.3, 0.4, 0.5])) ] with TestPipeline() as p: p.not_use_test_runner_api = True result = (p | beam.Create(test_chunks) | Enrichment(handler)) assert_that( result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent( actual, expected_chunks)) def test_filtered_search_with_bm25_full_text_and_batching(self): test_chunks = [ Chunk( id="query1", embedding=Embedding(sparse_embedding=None), content=Content(text="This is a test document")), Chunk( id="query2", embedding=Embedding(sparse_embedding=None), content=Content(text="Another test document")), Chunk( id="query3", embedding=Embedding(sparse_embedding=None), content=Content(text="وثيقة اختبار")) ] filter_condition = 'ARRAY_CONTAINS_ANY(tags, ["healthcare", "banking"])' anns_field = "sparse_embedding_bm25" addition_search_params = {"metric_type": KeywordSearchMetrics.BM25.value} keyword_search_parameters = KeywordSearchParameters( anns_field=anns_field, limit=10, filter=filter_condition, search_params=addition_search_params) search_parameters = MilvusSearchParameters( collection_name=MILVUS_IT_CONFIG["collection_name"], search_strategy=keyword_search_parameters, output_fields=["id", "content", "metadata"], round_decimal=1) collection_load_parameters = MilvusCollectionLoadParameters() # Force batching. min_batch_size, max_batch_size = 2, 2 handler = MilvusSearchEnrichmentHandler( connection_parameters=self._connection_params, search_parameters=search_parameters, collection_load_parameters=collection_load_parameters, min_batch_size=min_batch_size, max_batch_size=max_batch_size) expected_chunks = [ Chunk( id='query1', content=Content(text='This is a test document'), metadata={ 'enrichment_data': { 'id': [1], 'distance': [3.3], 'fields': [{ 'content': 'This is a test document', 'metadata': { 'language': 'en' }, 'id': 1 }] } }, embedding=Embedding()), Chunk( id='query2', content=Content(text='Another test document'), metadata={ 'enrichment_data': { 'id': [1], 'distance': [0.8], 'fields': [{ 'content': 'This is a test document', 'metadata': { 'language': 'en' }, 'id': 1 }] } }, embedding=Embedding()), Chunk( id='query3', content=Content(text='وثيقة اختبار'), metadata={ 'enrichment_data': { 'id': [3], 'distance': [2.3], 'fields': [{ 'content': 'وثيقة اختبار', 'metadata': { 'language': 'ar' }, 'id': 3 }] } }, embedding=Embedding()) ] with TestPipeline() as p: p.not_use_test_runner_api = True result = (p | beam.Create(test_chunks) | Enrichment(handler)) assert_that( result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent( actual, expected_chunks)) def test_vector_search_with_euclidean_distance(self): test_chunks = [ Chunk( id="query1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content()), Chunk( id="query2", embedding=Embedding(dense_embedding=[0.2, 0.3, 0.4]), content=Content()), Chunk( id="query3", embedding=Embedding(dense_embedding=[0.3, 0.4, 0.5]), content=Content()) ] anns_field = "dense_embedding_euclidean" addition_search_params = { "metric_type": VectorSearchMetrics.EUCLIDEAN_DISTANCE.value, "nprobe": 1 } vector_search_parameters = VectorSearchParameters( anns_field=anns_field, limit=10, search_params=addition_search_params) search_parameters = MilvusSearchParameters( collection_name=MILVUS_IT_CONFIG["collection_name"], search_strategy=vector_search_parameters, output_fields=["id", "content", "metadata"], round_decimal=1) collection_load_parameters = MilvusCollectionLoadParameters() handler = MilvusSearchEnrichmentHandler( connection_parameters=self._connection_params, search_parameters=search_parameters, collection_load_parameters=collection_load_parameters) expected_chunks = [ Chunk( id='query1', content=Content(), metadata={ 'enrichment_data': { 'id': [1, 2, 3], 'distance': [0.0, 0.0, 0.1], 'fields': [{ 'content': 'This is a test document', 'metadata': { 'language': 'en' }, 'id': 1 }, { 'content': 'Another test document', 'metadata': { 'language': 'en' }, 'id': 2 }, { 'content': 'وثيقة اختبار', 'metadata': { 'language': 'ar' }, 'id': 3 }] } }, embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3])), Chunk( id='query2', content=Content(), metadata={ 'enrichment_data': { 'id': [2, 3, 1], 'distance': [0.0, 0.0, 0.0], 'fields': [{ 'content': 'Another test document', 'metadata': { 'language': 'en' }, 'id': 2 }, { 'content': 'وثيقة اختبار', 'metadata': { 'language': 'ar' }, 'id': 3 }, { 'content': 'This is a test document', 'metadata': { 'language': 'en' }, 'id': 1 }] } }, embedding=Embedding(dense_embedding=[0.2, 0.3, 0.4])), Chunk( id='query3', content=Content(), metadata={ 'enrichment_data': { 'id': [3, 2, 1], 'distance': [0.0, 0.0, 0.1], 'fields': [{ 'content': 'وثيقة اختبار', 'metadata': { 'language': 'ar' }, 'id': 3 }, { 'content': 'Another test document', 'metadata': { 'language': 'en' }, 'id': 2 }, { 'content': 'This is a test document', 'metadata': { 'language': 'en' }, 'id': 1 }] } }, embedding=Embedding(dense_embedding=[0.3, 0.4, 0.5])) ] with TestPipeline() as p: p.not_use_test_runner_api = True result = (p | beam.Create(test_chunks) | Enrichment(handler)) assert_that( result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent( actual, expected_chunks)) def test_vector_search_with_inner_product_similarity(self): test_chunks = [ Chunk( id="query1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content()), Chunk( id="query2", embedding=Embedding(dense_embedding=[0.2, 0.3, 0.4]), content=Content()), Chunk( id="query3", embedding=Embedding(dense_embedding=[0.3, 0.4, 0.5]), content=Content()) ] anns_field = "dense_embedding_inner_product" addition_search_params = { "metric_type": VectorSearchMetrics.INNER_PRODUCT.value, "nprobe": 1 } vector_search_parameters = VectorSearchParameters( anns_field=anns_field, limit=10, search_params=addition_search_params) search_parameters = MilvusSearchParameters( collection_name=MILVUS_IT_CONFIG["collection_name"], search_strategy=vector_search_parameters, output_fields=["id", "content", "metadata"], round_decimal=1) collection_load_parameters = MilvusCollectionLoadParameters() handler = MilvusSearchEnrichmentHandler( connection_parameters=self._connection_params, search_parameters=search_parameters, collection_load_parameters=collection_load_parameters) expected_chunks = [ Chunk( id='query1', content=Content(), metadata={ 'enrichment_data': { 'id': [3, 2, 1], 'distance': [0.3, 0.2, 0.1], 'fields': [{ 'content': 'وثيقة اختبار', 'metadata': { 'language': 'ar' }, 'id': 3 }, { 'content': 'Another test document', 'metadata': { 'language': 'en' }, 'id': 2 }, { 'content': 'This is a test document', 'metadata': { 'language': 'en' }, 'id': 1 }] } }, embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3])), Chunk( id='query2', content=Content(), metadata={ 'enrichment_data': { 'id': [3, 2, 1], 'distance': [0.4, 0.3, 0.2], 'fields': [{ 'content': 'وثيقة اختبار', 'metadata': { 'language': 'ar' }, 'id': 3 }, { 'content': 'Another test document', 'metadata': { 'language': 'en' }, 'id': 2 }, { 'content': 'This is a test document', 'metadata': { 'language': 'en' }, 'id': 1 }] } }, embedding=Embedding(dense_embedding=[0.2, 0.3, 0.4])), Chunk( id='query3', content=Content(), metadata={ 'enrichment_data': { 'id': [3, 2, 1], 'distance': [0.5, 0.4, 0.3], 'fields': [{ 'content': 'وثيقة اختبار', 'metadata': { 'language': 'ar' }, 'id': 3 }, { 'content': 'Another test document', 'metadata': { 'language': 'en' }, 'id': 2 }, { 'content': 'This is a test document', 'metadata': { 'language': 'en' }, 'id': 1 }] } }, embedding=Embedding(dense_embedding=[0.3, 0.4, 0.5])) ] with TestPipeline() as p: p.not_use_test_runner_api = True result = (p | beam.Create(test_chunks) | Enrichment(handler)) assert_that( result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent( actual, expected_chunks)) def test_keyword_search_with_inner_product_sparse_embedding(self): test_chunks = [ Chunk( id="query1", embedding=Embedding( sparse_embedding=([1, 2, 3, 4], [0.05, 0.41, 0.05, 0.41])), content=Content()) ] anns_field = "sparse_embedding_inner_product" addition_search_params = { "metric_type": VectorSearchMetrics.INNER_PRODUCT.value, } keyword_search_parameters = KeywordSearchParameters( anns_field=anns_field, limit=3, search_params=addition_search_params) search_parameters = MilvusSearchParameters( collection_name=MILVUS_IT_CONFIG["collection_name"], search_strategy=keyword_search_parameters, output_fields=["id", "content", "metadata"], round_decimal=1) collection_load_parameters = MilvusCollectionLoadParameters() handler = MilvusSearchEnrichmentHandler( connection_parameters=self._connection_params, search_parameters=search_parameters, collection_load_parameters=collection_load_parameters) expected_chunks = [ Chunk( id='query1', content=Content(), metadata={ 'enrichment_data': { 'id': [1, 2], 'distance': [0.3, 0.2], 'fields': [{ 'content': 'This is a test document', 'metadata': { 'language': 'en' }, 'id': 1 }, { 'content': 'Another test document', 'metadata': { 'language': 'en' }, 'id': 2 }] } }, embedding=Embedding( sparse_embedding=([1, 2, 3, 4], [0.05, 0.41, 0.05, 0.41]))) ] with TestPipeline() as p: p.not_use_test_runner_api = True result = (p | beam.Create(test_chunks) | Enrichment(handler)) assert_that( result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent( actual, expected_chunks)) def test_hybrid_search(self): test_chunks = [ Chunk( id="query1", embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3]), content=Content(text="This is a test document")) ] anns_vector_field = "dense_embedding_cosine" addition_vector_search_params = { "metric_type": VectorSearchMetrics.COSINE.value, "nprobe": 1 } vector_search_parameters = VectorSearchParameters( anns_field=anns_vector_field, limit=10, search_params=addition_vector_search_params) anns_keyword_field = "sparse_embedding_bm25" addition_keyword_search_params = { "metric_type": KeywordSearchMetrics.BM25.value } keyword_search_parameters = KeywordSearchParameters( anns_field=anns_keyword_field, limit=10, search_params=addition_keyword_search_params) hybrid_search_parameters = HybridSearchParameters( vector=vector_search_parameters, keyword=keyword_search_parameters, ranker=RRFRanker(1), limit=1) search_parameters = MilvusSearchParameters( collection_name=MILVUS_IT_CONFIG["collection_name"], search_strategy=hybrid_search_parameters, output_fields=["id", "content", "metadata"], round_decimal=1) collection_load_parameters = MilvusCollectionLoadParameters() handler = MilvusSearchEnrichmentHandler( connection_parameters=self._connection_params, search_parameters=search_parameters, collection_load_parameters=collection_load_parameters) expected_chunks = [ Chunk( content=Content(text='This is a test document'), id='query1', metadata={ 'enrichment_data': { 'id': [1], 'distance': [1.0], 'fields': [{ 'content': 'This is a test document', 'metadata': { 'language': 'en' }, 'id': 1 }] } }, embedding=Embedding(dense_embedding=[0.1, 0.2, 0.3])) ] with TestPipeline() as p: p.not_use_test_runner_api = True result = (p | beam.Create(test_chunks) | Enrichment(handler)) assert_that( result, lambda actual: MilvusTestHelpers.assert_chunks_equivalent( actual, expected_chunks)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/enrichment/milvus_search_test.py0000644000175100017510000003360415206113032025455 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest from parameterized import parameterized try: from apache_beam.ml.rag.enrichment.milvus_search import HybridSearchParameters from apache_beam.ml.rag.enrichment.milvus_search import KeywordSearchParameters from apache_beam.ml.rag.enrichment.milvus_search import MilvusBaseRanker from apache_beam.ml.rag.enrichment.milvus_search import MilvusCollectionLoadParameters from apache_beam.ml.rag.enrichment.milvus_search import MilvusConnectionParameters from apache_beam.ml.rag.enrichment.milvus_search import MilvusSearchEnrichmentHandler from apache_beam.ml.rag.enrichment.milvus_search import MilvusSearchParameters from apache_beam.ml.rag.enrichment.milvus_search import VectorSearchParameters from apache_beam.ml.rag.enrichment.milvus_search import unpack_dataclass_with_kwargs from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import Embedding except ImportError as e: raise unittest.SkipTest(f'Milvus dependencies not installed: {str(e)}') class MockRanker(MilvusBaseRanker): def dict(self): return {"name": "mock_ranker"} class TestMilvusSearchEnrichment(unittest.TestCase): """Unit tests for general search functionality in the Enrichment Handler.""" def test_invalid_connection_parameters(self): """Test validation errors for invalid connection parameters.""" # Empty URI in connection parameters. with self.assertRaises(ValueError) as context: connection_params = MilvusConnectionParameters(uri="") search_params = MilvusSearchParameters( collection_name="test_collection", search_strategy=VectorSearchParameters(anns_field="embedding")) collection_load_params = MilvusCollectionLoadParameters() _ = MilvusSearchEnrichmentHandler( connection_parameters=connection_params, search_parameters=search_params, collection_load_parameters=collection_load_params) self.assertIn( "URI must be provided for Milvus connection", str(context.exception)) @parameterized.expand([ # Empty collection name. ( lambda: MilvusSearchParameters( collection_name="", search_strategy=VectorSearchParameters(anns_field="embedding")), "Collection name must be provided" ), # Missing search strategy. ( lambda: MilvusSearchParameters( collection_name="test_collection", search_strategy=None), # type: ignore[arg-type] "Search strategy must be provided" ), ]) def test_invalid_search_parameters(self, create_params, expected_error_msg): """Test validation errors for invalid general search parameters.""" with self.assertRaises(ValueError) as context: connection_params = MilvusConnectionParameters( uri="http://localhost:19530") search_params = create_params() collection_load_params = MilvusCollectionLoadParameters() _ = MilvusSearchEnrichmentHandler( connection_parameters=connection_params, search_parameters=search_params, collection_load_parameters=collection_load_params) self.assertIn(expected_error_msg, str(context.exception)) def test_unpack_dataclass_with_kwargs(self): """Test the unpack_dataclass_with_kwargs function.""" # Create a test dataclass instance. connection_params = MilvusConnectionParameters( uri="http://localhost:19530", user="test_user", kwargs={"custom_param": "value"}) # Call the actual function. result = unpack_dataclass_with_kwargs(connection_params) # Verify the function correctly unpacks the dataclass and merges kwargs. self.assertEqual(result["uri"], "http://localhost:19530") self.assertEqual(result["user"], "test_user") self.assertEqual(result["custom_param"], "value") # Verify that kwargs take precedence over existing attributes. connection_params_with_override = MilvusConnectionParameters( uri="http://localhost:19530", user="test_user", kwargs={"user": "override_user"}) result_with_override = unpack_dataclass_with_kwargs( connection_params_with_override) self.assertEqual(result_with_override["user"], "override_user") class TestMilvusVectorSearchEnrichment(unittest.TestCase): """Unit tests specific to vector search functionality""" @parameterized.expand([ # Negative limit in vector search parameters. ( lambda: VectorSearchParameters(anns_field="embedding", limit=-1), "Search limit must be positive, got -1" ), # Missing anns_field in vector search parameters. ( lambda: VectorSearchParameters(anns_field=None), # type: ignore[arg-type] "Approximate Nearest Neighbor Search (ANNS) field must be provided" ), ]) def test_invalid_search_parameters(self, create_params, expected_error_msg): """Test validation errors for invalid vector search parameters.""" with self.assertRaises(ValueError) as context: connection_params = MilvusConnectionParameters( uri="http://localhost:19530") vector_search_params = create_params() search_params = MilvusSearchParameters( collection_name="test_collection", search_strategy=vector_search_params) collection_load_params = MilvusCollectionLoadParameters() _ = MilvusSearchEnrichmentHandler( connection_parameters=connection_params, search_parameters=search_params, collection_load_parameters=collection_load_params) self.assertIn(expected_error_msg, str(context.exception)) def test_missing_dense_embedding(self): with self.assertRaises(ValueError) as context: chunk = Chunk( id=1, content=None, embedding=Embedding(dense_embedding=None)) connection_params = MilvusConnectionParameters( uri="http://localhost:19530") vector_search_params = VectorSearchParameters(anns_field="embedding") search_params = MilvusSearchParameters( collection_name="test_collection", search_strategy=vector_search_params) collection_load_params = MilvusCollectionLoadParameters() handler = MilvusSearchEnrichmentHandler( connection_parameters=connection_params, search_parameters=search_params, collection_load_parameters=collection_load_params) _ = handler._get_vector_search_data(chunk) err_msg = "Item 1 missing dense embedding required for vector search" self.assertIn(err_msg, str(context.exception)) class TestMilvusKeywordSearchEnrichment(unittest.TestCase): """Unit tests specific to keyword search functionality""" @parameterized.expand([ # Negative limit in keyword search parameters. ( lambda: KeywordSearchParameters( anns_field="sparse_embedding", limit=-1), "Search limit must be positive, got -1" ), # Missing anns_field in keyword search parameters. ( lambda: KeywordSearchParameters(anns_field=None), # type: ignore[arg-type] "Approximate Nearest Neighbor Search (ANNS) field must be provided" ), ]) def test_invalid_search_parameters(self, create_params, expected_error_msg): """Test validation errors for invalid keyword search parameters.""" with self.assertRaises(ValueError) as context: connection_params = MilvusConnectionParameters( uri="http://localhost:19530") keyword_search_params = create_params() search_params = MilvusSearchParameters( collection_name="test_collection", search_strategy=keyword_search_params) collection_load_params = MilvusCollectionLoadParameters() _ = MilvusSearchEnrichmentHandler( connection_parameters=connection_params, search_parameters=search_params, collection_load_parameters=collection_load_params) self.assertIn(expected_error_msg, str(context.exception)) def test_missing_text_content_and_sparse_embedding(self): with self.assertRaises(ValueError) as context: chunk = Chunk( id=1, content=Content(text=None), embedding=Embedding(sparse_embedding=None)) connection_params = MilvusConnectionParameters( uri="http://localhost:19530") vector_search_params = VectorSearchParameters(anns_field="embedding") search_params = MilvusSearchParameters( collection_name="test_collection", search_strategy=vector_search_params) collection_load_params = MilvusCollectionLoadParameters() handler = MilvusSearchEnrichmentHandler( connection_parameters=connection_params, search_parameters=search_params, collection_load_parameters=collection_load_params) _ = handler._get_keyword_search_data(chunk) err_msg = ( "Item 1 missing both text content and sparse embedding " "required for keyword search") self.assertIn(err_msg, str(context.exception)) def test_missing_text_content_only(self): try: chunk = Chunk( id=1, content=Content(text=None), embedding=Embedding( sparse_embedding=([1, 2, 3, 4], [0.05, 0.41, 0.05, 0.41]))) connection_params = MilvusConnectionParameters( uri="http://localhost:19530") vector_search_params = VectorSearchParameters(anns_field="embedding") search_params = MilvusSearchParameters( collection_name="test_collection", search_strategy=vector_search_params) collection_load_params = MilvusCollectionLoadParameters() handler = MilvusSearchEnrichmentHandler( connection_parameters=connection_params, search_parameters=search_params, collection_load_parameters=collection_load_params) _ = handler._get_keyword_search_data(chunk) except Exception as e: self.fail(f"raised an unexpected exception: {e}") def test_missing_sparse_embedding_only(self): try: chunk = Chunk( id=1, content=Content(text="what is apache beam?"), embedding=Embedding(sparse_embedding=None)) connection_params = MilvusConnectionParameters( uri="http://localhost:19530") vector_search_params = VectorSearchParameters(anns_field="embedding") search_params = MilvusSearchParameters( collection_name="test_collection", search_strategy=vector_search_params) collection_load_params = MilvusCollectionLoadParameters() handler = MilvusSearchEnrichmentHandler( connection_parameters=connection_params, search_parameters=search_params, collection_load_parameters=collection_load_params) _ = handler._get_keyword_search_data(chunk) except Exception as e: self.fail(f"raised an unexpected exception: {e}") pass class TestMilvusHybridSearchEnrichment(unittest.TestCase): """Tests specific to hybrid search functionality""" @parameterized.expand([ # Missing vector in hybrid search parameters. ( lambda: HybridSearchParameters( vector=None, # type: ignore[arg-type] keyword=KeywordSearchParameters(anns_field="sparse_embedding"), ranker=MockRanker()), "Vector and keyword search parameters must be provided for hybrid " "search" ), # Missing keyword in hybrid search parameters. ( lambda: HybridSearchParameters( vector=VectorSearchParameters(anns_field="embedding"), keyword=None, # type: ignore[arg-type] ranker=MockRanker()), "Vector and keyword search parameters must be provided for hybrid " "search" ), # Missing ranker in hybrid search parameters. ( lambda: HybridSearchParameters( vector=VectorSearchParameters(anns_field="embedding"), keyword=KeywordSearchParameters(anns_field="sparse_embedding"), ranker=None), # type: ignore[arg-type] "Ranker must be provided for hybrid search" ), # Negative limit in hybrid search parameters. ( lambda: HybridSearchParameters( vector=VectorSearchParameters(anns_field="embedding"), keyword=KeywordSearchParameters(anns_field="sparse_embedding"), ranker=MockRanker(), limit=-1), "Search limit must be positive, got -1" ), ]) def test_invalid_search_parameters(self, create_params, expected_error_msg): """Test validation errors for invalid hybrid search parameters.""" with self.assertRaises(ValueError) as context: connection_params = MilvusConnectionParameters( uri="http://localhost:19530") hybrid_search_params = create_params() search_params = MilvusSearchParameters( collection_name="test_collection", search_strategy=hybrid_search_params) collection_load_params = MilvusCollectionLoadParameters() _ = MilvusSearchEnrichmentHandler( connection_parameters=connection_params, search_parameters=search_params, collection_load_parameters=collection_load_params) self.assertIn(expected_error_msg, str(context.exception)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.514422 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/0000755000175100017510000000000015206113160021037 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/__init__.py0000644000175100017510000000162515206113032023152 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Vector storage ingestion components for RAG pipelines. This module provides components for storing vectors in RAG pipelines. """ ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/alloydb.py0000644000175100017510000001706015206113032023041 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from dataclasses import dataclass from typing import Any from typing import Optional from apache_beam.ml.rag.ingestion.jdbc_common import ConnectionConfig from apache_beam.ml.rag.ingestion.jdbc_common import WriteConfig from apache_beam.ml.rag.ingestion.postgres import ColumnSpecsBuilder from apache_beam.ml.rag.ingestion.postgres import PostgresVectorWriterConfig from apache_beam.ml.rag.ingestion.postgres_common import ColumnSpec from apache_beam.ml.rag.ingestion.postgres_common import ConflictResolution @dataclass class AlloyDBLanguageConnectorConfig: """Configuration options for AlloyDB language connector. Contains all parameters needed to configure a connection using the AlloyDB Java connector via JDBC. For details see https://github.com/GoogleCloudPlatform/alloydb-java-connector/blob/main/docs/jdbc.md Attributes: username: Database username. password: Database password. Can be empty string when using IAM. database_name: Name of the database to connect to. instance_name: Fullly qualified instance. Format: 'projects//locations//clusters//instances /' ip_type: IP type to use for connection. Either 'PRIVATE' (default), 'PUBLIC' 'PSC. enable_iam_auth: Whether to enable IAM authentication. Default is False target_principal: Optional service account to impersonate for connection. delegates: Optional comma-separated list of service accounts for delegated impersonation. admin_service_endpoint: Optional custom API service endpoint. quota_project: Optional project ID for quota and billing. connection_properties: Optional JDBC connection properties dict. Example: {'ssl': 'true'} additional_properties: Additional properties to be added to the JDBC url. Example: {'someProperty': 'true'} """ username: str password: str database_name: str instance_name: str ip_type: str = "PRIVATE" enable_iam_auth: bool = False target_principal: Optional[str] = None delegates: Optional[list[str]] = None admin_service_endpoint: Optional[str] = None quota_project: Optional[str] = None connection_properties: Optional[dict[str, str]] = None additional_properties: Optional[dict[str, Any]] = None def to_jdbc_url(self) -> str: """Convert options to a properly formatted JDBC URL. Returns: JDBC URL string configured with all options. """ # Base URL with database name url = f"jdbc:postgresql:///{self.database_name}?" # Add required properties properties = { "socketFactory": "com.google.cloud.alloydb.SocketFactory", "alloydbInstanceName": self.instance_name, "alloydbIpType": self.ip_type } if self.enable_iam_auth: properties["alloydbEnableIAMAuth"] = "true" if self.target_principal: properties["alloydbTargetPrincipal"] = self.target_principal if self.delegates: properties["alloydbDelegates"] = ",".join(self.delegates) if self.admin_service_endpoint: properties["alloydbAdminServiceEndpoint"] = self.admin_service_endpoint if self.quota_project: properties["alloydbQuotaProject"] = self.quota_project if self.additional_properties: properties.update(self.additional_properties) property_string = "&".join(f"{k}={v}" for k, v in properties.items()) return url + property_string def to_connection_config(self): return ConnectionConfig( jdbc_url=self.to_jdbc_url(), username=self.username, password=self.password, connection_properties=self.connection_properties, additional_jdbc_args=self.additional_jdbc_args()) def additional_jdbc_args(self) -> dict[str, list[Any]]: return { 'classpath': [ "org.postgresql:postgresql:42.2.16", "com.google.cloud:alloydb-jdbc-connector:1.2.0" ] } class AlloyDBVectorWriterConfig(PostgresVectorWriterConfig): def __init__( self, connection_config: AlloyDBLanguageConnectorConfig, table_name: str, *, # pylint: disable=dangerous-default-value write_config: WriteConfig = WriteConfig(), column_specs: list[ColumnSpec] = ColumnSpecsBuilder.with_defaults().build( ), conflict_resolution: Optional[ConflictResolution] = ConflictResolution( on_conflict_fields=[], action='IGNORE')): """Configuration for writing vectors to AlloyDB. Supports flexible schema configuration through column specifications and conflict resolution strategies. Args: connection_config: AlloyDB connection configuration. table_name: Target table name. write_config: JdbcIO :class:`~.jdbc_common.WriteConfig` to control batch sizes, authosharding, etc. column_specs: Use :class:`~.postgres_common.ColumnSpecsBuilder` to configure how embeddings and metadata are written a database schema. If None, uses default EmbeddableItem schema. conflict_resolution: Optional :class:`~.postgres_common.ConflictResolution` strategy for handling insert conflicts. ON CONFLICT DO NOTHING by default. Examples: Basic usage with default schema: >>> config = AlloyDBVectorWriterConfig( ... connection_config=AlloyDBConnectionConfig(...), ... table_name='embeddings' ... ) Simple case with default schema: >>> config = PostgresVectorWriterConfig( ... connection_config=ConnectionConfig(...), ... table_name='embeddings' ... ) Custom schema with metadata fields: >>> specs = (ColumnSpecsBuilder() ... .with_id_spec(column_name="my_id_column") ... .with_embedding_spec(column_name="embedding_vec") ... .add_metadata_field(field="source", column_name="src") ... .add_metadata_field( ... "timestamp", ... column_name="created_at", ... sql_typecast="::timestamp" ... ) ... .build()) Minimal schema (only ID + embedding written) >>> column_specs = (ColumnSpecsBuilder() ... .with_id_spec() ... .with_embedding_spec() ... .build()) >>> config = AlloyDBVectorWriterConfig( ... connection_config=AlloyDBConnectionConfig(...), ... table_name='embeddings', ... column_specs=specs ... ) """ super().__init__( connection_config=connection_config.to_connection_config(), write_config=write_config, table_name=table_name, column_specs=column_specs, conflict_resolution=conflict_resolution) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/alloydb_it_test.py0000644000175100017510000011441615206113032024577 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import json import logging import os import secrets import time import unittest from typing import NamedTuple import psycopg2 import pytest import apache_beam as beam from apache_beam.coders import registry from apache_beam.coders.row_coder import RowCoder from apache_beam.io.jdbc import ReadFromJdbc from apache_beam.ml.rag.ingestion import test_utils from apache_beam.ml.rag.ingestion.alloydb import AlloyDBLanguageConnectorConfig from apache_beam.ml.rag.ingestion.alloydb import AlloyDBVectorWriterConfig from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteTransform from apache_beam.ml.rag.ingestion.postgres_common import ColumnSpec from apache_beam.ml.rag.ingestion.postgres_common import ColumnSpecsBuilder from apache_beam.ml.rag.ingestion.postgres_common import ConflictResolution from apache_beam.ml.rag.ingestion.postgres_common import chunk_embedding_fn from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import Embedding from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to _LOGGER = logging.getLogger(__name__) CustomSpecsRow = NamedTuple( 'CustomSpecsRow', [ ('custom_id', str), # For id_spec test ('embedding_vec', list[float]), # For embedding_spec test ('content_col', str), # For content_spec test ('metadata', str) ]) registry.register_coder(CustomSpecsRow, RowCoder) MetadataConflictRow = NamedTuple( 'MetadataConflictRow', [ ('id', str), ('source', str), # For metadata_spec and composite key ('timestamp', str), # For metadata_spec and composite key ('content', str), ('embedding', list[float]), ('metadata', str) ]) registry.register_coder(MetadataConflictRow, RowCoder) @pytest.mark.uses_gcp_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_JARS'), "EXPANSION_JARS environment var is not provided, " "indicating that jars have not been built") @unittest.skipUnless( os.environ.get('ALLOYDB_PASSWORD'), "ALLOYDB_PASSWORD environment var is not provided") class AlloydbVectorWriterConfigTest(unittest.TestCase): ALLOYDB_TABLE_PREFIX = 'python_rag_postgres_' @classmethod def setUpClass(cls): cls.host = os.environ.get('ALLOYDB_HOST', '10.119.0.22') cls.port = os.environ.get('ALLOYDB_PORT', '5432') cls.database = os.environ.get('ALLOYDB_DATABASE', 'postgres') cls.username = os.environ.get('ALLOYDB_USERNAME', 'postgres') if not os.environ.get('ALLOYDB_PASSWORD'): raise ValueError('ALLOYDB_PASSWORD env not set') cls.password = os.environ.get('ALLOYDB_PASSWORD') # Create unique table name suffix cls.table_suffix = '%d%s' % (int(time.time()), secrets.token_hex(3)) # Setup database connection cls.conn = psycopg2.connect( host=cls.host, port=cls.port, database=cls.database, user=cls.username, password=cls.password) cls.conn.autocommit = True def skip_if_dataflow_runner(self): if self._runner and "dataflowrunner" in self._runner.lower(): self.skipTest( "Skipping some tests on Dataflow Runner to avoid bloat and timeouts") def setUp(self): self.write_test_pipeline = TestPipeline(is_integration_test=True) self.read_test_pipeline = TestPipeline(is_integration_test=True) self.write_test_pipeline2 = TestPipeline(is_integration_test=True) self.read_test_pipeline2 = TestPipeline(is_integration_test=True) self._runner = type(self.read_test_pipeline.runner).__name__ self.default_table_name = f"{self.ALLOYDB_TABLE_PREFIX}" \ f"{self.table_suffix}" self.custom_table_name = f"{self.ALLOYDB_TABLE_PREFIX}" \ f"_custom_{self.table_suffix}" self.metadata_conflicts_table = f"{self.ALLOYDB_TABLE_PREFIX}" \ f"_meta_conf_{self.table_suffix}" self.jdbc_url = f'jdbc:postgresql://{self.host}:{self.port}/{self.database}' # Create test table with self.conn.cursor() as cursor: cursor.execute( f""" CREATE TABLE {self.default_table_name} ( id TEXT PRIMARY KEY, embedding VECTOR({test_utils.VECTOR_SIZE}), content TEXT, metadata JSONB ) """) cursor.execute( f""" CREATE TABLE {self.custom_table_name} ( custom_id TEXT PRIMARY KEY, embedding_vec VECTOR(2), content_col TEXT, metadata JSONB ) """) cursor.execute( f""" CREATE TABLE {self.metadata_conflicts_table} ( id TEXT, source TEXT, timestamp TIMESTAMP, content TEXT, embedding VECTOR(2), PRIMARY KEY (id), UNIQUE (source, timestamp) ) """) _LOGGER.info("Created table %s", self.default_table_name) def tearDown(self): # Drop test table with self.conn.cursor() as cursor: cursor.execute(f"DROP TABLE IF EXISTS {self.default_table_name}") cursor.execute(f"DROP TABLE IF EXISTS {self.custom_table_name}") cursor.execute(f"DROP TABLE IF EXISTS {self.metadata_conflicts_table}") _LOGGER.info("Dropped table %s", self.default_table_name) @classmethod def tearDownClass(cls): if hasattr(cls, 'conn'): cls.conn.close() def test_default_schema(self): """Test basic write with default schema.""" connection_config = AlloyDBLanguageConnectorConfig( username=self.username, password=self.password, database_name=self.database, instance_name="projects/apache-beam-testing/locations/us-central1/\ clusters/testing-psc/instances/testing-psc-1", ip_type="PSC") config = AlloyDBVectorWriterConfig( connection_config=connection_config, table_name=self.default_table_name) # Create test chunks num_records = 1500 sample_size = min(500, num_records // 2) # Generate test chunks chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records) # Run pipeline and verify self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = (p | beam.Create(chunks) | config.create_write_transform()) self.read_test_pipeline.not_use_test_runner_api = True # Read pipeline to verify read_query = f""" SELECT CAST(id AS VARCHAR(255)), CAST(content AS VARCHAR(255)), CAST(embedding AS text), CAST(metadata AS text) FROM {self.default_table_name} """ # Read and verify pipeline with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=connection_config.to_connection_config().jdbc_url, username=self.username, password=self.password, query=read_query, classpath=connection_config.additional_jdbc_args()['classpath'])) count_result = rows | "Count All" >> beam.combiners.Count.Globally() assert_that(count_result, equal_to([num_records]), label='count_check') chunks = (rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk)) chunk_hashes = chunks | "Hash Chunks" >> beam.CombineGlobally( test_utils.HashingFn()) assert_that( chunk_hashes, equal_to([test_utils.generate_expected_hash(num_records)]), label='hash_check') # Sample validation first_n = ( chunks | "Key on Index" >> beam.Map(test_utils.key_on_id) | f"Get First {sample_size}" >> beam.transforms.combiners.Top.Of( sample_size, key=lambda x: x[0], reverse=True) | "Remove Keys 1" >> beam.Map(lambda xs: [x[1] for x in xs])) expected_first_n = test_utils.ChunkTestUtils.get_expected_values( 0, sample_size) assert_that( first_n, equal_to([expected_first_n]), label=f"first_{sample_size}_check") last_n = ( chunks | "Key on Index 2" >> beam.Map(test_utils.key_on_id) | f"Get Last {sample_size}" >> beam.transforms.combiners.Top.Of( sample_size, key=lambda x: x[0]) | "Remove Keys 2" >> beam.Map(lambda xs: [x[1] for x in xs])) expected_last_n = test_utils.ChunkTestUtils.get_expected_values( num_records - sample_size, num_records)[::-1] assert_that( last_n, equal_to([expected_last_n]), label=f"last_{sample_size}_check") def test_language_connector(self): """Test language connector.""" self.skip_if_dataflow_runner() connection_config = AlloyDBLanguageConnectorConfig( username=self.username, password=self.password, database_name=self.database, instance_name="projects/apache-beam-testing/locations/us-central1/\ clusters/testing-psc/instances/testing-psc-1", ip_type="PSC") writer_config = AlloyDBVectorWriterConfig( connection_config=connection_config, table_name=self.default_table_name) # Create test chunks num_records = 150 sample_size = min(500, num_records // 2) chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records) self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = ( p | beam.Create(chunks) | VectorDatabaseWriteTransform(writer_config)) self.read_test_pipeline.not_use_test_runner_api = True read_query = f""" SELECT CAST(id AS VARCHAR(255)), CAST(content AS VARCHAR(255)), CAST(embedding AS text), CAST(metadata AS text) FROM {self.default_table_name} """ with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=connection_config.to_connection_config().jdbc_url, username=self.username, password=self.password, query=read_query, classpath=connection_config.additional_jdbc_args()['classpath'])) count_result = rows | "Count All" >> beam.combiners.Count.Globally() assert_that(count_result, equal_to([num_records]), label='count_check') chunks = (rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk)) chunk_hashes = chunks | "Hash Chunks" >> beam.CombineGlobally( test_utils.HashingFn()) assert_that( chunk_hashes, equal_to([test_utils.generate_expected_hash(num_records)]), label='hash_check') # Sample validation first_n = ( chunks | "Key on Index" >> beam.Map(test_utils.key_on_id) | f"Get First {sample_size}" >> beam.transforms.combiners.Top.Of( sample_size, key=lambda x: x[0], reverse=True) | "Remove Keys 1" >> beam.Map(lambda xs: [x[1] for x in xs])) expected_first_n = test_utils.ChunkTestUtils.get_expected_values( 0, sample_size) assert_that( first_n, equal_to([expected_first_n]), label=f"first_{sample_size}_check") last_n = ( chunks | "Key on Index 2" >> beam.Map(test_utils.key_on_id) | f"Get Last {sample_size}" >> beam.transforms.combiners.Top.Of( sample_size, key=lambda x: x[0]) | "Remove Keys 2" >> beam.Map(lambda xs: [x[1] for x in xs])) expected_last_n = test_utils.ChunkTestUtils.get_expected_values( num_records - sample_size, num_records)[::-1] assert_that( last_n, equal_to([expected_last_n]), label=f"last_{sample_size}_check") def test_custom_specs(self): """Test custom specifications for ID, embedding, and content.""" self.skip_if_dataflow_runner() num_records = 20 specs = ( ColumnSpecsBuilder().add_custom_column_spec( ColumnSpec.text( column_name="custom_id", value_fn=lambda chunk: f"timestamp_{chunk.metadata.get('timestamp', '')}") ).add_custom_column_spec( ColumnSpec.vector( column_name="embedding_vec", value_fn=chunk_embedding_fn)).add_custom_column_spec( ColumnSpec.text( column_name="content_col", value_fn=lambda chunk: f"{len(chunk.content.text)}:{chunk.content.text}")). with_metadata_spec().build()) connection_config = AlloyDBLanguageConnectorConfig( username=self.username, password=self.password, database_name=self.database, instance_name="projects/apache-beam-testing/locations/us-central1/\ clusters/testing-psc/instances/testing-psc-1", ip_type="PSC") writer_config = AlloyDBVectorWriterConfig( connection_config=connection_config, table_name=self.custom_table_name, column_specs=specs) # Generate test chunks test_chunks = [ Chunk( id=str(i), content=Content(text=f"content_{i}"), embedding=Embedding(dense_embedding=[float(i), float(i + 1)]), metadata={"timestamp": f"2024-02-02T{i:02d}:00:00"}) for i in range(num_records) ] # Write pipeline self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = ( p | beam.Create(test_chunks) | VectorDatabaseWriteTransform(writer_config)) # Read and verify read_query = f""" SELECT CAST(custom_id AS VARCHAR(255)), CAST(embedding_vec AS text), CAST(content_col AS VARCHAR(255)), CAST(metadata AS text) FROM {self.custom_table_name} ORDER BY custom_id """ # Convert BeamRow back to Chunk def custom_row_to_chunk(row): # Extract timestamp from custom_id timestamp = row.custom_id.split('timestamp_')[1] # Extract index from timestamp i = int(timestamp.split('T')[1][:2]) # Parse embedding vector embedding_list = [ float(x) for x in row.embedding_vec.strip('[]').split(',') ] # Extract content from length-prefixed format content = row.content_col.split(':', 1)[1] return Chunk( id=str(i), content=Content(text=content), embedding=Embedding(dense_embedding=embedding_list), metadata=json.loads(row.metadata)) self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.custom_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=connection_config.to_connection_config().jdbc_url, username=self.username, password=self.password, query=read_query, classpath=connection_config.additional_jdbc_args()['classpath'])) # Verify count count_result = rows | "Count All" >> beam.combiners.Count.Globally() assert_that(count_result, equal_to([num_records]), label='count_check') chunks = rows | "To Chunks" >> beam.Map(custom_row_to_chunk) assert_that(chunks, equal_to(test_chunks), label='chunks_check') def test_defaults_with_args_specs(self): """Test custom specifications for ID, embedding, and content.""" self.skip_if_dataflow_runner() num_records = 20 specs = ( ColumnSpecsBuilder().with_id_spec( column_name="custom_id", python_type=int, convert_fn=lambda x: int(x), sql_typecast="::text").with_content_spec( column_name="content_col", convert_fn=lambda x: f"{len(x)}:{x}", ).with_embedding_spec( column_name="embedding_vec").with_metadata_spec().build()) connection_config = AlloyDBLanguageConnectorConfig( username=self.username, password=self.password, database_name=self.database, instance_name="projects/apache-beam-testing/locations/us-central1/\ clusters/testing-psc/instances/testing-psc-1", ip_type="PSC") writer_config = AlloyDBVectorWriterConfig( connection_config=connection_config, table_name=self.custom_table_name, column_specs=specs) # Generate test chunks test_chunks = [ Chunk( id=str(i), content=Content(text=f"content_{i}"), embedding=Embedding(dense_embedding=[float(i), float(i + 1)]), metadata={"timestamp": f"2024-02-02T{i:02d}:00:00"}) for i in range(num_records) ] # Write pipeline self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = ( p | beam.Create(test_chunks) | VectorDatabaseWriteTransform(writer_config)) # Read and verify read_query = f""" SELECT CAST(custom_id AS VARCHAR(255)), CAST(embedding_vec AS text), CAST(content_col AS VARCHAR(255)), CAST(metadata AS text) FROM {self.custom_table_name} ORDER BY custom_id """ # Convert BeamRow back to Chunk def custom_row_to_chunk(row): # Parse embedding vector embedding_list = [ float(x) for x in row.embedding_vec.strip('[]').split(',') ] # Extract content from length-prefixed format content = row.content_col.split(':', 1)[1] return Chunk( id=row.custom_id, content=Content(text=content), embedding=Embedding(dense_embedding=embedding_list), metadata=json.loads(row.metadata)) self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.custom_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=connection_config.to_connection_config().jdbc_url, username=self.username, password=self.password, query=read_query, classpath=connection_config.additional_jdbc_args()['classpath'])) # Verify count count_result = rows | "Count All" >> beam.combiners.Count.Globally() assert_that(count_result, equal_to([num_records]), label='count_check') chunks = rows | "To Chunks" >> beam.Map(custom_row_to_chunk) assert_that(chunks, equal_to(test_chunks), label='chunks_check') def test_default_id_embedding_specs(self): """Test with only default id and embedding specs, others set to None.""" self.skip_if_dataflow_runner() num_records = 20 connection_config = AlloyDBLanguageConnectorConfig( username=self.username, password=self.password, database_name=self.database, instance_name="projects/apache-beam-testing/locations/us-central1/\ clusters/testing-psc/instances/testing-psc-1", ip_type="PSC") specs = ( ColumnSpecsBuilder().with_id_spec() # Use default id spec .with_embedding_spec() # Use default embedding spec .build()) writer_config = AlloyDBVectorWriterConfig( connection_config=connection_config, table_name=self.default_table_name, column_specs=specs) # Generate test chunks test_chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records) # Write pipeline self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = ( p | beam.Create(test_chunks) | VectorDatabaseWriteTransform(writer_config)) # Read and verify only id and embedding read_query = f""" SELECT CAST(id AS VARCHAR(255)), CAST(embedding AS text) FROM {self.default_table_name} ORDER BY id """ self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=connection_config.to_connection_config().jdbc_url, username=self.username, password=self.password, query=read_query, classpath=connection_config.additional_jdbc_args()['classpath'])) chunks = rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk) # Create expected chunks with None values expected_chunks = test_utils.ChunkTestUtils.get_expected_values( 0, num_records) for chunk in expected_chunks: chunk.content.text = None chunk.metadata = {} assert_that(chunks, equal_to(expected_chunks), label='chunks_check') def test_metadata_spec_and_conflicts(self): """Test metadata specification and conflict resolution.""" self.skip_if_dataflow_runner() num_records = 20 specs = ( ColumnSpecsBuilder().with_id_spec().with_embedding_spec(). with_content_spec().add_metadata_field( field="source", column_name="source", python_type=str, sql_typecast=None # Plain text field ).add_metadata_field( field="timestamp", python_type=str, sql_typecast="::timestamp").build()) # Conflict resolution on source+timestamp conflict_resolution = ConflictResolution( on_conflict_fields=["source", "timestamp"], action="UPDATE", update_fields=["embedding", "content"]) connection_config = AlloyDBLanguageConnectorConfig( username=self.username, password=self.password, database_name=self.database, instance_name="projects/apache-beam-testing/locations/us-central1/\ clusters/testing-psc/instances/testing-psc-1", ip_type="PSC") writer_config = AlloyDBVectorWriterConfig( connection_config=connection_config, table_name=self.metadata_conflicts_table, column_specs=specs, conflict_resolution=conflict_resolution) # Generate initial test chunks initial_chunks = [ Chunk( id=str(i), content=Content(text=f"content_{i}"), embedding=Embedding(dense_embedding=[float(i), float(i + 1)]), metadata={ "source": "source_A", "timestamp": f"2024-02-02T{i:02d}:00:00" }) for i in range(num_records) ] # Write initial chunks self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = ( p | "Write Initial" >> beam.Create(initial_chunks) | VectorDatabaseWriteTransform(writer_config)) # Generate conflicting chunks (same source+timestamp, different content) conflicting_chunks = [ Chunk( id=f"new_{i}", content=Content(text=f"updated_content_{i}"), embedding=Embedding( dense_embedding=[float(i) * 2, float(i + 1) * 2]), metadata={ "source": "source_A", "timestamp": f"2024-02-02T{i:02d}:00:00" }) for i in range(num_records) ] # Write conflicting chunks self.write_test_pipeline2.not_use_test_runner_api = True with self.write_test_pipeline2 as p: _ = ( p | "Write Conflicts" >> beam.Create(conflicting_chunks) | VectorDatabaseWriteTransform(writer_config)) # Read and verify read_query = f""" SELECT CAST(id AS VARCHAR(255)), CAST(embedding AS text), CAST(content AS VARCHAR(255)), CAST(source AS VARCHAR(255)), CAST(timestamp AS VARCHAR(255)) FROM {self.metadata_conflicts_table} ORDER BY timestamp, id """ # Expected chunks after conflict resolution expected_chunks = [ Chunk( id=str(i), content=Content(text=f"updated_content_{i}"), embedding=Embedding( dense_embedding=[float(i) * 2, float(i + 1) * 2]), metadata={ "source": "source_A", "timestamp": f"2024-02-02T{i:02d}:00:00" }) for i in range(num_records) ] def metadata_row_to_chunk(row): return Chunk( id=row.id, content=Content(text=row.content), embedding=Embedding( dense_embedding=[ float(x) for x in row.embedding.strip('[]').split(',') ]), metadata={ "source": row.source, "timestamp": row.timestamp.replace(' ', 'T') }) self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.metadata_conflicts_table, driver_class_name="org.postgresql.Driver", jdbc_url=connection_config.to_connection_config().jdbc_url, username=self.username, password=self.password, query=read_query, classpath=connection_config.additional_jdbc_args()['classpath'])) chunks = rows | "To Chunks" >> beam.Map(metadata_row_to_chunk) assert_that(chunks, equal_to(expected_chunks), label='chunks_check') def test_conflict_resolution_update(self): """Test conflict resolution with UPDATE action.""" self.skip_if_dataflow_runner() num_records = 20 connection_config = AlloyDBLanguageConnectorConfig( username=self.username, password=self.password, database_name=self.database, instance_name="projects/apache-beam-testing/locations/us-central1/\ clusters/testing-psc/instances/testing-psc-1", ip_type="PSC") conflict_resolution = ConflictResolution( on_conflict_fields="id", action="UPDATE", update_fields=["embedding", "content"]) config = AlloyDBVectorWriterConfig( connection_config=connection_config, table_name=self.default_table_name, conflict_resolution=conflict_resolution) # Generate initial test chunks test_chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records) self.write_test_pipeline.not_use_test_runner_api = True # Insert initial test chunks with self.write_test_pipeline as p: _ = ( p | "Create initial chunks" >> beam.Create(test_chunks) | "Write initial chunks" >> config.create_write_transform()) read_query = f""" SELECT CAST(id AS VARCHAR(255)), CAST(content AS VARCHAR(255)), CAST(embedding AS text), CAST(metadata AS text) FROM {self.default_table_name} ORDER BY id desc """ self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=connection_config.to_connection_config().jdbc_url, username=self.username, password=self.password, query=read_query, classpath=connection_config.additional_jdbc_args()['classpath'])) chunks = ( rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk) | "Key on Index" >> beam.Map(test_utils.key_on_id) | "Get First 500" >> beam.transforms.combiners.Top.Of( num_records, key=lambda x: x[0], reverse=True) | "Remove Keys 1" >> beam.Map(lambda xs: [x[1] for x in xs])) assert_that( chunks, equal_to([test_chunks]), label='original_chunks_check') updated_chunks = test_utils.ChunkTestUtils.get_expected_values( 0, num_records, content_prefix="Newcontent", seed_multiplier=2) self.write_test_pipeline2.not_use_test_runner_api = True with self.write_test_pipeline2 as p: _ = ( p | "Create updated Chunks" >> beam.Create(updated_chunks) | "Write updated Chunks" >> config.create_write_transform()) self.read_test_pipeline2.not_use_test_runner_api = True with self.read_test_pipeline2 as p: rows = ( p | "Read Updated chunks" >> ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=connection_config.to_connection_config().jdbc_url, username=self.username, password=self.password, query=read_query, classpath=connection_config.additional_jdbc_args()['classpath'])) chunks = ( rows | "To Chunks 2" >> beam.Map(test_utils.row_to_chunk) | "Key on Index 2" >> beam.Map(test_utils.key_on_id) | "Get First 500 2" >> beam.transforms.combiners.Top.Of( num_records, key=lambda x: x[0], reverse=True) | "Remove Keys 2" >> beam.Map(lambda xs: [x[1] for x in xs])) assert_that( chunks, equal_to([updated_chunks]), label='updated_chunks_check') def test_conflict_resolution_default_ignore(self): """Test conflict resolution with default.""" self.skip_if_dataflow_runner() num_records = 20 connection_config = AlloyDBLanguageConnectorConfig( username=self.username, password=self.password, database_name=self.database, instance_name="projects/apache-beam-testing/locations/us-central1/\ clusters/testing-psc/instances/testing-psc-1", ip_type="PSC") config = AlloyDBVectorWriterConfig( connection_config=connection_config, table_name=self.default_table_name) # Generate initial test chunks test_chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records) self.write_test_pipeline.not_use_test_runner_api = True # Insert initial test chunks with self.write_test_pipeline as p: _ = ( p | "Create initial chunks" >> beam.Create(test_chunks) | "Write initial chunks" >> config.create_write_transform()) read_query = f""" SELECT CAST(id AS VARCHAR(255)), CAST(content AS VARCHAR(255)), CAST(embedding AS text), CAST(metadata AS text) FROM {self.default_table_name} ORDER BY id desc """ self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=connection_config.to_connection_config().jdbc_url, username=self.username, password=self.password, query=read_query, classpath=connection_config.additional_jdbc_args()['classpath'])) chunks = ( rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk) | "Key on Index" >> beam.Map(test_utils.key_on_id) | "Get First 500" >> beam.transforms.combiners.Top.Of( num_records, key=lambda x: x[0], reverse=True) | "Remove Keys 1" >> beam.Map(lambda xs: [x[1] for x in xs])) assert_that( chunks, equal_to([test_chunks]), label='original_chunks_check') updated_chunks = test_utils.ChunkTestUtils.get_expected_values( 0, num_records, content_prefix="Newcontent", seed_multiplier=2) self.write_test_pipeline2.not_use_test_runner_api = True with self.write_test_pipeline2 as p: _ = ( p | "Create updated Chunks" >> beam.Create(updated_chunks) | "Write updated Chunks" >> config.create_write_transform()) self.read_test_pipeline2.not_use_test_runner_api = True with self.read_test_pipeline2 as p: rows = ( p | "Read Updated chunks" >> ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=connection_config.to_connection_config().jdbc_url, username=self.username, password=self.password, query=read_query, classpath=connection_config.additional_jdbc_args()['classpath'])) chunks = ( rows | "To Chunks 2" >> beam.Map(test_utils.row_to_chunk) | "Key on Index 2" >> beam.Map(test_utils.key_on_id) | "Get First 500 2" >> beam.transforms.combiners.Top.Of( num_records, key=lambda x: x[0], reverse=True) | "Remove Keys 2" >> beam.Map(lambda xs: [x[1] for x in xs])) assert_that(chunks, equal_to([test_chunks]), label='updated_chunks_check') def test_conflict_resolution_default_update_fields(self): """Test conflict resolution with default update fields (all non-conflict fields).""" self.skip_if_dataflow_runner() num_records = 20 connection_config = AlloyDBLanguageConnectorConfig( username=self.username, password=self.password, database_name=self.database, instance_name="projects/apache-beam-testing/locations/us-central1/\ clusters/testing-psc/instances/testing-psc-1", ip_type="PSC") # Create a conflict resolution with only the conflict field specified # No update_fields specified - should default to all non-conflict fields conflict_resolution = ConflictResolution( on_conflict_fields="id", action="UPDATE") config = AlloyDBVectorWriterConfig( connection_config=connection_config, table_name=self.default_table_name, conflict_resolution=conflict_resolution) # Generate initial test chunks test_chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records) self.write_test_pipeline.not_use_test_runner_api = True # Insert initial test chunks with self.write_test_pipeline as p: _ = ( p | "Create initial chunks" >> beam.Create(test_chunks) | "Write initial chunks" >> config.create_write_transform()) # Verify initial data was written correctly read_query = f""" SELECT CAST(id AS VARCHAR(255)), CAST(content AS VARCHAR(255)), CAST(embedding AS text), CAST(metadata AS text) FROM {self.default_table_name} ORDER BY id desc """ self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=connection_config.to_connection_config().jdbc_url, username=self.username, password=self.password, query=read_query, classpath=connection_config.additional_jdbc_args()['classpath'])) chunks = ( rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk) | "Key on Index" >> beam.Map(test_utils.key_on_id) | "Get First 500" >> beam.transforms.combiners.Top.Of( num_records, key=lambda x: x[0], reverse=True) | "Remove Keys 1" >> beam.Map(lambda xs: [x[1] for x in xs])) assert_that( chunks, equal_to([test_chunks]), label='original_chunks_check') # Create updated chunks with same IDs but different content, embedding, and # metadata updated_chunks = [] for i in range(num_records): original_chunk = test_chunks[i] updated_chunk = Chunk( id=original_chunk.id, content=Content(text=f"Updated content {i}"), embedding=Embedding( dense_embedding=[float(i * 2), float(i * 2 + 1)] + [0.0] * (test_utils.VECTOR_SIZE - 2)), metadata={ "updated": "true", "timestamp": "2024-02-25" }) updated_chunks.append(updated_chunk) # Write updated chunks - should update all non-conflict fields self.write_test_pipeline2.not_use_test_runner_api = True with self.write_test_pipeline2 as p: _ = ( p | "Create updated Chunks" >> beam.Create(updated_chunks) | "Write updated Chunks" >> config.create_write_transform()) # Read and verify that all non-conflict fields were updated self.read_test_pipeline2.not_use_test_runner_api = True with self.read_test_pipeline2 as p: rows = ( p | "Read Updated chunks" >> ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=connection_config.to_connection_config().jdbc_url, username=self.username, password=self.password, query=read_query, classpath=connection_config.additional_jdbc_args()['classpath'])) chunks = ( rows | "To Chunks 2" >> beam.Map(test_utils.row_to_chunk) | "Key on Index 2" >> beam.Map(test_utils.key_on_id) | "Get First 500 2" >> beam.transforms.combiners.Top.Of( num_records, key=lambda x: x[0], reverse=True) | "Remove Keys 2" >> beam.Map(lambda xs: [x[1] for x in xs])) # Verify that all non-conflict fields were updated assert_that( chunks, equal_to([updated_chunks]), label='updated_chunks_check') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/base.py0000644000175100017510000001065315206113032022326 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from abc import ABC from abc import abstractmethod from typing import Any import apache_beam as beam from apache_beam.ml.rag.types import EmbeddableItem class VectorDatabaseWriteConfig(ABC): """Abstract base class for vector database configurations in RAG pipelines. VectorDatabaseWriteConfig defines the interface for configuring vector database writes in RAG pipelines. Implementations should provide database-specific configuration and create appropriate write transforms. The configuration flow: 1. Subclass provides database-specific configuration (table names, etc) 2. create_write_transform() creates appropriate PTransform for writing 3. Transform handles converting EmbeddableItems to database-specific format Example implementation: >>> class BigQueryVectorWriterConfig(VectorDatabaseWriteConfig): ... def __init__(self, table: str): ... self.embedding_column = embedding_column ... ... def create_write_transform(self): ... return beam.io.WriteToBigQuery( ... table=self.table ... ) """ @abstractmethod def create_write_transform(self) -> beam.PTransform[EmbeddableItem, Any]: """Creates a PTransform that writes embeddings to the vector database. Returns: A PTransform that accepts PCollection[EmbeddableItem] and writes the embeddings and metadata to the configured vector database. The transform should handle: - Converting EmbeddableItem format to database schema - Setting up database connection/client - Writing with appropriate batching/error handling """ raise NotImplementedError(type(self)) class VectorDatabaseWriteTransform(beam.PTransform): """A PTransform for writing embedded chunks to vector databases. This transform uses a VectorDatabaseWriteConfig to write chunks with embeddings to vector database. It handles validating the config and applying the database-specific write transform. Example usage: >>> config = BigQueryVectorConfig( ... table='project.dataset.embeddings', ... embedding_column='embedding' ... ) >>> >>> with beam.Pipeline() as p: ... items = p | beam.Create([...]) # PCollection[EmbeddableItem] ... items | VectorDatabaseWriteTransform(config) Args: database_config: Configuration for the target vector database. Must be a subclass of VectorDatabaseWriteConfig that implements create_write_transform(). Raises: TypeError: If database_config is not a VectorDatabaseWriteConfig instance. """ def __init__(self, database_config: VectorDatabaseWriteConfig): """Initialize transform with database config. Args: database_config: Configuration for target vector database. """ if not isinstance(database_config, VectorDatabaseWriteConfig): raise TypeError( f"database_config must be VectorDatabaseWriteConfig, " f"got {type(database_config)}") self.database_config = database_config def expand( self, pcoll: beam.PCollection[EmbeddableItem] ) -> beam.PTransform[EmbeddableItem, Any]: """Creates and applies the database-specific write transform. Args: pcoll: PCollection of EmbeddableItems with embeddings to write to the vector database. Each EmbeddableItem must have: - An embedding - An ID - Metadata used to filter results as specified by database config Returns: Result of writing to database (implementation specific). """ write_transform = self.database_config.create_write_transform() return pcoll | write_transform ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/base_test.py0000644000175100017510000000536115206113032023365 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import unittest import apache_beam as beam from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteConfig from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteTransform from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import Embedding from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class MockWriteTransform(beam.PTransform): """Mock transform that returns element.""" def expand(self, pcoll): return pcoll | beam.Map(lambda x: x) class MockDatabaseConfig(VectorDatabaseWriteConfig): """Mock database config for testing.""" def __init__(self): self.write_transform = MockWriteTransform() def create_write_transform(self) -> beam.PTransform: return self.write_transform class VectorDatabaseBaseTest(unittest.TestCase): def test_write_transform_creation(self): """Test that write transform is created correctly.""" config = MockDatabaseConfig() transform = VectorDatabaseWriteTransform(config) self.assertEqual(transform.database_config, config) def test_pipeline_integration(self): """Test writing through pipeline.""" test_data = [ Chunk( content=Content(text="foo"), id="1", embedding=Embedding(dense_embedding=[0.1, 0.2])), Chunk( content=Content(text="bar"), id="2", embedding=Embedding(dense_embedding=[0.3, 0.4])) ] with TestPipeline() as p: result = ( p | beam.Create(test_data) | VectorDatabaseWriteTransform(MockDatabaseConfig())) # Verify data was written assert_that(result, equal_to(test_data)) def test_invalid_config(self): """Test error handling for invalid config.""" with self.assertRaises(TypeError): VectorDatabaseWriteTransform(None) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/bigquery.py0000644000175100017510000001637515206113032023252 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import warnings from collections.abc import Callable from typing import Any from typing import Optional import apache_beam as beam from apache_beam.io.gcp.bigquery_tools import beam_row_from_dict from apache_beam.io.gcp.bigquery_tools import get_beam_typehints_from_tableschema from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteConfig from apache_beam.ml.rag.types import EmbeddableItem from apache_beam.typehints.row_type import RowTypeConstraint EmbeddableToDictFn = Callable[[EmbeddableItem], dict[str, any]] # Backward compatibility alias. ChunkToDictFn = EmbeddableToDictFn class SchemaConfig: def __init__( self, schema: dict, embeddable_to_dict_fn: Optional[EmbeddableToDictFn] = None, **kwargs): """Configuration for custom BigQuery schema and row conversion. Allows overriding the default schema and row conversion logic for BigQuery vector storage. This enables custom table schemas beyond the default id/embedding/content/metadata structure. Args: schema: BigQuery TableSchema dict defining the table structure. embeddable_to_dict_fn: Function that converts an EmbeddableItem to a dict matching the schema. Takes an EmbeddableItem and returns Dict[str, Any] with keys matching schema fields. Example with custom schema: >>> schema_config = SchemaConfig( ... schema={ ... 'fields': [ ... {'name': 'id', 'type': 'STRING'}, ... {'name': 'embedding', 'type': 'FLOAT64', 'mode': 'REPEATED'}, ... {'name': 'source_url', 'type': 'STRING'} ... ] ... }, ... embeddable_to_dict_fn=lambda item: { ... 'id': item.id, ... 'embedding': item.embedding.dense_embedding, ... 'source_url': item.metadata.get('url') ... } ... ) """ self.schema = schema if 'chunk_to_dict_fn' in kwargs: warnings.warn( "chunk_to_dict_fn is deprecated, use embeddable_to_dict_fn", DeprecationWarning, stacklevel=2) embeddable_to_dict_fn = kwargs.pop('chunk_to_dict_fn') if kwargs: raise TypeError(f"Unexpected keyword arguments: {', '.join(kwargs)}") if embeddable_to_dict_fn is None: raise TypeError("SchemaConfig requires embeddable_to_dict_fn") self.embeddable_to_dict_fn = embeddable_to_dict_fn class BigQueryVectorWriterConfig(VectorDatabaseWriteConfig): def __init__( self, write_config: dict[str, Any], *, # Force keyword arguments schema_config: Optional[SchemaConfig] = None): """Configuration for writing vectors to BigQuery using managed transforms. Supports both default schema (id, embedding, content, metadata columns) and custom schemas through SchemaConfig. Example with default schema: >>> config = BigQueryVectorWriterConfig( ... write_config={'table': 'project.dataset.embeddings'}) Example with custom schema: >>> schema_config = SchemaConfig( ... schema={ ... 'fields': [ ... {'name': 'id', 'type': 'STRING'}, ... {'name': 'embedding', 'type': 'FLOAT64', 'mode': 'REPEATED'}, ... {'name': 'source_url', 'type': 'STRING'} ... ] ... }, ... embeddable_to_dict_fn=lambda item: { ... 'id': item.id, ... 'embedding': item.embedding.dense_embedding, ... 'source_url': item.metadata.get('url') ... } ... ) >>> config = BigQueryVectorWriterConfig( ... write_config={'table': 'project.dataset.embeddings'}, ... schema_config=schema_config ... ) Args: write_config: BigQuery write configuration dict. Must include 'table'. Other options like create_disposition, write_disposition can be specified. schema_config: Optional configuration for custom schema and row conversion. If not provided, uses default schema with id, embedding, content and metadata columns. Raises: ValueError: If write_config doesn't include table specification. """ if 'table' not in write_config: raise ValueError("write_config must be provided with 'table' specified") self.write_config = write_config self.schema_config = schema_config def create_write_transform(self) -> beam.PTransform: """Creates transform to write to BigQuery.""" return _WriteToBigQueryVectorDatabase(self) def _default_embeddable_to_dict_fn(item: EmbeddableItem): if item.embedding is None or item.embedding.dense_embedding is None: raise ValueError("EmbeddableItem must contain dense embedding") return { 'id': item.id, 'embedding': item.embedding.dense_embedding, 'content': item.content_string, 'metadata': [{ "key": k, "value": str(v) } for k, v in item.metadata.items()] } def _default_schema(): return { 'fields': [{ 'name': 'id', 'type': 'STRING' }, { 'name': 'embedding', 'type': 'FLOAT64', 'mode': 'REPEATED' }, { 'name': 'content', 'type': 'STRING' }, { 'name': 'metadata', 'type': 'RECORD', 'mode': 'REPEATED', 'fields': [{ 'name': 'key', 'type': 'STRING' }, { 'name': 'value', 'type': 'STRING' }] }] } class _WriteToBigQueryVectorDatabase(beam.PTransform): """Implementation of BigQuery vector database write. """ def __init__(self, config: BigQueryVectorWriterConfig): self.config = config def expand(self, pcoll: beam.PCollection[EmbeddableItem]): schema = ( self.config.schema_config.schema if self.config.schema_config else _default_schema()) embeddable_to_dict_fn = ( self.config.schema_config.embeddable_to_dict_fn if self.config.schema_config else _default_embeddable_to_dict_fn) return ( pcoll | "EmbeddableItem to dict" >> beam.Map(embeddable_to_dict_fn) | "EmbeddableItem dict to schema'd row" >> beam.Map( lambda embeddable_item_dict: beam_row_from_dict( row=embeddable_item_dict, schema=schema)).with_output_types( RowTypeConstraint.from_fields( get_beam_typehints_from_tableschema(schema))) | "Write to BigQuery" >> beam.managed.Write( beam.managed.BIGQUERY, config=self.config.write_config)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/bigquery_it_test.py0000644000175100017510000002136615206113032025001 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import os import secrets import time import unittest import hamcrest as hc import pytest import apache_beam as beam from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultMatcher from apache_beam.ml.rag.ingestion.bigquery import BigQueryVectorWriterConfig from apache_beam.ml.rag.ingestion.bigquery import SchemaConfig from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import Embedding from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.transforms.periodicsequence import PeriodicImpulse @pytest.mark.uses_gcp_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_JARS'), "EXPANSION_JARS environment var is not provided, " "indicating that jars have not been built") class BigQueryVectorWriterConfigTest(unittest.TestCase): BIG_QUERY_DATASET_ID = 'python_rag_bigquery_' def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self._runner = type(self.test_pipeline.runner).__name__ self.project = self.test_pipeline.get_option('project') self.bigquery_client = BigQueryWrapper() self.dataset_id = '%s%d%s' % ( self.BIG_QUERY_DATASET_ID, int(time.time()), secrets.token_hex(3)) self.bigquery_client.get_or_create_dataset(self.project, self.dataset_id) _LOGGER = logging.getLogger(__name__) _LOGGER.info( "Created dataset %s in project %s", self.dataset_id, self.project) def tearDown(self): request = bigquery.BigqueryDatasetsDeleteRequest( projectId=self.project, datasetId=self.dataset_id, deleteContents=True) try: _LOGGER = logging.getLogger(__name__) _LOGGER.info( "Deleting dataset %s in project %s", self.dataset_id, self.project) self.bigquery_client.client.datasets.Delete(request) # Failing to delete a dataset should not cause a test failure. except Exception: _LOGGER = logging.getLogger(__name__) _LOGGER.debug( 'Failed to clean up dataset %s in project %s', self.dataset_id, self.project) def test_default_schema(self): table_name = 'python_default_schema_table' table_id = '{}.{}.{}'.format(self.project, self.dataset_id, table_name) config = BigQueryVectorWriterConfig(write_config={'table': table_id}) chunks = [ Chunk( id="1", embedding=Embedding(dense_embedding=[0.1, 0.2]), content=Content(text="foo"), metadata={"a": "b"}), Chunk( id="2", embedding=Embedding(dense_embedding=[0.3, 0.4]), content=Content(text="bar"), metadata={"c": "d"}) ] pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query="SELECT id, content, embedding, metadata FROM %s" % table_id, data=[("1", "foo", [0.1, 0.2], [{ "key": "a", "value": "b" }]), ("2", "bar", [0.3, 0.4], [{ "key": "c", "value": "d" }])]) ] args = self.test_pipeline.get_full_options_as_args( on_success_matcher=hc.all_of(*pipeline_verifiers)) with beam.Pipeline(argv=args) as p: _ = (p | beam.Create(chunks) | config.create_write_transform()) def test_default_schema_missing_embedding(self): table_name = 'python_default_schema_table' table_id = '{}.{}.{}'.format(self.project, self.dataset_id, table_name) config = BigQueryVectorWriterConfig(write_config={'table': table_id}) chunks = [ Chunk(id="1", content=Content(text="foo"), metadata={"a": "b"}), Chunk(id="2", content=Content(text="bar"), metadata={"c": "d"}) ] with self.assertRaisesRegex(Exception, "must contain dense embedding"): with beam.Pipeline() as p: _ = (p | beam.Create(chunks) | config.create_write_transform()) def test_custom_schema(self): table_name = 'python_custom_schema_table' table_id = '{}.{}.{}'.format(self.project, self.dataset_id, table_name) schema_config = SchemaConfig( schema={ 'fields': [{ 'name': 'id', 'type': 'STRING' }, { 'name': 'embedding', 'type': 'FLOAT64', 'mode': 'REPEATED' }, { 'name': 'source', 'type': 'STRING' }] }, embeddable_to_dict_fn=lambda chunk: { 'id': chunk.id, 'embedding': chunk.embedding.dense_embedding, 'source': chunk.metadata.get('source') }) config = BigQueryVectorWriterConfig( write_config={'table': table_id}, schema_config=schema_config) chunks = [ Chunk( id="1", embedding=Embedding(dense_embedding=[0.1, 0.2]), content=Content(text="foo content"), metadata={"source": "foo"}), Chunk( id="2", embedding=Embedding(dense_embedding=[0.3, 0.4]), content=Content(text="bar content"), metadata={"source": "bar"}) ] pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query="SELECT id, embedding, source FROM %s" % table_id, data=[("1", [0.1, 0.2], "foo"), ("2", [0.3, 0.4], "bar")]) ] args = self.test_pipeline.get_full_options_as_args( on_success_matcher=hc.all_of(*pipeline_verifiers)) with beam.Pipeline(argv=args) as p: _ = (p | beam.Create(chunks) | config.create_write_transform()) def test_streaming_default_schema(self): self.skip_if_not_dataflow_runner() table_name = 'python_streaming_default_schema_table' table_id = '{}.{}.{}'.format(self.project, self.dataset_id, table_name) config = BigQueryVectorWriterConfig(write_config={'table': table_id}) chunks = [ Chunk( id="1", embedding=Embedding(dense_embedding=[0.1, 0.2]), content=Content(text="foo"), metadata={"a": "b"}), Chunk( id="2", embedding=Embedding(dense_embedding=[0.3, 0.4]), content=Content(text="bar"), metadata={"c": "d"}), Chunk( id="3", embedding=Embedding(dense_embedding=[0.5, 0.6]), content=Content(text="foo"), metadata={"e": "f"}), Chunk( id="4", embedding=Embedding(dense_embedding=[0.7, 0.8]), content=Content(text="bar"), metadata={"g": "h"}) ] pipeline_verifiers = [ BigqueryFullResultMatcher( project=self.project, query="SELECT id, content, embedding, metadata FROM %s" % table_id, data=[("1", "foo", [0.1, 0.2], [{ "key": "a", "value": "b" }]), ("2", "bar", [0.3, 0.4], [{ "key": "c", "value": "d" }]), ("3", "foo", [0.5, 0.6], [{ "key": "e", "value": "f" }]), ("4", "bar", [0.7, 0.8], [{ "key": "g", "value": "h" }])]) ] args = self.test_pipeline.get_full_options_as_args( on_success_matcher=hc.all_of(*pipeline_verifiers), streaming=True, allow_unsafe_triggers=True) with beam.Pipeline(argv=args) as p: _ = ( p | PeriodicImpulse(0, 4, 1) | beam.Map(lambda t: chunks[t]) | config.create_write_transform()) def skip_if_not_dataflow_runner(self): # skip if dataflow runner is not specified if not self._runner or "dataflowrunner" not in self._runner.lower(): self.skipTest( "Streaming with exactly-once route has the requirement " "`beam:requirement:pardo:on_window_expiration:v1`, " "which is currently only supported by the Dataflow runner") if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/cloudsql.py0000644000175100017510000002264615206113032023247 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from dataclasses import asdict from dataclasses import dataclass from typing import Any from typing import Optional from apache_beam.ml.rag.ingestion import mysql from apache_beam.ml.rag.ingestion import mysql_common from apache_beam.ml.rag.ingestion import postgres from apache_beam.ml.rag.ingestion import postgres_common from apache_beam.ml.rag.ingestion.jdbc_common import ConnectionConfig from apache_beam.ml.rag.ingestion.jdbc_common import WriteConfig @dataclass class LanguageConnectorConfig: """Configuration options for CloudSQL Java language connector. Set parameters to connect connection to a CloudSQL instance using Java language connector connector. For details see https://github.com/GoogleCloudPlatform/cloud-sql-jdbc-socket-factory/blob/main/docs/jdbc.md Attributes: username: Database username. password: Database password. Can be empty string when using IAM. database_name: Name of the database to connect to. instance_name: Instance connection name. Format: '::' ip_type: Preferred order of IP types used to connect via a comma list of strings. enable_iam_auth: Whether to enable IAM authentication. Default is False target_principal: Optional service account to impersonate for connection. delegates: Optional list of service accounts for delegated impersonation. admin_service_endpoint: Optional custom API service endpoint. quota_project: Optional project ID for quota and billing. connection_properties: Optional JDBC connection properties dict. Example: {'ssl': 'true'} additional_properties: Additional properties to be added to the JDBC url. Example: {'someProperty': 'true'} """ username: str password: str database_name: str instance_name: str ip_types: Optional[list[str]] = None enable_iam_auth: bool = False target_principal: Optional[str] = None delegates: Optional[list[str]] = None quota_project: Optional[str] = None connection_properties: Optional[dict[str, str]] = None additional_properties: Optional[dict[str, Any]] = None def _base_jdbc_properties(self) -> dict[str, Any]: properties = {"cloudSqlInstance": self.instance_name} if self.ip_types: properties["ipTypes"] = ",".join(self.ip_types) if self.enable_iam_auth: properties["enableIamAuth"] = "true" if self.target_principal: properties["cloudSqlTargetPrincipal"] = self.target_principal if self.delegates: properties["cloudSqlDelegates"] = ",".join(self.delegates) if self.quota_project: properties["cloudSqlAdminQuotaProject"] = self.quota_project if self.additional_properties: properties.update(self.additional_properties) return properties def _build_jdbc_url(self, socketFactory, database_type): url = f"jdbc:{database_type}:///{self.database_name}?" properties = self._base_jdbc_properties() properties['socketFactory'] = socketFactory property_string = "&".join(f"{k}={v}" for k, v in properties.items()) return url + property_string def to_connection_config(self): return ConnectionConfig( jdbc_url=self.to_jdbc_url(), username=self.username, password=self.password, connection_properties=self.connection_properties, additional_jdbc_args=self.additional_jdbc_args()) def additional_jdbc_args(self) -> dict[str, list[Any]]: return {} @dataclass class _PostgresConnectorConfig(LanguageConnectorConfig): def to_jdbc_url(self) -> str: """Convert options to a properly formatted JDBC URL. Returns: JDBC URL string configured with all options. """ return self._build_jdbc_url( socketFactory="com.google.cloud.sql.postgres.SocketFactory", database_type="postgresql") def additional_jdbc_args(self) -> dict[str, list[Any]]: return { 'classpath': [ "org.postgresql:postgresql:42.2.16", "com.google.cloud.sql:postgres-socket-factory:1.25.0" ] } @classmethod def from_base_config(cls, config: LanguageConnectorConfig): return cls(**asdict(config)) class CloudSQLPostgresVectorWriterConfig(postgres.PostgresVectorWriterConfig): def __init__( self, connection_config: LanguageConnectorConfig, table_name: str, *, # pylint: disable=dangerous-default-value write_config: WriteConfig = WriteConfig(), column_specs: list[postgres_common.ColumnSpec] = postgres_common. ColumnSpecsBuilder.with_defaults().build(), conflict_resolution: Optional[ postgres_common.ConflictResolution] = postgres_common. ConflictResolution(on_conflict_fields=[], action='IGNORE')): """Configuration for writing vectors to ClouSQL Postgres. Supports flexible schema configuration through column specifications and conflict resolution strategies. Args: connection_config: :class:`LanguageConnectorConfig`. table_name: Target table name. write_config: JdbcIO :class:`~.jdbc_common.WriteConfig` to control batch sizes, authosharding, etc. column_specs: Use :class:`~.postgres_common.ColumnSpecsBuilder` to configure how embeddings and metadata are written a database schema. If None, uses default EmbeddableItem schema. conflict_resolution: Optional :class:`~.postgres_common.ConflictResolution` strategy for handling insert conflicts. ON CONFLICT DO NOTHING by default. Examples: Basic usage with default schema: >>> config = PostgresVectorWriterConfig( ... connection_config=PostgresConnectionConfig(...), ... table_name='embeddings' ... ) Simple case with default schema: >>> config = PostgresVectorWriterConfig( ... connection_config=ConnectionConfig(...), ... table_name='embeddings' ... ) Custom schema with metadata fields: >>> specs = (ColumnSpecsBuilder() ... .with_id_spec(column_name="my_id_column") ... .with_embedding_spec(column_name="embedding_vec") ... .add_metadata_field(field="source", column_name="src") ... .add_metadata_field( ... "timestamp", ... column_name="created_at", ... sql_typecast="::timestamp" ... ) ... .build()) Minimal schema (only ID + embedding written) >>> column_specs = (ColumnSpecsBuilder() ... .with_id_spec() ... .with_embedding_spec() ... .build()) >>> config = CloudSQLPostgresVectorWriterConfig( ... connection_config=PostgresConnectionConfig(...), ... table_name='embeddings', ... column_specs=specs ... ) """ self.connector_config = _PostgresConnectorConfig.from_base_config( connection_config) super().__init__( connection_config=self.connector_config.to_connection_config(), write_config=write_config, table_name=table_name, column_specs=column_specs, conflict_resolution=conflict_resolution) @dataclass class _MySQLConnectorConfig(LanguageConnectorConfig): def to_jdbc_url(self) -> str: """Convert options to a properly formatted MySQL JDBC URL.""" return self._build_jdbc_url( socketFactory="com.google.cloud.sql.mysql.SocketFactory", database_type="mysql") def additional_jdbc_args(self) -> dict[str, list[Any]]: return { 'classpath': [ "mysql:mysql-connector-java:8.0.22", "com.google.cloud.sql:mysql-socket-factory-connector-j-8:1.25.0" ] } @classmethod def from_base_config(cls, config: LanguageConnectorConfig): return cls(**asdict(config)) class CloudSQLMySQLVectorWriterConfig(mysql.MySQLVectorWriterConfig): def __init__( self, connection_config: LanguageConnectorConfig, table_name: str, *, write_config: WriteConfig = WriteConfig(), # pylint: disable=dangerous-default-value column_specs: list[mysql_common.ColumnSpec] = mysql_common. ColumnSpecsBuilder.with_defaults().build(), conflict_resolution: Optional[mysql_common.ConflictResolution] = None): self.connector_config = _MySQLConnectorConfig.from_base_config( connection_config) super().__init__( connection_config=self.connector_config.to_connection_config(), write_config=write_config, table_name=table_name, column_specs=column_specs, conflict_resolution=conflict_resolution) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/cloudsql_it_test.py0000644000175100017510000011442315206113032024775 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import json import logging import os import secrets import time import unittest from dataclasses import dataclass from typing import Any from typing import Literal from typing import Optional import pytest import sqlalchemy from google.cloud.sql.connector import Connector from parameterized import parameterized from sqlalchemy import text import apache_beam as beam from apache_beam.io.jdbc import ReadFromJdbc from apache_beam.ml.rag.ingestion import mysql_common from apache_beam.ml.rag.ingestion import postgres_common from apache_beam.ml.rag.ingestion import test_utils from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteTransform from apache_beam.ml.rag.ingestion.cloudsql import CloudSQLMySQLVectorWriterConfig from apache_beam.ml.rag.ingestion.cloudsql import CloudSQLPostgresVectorWriterConfig from apache_beam.ml.rag.ingestion.cloudsql import LanguageConnectorConfig from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import Embedding from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to _LOGGER = logging.getLogger(__name__) @dataclass class DatabaseTestConfig: """Database-specific test configuration.""" database_type: Literal["postgresql", "mysql"] writer_config_class: type jdbc_driver: str connector_module: Literal["pg8000", "pymysql"] table_prefix: str password_env_var: str username: str database: str instance_uri: str vector_column_type: str metadata_column_type: str common_module: Any id_column_type: str = "VARCHAR(255)" class DatabaseTestHelper: """Helper class to manage database setup, connections, and operations.""" def __init__(self, db_config: DatabaseTestConfig, table_suffix: str): self.db_config = db_config self.table_suffix = table_suffix self.connector = None self.engine = None self.connection_config = None self.default_table_name = f"{db_config.table_prefix}{table_suffix}" self.custom_table_name = f"{db_config.table_prefix}_custom_{table_suffix}" self.metadata_conflicts_table = f"{db_config.table_prefix}_meta_conf_" \ f"{table_suffix}" self._setup_read_queries() def _setup_read_queries(self): if self.db_config.database_type == "postgresql": self.read_queries = { self.default_table_name: f""" SELECT CAST(id AS VARCHAR(255)), CAST(content AS VARCHAR(255)), CAST(embedding AS text), CAST(metadata AS text) FROM {self.default_table_name} """, self.custom_table_name: f""" SELECT CAST(custom_id AS VARCHAR(255)), CAST(embedding_vec AS text), CAST(content_col AS VARCHAR(255)), CAST(metadata AS text) FROM {self.custom_table_name} ORDER BY custom_id """, self.metadata_conflicts_table: f""" SELECT CAST(id AS VARCHAR(255)), CAST(embedding AS text), CAST(content AS VARCHAR(255)), CAST(source AS VARCHAR(255)), CAST(timestamp AS VARCHAR(255)) FROM {self.metadata_conflicts_table} ORDER BY timestamp, id """ } elif self.db_config.database_type == "mysql": self.read_queries = { self.default_table_name: f""" SELECT CAST(id AS CHAR(255)) as id, CAST(content AS CHAR(255)) as content, vector_to_string(embedding) as embedding, CAST(metadata AS CHAR(10000)) as metadata FROM {self.default_table_name} """, self.custom_table_name: f""" SELECT CAST(custom_id AS CHAR(255)) as custom_id, vector_to_string(embedding_vec) as embedding_vec, CAST(content_col AS CHAR(255)) as content_col, CAST(metadata AS CHAR(10000)) as metadata FROM {self.custom_table_name} ORDER BY custom_id """, self.metadata_conflicts_table: f""" SELECT CAST(id AS CHAR(255)) as id, vector_to_string(embedding) as embedding, CAST(content AS CHAR(255)) as content, CAST(source AS CHAR(255)) as source, CAST(timestamp AS CHAR(255)) as timestamp FROM {self.metadata_conflicts_table} ORDER BY timestamp, id """ } def get_read_query(self, table_name: str) -> str: if table_name not in self.read_queries: raise ValueError(f"No read query defined for table: {table_name}") return self.read_queries[table_name] def setup_connection(self): """Set up database connection and engine.""" if not os.environ.get(self.db_config.password_env_var): raise ValueError("Password environment variable not set.") password = os.environ.get(self.db_config.password_env_var) self.connection_config = LanguageConnectorConfig( username=self.db_config.username, password=password, database_name=self.db_config.database, instance_name=self.db_config.instance_uri) self.connector = Connector(refresh_strategy="LAZY") def getconn(): return self.connector.connect( self.db_config.instance_uri, self.db_config.connector_module, user=self.db_config.username, password=password, db=self.db_config.database, ) dialect = "postgresql+pg8000" \ if self.db_config.database_type == "postgresql" else "mysql+pymysql" self.engine = sqlalchemy.create_engine(f"{dialect}://", creator=getconn) def create_all_tables(self): if not self.engine: raise ValueError("Engine not initialized. Call setup_connection() first.") vector_type_large = self.db_config.vector_column_type.format( size=test_utils.VECTOR_SIZE) vector_type_small = self.db_config.vector_column_type.format(size=2) metadata_type = self.db_config.metadata_column_type id_type = self.db_config.id_column_type with self.engine.connect() as connection: default_table_sql = f""" CREATE TABLE {self.default_table_name} ( id {id_type} PRIMARY KEY, embedding {vector_type_large}, content TEXT, metadata {metadata_type} ) """ connection.execute(text(default_table_sql)) custom_table_sql = f""" CREATE TABLE {self.custom_table_name} ( custom_id {id_type} PRIMARY KEY, embedding_vec {vector_type_small}, content_col TEXT, metadata {metadata_type} ) """ connection.execute(text(custom_table_sql)) if self.db_config.database_type == "postgresql": metadata_conflicts_sql = f""" CREATE TABLE {self.metadata_conflicts_table} ( id {id_type}, source TEXT, timestamp TIMESTAMP, content TEXT, embedding {vector_type_small}, PRIMARY KEY (id), UNIQUE (source, timestamp) ) """ elif self.db_config.database_type == "mysql": metadata_conflicts_sql = f""" CREATE TABLE {self.metadata_conflicts_table} ( id {id_type}, source TEXT, timestamp TIMESTAMP, content TEXT, embedding {vector_type_small}, PRIMARY KEY (id), UNIQUE KEY unique_source_timestamp (source(255), timestamp) ) """ connection.execute(text(metadata_conflicts_sql)) connection.commit() def create_writer_config( self, table_name: Optional[str] = None, column_specs=None, conflict_resolution=None): if not self.connection_config: raise ValueError( "Connection not initialized. Call setup_connection() first.") table_name = table_name or self.default_table_name kwargs = { 'connection_config': self.connection_config, 'table_name': table_name, } if column_specs is not None: kwargs['column_specs'] = column_specs if conflict_resolution is not None: kwargs['conflict_resolution'] = conflict_resolution return self.db_config.writer_config_class(**kwargs) def cleanup(self): if self.engine: table_names = [ self.default_table_name, self.custom_table_name, self.metadata_conflicts_table ] try: with self.engine.connect() as connection: for table_name in table_names: connection.execute(text(f"DROP TABLE IF EXISTS {table_name}")) connection.commit() _LOGGER.info( "Dropped %s tables: %s", self.db_config.database_type, ', '.join(table_names)) except Exception as e: _LOGGER.warning( "Error dropping %s tables: %s", self.db_config.database_type, e) if self.connector: try: self.connector.close() except Exception as e: _LOGGER.warning("Error closing connector: %s", e) if self.engine: try: self.engine.dispose() except Exception as e: _LOGGER.warning("Error disposing engine: %s", e) class PipelineVerificationHelper: """Helper class for common pipeline verification patterns.""" @staticmethod def build_jdbc_params(helper: DatabaseTestHelper, table_name: str) -> dict: """Build JDBC parameters dictionary for ReadFromJdbc.""" writer_config = helper.create_writer_config(table_name) return { 'table_name': table_name, 'driver_class_name': helper.db_config.jdbc_driver, 'jdbc_url': writer_config.connector_config.to_connection_config(). jdbc_url, 'username': helper.db_config.username, 'password': helper.connection_config.password, 'query': helper.get_read_query(table_name), 'classpath': writer_config.connector_config.additional_jdbc_args() ['classpath'] } @staticmethod def verify_standard_operations( pipeline, jdbc_params: dict, expected_chunks: list[Chunk]): num_records = len(expected_chunks) sample_size = min(500, num_records // 2) with pipeline as p: rows = (p | ReadFromJdbc(**jdbc_params)) # Count verification count_result = rows | "Count All" >> beam.combiners.Count.Globally() assert_that(count_result, equal_to([num_records]), label='count_check') # Hash verification chunks = (rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk)) chunk_hashes = chunks | "Hash Chunks" >> beam.CombineGlobally( test_utils.HashingFn()) expected_hash = test_utils.generate_expected_hash(num_records) assert_that(chunk_hashes, equal_to([expected_hash]), label='hash_check') # Sample validation - first N first_n = ( chunks | "Key on Index" >> beam.Map(test_utils.key_on_id) | f"Get First {sample_size}" >> beam.transforms.combiners.Top.Of( sample_size, key=lambda x: x[0], reverse=True) | "Remove Keys 1" >> beam.Map(lambda xs: [x[1] for x in xs])) expected_first_n = expected_chunks[:sample_size] assert_that( first_n, equal_to([expected_first_n]), label=f"first_{sample_size}_check") # Sample validation - last N last_n = ( chunks | "Key on Index 2" >> beam.Map(test_utils.key_on_id) | f"Get Last {sample_size}" >> beam.transforms.combiners.Top.Of( sample_size, key=lambda x: x[0]) | "Remove Keys 2" >> beam.Map(lambda xs: [x[1] for x in xs])) expected_last_n = expected_chunks[-sample_size:][::-1] assert_that( last_n, equal_to([expected_last_n]), label=f"last_{sample_size}_check") # Database configurations POSTGRES_CONFIG = DatabaseTestConfig( database_type="postgresql", writer_config_class=CloudSQLPostgresVectorWriterConfig, jdbc_driver="org.postgresql.Driver", connector_module="pg8000", table_prefix="python_rag_postgres_", password_env_var="ALLOYDB_PASSWORD", username="postgres", database="postgres", instance_uri="apache-beam-testing:us-central1:beam-integration-tests", vector_column_type="VECTOR({size})", metadata_column_type="JSONB", common_module=postgres_common) MYSQL_CONFIG = DatabaseTestConfig( database_type="mysql", writer_config_class=CloudSQLMySQLVectorWriterConfig, jdbc_driver="com.mysql.cj.jdbc.Driver", connector_module="pymysql", table_prefix="python_rag_mysql_", password_env_var="ALLOYDB_PASSWORD", username="mysql", database="embeddings", instance_uri="apache-beam-testing:us-central1:beam-integration-tests-mysql", vector_column_type="VECTOR({size}) USING VARBINARY", metadata_column_type="JSON", common_module=mysql_common) @pytest.mark.uses_gcp_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_JARS'), "EXPANSION_JARS environment var is not provided, " "indicating that jars have not been built") class CloudSQLVectorWriterConfigTest(unittest.TestCase): def setUp(self): self.write_test_pipeline = TestPipeline(is_integration_test=True) self.read_test_pipeline = TestPipeline(is_integration_test=True) self.write_test_pipeline2 = TestPipeline(is_integration_test=True) self.read_test_pipeline2 = TestPipeline(is_integration_test=True) self.write_test_pipeline.not_use_test_runner_api = True self.read_test_pipeline.not_use_test_runner_api = True self.write_test_pipeline2.not_use_test_runner_api = True self.read_test_pipeline2.not_use_test_runner_api = True self._runner = type(self.read_test_pipeline.runner).__name__ self.db_helpers = {} self.table_suffix = '%d%s' % (int(time.time()), secrets.token_hex(3)) # Set up database helpers for config in [POSTGRES_CONFIG, MYSQL_CONFIG]: helper = DatabaseTestHelper(config, self.table_suffix) helper.setup_connection() helper.create_all_tables() self.db_helpers[config.database_type] = helper _LOGGER.info("Successfully set up %s database", config.database_type) def tearDown(self): for helper in self.db_helpers.values(): helper.cleanup() def skip_if_dataflow_runner(self): if self._runner and "dataflowrunner" in self._runner.lower(): self.skipTest( "Skipping some tests on Dataflow Runner to avoid bloat and timeouts") @parameterized.expand([(POSTGRES_CONFIG), (MYSQL_CONFIG)]) def test_default_config(self, db_config): """Test basic write and read operations with default configuration. This test validates the most basic CloudSQL vector database functionality: - Default table schema: id (VARCHAR), content (TEXT), embedding (VECTOR), metadata (JSON/JSONB) - Default column specifications (no customization) - Default conflict resolution (IGNORE on primary key conflicts) - Write chunks to database and read them back - Verify data integrity through count, hash, and sample validation """ self.skip_if_dataflow_runner() helper = self.db_helpers[db_config.database_type] num_records = 150 # Create test data test_chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records) # Write test writer_config = helper.create_writer_config() self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = ( p | beam.Create(test_chunks) | VectorDatabaseWriteTransform(writer_config)) # Read and verify self.read_test_pipeline.not_use_test_runner_api = True jdbc_params = PipelineVerificationHelper.build_jdbc_params( helper, helper.default_table_name) PipelineVerificationHelper.verify_standard_operations( self.read_test_pipeline, jdbc_params, test_chunks) @parameterized.expand([ (POSTGRES_CONFIG, "UPDATE", ["embedding", "content"]), (MYSQL_CONFIG, "UPDATE", ["embedding", "content"]), (POSTGRES_CONFIG, "IGNORE", None), (MYSQL_CONFIG, "IGNORE", None), (POSTGRES_CONFIG, "UPDATE_ALL", None), # Default update fields (MYSQL_CONFIG, "UPDATE_ALL", None), ]) def test_conflict_resolution(self, db_config, action, update_fields): """Test conflict resolution strategies when primary key conflicts occur. This test validates different approaches to handling duplicate primary keys: UPDATE with specific fields: - When duplicate ID encountered, update only specified fields (embedding, content) - Other fields (metadata) remain unchanged from original record IGNORE: - When duplicate ID encountered, keep original record unchanged UPDATE_ALL (default update fields): - When duplicate ID encountered, update ALL non-key fields - This includes content, embedding, AND metadata Scenario for all strategies: 1. Insert initial records 2. Insert records with same IDs but different content/embeddings 3. Verify final state matches expected conflict resolution behavior """ self.skip_if_dataflow_runner() helper = self.db_helpers[db_config.database_type] num_records = 20 common_module = db_config.common_module if action == "IGNORE": if db_config.database_type == "mysql": conflict_resolution = common_module.ConflictResolution( action="IGNORE", primary_key_field="id") else: conflict_resolution = None # Default behavior for PostgreSQL elif action == "UPDATE": if db_config.database_type == "postgresql": conflict_resolution = common_module.ConflictResolution( on_conflict_fields="id", action="UPDATE", update_fields=update_fields) else: conflict_resolution = common_module.ConflictResolution( action="UPDATE", update_fields=update_fields) else: # UPDATE_ALL if db_config.database_type == "postgresql": conflict_resolution = common_module.ConflictResolution( on_conflict_fields="id", action="UPDATE") else: conflict_resolution = common_module.ConflictResolution(action="UPDATE") initial_chunks = test_utils.ChunkTestUtils.get_expected_values( 0, num_records) writer_config = helper.create_writer_config( conflict_resolution=conflict_resolution) self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = ( p | "Write Initial" >> beam.Create(initial_chunks) | VectorDatabaseWriteTransform(writer_config)) # Write conflicting data updated_chunks = test_utils.ChunkTestUtils.get_expected_values( 0, num_records, content_prefix="Updated", seed_multiplier=2) self.write_test_pipeline2.not_use_test_runner_api = True with self.write_test_pipeline2 as p: _ = ( p | "Write Conflicts" >> beam.Create(updated_chunks) | VectorDatabaseWriteTransform(writer_config)) jdbc_params = PipelineVerificationHelper.build_jdbc_params( helper, helper.default_table_name) expected_chunks = updated_chunks if action != "IGNORE" else initial_chunks self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = (p | ReadFromJdbc(**jdbc_params)) count_result = rows | "Count All" >> beam.combiners.Count.Globally() assert_that(count_result, equal_to([num_records]), label='count_check') chunks = rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk) assert_that(chunks, equal_to(expected_chunks), label='chunks_check') @parameterized.expand([(POSTGRES_CONFIG), (MYSQL_CONFIG)]) def test_custom_column_names_and_value_functions(self, db_config): """Test completely custom column specifications with custom value extraction. This test validates advanced customization of how chunk data is stored: Custom column names: - custom_id (instead of 'id') - embedding_vec (instead of 'embedding') - content_col (instead of 'content') Custom value extraction functions: - ID: Extract timestamp from metadata and prefix with "timestamp_" - Content: Prefix content with its character length "10:actual_content" - Embedding: Use custom embedding extraction function This tests the flexibility to completely reshape how chunk data maps to database columns, useful for integrating with existing database schemas or applying business-specific transformations. """ self.skip_if_dataflow_runner() helper = self.db_helpers[db_config.database_type] num_records = 20 common_module = db_config.common_module test_chunks = [ Chunk( id=str(i), content=Content(text=f"content_{i}"), embedding=Embedding(dense_embedding=[float(i), float(i + 1)]), metadata={"timestamp": f"2024-02-02T{i:02d}:00:00"}) for i in range(num_records) ] chunk_embedding_fn = common_module.chunk_embedding_fn specs = ( common_module.ColumnSpecsBuilder().add_custom_column_spec( common_module.ColumnSpec.text( column_name="custom_id", value_fn=lambda chunk: f"timestamp_{chunk.metadata.get('timestamp', '')}") ).add_custom_column_spec( common_module.ColumnSpec.vector( column_name="embedding_vec", value_fn=chunk_embedding_fn)).add_custom_column_spec( common_module.ColumnSpec.text( column_name="content_col", value_fn=lambda chunk: f"{len(chunk.content.text)}:{chunk.content.text}")). with_metadata_spec().build()) def custom_row_to_chunk(row): timestamp = row.custom_id.split('timestamp_')[1] i = int(timestamp.split('T')[1][:2]) embedding_list = [ float(x) for x in row.embedding_vec.strip('[]').split(',') ] content = row.content_col.split(':', 1)[1] return Chunk( id=str(i), content=Content(text=content), embedding=Embedding(dense_embedding=embedding_list), metadata=json.loads(row.metadata)) writer_config = helper.create_writer_config(helper.custom_table_name, specs) self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = ( p | beam.Create(test_chunks) | VectorDatabaseWriteTransform(writer_config)) jdbc_params = PipelineVerificationHelper.build_jdbc_params( helper, helper.custom_table_name) self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = (p | ReadFromJdbc(**jdbc_params)) count_result = rows | "Count All" >> beam.combiners.Count.Globally() assert_that(count_result, equal_to([num_records]), label='count_check') chunks = rows | "To Chunks" >> beam.Map(custom_row_to_chunk) assert_that(chunks, equal_to(test_chunks), label='chunks_check') @parameterized.expand([(POSTGRES_CONFIG), (MYSQL_CONFIG)]) def test_custom_type_conversion_with_default_columns(self, db_config): """Test custom type conversion and SQL typecasting with modified column names. This test validates data type handling and database-specific SQL features: Type conversion: - Convert string IDs to integers before storage - Apply length-prefix transformation to content SQL typecasting (database-specific): - PostgreSQL: Use ::text typecast for converted integers - MySQL: Rely on automatic type conversion (no explicit typecast) Column name customization: - Use custom names but with standard spec builders (not completely custom functions) This tests the ability to adapt data types for database constraints while maintaining the standard chunk-to-database mapping logic. """ self.skip_if_dataflow_runner() helper = self.db_helpers[db_config.database_type] num_records = 20 common_module = db_config.common_module test_chunks = [ Chunk( id=str(i), content=Content(text=f"content_{i}"), embedding=Embedding(dense_embedding=[float(i), float(i + 1)]), metadata={"timestamp": f"2024-02-02T{i:02d}:00:00"}) for i in range(num_records) ] if db_config.database_type == "postgresql": specs = ( common_module.ColumnSpecsBuilder().with_id_spec( column_name="custom_id", python_type=int, convert_fn=lambda x: int(x), sql_typecast="::text").with_content_spec( column_name="content_col", convert_fn=lambda x: f"{len(x)}:{x}" # Add length prefix ).with_embedding_spec( column_name="embedding_vec").with_metadata_spec().build()) else: # MySQL specs = ( common_module.ColumnSpecsBuilder().with_id_spec( column_name="custom_id", python_type=int, convert_fn=lambda x: int(x)).with_content_spec( column_name="content_col", convert_fn=lambda x: f"{len(x)}:{x}").with_embedding_spec( column_name="embedding_vec").with_metadata_spec().build()) def type_conversion_row_to_chunk(row): embedding_list = [ float(x) for x in row.embedding_vec.strip('[]').split(',') ] content = row.content_col.split(':', 1)[1] return Chunk( id=row.custom_id, # custom_id is the converted ID field content=Content(text=content), embedding=Embedding(dense_embedding=embedding_list), metadata=json.loads(row.metadata)) writer_config = helper.create_writer_config(helper.custom_table_name, specs) self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = ( p | beam.Create(test_chunks) | VectorDatabaseWriteTransform(writer_config)) jdbc_params = PipelineVerificationHelper.build_jdbc_params( helper, helper.custom_table_name) self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = (p | ReadFromJdbc(**jdbc_params)) count_result = rows | "Count All" >> beam.combiners.Count.Globally() assert_that(count_result, equal_to([num_records]), label='count_check') chunks = rows | "To Chunks" >> beam.Map(type_conversion_row_to_chunk) assert_that(chunks, equal_to(test_chunks), label='chunks_check') @parameterized.expand([(POSTGRES_CONFIG), (MYSQL_CONFIG)]) def test_default_id_embedding_specs(self, db_config): """Test minimal schema with only ID and embedding columns. This test validates the ability to create a minimal vector database schema: - Only stores id and embedding fields - content and metadata columns are excluded from the table - Tests that the system correctly handles missing/null fields Use case: When you only need vector similarity search without storing the original content or metadata (perhaps stored elsewhere). Validation: - Chunks written with content/metadata are stored with those fields as null - Reading back produces chunks with null content and empty metadata - Vector similarity operations still work normally """ self.skip_if_dataflow_runner() helper = self.db_helpers[db_config.database_type] num_records = 20 common_module = db_config.common_module specs = ( common_module.ColumnSpecsBuilder().with_id_spec().with_embedding_spec(). build()) writer_config = helper.create_writer_config(column_specs=specs) test_chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records) with self.write_test_pipeline as p: _ = ( p | beam.Create(test_chunks) | VectorDatabaseWriteTransform(writer_config)) expected_chunks = test_utils.ChunkTestUtils.get_expected_values( 0, num_records) for chunk in expected_chunks: chunk.content.text = None # Content column not included in schema chunk.metadata = {} # Metadata column not included in schema jdbc_params = PipelineVerificationHelper.build_jdbc_params( helper, helper.default_table_name) if db_config.database_type == "postgresql": jdbc_params['query'] = f""" SELECT CAST(id AS VARCHAR(255)), CAST(embedding AS text) FROM {helper.default_table_name} ORDER BY id """ elif db_config.database_type == "mysql": jdbc_params['query'] = f""" SELECT CAST(id AS CHAR(255)) as id, vector_to_string(embedding) as embedding FROM {helper.default_table_name} """ with self.read_test_pipeline as p: rows = (p | ReadFromJdbc(**jdbc_params)) chunks = rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk) assert_that(chunks, equal_to(expected_chunks), label='chunks_check') @parameterized.expand([(POSTGRES_CONFIG), (MYSQL_CONFIG)]) def test_metadata_field_extraction(self, db_config): """Test extracting specific metadata fields into separate database columns. This test validates the ability to: - Extract specific fields from the JSON metadata object - Map them to dedicated database columns (e.g., metadata.source -> source column) - Apply database-specific SQL typecasts (PostgreSQL ::timestamp vs MySQL default) - Store and retrieve the extracted fields correctly This is different from default metadata handling which stores the entire metadata object as JSON in a single column. """ self.skip_if_dataflow_runner() helper = self.db_helpers[db_config.database_type] num_records = 20 common_module = db_config.common_module if db_config.database_type == "postgresql": specs = ( common_module.ColumnSpecsBuilder().with_id_spec().with_embedding_spec( ).with_content_spec().add_metadata_field( field="source", column_name="source", python_type=str, sql_typecast=None).add_metadata_field( field="timestamp", python_type=str, sql_typecast="::timestamp").build()) else: specs = ( common_module.ColumnSpecsBuilder().with_id_spec().with_embedding_spec( ).with_content_spec().add_metadata_field( field="source", column_name="source", python_type=str).add_metadata_field( field="timestamp", python_type=str).build()) writer_config = helper.create_writer_config( helper.metadata_conflicts_table, specs, conflict_resolution=None) test_chunks = [ Chunk( id=str(i), content=Content(text=f"content_{i}"), embedding=Embedding(dense_embedding=[float(i), float(i + 1)]), metadata={ "source": f"source_{i % 3}", "timestamp": f"2024-02-02T{i:02d}:00:00" }) for i in range(num_records) ] self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = ( p | beam.Create(test_chunks) | VectorDatabaseWriteTransform(writer_config)) def metadata_row_to_chunk(row): embedding_list = [float(x) for x in row.embedding.strip('[]').split(',')] timestamp = row.timestamp.replace( ' ', 'T') if ' ' in row.timestamp else row.timestamp return Chunk( id=row.id, content=Content(text=row.content), embedding=Embedding(dense_embedding=embedding_list), metadata={ "source": row.source, "timestamp": timestamp }) jdbc_params = PipelineVerificationHelper.build_jdbc_params( helper, helper.metadata_conflicts_table) self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = (p | ReadFromJdbc(**jdbc_params)) chunks = rows | "To Chunks" >> beam.Map(metadata_row_to_chunk) assert_that(chunks, equal_to(test_chunks), label='chunks_check') @parameterized.expand([(POSTGRES_CONFIG), (MYSQL_CONFIG)]) def test_composite_unique_constraint_conflicts(self, db_config): """Test conflict resolution when unique constraints span multiple columns. This test validates conflict resolution when the unique constraint is NOT on the primary key, but on a combination of other columns (source + timestamp). Scenario: 1. Insert records with unique (source, timestamp) combinations 2. Attempt to insert records with same (source, timestamp) but different IDs and content 3. Verify that conflict resolution (UPDATE) works correctly based on composite key This is different from test_conflict_resolution which tests conflicts on the primary key field only. """ self.skip_if_dataflow_runner() helper = self.db_helpers[db_config.database_type] num_records = 5 common_module = db_config.common_module if db_config.database_type == "postgresql": specs = ( common_module.ColumnSpecsBuilder().with_id_spec().with_embedding_spec( ).with_content_spec().add_metadata_field( field="source", column_name="source", python_type=str, sql_typecast=None).add_metadata_field( field="timestamp", python_type=str, sql_typecast="::timestamp").build()) conflict_resolution = common_module.ConflictResolution( on_conflict_fields=["source", "timestamp"], action="UPDATE", update_fields=["embedding", "content"]) elif db_config.database_type == "mysql": specs = ( common_module.ColumnSpecsBuilder().with_id_spec().with_embedding_spec( ).with_content_spec().add_metadata_field( field="source", column_name="source", python_type=str).add_metadata_field( field="timestamp", python_type=str).build()) # MySQL conflict resolution - detects unique constraint automatically conflict_resolution = common_module.ConflictResolution( action="UPDATE", update_fields=["embedding", "content"]) writer_config = helper.create_writer_config( helper.metadata_conflicts_table, specs, conflict_resolution) initial_chunks = [ Chunk( id=str(i), content=Content(text=f"content_{i}"), embedding=Embedding(dense_embedding=[float(i), float(i + 1)]), metadata={ "source": "source_A", "timestamp": f"2024-02-02T{i:02d}:00:00" }) for i in range(num_records) ] with self.write_test_pipeline as p: _ = ( p | "Write Initial" >> beam.Create(initial_chunks) | VectorDatabaseWriteTransform(writer_config)) conflicting_chunks = [ Chunk( id=f"new_{i}", content=Content(text=f"updated_content_{i}"), embedding=Embedding( dense_embedding=[float(i) * 2, float(i + 1) * 2]), metadata={ "source": "source_A", "timestamp": f"2024-02-02T{i:02d}:00:00" }) for i in range(num_records) ] with self.write_test_pipeline2 as p: _ = ( p | "Write Conflicts" >> beam.Create(conflicting_chunks) | VectorDatabaseWriteTransform(writer_config)) expected_chunks = [ Chunk( id=str(i), content=Content(text=f"updated_content_{i}"), embedding=Embedding( dense_embedding=[float(i) * 2, float(i + 1) * 2]), metadata={ "source": "source_A", "timestamp": f"2024-02-02T{i:02d}:00:00" }) for i in range(num_records) ] def metadata_row_to_chunk(row): embedding_list = [float(x) for x in row.embedding.strip('[]').split(',')] timestamp = row.timestamp.replace( ' ', 'T') if ' ' in row.timestamp else row.timestamp return Chunk( id=row.id, content=Content(text=row.content), embedding=Embedding(dense_embedding=embedding_list), metadata={ "source": row.source, "timestamp": timestamp }) jdbc_params = PipelineVerificationHelper.build_jdbc_params( helper, helper.metadata_conflicts_table) with self.read_test_pipeline as p: rows = (p | ReadFromJdbc(**jdbc_params)) count_result = rows | "Count All" >> beam.combiners.Count.Globally() assert_that(count_result, equal_to([num_records]), label='count_check') chunks = rows | "To Chunks" >> beam.Map(metadata_row_to_chunk) assert_that(chunks, equal_to(expected_chunks), label='chunks_check') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/jdbc_common.py0000644000175100017510000000551715206113032023671 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from dataclasses import dataclass from dataclasses import field from typing import Any from typing import Optional @dataclass class ConnectionConfig: """Configuration for connecting to a JDBC database. Provides connection details and options for connecting to a database instance. Attributes: jdbc_url: JDBC URL for the database instance. Example: 'jdbc:postgresql://host:port/database' username: Database username. password: Database password. connection_properties: Optional JDBC connection properties dict. Example: {'ssl': 'true'} connection_init_sqls: Optional list of SQL statements to execute when connection is established. additional_jdbc_args: Additional arguments that will be passed to WriteToJdbc. These may include 'driver_jars', 'expansion_service', 'classpath', etc. See full set of args at :class:`~apache_beam.io.jdbc.WriteToJdbc` Example: >>> config = AlloyDBConnectionConfig( ... jdbc_url='jdbc:postgresql://localhost:5432/mydb', ... username='user', ... password='pass', ... connection_properties={'ssl': 'true'}, ... max_connections=10 ... ) """ jdbc_url: str username: str password: str connection_properties: Optional[dict[str, str]] = None connection_init_sqls: Optional[list[str]] = None additional_jdbc_args: dict[str, Any] = field(default_factory=dict) @dataclass class WriteConfig: """Configuration writing to JDBC database. Modifies the write behavior when writing via JdbcIO. Attributes: autosharding: Enable automatic re-sharding of bundles to scale the number of shards with workers. max_connections: Optional number of connections in the pool. Use negative for no limit. write_batch_size: Optional write batch size for bulk operations. """ autosharding: Optional[bool] = None max_connections: Optional[int] = None write_batch_size: Optional[int] = None ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/milvus_search.py0000644000175100017510000003100115206113032024246 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import logging from dataclasses import dataclass from dataclasses import field from typing import Any from typing import Callable from typing import Optional from pymilvus import MilvusClient from pymilvus.exceptions import MilvusException import apache_beam as beam from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteConfig from apache_beam.ml.rag.ingestion.jdbc_common import WriteConfig from apache_beam.ml.rag.ingestion.postgres_common import ColumnSpec from apache_beam.ml.rag.ingestion.postgres_common import ColumnSpecsBuilder from apache_beam.ml.rag.types import EmbeddableItem from apache_beam.ml.rag.utils import DEFAULT_WRITE_BATCH_SIZE from apache_beam.ml.rag.utils import MilvusConnectionParameters from apache_beam.ml.rag.utils import MilvusHelpers from apache_beam.ml.rag.utils import retry_with_backoff from apache_beam.ml.rag.utils import unpack_dataclass_with_kwargs from apache_beam.transforms import DoFn _LOGGER = logging.getLogger(__name__) @dataclass class MilvusWriteConfig: """Configuration parameters for writing data to Milvus collections. This class defines the parameters needed to write data to a Milvus collection, including collection targeting, batching behavior, and operation timeouts. Args: collection_name: Name of the target Milvus collection to write data to. Must be a non-empty string. partition_name: Name of the specific partition within the collection to write to. If empty, writes to the default partition. timeout: Maximum time in seconds to wait for write operations to complete. If None, uses the client's default timeout. write_config: Configuration for write operations including batch size and other write-specific settings. kwargs: Additional keyword arguments for write operations. Enables forward compatibility with future Milvus client parameters. """ collection_name: str partition_name: str = "" timeout: Optional[float] = None write_config: WriteConfig = field(default_factory=WriteConfig) kwargs: dict[str, Any] = field(default_factory=dict) def __post_init__(self): if not self.collection_name: raise ValueError("Collection name must be provided") @property def write_batch_size(self): """Returns the batch size for write operations. Returns: The configured batch size, or DEFAULT_WRITE_BATCH_SIZE if not specified. """ return self.write_config.write_batch_size or DEFAULT_WRITE_BATCH_SIZE @dataclass class MilvusVectorWriterConfig(VectorDatabaseWriteConfig): """Configuration for writing vector data to Milvus collections. This class extends VectorDatabaseWriteConfig to provide Milvus-specific configuration for ingesting vector embeddings and associated metadata. It defines how EmbeddableItem objects are converted to Milvus records and handles the write operation parameters. The configuration includes connection parameters, write settings, and column specifications that determine how data is mapped to Milvus fields. Args: connection_params: Configuration for connecting to the Milvus server, including URI, credentials, and connection options. write_config: Configuration for write operations including collection name, partition, batch size, and timeouts. column_specs: List of column specifications defining how EmbeddableItem fields are mapped to Milvus collection fields. Defaults to standard RAG fields (id, embedding, sparse_embedding, content, metadata). Example: config = MilvusVectorWriterConfig( connection_params=MilvusConnectionParameters( uri="http://localhost:19530"), write_config=MilvusWriteConfig(collection_name="my_collection"), column_specs=MilvusVectorWriterConfig.default_column_specs()) """ connection_params: MilvusConnectionParameters write_config: MilvusWriteConfig column_specs: list[ColumnSpec] = field( default_factory=lambda: MilvusVectorWriterConfig.default_column_specs()) def create_converter(self) -> Callable[[EmbeddableItem], dict[str, Any]]: """Creates a function to convert EmbeddableItem objects to Milvus records. Returns: A function that takes an EmbeddableItem and returns a dictionary representing a Milvus record with fields mapped according to column_specs. """ def convert(chunk: EmbeddableItem) -> dict[str, Any]: result = {} for col in self.column_specs: result[col.column_name] = col.value_fn(chunk) return result return convert def create_write_transform(self) -> beam.PTransform: """Creates the Apache Beam transform for writing to Milvus. Returns: A PTransform that can be applied to a PCollection of EmbeddableItem objects to write them to the configured Milvus collection. """ return _WriteToMilvusVectorDatabase(self) @staticmethod def default_column_specs() -> list[ColumnSpec]: """Returns default column specifications for RAG use cases. Creates column mappings for standard RAG fields: id, dense embedding, sparse embedding, content text, and metadata. These specifications define how EmbeddableItem fields are converted to Milvus-compatible formats. Returns: List of ColumnSpec objects defining the default field mappings. """ column_specs = ColumnSpecsBuilder() return column_specs\ .with_id_spec()\ .with_embedding_spec(convert_fn=lambda values: list(values))\ .with_sparse_embedding_spec(conv_fn=MilvusHelpers.sparse_embedding)\ .with_content_spec()\ .with_metadata_spec(convert_fn=lambda values: dict(values))\ .build() class _WriteToMilvusVectorDatabase(beam.PTransform): """Apache Beam PTransform for writing vector data to Milvus. This transform handles the conversion of EmbeddableItem objects to Milvus records and coordinates the write operations. It applies the configured converter function and uses a DoFn for batched writes to optimize performance. Args: config: MilvusVectorWriterConfig containing all necessary parameters for the write operation. """ def __init__(self, config: MilvusVectorWriterConfig): self.config = config def expand(self, pcoll: beam.PCollection[EmbeddableItem]): """Expands the PTransform to convert and write to Milvus. Args: pcoll: PCollection of EmbeddableItem objects to write to Milvus. Returns: PCollection of dictionaries representing the records written to Milvus. """ return ( pcoll | "Convert to Records" >> beam.Map(self.config.create_converter()) | beam.ParDo( _WriteMilvusFn( self.config.connection_params, self.config.write_config))) class _WriteMilvusFn(DoFn): """DoFn that handles batched writes to Milvus. This DoFn accumulates records in batches and flushes them to Milvus when the batch size is reached or when the bundle finishes. This approach optimizes performance by reducing the number of individual write operations. Args: connection_params: Configuration for connecting to the Milvus server. write_config: Configuration for write operations including batch size and collection details. """ def __init__( self, connection_params: MilvusConnectionParameters, write_config: MilvusWriteConfig): self._connection_params = connection_params self._write_config = write_config self.batch = [] def process(self, element, *args, **kwargs): """Processes individual records, batching them for efficient writes. Args: element: A dictionary representing a Milvus record to write. *args: Additional positional arguments. **kwargs: Additional keyword arguments. Yields: The original element after adding it to the batch. """ _ = args, kwargs # Unused parameters self.batch.append(element) if len(self.batch) >= self._write_config.write_batch_size: self._flush() yield element def finish_bundle(self): """Called when a bundle finishes processing. Flushes any remaining records in the batch to ensure all data is written. """ self._flush() def _flush(self): """Flushes the current batch of records to Milvus. Creates a MilvusSink connection and writes all batched records, then clears the batch for the next set of records. """ if len(self.batch) == 0: return with _MilvusSink(self._connection_params, self._write_config) as sink: sink.write(self.batch) self.batch = [] def display_data(self): """Returns display data for monitoring and debugging. Returns: Dictionary containing database, collection, and batch size information for display in the Apache Beam monitoring UI. """ res = super().display_data() res["database"] = self._connection_params.db_name res["collection"] = self._write_config.collection_name res["batch_size"] = self._write_config.write_batch_size return res class _MilvusSink: """Low-level sink for writing data directly to Milvus. This class handles the direct interaction with the Milvus client for upsert operations. It manages the connection lifecycle and provides context manager support for proper resource cleanup. Args: connection_params: Configuration for connecting to the Milvus server. write_config: Configuration for write operations including collection and partition targeting. """ def __init__( self, connection_params: MilvusConnectionParameters, write_config: MilvusWriteConfig): self._connection_params = connection_params self._write_config = write_config self._client = None def write(self, documents): """Writes a batch of documents to the Milvus collection. Performs an upsert operation to insert new documents or update existing ones based on primary key. After the upsert, flushes the collection to ensure data persistence. Args: documents: List of dictionaries representing Milvus records to write. Each dictionary should contain fields matching the collection schema. """ self._client = MilvusClient( **unpack_dataclass_with_kwargs(self._connection_params)) resp = self._client.upsert( collection_name=self._write_config.collection_name, partition_name=self._write_config.partition_name, data=documents, timeout=self._write_config.timeout, **self._write_config.kwargs) _LOGGER.debug( "Upserted into Milvus: upsert_count=%d, cost=%d", resp.get("upsert_count", 0), resp.get("cost", 0)) def __enter__(self): """Enters the context manager and establishes Milvus connection. Returns: Self, enabling use in 'with' statements. """ if not self._client: connection_params = unpack_dataclass_with_kwargs(self._connection_params) # Extract retry parameters from connection_params. max_retries = connection_params.pop('max_retries', 3) retry_delay = connection_params.pop('retry_delay', 1.0) retry_backoff_factor = connection_params.pop('retry_backoff_factor', 2.0) def create_client(): return MilvusClient(**connection_params) self._client = retry_with_backoff( create_client, max_retries=max_retries, retry_delay=retry_delay, retry_backoff_factor=retry_backoff_factor, operation_name="Milvus connection", exception_types=(MilvusException, )) return self def __exit__(self, exc_type, exc_val, exc_tb): """Exits the context manager and closes the Milvus connection. Args: exc_type: Exception type if an exception was raised. exc_val: Exception value if an exception was raised. exc_tb: Exception traceback if an exception was raised. """ _ = exc_type, exc_val, exc_tb # Unused parameters if self._client: self._client.close() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/milvus_search_it_test.py0000644000175100017510000005737115206113032026023 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import platform import unittest import uuid from typing import Callable from typing import cast import pytest import apache_beam as beam from apache_beam.ml.rag.ingestion.jdbc_common import WriteConfig from apache_beam.ml.rag.test_utils import MilvusTestHelpers from apache_beam.ml.rag.test_utils import VectorDBContainerInfo from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import Embedding from apache_beam.ml.rag.utils import MilvusConnectionParameters from apache_beam.ml.rag.utils import retry_with_backoff from apache_beam.ml.rag.utils import unpack_dataclass_with_kwargs from apache_beam.testing.test_pipeline import TestPipeline # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from pymilvus import CollectionSchema from pymilvus import DataType from pymilvus import FieldSchema from pymilvus import MilvusClient from pymilvus.exceptions import MilvusException from pymilvus.milvus_client import IndexParams from apache_beam.ml.rag.ingestion.milvus_search import MilvusVectorWriterConfig from apache_beam.ml.rag.ingestion.milvus_search import MilvusWriteConfig PYMILVUS_AVAILABLE = True except ImportError: PYMILVUS_AVAILABLE = False # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports def _construct_index_params(): index_params = IndexParams() # Dense vector index for dense embeddings. index_params.add_index( field_name="embedding", index_name="embedding_ivf_flat", index_type="IVF_FLAT", metric_type="COSINE", params={"nlist": 1}) # Sparse vector index for sparse embeddings. index_params.add_index( field_name="sparse_embedding", index_name="sparse_embedding_inverted_index", index_type="SPARSE_INVERTED_INDEX", metric_type="IP", params={"inverted_index_algo": "TAAT_NAIVE"}) return index_params MILVUS_INGESTION_IT_CONFIG = { "fields": [ FieldSchema( name="id", dtype=DataType.INT64, is_primary=True, auto_id=False), FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=1000), FieldSchema(name="metadata", dtype=DataType.JSON), FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=3), FieldSchema( name="sparse_embedding", dtype=DataType.SPARSE_FLOAT_VECTOR) ], "index": _construct_index_params, "corpus": [ Chunk( id=1, # type: ignore[arg-type] content=Content(text="Test document one"), metadata={"source": "test1"}, embedding=Embedding( dense_embedding=[0.1, 0.2, 0.3], sparse_embedding=([1, 2], [0.1, 0.2])), ), Chunk( id=2, # type: ignore[arg-type] content=Content(text="Test document two"), metadata={"source": "test2"}, embedding=Embedding( dense_embedding=[0.2, 0.3, 0.4], sparse_embedding=([2, 3], [0.3, 0.1]), ), ), Chunk( id=3, # type: ignore[arg-type] content=Content(text="Test document three"), metadata={"source": "test3"}, embedding=Embedding( dense_embedding=[0.3, 0.4, 0.5], sparse_embedding=([3, 4], [0.4, 0.2]), ), ) ] } def create_collection_with_partition( client: MilvusClient, collection_name: str, partition_name: str = '', fields=None): if fields is None: fields = MILVUS_INGESTION_IT_CONFIG["fields"] # Configure schema. schema = CollectionSchema(fields=fields) # Configure index. index_function: Callable[[], IndexParams] = cast( Callable[[], IndexParams], MILVUS_INGESTION_IT_CONFIG["index"]) # Create collection with schema. client.create_collection( collection_name=collection_name, schema=schema, index_params=index_function()) # Create partition within the collection. client.create_partition( collection_name=collection_name, partition_name=partition_name) msg = f"Expected collection '{collection_name}' to be created." assert client.has_collection(collection_name), msg msg = f"Expected partition '{partition_name}' to be created." assert client.has_partition(collection_name, partition_name), msg # Release the collection from memory. We don't need that on pure writing. client.release_collection(collection_name) def drop_collection(client: MilvusClient, collection_name: str): try: client.drop_collection(collection_name) assert not client.has_collection(collection_name) except Exception: # Silently ignore connection errors during cleanup. pass @pytest.mark.require_docker_in_docker @unittest.skipIf(not PYMILVUS_AVAILABLE, 'pymilvus is not installed.') @unittest.skipUnless( platform.system() == "Linux", "Test runs only on Linux due to lack of support, as yet, for nested " "virtualization in CI environments on Windows/macOS. Many CI providers run " "tests in virtualized environments, and nested virtualization " "(Docker inside a VM) is either unavailable or has several issues on " "non-Linux platforms.") class TestMilvusVectorWriterConfig(unittest.TestCase): """Integration tests for Milvus vector database ingestion functionality""" _db: VectorDBContainerInfo @classmethod def setUpClass(cls): cls._db = MilvusTestHelpers.start_db_container() cls._connection_config = MilvusConnectionParameters( uri=cls._db.uri, user=cls._db.user, password=cls._db.password, db_name=cls._db.id, token=cls._db.token) @classmethod def tearDownClass(cls): MilvusTestHelpers.stop_db_container(cls._db) cls._db = None def setUp(self): self.write_test_pipeline = TestPipeline() self._collection_name = f"test_collection_{self._testMethodName}" self._partition_name = f"test_partition_{self._testMethodName}" config = unpack_dataclass_with_kwargs(self._connection_config) config["alias"] = f"milvus_conn_{uuid.uuid4().hex[:8]}" # Use retry_with_backoff for test client connection. def create_client(): return MilvusClient(**config) self._test_client = retry_with_backoff( create_client, max_retries=5, retry_delay=2.0, operation_name="Test Milvus client connection", exception_types=(MilvusException, )) create_collection_with_partition( self._test_client, self._collection_name, self._partition_name) def tearDown(self): drop_collection(self._test_client, self._collection_name) self._test_client.close() def test_invalid_write_on_non_existent_collection(self): non_existent_collection = "nonexistent_collection" test_chunks = MILVUS_INGESTION_IT_CONFIG["corpus"] write_config = MilvusWriteConfig( collection_name=non_existent_collection, write_config=WriteConfig(write_batch_size=1)) config = MilvusVectorWriterConfig( connection_params=self._connection_config, write_config=write_config, ) # Write pipeline. with self.assertRaises(Exception) as context: with TestPipeline() as p: p.not_use_test_runner_api = True _ = (p | beam.Create(test_chunks) | config.create_write_transform()) # Assert on what should happen. self.assertIn("can't find collection", str(context.exception).lower()) def test_invalid_write_on_non_existent_partition(self): non_existent_partition = "nonexistent_partition" test_chunks = MILVUS_INGESTION_IT_CONFIG["corpus"] write_config = MilvusWriteConfig( collection_name=self._collection_name, partition_name=non_existent_partition, write_config=WriteConfig(write_batch_size=1)) config = MilvusVectorWriterConfig( connection_params=self._connection_config, write_config=write_config) # Write pipeline. with self.assertRaises(Exception) as context: with TestPipeline() as p: p.not_use_test_runner_api = True _ = (p | beam.Create(test_chunks) | config.create_write_transform()) # Assert on what should happen. self.assertIn("partition not found", str(context.exception).lower()) def test_invalid_write_on_missing_primary_key_in_entity(self): test_chunks = [ Chunk( content=Content(text="Test content without ID"), embedding=Embedding( dense_embedding=[0.1, 0.2, 0.3], sparse_embedding=([1, 2], [0.1, 0.2])), metadata={"source": "test"}) ] write_config = MilvusWriteConfig( collection_name=self._collection_name, partition_name=self._partition_name, write_config=WriteConfig(write_batch_size=1)) # Deliberately remove id primary key from the entity. specs = MilvusVectorWriterConfig.default_column_specs() for i, spec in enumerate(specs): if spec.column_name == "id": del specs[i] break config = MilvusVectorWriterConfig( connection_params=self._connection_config, write_config=write_config, column_specs=specs) # Write pipeline. with self.assertRaises(Exception) as context: with TestPipeline() as p: p.not_use_test_runner_api = True _ = (p | beam.Create(test_chunks) | config.create_write_transform()) # Assert on what should happen. self.assertIn( "insert missed an field `id` to collection", str(context.exception).lower()) def test_write_on_auto_id_primary_key(self): auto_id_collection = f"auto_id_collection_{self._testMethodName}" auto_id_partition = f"auto_id_partition_{self._testMethodName}" auto_id_fields = [ FieldSchema( name="id", dtype=DataType.INT64, is_primary=True, auto_id=True), FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=1000), FieldSchema(name="metadata", dtype=DataType.JSON), FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=3), FieldSchema( name="sparse_embedding", dtype=DataType.SPARSE_FLOAT_VECTOR) ] # Create collection with an auto id field. create_collection_with_partition( client=self._test_client, collection_name=auto_id_collection, partition_name=auto_id_partition, fields=auto_id_fields) test_chunks = [ Chunk( id=1, content=Content(text="Test content without ID"), embedding=Embedding( dense_embedding=[0.1, 0.2, 0.3], sparse_embedding=([1, 2], [0.1, 0.2])), metadata={"source": "test"}) ] write_config = MilvusWriteConfig( collection_name=auto_id_collection, partition_name=auto_id_partition, write_config=WriteConfig(write_batch_size=1)) config = MilvusVectorWriterConfig( connection_params=self._connection_config, write_config=write_config) self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = (p | beam.Create(test_chunks) | config.create_write_transform()) self._test_client.flush(auto_id_collection) self._test_client.load_collection(auto_id_collection) result = self._test_client.query( collection_name=auto_id_collection, partition_names=[auto_id_partition], limit=3) # Test there is only one item in the result and the ID is not equal to one. self.assertEqual(len(result), len(test_chunks)) result_item = dict(result[0]) self.assertNotEqual(result_item["id"], 1) def test_write_on_existent_collection_with_default_schema(self): test_chunks = MILVUS_INGESTION_IT_CONFIG["corpus"] write_config = MilvusWriteConfig( collection_name=self._collection_name, partition_name=self._partition_name, write_config=WriteConfig(write_batch_size=3)) config = MilvusVectorWriterConfig( connection_params=self._connection_config, write_config=write_config) self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = (p | beam.Create(test_chunks) | config.create_write_transform()) # Verify data was written successfully. self._test_client.flush(self._collection_name) self._test_client.load_collection(self._collection_name) result = self._test_client.query( collection_name=self._collection_name, partition_names=[self._partition_name], limit=10) self.assertEqual(len(result), len(test_chunks)) # Verify each chunk was written correctly. result_by_id = {item["id"]: item for item in result} for chunk in test_chunks: self.assertIn(chunk.id, result_by_id) result_item = result_by_id[chunk.id] self.assertEqual(result_item["content"], chunk.content.text) self.assertEqual(result_item["metadata"], chunk.metadata) # Verify embedding is present and has correct length. expected_embedding = chunk.embedding.dense_embedding actual_embedding = result_item["embedding"] self.assertIsNotNone(actual_embedding) self.assertEqual(len(actual_embedding), len(expected_embedding)) def test_write_with_custom_column_specifications(self): from apache_beam.ml.rag.ingestion.postgres_common import ColumnSpec from apache_beam.ml.rag.utils import MilvusHelpers custom_column_specs = [ ColumnSpec("id", int, lambda chunk: int(chunk.id) if chunk.id else 0), ColumnSpec("content", str, lambda chunk: chunk.content.text), ColumnSpec("metadata", dict, lambda chunk: chunk.metadata or {}), ColumnSpec( "embedding", list, lambda chunk: chunk.embedding.dense_embedding or []), ColumnSpec( "sparse_embedding", dict, lambda chunk: ( MilvusHelpers.sparse_embedding( chunk.embedding.sparse_embedding) if chunk.embedding and chunk.embedding.sparse_embedding else {})) ] test_chunks = [ Chunk( id=10, content=Content(text="Custom column spec test"), embedding=Embedding( dense_embedding=[0.8, 0.9, 1.0], sparse_embedding=([1, 3, 5], [0.8, 0.9, 1.0])), metadata={"custom": "spec_test"}) ] write_config = MilvusWriteConfig( collection_name=self._collection_name, partition_name=self._partition_name, write_config=WriteConfig(write_batch_size=1)) config = MilvusVectorWriterConfig( connection_params=self._connection_config, write_config=write_config, column_specs=custom_column_specs) self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = (p | beam.Create(test_chunks) | config.create_write_transform()) # Verify data was written successfully. self._test_client.flush(self._collection_name) self._test_client.load_collection(self._collection_name) result = self._test_client.query( collection_name=self._collection_name, partition_names=[self._partition_name], filter="id == 10", limit=1) self.assertEqual(len(result), 1) result_item = result[0] # Verify custom column specs worked correctly. self.assertEqual(result_item["id"], 10) self.assertEqual(result_item["content"], "Custom column spec test") self.assertEqual(result_item["metadata"], {"custom": "spec_test"}) # Verify embedding is present and has correct length. expected_embedding = [0.8, 0.9, 1.0] actual_embedding = result_item["embedding"] self.assertIsNotNone(actual_embedding) self.assertEqual(len(actual_embedding), len(expected_embedding)) # Verify sparse embedding was converted correctly - check keys are present. expected_sparse_keys = {1, 3, 5} actual_sparse = result_item["sparse_embedding"] self.assertIsNotNone(actual_sparse) self.assertEqual(set(actual_sparse.keys()), expected_sparse_keys) def test_write_with_batching(self): test_chunks = [ Chunk( id=i, content=Content(text=f"Batch test document {i}"), embedding=Embedding( dense_embedding=[0.1 * i, 0.2 * i, 0.3 * i], sparse_embedding=([i, i + 1], [0.1 * i, 0.2 * i])), metadata={"batch_id": i}) for i in range(1, 8) # 7 chunks ] # Set small batch size to force batching (7 chunks with batch size 3). batch_write_config = WriteConfig(write_batch_size=3) write_config = MilvusWriteConfig( collection_name=self._collection_name, partition_name=self._partition_name, write_config=batch_write_config) config = MilvusVectorWriterConfig( connection_params=self._connection_config, write_config=write_config) self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = (p | beam.Create(test_chunks) | config.create_write_transform()) # Verify all data was written successfully. # Flush to persist all data to disk, then load collection for querying. self._test_client.flush(self._collection_name) self._test_client.load_collection(self._collection_name) result = self._test_client.query( collection_name=self._collection_name, partition_names=[self._partition_name], limit=10) self.assertEqual(len(result), len(test_chunks)) # Verify each batch was written correctly. result_by_id = {item["id"]: item for item in result} for chunk in test_chunks: self.assertIn(chunk.id, result_by_id) result_item = result_by_id[chunk.id] # Verify content and metadata. self.assertEqual(result_item["content"], chunk.content.text) self.assertEqual(result_item["metadata"], chunk.metadata) # Verify embeddings are present and have correct length. expected_embedding = chunk.embedding.dense_embedding actual_embedding = result_item["embedding"] self.assertIsNotNone(actual_embedding) self.assertEqual(len(actual_embedding), len(expected_embedding)) # Verify sparse embedding keys are present. expected_sparse_keys = {chunk.id, chunk.id + 1} actual_sparse = result_item["sparse_embedding"] self.assertIsNotNone(actual_sparse) self.assertEqual(set(actual_sparse.keys()), expected_sparse_keys) def test_idempotent_write(self): # Step 1: Insert initial data that doesn't exist. initial_chunks = [ Chunk( id=100, content=Content(text="Initial document"), embedding=Embedding( dense_embedding=[1.0, 2.0, 3.0], sparse_embedding=([100, 101], [1.0, 2.0])), metadata={"version": 1}), Chunk( id=200, content=Content(text="Another initial document"), embedding=Embedding( dense_embedding=[2.0, 3.0, 4.0], sparse_embedding=([200, 201], [2.0, 3.0])), metadata={"version": 1}) ] write_config = MilvusWriteConfig( collection_name=self._collection_name, partition_name=self._partition_name, write_config=WriteConfig(write_batch_size=2)) config = MilvusVectorWriterConfig( connection_params=self._connection_config, write_config=write_config) # Insert initial data. with TestPipeline() as p: p.not_use_test_runner_api = True _ = ( p | "Create initial" >> beam.Create(initial_chunks) | "Write initial" >> config.create_write_transform()) # Verify initial data was inserted (not existed before). self._test_client.flush(self._collection_name) self._test_client.load_collection(self._collection_name) result = self._test_client.query( collection_name=self._collection_name, partition_names=[self._partition_name], limit=10) self.assertEqual(len(result), 2) result_by_id = {item["id"]: item for item in result} # Verify initial state. self.assertEqual(result_by_id[100]["content"], "Initial document") self.assertEqual(result_by_id[100]["metadata"], {"version": 1}) self.assertEqual(result_by_id[200]["content"], "Another initial document") self.assertEqual(result_by_id[200]["metadata"], {"version": 1}) # Step 2: Update existing data (same IDs, different content). updated_chunks = [ Chunk( id=100, content=Content(text="Updated document"), embedding=Embedding( dense_embedding=[1.1, 2.1, 3.1], sparse_embedding=([100, 102], [1.1, 2.1])), metadata={"version": 2}), Chunk( id=200, content=Content(text="Another updated document"), embedding=Embedding( dense_embedding=[2.1, 3.1, 4.1], sparse_embedding=([200, 202], [2.1, 3.1])), metadata={"version": 2}) ] # Perform first update. with TestPipeline() as p: p.not_use_test_runner_api = True _ = ( p | "Create update1" >> beam.Create(updated_chunks) | "Write update1" >> config.create_write_transform()) # Verify update worked. self._test_client.flush(self._collection_name) self._test_client.load_collection(self._collection_name) result = self._test_client.query( collection_name=self._collection_name, partition_names=[self._partition_name], limit=10) self.assertEqual(len(result), 2) # Still only 2 records. result_by_id = {item["id"]: item for item in result} # Verify updated state. self.assertEqual(result_by_id[100]["content"], "Updated document") self.assertEqual(result_by_id[100]["metadata"], {"version": 2}) self.assertEqual(result_by_id[200]["content"], "Another updated document") self.assertEqual(result_by_id[200]["metadata"], {"version": 2}) # Step 3: Repeat the same update operation 3 more times (idempotence test). for i in range(3): with TestPipeline() as p: p.not_use_test_runner_api = True _ = ( p | f"Create repeat{i+2}" >> beam.Create(updated_chunks) | f"Write repeat{i+2}" >> config.create_write_transform()) # Verify state hasn't changed after repeated updates. self._test_client.flush(self._collection_name) self._test_client.load_collection(self._collection_name) result = self._test_client.query( collection_name=self._collection_name, partition_names=[self._partition_name], limit=10) # Still only 2 records. self.assertEqual(len(result), 2) result_by_id = {item["id"]: item for item in result} # Final state should remain unchanged. self.assertEqual(result_by_id[100]["content"], "Updated document") self.assertEqual(result_by_id[100]["metadata"], {"version": 2}) self.assertEqual(result_by_id[200]["content"], "Another updated document") self.assertEqual(result_by_id[200]["metadata"], {"version": 2}) # Verify embeddings are still correct. self.assertIsNotNone(result_by_id[100]["embedding"]) self.assertEqual(len(result_by_id[100]["embedding"]), 3) self.assertIsNotNone(result_by_id[200]["embedding"]) self.assertEqual(len(result_by_id[200]["embedding"]), 3) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/milvus_search_test.py0000644000175100017510000001161415206113032025315 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest from parameterized import parameterized try: from apache_beam.ml.rag.ingestion.milvus_search import MilvusVectorWriterConfig from apache_beam.ml.rag.ingestion.milvus_search import MilvusWriteConfig from apache_beam.ml.rag.utils import MilvusConnectionParameters except ImportError as e: raise unittest.SkipTest(f'Milvus dependencies not installed: {str(e)}') class TestMilvusWriteConfig(unittest.TestCase): """Unit tests for MilvusWriteConfig validation errors.""" def test_empty_collection_name_raises_error(self): """Test that empty collection name raises ValueError.""" with self.assertRaises(ValueError) as context: MilvusWriteConfig(collection_name="") self.assertIn("Collection name must be provided", str(context.exception)) def test_none_collection_name_raises_error(self): """Test that None collection name raises ValueError.""" with self.assertRaises(ValueError) as context: MilvusWriteConfig(collection_name=None) self.assertIn("Collection name must be provided", str(context.exception)) class TestMilvusVectorWriterConfig(unittest.TestCase): """Unit tests for MilvusVectorWriterConfig validation and functionality.""" def test_valid_config_creation(self): """Test creation of valid MilvusVectorWriterConfig.""" connection_params = MilvusConnectionParameters(uri="http://localhost:19530") write_config = MilvusWriteConfig(collection_name="test_collection") config = MilvusVectorWriterConfig( connection_params=connection_params, write_config=write_config) self.assertEqual(config.connection_params, connection_params) self.assertEqual(config.write_config, write_config) self.assertIsNotNone(config.column_specs) def test_create_converter_returns_callable(self): """Test that create_converter returns a callable function.""" connection_params = MilvusConnectionParameters(uri="http://localhost:19530") write_config = MilvusWriteConfig(collection_name="test_collection") config = MilvusVectorWriterConfig( connection_params=connection_params, write_config=write_config) converter = config.create_converter() self.assertTrue(callable(converter)) def test_create_write_transform_returns_ptransform(self): """Test that create_write_transform returns a PTransform.""" connection_params = MilvusConnectionParameters(uri="http://localhost:19530") write_config = MilvusWriteConfig(collection_name="test_collection") config = MilvusVectorWriterConfig( connection_params=connection_params, write_config=write_config) transform = config.create_write_transform() self.assertIsNotNone(transform) def test_default_column_specs_has_expected_fields(self): """Test that default column specs include expected fields.""" column_specs = MilvusVectorWriterConfig.default_column_specs() self.assertIsInstance(column_specs, list) self.assertGreater(len(column_specs), 0) column_names = [spec.column_name for spec in column_specs] expected_fields = [ "id", "embedding", "sparse_embedding", "content", "metadata" ] for field in expected_fields: self.assertIn(field, column_names) @parameterized.expand([ # Invalid connection parameters - empty URI. ( lambda: ( MilvusConnectionParameters(uri=""), MilvusWriteConfig( collection_name="test_collection")), "URI must be provided"), # Invalid write config - empty collection name. ( lambda: ( MilvusConnectionParameters(uri="http://localhost:19530"), MilvusWriteConfig(collection_name="")), "Collection name must be provided"), ]) def test_invalid_configuration_parameters( self, create_params, expected_error_msg): """Test validation errors for invalid configuration parameters.""" with self.assertRaises(ValueError) as context: connection_params, write_config = create_params() MilvusVectorWriterConfig( connection_params=connection_params, write_config=write_config) self.assertIn(expected_error_msg, str(context.exception)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/mysql.py0000644000175100017510000002413715206113032022563 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging from abc import ABC from abc import abstractmethod from typing import Callable from typing import NamedTuple from typing import Optional import apache_beam as beam from apache_beam.coders import registry from apache_beam.coders.row_coder import RowCoder from apache_beam.io.jdbc import WriteToJdbc from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteConfig from apache_beam.ml.rag.ingestion.jdbc_common import ConnectionConfig from apache_beam.ml.rag.ingestion.jdbc_common import WriteConfig from apache_beam.ml.rag.ingestion.mysql_common import ColumnSpec from apache_beam.ml.rag.ingestion.mysql_common import ColumnSpecsBuilder from apache_beam.ml.rag.ingestion.mysql_common import ConflictResolution from apache_beam.ml.rag.types import EmbeddableItem _LOGGER = logging.getLogger(__name__) class _ConflictResolutionStrategy(ABC): """Abstract base class for conflict resolution strategies.""" @abstractmethod def get_conflict_clause(self, all_columns: list[str]) -> str: """Generate the MySQL conflict clause.""" pass class _NoConflictStrategy(_ConflictResolutionStrategy): """Strategy for when no conflict resolution is needed.""" def get_conflict_clause(self, all_columns: list[str]) -> str: return "" class _UpdateStrategy(_ConflictResolutionStrategy): """Strategy for UPDATE action on conflict.""" def __init__(self, update_fields: Optional[list[str]] = None): self.update_fields = update_fields def get_conflict_clause(self, all_columns: list[str]) -> str: # Use provided fields or default to all columns fields_to_update = self.update_fields or all_columns assert len(fields_to_update) > 0 updates = [f"{field} = VALUES({field})" for field in fields_to_update] return f"ON DUPLICATE KEY UPDATE {', '.join(updates)}" class _IgnoreStrategy(_ConflictResolutionStrategy): """Strategy for IGNORE action on conflict.""" def __init__(self, primary_key_field: str): self.primary_key_field = primary_key_field def get_conflict_clause(self, all_columns: list[str]) -> str: return f"ON DUPLICATE KEY UPDATE {self.primary_key_field}"\ f" = {self.primary_key_field}" def _create_conflict_strategy( conflict_resolution: Optional[ConflictResolution] ) -> _ConflictResolutionStrategy: if conflict_resolution is None: return _NoConflictStrategy() if conflict_resolution.action == "UPDATE": return _UpdateStrategy(conflict_resolution.update_fields) if conflict_resolution.action == "IGNORE": assert conflict_resolution.primary_key_field is not None return _IgnoreStrategy(conflict_resolution.primary_key_field) raise ValueError(f"Unknown conflict resolution {conflict_resolution.action}") class _MySQLQueryBuilder: def __init__( self, table_name: str, *, column_specs: list[ColumnSpec], conflict_resolution: Optional[ConflictResolution] = None): """Builds SQL queries for writing EmbeddableItems with Embeddings to MySQL. """ self.table_name = table_name self.column_specs = column_specs self.conflict_resolution_strategy = _create_conflict_strategy( conflict_resolution) names = [col.column_name for col in self.column_specs] duplicates = set(name for name in names if names.count(name) > 1) if duplicates: raise ValueError(f"Duplicate column names found: {duplicates}") fields = [(col.column_name, col.python_type) for col in self.column_specs] type_name = f"VectorRecord_{table_name}" self.record_type = NamedTuple(type_name, fields) # type: ignore registry.register_coder(self.record_type, RowCoder) def build_insert(self) -> str: fields = [col.column_name for col in self.column_specs] placeholders = [col.placeholder for col in self.column_specs] # Build base query query = f""" INSERT INTO {self.table_name} ({', '.join(fields)}) VALUES ({', '.join(placeholders)}) """ conflict_clause = self.conflict_resolution_strategy.get_conflict_clause( all_columns=fields) query += f" {conflict_clause}" _LOGGER.info("MySQL Query with placeholders %s", query) return query def create_converter(self) -> Callable[[EmbeddableItem], NamedTuple]: """Creates a function to convert EmbeddableItems to records.""" def convert(chunk: EmbeddableItem) -> self.record_type: # type: ignore return self.record_type( **{col.column_name: col.value_fn(chunk) for col in self.column_specs}) # type: ignore return convert class MySQLVectorWriterConfig(VectorDatabaseWriteConfig): def __init__( self, connection_config: ConnectionConfig, table_name: str, *, # pylint: disable=dangerous-default-value write_config: WriteConfig = WriteConfig(), column_specs: list[ColumnSpec] = ColumnSpecsBuilder.with_defaults().build( ), conflict_resolution: Optional[ConflictResolution] = None): """Configuration for writing vectors to MySQL using jdbc. Supports flexible schema configuration through column specifications and conflict resolution strategies with MySQL-specific syntax. Args: connection_config: :class:`~apache_beam.ml.rag.ingestion.jdbc_common.ConnectionConfig`. table_name: Target table name. write_config: JdbcIO :class:`~.jdbc_common.WriteConfig` to control batch sizes, authosharding, etc. column_specs: Use :class:`~.mysql_common.ColumnSpecsBuilder` to configure how embeddings and metadata are written to the database schema. If None, uses default EmbeddableItem schema with MySQL vector functions. conflict_resolution: Optional :class:`~.mysql_common.ConflictResolution` strategy for handling insert conflicts. ON DUPLICATE KEY UPDATE. None by default, meaning errors are thrown when attempting to insert duplicates. Examples: Simple case with default schema: >>> config = MySQLVectorWriterConfig( ... connection_config=ConnectionConfig(...), ... table_name='embeddings' ... ) Custom schema with metadata fields and MySQL functions: >>> specs = (ColumnSpecsBuilder() ... .with_id_spec(column_name="my_id_column") ... .with_embedding_spec( ... column_name="embedding_vec", ... placeholder="string_to_vector(?)" ... ) ... .add_metadata_field(field="source", column_name="src") ... .add_metadata_field( ... "timestamp", ... column_name="created_at", ... placeholder="STR_TO_DATE(?, '%Y-%m-%d %H:%i:%s')" ... ) ... .build()) Minimal schema (only ID + embedding written): >>> column_specs = (ColumnSpecsBuilder() ... .with_id_spec() ... .with_embedding_spec() ... .build()) >>> config = MySQLVectorWriterConfig( ... connection_config=ConnectionConfig(...), ... table_name='embeddings', ... column_specs=specs, ... conflict_resolution=ConflictResolution( ... on_conflict_fields=["id"], ... action="UPDATE", ... update_fields=["embedding", "content"] ... ) ... ) Using MySQL JSON functions: >>> specs = (ColumnSpecsBuilder() ... .with_id_spec() ... .with_embedding_spec() ... .with_metadata_spec( ... column_name="metadata_json", ... placeholder="CAST(? AS JSON)" ... ) ... .build()) """ self.connection_config = connection_config self.write_config = write_config # NamedTuple is created and registered here during pipeline construction self.query_builder = _MySQLQueryBuilder( table_name, column_specs=column_specs, conflict_resolution=conflict_resolution) def create_write_transform(self) -> beam.PTransform: return _WriteToMySQLVectorDatabase(self) class _WriteToMySQLVectorDatabase(beam.PTransform): """Implementation of MySQL vector database write.""" def __init__(self, config: MySQLVectorWriterConfig): self.config = config self.query_builder = config.query_builder self.connection_config = config.connection_config self.write_config = config.write_config def expand(self, pcoll: beam.PCollection[EmbeddableItem]): return ( pcoll | "Convert to Records" >> beam.Map(self.query_builder.create_converter()) | "Write to MySQL" >> WriteToJdbc( table_name=self.query_builder.table_name, driver_class_name="com.mysql.cj.jdbc.Driver", jdbc_url=self.connection_config.jdbc_url, username=self.connection_config.username, password=self.connection_config.password, statement=self.query_builder.build_insert(), connection_properties=self.connection_config.connection_properties, connection_init_sqls=self.connection_config.connection_init_sqls, autosharding=self.write_config.autosharding, max_connections=self.write_config.max_connections, write_batch_size=self.write_config.write_batch_size, **self.connection_config.additional_jdbc_args)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/mysql_common.py0000644000175100017510000003547315206113032024140 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import json from dataclasses import dataclass from typing import Any from typing import Callable from typing import Literal from typing import Optional from apache_beam.ml.rag.types import EmbeddableItem def chunk_embedding_fn(chunk: EmbeddableItem) -> str: """Convert embedding to MySQL vector string format. Formats dense embedding as a MySQL-compatible vector string. Example: [1.0, 2.0] -> '[1.0,2.0]' Args: chunk: Input EmbeddableItem object. Returns: str: MySQL vector string representation of the embedding. Raises: ValueError: If chunk has no dense embedding. """ if chunk.embedding is None or chunk.embedding.dense_embedding is None: raise ValueError(f'Expected chunk to contain embedding. {chunk}') return '[' + ','.join(str(x) for x in chunk.embedding.dense_embedding) + ']' @dataclass class ColumnSpec: """Mapping of EmbeddableItem fields to SQL columns for insertion. Defines how to extract and format values from EmbeddableItems into MySQL database columns, handling the full pipeline from Python value to SQL insertion. The insertion process works as follows: - value_fn extracts a value from the EmbeddableItem and formats it as needed - The value is stored in a NamedTuple field with the specified python_type - During SQL insertion, the value is bound to a ? placeholder Attributes: column_name: The column name in the database table. python_type: :class:`~apache_beam.coders.row_coder.RowCoder` compatible python type. value_fn: Function to extract and format the value from an EmbeddableItem. placeholder: Optional placeholder to apply typecasts or functions to value ? placeholder e.g. "string_to_vector(?)" for vector columns. Examples: Basic text column (uses standard JDBC type mapping): >>> ColumnSpec.text( ... column_name="content", ... value_fn=lambda chunk: chunk.content.text ... ) ... # Results in: INSERT INTO table (content) VALUES (?) Timestamp from metadata: >>> ColumnSpec( ... column_name="created_at", ... python_type=str, ... value_fn=lambda chunk: chunk.metadata.get("timestamp") ... ) ... # Results in: INSERT INTO table (created_at) VALUES (?) Factory Methods: text: Creates a text column specification. integer: Creates an integer column specification. float: Creates a float column specification. vector: Creates a vector column specification with string_to_vector(). json: Creates a JSON column specification. """ column_name: str python_type: type value_fn: Callable[[EmbeddableItem], Any] placeholder: str = '?' @classmethod def text( cls, column_name: str, value_fn: Callable[[EmbeddableItem], Any]) -> 'ColumnSpec': """Create a text column specification.""" return cls(column_name, str, value_fn) @classmethod def integer( cls, column_name: str, value_fn: Callable[[EmbeddableItem], Any]) -> 'ColumnSpec': """Create an integer column specification.""" return cls(column_name, int, value_fn) @classmethod def float( cls, column_name: str, value_fn: Callable[[EmbeddableItem], Any]) -> 'ColumnSpec': """Create a float column specification.""" return cls(column_name, float, value_fn) @classmethod def vector( cls, column_name: str, value_fn: Callable[[EmbeddableItem], Any] = chunk_embedding_fn ) -> 'ColumnSpec': """Create a vector column specification with string_to_vector() function.""" return cls(column_name, str, value_fn, "string_to_vector(?)") @classmethod def json( cls, column_name: str, value_fn: Callable[[EmbeddableItem], Any]) -> 'ColumnSpec': """Create a JSON column specification.""" return cls(column_name, str, value_fn) def embedding_to_string(embedding: list[float]) -> str: """Convert embedding to MySQL vector string format.""" return '[' + ','.join(str(x) for x in embedding) + ']' class ColumnSpecsBuilder: """Builder for :class:`.ColumnSpec`'s with chainable methods.""" def __init__(self): self._specs: list[ColumnSpec] = [] @staticmethod def with_defaults() -> 'ColumnSpecsBuilder': """Add all default column specifications.""" return ( ColumnSpecsBuilder().with_id_spec().with_embedding_spec(). with_content_spec().with_metadata_spec()) def with_id_spec( self, column_name: str = "id", python_type: type = str, convert_fn: Optional[Callable[[str], Any]] = None) -> 'ColumnSpecsBuilder': """Add ID :class:`.ColumnSpec` with optional type and conversion. Args: column_name: Name for the ID column (defaults to "id") python_type: Python type for the column (defaults to str) convert_fn: Optional function to convert the chunk ID If None, uses ID as-is Returns: Self for method chaining Example: >>> builder.with_id_spec( ... column_name="doc_id", ... python_type=int, ... convert_fn=lambda id: int(id.split('_')[1]) ... ) """ def value_fn(chunk: EmbeddableItem) -> Any: value = chunk.id return convert_fn(value) if convert_fn else value self._specs.append( ColumnSpec( column_name=column_name, python_type=python_type, value_fn=value_fn)) return self def with_content_spec( self, column_name: str = "content", python_type: type = str, convert_fn: Optional[Callable[[str], Any]] = None) -> 'ColumnSpecsBuilder': """Add content :class:`.ColumnSpec` with optional type and conversion. Args: column_name: Name for the content column (defaults to "content") python_type: Python type for the column (defaults to str) convert_fn: Optional function to convert the content text If None, uses content text as-is Returns: Self for method chaining Example: >>> builder.with_content_spec( ... column_name="content_length", ... python_type=int, ... convert_fn=len # Store content length instead of content ... ) """ def value_fn(chunk: EmbeddableItem) -> Any: value = chunk.content_string return convert_fn(value) if convert_fn else value self._specs.append( ColumnSpec( column_name=column_name, python_type=python_type, value_fn=value_fn)) return self def with_metadata_spec( self, column_name: str = "metadata", python_type: type = str, convert_fn: Optional[Callable[[dict[str, Any]], Any]] = None ) -> 'ColumnSpecsBuilder': """Add metadata :class:`.ColumnSpec` with optional type and conversion. Args: column_name: Name for the metadata column (defaults to "metadata") python_type: Python type for the column (defaults to str) convert_fn: Optional function to convert the metadata dictionary If None and python_type is str, converts to JSON string Returns: Self for method chaining Example: >>> builder.with_metadata_spec( ... column_name="meta_tags", ... python_type=str, ... convert_fn=lambda meta: ','.join(meta.keys()) ... ) """ def value_fn(chunk: EmbeddableItem) -> Any: if convert_fn: return convert_fn(chunk.metadata) return json.dumps( chunk.metadata) if python_type == str else chunk.metadata self._specs.append( ColumnSpec( column_name=column_name, python_type=python_type, value_fn=value_fn)) return self def with_embedding_spec( self, column_name: str = "embedding", convert_fn: Callable[[list[float]], Any] = embedding_to_string ) -> 'ColumnSpecsBuilder': """Add embedding :class:`.ColumnSpec` with optional conversion. Args: column_name: Name for the embedding column (defaults to "embedding") convert_fn: Optional function to convert the dense embedding values If None, uses default MySQL vector format Returns: Self for method chaining Example: >>> builder.with_embedding_spec( ... column_name="embedding_vector", ... convert_fn=lambda values: '[' + ','.join(f"{x:.4f}" ... for x in values) + ']' ... ) """ def value_fn(chunk: EmbeddableItem) -> Any: if chunk.embedding is None or chunk.embedding.dense_embedding is None: raise ValueError(f'Expected chunk to contain embedding. {chunk}') values = chunk.embedding.dense_embedding return convert_fn(values) self._specs.append( ColumnSpec.vector(column_name=column_name, value_fn=value_fn)) return self def add_metadata_field( self, field: str, python_type: type, column_name: Optional[str] = None, convert_fn: Optional[Callable[[Any], Any]] = None, default: Any = None) -> 'ColumnSpecsBuilder': """Add a :class:`.ColumnSpec` that extracts and converts a field from chunk metadata. Args: field: Key to extract from chunk metadata python_type: Python type for the column (e.g. str, int, float) column_name: Name for the column (defaults to metadata field name) convert_fn: Optional function to convert the extracted value to desired type. If None, value is used as-is default: Default value if field is missing from metadata Returns: Self for chaining Examples: Simple string field: >>> builder.add_metadata_field("source", str) Integer with default: >>> builder.add_metadata_field( ... field="count", ... python_type=int, ... column_name="item_count", ... default=0 ... ) Float with conversion and default: >>> builder.add_metadata_field( ... field="confidence", ... python_type=float, ... convert_fn=lambda x: round(float(x), 2), ... default=0.0 ... ) Timestamp with conversion: >>> builder.add_metadata_field( ... field="created_at", ... python_type=str, ... convert_fn=lambda ts: ts.replace('T', ' ') ... ) """ name = column_name or field def value_fn(chunk: EmbeddableItem) -> Any: value = chunk.metadata.get(field, default) if value is not None and convert_fn is not None: value = convert_fn(value) return value spec = ColumnSpec( column_name=name, python_type=python_type, value_fn=value_fn) self._specs.append(spec) return self def add_custom_column_spec(self, spec: ColumnSpec) -> 'ColumnSpecsBuilder': """Add a custom :class:`.ColumnSpec` to the builder. Use this method when you need complete control over the :class:`.ColumnSpec`, including custom value extraction and type handling. Args: spec: A :class:`.ColumnSpec` instance defining the column name, type, value extraction, and optional MySQL function. Returns: Self for method chaining Examples: Custom text column from chunk metadata: >>> builder.add_custom_column_spec( ... ColumnSpec.text( ... column_name="source_and_id", ... value_fn=lambda chunk: ... f"{chunk.metadata.get('source')}_{chunk.id}" ... ) ... ) """ self._specs.append(spec) return self def build(self) -> list[ColumnSpec]: """Build the final list of column specifications.""" return self._specs.copy() @dataclass class ConflictResolution: """Specification for how to handle conflicts during insert. Configures conflict handling behavior when inserting records that may violate unique constraints using MySQL's ON DUPLICATE KEY UPDATE syntax. MySQL automatically detects conflicts based on PRIMARY KEY or UNIQUE constraints defined on the table. Attributes: action: How to handle conflicts - either "UPDATE" or "IGNORE". UPDATE: Updates existing record with new values. IGNORE: Skips conflicting records (uses no-op update). update_fields: Optional list of fields to update on conflict. If None, all fields are updated (for UPDATE action only). primary_key_field: Required for IGNORE action. The primary key field name to use for the no-op update. Examples: Update all fields on conflict: >>> ConflictResolution(action="UPDATE") Update specific fields on conflict: >>> ConflictResolution( ... action="UPDATE", ... update_fields=["embedding", "content"] ... ) Ignore conflicts with explicit primary key: >>> ConflictResolution( ... action="IGNORE", ... primary_key_field="id" ... ) Ignore conflicts with custom primary key: >>> ConflictResolution( ... action="IGNORE", ... primary_key_field="custom_id" ... ) """ action: Literal["UPDATE", "IGNORE"] = "UPDATE" update_fields: Optional[list[str]] = None primary_key_field: Optional[str] = None def __post_init__(self): """Validate configuration after initialization.""" if self.action == "IGNORE" and self.primary_key_field is None: raise ValueError("primary_key_field is required when action='IGNORE'") ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/postgres.py0000644000175100017510000001765715206113032023275 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import logging from typing import Callable from typing import NamedTuple from typing import Optional from typing import Union import apache_beam as beam from apache_beam.coders import registry from apache_beam.coders.row_coder import RowCoder from apache_beam.io.jdbc import WriteToJdbc from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteConfig from apache_beam.ml.rag.ingestion.jdbc_common import ConnectionConfig from apache_beam.ml.rag.ingestion.jdbc_common import WriteConfig from apache_beam.ml.rag.ingestion.postgres_common import ColumnSpec from apache_beam.ml.rag.ingestion.postgres_common import ColumnSpecsBuilder from apache_beam.ml.rag.ingestion.postgres_common import ConflictResolution from apache_beam.ml.rag.types import EmbeddableItem _LOGGER = logging.getLogger(__name__) MetadataSpec = Union[ColumnSpec, dict[str, ColumnSpec]] class _PostgresQueryBuilder: def __init__( self, table_name: str, *, column_specs: list[ColumnSpec], conflict_resolution: Optional[ConflictResolution] = None): """Builds SQL queries for writing EmbeddableItems to Postgres. """ self.table_name = table_name self.column_specs = column_specs self.conflict_resolution = conflict_resolution # Validate no duplicate column names names = [col.column_name for col in self.column_specs] duplicates = set(name for name in names if names.count(name) > 1) if duplicates: raise ValueError(f"Duplicate column names found: {duplicates}") # Create NamedTuple type fields = [(col.column_name, col.python_type) for col in self.column_specs] type_name = f"VectorRecord_{table_name}" self.record_type = NamedTuple(type_name, fields) # type: ignore # Register coder registry.register_coder(self.record_type, RowCoder) # Set default update fields to all non-conflict fields if update fields are # not specified if self.conflict_resolution: self.conflict_resolution.maybe_set_default_update_fields( [col.column_name for col in self.column_specs if col.column_name]) def build_insert(self) -> str: """Build INSERT query with proper type casting.""" # Get column names and placeholders fields = [col.column_name for col in self.column_specs] placeholders = [col.placeholder for col in self.column_specs] # Build base query query = f""" INSERT INTO {self.table_name} ({', '.join(fields)}) VALUES ({', '.join(placeholders)}) """ # Add conflict handling if configured if self.conflict_resolution: query += f" {self.conflict_resolution.get_conflict_clause()}" _LOGGER.info("Query with placeholders %s", query) return query def create_converter(self) -> Callable[[EmbeddableItem], NamedTuple]: """Creates a function to convert EmbeddableItems to records.""" def convert(chunk: EmbeddableItem) -> self.record_type: # type: ignore return self.record_type( **{col.column_name: col.value_fn(chunk) for col in self.column_specs}) # type: ignore return convert class PostgresVectorWriterConfig(VectorDatabaseWriteConfig): def __init__( self, connection_config: ConnectionConfig, table_name: str, *, # pylint: disable=dangerous-default-value write_config: WriteConfig = WriteConfig(), column_specs: list[ColumnSpec] = ColumnSpecsBuilder.with_defaults().build( ), conflict_resolution: Optional[ConflictResolution] = ConflictResolution( on_conflict_fields=[], action='IGNORE')): """Configuration for writing vectors to Postgres using jdbc. Supports flexible schema configuration through column specifications and conflict resolution strategies. Args: connection_config: :class:`~apache_beam.ml.rag.ingestion.jdbc_common.ConnectionConfig`. table_name: Target table name. write_config: JdbcIO :class:`~.jdbc_common.WriteConfig` to control batch sizes, authosharding, etc. column_specs: Use :class:`~.postgres_common.ColumnSpecsBuilder` to configure how embeddings and metadata are written a database schema. If None, uses default EmbeddableItem schema. conflict_resolution: Optional :class:`~.postgres_common.ConflictResolution` strategy for handling insert conflicts. ON CONFLICT DO NOTHING by default. Examples: Simple case with default schema: >>> config = PostgresVectorWriterConfig( ... connection_config=ConnectionConfig(...), ... table_name='embeddings' ... ) Custom schema with metadata fields: >>> specs = (ColumnSpecsBuilder() ... .with_id_spec(column_name="my_id_column") ... .with_embedding_spec(column_name="embedding_vec") ... .add_metadata_field(field="source", column_name="src") ... .add_metadata_field( ... "timestamp", ... column_name="created_at", ... sql_typecast="::timestamp" ... ) ... .build()) Minimal schema (only ID + embedding written) >>> column_specs = (ColumnSpecsBuilder() ... .with_id_spec() ... .with_embedding_spec() ... .build()) >>> config = PostgresVectorWriterConfig( ... connection_config=ConnectionConfig(...), ... table_name='embeddings', ... column_specs=specs ... ) """ self.connection_config = connection_config self.write_config = write_config # NamedTuple is created and registered here during pipeline construction self.query_builder = _PostgresQueryBuilder( table_name, column_specs=column_specs, conflict_resolution=conflict_resolution) def create_write_transform(self) -> beam.PTransform: return _WriteToPostgresVectorDatabase(self) class _WriteToPostgresVectorDatabase(beam.PTransform): """Implementation of Postgres vector database write. """ def __init__(self, config: PostgresVectorWriterConfig): self.config = config self.query_builder = config.query_builder self.connection_config = config.connection_config self.write_config = config.write_config def expand(self, pcoll: beam.PCollection[EmbeddableItem]): return ( pcoll | "Convert to Records" >> beam.Map(self.query_builder.create_converter()) | "Write to Postgres" >> WriteToJdbc( table_name=self.query_builder.table_name, driver_class_name="org.postgresql.Driver", jdbc_url=self.connection_config.jdbc_url, username=self.connection_config.username, password=self.connection_config.password, statement=self.query_builder.build_insert(), connection_properties=self.connection_config.connection_properties, connection_init_sqls=self.connection_config.connection_init_sqls, autosharding=self.write_config.autosharding, max_connections=self.write_config.max_connections, write_batch_size=self.write_config.write_batch_size, **self.connection_config.additional_jdbc_args)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/postgres_common.py0000644000175100017510000004370715206113032024640 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import json from dataclasses import dataclass from typing import Any from typing import Callable from typing import Literal from typing import Optional from typing import Union from apache_beam.ml.rag.types import EmbeddableItem def chunk_embedding_fn(chunk: EmbeddableItem) -> str: """Convert embedding to PostgreSQL array string. Formats dense embedding as a PostgreSQL-compatible array string. Example: [1.0, 2.0] -> '{1.0,2.0}' Args: chunk: Input EmbeddableItem object. Returns: str: PostgreSQL array string representation of the embedding. Raises: ValueError: If chunk has no dense embedding. """ if chunk.embedding is None or chunk.embedding.dense_embedding is None: raise ValueError(f'Expected chunk to contain embedding. {chunk}') return '{' + ','.join(str(x) for x in chunk.embedding.dense_embedding) + '}' @dataclass class ColumnSpec: """Mapping of EmbeddableItem fields to SQL columns for insertion. Defines how to extract and format values from EmbeddableItems into database columns, handling the full pipeline from Python value to SQL insertion. The insertion process works as follows: - value_fn extracts a value from the EmbeddableItem and formats it as needed - The value is stored in a NamedTuple field with the specified python_type - During SQL insertion, the value is bound to a ? placeholder Attributes: column_name: The column name in the database table. python_type: :class:`~apache_beam.coders.row_coder.RowCoder` compatible python type. value_fn: Function to extract and format the value from an EmbeddableItem. sql_typecast: Optional SQL type cast to append to the ? placeholder. Common examples: - "::float[]" for vector arrays - "::jsonb" for JSON data Examples: Basic text column (uses standard JDBC type mapping): >>> ColumnSpec.text( ... column_name="content", ... value_fn=lambda chunk: chunk.content.text ... ) # Results in: INSERT INTO table (content) VALUES (?) Vector column with explicit array casting: >>> ColumnSpec.vector( ... column_name="embedding", ... value_fn=lambda chunk: '{' + ... ','.join(map(str, chunk.embedding.dense_embedding)) + '}' ... ) # Results in: INSERT INTO table (embedding) VALUES (?::float[]) # The value_fn formats [1.0, 2.0] as '{1.0,2.0}' for PostgreSQL array Timestamp from metadata with explicit casting: >>> ColumnSpec( ... column_name="created_at", ... python_type=str, ... value_fn=lambda chunk: chunk.metadata.get("timestamp"), ... sql_typecast="::timestamp" ... ) # Results in: INSERT INTO table (created_at) VALUES (?::timestamp) # Allows inserting string timestamps with proper PostgreSQL casting Factory Methods: text: Creates a text column specification (no type cast). integer: Creates an integer column specification (no type cast). float: Creates a float column specification (no type cast). vector: Creates a vector column specification with float[] casting. jsonb: Creates a JSONB column specification with jsonb casting. """ column_name: str python_type: type value_fn: Callable[[EmbeddableItem], Any] sql_typecast: Optional[str] = None @property def placeholder(self) -> str: """Get SQL placeholder with optional typecast.""" return f"?{self.sql_typecast or ''}" @classmethod def text( cls, column_name: str, value_fn: Callable[[EmbeddableItem], Any]) -> 'ColumnSpec': """Create a text column specification.""" return cls(column_name, str, value_fn) @classmethod def integer( cls, column_name: str, value_fn: Callable[[EmbeddableItem], Any]) -> 'ColumnSpec': """Create an integer column specification.""" return cls(column_name, int, value_fn) @classmethod def float( cls, column_name: str, value_fn: Callable[[EmbeddableItem], Any]) -> 'ColumnSpec': """Create a float column specification.""" return cls(column_name, float, value_fn) @classmethod def vector( cls, column_name: str, value_fn: Callable[[EmbeddableItem], Any] = chunk_embedding_fn ) -> 'ColumnSpec': """Create a vector column specification.""" return cls(column_name, str, value_fn, "::float[]") @classmethod def jsonb( cls, column_name: str, value_fn: Callable[[EmbeddableItem], Any]) -> 'ColumnSpec': """Create a JSONB column specification.""" return cls(column_name, str, value_fn, "::jsonb") class ColumnSpecsBuilder: """Builder for :class:`.ColumnSpec`'s with chainable methods.""" def __init__(self): self._specs: list[ColumnSpec] = [] @staticmethod def with_defaults() -> 'ColumnSpecsBuilder': """Add all default column specifications.""" return ( ColumnSpecsBuilder().with_id_spec().with_embedding_spec(). with_content_spec().with_metadata_spec()) def with_id_spec( self, column_name: str = "id", python_type: type = str, convert_fn: Optional[Callable[[str], Any]] = None, sql_typecast: Optional[str] = None) -> 'ColumnSpecsBuilder': """Add ID :class:`.ColumnSpec` with optional type and conversion. Args: column_name: Name for the ID column (defaults to "id") python_type: Python type for the column (defaults to str) convert_fn: Optional function to convert the chunk ID If None, uses ID as-is sql_typecast: Optional SQL type cast Returns: Self for method chaining Example: >>> builder.with_id_spec( ... column_name="doc_id", ... python_type=int, ... convert_fn=lambda id: int(id.split('_')[1]) ... ) """ def value_fn(chunk: EmbeddableItem) -> Any: value = chunk.id return convert_fn(value) if convert_fn else value self._specs.append( ColumnSpec( column_name=column_name, python_type=python_type, value_fn=value_fn, sql_typecast=sql_typecast)) return self def with_content_spec( self, column_name: str = "content", python_type: type = str, convert_fn: Optional[Callable[[str], Any]] = None, sql_typecast: Optional[str] = None) -> 'ColumnSpecsBuilder': """Add content :class:`.ColumnSpec` with optional type and conversion. Args: column_name: Name for the content column (defaults to "content") python_type: Python type for the column (defaults to str) convert_fn: Optional function to convert the content text If None, uses content text as-is sql_typecast: Optional SQL type cast Returns: Self for method chaining Example: >>> builder.with_content_spec( ... column_name="content_length", ... python_type=int, ... convert_fn=len # Store content length instead of content ... ) """ def value_fn(chunk: EmbeddableItem) -> Any: value = chunk.content_string return convert_fn(value) if convert_fn else value self._specs.append( ColumnSpec( column_name=column_name, python_type=python_type, value_fn=value_fn, sql_typecast=sql_typecast)) return self def with_metadata_spec( self, column_name: str = "metadata", python_type: type = str, convert_fn: Optional[Callable[[dict[str, Any]], Any]] = None, sql_typecast: Optional[str] = "::jsonb") -> 'ColumnSpecsBuilder': """Add metadata :class:`.ColumnSpec` with optional type and conversion. Args: column_name: Name for the metadata column (defaults to "metadata") python_type: Python type for the column (defaults to str) convert_fn: Optional function to convert the metadata dictionary If None and python_type is str, converts to JSON string sql_typecast: Optional SQL type cast (defaults to "::jsonb") Returns: Self for method chaining Example: >>> builder.with_metadata_spec( ... column_name="meta_tags", ... python_type=list, ... convert_fn=lambda meta: list(meta.keys()), ... sql_typecast="::text[]" ... ) """ def value_fn(chunk: EmbeddableItem) -> Any: if convert_fn: return convert_fn(chunk.metadata) return json.dumps( chunk.metadata) if python_type == str else chunk.metadata self._specs.append( ColumnSpec( column_name=column_name, python_type=python_type, value_fn=value_fn, sql_typecast=sql_typecast)) return self def with_embedding_spec( self, column_name: str = "embedding", convert_fn: Optional[Callable[[list[float]], Any]] = None ) -> 'ColumnSpecsBuilder': """Add embedding :class:`.ColumnSpec` with optional conversion. Args: column_name: Name for the embedding column (defaults to "embedding") convert_fn: Optional function to convert the dense embedding values If None, uses default PostgreSQL array format Returns: Self for method chaining Example: >>> builder.with_embedding_spec( ... column_name="embedding_vector", ... convert_fn=lambda values: '{' + ','.join(f"{x:.4f}" ... for x in values) + '}' ... ) """ def value_fn(chunk: EmbeddableItem) -> Any: if chunk.embedding is None or chunk.embedding.dense_embedding is None: raise ValueError(f'Expected chunk to contain embedding. {chunk}') values = chunk.embedding.dense_embedding if convert_fn: return convert_fn(values) return '{' + ','.join(str(x) for x in values) + '}' self._specs.append( ColumnSpec.vector(column_name=column_name, value_fn=value_fn)) return self def with_sparse_embedding_spec( self, column_name: str = "sparse_embedding", conv_fn: Optional[Callable[[tuple[list[int], list[float]]], Any]] = None ) -> 'ColumnSpecsBuilder': """Add sparse embedding :class:`.ColumnSpec` with optional conversion. Args: column_name: Name for the sparse embedding column (defaults to "sparse_embedding") conv_fn: Optional function to convert the sparse embedding tuple If None, converts to PostgreSQL-compatible JSON format Returns: Self for method chaining Example: >>> builder.with_sparse_embedding_spec( ... column_name="sparse_vector", ... convert_fn=lambda sparse: dict(zip(sparse[0], sparse[1])) ... ) """ def value_fn(chunk: EmbeddableItem) -> Any: if chunk.embedding is None or chunk.embedding.sparse_embedding is None: raise ValueError(f'Expected chunk to contain sparse embedding. {chunk}') sparse_embedding = chunk.embedding.sparse_embedding if conv_fn: return conv_fn(sparse_embedding) # Default: convert to dict format for JSON storage. indices, values = sparse_embedding return json.dumps(dict(zip(indices, values))) self._specs.append( ColumnSpec.jsonb(column_name=column_name, value_fn=value_fn)) return self def add_metadata_field( self, field: str, python_type: type, column_name: Optional[str] = None, convert_fn: Optional[Callable[[Any], Any]] = None, default: Any = None, sql_typecast: Optional[str] = None) -> 'ColumnSpecsBuilder': """""Add a :class:`.ColumnSpec` that extracts and converts a field from chunk metadata. Args: field: Key to extract from chunk metadata python_type: Python type for the column (e.g. str, int, float) column_name: Name for the column (defaults to metadata field name) convert_fn: Optional function to convert the extracted value to desired type. If None, value is used as-is default: Default value if field is missing from metadata sql_typecast: Optional SQL type cast (e.g. "::timestamp") Returns: Self for chaining Examples: Simple string field: >>> builder.add_metadata_field("source", str) Integer with default: >>> builder.add_metadata_field( ... field="count", ... python_type=int, ... column_name="item_count", ... default=0 ... ) Float with conversion and default: >>> builder.add_metadata_field( ... field="confidence", ... python_type=float, ... convert_fn=lambda x: round(float(x), 2), ... default=0.0 ... ) Timestamp with conversion and type cast: >>> builder.add_metadata_field( ... field="created_at", ... python_type=str, ... convert_fn=lambda ts: ts.replace('T', ' '), ... sql_typecast="::timestamp" ... ) """ name = column_name or field def value_fn(chunk: EmbeddableItem) -> Any: value = chunk.metadata.get(field, default) if value is not None and convert_fn is not None: value = convert_fn(value) return value spec = ColumnSpec( column_name=name, python_type=python_type, value_fn=value_fn, sql_typecast=sql_typecast) self._specs.append(spec) return self def add_custom_column_spec(self, spec: ColumnSpec) -> 'ColumnSpecsBuilder': """Add a custom :class:`.ColumnSpec` to the builder. Use this method when you need complete control over the :class:`.ColumnSpec` , including custom value extraction and type handling. Args: spec: A :class:`.ColumnSpec` instance defining the column name, type, value extraction, and optional SQL type casting. Returns: Self for method chaining Examples: Custom text column from chunk metadata: >>> builder.add_custom_column_spec( ... ColumnSpec.text( ... name="source_and_id", ... value_fn=lambda chunk: \ ... f"{chunk.metadata.get('source')}_{chunk.id}" ... ) ... ) """ self._specs.append(spec) return self def build(self) -> list[ColumnSpec]: """Build the final list of column specifications.""" return self._specs.copy() @dataclass class ConflictResolution: """Specification for how to handle conflicts during insert. Configures conflict handling behavior when inserting records that may violate unique constraints. Attributes: on_conflict_fields: Field(s) that determine uniqueness. Can be a single field name or list of field names for composite constraints. action: How to handle conflicts - either "UPDATE" or "IGNORE". UPDATE: Updates existing record with new values. IGNORE: Skips conflicting records. update_fields: Optional list of fields to update on conflict. If None, all non-conflict fields are updated. Examples: Simple primary key: >>> ConflictResolution("id") Composite key with specific update fields: >>> ConflictResolution( ... on_conflict_fields=["source", "timestamp"], ... action="UPDATE", ... update_fields=["embedding", "content"] ... ) Ignore conflicts: >>> ConflictResolution( ... on_conflict_fields="id", ... action="IGNORE" ... ) """ on_conflict_fields: Union[str, list[str]] action: Literal["UPDATE", "IGNORE"] = "UPDATE" update_fields: Optional[list[str]] = None def maybe_set_default_update_fields(self, columns: list[str]): if self.action != "UPDATE": return if self.update_fields is not None: return conflict_fields = ([self.on_conflict_fields] if isinstance( self.on_conflict_fields, str) else self.on_conflict_fields) self.update_fields = [col for col in columns if col not in conflict_fields] def get_conflict_clause(self) -> str: """Get conflict clause with update fields.""" conflict_fields = [self.on_conflict_fields] \ if isinstance(self.on_conflict_fields, str) \ else self.on_conflict_fields if self.action == "IGNORE": conflict_fields_string = f"({', '.join(conflict_fields)})" \ if len(conflict_fields) > 0 else "" return f"ON CONFLICT {conflict_fields_string} DO NOTHING" # update_fields should be set by query builder before this is called assert self.update_fields is not None, \ "update_fields must be set before generating conflict clause" updates = [f"{field} = EXCLUDED.{field}" for field in self.update_fields] return f"ON CONFLICT " \ f"({', '.join(conflict_fields)}) DO UPDATE SET {', '.join(updates)}" ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/postgres_it_test.py0000644000175100017510000010102415206113032025006 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import json import logging import os import secrets import time import unittest from typing import NamedTuple import psycopg2 import pytest import apache_beam as beam from apache_beam.coders import registry from apache_beam.coders.row_coder import RowCoder from apache_beam.io.jdbc import ReadFromJdbc from apache_beam.ml.rag.ingestion import test_utils from apache_beam.ml.rag.ingestion.jdbc_common import ConnectionConfig from apache_beam.ml.rag.ingestion.jdbc_common import WriteConfig from apache_beam.ml.rag.ingestion.postgres import PostgresVectorWriterConfig from apache_beam.ml.rag.ingestion.postgres_common import ColumnSpec from apache_beam.ml.rag.ingestion.postgres_common import ColumnSpecsBuilder from apache_beam.ml.rag.ingestion.postgres_common import ConflictResolution from apache_beam.ml.rag.ingestion.postgres_common import chunk_embedding_fn from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import Embedding from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to CustomSpecsRow = NamedTuple( 'CustomSpecsRow', [ ('custom_id', str), # For id_spec test ('embedding_vec', list[float]), # For embedding_spec test ('content_col', str), # For content_spec test ('metadata', str) ]) registry.register_coder(CustomSpecsRow, RowCoder) MetadataConflictRow = NamedTuple( 'MetadataConflictRow', [ ('id', str), ('source', str), # For metadata_spec and composite key ('timestamp', str), # For metadata_spec and composite key ('content', str), ('embedding', list[float]), ('metadata', str) ]) registry.register_coder(MetadataConflictRow, RowCoder) @pytest.mark.uses_gcp_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_JARS'), "EXPANSION_JARS environment var is not provided, " "indicating that jars have not been built") @unittest.skipUnless( os.environ.get('ALLOYDB_PASSWORD'), "ALLOYDB_PASSWORD environment var is not provided") class PostgresVectorWriterConfigTest(unittest.TestCase): POSTGRES_TABLE_PREFIX = 'python_rag_postgres_' @classmethod def setUpClass(cls): cls.host = os.environ.get('ALLOYDB_HOST', '10.119.0.22') cls.port = os.environ.get('ALLOYDB_PORT', '5432') cls.database = os.environ.get('ALLOYDB_DATABASE', 'postgres') cls.username = os.environ.get('ALLOYDB_USERNAME', 'postgres') if not os.environ.get('ALLOYDB_PASSWORD'): raise ValueError('ALLOYDB_PASSWORD env not set') cls.password = os.environ.get('ALLOYDB_PASSWORD') # Create unique table name suffix cls.table_suffix = '%d%s' % (int(time.time()), secrets.token_hex(3)) # Setup database connection cls.conn = psycopg2.connect( host=cls.host, port=cls.port, database=cls.database, user=cls.username, password=cls.password) cls.conn.autocommit = True def skip_if_dataflow_runner(self): if self._runner and "dataflowrunner" in self._runner.lower(): self.skipTest( "Skipping some tests on Dataflow Runner to avoid bloat and timeouts") def setUp(self): self.write_test_pipeline = TestPipeline(is_integration_test=True) self.read_test_pipeline = TestPipeline(is_integration_test=True) self.write_test_pipeline2 = TestPipeline(is_integration_test=True) self.read_test_pipeline2 = TestPipeline(is_integration_test=True) self._runner = type(self.read_test_pipeline.runner).__name__ self.default_table_name = f"{self.POSTGRES_TABLE_PREFIX}" \ f"{self.table_suffix}" self.default_table_name = f"{self.POSTGRES_TABLE_PREFIX}" \ f"{self.table_suffix}" self.custom_table_name = f"{self.POSTGRES_TABLE_PREFIX}" \ f"_custom_{self.table_suffix}" self.metadata_conflicts_table = f"{self.POSTGRES_TABLE_PREFIX}" \ f"_meta_conf_{self.table_suffix}" self.jdbc_url = f'jdbc:postgresql://{self.host}:{self.port}/{self.database}' # Create test table with self.conn.cursor() as cursor: cursor.execute( f""" CREATE TABLE {self.default_table_name} ( id TEXT PRIMARY KEY, embedding VECTOR({test_utils.VECTOR_SIZE}), content TEXT, metadata JSONB ) """) cursor.execute( f""" CREATE TABLE {self.custom_table_name} ( custom_id TEXT PRIMARY KEY, embedding_vec VECTOR(2), content_col TEXT, metadata JSONB ) """) cursor.execute( f""" CREATE TABLE {self.metadata_conflicts_table} ( id TEXT, source TEXT, timestamp TIMESTAMP, content TEXT, embedding VECTOR(2), PRIMARY KEY (id), UNIQUE (source, timestamp) ) """) _LOGGER = logging.getLogger(__name__) _LOGGER.info("Created table %s", self.default_table_name) def tearDown(self): # Drop test table with self.conn.cursor() as cursor: cursor.execute(f"DROP TABLE IF EXISTS {self.default_table_name}") cursor.execute(f"DROP TABLE IF EXISTS {self.custom_table_name}") cursor.execute(f"DROP TABLE IF EXISTS {self.metadata_conflicts_table}") _LOGGER = logging.getLogger(__name__) _LOGGER.info("Dropped table %s", self.default_table_name) @classmethod def tearDownClass(cls): if hasattr(cls, 'conn'): cls.conn.close() def test_default_schema(self): """Test basic write with default schema.""" jdbc_url = f'jdbc:postgresql://{self.host}:{self.port}/{self.database}' connection_config = ConnectionConfig( jdbc_url=jdbc_url, username=self.username, password=self.password) config = PostgresVectorWriterConfig( connection_config=connection_config, table_name=self.default_table_name, write_config=WriteConfig(write_batch_size=100)) # Create test chunks num_records = 1500 sample_size = min(500, num_records // 2) # Generate test chunks chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records) # Run pipeline and verify self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = (p | beam.Create(chunks) | config.create_write_transform()) self.read_test_pipeline.not_use_test_runner_api = True # Read pipeline to verify read_query = f""" SELECT CAST(id AS VARCHAR(255)), CAST(content AS VARCHAR(255)), CAST(embedding AS text), CAST(metadata AS text) FROM {self.default_table_name} """ # Read and verify pipeline with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=jdbc_url, username=self.username, password=self.password, query=read_query)) count_result = rows | "Count All" >> beam.combiners.Count.Globally() assert_that(count_result, equal_to([num_records]), label='count_check') chunks = (rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk)) chunk_hashes = chunks | "Hash Chunks" >> beam.CombineGlobally( test_utils.HashingFn()) assert_that( chunk_hashes, equal_to([test_utils.generate_expected_hash(num_records)]), label='hash_check') # Sample validation first_n = ( chunks | "Key on Index" >> beam.Map(test_utils.key_on_id) | f"Get First {sample_size}" >> beam.transforms.combiners.Top.Of( sample_size, key=lambda x: x[0], reverse=True) | "Remove Keys 1" >> beam.Map(lambda xs: [x[1] for x in xs])) expected_first_n = test_utils.ChunkTestUtils.get_expected_values( 0, sample_size) assert_that( first_n, equal_to([expected_first_n]), label=f"first_{sample_size}_check") last_n = ( chunks | "Key on Index 2" >> beam.Map(test_utils.key_on_id) | f"Get Last {sample_size}" >> beam.transforms.combiners.Top.Of( sample_size, key=lambda x: x[0]) | "Remove Keys 2" >> beam.Map(lambda xs: [x[1] for x in xs])) expected_last_n = test_utils.ChunkTestUtils.get_expected_values( num_records - sample_size, num_records)[::-1] assert_that( last_n, equal_to([expected_last_n]), label=f"last_{sample_size}_check") def test_custom_specs(self): """Test custom specifications for ID, embedding, and content.""" self.skip_if_dataflow_runner() num_records = 20 specs = ( ColumnSpecsBuilder().add_custom_column_spec( ColumnSpec.text( column_name="custom_id", value_fn=lambda chunk: f"timestamp_{chunk.metadata.get('timestamp', '')}") ).add_custom_column_spec( ColumnSpec.vector( column_name="embedding_vec", value_fn=chunk_embedding_fn)).add_custom_column_spec( ColumnSpec.text( column_name="content_col", value_fn=lambda chunk: f"{len(chunk.content.text)}:{chunk.content.text}")). with_metadata_spec().build()) connection_config = ConnectionConfig( jdbc_url=self.jdbc_url, username=self.username, password=self.password) writer_config = PostgresVectorWriterConfig( connection_config=connection_config, table_name=self.custom_table_name, column_specs=specs) # Generate test chunks test_chunks = [ Chunk( id=str(i), content=Content(text=f"content_{i}"), embedding=Embedding(dense_embedding=[float(i), float(i + 1)]), metadata={"timestamp": f"2024-02-02T{i:02d}:00:00"}) for i in range(num_records) ] # Write pipeline self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = ( p | beam.Create(test_chunks) | writer_config.create_write_transform()) # Read and verify read_query = f""" SELECT CAST(custom_id AS VARCHAR(255)), CAST(embedding_vec AS text), CAST(content_col AS VARCHAR(255)), CAST(metadata AS text) FROM {self.custom_table_name} ORDER BY custom_id """ # Convert BeamRow back to Chunk def custom_row_to_chunk(row): # Extract timestamp from custom_id timestamp = row.custom_id.split('timestamp_')[1] # Extract index from timestamp i = int(timestamp.split('T')[1][:2]) # Parse embedding vector embedding_list = [ float(x) for x in row.embedding_vec.strip('[]').split(',') ] # Extract content from length-prefixed format content = row.content_col.split(':', 1)[1] return Chunk( id=str(i), content=Content(text=content), embedding=Embedding(dense_embedding=embedding_list), metadata=json.loads(row.metadata)) self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.custom_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=self.jdbc_url, username=self.username, password=self.password, query=read_query)) # Verify count count_result = rows | "Count All" >> beam.combiners.Count.Globally() assert_that(count_result, equal_to([num_records]), label='count_check') chunks = rows | "To Chunks" >> beam.Map(custom_row_to_chunk) assert_that(chunks, equal_to(test_chunks), label='chunks_check') def test_defaults_with_args_specs(self): """Test custom specifications for ID, embedding, and content.""" self.skip_if_dataflow_runner() num_records = 20 specs = ( ColumnSpecsBuilder().with_id_spec( column_name="custom_id", python_type=int, convert_fn=lambda x: int(x), sql_typecast="::text").with_content_spec( column_name="content_col", convert_fn=lambda x: f"{len(x)}:{x}", ).with_embedding_spec( column_name="embedding_vec").with_metadata_spec().build()) connection_config = ConnectionConfig( jdbc_url=self.jdbc_url, username=self.username, password=self.password) writer_config = PostgresVectorWriterConfig( connection_config=connection_config, table_name=self.custom_table_name, column_specs=specs) # Generate test chunks test_chunks = [ Chunk( id=str(i), content=Content(text=f"content_{i}"), embedding=Embedding(dense_embedding=[float(i), float(i + 1)]), metadata={"timestamp": f"2024-02-02T{i:02d}:00:00"}) for i in range(num_records) ] # Write pipeline self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = ( p | beam.Create(test_chunks) | writer_config.create_write_transform()) # Read and verify read_query = f""" SELECT CAST(custom_id AS VARCHAR(255)), CAST(embedding_vec AS text), CAST(content_col AS VARCHAR(255)), CAST(metadata AS text) FROM {self.custom_table_name} ORDER BY custom_id """ # Convert BeamRow back to Chunk def custom_row_to_chunk(row): # Parse embedding vector embedding_list = [ float(x) for x in row.embedding_vec.strip('[]').split(',') ] # Extract content from length-prefixed format content = row.content_col.split(':', 1)[1] return Chunk( id=row.custom_id, content=Content(text=content), embedding=Embedding(dense_embedding=embedding_list), metadata=json.loads(row.metadata)) self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.custom_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=self.jdbc_url, username=self.username, password=self.password, query=read_query)) # Verify count count_result = rows | "Count All" >> beam.combiners.Count.Globally() assert_that(count_result, equal_to([num_records]), label='count_check') chunks = rows | "To Chunks" >> beam.Map(custom_row_to_chunk) assert_that(chunks, equal_to(test_chunks), label='chunks_check') def test_default_id_embedding_specs(self): """Test with only default id and embedding specs, others set to None.""" self.skip_if_dataflow_runner() num_records = 20 connection_config = ConnectionConfig( jdbc_url=self.jdbc_url, username=self.username, password=self.password) specs = ( ColumnSpecsBuilder().with_id_spec() # Use default id spec .with_embedding_spec() # Use default embedding spec .build()) writer_config = PostgresVectorWriterConfig( connection_config=connection_config, table_name=self.default_table_name, column_specs=specs) # Generate test chunks test_chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records) # Write pipeline self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = ( p | beam.Create(test_chunks) | writer_config.create_write_transform()) # Read and verify only id and embedding read_query = f""" SELECT CAST(id AS VARCHAR(255)), CAST(embedding AS text) FROM {self.default_table_name} ORDER BY id """ self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=self.jdbc_url, username=self.username, password=self.password, query=read_query)) chunks = rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk) # Create expected chunks with None values expected_chunks = test_utils.ChunkTestUtils.get_expected_values( 0, num_records) for chunk in expected_chunks: chunk.content.text = None chunk.metadata = {} assert_that(chunks, equal_to(expected_chunks), label='chunks_check') def test_metadata_spec_and_conflicts(self): """Test metadata specification and conflict resolution.""" self.skip_if_dataflow_runner() num_records = 20 specs = ( ColumnSpecsBuilder().with_id_spec().with_embedding_spec(). with_content_spec().add_metadata_field( field="source", column_name="source", python_type=str, sql_typecast=None # Plain text field ).add_metadata_field( field="timestamp", python_type=str, sql_typecast="::timestamp").build()) # Conflict resolution on source+timestamp conflict_resolution = ConflictResolution( on_conflict_fields=["source", "timestamp"], action="UPDATE", update_fields=["embedding", "content"]) connection_config = ConnectionConfig( jdbc_url=self.jdbc_url, username=self.username, password=self.password) writer_config = PostgresVectorWriterConfig( connection_config=connection_config, table_name=self.metadata_conflicts_table, column_specs=specs, conflict_resolution=conflict_resolution) # Generate initial test chunks initial_chunks = [ Chunk( id=str(i), content=Content(text=f"content_{i}"), embedding=Embedding(dense_embedding=[float(i), float(i + 1)]), metadata={ "source": "source_A", "timestamp": f"2024-02-02T{i:02d}:00:00" }) for i in range(num_records) ] # Write initial chunks self.write_test_pipeline.not_use_test_runner_api = True with self.write_test_pipeline as p: _ = ( p | "Write Initial" >> beam.Create(initial_chunks) | writer_config.create_write_transform()) # Generate conflicting chunks (same source+timestamp, different content) conflicting_chunks = [ Chunk( id=f"new_{i}", content=Content(text=f"updated_content_{i}"), embedding=Embedding( dense_embedding=[float(i) * 2, float(i + 1) * 2]), metadata={ "source": "source_A", "timestamp": f"2024-02-02T{i:02d}:00:00" }) for i in range(num_records) ] # Write conflicting chunks self.write_test_pipeline2.not_use_test_runner_api = True with self.write_test_pipeline2 as p: _ = ( p | "Write Conflicts" >> beam.Create(conflicting_chunks) | writer_config.create_write_transform()) # Read and verify read_query = f""" SELECT CAST(id AS VARCHAR(255)), CAST(embedding AS text), CAST(content AS VARCHAR(255)), CAST(source AS VARCHAR(255)), CAST(timestamp AS VARCHAR(255)) FROM {self.metadata_conflicts_table} ORDER BY timestamp, id """ # Expected chunks after conflict resolution expected_chunks = [ Chunk( id=str(i), content=Content(text=f"updated_content_{i}"), embedding=Embedding( dense_embedding=[float(i) * 2, float(i + 1) * 2]), metadata={ "source": "source_A", "timestamp": f"2024-02-02T{i:02d}:00:00" }) for i in range(num_records) ] def metadata_row_to_chunk(row): return Chunk( id=row.id, content=Content(text=row.content), embedding=Embedding( dense_embedding=[ float(x) for x in row.embedding.strip('[]').split(',') ]), metadata={ "source": row.source, "timestamp": row.timestamp.replace(' ', 'T') }) self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.metadata_conflicts_table, driver_class_name="org.postgresql.Driver", jdbc_url=self.jdbc_url, username=self.username, password=self.password, query=read_query)) chunks = rows | "To Chunks" >> beam.Map(metadata_row_to_chunk) assert_that(chunks, equal_to(expected_chunks), label='chunks_check') def test_conflict_resolution_update(self): """Test conflict resolution with UPDATE action.""" self.skip_if_dataflow_runner() num_records = 20 connection_config = ConnectionConfig( jdbc_url=self.jdbc_url, username=self.username, password=self.password) conflict_resolution = ConflictResolution( on_conflict_fields="id", action="UPDATE", update_fields=["embedding", "content"]) config = PostgresVectorWriterConfig( connection_config=connection_config, table_name=self.default_table_name, conflict_resolution=conflict_resolution) # Generate initial test chunks test_chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records) self.write_test_pipeline.not_use_test_runner_api = True # Insert initial test chunks with self.write_test_pipeline as p: _ = ( p | "Create initial chunks" >> beam.Create(test_chunks) | "Write initial chunks" >> config.create_write_transform()) read_query = f""" SELECT CAST(id AS VARCHAR(255)), CAST(content AS VARCHAR(255)), CAST(embedding AS text), CAST(metadata AS text) FROM {self.default_table_name} ORDER BY id desc """ self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=self.jdbc_url, username=self.username, password=self.password, query=read_query)) chunks = ( rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk) | "Key on Index" >> beam.Map(test_utils.key_on_id) | "Get First 500" >> beam.transforms.combiners.Top.Of( num_records, key=lambda x: x[0], reverse=True) | "Remove Keys 1" >> beam.Map(lambda xs: [x[1] for x in xs])) assert_that( chunks, equal_to([test_chunks]), label='original_chunks_check') updated_chunks = test_utils.ChunkTestUtils.get_expected_values( 0, num_records, content_prefix="Newcontent", seed_multiplier=2) self.write_test_pipeline2.not_use_test_runner_api = True with self.write_test_pipeline2 as p: _ = ( p | "Create updated Chunks" >> beam.Create(updated_chunks) | "Write updated Chunks" >> config.create_write_transform()) self.read_test_pipeline2.not_use_test_runner_api = True with self.read_test_pipeline2 as p: rows = ( p | "Read Updated chunks" >> ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=self.jdbc_url, username=self.username, password=self.password, query=read_query)) chunks = ( rows | "To Chunks 2" >> beam.Map(test_utils.row_to_chunk) | "Key on Index 2" >> beam.Map(test_utils.key_on_id) | "Get First 500 2" >> beam.transforms.combiners.Top.Of( num_records, key=lambda x: x[0], reverse=True) | "Remove Keys 2" >> beam.Map(lambda xs: [x[1] for x in xs])) assert_that( chunks, equal_to([updated_chunks]), label='updated_chunks_check') def test_conflict_resolution_default_ignore(self): """Test conflict resolution with default.""" self.skip_if_dataflow_runner() num_records = 20 connection_config = ConnectionConfig( jdbc_url=self.jdbc_url, username=self.username, password=self.password) config = PostgresVectorWriterConfig( connection_config=connection_config, table_name=self.default_table_name) # Generate initial test chunks test_chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records) self.write_test_pipeline.not_use_test_runner_api = True # Insert initial test chunks with self.write_test_pipeline as p: _ = ( p | "Create initial chunks" >> beam.Create(test_chunks) | "Write initial chunks" >> config.create_write_transform()) read_query = f""" SELECT CAST(id AS VARCHAR(255)), CAST(content AS VARCHAR(255)), CAST(embedding AS text), CAST(metadata AS text) FROM {self.default_table_name} ORDER BY id desc """ self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=self.jdbc_url, username=self.username, password=self.password, query=read_query)) chunks = ( rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk) | "Key on Index" >> beam.Map(test_utils.key_on_id) | "Get First 500" >> beam.transforms.combiners.Top.Of( num_records, key=lambda x: x[0], reverse=True) | "Remove Keys 1" >> beam.Map(lambda xs: [x[1] for x in xs])) assert_that( chunks, equal_to([test_chunks]), label='original_chunks_check') updated_chunks = test_utils.ChunkTestUtils.get_expected_values( 0, num_records, content_prefix="Newcontent", seed_multiplier=2) self.write_test_pipeline2.not_use_test_runner_api = True with self.write_test_pipeline2 as p: _ = ( p | "Create updated Chunks" >> beam.Create(updated_chunks) | "Write updated Chunks" >> config.create_write_transform()) self.read_test_pipeline2.not_use_test_runner_api = True with self.read_test_pipeline2 as p: rows = ( p | "Read Updated chunks" >> ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=self.jdbc_url, username=self.username, password=self.password, query=read_query)) chunks = ( rows | "To Chunks 2" >> beam.Map(test_utils.row_to_chunk) | "Key on Index 2" >> beam.Map(test_utils.key_on_id) | "Get First 500 2" >> beam.transforms.combiners.Top.Of( num_records, key=lambda x: x[0], reverse=True) | "Remove Keys 2" >> beam.Map(lambda xs: [x[1] for x in xs])) assert_that(chunks, equal_to([test_chunks]), label='updated_chunks_check') def test_conflict_resolution_default_update_fields(self): """Test conflict resolution with default update fields (all non-conflict fields).""" self.skip_if_dataflow_runner() num_records = 20 connection_config = ConnectionConfig( jdbc_url=self.jdbc_url, username=self.username, password=self.password) # Create a conflict resolution with only the conflict field specified # No update_fields specified - should default to all non-conflict fields conflict_resolution = ConflictResolution( on_conflict_fields="id", action="UPDATE") config = PostgresVectorWriterConfig( connection_config=connection_config, table_name=self.default_table_name, conflict_resolution=conflict_resolution) # Generate initial test chunks test_chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records) self.write_test_pipeline.not_use_test_runner_api = True # Insert initial test chunks with self.write_test_pipeline as p: _ = ( p | "Create initial chunks" >> beam.Create(test_chunks) | "Write initial chunks" >> config.create_write_transform()) # Verify initial data was written correctly read_query = f""" SELECT CAST(id AS VARCHAR(255)), CAST(content AS VARCHAR(255)), CAST(embedding AS text), CAST(metadata AS text) FROM {self.default_table_name} ORDER BY id desc """ self.read_test_pipeline.not_use_test_runner_api = True with self.read_test_pipeline as p: rows = ( p | ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=self.jdbc_url, username=self.username, password=self.password, query=read_query)) chunks = ( rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk) | "Key on Index" >> beam.Map(test_utils.key_on_id) | "Get First 500" >> beam.transforms.combiners.Top.Of( num_records, key=lambda x: x[0], reverse=True) | "Remove Keys 1" >> beam.Map(lambda xs: [x[1] for x in xs])) assert_that( chunks, equal_to([test_chunks]), label='original_chunks_check') # Create updated chunks with same IDs but different content, embedding, and # metadata updated_chunks = [] for i in range(num_records): original_chunk = test_chunks[i] updated_chunk = Chunk( id=original_chunk.id, content=Content(text=f"Updated content {i}"), embedding=Embedding( dense_embedding=[float(i * 2), float(i * 2 + 1)] + [0.0] * (test_utils.VECTOR_SIZE - 2)), metadata={ "updated": "true", "timestamp": "2024-02-25" }) updated_chunks.append(updated_chunk) # Write updated chunks - should update all non-conflict fields self.write_test_pipeline2.not_use_test_runner_api = True with self.write_test_pipeline2 as p: _ = ( p | "Create updated Chunks" >> beam.Create(updated_chunks) | "Write updated Chunks" >> config.create_write_transform()) # Read and verify that all non-conflict fields were updated self.read_test_pipeline2.not_use_test_runner_api = True with self.read_test_pipeline2 as p: rows = ( p | "Read Updated chunks" >> ReadFromJdbc( table_name=self.default_table_name, driver_class_name="org.postgresql.Driver", jdbc_url=self.jdbc_url, username=self.username, password=self.password, query=read_query)) chunks = ( rows | "To Chunks 2" >> beam.Map(test_utils.row_to_chunk) | "Key on Index 2" >> beam.Map(test_utils.key_on_id) | "Get First 500 2" >> beam.transforms.combiners.Top.Of( num_records, key=lambda x: x[0], reverse=True) | "Remove Keys 2" >> beam.Map(lambda xs: [x[1] for x in xs])) # Verify that all non-conflict fields were updated assert_that( chunks, equal_to([updated_chunks]), label='updated_chunks_check') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/spanner.py0000644000175100017510000005222215206113032023060 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Cloud Spanner vector store writer for RAG pipelines. This module provides a writer for storing embeddings and associated metadata in Google Cloud Spanner. It supports flexible schema configuration with the ability to flatten metadata fields into dedicated columns. Example usage: Default schema (id, embedding, content, metadata): >>> config = SpannerVectorWriterConfig( ... project_id="my-project", ... instance_id="my-instance", ... database_id="my-db", ... table_name="embeddings" ... ) Flattened metadata fields: >>> specs = ( ... SpannerColumnSpecsBuilder() ... .with_id_spec() ... .with_embedding_spec() ... .with_content_spec() ... .add_metadata_field("source", str) ... .add_metadata_field("page_number", int, default=0) ... .with_metadata_spec() ... .build() ... ) >>> config = SpannerVectorWriterConfig( ... project_id="my-project", ... instance_id="my-instance", ... database_id="my-db", ... table_name="embeddings", ... column_specs=specs ... ) Spanner schema example: CREATE TABLE embeddings ( id STRING(1024) NOT NULL, embedding ARRAY(vector_length=>768), content STRING(MAX), source STRING(MAX), page_number INT64, metadata JSON ) PRIMARY KEY (id) """ import functools import json from dataclasses import dataclass from typing import Any from typing import Callable from typing import Literal from typing import NamedTuple from typing import Optional import apache_beam as beam from apache_beam.coders import registry from apache_beam.coders.row_coder import RowCoder from apache_beam.io.gcp import spanner from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteConfig from apache_beam.ml.rag.types import EmbeddableItem @dataclass class SpannerColumnSpec: """Column specification for Spanner vector writes. Defines how to extract and format values from EmbeddableItems for insertion into Spanner table columns. Each spec maps to one column in the target table. Attributes: column_name: Name of the Spanner table column python_type: Python type for the NamedTuple field (required for RowCoder) value_fn: Function to extract value from an EmbeddableItem Examples: String column: >>> SpannerColumnSpec( ... column_name="id", ... python_type=str, ... value_fn=lambda embeddable: embeddable.id ... ) Array column with conversion: >>> SpannerColumnSpec( ... column_name="embedding", ... python_type=List[float], ... value_fn=lambda embeddable: embeddable.embedding.dense_embedding ... ) """ column_name: str python_type: type value_fn: Callable[[EmbeddableItem], Any] def _extract_and_convert(extract_fn, convert_fn, embeddable): if convert_fn: return convert_fn(extract_fn(embeddable)) return extract_fn(embeddable) class SpannerColumnSpecsBuilder: """Builder for creating Spanner column specifications. Provides a fluent API for defining table schemas and how to populate them from EmbeddableItem objects. Supports standard EmbeddableItem fields (id, embedding, content, metadata) and flattening metadata fields into dedicated columns. Example: >>> specs = ( ... SpannerColumnSpecsBuilder() ... .with_id_spec() ... .with_embedding_spec() ... .with_content_spec() ... .add_metadata_field("source", str) ... .with_metadata_spec() ... .build() ... ) """ def __init__(self): self._specs: list[SpannerColumnSpec] = [] @staticmethod def with_defaults() -> 'SpannerColumnSpecsBuilder': """Create builder with default schema. Default schema includes: - id (STRING): EmbeddableItem ID - embedding (ARRAY): Dense embedding vector - content (STRING): EmbeddableItem content text - metadata (JSON): Full metadata as JSON Returns: Builder with default column specifications """ return ( SpannerColumnSpecsBuilder().with_id_spec().with_embedding_spec(). with_content_spec().with_metadata_spec()) def with_id_spec( self, column_name: str = "id", python_type: type = str, convert_fn: Optional[Callable[[str], Any]] = None ) -> 'SpannerColumnSpecsBuilder': """Add ID column specification. Args: column_name: Column name (default: "id") python_type: Python type (default: str) convert_fn: Optional converter (e.g., to cast to int) Returns: Self for method chaining Examples: Default string ID: >>> builder.with_id_spec() Integer ID with conversion: >>> builder.with_id_spec( ... python_type=int, ... convert_fn=lambda id: int(id.split('_')[1]) ... ) """ self._specs.append( SpannerColumnSpec( column_name=column_name, python_type=python_type, value_fn=functools.partial( _extract_and_convert, lambda embeddable: embeddable.id, convert_fn))) return self def with_embedding_spec( self, column_name: str = "embedding", convert_fn: Optional[Callable[[list[float]], list[float]]] = None ) -> 'SpannerColumnSpecsBuilder': """Add embedding array column (ARRAY or ARRAY). Args: column_name: Column name (default: "embedding") convert_fn: Optional converter (e.g., normalize, quantize) Returns: Self for method chaining Examples: Default embedding: >>> builder.with_embedding_spec() Normalized embedding: >>> def normalize(vec): ... norm = (sum(x**2 for x in vec) ** 0.5) or 1.0 ... return [x/norm for x in vec] >>> builder.with_embedding_spec(convert_fn=normalize) Rounded precision: >>> builder.with_embedding_spec( ... convert_fn=lambda vec: [round(x, 4) for x in vec] ... ) """ def extract_fn(embeddable: EmbeddableItem) -> list[float]: if not embeddable.dense_embedding: raise ValueError(f'EmbeddableItem must contain embedding: {embeddable}') return embeddable.dense_embedding self._specs.append( SpannerColumnSpec( column_name=column_name, python_type=list[float], value_fn=functools.partial( _extract_and_convert, extract_fn, convert_fn))) return self def with_content_spec( self, column_name: str = "content", python_type: type = str, convert_fn: Optional[Callable[[str], Any]] = None ) -> 'SpannerColumnSpecsBuilder': """Add content column. Args: column_name: Column name (default: "content") python_type: Python type (default: str) convert_fn: Optional converter Returns: Self for method chaining Examples: Default text content: >>> builder.with_content_spec() Content length as integer: >>> builder.with_content_spec( ... column_name="content_length", ... python_type=int, ... convert_fn=lambda text: len(text.split()) ... ) Truncated content: >>> builder.with_content_spec( ... convert_fn=lambda text: text[:1000] ... ) """ def extract_fn(embeddable: EmbeddableItem) -> str: return embeddable.content_string self._specs.append( SpannerColumnSpec( column_name=column_name, python_type=python_type, value_fn=functools.partial( _extract_and_convert, extract_fn, convert_fn))) return self def with_metadata_spec( self, column_name: str = "metadata") -> 'SpannerColumnSpecsBuilder': """Add metadata JSON column. Stores the full metadata dictionary as a JSON string in Spanner. Args: column_name: Column name (default: "metadata") Returns: Self for method chaining Note: Metadata is automatically converted to JSON string using json.dumps() """ value_fn = lambda embeddable: json.dumps(embeddable.metadata) self._specs.append( SpannerColumnSpec( column_name=column_name, python_type=str, value_fn=value_fn)) return self def add_metadata_field( self, field: str, python_type: type, column_name: Optional[str] = None, convert_fn: Optional[Callable[[Any], Any]] = None, default: Any = None) -> 'SpannerColumnSpecsBuilder': """Flatten a metadata field into its own column. Extracts a specific field from embeddable.metadata and stores it in a dedicated table column. Args: field: Key in embeddable.metadata to extract python_type: Python type (must be explicitly specified) column_name: Column name (default: same as field) convert_fn: Optional converter for type casting/transformation default: Default value if field is missing from metadata Returns: Self for method chaining Examples: String field: >>> builder.add_metadata_field("source", str) Integer with default: >>> builder.add_metadata_field( ... "page_number", ... int, ... default=0 ... ) Float with conversion: >>> builder.add_metadata_field( ... "confidence", ... float, ... convert_fn=lambda x: round(float(x), 2), ... default=0.0 ... ) List of strings: >>> builder.add_metadata_field( ... "tags", ... List[str], ... default=[] ... ) Timestamp with conversion: >>> builder.add_metadata_field( ... "created_at", ... str, ... convert_fn=lambda ts: ts.isoformat() ... ) """ name = column_name or field def value_fn(embeddable: EmbeddableItem) -> Any: return embeddable.metadata.get(field, default) self._specs.append( SpannerColumnSpec( column_name=name, python_type=python_type, value_fn=functools.partial( _extract_and_convert, value_fn, convert_fn))) return self def add_column( self, column_name: str, python_type: type, value_fn: Callable[[EmbeddableItem], Any]) -> 'SpannerColumnSpecsBuilder': """Add a custom column with full control. Args: column_name: Column name python_type: Python type (required) value_fn: Value extraction function Returns: Self for method chaining Examples: Boolean flag: >>> builder.add_column( ... column_name="has_code", ... python_type=bool, ... value_fn=lambda embeddable: "```" in embeddable.content.text ... ) Computed value: >>> builder.add_column( ... column_name="word_count", ... python_type=int, ... value_fn=lambda embeddable: len(embeddable.content.text.split()) ... ) """ self._specs.append( SpannerColumnSpec( column_name=column_name, python_type=python_type, value_fn=value_fn)) return self def build(self) -> list[SpannerColumnSpec]: """Build the final list of column specifications. Returns: List of SpannerColumnSpec objects """ return self._specs.copy() class _SpannerSchemaBuilder: """Internal: Builds NamedTuple schema and registers RowCoder. Creates a NamedTuple type from column specifications and registers it with Beam's RowCoder for serialization. """ def __init__(self, table_name: str, column_specs: list[SpannerColumnSpec]): """Initialize schema builder. Args: table_name: Table name (used in NamedTuple type name) column_specs: List of column specifications Raises: ValueError: If duplicate column names are found """ self.table_name = table_name self.column_specs = column_specs # Validate no duplicates names = [col.column_name for col in column_specs] duplicates = set(name for name in names if names.count(name) > 1) if duplicates: raise ValueError(f"Duplicate column names: {duplicates}") # Create NamedTuple type fields = [(col.column_name, col.python_type) for col in column_specs] type_name = f"SpannerVectorRecord_{table_name}" self.record_type = NamedTuple(type_name, fields) # type: ignore # Register coder registry.register_coder(self.record_type, RowCoder) def create_converter(self) -> Callable[[EmbeddableItem], NamedTuple]: """Create converter function from EmbeddableItem to NamedTuple record. Returns: Function that converts an EmbeddableItem to a NamedTuple record """ def convert(embeddable: EmbeddableItem) -> self.record_type: # type: ignore values = { col.column_name: col.value_fn(embeddable) for col in self.column_specs } return self.record_type(**values) # type: ignore return convert class SpannerVectorWriterConfig(VectorDatabaseWriteConfig): """Configuration for writing vectors to Cloud Spanner. Supports flexible schema configuration through column specifications and provides control over Spanner-specific write parameters. Examples: Default schema: >>> config = SpannerVectorWriterConfig( ... project_id="my-project", ... instance_id="my-instance", ... database_id="my-db", ... table_name="embeddings" ... ) Custom schema with flattened metadata: >>> specs = ( ... SpannerColumnSpecsBuilder() ... .with_id_spec() ... .with_embedding_spec() ... .with_content_spec() ... .add_metadata_field("source", str) ... .add_metadata_field("page_number", int, default=0) ... .with_metadata_spec() ... .build() ... ) >>> config = SpannerVectorWriterConfig( ... project_id="my-project", ... instance_id="my-instance", ... database_id="my-db", ... table_name="embeddings", ... column_specs=specs ... ) With emulator: >>> config = SpannerVectorWriterConfig( ... project_id="test-project", ... instance_id="test-instance", ... database_id="test-db", ... table_name="embeddings", ... emulator_host="http://localhost:9010" ... ) """ def __init__( self, project_id: str, instance_id: str, database_id: str, table_name: str, *, # Schema configuration column_specs: Optional[list[SpannerColumnSpec]] = None, # Write operation type write_mode: Literal["INSERT", "UPDATE", "REPLACE", "INSERT_OR_UPDATE"] = "INSERT_OR_UPDATE", # Batching configuration max_batch_size_bytes: Optional[int] = None, max_number_mutations: Optional[int] = None, max_number_rows: Optional[int] = None, grouping_factor: Optional[int] = None, # Networking host: Optional[str] = None, emulator_host: Optional[str] = None, expansion_service: Optional[str] = None, # Retry/deadline configuration commit_deadline: Optional[int] = None, max_cumulative_backoff: Optional[int] = None, # Error handling failure_mode: Optional[ spanner.FailureMode] = spanner.FailureMode.REPORT_FAILURES, high_priority: bool = False, # Additional Spanner arguments **spanner_kwargs): """Initialize Spanner vector writer configuration. Args: project_id: GCP project ID instance_id: Spanner instance ID database_id: Spanner database ID table_name: Target table name column_specs: Schema configuration using SpannerColumnSpecsBuilder. If None, uses default schema (id, embedding, content, metadata) write_mode: Spanner write operation type: - INSERT: Fail if row exists - UPDATE: Fail if row doesn't exist - REPLACE: Delete then insert - INSERT_OR_UPDATE: Insert or update if exists (default) max_batch_size_bytes: Maximum bytes per mutation batch (default: 1MB) max_number_mutations: Maximum cell mutations per batch (default: 5000) max_number_rows: Maximum rows per batch (default: 500) grouping_factor: Multiple of max mutation for sorting (default: 1000) host: Spanner host URL (usually not needed) emulator_host: Spanner emulator host (e.g., "http://localhost:9010") expansion_service: Java expansion service address (host:port) commit_deadline: Commit API deadline in seconds (default: 15) max_cumulative_backoff: Max retry backoff seconds (default: 900) failure_mode: Error handling strategy: - FAIL_FAST: Throw exception for any failure - REPORT_FAILURES: Continue processing (default) high_priority: Use high priority for operations (default: False) **spanner_kwargs: Additional keyword arguments to pass to the underlying Spanner write transform. Use this to pass any Spanner-specific parameters not explicitly exposed by this config. """ self.project_id = project_id self.instance_id = instance_id self.database_id = database_id self.table_name = table_name self.write_mode = write_mode self.max_batch_size_bytes = max_batch_size_bytes self.max_number_mutations = max_number_mutations self.max_number_rows = max_number_rows self.grouping_factor = grouping_factor self.host = host self.emulator_host = emulator_host self.expansion_service = expansion_service self.commit_deadline = commit_deadline self.max_cumulative_backoff = max_cumulative_backoff self.failure_mode = failure_mode self.high_priority = high_priority self.spanner_kwargs = spanner_kwargs # Use defaults if not provided specs = column_specs or SpannerColumnSpecsBuilder.with_defaults().build() # Create schema builder (NamedTuple + RowCoder registration) self.schema_builder = _SpannerSchemaBuilder(table_name, specs) def create_write_transform(self) -> beam.PTransform: """Create the Spanner write PTransform. Returns: PTransform for writing to Spanner """ return _WriteToSpannerVectorDatabase(self) class _WriteToSpannerVectorDatabase(beam.PTransform): """Internal: PTransform for writing to Spanner vector database.""" def __init__(self, config: SpannerVectorWriterConfig): """Initialize write transform. Args: config: Spanner writer configuration """ self.config = config self.schema_builder = config.schema_builder def expand(self, pcoll: beam.PCollection[EmbeddableItem]): """Expand the transform. Args: pcoll: PCollection of EmbeddableItems to write """ # Select appropriate Spanner write transform based on write_mode write_transform_class = { "INSERT": spanner.SpannerInsert, "UPDATE": spanner.SpannerUpdate, "REPLACE": spanner.SpannerReplace, "INSERT_OR_UPDATE": spanner.SpannerInsertOrUpdate, }[self.config.write_mode] return ( pcoll | "Convert to Records" >> beam.Map( self.schema_builder.create_converter()).with_output_types( self.schema_builder.record_type) | "Write to Spanner" >> write_transform_class( project_id=self.config.project_id, instance_id=self.config.instance_id, database_id=self.config.database_id, table=self.config.table_name, max_batch_size_bytes=self.config.max_batch_size_bytes, max_number_mutations=self.config.max_number_mutations, max_number_rows=self.config.max_number_rows, grouping_factor=self.config.grouping_factor, host=self.config.host, emulator_host=self.config.emulator_host, commit_deadline=self.config.commit_deadline, max_cumulative_backoff=self.config.max_cumulative_backoff, failure_mode=self.config.failure_mode, expansion_service=self.config.expansion_service, high_priority=self.config.high_priority, **self.config.spanner_kwargs)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/spanner_it_test.py0000644000175100017510000004766715206113032024634 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration tests for Spanner vector store writer.""" import logging import os import time import unittest import uuid import pytest import apache_beam as beam from apache_beam.ml.rag.ingestion.spanner import SpannerVectorWriterConfig from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import Embedding from apache_beam.testing.test_pipeline import TestPipeline # pylint: disable=wrong-import-order, wrong-import-position try: from google.cloud import spanner except ImportError: spanner = None # type: ignore[assignment] try: from testcontainers.core.container import DockerContainer except ImportError: DockerContainer = None # pylint: enable=wrong-import-order, wrong-import-position def retry(fn, retries, err_msg, *args, **kwargs): """Retry a function with exponential backoff.""" for _ in range(retries): try: return fn(*args, **kwargs) except: # pylint: disable=bare-except time.sleep(1) logging.error(err_msg) raise RuntimeError(err_msg) class SpannerEmulatorHelper: """Helper for managing Spanner emulator lifecycle.""" def __init__(self, project_id: str, instance_id: str, table_name: str): self.project_id = project_id self.instance_id = instance_id self.table_name = table_name self.host = None # Start emulator self.emulator = DockerContainer( 'gcr.io/cloud-spanner-emulator/emulator:latest').with_exposed_ports( 9010, 9020) retry(self.emulator.start, 3, 'Could not start spanner emulator.') time.sleep(3) self.host = f'{self.emulator.get_container_host_ip()}:' \ f'{self.emulator.get_exposed_port(9010)}' os.environ['SPANNER_EMULATOR_HOST'] = self.host # Create client and instance self.client = spanner.Client(project_id) self.instance = self.client.instance(instance_id) self.create_instance() def create_instance(self): """Create Spanner instance in emulator.""" self.instance.create().result(120) def create_database(self, database_id: str): """Create database with default vector table schema.""" database = self.instance.database( database_id, ddl_statements=[ f''' CREATE TABLE {self.table_name} ( id STRING(1024) NOT NULL, embedding ARRAY(vector_length=>3), content STRING(MAX), metadata JSON ) PRIMARY KEY (id)''' ]) database.create().result(120) def read_data(self, database_id: str): """Read all data from the table.""" database = self.instance.database(database_id) with database.snapshot() as snapshot: results = snapshot.execute_sql( f'SELECT * FROM {self.table_name} ORDER BY id') return list(results) if results else [] def drop_database(self, database_id: str): """Drop the database.""" database = self.instance.database(database_id) database.drop() def shutdown(self): """Stop the emulator.""" if self.emulator: try: self.emulator.stop() except: # pylint: disable=bare-except logging.error('Could not stop Spanner emulator.') def get_emulator_host(self) -> str: """Get the emulator host URL.""" return f'http://{self.host}' @pytest.mark.uses_gcp_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_JARS'), "EXPANSION_JARS environment var is not provided, " "indicating that jars have not been built") @unittest.skipIf(spanner is None, 'GCP dependencies are not installed.') @unittest.skipIf( DockerContainer is None, 'testcontainers package is not installed.') class SpannerVectorWriterTest(unittest.TestCase): """Integration tests for Spanner vector writer.""" @classmethod def setUpClass(cls): """Set up Spanner emulator for all tests.""" pipeline = TestPipeline(is_integration_test=True) runner_name = type(pipeline.runner).__name__ if 'DataflowRunner' in runner_name: pytest.skip("Spanner emulator not compatible with dataflow runner.") cls.project_id = 'test-project' cls.instance_id = 'test-instance' cls.table_name = 'embeddings' cls.spanner_helper = SpannerEmulatorHelper( cls.project_id, cls.instance_id, cls.table_name) @classmethod def tearDownClass(cls): """Tear down Spanner emulator.""" cls.spanner_helper.shutdown() def setUp(self): """Create a unique database for each test.""" self.database_id = f'test_db_{uuid.uuid4().hex}'[:30] self.spanner_helper.create_database(self.database_id) def tearDown(self): """Drop the test database.""" self.spanner_helper.drop_database(self.database_id) def test_write_default_schema(self): """Test writing with default schema (id, embedding, content, metadata).""" # Create test chunks chunks = [ Chunk( id='doc1', embedding=Embedding(dense_embedding=[1.0, 2.0, 3.0]), content=Content(text='First document'), metadata={ 'source': 'test', 'page': 1 }), Chunk( id='doc2', embedding=Embedding(dense_embedding=[4.0, 5.0, 6.0]), content=Content(text='Second document'), metadata={ 'source': 'test', 'page': 2 }), ] # Create config with default schema config = SpannerVectorWriterConfig( project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id, table_name=self.table_name, emulator_host=self.spanner_helper.get_emulator_host(), ) # Write chunks with TestPipeline() as p: p.not_use_test_runner_api = True _ = (p | beam.Create(chunks) | config.create_write_transform()) # Verify data was written results = self.spanner_helper.read_data(self.database_id) self.assertEqual(len(results), 2) # Check first row row1 = results[0] self.assertEqual(row1[0], 'doc1') # id self.assertEqual(list(row1[1]), [1.0, 2.0, 3.0]) # embedding self.assertEqual(row1[2], 'First document') # content # metadata is JSON metadata1 = row1[3] self.assertEqual(metadata1['source'], 'test') self.assertEqual(metadata1['page'], 1) # Check second row row2 = results[1] self.assertEqual(row2[0], 'doc2') self.assertEqual(list(row2[1]), [4.0, 5.0, 6.0]) self.assertEqual(row2[2], 'Second document') def test_write_flattened_metadata(self): """Test writing with flattened metadata fields.""" from apache_beam.ml.rag.ingestion.spanner import SpannerColumnSpecsBuilder # Create custom database with flattened columns self.spanner_helper.drop_database(self.database_id) database = self.spanner_helper.instance.database( self.database_id, ddl_statements=[ f''' CREATE TABLE {self.table_name} ( id STRING(1024) NOT NULL, embedding ARRAY(vector_length=>3), content STRING(MAX), source STRING(MAX), page_number INT64, metadata JSON ) PRIMARY KEY (id)''' ]) database.create().result(120) # Create test chunks chunks = [ Chunk( id='doc1', embedding=Embedding(dense_embedding=[1.0, 2.0, 3.0]), content=Content(text='First document'), metadata={ 'source': 'book.pdf', 'page': 10, 'author': 'John' }), Chunk( id='doc2', embedding=Embedding(dense_embedding=[4.0, 5.0, 6.0]), content=Content(text='Second document'), metadata={ 'source': 'article.txt', 'page': 5, 'author': 'Jane' }), ] # Create config with flattened metadata specs = ( SpannerColumnSpecsBuilder().with_id_spec().with_embedding_spec(). with_content_spec().add_metadata_field( 'source', str, column_name='source').add_metadata_field( 'page', int, column_name='page_number').with_metadata_spec().build()) config = SpannerVectorWriterConfig( project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id, table_name=self.table_name, column_specs=specs, emulator_host=self.spanner_helper.get_emulator_host(), ) # Write chunks with TestPipeline() as p: p.not_use_test_runner_api = True _ = (p | beam.Create(chunks) | config.create_write_transform()) # Verify data database = self.spanner_helper.instance.database(self.database_id) with database.snapshot() as snapshot: results = snapshot.execute_sql( f'SELECT id, embedding, content, source, page_number, metadata ' f'FROM {self.table_name} ORDER BY id') rows = list(results) self.assertEqual(len(rows), 2) # Check first row self.assertEqual(rows[0][0], 'doc1') self.assertEqual(list(rows[0][1]), [1.0, 2.0, 3.0]) self.assertEqual(rows[0][2], 'First document') self.assertEqual(rows[0][3], 'book.pdf') # flattened source self.assertEqual(rows[0][4], 10) # flattened page_number metadata1 = rows[0][5] self.assertEqual(metadata1['author'], 'John') def test_write_minimal_schema(self): """Test writing with minimal schema (only id and embedding).""" from apache_beam.ml.rag.ingestion.spanner import SpannerColumnSpecsBuilder # Create custom database with minimal schema self.spanner_helper.drop_database(self.database_id) database = self.spanner_helper.instance.database( self.database_id, ddl_statements=[ f''' CREATE TABLE {self.table_name} ( id STRING(1024) NOT NULL, embedding ARRAY(vector_length=>3) ) PRIMARY KEY (id)''' ]) database.create().result(120) # Create test chunks chunks = [ Chunk( id='doc1', embedding=Embedding(dense_embedding=[1.0, 2.0, 3.0]), content=Content(text='First document'), metadata={'source': 'test'}), Chunk( id='doc2', embedding=Embedding(dense_embedding=[4.0, 5.0, 6.0]), content=Content(text='Second document'), metadata={'source': 'test'}), ] # Create config with minimal schema specs = ( SpannerColumnSpecsBuilder().with_id_spec().with_embedding_spec().build( )) config = SpannerVectorWriterConfig( project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id, table_name=self.table_name, column_specs=specs, emulator_host=self.spanner_helper.get_emulator_host(), ) # Write chunks with TestPipeline() as p: p.not_use_test_runner_api = True _ = (p | beam.Create(chunks) | config.create_write_transform()) # Verify data results = self.spanner_helper.read_data(self.database_id) self.assertEqual(len(results), 2) self.assertEqual(results[0][0], 'doc1') self.assertEqual(list(results[0][1]), [1.0, 2.0, 3.0]) def test_write_with_converter(self): """Test writing with custom converter function.""" from apache_beam.ml.rag.ingestion.spanner import SpannerColumnSpecsBuilder # Create test chunks with embeddings that need normalization chunks = [ Chunk( id='doc1', embedding=Embedding(dense_embedding=[3.0, 4.0, 0.0]), content=Content(text='First document'), metadata={'source': 'test'}), ] # Define normalizer def normalize(vec): norm = (sum(x**2 for x in vec)**0.5) or 1.0 return [x / norm for x in vec] # Create config with normalized embeddings specs = ( SpannerColumnSpecsBuilder().with_id_spec().with_embedding_spec( convert_fn=normalize).with_content_spec().with_metadata_spec(). build()) config = SpannerVectorWriterConfig( project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id, table_name=self.table_name, column_specs=specs, emulator_host=self.spanner_helper.get_emulator_host(), ) # Write chunks with TestPipeline() as p: p.not_use_test_runner_api = True _ = (p | beam.Create(chunks) | config.create_write_transform()) # Verify data - embedding should be normalized results = self.spanner_helper.read_data(self.database_id) self.assertEqual(len(results), 1) embedding = list(results[0][1]) # Original was [3.0, 4.0, 0.0], normalized should be [0.6, 0.8, 0.0] self.assertAlmostEqual(embedding[0], 0.6, places=5) self.assertAlmostEqual(embedding[1], 0.8, places=5) self.assertAlmostEqual(embedding[2], 0.0, places=5) # Check norm is 1.0 norm = sum(x**2 for x in embedding)**0.5 self.assertAlmostEqual(norm, 1.0, places=5) def test_write_update_mode(self): """Test writing with UPDATE mode.""" # First insert data chunks_insert = [ Chunk( id='doc1', embedding=Embedding(dense_embedding=[1.0, 2.0, 3.0]), content=Content(text='Original content'), metadata={'version': 1}), ] config_insert = SpannerVectorWriterConfig( project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id, table_name=self.table_name, write_mode='INSERT', emulator_host=self.spanner_helper.get_emulator_host(), ) with TestPipeline() as p: p.not_use_test_runner_api = True _ = ( p | beam.Create(chunks_insert) | config_insert.create_write_transform()) # Update existing row chunks_update = [ Chunk( id='doc1', embedding=Embedding(dense_embedding=[4.0, 5.0, 6.0]), content=Content(text='Updated content'), metadata={'version': 2}), ] config_update = SpannerVectorWriterConfig( project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id, table_name=self.table_name, write_mode='UPDATE', emulator_host=self.spanner_helper.get_emulator_host(), ) with TestPipeline() as p: p.not_use_test_runner_api = True _ = ( p | beam.Create(chunks_update) | config_update.create_write_transform()) # Verify update succeeded results = self.spanner_helper.read_data(self.database_id) self.assertEqual(len(results), 1) self.assertEqual(results[0][0], 'doc1') self.assertEqual(list(results[0][1]), [4.0, 5.0, 6.0]) self.assertEqual(results[0][2], 'Updated content') metadata = results[0][3] self.assertEqual(metadata['version'], 2) def test_write_custom_column(self): """Test writing with custom computed column.""" from apache_beam.ml.rag.ingestion.spanner import SpannerColumnSpecsBuilder # Create custom database with computed column self.spanner_helper.drop_database(self.database_id) database = self.spanner_helper.instance.database( self.database_id, ddl_statements=[ f''' CREATE TABLE {self.table_name} ( id STRING(1024) NOT NULL, embedding ARRAY(vector_length=>3), content STRING(MAX), word_count INT64, metadata JSON ) PRIMARY KEY (id)''' ]) database.create().result(120) # Create test chunks chunks = [ Chunk( id='doc1', embedding=Embedding(dense_embedding=[1.0, 2.0, 3.0]), content=Content(text='Hello world test'), metadata={}), Chunk( id='doc2', embedding=Embedding(dense_embedding=[4.0, 5.0, 6.0]), content=Content(text='This is a longer test document'), metadata={}), ] # Create config with custom word_count column specs = ( SpannerColumnSpecsBuilder().with_id_spec().with_embedding_spec( ).with_content_spec().add_column( column_name='word_count', python_type=int, value_fn=lambda chunk: len(chunk.content.text.split())). with_metadata_spec().build()) config = SpannerVectorWriterConfig( project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id, table_name=self.table_name, column_specs=specs, emulator_host=self.spanner_helper.get_emulator_host(), ) # Write chunks with TestPipeline() as p: p.not_use_test_runner_api = True _ = (p | beam.Create(chunks) | config.create_write_transform()) # Verify data database = self.spanner_helper.instance.database(self.database_id) with database.snapshot() as snapshot: results = snapshot.execute_sql( f'SELECT id, word_count FROM {self.table_name} ORDER BY id') rows = list(results) self.assertEqual(len(rows), 2) self.assertEqual(rows[0][1], 3) # "Hello world test" = 3 words self.assertEqual(rows[1][1], 6) # 6 words def test_write_with_timestamp(self): """Test writing with timestamp columns.""" from apache_beam.ml.rag.ingestion.spanner import SpannerColumnSpecsBuilder # Create database with timestamp column self.spanner_helper.drop_database(self.database_id) database = self.spanner_helper.instance.database( self.database_id, ddl_statements=[ f''' CREATE TABLE {self.table_name} ( id STRING(1024) NOT NULL, embedding ARRAY(vector_length=>3), content STRING(MAX), created_at TIMESTAMP, metadata JSON ) PRIMARY KEY (id)''' ]) database.create().result(120) # Create chunks with timestamp timestamp_str = "2025-10-28T09:45:00.123456Z" chunks = [ Chunk( id='doc1', embedding=Embedding(dense_embedding=[1.0, 2.0, 3.0]), content=Content(text='Document with timestamp'), metadata={'created_at': timestamp_str}), ] # Create config with timestamp field specs = ( SpannerColumnSpecsBuilder().with_id_spec().with_embedding_spec(). with_content_spec().add_metadata_field( 'created_at', str, column_name='created_at').with_metadata_spec().build()) config = SpannerVectorWriterConfig( project_id=self.project_id, instance_id=self.instance_id, database_id=self.database_id, table_name=self.table_name, column_specs=specs, emulator_host=self.spanner_helper.get_emulator_host(), ) # Write chunks with TestPipeline() as p: p.not_use_test_runner_api = True _ = (p | beam.Create(chunks) | config.create_write_transform()) # Verify timestamp was written database = self.spanner_helper.instance.database(self.database_id) with database.snapshot() as snapshot: results = snapshot.execute_sql( f'SELECT id, created_at FROM {self.table_name}') rows = list(results) self.assertEqual(len(rows), 1) self.assertEqual(rows[0][0], 'doc1') # Timestamp is returned as datetime object by Spanner client self.assertIsNotNone(rows[0][1]) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/ingestion/test_utils.py0000644000175100017510000000616215206113032023613 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import hashlib import json import apache_beam as beam from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import Embedding VECTOR_SIZE = 768 def row_to_chunk(row) -> Chunk: # Parse embedding string back to float list embedding_list = [float(x) for x in row.embedding.strip('[]').split(',')] return Chunk( id=row.id, content=Content(text=row.content if hasattr(row, 'content') else None), embedding=Embedding(dense_embedding=embedding_list), metadata=json.loads(row.metadata) if hasattr(row, 'metadata') else {}) class ChunkTestUtils: """Helper functions for generating test Chunks.""" @staticmethod def from_seed(seed: int, content_prefix: str, seed_multiplier: int) -> Chunk: """Creates a deterministic Chunk from a seed value.""" return Chunk( id=f"id_{seed}", content=Content(text=f"{content_prefix}{seed}"), embedding=Embedding( dense_embedding=[ float(seed + i * seed_multiplier) / 100 for i in range(VECTOR_SIZE) ]), metadata={"seed": str(seed)}) @staticmethod def get_expected_values( range_start: int, range_end: int, content_prefix: str = "Testval", seed_multiplier: int = 1) -> list[Chunk]: """Returns a range of test Chunks.""" return [ ChunkTestUtils.from_seed(i, content_prefix, seed_multiplier) for i in range(range_start, range_end) ] class HashingFn(beam.CombineFn): """Hashing function for verification.""" def create_accumulator(self): return [] def add_input(self, accumulator, input): accumulator.append(input.content.text if input.content.text else "") return accumulator def merge_accumulators(self, accumulators): merged = [] for acc in accumulators: merged.extend(acc) return merged def extract_output(self, accumulator): sorted_values = sorted(accumulator) return hashlib.md5(''.join(sorted_values).encode()).hexdigest() def generate_expected_hash(num_records: int) -> str: chunks = ChunkTestUtils.get_expected_values(0, num_records) values = sorted( chunk.content.text if chunk.content.text else "" for chunk in chunks) return hashlib.md5(''.join(values).encode()).hexdigest() def key_on_id(chunk): return (int(chunk.id.split('_')[1]), chunk) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/test_utils.py0000644000175100017510000004073115206113032021614 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import contextlib import logging import os import socket import tempfile import unittest from dataclasses import dataclass from typing import Callable from typing import Optional from typing import cast from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.utils import retry_with_backoff # pylint: disable=ungrouped-imports try: import yaml from pymilvus import CollectionSchema from pymilvus import FieldSchema from pymilvus import MilvusClient from pymilvus.exceptions import MilvusException from pymilvus.milvus_client import IndexParams from testcontainers.core.config import testcontainers_config from testcontainers.core.generic import DbContainer from testcontainers.milvus import MilvusContainer from apache_beam.ml.rag.enrichment.milvus_search import MilvusConnectionParameters except ImportError as e: raise unittest.SkipTest(f'RAG test util dependencies not installed: {str(e)}') _LOGGER = logging.getLogger(__name__) @dataclass class VectorDBContainerInfo: """Container information for vector database test instances. Holds connection details and container reference for testing with vector databases like Milvus in containerized environments. """ container: DbContainer host: str port: int user: str = "" password: str = "" token: str = "" id: str = "default" @property def uri(self) -> str: return f"http://{self.host}:{self.port}" class TestHelpers: @staticmethod def find_free_port(): """Find a free port on the local machine.""" with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: # Bind to port 0, which asks OS to assign a free port. s.bind(('', 0)) s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) # Return the port number assigned by OS. return s.getsockname()[1] class CustomMilvusContainer(MilvusContainer): """Custom Milvus container with configurable ports and environment setup. Extends MilvusContainer to provide custom port binding and environment configuration for testing with standalone Milvus instances. """ def __init__( # pylint: disable=bad-super-call self, image: str, service_container_port, healthcheck_container_port, **kwargs, ) -> None: # Skip the parent class's constructor and go straight to # GenericContainer. super( MilvusContainer, self, ).__init__( image=image, **kwargs) self.port = service_container_port self.healthcheck_port = healthcheck_container_port self.with_exposed_ports(service_container_port, healthcheck_container_port) # Get free host ports. service_host_port = TestHelpers.find_free_port() healthcheck_host_port = TestHelpers.find_free_port() # Bind container and host ports. self.with_bind_ports(service_container_port, service_host_port) self.with_bind_ports(healthcheck_container_port, healthcheck_host_port) self.cmd = "milvus run standalone" # Set environment variables needed for Milvus. envs = { "ETCD_USE_EMBED": "true", "ETCD_DATA_DIR": "/var/lib/milvus/etcd", "COMMON_STORAGETYPE": "local", "METRICS_PORT": str(healthcheck_container_port) } for env, value in envs.items(): self.with_env(env, value) class MilvusTestHelpers: """Helper utilities for testing Milvus vector database operations. Provides static methods for managing test containers, configuration files, and chunk comparison utilities for Milvus-based integration tests. """ # IMPORTANT: When upgrading the Milvus server version, ensure the pymilvus # Python SDK client in setup.py is updated to match. Referring to the Milvus # release notes compatibility matrix at # https://milvus.io/docs/release_notes.md or PyPI at # https://pypi.org/project/pymilvus/ for version compatibility. # Example: Milvus v2.6.0 requires pymilvus==2.6.0 (exact match required). @staticmethod def _wait_for_milvus_grpc(uri: str) -> None: """Wait until Milvus accepts RPCs. Docker may report started before gRPC is ready. """ def list_collections_probe(): client = MilvusClient(uri=uri) try: client.list_collections() finally: client.close() retry_with_backoff( list_collections_probe, max_retries=25, retry_delay=2.0, retry_backoff_factor=1.2, operation_name="Milvus client connection after container start", exception_types=(MilvusException, )) @staticmethod def start_db_container( image="milvusdb/milvus:v2.5.10", max_vec_fields=5, vector_client_max_retries=3, tc_max_retries=None) -> Optional[VectorDBContainerInfo]: service_container_port = TestHelpers.find_free_port() healthcheck_container_port = TestHelpers.find_free_port() user_yaml_creator = MilvusTestHelpers.create_user_yaml with user_yaml_creator(service_container_port, max_vec_fields) as cfg: info = None original_tc_max_tries = testcontainers_config.max_tries if tc_max_retries is not None: testcontainers_config.max_tries = tc_max_retries for i in range(vector_client_max_retries): vector_db_container: Optional[CustomMilvusContainer] = None try: vector_db_container = CustomMilvusContainer( image=image, service_container_port=service_container_port, healthcheck_container_port=healthcheck_container_port) mapped_container = vector_db_container.with_volume_mapping( cfg, "/milvus/configs/user.yaml") assert mapped_container is not None running_container: CustomMilvusContainer = mapped_container vector_db_container = running_container running_container.start() host = running_container.get_container_host_ip() port = running_container.get_exposed_port(service_container_port) info = VectorDBContainerInfo(running_container, host, port) MilvusTestHelpers._wait_for_milvus_grpc(info.uri) _LOGGER.info( "milvus db container started successfully on %s.", info.uri) break except Exception as e: stdout_logs = stderr_logs = "" if vector_db_container is not None: raw_out, raw_err = vector_db_container.get_logs() stdout_logs = raw_out.decode("utf-8") stderr_logs = raw_err.decode("utf-8") _LOGGER.warning( "Retry %d/%d: Failed to start Milvus DB container. Reason: %s. " "STDOUT logs:\n%s\nSTDERR logs:\n%s", i + 1, vector_client_max_retries, e, stdout_logs, stderr_logs) if i == vector_client_max_retries - 1: _LOGGER.error( "Unable to start milvus db container for I/O tests after %d " "retries. Tests cannot proceed. STDOUT logs:\n%s\n" "STDERR logs:\n%s", vector_client_max_retries, stdout_logs, stderr_logs) raise e finally: testcontainers_config.max_tries = original_tc_max_tries return info @staticmethod def stop_db_container(db_info: VectorDBContainerInfo): if db_info is None: _LOGGER.warning("Milvus db info is None. Skipping stop operation.") return _LOGGER.debug("Stopping milvus db container.") db_info.container.stop() _LOGGER.info("milvus db container stopped successfully.") @staticmethod def initialize_db_with_data( connc_params: MilvusConnectionParameters, config: dict): # Open the connection to the milvus db with retry. def create_client(): return MilvusClient(**connc_params.__dict__) client = retry_with_backoff( create_client, max_retries=5, retry_delay=2.0, operation_name="Test Milvus client connection", exception_types=(MilvusException, )) # Configure schema. field_schemas: list[FieldSchema] = cast(list[FieldSchema], config["fields"]) schema = CollectionSchema( fields=field_schemas, functions=config["functions"]) # Create collection with the schema. collection_name = config["collection_name"] index_function: Callable[[], IndexParams] = cast( Callable[[], IndexParams], config["index"]) client.create_collection( collection_name=collection_name, schema=schema, index_params=index_function()) # Assert that collection was created. collection_error = f"Expected collection '{collection_name}' to be created." assert client.has_collection(collection_name), collection_error # Gather all fields we have excluding 'sparse_embedding_bm25' special field. fields = list(map(lambda field: field.name, field_schemas)) # Prep data for indexing. Currently we can't insert sparse vectors for BM25 # sparse embedding field as it would be automatically generated by Milvus # through the registered BM25 function. data_ready_to_index = [] for doc in config["corpus"]: item = {} for field in fields: if field.startswith("dense_embedding"): item[field] = doc["dense_embedding"] elif field == "sparse_embedding_inner_product": item[field] = doc["sparse_embedding"] elif field == "sparse_embedding_bm25": # It is automatically generated by Milvus from the content field. continue else: item[field] = doc[field] data_ready_to_index.append(item) # Index data. result = client.insert( collection_name=collection_name, data=data_ready_to_index) # Assert that the intended data has been properly indexed. insertion_err = f'failed to insert the {result["insert_count"]} data points' assert result["insert_count"] == len(data_ready_to_index), insertion_err # Release the collection from memory. It will be loaded lazily when the # enrichment handler is invoked. client.release_collection(collection_name) # Close the connection to the Milvus database, as no further preparation # operations are needed before executing the enrichment handler. client.close() return collection_name @staticmethod @contextlib.contextmanager def create_user_yaml(service_port: int, max_vector_field_num=5): """Creates a temporary user.yaml file for Milvus configuration. This user yaml file overrides Milvus default configurations. It sets the Milvus service port to the specified container service port. The default for maxVectorFieldNum is 4, but we need 5 (one unique field for each metric). Args: service_port: Port number for the Milvus service. max_vector_field_num: Max number of vec fields allowed per collection. Yields: str: Path to the created temporary yaml file. """ with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as temp_file: # Define the content for user.yaml. user_config = { 'proxy': { 'maxVectorFieldNum': max_vector_field_num, 'port': service_port }, 'etcd': { 'use': { 'embed': True }, 'data': { 'dir': '/var/lib/milvus/etcd' } } } # Write the content to the file. yaml.dump(user_config, temp_file, default_flow_style=False) path = temp_file.name try: yield path finally: if os.path.exists(path): os.remove(path) @staticmethod def assert_chunks_equivalent( actual_chunks: list[Chunk], expected_chunks: list[Chunk]): """assert_chunks_equivalent checks for presence rather than exact match""" # Sort both lists by ID to ensure consistent ordering. actual_sorted = sorted(actual_chunks, key=lambda c: c.id) expected_sorted = sorted(expected_chunks, key=lambda c: c.id) actual_len = len(actual_sorted) expected_len = len(expected_sorted) err_msg = ( f"Different number of chunks, actual: {actual_len}, " f"expected: {expected_len}") assert actual_len == expected_len, err_msg for actual, expected in zip(actual_sorted, expected_sorted): # Assert that IDs match. assert actual.id == expected.id # Assert that dense embeddings match. err_msg = f"Dense embedding mismatch for chunk {actual.id}" assert actual.dense_embedding == expected.dense_embedding, err_msg # Assert that sparse embeddings match. err_msg = f"Sparse embedding mismatch for chunk {actual.id}" assert actual.sparse_embedding == expected.sparse_embedding, err_msg # Assert that text content match. err_msg = f"Text Content mismatch for chunk {actual.id}" assert actual.content.text == expected.content.text, err_msg # For enrichment_data, be more flexible. # If "expected" has values for enrichment_data but actual doesn't, that's # acceptable since vector search results can vary based on many factors # including implementation details, vector database state, and slight # variations in similarity calculations. # First ensure the enrichment data key exists. err_msg = f"Missing enrichment_data key in chunk {actual.id}" assert 'enrichment_data' in actual.metadata, err_msg # For enrichment_data, ensure consistent ordering of results. actual_data = actual.metadata['enrichment_data'] expected_data = expected.metadata['enrichment_data'] # If actual has enrichment data, then perform detailed validation. if actual_data: # Ensure the id key exist. err_msg = f"Missing id key in metadata {actual.id}" assert 'id' in actual_data, err_msg # Validate IDs have consistent ordering. actual_ids = sorted(actual_data['id']) expected_ids = sorted(expected_data['id']) err_msg = f"IDs in enrichment_data don't match for chunk {actual.id}" assert actual_ids == expected_ids, err_msg # Ensure the distance key exist. err_msg = f"Missing distance key in metadata {actual.id}" assert 'distance' in actual_data, err_msg # Validate distances exist and have same length as IDs. actual_distances = actual_data['distance'] expected_distances = expected_data['distance'] err_msg = ( "Number of distances doesn't match number of IDs for " f"chunk {actual.id}") assert len(actual_distances) == len(expected_distances), err_msg # Ensure the fields key exist. err_msg = f"Missing fields key in metadata {actual.id}" assert 'fields' in actual_data, err_msg # Validate fields have consistent content. # Sort fields by 'id' to ensure consistent ordering. actual_fields_sorted = sorted( actual_data['fields'], key=lambda f: f.get('id', 0)) expected_fields_sorted = sorted( expected_data['fields'], key=lambda f: f.get('id', 0)) # Compare field IDs. actual_field_ids = [f.get('id') for f in actual_fields_sorted] expected_field_ids = [f.get('id') for f in expected_fields_sorted] err_msg = f"Field IDs don't match for chunk {actual.id}" assert actual_field_ids == expected_field_ids, err_msg # Compare field content. for a_f, e_f in zip(actual_fields_sorted, expected_fields_sorted): # Ensure the id key exist. err_msg = f"Missing id key in metadata.fields {actual.id}" assert 'id' in a_f, err_msg err_msg = f"Field ID mismatch chunk {actual.id}" assert a_f['id'] == e_f['id'], err_msg # Validate field metadata. err_msg = f"Field Metadata doesn't match for chunk {actual.id}" assert a_f['metadata'] == e_f['metadata'], err_msg if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/types.py0000644000175100017510000001220515206113032020554 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Core types for RAG pipelines. This module contains the core dataclasses used throughout the RAG pipeline implementation. The primary type is EmbeddableItem, which represents any content that can be embedded and stored in a vector database. Types: - Content: Container for embeddable content - Embedding: Vector embedding with optional metadata - EmbeddableItem: Universal container for embeddable content - Chunk: Alias for EmbeddableItem (backward compatibility) """ import uuid from dataclasses import dataclass from dataclasses import field from typing import Any from typing import Optional from typing import Union @dataclass class Content: """Container for embeddable content. Args: text: Text content to be embedded. image: Image as bytes or path/URI (e.g., 'gs://bucket/img.jpg'). """ text: Optional[str] = None image: Optional[Union[bytes, str]] = None @dataclass class Embedding: """Represents vector embeddings with optional metadata. Args: dense_embedding: Dense vector representation. sparse_embedding: Optional sparse vector representation for hybrid search. """ dense_embedding: Optional[list[float]] = None sparse_embedding: Optional[tuple[list[int], list[float]]] = None @dataclass class EmbeddableItem: """Universal container for embeddable content. Represents any content that can be embedded and stored in a vector database. Use factory methods for convenient construction, or construct directly with a Content object. Examples: Text (via factory): item = EmbeddableItem.from_text( "hello world", metadata={'src': 'doc'}) Image (via factory): item = EmbeddableItem.from_image( 'gs://bucket/img.jpg') Text (direct, equivalent to old Chunk usage): item = EmbeddableItem( content=Content(text="hello"), index=3) Args: content: The content to embed. id: Unique identifier. index: Position within source document (for chunking use cases). metadata: Additional metadata (e.g., document source, language). embedding: Embedding populated by the embedding step. """ content: Content id: str = field(default_factory=lambda: str(uuid.uuid4())) index: int = 0 metadata: dict[str, Any] = field(default_factory=dict) embedding: Optional[Embedding] = None @classmethod def from_text( cls, text: str, *, id: Optional[str] = None, index: int = 0, metadata: Optional[dict[str, Any]] = None, ) -> 'EmbeddableItem': """Create an EmbeddableItem with text content. Args: text: The text content to embed id: Unique identifier (auto-generated if not provided) index: Position within source document (for chunking) metadata: Additional metadata """ return cls( content=Content(text=text), id=id or str(uuid.uuid4()), index=index, metadata=metadata or {}, ) @classmethod def from_image( cls, image: Union[bytes, str], *, id: Optional[str] = None, metadata: Optional[dict[str, Any]] = None, ) -> 'EmbeddableItem': """Create an EmbeddableItem with image content. Args: image: Image bytes or path/URI (e.g. GCS path) id: Unique identifier (auto-generated if not provided) metadata: Additional metadata """ return cls( content=Content(image=image), id=id or str(uuid.uuid4()), metadata=metadata or {}, ) @property def dense_embedding(self) -> Optional[list[float]]: return self.embedding.dense_embedding if self.embedding else None @property def sparse_embedding(self) -> Optional[tuple[list[int], list[float]]]: return self.embedding.sparse_embedding if self.embedding else None @property def content_string(self) -> str: """Returns storable string content for ingestion. Falls back through content fields in priority order: text > image URI. """ if self.content.text is not None: return self.content.text if isinstance(self.content.image, str): return self.content.image raise ValueError( f'EmbeddableItem does not contain storable string content' f' (text or image URI). {self}') # Backward compatibility alias. Existing code using Chunk continues to work # unchanged since Chunk IS EmbeddableItem. Chunk = EmbeddableItem ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/rag/utils.py0000644000175100017510000001732415206113032020557 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import re import time import uuid from collections import defaultdict from dataclasses import dataclass from dataclasses import field from typing import Any from typing import Callable from typing import Optional from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.rag.types import Embedding _LOGGER = logging.getLogger(__name__) # Default batch size for writing data to Milvus, matching # JdbcIO.DEFAULT_BATCH_SIZE. DEFAULT_WRITE_BATCH_SIZE = 1000 @dataclass class MilvusConnectionParameters: """Configurations for establishing connections to Milvus servers. Args: uri: URI endpoint for connecting to Milvus server in the format "http(s)://hostname:port". user: Username for authentication. Required if authentication is enabled and not using token authentication. password: Password for authentication. Required if authentication is enabled and not using token authentication. db_name: Database Name to connect to. Specifies which Milvus database to use. Defaults to 'default'. token: Authentication token as an alternative to username/password. timeout: Connection timeout in seconds. Uses client default if None. kwargs: Optional keyword arguments for additional connection parameters. Enables forward compatibility. """ uri: str user: str = field(default_factory=str) password: str = field(default_factory=str) db_name: str = "default" token: str = field(default_factory=str) timeout: Optional[float] = None kwargs: dict[str, Any] = field(default_factory=dict) def __post_init__(self): if not self.uri: raise ValueError("URI must be provided for Milvus connection") # Generate unique alias if not provided. One-to-one mapping between alias # and connection - each alias represents exactly one Milvus connection. if "alias" not in self.kwargs: alias = f"milvus_conn_{uuid.uuid4().hex[:8]}" self.kwargs["alias"] = alias class MilvusHelpers: """Utility class providing helper methods for Milvus vector db operations.""" @staticmethod def sparse_embedding( sparse_vector: Optional[tuple[list[int], list[float]]] ) -> Optional[dict[int, float]]: if not sparse_vector: return None # Converts sparse embedding from (indices, values) tuple format to # Milvus-compatible values dict format {dimension_index: value, ...}. indices, values = sparse_vector return {int(idx): float(val) for idx, val in zip(indices, values)} def parse_chunk_strings(chunk_str_list: list[str]) -> list[Chunk]: parsed_chunks = [] # Define safe globals and disable built-in functions for safety. safe_globals = { 'Chunk': Chunk, 'Content': Content, 'Embedding': Embedding, 'defaultdict': defaultdict, 'list': list, '__builtins__': {} } for raw_str in chunk_str_list: try: # replace "" with actual list reference. cleaned_str = re.sub( r"defaultdict\(", "defaultdict(list", raw_str) # Evaluate string in restricted environment. chunk = eval(cleaned_str, safe_globals) # pylint: disable=eval-used if isinstance(chunk, Chunk): parsed_chunks.append(chunk) else: raise ValueError("Parsed object is not a Chunk instance") except Exception as e: raise ValueError(f"Error parsing string:\n{raw_str}\n{e}") return parsed_chunks def unpack_dataclass_with_kwargs(dataclass_instance): """Unpacks dataclass fields into a flat dict, merging kwargs with precedence. Args: dataclass_instance: Dataclass instance to unpack. Returns: dict: Flattened dictionary with kwargs taking precedence over fields. """ # Create a copy of the dataclass's __dict__. params_dict: dict = dataclass_instance.__dict__.copy() # Extract the nested kwargs dictionary. nested_kwargs = params_dict.pop('kwargs', {}) # Merge the dictionaries, with nested_kwargs taking precedence # in case of duplicate keys. return {**params_dict, **nested_kwargs} def retry_with_backoff( operation: Callable[[], Any], max_retries: int = 3, retry_delay: float = 1.0, retry_backoff_factor: float = 2.0, operation_name: str = "operation", exception_types: tuple[type[BaseException], ...] = (Exception, ) ) -> Any: """Executes an operation with retry logic and exponential backoff. This is a generic retry utility that can be used for any operation that may fail transiently. It retries the operation with exponential backoff between attempts. Note: This utility is designed for one-time setup operations and complements Apache Beam's RequestResponseIO pattern. Use retry_with_backoff() for: * Establishing client connections in __enter__() methods (e.g., creating MilvusClient instances, database connections) before processing elements * One-time setup/teardown operations in DoFn lifecycle methods * Operations outside of per-element processing where retry is needed For per-element operations (e.g., API calls within Caller.__call__), use RequestResponseIO which already provides automatic retry with exponential backoff, failure handling, caching, and other features. See: https://beam.apache.org/documentation/io/built-in/webapis/ Args: operation: Callable that performs the operation to retry. Should return the result of the operation. max_retries: Maximum number of retry attempts. Default is 3. retry_delay: Initial delay in seconds between retries. Default is 1.0. retry_backoff_factor: Multiplier for the delay after each retry. Default is 2.0 (exponential backoff). operation_name: Name of the operation for logging purposes. Default is "operation". exception_types: Tuple of exception types to catch and retry. Default is (Exception,) which catches all exceptions. Returns: The result of the operation if successful. Raises: The last exception encountered if all retry attempts fail. Example: >>> def connect_to_service(): ... return service.connect(host="localhost") >>> client = retry_with_backoff( ... connect_to_service, ... max_retries=5, ... retry_delay=2.0, ... operation_name="service connection") """ last_exception = None for attempt in range(max_retries + 1): try: result = operation() _LOGGER.info( "Successfully completed %s on attempt %d", operation_name, attempt + 1) return result except exception_types as e: last_exception = e if attempt < max_retries: delay = retry_delay * (retry_backoff_factor**attempt) _LOGGER.warning( "%s attempt %d failed: %s. Retrying in %.2f seconds...", operation_name, attempt + 1, e, delay) time.sleep(delay) else: _LOGGER.error( "Failed %s after %d attempts", operation_name, max_retries + 1) raise last_exception ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5184224 apache_beam-2.74.0/apache_beam/ml/transforms/0000755000175100017510000000000015206113160020465 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/__init__.py0000644000175100017510000000142015206113032022571 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/base.py0000644000175100017510000007701615206113032021762 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import abc import functools import logging import os import tempfile import uuid from collections.abc import Callable from collections.abc import Mapping from collections.abc import Sequence from dataclasses import dataclass from typing import Any from typing import Generic from typing import Iterable from typing import Optional from typing import TypeVar from typing import Union from typing import cast import jsonpickle import numpy as np import apache_beam as beam from apache_beam.io.filesystems import FileSystems from apache_beam.metrics.metric import Metrics from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import ModelT from apache_beam.ml.inference.base import RunInferenceDLQ from apache_beam.options.pipeline_options import PipelineOptions _LOGGER = logging.getLogger(__name__) _ATTRIBUTE_FILE_NAME = 'attributes.json' __all__ = [ 'MLTransform', 'ProcessHandler', 'MLTransformProvider', 'BaseOperation', 'EmbeddingsManager' ] TransformedDatasetT = TypeVar('TransformedDatasetT') TransformedMetadataT = TypeVar('TransformedMetadataT') # Input/Output types to the MLTransform. MLTransformOutputT = TypeVar('MLTransformOutputT') ExampleT = TypeVar('ExampleT') # Input to the apply() method of BaseOperation. OperationInputT = TypeVar('OperationInputT') # Output of the apply() method of BaseOperation. OperationOutputT = TypeVar('OperationOutputT') # Input to the EmbeddingTypeAdapter input_fn EmbeddingTypeAdapterInputT = TypeVar( 'EmbeddingTypeAdapterInputT') # e.g., Chunk # Output of the EmbeddingTypeAdapter output_fn EmbeddingTypeAdapterOutputT = TypeVar( 'EmbeddingTypeAdapterOutputT') # e.g., Embedding @dataclass class EmbeddingTypeAdapter(Generic[EmbeddingTypeAdapterInputT, EmbeddingTypeAdapterOutputT]): """Adapts input types to text for embedding and converts output embeddings. Args: input_fn: Function to extract text for embedding from input type output_fn: Function to create output type from input and embeddings """ input_fn: Callable[[Sequence[EmbeddingTypeAdapterInputT]], list[str]] output_fn: Callable[[Sequence[EmbeddingTypeAdapterInputT], Sequence[Any]], list[EmbeddingTypeAdapterOutputT]] def __reduce__(self): """Custom serialization that preserves type information during jsonpickle.""" return (self.__class__, (self.input_fn, self.output_fn)) def _map_errors_to_beam_row(element, cls_name=None): row_elements = { 'element': element[0], 'msg': str(element[1][1]), 'stack': str(element[1][2]), } if cls_name is not None: row_elements['transform_name'] = cls_name return beam.Row(**row_elements) class ArtifactMode(object): PRODUCE = 'produce' CONSUME = 'consume' class MLTransformProvider: """ Data processing transforms that are intended to be used with MLTransform should subclass MLTransformProvider and implement get_ptransform_for_processing(). get_ptransform_for_processing() method should return a PTransform that can be used to process the data. """ @abc.abstractmethod def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: """ Returns a PTransform that can be used to process the data. """ def get_counter(self): """ Returns the counter name for the data processing transform. """ counter_name = self.__class__.__name__ return Metrics.counter(MLTransform, f'BeamML_{counter_name}') class BaseOperation(Generic[OperationInputT, OperationOutputT], MLTransformProvider, abc.ABC): def __init__(self, columns: list[str]) -> None: """ Base Opertation class data processing transformations. Args: columns: List of column names to apply the transformation. """ self.columns = columns @abc.abstractmethod def apply_transform(self, data: OperationInputT, output_column_name: str) -> dict[str, OperationOutputT]: """ Define any processing logic in the apply_transform() method. processing logics are applied on inputs and returns a transformed output. Args: inputs: input data. """ def __call__(self, data: OperationInputT, output_column_name: str) -> dict[str, OperationOutputT]: """ This method is called when the instance of the class is called. This method will invoke the apply() method of the class. """ transformed_data = self.apply_transform(data, output_column_name) return transformed_data class ProcessHandler( beam.PTransform[beam.PCollection[ExampleT], Union[beam.PCollection[MLTransformOutputT], tuple[beam.PCollection[MLTransformOutputT], beam.PCollection[beam.Row]]]], abc.ABC): """ Only for internal use. No backwards compatibility guarantees. """ @abc.abstractmethod def append_transform(self, transform: BaseOperation): """ Append transforms to the ProcessHandler. """ def _dict_input_fn( columns: Sequence[str], batch: Sequence[Union[dict[str, Any], beam.Row]]) -> list[str]: """Extract text from specified columns in batch.""" if batch and hasattr(batch[0], '_asdict'): batch = [row._asdict() if hasattr(row, '_asdict') else row for row in batch] if not batch or not isinstance(batch[0], dict): raise TypeError( 'Expected data to be dicts, got ' f'{type(batch[0])} instead.') result = [] expected_keys = set(batch[0].keys()) expected_columns = set(columns) # Process one batch item at a time for item in batch: item_keys = item.keys() if isinstance(item, dict) else set() if set(item_keys) != expected_keys: extra_keys = item_keys - expected_keys missing_keys = expected_keys - item_keys raise RuntimeError( f'All dicts in batch must have the same keys. ' f'extra keys: {extra_keys}, ' f'missing keys: {missing_keys}') missing_columns = expected_columns - item_keys if (missing_columns): raise RuntimeError( f'Data does not contain the following columns ' f': {missing_columns}.') # Get all columns for this item for col in columns: if isinstance(item, dict): result.append(item[col]) return result def _dict_output_fn( columns: Sequence[str], batch: Sequence[Union[dict[str, Any], beam.Row]], embeddings: Sequence[Any]) -> list[Union[dict[str, Any], beam.Row]]: """Map embeddings back to columns in batch.""" is_beam_row = False if batch and hasattr(batch[0], '_asdict'): is_beam_row = True batch = [row._asdict() if hasattr(row, '_asdict') else row for row in batch] result = [] for batch_idx, item in enumerate(batch): for col_idx, col in enumerate(columns): embedding_idx = batch_idx * len(columns) + col_idx if isinstance(item, dict): item[col] = embeddings[embedding_idx] result.append(item) if is_beam_row: result = [beam.Row(**item) for item in result if isinstance(item, dict)] return result def _create_dict_adapter( columns: list[str]) -> EmbeddingTypeAdapter[dict[str, Any], dict[str, Any]]: """Create adapter for dict-based processing.""" return EmbeddingTypeAdapter[dict[str, Any], dict[str, Any]]( input_fn=cast( Callable[[Sequence[dict[str, Any]]], list[str]], functools.partial(_dict_input_fn, columns)), output_fn=cast( Callable[[Sequence[dict[str, Any]], Sequence[Any]], list[dict[str, Any]]], functools.partial(_dict_output_fn, columns))) # TODO:https://github.com/apache/beam/issues/29356 # Add support for inference_fn class EmbeddingsManager(MLTransformProvider): def __init__( self, *, columns: Optional[list[str]] = None, type_adapter: Optional[EmbeddingTypeAdapter] = None, # common args for all ModelHandlers. load_model_args: Optional[dict[str, Any]] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, large_model: bool = False, **kwargs): self.load_model_args = load_model_args or {} self.min_batch_size = min_batch_size self.max_batch_size = max_batch_size self.large_model = large_model self.columns = columns if columns is not None: self.type_adapter = _create_dict_adapter(columns) elif type_adapter is not None: self.type_adapter = type_adapter else: raise ValueError("Either columns or type_adapter must be specified") self.inference_args = kwargs.pop('inference_args', {}) if kwargs: _LOGGER.warning("Ignoring the following arguments: %s", kwargs.keys()) # TODO:https://github.com/apache/beam/pull/29564 add set_model_handler method @abc.abstractmethod def get_model_handler(self) -> ModelHandler: """ Return framework specific model handler. """ def get_columns_to_apply(self): return self.columns class MLTransform( beam.PTransform[beam.PCollection[ExampleT], Union[beam.PCollection[MLTransformOutputT], tuple[beam.PCollection[MLTransformOutputT], beam.PCollection[beam.Row]]]], Generic[ExampleT, MLTransformOutputT]): def __init__( self, *, write_artifact_location: Optional[str] = None, read_artifact_location: Optional[str] = None, transforms: Optional[list[MLTransformProvider]] = None): """ MLTransform is a Beam PTransform that can be used to apply transformations to the data. MLTransform is used to wrap the data processing transforms provided by Apache Beam. MLTransform works in two modes: write and read. In the write mode, MLTransform will apply the transforms to the data and store the artifacts in the write_artifact_location. In the read mode, MLTransform will read the artifacts from the read_artifact_location and apply the transforms to the data. The artifact location should be a valid storage path where the artifacts can be written to or read from. Note that when consuming artifacts, it is not necessary to pass the transforms since they are inherently stored within the artifacts themselves. Args: write_artifact_location: A storage location for artifacts resulting from MLTransform. These artifacts include transformations applied to the dataset and generated values like min, max from ScaleTo01, and mean, var from ScaleToZScore. Artifacts are produced and written to this location when using `write_artifact_mode`. Later MLTransforms can reuse produced artifacts by setting `read_artifact_mode` instead of `write_artifact_mode`. The value assigned to `write_artifact_location` should be a valid storage directory that the artifacts from this transform can be written to. If no directory exists at this location, one will be created. This will overwrite any artifacts already in this location, so distinct locations should be used for each instance of MLTransform. Only one of write_artifact_location and read_artifact_location should be specified. read_artifact_location: A storage location to read artifacts resulting froma previous MLTransform. These artifacts include transformations applied to the dataset and generated values like min, max from ScaleTo01, and mean, var from ScaleToZScore. Note that when consuming artifacts, it is not necessary to pass the transforms since they are inherently stored within the artifacts themselves. The value assigned to `read_artifact_location` should be a valid storage path where the artifacts can be read from. Only one of write_artifact_location and read_artifact_location should be specified. transforms: A list of transforms to apply to the data. All the transforms are applied in the order they are specified. The input of the i-th transform is the output of the (i-1)-th transform. Multi-input transforms are not supported yet. """ if read_artifact_location and write_artifact_location: raise ValueError( 'Only one of read_artifact_location or write_artifact_location can ' 'be specified to initialize MLTransform') if not read_artifact_location and not write_artifact_location: raise ValueError( 'Either a read_artifact_location or write_artifact_location must be ' 'specified to initialize MLTransform') if read_artifact_location: artifact_location = read_artifact_location artifact_mode = ArtifactMode.CONSUME if transforms: raise ValueError( 'Transforms should not be passed in read mode. In read mode, ' 'the transforms are read from the artifact location.') else: artifact_location = write_artifact_location # type: ignore[assignment] artifact_mode = ArtifactMode.PRODUCE self._parent_artifact_location = artifact_location self._artifact_mode = artifact_mode self.transforms = transforms or [] self._counter = Metrics.counter( MLTransform, f'BeamML_{self.__class__.__name__}') self._with_exception_handling = False self._exception_handling_args: dict[str, Any] = {} def expand( self, pcoll: beam.PCollection[ExampleT] ) -> Union[beam.PCollection[MLTransformOutputT], tuple[beam.PCollection[MLTransformOutputT], beam.PCollection[beam.Row]]]: """ This is the entrypoint for the MLTransform. This method will invoke the process_data() method of the ProcessHandler instance to process the incoming data. process_data takes in a PCollection and applies the PTransforms necessary to process the data and returns a PCollection of transformed data. Args: pcoll: A PCollection of ExampleT type. Returns: A PCollection of MLTransformOutputT type """ upstream_errors = [] _ = [self._validate_transform(transform) for transform in self.transforms] if self._artifact_mode == ArtifactMode.PRODUCE: ptransform_partitioner = _MLTransformToPTransformMapper( transforms=self.transforms, artifact_location=self._parent_artifact_location, artifact_mode=self._artifact_mode, pipeline_options=pcoll.pipeline.options) ptransform_list = ptransform_partitioner.create_and_save_ptransform_list() else: ptransform_list = ( _MLTransformToPTransformMapper.load_transforms_from_artifact_location( self._parent_artifact_location)) # the saved transforms has artifact mode set to PRODUCE. # set the artifact mode to CONSUME. for i in range(len(ptransform_list)): if hasattr(ptransform_list[i], 'artifact_mode'): ptransform_list[i].artifact_mode = self._artifact_mode transform_name = None for ptransform in ptransform_list: if self._with_exception_handling: if hasattr(ptransform, 'with_exception_handling'): ptransform = ptransform.with_exception_handling( **self._exception_handling_args) pcoll, bad_results = pcoll | ptransform # RunInference outputs a RunInferenceDLQ instead of a PCollection. # since TFTProcessHandler and RunInferene are supported, try to infer # the type of bad_results and append it to the list of errors. if isinstance(bad_results, RunInferenceDLQ): bad_results = bad_results.failed_inferences transform_name = ptransform.annotations()['model_handler'] elif not isinstance(bad_results, beam.PCollection): raise NotImplementedError( f'Unexpected type for bad_results: {type(bad_results)}') bad_results = bad_results | beam.Map( lambda x: _map_errors_to_beam_row(x, transform_name)) upstream_errors.append(bad_results) else: pcoll = pcoll | ptransform _ = ( pcoll.pipeline | "MLTransformMetricsUsage" >> MLTransformMetricsUsage(self)) if self._with_exception_handling: bad_pcoll = (upstream_errors | beam.Flatten()) return pcoll, bad_pcoll # type: ignore[return-value] return pcoll # type: ignore[return-value] def with_transform(self, transform: MLTransformProvider): """ Add a transform to the MLTransform pipeline. Args: transform: A BaseOperation instance. Returns: A MLTransform instance. """ self._validate_transform(transform) self.transforms.append(transform) return self def _validate_transform(self, transform): if not isinstance(transform, MLTransformProvider): raise TypeError( 'transform must be a subclass of BaseOperation. ' 'Got: %s instead.' % type(transform)) def with_exception_handling( self, *, exc_class=Exception, use_subprocess=False, threshold=1): self._with_exception_handling = True self._exception_handling_args = { 'exc_class': exc_class, 'use_subprocess': use_subprocess, 'threshold': threshold } return self class MLTransformMetricsUsage(beam.PTransform): def __init__(self, ml_transform: MLTransform): self._ml_transform = ml_transform self._ml_transform._counter.inc() def expand(self, pipeline): def _increment_counters(): # increment for MLTransform. self._ml_transform._counter.inc() # increment if data processing transforms are passed. transforms = self._ml_transform.transforms if transforms: for transform in transforms: transform.get_counter().inc() _ = ( pipeline | beam.Create([None]) | beam.Map(lambda _: _increment_counters())) class _TransformAttributeManager: """ Base class used for saving and loading the attributes. """ @staticmethod def save_attributes(artifact_location): """ Save the attributes to json file using stdlib json. """ raise NotImplementedError @staticmethod def load_attributes(artifact_location): """ Load the attributes from json file. """ raise NotImplementedError class _JsonPickleTransformAttributeManager(_TransformAttributeManager): """ Use Jsonpickle to save and load the attributes. Here the attributes refer to the list of PTransforms that are used to process the data. jsonpickle is used to serialize the PTransforms and save it to a json file and is compatible across python versions. """ @staticmethod def _is_remote_path(path): is_gcs = path.find('gs://') != -1 # TODO:https://github.com/apache/beam/issues/29356 # Add support for other remote paths. if not is_gcs and path.find('://') != -1: raise RuntimeError( "Artifact locations are currently supported for only available for " "local paths and GCS paths. Got: %s" % path) return is_gcs @staticmethod def save_attributes( ptransform_list, artifact_location, **kwargs, ): # if an artifact location is present, instead of overwriting the # existing file, raise an error since the same artifact location # can be used by multiple beam jobs and this could result in undesired # behavior. if FileSystems.exists(FileSystems.join(artifact_location, _ATTRIBUTE_FILE_NAME)): raise FileExistsError( "The artifact location %s already exists and contains %s. Please " "specify a different location." % (artifact_location, _ATTRIBUTE_FILE_NAME)) if _JsonPickleTransformAttributeManager._is_remote_path(artifact_location): temp_dir = tempfile.mkdtemp() temp_json_file = os.path.join(temp_dir, _ATTRIBUTE_FILE_NAME) with open(temp_json_file, 'w+') as f: f.write(jsonpickle.encode(ptransform_list)) with open(temp_json_file, 'rb') as f: from apache_beam.runners.dataflow.internal import apiclient _LOGGER.info('Creating artifact location: %s', artifact_location) # pipeline options required to for the client to configure project. options = kwargs.get('options') try: apiclient.DataflowApplicationClient(options=options).stage_file( gcs_or_local_path=artifact_location, file_name=_ATTRIBUTE_FILE_NAME, stream=f, mime_type='application/json') except Exception as exc: if not options: raise RuntimeError( "Failed to create Dataflow client. " "Pipeline options are required to save the attributes." "in the artifact location %s" % artifact_location) from exc raise else: if not FileSystems.exists(artifact_location): FileSystems.mkdirs(artifact_location) # FileSystems.open() fails if the file does not exist. with open(os.path.join(artifact_location, _ATTRIBUTE_FILE_NAME), 'w+') as f: f.write(jsonpickle.encode(ptransform_list)) @staticmethod def load_attributes(artifact_location): with FileSystems.open(os.path.join(artifact_location, _ATTRIBUTE_FILE_NAME), 'rb') as f: return jsonpickle.decode(f.read()) _transform_attribute_manager = _JsonPickleTransformAttributeManager class _MLTransformToPTransformMapper: """ This class takes in a list of data processing transforms compatible to be wrapped around MLTransform and returns a list of PTransforms that are used to run the data processing transforms. The _MLTransformToPTransformMapper is responsible for loading and saving the PTransforms or attributes of PTransforms to the artifact location to seal the gap between the training and inference pipelines. """ def __init__( self, transforms: list[MLTransformProvider], artifact_location: str, artifact_mode: str = ArtifactMode.PRODUCE, pipeline_options: Optional[PipelineOptions] = None, ): self.transforms = transforms self._parent_artifact_location = artifact_location self.artifact_mode = artifact_mode self.pipeline_options = pipeline_options def create_and_save_ptransform_list(self): ptransform_list = self.create_ptransform_list() self.save_transforms_in_artifact_location(ptransform_list) return ptransform_list def create_ptransform_list(self): previous_ptransform_type = None current_ptransform = None ptransform_list = [] for transform in self.transforms: if not isinstance(transform, MLTransformProvider): raise RuntimeError( 'Transforms must be instances of MLTransformProvider and ' 'implement get_ptransform_for_processing() method.') # for each instance of PTransform, create a new artifact location current_ptransform = transform.get_ptransform_for_processing( artifact_location=os.path.join( self._parent_artifact_location, uuid.uuid4().hex[:6]), artifact_mode=self.artifact_mode) append_transform = hasattr(current_ptransform, 'append_transform') if (type(current_ptransform) != previous_ptransform_type) or not append_transform: ptransform_list.append(current_ptransform) previous_ptransform_type = type(current_ptransform) # If different PTransform is appended to the list and the PTransform # supports append_transform, append the transform to the PTransform. if append_transform: ptransform_list[-1].append_transform(transform) return ptransform_list def save_transforms_in_artifact_location(self, ptransform_list): """ Save the ptransform references to json file. """ _transform_attribute_manager.save_attributes( ptransform_list=ptransform_list, artifact_location=self._parent_artifact_location, options=self.pipeline_options) @staticmethod def load_transforms_from_artifact_location(artifact_location): return _transform_attribute_manager.load_attributes(artifact_location) class _EmbeddingHandler(ModelHandler): """ A ModelHandler intended to be work on list[dict[str, Any]] inputs. The inputs to the model handler are expected to be a list of dicts. For example, if the original mode is used with RunInference to take a PCollection[E] to a PCollection[P], this ModelHandler would take a PCollection[dict[str, E]] to a PCollection[dict[str, P]]. _EmbeddingHandler will accept an EmbeddingsManager instance, which contains the details of the model to be loaded and the inference_fn to be used. The purpose of _EmbeddingHandler is to generate embeddings for general inputs using the EmbeddingsManager instance. This is an internal class and offers no backwards compatibility guarantees. Args: embeddings_manager: An EmbeddingsManager instance. """ def __init__(self, embeddings_manager: EmbeddingsManager): self.embedding_config = embeddings_manager self._underlying = self.embedding_config.get_model_handler() self.columns = self.embedding_config.get_columns_to_apply() def load_model(self): model = self._underlying.load_model() return model def _validate_column_data(self, batch): pass def run_inference( self, batch: Sequence[dict[str, list[str]]], model: ModelT, inference_args: Optional[dict[str, Any]] = None, ) -> list[dict[str, Union[list[float], list[str]]]]: """ Runs inference on a batch of text inputs. The inputs are expected to be a list of dicts. Each dict should have the same keys, and the shape should be of the same size for a single key across the batch. """ embedding_input = self.embedding_config.type_adapter.input_fn(batch) self._validate_column_data(batch=embedding_input) prediction = self._underlying.run_inference( embedding_input, model, inference_args) # Convert prediction to Sequence[Any] if isinstance(prediction, np.ndarray): prediction_seq = prediction.tolist() elif isinstance(prediction, Iterable) and not isinstance(prediction, (str, bytes)): prediction_seq = list(prediction) else: prediction_seq = [prediction] return self.embedding_config.type_adapter.output_fn(batch, prediction_seq) def get_metrics_namespace(self) -> str: return ( self._underlying.get_metrics_namespace() or 'BeamML_EmbeddingHandler') def batch_elements_kwargs(self) -> Mapping[str, Any]: batch_sizes_map = {} if self.embedding_config.max_batch_size: batch_sizes_map['max_batch_size'] = self.embedding_config.max_batch_size if self.embedding_config.min_batch_size: batch_sizes_map['min_batch_size'] = self.embedding_config.min_batch_size return (self._underlying.batch_elements_kwargs() or batch_sizes_map) def __repr__(self): return self._underlying.__repr__() def validate_inference_args(self, _): pass class _TextEmbeddingHandler(_EmbeddingHandler): """ A ModelHandler intended to be work on list[dict[str, str]] inputs. The inputs to the model handler are expected to be a list of dicts. For example, if the original mode is used with RunInference to take a PCollection[E] to a PCollection[P], this ModelHandler would take a PCollection[dict[str, E]] to a PCollection[dict[str, P]]. _TextEmbeddingHandler will accept an EmbeddingsManager instance, which contains the details of the model to be loaded and the inference_fn to be used. The purpose of _TextEmbeddingHandler is to generate embeddings for text inputs using the EmbeddingsManager instance. If the input is not a text column, a RuntimeError will be raised. This is an internal class and offers no backwards compatibility guarantees. Args: embeddings_manager: An EmbeddingsManager instance. """ def _validate_column_data(self, batch): if not isinstance(batch[0], (str, bytes)): raise TypeError( 'Embeddings can only be generated on dict[str, str].' f'Got dict[str, {type(batch[0])}] instead.') def get_metrics_namespace(self) -> str: return ( self._underlying.get_metrics_namespace() or 'BeamML_TextEmbeddingHandler') class _ImageEmbeddingHandler(_EmbeddingHandler): """ A ModelHandler intended to be work on list[dict[str, Image]] inputs. The inputs to the model handler are expected to be a list of dicts. For example, if the original mode is used with RunInference to take a PCollection[E] to a PCollection[P], this ModelHandler would take a PCollection[dict[str, E]] to a PCollection[dict[str, P]]. _ImageEmbeddingHandler will accept an EmbeddingsManager instance, which contains the details of the model to be loaded and the inference_fn to be used. The purpose of _ImageEmbeddingHandler is to generate embeddings for image inputs using the EmbeddingsManager instance. If the input is not an Image representation column, a RuntimeError will be raised. This is an internal class and offers no backwards compatibility guarantees. Args: embeddings_manager: An EmbeddingsManager instance. """ def _validate_column_data(self, batch): # Don't want to require framework-specific imports # here, so just catch columns of primatives for now. if isinstance(batch[0], (int, str, float, bool)): raise TypeError( 'Embeddings can only be generated on dict[str, Image].' f'Got dict[str, {type(batch[0])}] instead.') def get_metrics_namespace(self) -> str: return ( self._underlying.get_metrics_namespace() or 'BeamML_ImageEmbeddingHandler') class _MultiModalEmbeddingHandler(_EmbeddingHandler): """ A ModelHandler intended to be work on list[dict[str, TypedDict(Image, Video, str)]] inputs. The inputs to the model handler are expected to be a list of dicts. For example, if the original mode is used with RunInference to take a PCollection[E] to a PCollection[P], this ModelHandler would take a PCollection[dict[str, E]] to a PCollection[dict[str, P]]. _MultiModalEmbeddingHandler will accept an EmbeddingsManager instance, which contains the details of the model to be loaded and the inference_fn to be used. The purpose of _MultiMOdalEmbeddingHandler is to generate embeddings for image, video, and text inputs using the EmbeddingsManager instance. If the input is not an Image representation column, a RuntimeError will be raised. This is an internal class and offers no backwards compatibility guarantees. Args: embeddings_manager: An EmbeddingsManager instance. """ def _validate_column_data(self, batch): # Don't want to require framework-specific imports # here, so just catch columns of primatives for now. if isinstance(batch[0], (int, str, float, bool)): raise TypeError( 'Embeddings can only be generated on ' ' dict[str, dataclass] types. ' f'Got dict[str, {type(batch[0])}] instead.') def get_metrics_namespace(self) -> str: return ( self._underlying.get_metrics_namespace() or 'BeamML_MultiModalEmbeddingHandler') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/base_test.py0000644000175100017510000007460215206113032023017 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import os import secrets import shutil import tempfile import time import unittest from collections.abc import Sequence from dataclasses import dataclass from typing import Any from typing import Optional import numpy as np from parameterized import param from parameterized import parameterized import apache_beam as beam from apache_beam.metrics.metric import MetricsFilter from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import RunInference from apache_beam.ml.transforms import base from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from apache_beam.ml.transforms import tft from apache_beam.ml.transforms.handlers import TFTProcessHandler from apache_beam.ml.transforms.tft import TFTOperation except ImportError: tft = None # type: ignore try: import PIL from PIL.Image import Image as PIL_Image except ImportError: PIL = None # type: ignore[assignment] PIL_Image = Any # type: ignore[misc, assignment] try: class _FakeOperation(TFTOperation): def __init__(self, name, *args, **kwargs): super().__init__(*args, **kwargs) self.name = name def apply_transform(self, inputs, output_column_name, **kwargs): return {output_column_name: inputs} except: # pylint: disable=bare-except pass try: from apache_beam.runners.dataflow.internal import apiclient except ImportError: apiclient = None # type: ignore class BaseMLTransformTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) def test_ml_transform_no_read_or_write_artifact_lcoation(self): with self.assertRaises(ValueError): _ = base.MLTransform(transforms=[]) @unittest.skipIf(tft is None, 'tft module is not installed.') def test_ml_transform_appends_transforms_to_process_handler_correctly(self): fake_fn_1 = _FakeOperation(name='fake_fn_1', columns=['x']) transforms = [fake_fn_1] ml_transform = base.MLTransform( transforms=transforms, write_artifact_location=self.artifact_location) ml_transform = ml_transform.with_transform( transform=_FakeOperation(name='fake_fn_2', columns=['x'])) self.assertEqual(len(ml_transform.transforms), 2) self.assertEqual(ml_transform.transforms[0].name, 'fake_fn_1') self.assertEqual(ml_transform.transforms[1].name, 'fake_fn_2') @unittest.skipIf(tft is None, 'tft module is not installed.') def test_ml_transform_on_dict(self): transforms = [tft.ScaleTo01(columns=['x'])] data = [{'x': 1}, {'x': 2}] with beam.Pipeline() as p: result = ( p | beam.Create(data) | base.MLTransform( write_artifact_location=self.artifact_location, transforms=transforms)) expected_output = [ np.array([0.0], dtype=np.float32), np.array([1.0], dtype=np.float32), ] actual_output = result | beam.Map(lambda x: x.x) assert_that( actual_output, equal_to(expected_output, equals_fn=np.array_equal)) @unittest.skipIf(tft is None, 'tft module is not installed.') def test_ml_transform_on_list_dict(self): transforms = [tft.ScaleTo01(columns=['x'])] data = [{'x': [1, 2, 3]}, {'x': [4, 5, 6]}] with beam.Pipeline() as p: result = ( p | beam.Create(data) | base.MLTransform( transforms=transforms, write_artifact_location=self.artifact_location)) expected_output = [ np.array([0, 0.2, 0.4], dtype=np.float32), np.array([0.6, 0.8, 1], dtype=np.float32), ] actual_output = result | beam.Map(lambda x: x.x) assert_that( actual_output, equal_to(expected_output, equals_fn=np.array_equal)) @parameterized.expand([ param( input_data=[{ 'x': 1, 'y': 2.0, }], input_types={ 'x': int, 'y': float }, expected_dtype={ 'x': Sequence[np.float32], 'y': Sequence[np.float32], }, ), param( input_data=[{ 'x': np.array([1], dtype=np.int64), 'y': np.array([2.0], dtype=np.float32), }], input_types={ 'x': np.int32, 'y': np.float32 }, expected_dtype={ 'x': Sequence[np.float32], 'y': Sequence[np.float32], }, ), param( input_data=[{ 'x': [1, 2, 3], 'y': [2.0, 3.0, 4.0] }], input_types={ 'x': list[int], 'y': list[float] }, expected_dtype={ 'x': Sequence[np.float32], 'y': Sequence[np.float32], }, ), param( input_data=[{ 'x': [1, 2, 3], 'y': [2.0, 3.0, 4.0] }], input_types={ 'x': Sequence[int], 'y': Sequence[float], }, expected_dtype={ 'x': Sequence[np.float32], 'y': Sequence[np.float32], }, ), ]) @unittest.skipIf(tft is None, 'tft module is not installed.') def test_ml_transform_dict_output_pcoll_schema( self, input_data, input_types, expected_dtype): transforms = [tft.ScaleTo01(columns=['x'])] with beam.Pipeline() as p: schema_data = ( p | beam.Create(input_data) | beam.Map(lambda x: beam.Row(**x)).with_output_types( beam.row_type.RowTypeConstraint.from_fields( list(input_types.items())))) transformed_data = schema_data | base.MLTransform( write_artifact_location=self.artifact_location, transforms=transforms) for name, typ in transformed_data.element_type._fields: if name in expected_dtype: self.assertEqual(expected_dtype[name], typ) @unittest.skipIf(tft is None, 'tft module is not installed.') def test_ml_transform_fail_for_non_global_windows_in_produce_mode(self): transforms = [tft.ScaleTo01(columns=['x'])] with beam.Pipeline() as p: with self.assertRaises(RuntimeError): _ = ( p | beam.Create([{ 'x': 1, 'y': 2.0 }]) | beam.WindowInto(beam.window.FixedWindows(1)) | base.MLTransform( transforms=transforms, write_artifact_location=self.artifact_location, )) @unittest.skipIf(tft is None, 'tft module is not installed.') def test_ml_transform_on_multiple_columns_single_transform(self): transforms = [tft.ScaleTo01(columns=['x', 'y'])] data = [{'x': [1, 2, 3], 'y': [1.0, 10.0, 20.0]}] with beam.Pipeline() as p: result = ( p | beam.Create(data) | base.MLTransform( transforms=transforms, write_artifact_location=self.artifact_location)) expected_output_x = [ np.array([0, 0.5, 1], dtype=np.float32), ] expected_output_y = [np.array([0, 0.47368422, 1], dtype=np.float32)] actual_output_x = result | beam.Map(lambda x: x.x) actual_output_y = result | beam.Map(lambda x: x.y) assert_that( actual_output_x, equal_to(expected_output_x, equals_fn=np.array_equal)) assert_that( actual_output_y, equal_to(expected_output_y, equals_fn=np.array_equal), label='y') @unittest.skipIf(tft is None, 'tft module is not installed.') def test_ml_transforms_on_multiple_columns_multiple_transforms(self): transforms = [ tft.ScaleTo01(columns=['x']), tft.ComputeAndApplyVocabulary(columns=['y']) ] data = [{'x': [1, 2, 3], 'y': ['a', 'b', 'c']}] with beam.Pipeline() as p: result = ( p | beam.Create(data) | base.MLTransform( transforms=transforms, write_artifact_location=self.artifact_location)) expected_output_x = [ np.array([0, 0.5, 1], dtype=np.float32), ] expected_output_y = [np.array([2, 1, 0])] actual_output_x = result | beam.Map(lambda x: x.x) actual_output_y = result | beam.Map(lambda x: x.y) assert_that( actual_output_x, equal_to(expected_output_x, equals_fn=np.array_equal)) assert_that( actual_output_y, equal_to(expected_output_y, equals_fn=np.array_equal), label='actual_output_y') @unittest.skipIf(tft is None, 'tft module is not installed.') def test_mltransform_with_counter(self): transforms = [ tft.ComputeAndApplyVocabulary(columns=['y']), tft.ScaleTo01(columns=['x']) ] data = [{'x': [1, 2, 3], 'y': ['a', 'b', 'c']}] with beam.Pipeline() as p: _ = ( p | beam.Create(data) | base.MLTransform( transforms=transforms, write_artifact_location=self.artifact_location)) scale_to_01_counter = MetricsFilter().with_name('BeamML_ScaleTo01') vocab_counter = MetricsFilter().with_name( 'BeamML_ComputeAndApplyVocabulary') mltransform_counter = MetricsFilter().with_name('BeamML_MLTransform') result = p.result self.assertEqual( result.metrics().query(scale_to_01_counter)['counters'][0].result, 1) self.assertEqual( result.metrics().query(vocab_counter)['counters'][0].result, 1) self.assertEqual( result.metrics().query(mltransform_counter)['counters'][0].result, 1) def test_non_ptransfrom_provider_class_to_mltransform(self): class Add: def __call__(self, x): return x + 1 with self.assertRaisesRegex(TypeError, 'transform must be a subclass of'): with beam.Pipeline() as p: _ = ( p | beam.Create([{ 'x': 1 }]) | base.MLTransform( write_artifact_location=self.artifact_location).with_transform( Add())) def test_read_mode_with_transforms(self): with self.assertRaises(ValueError): _ = base.MLTransform( # fake callable transforms=[lambda x: x], read_artifact_location=self.artifact_location) class FakeModel: def __call__(self, example: list[str]) -> list[str]: for i in range(len(example)): if not isinstance(example[i], str): raise TypeError('Input must be a string') example[i] = example[i][::-1] return example class FakeModelHandler(ModelHandler): def run_inference( self, batch: Sequence[str], model: Any, inference_args: Optional[dict[str, Any]] = None): return model(batch) def load_model(self): return FakeModel() class FakeEmbeddingsManager(base.EmbeddingsManager): def __init__(self, columns, **kwargs): super().__init__(columns=columns, **kwargs) def get_model_handler(self) -> ModelHandler: FakeModelHandler.__repr__ = lambda x: 'FakeEmbeddingsManager' # type: ignore[method-assign] return FakeModelHandler() def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: return (RunInference(model_handler=base._TextEmbeddingHandler(self))) def __repr__(self): return 'FakeEmbeddingsManager' class InvalidEmbeddingsManager(base.EmbeddingsManager): def __init__(self, **kwargs): super().__init__(**kwargs) def get_model_handler(self) -> ModelHandler: InvalidEmbeddingsManager.__repr__ = lambda x: 'InvalidEmbeddingsManager' # type: ignore[method-assign] return FakeModelHandler() def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: return (RunInference(model_handler=base._TextEmbeddingHandler(self))) def __repr__(self): return 'InvalidEmbeddingsManager' class TextEmbeddingHandlerTest(unittest.TestCase): def setUp(self) -> None: self.embedding_conig = FakeEmbeddingsManager(columns=['x']) self.artifact_location = tempfile.mkdtemp() def tearDown(self) -> None: shutil.rmtree(self.artifact_location) def test_no_columns_or_type_adapter(self): with self.assertRaises(ValueError): _ = InvalidEmbeddingsManager() def test_handler_with_incompatible_datatype(self): text_handler = base._TextEmbeddingHandler( embeddings_manager=self.embedding_conig) data = [ ('x', 1), ('x', 2), ('x', 3), ] with self.assertRaises(TypeError): text_handler.run_inference(data, None, None) def test_handler_with_dict_inputs(self): data = [ { 'x': "Hello world" }, { 'x': "Apache Beam" }, ] expected_data = [{ key: value[::-1] for key, value in d.items() } for d in data] with beam.Pipeline() as p: result = ( p | beam.Create(data) | base.MLTransform( write_artifact_location=self.artifact_location).with_transform( self.embedding_conig)) assert_that( result, equal_to(expected_data), ) def test_handler_with_batch_sizes(self): self.embedding_conig.max_batch_size = 100 self.embedding_conig.min_batch_size = 10 data = [ { 'x': "Hello world" }, { 'x': "Apache Beam" }, ] * 100 expected_data = [{ key: value[::-1] for key, value in d.items() } for d in data] with beam.Pipeline() as p: result = ( p | beam.Create(data) | base.MLTransform( write_artifact_location=self.artifact_location).with_transform( self.embedding_conig)) assert_that( result, equal_to(expected_data), ) def test_handler_on_multiple_columns(self): data = [ { 'x': "Hello world", 'y': "Apache Beam", 'z': 'unchanged' }, { 'x': "Apache Beam", 'y': "Hello world", 'z': 'unchanged' }, ] embedding_config = FakeEmbeddingsManager(columns=['x', 'y']) expected_data = [{ key: (value[::-1] if key in embedding_config.columns else value) for key, value in d.items() } for d in data] with beam.Pipeline() as p: result = ( p | beam.Create(data) | base.MLTransform(write_artifact_location=self.artifact_location). with_transform(embedding_config)) assert_that( result, equal_to(expected_data), ) def test_handler_on_columns_not_exist_in_input_data(self): data = [ { 'x': "Hello world", 'y': "Apache Beam" }, { 'x': "Apache Beam", 'y': "Hello world" }, ] embedding_config = FakeEmbeddingsManager(columns=['x', 'y', 'a']) with self.assertRaises(RuntimeError): with beam.Pipeline() as p: _ = ( p | beam.Create(data) | base.MLTransform(write_artifact_location=self.artifact_location). with_transform(embedding_config)) def test_handler_with_list_data(self): data = [{ 'x': ['Hello world', 'Apache Beam'], }, { 'x': ['Apache Beam', 'Hello world'], }] with self.assertRaisesRegex(Exception, "Embeddings can only be generated"): with beam.Pipeline() as p: _ = ( p | beam.Create(data) | base.MLTransform( write_artifact_location=self.artifact_location).with_transform( self.embedding_conig)) def test_handler_with_inconsistent_keys(self): data = [ { 'x': 'foo', 'y': 'bar', 'z': 'baz' }, { 'x': 'foo2', 'y': 'bar2' }, { 'x': 'foo3', 'y': 'bar3', 'z': 'baz3' }, ] self.embedding_conig.min_batch_size = 2 with self.assertRaises(RuntimeError): with beam.Pipeline() as p: _ = ( p | beam.Create(data) | base.MLTransform( write_artifact_location=self.artifact_location).with_transform( self.embedding_conig)) class FakeImageModel: def __call__(self, example: list[PIL_Image]) -> list[PIL_Image]: for i in range(len(example)): if not isinstance(example[i], PIL_Image): raise TypeError('Input must be an Image') return example class FakeImageModelHandler(ModelHandler): def run_inference( self, batch: Sequence[PIL_Image], model: Any, inference_args: Optional[dict[str, Any]] = None): return model(batch) def load_model(self): return FakeImageModel() class FakeImageEmbeddingsManager(base.EmbeddingsManager): def __init__(self, columns, **kwargs): super().__init__(columns=columns, **kwargs) def get_model_handler(self) -> ModelHandler: FakeModelHandler.__repr__ = lambda x: 'FakeImageEmbeddingsManager' # type: ignore[method-assign] return FakeImageModelHandler() def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: return (RunInference(model_handler=base._ImageEmbeddingHandler(self))) def __repr__(self): return 'FakeImageEmbeddingsManager' class TestImageEmbeddingHandler(unittest.TestCase): def setUp(self) -> None: self.embedding_config = FakeImageEmbeddingsManager(columns=['x']) self.artifact_location = tempfile.mkdtemp() def tearDown(self) -> None: shutil.rmtree(self.artifact_location) @unittest.skipIf(PIL is None, 'PIL module is not installed.') def test_handler_with_non_dict_datatype(self): image_handler = base._ImageEmbeddingHandler( embeddings_manager=self.embedding_config) data = [ ('x', 'hi there'), ('x', 'not an image'), ('x', 'image_path.jpg'), ] with self.assertRaises(TypeError): image_handler.run_inference(data, None, None) @unittest.skipIf(PIL is None, 'PIL module is not installed.') def test_handler_with_non_image_datatype(self): image_handler = base._ImageEmbeddingHandler( embeddings_manager=self.embedding_config) data = [ { 'x': 'hi there' }, { 'x': 'not an image' }, { 'x': 'image_path.jpg' }, ] with self.assertRaises(TypeError): image_handler.run_inference(data, None, None) @unittest.skipIf(PIL is None, 'PIL module is not installed.') def test_handler_with_dict_inputs(self): img_one = PIL.Image.new(mode='RGB', size=(1, 1)) img_two = PIL.Image.new(mode='RGB', size=(1, 1)) data = [ { 'x': img_one }, { 'x': img_two }, ] expected_data = [{key: value for key, value in d.items()} for d in data] with beam.Pipeline() as p: result = ( p | beam.Create(data) | base.MLTransform( write_artifact_location=self.artifact_location).with_transform( self.embedding_config)) assert_that( result, equal_to(expected_data), ) @dataclass class FakeMultiModalInput: image: Optional[PIL_Image] = None video: Optional[Any] = None text: Optional[str] = None class FakeMultiModalModel: def __call__(self, example: list[FakeMultiModalInput]) -> list[FakeMultiModalInput]: for i in range(len(example)): if not isinstance(example[i], FakeMultiModalInput): raise TypeError('Input must be a MultiModalInput') return example class FakeMultiModalModelHandler(ModelHandler): def run_inference( self, batch: Sequence[FakeMultiModalInput], model: Any, inference_args: Optional[dict[str, Any]] = None): return model(batch) def load_model(self): return FakeMultiModalModel() class FakeMultiModalEmbeddingsManager(base.EmbeddingsManager): def __init__(self, columns, **kwargs): super().__init__(columns=columns, **kwargs) def get_model_handler(self) -> ModelHandler: FakeModelHandler.__repr__ = lambda x: 'FakeMultiModalEmbeddingsManager' # type: ignore[method-assign] return FakeMultiModalModelHandler() def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: return (RunInference(model_handler=base._MultiModalEmbeddingHandler(self))) def __repr__(self): return 'FakeMultiModalEmbeddingsManager' class TestMultiModalEmbeddingHandler(unittest.TestCase): def setUp(self) -> None: self.embedding_config = FakeMultiModalEmbeddingsManager(columns=['x']) self.artifact_location = tempfile.mkdtemp() def tearDown(self) -> None: shutil.rmtree(self.artifact_location) @unittest.skipIf(PIL is None, 'PIL module is not installed.') def test_handler_with_non_dict_datatype(self): image_handler = base._MultiModalEmbeddingHandler( embeddings_manager=self.embedding_config) data = [ ('x', 'hi there'), ('x', 'not an image'), ('x', 'image_path.jpg'), ] with self.assertRaises(TypeError): image_handler.run_inference(data, None, None) @unittest.skipIf(PIL is None, 'PIL module is not installed.') def test_handler_with_incorrect_datatype(self): image_handler = base._MultiModalEmbeddingHandler( embeddings_manager=self.embedding_config) data = [ { 'x': 'hi there' }, { 'x': 'not an image' }, { 'x': 'image_path.jpg' }, ] with self.assertRaises(TypeError): image_handler.run_inference(data, None, None) @unittest.skipIf(PIL is None, 'PIL module is not installed.') def test_handler_with_dict_inputs(self): input_one = FakeMultiModalInput( image=PIL.Image.new(mode='RGB', size=(1, 1)), text="test image one") input_two = FakeMultiModalInput( image=PIL.Image.new(mode='RGB', size=(1, 1)), text="test image two") input_three = FakeMultiModalInput( image=PIL.Image.new(mode='RGB', size=(1, 1)), video=bytes.fromhex('2Ef0 F1f2 '), text="test image three with video") data = [ { 'x': input_one }, { 'x': input_two }, { 'x': input_three }, ] expected_data = [{key: value for key, value in d.items()} for d in data] with beam.Pipeline() as p: result = ( p | beam.Create(data) | base.MLTransform( write_artifact_location=self.artifact_location).with_transform( self.embedding_config)) assert_that( result, equal_to(expected_data), ) class TestUtilFunctions(unittest.TestCase): def test_dict_input_fn_normal(self): input_list = [{'a': 1, 'b': 2}, {'a': 3, 'b': 4}] columns = ['a', 'b'] expected_output = [1, 2, 3, 4] self.assertEqual(base._dict_input_fn(columns, input_list), expected_output) def test_dict_output_fn_normal(self): input_list = [{'a': 1, 'b': 2}, {'a': 3, 'b': 4}] columns = ['a', 'b'] embeddings = [1.1, 2.2, 3.3, 4.4] expected_output = [{'a': 1.1, 'b': 2.2}, {'a': 3.3, 'b': 4.4}] self.assertEqual( base._dict_output_fn(columns, input_list, embeddings), expected_output) def test_dict_input_fn_on_list_inputs(self): input_list = [{'a': [1, 2, 10], 'b': 3}, {'a': [1], 'b': 5}] columns = ['a', 'b'] expected_output = [[1, 2, 10], 3, [1], 5] self.assertEqual(base._dict_input_fn(columns, input_list), expected_output) def test_dict_output_fn_on_list_inputs(self): input_list = [{'a': [1, 2, 10], 'b': 3}, {'a': [1], 'b': 5}] columns = ['a', 'b'] embeddings = [1.1, 2.2, 3.3, 4.4] expected_output = [{'a': 1.1, 'b': 2.2}, {'a': 3.3, 'b': 4.4}] self.assertEqual( base._dict_output_fn(columns, input_list, embeddings), expected_output) class TestJsonPickleTransformAttributeManager(unittest.TestCase): def setUp(self): self.attribute_manager = base._transform_attribute_manager self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) @unittest.skipIf(tft is None, 'tft module is not installed.') def test_save_tft_process_handler(self): transforms = [ tft.ScaleTo01(columns=['x']), tft.ComputeAndApplyVocabulary(columns=['y']) ] process_handler = TFTProcessHandler( transforms=transforms, artifact_location=self.artifact_location, ) self.attribute_manager.save_attributes( ptransform_list=[process_handler], artifact_location=self.artifact_location, ) files = os.listdir(self.artifact_location) self.assertTrue(len(files) == 1) self.assertTrue(files[0] == base._ATTRIBUTE_FILE_NAME) def test_save_run_inference(self): self.attribute_manager.save_attributes( ptransform_list=[RunInference(model_handler=FakeModelHandler())], artifact_location=self.artifact_location, ) files = os.listdir(self.artifact_location) self.assertTrue(len(files) == 1) self.assertTrue(files[0] == base._ATTRIBUTE_FILE_NAME) def test_save_and_load_run_inference(self): ptransform_list = [RunInference(model_handler=FakeModelHandler())] self.attribute_manager.save_attributes( ptransform_list=ptransform_list, artifact_location=self.artifact_location, ) loaded_ptransform_list = self.attribute_manager.load_attributes( artifact_location=self.artifact_location, ) self.assertTrue(len(loaded_ptransform_list) == len(ptransform_list)) self.assertListEqual( list(loaded_ptransform_list[0].__dict__.keys()), list(ptransform_list[0].__dict__.keys())) get_keys = lambda x: list(x.__dict__.keys()) for i, transform in enumerate(ptransform_list): self.assertListEqual( get_keys(transform), get_keys(loaded_ptransform_list[i])) if hasattr(transform, 'model_handler'): model_handler = transform.model_handler loaded_model_handler = loaded_ptransform_list[i].model_handler self.assertListEqual( get_keys(model_handler), get_keys(loaded_model_handler)) def test_mltransform_to_ptransform_wrapper(self): transforms = [ FakeEmbeddingsManager(columns=['x']), FakeEmbeddingsManager(columns=['y', 'z']), ] ptransform_mapper = base._MLTransformToPTransformMapper( transforms=transforms, artifact_location=self.artifact_location, artifact_mode=None) ptransform_list = ptransform_mapper.create_ptransform_list() self.assertTrue(len(ptransform_list) == 2) self.assertEqual(type(ptransform_list[0]), RunInference) expected_columns = [['x'], ['y', 'z']] for i in range(len(ptransform_list)): self.assertEqual(type(ptransform_list[i]), RunInference) self.assertEqual( type(ptransform_list[i]._model_handler), base._TextEmbeddingHandler) self.assertEqual( ptransform_list[i]._model_handler.columns, expected_columns[i]) @unittest.skipIf(apiclient is None, 'apache_beam[gcp] is not installed.') def test_with_gcs_location_with_none_options(self): path = f"gs://fake_path_{secrets.token_hex(3)}_{int(time.time())}" with self.assertRaises(RuntimeError): self.attribute_manager.save_attributes( ptransform_list=[], artifact_location=path, options=None) with self.assertRaises(RuntimeError): self.attribute_manager.save_attributes( ptransform_list=[], artifact_location=path) def test_with_same_local_artifact_location(self): artifact_location = self.artifact_location attribute_manager = base._JsonPickleTransformAttributeManager() ptransform_list = [RunInference(model_handler=FakeModelHandler())] attribute_manager.save_attributes( ptransform_list, artifact_location=artifact_location) with self.assertRaises(FileExistsError): attribute_manager.save_attributes([lambda x: x], artifact_location=artifact_location) class MLTransformDLQTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) def test_dlq_with_embeddings(self): with beam.Pipeline() as p: good, bad = ( p | beam.Create([{ 'x': 1 }, { 'x': 3, }, { 'x': 'Hello' } ], ) | base.MLTransform( write_artifact_location=self.artifact_location).with_transform( FakeEmbeddingsManager( columns=['x'])).with_exception_handling()) good_expected_elements = [{'x': 'olleH'}] assert_that( good, equal_to(good_expected_elements), label='good', ) # batching happens in RunInference hence elements # are in lists in the bad pcoll. bad_expected_elements = [[{'x': 1}], [{'x': 3}]] assert_that( bad | beam.Map(lambda x: x.element), equal_to(bad_expected_elements), label='bad', ) def test_mltransform_with_dlq_and_extract_tranform_name(self): with beam.Pipeline() as p: good, bad = ( p | beam.Create([{ 'x': 1 }, { 'x': 3, }, { 'x': 'Hello' } ], ) | base.MLTransform( write_artifact_location=self.artifact_location).with_transform( FakeEmbeddingsManager( columns=['x'])).with_exception_handling()) good_expected_elements = [{'x': 'olleH'}] assert_that( good, equal_to(good_expected_elements), label='good', ) bad_expected_transform_name = [ 'FakeEmbeddingsManager', 'FakeEmbeddingsManager' ] assert_that( bad | beam.Map(lambda x: x.transform_name), equal_to(bad_expected_transform_name), ) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5214226 apache_beam-2.74.0/apache_beam/ml/transforms/embeddings/0000755000175100017510000000000015206113160022566 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/embeddings/__init__.py0000644000175100017510000000166415206113032024704 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # TODO: Add dead letter queue for RunInference transforms. """ This module contains embedding configs that can be used to generate embeddings using MLTransform. """ ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/embeddings/huggingface.py0000644000175100017510000002220715206113032025410 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. __all__ = ["SentenceTransformerEmbeddings", "InferenceAPIEmbeddings"] import logging import os from collections.abc import Callable from collections.abc import Mapping from collections.abc import Sequence from typing import Any from typing import Optional import requests import apache_beam as beam from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import RunInference from apache_beam.ml.transforms.base import EmbeddingsManager from apache_beam.ml.transforms.base import _ImageEmbeddingHandler from apache_beam.ml.transforms.base import _TextEmbeddingHandler try: from sentence_transformers import SentenceTransformer except ImportError: SentenceTransformer = None _LOGGER = logging.getLogger(__name__) # TODO: https://github.com/apache/beam/issues/29621 # Use HuggingFaceModelHandlerTensor once the import issue is fixed. # Right now, the hugging face model handler import torch and tensorflow # at the same time, which adds too much weigth to the container unnecessarily. class _SentenceTransformerModelHandler(ModelHandler): """ Note: Intended for internal use and guarantees no backwards compatibility. """ def __init__( self, model_name: str, model_class: Callable, load_model_args: Optional[dict] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_seq_length: Optional[int] = None, large_model: bool = False, **kwargs): self._max_seq_length = max_seq_length self.model_name = model_name self._model_class = model_class self._load_model_args = load_model_args self._min_batch_size = min_batch_size self._max_batch_size = max_batch_size self._large_model = large_model self._kwargs = kwargs if not SentenceTransformer: raise ImportError( "sentence-transformers is required to use " "SentenceTransformerEmbeddings." "Please install it with using `pip install sentence-transformers`.") def run_inference( self, batch: Sequence[str], model: SentenceTransformer, inference_args: Optional[dict[str, Any]] = None, ): inference_args = inference_args or {} return model.encode(batch, **inference_args) def load_model(self): model = self._model_class(self.model_name, **self._load_model_args) if self._max_seq_length: model.max_seq_length = self._max_seq_length return model def share_model_across_processes(self) -> bool: return self._large_model def batch_elements_kwargs(self) -> Mapping[str, Any]: batch_sizes = {} if self._min_batch_size: batch_sizes["min_batch_size"] = self._min_batch_size if self._max_batch_size: batch_sizes["max_batch_size"] = self._max_batch_size return batch_sizes def __repr__(self) -> str: # ModelHandler is internal to the user and is not exposed. # Hence we need to override the __repr__ method to expose # the name of the class. return 'SentenceTransformerEmbeddings' class SentenceTransformerEmbeddings(EmbeddingsManager): def __init__( self, model_name: str, columns: list[str], max_seq_length: Optional[int] = None, image_model: bool = False, **kwargs): """ Embedding config for sentence-transformers. This config can be used with MLTransform to embed text data. Models are loaded using the RunInference PTransform with the help of ModelHandler. Args: model_name: Name of the model to use. The model should be hosted on HuggingFace Hub or compatible with sentence_transformers. For image embedding models, see https://www.sbert.net/docs/sentence_transformer/pretrained_models.html#image-text-models # pylint: disable=line-too-long for a list of available sentence_transformers models. columns: List of columns to be embedded. max_seq_length: Max sequence length to use for the model if applicable. image_model: Whether the model is generating image embeddings. min_batch_size: The minimum batch size to be used for inference. max_batch_size: The maximum batch size to be used for inference. large_model: Whether to share the model across processes. """ super().__init__(columns=columns, **kwargs) self.model_name = model_name self.max_seq_length = max_seq_length self.image_model = image_model def get_model_handler(self): return _SentenceTransformerModelHandler( model_class=SentenceTransformer, max_seq_length=self.max_seq_length, model_name=self.model_name, load_model_args=self.load_model_args, min_batch_size=self.min_batch_size, max_batch_size=self.max_batch_size, large_model=self.large_model) def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: # wrap the model handler in an appropriate embedding handler to provide # some type checking. if self.image_model: return ( RunInference( model_handler=_ImageEmbeddingHandler(self), inference_args=self.inference_args, )) return ( RunInference( model_handler=_TextEmbeddingHandler(self), inference_args=self.inference_args, )) class _InferenceAPIHandler(ModelHandler): def __init__(self, config: 'InferenceAPIEmbeddings'): super().__init__() self._config = config def load_model(self): session = requests.Session() # if the token is not provided during construction time, it might have # been provided with custom container, which we can get it during runtume. if not self._config.hf_token: hf_token = os.environ.get("HF_TOKEN") if not hf_token: raise ValueError( 'HF_TOKEN environment variable not set. ' 'Please set the environment variable or pass the token as an ' 'argument.') session.headers.update({"Authorization": f"Bearer {hf_token}"}) return session session.headers.update(self._config.authorization_token) return session def run_inference( self, batch, session: requests.Session, inference_args=None): response = session.post( self._config.api_url, headers=self._config.authorization_token, json={ "inputs": batch, "options": inference_args }) return response.json() class InferenceAPIEmbeddings(EmbeddingsManager): """ Feature extraction using HuggingFace's Inference API. Intended to be used for feature-extraction. For other tasks, please refer to https://huggingface.co/inference-api. Args: hf_token: HuggingFace token. columns: List of columns to be embedded. model_name: Model name used for feature extraction. api_url: API url for feature extraction. If specified, model_name will be ignored. If none, the default url for feature extraction will be used. """ def __init__( self, hf_token: Optional[str], columns: list[str], model_name: Optional[str] = None, # example: "sentence-transformers/all-MiniLM-l6-v2" # pylint: disable=line-too-long api_url: Optional[str] = None, **kwargs, ): super().__init__(columns=columns, **kwargs) self._authorization_token = {"Authorization": f"Bearer {hf_token}"} self._model_name = model_name self.hf_token = hf_token if not api_url: if not self._model_name: raise ValueError("Either api_url or model_name must be provided.") self._api_url = ( f"https://router.huggingface.co/hf-inference/models/{self._model_name}/pipeline/feature-extraction" # pylint: disable=line-too-long ) else: self._api_url = api_url _LOGGER.info("HuggingFace API URL: %s") def get_token(self): return os.environ.get('HF_TOKEN') @property def api_url(self): return self._api_url @property def authorization_token(self): return self._authorization_token def get_model_handler(self): return _InferenceAPIHandler(self) def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: options = { # sometimes the model is not ready and returns an error response # instead of waiting. So we wait for the model to be ready. 'wait_for_model': True, 'use_cache': True, } self.inference_args.update(options) return ( RunInference( model_handler=_TextEmbeddingHandler(self), inference_args=self.inference_args)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/embeddings/huggingface_test.py0000644000175100017510000003271515206113032026454 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os import shutil import tempfile import unittest import uuid import numpy as np import pytest from parameterized import parameterized import apache_beam as beam from apache_beam.ml.inference.base import RunInference from apache_beam.ml.transforms import base from apache_beam.ml.transforms.base import MLTransform from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # pylint: disable=ungrouped-imports try: import torch from PIL import Image from apache_beam.ml.transforms.embeddings.huggingface import InferenceAPIEmbeddings from apache_beam.ml.transforms.embeddings.huggingface import SentenceTransformerEmbeddings except ImportError: SentenceTransformerEmbeddings = None # type: ignore # pylint: disable=ungrouped-imports try: import tensorflow_transform as tft from apache_beam.ml.transforms.tft import ScaleTo01 except ImportError: tft = None # pylint: disable=ungrouped-imports try: from PIL import Image except ImportError: Image = None # type: ignore[assignment] _HF_TOKEN = os.environ.get('HF_INFERENCE_TOKEN') test_query = "This is a test" test_query_column = "feature_1" DEFAULT_MODEL_NAME = "sentence-transformers/all-mpnet-base-v2" IMAGE_MODEL_NAME = "clip-ViT-B-32" _parameterized_inputs = [ ([{ test_query_column: 'That is a happy person' }, { test_query_column: 'That is a very happy person' }], 'thenlper/gte-base', [0.11, 0.11]), ([{ test_query_column: test_query, }], DEFAULT_MODEL_NAME, [0.13]), ( [{ test_query_column: 'query: how much protein should a female eat', }, { test_query_column: ( "passage: As a general guideline, the CDC's " "average requirement of protein for women " "ages 19 to 70 is 46 grams per day. But, " "as you can see from this chart, you'll need " "to increase that if you're expecting or training" " for a marathon. Check out the chart below " "to see how much protein " "you should be eating each day.") }], 'intfloat/e5-base-v2', # this model requires inputs to be specified as query: and passage: [0.1, 0.1]), ] @pytest.mark.no_xdist @unittest.skipIf( SentenceTransformerEmbeddings is None, 'sentence-transformers is not installed.') class SentenceTransformerEmbeddingsTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp(prefix='sentence_transformers_') # this bucket has TTL and will be deleted periodically self.gcs_artifact_location = os.path.join( 'gs://temp-storage-for-perf-tests/sentence_transformers', uuid.uuid4().hex) def tearDown(self) -> None: shutil.rmtree(self.artifact_location) def test_sentence_transformer_embeddings(self): model_name = DEFAULT_MODEL_NAME embedding_config = SentenceTransformerEmbeddings( model_name=model_name, columns=[test_query_column]) with beam.Pipeline() as pipeline: result_pcoll = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: test_query }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) def assert_element(element): assert len(element[test_query_column]) == 768 _ = (result_pcoll | beam.Map(assert_element)) @unittest.skipIf(tft is None, 'Tensorflow Transform is not installed.') def test_embeddings_with_scale_to_0_1(self): model_name = DEFAULT_MODEL_NAME embedding_config = SentenceTransformerEmbeddings( model_name=model_name, columns=[test_query_column], ) with beam.Pipeline() as pipeline: transformed_pcoll = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: test_query }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config).with_transform( ScaleTo01(columns=[test_query_column]))) def assert_element(element): assert max(element.feature_1) == 1 _ = (transformed_pcoll | beam.Map(assert_element)) @parameterized.expand(_parameterized_inputs) def test_embeddings_with_read_artifact_location( self, inputs, model_name, output): embedding_config = SentenceTransformerEmbeddings( model_name=model_name, columns=[test_query_column]) with beam.Pipeline() as p: result_pcoll = ( p | "CreateData" >> beam.Create(inputs) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) max_ele_pcoll = ( result_pcoll | beam.Map(lambda x: round(max(x[test_query_column]), 2))) assert_that(max_ele_pcoll, equal_to(output)) with beam.Pipeline() as p: result_pcoll = ( p | "CreateData" >> beam.Create(inputs) | "MLTransform" >> MLTransform(read_artifact_location=self.artifact_location)) max_ele_pcoll = ( result_pcoll | beam.Map(lambda x: round(max(x[test_query_column]), 2))) assert_that(max_ele_pcoll, equal_to(output)) def test_sentence_transformer_with_int_data_types(self): model_name = DEFAULT_MODEL_NAME embedding_config = SentenceTransformerEmbeddings( model_name=model_name, columns=[test_query_column]) with self.assertRaisesRegex(Exception, "Embeddings can only be generated"): with beam.Pipeline() as pipeline: _ = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: 1 }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) @parameterized.expand(_parameterized_inputs) def test_with_gcs_artifact_location(self, inputs, model_name, output): embedding_config = SentenceTransformerEmbeddings( model_name=model_name, columns=[test_query_column]) with beam.Pipeline() as p: result_pcoll = ( p | "CreateData" >> beam.Create(inputs) | "MLTransform" >> MLTransform(write_artifact_location=self.gcs_artifact_location ).with_transform(embedding_config)) max_ele_pcoll = ( result_pcoll | beam.Map(lambda x: round(np.max(x[test_query_column]), 2))) assert_that(max_ele_pcoll, equal_to(output)) with beam.Pipeline() as p: result_pcoll = ( p | "CreateData" >> beam.Create(inputs) | "MLTransform" >> MLTransform(read_artifact_location=self.gcs_artifact_location)) max_ele_pcoll = ( result_pcoll | beam.Map(lambda x: round(np.max(x[test_query_column]), 2))) assert_that(max_ele_pcoll, equal_to(output)) def test_embeddings_with_inference_args(self): model_name = DEFAULT_MODEL_NAME inference_args = {'convert_to_numpy': False} embedding_config = SentenceTransformerEmbeddings( model_name=model_name, columns=[test_query_column], inference_args=inference_args) with beam.Pipeline() as pipeline: result_pcoll = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: test_query }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) def assert_element(element): assert type(element) == torch.Tensor _ = ( result_pcoll | beam.Map(lambda x: x[test_query_column]) | beam.Map(assert_element)) def test_mltransform_to_ptransform_with_sentence_transformer(self): model_name = '' transforms = [ SentenceTransformerEmbeddings(columns=['x'], model_name=model_name), SentenceTransformerEmbeddings( columns=['y', 'z'], model_name=model_name) ] ptransform_mapper = base._MLTransformToPTransformMapper( transforms=transforms, artifact_location=self.artifact_location, artifact_mode=None) ptransform_list = ptransform_mapper.create_and_save_ptransform_list() self.assertTrue(len(ptransform_list) == 2) self.assertEqual(type(ptransform_list[0]), RunInference) expected_columns = [['x'], ['y', 'z']] for i in range(len(ptransform_list)): self.assertEqual(type(ptransform_list[i]), RunInference) self.assertEqual( type(ptransform_list[i]._model_handler), base._TextEmbeddingHandler) self.assertEqual( ptransform_list[i]._model_handler.columns, expected_columns[i]) self.assertEqual( ptransform_list[i]._model_handler._underlying.model_name, model_name) ptransform_list = ( base._MLTransformToPTransformMapper. load_transforms_from_artifact_location(self.artifact_location)) for i in range(len(ptransform_list)): self.assertEqual(type(ptransform_list[i]), RunInference) self.assertEqual( type(ptransform_list[i]._model_handler), base._TextEmbeddingHandler) self.assertEqual( ptransform_list[i]._model_handler.columns, expected_columns[i]) self.assertEqual( ptransform_list[i]._model_handler._underlying.model_name, model_name) def generateRandomImage(self, size: int): imarray = np.random.rand(size, size, 3) * 255 return Image.fromarray(imarray.astype('uint8')).convert('RGBA') @unittest.skipIf(Image is None, 'Pillow is not installed.') def test_sentence_transformer_image_embeddings(self): embedding_config = SentenceTransformerEmbeddings( model_name=IMAGE_MODEL_NAME, columns=[test_query_column], image_model=True) img = self.generateRandomImage(256) with beam.Pipeline() as pipeline: result_pcoll = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: img }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) def assert_element(element): assert len(element[test_query_column]) == 512 _ = (result_pcoll | beam.Map(assert_element)) def test_sentence_transformer_images_with_str_data_types(self): embedding_config = SentenceTransformerEmbeddings( model_name=IMAGE_MODEL_NAME, columns=[test_query_column], image_model=True) with self.assertRaisesRegex(Exception, "Embeddings can only be generated"): with beam.Pipeline() as pipeline: _ = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: "image.jpg" }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) @unittest.skipIf(not _HF_TOKEN, 'HF_TOKEN environment variable not set.') class HuggingfaceInferenceAPITest(unittest.TestCase): def setUp(self): self.artifact_location = tempfile.mkdtemp() self.inputs = [{test_query_column: test_query}] def tearDown(self): shutil.rmtree(self.artifact_location) def test_get_api_url_and_when_model_name_not_provided(self): with self.assertRaises(ValueError): inference_embeddings = InferenceAPIEmbeddings( hf_token=_HF_TOKEN, columns=[test_query_column], ) _ = inference_embeddings.api_url def test_embeddings_with_inference_api(self): embedding_config = InferenceAPIEmbeddings( hf_token=_HF_TOKEN, model_name=DEFAULT_MODEL_NAME, columns=[test_query_column], ) expected_output = [0.13] with beam.Pipeline() as p: result_pcoll = ( p | "CreateData" >> beam.Create(self.inputs) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) max_ele_pcoll = ( result_pcoll | beam.Map(lambda x: round(np.max(x[test_query_column]), 2))) assert_that(max_ele_pcoll, equal_to(expected_output)) @unittest.skipIf(not _HF_TOKEN, 'HF_TOKEN environment variable not set.') class HuggingfaceInferenceAPIGCSLocationTest(HuggingfaceInferenceAPITest): def setUp(self): self.artifact_location = self.gcs_artifact_location = os.path.join( 'gs://temp-storage-for-perf-tests/tft_handler', uuid.uuid4().hex) self.inputs = [{test_query_column: test_query}] def tearDown(self): pass if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/embeddings/open_ai.py0000644000175100017510000001325015206113032024551 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import logging from collections.abc import Iterable from collections.abc import Sequence from typing import Any from typing import Optional from typing import TypeVar from typing import Union import openai from openai import APIError from openai import RateLimitError import apache_beam as beam from apache_beam.ml.inference.base import RemoteModelHandler from apache_beam.ml.inference.base import RunInference from apache_beam.ml.transforms.base import EmbeddingsManager from apache_beam.ml.transforms.base import _TextEmbeddingHandler from apache_beam.pvalue import PCollection from apache_beam.pvalue import Row __all__ = ["OpenAITextEmbeddings"] # Define a type variable for the output MLTransformOutputT = TypeVar('MLTransformOutputT') # Default batch size for OpenAI API requests _DEFAULT_BATCH_SIZE = 20 LOGGER = logging.getLogger("OpenAIEmbeddings") def _retry_on_appropriate_openai_error(exception): """ Retry filter that returns True for rate limit (429) or server (5xx) errors. Args: exception: the returned exception encountered during the request/response loop. Returns: boolean indication whether or not the exception is a Server Error (5xx) or a RateLimitError (429) error. """ return isinstance(exception, (RateLimitError, APIError)) class _OpenAITextEmbeddingHandler(RemoteModelHandler): """ Note: Intended for internal use and guarantees no backwards compatibility. """ def __init__( self, model_name: str, api_key: Optional[str] = None, organization: Optional[str] = None, dimensions: Optional[int] = None, user: Optional[str] = None, max_batch_size: Optional[int] = None, ): super().__init__( namespace="OpenAITextEmbeddings", num_retries=5, throttle_delay_secs=5, retry_filter=_retry_on_appropriate_openai_error) self.model_name = model_name self.api_key = api_key self.organization = organization self.dimensions = dimensions self.user = user self.max_batch_size = max_batch_size or _DEFAULT_BATCH_SIZE def create_client(self): """Creates and returns an OpenAI client.""" if self.api_key: client = openai.OpenAI( api_key=self.api_key, organization=self.organization, ) else: client = openai.OpenAI(organization=self.organization) return client def request( self, batch: Sequence[str], model: Any, inference_args: Optional[dict[str, Any]] = None, ) -> Iterable: """Makes a request to OpenAI embedding API and returns embeddings.""" # Prepare arguments for the API call kwargs = { "model": self.model_name, "input": batch, } if self.dimensions: kwargs["dimensions"] = [str(self.dimensions)] if self.user: kwargs["user"] = self.user # Make the API call - let RemoteModelHandler handle retries and exceptions response = model.embeddings.create(**kwargs) return [item.embedding for item in response.data] def batch_elements_kwargs(self) -> dict[str, Any]: """Return kwargs suitable for BatchElements with appropriate batch size""" return {'max_batch_size': self.max_batch_size} def __repr__(self): return 'OpenAITextEmbeddings' class OpenAITextEmbeddings(EmbeddingsManager): @beam.typehints.with_output_types(PCollection[Union[MLTransformOutputT, Row]]) def __init__( self, model_name: str, columns: list[str], api_key: Optional[str] = None, organization: Optional[str] = None, dimensions: Optional[int] = None, user: Optional[str] = None, max_batch_size: Optional[int] = None, **kwargs): """ Embedding Config for OpenAI Text Embedding models. Text Embeddings are generated for a batch of text using the OpenAI API. Args: model_name: Name of the OpenAI embedding model columns: The columns where the embeddings will be stored in the output api_key: OpenAI API key organization: OpenAI organization ID dimensions: Specific embedding dimensions to use (if model supports it) user: End-user identifier for tracking and rate limit calculations max_batch_size: Maximum batch size for requests to OpenAI API """ self.model_name = model_name self.api_key = api_key self.organization = organization self.dimensions = dimensions self.user = user self.max_batch_size = max_batch_size super().__init__(columns=columns, **kwargs) def get_model_handler(self) -> RemoteModelHandler: return _OpenAITextEmbeddingHandler( model_name=self.model_name, api_key=self.api_key, organization=self.organization, dimensions=self.dimensions, user=self.user, max_batch_size=self.max_batch_size, ) def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: return RunInference( model_handler=_TextEmbeddingHandler(self), inference_args=self.inference_args) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/embeddings/open_ai_it_test.py0000644000175100017510000002355115206113032026311 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os import shutil import tempfile import unittest import apache_beam as beam from apache_beam.ml.inference.base import RunInference from apache_beam.ml.transforms import base from apache_beam.ml.transforms.base import MLTransform try: from sdks.python.apache_beam.ml.transforms.embeddings.open_ai import OpenAITextEmbeddings except ImportError: OpenAITextEmbeddings = None # pylint: disable=ungrouped-imports try: import tensorflow_transform as tft except ImportError: tft = None test_query = "This is a test" test_query_column = "embedding" model_name: str = "text-embedding-3-small" @unittest.skipIf( OpenAITextEmbeddings is None, 'OpenAI Python SDK is not installed.') class OpenAIEmbeddingsTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp(prefix='_openai_test') self.api_key = os.environ.get('OPENAI_API_KEY') def tearDown(self) -> None: shutil.rmtree(self.artifact_location) def test_openai_text_embeddings(self): embedding_config = OpenAITextEmbeddings( model_name=model_name, columns=["embedding"], api_key=self.api_key) with beam.Pipeline() as pipeline: transformed_pcoll = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: test_query }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) def assert_element(element): # OpenAI text-embedding-3-small produces 1536-dimensional embeddings assert len(element[test_query_column]) == 1536 _ = (transformed_pcoll | beam.Map(assert_element)) @unittest.skipIf(tft is None, 'Tensorflow Transform is not installed.') def test_embeddings_with_scale_to_0_1(self): embedding_config = OpenAITextEmbeddings( model_name=model_name, columns=[test_query_column], api_key=self.api_key, ) with beam.Pipeline() as pipeline: transformed_pcoll = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: test_query }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) def assert_element(element): assert max(element.feature_1) == 1 _ = (transformed_pcoll | beam.Map(assert_element)) def test_embeddings_with_dimensions(self): embedding_config = OpenAITextEmbeddings( model_name=model_name, columns=[test_query_column], api_key=self.api_key, dimensions=512) with beam.Pipeline() as pipeline: transformed_pcoll = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: test_query }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) def assert_element(element): # Check that we get 512-dimensional embeddings as requested assert len(element[test_query_column]) == 512 _ = (transformed_pcoll | beam.Map(assert_element)) def pipeline_with_configurable_artifact_location( self, pipeline, embedding_config=None, read_artifact_location=None, write_artifact_location=None): if write_artifact_location: return ( pipeline | MLTransform(write_artifact_location=write_artifact_location). with_transform(embedding_config)) elif read_artifact_location: return ( pipeline | MLTransform(read_artifact_location=read_artifact_location)) else: raise NotImplementedError def test_embeddings_with_read_artifact_location(self): embedding_config = OpenAITextEmbeddings( model_name=model_name, columns=[test_query_column], api_key=self.api_key) with beam.Pipeline() as p: data = ( p | "CreateData" >> beam.Create([{ test_query_column: test_query }])) _ = self.pipeline_with_configurable_artifact_location( pipeline=data, embedding_config=embedding_config, write_artifact_location=self.artifact_location) with beam.Pipeline() as p: data = ( p | "CreateData" >> beam.Create([{ test_query_column: test_query }, { test_query_column: test_query }])) result_pcoll = self.pipeline_with_configurable_artifact_location( pipeline=data, read_artifact_location=self.artifact_location) # Since we don't know the exact values of the embeddings, # we just check that they are within a reasonable range def assert_element(element): # Embeddings should be normalized and generally small values assert -1 <= element <= 1 _ = ( result_pcoll | beam.Map(lambda x: max(x[test_query_column])) | beam.Map(assert_element)) def test_with_int_data_types(self): embedding_config = OpenAITextEmbeddings( model_name=model_name, columns=[test_query_column], api_key=self.api_key) with self.assertRaises(Exception): with beam.Pipeline() as pipeline: _ = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: 1 }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) def test_with_artifact_location(self): # pylint: disable=line-too-long """Local artifact location test""" secondary_artifact_location = tempfile.mkdtemp( prefix='_openai_secondary_test') try: embedding_config = OpenAITextEmbeddings( model_name=model_name, columns=[test_query_column], api_key=self.api_key) with beam.Pipeline() as p: data = ( p | "CreateData" >> beam.Create([{ test_query_column: test_query }])) _ = self.pipeline_with_configurable_artifact_location( pipeline=data, embedding_config=embedding_config, write_artifact_location=secondary_artifact_location) with beam.Pipeline() as p: data = ( p | "CreateData" >> beam.Create([{ test_query_column: test_query }, { test_query_column: test_query }])) result_pcoll = self.pipeline_with_configurable_artifact_location( pipeline=data, read_artifact_location=secondary_artifact_location) def assert_element(element): # Embeddings should be normalized and generally small values assert -1 <= element <= 1 _ = ( result_pcoll | beam.Map(lambda x: max(x[test_query_column])) | beam.Map(assert_element)) finally: # Clean up the temporary directory shutil.rmtree(secondary_artifact_location) def test_mltransform_to_ptransform_with_openai(self): # pylint: disable=line-too-long transforms = [ OpenAITextEmbeddings( columns=['x'], model_name=model_name, api_key=self.api_key, dimensions=512), OpenAITextEmbeddings( columns=['y', 'z'], model_name=model_name, api_key=self.api_key) ] ptransform_mapper = base._MLTransformToPTransformMapper( transforms=transforms, artifact_location=self.artifact_location, artifact_mode=None) ptransform_list = ptransform_mapper.create_and_save_ptransform_list() self.assertTrue(len(ptransform_list) == 2) self.assertEqual(type(ptransform_list[0]), RunInference) expected_columns = [['x'], ['y', 'z']] expected_dimensions = [512, None] for i in range(len(ptransform_list)): self.assertEqual(type(ptransform_list[i]), RunInference) self.assertEqual( type(ptransform_list[i]._model_handler), base._TextEmbeddingHandler) self.assertEqual( ptransform_list[i]._model_handler.columns, expected_columns[i]) self.assertEqual( ptransform_list[i]._model_handler._underlying.model_name, model_name) if expected_dimensions[i]: self.assertEqual( ptransform_list[i]._model_handler._underlying.dimensions, expected_dimensions[i]) ptransform_list = ( base._MLTransformToPTransformMapper. load_transforms_from_artifact_location(self.artifact_location)) for i in range(len(ptransform_list)): self.assertEqual(type(ptransform_list[i]), RunInference) self.assertEqual( type(ptransform_list[i]._model_handler), base._TextEmbeddingHandler) self.assertEqual( ptransform_list[i]._model_handler.columns, expected_columns[i]) self.assertEqual( ptransform_list[i]._model_handler._underlying.model_name, model_name) if expected_dimensions[i]: self.assertEqual( ptransform_list[i]._model_handler._underlying.dimensions, expected_dimensions[i]) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/embeddings/tensorflow_hub.py0000644000175100017510000001573315206113032026207 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from collections.abc import Iterable from typing import Optional import tensorflow as tf import tensorflow_hub as hub import tensorflow_text as text # required to register TF ops. # pylint: disable=unused-import import apache_beam as beam from apache_beam.ml.inference import utils from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import PredictionResult from apache_beam.ml.inference.base import RunInference from apache_beam.ml.inference.tensorflow_inference import TFModelHandlerTensor from apache_beam.ml.inference.tensorflow_inference import default_tensor_inference_fn from apache_beam.ml.transforms.base import EmbeddingsManager from apache_beam.ml.transforms.base import _ImageEmbeddingHandler from apache_beam.ml.transforms.base import _TextEmbeddingHandler __all__ = ['TensorflowHubTextEmbeddings', 'TensorflowHubImageEmbeddings'] # TODO: https://github.com/apache/beam/issues/30288 # Replace with TFModelHandlerTensor when load_model() supports TFHUB models. class _TensorflowHubModelHandler(TFModelHandlerTensor): """ Note: Intended for internal use only. No backwards compatibility guarantees. """ def __init__(self, preprocessing_url: Optional[str], *args, **kwargs): self.preprocessing_url = preprocessing_url super().__init__(*args, **kwargs) def load_model(self): # unable to load the models with tf.keras.models.load_model so # using hub.KerasLayer instead model = hub.KerasLayer(self._model_uri, **self._load_model_args) return model def _convert_prediction_result_to_list( self, predictions: Iterable[PredictionResult]): result = [] for prediction in predictions: inference = prediction.inference.numpy().tolist() result.append(inference) return result def run_inference(self, batch, model, inference_args, model_id=None): if not inference_args: inference_args = {} if not self.preprocessing_url: predictions = default_tensor_inference_fn( model=model, batch=batch, inference_args=inference_args, model_id=model_id) return self._convert_prediction_result_to_list(predictions) vectorized_batch = tf.stack(batch, axis=0) preprocessor_fn = hub.KerasLayer(self.preprocessing_url) vectorized_batch = preprocessor_fn(vectorized_batch) predictions = model(vectorized_batch) # https://www.tensorflow.org/text/tutorials/classify_text_with_bert#using_the_bert_model # pylint: disable=line-too-long # pooled_output -> represents the text as a whole. This is an embeddings # of the whole text. The shape is [batch_size, embedding_dimension] # sequence_output -> represents the text as a sequence of tokens. This is # an embeddings of each token in the text. The shape is # [batch_size, max_sequence_length, embedding_dimension] # pooled output is the embeedings as per the documentation. so let's use # that. embeddings = predictions['pooled_output'] predictions = utils._convert_to_result(batch, embeddings, model_id) return self._convert_prediction_result_to_list(predictions) class TensorflowHubTextEmbeddings(EmbeddingsManager): def __init__( self, columns: list[str], hub_url: str, preprocessing_url: Optional[str] = None, **kwargs): """ Embedding config for tensorflow hub models. This config can be used with MLTransform to embed text data. Models are loaded using the RunInference PTransform with the help of a ModelHandler. Args: columns: The columns containing the text to be embedded. hub_url: The url of the tensorflow hub model. preprocessing_url: The url of the preprocessing model. This is optional. If provided, the preprocessing model will be used to preprocess the text before feeding it to the main model. min_batch_size: The minimum batch size to be used for inference. max_batch_size: The maximum batch size to be used for inference. large_model: Whether to share the model across processes. """ super().__init__(columns=columns, **kwargs) self.model_uri = hub_url self.preprocessing_url = preprocessing_url def get_model_handler(self) -> ModelHandler: # override the default inference function return _TensorflowHubModelHandler( model_uri=self.model_uri, preprocessing_url=self.preprocessing_url, min_batch_size=self.min_batch_size, max_batch_size=self.max_batch_size, large_model=self.large_model, ) def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: """ Returns a RunInference object that is used to run inference on the text input using _TextEmbeddingHandler. """ return ( RunInference( model_handler=_TextEmbeddingHandler(self), inference_args=self.inference_args, )) class TensorflowHubImageEmbeddings(EmbeddingsManager): def __init__(self, columns: list[str], hub_url: str, **kwargs): """ Embedding config for tensorflow hub models. This config can be used with MLTransform to embed image data. Models are loaded using the RunInference PTransform with the help of a ModelHandler. Args: columns: The columns containing the images to be embedded. hub_url: The url of the tensorflow hub model. min_batch_size: The minimum batch size to be used for inference. max_batch_size: The maximum batch size to be used for inference. large_model: Whether to share the model across processes. """ super().__init__(columns=columns, **kwargs) self.model_uri = hub_url def get_model_handler(self) -> ModelHandler: # override the default inference function return _TensorflowHubModelHandler( model_uri=self.model_uri, preprocessing_url=None, min_batch_size=self.min_batch_size, max_batch_size=self.max_batch_size, large_model=self.large_model, ) def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: """ Returns a RunInference object that is used to run inference on the text input using _ImageEmbeddingHandler. """ return ( RunInference( model_handler=_ImageEmbeddingHandler(self), inference_args=self.inference_args, )) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py0000644000175100017510000001747615206113032027254 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os import shutil import tempfile import unittest import uuid import numpy as np import apache_beam as beam from apache_beam.ml.transforms.base import MLTransform hub_url = 'https://tfhub.dev/google/nnlm-en-dim128/2' hub_img_url = 'https://www.kaggle.com/models/google/resnet-v2/TensorFlow2/101-feature-vector/2' # pylint: disable=line-too-long test_query_column = 'test_query' test_query = 'This is a test query' # pylint: disable=ungrouped-imports try: from apache_beam.ml.transforms.embeddings.tensorflow_hub import TensorflowHubTextEmbeddings except ImportError: TensorflowHubTextEmbeddings = None # type: ignore tf = None # pylint: disable=ungrouped-imports try: import tensorflow_transform as tft from apache_beam.ml.transforms.tft import ScaleTo01 except ImportError: tft = None # pylint: disable=ungrouped-imports try: from PIL import Image from apache_beam.ml.transforms.embeddings.tensorflow_hub import TensorflowHubImageEmbeddings except ImportError: TensorflowHubImageEmbeddings = None # type: ignore Image = None # type: ignore[assignment] @unittest.skipIf( TensorflowHubTextEmbeddings is None, 'Tensorflow is not installed.') class TFHubEmbeddingsTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self) -> None: shutil.rmtree(self.artifact_location) def test_tfhub_text_embeddings(self): embedding_config = TensorflowHubTextEmbeddings( hub_url=hub_url, columns=[test_query_column]) with beam.Pipeline() as pipeline: transformed_pcoll = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: test_query }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) def assert_element(element): assert len(element[test_query_column]) == 128 _ = (transformed_pcoll | beam.Map(assert_element)) @unittest.skipIf(tft is None, 'Tensorflow Transform is not installed.') def test_embeddings_with_scale_to_0_1(self): embedding_config = TensorflowHubTextEmbeddings( hub_url=hub_url, columns=[test_query_column], ) with beam.Pipeline() as pipeline: transformed_pcoll = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: test_query }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config).with_transform( ScaleTo01(columns=[test_query_column]))) def assert_element(element): assert max(element[test_query_column]) == 1 _ = ( transformed_pcoll | beam.Map(lambda x: x.as_dict()) | beam.Map(assert_element)) def pipeline_with_configurable_artifact_location( self, pipeline, embedding_config=None, read_artifact_location=None, write_artifact_location=None): if write_artifact_location: return ( pipeline | MLTransform(write_artifact_location=write_artifact_location). with_transform(embedding_config)) elif read_artifact_location: return ( pipeline | MLTransform(read_artifact_location=read_artifact_location)) else: raise NotImplementedError def test_embeddings_with_read_artifact_location(self): with beam.Pipeline() as p: embedding_config = TensorflowHubTextEmbeddings( hub_url=hub_url, columns=[test_query_column]) with beam.Pipeline() as p: data = ( p | "CreateData" >> beam.Create([{ test_query_column: test_query }])) _ = self.pipeline_with_configurable_artifact_location( pipeline=data, embedding_config=embedding_config, write_artifact_location=self.artifact_location) with beam.Pipeline() as p: data = ( p | "CreateData" >> beam.Create([{ test_query_column: test_query }, { test_query_column: test_query }])) result_pcoll = self.pipeline_with_configurable_artifact_location( pipeline=data, read_artifact_location=self.artifact_location) def assert_element(element): # 0.29836970567703247 assert round(element, 2) == 0.3 _ = ( result_pcoll | beam.Map(lambda x: max(x[test_query_column])) | beam.Map(assert_element)) def test_with_int_data_types(self): embedding_config = TensorflowHubTextEmbeddings( hub_url=hub_url, columns=[test_query_column]) with self.assertRaises(Exception): with beam.Pipeline() as pipeline: _ = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: 1 }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) @unittest.skipIf( TensorflowHubImageEmbeddings is None, 'Tensorflow is not installed.') class TFHubImageEmbeddingsTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self) -> None: shutil.rmtree(self.artifact_location) def generateRandomImage(self, size: int): imarray = np.random.rand(size, size, 3) * 255 return imarray / 255.0 @unittest.skipIf(Image is None, 'Pillow is not installed.') def test_sentence_transformer_image_embeddings(self): embedding_config = TensorflowHubImageEmbeddings( hub_url=hub_img_url, columns=[test_query_column]) img = self.generateRandomImage(224) with beam.Pipeline() as pipeline: result_pcoll = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: img }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) def assert_element(element): assert len(element[test_query_column]) == 2048 _ = (result_pcoll | beam.Map(assert_element)) def test_with_str_data_types(self): embedding_config = TensorflowHubImageEmbeddings( hub_url=hub_img_url, columns=[test_query_column]) with self.assertRaises(TypeError): with beam.Pipeline() as pipeline: _ = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: "img.jpg" }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) @unittest.skipIf( TensorflowHubTextEmbeddings is None, 'Tensorflow is not installed.') class TFHubEmbeddingsGCSArtifactLocationTest(TFHubEmbeddingsTest): def setUp(self): self.artifact_location = os.path.join( 'gs://temp-storage-for-perf-tests/tfhub', uuid.uuid4().hex) def tearDown(self): pass if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/embeddings/vertex_ai.py0000644000175100017510000004530015206113032025126 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # Vertex AI Python SDK is required for this module. # Follow https://cloud.google.com/vertex-ai/docs/python-sdk/use-vertex-ai-python-sdk # pylint: disable=line-too-long # to install Vertex AI Python SDK. import functools import logging from collections.abc import Callable from collections.abc import Sequence from dataclasses import dataclass from typing import Any from typing import Optional from typing import cast import vertexai from google.api_core.exceptions import ServerError from google.api_core.exceptions import TooManyRequests from google.auth.credentials import Credentials from vertexai.language_models import TextEmbeddingInput from vertexai.language_models import TextEmbeddingModel from vertexai.vision_models import Image from vertexai.vision_models import MultiModalEmbeddingModel from vertexai.vision_models import MultiModalEmbeddingResponse from vertexai.vision_models import Video from vertexai.vision_models import VideoEmbedding from vertexai.vision_models import VideoSegmentConfig import apache_beam as beam from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import RemoteModelHandler from apache_beam.ml.inference.base import RunInference from apache_beam.ml.rag.types import EmbeddableItem from apache_beam.ml.rag.types import Embedding from apache_beam.ml.transforms.base import EmbeddingsManager from apache_beam.ml.transforms.base import EmbeddingTypeAdapter from apache_beam.ml.transforms.base import _ImageEmbeddingHandler from apache_beam.ml.transforms.base import _MultiModalEmbeddingHandler from apache_beam.ml.transforms.base import _TextEmbeddingHandler __all__ = [ "VertexAITextEmbeddings", "VertexAIImageEmbeddings", "VertexAIMultiModalEmbeddings", "VertexAIMultiModalInput", ] DEFAULT_TASK_TYPE = "RETRIEVAL_DOCUMENT" # TODO: https://github.com/apache/beam/issues/29356 # Can this list be automatically pulled from Vertex SDK? TASK_TYPE_INPUTS = [ "RETRIEVAL_DOCUMENT", "RETRIEVAL_QUERY", "SEMANTIC_SIMILARITY", "CLASSIFICATION", "CLUSTERING" ] _BATCH_SIZE = 5 # Vertex AI limits requests to 5 at a time. LOGGER = logging.getLogger("VertexAIEmbeddings") def _retry_on_appropriate_gcp_error(exception): """ Retry filter that returns True if a returned HTTP error code is 5xx or 429. This is used to retry remote requests that fail, most notably 429 (TooManyRequests.) Args: exception: the returned exception encountered during the request/response loop. Returns: boolean indication whether or not the exception is a Server Error (5xx) or a TooManyRequests (429) error. """ return isinstance(exception, (TooManyRequests, ServerError)) class _VertexAITextEmbeddingHandler(RemoteModelHandler): """ Note: Intended for internal use and guarantees no backwards compatibility. """ def __init__( self, model_name: str, title: Optional[str] = None, task_type: str = DEFAULT_TASK_TYPE, project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[Credentials] = None, **kwargs): vertexai.init(project=project, location=location, credentials=credentials) self.model_name = model_name if task_type not in TASK_TYPE_INPUTS: raise ValueError( f"task_type must be one of {TASK_TYPE_INPUTS}, got {task_type}") self.task_type = task_type self.title = title super().__init__( namespace='VertexAITextEmbeddingHandler', retry_filter=_retry_on_appropriate_gcp_error, **kwargs) def request( self, batch: Sequence[str], model: TextEmbeddingModel, inference_args: Optional[dict[str, Any]] = None): embeddings = [] batch_size = _BATCH_SIZE for i in range(0, len(batch), batch_size): text_batch_strs = batch[i:i + batch_size] text_batch = [ TextEmbeddingInput( text=text, title=self.title, task_type=self.task_type) for text in text_batch_strs ] embeddings_batch = model.get_embeddings(list(text_batch)) embeddings.extend([el.values for el in embeddings_batch]) return embeddings def create_client(self) -> TextEmbeddingModel: model = TextEmbeddingModel.from_pretrained(self.model_name) return model def __repr__(self): # ModelHandler is internal to the user and is not exposed. # Hence we need to override the __repr__ method to expose # the name of the class. return 'VertexAITextEmbeddings' class VertexAITextEmbeddings(EmbeddingsManager): def __init__( self, model_name: str, columns: list[str], title: Optional[str] = None, task_type: str = DEFAULT_TASK_TYPE, project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[Credentials] = None, **kwargs): """ Embedding Config for Vertex AI Text Embedding models following https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings # pylint: disable=line-too-long Text Embeddings are generated for a batch of text using the Vertex AI SDK. Embeddings are returned in a list for each text in the batch. Look at https://cloud.google.com/vertex-ai/docs/generative-ai/learn/model-versioning#stable-versions-available.md # pylint: disable=line-too-long for more information on model versions and lifecycle. Args: model_name: The name of the Vertex AI Text Embedding model. columns: The columns containing the text to be embedded. task_type: The downstream task for the embeddings. Valid values are RETRIEVAL_QUERY, RETRIEVAL_DOCUMENT, SEMANTIC_SIMILARITY, CLASSIFICATION, CLUSTERING. For more information on the task type, look at https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings # pylint: disable=line-too-long title: Identifier of the text content. project: The default GCP project for API calls. location: The default location for API calls. credentials: Custom credentials for API calls. Defaults to environment credentials. """ self.model_name = model_name self.project = project self.location = location self.credentials = credentials self.title = title self.task_type = task_type self.kwargs = kwargs super().__init__(columns=columns, **kwargs) def get_model_handler(self) -> ModelHandler: return _VertexAITextEmbeddingHandler( model_name=self.model_name, project=self.project, location=self.location, credentials=self.credentials, title=self.title, task_type=self.task_type, **self.kwargs) def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: return RunInference( model_handler=_TextEmbeddingHandler(self), inference_args=self.inference_args) class _VertexAIImageEmbeddingHandler(RemoteModelHandler): def __init__( self, model_name: str, dimension: Optional[int] = None, project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[Credentials] = None, **kwargs): vertexai.init(project=project, location=location, credentials=credentials) self.model_name = model_name self.dimension = dimension super().__init__( namespace='VertexAIImageEmbeddingHandler', retry_filter=_retry_on_appropriate_gcp_error, **kwargs) def request( self, imgs: Sequence[Image], model: MultiModalEmbeddingModel, inference_args: Optional[dict[str, Any]] = None): embeddings = [] # Max request size for multi-modal embedding models is 1 for img in imgs: prediction = model.get_embeddings(image=img, dimension=self.dimension) embeddings.append(prediction.image_embedding) return embeddings def create_client(self): model = MultiModalEmbeddingModel.from_pretrained(self.model_name) return model def __repr__(self): # ModelHandler is internal to the user and is not exposed. # Hence we need to override the __repr__ method to expose # the name of the class. return 'VertexAIImageEmbeddings' class VertexAIImageEmbeddings(EmbeddingsManager): def __init__( self, model_name: str, columns: list[str], dimension: Optional[int], project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[Credentials] = None, **kwargs): """ Embedding Config for Vertex AI Image Embedding models following https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-multimodal-embeddings # pylint: disable=line-too-long Image Embeddings are generated for a batch of images using the Vertex AI API. Embeddings are returned in a list for each image in the batch. This transform makes remote calls to the Vertex AI service and may incur costs for use. Args: model_name: The name of the Vertex AI Multi-Modal Embedding model. columns: The columns containing the image to be embedded. dimension: The length of the embedding vector to generate. Must be one of 128, 256, 512, or 1408. If not set, Vertex AI's default value is 1408. project: The default GCP project for API calls. location: The default location for API calls. credentials: Custom credentials for API calls. Defaults to environment credentials. """ self.model_name = model_name self.project = project self.location = location self.credentials = credentials self.kwargs = kwargs if dimension is not None and dimension not in (128, 256, 512, 1408): raise ValueError( "dimension argument must be one of 128, 256, 512, or 1408") self.dimension = dimension super().__init__(columns=columns, **kwargs) def get_model_handler(self) -> ModelHandler: return _VertexAIImageEmbeddingHandler( model_name=self.model_name, dimension=self.dimension, project=self.project, location=self.location, credentials=self.credentials, **self.kwargs) def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: return RunInference( model_handler=_ImageEmbeddingHandler(self), inference_args=self.inference_args) @dataclass class VertexImage: image_content: Image embedding: Optional[list[float]] = None @dataclass class VertexVideo: video_content: Video config: VideoSegmentConfig embeddings: Optional[list[VideoEmbedding]] = None @dataclass class VertexAIMultiModalInput: image: Optional[VertexImage] = None video: Optional[VertexVideo] = None contextual_text: Optional[EmbeddableItem] = None class _VertexAIMultiModalEmbeddingHandler(RemoteModelHandler): def __init__( self, model_name: str, dimension: Optional[int] = None, project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[Credentials] = None, **kwargs): vertexai.init(project=project, location=location, credentials=credentials) self.model_name = model_name self.dimension = dimension super().__init__( namespace='VertexAIMultiModelEmbeddingHandler', retry_filter=_retry_on_appropriate_gcp_error, **kwargs) def request( self, batch: Sequence[VertexAIMultiModalInput], model: MultiModalEmbeddingModel, inference_args: Optional[dict[str, Any]] = None): embeddings = [] # Max request size for multi-modal embedding models is 1 for input in batch: image_content: Optional[Image] = None video_content: Optional[Video] = None text_content: Optional[str] = None video_config: Optional[VideoSegmentConfig] = None if input.image: image_content = input.image.image_content if input.video: video_content = input.video.video_content video_config = input.video.config if input.contextual_text: text_content = input.contextual_text.content.text prediction = model.get_embeddings( image=image_content, video=video_content, contextual_text=text_content, dimension=self.dimension, video_segment_config=video_config) embeddings.append(prediction) return embeddings def create_client(self) -> MultiModalEmbeddingModel: model = MultiModalEmbeddingModel.from_pretrained(self.model_name) return model def __repr__(self): # ModelHandler is internal to the user and is not exposed. # Hence we need to override the __repr__ method to expose # the name of the class. return 'VertexAIMultiModalEmbeddings' def _multimodal_dict_input_fn( image_column: Optional[str], video_column: Optional[str], text_column: Optional[str], batch: Sequence[dict[str, Any]]) -> list[VertexAIMultiModalInput]: multimodal_inputs: list[VertexAIMultiModalInput] = [] for item in batch: img: Optional[VertexImage] = None vid: Optional[VertexVideo] = None text: Optional[EmbeddableItem] = None if image_column: img = item[image_column] if video_column: vid = item[video_column] if text_column: text = item[text_column] multimodal_inputs.append( VertexAIMultiModalInput(image=img, video=vid, contextual_text=text)) return multimodal_inputs def _multimodal_dict_output_fn( image_column: Optional[str], video_column: Optional[str], text_column: Optional[str], batch: Sequence[dict[str, Any]], embeddings: Sequence[MultiModalEmbeddingResponse]) -> list[dict[str, Any]]: results = [] for batch_idx, item in enumerate(batch): mm_embedding = embeddings[batch_idx] if image_column: item[image_column].embedding = mm_embedding.image_embedding if video_column: item[video_column].embeddings = mm_embedding.video_embeddings if text_column: item[text_column].embedding = Embedding( dense_embedding=mm_embedding.text_embedding) results.append(item) return results def _create_multimodal_dict_adapter( image_column: Optional[str], video_column: Optional[str], text_column: Optional[str] ) -> EmbeddingTypeAdapter[dict[str, Any], dict[str, Any]]: return EmbeddingTypeAdapter[dict[str, Any], dict[str, Any]]( input_fn=cast( Callable[[Sequence[dict[str, Any]]], list[str]], functools.partial( _multimodal_dict_input_fn, image_column, video_column, text_column)), output_fn=cast( Callable[[Sequence[dict[str, Any]], Sequence[Any]], list[dict[str, Any]]], functools.partial( _multimodal_dict_output_fn, image_column, video_column, text_column))) class VertexAIMultiModalEmbeddings(EmbeddingsManager): def __init__( self, model_name: str, image_column: Optional[str] = None, video_column: Optional[str] = None, text_column: Optional[str] = None, dimension: Optional[int] = None, project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[Credentials] = None, **kwargs): """ Embedding Config for Vertex AI Multi-Modal Embedding models following https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-multimodal-embeddings # pylint: disable=line-too-long Multi-Modal Embeddings are generated for a batch of image, video, and string groupings using the Vertex AI API. Embeddings are returned in a list for each image in the batch as MultiModalEmbeddingResponses. This transform makes remote calls to the Vertex AI service and may incur costs for use. Args: model_name: The name of the Vertex AI Multi-Modal Embedding model. image_column: The column containing image data to be embedded. This data is expected to be formatted as VertexImage objects, containing a Vertex Image object. video_column: The column containing video data to be embedded. This data is expected to be formatted as VertexVideo objects, containing a Vertex Video object an a VideoSegmentConfig object. text_column: The column containing text data to be embedded. This data is expected to be formatted as EmbeddableItem objects, containing the string to be embedded in the item's content field. dimension: The length of the embedding vector to generate. Must be one of 128, 256, 512, or 1408. If not set, Vertex AI's default value is 1408. If submitting video content, dimension *musst* be 1408. project: The default GCP project for API calls. location: The default location for API calls. credentials: Custom credentials for API calls. Defaults to environment credentials. """ self.model_name = model_name self.project = project self.location = location self.credentials = credentials self.kwargs = kwargs if dimension is not None and dimension not in (128, 256, 512, 1408): raise ValueError( "dimension argument must be one of 128, 256, 512, or 1408") self.dimension = dimension if not image_column and not video_column and not text_column: raise ValueError("at least one input column must be specified") if video_column is not None and dimension != 1408: raise ValueError( "Vertex AI does not support custom dimensions for video input, want dimension = 1408, got ", dimension) self.type_adapter = _create_multimodal_dict_adapter( image_column=image_column, video_column=video_column, text_column=text_column) super().__init__(type_adapter=self.type_adapter, **kwargs) def get_model_handler(self) -> ModelHandler: return _VertexAIMultiModalEmbeddingHandler( model_name=self.model_name, dimension=self.dimension, project=self.project, location=self.location, credentials=self.credentials, **self.kwargs) def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: return RunInference( model_handler=_MultiModalEmbeddingHandler(self), inference_args=self.inference_args) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/embeddings/vertex_ai_it_test.py0000644000175100017510000003516215206113032026666 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os import shutil import tempfile import unittest import uuid import pytest import apache_beam as beam from apache_beam.ml.inference.base import RunInference from apache_beam.ml.transforms import base from apache_beam.ml.transforms.base import MLTransform pytest.importorskip("vertexai", reason="Vertex AI dependencies not available") # pylint: disable=ungrouped-imports # isort: off try: from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content from apache_beam.ml.transforms.embeddings.vertex_ai import VertexAIImageEmbeddings from apache_beam.ml.transforms.embeddings.vertex_ai import VertexAIMultiModalEmbeddings from apache_beam.ml.transforms.embeddings.vertex_ai import VertexAITextEmbeddings from apache_beam.ml.transforms.embeddings.vertex_ai import VertexImage from apache_beam.ml.transforms.embeddings.vertex_ai import VertexVideo # Load the Vertex dependencies last so type resolution still pulls in RAG # types. from vertexai.vision_models import Image from vertexai.vision_models import Video from vertexai.vision_models import VideoSegmentConfig except ImportError: VertexAIMultiModalEmbeddings = None # type: ignore VertexAITextEmbeddings = None # type: ignore VertexAIImageEmbeddings = None # type: ignore try: import tensorflow_transform as tft from apache_beam.ml.transforms.tft import ScaleTo01 except ImportError: tft = None test_query = "This is a test" test_query_column = "feature_1" model_name: str = "text-embedding-005" @pytest.mark.vertex_ai_postcommit class VertexAIEmbeddingsTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp(prefix='_vertex_ai_test') self.gcs_artifact_location = os.path.join( 'gs://temp-storage-for-perf-tests/vertex_ai', uuid.uuid4().hex) def tearDown(self) -> None: shutil.rmtree(self.artifact_location) def test_vertex_ai_text_embeddings(self): embedding_config = VertexAITextEmbeddings( model_name=model_name, columns=[test_query_column]) with beam.Pipeline() as pipeline: transformed_pcoll = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: test_query }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) def assert_element(element): assert len(element[test_query_column]) == 768 _ = (transformed_pcoll | beam.Map(assert_element)) @unittest.skipIf(tft is None, 'Tensorflow Transform is not installed.') def test_embeddings_with_scale_to_0_1(self): embedding_config = VertexAITextEmbeddings( model_name=model_name, columns=[test_query_column], ) with beam.Pipeline() as pipeline: transformed_pcoll = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: test_query }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config).with_transform( ScaleTo01(columns=[test_query_column]))) def assert_element(element): assert max(element.feature_1) == 1 _ = (transformed_pcoll | beam.Map(assert_element)) def pipeline_with_configurable_artifact_location( self, pipeline, embedding_config=None, read_artifact_location=None, write_artifact_location=None): if write_artifact_location: return ( pipeline | MLTransform(write_artifact_location=write_artifact_location). with_transform(embedding_config)) elif read_artifact_location: return ( pipeline | MLTransform(read_artifact_location=read_artifact_location)) else: raise NotImplementedError def test_embeddings_with_read_artifact_location(self): with beam.Pipeline() as p: embedding_config = VertexAITextEmbeddings( model_name=model_name, columns=[test_query_column]) with beam.Pipeline() as p: data = ( p | "CreateData" >> beam.Create([{ test_query_column: test_query }])) _ = self.pipeline_with_configurable_artifact_location( pipeline=data, embedding_config=embedding_config, write_artifact_location=self.artifact_location) with beam.Pipeline() as p: data = ( p | "CreateData" >> beam.Create([{ test_query_column: test_query }, { test_query_column: test_query }])) result_pcoll = self.pipeline_with_configurable_artifact_location( pipeline=data, read_artifact_location=self.artifact_location) def assert_element(element): assert round(element, 2) == 0.11 _ = ( result_pcoll | beam.Map(lambda x: max(x[test_query_column])) # 0.14797046780586243 | beam.Map(assert_element)) def test_with_int_data_types(self): embedding_config = VertexAITextEmbeddings( model_name=model_name, columns=[test_query_column]) with self.assertRaisesRegex(Exception, "Embeddings can only be generated"): with beam.Pipeline() as pipeline: _ = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: 1 }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) def test_with_gcs_artifact_location(self): with beam.Pipeline() as p: embedding_config = VertexAITextEmbeddings( model_name=model_name, columns=[test_query_column]) with beam.Pipeline() as p: data = ( p | "CreateData" >> beam.Create([{ test_query_column: test_query }])) _ = self.pipeline_with_configurable_artifact_location( pipeline=data, embedding_config=embedding_config, write_artifact_location=self.gcs_artifact_location) with beam.Pipeline() as p: data = ( p | "CreateData" >> beam.Create([{ test_query_column: test_query }, { test_query_column: test_query }])) result_pcoll = self.pipeline_with_configurable_artifact_location( pipeline=data, read_artifact_location=self.gcs_artifact_location) def assert_element(element): assert round(element, 2) == 0.11 _ = ( result_pcoll | beam.Map(lambda x: max(x[test_query_column])) # 0.14797046780586243 | beam.Map(assert_element)) def test_mltransform_to_ptransform_with_vertex(self): model_name = 'text-embedding-005' transforms = [ VertexAITextEmbeddings( columns=['x'], model_name=model_name, task_type='RETRIEVAL_DOCUMENT'), VertexAITextEmbeddings( columns=['y', 'z'], model_name=model_name, task_type='CLUSTERING') ] ptransform_mapper = base._MLTransformToPTransformMapper( transforms=transforms, artifact_location=self.artifact_location, artifact_mode=None) ptransform_list = ptransform_mapper.create_and_save_ptransform_list() self.assertTrue(len(ptransform_list) == 2) self.assertEqual(type(ptransform_list[0]), RunInference) expected_columns = [['x'], ['y', 'z']] expected_task_type = ['RETRIEVAL_DOCUMENT', 'CLUSTERING'] for i in range(len(ptransform_list)): self.assertEqual(type(ptransform_list[i]), RunInference) self.assertEqual( type(ptransform_list[i]._model_handler), base._TextEmbeddingHandler) self.assertEqual( ptransform_list[i]._model_handler.columns, expected_columns[i]) self.assertEqual( ptransform_list[i]._model_handler._underlying.task_type, expected_task_type[i]) self.assertEqual( ptransform_list[i]._model_handler._underlying.model_name, model_name) ptransform_list = ( base._MLTransformToPTransformMapper. load_transforms_from_artifact_location(self.artifact_location)) for i in range(len(ptransform_list)): self.assertEqual(type(ptransform_list[i]), RunInference) self.assertEqual( type(ptransform_list[i]._model_handler), base._TextEmbeddingHandler) self.assertEqual( ptransform_list[i]._model_handler.columns, expected_columns[i]) self.assertEqual( ptransform_list[i]._model_handler._underlying.task_type, expected_task_type[i]) self.assertEqual( ptransform_list[i]._model_handler._underlying.model_name, model_name) @pytest.mark.vertex_ai_postcommit class VertexAIImageEmbeddingsTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp(prefix='_vertex_ai_image_test') self.gcs_artifact_location = os.path.join( 'gs://temp-storage-for-perf-tests/vertex_ai_image', uuid.uuid4().hex) self.model_name = "multimodalembedding" self.image_path = "gs://apache-beam-ml/testing/inputs/vertex_images/sunflowers/1008566138_6927679c8a.jpg" # pylint: disable=line-too-long def tearDown(self) -> None: shutil.rmtree(self.artifact_location) def test_vertex_ai_image_embedding(self): embedding_config = VertexAIImageEmbeddings( model_name=self.model_name, columns=[test_query_column], dimension=128) with beam.Pipeline() as pipeline: transformed_pcoll = ( pipeline | "CreateData" >> beam.Create([{ test_query_column: Image(gcs_uri=self.image_path) }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) def assert_element(element): assert len(element[test_query_column]) == 128 _ = (transformed_pcoll | beam.Map(assert_element)) def test_improper_dimension(self): with self.assertRaises(ValueError): _ = VertexAIImageEmbeddings( model_name=self.model_name, columns=[test_query_column], dimension=127) image_feature_column: str = "img_feature" text_feature_column: str = "txt_feature" video_feature_column: str = "vid_feature" def _make_text_chunk(input: str) -> Chunk: return Chunk(content=Content(text=input)) @pytest.mark.vertex_ai_postcommit class VertexAIMultiModalEmbeddingsTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp( prefix='_vertex_ai_multi_modal_test') self.gcs_artifact_location = os.path.join( 'gs://temp-storage-for-perf-tests/vertex_ai_multi_modal', uuid.uuid4().hex) self.model_name = "multimodalembedding" self.image_path = "gs://apache-beam-ml/testing/inputs/vertex_images/sunflowers/1008566138_6927679c8a.jpg" # pylint: disable=line-too-long self.video_path = "gs://cloud-samples-data/vertex-ai-vision/highway_vehicles.mp4" # pylint: disable=line-too-long self.video_segment_config = VideoSegmentConfig(end_offset_sec=1) def tearDown(self) -> None: shutil.rmtree(self.artifact_location) def test_vertex_ai_multimodal_embedding_img_and_text(self): embedding_config = VertexAIMultiModalEmbeddings( model_name=self.model_name, image_column=image_feature_column, text_column=text_feature_column, dimension=128, project="apache-beam-testing", location="us-central1") with beam.Pipeline() as pipeline: transformed_pcoll = ( pipeline | "CreateData" >> beam.Create([{ image_feature_column: VertexImage( image_content=Image(gcs_uri=self.image_path)), text_feature_column: _make_text_chunk("an image of sunflowers"), }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) def assert_element(element): assert len(element[image_feature_column].embedding) == 128 assert len( element[text_feature_column].embedding.dense_embedding) == 128 _ = (transformed_pcoll | beam.Map(assert_element)) def test_vertex_ai_multimodal_embedding_video(self): embedding_config = VertexAIMultiModalEmbeddings( model_name=self.model_name, video_column=video_feature_column, dimension=1408, project="apache-beam-testing", location="us-central1") with beam.Pipeline() as pipeline: transformed_pcoll = ( pipeline | "CreateData" >> beam.Create([{ video_feature_column: VertexVideo( video_content=Video(gcs_uri=self.video_path), config=self.video_segment_config) }]) | "MLTransform" >> MLTransform( write_artifact_location=self.artifact_location).with_transform( embedding_config)) def assert_element(element): # Videos are returned in VideoEmbedding objects, must unroll # for each segment. for segment in element[video_feature_column].embeddings: assert len(segment.embedding) == 1408 _ = (transformed_pcoll | beam.Map(assert_element)) def test_improper_dimension(self): with self.assertRaises(ValueError): _ = VertexAIMultiModalEmbeddings( model_name=self.model_name, image_column="fake_img_column", dimension=127) def test_missing_columns(self): with self.assertRaises(ValueError): _ = VertexAIMultiModalEmbeddings( model_name=self.model_name, dimension=128) def test_improper_video_dimension(self): with self.assertRaises(ValueError): _ = VertexAIMultiModalEmbeddings( model_name=self.model_name, video_column=video_feature_column, dimension=128) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/handlers.py0000644000175100017510000004534715206113032022652 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import collections import copy import os import typing from collections.abc import Sequence from typing import Any from typing import NamedTuple from typing import Optional from typing import Union import numpy as np import tensorflow as tf import tensorflow_transform.beam as tft_beam from tensorflow_metadata.proto.v0 import schema_pb2 from tensorflow_transform import common_types from tensorflow_transform.beam.tft_beam_io import beam_metadata_io from tensorflow_transform.beam.tft_beam_io import transform_fn_io from tensorflow_transform.tf_metadata import dataset_metadata from tensorflow_transform.tf_metadata import metadata_io from tensorflow_transform.tf_metadata import schema_utils import apache_beam as beam from apache_beam import coders from apache_beam.io.filesystems import FileSystems from apache_beam.ml.transforms.base import ArtifactMode from apache_beam.ml.transforms.base import ProcessHandler from apache_beam.ml.transforms.tft import _EXPECTED_TYPES from apache_beam.ml.transforms.tft import TFTOperation from apache_beam.typehints import native_type_compatibility from apache_beam.typehints.row_type import RowTypeConstraint __all__ = [ 'TFTProcessHandler', ] _TEMP_KEY = 'CODED_SAMPLE' # key for the encoded sample RAW_DATA_METADATA_DIR = 'raw_data_metadata' SCHEMA_FILE = 'schema.pbtxt' # tensorflow transform doesn't support the types other than tf.int64, # tf.float32 and tf.string. _default_type_to_tensor_type_map = { int: tf.int64, float: tf.float32, str: tf.string, bytes: tf.string, np.int64: tf.int64, np.int32: tf.int64, np.float32: tf.float32, np.float64: tf.float32, np.bytes_: tf.string, np.str_: tf.string, } _primitive_types_to_typing_container_type = { int: list[int], float: list[float], str: list[str], bytes: list[bytes] } tft_process_handler_input_type = Union[NamedTuple, beam.Row, dict[str, Union[str, float, int, bytes, np.ndarray]]] tft_process_handler_output_type = Union[beam.Row, dict[str, np.ndarray]] class _DataCoder: def __init__( self, exclude_columns, coder=coders.registry.get_coder(Any), ): """ Encodes/decodes items of a dictionary into a single element. Args: exclude_columns: list of columns to exclude from the encoding. """ self.coder = coder self.exclude_columns = exclude_columns def encode(self, element): data_to_encode = element.copy() element_to_return = element.copy() for key in self.exclude_columns: if key in data_to_encode: del data_to_encode[key] element_to_return[_TEMP_KEY] = self.coder.encode(data_to_encode) return element_to_return def decode(self, element): clone = copy.copy(element) clone.update(self.coder.decode(clone[_TEMP_KEY].item())) del clone[_TEMP_KEY] return clone class _ConvertScalarValuesToListValues(beam.DoFn): def process( self, element, ): new_dict = {} for key, value in element.items(): if isinstance(value, tuple(_primitive_types_to_typing_container_type.keys())): new_dict[key] = [value] else: new_dict[key] = value yield new_dict class _ConvertNamedTupleToDict( beam.PTransform[beam.PCollection[Union[beam.Row, NamedTuple]], beam.PCollection[dict[str, common_types.InstanceDictType]]]): """ A PTransform that converts a collection of NamedTuples or Rows into a collection of dictionaries. """ def expand( self, pcoll: beam.PCollection[Union[beam.Row, NamedTuple]] ) -> beam.PCollection[common_types.InstanceDictType]: """ Args: pcoll: A PCollection of NamedTuples or Rows. Returns: A PCollection of dictionaries. """ return pcoll | beam.Map(lambda x: x._asdict()) class TFTProcessHandler(ProcessHandler[tft_process_handler_input_type, tft_process_handler_output_type]): def __init__( self, *, artifact_location: str, transforms: Optional[Sequence[TFTOperation]] = None, artifact_mode: str = ArtifactMode.PRODUCE): """ A handler class for processing data with TensorFlow Transform (TFT) operations. """ self.transforms = transforms if transforms else [] self.transformed_schema: dict[str, type] = {} self.artifact_location = artifact_location self.artifact_mode = artifact_mode if artifact_mode not in ['produce', 'consume']: raise ValueError('artifact_mode must be either `produce` or `consume`.') def append_transform(self, transform): self.transforms.append(transform) def _map_column_names_to_types(self, row_type): """ Return a dictionary of column names and types. Args: element_type: A type of the element. This could be a NamedTuple or a Row. Returns: A dictionary of column names and types. """ try: if not isinstance(row_type, RowTypeConstraint): row_type = RowTypeConstraint.from_user_type(row_type) inferred_types = {name: typ for name, typ in row_type._fields} for k, t in inferred_types.items(): if t in _primitive_types_to_typing_container_type: inferred_types[k] = _primitive_types_to_typing_container_type[t] # sometimes a numpy type can be provided as np.dtype('int64'). # convert numpy.dtype to numpy type since both are same. for name, typ in inferred_types.items(): if isinstance(typ, np.dtype): inferred_types[name] = typ.type return inferred_types except: # pylint: disable=bare-except return {} def _map_column_names_to_types_from_transforms(self): column_type_mapping = {} for transform in self.transforms: for col in transform.columns: if col not in column_type_mapping: # we just need to dtype of first occurance of column in transforms. class_name = transform.__class__.__name__ if class_name not in _EXPECTED_TYPES: raise KeyError( f"Transform {class_name} is not registered with a supported " "type. Please register the transform with a supported type " "using register_input_dtype decorator.") column_type_mapping[col] = _EXPECTED_TYPES[ transform.__class__.__name__] return column_type_mapping def get_raw_data_feature_spec( self, input_types: dict[str, type]) -> dict[str, tf.io.VarLenFeature]: """ Return a DatasetMetadata object to be used with tft_beam.AnalyzeAndTransformDataset. Args: input_types: A dictionary of column names and types. Returns: A DatasetMetadata object. """ raw_data_feature_spec = {} for key, value in input_types.items(): raw_data_feature_spec[key] = self._get_raw_data_feature_spec_per_column( typ=value, col_name=key) return raw_data_feature_spec def convert_raw_data_feature_spec_to_dataset_metadata( self, raw_data_feature_spec) -> dataset_metadata.DatasetMetadata: raw_data_metadata = dataset_metadata.DatasetMetadata( schema_utils.schema_from_feature_spec(raw_data_feature_spec)) return raw_data_metadata def _get_raw_data_feature_spec_per_column( self, typ: type, col_name: str) -> tf.io.VarLenFeature: """ Return a FeatureSpec object to be used with tft_beam.AnalyzeAndTransformDataset Args: typ: A type of the column. col_name: A name of the column. Returns: A FeatureSpec object. """ # lets conver the builtin types to typing types for consistency. typ = native_type_compatibility.convert_builtin_to_typing(typ) primitive_containers_type = ( list, collections.abc.Sequence, ) is_primitive_container = ( typing.get_origin(typ) in primitive_containers_type) if is_primitive_container: dtype = typing.get_args(typ)[0] if len(typing.get_args(typ)) > 1 or typing.get_origin(dtype) == Union: raise RuntimeError( f"Union type is not supported for column: {col_name}. " f"Please pass a PCollection with valid schema for column " f"{col_name} by passing a single type " "in container. For example, list[int].") elif issubclass(typ, np.generic) or typ in _default_type_to_tensor_type_map: dtype = typ else: raise TypeError( f"Unable to identify type: {typ} specified on column: {col_name}. " f"Please provide a valid type from the following: " f"{_default_type_to_tensor_type_map.keys()}") return tf.io.VarLenFeature(_default_type_to_tensor_type_map[dtype]) def get_raw_data_metadata( self, input_types: dict[str, type]) -> dataset_metadata.DatasetMetadata: raw_data_feature_spec = self.get_raw_data_feature_spec(input_types) raw_data_feature_spec[_TEMP_KEY] = tf.io.VarLenFeature(dtype=tf.string) return self.convert_raw_data_feature_spec_to_dataset_metadata( raw_data_feature_spec) def write_transform_artifacts(self, transform_fn, location): """ Write transform artifacts to the given location. Args: transform_fn: A transform_fn object. location: A location to write the artifacts. Returns: A PCollection of WriteTransformFn writing a TF transform graph. """ return ( transform_fn | 'Write Transform Artifacts' >> transform_fn_io.WriteTransformFn(location)) def _fail_on_non_default_windowing(self, pcoll: beam.PCollection): if not pcoll.windowing.is_default(): raise RuntimeError( "MLTransform only supports GlobalWindows when producing " "artifacts such as min, max, variance etc over the dataset." "Please use beam.WindowInto(beam.transforms.window.GlobalWindows()) " "to convert your PCollection to GlobalWindow.") def process_data_fn( self, inputs: dict[str, common_types.ConsistentTensorType] ) -> dict[str, common_types.ConsistentTensorType]: """ This method is used in the AnalyzeAndTransformDataset step. It applies the transforms to the `inputs` in sequential order on the columns provided for a given transform. Args: inputs: A dictionary of column names and data. Returns: A dictionary of column names and transformed data. """ outputs = inputs.copy() for transform in self.transforms: columns = transform.columns for col in columns: intermediate_result = transform(outputs[col], output_column_name=col) for key, value in intermediate_result.items(): outputs[key] = value return outputs def _get_transformed_data_schema( self, metadata: dataset_metadata.DatasetMetadata, ): schema = metadata._schema transformed_types = {} for feature in schema.feature: name = feature.name feature_type = feature.type if feature_type == schema_pb2.FeatureType.FLOAT: transformed_types[name] = Sequence[np.float32] elif feature_type == schema_pb2.FeatureType.INT: transformed_types[name] = Sequence[np.int64] # type: ignore[assignment] else: transformed_types[name] = Sequence[bytes] # type: ignore[assignment] return transformed_types def expand( self, raw_data: beam.PCollection[tft_process_handler_input_type] ) -> beam.PCollection[tft_process_handler_output_type]: """ This method also computes the required dataset metadata for the tft AnalyzeDataset/TransformDataset step. This method uses tensorflow_transform's Analyze step to produce the artifacts and Transform step to apply the transforms on the data. Artifacts are only produced if the artifact_mode is set to `produce`. If artifact_mode is set to `consume`, then the artifacts are read from the artifact_location, which was previously used to store the produced artifacts. """ if self.artifact_mode == ArtifactMode.PRODUCE: # If we are computing artifacts, we should fail for windows other than # default windowing since for example, for a fixed window, each window can # be treated as a separate dataset and we might need to compute artifacts # for each window. This is not supported yet. self._fail_on_non_default_windowing(raw_data) element_type = raw_data.element_type column_type_mapping = {} if (isinstance(element_type, RowTypeConstraint) or native_type_compatibility.match_is_named_tuple(element_type)): column_type_mapping = self._map_column_names_to_types( row_type=element_type) # convert Row or NamedTuple to Dict raw_data = ( raw_data | _ConvertNamedTupleToDict().with_output_types( dict[str, Union[tuple(column_type_mapping.values())]])) # type: ignore # AnalyzeAndTransformDataset raise type hint since this is # schema'd PCollection and the current output type would be a # custom type(NamedTuple) or a beam.Row type. else: column_type_mapping = self._map_column_names_to_types_from_transforms() # Add id so TFT can output id as output but as a no-op. raw_data_metadata = self.get_raw_data_metadata( input_types=column_type_mapping) # Write untransformed metadata to a file so that it can be re-used # during Transform step. metadata_io.write_metadata( metadata=raw_data_metadata, path=os.path.join(self.artifact_location, RAW_DATA_METADATA_DIR)) else: # Read the metadata from the artifact_location. if not FileSystems.exists(os.path.join( self.artifact_location, RAW_DATA_METADATA_DIR, SCHEMA_FILE)): raise FileNotFoundError( "Artifacts not found at location: %s when using " "read_artifact_location. Make sure you've run the pipeline with " "write_artifact_location using this artifact location before " "running with read_artifact_location set." % os.path.join(self.artifact_location, RAW_DATA_METADATA_DIR)) raw_data_metadata = metadata_io.read_metadata( os.path.join(self.artifact_location, RAW_DATA_METADATA_DIR)) element_type = raw_data.element_type if (isinstance(element_type, RowTypeConstraint) or native_type_compatibility.match_is_named_tuple(element_type)): # convert Row or NamedTuple to Dict column_type_mapping = self._map_column_names_to_types( row_type=element_type) raw_data = ( raw_data | _ConvertNamedTupleToDict().with_output_types( dict[str, Union[tuple(column_type_mapping.values())]])) # type: ignore feature_set = [feature.name for feature in raw_data_metadata.schema.feature] # TFT ignores columns in the input data that aren't explicitly defined # in the schema. This is because TFT operations # are designed to work with a predetermined schema. # To preserve these extra columns without disrupting TFT processing, # they are temporarily encoded as bytes and added to the PCollection with # a unique identifier data_coder = _DataCoder(exclude_columns=feature_set) data_with_encoded_columns = ( raw_data | "EncodeUnmodifiedColumns" >> beam.Map(lambda elem: data_coder.encode(elem))) # To maintain consistency by outputting numpy array all the time, # whether a scalar value or list or np array is passed as input, # we will convert scalar values to list values and TFT will ouput # numpy array all the time. data_list = data_with_encoded_columns | beam.ParDo( _ConvertScalarValuesToListValues()) with tft_beam.Context(temp_dir=self.artifact_location): data = (data_list, raw_data_metadata) if self.artifact_mode == ArtifactMode.PRODUCE: transform_fn = ( data | "AnalyzeDataset" >> tft_beam.AnalyzeDataset(self.process_data_fn)) # TODO: Remove the 'id' column from the transformed # dataset schema generated by TFT. self.write_transform_artifacts(transform_fn, self.artifact_location) else: transform_fn = ( data_list.pipeline | "ReadTransformFn" >> tft_beam.ReadTransformFn( self.artifact_location)) (transformed_dataset, transformed_metadata) = ( (data, transform_fn) | "TransformDataset" >> tft_beam.TransformDataset()) if isinstance(transformed_metadata, beam_metadata_io.BeamDatasetMetadata): self.transformed_schema = self._get_transformed_data_schema( metadata=transformed_metadata.dataset_metadata) else: self.transformed_schema = self._get_transformed_data_schema( transformed_metadata) # We will a pass a schema'd PCollection to the next step. # So we will use a RowTypeConstraint to create a schema'd PCollection. # this is needed since new columns are included in the # transformed_dataset. del self.transformed_schema[_TEMP_KEY] row_type = RowTypeConstraint.from_fields( list(self.transformed_schema.items())) # Decode the extra columns that were encoded as bytes. transformed_dataset = ( transformed_dataset | "DecodeUnmodifiedColumns" >> beam.Map(lambda x: data_coder.decode(x))) # The schema only contains the columns that are transformed. transformed_dataset = ( transformed_dataset | "ConvertToRowType" >> beam.Map(lambda x: beam.Row(**x)).with_output_types(row_type)) return transformed_dataset def with_exception_handling(self): raise NotImplementedError( "with_exception_handling with TensorFlow Transform-based MLTransform " "operations is not supported. To enable exception handling for those " "operations, please create a separate MLTransform instance") ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/handlers_test.py0000644000175100017510000005344315206113032023705 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import os import shutil import sys import tempfile import unittest import uuid from collections.abc import Sequence from typing import NamedTuple from typing import Union import numpy as np from parameterized import parameterized import apache_beam as beam from apache_beam.io.filesystems import FileSystems # pylint: disable=wrong-import-position, ungrouped-imports try: import tensorflow as tf from tensorflow_transform.tf_metadata import dataset_metadata from tensorflow_transform.tf_metadata import schema_utils from apache_beam.ml.transforms import handlers from apache_beam.ml.transforms import tft from apache_beam.ml.transforms.tft import TFTOperation from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to except ImportError: tft = None # type: ignore[assignment] if not tft: raise unittest.SkipTest('tensorflow_transform is not installed.') class _AddOperation(TFTOperation): def apply_transform(self, inputs, output_column_name, **kwargs): return {output_column_name: inputs + 1} class _MultiplyOperation(TFTOperation): def apply_transform(self, inputs, output_column_name, **kwargs): return {output_column_name: inputs * 10} class IntType(NamedTuple): x: int class ListIntType(NamedTuple): x: list[int] class NumpyType(NamedTuple): x: np.int64 class TFTProcessHandlerTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) @parameterized.expand([ ({ 'x': 1, 'y': 2 }, ['x'], { 'x': 20, 'y': 2 }), ({ 'x': 1, 'y': 2 }, ['x', 'y'], { 'x': 20, 'y': 30 }), ]) def test_tft_operation_preprocessing_fn( self, inputs, columns, expected_result): add_fn = _AddOperation(columns=columns) mul_fn = _MultiplyOperation(columns=columns) process_handler = handlers.TFTProcessHandler( transforms=[add_fn, mul_fn], artifact_location=self.artifact_location) actual_result = process_handler.process_data_fn(inputs) self.assertDictEqual(actual_result, expected_result) def test_input_type_from_schema_named_tuple_pcoll(self): data = [{'x': 1}] with beam.Pipeline() as p: data = ( p | beam.Create(data) | beam.Map(lambda x: IntType(**x)).with_output_types(IntType)) element_type = data.element_type process_handler = handlers.TFTProcessHandler( artifact_location=self.artifact_location) inferred_input_type = process_handler._map_column_names_to_types( element_type) expected_input_type = dict(x=list[int]) self.assertEqual(inferred_input_type, expected_input_type) def test_input_type_from_schema_named_tuple_pcoll_list(self): data = [{'x': [1, 2, 3]}, {'x': [4, 5, 6]}] with beam.Pipeline() as p: data = ( p | beam.Create(data) | beam.Map(lambda x: ListIntType(**x)).with_output_types(ListIntType)) element_type = data.element_type process_handler = handlers.TFTProcessHandler( artifact_location=self.artifact_location) inferred_input_type = process_handler._map_column_names_to_types( element_type) expected_input_type = dict(x=list[int]) self.assertEqual(inferred_input_type, expected_input_type) def test_input_type_from_row_type_pcoll(self): data = [{'x': 1}] with beam.Pipeline() as p: data = ( p | beam.Create(data) | beam.Map(lambda ele: beam.Row(x=int(ele['x'])))) element_type = data.element_type process_handler = handlers.TFTProcessHandler( artifact_location=self.artifact_location) inferred_input_type = process_handler._map_column_names_to_types( element_type) expected_input_type = dict(x=list[int]) self.assertEqual(inferred_input_type, expected_input_type) def test_input_type_from_row_type_pcoll_list(self): data = [{'x': [1, 2, 3]}, {'x': [4, 5, 6]}] with beam.Pipeline() as p: data = ( p | beam.Create(data) | beam.Map(lambda ele: beam.Row(x=list(ele['x']))).with_output_types( beam.row_type.RowTypeConstraint.from_fields([('x', list[int])]))) element_type = data.element_type process_handler = handlers.TFTProcessHandler( artifact_location=self.artifact_location) inferred_input_type = process_handler._map_column_names_to_types( element_type) expected_input_type = dict(x=list[int]) self.assertEqual(inferred_input_type, expected_input_type) def test_input_type_from_named_tuple_pcoll_numpy(self): np_data = [{ 'x': np.array([1, 2, 3], dtype=np.int64) }, { 'x': np.array([4, 5, 6], dtype=np.int64) }] with beam.Pipeline() as p: data = ( p | beam.Create(np_data) | beam.Map(lambda x: NumpyType(**x)).with_output_types(NumpyType)) element_type = data.element_type process_handler = handlers.TFTProcessHandler( artifact_location=self.artifact_location) inferred_input_type = process_handler._map_column_names_to_types( element_type) expected_type = dict(x=np.int64) self.assertEqual(inferred_input_type, expected_type) def test_tensorflow_raw_data_metadata_primitive_types(self): input_types = dict(x=int, y=float, k=bytes, l=str) process_handler = handlers.TFTProcessHandler( artifact_location=self.artifact_location) for col_name, typ in input_types.items(): feature_spec = process_handler._get_raw_data_feature_spec_per_column( typ=typ, col_name=col_name) self.assertEqual( handlers._default_type_to_tensor_type_map[typ], feature_spec.dtype) self.assertIsInstance(feature_spec, tf.io.VarLenFeature) def test_tensorflow_raw_data_metadata_primitive_types_in_containers(self): input_types = dict([("x", list[int]), ("y", list[float]), ("k", list[bytes]), ("l", list[str])]) process_handler = handlers.TFTProcessHandler( artifact_location=self.artifact_location) for col_name, typ in input_types.items(): feature_spec = process_handler._get_raw_data_feature_spec_per_column( typ=typ, col_name=col_name) self.assertIsInstance(feature_spec, tf.io.VarLenFeature) @unittest.skipIf(sys.version_info < (3, 9), "not supported in python<3.9") def test_tensorflow_raw_data_metadata_primitive_native_container_types(self): input_types = dict([("x", list[int]), ("y", list[float]), ("k", list[bytes]), ("l", list[str])]) process_handler = handlers.TFTProcessHandler( artifact_location=self.artifact_location) for col_name, typ in input_types.items(): feature_spec = process_handler._get_raw_data_feature_spec_per_column( typ=typ, col_name=col_name) self.assertIsInstance(feature_spec, tf.io.VarLenFeature) def test_tensorflow_raw_data_metadata_numpy_types(self): input_types = dict(x=np.int64, y=np.float32, z=list[np.int64]) process_handler = handlers.TFTProcessHandler( artifact_location=self.artifact_location) for col_name, typ in input_types.items(): feature_spec = process_handler._get_raw_data_feature_spec_per_column( typ=typ, col_name=col_name) self.assertIsInstance(feature_spec, tf.io.VarLenFeature) def test_tensorflow_raw_data_metadata_union_type_in_single_column(self): input_types = dict(x=Union[int, float]) process_handler = handlers.TFTProcessHandler( artifact_location=self.artifact_location) with self.assertRaises(TypeError): for col_name, typ in input_types.items(): _ = process_handler._get_raw_data_feature_spec_per_column( typ=typ, col_name=col_name) def test_tensorflow_raw_data_metadata_dtypes(self): input_types = dict(x=np.int32, y=np.float64) expected_dtype = dict(x=np.int64, y=np.float32) process_handler = handlers.TFTProcessHandler( artifact_location=self.artifact_location) for col_name, typ in input_types.items(): feature_spec = process_handler._get_raw_data_feature_spec_per_column( typ=typ, col_name=col_name) self.assertEqual(expected_dtype[col_name], feature_spec.dtype) def test_tft_process_handler_default_transform_types(self): transforms = [ tft.ScaleTo01(columns=['x']), tft.ScaleToZScore(columns=['y']), tft.Bucketize(columns=['z'], num_buckets=2), tft.ComputeAndApplyVocabulary(columns=['w']) ] process_handler = handlers.TFTProcessHandler( transforms=transforms, artifact_location=self.artifact_location) column_type_mapping = ( process_handler._map_column_names_to_types_from_transforms()) expected_column_type_mapping = { 'x': float, 'y': float, 'z': float, 'w': str } self.assertDictEqual(column_type_mapping, expected_column_type_mapping) expected_tft_raw_data_feature_spec = { 'x': tf.io.VarLenFeature(tf.float32), 'y': tf.io.VarLenFeature(tf.float32), 'z': tf.io.VarLenFeature(tf.float32), 'w': tf.io.VarLenFeature(tf.string) } actual_tft_raw_data_feature_spec = ( process_handler.get_raw_data_feature_spec(column_type_mapping)) self.assertDictEqual( actual_tft_raw_data_feature_spec, expected_tft_raw_data_feature_spec) def test_tft_process_handler_transformed_data_schema(self): process_handler = handlers.TFTProcessHandler( artifact_location=self.artifact_location) raw_data_feature_spec = { 'x': tf.io.VarLenFeature(tf.float32), 'y': tf.io.VarLenFeature(tf.float32), 'z': tf.io.VarLenFeature(tf.string), } raw_data_metadata = dataset_metadata.DatasetMetadata( schema_utils.schema_from_feature_spec(raw_data_feature_spec)) expected_transformed_data_schema = { 'x': Sequence[np.float32], 'y': Sequence[np.float32], 'z': Sequence[bytes] } actual_transformed_data_schema = ( process_handler._get_transformed_data_schema(raw_data_metadata)) self.assertDictEqual( actual_transformed_data_schema, expected_transformed_data_schema) def test_tft_process_handler_verify_artifacts(self): with beam.Pipeline() as p: raw_data = ( p | beam.Create([{ 'x': np.array([1, 3]) }, { 'x': np.array([4, 6]) }])) process_handler = handlers.TFTProcessHandler( transforms=[tft.ScaleTo01(columns=['x'])], artifact_location=self.artifact_location, ) _ = raw_data | process_handler self.assertTrue( FileSystems.exists( # To check the gcs directory with FileSystems, the dir path must # end with / os.path.join( self.artifact_location, handlers.RAW_DATA_METADATA_DIR + '/'))) self.assertTrue( FileSystems.exists( os.path.join( self.artifact_location, handlers.RAW_DATA_METADATA_DIR, handlers.SCHEMA_FILE))) with beam.Pipeline() as p: raw_data = (p | beam.Create([{'x': np.array([2, 5])}])) process_handler = handlers.TFTProcessHandler( artifact_location=self.artifact_location, artifact_mode='consume') transformed_data = raw_data | process_handler transformed_data |= beam.Map(lambda x: x.x) # the previous min is 1 and max is 6. So this should scale by (1, 6) assert_that( transformed_data, equal_to([np.array([0.2, 0.8], dtype=np.float32)], equals_fn=np.array_equal)) def test_tft_process_handler_unused_column(self): data = [ { 'x': [5, 1], 'y': 2 }, { 'x': [8, 4], 'y': 5 }, { 'x': [9, 5], 'y': 6 }, { 'x': [10, 6], 'y': 7 }, { 'x': [11, 7], 'y': 8 }, { 'x': [12, 8], 'y': 9 }, { 'x': [13, 9], 'y': 10 }, { 'x': [14, 10], 'y': 11 }, { 'x': [15, 11], 'y': 12 }, { 'x': [16, 12], 'y': 13 }, { 'x': [17, 13], 'y': 14 }, { 'x': [18, 14], 'y': 15 }, { 'x': [19, 15], 'y': 16 }, { 'x': [20, 16], 'y': 17 }, { 'x': [21, 17], 'y': 18 }, { 'x': [22, 18], 'y': 19 }, { 'x': [23, 19], 'y': 20 }, { 'x': [24, 20], 'y': 21 }, { 'x': [25, 21], 'y': 22 }, ] # pylint: disable=line-too-long expected_data = [ beam.Row( x=np.array([0.16666667, 0.], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=2), beam.Row( x=np.array([0.29166666, 0.125], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=5), beam.Row( x=np.array([0.33333334, 0.16666667], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=6), beam.Row( x=np.array([0.375, 0.20833333], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=7), beam.Row( x=np.array([0.41666666, 0.25], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=8), beam.Row( x=np.array([0.45833334, 0.29166666], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=9), beam.Row( x=np.array([0.5, 0.33333334], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=10), beam.Row( x=np.array([0.5416667, 0.375], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=11), beam.Row( x=np.array([0.5833333, 0.41666666], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=12), beam.Row( x=np.array([0.625, 0.45833334], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=13), beam.Row( x=np.array([0.6666667, 0.5], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=14), beam.Row( x=np.array([0.7083333, 0.5416667], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=15), beam.Row( x=np.array([0.75, 0.5833333], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=16), beam.Row( x=np.array([0.7916667, 0.625], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=17), beam.Row( x=np.array([0.8333333, 0.6666667], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=18), beam.Row( x=np.array([0.875, 0.7083333], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=19), beam.Row( x=np.array([0.9166667, 0.75], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=20), beam.Row( x=np.array([0.9583333, 0.7916667], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=21), beam.Row( x=np.array([1., 0.8333333], dtype=np.float32), x_max=np.array([25.], dtype=np.float32), x_min=np.array([1.], dtype=np.float32), y=22), ] # pylint: enable=line-too-long expected_data_x = [row.x for row in expected_data] expected_data_y = [row.y for row in expected_data] scale_to_0_1_fn = tft.ScaleTo01(columns=['x']) with beam.Pipeline() as p: raw_data = p | beam.Create(data) process_handler = handlers.TFTProcessHandler( transforms=[scale_to_0_1_fn], artifact_location=self.artifact_location, ) transformed_pcoll = raw_data | process_handler transformed_pcoll_x = transformed_pcoll | beam.Map(lambda x: x.x) transformed_pcoll_y = transformed_pcoll | beam.Map(lambda x: x.y) assert_that( transformed_pcoll_x, equal_to(expected_data_x, equals_fn=np.array_equal), label='transformed data') assert_that( transformed_pcoll_y, equal_to(expected_data_y, equals_fn=np.array_equal), label='unused column') def test_consume_mode_with_extra_columns_in_the_input(self): with beam.Pipeline() as p: raw_data = ( p | beam.Create([{ 'x': np.array([1, 3]) }, { 'x': np.array([4, 6]) }])) process_handler = handlers.TFTProcessHandler( transforms=[tft.ScaleTo01(columns=['x'])], artifact_location=self.artifact_location, ) _ = raw_data | process_handler test_data = [{ 'x': np.array([2, 5]), 'y': np.array([1, 2]), 'z': 'fake_string' }, { 'x': np.array([1, 10]), 'y': np.array([2, 3]), 'z': 'fake_string2' }] expected_test_data = [{ 'x': np.array([0.2, 0.8], dtype=np.float32), 'y': np.array([1, 2]), 'z': 'fake_string' }, { 'x': np.array([0., 1.8], dtype=np.float32), 'y': np.array([2, 3]), 'z': 'fake_string2' }] expected_test_data_x = [row['x'] for row in expected_test_data] expected_test_data_y = [row['y'] for row in expected_test_data] expected_test_data_z = [row['z'] for row in expected_test_data] with beam.Pipeline() as p: raw_data = (p | beam.Create(test_data)) process_handler = handlers.TFTProcessHandler( artifact_location=self.artifact_location, artifact_mode='consume') transformed_data = raw_data | process_handler transformed_data_x = transformed_data | beam.Map(lambda x: x.x) transformed_data_y = transformed_data | beam.Map(lambda x: x.y) transformed_data_z = transformed_data | beam.Map(lambda x: x.z) assert_that( transformed_data_x, equal_to(expected_test_data_x, equals_fn=np.array_equal), label='transformed data') assert_that( transformed_data_y, equal_to(expected_test_data_y, equals_fn=np.array_equal), label='unused column: y') assert_that( transformed_data_z, equal_to(expected_test_data_z, equals_fn=np.array_equal), label='unused column: z') def test_handler_with_same_input_elements(self): with beam.Pipeline() as p: data = [ { 'x': 'I' }, { 'x': 'love' }, { 'x': 'Beam' }, { 'x': 'Beam' }, { 'x': 'is' }, { 'x': 'awesome' }, ] raw_data = (p | beam.Create(data)) process_handler = handlers.TFTProcessHandler( transforms=[tft.ComputeAndApplyVocabulary(columns=['x'])], artifact_location=self.artifact_location, ) transformed_data = raw_data | process_handler expected_data = [ beam.Row(x=np.array([4])), beam.Row(x=np.array([1])), beam.Row(x=np.array([0])), beam.Row(x=np.array([0])), beam.Row(x=np.array([2])), beam.Row(x=np.array([3])), ] expected_data_x = [row.x for row in expected_data] actual_data_x = transformed_data | beam.Map(lambda x: x.x) assert_that( actual_data_x, equal_to(expected_data_x, equals_fn=np.array_equal), label='transformed data') class TFTProcessHandlerTestWithGCSLocation(TFTProcessHandlerTest): def setUp(self) -> None: self.artifact_location = self.gcs_artifact_location = os.path.join( 'gs://temp-storage-for-perf-tests/tft_handler', uuid.uuid4().hex) def tearDown(self): pass if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/tft.py0000644000175100017510000006557615206113032021655 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ This module defines a set of data processing transforms that can be used to perform common data transformations on a dataset. These transforms are implemented using the TensorFlow Transform (TFT) library. The transforms in this module are intended to be used in conjunction with the MLTransform class, which provides a convenient interface for applying a sequence of data processing transforms to a dataset. See the documentation for MLTransform for more details. Note: The data processing transforms defined in this module don't perform the transformation immediately. Instead, it returns a configured operation object, which encapsulates the details of the transformation. The actual computation takes place later in the Apache Beam pipeline, after all transformations are set up and the pipeline is run. """ # pytype: skip-file import logging from collections.abc import Iterable from typing import Any from typing import Optional from typing import Union import tensorflow as tf import tensorflow_transform as tft from tensorflow_transform import common_types import apache_beam as beam from apache_beam.ml.transforms.base import BaseOperation __all__ = [ 'ComputeAndApplyVocabulary', 'ScaleToZScore', 'ScaleTo01', 'ScaleToGaussian', 'ApplyBuckets', 'ApplyBucketsWithInterpolation', 'Bucketize', 'TFIDF', 'TFTOperation', 'ScaleByMinMax', 'NGrams', 'BagOfWords', 'HashStrings', 'DeduplicateTensorPerRow', ] # Register the expected input types for each operation # this will be used to determine schema for the tft.AnalyzeDataset _EXPECTED_TYPES: dict[str, Union[int, str, float]] = {} _LOGGER = logging.getLogger(__name__) def register_input_dtype(type): def wrapper(fn): _EXPECTED_TYPES[fn.__name__] = type return fn return wrapper # TODO: https://github.com/apache/beam/pull/29016 # Add support for outputting artifacts to a text file in human readable form. class TFTOperation(BaseOperation[common_types.TensorType, common_types.TensorType]): def __init__(self, columns: list[str]) -> None: """ Base Operation class for TFT data processing transformations. Processing logic for the transformation is defined in the apply_transform() method. If you have a custom transformation that is not supported by the existing transforms, you can extend this class and implement the apply_transform() method. Args: columns: List of column names to apply the transformation. """ super().__init__(columns) if not columns: raise RuntimeError( "Columns are not specified. Please specify the column for the " " op %s" % self.__class__.__name__) def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: from apache_beam.ml.transforms.handlers import TFTProcessHandler params = {} artifact_location = kwargs.get('artifact_location') if not artifact_location: raise RuntimeError( "artifact_location is not specified. Please specify the " "artifact_location for the op %s" % self.__class__.__name__) artifact_mode = kwargs.get('artifact_mode') if artifact_mode: params['artifact_mode'] = artifact_mode return TFTProcessHandler(artifact_location=artifact_location, **params) @tf.function def _split_string_with_delimiter(self, data, delimiter): """ only applicable to string columns. """ data = tf.sparse.to_dense(data) # this method acts differently compared to tf.strings.split # this will split the string based on multiple delimiters while # the latter will split the string based on a single delimiter. fn = lambda data: tf.compat.v1.string_split( data, delimiter, result_type='RaggedTensor') # tf.compat.v1.string_split works on a single string. Use tf.map_fn # to apply the function on each element of the input data. data = tf.map_fn( fn, data, fn_output_signature=tf.RaggedTensorSpec( tf.TensorShape([None, None]), tf.string)) data = data.values.to_sparse() # the columns of the sparse tensor are suffixed with $indices, $values # related to sparse tensor. Create a new sparse tensor by extracting # the indices, values and dense_shape from the original sparse tensor # to preserve the original column name. data = tf.sparse.SparseTensor( indices=data.indices, values=data.values, dense_shape=data.dense_shape) # for list of string, batch dimensions becomes inverted after tf.map_fn, # transpose the data to get the original shape. if tf.shape(data)[1] == 1: data = tf.sparse.transpose(data) return data @register_input_dtype(str) class ComputeAndApplyVocabulary(TFTOperation): def __init__( self, columns: list[str], split_string_by_delimiter: Optional[str] = None, *, default_value: Any = -1, top_k: Optional[int] = None, frequency_threshold: Optional[int] = None, num_oov_buckets: int = 0, vocab_filename: Optional[str] = None, name: Optional[str] = None): """ This function computes the vocabulary for the given columns of incoming data. The transformation converts the input values to indices of the vocabulary. Args: columns: List of column names to apply the transformation. split_string_by_delimiter: (Optional) A string that specifies the delimiter to split strings. default_value: (Optional) The value to use for out-of-vocabulary values. top_k: (Optional) The number of most frequent tokens to keep. frequency_threshold: (Optional) Limit the generated vocabulary only to elements whose absolute frequency is >= to the supplied threshold. If set to None, the full vocabulary is generated. num_oov_buckets: Any lookup of an out-of-vocabulary token will return a bucket ID based on its hash if `num_oov_buckets` is greater than zero. Otherwise it is assigned the `default_value`. vocab_filename: The file name for the vocabulary file. The vocab file will be suffixed with the column name. NOTE in order to make your pipelines resilient to implementation details please set `vocab_filename` when you are using the vocab_filename on a downstream component. """ super().__init__(columns) self._default_value = default_value self._top_k = top_k self._frequency_threshold = frequency_threshold self._num_oov_buckets = num_oov_buckets self._vocab_filename = vocab_filename self._name = name self.split_string_by_delimiter = split_string_by_delimiter def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> dict[str, common_types.TensorType]: if self.split_string_by_delimiter: data = self._split_string_with_delimiter( data, self.split_string_by_delimiter) vocab_filename = self._vocab_filename if vocab_filename: vocab_filename = vocab_filename + f'_{output_column_name}' return { output_column_name: tft.compute_and_apply_vocabulary( x=data, default_value=self._default_value, top_k=self._top_k, frequency_threshold=self._frequency_threshold, num_oov_buckets=self._num_oov_buckets, vocab_filename=vocab_filename, name=self._name) } @register_input_dtype(float) class ScaleToZScore(TFTOperation): def __init__( self, columns: list[str], *, elementwise: bool = False, name: Optional[str] = None): """ This function performs a scaling transformation on the specified columns of the incoming data. It processes the input data such that it's normalized to have a mean of 0 and a variance of 1. The transformation achieves this by subtracting the mean from the input data and then dividing it by the square root of the variance. Args: columns: A list of column names to apply the transformation on. elementwise: If True, the transformation is applied elementwise. Otherwise, the transformation is applied on the entire column. name: A name for the operation (optional). scale_to_z_score also outputs additional artifacts. The artifacts are mean, which is the mean value in the column, and var, which is the variance in the column. The artifacts are stored in the column named with the suffix _mean and _var respectively. """ super().__init__(columns) self.elementwise = elementwise self.name = name def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> dict[str, common_types.TensorType]: output_dict = { output_column_name: tft.scale_to_z_score( x=data, elementwise=self.elementwise, name=self.name) } return output_dict @register_input_dtype(float) class ScaleTo01(TFTOperation): def __init__( self, columns: list[str], elementwise: bool = False, name: Optional[str] = None): """ This function applies a scaling transformation on the given columns of incoming data. The transformation scales the input values to the range [0, 1] by dividing each value by the maximum value in the column. Args: columns: A list of column names to apply the transformation on. elementwise: If True, the transformation is applied elementwise. Otherwise, the transformation is applied on the entire column. name: A name for the operation (optional). ScaleTo01 also outputs additional artifacts. The artifacts are max, which is the maximum value in the column, and min, which is the minimum value in the column. The artifacts are stored in the column named with the suffix _min and _max respectively. """ super().__init__(columns) self.elementwise = elementwise self.name = name def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> dict[str, common_types.TensorType]: output = tft.scale_to_0_1( x=data, elementwise=self.elementwise, name=self.name) output_dict = {output_column_name: output} return output_dict @register_input_dtype(float) class ScaleToGaussian(TFTOperation): def __init__( self, columns: list[str], elementwise: bool = False, name: Optional[str] = None): """ This operation scales the given input column values to an approximately normal distribution with mean 0 and variance of 1. The Gaussian transformation is only applied if the column has long tails; otherwise, the transformation is the same as normalizing to z scores. For more information, see: https://www.tensorflow.org/tfx/transform/api_docs/python/tft/scale_to_gaussian Args: columns: A list of column names to apply the transformation on. elementwise: If True, the transformation is applied elementwise. Otherwise, the transformation is applied on the entire column. name: A name for the operation (optional). """ super().__init__(columns) self.elementwise = elementwise self.name = name def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> dict[str, common_types.TensorType]: output_dict = { output_column_name: tft.scale_to_gaussian( x=data, elementwise=self.elementwise, name=self.name) } return output_dict @register_input_dtype(float) class ApplyBuckets(TFTOperation): def __init__( self, columns: list[str], bucket_boundaries: Iterable[Union[int, float]], name: Optional[str] = None): """ This functions is used to map the element to a positive index i for which `bucket_boundaries[i-1] <= element < bucket_boundaries[i]`, if it exists. If `input < bucket_boundaries[0]`, then element is mapped to 0. If `element >= bucket_boundaries[-1]`, then element is mapped to len(bucket_boundaries). NaNs are mapped to len(bucket_boundaries). Args: columns: A list of column names to apply the transformation on. bucket_boundaries: An iterable of ints or floats representing the bucket boundaries. Must be sorted in ascending order. name: (Optional) A string that specifies the name of the operation. """ super().__init__(columns) self.bucket_boundaries = [bucket_boundaries] self.name = name def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> dict[str, common_types.TensorType]: output = { output_column_name: tft.apply_buckets( x=data, bucket_boundaries=self.bucket_boundaries, name=self.name) } return output @register_input_dtype(float) class ApplyBucketsWithInterpolation(TFTOperation): def __init__( self, columns: list[str], bucket_boundaries: Iterable[Union[int, float]], name: Optional[str] = None): """Interpolates values within the provided buckets and then normalizes to [0, 1]. Input values are bucketized based on the provided boundaries such that the input is mapped to a positive index i for which `bucket_boundaries[i-1] <= element < bucket_boundaries[i]`, if it exists. The values are then normalized to the range [0,1] within the bucket, with NaN values being mapped to 0.5. For more information, see: https://www.tensorflow.org/tfx/transform/api_docs/python/tft/apply_buckets_with_interpolation Args: columns: A list of column names to apply the transformation on. bucket_boundaries: An iterable of ints or floats representing the bucket boundaries sorted in ascending order. name: (Optional) A string that specifies the name of the operation. """ super().__init__(columns) self.bucket_boundaries = [bucket_boundaries] self.name = name def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> dict[str, common_types.TensorType]: output = { output_column_name: tft.apply_buckets_with_interpolation( x=data, bucket_boundaries=self.bucket_boundaries, name=self.name) } return output @register_input_dtype(float) class Bucketize(TFTOperation): def __init__( self, columns: list[str], num_buckets: int, *, epsilon: Optional[float] = None, elementwise: bool = False, name: Optional[str] = None): """ This function applies a bucketizing transformation on the given columns of incoming data. The transformation splits the input data range into a set of consecutive bins/buckets, and converts the input values to bucket IDs (integers) where each ID corresponds to a particular bin. Args: columns: List of column names to apply the transformation. num_buckets: Number of buckets to be created. epsilon: (Optional) A float number that specifies the error tolerance when computing quantiles, so that we guarantee that any value x will have a quantile q such that x is in the interval [q - epsilon, q + epsilon] (or the symmetric interval for even num_buckets). Must be greater than 0.0. elementwise: (Optional) A boolean that specifies whether the quantiles should be computed on an element-wise basis. If False, the quantiles are computed globally. name: (Optional) A string that specifies the name of the operation. """ super().__init__(columns) self.num_buckets = num_buckets self.epsilon = epsilon self.elementwise = elementwise self.name = name def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> dict[str, common_types.TensorType]: output = { output_column_name: tft.bucketize( x=data, num_buckets=self.num_buckets, epsilon=self.epsilon, elementwise=self.elementwise, name=self.name) } return output @register_input_dtype(float) class TFIDF(TFTOperation): def __init__( self, columns: list[str], vocab_size: Optional[int] = None, smooth: bool = True, name: Optional[str] = None, ): """ This function applies a tf-idf transformation on the given columns of incoming data. TFIDF outputs two artifacts for each column: the vocabulary index and the tfidf weight. The vocabulary index is a mapping from the original vocabulary to the new vocabulary. The tfidf weight is a mapping from the original vocabulary to the tfidf score. Input passed to the TFIDF is not modified and used to calculate the required artifacts. Args: columns: List of column names to apply the transformation. vocab_size: (Optional) An integer that specifies the size of the vocabulary. Defaults to None. If vocab_size is None, then the size of the vocabulary is determined by `tft.get_num_buckets_for_transformed_feature`. smooth: (Optional) A boolean that specifies whether to apply smoothing to the tf-idf score. Defaults to True. name: (Optional) A string that specifies the name of the operation. """ super().__init__(columns) self.vocab_size = vocab_size self.smooth = smooth self.name = name self.tfidf_weight = None def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> common_types.TensorType: if self.vocab_size is None: try: _LOGGER.info( 'vocab_size is not specified. Trying to infer vocab_size ' 'from the input data using ' 'tft.get_num_buckets_for_transformed_feature.') vocab_size = tft.get_num_buckets_for_transformed_feature(data) except RuntimeError: raise RuntimeError( 'vocab_size is not specified. Tried to infer vocab_size from the ' 'input data using tft.get_num_buckets_for_transformed_feature, but ' 'failed. Please specify vocab_size explicitly.') else: vocab_size = self.vocab_size vocab_index, tfidf_weight = tft.tfidf( data, vocab_size, self.smooth, self.name ) output = { output_column_name + '_vocab_index': vocab_index, output_column_name + '_tfidf_weight': tfidf_weight } return output @register_input_dtype(float) class ScaleByMinMax(TFTOperation): def __init__( self, columns: list[str], min_value: float = 0.0, max_value: float = 1.0, name: Optional[str] = None): """ This function applies a scaling transformation on the given columns of incoming data. The transformation scales the input values to the range [min_value, max_value]. Args: columns: A list of column names to apply the transformation on. min_value: The minimum value of the output range. max_value: The maximum value of the output range. name: A name for the operation (optional). """ super().__init__(columns) self.min_value = min_value self.max_value = max_value self.name = name if self.max_value <= self.min_value: raise ValueError('max_value must be greater than min_value') def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> common_types.TensorType: output = tft.scale_by_min_max( x=data, output_min=self.min_value, output_max=self.max_value) return {output_column_name: output} @register_input_dtype(str) class NGrams(TFTOperation): def __init__( self, columns: list[str], split_string_by_delimiter: Optional[str] = None, *, ngram_range: tuple[int, int] = (1, 1), ngrams_separator: Optional[str] = None, name: Optional[str] = None): """ An n-gram is a contiguous sequence of n items from a given sample of text or speech. This operation applies an n-gram transformation to specified columns of incoming data, splitting the input data into a set of consecutive n-grams. Args: columns: A list of column names to apply the transformation on. split_string_by_delimiter: (Optional) A string that specifies the delimiter to split the input strings before computing ngrams. ngram_range: A tuple of integers(inclusive) specifying the range of n-gram sizes. ngrams_separator: A string that will be inserted between each ngram. name: A name for the operation (optional). """ super().__init__(columns) self.ngram_range = ngram_range self.ngrams_separator = ngrams_separator self.name = name self.split_string_by_delimiter = split_string_by_delimiter if ngram_range != (1, 1) and not ngrams_separator: raise ValueError( 'ngrams_separator must be specified when ngram_range is not (1, 1)') def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> dict[str, common_types.TensorType]: if self.split_string_by_delimiter: data = self._split_string_with_delimiter( data, self.split_string_by_delimiter) output = tft.ngrams(data, self.ngram_range, self.ngrams_separator) return {output_column_name: output} @register_input_dtype(str) class BagOfWords(TFTOperation): def __init__( self, columns: list[str], split_string_by_delimiter: Optional[str] = None, *, ngram_range: tuple[int, int] = (1, 1), ngrams_separator: Optional[str] = None, compute_word_count: bool = False, key_vocab_filename: Optional[str] = None, name: Optional[str] = None, ): """ Bag of words contains the unique words present in the input text. This operation applies a bag of words transformation to specified columns of incoming data. Also, the transformation accepts a Tuple of integers specifying the range of n-gram sizes. The transformation splits the input data into a set of consecutive n-grams if ngram_range is specified. The n-grams are then converted to a bag of words. Also, you can specify a seperator string that will be inserted between each ngram. Args: columns: A list of column names to apply the transformation on. split_string_by_delimiter: (Optional) A string that specifies the delimiter to split the input strings before computing ngrams. ngram_range: A tuple of integers(inclusive) specifying the range of n-gram sizes. seperator: A string that will be inserted between each ngram. compute_word_count: A boolean that specifies whether to compute the unique word count over the entire dataset. Defaults to False. key_vocab_filename: The file name for the key vocabulary file when compute_word_count is True. If empty, a file name will be chosen based on the current scope. If provided, the vocab file will be suffixed with the column name. name: A name for the operation (optional). Note that original order of the input may not be preserved. """ self.columns = columns self.ngram_range = ngram_range self.ngrams_separator = ngrams_separator self.name = name self.split_string_by_delimiter = split_string_by_delimiter self.key_vocab_filename = key_vocab_filename if compute_word_count: self.compute_word_count_fn = count_unique_words else: self.compute_word_count_fn = lambda *args, **kwargs: None if ngram_range != (1, 1) and not ngrams_separator: raise ValueError( 'ngrams_separator must be specified when ngram_range is not (1, 1)') def apply_transform(self, data: tf.SparseTensor, output_col_name: str): if self.split_string_by_delimiter: data = self._split_string_with_delimiter( data, self.split_string_by_delimiter) output = tft.bag_of_words( data, self.ngram_range, self.ngrams_separator, self.name) # word counts are written to the file only if compute_word_count is True key_vocab_filename = self.key_vocab_filename if key_vocab_filename: key_vocab_filename = key_vocab_filename + f'_{output_col_name}' self.compute_word_count_fn(data, key_vocab_filename) return {output_col_name: output} def count_unique_words( data: tf.SparseTensor, output_vocab_name: Optional[str]) -> None: tft.count_per_key(data, key_vocabulary_filename=output_vocab_name) @register_input_dtype(str) class HashStrings(TFTOperation): def __init__( self, columns: list[str], hash_buckets: int, key: Optional[tuple[int, int]] = None, name: Optional[str] = None): '''Hashes strings into the provided number of buckets. Args: columns: A list of the column names to apply the transformation on. hash_buckets: the number of buckets to hash the strings into. key: optional. An array of two Python `uint64`. If passed, output will be a deterministic function of `strings` and `key`. Note that hashing will be slower if this value is specified. name: optional. A name for this operation. Raises: ValueError if `hash_buckets` is not a positive and non-zero integer. ''' self.hash_buckets = hash_buckets self.key = key self.name = name if hash_buckets < 1: raise ValueError( 'number of hash buckets must be positive, got ', hash_buckets) super().__init__(columns) def apply_transform( self, data: common_types.TensorType, output_col_name: str) -> dict[str, common_types.TensorType]: output_dict = { output_col_name: tft.hash_strings( strings=data, hash_buckets=self.hash_buckets, key=self.key, name=self.name) } return output_dict @register_input_dtype(str) class DeduplicateTensorPerRow(TFTOperation): def __init__(self, columns: list[str], name: Optional[str] = None): """ Deduplicates each row (0th dimension) of the provided tensor. Args: columns: A list of the columns to apply the transformation on. name: optional. A name for this operation. """ self.name = name super().__init__(columns) def apply_transform( self, data: common_types.TensorType, output_col_name: str) -> dict[str, common_types.TensorType]: output_dict = { output_col_name: tft.deduplicate_tensor_per_row( input_tensor=data, name=self.name) } return output_dict ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/tft_test.py0000644000175100017510000011062315206113032022674 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import os import shutil import tempfile import unittest import numpy as np from parameterized import parameterized import apache_beam as beam from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.ml.transforms import base from apache_beam.ml.transforms import tft except ImportError: tft = None # type: ignore[assignment] if not tft: raise unittest.SkipTest('tensorflow_transform is not installed.') class ScaleZScoreTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) def test_z_score(self): data = [ { 'x': 1 }, { 'x': 2 }, { 'x': 3 }, { 'x': 4 }, { 'x': 5 }, { 'x': 6 }, ] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ScaleToZScore(columns=['x']))) expected_data = [ np.array([-1.46385], dtype=np.float32), np.array([-0.87831], dtype=np.float32), np.array([-0.29277], dtype=np.float32), np.array([0.29277], dtype=np.float32), np.array([0.87831], dtype=np.float32), np.array([1.46385], dtype=np.float32), ] actual_data = (result | beam.Map(lambda x: x.x)) assert_that( actual_data, equal_to(expected_data, equals_fn=np.array_equal)) def test_z_score_list_data(self): list_data = [{'x': [1, 2, 3]}, {'x': [4, 5, 6]}] with beam.Pipeline() as p: list_result = ( p | "listCreate" >> beam.Create(list_data) | "listMLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ScaleToZScore(columns=['x']))) expected_data = [ np.array([-1.46385, -0.87831, -0.29277], dtype=np.float32), np.array([0.29277, 0.87831, 1.46385], dtype=np.float32) ] actual_data = (list_result | beam.Map(lambda x: x.x)) assert_that( actual_data, equal_to(expected_data, equals_fn=np.array_equal)) class ScaleTo01Test(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) def test_ScaleTo01_list(self): list_data = [{'x': [1, 2, 3]}, {'x': [4, 5, 6]}] with beam.Pipeline() as p: list_result = ( p | "listCreate" >> beam.Create(list_data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ScaleTo01(columns=['x']))) expected_output = [ np.array([0, 0.2, 0.4], dtype=np.float32), np.array([0.6, 0.8, 1], dtype=np.float32) ] actual_output = (list_result | beam.Map(lambda x: x.x)) assert_that( actual_output, equal_to(expected_output, equals_fn=np.array_equal)) def test_ScaleTo01(self): data = [{'x': 1}, {'x': 2}, {'x': 3}, {'x': 4}, {'x': 5}, {'x': 6}] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ScaleTo01(columns=['x']))) expected_output = ( np.array([0], dtype=np.float32), np.array([0.2], dtype=np.float32), np.array([0.4], dtype=np.float32), np.array([0.6], dtype=np.float32), np.array([0.8], dtype=np.float32), np.array([1], dtype=np.float32)) actual_output = (result | beam.Map(lambda x: x.x)) assert_that( actual_output, equal_to(expected_output, equals_fn=np.array_equal)) class ScaleToGaussianTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) def test_gaussian_list_uniform_distribution(self): list_data = [{'x': [1, 2, 3]}, {'x': [4, 5, 6]}] with beam.Pipeline() as p: list_result = ( p | "listCreate" >> beam.Create(list_data) | "listMLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ScaleToGaussian(columns=['x']))) expected_data = [ np.array([-1.46385, -0.87831, -0.29277], dtype=np.float32), np.array([0.29277, 0.87831, 1.46385], dtype=np.float32) ] actual_data = (list_result | beam.Map(lambda x: x.x)) assert_that( actual_data, equal_to(expected_data, equals_fn=np.array_equal)) def test_gaussian_list_skewed_distribution(self): list_data = [{'x': [1, 2, 4]}, {'x': [8, 16, 32]}] with beam.Pipeline() as p: list_result = ( p | "listCreate" >> beam.Create(list_data) | "listMLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ScaleToGaussian(columns=['x']))) expected_data = [ np.array([-0.87733847, -0.78498703, -0.6002842], dtype=np.float32), np.array([-0.23087855, 0.5079328, 1.9855555], dtype=np.float32) ] actual_data = (list_result | beam.Map(lambda x: x.x)) assert_that(actual_data, equal_to(expected_data, equals_fn=np.allclose)) def test_gaussian_uniform(self): data = [ { 'x': 1 }, { 'x': 2 }, { 'x': 3 }, { 'x': 4 }, { 'x': 5 }, { 'x': 6 }, ] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ScaleToZScore(columns=['x']))) expected_data = [ np.array([-1.46385], dtype=np.float32), np.array([-0.87831], dtype=np.float32), np.array([-0.29277], dtype=np.float32), np.array([0.29277], dtype=np.float32), np.array([0.87831], dtype=np.float32), np.array([1.46385], dtype=np.float32), ] actual_data = (result | beam.Map(lambda x: x.x)) assert_that( actual_data, equal_to(expected_data, equals_fn=np.array_equal)) def test_gaussian_skewed(self): data = [ { 'x': 1 }, { 'x': 2 }, { 'x': 4 }, { 'x': 8 }, { 'x': 16 }, { 'x': 32 }, ] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ScaleToZScore(columns=['x']))) expected_data = [ np.array([-0.87733847], dtype=np.float32), np.array([-0.78498703], dtype=np.float32), np.array([-0.6002842], dtype=np.float32), np.array([-0.23087855], dtype=np.float32), np.array([0.5079328], dtype=np.float32), np.array([1.9855555], dtype=np.float32), ] actual_data = (result | beam.Map(lambda x: x.x)) assert_that(actual_data, equal_to(expected_data, equals_fn=np.allclose)) class BucketizeTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) def test_bucketize(self): data = [{'x': 1}, {'x': 2}, {'x': 3}, {'x': 4}, {'x': 5}, {'x': 6}] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.Bucketize(columns=['x'], num_buckets=3))) transformed_data = (result | beam.Map(lambda x: x.x)) expected_data = [ np.array([0]), np.array([0]), np.array([1]), np.array([1]), np.array([2]), np.array([2]) ] assert_that( transformed_data, equal_to(expected_data, equals_fn=np.array_equal)) def test_bucketize_list(self): list_data = [{'x': [1, 2, 3]}, {'x': [4, 5, 6]}] with beam.Pipeline() as p: list_result = ( p | "Create" >> beam.Create(list_data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.Bucketize(columns=['x'], num_buckets=3))) transformed_data = ( list_result | "TransformedColumnX" >> beam.Map(lambda ele: ele.x)) expected_data = [ np.array([0, 0, 1], dtype=np.int64), np.array([1, 2, 2], dtype=np.int64) ] assert_that( transformed_data, equal_to(expected_data, equals_fn=np.array_equal)) class ApplyBucketsTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) @parameterized.expand([ (range(1, 100), [25, 50, 75]), (range(1, 100, 2), [25, 51, 75]), ]) def test_apply_buckets(self, test_inputs, bucket_boundaries): with beam.Pipeline() as p: data = [{'x': [i]} for i in test_inputs] result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ApplyBuckets( columns=['x'], bucket_boundaries=bucket_boundaries))) expected_output = [] bucket = 0 for x in sorted(test_inputs): # Increment the bucket number when crossing the boundary if (bucket < len(bucket_boundaries) and x >= bucket_boundaries[bucket]): bucket += 1 expected_output.append(np.array([bucket])) actual_output = (result | beam.Map(lambda x: x.x)) assert_that( actual_output, equal_to(expected_output, equals_fn=np.array_equal)) class ApplyBucketsWithInterpolationTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) @parameterized.expand([ ([-1, 9, 10, 11], [10], [0., 0., 1., 1.]), ([15, 20, 25], [10, 20], [.5, 1, 1]), ]) def test_apply_buckets(self, test_inputs, bucket_boundaries, expected_values): with beam.Pipeline() as p: data = [{'x': [i]} for i in test_inputs] result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ApplyBucketsWithInterpolation( columns=['x'], bucket_boundaries=bucket_boundaries))) expected_output = [] for x in expected_values: expected_output.append(np.array(x)) actual_output = (result | beam.Map(lambda x: x.x)) assert_that( actual_output, equal_to(expected_output, equals_fn=np.allclose)) class ComputeAndApplyVocabTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) def test_compute_and_apply_vocabulary_inputs(self): num_elements = 100 num_instances = num_elements + 1 input_data = [ { 'x': '%.10i' % i, # Front-padded to facilitate lexicographic sorting. } for i in range(num_instances) ] expected_data = [ { 'x': (len(input_data) - 1) - i, # Due to reverse lexicographic sorting. } for i in range(len(input_data)) ] with beam.Pipeline() as p: actual_data = ( p | "Create" >> beam.Create(input_data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ComputeAndApplyVocabulary(columns=['x']))) actual_data |= beam.Map(lambda x: x.as_dict()) assert_that(actual_data, equal_to(expected_data)) def test_compute_and_apply_vocabulary(self): num_elements = 100 num_instances = num_elements + 1 input_data = [ { 'x': ['%.10i' % i, '%.10i' % (i + 1), '%.10i' % (i + 2)], # Front-padded to facilitate lexicographic sorting. } for i in range(0, num_instances, 3) ] # since we have 3 elements in a single list, multiply with 3 for # each iteration i on the expected output. excepted_data = [ np.array([(len(input_data) * 3 - 1) - i, (len(input_data) * 3 - 1) - i - 1, (len(input_data) * 3 - 1) - i - 2], dtype=np.int64) # Front-padded to facilitate lexicographic # sorting. for i in range(0, len(input_data) * 3, 3) ] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(input_data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ComputeAndApplyVocabulary(columns=['x']))) actual_output = (result | beam.Map(lambda x: x.x)) assert_that( actual_output, equal_to(excepted_data, equals_fn=np.array_equal)) def test_with_basic_example_list(self): data = [{ 'x': ['I', 'like', 'pie'], }, { 'x': ['yum', 'yum', 'pie'], }] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ComputeAndApplyVocabulary(columns=['x']))) result = result | beam.Map(lambda x: x.x) expected_result = [np.array([3, 2, 1]), np.array([0, 0, 1])] assert_that(result, equal_to(expected_result, equals_fn=np.array_equal)) def test_string_split_with_single_delimiter(self): data = [{ 'x': 'I like pie', }, { 'x': 'yum yum pie' }] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ComputeAndApplyVocabulary( columns=['x'], split_string_by_delimiter=' '))) result = result | beam.Map(lambda x: x.x) expected_result = [np.array([3, 2, 1]), np.array([0, 0, 1])] assert_that(result, equal_to(expected_result, equals_fn=np.array_equal)) def test_string_split_with_multiple_delimiters(self): data = [{ 'x': 'I like pie', }, { 'x': 'yum;yum;pie' }, { 'x': 'yum yum pie' }] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ComputeAndApplyVocabulary( columns=['x'], split_string_by_delimiter=' ;'))) result = result | beam.Map(lambda x: x.x) expected_result = [ np.array([3, 2, 1]), np.array([0, 0, 1]), np.array([0, 0, 1]) ] assert_that(result, equal_to(expected_result, equals_fn=np.array_equal)) def test_multiple_columns_with_default_vocab_name(self): data = [{ 'x': ['I', 'like', 'pie'], 'y': ['Apach', 'Beam', 'is', 'awesome'] }, { 'x': ['yum', 'yum', 'pie'], 'y': ['Beam', 'is', 'a', 'unified', 'model'] }] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ComputeAndApplyVocabulary(columns=['x', 'y']))) expected_data_x = [np.array([3, 2, 1]), np.array([0, 0, 1])] expected_data_y = [np.array([6, 1, 0, 4]), np.array([1, 0, 5, 2, 3])] actual_data_x = (result | beam.Map(lambda x: x.x)) actual_data_y = (result | beam.Map(lambda x: x.y)) assert_that( actual_data_x, equal_to(expected_data_x, equals_fn=np.array_equal), label='x') assert_that( actual_data_y, equal_to(expected_data_y, equals_fn=np.array_equal), label='y') files = os.listdir(self.artifact_location) files.remove(base._ATTRIBUTE_FILE_NAME) assert len(files) == 1 tft_vocab_assets = os.listdir( os.path.join( self.artifact_location, files[0], 'transform_fn', 'assets')) assert len(tft_vocab_assets) == 2 def test_multiple_columns_with_vocab_name(self): data = [{ 'x': ['I', 'like', 'pie'], 'y': ['Apach', 'Beam', 'is', 'awesome'] }, { 'x': ['yum', 'yum', 'pie'], 'y': ['Beam', 'is', 'a', 'unified', 'model'] }] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ComputeAndApplyVocabulary( columns=['x', 'y'], vocab_filename='my_vocab'))) expected_data_x = [np.array([3, 2, 1]), np.array([0, 0, 1])] expected_data_y = [np.array([6, 1, 0, 4]), np.array([1, 0, 5, 2, 3])] actual_data_x = (result | beam.Map(lambda x: x.x)) actual_data_y = (result | beam.Map(lambda x: x.y)) assert_that( actual_data_x, equal_to(expected_data_x, equals_fn=np.array_equal), label='x') assert_that( actual_data_y, equal_to(expected_data_y, equals_fn=np.array_equal), label='y') files = os.listdir(self.artifact_location) files.remove(base._ATTRIBUTE_FILE_NAME) assert len(files) == 1 tft_vocab_assets = os.listdir( os.path.join( self.artifact_location, files[0], 'transform_fn', 'assets')) assert len(tft_vocab_assets) == 2 class TFIDIFTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) def test_tfidf_compute_vocab_size_during_runtime(self): raw_data = [ dict(x=["I", "like", "pie", "pie", "pie"]), dict(x=["yum", "yum", "pie"]) ] with beam.Pipeline() as p: transforms = [ tft.ComputeAndApplyVocabulary(columns=['x']), tft.TFIDF(columns=['x']) ] actual_output = ( p | "Create" >> beam.Create(raw_data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location, transforms=transforms)) actual_output |= beam.Map(lambda x: x.as_dict()) def equals_fn(a, b): is_equal = True for key, value in a.items(): value_b = a[key] is_equal = is_equal and np.array_equal(value, value_b) return is_equal expected_output = ([{ 'x': np.array([3, 2, 0, 0, 0]), 'x_tfidf_weight': np.array([0.6, 0.28109303, 0.28109303], dtype=np.float32), 'x_vocab_index': np.array([0, 2, 3], dtype=np.int64) }, { 'x': np.array([1, 1, 0]), 'x_tfidf_weight': np.array( [0.33333334, 0.9369768], dtype=np.float32), 'x_vocab_index': np.array([0, 1], dtype=np.int32) }]) assert_that(actual_output, equal_to(expected_output, equals_fn=equals_fn)) class ScaleToMinMaxTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) def test_scale_to_min_max(self): data = [{ 'x': 4, }, { 'x': 1, }, { 'x': 5, }, { 'x': 2, }] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.ScaleByMinMax( columns=['x'], min_value=-1, max_value=1, ), )) result = result | beam.Map(lambda x: x.as_dict()) expected_data = [{ 'x': np.array([0.5], dtype=np.float32) }, { 'x': np.array([-1.0], dtype=np.float32) }, { 'x': np.array([1.0], dtype=np.float32) }, { 'x': np.array([-0.5], dtype=np.float32) }] assert_that(result, equal_to(expected_data)) def test_fail_max_value_less_than_min(self): with self.assertRaises(ValueError): tft.ScaleByMinMax(columns=['x'], min_value=10, max_value=0) class NGramsTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) def test_ngrams_on_list_separated_words_default_args(self): data = [{ 'x': ['I', 'like', 'pie'], }, { 'x': ['yum', 'yum', 'pie'] }] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location, transforms=[tft.NGrams(columns=['x'])])) result = result | beam.Map(lambda x: x.x) expected_data = [ np.array([b'I', b'like', b'pie'], dtype=object), np.array([b'yum', b'yum', b'pie'], dtype=object) ] assert_that(result, equal_to(expected_data, equals_fn=np.array_equal)) def test_ngrams_on_list_separated_words(self): data = [{ 'x': ['I', 'like', 'pie'], }, { 'x': ['yum', 'yum', 'pie'] }] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location, transforms=[ tft.NGrams( columns=['x'], ngram_range=(1, 3), ngrams_separator=' ') ])) result = result | beam.Map(lambda x: x.x) expected_data = [ np.array( [b'I', b'I like', b'I like pie', b'like', b'like pie', b'pie'], dtype=object), np.array( [b'yum', b'yum yum', b'yum yum pie', b'yum', b'yum pie', b'pie'], dtype=object) ] assert_that(result, equal_to(expected_data, equals_fn=np.array_equal)) def test_with_string_split_delimiter(self): data = [{ 'x': 'I like pie', }, { 'x': 'yum yum pie' }] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location, transforms=[ tft.NGrams( columns=['x'], split_string_by_delimiter=' ', ngram_range=(1, 3), ngrams_separator=' ') ])) result = result | beam.Map(lambda x: x.x) expected_data = [ np.array( [b'I', b'I like', b'I like pie', b'like', b'like pie', b'pie'], dtype=object), np.array( [b'yum', b'yum yum', b'yum yum pie', b'yum', b'yum pie', b'pie'], dtype=object) ] assert_that(result, equal_to(expected_data, equals_fn=np.array_equal)) def test_with_multiple_string_delimiters(self): data = [{ 'x': 'I?like?pie', }, { 'x': 'yum yum pie' }] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location, transforms=[ tft.NGrams( columns=['x'], split_string_by_delimiter=' ?', ngram_range=(1, 3), ngrams_separator=' ') ])) result = result | beam.Map(lambda x: x.x) expected_data = [ np.array( [b'I', b'I like', b'I like pie', b'like', b'like pie', b'pie'], dtype=object), np.array( [b'yum', b'yum yum', b'yum yum pie', b'yum', b'yum pie', b'pie'], dtype=object) ] assert_that(result, equal_to(expected_data, equals_fn=np.array_equal)) class BagOfWordsTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) def test_bag_of_words_on_list_seperated_words_default_ngrams(self): data = [{ 'x': ['I', 'like', 'pie', 'pie', 'pie'], }, { 'x': ['yum', 'yum', 'pie'] }] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location, transforms=[tft.BagOfWords(columns=['x'])])) result = result | beam.Map(lambda x: x.x) expected_data = [ np.array([b'I', b'like', b'pie'], dtype=object), np.array([b'yum', b'pie'], dtype=object) ] assert_that(result, equal_to(expected_data, equals_fn=np.array_equal)) def test_bag_of_words_on_list_seperated_words_custom_ngrams(self): data = [{ 'x': ['I', 'like', 'pie', 'I', 'like', 'pie'], }, { 'x': ['yum', 'yum', 'pie'] }] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location, transforms=[ tft.BagOfWords( columns=['x'], ngram_range=(1, 3), ngrams_separator=' ') ])) result = result | beam.Map(lambda x: x.x) expected_data = [[ b'I', b'I like', b'I like pie', b'like', b'like pie', b'like pie I', b'pie', b'pie I', b'pie I like' ], [b'yum', b'yum yum', b'yum yum pie', b'yum pie', b'pie']] assert_that(result, equal_to(expected_data, equals_fn=np.array_equal)) def test_bag_of_words_on_numpy_data(self): data = [{ 'x': np.array(['I', 'like', 'pie', 'I', 'like', 'pie'], dtype=object), }, { 'x': np.array(['yum', 'yum', 'pie'], dtype=object) }] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location, transforms=[ tft.BagOfWords( columns=['x'], ngram_range=(1, 3), ngrams_separator=' ') ])) result = result | beam.Map(lambda x: x.x) expected_data = [[ b'I', b'I like', b'I like pie', b'like', b'like pie', b'like pie I', b'pie', b'pie I', b'pie I like' ], [b'yum', b'yum yum', b'yum yum pie', b'yum pie', b'pie']] assert_that(result, equal_to(expected_data, equals_fn=np.array_equal)) def test_bag_of_words_on_by_splitting_input_text(self): data = [{'x': 'I like pie I like pie'}, {'x': 'yum yum pie'}] with beam.Pipeline() as p: result = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location, transforms=[ tft.BagOfWords( columns=['x'], split_string_by_delimiter=' ', ngram_range=(1, 3), ngrams_separator=' ') ])) result = result | beam.Map(lambda x: x.x) expected_data = [[ b'I', b'I like', b'I like pie', b'like', b'like pie', b'like pie I', b'pie', b'pie I', b'pie I like' ], [b'yum', b'yum yum', b'yum yum pie', b'yum pie', b'pie']] assert_that(result, equal_to(expected_data, equals_fn=np.array_equal)) def test_count_per_key_on_list(self): data = [{ 'x': ['I', 'like', 'pie', 'pie', 'pie'], }, { 'x': ['yum', 'yum', 'pie'] }, { 'x': ['Banana', 'Banana', 'Apple', 'Apple', 'Apple', 'Apple'] }] with beam.Pipeline() as p: _ = ( p | "Create" >> beam.Create(data) | "MLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location, transforms=[ tft.BagOfWords( columns=['x'], compute_word_count=True, key_vocab_filename='my_vocab') ])) def validate_count_per_key(key_vocab_filename): files = os.listdir(self.artifact_location) files.remove(base._ATTRIBUTE_FILE_NAME) key_vocab_location = os.path.join( self.artifact_location, files[0], 'transform_fn/assets', key_vocab_filename + '_x') with open(key_vocab_location, 'r') as f: key_vocab_list = [line.strip() for line in f] return key_vocab_list expected_data = ['2 yum', '4 Apple', '1 like', '1 I', '4 pie', '2 Banana'] actual_data = validate_count_per_key('my_vocab') self.assertEqual(expected_data, actual_data) class HashStringsTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) def test_single_bucket(self): list_data = [{'x': 'this is a test string'}] expected_values = [np.array([0])] with beam.Pipeline() as p: list_result = ( p | "listCreate" >> beam.Create(list_data) | "listMLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.HashStrings(columns=['x'], hash_buckets=1))) result = (list_result | beam.Map(lambda x: x.x)) assert_that(result, equal_to(expected_values, equals_fn=np.array_equal)) def test_multi_bucket_one_string(self): list_data = [{'x': 'this is a test string'}] expected_values = [np.array([1])] with beam.Pipeline() as p: list_result = ( p | "listCreate" >> beam.Create(list_data) | "listMLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.HashStrings(columns=['x'], hash_buckets=2))) result = (list_result | beam.Map(lambda x: x.x)) assert_that(result, equal_to(expected_values, equals_fn=np.array_equal)) def test_one_bucket_multi_string(self): list_data = [{'x': ['these', 'are', 'test', 'strings']}] expected_values = [np.array([0, 0, 0, 0])] with beam.Pipeline() as p: list_result = ( p | "listCreate" >> beam.Create(list_data) | "listMLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.HashStrings(columns=['x'], hash_buckets=1))) result = (list_result | beam.Map(lambda x: x.x)) assert_that(result, equal_to(expected_values, equals_fn=np.array_equal)) def test_two_bucket_multi_string(self): list_data = [{'x': ['these', 'are', 'test', 'strings']}] expected_values = [np.array([1, 0, 1, 0])] with beam.Pipeline() as p: list_result = ( p | "listCreate" >> beam.Create(list_data) | "listMLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.HashStrings( columns=['x'], hash_buckets=2, key=[123, 456]))) result = (list_result | beam.Map(lambda x: x.x)) assert_that(result, equal_to(expected_values, equals_fn=np.array_equal)) def test_multi_buckets_multi_string(self): # This is a recreation of one of the TFT test cases from # https://github.com/tensorflow/transform/blob/master/tensorflow_transform/mappers_test.py list_data = [{'x': ['Cake', 'Pie', 'Sundae']}] expected_values = [np.array([6, 5, 6])] with beam.Pipeline() as p: list_result = ( p | "listCreate" >> beam.Create(list_data) | "listMLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.HashStrings( columns=['x'], hash_buckets=11, key=[123, 456]))) result = (list_result | beam.Map(lambda x: x.x)) assert_that(result, equal_to(expected_values, equals_fn=np.array_equal)) class DeduplicateTensorPerRowTest(unittest.TestCase): def setUp(self) -> None: self.artifact_location = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.artifact_location) def test_deduplicate(self): values = [{ 'x': [b'a', b'b', b'a', b'b'], }, { 'x': [b'b', b'c', b'b', b'c'] }] expected_output = [np.array([b'a', b'b']), np.array([b'b', b'c'])] with beam.Pipeline() as p: list_result = ( p | "listCreate" >> beam.Create(values) | "listMLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.DeduplicateTensorPerRow(columns=['x']))) result = (list_result | beam.Map(lambda x: x.x)) assert_that(result, equal_to(expected_output, equals_fn=np.array_equal)) def test_deduplicate_no_op(self): values = [{ 'x': [b'a', b'b'], }, { 'x': [b'c', b'd'] }] expected_output = [np.array([b'a', b'b']), np.array([b'c', b'd'])] with beam.Pipeline() as p: list_result = ( p | "listCreate" >> beam.Create(values) | "listMLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.DeduplicateTensorPerRow(columns=['x']))) result = (list_result | beam.Map(lambda x: x.x)) assert_that(result, equal_to(expected_output, equals_fn=np.array_equal)) def test_deduplicate_different_output_sizes(self): values = [{ 'x': [b'a', b'b', b'a', b'b'], }, { 'x': [b'c', b'a', b'd', b'd'] }] expected_output = [np.array([b'a', b'b']), np.array([b'c', b'a', b'd'])] with beam.Pipeline() as p: list_result = ( p | "listCreate" >> beam.Create(values) | "listMLTransform" >> base.MLTransform( write_artifact_location=self.artifact_location).with_transform( tft.DeduplicateTensorPerRow(columns=['x']))) result = (list_result | beam.Map(lambda x: x.x)) assert_that(result, equal_to(expected_output, equals_fn=np.array_equal)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/transforms/utils.py0000644000175100017510000000713115206113032022177 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # __all__ = ['ArtifactsFetcher'] import os import tempfile import tensorflow_transform as tft from google.cloud.storage import Client from google.cloud.storage import transfer_manager from apache_beam.ml.transforms import base def download_artifacts_from_gcs(bucket_name, prefix, local_path): """Downloads artifacts from GCS to the local file system. Args: bucket_name: The name of the GCS bucket to download from. prefix: Prefix of GCS objects to download. local_path: The local path to download the folder to. """ client = Client() bucket = client.get_bucket(bucket_name) blobs = [blob.name for blob in bucket.list_blobs(prefix=prefix)] _ = transfer_manager.download_many_to_path( bucket, blobs, destination_directory=local_path) class ArtifactsFetcher: """ Utility class used to fetch artifacts from the artifact_location passed to the TFTProcessHandlers in MLTransform. This is intended to be used for testing purposes only. """ def __init__(self, artifact_location: str): tempdir = tempfile.mkdtemp() if artifact_location.startswith('gs://'): parts = artifact_location[5:].split('/') bucket_name = parts[0] prefix = '/'.join(parts[1:]) download_artifacts_from_gcs(bucket_name, prefix, tempdir) assert os.listdir(tempdir), f"No files found in {artifact_location}" artifact_location = os.path.join(tempdir, prefix) files = os.listdir(artifact_location) files.remove(base._ATTRIBUTE_FILE_NAME) # TODO: https://github.com/apache/beam/issues/29356 # Integrate ArtifactFetcher into MLTransform. if len(files) > 1: raise NotImplementedError( "MLTransform may have been utilized alongside transforms written " "in TensorFlow Transform, in conjunction with those from different " "frameworks. Currently, retrieving artifacts from this " "multi-framework setup is not supported.") self._artifact_location = os.path.join(artifact_location, files[0]) self.transform_output = tft.TFTransformOutput(self._artifact_location) def get_vocab_list(self, vocab_filename: str) -> list[bytes]: """ Returns list of vocabulary terms created during MLTransform. """ try: vocab_list = self.transform_output.vocabulary_by_name(vocab_filename) except ValueError as e: raise ValueError( 'Vocabulary file {} not found in artifact location'.format( vocab_filename)) from e return [x.decode('utf-8') for x in vocab_list] def get_vocab_filepath(self, vocab_filename: str) -> str: """ Return the path to the vocabulary file created during MLTransform. """ return self.transform_output.vocabulary_file_by_name(vocab_filename) def get_vocab_size(self, vocab_filename: str) -> int: return self.transform_output.vocabulary_size_by_name(vocab_filename) ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5224228 apache_beam-2.74.0/apache_beam/ml/ts/0000755000175100017510000000000015206113160016715 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/ml/ts/__init__.py0000644000175100017510000000142015206113032021021 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5264232 apache_beam-2.74.0/apache_beam/options/0000755000175100017510000000000015206113160017352 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/options/__init__.py0000644000175100017510000000142015206113032021456 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/options/pipeline_options.py0000644000175100017510000025402415206113032023311 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Pipeline options obtained from command line parsing.""" # pytype: skip-file import argparse import difflib import json import logging import os from typing import Any from typing import Callable from typing import Optional from typing import Sequence from typing import TypeVar import apache_beam as beam from apache_beam.options.value_provider import RuntimeValueProvider from apache_beam.options.value_provider import StaticValueProvider from apache_beam.options.value_provider import ValueProvider from apache_beam.transforms.display import HasDisplayData from apache_beam.utils import logger from apache_beam.utils import proto_utils __all__ = [ 'PipelineOptions', 'StandardOptions', 'TypeOptions', 'DirectOptions', 'GoogleCloudOptions', 'AzureOptions', 'HadoopFileSystemOptions', 'WorkerOptions', 'DebugOptions', 'ProfilingOptions', 'SetupOptions', 'TestOptions', 'S3Options' ] PipelineOptionsT = TypeVar('PipelineOptionsT', bound='PipelineOptions') _LOGGER = logging.getLogger(__name__) # Map defined with option names to flag names for boolean options # that have a destination(dest) in parser.add_argument() different # from the flag name and whose default value is `None`. _FLAG_THAT_SETS_FALSE_VALUE = { 'use_public_ips': 'no_use_public_ips', 'save_main_session': 'no_save_main_session' } # Set of options which should not be overriden when applying options from a # different language. This is relevant when using x-lang transforms where the # expansion service is started up with some pipeline options, and will # impact which options are passed in to expanded transforms' expand functions. _NON_OVERIDABLE_XLANG_OPTIONS = ['runner', 'experiments'] def _static_value_provider_of(value_type): """Helper function to plug a ValueProvider into argparse. Args: value_type: the type of the value. Since the type param of argparse's add_argument will always be ValueProvider, we need to preserve the type of the actual value. Returns: A partially constructed StaticValueProvider in the form of a function. """ def _f(value): _f.__name__ = value_type.__name__ return StaticValueProvider(value_type, value) return _f class _BeamArgumentParser(argparse.ArgumentParser): """An ArgumentParser that supports ValueProvider options. Example Usage:: class TemplateUserOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument('--vp_arg1', default='start') parser.add_value_provider_argument('--vp_arg2') parser.add_argument('--non_vp_arg') """ def add_value_provider_argument(self, *args, **kwargs): """ValueProvider arguments can be either of type keyword or positional. At runtime, even positional arguments will need to be supplied in the key/value form. """ # Extract the option name from positional argument ['pos_arg'] assert args and len(args[0]) >= 1 if args[0][0] != '-': option_name = args[0] if kwargs.get('nargs') is None: # make them optionally templated kwargs['nargs'] = '?' else: # or keyword arguments like [--kw_arg, -k, -w] or [--kw-arg] option_name = [i.replace('--', '') for i in args if i[:2] == '--'][0] # reassign the type to make room for using # StaticValueProvider as the type for add_argument value_type = kwargs.get('type') or str kwargs['type'] = _static_value_provider_of(value_type) # reassign default to default_value to make room for using # RuntimeValueProvider as the default for add_argument default_value = kwargs.get('default') kwargs['default'] = RuntimeValueProvider( option_name=option_name, value_type=value_type, default_value=default_value) # have add_argument do most of the work self.add_argument(*args, **kwargs) # The argparse package by default tries to autocomplete option names. This # results in an "ambiguous option" error from argparse when an unknown option # matching multiple known ones are used. This suppresses that behavior. def error(self, message): if message.startswith('ambiguous option: '): return super().error(message) class _DictUnionAction(argparse.Action): """ argparse Action take union of json loads values. If a key is specified in more than one of the values, the last value takes precedence. """ def __call__(self, parser, namespace, values, option_string=None): if not hasattr(namespace, self.dest) or getattr(namespace, self.dest) is None: setattr(namespace, self.dest, {}) getattr(namespace, self.dest).update(values) class _GcsCustomAuditEntriesAction(argparse.Action): """ Argparse Action for GCS custom audit entries. According to Google Cloud Storage audit logging documentation (https://cloud.google.com/storage/docs/audit-logging#add-custom-metadata), the following limitations apply: - keys must be 64 characters or less, - values 1,200 characters or less, and - a maximum of four custom metadata entries are permitted per request. """ MAX_KEY_LENGTH = 64 MAX_VALUE_LENGTH = 1200 MAX_ENTRIES = 4 GCS_AUDIT_PREFIX = 'x-goog-custom-audit-' def _exceed_entry_limit(self): job_audit_entry = _GcsCustomAuditEntriesAction.GCS_AUDIT_PREFIX + 'job' if job_audit_entry in self._custom_audit_entries: return len( self._custom_audit_entries) > _GcsCustomAuditEntriesAction.MAX_ENTRIES else: return len(self._custom_audit_entries) > ( _GcsCustomAuditEntriesAction.MAX_ENTRIES - 1) def _add_entry(self, key, value): if len(key) > _GcsCustomAuditEntriesAction.MAX_KEY_LENGTH: raise argparse.ArgumentError( None, "The key '%s' in GCS custom audit entries exceeds the %d-character limit." # pylint: disable=line-too-long % (key, _GcsCustomAuditEntriesAction.MAX_KEY_LENGTH)) if len(value) > _GcsCustomAuditEntriesAction.MAX_VALUE_LENGTH: raise argparse.ArgumentError( None, "The value '%s' in GCS custom audit entries exceeds the %d-character limit." # pylint: disable=line-too-long % (value, _GcsCustomAuditEntriesAction.MAX_VALUE_LENGTH)) if key.startswith(_GcsCustomAuditEntriesAction.GCS_AUDIT_PREFIX): self._custom_audit_entries[key] = value else: self._custom_audit_entries[_GcsCustomAuditEntriesAction.GCS_AUDIT_PREFIX + key] = value def __call__(self, parser, namespace, values, option_string=None): if not hasattr(namespace, self.dest) or getattr(namespace, self.dest) is None: setattr(namespace, self.dest, {}) self._custom_audit_entries = getattr(namespace, self.dest) if option_string == '--gcs_custom_audit_entries': # in the format of {"key": "value"} assert (isinstance(values, str)) sub_entries = json.loads(values) for key, value in sub_entries.items(): self._add_entry(key, value) else: # option_string == '--gcs_custom_audit_entry' # in the format of 'key=value' assert (isinstance(values, str)) parts = values.split('=', 1) key = parts[0] value = parts[1] if len(parts) > 1 else '' self._add_entry(key, value) if self._exceed_entry_limit(): raise argparse.ArgumentError( None, "The maximum allowed number of GCS custom audit entries (including the default x-goo-custom-audit-job) is %d." # pylint: disable=line-too-long % _GcsCustomAuditEntriesAction.MAX_ENTRIES) class _CommaSeparatedListAction(argparse.Action): """ Argparse Action that splits comma-separated values and appends them to a list. This allows options like --experiments=abc,def to be treated as separate experiments 'abc' and 'def', similar to how Java SDK handles them. If there are key=value experiments in a raw argument, the remaining part of the argument are treated as values and won't split further. For example: 'abc,def,master_key=k1=v1,k2=v2' becomes ['abc', 'def', 'master_key=k1=v1,k2=v2']. """ def __call__(self, parser, namespace, values, option_string=None): if not hasattr(namespace, self.dest) or getattr(namespace, self.dest) is None: setattr(namespace, self.dest, []) # Split comma-separated values and extend the list if isinstance(values, str): # Smart splitting: only split at commas that are not part of # key=value pairs split_values = self._smart_split(values) getattr(namespace, self.dest).extend(split_values) else: # If values is not a string, just append it getattr(namespace, self.dest).append(values) def _smart_split(self, values): """Split comma-separated values, but preserve commas within key=value pairs.""" result = [] current = [] equals_depth = 0 i = 0 while i < len(values): char = values[i] if char == '=': equals_depth += 1 current.append(char) elif char == ',' and equals_depth <= 1: # This comma is a top-level separator (not inside a complex value) if current: result.append(''.join(current).strip()) current = [] equals_depth = 0 elif char == ',' and equals_depth > 1: # This comma is inside a complex value, keep it current.append(char) elif char == ' ' and not current: # Skip leading spaces pass else: current.append(char) i += 1 # Add the last item if current: result.append(''.join(current).strip()) return [v for v in result if v] # Filter out empty values class PipelineOptions(HasDisplayData): # Set of options which should not be overriden when pipeline options are # being merged (see from_runner_api). This primarily comes up when expanding # the Python expansion service """This class and subclasses are used as containers for command line options. These classes are wrappers over the standard argparse Python module (see https://docs.python.org/3/library/argparse.html). To define one option or a group of options, create a subclass from PipelineOptions. Example Usage:: class XyzOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--abc', default='start') parser.add_argument('--xyz', default='end') The arguments for the add_argument() method are exactly the ones described in the argparse public documentation. Pipeline objects require an options object during initialization. This is obtained simply by initializing an options class as defined above. Example Usage:: p = Pipeline(options=XyzOptions()) if p.options.xyz == 'end': raise ValueError('Option xyz has an invalid value.') Instances of PipelineOptions or any of its subclass have access to values defined by other PipelineOption subclasses (see get_all_options()), and can be converted to an instance of another PipelineOptions subclass (see view_as()). All views share the underlying data structure that stores option key-value pairs. By default the options classes will use command line arguments to initialize the options. """ def __init__(self, flags: Optional[Sequence[str]] = None, **kwargs) -> None: """Initialize an options class. The initializer will traverse all subclasses, add all their argparse arguments and then parse the command line specified by flags or by default the one obtained from sys.argv. The subclasses of PipelineOptions do not need to redefine __init__. Args: flags: An iterable of command line arguments to be used. If not specified then sys.argv will be used as input for parsing arguments. **kwargs: Add overrides for arguments passed in flags. For overrides of arguments, please pass the `option names` instead of flag names. Option names: These are defined as dest in the parser.add_argument() for each flag. Passing flags like {no_use_public_ips: True}, for which the dest is defined to a different flag name in the parser, would be discarded. Instead, pass the dest of the flag (dest of no_use_public_ips is use_public_ips). """ # Initializing logging configuration in case the user did not set it up. logging.basicConfig() # self._flags stores a list of not yet parsed arguments, typically, # command-line flags. This list is shared across different views. # See: view_as(). if isinstance(flags, str): # Unfortunately a single str passes the Iterable[str] test, as it is # an iterable of single characters. This is almost certainly not the # intent... raise ValueError( "Flags must be an iterable of of strings, not a single string.") self._flags = flags # Build parser that will parse options recognized by the [sub]class of # PipelineOptions whose object is being instantiated. parser = _BeamArgumentParser(allow_abbrev=False) for cls in type(self).mro(): if cls == PipelineOptions: break elif '_add_argparse_args' in cls.__dict__: cls._add_argparse_args(parser) # type: ignore # The _visible_options attribute will contain options that were recognized # by the parser. self._visible_options, _ = parser.parse_known_args(flags) # self._all_options is initialized with overrides to flag values, # provided in kwargs, and will store key-value pairs for options recognized # by current PipelineOptions [sub]class and its views that may be created. # See: view_as(). # This dictionary is shared across different views, and is lazily updated # as each new views are created. # Users access this dictionary store via __getattr__ / __setattr__ methods. self._all_options = kwargs # Initialize values of keys defined by this class. for option_name in self._visible_option_list(): # Note that options specified in kwargs will not be overwritten. if option_name not in self._all_options: self._all_options[option_name] = getattr( self._visible_options, option_name) def __getstate__(self): # The impersonate_service_account option must be used only at submission of # a Beam job. However, Beam IOs might store pipeline options # within transform implementation that becomes serialized in RunnerAPI, # causing this option to be inadvertently used at runtime. # This serialization hook removes it. if self.view_as(GoogleCloudOptions).impersonate_service_account: dict_copy = dict(self.__dict__) dict_copy['_all_options'] = dict(dict_copy['_all_options']) dict_copy['_all_options']['impersonate_service_account'] = None return dict_copy else: return self.__dict__ @classmethod def _add_argparse_args(cls, parser: _BeamArgumentParser) -> None: # Override this in subclasses to provide options. pass @classmethod def from_dictionary(cls, options): """Returns a PipelineOptions from a dictionary of arguments. Args: options: Dictionary of argument value pairs. Returns: A PipelineOptions object representing the given arguments. """ flags = [] for k, v in options.items(): # Note: If a boolean flag is True in the dictionary, # implicitly the method assumes the boolean flag is # specified as a command line argument. If the # boolean flag is False, this method simply discards them. # Eg: {no_auth: True} is similar to python your_file.py --no_auth # {no_auth: False} is similar to python your_file.py. if isinstance(v, bool): if v: flags.append('--%s' % k) elif k in _FLAG_THAT_SETS_FALSE_VALUE: # Capture overriding flags, which have a different dest # from the flag name defined in the parser.add_argument # Eg: no_use_public_ips, which has the dest=use_public_ips # different from flag name flag_that_disables_the_option = (_FLAG_THAT_SETS_FALSE_VALUE[k]) flags.append('--%s' % flag_that_disables_the_option) elif isinstance(v, list): for i in v: flags.append('--%s=%s' % (k, i)) elif isinstance(v, dict): flags.append('--%s=%s' % (k, json.dumps(v))) elif v is None: # Don't process None type args here, they will be treated # as strings when parsed by BeamArgumentParser.. logging.warning('Not setting flag with value None: %s', k) else: flags.append('--%s=%s' % (k, v)) return cls(flags) @staticmethod def _warn_on_unknown_options(unknown_args, parser): if not unknown_args: return all_known_options = [ opt for action in parser._actions for opt in action.option_strings ] for arg in unknown_args: msg = f"Unparseable argument: {arg}" if arg.startswith('--'): arg_name = arg.split('=', 1)[0] suggestions = difflib.get_close_matches(arg_name, all_known_options) if suggestions: msg += f". Did you mean '{suggestions[0]}'?'" logger.log_first_n(logging.WARN, msg, key="message") def get_all_options( self, drop_default=False, add_extra_args_fn: Optional[Callable[[_BeamArgumentParser], None]] = None, retain_unknown_options=False, display_warnings=False, current_only=False, ) -> dict[str, Any]: """Returns a dictionary of all defined arguments. Returns a dictionary of all defined arguments into a dictionary. Args: drop_default: If set to true, options that are equal to their default values, are not returned as part of the result dictionary. add_extra_args_fn: Callback to populate additional arguments, can be used by runner to supply otherwise unknown args. retain_unknown_options: If set to true, options not recognized by any known pipeline options class will still be included in the result. If set to false, they will be discarded. current_only: If set to true, only returns options defined in this class. Otherwise, arguments that are defined in any subclass of PipelineOptions are returned (default). Returns: Dictionary of all args and values. """ # TODO(https://github.com/apache/beam/issues/18197): PipelineOption # sub-classes in the main session might be repeated. Pick last unique # instance of each subclass to avoid conflicts. subset = {} parser = _BeamArgumentParser(allow_abbrev=False) if current_only: subset.setdefault(str(type(self)), type(self)) else: for cls in PipelineOptions.__subclasses__(): subset.setdefault(str(cls), cls) for cls in subset.values(): cls._add_argparse_args(parser) # pylint: disable=protected-access if add_extra_args_fn: add_extra_args_fn(parser) known_args, unknown_args = parser.parse_known_args(self._flags) if display_warnings: self._warn_on_unknown_options(unknown_args, parser) if retain_unknown_options: seen = set() def add_new_arg(arg, **kwargs): if arg not in seen: parser.add_argument(arg, **kwargs) seen.add(arg) i = 0 while i < len(unknown_args): # End of argument parsing. if unknown_args[i] == '--': break # Treat all unary flags as booleans, and all binary argument values as # strings. if not unknown_args[i].startswith('-'): i += 1 continue if i + 1 >= len(unknown_args) or unknown_args[i + 1].startswith('-'): split = unknown_args[i].split('=', 1) if len(split) == 1: add_new_arg(unknown_args[i], action='store_true') else: add_new_arg(split[0], type=str) i += 1 elif unknown_args[i].startswith('--'): add_new_arg(unknown_args[i], type=str) i += 2 else: # skip all binary flags used with '-' and not '--'. # ex: using -f instead of --f (or --flexrs_goal) will prevent # argument validation before job submission and can be incorrectly # submitted to job. _LOGGER.warning( "Discarding flag %s, single dash flags are not allowed.", unknown_args[i]) i += 2 continue parsed_args, _ = parser.parse_known_args(self._flags) else: if unknown_args and not current_only: _LOGGER.warning("Discarding unparseable args: %s", unknown_args) parsed_args = known_args result = vars(parsed_args) overrides = self._all_options.copy() # Apply the overrides if any for k in list(result): overrides.pop(k, None) if k in self._all_options: result[k] = self._all_options[k] if (drop_default and parser.get_default(k) == result[k] and not isinstance(parser.get_default(k), ValueProvider)): del result[k] if overrides: if retain_unknown_options: result.update(overrides) elif not current_only: _LOGGER.warning("Discarding invalid overrides: %s", overrides) return result def to_runner_api(self): def to_struct_value(o): if isinstance(o, (bool, int, str)): return o elif isinstance(o, (tuple, list)): return [to_struct_value(e) for e in o] elif isinstance(o, dict): return {str(k): to_struct_value(v) for k, v in o.items()} else: return str(o) # Best effort. return proto_utils.pack_Struct( **{ f'beam:option:{k}:v1': to_struct_value(v) for (k, v) in self.get_all_options( drop_default=True, retain_unknown_options=True).items() if v is not None }) @classmethod def from_runner_api(cls, proto_options, original_options=None): def from_urn(key): assert key.startswith('beam:option:') assert key.endswith(':v1') return key[12:-3] parsed = {from_urn(key): value for (key, value) in proto_options.items()} if original_options is None: return cls(**parsed) for (key, value) in parsed.items(): if value and key not in _NON_OVERIDABLE_XLANG_OPTIONS: original_options._all_options[key] = value return original_options def display_data(self): return self.get_all_options(drop_default=True, retain_unknown_options=True) def view_as(self, cls: type[PipelineOptionsT]) -> PipelineOptionsT: """Returns a view of current object as provided PipelineOption subclass. Example Usage:: options = PipelineOptions(['--runner', 'Direct', '--streaming']) standard_options = options.view_as(StandardOptions) if standard_options.streaming: # ... start a streaming job ... Note that options objects may have multiple views, and modifications of values in any view-object will apply to current object and other view-objects. Args: cls: PipelineOptions class or any of its subclasses. Returns: An instance of cls that is initialized using options contained in current object. """ view = cls(self._flags) for option_name in view._visible_option_list(): # Initialize values of keys defined by a cls. # # Note that we do initialization only once per key to make sure that # values in _all_options dict are not-recreated with each new view. # This is important to make sure that values of multi-options keys are # backed by the same list across multiple views, and that any overrides of # pipeline options already stored in _all_options are preserved. if option_name not in self._all_options: self._all_options[option_name] = getattr( view._visible_options, option_name) # Note that views will still store _all_options of the source object. view._all_options = self._all_options return view def is_compat_version_prior_to(self, breaking_change_version): """Check if update_compatibility_version is prior to a breaking change. Returns True if the pipeline should use old behavior (i.e., the update_compatibility_version is set and is earlier than the given version). Returns False if update_compatibility_version is not set or is >= the breaking change version. Args: breaking_change_version: Version string (e.g., "2.72.0") at which the breaking change was introduced. """ v1 = self.view_as(StreamingOptions).update_compatibility_version if v1 is None: return False v1_parts = (v1.split('.') + ['0', '0', '0'])[:3] v2_parts = (breaking_change_version.split('.') + ['0', '0', '0'])[:3] return tuple(map(int, v1_parts)) < tuple(map(int, v2_parts)) def _visible_option_list(self) -> list[str]: return sorted( option for option in dir(self._visible_options) if option[0] != '_') def __dir__(self) -> list[str]: return sorted( dir(type(self)) + list(self.__dict__) + self._visible_option_list()) def __getattr__(self, name): # Special methods which may be accessed before the object is # fully constructed (e.g. in unpickling). if name[:2] == name[-2:] == '__': return object.__getattribute__(self, name) elif name in self._visible_option_list(): return self._all_options[name] else: raise AttributeError( "'%s' object has no attribute '%s'" % (type(self).__name__, name)) def __setattr__(self, name, value): if name in ('_flags', '_all_options', '_visible_options'): super().__setattr__(name, value) elif name in self._visible_option_list(): self._all_options[name] = value else: raise AttributeError( "'%s' object has no attribute '%s'" % (type(self).__name__, name)) def __str__(self): return '%s(%s)' % ( type(self).__name__, ', '.join( '%s=%s' % (option, getattr(self, option)) for option in self._visible_option_list())) class StandardOptions(PipelineOptions): DEFAULT_RUNNER = 'DirectRunner' ALL_KNOWN_RUNNERS = ( 'apache_beam.runners.dataflow.dataflow_runner.DataflowRunner', 'apache_beam.runners.direct.direct_runner.BundleBasedDirectRunner', 'apache_beam.runners.direct.direct_runner.DirectRunner', 'apache_beam.runners.direct.direct_runner.SwitchingDirectRunner', 'apache_beam.runners.interactive.interactive_runner.InteractiveRunner', 'apache_beam.runners.portability.flink_runner.FlinkRunner', 'apache_beam.runners.portability.fn_api_runner.FnApiRunner', 'apache_beam.runners.portability.portable_runner.PortableRunner', 'apache_beam.runners.portability.prism_runner.PrismRunner', 'apache_beam.runners.portability.spark_runner.SparkRunner', 'apache_beam.runners.test.TestDirectRunner', 'apache_beam.runners.test.TestDataflowRunner', ) LOCAL_RUNNERS = ( 'BundleBasedDirectRunner', 'DirectRunner', 'SwitchingDirectRunner', 'FnApiRunner', 'PrismRunner', 'TestDirectRunner', ) KNOWN_RUNNER_NAMES = [path.split('.')[-1] for path in ALL_KNOWN_RUNNERS] @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--runner', help=( 'Pipeline runner used to execute the workflow. Valid values are ' 'one of %s, or the fully qualified name of a PipelineRunner ' 'subclass. If unspecified, defaults to %s.' % (', '.join(cls.KNOWN_RUNNER_NAMES), cls.DEFAULT_RUNNER))) # Whether to enable streaming mode. parser.add_argument( '--streaming', default=False, action='store_true', help='Whether to enable streaming mode.') parser.add_argument( '--resource_hint', '--resource_hints', dest='resource_hints', action='append', default=[], help=( 'Resource hint to set in the pipeline execution environment.' 'Hints specified via this option override hints specified ' 'at transform level. Interpretation of hints is defined by ' 'Beam runners.')) parser.add_argument( '--auto_unique_labels', default=False, action='store_true', help='Whether to automatically generate unique transform labels ' 'for every transform. The default behavior is to raise an ' 'exception if a transform is created with a non-unique label. ' 'Using --auto_unique_labels could cause data loss when ' 'updating a pipeline or reloading the job state. ' 'This is not recommended for streaming jobs.') parser.add_argument( '--no_wait_until_finish', default=False, action='store_true', help='By default, the "with" statement waits for the job to ' 'complete. Set this flag to bypass this behavior and continue ' 'execution immediately') class StreamingOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--update_compatibility_version', default=None, nargs='?', const=None, help='Attempt to produce a pipeline compatible with the given prior ' 'version of the Beam SDK. ' 'See for example, https://cloud.google.com/dataflow/docs/guides/' 'updating-a-pipeline') class CrossLanguageOptions(PipelineOptions): @staticmethod def _beam_services_from_enviroment(): return json.loads(os.environ.get('BEAM_SERVICE_OVERRIDES') or '{}') @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--beam_services', type=lambda s: { **cls._beam_services_from_enviroment(), **json.loads(s) }, default=cls._beam_services_from_enviroment(), help=( 'For convenience, Beam provides the ability to automatically ' 'download and start various services (such as expansion services) ' 'used at pipeline construction and execution. These services are ' 'identified by gradle target. This option provides the ability to ' 'use pre-started services or non-default pre-existing artifacts to ' 'start the given service. ' 'Should be a json mapping of gradle build targets to pre-built ' 'artifacts (e.g. jar files) or expansion endpoints ' '(e.g. host:port). Defaults to the value of BEAM_SERVICE_OVERRIDES ' 'from the environment.')) parser.add_argument( '--use_transform_service', default=False, action='store_true', help='Use the Docker-composed-based transform service when expanding ' 'cross-language transforms.') def additional_option_ptransform_fn(): beam.transforms.ptransform.ptransform_fn_typehints_enabled = True # Optional type checks that aren't enabled by default. additional_type_checks: dict[str, Callable[[], None]] = { 'ptransform_fn': additional_option_ptransform_fn, } def enable_all_additional_type_checks(): """Same as passing --type_check_additional=all.""" for f in additional_type_checks.values(): f() class TypeOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): # TODO(laolu): Add a type inferencing option here once implemented. parser.add_argument( '--type_check_strictness', default='DEFAULT_TO_ANY', choices=['ALL_REQUIRED', 'DEFAULT_TO_ANY'], help='The level of exhaustive manual type-hint ' 'annotation required') parser.add_argument( '--type_check_additional', default='', help='Comma separated list of additional type checking features to ' 'enable. Options: all, ptransform_fn. For details see:' 'https://beam.apache.org/documentation/sdks/python-type-safety/') parser.add_argument( '--no_pipeline_type_check', dest='pipeline_type_check', action='store_false', help='Disable type checking at pipeline construction ' 'time') parser.add_argument( '--disable_beartype', default=False, action='store_true', help='Disable the use of beartype for type checking.') parser.add_argument( '--runtime_type_check', default=False, action='store_true', help='Enable type checking at pipeline execution ' 'time. NOTE: only supported with the ' 'DirectRunner') parser.add_argument( '--performance_runtime_type_check', default=False, action='store_true', help='Enable faster type checking via sampling at pipeline execution ' 'time. NOTE: only supported with portable runners ' '(including the DirectRunner)') parser.add_argument( '--allow_non_deterministic_key_coders', default=False, action='store_true', help='Use non-deterministic coders (such as pickling) for key-grouping ' 'operations such as GroupByKey. This is unsafe, as runners may group ' 'keys based on their encoded bytes, but is available for backwards ' 'compatibility. See BEAM-11719.') parser.add_argument( '--allow_unsafe_triggers', default=False, action='store_true', help='Allow the use of unsafe triggers. Unsafe triggers have the ' 'potential to cause data loss due to finishing and/or never having ' 'their condition met. Some operations, such as GroupByKey, disallow ' 'this. This exists for cases where such loss is acceptable and for ' 'backwards compatibility. See BEAM-9487.') parser.add_argument( '--force_cloudpickle_deterministic_coders', default=False, action='store_true', help=( 'Force the use of cloudpickle-based deterministic coders ' 'instead of dill-based coders, even when ' 'update_compatibility_version would normally trigger dill usage ' 'for backward compatibility. This flag overrides automatic coder ' 'selection to always use the modern cloudpickle serialization ' ' path. Warning: May break pipeline update compatibility with ' ' SDK versions prior to 2.68.0.')) def validate(self, unused_validator): errors = [] if beam.version.__version__ >= '3': errors.append( 'Update --type_check_additional default to include all ' 'available additional checks at Beam 3.0 release time.') keys = self.type_check_additional.split(',') for key in keys: if not key: continue elif key == 'all': enable_all_additional_type_checks() elif key in additional_type_checks: additional_type_checks[key]() else: errors.append('Unrecognized --type_check_additional feature: %s' % key) return errors class DirectOptions(PipelineOptions): """DirectRunner-specific execution options.""" @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--no_direct_runner_use_stacked_bundle', action='store_false', dest='direct_runner_use_stacked_bundle', help='DirectRunner uses stacked WindowedValues within a Bundle for ' 'memory optimization. Set --no_direct_runner_use_stacked_bundle to ' 'avoid it.') parser.add_argument( '--direct_runner_bundle_repeat', type=int, default=0, help='replay every bundle this many extra times, for profiling' 'and debugging') parser.add_argument( '--direct_num_workers', type=int, default=1, help='number of parallel running workers.') parser.add_argument( '--direct_running_mode', default='in_memory', choices=['in_memory', 'multi_threading', 'multi_processing'], help='Workers running environment.') parser.add_argument( '--direct_embed_docker_python', default=False, action='store_true', dest='direct_embed_docker_python', help='DirectRunner uses the embedded Python environment when ' 'the default Python docker environment is specified.') parser.add_argument( '--direct_test_splits', default={}, type=json.loads, help='Split test configuration of the json form ' '{"step_name": {"timings": [...], "fractions": [...]}, ...} ' 'where step_name is the name of a step controlling the stage to which ' 'splits will be sent, timings is a list of floating-point times ' '(in seconds) at which the split requests will be sent, and ' 'fractions is a corresponding list of floating points to use in the ' 'split requests themselves.') class GoogleCloudOptions(PipelineOptions): """Google Cloud Dataflow service execution options.""" BIGQUERY_API_SERVICE = 'bigquery.googleapis.com' COMPUTE_API_SERVICE = 'compute.googleapis.com' STORAGE_API_SERVICE = 'storage.googleapis.com' DATAFLOW_ENDPOINT = 'https://dataflow.googleapis.com' OAUTH_SCOPES = [ 'https://www.googleapis.com/auth/bigquery', 'https://www.googleapis.com/auth/cloud-platform', 'https://www.googleapis.com/auth/devstorage.full_control', 'https://www.googleapis.com/auth/userinfo.email', 'https://www.googleapis.com/auth/datastore', 'https://www.googleapis.com/auth/spanner.admin', 'https://www.googleapis.com/auth/spanner.data' ] @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--dataflow_endpoint', default=cls.DATAFLOW_ENDPOINT, help=( 'The URL for the Dataflow API. If not set, the default public URL ' 'will be used.')) # Remote execution must check that this option is not None. parser.add_argument( '--project', default=None, help='Name of the Cloud project owning the Dataflow ' 'job.') # Remote execution must check that this option is not None. parser.add_argument( '--job_name', default=None, help='Name of the Cloud Dataflow job.') # Remote execution must check that this option is not None. parser.add_argument( '--staging_location', default=None, help='GCS path for staging code packages needed by ' 'workers.') # Remote execution must check that this option is not None. # If staging_location is not set, it defaults to temp_location. parser.add_argument( '--temp_location', default=None, help='GCS path for saving temporary workflow jobs.') # The Google Compute Engine region for creating Dataflow jobs. See # https://cloud.google.com/compute/docs/regions-zones/regions-zones for a # list of valid options. parser.add_argument( '--region', default=None, help='The Google Compute Engine region for creating ' 'Dataflow job.') parser.add_argument( '--service_account_email', default=None, help='Identity to run virtual machines as.') parser.add_argument( '--no_auth', dest='no_auth', action='store_true', default=False, help='Skips authorizing credentials with Google Cloud.') # Option to run templated pipelines parser.add_argument( '--template_location', default=None, help='Save job to specified local or GCS location.') parser.add_argument( '--label', '--labels', dest='labels', action='append', default=None, help='Labels to be applied to this Dataflow job. ' 'Labels are key value pairs separated by = ' '(e.g. --label key=value) or ' '(--labels=\'{ "key": "value", "mass": "1_3kg", "count": "3" }\').') parser.add_argument( '--update', default=False, action='store_true', help='Update an existing streaming Cloud Dataflow job. ' 'See https://cloud.google.com/dataflow/docs/guides/' 'updating-a-pipeline') parser.add_argument( '--transform_name_mapping', default=None, type=json.loads, help='The transform mapping that maps the named ' 'transforms in your prior pipeline code to names ' 'in your replacement pipeline code.' 'See https://cloud.google.com/dataflow/docs/guides/' 'updating-a-pipeline') parser.add_argument( '--enable_streaming_engine', default=True, action='store_true', help='Deprecated. All Python streaming pipelines on Dataflow' 'use Streaming Engine.') parser.add_argument( '--dataflow_kms_key', default=None, help='Set a Google Cloud KMS key name to be used in ' 'Dataflow state operations (GBK, Streaming).') parser.add_argument( '--create_from_snapshot', default=None, help='The snapshot from which the job should be created.') parser.add_argument( '--flexrs_goal', default=None, choices=['COST_OPTIMIZED', 'SPEED_OPTIMIZED'], help='Set the Flexible Resource Scheduling mode') parser.add_argument( '--dataflow_service_option', '--dataflow_service_options', dest='dataflow_service_options', action=_CommaSeparatedListAction, default=None, help=( 'Options to configure the Dataflow service. These ' 'options decouple service side feature availability ' 'from the Apache Beam release cycle.' 'Note: If set programmatically, must be set as a ' 'list of strings')) parser.add_argument( '--enable_hot_key_logging', default=False, action='store_true', help='When true, will enable the direct logging of any detected hot ' 'keys into Cloud Logging. Warning: this will log the literal key as an ' 'unobfuscated string.') parser.add_argument( '--enable_artifact_caching', default=False, action='store_true', help='When true, artifacts will be cached across job submissions in ' 'the GCS staging bucket') parser.add_argument( '--impersonate_service_account', default=None, help='All API requests will be made as the given service account or ' 'target service account in an impersonation delegation chain ' 'instead of the currently selected account. You can specify ' 'either a single service account as the impersonator, or a ' 'comma-separated list of service accounts to create an ' 'impersonation delegation chain.') parser.add_argument( '--gcp_oauth_scope', '--gcp_oauth_scopes', dest='gcp_oauth_scopes', action='append', default=cls.OAUTH_SCOPES, help=( 'Controls the OAuth scopes that will be requested when creating ' 'GCP credentials. Note: If set programmatically, must be set as a ' 'list of strings')) parser.add_argument( '--enable_bucket_read_metric_counter', default=False, action='store_true', help='Create metrics reporting the approximate number of bytes read ' 'per bucket.') parser.add_argument( '--enable_bucket_write_metric_counter', default=False, action='store_true', help= 'Create metrics reporting the approximate number of bytes written per ' 'bucket.') parser.add_argument( '--no_gcsio_throttling_counter', default=False, action='store_true', help='Throttling counter in GcsIO is enabled by default. Set ' '--no_gcsio_throttling_counter to avoid it.') parser.add_argument( '--enable_gcsio_blob_generation', default=False, action='store_true', help='Use blob generation when mutating blobs in GCSIO to ' 'mitigate race conditions at the cost of more HTTP requests.') parser.add_argument( '--gcs_custom_audit_entry', '--gcs_custom_audit_entries', dest='gcs_custom_audit_entries', action=_GcsCustomAuditEntriesAction, default=None, help='Custom information to be attached to audit logs. ' 'Entries are key value pairs separated by = ' '(e.g. --gcs_custom_audit_entry key=value) or a JSON string ' '(e.g. --gcs_custom_audit_entries=\'{ "user": "test", "id": "12" }\').') def _create_default_gcs_bucket(self): try: from apache_beam.io.gcp import gcsio except ImportError: _LOGGER.warning('Unable to create default GCS bucket.') return None bucket = gcsio.get_or_create_default_gcs_bucket(self) if bucket: return 'gs://%s/' % bucket.id else: return None # Log warning if soft delete policy is enabled in a gcs bucket # that is specified in an argument. def _warn_if_soft_delete_policy_enabled(self, arg_name): # skip the check if it is in dry-run mode because the later step requires # internet connection to access GCS if self.view_as(TestOptions).dry_run: return gcs_path = getattr(self, arg_name, None) try: from apache_beam.io.gcp import gcsio if gcsio.GcsIO().is_soft_delete_enabled(gcs_path): logger.log_first_n( logging.WARN, "Bucket %s used as %s has soft-delete policy enabled." " To avoid being billed for unnecessary storage costs, turn" " off the soft delete feature on buckets that your Dataflow" " jobs use for temporary and staging storage. For more" " information, see" " https://cloud.google.com/storage/docs/use-soft-delete" "#remove-soft-delete-policy.", gcs_path, arg_name, n=1, key="message") except ImportError: _LOGGER.warning('Unable to check soft delete policy due to import error.') # If either temp or staging location has an issue, we use the valid one for # both locations. If both are bad we return an error. def _handle_temp_and_staging_locations(self, validator): temp_errors = validator.validate_gcs_path(self, 'temp_location') staging_errors = validator.validate_gcs_path(self, 'staging_location') if temp_errors and not staging_errors: setattr(self, 'temp_location', getattr(self, 'staging_location')) self._warn_if_soft_delete_policy_enabled('staging_location') return [] elif staging_errors and not temp_errors: setattr(self, 'staging_location', getattr(self, 'temp_location')) self._warn_if_soft_delete_policy_enabled('temp_location') return [] elif not staging_errors and not temp_errors: self._warn_if_soft_delete_policy_enabled('temp_location') self._warn_if_soft_delete_policy_enabled('staging_location') return [] # Both staging and temp locations are bad, try to use default bucket. else: default_bucket = self._create_default_gcs_bucket() if default_bucket is None: temp_errors.extend(staging_errors) return temp_errors else: setattr(self, 'temp_location', default_bucket) setattr(self, 'staging_location', default_bucket) self._warn_if_soft_delete_policy_enabled('temp_location') self._warn_if_soft_delete_policy_enabled('staging_location') return [] def validate(self, validator): errors = [] if validator.is_service_runner(): errors.extend(self._handle_temp_and_staging_locations(validator)) errors.extend(validator.validate_cloud_options(self)) if self.view_as(DebugOptions).dataflow_job_file: if self.view_as(GoogleCloudOptions).template_location: errors.append( '--dataflow_job_file and --template_location ' 'are mutually exclusive.') # Validate that dataflow_service_options is a list if self.dataflow_service_options: errors.extend( validator.validate_repeatable_argument_passed_as_list( self, 'dataflow_service_options')) return errors def get_cloud_profiler_service_name(self): _ENABLE_GOOGLE_CLOUD_PROFILER = 'enable_google_cloud_profiler' if self.dataflow_service_options: if _ENABLE_GOOGLE_CLOUD_PROFILER in self.dataflow_service_options: return os.environ["JOB_NAME"] for option_name in self.dataflow_service_options: if option_name.startswith(_ENABLE_GOOGLE_CLOUD_PROFILER + '='): return option_name.split('=', 1)[1] experiments = self.view_as(DebugOptions).experiments or [] if _ENABLE_GOOGLE_CLOUD_PROFILER in experiments: return os.environ["JOB_NAME"] return None class AzureOptions(PipelineOptions): """Azure Blob Storage options.""" @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--azure_connection_string', default=None, help='Connection string of the Azure Blob Storage Account.') parser.add_argument( '--blob_service_endpoint', default=None, help='URL of the Azure Blob Storage Account.') parser.add_argument( '--azure_managed_identity_client_id', default=None, help='Client ID of a user-assigned managed identity.') def validate(self, validator): errors = [] if self.azure_connection_string: if self.blob_service_endpoint: errors.append( '--azure_connection_string and ' '--blob_service_endpoint are mutually exclusive.') return errors class HadoopFileSystemOptions(PipelineOptions): """``HadoopFileSystem`` connection options.""" @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--hdfs_host', default=None, help=('Hostname or address of the HDFS namenode.')) parser.add_argument( '--hdfs_port', default=None, help=('Port of the HDFS namenode.')) parser.add_argument( '--hdfs_user', default=None, help=('HDFS username to use.')) parser.add_argument( '--hdfs_full_urls', default=False, action='store_true', help=( 'If set, URLs will be parsed as "hdfs://server/path/...", instead ' 'of "hdfs://path/...". The "server" part will be unused (use ' '--hdfs_host and --hdfs_port).')) def validate(self, validator): errors = [] errors.extend(validator.validate_optional_argument_positive(self, 'port')) return errors # Command line options controlling the worker pool configuration. # TODO(silviuc): Update description when autoscaling options are in. class WorkerOptions(PipelineOptions): """Worker pool configuration options.""" @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--num_workers', type=int, default=None, help=( 'Number of workers to use when executing the Dataflow job. If not ' 'set, the Dataflow service will use a reasonable default.')) parser.add_argument( '--max_num_workers', type=int, default=None, help=( 'Maximum number of workers to use when executing the Dataflow job.' )) parser.add_argument( '--autoscaling_algorithm', type=str, choices=['NONE', 'THROUGHPUT_BASED'], default=None, # Meaning unset, distinct from 'NONE' meaning don't scale help=('If and how to autoscale the workerpool.')) parser.add_argument( '--worker_machine_type', '--machine_type', dest='machine_type', default=None, help=( 'Machine type to create Dataflow worker VMs as. See ' 'https://cloud.google.com/compute/docs/machine-types ' 'for a list of valid options. If not set, ' 'the Dataflow service will choose a reasonable ' 'default.')) parser.add_argument( '--disk_size_gb', type=int, default=None, help=( 'Remote worker disk size, in gigabytes, or 0 to use the default ' 'size. If not set, the Dataflow service will use a reasonable ' 'default.')) parser.add_argument( '--worker_disk_type', '--disk_type', dest='disk_type', default=None, help=('Specifies what type of persistent disk should be used.')) parser.add_argument( '--disk_provisioned_iops', type=int, default=None, dest='disk_provisioned_iops', help=( 'The provisioned IOPS of the disk. If not set, the Dataflow service' ' will choose a reasonable default.'), ) parser.add_argument( '--disk_provisioned_throughput_mibps', type=int, default=None, dest='disk_provisioned_throughput_mibps', help=( 'The provisioned throughput of the disk in MiB/s. If not set, the' ' Dataflow service will choose a reasonable default.'), ) parser.add_argument( '--worker_region', default=None, help=( 'The Compute Engine region (https://cloud.google.com/compute/docs/' 'regions-zones/regions-zones) in which worker processing should ' 'occur, e.g. "us-west1". Mutually exclusive with worker_zone. If ' 'neither worker_region nor worker_zone is specified, default to ' 'same value as --region.')) parser.add_argument( '--worker_zone', default=None, help=( 'The Compute Engine zone (https://cloud.google.com/compute/docs/' 'regions-zones/regions-zones) in which worker processing should ' 'occur, e.g. "us-west1-a". Mutually exclusive with worker_region. ' 'If neither worker_region nor worker_zone is specified, the ' 'Dataflow service will choose a zone in --region based on ' 'available capacity.')) parser.add_argument( '--zone', default=None, help=( 'GCE availability zone for launching workers. Default is up to the ' 'Dataflow service. This flag is deprecated, and will be replaced ' 'by worker_zone.')) parser.add_argument( '--network', default=None, help=( 'GCE network for launching workers. Default is up to the Dataflow ' 'service.')) parser.add_argument( '--subnetwork', default=None, help=( 'GCE subnetwork for launching workers. Default is up to the ' 'Dataflow service. Expected format is ' 'regions/REGION/subnetworks/SUBNETWORK or the fully qualified ' 'subnetwork name. For more information, see ' 'https://cloud.google.com/compute/docs/vpc/')) parser.add_argument( '--worker_harness_container_image', default=None, help=( 'Docker registry location of container image to use for the ' 'worker harness. If not set, an appropriate approved Google Cloud ' 'Dataflow image will be used based on the version of the ' 'SDK. Note: This flag is deprecated and only supports ' 'approved Google Cloud Dataflow container images. To provide a ' 'custom container image, use sdk_container_image instead.')) parser.add_argument( '--sdk_container_image', default=None, help=( 'Docker registry location of container image to use for the ' 'worker harness. If not set, an appropriate approved Google Cloud ' 'Dataflow image will be used based on the version of the ' 'SDK. If set for a non-portable pipeline, only official ' 'Google Cloud Dataflow container images may be used here.')) parser.add_argument( '--sdk_harness_container_image_overrides', action='append', default=None, help=( 'Overrides for SDK harness container images. Could be for the ' 'local SDK or for a remote SDK that pipeline has to support due ' 'to a cross-language transform. Each entry consist of two values ' 'separated by a comma where first value gives a regex to ' 'identify the container image to override and the second value ' 'gives the replacement container image.')) parser.add_argument( '--default_sdk_harness_log_level', default=None, help=( 'Controls the default log level of all loggers without a log level ' 'override. Values can be either a labeled level or a number ' '(See https://docs.python.org/3/library/logging.html#levels). ' 'Default log level is INFO.')) parser.add_argument( '--sdk_harness_log_level_overrides', type=json.loads, action=_DictUnionAction, default=None, help=( 'Controls the log levels for specifically named loggers. The ' 'expected format is a json string: \'{"module":"log_level",...}\'. ' 'For example, by specifying the value \'{"a.b.c":"DEBUG"}\', ' 'the logger underneath the module "a.b.c" will be configured to ' 'output logs at the DEBUG level. Similarly, by specifying the ' 'value \'{"a.b.c":"WARNING"}\' all loggers underneath the "a.b.c" ' 'module will be configured to output logs at the WARNING level. ' 'Also, note that when multiple overrides are specified, the exact ' 'name followed by the closest parent takes precedence.')) parser.add_argument( '--use_public_ips', default=None, action='store_true', help='Whether to assign public IP addresses to the worker VMs.') parser.add_argument( '--no_use_public_ips', dest='use_public_ips', default=None, action='store_false', help='Whether to assign only private IP addresses to the worker VMs.') parser.add_argument( '--min_cpu_platform', dest='min_cpu_platform', type=str, help='GCE minimum CPU platform. Default is determined by GCP.') parser.add_argument( '--max_cache_memory_usage_mb', dest='max_cache_memory_usage_mb', type=int, default=0, help=( 'Size of the SDK Harness cache to store user state and side ' 'inputs in MB. The cache is disabled by default. Increasing ' 'cache size might improve performance of some pipelines, such as ' 'pipelines that use iterable side input views, but can ' 'lead to an increase in memory consumption and OOM errors if ' 'workers are not appropriately provisioned. ' 'Using the cache might decrease performance pipelines using ' 'materialized side inputs. ' 'If the cache is full, least ' 'recently used elements will be evicted. This cache is per ' 'each SDK Harness instance. SDK Harness is a component ' 'responsible for executing the user code and communicating with ' 'the runner. Depending on the runner, there may be more than one ' 'SDK Harness process running on the same worker node.')) parser.add_argument( '--element_processing_timeout_minutes', type=int, default=None, help=( 'The time limit (in minutes) for any PTransform to finish ' 'processing a single element. If exceeded, the SDK worker ' 'process self-terminates and processing may be restarted ' 'by a runner.')) def validate(self, validator): errors = [] errors.extend(validator.validate_sdk_container_image_options(self)) if validator.is_service_runner(): errors.extend(validator.validate_num_workers(self)) errors.extend(validator.validate_worker_region_zone(self)) return errors class DebugOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--dataflow_job_file', default=None, help='Debug file to write the workflow specification.') parser.add_argument( '--experiment', '--experiments', dest='experiments', action=_CommaSeparatedListAction, default=None, help=( 'Runners may provide a number of experimental features that can be ' 'enabled with this flag. Please sync with the owners of the runner ' 'before enabling any experiments.')) parser.add_argument( '--number_of_worker_harness_threads', type=int, default=None, help=( 'Number of threads per worker to use on the runner. If left ' 'unspecified, the runner will compute an appropriate number of ' 'threads to use.')) def add_experiment(self, experiment): # pylint: disable=access-member-before-definition if self.experiments is None: self.experiments = [] if experiment not in self.experiments: self.experiments.append(experiment) def lookup_experiment(self, key, default=None): if not self.experiments: return default elif key in self.experiments: return True for experiment in self.experiments: if experiment.startswith(key + '='): return experiment.split('=', 1)[1] return default def validate(self, validator): errors = [] if self.experiments: errors.extend( validator.validate_repeatable_argument_passed_as_list( self, 'experiments')) return errors class ProfilingOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--profile_cpu', action='store_true', help='Enable work item CPU profiling.') parser.add_argument( '--profile_memory', action='store_true', help='Enable work item heap profiling.') parser.add_argument( '--profile_location', default=None, help='path for saving profiler data.') parser.add_argument( '--profile_sample_rate', type=float, default=1.0, help='A number between 0 and 1 indicating the ratio ' 'of bundles that should be profiled.') class SetupOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): # Options for installing dependencies in the worker. parser.add_argument( '--requirements_file', default=None, help=( 'Path to a requirements file containing package dependencies. ' 'Typically it is produced by a pip freeze command. More details: ' 'https://pip.pypa.io/en/latest/reference/pip_freeze.html. ' 'If used, all the packages specified will be downloaded, ' 'cached (use --requirements_cache to change default location), ' 'and then staged so that they can be automatically installed in ' 'workers during startup. The cache is refreshed as needed ' 'avoiding extra downloads for existing packages. Typically the ' 'file is named requirements.txt.')) parser.add_argument( '--requirements_cache', default=None, help=( 'Path to a folder to cache the packages specified in ' 'the requirements file using the --requirements_file option.' 'If you want to skip populating requirements cache, please ' 'specify --requirements_cache="skip".')) parser.add_argument( '--requirements_cache_only_sources', action='store_true', help=( 'Enable this flag to populate requirements cache only ' 'with Source distributions(sdists) of the dependencies ' 'mentioned in the --requirements_file' 'Note: (BEAM-4032): This flag may significantly slow down ' 'the pipeline submission. It is added to preserve the requirements' ' cache behavior prior to 2.37.0 and will likely be removed in ' 'future releases.')) parser.add_argument( '--setup_file', default=None, help=( 'Path to a setup Python file containing package dependencies. If ' 'specified, the file\'s containing folder is assumed to have the ' 'structure required for a setuptools setup package. The file must ' 'be named setup.py. More details: ' 'https://pythonhosted.org/an_example_pypi_project/setuptools.html ' 'During job submission a source distribution will be built and ' 'the worker will install the resulting package before running any ' 'custom code.')) parser.add_argument( '--beam_plugin', '--beam_plugins', dest='beam_plugins', action='append', default=None, help=( 'Bootstrap the python process before executing any code by ' 'importing all the plugins used in the pipeline. Please pass a ' 'comma separated list of import paths to be included. This is ' 'currently an experimental flag and provides no stability. ' 'Multiple --beam_plugin options can be specified if more than ' 'one plugin is needed.')) parser.add_argument( '--pickle_library', default='default', help=( 'Chooses which pickle library to use. Options are dill, ' 'cloudpickle or default.'), choices=['cloudpickle', 'default', 'dill', 'dill_unsafe']) parser.add_argument( '--save_main_session', default=None, action='store_true', help=( 'Save the main session state so that pickled functions and classes ' 'defined in __main__ (e.g. interactive session) can be unpickled. ' 'Some workflows do not need the session state if for instance all ' 'their functions/classes are defined in proper modules ' '(not __main__) and the modules are importable in the worker. ' 'It is disabled by default except for cloudpickle as pickle ' 'library on Dataflow runner.')) parser.add_argument( '--no_save_main_session', default=None, action='store_false', dest='save_main_session', help=( 'Disable saving the main session state. See "save_main_session".')) parser.add_argument( '--sdk_location', default='default', help=( 'Path to a custom Beam SDK package to install and use on the' 'runner. It can be a URL, a GCS path, or a local path to an ' 'SDK tarball. Workflow submissions will download or copy an SDK ' 'tarball from here. If set to "default", ' 'runners will use the SDK provided in the default environment.' 'Use this flag when running pipelines with an unreleased or ' 'manually patched version of Beam SDK.')) parser.add_argument( '--extra_package', '--extra_packages', dest='extra_packages', action='append', default=None, help=( 'Local path to a Python package file. The file is expected to be ' '(1) a package tarball (".tar"), (2) a compressed package tarball ' '(".tar.gz"), (3) a Wheel file (".whl") or (4) a compressed ' 'package zip file (".zip") which can be installed using the ' '"pip install" command of the standard pip package. Multiple ' '--extra_package options can be specified if more than one ' 'package is needed. During job submission, the files will be ' 'staged in the staging area (--staging_location option) and the ' 'workers will install them in same order they were specified on ' 'the command line.')) parser.add_argument( '--files_to_stage', dest='files_to_stage', action='append', default=None, help=( 'Local path to a file. During job submission, the files will be ' 'staged in the staging area (--staging_location option) and then ' 'workers will upload them to the worker specific staging location ' '(e.g. $SEMI_PERSISTENT_DIRECTORY/staged/ for portable runner.')) parser.add_argument( '--prebuild_sdk_container_engine', help=( 'Prebuild sdk worker container image before job submission. If ' 'enabled, SDK invokes the boot sequence in SDK worker ' 'containers to install all pipeline dependencies in the ' 'container, and uses the prebuilt image in the pipeline ' 'environment. This may speed up pipeline execution. To enable, ' 'select the Docker build engine: local_docker using ' 'locally-installed Docker or cloud_build for using Google Cloud ' 'Build (requires a GCP project with Cloud Build API enabled). You ' 'can also subclass SdkContainerImageBuilder and use that to build ' 'in other environments.')) parser.add_argument( '--prebuild_sdk_container_base_image', default=None, help=('Deprecated. Use --sdk_container_image instead.')) parser.add_argument( '--cloud_build_machine_type', default=None, help=( 'If specified, use the machine type explicitly when prebuilding' 'SDK container image on Google Cloud Build.')) parser.add_argument( '--docker_registry_push_url', default=None, help=( 'Docker registry url to use for tagging and pushing the prebuilt ' 'sdk worker container image.')) parser.add_argument( '--gbek', default=None, help=( 'When set, will replace all GroupByKey transforms in the pipeline ' 'with EncryptedGroupByKey transforms using the secret passed in ' 'the option. Beam will infer the secret type and value based on ' 'secret itself. This guarantees that any data at rest during the ' 'GBK will be encrypted. Many runners only store data at rest when ' 'performing a GBK, so this can be used to guarantee that data is ' 'not unencrypted. The secret should be a url safe base64 encoded ' '32 byte value. To generate a secret in this format, you can use ' 'Secret.generate_secret_bytes(). For an example of this, see ' 'https://github.com/apache/beam/blob/c8df4da229da49d533491857e1bb4ab5dbf4fd37/sdks/python/apache_beam/transforms/util_test.py#L356. ' # pylint: disable=line-too-long 'Runners with this behavior include the Dataflow, ' 'Flink, and Spark runners. The option should be ' 'structured like: ' '--gbek=type:;:, for example ' '--gbek=type:GcpSecret;version_name:my_secret/versions/latest')) parser.add_argument( '--user_agent', default=None, help=( 'A user agent string describing the pipeline to external services. ' 'The format should follow RFC2616.')) parser.add_argument( '--maven_repository_url', default=None, help=( 'Custom Maven repository URL to use for downloading JAR files. ' 'If not specified, the default Maven Central repository will be ' 'used.')) def _handle_load_main_session(self, validator): save_main_session = getattr(self, 'save_main_session') if save_main_session is None: if not validator.is_service_runner(): setattr(self, 'save_main_session', False) else: # save_main_session default to False for dill, while default to true # for cloudpickle on service runner pickle_library = getattr(self, 'pickle_library') if pickle_library == 'default': from apache_beam.internal.pickler import DEFAULT_PICKLE_LIB pickle_library = DEFAULT_PICKLE_LIB if pickle_library == 'cloudpickle': setattr(self, 'save_main_session', True) else: setattr(self, 'save_main_session', False) return [] def validate(self, validator): errors = [] errors.extend(validator.validate_container_prebuilding_options(self)) errors.extend(validator.validate_pickle_library(self)) errors.extend(self._handle_load_main_session(validator)) return errors class PortableOptions(PipelineOptions): """Portable options are common options expected to be understood by most of the portable runners. Should generally be kept in sync with PortablePipelineOptions.java. """ @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--job_endpoint', default=None, help=( 'Job service endpoint to use. Should be in the form of host ' 'and port, e.g. localhost:8099.')) parser.add_argument( '--artifact_endpoint', default=None, help=( 'Artifact staging endpoint to use. Should be in the form of host ' 'and port, e.g. localhost:8098. If none is specified, the ' 'artifact endpoint sent from the job server is used.')) parser.add_argument( '--job_server_timeout', '--job-server-timeout', # For backwards compatibility. default=300, type=int, help=( 'Job service request timeout in seconds. The timeout ' 'determines the max time the driver program will wait to ' 'get a response from the job server. NOTE: the timeout does not ' 'apply to the actual pipeline run time. The driver program can ' 'still wait for job completion indefinitely.')) parser.add_argument( '--environment_type', default=None, help=( 'Set the default environment type for running ' 'user code. DOCKER (default) runs user code in a container. ' 'PROCESS runs user code in processes that are automatically ' 'started on each worker node. LOOPBACK runs user code on the ' 'same process that originally submitted the job.')) parser.add_argument( '--environment_config', default=None, help=( 'Set environment configuration for running the user code.\n For ' 'DOCKER: Url for the docker image.\n For PROCESS: json of the ' 'form {"os": "", "arch": "", "command": ' '"", "env":{"": ' '""} }. All fields in the json are optional except ' 'command.\n\nPrefer using --environment_options instead.')) parser.add_argument( '--environment_option', '--environment_options', dest='environment_options', action='append', default=None, help=( 'Environment configuration for running the user code. ' 'Recognized options depend on --environment_type.\n ' 'For DOCKER: docker_container_image (optional)\n ' 'For PROCESS: process_command (required), process_variables ' '(optional, comma-separated)\n ' 'For EXTERNAL: external_service_address (required)')) parser.add_argument( '--sdk_worker_parallelism', default=1, help=( 'Sets the number of sdk worker processes that will run on each ' 'worker node. Default is 1. If 0, a value will be chosen by the ' 'runner.')) parser.add_argument( '--environment_cache_millis', default=0, help=( 'Duration in milliseconds for environment cache within a job. ' '0 means no caching.')) parser.add_argument( '--output_executable_path', default=None, help=( 'Create an executable jar at this path rather than running ' 'the pipeline.')) def validate(self, validator): return validator.validate_environment_options(self) def add_environment_option(self, option): # pylint: disable=access-member-before-definition if self.environment_options is None: self.environment_options = [] if option not in self.environment_options: self.environment_options.append(option) def lookup_environment_option(self, key, default=None): if not self.environment_options: return default elif key in self.environment_options: return True for option in self.environment_options: if option.startswith(key + '='): return option.split('=', 1)[1] return default class JobServerOptions(PipelineOptions): """Options for starting a Beam job server. Roughly corresponds to JobServerDriver.ServerConfiguration in Java. """ @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--artifacts_dir', default=None, help='The location to store staged artifact files. ' 'Any Beam-supported file system is allowed. ' 'If unset, the local temp dir will be used.') parser.add_argument( '--job_port', default=0, type=int, help='Port to use for the job service. 0 to use a ' 'dynamic port.') parser.add_argument( '--artifact_port', default=0, type=int, help='Port to use for artifact staging. 0 to use a ' 'dynamic port.') parser.add_argument( '--expansion_port', default=0, type=int, help='Port to use for artifact staging. 0 to use a ' 'dynamic port.') parser.add_argument( '--job_server_java_launcher', default='java', help='The Java Application Launcher executable file to use for ' 'starting a Java job server. If unset, `java` from the ' 'environment\'s $PATH is used.') parser.add_argument( '--job_server_jvm_properties', '--job_server_jvm_property', dest='job_server_jvm_properties', action='append', default=[], help='JVM properties to pass to a Java job server.') parser.add_argument( '--jar_cache_dir', default=None, help='The location to store jar cache for job server.') class FlinkRunnerOptions(PipelineOptions): # These should stay in sync with gradle.properties. PUBLISHED_FLINK_VERSIONS = ['1.17', '1.18', '1.19', '1.20', '2.0'] @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--flink_master', default='[auto]', help='Flink master address (http://host:port)' ' Use "[local]" to start a local cluster' ' for the execution. Use "[auto]" if you' ' plan to either execute locally or let the' ' Flink job server infer the cluster address.') parser.add_argument( '--flink_version', default=cls.PUBLISHED_FLINK_VERSIONS[-1], choices=cls.PUBLISHED_FLINK_VERSIONS, help='Flink version to use.') parser.add_argument( '--flink_job_server_jar', help='Path or URL to a flink jobserver jar.') parser.add_argument( '--flink_submit_uber_jar', default=False, action='store_true', help='Create and upload an uberjar to the flink master' ' directly, rather than starting up a job server.' ' Only applies when flink_master is set to a' ' cluster address. Requires Python 3.6+.') parser.add_argument( '--parallelism', type=int, default=-1, help='The degree of parallelism to be used when distributing' ' operations onto workers. If the parallelism is not set, the' ' configured Flink default is used, or 1 if none can be found.') parser.add_argument( '--max_parallelism', type=int, default=-1, help='The pipeline wide maximum degree of parallelism to be used. The' ' maximum parallelism specifies the upper limit for dynamic scaling' ' and the number of key groups used for partitioned state.') class SparkRunnerOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--spark_master_url', default='local[4]', help='Spark master URL (spark://HOST:PORT). ' 'Use "local" (single-threaded) or "local[*]" ' '(multi-threaded) to start a local cluster for ' 'the execution.') parser.add_argument( '--spark_job_server_jar', help='Path or URL to a Beam Spark job server jar. ' 'Overrides --spark_version.') parser.add_argument( '--spark_submit_uber_jar', default=False, action='store_true', help='Create and upload an uber jar to the Spark REST' ' endpoint, rather than starting up a job server.' ' Requires Python 3.6+.') parser.add_argument( '--spark_rest_url', help='URL for the Spark REST endpoint. ' 'Only required when using spark_submit_uber_jar. ' 'For example, http://hostname:6066') parser.add_argument( '--spark_version', default='3', choices=['3'], help='Spark major version to use.') class PrismRunnerOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--prism_location', help='Path or URL to a prism binary, or zipped binary for the current ' 'platform (Operating System and Architecture). May also be an Apache ' 'Beam Github Release page URL, with a matching beam_version_override ' 'set. This option overrides all others for finding a prism binary.') parser.add_argument( '--prism_beam_version_override', help= 'Override the SDK\'s version for deriving the Github Release URLs for ' 'downloading a zipped prism binary, for the current platform. If ' 'prism_location is set to a Github Release page URL, them it will use ' 'that release page as a base when constructing the download URL.') parser.add_argument( '--prism_log_level', default="info", choices=["debug", "info", "warn", "error"], help=( 'Controls the log level in Prism. Values can be "debug", "info", ' '"warn", and "error". Default log level is "info".')) parser.add_argument( '--prism_log_kind', default="console", choices=["dev", "json", "text", "console"], help=( 'Controls the log format in Prism. Values can be "dev", "json", ' '"text", and "console". Default log format is "console".')) class TestOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): # Options for e2e test pipeline. parser.add_argument( '--on_success_matcher', default=None, help=( 'Verify state/output of e2e test pipeline. This is pickled ' 'version of the matcher which should extends ' 'hamcrest.core.base_matcher.BaseMatcher.')) parser.add_argument( '--dry_run', default=False, help=( 'Used in unit testing runners without submitting the ' 'actual job.')) parser.add_argument( '--wait_until_finish_duration', default=None, type=int, help='The time to wait (in milliseconds) for test pipeline to finish. ' 'If it is set to None, it will wait indefinitely until the job ' 'is finished.') def validate(self, validator): errors = [] if self.view_as(TestOptions).on_success_matcher: errors.extend(validator.validate_test_matcher(self, 'on_success_matcher')) return errors class TestDataflowOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): # This option is passed to Dataflow Runner's Pub/Sub client. The camelCase # style in 'dest' matches the runner's. parser.add_argument( '--pubsub_root_url', dest='pubsubRootUrl', default=None, help='Root URL for use with the Google Cloud Pub/Sub API.', ) # TODO(silviuc): Non-standard options. Keep them? If yes, add help too! # Remote execution must check that this option is not None. class OptionsContext(object): """Set default pipeline options for pipelines created in this block. This is particularly useful for pipelines implicitly created with the [python list] | PTransform construct. Can also be used as a decorator. """ overrides: list[dict[str, Any]] = [] def __init__(self, **options): self.options = options def __enter__(self): self.overrides.append(self.options) def __exit__(self, *exn_info): self.overrides.pop() def __call__(self, f, *args, **kwargs): def wrapper(*args, **kwargs): with self: f(*args, **kwargs) return wrapper @classmethod def augment_options(cls, options): for override in cls.overrides: for name, value in override.items(): setattr(options, name, value) return options class S3Options(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): # These options are passed to the S3 IO Client parser.add_argument( '--s3_access_key_id', default=None, help='The secret key to use when creating the s3 client.') parser.add_argument( '--s3_secret_access_key', default=None, help='The secret key to use when creating the s3 client.') parser.add_argument( '--s3_session_token', default=None, help='The session token to use when creating the s3 client.') parser.add_argument( '--s3_endpoint_url', default=None, help='The complete URL to use for the constructed s3 client.') parser.add_argument( '--s3_region_name', default=None, help='The name of the region associated with the s3 client.') parser.add_argument( '--s3_api_version', default=None, help='The API version to use with the s3 client.') parser.add_argument( '--s3_verify', default=None, help='Whether or not to verify SSL certificates with the s3 client.') parser.add_argument( '--s3_disable_ssl', default=False, action='store_true', help=( 'Whether or not to use SSL with the s3 client. ' 'By default, SSL is used.')) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/options/pipeline_options_context.py0000644000175100017510000000450115206113032025046 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Context-scoped access to pipeline options during graph construction and translation. This module provides thread-safe and async-safe access to globally-available instances of PipelineOptions using contextvars, scoped by the current pipeline. It allows components like transforms and coders to access the pipeline's configuration without requiring explicit parameter passing through every level of the call stack. For internal use only; no backwards-compatibility guarantees. """ from contextlib import contextmanager from contextvars import ContextVar from typing import TYPE_CHECKING from typing import Optional if TYPE_CHECKING: from apache_beam.options.pipeline_options import PipelineOptions # The contextvar holding the current pipeline's options. # Each thread and each asyncio task gets its own isolated copy. _pipeline_options: ContextVar[Optional['PipelineOptions']] = ContextVar( 'pipeline_options', default=None) def get_pipeline_options() -> Optional['PipelineOptions']: """Get the current pipeline's options from the context. Returns: The PipelineOptions for the currently executing pipeline operation, or None if called outside of a pipeline context. """ return _pipeline_options.get() @contextmanager def scoped_pipeline_options(options: Optional['PipelineOptions']): """Context manager that sets pipeline options for the duration of a block. Args: options: The PipelineOptions to make available during this scope. """ token = _pipeline_options.set(options) try: yield finally: _pipeline_options.reset(token) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/options/pipeline_options_context_test.py0000644000175100017510000002755015206113032026116 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for pipeline_options_context module. These tests verify that the contextvar-based approach properly isolates pipeline options across threads and async tasks, preventing race conditions. """ import asyncio import threading import unittest import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options_context import get_pipeline_options from apache_beam.options.pipeline_options_context import scoped_pipeline_options class PipelineConstructionOptionsTest(unittest.TestCase): def test_nested_scoping(self): """Test that nested scopes properly restore outer options.""" outer_options = PipelineOptions(['--runner=DirectRunner']) inner_options = PipelineOptions(['--runner=DataflowRunner']) with scoped_pipeline_options(outer_options): self.assertIs(get_pipeline_options(), outer_options) with scoped_pipeline_options(inner_options): self.assertIs(get_pipeline_options(), inner_options) self.assertIs(get_pipeline_options(), outer_options) self.assertIsNone(get_pipeline_options()) def test_exception_in_scope_restores_options(self): """Test that options are restored even when an exception is raised.""" outer_options = PipelineOptions(['--runner=DirectRunner']) inner_options = PipelineOptions(['--runner=DataflowRunner']) with scoped_pipeline_options(outer_options): try: with scoped_pipeline_options(inner_options): self.assertIs(get_pipeline_options(), inner_options) raise ValueError("Test exception") except ValueError: pass self.assertIs(get_pipeline_options(), outer_options) def test_different_threads_see_their_own_isolated_options(self): """Test that different threads see their own isolated options.""" results = {} errors = [] barrier = threading.Barrier(2) def thread_worker(thread_id, runner_name): try: options = PipelineOptions([f'--runner={runner_name}']) with scoped_pipeline_options(options): barrier.wait(timeout=5) current = get_pipeline_options() results[thread_id] = current.get_all_options()['runner'] import time time.sleep(0.01) current_after = get_pipeline_options() if current_after is not current: errors.append( f"Thread {thread_id}: options changed during execution") except Exception as e: errors.append(f"Thread {thread_id}: {e}") thread1 = threading.Thread(target=thread_worker, args=(1, 'DirectRunner')) thread2 = threading.Thread(target=thread_worker, args=(2, 'DataflowRunner')) thread1.start() thread2.start() thread1.join(timeout=5) thread2.join(timeout=5) self.assertEqual(errors, []) self.assertEqual(results[1], 'DirectRunner') self.assertEqual(results[2], 'DataflowRunner') def test_asyncio_task_isolation(self): """Test that different asyncio tasks see their own isolated options.""" async def async_worker( task_id, runner_name, results, ready_event, go_event): options = PipelineOptions([f'--runner={runner_name}']) with scoped_pipeline_options(options): ready_event.set() await go_event.wait() current = get_pipeline_options() results[task_id] = current.get_all_options()['runner'] await asyncio.sleep(0.01) current_after = get_pipeline_options() assert current_after is current, \ f"Task {task_id}: options changed during execution" async def run_test(): results = {} ready_events = [asyncio.Event() for _ in range(2)] go_event = asyncio.Event() task1 = asyncio.create_task( async_worker(1, 'DirectRunner', results, ready_events[0], go_event)) task2 = asyncio.create_task( async_worker(2, 'DataflowRunner', results, ready_events[1], go_event)) # Wait for both tasks to be ready await asyncio.gather(*[e.wait() for e in ready_events]) # Signal all tasks to proceed go_event.set() await asyncio.gather(task1, task2) return results results = asyncio.run(run_test()) self.assertEqual(results[1], 'DirectRunner') self.assertEqual(results[2], 'DataflowRunner') def test_transform_sees_pipeline_options(self): """Test that a transform can access pipeline options during expand().""" class OptionsCapturingTransform(beam.PTransform): """Transform that captures pipeline options during expand().""" def __init__(self, expected_job_name): self.expected_job_name = expected_job_name self.captured_options = None def expand(self, pcoll): # This runs during pipeline construction self.captured_options = get_pipeline_options() return pcoll | beam.Map(lambda x: x) options = PipelineOptions(['--job_name=test_job_123']) transform = OptionsCapturingTransform('test_job_123') with beam.Pipeline(options=options) as p: _ = p | beam.Create([1, 2, 3]) | transform # Verify the transform saw the correct options self.assertIsNotNone(transform.captured_options) self.assertEqual( transform.captured_options.get_all_options()['job_name'], 'test_job_123') def test_coder_sees_correct_options_during_run(self): """Test that coders see correct pipeline options during proto conversion. This tests the run path where as_deterministic_coder() is called during to_runner_api() proto conversion. """ from apache_beam.coders import coders from apache_beam.utils import shared errors = [] class WeakRefDict(dict): pass class TestKey: def __init__(self, value): self.value = value def __eq__(self, other): return isinstance(other, TestKey) and self.value == other.value def __hash__(self): return hash(self.value) class OptionsCapturingKeyCoder(coders.Coder): """Coder that captures pipeline options in as_deterministic_coder.""" shared_handle = shared.Shared() def encode(self, value): return str(value.value).encode('utf-8') def decode(self, encoded): return TestKey(encoded.decode('utf-8')) def is_deterministic(self): return False def as_deterministic_coder(self, step_label, error_message=None): opts = get_pipeline_options() if opts is not None: results = OptionsCapturingKeyCoder.shared_handle.acquire(WeakRefDict) job_name = opts.get_all_options().get('job_name') results['Worker1'] = job_name return self beam.coders.registry.register_coder(TestKey, OptionsCapturingKeyCoder) results = OptionsCapturingKeyCoder.shared_handle.acquire(WeakRefDict) job_name = 'gbk_job' options = PipelineOptions([f'--job_name={job_name}']) with beam.Pipeline(options=options) as p: _ = ( p | beam.Create([(TestKey(1), 'a'), (TestKey(2), 'b')]) | beam.GroupByKey()) self.assertEqual(errors, [], f"Errors occurred: {errors}") self.assertEqual( results.get('Worker1'), job_name, f"Worker1 saw wrong options: {results}") self.assertFalse(get_pipeline_options() == options) def test_barrier_inside_default_type_hints(self): """Test race condition detection with barrier inside default_type_hints. This test reliably detects race conditions because: 1. Both threads start pipeline construction simultaneously 2. Inside default_type_hints (which is called during Pipeline.apply()), both threads hit a barrier and wait for each other 3. At this point, BOTH threads are inside scoped_pipeline_options 4. When they continue, they read options - with a global var, they'd see the wrong values because the last thread to set options would win """ results = {} errors = [] inner_barrier = threading.Barrier(2) class BarrierTransform(beam.PTransform): """Transform that synchronizes threads INSIDE default_type_hints.""" def __init__(self, worker_id, results_dict, barrier): self.worker_id = worker_id self.results_dict = results_dict self.barrier = barrier def expand(self, pcoll): return pcoll | beam.Map(lambda x: x) def default_type_hints(self): self.barrier.wait(timeout=5) opts = get_pipeline_options() if opts is not None: job_name = opts.get_all_options().get('job_name') if self.worker_id not in self.results_dict: self.results_dict[self.worker_id] = job_name return super().default_type_hints() def construct_pipeline(worker_id): try: job_name = f'barrier_job_{worker_id}' options = PipelineOptions([f'--job_name={job_name}']) transform = BarrierTransform(worker_id, results, inner_barrier) with beam.Pipeline(options=options) as p: _ = p | beam.Create([1, 2, 3]) | transform except Exception as e: import traceback errors.append(f"Worker {worker_id}: {e}\n{traceback.format_exc()}") thread1 = threading.Thread( target=construct_pipeline, args=(1, )) thread2 = threading.Thread( target=construct_pipeline, args=(2, )) thread1.start() thread2.start() thread1.join(timeout=10) thread2.join(timeout=10) self.assertEqual(errors, [], f"Errors occurred: {errors}") self.assertEqual( results.get(1), 'barrier_job_1', f"Worker 1 saw wrong options: {results}") self.assertEqual( results.get(2), 'barrier_job_2', f"Worker 2 saw wrong options: {results}") class PipelineSubclassApplyTest(unittest.TestCase): def test_subclass_apply_called_on_recursive_paths(self): """Test that Pipeline subclass overrides of apply() are respected. _apply_internal's recursive calls must go through self.apply(), not self._apply_internal(), so that subclass interceptions are not skipped. """ apply_calls = [] class TrackingPipeline(beam.Pipeline): def apply(self, transform, pvalueish=None, label=None): apply_calls.append(label or transform.label) return super().apply(transform, pvalueish, label) options = PipelineOptions(['--job_name=subclass_test']) with TrackingPipeline(options=options) as p: # "my_label" >> transform creates a _NamedPTransform, which triggers # two recursive apply() calls: one to unwrap _NamedPTransform, and # one to handle the label argument. _ = p | beam.Create([1, 2, 3]) | "my_label" >> beam.Map(lambda x: x) # beam.Create goes through apply() once (no recursion). # "my_label" >> Map triggers: apply(_NamedPTransform) -> apply(Map, # label="my_label") -> apply(Map). That's 3 calls through apply(). # Total: 1 (Create) + 3 (Map) = 4 calls minimum. map_calls = [c for c in apply_calls if c == 'my_label' or c == 'Map'] self.assertGreaterEqual( len(map_calls), 3, f"Expected at least 3 apply() calls for the Map transform " f"(NamedPTransform unwrap + label handling + final), " f"got {len(map_calls)}. All calls: {apply_calls}") if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/options/pipeline_options_test.py0000644000175100017510000012171015206113032024343 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the pipeline options module.""" # pytype: skip-file import json import logging import os import unittest import hamcrest as hc import mock from parameterized import parameterized from apache_beam.options.pipeline_options import CrossLanguageOptions from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import JobServerOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import ProfilingOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import TypeOptions from apache_beam.options.pipeline_options import WorkerOptions from apache_beam.options.pipeline_options import _BeamArgumentParser from apache_beam.options.pipeline_options_validator import PipelineOptionsValidator from apache_beam.options.value_provider import RuntimeValueProvider from apache_beam.options.value_provider import StaticValueProvider from apache_beam.transforms.display import DisplayData from apache_beam.transforms.display_test import DisplayDataItemMatcher _LOGGER = logging.getLogger(__name__) try: import apache_beam.io.gcp.gcsio # pylint: disable=unused-import has_gcsio = True except ImportError: has_gcsio = False # Mock runners to use for validations. class MockRunners(object): class DataflowRunner(object): def get_default_gcp_region(self): # Return a default so we don't have to specify --region in every test # (unless specifically testing it). return 'us-central1' class MockGoogleCloudOptionsNoBucket(GoogleCloudOptions): def _create_default_gcs_bucket(self): return None class MockGoogleCloudOptionsWithBucket(GoogleCloudOptions): def _create_default_gcs_bucket(self): return "gs://default/bucket" class PipelineOptionsTest(unittest.TestCase): def setUp(self): # Reset runtime options to avoid side-effects caused by other tests. # Note that is_accessible assertions require runtime_options to # be uninitialized. RuntimeValueProvider.set_runtime_options(None) def tearDown(self): # Reset runtime options to avoid side-effects in other tests. RuntimeValueProvider.set_runtime_options(None) TEST_CASES = [ (['--num_workers', '5'], { 'num_workers': 5, 'mock_flag': False, 'mock_option': None, 'mock_multi_option': None }, [DisplayDataItemMatcher('num_workers', 5)]), (['--direct_num_workers', '5'], { 'direct_num_workers': 5, 'mock_flag': False, 'mock_option': None, 'mock_multi_option': None }, [DisplayDataItemMatcher('direct_num_workers', 5)]), (['--direct_running_mode', 'multi_threading'], { 'direct_running_mode': 'multi_threading', 'mock_flag': False, 'mock_option': None, 'mock_multi_option': None }, [DisplayDataItemMatcher('direct_running_mode', 'multi_threading')]), (['--direct_running_mode', 'multi_processing'], { 'direct_running_mode': 'multi_processing', 'mock_flag': False, 'mock_option': None, 'mock_multi_option': None }, [DisplayDataItemMatcher('direct_running_mode', 'multi_processing')]), (['--profile_cpu', '--profile_location', 'gs://bucket/', 'ignored'], { 'profile_cpu': True, 'profile_location': 'gs://bucket/', 'mock_flag': False, 'mock_option': None, 'mock_multi_option': None }, [ DisplayDataItemMatcher('profile_cpu', True), DisplayDataItemMatcher('profile_location', 'gs://bucket/') ]), (['--num_workers', '5', '--mock_flag'], { 'num_workers': 5, 'mock_flag': True, 'mock_option': None, 'mock_multi_option': None }, [ DisplayDataItemMatcher('num_workers', 5), DisplayDataItemMatcher('mock_flag', True) ]), (['--mock_option', 'abc'], { 'mock_flag': False, 'mock_option': 'abc', 'mock_multi_option': None }, [DisplayDataItemMatcher('mock_option', 'abc')]), (['--mock_option', ' abc def '], { 'mock_flag': False, 'mock_option': ' abc def ', 'mock_multi_option': None }, [DisplayDataItemMatcher('mock_option', ' abc def ')]), (['--mock_option= abc xyz '], { 'mock_flag': False, 'mock_option': ' abc xyz ', 'mock_multi_option': None }, [DisplayDataItemMatcher('mock_option', ' abc xyz ')]), ([ '--mock_option=gs://my bucket/my folder/my file', '--mock_multi_option=op1', '--mock_multi_option=op2' ], { 'mock_flag': False, 'mock_option': 'gs://my bucket/my folder/my file', 'mock_multi_option': ['op1', 'op2'] }, [ DisplayDataItemMatcher( 'mock_option', 'gs://my bucket/my folder/my file'), DisplayDataItemMatcher('mock_multi_option', ['op1', 'op2']) ]), (['--mock_multi_option=op1', '--mock_multi_option=op2'], { 'mock_flag': False, 'mock_option': None, 'mock_multi_option': ['op1', 'op2'] }, [DisplayDataItemMatcher('mock_multi_option', ['op1', 'op2'])]), (['--mock_json_option={"11a": 0, "37a": 1}'], { 'mock_flag': False, 'mock_option': None, 'mock_multi_option': None, 'mock_json_option': { '11a': 0, '37a': 1 }, }, [DisplayDataItemMatcher('mock_json_option', { '11a': 0, '37a': 1 })]), ] # Used for testing newly added flags. class MockOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--mock_flag', action='store_true', help='mock flag') parser.add_argument('--mock_option', help='mock option') parser.add_argument( '--mock_multi_option', action='append', help='mock multi option') parser.add_argument('--option with space', help='mock option with space') parser.add_argument('--mock_json_option', type=json.loads, default={}) # Use with MockOptions in test cases where multiple option classes are needed. class FakeOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--fake_flag', action='store_true', help='fake flag') parser.add_argument('--fake_option', help='fake option') parser.add_argument( '--fake_multi_option', action='append', help='fake multi option') @parameterized.expand(TEST_CASES) def test_display_data(self, flags, _, display_data): options = PipelineOptions(flags=flags) dd = DisplayData.create_from(options) hc.assert_that(dd.items, hc.contains_inanyorder(*display_data)) @parameterized.expand(TEST_CASES) def test_get_all_options_subclass(self, flags, expected, _): options = PipelineOptionsTest.MockOptions(flags=flags) self.assertLessEqual(expected.items(), options.get_all_options().items()) self.assertEqual( options.view_as(PipelineOptionsTest.MockOptions).mock_flag, expected['mock_flag']) self.assertEqual( options.view_as(PipelineOptionsTest.MockOptions).mock_option, expected['mock_option']) self.assertEqual( options.view_as(PipelineOptionsTest.MockOptions).mock_multi_option, expected['mock_multi_option']) @parameterized.expand(TEST_CASES) def test_get_all_options(self, flags, expected, _): options = PipelineOptions(flags=flags) self.assertLessEqual(expected.items(), options.get_all_options().items()) self.assertEqual( options.view_as(PipelineOptionsTest.MockOptions).mock_flag, expected['mock_flag']) self.assertEqual( options.view_as(PipelineOptionsTest.MockOptions).mock_option, expected['mock_option']) self.assertEqual( options.view_as(PipelineOptionsTest.MockOptions).mock_multi_option, expected['mock_multi_option']) def test_get_superclass_options(self): flags = ["--mock_option", "mock", "--fake_option", "fake"] options = PipelineOptions(flags=flags).view_as( PipelineOptionsTest.FakeOptions) items = options.get_all_options(current_only=True).items() print(items) self.assertTrue(('fake_option', 'fake') in items) self.assertFalse(('mock_option', 'mock') in items) items = options.view_as(PipelineOptionsTest.MockOptions).get_all_options( current_only=True).items() self.assertFalse(('fake_option', 'fake') in items) self.assertTrue(('mock_option', 'mock') in items) @parameterized.expand(TEST_CASES) def test_subclasses_of_pipeline_options_can_be_instantiated( self, flags, expected, _): mock_options = PipelineOptionsTest.MockOptions(flags=flags) self.assertEqual(mock_options.mock_flag, expected['mock_flag']) self.assertEqual(mock_options.mock_option, expected['mock_option']) self.assertEqual( mock_options.mock_multi_option, expected['mock_multi_option']) @parameterized.expand(TEST_CASES) def test_views_can_be_constructed_from_pipeline_option_subclasses( self, flags, expected, _): fake_options = PipelineOptionsTest.FakeOptions(flags=flags) mock_options = fake_options.view_as(PipelineOptionsTest.MockOptions) self.assertEqual(mock_options.mock_flag, expected['mock_flag']) self.assertEqual(mock_options.mock_option, expected['mock_option']) self.assertEqual( mock_options.mock_multi_option, expected['mock_multi_option']) def test_views_do_not_expose_options_defined_by_other_views(self): flags = ['--mock_option=mock_value', '--fake_option=fake_value'] options = PipelineOptions(flags) assert options.view_as( PipelineOptionsTest.MockOptions).mock_option == 'mock_value' assert options.view_as( PipelineOptionsTest.FakeOptions).fake_option == 'fake_value' assert options.view_as(PipelineOptionsTest.MockOptions).view_as( PipelineOptionsTest.FakeOptions).fake_option == 'fake_value' self.assertRaises( AttributeError, lambda: options.view_as(PipelineOptionsTest.MockOptions).fake_option) self.assertRaises( AttributeError, lambda: options.view_as(PipelineOptionsTest.MockOptions).view_as( PipelineOptionsTest.FakeOptions).view_as( PipelineOptionsTest.MockOptions).fake_option) @parameterized.expand(TEST_CASES) def test_from_dictionary(self, flags, expected, _): options = PipelineOptions(flags=flags) all_options_dict = options.get_all_options() options_from_dict = PipelineOptions.from_dictionary(all_options_dict) self.assertEqual( options_from_dict.view_as(PipelineOptionsTest.MockOptions).mock_flag, expected['mock_flag']) self.assertEqual( options.view_as(PipelineOptionsTest.MockOptions).mock_option, expected['mock_option']) self.assertEqual( options.view_as(PipelineOptionsTest.MockOptions).mock_multi_option, expected['mock_multi_option']) self.assertEqual( options.view_as(PipelineOptionsTest.MockOptions).mock_json_option, expected.get('mock_json_option', {})) def test_none_from_dictionary(self): class NoneDefaultOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--test_arg_none', default=None, type=int) parser.add_argument('--test_arg_int', default=1, type=int) options_dict = {'test_arg_none': None, 'test_arg_int': 5} options_from_dict = NoneDefaultOptions.from_dictionary(options_dict) result = options_from_dict.get_all_options() self.assertEqual(result['test_arg_int'], 5) self.assertEqual(result['test_arg_none'], None) def test_merging_options(self): opts = PipelineOptions(flags=['--num_workers', '5']) actual_opts = PipelineOptions.from_runner_api(opts.to_runner_api()) actual = actual_opts.view_as(WorkerOptions).num_workers self.assertEqual(5, actual) def test_merging_options_with_overriden_options(self): opts = PipelineOptions(flags=['--num_workers', '5']) base = PipelineOptions(flags=['--num_workers', '2']) actual_opts = PipelineOptions.from_runner_api(opts.to_runner_api(), base) actual = actual_opts.view_as(WorkerOptions).num_workers self.assertEqual(5, actual) def test_merging_options_with_overriden_runner(self): opts = PipelineOptions(flags=['--runner', 'FnApiRunner']) base = PipelineOptions(flags=['--runner', 'Direct']) actual_opts = PipelineOptions.from_runner_api(opts.to_runner_api(), base) actual = actual_opts.view_as(StandardOptions).runner self.assertEqual('Direct', actual) def test_from_kwargs(self): class MyOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--test_arg') # kwarg takes precedence over parsed flag options = PipelineOptions(flags=['--test_arg=A'], test_arg='B') self.assertEqual(options.view_as(MyOptions).test_arg, 'B') self.assertEqual(options.get_all_options()['test_arg'], 'B') def test_option_with_space(self): options = PipelineOptions(flags=['--option with space= value with space']) self.assertEqual( getattr( options.view_as(PipelineOptionsTest.MockOptions), 'option with space'), ' value with space') options_from_dict = PipelineOptions.from_dictionary( options.get_all_options()) self.assertEqual( getattr( options_from_dict.view_as(PipelineOptionsTest.MockOptions), 'option with space'), ' value with space') def test_retain_unknown_options_binary_store_string(self): options = PipelineOptions(['--unknown_option', 'some_value']) result = options.get_all_options(retain_unknown_options=True) self.assertEqual(result['unknown_option'], 'some_value') def test_retain_unknown_options_binary_equals_store_string(self): options = PipelineOptions(['--unknown_option=some_value']) result = options.get_all_options(retain_unknown_options=True) self.assertEqual(result['unknown_option'], 'some_value') def test_retain_unknown_options_binary_multi_equals_store_string(self): options = PipelineOptions(['--unknown_option=expr = "2 + 2 = 5"']) result = options.get_all_options(retain_unknown_options=True) self.assertEqual(result['unknown_option'], 'expr = "2 + 2 = 5"') def test_retain_unknown_options_binary_single_dash_store_string(self): options = PipelineOptions(['-i', 'some_value']) with self.assertRaises(KeyError): _ = options.get_all_options(retain_unknown_options=True)['i'] def test_retain_unknown_options_unary_store_true(self): options = PipelineOptions(['--unknown_option']) result = options.get_all_options(retain_unknown_options=True) self.assertEqual(result['unknown_option'], True) def test_retain_unknown_options_consecutive_unary_store_true(self): options = PipelineOptions(['--option_foo', '--option_bar']) result = options.get_all_options(retain_unknown_options=True) self.assertEqual(result['option_foo'], True) self.assertEqual(result['option_bar'], True) def test_retain_unknown_options_unary_single_dash_store_true(self): options = PipelineOptions(['-i']) result = options.get_all_options(retain_unknown_options=True) self.assertEqual(result['i'], True) def test_override_options(self): base_flags = ['--num_workers', '5'] options = PipelineOptions(base_flags) self.assertEqual(options.get_all_options()['num_workers'], 5) self.assertEqual(options.get_all_options()['mock_flag'], False) options.view_as(PipelineOptionsTest.MockOptions).mock_flag = True self.assertEqual(options.get_all_options()['num_workers'], 5) self.assertTrue(options.get_all_options()['mock_flag']) def test_override_init_options(self): base_flags = ['--num_workers', '5'] options = PipelineOptions(base_flags, mock_flag=True) self.assertEqual(options.get_all_options()['num_workers'], 5) self.assertEqual(options.get_all_options()['mock_flag'], True) def test_invalid_override_init_options(self): base_flags = ['--num_workers', '5'] options = PipelineOptions(base_flags, mock_invalid_flag=True) self.assertEqual(options.get_all_options()['num_workers'], 5) self.assertEqual(options.get_all_options()['mock_flag'], False) def test_experiments(self): options = PipelineOptions(['--experiment', 'abc', '--experiment', 'def']) self.assertEqual( sorted(options.get_all_options()['experiments']), ['abc', 'def']) options = PipelineOptions(['--experiments', 'abc', '--experiments', 'def']) self.assertEqual( sorted(options.get_all_options()['experiments']), ['abc', 'def']) options = PipelineOptions(flags=['']) self.assertEqual(options.get_all_options()['experiments'], None) def test_worker_options(self): options = PipelineOptions([ '--machine_type', 'abc', '--disk_type', 'def', '--disk_provisioned_iops', '4000', '--disk_provisioned_throughput_mibps', '200', '--element_processing_timeout_minutes', '10', ]) worker_options = options.view_as(WorkerOptions) self.assertEqual(worker_options.machine_type, 'abc') self.assertEqual(worker_options.disk_type, 'def') self.assertEqual(worker_options.disk_provisioned_iops, 4000) self.assertEqual(worker_options.disk_provisioned_throughput_mibps, 200) self.assertEqual(worker_options.element_processing_timeout_minutes, 10) options = PipelineOptions( ['--worker_machine_type', 'abc', '--worker_disk_type', 'def']) worker_options = options.view_as(WorkerOptions) self.assertEqual(worker_options.machine_type, 'abc') self.assertEqual(worker_options.disk_type, 'def') def test_beam_services_empty(self): with mock.patch.dict(os.environ, {}, clear=True): options = PipelineOptions().view_as(CrossLanguageOptions) self.assertEqual(options.beam_services, {}) def test_beam_services_from_env(self): with mock.patch.dict(os.environ, {'BEAM_SERVICE_OVERRIDES': '{"foo": "bar"}'}, clear=True): options = PipelineOptions().view_as(CrossLanguageOptions) self.assertEqual(options.beam_services, {'foo': 'bar'}) def test_beam_services_from_flag(self): with mock.patch.dict(os.environ, {}, clear=True): options = PipelineOptions(['--beam_services={"foo": "bar"}' ]).view_as(CrossLanguageOptions) self.assertEqual(options.beam_services, {'foo': 'bar'}) def test_beam_services_from_env_and_flag(self): with mock.patch.dict( os.environ, {'BEAM_SERVICE_OVERRIDES': '{"foo": "bar", "other": "zzz"}'}, clear=True): options = PipelineOptions(['--beam_services={"foo": "override"}' ]).view_as(CrossLanguageOptions) self.assertEqual( options.beam_services, { 'foo': 'override', 'other': 'zzz' }) def test_option_modifications_are_shared_between_views(self): pipeline_options = PipelineOptions([ '--mock_option', 'value', '--mock_flag', '--mock_multi_option', 'value1', '--mock_multi_option', 'value2', ]) mock_options = PipelineOptionsTest.MockOptions([ '--mock_option', 'value', '--mock_flag', '--mock_multi_option', 'value1', '--mock_multi_option', 'value2', ]) for options in [pipeline_options, mock_options]: view1 = options.view_as(PipelineOptionsTest.MockOptions) view2 = options.view_as(PipelineOptionsTest.MockOptions) view1.mock_option = 'new_value' view1.mock_flag = False view1.mock_multi_option.append('value3') view3 = options.view_as(PipelineOptionsTest.MockOptions) view4 = view1.view_as(PipelineOptionsTest.MockOptions) view5 = options.view_as(TypeOptions).view_as( PipelineOptionsTest.MockOptions) for view in [view1, view2, view3, view4, view5]: self.assertEqual('new_value', view.mock_option) self.assertFalse(view.mock_flag) self.assertEqual(['value1', 'value2', 'value3'], view.mock_multi_option) def test_uninitialized_option_modifications_are_shared_between_views(self): options = PipelineOptions([]) view1 = options.view_as(PipelineOptionsTest.MockOptions) view2 = options.view_as(PipelineOptionsTest.MockOptions) view1.mock_option = 'some_value' view1.mock_flag = False view1.mock_multi_option = ['value1', 'value2'] view3 = options.view_as(PipelineOptionsTest.MockOptions) view4 = view1.view_as(PipelineOptionsTest.MockOptions) view5 = options.view_as(TypeOptions).view_as( PipelineOptionsTest.MockOptions) for view in [view1, view2, view3, view4, view5]: self.assertEqual('some_value', view.mock_option) self.assertFalse(view.mock_flag) self.assertEqual(['value1', 'value2'], view.mock_multi_option) def test_extra_package(self): options = PipelineOptions([ '--extra_package', 'abc', '--extra_packages', 'def', '--extra_packages', 'ghi' ]) self.assertEqual( sorted(options.get_all_options()['extra_packages']), ['abc', 'def', 'ghi']) options = PipelineOptions(flags=['']) self.assertEqual(options.get_all_options()['extra_packages'], None) def test_dataflow_job_file(self): options = PipelineOptions(['--dataflow_job_file', 'abc']) self.assertEqual(options.get_all_options()['dataflow_job_file'], 'abc') options = PipelineOptions(flags=['']) self.assertEqual(options.get_all_options()['dataflow_job_file'], None) def test_template_location(self): options = PipelineOptions(['--template_location', 'abc']) self.assertEqual(options.get_all_options()['template_location'], 'abc') options = PipelineOptions(flags=['']) self.assertEqual(options.get_all_options()['template_location'], None) def test_redefine_options(self): class TestRedefinedOptions(PipelineOptions): # pylint: disable=unused-variable @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--redefined_flag', action='store_true') class TestRedefinedOptions(PipelineOptions): # pylint: disable=function-redefined @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--redefined_flag', action='store_true') options = PipelineOptions(['--redefined_flag']) self.assertTrue(options.get_all_options()['redefined_flag']) # TODO(https://github.com/apache/beam/issues/18197): Require unique names # only within a test. For now, _vp_arg will be # the convention to name value-provider arguments in tests, as opposed to # _non_vp_arg for non-value-provider arguments. # The number will grow per file as tests are added. def test_value_provider_options(self): class UserOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument( '--pot_vp_arg1', help='This flag is a value provider') parser.add_value_provider_argument('--pot_vp_arg2', default=1, type=int) parser.add_argument('--pot_non_vp_arg1', default=1, type=int) # Provide values: if not provided, the option becomes of the type runtime vp options = UserOptions(['--pot_vp_arg1', 'hello']) self.assertIsInstance(options.pot_vp_arg1, StaticValueProvider) self.assertIsInstance(options.pot_vp_arg2, RuntimeValueProvider) self.assertIsInstance(options.pot_non_vp_arg1, int) # Values can be overwritten options = UserOptions( pot_vp_arg1=5, pot_vp_arg2=StaticValueProvider(value_type=str, value='bye'), pot_non_vp_arg1=RuntimeValueProvider( option_name='foo', value_type=int, default_value=10)) self.assertEqual(options.pot_vp_arg1, 5) self.assertTrue( options.pot_vp_arg2.is_accessible(), '%s is not accessible' % options.pot_vp_arg2) self.assertEqual(options.pot_vp_arg2.get(), 'bye') self.assertFalse(options.pot_non_vp_arg1.is_accessible()) with self.assertRaises(RuntimeError): options.pot_non_vp_arg1.get() # Converts extra arguments to list value. def test_extra_args(self): options = PipelineOptions([ '--extra_arg', 'val1', '--extra_arg', 'val2', '--extra_arg=val3', '--unknown_arg', 'val4' ]) def add_extra_options(parser): parser.add_argument("--extra_arg", action='append') self.assertEqual( options.get_all_options( add_extra_args_fn=add_extra_options)['extra_arg'], ['val1', 'val2', 'val3']) # The argparse package by default tries to autocomplete option names. This # results in an "ambiguous option" error from argparse when an unknown option # matching multiple known ones are used. This tests that we suppress this # error. def test_unknown_option_prefix(self): # Test that the "ambiguous option" error is suppressed. options = PipelineOptions(['--profi', 'val1']) options.view_as(ProfilingOptions) # Test that valid errors are not suppressed. with self.assertRaises(SystemExit): # Invalid option choice. options = PipelineOptions(['--type_check_strictness', 'blahblah']) options.view_as(TypeOptions) def test_add_experiment(self): options = PipelineOptions([]) options.view_as(DebugOptions).add_experiment('new_experiment') self.assertEqual(['new_experiment'], options.view_as(DebugOptions).experiments) def test_add_experiment_preserves_existing_experiments(self): options = PipelineOptions(['--experiment=existing_experiment']) options.view_as(DebugOptions).add_experiment('new_experiment') self.assertEqual(['existing_experiment', 'new_experiment'], options.view_as(DebugOptions).experiments) def test_lookup_experiments(self): options = PipelineOptions([ '--experiment=existing_experiment', '--experiment', 'key=value', '--experiment', 'master_key=k1=v1,k2=v2', ]) debug_options = options.view_as(DebugOptions) self.assertEqual( 'default_value', debug_options.lookup_experiment('nonexistent', 'default_value')) self.assertEqual( 'value', debug_options.lookup_experiment('key', 'default_value')) self.assertEqual( 'k1=v1,k2=v2', debug_options.lookup_experiment('master_key')) self.assertEqual( True, debug_options.lookup_experiment('existing_experiment')) def test_transform_name_mapping(self): options = PipelineOptions(['--transform_name_mapping={\"from\":\"to\"}']) mapping = options.view_as(GoogleCloudOptions).transform_name_mapping self.assertEqual(mapping['from'], 'to') def test_jar_cache_dir(self): options = PipelineOptions(['--jar_cache_dir=/path/to/jar_cache_dir']) jar_cache_dir = options.view_as(JobServerOptions).jar_cache_dir self.assertEqual(jar_cache_dir, '/path/to/jar_cache_dir') def test_dataflow_service_options(self): options = PipelineOptions([ '--dataflow_service_option', 'whizz=bang', '--dataflow_service_option', 'beep=boop' ]) self.assertEqual( sorted(options.get_all_options()['dataflow_service_options']), ['beep=boop', 'whizz=bang']) options = PipelineOptions([ '--dataflow_service_options', 'whizz=bang', '--dataflow_service_options', 'beep=boop' ]) self.assertEqual( sorted(options.get_all_options()['dataflow_service_options']), ['beep=boop', 'whizz=bang']) options = PipelineOptions(flags=['']) self.assertEqual( options.get_all_options()['dataflow_service_options'], None) def test_options_store_false_with_different_dest(self): parser = _BeamArgumentParser() for cls in PipelineOptions.__subclasses__(): cls._add_argparse_args(parser) actions = parser._actions.copy() options_to_flags = {} options_diff_dest_store_true = {} for i in range(len(actions)): flag_names = actions[i].option_strings option_name = actions[i].dest if isinstance(actions[i].const, bool): for flag_name in flag_names: flag_name = flag_name.strip('-') if flag_name != option_name: # Capture flags which has store_action=True and has a # different dest. This behavior would be confusing. if actions[i].const: options_diff_dest_store_true[flag_name] = option_name continue # check the flags like no_use_public_ips # default is None, action is {True, False} if actions[i].default is None: options_to_flags[option_name] = flag_name self.assertEqual( len(options_diff_dest_store_true), 0, _LOGGER.error( "There should be no flags that have a dest " "different from flag name and action as " "store_true. It would be confusing " "to the user. Please specify the dest as the " "flag_name instead.")) from apache_beam.options.pipeline_options import _FLAG_THAT_SETS_FALSE_VALUE self.assertDictEqual( _FLAG_THAT_SETS_FALSE_VALUE, options_to_flags, "If you are adding a new boolean flag with default=None," " with different dest/option_name from the flag name, please add " "the dest and the flag name to the map " "_FLAG_THAT_SETS_FALSE_VALUE in PipelineOptions.py") def test_gcs_custom_audit_entries(self): options = PipelineOptions([ '--gcs_custom_audit_entry=user=test-user', '--gcs_custom_audit_entry=work=test-work', '--gcs_custom_audit_entries={"job":"test-job", "id":"1234"}' ]) entries = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries self.assertDictEqual( entries, { 'x-goog-custom-audit-user': 'test-user', 'x-goog-custom-audit-work': 'test-work', 'x-goog-custom-audit-job': 'test-job', 'x-goog-custom-audit-id': '1234' }) def test_gcs_custom_audit_entries_wo_duplicated_prefix(self): options = PipelineOptions([ '--gcs_custom_audit_entry=x-goog-custom-audit-user=test-user', '--gcs_custom_audit_entries={"job":"test-job", "id":"1234"}' ]) entries = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries self.assertDictEqual( entries, { 'x-goog-custom-audit-user': 'test-user', 'x-goog-custom-audit-job': 'test-job', 'x-goog-custom-audit-id': '1234' }) @mock.patch('apache_beam.options.pipeline_options._BeamArgumentParser.error') def test_gcs_custom_audit_entries_with_errors(self, mock_error): long_key = 'a' * 65 options = PipelineOptions([f'--gcs_custom_audit_entry={long_key}=1']) _ = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries self.assertRegex( mock_error.call_args[0][0], 'The key .* exceeds the 64-character limit.') mock_error.reset_mock() long_value = 'b' * 1201 options = PipelineOptions([f'--gcs_custom_audit_entry=key={long_value}']) _ = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries self.assertRegex( mock_error.call_args[0][0], 'The value .* exceeds the 1200-character limit.') mock_error.reset_mock() options = PipelineOptions([ '--gcs_custom_audit_entry=a=1', '--gcs_custom_audit_entry=b=2', '--gcs_custom_audit_entry=c=3', '--gcs_custom_audit_entry=d=4', '--gcs_custom_audit_entry=job=test-job' ]) _ = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries self.assertRegex( mock_error.call_args[0][0], 'The maximum allowed number of GCS custom audit entries .*') mock_error.reset_mock() options = PipelineOptions([ '--gcs_custom_audit_entry=a=1', '--gcs_custom_audit_entry=b=2', '--gcs_custom_audit_entry=c=3', '--gcs_custom_audit_entry=d=4' ]) _ = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries self.assertRegex( mock_error.call_args[0][0], 'The maximum allowed number of GCS custom audit entries .*') def _check_errors(self, options, validator, expected): if has_gcsio: with mock.patch('apache_beam.io.gcp.gcsio.GcsIO.is_soft_delete_enabled', return_value=False): errors = options._handle_temp_and_staging_locations(validator) self.assertEqual(errors, expected) else: errors = options._handle_temp_and_staging_locations(validator) self.assertEqual(errors, expected) def test_validation_good_stg_good_temp(self): runner = MockRunners.DataflowRunner() options = GoogleCloudOptions([ '--project=myproject', '--staging_location=gs://beam/stg', '--temp_location=gs://beam/tmp' ]) validator = PipelineOptionsValidator(options, runner) self._check_errors(options, validator, []) self.assertEqual( options.get_all_options()['staging_location'], "gs://beam/stg") self.assertEqual( options.get_all_options()['temp_location'], "gs://beam/tmp") def test_validation_bad_stg_good_temp(self): runner = MockRunners.DataflowRunner() options = GoogleCloudOptions([ '--project=myproject', '--staging_location=badGSpath', '--temp_location=gs://beam/tmp' ]) validator = PipelineOptionsValidator(options, runner) self._check_errors(options, validator, []) self.assertEqual( options.get_all_options()['staging_location'], "gs://beam/tmp") self.assertEqual( options.get_all_options()['temp_location'], "gs://beam/tmp") def test_validation_good_stg_bad_temp(self): runner = MockRunners.DataflowRunner() options = GoogleCloudOptions([ '--project=myproject', '--staging_location=gs://beam/stg', '--temp_location=badGSpath' ]) validator = PipelineOptionsValidator(options, runner) self._check_errors(options, validator, []) self.assertEqual( options.get_all_options()['staging_location'], "gs://beam/stg") self.assertEqual( options.get_all_options()['temp_location'], "gs://beam/stg") def test_validation_bad_stg_bad_temp_with_default(self): runner = MockRunners.DataflowRunner() options = MockGoogleCloudOptionsWithBucket([ '--project=myproject', '--staging_location=badGSpath', '--temp_location=badGSpath' ]) validator = PipelineOptionsValidator(options, runner) self._check_errors(options, validator, []) self.assertEqual( options.get_all_options()['staging_location'], "gs://default/bucket") self.assertEqual( options.get_all_options()['temp_location'], "gs://default/bucket") def test_validation_bad_stg_bad_temp_no_default(self): runner = MockRunners.DataflowRunner() options = MockGoogleCloudOptionsNoBucket([ '--project=myproject', '--staging_location=badGSpath', '--temp_location=badGSpath' ]) validator = PipelineOptionsValidator(options, runner) self._check_errors( options, validator, [ 'Invalid GCS path (badGSpath), given for the option: ' \ 'temp_location.', 'Invalid GCS path (badGSpath), given for the option: ' \ 'staging_location.' ]) def test_comma_separated_experiments(self): """Test that comma-separated experiments are parsed correctly.""" # Test single experiment options = PipelineOptions(['--experiments=abc']) self.assertEqual(['abc'], options.get_all_options()['experiments']) # Test comma-separated experiments options = PipelineOptions(['--experiments=abc,def,ghi']) self.assertEqual(['abc', 'def', 'ghi'], options.get_all_options()['experiments']) # Test multiple flags with comma-separated values options = PipelineOptions( ['--experiments=abc,def', '--experiments=ghi,jkl']) self.assertEqual(['abc', 'def', 'ghi', 'jkl'], options.get_all_options()['experiments']) # Test with spaces around commas options = PipelineOptions(['--experiments=abc, def , ghi']) self.assertEqual(['abc', 'def', 'ghi'], options.get_all_options()['experiments']) # Test empty values are filtered out options = PipelineOptions(['--experiments=abc,,def,']) self.assertEqual(['abc', 'def'], options.get_all_options()['experiments']) def test_comma_separated_dataflow_service_options(self): """Test that comma-separated dataflow service options are parsed correctly.""" # Test single option options = PipelineOptions(['--dataflow_service_options=option1=value1']) self.assertEqual(['option1=value1'], options.get_all_options()['dataflow_service_options']) # Test comma-separated options options = PipelineOptions([ '--dataflow_service_options=option1=value1,option2=value2,' 'option3=value3' ]) self.assertEqual(['option1=value1', 'option2=value2', 'option3=value3'], options.get_all_options()['dataflow_service_options']) # Test multiple flags with comma-separated values options = PipelineOptions([ '--dataflow_service_options=option1=value1,option2=value2', '--dataflow_service_options=option3=value3,option4=value4' ]) self.assertEqual([ 'option1=value1', 'option2=value2', 'option3=value3', 'option4=value4' ], options.get_all_options()['dataflow_service_options']) class CompatVersionTest(unittest.TestCase): def test_is_compat_version_prior_to(self): test_cases = [ # Basic comparison cases ("1.0.0", "2.0.0", True), # v1 < v2 in major ("2.0.0", "1.0.0", False), # v1 > v2 in major ("1.1.0", "1.2.0", True), # v1 < v2 in minor ("1.2.0", "1.1.0", False), # v1 > v2 in minor ("1.0.1", "1.0.2", True), # v1 < v2 in patch ("1.0.2", "1.0.1", False), # v1 > v2 in patch # Equal versions ("1.0.0", "1.0.0", False), # Identical ("0.0.0", "0.0.0", False), # Both zero # Different lengths - shorter vs longer ("1.0", "1.0.0", False), # Should be equal (1.0 = 1.0.0) ("1.0", "1.0.1", True), # 1.0.0 < 1.0.1 ("1.2", "1.2.0", False), # Should be equal (1.2 = 1.2.0) ("1.2", "1.2.3", True), # 1.2.0 < 1.2.3 ("2", "2.0.0", False), # Should be equal (2 = 2.0.0) ("2", "2.0.1", True), # 2.0.0 < 2.0.1 ("1", "2.0", True), # 1.0.0 < 2.0.0 # Different lengths - longer vs shorter ("1.0.0", "1.0", False), # Should be equal ("1.0.1", "1.0", False), # 1.0.1 > 1.0.0 ("1.2.0", "1.2", False), # Should be equal ("1.2.3", "1.2", False), # 1.2.3 > 1.2.0 ("2.0.0", "2", False), # Should be equal ("2.0.1", "2", False), # 2.0.1 > 2.0.0 ("2.0", "1", False), # 2.0.0 > 1.0.0 # Mixed length comparisons ("1.0", "2.0.0", True), # 1.0.0 < 2.0.0 ("2.0", "1.0.0", False), # 2.0.0 > 1.0.0 ("1", "1.0.1", True), # 1.0.0 < 1.0.1 ("1.1", "1.0.9", False), # 1.1.0 > 1.0.9 # Large numbers ("1.9.9", "2.0.0", True), # 1.9.9 < 2.0.0 ("10.0.0", "9.9.9", False), # 10.0.0 > 9.9.9 ("1.10.0", "1.9.0", False), # 1.10.0 > 1.9.0 ("1.2.10", "1.2.9", False), # 1.2.10 > 1.2.9 # Sequential versions ("1.0.0", "1.0.1", True), ("1.0.1", "1.0.2", True), ("1.0.9", "1.1.0", True), ("1.9.9", "2.0.0", True), ] for v1, v2, expected in test_cases: options = PipelineOptions(update_compatibility_version=v1) self.assertEqual( options.is_compat_version_prior_to(v2), expected, msg=f"Failed {v1} < {v2} == {expected}") # None case: no update_compatibility_version set options_no_compat = PipelineOptions() self.assertFalse( options_no_compat.is_compat_version_prior_to("1.0.0"), msg="Should return False when update_compatibility_version is not set") if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/options/pipeline_options_validator.py0000644000175100017510000004340115206113032025351 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Pipeline options validator. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file import logging import re import string from urllib.parse import urlparse from apache_beam.internal import pickler from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PortableOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import TestOptions from apache_beam.options.pipeline_options import TypeOptions from apache_beam.options.pipeline_options import WorkerOptions _LOGGER = logging.getLogger(__name__) class PipelineOptionsValidator(object): """Validates PipelineOptions. Goes through a list of known PipelineOption subclassess and calls:: validate(validator) if one is implemented. Aggregates a list of validation errors from all and returns an aggregated list. """ # Validator will call validate on these subclasses of PipelineOptions OPTIONS = [ DebugOptions, GoogleCloudOptions, PortableOptions, SetupOptions, StandardOptions, TestOptions, TypeOptions, WorkerOptions ] # Mutually exclusive options for different types of portable environments. REQUIRED_ENVIRONMENT_OPTIONS = { 'DOCKER': [], 'PROCESS': ['process_command'], 'EXTERNAL': ['external_service_address'], 'LOOPBACK': [] } OPTIONAL_ENVIRONMENT_OPTIONS = { 'DOCKER': ['docker_container_image'], 'PROCESS': ['process_variables'], 'EXTERNAL': [], 'LOOPBACK': [] } # Possible validation errors. ERR_MISSING_OPTION = 'Missing required option: %s.' ERR_MISSING_GCS_PATH = 'Missing GCS path option: %s.' ERR_INVALID_GCS_PATH = 'Invalid GCS path (%s), given for the option: %s.' ERR_INVALID_GCS_BUCKET = ( 'Invalid GCS bucket (%s), given for the option: %s. See ' 'https://developers.google.com/storage/docs/bucketnaming ' 'for more details.') ERR_INVALID_GCS_OBJECT = 'Invalid GCS object (%s), given for the option: %s.' ERR_INVALID_JOB_NAME = ( 'Invalid job_name (%s); the name must consist of only the characters ' '[-a-z0-9], starting with a letter and ending with a letter or number') ERR_INVALID_PROJECT_NUMBER = ( 'Invalid Project ID (%s). Please make sure you specified the Project ID, ' 'not project number.') ERR_INVALID_PROJECT_ID = ( 'Invalid Project ID (%s). Please make sure you specified the Project ID, ' 'not project description.') ERR_INVALID_ENDPOINT = ( 'Invalid url (%s) for dataflow endpoint. Please provide a valid url.') ERR_INVALID_NOT_POSITIVE = ( 'Invalid value (%s) for option: %s. Value needs ' 'to be positive.') ERR_INVALID_TEST_MATCHER_TYPE = ( 'Invalid value (%s) for option: %s. Please extend your matcher object ' 'from hamcrest.core.base_matcher.BaseMatcher.') ERR_INVALID_TEST_MATCHER_UNPICKLABLE = ( 'Invalid value (%s) for option: %s. Please make sure the test matcher ' 'is unpicklable.') ERR_INVALID_TRANSFORM_NAME_MAPPING = ( 'Invalid transform name mapping format. Please make sure the mapping is ' 'string key-value pairs. Invalid pair: (%s:%s)') ERR_INVALID_ENVIRONMENT = ( 'Option %s is not compatible with environment type %s.') ERR_ENVIRONMENT_CONFIG = ( 'Option environment_config is incompatible with option(s) %s.') ERR_MISSING_REQUIRED_ENVIRONMENT_OPTION = ( 'Option %s is required for environment type %s.') ERR_NUM_WORKERS_TOO_HIGH = ( 'num_workers (%s) cannot exceed max_num_workers (%s)') ERR_REPEATABLE_OPTIONS_NOT_SET_AS_LIST = ( '(%s) is a string. Programmatically set PipelineOptions like (%s) ' 'options need to be specified as a list.') ERR_DILL_NOT_INSTALLED = ( 'Option pickle_library=dill requires dill==0.3.1.1. Install apache-beam ' 'with the dill extra e.g. apache-beam[gcp, dill]. Dill package was not ' 'found') ERR_UNSAFE_DILL_VERSION = ( 'Dill version 0.3.1.1 is required when using pickle_library=dill. Other ' 'versions of dill are untested with Apache Beam. To install the supported' ' dill version instal apache-beam[dill] extra. To use an unsupported ' 'dill version, use pickle_library=dill_unsafe. %s') # GCS path specific patterns. GCS_URI = '(?P[^:]+)://(?P[^/]+)(/(?P.*))?' GCS_BUCKET = '^[a-z0-9][-_a-z0-9.]+[a-z0-9]$' GCS_SCHEME = 'gs' # GoogleCloudOptions specific patterns. JOB_PATTERN = '[a-z]([-a-z0-9]*[a-z0-9])?' PROJECT_ID_PATTERN = '[a-z][-a-z0-9:.]+[a-z0-9]' PROJECT_NUMBER_PATTERN = '[0-9]*' def __init__(self, options, runner): self.options = options self.runner = runner def validate(self): """Calls validate on subclassess and returns a list of errors. validate will call validate method on subclasses, accumulate the returned list of errors, and returns the aggregate list. Returns: Aggregate list of errors after all calling all possible validate methods. """ errors = [] for cls in self.OPTIONS: if 'validate' in cls.__dict__ and callable(cls.__dict__['validate']): errors.extend(self.options.view_as(cls).validate(self)) return errors def is_service_runner(self): """True if pipeline will execute on the Google Cloud Dataflow service.""" is_service_runner = ( self.runner is not None and type(self.runner).__name__ in ['DataflowRunner', 'TestDataflowRunner']) dataflow_endpoint = ( self.options.view_as(GoogleCloudOptions).dataflow_endpoint) if dataflow_endpoint is None: return False else: endpoint_parts = urlparse(dataflow_endpoint, allow_fragments=False) if endpoint_parts.netloc.startswith("localhost"): return False return is_service_runner def is_full_string_match(self, pattern, string): """Returns True if the pattern matches the whole string.""" pattern = '^%s$' % pattern return re.search(pattern, string) is not None def _validate_error(self, err, *args): return [err % args] def validate_gcs_path(self, view, arg_name): """Validates a GCS path against gs://bucket/object URI format.""" arg = getattr(view, arg_name, None) if arg is None: return self._validate_error(self.ERR_MISSING_GCS_PATH, arg_name) match = re.match(self.GCS_URI, arg, re.DOTALL) if match is None: return self._validate_error(self.ERR_INVALID_GCS_PATH, arg, arg_name) scheme = match.group('SCHEME') bucket = match.group('BUCKET') gcs_object = match.group('OBJECT') if ((scheme is None) or (scheme.lower() != self.GCS_SCHEME) or (bucket is None)): return self._validate_error(self.ERR_INVALID_GCS_PATH, arg, arg_name) if not self.is_full_string_match(self.GCS_BUCKET, bucket): return self._validate_error(self.ERR_INVALID_GCS_BUCKET, arg, arg_name) if gcs_object is None or '\n' in gcs_object or '\r' in gcs_object: return self._validate_error(self.ERR_INVALID_GCS_OBJECT, arg, arg_name) return [] def validate_pickle_library(self, view): """Validates the pickle_library option.""" if view.pickle_library == 'default' or view.pickle_library == 'cloudpickle': return [] if view.pickle_library == 'dill_unsafe': return [] if view.pickle_library == 'dill': try: import dill if dill.__version__ != "0.3.1.1": return self._validate_error( self.ERR_UNSAFE_DILL_VERSION, f"Dill version found {dill.__version__}") except ImportError: return self._validate_error(self.ERR_DILL_NOT_INSTALLED) return [] def validate_cloud_options(self, view): """Validates job_name and project arguments.""" errors = [] if (view.job_name and not self.is_full_string_match(self.JOB_PATTERN, view.job_name)): errors.extend( self._validate_error(self.ERR_INVALID_JOB_NAME, view.job_name)) project = view.project if project is None: errors.extend(self._validate_error(self.ERR_MISSING_OPTION, 'project')) else: if self.is_full_string_match(self.PROJECT_NUMBER_PATTERN, project): errors.extend( self._validate_error(self.ERR_INVALID_PROJECT_NUMBER, project)) elif not self.is_full_string_match(self.PROJECT_ID_PATTERN, project): errors.extend( self._validate_error(self.ERR_INVALID_PROJECT_ID, project)) if view.update: if not view.job_name: errors.extend( self._validate_error( 'Existing job name must be provided when updating a pipeline.')) if view.transform_name_mapping: if not view.update or not self.options.view_as(StandardOptions).streaming: errors.append( 'Transform name mapping option is only useful when ' '--update and --streaming is specified') for _, (key, value) in enumerate(view.transform_name_mapping.items()): if not isinstance(key, str) or not isinstance(value, str): errors.extend( self._validate_error( self.ERR_INVALID_TRANSFORM_NAME_MAPPING, key, value)) break if view.region is None and self.is_service_runner(): default_region = self.runner.get_default_gcp_region() if default_region is None: errors.extend(self._validate_error(self.ERR_MISSING_OPTION, 'region')) else: view.region = default_region dataflow_endpoint = view.dataflow_endpoint if dataflow_endpoint is None: errors.extend( self._validate_error(self.ERR_MISSING_OPTION, dataflow_endpoint)) else: valid_endpoint = self.validate_endpoint_url(dataflow_endpoint) if valid_endpoint is False: errors.extend( self._validate_error(self.ERR_INVALID_ENDPOINT, dataflow_endpoint)) return errors def validate_sdk_container_image_options(self, view): errors = [] if view.sdk_container_image and view.worker_harness_container_image: # To be fully backwards-compatible, these options will be set to the same # value. Check that the values are different. if view.sdk_container_image != view.worker_harness_container_image: errors.extend( self._validate_error( 'Cannot use legacy flag --worker_harness_container_image along ' 'with view.sdk_container_image')) elif view.worker_harness_container_image: # Warn about legacy flag and set new flag to value of old flag. _LOGGER.warning( 'Setting sdk_container_image to value of legacy flag ' 'worker_harness_container_image.') view.sdk_container_image = view.worker_harness_container_image elif view.sdk_container_image: # Set legacy option to value of new option. view.worker_harness_container_image = view.sdk_container_image return errors def validate_container_prebuilding_options(self, view): errors = [] custom_image = self.options.view_as(WorkerOptions).sdk_container_image if (view.prebuild_sdk_container_base_image is not None and custom_image != view.prebuild_sdk_container_base_image): errors.extend( self._validate_error( 'Don\'t use the deprecated option ' '--prebuild_sdk_container_base_image. Use --sdk_container_image ' 'instead.')) return errors def validate_num_workers(self, view): """Validates that Dataflow worker number is valid.""" errors = self.validate_optional_argument_positive(view, 'num_workers') errors.extend( self.validate_optional_argument_positive(view, 'max_num_workers')) num_workers = view.num_workers max_num_workers = view.max_num_workers if (num_workers is not None and max_num_workers is not None and num_workers > max_num_workers): errors.extend( self._validate_error( self.ERR_NUM_WORKERS_TOO_HIGH, num_workers, max_num_workers)) return errors def validate_worker_region_zone(self, view): """Validates Dataflow worker region and zone arguments are consistent.""" errors = [] if view.zone and (view.worker_region or view.worker_zone): errors.extend( self._validate_error( 'Cannot use deprecated flag --zone along with worker_region or ' 'worker_zone.')) if self.options.view_as(DebugOptions).lookup_experiment('worker_region') \ and (view.worker_region or view.worker_zone): errors.extend( self._validate_error( 'Cannot use deprecated experiment worker_region along with ' 'worker_region or worker_zone.')) if view.worker_region and view.worker_zone: errors.extend( self._validate_error( 'worker_region and worker_zone are mutually exclusive.')) if view.zone: _LOGGER.warning( 'Option --zone is deprecated. Please use --worker_zone instead.') view.worker_zone = view.zone view.zone = None return errors def validate_optional_argument_positive(self, view, arg_name): """Validates that an optional argument (if set) has a positive value.""" arg = getattr(view, arg_name, None) if arg is not None and int(arg) <= 0: return self._validate_error(self.ERR_INVALID_NOT_POSITIVE, arg, arg_name) return [] def validate_test_matcher(self, view, arg_name): """Validates that on_success_matcher argument if set. Validates that on_success_matcher is unpicklable and is instance of `hamcrest.core.base_matcher.BaseMatcher`. """ # This is a test only method and requires hamcrest from hamcrest.core.base_matcher import BaseMatcher pickled_matcher = view.on_success_matcher errors = [] try: matcher = pickler.loads(pickled_matcher) if not isinstance(matcher, BaseMatcher): errors.extend( self._validate_error( self.ERR_INVALID_TEST_MATCHER_TYPE, matcher, arg_name)) except: # pylint: disable=bare-except errors.extend( self._validate_error( self.ERR_INVALID_TEST_MATCHER_UNPICKLABLE, pickled_matcher, arg_name)) return errors def validate_environment_options(self, view): """Validates portable environment options.""" errors = [] actual_environment_type = ( view.environment_type.upper() if view.environment_type else None) for environment_type, required in self.REQUIRED_ENVIRONMENT_OPTIONS.items(): found_required_options = [ opt for opt in required if view.lookup_environment_option(opt) is not None ] found_optional_options = [ opt for opt in self.OPTIONAL_ENVIRONMENT_OPTIONS[environment_type] if view.lookup_environment_option(opt) is not None ] found_options = found_required_options + found_optional_options if environment_type == actual_environment_type: if view.environment_config: if found_options: errors.extend( self._validate_error( self.ERR_ENVIRONMENT_CONFIG, ', '.join(found_options))) else: missing_options = set(required).difference( set(found_required_options)) for opt in missing_options: errors.extend( self._validate_error( self.ERR_MISSING_REQUIRED_ENVIRONMENT_OPTION, opt, environment_type)) else: # Environment options classes are mutually exclusive. for opt in found_options: errors.extend( self._validate_error( self.ERR_INVALID_ENVIRONMENT, opt, actual_environment_type)) if actual_environment_type == 'LOOPBACK' and view.environment_config: errors.extend( self._validate_error( self.ERR_INVALID_ENVIRONMENT, 'environment_config', 'LOOPBACK')) return errors def validate_repeatable_argument_passed_as_list(self, view, arg_name): """Validates that repeatable PipelineOptions like dataflow_service_options or experiments are specified as a list when set programmatically. This way, users do not inadvertently specify it as a string, mirroring the way they are set via the command lineRepeatable options, which are as passed a list. """ arg = getattr(view, arg_name, None) if not isinstance(arg, list): return self._validate_error( self.ERR_REPEATABLE_OPTIONS_NOT_SET_AS_LIST, arg, arg_name) return [] # Minimally validates the endpoint url. This is not a strict application # of http://www.faqs.org/rfcs/rfc1738.html. # If the url matches localhost, set def validate_endpoint_url(self, endpoint_url): url_parts = urlparse(endpoint_url, allow_fragments=False) if not url_parts.scheme or not url_parts.netloc: return False if url_parts.scheme not in ['http', 'https']: return False if set(url_parts.netloc) <= set(string.ascii_letters + string.digits + '-.'): return True return False ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/options/pipeline_options_validator_test.py0000644000175100017510000007041315206113032026413 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the pipeline options validator module.""" # pytype: skip-file import logging import unittest import pytest from hamcrest import assert_that from hamcrest import contains_string from hamcrest import only_contains from hamcrest.core.base_matcher import BaseMatcher from parameterized import parameterized from apache_beam.internal import pickler from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import WorkerOptions from apache_beam.options.pipeline_options_validator import PipelineOptionsValidator # Mock runners to use for validations. class MockRunners(object): class DataflowRunner(object): def get_default_gcp_region(self): # Return a default so we don't have to specify --region in every test # (unless specifically testing it). return 'us-central1' class TestDataflowRunner(DataflowRunner): pass class OtherRunner(object): pass # Matcher that always passes for testing on_success_matcher option class AlwaysPassMatcher(BaseMatcher): def _matches(self, item): return True class SetupTest(unittest.TestCase): def check_errors_for_arguments(self, errors, args): """Checks that there is exactly one error for each given argument.""" missing = [] remaining = list(errors) for arg in args: found = False for error in remaining: if arg in error: remaining.remove(error) found = True break if not found: missing.append('Missing error for: %s.' % arg) # Return missing and remaining (not matched) errors. return missing + remaining def test_local_runner(self): runner = MockRunners.OtherRunner() options = PipelineOptions([]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 0) def test_missing_required_options(self): options = PipelineOptions(['']) runner = MockRunners.DataflowRunner() # Remove default region for this test. runner.get_default_gcp_region = lambda: None validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual( self.check_errors_for_arguments( errors, ['project', 'staging_location', 'temp_location', 'region']), []) @parameterized.expand([ (None, 'gs://foo/bar', []), (None, None, ['staging_location', 'temp_location']), ('gs://foo/bar', None, []), ('gs://foo/bar', 'gs://ABC/bar', []), ('gcs:/foo/bar', 'gs://foo/bar', []), ('gs:/foo/bar', 'gs://foo/bar', []), ('gs://ABC/bar', 'gs://foo/bar', []), ('gs://ABC/bar', 'gs://BCD/bar', ['temp_location', 'staging_location' ]), ('gs://foo', 'gs://foo/bar', []), ('gs://foo/', 'gs://foo/bar', []), ('gs://foo/bar', 'gs://foo/bar', []) ]) @unittest.skip('Not compatible with new GCS client. See GH issue #26335.') def test_gcs_path(self, temp_location, staging_location, expected_error_args): def get_validator(_temp_location, _staging_location): options = ['--project=example:example', '--job_name=job'] if _temp_location is not None: options.append('--temp_location=' + _temp_location) if _staging_location is not None: options.append('--staging_location=' + _staging_location) pipeline_options = PipelineOptions(options) runner = MockRunners.DataflowRunner() validator = PipelineOptionsValidator(pipeline_options, runner) return validator errors = get_validator(temp_location, staging_location).validate() self.assertEqual( self.check_errors_for_arguments(errors, expected_error_args), []) @parameterized.expand([(None, ['project']), ('12345', ['project']), ('FOO', ['project']), ('foo:BAR', ['project']), ('fo', ['project']), ('foo', []), ('foo:bar', [])]) def test_project(self, project, expected_error_args): def get_validator(_project): options = [ '--job_name=job', '--staging_location=gs://foo/bar', '--temp_location=gs://foo/bar' ] if _project is not None: options.append('--project=' + _project) pipeline_options = PipelineOptions(options) runner = MockRunners.DataflowRunner() validator = PipelineOptionsValidator(pipeline_options, runner) return validator errors = get_validator(project).validate() self.assertEqual( self.check_errors_for_arguments(errors, expected_error_args), []) @parameterized.expand([(None, []), ('12345', ['job_name']), ('FOO', ['job_name']), ('foo:bar', ['job_name']), ('fo', []), ('foo', [])]) def test_job_name(self, job_name, expected_error_args): def get_validator(_job_name): options = [ '--project=example:example', '--staging_location=gs://foo/bar', '--temp_location=gs://foo/bar' ] if _job_name is not None: options.append('--job_name=' + _job_name) pipeline_options = PipelineOptions(options) runner = MockRunners.DataflowRunner() validator = PipelineOptionsValidator(pipeline_options, runner) return validator errors = get_validator(job_name).validate() self.assertEqual( self.check_errors_for_arguments(errors, expected_error_args), []) @parameterized.expand([(None, []), ('1', []), ('0', ['num_workers']), ('-1', ['num_workers'])]) def test_num_workers(self, num_workers, expected_error_args): def get_validator(_num_workers): options = [ '--project=example:example', '--job_name=job', '--staging_location=gs://foo/bar', '--temp_location=gs://foo/bar' ] if _num_workers is not None: options.append('--num_workers=' + _num_workers) pipeline_options = PipelineOptions(options) runner = MockRunners.DataflowRunner() validator = PipelineOptionsValidator(pipeline_options, runner) return validator errors = get_validator(num_workers).validate() self.assertEqual( self.check_errors_for_arguments(errors, expected_error_args), []) @parameterized.expand([ ( MockRunners.OtherRunner(), [], False, ), ( MockRunners.OtherRunner(), ['--dataflow_endpoint=https://dataflow.googleapis.com'], False, ), ( MockRunners.OtherRunner(), ['--dataflow_endpoint=https://dataflow.googleapis.com/'], False, ), ( MockRunners.DataflowRunner(), [], True, ), ( MockRunners.DataflowRunner(), ['--dataflow_endpoint=https://another.service.com'], True, ), ( MockRunners.DataflowRunner(), ['--dataflow_endpoint=https://dataflow.googleapis.com'], True, ), ( MockRunners.DataflowRunner(), ['--dataflow_endpoint=http://localhost:1000'], False, ), ( MockRunners.DataflowRunner(), ['--dataflow_endpoint=foo: //dataflow. googleapis. com'], True, ), ( MockRunners.DataflowRunner(), [], True, ) ]) def test_is_service_runner(self, runner, options, expected): validator = PipelineOptionsValidator(PipelineOptions(options), runner) self.assertEqual(validator.is_service_runner(), expected) def test_pickle_library_dill_not_installed_returns_error(self): runner = MockRunners.OtherRunner() # Remove default region for this test. options = PipelineOptions(['--pickle_library=dill']) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 1, errors) self.assertIn("Option pickle_library=dill requires dill", errors[0]) @pytest.mark.uses_dill def test_pickle_library_dill_installed_returns_no_error(self): pytest.importorskip("dill") runner = MockRunners.OtherRunner() # Remove default region for this test. options = PipelineOptions(['--pickle_library=dill']) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 0, errors) @pytest.mark.uses_dill def test_pickle_library_dill_installed_returns_wrong_version(self): pytest.importorskip("dill") with unittest.mock.patch('dill.__version__', '0.3.6'): runner = MockRunners.OtherRunner() # Remove default region for this test. options = PipelineOptions(['--pickle_library=dill']) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 1, errors) self.assertIn("Dill version 0.3.1.1 is required when using ", errors[0]) @pytest.mark.uses_dill def test_pickle_library_dill_unsafe_no_error(self): pytest.importorskip("dill") with unittest.mock.patch('dill.__version__', '0.3.6'): runner = MockRunners.OtherRunner() # Remove default region for this test. options = PipelineOptions(['--pickle_library=dill_unsafe']) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 0, errors) def test_dataflow_job_file_and_template_location_mutually_exclusive(self): runner = MockRunners.OtherRunner() options = PipelineOptions( ['--template_location', 'abc', '--dataflow_job_file', 'def']) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertTrue(errors) def test_validate_template_location(self): runner = MockRunners.OtherRunner() options = PipelineOptions([ '--template_location', 'abc', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertFalse(errors) def test_validate_dataflow_job_file(self): runner = MockRunners.OtherRunner() options = PipelineOptions(['--dataflow_job_file', 'abc']) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertFalse(errors) def test_num_workers_is_positive(self): runner = MockRunners.DataflowRunner() options = PipelineOptions([ '--num_workers=-1', '--worker_region=us-east1', '--project=example:example', '--temp_location=gs://foo/bar', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 1) self.assertIn('num_workers', errors[0]) self.assertIn('-1', errors[0]) def test_max_num_workers_is_positive(self): runner = MockRunners.DataflowRunner() options = PipelineOptions([ '--max_num_workers=-1', '--worker_region=us-east1', '--project=example:example', '--temp_location=gs://foo/bar', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 1) self.assertIn('max_num_workers', errors[0]) self.assertIn('-1', errors[0]) def test_num_workers_cannot_exceed_max_num_workers(self): runner = MockRunners.DataflowRunner() options = PipelineOptions([ '--num_workers=43', '--max_num_workers=42', '--worker_region=us-east1', '--project=example:example', '--temp_location=gs://foo/bar', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 1) self.assertIn('num_workers', errors[0]) self.assertIn('43', errors[0]) self.assertIn('max_num_workers', errors[0]) self.assertIn('42', errors[0]) def test_num_workers_can_equal_max_num_workers(self): runner = MockRunners.DataflowRunner() options = PipelineOptions([ '--num_workers=42', '--max_num_workers=42', '--worker_region=us-east1', '--project=example:example', '--temp_location=gs://foo/bar', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 0) def test_zone_and_worker_region_mutually_exclusive(self): runner = MockRunners.DataflowRunner() options = PipelineOptions([ '--zone', 'us-east1-b', '--worker_region', 'us-east1', '--project=example:example', '--temp_location=gs://foo/bar', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 1) self.assertIn('zone', errors[0]) self.assertIn('worker_region', errors[0]) def test_zone_and_worker_zone_mutually_exclusive(self): runner = MockRunners.DataflowRunner() options = PipelineOptions([ '--zone', 'us-east1-b', '--worker_zone', 'us-east1-c', '--project=example:example', '--temp_location=gs://foo/bar', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 1) self.assertIn('zone', errors[0]) self.assertIn('worker_zone', errors[0]) def test_experiment_region_and_worker_region_mutually_exclusive(self): runner = MockRunners.DataflowRunner() options = PipelineOptions([ '--experiments', 'worker_region=us-west1', '--worker_region', 'us-east1', '--project=example:example', '--temp_location=gs://foo/bar', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 1) self.assertIn('experiment', errors[0]) self.assertIn('worker_region', errors[0]) def test_region_and_worker_zone_mutually_exclusive(self): runner = MockRunners.DataflowRunner() options = PipelineOptions([ '--worker_region=us-west1', '--worker_zone=us-east1-b', '--project=example:example', '--temp_location=gs://foo/bar', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 1) self.assertIn('worker_region', errors[0]) self.assertIn('worker_zone', errors[0]) def test_programmatically_set_experiment_passed_as_string(self): runner = MockRunners.DataflowRunner() options = PipelineOptions( project='example.com:example', temp_location='gs://foo/bar/', experiments='enable_prime', dataflow_service_options='use_runner_v2', ) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 2) self.assertIn('experiments', errors[0]) self.assertIn('dataflow_service_options', errors[1]) def test_programmatically_set_experiment_passed_as_list(self): runner = MockRunners.DataflowRunner() options = PipelineOptions( project='example.com:example', temp_location='gs://foo/bar/', experiments=['enable_prime'], dataflow_service_options=['use_runner_v2'], ) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 0) self.assertEqual( options.view_as(DebugOptions).experiments, ['enable_prime']) self.assertEqual( options.view_as(GoogleCloudOptions).dataflow_service_options, ['use_runner_v2']) def test_worker_region_and_worker_zone_mutually_exclusive(self): runner = MockRunners.DataflowRunner() options = PipelineOptions([ '--worker_region', 'us-east1', '--worker_zone', 'us-east1-b', '--project=example:example', '--temp_location=gs://foo/bar', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 1) self.assertIn('worker_region', errors[0]) self.assertIn('worker_zone', errors[0]) def test_zone_alias_worker_zone(self): runner = MockRunners.DataflowRunner() options = PipelineOptions([ '--zone=us-east1-b', '--project=example:example', '--temp_location=gs://foo/bar', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 0) self.assertIsNone(options.view_as(WorkerOptions).zone) self.assertEqual(options.view_as(WorkerOptions).worker_zone, 'us-east1-b') def test_region_optional_for_non_service_runner(self): runner = MockRunners.OtherRunner() # Remove default region for this test. runner.get_default_gcp_region = lambda: None options = PipelineOptions([ '--project=example:example', '--temp_location=gs://foo/bar', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 0) def test_dataflow_endpoint_is_a_url(self): runner = MockRunners.DataflowRunner() # Remove default region for this test. options = PipelineOptions([ '--project=example:example', '--temp_location=gs://foo/bar', '--staging_location=gs://foo/baz', '--dataflow_endpoint=foo and bar' ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 1, errors) self.assertIn("Invalid url (foo and bar)", errors[0]) def test_alias_sdk_container_to_worker_harness(self): runner = MockRunners.DataflowRunner() test_image = "SDK_IMAGE" options = PipelineOptions([ '--sdk_container_image=%s' % test_image, '--project=example:example', '--temp_location=gs://foo/bar', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 0) self.assertEqual( options.view_as(WorkerOptions).worker_harness_container_image, test_image) self.assertEqual( options.view_as(WorkerOptions).sdk_container_image, test_image) def test_alias_worker_harness_sdk_container_image(self): runner = MockRunners.DataflowRunner() test_image = "WORKER_HARNESS" options = PipelineOptions([ '--worker_harness_container_image=%s' % test_image, '--project=example:example', '--temp_location=gs://foo/bar', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 0) self.assertEqual( options.view_as(WorkerOptions).worker_harness_container_image, test_image) self.assertEqual( options.view_as(WorkerOptions).sdk_container_image, test_image) def test_worker_harness_sdk_container_image_mutually_exclusive(self): runner = MockRunners.DataflowRunner() options = PipelineOptions([ '--worker_harness_container_image=WORKER', '--sdk_container_image=SDK_ONLY', '--project=example:example', '--temp_location=gs://foo/bar', ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 1) self.assertIn('sdk_container_image', errors[0]) self.assertIn('worker_harness_container_image', errors[0]) def test_prebuild_sdk_container_base_image_disallowed(self): runner = MockRunners.DataflowRunner() options = PipelineOptions([ '--project=example:example', '--temp_location=gs://foo/bar', '--prebuild_sdk_container_base_image=gcr.io/foo:bar' ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 1) self.assertIn('prebuild_sdk_container_base_image', errors[0]) self.assertIn('sdk_container_image', errors[0]) def test_prebuild_sdk_container_base_allowed_if_matches_custom_image(self): runner = MockRunners.DataflowRunner() options = PipelineOptions([ '--project=example:example', '--temp_location=gs://foo/bar', '--sdk_container_image=gcr.io/foo:bar', '--prebuild_sdk_container_base_image=gcr.io/foo:bar' ]) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertEqual(len(errors), 0) @parameterized.expand([(None, []), (pickler.dumps(AlwaysPassMatcher()), []), (b'abc', ['on_success_matcher']), (pickler.dumps(object), ['on_success_matcher'])]) def test_test_matcher(self, on_success_matcher, errors): def get_validator(matcher): options = [ '--project=example:example', '--job_name=job', '--staging_location=gs://foo/bar', '--temp_location=gs://foo/bar', ] if matcher: options.append('%s=%s' % ('--on_success_matcher', matcher.decode())) pipeline_options = PipelineOptions(options) runner = MockRunners.TestDataflowRunner() return PipelineOptionsValidator(pipeline_options, runner) errors = get_validator(on_success_matcher).validate() self.assertEqual(self.check_errors_for_arguments(errors, errors), []) def test_transform_name_mapping_without_update(self): options = [ '--project=example:example', '--staging_location=gs://foo/bar', '--temp_location=gs://foo/bar', '--transform_name_mapping={\"fromPardo\":\"toPardo\"}' ] pipeline_options = PipelineOptions(options) runner = MockRunners.DataflowRunner() validator = PipelineOptionsValidator(pipeline_options, runner) errors = validator.validate() assert_that( errors, only_contains( contains_string( 'Transform name mapping option is only useful when ' '--update and --streaming is specified'))) def test_transform_name_mapping_invalid_format(self): options = [ '--project=example:example', '--staging_location=gs://foo/bar', '--temp_location=gs://foo/bar', '--update', '--job_name=test', '--streaming', '--transform_name_mapping={\"fromPardo\":123}' ] pipeline_options = PipelineOptions(options) runner = MockRunners.DataflowRunner() validator = PipelineOptionsValidator(pipeline_options, runner) errors = validator.validate() assert_that( errors, only_contains( contains_string('Invalid transform name mapping format.'))) def test_type_check_additional(self): runner = MockRunners.OtherRunner() options = PipelineOptions(['--type_check_additional=all']) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertFalse(errors) options = PipelineOptions(['--type_check_additional=']) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertFalse(errors) def test_type_check_additional_unrecognized_feature(self): runner = MockRunners.OtherRunner() options = PipelineOptions(['--type_check_additional=all,dfgdf']) validator = PipelineOptionsValidator(options, runner) errors = validator.validate() self.assertTrue(errors) @parameterized.expand([ (['--environment_type=dOcKeR'], []), ([ '--environment_type=dOcKeR', '--environment_options=docker_container_image=foo' ], []), (['--environment_type=dOcKeR', '--environment_config=foo'], []), ([ '--environment_type=dOcKeR', '--environment_options=docker_container_image=foo', '--environment_config=foo' ], ['environment_config']), ([ '--environment_type=dOcKeR', '--environment_options=process_command=foo', '--environment_options=process_variables=foo=bar', '--environment_options=external_service_address=foo' ], ['process_command', 'process_variables', 'external_service_address']), (['--environment_type=pRoCeSs'], ['process_command']), ([ '--environment_type=pRoCeSs', '--environment_options=process_command=foo' ], []), (['--environment_type=pRoCeSs', '--environment_config=foo'], []), ([ '--environment_type=pRoCeSs', '--environment_options=process_command=foo', '--environment_config=foo' ], ['environment_config']), ([ '--environment_type=pRoCeSs', '--environment_options=process_command=foo', '--environment_options=process_variables=foo=bar', '--environment_options=docker_container_image=foo', '--environment_options=external_service_address=foo' ], ['docker_container_image', 'external_service_address']), (['--environment_type=eXtErNaL'], ['external_service_address']), ([ '--environment_type=eXtErNaL', '--environment_options=external_service_address=foo' ], []), (['--environment_type=eXtErNaL', '--environment_config=foo'], []), ([ '--environment_type=eXtErNaL', '--environment_options=external_service_address=foo', '--environment_config=foo' ], ['environment_config']), ([ '--environment_type=eXtErNaL', '--environment_options=external_service_address=foo', '--environment_options=process_command=foo', '--environment_options=process_variables=foo=bar', '--environment_options=docker_container_image=foo' ], ['process_command', 'process_variables', 'docker_container_image']), (['--environment_type=lOoPbACk'], []), (['--environment_type=lOoPbACk', '--environment_config=foo'], ['environment_config']), ([ '--environment_type=lOoPbACk', '--environment_options=docker_container_image=foo', '--environment_options=process_command=foo', '--environment_options=process_variables=foo=bar', '--environment_options=external_service_address=foo' ], [ 'docker_container_image', 'process_command', 'process_variables', 'external_service_address' ]), (['--environment_type=beam:env:foo:v1'], []), (['--environment_type=beam:env:foo:v1', '--environment_config=foo'], []), ([ '--environment_type=beam:env:foo:v1', '--environment_options=docker_container_image=foo', '--environment_options=process_command=foo', '--environment_options=process_variables=foo=bar', '--environment_options=external_service_address=foo' ], [ 'docker_container_image', 'process_command', 'process_variables', 'external_service_address' ]), ([ '--environment_options=docker_container_image=foo', '--environment_options=process_command=foo', '--environment_options=process_variables=foo=bar', '--environment_options=external_service_address=foo' ], [ 'docker_container_image', 'process_command', 'process_variables', 'external_service_address' ]) ]) def test_environment_options(self, options, expected_error_args): errors = [] validator = PipelineOptionsValidator( PipelineOptions(options), MockRunners.OtherRunner()) validation_result = validator.validate() validation_errors = self.check_errors_for_arguments( validation_result, expected_error_args) if validation_errors: errors.append( 'Options "%s" had unexpected validation results: "%s"' % (' '.join(options), ' '.join(validation_errors))) self.assertEqual(errors, [], expected_error_args) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/options/value_provider.py0000644000175100017510000001324415206113032022754 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A ValueProvider abstracts the notion of fetching a value that may or may not be currently available. This can be used to parameterize transforms that only read values in at runtime, for example. """ # pytype: skip-file from functools import wraps from apache_beam import error __all__ = [ 'ValueProvider', 'StaticValueProvider', 'RuntimeValueProvider', 'NestedValueProvider', 'check_accessible', ] class ValueProvider(object): """Base class that all other ValueProviders must implement. """ def is_accessible(self): """Whether the contents of this ValueProvider is available to routines that run at graph construction time. """ raise NotImplementedError( 'ValueProvider.is_accessible implemented in derived classes') def get(self): """Return the value wrapped by this ValueProvider. """ raise NotImplementedError( 'ValueProvider.get implemented in derived classes') class StaticValueProvider(ValueProvider): """StaticValueProvider is an implementation of ValueProvider that allows for a static value to be provided. """ def __init__(self, value_type, value): """ Args: value_type: Type of the static value value: Static value """ self.value_type = value_type self.value = value_type(value) def is_accessible(self): return True def get(self): return self.value def __str__(self): return str(self.value) def __eq__(self, other): if self.value == other: return True if isinstance(other, StaticValueProvider): if (self.value_type == other.value_type and self.value == other.value): return True return False def __hash__(self): return hash((type(self), self.value_type, self.value)) class RuntimeValueProvider(ValueProvider): """RuntimeValueProvider is an implementation of ValueProvider that allows for a value to be provided at execution time rather than at graph construction time. """ runtime_options = None experiments: set[str] = set() def __init__(self, option_name, value_type, default_value): self.option_name = option_name self.default_value = default_value self.value_type = value_type def is_accessible(self): return RuntimeValueProvider.runtime_options is not None @classmethod def get_value(cls, option_name, value_type, default_value): if not RuntimeValueProvider.runtime_options: return default_value candidate = RuntimeValueProvider.runtime_options.get(option_name) if candidate: return value_type(candidate) else: return default_value def get(self): if RuntimeValueProvider.runtime_options is None: raise error.RuntimeValueProviderError( '%s.get() not called from a runtime context' % self) return RuntimeValueProvider.get_value( self.option_name, self.value_type, self.default_value) @classmethod def set_runtime_options(cls, pipeline_options): RuntimeValueProvider.runtime_options = pipeline_options RuntimeValueProvider.experiments = RuntimeValueProvider.get_value( 'experiments', set, set()) def __str__(self): return '%s(option: %s, type: %s, default_value: %s)' % ( self.__class__.__name__, self.option_name, self.value_type.__name__, repr(self.default_value)) class NestedValueProvider(ValueProvider): """NestedValueProvider is an implementation of ValueProvider that allows for wrapping another ValueProvider object. """ def __init__(self, value, translator): """Creates a NestedValueProvider that wraps the provided ValueProvider. Args: value: ValueProvider object to wrap translator: function that is applied to the ValueProvider Raises: ``RuntimeValueProviderError``: if any of the provided objects are not accessible. """ self.value = value self.translator = translator def is_accessible(self): return self.value.is_accessible() def get(self): try: return self.cached_value except AttributeError: self.cached_value = self.translator(self.value.get()) return self.cached_value def __str__(self): return "%s(value: %s, translator: %s)" % ( self.__class__.__name__, self.value, self.translator.__name__, ) def check_accessible(value_provider_list): """A decorator that checks accessibility of a list of ValueProvider objects. Args: value_provider_list: list of ValueProvider objects Raises: ``RuntimeValueProviderError``: if any of the provided objects are not accessible. """ assert isinstance(value_provider_list, list) def _check_accessible(fnc): @wraps(fnc) def _f(self, *args, **kwargs): for obj in [getattr(self, vp) for vp in value_provider_list]: if not obj.is_accessible(): raise error.RuntimeValueProviderError('%s not accessible' % obj) return fnc(self, *args, **kwargs) return _f return _check_accessible ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/options/value_provider_test.py0000644000175100017510000002437115206113032024016 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the ValueProvider class.""" # pytype: skip-file import logging import unittest from mock import Mock from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.value_provider import NestedValueProvider from apache_beam.options.value_provider import RuntimeValueProvider from apache_beam.options.value_provider import StaticValueProvider # TODO(https://github.com/apache/beam/issues/18197): Require unique names only # within a test. For now, _vp_arg will be the # convention to name value-provider arguments in tests, as opposed to # _non_vp_arg for non-value-provider arguments. # The number will grow per file as tests are added. class ValueProviderTests(unittest.TestCase): def setUp(self): # Reset runtime options to avoid side-effects caused by other tests. # Note that is_accessible assertions require runtime_options to # be uninitialized. RuntimeValueProvider.set_runtime_options(None) def tearDown(self): # Reset runtime options to avoid side-effects in other tests. RuntimeValueProvider.set_runtime_options(None) def test_static_value_provider_keyword_argument(self): class UserDefinedOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument( '--vpt_vp_arg1', help='This keyword argument is a value provider', default='some value') options = UserDefinedOptions(['--vpt_vp_arg1', 'abc']) self.assertTrue(isinstance(options.vpt_vp_arg1, StaticValueProvider)) self.assertTrue(options.vpt_vp_arg1.is_accessible()) self.assertEqual(options.vpt_vp_arg1.get(), 'abc') def test_runtime_value_provider_keyword_argument(self): class UserDefinedOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument( '--vpt_vp_arg2', help='This keyword argument is a value provider') options = UserDefinedOptions() self.assertTrue(isinstance(options.vpt_vp_arg2, RuntimeValueProvider)) self.assertFalse(options.vpt_vp_arg2.is_accessible()) with self.assertRaises(RuntimeError): options.vpt_vp_arg2.get() def test_static_value_provider_positional_argument(self): class UserDefinedOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument( 'vpt_vp_arg3', help='This positional argument is a value provider', default='some value') options = UserDefinedOptions(['abc']) self.assertTrue(isinstance(options.vpt_vp_arg3, StaticValueProvider)) self.assertTrue(options.vpt_vp_arg3.is_accessible()) self.assertEqual(options.vpt_vp_arg3.get(), 'abc') def test_runtime_value_provider_positional_argument(self): class UserDefinedOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument( 'vpt_vp_arg4', help='This positional argument is a value provider') options = UserDefinedOptions([]) self.assertTrue(isinstance(options.vpt_vp_arg4, RuntimeValueProvider)) self.assertFalse(options.vpt_vp_arg4.is_accessible()) with self.assertRaises(RuntimeError): options.vpt_vp_arg4.get() def test_static_value_provider_type_cast(self): class UserDefinedOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument( '--vpt_vp_arg5', type=int, help='This flag is a value provider') options = UserDefinedOptions(['--vpt_vp_arg5', '123']) self.assertTrue(isinstance(options.vpt_vp_arg5, StaticValueProvider)) self.assertTrue(options.vpt_vp_arg5.is_accessible()) self.assertEqual(options.vpt_vp_arg5.get(), 123) def test_set_runtime_option(self): # define ValueProvider options, with and without default values class UserDefinedOptions1(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument( '--vpt_vp_arg6', help='This keyword argument is a value provider') # set at runtime parser.add_value_provider_argument( # not set, had default int '-v', '--vpt_vp_arg7', # with short form default=123, type=int) parser.add_value_provider_argument( # not set, had default str '--vpt_vp-arg8', # with dash in name default='123', type=str) parser.add_value_provider_argument( # not set and no default '--vpt_vp_arg9', type=float) parser.add_value_provider_argument( # positional argument set 'vpt_vp_arg10', # default & runtime ignored help='This positional argument is a value provider', type=float, default=5.4) # provide values at graph-construction time # (options not provided here become of the type RuntimeValueProvider) options = UserDefinedOptions1(['1.2']) self.assertFalse(options.vpt_vp_arg6.is_accessible()) self.assertFalse(options.vpt_vp_arg7.is_accessible()) self.assertFalse(options.vpt_vp_arg8.is_accessible()) self.assertFalse(options.vpt_vp_arg9.is_accessible()) self.assertTrue(options.vpt_vp_arg10.is_accessible()) # provide values at job-execution time # (options not provided here will use their default, if they have one) RuntimeValueProvider.set_runtime_options({ 'vpt_vp_arg6': 'abc', 'vpt_vp_arg10': '3.2' }) self.assertTrue(options.vpt_vp_arg6.is_accessible()) self.assertEqual(options.vpt_vp_arg6.get(), 'abc') self.assertTrue(options.vpt_vp_arg7.is_accessible()) self.assertEqual(options.vpt_vp_arg7.get(), 123) self.assertTrue(options.vpt_vp_arg8.is_accessible()) self.assertEqual(options.vpt_vp_arg8.get(), '123') self.assertTrue(options.vpt_vp_arg9.is_accessible()) self.assertIsNone(options.vpt_vp_arg9.get()) self.assertTrue(options.vpt_vp_arg10.is_accessible()) self.assertEqual(options.vpt_vp_arg10.get(), 1.2) def test_choices(self): class UserDefinedOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--vpt_vp_arg11', choices=['a', 'b'], help='This flag is a value provider with concrete choices') parser.add_argument( '--vpt_vp_arg12', choices=[1, 2], type=int, help='This flag is a value provider with concrete choices') options = UserDefinedOptions(['--vpt_vp_arg11', 'a', '--vpt_vp_arg12', '2']) self.assertEqual(options.vpt_vp_arg11, 'a') self.assertEqual(options.vpt_vp_arg12, 2) def test_static_value_provider_choices(self): class UserDefinedOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument( '--vpt_vp_arg13', choices=['a', 'b'], help='This flag is a value provider with concrete choices') parser.add_value_provider_argument( '--vpt_vp_arg14', choices=[1, 2], type=int, help='This flag is a value provider with concrete choices') options = UserDefinedOptions(['--vpt_vp_arg13', 'a', '--vpt_vp_arg14', '2']) self.assertEqual(options.vpt_vp_arg13.get(), 'a') self.assertEqual(options.vpt_vp_arg14.get(), 2) def test_experiments_setup(self): self.assertFalse('feature_1' in RuntimeValueProvider.experiments) RuntimeValueProvider.set_runtime_options( {'experiments': ['feature_1', 'feature_2']}) self.assertTrue(isinstance(RuntimeValueProvider.experiments, set)) self.assertTrue('feature_1' in RuntimeValueProvider.experiments) self.assertTrue('feature_2' in RuntimeValueProvider.experiments) def test_experiments_options_setup(self): options = PipelineOptions(['--experiments', 'a', '--experiments', 'b,c']) options = options.view_as(DebugOptions) self.assertIn('a', options.experiments) self.assertIn('b', options.experiments) self.assertIn('c', options.experiments) def test_nested_value_provider_wrap_static(self): vp = NestedValueProvider(StaticValueProvider(int, 1), lambda x: x + 1) self.assertTrue(vp.is_accessible()) self.assertEqual(vp.get(), 2) def test_nested_value_provider_caches_value(self): mock_fn = Mock() def translator(x): mock_fn() return x vp = NestedValueProvider(StaticValueProvider(int, 1), translator) vp.get() self.assertEqual(mock_fn.call_count, 1) vp.get() self.assertEqual(mock_fn.call_count, 1) def test_nested_value_provider_wrap_runtime(self): class UserDefinedOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument( '--vpt_vp_arg15', help='This keyword argument is a value provider') # set at runtime options = UserDefinedOptions([]) vp = NestedValueProvider(options.vpt_vp_arg15, lambda x: x + x) self.assertFalse(vp.is_accessible()) RuntimeValueProvider.set_runtime_options({'vpt_vp_arg15': 'abc'}) self.assertTrue(vp.is_accessible()) self.assertEqual(vp.get(), 'abcabc') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/pipeline.py0000644000175100017510000021552315206113032020044 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Pipeline, the top-level Beam object. A pipeline holds a DAG of data transforms. Conceptually the nodes of the DAG are transforms (:class:`~apache_beam.transforms.ptransform.PTransform` objects) and the edges are values (mostly :class:`~apache_beam.pvalue.PCollection` objects). The transforms take as inputs one or more PValues and output one or more :class:`~apache_beam.pvalue.PValue` s. The pipeline offers functionality to traverse the graph. The actual operation to be executed for each node visited is specified through a runner object. Typical usage:: # Create a pipeline object using a local runner for execution. with beam.Pipeline('DirectRunner') as p: # Add to the pipeline a "Create" transform. When executed this # transform will produce a PCollection object with the specified values. pcoll = p | 'Create' >> beam.Create([1, 2, 3]) # Another transform could be applied to pcoll, e.g., writing to a text file. # For other transforms, refer to transforms/ directory. pcoll | 'Write' >> beam.io.WriteToText('./output') # run() will execute the DAG stored in the pipeline. The execution of the # nodes visited is done using the specified local runner. """ # pytype: skip-file # mypy: disallow-untyped-defs import abc import contextlib import copy import logging import os import re import shutil import tempfile import threading import unicodedata import uuid from collections import defaultdict from collections.abc import Iterable from collections.abc import Mapping from collections.abc import Sequence from typing import TYPE_CHECKING from typing import Any from typing import Optional from typing import Union from google.protobuf import message from apache_beam import pvalue from apache_beam.coders import typecoders from apache_beam.internal import pickler from apache_beam.io.filesystems import FileSystems from apache_beam.options.pipeline_options import CrossLanguageOptions from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import StreamingOptions from apache_beam.options.pipeline_options import TypeOptions from apache_beam.options.pipeline_options_context import scoped_pipeline_options from apache_beam.options.pipeline_options_validator import PipelineOptionsValidator from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners import PipelineRunner from apache_beam.runners import create_runner from apache_beam.runners import pipeline_utils from apache_beam.transforms import ParDo from apache_beam.transforms import ptransform from apache_beam.transforms.display import DisplayData from apache_beam.transforms.display import HasDisplayData from apache_beam.transforms.resources import merge_resource_hints from apache_beam.transforms.resources import resource_hints_from_options from apache_beam.transforms.sideinputs import get_sideinput_index from apache_beam.typehints import TypeCheckError from apache_beam.typehints import typehints from apache_beam.utils import proto_utils from apache_beam.utils import subprocess_server from apache_beam.utils.annotations import deprecated from apache_beam.utils.interactive_utils import alter_label_if_ipython from apache_beam.utils.interactive_utils import is_in_ipython if TYPE_CHECKING: from types import TracebackType from apache_beam.runners.pipeline_context import PipelineContext from apache_beam.runners.runner import PipelineResult from apache_beam.transforms import environments __all__ = ['Pipeline', 'transform_annotations'] class Pipeline(HasDisplayData): """A pipeline object that manages a DAG of :class:`~apache_beam.transforms.ptransform.PTransform` s and their :class:`~apache_beam.pvalue.PValue` s. Conceptually the :class:`~apache_beam.transforms.ptransform.PTransform` s are the DAG's nodes and the :class:`~apache_beam.pvalue.PValue` s are the edges. All the transforms applied to the pipeline must have distinct full labels. If same transform instance needs to be applied then the right shift operator should be used to designate new names (e.g. ``input | "label" >> my_transform``). """ @classmethod def runner_implemented_transforms(cls) -> frozenset[str]: # This set should only contain transforms which are required to be # implemented by a runner. return frozenset([ common_urns.primitives.GROUP_BY_KEY.urn, common_urns.primitives.IMPULSE.urn, ]) def __init__( self, runner: Optional[Union[str, PipelineRunner]] = None, options: Optional[PipelineOptions] = None, argv: Optional[list[str]] = None, display_data: Optional[dict[str, Any]] = None): """Initialize a pipeline object. Args: runner (~apache_beam.runners.runner.PipelineRunner): An object of type :class:`~apache_beam.runners.runner.PipelineRunner` that will be used to execute the pipeline. For registered runners, the runner name can be specified, otherwise a runner object must be supplied. options (~apache_beam.options.pipeline_options.PipelineOptions): A configured :class:`~apache_beam.options.pipeline_options.PipelineOptions` object containing arguments that should be used for running the Beam job. argv (list[str]): a list of arguments (such as :data:`sys.argv`) to be used for building a :class:`~apache_beam.options.pipeline_options.PipelineOptions` object. This will only be used if argument **options** is :data:`None`. display_data (dict[str, Any]): a dictionary of static data associated with this pipeline that can be displayed when it runs. Raises: ValueError: if either the runner or options argument is not of the expected type. """ # Initializing logging configuration in case the user did not set it up. logging.basicConfig() if options is not None: if isinstance(options, PipelineOptions): # Make a deep copy of options since they could be overwritten in later # steps. However, the 'runner' object within 'options' is excluded from # the deep copy (it is shallow copied) due to potential issues with deep # copying specific runner instances, such as FlumeRunner. saved_runner = options.view_as(StandardOptions).runner options.view_as(StandardOptions).runner = None self._options = copy.deepcopy(options) self._options.view_as(StandardOptions).runner = saved_runner options.view_as(StandardOptions).runner = saved_runner else: raise ValueError( 'Parameter options, if specified, must be of type PipelineOptions. ' 'Received : %r' % options) elif argv is not None: if isinstance(argv, list): self._options = PipelineOptions(argv) else: raise ValueError( 'Parameter argv, if specified, must be a list. Received : %r' % argv) else: self._options = PipelineOptions([]) FileSystems.set_options(self._options) if runner is None: runner = self._options.view_as(StandardOptions).runner if runner is None: runner = StandardOptions.DEFAULT_RUNNER logging.info(( 'Missing pipeline option (runner). Executing pipeline ' 'using the default runner: %s.'), runner) if isinstance(runner, str): runner = create_runner(runner) elif not isinstance(runner, PipelineRunner): raise TypeError( 'Runner %s is not a PipelineRunner object or the ' 'name of a registered runner.' % runner) # Runner can override the default pickler to be used. if (self._options.view_as(SetupOptions).pickle_library == 'default' and runner.default_pickle_library_override()): logging.info( "Runner defaulting to pickling library: %s.", runner.default_pickle_library_override()) self._options.view_as( SetupOptions).pickle_library = runner.default_pickle_library_override( ) pickler.set_library(self._options.view_as(SetupOptions).pickle_library) # Validate pipeline options errors = PipelineOptionsValidator(self._options, runner).validate() if errors: raise ValueError( 'Pipeline has validations errors: \n' + '\n'.join(errors)) typecoders.registry.update_compatibility_version = self._options.view_as( StreamingOptions).update_compatibility_version # set default experiments for portable runners # (needs to occur prior to pipeline construction) if runner.is_fnapi_compatible(): experiments = (self._options.view_as(DebugOptions).experiments or []) if not 'beam_fn_api' in experiments: experiments.append('beam_fn_api') self._options.view_as(DebugOptions).experiments = experiments self.local_tempdir = tempfile.mkdtemp(prefix='beam-pipeline-temp') # Default runner to be used. self.runner = runner # Stack of transforms generated by nested apply() calls. The stack will # contain a root node as an enclosing (parent) node for top transforms. self.transforms_stack = [ AppliedPTransform(None, None, '', None, None, None) ] # Set of transform labels (full labels) applied to the pipeline. # If a transform is applied and the full label is already in the set # then the transform will have to be cloned with a new label. self.applied_labels: set[str] = set() # Hints supplied via pipeline options are considered the outermost hints. self._root_transform().resource_hints = resource_hints_from_options(options) # Create a ComponentIdMap for assigning IDs to components. Ensures that any # components that receive an ID during pipeline construction (for example in # ExternalTransform), will receive the same component ID when generating the # full pipeline proto. self.component_id_map = ComponentIdMap() # Records whether this pipeline contains any external transforms. self.contains_external_transforms = False self._display_data = display_data or {} self._error_handlers = [] self._annotations_stack = [{}] def display_data(self) -> dict[str, Any]: return self._display_data @property # type: ignore[misc] # decorated property not supported def options(self) -> PipelineOptions: return self._options @property def allow_unsafe_triggers(self) -> bool: return self._options.view_as(TypeOptions).allow_unsafe_triggers def _register_error_handler(self, error_handler): self._error_handlers.append(error_handler) def _current_transform(self) -> 'AppliedPTransform': """Returns the transform currently on the top of the stack.""" return self.transforms_stack[-1] @contextlib.contextmanager def transform_annotations(self, **annotations): """A context manager for attaching annotations to a set of transforms. All transforms applied while this context is active will have these annotations attached. This includes sub-transforms applied within composite transforms. """ self._annotations_stack.append({ **self._annotations_stack[-1], **encode_annotations(annotations) }) yield self._annotations_stack.pop() def _current_annotations(self): """Returns the set of annotations that should be used on apply.""" return {**_global_annotations_stack()[-1], **self._annotations_stack[-1]} def _root_transform(self) -> 'AppliedPTransform': """Returns the root transform of the transform stack.""" return self.transforms_stack[0] def _remove_labels_recursively( self, applied_transform: 'AppliedPTransform') -> None: for part in applied_transform.parts: if part.full_label in self.applied_labels: self.applied_labels.remove(part.full_label) self._remove_labels_recursively(part) def _replace(self, override: 'PTransformOverride') -> None: assert isinstance(override, PTransformOverride) # From original transform output --> replacement transform output output_map: dict[pvalue.PValue, pvalue.PValue] = {} output_replacements: dict[AppliedPTransform, list[tuple[pvalue.PValue, Optional[str]]]] = {} input_replacements: dict[AppliedPTransform, Mapping[str, Union[pvalue.PBegin, pvalue.PCollection]]] = {} side_input_replacements: dict[AppliedPTransform, list[pvalue.AsSideInput]] = {} class TransformUpdater(PipelineVisitor): # pylint: disable=used-before-assignment """"A visitor that replaces the matching PTransforms.""" def __init__(self, pipeline: Pipeline) -> None: self.pipeline = pipeline def _replace_if_needed( self, original_transform_node: AppliedPTransform) -> None: if override.matches(original_transform_node): assert isinstance(original_transform_node, AppliedPTransform) replacement_transform = ( override.get_replacement_transform_for_applied_ptransform( original_transform_node)) if replacement_transform is original_transform_node.transform: return replacement_transform.side_inputs = tuple( getattr(original_transform_node.transform, 'side_inputs', ())) replacement_transform_node = AppliedPTransform( original_transform_node.parent, replacement_transform, original_transform_node.full_label, original_transform_node.main_inputs, None, annotations=original_transform_node.annotations) # TODO(https://github.com/apache/beam/issues/21178): Merge rather # than override. replacement_transform_node.resource_hints = ( original_transform_node.resource_hints) # Transform execution could depend on order in which nodes are # considered. Hence we insert the replacement transform node to same # index as the original transform node. Note that this operation # removes the original transform node. if original_transform_node.parent: assert isinstance(original_transform_node.parent, AppliedPTransform) parent_parts = original_transform_node.parent.parts parent_parts[parent_parts.index(original_transform_node)] = ( replacement_transform_node) else: # Original transform has to be a root. roots = self.pipeline.transforms_stack[0].parts assert original_transform_node in roots roots[roots.index(original_transform_node)] = ( replacement_transform_node) inputs = override.get_replacement_inputs(original_transform_node) if len(inputs) > 1: transform_input = inputs elif len(inputs) == 1: transform_input = inputs[0] elif len(inputs) == 0: transform_input = pvalue.PBegin(self.pipeline) try: # We have to add the new AppliedTransform to the stack before # expand() and pop it out later to make sure that parts get added # correctly. self.pipeline.transforms_stack.append(replacement_transform_node) # Keeping the same label for the replaced node but recursively # removing labels of child transforms of original transform since # they will be replaced during the expand below. This is needed in # case the replacement contains children that have labels that # conflicts with labels of the children of the original. self.pipeline._remove_labels_recursively(original_transform_node) new_output = replacement_transform.expand(transform_input) assert isinstance( new_output, (dict, pvalue.PValue, pvalue.DoOutputsTuple)) if isinstance(new_output, pvalue.PValue): new_output.element_type = None self.pipeline._infer_result_type( replacement_transform, inputs, new_output) if isinstance(new_output, dict): for new_tag, new_pcoll in new_output.items(): replacement_transform_node.add_output(new_pcoll, new_tag) elif isinstance(new_output, pvalue.DoOutputsTuple): replacement_transform_node.add_output( new_output, new_output._main_tag) else: replacement_transform_node.add_output(new_output, new_output.tag) # Recording updated outputs. This cannot be done in the same # visitor since if we dynamically update output type here, we'll # run into errors when visiting child nodes. # # NOTE: When replacing multiple outputs, the replacement # PCollection tags must have a matching tag in the original # transform. if isinstance(new_output, pvalue.PValue): if not new_output.producer: new_output.producer = replacement_transform_node output_map[original_transform_node.outputs[new_output.tag]] = \ new_output elif isinstance(new_output, (pvalue.DoOutputsTuple, tuple)): for pcoll in new_output: if not pcoll.producer: pcoll.producer = replacement_transform_node output_map[original_transform_node.outputs[pcoll.tag]] = pcoll elif isinstance(new_output, dict): for tag, pcoll in new_output.items(): if not pcoll.producer: pcoll.producer = replacement_transform_node output_map[original_transform_node.outputs[tag]] = pcoll finally: self.pipeline.transforms_stack.pop() def enter_composite_transform( self, transform_node: AppliedPTransform) -> None: self._replace_if_needed(transform_node) def visit_transform(self, transform_node: AppliedPTransform) -> None: self._replace_if_needed(transform_node) self.visit(TransformUpdater(self)) # Ensure no type information is lost. for old, new in output_map.items(): if new.element_type == typehints.Any: # TODO(robertwb): Perhaps take the intersection? new.element_type = old.element_type # Adjusting inputs and outputs class InputOutputUpdater(PipelineVisitor): # pylint: disable=used-before-assignment """"A visitor that records input and output values to be replaced. Input and output values that should be updated are recorded in maps input_replacements and output_replacements respectively. We cannot update input and output values while visiting since that results in validation errors. """ def __init__(self, pipeline: Pipeline) -> None: self.pipeline = pipeline def enter_composite_transform( self, transform_node: AppliedPTransform) -> None: self.visit_transform(transform_node) def visit_transform(self, transform_node: AppliedPTransform) -> None: replace_output = False for tag in transform_node.outputs: if transform_node.outputs[tag] in output_map: replace_output = True break replace_input = False for input in transform_node.inputs: if input in output_map: replace_input = True break replace_side_inputs = False for side_input in transform_node.side_inputs: if side_input.pvalue in output_map: replace_side_inputs = True break if replace_output: output_replacements[transform_node] = [] for original, replacement in output_map.items(): for tag, output in transform_node.outputs.items(): if output == original: output_replacements[transform_node].append((tag, replacement)) if replace_input: new_inputs = { tag: input if not input in output_map else output_map[input] for (tag, input) in transform_node.main_inputs.items() } input_replacements[transform_node] = new_inputs if replace_side_inputs: new_side_inputs = [] for side_input in transform_node.side_inputs: if side_input.pvalue in output_map: side_input.pvalue = output_map[side_input.pvalue] new_side_inputs.append(side_input) else: new_side_inputs.append(side_input) side_input_replacements[transform_node] = new_side_inputs self.visit(InputOutputUpdater(self)) for transform, output_replacement in output_replacements.items(): for tag, output in output_replacement: transform.replace_output(output, tag=tag) for transform, input_replacement in input_replacements.items(): transform.replace_inputs(input_replacement) for transform, side_input_replacement in side_input_replacements.items(): transform.replace_side_inputs(side_input_replacement) def _check_replacement(self, override: 'PTransformOverride') -> None: class ReplacementValidator(PipelineVisitor): def visit_transform(self, transform_node: AppliedPTransform) -> None: if override.matches(transform_node): raise RuntimeError( 'Transform node %r was not replaced as expected.' % transform_node) self.visit(ReplacementValidator()) def replace_all(self, replacements: Iterable['PTransformOverride']) -> None: """ Dynamically replaces PTransforms in the currently populated hierarchy. Currently this only works for replacements where input and output types are exactly the same. TODO: Update this to also work for transform overrides where input and output types are different. Args: replacements (list[~apache_beam.pipeline.PTransformOverride]): a list of :class:`~apache_beam.pipeline.PTransformOverride` objects. """ for override in replacements: assert isinstance(override, PTransformOverride) self._replace(override) # Checking if the PTransforms have been successfully replaced. This will # result in a failure if a PTransform that was replaced in a given override # gets re-added in a subsequent override. This is not allowed and ordering # of PTransformOverride objects in 'replacements' is important. for override in replacements: self._check_replacement(override) def run(self, test_runner_api: Union[bool, str] = 'AUTO') -> 'PipelineResult': """Runs the pipeline. Returns whatever our runner returns after running.""" with scoped_pipeline_options(self._options): return self._run_internal(test_runner_api) def _run_internal( self, test_runner_api: Union[bool, str] = 'AUTO') -> 'PipelineResult': """Internal implementation of run(), called within scoped options.""" # All pipeline options are finalized at this point. # Call get_all_options to print warnings on invalid options. self.options.get_all_options( retain_unknown_options=True, display_warnings=True) for error_handler in self._error_handlers: error_handler.verify_closed() # Records whether this pipeline contains any cross-language transforms. self.contains_external_transforms = ( ExternalTransformFinder.contains_external_transforms(self)) try: if test_runner_api == 'AUTO': # Don't pay the cost of a round-trip if we're going to be going through # the FnApi anyway... is_fnapi_compatible = self.runner.is_fnapi_compatible() or ( # DirectRunner uses the Fn API for batch only self.runner.__class__.__name__ == 'SwitchingDirectRunner' and not self._options.view_as(StandardOptions).streaming) # Multi-language pipelines that contain external pipeline segments may # not be able to create a Python pipeline object graph. Hence following # runner API check should be skipped for such pipelines. # The InteractiveRunner relies on a constant pipeline reference, skip # it. test_runner_api = ( not is_fnapi_compatible and not self.contains_external_transforms and self.runner.__class__.__name__ != 'InteractiveRunner') # When possible, invoke a round trip through the runner API. if test_runner_api and self._verify_runner_api_compatible(): return Pipeline.from_runner_api( self.to_runner_api(use_fake_coders=True), self.runner, self._options).run(False) if (self._options.view_as(TypeOptions).runtime_type_check and self._options.view_as(TypeOptions).performance_runtime_type_check): raise RuntimeError( 'You cannot turn on runtime_type_check ' 'and performance_runtime_type_check simultaneously. ' 'Pick one or the other.') if self._options.view_as(TypeOptions).runtime_type_check: from apache_beam.typehints import typecheck self.visit(typecheck.TypeCheckVisitor()) if self._options.view_as(TypeOptions).performance_runtime_type_check: from apache_beam.typehints import typecheck self.visit(typecheck.PerformanceTypeCheckVisitor()) if self._options.view_as(SetupOptions).save_main_session: # If this option is chosen, verify we can pickle the main session early. tmpdir = tempfile.mkdtemp() try: pickler.dump_session(os.path.join(tmpdir, 'main_session.pickle')) finally: shutil.rmtree(tmpdir) return self.runner.run_pipeline(self, self._options) finally: if not is_in_ipython(): shutil.rmtree(self.local_tempdir, ignore_errors=True) # else interactive beam handles the cleanup. def __enter__(self) -> 'Pipeline': self._extra_context = contextlib.ExitStack() self._extra_context.enter_context( subprocess_server.JavaJarServer.beam_services( self._options.view_as(CrossLanguageOptions).beam_services)) self._extra_context.enter_context( subprocess_server.SubprocessServer.cache_subprocesses()) return self def __exit__( self, exc_type: Optional[type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional['TracebackType']) -> None: try: if not exc_type: self.result = self.run() if not self._options.view_as(StandardOptions).no_wait_until_finish: self.result.wait_until_finish() else: logging.info( 'Job execution continues without waiting for completion.' ' Use "wait_until_finish" in PipelineResult to block' ' until finished.') finally: self._extra_context.__exit__(exc_type, exc_val, exc_tb) def visit(self, visitor: 'PipelineVisitor') -> None: """Visits depth-first every node of a pipeline's DAG. Runner-internal implementation detail; no backwards-compatibility guarantees Args: visitor (~apache_beam.pipeline.PipelineVisitor): :class:`~apache_beam.pipeline.PipelineVisitor` object whose callbacks will be called for each node visited. See :class:`~apache_beam.pipeline.PipelineVisitor` comments. Raises: TypeError: if node is specified and is not a :class:`~apache_beam.pvalue.PValue`. ~apache_beam.error.PipelineError: if node is specified and does not belong to this pipeline instance. """ visited: set[pvalue.PValue] = set() self._root_transform().visit(visitor, self, visited) def apply( self, transform: ptransform.PTransform, pvalueish: Optional[pvalue.PValue] = None, label: Optional[str] = None) -> pvalue.PValue: """Applies a custom transform using the pvalueish specified. Args: transform (~apache_beam.transforms.ptransform.PTransform): the :class:`~apache_beam.transforms.ptransform.PTransform` to apply. pvalueish (~apache_beam.pvalue.PCollection): the input for the :class:`~apache_beam.transforms.ptransform.PTransform` (typically a :class:`~apache_beam.pvalue.PCollection`). label (str): label of the :class:`~apache_beam.transforms.ptransform.PTransform`. Raises: TypeError: if the transform object extracted from the argument list is not a :class:`~apache_beam.transforms.ptransform.PTransform`. RuntimeError: if the transform object was already applied to this pipeline and needs to be cloned in order to apply again. """ with scoped_pipeline_options(self._options): return self._apply_internal(transform, pvalueish, label) def _apply_internal( self, transform: ptransform.PTransform, pvalueish: Optional[pvalue.PValue] = None, label: Optional[str] = None) -> pvalue.PValue: """Internal implementation of apply(), called within scoped options.""" if isinstance(transform, ptransform._NamedPTransform): return self.apply( transform.transform, pvalueish, label or transform.label) if not label and isinstance(transform, ptransform._PTransformFnPTransform): # This must be set before label is inspected. transform.set_options(self._options) if not isinstance(transform, ptransform.PTransform): raise TypeError("Expected a PTransform object, got %s" % transform) if label: # Fix self.label as it is inspected by some PTransform operations # (e.g. to produce error messages for type hint violations). old_label, transform.label = transform.label, label try: return self.apply(transform, pvalueish) finally: transform.label = old_label # Attempts to alter the label of the transform to be applied only when it's # a top-level transform so that the cell number will not be prepended to # every child transform in a composite. if self._current_transform() is self._root_transform(): alter_label_if_ipython(transform, pvalueish) full_label = '/'.join( [self._current_transform().full_label, transform.label]).lstrip('/') if full_label in self.applied_labels: auto_unique_labels = self._options.view_as( StandardOptions).auto_unique_labels if auto_unique_labels: # If auto_unique_labels is set, we will append a unique suffix to the # label to make it unique. logging.warning( 'Using --auto_unique_labels could cause data loss when ' 'updating a pipeline or reloading the job state. ' 'This is not recommended for streaming jobs.') unique_label = self._generate_unique_label(transform) return self.apply(transform, pvalueish, unique_label) else: raise RuntimeError( 'A transform with label "%s" already exists in the pipeline. ' 'To apply a transform with a specified label, write ' 'pvalue | "label" >> transform or use the option ' '"auto_unique_labels" to automatically generate unique ' 'transform labels. Note "auto_unique_labels" ' 'could cause data loss when updating a pipeline or ' 'reloading the job state. This is not recommended for ' 'streaming jobs.' % full_label) self.applied_labels.add(full_label) if pvalueish is None: full_label = self._current_transform().full_label raise TypeCheckError( f'Transform "{full_label}" was applied to the output of ' f'an object of type None.') pvalueish, inputs = transform._extract_input_pvalues(pvalueish) try: if not isinstance(inputs, dict): inputs = {str(ix): input for (ix, input) in enumerate(inputs)} except TypeError: raise NotImplementedError( 'Unable to extract PValue inputs from %s; either %s does not accept ' 'inputs of this format, or it does not properly override ' '_extract_input_pvalues' % (pvalueish, transform)) for t, leaf_input in inputs.items(): if not isinstance(leaf_input, pvalue.PValue) or not isinstance(t, str): raise NotImplementedError( '%s does not properly override _extract_input_pvalues, ' 'returned %s from %s' % (transform, inputs, pvalueish)) current = AppliedPTransform( self._current_transform(), transform, full_label, inputs, None, annotations=self._current_annotations()) self._current_transform().add_part(current) try: self.transforms_stack.append(current) type_options = self._options.view_as(TypeOptions) if type_options.pipeline_type_check: transform.type_check_inputs(pvalueish) if isinstance(pvalueish, pvalue.PBegin) and isinstance(transform, ParDo): full_label = self._current_transform().full_label raise TypeCheckError( f"Transform '{full_label}' expects a PCollection as input. " "Got a PBegin/Pipeline instead.") self._assert_not_applying_PDone(pvalueish, transform) pvalueish_result = self.runner.apply(transform, pvalueish, self._options) if type_options is not None and type_options.pipeline_type_check: transform.type_check_outputs(pvalueish_result) for tag, result in ptransform.get_named_nested_pvalues(pvalueish_result): assert isinstance(result, (pvalue.PValue, pvalue.DoOutputsTuple)) # Make sure we set the producer only for a leaf node in the transform # DAG. This way we preserve the last transform of a composite transform # as being the real producer of the result. if result.producer is None: result.producer = current # TODO(BEAM-1833): Pass full tuples dict. self._infer_result_type(transform, tuple(inputs.values()), result) assert isinstance(result.producer.inputs, tuple) if isinstance(result, pvalue.DoOutputsTuple): for tag, pc in list(result._pcolls.items()): if tag not in current.outputs: current.add_output(pc, tag) continue # If there is already a tag with the same name, increase a counter for # the name. This can happen, for example, when a composite outputs a # list of PCollections where all the tags are None. base = tag counter = 0 while tag in current.outputs: counter += 1 tag = '%s_%d' % (base, counter) current.add_output(result, tag) if (type_options is not None and type_options.type_check_strictness == 'ALL_REQUIRED' and transform.get_type_hints().output_types is None): ptransform_name = '%s(%s)' % (transform.__class__.__name__, full_label) raise TypeCheckError( 'Pipeline type checking is enabled, however no ' 'output type-hint was found for the ' 'PTransform %s' % ptransform_name) finally: self.transforms_stack.pop() return pvalueish_result def _assert_not_applying_PDone( self, pvalueish: Optional[pvalue.PValue], transform: ptransform.PTransform): if isinstance(pvalueish, pvalue.PDone) and isinstance(transform, ParDo): # If the input is a PDone, we cannot apply a ParDo transform. full_label = self._current_transform().full_label producer_label = pvalueish.producer.full_label raise TypeCheckError( f'Transform "{full_label}" was applied to the output of ' f'"{producer_label}" but "{producer_label.split("/")[-1]}" ' 'produces no PCollections.') def _generate_unique_label(self, transform: str) -> str: """ Given a transform, generate a unique label for it based on current label. """ unique_suffix = uuid.uuid4().hex[:6] return '%s_%s' % (transform.label, unique_suffix) def _infer_result_type( self, transform: ptransform.PTransform, inputs: Sequence[Union[pvalue.PBegin, pvalue.PCollection]], result_pcollection: Union[pvalue.PValue, pvalue.DoOutputsTuple]) -> None: """Infer and set the output element type for a PCollection. This function determines the output types of transforms by combining: 1. Concrete input types from previous transforms 2. Type hints declared on the current transform 3. Type variable binding and substitution TYPE VARIABLE BINDING --------------------- Type variables (K, V, T, etc.) act as placeholders that get bound to concrete types through pattern matching. This requires both an input pattern and an output template: Input Pattern (from .with_input_types()): Defines where in the input to find each type variable Example: Tuple[K, V] means "K is the first element, V is the second" Output Template (from .with_output_types()): Defines how to use the bound variables in the output Example: Tuple[V, K] means "swap the positions" CONCRETE TYPES VS TYPE VARIABLES --------------------------------- The system handles these differently: Concrete Types (e.g., str, int, Tuple[str, int]): - Used as-is without any binding - Do not fall back to Any - Example: .with_output_types(Tuple[str, int]) → Tuple[str, int] Type Variables (e.g., K, V, T): - Must be bound through pattern matching - Require .with_input_types() to provide the pattern - Fall back to Any if not bound - Example without pattern: Tuple[K, V] → Tuple[Any, Any] - Example with pattern: Tuple[K, V] → Tuple[str, int] BINDING ALGORITHM ----------------- 1. Match: Compare input pattern to concrete input Pattern: Tuple[K, V] Concrete: Tuple[str, int] Result: {K: str, V: int} ← Bindings created 2. Substitute: Apply bindings to output template Template: Tuple[V, K] ← Note: swapped! Bindings: {K: str, V: int} Result: Tuple[int, str] ← Swapped concrete types Each transform operates in its own type inference scope. Type variables declared in a parent composite transform do NOT automatically propagate to child transforms. Parent scope (composite): @with_input_types(Tuple[K, V]) ← K, V defined here class MyComposite(PTransform): def expand(self, pcoll): # Child scope - parent's K, V are NOT available return pcoll | ChildTransform() Type variables that remain unbound after inference fall back to Any: EXAMPLES -------- Example 1: Concrete types (no variables) Input: Tuple[str, int] Transform: .with_output_types(Tuple[str, int]) Output: Tuple[str, int] ← Used as-is Example 2: Type variables with pattern (correct) Input: Tuple[str, int] Transform: .with_input_types(Tuple[K, V]) .with_output_types(Tuple[V, K]) Binding: {K: str, V: int} Output: Tuple[int, str] ← Swapped! Example 3: Type variables without pattern (falls back to Any) Input: Tuple[str, int] Transform: .with_output_types(Tuple[K, V]) ← No input pattern! Binding: None (can't match) Output: Tuple[Any, Any] ← Fallback Example 4: Mixed concrete and variables Input: Tuple[str, int] Transform: .with_input_types(Tuple[str, V]) .with_output_types(Tuple[str, V]) Binding: {V: int} ← Only V needs binding Output: Tuple[str, int] ← str passed through, V bound to int Args: transform: The PTransform being applied inputs: Input PCollections (provides concrete types) result_pcollection: Output PCollection to set type on """ # TODO(robertwb): Multi-input inference. type_options = self._options.view_as(TypeOptions) if type_options is None or not type_options.pipeline_type_check: return if (isinstance(result_pcollection, pvalue.PCollection) and (not result_pcollection.element_type # TODO(robertwb): Ideally we'd do intersection here. or result_pcollection.element_type == typehints.Any)): # {Single, multi}-input, single-output inference. input_element_types_tuple = tuple(i.element_type for i in inputs) input_element_type = ( input_element_types_tuple[0] if len(input_element_types_tuple) == 1 else typehints.Union[input_element_types_tuple]) type_hints = transform.get_type_hints() if not result_pcollection.tag: declared_output_type = type_hints.simple_output_type(transform.label) else: declared_output_type = type_hints.tagged_output_types().get( result_pcollection.tag, typehints.Any) if declared_output_type: input_types = type_hints.input_types if input_types and input_types[0]: declared_input_type = input_types[0][0] result_element_type = typehints.bind_type_variables( declared_output_type, typehints.match_type_variables( declared_input_type, input_element_type)) else: result_element_type = declared_output_type else: result_element_type = transform.infer_output_type(input_element_type) # Any remaining type variables have no bindings higher than this scope. result_pcollection.element_type = typehints.bind_type_variables( result_element_type, {'*': typehints.Any}) elif isinstance(result_pcollection, pvalue.DoOutputsTuple): # {Single, multi}-input, multi-output inference. # TODO(https://github.com/apache/beam/issues/18957): Add support for # tagged type hints. # https://github.com/apache/beam/pull/9810#discussion_r338765251 for pcoll in result_pcollection: if pcoll.element_type is None: pcoll.element_type = typehints.Any def __reduce__(self) -> tuple[type, tuple[str, ...]]: # Some transforms contain a reference to their enclosing pipeline, # which in turn reference all other transforms (resulting in quadratic # time/space to pickle each transform individually). As we don't # require pickled pipelines to be executable, break the chain here. return str, ('Pickled pipeline stub.', ) def _verify_runner_api_compatible(self) -> bool: if self._options.view_as(TypeOptions).runtime_type_check: # This option is incompatible with the runner API as it requires # the runner to inspect non-serialized hints on the transform # itself. return False class Visitor(PipelineVisitor): # pylint: disable=used-before-assignment ok = True # Really a nonlocal. def enter_composite_transform( self, transform_node: AppliedPTransform) -> None: pass def visit_transform(self, transform_node: AppliedPTransform) -> None: try: # Transforms must be picklable. pickler.loads( pickler.dumps(transform_node.transform, enable_trace=False), enable_trace=False) except Exception: Visitor.ok = False def visit_value(self, value: pvalue.PValue, _: AppliedPTransform) -> None: if isinstance(value, pvalue.PDone): Visitor.ok = False self.visit(Visitor()) return Visitor.ok def to_runner_api( self, return_context: bool = False, context: Optional['PipelineContext'] = None, use_fake_coders: bool = False, default_environment: Optional['environments.Environment'] = None ) -> beam_runner_api_pb2.Pipeline: """For internal use only; no backwards-compatibility guarantees.""" from apache_beam.runners import pipeline_context if context is None: context = pipeline_context.PipelineContext( use_fake_coders=use_fake_coders, component_id_map=self.component_id_map, default_environment=default_environment) elif default_environment is not None: raise ValueError( 'Only one of context or default_environment may be specified.') # The FlumeRunner is the only runner setting this option. Use getattr # because other runners do not have this option. context.enable_best_effort_deterministic_pickling = getattr( self.runner, 'enable_best_effort_deterministic_pickling', False) # The RunnerAPI spec requires certain transforms and side-inputs to have KV # inputs (and corresponding outputs). # Currently we only upgrade to KV pairs. If there is a need for more # general shapes, potential conflicts will have to be resolved. # We also only handle single-input, and (for fixing the output) single # output, which is sufficient. # Also marks such values as requiring deterministic key coders. deterministic_key_coders = not self._options.view_as( TypeOptions).allow_non_deterministic_key_coders class ForceKvInputTypes(PipelineVisitor): def enter_composite_transform( self, transform_node: AppliedPTransform) -> None: self.visit_transform(transform_node) def visit_transform(self, transform_node: AppliedPTransform) -> None: if not transform_node.transform: return if hasattr( transform_node.transform, 'runner_api_requires_keyed_input' ) and transform_node.transform.runner_api_requires_keyed_input(): pcoll = transform_node.inputs[0] pcoll.element_type = typehints.coerce_to_kv_type( pcoll.element_type, transform_node.full_label) pcoll.requires_deterministic_key_coder = ( deterministic_key_coders and transform_node.full_label) if len(transform_node.outputs) == 1: # The runner often has expectations about the output types as well. output, = transform_node.outputs.values() if not output.element_type: output.element_type = transform_node.transform.infer_output_type( pcoll.element_type) if (isinstance(output.element_type, typehints.TupleHint.TupleConstraint) and len(output.element_type.tuple_types) == 2 and pcoll.element_type.tuple_types[0] == output.element_type.tuple_types[0]): output.requires_deterministic_key_coder = ( deterministic_key_coders and transform_node.full_label) for side_input in getattr(transform_node.transform, 'side_inputs', []): if side_input.requires_keyed_input(): side_input.pvalue.element_type = typehints.coerce_to_kv_type( side_input.pvalue.element_type, transform_node.full_label, side_input_producer=side_input.pvalue.producer.full_label) side_input.pvalue.requires_deterministic_key_coder = ( deterministic_key_coders and transform_node.full_label) self.visit(ForceKvInputTypes()) # Mutates context; placing inline would force dependence on # argument evaluation order. root_transform_id = context.transforms.get_id(self._root_transform()) proto = beam_runner_api_pb2.Pipeline( root_transform_ids=[root_transform_id], components=context.to_runner_api(), requirements=context.requirements(), display_data=DisplayData('', self._display_data).to_proto()) proto.components.transforms[root_transform_id].unique_name = ( root_transform_id) self.merge_compatible_environments(proto) if return_context: return proto, context # type: ignore # too complicated for now else: return proto @staticmethod def merge_compatible_environments(proto): """Tries to minimize the number of distinct environments by merging those that are compatible (currently defined as identical). Mutates proto as contexts may have references to proto.components. """ pipeline_utils.merge_common_environments(proto, inplace=True) @staticmethod def from_runner_api( proto: beam_runner_api_pb2.Pipeline, runner: PipelineRunner, options: PipelineOptions, return_context: bool = False, ) -> 'Pipeline': """For internal use only; no backwards-compatibility guarantees.""" p = Pipeline( runner=runner, options=options, display_data={ str(ix): d for ix, d in enumerate(proto.display_data) }) from apache_beam.runners import pipeline_context context = pipeline_context.PipelineContext( proto.components, requirements=proto.requirements) if proto.root_transform_ids: root_transform_id, = proto.root_transform_ids p.transforms_stack = [context.transforms.get_by_id(root_transform_id)] else: p.transforms_stack = [AppliedPTransform(None, None, '', None, None, None)] # TODO(robertwb): These are only needed to continue construction. Omit? p.applied_labels = { t.unique_name for t in proto.components.transforms.values() } for id in proto.components.pcollections: pcollection = context.pcollections.get_by_id(id) pcollection.pipeline = p if not pcollection.producer: raise ValueError('No producer for %s' % id) # Inject PBegin input where necessary. from apache_beam.io.iobase import Read from apache_beam.transforms.core import Create has_pbegin = [Read, Create] for id in proto.components.transforms: transform = context.transforms.get_by_id(id) if not transform.inputs and transform.transform.__class__ in has_pbegin: transform.main_inputs = {'None': pvalue.PBegin(p)} if return_context: return p, context # type: ignore # too complicated for now else: return p class PipelineVisitor(object): """For internal use only; no backwards-compatibility guarantees. Visitor pattern class used to traverse a DAG of transforms (used internally by Pipeline for bookkeeping purposes). """ def visit_value( self, value: pvalue.PValue, producer_node: 'AppliedPTransform') -> None: """Callback for visiting a PValue in the pipeline DAG. Args: value: PValue visited (typically a PCollection instance). producer_node: AppliedPTransform object whose transform produced the pvalue. """ pass def visit_transform(self, transform_node: 'AppliedPTransform') -> None: """Callback for visiting a transform leaf node in the pipeline DAG.""" pass def enter_composite_transform( self, transform_node: 'AppliedPTransform') -> None: """Callback for entering traversal of a composite transform node.""" pass def leave_composite_transform( self, transform_node: 'AppliedPTransform') -> None: """Callback for leaving traversal of a composite transform node.""" pass class ExternalTransformFinder(PipelineVisitor): """Looks for any external transforms in the pipeline and if found records it. """ def __init__(self): self._contains_external_transforms = False @staticmethod def contains_external_transforms(pipeline): visitor = ExternalTransformFinder() pipeline.visit(visitor) return visitor._contains_external_transforms def _perform_exernal_transform_test(self, transform): if not transform: return from apache_beam.transforms import ExternalTransform if isinstance(transform, ExternalTransform): self._contains_external_transforms = True def visit_transform(self, transform_node: 'AppliedPTransform') -> None: self._perform_exernal_transform_test(transform_node.transform) def enter_composite_transform( self, transform_node: 'AppliedPTransform') -> None: # Python SDK object graph may represent an external transform that is a leaf # of the pipeline graph as a composite without sub-transforms. # Note that this visitor is just used to identify pipelines with external # transforms. A Runner API pipeline proto generated from the Pipeline object # will include external sub-transform. self._perform_exernal_transform_test(transform_node.transform) class AppliedPTransform(object): """For internal use only; no backwards-compatibility guarantees. A transform node representing an instance of applying a PTransform (used internally by Pipeline for bookeeping purposes). """ def __init__( self, parent: Optional['AppliedPTransform'], transform: Optional[ptransform.PTransform], full_label: str, main_inputs: Optional[Mapping[str, Union[pvalue.PBegin, pvalue.PCollection]]], environment_id: Optional[str], annotations: Optional[dict[str, bytes]], ) -> None: self.parent = parent self.transform = transform # Note that we want the PipelineVisitor classes to use the full_label, # inputs, side_inputs, and outputs fields from this instance instead of the # ones of the PTransform instance associated with it. Doing this permits # reusing PTransform instances in different contexts (apply() calls) without # any interference. This is particularly useful for composite transforms. self.full_label = full_label self.main_inputs = dict(main_inputs or {}) self.side_inputs = ( tuple() if transform is None else getattr( transform, 'side_inputs', tuple())) self.outputs: dict[Union[str, int, None], pvalue.PValue] = {} self.parts: list[AppliedPTransform] = [] self.environment_id: Optional[ str] = environment_id if environment_id else None # We may need to merge the hints with environment-provided hints here # once environment is a first-class citizen in Beam graph and we have # access to actual environment, not just an id. self.resource_hints: dict[str, bytes] = dict( transform.get_resource_hints()) if transform and hasattr( transform, 'get_resource_hints') else {} if transform and hasattr(transform, 'annotations'): annotations = { **(annotations or {}), **encode_annotations(transform.annotations()) } self.annotations = annotations self.display_data = {} @property def inputs(self): return tuple(self.main_inputs.values()) def __repr__(self) -> str: return "%s(%s, %s)" % ( self.__class__.__name__, self.full_label, type(self.transform).__name__) def replace_output( self, output: Union[pvalue.PValue, pvalue.DoOutputsTuple], tag: Union[str, int, None] = None) -> None: """Replaces the output defined by the given tag with the given output. Args: output: replacement output tag: tag of the output to be replaced. """ if isinstance(output, pvalue.DoOutputsTuple): self.replace_output(output[output._main_tag]) elif isinstance(output, pvalue.PValue): self.outputs[tag] = output elif isinstance(output, dict): for output_tag, out in output.items(): self.outputs[output_tag] = out else: raise TypeError("Unexpected output type: %s" % output) # Importing locally to prevent circular dependency issues. from apache_beam.transforms import external if isinstance(self.transform, external.ExternalTransform): self.transform.replace_named_outputs(self.named_outputs()) def replace_inputs(self, main_inputs): self.main_inputs = main_inputs # Importing locally to prevent circular dependency issues. from apache_beam.transforms import external if isinstance(self.transform, external.ExternalTransform): self.transform.replace_named_inputs(self.named_inputs()) def replace_side_inputs(self, side_inputs): self.side_inputs = side_inputs # Importing locally to prevent circular dependency issues. from apache_beam.transforms import external if isinstance(self.transform, external.ExternalTransform): self.transform.replace_named_inputs(self.named_inputs()) def add_output( self, output: Union[pvalue.DoOutputsTuple, pvalue.PValue], tag: Union[str, int, None]) -> None: if isinstance(output, pvalue.DoOutputsTuple): self.add_output(output[tag], tag) elif isinstance(output, pvalue.PValue): assert tag not in self.outputs self.outputs[tag] = output else: raise TypeError("Unexpected output type: %s" % output) def add_part(self, part: 'AppliedPTransform') -> None: assert isinstance(part, AppliedPTransform) part._merge_outer_resource_hints() self.parts.append(part) def is_composite(self) -> bool: """Returns whether this is a composite transform. A composite transform has parts (inner transforms) or isn't the producer for any of its outputs. (An example of a transform that is not a producer is one that returns its inputs instead.) """ return bool(self.parts) or all( pval.producer is not self for pval in self.outputs.values()) def visit( self, visitor: PipelineVisitor, pipeline: Pipeline, visited: set[pvalue.PValue]) -> None: """Visits all nodes reachable from the current node.""" for in_pval in self.inputs: if in_pval not in visited and not isinstance(in_pval, pvalue.PBegin): if in_pval.producer is not None: in_pval.producer.visit(visitor, pipeline, visited) # The value should be visited now since we visit outputs too. assert in_pval in visited, in_pval # Visit side inputs. for side_input in self.side_inputs: if isinstance(side_input, pvalue.AsSideInput) \ and side_input.pvalue not in visited: pval = side_input.pvalue # Unpack marker-object-wrapped pvalue. if pval.producer is not None: pval.producer.visit(visitor, pipeline, visited) # The value should be visited now since we visit outputs too. assert pval in visited # TODO(silviuc): Is there a way to signal that we are visiting a side # value? The issue is that the same PValue can be reachable through # multiple paths and therefore it is not guaranteed that the value # will be visited as a side value. # Visit a composite or primitive transform. if self.is_composite(): visitor.enter_composite_transform(self) for part in self.parts: part.visit(visitor, pipeline, visited) visitor.leave_composite_transform(self) else: visitor.visit_transform(self) # Visit the outputs (one or more). It is essential to mark as visited the # tagged PCollections of the DoOutputsTuple object. A tagged PCollection is # connected directly with its producer (a multi-output ParDo), but the # output of such a transform is the containing DoOutputsTuple, not the # PCollection inside it. Without the code below a tagged PCollection will # not be marked as visited while visiting its producer. for out_pval in self.outputs.values(): if isinstance(out_pval, pvalue.DoOutputsTuple): pvals = (v for v in out_pval) else: pvals = (out_pval, ) for v in pvals: if v not in visited: visited.add(v) visitor.visit_value(v, self) def named_inputs(self) -> dict[str, pvalue.PValue]: if self.transform is None: assert not self.main_inputs and not self.side_inputs return {} else: if hasattr(self.transform, '_named_inputs'): named_inputs = self.transform._named_inputs( self.main_inputs, self.side_inputs) else: named_inputs = {} if not self.parts: for name, pc_out in self.outputs.items(): if pc_out.producer is not self and pc_out not in named_inputs.values( ): named_inputs[f'__implicit_input_{name}'] = pc_out return named_inputs def named_outputs(self) -> dict[str, pvalue.PCollection]: if self.transform is None: assert not self.outputs return {} else: if hasattr(self.transform, '_named_outputs'): return self.transform._named_outputs(self.outputs) else: return {} def to_runner_api( self, context: 'PipelineContext') -> beam_runner_api_pb2.PTransform: # External transforms require more splicing than just setting the spec. from apache_beam.transforms import external if isinstance(self.transform, external.ExternalTransform): # TODO(https://github.com/apache/beam/issues/18371): Support resource # hints in XLang transforms. In particular, make sure hints on composites # are properly propagated. return self.transform.to_runner_api_transform(context, self.full_label) def transform_to_runner_api( transform: Optional[ptransform.PTransform], context: 'PipelineContext' ) -> Optional[beam_runner_api_pb2.FunctionSpec]: if transform is None: return None else: # We only populate inputs information to ParDo in order to expose # key_coder and window_coder to stateful DoFn. if isinstance(transform, ParDo): return transform.to_runner_api( context, has_parts=bool(self.parts), named_inputs=self.named_inputs()) elif hasattr(transform, 'to_runner_api'): return transform.to_runner_api(context, has_parts=bool(self.parts)) return None # Iterate over inputs and outputs by sorted key order, so that ids are # consistently generated for multiple runs of the same pipeline. try: transform_spec = transform_to_runner_api(self.transform, context) except Exception as exn: raise RuntimeError(f'Unable to translate {self.full_label}') from exn environment_id = self.environment_id transform_urn = transform_spec.urn if transform_spec else None if (not environment_id and (transform_urn not in Pipeline.runner_implemented_transforms())): environment_id = context.get_environment_id_for_resource_hints( self.resource_hints) if self.transform is not None: display_data = DisplayData.create_from( self.transform, extra_items=self.display_data) else: display_data = None return beam_runner_api_pb2.PTransform( unique_name=self.full_label, spec=transform_spec, subtransforms=[ context.transforms.get_id(part, label=part.full_label) for part in self.parts ], inputs={ tag: context.pcollections.get_id(pc) for tag, pc in sorted(self.named_inputs().items()) }, outputs={ tag: context.pcollections.get_id(out) for tag, out in sorted(self.named_outputs().items()) }, environment_id=environment_id, annotations=self.annotations, # TODO(https://github.com/apache/beam/issues/18012): Add display_data. display_data=display_data.to_proto() if display_data else None) @staticmethod def from_runner_api( proto: beam_runner_api_pb2.PTransform, context: 'PipelineContext') -> 'AppliedPTransform': if common_urns.primitives.PAR_DO.urn == proto.spec.urn: # Preserving side input tags. pardo_payload = ( proto_utils.parse_Bytes( proto.spec.payload, beam_runner_api_pb2.ParDoPayload)) side_input_tags = list(pardo_payload.side_inputs.keys()) else: pardo_payload = None side_input_tags = [] main_inputs = { tag: context.pcollections.get_by_id(id) for (tag, id) in proto.inputs.items() if tag not in side_input_tags } transform = ptransform.PTransform.from_runner_api(proto, context) if transform and proto.environment_id: resource_hints = context.environments.get_by_id( proto.environment_id).resource_hints() if resource_hints: transform._resource_hints = dict(resource_hints) # Ordering is important here. # TODO(https://github.com/apache/beam/issues/20136): use key, value pairs # instead of depending on tags with index as a suffix. indexed_side_inputs = [ (get_sideinput_index(tag), context.pcollections.get_by_id(id)) for tag, id in proto.inputs.items() if tag in side_input_tags ] side_inputs = [si for _, si in sorted(indexed_side_inputs)] result = AppliedPTransform( parent=None, transform=transform, full_label=proto.unique_name, main_inputs=main_inputs, environment_id=None, annotations=proto.annotations) if result.transform and hasattr( result.transform, 'side_inputs') and result.transform.side_inputs: for si, pcoll in zip(result.transform.side_inputs, side_inputs): si.pvalue = pcoll result.side_inputs = tuple(result.transform.side_inputs) result.parts = [] for transform_id in proto.subtransforms: part = context.transforms.get_by_id(transform_id) part.parent = result result.add_part(part) result.outputs = { None if tag == 'None' else tag: context.pcollections.get_by_id(id) for tag, id in proto.outputs.items() } # This annotation is expected by some runners. if proto.spec.urn == common_urns.primitives.PAR_DO.urn: result.transform.output_tags = set(proto.outputs.keys()).difference( {'None'}) if not result.parts: for tag, pcoll_id in proto.outputs.items(): if pcoll_id not in proto.inputs.values(): pc = context.pcollections.get_by_id(pcoll_id) pc.producer = result pc.tag = None if tag == 'None' else tag return result def _merge_outer_resource_hints(self): if (self.parent is not None and self.parent.resource_hints): self.resource_hints = merge_resource_hints( outer_hints=self.parent.resource_hints, inner_hints=self.resource_hints) if self.resource_hints: for part in self.parts: part._merge_outer_resource_hints() def encode_annotations(annotations: Optional[dict[str, Any]]): """Encodes non-byte annotation values as bytes.""" if not annotations: return {} def annotation_to_bytes(key, a: Any) -> bytes: if isinstance(a, bytes): return a elif isinstance(a, str): return a.encode('ascii') elif isinstance(a, message.Message): return a.SerializeToString() else: raise TypeError( 'Unknown annotation type %r (type %s) for %s' % (a, type(a), key)) return {key: annotation_to_bytes(key, a) for (key, a) in annotations.items()} _global_annotations_stack_data = threading.local() def _global_annotations_stack(): try: return _global_annotations_stack_data.stack except AttributeError: _global_annotations_stack_data.stack = [{}] return _global_annotations_stack_data.stack @contextlib.contextmanager def transform_annotations(**annotations): """A context manager for attaching annotations to a set of transforms. All transforms applied while this context is active will have these annotations attached. This includes sub-transforms applied within composite transforms. """ cur_stack = _global_annotations_stack() cur_stack.append({**cur_stack[-1], **encode_annotations(annotations)}) yield cur_stack.pop() class PTransformOverride(metaclass=abc.ABCMeta): """For internal use only; no backwards-compatibility guarantees. Gives a matcher and replacements for matching PTransforms. TODO: Update this to support cases where input and/our output types are different. """ @abc.abstractmethod def matches(self, applied_ptransform: AppliedPTransform) -> bool: """Determines whether the given AppliedPTransform matches. Note that the matching will happen *after* Runner API proto translation. If matching is done via type checks, to/from_runner_api[_parameter] methods must be implemented to preserve the type (and other data) through proto serialization. Consider URN-based translation instead. Args: applied_ptransform: AppliedPTransform to be matched. Returns: a bool indicating whether the given AppliedPTransform is a match. """ raise NotImplementedError def get_replacement_transform_for_applied_ptransform( self, applied_ptransform: AppliedPTransform) -> ptransform.PTransform: """Provides a runner specific override for a given `AppliedPTransform`. Args: applied_ptransform: `AppliedPTransform` containing the `PTransform` to be replaced. Returns: A `PTransform` that will be the replacement for the `PTransform` inside the `AppliedPTransform` given as an argument. """ # Returns a PTransformReplacement return self.get_replacement_transform(applied_ptransform.transform) @deprecated( since='2.24', current='get_replacement_transform_for_applied_ptransform') def get_replacement_transform( self, ptransform: Optional[ptransform.PTransform]) -> ptransform.PTransform: """Provides a runner specific override for a given PTransform. Args: ptransform: PTransform to be replaced. Returns: A PTransform that will be the replacement for the PTransform given as an argument. """ # Returns a PTransformReplacement raise NotImplementedError def get_replacement_inputs( self, applied_ptransform: AppliedPTransform) -> Iterable[pvalue.PValue]: """Provides inputs that will be passed to the replacement PTransform. Args: applied_ptransform: Original AppliedPTransform containing the PTransform to be replaced. Returns: An iterable of PValues that will be passed to the expand() method of the replacement PTransform. """ return tuple(applied_ptransform.inputs) + tuple( side_input.pvalue for side_input in applied_ptransform.side_inputs) class ComponentIdMap(object): """A utility for assigning unique component ids to Beam components. Component ID assignments are only guaranteed to be unique and consistent within the scope of a ComponentIdMap instance. """ def __init__(self, namespace="ref"): self.namespace = namespace self._counters: dict[type, int] = defaultdict(lambda: 0) self._obj_to_id: dict[Any, str] = {} def get_or_assign(self, obj=None, obj_type=None, label=None): if obj not in self._obj_to_id: self._obj_to_id[obj] = self._unique_ref(obj, obj_type, label) return self._obj_to_id[obj] def _normalize(self, str_value): str_value = unicodedata.normalize('NFC', str_value) return re.sub(r'[^a-zA-Z0-9-_]+', '-', str_value) def _unique_ref(self, obj=None, obj_type=None, label=None): # Normalize, trim, and uniqify. prefix = self._normalize( '%s_%s_%s' % (self.namespace, obj_type.__name__, label or type(obj).__name__))[0:100] if isinstance(obj, typecoders.coders.Coder) and obj.version_tag(): prefix = "%s_%s" % (prefix, obj.version_tag()) self._counters[obj_type] += 1 return '%s_%d' % (prefix, self._counters[obj_type]) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/pipeline_test.py0000644000175100017510000017103215206113032021077 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the Pipeline class.""" # pytype: skip-file import copy import platform import unittest import uuid import mock import pytest import apache_beam as beam from apache_beam import coders from apache_beam import typehints from apache_beam.coders import BytesCoder from apache_beam.io import Read from apache_beam.io.iobase import SourceBase from apache_beam.options.pipeline_options import PortableOptions from apache_beam.pipeline import Pipeline from apache_beam.pipeline import PipelineOptions from apache_beam.pipeline import PipelineVisitor from apache_beam.pipeline import PTransformOverride from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.pvalue import AsSingleton from apache_beam.pvalue import TaggedOutput from apache_beam.runners.runner import PipelineRunner from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import CombineGlobally from apache_beam.transforms import Create from apache_beam.transforms import DoFn from apache_beam.transforms import FlatMap from apache_beam.transforms import Map from apache_beam.transforms import ParDo from apache_beam.transforms import PTransform from apache_beam.transforms import WindowInto from apache_beam.transforms.display import DisplayDataItem from apache_beam.transforms.environments import ProcessEnvironment from apache_beam.transforms.resources import ResourceHint from apache_beam.transforms.userstate import BagStateSpec from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import IntervalWindow from apache_beam.transforms.window import SlidingWindows from apache_beam.transforms.window import TimestampedValue from apache_beam.typehints import TypeCheckError from apache_beam.utils import windowed_value from apache_beam.utils.timestamp import MIN_TIMESTAMP class FakeUnboundedSource(SourceBase): """Fake unbounded source. Does not work at runtime""" def is_bounded(self): return False class DoubleParDo(beam.PTransform): def expand(self, input): return input | 'Inner' >> beam.Map(lambda a: a * 2) def to_runner_api_parameter(self, context): return self.to_runner_api_pickled(context) class TripleParDo(beam.PTransform): def expand(self, input): # Keeping labels the same intentionally to make sure that there is no label # conflict due to replacement. return input | 'Inner' >> beam.Map(lambda a: a * 3) class ToStringParDo(beam.PTransform): def expand(self, input): # We use copy.copy() here to make sure the typehint mechanism doesn't # automatically infer that the output type is str. return input | 'Inner' >> beam.Map(lambda a: copy.copy(str(a))) class FlattenAndDouble(beam.PTransform): def expand(self, pcolls): return pcolls | beam.Flatten() | 'Double' >> DoubleParDo() class FlattenAndTriple(beam.PTransform): def expand(self, pcolls): return pcolls | beam.Flatten() | 'Triple' >> TripleParDo() class AddWithProductDoFn(beam.DoFn): def process(self, input, a, b): yield input + a * b class AddThenMultiplyDoFn(beam.DoFn): def process(self, input, a, b): yield (input + a) * b class AddThenMultiply(beam.PTransform): def expand(self, pvalues): return pvalues[0] | beam.ParDo( AddThenMultiplyDoFn(), AsSingleton(pvalues[1]), AsSingleton(pvalues[2])) class PipelineTest(unittest.TestCase): @staticmethod def custom_callable(pcoll): return pcoll | '+1' >> FlatMap(lambda x: [x + 1]) # Some of these tests designate a runner by name, others supply a runner. # This variation is just to verify that both means of runner specification # work and is not related to other aspects of the tests. class CustomTransform(PTransform): def expand(self, pcoll): return pcoll | '+1' >> FlatMap(lambda x: [x + 1]) class Visitor(PipelineVisitor): def __init__(self, visited): self.visited = visited self.enter_composite = [] self.leave_composite = [] def visit_value(self, value, _): self.visited.append(value) def enter_composite_transform(self, transform_node): self.enter_composite.append(transform_node) def leave_composite_transform(self, transform_node): self.leave_composite.append(transform_node) def test_create(self): with TestPipeline() as pipeline: pcoll = pipeline | 'label1' >> Create([1, 2, 3]) assert_that(pcoll, equal_to([1, 2, 3])) # Test if initial value is an iterator object. pcoll2 = pipeline | 'label2' >> Create(iter((4, 5, 6))) pcoll3 = pcoll2 | 'do' >> FlatMap(lambda x: [x + 10]) assert_that(pcoll3, equal_to([14, 15, 16]), label='pcoll3') def test_unexpected_PDone_errmsg(self): """ Test that a nice error message is raised if a transform that returns None (i.e. produces no PCollection) is used as input to a PTransform. """ class DoNothingTransform(PTransform): def expand(self, pcoll): return None class ParentTransform(PTransform): def expand(self, pcoll): return pcoll | DoNothingTransform() with pytest.raises(TypeCheckError, match=r".*applied to the output"): with TestPipeline() as pipeline: _ = pipeline | ParentTransform() | beam.Map(lambda x: x + 1) @mock.patch('logging.info') @pytest.mark.uses_dill def test_runner_overrides_default_pickler(self, mock_info): pytest.importorskip("dill") with mock.patch.object(PipelineRunner, 'default_pickle_library_override') as mock_fn: mock_fn.return_value = 'dill' with TestPipeline() as pipeline: pcoll = pipeline | 'label1' >> Create([1, 2, 3]) assert_that(pcoll, equal_to([1, 2, 3])) from apache_beam.internal import dill_pickler from apache_beam.internal import pickler self.assertIs(pickler.desired_pickle_lib, dill_pickler) mock_info.assert_any_call( 'Runner defaulting to pickling library: %s.', 'dill') def test_flatmap_builtin(self): with TestPipeline() as pipeline: pcoll = pipeline | 'label1' >> Create([1, 2, 3]) assert_that(pcoll, equal_to([1, 2, 3])) pcoll2 = pcoll | 'do' >> FlatMap(lambda x: [x + 10]) assert_that(pcoll2, equal_to([11, 12, 13]), label='pcoll2') pcoll3 = pcoll2 | 'm1' >> Map(lambda x: [x, 12]) assert_that( pcoll3, equal_to([[11, 12], [12, 12], [13, 12]]), label='pcoll3') pcoll4 = pcoll3 | 'do2' >> FlatMap(set) assert_that(pcoll4, equal_to([11, 12, 12, 12, 13]), label='pcoll4') def test_maptuple_builtin(self): with TestPipeline() as pipeline: pcoll = pipeline | Create([('e1', 'e2')]) side1 = beam.pvalue.AsSingleton(pipeline | 'side1' >> Create(['s1'])) side2 = beam.pvalue.AsSingleton(pipeline | 'side2' >> Create(['s2'])) # A test function with a tuple input, an auxiliary parameter, # and some side inputs. fn = lambda e1, e2, t=DoFn.TimestampParam, s1=None, s2=None: ( e1, e2, t, s1, s2) assert_that( pcoll | 'NoSides' >> beam.core.MapTuple(fn), equal_to([('e1', 'e2', MIN_TIMESTAMP, None, None)]), label='NoSidesCheck') assert_that( pcoll | 'StaticSides' >> beam.core.MapTuple(fn, 's1', 's2'), equal_to([('e1', 'e2', MIN_TIMESTAMP, 's1', 's2')]), label='StaticSidesCheck') assert_that( pcoll | 'DynamicSides' >> beam.core.MapTuple(fn, side1, side2), equal_to([('e1', 'e2', MIN_TIMESTAMP, 's1', 's2')]), label='DynamicSidesCheck') assert_that( pcoll | 'MixedSides' >> beam.core.MapTuple(fn, s2=side2), equal_to([('e1', 'e2', MIN_TIMESTAMP, None, 's2')]), label='MixedSidesCheck') def test_flatmaptuple_builtin(self): with TestPipeline() as pipeline: pcoll = pipeline | Create([('e1', 'e2')]) side1 = beam.pvalue.AsSingleton(pipeline | 'side1' >> Create(['s1'])) side2 = beam.pvalue.AsSingleton(pipeline | 'side2' >> Create(['s2'])) # A test function with a tuple input, an auxiliary parameter, # and some side inputs. fn = lambda e1, e2, t=DoFn.TimestampParam, s1=None, s2=None: ( e1, e2, t, s1, s2) assert_that( pcoll | 'NoSides' >> beam.core.FlatMapTuple(fn), equal_to(['e1', 'e2', MIN_TIMESTAMP, None, None]), label='NoSidesCheck') assert_that( pcoll | 'StaticSides' >> beam.core.FlatMapTuple(fn, 's1', 's2'), equal_to(['e1', 'e2', MIN_TIMESTAMP, 's1', 's2']), label='StaticSidesCheck') assert_that( pcoll | 'DynamicSides' >> beam.core.FlatMapTuple(fn, side1, side2), equal_to(['e1', 'e2', MIN_TIMESTAMP, 's1', 's2']), label='DynamicSidesCheck') assert_that( pcoll | 'MixedSides' >> beam.core.FlatMapTuple(fn, s2=side2), equal_to(['e1', 'e2', MIN_TIMESTAMP, None, 's2']), label='MixedSidesCheck') def test_create_singleton_pcollection(self): with TestPipeline() as pipeline: pcoll = pipeline | 'label' >> Create([[1, 2, 3]]) assert_that(pcoll, equal_to([[1, 2, 3]])) def test_visit_entire_graph(self): pipeline = Pipeline() pcoll1 = pipeline | 'pcoll' >> beam.Impulse() pcoll2 = pcoll1 | 'do1' >> FlatMap(lambda x: [x + 1]) pcoll3 = pcoll2 | 'do2' >> FlatMap(lambda x: [x + 1]) pcoll4 = pcoll2 | 'do3' >> FlatMap(lambda x: [x + 1]) transform = PipelineTest.CustomTransform() pcoll5 = pcoll4 | transform visitor = PipelineTest.Visitor(visited=[]) pipeline.visit(visitor) self.assertEqual({pcoll1, pcoll2, pcoll3, pcoll4, pcoll5}, set(visitor.visited)) self.assertEqual(set(visitor.enter_composite), set(visitor.leave_composite)) self.assertEqual(2, len(visitor.enter_composite)) self.assertEqual(visitor.enter_composite[1].transform, transform) self.assertEqual(visitor.leave_composite[0].transform, transform) def test_apply_custom_transform(self): with TestPipeline() as pipeline: pcoll = pipeline | 'pcoll' >> Create([1, 2, 3]) result = pcoll | PipelineTest.CustomTransform() assert_that(result, equal_to([2, 3, 4])) def test_reuse_custom_transform_instance(self): pipeline = Pipeline() pcoll1 = pipeline | 'pcoll1' >> Create([1, 2, 3]) pcoll2 = pipeline | 'pcoll2' >> Create([4, 5, 6]) transform = PipelineTest.CustomTransform() pcoll1 | transform with self.assertRaises(RuntimeError) as cm: pipeline.apply(transform, pcoll2) self.assertEqual( cm.exception.args[0], 'A transform with label "CustomTransform" already exists in the ' 'pipeline. To apply a transform with a specified label, write ' 'pvalue | "label" >> transform or use the option ' '"auto_unique_labels" to automatically generate unique ' 'transform labels. Note "auto_unique_labels" ' 'could cause data loss when updating a pipeline or ' 'reloading the job state. This is not recommended for ' 'streaming jobs.') @mock.patch('logging.info') # Mock the logging.info function def test_no_wait_until_finish(self, mock_info): with Pipeline(runner='DirectRunner', options=PipelineOptions(["--no_wait_until_finish"])) as p: _ = p | beam.Create(['test']) mock_info.assert_any_call( 'Job execution continues without waiting for completion. ' 'Use "wait_until_finish" in PipelineResult to block until finished.') p.result.wait_until_finish() def test_auto_unique_labels(self): opts = PipelineOptions(["--auto_unique_labels"]) mock_uuids = [mock.Mock(hex='UUID01XXX'), mock.Mock(hex='UUID02XXX')] mock_uuid_gen = mock.Mock(side_effect=mock_uuids) original_generate_unique_label = Pipeline._generate_unique_label def patched_generate_unique_label(self, transform): with mock.patch.object(uuid, 'uuid4', return_value=mock_uuid_gen()): return original_generate_unique_label(self, transform) with mock.patch.object(Pipeline, '_generate_unique_label', patched_generate_unique_label): with TestPipeline(options=opts) as pipeline: pcoll = pipeline | 'pcoll' >> Create([1, 2, 3]) def identity(x): return x pcoll2 = pcoll | Map(identity) pcoll3 = pcoll2 | Map(identity) pcoll4 = pcoll3 | Map(identity) assert_that(pcoll4, equal_to([1, 2, 3])) map_id_full_labels = { label for label in pipeline.applied_labels if "Map(identity)" in label } map_id_leaf_labels = {label.split(":")[-1] for label in map_id_full_labels} # Only the first 6 chars of the UUID hex should be used assert map_id_leaf_labels == set( ["Map(identity)", "Map(identity)_UUID01", "Map(identity)_UUID02"]) def test_reuse_cloned_custom_transform_instance(self): with TestPipeline() as pipeline: pcoll1 = pipeline | 'pc1' >> Create([1, 2, 3]) pcoll2 = pipeline | 'pc2' >> Create([4, 5, 6]) transform = PipelineTest.CustomTransform() result1 = pcoll1 | transform result2 = pcoll2 | 'new_label' >> transform assert_that(result1, equal_to([2, 3, 4]), label='r1') assert_that(result2, equal_to([5, 6, 7]), label='r2') def test_transform_no_super_init(self): class AddSuffix(PTransform): def __init__(self, suffix): # No call to super(...).__init__ self.suffix = suffix def expand(self, pcoll): return pcoll | Map(lambda x: x + self.suffix) self.assertEqual(['a-x', 'b-x', 'c-x'], sorted(['a', 'b', 'c'] | 'AddSuffix' >> AddSuffix('-x'))) @unittest.skip("Fails on some platforms with new urllib3.") def test_memory_usage(self): try: import resource except ImportError: # Skip the test if resource module is not available (e.g. non-Unix os). self.skipTest('resource module not available.') if platform.mac_ver()[0]: # Skip the test on macos, depending on version it returns ru_maxrss in # different units. self.skipTest('ru_maxrss is not in standard units.') def get_memory_usage_in_bytes(): return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * (2**10) def check_memory(value, memory_threshold): memory_usage = get_memory_usage_in_bytes() if memory_usage > memory_threshold: raise RuntimeError( 'High memory usage: %d > %d' % (memory_usage, memory_threshold)) return value len_elements = 1000000 num_elements = 10 num_maps = 100 # TODO(robertwb): reduce memory usage of FnApiRunner so that this test # passes. with TestPipeline(runner='BundleBasedDirectRunner') as pipeline: # Consumed memory should not be proportional to the number of maps. memory_threshold = ( get_memory_usage_in_bytes() + (5 * len_elements * num_elements)) # Plus small additional slack for memory fluctuations during the test. memory_threshold += 10 * (2**20) biglist = pipeline | 'oom:create' >> Create( ['x' * len_elements] * num_elements) for i in range(num_maps): biglist = biglist | ('oom:addone-%d' % i) >> Map(lambda x: x + 'y') result = biglist | 'oom:check' >> Map(check_memory, memory_threshold) assert_that( result, equal_to(['x' * len_elements + 'y' * num_maps] * num_elements)) def test_aggregator_empty_input(self): actual = [] | CombineGlobally(max).without_defaults() self.assertEqual(actual, []) def test_pipeline_as_context(self): def raise_exception(exn): raise exn with self.assertRaises(Exception): with Pipeline() as p: # pylint: disable=expression-not-assigned p | Create([ValueError('msg')]) | Map(raise_exception) def test_ptransform_overrides(self): class MyParDoOverride(PTransformOverride): def matches(self, applied_ptransform): return isinstance(applied_ptransform.transform, DoubleParDo) def get_replacement_transform_for_applied_ptransform( self, applied_ptransform): ptransform = applied_ptransform.transform if isinstance(ptransform, DoubleParDo): return TripleParDo() raise ValueError('Unsupported type of transform: %r' % ptransform) p = Pipeline() pcoll = p | beam.Create([1, 2, 3]) | 'Multiply' >> DoubleParDo() assert_that(pcoll, equal_to([3, 6, 9])) p.replace_all([MyParDoOverride()]) p.run() def test_ptransform_override_type_hints(self): class NoTypeHintOverride(PTransformOverride): def matches(self, applied_ptransform): return isinstance(applied_ptransform.transform, DoubleParDo) def get_replacement_transform_for_applied_ptransform( self, applied_ptransform): return ToStringParDo() class WithTypeHintOverride(PTransformOverride): def matches(self, applied_ptransform): return isinstance(applied_ptransform.transform, DoubleParDo) def get_replacement_transform_for_applied_ptransform( self, applied_ptransform): return ToStringParDo().with_input_types(int).with_output_types(str) for override, expected_type in [(NoTypeHintOverride(), int), (WithTypeHintOverride(), str)]: p = TestPipeline() pcoll = ( p | beam.Create([1, 2, 3]) | 'Operate' >> DoubleParDo() | 'NoOp' >> beam.Map(lambda x: x)) p.replace_all([override]) self.assertEqual(pcoll.producer.inputs[0].element_type, expected_type) def test_ptransform_override_multiple_inputs(self): class MyParDoOverride(PTransformOverride): def matches(self, applied_ptransform): return isinstance(applied_ptransform.transform, FlattenAndDouble) def get_replacement_transform(self, applied_ptransform): return FlattenAndTriple() p = Pipeline() pcoll1 = p | 'pc1' >> beam.Create([1, 2, 3]) pcoll2 = p | 'pc2' >> beam.Create([4, 5, 6]) pcoll3 = (pcoll1, pcoll2) | 'FlattenAndMultiply' >> FlattenAndDouble() assert_that(pcoll3, equal_to([3, 6, 9, 12, 15, 18])) p.replace_all([MyParDoOverride()]) p.run() def test_ptransform_override_side_inputs(self): class MyParDoOverride(PTransformOverride): def matches(self, applied_ptransform): return ( isinstance(applied_ptransform.transform, ParDo) and isinstance(applied_ptransform.transform.fn, AddWithProductDoFn)) def get_replacement_transform(self, transform): return AddThenMultiply() p = Pipeline() pcoll1 = p | 'pc1' >> beam.Create([2]) pcoll2 = p | 'pc2' >> beam.Create([3]) pcoll3 = p | 'pc3' >> beam.Create([4, 5, 6]) result = pcoll3 | 'Operate' >> beam.ParDo( AddWithProductDoFn(), AsSingleton(pcoll1), AsSingleton(pcoll2)) assert_that(result, equal_to([18, 21, 24])) p.replace_all([MyParDoOverride()]) p.run() def test_ptransform_override_replacement_inputs(self): class MyParDoOverride(PTransformOverride): def matches(self, applied_ptransform): return ( isinstance(applied_ptransform.transform, ParDo) and isinstance(applied_ptransform.transform.fn, AddWithProductDoFn)) def get_replacement_transform(self, transform): return AddThenMultiply() def get_replacement_inputs(self, applied_ptransform): assert len(applied_ptransform.inputs) == 1 assert len(applied_ptransform.side_inputs) == 2 # Swap the order of the two side inputs return ( applied_ptransform.inputs[0], applied_ptransform.side_inputs[1].pvalue, applied_ptransform.side_inputs[0].pvalue) p = Pipeline() pcoll1 = p | 'pc1' >> beam.Create([2]) pcoll2 = p | 'pc2' >> beam.Create([3]) pcoll3 = p | 'pc3' >> beam.Create([4, 5, 6]) result = pcoll3 | 'Operate' >> beam.ParDo( AddWithProductDoFn(), AsSingleton(pcoll1), AsSingleton(pcoll2)) assert_that(result, equal_to([14, 16, 18])) p.replace_all([MyParDoOverride()]) p.run() def test_ptransform_override_multiple_outputs(self): class MultiOutputComposite(PTransform): def __init__(self): self.output_tags = set() def expand(self, pcoll): def mux_input(x): x = x * 2 if isinstance(x, int): yield TaggedOutput('numbers', x) else: yield TaggedOutput('letters', x) multi = pcoll | 'MyReplacement' >> beam.ParDo(mux_input).with_outputs() letters = multi.letters | 'LettersComposite' >> beam.Map( lambda x: x * 3) numbers = multi.numbers | 'NumbersComposite' >> beam.Map( lambda x: x * 5) return { 'letters': letters, 'numbers': numbers, } class MultiOutputOverride(PTransformOverride): def matches(self, applied_ptransform): return applied_ptransform.full_label == 'MyMultiOutput' def get_replacement_transform_for_applied_ptransform( self, applied_ptransform): return MultiOutputComposite() def mux_input(x): if isinstance(x, int): yield TaggedOutput('numbers', x) else: yield TaggedOutput('letters', x) with TestPipeline() as p: multi = ( p | beam.Create([1, 2, 3, 'a', 'b', 'c']) | 'MyMultiOutput' >> beam.ParDo(mux_input).with_outputs()) letters = multi.letters | 'MyLetters' >> beam.Map(lambda x: x) numbers = multi.numbers | 'MyNumbers' >> beam.Map(lambda x: x) # Assert that the PCollection replacement worked correctly and that # elements are flowing through. The replacement transform first # multiples by 2 then the leaf nodes inside the composite multiply by # an additional 3 and 5. Use prime numbers to ensure that each # transform is getting executed once. assert_that( letters, equal_to(['a' * 2 * 3, 'b' * 2 * 3, 'c' * 2 * 3]), label='assert letters') assert_that( numbers, equal_to([1 * 2 * 5, 2 * 2 * 5, 3 * 2 * 5]), label='assert numbers') # Do the replacement and run the element assertions. p.replace_all([MultiOutputOverride()]) # The following checks the graph to make sure the replacement occurred. visitor = PipelineTest.Visitor(visited=[]) p.visit(visitor) pcollections = visitor.visited composites = visitor.enter_composite # Assert the replacement is in the composite list and retrieve the # AppliedPTransform. self.assertIn( MultiOutputComposite, [t.transform.__class__ for t in composites]) multi_output_composite = list( filter( lambda t: t.transform.__class__ == MultiOutputComposite, composites))[0] # Assert that all of the replacement PCollections are in the graph. for output in multi_output_composite.outputs.values(): self.assertIn(output, pcollections) # Assert that all of the "old"/replaced PCollections are not in the graph. self.assertNotIn(multi[None], visitor.visited) self.assertNotIn(multi.letters, visitor.visited) self.assertNotIn(multi.numbers, visitor.visited) def test_filter_typehint(self): # Check input type hint and output type hint are both specified. def always_true_with_all_typehints(x: int) -> bool: return True # Check only input type hint is specified. def always_true_only_inptype(x: int): return True # Check only output type hint is specified. def always_true_only_outptype(x) -> bool: return True # Check if inp type hint is Any that we can still infer # from the input pcollection type def always_true_any_inptype(x: typehints.Any) -> bool: return True for filter_fn in [always_true_with_all_typehints, always_true_only_inptype, always_true_only_outptype, always_true_any_inptype]: with TestPipeline() as p: pcoll = ( p | beam.Create([1, 2, 3]).with_input_types(int) | beam.Filter(filter_fn)) self.assertEqual(pcoll.element_type, int) def test_kv_ptransform_honor_type_hints(self): # The return type of this DoFn cannot be inferred by the default # Beam type inference class StatefulDoFn(DoFn): BYTES_STATE = BagStateSpec('bytes', BytesCoder()) def return_recursive(self, count): if count == 0: return ["some string"] else: self.return_recursive(count - 1) def process(self, element, counter=DoFn.StateParam(BYTES_STATE)): return self.return_recursive(1) with TestPipeline() as p: pcoll = ( p | beam.Create([(1, 1), (2, 2), (3, 3)]) | beam.GroupByKey() | beam.ParDo(StatefulDoFn())) self.assertEqual(pcoll.element_type, typehints.Any) with TestPipeline() as p: pcoll = ( p | beam.Create([(1, 1), (2, 2), (3, 3)]) | beam.GroupByKey() | beam.ParDo(StatefulDoFn()).with_output_types(str)) self.assertEqual(pcoll.element_type, str) def test_track_pcoll_unbounded(self): pipeline = TestPipeline() pcoll1 = pipeline | 'read' >> Read(FakeUnboundedSource()) pcoll2 = pcoll1 | 'do1' >> FlatMap(lambda x: [x + 1]) pcoll3 = pcoll2 | 'do2' >> FlatMap(lambda x: [x + 1]) self.assertIs(pcoll1.is_bounded, False) self.assertIs(pcoll2.is_bounded, False) self.assertIs(pcoll3.is_bounded, False) def test_track_pcoll_bounded(self): pipeline = TestPipeline() pcoll1 = pipeline | 'label1' >> Create([1, 2, 3]) pcoll2 = pcoll1 | 'do1' >> FlatMap(lambda x: [x + 1]) pcoll3 = pcoll2 | 'do2' >> FlatMap(lambda x: [x + 1]) self.assertIs(pcoll1.is_bounded, True) self.assertIs(pcoll2.is_bounded, True) self.assertIs(pcoll3.is_bounded, True) def test_track_pcoll_bounded_flatten(self): pipeline = TestPipeline() pcoll1_a = pipeline | 'label_a' >> Create([1, 2, 3]) pcoll2_a = pcoll1_a | 'do_a' >> FlatMap(lambda x: [x + 1]) pcoll1_b = pipeline | 'label_b' >> Create([1, 2, 3]) pcoll2_b = pcoll1_b | 'do_b' >> FlatMap(lambda x: [x + 1]) merged = (pcoll2_a, pcoll2_b) | beam.Flatten() self.assertIs(pcoll1_a.is_bounded, True) self.assertIs(pcoll2_a.is_bounded, True) self.assertIs(pcoll1_b.is_bounded, True) self.assertIs(pcoll2_b.is_bounded, True) self.assertIs(merged.is_bounded, True) def test_track_pcoll_unbounded_flatten(self): pipeline = TestPipeline() pcoll1_bounded = pipeline | 'label1' >> Create([1, 2, 3]) pcoll2_bounded = pcoll1_bounded | 'do1' >> FlatMap(lambda x: [x + 1]) pcoll1_unbounded = pipeline | 'read' >> Read(FakeUnboundedSource()) pcoll2_unbounded = pcoll1_unbounded | 'do2' >> FlatMap(lambda x: [x + 1]) merged = (pcoll2_bounded, pcoll2_unbounded) | beam.Flatten() self.assertIs(pcoll1_bounded.is_bounded, True) self.assertIs(pcoll2_bounded.is_bounded, True) self.assertIs(pcoll1_unbounded.is_bounded, False) self.assertIs(pcoll2_unbounded.is_bounded, False) self.assertIs(merged.is_bounded, False) def test_incompatible_pcollection_errmsg(self): with pytest.raises(Exception, match=r".*Map\(print\).*Got a PBegin/Pipeline instead."): with beam.Pipeline() as pipeline: _ = (pipeline | beam.Map(print)) class ParentTransform(PTransform): def expand(self, pcoll): return pcoll | beam.Map(print) with pytest.raises( Exception, match=r".*ParentTransform/Map\(print\).*Got a PBegin/Pipeline instead." ): with beam.Pipeline() as pipeline: _ = (pipeline | ParentTransform()) def test_incompatible_submission_and_runtime_envs_fail_pipeline(self): with mock.patch( 'apache_beam.transforms.environments.sdk_base_version_capability' ) as base_version: base_version.side_effect = [ f"beam:version:sdk_base:apache/beam_python3.5_sdk:2.{i}.0" for i in range(100) ] with self.assertRaisesRegex( RuntimeError, 'Pipeline construction environment and pipeline runtime ' 'environment are not compatible.'): # TODO(https://github.com/apache/beam/issues/34549): Prism doesn't # pass through capabilities as part of the ProcessBundleDescriptor. with TestPipeline('FnApiRunner') as p: _ = p | Create([None]) class DoFnTest(unittest.TestCase): def test_element(self): class TestDoFn(DoFn): def process(self, element): yield element + 10 with TestPipeline() as pipeline: pcoll = pipeline | 'Create' >> Create([1, 2]) | 'Do' >> ParDo(TestDoFn()) assert_that(pcoll, equal_to([11, 12])) def test_side_input_no_tag(self): class TestDoFn(DoFn): def process(self, element, prefix, suffix): return ['%s-%s-%s' % (prefix, element, suffix)] with TestPipeline() as pipeline: words_list = ['aa', 'bb', 'cc'] words = pipeline | 'SomeWords' >> Create(words_list) prefix = 'zyx' suffix = pipeline | 'SomeString' >> Create(['xyz']) # side in result = words | 'DecorateWordsDoFnNoTag' >> ParDo( TestDoFn(), prefix, suffix=AsSingleton(suffix)) assert_that(result, equal_to(['zyx-%s-xyz' % x for x in words_list])) def test_side_input_tagged(self): class TestDoFn(DoFn): def process(self, element, prefix, suffix=DoFn.SideInputParam): return ['%s-%s-%s' % (prefix, element, suffix)] with TestPipeline() as pipeline: words_list = ['aa', 'bb', 'cc'] words = pipeline | 'SomeWords' >> Create(words_list) prefix = 'zyx' suffix = pipeline | 'SomeString' >> Create(['xyz']) # side in result = words | 'DecorateWordsDoFnNoTag' >> ParDo( TestDoFn(), prefix, suffix=AsSingleton(suffix)) assert_that(result, equal_to(['zyx-%s-xyz' % x for x in words_list])) @pytest.mark.it_validatesrunner def test_element_param(self): pipeline = TestPipeline() input = [1, 2] pcoll = ( pipeline | 'Create' >> Create(input) | 'Ele param' >> Map(lambda element=DoFn.ElementParam: element)) assert_that(pcoll, equal_to(input)) pipeline.run() @pytest.mark.it_validatesrunner def test_key_param(self): pipeline = TestPipeline() pcoll = ( pipeline | 'Create' >> Create([('a', 1), ('b', 2)]) | 'Key param' >> Map(lambda _, key=DoFn.KeyParam: key)) assert_that(pcoll, equal_to(['a', 'b'])) pipeline.run() def test_window_param(self): class TestDoFn(DoFn): def process(self, element, window=DoFn.WindowParam): yield (element, (float(window.start), float(window.end))) with TestPipeline() as pipeline: pcoll = ( pipeline | Create([1, 7]) | Map(lambda x: TimestampedValue(x, x)) | WindowInto(windowfn=SlidingWindows(10, 5)) | ParDo(TestDoFn())) assert_that( pcoll, equal_to([(1, (-5, 5)), (1, (0, 10)), (7, (0, 10)), (7, (5, 15))])) pcoll2 = pcoll | 'Again' >> ParDo(TestDoFn()) assert_that( pcoll2, equal_to([((1, (-5, 5)), (-5, 5)), ((1, (0, 10)), (0, 10)), ((7, (0, 10)), (0, 10)), ((7, (5, 15)), (5, 15))]), label='doubled windows') def test_windowed_value_param(self): with TestPipeline() as pipeline: pcoll = ( pipeline | Create([1, 7]) | Map(lambda x: TimestampedValue(x, x)) | WindowInto(windowfn=FixedWindows(5)) | Map(lambda _, wv=DoFn.WindowedValueParam: (wv.value, wv.windows))) assert_that( pcoll, equal_to([(1, [IntervalWindow(0, 5)]), (7, [IntervalWindow(5, 10)])])) # pylint: disable=too-many-function-args def test_timestamp_param(self): class TestDoFn(DoFn): def process(self, element, timestamp=DoFn.TimestampParam): yield timestamp with TestPipeline() as pipeline: pcoll = pipeline | 'Create' >> Create([1, 2]) | 'Do' >> ParDo(TestDoFn()) assert_that(pcoll, equal_to([MIN_TIMESTAMP, MIN_TIMESTAMP])) def test_timestamp_param_map(self): with TestPipeline() as p: assert_that( p | Create([1, 2]) | beam.Map(lambda _, t=DoFn.TimestampParam: t), equal_to([MIN_TIMESTAMP, MIN_TIMESTAMP])) def test_pane_info_param(self): with TestPipeline() as p: pc = p | Create([(None, None)]) assert_that( pc | beam.Map(lambda _, p=DoFn.PaneInfoParam: p), equal_to([windowed_value.PANE_INFO_UNKNOWN]), label='CheckUngrouped') assert_that( pc | beam.GroupByKey() | beam.Map(lambda _, p=DoFn.PaneInfoParam: p), equal_to([ windowed_value.PaneInfo( is_first=True, is_last=True, timing=windowed_value.PaneInfoTiming.ON_TIME, index=0, nonspeculative_index=0) ]), label='CheckGrouped') def test_context_params(self): def test_map( x, context_a=DoFn.BundleContextParam(_TestContext, args=('a')), context_b=DoFn.BundleContextParam(_TestContext, args=('b')), context_c=DoFn.SetupContextParam(_TestContext, args=('c'))): return (x, context_a, context_b, context_c) self.assertEqual(_TestContext.live_contexts, 0) with TestPipeline('FnApiRunner') as p: pcoll = p | Create([1, 2]) | beam.Map(test_map) assert_that(pcoll, equal_to([(1, 'a', 'b', 'c'), (2, 'a', 'b', 'c')])) self.assertEqual(_TestContext.live_contexts, 0) def test_incomparable_default(self): class IncomparableType(object): def __eq__(self, other): raise RuntimeError() def __ne__(self, other): raise RuntimeError() def __hash__(self): raise RuntimeError() # Ensure that we don't use default values in a context where they must be # comparable (see BEAM-8301). with TestPipeline() as pipeline: pcoll = ( pipeline | beam.Create([None]) | Map(lambda e, x=IncomparableType(): (e, type(x).__name__))) assert_that(pcoll, equal_to([(None, 'IncomparableType')])) class _TestContext: live_contexts = 0 def __init__(self, value): self._value = value def __enter__(self): _TestContext.live_contexts += 1 return self._value def __exit__(self, *args): _TestContext.live_contexts -= 1 class Bacon(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--slices', type=int) class Eggs(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--style', default='scrambled') class Breakfast(Bacon, Eggs): pass class PipelineOptionsTest(unittest.TestCase): def test_flag_parsing(self): options = Breakfast(['--slices=3', '--style=sunny side up', '--ignored']) self.assertEqual(3, options.slices) self.assertEqual('sunny side up', options.style) def test_keyword_parsing(self): options = Breakfast(['--slices=3', '--style=sunny side up', '--ignored'], slices=10) self.assertEqual(10, options.slices) self.assertEqual('sunny side up', options.style) def test_attribute_setting(self): options = Breakfast(slices=10) self.assertEqual(10, options.slices) options.slices = 20 self.assertEqual(20, options.slices) def test_view_as(self): generic_options = PipelineOptions(['--slices=3']) self.assertEqual(3, generic_options.view_as(Bacon).slices) self.assertEqual(3, generic_options.view_as(Breakfast).slices) generic_options.view_as(Breakfast).slices = 10 self.assertEqual(10, generic_options.view_as(Bacon).slices) with self.assertRaises(AttributeError): generic_options.slices # pylint: disable=pointless-statement with self.assertRaises(AttributeError): generic_options.view_as(Eggs).slices # pylint: disable=expression-not-assigned def test_defaults(self): options = Breakfast(['--slices=3']) self.assertEqual(3, options.slices) self.assertEqual('scrambled', options.style) def test_dir(self): options = Breakfast() self.assertEqual({ 'from_dictionary', 'get_all_options', 'is_compat_version_prior_to', 'slices', 'style', 'view_as', 'display_data', 'from_runner_api', 'to_runner_api', }, { attr for attr in dir(options) if not attr.startswith('_') and attr != 'next' }) self.assertEqual({ 'from_dictionary', 'get_all_options', 'is_compat_version_prior_to', 'style', 'view_as', 'display_data', 'from_runner_api', 'to_runner_api', }, { attr for attr in dir(options.view_as(Eggs)) if not attr.startswith('_') and attr != 'next' }) class RunnerApiTest(unittest.TestCase): def test_parent_pointer(self): class MyPTransform(beam.PTransform): def expand(self, p): self.p = p return p | beam.Create([None]) p = beam.Pipeline() p | MyPTransform() # pylint: disable=expression-not-assigned p = Pipeline.from_runner_api( Pipeline.to_runner_api(p, use_fake_coders=True), None, None) self.assertIsNotNone(p.transforms_stack[0].parts[0].parent) self.assertEqual( p.transforms_stack[0].parts[0].parent, p.transforms_stack[0]) def test_requirements(self): p = beam.Pipeline() _ = ( p | beam.Create([]) | beam.ParDo(lambda x, finalize=beam.DoFn.BundleFinalizerParam: None)) proto = p.to_runner_api() self.assertTrue( common_urns.requirements.REQUIRES_BUNDLE_FINALIZATION.urn, proto.requirements) def test_coder_version_tag_included_in_runner_api_key(self): class MyClass: def __init__(self, value: int): self.value = value class VersionedCoder(coders.Coder): def encode(self, value): return str(value.value).encode() def decode(self, encoded): return MyClass(int(encoded.decode())) def version_tag(self): return "v269" def to_type_hint(self): return MyClass coders.registry.register_coder(MyClass, VersionedCoder) p = beam.Pipeline() _ = (p | beam.Impulse() | beam.Map(lambda _: MyClass(1))) pipeline_proto = p.to_runner_api() coder_keys = sorted(list(pipeline_proto.components.coders.keys())) self.assertListEqual( coder_keys, [ 'ref_Coder_BytesCoder_1', 'ref_Coder_GlobalWindowCoder_2', 'ref_Coder_VersionedCoder_v269_3' ]) def test_annotations(self): some_proto = BytesCoder().to_runner_api(None) class EmptyTransform(beam.PTransform): def expand(self, pcoll): return pcoll def annotations(self): return {'foo': 'some_string'} class NonEmptyTransform(beam.PTransform): def expand(self, pcoll): return pcoll | beam.Map(lambda x: x) def annotations(self): return { 'foo': b'some_bytes', 'proto': some_proto, } p = beam.Pipeline() _ = p | beam.Create([]) | EmptyTransform() | NonEmptyTransform() proto = p.to_runner_api() seen = 0 for transform in proto.components.transforms.values(): if transform.unique_name == 'EmptyTransform': seen += 1 self.assertEqual(transform.annotations['foo'], b'some_string') elif transform.unique_name == 'NonEmptyTransform': seen += 1 self.assertEqual(transform.annotations['foo'], b'some_bytes') self.assertEqual( transform.annotations['proto'], some_proto.SerializeToString()) self.assertEqual(seen, 2) def assertHasAnnotation(self, pipeline_proto, transform, key, value): for transform_proto in pipeline_proto.components.transforms.values(): if transform_proto.unique_name == transform: self.assertIn(key, transform_proto.annotations.keys()) self.assertEqual(transform_proto.annotations[key], value) break else: self.fail( "Unknown transform: %r not in %s" % ( transform, sorted([ t.unique_name for t in pipeline_proto.components.transforms.values() ]))) def test_pipeline_context_annotations(self): p = beam.Pipeline() with p.transform_annotations(foo='first'): pcoll = p | beam.Create([1, 2, 3]) | 'First' >> beam.Map(lambda x: x + 1) with p.transform_annotations(foo='second'): _ = pcoll | 'Second' >> beam.Map(lambda x: x * 2) with p.transform_annotations(foo='nested', another='more'): _ = pcoll | 'Nested' >> beam.Map(lambda x: x * 3) proto = p.to_runner_api() self.assertHasAnnotation(proto, 'First', 'foo', b'first') self.assertHasAnnotation(proto, 'Second', 'foo', b'second') self.assertHasAnnotation(proto, 'Nested', 'foo', b'nested') self.assertHasAnnotation(proto, 'Nested', 'another', b'more') def test_beam_context_annotations(self): p = beam.Pipeline() with beam.transform_annotations(foo='first'): pcoll = p | beam.Create([1, 2, 3]) | 'First' >> beam.Map(lambda x: x + 1) with beam.transform_annotations(foo='second'): _ = pcoll | 'Second' >> beam.Map(lambda x: x * 2) with beam.transform_annotations(foo='nested', another='more'): _ = pcoll | 'Nested' >> beam.Map(lambda x: x * 3) proto = p.to_runner_api() self.assertHasAnnotation(proto, 'First', 'foo', b'first') self.assertHasAnnotation(proto, 'Second', 'foo', b'second') self.assertHasAnnotation(proto, 'Nested', 'foo', b'nested') self.assertHasAnnotation(proto, 'Nested', 'another', b'more') def test_transform_ids(self): class MyPTransform(beam.PTransform): def expand(self, p): self.p = p return p | beam.Create([None]) p = beam.Pipeline() p | MyPTransform() # pylint: disable=expression-not-assigned runner_api_proto = Pipeline.to_runner_api(p) for transform_id in runner_api_proto.components.transforms: self.assertRegex(transform_id, r'[a-zA-Z0-9-_]+') def test_input_names(self): class MyPTransform(beam.PTransform): def expand(self, pcolls): return pcolls.values() | beam.Flatten() p = beam.Pipeline() input_names = set('ABC') inputs = {x: p | x >> beam.Create([x]) for x in input_names} inputs | MyPTransform() # pylint: disable=expression-not-assigned runner_api_proto = Pipeline.to_runner_api(p) for transform_proto in runner_api_proto.components.transforms.values(): if transform_proto.unique_name == 'MyPTransform': self.assertEqual(set(transform_proto.inputs.keys()), input_names) break else: self.fail('Unable to find transform.') def test_display_data(self): class MyParentTransform(beam.PTransform): def expand(self, p): self.p = p return p | beam.Create([None]) def display_data(self) -> dict: parent_dd = super().display_data() parent_dd['p_dd_string'] = DisplayDataItem( 'p_dd_string_value', label='p_dd_string_label') parent_dd['p_dd_string_2'] = DisplayDataItem('p_dd_string_value_2') parent_dd['p_dd_bool'] = DisplayDataItem(True, label='p_dd_bool_label') parent_dd['p_dd_int'] = DisplayDataItem(1, label='p_dd_int_label') return parent_dd class MyPTransform(MyParentTransform): def expand(self, p): self.p = p return p | beam.Create([None]) def display_data(self) -> dict: parent_dd = super().display_data() parent_dd['dd_string'] = DisplayDataItem( 'dd_string_value', label='dd_string_label') parent_dd['dd_string_2'] = DisplayDataItem('dd_string_value_2') parent_dd['dd_bool'] = DisplayDataItem(False, label='dd_bool_label') parent_dd['dd_double'] = DisplayDataItem(1.1, label='dd_double_label') return parent_dd p = beam.Pipeline() p | MyPTransform() # pylint: disable=expression-not-assigned proto_pipeline = Pipeline.to_runner_api(p, use_fake_coders=True) my_transform, = [ transform for transform in proto_pipeline.components.transforms.values() if transform.unique_name == 'MyPTransform' ] self.assertIsNotNone(my_transform) self.assertListEqual( list(my_transform.display_data), [ beam_runner_api_pb2.DisplayData( urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn, payload=beam_runner_api_pb2.LabelledPayload( label='p_dd_string_label', key='p_dd_string', namespace='apache_beam.pipeline_test.MyPTransform', string_value='p_dd_string_value').SerializeToString()), beam_runner_api_pb2.DisplayData( urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn, payload=beam_runner_api_pb2.LabelledPayload( label='p_dd_string_2', key='p_dd_string_2', namespace='apache_beam.pipeline_test.MyPTransform', string_value='p_dd_string_value_2').SerializeToString()), beam_runner_api_pb2.DisplayData( urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn, payload=beam_runner_api_pb2.LabelledPayload( label='p_dd_bool_label', key='p_dd_bool', namespace='apache_beam.pipeline_test.MyPTransform', bool_value=True).SerializeToString()), beam_runner_api_pb2.DisplayData( urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn, payload=beam_runner_api_pb2.LabelledPayload( label='p_dd_int_label', key='p_dd_int', namespace='apache_beam.pipeline_test.MyPTransform', int_value=1).SerializeToString()), beam_runner_api_pb2.DisplayData( urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn, payload=beam_runner_api_pb2.LabelledPayload( label='dd_string_label', key='dd_string', namespace='apache_beam.pipeline_test.MyPTransform', string_value='dd_string_value').SerializeToString()), beam_runner_api_pb2.DisplayData( urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn, payload=beam_runner_api_pb2.LabelledPayload( label='dd_string_2', key='dd_string_2', namespace='apache_beam.pipeline_test.MyPTransform', string_value='dd_string_value_2').SerializeToString()), beam_runner_api_pb2.DisplayData( urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn, payload=beam_runner_api_pb2.LabelledPayload( label='dd_bool_label', key='dd_bool', namespace='apache_beam.pipeline_test.MyPTransform', bool_value=False).SerializeToString()), beam_runner_api_pb2.DisplayData( urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn, payload=beam_runner_api_pb2.LabelledPayload( label='dd_double_label', key='dd_double', namespace='apache_beam.pipeline_test.MyPTransform', double_value=1.1).SerializeToString()), ]) def test_runner_api_roundtrip_preserves_resource_hints(self): p = beam.Pipeline() _ = ( p | beam.Create([1, 2]) | beam.Map(lambda x: x + 1).with_resource_hints(accelerator='gpu')) self.assertEqual( p.transforms_stack[0].parts[1].transform.get_resource_hints(), {common_urns.resource_hints.ACCELERATOR.urn: b'gpu'}) for _ in range(3): # Verify that DEFAULT environments are recreated during multiple RunnerAPI # translation and hints don't get lost. p = Pipeline.from_runner_api(Pipeline.to_runner_api(p), None, None) self.assertEqual( p.transforms_stack[0].parts[1].transform.get_resource_hints(), {common_urns.resource_hints.ACCELERATOR.urn: b'gpu'}) def test_hints_on_composite_transforms_are_propagated_to_subtransforms(self): class FooHint(ResourceHint): urn = 'foo_urn' class BarHint(ResourceHint): urn = 'bar_urn' class BazHint(ResourceHint): urn = 'baz_urn' class QuxHint(ResourceHint): urn = 'qux_urn' class UseMaxValueHint(ResourceHint): urn = 'use_max_value_urn' @classmethod def get_merged_value( cls, outer_value: bytes, inner_value: bytes) -> bytes: return ResourceHint._use_max(outer_value, inner_value) ResourceHint.register_resource_hint('foo_hint', FooHint) ResourceHint.register_resource_hint('bar_hint', BarHint) ResourceHint.register_resource_hint('baz_hint', BazHint) ResourceHint.register_resource_hint('qux_hint', QuxHint) ResourceHint.register_resource_hint('use_max_value_hint', UseMaxValueHint) @beam.ptransform_fn def SubTransform(pcoll): return pcoll | beam.Map(lambda x: x + 1).with_resource_hints( foo_hint='set_on_subtransform', use_max_value_hint='10') @beam.ptransform_fn def CompositeTransform(pcoll): return pcoll | beam.Map(lambda x: x * 2) | SubTransform() p = beam.Pipeline() _ = ( p | beam.Create([1, 2]) | CompositeTransform().with_resource_hints( foo_hint='should_be_overriden_by_subtransform', bar_hint='set_on_composite', baz_hint='set_on_composite', use_max_value_hint='100')) options = PortableOptions([ '--resource_hint=baz_hint=should_be_overriden_by_composite', '--resource_hint=qux_hint=set_via_options', '--environment_type=PROCESS', '--environment_option=process_command=foo', '--sdk_location=container', ]) environment = ProcessEnvironment.from_options(options) proto = Pipeline.to_runner_api(p, default_environment=environment) for t in proto.components.transforms.values(): if "CompositeTransform/SubTransform/Map" in t.unique_name: environment = proto.components.environments.get(t.environment_id) self.assertEqual( environment.resource_hints.get('foo_urn'), b'set_on_subtransform') self.assertEqual( environment.resource_hints.get('bar_urn'), b'set_on_composite') self.assertEqual( environment.resource_hints.get('baz_urn'), b'set_on_composite') self.assertEqual( environment.resource_hints.get('qux_urn'), b'set_via_options') self.assertEqual( environment.resource_hints.get('use_max_value_urn'), b'100') found = True assert found def test_environments_with_same_resource_hints_are_reused(self): class HintX(ResourceHint): urn = 'X_urn' class HintY(ResourceHint): urn = 'Y_urn' class HintIsOdd(ResourceHint): urn = 'IsOdd_urn' ResourceHint.register_resource_hint('X', HintX) ResourceHint.register_resource_hint('Y', HintY) ResourceHint.register_resource_hint('IsOdd', HintIsOdd) p = beam.Pipeline() num_iter = 4 for i in range(num_iter): _ = ( p | f'NoHintCreate_{i}' >> beam.Create([1, 2]) | f'NoHint_{i}' >> beam.Map(lambda x: x + 1)) _ = ( p | f'XCreate_{i}' >> beam.Create([1, 2]) | f'HintX_{i}' >> beam.Map(lambda x: x + 1).with_resource_hints(X='X')) _ = ( p | f'XYCreate_{i}' >> beam.Create([1, 2]) | f'HintXY_{i}' >> beam.Map(lambda x: x + 1).with_resource_hints( X='X', Y='Y')) _ = ( p | f'IsOddCreate_{i}' >> beam.Create([1, 2]) | f'IsOdd_{i}' >> beam.Map(lambda x: x + 1).with_resource_hints(IsOdd=str(i % 2 != 0))) proto = Pipeline.to_runner_api(p) count_x = count_xy = count_is_odd = count_no_hints = 0 env_ids = set() for _, t in proto.components.transforms.items(): env = proto.components.environments[t.environment_id] if t.unique_name.startswith('HintX_'): count_x += 1 env_ids.add(t.environment_id) self.assertEqual(env.resource_hints, {'X_urn': b'X'}) if t.unique_name.startswith('HintXY_'): count_xy += 1 env_ids.add(t.environment_id) self.assertEqual(env.resource_hints, {'X_urn': b'X', 'Y_urn': b'Y'}) if t.unique_name.startswith('NoHint_'): count_no_hints += 1 env_ids.add(t.environment_id) self.assertEqual(env.resource_hints, {}) if t.unique_name.startswith('IsOdd_'): count_is_odd += 1 env_ids.add(t.environment_id) self.assertTrue( env.resource_hints == {'IsOdd_urn': b'True'} or env.resource_hints == {'IsOdd_urn': b'False'}) assert count_x == count_is_odd == count_xy == count_no_hints == num_iter assert num_iter > 1 self.assertEqual(len(env_ids), 5) def test_multiple_application_of_the_same_transform_set_different_hints(self): class FooHint(ResourceHint): urn = 'foo_urn' class UseMaxValueHint(ResourceHint): urn = 'use_max_value_urn' @classmethod def get_merged_value( cls, outer_value: bytes, inner_value: bytes) -> bytes: return ResourceHint._use_max(outer_value, inner_value) ResourceHint.register_resource_hint('foo_hint', FooHint) ResourceHint.register_resource_hint('use_max_value_hint', UseMaxValueHint) @beam.ptransform_fn def SubTransform(pcoll): return pcoll | beam.Map(lambda x: x + 1) @beam.ptransform_fn def CompositeTransform(pcoll): sub = SubTransform() return ( pcoll | 'first' >> sub.with_resource_hints(foo_hint='first_application') | 'second' >> sub.with_resource_hints(foo_hint='second_application')) p = beam.Pipeline() _ = (p | beam.Create([1, 2]) | CompositeTransform()) proto = Pipeline.to_runner_api(p) count = 0 for t in proto.components.transforms.values(): if "CompositeTransform/first/Map" in t.unique_name: environment = proto.components.environments.get(t.environment_id) self.assertEqual( b'first_application', environment.resource_hints.get('foo_urn')) count += 1 if "CompositeTransform/second/Map" in t.unique_name: environment = proto.components.environments.get(t.environment_id) self.assertEqual( b'second_application', environment.resource_hints.get('foo_urn')) count += 1 assert count == 2 def test_environments_are_deduplicated(self): def file_artifact(path, hash, staged_name): return beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2.ArtifactFilePayload( path=path, sha256=hash).SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2.ArtifactStagingToRolePayload( staged_name=staged_name).SerializeToString(), ) proto = beam_runner_api_pb2.Pipeline( components=beam_runner_api_pb2.Components( transforms={ f'transform{ix}': beam_runner_api_pb2.PTransform( environment_id=f'e{ix}') for ix in range(8) }, environments={ # Same hash and destination. 'e1': beam_runner_api_pb2.Environment( dependencies=[file_artifact('a1', 'x', 'dest')]), 'e2': beam_runner_api_pb2.Environment( dependencies=[file_artifact('a2', 'x', 'dest')]), 'e3': beam_runner_api_pb2.Environment( dependencies=[file_artifact('a3', 'y', 'dest') ]), # Different hash. 'e4': beam_runner_api_pb2.Environment( dependencies=[file_artifact('a4', 'y', 'dest2') ]), # Different destination. 'e5': beam_runner_api_pb2.Environment( dependencies=[ file_artifact('a1', 'x', 'dest'), file_artifact('b1', 'xb', 'destB') ]), # Multiple files with same hash and destinations. 'e6': beam_runner_api_pb2.Environment( dependencies=[ file_artifact('a2', 'x', 'dest'), file_artifact('b2', 'xb', 'destB') ]), 'e7': beam_runner_api_pb2.Environment( dependencies=[ file_artifact('a1', 'x', 'dest'), file_artifact('b2', 'y', 'destB') ]), # Overlapping, but not identical, files. 'e0': beam_runner_api_pb2.Environment( resource_hints={'hint': b'value'}, dependencies=[file_artifact('a1', 'x', 'dest')] ), # Same files as first, but differing other properties. })) Pipeline.merge_compatible_environments(proto) # These environments are equivalent. self.assertEqual( proto.components.transforms['transform1'].environment_id, proto.components.transforms['transform2'].environment_id) self.assertEqual( proto.components.transforms['transform5'].environment_id, proto.components.transforms['transform6'].environment_id) # These are not. self.assertNotEqual( proto.components.transforms['transform1'].environment_id, proto.components.transforms['transform3'].environment_id) self.assertNotEqual( proto.components.transforms['transform4'].environment_id, proto.components.transforms['transform3'].environment_id) self.assertNotEqual( proto.components.transforms['transform6'].environment_id, proto.components.transforms['transform7'].environment_id) self.assertNotEqual( proto.components.transforms['transform1'].environment_id, proto.components.transforms['transform0'].environment_id) self.assertEqual(len(proto.components.environments), 6) def test_multiple_outputs_composite_ptransform(self): """ Test that a composite PTransform with multiple outputs is represented correctly in the pipeline proto. """ class SalesSplitter(beam.DoFn): def process(self, element): price = element['price'] if price > 100: yield beam.pvalue.TaggedOutput('premium_sales', element) else: yield beam.pvalue.TaggedOutput('standard_sales', element) class ParentSalesSplitter(beam.PTransform): def expand(self, pcoll): return pcoll | beam.ParDo(SalesSplitter()).with_outputs( 'premium_sales', 'standard_sales') sales_data = [ { 'item': 'Laptop', 'price': 1200 }, { 'item': 'Mouse', 'price': 25 }, { 'item': 'Keyboard', 'price': 75 }, { 'item': 'Monitor', 'price': 350 }, { 'item': 'Headphones', 'price': 90 }, ] with beam.Pipeline() as pipeline: sales_records = pipeline | 'Create Sales' >> beam.Create(sales_data) _ = sales_records | 'Split Sales' >> ParentSalesSplitter() current_transforms = list(pipeline.transforms_stack) all_applied_transforms = { xform.full_label: xform for xform in current_transforms } while current_transforms: xform = current_transforms.pop() all_applied_transforms[xform.full_label] = xform current_transforms.extend(xform.parts) xform = all_applied_transforms['Split Sales'] # Confirm that Split Sales correctly has two outputs as specified by # ParDo.with_outputs in ParentSalesSplitter. assert len(xform.outputs) == 2 if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5274234 apache_beam-2.74.0/apache_beam/portability/0000755000175100017510000000000015206113160020221 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/portability/__init__.py0000644000175100017510000000152515206113032022333 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """For internal use only; no backwards-compatibility guarantees.""" ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5284233 apache_beam-2.74.0/apache_beam/portability/api/0000755000175100017510000000000015206113160020772 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/__init__.py0000644000175100017510000000523715206113151023112 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ For internal use only; no backwards-compatibility guarantees. Automatically generated when running python -m build. """ from .org.apache.beam.model.interactive import v1 from .org.apache.beam.model.job_management import v1 as v1_2 from .org.apache.beam.model.pipeline import v1 as v1_3 from .org.apache.beam.model.fn_execution import v1 as v1_4 beam_interactive_api_pb2 = v1.beam_interactive_api_pb2 beam_interactive_api_pb2_grpc = v1.beam_interactive_api_pb2_grpc beam_artifact_api_pb2_grpc = v1_2.beam_artifact_api_pb2_grpc beam_job_api_pb2_grpc = v1_2.beam_job_api_pb2_grpc beam_expansion_api_pb2_grpc = v1_2.beam_expansion_api_pb2_grpc beam_job_api_pb2 = v1_2.beam_job_api_pb2 beam_expansion_api_pb2 = v1_2.beam_expansion_api_pb2 beam_artifact_api_pb2_urns = v1_2.beam_artifact_api_pb2_urns beam_artifact_api_pb2 = v1_2.beam_artifact_api_pb2 external_transforms_pb2_urns = v1_3.external_transforms_pb2_urns beam_runner_api_pb2_grpc = v1_3.beam_runner_api_pb2_grpc schema_pb2_grpc = v1_3.schema_pb2_grpc external_transforms_pb2_grpc = v1_3.external_transforms_pb2_grpc metrics_pb2_grpc = v1_3.metrics_pb2_grpc standard_window_fns_pb2_urns = v1_3.standard_window_fns_pb2_urns beam_runner_api_pb2 = v1_3.beam_runner_api_pb2 schema_pb2 = v1_3.schema_pb2 standard_window_fns_pb2 = v1_3.standard_window_fns_pb2 metrics_pb2_urns = v1_3.metrics_pb2_urns standard_window_fns_pb2_grpc = v1_3.standard_window_fns_pb2_grpc endpoints_pb2_grpc = v1_3.endpoints_pb2_grpc schema_pb2_urns = v1_3.schema_pb2_urns external_transforms_pb2 = v1_3.external_transforms_pb2 metrics_pb2 = v1_3.metrics_pb2 endpoints_pb2 = v1_3.endpoints_pb2 beam_runner_api_pb2_urns = v1_3.beam_runner_api_pb2_urns beam_provision_api_pb2 = v1_4.beam_provision_api_pb2 beam_provision_api_pb2_grpc = v1_4.beam_provision_api_pb2_grpc beam_fn_api_pb2_grpc = v1_4.beam_fn_api_pb2_grpc beam_fn_api_pb2_urns = v1_4.beam_fn_api_pb2_urns beam_fn_api_pb2 = v1_4.beam_fn_api_pb2 ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5284233 apache_beam-2.74.0/apache_beam/portability/api/org/0000755000175100017510000000000015206113160021561 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/__init__.py0000644000175100017510000000142015206113151023667 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5294235 apache_beam-2.74.0/apache_beam/portability/api/org/apache/0000755000175100017510000000000015206113160023002 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/__init__.py0000644000175100017510000000142015206113151025110 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5294235 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/0000755000175100017510000000000015206113160023706 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/__init__.py0000644000175100017510000000142015206113151026014 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5294235 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/0000755000175100017510000000000015206113160025006 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/__init__.py0000644000175100017510000000142015206113151027114 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5294235 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/0000755000175100017510000000000015206113160027474 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/__init__.py0000644000175100017510000000142015206113151031602 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.533424 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/0000755000175100017510000000000015206113160030022 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/__init__.py0000644000175100017510000000170315206113151032134 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from . import beam_provision_api_pb2 from . import beam_provision_api_pb2_grpc from . import beam_fn_api_pb2_grpc from . import beam_fn_api_pb2_urns from . import beam_fn_api_pb2 ././@PaxHeader0000000000000000000000000000020700000000000010214 xustar00113 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_fn_api_pb2.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_fn_api_pb20000644000175100017510000010751615206113151032742 0ustar00runnerrunner# -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # source: org/apache/beam/model/fn_execution/v1/beam_fn_api.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() from ...pipeline.v1 import beam_runner_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_beam__runner__api__pb2 from ...pipeline.v1 import endpoints_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_endpoints__pb2 from ...pipeline.v1 import metrics_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_metrics__pb2 from google.protobuf import descriptor_pb2 as google_dot_protobuf_dot_descriptor__pb2 from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2 from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2 from google.protobuf import duration_pb2 as google_dot_protobuf_dot_duration__pb2 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n7org/apache/beam/model/fn_execution/v1/beam_fn_api.proto\x12%org.apache.beam.model.fn_execution.v1\x1a\x37org/apache/beam/model/pipeline/v1/beam_runner_api.proto\x1a\x31org/apache/beam/model/pipeline/v1/endpoints.proto\x1a/org/apache/beam/model/pipeline/v1/metrics.proto\x1a google/protobuf/descriptor.proto\x1a\x1cgoogle/protobuf/struct.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1egoogle/protobuf/duration.proto\"s\n\x0f\x46nApiTransforms\"`\n\x06Runner\x12,\n\x0b\x44\x41TA_SOURCE\x10\x00\x1a\x1b\xa2\xb4\xfa\xc2\x05\x15\x62\x65\x61m:runner:source:v1\x12(\n\tDATA_SINK\x10\x01\x1a\x19\xa2\xb4\xfa\xc2\x05\x13\x62\x65\x61m:runner:sink:v1\"{\n\x0eRemoteGrpcPort\x12W\n\x16\x61pi_service_descriptor\x18\x01 \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.ApiServiceDescriptor\x12\x10\n\x08\x63oder_id\x18\x02 \x01(\t\"I\n!GetProcessBundleDescriptorRequest\x12$\n\x1cprocess_bundle_descriptor_id\x18\x01 \x01(\t\"\xa3\x06\n\x12InstructionRequest\x12\x16\n\x0einstruction_id\x18\x01 \x01(\t\x12V\n\x0eprocess_bundle\x18\xe9\x07 \x01(\x0b\x32;.org.apache.beam.model.fn_execution.v1.ProcessBundleRequestH\x00\x12g\n\x17process_bundle_progress\x18\xea\x07 \x01(\x0b\x32\x43.org.apache.beam.model.fn_execution.v1.ProcessBundleProgressRequestH\x00\x12\x61\n\x14process_bundle_split\x18\xeb\x07 \x01(\x0b\x32@.org.apache.beam.model.fn_execution.v1.ProcessBundleSplitRequestH\x00\x12X\n\x0f\x66inalize_bundle\x18\xec\x07 \x01(\x0b\x32<.org.apache.beam.model.fn_execution.v1.FinalizeBundleRequestH\x00\x12\x62\n\x10monitoring_infos\x18\xed\x07 \x01(\x0b\x32\x45.org.apache.beam.model.fn_execution.v1.MonitoringInfosMetadataRequestH\x00\x12i\n\x18harness_monitoring_infos\x18\xee\x07 \x01(\x0b\x32\x44.org.apache.beam.model.fn_execution.v1.HarnessMonitoringInfosRequestH\x00\x12P\n\x0bsample_data\x18\xef\x07 \x01(\x0b\x32\x38.org.apache.beam.model.fn_execution.v1.SampleDataRequestH\x00\x12K\n\x08register\x18\xe8\x07 \x01(\x0b\x32\x36.org.apache.beam.model.fn_execution.v1.RegisterRequestH\x00\x42\t\n\x07request\"\xbc\x06\n\x13InstructionResponse\x12\x16\n\x0einstruction_id\x18\x01 \x01(\t\x12\r\n\x05\x65rror\x18\x02 \x01(\t\x12W\n\x0eprocess_bundle\x18\xe9\x07 \x01(\x0b\x32<.org.apache.beam.model.fn_execution.v1.ProcessBundleResponseH\x00\x12h\n\x17process_bundle_progress\x18\xea\x07 \x01(\x0b\x32\x44.org.apache.beam.model.fn_execution.v1.ProcessBundleProgressResponseH\x00\x12\x62\n\x14process_bundle_split\x18\xeb\x07 \x01(\x0b\x32\x41.org.apache.beam.model.fn_execution.v1.ProcessBundleSplitResponseH\x00\x12Y\n\x0f\x66inalize_bundle\x18\xec\x07 \x01(\x0b\x32=.org.apache.beam.model.fn_execution.v1.FinalizeBundleResponseH\x00\x12\x63\n\x10monitoring_infos\x18\xed\x07 \x01(\x0b\x32\x46.org.apache.beam.model.fn_execution.v1.MonitoringInfosMetadataResponseH\x00\x12j\n\x18harness_monitoring_infos\x18\xee\x07 \x01(\x0b\x32\x45.org.apache.beam.model.fn_execution.v1.HarnessMonitoringInfosResponseH\x00\x12Q\n\x0bsample_data\x18\xef\x07 \x01(\x0b\x32\x39.org.apache.beam.model.fn_execution.v1.SampleDataResponseH\x00\x12L\n\x08register\x18\xe8\x07 \x01(\x0b\x32\x37.org.apache.beam.model.fn_execution.v1.RegisterResponseH\x00\x42\n\n\x08response\",\n\x11SampleDataRequest\x12\x17\n\x0fpcollection_ids\x18\x01 \x03(\t\"\xf5\x01\n\x0eSampledElement\x12\x0f\n\x07\x65lement\x18\x01 \x01(\x0c\x12\x34\n\x10sample_timestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12R\n\texception\x18\x03 \x01(\x0b\x32?.org.apache.beam.model.fn_execution.v1.SampledElement.Exception\x1aH\n\tException\x12\x16\n\x0einstruction_id\x18\x01 \x01(\t\x12\x14\n\x0ctransform_id\x18\x02 \x01(\t\x12\r\n\x05\x65rror\x18\x03 \x01(\t\"\xd2\x02\n\x12SampleDataResponse\x12\x66\n\x0f\x65lement_samples\x18\x01 \x03(\x0b\x32M.org.apache.beam.model.fn_execution.v1.SampleDataResponse.ElementSamplesEntry\x1aV\n\x0b\x45lementList\x12G\n\x08\x65lements\x18\x01 \x03(\x0b\x32\x35.org.apache.beam.model.fn_execution.v1.SampledElement\x1a|\n\x13\x45lementSamplesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12T\n\x05value\x18\x02 \x01(\x0b\x32\x45.org.apache.beam.model.fn_execution.v1.SampleDataResponse.ElementList:\x02\x38\x01\"\x1f\n\x1dHarnessMonitoringInfosRequest\"\xcb\x01\n\x1eHarnessMonitoringInfosResponse\x12r\n\x0fmonitoring_data\x18\x01 \x03(\x0b\x32Y.org.apache.beam.model.fn_execution.v1.HarnessMonitoringInfosResponse.MonitoringDataEntry\x1a\x35\n\x13MonitoringDataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x0c:\x02\x38\x01\"t\n\x0fRegisterRequest\x12\x61\n\x19process_bundle_descriptor\x18\x01 \x03(\x0b\x32>.org.apache.beam.model.fn_execution.v1.ProcessBundleDescriptor\"\x12\n\x10RegisterResponse\"\xe1\t\n\x17ProcessBundleDescriptor\x12\n\n\x02id\x18\x01 \x01(\t\x12\x62\n\ntransforms\x18\x02 \x03(\x0b\x32N.org.apache.beam.model.fn_execution.v1.ProcessBundleDescriptor.TransformsEntry\x12\x66\n\x0cpcollections\x18\x03 \x03(\x0b\x32P.org.apache.beam.model.fn_execution.v1.ProcessBundleDescriptor.PcollectionsEntry\x12u\n\x14windowing_strategies\x18\x04 \x03(\x0b\x32W.org.apache.beam.model.fn_execution.v1.ProcessBundleDescriptor.WindowingStrategiesEntry\x12Z\n\x06\x63oders\x18\x05 \x03(\x0b\x32J.org.apache.beam.model.fn_execution.v1.ProcessBundleDescriptor.CodersEntry\x12\x66\n\x0c\x65nvironments\x18\x06 \x03(\x0b\x32P.org.apache.beam.model.fn_execution.v1.ProcessBundleDescriptor.EnvironmentsEntry\x12]\n\x1cstate_api_service_descriptor\x18\x07 \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.ApiServiceDescriptor\x12]\n\x1ctimer_api_service_descriptor\x18\x08 \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.ApiServiceDescriptor\x1a`\n\x0fTransformsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12<\n\x05value\x18\x02 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.PTransform:\x02\x38\x01\x1a\x63\n\x11PcollectionsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12=\n\x05value\x18\x02 \x01(\x0b\x32..org.apache.beam.model.pipeline.v1.PCollection:\x02\x38\x01\x1ap\n\x18WindowingStrategiesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x43\n\x05value\x18\x02 \x01(\x0b\x32\x34.org.apache.beam.model.pipeline.v1.WindowingStrategy:\x02\x38\x01\x1aW\n\x0b\x43odersEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x37\n\x05value\x18\x02 \x01(\x0b\x32(.org.apache.beam.model.pipeline.v1.Coder:\x02\x38\x01\x1a\x63\n\x11\x45nvironmentsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12=\n\x05value\x18\x02 \x01(\x0b\x32..org.apache.beam.model.pipeline.v1.Environment:\x02\x38\x01\"\xd3\x02\n\x11\x42undleApplication\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x10\n\x08input_id\x18\x02 \x01(\t\x12\x0f\n\x07\x65lement\x18\x03 \x01(\x0c\x12i\n\x11output_watermarks\x18\x04 \x03(\x0b\x32N.org.apache.beam.model.fn_execution.v1.BundleApplication.OutputWatermarksEntry\x12\x45\n\nis_bounded\x18\x05 \x01(\x0e\x32\x31.org.apache.beam.model.pipeline.v1.IsBounded.Enum\x1aS\n\x15OutputWatermarksEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12)\n\x05value\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp:\x02\x38\x01\"\xa2\x01\n\x18\x44\x65layedBundleApplication\x12M\n\x0b\x61pplication\x18\x01 \x01(\x0b\x32\x38.org.apache.beam.model.fn_execution.v1.BundleApplication\x12\x37\n\x14requested_time_delay\x18\x02 \x01(\x0b\x32\x19.google.protobuf.Duration\"\xce\x04\n\x14ProcessBundleRequest\x12$\n\x1cprocess_bundle_descriptor_id\x18\x01 \x01(\t\x12\\\n\x0c\x63\x61\x63he_tokens\x18\x02 \x03(\x0b\x32\x46.org.apache.beam.model.fn_execution.v1.ProcessBundleRequest.CacheToken\x12\x41\n\x08\x65lements\x18\x03 \x01(\x0b\x32/.org.apache.beam.model.fn_execution.v1.Elements\x12\x14\n\x0chas_no_state\x18\x04 \x01(\x08\x12\x1c\n\x14only_bundle_for_keys\x18\x05 \x01(\x08\x1a\xba\x02\n\nCacheToken\x12\x66\n\nuser_state\x18\x01 \x01(\x0b\x32P.org.apache.beam.model.fn_execution.v1.ProcessBundleRequest.CacheToken.UserStateH\x00\x12\x66\n\nside_input\x18\x02 \x01(\x0b\x32P.org.apache.beam.model.fn_execution.v1.ProcessBundleRequest.CacheToken.SideInputH\x00\x12\r\n\x05token\x18\n \x01(\x0c\x1a\x0b\n\tUserState\x1a\x38\n\tSideInput\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x15\n\rside_input_id\x18\x02 \x01(\tB\x06\n\x04type\"\xc7\x03\n\x15ProcessBundleResponse\x12W\n\x0eresidual_roots\x18\x02 \x03(\x0b\x32?.org.apache.beam.model.fn_execution.v1.DelayedBundleApplication\x12K\n\x10monitoring_infos\x18\x03 \x03(\x0b\x32\x31.org.apache.beam.model.pipeline.v1.MonitoringInfo\x12\x1d\n\x15requires_finalization\x18\x04 \x01(\x08\x12i\n\x0fmonitoring_data\x18\x05 \x03(\x0b\x32P.org.apache.beam.model.fn_execution.v1.ProcessBundleResponse.MonitoringDataEntry\x12\x41\n\x08\x65lements\x18\x06 \x01(\x0b\x32/.org.apache.beam.model.fn_execution.v1.Elements\x1a\x35\n\x13MonitoringDataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x0c:\x02\x38\x01J\x04\x08\x01\x10\x02\"6\n\x1cProcessBundleProgressRequest\x12\x16\n\x0einstruction_id\x18\x01 \x01(\t\"<\n\x1eMonitoringInfosMetadataRequest\x12\x1a\n\x12monitoring_info_id\x18\x01 \x03(\t\"\xea\x02\n\x1dProcessBundleProgressResponse\x12K\n\x10monitoring_infos\x18\x03 \x03(\x0b\x32\x31.org.apache.beam.model.pipeline.v1.MonitoringInfo\x12q\n\x0fmonitoring_data\x18\x05 \x03(\x0b\x32X.org.apache.beam.model.fn_execution.v1.ProcessBundleProgressResponse.MonitoringDataEntry\x12$\n\x17\x63onsuming_received_data\x18\x06 \x01(\x08H\x00\x88\x01\x01\x1a\x35\n\x13MonitoringDataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x0c:\x02\x38\x01\x42\x1a\n\x18_consuming_received_dataJ\x04\x08\x01\x10\x02J\x04\x08\x02\x10\x03J\x04\x08\x04\x10\x05\"\x80\x02\n\x1fMonitoringInfosMetadataResponse\x12s\n\x0fmonitoring_info\x18\x01 \x03(\x0b\x32Z.org.apache.beam.model.fn_execution.v1.MonitoringInfosMetadataResponse.MonitoringInfoEntry\x1ah\n\x13MonitoringInfoEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12@\n\x05value\x18\x02 \x01(\x0b\x32\x31.org.apache.beam.model.pipeline.v1.MonitoringInfo:\x02\x38\x01\"\x95\x03\n\x19ProcessBundleSplitRequest\x12\x16\n\x0einstruction_id\x18\x01 \x01(\t\x12k\n\x0e\x64\x65sired_splits\x18\x03 \x03(\x0b\x32S.org.apache.beam.model.fn_execution.v1.ProcessBundleSplitRequest.DesiredSplitsEntry\x1am\n\x0c\x44\x65siredSplit\x12\x1d\n\x15\x66raction_of_remainder\x18\x01 \x01(\x01\x12\x1c\n\x14\x61llowed_split_points\x18\x03 \x03(\x03\x12 \n\x18\x65stimated_input_elements\x18\x02 \x01(\x03\x1a\x83\x01\n\x12\x44\x65siredSplitsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\\\n\x05value\x18\x02 \x01(\x0b\x32M.org.apache.beam.model.fn_execution.v1.ProcessBundleSplitRequest.DesiredSplit:\x02\x38\x01\"\x92\x03\n\x1aProcessBundleSplitResponse\x12O\n\rprimary_roots\x18\x01 \x03(\x0b\x32\x38.org.apache.beam.model.fn_execution.v1.BundleApplication\x12W\n\x0eresidual_roots\x18\x02 \x03(\x0b\x32?.org.apache.beam.model.fn_execution.v1.DelayedBundleApplication\x12\x66\n\x0e\x63hannel_splits\x18\x03 \x03(\x0b\x32N.org.apache.beam.model.fn_execution.v1.ProcessBundleSplitResponse.ChannelSplit\x1a\x62\n\x0c\x43hannelSplit\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x1c\n\x14last_primary_element\x18\x02 \x01(\x03\x12\x1e\n\x16\x66irst_residual_element\x18\x03 \x01(\x03\"/\n\x15\x46inalizeBundleRequest\x12\x16\n\x0einstruction_id\x18\x01 \x01(\t\"\x18\n\x16\x46inalizeBundleResponse\"\xc1\x06\n\x08\x45lements\x12\x42\n\x04\x64\x61ta\x18\x01 \x03(\x0b\x32\x34.org.apache.beam.model.fn_execution.v1.Elements.Data\x12\x46\n\x06timers\x18\x02 \x03(\x0b\x32\x36.org.apache.beam.model.fn_execution.v1.Elements.Timers\x1aW\n\x04\x44\x61ta\x12\x16\n\x0einstruction_id\x18\x01 \x01(\t\x12\x14\n\x0ctransform_id\x18\x02 \x01(\t\x12\x10\n\x04\x64\x61ta\x18\x03 \x01(\x0c\x42\x02\x08\x01\x12\x0f\n\x07is_last\x18\x04 \x01(\x08\x1a\x44\n\tDrainMode\"7\n\x04\x45num\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x10\n\x0cNOT_DRAINING\x10\x01\x12\x0c\n\x08\x44RAINING\x10\x02\x1al\n\tValueKind\"_\n\x04\x45num\x12\x1a\n\x16VALUE_KIND_UNSPECIFIED\x10\x00\x12\n\n\x06INSERT\x10\x01\x12\x11\n\rUPDATE_BEFORE\x10\x02\x12\x10\n\x0cUPDATE_AFTER\x10\x03\x12\n\n\x06\x44\x45LETE\x10\x04\x1a\xa9\x02\n\x0f\x45lementMetadata\x12R\n\x05\x64rain\x18\x01 \x01(\x0e\x32>.org.apache.beam.model.fn_execution.v1.Elements.DrainMode.EnumH\x00\x88\x01\x01\x12\x18\n\x0btraceparent\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x17\n\ntracestate\x18\x03 \x01(\tH\x02\x88\x01\x01\x12W\n\nvalue_kind\x18\x04 \x01(\x0e\x32>.org.apache.beam.model.fn_execution.v1.Elements.ValueKind.EnumH\x03\x88\x01\x01\x42\x08\n\x06_drainB\x0e\n\x0c_traceparentB\r\n\x0b_tracestateB\r\n\x0b_value_kind\x1ap\n\x06Timers\x12\x16\n\x0einstruction_id\x18\x01 \x01(\t\x12\x14\n\x0ctransform_id\x18\x02 \x01(\t\x12\x17\n\x0ftimer_family_id\x18\x03 \x01(\t\x12\x0e\n\x06timers\x18\x04 \x01(\x0c\x12\x0f\n\x07is_last\x18\x05 \x01(\x08\"\xe3\x02\n\x0cStateRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x16\n\x0einstruction_id\x18\x02 \x01(\t\x12\x42\n\tstate_key\x18\x03 \x01(\x0b\x32/.org.apache.beam.model.fn_execution.v1.StateKey\x12\x46\n\x03get\x18\xe8\x07 \x01(\x0b\x32\x36.org.apache.beam.model.fn_execution.v1.StateGetRequestH\x00\x12L\n\x06\x61ppend\x18\xe9\x07 \x01(\x0b\x32\x39.org.apache.beam.model.fn_execution.v1.StateAppendRequestH\x00\x12J\n\x05\x63lear\x18\xea\x07 \x01(\x0b\x32\x38.org.apache.beam.model.fn_execution.v1.StateClearRequestH\x00\x42\t\n\x07request\"\x9b\x02\n\rStateResponse\x12\n\n\x02id\x18\x01 \x01(\t\x12\r\n\x05\x65rror\x18\x02 \x01(\t\x12G\n\x03get\x18\xe8\x07 \x01(\x0b\x32\x37.org.apache.beam.model.fn_execution.v1.StateGetResponseH\x00\x12M\n\x06\x61ppend\x18\xe9\x07 \x01(\x0b\x32:.org.apache.beam.model.fn_execution.v1.StateAppendResponseH\x00\x12K\n\x05\x63lear\x18\xea\x07 \x01(\x0b\x32\x39.org.apache.beam.model.fn_execution.v1.StateClearResponseH\x00\x42\n\n\x08response\"\xba\x0f\n\x08StateKey\x12H\n\x06runner\x18\x01 \x01(\x0b\x32\x36.org.apache.beam.model.fn_execution.v1.StateKey.RunnerH\x00\x12`\n\x13multimap_side_input\x18\x02 \x01(\x0b\x32\x41.org.apache.beam.model.fn_execution.v1.StateKey.MultimapSideInputH\x00\x12V\n\x0e\x62\x61g_user_state\x18\x03 \x01(\x0b\x32<.org.apache.beam.model.fn_execution.v1.StateKey.BagUserStateH\x00\x12`\n\x13iterable_side_input\x18\x04 \x01(\x0b\x32\x41.org.apache.beam.model.fn_execution.v1.StateKey.IterableSideInputH\x00\x12i\n\x18multimap_keys_side_input\x18\x05 \x01(\x0b\x32\x45.org.apache.beam.model.fn_execution.v1.StateKey.MultimapKeysSideInputH\x00\x12v\n\x1fmultimap_keys_values_side_input\x18\x08 \x01(\x0b\x32K.org.apache.beam.model.fn_execution.v1.StateKey.MultimapKeysValuesSideInputH\x00\x12i\n\x18multimap_keys_user_state\x18\x06 \x01(\x0b\x32\x45.org.apache.beam.model.fn_execution.v1.StateKey.MultimapKeysUserStateH\x00\x12o\n\x1bmultimap_entries_user_state\x18\n \x01(\x0b\x32H.org.apache.beam.model.fn_execution.v1.StateKey.MultimapEntriesUserStateH\x00\x12`\n\x13multimap_user_state\x18\x07 \x01(\x0b\x32\x41.org.apache.beam.model.fn_execution.v1.StateKey.MultimapUserStateH\x00\x12g\n\x17ordered_list_user_state\x18\t \x01(\x0b\x32\x44.org.apache.beam.model.fn_execution.v1.StateKey.OrderedListUserStateH\x00\x1a\x15\n\x06Runner\x12\x0b\n\x03key\x18\x01 \x01(\x0c\x1aP\n\x11IterableSideInput\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x15\n\rside_input_id\x18\x02 \x01(\t\x12\x0e\n\x06window\x18\x03 \x01(\x0c\x1a]\n\x11MultimapSideInput\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x15\n\rside_input_id\x18\x02 \x01(\t\x12\x0e\n\x06window\x18\x03 \x01(\x0c\x12\x0b\n\x03key\x18\x04 \x01(\x0c\x1aT\n\x15MultimapKeysSideInput\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x15\n\rside_input_id\x18\x02 \x01(\t\x12\x0e\n\x06window\x18\x03 \x01(\x0c\x1aZ\n\x1bMultimapKeysValuesSideInput\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x15\n\rside_input_id\x18\x02 \x01(\t\x12\x0e\n\x06window\x18\x03 \x01(\x0c\x1aX\n\x0c\x42\x61gUserState\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x15\n\ruser_state_id\x18\x02 \x01(\t\x12\x0e\n\x06window\x18\x03 \x01(\x0c\x12\x0b\n\x03key\x18\x04 \x01(\x0c\x1a\x61\n\x15MultimapKeysUserState\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x15\n\ruser_state_id\x18\x02 \x01(\t\x12\x0e\n\x06window\x18\x03 \x01(\x0c\x12\x0b\n\x03key\x18\x04 \x01(\x0c\x1a\x64\n\x18MultimapEntriesUserState\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x15\n\ruser_state_id\x18\x02 \x01(\t\x12\x0e\n\x06window\x18\x03 \x01(\x0c\x12\x0b\n\x03key\x18\x04 \x01(\x0c\x1an\n\x11MultimapUserState\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x15\n\ruser_state_id\x18\x02 \x01(\t\x12\x0e\n\x06window\x18\x03 \x01(\x0c\x12\x0b\n\x03key\x18\x04 \x01(\x0c\x12\x0f\n\x07map_key\x18\x05 \x01(\x0c\x1a\xa8\x01\n\x14OrderedListUserState\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x15\n\ruser_state_id\x18\x02 \x01(\t\x12\x0e\n\x06window\x18\x03 \x01(\x0c\x12\x0b\n\x03key\x18\x04 \x01(\x0c\x12\x46\n\x05range\x18\x05 \x01(\x0b\x32\x37.org.apache.beam.model.fn_execution.v1.OrderedListRangeB\x06\n\x04type\"-\n\x0fStateGetRequest\x12\x1a\n\x12\x63ontinuation_token\x18\x01 \x01(\x0c\"<\n\x10StateGetResponse\x12\x1a\n\x12\x63ontinuation_token\x18\x01 \x01(\x0c\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"\"\n\x12StateAppendRequest\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\"\x15\n\x13StateAppendResponse\"\x13\n\x11StateClearRequest\"\x14\n\x12StateClearResponse\".\n\x10OrderedListRange\x12\r\n\x05start\x18\x01 \x01(\x03\x12\x0b\n\x03\x65nd\x18\x02 \x01(\x03\"\xee\x03\n\x08LogEntry\x12O\n\x08severity\x18\x01 \x01(\x0e\x32=.org.apache.beam.model.fn_execution.v1.LogEntry.Severity.Enum\x12-\n\ttimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0f\n\x07message\x18\x03 \x01(\t\x12\r\n\x05trace\x18\x04 \x01(\t\x12\x16\n\x0einstruction_id\x18\x05 \x01(\t\x12\x14\n\x0ctransform_id\x18\x06 \x01(\t\x12\x14\n\x0clog_location\x18\x07 \x01(\t\x12\x0e\n\x06thread\x18\x08 \x01(\t\x12,\n\x0b\x63ustom_data\x18\t \x01(\x0b\x32\x17.google.protobuf.Struct\x1aL\n\x04List\x12\x44\n\x0blog_entries\x18\x01 \x03(\x0b\x32/.org.apache.beam.model.fn_execution.v1.LogEntry\x1ar\n\x08Severity\"f\n\x04\x45num\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\t\n\x05TRACE\x10\x01\x12\t\n\x05\x44\x45\x42UG\x10\x02\x12\x08\n\x04INFO\x10\x03\x12\n\n\x06NOTICE\x10\x04\x12\x08\n\x04WARN\x10\x05\x12\t\n\x05\x45RROR\x10\x06\x12\x0c\n\x08\x43RITICAL\x10\x07\"\x0c\n\nLogControl\"\xfc\x03\n\x12StartWorkerRequest\x12\x11\n\tworker_id\x18\x01 \x01(\t\x12Q\n\x10\x63ontrol_endpoint\x18\x02 \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.ApiServiceDescriptor\x12Q\n\x10logging_endpoint\x18\x03 \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.ApiServiceDescriptor\x12R\n\x11\x61rtifact_endpoint\x18\x04 \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.ApiServiceDescriptor\x12S\n\x12provision_endpoint\x18\x05 \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.ApiServiceDescriptor\x12U\n\x06params\x18\n \x03(\x0b\x32\x45.org.apache.beam.model.fn_execution.v1.StartWorkerRequest.ParamsEntry\x1a-\n\x0bParamsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"$\n\x13StartWorkerResponse\x12\r\n\x05\x65rror\x18\x01 \x01(\t\"&\n\x11StopWorkerRequest\x12\x11\n\tworker_id\x18\x01 \x01(\t\"#\n\x12StopWorkerResponse\x12\r\n\x05\x65rror\x18\x01 \x01(\t\"!\n\x13WorkerStatusRequest\x12\n\n\x02id\x18\x01 \x01(\t\"F\n\x14WorkerStatusResponse\x12\n\n\x02id\x18\x01 \x01(\t\x12\r\n\x05\x65rror\x18\x02 \x01(\t\x12\x13\n\x0bstatus_info\x18\x03 \x01(\t2\xc3\x02\n\rBeamFnControl\x12\x86\x01\n\x07\x43ontrol\x12:.org.apache.beam.model.fn_execution.v1.InstructionResponse\x1a\x39.org.apache.beam.model.fn_execution.v1.InstructionRequest\"\x00(\x01\x30\x01\x12\xa8\x01\n\x1aGetProcessBundleDescriptor\x12H.org.apache.beam.model.fn_execution.v1.GetProcessBundleDescriptorRequest\x1a>.org.apache.beam.model.fn_execution.v1.ProcessBundleDescriptor\"\x00\x32|\n\nBeamFnData\x12n\n\x04\x44\x61ta\x12/.org.apache.beam.model.fn_execution.v1.Elements\x1a/.org.apache.beam.model.fn_execution.v1.Elements\"\x00(\x01\x30\x01\x32\x87\x01\n\x0b\x42\x65\x61mFnState\x12x\n\x05State\x12\x33.org.apache.beam.model.fn_execution.v1.StateRequest\x1a\x34.org.apache.beam.model.fn_execution.v1.StateResponse\"\x00(\x01\x30\x01\x32\x89\x01\n\rBeamFnLogging\x12x\n\x07Logging\x12\x34.org.apache.beam.model.fn_execution.v1.LogEntry.List\x1a\x31.org.apache.beam.model.fn_execution.v1.LogControl\"\x00(\x01\x30\x01\x32\xa9\x02\n\x18\x42\x65\x61mFnExternalWorkerPool\x12\x86\x01\n\x0bStartWorker\x12\x39.org.apache.beam.model.fn_execution.v1.StartWorkerRequest\x1a:.org.apache.beam.model.fn_execution.v1.StartWorkerResponse\"\x00\x12\x83\x01\n\nStopWorker\x12\x38.org.apache.beam.model.fn_execution.v1.StopWorkerRequest\x1a\x39.org.apache.beam.model.fn_execution.v1.StopWorkerResponse\"\x00\x32\xa4\x01\n\x12\x42\x65\x61mFnWorkerStatus\x12\x8d\x01\n\x0cWorkerStatus\x12;.org.apache.beam.model.fn_execution.v1.WorkerStatusResponse\x1a:.org.apache.beam.model.fn_execution.v1.WorkerStatusRequest\"\x00(\x01\x30\x01\x42\x81\x01\n$org.apache.beam.model.fnexecution.v1B\tBeamFnApiZNgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1;fnexecution_v1b\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'org.apache.beam.model.fn_execution.v1.beam_fn_api_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n$org.apache.beam.model.fnexecution.v1B\tBeamFnApiZNgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1;fnexecution_v1' _globals['_FNAPITRANSFORMS_RUNNER'].values_by_name["DATA_SOURCE"]._options = None _globals['_FNAPITRANSFORMS_RUNNER'].values_by_name["DATA_SOURCE"]._serialized_options = b'\242\264\372\302\005\025beam:runner:source:v1' _globals['_FNAPITRANSFORMS_RUNNER'].values_by_name["DATA_SINK"]._options = None _globals['_FNAPITRANSFORMS_RUNNER'].values_by_name["DATA_SINK"]._serialized_options = b'\242\264\372\302\005\023beam:runner:sink:v1' _globals['_SAMPLEDATARESPONSE_ELEMENTSAMPLESENTRY']._options = None _globals['_SAMPLEDATARESPONSE_ELEMENTSAMPLESENTRY']._serialized_options = b'8\001' _globals['_HARNESSMONITORINGINFOSRESPONSE_MONITORINGDATAENTRY']._options = None _globals['_HARNESSMONITORINGINFOSRESPONSE_MONITORINGDATAENTRY']._serialized_options = b'8\001' _globals['_PROCESSBUNDLEDESCRIPTOR_TRANSFORMSENTRY']._options = None _globals['_PROCESSBUNDLEDESCRIPTOR_TRANSFORMSENTRY']._serialized_options = b'8\001' _globals['_PROCESSBUNDLEDESCRIPTOR_PCOLLECTIONSENTRY']._options = None _globals['_PROCESSBUNDLEDESCRIPTOR_PCOLLECTIONSENTRY']._serialized_options = b'8\001' _globals['_PROCESSBUNDLEDESCRIPTOR_WINDOWINGSTRATEGIESENTRY']._options = None _globals['_PROCESSBUNDLEDESCRIPTOR_WINDOWINGSTRATEGIESENTRY']._serialized_options = b'8\001' _globals['_PROCESSBUNDLEDESCRIPTOR_CODERSENTRY']._options = None _globals['_PROCESSBUNDLEDESCRIPTOR_CODERSENTRY']._serialized_options = b'8\001' _globals['_PROCESSBUNDLEDESCRIPTOR_ENVIRONMENTSENTRY']._options = None _globals['_PROCESSBUNDLEDESCRIPTOR_ENVIRONMENTSENTRY']._serialized_options = b'8\001' _globals['_BUNDLEAPPLICATION_OUTPUTWATERMARKSENTRY']._options = None _globals['_BUNDLEAPPLICATION_OUTPUTWATERMARKSENTRY']._serialized_options = b'8\001' _globals['_PROCESSBUNDLERESPONSE_MONITORINGDATAENTRY']._options = None _globals['_PROCESSBUNDLERESPONSE_MONITORINGDATAENTRY']._serialized_options = b'8\001' _globals['_PROCESSBUNDLEPROGRESSRESPONSE_MONITORINGDATAENTRY']._options = None _globals['_PROCESSBUNDLEPROGRESSRESPONSE_MONITORINGDATAENTRY']._serialized_options = b'8\001' _globals['_MONITORINGINFOSMETADATARESPONSE_MONITORINGINFOENTRY']._options = None _globals['_MONITORINGINFOSMETADATARESPONSE_MONITORINGINFOENTRY']._serialized_options = b'8\001' _globals['_PROCESSBUNDLESPLITREQUEST_DESIREDSPLITSENTRY']._options = None _globals['_PROCESSBUNDLESPLITREQUEST_DESIREDSPLITSENTRY']._serialized_options = b'8\001' _globals['_ELEMENTS_DATA'].fields_by_name['data']._options = None _globals['_ELEMENTS_DATA'].fields_by_name['data']._serialized_options = b'\010\001' _globals['_STARTWORKERREQUEST_PARAMSENTRY']._options = None _globals['_STARTWORKERREQUEST_PARAMSENTRY']._serialized_options = b'8\001' _globals['_FNAPITRANSFORMS']._serialized_start=384 _globals['_FNAPITRANSFORMS']._serialized_end=499 _globals['_FNAPITRANSFORMS_RUNNER']._serialized_start=403 _globals['_FNAPITRANSFORMS_RUNNER']._serialized_end=499 _globals['_REMOTEGRPCPORT']._serialized_start=501 _globals['_REMOTEGRPCPORT']._serialized_end=624 _globals['_GETPROCESSBUNDLEDESCRIPTORREQUEST']._serialized_start=626 _globals['_GETPROCESSBUNDLEDESCRIPTORREQUEST']._serialized_end=699 _globals['_INSTRUCTIONREQUEST']._serialized_start=702 _globals['_INSTRUCTIONREQUEST']._serialized_end=1505 _globals['_INSTRUCTIONRESPONSE']._serialized_start=1508 _globals['_INSTRUCTIONRESPONSE']._serialized_end=2336 _globals['_SAMPLEDATAREQUEST']._serialized_start=2338 _globals['_SAMPLEDATAREQUEST']._serialized_end=2382 _globals['_SAMPLEDELEMENT']._serialized_start=2385 _globals['_SAMPLEDELEMENT']._serialized_end=2630 _globals['_SAMPLEDELEMENT_EXCEPTION']._serialized_start=2558 _globals['_SAMPLEDELEMENT_EXCEPTION']._serialized_end=2630 _globals['_SAMPLEDATARESPONSE']._serialized_start=2633 _globals['_SAMPLEDATARESPONSE']._serialized_end=2971 _globals['_SAMPLEDATARESPONSE_ELEMENTLIST']._serialized_start=2759 _globals['_SAMPLEDATARESPONSE_ELEMENTLIST']._serialized_end=2845 _globals['_SAMPLEDATARESPONSE_ELEMENTSAMPLESENTRY']._serialized_start=2847 _globals['_SAMPLEDATARESPONSE_ELEMENTSAMPLESENTRY']._serialized_end=2971 _globals['_HARNESSMONITORINGINFOSREQUEST']._serialized_start=2973 _globals['_HARNESSMONITORINGINFOSREQUEST']._serialized_end=3004 _globals['_HARNESSMONITORINGINFOSRESPONSE']._serialized_start=3007 _globals['_HARNESSMONITORINGINFOSRESPONSE']._serialized_end=3210 _globals['_HARNESSMONITORINGINFOSRESPONSE_MONITORINGDATAENTRY']._serialized_start=3157 _globals['_HARNESSMONITORINGINFOSRESPONSE_MONITORINGDATAENTRY']._serialized_end=3210 _globals['_REGISTERREQUEST']._serialized_start=3212 _globals['_REGISTERREQUEST']._serialized_end=3328 _globals['_REGISTERRESPONSE']._serialized_start=3330 _globals['_REGISTERRESPONSE']._serialized_end=3348 _globals['_PROCESSBUNDLEDESCRIPTOR']._serialized_start=3351 _globals['_PROCESSBUNDLEDESCRIPTOR']._serialized_end=4600 _globals['_PROCESSBUNDLEDESCRIPTOR_TRANSFORMSENTRY']._serialized_start=4099 _globals['_PROCESSBUNDLEDESCRIPTOR_TRANSFORMSENTRY']._serialized_end=4195 _globals['_PROCESSBUNDLEDESCRIPTOR_PCOLLECTIONSENTRY']._serialized_start=4197 _globals['_PROCESSBUNDLEDESCRIPTOR_PCOLLECTIONSENTRY']._serialized_end=4296 _globals['_PROCESSBUNDLEDESCRIPTOR_WINDOWINGSTRATEGIESENTRY']._serialized_start=4298 _globals['_PROCESSBUNDLEDESCRIPTOR_WINDOWINGSTRATEGIESENTRY']._serialized_end=4410 _globals['_PROCESSBUNDLEDESCRIPTOR_CODERSENTRY']._serialized_start=4412 _globals['_PROCESSBUNDLEDESCRIPTOR_CODERSENTRY']._serialized_end=4499 _globals['_PROCESSBUNDLEDESCRIPTOR_ENVIRONMENTSENTRY']._serialized_start=4501 _globals['_PROCESSBUNDLEDESCRIPTOR_ENVIRONMENTSENTRY']._serialized_end=4600 _globals['_BUNDLEAPPLICATION']._serialized_start=4603 _globals['_BUNDLEAPPLICATION']._serialized_end=4942 _globals['_BUNDLEAPPLICATION_OUTPUTWATERMARKSENTRY']._serialized_start=4859 _globals['_BUNDLEAPPLICATION_OUTPUTWATERMARKSENTRY']._serialized_end=4942 _globals['_DELAYEDBUNDLEAPPLICATION']._serialized_start=4945 _globals['_DELAYEDBUNDLEAPPLICATION']._serialized_end=5107 _globals['_PROCESSBUNDLEREQUEST']._serialized_start=5110 _globals['_PROCESSBUNDLEREQUEST']._serialized_end=5700 _globals['_PROCESSBUNDLEREQUEST_CACHETOKEN']._serialized_start=5386 _globals['_PROCESSBUNDLEREQUEST_CACHETOKEN']._serialized_end=5700 _globals['_PROCESSBUNDLEREQUEST_CACHETOKEN_USERSTATE']._serialized_start=5623 _globals['_PROCESSBUNDLEREQUEST_CACHETOKEN_USERSTATE']._serialized_end=5634 _globals['_PROCESSBUNDLEREQUEST_CACHETOKEN_SIDEINPUT']._serialized_start=5636 _globals['_PROCESSBUNDLEREQUEST_CACHETOKEN_SIDEINPUT']._serialized_end=5692 _globals['_PROCESSBUNDLERESPONSE']._serialized_start=5703 _globals['_PROCESSBUNDLERESPONSE']._serialized_end=6158 _globals['_PROCESSBUNDLERESPONSE_MONITORINGDATAENTRY']._serialized_start=3157 _globals['_PROCESSBUNDLERESPONSE_MONITORINGDATAENTRY']._serialized_end=3210 _globals['_PROCESSBUNDLEPROGRESSREQUEST']._serialized_start=6160 _globals['_PROCESSBUNDLEPROGRESSREQUEST']._serialized_end=6214 _globals['_MONITORINGINFOSMETADATAREQUEST']._serialized_start=6216 _globals['_MONITORINGINFOSMETADATAREQUEST']._serialized_end=6276 _globals['_PROCESSBUNDLEPROGRESSRESPONSE']._serialized_start=6279 _globals['_PROCESSBUNDLEPROGRESSRESPONSE']._serialized_end=6641 _globals['_PROCESSBUNDLEPROGRESSRESPONSE_MONITORINGDATAENTRY']._serialized_start=3157 _globals['_PROCESSBUNDLEPROGRESSRESPONSE_MONITORINGDATAENTRY']._serialized_end=3210 _globals['_MONITORINGINFOSMETADATARESPONSE']._serialized_start=6644 _globals['_MONITORINGINFOSMETADATARESPONSE']._serialized_end=6900 _globals['_MONITORINGINFOSMETADATARESPONSE_MONITORINGINFOENTRY']._serialized_start=6796 _globals['_MONITORINGINFOSMETADATARESPONSE_MONITORINGINFOENTRY']._serialized_end=6900 _globals['_PROCESSBUNDLESPLITREQUEST']._serialized_start=6903 _globals['_PROCESSBUNDLESPLITREQUEST']._serialized_end=7308 _globals['_PROCESSBUNDLESPLITREQUEST_DESIREDSPLIT']._serialized_start=7065 _globals['_PROCESSBUNDLESPLITREQUEST_DESIREDSPLIT']._serialized_end=7174 _globals['_PROCESSBUNDLESPLITREQUEST_DESIREDSPLITSENTRY']._serialized_start=7177 _globals['_PROCESSBUNDLESPLITREQUEST_DESIREDSPLITSENTRY']._serialized_end=7308 _globals['_PROCESSBUNDLESPLITRESPONSE']._serialized_start=7311 _globals['_PROCESSBUNDLESPLITRESPONSE']._serialized_end=7713 _globals['_PROCESSBUNDLESPLITRESPONSE_CHANNELSPLIT']._serialized_start=7615 _globals['_PROCESSBUNDLESPLITRESPONSE_CHANNELSPLIT']._serialized_end=7713 _globals['_FINALIZEBUNDLEREQUEST']._serialized_start=7715 _globals['_FINALIZEBUNDLEREQUEST']._serialized_end=7762 _globals['_FINALIZEBUNDLERESPONSE']._serialized_start=7764 _globals['_FINALIZEBUNDLERESPONSE']._serialized_end=7788 _globals['_ELEMENTS']._serialized_start=7791 _globals['_ELEMENTS']._serialized_end=8624 _globals['_ELEMENTS_DATA']._serialized_start=7943 _globals['_ELEMENTS_DATA']._serialized_end=8030 _globals['_ELEMENTS_DRAINMODE']._serialized_start=8032 _globals['_ELEMENTS_DRAINMODE']._serialized_end=8100 _globals['_ELEMENTS_DRAINMODE_ENUM']._serialized_start=8045 _globals['_ELEMENTS_DRAINMODE_ENUM']._serialized_end=8100 _globals['_ELEMENTS_VALUEKIND']._serialized_start=8102 _globals['_ELEMENTS_VALUEKIND']._serialized_end=8210 _globals['_ELEMENTS_VALUEKIND_ENUM']._serialized_start=8115 _globals['_ELEMENTS_VALUEKIND_ENUM']._serialized_end=8210 _globals['_ELEMENTS_ELEMENTMETADATA']._serialized_start=8213 _globals['_ELEMENTS_ELEMENTMETADATA']._serialized_end=8510 _globals['_ELEMENTS_TIMERS']._serialized_start=8512 _globals['_ELEMENTS_TIMERS']._serialized_end=8624 _globals['_STATEREQUEST']._serialized_start=8627 _globals['_STATEREQUEST']._serialized_end=8982 _globals['_STATERESPONSE']._serialized_start=8985 _globals['_STATERESPONSE']._serialized_end=9268 _globals['_STATEKEY']._serialized_start=9271 _globals['_STATEKEY']._serialized_end=11249 _globals['_STATEKEY_RUNNER']._serialized_start=10291 _globals['_STATEKEY_RUNNER']._serialized_end=10312 _globals['_STATEKEY_ITERABLESIDEINPUT']._serialized_start=10314 _globals['_STATEKEY_ITERABLESIDEINPUT']._serialized_end=10394 _globals['_STATEKEY_MULTIMAPSIDEINPUT']._serialized_start=10396 _globals['_STATEKEY_MULTIMAPSIDEINPUT']._serialized_end=10489 _globals['_STATEKEY_MULTIMAPKEYSSIDEINPUT']._serialized_start=10491 _globals['_STATEKEY_MULTIMAPKEYSSIDEINPUT']._serialized_end=10575 _globals['_STATEKEY_MULTIMAPKEYSVALUESSIDEINPUT']._serialized_start=10577 _globals['_STATEKEY_MULTIMAPKEYSVALUESSIDEINPUT']._serialized_end=10667 _globals['_STATEKEY_BAGUSERSTATE']._serialized_start=10669 _globals['_STATEKEY_BAGUSERSTATE']._serialized_end=10757 _globals['_STATEKEY_MULTIMAPKEYSUSERSTATE']._serialized_start=10759 _globals['_STATEKEY_MULTIMAPKEYSUSERSTATE']._serialized_end=10856 _globals['_STATEKEY_MULTIMAPENTRIESUSERSTATE']._serialized_start=10858 _globals['_STATEKEY_MULTIMAPENTRIESUSERSTATE']._serialized_end=10958 _globals['_STATEKEY_MULTIMAPUSERSTATE']._serialized_start=10960 _globals['_STATEKEY_MULTIMAPUSERSTATE']._serialized_end=11070 _globals['_STATEKEY_ORDEREDLISTUSERSTATE']._serialized_start=11073 _globals['_STATEKEY_ORDEREDLISTUSERSTATE']._serialized_end=11241 _globals['_STATEGETREQUEST']._serialized_start=11251 _globals['_STATEGETREQUEST']._serialized_end=11296 _globals['_STATEGETRESPONSE']._serialized_start=11298 _globals['_STATEGETRESPONSE']._serialized_end=11358 _globals['_STATEAPPENDREQUEST']._serialized_start=11360 _globals['_STATEAPPENDREQUEST']._serialized_end=11394 _globals['_STATEAPPENDRESPONSE']._serialized_start=11396 _globals['_STATEAPPENDRESPONSE']._serialized_end=11417 _globals['_STATECLEARREQUEST']._serialized_start=11419 _globals['_STATECLEARREQUEST']._serialized_end=11438 _globals['_STATECLEARRESPONSE']._serialized_start=11440 _globals['_STATECLEARRESPONSE']._serialized_end=11460 _globals['_ORDEREDLISTRANGE']._serialized_start=11462 _globals['_ORDEREDLISTRANGE']._serialized_end=11508 _globals['_LOGENTRY']._serialized_start=11511 _globals['_LOGENTRY']._serialized_end=12005 _globals['_LOGENTRY_LIST']._serialized_start=11813 _globals['_LOGENTRY_LIST']._serialized_end=11889 _globals['_LOGENTRY_SEVERITY']._serialized_start=11891 _globals['_LOGENTRY_SEVERITY']._serialized_end=12005 _globals['_LOGENTRY_SEVERITY_ENUM']._serialized_start=11903 _globals['_LOGENTRY_SEVERITY_ENUM']._serialized_end=12005 _globals['_LOGCONTROL']._serialized_start=12007 _globals['_LOGCONTROL']._serialized_end=12019 _globals['_STARTWORKERREQUEST']._serialized_start=12022 _globals['_STARTWORKERREQUEST']._serialized_end=12530 _globals['_STARTWORKERREQUEST_PARAMSENTRY']._serialized_start=12485 _globals['_STARTWORKERREQUEST_PARAMSENTRY']._serialized_end=12530 _globals['_STARTWORKERRESPONSE']._serialized_start=12532 _globals['_STARTWORKERRESPONSE']._serialized_end=12568 _globals['_STOPWORKERREQUEST']._serialized_start=12570 _globals['_STOPWORKERREQUEST']._serialized_end=12608 _globals['_STOPWORKERRESPONSE']._serialized_start=12610 _globals['_STOPWORKERRESPONSE']._serialized_end=12645 _globals['_WORKERSTATUSREQUEST']._serialized_start=12647 _globals['_WORKERSTATUSREQUEST']._serialized_end=12680 _globals['_WORKERSTATUSRESPONSE']._serialized_start=12682 _globals['_WORKERSTATUSRESPONSE']._serialized_end=12752 _globals['_BEAMFNCONTROL']._serialized_start=12755 _globals['_BEAMFNCONTROL']._serialized_end=13078 _globals['_BEAMFNDATA']._serialized_start=13080 _globals['_BEAMFNDATA']._serialized_end=13204 _globals['_BEAMFNSTATE']._serialized_start=13207 _globals['_BEAMFNSTATE']._serialized_end=13342 _globals['_BEAMFNLOGGING']._serialized_start=13345 _globals['_BEAMFNLOGGING']._serialized_end=13482 _globals['_BEAMFNEXTERNALWORKERPOOL']._serialized_start=13485 _globals['_BEAMFNEXTERNALWORKERPOOL']._serialized_end=13782 _globals['_BEAMFNWORKERSTATUS']._serialized_start=13785 _globals['_BEAMFNWORKERSTATUS']._serialized_end=13949 # @@protoc_insertion_point(module_scope) ././@PaxHeader0000000000000000000000000000021000000000000010206 xustar00114 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_fn_api_pb2.pyi 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_fn_api_pb20000644000175100017510000036570115206113151032744 0ustar00runnerrunner""" @generated by mypy-protobuf. Do not edit manually! isort:skip_file TODO: Consider consolidating common components in another package and language namespaces for re-use with Runner Api. """ import builtins import collections.abc import google.protobuf.descriptor import google.protobuf.duration_pb2 import google.protobuf.internal.containers import google.protobuf.internal.enum_type_wrapper import google.protobuf.message import google.protobuf.struct_pb2 import google.protobuf.timestamp_pb2 import org.apache.beam.model.pipeline.v1.beam_runner_api_pb2 import org.apache.beam.model.pipeline.v1.endpoints_pb2 import org.apache.beam.model.pipeline.v1.metrics_pb2 import sys import typing if sys.version_info >= (3, 10): import typing as typing_extensions else: import typing_extensions DESCRIPTOR: google.protobuf.descriptor.FileDescriptor @typing_extensions.final class FnApiTransforms(google.protobuf.message.Message): """Describes transforms necessary to execute Beam over the FnAPI but are implementation details rather than part of the core model. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Runner: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _RunnerEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[FnApiTransforms._Runner.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor DATA_SOURCE: FnApiTransforms._Runner.ValueType # 0 """DataSource is a Root Transform, and a source of data for downstream transforms in the same ProcessBundleDescriptor. It represents a stream of values coming in from an external source/over a data channel, typically from the runner. It's not the PCollection itself but a description of how to get the portion of the PCollection for a given bundle. The DataSource transform is implemented in each SDK and not explicitly provided during pipeline construction. A runner inserts the transform in ProcessBundleDescriptors to indicate where the bundle can retrieve data for an associated ProcessBundleRequest. Data for the same request will be retrieved with the matching instruction ID, and transform ID determined by the runner. The DataSource transform will take a stream of bytes from the remote source for the matching instruction ID and decode them as windowed values using the provided coder ID, which must be a windowed value coder. Payload: RemoteGrpcPort """ DATA_SINK: FnApiTransforms._Runner.ValueType # 1 """DataSink is a transform that sends PCollection elements to a remote port using the Data API. The DataSink transform is implemented in each SDK and not explicitly provided during pipeline construction. A runner inserts the transform in ProcessBundleDescriptors to indicate where the bundle can send data for each associated ProcessBundleRequest. Data for the same request will be sent with the matching instruction ID and transform ID. Each PCollection that exits the ProcessBundleDescriptor subgraph will have it's own DataSink, keyed by a transform ID determined by the runner. The DataSink will take in a stream of elements for a given instruction ID and encode them for transmission to the remote sink. The coder ID must be for a windowed value coder. Payload: RemoteGrpcPort """ class Runner(_Runner, metaclass=_RunnerEnumTypeWrapper): ... DATA_SOURCE: FnApiTransforms.Runner.ValueType # 0 """DataSource is a Root Transform, and a source of data for downstream transforms in the same ProcessBundleDescriptor. It represents a stream of values coming in from an external source/over a data channel, typically from the runner. It's not the PCollection itself but a description of how to get the portion of the PCollection for a given bundle. The DataSource transform is implemented in each SDK and not explicitly provided during pipeline construction. A runner inserts the transform in ProcessBundleDescriptors to indicate where the bundle can retrieve data for an associated ProcessBundleRequest. Data for the same request will be retrieved with the matching instruction ID, and transform ID determined by the runner. The DataSource transform will take a stream of bytes from the remote source for the matching instruction ID and decode them as windowed values using the provided coder ID, which must be a windowed value coder. Payload: RemoteGrpcPort """ DATA_SINK: FnApiTransforms.Runner.ValueType # 1 """DataSink is a transform that sends PCollection elements to a remote port using the Data API. The DataSink transform is implemented in each SDK and not explicitly provided during pipeline construction. A runner inserts the transform in ProcessBundleDescriptors to indicate where the bundle can send data for each associated ProcessBundleRequest. Data for the same request will be sent with the matching instruction ID and transform ID. Each PCollection that exits the ProcessBundleDescriptor subgraph will have it's own DataSink, keyed by a transform ID determined by the runner. The DataSink will take in a stream of elements for a given instruction ID and encode them for transmission to the remote sink. The coder ID must be for a windowed value coder. Payload: RemoteGrpcPort """ def __init__( self, ) -> None: ... global___FnApiTransforms = FnApiTransforms @typing_extensions.final class RemoteGrpcPort(google.protobuf.message.Message): """A descriptor for connecting to a remote port using the Beam Fn Data API. Allows for communication between two environments (for example between the runner and the SDK). Stable """ DESCRIPTOR: google.protobuf.descriptor.Descriptor API_SERVICE_DESCRIPTOR_FIELD_NUMBER: builtins.int CODER_ID_FIELD_NUMBER: builtins.int @property def api_service_descriptor(self) -> org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor: """(Required) An API descriptor which describes where to connect to including any authentication that is required. """ coder_id: builtins.str """(Required) The ID of the Coder that will be used to encode and decode data sent over this port. """ def __init__( self, *, api_service_descriptor: org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor | None = ..., coder_id: builtins.str | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["api_service_descriptor", b"api_service_descriptor"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["api_service_descriptor", b"api_service_descriptor", "coder_id", b"coder_id"]) -> None: ... global___RemoteGrpcPort = RemoteGrpcPort @typing_extensions.final class GetProcessBundleDescriptorRequest(google.protobuf.message.Message): """Requests the ProcessBundleDescriptor with the given id.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor PROCESS_BUNDLE_DESCRIPTOR_ID_FIELD_NUMBER: builtins.int process_bundle_descriptor_id: builtins.str def __init__( self, *, process_bundle_descriptor_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["process_bundle_descriptor_id", b"process_bundle_descriptor_id"]) -> None: ... global___GetProcessBundleDescriptorRequest = GetProcessBundleDescriptorRequest @typing_extensions.final class InstructionRequest(google.protobuf.message.Message): """A request sent by a runner which the SDK is asked to fulfill. For any unsupported request type, an error should be returned with a matching instruction id. Stable """ DESCRIPTOR: google.protobuf.descriptor.Descriptor INSTRUCTION_ID_FIELD_NUMBER: builtins.int PROCESS_BUNDLE_FIELD_NUMBER: builtins.int PROCESS_BUNDLE_PROGRESS_FIELD_NUMBER: builtins.int PROCESS_BUNDLE_SPLIT_FIELD_NUMBER: builtins.int FINALIZE_BUNDLE_FIELD_NUMBER: builtins.int MONITORING_INFOS_FIELD_NUMBER: builtins.int HARNESS_MONITORING_INFOS_FIELD_NUMBER: builtins.int SAMPLE_DATA_FIELD_NUMBER: builtins.int REGISTER_FIELD_NUMBER: builtins.int instruction_id: builtins.str """(Required) A unique identifier provided by the runner which represents this requests execution. The InstructionResponse MUST have the matching id. """ @property def process_bundle(self) -> global___ProcessBundleRequest: ... @property def process_bundle_progress(self) -> global___ProcessBundleProgressRequest: ... @property def process_bundle_split(self) -> global___ProcessBundleSplitRequest: ... @property def finalize_bundle(self) -> global___FinalizeBundleRequest: ... @property def monitoring_infos(self) -> global___MonitoringInfosMetadataRequest: ... @property def harness_monitoring_infos(self) -> global___HarnessMonitoringInfosRequest: ... @property def sample_data(self) -> global___SampleDataRequest: ... @property def register(self) -> global___RegisterRequest: """DEPRECATED""" def __init__( self, *, instruction_id: builtins.str | None = ..., process_bundle: global___ProcessBundleRequest | None = ..., process_bundle_progress: global___ProcessBundleProgressRequest | None = ..., process_bundle_split: global___ProcessBundleSplitRequest | None = ..., finalize_bundle: global___FinalizeBundleRequest | None = ..., monitoring_infos: global___MonitoringInfosMetadataRequest | None = ..., harness_monitoring_infos: global___HarnessMonitoringInfosRequest | None = ..., sample_data: global___SampleDataRequest | None = ..., register: global___RegisterRequest | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["finalize_bundle", b"finalize_bundle", "harness_monitoring_infos", b"harness_monitoring_infos", "monitoring_infos", b"monitoring_infos", "process_bundle", b"process_bundle", "process_bundle_progress", b"process_bundle_progress", "process_bundle_split", b"process_bundle_split", "register", b"register", "request", b"request", "sample_data", b"sample_data"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["finalize_bundle", b"finalize_bundle", "harness_monitoring_infos", b"harness_monitoring_infos", "instruction_id", b"instruction_id", "monitoring_infos", b"monitoring_infos", "process_bundle", b"process_bundle", "process_bundle_progress", b"process_bundle_progress", "process_bundle_split", b"process_bundle_split", "register", b"register", "request", b"request", "sample_data", b"sample_data"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["request", b"request"]) -> typing_extensions.Literal["process_bundle", "process_bundle_progress", "process_bundle_split", "finalize_bundle", "monitoring_infos", "harness_monitoring_infos", "sample_data", "register"] | None: ... global___InstructionRequest = InstructionRequest @typing_extensions.final class InstructionResponse(google.protobuf.message.Message): """The response for an associated request the SDK had been asked to fulfill. Stable """ DESCRIPTOR: google.protobuf.descriptor.Descriptor INSTRUCTION_ID_FIELD_NUMBER: builtins.int ERROR_FIELD_NUMBER: builtins.int PROCESS_BUNDLE_FIELD_NUMBER: builtins.int PROCESS_BUNDLE_PROGRESS_FIELD_NUMBER: builtins.int PROCESS_BUNDLE_SPLIT_FIELD_NUMBER: builtins.int FINALIZE_BUNDLE_FIELD_NUMBER: builtins.int MONITORING_INFOS_FIELD_NUMBER: builtins.int HARNESS_MONITORING_INFOS_FIELD_NUMBER: builtins.int SAMPLE_DATA_FIELD_NUMBER: builtins.int REGISTER_FIELD_NUMBER: builtins.int instruction_id: builtins.str """(Required) A reference provided by the runner which represents a requests execution. The InstructionResponse MUST have the matching id when responding to the runner. """ error: builtins.str """If this is specified, then this instruction has failed. A human readable string representing the reason as to why processing has failed. """ @property def process_bundle(self) -> global___ProcessBundleResponse: ... @property def process_bundle_progress(self) -> global___ProcessBundleProgressResponse: ... @property def process_bundle_split(self) -> global___ProcessBundleSplitResponse: ... @property def finalize_bundle(self) -> global___FinalizeBundleResponse: ... @property def monitoring_infos(self) -> global___MonitoringInfosMetadataResponse: ... @property def harness_monitoring_infos(self) -> global___HarnessMonitoringInfosResponse: ... @property def sample_data(self) -> global___SampleDataResponse: ... @property def register(self) -> global___RegisterResponse: """DEPRECATED""" def __init__( self, *, instruction_id: builtins.str | None = ..., error: builtins.str | None = ..., process_bundle: global___ProcessBundleResponse | None = ..., process_bundle_progress: global___ProcessBundleProgressResponse | None = ..., process_bundle_split: global___ProcessBundleSplitResponse | None = ..., finalize_bundle: global___FinalizeBundleResponse | None = ..., monitoring_infos: global___MonitoringInfosMetadataResponse | None = ..., harness_monitoring_infos: global___HarnessMonitoringInfosResponse | None = ..., sample_data: global___SampleDataResponse | None = ..., register: global___RegisterResponse | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["finalize_bundle", b"finalize_bundle", "harness_monitoring_infos", b"harness_monitoring_infos", "monitoring_infos", b"monitoring_infos", "process_bundle", b"process_bundle", "process_bundle_progress", b"process_bundle_progress", "process_bundle_split", b"process_bundle_split", "register", b"register", "response", b"response", "sample_data", b"sample_data"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["error", b"error", "finalize_bundle", b"finalize_bundle", "harness_monitoring_infos", b"harness_monitoring_infos", "instruction_id", b"instruction_id", "monitoring_infos", b"monitoring_infos", "process_bundle", b"process_bundle", "process_bundle_progress", b"process_bundle_progress", "process_bundle_split", b"process_bundle_split", "register", b"register", "response", b"response", "sample_data", b"sample_data"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["response", b"response"]) -> typing_extensions.Literal["process_bundle", "process_bundle_progress", "process_bundle_split", "finalize_bundle", "monitoring_infos", "harness_monitoring_infos", "sample_data", "register"] | None: ... global___InstructionResponse = InstructionResponse @typing_extensions.final class SampleDataRequest(google.protobuf.message.Message): """If supported, the `SampleDataRequest` will respond with a `SampleDataResponse`. The SDK being queried must have the "beam:protocol:data_sampling:v1" capability. Samples are taken only from the specified PCollection ids. An empty list will return everything. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor PCOLLECTION_IDS_FIELD_NUMBER: builtins.int @property def pcollection_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """(Optional) The PCollection ids to filter for.""" def __init__( self, *, pcollection_ids: collections.abc.Iterable[builtins.str] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["pcollection_ids", b"pcollection_ids"]) -> None: ... global___SampleDataRequest = SampleDataRequest @typing_extensions.final class SampledElement(google.protobuf.message.Message): """An element sampled when the SDK is processing a bundle. This is a proto message to allow for additional per-element metadata. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class Exception(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor INSTRUCTION_ID_FIELD_NUMBER: builtins.int TRANSFORM_ID_FIELD_NUMBER: builtins.int ERROR_FIELD_NUMBER: builtins.int instruction_id: builtins.str """(Required) The instruction ID of the associated ProcessBundleRequest.""" transform_id: builtins.str """(Required) The transform ID of the executing PTransform during the exception. """ error: builtins.str """(Required) The error message to be displayed to the user. Can use the other fields to query for contextual logs. """ def __init__( self, *, instruction_id: builtins.str | None = ..., transform_id: builtins.str | None = ..., error: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["error", b"error", "instruction_id", b"instruction_id", "transform_id", b"transform_id"]) -> None: ... ELEMENT_FIELD_NUMBER: builtins.int SAMPLE_TIMESTAMP_FIELD_NUMBER: builtins.int EXCEPTION_FIELD_NUMBER: builtins.int element: builtins.bytes """(Required) Sampled raw bytes for an element. This is a single encoded element in the nested context. """ @property def sample_timestamp(self) -> google.protobuf.timestamp_pb2.Timestamp: """(Required) Timestamp of when the sample was taken.""" @property def exception(self) -> global___SampledElement.Exception: """(Optional) This will be set if this element was sampled because of a user exception. """ def __init__( self, *, element: builtins.bytes | None = ..., sample_timestamp: google.protobuf.timestamp_pb2.Timestamp | None = ..., exception: global___SampledElement.Exception | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["exception", b"exception", "sample_timestamp", b"sample_timestamp"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["element", b"element", "exception", b"exception", "sample_timestamp", b"sample_timestamp"]) -> None: ... global___SampledElement = SampledElement @typing_extensions.final class SampleDataResponse(google.protobuf.message.Message): """If supported, the `SampleDataResponse` will contain samples from PCollections based upon the filters specified in the request. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class ElementList(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ELEMENTS_FIELD_NUMBER: builtins.int @property def elements(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___SampledElement]: """Required. The individual elements sampled from a PCollection.""" def __init__( self, *, elements: collections.abc.Iterable[global___SampledElement] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["elements", b"elements"]) -> None: ... @typing_extensions.final class ElementSamplesEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> global___SampleDataResponse.ElementList: ... def __init__( self, *, key: builtins.str | None = ..., value: global___SampleDataResponse.ElementList | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... ELEMENT_SAMPLES_FIELD_NUMBER: builtins.int @property def element_samples(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___SampleDataResponse.ElementList]: """Map from PCollection id to sampled elements.""" def __init__( self, *, element_samples: collections.abc.Mapping[builtins.str, global___SampleDataResponse.ElementList] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["element_samples", b"element_samples"]) -> None: ... global___SampleDataResponse = SampleDataResponse @typing_extensions.final class HarnessMonitoringInfosRequest(google.protobuf.message.Message): """A request to provide full MonitoringInfo associated with the entire SDK harness process, not specific to a bundle. An SDK can report metrics using an identifier that only contains the associated payload. A runner who wants to receive the full metrics information can request all the monitoring metadata via a MonitoringInfosMetadataRequest providing a list of ids as necessary. The SDK is allowed to reuse the identifiers for the lifetime of the associated control connection as long as the MonitoringInfo could be reconstructed fully by overwriting its payload field with the bytes specified here. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... global___HarnessMonitoringInfosRequest = HarnessMonitoringInfosRequest @typing_extensions.final class HarnessMonitoringInfosResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class MonitoringDataEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str value: builtins.bytes def __init__( self, *, key: builtins.str | None = ..., value: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... MONITORING_DATA_FIELD_NUMBER: builtins.int @property def monitoring_data(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.bytes]: """An identifier to MonitoringInfo.payload mapping containing Metrics associated with the SDK harness, not a specific bundle. An SDK can report metrics using an identifier that only contains the associated payload. A runner who wants to receive the full metrics information can request all the monitoring metadata via a MonitoringInfosMetadataRequest providing a list of ids as necessary. The SDK is allowed to reuse the identifiers for the lifetime of the associated control connection as long as the MonitoringInfo could be reconstructed fully by overwriting its payload field with the bytes specified here. """ def __init__( self, *, monitoring_data: collections.abc.Mapping[builtins.str, builtins.bytes] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["monitoring_data", b"monitoring_data"]) -> None: ... global___HarnessMonitoringInfosResponse = HarnessMonitoringInfosResponse @typing_extensions.final class RegisterRequest(google.protobuf.message.Message): """A list of objects which can be referred to by the runner in future requests. Stable """ DESCRIPTOR: google.protobuf.descriptor.Descriptor PROCESS_BUNDLE_DESCRIPTOR_FIELD_NUMBER: builtins.int @property def process_bundle_descriptor(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ProcessBundleDescriptor]: """(Optional) The set of descriptors used to process bundles.""" def __init__( self, *, process_bundle_descriptor: collections.abc.Iterable[global___ProcessBundleDescriptor] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["process_bundle_descriptor", b"process_bundle_descriptor"]) -> None: ... global___RegisterRequest = RegisterRequest @typing_extensions.final class RegisterResponse(google.protobuf.message.Message): """Stable""" DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... global___RegisterResponse = RegisterResponse @typing_extensions.final class ProcessBundleDescriptor(google.protobuf.message.Message): """Definitions that should be used to construct the bundle processing graph.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class TransformsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.PTransform: ... def __init__( self, *, key: builtins.str | None = ..., value: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.PTransform | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... @typing_extensions.final class PcollectionsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.PCollection: ... def __init__( self, *, key: builtins.str | None = ..., value: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.PCollection | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... @typing_extensions.final class WindowingStrategiesEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.WindowingStrategy: ... def __init__( self, *, key: builtins.str | None = ..., value: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.WindowingStrategy | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... @typing_extensions.final class CodersEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Coder: ... def __init__( self, *, key: builtins.str | None = ..., value: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Coder | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... @typing_extensions.final class EnvironmentsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Environment: ... def __init__( self, *, key: builtins.str | None = ..., value: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Environment | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... ID_FIELD_NUMBER: builtins.int TRANSFORMS_FIELD_NUMBER: builtins.int PCOLLECTIONS_FIELD_NUMBER: builtins.int WINDOWING_STRATEGIES_FIELD_NUMBER: builtins.int CODERS_FIELD_NUMBER: builtins.int ENVIRONMENTS_FIELD_NUMBER: builtins.int STATE_API_SERVICE_DESCRIPTOR_FIELD_NUMBER: builtins.int TIMER_API_SERVICE_DESCRIPTOR_FIELD_NUMBER: builtins.int id: builtins.str """(Required) A pipeline level unique id which can be used as a reference to refer to this. """ @property def transforms(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.PTransform]: """(Required) A map from pipeline-scoped id to PTransform.""" @property def pcollections(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.PCollection]: """(Required) A map from pipeline-scoped id to PCollection.""" @property def windowing_strategies(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.WindowingStrategy]: """(Required) A map from pipeline-scoped id to WindowingStrategy.""" @property def coders(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Coder]: """(Required) A map from pipeline-scoped id to Coder.""" @property def environments(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Environment]: """(Required) A map from pipeline-scoped id to Environment.""" @property def state_api_service_descriptor(self) -> org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor: """A descriptor describing the end point to use for State API calls. Required if the Runner intends to send remote references over the data plane or if any of the transforms rely on user state or side inputs. """ @property def timer_api_service_descriptor(self) -> org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor: """A descriptor describing the end point to use for Data API for user timers. Required if the ProcessBundleDescriptor contains any transforms that have user timers. """ def __init__( self, *, id: builtins.str | None = ..., transforms: collections.abc.Mapping[builtins.str, org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.PTransform] | None = ..., pcollections: collections.abc.Mapping[builtins.str, org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.PCollection] | None = ..., windowing_strategies: collections.abc.Mapping[builtins.str, org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.WindowingStrategy] | None = ..., coders: collections.abc.Mapping[builtins.str, org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Coder] | None = ..., environments: collections.abc.Mapping[builtins.str, org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Environment] | None = ..., state_api_service_descriptor: org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor | None = ..., timer_api_service_descriptor: org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["state_api_service_descriptor", b"state_api_service_descriptor", "timer_api_service_descriptor", b"timer_api_service_descriptor"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["coders", b"coders", "environments", b"environments", "id", b"id", "pcollections", b"pcollections", "state_api_service_descriptor", b"state_api_service_descriptor", "timer_api_service_descriptor", b"timer_api_service_descriptor", "transforms", b"transforms", "windowing_strategies", b"windowing_strategies"]) -> None: ... global___ProcessBundleDescriptor = ProcessBundleDescriptor @typing_extensions.final class BundleApplication(google.protobuf.message.Message): """One of the applications specifying the scope of work for a bundle. See https://docs.google.com/document/d/1tUDb45sStdR8u7-jBkGdw3OGFK7aa2-V7eo86zYSE_4/edit#heading=h.9g3g5weg2u9 for further details. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class OutputWatermarksEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> google.protobuf.timestamp_pb2.Timestamp: ... def __init__( self, *, key: builtins.str | None = ..., value: google.protobuf.timestamp_pb2.Timestamp | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... TRANSFORM_ID_FIELD_NUMBER: builtins.int INPUT_ID_FIELD_NUMBER: builtins.int ELEMENT_FIELD_NUMBER: builtins.int OUTPUT_WATERMARKS_FIELD_NUMBER: builtins.int IS_BOUNDED_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The transform to which to pass the element""" input_id: builtins.str """(Required) Name of the transform's input to which to pass the element.""" element: builtins.bytes """(Required) The encoded element to pass to the transform.""" @property def output_watermarks(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, google.protobuf.timestamp_pb2.Timestamp]: """The map is keyed by the local output name of the PTransform. Each value represents a lower bound on the timestamps of elements that are produced by this PTransform into each of its output PCollections when invoked with this application. If there is no watermark reported from RestrictionTracker, the runner will use MIN_TIMESTAMP by default. """ is_bounded: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.IsBounded.Enum.ValueType """Whether this application potentially produces an unbounded amount of data. Note that this should only be set to BOUNDED if and only if the application is known to produce a finite amount of output. """ def __init__( self, *, transform_id: builtins.str | None = ..., input_id: builtins.str | None = ..., element: builtins.bytes | None = ..., output_watermarks: collections.abc.Mapping[builtins.str, google.protobuf.timestamp_pb2.Timestamp] | None = ..., is_bounded: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.IsBounded.Enum.ValueType | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["element", b"element", "input_id", b"input_id", "is_bounded", b"is_bounded", "output_watermarks", b"output_watermarks", "transform_id", b"transform_id"]) -> None: ... global___BundleApplication = BundleApplication @typing_extensions.final class DelayedBundleApplication(google.protobuf.message.Message): """An Application should be scheduled for execution after a delay. Either an absolute timestamp or a relative timestamp can represent a scheduled execution time. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor APPLICATION_FIELD_NUMBER: builtins.int REQUESTED_TIME_DELAY_FIELD_NUMBER: builtins.int @property def application(self) -> global___BundleApplication: """(Required) The application that should be scheduled.""" @property def requested_time_delay(self) -> google.protobuf.duration_pb2.Duration: """Recommended time delay at which the application should be scheduled to execute by the runner. Time delay that equals 0 may be scheduled to execute immediately. The unit of time delay should be microsecond. """ def __init__( self, *, application: global___BundleApplication | None = ..., requested_time_delay: google.protobuf.duration_pb2.Duration | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["application", b"application", "requested_time_delay", b"requested_time_delay"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["application", b"application", "requested_time_delay", b"requested_time_delay"]) -> None: ... global___DelayedBundleApplication = DelayedBundleApplication @typing_extensions.final class ProcessBundleRequest(google.protobuf.message.Message): """A request to process a given bundle. Stable """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class CacheToken(google.protobuf.message.Message): """Contains the cache token and also defines the scope of what the token applies to. See https://s.apache.org/beam-fn-state-api-and-bundle-processing#heading=h.7ghoih5aig5m for additional details on how to use the cache token with the State API to cache data across bundle boundaries. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class UserState(google.protobuf.message.Message): """A flag to indicate a cache token is valid for all user state.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... @typing_extensions.final class SideInput(google.protobuf.message.Message): """A flag to indicate a cache token is valid for a side input.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int SIDE_INPUT_ID_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The id of the PTransform containing a side input.""" side_input_id: builtins.str """(Required) The id of the side input.""" def __init__( self, *, transform_id: builtins.str | None = ..., side_input_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["side_input_id", b"side_input_id", "transform_id", b"transform_id"]) -> None: ... USER_STATE_FIELD_NUMBER: builtins.int SIDE_INPUT_FIELD_NUMBER: builtins.int TOKEN_FIELD_NUMBER: builtins.int @property def user_state(self) -> global___ProcessBundleRequest.CacheToken.UserState: ... @property def side_input(self) -> global___ProcessBundleRequest.CacheToken.SideInput: ... token: builtins.bytes """An opaque token used with the StateKey to create a globally unique identifier. """ def __init__( self, *, user_state: global___ProcessBundleRequest.CacheToken.UserState | None = ..., side_input: global___ProcessBundleRequest.CacheToken.SideInput | None = ..., token: builtins.bytes | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["side_input", b"side_input", "type", b"type", "user_state", b"user_state"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["side_input", b"side_input", "token", b"token", "type", b"type", "user_state", b"user_state"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["type", b"type"]) -> typing_extensions.Literal["user_state", "side_input"] | None: ... PROCESS_BUNDLE_DESCRIPTOR_ID_FIELD_NUMBER: builtins.int CACHE_TOKENS_FIELD_NUMBER: builtins.int ELEMENTS_FIELD_NUMBER: builtins.int HAS_NO_STATE_FIELD_NUMBER: builtins.int ONLY_BUNDLE_FOR_KEYS_FIELD_NUMBER: builtins.int process_bundle_descriptor_id: builtins.str """(Required) A reference to the process bundle descriptor that must be instantiated and executed by the SDK harness. """ @property def cache_tokens(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ProcessBundleRequest.CacheToken]: """(Optional) A list of cache tokens that can be used by an SDK to reuse cached data returned by the State API across multiple bundles. Note that SDKs that can efficiently consume this field should declare the beam:protocol:state_caching:v1 capability enabling runners to reduce the amount of memory used. See https://s.apache.org/beam-fn-state-api-and-bundle-processing#heading=h.7ghoih5aig5m for additional details on how to use the cache token with the State API to cache data across bundle boundaries. """ @property def elements(self) -> global___Elements: """(Optional) Elements to be processed with the bundle. Either all or none of the bundle elements should be included in the ProcessBundleRequest. This embedding is to achieve better efficiency for bundles that contain only small amounts of data and are cheap to be processed on the SDK harness side. This field can be set only if the SDK declares that it supports the beam:protocol:control_request_elements_embedding:v1 capability. See more at https://s.apache.org/beam-fn-api-control-data-embedding. """ has_no_state: builtins.bool """indicates that the runner has no stare for the keys in this bundle so SDk can safely begin stateful processing with a locally-generated initial empty state """ only_bundle_for_keys: builtins.bool """indicates that the runner will never process another bundle for the keys in this bundle so state need not be included in the bundle commit. """ def __init__( self, *, process_bundle_descriptor_id: builtins.str | None = ..., cache_tokens: collections.abc.Iterable[global___ProcessBundleRequest.CacheToken] | None = ..., elements: global___Elements | None = ..., has_no_state: builtins.bool | None = ..., only_bundle_for_keys: builtins.bool | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["elements", b"elements"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["cache_tokens", b"cache_tokens", "elements", b"elements", "has_no_state", b"has_no_state", "only_bundle_for_keys", b"only_bundle_for_keys", "process_bundle_descriptor_id", b"process_bundle_descriptor_id"]) -> None: ... global___ProcessBundleRequest = ProcessBundleRequest @typing_extensions.final class ProcessBundleResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class MonitoringDataEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str value: builtins.bytes def __init__( self, *, key: builtins.str | None = ..., value: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... RESIDUAL_ROOTS_FIELD_NUMBER: builtins.int MONITORING_INFOS_FIELD_NUMBER: builtins.int REQUIRES_FINALIZATION_FIELD_NUMBER: builtins.int MONITORING_DATA_FIELD_NUMBER: builtins.int ELEMENTS_FIELD_NUMBER: builtins.int @property def residual_roots(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___DelayedBundleApplication]: """(Optional) Specifies that the bundle has not been completed and the following applications need to be scheduled and executed in the future. A runner that does not yet support residual roots MUST still check that this is empty for correctness. Note that these residual roots must not have been returned as part of a prior split for this bundle. """ @property def monitoring_infos(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[org.apache.beam.model.pipeline.v1.metrics_pb2.MonitoringInfo]: """DEPRECATED (Required) The list of metrics or other MonitoredState collected while processing this bundle. """ requires_finalization: builtins.bool """(Optional) Specifies that the runner must callback to this worker once the output of the bundle is committed. The Runner must send a FinalizeBundleRequest with the instruction id of the ProcessBundleRequest that is related to this ProcessBundleResponse. """ @property def monitoring_data(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.bytes]: """An identifier to MonitoringInfo.payload mapping. An SDK can report metrics using an identifier that only contains the associated payload. A runner who wants to receive the full metrics information can request all the monitoring metadata via a MonitoringInfosMetadataRequest providing a list of ids as necessary. The SDK is allowed to reuse the identifiers across multiple bundles as long as the MonitoringInfo could be reconstructed fully by overwriting its payload field with the bytes specified here. """ @property def elements(self) -> global___Elements: """(Optional) Output elements of the processed bundle. Either all or none of the bundle elements should be included in the ProcessBundleResponse. This embedding is to achieve better efficiency for bundles that only contain small amounts of data. his field can be set only if the runner declares that it supports the beam:protocol:control_request_elements_embedding:v1 capability. See more at https://s.apache.org/beam-fn-api-control-data-embedding. """ def __init__( self, *, residual_roots: collections.abc.Iterable[global___DelayedBundleApplication] | None = ..., monitoring_infos: collections.abc.Iterable[org.apache.beam.model.pipeline.v1.metrics_pb2.MonitoringInfo] | None = ..., requires_finalization: builtins.bool | None = ..., monitoring_data: collections.abc.Mapping[builtins.str, builtins.bytes] | None = ..., elements: global___Elements | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["elements", b"elements"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["elements", b"elements", "monitoring_data", b"monitoring_data", "monitoring_infos", b"monitoring_infos", "requires_finalization", b"requires_finalization", "residual_roots", b"residual_roots"]) -> None: ... global___ProcessBundleResponse = ProcessBundleResponse @typing_extensions.final class ProcessBundleProgressRequest(google.protobuf.message.Message): """A request to report progress information for a given bundle. This is an optional request to be handled and is used to support advanced SDK features such as SplittableDoFn, user level metrics etc. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor INSTRUCTION_ID_FIELD_NUMBER: builtins.int instruction_id: builtins.str """(Required) A reference to an active process bundle request with the given instruction id. """ def __init__( self, *, instruction_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["instruction_id", b"instruction_id"]) -> None: ... global___ProcessBundleProgressRequest = ProcessBundleProgressRequest @typing_extensions.final class MonitoringInfosMetadataRequest(google.protobuf.message.Message): """A request to provide full MonitoringInfo for a set of provided ids. An SDK can report metrics using an identifier that only contains the associated payload. A runner who wants to receive the full metrics information can request all the monitoring metadata via a MonitoringInfosMetadataRequest providing a list of ids as necessary. The SDK is allowed to reuse the identifiers for the lifetime of the associated control connection as long as the MonitoringInfo could be reconstructed fully by overwriting its payload field with the bytes specified here. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor MONITORING_INFO_ID_FIELD_NUMBER: builtins.int @property def monitoring_info_id(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """A list of ids for which the full MonitoringInfo is requested for.""" def __init__( self, *, monitoring_info_id: collections.abc.Iterable[builtins.str] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["monitoring_info_id", b"monitoring_info_id"]) -> None: ... global___MonitoringInfosMetadataRequest = MonitoringInfosMetadataRequest @typing_extensions.final class ProcessBundleProgressResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class MonitoringDataEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str value: builtins.bytes def __init__( self, *, key: builtins.str | None = ..., value: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... MONITORING_INFOS_FIELD_NUMBER: builtins.int MONITORING_DATA_FIELD_NUMBER: builtins.int CONSUMING_RECEIVED_DATA_FIELD_NUMBER: builtins.int @property def monitoring_infos(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[org.apache.beam.model.pipeline.v1.metrics_pb2.MonitoringInfo]: """DEPRECATED (Required) The list of metrics or other MonitoredState collected while processing this bundle. """ @property def monitoring_data(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.bytes]: """An identifier to MonitoringInfo.payload mapping. An SDK can report metrics using an identifier that only contains the associated payload. A runner who wants to receive the full metrics information can request all the monitoring metadata via a MonitoringInfosMetadataRequest providing a list of ids as necessary. The SDK is allowed to reuse the identifiers for the lifetime of the associated control connection as long as the MonitoringInfo could be reconstructed fully by overwriting its payload field with the bytes specified here. """ consuming_received_data: builtins.bool """Indicates that the SDK is still busy consuming the data that as already been received on the data channel. If this is set, a runner may abstain from sending further data on the data channel until this field becomes unset. This field is currently used during shuffle reads on large elements. """ def __init__( self, *, monitoring_infos: collections.abc.Iterable[org.apache.beam.model.pipeline.v1.metrics_pb2.MonitoringInfo] | None = ..., monitoring_data: collections.abc.Mapping[builtins.str, builtins.bytes] | None = ..., consuming_received_data: builtins.bool | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["_consuming_received_data", b"_consuming_received_data", "consuming_received_data", b"consuming_received_data"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["_consuming_received_data", b"_consuming_received_data", "consuming_received_data", b"consuming_received_data", "monitoring_data", b"monitoring_data", "monitoring_infos", b"monitoring_infos"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["_consuming_received_data", b"_consuming_received_data"]) -> typing_extensions.Literal["consuming_received_data"] | None: ... global___ProcessBundleProgressResponse = ProcessBundleProgressResponse @typing_extensions.final class MonitoringInfosMetadataResponse(google.protobuf.message.Message): """A response that contains the full mapping information associated with a specified set of identifiers. An SDK can report metrics using an identifier that only contains the associated payload. A runner who wants to receive the full metrics information can request all the monitoring metadata via a MonitoringInfosMetadataRequest providing a list of ids as necessary. The SDK is allowed to reuse the identifiers for the lifetime of the associated control connection as long as the MonitoringInfo could be reconstructed fully by overwriting its payload field with the bytes specified here. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class MonitoringInfoEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> org.apache.beam.model.pipeline.v1.metrics_pb2.MonitoringInfo: ... def __init__( self, *, key: builtins.str | None = ..., value: org.apache.beam.model.pipeline.v1.metrics_pb2.MonitoringInfo | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... MONITORING_INFO_FIELD_NUMBER: builtins.int @property def monitoring_info(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, org.apache.beam.model.pipeline.v1.metrics_pb2.MonitoringInfo]: """A mapping from an identifier to the full metrics information.""" def __init__( self, *, monitoring_info: collections.abc.Mapping[builtins.str, org.apache.beam.model.pipeline.v1.metrics_pb2.MonitoringInfo] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["monitoring_info", b"monitoring_info"]) -> None: ... global___MonitoringInfosMetadataResponse = MonitoringInfosMetadataResponse @typing_extensions.final class ProcessBundleSplitRequest(google.protobuf.message.Message): """Represents a request to the SDK to split a currently active bundle.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class DesiredSplit(google.protobuf.message.Message): """A message specifying the desired split for a single transform.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor FRACTION_OF_REMAINDER_FIELD_NUMBER: builtins.int ALLOWED_SPLIT_POINTS_FIELD_NUMBER: builtins.int ESTIMATED_INPUT_ELEMENTS_FIELD_NUMBER: builtins.int fraction_of_remainder: builtins.float """(Required) The fraction of known work remaining in this bundle for this transform that should be kept by the SDK after this split. Set to 0 to "checkpoint" as soon as possible (keeping as little work as possible and returning the remainder). """ @property def allowed_split_points(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: """(Optional) A set of allowed element indices where the SDK may split. When this is empty, there are no constraints on where to split. """ estimated_input_elements: builtins.int """(Required for gRPC Read operation transforms) Number of total elements expected to be sent to this GrpcRead operation, required to correctly account for unreceived data when determining where to split. """ def __init__( self, *, fraction_of_remainder: builtins.float | None = ..., allowed_split_points: collections.abc.Iterable[builtins.int] | None = ..., estimated_input_elements: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["allowed_split_points", b"allowed_split_points", "estimated_input_elements", b"estimated_input_elements", "fraction_of_remainder", b"fraction_of_remainder"]) -> None: ... @typing_extensions.final class DesiredSplitsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> global___ProcessBundleSplitRequest.DesiredSplit: ... def __init__( self, *, key: builtins.str | None = ..., value: global___ProcessBundleSplitRequest.DesiredSplit | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... INSTRUCTION_ID_FIELD_NUMBER: builtins.int DESIRED_SPLITS_FIELD_NUMBER: builtins.int instruction_id: builtins.str """(Required) A reference to an active process bundle request with the given instruction id. """ @property def desired_splits(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___ProcessBundleSplitRequest.DesiredSplit]: """(Required) Specifies the desired split for each transform. Currently only splits at gRPC read operations are supported. This may, of course, limit the amount of work downstream operations receive. """ def __init__( self, *, instruction_id: builtins.str | None = ..., desired_splits: collections.abc.Mapping[builtins.str, global___ProcessBundleSplitRequest.DesiredSplit] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["desired_splits", b"desired_splits", "instruction_id", b"instruction_id"]) -> None: ... global___ProcessBundleSplitRequest = ProcessBundleSplitRequest @typing_extensions.final class ProcessBundleSplitResponse(google.protobuf.message.Message): """Represents a partition of the bundle: a "primary" and a "residual", with the following properties: - The work in primary and residual doesn't overlap, and combined, adds up to the work in the current bundle if the split hadn't happened. - The current bundle, if it keeps executing, will have done exactly none of the work under residual_roots and none of the elements at and beyond the first_residual_element. - The current bundle, if no further splits happen, will have done exactly the work under primary_roots and all elements up to and including the channel split's last_primary_element. This allows the SDK to relinquish ownership of and commit to not process some of the elements that it may have been sent (the residual) while retaining ownership and commitment to finish the other portion (the primary). Example with three splits of a single bundle: Let's say the SDK is processing elements [A B C D E]. These elements make up the 0-indexed channel. ** First Split ** Channel Split = [ A B C D <> E ] Primary Roots = [] (No elements were split) Residual Roots = [] Say a split request comes in. The SDK could return a response with a channel split representing a last_primary_element of 3 (D) and first_residual_element of 4 (E). The SDK is now responsible for processing A B C D and the runner must process E in the future. (A B C D) | (E) ** Second Split ** Channel Split = [ A < B C > D E ] Primary Roots = [B1 C1] Residual Roots = [B2 C2] A future split request could have the SDK split the elements B into B1 and B2 and C into C1 and C2 representing their primary and residual roots. The (A B1 C1) | (B2 C2 D) SDK would return a response with a channel split representing a last_primary_element of 0 (A) and first_residual_element of 3 (D) with primary_roots (B1, C1) and residual_roots (B2, C2). The SDK is now responsible for processing A B1 C1 and the runner must process B2 C2 D (and E from the prior split) in the future. ** Third Split ** Channel Split = [ A < B C > D E ] Primary Roots = [B1a] Residual Roots [B1b C1] Yet another future split request could have the SDK could split B1 further into B1a and B1b primary and residuals and return C1 as a residual (assuming C1 was left unprocessed). The SDK would return a response with a channel split representing a last_primary_element of 0 (A) and first_residual_element of 3 (E) with primary_roots (B1a) and residual_roots (B1b, C1). The SDK is now responsible for processing A B1a the runner must process B1b C1 (in addition to C2, D, E from prior splits) in the future. (A B1a) | (B1b C1) For more rigorous definitions see https://s.apache.org/beam-breaking-fusion """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class ChannelSplit(google.protobuf.message.Message): """Represents contiguous portions of the data channel that are either entirely processed or entirely unprocessed and belong to the primary or residual respectively. This affords both a more efficient representation over the FnAPI (if the bundle is large) and often a more efficient representation on the runner side (e.g. if the set of elements can be represented as some range in an underlying dataset). Note that for a split the following properties must hold: - last_primary_element < first_residual_element - primary roots and residual roots can only be specified if the last_primary_element + 1 < first_residual_element (typically there is one primary and residual root per element in the range (last_primary_element, first_residual_element)) - primary roots and residual roots must represent a disjoint but full coverage of work represented by the elements between last_primary_element and first_residual_element Note that subsequent splits of the same bundle must ensure that: - the first_residual_element does not increase - the first_residual_element does not decrease if there were residual or primary roots returned in a prior split. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int LAST_PRIMARY_ELEMENT_FIELD_NUMBER: builtins.int FIRST_RESIDUAL_ELEMENT_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The grpc read transform reading this channel.""" last_primary_element: builtins.int """(Required) The last element of the input channel that should be entirely considered part of the primary, identified by its absolute zero-based index in the (ordered) channel. """ first_residual_element: builtins.int """(Required) The first element of the input channel that should be entirely considered part of the residual, identified by its absolute zero-based index in the (ordered) channel. """ def __init__( self, *, transform_id: builtins.str | None = ..., last_primary_element: builtins.int | None = ..., first_residual_element: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["first_residual_element", b"first_residual_element", "last_primary_element", b"last_primary_element", "transform_id", b"transform_id"]) -> None: ... PRIMARY_ROOTS_FIELD_NUMBER: builtins.int RESIDUAL_ROOTS_FIELD_NUMBER: builtins.int CHANNEL_SPLITS_FIELD_NUMBER: builtins.int @property def primary_roots(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___BundleApplication]: """(Optional) Root applications that should replace the current bundle. Note that primary roots can only be specified if a channel split's last_primary_element + 1 < first_residual_element Note that there must be a corresponding residual root contained within residual_roots representing the remainder of processing for the original element this this primary root represents a fraction of. """ @property def residual_roots(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___DelayedBundleApplication]: """(Optional) Root applications that have been removed from the current bundle and have to be executed in a separate bundle (e.g. in parallel on a different worker, or after the current bundle completes, etc.) Note that residual roots can only be specified if a channel split's last_primary_element + 1 < first_residual_element Note that there must be a corresponding primary root contained within primary_roots representing the remainder of processing for the original element this this residual root represents a fraction of. Note that subsequent splits must not return prior residual roots. """ @property def channel_splits(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ProcessBundleSplitResponse.ChannelSplit]: """(Required) Partitions of input data channels into primary and residual elements, if any. Must not include any elements represented in the bundle applications roots above of the current split or any prior split of the same bundle. """ def __init__( self, *, primary_roots: collections.abc.Iterable[global___BundleApplication] | None = ..., residual_roots: collections.abc.Iterable[global___DelayedBundleApplication] | None = ..., channel_splits: collections.abc.Iterable[global___ProcessBundleSplitResponse.ChannelSplit] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["channel_splits", b"channel_splits", "primary_roots", b"primary_roots", "residual_roots", b"residual_roots"]) -> None: ... global___ProcessBundleSplitResponse = ProcessBundleSplitResponse @typing_extensions.final class FinalizeBundleRequest(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor INSTRUCTION_ID_FIELD_NUMBER: builtins.int instruction_id: builtins.str """(Required) A reference to a completed process bundle request with the given instruction id. """ def __init__( self, *, instruction_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["instruction_id", b"instruction_id"]) -> None: ... global___FinalizeBundleRequest = FinalizeBundleRequest @typing_extensions.final class FinalizeBundleResponse(google.protobuf.message.Message): """Empty""" DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... global___FinalizeBundleResponse = FinalizeBundleResponse @typing_extensions.final class Elements(google.protobuf.message.Message): """ Data Plane API Messages used to represent logical byte streams. Stable """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class Data(google.protobuf.message.Message): """Represents multiple encoded elements in nested context for a given named instruction and transform. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor INSTRUCTION_ID_FIELD_NUMBER: builtins.int TRANSFORM_ID_FIELD_NUMBER: builtins.int DATA_FIELD_NUMBER: builtins.int IS_LAST_FIELD_NUMBER: builtins.int instruction_id: builtins.str """(Required) A reference to an active instruction request with the given instruction id. """ transform_id: builtins.str """(Required) A definition representing a consumer or producer of this data. If received by a harness, this represents the consumer within that harness that should consume these bytes. If sent by a harness, this represents the producer of these bytes. Note that a single element may span multiple Data messages. Note that a sending/receiving pair should share the same identifier. """ data: builtins.bytes """(Optional) Represents a part of a logical byte stream. Elements within the logical byte stream are encoded in the nested context and concatenated together. """ is_last: builtins.bool """(Optional) Set this bit to indicate the this is the last data block for the given instruction and transform, ending the stream. """ def __init__( self, *, instruction_id: builtins.str | None = ..., transform_id: builtins.str | None = ..., data: builtins.bytes | None = ..., is_last: builtins.bool | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["data", b"data", "instruction_id", b"instruction_id", "is_last", b"is_last", "transform_id", b"transform_id"]) -> None: ... @typing_extensions.final class DrainMode(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[Elements.DrainMode._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor UNSPECIFIED: Elements.DrainMode._Enum.ValueType # 0 NOT_DRAINING: Elements.DrainMode._Enum.ValueType # 1 DRAINING: Elements.DrainMode._Enum.ValueType # 2 class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... UNSPECIFIED: Elements.DrainMode.Enum.ValueType # 0 NOT_DRAINING: Elements.DrainMode.Enum.ValueType # 1 DRAINING: Elements.DrainMode.Enum.ValueType # 2 def __init__( self, ) -> None: ... @typing_extensions.final class ValueKind(google.protobuf.message.Message): """The type of change operation represented by a Change Data Capture (CDC) record""" DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[Elements.ValueKind._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor VALUE_KIND_UNSPECIFIED: Elements.ValueKind._Enum.ValueType # 0 """Unspecified ValueKind. Will be treated as an INSERT""" INSERT: Elements.ValueKind._Enum.ValueType # 1 """Indicates a new record was created in the source system.""" UPDATE_BEFORE: Elements.ValueKind._Enum.ValueType # 2 """Indicates the state of a record immediately before an update occurred. This is typically used to identify the previous values of modified columns or to locate the record via its primary key. """ UPDATE_AFTER: Elements.ValueKind._Enum.ValueType # 3 """Indicates the state of a record immediately after an update occurred. Represents the current, valid state of the record following the change. """ DELETE: Elements.ValueKind._Enum.ValueType # 4 """Indicates that an existing record was removed from the source system.""" class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... VALUE_KIND_UNSPECIFIED: Elements.ValueKind.Enum.ValueType # 0 """Unspecified ValueKind. Will be treated as an INSERT""" INSERT: Elements.ValueKind.Enum.ValueType # 1 """Indicates a new record was created in the source system.""" UPDATE_BEFORE: Elements.ValueKind.Enum.ValueType # 2 """Indicates the state of a record immediately before an update occurred. This is typically used to identify the previous values of modified columns or to locate the record via its primary key. """ UPDATE_AFTER: Elements.ValueKind.Enum.ValueType # 3 """Indicates the state of a record immediately after an update occurred. Represents the current, valid state of the record following the change. """ DELETE: Elements.ValueKind.Enum.ValueType # 4 """Indicates that an existing record was removed from the source system.""" def __init__( self, ) -> None: ... @typing_extensions.final class ElementMetadata(google.protobuf.message.Message): """Element metadata passed as part of WindowedValue to make WindowedValue extensible and backward compatible """ DESCRIPTOR: google.protobuf.descriptor.Descriptor DRAIN_FIELD_NUMBER: builtins.int TRACEPARENT_FIELD_NUMBER: builtins.int TRACESTATE_FIELD_NUMBER: builtins.int VALUE_KIND_FIELD_NUMBER: builtins.int drain: global___Elements.DrainMode.Enum.ValueType traceparent: builtins.str """(Optional) As part of https://www.w3.org/TR/trace-context/ we are forwarding a trace and participating in it. Traceparent header represents the incoming request in a tracing system in a common format. Example value: 00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01 """ tracestate: builtins.str """(Optional) tracestate extends traceparent with open telemetry data represented by a set of name/value pairs. Format specified https://www.w3.org/TR/trace-context/#list for interoperability and commonly used across IOs - Kafka, PubSub, http. Example value: congo=t61rcWkgMzE """ value_kind: global___Elements.ValueKind.Enum.ValueType """(Optional) The kind of value for CDC metadata. If missing or unspecified, implies INSERT for backwards compatibility. """ def __init__( self, *, drain: global___Elements.DrainMode.Enum.ValueType | None = ..., traceparent: builtins.str | None = ..., tracestate: builtins.str | None = ..., value_kind: global___Elements.ValueKind.Enum.ValueType | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["_drain", b"_drain", "_traceparent", b"_traceparent", "_tracestate", b"_tracestate", "_value_kind", b"_value_kind", "drain", b"drain", "traceparent", b"traceparent", "tracestate", b"tracestate", "value_kind", b"value_kind"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["_drain", b"_drain", "_traceparent", b"_traceparent", "_tracestate", b"_tracestate", "_value_kind", b"_value_kind", "drain", b"drain", "traceparent", b"traceparent", "tracestate", b"tracestate", "value_kind", b"value_kind"]) -> None: ... @typing.overload def WhichOneof(self, oneof_group: typing_extensions.Literal["_drain", b"_drain"]) -> typing_extensions.Literal["drain"] | None: ... @typing.overload def WhichOneof(self, oneof_group: typing_extensions.Literal["_traceparent", b"_traceparent"]) -> typing_extensions.Literal["traceparent"] | None: ... @typing.overload def WhichOneof(self, oneof_group: typing_extensions.Literal["_tracestate", b"_tracestate"]) -> typing_extensions.Literal["tracestate"] | None: ... @typing.overload def WhichOneof(self, oneof_group: typing_extensions.Literal["_value_kind", b"_value_kind"]) -> typing_extensions.Literal["value_kind"] | None: ... @typing_extensions.final class Timers(google.protobuf.message.Message): """Represent the encoded user timer for a given instruction, transform and timer id. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor INSTRUCTION_ID_FIELD_NUMBER: builtins.int TRANSFORM_ID_FIELD_NUMBER: builtins.int TIMER_FAMILY_ID_FIELD_NUMBER: builtins.int TIMERS_FIELD_NUMBER: builtins.int IS_LAST_FIELD_NUMBER: builtins.int instruction_id: builtins.str """(Required) A reference to an active instruction request with the given instruction id. """ transform_id: builtins.str """(Required) A definition representing a consumer or producer of this data. If received by a harness, this represents the consumer within that harness that should consume these timers. If sent by a harness, this represents the producer of these timers. """ timer_family_id: builtins.str """(Required) The local timer family name used to identify the associated timer family specification """ timers: builtins.bytes """(Optional) Represents a logical byte stream of timers. Encoded according to the coder in the timer spec. """ is_last: builtins.bool """(Optional) Set this bit to indicate the this is the last data block for the given instruction and transform, ending the stream. """ def __init__( self, *, instruction_id: builtins.str | None = ..., transform_id: builtins.str | None = ..., timer_family_id: builtins.str | None = ..., timers: builtins.bytes | None = ..., is_last: builtins.bool | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["instruction_id", b"instruction_id", "is_last", b"is_last", "timer_family_id", b"timer_family_id", "timers", b"timers", "transform_id", b"transform_id"]) -> None: ... DATA_FIELD_NUMBER: builtins.int TIMERS_FIELD_NUMBER: builtins.int @property def data(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Elements.Data]: """(Optional) A list containing parts of logical byte streams.""" @property def timers(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Elements.Timers]: """(Optional) A list of timer byte streams.""" def __init__( self, *, data: collections.abc.Iterable[global___Elements.Data] | None = ..., timers: collections.abc.Iterable[global___Elements.Timers] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["data", b"data", "timers", b"timers"]) -> None: ... global___Elements = Elements @typing_extensions.final class StateRequest(google.protobuf.message.Message): """ State API """ DESCRIPTOR: google.protobuf.descriptor.Descriptor ID_FIELD_NUMBER: builtins.int INSTRUCTION_ID_FIELD_NUMBER: builtins.int STATE_KEY_FIELD_NUMBER: builtins.int GET_FIELD_NUMBER: builtins.int APPEND_FIELD_NUMBER: builtins.int CLEAR_FIELD_NUMBER: builtins.int id: builtins.str """(Required) A unique identifier provided by the SDK which represents this requests execution. The StateResponse corresponding with this request will have the matching id. """ instruction_id: builtins.str """(Required) The associated instruction id of the work that is currently being processed. This allows for the runner to associate any modifications to state to be committed with the appropriate work execution. """ @property def state_key(self) -> global___StateKey: """(Required) The state key this request is for.""" @property def get(self) -> global___StateGetRequest: """A request to get state.""" @property def append(self) -> global___StateAppendRequest: """A request to append to state.""" @property def clear(self) -> global___StateClearRequest: """A request to clear state.""" def __init__( self, *, id: builtins.str | None = ..., instruction_id: builtins.str | None = ..., state_key: global___StateKey | None = ..., get: global___StateGetRequest | None = ..., append: global___StateAppendRequest | None = ..., clear: global___StateClearRequest | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["append", b"append", "clear", b"clear", "get", b"get", "request", b"request", "state_key", b"state_key"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["append", b"append", "clear", b"clear", "get", b"get", "id", b"id", "instruction_id", b"instruction_id", "request", b"request", "state_key", b"state_key"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["request", b"request"]) -> typing_extensions.Literal["get", "append", "clear"] | None: ... global___StateRequest = StateRequest @typing_extensions.final class StateResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ID_FIELD_NUMBER: builtins.int ERROR_FIELD_NUMBER: builtins.int GET_FIELD_NUMBER: builtins.int APPEND_FIELD_NUMBER: builtins.int CLEAR_FIELD_NUMBER: builtins.int id: builtins.str """(Required) A reference provided by the SDK which represents a requests execution. The StateResponse must have the matching id when responding to the SDK. """ error: builtins.str """(Optional) If this is specified, then the state request has failed. A human readable string representing the reason as to why the request failed. """ @property def get(self) -> global___StateGetResponse: """A response to getting state.""" @property def append(self) -> global___StateAppendResponse: """A response to appending to state.""" @property def clear(self) -> global___StateClearResponse: """A response to clearing state.""" def __init__( self, *, id: builtins.str | None = ..., error: builtins.str | None = ..., get: global___StateGetResponse | None = ..., append: global___StateAppendResponse | None = ..., clear: global___StateClearResponse | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["append", b"append", "clear", b"clear", "get", b"get", "response", b"response"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["append", b"append", "clear", b"clear", "error", b"error", "get", b"get", "id", b"id", "response", b"response"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["response", b"response"]) -> typing_extensions.Literal["get", "append", "clear"] | None: ... global___StateResponse = StateResponse @typing_extensions.final class StateKey(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class Runner(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int key: builtins.bytes """(Required) Opaque information supplied by the runner. Used to support remote references. https://s.apache.org/beam-fn-api-send-and-receive-data Used by state backed iterable. And in this use case, request type can only be of type get. Details see: https://s.apache.org/beam-fn-api-state-backed-iterables """ def __init__( self, *, key: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key"]) -> None: ... @typing_extensions.final class IterableSideInput(google.protobuf.message.Message): """Represents a request for the values associated with a specified window in a PCollection. See https://s.apache.org/beam-fn-state-api-and-bundle-processing for further details. Can only be used to perform StateGetRequests on side inputs of the URN beam:side_input:iterable:v1. For a PCollection, the response data stream will be a concatenation of all V's. See https://s.apache.org/beam-fn-api-send-and-receive-data for further details. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int SIDE_INPUT_ID_FIELD_NUMBER: builtins.int WINDOW_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The id of the PTransform containing a side input.""" side_input_id: builtins.str """(Required) The id of the side input.""" window: builtins.bytes """(Required) The window (after mapping the currently executing elements window into the side input windows domain) encoded in a nested context. """ def __init__( self, *, transform_id: builtins.str | None = ..., side_input_id: builtins.str | None = ..., window: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["side_input_id", b"side_input_id", "transform_id", b"transform_id", "window", b"window"]) -> None: ... @typing_extensions.final class MultimapSideInput(google.protobuf.message.Message): """Represents a request for the values associated with a specified user key and window in a PCollection. See https://s.apache.org/beam-fn-state-api-and-bundle-processing for further details. Can only be used to perform StateGetRequests on side inputs of the URN beam:side_input:multimap:v1. For a PCollection>, the response data stream will be a concatenation of all V's associated with the specified key K. See https://s.apache.org/beam-fn-api-send-and-receive-data for further details. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int SIDE_INPUT_ID_FIELD_NUMBER: builtins.int WINDOW_FIELD_NUMBER: builtins.int KEY_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The id of the PTransform containing a side input.""" side_input_id: builtins.str """(Required) The id of the side input.""" window: builtins.bytes """(Required) The window (after mapping the currently executing elements window into the side input windows domain) encoded in a nested context. """ key: builtins.bytes """(Required) The key encoded in a nested context.""" def __init__( self, *, transform_id: builtins.str | None = ..., side_input_id: builtins.str | None = ..., window: builtins.bytes | None = ..., key: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "side_input_id", b"side_input_id", "transform_id", b"transform_id", "window", b"window"]) -> None: ... @typing_extensions.final class MultimapKeysSideInput(google.protobuf.message.Message): """Represents a request for the keys associated with a specified window in a PCollection. See https://s.apache.org/beam-fn-state-api-and-bundle-processing for further details. Can only be used to perform StateGetRequests on side inputs of the URN beam:side_input:multimap:v1. For a PCollection>, the response data stream will be a concatenation of all K's associated with the specified window. See https://s.apache.org/beam-fn-api-send-and-receive-data for further details. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int SIDE_INPUT_ID_FIELD_NUMBER: builtins.int WINDOW_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The id of the PTransform containing a side input.""" side_input_id: builtins.str """(Required) The id of the side input.""" window: builtins.bytes """(Required) The window (after mapping the currently executing elements window into the side input windows domain) encoded in a nested context. """ def __init__( self, *, transform_id: builtins.str | None = ..., side_input_id: builtins.str | None = ..., window: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["side_input_id", b"side_input_id", "transform_id", b"transform_id", "window", b"window"]) -> None: ... @typing_extensions.final class MultimapKeysValuesSideInput(google.protobuf.message.Message): """Represents a request for the keys and values associated with a specified window in a PCollection. See https://s.apache.org/beam-fn-state-api-and-bundle-processing for further details. This is expected to be more efficient than iterating over they keys and looking up the values one at a time. If a runner chooses not to implement this protocol, or a key has too many values to fit into a single response, the runner is free to fail the request and a fallback to point lookups will be performed by the SDK. Can only be used to perform StateGetRequests on side inputs of the URN beam:side_input:multimap:v1. For a PCollection>, the response data stream will be a concatenation of all KVs associated with the specified window, encoded with the the KV> coder. See https://s.apache.org/beam-fn-api-send-and-receive-data for further details. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int SIDE_INPUT_ID_FIELD_NUMBER: builtins.int WINDOW_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The id of the PTransform containing a side input.""" side_input_id: builtins.str """(Required) The id of the side input.""" window: builtins.bytes """(Required) The window (after mapping the currently executing elements window into the side input windows domain) encoded in a nested context. """ def __init__( self, *, transform_id: builtins.str | None = ..., side_input_id: builtins.str | None = ..., window: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["side_input_id", b"side_input_id", "transform_id", b"transform_id", "window", b"window"]) -> None: ... @typing_extensions.final class BagUserState(google.protobuf.message.Message): """Represents a request for an unordered set of values associated with a specified user key and window for a PTransform. See https://s.apache.org/beam-fn-state-api-and-bundle-processing for further details. The response data stream will be a concatenation of all V's associated with the specified user key and window. See https://s.apache.org/beam-fn-api-send-and-receive-data for further details. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int USER_STATE_ID_FIELD_NUMBER: builtins.int WINDOW_FIELD_NUMBER: builtins.int KEY_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The id of the PTransform containing user state.""" user_state_id: builtins.str """(Required) The id of the user state.""" window: builtins.bytes """(Required) The window encoded in a nested context.""" key: builtins.bytes """(Required) The key of the currently executing element encoded in a nested context. """ def __init__( self, *, transform_id: builtins.str | None = ..., user_state_id: builtins.str | None = ..., window: builtins.bytes | None = ..., key: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "transform_id", b"transform_id", "user_state_id", b"user_state_id", "window", b"window"]) -> None: ... @typing_extensions.final class MultimapKeysUserState(google.protobuf.message.Message): """Represents a request for the keys of a multimap associated with a specified user key and window for a PTransform. See https://s.apache.org/beam-fn-state-api-and-bundle-processing for further details. Can only be used to perform StateGetRequests and StateClearRequests on the user state. The response data stream will be a concatenation of all K's associated with the specified user key and window. See https://s.apache.org/beam-fn-api-send-and-receive-data for further details. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int USER_STATE_ID_FIELD_NUMBER: builtins.int WINDOW_FIELD_NUMBER: builtins.int KEY_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The id of the PTransform containing user state.""" user_state_id: builtins.str """(Required) The id of the user state.""" window: builtins.bytes """(Required) The window encoded in a nested context.""" key: builtins.bytes """(Required) The key of the currently executing element encoded in a nested context. """ def __init__( self, *, transform_id: builtins.str | None = ..., user_state_id: builtins.str | None = ..., window: builtins.bytes | None = ..., key: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "transform_id", b"transform_id", "user_state_id", b"user_state_id", "window", b"window"]) -> None: ... @typing_extensions.final class MultimapEntriesUserState(google.protobuf.message.Message): """Represents a request for all of the entries of a multimap associated with a specified user key and window for a PTransform. See https://s.apache.org/beam-fn-state-api-and-bundle-processing for further details. Can only be used to perform StateGetRequests and StateClearRequests on the user state. The response data stream will be a concatenation of pairs, where the first component is the map key and the second component is a concatenation of values associated with that map key. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int USER_STATE_ID_FIELD_NUMBER: builtins.int WINDOW_FIELD_NUMBER: builtins.int KEY_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The id of the PTransform containing user state.""" user_state_id: builtins.str """(Required) The id of the user state.""" window: builtins.bytes """(Required) The window encoded in a nested context.""" key: builtins.bytes """(Required) The key of the currently executing element encoded in a nested context. """ def __init__( self, *, transform_id: builtins.str | None = ..., user_state_id: builtins.str | None = ..., window: builtins.bytes | None = ..., key: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "transform_id", b"transform_id", "user_state_id", b"user_state_id", "window", b"window"]) -> None: ... @typing_extensions.final class MultimapUserState(google.protobuf.message.Message): """Represents a request for the values of the map key associated with a specified user key and window for a PTransform. See https://s.apache.org/beam-fn-state-api-and-bundle-processing for further details. The response data stream will be a concatenation of all V's associated with the specified map key, user key, and window. See https://s.apache.org/beam-fn-api-send-and-receive-data for further details. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int USER_STATE_ID_FIELD_NUMBER: builtins.int WINDOW_FIELD_NUMBER: builtins.int KEY_FIELD_NUMBER: builtins.int MAP_KEY_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The id of the PTransform containing user state.""" user_state_id: builtins.str """(Required) The id of the user state.""" window: builtins.bytes """(Required) The window encoded in a nested context.""" key: builtins.bytes """(Required) The key of the currently executing element encoded in a nested context. """ map_key: builtins.bytes """(Required) The map key encoded in a nested context.""" def __init__( self, *, transform_id: builtins.str | None = ..., user_state_id: builtins.str | None = ..., window: builtins.bytes | None = ..., key: builtins.bytes | None = ..., map_key: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "map_key", b"map_key", "transform_id", b"transform_id", "user_state_id", b"user_state_id", "window", b"window"]) -> None: ... @typing_extensions.final class OrderedListUserState(google.protobuf.message.Message): """Represents a request for an ordered list of values associated with a specified user key and window for a PTransform. See https://s.apache.org/beam-fn-state-api-and-bundle-processing for further details. The response data stream will be a concatenation of all entries of sort key and V's associated with the specified user key and window. See https://s.apache.org/beam-fn-api-send-and-receive-data for further details. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int USER_STATE_ID_FIELD_NUMBER: builtins.int WINDOW_FIELD_NUMBER: builtins.int KEY_FIELD_NUMBER: builtins.int RANGE_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The id of the PTransform containing user state.""" user_state_id: builtins.str """(Required) The id of the user state.""" window: builtins.bytes """(Required) The window encoded in a nested context.""" key: builtins.bytes """(Required) The key of the currently executing element encoded in a nested context. """ @property def range(self) -> global___OrderedListRange: """(Required) The sort range encoded in a nested context.""" def __init__( self, *, transform_id: builtins.str | None = ..., user_state_id: builtins.str | None = ..., window: builtins.bytes | None = ..., key: builtins.bytes | None = ..., range: global___OrderedListRange | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["range", b"range"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "range", b"range", "transform_id", b"transform_id", "user_state_id", b"user_state_id", "window", b"window"]) -> None: ... RUNNER_FIELD_NUMBER: builtins.int MULTIMAP_SIDE_INPUT_FIELD_NUMBER: builtins.int BAG_USER_STATE_FIELD_NUMBER: builtins.int ITERABLE_SIDE_INPUT_FIELD_NUMBER: builtins.int MULTIMAP_KEYS_SIDE_INPUT_FIELD_NUMBER: builtins.int MULTIMAP_KEYS_VALUES_SIDE_INPUT_FIELD_NUMBER: builtins.int MULTIMAP_KEYS_USER_STATE_FIELD_NUMBER: builtins.int MULTIMAP_ENTRIES_USER_STATE_FIELD_NUMBER: builtins.int MULTIMAP_USER_STATE_FIELD_NUMBER: builtins.int ORDERED_LIST_USER_STATE_FIELD_NUMBER: builtins.int @property def runner(self) -> global___StateKey.Runner: ... @property def multimap_side_input(self) -> global___StateKey.MultimapSideInput: ... @property def bag_user_state(self) -> global___StateKey.BagUserState: ... @property def iterable_side_input(self) -> global___StateKey.IterableSideInput: ... @property def multimap_keys_side_input(self) -> global___StateKey.MultimapKeysSideInput: ... @property def multimap_keys_values_side_input(self) -> global___StateKey.MultimapKeysValuesSideInput: ... @property def multimap_keys_user_state(self) -> global___StateKey.MultimapKeysUserState: ... @property def multimap_entries_user_state(self) -> global___StateKey.MultimapEntriesUserState: ... @property def multimap_user_state(self) -> global___StateKey.MultimapUserState: ... @property def ordered_list_user_state(self) -> global___StateKey.OrderedListUserState: ... def __init__( self, *, runner: global___StateKey.Runner | None = ..., multimap_side_input: global___StateKey.MultimapSideInput | None = ..., bag_user_state: global___StateKey.BagUserState | None = ..., iterable_side_input: global___StateKey.IterableSideInput | None = ..., multimap_keys_side_input: global___StateKey.MultimapKeysSideInput | None = ..., multimap_keys_values_side_input: global___StateKey.MultimapKeysValuesSideInput | None = ..., multimap_keys_user_state: global___StateKey.MultimapKeysUserState | None = ..., multimap_entries_user_state: global___StateKey.MultimapEntriesUserState | None = ..., multimap_user_state: global___StateKey.MultimapUserState | None = ..., ordered_list_user_state: global___StateKey.OrderedListUserState | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["bag_user_state", b"bag_user_state", "iterable_side_input", b"iterable_side_input", "multimap_entries_user_state", b"multimap_entries_user_state", "multimap_keys_side_input", b"multimap_keys_side_input", "multimap_keys_user_state", b"multimap_keys_user_state", "multimap_keys_values_side_input", b"multimap_keys_values_side_input", "multimap_side_input", b"multimap_side_input", "multimap_user_state", b"multimap_user_state", "ordered_list_user_state", b"ordered_list_user_state", "runner", b"runner", "type", b"type"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["bag_user_state", b"bag_user_state", "iterable_side_input", b"iterable_side_input", "multimap_entries_user_state", b"multimap_entries_user_state", "multimap_keys_side_input", b"multimap_keys_side_input", "multimap_keys_user_state", b"multimap_keys_user_state", "multimap_keys_values_side_input", b"multimap_keys_values_side_input", "multimap_side_input", b"multimap_side_input", "multimap_user_state", b"multimap_user_state", "ordered_list_user_state", b"ordered_list_user_state", "runner", b"runner", "type", b"type"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["type", b"type"]) -> typing_extensions.Literal["runner", "multimap_side_input", "bag_user_state", "iterable_side_input", "multimap_keys_side_input", "multimap_keys_values_side_input", "multimap_keys_user_state", "multimap_entries_user_state", "multimap_user_state", "ordered_list_user_state"] | None: ... global___StateKey = StateKey @typing_extensions.final class StateGetRequest(google.protobuf.message.Message): """A request to get state.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor CONTINUATION_TOKEN_FIELD_NUMBER: builtins.int continuation_token: builtins.bytes """(Optional) If specified, signals to the runner that the response should resume from the following continuation token. If unspecified, signals to the runner that the response should start from the beginning of the logical continuable stream. """ def __init__( self, *, continuation_token: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["continuation_token", b"continuation_token"]) -> None: ... global___StateGetRequest = StateGetRequest @typing_extensions.final class StateGetResponse(google.protobuf.message.Message): """A response to get state representing a logical byte stream which can be continued using the state API. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor CONTINUATION_TOKEN_FIELD_NUMBER: builtins.int DATA_FIELD_NUMBER: builtins.int continuation_token: builtins.bytes """(Optional) If specified, represents a token which can be used with the state API to get the next chunk of this logical byte stream. The end of the logical byte stream is signalled by this field being unset. """ data: builtins.bytes """Represents a part of a logical byte stream. Elements within the logical byte stream are encoded in the nested context and concatenated together. See also the note about OrderedListState in StateAppendRequest. """ def __init__( self, *, continuation_token: builtins.bytes | None = ..., data: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["continuation_token", b"continuation_token", "data", b"data"]) -> None: ... global___StateGetResponse = StateGetResponse @typing_extensions.final class StateAppendRequest(google.protobuf.message.Message): """A request to append state.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor DATA_FIELD_NUMBER: builtins.int data: builtins.bytes """Represents a part of a logical byte stream. Elements within the logical byte stream are encoded in the nested context and multiple append requests are concatenated together. For OrderedListState, elements of should be encoded with the beam:coder:kv:v1 coder, where the first (key) component must be a beam:coder:varint:v1 and the second (value) component must be encoded with a beam:coder:length_prefix:v1 coder. """ def __init__( self, *, data: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["data", b"data"]) -> None: ... global___StateAppendRequest = StateAppendRequest @typing_extensions.final class StateAppendResponse(google.protobuf.message.Message): """A response to append state.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... global___StateAppendResponse = StateAppendResponse @typing_extensions.final class StateClearRequest(google.protobuf.message.Message): """A request to clear state.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... global___StateClearRequest = StateClearRequest @typing_extensions.final class StateClearResponse(google.protobuf.message.Message): """A response to clear state.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... global___StateClearResponse = StateClearResponse @typing_extensions.final class OrderedListRange(google.protobuf.message.Message): """A message describes a sort key range [start, end).""" DESCRIPTOR: google.protobuf.descriptor.Descriptor START_FIELD_NUMBER: builtins.int END_FIELD_NUMBER: builtins.int start: builtins.int end: builtins.int def __init__( self, *, start: builtins.int | None = ..., end: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["end", b"end", "start", b"start"]) -> None: ... global___OrderedListRange = OrderedListRange @typing_extensions.final class LogEntry(google.protobuf.message.Message): """ Logging API This is very stable. There can be some changes to how we define a LogEntry, to increase/decrease the severity types, the way we format an exception/stack trace, or the log site. A log entry """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class List(google.protobuf.message.Message): """A list of log entries, enables buffering and batching of multiple log messages using the logging API. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor LOG_ENTRIES_FIELD_NUMBER: builtins.int @property def log_entries(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___LogEntry]: """(Required) One or or more log messages.""" def __init__( self, *, log_entries: collections.abc.Iterable[global___LogEntry] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["log_entries", b"log_entries"]) -> None: ... @typing_extensions.final class Severity(google.protobuf.message.Message): """The severity of the event described in a log entry, expressed as one of the severity levels listed below. For your reference, the levels are assigned the listed numeric values. The effect of using numeric values other than those listed is undefined. If you are writing log entries, you should map other severity encodings to one of these standard levels. For example, you might map all of Java's FINE, FINER, and FINEST levels to `Severity.DEBUG`. This list is intentionally not comprehensive; the intent is to provide a common set of "good enough" severity levels so that logging front ends can provide filtering and searching across log types. Users of the API are free not to use all severity levels in their log messages. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[LogEntry.Severity._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor UNSPECIFIED: LogEntry.Severity._Enum.ValueType # 0 """Unspecified level information. Will be logged at the TRACE level.""" TRACE: LogEntry.Severity._Enum.ValueType # 1 DEBUG: LogEntry.Severity._Enum.ValueType # 2 """Debugging information.""" INFO: LogEntry.Severity._Enum.ValueType # 3 """Normal events.""" NOTICE: LogEntry.Severity._Enum.ValueType # 4 """Normal but significant events, such as start up, shut down, or configuration. """ WARN: LogEntry.Severity._Enum.ValueType # 5 """Warning events might cause problems.""" ERROR: LogEntry.Severity._Enum.ValueType # 6 """Error events are likely to cause problems.""" CRITICAL: LogEntry.Severity._Enum.ValueType # 7 """Critical events cause severe problems or brief outages and may indicate that a person must take action. """ class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... UNSPECIFIED: LogEntry.Severity.Enum.ValueType # 0 """Unspecified level information. Will be logged at the TRACE level.""" TRACE: LogEntry.Severity.Enum.ValueType # 1 DEBUG: LogEntry.Severity.Enum.ValueType # 2 """Debugging information.""" INFO: LogEntry.Severity.Enum.ValueType # 3 """Normal events.""" NOTICE: LogEntry.Severity.Enum.ValueType # 4 """Normal but significant events, such as start up, shut down, or configuration. """ WARN: LogEntry.Severity.Enum.ValueType # 5 """Warning events might cause problems.""" ERROR: LogEntry.Severity.Enum.ValueType # 6 """Error events are likely to cause problems.""" CRITICAL: LogEntry.Severity.Enum.ValueType # 7 """Critical events cause severe problems or brief outages and may indicate that a person must take action. """ def __init__( self, ) -> None: ... SEVERITY_FIELD_NUMBER: builtins.int TIMESTAMP_FIELD_NUMBER: builtins.int MESSAGE_FIELD_NUMBER: builtins.int TRACE_FIELD_NUMBER: builtins.int INSTRUCTION_ID_FIELD_NUMBER: builtins.int TRANSFORM_ID_FIELD_NUMBER: builtins.int LOG_LOCATION_FIELD_NUMBER: builtins.int THREAD_FIELD_NUMBER: builtins.int CUSTOM_DATA_FIELD_NUMBER: builtins.int severity: global___LogEntry.Severity.Enum.ValueType """(Required) The severity of the log statement.""" @property def timestamp(self) -> google.protobuf.timestamp_pb2.Timestamp: """(Required) The time at which this log statement occurred.""" message: builtins.str """(Required) A human readable message.""" trace: builtins.str """(Optional) An optional trace of the functions involved. For example, in Java this can include multiple causes and multiple suppressed exceptions. """ instruction_id: builtins.str """(Optional) A reference to the instruction this log statement is associated with. """ transform_id: builtins.str """(Optional) A reference to the transform this log statement is associated with. """ log_location: builtins.str """(Optional) Human-readable name of the function or method being invoked, with optional context such as the class or package name. The format can vary by language. For example: qual.if.ied.Class.method (Java) dir/package.func (Go) module.function (Python) file.cc:382 (C++) """ thread: builtins.str """(Optional) The name of the thread this log statement is associated with.""" @property def custom_data(self) -> google.protobuf.struct_pb2.Struct: """(Optional) Additional structured data to log. Keys are limited to these characters: [a-zA-Z_-] """ def __init__( self, *, severity: global___LogEntry.Severity.Enum.ValueType | None = ..., timestamp: google.protobuf.timestamp_pb2.Timestamp | None = ..., message: builtins.str | None = ..., trace: builtins.str | None = ..., instruction_id: builtins.str | None = ..., transform_id: builtins.str | None = ..., log_location: builtins.str | None = ..., thread: builtins.str | None = ..., custom_data: google.protobuf.struct_pb2.Struct | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["custom_data", b"custom_data", "timestamp", b"timestamp"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["custom_data", b"custom_data", "instruction_id", b"instruction_id", "log_location", b"log_location", "message", b"message", "severity", b"severity", "thread", b"thread", "timestamp", b"timestamp", "trace", b"trace", "transform_id", b"transform_id"]) -> None: ... global___LogEntry = LogEntry @typing_extensions.final class LogControl(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... global___LogControl = LogControl @typing_extensions.final class StartWorkerRequest(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class ParamsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str value: builtins.str def __init__( self, *, key: builtins.str | None = ..., value: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... WORKER_ID_FIELD_NUMBER: builtins.int CONTROL_ENDPOINT_FIELD_NUMBER: builtins.int LOGGING_ENDPOINT_FIELD_NUMBER: builtins.int ARTIFACT_ENDPOINT_FIELD_NUMBER: builtins.int PROVISION_ENDPOINT_FIELD_NUMBER: builtins.int PARAMS_FIELD_NUMBER: builtins.int worker_id: builtins.str @property def control_endpoint(self) -> org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor: ... @property def logging_endpoint(self) -> org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor: ... @property def artifact_endpoint(self) -> org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor: ... @property def provision_endpoint(self) -> org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor: ... @property def params(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: ... def __init__( self, *, worker_id: builtins.str | None = ..., control_endpoint: org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor | None = ..., logging_endpoint: org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor | None = ..., artifact_endpoint: org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor | None = ..., provision_endpoint: org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor | None = ..., params: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["artifact_endpoint", b"artifact_endpoint", "control_endpoint", b"control_endpoint", "logging_endpoint", b"logging_endpoint", "provision_endpoint", b"provision_endpoint"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["artifact_endpoint", b"artifact_endpoint", "control_endpoint", b"control_endpoint", "logging_endpoint", b"logging_endpoint", "params", b"params", "provision_endpoint", b"provision_endpoint", "worker_id", b"worker_id"]) -> None: ... global___StartWorkerRequest = StartWorkerRequest @typing_extensions.final class StartWorkerResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ERROR_FIELD_NUMBER: builtins.int error: builtins.str def __init__( self, *, error: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["error", b"error"]) -> None: ... global___StartWorkerResponse = StartWorkerResponse @typing_extensions.final class StopWorkerRequest(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor WORKER_ID_FIELD_NUMBER: builtins.int worker_id: builtins.str def __init__( self, *, worker_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["worker_id", b"worker_id"]) -> None: ... global___StopWorkerRequest = StopWorkerRequest @typing_extensions.final class StopWorkerResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ERROR_FIELD_NUMBER: builtins.int error: builtins.str def __init__( self, *, error: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["error", b"error"]) -> None: ... global___StopWorkerResponse = StopWorkerResponse @typing_extensions.final class WorkerStatusRequest(google.protobuf.message.Message): """Request from runner to SDK Harness asking for its status. For more details see https://s.apache.org/beam-fn-api-harness-status """ DESCRIPTOR: google.protobuf.descriptor.Descriptor ID_FIELD_NUMBER: builtins.int id: builtins.str """(Required) Unique ID identifying this request.""" def __init__( self, *, id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["id", b"id"]) -> None: ... global___WorkerStatusRequest = WorkerStatusRequest @typing_extensions.final class WorkerStatusResponse(google.protobuf.message.Message): """Response from SDK Harness to runner containing the debug related status info.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor ID_FIELD_NUMBER: builtins.int ERROR_FIELD_NUMBER: builtins.int STATUS_INFO_FIELD_NUMBER: builtins.int id: builtins.str """(Required) Unique ID from the original request.""" error: builtins.str """(Optional) Error message if exception encountered generating the status response.""" status_info: builtins.str """(Optional) Status debugging info reported by SDK harness worker. Content and format is not strongly enforced but should be print-friendly and appropriate as an HTTP response body for end user. For details of the preferred info to include in the message see https://s.apache.org/beam-fn-api-harness-status """ def __init__( self, *, id: builtins.str | None = ..., error: builtins.str | None = ..., status_info: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["error", b"error", "id", b"id", "status_info", b"status_info"]) -> None: ... global___WorkerStatusResponse = WorkerStatusResponse ././@PaxHeader0000000000000000000000000000021400000000000010212 xustar00118 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_fn_api_pb2_grpc.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_fn_api_pb20000644000175100017510000005236215206113151032740 0ustar00runnerrunner# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc from . import beam_fn_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2 class BeamFnControlStub(object): """ Control Plane API Progress reporting and splitting still need further vetting. Also, this may change with the addition of new types of instructions/responses related to metrics. An API that describes the work that a SDK harness is meant to do. Stable """ def __init__(self, channel): """Constructor. Args: channel: A grpc.Channel. """ self.Control = channel.stream_stream( '/org.apache.beam.model.fn_execution.v1.BeamFnControl/Control', request_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.InstructionResponse.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.InstructionRequest.FromString, ) self.GetProcessBundleDescriptor = channel.unary_unary( '/org.apache.beam.model.fn_execution.v1.BeamFnControl/GetProcessBundleDescriptor', request_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.GetProcessBundleDescriptorRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.ProcessBundleDescriptor.FromString, ) class BeamFnControlServicer(object): """ Control Plane API Progress reporting and splitting still need further vetting. Also, this may change with the addition of new types of instructions/responses related to metrics. An API that describes the work that a SDK harness is meant to do. Stable """ def Control(self, request_iterator, context): """Instructions sent by the runner to the SDK requesting different types of work. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def GetProcessBundleDescriptor(self, request, context): """Used to get the full process bundle descriptors for bundles one is asked to process. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def add_BeamFnControlServicer_to_server(servicer, server): rpc_method_handlers = { 'Control': grpc.stream_stream_rpc_method_handler( servicer.Control, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.InstructionResponse.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.InstructionRequest.SerializeToString, ), 'GetProcessBundleDescriptor': grpc.unary_unary_rpc_method_handler( servicer.GetProcessBundleDescriptor, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.GetProcessBundleDescriptorRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.ProcessBundleDescriptor.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( 'org.apache.beam.model.fn_execution.v1.BeamFnControl', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) # This class is part of an EXPERIMENTAL API. class BeamFnControl(object): """ Control Plane API Progress reporting and splitting still need further vetting. Also, this may change with the addition of new types of instructions/responses related to metrics. An API that describes the work that a SDK harness is meant to do. Stable """ @staticmethod def Control(request_iterator, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.stream_stream(request_iterator, target, '/org.apache.beam.model.fn_execution.v1.BeamFnControl/Control', org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.InstructionResponse.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.InstructionRequest.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def GetProcessBundleDescriptor(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.fn_execution.v1.BeamFnControl/GetProcessBundleDescriptor', org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.GetProcessBundleDescriptorRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.ProcessBundleDescriptor.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) class BeamFnDataStub(object): """Stable """ def __init__(self, channel): """Constructor. Args: channel: A grpc.Channel. """ self.Data = channel.stream_stream( '/org.apache.beam.model.fn_execution.v1.BeamFnData/Data', request_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.Elements.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.Elements.FromString, ) class BeamFnDataServicer(object): """Stable """ def Data(self, request_iterator, context): """Used to send data between harnesses. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def add_BeamFnDataServicer_to_server(servicer, server): rpc_method_handlers = { 'Data': grpc.stream_stream_rpc_method_handler( servicer.Data, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.Elements.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.Elements.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( 'org.apache.beam.model.fn_execution.v1.BeamFnData', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) # This class is part of an EXPERIMENTAL API. class BeamFnData(object): """Stable """ @staticmethod def Data(request_iterator, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.stream_stream(request_iterator, target, '/org.apache.beam.model.fn_execution.v1.BeamFnData/Data', org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.Elements.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.Elements.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) class BeamFnStateStub(object): """Missing associated documentation comment in .proto file.""" def __init__(self, channel): """Constructor. Args: channel: A grpc.Channel. """ self.State = channel.stream_stream( '/org.apache.beam.model.fn_execution.v1.BeamFnState/State', request_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StateRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StateResponse.FromString, ) class BeamFnStateServicer(object): """Missing associated documentation comment in .proto file.""" def State(self, request_iterator, context): """Used to get/append/clear state stored by the runner on behalf of the SDK. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def add_BeamFnStateServicer_to_server(servicer, server): rpc_method_handlers = { 'State': grpc.stream_stream_rpc_method_handler( servicer.State, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StateRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StateResponse.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( 'org.apache.beam.model.fn_execution.v1.BeamFnState', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) # This class is part of an EXPERIMENTAL API. class BeamFnState(object): """Missing associated documentation comment in .proto file.""" @staticmethod def State(request_iterator, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.stream_stream(request_iterator, target, '/org.apache.beam.model.fn_execution.v1.BeamFnState/State', org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StateRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StateResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) class BeamFnLoggingStub(object): """Stable """ def __init__(self, channel): """Constructor. Args: channel: A grpc.Channel. """ self.Logging = channel.stream_stream( '/org.apache.beam.model.fn_execution.v1.BeamFnLogging/Logging', request_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.LogEntry.List.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.LogControl.FromString, ) class BeamFnLoggingServicer(object): """Stable """ def Logging(self, request_iterator, context): """Allows for the SDK to emit log entries which the runner can associate with the active job. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def add_BeamFnLoggingServicer_to_server(servicer, server): rpc_method_handlers = { 'Logging': grpc.stream_stream_rpc_method_handler( servicer.Logging, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.LogEntry.List.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.LogControl.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( 'org.apache.beam.model.fn_execution.v1.BeamFnLogging', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) # This class is part of an EXPERIMENTAL API. class BeamFnLogging(object): """Stable """ @staticmethod def Logging(request_iterator, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.stream_stream(request_iterator, target, '/org.apache.beam.model.fn_execution.v1.BeamFnLogging/Logging', org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.LogEntry.List.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.LogControl.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) class BeamFnExternalWorkerPoolStub(object): """Missing associated documentation comment in .proto file.""" def __init__(self, channel): """Constructor. Args: channel: A grpc.Channel. """ self.StartWorker = channel.unary_unary( '/org.apache.beam.model.fn_execution.v1.BeamFnExternalWorkerPool/StartWorker', request_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StartWorkerRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StartWorkerResponse.FromString, ) self.StopWorker = channel.unary_unary( '/org.apache.beam.model.fn_execution.v1.BeamFnExternalWorkerPool/StopWorker', request_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StopWorkerRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StopWorkerResponse.FromString, ) class BeamFnExternalWorkerPoolServicer(object): """Missing associated documentation comment in .proto file.""" def StartWorker(self, request, context): """Start the SDK worker with the given ID. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def StopWorker(self, request, context): """Stop the SDK worker. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def add_BeamFnExternalWorkerPoolServicer_to_server(servicer, server): rpc_method_handlers = { 'StartWorker': grpc.unary_unary_rpc_method_handler( servicer.StartWorker, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StartWorkerRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StartWorkerResponse.SerializeToString, ), 'StopWorker': grpc.unary_unary_rpc_method_handler( servicer.StopWorker, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StopWorkerRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StopWorkerResponse.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( 'org.apache.beam.model.fn_execution.v1.BeamFnExternalWorkerPool', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) # This class is part of an EXPERIMENTAL API. class BeamFnExternalWorkerPool(object): """Missing associated documentation comment in .proto file.""" @staticmethod def StartWorker(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.fn_execution.v1.BeamFnExternalWorkerPool/StartWorker', org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StartWorkerRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StartWorkerResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def StopWorker(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.fn_execution.v1.BeamFnExternalWorkerPool/StopWorker', org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StopWorkerRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.StopWorkerResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) class BeamFnWorkerStatusStub(object): """API for SDKs to report debug-related statuses to runner during pipeline execution. """ def __init__(self, channel): """Constructor. Args: channel: A grpc.Channel. """ self.WorkerStatus = channel.stream_stream( '/org.apache.beam.model.fn_execution.v1.BeamFnWorkerStatus/WorkerStatus', request_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.WorkerStatusResponse.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.WorkerStatusRequest.FromString, ) class BeamFnWorkerStatusServicer(object): """API for SDKs to report debug-related statuses to runner during pipeline execution. """ def WorkerStatus(self, request_iterator, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def add_BeamFnWorkerStatusServicer_to_server(servicer, server): rpc_method_handlers = { 'WorkerStatus': grpc.stream_stream_rpc_method_handler( servicer.WorkerStatus, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.WorkerStatusResponse.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.WorkerStatusRequest.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( 'org.apache.beam.model.fn_execution.v1.BeamFnWorkerStatus', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) # This class is part of an EXPERIMENTAL API. class BeamFnWorkerStatus(object): """API for SDKs to report debug-related statuses to runner during pipeline execution. """ @staticmethod def WorkerStatus(request_iterator, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.stream_stream(request_iterator, target, '/org.apache.beam.model.fn_execution.v1.BeamFnWorkerStatus/WorkerStatus', org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.WorkerStatusResponse.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__fn__api__pb2.WorkerStatusRequest.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) ././@PaxHeader0000000000000000000000000000021400000000000010212 xustar00118 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_fn_api_pb2_urns.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_fn_api_pb20000644000175100017510000000104415206113151032727 0ustar00runnerrunnerfrom ........utils import PropertiesFromEnumValue from ...pipeline.v1 import metrics_pb2 EMPTY_MONITORING_INFO_LABEL_PROPS = metrics_pb2.MonitoringInfoLabelProps() EMPTY_MONITORING_INFO_SPEC = metrics_pb2.MonitoringInfoSpec() class FnApiTransforms(object): class Runner(object): DATA_SOURCE = PropertiesFromEnumValue('beam:runner:source:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) DATA_SINK = PropertiesFromEnumValue('beam:runner:sink:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ././@PaxHeader0000000000000000000000000000021600000000000010214 xustar00120 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_provision_api_pb2.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_provision_0000644000175100017510000000756615206113151033136 0ustar00runnerrunner# -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # source: org/apache/beam/model/fn_execution/v1/beam_provision_api.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() from ...pipeline.v1 import beam_runner_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_beam__runner__api__pb2 from ...pipeline.v1 import endpoints_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_endpoints__pb2 from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n>org/apache/beam/model/fn_execution/v1/beam_provision_api.proto\x12%org.apache.beam.model.fn_execution.v1\x1a\x37org/apache/beam/model/pipeline/v1/beam_runner_api.proto\x1a\x31org/apache/beam/model/pipeline/v1/endpoints.proto\x1a\x1cgoogle/protobuf/struct.proto\"\x19\n\x17GetProvisionInfoRequest\"^\n\x18GetProvisionInfoResponse\x12\x42\n\x04info\x18\x01 \x01(\x0b\x32\x34.org.apache.beam.model.fn_execution.v1.ProvisionInfo\"\xb5\x05\n\rProvisionInfo\x12\x31\n\x10pipeline_options\x18\x03 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x17\n\x0fretrieval_token\x18\x06 \x01(\t\x12P\n\x0fstatus_endpoint\x18\x07 \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.ApiServiceDescriptor\x12Q\n\x10logging_endpoint\x18\x08 \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.ApiServiceDescriptor\x12R\n\x11\x61rtifact_endpoint\x18\t \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.ApiServiceDescriptor\x12Q\n\x10\x63ontrol_endpoint\x18\n \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.ApiServiceDescriptor\x12L\n\x0c\x64\x65pendencies\x18\x0b \x03(\x0b\x32\x36.org.apache.beam.model.pipeline.v1.ArtifactInformation\x12\x1b\n\x13runner_capabilities\x18\x0c \x03(\t\x12T\n\x08metadata\x18\r \x03(\x0b\x32\x42.org.apache.beam.model.fn_execution.v1.ProvisionInfo.MetadataEntry\x12\x1a\n\x12sibling_worker_ids\x18\x0e \x03(\t\x1a/\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\x32\xa8\x01\n\x10ProvisionService\x12\x93\x01\n\x10GetProvisionInfo\x12>.org.apache.beam.model.fn_execution.v1.GetProvisionInfoRequest\x1a?.org.apache.beam.model.fn_execution.v1.GetProvisionInfoResponseB\x84\x01\n$org.apache.beam.model.fnexecution.v1B\x0cProvisionApiZNgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1;fnexecution_v1b\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'org.apache.beam.model.fn_execution.v1.beam_provision_api_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n$org.apache.beam.model.fnexecution.v1B\014ProvisionApiZNgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1;fnexecution_v1' _globals['_PROVISIONINFO_METADATAENTRY']._options = None _globals['_PROVISIONINFO_METADATAENTRY']._serialized_options = b'8\001' _globals['_GETPROVISIONINFOREQUEST']._serialized_start=243 _globals['_GETPROVISIONINFOREQUEST']._serialized_end=268 _globals['_GETPROVISIONINFORESPONSE']._serialized_start=270 _globals['_GETPROVISIONINFORESPONSE']._serialized_end=364 _globals['_PROVISIONINFO']._serialized_start=367 _globals['_PROVISIONINFO']._serialized_end=1060 _globals['_PROVISIONINFO_METADATAENTRY']._serialized_start=1013 _globals['_PROVISIONINFO_METADATAENTRY']._serialized_end=1060 _globals['_PROVISIONSERVICE']._serialized_start=1063 _globals['_PROVISIONSERVICE']._serialized_end=1231 # @@protoc_insertion_point(module_scope) ././@PaxHeader0000000000000000000000000000021700000000000010215 xustar00121 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_provision_api_pb2.pyi 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_provision_0000644000175100017510000001632715206113151033131 0ustar00runnerrunner""" @generated by mypy-protobuf. Do not edit manually! isort:skip_file Protocol Buffers describing the Provision API, for communicating with a runner for job and environment provisioning information over GRPC. """ import builtins import collections.abc import google.protobuf.descriptor import google.protobuf.internal.containers import google.protobuf.message import google.protobuf.struct_pb2 import org.apache.beam.model.pipeline.v1.beam_runner_api_pb2 import org.apache.beam.model.pipeline.v1.endpoints_pb2 import sys if sys.version_info >= (3, 8): import typing as typing_extensions else: import typing_extensions DESCRIPTOR: google.protobuf.descriptor.FileDescriptor @typing_extensions.final class GetProvisionInfoRequest(google.protobuf.message.Message): """A request to get the provision info of a SDK harness worker instance.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... global___GetProvisionInfoRequest = GetProvisionInfoRequest @typing_extensions.final class GetProvisionInfoResponse(google.protobuf.message.Message): """A response containing the provision info of a SDK harness worker instance.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor INFO_FIELD_NUMBER: builtins.int @property def info(self) -> global___ProvisionInfo: ... def __init__( self, *, info: global___ProvisionInfo | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["info", b"info"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["info", b"info"]) -> None: ... global___GetProvisionInfoResponse = GetProvisionInfoResponse @typing_extensions.final class ProvisionInfo(google.protobuf.message.Message): """Runtime provisioning information for a SDK harness worker instance, such as pipeline options, resource constraints and other job metadata """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class MetadataEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str value: builtins.str def __init__( self, *, key: builtins.str | None = ..., value: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... PIPELINE_OPTIONS_FIELD_NUMBER: builtins.int RETRIEVAL_TOKEN_FIELD_NUMBER: builtins.int STATUS_ENDPOINT_FIELD_NUMBER: builtins.int LOGGING_ENDPOINT_FIELD_NUMBER: builtins.int ARTIFACT_ENDPOINT_FIELD_NUMBER: builtins.int CONTROL_ENDPOINT_FIELD_NUMBER: builtins.int DEPENDENCIES_FIELD_NUMBER: builtins.int RUNNER_CAPABILITIES_FIELD_NUMBER: builtins.int METADATA_FIELD_NUMBER: builtins.int SIBLING_WORKER_IDS_FIELD_NUMBER: builtins.int @property def pipeline_options(self) -> google.protobuf.struct_pb2.Struct: """(required) Pipeline options. For non-template jobs, the options are identical to what is passed to job submission. """ retrieval_token: builtins.str """(required) The artifact retrieval token produced by LegacyArtifactStagingService.CommitManifestResponse. """ @property def status_endpoint(self) -> org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor: """(optional) The endpoint that the runner is hosting for the SDK to submit status reports to during pipeline execution. This field will only be populated if the runner supports SDK status reports. For more details see https://s.apache.org/beam-fn-api-harness-status """ @property def logging_endpoint(self) -> org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor: """(optional) The logging endpoint this SDK should use.""" @property def artifact_endpoint(self) -> org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor: """(optional) The artifact retrieval endpoint this SDK should use.""" @property def control_endpoint(self) -> org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor: """(optional) The control endpoint this SDK should use.""" @property def dependencies(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.ArtifactInformation]: """The set of dependencies that should be staged into this environment.""" @property def runner_capabilities(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """(optional) A set of capabilities that this SDK is allowed to use in its interactions with this runner. """ @property def metadata(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: """(optional) Runtime environment metadata that are static throughout the pipeline execution. """ @property def sibling_worker_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """(optional) If this environment supports SIBLING_WORKERS, used to indicate the ids of sibling workers, if any, that should be started in addition to this worker (which already has its own worker id). """ def __init__( self, *, pipeline_options: google.protobuf.struct_pb2.Struct | None = ..., retrieval_token: builtins.str | None = ..., status_endpoint: org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor | None = ..., logging_endpoint: org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor | None = ..., artifact_endpoint: org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor | None = ..., control_endpoint: org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor | None = ..., dependencies: collections.abc.Iterable[org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.ArtifactInformation] | None = ..., runner_capabilities: collections.abc.Iterable[builtins.str] | None = ..., metadata: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., sibling_worker_ids: collections.abc.Iterable[builtins.str] | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["artifact_endpoint", b"artifact_endpoint", "control_endpoint", b"control_endpoint", "logging_endpoint", b"logging_endpoint", "pipeline_options", b"pipeline_options", "status_endpoint", b"status_endpoint"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["artifact_endpoint", b"artifact_endpoint", "control_endpoint", b"control_endpoint", "dependencies", b"dependencies", "logging_endpoint", b"logging_endpoint", "metadata", b"metadata", "pipeline_options", b"pipeline_options", "retrieval_token", b"retrieval_token", "runner_capabilities", b"runner_capabilities", "sibling_worker_ids", b"sibling_worker_ids", "status_endpoint", b"status_endpoint"]) -> None: ... global___ProvisionInfo = ProvisionInfo ././@PaxHeader0000000000000000000000000000022300000000000010212 xustar00125 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_provision_api_pb2_grpc.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_provision_0000644000175100017510000000721715206113151033127 0ustar00runnerrunner# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc from . import beam_provision_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__provision__api__pb2 class ProvisionServiceStub(object): """A service to provide runtime provisioning information to the SDK harness worker instances -- such as pipeline options, resource constraints and other job metadata -- needed by an SDK harness instance to initialize. """ def __init__(self, channel): """Constructor. Args: channel: A grpc.Channel. """ self.GetProvisionInfo = channel.unary_unary( '/org.apache.beam.model.fn_execution.v1.ProvisionService/GetProvisionInfo', request_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__provision__api__pb2.GetProvisionInfoRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__provision__api__pb2.GetProvisionInfoResponse.FromString, ) class ProvisionServiceServicer(object): """A service to provide runtime provisioning information to the SDK harness worker instances -- such as pipeline options, resource constraints and other job metadata -- needed by an SDK harness instance to initialize. """ def GetProvisionInfo(self, request, context): """Get provision information for the SDK harness worker instance. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def add_ProvisionServiceServicer_to_server(servicer, server): rpc_method_handlers = { 'GetProvisionInfo': grpc.unary_unary_rpc_method_handler( servicer.GetProvisionInfo, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__provision__api__pb2.GetProvisionInfoRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__provision__api__pb2.GetProvisionInfoResponse.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( 'org.apache.beam.model.fn_execution.v1.ProvisionService', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) # This class is part of an EXPERIMENTAL API. class ProvisionService(object): """A service to provide runtime provisioning information to the SDK harness worker instances -- such as pipeline options, resource constraints and other job metadata -- needed by an SDK harness instance to initialize. """ @staticmethod def GetProvisionInfo(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.fn_execution.v1.ProvisionService/GetProvisionInfo', org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__provision__api__pb2.GetProvisionInfoRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_fn__execution_dot_v1_dot_beam__provision__api__pb2.GetProvisionInfoResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.533424 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/interactive/0000755000175100017510000000000015206113160027323 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/interactive/__init__.py0000644000175100017510000000142015206113151031431 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5354242 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/interactive/v1/0000755000175100017510000000000015206113160027651 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/interactive/v1/__init__.py0000644000175100017510000000154315206113151031765 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from . import beam_interactive_api_pb2 from . import beam_interactive_api_pb2_grpc ././@PaxHeader0000000000000000000000000000021700000000000010215 xustar00121 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/interactive/v1/beam_interactive_api_pb2.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/interactive/v1/beam_interactive0000644000175100017510000000465115206113151033103 0ustar00runnerrunner# -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # source: org/apache/beam/model/interactive/v1/beam_interactive_api.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() from ...pipeline.v1 import beam_runner_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_beam__runner__api__pb2 from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n?org/apache/beam/model/interactive/v1/beam_interactive_api.proto\x12$org.apache.beam.model.interactive.v1\x1a\x37org/apache/beam/model/pipeline/v1/beam_runner_api.proto\x1a\x1fgoogle/protobuf/timestamp.proto\"#\n\x14TestStreamFileHeader\x12\x0b\n\x03tag\x18\x01 \x01(\t\"j\n\x14TestStreamFileRecord\x12R\n\x0erecorded_event\x18\x01 \x01(\x0b\x32:.org.apache.beam.model.pipeline.v1.TestStreamPayload.Event\"#\n\rEventsRequest\x12\x12\n\noutput_ids\x18\x01 \x03(\t2\x92\x01\n\x11TestStreamService\x12}\n\x06\x45vents\x12\x33.org.apache.beam.model.interactive.v1.EventsRequest\x1a:.org.apache.beam.model.pipeline.v1.TestStreamPayload.Event\"\x00\x30\x01\x42\x46\n$org.apache.beam.model.interactive.v1B\x0eInteractiveApiZ\x0einteractive_v1b\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'org.apache.beam.model.interactive.v1.beam_interactive_api_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n$org.apache.beam.model.interactive.v1B\016InteractiveApiZ\016interactive_v1' _globals['_TESTSTREAMFILEHEADER']._serialized_start=195 _globals['_TESTSTREAMFILEHEADER']._serialized_end=230 _globals['_TESTSTREAMFILERECORD']._serialized_start=232 _globals['_TESTSTREAMFILERECORD']._serialized_end=338 _globals['_EVENTSREQUEST']._serialized_start=340 _globals['_EVENTSREQUEST']._serialized_end=375 _globals['_TESTSTREAMSERVICE']._serialized_start=378 _globals['_TESTSTREAMSERVICE']._serialized_end=524 # @@protoc_insertion_point(module_scope) ././@PaxHeader0000000000000000000000000000022000000000000010207 xustar00122 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/interactive/v1/beam_interactive_api_pb2.pyi 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/interactive/v1/beam_interactive0000644000175100017510000000620315206113151033076 0ustar00runnerrunner""" @generated by mypy-protobuf. Do not edit manually! isort:skip_file Protocol Buffers describing any Interactive Beam specific protos. TestStreamFileHeader: A record written at the top of a cached of TestStream events. This is used as metadata for the entire stream. """ import builtins import collections.abc import google.protobuf.descriptor import google.protobuf.internal.containers import google.protobuf.message import org.apache.beam.model.pipeline.v1.beam_runner_api_pb2 import sys if sys.version_info >= (3, 8): import typing as typing_extensions else: import typing_extensions DESCRIPTOR: google.protobuf.descriptor.FileDescriptor @typing_extensions.final class TestStreamFileHeader(google.protobuf.message.Message): """The first record. This contains metadata about the stream and how to properly process it. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TAG_FIELD_NUMBER: builtins.int tag: builtins.str """The PCollection tag this stream is associated with.""" def __init__( self, *, tag: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["tag", b"tag"]) -> None: ... global___TestStreamFileHeader = TestStreamFileHeader @typing_extensions.final class TestStreamFileRecord(google.protobuf.message.Message): """A record is a recorded element that a source produced. Its function is to give enough information to create a faithful recreation of the original stream of data. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor RECORDED_EVENT_FIELD_NUMBER: builtins.int @property def recorded_event(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.TestStreamPayload.Event: """The recorded event from an element stream.""" def __init__( self, *, recorded_event: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.TestStreamPayload.Event | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["recorded_event", b"recorded_event"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["recorded_event", b"recorded_event"]) -> None: ... global___TestStreamFileRecord = TestStreamFileRecord @typing_extensions.final class EventsRequest(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor OUTPUT_IDS_FIELD_NUMBER: builtins.int @property def output_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """The set of PCollections to read from. These are the PTransform outputs local names. These are a subset of the TestStream's outputs. This allows Interactive Beam to cache many PCollections from a pipeline then replay a subset of them. """ def __init__( self, *, output_ids: collections.abc.Iterable[builtins.str] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["output_ids", b"output_ids"]) -> None: ... global___EventsRequest = EventsRequest ././@PaxHeader0000000000000000000000000000022400000000000010213 xustar00126 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/interactive/v1/beam_interactive_api_pb2_grpc.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/interactive/v1/beam_interactive0000644000175100017510000000622115206113151033076 0ustar00runnerrunner# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc from . import beam_interactive_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_interactive_dot_v1_dot_beam__interactive__api__pb2 from ...pipeline.v1 import beam_runner_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_beam__runner__api__pb2 class TestStreamServiceStub(object): """Missing associated documentation comment in .proto file.""" def __init__(self, channel): """Constructor. Args: channel: A grpc.Channel. """ self.Events = channel.unary_stream( '/org.apache.beam.model.interactive.v1.TestStreamService/Events', request_serializer=org_dot_apache_dot_beam_dot_model_dot_interactive_dot_v1_dot_beam__interactive__api__pb2.EventsRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_beam__runner__api__pb2.TestStreamPayload.Event.FromString, ) class TestStreamServiceServicer(object): """Missing associated documentation comment in .proto file.""" def Events(self, request, context): """A TestStream will request for events using this RPC. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def add_TestStreamServiceServicer_to_server(servicer, server): rpc_method_handlers = { 'Events': grpc.unary_stream_rpc_method_handler( servicer.Events, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_interactive_dot_v1_dot_beam__interactive__api__pb2.EventsRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_beam__runner__api__pb2.TestStreamPayload.Event.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( 'org.apache.beam.model.interactive.v1.TestStreamService', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) # This class is part of an EXPERIMENTAL API. class TestStreamService(object): """Missing associated documentation comment in .proto file.""" @staticmethod def Events(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_stream(request, target, '/org.apache.beam.model.interactive.v1.TestStreamService/Events', org_dot_apache_dot_beam_dot_model_dot_interactive_dot_v1_dot_beam__interactive__api__pb2.EventsRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_beam__runner__api__pb2.TestStreamPayload.Event.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5354242 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/0000755000175100017510000000000015206113160027754 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/__init__.py0000644000175100017510000000142015206113151032062 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5404246 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/0000755000175100017510000000000015206113160030302 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/__init__.py0000644000175100017510000000203015206113151032406 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from . import beam_artifact_api_pb2_grpc from . import beam_job_api_pb2_grpc from . import beam_expansion_api_pb2_grpc from . import beam_job_api_pb2 from . import beam_expansion_api_pb2 from . import beam_artifact_api_pb2_urns from . import beam_artifact_api_pb2 ././@PaxHeader0000000000000000000000000000021700000000000010215 xustar00121 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_artifact_api_pb2.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_artifact0000644000175100017510000002160015206113151033005 0ustar00runnerrunner# -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # source: org/apache/beam/model/job_management/v1/beam_artifact_api.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() from ...pipeline.v1 import beam_runner_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_beam__runner__api__pb2 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n?org/apache/beam/model/job_management/v1/beam_artifact_api.proto\x12\'org.apache.beam.model.job_management.v1\x1a\x37org/apache/beam/model/pipeline/v1/beam_runner_api.proto\"|\n\x17ResolveArtifactsRequest\x12I\n\tartifacts\x18\x01 \x03(\x0b\x32\x36.org.apache.beam.model.pipeline.v1.ArtifactInformation\x12\x16\n\x0epreferred_urns\x18\x02 \x03(\t\"h\n\x18ResolveArtifactsResponse\x12L\n\x0creplacements\x18\x01 \x03(\x0b\x32\x36.org.apache.beam.model.pipeline.v1.ArtifactInformation\"^\n\x12GetArtifactRequest\x12H\n\x08\x61rtifact\x18\x01 \x01(\x0b\x32\x36.org.apache.beam.model.pipeline.v1.ArtifactInformation\"#\n\x13GetArtifactResponse\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\"\xd8\x01\n\x16\x41rtifactRequestWrapper\x12]\n\x10resolve_artifact\x18\xe8\x07 \x01(\x0b\x32@.org.apache.beam.model.job_management.v1.ResolveArtifactsRequestH\x00\x12T\n\x0cget_artifact\x18\xe9\x07 \x01(\x0b\x32;.org.apache.beam.model.job_management.v1.GetArtifactRequestH\x00\x42\t\n\x07request\"\x96\x02\n\x17\x41rtifactResponseWrapper\x12\x15\n\rstaging_token\x18\x01 \x01(\t\x12\x0f\n\x07is_last\x18\x02 \x01(\x08\x12g\n\x19resolve_artifact_response\x18\xe8\x07 \x01(\x0b\x32\x41.org.apache.beam.model.job_management.v1.ResolveArtifactsResponseH\x00\x12^\n\x15get_artifact_response\x18\xe9\x07 \x01(\x0b\x32<.org.apache.beam.model.job_management.v1.GetArtifactResponseH\x00\x42\n\n\x08response\"E\n\x10\x41rtifactMetadata\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x13\n\x0bpermissions\x18\x02 \x01(\r\x12\x0e\n\x06sha256\x18\x04 \x01(\t\"W\n\x08Manifest\x12K\n\x08\x61rtifact\x18\x01 \x03(\x0b\x32\x39.org.apache.beam.model.job_management.v1.ArtifactMetadata\"\xce\x01\n\rProxyManifest\x12\x43\n\x08manifest\x18\x01 \x01(\x0b\x32\x31.org.apache.beam.model.job_management.v1.Manifest\x12Q\n\x08location\x18\x02 \x03(\x0b\x32?.org.apache.beam.model.job_management.v1.ProxyManifest.Location\x1a%\n\x08Location\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0b\n\x03uri\x18\x02 \x01(\t\"-\n\x12GetManifestRequest\x12\x17\n\x0fretrieval_token\x18\x01 \x01(\t\"Z\n\x13GetManifestResponse\x12\x43\n\x08manifest\x18\x01 \x01(\x0b\x32\x31.org.apache.beam.model.job_management.v1.Manifest\"A\n\x18LegacyGetArtifactRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x17\n\x0fretrieval_token\x18\x02 \x01(\t\"\x1d\n\rArtifactChunk\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\"\x81\x01\n\x13PutArtifactMetadata\x12\x1d\n\x15staging_session_token\x18\x01 \x01(\t\x12K\n\x08metadata\x18\x02 \x01(\x0b\x32\x39.org.apache.beam.model.job_management.v1.ArtifactMetadata\"\xb9\x01\n\x12PutArtifactRequest\x12P\n\x08metadata\x18\x01 \x01(\x0b\x32<.org.apache.beam.model.job_management.v1.PutArtifactMetadataH\x00\x12\x46\n\x04\x64\x61ta\x18\x02 \x01(\x0b\x32\x36.org.apache.beam.model.job_management.v1.ArtifactChunkH\x00\x42\t\n\x07\x63ontent\"\x15\n\x13PutArtifactResponse\"{\n\x15\x43ommitManifestRequest\x12\x43\n\x08manifest\x18\x01 \x01(\x0b\x32\x31.org.apache.beam.model.job_management.v1.Manifest\x12\x1d\n\x15staging_session_token\x18\x02 \x01(\t\"|\n\x16\x43ommitManifestResponse\x12\x17\n\x0fretrieval_token\x18\x01 \x01(\t\"I\n\tConstants\x12<\n\x19NO_ARTIFACTS_STAGED_TOKEN\x10\x00\x1a\x1d\xaa\xb4\xfa\xc2\x05\x17__no_artifacts_staged__2\xc1\x02\n\x18\x41rtifactRetrievalService\x12\x97\x01\n\x10ResolveArtifacts\x12@.org.apache.beam.model.job_management.v1.ResolveArtifactsRequest\x1a\x41.org.apache.beam.model.job_management.v1.ResolveArtifactsResponse\x12\x8a\x01\n\x0bGetArtifact\x12;.org.apache.beam.model.job_management.v1.GetArtifactRequest\x1a<.org.apache.beam.model.job_management.v1.GetArtifactResponse0\x01\x32\xc3\x01\n\x16\x41rtifactStagingService\x12\xa8\x01\n\x1fReverseArtifactRetrievalService\x12@.org.apache.beam.model.job_management.v1.ArtifactResponseWrapper\x1a?.org.apache.beam.model.job_management.v1.ArtifactRequestWrapper(\x01\x30\x01\x32\xbf\x02\n\x1cLegacyArtifactStagingService\x12\x8a\x01\n\x0bPutArtifact\x12;.org.apache.beam.model.job_management.v1.PutArtifactRequest\x1a<.org.apache.beam.model.job_management.v1.PutArtifactResponse(\x01\x12\x91\x01\n\x0e\x43ommitManifest\x12>.org.apache.beam.model.job_management.v1.CommitManifestRequest\x1a?.org.apache.beam.model.job_management.v1.CommitManifestResponse2\xb8\x02\n\x1eLegacyArtifactRetrievalService\x12\x88\x01\n\x0bGetManifest\x12;.org.apache.beam.model.job_management.v1.GetManifestRequest\x1a<.org.apache.beam.model.job_management.v1.GetManifestResponse\x12\x8a\x01\n\x0bGetArtifact\x12\x41.org.apache.beam.model.job_management.v1.LegacyGetArtifactRequest\x1a\x36.org.apache.beam.model.job_management.v1.ArtifactChunk0\x01\x42\x89\x01\n&org.apache.beam.model.jobmanagement.v1B\x0b\x41rtifactApiZRgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/jobmanagement_v1;jobmanagement_v1b\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'org.apache.beam.model.job_management.v1.beam_artifact_api_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n&org.apache.beam.model.jobmanagement.v1B\013ArtifactApiZRgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/jobmanagement_v1;jobmanagement_v1' _globals['_COMMITMANIFESTRESPONSE_CONSTANTS'].values_by_name["NO_ARTIFACTS_STAGED_TOKEN"]._options = None _globals['_COMMITMANIFESTRESPONSE_CONSTANTS'].values_by_name["NO_ARTIFACTS_STAGED_TOKEN"]._serialized_options = b'\252\264\372\302\005\027__no_artifacts_staged__' _globals['_RESOLVEARTIFACTSREQUEST']._serialized_start=165 _globals['_RESOLVEARTIFACTSREQUEST']._serialized_end=289 _globals['_RESOLVEARTIFACTSRESPONSE']._serialized_start=291 _globals['_RESOLVEARTIFACTSRESPONSE']._serialized_end=395 _globals['_GETARTIFACTREQUEST']._serialized_start=397 _globals['_GETARTIFACTREQUEST']._serialized_end=491 _globals['_GETARTIFACTRESPONSE']._serialized_start=493 _globals['_GETARTIFACTRESPONSE']._serialized_end=528 _globals['_ARTIFACTREQUESTWRAPPER']._serialized_start=531 _globals['_ARTIFACTREQUESTWRAPPER']._serialized_end=747 _globals['_ARTIFACTRESPONSEWRAPPER']._serialized_start=750 _globals['_ARTIFACTRESPONSEWRAPPER']._serialized_end=1028 _globals['_ARTIFACTMETADATA']._serialized_start=1030 _globals['_ARTIFACTMETADATA']._serialized_end=1099 _globals['_MANIFEST']._serialized_start=1101 _globals['_MANIFEST']._serialized_end=1188 _globals['_PROXYMANIFEST']._serialized_start=1191 _globals['_PROXYMANIFEST']._serialized_end=1397 _globals['_PROXYMANIFEST_LOCATION']._serialized_start=1360 _globals['_PROXYMANIFEST_LOCATION']._serialized_end=1397 _globals['_GETMANIFESTREQUEST']._serialized_start=1399 _globals['_GETMANIFESTREQUEST']._serialized_end=1444 _globals['_GETMANIFESTRESPONSE']._serialized_start=1446 _globals['_GETMANIFESTRESPONSE']._serialized_end=1536 _globals['_LEGACYGETARTIFACTREQUEST']._serialized_start=1538 _globals['_LEGACYGETARTIFACTREQUEST']._serialized_end=1603 _globals['_ARTIFACTCHUNK']._serialized_start=1605 _globals['_ARTIFACTCHUNK']._serialized_end=1634 _globals['_PUTARTIFACTMETADATA']._serialized_start=1637 _globals['_PUTARTIFACTMETADATA']._serialized_end=1766 _globals['_PUTARTIFACTREQUEST']._serialized_start=1769 _globals['_PUTARTIFACTREQUEST']._serialized_end=1954 _globals['_PUTARTIFACTRESPONSE']._serialized_start=1956 _globals['_PUTARTIFACTRESPONSE']._serialized_end=1977 _globals['_COMMITMANIFESTREQUEST']._serialized_start=1979 _globals['_COMMITMANIFESTREQUEST']._serialized_end=2102 _globals['_COMMITMANIFESTRESPONSE']._serialized_start=2104 _globals['_COMMITMANIFESTRESPONSE']._serialized_end=2228 _globals['_COMMITMANIFESTRESPONSE_CONSTANTS']._serialized_start=2155 _globals['_COMMITMANIFESTRESPONSE_CONSTANTS']._serialized_end=2228 _globals['_ARTIFACTRETRIEVALSERVICE']._serialized_start=2231 _globals['_ARTIFACTRETRIEVALSERVICE']._serialized_end=2552 _globals['_ARTIFACTSTAGINGSERVICE']._serialized_start=2555 _globals['_ARTIFACTSTAGINGSERVICE']._serialized_end=2750 _globals['_LEGACYARTIFACTSTAGINGSERVICE']._serialized_start=2753 _globals['_LEGACYARTIFACTSTAGINGSERVICE']._serialized_end=3072 _globals['_LEGACYARTIFACTRETRIEVALSERVICE']._serialized_start=3075 _globals['_LEGACYARTIFACTRETRIEVALSERVICE']._serialized_end=3387 # @@protoc_insertion_point(module_scope) ././@PaxHeader0000000000000000000000000000022000000000000010207 xustar00122 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_artifact_api_pb2.pyi 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_artifact0000644000175100017510000004454715206113151033024 0ustar00runnerrunner""" @generated by mypy-protobuf. Do not edit manually! isort:skip_file Protocol Buffers describing the Artifact API, for communicating with a runner for artifact staging and retrieval over GRPC. """ import builtins import collections.abc import google.protobuf.descriptor import google.protobuf.internal.containers import google.protobuf.internal.enum_type_wrapper import google.protobuf.message import org.apache.beam.model.pipeline.v1.beam_runner_api_pb2 import sys import typing if sys.version_info >= (3, 10): import typing as typing_extensions else: import typing_extensions DESCRIPTOR: google.protobuf.descriptor.FileDescriptor @typing_extensions.final class ResolveArtifactsRequest(google.protobuf.message.Message): """A request for artifact resolution.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor ARTIFACTS_FIELD_NUMBER: builtins.int PREFERRED_URNS_FIELD_NUMBER: builtins.int @property def artifacts(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.ArtifactInformation]: """An (ordered) set of artifacts to (jointly) resolve.""" @property def preferred_urns(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """A set of artifact type urns that are understood by the requester. An attempt should be made to resolve the artifacts in terms of these URNs, but other URNs may be used as well with the understanding that they must be fetch-able as bytes via GetArtifact. """ def __init__( self, *, artifacts: collections.abc.Iterable[org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.ArtifactInformation] | None = ..., preferred_urns: collections.abc.Iterable[builtins.str] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["artifacts", b"artifacts", "preferred_urns", b"preferred_urns"]) -> None: ... global___ResolveArtifactsRequest = ResolveArtifactsRequest @typing_extensions.final class ResolveArtifactsResponse(google.protobuf.message.Message): """A response for artifact resolution.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor REPLACEMENTS_FIELD_NUMBER: builtins.int @property def replacements(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.ArtifactInformation]: """A full (ordered) set of replacements for the set of requested artifacts, preferably in terms of the requested type URNs. If there is no better resolution, the original list is returned. """ def __init__( self, *, replacements: collections.abc.Iterable[org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.ArtifactInformation] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["replacements", b"replacements"]) -> None: ... global___ResolveArtifactsResponse = ResolveArtifactsResponse @typing_extensions.final class GetArtifactRequest(google.protobuf.message.Message): """A request to get an artifact.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor ARTIFACT_FIELD_NUMBER: builtins.int @property def artifact(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.ArtifactInformation: ... def __init__( self, *, artifact: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.ArtifactInformation | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["artifact", b"artifact"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["artifact", b"artifact"]) -> None: ... global___GetArtifactRequest = GetArtifactRequest @typing_extensions.final class GetArtifactResponse(google.protobuf.message.Message): """Part of a response to getting an artifact.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor DATA_FIELD_NUMBER: builtins.int data: builtins.bytes def __init__( self, *, data: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["data", b"data"]) -> None: ... global___GetArtifactResponse = GetArtifactResponse @typing_extensions.final class ArtifactRequestWrapper(google.protobuf.message.Message): """Wraps an ArtifactRetrievalService request for use in ReverseArtifactRetrievalService.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor RESOLVE_ARTIFACT_FIELD_NUMBER: builtins.int GET_ARTIFACT_FIELD_NUMBER: builtins.int @property def resolve_artifact(self) -> global___ResolveArtifactsRequest: ... @property def get_artifact(self) -> global___GetArtifactRequest: ... def __init__( self, *, resolve_artifact: global___ResolveArtifactsRequest | None = ..., get_artifact: global___GetArtifactRequest | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["get_artifact", b"get_artifact", "request", b"request", "resolve_artifact", b"resolve_artifact"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["get_artifact", b"get_artifact", "request", b"request", "resolve_artifact", b"resolve_artifact"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["request", b"request"]) -> typing_extensions.Literal["resolve_artifact", "get_artifact"] | None: ... global___ArtifactRequestWrapper = ArtifactRequestWrapper @typing_extensions.final class ArtifactResponseWrapper(google.protobuf.message.Message): """Wraps an ArtifactRetrievalService response for use in ReverseArtifactRetrievalService.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor STAGING_TOKEN_FIELD_NUMBER: builtins.int IS_LAST_FIELD_NUMBER: builtins.int RESOLVE_ARTIFACT_RESPONSE_FIELD_NUMBER: builtins.int GET_ARTIFACT_RESPONSE_FIELD_NUMBER: builtins.int staging_token: builtins.str """A token indicating which job these artifacts are being staged for.""" is_last: builtins.bool """Whether this is the last response for this request (for those responses that would typically be terminated by the end of the response stream.) """ @property def resolve_artifact_response(self) -> global___ResolveArtifactsResponse: ... @property def get_artifact_response(self) -> global___GetArtifactResponse: ... def __init__( self, *, staging_token: builtins.str | None = ..., is_last: builtins.bool | None = ..., resolve_artifact_response: global___ResolveArtifactsResponse | None = ..., get_artifact_response: global___GetArtifactResponse | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["get_artifact_response", b"get_artifact_response", "resolve_artifact_response", b"resolve_artifact_response", "response", b"response"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["get_artifact_response", b"get_artifact_response", "is_last", b"is_last", "resolve_artifact_response", b"resolve_artifact_response", "response", b"response", "staging_token", b"staging_token"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["response", b"response"]) -> typing_extensions.Literal["resolve_artifact_response", "get_artifact_response"] | None: ... global___ArtifactResponseWrapper = ArtifactResponseWrapper @typing_extensions.final class ArtifactMetadata(google.protobuf.message.Message): """An artifact identifier and associated metadata.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor NAME_FIELD_NUMBER: builtins.int PERMISSIONS_FIELD_NUMBER: builtins.int SHA256_FIELD_NUMBER: builtins.int name: builtins.str """(Required) The name of the artifact.""" permissions: builtins.int """(Optional) The Unix-like permissions of the artifact""" sha256: builtins.str """(Optional) The hex-encoded sha256 checksum of the artifact. Used, among other things, by harness boot code to validate the integrity of the artifact. """ def __init__( self, *, name: builtins.str | None = ..., permissions: builtins.int | None = ..., sha256: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["name", b"name", "permissions", b"permissions", "sha256", b"sha256"]) -> None: ... global___ArtifactMetadata = ArtifactMetadata @typing_extensions.final class Manifest(google.protobuf.message.Message): """A collection of artifacts.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor ARTIFACT_FIELD_NUMBER: builtins.int @property def artifact(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ArtifactMetadata]: ... def __init__( self, *, artifact: collections.abc.Iterable[global___ArtifactMetadata] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["artifact", b"artifact"]) -> None: ... global___Manifest = Manifest @typing_extensions.final class ProxyManifest(google.protobuf.message.Message): """A manifest with location information.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class Location(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor NAME_FIELD_NUMBER: builtins.int URI_FIELD_NUMBER: builtins.int name: builtins.str uri: builtins.str def __init__( self, *, name: builtins.str | None = ..., uri: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["name", b"name", "uri", b"uri"]) -> None: ... MANIFEST_FIELD_NUMBER: builtins.int LOCATION_FIELD_NUMBER: builtins.int @property def manifest(self) -> global___Manifest: ... @property def location(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ProxyManifest.Location]: ... def __init__( self, *, manifest: global___Manifest | None = ..., location: collections.abc.Iterable[global___ProxyManifest.Location] | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["manifest", b"manifest"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["location", b"location", "manifest", b"manifest"]) -> None: ... global___ProxyManifest = ProxyManifest @typing_extensions.final class GetManifestRequest(google.protobuf.message.Message): """A request to get the manifest of a Job.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor RETRIEVAL_TOKEN_FIELD_NUMBER: builtins.int retrieval_token: builtins.str """(Required) An opaque token representing the entirety of the staged artifacts. Returned in CommitManifestResponse. """ def __init__( self, *, retrieval_token: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["retrieval_token", b"retrieval_token"]) -> None: ... global___GetManifestRequest = GetManifestRequest @typing_extensions.final class GetManifestResponse(google.protobuf.message.Message): """A response containing a job manifest.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor MANIFEST_FIELD_NUMBER: builtins.int @property def manifest(self) -> global___Manifest: ... def __init__( self, *, manifest: global___Manifest | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["manifest", b"manifest"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["manifest", b"manifest"]) -> None: ... global___GetManifestResponse = GetManifestResponse @typing_extensions.final class LegacyGetArtifactRequest(google.protobuf.message.Message): """A request to get an artifact. The artifact must be present in the manifest for the job.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor NAME_FIELD_NUMBER: builtins.int RETRIEVAL_TOKEN_FIELD_NUMBER: builtins.int name: builtins.str """(Required) The name of the artifact to retrieve.""" retrieval_token: builtins.str """(Required) An opaque token representing the entirety of the staged artifacts. Returned in CommitManifestResponse. """ def __init__( self, *, name: builtins.str | None = ..., retrieval_token: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["name", b"name", "retrieval_token", b"retrieval_token"]) -> None: ... global___LegacyGetArtifactRequest = LegacyGetArtifactRequest @typing_extensions.final class ArtifactChunk(google.protobuf.message.Message): """Part of an artifact.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor DATA_FIELD_NUMBER: builtins.int data: builtins.bytes def __init__( self, *, data: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["data", b"data"]) -> None: ... global___ArtifactChunk = ArtifactChunk @typing_extensions.final class PutArtifactMetadata(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor STAGING_SESSION_TOKEN_FIELD_NUMBER: builtins.int METADATA_FIELD_NUMBER: builtins.int staging_session_token: builtins.str """(Required) A token for artifact staging session. This token can be obtained from PrepareJob request in JobService """ @property def metadata(self) -> global___ArtifactMetadata: """(Required) The Artifact metadata.""" def __init__( self, *, staging_session_token: builtins.str | None = ..., metadata: global___ArtifactMetadata | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["metadata", b"metadata"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["metadata", b"metadata", "staging_session_token", b"staging_session_token"]) -> None: ... global___PutArtifactMetadata = PutArtifactMetadata @typing_extensions.final class PutArtifactRequest(google.protobuf.message.Message): """A request to stage an artifact.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor METADATA_FIELD_NUMBER: builtins.int DATA_FIELD_NUMBER: builtins.int @property def metadata(self) -> global___PutArtifactMetadata: """The first message in a PutArtifact call must contain this field.""" @property def data(self) -> global___ArtifactChunk: """A chunk of the artifact. All messages after the first in a PutArtifact call must contain a chunk. """ def __init__( self, *, metadata: global___PutArtifactMetadata | None = ..., data: global___ArtifactChunk | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["content", b"content", "data", b"data", "metadata", b"metadata"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["content", b"content", "data", b"data", "metadata", b"metadata"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["content", b"content"]) -> typing_extensions.Literal["metadata", "data"] | None: ... global___PutArtifactRequest = PutArtifactRequest @typing_extensions.final class PutArtifactResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... global___PutArtifactResponse = PutArtifactResponse @typing_extensions.final class CommitManifestRequest(google.protobuf.message.Message): """A request to commit the manifest for a Job. All artifacts must have been successfully uploaded before this call is made. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor MANIFEST_FIELD_NUMBER: builtins.int STAGING_SESSION_TOKEN_FIELD_NUMBER: builtins.int @property def manifest(self) -> global___Manifest: """(Required) The manifest to commit.""" staging_session_token: builtins.str """(Required) A token for artifact staging session. This token can be obtained from PrepareJob request in JobService """ def __init__( self, *, manifest: global___Manifest | None = ..., staging_session_token: builtins.str | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["manifest", b"manifest"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["manifest", b"manifest", "staging_session_token", b"staging_session_token"]) -> None: ... global___CommitManifestRequest = CommitManifestRequest @typing_extensions.final class CommitManifestResponse(google.protobuf.message.Message): """The result of committing a manifest.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Constants: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _ConstantsEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[CommitManifestResponse._Constants.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor NO_ARTIFACTS_STAGED_TOKEN: CommitManifestResponse._Constants.ValueType # 0 """Token indicating that no artifacts were staged and therefore no retrieval attempt is necessary.""" class Constants(_Constants, metaclass=_ConstantsEnumTypeWrapper): ... NO_ARTIFACTS_STAGED_TOKEN: CommitManifestResponse.Constants.ValueType # 0 """Token indicating that no artifacts were staged and therefore no retrieval attempt is necessary.""" RETRIEVAL_TOKEN_FIELD_NUMBER: builtins.int retrieval_token: builtins.str """(Required) An opaque token representing the entirety of the staged artifacts. This can be used to retrieve the manifest and artifacts from an associated LegacyArtifactRetrievalService. """ def __init__( self, *, retrieval_token: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["retrieval_token", b"retrieval_token"]) -> None: ... global___CommitManifestResponse = CommitManifestResponse ././@PaxHeader0000000000000000000000000000022400000000000010213 xustar00126 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_artifact_api_pb2_grpc.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_artifact0000644000175100017510000004707115206113151033017 0ustar00runnerrunner# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc from . import beam_artifact_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2 class ArtifactRetrievalServiceStub(object): """A service to retrieve artifacts for use in a Job. """ def __init__(self, channel): """Constructor. Args: channel: A grpc.Channel. """ self.ResolveArtifacts = channel.unary_unary( '/org.apache.beam.model.job_management.v1.ArtifactRetrievalService/ResolveArtifacts', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ResolveArtifactsRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ResolveArtifactsResponse.FromString, ) self.GetArtifact = channel.unary_stream( '/org.apache.beam.model.job_management.v1.ArtifactRetrievalService/GetArtifact', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.GetArtifactRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.GetArtifactResponse.FromString, ) class ArtifactRetrievalServiceServicer(object): """A service to retrieve artifacts for use in a Job. """ def ResolveArtifacts(self, request, context): """Resolves the given artifact references into one or more replacement artifact references (e.g. a Maven dependency into a (transitive) set of jars. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def GetArtifact(self, request, context): """Retrieves the given artifact as a stream of bytes. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def add_ArtifactRetrievalServiceServicer_to_server(servicer, server): rpc_method_handlers = { 'ResolveArtifacts': grpc.unary_unary_rpc_method_handler( servicer.ResolveArtifacts, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ResolveArtifactsRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ResolveArtifactsResponse.SerializeToString, ), 'GetArtifact': grpc.unary_stream_rpc_method_handler( servicer.GetArtifact, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.GetArtifactRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.GetArtifactResponse.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( 'org.apache.beam.model.job_management.v1.ArtifactRetrievalService', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) # This class is part of an EXPERIMENTAL API. class ArtifactRetrievalService(object): """A service to retrieve artifacts for use in a Job. """ @staticmethod def ResolveArtifacts(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.job_management.v1.ArtifactRetrievalService/ResolveArtifacts', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ResolveArtifactsRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ResolveArtifactsResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def GetArtifact(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_stream(request, target, '/org.apache.beam.model.job_management.v1.ArtifactRetrievalService/GetArtifact', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.GetArtifactRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.GetArtifactResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) class ArtifactStagingServiceStub(object): """A service that allows the client to act as an ArtifactRetrievalService, for a particular job with the server initiating requests and receiving responses. A client calls the service with an ArtifactResponseWrapper that has the staging token set, and thereafter responds to the server's requests. """ def __init__(self, channel): """Constructor. Args: channel: A grpc.Channel. """ self.ReverseArtifactRetrievalService = channel.stream_stream( '/org.apache.beam.model.job_management.v1.ArtifactStagingService/ReverseArtifactRetrievalService', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ArtifactResponseWrapper.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ArtifactRequestWrapper.FromString, ) class ArtifactStagingServiceServicer(object): """A service that allows the client to act as an ArtifactRetrievalService, for a particular job with the server initiating requests and receiving responses. A client calls the service with an ArtifactResponseWrapper that has the staging token set, and thereafter responds to the server's requests. """ def ReverseArtifactRetrievalService(self, request_iterator, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def add_ArtifactStagingServiceServicer_to_server(servicer, server): rpc_method_handlers = { 'ReverseArtifactRetrievalService': grpc.stream_stream_rpc_method_handler( servicer.ReverseArtifactRetrievalService, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ArtifactResponseWrapper.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ArtifactRequestWrapper.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( 'org.apache.beam.model.job_management.v1.ArtifactStagingService', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) # This class is part of an EXPERIMENTAL API. class ArtifactStagingService(object): """A service that allows the client to act as an ArtifactRetrievalService, for a particular job with the server initiating requests and receiving responses. A client calls the service with an ArtifactResponseWrapper that has the staging token set, and thereafter responds to the server's requests. """ @staticmethod def ReverseArtifactRetrievalService(request_iterator, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.stream_stream(request_iterator, target, '/org.apache.beam.model.job_management.v1.ArtifactStagingService/ReverseArtifactRetrievalService', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ArtifactResponseWrapper.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ArtifactRequestWrapper.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) class LegacyArtifactStagingServiceStub(object): """Legacy artifact staging service for pipeline-level artifacts. A service to stage artifacts for use in a Job. """ def __init__(self, channel): """Constructor. Args: channel: A grpc.Channel. """ self.PutArtifact = channel.stream_unary( '/org.apache.beam.model.job_management.v1.LegacyArtifactStagingService/PutArtifact', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.PutArtifactRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.PutArtifactResponse.FromString, ) self.CommitManifest = channel.unary_unary( '/org.apache.beam.model.job_management.v1.LegacyArtifactStagingService/CommitManifest', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.CommitManifestRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.CommitManifestResponse.FromString, ) class LegacyArtifactStagingServiceServicer(object): """Legacy artifact staging service for pipeline-level artifacts. A service to stage artifacts for use in a Job. """ def PutArtifact(self, request_iterator, context): """Stage an artifact to be available during job execution. The first request must contain the name of the artifact. All future requests must contain sequential chunks of the content of the artifact. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def CommitManifest(self, request, context): """Commit the manifest for a Job. All artifacts must have been successfully uploaded before this call is made. Throws error INVALID_ARGUMENT if not all of the members of the manifest are present """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def add_LegacyArtifactStagingServiceServicer_to_server(servicer, server): rpc_method_handlers = { 'PutArtifact': grpc.stream_unary_rpc_method_handler( servicer.PutArtifact, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.PutArtifactRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.PutArtifactResponse.SerializeToString, ), 'CommitManifest': grpc.unary_unary_rpc_method_handler( servicer.CommitManifest, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.CommitManifestRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.CommitManifestResponse.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( 'org.apache.beam.model.job_management.v1.LegacyArtifactStagingService', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) # This class is part of an EXPERIMENTAL API. class LegacyArtifactStagingService(object): """Legacy artifact staging service for pipeline-level artifacts. A service to stage artifacts for use in a Job. """ @staticmethod def PutArtifact(request_iterator, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.stream_unary(request_iterator, target, '/org.apache.beam.model.job_management.v1.LegacyArtifactStagingService/PutArtifact', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.PutArtifactRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.PutArtifactResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def CommitManifest(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.job_management.v1.LegacyArtifactStagingService/CommitManifest', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.CommitManifestRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.CommitManifestResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) class LegacyArtifactRetrievalServiceStub(object): """A service to retrieve artifacts for use in a Job. """ def __init__(self, channel): """Constructor. Args: channel: A grpc.Channel. """ self.GetManifest = channel.unary_unary( '/org.apache.beam.model.job_management.v1.LegacyArtifactRetrievalService/GetManifest', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.GetManifestRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.GetManifestResponse.FromString, ) self.GetArtifact = channel.unary_stream( '/org.apache.beam.model.job_management.v1.LegacyArtifactRetrievalService/GetArtifact', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.LegacyGetArtifactRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ArtifactChunk.FromString, ) class LegacyArtifactRetrievalServiceServicer(object): """A service to retrieve artifacts for use in a Job. """ def GetManifest(self, request, context): """Get the manifest for the job """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def GetArtifact(self, request, context): """Get an artifact staged for the job. The requested artifact must be within the manifest """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def add_LegacyArtifactRetrievalServiceServicer_to_server(servicer, server): rpc_method_handlers = { 'GetManifest': grpc.unary_unary_rpc_method_handler( servicer.GetManifest, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.GetManifestRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.GetManifestResponse.SerializeToString, ), 'GetArtifact': grpc.unary_stream_rpc_method_handler( servicer.GetArtifact, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.LegacyGetArtifactRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ArtifactChunk.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( 'org.apache.beam.model.job_management.v1.LegacyArtifactRetrievalService', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) # This class is part of an EXPERIMENTAL API. class LegacyArtifactRetrievalService(object): """A service to retrieve artifacts for use in a Job. """ @staticmethod def GetManifest(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.job_management.v1.LegacyArtifactRetrievalService/GetManifest', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.GetManifestRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.GetManifestResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def GetArtifact(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_stream(request, target, '/org.apache.beam.model.job_management.v1.LegacyArtifactRetrievalService/GetArtifact', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.LegacyGetArtifactRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__artifact__api__pb2.ArtifactChunk.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) ././@PaxHeader0000000000000000000000000000022400000000000010213 xustar00126 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_artifact_api_pb2_urns.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_artifact0000644000175100017510000000067415206113151033015 0ustar00runnerrunnerfrom ........utils import PropertiesFromEnumValue from ...pipeline.v1 import metrics_pb2 EMPTY_MONITORING_INFO_LABEL_PROPS = metrics_pb2.MonitoringInfoLabelProps() EMPTY_MONITORING_INFO_SPEC = metrics_pb2.MonitoringInfoSpec() class CommitManifestResponse(object): class Constants(object): NO_ARTIFACTS_STAGED_TOKEN = PropertiesFromEnumValue('', '__no_artifacts_staged__', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ././@PaxHeader0000000000000000000000000000022000000000000010207 xustar00122 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_expansion_api_pb2.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_expansio0000644000175100017510000001230315206113151033036 0ustar00runnerrunner# -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # source: org/apache/beam/model/job_management/v1/beam_expansion_api.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2 from ...pipeline.v1 import beam_runner_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_beam__runner__api__pb2 from ...pipeline.v1 import schema_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_schema__pb2 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n@org/apache/beam/model/job_management/v1/beam_expansion_api.proto\x12\"org.apache.beam.model.expansion.v1\x1a\x1cgoogle/protobuf/struct.proto\x1a\x37org/apache/beam/model/pipeline/v1/beam_runner_api.proto\x1a.org/apache/beam/model/pipeline/v1/schema.proto\"\x9d\x03\n\x10\x45xpansionRequest\x12\x41\n\ncomponents\x18\x01 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.Components\x12@\n\ttransform\x18\x02 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.PTransform\x12\x11\n\tnamespace\x18\x03 \x01(\t\x12l\n\x15output_coder_requests\x18\x04 \x03(\x0b\x32M.org.apache.beam.model.expansion.v1.ExpansionRequest.OutputCoderRequestsEntry\x12\x14\n\x0crequirements\x18\x05 \x03(\t\x12\x31\n\x10pipeline_options\x18\x06 \x01(\x0b\x32\x17.google.protobuf.Struct\x1a:\n\x18OutputCoderRequestsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\xbd\x01\n\x11\x45xpansionResponse\x12\x41\n\ncomponents\x18\x01 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.Components\x12@\n\ttransform\x18\x02 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.PTransform\x12\x14\n\x0crequirements\x18\x03 \x03(\t\x12\r\n\x05\x65rror\x18\n \x01(\t\" \n\x1e\x44iscoverSchemaTransformRequest\"\xb1\x01\n\x15SchemaTransformConfig\x12@\n\rconfig_schema\x18\x01 \x01(\x0b\x32).org.apache.beam.model.pipeline.v1.Schema\x12\x1f\n\x17input_pcollection_names\x18\x02 \x03(\t\x12 \n\x18output_pcollection_names\x18\x03 \x03(\t\x12\x13\n\x0b\x64\x65scription\x18\x04 \x01(\t\"\xae\x02\n\x1f\x44iscoverSchemaTransformResponse\x12\x81\x01\n\x18schema_transform_configs\x18\x01 \x03(\x0b\x32_.org.apache.beam.model.expansion.v1.DiscoverSchemaTransformResponse.SchemaTransformConfigsEntry\x12\r\n\x05\x65rror\x18\x02 \x01(\t\x1ax\n\x1bSchemaTransformConfigsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12H\n\x05value\x18\x02 \x01(\x0b\x32\x39.org.apache.beam.model.expansion.v1.SchemaTransformConfig:\x02\x38\x01\x32\xae\x02\n\x10\x45xpansionService\x12u\n\x06\x45xpand\x12\x34.org.apache.beam.model.expansion.v1.ExpansionRequest\x1a\x35.org.apache.beam.model.expansion.v1.ExpansionResponse\x12\xa2\x01\n\x17\x44iscoverSchemaTransform\x12\x42.org.apache.beam.model.expansion.v1.DiscoverSchemaTransformRequest\x1a\x43.org.apache.beam.model.expansion.v1.DiscoverSchemaTransformResponseB\x86\x01\n\"org.apache.beam.model.expansion.v1B\x0c\x45xpansionApiZRgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/jobmanagement_v1;jobmanagement_v1b\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'org.apache.beam.model.job_management.v1.beam_expansion_api_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n\"org.apache.beam.model.expansion.v1B\014ExpansionApiZRgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/jobmanagement_v1;jobmanagement_v1' _globals['_EXPANSIONREQUEST_OUTPUTCODERREQUESTSENTRY']._options = None _globals['_EXPANSIONREQUEST_OUTPUTCODERREQUESTSENTRY']._serialized_options = b'8\001' _globals['_DISCOVERSCHEMATRANSFORMRESPONSE_SCHEMATRANSFORMCONFIGSENTRY']._options = None _globals['_DISCOVERSCHEMATRANSFORMRESPONSE_SCHEMATRANSFORMCONFIGSENTRY']._serialized_options = b'8\001' _globals['_EXPANSIONREQUEST']._serialized_start=240 _globals['_EXPANSIONREQUEST']._serialized_end=653 _globals['_EXPANSIONREQUEST_OUTPUTCODERREQUESTSENTRY']._serialized_start=595 _globals['_EXPANSIONREQUEST_OUTPUTCODERREQUESTSENTRY']._serialized_end=653 _globals['_EXPANSIONRESPONSE']._serialized_start=656 _globals['_EXPANSIONRESPONSE']._serialized_end=845 _globals['_DISCOVERSCHEMATRANSFORMREQUEST']._serialized_start=847 _globals['_DISCOVERSCHEMATRANSFORMREQUEST']._serialized_end=879 _globals['_SCHEMATRANSFORMCONFIG']._serialized_start=882 _globals['_SCHEMATRANSFORMCONFIG']._serialized_end=1059 _globals['_DISCOVERSCHEMATRANSFORMRESPONSE']._serialized_start=1062 _globals['_DISCOVERSCHEMATRANSFORMRESPONSE']._serialized_end=1364 _globals['_DISCOVERSCHEMATRANSFORMRESPONSE_SCHEMATRANSFORMCONFIGSENTRY']._serialized_start=1244 _globals['_DISCOVERSCHEMATRANSFORMRESPONSE_SCHEMATRANSFORMCONFIGSENTRY']._serialized_end=1364 _globals['_EXPANSIONSERVICE']._serialized_start=1367 _globals['_EXPANSIONSERVICE']._serialized_end=1669 # @@protoc_insertion_point(module_scope) ././@PaxHeader0000000000000000000000000000022100000000000010210 xustar00123 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_expansion_api_pb2.pyi 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_expansio0000644000175100017510000002374715206113151033054 0ustar00runnerrunner""" @generated by mypy-protobuf. Do not edit manually! isort:skip_file Protocol Buffers describing the Expansion API, an api for expanding transforms in a remote SDK. """ import builtins import collections.abc import google.protobuf.descriptor import google.protobuf.internal.containers import google.protobuf.message import google.protobuf.struct_pb2 import org.apache.beam.model.pipeline.v1.beam_runner_api_pb2 import org.apache.beam.model.pipeline.v1.schema_pb2 import sys if sys.version_info >= (3, 8): import typing as typing_extensions else: import typing_extensions DESCRIPTOR: google.protobuf.descriptor.FileDescriptor @typing_extensions.final class ExpansionRequest(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class OutputCoderRequestsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str value: builtins.str def __init__( self, *, key: builtins.str | None = ..., value: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... COMPONENTS_FIELD_NUMBER: builtins.int TRANSFORM_FIELD_NUMBER: builtins.int NAMESPACE_FIELD_NUMBER: builtins.int OUTPUT_CODER_REQUESTS_FIELD_NUMBER: builtins.int REQUIREMENTS_FIELD_NUMBER: builtins.int PIPELINE_OPTIONS_FIELD_NUMBER: builtins.int @property def components(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Components: """Set of components needed to interpret the transform, or which may be useful for its expansion. This includes the input PCollections (if any) to the to-be-expanded transform, along with their coders and windowing strategies. """ @property def transform(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.PTransform: """The actual PTransform to be expaneded according to its spec. Its input should be set, but its subtransforms and outputs should not be. """ namespace: builtins.str """A namespace (prefix) to use for the id of any newly created components. """ @property def output_coder_requests(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: """(Optional) Map from a local output tag to a coder id. If it is set, asks the expansion service to use the given coders for the output PCollections. Note that the request may not be fulfilled. """ @property def requirements(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """A set of requirements that must be used by the expansion service to interpret the components provided with this request. """ @property def pipeline_options(self) -> google.protobuf.struct_pb2.Struct: """(Optional) A set of Pipeline Options that should be used when expanding this transform. """ def __init__( self, *, components: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Components | None = ..., transform: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.PTransform | None = ..., namespace: builtins.str | None = ..., output_coder_requests: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., requirements: collections.abc.Iterable[builtins.str] | None = ..., pipeline_options: google.protobuf.struct_pb2.Struct | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["components", b"components", "pipeline_options", b"pipeline_options", "transform", b"transform"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["components", b"components", "namespace", b"namespace", "output_coder_requests", b"output_coder_requests", "pipeline_options", b"pipeline_options", "requirements", b"requirements", "transform", b"transform"]) -> None: ... global___ExpansionRequest = ExpansionRequest @typing_extensions.final class ExpansionResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor COMPONENTS_FIELD_NUMBER: builtins.int TRANSFORM_FIELD_NUMBER: builtins.int REQUIREMENTS_FIELD_NUMBER: builtins.int ERROR_FIELD_NUMBER: builtins.int @property def components(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Components: """Set of components needed to execute the expanded transform, including the (original) inputs, outputs, and subtransforms. """ @property def transform(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.PTransform: """The expanded transform itself, with references to its outputs and subtransforms. """ @property def requirements(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """A set of requirements that must be appended to this pipeline's requirements. """ error: builtins.str """(Optional) An string representation of any error encountered while attempting to expand this transform. """ def __init__( self, *, components: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Components | None = ..., transform: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.PTransform | None = ..., requirements: collections.abc.Iterable[builtins.str] | None = ..., error: builtins.str | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["components", b"components", "transform", b"transform"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["components", b"components", "error", b"error", "requirements", b"requirements", "transform", b"transform"]) -> None: ... global___ExpansionResponse = ExpansionResponse @typing_extensions.final class DiscoverSchemaTransformRequest(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... global___DiscoverSchemaTransformRequest = DiscoverSchemaTransformRequest @typing_extensions.final class SchemaTransformConfig(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor CONFIG_SCHEMA_FIELD_NUMBER: builtins.int INPUT_PCOLLECTION_NAMES_FIELD_NUMBER: builtins.int OUTPUT_PCOLLECTION_NAMES_FIELD_NUMBER: builtins.int DESCRIPTION_FIELD_NUMBER: builtins.int @property def config_schema(self) -> org.apache.beam.model.pipeline.v1.schema_pb2.Schema: """Config schema of the SchemaTransform""" @property def input_pcollection_names(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """Names of input PCollections""" @property def output_pcollection_names(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """Names of output PCollections""" description: builtins.str """Description of this transform and usage used for documentation. May be markdown formatted. Note that configuration parameters may also have documentation attached as part of the config_schema. """ def __init__( self, *, config_schema: org.apache.beam.model.pipeline.v1.schema_pb2.Schema | None = ..., input_pcollection_names: collections.abc.Iterable[builtins.str] | None = ..., output_pcollection_names: collections.abc.Iterable[builtins.str] | None = ..., description: builtins.str | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["config_schema", b"config_schema"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["config_schema", b"config_schema", "description", b"description", "input_pcollection_names", b"input_pcollection_names", "output_pcollection_names", b"output_pcollection_names"]) -> None: ... global___SchemaTransformConfig = SchemaTransformConfig @typing_extensions.final class DiscoverSchemaTransformResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class SchemaTransformConfigsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> global___SchemaTransformConfig: ... def __init__( self, *, key: builtins.str | None = ..., value: global___SchemaTransformConfig | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... SCHEMA_TRANSFORM_CONFIGS_FIELD_NUMBER: builtins.int ERROR_FIELD_NUMBER: builtins.int @property def schema_transform_configs(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___SchemaTransformConfig]: """A mapping from SchemaTransform ID to schema transform config of discovered SchemaTransforms """ error: builtins.str """If list of identifies are empty, this may contain an error.""" def __init__( self, *, schema_transform_configs: collections.abc.Mapping[builtins.str, global___SchemaTransformConfig] | None = ..., error: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["error", b"error", "schema_transform_configs", b"schema_transform_configs"]) -> None: ... global___DiscoverSchemaTransformResponse = DiscoverSchemaTransformResponse ././@PaxHeader0000000000000000000000000000022500000000000010214 xustar00127 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_expansion_api_pb2_grpc.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_expansio0000644000175100017510000001235715206113151033047 0ustar00runnerrunner# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc from . import beam_expansion_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__expansion__api__pb2 class ExpansionServiceStub(object): """Job Service for constructing pipelines """ def __init__(self, channel): """Constructor. Args: channel: A grpc.Channel. """ self.Expand = channel.unary_unary( '/org.apache.beam.model.expansion.v1.ExpansionService/Expand', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__expansion__api__pb2.ExpansionRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__expansion__api__pb2.ExpansionResponse.FromString, ) self.DiscoverSchemaTransform = channel.unary_unary( '/org.apache.beam.model.expansion.v1.ExpansionService/DiscoverSchemaTransform', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__expansion__api__pb2.DiscoverSchemaTransformRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__expansion__api__pb2.DiscoverSchemaTransformResponse.FromString, ) class ExpansionServiceServicer(object): """Job Service for constructing pipelines """ def Expand(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def DiscoverSchemaTransform(self, request, context): """A RPC to discover already registered SchemaTransformProviders. See https://s.apache.org/easy-multi-language for more details. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def add_ExpansionServiceServicer_to_server(servicer, server): rpc_method_handlers = { 'Expand': grpc.unary_unary_rpc_method_handler( servicer.Expand, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__expansion__api__pb2.ExpansionRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__expansion__api__pb2.ExpansionResponse.SerializeToString, ), 'DiscoverSchemaTransform': grpc.unary_unary_rpc_method_handler( servicer.DiscoverSchemaTransform, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__expansion__api__pb2.DiscoverSchemaTransformRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__expansion__api__pb2.DiscoverSchemaTransformResponse.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( 'org.apache.beam.model.expansion.v1.ExpansionService', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) # This class is part of an EXPERIMENTAL API. class ExpansionService(object): """Job Service for constructing pipelines """ @staticmethod def Expand(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.expansion.v1.ExpansionService/Expand', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__expansion__api__pb2.ExpansionRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__expansion__api__pb2.ExpansionResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def DiscoverSchemaTransform(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.expansion.v1.ExpansionService/DiscoverSchemaTransform', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__expansion__api__pb2.DiscoverSchemaTransformRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__expansion__api__pb2.DiscoverSchemaTransformResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) ././@PaxHeader0000000000000000000000000000021200000000000010210 xustar00116 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_job_api_pb2.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_job_api_0000644000175100017510000002641215206113151032760 0ustar00runnerrunner# -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # source: org/apache/beam/model/job_management/v1/beam_job_api.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() from ...pipeline.v1 import beam_runner_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_beam__runner__api__pb2 from ...pipeline.v1 import endpoints_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_endpoints__pb2 from ...pipeline.v1 import metrics_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_metrics__pb2 from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2 from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n:org/apache/beam/model/job_management/v1/beam_job_api.proto\x12\'org.apache.beam.model.job_management.v1\x1a\x37org/apache/beam/model/pipeline/v1/beam_runner_api.proto\x1a\x31org/apache/beam/model/pipeline/v1/endpoints.proto\x1a/org/apache/beam/model/pipeline/v1/metrics.proto\x1a\x1cgoogle/protobuf/struct.proto\x1a\x1fgoogle/protobuf/timestamp.proto\"\x97\x01\n\x11PrepareJobRequest\x12=\n\x08pipeline\x18\x01 \x01(\x0b\x32+.org.apache.beam.model.pipeline.v1.Pipeline\x12\x31\n\x10pipeline_options\x18\x02 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x10\n\x08job_name\x18\x03 \x01(\t\"\xa7\x01\n\x12PrepareJobResponse\x12\x16\n\x0epreparation_id\x18\x01 \x01(\t\x12Z\n\x19\x61rtifact_staging_endpoint\x18\x02 \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.ApiServiceDescriptor\x12\x1d\n\x15staging_session_token\x18\x03 \x01(\t\"@\n\rRunJobRequest\x12\x16\n\x0epreparation_id\x18\x01 \x01(\t\x12\x17\n\x0fretrieval_token\x18\x02 \x01(\t\" \n\x0eRunJobResponse\x12\x0e\n\x06job_id\x18\x01 \x01(\t\"\"\n\x10\x43\x61ncelJobRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t\"Z\n\x11\x43\x61ncelJobResponse\x12\x45\n\x05state\x18\x01 \x01(\x0e\x32\x36.org.apache.beam.model.job_management.v1.JobState.Enum\"!\n\x0f\x44rainJobRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t\"Y\n\x10\x44rainJobResponse\x12\x45\n\x05state\x18\x01 \x01(\x0e\x32\x36.org.apache.beam.model.job_management.v1.JobState.Enum\"\xa5\x01\n\x07JobInfo\x12\x0e\n\x06job_id\x18\x01 \x01(\t\x12\x10\n\x08job_name\x18\x02 \x01(\t\x12\x31\n\x10pipeline_options\x18\x03 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x45\n\x05state\x18\x04 \x01(\x0e\x32\x36.org.apache.beam.model.job_management.v1.JobState.Enum\"\x10\n\x0eGetJobsRequest\"U\n\x0fGetJobsResponse\x12\x42\n\x08job_info\x18\x01 \x03(\x0b\x32\x30.org.apache.beam.model.job_management.v1.JobInfo\"$\n\x12GetJobStateRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t\"\x85\x01\n\rJobStateEvent\x12\x45\n\x05state\x18\x01 \x01(\x0e\x32\x36.org.apache.beam.model.job_management.v1.JobState.Enum\x12-\n\ttimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\"\'\n\x15GetJobPipelineRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t\"W\n\x16GetJobPipelineResponse\x12=\n\x08pipeline\x18\x01 \x01(\x0b\x32+.org.apache.beam.model.pipeline.v1.Pipeline\"$\n\x12JobMessagesRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t\"\xd1\x02\n\nJobMessage\x12\x12\n\nmessage_id\x18\x01 \x01(\t\x12\x0c\n\x04time\x18\x02 \x01(\t\x12Y\n\nimportance\x18\x03 \x01(\x0e\x32\x45.org.apache.beam.model.job_management.v1.JobMessage.MessageImportance\x12\x14\n\x0cmessage_text\x18\x04 \x01(\t\"\xaf\x01\n\x11MessageImportance\x12\"\n\x1eMESSAGE_IMPORTANCE_UNSPECIFIED\x10\x00\x12\x15\n\x11JOB_MESSAGE_DEBUG\x10\x01\x12\x18\n\x14JOB_MESSAGE_DETAILED\x10\x02\x12\x15\n\x11JOB_MESSAGE_BASIC\x10\x03\x12\x17\n\x13JOB_MESSAGE_WARNING\x10\x04\x12\x15\n\x11JOB_MESSAGE_ERROR\x10\x05\"\xc4\x01\n\x13JobMessagesResponse\x12O\n\x10message_response\x18\x01 \x01(\x0b\x32\x33.org.apache.beam.model.job_management.v1.JobMessageH\x00\x12P\n\x0estate_response\x18\x02 \x01(\x0b\x32\x36.org.apache.beam.model.job_management.v1.JobStateEventH\x00\x42\n\n\x08response\"\xb7\x01\n\x08JobState\"\xaa\x01\n\x04\x45num\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x0b\n\x07STOPPED\x10\x01\x12\x0b\n\x07RUNNING\x10\x02\x12\x08\n\x04\x44ONE\x10\x03\x12\n\n\x06\x46\x41ILED\x10\x04\x12\r\n\tCANCELLED\x10\x05\x12\x0b\n\x07UPDATED\x10\x06\x12\x0c\n\x08\x44RAINING\x10\x07\x12\x0b\n\x07\x44RAINED\x10\x08\x12\x0c\n\x08STARTING\x10\t\x12\x0e\n\nCANCELLING\x10\n\x12\x0c\n\x08UPDATING\x10\x0b\"&\n\x14GetJobMetricsRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t\"`\n\x15GetJobMetricsResponse\x12G\n\x07metrics\x18\x01 \x01(\x0b\x32\x36.org.apache.beam.model.job_management.v1.MetricResults\"\x9b\x01\n\rMetricResults\x12\x44\n\tattempted\x18\x01 \x03(\x0b\x32\x31.org.apache.beam.model.pipeline.v1.MonitoringInfo\x12\x44\n\tcommitted\x18\x02 \x03(\x0b\x32\x31.org.apache.beam.model.pipeline.v1.MonitoringInfo\" \n\x1e\x44\x65scribePipelineOptionsRequest\"e\n\x12PipelineOptionType\"O\n\x04\x45num\x12\n\n\x06STRING\x10\x00\x12\x0b\n\x07\x42OOLEAN\x10\x01\x12\x0b\n\x07INTEGER\x10\x02\x12\n\n\x06NUMBER\x10\x03\x12\t\n\x05\x41RRAY\x10\x04\x12\n\n\x06OBJECT\x10\x05\"\xb3\x01\n\x18PipelineOptionDescriptor\x12\x0c\n\x04name\x18\x01 \x01(\t\x12N\n\x04type\x18\x02 \x01(\x0e\x32@.org.apache.beam.model.job_management.v1.PipelineOptionType.Enum\x12\x13\n\x0b\x64\x65scription\x18\x03 \x01(\t\x12\x15\n\rdefault_value\x18\x04 \x01(\t\x12\r\n\x05group\x18\x05 \x01(\t\"u\n\x1f\x44\x65scribePipelineOptionsResponse\x12R\n\x07options\x18\x01 \x03(\x0b\x32\x41.org.apache.beam.model.job_management.v1.PipelineOptionDescriptor2\xf4\x0b\n\nJobService\x12\x82\x01\n\x07Prepare\x12:.org.apache.beam.model.job_management.v1.PrepareJobRequest\x1a;.org.apache.beam.model.job_management.v1.PrepareJobResponse\x12v\n\x03Run\x12\x36.org.apache.beam.model.job_management.v1.RunJobRequest\x1a\x37.org.apache.beam.model.job_management.v1.RunJobResponse\x12|\n\x07GetJobs\x12\x37.org.apache.beam.model.job_management.v1.GetJobsRequest\x1a\x38.org.apache.beam.model.job_management.v1.GetJobsResponse\x12\x7f\n\x08GetState\x12;.org.apache.beam.model.job_management.v1.GetJobStateRequest\x1a\x36.org.apache.beam.model.job_management.v1.JobStateEvent\x12\x8e\x01\n\x0bGetPipeline\x12>.org.apache.beam.model.job_management.v1.GetJobPipelineRequest\x1a?.org.apache.beam.model.job_management.v1.GetJobPipelineResponse\x12\x7f\n\x06\x43\x61ncel\x12\x39.org.apache.beam.model.job_management.v1.CancelJobRequest\x1a:.org.apache.beam.model.job_management.v1.CancelJobResponse\x12|\n\x05\x44rain\x12\x38.org.apache.beam.model.job_management.v1.DrainJobRequest\x1a\x39.org.apache.beam.model.job_management.v1.DrainJobResponse\x12\x87\x01\n\x0eGetStateStream\x12;.org.apache.beam.model.job_management.v1.GetJobStateRequest\x1a\x36.org.apache.beam.model.job_management.v1.JobStateEvent0\x01\x12\x8f\x01\n\x10GetMessageStream\x12;.org.apache.beam.model.job_management.v1.JobMessagesRequest\x1a<.org.apache.beam.model.job_management.v1.JobMessagesResponse0\x01\x12\x8e\x01\n\rGetJobMetrics\x12=.org.apache.beam.model.job_management.v1.GetJobMetricsRequest\x1a>.org.apache.beam.model.job_management.v1.GetJobMetricsResponse\x12\xac\x01\n\x17\x44\x65scribePipelineOptions\x12G.org.apache.beam.model.job_management.v1.DescribePipelineOptionsRequest\x1aH.org.apache.beam.model.job_management.v1.DescribePipelineOptionsResponseB\x84\x01\n&org.apache.beam.model.jobmanagement.v1B\x06JobApiZRgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/jobmanagement_v1;jobmanagement_v1b\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'org.apache.beam.model.job_management.v1.beam_job_api_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n&org.apache.beam.model.jobmanagement.v1B\006JobApiZRgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/jobmanagement_v1;jobmanagement_v1' _globals['_PREPAREJOBREQUEST']._serialized_start=324 _globals['_PREPAREJOBREQUEST']._serialized_end=475 _globals['_PREPAREJOBRESPONSE']._serialized_start=478 _globals['_PREPAREJOBRESPONSE']._serialized_end=645 _globals['_RUNJOBREQUEST']._serialized_start=647 _globals['_RUNJOBREQUEST']._serialized_end=711 _globals['_RUNJOBRESPONSE']._serialized_start=713 _globals['_RUNJOBRESPONSE']._serialized_end=745 _globals['_CANCELJOBREQUEST']._serialized_start=747 _globals['_CANCELJOBREQUEST']._serialized_end=781 _globals['_CANCELJOBRESPONSE']._serialized_start=783 _globals['_CANCELJOBRESPONSE']._serialized_end=873 _globals['_DRAINJOBREQUEST']._serialized_start=875 _globals['_DRAINJOBREQUEST']._serialized_end=908 _globals['_DRAINJOBRESPONSE']._serialized_start=910 _globals['_DRAINJOBRESPONSE']._serialized_end=999 _globals['_JOBINFO']._serialized_start=1002 _globals['_JOBINFO']._serialized_end=1167 _globals['_GETJOBSREQUEST']._serialized_start=1169 _globals['_GETJOBSREQUEST']._serialized_end=1185 _globals['_GETJOBSRESPONSE']._serialized_start=1187 _globals['_GETJOBSRESPONSE']._serialized_end=1272 _globals['_GETJOBSTATEREQUEST']._serialized_start=1274 _globals['_GETJOBSTATEREQUEST']._serialized_end=1310 _globals['_JOBSTATEEVENT']._serialized_start=1313 _globals['_JOBSTATEEVENT']._serialized_end=1446 _globals['_GETJOBPIPELINEREQUEST']._serialized_start=1448 _globals['_GETJOBPIPELINEREQUEST']._serialized_end=1487 _globals['_GETJOBPIPELINERESPONSE']._serialized_start=1489 _globals['_GETJOBPIPELINERESPONSE']._serialized_end=1576 _globals['_JOBMESSAGESREQUEST']._serialized_start=1578 _globals['_JOBMESSAGESREQUEST']._serialized_end=1614 _globals['_JOBMESSAGE']._serialized_start=1617 _globals['_JOBMESSAGE']._serialized_end=1954 _globals['_JOBMESSAGE_MESSAGEIMPORTANCE']._serialized_start=1779 _globals['_JOBMESSAGE_MESSAGEIMPORTANCE']._serialized_end=1954 _globals['_JOBMESSAGESRESPONSE']._serialized_start=1957 _globals['_JOBMESSAGESRESPONSE']._serialized_end=2153 _globals['_JOBSTATE']._serialized_start=2156 _globals['_JOBSTATE']._serialized_end=2339 _globals['_JOBSTATE_ENUM']._serialized_start=2169 _globals['_JOBSTATE_ENUM']._serialized_end=2339 _globals['_GETJOBMETRICSREQUEST']._serialized_start=2341 _globals['_GETJOBMETRICSREQUEST']._serialized_end=2379 _globals['_GETJOBMETRICSRESPONSE']._serialized_start=2381 _globals['_GETJOBMETRICSRESPONSE']._serialized_end=2477 _globals['_METRICRESULTS']._serialized_start=2480 _globals['_METRICRESULTS']._serialized_end=2635 _globals['_DESCRIBEPIPELINEOPTIONSREQUEST']._serialized_start=2637 _globals['_DESCRIBEPIPELINEOPTIONSREQUEST']._serialized_end=2669 _globals['_PIPELINEOPTIONTYPE']._serialized_start=2671 _globals['_PIPELINEOPTIONTYPE']._serialized_end=2772 _globals['_PIPELINEOPTIONTYPE_ENUM']._serialized_start=2693 _globals['_PIPELINEOPTIONTYPE_ENUM']._serialized_end=2772 _globals['_PIPELINEOPTIONDESCRIPTOR']._serialized_start=2775 _globals['_PIPELINEOPTIONDESCRIPTOR']._serialized_end=2954 _globals['_DESCRIBEPIPELINEOPTIONSRESPONSE']._serialized_start=2956 _globals['_DESCRIBEPIPELINEOPTIONSRESPONSE']._serialized_end=3073 _globals['_JOBSERVICE']._serialized_start=3076 _globals['_JOBSERVICE']._serialized_end=4600 # @@protoc_insertion_point(module_scope) ././@PaxHeader0000000000000000000000000000021300000000000010211 xustar00117 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_job_api_pb2.pyi 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_job_api_0000644000175100017510000006544415206113151032770 0ustar00runnerrunner""" @generated by mypy-protobuf. Do not edit manually! isort:skip_file Protocol Buffers describing the Job API, api for communicating with a runner for job submission over GRPC. """ import builtins import collections.abc import google.protobuf.descriptor import google.protobuf.internal.containers import google.protobuf.internal.enum_type_wrapper import google.protobuf.message import google.protobuf.struct_pb2 import google.protobuf.timestamp_pb2 import org.apache.beam.model.pipeline.v1.beam_runner_api_pb2 import org.apache.beam.model.pipeline.v1.endpoints_pb2 import org.apache.beam.model.pipeline.v1.metrics_pb2 import sys import typing if sys.version_info >= (3, 10): import typing as typing_extensions else: import typing_extensions DESCRIPTOR: google.protobuf.descriptor.FileDescriptor @typing_extensions.final class PrepareJobRequest(google.protobuf.message.Message): """Prepare is a synchronous request that returns a preparationId back Throws error GRPC_STATUS_UNAVAILABLE if server is down Throws error ALREADY_EXISTS if the jobName is reused. Runners are permitted to deduplicate based on the name of the job. Throws error UNKNOWN for all other issues """ DESCRIPTOR: google.protobuf.descriptor.Descriptor PIPELINE_FIELD_NUMBER: builtins.int PIPELINE_OPTIONS_FIELD_NUMBER: builtins.int JOB_NAME_FIELD_NUMBER: builtins.int @property def pipeline(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Pipeline: """(required)""" @property def pipeline_options(self) -> google.protobuf.struct_pb2.Struct: """(required)""" job_name: builtins.str """(required)""" def __init__( self, *, pipeline: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Pipeline | None = ..., pipeline_options: google.protobuf.struct_pb2.Struct | None = ..., job_name: builtins.str | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["pipeline", b"pipeline", "pipeline_options", b"pipeline_options"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["job_name", b"job_name", "pipeline", b"pipeline", "pipeline_options", b"pipeline_options"]) -> None: ... global___PrepareJobRequest = PrepareJobRequest @typing_extensions.final class PrepareJobResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor PREPARATION_ID_FIELD_NUMBER: builtins.int ARTIFACT_STAGING_ENDPOINT_FIELD_NUMBER: builtins.int STAGING_SESSION_TOKEN_FIELD_NUMBER: builtins.int preparation_id: builtins.str """(required) The ID used to associate calls made while preparing the job. preparationId is used to run the job. """ @property def artifact_staging_endpoint(self) -> org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor: """An endpoint which exposes the Beam Artifact Staging API. Artifacts used by the job should be staged to this endpoint, and will be available during job execution. """ staging_session_token: builtins.str """(required) Token for the artifact staging. This token also represent an artifact staging session with the artifact staging service. """ def __init__( self, *, preparation_id: builtins.str | None = ..., artifact_staging_endpoint: org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor | None = ..., staging_session_token: builtins.str | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["artifact_staging_endpoint", b"artifact_staging_endpoint"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["artifact_staging_endpoint", b"artifact_staging_endpoint", "preparation_id", b"preparation_id", "staging_session_token", b"staging_session_token"]) -> None: ... global___PrepareJobResponse = PrepareJobResponse @typing_extensions.final class RunJobRequest(google.protobuf.message.Message): """Run is a synchronous request that returns a jobId back. Throws error GRPC_STATUS_UNAVAILABLE if server is down Throws error NOT_FOUND if the preparation ID does not exist Throws error UNKNOWN for all other issues """ DESCRIPTOR: google.protobuf.descriptor.Descriptor PREPARATION_ID_FIELD_NUMBER: builtins.int RETRIEVAL_TOKEN_FIELD_NUMBER: builtins.int preparation_id: builtins.str """(required) The ID provided by an earlier call to prepare. Runs the job. All prerequisite tasks must have been completed. """ retrieval_token: builtins.str """(optional) If any artifacts have been staged for this job, contains the retrieval_token returned from the CommitManifestResponse. """ def __init__( self, *, preparation_id: builtins.str | None = ..., retrieval_token: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["preparation_id", b"preparation_id", "retrieval_token", b"retrieval_token"]) -> None: ... global___RunJobRequest = RunJobRequest @typing_extensions.final class RunJobResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor JOB_ID_FIELD_NUMBER: builtins.int job_id: builtins.str """(required) The ID for the executing job""" def __init__( self, *, job_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["job_id", b"job_id"]) -> None: ... global___RunJobResponse = RunJobResponse @typing_extensions.final class CancelJobRequest(google.protobuf.message.Message): """Cancel is a synchronus request that returns a job state back Throws error GRPC_STATUS_UNAVAILABLE if server is down Throws error NOT_FOUND if the jobId is not found """ DESCRIPTOR: google.protobuf.descriptor.Descriptor JOB_ID_FIELD_NUMBER: builtins.int job_id: builtins.str """(required)""" def __init__( self, *, job_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["job_id", b"job_id"]) -> None: ... global___CancelJobRequest = CancelJobRequest @typing_extensions.final class CancelJobResponse(google.protobuf.message.Message): """Valid responses include any terminal state or CANCELLING""" DESCRIPTOR: google.protobuf.descriptor.Descriptor STATE_FIELD_NUMBER: builtins.int state: global___JobState.Enum.ValueType """(required)""" def __init__( self, *, state: global___JobState.Enum.ValueType | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["state", b"state"]) -> None: ... global___CancelJobResponse = CancelJobResponse @typing_extensions.final class DrainJobRequest(google.protobuf.message.Message): """Drain is a request to: - stop a pipeline Job from ingesting new data from input sources - process remaining data in the pipeline - truncate Splittable DoFns """ DESCRIPTOR: google.protobuf.descriptor.Descriptor JOB_ID_FIELD_NUMBER: builtins.int job_id: builtins.str """(required)""" def __init__( self, *, job_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["job_id", b"job_id"]) -> None: ... global___DrainJobRequest = DrainJobRequest @typing_extensions.final class DrainJobResponse(google.protobuf.message.Message): """Valid responses include any terminal state or DRAINING""" DESCRIPTOR: google.protobuf.descriptor.Descriptor STATE_FIELD_NUMBER: builtins.int state: global___JobState.Enum.ValueType """(required)""" def __init__( self, *, state: global___JobState.Enum.ValueType | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["state", b"state"]) -> None: ... global___DrainJobResponse = DrainJobResponse @typing_extensions.final class JobInfo(google.protobuf.message.Message): """A subset of info provided by ProvisionApi.ProvisionInfo""" DESCRIPTOR: google.protobuf.descriptor.Descriptor JOB_ID_FIELD_NUMBER: builtins.int JOB_NAME_FIELD_NUMBER: builtins.int PIPELINE_OPTIONS_FIELD_NUMBER: builtins.int STATE_FIELD_NUMBER: builtins.int job_id: builtins.str """(required)""" job_name: builtins.str """(required)""" @property def pipeline_options(self) -> google.protobuf.struct_pb2.Struct: """(required)""" state: global___JobState.Enum.ValueType """(required)""" def __init__( self, *, job_id: builtins.str | None = ..., job_name: builtins.str | None = ..., pipeline_options: google.protobuf.struct_pb2.Struct | None = ..., state: global___JobState.Enum.ValueType | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["pipeline_options", b"pipeline_options"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["job_id", b"job_id", "job_name", b"job_name", "pipeline_options", b"pipeline_options", "state", b"state"]) -> None: ... global___JobInfo = JobInfo @typing_extensions.final class GetJobsRequest(google.protobuf.message.Message): """GetJobs is a synchronus request that returns a list of invoked jobs back Throws error GRPC_STATUS_UNAVAILABLE if server is down """ DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... global___GetJobsRequest = GetJobsRequest @typing_extensions.final class GetJobsResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor JOB_INFO_FIELD_NUMBER: builtins.int @property def job_info(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___JobInfo]: """(required)""" def __init__( self, *, job_info: collections.abc.Iterable[global___JobInfo] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["job_info", b"job_info"]) -> None: ... global___GetJobsResponse = GetJobsResponse @typing_extensions.final class GetJobStateRequest(google.protobuf.message.Message): """GetState is a synchronus request that returns a job state back Throws error GRPC_STATUS_UNAVAILABLE if server is down Throws error NOT_FOUND if the jobId is not found """ DESCRIPTOR: google.protobuf.descriptor.Descriptor JOB_ID_FIELD_NUMBER: builtins.int job_id: builtins.str """(required)""" def __init__( self, *, job_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["job_id", b"job_id"]) -> None: ... global___GetJobStateRequest = GetJobStateRequest @typing_extensions.final class JobStateEvent(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor STATE_FIELD_NUMBER: builtins.int TIMESTAMP_FIELD_NUMBER: builtins.int state: global___JobState.Enum.ValueType """(required)""" @property def timestamp(self) -> google.protobuf.timestamp_pb2.Timestamp: """(required)""" def __init__( self, *, state: global___JobState.Enum.ValueType | None = ..., timestamp: google.protobuf.timestamp_pb2.Timestamp | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["timestamp", b"timestamp"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["state", b"state", "timestamp", b"timestamp"]) -> None: ... global___JobStateEvent = JobStateEvent @typing_extensions.final class GetJobPipelineRequest(google.protobuf.message.Message): """GetPipeline is a synchronus request that returns a pipeline back Throws error GRPC_STATUS_UNAVAILABLE if server is down Throws error NOT_FOUND if the jobId is not found """ DESCRIPTOR: google.protobuf.descriptor.Descriptor JOB_ID_FIELD_NUMBER: builtins.int job_id: builtins.str """(required)""" def __init__( self, *, job_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["job_id", b"job_id"]) -> None: ... global___GetJobPipelineRequest = GetJobPipelineRequest @typing_extensions.final class GetJobPipelineResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor PIPELINE_FIELD_NUMBER: builtins.int @property def pipeline(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Pipeline: """(required)""" def __init__( self, *, pipeline: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.Pipeline | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["pipeline", b"pipeline"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["pipeline", b"pipeline"]) -> None: ... global___GetJobPipelineResponse = GetJobPipelineResponse @typing_extensions.final class JobMessagesRequest(google.protobuf.message.Message): """GetJobMessages is a streaming api for streaming job messages from the service One request will connect you to the job and you'll get a stream of job state and job messages back; one is used for logging and the other for detecting the job ended. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor JOB_ID_FIELD_NUMBER: builtins.int job_id: builtins.str """(required)""" def __init__( self, *, job_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["job_id", b"job_id"]) -> None: ... global___JobMessagesRequest = JobMessagesRequest @typing_extensions.final class JobMessage(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor class _MessageImportance: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _MessageImportanceEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[JobMessage._MessageImportance.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor MESSAGE_IMPORTANCE_UNSPECIFIED: JobMessage._MessageImportance.ValueType # 0 JOB_MESSAGE_DEBUG: JobMessage._MessageImportance.ValueType # 1 JOB_MESSAGE_DETAILED: JobMessage._MessageImportance.ValueType # 2 JOB_MESSAGE_BASIC: JobMessage._MessageImportance.ValueType # 3 JOB_MESSAGE_WARNING: JobMessage._MessageImportance.ValueType # 4 JOB_MESSAGE_ERROR: JobMessage._MessageImportance.ValueType # 5 class MessageImportance(_MessageImportance, metaclass=_MessageImportanceEnumTypeWrapper): ... MESSAGE_IMPORTANCE_UNSPECIFIED: JobMessage.MessageImportance.ValueType # 0 JOB_MESSAGE_DEBUG: JobMessage.MessageImportance.ValueType # 1 JOB_MESSAGE_DETAILED: JobMessage.MessageImportance.ValueType # 2 JOB_MESSAGE_BASIC: JobMessage.MessageImportance.ValueType # 3 JOB_MESSAGE_WARNING: JobMessage.MessageImportance.ValueType # 4 JOB_MESSAGE_ERROR: JobMessage.MessageImportance.ValueType # 5 MESSAGE_ID_FIELD_NUMBER: builtins.int TIME_FIELD_NUMBER: builtins.int IMPORTANCE_FIELD_NUMBER: builtins.int MESSAGE_TEXT_FIELD_NUMBER: builtins.int message_id: builtins.str time: builtins.str importance: global___JobMessage.MessageImportance.ValueType message_text: builtins.str def __init__( self, *, message_id: builtins.str | None = ..., time: builtins.str | None = ..., importance: global___JobMessage.MessageImportance.ValueType | None = ..., message_text: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["importance", b"importance", "message_id", b"message_id", "message_text", b"message_text", "time", b"time"]) -> None: ... global___JobMessage = JobMessage @typing_extensions.final class JobMessagesResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor MESSAGE_RESPONSE_FIELD_NUMBER: builtins.int STATE_RESPONSE_FIELD_NUMBER: builtins.int @property def message_response(self) -> global___JobMessage: ... @property def state_response(self) -> global___JobStateEvent: ... def __init__( self, *, message_response: global___JobMessage | None = ..., state_response: global___JobStateEvent | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["message_response", b"message_response", "response", b"response", "state_response", b"state_response"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["message_response", b"message_response", "response", b"response", "state_response", b"state_response"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["response", b"response"]) -> typing_extensions.Literal["message_response", "state_response"] | None: ... global___JobMessagesResponse = JobMessagesResponse @typing_extensions.final class JobState(google.protobuf.message.Message): """Enumeration of all JobStates The state transition diagram is: STOPPED -> STARTING -> RUNNING -> DONE \\> FAILED \\> CANCELLING -> CANCELLED \\> UPDATING -> UPDATED \\> DRAINING -> DRAINED Transitions are optional such that a job may go from STOPPED to RUNNING without needing to pass through STARTING. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[JobState._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor UNSPECIFIED: JobState._Enum.ValueType # 0 """The job state reported by a runner cannot be interpreted by the SDK.""" STOPPED: JobState._Enum.ValueType # 1 """The job has not yet started.""" RUNNING: JobState._Enum.ValueType # 2 """The job is currently running.""" DONE: JobState._Enum.ValueType # 3 """The job has successfully completed. (terminal)""" FAILED: JobState._Enum.ValueType # 4 """The job has failed. (terminal)""" CANCELLED: JobState._Enum.ValueType # 5 """The job has been explicitly cancelled. (terminal)""" UPDATED: JobState._Enum.ValueType # 6 """The job has been updated. (terminal)""" DRAINING: JobState._Enum.ValueType # 7 """The job is draining its data. (optional)""" DRAINED: JobState._Enum.ValueType # 8 """The job has completed draining its data. (terminal)""" STARTING: JobState._Enum.ValueType # 9 """The job is starting up.""" CANCELLING: JobState._Enum.ValueType # 10 """The job is cancelling. (optional)""" UPDATING: JobState._Enum.ValueType # 11 """The job is in the process of being updated. (optional)""" class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... UNSPECIFIED: JobState.Enum.ValueType # 0 """The job state reported by a runner cannot be interpreted by the SDK.""" STOPPED: JobState.Enum.ValueType # 1 """The job has not yet started.""" RUNNING: JobState.Enum.ValueType # 2 """The job is currently running.""" DONE: JobState.Enum.ValueType # 3 """The job has successfully completed. (terminal)""" FAILED: JobState.Enum.ValueType # 4 """The job has failed. (terminal)""" CANCELLED: JobState.Enum.ValueType # 5 """The job has been explicitly cancelled. (terminal)""" UPDATED: JobState.Enum.ValueType # 6 """The job has been updated. (terminal)""" DRAINING: JobState.Enum.ValueType # 7 """The job is draining its data. (optional)""" DRAINED: JobState.Enum.ValueType # 8 """The job has completed draining its data. (terminal)""" STARTING: JobState.Enum.ValueType # 9 """The job is starting up.""" CANCELLING: JobState.Enum.ValueType # 10 """The job is cancelling. (optional)""" UPDATING: JobState.Enum.ValueType # 11 """The job is in the process of being updated. (optional)""" def __init__( self, ) -> None: ... global___JobState = JobState @typing_extensions.final class GetJobMetricsRequest(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor JOB_ID_FIELD_NUMBER: builtins.int job_id: builtins.str """(required)""" def __init__( self, *, job_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["job_id", b"job_id"]) -> None: ... global___GetJobMetricsRequest = GetJobMetricsRequest @typing_extensions.final class GetJobMetricsResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor METRICS_FIELD_NUMBER: builtins.int @property def metrics(self) -> global___MetricResults: ... def __init__( self, *, metrics: global___MetricResults | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["metrics", b"metrics"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["metrics", b"metrics"]) -> None: ... global___GetJobMetricsResponse = GetJobMetricsResponse @typing_extensions.final class MetricResults(google.protobuf.message.Message): """All metrics for a given job. Runners may support one or the other or both.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor ATTEMPTED_FIELD_NUMBER: builtins.int COMMITTED_FIELD_NUMBER: builtins.int @property def attempted(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[org.apache.beam.model.pipeline.v1.metrics_pb2.MonitoringInfo]: ... @property def committed(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[org.apache.beam.model.pipeline.v1.metrics_pb2.MonitoringInfo]: ... def __init__( self, *, attempted: collections.abc.Iterable[org.apache.beam.model.pipeline.v1.metrics_pb2.MonitoringInfo] | None = ..., committed: collections.abc.Iterable[org.apache.beam.model.pipeline.v1.metrics_pb2.MonitoringInfo] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["attempted", b"attempted", "committed", b"committed"]) -> None: ... global___MetricResults = MetricResults @typing_extensions.final class DescribePipelineOptionsRequest(google.protobuf.message.Message): """DescribePipelineOptions provides metadata about the options supported by a runner. It will be used by the SDK client to validate the options specified by or list available options to the user. Throws error GRPC_STATUS_UNAVAILABLE if server is down """ DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... global___DescribePipelineOptionsRequest = DescribePipelineOptionsRequest @typing_extensions.final class PipelineOptionType(google.protobuf.message.Message): """Type for pipeline options. Types mirror those of JSON, since that's how pipeline options are serialized. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[PipelineOptionType._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor STRING: PipelineOptionType._Enum.ValueType # 0 BOOLEAN: PipelineOptionType._Enum.ValueType # 1 INTEGER: PipelineOptionType._Enum.ValueType # 2 """whole numbers, see https://json-schema.org/understanding-json-schema/reference/numeric.html""" NUMBER: PipelineOptionType._Enum.ValueType # 3 ARRAY: PipelineOptionType._Enum.ValueType # 4 OBJECT: PipelineOptionType._Enum.ValueType # 5 class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... STRING: PipelineOptionType.Enum.ValueType # 0 BOOLEAN: PipelineOptionType.Enum.ValueType # 1 INTEGER: PipelineOptionType.Enum.ValueType # 2 """whole numbers, see https://json-schema.org/understanding-json-schema/reference/numeric.html""" NUMBER: PipelineOptionType.Enum.ValueType # 3 ARRAY: PipelineOptionType.Enum.ValueType # 4 OBJECT: PipelineOptionType.Enum.ValueType # 5 def __init__( self, ) -> None: ... global___PipelineOptionType = PipelineOptionType @typing_extensions.final class PipelineOptionDescriptor(google.protobuf.message.Message): """Metadata for a pipeline option.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor NAME_FIELD_NUMBER: builtins.int TYPE_FIELD_NUMBER: builtins.int DESCRIPTION_FIELD_NUMBER: builtins.int DEFAULT_VALUE_FIELD_NUMBER: builtins.int GROUP_FIELD_NUMBER: builtins.int name: builtins.str """(Required) The option name.""" type: global___PipelineOptionType.Enum.ValueType """(Required) Type of option.""" description: builtins.str """(Optional) Description suitable for display / help text.""" default_value: builtins.str """(Optional) Default value.""" group: builtins.str """(Required) The group this option belongs to.""" def __init__( self, *, name: builtins.str | None = ..., type: global___PipelineOptionType.Enum.ValueType | None = ..., description: builtins.str | None = ..., default_value: builtins.str | None = ..., group: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["default_value", b"default_value", "description", b"description", "group", b"group", "name", b"name", "type", b"type"]) -> None: ... global___PipelineOptionDescriptor = PipelineOptionDescriptor @typing_extensions.final class DescribePipelineOptionsResponse(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor OPTIONS_FIELD_NUMBER: builtins.int @property def options(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___PipelineOptionDescriptor]: """List of pipeline option descriptors.""" def __init__( self, *, options: collections.abc.Iterable[global___PipelineOptionDescriptor] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["options", b"options"]) -> None: ... global___DescribePipelineOptionsResponse = DescribePipelineOptionsResponse ././@PaxHeader0000000000000000000000000000021700000000000010215 xustar00121 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_job_api_pb2_grpc.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_job_api_0000644000175100017510000005531115206113151032760 0ustar00runnerrunner# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc from . import beam_job_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2 class JobServiceStub(object): """Job Service for running RunnerAPI pipelines """ def __init__(self, channel): """Constructor. Args: channel: A grpc.Channel. """ self.Prepare = channel.unary_unary( '/org.apache.beam.model.job_management.v1.JobService/Prepare', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.PrepareJobRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.PrepareJobResponse.FromString, ) self.Run = channel.unary_unary( '/org.apache.beam.model.job_management.v1.JobService/Run', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.RunJobRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.RunJobResponse.FromString, ) self.GetJobs = channel.unary_unary( '/org.apache.beam.model.job_management.v1.JobService/GetJobs', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobsRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobsResponse.FromString, ) self.GetState = channel.unary_unary( '/org.apache.beam.model.job_management.v1.JobService/GetState', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobStateRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.JobStateEvent.FromString, ) self.GetPipeline = channel.unary_unary( '/org.apache.beam.model.job_management.v1.JobService/GetPipeline', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobPipelineRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobPipelineResponse.FromString, ) self.Cancel = channel.unary_unary( '/org.apache.beam.model.job_management.v1.JobService/Cancel', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.CancelJobRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.CancelJobResponse.FromString, ) self.Drain = channel.unary_unary( '/org.apache.beam.model.job_management.v1.JobService/Drain', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.DrainJobRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.DrainJobResponse.FromString, ) self.GetStateStream = channel.unary_stream( '/org.apache.beam.model.job_management.v1.JobService/GetStateStream', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobStateRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.JobStateEvent.FromString, ) self.GetMessageStream = channel.unary_stream( '/org.apache.beam.model.job_management.v1.JobService/GetMessageStream', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.JobMessagesRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.JobMessagesResponse.FromString, ) self.GetJobMetrics = channel.unary_unary( '/org.apache.beam.model.job_management.v1.JobService/GetJobMetrics', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobMetricsRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobMetricsResponse.FromString, ) self.DescribePipelineOptions = channel.unary_unary( '/org.apache.beam.model.job_management.v1.JobService/DescribePipelineOptions', request_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.DescribePipelineOptionsRequest.SerializeToString, response_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.DescribePipelineOptionsResponse.FromString, ) class JobServiceServicer(object): """Job Service for running RunnerAPI pipelines """ def Prepare(self, request, context): """Prepare a job for execution. The job will not be executed until a call is made to run with the returned preparationId. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def Run(self, request, context): """Submit the job for execution """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def GetJobs(self, request, context): """Get a list of all invoked jobs """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def GetState(self, request, context): """Get the current state of the job """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def GetPipeline(self, request, context): """Get the job's pipeline """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def Cancel(self, request, context): """Cancel the job """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def Drain(self, request, context): """Drain the job """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def GetStateStream(self, request, context): """Subscribe to a stream of state changes of the job, will immediately return the current state of the job as the first response. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def GetMessageStream(self, request, context): """Subscribe to a stream of state changes and messages from the job """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def GetJobMetrics(self, request, context): """Fetch metrics for a given job """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def DescribePipelineOptions(self, request, context): """Get the supported pipeline options of the runner """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') def add_JobServiceServicer_to_server(servicer, server): rpc_method_handlers = { 'Prepare': grpc.unary_unary_rpc_method_handler( servicer.Prepare, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.PrepareJobRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.PrepareJobResponse.SerializeToString, ), 'Run': grpc.unary_unary_rpc_method_handler( servicer.Run, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.RunJobRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.RunJobResponse.SerializeToString, ), 'GetJobs': grpc.unary_unary_rpc_method_handler( servicer.GetJobs, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobsRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobsResponse.SerializeToString, ), 'GetState': grpc.unary_unary_rpc_method_handler( servicer.GetState, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobStateRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.JobStateEvent.SerializeToString, ), 'GetPipeline': grpc.unary_unary_rpc_method_handler( servicer.GetPipeline, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobPipelineRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobPipelineResponse.SerializeToString, ), 'Cancel': grpc.unary_unary_rpc_method_handler( servicer.Cancel, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.CancelJobRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.CancelJobResponse.SerializeToString, ), 'Drain': grpc.unary_unary_rpc_method_handler( servicer.Drain, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.DrainJobRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.DrainJobResponse.SerializeToString, ), 'GetStateStream': grpc.unary_stream_rpc_method_handler( servicer.GetStateStream, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobStateRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.JobStateEvent.SerializeToString, ), 'GetMessageStream': grpc.unary_stream_rpc_method_handler( servicer.GetMessageStream, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.JobMessagesRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.JobMessagesResponse.SerializeToString, ), 'GetJobMetrics': grpc.unary_unary_rpc_method_handler( servicer.GetJobMetrics, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobMetricsRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobMetricsResponse.SerializeToString, ), 'DescribePipelineOptions': grpc.unary_unary_rpc_method_handler( servicer.DescribePipelineOptions, request_deserializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.DescribePipelineOptionsRequest.FromString, response_serializer=org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.DescribePipelineOptionsResponse.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( 'org.apache.beam.model.job_management.v1.JobService', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) # This class is part of an EXPERIMENTAL API. class JobService(object): """Job Service for running RunnerAPI pipelines """ @staticmethod def Prepare(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.job_management.v1.JobService/Prepare', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.PrepareJobRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.PrepareJobResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def Run(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.job_management.v1.JobService/Run', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.RunJobRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.RunJobResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def GetJobs(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.job_management.v1.JobService/GetJobs', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobsRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobsResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def GetState(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.job_management.v1.JobService/GetState', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobStateRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.JobStateEvent.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def GetPipeline(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.job_management.v1.JobService/GetPipeline', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobPipelineRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobPipelineResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def Cancel(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.job_management.v1.JobService/Cancel', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.CancelJobRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.CancelJobResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def Drain(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.job_management.v1.JobService/Drain', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.DrainJobRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.DrainJobResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def GetStateStream(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_stream(request, target, '/org.apache.beam.model.job_management.v1.JobService/GetStateStream', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobStateRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.JobStateEvent.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def GetMessageStream(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_stream(request, target, '/org.apache.beam.model.job_management.v1.JobService/GetMessageStream', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.JobMessagesRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.JobMessagesResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def GetJobMetrics(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.job_management.v1.JobService/GetJobMetrics', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobMetricsRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.GetJobMetricsResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) @staticmethod def DescribePipelineOptions(request, target, options=(), channel_credentials=None, call_credentials=None, insecure=False, compression=None, wait_for_ready=None, timeout=None, metadata=None): return grpc.experimental.unary_unary(request, target, '/org.apache.beam.model.job_management.v1.JobService/DescribePipelineOptions', org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.DescribePipelineOptionsRequest.SerializeToString, org_dot_apache_dot_beam_dot_model_dot_job__management_dot_v1_dot_beam__job__api__pb2.DescribePipelineOptionsResponse.FromString, options, channel_credentials, insecure, call_credentials, compression, wait_for_ready, timeout, metadata) ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5404246 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/0000755000175100017510000000000015206113160026613 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/__init__.py0000644000175100017510000000142015206113151030721 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.551426 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/0000755000175100017510000000000015206113160027141 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/__init__.py0000644000175100017510000000254215206113151031255 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from . import external_transforms_pb2_urns from . import beam_runner_api_pb2_grpc from . import schema_pb2_grpc from . import external_transforms_pb2_grpc from . import metrics_pb2_grpc from . import standard_window_fns_pb2_urns from . import beam_runner_api_pb2 from . import schema_pb2 from . import standard_window_fns_pb2 from . import metrics_pb2_urns from . import standard_window_fns_pb2_grpc from . import endpoints_pb2_grpc from . import schema_pb2_urns from . import external_transforms_pb2 from . import metrics_pb2 from . import endpoints_pb2 from . import beam_runner_api_pb2_urns ././@PaxHeader0000000000000000000000000000020700000000000010214 xustar00113 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/beam_runner_api_pb2.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/beam_runner_api_pb20000644000175100017510000022364315206113151032767 0ustar00runnerrunner# -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # source: org/apache/beam/model/pipeline/v1/beam_runner_api.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() from . import endpoints_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_endpoints__pb2 from google.protobuf import descriptor_pb2 as google_dot_protobuf_dot_descriptor__pb2 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n7org/apache/beam/model/pipeline/v1/beam_runner_api.proto\x12!org.apache.beam.model.pipeline.v1\x1a\x31org/apache/beam/model/pipeline/v1/endpoints.proto\x1a google/protobuf/descriptor.proto\"\xc2\x01\n\rBeamConstants\"\xb0\x01\n\tConstants\x12\x31\n\x14MIN_TIMESTAMP_MILLIS\x10\x00\x1a\x17\xaa\xb4\xfa\xc2\x05\x11-9223372036854775\x12\x30\n\x14MAX_TIMESTAMP_MILLIS\x10\x01\x1a\x16\xaa\xb4\xfa\xc2\x05\x10\x39\x32\x32\x33\x33\x37\x32\x30\x33\x36\x38\x35\x34\x37\x37\x35\x12>\n\"GLOBAL_WINDOW_MAX_TIMESTAMP_MILLIS\x10\x02\x1a\x16\xaa\xb4\xfa\xc2\x05\x10\x39\x32\x32\x33\x33\x37\x31\x39\x35\x30\x34\x35\x34\x37\x37\x35\"\xb5\x07\n\nComponents\x12Q\n\ntransforms\x18\x01 \x03(\x0b\x32=.org.apache.beam.model.pipeline.v1.Components.TransformsEntry\x12U\n\x0cpcollections\x18\x02 \x03(\x0b\x32?.org.apache.beam.model.pipeline.v1.Components.PcollectionsEntry\x12\x64\n\x14windowing_strategies\x18\x03 \x03(\x0b\x32\x46.org.apache.beam.model.pipeline.v1.Components.WindowingStrategiesEntry\x12I\n\x06\x63oders\x18\x04 \x03(\x0b\x32\x39.org.apache.beam.model.pipeline.v1.Components.CodersEntry\x12U\n\x0c\x65nvironments\x18\x05 \x03(\x0b\x32?.org.apache.beam.model.pipeline.v1.Components.EnvironmentsEntry\x1a`\n\x0fTransformsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12<\n\x05value\x18\x02 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.PTransform:\x02\x38\x01\x1a\x63\n\x11PcollectionsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12=\n\x05value\x18\x02 \x01(\x0b\x32..org.apache.beam.model.pipeline.v1.PCollection:\x02\x38\x01\x1ap\n\x18WindowingStrategiesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x43\n\x05value\x18\x02 \x01(\x0b\x32\x34.org.apache.beam.model.pipeline.v1.WindowingStrategy:\x02\x38\x01\x1aW\n\x0b\x43odersEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x37\n\x05value\x18\x02 \x01(\x0b\x32(.org.apache.beam.model.pipeline.v1.Coder:\x02\x38\x01\x1a\x63\n\x11\x45nvironmentsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12=\n\x05value\x18\x02 \x01(\x0b\x32..org.apache.beam.model.pipeline.v1.Environment:\x02\x38\x01\"\xc5\x01\n\x08Pipeline\x12\x41\n\ncomponents\x18\x01 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.Components\x12\x1a\n\x12root_transform_ids\x18\x02 \x03(\t\x12\x44\n\x0c\x64isplay_data\x18\x03 \x03(\x0b\x32..org.apache.beam.model.pipeline.v1.DisplayData\x12\x14\n\x0crequirements\x18\x04 \x03(\t\"\xd5\x04\n\nPTransform\x12\x13\n\x0bunique_name\x18\x05 \x01(\t\x12=\n\x04spec\x18\x01 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\x12\x15\n\rsubtransforms\x18\x02 \x03(\t\x12I\n\x06inputs\x18\x03 \x03(\x0b\x32\x39.org.apache.beam.model.pipeline.v1.PTransform.InputsEntry\x12K\n\x07outputs\x18\x04 \x03(\x0b\x32:.org.apache.beam.model.pipeline.v1.PTransform.OutputsEntry\x12\x44\n\x0c\x64isplay_data\x18\x06 \x03(\x0b\x32..org.apache.beam.model.pipeline.v1.DisplayData\x12\x16\n\x0e\x65nvironment_id\x18\x07 \x01(\t\x12S\n\x0b\x61nnotations\x18\x08 \x03(\x0b\x32>.org.apache.beam.model.pipeline.v1.PTransform.AnnotationsEntry\x1a-\n\x0bInputsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\x1a.\n\x0cOutputsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\x1a\x32\n\x10\x41nnotationsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x0c:\x02\x38\x01\"\xdb\x11\n\x13StandardPTransforms\"\xdb\x03\n\nPrimitives\x12)\n\x06PAR_DO\x10\x00\x1a\x1d\xa2\xb4\xfa\xc2\x05\x17\x62\x65\x61m:transform:pardo:v1\x12,\n\x07\x46LATTEN\x10\x01\x1a\x1f\xa2\xb4\xfa\xc2\x05\x19\x62\x65\x61m:transform:flatten:v1\x12\x36\n\x0cGROUP_BY_KEY\x10\x02\x1a$\xa2\xb4\xfa\xc2\x05\x1e\x62\x65\x61m:transform:group_by_key:v1\x12,\n\x07IMPULSE\x10\x03\x1a\x1f\xa2\xb4\xfa\xc2\x05\x19\x62\x65\x61m:transform:impulse:v1\x12\x37\n\x0e\x41SSIGN_WINDOWS\x10\x04\x1a#\xa2\xb4\xfa\xc2\x05\x1d\x62\x65\x61m:transform:window_into:v1\x12\x33\n\x0bTEST_STREAM\x10\x05\x1a\"\xa2\xb4\xfa\xc2\x05\x1c\x62\x65\x61m:transform:teststream:v1\x12\x34\n\x0bMAP_WINDOWS\x10\x06\x1a#\xa2\xb4\xfa\xc2\x05\x1d\x62\x65\x61m:transform:map_windows:v1\x12\x38\n\rMERGE_WINDOWS\x10\x07\x1a%\xa2\xb4\xfa\xc2\x05\x1f\x62\x65\x61m:transform:merge_windows:v1\x12\x30\n\tTO_STRING\x10\x08\x1a!\xa2\xb4\xfa\xc2\x05\x1b\x62\x65\x61m:transform:to_string:v1\"t\n\x14\x44\x65precatedPrimitives\x12&\n\x04READ\x10\x00\x1a\x1c\xa2\xb4\xfa\xc2\x05\x16\x62\x65\x61m:transform:read:v1\x12\x34\n\x0b\x43REATE_VIEW\x10\x01\x1a#\xa2\xb4\xfa\xc2\x05\x1d\x62\x65\x61m:transform:create_view:v1\"\x97\x05\n\nComposites\x12<\n\x0f\x43OMBINE_PER_KEY\x10\x00\x1a\'\xa2\xb4\xfa\xc2\x05!beam:transform:combine_per_key:v1\x12>\n\x10\x43OMBINE_GLOBALLY\x10\x01\x1a(\xa2\xb4\xfa\xc2\x05\"beam:transform:combine_globally:v1\x12\x30\n\tRESHUFFLE\x10\x02\x1a!\xa2\xb4\xfa\xc2\x05\x1b\x62\x65\x61m:transform:reshuffle:v1\x12\x34\n\x0bWRITE_FILES\x10\x03\x1a#\xa2\xb4\xfa\xc2\x05\x1d\x62\x65\x61m:transform:write_files:v1\x12\x34\n\x0bPUBSUB_READ\x10\x04\x1a#\xa2\xb4\xfa\xc2\x05\x1d\x62\x65\x61m:transform:pubsub_read:v1\x12\x36\n\x0cPUBSUB_WRITE\x10\x05\x1a$\xa2\xb4\xfa\xc2\x05\x1e\x62\x65\x61m:transform:pubsub_write:v1\x12\x39\n\x0fPUBSUB_WRITE_V2\x10\x07\x1a$\xa2\xb4\xfa\xc2\x05\x1e\x62\x65\x61m:transform:pubsub_write:v2\x12\x64\n#GROUP_INTO_BATCHES_WITH_SHARDED_KEY\x10\x06\x1a;\xa2\xb4\xfa\xc2\x05\x35\x62\x65\x61m:transform:group_into_batches_with_sharded_key:v1\x12\x44\n\x13REDISTRIBUTE_BY_KEY\x10\x08\x1a+\xa2\xb4\xfa\xc2\x05%beam:transform:redistribute_by_key:v1\x12N\n\x18REDISTRIBUTE_ARBITRARILY\x10\t\x1a\x30\xa2\xb4\xfa\xc2\x05*beam:transform:redistribute_arbitrarily:v1\"\xe3\x03\n\x11\x43ombineComponents\x12R\n\x1a\x43OMBINE_PER_KEY_PRECOMBINE\x10\x00\x1a\x32\xa2\xb4\xfa\xc2\x05,beam:transform:combine_per_key_precombine:v1\x12\x62\n\"COMBINE_PER_KEY_MERGE_ACCUMULATORS\x10\x01\x1a:\xa2\xb4\xfa\xc2\x05\x34\x62\x65\x61m:transform:combine_per_key_merge_accumulators:v1\x12\\\n\x1f\x43OMBINE_PER_KEY_EXTRACT_OUTPUTS\x10\x02\x1a\x37\xa2\xb4\xfa\xc2\x05\x31\x62\x65\x61m:transform:combine_per_key_extract_outputs:v1\x12J\n\x16\x43OMBINE_GROUPED_VALUES\x10\x03\x1a.\xa2\xb4\xfa\xc2\x05(beam:transform:combine_grouped_values:v1\x12l\n\'COMBINE_PER_KEY_CONVERT_TO_ACCUMULATORS\x10\x04\x1a?\xa2\xb4\xfa\xc2\x05\x39\x62\x65\x61m:transform:combine_per_key_convert_to_accumulators:v1\"\x8d\x03\n\x19SplittableParDoComponents\x12L\n\x15PAIR_WITH_RESTRICTION\x10\x00\x1a\x31\xa2\xb4\xfa\xc2\x05+beam:transform:sdf_pair_with_restriction:v1\x12X\n\x1bSPLIT_AND_SIZE_RESTRICTIONS\x10\x01\x1a\x37\xa2\xb4\xfa\xc2\x05\x31\x62\x65\x61m:transform:sdf_split_and_size_restrictions:v1\x12o\n\'PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS\x10\x02\x1a\x42\xa2\xb4\xfa\xc2\x05\n\x05\x64o_fn\x18\x01 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\x12T\n\x0bside_inputs\x18\x03 \x03(\x0b\x32?.org.apache.beam.model.pipeline.v1.ParDoPayload.SideInputsEntry\x12T\n\x0bstate_specs\x18\x04 \x03(\x0b\x32?.org.apache.beam.model.pipeline.v1.ParDoPayload.StateSpecsEntry\x12\x61\n\x12timer_family_specs\x18\t \x03(\x0b\x32\x45.org.apache.beam.model.pipeline.v1.ParDoPayload.TimerFamilySpecsEntry\x12\x1c\n\x14restriction_coder_id\x18\x07 \x01(\t\x12\x1d\n\x15requests_finalization\x18\x08 \x01(\x08\x12\"\n\x1arequires_time_sorted_input\x18\n \x01(\x08\x12\x1d\n\x15requires_stable_input\x18\x0b \x01(\x08\x12.\n&on_window_expiration_timer_family_spec\x18\x0c \x01(\t\x1a_\n\x0fSideInputsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12;\n\x05value\x18\x02 \x01(\x0b\x32,.org.apache.beam.model.pipeline.v1.SideInput:\x02\x38\x01\x1a_\n\x0fStateSpecsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12;\n\x05value\x18\x02 \x01(\x0b\x32,.org.apache.beam.model.pipeline.v1.StateSpec:\x02\x38\x01\x1ak\n\x15TimerFamilySpecsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x41\n\x05value\x18\x02 \x01(\x0b\x32\x32.org.apache.beam.model.pipeline.v1.TimerFamilySpec:\x02\x38\x01J\x04\x08\x06\x10\x07\"\xfa\x04\n\tStateSpec\x12]\n\x16read_modify_write_spec\x18\x01 \x01(\x0b\x32;.org.apache.beam.model.pipeline.v1.ReadModifyWriteStateSpecH\x00\x12\x43\n\x08\x62\x61g_spec\x18\x02 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.BagStateSpecH\x00\x12O\n\x0e\x63ombining_spec\x18\x03 \x01(\x0b\x32\x35.org.apache.beam.model.pipeline.v1.CombiningStateSpecH\x00\x12\x43\n\x08map_spec\x18\x04 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.MapStateSpecH\x00\x12\x43\n\x08set_spec\x18\x05 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.SetStateSpecH\x00\x12T\n\x11ordered_list_spec\x18\x06 \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.OrderedListStateSpecH\x00\x12M\n\rmultimap_spec\x18\x08 \x01(\x0b\x32\x34.org.apache.beam.model.pipeline.v1.MultimapStateSpecH\x00\x12\x41\n\x08protocol\x18\x07 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpecB\x06\n\x04spec\",\n\x18ReadModifyWriteStateSpec\x12\x10\n\x08\x63oder_id\x18\x01 \x01(\t\"(\n\x0c\x42\x61gStateSpec\x12\x18\n\x10\x65lement_coder_id\x18\x01 \x01(\t\"0\n\x14OrderedListStateSpec\x12\x18\n\x10\x65lement_coder_id\x18\x01 \x01(\t\"w\n\x12\x43ombiningStateSpec\x12\x1c\n\x14\x61\x63\x63umulator_coder_id\x18\x01 \x01(\t\x12\x43\n\ncombine_fn\x18\x02 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\"<\n\x0cMapStateSpec\x12\x14\n\x0ckey_coder_id\x18\x01 \x01(\t\x12\x16\n\x0evalue_coder_id\x18\x02 \x01(\t\"A\n\x11MultimapStateSpec\x12\x14\n\x0ckey_coder_id\x18\x01 \x01(\t\x12\x16\n\x0evalue_coder_id\x18\x02 \x01(\t\"(\n\x0cSetStateSpec\x12\x18\n\x10\x65lement_coder_id\x18\x01 \x01(\t\"y\n\x0fTimerFamilySpec\x12G\n\x0btime_domain\x18\x01 \x01(\x0e\x32\x32.org.apache.beam.model.pipeline.v1.TimeDomain.Enum\x12\x1d\n\x15timer_family_coder_id\x18\x02 \x01(\t\"@\n\tIsBounded\"3\n\x04\x45num\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\r\n\tUNBOUNDED\x10\x01\x12\x0b\n\x07\x42OUNDED\x10\x02\"\x95\x01\n\x0bReadPayload\x12?\n\x06source\x18\x01 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\x12\x45\n\nis_bounded\x18\x02 \x01(\x0e\x32\x31.org.apache.beam.model.pipeline.v1.IsBounded.Enum\"W\n\x11WindowIntoPayload\x12\x42\n\twindow_fn\x18\x01 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\"s\n\x0e\x43ombinePayload\x12\x43\n\ncombine_fn\x18\x01 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\x12\x1c\n\x14\x61\x63\x63umulator_coder_id\x18\x02 \x01(\t\"\xaf\x06\n\x11TestStreamPayload\x12\x10\n\x08\x63oder_id\x18\x01 \x01(\t\x12J\n\x06\x65vents\x18\x02 \x03(\x0b\x32:.org.apache.beam.model.pipeline.v1.TestStreamPayload.Event\x12I\n\x08\x65ndpoint\x18\x03 \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.ApiServiceDescriptor\x1a\xae\x04\n\x05\x45vent\x12\x66\n\x0fwatermark_event\x18\x01 \x01(\x0b\x32K.org.apache.beam.model.pipeline.v1.TestStreamPayload.Event.AdvanceWatermarkH\x00\x12q\n\x15processing_time_event\x18\x02 \x01(\x0b\x32P.org.apache.beam.model.pipeline.v1.TestStreamPayload.Event.AdvanceProcessingTimeH\x00\x12_\n\relement_event\x18\x03 \x01(\x0b\x32\x46.org.apache.beam.model.pipeline.v1.TestStreamPayload.Event.AddElementsH\x00\x1a\x36\n\x10\x41\x64vanceWatermark\x12\x15\n\rnew_watermark\x18\x01 \x01(\x03\x12\x0b\n\x03tag\x18\x02 \x01(\t\x1a\x31\n\x15\x41\x64vanceProcessingTime\x12\x18\n\x10\x61\x64vance_duration\x18\x01 \x01(\x03\x1au\n\x0b\x41\x64\x64\x45lements\x12Y\n\x08\x65lements\x18\x01 \x03(\x0b\x32G.org.apache.beam.model.pipeline.v1.TestStreamPayload.TimestampedElement\x12\x0b\n\x03tag\x18\x03 \x01(\tB\x07\n\x05\x65vent\x1a@\n\x12TimestampedElement\x12\x17\n\x0f\x65ncoded_element\x18\x01 \x01(\x0c\x12\x11\n\ttimestamp\x18\x02 \x01(\x03\"\xab\x03\n\x11WriteFilesPayload\x12=\n\x04sink\x18\x01 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\x12H\n\x0f\x66ormat_function\x18\x02 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\x12\x17\n\x0fwindowed_writes\x18\x03 \x01(\x08\x12\"\n\x1arunner_determined_sharding\x18\x04 \x01(\x08\x12Y\n\x0bside_inputs\x18\x05 \x03(\x0b\x32\x44.org.apache.beam.model.pipeline.v1.WriteFilesPayload.SideInputsEntry\x12\x14\n\x0c\x61uto_sharded\x18\x06 \x01(\x08\x1a_\n\x0fSideInputsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12;\n\x05value\x18\x02 \x01(\x0b\x32,.org.apache.beam.model.pipeline.v1.SideInput:\x02\x38\x01\"\xcf\x01\n\x11PubSubReadPayload\x12\r\n\x05topic\x18\x01 \x01(\t\x12\x14\n\x0csubscription\x18\x02 \x01(\t\x12\x1b\n\x13timestamp_attribute\x18\x03 \x01(\t\x12\x14\n\x0cid_attribute\x18\x04 \x01(\t\x12\x17\n\x0fwith_attributes\x18\x05 \x01(\x08\x12 \n\x18topic_runtime_overridden\x18\x06 \x01(\t\x12\'\n\x1fsubscription_runtime_overridden\x18\x07 \x01(\t\"x\n\x12PubSubWritePayload\x12\r\n\x05topic\x18\x01 \x01(\t\x12\x1b\n\x13timestamp_attribute\x18\x02 \x01(\t\x12\x14\n\x0cid_attribute\x18\x03 \x01(\t\x12 \n\x18topic_runtime_overridden\x18\x04 \x01(\t\"n\n\x17GroupIntoBatchesPayload\x12\x12\n\nbatch_size\x18\x01 \x01(\x03\x12\x18\n\x10\x62\x61tch_size_bytes\x18\x03 \x01(\x03\x12%\n\x1dmax_buffering_duration_millis\x18\x02 \x01(\x03\"/\n\x13RedistributePayload\x12\x18\n\x10\x61llow_duplicates\x18\x01 \x01(\x08\"c\n\x05\x43oder\x12=\n\x04spec\x18\x01 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\x12\x1b\n\x13\x63omponent_coder_ids\x18\x02 \x03(\t\"\x9d\x07\n\x0eStandardCoders\"\x8a\x07\n\x04\x45num\x12$\n\x05\x42YTES\x10\x00\x1a\x19\xa2\xb4\xfa\xc2\x05\x13\x62\x65\x61m:coder:bytes:v1\x12\x30\n\x0bSTRING_UTF8\x10\n\x1a\x1f\xa2\xb4\xfa\xc2\x05\x19\x62\x65\x61m:coder:string_utf8:v1\x12\x1e\n\x02KV\x10\x01\x1a\x16\xa2\xb4\xfa\xc2\x05\x10\x62\x65\x61m:coder:kv:v1\x12\"\n\x04\x42OOL\x10\x0c\x1a\x18\xa2\xb4\xfa\xc2\x05\x12\x62\x65\x61m:coder:bool:v1\x12&\n\x06VARINT\x10\x02\x1a\x1a\xa2\xb4\xfa\xc2\x05\x14\x62\x65\x61m:coder:varint:v1\x12&\n\x06\x44OUBLE\x10\x0b\x1a\x1a\xa2\xb4\xfa\xc2\x05\x14\x62\x65\x61m:coder:double:v1\x12*\n\x08ITERABLE\x10\x03\x1a\x1c\xa2\xb4\xfa\xc2\x05\x16\x62\x65\x61m:coder:iterable:v1\x12$\n\x05TIMER\x10\x04\x1a\x19\xa2\xb4\xfa\xc2\x05\x13\x62\x65\x61m:coder:timer:v1\x12\x38\n\x0fINTERVAL_WINDOW\x10\x05\x1a#\xa2\xb4\xfa\xc2\x05\x1d\x62\x65\x61m:coder:interval_window:v1\x12\x34\n\rLENGTH_PREFIX\x10\x06\x1a!\xa2\xb4\xfa\xc2\x05\x1b\x62\x65\x61m:coder:length_prefix:v1\x12\x34\n\rGLOBAL_WINDOW\x10\x07\x1a!\xa2\xb4\xfa\xc2\x05\x1b\x62\x65\x61m:coder:global_window:v1\x12\x36\n\x0eWINDOWED_VALUE\x10\x08\x1a\"\xa2\xb4\xfa\xc2\x05\x1c\x62\x65\x61m:coder:windowed_value:v1\x12\x42\n\x14PARAM_WINDOWED_VALUE\x10\x0e\x1a(\xa2\xb4\xfa\xc2\x05\"beam:coder:param_windowed_value:v1\x12\x44\n\x15STATE_BACKED_ITERABLE\x10\t\x1a)\xa2\xb4\xfa\xc2\x05#beam:coder:state_backed_iterable:v1\x12\x34\n\rCUSTOM_WINDOW\x10\x10\x1a!\xa2\xb4\xfa\xc2\x05\x1b\x62\x65\x61m:coder:custom_window:v1\x12 \n\x03ROW\x10\r\x1a\x17\xa2\xb4\xfa\xc2\x05\x11\x62\x65\x61m:coder:row:v1\x12&\n\x06SCHEMA\x10\x12\x1a\x1a\xa2\xb4\xfa\xc2\x05\x14\x62\x65\x61m:coder:schema:v1\x12\x30\n\x0bSHARDED_KEY\x10\x0f\x1a\x1f\xa2\xb4\xfa\xc2\x05\x19\x62\x65\x61m:coder:sharded_key:v1\x12*\n\x08NULLABLE\x10\x11\x1a\x1c\xa2\xb4\xfa\xc2\x05\x16\x62\x65\x61m:coder:nullable:v1\"\x8c\x05\n\x11WindowingStrategy\x12\x42\n\twindow_fn\x18\x01 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\x12I\n\x0cmerge_status\x18\x02 \x01(\x0e\x32\x33.org.apache.beam.model.pipeline.v1.MergeStatus.Enum\x12\x17\n\x0fwindow_coder_id\x18\x03 \x01(\t\x12;\n\x07trigger\x18\x04 \x01(\x0b\x32*.org.apache.beam.model.pipeline.v1.Trigger\x12S\n\x11\x61\x63\x63umulation_mode\x18\x05 \x01(\x0e\x32\x38.org.apache.beam.model.pipeline.v1.AccumulationMode.Enum\x12G\n\x0boutput_time\x18\x06 \x01(\x0e\x32\x32.org.apache.beam.model.pipeline.v1.OutputTime.Enum\x12Q\n\x10\x63losing_behavior\x18\x07 \x01(\x0e\x32\x37.org.apache.beam.model.pipeline.v1.ClosingBehavior.Enum\x12\x18\n\x10\x61llowed_lateness\x18\x08 \x01(\x03\x12P\n\x10on_time_behavior\x18\t \x01(\x0e\x32\x36.org.apache.beam.model.pipeline.v1.OnTimeBehavior.Enum\x12\x1d\n\x15\x61ssigns_to_one_window\x18\n \x01(\x08\x12\x16\n\x0e\x65nvironment_id\x18\x0b \x01(\t\"\\\n\x0bMergeStatus\"M\n\x04\x45num\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x0f\n\x0bNON_MERGING\x10\x01\x12\x0f\n\x0bNEEDS_MERGE\x10\x02\x12\x12\n\x0e\x41LREADY_MERGED\x10\x03\"]\n\x10\x41\x63\x63umulationMode\"I\n\x04\x45num\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x0e\n\nDISCARDING\x10\x01\x12\x10\n\x0c\x41\x43\x43UMULATING\x10\x02\x12\x0e\n\nRETRACTING\x10\x03\"Q\n\x0f\x43losingBehavior\">\n\x04\x45num\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x0f\n\x0b\x45MIT_ALWAYS\x10\x01\x12\x14\n\x10\x45MIT_IF_NONEMPTY\x10\x02\"P\n\x0eOnTimeBehavior\">\n\x04\x45num\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x0f\n\x0b\x46IRE_ALWAYS\x10\x01\x12\x14\n\x10\x46IRE_IF_NONEMPTY\x10\x02\"b\n\nOutputTime\"T\n\x04\x45num\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x11\n\rEND_OF_WINDOW\x10\x01\x12\x12\n\x0eLATEST_IN_PANE\x10\x02\x12\x14\n\x10\x45\x41RLIEST_IN_PANE\x10\x03\"n\n\nTimeDomain\"`\n\x04\x45num\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x0e\n\nEVENT_TIME\x10\x01\x12\x13\n\x0fPROCESSING_TIME\x10\x02\"\x04\x08\x03\x10\x03*\x1cSYNCHRONIZED_PROCESSING_TIME\"\x82\x0e\n\x07Trigger\x12H\n\tafter_all\x18\x01 \x01(\x0b\x32\x33.org.apache.beam.model.pipeline.v1.Trigger.AfterAllH\x00\x12H\n\tafter_any\x18\x02 \x01(\x0b\x32\x33.org.apache.beam.model.pipeline.v1.Trigger.AfterAnyH\x00\x12J\n\nafter_each\x18\x03 \x01(\x0b\x32\x34.org.apache.beam.model.pipeline.v1.Trigger.AfterEachH\x00\x12Z\n\x13\x61\x66ter_end_of_window\x18\x04 \x01(\x0b\x32;.org.apache.beam.model.pipeline.v1.Trigger.AfterEndOfWindowH\x00\x12_\n\x15\x61\x66ter_processing_time\x18\x05 \x01(\x0b\x32>.org.apache.beam.model.pipeline.v1.Trigger.AfterProcessingTimeH\x00\x12x\n\"after_synchronized_processing_time\x18\x06 \x01(\x0b\x32J.org.apache.beam.model.pipeline.v1.Trigger.AfterSynchronizedProcessingTimeH\x00\x12\x43\n\x06\x61lways\x18\x0c \x01(\x0b\x32\x31.org.apache.beam.model.pipeline.v1.Trigger.AlwaysH\x00\x12\x45\n\x07\x64\x65\x66\x61ult\x18\x07 \x01(\x0b\x32\x32.org.apache.beam.model.pipeline.v1.Trigger.DefaultH\x00\x12P\n\relement_count\x18\x08 \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.Trigger.ElementCountH\x00\x12\x41\n\x05never\x18\t \x01(\x0b\x32\x30.org.apache.beam.model.pipeline.v1.Trigger.NeverH\x00\x12J\n\nor_finally\x18\n \x01(\x0b\x32\x34.org.apache.beam.model.pipeline.v1.Trigger.OrFinallyH\x00\x12\x43\n\x06repeat\x18\x0b \x01(\x0b\x32\x31.org.apache.beam.model.pipeline.v1.Trigger.RepeatH\x00\x1aK\n\x08\x41\x66terAll\x12?\n\x0bsubtriggers\x18\x01 \x03(\x0b\x32*.org.apache.beam.model.pipeline.v1.Trigger\x1aK\n\x08\x41\x66terAny\x12?\n\x0bsubtriggers\x18\x01 \x03(\x0b\x32*.org.apache.beam.model.pipeline.v1.Trigger\x1aL\n\tAfterEach\x12?\n\x0bsubtriggers\x18\x01 \x03(\x0b\x32*.org.apache.beam.model.pipeline.v1.Trigger\x1a\x97\x01\n\x10\x41\x66terEndOfWindow\x12\x41\n\rearly_firings\x18\x01 \x01(\x0b\x32*.org.apache.beam.model.pipeline.v1.Trigger\x12@\n\x0clate_firings\x18\x02 \x01(\x0b\x32*.org.apache.beam.model.pipeline.v1.Trigger\x1aj\n\x13\x41\x66terProcessingTime\x12S\n\x14timestamp_transforms\x18\x01 \x03(\x0b\x32\x35.org.apache.beam.model.pipeline.v1.TimestampTransform\x1a!\n\x1f\x41\x66terSynchronizedProcessingTime\x1a\t\n\x07\x44\x65\x66\x61ult\x1a%\n\x0c\x45lementCount\x12\x15\n\relement_count\x18\x01 \x01(\x05\x1a\x07\n\x05Never\x1a\x08\n\x06\x41lways\x1a\x82\x01\n\tOrFinally\x12\x38\n\x04main\x18\x01 \x01(\x0b\x32*.org.apache.beam.model.pipeline.v1.Trigger\x12;\n\x07\x66inally\x18\x02 \x01(\x0b\x32*.org.apache.beam.model.pipeline.v1.Trigger\x1aH\n\x06Repeat\x12>\n\nsubtrigger\x18\x01 \x01(\x0b\x32*.org.apache.beam.model.pipeline.v1.TriggerB\t\n\x07trigger\"\x96\x02\n\x12TimestampTransform\x12L\n\x05\x64\x65lay\x18\x01 \x01(\x0b\x32;.org.apache.beam.model.pipeline.v1.TimestampTransform.DelayH\x00\x12Q\n\x08\x61lign_to\x18\x02 \x01(\x0b\x32=.org.apache.beam.model.pipeline.v1.TimestampTransform.AlignToH\x00\x1a\x1d\n\x05\x44\x65lay\x12\x14\n\x0c\x64\x65lay_millis\x18\x01 \x01(\x03\x1a)\n\x07\x41lignTo\x12\x0e\n\x06period\x18\x03 \x01(\x03\x12\x0e\n\x06offset\x18\x04 \x01(\x03\x42\x15\n\x13timestamp_transform\"\xe2\x01\n\tSideInput\x12G\n\x0e\x61\x63\x63\x65ss_pattern\x18\x01 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\x12@\n\x07view_fn\x18\x02 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\x12J\n\x11window_mapping_fn\x18\x03 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\"\x89\x04\n\x11StandardArtifacts\"\x9f\x02\n\x05Types\x12*\n\x04\x46ILE\x10\x00\x1a \xa2\xb4\xfa\xc2\x05\x1a\x62\x65\x61m:artifact:type:file:v1\x12(\n\x03URL\x10\x01\x1a\x1f\xa2\xb4\xfa\xc2\x05\x19\x62\x65\x61m:artifact:type:url:v1\x12\x32\n\x08\x45MBEDDED\x10\x02\x1a$\xa2\xb4\xfa\xc2\x05\x1e\x62\x65\x61m:artifact:type:embedded:v1\x12*\n\x04PYPI\x10\x03\x1a \xa2\xb4\xfa\xc2\x05\x1a\x62\x65\x61m:artifact:type:pypi:v1\x12,\n\x05MAVEN\x10\x04\x1a!\xa2\xb4\xfa\xc2\x05\x1b\x62\x65\x61m:artifact:type:maven:v1\x12\x32\n\x08\x44\x45\x46\x45RRED\x10\x05\x1a$\xa2\xb4\xfa\xc2\x05\x1e\x62\x65\x61m:artifact:type:deferred:v1\"\xd1\x01\n\x05Roles\x12\x36\n\nSTAGING_TO\x10\x00\x1a&\xa2\xb4\xfa\xc2\x05 beam:artifact:role:staging_to:v1\x12L\n\x15PIP_REQUIREMENTS_FILE\x10\x01\x1a\x31\xa2\xb4\xfa\xc2\x05+beam:artifact:role:pip_requirements_file:v1\x12\x42\n\x10GO_WORKER_BINARY\x10\x02\x1a,\xa2\xb4\xfa\xc2\x05&beam:artifact:role:go_worker_binary:v1\"3\n\x13\x41rtifactFilePayload\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x0e\n\x06sha256\x18\x02 \x01(\t\"1\n\x12\x41rtifactUrlPayload\x12\x0b\n\x03url\x18\x01 \x01(\t\x12\x0e\n\x06sha256\x18\x02 \x01(\t\"#\n\x13\x45mbeddedFilePayload\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\"3\n\x0bPyPIPayload\x12\x13\n\x0b\x61rtifact_id\x18\x01 \x01(\t\x12\x0f\n\x07version\x18\x02 \x01(\t\"8\n\x0cMavenPayload\x12\x10\n\x08\x61rtifact\x18\x01 \x01(\t\x12\x16\n\x0erepository_url\x18\x02 \x01(\t\"4\n\x17\x44\x65\x66\x65rredArtifactPayload\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"3\n\x1c\x41rtifactStagingToRolePayload\x12\x13\n\x0bstaged_name\x18\x01 \x01(\t\"e\n\x13\x41rtifactInformation\x12\x10\n\x08type_urn\x18\x01 \x01(\t\x12\x14\n\x0ctype_payload\x18\x02 \x01(\x0c\x12\x10\n\x08role_urn\x18\x03 \x01(\t\x12\x14\n\x0crole_payload\x18\x04 \x01(\x0c\"\xec\x02\n\x0b\x45nvironment\x12\x0b\n\x03urn\x18\x02 \x01(\t\x12\x0f\n\x07payload\x18\x03 \x01(\x0c\x12\x44\n\x0c\x64isplay_data\x18\x04 \x03(\x0b\x32..org.apache.beam.model.pipeline.v1.DisplayData\x12\x14\n\x0c\x63\x61pabilities\x18\x05 \x03(\t\x12L\n\x0c\x64\x65pendencies\x18\x06 \x03(\x0b\x32\x36.org.apache.beam.model.pipeline.v1.ArtifactInformation\x12Y\n\x0eresource_hints\x18\x07 \x03(\x0b\x32\x41.org.apache.beam.model.pipeline.v1.Environment.ResourceHintsEntry\x1a\x34\n\x12ResourceHintsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x0c:\x02\x38\x01J\x04\x08\x01\x10\x02\"\xeb\x01\n\x14StandardEnvironments\"\xd2\x01\n\x0c\x45nvironments\x12$\n\x06\x44OCKER\x10\x00\x1a\x18\xa2\xb4\xfa\xc2\x05\x12\x62\x65\x61m:env:docker:v1\x12&\n\x07PROCESS\x10\x01\x1a\x19\xa2\xb4\xfa\xc2\x05\x13\x62\x65\x61m:env:process:v1\x12(\n\x08\x45XTERNAL\x10\x02\x1a\x1a\xa2\xb4\xfa\xc2\x05\x14\x62\x65\x61m:env:external:v1\x12&\n\x07\x44\x45\x46\x41ULT\x10\x03\x1a\x19\xa2\xb4\xfa\xc2\x05\x13\x62\x65\x61m:env:default:v1\x12\"\n\x05\x41NYOF\x10\x04\x1a\x17\xa2\xb4\xfa\xc2\x05\x11\x62\x65\x61m:env:anyof:v1\"(\n\rDockerPayload\x12\x17\n\x0f\x63ontainer_image\x18\x01 \x01(\t\"\xb0\x01\n\x0eProcessPayload\x12\n\n\x02os\x18\x01 \x01(\t\x12\x0c\n\x04\x61rch\x18\x02 \x01(\t\x12\x0f\n\x07\x63ommand\x18\x03 \x01(\t\x12G\n\x03\x65nv\x18\x04 \x03(\x0b\x32:.org.apache.beam.model.pipeline.v1.ProcessPayload.EnvEntry\x1a*\n\x08\x45nvEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\xdb\x01\n\x0f\x45xternalPayload\x12I\n\x08\x65ndpoint\x18\x01 \x01(\x0b\x32\x37.org.apache.beam.model.pipeline.v1.ApiServiceDescriptor\x12N\n\x06params\x18\x02 \x03(\x0b\x32>.org.apache.beam.model.pipeline.v1.ExternalPayload.ParamsEntry\x1a-\n\x0bParamsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"_\n\x17\x41nyOfEnvironmentPayload\x12\x44\n\x0c\x65nvironments\x18\x01 \x03(\x0b\x32..org.apache.beam.model.pipeline.v1.Environment\"\xac\x07\n\x11StandardProtocols\"\x96\x07\n\x04\x45num\x12H\n\x19LEGACY_PROGRESS_REPORTING\x10\x00\x1a)\xa2\xb4\xfa\xc2\x05#beam:protocol:progress_reporting:v0\x12\x41\n\x12PROGRESS_REPORTING\x10\x01\x1a)\xa2\xb4\xfa\xc2\x05#beam:protocol:progress_reporting:v1\x12\x37\n\rWORKER_STATUS\x10\x02\x1a$\xa2\xb4\xfa\xc2\x05\x1e\x62\x65\x61m:protocol:worker_status:v1\x12U\n\x1cMULTI_CORE_BUNDLE_PROCESSING\x10\x03\x1a\x33\xa2\xb4\xfa\xc2\x05-beam:protocol:multi_core_bundle_processing:v1\x12;\n\x0fSIBLING_WORKERS\x10\x05\x1a&\xa2\xb4\xfa\xc2\x05 beam:protocol:sibling_workers:v1\x12M\n\x18HARNESS_MONITORING_INFOS\x10\x04\x1a/\xa2\xb4\xfa\xc2\x05)beam:protocol:harness_monitoring_infos:v1\x12\x61\n\"CONTROL_REQUEST_ELEMENTS_EMBEDDING\x10\x06\x1a\x39\xa2\xb4\xfa\xc2\x05\x33\x62\x65\x61m:protocol:control_request_elements_embedding:v1\x12\x37\n\rSTATE_CACHING\x10\x07\x1a$\xa2\xb4\xfa\xc2\x05\x1e\x62\x65\x61m:protocol:state_caching:v1\x12\x37\n\rDATA_SAMPLING\x10\x08\x1a$\xa2\xb4\xfa\xc2\x05\x1e\x62\x65\x61m:protocol:data_sampling:v1\x12S\n\x1bSDK_CONSUMING_RECEIVED_DATA\x10\t\x1a\x32\xa2\xb4\xfa\xc2\x05,beam:protocol:sdk_consuming_received_data:v1\x12\x41\n\x12ORDERED_LIST_STATE\x10\n\x1a)\xa2\xb4\xfa\xc2\x05#beam:protocol:ordered_list_state:v1\x12=\n\x10\x45LEMENT_METADATA\x10\x0b\x1a\'\xa2\xb4\xfa\xc2\x05!beam:protocol:element_metadata:v1\x12\x39\n\x0eMULTIMAP_STATE\x10\x0c\x1a%\xa2\xb4\xfa\xc2\x05\x1f\x62\x65\x61m:protocol:multimap_state:v1\"\xb5\x02\n\x17StandardRunnerProtocols\"\x99\x02\n\x04\x45num\x12O\n\x19MONITORING_INFO_SHORT_IDS\x10\x00\x1a\x30\xa2\xb4\xfa\xc2\x05*beam:protocol:monitoring_info_short_ids:v1\x12\x63\n#CONTROL_RESPONSE_ELEMENTS_EMBEDDING\x10\x06\x1a:\xa2\xb4\xfa\xc2\x05\x34\x62\x65\x61m:protocol:control_response_elements_embedding:v1\x12[\n\x1fMULTIMAP_KEYS_VALUES_SIDE_INPUT\x10\x07\x1a\x36\xa2\xb4\xfa\xc2\x05\x30\x62\x65\x61m:protocol:multimap_keys_values_side_input:v1\"\xff\x03\n\x14StandardRequirements\"\xe6\x03\n\x04\x45num\x12J\n\x1cREQUIRES_STATEFUL_PROCESSING\x10\x00\x1a(\xa2\xb4\xfa\xc2\x05\"beam:requirement:pardo:stateful:v1\x12N\n\x1cREQUIRES_BUNDLE_FINALIZATION\x10\x01\x1a,\xa2\xb4\xfa\xc2\x05&beam:requirement:pardo:finalization:v1\x12G\n\x15REQUIRES_STABLE_INPUT\x10\x02\x1a,\xa2\xb4\xfa\xc2\x05&beam:requirement:pardo:stable_input:v1\x12Q\n\x1aREQUIRES_TIME_SORTED_INPUT\x10\x03\x1a\x31\xa2\xb4\xfa\xc2\x05+beam:requirement:pardo:time_sorted_input:v1\x12M\n\x18REQUIRES_SPLITTABLE_DOFN\x10\x04\x1a/\xa2\xb4\xfa\xc2\x05)beam:requirement:pardo:splittable_dofn:v1\x12W\n\x1dREQUIRES_ON_WINDOW_EXPIRATION\x10\x05\x1a\x34\xa2\xb4\xfa\xc2\x05.beam:requirement:pardo:on_window_expiration:v1\",\n\x0c\x46unctionSpec\x12\x0b\n\x03urn\x18\x01 \x01(\t\x12\x0f\n\x07payload\x18\x03 \x01(\x0c\"W\n\x13StandardDisplayData\"@\n\x0b\x44isplayData\x12\x31\n\x08LABELLED\x10\x00\x1a#\xa2\xb4\xfa\xc2\x05\x1d\x62\x65\x61m:display_data:labelled:v1\"\xa4\x01\n\x0fLabelledPayload\x12\r\n\x05label\x18\x01 \x01(\t\x12\x16\n\x0cstring_value\x18\x02 \x01(\tH\x00\x12\x14\n\nbool_value\x18\x03 \x01(\x08H\x00\x12\x16\n\x0c\x64ouble_value\x18\x04 \x01(\x01H\x00\x12\x13\n\tint_value\x18\x05 \x01(\x03H\x00\x12\x0b\n\x03key\x18\x06 \x01(\t\x12\x11\n\tnamespace\x18\x07 \x01(\tB\x07\n\x05value\"+\n\x0b\x44isplayData\x12\x0b\n\x03urn\x18\x01 \x01(\t\x12\x0f\n\x07payload\x18\x02 \x01(\x0c\"\xc1\x06\n\x15MessageWithComponents\x12\x41\n\ncomponents\x18\x01 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.Components\x12\x39\n\x05\x63oder\x18\x02 \x01(\x0b\x32(.org.apache.beam.model.pipeline.v1.CoderH\x00\x12L\n\x0f\x63ombine_payload\x18\x03 \x01(\x0b\x32\x31.org.apache.beam.model.pipeline.v1.CombinePayloadH\x00\x12H\n\rfunction_spec\x18\x04 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpecH\x00\x12I\n\x0epar_do_payload\x18\x06 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.ParDoPayloadH\x00\x12\x43\n\nptransform\x18\x07 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.PTransformH\x00\x12\x45\n\x0bpcollection\x18\x08 \x01(\x0b\x32..org.apache.beam.model.pipeline.v1.PCollectionH\x00\x12\x46\n\x0cread_payload\x18\t \x01(\x0b\x32..org.apache.beam.model.pipeline.v1.ReadPayloadH\x00\x12\x42\n\nside_input\x18\x0b \x01(\x0b\x32,.org.apache.beam.model.pipeline.v1.SideInputH\x00\x12S\n\x13window_into_payload\x18\x0c \x01(\x0b\x32\x34.org.apache.beam.model.pipeline.v1.WindowIntoPayloadH\x00\x12R\n\x12windowing_strategy\x18\r \x01(\x0b\x32\x34.org.apache.beam.model.pipeline.v1.WindowingStrategyH\x00\x42\x06\n\x04root\"\xb9\x08\n\x16\x45xecutableStagePayload\x12\x43\n\x0b\x65nvironment\x18\x01 \x01(\x0b\x32..org.apache.beam.model.pipeline.v1.Environment\x12g\n\x13wire_coder_settings\x18\t \x03(\x0b\x32J.org.apache.beam.model.pipeline.v1.ExecutableStagePayload.WireCoderSetting\x12\r\n\x05input\x18\x02 \x01(\t\x12Z\n\x0bside_inputs\x18\x03 \x03(\x0b\x32\x45.org.apache.beam.model.pipeline.v1.ExecutableStagePayload.SideInputId\x12\x12\n\ntransforms\x18\x04 \x03(\t\x12\x0f\n\x07outputs\x18\x05 \x03(\t\x12\x41\n\ncomponents\x18\x06 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.Components\x12Z\n\x0buser_states\x18\x07 \x03(\x0b\x32\x45.org.apache.beam.model.pipeline.v1.ExecutableStagePayload.UserStateId\x12Q\n\x06timers\x18\x08 \x03(\x0b\x32\x41.org.apache.beam.model.pipeline.v1.ExecutableStagePayload.TimerId\x12^\n\rtimerFamilies\x18\n \x03(\x0b\x32G.org.apache.beam.model.pipeline.v1.ExecutableStagePayload.TimerFamilyId\x1a\x37\n\x0bSideInputId\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x12\n\nlocal_name\x18\x02 \x01(\t\x1a\x37\n\x0bUserStateId\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x12\n\nlocal_name\x18\x02 \x01(\t\x1a\x33\n\x07TimerId\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x12\n\nlocal_name\x18\x02 \x01(\t\x1a\x39\n\rTimerFamilyId\x12\x14\n\x0ctransform_id\x18\x01 \x01(\t\x12\x12\n\nlocal_name\x18\x02 \x01(\t\x1a\xac\x01\n\x10WireCoderSetting\x12\x0b\n\x03urn\x18\x01 \x01(\t\x12\x0f\n\x07payload\x18\x02 \x01(\x0c\x12\x1c\n\x12input_or_output_id\x18\x03 \x01(\tH\x00\x12R\n\x05timer\x18\x04 \x01(\x0b\x32\x41.org.apache.beam.model.pipeline.v1.ExecutableStagePayload.TimerIdH\x00\x42\x08\n\x06target\"\x9c\x02\n\x15StandardResourceHints\"\x82\x02\n\x04\x45num\x12\x34\n\x0b\x41\x43\x43\x45LERATOR\x10\x00\x1a#\xa2\xb4\xfa\xc2\x05\x1d\x62\x65\x61m:resources:accelerator:v1\x12\x38\n\rMIN_RAM_BYTES\x10\x01\x1a%\xa2\xb4\xfa\xc2\x05\x1f\x62\x65\x61m:resources:min_ram_bytes:v1\x12\x30\n\tCPU_COUNT\x10\x02\x1a!\xa2\xb4\xfa\xc2\x05\x1b\x62\x65\x61m:resources:cpu_count:v1\x12X\n\x1dMAX_ACTIVE_BUNDLES_PER_WORKER\x10\x03\x1a\x35\xa2\xb4\xfa\xc2\x05/beam:resources:max_active_bundles_per_worker:v1:6\n\x08\x62\x65\x61m_urn\x12!.google.protobuf.EnumValueOptions\x18\xc4\xa6\xafX \x01(\t:;\n\rbeam_constant\x12!.google.protobuf.EnumValueOptions\x18\xc5\xa6\xafX \x01(\tBx\n!org.apache.beam.model.pipeline.v1B\tRunnerApiZHgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1;pipeline_v1b\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'org.apache.beam.model.pipeline.v1.beam_runner_api_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n!org.apache.beam.model.pipeline.v1B\tRunnerApiZHgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1;pipeline_v1' _globals['_BEAMCONSTANTS_CONSTANTS'].values_by_name["MIN_TIMESTAMP_MILLIS"]._options = None _globals['_BEAMCONSTANTS_CONSTANTS'].values_by_name["MIN_TIMESTAMP_MILLIS"]._serialized_options = b'\252\264\372\302\005\021-9223372036854775' _globals['_BEAMCONSTANTS_CONSTANTS'].values_by_name["MAX_TIMESTAMP_MILLIS"]._options = None _globals['_BEAMCONSTANTS_CONSTANTS'].values_by_name["MAX_TIMESTAMP_MILLIS"]._serialized_options = b'\252\264\372\302\005\0209223372036854775' _globals['_BEAMCONSTANTS_CONSTANTS'].values_by_name["GLOBAL_WINDOW_MAX_TIMESTAMP_MILLIS"]._options = None _globals['_BEAMCONSTANTS_CONSTANTS'].values_by_name["GLOBAL_WINDOW_MAX_TIMESTAMP_MILLIS"]._serialized_options = b'\252\264\372\302\005\0209223371950454775' _globals['_COMPONENTS_TRANSFORMSENTRY']._options = None _globals['_COMPONENTS_TRANSFORMSENTRY']._serialized_options = b'8\001' _globals['_COMPONENTS_PCOLLECTIONSENTRY']._options = None _globals['_COMPONENTS_PCOLLECTIONSENTRY']._serialized_options = b'8\001' _globals['_COMPONENTS_WINDOWINGSTRATEGIESENTRY']._options = None _globals['_COMPONENTS_WINDOWINGSTRATEGIESENTRY']._serialized_options = b'8\001' _globals['_COMPONENTS_CODERSENTRY']._options = None _globals['_COMPONENTS_CODERSENTRY']._serialized_options = b'8\001' _globals['_COMPONENTS_ENVIRONMENTSENTRY']._options = None _globals['_COMPONENTS_ENVIRONMENTSENTRY']._serialized_options = b'8\001' _globals['_PTRANSFORM_INPUTSENTRY']._options = None _globals['_PTRANSFORM_INPUTSENTRY']._serialized_options = b'8\001' _globals['_PTRANSFORM_OUTPUTSENTRY']._options = None _globals['_PTRANSFORM_OUTPUTSENTRY']._serialized_options = b'8\001' _globals['_PTRANSFORM_ANNOTATIONSENTRY']._options = None _globals['_PTRANSFORM_ANNOTATIONSENTRY']._serialized_options = b'8\001' _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["PAR_DO"]._options = None _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["PAR_DO"]._serialized_options = b'\242\264\372\302\005\027beam:transform:pardo:v1' _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["FLATTEN"]._options = None _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["FLATTEN"]._serialized_options = b'\242\264\372\302\005\031beam:transform:flatten:v1' _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["GROUP_BY_KEY"]._options = None _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["GROUP_BY_KEY"]._serialized_options = b'\242\264\372\302\005\036beam:transform:group_by_key:v1' _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["IMPULSE"]._options = None _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["IMPULSE"]._serialized_options = b'\242\264\372\302\005\031beam:transform:impulse:v1' _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["ASSIGN_WINDOWS"]._options = None _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["ASSIGN_WINDOWS"]._serialized_options = b'\242\264\372\302\005\035beam:transform:window_into:v1' _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["TEST_STREAM"]._options = None _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["TEST_STREAM"]._serialized_options = b'\242\264\372\302\005\034beam:transform:teststream:v1' _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["MAP_WINDOWS"]._options = None _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["MAP_WINDOWS"]._serialized_options = b'\242\264\372\302\005\035beam:transform:map_windows:v1' _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["MERGE_WINDOWS"]._options = None _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["MERGE_WINDOWS"]._serialized_options = b'\242\264\372\302\005\037beam:transform:merge_windows:v1' _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["TO_STRING"]._options = None _globals['_STANDARDPTRANSFORMS_PRIMITIVES'].values_by_name["TO_STRING"]._serialized_options = b'\242\264\372\302\005\033beam:transform:to_string:v1' _globals['_STANDARDPTRANSFORMS_DEPRECATEDPRIMITIVES'].values_by_name["READ"]._options = None _globals['_STANDARDPTRANSFORMS_DEPRECATEDPRIMITIVES'].values_by_name["READ"]._serialized_options = b'\242\264\372\302\005\026beam:transform:read:v1' _globals['_STANDARDPTRANSFORMS_DEPRECATEDPRIMITIVES'].values_by_name["CREATE_VIEW"]._options = None _globals['_STANDARDPTRANSFORMS_DEPRECATEDPRIMITIVES'].values_by_name["CREATE_VIEW"]._serialized_options = b'\242\264\372\302\005\035beam:transform:create_view:v1' _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["COMBINE_PER_KEY"]._options = None _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["COMBINE_PER_KEY"]._serialized_options = b'\242\264\372\302\005!beam:transform:combine_per_key:v1' _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["COMBINE_GLOBALLY"]._options = None _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["COMBINE_GLOBALLY"]._serialized_options = b'\242\264\372\302\005\"beam:transform:combine_globally:v1' _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["RESHUFFLE"]._options = None _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["RESHUFFLE"]._serialized_options = b'\242\264\372\302\005\033beam:transform:reshuffle:v1' _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["WRITE_FILES"]._options = None _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["WRITE_FILES"]._serialized_options = b'\242\264\372\302\005\035beam:transform:write_files:v1' _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["PUBSUB_READ"]._options = None _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["PUBSUB_READ"]._serialized_options = b'\242\264\372\302\005\035beam:transform:pubsub_read:v1' _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["PUBSUB_WRITE"]._options = None _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["PUBSUB_WRITE"]._serialized_options = b'\242\264\372\302\005\036beam:transform:pubsub_write:v1' _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["PUBSUB_WRITE_V2"]._options = None _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["PUBSUB_WRITE_V2"]._serialized_options = b'\242\264\372\302\005\036beam:transform:pubsub_write:v2' _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["GROUP_INTO_BATCHES_WITH_SHARDED_KEY"]._options = None _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["GROUP_INTO_BATCHES_WITH_SHARDED_KEY"]._serialized_options = b'\242\264\372\302\0055beam:transform:group_into_batches_with_sharded_key:v1' _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["REDISTRIBUTE_BY_KEY"]._options = None _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["REDISTRIBUTE_BY_KEY"]._serialized_options = b'\242\264\372\302\005%beam:transform:redistribute_by_key:v1' _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["REDISTRIBUTE_ARBITRARILY"]._options = None _globals['_STANDARDPTRANSFORMS_COMPOSITES'].values_by_name["REDISTRIBUTE_ARBITRARILY"]._serialized_options = b'\242\264\372\302\005*beam:transform:redistribute_arbitrarily:v1' _globals['_STANDARDPTRANSFORMS_COMBINECOMPONENTS'].values_by_name["COMBINE_PER_KEY_PRECOMBINE"]._options = None _globals['_STANDARDPTRANSFORMS_COMBINECOMPONENTS'].values_by_name["COMBINE_PER_KEY_PRECOMBINE"]._serialized_options = b'\242\264\372\302\005,beam:transform:combine_per_key_precombine:v1' _globals['_STANDARDPTRANSFORMS_COMBINECOMPONENTS'].values_by_name["COMBINE_PER_KEY_MERGE_ACCUMULATORS"]._options = None _globals['_STANDARDPTRANSFORMS_COMBINECOMPONENTS'].values_by_name["COMBINE_PER_KEY_MERGE_ACCUMULATORS"]._serialized_options = b'\242\264\372\302\0054beam:transform:combine_per_key_merge_accumulators:v1' _globals['_STANDARDPTRANSFORMS_COMBINECOMPONENTS'].values_by_name["COMBINE_PER_KEY_EXTRACT_OUTPUTS"]._options = None _globals['_STANDARDPTRANSFORMS_COMBINECOMPONENTS'].values_by_name["COMBINE_PER_KEY_EXTRACT_OUTPUTS"]._serialized_options = b'\242\264\372\302\0051beam:transform:combine_per_key_extract_outputs:v1' _globals['_STANDARDPTRANSFORMS_COMBINECOMPONENTS'].values_by_name["COMBINE_GROUPED_VALUES"]._options = None _globals['_STANDARDPTRANSFORMS_COMBINECOMPONENTS'].values_by_name["COMBINE_GROUPED_VALUES"]._serialized_options = b'\242\264\372\302\005(beam:transform:combine_grouped_values:v1' _globals['_STANDARDPTRANSFORMS_COMBINECOMPONENTS'].values_by_name["COMBINE_PER_KEY_CONVERT_TO_ACCUMULATORS"]._options = None _globals['_STANDARDPTRANSFORMS_COMBINECOMPONENTS'].values_by_name["COMBINE_PER_KEY_CONVERT_TO_ACCUMULATORS"]._serialized_options = b'\242\264\372\302\0059beam:transform:combine_per_key_convert_to_accumulators:v1' _globals['_STANDARDPTRANSFORMS_SPLITTABLEPARDOCOMPONENTS'].values_by_name["PAIR_WITH_RESTRICTION"]._options = None _globals['_STANDARDPTRANSFORMS_SPLITTABLEPARDOCOMPONENTS'].values_by_name["PAIR_WITH_RESTRICTION"]._serialized_options = b'\242\264\372\302\005+beam:transform:sdf_pair_with_restriction:v1' _globals['_STANDARDPTRANSFORMS_SPLITTABLEPARDOCOMPONENTS'].values_by_name["SPLIT_AND_SIZE_RESTRICTIONS"]._options = None _globals['_STANDARDPTRANSFORMS_SPLITTABLEPARDOCOMPONENTS'].values_by_name["SPLIT_AND_SIZE_RESTRICTIONS"]._serialized_options = b'\242\264\372\302\0051beam:transform:sdf_split_and_size_restrictions:v1' _globals['_STANDARDPTRANSFORMS_SPLITTABLEPARDOCOMPONENTS'].values_by_name["PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS"]._options = None _globals['_STANDARDPTRANSFORMS_SPLITTABLEPARDOCOMPONENTS'].values_by_name["PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS"]._serialized_options = b'\242\264\372\302\005= (3, 10): import typing as typing_extensions else: import typing_extensions DESCRIPTOR: google.protobuf.descriptor.FileDescriptor @typing_extensions.final class BeamConstants(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Constants: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _ConstantsEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[BeamConstants._Constants.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor MIN_TIMESTAMP_MILLIS: BeamConstants._Constants.ValueType # 0 """All timestamps of elements or window boundaries must be within the interval [MIN_TIMESTAMP_MILLIS, MAX_TIMESTAMP_MILLIS]. The smallest representable timestamp of an element or a window boundary. """ MAX_TIMESTAMP_MILLIS: BeamConstants._Constants.ValueType # 1 """The largest representable timestamp of an element or a window boundary.""" GLOBAL_WINDOW_MAX_TIMESTAMP_MILLIS: BeamConstants._Constants.ValueType # 2 """The maximum timestamp for the global window. Triggers use max timestamp to set timers' timestamp. Timers fire when the watermark passes their timestamps. So, the timestamp needs to be smaller than the MAX_TIMESTAMP_MILLIS. One standard day is subtracted from MAX_TIMESTAMP_MILLIS to make sure the max timestamp is smaller than MAX_TIMESTAMP_MILLIS even after rounding up to seconds or minutes. """ class Constants(_Constants, metaclass=_ConstantsEnumTypeWrapper): """All timestamps in milliseconds since Jan 1, 1970.""" MIN_TIMESTAMP_MILLIS: BeamConstants.Constants.ValueType # 0 """All timestamps of elements or window boundaries must be within the interval [MIN_TIMESTAMP_MILLIS, MAX_TIMESTAMP_MILLIS]. The smallest representable timestamp of an element or a window boundary. """ MAX_TIMESTAMP_MILLIS: BeamConstants.Constants.ValueType # 1 """The largest representable timestamp of an element or a window boundary.""" GLOBAL_WINDOW_MAX_TIMESTAMP_MILLIS: BeamConstants.Constants.ValueType # 2 """The maximum timestamp for the global window. Triggers use max timestamp to set timers' timestamp. Timers fire when the watermark passes their timestamps. So, the timestamp needs to be smaller than the MAX_TIMESTAMP_MILLIS. One standard day is subtracted from MAX_TIMESTAMP_MILLIS to make sure the max timestamp is smaller than MAX_TIMESTAMP_MILLIS even after rounding up to seconds or minutes. """ def __init__( self, ) -> None: ... global___BeamConstants = BeamConstants @typing_extensions.final class Components(google.protobuf.message.Message): """A set of mappings from id to message. This is included as an optional field on any proto message that may contain references needing resolution. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class TransformsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> global___PTransform: ... def __init__( self, *, key: builtins.str | None = ..., value: global___PTransform | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... @typing_extensions.final class PcollectionsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> global___PCollection: ... def __init__( self, *, key: builtins.str | None = ..., value: global___PCollection | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... @typing_extensions.final class WindowingStrategiesEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> global___WindowingStrategy: ... def __init__( self, *, key: builtins.str | None = ..., value: global___WindowingStrategy | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... @typing_extensions.final class CodersEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> global___Coder: ... def __init__( self, *, key: builtins.str | None = ..., value: global___Coder | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... @typing_extensions.final class EnvironmentsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> global___Environment: ... def __init__( self, *, key: builtins.str | None = ..., value: global___Environment | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... TRANSFORMS_FIELD_NUMBER: builtins.int PCOLLECTIONS_FIELD_NUMBER: builtins.int WINDOWING_STRATEGIES_FIELD_NUMBER: builtins.int CODERS_FIELD_NUMBER: builtins.int ENVIRONMENTS_FIELD_NUMBER: builtins.int @property def transforms(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___PTransform]: """(Required) A map from pipeline-scoped id to PTransform. Keys of the transforms map may be used by runners to identify pipeline steps. Hence it's recommended to use strings that are not too long that match regex '[A-Za-z0-9-_]+'. """ @property def pcollections(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___PCollection]: """(Required) A map from pipeline-scoped id to PCollection.""" @property def windowing_strategies(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___WindowingStrategy]: """(Required) A map from pipeline-scoped id to WindowingStrategy.""" @property def coders(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___Coder]: """(Required) A map from pipeline-scoped id to Coder.""" @property def environments(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___Environment]: """(Required) A map from pipeline-scoped id to Environment.""" def __init__( self, *, transforms: collections.abc.Mapping[builtins.str, global___PTransform] | None = ..., pcollections: collections.abc.Mapping[builtins.str, global___PCollection] | None = ..., windowing_strategies: collections.abc.Mapping[builtins.str, global___WindowingStrategy] | None = ..., coders: collections.abc.Mapping[builtins.str, global___Coder] | None = ..., environments: collections.abc.Mapping[builtins.str, global___Environment] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["coders", b"coders", "environments", b"environments", "pcollections", b"pcollections", "transforms", b"transforms", "windowing_strategies", b"windowing_strategies"]) -> None: ... global___Components = Components @typing_extensions.final class Pipeline(google.protobuf.message.Message): """A Pipeline is a hierarchical graph of PTransforms, linked by PCollections. A typical graph may look like: Impulse -> PCollection -> ParDo -> PCollection -> GroupByKey -> ... \\> PCollection -> ParDo -> ... \\> ParDo -> ... Impulse -> PCollection -> ParDo -> PCollection -> ... This is represented by a number of by-reference maps to transforms, PCollections, SDK environments, coders, etc., for supporting compact reuse and arbitrary graph structure. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor COMPONENTS_FIELD_NUMBER: builtins.int ROOT_TRANSFORM_IDS_FIELD_NUMBER: builtins.int DISPLAY_DATA_FIELD_NUMBER: builtins.int REQUIREMENTS_FIELD_NUMBER: builtins.int @property def components(self) -> global___Components: """(Required) The coders, UDFs, graph nodes, etc, that make up this pipeline. """ @property def root_transform_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """(Required) The ids of all PTransforms that are not contained within another PTransform. These must be in shallow topological order, so that traversing them recursively in this order yields a recursively topological traversal. """ @property def display_data(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___DisplayData]: """(Optional) Static display data for the pipeline. If there is none, it may be omitted. """ @property def requirements(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """(Optional) A set of requirements that the runner MUST understand and be able to faithfully provide in order to execute this pipeline. These may indicate that a runner must inspect new fields on a component or provide additional guarantees when processing specific transforms. A runner should reject any pipelines with unknown requirements. """ def __init__( self, *, components: global___Components | None = ..., root_transform_ids: collections.abc.Iterable[builtins.str] | None = ..., display_data: collections.abc.Iterable[global___DisplayData] | None = ..., requirements: collections.abc.Iterable[builtins.str] | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["components", b"components"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["components", b"components", "display_data", b"display_data", "requirements", b"requirements", "root_transform_ids", b"root_transform_ids"]) -> None: ... global___Pipeline = Pipeline @typing_extensions.final class PTransform(google.protobuf.message.Message): """Transforms are the operations in your pipeline, and provide a generic processing framework. You provide processing logic in the form of a function object (colloquially referred to as “user code”), and your user code is applied to each element of an input PCollection (or more than one PCollection). Depending on the pipeline runner and back-end that you choose, many different workers across a cluster may execute instances of your user code in parallel. The user code running on each worker generates the output elements that are ultimately added to the final output PCollection that the transform produces. The Beam SDKs contain a number of different transforms that you can apply to your pipeline’s PCollections. These include general-purpose core transforms, such as ParDo or Combine. There are also pre-written composite transforms included in the SDKs, which combine one or more of the core transforms in a useful processing pattern, such as counting or combining elements in a collection. You can also define your own more complex composite transforms to fit your pipeline’s exact use case. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class InputsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str value: builtins.str def __init__( self, *, key: builtins.str | None = ..., value: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... @typing_extensions.final class OutputsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str value: builtins.str def __init__( self, *, key: builtins.str | None = ..., value: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... @typing_extensions.final class AnnotationsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str value: builtins.bytes def __init__( self, *, key: builtins.str | None = ..., value: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... UNIQUE_NAME_FIELD_NUMBER: builtins.int SPEC_FIELD_NUMBER: builtins.int SUBTRANSFORMS_FIELD_NUMBER: builtins.int INPUTS_FIELD_NUMBER: builtins.int OUTPUTS_FIELD_NUMBER: builtins.int DISPLAY_DATA_FIELD_NUMBER: builtins.int ENVIRONMENT_ID_FIELD_NUMBER: builtins.int ANNOTATIONS_FIELD_NUMBER: builtins.int unique_name: builtins.str """(Required) A unique name for the application node. Ideally, this should be stable over multiple evolutions of a pipeline for the purposes of logging and associating pipeline state with a node, etc. If it is not stable, then the runner decides what will happen. But, most importantly, it must always be here and be unique, even if it is autogenerated. """ @property def spec(self) -> global___FunctionSpec: """(Optional) A URN and payload that, together, fully defined the semantics of this transform. If absent, this must be an "anonymous" composite transform. For primitive transform in the Runner API, this is required, and the payloads are well-defined messages. When the URN indicates ParDo it is a ParDoPayload, and so on. For some special composite transforms, the payload is also officially defined. See StandardPTransforms for details. """ @property def subtransforms(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """(Optional) A list of the ids of transforms that it contains. Primitive transforms (see StandardPTransforms.Primitives) are not allowed to specify subtransforms. Note that a composite transform may have zero subtransforms as long as it only outputs PCollections that are in its inputs. """ @property def inputs(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: """(Required) A map from local names of inputs (unique only with this map, and likely embedded in the transform payload and serialized user code) to PCollection ids. The payload for this transform may clarify the relationship of these inputs. For example: - for a Flatten transform they are merged - for a ParDo transform, some may be side inputs All inputs are recorded here so that the topological ordering of the graph is consistent whether or not the payload is understood. """ @property def outputs(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: """(Required) A map from local names of outputs (unique only within this map, and likely embedded in the transform payload and serialized user code) to PCollection ids. The URN or payload for this transform node may clarify the type and relationship of these outputs. For example: - for a ParDo transform, these are tags on PCollections, which will be embedded in the DoFn. """ @property def display_data(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___DisplayData]: """(Optional) Static display data for this PTransform application. If there is none, it may be omitted. """ environment_id: builtins.str """Environment where the current PTransform should be executed in. Transforms that are required to be implemented by a runner must omit this. All other transforms are required to specify this. """ @property def annotations(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.bytes]: """(Optional) A map from URNs designating a type of annotation, to the annotation in binary format. For example, an annotation could indicate that this PTransform has specific privacy properties. A runner MAY ignore types of annotations it doesn't understand. Therefore annotations MUST NOT be used for metadata that can affect correct execution of the transform. """ def __init__( self, *, unique_name: builtins.str | None = ..., spec: global___FunctionSpec | None = ..., subtransforms: collections.abc.Iterable[builtins.str] | None = ..., inputs: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., outputs: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., display_data: collections.abc.Iterable[global___DisplayData] | None = ..., environment_id: builtins.str | None = ..., annotations: collections.abc.Mapping[builtins.str, builtins.bytes] | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["spec", b"spec"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["annotations", b"annotations", "display_data", b"display_data", "environment_id", b"environment_id", "inputs", b"inputs", "outputs", b"outputs", "spec", b"spec", "subtransforms", b"subtransforms", "unique_name", b"unique_name"]) -> None: ... global___PTransform = PTransform @typing_extensions.final class StandardPTransforms(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Primitives: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _PrimitivesEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardPTransforms._Primitives.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor PAR_DO: StandardPTransforms._Primitives.ValueType # 0 """ParDo is a Beam transform for generic parallel processing. The ParDo processing paradigm is similar to the “Map” phase of a Map/Shuffle/Reduce-style algorithm: a ParDo transform considers each element in the input PCollection, performs some processing function (your user code) on that element, and emits zero, one, or multiple elements to an output PCollection. See https://beam.apache.org/documentation/programming-guide/#pardo for additional details. Payload: ParDoPayload """ FLATTEN: StandardPTransforms._Primitives.ValueType # 1 """Flatten is a Beam transform for PCollection objects that store the same data type. Flatten merges multiple PCollection objects into a single logical PCollection. See https://beam.apache.org/documentation/programming-guide/#flatten for additional details. Payload: None """ GROUP_BY_KEY: StandardPTransforms._Primitives.ValueType # 2 """GroupByKey is a Beam transform for processing collections of key/value pairs. It’s a parallel reduction operation, analogous to the Shuffle phase of a Map/Shuffle/Reduce-style algorithm. The input to GroupByKey is a collection of key/value pairs that represents a multimap, where the collection contains multiple pairs that have the same key, but different values. Given such a collection, you use GroupByKey to collect all of the values associated with each unique key. See https://beam.apache.org/documentation/programming-guide/#groupbykey for additional details. Never defines an environment as the runner is required to implement this transform. Payload: None """ IMPULSE: StandardPTransforms._Primitives.ValueType # 3 """A transform which produces a single empty byte array at the minimum timestamp in the GlobalWindow. Never defines an environment as the runner is required to implement this transform. Payload: None """ ASSIGN_WINDOWS: StandardPTransforms._Primitives.ValueType # 4 """Windowing subdivides a PCollection according to the timestamps of its individual elements. Transforms that aggregate multiple elements, such as GroupByKey and Combine, work implicitly on a per-window basis — they process each PCollection as a succession of multiple, finite windows, though the entire collection itself may be of unbounded size. See https://beam.apache.org/documentation/programming-guide/#windowing for additional details. Payload: WindowIntoPayload """ TEST_STREAM: StandardPTransforms._Primitives.ValueType # 5 """A testing input that generates an unbounded {@link PCollection} of elements, advancing the watermark and processing time as elements are emitted. After all of the specified elements are emitted, ceases to produce output. See https://beam.apache.org/blog/2016/10/20/test-stream.html for additional details. Payload: TestStreamPayload """ MAP_WINDOWS: StandardPTransforms._Primitives.ValueType # 6 """Represents mapping of main input window onto side input window. Side input window mapping function: Input: KV Output: KV For each main input window, the side input window is returned. The nonce is used by a runner to associate each input with its output. The nonce is represented as an opaque set of bytes. Payload: SideInput#window_mapping_fn FunctionSpec """ MERGE_WINDOWS: StandardPTransforms._Primitives.ValueType # 7 """Used to merge windows during a GroupByKey. Window merging function: Input: KV> Output: KV, iterable>>> For each set of original windows, a list of all unmerged windows is output alongside a map of merged window to set of consumed windows. All original windows must be contained in either the unmerged original window set or one of the consumed original window sets. Each original window can only be part of one output set. The nonce is used by a runner to associate each input with its output. The nonce is represented as an opaque set of bytes. Payload: WindowingStrategy#window_fn FunctionSpec """ TO_STRING: StandardPTransforms._Primitives.ValueType # 8 """A transform that translates a given element to its human-readable representation. Input: KV Output: KV For each given element, the implementation returns the best-effort human-readable representation. When possible, the implementation could call a user-overridable method on the type. For example, Java could call `toString()`, Python could call `str()`, Golang could call `String()`. The nonce is used by a runner to associate each input with its output. The nonce is represented as an opaque set of bytes. Payload: none """ class Primitives(_Primitives, metaclass=_PrimitivesEnumTypeWrapper): """Primitive transforms may not specify composite sub-transforms.""" PAR_DO: StandardPTransforms.Primitives.ValueType # 0 """ParDo is a Beam transform for generic parallel processing. The ParDo processing paradigm is similar to the “Map” phase of a Map/Shuffle/Reduce-style algorithm: a ParDo transform considers each element in the input PCollection, performs some processing function (your user code) on that element, and emits zero, one, or multiple elements to an output PCollection. See https://beam.apache.org/documentation/programming-guide/#pardo for additional details. Payload: ParDoPayload """ FLATTEN: StandardPTransforms.Primitives.ValueType # 1 """Flatten is a Beam transform for PCollection objects that store the same data type. Flatten merges multiple PCollection objects into a single logical PCollection. See https://beam.apache.org/documentation/programming-guide/#flatten for additional details. Payload: None """ GROUP_BY_KEY: StandardPTransforms.Primitives.ValueType # 2 """GroupByKey is a Beam transform for processing collections of key/value pairs. It’s a parallel reduction operation, analogous to the Shuffle phase of a Map/Shuffle/Reduce-style algorithm. The input to GroupByKey is a collection of key/value pairs that represents a multimap, where the collection contains multiple pairs that have the same key, but different values. Given such a collection, you use GroupByKey to collect all of the values associated with each unique key. See https://beam.apache.org/documentation/programming-guide/#groupbykey for additional details. Never defines an environment as the runner is required to implement this transform. Payload: None """ IMPULSE: StandardPTransforms.Primitives.ValueType # 3 """A transform which produces a single empty byte array at the minimum timestamp in the GlobalWindow. Never defines an environment as the runner is required to implement this transform. Payload: None """ ASSIGN_WINDOWS: StandardPTransforms.Primitives.ValueType # 4 """Windowing subdivides a PCollection according to the timestamps of its individual elements. Transforms that aggregate multiple elements, such as GroupByKey and Combine, work implicitly on a per-window basis — they process each PCollection as a succession of multiple, finite windows, though the entire collection itself may be of unbounded size. See https://beam.apache.org/documentation/programming-guide/#windowing for additional details. Payload: WindowIntoPayload """ TEST_STREAM: StandardPTransforms.Primitives.ValueType # 5 """A testing input that generates an unbounded {@link PCollection} of elements, advancing the watermark and processing time as elements are emitted. After all of the specified elements are emitted, ceases to produce output. See https://beam.apache.org/blog/2016/10/20/test-stream.html for additional details. Payload: TestStreamPayload """ MAP_WINDOWS: StandardPTransforms.Primitives.ValueType # 6 """Represents mapping of main input window onto side input window. Side input window mapping function: Input: KV Output: KV For each main input window, the side input window is returned. The nonce is used by a runner to associate each input with its output. The nonce is represented as an opaque set of bytes. Payload: SideInput#window_mapping_fn FunctionSpec """ MERGE_WINDOWS: StandardPTransforms.Primitives.ValueType # 7 """Used to merge windows during a GroupByKey. Window merging function: Input: KV> Output: KV, iterable>>> For each set of original windows, a list of all unmerged windows is output alongside a map of merged window to set of consumed windows. All original windows must be contained in either the unmerged original window set or one of the consumed original window sets. Each original window can only be part of one output set. The nonce is used by a runner to associate each input with its output. The nonce is represented as an opaque set of bytes. Payload: WindowingStrategy#window_fn FunctionSpec """ TO_STRING: StandardPTransforms.Primitives.ValueType # 8 """A transform that translates a given element to its human-readable representation. Input: KV Output: KV For each given element, the implementation returns the best-effort human-readable representation. When possible, the implementation could call a user-overridable method on the type. For example, Java could call `toString()`, Python could call `str()`, Golang could call `String()`. The nonce is used by a runner to associate each input with its output. The nonce is represented as an opaque set of bytes. Payload: none """ class _DeprecatedPrimitives: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _DeprecatedPrimitivesEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardPTransforms._DeprecatedPrimitives.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor READ: StandardPTransforms._DeprecatedPrimitives.ValueType # 0 """Represents the operation to read a Bounded or Unbounded source. Payload: ReadPayload. """ CREATE_VIEW: StandardPTransforms._DeprecatedPrimitives.ValueType # 1 """Runners should move away from translating `CreatePCollectionView` and treat this as part of the translation for a `ParDo` side input. """ class DeprecatedPrimitives(_DeprecatedPrimitives, metaclass=_DeprecatedPrimitivesEnumTypeWrapper): ... READ: StandardPTransforms.DeprecatedPrimitives.ValueType # 0 """Represents the operation to read a Bounded or Unbounded source. Payload: ReadPayload. """ CREATE_VIEW: StandardPTransforms.DeprecatedPrimitives.ValueType # 1 """Runners should move away from translating `CreatePCollectionView` and treat this as part of the translation for a `ParDo` side input. """ class _Composites: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _CompositesEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardPTransforms._Composites.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor COMBINE_PER_KEY: StandardPTransforms._Composites.ValueType # 0 """Represents the Combine.perKey() operation. If this is produced by an SDK, it is assumed that the SDK understands each of CombineComponents. Payload: CombinePayload """ COMBINE_GLOBALLY: StandardPTransforms._Composites.ValueType # 1 """Represents the Combine.globally() operation. If this is produced by an SDK, it is assumed that the SDK understands each of CombineComponents. Payload: CombinePayload """ RESHUFFLE: StandardPTransforms._Composites.ValueType # 2 """Represents the Reshuffle operation.""" WRITE_FILES: StandardPTransforms._Composites.ValueType # 3 """Less well-known. Payload: WriteFilesPayload.""" PUBSUB_READ: StandardPTransforms._Composites.ValueType # 4 """Payload: PubSubReadPayload.""" PUBSUB_WRITE: StandardPTransforms._Composites.ValueType # 5 """Payload: PubSubWritePayload.""" PUBSUB_WRITE_V2: StandardPTransforms._Composites.ValueType # 7 """Used for pubsub dynamic destinations. Payload: PubSubWritePayload. """ GROUP_INTO_BATCHES_WITH_SHARDED_KEY: StandardPTransforms._Composites.ValueType # 6 """Represents the GroupIntoBatches.WithShardedKey operation. Payload: GroupIntoBatchesPayload """ REDISTRIBUTE_BY_KEY: StandardPTransforms._Composites.ValueType # 8 REDISTRIBUTE_ARBITRARILY: StandardPTransforms._Composites.ValueType # 9 class Composites(_Composites, metaclass=_CompositesEnumTypeWrapper): ... COMBINE_PER_KEY: StandardPTransforms.Composites.ValueType # 0 """Represents the Combine.perKey() operation. If this is produced by an SDK, it is assumed that the SDK understands each of CombineComponents. Payload: CombinePayload """ COMBINE_GLOBALLY: StandardPTransforms.Composites.ValueType # 1 """Represents the Combine.globally() operation. If this is produced by an SDK, it is assumed that the SDK understands each of CombineComponents. Payload: CombinePayload """ RESHUFFLE: StandardPTransforms.Composites.ValueType # 2 """Represents the Reshuffle operation.""" WRITE_FILES: StandardPTransforms.Composites.ValueType # 3 """Less well-known. Payload: WriteFilesPayload.""" PUBSUB_READ: StandardPTransforms.Composites.ValueType # 4 """Payload: PubSubReadPayload.""" PUBSUB_WRITE: StandardPTransforms.Composites.ValueType # 5 """Payload: PubSubWritePayload.""" PUBSUB_WRITE_V2: StandardPTransforms.Composites.ValueType # 7 """Used for pubsub dynamic destinations. Payload: PubSubWritePayload. """ GROUP_INTO_BATCHES_WITH_SHARDED_KEY: StandardPTransforms.Composites.ValueType # 6 """Represents the GroupIntoBatches.WithShardedKey operation. Payload: GroupIntoBatchesPayload """ REDISTRIBUTE_BY_KEY: StandardPTransforms.Composites.ValueType # 8 REDISTRIBUTE_ARBITRARILY: StandardPTransforms.Composites.ValueType # 9 class _CombineComponents: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _CombineComponentsEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardPTransforms._CombineComponents.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor COMBINE_PER_KEY_PRECOMBINE: StandardPTransforms._CombineComponents.ValueType # 0 """Represents the Pre-Combine part of a lifted Combine Per Key, as described in the following document: https://s.apache.org/beam-runner-api-combine-model#heading=h.ta0g6ase8z07 Payload: CombinePayload """ COMBINE_PER_KEY_MERGE_ACCUMULATORS: StandardPTransforms._CombineComponents.ValueType # 1 """Represents the Merge Accumulators part of a lifted Combine Per Key, as described in the following document: https://s.apache.org/beam-runner-api-combine-model#heading=h.jco9rvatld5m Payload: CombinePayload """ COMBINE_PER_KEY_EXTRACT_OUTPUTS: StandardPTransforms._CombineComponents.ValueType # 2 """Represents the Extract Outputs part of a lifted Combine Per Key, as described in the following document: https://s.apache.org/beam-runner-api-combine-model#heading=h.i9i6p8gtl6ku Payload: CombinePayload """ COMBINE_GROUPED_VALUES: StandardPTransforms._CombineComponents.ValueType # 3 """Represents the Combine Grouped Values transform, as described in the following document: https://s.apache.org/beam-runner-api-combine-model#heading=h.aj86ew4v1wk Payload: CombinePayload """ COMBINE_PER_KEY_CONVERT_TO_ACCUMULATORS: StandardPTransforms._CombineComponents.ValueType # 4 """Represents the Convert To Accumulators transform, as described in the following document: https://s.apache.org/beam-runner-api-combine-model#heading=h.h5697l1scd9x Payload: CombinePayload """ class CombineComponents(_CombineComponents, metaclass=_CombineComponentsEnumTypeWrapper): """Payload for all of these: CombinePayload""" COMBINE_PER_KEY_PRECOMBINE: StandardPTransforms.CombineComponents.ValueType # 0 """Represents the Pre-Combine part of a lifted Combine Per Key, as described in the following document: https://s.apache.org/beam-runner-api-combine-model#heading=h.ta0g6ase8z07 Payload: CombinePayload """ COMBINE_PER_KEY_MERGE_ACCUMULATORS: StandardPTransforms.CombineComponents.ValueType # 1 """Represents the Merge Accumulators part of a lifted Combine Per Key, as described in the following document: https://s.apache.org/beam-runner-api-combine-model#heading=h.jco9rvatld5m Payload: CombinePayload """ COMBINE_PER_KEY_EXTRACT_OUTPUTS: StandardPTransforms.CombineComponents.ValueType # 2 """Represents the Extract Outputs part of a lifted Combine Per Key, as described in the following document: https://s.apache.org/beam-runner-api-combine-model#heading=h.i9i6p8gtl6ku Payload: CombinePayload """ COMBINE_GROUPED_VALUES: StandardPTransforms.CombineComponents.ValueType # 3 """Represents the Combine Grouped Values transform, as described in the following document: https://s.apache.org/beam-runner-api-combine-model#heading=h.aj86ew4v1wk Payload: CombinePayload """ COMBINE_PER_KEY_CONVERT_TO_ACCUMULATORS: StandardPTransforms.CombineComponents.ValueType # 4 """Represents the Convert To Accumulators transform, as described in the following document: https://s.apache.org/beam-runner-api-combine-model#heading=h.h5697l1scd9x Payload: CombinePayload """ class _SplittableParDoComponents: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _SplittableParDoComponentsEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardPTransforms._SplittableParDoComponents.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor PAIR_WITH_RESTRICTION: StandardPTransforms._SplittableParDoComponents.ValueType # 0 """Pairs the input element with its initial restriction. Input: element; output: KV(element, restriction). """ SPLIT_AND_SIZE_RESTRICTIONS: StandardPTransforms._SplittableParDoComponents.ValueType # 1 """Splits the restriction of each element/restriction pair and returns the resulting splits, with a corresponding floating point size estimation for each. A reasonable value for size is the number of bytes expected to be produced by this (element, restriction) pair. Input: KV(element, restriction) Output: KV(KV(element, restriction), size)) """ PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS: StandardPTransforms._SplittableParDoComponents.ValueType # 2 """Applies the DoFn to every element and restriction. All primary and residuals returned from checkpointing or splitting must have the same type as the input to this transform. Input: KV(KV(element, restriction), size); output: DoFn's output. """ TRUNCATE_SIZED_RESTRICTION: StandardPTransforms._SplittableParDoComponents.ValueType # 3 """Truncates the restriction of each element/restriction pair and returns the finite restriction which will be processed when a pipeline is drained. See https://docs.google.com/document/d/1NExwHlj-2q2WUGhSO4jTu8XGhDPmm3cllSN8IMmWci8/edit#. for additional details about drain. Input: KV(KV(element, restriction), size); Output: KV(KV(element, restriction), size). """ class SplittableParDoComponents(_SplittableParDoComponents, metaclass=_SplittableParDoComponentsEnumTypeWrapper): """Payload for all of these: ParDoPayload containing the user's SDF""" PAIR_WITH_RESTRICTION: StandardPTransforms.SplittableParDoComponents.ValueType # 0 """Pairs the input element with its initial restriction. Input: element; output: KV(element, restriction). """ SPLIT_AND_SIZE_RESTRICTIONS: StandardPTransforms.SplittableParDoComponents.ValueType # 1 """Splits the restriction of each element/restriction pair and returns the resulting splits, with a corresponding floating point size estimation for each. A reasonable value for size is the number of bytes expected to be produced by this (element, restriction) pair. Input: KV(element, restriction) Output: KV(KV(element, restriction), size)) """ PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS: StandardPTransforms.SplittableParDoComponents.ValueType # 2 """Applies the DoFn to every element and restriction. All primary and residuals returned from checkpointing or splitting must have the same type as the input to this transform. Input: KV(KV(element, restriction), size); output: DoFn's output. """ TRUNCATE_SIZED_RESTRICTION: StandardPTransforms.SplittableParDoComponents.ValueType # 3 """Truncates the restriction of each element/restriction pair and returns the finite restriction which will be processed when a pipeline is drained. See https://docs.google.com/document/d/1NExwHlj-2q2WUGhSO4jTu8XGhDPmm3cllSN8IMmWci8/edit#. for additional details about drain. Input: KV(KV(element, restriction), size); Output: KV(KV(element, restriction), size). """ class _GroupIntoBatchesComponents: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _GroupIntoBatchesComponentsEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardPTransforms._GroupIntoBatchesComponents.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor GROUP_INTO_BATCHES: StandardPTransforms._GroupIntoBatchesComponents.ValueType # 0 class GroupIntoBatchesComponents(_GroupIntoBatchesComponents, metaclass=_GroupIntoBatchesComponentsEnumTypeWrapper): """Payload for all of these: GroupIntoBatchesPayload""" GROUP_INTO_BATCHES: StandardPTransforms.GroupIntoBatchesComponents.ValueType # 0 def __init__( self, ) -> None: ... global___StandardPTransforms = StandardPTransforms @typing_extensions.final class StandardSideInputTypes(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardSideInputTypes._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor ITERABLE: StandardSideInputTypes._Enum.ValueType # 0 """Represents a view over a PCollection. StateGetRequests performed on this side input must use StateKey.IterableSideInput. """ MULTIMAP: StandardSideInputTypes._Enum.ValueType # 1 """Represents a view over a PCollection>. StateGetRequests performed on this side input must use StateKey.MultimapKeysSideInput or StateKey.MultimapSideInput. """ class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... ITERABLE: StandardSideInputTypes.Enum.ValueType # 0 """Represents a view over a PCollection. StateGetRequests performed on this side input must use StateKey.IterableSideInput. """ MULTIMAP: StandardSideInputTypes.Enum.ValueType # 1 """Represents a view over a PCollection>. StateGetRequests performed on this side input must use StateKey.MultimapKeysSideInput or StateKey.MultimapSideInput. """ def __init__( self, ) -> None: ... global___StandardSideInputTypes = StandardSideInputTypes @typing_extensions.final class StandardUserStateTypes(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardUserStateTypes._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor BAG: StandardUserStateTypes._Enum.ValueType # 0 """Represents a user state specification that supports a bag. StateRequests performed on this user state must use StateKey.BagUserState. """ MULTIMAP: StandardUserStateTypes._Enum.ValueType # 1 """Represents a user state specification that supports a multimap. StateRequests performed on this user state must use StateKey.MultimapKeysUserState or StateKey.MultimapUserState. """ ORDERED_LIST: StandardUserStateTypes._Enum.ValueType # 2 """Represents a user state specification that supports an ordered list. StateRequests performed on this user state must use StateKey.OrderedListUserState. """ class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... BAG: StandardUserStateTypes.Enum.ValueType # 0 """Represents a user state specification that supports a bag. StateRequests performed on this user state must use StateKey.BagUserState. """ MULTIMAP: StandardUserStateTypes.Enum.ValueType # 1 """Represents a user state specification that supports a multimap. StateRequests performed on this user state must use StateKey.MultimapKeysUserState or StateKey.MultimapUserState. """ ORDERED_LIST: StandardUserStateTypes.Enum.ValueType # 2 """Represents a user state specification that supports an ordered list. StateRequests performed on this user state must use StateKey.OrderedListUserState. """ def __init__( self, ) -> None: ... global___StandardUserStateTypes = StandardUserStateTypes @typing_extensions.final class PCollection(google.protobuf.message.Message): """A PCollection!""" DESCRIPTOR: google.protobuf.descriptor.Descriptor UNIQUE_NAME_FIELD_NUMBER: builtins.int CODER_ID_FIELD_NUMBER: builtins.int IS_BOUNDED_FIELD_NUMBER: builtins.int WINDOWING_STRATEGY_ID_FIELD_NUMBER: builtins.int DISPLAY_DATA_FIELD_NUMBER: builtins.int unique_name: builtins.str """(Required) A unique name for the PCollection. Ideally, this should be stable over multiple evolutions of a pipeline for the purposes of logging and associating pipeline state with a node, etc. If it is not stable, then the runner decides what will happen. But, most importantly, it must always be here, even if it is autogenerated. """ coder_id: builtins.str """(Required) The id of the Coder for this PCollection.""" is_bounded: global___IsBounded.Enum.ValueType """(Required) Whether this PCollection is bounded or unbounded""" windowing_strategy_id: builtins.str """(Required) The id of the windowing strategy for this PCollection.""" @property def display_data(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___DisplayData]: """(Optional) Static display data for the PCollection. If there is none, it may be omitted. """ def __init__( self, *, unique_name: builtins.str | None = ..., coder_id: builtins.str | None = ..., is_bounded: global___IsBounded.Enum.ValueType | None = ..., windowing_strategy_id: builtins.str | None = ..., display_data: collections.abc.Iterable[global___DisplayData] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["coder_id", b"coder_id", "display_data", b"display_data", "is_bounded", b"is_bounded", "unique_name", b"unique_name", "windowing_strategy_id", b"windowing_strategy_id"]) -> None: ... global___PCollection = PCollection @typing_extensions.final class ParDoPayload(google.protobuf.message.Message): """The payload for the primitive ParDo transform.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class SideInputsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> global___SideInput: ... def __init__( self, *, key: builtins.str | None = ..., value: global___SideInput | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... @typing_extensions.final class StateSpecsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> global___StateSpec: ... def __init__( self, *, key: builtins.str | None = ..., value: global___StateSpec | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... @typing_extensions.final class TimerFamilySpecsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> global___TimerFamilySpec: ... def __init__( self, *, key: builtins.str | None = ..., value: global___TimerFamilySpec | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... DO_FN_FIELD_NUMBER: builtins.int SIDE_INPUTS_FIELD_NUMBER: builtins.int STATE_SPECS_FIELD_NUMBER: builtins.int TIMER_FAMILY_SPECS_FIELD_NUMBER: builtins.int RESTRICTION_CODER_ID_FIELD_NUMBER: builtins.int REQUESTS_FINALIZATION_FIELD_NUMBER: builtins.int REQUIRES_TIME_SORTED_INPUT_FIELD_NUMBER: builtins.int REQUIRES_STABLE_INPUT_FIELD_NUMBER: builtins.int ON_WINDOW_EXPIRATION_TIMER_FAMILY_SPEC_FIELD_NUMBER: builtins.int @property def do_fn(self) -> global___FunctionSpec: """(Required) The FunctionSpec of the DoFn.""" @property def side_inputs(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___SideInput]: """(Optional) A mapping of local input names to side inputs, describing the expected access pattern. """ @property def state_specs(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___StateSpec]: """(Optional) A mapping of local state names to state specifications. If this is set, the stateful processing requirement should also be placed in the pipeline requirements. """ @property def timer_family_specs(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___TimerFamilySpec]: """(Optional) A mapping of local timer family names to timer family specifications. If this is set, the stateful processing requirement should also be placed in the pipeline requirements. """ restriction_coder_id: builtins.str """(Optional) Only set when this ParDo contains a splittable DoFn. If this is set, the corresponding standard requirement should also be placed in the pipeline requirements. """ requests_finalization: builtins.bool """(Optional) Only set when this ParDo can request bundle finalization. If this is set, the corresponding standard requirement should also be placed in the pipeline requirements. """ requires_time_sorted_input: builtins.bool """Whether this stage requires time sorted input. If this is set, the corresponding standard requirement should also be placed in the pipeline requirements. """ requires_stable_input: builtins.bool """Whether this stage requires stable input. If this is set, the corresponding standard requirement should also be placed in the pipeline requirements. """ on_window_expiration_timer_family_spec: builtins.str """If populated, the name of the timer family spec which should be notified on each window expiry. If this is set, the corresponding standard requirement should also be placed in the pipeline requirements. """ def __init__( self, *, do_fn: global___FunctionSpec | None = ..., side_inputs: collections.abc.Mapping[builtins.str, global___SideInput] | None = ..., state_specs: collections.abc.Mapping[builtins.str, global___StateSpec] | None = ..., timer_family_specs: collections.abc.Mapping[builtins.str, global___TimerFamilySpec] | None = ..., restriction_coder_id: builtins.str | None = ..., requests_finalization: builtins.bool | None = ..., requires_time_sorted_input: builtins.bool | None = ..., requires_stable_input: builtins.bool | None = ..., on_window_expiration_timer_family_spec: builtins.str | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["do_fn", b"do_fn"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["do_fn", b"do_fn", "on_window_expiration_timer_family_spec", b"on_window_expiration_timer_family_spec", "requests_finalization", b"requests_finalization", "requires_stable_input", b"requires_stable_input", "requires_time_sorted_input", b"requires_time_sorted_input", "restriction_coder_id", b"restriction_coder_id", "side_inputs", b"side_inputs", "state_specs", b"state_specs", "timer_family_specs", b"timer_family_specs"]) -> None: ... global___ParDoPayload = ParDoPayload @typing_extensions.final class StateSpec(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor READ_MODIFY_WRITE_SPEC_FIELD_NUMBER: builtins.int BAG_SPEC_FIELD_NUMBER: builtins.int COMBINING_SPEC_FIELD_NUMBER: builtins.int MAP_SPEC_FIELD_NUMBER: builtins.int SET_SPEC_FIELD_NUMBER: builtins.int ORDERED_LIST_SPEC_FIELD_NUMBER: builtins.int MULTIMAP_SPEC_FIELD_NUMBER: builtins.int PROTOCOL_FIELD_NUMBER: builtins.int @property def read_modify_write_spec(self) -> global___ReadModifyWriteStateSpec: ... @property def bag_spec(self) -> global___BagStateSpec: ... @property def combining_spec(self) -> global___CombiningStateSpec: ... @property def map_spec(self) -> global___MapStateSpec: ... @property def set_spec(self) -> global___SetStateSpec: ... @property def ordered_list_spec(self) -> global___OrderedListStateSpec: ... @property def multimap_spec(self) -> global___MultimapStateSpec: ... @property def protocol(self) -> global___FunctionSpec: """(Required) URN of the protocol required by this state specification to present the desired SDK-specific interface to a UDF. This protocol defines the SDK harness <-> Runner Harness RPC interface for accessing and mutating user state. See StandardUserStateTypes for an enumeration of all user state types defined. """ def __init__( self, *, read_modify_write_spec: global___ReadModifyWriteStateSpec | None = ..., bag_spec: global___BagStateSpec | None = ..., combining_spec: global___CombiningStateSpec | None = ..., map_spec: global___MapStateSpec | None = ..., set_spec: global___SetStateSpec | None = ..., ordered_list_spec: global___OrderedListStateSpec | None = ..., multimap_spec: global___MultimapStateSpec | None = ..., protocol: global___FunctionSpec | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["bag_spec", b"bag_spec", "combining_spec", b"combining_spec", "map_spec", b"map_spec", "multimap_spec", b"multimap_spec", "ordered_list_spec", b"ordered_list_spec", "protocol", b"protocol", "read_modify_write_spec", b"read_modify_write_spec", "set_spec", b"set_spec", "spec", b"spec"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["bag_spec", b"bag_spec", "combining_spec", b"combining_spec", "map_spec", b"map_spec", "multimap_spec", b"multimap_spec", "ordered_list_spec", b"ordered_list_spec", "protocol", b"protocol", "read_modify_write_spec", b"read_modify_write_spec", "set_spec", b"set_spec", "spec", b"spec"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["spec", b"spec"]) -> typing_extensions.Literal["read_modify_write_spec", "bag_spec", "combining_spec", "map_spec", "set_spec", "ordered_list_spec", "multimap_spec"] | None: ... global___StateSpec = StateSpec @typing_extensions.final class ReadModifyWriteStateSpec(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor CODER_ID_FIELD_NUMBER: builtins.int coder_id: builtins.str def __init__( self, *, coder_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["coder_id", b"coder_id"]) -> None: ... global___ReadModifyWriteStateSpec = ReadModifyWriteStateSpec @typing_extensions.final class BagStateSpec(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ELEMENT_CODER_ID_FIELD_NUMBER: builtins.int element_coder_id: builtins.str def __init__( self, *, element_coder_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["element_coder_id", b"element_coder_id"]) -> None: ... global___BagStateSpec = BagStateSpec @typing_extensions.final class OrderedListStateSpec(google.protobuf.message.Message): """OrderedListState values are encoded with the var int encoded millis-since-epoch followed by the value encoded by the provided coder. Be aware this is not the standard timestamp value encoding. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor ELEMENT_CODER_ID_FIELD_NUMBER: builtins.int element_coder_id: builtins.str def __init__( self, *, element_coder_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["element_coder_id", b"element_coder_id"]) -> None: ... global___OrderedListStateSpec = OrderedListStateSpec @typing_extensions.final class CombiningStateSpec(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ACCUMULATOR_CODER_ID_FIELD_NUMBER: builtins.int COMBINE_FN_FIELD_NUMBER: builtins.int accumulator_coder_id: builtins.str @property def combine_fn(self) -> global___FunctionSpec: ... def __init__( self, *, accumulator_coder_id: builtins.str | None = ..., combine_fn: global___FunctionSpec | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["combine_fn", b"combine_fn"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["accumulator_coder_id", b"accumulator_coder_id", "combine_fn", b"combine_fn"]) -> None: ... global___CombiningStateSpec = CombiningStateSpec @typing_extensions.final class MapStateSpec(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_CODER_ID_FIELD_NUMBER: builtins.int VALUE_CODER_ID_FIELD_NUMBER: builtins.int key_coder_id: builtins.str value_coder_id: builtins.str def __init__( self, *, key_coder_id: builtins.str | None = ..., value_coder_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key_coder_id", b"key_coder_id", "value_coder_id", b"value_coder_id"]) -> None: ... global___MapStateSpec = MapStateSpec @typing_extensions.final class MultimapStateSpec(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_CODER_ID_FIELD_NUMBER: builtins.int VALUE_CODER_ID_FIELD_NUMBER: builtins.int key_coder_id: builtins.str value_coder_id: builtins.str def __init__( self, *, key_coder_id: builtins.str | None = ..., value_coder_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key_coder_id", b"key_coder_id", "value_coder_id", b"value_coder_id"]) -> None: ... global___MultimapStateSpec = MultimapStateSpec @typing_extensions.final class SetStateSpec(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ELEMENT_CODER_ID_FIELD_NUMBER: builtins.int element_coder_id: builtins.str def __init__( self, *, element_coder_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["element_coder_id", b"element_coder_id"]) -> None: ... global___SetStateSpec = SetStateSpec @typing_extensions.final class TimerFamilySpec(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor TIME_DOMAIN_FIELD_NUMBER: builtins.int TIMER_FAMILY_CODER_ID_FIELD_NUMBER: builtins.int time_domain: global___TimeDomain.Enum.ValueType timer_family_coder_id: builtins.str def __init__( self, *, time_domain: global___TimeDomain.Enum.ValueType | None = ..., timer_family_coder_id: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["time_domain", b"time_domain", "timer_family_coder_id", b"timer_family_coder_id"]) -> None: ... global___TimerFamilySpec = TimerFamilySpec @typing_extensions.final class IsBounded(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[IsBounded._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor UNSPECIFIED: IsBounded._Enum.ValueType # 0 UNBOUNDED: IsBounded._Enum.ValueType # 1 BOUNDED: IsBounded._Enum.ValueType # 2 class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... UNSPECIFIED: IsBounded.Enum.ValueType # 0 UNBOUNDED: IsBounded.Enum.ValueType # 1 BOUNDED: IsBounded.Enum.ValueType # 2 def __init__( self, ) -> None: ... global___IsBounded = IsBounded @typing_extensions.final class ReadPayload(google.protobuf.message.Message): """The payload for the primitive Read transform.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor SOURCE_FIELD_NUMBER: builtins.int IS_BOUNDED_FIELD_NUMBER: builtins.int @property def source(self) -> global___FunctionSpec: """(Required) The FunctionSpec of the source for this Read.""" is_bounded: global___IsBounded.Enum.ValueType """(Required) Whether the source is bounded or unbounded""" def __init__( self, *, source: global___FunctionSpec | None = ..., is_bounded: global___IsBounded.Enum.ValueType | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["source", b"source"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["is_bounded", b"is_bounded", "source", b"source"]) -> None: ... global___ReadPayload = ReadPayload @typing_extensions.final class WindowIntoPayload(google.protobuf.message.Message): """The payload for the WindowInto transform.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor WINDOW_FN_FIELD_NUMBER: builtins.int @property def window_fn(self) -> global___FunctionSpec: """(Required) The FunctionSpec of the WindowFn.""" def __init__( self, *, window_fn: global___FunctionSpec | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["window_fn", b"window_fn"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["window_fn", b"window_fn"]) -> None: ... global___WindowIntoPayload = WindowIntoPayload @typing_extensions.final class CombinePayload(google.protobuf.message.Message): """The payload for the special-but-not-primitive Combine transform.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor COMBINE_FN_FIELD_NUMBER: builtins.int ACCUMULATOR_CODER_ID_FIELD_NUMBER: builtins.int @property def combine_fn(self) -> global___FunctionSpec: """(Required) The FunctionSpec of the CombineFn.""" accumulator_coder_id: builtins.str """(Required) A reference to the Coder to use for accumulators of the CombineFn""" def __init__( self, *, combine_fn: global___FunctionSpec | None = ..., accumulator_coder_id: builtins.str | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["combine_fn", b"combine_fn"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["accumulator_coder_id", b"accumulator_coder_id", "combine_fn", b"combine_fn"]) -> None: ... global___CombinePayload = CombinePayload @typing_extensions.final class TestStreamPayload(google.protobuf.message.Message): """The payload for the test-only primitive TestStream""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class Event(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class AdvanceWatermark(google.protobuf.message.Message): """Advances the watermark to the specified timestamp.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor NEW_WATERMARK_FIELD_NUMBER: builtins.int TAG_FIELD_NUMBER: builtins.int new_watermark: builtins.int """(Required) The watermark in millisecond to advance to.""" tag: builtins.str """(Optional) The output watermark tag for a PCollection. If unspecified or with an empty string, this will default to the Main PCollection Output """ def __init__( self, *, new_watermark: builtins.int | None = ..., tag: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["new_watermark", b"new_watermark", "tag", b"tag"]) -> None: ... @typing_extensions.final class AdvanceProcessingTime(google.protobuf.message.Message): """Advances the processing time clock by the specified amount.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor ADVANCE_DURATION_FIELD_NUMBER: builtins.int advance_duration: builtins.int """(Required) The duration in millisecond to advance by.""" def __init__( self, *, advance_duration: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["advance_duration", b"advance_duration"]) -> None: ... @typing_extensions.final class AddElements(google.protobuf.message.Message): """Adds elements to the stream to be emitted.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor ELEMENTS_FIELD_NUMBER: builtins.int TAG_FIELD_NUMBER: builtins.int @property def elements(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___TestStreamPayload.TimestampedElement]: """(Required) The elements to add to the TestStream.""" tag: builtins.str """(Optional) The output PCollection tag to add these elements to. If unspecified or with an empty string, this will default to the Main PCollection Output. """ def __init__( self, *, elements: collections.abc.Iterable[global___TestStreamPayload.TimestampedElement] | None = ..., tag: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["elements", b"elements", "tag", b"tag"]) -> None: ... WATERMARK_EVENT_FIELD_NUMBER: builtins.int PROCESSING_TIME_EVENT_FIELD_NUMBER: builtins.int ELEMENT_EVENT_FIELD_NUMBER: builtins.int @property def watermark_event(self) -> global___TestStreamPayload.Event.AdvanceWatermark: ... @property def processing_time_event(self) -> global___TestStreamPayload.Event.AdvanceProcessingTime: ... @property def element_event(self) -> global___TestStreamPayload.Event.AddElements: ... def __init__( self, *, watermark_event: global___TestStreamPayload.Event.AdvanceWatermark | None = ..., processing_time_event: global___TestStreamPayload.Event.AdvanceProcessingTime | None = ..., element_event: global___TestStreamPayload.Event.AddElements | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["element_event", b"element_event", "event", b"event", "processing_time_event", b"processing_time_event", "watermark_event", b"watermark_event"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["element_event", b"element_event", "event", b"event", "processing_time_event", b"processing_time_event", "watermark_event", b"watermark_event"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["event", b"event"]) -> typing_extensions.Literal["watermark_event", "processing_time_event", "element_event"] | None: ... @typing_extensions.final class TimestampedElement(google.protobuf.message.Message): """A single element inside of the TestStream.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor ENCODED_ELEMENT_FIELD_NUMBER: builtins.int TIMESTAMP_FIELD_NUMBER: builtins.int encoded_element: builtins.bytes """(Required) The element encoded. Currently the TestStream only supports encoding primitives. """ timestamp: builtins.int """(Required) The event timestamp in millisecond of this element.""" def __init__( self, *, encoded_element: builtins.bytes | None = ..., timestamp: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["encoded_element", b"encoded_element", "timestamp", b"timestamp"]) -> None: ... CODER_ID_FIELD_NUMBER: builtins.int EVENTS_FIELD_NUMBER: builtins.int ENDPOINT_FIELD_NUMBER: builtins.int coder_id: builtins.str """(Required) the coder for elements in the TestStream events""" @property def events(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___TestStreamPayload.Event]: """(Optional) If specified, the TestStream will replay these events.""" @property def endpoint(self) -> org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor: """(Optional) If specified, points to a TestStreamService to be used to retrieve events. """ def __init__( self, *, coder_id: builtins.str | None = ..., events: collections.abc.Iterable[global___TestStreamPayload.Event] | None = ..., endpoint: org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["endpoint", b"endpoint"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["coder_id", b"coder_id", "endpoint", b"endpoint", "events", b"events"]) -> None: ... global___TestStreamPayload = TestStreamPayload @typing_extensions.final class WriteFilesPayload(google.protobuf.message.Message): """The payload for the special-but-not-primitive WriteFiles transform.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class SideInputsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> global___SideInput: ... def __init__( self, *, key: builtins.str | None = ..., value: global___SideInput | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... SINK_FIELD_NUMBER: builtins.int FORMAT_FUNCTION_FIELD_NUMBER: builtins.int WINDOWED_WRITES_FIELD_NUMBER: builtins.int RUNNER_DETERMINED_SHARDING_FIELD_NUMBER: builtins.int SIDE_INPUTS_FIELD_NUMBER: builtins.int AUTO_SHARDED_FIELD_NUMBER: builtins.int @property def sink(self) -> global___FunctionSpec: """(Required) The FunctionSpec of the FileBasedSink.""" @property def format_function(self) -> global___FunctionSpec: """(Required) The format function.""" windowed_writes: builtins.bool runner_determined_sharding: builtins.bool @property def side_inputs(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___SideInput]: ... auto_sharded: builtins.bool """This is different from runner based sharding. This is done by the runner backend, where as runner_determined_sharding is by the runner translator """ def __init__( self, *, sink: global___FunctionSpec | None = ..., format_function: global___FunctionSpec | None = ..., windowed_writes: builtins.bool | None = ..., runner_determined_sharding: builtins.bool | None = ..., side_inputs: collections.abc.Mapping[builtins.str, global___SideInput] | None = ..., auto_sharded: builtins.bool | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["format_function", b"format_function", "sink", b"sink"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["auto_sharded", b"auto_sharded", "format_function", b"format_function", "runner_determined_sharding", b"runner_determined_sharding", "side_inputs", b"side_inputs", "sink", b"sink", "windowed_writes", b"windowed_writes"]) -> None: ... global___WriteFilesPayload = WriteFilesPayload @typing_extensions.final class PubSubReadPayload(google.protobuf.message.Message): """Payload used by Google Cloud Pub/Sub read transform. This can be used by runners that wish to override Beam Pub/Sub read transform with a native implementation. The SDK should guarantee that only one of topic, subscription, topic_runtime_overridden and subscription_runtime_overridden is set. The output of PubSubReadPayload should be bytes of serialized PubsubMessage proto if with_attributes == true. Otherwise, the bytes is the raw payload. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TOPIC_FIELD_NUMBER: builtins.int SUBSCRIPTION_FIELD_NUMBER: builtins.int TIMESTAMP_ATTRIBUTE_FIELD_NUMBER: builtins.int ID_ATTRIBUTE_FIELD_NUMBER: builtins.int WITH_ATTRIBUTES_FIELD_NUMBER: builtins.int TOPIC_RUNTIME_OVERRIDDEN_FIELD_NUMBER: builtins.int SUBSCRIPTION_RUNTIME_OVERRIDDEN_FIELD_NUMBER: builtins.int topic: builtins.str """Topic to read from. Exactly one of topic or subscription should be set. Topic format is: /topics/project_id/subscription_name """ subscription: builtins.str """Subscription to read from. Exactly one of topic or subscription should be set. Subscription format is: /subscriptions/project_id/subscription_name """ timestamp_attribute: builtins.str """Attribute that provides element timestamps.""" id_attribute: builtins.str """Attribute to be used for uniquely identifying messages.""" with_attributes: builtins.bool """If true, reads Pub/Sub payload as well as attributes. If false, reads only the payload.""" topic_runtime_overridden: builtins.str """If set, the topic is expected to be provided during runtime.""" subscription_runtime_overridden: builtins.str """If set, the subscription that is expected to be provided during runtime.""" def __init__( self, *, topic: builtins.str | None = ..., subscription: builtins.str | None = ..., timestamp_attribute: builtins.str | None = ..., id_attribute: builtins.str | None = ..., with_attributes: builtins.bool | None = ..., topic_runtime_overridden: builtins.str | None = ..., subscription_runtime_overridden: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["id_attribute", b"id_attribute", "subscription", b"subscription", "subscription_runtime_overridden", b"subscription_runtime_overridden", "timestamp_attribute", b"timestamp_attribute", "topic", b"topic", "topic_runtime_overridden", b"topic_runtime_overridden", "with_attributes", b"with_attributes"]) -> None: ... global___PubSubReadPayload = PubSubReadPayload @typing_extensions.final class PubSubWritePayload(google.protobuf.message.Message): """Payload used by Google Cloud Pub/Sub write transform. This can be used by runners that wish to override Beam Pub/Sub write transform with a native implementation. The SDK should guarantee that only one of topic and topic_runtime_overridden is set. The output of PubSubWritePayload should be bytes if serialized PubsubMessage proto. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TOPIC_FIELD_NUMBER: builtins.int TIMESTAMP_ATTRIBUTE_FIELD_NUMBER: builtins.int ID_ATTRIBUTE_FIELD_NUMBER: builtins.int TOPIC_RUNTIME_OVERRIDDEN_FIELD_NUMBER: builtins.int topic: builtins.str """Topic to write to. Topic format is: /topics/project_id/subscription_name """ timestamp_attribute: builtins.str """Attribute that provides element timestamps.""" id_attribute: builtins.str """Attribute that uniquely identify messages.""" topic_runtime_overridden: builtins.str """If set, the topic is expected to be provided during runtime.""" def __init__( self, *, topic: builtins.str | None = ..., timestamp_attribute: builtins.str | None = ..., id_attribute: builtins.str | None = ..., topic_runtime_overridden: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["id_attribute", b"id_attribute", "timestamp_attribute", b"timestamp_attribute", "topic", b"topic", "topic_runtime_overridden", b"topic_runtime_overridden"]) -> None: ... global___PubSubWritePayload = PubSubWritePayload @typing_extensions.final class GroupIntoBatchesPayload(google.protobuf.message.Message): """Payload for GroupIntoBatches composite transform.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor BATCH_SIZE_FIELD_NUMBER: builtins.int BATCH_SIZE_BYTES_FIELD_NUMBER: builtins.int MAX_BUFFERING_DURATION_MILLIS_FIELD_NUMBER: builtins.int batch_size: builtins.int """Max size of a batch.""" batch_size_bytes: builtins.int """Max byte size of a batch in element.""" max_buffering_duration_millis: builtins.int """(Optional) Max duration a batch is allowed to be cached in states.""" def __init__( self, *, batch_size: builtins.int | None = ..., batch_size_bytes: builtins.int | None = ..., max_buffering_duration_millis: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["batch_size", b"batch_size", "batch_size_bytes", b"batch_size_bytes", "max_buffering_duration_millis", b"max_buffering_duration_millis"]) -> None: ... global___GroupIntoBatchesPayload = GroupIntoBatchesPayload @typing_extensions.final class RedistributePayload(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ALLOW_DUPLICATES_FIELD_NUMBER: builtins.int allow_duplicates: builtins.bool def __init__( self, *, allow_duplicates: builtins.bool | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["allow_duplicates", b"allow_duplicates"]) -> None: ... global___RedistributePayload = RedistributePayload @typing_extensions.final class Coder(google.protobuf.message.Message): """A coder, the binary format for serialization and deserialization of data in a pipeline. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor SPEC_FIELD_NUMBER: builtins.int COMPONENT_CODER_IDS_FIELD_NUMBER: builtins.int @property def spec(self) -> global___FunctionSpec: """(Required) A specification for the coder, as a URN plus parameters. This may be a cross-language agreed-upon format, or it may be a "custom coder" that can only be used by a particular SDK. It does not include component coders, as it is beneficial for these to be comprehensible to a runner regardless of whether the binary format is agreed-upon. """ @property def component_coder_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """(Optional) If this coder is parametric, such as ListCoder(VarIntCoder), this is a list of the components. In order for encodings to be identical, the FunctionSpec and all components must be identical, recursively. """ def __init__( self, *, spec: global___FunctionSpec | None = ..., component_coder_ids: collections.abc.Iterable[builtins.str] | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["spec", b"spec"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["component_coder_ids", b"component_coder_ids", "spec", b"spec"]) -> None: ... global___Coder = Coder @typing_extensions.final class StandardCoders(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardCoders._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor BYTES: StandardCoders._Enum.ValueType # 0 """Components: None""" STRING_UTF8: StandardCoders._Enum.ValueType # 10 """Components: None""" KV: StandardCoders._Enum.ValueType # 1 """Components: The key and value coder, in that order.""" BOOL: StandardCoders._Enum.ValueType # 12 """Components: None""" VARINT: StandardCoders._Enum.ValueType # 2 """Variable length Encodes a 64-bit integer. Components: None """ DOUBLE: StandardCoders._Enum.ValueType # 11 """Encodes the floating point value as a big-endian 64-bit integer according to the IEEE 754 double format bit layout. Components: None """ ITERABLE: StandardCoders._Enum.ValueType # 3 """Encodes an iterable of elements. The encoding for an iterable [e1...eN] of known length N is fixed32(N) encode(e1) encode(e2) encode(e3) ... encode(eN) If the length is unknown, it is batched up into groups of size b1..bM and encoded as fixed32(-1) varInt64(b1) encode(e1) encode(e2) ... encode(e_b1) varInt64(b2) encode(e_(b1+1)) encode(e_(b1+2)) ... encode(e_(b1+b2)) ... varInt64(bM) encode(e_(N-bM+1)) encode(e_(N-bM+2)) ... encode(eN) varInt64(0) Components: Coder for a single element. """ TIMER: StandardCoders._Enum.ValueType # 4 """Encodes a timer containing a user key, a dynamic timer tag, a clear bit, a fire timestamp, a hold timestamp, the windows and the paneinfo. The encoding is represented as: user key - user defined key, uses the component coder. dynamic timer tag - a string which identifies a timer. windows - uses component coders. clear bit - a boolean set for clearing the timer. fire timestamp - a big endian 8 byte integer representing millis-since-epoch. The encoded representation is shifted so that the byte representation of negative values are lexicographically ordered before the byte representation of positive values. This is typically done by subtracting -9223372036854775808 from the value and encoding it as a signed big endian integer. Example values: -9223372036854775808: 00 00 00 00 00 00 00 00 -255: 7F FF FF FF FF FF FF 01 -1: 7F FF FF FF FF FF FF FF 0: 80 00 00 00 00 00 00 00 1: 80 00 00 00 00 00 00 01 256: 80 00 00 00 00 00 01 00 9223372036854775807: FF FF FF FF FF FF FF FF hold timestamp - similar to the fire timestamp. paneinfo - similar to the paneinfo of the windowed_value. Components: Coder for the key and windows. """ INTERVAL_WINDOW: StandardCoders._Enum.ValueType # 5 """ The following coders are typically not specified manually by the user, but are used at runtime and must be supported by every SDK. Components: None """ LENGTH_PREFIX: StandardCoders._Enum.ValueType # 6 """Components: The coder to attach a length prefix to""" GLOBAL_WINDOW: StandardCoders._Enum.ValueType # 7 """Components: None""" WINDOWED_VALUE: StandardCoders._Enum.ValueType # 8 """Encodes an element, the windows it is in, the timestamp of the element, and the pane of the element. The encoding is represented as: timestamp windows pane element timestamp - A big endian 8 byte integer representing millis-since-epoch. The encoded representation is shifted so that the byte representation of negative values are lexicographically ordered before the byte representation of positive values. This is typically done by subtracting -9223372036854775808 from the value and encoding it as a signed big endian integer. Example values: -9223372036854775808: 00 00 00 00 00 00 00 00 -255: 7F FF FF FF FF FF FF 01 -1: 7F FF FF FF FF FF FF FF 0: 80 00 00 00 00 00 00 00 1: 80 00 00 00 00 00 00 01 256: 80 00 00 00 00 00 01 00 9223372036854775807: FF FF FF FF FF FF FF FF windows - The windows are encoded using the beam:coder:iterable:v1 format, where the windows are encoded using the supplied window coder. pane - The first byte of the pane info determines which type of encoding is used, as well as the is_first, is_last and timing fields. If this byte is bits [0 1 2 3 4 5 6 7], then: * bits [0 1 2 3] determine the encoding as follows: 0000 - The entire pane info is encoded as a single byte. The is_first, is_last, and timing fields are encoded as below, and the index and non-speculative index are both zero (and hence are not encoded here). 0001 - The pane info is encoded as this byte plus a single VarInt encoded integer representing the pane index. The non-speculative index can be derived as follows: -1 if the pane is early, otherwise equal to index. 0010 - The pane info is encoded as this byte plus two VarInt encoded integers representing the pane index and non-speculative index respectively. * bits [4 5] encode the timing as follows: 00 - early 01 - on time 10 - late 11 - unknown * bit 6 is 1 if this is the last pane, 0 otherwise. Commonly set with `byte |= 0x02` * bit 7 is 1 if this is the first pane, 0 otherwise. Commonly set with `byte |= 0x01` element - The element incoded using the supplied element coder. Components: The element coder and the window coder, in that order. """ PARAM_WINDOWED_VALUE: StandardCoders._Enum.ValueType # 14 """A windowed value coder with parameterized timestamp, windows and pane info. Encodes an element with only the value of the windowed value. Decodes the value and assigns the parameterized timestamp, windows and pane info to the windowed value. Components: The element coder and the window coder, in that order The payload of this coder is an encoded windowed value using the beam:coder:windowed_value:v1 coder parameterized by a beam:coder:bytes:v1 element coder and the window coder that this param_windowed_value coder uses. """ STATE_BACKED_ITERABLE: StandardCoders._Enum.ValueType # 9 """Encodes an iterable of elements, some of which may be stored elsewhere. The encoding for a state-backed iterable is the same as that for an iterable, but the final varInt64(0) terminating the set of batches may instead be replaced by varInt64(-1) varInt64(len(token)) token where token is an opaque byte string that can be used to fetch the remainder of the iterable (e.g. over the state API). Components: Coder for a single element. """ CUSTOM_WINDOW: StandardCoders._Enum.ValueType # 16 """Encodes an arbitrary user defined window and its max timestamp (inclusive). The encoding format is: maxTimestamp window maxTimestamp - A big endian 8 byte integer representing millis-since-epoch. The encoded representation is shifted so that the byte representation of negative values are lexicographically ordered before the byte representation of positive values. This is typically done by subtracting -9223372036854775808 from the value and encoding it as a signed big endian integer. Example values: -9223372036854775808: 00 00 00 00 00 00 00 00 -255: 7F FF FF FF FF FF FF 01 -1: 7F FF FF FF FF FF FF FF 0: 80 00 00 00 00 00 00 00 1: 80 00 00 00 00 00 00 01 256: 80 00 00 00 00 00 01 00 9223372036854775807: FF FF FF FF FF FF FF FF window - the window is encoded using the supplied window coder. Components: Coder for the custom window type. """ ROW: StandardCoders._Enum.ValueType # 13 """Additional Standard Coders -------------------------- The following coders are not required to be implemented for an SDK or runner to support the Beam model, but enable users to take advantage of schema-aware transforms. Encodes a "row", an element with a known schema, defined by an instance of Schema from schema.proto. A row is encoded as the concatenation of: - The number of attributes in the schema, encoded with beam:coder:varint:v1. This makes it possible to detect certain allowed schema changes (appending or removing columns) in long-running streaming pipelines. - A byte array representing a packed bitset indicating null fields (a 1 indicating a null) encoded with beam:coder:bytes:v1. The unused bits in the last byte must be set to 0. If there are no nulls an empty byte array is encoded. The two-byte bitset (not including the lenghth-prefix) for the row [NULL, 0, 0, 0, NULL, 0, 0, NULL, 0, NULL] would be [0b10010001, 0b00000010] - An encoding for each non-null field, concatenated together. Schema types are mapped to coders as follows: AtomicType: BYTE: not yet a standard coder (https://github.com/apache/beam/issues/19815) INT16: not yet a standard coder (https://github.com/apache/beam/issues/19815) INT32: beam:coder:varint:v1 INT64: beam:coder:varint:v1 FLOAT: not yet a standard coder (https://github.com/apache/beam/issues/19815) DOUBLE: beam:coder:double:v1 STRING: beam:coder:string_utf8:v1 BOOLEAN: beam:coder:bool:v1 BYTES: beam:coder:bytes:v1 ArrayType: beam:coder:iterable:v1 (always has a known length) MapType: not a standard coder, specification defined below. RowType: beam:coder:row:v1 LogicalType: Uses the coder for its representation. The MapType is encoded by: - An INT32 representing the size of the map (N) - Followed by N interleaved keys and values, encoded with their corresponding coder. Nullable types in container types (ArrayType, MapType) per the encoding described for general Nullable types below. Logical types understood by all SDKs should be defined in schema.proto. Example of well known logical types: beam:logical_type:schema:v1 - Representation type: BYTES - A Beam Schema stored as a serialized proto. The payload for RowCoder is an instance of Schema. Components: None """ SCHEMA: StandardCoders._Enum.ValueType # 18 """Similar to ROW above, but for arbitrary types that can be converted to and from row objects, which can then be encoded with a schema. """ SHARDED_KEY: StandardCoders._Enum.ValueType # 15 """Encodes a user key and a shard id which is an opaque byte string. The encoding for a sharded key consists of a shard id byte string and the encoded user key in the following order: - shard id using beam:coder:bytes:v1 - encoded user key Examples: user key with an empty shard id 0x00 encode(user_key) user key with a shard id taking up two bytes. 0x02 0x11 0x22 encode(user_key) Components: the user key coder. """ NULLABLE: StandardCoders._Enum.ValueType # 17 """Wraps a coder of a potentially null value A Nullable Type is encoded by: - A one byte null indicator, 0x00 for null values, or 0x01 for present values. - For present values the null indicator is followed by the value encoded with it's corresponding coder. Components: single coder for the value """ class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... BYTES: StandardCoders.Enum.ValueType # 0 """Components: None""" STRING_UTF8: StandardCoders.Enum.ValueType # 10 """Components: None""" KV: StandardCoders.Enum.ValueType # 1 """Components: The key and value coder, in that order.""" BOOL: StandardCoders.Enum.ValueType # 12 """Components: None""" VARINT: StandardCoders.Enum.ValueType # 2 """Variable length Encodes a 64-bit integer. Components: None """ DOUBLE: StandardCoders.Enum.ValueType # 11 """Encodes the floating point value as a big-endian 64-bit integer according to the IEEE 754 double format bit layout. Components: None """ ITERABLE: StandardCoders.Enum.ValueType # 3 """Encodes an iterable of elements. The encoding for an iterable [e1...eN] of known length N is fixed32(N) encode(e1) encode(e2) encode(e3) ... encode(eN) If the length is unknown, it is batched up into groups of size b1..bM and encoded as fixed32(-1) varInt64(b1) encode(e1) encode(e2) ... encode(e_b1) varInt64(b2) encode(e_(b1+1)) encode(e_(b1+2)) ... encode(e_(b1+b2)) ... varInt64(bM) encode(e_(N-bM+1)) encode(e_(N-bM+2)) ... encode(eN) varInt64(0) Components: Coder for a single element. """ TIMER: StandardCoders.Enum.ValueType # 4 """Encodes a timer containing a user key, a dynamic timer tag, a clear bit, a fire timestamp, a hold timestamp, the windows and the paneinfo. The encoding is represented as: user key - user defined key, uses the component coder. dynamic timer tag - a string which identifies a timer. windows - uses component coders. clear bit - a boolean set for clearing the timer. fire timestamp - a big endian 8 byte integer representing millis-since-epoch. The encoded representation is shifted so that the byte representation of negative values are lexicographically ordered before the byte representation of positive values. This is typically done by subtracting -9223372036854775808 from the value and encoding it as a signed big endian integer. Example values: -9223372036854775808: 00 00 00 00 00 00 00 00 -255: 7F FF FF FF FF FF FF 01 -1: 7F FF FF FF FF FF FF FF 0: 80 00 00 00 00 00 00 00 1: 80 00 00 00 00 00 00 01 256: 80 00 00 00 00 00 01 00 9223372036854775807: FF FF FF FF FF FF FF FF hold timestamp - similar to the fire timestamp. paneinfo - similar to the paneinfo of the windowed_value. Components: Coder for the key and windows. """ INTERVAL_WINDOW: StandardCoders.Enum.ValueType # 5 """ The following coders are typically not specified manually by the user, but are used at runtime and must be supported by every SDK. Components: None """ LENGTH_PREFIX: StandardCoders.Enum.ValueType # 6 """Components: The coder to attach a length prefix to""" GLOBAL_WINDOW: StandardCoders.Enum.ValueType # 7 """Components: None""" WINDOWED_VALUE: StandardCoders.Enum.ValueType # 8 """Encodes an element, the windows it is in, the timestamp of the element, and the pane of the element. The encoding is represented as: timestamp windows pane element timestamp - A big endian 8 byte integer representing millis-since-epoch. The encoded representation is shifted so that the byte representation of negative values are lexicographically ordered before the byte representation of positive values. This is typically done by subtracting -9223372036854775808 from the value and encoding it as a signed big endian integer. Example values: -9223372036854775808: 00 00 00 00 00 00 00 00 -255: 7F FF FF FF FF FF FF 01 -1: 7F FF FF FF FF FF FF FF 0: 80 00 00 00 00 00 00 00 1: 80 00 00 00 00 00 00 01 256: 80 00 00 00 00 00 01 00 9223372036854775807: FF FF FF FF FF FF FF FF windows - The windows are encoded using the beam:coder:iterable:v1 format, where the windows are encoded using the supplied window coder. pane - The first byte of the pane info determines which type of encoding is used, as well as the is_first, is_last and timing fields. If this byte is bits [0 1 2 3 4 5 6 7], then: * bits [0 1 2 3] determine the encoding as follows: 0000 - The entire pane info is encoded as a single byte. The is_first, is_last, and timing fields are encoded as below, and the index and non-speculative index are both zero (and hence are not encoded here). 0001 - The pane info is encoded as this byte plus a single VarInt encoded integer representing the pane index. The non-speculative index can be derived as follows: -1 if the pane is early, otherwise equal to index. 0010 - The pane info is encoded as this byte plus two VarInt encoded integers representing the pane index and non-speculative index respectively. * bits [4 5] encode the timing as follows: 00 - early 01 - on time 10 - late 11 - unknown * bit 6 is 1 if this is the last pane, 0 otherwise. Commonly set with `byte |= 0x02` * bit 7 is 1 if this is the first pane, 0 otherwise. Commonly set with `byte |= 0x01` element - The element incoded using the supplied element coder. Components: The element coder and the window coder, in that order. """ PARAM_WINDOWED_VALUE: StandardCoders.Enum.ValueType # 14 """A windowed value coder with parameterized timestamp, windows and pane info. Encodes an element with only the value of the windowed value. Decodes the value and assigns the parameterized timestamp, windows and pane info to the windowed value. Components: The element coder and the window coder, in that order The payload of this coder is an encoded windowed value using the beam:coder:windowed_value:v1 coder parameterized by a beam:coder:bytes:v1 element coder and the window coder that this param_windowed_value coder uses. """ STATE_BACKED_ITERABLE: StandardCoders.Enum.ValueType # 9 """Encodes an iterable of elements, some of which may be stored elsewhere. The encoding for a state-backed iterable is the same as that for an iterable, but the final varInt64(0) terminating the set of batches may instead be replaced by varInt64(-1) varInt64(len(token)) token where token is an opaque byte string that can be used to fetch the remainder of the iterable (e.g. over the state API). Components: Coder for a single element. """ CUSTOM_WINDOW: StandardCoders.Enum.ValueType # 16 """Encodes an arbitrary user defined window and its max timestamp (inclusive). The encoding format is: maxTimestamp window maxTimestamp - A big endian 8 byte integer representing millis-since-epoch. The encoded representation is shifted so that the byte representation of negative values are lexicographically ordered before the byte representation of positive values. This is typically done by subtracting -9223372036854775808 from the value and encoding it as a signed big endian integer. Example values: -9223372036854775808: 00 00 00 00 00 00 00 00 -255: 7F FF FF FF FF FF FF 01 -1: 7F FF FF FF FF FF FF FF 0: 80 00 00 00 00 00 00 00 1: 80 00 00 00 00 00 00 01 256: 80 00 00 00 00 00 01 00 9223372036854775807: FF FF FF FF FF FF FF FF window - the window is encoded using the supplied window coder. Components: Coder for the custom window type. """ ROW: StandardCoders.Enum.ValueType # 13 """Additional Standard Coders -------------------------- The following coders are not required to be implemented for an SDK or runner to support the Beam model, but enable users to take advantage of schema-aware transforms. Encodes a "row", an element with a known schema, defined by an instance of Schema from schema.proto. A row is encoded as the concatenation of: - The number of attributes in the schema, encoded with beam:coder:varint:v1. This makes it possible to detect certain allowed schema changes (appending or removing columns) in long-running streaming pipelines. - A byte array representing a packed bitset indicating null fields (a 1 indicating a null) encoded with beam:coder:bytes:v1. The unused bits in the last byte must be set to 0. If there are no nulls an empty byte array is encoded. The two-byte bitset (not including the lenghth-prefix) for the row [NULL, 0, 0, 0, NULL, 0, 0, NULL, 0, NULL] would be [0b10010001, 0b00000010] - An encoding for each non-null field, concatenated together. Schema types are mapped to coders as follows: AtomicType: BYTE: not yet a standard coder (https://github.com/apache/beam/issues/19815) INT16: not yet a standard coder (https://github.com/apache/beam/issues/19815) INT32: beam:coder:varint:v1 INT64: beam:coder:varint:v1 FLOAT: not yet a standard coder (https://github.com/apache/beam/issues/19815) DOUBLE: beam:coder:double:v1 STRING: beam:coder:string_utf8:v1 BOOLEAN: beam:coder:bool:v1 BYTES: beam:coder:bytes:v1 ArrayType: beam:coder:iterable:v1 (always has a known length) MapType: not a standard coder, specification defined below. RowType: beam:coder:row:v1 LogicalType: Uses the coder for its representation. The MapType is encoded by: - An INT32 representing the size of the map (N) - Followed by N interleaved keys and values, encoded with their corresponding coder. Nullable types in container types (ArrayType, MapType) per the encoding described for general Nullable types below. Logical types understood by all SDKs should be defined in schema.proto. Example of well known logical types: beam:logical_type:schema:v1 - Representation type: BYTES - A Beam Schema stored as a serialized proto. The payload for RowCoder is an instance of Schema. Components: None """ SCHEMA: StandardCoders.Enum.ValueType # 18 """Similar to ROW above, but for arbitrary types that can be converted to and from row objects, which can then be encoded with a schema. """ SHARDED_KEY: StandardCoders.Enum.ValueType # 15 """Encodes a user key and a shard id which is an opaque byte string. The encoding for a sharded key consists of a shard id byte string and the encoded user key in the following order: - shard id using beam:coder:bytes:v1 - encoded user key Examples: user key with an empty shard id 0x00 encode(user_key) user key with a shard id taking up two bytes. 0x02 0x11 0x22 encode(user_key) Components: the user key coder. """ NULLABLE: StandardCoders.Enum.ValueType # 17 """Wraps a coder of a potentially null value A Nullable Type is encoded by: - A one byte null indicator, 0x00 for null values, or 0x01 for present values. - For present values the null indicator is followed by the value encoded with it's corresponding coder. Components: single coder for the value """ def __init__( self, ) -> None: ... global___StandardCoders = StandardCoders @typing_extensions.final class WindowingStrategy(google.protobuf.message.Message): """A windowing strategy describes the window function, triggering, allowed lateness, and accumulation mode for a PCollection. TODO: consider inlining field on PCollection """ DESCRIPTOR: google.protobuf.descriptor.Descriptor WINDOW_FN_FIELD_NUMBER: builtins.int MERGE_STATUS_FIELD_NUMBER: builtins.int WINDOW_CODER_ID_FIELD_NUMBER: builtins.int TRIGGER_FIELD_NUMBER: builtins.int ACCUMULATION_MODE_FIELD_NUMBER: builtins.int OUTPUT_TIME_FIELD_NUMBER: builtins.int CLOSING_BEHAVIOR_FIELD_NUMBER: builtins.int ALLOWED_LATENESS_FIELD_NUMBER: builtins.int ON_TIME_BEHAVIOR_FIELD_NUMBER: builtins.int ASSIGNS_TO_ONE_WINDOW_FIELD_NUMBER: builtins.int ENVIRONMENT_ID_FIELD_NUMBER: builtins.int @property def window_fn(self) -> global___FunctionSpec: """(Required) The FunctionSpec of the UDF that assigns windows, merges windows, and shifts timestamps before they are combined according to the OutputTime. """ merge_status: global___MergeStatus.Enum.ValueType """(Required) Whether or not the window fn is merging. This knowledge is required for many optimizations. """ window_coder_id: builtins.str """(Required) The coder for the windows of this PCollection.""" @property def trigger(self) -> global___Trigger: """(Required) The trigger to use when grouping this PCollection.""" accumulation_mode: global___AccumulationMode.Enum.ValueType """(Required) The accumulation mode indicates whether new panes are a full replacement for prior panes or whether they are deltas to be combined with other panes (the combine should correspond to whatever the upstream grouping transform is). """ output_time: global___OutputTime.Enum.ValueType """(Required) The OutputTime specifies, for a grouping transform, how to compute the aggregate timestamp. The window_fn will first possibly shift it later, then the OutputTime takes the max, min, or ignores it and takes the end of window. This is actually only for input to grouping transforms, but since they may be introduced in runner-specific ways, it is carried along with the windowing strategy. """ closing_behavior: global___ClosingBehavior.Enum.ValueType """(Required) Indicate when output should be omitted upon window expiration.""" allowed_lateness: builtins.int """(Required) The duration, in milliseconds, beyond the end of a window at which the window becomes droppable. """ on_time_behavior: global___OnTimeBehavior.Enum.ValueType """(Required) Indicate whether empty on-time panes should be omitted.""" assigns_to_one_window: builtins.bool """(Required) Whether or not the window fn assigns inputs to exactly one window This knowledge is required for some optimizations """ environment_id: builtins.str """(Optional) Environment where the current window_fn should be applied in. Runner that executes the pipeline may choose to override this if needed. If not specified, environment will be decided by the runner. """ def __init__( self, *, window_fn: global___FunctionSpec | None = ..., merge_status: global___MergeStatus.Enum.ValueType | None = ..., window_coder_id: builtins.str | None = ..., trigger: global___Trigger | None = ..., accumulation_mode: global___AccumulationMode.Enum.ValueType | None = ..., output_time: global___OutputTime.Enum.ValueType | None = ..., closing_behavior: global___ClosingBehavior.Enum.ValueType | None = ..., allowed_lateness: builtins.int | None = ..., on_time_behavior: global___OnTimeBehavior.Enum.ValueType | None = ..., assigns_to_one_window: builtins.bool | None = ..., environment_id: builtins.str | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["trigger", b"trigger", "window_fn", b"window_fn"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["accumulation_mode", b"accumulation_mode", "allowed_lateness", b"allowed_lateness", "assigns_to_one_window", b"assigns_to_one_window", "closing_behavior", b"closing_behavior", "environment_id", b"environment_id", "merge_status", b"merge_status", "on_time_behavior", b"on_time_behavior", "output_time", b"output_time", "trigger", b"trigger", "window_coder_id", b"window_coder_id", "window_fn", b"window_fn"]) -> None: ... global___WindowingStrategy = WindowingStrategy @typing_extensions.final class MergeStatus(google.protobuf.message.Message): """Whether or not a PCollection's WindowFn is non-merging, merging, or merging-but-already-merged, in which case a subsequent GroupByKey is almost always going to do something the user does not want """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[MergeStatus._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor UNSPECIFIED: MergeStatus._Enum.ValueType # 0 NON_MERGING: MergeStatus._Enum.ValueType # 1 """The WindowFn does not require merging. Examples: global window, FixedWindows, SlidingWindows """ NEEDS_MERGE: MergeStatus._Enum.ValueType # 2 """The WindowFn is merging and the PCollection has not had merging performed. Example: Sessions prior to a GroupByKey """ ALREADY_MERGED: MergeStatus._Enum.ValueType # 3 """The WindowFn is merging and the PCollection has had merging occur already. Example: Sessions after a GroupByKey """ class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... UNSPECIFIED: MergeStatus.Enum.ValueType # 0 NON_MERGING: MergeStatus.Enum.ValueType # 1 """The WindowFn does not require merging. Examples: global window, FixedWindows, SlidingWindows """ NEEDS_MERGE: MergeStatus.Enum.ValueType # 2 """The WindowFn is merging and the PCollection has not had merging performed. Example: Sessions prior to a GroupByKey """ ALREADY_MERGED: MergeStatus.Enum.ValueType # 3 """The WindowFn is merging and the PCollection has had merging occur already. Example: Sessions after a GroupByKey """ def __init__( self, ) -> None: ... global___MergeStatus = MergeStatus @typing_extensions.final class AccumulationMode(google.protobuf.message.Message): """Whether or not subsequent outputs of aggregations should be entire replacement values or just the aggregation of inputs received since the prior output. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[AccumulationMode._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor UNSPECIFIED: AccumulationMode._Enum.ValueType # 0 DISCARDING: AccumulationMode._Enum.ValueType # 1 """The aggregation is discarded when it is output""" ACCUMULATING: AccumulationMode._Enum.ValueType # 2 """The aggregation is accumulated across outputs""" RETRACTING: AccumulationMode._Enum.ValueType # 3 """The aggregation emits retractions when it is output""" class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... UNSPECIFIED: AccumulationMode.Enum.ValueType # 0 DISCARDING: AccumulationMode.Enum.ValueType # 1 """The aggregation is discarded when it is output""" ACCUMULATING: AccumulationMode.Enum.ValueType # 2 """The aggregation is accumulated across outputs""" RETRACTING: AccumulationMode.Enum.ValueType # 3 """The aggregation emits retractions when it is output""" def __init__( self, ) -> None: ... global___AccumulationMode = AccumulationMode @typing_extensions.final class ClosingBehavior(google.protobuf.message.Message): """Controls whether or not an aggregating transform should output data when a window expires. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ClosingBehavior._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor UNSPECIFIED: ClosingBehavior._Enum.ValueType # 0 EMIT_ALWAYS: ClosingBehavior._Enum.ValueType # 1 """Emit output when a window expires, whether or not there has been any new data since the last output. """ EMIT_IF_NONEMPTY: ClosingBehavior._Enum.ValueType # 2 """Only emit output when new data has arrives since the last output""" class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... UNSPECIFIED: ClosingBehavior.Enum.ValueType # 0 EMIT_ALWAYS: ClosingBehavior.Enum.ValueType # 1 """Emit output when a window expires, whether or not there has been any new data since the last output. """ EMIT_IF_NONEMPTY: ClosingBehavior.Enum.ValueType # 2 """Only emit output when new data has arrives since the last output""" def __init__( self, ) -> None: ... global___ClosingBehavior = ClosingBehavior @typing_extensions.final class OnTimeBehavior(google.protobuf.message.Message): """Controls whether or not an aggregating transform should output data when an on-time pane is empty. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[OnTimeBehavior._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor UNSPECIFIED: OnTimeBehavior._Enum.ValueType # 0 FIRE_ALWAYS: OnTimeBehavior._Enum.ValueType # 1 """Always fire the on-time pane. Even if there is no new data since the previous firing, an element will be produced. """ FIRE_IF_NONEMPTY: OnTimeBehavior._Enum.ValueType # 2 """Only fire the on-time pane if there is new data since the previous firing.""" class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... UNSPECIFIED: OnTimeBehavior.Enum.ValueType # 0 FIRE_ALWAYS: OnTimeBehavior.Enum.ValueType # 1 """Always fire the on-time pane. Even if there is no new data since the previous firing, an element will be produced. """ FIRE_IF_NONEMPTY: OnTimeBehavior.Enum.ValueType # 2 """Only fire the on-time pane if there is new data since the previous firing.""" def __init__( self, ) -> None: ... global___OnTimeBehavior = OnTimeBehavior @typing_extensions.final class OutputTime(google.protobuf.message.Message): """When a number of windowed, timestamped inputs are aggregated, the timestamp for the resulting output. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[OutputTime._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor UNSPECIFIED: OutputTime._Enum.ValueType # 0 END_OF_WINDOW: OutputTime._Enum.ValueType # 1 """The output has the timestamp of the end of the window.""" LATEST_IN_PANE: OutputTime._Enum.ValueType # 2 """The output has the latest timestamp of the input elements since the last output. """ EARLIEST_IN_PANE: OutputTime._Enum.ValueType # 3 """The output has the earliest timestamp of the input elements since the last output. """ class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... UNSPECIFIED: OutputTime.Enum.ValueType # 0 END_OF_WINDOW: OutputTime.Enum.ValueType # 1 """The output has the timestamp of the end of the window.""" LATEST_IN_PANE: OutputTime.Enum.ValueType # 2 """The output has the latest timestamp of the input elements since the last output. """ EARLIEST_IN_PANE: OutputTime.Enum.ValueType # 3 """The output has the earliest timestamp of the input elements since the last output. """ def __init__( self, ) -> None: ... global___OutputTime = OutputTime @typing_extensions.final class TimeDomain(google.protobuf.message.Message): """The different time domains in the Beam model.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[TimeDomain._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor UNSPECIFIED: TimeDomain._Enum.ValueType # 0 EVENT_TIME: TimeDomain._Enum.ValueType # 1 """Event time is time from the perspective of the data""" PROCESSING_TIME: TimeDomain._Enum.ValueType # 2 """Processing time is time from the perspective of the execution of your pipeline """ class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... UNSPECIFIED: TimeDomain.Enum.ValueType # 0 EVENT_TIME: TimeDomain.Enum.ValueType # 1 """Event time is time from the perspective of the data""" PROCESSING_TIME: TimeDomain.Enum.ValueType # 2 """Processing time is time from the perspective of the execution of your pipeline """ def __init__( self, ) -> None: ... global___TimeDomain = TimeDomain @typing_extensions.final class Trigger(google.protobuf.message.Message): """A small DSL for expressing when to emit new aggregations from a GroupByKey or CombinePerKey A trigger is described in terms of when it is _ready_ to permit output. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class AfterAll(google.protobuf.message.Message): """Ready when all subtriggers are ready.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor SUBTRIGGERS_FIELD_NUMBER: builtins.int @property def subtriggers(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Trigger]: ... def __init__( self, *, subtriggers: collections.abc.Iterable[global___Trigger] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["subtriggers", b"subtriggers"]) -> None: ... @typing_extensions.final class AfterAny(google.protobuf.message.Message): """Ready when any subtrigger is ready.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor SUBTRIGGERS_FIELD_NUMBER: builtins.int @property def subtriggers(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Trigger]: ... def __init__( self, *, subtriggers: collections.abc.Iterable[global___Trigger] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["subtriggers", b"subtriggers"]) -> None: ... @typing_extensions.final class AfterEach(google.protobuf.message.Message): """Starting with the first subtrigger, ready when the _current_ subtrigger is ready. After output, advances the current trigger by one. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor SUBTRIGGERS_FIELD_NUMBER: builtins.int @property def subtriggers(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Trigger]: ... def __init__( self, *, subtriggers: collections.abc.Iterable[global___Trigger] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["subtriggers", b"subtriggers"]) -> None: ... @typing_extensions.final class AfterEndOfWindow(google.protobuf.message.Message): """Ready after the input watermark is past the end of the window. May have implicitly-repeated subtriggers for early and late firings. When the end of the window is reached, the trigger transitions between the subtriggers. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor EARLY_FIRINGS_FIELD_NUMBER: builtins.int LATE_FIRINGS_FIELD_NUMBER: builtins.int @property def early_firings(self) -> global___Trigger: """(Optional) A trigger governing output prior to the end of the window.""" @property def late_firings(self) -> global___Trigger: """(Optional) A trigger governing output after the end of the window.""" def __init__( self, *, early_firings: global___Trigger | None = ..., late_firings: global___Trigger | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["early_firings", b"early_firings", "late_firings", b"late_firings"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["early_firings", b"early_firings", "late_firings", b"late_firings"]) -> None: ... @typing_extensions.final class AfterProcessingTime(google.protobuf.message.Message): """After input arrives, ready when the specified delay has passed.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor TIMESTAMP_TRANSFORMS_FIELD_NUMBER: builtins.int @property def timestamp_transforms(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___TimestampTransform]: """(Required) The transforms to apply to an arriving element's timestamp, in order """ def __init__( self, *, timestamp_transforms: collections.abc.Iterable[global___TimestampTransform] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["timestamp_transforms", b"timestamp_transforms"]) -> None: ... @typing_extensions.final class AfterSynchronizedProcessingTime(google.protobuf.message.Message): """Ready whenever upstream processing time has all caught up with the arrival time of an input element """ DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... @typing_extensions.final class Default(google.protobuf.message.Message): """The default trigger. Equivalent to AfterEndOfWindow { Late: Always }} but specially denoted to indicate the user did not alter the triggering. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... @typing_extensions.final class ElementCount(google.protobuf.message.Message): """Ready whenever the requisite number of input elements have arrived""" DESCRIPTOR: google.protobuf.descriptor.Descriptor ELEMENT_COUNT_FIELD_NUMBER: builtins.int element_count: builtins.int def __init__( self, *, element_count: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["element_count", b"element_count"]) -> None: ... @typing_extensions.final class Never(google.protobuf.message.Message): """Never ready. There will only be an ON_TIME final output at window expiration. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... @typing_extensions.final class Always(google.protobuf.message.Message): """Always ready. This can also be expressed as Repeat{ ElementCount(1) } but is more explicit. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor def __init__( self, ) -> None: ... @typing_extensions.final class OrFinally(google.protobuf.message.Message): """Ready whenever either of its subtriggers are ready, but finishes output when the finally subtrigger fires. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor MAIN_FIELD_NUMBER: builtins.int FINALLY_FIELD_NUMBER: builtins.int @property def main(self) -> global___Trigger: """(Required) Trigger governing main output; may fire repeatedly.""" def __init__( self, *, main: global___Trigger | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["finally", b"finally", "main", b"main"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["finally", b"finally", "main", b"main"]) -> None: ... @typing_extensions.final class Repeat(google.protobuf.message.Message): """Ready whenever the subtrigger is ready; resets state when the subtrigger completes. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor SUBTRIGGER_FIELD_NUMBER: builtins.int @property def subtrigger(self) -> global___Trigger: """(Require) Trigger that is run repeatedly.""" def __init__( self, *, subtrigger: global___Trigger | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["subtrigger", b"subtrigger"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["subtrigger", b"subtrigger"]) -> None: ... AFTER_ALL_FIELD_NUMBER: builtins.int AFTER_ANY_FIELD_NUMBER: builtins.int AFTER_EACH_FIELD_NUMBER: builtins.int AFTER_END_OF_WINDOW_FIELD_NUMBER: builtins.int AFTER_PROCESSING_TIME_FIELD_NUMBER: builtins.int AFTER_SYNCHRONIZED_PROCESSING_TIME_FIELD_NUMBER: builtins.int ALWAYS_FIELD_NUMBER: builtins.int DEFAULT_FIELD_NUMBER: builtins.int ELEMENT_COUNT_FIELD_NUMBER: builtins.int NEVER_FIELD_NUMBER: builtins.int OR_FINALLY_FIELD_NUMBER: builtins.int REPEAT_FIELD_NUMBER: builtins.int @property def after_all(self) -> global___Trigger.AfterAll: ... @property def after_any(self) -> global___Trigger.AfterAny: ... @property def after_each(self) -> global___Trigger.AfterEach: ... @property def after_end_of_window(self) -> global___Trigger.AfterEndOfWindow: ... @property def after_processing_time(self) -> global___Trigger.AfterProcessingTime: ... @property def after_synchronized_processing_time(self) -> global___Trigger.AfterSynchronizedProcessingTime: ... @property def always(self) -> global___Trigger.Always: ... @property def default(self) -> global___Trigger.Default: ... @property def element_count(self) -> global___Trigger.ElementCount: ... @property def never(self) -> global___Trigger.Never: ... @property def or_finally(self) -> global___Trigger.OrFinally: ... @property def repeat(self) -> global___Trigger.Repeat: ... def __init__( self, *, after_all: global___Trigger.AfterAll | None = ..., after_any: global___Trigger.AfterAny | None = ..., after_each: global___Trigger.AfterEach | None = ..., after_end_of_window: global___Trigger.AfterEndOfWindow | None = ..., after_processing_time: global___Trigger.AfterProcessingTime | None = ..., after_synchronized_processing_time: global___Trigger.AfterSynchronizedProcessingTime | None = ..., always: global___Trigger.Always | None = ..., default: global___Trigger.Default | None = ..., element_count: global___Trigger.ElementCount | None = ..., never: global___Trigger.Never | None = ..., or_finally: global___Trigger.OrFinally | None = ..., repeat: global___Trigger.Repeat | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["after_all", b"after_all", "after_any", b"after_any", "after_each", b"after_each", "after_end_of_window", b"after_end_of_window", "after_processing_time", b"after_processing_time", "after_synchronized_processing_time", b"after_synchronized_processing_time", "always", b"always", "default", b"default", "element_count", b"element_count", "never", b"never", "or_finally", b"or_finally", "repeat", b"repeat", "trigger", b"trigger"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["after_all", b"after_all", "after_any", b"after_any", "after_each", b"after_each", "after_end_of_window", b"after_end_of_window", "after_processing_time", b"after_processing_time", "after_synchronized_processing_time", b"after_synchronized_processing_time", "always", b"always", "default", b"default", "element_count", b"element_count", "never", b"never", "or_finally", b"or_finally", "repeat", b"repeat", "trigger", b"trigger"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["trigger", b"trigger"]) -> typing_extensions.Literal["after_all", "after_any", "after_each", "after_end_of_window", "after_processing_time", "after_synchronized_processing_time", "always", "default", "element_count", "never", "or_finally", "repeat"] | None: ... global___Trigger = Trigger @typing_extensions.final class TimestampTransform(google.protobuf.message.Message): """A specification for a transformation on a timestamp. Primarily used by AfterProcessingTime triggers to transform the arrival time of input to a target time for firing. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class Delay(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor DELAY_MILLIS_FIELD_NUMBER: builtins.int delay_millis: builtins.int """(Required) The delay, in milliseconds.""" def __init__( self, *, delay_millis: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["delay_millis", b"delay_millis"]) -> None: ... @typing_extensions.final class AlignTo(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor PERIOD_FIELD_NUMBER: builtins.int OFFSET_FIELD_NUMBER: builtins.int period: builtins.int """(Required) A duration to which delays should be quantized in milliseconds. """ offset: builtins.int """(Required) An offset from 0 for the quantization specified by alignment_size, in milliseconds """ def __init__( self, *, period: builtins.int | None = ..., offset: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["offset", b"offset", "period", b"period"]) -> None: ... DELAY_FIELD_NUMBER: builtins.int ALIGN_TO_FIELD_NUMBER: builtins.int @property def delay(self) -> global___TimestampTransform.Delay: ... @property def align_to(self) -> global___TimestampTransform.AlignTo: ... def __init__( self, *, delay: global___TimestampTransform.Delay | None = ..., align_to: global___TimestampTransform.AlignTo | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["align_to", b"align_to", "delay", b"delay", "timestamp_transform", b"timestamp_transform"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["align_to", b"align_to", "delay", b"delay", "timestamp_transform", b"timestamp_transform"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["timestamp_transform", b"timestamp_transform"]) -> typing_extensions.Literal["delay", "align_to"] | None: ... global___TimestampTransform = TimestampTransform @typing_extensions.final class SideInput(google.protobuf.message.Message): """A specification for how to "side input" a PCollection.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor ACCESS_PATTERN_FIELD_NUMBER: builtins.int VIEW_FN_FIELD_NUMBER: builtins.int WINDOW_MAPPING_FN_FIELD_NUMBER: builtins.int @property def access_pattern(self) -> global___FunctionSpec: """(Required) URN of the access pattern required by the `view_fn` to present the desired SDK-specific interface to a UDF. This access pattern defines the SDK harness <-> Runner Harness RPC interface for accessing a side input. See StandardSideInputTypes for an enumeration of all side input types defined. """ @property def view_fn(self) -> global___FunctionSpec: """(Required) The FunctionSpec of the UDF that adapts a particular access_pattern to a user-facing view type. For example, View.asSingleton() may include a `view_fn` that adapts a specially-designed multimap to a single value per window. """ @property def window_mapping_fn(self) -> global___FunctionSpec: """(Required) The FunctionSpec of the UDF that maps a main input window to a side input window. For example, when the main input is in fixed windows of one hour, this can specify that the side input should be accessed according to the day in which that hour falls. """ def __init__( self, *, access_pattern: global___FunctionSpec | None = ..., view_fn: global___FunctionSpec | None = ..., window_mapping_fn: global___FunctionSpec | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["access_pattern", b"access_pattern", "view_fn", b"view_fn", "window_mapping_fn", b"window_mapping_fn"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["access_pattern", b"access_pattern", "view_fn", b"view_fn", "window_mapping_fn", b"window_mapping_fn"]) -> None: ... global___SideInput = SideInput @typing_extensions.final class StandardArtifacts(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Types: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _TypesEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardArtifacts._Types.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor FILE: StandardArtifacts._Types.ValueType # 0 """A URN for locally-accessible artifact files. payload: ArtifactFilePayload """ URL: StandardArtifacts._Types.ValueType # 1 """A URN for artifacts described by URLs. payload: ArtifactUrlPayload """ EMBEDDED: StandardArtifacts._Types.ValueType # 2 """A URN for artifacts embedded in ArtifactInformation proto. payload: EmbeddedFilePayload. """ PYPI: StandardArtifacts._Types.ValueType # 3 """A URN for Python artifacts hosted on PYPI. payload: PypiPayload """ MAVEN: StandardArtifacts._Types.ValueType # 4 """A URN for Java artifacts hosted on a Maven repository. payload: MavenPayload """ DEFERRED: StandardArtifacts._Types.ValueType # 5 """A URN for deferred artifacts. payload: DeferredArtifactPayload """ class Types(_Types, metaclass=_TypesEnumTypeWrapper): ... FILE: StandardArtifacts.Types.ValueType # 0 """A URN for locally-accessible artifact files. payload: ArtifactFilePayload """ URL: StandardArtifacts.Types.ValueType # 1 """A URN for artifacts described by URLs. payload: ArtifactUrlPayload """ EMBEDDED: StandardArtifacts.Types.ValueType # 2 """A URN for artifacts embedded in ArtifactInformation proto. payload: EmbeddedFilePayload. """ PYPI: StandardArtifacts.Types.ValueType # 3 """A URN for Python artifacts hosted on PYPI. payload: PypiPayload """ MAVEN: StandardArtifacts.Types.ValueType # 4 """A URN for Java artifacts hosted on a Maven repository. payload: MavenPayload """ DEFERRED: StandardArtifacts.Types.ValueType # 5 """A URN for deferred artifacts. payload: DeferredArtifactPayload """ class _Roles: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _RolesEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardArtifacts._Roles.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor STAGING_TO: StandardArtifacts._Roles.ValueType # 0 """A URN for staging-to role. payload: ArtifactStagingToRolePayload """ PIP_REQUIREMENTS_FILE: StandardArtifacts._Roles.ValueType # 1 """A URN for pip-requirements-file role. payload: None """ GO_WORKER_BINARY: StandardArtifacts._Roles.ValueType # 2 """A URN for the Go worker binary role. This represents the executable for a Go SDK environment. A Go environment may have one such artifact with this role. payload: None """ class Roles(_Roles, metaclass=_RolesEnumTypeWrapper): ... STAGING_TO: StandardArtifacts.Roles.ValueType # 0 """A URN for staging-to role. payload: ArtifactStagingToRolePayload """ PIP_REQUIREMENTS_FILE: StandardArtifacts.Roles.ValueType # 1 """A URN for pip-requirements-file role. payload: None """ GO_WORKER_BINARY: StandardArtifacts.Roles.ValueType # 2 """A URN for the Go worker binary role. This represents the executable for a Go SDK environment. A Go environment may have one such artifact with this role. payload: None """ def __init__( self, ) -> None: ... global___StandardArtifacts = StandardArtifacts @typing_extensions.final class ArtifactFilePayload(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor PATH_FIELD_NUMBER: builtins.int SHA256_FIELD_NUMBER: builtins.int path: builtins.str """a string for an artifact file path e.g. "/tmp/foo.jar" """ sha256: builtins.str """The hex-encoded sha256 checksum of the artifact.""" def __init__( self, *, path: builtins.str | None = ..., sha256: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["path", b"path", "sha256", b"sha256"]) -> None: ... global___ArtifactFilePayload = ArtifactFilePayload @typing_extensions.final class ArtifactUrlPayload(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor URL_FIELD_NUMBER: builtins.int SHA256_FIELD_NUMBER: builtins.int url: builtins.str """a string for an artifact URL e.g. "https://.../foo.jar" or "gs://tmp/foo.jar" """ sha256: builtins.str """(Optional) The hex-encoded sha256 checksum of the artifact if available.""" def __init__( self, *, url: builtins.str | None = ..., sha256: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["sha256", b"sha256", "url", b"url"]) -> None: ... global___ArtifactUrlPayload = ArtifactUrlPayload @typing_extensions.final class EmbeddedFilePayload(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor DATA_FIELD_NUMBER: builtins.int data: builtins.bytes """raw data bytes for an embedded artifact""" def __init__( self, *, data: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["data", b"data"]) -> None: ... global___EmbeddedFilePayload = EmbeddedFilePayload @typing_extensions.final class PyPIPayload(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ARTIFACT_ID_FIELD_NUMBER: builtins.int VERSION_FIELD_NUMBER: builtins.int artifact_id: builtins.str """Pypi compatible artifact id e.g. "apache-beam" """ version: builtins.str """Pypi compatible version string.""" def __init__( self, *, artifact_id: builtins.str | None = ..., version: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["artifact_id", b"artifact_id", "version", b"version"]) -> None: ... global___PyPIPayload = PyPIPayload @typing_extensions.final class MavenPayload(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ARTIFACT_FIELD_NUMBER: builtins.int REPOSITORY_URL_FIELD_NUMBER: builtins.int artifact: builtins.str """A string specifying Maven artifact. The standard format is "groupId:artifactId:version[:packaging[:classifier]]" """ repository_url: builtins.str """(Optional) Repository URL. If not specified, Maven central is used by default.""" def __init__( self, *, artifact: builtins.str | None = ..., repository_url: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["artifact", b"artifact", "repository_url", b"repository_url"]) -> None: ... global___MavenPayload = MavenPayload @typing_extensions.final class DeferredArtifactPayload(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int DATA_FIELD_NUMBER: builtins.int key: builtins.str """A unique string identifier assigned by the creator of this payload. The creator may use this key to confirm whether they can parse the data. """ data: builtins.bytes """Data for deferred artifacts. Interpretation of bytes is delegated to the creator of this payload.""" def __init__( self, *, key: builtins.str | None = ..., data: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["data", b"data", "key", b"key"]) -> None: ... global___DeferredArtifactPayload = DeferredArtifactPayload @typing_extensions.final class ArtifactStagingToRolePayload(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor STAGED_NAME_FIELD_NUMBER: builtins.int staged_name: builtins.str """A generated staged name (relative path under staging directory).""" def __init__( self, *, staged_name: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["staged_name", b"staged_name"]) -> None: ... global___ArtifactStagingToRolePayload = ArtifactStagingToRolePayload @typing_extensions.final class ArtifactInformation(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor TYPE_URN_FIELD_NUMBER: builtins.int TYPE_PAYLOAD_FIELD_NUMBER: builtins.int ROLE_URN_FIELD_NUMBER: builtins.int ROLE_PAYLOAD_FIELD_NUMBER: builtins.int type_urn: builtins.str """A URN that describes the type of artifact""" type_payload: builtins.bytes role_urn: builtins.str """A URN that describes the role of artifact""" role_payload: builtins.bytes def __init__( self, *, type_urn: builtins.str | None = ..., type_payload: builtins.bytes | None = ..., role_urn: builtins.str | None = ..., role_payload: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["role_payload", b"role_payload", "role_urn", b"role_urn", "type_payload", b"type_payload", "type_urn", b"type_urn"]) -> None: ... global___ArtifactInformation = ArtifactInformation @typing_extensions.final class Environment(google.protobuf.message.Message): """An environment for executing UDFs. By default, an SDK container URL, but can also be a process forked by a command, or an externally managed process. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class ResourceHintsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str value: builtins.bytes def __init__( self, *, key: builtins.str | None = ..., value: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... URN_FIELD_NUMBER: builtins.int PAYLOAD_FIELD_NUMBER: builtins.int DISPLAY_DATA_FIELD_NUMBER: builtins.int CAPABILITIES_FIELD_NUMBER: builtins.int DEPENDENCIES_FIELD_NUMBER: builtins.int RESOURCE_HINTS_FIELD_NUMBER: builtins.int urn: builtins.str """(Required) The URN of the payload""" payload: builtins.bytes """(Optional) The data specifying any parameters to the URN. If the URN does not require any arguments, this may be omitted. """ @property def display_data(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___DisplayData]: """(Optional) Static display data for the environment. If there is none, it may be omitted. """ @property def capabilities(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """(Optional) A set of capabilities this environment supports. This is typically a list of common URNs designating coders, transforms, etc. that this environment understands (and a runner MAY use) despite not appearing in the pipeline proto. This may also be used to indicate support of optional protocols not tied to a concrete component. """ @property def dependencies(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ArtifactInformation]: """(Optional) artifact dependency information used for executing UDFs in this environment.""" @property def resource_hints(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.bytes]: """(Optional) A mapping of resource URNs to requested values. The encoding of the values is specified by the URN. Resource hints are advisory; a runner is free to ignore resource hints that it does not understand. """ def __init__( self, *, urn: builtins.str | None = ..., payload: builtins.bytes | None = ..., display_data: collections.abc.Iterable[global___DisplayData] | None = ..., capabilities: collections.abc.Iterable[builtins.str] | None = ..., dependencies: collections.abc.Iterable[global___ArtifactInformation] | None = ..., resource_hints: collections.abc.Mapping[builtins.str, builtins.bytes] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["capabilities", b"capabilities", "dependencies", b"dependencies", "display_data", b"display_data", "payload", b"payload", "resource_hints", b"resource_hints", "urn", b"urn"]) -> None: ... global___Environment = Environment @typing_extensions.final class StandardEnvironments(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Environments: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnvironmentsEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardEnvironments._Environments.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor DOCKER: StandardEnvironments._Environments.ValueType # 0 """A managed docker container to run user code. Payload should be DockerPayload. """ PROCESS: StandardEnvironments._Environments.ValueType # 1 """A managed native process to run user code. Payload should be ProcessPayload. """ EXTERNAL: StandardEnvironments._Environments.ValueType # 2 """An external non managed process to run user code. Payload should be ExternalPayload. """ DEFAULT: StandardEnvironments._Environments.ValueType # 3 """Used as a stub when context is missing a runner-provided default environment.""" ANYOF: StandardEnvironments._Environments.ValueType # 4 """A selection of equivalent fully-specified environments a runner may use. Note that this environment itself does not declare any dependencies or capabilities, as those may differ among the several alternatives. Payload should be AnyOfEnvironmentPayload. """ class Environments(_Environments, metaclass=_EnvironmentsEnumTypeWrapper): ... DOCKER: StandardEnvironments.Environments.ValueType # 0 """A managed docker container to run user code. Payload should be DockerPayload. """ PROCESS: StandardEnvironments.Environments.ValueType # 1 """A managed native process to run user code. Payload should be ProcessPayload. """ EXTERNAL: StandardEnvironments.Environments.ValueType # 2 """An external non managed process to run user code. Payload should be ExternalPayload. """ DEFAULT: StandardEnvironments.Environments.ValueType # 3 """Used as a stub when context is missing a runner-provided default environment.""" ANYOF: StandardEnvironments.Environments.ValueType # 4 """A selection of equivalent fully-specified environments a runner may use. Note that this environment itself does not declare any dependencies or capabilities, as those may differ among the several alternatives. Payload should be AnyOfEnvironmentPayload. """ def __init__( self, ) -> None: ... global___StandardEnvironments = StandardEnvironments @typing_extensions.final class DockerPayload(google.protobuf.message.Message): """The payload of a Docker image""" DESCRIPTOR: google.protobuf.descriptor.Descriptor CONTAINER_IMAGE_FIELD_NUMBER: builtins.int container_image: builtins.str """implicitly linux_amd64.""" def __init__( self, *, container_image: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["container_image", b"container_image"]) -> None: ... global___DockerPayload = DockerPayload @typing_extensions.final class ProcessPayload(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class EnvEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str value: builtins.str def __init__( self, *, key: builtins.str | None = ..., value: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... OS_FIELD_NUMBER: builtins.int ARCH_FIELD_NUMBER: builtins.int COMMAND_FIELD_NUMBER: builtins.int ENV_FIELD_NUMBER: builtins.int os: builtins.str """"linux", "darwin", ..""" arch: builtins.str """"amd64", ..""" command: builtins.str """process to execute""" @property def env(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: """Environment variables""" def __init__( self, *, os: builtins.str | None = ..., arch: builtins.str | None = ..., command: builtins.str | None = ..., env: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["arch", b"arch", "command", b"command", "env", b"env", "os", b"os"]) -> None: ... global___ProcessPayload = ProcessPayload @typing_extensions.final class ExternalPayload(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class ParamsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str value: builtins.str def __init__( self, *, key: builtins.str | None = ..., value: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... ENDPOINT_FIELD_NUMBER: builtins.int PARAMS_FIELD_NUMBER: builtins.int @property def endpoint(self) -> org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor: """Serving BeamFnExternalWorkerPool API.""" @property def params(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: """Arbitrary extra parameters to pass""" def __init__( self, *, endpoint: org.apache.beam.model.pipeline.v1.endpoints_pb2.ApiServiceDescriptor | None = ..., params: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["endpoint", b"endpoint"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["endpoint", b"endpoint", "params", b"params"]) -> None: ... global___ExternalPayload = ExternalPayload @typing_extensions.final class AnyOfEnvironmentPayload(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ENVIRONMENTS_FIELD_NUMBER: builtins.int @property def environments(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Environment]: """Each is fully contained (with their own dependencies, capabilities, etc.)""" def __init__( self, *, environments: collections.abc.Iterable[global___Environment] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["environments", b"environments"]) -> None: ... global___AnyOfEnvironmentPayload = AnyOfEnvironmentPayload @typing_extensions.final class StandardProtocols(google.protobuf.message.Message): """These URNs are used to indicate capabilities of environments that cannot simply be expressed as a component (such as a Coder or PTransform) that this environment understands. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardProtocols._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor LEGACY_PROGRESS_REPORTING: StandardProtocols._Enum.ValueType # 0 """Indicates support for progress reporting via the legacy Metrics proto.""" PROGRESS_REPORTING: StandardProtocols._Enum.ValueType # 1 """Indicates support for progress reporting via the new MonitoringInfo proto.""" WORKER_STATUS: StandardProtocols._Enum.ValueType # 2 """Indicates support for worker status protocol defined at https://s.apache.org/beam-fn-api-harness-status. """ MULTI_CORE_BUNDLE_PROCESSING: StandardProtocols._Enum.ValueType # 3 """Indicates this SDK can take advantage of multiple cores when processing concurrent process bundle requests. (Note that all SDKs must process an unbounded number of concurrent process bundle requests; this capability simply indicates this SDK can actually parallelize the work across multiple cores. """ SIBLING_WORKERS: StandardProtocols._Enum.ValueType # 5 """Indicates this SDK can cheaply spawn sibling workers (e.g. within the same container) to work around the fact that it cannot take advantage of multiple cores (i.e. MULTI_CORE_BUNDLE_PROCESSING is not set). """ HARNESS_MONITORING_INFOS: StandardProtocols._Enum.ValueType # 4 """Indicates that this SDK handles the InstructionRequest of type HarnessMonitoringInfosRequest. A request to provide full MonitoringInfo data associated with the entire SDK harness process, not specific to a bundle. """ CONTROL_REQUEST_ELEMENTS_EMBEDDING: StandardProtocols._Enum.ValueType # 6 """Indicates that this SDK can process elements embedded in the ProcessBundleRequest. See more about the protocol at https://s.apache.org/beam-fn-api-control-data-embedding """ STATE_CACHING: StandardProtocols._Enum.ValueType # 7 """Indicates that this SDK can cache user state and side inputs across bundle boundaries. This is a hint to runners that runners can rely on the SDKs ability to store the data in memory reducing the amount of memory used overall. """ DATA_SAMPLING: StandardProtocols._Enum.ValueType # 8 """Indicates that this SDK can sample in-flight elements. These samples can then be queried using the SampleDataRequest. Samples are uniquely associated with a PCollection. Meaning, samples are taken for each PCollection during bundle processing. This is disabled by default and enabled with the `enable_data_sampling` experiment. """ SDK_CONSUMING_RECEIVED_DATA: StandardProtocols._Enum.ValueType # 9 """Indicates whether the SDK sets the consuming_received_data field on the ProcessBundleProgressResponse messages. """ ORDERED_LIST_STATE: StandardProtocols._Enum.ValueType # 10 """Indicates whether the SDK supports ordered list state.""" ELEMENT_METADATA: StandardProtocols._Enum.ValueType # 11 """Indicates support for reading, writing and propagating Element's metadata""" MULTIMAP_STATE: StandardProtocols._Enum.ValueType # 12 """Indicates whether the SDK supports multimap state.""" class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... LEGACY_PROGRESS_REPORTING: StandardProtocols.Enum.ValueType # 0 """Indicates support for progress reporting via the legacy Metrics proto.""" PROGRESS_REPORTING: StandardProtocols.Enum.ValueType # 1 """Indicates support for progress reporting via the new MonitoringInfo proto.""" WORKER_STATUS: StandardProtocols.Enum.ValueType # 2 """Indicates support for worker status protocol defined at https://s.apache.org/beam-fn-api-harness-status. """ MULTI_CORE_BUNDLE_PROCESSING: StandardProtocols.Enum.ValueType # 3 """Indicates this SDK can take advantage of multiple cores when processing concurrent process bundle requests. (Note that all SDKs must process an unbounded number of concurrent process bundle requests; this capability simply indicates this SDK can actually parallelize the work across multiple cores. """ SIBLING_WORKERS: StandardProtocols.Enum.ValueType # 5 """Indicates this SDK can cheaply spawn sibling workers (e.g. within the same container) to work around the fact that it cannot take advantage of multiple cores (i.e. MULTI_CORE_BUNDLE_PROCESSING is not set). """ HARNESS_MONITORING_INFOS: StandardProtocols.Enum.ValueType # 4 """Indicates that this SDK handles the InstructionRequest of type HarnessMonitoringInfosRequest. A request to provide full MonitoringInfo data associated with the entire SDK harness process, not specific to a bundle. """ CONTROL_REQUEST_ELEMENTS_EMBEDDING: StandardProtocols.Enum.ValueType # 6 """Indicates that this SDK can process elements embedded in the ProcessBundleRequest. See more about the protocol at https://s.apache.org/beam-fn-api-control-data-embedding """ STATE_CACHING: StandardProtocols.Enum.ValueType # 7 """Indicates that this SDK can cache user state and side inputs across bundle boundaries. This is a hint to runners that runners can rely on the SDKs ability to store the data in memory reducing the amount of memory used overall. """ DATA_SAMPLING: StandardProtocols.Enum.ValueType # 8 """Indicates that this SDK can sample in-flight elements. These samples can then be queried using the SampleDataRequest. Samples are uniquely associated with a PCollection. Meaning, samples are taken for each PCollection during bundle processing. This is disabled by default and enabled with the `enable_data_sampling` experiment. """ SDK_CONSUMING_RECEIVED_DATA: StandardProtocols.Enum.ValueType # 9 """Indicates whether the SDK sets the consuming_received_data field on the ProcessBundleProgressResponse messages. """ ORDERED_LIST_STATE: StandardProtocols.Enum.ValueType # 10 """Indicates whether the SDK supports ordered list state.""" ELEMENT_METADATA: StandardProtocols.Enum.ValueType # 11 """Indicates support for reading, writing and propagating Element's metadata""" MULTIMAP_STATE: StandardProtocols.Enum.ValueType # 12 """Indicates whether the SDK supports multimap state.""" def __init__( self, ) -> None: ... global___StandardProtocols = StandardProtocols @typing_extensions.final class StandardRunnerProtocols(google.protobuf.message.Message): """These URNs are used to indicate capabilities of runner that an environment may take advantage of when interacting with this runner. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardRunnerProtocols._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor MONITORING_INFO_SHORT_IDS: StandardRunnerProtocols._Enum.ValueType # 0 """Indicates suport the MonitoringInfo short id protocol.""" CONTROL_RESPONSE_ELEMENTS_EMBEDDING: StandardRunnerProtocols._Enum.ValueType # 6 """Indicates that this runner can process elements embedded in the ProcessBundleResponse. See more about the protocol at https://s.apache.org/beam-fn-api-control-data-embedding """ MULTIMAP_KEYS_VALUES_SIDE_INPUT: StandardRunnerProtocols._Enum.ValueType # 7 """Indicates that this runner can handle the multimap_keys_values_side_input style read of a multimap side input. """ class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... MONITORING_INFO_SHORT_IDS: StandardRunnerProtocols.Enum.ValueType # 0 """Indicates suport the MonitoringInfo short id protocol.""" CONTROL_RESPONSE_ELEMENTS_EMBEDDING: StandardRunnerProtocols.Enum.ValueType # 6 """Indicates that this runner can process elements embedded in the ProcessBundleResponse. See more about the protocol at https://s.apache.org/beam-fn-api-control-data-embedding """ MULTIMAP_KEYS_VALUES_SIDE_INPUT: StandardRunnerProtocols.Enum.ValueType # 7 """Indicates that this runner can handle the multimap_keys_values_side_input style read of a multimap side input. """ def __init__( self, ) -> None: ... global___StandardRunnerProtocols = StandardRunnerProtocols @typing_extensions.final class StandardRequirements(google.protobuf.message.Message): """These URNs are used to indicate requirements of a pipeline that cannot simply be expressed as a component (such as a Coder or PTransform) that the runner must understand. In many cases, this indicates a particular field of a transform must be inspected and respected (which allows new fields to be added in a forwards-compatible way). """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardRequirements._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor REQUIRES_STATEFUL_PROCESSING: StandardRequirements._Enum.ValueType # 0 """This requirement indicates the state_specs and timer_family_specs fields of ParDo transform payloads must be inspected. """ REQUIRES_BUNDLE_FINALIZATION: StandardRequirements._Enum.ValueType # 1 """This requirement indicates the requests_finalization field of ParDo transform payloads must be inspected. """ REQUIRES_STABLE_INPUT: StandardRequirements._Enum.ValueType # 2 """This requirement indicates the requires_stable_input field of ParDo transform payloads must be inspected. """ REQUIRES_TIME_SORTED_INPUT: StandardRequirements._Enum.ValueType # 3 """This requirement indicates the requires_time_sorted_input field of ParDo transform payloads must be inspected. """ REQUIRES_SPLITTABLE_DOFN: StandardRequirements._Enum.ValueType # 4 """This requirement indicates the restriction_coder_id field of ParDo transform payloads must be inspected. """ REQUIRES_ON_WINDOW_EXPIRATION: StandardRequirements._Enum.ValueType # 5 """This requirement indicates that the on_window_expiration_timer_family_spec field of ParDo transform payloads must be inspected. """ class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... REQUIRES_STATEFUL_PROCESSING: StandardRequirements.Enum.ValueType # 0 """This requirement indicates the state_specs and timer_family_specs fields of ParDo transform payloads must be inspected. """ REQUIRES_BUNDLE_FINALIZATION: StandardRequirements.Enum.ValueType # 1 """This requirement indicates the requests_finalization field of ParDo transform payloads must be inspected. """ REQUIRES_STABLE_INPUT: StandardRequirements.Enum.ValueType # 2 """This requirement indicates the requires_stable_input field of ParDo transform payloads must be inspected. """ REQUIRES_TIME_SORTED_INPUT: StandardRequirements.Enum.ValueType # 3 """This requirement indicates the requires_time_sorted_input field of ParDo transform payloads must be inspected. """ REQUIRES_SPLITTABLE_DOFN: StandardRequirements.Enum.ValueType # 4 """This requirement indicates the restriction_coder_id field of ParDo transform payloads must be inspected. """ REQUIRES_ON_WINDOW_EXPIRATION: StandardRequirements.Enum.ValueType # 5 """This requirement indicates that the on_window_expiration_timer_family_spec field of ParDo transform payloads must be inspected. """ def __init__( self, ) -> None: ... global___StandardRequirements = StandardRequirements @typing_extensions.final class FunctionSpec(google.protobuf.message.Message): """A URN along with a parameter object whose schema is determined by the URN. This structure is reused in two distinct, but compatible, ways: 1. This can be a specification of the function over PCollections that a PTransform computes. 2. This can be a specification of a user-defined function, possibly SDK-specific. (external to this message must be adequate context to indicate the environment in which the UDF can be understood). Though not explicit in this proto, there are two possibilities for the relationship of a runner to this specification that one should bear in mind: 1. The runner understands the URN. For example, it might be a well-known URN like "beam:transform:Top" or "beam:window_fn:FixedWindows" with an agreed-upon payload (e.g. a number or duration, respectively). 2. The runner does not understand the URN. It might be an SDK specific URN such as "beam:dofn:javasdk:1.0" that indicates to the SDK what the payload is, such as a serialized Java DoFn from a particular version of the Beam Java SDK. The payload will often then be an opaque message such as bytes in a language-specific serialization format. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor URN_FIELD_NUMBER: builtins.int PAYLOAD_FIELD_NUMBER: builtins.int urn: builtins.str """(Required) A URN that describes the accompanying payload. For any URN that is not recognized (by whomever is inspecting it) the parameter payload should be treated as opaque and passed as-is. """ payload: builtins.bytes """(Optional) The data specifying any parameters to the URN. If the URN does not require any arguments, this may be omitted. """ def __init__( self, *, urn: builtins.str | None = ..., payload: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["payload", b"payload", "urn", b"urn"]) -> None: ... global___FunctionSpec = FunctionSpec @typing_extensions.final class StandardDisplayData(google.protobuf.message.Message): """A set of well known URNs describing display data. All descriptions must contain how the value should be classified and how it is encoded. Note that some types are logical types which convey contextual information about the pipeline in addition to an encoding while others only specify the encoding itself. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _DisplayData: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _DisplayDataEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardDisplayData._DisplayData.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor LABELLED: StandardDisplayData._DisplayData.ValueType # 0 """A string label and value. Has a payload containing an encoded LabelledPayload. """ class DisplayData(_DisplayData, metaclass=_DisplayDataEnumTypeWrapper): ... LABELLED: StandardDisplayData.DisplayData.ValueType # 0 """A string label and value. Has a payload containing an encoded LabelledPayload. """ def __init__( self, ) -> None: ... global___StandardDisplayData = StandardDisplayData @typing_extensions.final class LabelledPayload(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor LABEL_FIELD_NUMBER: builtins.int STRING_VALUE_FIELD_NUMBER: builtins.int BOOL_VALUE_FIELD_NUMBER: builtins.int DOUBLE_VALUE_FIELD_NUMBER: builtins.int INT_VALUE_FIELD_NUMBER: builtins.int KEY_FIELD_NUMBER: builtins.int NAMESPACE_FIELD_NUMBER: builtins.int label: builtins.str """(Required) A human readable label for the value.""" string_value: builtins.str bool_value: builtins.bool double_value: builtins.float int_value: builtins.int key: builtins.str """(Required) The key identifies the actual content of the metadata.""" namespace: builtins.str """(Required) The namespace describes the context that specified the key.""" def __init__( self, *, label: builtins.str | None = ..., string_value: builtins.str | None = ..., bool_value: builtins.bool | None = ..., double_value: builtins.float | None = ..., int_value: builtins.int | None = ..., key: builtins.str | None = ..., namespace: builtins.str | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["bool_value", b"bool_value", "double_value", b"double_value", "int_value", b"int_value", "string_value", b"string_value", "value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["bool_value", b"bool_value", "double_value", b"double_value", "int_value", b"int_value", "key", b"key", "label", b"label", "namespace", b"namespace", "string_value", b"string_value", "value", b"value"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["value", b"value"]) -> typing_extensions.Literal["string_value", "bool_value", "double_value", "int_value"] | None: ... global___LabelledPayload = LabelledPayload @typing_extensions.final class DisplayData(google.protobuf.message.Message): """Static display data associated with a pipeline component. Display data is useful for pipeline runners IOs and diagnostic dashboards to display details about annotated components. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor URN_FIELD_NUMBER: builtins.int PAYLOAD_FIELD_NUMBER: builtins.int urn: builtins.str """A key used to describe the type of display data. See StandardDisplayData for the set of well known urns describing how the payload is meant to be interpreted. """ payload: builtins.bytes """(Optional) The data specifying any parameters to the URN. If the URN does not require any arguments, this may be omitted. """ def __init__( self, *, urn: builtins.str | None = ..., payload: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["payload", b"payload", "urn", b"urn"]) -> None: ... global___DisplayData = DisplayData @typing_extensions.final class MessageWithComponents(google.protobuf.message.Message): """The following transforms are not part of the RunnerApi specification, but may be useful for graph construction and manipulation. A disjoint union of all the things that may contain references that require Components to resolve. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor COMPONENTS_FIELD_NUMBER: builtins.int CODER_FIELD_NUMBER: builtins.int COMBINE_PAYLOAD_FIELD_NUMBER: builtins.int FUNCTION_SPEC_FIELD_NUMBER: builtins.int PAR_DO_PAYLOAD_FIELD_NUMBER: builtins.int PTRANSFORM_FIELD_NUMBER: builtins.int PCOLLECTION_FIELD_NUMBER: builtins.int READ_PAYLOAD_FIELD_NUMBER: builtins.int SIDE_INPUT_FIELD_NUMBER: builtins.int WINDOW_INTO_PAYLOAD_FIELD_NUMBER: builtins.int WINDOWING_STRATEGY_FIELD_NUMBER: builtins.int @property def components(self) -> global___Components: """(Optional) The by-reference components of the root message, enabling a standalone message. If this is absent, it is expected that there are no references. """ @property def coder(self) -> global___Coder: ... @property def combine_payload(self) -> global___CombinePayload: ... @property def function_spec(self) -> global___FunctionSpec: ... @property def par_do_payload(self) -> global___ParDoPayload: ... @property def ptransform(self) -> global___PTransform: ... @property def pcollection(self) -> global___PCollection: ... @property def read_payload(self) -> global___ReadPayload: ... @property def side_input(self) -> global___SideInput: ... @property def window_into_payload(self) -> global___WindowIntoPayload: ... @property def windowing_strategy(self) -> global___WindowingStrategy: ... def __init__( self, *, components: global___Components | None = ..., coder: global___Coder | None = ..., combine_payload: global___CombinePayload | None = ..., function_spec: global___FunctionSpec | None = ..., par_do_payload: global___ParDoPayload | None = ..., ptransform: global___PTransform | None = ..., pcollection: global___PCollection | None = ..., read_payload: global___ReadPayload | None = ..., side_input: global___SideInput | None = ..., window_into_payload: global___WindowIntoPayload | None = ..., windowing_strategy: global___WindowingStrategy | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["coder", b"coder", "combine_payload", b"combine_payload", "components", b"components", "function_spec", b"function_spec", "par_do_payload", b"par_do_payload", "pcollection", b"pcollection", "ptransform", b"ptransform", "read_payload", b"read_payload", "root", b"root", "side_input", b"side_input", "window_into_payload", b"window_into_payload", "windowing_strategy", b"windowing_strategy"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["coder", b"coder", "combine_payload", b"combine_payload", "components", b"components", "function_spec", b"function_spec", "par_do_payload", b"par_do_payload", "pcollection", b"pcollection", "ptransform", b"ptransform", "read_payload", b"read_payload", "root", b"root", "side_input", b"side_input", "window_into_payload", b"window_into_payload", "windowing_strategy", b"windowing_strategy"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["root", b"root"]) -> typing_extensions.Literal["coder", "combine_payload", "function_spec", "par_do_payload", "ptransform", "pcollection", "read_payload", "side_input", "window_into_payload", "windowing_strategy"] | None: ... global___MessageWithComponents = MessageWithComponents @typing_extensions.final class ExecutableStagePayload(google.protobuf.message.Message): """The payload for an executable stage. This will eventually be passed to an SDK in the form of a ProcessBundleDescriptor. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class SideInputId(google.protobuf.message.Message): """A reference to a side input. Side inputs are uniquely identified by PTransform id and local name. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int LOCAL_NAME_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The id of the PTransform that references this side input.""" local_name: builtins.str """(Required) The local name of this side input from the PTransform that references it.""" def __init__( self, *, transform_id: builtins.str | None = ..., local_name: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["local_name", b"local_name", "transform_id", b"transform_id"]) -> None: ... @typing_extensions.final class UserStateId(google.protobuf.message.Message): """A reference to user state. User states are uniquely identified by PTransform id and local name. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int LOCAL_NAME_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The id of the PTransform that references this user state.""" local_name: builtins.str """(Required) The local name of this user state for the PTransform that references it.""" def __init__( self, *, transform_id: builtins.str | None = ..., local_name: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["local_name", b"local_name", "transform_id", b"transform_id"]) -> None: ... @typing_extensions.final class TimerId(google.protobuf.message.Message): """A reference to a timer. Timers are uniquely identified by PTransform id and local name. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int LOCAL_NAME_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The id of the PTransform that references this timer.""" local_name: builtins.str """(Required) The local name of this timer for the PTransform that references it.""" def __init__( self, *, transform_id: builtins.str | None = ..., local_name: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["local_name", b"local_name", "transform_id", b"transform_id"]) -> None: ... @typing_extensions.final class TimerFamilyId(google.protobuf.message.Message): """A reference to a timer. Timers are uniquely identified by PTransform id and local name. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor TRANSFORM_ID_FIELD_NUMBER: builtins.int LOCAL_NAME_FIELD_NUMBER: builtins.int transform_id: builtins.str """(Required) The id of the PTransform that references this timer family.""" local_name: builtins.str """(Required) The local name of this timer family for the PTransform that references it.""" def __init__( self, *, transform_id: builtins.str | None = ..., local_name: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["local_name", b"local_name", "transform_id", b"transform_id"]) -> None: ... @typing_extensions.final class WireCoderSetting(google.protobuf.message.Message): """Settings that decide the coder type of wire coder.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor URN_FIELD_NUMBER: builtins.int PAYLOAD_FIELD_NUMBER: builtins.int INPUT_OR_OUTPUT_ID_FIELD_NUMBER: builtins.int TIMER_FIELD_NUMBER: builtins.int urn: builtins.str """(Required) The URN of the wire coder. Note that only windowed value coder or parameterized windowed value coder are supported. """ payload: builtins.bytes """(Optional) The data specifying any parameters to the URN. If the URN is beam:coder:windowed_value:v1, this may be omitted. If the URN is beam:coder:param_windowed_value:v1, the payload is an encoded windowed value using the beam:coder:windowed_value:v1 coder parameterized by a beam:coder:bytes:v1 element coder and the window coder that this param_windowed_value coder uses. """ input_or_output_id: builtins.str """The input or output PCollection id this setting applies to.""" @property def timer(self) -> global___ExecutableStagePayload.TimerId: """The timer id this setting applies to.""" def __init__( self, *, urn: builtins.str | None = ..., payload: builtins.bytes | None = ..., input_or_output_id: builtins.str | None = ..., timer: global___ExecutableStagePayload.TimerId | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["input_or_output_id", b"input_or_output_id", "target", b"target", "timer", b"timer"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["input_or_output_id", b"input_or_output_id", "payload", b"payload", "target", b"target", "timer", b"timer", "urn", b"urn"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["target", b"target"]) -> typing_extensions.Literal["input_or_output_id", "timer"] | None: ... ENVIRONMENT_FIELD_NUMBER: builtins.int WIRE_CODER_SETTINGS_FIELD_NUMBER: builtins.int INPUT_FIELD_NUMBER: builtins.int SIDE_INPUTS_FIELD_NUMBER: builtins.int TRANSFORMS_FIELD_NUMBER: builtins.int OUTPUTS_FIELD_NUMBER: builtins.int COMPONENTS_FIELD_NUMBER: builtins.int USER_STATES_FIELD_NUMBER: builtins.int TIMERS_FIELD_NUMBER: builtins.int TIMERFAMILIES_FIELD_NUMBER: builtins.int @property def environment(self) -> global___Environment: """(Required) Environment in which this stage executes. We use an environment rather than environment id because ExecutableStages use environments directly. This may change in the future. """ @property def wire_coder_settings(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ExecutableStagePayload.WireCoderSetting]: """The wire coder settings of this executable stage""" input: builtins.str """(Required) Input PCollection id. This must be present as a value in the inputs of any PTransform the ExecutableStagePayload is the payload of. """ @property def side_inputs(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ExecutableStagePayload.SideInputId]: """The side inputs required for this executable stage. Each side input of each PTransform within this ExecutableStagePayload must be represented within this field. """ @property def transforms(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """PTransform ids contained within this executable stage. This must contain at least one PTransform id. """ @property def outputs(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """Output PCollection ids. This must be equal to the values of the outputs of any PTransform the ExecutableStagePayload is the payload of. """ @property def components(self) -> global___Components: """(Required) The components for the Executable Stage. This must contain all of the Transforms in transforms, and the closure of all of the components they recognize. """ @property def user_states(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ExecutableStagePayload.UserStateId]: """The user states required for this executable stage. Each user state of each PTransform within this ExecutableStagePayload must be represented within this field. """ @property def timers(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ExecutableStagePayload.TimerId]: """The timers required for this executable stage. Each timer of each PTransform within this ExecutableStagePayload must be represented within this field. """ @property def timerFamilies(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ExecutableStagePayload.TimerFamilyId]: """The timerfamilies required for this executable stage. Each timer familyof each PTransform within this ExecutableStagePayload must be represented within this field. """ def __init__( self, *, environment: global___Environment | None = ..., wire_coder_settings: collections.abc.Iterable[global___ExecutableStagePayload.WireCoderSetting] | None = ..., input: builtins.str | None = ..., side_inputs: collections.abc.Iterable[global___ExecutableStagePayload.SideInputId] | None = ..., transforms: collections.abc.Iterable[builtins.str] | None = ..., outputs: collections.abc.Iterable[builtins.str] | None = ..., components: global___Components | None = ..., user_states: collections.abc.Iterable[global___ExecutableStagePayload.UserStateId] | None = ..., timers: collections.abc.Iterable[global___ExecutableStagePayload.TimerId] | None = ..., timerFamilies: collections.abc.Iterable[global___ExecutableStagePayload.TimerFamilyId] | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["components", b"components", "environment", b"environment"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["components", b"components", "environment", b"environment", "input", b"input", "outputs", b"outputs", "side_inputs", b"side_inputs", "timerFamilies", b"timerFamilies", "timers", b"timers", "transforms", b"transforms", "user_states", b"user_states", "wire_coder_settings", b"wire_coder_settings"]) -> None: ... global___ExecutableStagePayload = ExecutableStagePayload @typing_extensions.final class StandardResourceHints(google.protobuf.message.Message): """See https://beam.apache.org/documentation/runtime/resource-hints/ for additional documentation on the behavior of StandardResourceHint. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[StandardResourceHints._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor ACCELERATOR: StandardResourceHints._Enum.ValueType # 0 """Describes hardware accelerators that are desired to have in the execution environment. Payload: ASCII encoded string with the following format: "type:;count:;" where type is an accelerator sku, count is the number of accelerators per worker, and options are related options flags. """ MIN_RAM_BYTES: StandardResourceHints._Enum.ValueType # 1 """Describes desired minimal available RAM size in transform's execution environment. SDKs should convert the size to bytes, but can allow users to specify human-friendly units (e.g. GiB). Payload: ASCII encoded string of the base 10 representation of an integer number of bytes. """ CPU_COUNT: StandardResourceHints._Enum.ValueType # 2 """Describes desired number of CPUs available in transform's execution environment. SDKs should accept and validate a positive integer count. Payload: ASCII encoded string of the base 10 representation of an integer number of CPUs. """ MAX_ACTIVE_BUNDLES_PER_WORKER: StandardResourceHints._Enum.ValueType # 3 """Describes max number of active bundles per worker in transform's execution environment. SDKs should accept and validate a positive integer count. Payload: ASCII encoded string of the base 10 representation of an integer number of active bundles. """ class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... ACCELERATOR: StandardResourceHints.Enum.ValueType # 0 """Describes hardware accelerators that are desired to have in the execution environment. Payload: ASCII encoded string with the following format: "type:;count:;" where type is an accelerator sku, count is the number of accelerators per worker, and options are related options flags. """ MIN_RAM_BYTES: StandardResourceHints.Enum.ValueType # 1 """Describes desired minimal available RAM size in transform's execution environment. SDKs should convert the size to bytes, but can allow users to specify human-friendly units (e.g. GiB). Payload: ASCII encoded string of the base 10 representation of an integer number of bytes. """ CPU_COUNT: StandardResourceHints.Enum.ValueType # 2 """Describes desired number of CPUs available in transform's execution environment. SDKs should accept and validate a positive integer count. Payload: ASCII encoded string of the base 10 representation of an integer number of CPUs. """ MAX_ACTIVE_BUNDLES_PER_WORKER: StandardResourceHints.Enum.ValueType # 3 """Describes max number of active bundles per worker in transform's execution environment. SDKs should accept and validate a positive integer count. Payload: ASCII encoded string of the base 10 representation of an integer number of active bundles. """ def __init__( self, ) -> None: ... global___StandardResourceHints = StandardResourceHints BEAM_URN_FIELD_NUMBER: builtins.int BEAM_CONSTANT_FIELD_NUMBER: builtins.int beam_urn: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.EnumValueOptions, builtins.str] """An extension to be used for specifying the standard URN of various pipeline entities, e.g. transforms, functions, coders etc. Code should refer to the URNs of those entities by extracting it from the (beam_urn) extension, rather than by hard-coding the URN. The recommended pattern for declaring it is (exemplified by coders): message StandardCoders { enum Enum { BYTES = 0 [(beam_urn) = "beam:coder:bytes:v1"]; ... } } If there are multiple categories of entities of this type, use the following pattern (exemplified by PTransforms): message StandardPTransforms { enum Primitives { ... } enum Composites { ... } } """ beam_constant: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.EnumValueOptions, builtins.str] """A value to store other constants""" ././@PaxHeader0000000000000000000000000000021400000000000010212 xustar00118 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/beam_runner_api_pb2_grpc.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/beam_runner_api_pb20000644000175100017510000000023715206113151032757 0ustar00runnerrunner# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc ././@PaxHeader0000000000000000000000000000021400000000000010212 xustar00118 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/beam_runner_api_pb2_urns.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/beam_runner_api_pb20000644000175100017510000003713015206113151032761 0ustar00runnerrunnerfrom ........utils import PropertiesFromEnumValue from . import metrics_pb2 EMPTY_MONITORING_INFO_LABEL_PROPS = metrics_pb2.MonitoringInfoLabelProps() EMPTY_MONITORING_INFO_SPEC = metrics_pb2.MonitoringInfoSpec() class BeamConstants(object): class Constants(object): MIN_TIMESTAMP_MILLIS = PropertiesFromEnumValue('', '-9223372036854775', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MAX_TIMESTAMP_MILLIS = PropertiesFromEnumValue('', '9223372036854775', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) GLOBAL_WINDOW_MAX_TIMESTAMP_MILLIS = PropertiesFromEnumValue('', '9223371950454775', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class StandardArtifacts(object): class Types(object): FILE = PropertiesFromEnumValue('beam:artifact:type:file:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) URL = PropertiesFromEnumValue('beam:artifact:type:url:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) EMBEDDED = PropertiesFromEnumValue('beam:artifact:type:embedded:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) PYPI = PropertiesFromEnumValue('beam:artifact:type:pypi:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MAVEN = PropertiesFromEnumValue('beam:artifact:type:maven:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) DEFERRED = PropertiesFromEnumValue('beam:artifact:type:deferred:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class Roles(object): STAGING_TO = PropertiesFromEnumValue('beam:artifact:role:staging_to:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) PIP_REQUIREMENTS_FILE = PropertiesFromEnumValue('beam:artifact:role:pip_requirements_file:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) GO_WORKER_BINARY = PropertiesFromEnumValue('beam:artifact:role:go_worker_binary:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class StandardCoders(object): class Enum(object): BYTES = PropertiesFromEnumValue('beam:coder:bytes:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) STRING_UTF8 = PropertiesFromEnumValue('beam:coder:string_utf8:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) KV = PropertiesFromEnumValue('beam:coder:kv:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) BOOL = PropertiesFromEnumValue('beam:coder:bool:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) VARINT = PropertiesFromEnumValue('beam:coder:varint:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) DOUBLE = PropertiesFromEnumValue('beam:coder:double:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ITERABLE = PropertiesFromEnumValue('beam:coder:iterable:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) TIMER = PropertiesFromEnumValue('beam:coder:timer:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) INTERVAL_WINDOW = PropertiesFromEnumValue('beam:coder:interval_window:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) LENGTH_PREFIX = PropertiesFromEnumValue('beam:coder:length_prefix:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) GLOBAL_WINDOW = PropertiesFromEnumValue('beam:coder:global_window:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) WINDOWED_VALUE = PropertiesFromEnumValue('beam:coder:windowed_value:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) PARAM_WINDOWED_VALUE = PropertiesFromEnumValue('beam:coder:param_windowed_value:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) STATE_BACKED_ITERABLE = PropertiesFromEnumValue('beam:coder:state_backed_iterable:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) CUSTOM_WINDOW = PropertiesFromEnumValue('beam:coder:custom_window:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ROW = PropertiesFromEnumValue('beam:coder:row:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) SCHEMA = PropertiesFromEnumValue('beam:coder:schema:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) SHARDED_KEY = PropertiesFromEnumValue('beam:coder:sharded_key:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) NULLABLE = PropertiesFromEnumValue('beam:coder:nullable:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class StandardDisplayData(object): class DisplayData(object): LABELLED = PropertiesFromEnumValue('beam:display_data:labelled:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class StandardEnvironments(object): class Environments(object): DOCKER = PropertiesFromEnumValue('beam:env:docker:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) PROCESS = PropertiesFromEnumValue('beam:env:process:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) EXTERNAL = PropertiesFromEnumValue('beam:env:external:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) DEFAULT = PropertiesFromEnumValue('beam:env:default:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ANYOF = PropertiesFromEnumValue('beam:env:anyof:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class StandardPTransforms(object): class Primitives(object): PAR_DO = PropertiesFromEnumValue('beam:transform:pardo:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) FLATTEN = PropertiesFromEnumValue('beam:transform:flatten:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) GROUP_BY_KEY = PropertiesFromEnumValue('beam:transform:group_by_key:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) IMPULSE = PropertiesFromEnumValue('beam:transform:impulse:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ASSIGN_WINDOWS = PropertiesFromEnumValue('beam:transform:window_into:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) TEST_STREAM = PropertiesFromEnumValue('beam:transform:teststream:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MAP_WINDOWS = PropertiesFromEnumValue('beam:transform:map_windows:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MERGE_WINDOWS = PropertiesFromEnumValue('beam:transform:merge_windows:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) TO_STRING = PropertiesFromEnumValue('beam:transform:to_string:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class DeprecatedPrimitives(object): READ = PropertiesFromEnumValue('beam:transform:read:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) CREATE_VIEW = PropertiesFromEnumValue('beam:transform:create_view:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class Composites(object): COMBINE_PER_KEY = PropertiesFromEnumValue('beam:transform:combine_per_key:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) COMBINE_GLOBALLY = PropertiesFromEnumValue('beam:transform:combine_globally:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) RESHUFFLE = PropertiesFromEnumValue('beam:transform:reshuffle:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) WRITE_FILES = PropertiesFromEnumValue('beam:transform:write_files:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) PUBSUB_READ = PropertiesFromEnumValue('beam:transform:pubsub_read:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) PUBSUB_WRITE = PropertiesFromEnumValue('beam:transform:pubsub_write:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) PUBSUB_WRITE_V2 = PropertiesFromEnumValue('beam:transform:pubsub_write:v2', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) GROUP_INTO_BATCHES_WITH_SHARDED_KEY = PropertiesFromEnumValue('beam:transform:group_into_batches_with_sharded_key:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) REDISTRIBUTE_BY_KEY = PropertiesFromEnumValue('beam:transform:redistribute_by_key:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) REDISTRIBUTE_ARBITRARILY = PropertiesFromEnumValue('beam:transform:redistribute_arbitrarily:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class CombineComponents(object): COMBINE_PER_KEY_PRECOMBINE = PropertiesFromEnumValue('beam:transform:combine_per_key_precombine:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) COMBINE_PER_KEY_MERGE_ACCUMULATORS = PropertiesFromEnumValue('beam:transform:combine_per_key_merge_accumulators:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) COMBINE_PER_KEY_EXTRACT_OUTPUTS = PropertiesFromEnumValue('beam:transform:combine_per_key_extract_outputs:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) COMBINE_GROUPED_VALUES = PropertiesFromEnumValue('beam:transform:combine_grouped_values:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) COMBINE_PER_KEY_CONVERT_TO_ACCUMULATORS = PropertiesFromEnumValue('beam:transform:combine_per_key_convert_to_accumulators:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class SplittableParDoComponents(object): PAIR_WITH_RESTRICTION = PropertiesFromEnumValue('beam:transform:sdf_pair_with_restriction:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) SPLIT_AND_SIZE_RESTRICTIONS = PropertiesFromEnumValue('beam:transform:sdf_split_and_size_restrictions:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS = PropertiesFromEnumValue('beam:transform:sdf_process_sized_element_and_restrictions:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) TRUNCATE_SIZED_RESTRICTION = PropertiesFromEnumValue('beam:transform:sdf_truncate_sized_restrictions:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class GroupIntoBatchesComponents(object): GROUP_INTO_BATCHES = PropertiesFromEnumValue('beam:transform:group_into_batches:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class StandardProtocols(object): class Enum(object): LEGACY_PROGRESS_REPORTING = PropertiesFromEnumValue('beam:protocol:progress_reporting:v0', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) PROGRESS_REPORTING = PropertiesFromEnumValue('beam:protocol:progress_reporting:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) WORKER_STATUS = PropertiesFromEnumValue('beam:protocol:worker_status:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MULTI_CORE_BUNDLE_PROCESSING = PropertiesFromEnumValue('beam:protocol:multi_core_bundle_processing:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) SIBLING_WORKERS = PropertiesFromEnumValue('beam:protocol:sibling_workers:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) HARNESS_MONITORING_INFOS = PropertiesFromEnumValue('beam:protocol:harness_monitoring_infos:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) CONTROL_REQUEST_ELEMENTS_EMBEDDING = PropertiesFromEnumValue('beam:protocol:control_request_elements_embedding:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) STATE_CACHING = PropertiesFromEnumValue('beam:protocol:state_caching:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) DATA_SAMPLING = PropertiesFromEnumValue('beam:protocol:data_sampling:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) SDK_CONSUMING_RECEIVED_DATA = PropertiesFromEnumValue('beam:protocol:sdk_consuming_received_data:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ORDERED_LIST_STATE = PropertiesFromEnumValue('beam:protocol:ordered_list_state:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ELEMENT_METADATA = PropertiesFromEnumValue('beam:protocol:element_metadata:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MULTIMAP_STATE = PropertiesFromEnumValue('beam:protocol:multimap_state:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class StandardRequirements(object): class Enum(object): REQUIRES_STATEFUL_PROCESSING = PropertiesFromEnumValue('beam:requirement:pardo:stateful:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) REQUIRES_BUNDLE_FINALIZATION = PropertiesFromEnumValue('beam:requirement:pardo:finalization:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) REQUIRES_STABLE_INPUT = PropertiesFromEnumValue('beam:requirement:pardo:stable_input:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) REQUIRES_TIME_SORTED_INPUT = PropertiesFromEnumValue('beam:requirement:pardo:time_sorted_input:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) REQUIRES_SPLITTABLE_DOFN = PropertiesFromEnumValue('beam:requirement:pardo:splittable_dofn:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) REQUIRES_ON_WINDOW_EXPIRATION = PropertiesFromEnumValue('beam:requirement:pardo:on_window_expiration:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class StandardResourceHints(object): class Enum(object): ACCELERATOR = PropertiesFromEnumValue('beam:resources:accelerator:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MIN_RAM_BYTES = PropertiesFromEnumValue('beam:resources:min_ram_bytes:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) CPU_COUNT = PropertiesFromEnumValue('beam:resources:cpu_count:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MAX_ACTIVE_BUNDLES_PER_WORKER = PropertiesFromEnumValue('beam:resources:max_active_bundles_per_worker:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class StandardRunnerProtocols(object): class Enum(object): MONITORING_INFO_SHORT_IDS = PropertiesFromEnumValue('beam:protocol:monitoring_info_short_ids:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) CONTROL_RESPONSE_ELEMENTS_EMBEDDING = PropertiesFromEnumValue('beam:protocol:control_response_elements_embedding:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MULTIMAP_KEYS_VALUES_SIDE_INPUT = PropertiesFromEnumValue('beam:protocol:multimap_keys_values_side_input:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class StandardSideInputTypes(object): class Enum(object): ITERABLE = PropertiesFromEnumValue('beam:side_input:iterable:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MULTIMAP = PropertiesFromEnumValue('beam:side_input:multimap:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class StandardUserStateTypes(object): class Enum(object): BAG = PropertiesFromEnumValue('beam:user_state:bag:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MULTIMAP = PropertiesFromEnumValue('beam:user_state:multimap:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ORDERED_LIST = PropertiesFromEnumValue('beam:user_state:ordered_list:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/endpoints_pb2.py0000644000175100017510000000343015206113151032261 0ustar00runnerrunner# -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # source: org/apache/beam/model/pipeline/v1/endpoints.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n1org/apache/beam/model/pipeline/v1/endpoints.proto\x12!org.apache.beam.model.pipeline.v1\"r\n\x14\x41piServiceDescriptor\x12\x0b\n\x03url\x18\x01 \x01(\t\x12M\n\x0e\x61uthentication\x18\x02 \x01(\x0b\x32\x35.org.apache.beam.model.pipeline.v1.AuthenticationSpec\"2\n\x12\x41uthenticationSpec\x12\x0b\n\x03urn\x18\x01 \x01(\t\x12\x0f\n\x07payload\x18\x02 \x01(\x0c\x42x\n!org.apache.beam.model.pipeline.v1B\tEndpointsZHgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1;pipeline_v1b\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'org.apache.beam.model.pipeline.v1.endpoints_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n!org.apache.beam.model.pipeline.v1B\tEndpointsZHgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1;pipeline_v1' _globals['_APISERVICEDESCRIPTOR']._serialized_start=88 _globals['_APISERVICEDESCRIPTOR']._serialized_end=202 _globals['_AUTHENTICATIONSPEC']._serialized_start=204 _globals['_AUTHENTICATIONSPEC']._serialized_end=254 # @@protoc_insertion_point(module_scope) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/endpoints_pb2.pyi0000644000175100017510000000457415206113151032444 0ustar00runnerrunner""" @generated by mypy-protobuf. Do not edit manually! isort:skip_file Protocol Buffers describing endpoints containing a service. """ import builtins import google.protobuf.descriptor import google.protobuf.message import sys if sys.version_info >= (3, 8): import typing as typing_extensions else: import typing_extensions DESCRIPTOR: google.protobuf.descriptor.FileDescriptor @typing_extensions.final class ApiServiceDescriptor(google.protobuf.message.Message): """A description of how to connect to a Beam API endpoint.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor URL_FIELD_NUMBER: builtins.int AUTHENTICATION_FIELD_NUMBER: builtins.int url: builtins.str """(Required) The URL to connect to.""" @property def authentication(self) -> global___AuthenticationSpec: """(Optional) The method for authentication. If unspecified, access to the url is already being performed in a trusted context (e.g. localhost, private network). """ def __init__( self, *, url: builtins.str | None = ..., authentication: global___AuthenticationSpec | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["authentication", b"authentication"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["authentication", b"authentication", "url", b"url"]) -> None: ... global___ApiServiceDescriptor = ApiServiceDescriptor @typing_extensions.final class AuthenticationSpec(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor URN_FIELD_NUMBER: builtins.int PAYLOAD_FIELD_NUMBER: builtins.int urn: builtins.str """(Required) A URN that describes the accompanying payload. For any URN that is not recognized (by whomever is inspecting it) the parameter payload should be treated as opaque and passed as-is. """ payload: builtins.bytes """(Optional) The data specifying any parameters to the URN. If the URN does not require any arguments, this may be omitted. """ def __init__( self, *, urn: builtins.str | None = ..., payload: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["payload", b"payload", "urn", b"urn"]) -> None: ... global___AuthenticationSpec = AuthenticationSpec ././@PaxHeader0000000000000000000000000000020600000000000010213 xustar00112 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/endpoints_pb2_grpc.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/endpoints_pb2_grpc.0000644000175100017510000000023715206113151032725 0ustar00runnerrunner# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc ././@PaxHeader0000000000000000000000000000021300000000000010211 xustar00117 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/external_transforms_pb2.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/external_transforms0000644000175100017510000002523415206113151033172 0ustar00runnerrunner# -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # source: org/apache/beam/model/pipeline/v1/external_transforms.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() from . import schema_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_schema__pb2 from . import beam_runner_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_beam__runner__api__pb2 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n;org/apache/beam/model/pipeline/v1/external_transforms.proto\x12!org.apache.beam.model.pipeline.v1\x1a.org/apache/beam/model/pipeline/v1/schema.proto\x1a\x37org/apache/beam/model/pipeline/v1/beam_runner_api.proto\"j\n\x1c\x45xternalConfigurationPayload\x12\x39\n\x06schema\x18\x01 \x01(\x0b\x32).org.apache.beam.model.pipeline.v1.Schema\x12\x0f\n\x07payload\x18\x02 \x01(\x0c\"\xac\x01\n\x10\x45xpansionMethods\"\x97\x01\n\x04\x45num\x12H\n\x11JAVA_CLASS_LOOKUP\x10\x00\x1a\x31\xa2\xb4\xfa\xc2\x05+beam:expansion:payload:java_class_lookup:v1\x12\x45\n\x10SCHEMA_TRANSFORM\x10\x01\x1a/\xa2\xb4\xfa\xc2\x05)beam:expansion:payload:schematransform:v1\"\xb2\x08\n\x11ManagedTransforms\"\x9c\x08\n\x04Urns\x12L\n\x0cICEBERG_READ\x10\x00\x1a:\xa2\xb4\xfa\xc2\x05\x34\x62\x65\x61m:schematransform:org.apache.beam:iceberg_read:v1\x12N\n\rICEBERG_WRITE\x10\x01\x1a;\xa2\xb4\xfa\xc2\x05\x35\x62\x65\x61m:schematransform:org.apache.beam:iceberg_write:v1\x12H\n\nKAFKA_READ\x10\x02\x1a\x38\xa2\xb4\xfa\xc2\x05\x32\x62\x65\x61m:schematransform:org.apache.beam:kafka_read:v1\x12J\n\x0bKAFKA_WRITE\x10\x03\x1a\x39\xa2\xb4\xfa\xc2\x05\x33\x62\x65\x61m:schematransform:org.apache.beam:kafka_write:v1\x12V\n\rBIGQUERY_READ\x10\x04\x1a\x43\xa2\xb4\xfa\xc2\x05=beam:schematransform:org.apache.beam:bigquery_storage_read:v1\x12P\n\x0e\x42IGQUERY_WRITE\x10\x05\x1a<\xa2\xb4\xfa\xc2\x05\x36\x62\x65\x61m:schematransform:org.apache.beam:bigquery_write:v1\x12T\n\x10ICEBERG_CDC_READ\x10\x06\x1a>\xa2\xb4\xfa\xc2\x05\x38\x62\x65\x61m:schematransform:org.apache.beam:iceberg_cdc_read:v1\x12N\n\rPOSTGRES_READ\x10\x07\x1a;\xa2\xb4\xfa\xc2\x05\x35\x62\x65\x61m:schematransform:org.apache.beam:postgres_read:v1\x12P\n\x0ePOSTGRES_WRITE\x10\x08\x1a<\xa2\xb4\xfa\xc2\x05\x36\x62\x65\x61m:schematransform:org.apache.beam:postgres_write:v1\x12H\n\nMYSQL_READ\x10\t\x1a\x38\xa2\xb4\xfa\xc2\x05\x32\x62\x65\x61m:schematransform:org.apache.beam:mysql_read:v1\x12J\n\x0bMYSQL_WRITE\x10\n\x1a\x39\xa2\xb4\xfa\xc2\x05\x33\x62\x65\x61m:schematransform:org.apache.beam:mysql_write:v1\x12R\n\x0fSQL_SERVER_READ\x10\x0b\x1a=\xa2\xb4\xfa\xc2\x05\x37\x62\x65\x61m:schematransform:org.apache.beam:sql_server_read:v1\x12T\n\x10SQL_SERVER_WRITE\x10\x0c\x1a>\xa2\xb4\xfa\xc2\x05\x38\x62\x65\x61m:schematransform:org.apache.beam:sql_server_write:v1\"\xf7\x01\n\x16JavaClassLookupPayload\x12\x12\n\nclass_name\x18\x01 \x01(\t\x12\x1a\n\x12\x63onstructor_method\x18\x02 \x01(\t\x12\x45\n\x12\x63onstructor_schema\x18\x03 \x01(\x0b\x32).org.apache.beam.model.pipeline.v1.Schema\x12\x1b\n\x13\x63onstructor_payload\x18\x04 \x01(\x0c\x12I\n\x0f\x62uilder_methods\x18\x05 \x03(\x0b\x32\x30.org.apache.beam.model.pipeline.v1.BuilderMethod\"i\n\rBuilderMethod\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x39\n\x06schema\x18\x02 \x01(\x0b\x32).org.apache.beam.model.pipeline.v1.Schema\x12\x0f\n\x07payload\x18\x03 \x01(\x0c\"\xfa\x01\n\x0b\x41nnotations\"\xea\x01\n\x04\x45num\x12$\n\x0e\x43ONFIG_ROW_KEY\x10\x00\x1a\x10\xaa\xb4\xfa\xc2\x05\nconfig_row\x12\x32\n\x15\x43ONFIG_ROW_SCHEMA_KEY\x10\x01\x1a\x17\xaa\xb4\xfa\xc2\x05\x11\x63onfig_row_schema\x12\x36\n\x17SCHEMATRANSFORM_URN_KEY\x10\x02\x1a\x19\xaa\xb4\xfa\xc2\x05\x13schematransform_urn\x12P\n$MANAGED_UNDERLYING_TRANSFORM_URN_KEY\x10\x03\x1a&\xaa\xb4\xfa\xc2\x05 managed_underlying_transform_urn\"\x90\x01\n\x16SchemaTransformPayload\x12\x12\n\nidentifier\x18\x01 \x01(\t\x12G\n\x14\x63onfiguration_schema\x18\x02 \x01(\x0b\x32).org.apache.beam.model.pipeline.v1.Schema\x12\x19\n\x11\x63onfiguration_row\x18\x03 \x01(\x0c\x42\x81\x01\n!org.apache.beam.model.pipeline.v1B\x12\x45xternalTransformsZHgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1;pipeline_v1b\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'org.apache.beam.model.pipeline.v1.external_transforms_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n!org.apache.beam.model.pipeline.v1B\022ExternalTransformsZHgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1;pipeline_v1' _globals['_EXPANSIONMETHODS_ENUM'].values_by_name["JAVA_CLASS_LOOKUP"]._options = None _globals['_EXPANSIONMETHODS_ENUM'].values_by_name["JAVA_CLASS_LOOKUP"]._serialized_options = b'\242\264\372\302\005+beam:expansion:payload:java_class_lookup:v1' _globals['_EXPANSIONMETHODS_ENUM'].values_by_name["SCHEMA_TRANSFORM"]._options = None _globals['_EXPANSIONMETHODS_ENUM'].values_by_name["SCHEMA_TRANSFORM"]._serialized_options = b'\242\264\372\302\005)beam:expansion:payload:schematransform:v1' _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["ICEBERG_READ"]._options = None _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["ICEBERG_READ"]._serialized_options = b'\242\264\372\302\0054beam:schematransform:org.apache.beam:iceberg_read:v1' _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["ICEBERG_WRITE"]._options = None _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["ICEBERG_WRITE"]._serialized_options = b'\242\264\372\302\0055beam:schematransform:org.apache.beam:iceberg_write:v1' _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["KAFKA_READ"]._options = None _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["KAFKA_READ"]._serialized_options = b'\242\264\372\302\0052beam:schematransform:org.apache.beam:kafka_read:v1' _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["KAFKA_WRITE"]._options = None _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["KAFKA_WRITE"]._serialized_options = b'\242\264\372\302\0053beam:schematransform:org.apache.beam:kafka_write:v1' _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["BIGQUERY_READ"]._options = None _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["BIGQUERY_READ"]._serialized_options = b'\242\264\372\302\005=beam:schematransform:org.apache.beam:bigquery_storage_read:v1' _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["BIGQUERY_WRITE"]._options = None _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["BIGQUERY_WRITE"]._serialized_options = b'\242\264\372\302\0056beam:schematransform:org.apache.beam:bigquery_write:v1' _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["ICEBERG_CDC_READ"]._options = None _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["ICEBERG_CDC_READ"]._serialized_options = b'\242\264\372\302\0058beam:schematransform:org.apache.beam:iceberg_cdc_read:v1' _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["POSTGRES_READ"]._options = None _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["POSTGRES_READ"]._serialized_options = b'\242\264\372\302\0055beam:schematransform:org.apache.beam:postgres_read:v1' _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["POSTGRES_WRITE"]._options = None _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["POSTGRES_WRITE"]._serialized_options = b'\242\264\372\302\0056beam:schematransform:org.apache.beam:postgres_write:v1' _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["MYSQL_READ"]._options = None _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["MYSQL_READ"]._serialized_options = b'\242\264\372\302\0052beam:schematransform:org.apache.beam:mysql_read:v1' _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["MYSQL_WRITE"]._options = None _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["MYSQL_WRITE"]._serialized_options = b'\242\264\372\302\0053beam:schematransform:org.apache.beam:mysql_write:v1' _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["SQL_SERVER_READ"]._options = None _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["SQL_SERVER_READ"]._serialized_options = b'\242\264\372\302\0057beam:schematransform:org.apache.beam:sql_server_read:v1' _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["SQL_SERVER_WRITE"]._options = None _globals['_MANAGEDTRANSFORMS_URNS'].values_by_name["SQL_SERVER_WRITE"]._serialized_options = b'\242\264\372\302\0058beam:schematransform:org.apache.beam:sql_server_write:v1' _globals['_ANNOTATIONS_ENUM'].values_by_name["CONFIG_ROW_KEY"]._options = None _globals['_ANNOTATIONS_ENUM'].values_by_name["CONFIG_ROW_KEY"]._serialized_options = b'\252\264\372\302\005\nconfig_row' _globals['_ANNOTATIONS_ENUM'].values_by_name["CONFIG_ROW_SCHEMA_KEY"]._options = None _globals['_ANNOTATIONS_ENUM'].values_by_name["CONFIG_ROW_SCHEMA_KEY"]._serialized_options = b'\252\264\372\302\005\021config_row_schema' _globals['_ANNOTATIONS_ENUM'].values_by_name["SCHEMATRANSFORM_URN_KEY"]._options = None _globals['_ANNOTATIONS_ENUM'].values_by_name["SCHEMATRANSFORM_URN_KEY"]._serialized_options = b'\252\264\372\302\005\023schematransform_urn' _globals['_ANNOTATIONS_ENUM'].values_by_name["MANAGED_UNDERLYING_TRANSFORM_URN_KEY"]._options = None _globals['_ANNOTATIONS_ENUM'].values_by_name["MANAGED_UNDERLYING_TRANSFORM_URN_KEY"]._serialized_options = b'\252\264\372\302\005 managed_underlying_transform_urn' _globals['_EXTERNALCONFIGURATIONPAYLOAD']._serialized_start=203 _globals['_EXTERNALCONFIGURATIONPAYLOAD']._serialized_end=309 _globals['_EXPANSIONMETHODS']._serialized_start=312 _globals['_EXPANSIONMETHODS']._serialized_end=484 _globals['_EXPANSIONMETHODS_ENUM']._serialized_start=333 _globals['_EXPANSIONMETHODS_ENUM']._serialized_end=484 _globals['_MANAGEDTRANSFORMS']._serialized_start=487 _globals['_MANAGEDTRANSFORMS']._serialized_end=1561 _globals['_MANAGEDTRANSFORMS_URNS']._serialized_start=509 _globals['_MANAGEDTRANSFORMS_URNS']._serialized_end=1561 _globals['_JAVACLASSLOOKUPPAYLOAD']._serialized_start=1564 _globals['_JAVACLASSLOOKUPPAYLOAD']._serialized_end=1811 _globals['_BUILDERMETHOD']._serialized_start=1813 _globals['_BUILDERMETHOD']._serialized_end=1918 _globals['_ANNOTATIONS']._serialized_start=1921 _globals['_ANNOTATIONS']._serialized_end=2171 _globals['_ANNOTATIONS_ENUM']._serialized_start=1937 _globals['_ANNOTATIONS_ENUM']._serialized_end=2171 _globals['_SCHEMATRANSFORMPAYLOAD']._serialized_start=2174 _globals['_SCHEMATRANSFORMPAYLOAD']._serialized_end=2318 # @@protoc_insertion_point(module_scope) ././@PaxHeader0000000000000000000000000000021400000000000010212 xustar00118 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/external_transforms_pb2.pyi 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/external_transforms0000644000175100017510000003262415206113151033173 0ustar00runnerrunner""" @generated by mypy-protobuf. Do not edit manually! isort:skip_file Protocol Buffers describing the external transforms available. """ import builtins import collections.abc import google.protobuf.descriptor import google.protobuf.internal.containers import google.protobuf.internal.enum_type_wrapper import google.protobuf.message import org.apache.beam.model.pipeline.v1.schema_pb2 import sys import typing if sys.version_info >= (3, 10): import typing as typing_extensions else: import typing_extensions DESCRIPTOR: google.protobuf.descriptor.FileDescriptor @typing_extensions.final class ExternalConfigurationPayload(google.protobuf.message.Message): """A configuration payload for an external transform. Used as the payload of ExternalTransform as part of an ExpansionRequest. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor SCHEMA_FIELD_NUMBER: builtins.int PAYLOAD_FIELD_NUMBER: builtins.int @property def schema(self) -> org.apache.beam.model.pipeline.v1.schema_pb2.Schema: """A schema for use in beam:coder:row:v1""" payload: builtins.bytes """A payload which can be decoded using beam:coder:row:v1 and the given schema. """ def __init__( self, *, schema: org.apache.beam.model.pipeline.v1.schema_pb2.Schema | None = ..., payload: builtins.bytes | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["schema", b"schema"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["payload", b"payload", "schema", b"schema"]) -> None: ... global___ExternalConfigurationPayload = ExternalConfigurationPayload @typing_extensions.final class ExpansionMethods(google.protobuf.message.Message): """Defines specific expansion methods that may be used to expand cross-language transforms. Has to be set as the URN of the transform of the expansion request. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ExpansionMethods._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor JAVA_CLASS_LOOKUP: ExpansionMethods._Enum.ValueType # 0 """Expand a Java transform using specified constructor and builder methods. Transform payload will be of type JavaClassLookupPayload. """ SCHEMA_TRANSFORM: ExpansionMethods._Enum.ValueType # 1 """Expanding a SchemaTransform identified by the expansion service. Transform payload will be of type SchemaTransformPayload. """ class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... JAVA_CLASS_LOOKUP: ExpansionMethods.Enum.ValueType # 0 """Expand a Java transform using specified constructor and builder methods. Transform payload will be of type JavaClassLookupPayload. """ SCHEMA_TRANSFORM: ExpansionMethods.Enum.ValueType # 1 """Expanding a SchemaTransform identified by the expansion service. Transform payload will be of type SchemaTransformPayload. """ def __init__( self, ) -> None: ... global___ExpansionMethods = ExpansionMethods @typing_extensions.final class ManagedTransforms(google.protobuf.message.Message): """Defines the URNs for managed transforms.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Urns: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _UrnsEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ManagedTransforms._Urns.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor ICEBERG_READ: ManagedTransforms._Urns.ValueType # 0 ICEBERG_WRITE: ManagedTransforms._Urns.ValueType # 1 KAFKA_READ: ManagedTransforms._Urns.ValueType # 2 KAFKA_WRITE: ManagedTransforms._Urns.ValueType # 3 BIGQUERY_READ: ManagedTransforms._Urns.ValueType # 4 BIGQUERY_WRITE: ManagedTransforms._Urns.ValueType # 5 ICEBERG_CDC_READ: ManagedTransforms._Urns.ValueType # 6 POSTGRES_READ: ManagedTransforms._Urns.ValueType # 7 POSTGRES_WRITE: ManagedTransforms._Urns.ValueType # 8 MYSQL_READ: ManagedTransforms._Urns.ValueType # 9 MYSQL_WRITE: ManagedTransforms._Urns.ValueType # 10 SQL_SERVER_READ: ManagedTransforms._Urns.ValueType # 11 SQL_SERVER_WRITE: ManagedTransforms._Urns.ValueType # 12 class Urns(_Urns, metaclass=_UrnsEnumTypeWrapper): ... ICEBERG_READ: ManagedTransforms.Urns.ValueType # 0 ICEBERG_WRITE: ManagedTransforms.Urns.ValueType # 1 KAFKA_READ: ManagedTransforms.Urns.ValueType # 2 KAFKA_WRITE: ManagedTransforms.Urns.ValueType # 3 BIGQUERY_READ: ManagedTransforms.Urns.ValueType # 4 BIGQUERY_WRITE: ManagedTransforms.Urns.ValueType # 5 ICEBERG_CDC_READ: ManagedTransforms.Urns.ValueType # 6 POSTGRES_READ: ManagedTransforms.Urns.ValueType # 7 POSTGRES_WRITE: ManagedTransforms.Urns.ValueType # 8 MYSQL_READ: ManagedTransforms.Urns.ValueType # 9 MYSQL_WRITE: ManagedTransforms.Urns.ValueType # 10 SQL_SERVER_READ: ManagedTransforms.Urns.ValueType # 11 SQL_SERVER_WRITE: ManagedTransforms.Urns.ValueType # 12 def __init__( self, ) -> None: ... global___ManagedTransforms = ManagedTransforms @typing_extensions.final class JavaClassLookupPayload(google.protobuf.message.Message): """A configuration payload for an external transform. Used to define a Java transform that can be directly instantiated by a Java expansion service. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor CLASS_NAME_FIELD_NUMBER: builtins.int CONSTRUCTOR_METHOD_FIELD_NUMBER: builtins.int CONSTRUCTOR_SCHEMA_FIELD_NUMBER: builtins.int CONSTRUCTOR_PAYLOAD_FIELD_NUMBER: builtins.int BUILDER_METHODS_FIELD_NUMBER: builtins.int class_name: builtins.str """Name of the Java transform class.""" constructor_method: builtins.str """A static method to construct the initial instance of the transform. If not provided, the transform should be instantiated using a class constructor. """ @property def constructor_schema(self) -> org.apache.beam.model.pipeline.v1.schema_pb2.Schema: """The top level fields of the schema represent the method parameters in order. If able, top level field names are also verified against the method parameters for a match. Any field names in the form 'ignore[0-9]+' will not be used for validation hence that format can be used to represent arbitrary field names. """ constructor_payload: builtins.bytes """A payload which can be decoded using beam:coder:row:v1 and the provided constructor schema. """ @property def builder_methods(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___BuilderMethod]: """Set of builder methods and corresponding parameters to apply after the transform object is constructed. When constructing the transform object, given builder methods will be applied in order. """ def __init__( self, *, class_name: builtins.str | None = ..., constructor_method: builtins.str | None = ..., constructor_schema: org.apache.beam.model.pipeline.v1.schema_pb2.Schema | None = ..., constructor_payload: builtins.bytes | None = ..., builder_methods: collections.abc.Iterable[global___BuilderMethod] | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["constructor_schema", b"constructor_schema"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["builder_methods", b"builder_methods", "class_name", b"class_name", "constructor_method", b"constructor_method", "constructor_payload", b"constructor_payload", "constructor_schema", b"constructor_schema"]) -> None: ... global___JavaClassLookupPayload = JavaClassLookupPayload @typing_extensions.final class BuilderMethod(google.protobuf.message.Message): """This represents a builder method of the transform class that should be applied in-order after instantiating the initial transform object. Each builder method may take one or more parameters and has to return an instance of the transform object. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor NAME_FIELD_NUMBER: builtins.int SCHEMA_FIELD_NUMBER: builtins.int PAYLOAD_FIELD_NUMBER: builtins.int name: builtins.str """Name of the builder method""" @property def schema(self) -> org.apache.beam.model.pipeline.v1.schema_pb2.Schema: """The top level fields of the schema represent the method parameters in order. If able, top level field names are also verified against the method parameters for a match. Any field names in the form 'ignore[0-9]+' will not be used for validation hence that format can be used to represent arbitrary field names. """ payload: builtins.bytes """A payload which can be decoded using beam:coder:row:v1 and the builder method schema. """ def __init__( self, *, name: builtins.str | None = ..., schema: org.apache.beam.model.pipeline.v1.schema_pb2.Schema | None = ..., payload: builtins.bytes | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["schema", b"schema"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["name", b"name", "payload", b"payload", "schema", b"schema"]) -> None: ... global___BuilderMethod = BuilderMethod @typing_extensions.final class Annotations(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[Annotations._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor CONFIG_ROW_KEY: Annotations._Enum.ValueType # 0 """The annotation key for the encoded configuration Row used to build a transform""" CONFIG_ROW_SCHEMA_KEY: Annotations._Enum.ValueType # 1 """The annotation key for the configuration Schema used to decode the configuration Row""" SCHEMATRANSFORM_URN_KEY: Annotations._Enum.ValueType # 2 """If ths transform is a SchemaTransform, this is the annotation key for the SchemaTransform's URN""" MANAGED_UNDERLYING_TRANSFORM_URN_KEY: Annotations._Enum.ValueType # 3 """If the transform is a ManagedSchemaTransform, this is the annotation key for the underlying SchemaTransform's URN""" class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... CONFIG_ROW_KEY: Annotations.Enum.ValueType # 0 """The annotation key for the encoded configuration Row used to build a transform""" CONFIG_ROW_SCHEMA_KEY: Annotations.Enum.ValueType # 1 """The annotation key for the configuration Schema used to decode the configuration Row""" SCHEMATRANSFORM_URN_KEY: Annotations.Enum.ValueType # 2 """If ths transform is a SchemaTransform, this is the annotation key for the SchemaTransform's URN""" MANAGED_UNDERLYING_TRANSFORM_URN_KEY: Annotations.Enum.ValueType # 3 """If the transform is a ManagedSchemaTransform, this is the annotation key for the underlying SchemaTransform's URN""" def __init__( self, ) -> None: ... global___Annotations = Annotations @typing_extensions.final class SchemaTransformPayload(google.protobuf.message.Message): """Payload for a Schema-aware PTransform. This is a transform that is aware of its input and output PCollection schemas and is configured using Beam Schema-compatible parameters. The information available in the payload can be used to instantiate the schema-aware transform. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor IDENTIFIER_FIELD_NUMBER: builtins.int CONFIGURATION_SCHEMA_FIELD_NUMBER: builtins.int CONFIGURATION_ROW_FIELD_NUMBER: builtins.int identifier: builtins.str """The identifier of the SchemaTransform (typically a URN).""" @property def configuration_schema(self) -> org.apache.beam.model.pipeline.v1.schema_pb2.Schema: """The configuration schema of the SchemaTransform.""" configuration_row: builtins.bytes """The configuration of the SchemaTransform. Should be decodable via beam:coder:row:v1. The schema of the Row should be compatible with the schema of the SchemaTransform. """ def __init__( self, *, identifier: builtins.str | None = ..., configuration_schema: org.apache.beam.model.pipeline.v1.schema_pb2.Schema | None = ..., configuration_row: builtins.bytes | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["configuration_schema", b"configuration_schema"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["configuration_row", b"configuration_row", "configuration_schema", b"configuration_schema", "identifier", b"identifier"]) -> None: ... global___SchemaTransformPayload = SchemaTransformPayload ././@PaxHeader0000000000000000000000000000022000000000000010207 xustar00122 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/external_transforms_pb2_grpc.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/external_transforms0000644000175100017510000000023715206113151033166 0ustar00runnerrunner# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc ././@PaxHeader0000000000000000000000000000022000000000000010207 xustar00122 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/external_transforms_pb2_urns.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/external_transforms0000644000175100017510000000661715206113151033176 0ustar00runnerrunnerfrom ........utils import PropertiesFromEnumValue from . import metrics_pb2 EMPTY_MONITORING_INFO_LABEL_PROPS = metrics_pb2.MonitoringInfoLabelProps() EMPTY_MONITORING_INFO_SPEC = metrics_pb2.MonitoringInfoSpec() class Annotations(object): class Enum(object): CONFIG_ROW_KEY = PropertiesFromEnumValue('', 'config_row', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) CONFIG_ROW_SCHEMA_KEY = PropertiesFromEnumValue('', 'config_row_schema', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) SCHEMATRANSFORM_URN_KEY = PropertiesFromEnumValue('', 'schematransform_urn', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MANAGED_UNDERLYING_TRANSFORM_URN_KEY = PropertiesFromEnumValue('', 'managed_underlying_transform_urn', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class ExpansionMethods(object): class Enum(object): JAVA_CLASS_LOOKUP = PropertiesFromEnumValue('beam:expansion:payload:java_class_lookup:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) SCHEMA_TRANSFORM = PropertiesFromEnumValue('beam:expansion:payload:schematransform:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class ManagedTransforms(object): class Urns(object): ICEBERG_READ = PropertiesFromEnumValue('beam:schematransform:org.apache.beam:iceberg_read:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ICEBERG_WRITE = PropertiesFromEnumValue('beam:schematransform:org.apache.beam:iceberg_write:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) KAFKA_READ = PropertiesFromEnumValue('beam:schematransform:org.apache.beam:kafka_read:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) KAFKA_WRITE = PropertiesFromEnumValue('beam:schematransform:org.apache.beam:kafka_write:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) BIGQUERY_READ = PropertiesFromEnumValue('beam:schematransform:org.apache.beam:bigquery_storage_read:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) BIGQUERY_WRITE = PropertiesFromEnumValue('beam:schematransform:org.apache.beam:bigquery_write:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ICEBERG_CDC_READ = PropertiesFromEnumValue('beam:schematransform:org.apache.beam:iceberg_cdc_read:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) POSTGRES_READ = PropertiesFromEnumValue('beam:schematransform:org.apache.beam:postgres_read:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) POSTGRES_WRITE = PropertiesFromEnumValue('beam:schematransform:org.apache.beam:postgres_write:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MYSQL_READ = PropertiesFromEnumValue('beam:schematransform:org.apache.beam:mysql_read:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MYSQL_WRITE = PropertiesFromEnumValue('beam:schematransform:org.apache.beam:mysql_write:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) SQL_SERVER_READ = PropertiesFromEnumValue('beam:schematransform:org.apache.beam:sql_server_read:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) SQL_SERVER_WRITE = PropertiesFromEnumValue('beam:schematransform:org.apache.beam:sql_server_write:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/metrics_pb2.py0000644000175100017510000011467115206113151031736 0ustar00runnerrunner# -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # source: org/apache/beam/model/pipeline/v1/metrics.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() from . import beam_runner_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_beam__runner__api__pb2 from google.protobuf import descriptor_pb2 as google_dot_protobuf_dot_descriptor__pb2 from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n/org/apache/beam/model/pipeline/v1/metrics.proto\x12!org.apache.beam.model.pipeline.v1\x1a\x37org/apache/beam/model/pipeline/v1/beam_runner_api.proto\x1a google/protobuf/descriptor.proto\x1a\x1fgoogle/protobuf/timestamp.proto\"\x8c\x01\n\x12MonitoringInfoSpec\x12\x0b\n\x03urn\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\x17\n\x0frequired_labels\x18\x03 \x03(\t\x12\x42\n\x0b\x61nnotations\x18\x04 \x03(\x0b\x32-.org.apache.beam.model.pipeline.v1.Annotation\"(\n\nAnnotation\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t\"\xec\'\n\x13MonitoringInfoSpecs\"\xd4\'\n\x04\x45num\x12\xa7\x01\n\x0eUSER_SUM_INT64\x10\x00\x1a\x92\x01\xd2\xa7\xa7\x96\x06\x8b\x01\n\x1d\x62\x65\x61m:metric:user:sum_int64:v1\x12\x19\x62\x65\x61m:metrics:sum_int64:v1\x1a\nPTRANSFORM\x1a\tNAMESPACE\x1a\x04NAME\"2\n\x0b\x64\x65scription\x12#URN utilized to report user metric.\x12\xaa\x01\n\x0fUSER_SUM_DOUBLE\x10\x01\x1a\x94\x01\xd2\xa7\xa7\x96\x06\x8d\x01\n\x1e\x62\x65\x61m:metric:user:sum_double:v1\x12\x1a\x62\x65\x61m:metrics:sum_double:v1\x1a\nPTRANSFORM\x1a\tNAMESPACE\x1a\x04NAME\"2\n\x0b\x64\x65scription\x12#URN utilized to report user metric.\x12\xc2\x01\n\x17USER_DISTRIBUTION_INT64\x10\x02\x1a\xa4\x01\xd2\xa7\xa7\x96\x06\x9d\x01\n&beam:metric:user:distribution_int64:v1\x12\"beam:metrics:distribution_int64:v1\x1a\nPTRANSFORM\x1a\tNAMESPACE\x1a\x04NAME\"2\n\x0b\x64\x65scription\x12#URN utilized to report user metric.\x12\xc5\x01\n\x18USER_DISTRIBUTION_DOUBLE\x10\x03\x1a\xa6\x01\xd2\xa7\xa7\x96\x06\x9f\x01\n\'beam:metric:user:distribution_double:v1\x12#beam:metrics:distribution_double:v1\x1a\nPTRANSFORM\x1a\tNAMESPACE\x1a\x04NAME\"2\n\x0b\x64\x65scription\x12#URN utilized to report user metric.\x12\xb0\x01\n\x11USER_LATEST_INT64\x10\x04\x1a\x98\x01\xd2\xa7\xa7\x96\x06\x91\x01\n beam:metric:user:latest_int64:v1\x12\x1c\x62\x65\x61m:metrics:latest_int64:v1\x1a\nPTRANSFORM\x1a\tNAMESPACE\x1a\x04NAME\"2\n\x0b\x64\x65scription\x12#URN utilized to report user metric.\x12\xb3\x01\n\x12USER_LATEST_DOUBLE\x10\x05\x1a\x9a\x01\xd2\xa7\xa7\x96\x06\x93\x01\n!beam:metric:user:latest_double:v1\x12\x1d\x62\x65\x61m:metrics:latest_double:v1\x1a\nPTRANSFORM\x1a\tNAMESPACE\x1a\x04NAME\"2\n\x0b\x64\x65scription\x12#URN utilized to report user metric.\x12\xad\x01\n\x10USER_TOP_N_INT64\x10\x06\x1a\x96\x01\xd2\xa7\xa7\x96\x06\x8f\x01\n\x1f\x62\x65\x61m:metric:user:top_n_int64:v1\x12\x1b\x62\x65\x61m:metrics:top_n_int64:v1\x1a\nPTRANSFORM\x1a\tNAMESPACE\x1a\x04NAME\"2\n\x0b\x64\x65scription\x12#URN utilized to report user metric.\x12\xb0\x01\n\x11USER_TOP_N_DOUBLE\x10\x07\x1a\x98\x01\xd2\xa7\xa7\x96\x06\x91\x01\n beam:metric:user:top_n_double:v1\x12\x1c\x62\x65\x61m:metrics:top_n_double:v1\x1a\nPTRANSFORM\x1a\tNAMESPACE\x1a\x04NAME\"2\n\x0b\x64\x65scription\x12#URN utilized to report user metric.\x12\xb6\x01\n\x13USER_BOTTOM_N_INT64\x10\x08\x1a\x9c\x01\xd2\xa7\xa7\x96\x06\x95\x01\n\"beam:metric:user:bottom_n_int64:v1\x12\x1e\x62\x65\x61m:metrics:bottom_n_int64:v1\x1a\nPTRANSFORM\x1a\tNAMESPACE\x1a\x04NAME\"2\n\x0b\x64\x65scription\x12#URN utilized to report user metric.\x12\xb9\x01\n\x14USER_BOTTOM_N_DOUBLE\x10\t\x1a\x9e\x01\xd2\xa7\xa7\x96\x06\x97\x01\n#beam:metric:user:bottom_n_double:v1\x12\x1f\x62\x65\x61m:metrics:bottom_n_double:v1\x1a\nPTRANSFORM\x1a\tNAMESPACE\x1a\x04NAME\"2\n\x0b\x64\x65scription\x12#URN utilized to report user metric.\x12\xaa\x01\n\x0fUSER_SET_STRING\x10\x15\x1a\x94\x01\xd2\xa7\xa7\x96\x06\x8d\x01\n\x1e\x62\x65\x61m:metric:user:set_string:v1\x12\x1a\x62\x65\x61m:metrics:set_string:v1\x1a\nPTRANSFORM\x1a\tNAMESPACE\x1a\x04NAME\"2\n\x0b\x64\x65scription\x12#URN utilized to report user metric.\x12\xb0\x01\n\x11USER_BOUNDED_TRIE\x10\x16\x1a\x98\x01\xd2\xa7\xa7\x96\x06\x91\x01\n beam:metric:user:bounded_trie:v1\x12\x1c\x62\x65\x61m:metrics:bounded_trie:v1\x1a\nPTRANSFORM\x1a\tNAMESPACE\x1a\x04NAME\"2\n\x0b\x64\x65scription\x12#URN utilized to report user metric.\x12\xad\x01\n\rELEMENT_COUNT\x10\n\x1a\x99\x01\xd2\xa7\xa7\x96\x06\x92\x01\n\x1c\x62\x65\x61m:metric:element_count:v1\x12\x19\x62\x65\x61m:metrics:sum_int64:v1\x1a\x0bPCOLLECTION\"J\n\x0b\x64\x65scription\x12;The total elements output to a Pcollection by a PTransform.\x12\xcd\x02\n\x11SAMPLED_BYTE_SIZE\x10\x0b\x1a\xb5\x02\xd2\xa7\xa7\x96\x06\xae\x02\n beam:metric:sampled_byte_size:v1\x12\"beam:metrics:distribution_int64:v1\x1a\x0bPCOLLECTION\"\xd8\x01\n\x0b\x64\x65scription\x12\xc8\x01The total byte size and count of a sampled set (or all) of elements in the pcollection. Sampling is used because calculating the byte count involves serializing the elements which is CPU intensive.\x12\xd9\x01\n\x12START_BUNDLE_MSECS\x10\x0c\x1a\xc0\x01\xd2\xa7\xa7\x96\x06\xb9\x01\n6beam:metric:pardo_execution_time:start_bundle_msecs:v1\x12\x19\x62\x65\x61m:metrics:sum_int64:v1\x1a\nPTRANSFORM\"X\n\x0b\x64\x65scription\x12IThe total estimated execution time of the start bundlefunction in a pardo\x12\xdf\x01\n\x14PROCESS_BUNDLE_MSECS\x10\r\x1a\xc4\x01\xd2\xa7\xa7\x96\x06\xbd\x01\n8beam:metric:pardo_execution_time:process_bundle_msecs:v1\x12\x19\x62\x65\x61m:metrics:sum_int64:v1\x1a\nPTRANSFORM\"Z\n\x0b\x64\x65scription\x12KThe total estimated execution time of the process bundlefunction in a pardo\x12\xdd\x01\n\x13\x46INISH_BUNDLE_MSECS\x10\x0e\x1a\xc3\x01\xd2\xa7\xa7\x96\x06\xbc\x01\n7beam:metric:pardo_execution_time:finish_bundle_msecs:v1\x12\x19\x62\x65\x61m:metrics:sum_int64:v1\x1a\nPTRANSFORM\"Z\n\x0b\x64\x65scription\x12KThe total estimated execution time of the finish bundle function in a pardo\x12\xbb\x01\n\x0bTOTAL_MSECS\x10\x0f\x1a\xa9\x01\xd2\xa7\xa7\x96\x06\xa2\x01\n4beam:metric:ptransform_execution_time:total_msecs:v1\x12\x19\x62\x65\x61m:metrics:sum_int64:v1\x1a\nPTRANSFORM\"C\n\x0b\x64\x65scription\x12\x34The total estimated execution time of the ptransform\x12\x9f\x02\n\x0eWORK_REMAINING\x10\x10\x1a\x8a\x02\xd2\xa7\xa7\x96\x06\x83\x02\n,beam:metric:ptransform_progress:remaining:v1\x12\x18\x62\x65\x61m:metrics:progress:v1\x1a\nPTRANSFORM\"\xac\x01\n\x0b\x64\x65scription\x12\x9c\x01The remaining amount of work for each active element. Each active element represents an independent amount of work not shared with any other active element.\x12\x9f\x02\n\x0eWORK_COMPLETED\x10\x11\x1a\x8a\x02\xd2\xa7\xa7\x96\x06\x83\x02\n,beam:metric:ptransform_progress:completed:v1\x12\x18\x62\x65\x61m:metrics:progress:v1\x1a\nPTRANSFORM\"\xac\x01\n\x0b\x64\x65scription\x12\x9c\x01The remaining amount of work for each active element. Each active element represents an independent amount of work not shared with any other active element.\x12\xa8\x01\n\x17\x44\x41TA_CHANNEL_READ_INDEX\x10\x12\x1a\x8a\x01\xd2\xa7\xa7\x96\x06\x83\x01\n&beam:metric:data_channel:read_index:v1\x12\x19\x62\x65\x61m:metrics:sum_int64:v1\x1a\nPTRANSFORM\"2\n\x0b\x64\x65scription\x12#The read index of the data channel.\x12\x8a\x02\n\x11\x41PI_REQUEST_COUNT\x10\x13\x1a\xf2\x01\xd2\xa7\xa7\x96\x06\xeb\x01\n#beam:metric:io:api_request_count:v1\x12\x19\x62\x65\x61m:metrics:sum_int64:v1\x1a\x07SERVICE\x1a\x06METHOD\x1a\x08RESOURCE\x1a\nPTRANSFORM\x1a\x06STATUS\"b\n\x0b\x64\x65scription\x12SRequest counts with status made to IO service APIs to batch read or write elements.\"\x16\n\x0eprocess_metric\x12\x04true\x12\xb3\x02\n\x15\x41PI_REQUEST_LATENCIES\x10\x14\x1a\x97\x02\xd2\xa7\xa7\x96\x06\x90\x02\n\'beam:metric:io:api_request_latencies:v1\x12\x1f\x62\x65\x61m:metrics:histogram_int64:v1\x1a\x07SERVICE\x1a\x06METHOD\x1a\x08RESOURCE\x1a\nPTRANSFORM\"n\n\x0b\x64\x65scription\x12_Histogram counts for request latencies made to IO service APIs to batch read or write elements.\"\x15\n\x05units\x12\x0cMilliseconds\"\x16\n\x0eprocess_metric\x12\x04true\x12\xb8\x01\n\x0eUSER_HISTOGRAM\x10\x17\x1a\xa3\x01\xd2\xa7\xa7\x96\x06\x9c\x01\n#beam:metric:user:histogram_int64:v1\x12\x1f\x62\x65\x61m:metrics:histogram_int64:v1\x1a\nPTRANSFORM\x1a\tNAMESPACE\x1a\x04NAME\"7\n\x0b\x64\x65scription\x12(URN utilized to report histogram metric.\"(\n\x18MonitoringInfoLabelProps\x12\x0c\n\x04name\x18\x01 \x01(\t\"\xcc\x0b\n\x0eMonitoringInfo\x12\x0b\n\x03urn\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\x0f\n\x07payload\x18\x03 \x01(\x0c\x12M\n\x06labels\x18\x04 \x03(\x0b\x32=.org.apache.beam.model.pipeline.v1.MonitoringInfo.LabelsEntry\x12.\n\nstart_time\x18\x05 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x1a-\n\x0bLabelsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\xdf\t\n\x14MonitoringInfoLabels\x12!\n\tTRANSFORM\x10\x00\x1a\x12\xa2\xd4\xe0\xe5\x03\x0c\n\nPTRANSFORM\x12$\n\x0bPCOLLECTION\x10\x01\x1a\x13\xa2\xd4\xe0\xe5\x03\r\n\x0bPCOLLECTION\x12\x32\n\x12WINDOWING_STRATEGY\x10\x02\x1a\x1a\xa2\xd4\xe0\xe5\x03\x14\n\x12WINDOWING_STRATEGY\x12\x18\n\x05\x43ODER\x10\x03\x1a\r\xa2\xd4\xe0\xe5\x03\x07\n\x05\x43ODER\x12$\n\x0b\x45NVIRONMENT\x10\x04\x1a\x13\xa2\xd4\xe0\xe5\x03\r\n\x0b\x45NVIRONMENT\x12 \n\tNAMESPACE\x10\x05\x1a\x11\xa2\xd4\xe0\xe5\x03\x0b\n\tNAMESPACE\x12\x16\n\x04NAME\x10\x06\x1a\x0c\xa2\xd4\xe0\xe5\x03\x06\n\x04NAME\x12\x1c\n\x07SERVICE\x10\x07\x1a\x0f\xa2\xd4\xe0\xe5\x03\t\n\x07SERVICE\x12\x1a\n\x06METHOD\x10\x08\x1a\x0e\xa2\xd4\xe0\xe5\x03\x08\n\x06METHOD\x12\x1e\n\x08RESOURCE\x10\t\x1a\x10\xa2\xd4\xe0\xe5\x03\n\n\x08RESOURCE\x12\x1a\n\x06STATUS\x10\n\x1a\x0e\xa2\xd4\xe0\xe5\x03\x08\n\x06STATUS\x12\x34\n\x13\x42IGQUERY_PROJECT_ID\x10\x0b\x1a\x1b\xa2\xd4\xe0\xe5\x03\x15\n\x13\x42IGQUERY_PROJECT_ID\x12.\n\x10\x42IGQUERY_DATASET\x10\x0c\x1a\x18\xa2\xd4\xe0\xe5\x03\x12\n\x10\x42IGQUERY_DATASET\x12*\n\x0e\x42IGQUERY_TABLE\x10\r\x1a\x16\xa2\xd4\xe0\xe5\x03\x10\n\x0e\x42IGQUERY_TABLE\x12(\n\rBIGQUERY_VIEW\x10\x0e\x1a\x15\xa2\xd4\xe0\xe5\x03\x0f\n\rBIGQUERY_VIEW\x12\x34\n\x13\x42IGQUERY_QUERY_NAME\x10\x0f\x1a\x1b\xa2\xd4\xe0\xe5\x03\x15\n\x13\x42IGQUERY_QUERY_NAME\x12\"\n\nGCS_BUCKET\x10\x10\x1a\x12\xa2\xd4\xe0\xe5\x03\x0c\n\nGCS_BUCKET\x12*\n\x0eGCS_PROJECT_ID\x10\x11\x1a\x16\xa2\xd4\xe0\xe5\x03\x10\n\x0eGCS_PROJECT_ID\x12\x30\n\x11\x44\x41TASTORE_PROJECT\x10\x12\x1a\x19\xa2\xd4\xe0\xe5\x03\x13\n\x11\x44\x41TASTORE_PROJECT\x12\x34\n\x13\x44\x41TASTORE_NAMESPACE\x10\x13\x1a\x1b\xa2\xd4\xe0\xe5\x03\x15\n\x13\x44\x41TASTORE_NAMESPACE\x12\x34\n\x13\x42IGTABLE_PROJECT_ID\x10\x14\x1a\x1b\xa2\xd4\xe0\xe5\x03\x15\n\x13\x42IGTABLE_PROJECT_ID\x12$\n\x0bINSTANCE_ID\x10\x15\x1a\x13\xa2\xd4\xe0\xe5\x03\r\n\x0bINSTANCE_ID\x12\x1e\n\x08TABLE_ID\x10\x16\x1a\x10\xa2\xd4\xe0\xe5\x03\n\n\x08TABLE_ID\x12\x32\n\x12SPANNER_PROJECT_ID\x10\x17\x1a\x1a\xa2\xd4\xe0\xe5\x03\x14\n\x12SPANNER_PROJECT_ID\x12\x34\n\x13SPANNER_DATABASE_ID\x10\x18\x1a\x1b\xa2\xd4\xe0\xe5\x03\x15\n\x13SPANNER_DATABASE_ID\x12.\n\x10SPANNER_TABLE_ID\x10\x19\x1a\x18\xa2\xd4\xe0\xe5\x03\x12\n\x10SPANNER_TABLE_ID\x12\x34\n\x13SPANNER_INSTANCE_ID\x10\x1a\x1a\x1b\xa2\xd4\xe0\xe5\x03\x15\n\x13SPANNER_INSTANCE_ID\x12\x32\n\x12SPANNER_QUERY_NAME\x10\x1b\x1a\x1a\xa2\xd4\xe0\xe5\x03\x14\n\x12SPANNER_QUERY_NAME\x12\x30\n\x11PER_WORKER_METRIC\x10\x1c\x1a\x19\xa2\xd4\xe0\xe5\x03\x13\n\x11PER_WORKER_METRIC\"\xe4\x06\n\x16MonitoringInfoTypeUrns\"\xc9\x06\n\x04\x45num\x12\x33\n\x0eSUM_INT64_TYPE\x10\x00\x1a\x1f\xa2\xb4\xfa\xc2\x05\x19\x62\x65\x61m:metrics:sum_int64:v1\x12\x35\n\x0fSUM_DOUBLE_TYPE\x10\x01\x1a \xa2\xb4\xfa\xc2\x05\x1a\x62\x65\x61m:metrics:sum_double:v1\x12\x45\n\x17\x44ISTRIBUTION_INT64_TYPE\x10\x02\x1a(\xa2\xb4\xfa\xc2\x05\"beam:metrics:distribution_int64:v1\x12G\n\x18\x44ISTRIBUTION_DOUBLE_TYPE\x10\x03\x1a)\xa2\xb4\xfa\xc2\x05#beam:metrics:distribution_double:v1\x12\x39\n\x11LATEST_INT64_TYPE\x10\x04\x1a\"\xa2\xb4\xfa\xc2\x05\x1c\x62\x65\x61m:metrics:latest_int64:v1\x12;\n\x12LATEST_DOUBLE_TYPE\x10\x05\x1a#\xa2\xb4\xfa\xc2\x05\x1d\x62\x65\x61m:metrics:latest_double:v1\x12\x37\n\x10TOP_N_INT64_TYPE\x10\x06\x1a!\xa2\xb4\xfa\xc2\x05\x1b\x62\x65\x61m:metrics:top_n_int64:v1\x12\x39\n\x11TOP_N_DOUBLE_TYPE\x10\x07\x1a\"\xa2\xb4\xfa\xc2\x05\x1c\x62\x65\x61m:metrics:top_n_double:v1\x12=\n\x13\x42OTTOM_N_INT64_TYPE\x10\x08\x1a$\xa2\xb4\xfa\xc2\x05\x1e\x62\x65\x61m:metrics:bottom_n_int64:v1\x12?\n\x14\x42OTTOM_N_DOUBLE_TYPE\x10\t\x1a%\xa2\xb4\xfa\xc2\x05\x1f\x62\x65\x61m:metrics:bottom_n_double:v1\x12\x31\n\rPROGRESS_TYPE\x10\n\x1a\x1e\xa2\xb4\xfa\xc2\x05\x18\x62\x65\x61m:metrics:progress:v1\x12\x35\n\x0fSET_STRING_TYPE\x10\x0b\x1a \xa2\xb4\xfa\xc2\x05\x1a\x62\x65\x61m:metrics:set_string:v1\x12\x39\n\x11\x42OUNDED_TRIE_TYPE\x10\x0c\x1a\"\xa2\xb4\xfa\xc2\x05\x1c\x62\x65\x61m:metrics:bounded_trie:v1\x12\x34\n\tHISTOGRAM\x10\r\x1a%\xa2\xb4\xfa\xc2\x05\x1f\x62\x65\x61m:metrics:histogram_int64:v1\"\xdd\x01\n\x0f\x42oundedTrieNode\x12\x11\n\ttruncated\x18\x01 \x01(\x08\x12R\n\x08\x63hildren\x18\x02 \x03(\x0b\x32@.org.apache.beam.model.pipeline.v1.BoundedTrieNode.ChildrenEntry\x1a\x63\n\rChildrenEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x41\n\x05value\x18\x02 \x01(\x0b\x32\x32.org.apache.beam.model.pipeline.v1.BoundedTrieNode:\x02\x38\x01\"q\n\x0b\x42oundedTrie\x12\r\n\x05\x62ound\x18\x01 \x01(\x05\x12@\n\x04root\x18\x02 \x01(\x0b\x32\x32.org.apache.beam.model.pipeline.v1.BoundedTrieNode\x12\x11\n\tsingleton\x18\x03 \x03(\t\"\xf8\x04\n\x0eHistogramValue\x12\x12\n\x05\x63ount\x18\x01 \x01(\x03H\x00\x88\x01\x01\x12\\\n\x0e\x62ucket_options\x18\x02 \x01(\x0b\x32?.org.apache.beam.model.pipeline.v1.HistogramValue.BucketOptionsH\x01\x88\x01\x01\x12\x15\n\rbucket_counts\x18\x03 \x03(\x03\x1a\xbf\x03\n\rBucketOptions\x12X\n\x06linear\x18\x01 \x01(\x0b\x32\x46.org.apache.beam.model.pipeline.v1.HistogramValue.BucketOptions.LinearH\x00\x12\x64\n\x0b\x65xponential\x18\x02 \x01(\x0b\x32M.org.apache.beam.model.pipeline.v1.HistogramValue.BucketOptions.Base2ExponentH\x00\x1az\n\x06Linear\x12\x1e\n\x11number_of_buckets\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12\x12\n\x05width\x18\x02 \x01(\x01H\x01\x88\x01\x01\x12\x12\n\x05start\x18\x03 \x01(\x01H\x02\x88\x01\x01\x42\x14\n\x12_number_of_bucketsB\x08\n\x06_widthB\x08\n\x06_start\x1a\x63\n\rBase2Exponent\x12\x1e\n\x11number_of_buckets\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12\x12\n\x05scale\x18\x02 \x01(\x05H\x01\x88\x01\x01\x42\x14\n\x12_number_of_bucketsB\x08\n\x06_scaleB\r\n\x0b\x62ucket_typeB\x08\n\x06_countB\x11\n\x0f_bucket_options:v\n\x0blabel_props\x12!.google.protobuf.EnumValueOptions\x18\xc4\x8a\xdc< \x01(\x0b\x32;.org.apache.beam.model.pipeline.v1.MonitoringInfoLabelProps:y\n\x14monitoring_info_spec\x12!.google.protobuf.EnumValueOptions\x18\xfa\xf4\xe4\x62 \x01(\x0b\x32\x35.org.apache.beam.model.pipeline.v1.MonitoringInfoSpecBy\n!org.apache.beam.model.pipeline.v1B\nMetricsApiZHgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1;pipeline_v1b\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'org.apache.beam.model.pipeline.v1.metrics_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n!org.apache.beam.model.pipeline.v1B\nMetricsApiZHgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1;pipeline_v1' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_SUM_INT64"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_SUM_INT64"]._serialized_options = b'\322\247\247\226\006\213\001\n\035beam:metric:user:sum_int64:v1\022\031beam:metrics:sum_int64:v1\032\nPTRANSFORM\032\tNAMESPACE\032\004NAME\"2\n\013description\022#URN utilized to report user metric.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_SUM_DOUBLE"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_SUM_DOUBLE"]._serialized_options = b'\322\247\247\226\006\215\001\n\036beam:metric:user:sum_double:v1\022\032beam:metrics:sum_double:v1\032\nPTRANSFORM\032\tNAMESPACE\032\004NAME\"2\n\013description\022#URN utilized to report user metric.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_DISTRIBUTION_INT64"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_DISTRIBUTION_INT64"]._serialized_options = b'\322\247\247\226\006\235\001\n&beam:metric:user:distribution_int64:v1\022\"beam:metrics:distribution_int64:v1\032\nPTRANSFORM\032\tNAMESPACE\032\004NAME\"2\n\013description\022#URN utilized to report user metric.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_DISTRIBUTION_DOUBLE"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_DISTRIBUTION_DOUBLE"]._serialized_options = b'\322\247\247\226\006\237\001\n\'beam:metric:user:distribution_double:v1\022#beam:metrics:distribution_double:v1\032\nPTRANSFORM\032\tNAMESPACE\032\004NAME\"2\n\013description\022#URN utilized to report user metric.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_LATEST_INT64"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_LATEST_INT64"]._serialized_options = b'\322\247\247\226\006\221\001\n beam:metric:user:latest_int64:v1\022\034beam:metrics:latest_int64:v1\032\nPTRANSFORM\032\tNAMESPACE\032\004NAME\"2\n\013description\022#URN utilized to report user metric.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_LATEST_DOUBLE"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_LATEST_DOUBLE"]._serialized_options = b'\322\247\247\226\006\223\001\n!beam:metric:user:latest_double:v1\022\035beam:metrics:latest_double:v1\032\nPTRANSFORM\032\tNAMESPACE\032\004NAME\"2\n\013description\022#URN utilized to report user metric.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_TOP_N_INT64"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_TOP_N_INT64"]._serialized_options = b'\322\247\247\226\006\217\001\n\037beam:metric:user:top_n_int64:v1\022\033beam:metrics:top_n_int64:v1\032\nPTRANSFORM\032\tNAMESPACE\032\004NAME\"2\n\013description\022#URN utilized to report user metric.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_TOP_N_DOUBLE"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_TOP_N_DOUBLE"]._serialized_options = b'\322\247\247\226\006\221\001\n beam:metric:user:top_n_double:v1\022\034beam:metrics:top_n_double:v1\032\nPTRANSFORM\032\tNAMESPACE\032\004NAME\"2\n\013description\022#URN utilized to report user metric.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_BOTTOM_N_INT64"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_BOTTOM_N_INT64"]._serialized_options = b'\322\247\247\226\006\225\001\n\"beam:metric:user:bottom_n_int64:v1\022\036beam:metrics:bottom_n_int64:v1\032\nPTRANSFORM\032\tNAMESPACE\032\004NAME\"2\n\013description\022#URN utilized to report user metric.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_BOTTOM_N_DOUBLE"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_BOTTOM_N_DOUBLE"]._serialized_options = b'\322\247\247\226\006\227\001\n#beam:metric:user:bottom_n_double:v1\022\037beam:metrics:bottom_n_double:v1\032\nPTRANSFORM\032\tNAMESPACE\032\004NAME\"2\n\013description\022#URN utilized to report user metric.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_SET_STRING"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_SET_STRING"]._serialized_options = b'\322\247\247\226\006\215\001\n\036beam:metric:user:set_string:v1\022\032beam:metrics:set_string:v1\032\nPTRANSFORM\032\tNAMESPACE\032\004NAME\"2\n\013description\022#URN utilized to report user metric.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_BOUNDED_TRIE"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_BOUNDED_TRIE"]._serialized_options = b'\322\247\247\226\006\221\001\n beam:metric:user:bounded_trie:v1\022\034beam:metrics:bounded_trie:v1\032\nPTRANSFORM\032\tNAMESPACE\032\004NAME\"2\n\013description\022#URN utilized to report user metric.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["ELEMENT_COUNT"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["ELEMENT_COUNT"]._serialized_options = b'\322\247\247\226\006\222\001\n\034beam:metric:element_count:v1\022\031beam:metrics:sum_int64:v1\032\013PCOLLECTION\"J\n\013description\022;The total elements output to a Pcollection by a PTransform.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["SAMPLED_BYTE_SIZE"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["SAMPLED_BYTE_SIZE"]._serialized_options = b'\322\247\247\226\006\256\002\n beam:metric:sampled_byte_size:v1\022\"beam:metrics:distribution_int64:v1\032\013PCOLLECTION\"\330\001\n\013description\022\310\001The total byte size and count of a sampled set (or all) of elements in the pcollection. Sampling is used because calculating the byte count involves serializing the elements which is CPU intensive.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["START_BUNDLE_MSECS"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["START_BUNDLE_MSECS"]._serialized_options = b'\322\247\247\226\006\271\001\n6beam:metric:pardo_execution_time:start_bundle_msecs:v1\022\031beam:metrics:sum_int64:v1\032\nPTRANSFORM\"X\n\013description\022IThe total estimated execution time of the start bundlefunction in a pardo' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["PROCESS_BUNDLE_MSECS"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["PROCESS_BUNDLE_MSECS"]._serialized_options = b'\322\247\247\226\006\275\001\n8beam:metric:pardo_execution_time:process_bundle_msecs:v1\022\031beam:metrics:sum_int64:v1\032\nPTRANSFORM\"Z\n\013description\022KThe total estimated execution time of the process bundlefunction in a pardo' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["FINISH_BUNDLE_MSECS"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["FINISH_BUNDLE_MSECS"]._serialized_options = b'\322\247\247\226\006\274\001\n7beam:metric:pardo_execution_time:finish_bundle_msecs:v1\022\031beam:metrics:sum_int64:v1\032\nPTRANSFORM\"Z\n\013description\022KThe total estimated execution time of the finish bundle function in a pardo' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["TOTAL_MSECS"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["TOTAL_MSECS"]._serialized_options = b'\322\247\247\226\006\242\001\n4beam:metric:ptransform_execution_time:total_msecs:v1\022\031beam:metrics:sum_int64:v1\032\nPTRANSFORM\"C\n\013description\0224The total estimated execution time of the ptransform' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["WORK_REMAINING"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["WORK_REMAINING"]._serialized_options = b'\322\247\247\226\006\203\002\n,beam:metric:ptransform_progress:remaining:v1\022\030beam:metrics:progress:v1\032\nPTRANSFORM\"\254\001\n\013description\022\234\001The remaining amount of work for each active element. Each active element represents an independent amount of work not shared with any other active element.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["WORK_COMPLETED"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["WORK_COMPLETED"]._serialized_options = b'\322\247\247\226\006\203\002\n,beam:metric:ptransform_progress:completed:v1\022\030beam:metrics:progress:v1\032\nPTRANSFORM\"\254\001\n\013description\022\234\001The remaining amount of work for each active element. Each active element represents an independent amount of work not shared with any other active element.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["DATA_CHANNEL_READ_INDEX"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["DATA_CHANNEL_READ_INDEX"]._serialized_options = b'\322\247\247\226\006\203\001\n&beam:metric:data_channel:read_index:v1\022\031beam:metrics:sum_int64:v1\032\nPTRANSFORM\"2\n\013description\022#The read index of the data channel.' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["API_REQUEST_COUNT"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["API_REQUEST_COUNT"]._serialized_options = b'\322\247\247\226\006\353\001\n#beam:metric:io:api_request_count:v1\022\031beam:metrics:sum_int64:v1\032\007SERVICE\032\006METHOD\032\010RESOURCE\032\nPTRANSFORM\032\006STATUS\"b\n\013description\022SRequest counts with status made to IO service APIs to batch read or write elements.\"\026\n\016process_metric\022\004true' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["API_REQUEST_LATENCIES"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["API_REQUEST_LATENCIES"]._serialized_options = b'\322\247\247\226\006\220\002\n\'beam:metric:io:api_request_latencies:v1\022\037beam:metrics:histogram_int64:v1\032\007SERVICE\032\006METHOD\032\010RESOURCE\032\nPTRANSFORM\"n\n\013description\022_Histogram counts for request latencies made to IO service APIs to batch read or write elements.\"\025\n\005units\022\014Milliseconds\"\026\n\016process_metric\022\004true' _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_HISTOGRAM"]._options = None _globals['_MONITORINGINFOSPECS_ENUM'].values_by_name["USER_HISTOGRAM"]._serialized_options = b'\322\247\247\226\006\234\001\n#beam:metric:user:histogram_int64:v1\022\037beam:metrics:histogram_int64:v1\032\nPTRANSFORM\032\tNAMESPACE\032\004NAME\"7\n\013description\022(URN utilized to report histogram metric.' _globals['_MONITORINGINFO_LABELSENTRY']._options = None _globals['_MONITORINGINFO_LABELSENTRY']._serialized_options = b'8\001' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["TRANSFORM"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["TRANSFORM"]._serialized_options = b'\242\324\340\345\003\014\n\nPTRANSFORM' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["PCOLLECTION"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["PCOLLECTION"]._serialized_options = b'\242\324\340\345\003\r\n\013PCOLLECTION' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["WINDOWING_STRATEGY"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["WINDOWING_STRATEGY"]._serialized_options = b'\242\324\340\345\003\024\n\022WINDOWING_STRATEGY' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["CODER"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["CODER"]._serialized_options = b'\242\324\340\345\003\007\n\005CODER' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["ENVIRONMENT"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["ENVIRONMENT"]._serialized_options = b'\242\324\340\345\003\r\n\013ENVIRONMENT' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["NAMESPACE"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["NAMESPACE"]._serialized_options = b'\242\324\340\345\003\013\n\tNAMESPACE' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["NAME"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["NAME"]._serialized_options = b'\242\324\340\345\003\006\n\004NAME' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["SERVICE"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["SERVICE"]._serialized_options = b'\242\324\340\345\003\t\n\007SERVICE' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["METHOD"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["METHOD"]._serialized_options = b'\242\324\340\345\003\010\n\006METHOD' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["RESOURCE"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["RESOURCE"]._serialized_options = b'\242\324\340\345\003\n\n\010RESOURCE' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["STATUS"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["STATUS"]._serialized_options = b'\242\324\340\345\003\010\n\006STATUS' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["BIGQUERY_PROJECT_ID"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["BIGQUERY_PROJECT_ID"]._serialized_options = b'\242\324\340\345\003\025\n\023BIGQUERY_PROJECT_ID' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["BIGQUERY_DATASET"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["BIGQUERY_DATASET"]._serialized_options = b'\242\324\340\345\003\022\n\020BIGQUERY_DATASET' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["BIGQUERY_TABLE"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["BIGQUERY_TABLE"]._serialized_options = b'\242\324\340\345\003\020\n\016BIGQUERY_TABLE' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["BIGQUERY_VIEW"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["BIGQUERY_VIEW"]._serialized_options = b'\242\324\340\345\003\017\n\rBIGQUERY_VIEW' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["BIGQUERY_QUERY_NAME"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["BIGQUERY_QUERY_NAME"]._serialized_options = b'\242\324\340\345\003\025\n\023BIGQUERY_QUERY_NAME' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["GCS_BUCKET"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["GCS_BUCKET"]._serialized_options = b'\242\324\340\345\003\014\n\nGCS_BUCKET' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["GCS_PROJECT_ID"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["GCS_PROJECT_ID"]._serialized_options = b'\242\324\340\345\003\020\n\016GCS_PROJECT_ID' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["DATASTORE_PROJECT"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["DATASTORE_PROJECT"]._serialized_options = b'\242\324\340\345\003\023\n\021DATASTORE_PROJECT' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["DATASTORE_NAMESPACE"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["DATASTORE_NAMESPACE"]._serialized_options = b'\242\324\340\345\003\025\n\023DATASTORE_NAMESPACE' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["BIGTABLE_PROJECT_ID"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["BIGTABLE_PROJECT_ID"]._serialized_options = b'\242\324\340\345\003\025\n\023BIGTABLE_PROJECT_ID' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["INSTANCE_ID"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["INSTANCE_ID"]._serialized_options = b'\242\324\340\345\003\r\n\013INSTANCE_ID' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["TABLE_ID"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["TABLE_ID"]._serialized_options = b'\242\324\340\345\003\n\n\010TABLE_ID' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["SPANNER_PROJECT_ID"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["SPANNER_PROJECT_ID"]._serialized_options = b'\242\324\340\345\003\024\n\022SPANNER_PROJECT_ID' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["SPANNER_DATABASE_ID"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["SPANNER_DATABASE_ID"]._serialized_options = b'\242\324\340\345\003\025\n\023SPANNER_DATABASE_ID' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["SPANNER_TABLE_ID"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["SPANNER_TABLE_ID"]._serialized_options = b'\242\324\340\345\003\022\n\020SPANNER_TABLE_ID' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["SPANNER_INSTANCE_ID"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["SPANNER_INSTANCE_ID"]._serialized_options = b'\242\324\340\345\003\025\n\023SPANNER_INSTANCE_ID' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["SPANNER_QUERY_NAME"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["SPANNER_QUERY_NAME"]._serialized_options = b'\242\324\340\345\003\024\n\022SPANNER_QUERY_NAME' _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["PER_WORKER_METRIC"]._options = None _globals['_MONITORINGINFO_MONITORINGINFOLABELS'].values_by_name["PER_WORKER_METRIC"]._serialized_options = b'\242\324\340\345\003\023\n\021PER_WORKER_METRIC' _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["SUM_INT64_TYPE"]._options = None _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["SUM_INT64_TYPE"]._serialized_options = b'\242\264\372\302\005\031beam:metrics:sum_int64:v1' _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["SUM_DOUBLE_TYPE"]._options = None _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["SUM_DOUBLE_TYPE"]._serialized_options = b'\242\264\372\302\005\032beam:metrics:sum_double:v1' _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["DISTRIBUTION_INT64_TYPE"]._options = None _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["DISTRIBUTION_INT64_TYPE"]._serialized_options = b'\242\264\372\302\005\"beam:metrics:distribution_int64:v1' _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["DISTRIBUTION_DOUBLE_TYPE"]._options = None _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["DISTRIBUTION_DOUBLE_TYPE"]._serialized_options = b'\242\264\372\302\005#beam:metrics:distribution_double:v1' _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["LATEST_INT64_TYPE"]._options = None _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["LATEST_INT64_TYPE"]._serialized_options = b'\242\264\372\302\005\034beam:metrics:latest_int64:v1' _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["LATEST_DOUBLE_TYPE"]._options = None _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["LATEST_DOUBLE_TYPE"]._serialized_options = b'\242\264\372\302\005\035beam:metrics:latest_double:v1' _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["TOP_N_INT64_TYPE"]._options = None _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["TOP_N_INT64_TYPE"]._serialized_options = b'\242\264\372\302\005\033beam:metrics:top_n_int64:v1' _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["TOP_N_DOUBLE_TYPE"]._options = None _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["TOP_N_DOUBLE_TYPE"]._serialized_options = b'\242\264\372\302\005\034beam:metrics:top_n_double:v1' _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["BOTTOM_N_INT64_TYPE"]._options = None _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["BOTTOM_N_INT64_TYPE"]._serialized_options = b'\242\264\372\302\005\036beam:metrics:bottom_n_int64:v1' _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["BOTTOM_N_DOUBLE_TYPE"]._options = None _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["BOTTOM_N_DOUBLE_TYPE"]._serialized_options = b'\242\264\372\302\005\037beam:metrics:bottom_n_double:v1' _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["PROGRESS_TYPE"]._options = None _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["PROGRESS_TYPE"]._serialized_options = b'\242\264\372\302\005\030beam:metrics:progress:v1' _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["SET_STRING_TYPE"]._options = None _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["SET_STRING_TYPE"]._serialized_options = b'\242\264\372\302\005\032beam:metrics:set_string:v1' _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["BOUNDED_TRIE_TYPE"]._options = None _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["BOUNDED_TRIE_TYPE"]._serialized_options = b'\242\264\372\302\005\034beam:metrics:bounded_trie:v1' _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["HISTOGRAM"]._options = None _globals['_MONITORINGINFOTYPEURNS_ENUM'].values_by_name["HISTOGRAM"]._serialized_options = b'\242\264\372\302\005\037beam:metrics:histogram_int64:v1' _globals['_BOUNDEDTRIENODE_CHILDRENENTRY']._options = None _globals['_BOUNDEDTRIENODE_CHILDRENENTRY']._serialized_options = b'8\001' _globals['_MONITORINGINFOSPEC']._serialized_start=211 _globals['_MONITORINGINFOSPEC']._serialized_end=351 _globals['_ANNOTATION']._serialized_start=353 _globals['_ANNOTATION']._serialized_end=393 _globals['_MONITORINGINFOSPECS']._serialized_start=396 _globals['_MONITORINGINFOSPECS']._serialized_end=5496 _globals['_MONITORINGINFOSPECS_ENUM']._serialized_start=420 _globals['_MONITORINGINFOSPECS_ENUM']._serialized_end=5496 _globals['_MONITORINGINFOLABELPROPS']._serialized_start=5498 _globals['_MONITORINGINFOLABELPROPS']._serialized_end=5538 _globals['_MONITORINGINFO']._serialized_start=5541 _globals['_MONITORINGINFO']._serialized_end=7025 _globals['_MONITORINGINFO_LABELSENTRY']._serialized_start=5730 _globals['_MONITORINGINFO_LABELSENTRY']._serialized_end=5775 _globals['_MONITORINGINFO_MONITORINGINFOLABELS']._serialized_start=5778 _globals['_MONITORINGINFO_MONITORINGINFOLABELS']._serialized_end=7025 _globals['_MONITORINGINFOTYPEURNS']._serialized_start=7028 _globals['_MONITORINGINFOTYPEURNS']._serialized_end=7896 _globals['_MONITORINGINFOTYPEURNS_ENUM']._serialized_start=7055 _globals['_MONITORINGINFOTYPEURNS_ENUM']._serialized_end=7896 _globals['_BOUNDEDTRIENODE']._serialized_start=7899 _globals['_BOUNDEDTRIENODE']._serialized_end=8120 _globals['_BOUNDEDTRIENODE_CHILDRENENTRY']._serialized_start=8021 _globals['_BOUNDEDTRIENODE_CHILDRENENTRY']._serialized_end=8120 _globals['_BOUNDEDTRIE']._serialized_start=8122 _globals['_BOUNDEDTRIE']._serialized_end=8235 _globals['_HISTOGRAMVALUE']._serialized_start=8238 _globals['_HISTOGRAMVALUE']._serialized_end=8870 _globals['_HISTOGRAMVALUE_BUCKETOPTIONS']._serialized_start=8394 _globals['_HISTOGRAMVALUE_BUCKETOPTIONS']._serialized_end=8841 _globals['_HISTOGRAMVALUE_BUCKETOPTIONS_LINEAR']._serialized_start=8603 _globals['_HISTOGRAMVALUE_BUCKETOPTIONS_LINEAR']._serialized_end=8725 _globals['_HISTOGRAMVALUE_BUCKETOPTIONS_BASE2EXPONENT']._serialized_start=8727 _globals['_HISTOGRAMVALUE_BUCKETOPTIONS_BASE2EXPONENT']._serialized_end=8826 # @@protoc_insertion_point(module_scope) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/metrics_pb2.pyi0000644000175100017510000012640015206113151032100 0ustar00runnerrunner""" @generated by mypy-protobuf. Do not edit manually! isort:skip_file Protocol Buffers for metrics classes, used in the Fn API, Job API, and by SDKs. """ import builtins import collections.abc import google.protobuf.descriptor import google.protobuf.descriptor_pb2 import google.protobuf.internal.containers import google.protobuf.internal.enum_type_wrapper import google.protobuf.internal.extension_dict import google.protobuf.message import google.protobuf.timestamp_pb2 import sys import typing if sys.version_info >= (3, 10): import typing as typing_extensions else: import typing_extensions DESCRIPTOR: google.protobuf.descriptor.FileDescriptor @typing_extensions.final class MonitoringInfoSpec(google.protobuf.message.Message): """A specification for describing a well known MonitoringInfo. All specifications are uniquely identified by the urn. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor URN_FIELD_NUMBER: builtins.int TYPE_FIELD_NUMBER: builtins.int REQUIRED_LABELS_FIELD_NUMBER: builtins.int ANNOTATIONS_FIELD_NUMBER: builtins.int urn: builtins.str """Defines the semantic meaning of the metric or monitored state. See MonitoringInfoSpecs.Enum for the set of well known metrics/monitored state. """ type: builtins.str """Defines the required encoding and aggregation method for the payload. See MonitoringInfoTypeUrns.Enum for the set of well known types. """ @property def required_labels(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """The list of required labels for the specified urn and type.""" @property def annotations(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Annotation]: """Extra non functional parts of the spec for descriptive purposes. i.e. description, units, etc. """ def __init__( self, *, urn: builtins.str | None = ..., type: builtins.str | None = ..., required_labels: collections.abc.Iterable[builtins.str] | None = ..., annotations: collections.abc.Iterable[global___Annotation] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["annotations", b"annotations", "required_labels", b"required_labels", "type", b"type", "urn", b"urn"]) -> None: ... global___MonitoringInfoSpec = MonitoringInfoSpec @typing_extensions.final class Annotation(google.protobuf.message.Message): """The key name and value string of MonitoringInfo annotations.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str value: builtins.str def __init__( self, *, key: builtins.str | None = ..., value: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... global___Annotation = Annotation @typing_extensions.final class MonitoringInfoSpecs(google.protobuf.message.Message): """A set of well known MonitoringInfo specifications.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[MonitoringInfoSpecs._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor USER_SUM_INT64: MonitoringInfoSpecs._Enum.ValueType # 0 """Represents an integer counter where values are summed across bundles.""" USER_SUM_DOUBLE: MonitoringInfoSpecs._Enum.ValueType # 1 """Represents a double counter where values are summed across bundles.""" USER_DISTRIBUTION_INT64: MonitoringInfoSpecs._Enum.ValueType # 2 """Represents a distribution of an integer value where: - count: represents the number of values seen across all bundles - sum: represents the total of the value across all bundles - min: represents the smallest value seen across all bundles - max: represents the largest value seen across all bundles """ USER_DISTRIBUTION_DOUBLE: MonitoringInfoSpecs._Enum.ValueType # 3 """Represents a distribution of a double value where: - count: represents the number of values seen across all bundles - sum: represents the total of the value across all bundles - min: represents the smallest value seen across all bundles - max: represents the largest value seen across all bundles """ USER_LATEST_INT64: MonitoringInfoSpecs._Enum.ValueType # 4 """Represents the latest seen integer value. The timestamp is used to provide an "ordering" over multiple values to determine which is the latest. """ USER_LATEST_DOUBLE: MonitoringInfoSpecs._Enum.ValueType # 5 """Represents the latest seen double value. The timestamp is used to provide an "ordering" over multiple values to determine which is the latest. """ USER_TOP_N_INT64: MonitoringInfoSpecs._Enum.ValueType # 6 """Represents the largest set of integer values seen across bundles.""" USER_TOP_N_DOUBLE: MonitoringInfoSpecs._Enum.ValueType # 7 """Represents the largest set of double values seen across bundles.""" USER_BOTTOM_N_INT64: MonitoringInfoSpecs._Enum.ValueType # 8 """Represents the smallest set of integer values seen across bundles.""" USER_BOTTOM_N_DOUBLE: MonitoringInfoSpecs._Enum.ValueType # 9 """Represents the smallest set of double values seen across bundles.""" USER_SET_STRING: MonitoringInfoSpecs._Enum.ValueType # 21 """Represents a set of strings seen across bundles.""" USER_BOUNDED_TRIE: MonitoringInfoSpecs._Enum.ValueType # 22 """Represents a set of strings seen across bundles.""" ELEMENT_COUNT: MonitoringInfoSpecs._Enum.ValueType # 10 """General monitored state information which contains structured information which does not fit into a typical metric format. See MonitoringTableData for more details. TODO(https://github.com/apache/beam/issues/20122): Support monitored state. USER_MONITORING_TABLE = XX [(monitoring_info_spec) = { urn: "beam:metric:user:v1", type: "beam:metrics:monitoring_table:v1", required_labels: ["PTRANSFORM", "NAMESPACE", "NAME"], annotations: [{ key: "description", value: "URN utilized to report user monitoring data." }] }]; """ SAMPLED_BYTE_SIZE: MonitoringInfoSpecs._Enum.ValueType # 11 START_BUNDLE_MSECS: MonitoringInfoSpecs._Enum.ValueType # 12 PROCESS_BUNDLE_MSECS: MonitoringInfoSpecs._Enum.ValueType # 13 FINISH_BUNDLE_MSECS: MonitoringInfoSpecs._Enum.ValueType # 14 TOTAL_MSECS: MonitoringInfoSpecs._Enum.ValueType # 15 WORK_REMAINING: MonitoringInfoSpecs._Enum.ValueType # 16 """All values reported across all beam:metric:ptransform_progress:.*:v1 metrics are of the same magnitude. """ WORK_COMPLETED: MonitoringInfoSpecs._Enum.ValueType # 17 """All values reported across all beam:metric:ptransform_progress:.*:v1 metrics are of the same magnitude. """ DATA_CHANNEL_READ_INDEX: MonitoringInfoSpecs._Enum.ValueType # 18 """The (0-based) index of the latest item processed from the data channel. This gives an indication of the SDKs progress through the data channel, and is a lower bound on where it is able to split. For an SDK that processes items sequentially, this is equivalently the number of items fully processed (or -1 if processing has not yet started). """ API_REQUEST_COUNT: MonitoringInfoSpecs._Enum.ValueType # 19 API_REQUEST_LATENCIES: MonitoringInfoSpecs._Enum.ValueType # 20 USER_HISTOGRAM: MonitoringInfoSpecs._Enum.ValueType # 23 class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... USER_SUM_INT64: MonitoringInfoSpecs.Enum.ValueType # 0 """Represents an integer counter where values are summed across bundles.""" USER_SUM_DOUBLE: MonitoringInfoSpecs.Enum.ValueType # 1 """Represents a double counter where values are summed across bundles.""" USER_DISTRIBUTION_INT64: MonitoringInfoSpecs.Enum.ValueType # 2 """Represents a distribution of an integer value where: - count: represents the number of values seen across all bundles - sum: represents the total of the value across all bundles - min: represents the smallest value seen across all bundles - max: represents the largest value seen across all bundles """ USER_DISTRIBUTION_DOUBLE: MonitoringInfoSpecs.Enum.ValueType # 3 """Represents a distribution of a double value where: - count: represents the number of values seen across all bundles - sum: represents the total of the value across all bundles - min: represents the smallest value seen across all bundles - max: represents the largest value seen across all bundles """ USER_LATEST_INT64: MonitoringInfoSpecs.Enum.ValueType # 4 """Represents the latest seen integer value. The timestamp is used to provide an "ordering" over multiple values to determine which is the latest. """ USER_LATEST_DOUBLE: MonitoringInfoSpecs.Enum.ValueType # 5 """Represents the latest seen double value. The timestamp is used to provide an "ordering" over multiple values to determine which is the latest. """ USER_TOP_N_INT64: MonitoringInfoSpecs.Enum.ValueType # 6 """Represents the largest set of integer values seen across bundles.""" USER_TOP_N_DOUBLE: MonitoringInfoSpecs.Enum.ValueType # 7 """Represents the largest set of double values seen across bundles.""" USER_BOTTOM_N_INT64: MonitoringInfoSpecs.Enum.ValueType # 8 """Represents the smallest set of integer values seen across bundles.""" USER_BOTTOM_N_DOUBLE: MonitoringInfoSpecs.Enum.ValueType # 9 """Represents the smallest set of double values seen across bundles.""" USER_SET_STRING: MonitoringInfoSpecs.Enum.ValueType # 21 """Represents a set of strings seen across bundles.""" USER_BOUNDED_TRIE: MonitoringInfoSpecs.Enum.ValueType # 22 """Represents a set of strings seen across bundles.""" ELEMENT_COUNT: MonitoringInfoSpecs.Enum.ValueType # 10 """General monitored state information which contains structured information which does not fit into a typical metric format. See MonitoringTableData for more details. TODO(https://github.com/apache/beam/issues/20122): Support monitored state. USER_MONITORING_TABLE = XX [(monitoring_info_spec) = { urn: "beam:metric:user:v1", type: "beam:metrics:monitoring_table:v1", required_labels: ["PTRANSFORM", "NAMESPACE", "NAME"], annotations: [{ key: "description", value: "URN utilized to report user monitoring data." }] }]; """ SAMPLED_BYTE_SIZE: MonitoringInfoSpecs.Enum.ValueType # 11 START_BUNDLE_MSECS: MonitoringInfoSpecs.Enum.ValueType # 12 PROCESS_BUNDLE_MSECS: MonitoringInfoSpecs.Enum.ValueType # 13 FINISH_BUNDLE_MSECS: MonitoringInfoSpecs.Enum.ValueType # 14 TOTAL_MSECS: MonitoringInfoSpecs.Enum.ValueType # 15 WORK_REMAINING: MonitoringInfoSpecs.Enum.ValueType # 16 """All values reported across all beam:metric:ptransform_progress:.*:v1 metrics are of the same magnitude. """ WORK_COMPLETED: MonitoringInfoSpecs.Enum.ValueType # 17 """All values reported across all beam:metric:ptransform_progress:.*:v1 metrics are of the same magnitude. """ DATA_CHANNEL_READ_INDEX: MonitoringInfoSpecs.Enum.ValueType # 18 """The (0-based) index of the latest item processed from the data channel. This gives an indication of the SDKs progress through the data channel, and is a lower bound on where it is able to split. For an SDK that processes items sequentially, this is equivalently the number of items fully processed (or -1 if processing has not yet started). """ API_REQUEST_COUNT: MonitoringInfoSpecs.Enum.ValueType # 19 API_REQUEST_LATENCIES: MonitoringInfoSpecs.Enum.ValueType # 20 USER_HISTOGRAM: MonitoringInfoSpecs.Enum.ValueType # 23 def __init__( self, ) -> None: ... global___MonitoringInfoSpecs = MonitoringInfoSpecs @typing_extensions.final class MonitoringInfoLabelProps(google.protobuf.message.Message): """A set of properties for the MonitoringInfoLabel, this is useful to obtain the proper label string for the MonitoringInfoLabel. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor NAME_FIELD_NUMBER: builtins.int name: builtins.str """The label key to use in the MonitoringInfo labels map.""" def __init__( self, *, name: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["name", b"name"]) -> None: ... global___MonitoringInfoLabelProps = MonitoringInfoLabelProps @typing_extensions.final class MonitoringInfo(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor class _MonitoringInfoLabels: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _MonitoringInfoLabelsEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[MonitoringInfo._MonitoringInfoLabels.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor TRANSFORM: MonitoringInfo._MonitoringInfoLabels.ValueType # 0 """The values used for TRANSFORM, PCOLLECTION, WINDOWING_STRATEGY CODER, ENVIRONMENT, etc. must always match the keys used to refer to them. For actively processed bundles, these should match the values within the ProcessBundleDescriptor. For job management APIs, these should match values within the original pipeline representation. """ PCOLLECTION: MonitoringInfo._MonitoringInfoLabels.ValueType # 1 WINDOWING_STRATEGY: MonitoringInfo._MonitoringInfoLabels.ValueType # 2 CODER: MonitoringInfo._MonitoringInfoLabels.ValueType # 3 ENVIRONMENT: MonitoringInfo._MonitoringInfoLabels.ValueType # 4 NAMESPACE: MonitoringInfo._MonitoringInfoLabels.ValueType # 5 NAME: MonitoringInfo._MonitoringInfoLabels.ValueType # 6 SERVICE: MonitoringInfo._MonitoringInfoLabels.ValueType # 7 METHOD: MonitoringInfo._MonitoringInfoLabels.ValueType # 8 RESOURCE: MonitoringInfo._MonitoringInfoLabels.ValueType # 9 STATUS: MonitoringInfo._MonitoringInfoLabels.ValueType # 10 BIGQUERY_PROJECT_ID: MonitoringInfo._MonitoringInfoLabels.ValueType # 11 BIGQUERY_DATASET: MonitoringInfo._MonitoringInfoLabels.ValueType # 12 BIGQUERY_TABLE: MonitoringInfo._MonitoringInfoLabels.ValueType # 13 BIGQUERY_VIEW: MonitoringInfo._MonitoringInfoLabels.ValueType # 14 BIGQUERY_QUERY_NAME: MonitoringInfo._MonitoringInfoLabels.ValueType # 15 GCS_BUCKET: MonitoringInfo._MonitoringInfoLabels.ValueType # 16 GCS_PROJECT_ID: MonitoringInfo._MonitoringInfoLabels.ValueType # 17 DATASTORE_PROJECT: MonitoringInfo._MonitoringInfoLabels.ValueType # 18 DATASTORE_NAMESPACE: MonitoringInfo._MonitoringInfoLabels.ValueType # 19 BIGTABLE_PROJECT_ID: MonitoringInfo._MonitoringInfoLabels.ValueType # 20 INSTANCE_ID: MonitoringInfo._MonitoringInfoLabels.ValueType # 21 TABLE_ID: MonitoringInfo._MonitoringInfoLabels.ValueType # 22 SPANNER_PROJECT_ID: MonitoringInfo._MonitoringInfoLabels.ValueType # 23 SPANNER_DATABASE_ID: MonitoringInfo._MonitoringInfoLabels.ValueType # 24 SPANNER_TABLE_ID: MonitoringInfo._MonitoringInfoLabels.ValueType # 25 SPANNER_INSTANCE_ID: MonitoringInfo._MonitoringInfoLabels.ValueType # 26 SPANNER_QUERY_NAME: MonitoringInfo._MonitoringInfoLabels.ValueType # 27 PER_WORKER_METRIC: MonitoringInfo._MonitoringInfoLabels.ValueType # 28 """Label which if has a "true" value indicates that the metric is intended to be aggregated per-worker. """ class MonitoringInfoLabels(_MonitoringInfoLabels, metaclass=_MonitoringInfoLabelsEnumTypeWrapper): ... TRANSFORM: MonitoringInfo.MonitoringInfoLabels.ValueType # 0 """The values used for TRANSFORM, PCOLLECTION, WINDOWING_STRATEGY CODER, ENVIRONMENT, etc. must always match the keys used to refer to them. For actively processed bundles, these should match the values within the ProcessBundleDescriptor. For job management APIs, these should match values within the original pipeline representation. """ PCOLLECTION: MonitoringInfo.MonitoringInfoLabels.ValueType # 1 WINDOWING_STRATEGY: MonitoringInfo.MonitoringInfoLabels.ValueType # 2 CODER: MonitoringInfo.MonitoringInfoLabels.ValueType # 3 ENVIRONMENT: MonitoringInfo.MonitoringInfoLabels.ValueType # 4 NAMESPACE: MonitoringInfo.MonitoringInfoLabels.ValueType # 5 NAME: MonitoringInfo.MonitoringInfoLabels.ValueType # 6 SERVICE: MonitoringInfo.MonitoringInfoLabels.ValueType # 7 METHOD: MonitoringInfo.MonitoringInfoLabels.ValueType # 8 RESOURCE: MonitoringInfo.MonitoringInfoLabels.ValueType # 9 STATUS: MonitoringInfo.MonitoringInfoLabels.ValueType # 10 BIGQUERY_PROJECT_ID: MonitoringInfo.MonitoringInfoLabels.ValueType # 11 BIGQUERY_DATASET: MonitoringInfo.MonitoringInfoLabels.ValueType # 12 BIGQUERY_TABLE: MonitoringInfo.MonitoringInfoLabels.ValueType # 13 BIGQUERY_VIEW: MonitoringInfo.MonitoringInfoLabels.ValueType # 14 BIGQUERY_QUERY_NAME: MonitoringInfo.MonitoringInfoLabels.ValueType # 15 GCS_BUCKET: MonitoringInfo.MonitoringInfoLabels.ValueType # 16 GCS_PROJECT_ID: MonitoringInfo.MonitoringInfoLabels.ValueType # 17 DATASTORE_PROJECT: MonitoringInfo.MonitoringInfoLabels.ValueType # 18 DATASTORE_NAMESPACE: MonitoringInfo.MonitoringInfoLabels.ValueType # 19 BIGTABLE_PROJECT_ID: MonitoringInfo.MonitoringInfoLabels.ValueType # 20 INSTANCE_ID: MonitoringInfo.MonitoringInfoLabels.ValueType # 21 TABLE_ID: MonitoringInfo.MonitoringInfoLabels.ValueType # 22 SPANNER_PROJECT_ID: MonitoringInfo.MonitoringInfoLabels.ValueType # 23 SPANNER_DATABASE_ID: MonitoringInfo.MonitoringInfoLabels.ValueType # 24 SPANNER_TABLE_ID: MonitoringInfo.MonitoringInfoLabels.ValueType # 25 SPANNER_INSTANCE_ID: MonitoringInfo.MonitoringInfoLabels.ValueType # 26 SPANNER_QUERY_NAME: MonitoringInfo.MonitoringInfoLabels.ValueType # 27 PER_WORKER_METRIC: MonitoringInfo.MonitoringInfoLabels.ValueType # 28 """Label which if has a "true" value indicates that the metric is intended to be aggregated per-worker. """ @typing_extensions.final class LabelsEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str value: builtins.str def __init__( self, *, key: builtins.str | None = ..., value: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... URN_FIELD_NUMBER: builtins.int TYPE_FIELD_NUMBER: builtins.int PAYLOAD_FIELD_NUMBER: builtins.int LABELS_FIELD_NUMBER: builtins.int START_TIME_FIELD_NUMBER: builtins.int urn: builtins.str """(Required) Defines the semantic meaning of the metric or monitored state. See MonitoringInfoSpecs.Enum for the set of well known metrics/monitored state. """ type: builtins.str """(Required) Defines the encoding and aggregation method for the payload. See MonitoringInfoTypeUrns.Enum for the set of well known types. """ payload: builtins.bytes """(Required) The metric or monitored state encoded as per the specification defined by the type. """ @property def labels(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: """A set of key and value labels which define the scope of the metric. For well known URNs, the set of required labels is provided by the associated MonitoringInfoSpec. Either a well defined entity id for matching the enum names in the MonitoringInfoLabels enum or any arbitrary label set by a custom metric or user metric. A monitoring system is expected to be able to aggregate the metrics together for all updates having the same URN and labels. Some systems such as Stackdriver will be able to aggregate the metrics using a subset of the provided labels """ @property def start_time(self) -> google.protobuf.timestamp_pb2.Timestamp: """This indicates the start of the time range over which this value was measured. This is needed by some external metric aggregation services to indicate when the reporter of the metric first began collecting the cumulative value for the timeseries. If the SDK Harness restarts, it should reset the start_time, and reset the collection of cumulative metrics (i.e. start to count again from 0). HarnessMonitoringInfos should set this start_time once, when the MonitoringInfo is first reported. ProcessBundle MonitoringInfos should set a start_time for each bundle. """ def __init__( self, *, urn: builtins.str | None = ..., type: builtins.str | None = ..., payload: builtins.bytes | None = ..., labels: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., start_time: google.protobuf.timestamp_pb2.Timestamp | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["start_time", b"start_time"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["labels", b"labels", "payload", b"payload", "start_time", b"start_time", "type", b"type", "urn", b"urn"]) -> None: ... global___MonitoringInfo = MonitoringInfo @typing_extensions.final class MonitoringInfoTypeUrns(google.protobuf.message.Message): """A set of well known URNs that specify the encoding and aggregation method.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[MonitoringInfoTypeUrns._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor SUM_INT64_TYPE: MonitoringInfoTypeUrns._Enum.ValueType # 0 """Represents an integer counter where values are summed across bundles. Encoding: - value: beam:coder:varint:v1 """ SUM_DOUBLE_TYPE: MonitoringInfoTypeUrns._Enum.ValueType # 1 """Represents a double counter where values are summed across bundles. Encoding: value: beam:coder:double:v1 """ DISTRIBUTION_INT64_TYPE: MonitoringInfoTypeUrns._Enum.ValueType # 2 """Represents a distribution of an integer value where: - count: represents the number of values seen across all bundles - sum: represents the total of the value across all bundles - min: represents the smallest value seen across all bundles - max: represents the largest value seen across all bundles Encoding: - count: beam:coder:varint:v1 - sum: beam:coder:varint:v1 - min: beam:coder:varint:v1 - max: beam:coder:varint:v1 Note that when count is zero, the SDK may not send sum, min, and max in the response. If those fields are included in the payload, runners should omit them. """ DISTRIBUTION_DOUBLE_TYPE: MonitoringInfoTypeUrns._Enum.ValueType # 3 """Represents a distribution of a double value where: - count: represents the number of values seen across all bundles - sum: represents the total of the value across all bundles - min: represents the smallest value seen across all bundles - max: represents the largest value seen across all bundles Encoding: - count: beam:coder:varint:v1 - sum: beam:coder:double:v1 - min: beam:coder:double:v1 - max: beam:coder:double:v1 Note that when count is zero, the SDK may not send sum, min, and max in the response. If those fields are included in the payload, runners should omit them. """ LATEST_INT64_TYPE: MonitoringInfoTypeUrns._Enum.ValueType # 4 """Represents the latest seen integer value. The timestamp is used to provide an "ordering" over multiple values to determine which is the latest. Encoding: - timestamp: beam:coder:varint:v1 (milliseconds since epoch) - value: beam:coder:varint:v1 """ LATEST_DOUBLE_TYPE: MonitoringInfoTypeUrns._Enum.ValueType # 5 """Represents the latest seen double value. The timestamp is used to provide an "ordering" over multiple values to determine which is the latest. Encoding: - timestamp: beam:coder:varint:v1 (milliseconds since epoch) - value: beam:coder:double:v1 """ TOP_N_INT64_TYPE: MonitoringInfoTypeUrns._Enum.ValueType # 6 """Represents the largest set of integer values seen across bundles. Encoding: ... - iter: beam:coder:iterable:v1 - valueX: beam:coder:varint:v1 """ TOP_N_DOUBLE_TYPE: MonitoringInfoTypeUrns._Enum.ValueType # 7 """Represents the largest set of double values seen across bundles. Encoding: ... - iter: beam:coder:iterable:v1 - valueX: beam:coder... - iter: beam:coder:iterable:v1 - valueX: beam:coder:varint:v1 """ BOTTOM_N_DOUBLE_TYPE: MonitoringInfoTypeUrns._Enum.ValueType # 9 """Represents the smallest set of double values seen across bundles. Encoding: ... - iter: beam:coder:iterable:v1 - valueX: beam:coder:double:v1 """ PROGRESS_TYPE: MonitoringInfoTypeUrns._Enum.ValueType # 10 """Encoding: ... - iter: beam:coder:iterable:v1 - valueX: beam:coder:double:v1 """ SET_STRING_TYPE: MonitoringInfoTypeUrns._Enum.ValueType # 11 """Represents a set of strings. Encoding: ... - iter: beam:coder:iterable:v1 - valueX: beam:coder:stringutf8:v1 """ BOUNDED_TRIE_TYPE: MonitoringInfoTypeUrns._Enum.ValueType # 12 """Represents a bounded trie of strings. Encoding: BoundedTrie proto """ HISTOGRAM: MonitoringInfoTypeUrns._Enum.ValueType # 13 """Represents a histogram.""" class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... SUM_INT64_TYPE: MonitoringInfoTypeUrns.Enum.ValueType # 0 """Represents an integer counter where values are summed across bundles. Encoding: - value: beam:coder:varint:v1 """ SUM_DOUBLE_TYPE: MonitoringInfoTypeUrns.Enum.ValueType # 1 """Represents a double counter where values are summed across bundles. Encoding: value: beam:coder:double:v1 """ DISTRIBUTION_INT64_TYPE: MonitoringInfoTypeUrns.Enum.ValueType # 2 """Represents a distribution of an integer value where: - count: represents the number of values seen across all bundles - sum: represents the total of the value across all bundles - min: represents the smallest value seen across all bundles - max: represents the largest value seen across all bundles Encoding: - count: beam:coder:varint:v1 - sum: beam:coder:varint:v1 - min: beam:coder:varint:v1 - max: beam:coder:varint:v1 Note that when count is zero, the SDK may not send sum, min, and max in the response. If those fields are included in the payload, runners should omit them. """ DISTRIBUTION_DOUBLE_TYPE: MonitoringInfoTypeUrns.Enum.ValueType # 3 """Represents a distribution of a double value where: - count: represents the number of values seen across all bundles - sum: represents the total of the value across all bundles - min: represents the smallest value seen across all bundles - max: represents the largest value seen across all bundles Encoding: - count: beam:coder:varint:v1 - sum: beam:coder:double:v1 - min: beam:coder:double:v1 - max: beam:coder:double:v1 Note that when count is zero, the SDK may not send sum, min, and max in the response. If those fields are included in the payload, runners should omit them. """ LATEST_INT64_TYPE: MonitoringInfoTypeUrns.Enum.ValueType # 4 """Represents the latest seen integer value. The timestamp is used to provide an "ordering" over multiple values to determine which is the latest. Encoding: - timestamp: beam:coder:varint:v1 (milliseconds since epoch) - value: beam:coder:varint:v1 """ LATEST_DOUBLE_TYPE: MonitoringInfoTypeUrns.Enum.ValueType # 5 """Represents the latest seen double value. The timestamp is used to provide an "ordering" over multiple values to determine which is the latest. Encoding: - timestamp: beam:coder:varint:v1 (milliseconds since epoch) - value: beam:coder:double:v1 """ TOP_N_INT64_TYPE: MonitoringInfoTypeUrns.Enum.ValueType # 6 """Represents the largest set of integer values seen across bundles. Encoding: ... - iter: beam:coder:iterable:v1 - valueX: beam:coder:varint:v1 """ TOP_N_DOUBLE_TYPE: MonitoringInfoTypeUrns.Enum.ValueType # 7 """Represents the largest set of double values seen across bundles. Encoding: ... - iter: beam:coder:iterable:v1 - valueX: beam:coder... - iter: beam:coder:iterable:v1 - valueX: beam:coder:varint:v1 """ BOTTOM_N_DOUBLE_TYPE: MonitoringInfoTypeUrns.Enum.ValueType # 9 """Represents the smallest set of double values seen across bundles. Encoding: ... - iter: beam:coder:iterable:v1 - valueX: beam:coder:double:v1 """ PROGRESS_TYPE: MonitoringInfoTypeUrns.Enum.ValueType # 10 """Encoding: ... - iter: beam:coder:iterable:v1 - valueX: beam:coder:double:v1 """ SET_STRING_TYPE: MonitoringInfoTypeUrns.Enum.ValueType # 11 """Represents a set of strings. Encoding: ... - iter: beam:coder:iterable:v1 - valueX: beam:coder:stringutf8:v1 """ BOUNDED_TRIE_TYPE: MonitoringInfoTypeUrns.Enum.ValueType # 12 """Represents a bounded trie of strings. Encoding: BoundedTrie proto """ HISTOGRAM: MonitoringInfoTypeUrns.Enum.ValueType # 13 """Represents a histogram.""" def __init__( self, ) -> None: ... global___MonitoringInfoTypeUrns = MonitoringInfoTypeUrns @typing_extensions.final class BoundedTrieNode(google.protobuf.message.Message): """A single node in a BoundedTrie.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class ChildrenEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.str @property def value(self) -> global___BoundedTrieNode: ... def __init__( self, *, key: builtins.str | None = ..., value: global___BoundedTrieNode | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... TRUNCATED_FIELD_NUMBER: builtins.int CHILDREN_FIELD_NUMBER: builtins.int truncated: builtins.bool """Whether this node has been truncated. A truncated leaf represents possibly many children with the same prefix. """ @property def children(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___BoundedTrieNode]: """Children of this node. Must be empty if truncated is true.""" def __init__( self, *, truncated: builtins.bool | None = ..., children: collections.abc.Mapping[builtins.str, global___BoundedTrieNode] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["children", b"children", "truncated", b"truncated"]) -> None: ... global___BoundedTrieNode = BoundedTrieNode @typing_extensions.final class BoundedTrie(google.protobuf.message.Message): """The message type used for encoding metrics of type bounded trie.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor BOUND_FIELD_NUMBER: builtins.int ROOT_FIELD_NUMBER: builtins.int SINGLETON_FIELD_NUMBER: builtins.int bound: builtins.int """The maximum number of elements to store before truncation.""" @property def root(self) -> global___BoundedTrieNode: """A compact representation of all the elements in this trie.""" @property def singleton(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """A more efficient representation for metrics consisting of a single value.""" def __init__( self, *, bound: builtins.int | None = ..., root: global___BoundedTrieNode | None = ..., singleton: collections.abc.Iterable[builtins.str] | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["root", b"root"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["bound", b"bound", "root", b"root", "singleton", b"singleton"]) -> None: ... global___BoundedTrie = BoundedTrie @typing_extensions.final class HistogramValue(google.protobuf.message.Message): """The message type used for encoding Histogram Data""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class BucketOptions(google.protobuf.message.Message): """`BucketOptions` describes the bucket boundaries used in the histogram.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class Linear(google.protobuf.message.Message): """Linear buckets with the following boundaries for indices in 0 to n-1. - i in [0, n-1]: [start + (i)*width, start + (i+1)*width) """ DESCRIPTOR: google.protobuf.descriptor.Descriptor NUMBER_OF_BUCKETS_FIELD_NUMBER: builtins.int WIDTH_FIELD_NUMBER: builtins.int START_FIELD_NUMBER: builtins.int number_of_buckets: builtins.int """Must be greater than 0. (-- api-linter: core::0140::prepositions=disabled aip.dev/not-precedent: `bucket_count` would cause confusion with `bucket_counts` field --) """ width: builtins.float """Distance between bucket boundaries. Must be greater than 0.""" start: builtins.float """Lower bound of the first bucket.""" def __init__( self, *, number_of_buckets: builtins.int | None = ..., width: builtins.float | None = ..., start: builtins.float | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["_number_of_buckets", b"_number_of_buckets", "_start", b"_start", "_width", b"_width", "number_of_buckets", b"number_of_buckets", "start", b"start", "width", b"width"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["_number_of_buckets", b"_number_of_buckets", "_start", b"_start", "_width", b"_width", "number_of_buckets", b"number_of_buckets", "start", b"start", "width", b"width"]) -> None: ... @typing.overload def WhichOneof(self, oneof_group: typing_extensions.Literal["_number_of_buckets", b"_number_of_buckets"]) -> typing_extensions.Literal["number_of_buckets"] | None: ... @typing.overload def WhichOneof(self, oneof_group: typing_extensions.Literal["_start", b"_start"]) -> typing_extensions.Literal["start"] | None: ... @typing.overload def WhichOneof(self, oneof_group: typing_extensions.Literal["_width", b"_width"]) -> typing_extensions.Literal["width"] | None: ... @typing_extensions.final class Base2Exponent(google.protobuf.message.Message): """Exponential buckets where the growth factor between buckets is `2**(2**-scale)`. e.g. for `scale=1` growth factor is `2**(2**(-1))=sqrt(2)`. `n` buckets will have the following boundaries. - 0th: [0, gf) - i in [1, n-1]: [gf^(i), gf^(i+1)) """ DESCRIPTOR: google.protobuf.descriptor.Descriptor NUMBER_OF_BUCKETS_FIELD_NUMBER: builtins.int SCALE_FIELD_NUMBER: builtins.int number_of_buckets: builtins.int """Must be greater than 0. (-- api-linter: core::0140::prepositions=disabled aip.dev/not-precedent: `bucket_count` would cause confusion with `bucket_counts` field --) """ scale: builtins.int """Must be between -3 and 3. This forces the growth factor of the bucket boundaries to be between `2^(1/8)` and `256`. """ def __init__( self, *, number_of_buckets: builtins.int | None = ..., scale: builtins.int | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["_number_of_buckets", b"_number_of_buckets", "_scale", b"_scale", "number_of_buckets", b"number_of_buckets", "scale", b"scale"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["_number_of_buckets", b"_number_of_buckets", "_scale", b"_scale", "number_of_buckets", b"number_of_buckets", "scale", b"scale"]) -> None: ... @typing.overload def WhichOneof(self, oneof_group: typing_extensions.Literal["_number_of_buckets", b"_number_of_buckets"]) -> typing_extensions.Literal["number_of_buckets"] | None: ... @typing.overload def WhichOneof(self, oneof_group: typing_extensions.Literal["_scale", b"_scale"]) -> typing_extensions.Literal["scale"] | None: ... LINEAR_FIELD_NUMBER: builtins.int EXPONENTIAL_FIELD_NUMBER: builtins.int @property def linear(self) -> global___HistogramValue.BucketOptions.Linear: """Bucket boundaries grow linearly.""" @property def exponential(self) -> global___HistogramValue.BucketOptions.Base2Exponent: """Bucket boundaries grow exponentially.""" def __init__( self, *, linear: global___HistogramValue.BucketOptions.Linear | None = ..., exponential: global___HistogramValue.BucketOptions.Base2Exponent | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["bucket_type", b"bucket_type", "exponential", b"exponential", "linear", b"linear"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["bucket_type", b"bucket_type", "exponential", b"exponential", "linear", b"linear"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["bucket_type", b"bucket_type"]) -> typing_extensions.Literal["linear", "exponential"] | None: ... COUNT_FIELD_NUMBER: builtins.int BUCKET_OPTIONS_FIELD_NUMBER: builtins.int BUCKET_COUNTS_FIELD_NUMBER: builtins.int count: builtins.int """Number of values recorded in this histogram.""" @property def bucket_options(self) -> global___HistogramValue.BucketOptions: """Describes the bucket boundaries used in the histogram.""" @property def bucket_counts(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: """The number of values in each bucket of the histogram, as described in `bucket_options`. `bucket_counts` should contain N values, where N is the number of buckets specified in `bucket_options`. If `bucket_counts` has fewer than N values, the remaining values are assumed to be 0. """ def __init__( self, *, count: builtins.int | None = ..., bucket_options: global___HistogramValue.BucketOptions | None = ..., bucket_counts: collections.abc.Iterable[builtins.int] | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["_bucket_options", b"_bucket_options", "_count", b"_count", "bucket_options", b"bucket_options", "count", b"count"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["_bucket_options", b"_bucket_options", "_count", b"_count", "bucket_counts", b"bucket_counts", "bucket_options", b"bucket_options", "count", b"count"]) -> None: ... @typing.overload def WhichOneof(self, oneof_group: typing_extensions.Literal["_bucket_options", b"_bucket_options"]) -> typing_extensions.Literal["bucket_options"] | None: ... @typing.overload def WhichOneof(self, oneof_group: typing_extensions.Literal["_count", b"_count"]) -> typing_extensions.Literal["count"] | None: ... global___HistogramValue = HistogramValue LABEL_PROPS_FIELD_NUMBER: builtins.int MONITORING_INFO_SPEC_FIELD_NUMBER: builtins.int label_props: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.EnumValueOptions, global___MonitoringInfoLabelProps] """From: commit 0x7970544.""" monitoring_info_spec: google.protobuf.internal.extension_dict._ExtensionFieldDescriptor[google.protobuf.descriptor_pb2.EnumValueOptions, global___MonitoringInfoSpec] """Enum extension to store the MonitoringInfoSpecs.""" ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/metrics_pb2_grpc.py0000644000175100017510000000023715206113151032741 0ustar00runnerrunner# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/metrics_pb2_urns.py0000644000175100017510000003704315206113151033002 0ustar00runnerrunnerfrom ........utils import PropertiesFromEnumValue from . import metrics_pb2 EMPTY_MONITORING_INFO_LABEL_PROPS = metrics_pb2.MonitoringInfoLabelProps() EMPTY_MONITORING_INFO_SPEC = metrics_pb2.MonitoringInfoSpec() class MonitoringInfo(object): class MonitoringInfoLabels(object): TRANSFORM = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='PTRANSFORM')) PCOLLECTION = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='PCOLLECTION')) WINDOWING_STRATEGY = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='WINDOWING_STRATEGY')) CODER = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='CODER')) ENVIRONMENT = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='ENVIRONMENT')) NAMESPACE = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='NAMESPACE')) NAME = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='NAME')) SERVICE = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='SERVICE')) METHOD = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='METHOD')) RESOURCE = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='RESOURCE')) STATUS = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='STATUS')) BIGQUERY_PROJECT_ID = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='BIGQUERY_PROJECT_ID')) BIGQUERY_DATASET = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='BIGQUERY_DATASET')) BIGQUERY_TABLE = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='BIGQUERY_TABLE')) BIGQUERY_VIEW = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='BIGQUERY_VIEW')) BIGQUERY_QUERY_NAME = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='BIGQUERY_QUERY_NAME')) GCS_BUCKET = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='GCS_BUCKET')) GCS_PROJECT_ID = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='GCS_PROJECT_ID')) DATASTORE_PROJECT = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='DATASTORE_PROJECT')) DATASTORE_NAMESPACE = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='DATASTORE_NAMESPACE')) BIGTABLE_PROJECT_ID = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='BIGTABLE_PROJECT_ID')) INSTANCE_ID = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='INSTANCE_ID')) TABLE_ID = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='TABLE_ID')) SPANNER_PROJECT_ID = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='SPANNER_PROJECT_ID')) SPANNER_DATABASE_ID = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='SPANNER_DATABASE_ID')) SPANNER_TABLE_ID = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='SPANNER_TABLE_ID')) SPANNER_INSTANCE_ID = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='SPANNER_INSTANCE_ID')) SPANNER_QUERY_NAME = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='SPANNER_QUERY_NAME')) PER_WORKER_METRIC = PropertiesFromEnumValue('', '', EMPTY_MONITORING_INFO_SPEC, metrics_pb2.MonitoringInfoLabelProps(name='PER_WORKER_METRIC')) class MonitoringInfoSpecs(object): class Enum(object): USER_SUM_INT64 = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:user:sum_int64:v1', type='beam:metrics:sum_int64:v1', required_labels=['PTRANSFORM', 'NAMESPACE', 'NAME'], annotations=[metrics_pb2.Annotation(key='description', value='URN utilized to report user metric.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) USER_SUM_DOUBLE = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:user:sum_double:v1', type='beam:metrics:sum_double:v1', required_labels=['PTRANSFORM', 'NAMESPACE', 'NAME'], annotations=[metrics_pb2.Annotation(key='description', value='URN utilized to report user metric.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) USER_DISTRIBUTION_INT64 = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:user:distribution_int64:v1', type='beam:metrics:distribution_int64:v1', required_labels=['PTRANSFORM', 'NAMESPACE', 'NAME'], annotations=[metrics_pb2.Annotation(key='description', value='URN utilized to report user metric.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) USER_DISTRIBUTION_DOUBLE = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:user:distribution_double:v1', type='beam:metrics:distribution_double:v1', required_labels=['PTRANSFORM', 'NAMESPACE', 'NAME'], annotations=[metrics_pb2.Annotation(key='description', value='URN utilized to report user metric.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) USER_LATEST_INT64 = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:user:latest_int64:v1', type='beam:metrics:latest_int64:v1', required_labels=['PTRANSFORM', 'NAMESPACE', 'NAME'], annotations=[metrics_pb2.Annotation(key='description', value='URN utilized to report user metric.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) USER_LATEST_DOUBLE = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:user:latest_double:v1', type='beam:metrics:latest_double:v1', required_labels=['PTRANSFORM', 'NAMESPACE', 'NAME'], annotations=[metrics_pb2.Annotation(key='description', value='URN utilized to report user metric.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) USER_TOP_N_INT64 = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:user:top_n_int64:v1', type='beam:metrics:top_n_int64:v1', required_labels=['PTRANSFORM', 'NAMESPACE', 'NAME'], annotations=[metrics_pb2.Annotation(key='description', value='URN utilized to report user metric.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) USER_TOP_N_DOUBLE = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:user:top_n_double:v1', type='beam:metrics:top_n_double:v1', required_labels=['PTRANSFORM', 'NAMESPACE', 'NAME'], annotations=[metrics_pb2.Annotation(key='description', value='URN utilized to report user metric.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) USER_BOTTOM_N_INT64 = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:user:bottom_n_int64:v1', type='beam:metrics:bottom_n_int64:v1', required_labels=['PTRANSFORM', 'NAMESPACE', 'NAME'], annotations=[metrics_pb2.Annotation(key='description', value='URN utilized to report user metric.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) USER_BOTTOM_N_DOUBLE = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:user:bottom_n_double:v1', type='beam:metrics:bottom_n_double:v1', required_labels=['PTRANSFORM', 'NAMESPACE', 'NAME'], annotations=[metrics_pb2.Annotation(key='description', value='URN utilized to report user metric.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) USER_SET_STRING = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:user:set_string:v1', type='beam:metrics:set_string:v1', required_labels=['PTRANSFORM', 'NAMESPACE', 'NAME'], annotations=[metrics_pb2.Annotation(key='description', value='URN utilized to report user metric.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) USER_BOUNDED_TRIE = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:user:bounded_trie:v1', type='beam:metrics:bounded_trie:v1', required_labels=['PTRANSFORM', 'NAMESPACE', 'NAME'], annotations=[metrics_pb2.Annotation(key='description', value='URN utilized to report user metric.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) ELEMENT_COUNT = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:element_count:v1', type='beam:metrics:sum_int64:v1', required_labels=['PCOLLECTION'], annotations=[metrics_pb2.Annotation(key='description', value='The total elements output to a Pcollection by a PTransform.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) SAMPLED_BYTE_SIZE = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:sampled_byte_size:v1', type='beam:metrics:distribution_int64:v1', required_labels=['PCOLLECTION'], annotations=[metrics_pb2.Annotation(key='description', value='The total byte size and count of a sampled set (or all) of elements in the pcollection. Sampling is used because calculating the byte count involves serializing the elements which is CPU intensive.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) START_BUNDLE_MSECS = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:pardo_execution_time:start_bundle_msecs:v1', type='beam:metrics:sum_int64:v1', required_labels=['PTRANSFORM'], annotations=[metrics_pb2.Annotation(key='description', value='The total estimated execution time of the start bundlefunction in a pardo')]), EMPTY_MONITORING_INFO_LABEL_PROPS) PROCESS_BUNDLE_MSECS = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:pardo_execution_time:process_bundle_msecs:v1', type='beam:metrics:sum_int64:v1', required_labels=['PTRANSFORM'], annotations=[metrics_pb2.Annotation(key='description', value='The total estimated execution time of the process bundlefunction in a pardo')]), EMPTY_MONITORING_INFO_LABEL_PROPS) FINISH_BUNDLE_MSECS = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:pardo_execution_time:finish_bundle_msecs:v1', type='beam:metrics:sum_int64:v1', required_labels=['PTRANSFORM'], annotations=[metrics_pb2.Annotation(key='description', value='The total estimated execution time of the finish bundle function in a pardo')]), EMPTY_MONITORING_INFO_LABEL_PROPS) TOTAL_MSECS = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:ptransform_execution_time:total_msecs:v1', type='beam:metrics:sum_int64:v1', required_labels=['PTRANSFORM'], annotations=[metrics_pb2.Annotation(key='description', value='The total estimated execution time of the ptransform')]), EMPTY_MONITORING_INFO_LABEL_PROPS) WORK_REMAINING = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:ptransform_progress:remaining:v1', type='beam:metrics:progress:v1', required_labels=['PTRANSFORM'], annotations=[metrics_pb2.Annotation(key='description', value='The remaining amount of work for each active element. Each active element represents an independent amount of work not shared with any other active element.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) WORK_COMPLETED = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:ptransform_progress:completed:v1', type='beam:metrics:progress:v1', required_labels=['PTRANSFORM'], annotations=[metrics_pb2.Annotation(key='description', value='The remaining amount of work for each active element. Each active element represents an independent amount of work not shared with any other active element.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) DATA_CHANNEL_READ_INDEX = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:data_channel:read_index:v1', type='beam:metrics:sum_int64:v1', required_labels=['PTRANSFORM'], annotations=[metrics_pb2.Annotation(key='description', value='The read index of the data channel.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) API_REQUEST_COUNT = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:io:api_request_count:v1', type='beam:metrics:sum_int64:v1', required_labels=['SERVICE', 'METHOD', 'RESOURCE', 'PTRANSFORM', 'STATUS'], annotations=[metrics_pb2.Annotation(key='description', value='Request counts with status made to IO service APIs to batch read or write elements.'), metrics_pb2.Annotation(key='process_metric', value='true')]), EMPTY_MONITORING_INFO_LABEL_PROPS) API_REQUEST_LATENCIES = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:io:api_request_latencies:v1', type='beam:metrics:histogram_int64:v1', required_labels=['SERVICE', 'METHOD', 'RESOURCE', 'PTRANSFORM'], annotations=[metrics_pb2.Annotation(key='description', value='Histogram counts for request latencies made to IO service APIs to batch read or write elements.'), metrics_pb2.Annotation(key='units', value='Milliseconds'), metrics_pb2.Annotation(key='process_metric', value='true')]), EMPTY_MONITORING_INFO_LABEL_PROPS) USER_HISTOGRAM = PropertiesFromEnumValue('', '', metrics_pb2.MonitoringInfoSpec(urn='beam:metric:user:histogram_int64:v1', type='beam:metrics:histogram_int64:v1', required_labels=['PTRANSFORM', 'NAMESPACE', 'NAME'], annotations=[metrics_pb2.Annotation(key='description', value='URN utilized to report histogram metric.')]), EMPTY_MONITORING_INFO_LABEL_PROPS) class MonitoringInfoTypeUrns(object): class Enum(object): SUM_INT64_TYPE = PropertiesFromEnumValue('beam:metrics:sum_int64:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) SUM_DOUBLE_TYPE = PropertiesFromEnumValue('beam:metrics:sum_double:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) DISTRIBUTION_INT64_TYPE = PropertiesFromEnumValue('beam:metrics:distribution_int64:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) DISTRIBUTION_DOUBLE_TYPE = PropertiesFromEnumValue('beam:metrics:distribution_double:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) LATEST_INT64_TYPE = PropertiesFromEnumValue('beam:metrics:latest_int64:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) LATEST_DOUBLE_TYPE = PropertiesFromEnumValue('beam:metrics:latest_double:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) TOP_N_INT64_TYPE = PropertiesFromEnumValue('beam:metrics:top_n_int64:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) TOP_N_DOUBLE_TYPE = PropertiesFromEnumValue('beam:metrics:top_n_double:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) BOTTOM_N_INT64_TYPE = PropertiesFromEnumValue('beam:metrics:bottom_n_int64:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) BOTTOM_N_DOUBLE_TYPE = PropertiesFromEnumValue('beam:metrics:bottom_n_double:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) PROGRESS_TYPE = PropertiesFromEnumValue('beam:metrics:progress:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) SET_STRING_TYPE = PropertiesFromEnumValue('beam:metrics:set_string:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) BOUNDED_TRIE_TYPE = PropertiesFromEnumValue('beam:metrics:bounded_trie:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) HISTOGRAM = PropertiesFromEnumValue('beam:metrics:histogram_int64:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/schema_pb2.py0000644000175100017510000002666515206113151031535 0ustar00runnerrunner# -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # source: org/apache/beam/model/pipeline/v1/schema.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() from . import beam_runner_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_beam__runner__api__pb2 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n.org/apache/beam/model/pipeline/v1/schema.proto\x12!org.apache.beam.model.pipeline.v1\x1a\x37org/apache/beam/model/pipeline/v1/beam_runner_api.proto\"\xaa\x01\n\x06Schema\x12\x38\n\x06\x66ields\x18\x01 \x03(\x0b\x32(.org.apache.beam.model.pipeline.v1.Field\x12\n\n\x02id\x18\x02 \x01(\t\x12:\n\x07options\x18\x03 \x03(\x0b\x32).org.apache.beam.model.pipeline.v1.Option\x12\x1e\n\x16\x65ncoding_positions_set\x18\x04 \x01(\x08\"\xc9\x01\n\x05\x46ield\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x02 \x01(\t\x12:\n\x04type\x18\x03 \x01(\x0b\x32,.org.apache.beam.model.pipeline.v1.FieldType\x12\n\n\x02id\x18\x04 \x01(\x05\x12\x19\n\x11\x65ncoding_position\x18\x05 \x01(\x05\x12:\n\x07options\x18\x06 \x03(\x0b\x32).org.apache.beam.model.pipeline.v1.Option\"\xc6\x03\n\tFieldType\x12\x10\n\x08nullable\x18\x01 \x01(\x08\x12\x44\n\x0b\x61tomic_type\x18\x02 \x01(\x0e\x32-.org.apache.beam.model.pipeline.v1.AtomicTypeH\x00\x12\x42\n\narray_type\x18\x03 \x01(\x0b\x32,.org.apache.beam.model.pipeline.v1.ArrayTypeH\x00\x12H\n\riterable_type\x18\x04 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.IterableTypeH\x00\x12>\n\x08map_type\x18\x05 \x01(\x0b\x32*.org.apache.beam.model.pipeline.v1.MapTypeH\x00\x12>\n\x08row_type\x18\x06 \x01(\x0b\x32*.org.apache.beam.model.pipeline.v1.RowTypeH\x00\x12\x46\n\x0clogical_type\x18\x07 \x01(\x0b\x32..org.apache.beam.model.pipeline.v1.LogicalTypeH\x00\x42\x0b\n\ttype_info\"O\n\tArrayType\x12\x42\n\x0c\x65lement_type\x18\x01 \x01(\x0b\x32,.org.apache.beam.model.pipeline.v1.FieldType\"R\n\x0cIterableType\x12\x42\n\x0c\x65lement_type\x18\x01 \x01(\x0b\x32,.org.apache.beam.model.pipeline.v1.FieldType\"\x8b\x01\n\x07MapType\x12>\n\x08key_type\x18\x01 \x01(\x0b\x32,.org.apache.beam.model.pipeline.v1.FieldType\x12@\n\nvalue_type\x18\x02 \x01(\x0b\x32,.org.apache.beam.model.pipeline.v1.FieldType\"D\n\x07RowType\x12\x39\n\x06schema\x18\x01 \x01(\x0b\x32).org.apache.beam.model.pipeline.v1.Schema\"\xf7\x01\n\x0bLogicalType\x12\x0b\n\x03urn\x18\x01 \x01(\t\x12\x0f\n\x07payload\x18\x02 \x01(\x0c\x12\x44\n\x0erepresentation\x18\x03 \x01(\x0b\x32,.org.apache.beam.model.pipeline.v1.FieldType\x12\x43\n\rargument_type\x18\x04 \x01(\x0b\x32,.org.apache.beam.model.pipeline.v1.FieldType\x12?\n\x08\x61rgument\x18\x05 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.FieldValue\"\x8a\x04\n\x0cLogicalTypes\"\xf9\x03\n\x04\x45num\x12?\n\x0fPYTHON_CALLABLE\x10\x00\x1a*\xa2\xb4\xfa\xc2\x05$beam:logical_type:python_callable:v1\x12=\n\x0eMICROS_INSTANT\x10\x01\x1a)\xa2\xb4\xfa\xc2\x05#beam:logical_type:micros_instant:v1\x12=\n\x0eMILLIS_INSTANT\x10\x02\x1a)\xa2\xb4\xfa\xc2\x05#beam:logical_type:millis_instant:v1\x12/\n\x07\x44\x45\x43IMAL\x10\x03\x1a\"\xa2\xb4\xfa\xc2\x05\x1c\x62\x65\x61m:logical_type:decimal:v1\x12\x37\n\x0b\x46IXED_BYTES\x10\x04\x1a&\xa2\xb4\xfa\xc2\x05 beam:logical_type:fixed_bytes:v1\x12\x33\n\tVAR_BYTES\x10\x05\x1a$\xa2\xb4\xfa\xc2\x05\x1e\x62\x65\x61m:logical_type:var_bytes:v1\x12\x35\n\nFIXED_CHAR\x10\x06\x1a%\xa2\xb4\xfa\xc2\x05\x1f\x62\x65\x61m:logical_type:fixed_char:v1\x12\x31\n\x08VAR_CHAR\x10\x07\x1a#\xa2\xb4\xfa\xc2\x05\x1d\x62\x65\x61m:logical_type:var_char:v1\x12)\n\x04\x44\x41TE\x10\x08\x1a\x1f\xa2\xb4\xfa\xc2\x05\x19\x62\x65\x61m:logical_type:date:v1\"\x90\x01\n\x06Option\x12\x0c\n\x04name\x18\x01 \x01(\t\x12:\n\x04type\x18\x02 \x01(\x0b\x32,.org.apache.beam.model.pipeline.v1.FieldType\x12<\n\x05value\x18\x03 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.FieldValue\"D\n\x03Row\x12=\n\x06values\x18\x01 \x03(\x0b\x32-.org.apache.beam.model.pipeline.v1.FieldValue\"\xd7\x03\n\nFieldValue\x12J\n\x0c\x61tomic_value\x18\x01 \x01(\x0b\x32\x32.org.apache.beam.model.pipeline.v1.AtomicTypeValueH\x00\x12H\n\x0b\x61rray_value\x18\x02 \x01(\x0b\x32\x31.org.apache.beam.model.pipeline.v1.ArrayTypeValueH\x00\x12N\n\x0eiterable_value\x18\x03 \x01(\x0b\x32\x34.org.apache.beam.model.pipeline.v1.IterableTypeValueH\x00\x12\x44\n\tmap_value\x18\x04 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.MapTypeValueH\x00\x12;\n\trow_value\x18\x05 \x01(\x0b\x32&.org.apache.beam.model.pipeline.v1.RowH\x00\x12Q\n\x12logical_type_value\x18\x06 \x01(\x0b\x32\x33.org.apache.beam.model.pipeline.v1.LogicalTypeValueH\x00\x42\r\n\x0b\x66ield_value\"\xb6\x01\n\x0f\x41tomicTypeValue\x12\x0e\n\x04\x62yte\x18\x01 \x01(\x05H\x00\x12\x0f\n\x05int16\x18\x02 \x01(\x05H\x00\x12\x0f\n\x05int32\x18\x03 \x01(\x05H\x00\x12\x0f\n\x05int64\x18\x04 \x01(\x03H\x00\x12\x0f\n\x05\x66loat\x18\x05 \x01(\x02H\x00\x12\x10\n\x06\x64ouble\x18\x06 \x01(\x01H\x00\x12\x10\n\x06string\x18\x07 \x01(\tH\x00\x12\x11\n\x07\x62oolean\x18\x08 \x01(\x08H\x00\x12\x0f\n\x05\x62ytes\x18\t \x01(\x0cH\x00\x42\x07\n\x05value\"P\n\x0e\x41rrayTypeValue\x12>\n\x07\x65lement\x18\x01 \x03(\x0b\x32-.org.apache.beam.model.pipeline.v1.FieldValue\"S\n\x11IterableTypeValue\x12>\n\x07\x65lement\x18\x01 \x03(\x0b\x32-.org.apache.beam.model.pipeline.v1.FieldValue\"P\n\x0cMapTypeValue\x12@\n\x07\x65ntries\x18\x01 \x03(\x0b\x32/.org.apache.beam.model.pipeline.v1.MapTypeEntry\"\x88\x01\n\x0cMapTypeEntry\x12:\n\x03key\x18\x01 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.FieldValue\x12<\n\x05value\x18\x02 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.FieldValue\"P\n\x10LogicalTypeValue\x12<\n\x05value\x18\x01 \x01(\x0b\x32-.org.apache.beam.model.pipeline.v1.FieldValue\"\xf9\x02\n\x12SchemaCoderPayload\x12\x39\n\x06schema\x18\x01 \x01(\x0b\x32).org.apache.beam.model.pipeline.v1.Schema\x12\x42\n\tto_row_fn\x18\x02 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\x12\x44\n\x0b\x66rom_row_fn\x18\x03 \x01(\x0b\x32/.org.apache.beam.model.pipeline.v1.FunctionSpec\x12i\n\x16\x61\x64\x64itional_coder_infos\x18\x04 \x03(\x0b\x32I.org.apache.beam.model.pipeline.v1.SchemaCoderPayload.AdditionalCoderInfo\x1a\x33\n\x13\x41\x64\x64itionalCoderInfo\x12\x0b\n\x03urn\x18\x01 \x01(\t\x12\x0f\n\x07payload\x18\x02 \x01(\x0c*\x83\x01\n\nAtomicType\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x08\n\x04\x42YTE\x10\x01\x12\t\n\x05INT16\x10\x02\x12\t\n\x05INT32\x10\x03\x12\t\n\x05INT64\x10\x04\x12\t\n\x05\x46LOAT\x10\x05\x12\n\n\x06\x44OUBLE\x10\x06\x12\n\n\x06STRING\x10\x07\x12\x0b\n\x07\x42OOLEAN\x10\x08\x12\t\n\x05\x42YTES\x10\tBx\n!org.apache.beam.model.pipeline.v1B\tSchemaApiZHgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1;pipeline_v1b\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'org.apache.beam.model.pipeline.v1.schema_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n!org.apache.beam.model.pipeline.v1B\tSchemaApiZHgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1;pipeline_v1' _globals['_LOGICALTYPES_ENUM'].values_by_name["PYTHON_CALLABLE"]._options = None _globals['_LOGICALTYPES_ENUM'].values_by_name["PYTHON_CALLABLE"]._serialized_options = b'\242\264\372\302\005$beam:logical_type:python_callable:v1' _globals['_LOGICALTYPES_ENUM'].values_by_name["MICROS_INSTANT"]._options = None _globals['_LOGICALTYPES_ENUM'].values_by_name["MICROS_INSTANT"]._serialized_options = b'\242\264\372\302\005#beam:logical_type:micros_instant:v1' _globals['_LOGICALTYPES_ENUM'].values_by_name["MILLIS_INSTANT"]._options = None _globals['_LOGICALTYPES_ENUM'].values_by_name["MILLIS_INSTANT"]._serialized_options = b'\242\264\372\302\005#beam:logical_type:millis_instant:v1' _globals['_LOGICALTYPES_ENUM'].values_by_name["DECIMAL"]._options = None _globals['_LOGICALTYPES_ENUM'].values_by_name["DECIMAL"]._serialized_options = b'\242\264\372\302\005\034beam:logical_type:decimal:v1' _globals['_LOGICALTYPES_ENUM'].values_by_name["FIXED_BYTES"]._options = None _globals['_LOGICALTYPES_ENUM'].values_by_name["FIXED_BYTES"]._serialized_options = b'\242\264\372\302\005 beam:logical_type:fixed_bytes:v1' _globals['_LOGICALTYPES_ENUM'].values_by_name["VAR_BYTES"]._options = None _globals['_LOGICALTYPES_ENUM'].values_by_name["VAR_BYTES"]._serialized_options = b'\242\264\372\302\005\036beam:logical_type:var_bytes:v1' _globals['_LOGICALTYPES_ENUM'].values_by_name["FIXED_CHAR"]._options = None _globals['_LOGICALTYPES_ENUM'].values_by_name["FIXED_CHAR"]._serialized_options = b'\242\264\372\302\005\037beam:logical_type:fixed_char:v1' _globals['_LOGICALTYPES_ENUM'].values_by_name["VAR_CHAR"]._options = None _globals['_LOGICALTYPES_ENUM'].values_by_name["VAR_CHAR"]._serialized_options = b'\242\264\372\302\005\035beam:logical_type:var_char:v1' _globals['_LOGICALTYPES_ENUM'].values_by_name["DATE"]._options = None _globals['_LOGICALTYPES_ENUM'].values_by_name["DATE"]._serialized_options = b'\242\264\372\302\005\031beam:logical_type:date:v1' _globals['_ATOMICTYPE']._serialized_start=3855 _globals['_ATOMICTYPE']._serialized_end=3986 _globals['_SCHEMA']._serialized_start=143 _globals['_SCHEMA']._serialized_end=313 _globals['_FIELD']._serialized_start=316 _globals['_FIELD']._serialized_end=517 _globals['_FIELDTYPE']._serialized_start=520 _globals['_FIELDTYPE']._serialized_end=974 _globals['_ARRAYTYPE']._serialized_start=976 _globals['_ARRAYTYPE']._serialized_end=1055 _globals['_ITERABLETYPE']._serialized_start=1057 _globals['_ITERABLETYPE']._serialized_end=1139 _globals['_MAPTYPE']._serialized_start=1142 _globals['_MAPTYPE']._serialized_end=1281 _globals['_ROWTYPE']._serialized_start=1283 _globals['_ROWTYPE']._serialized_end=1351 _globals['_LOGICALTYPE']._serialized_start=1354 _globals['_LOGICALTYPE']._serialized_end=1601 _globals['_LOGICALTYPES']._serialized_start=1604 _globals['_LOGICALTYPES']._serialized_end=2126 _globals['_LOGICALTYPES_ENUM']._serialized_start=1621 _globals['_LOGICALTYPES_ENUM']._serialized_end=2126 _globals['_OPTION']._serialized_start=2129 _globals['_OPTION']._serialized_end=2273 _globals['_ROW']._serialized_start=2275 _globals['_ROW']._serialized_end=2343 _globals['_FIELDVALUE']._serialized_start=2346 _globals['_FIELDVALUE']._serialized_end=2817 _globals['_ATOMICTYPEVALUE']._serialized_start=2820 _globals['_ATOMICTYPEVALUE']._serialized_end=3002 _globals['_ARRAYTYPEVALUE']._serialized_start=3004 _globals['_ARRAYTYPEVALUE']._serialized_end=3084 _globals['_ITERABLETYPEVALUE']._serialized_start=3086 _globals['_ITERABLETYPEVALUE']._serialized_end=3169 _globals['_MAPTYPEVALUE']._serialized_start=3171 _globals['_MAPTYPEVALUE']._serialized_end=3251 _globals['_MAPTYPEENTRY']._serialized_start=3254 _globals['_MAPTYPEENTRY']._serialized_end=3390 _globals['_LOGICALTYPEVALUE']._serialized_start=3392 _globals['_LOGICALTYPEVALUE']._serialized_end=3472 _globals['_SCHEMACODERPAYLOAD']._serialized_start=3475 _globals['_SCHEMACODERPAYLOAD']._serialized_end=3852 _globals['_SCHEMACODERPAYLOAD_ADDITIONALCODERINFO']._serialized_start=3801 _globals['_SCHEMACODERPAYLOAD_ADDITIONALCODERINFO']._serialized_end=3852 # @@protoc_insertion_point(module_scope) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/schema_pb2.pyi0000644000175100017510000006763715206113151031712 0ustar00runnerrunner""" @generated by mypy-protobuf. Do not edit manually! isort:skip_file Protocol Buffers describing Beam Schemas, a portable representation for complex types. The primary application of Schema is as the payload for the standard coder "beam:coder:row:v1", defined in beam_runner_api.proto """ import builtins import collections.abc import google.protobuf.descriptor import google.protobuf.internal.containers import google.protobuf.internal.enum_type_wrapper import google.protobuf.message import org.apache.beam.model.pipeline.v1.beam_runner_api_pb2 import sys import typing if sys.version_info >= (3, 10): import typing as typing_extensions else: import typing_extensions DESCRIPTOR: google.protobuf.descriptor.FileDescriptor class _AtomicType: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _AtomicTypeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_AtomicType.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor UNSPECIFIED: _AtomicType.ValueType # 0 BYTE: _AtomicType.ValueType # 1 INT16: _AtomicType.ValueType # 2 INT32: _AtomicType.ValueType # 3 INT64: _AtomicType.ValueType # 4 FLOAT: _AtomicType.ValueType # 5 DOUBLE: _AtomicType.ValueType # 6 STRING: _AtomicType.ValueType # 7 BOOLEAN: _AtomicType.ValueType # 8 BYTES: _AtomicType.ValueType # 9 class AtomicType(_AtomicType, metaclass=_AtomicTypeEnumTypeWrapper): ... UNSPECIFIED: AtomicType.ValueType # 0 BYTE: AtomicType.ValueType # 1 INT16: AtomicType.ValueType # 2 INT32: AtomicType.ValueType # 3 INT64: AtomicType.ValueType # 4 FLOAT: AtomicType.ValueType # 5 DOUBLE: AtomicType.ValueType # 6 STRING: AtomicType.ValueType # 7 BOOLEAN: AtomicType.ValueType # 8 BYTES: AtomicType.ValueType # 9 global___AtomicType = AtomicType @typing_extensions.final class Schema(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor FIELDS_FIELD_NUMBER: builtins.int ID_FIELD_NUMBER: builtins.int OPTIONS_FIELD_NUMBER: builtins.int ENCODING_POSITIONS_SET_FIELD_NUMBER: builtins.int @property def fields(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Field]: """List of fields for this schema. Two fields may not share a name.""" id: builtins.str """REQUIRED. An RFC 4122 UUID.""" @property def options(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Option]: ... encoding_positions_set: builtins.bool """Indicates that encoding positions have been overridden.""" def __init__( self, *, fields: collections.abc.Iterable[global___Field] | None = ..., id: builtins.str | None = ..., options: collections.abc.Iterable[global___Option] | None = ..., encoding_positions_set: builtins.bool | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["encoding_positions_set", b"encoding_positions_set", "fields", b"fields", "id", b"id", "options", b"options"]) -> None: ... global___Schema = Schema @typing_extensions.final class Field(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor NAME_FIELD_NUMBER: builtins.int DESCRIPTION_FIELD_NUMBER: builtins.int TYPE_FIELD_NUMBER: builtins.int ID_FIELD_NUMBER: builtins.int ENCODING_POSITION_FIELD_NUMBER: builtins.int OPTIONS_FIELD_NUMBER: builtins.int name: builtins.str """REQUIRED. Name of this field within the schema.""" description: builtins.str """OPTIONAL. Human readable description of this field, such as the query that generated it.""" @property def type(self) -> global___FieldType: ... id: builtins.int encoding_position: builtins.int """OPTIONAL. The position of this field's data when encoded, e.g. with beam:coder:row:v1. Either no fields in a given row are have encoding position populated, or all of them are. Used to support backwards compatibility with schema changes. If no fields have encoding position populated the order of encoding is the same as the order in the Schema. If this Field is part of a Schema where encoding_positions_set is True then encoding_position must be defined, otherwise this field is ignored. """ @property def options(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Option]: ... def __init__( self, *, name: builtins.str | None = ..., description: builtins.str | None = ..., type: global___FieldType | None = ..., id: builtins.int | None = ..., encoding_position: builtins.int | None = ..., options: collections.abc.Iterable[global___Option] | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["type", b"type"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["description", b"description", "encoding_position", b"encoding_position", "id", b"id", "name", b"name", "options", b"options", "type", b"type"]) -> None: ... global___Field = Field @typing_extensions.final class FieldType(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor NULLABLE_FIELD_NUMBER: builtins.int ATOMIC_TYPE_FIELD_NUMBER: builtins.int ARRAY_TYPE_FIELD_NUMBER: builtins.int ITERABLE_TYPE_FIELD_NUMBER: builtins.int MAP_TYPE_FIELD_NUMBER: builtins.int ROW_TYPE_FIELD_NUMBER: builtins.int LOGICAL_TYPE_FIELD_NUMBER: builtins.int nullable: builtins.bool atomic_type: global___AtomicType.ValueType @property def array_type(self) -> global___ArrayType: ... @property def iterable_type(self) -> global___IterableType: ... @property def map_type(self) -> global___MapType: ... @property def row_type(self) -> global___RowType: ... @property def logical_type(self) -> global___LogicalType: ... def __init__( self, *, nullable: builtins.bool | None = ..., atomic_type: global___AtomicType.ValueType | None = ..., array_type: global___ArrayType | None = ..., iterable_type: global___IterableType | None = ..., map_type: global___MapType | None = ..., row_type: global___RowType | None = ..., logical_type: global___LogicalType | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["array_type", b"array_type", "atomic_type", b"atomic_type", "iterable_type", b"iterable_type", "logical_type", b"logical_type", "map_type", b"map_type", "row_type", b"row_type", "type_info", b"type_info"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["array_type", b"array_type", "atomic_type", b"atomic_type", "iterable_type", b"iterable_type", "logical_type", b"logical_type", "map_type", b"map_type", "nullable", b"nullable", "row_type", b"row_type", "type_info", b"type_info"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["type_info", b"type_info"]) -> typing_extensions.Literal["atomic_type", "array_type", "iterable_type", "map_type", "row_type", "logical_type"] | None: ... global___FieldType = FieldType @typing_extensions.final class ArrayType(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ELEMENT_TYPE_FIELD_NUMBER: builtins.int @property def element_type(self) -> global___FieldType: ... def __init__( self, *, element_type: global___FieldType | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["element_type", b"element_type"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["element_type", b"element_type"]) -> None: ... global___ArrayType = ArrayType @typing_extensions.final class IterableType(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ELEMENT_TYPE_FIELD_NUMBER: builtins.int @property def element_type(self) -> global___FieldType: ... def __init__( self, *, element_type: global___FieldType | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["element_type", b"element_type"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["element_type", b"element_type"]) -> None: ... global___IterableType = IterableType @typing_extensions.final class MapType(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_TYPE_FIELD_NUMBER: builtins.int VALUE_TYPE_FIELD_NUMBER: builtins.int @property def key_type(self) -> global___FieldType: ... @property def value_type(self) -> global___FieldType: ... def __init__( self, *, key_type: global___FieldType | None = ..., value_type: global___FieldType | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["key_type", b"key_type", "value_type", b"value_type"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key_type", b"key_type", "value_type", b"value_type"]) -> None: ... global___MapType = MapType @typing_extensions.final class RowType(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor SCHEMA_FIELD_NUMBER: builtins.int @property def schema(self) -> global___Schema: ... def __init__( self, *, schema: global___Schema | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["schema", b"schema"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["schema", b"schema"]) -> None: ... global___RowType = RowType @typing_extensions.final class LogicalType(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor URN_FIELD_NUMBER: builtins.int PAYLOAD_FIELD_NUMBER: builtins.int REPRESENTATION_FIELD_NUMBER: builtins.int ARGUMENT_TYPE_FIELD_NUMBER: builtins.int ARGUMENT_FIELD_NUMBER: builtins.int urn: builtins.str payload: builtins.bytes @property def representation(self) -> global___FieldType: ... @property def argument_type(self) -> global___FieldType: ... @property def argument(self) -> global___FieldValue: ... def __init__( self, *, urn: builtins.str | None = ..., payload: builtins.bytes | None = ..., representation: global___FieldType | None = ..., argument_type: global___FieldType | None = ..., argument: global___FieldValue | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["argument", b"argument", "argument_type", b"argument_type", "representation", b"representation"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["argument", b"argument", "argument_type", b"argument_type", "payload", b"payload", "representation", b"representation", "urn", b"urn"]) -> None: ... global___LogicalType = LogicalType @typing_extensions.final class LogicalTypes(google.protobuf.message.Message): """Universally defined Logical types for Row schemas. These logical types are supposed to be understood by all SDKs. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[LogicalTypes._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor PYTHON_CALLABLE: LogicalTypes._Enum.ValueType # 0 """A URN for Python Callable logical type - Representation type: STRING - Language type: In Python SDK, PythonCallableWithSource. In any other SDKs, a wrapper object for a string which can be evaluated to a Python Callable object. """ MICROS_INSTANT: LogicalTypes._Enum.ValueType # 1 """A URN for MicrosInstant type - Representation type: ROW - A timestamp without a timezone where seconds + micros represents the amount of time since the epoch. """ MILLIS_INSTANT: LogicalTypes._Enum.ValueType # 2 """A URN for MillisInstant type - Representation type: INT64 - A timestamp without a timezone represented by the number of milliseconds since the epoch. The INT64 value is encoded with big-endian shifted such that lexicographic ordering of the bytes corresponds to chronological order. """ DECIMAL: LogicalTypes._Enum.ValueType # 3 """A URN for Decimal type - Representation type: BYTES - A decimal number with variable scale. Its BYTES representation consists of an integer (INT32) scale followed by a two's complement encoded big integer. """ FIXED_BYTES: LogicalTypes._Enum.ValueType # 4 """A URN for FixedLengthBytes type - Representation type: BYTES - Argument type: INT32. A fixed-length bytes with its length as the argument. """ VAR_BYTES: LogicalTypes._Enum.ValueType # 5 """A URN for VariableLengthBytes type - Representation type: BYTES - Argument type: INT32. A variable-length bytes with its maximum length as the argument. """ FIXED_CHAR: LogicalTypes._Enum.ValueType # 6 """A URN for FixedLengthString type - Representation type: STRING - Argument type: INT32. A fixed-length string with its length as the argument. """ VAR_CHAR: LogicalTypes._Enum.ValueType # 7 """A URN for VariableLengthString type - Representation type: STRING - Argument type: INT32. A variable-length string with its maximum length as the argument. """ DATE: LogicalTypes._Enum.ValueType # 8 """A URN for Date type - Representation type: INT64 - A date without a timezone, represented by the number of days since the epoch. """ class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... PYTHON_CALLABLE: LogicalTypes.Enum.ValueType # 0 """A URN for Python Callable logical type - Representation type: STRING - Language type: In Python SDK, PythonCallableWithSource. In any other SDKs, a wrapper object for a string which can be evaluated to a Python Callable object. """ MICROS_INSTANT: LogicalTypes.Enum.ValueType # 1 """A URN for MicrosInstant type - Representation type: ROW - A timestamp without a timezone where seconds + micros represents the amount of time since the epoch. """ MILLIS_INSTANT: LogicalTypes.Enum.ValueType # 2 """A URN for MillisInstant type - Representation type: INT64 - A timestamp without a timezone represented by the number of milliseconds since the epoch. The INT64 value is encoded with big-endian shifted such that lexicographic ordering of the bytes corresponds to chronological order. """ DECIMAL: LogicalTypes.Enum.ValueType # 3 """A URN for Decimal type - Representation type: BYTES - A decimal number with variable scale. Its BYTES representation consists of an integer (INT32) scale followed by a two's complement encoded big integer. """ FIXED_BYTES: LogicalTypes.Enum.ValueType # 4 """A URN for FixedLengthBytes type - Representation type: BYTES - Argument type: INT32. A fixed-length bytes with its length as the argument. """ VAR_BYTES: LogicalTypes.Enum.ValueType # 5 """A URN for VariableLengthBytes type - Representation type: BYTES - Argument type: INT32. A variable-length bytes with its maximum length as the argument. """ FIXED_CHAR: LogicalTypes.Enum.ValueType # 6 """A URN for FixedLengthString type - Representation type: STRING - Argument type: INT32. A fixed-length string with its length as the argument. """ VAR_CHAR: LogicalTypes.Enum.ValueType # 7 """A URN for VariableLengthString type - Representation type: STRING - Argument type: INT32. A variable-length string with its maximum length as the argument. """ DATE: LogicalTypes.Enum.ValueType # 8 """A URN for Date type - Representation type: INT64 - A date without a timezone, represented by the number of days since the epoch. """ def __init__( self, ) -> None: ... global___LogicalTypes = LogicalTypes @typing_extensions.final class Option(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor NAME_FIELD_NUMBER: builtins.int TYPE_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int name: builtins.str """REQUIRED. Identifier for the option.""" @property def type(self) -> global___FieldType: """REQUIRED. Type specifier for the structure of value. Conventionally, options that don't require additional configuration should use a boolean type, with the value set to true. """ @property def value(self) -> global___FieldValue: ... def __init__( self, *, name: builtins.str | None = ..., type: global___FieldType | None = ..., value: global___FieldValue | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["type", b"type", "value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["name", b"name", "type", b"type", "value", b"value"]) -> None: ... global___Option = Option @typing_extensions.final class Row(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor VALUES_FIELD_NUMBER: builtins.int @property def values(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___FieldValue]: ... def __init__( self, *, values: collections.abc.Iterable[global___FieldValue] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["values", b"values"]) -> None: ... global___Row = Row @typing_extensions.final class FieldValue(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ATOMIC_VALUE_FIELD_NUMBER: builtins.int ARRAY_VALUE_FIELD_NUMBER: builtins.int ITERABLE_VALUE_FIELD_NUMBER: builtins.int MAP_VALUE_FIELD_NUMBER: builtins.int ROW_VALUE_FIELD_NUMBER: builtins.int LOGICAL_TYPE_VALUE_FIELD_NUMBER: builtins.int @property def atomic_value(self) -> global___AtomicTypeValue: ... @property def array_value(self) -> global___ArrayTypeValue: ... @property def iterable_value(self) -> global___IterableTypeValue: ... @property def map_value(self) -> global___MapTypeValue: ... @property def row_value(self) -> global___Row: ... @property def logical_type_value(self) -> global___LogicalTypeValue: ... def __init__( self, *, atomic_value: global___AtomicTypeValue | None = ..., array_value: global___ArrayTypeValue | None = ..., iterable_value: global___IterableTypeValue | None = ..., map_value: global___MapTypeValue | None = ..., row_value: global___Row | None = ..., logical_type_value: global___LogicalTypeValue | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["array_value", b"array_value", "atomic_value", b"atomic_value", "field_value", b"field_value", "iterable_value", b"iterable_value", "logical_type_value", b"logical_type_value", "map_value", b"map_value", "row_value", b"row_value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["array_value", b"array_value", "atomic_value", b"atomic_value", "field_value", b"field_value", "iterable_value", b"iterable_value", "logical_type_value", b"logical_type_value", "map_value", b"map_value", "row_value", b"row_value"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["field_value", b"field_value"]) -> typing_extensions.Literal["atomic_value", "array_value", "iterable_value", "map_value", "row_value", "logical_type_value"] | None: ... global___FieldValue = FieldValue @typing_extensions.final class AtomicTypeValue(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor BYTE_FIELD_NUMBER: builtins.int INT16_FIELD_NUMBER: builtins.int INT32_FIELD_NUMBER: builtins.int INT64_FIELD_NUMBER: builtins.int FLOAT_FIELD_NUMBER: builtins.int DOUBLE_FIELD_NUMBER: builtins.int STRING_FIELD_NUMBER: builtins.int BOOLEAN_FIELD_NUMBER: builtins.int BYTES_FIELD_NUMBER: builtins.int byte: builtins.int int16: builtins.int int32: builtins.int int64: builtins.int float: builtins.float double: builtins.float string: builtins.str boolean: builtins.bool bytes: builtins.bytes def __init__( self, *, byte: builtins.int | None = ..., int16: builtins.int | None = ..., int32: builtins.int | None = ..., int64: builtins.int | None = ..., float: builtins.float | None = ..., double: builtins.float | None = ..., string: builtins.str | None = ..., boolean: builtins.bool | None = ..., bytes: builtins.bytes | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["boolean", b"boolean", "byte", b"byte", "bytes", b"bytes", "double", b"double", "float", b"float", "int16", b"int16", "int32", b"int32", "int64", b"int64", "string", b"string", "value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["boolean", b"boolean", "byte", b"byte", "bytes", b"bytes", "double", b"double", "float", b"float", "int16", b"int16", "int32", b"int32", "int64", b"int64", "string", b"string", "value", b"value"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["value", b"value"]) -> typing_extensions.Literal["byte", "int16", "int32", "int64", "float", "double", "string", "boolean", "bytes"] | None: ... global___AtomicTypeValue = AtomicTypeValue @typing_extensions.final class ArrayTypeValue(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ELEMENT_FIELD_NUMBER: builtins.int @property def element(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___FieldValue]: ... def __init__( self, *, element: collections.abc.Iterable[global___FieldValue] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["element", b"element"]) -> None: ... global___ArrayTypeValue = ArrayTypeValue @typing_extensions.final class IterableTypeValue(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ELEMENT_FIELD_NUMBER: builtins.int @property def element(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___FieldValue]: ... def __init__( self, *, element: collections.abc.Iterable[global___FieldValue] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["element", b"element"]) -> None: ... global___IterableTypeValue = IterableTypeValue @typing_extensions.final class MapTypeValue(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor ENTRIES_FIELD_NUMBER: builtins.int @property def entries(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___MapTypeEntry]: ... def __init__( self, *, entries: collections.abc.Iterable[global___MapTypeEntry] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["entries", b"entries"]) -> None: ... global___MapTypeValue = MapTypeValue @typing_extensions.final class MapTypeEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int @property def key(self) -> global___FieldValue: ... @property def value(self) -> global___FieldValue: ... def __init__( self, *, key: global___FieldValue | None = ..., value: global___FieldValue | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... global___MapTypeEntry = MapTypeEntry @typing_extensions.final class LogicalTypeValue(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor VALUE_FIELD_NUMBER: builtins.int @property def value(self) -> global___FieldValue: ... def __init__( self, *, value: global___FieldValue | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["value", b"value"]) -> None: ... global___LogicalTypeValue = LogicalTypeValue @typing_extensions.final class SchemaCoderPayload(google.protobuf.message.Message): """Information needed to represent a Coder of type SCHEMA.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class AdditionalCoderInfo(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor URN_FIELD_NUMBER: builtins.int PAYLOAD_FIELD_NUMBER: builtins.int urn: builtins.str payload: builtins.bytes def __init__( self, *, urn: builtins.str | None = ..., payload: builtins.bytes | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["payload", b"payload", "urn", b"urn"]) -> None: ... SCHEMA_FIELD_NUMBER: builtins.int TO_ROW_FN_FIELD_NUMBER: builtins.int FROM_ROW_FN_FIELD_NUMBER: builtins.int ADDITIONAL_CODER_INFOS_FIELD_NUMBER: builtins.int @property def schema(self) -> global___Schema: """The schema to use for encoding corresponding Row types.""" @property def to_row_fn(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.FunctionSpec: """Function mapping from underlying object to Row type.""" @property def from_row_fn(self) -> org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.FunctionSpec: """Function mapping from Row type to underlying object.""" @property def additional_coder_infos(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___SchemaCoderPayload.AdditionalCoderInfo]: """Any additional information SDKs need to encode/decode elements.""" def __init__( self, *, schema: global___Schema | None = ..., to_row_fn: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.FunctionSpec | None = ..., from_row_fn: org.apache.beam.model.pipeline.v1.beam_runner_api_pb2.FunctionSpec | None = ..., additional_coder_infos: collections.abc.Iterable[global___SchemaCoderPayload.AdditionalCoderInfo] | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["from_row_fn", b"from_row_fn", "schema", b"schema", "to_row_fn", b"to_row_fn"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["additional_coder_infos", b"additional_coder_infos", "from_row_fn", b"from_row_fn", "schema", b"schema", "to_row_fn", b"to_row_fn"]) -> None: ... global___SchemaCoderPayload = SchemaCoderPayload ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/schema_pb2_grpc.py0000644000175100017510000000023715206113151032533 0ustar00runnerrunner# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/schema_pb2_urns.py0000644000175100017510000000302515206113151032565 0ustar00runnerrunnerfrom ........utils import PropertiesFromEnumValue from . import metrics_pb2 EMPTY_MONITORING_INFO_LABEL_PROPS = metrics_pb2.MonitoringInfoLabelProps() EMPTY_MONITORING_INFO_SPEC = metrics_pb2.MonitoringInfoSpec() class LogicalTypes(object): class Enum(object): PYTHON_CALLABLE = PropertiesFromEnumValue('beam:logical_type:python_callable:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MICROS_INSTANT = PropertiesFromEnumValue('beam:logical_type:micros_instant:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) MILLIS_INSTANT = PropertiesFromEnumValue('beam:logical_type:millis_instant:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) DECIMAL = PropertiesFromEnumValue('beam:logical_type:decimal:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) FIXED_BYTES = PropertiesFromEnumValue('beam:logical_type:fixed_bytes:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) VAR_BYTES = PropertiesFromEnumValue('beam:logical_type:var_bytes:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) FIXED_CHAR = PropertiesFromEnumValue('beam:logical_type:fixed_char:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) VAR_CHAR = PropertiesFromEnumValue('beam:logical_type:var_char:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) DATE = PropertiesFromEnumValue('beam:logical_type:date:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ././@PaxHeader0000000000000000000000000000021300000000000010211 xustar00117 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/standard_window_fns_pb2.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/standard_window_fns0000644000175100017510000001111115206113151033114 0ustar00runnerrunner# -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # source: org/apache/beam/model/pipeline/v1/standard_window_fns.proto # Protobuf Python Version: 4.25.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() from . import beam_runner_api_pb2 as org_dot_apache_dot_beam_dot_model_dot_pipeline_dot_v1_dot_beam__runner__api__pb2 from google.protobuf import duration_pb2 as google_dot_protobuf_dot_duration__pb2 from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n;org/apache/beam/model/pipeline/v1/standard_window_fns.proto\x12!org.apache.beam.model.pipeline.v1\x1a\x37org/apache/beam/model/pipeline/v1/beam_runner_api.proto\x1a\x1egoogle/protobuf/duration.proto\x1a\x1fgoogle/protobuf/timestamp.proto\"V\n\x14GlobalWindowsPayload\">\n\x04\x45num\x12\x36\n\nPROPERTIES\x10\x00\x1a&\xa2\xb4\xfa\xc2\x05 beam:window_fn:global_windows:v1\"\xa9\x01\n\x13\x46ixedWindowsPayload\x12\'\n\x04size\x18\x01 \x01(\x0b\x32\x19.google.protobuf.Duration\x12*\n\x06offset\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\"=\n\x04\x45num\x12\x35\n\nPROPERTIES\x10\x00\x1a%\xa2\xb4\xfa\xc2\x05\x1f\x62\x65\x61m:window_fn:fixed_windows:v1\"\xd8\x01\n\x15SlidingWindowsPayload\x12\'\n\x04size\x18\x01 \x01(\x0b\x32\x19.google.protobuf.Duration\x12*\n\x06offset\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12)\n\x06period\x18\x03 \x01(\x0b\x32\x19.google.protobuf.Duration\"?\n\x04\x45num\x12\x37\n\nPROPERTIES\x10\x00\x1a\'\xa2\xb4\xfa\xc2\x05!beam:window_fn:sliding_windows:v1\"\x85\x01\n\x15SessionWindowsPayload\x12+\n\x08gap_size\x18\x01 \x01(\x0b\x32\x19.google.protobuf.Duration\"?\n\x04\x45num\x12\x37\n\nPROPERTIES\x10\x00\x1a\'\xa2\xb4\xfa\xc2\x05!beam:window_fn:session_windows:v1B\x80\x01\n!org.apache.beam.model.pipeline.v1B\x11StandardWindowFnsZHgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1;pipeline_v1b\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'org.apache.beam.model.pipeline.v1.standard_window_fns_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n!org.apache.beam.model.pipeline.v1B\021StandardWindowFnsZHgithub.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1;pipeline_v1' _globals['_GLOBALWINDOWSPAYLOAD_ENUM'].values_by_name["PROPERTIES"]._options = None _globals['_GLOBALWINDOWSPAYLOAD_ENUM'].values_by_name["PROPERTIES"]._serialized_options = b'\242\264\372\302\005 beam:window_fn:global_windows:v1' _globals['_FIXEDWINDOWSPAYLOAD_ENUM'].values_by_name["PROPERTIES"]._options = None _globals['_FIXEDWINDOWSPAYLOAD_ENUM'].values_by_name["PROPERTIES"]._serialized_options = b'\242\264\372\302\005\037beam:window_fn:fixed_windows:v1' _globals['_SLIDINGWINDOWSPAYLOAD_ENUM'].values_by_name["PROPERTIES"]._options = None _globals['_SLIDINGWINDOWSPAYLOAD_ENUM'].values_by_name["PROPERTIES"]._serialized_options = b'\242\264\372\302\005!beam:window_fn:sliding_windows:v1' _globals['_SESSIONWINDOWSPAYLOAD_ENUM'].values_by_name["PROPERTIES"]._options = None _globals['_SESSIONWINDOWSPAYLOAD_ENUM'].values_by_name["PROPERTIES"]._serialized_options = b'\242\264\372\302\005!beam:window_fn:session_windows:v1' _globals['_GLOBALWINDOWSPAYLOAD']._serialized_start=220 _globals['_GLOBALWINDOWSPAYLOAD']._serialized_end=306 _globals['_GLOBALWINDOWSPAYLOAD_ENUM']._serialized_start=244 _globals['_GLOBALWINDOWSPAYLOAD_ENUM']._serialized_end=306 _globals['_FIXEDWINDOWSPAYLOAD']._serialized_start=309 _globals['_FIXEDWINDOWSPAYLOAD']._serialized_end=478 _globals['_FIXEDWINDOWSPAYLOAD_ENUM']._serialized_start=417 _globals['_FIXEDWINDOWSPAYLOAD_ENUM']._serialized_end=478 _globals['_SLIDINGWINDOWSPAYLOAD']._serialized_start=481 _globals['_SLIDINGWINDOWSPAYLOAD']._serialized_end=697 _globals['_SLIDINGWINDOWSPAYLOAD_ENUM']._serialized_start=634 _globals['_SLIDINGWINDOWSPAYLOAD_ENUM']._serialized_end=697 _globals['_SESSIONWINDOWSPAYLOAD']._serialized_start=700 _globals['_SESSIONWINDOWSPAYLOAD']._serialized_end=833 _globals['_SESSIONWINDOWSPAYLOAD_ENUM']._serialized_start=770 _globals['_SESSIONWINDOWSPAYLOAD_ENUM']._serialized_end=833 # @@protoc_insertion_point(module_scope) ././@PaxHeader0000000000000000000000000000021400000000000010212 xustar00118 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/standard_window_fns_pb2.pyi 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/standard_window_fns0000644000175100017510000001761015206113151033126 0ustar00runnerrunner""" @generated by mypy-protobuf. Do not edit manually! isort:skip_file Protocol Buffers describing the Runner API, which is the runner-independent, SDK-independent definition of the Beam model. """ import builtins import google.protobuf.descriptor import google.protobuf.duration_pb2 import google.protobuf.internal.enum_type_wrapper import google.protobuf.message import google.protobuf.timestamp_pb2 import sys import typing if sys.version_info >= (3, 10): import typing as typing_extensions else: import typing_extensions DESCRIPTOR: google.protobuf.descriptor.FileDescriptor @typing_extensions.final class GlobalWindowsPayload(google.protobuf.message.Message): """By default, all data in a PCollection is assigned to the single global window. See BeamConstants for the time span this window encompasses. See https://beam.apache.org/documentation/programming-guide/#single-global-window for additional details. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[GlobalWindowsPayload._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor PROPERTIES: GlobalWindowsPayload._Enum.ValueType # 0 class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... PROPERTIES: GlobalWindowsPayload.Enum.ValueType # 0 def __init__( self, ) -> None: ... global___GlobalWindowsPayload = GlobalWindowsPayload @typing_extensions.final class FixedWindowsPayload(google.protobuf.message.Message): """A fixed time window represents a consistent duration size, non overlapping time interval in the data stream. See https://beam.apache.org/documentation/programming-guide/#fixed-time-windows for additional details. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[FixedWindowsPayload._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor PROPERTIES: FixedWindowsPayload._Enum.ValueType # 0 class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... PROPERTIES: FixedWindowsPayload.Enum.ValueType # 0 SIZE_FIELD_NUMBER: builtins.int OFFSET_FIELD_NUMBER: builtins.int @property def size(self) -> google.protobuf.duration_pb2.Duration: """(Required) Represents the size of the window.""" @property def offset(self) -> google.protobuf.timestamp_pb2.Timestamp: """(Required) Represents the timestamp of when the first window begins. Window N will start at offset + N * size. """ def __init__( self, *, size: google.protobuf.duration_pb2.Duration | None = ..., offset: google.protobuf.timestamp_pb2.Timestamp | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["offset", b"offset", "size", b"size"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["offset", b"offset", "size", b"size"]) -> None: ... global___FixedWindowsPayload = FixedWindowsPayload @typing_extensions.final class SlidingWindowsPayload(google.protobuf.message.Message): """A sliding time window represents time intervals in the data stream that can overlap. For example, each window might capture 60 seconds worth of data, but a new window starts every 30 seconds. The frequency with which sliding windows begin is called the period. Therefore, our example would have a window size of 60 seconds and a period of 30 seconds. Because multiple windows overlap, most elements in a data set will belong to more than one window. This kind of windowing is useful for taking running averages of data; using sliding time windows, you can compute a running average of the past 60 seconds’ worth of data, updated every 30 seconds, in our example. See https://beam.apache.org/documentation/programming-guide/#sliding-time-windows for additional details. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[SlidingWindowsPayload._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor PROPERTIES: SlidingWindowsPayload._Enum.ValueType # 0 class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... PROPERTIES: SlidingWindowsPayload.Enum.ValueType # 0 SIZE_FIELD_NUMBER: builtins.int OFFSET_FIELD_NUMBER: builtins.int PERIOD_FIELD_NUMBER: builtins.int @property def size(self) -> google.protobuf.duration_pb2.Duration: """(Required) Represents the size of the window.""" @property def offset(self) -> google.protobuf.timestamp_pb2.Timestamp: """(Required) Represents the timestamp of when the first window begins. Window N will start at offset + N * period. """ @property def period(self) -> google.protobuf.duration_pb2.Duration: """(Required) Represents the amount of time between each start of a window.""" def __init__( self, *, size: google.protobuf.duration_pb2.Duration | None = ..., offset: google.protobuf.timestamp_pb2.Timestamp | None = ..., period: google.protobuf.duration_pb2.Duration | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["offset", b"offset", "period", b"period", "size", b"size"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["offset", b"offset", "period", b"period", "size", b"size"]) -> None: ... global___SlidingWindowsPayload = SlidingWindowsPayload @typing_extensions.final class SessionWindowsPayload(google.protobuf.message.Message): """A session window function defines windows that contain elements that are within a certain gap size of another element. Session windowing applies on a per-key basis and is useful for data that is irregularly distributed with respect to time. For example, a data stream representing user mouse activity may have long periods of idle time interspersed with high concentrations of clicks. If data arrives after the minimum specified gap size duration, this initiates the start of a new window. See https://beam.apache.org/documentation/programming-guide/#session-windows for additional details. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Enum: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _EnumEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[SessionWindowsPayload._Enum.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor PROPERTIES: SessionWindowsPayload._Enum.ValueType # 0 class Enum(_Enum, metaclass=_EnumEnumTypeWrapper): ... PROPERTIES: SessionWindowsPayload.Enum.ValueType # 0 GAP_SIZE_FIELD_NUMBER: builtins.int @property def gap_size(self) -> google.protobuf.duration_pb2.Duration: """(Required) Minimum duration of gaps between sessions.""" def __init__( self, *, gap_size: google.protobuf.duration_pb2.Duration | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["gap_size", b"gap_size"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["gap_size", b"gap_size"]) -> None: ... global___SessionWindowsPayload = SessionWindowsPayload ././@PaxHeader0000000000000000000000000000022000000000000010207 xustar00122 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/standard_window_fns_pb2_grpc.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/standard_window_fns0000644000175100017510000000023715206113151033123 0ustar00runnerrunner# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc ././@PaxHeader0000000000000000000000000000022000000000000010207 xustar00122 path=apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/standard_window_fns_pb2_urns.py 22 mtime=1779996265.0 apache_beam-2.74.0/apache_beam/portability/api/org/apache/beam/model/pipeline/v1/standard_window_fns0000644000175100017510000000201315206113151033115 0ustar00runnerrunnerfrom ........utils import PropertiesFromEnumValue from . import metrics_pb2 EMPTY_MONITORING_INFO_LABEL_PROPS = metrics_pb2.MonitoringInfoLabelProps() EMPTY_MONITORING_INFO_SPEC = metrics_pb2.MonitoringInfoSpec() class FixedWindowsPayload(object): class Enum(object): PROPERTIES = PropertiesFromEnumValue('beam:window_fn:fixed_windows:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class GlobalWindowsPayload(object): class Enum(object): PROPERTIES = PropertiesFromEnumValue('beam:window_fn:global_windows:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class SessionWindowsPayload(object): class Enum(object): PROPERTIES = PropertiesFromEnumValue('beam:window_fn:session_windows:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) class SlidingWindowsPayload(object): class Enum(object): PROPERTIES = PropertiesFromEnumValue('beam:window_fn:sliding_windows:v1', '', EMPTY_MONITORING_INFO_SPEC, EMPTY_MONITORING_INFO_LABEL_PROPS) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996183.0 apache_beam-2.74.0/apache_beam/portability/api/standard_coders.yaml0000644000175100017510000005563515206113027025035 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # This file is broken into multiple sections delimited by ---. Each section specifies a set of # reference encodings for a single standardized coder used in a specific context. # # Each section contains up to 3 properties: # # coder: a common coder spec. Currently, a URN and URNs for component coders as necessary. # nested: a boolean meaning whether the coder was used in the nested context. Missing means to # test both contexts, a shorthand for when the coder is invariant across context. # examples: a map of {encoded bytes: original JSON object} encoded with the coder in the context. # The LHS (key) is a byte array encoded as a JSON-escaped string. The RHS (value) is # one of a few standard JSON types such as numbers, strings, dicts that map naturally # to the type encoded by the coder. # # Java code snippet to generate example bytes: # Coder> coder = Timer.Coder.of(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); # Instant now = new Instant(1000L); # Timer timer = Timer.of( # "key", # "tag", # Collections.singletonList(GlobalWindow.INSTANCE), # now, # now, # PaneInfo.NO_FIRING); # byte[] bytes = CoderUtils.encodeToByteArray(coder, timer); # String str = new String(bytes, java.nio.charset.StandardCharsets.ISO_8859_1); # String example = ""; # for(int i = 0; i < str.length(); i++){ # example += CharUtils.unicodeEscaped(str.charAt(i)); # } # # These choices were made to strike a balance between portability, ease of use, and simple # legibility of this file itself. # # It is expected that future work will move the `coder` field into a format that it would be # represented by the Runner API, so that it can be understood by all SDKs and harnesses. # # If a coder is marked non-deterministic in the coder spec, then only the decoding should be validated. coder: urn: "beam:coder:bytes:v1" nested: false examples: "abc": abc "ab\0c": "ab\0c" --- coder: urn: "beam:coder:bytes:v1" nested: true examples: "\u0003abc": abc "\u0004ab\0c": "ab\0c" "\u00c8\u0001 10| 20| 30| 40| 50| 60| 70| 80| 90| 100| 110| 120| 130| 140| 150| 160| 170| 180| 190| 200|": " 10| 20| 30| 40| 50| 60| 70| 80| 90| 100| 110| 120| 130| 140| 150| 160| 170| 180| 190| 200|" --- coder: urn: "beam:coder:bool:v1" examples: "\0": false "\u0001": true --- coder: urn: "beam:coder:string_utf8:v1" nested: false examples: "abc": abc "ab\0c": "ab\0c" "\u00c3\u00bf": "\u00ff" "\u00e5\u0085\u0089\u00e7\u00ba\u00bf": "光线" --- coder: urn: "beam:coder:string_utf8:v1" nested: true examples: "\u0003abc": abc "\u0004ab\0c": "ab\0c" "\u0002\u00c3\u00bf": "\u00ff" "\u0006\u00e5\u0085\u0089\u00e7\u00ba\u00bf": "光线" "\u00c8\u0001 10| 20| 30| 40| 50| 60| 70| 80| 90| 100| 110| 120| 130| 140| 150| 160| 170| 180| 190| 200|": " 10| 20| 30| 40| 50| 60| 70| 80| 90| 100| 110| 120| 130| 140| 150| 160| 170| 180| 190| 200|" --- coder: urn: "beam:coder:varint:v1" examples: "\0": 0 "\u0001": 1 "\u000A": 10 "\u00c8\u0001": 200 "\u00e8\u0007": 1000 "\u00a9\u0046": 9001 "\u00ff\u00ff\u00ff\u00ff\u00ff\u00ff\u00ff\u00ff\u00ff\u0001": -1 --- coder: urn: "beam:coder:kv:v1" components: [{urn: "beam:coder:bytes:v1"}, {urn: "beam:coder:varint:v1"}] examples: "\u0003abc\0": {key: abc, value: 0} "\u0004ab\0c\u000A": {key: "ab\0c", value: 10} --- coder: urn: "beam:coder:kv:v1" components: [{urn: "beam:coder:bytes:v1"}, {urn: "beam:coder:bytes:v1"}] nested: false examples: "\u0003abcdef": {key: abc, value: def} "\u0004ab\0cde\0f": {key: "ab\0c", value: "de\0f"} --- coder: urn: "beam:coder:kv:v1" components: [{urn: "beam:coder:bytes:v1"}, {urn: "beam:coder:bytes:v1"}] nested: true examples: "\u0003abc\u0003def": {key: abc, value: def} "\u0004ab\0c\u0004de\0f": {key: "ab\0c", value: "de\0f"} --- coder: urn: "beam:coder:kv:v1" components: [{urn: "beam:coder:bytes:v1"}, {urn: "beam:coder:bool:v1"}] examples: "\u0003abc\u0001": {key: abc, value: true} "\u0004ab\0c\0": {key: "ab\0c", value: false} --- coder: urn: "beam:coder:interval_window:v1" examples: "\u0080\u0000\u0001\u0052\u009a\u00a4\u009b\u0068\u0080\u00dd\u00db\u0001" : {end: 1454293425000, span: 3600000} "\u0080\u0000\u0001\u0053\u0034\u00ec\u0074\u00e8\u0080\u0090\u00fb\u00d3\u0009" : {end: 1456881825000, span: 2592000000} "\u007f\u00df\u003b\u0064\u005a\u001c\u00ad\u0076\u00ed\u0002" : {end: -9223372036854410, span: 365} "\u0080\u0020\u00c4\u009b\u00a5\u00e3\u0053\u00f7\u0000" : {end: 9223372036854775, span: 0} --- coder: urn: "beam:coder:iterable:v1" components: [{urn: "beam:coder:varint:v1"}] examples: "\0\0\0\u0001\0": [0] "\0\0\0\u0004\u0001\n\u00c8\u0001\u00e8\u0007": [1, 10, 200, 1000] # 130 items. "\0\0\0\u0082A-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA-A": [65, 45, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 65, 45, 65] "\0\0\0\0": [] --- coder: urn: "beam:coder:iterable:v1" components: [{urn: "beam:coder:bytes:v1"}] examples: "\0\0\0\u0001\u0003abc": ["abc"] "\0\0\0\u0002\u0004ab\0c\u0004de\0f": ["ab\0c", "de\0f"] "\0\0\0\0": [] --- coder: urn: "beam:coder:iterable:v1" components: [{urn: "beam:coder:bool:v1"}] examples: "\0\0\0\u0001\u0001": [true] "\0\0\0\u0002\u0001\0": [true, false] "\0\0\0\0": [] --- coder: urn: "beam:coder:iterable:v1" components: [{urn: "beam:coder:bytes:v1"}] # This is for iterables of unknown length, where the encoding is not # deterministic. non_deterministic: true examples: "\u00ff\u00ff\u00ff\u00ff\u0000": [] "\u00ff\u00ff\u00ff\u00ff\u0001\u0003abc\u0000": ["abc"] "\u00ff\u00ff\u00ff\u00ff\u0002\u0004ab\u0000c\u0004de\u0000f\u0000": ["ab\0c", "de\0f"] --- coder: urn: "beam:coder:iterable:v1" components: [{urn: "beam:coder:global_window:v1"}] examples: "\0\0\0\u0001": [""] --- coder: urn: "beam:coder:timer:v1" components: [{urn: "beam:coder:string_utf8:v1"}, {urn: "beam:coder:global_window:v1"}] examples: "\u0003\u006b\u0065\u0079\u0003\u0074\u0061\u0067\u0000\u0000\u0000\u0001\u0000\u0080\u0000\u0000\u0000\u0000\u0000\u0004\u00d2\u0080\u0000\u0000\u0000\u0000\u0000\u0016\u002e\u000f": { userKey: key, dynamicTimerTag: tag, windows: ["global"], clearBit: false, fireTimestamp: 1234, holdTimestamp: 5678, pane: {is_first: true, is_last: true, timing: UNKNOWN, index: 0, on_time_index: 0}, } "\u0003\u006b\u0065\u0079\u0003\u0074\u0061\u0067\u0000\u0000\u0000\u0001\u0001": { userKey: key, dynamicTimerTag: tag, windows: ["global"], clearBit: true, } --- coder: urn: "beam:coder:global_window:v1" examples: "": "" --- # All windowed values consist of pane infos that represent NO_FIRING until full support is added # in the Python SDK (BEAM-1522). coder: urn: "beam:coder:windowed_value:v1" components: [{urn: "beam:coder:varint:v1"}, {urn: "beam:coder:global_window:v1"}] examples: "\u0080\0\u0001R\u009a\u00a4\u009bh\0\0\0\u0001\u000f\u0002": { value: 2, timestamp: 1454293425000, pane: {is_first: true, is_last: true, timing: UNKNOWN, index: 0, on_time_index: 0}, windows: ["global"] } --- coder: urn: "beam:coder:windowed_value:v1" components: [{urn: "beam:coder:varint:v1"}, {urn: "beam:coder:interval_window:v1"}] examples: "\u007f\u00ff\u00ff\u00ff\u00ff\u00f9\u00e5\u0080\0\0\0\u0001\u0080\0\u0001R\u009a\u00a4\u009bh\u00c0\u008b\u0011\u000f\u0004": { value: 4, timestamp: -400000, pane: {is_first: true, is_last: true, timing: UNKNOWN, index: 0, on_time_index: 0}, windows: [{end: 1454293425000, span: 280000}] } "\u007f\u00ff\u00ff\u00ff\u00ff\u00ff\u00ff\u009c\0\0\0\u0002\u0080\0\u0001R\u009a\u00a4\u009bh\u0080\u00dd\u00db\u0001\u007f\u00df;dZ\u001c\u00adv\u00ed\u0002\u000f\u0002": { value: 2, timestamp: -100, pane: {is_first: true, is_last: true, timing: UNKNOWN, index: 0, on_time_index: 0}, windows: [{end: 1454293425000, span: 3600000}, {end: -9223372036854410, span: 365}] } --- # ParamWindowedValueCoder with constant value of: # timestamp: Instant.ofEpochMilli(1000) # windows: [IntervalWindow(10, 20)] # pane info: PaneInfo(false, true, PaneInfo.Timing.ON_TIME, 30, 40) coder: urn: "beam:coder:param_windowed_value:v1" payload: "\x80\x00\x00\x00\x00\x00\x03è\x00\x00\x00\x01\x80\x00\x00\x00\x00\x00\x00\x14\n&\x1E(\x00" components: [{urn: "beam:coder:varint:v1"}, {urn: "beam:coder:interval_window:v1"}] examples: "\u0002": { value: 2, timestamp: 1000, pane: {is_first: false, is_last: true, timing: ON_TIME, index: 30, on_time_index: 40}, windows: [{end: 20, span: 10}] } --- # ParamWindowedValueCoder with constant value of: # timestamp: BoundedWindow.TIMESTAMP_MIN_VALUE # windows: [GlobalWindow.INSTANCE] # pane info: PaneInfo.NO_FIRING coder: urn: "beam:coder:param_windowed_value:v1" payload: "\x7Fß;dZ\x1C¬\t\x00\x00\x00\x01\x0F\x00" components: [{urn: "beam:coder:varint:v1"}, {urn: "beam:coder:global_window:v1"}] examples: "\u0002": { value: 2, timestamp: -9223372036854775, pane: {is_first: true, is_last: true, timing: UNKNOWN, index: 0, on_time_index: 0}, windows: ["global"] } --- coder: urn: "beam:coder:double:v1" examples: "\0\0\0\0\0\0\0\0": "0" "\u0080\0\0\0\0\0\0\0": "-0" "\u003f\u00b9\u0099\u0099\u0099\u0099\u0099\u009a": "0.1" "\u00bf\u00b9\u0099\u0099\u0099\u0099\u0099\u009a": "-0.1" "\0\0\0\0\0\0\0\u0001": "4.9e-324" "\0\u0001\0\0\0\0\0\0": "1.390671161567e-309" "\u007f\u00ef\u00ff\u00ff\u00ff\u00ff\u00ff\u00ff": "1.7976931348623157e308" "\u007f\u00f0\0\0\0\0\0\0": "Infinity" "\u00ff\u00f0\0\0\0\0\0\0": "-Infinity" "\u007f\u00f8\0\0\0\0\0\0": "NaN" --- coder: urn: "beam:coder:row:v1" # str: string, i32: int32, f64: float64, arr: array[string] payload: "\n\t\n\x03str\x1a\x02\x10\x07\n\t\n\x03i32\x1a\x02\x10\x03\n\t\n\x03f64\x1a\x02\x10\x06\n\r\n\x03arr\x1a\x06\x1a\x04\n\x02\x10\x07\x12$4e5e554c-d4c1-4a5d-b5e1-f3293a6b9f05" examples: "\u0004\u0000\u0003foo\u00a9\u0046\u003f\u00b9\u0099\u0099\u0099\u0099\u0099\u009a\0\0\0\u0003\u0003foo\u0003bar\u0003baz": {str: "foo", i32: 9001, f64: "0.1", arr: ["foo", "bar", "baz"]} --- coder: urn: "beam:coder:row:v1" # str: nullable string, i32: nullable int32, f64: nullable float64 payload: "\n\x0b\n\x03str\x1a\x04\x08\x01\x10\x07\n\x0b\n\x03i32\x1a\x04\x08\x01\x10\x03\n\x0b\n\x03f64\x1a\x04\x08\x01\x10\x06\x12$b20c6545-57af-4bc8-b2a9-51ace21c7393" examples: "\u0003\u0001\u0007": {str: null, i32: null, f64: null} "\u0003\u0001\u0004\u0003foo\u00a9\u0046": {str: "foo", i32: 9001, f64: null} "\u0003\u0000\u0003foo\u00a9\u0046\u003f\u00b9\u0099\u0099\u0099\u0099\u0099\u009a": {str: "foo", i32: 9001, f64: "0.1"} --- # Binary data generated with the python SDK: # # import typing # import apache_beam as beam # class Test(typing.NamedTuple): # f_bool: bool # f_byte: typing.Optional[bytes] # schema = beam.typehints.schemas.named_tuple_to_schema(Test) # coder = beam.coders.row_coder.RowCoder(schema) # print("payload = %s" % schema.SerializeToString()) # print("example = %s" % coder.encode(Test(f_bool=True, f_nullable=None))) coder: urn: "beam:coder:row:v1" # f_bool: boolean, f_bytes: nullable bytes payload: "\n\x0c\n\x06f_bool\x1a\x02\x10\x08\n\x0f\n\x07f_bytes\x1a\x04\x08\x01\x10\t\x12$eea1b747-7571-43d3-aafa-9255afdceafb" examples: "\x02\x01\x02\x01": {f_bool: true, f_bytes: null} "\x02\x00\x00\x04ab\x00c": {f_bool: false, f_bytes: "ab\0c"} --- # Binary data generated with the python SDK: # # import typing # import apache_beam as beam # class Test(typing.NamedTuple): # f_map: typing.Mapping[str,typing.Optional[int]] # schema = beam.typehints.schemas.named_tuple_to_schema(Test) # coder = beam.coders.row_coder.RowCoder(schema) # print("payload = %s" % schema.SerializeToString()) # examples = (Test(f_map={}), # Test(f_map={"foo": 9001, "bar": 9223372036854775807}), # Test(f_map={"everything": None, "is": None, "null!": None, "¯\_(ツ)_/¯": None})) # for example in examples: # print("example = %s" % coder.encode(example)) coder: urn: "beam:coder:row:v1" # f_map: map payload: "\n\x15\n\x05f_map\x1a\x0c*\n\n\x02\x10\x07\x12\x04\x08\x01\x10\x04\x12$d8c8f969-14e6-457f-a8b5-62a1aec7f1cd" # map ordering is non-deterministic non_deterministic: true examples: "\x01\x00\x00\x00\x00\x00": {f_map: {}} "\x01\x00\x00\x00\x00\x02\x03foo\x01\xa9F\x03bar\x01\xff\xff\xff\xff\xff\xff\xff\xff\x7f": {f_map: {"foo": 9001, "bar": 9223372036854775807}} "\x01\x00\x00\x00\x00\x04\neverything\x00\x02is\x00\x05null!\x00\r\xc2\xaf\\_(\xe3\x83\x84)_/\xc2\xaf\x00": {f_map: {"everything": null, "is": null, "null!": null, "¯\\_(ツ)_/¯": null}} --- # Binary data generated with the python SDK: # schema1 = schema_pb2.Schema( # id="30ea5a25-dcd8-4cdb-abeb-5332d15ab4b9", # fields=[ # schema_pb2.Field( # name="str", # type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING), # encoding_position=1), # schema_pb2.Field( # name="f_bool", # type=schema_pb2.FieldType(atomic_type=schema_pb2.BOOLEAN), # encoding_position=2), # schema_pb2.Field( # name="i32", # type=schema_pb2.FieldType( # atomic_type=schema_pb2.INT32, nullable=True), # encoding_position=0) # ], # encoding_positions_set=True) # # coder = RowCoder(schema1) # c = coder.schema.SerializeToString() # print("payload = %s" % c) # test = typing.NamedTuple("test", [ ("f_bool", bool), ("i32", np.int32), ("str", str) ]) # example = coder.encode(test(False,21,"str2")) # print("example = %s" % example) coder: urn: "beam:coder:row:v1" payload: "\n\x0b\n\x03str\x1a\x02\x10\x07(\x01\n\x0e\n\x06f_bool\x1a\x02\x10\x08(\x02\n\x0b\n\x03i32\x1a\x04\x08\x01\x10\x03\x12$30ea5a25-dcd8-4cdb-abeb-5332d15ab4b9 \x01" examples: "\x03\x00\x15\x04str2\x00": {f_bool: false, i32: 21, str: "str2"} --- coder: urn: "beam:coder:row:v1" # f_timestamp: logical(micros_instant), f_string: string, f_int: int64 payload: "\n\x7f\n\x0bf_timestamp\x1ap:n\n#beam:logical_type:micros_instant:v1\x1aG2E\nC\n\r\n\x07seconds\x1a\x02\x10\x04\n\x0c\n\x06micros\x1a\x02\x10\x04\x12$4d3f6e8f-7412-4ad7-bfd9-b424a1664aef\n\x0e\n\x08f_string\x1a\x02\x10\x07\n\x0b\n\x05f_int\x1a\x02\x10\x04\x12$33dafd37-397c-4083-a84e-42177d122221" examples: "\x03\x00\x02\x00\xb6\x95\xd5\xf9\x05\xc0\xc4\x07\x1b2020-08-13T14:14:14.123456Z\xc0\xf7\x85\xda\xae\x98\xeb\x02": {f_timestamp: {seconds: 1597328054, micros: 123456}, f_string: "2020-08-13T14:14:14.123456Z", f_int: 1597328054123456} --- coder: urn: "beam:coder:row:v1" # f_timestamp: logical(millis_instant), f_string: string, f_int: int64 payload: "\n:\n\x0bf_timestamp\x1a+:)\n#beam:logical_type:millis_instant:v1\x1a\x02\x10\x04\n\x0e\n\x08f_string\x1a\x02\x10\x07\n\x0b\n\x05f_int\x1a\x02\x10\x04\x12$80be749a-5700-4ede-89d8-dd9a4433a3f8" examples: "\x03\x00\x80\x00\x01s\xe8+\xd7k\x182020-08-13T14:14:14.123Z\xeb\xae\xaf\xc1\xbe.": {f_timestamp: -9223370439526721685, f_string: "2020-08-13T14:14:14.123Z", f_int: 1597328054123} --- coder: urn: "beam:coder:row:v1" # f_float: float32 payload: "\n\r\n\x07f_float\x1a\x02\x10\x05\x12$8c97b6c5-69e5-4733-907b-26cd8edae612" examples: "\x01\x00\x00\x00\x00\x00": {f_float: "0.0"} "\x01\x00?\x80\x00\x00": {f_float: "1.0"} "\x01\x00@I\x0eV": {f_float: "3.1415"} --- coder: urn: "beam:coder:row:v1" # f_float: float32, f_decimal: logical(decimal) payload: "\n\r\n\x07f_float\x1a\x02\x10\x05\n1\n\tf_decimal\x1a$:\"\n\x1cbeam:logical_type:decimal:v1\x1a\x02\x10\t\x12$800c44ae-a1b7-4def-bbf6-6217cca89ec4" examples: "\x02\x00\x00\x00\x00\x00\x01\x01\x00": {f_float: "0.0", f_decimal: "0.0"} "\x02\x00?\x80\x00\x00\x01\x01\n": {f_float: "1.0", f_decimal: "1.0"} "\x02\x00@I\x0eV\x04\x02z\xb7": {f_float: "3.1415", f_decimal: "3.1415"} "\x02\x00\xc2\xc8>\xfa\x03\x03\xfex\xe5": {f_float: "-100.123", f_decimal: "-100.123"} --- coder: urn: "beam:coder:row:v1" # f_char: logical(fixed_char(5)), f_varchar: logical(var_char(5)), f_bytes: logical(fixed_bytes(5)), f_varbytes: logical(var_bytes(5)) payload: "\n=\n\x06f_char\x1a3\x08\x01:/\n\x1fbeam:logical_type:fixed_char:v1\x1a\x02\x10\x07\"\x02\x10\x03*\x04\n\x02\x18\x05\nB\n\tf_varchar\x1a1\x08\x01:-\n\x1dbeam:logical_type:var_char:v1\x1a\x02\x10\x07\"\x02\x10\x03*\x04\n\x02\x18\n \x01(\x01\nC\n\x07f_bytes\x1a4\x08\x01:0\n beam:logical_type:fixed_bytes:v1\x1a\x02\x10\t\"\x02\x10\x03*\x04\n\x02\x18\x05 \x02(\x02\nD\n\nf_varbytes\x1a2\x08\x01:.\n\x1ebeam:logical_type:var_bytes:v1\x1a\x02\x10\t\"\x02\x10\x03*\x04\n\x02\x18\n \x03(\x03\x12$f0ffb3a4-f46f-41ca-a942-85e3e939452a" examples: "\x04\x00\x05ABCDE\x05ABCDE\x05ABCDE\x05ABCDE": {f_char: "ABCDE", f_varchar: "ABCDE", f_bytes: "ABCDE", f_varbytes: "ABCDE"} "\x04\x00\x05A\n \x02A\n\x05A\n\x00\x00\x00\x02A\n": {f_char: "A\n ", f_varchar: "A\n", f_bytes: "A\n\x00\x00\x00", f_varbytes: "A\n"} "\x04\x01\x06\x05null?\x04null": {f_char: "null?", f_varchar: null, f_bytes: null, f_varbytes: "null"} --- coder: urn: "beam:coder:sharded_key:v1" components: [{urn: "beam:coder:string_utf8:v1"}] examples: "\u0000\u0000": { shardId: "", key: "" } "\u0008\u0073\u0068\u0061\u0072\u0064\u005f\u0069\u0064\u0000": { shardId: shard_id, key: "" } "\u0008\u0073\u0068\u0061\u0072\u0064\u005f\u0069\u0064\u0003\u006b\u0065\u0079": { shardId: shard_id, key: "key" } "\u0000\u0003\u006b\u0065\u0079": { shardId: "", key: "key" } --- coder: urn: "beam:coder:state_backed_iterable:v1" components: [{urn: "beam:coder:varint:v1"}] examples: "\0\0\0\u0001\0": [0] "\0\0\0\u0004\u0001\n\u00c8\u0001\u00e8\u0007": [1, 10, 200, 1000] "\0\0\0\0": [] --- coder: urn: "beam:coder:state_backed_iterable:v1" components: [{urn: "beam:coder:bytes:v1"}] examples: "\0\0\0\u0001\u0003abc": ["abc"] "\0\0\0\u0002\u0004ab\0c\u0004de\0f": ["ab\0c", "de\0f"] "\0\0\0\0": [] --- coder: urn: "beam:coder:state_backed_iterable:v1" components: [{urn: "beam:coder:bool:v1"}] examples: "\0\0\0\u0001\u0001": [true] "\0\0\0\u0002\u0001\0": [true, false] "\0\0\0\0": [] --- coder: urn: "beam:coder:state_backed_iterable:v1" components: [{urn: "beam:coder:bytes:v1"}] # This is for iterables of unknown length, where the encoding is not # deterministic. non_deterministic: true examples: "\u00ff\u00ff\u00ff\u00ff\u0000": [] "\u00ff\u00ff\u00ff\u00ff\u0001\u0003abc\u0000": ["abc"] "\u00ff\u00ff\u00ff\u00ff\u0002\u0004ab\u0000c\u0004de\u0000f\u0000": ["ab\0c", "de\0f"] --- coder: urn: "beam:coder:state_backed_iterable:v1" components: [{urn: "beam:coder:global_window:v1"}] examples: "\0\0\0\u0001": [""] --- coder: urn: "beam:coder:state_backed_iterable:v1" components: [{urn: "beam:coder:varint:v1"}] state: { emptySuffix: "", nonEmptySuffix: "\u0001\n" # Represents the values [1, 10] } # This is for iterables of unknown length, where the encoding is not # deterministic. non_deterministic: true examples: "\u00FF\u00FF\u00FF\u00FF\0": [] # Unknown length with empty prefix and no token "\u00FF\u00FF\u00FF\u00FF\u0002\u00C8\u0001\u00E8\u0007\0": [200, 1000] # Unknown length with non-empty prefix and no token "\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u0001\u000BemptySuffix": [] # Unknown length with empty prefix and token with empty suffix "\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u0001\u000EnonEmptySuffix": [1, 10] # Unknown length with empty prefix and token with non-empty suffix "\u00FF\u00FF\u00FF\u00FF\u0002\u00C8\u0001\u00E8\u0007\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u0001\u000BemptySuffix": [200, 1000] # Unknown length with non-empty prefix and token with empty suffix "\u00FF\u00FF\u00FF\u00FF\u0002\u00C8\u0001\u00E8\u0007\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u00FF\u0001\u000EnonEmptySuffix": [200, 1000, 1, 10] # Unknown length with non-empty prefix and token with non-empty suffix --- # Java code snippet to generate example bytes: # TimestampPrefixingWindowCoder coder = TimestampPrefixingWindowCoder.of(IntervalWindowCoder.of()); # Instant end = new Instant(-9223372036854410L); # Duration span = Duration.millis(365L); # IntervalWindow window = new IntervalWindow(end.minus(span), span); # byte[] bytes = CoderUtils.encodeToByteArray(coder, window); # String str = new String(bytes, java.nio.charset.StandardCharsets.ISO_8859_1); # String example = ""; # for(int i = 0; i < str.length(); i++){ # example += CharUtils.unicodeEscaped(str.charAt(i)); # } # System.out.println(example); coder: urn: "beam:coder:custom_window:v1" components: [{urn: "beam:coder:interval_window:v1"}] examples: "\u0080\u0000\u0001\u0052\u009a\u00a4\u009b\u0067\u0080\u0000\u0001\u0052\u009a\u00a4\u009b\u0068\u0080\u00dd\u00db\u0001" : {window: {end: 1454293425000, span: 3600000}} "\u007f\u00df\u003b\u0064\u005a\u001c\u00ad\u0075\u007f\u00df\u003b\u0064\u005a\u001c\u00ad\u0076\u00ed\u0002" : {window: {end: -9223372036854410, span: 365}} --- coder: urn: "beam:coder:nullable:v1" components: [{urn: "beam:coder:bytes:v1"}] nested: true examples: "\u0001\u0003\u0061\u0062\u0063" : "abc" "\u0001\u000a\u006d\u006f\u0072\u0065\u0020\u0062\u0079\u0074\u0065\u0073" : "more bytes" "\u0000" : null ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/portability/common_urns.py0000644000175100017510000001032315206113032023127 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Accessors for URNs of common Beam entities. """ # pytype: skip-file from .api import beam_runner_api_pb2_urns from .api import external_transforms_pb2_urns from .api import metrics_pb2_urns from .api import schema_pb2_urns from .api import standard_window_fns_pb2_urns BeamConstants = beam_runner_api_pb2_urns.BeamConstants StandardArtifacts = beam_runner_api_pb2_urns.StandardArtifacts StandardCoders = beam_runner_api_pb2_urns.StandardCoders StandardDisplayData = beam_runner_api_pb2_urns.StandardDisplayData StandardEnvironments = beam_runner_api_pb2_urns.StandardEnvironments StandardProtocols = beam_runner_api_pb2_urns.StandardProtocols StandardPTransforms = beam_runner_api_pb2_urns.StandardPTransforms StandardRequirements = beam_runner_api_pb2_urns.StandardRequirements StandardResourceHints = beam_runner_api_pb2_urns.StandardResourceHints StandardRunnerProtocols = beam_runner_api_pb2_urns.StandardRunnerProtocols StandardSideInputTypes = beam_runner_api_pb2_urns.StandardSideInputTypes StandardUserStateTypes = beam_runner_api_pb2_urns.StandardUserStateTypes ExpansionMethods = external_transforms_pb2_urns.ExpansionMethods ManagedTransforms = external_transforms_pb2_urns.ManagedTransforms MonitoringInfo = metrics_pb2_urns.MonitoringInfo MonitoringInfoSpecs = metrics_pb2_urns.MonitoringInfoSpecs MonitoringInfoTypeUrns = metrics_pb2_urns.MonitoringInfoTypeUrns LogicalTypes = schema_pb2_urns.LogicalTypes FixedWindowsPayload = standard_window_fns_pb2_urns.FixedWindowsPayload GlobalWindowsPayload = standard_window_fns_pb2_urns.GlobalWindowsPayload SessionWindowsPayload = standard_window_fns_pb2_urns.SessionWindowsPayload SlidingWindowsPayload = standard_window_fns_pb2_urns.SlidingWindowsPayload primitives = StandardPTransforms.Primitives deprecated_primitives = StandardPTransforms.DeprecatedPrimitives composites = StandardPTransforms.Composites combine_components = StandardPTransforms.CombineComponents sdf_components = StandardPTransforms.SplittableParDoComponents group_into_batches_components = StandardPTransforms.GroupIntoBatchesComponents executable_stage = "beam:runner:executable_stage:v1" user_state = StandardUserStateTypes.Enum side_inputs = StandardSideInputTypes.Enum coders = StandardCoders.Enum constants = BeamConstants.Constants environments = StandardEnvironments.Environments artifact_types = StandardArtifacts.Types artifact_roles = StandardArtifacts.Roles resource_hints = StandardResourceHints.Enum global_windows = GlobalWindowsPayload.Enum.PROPERTIES fixed_windows = FixedWindowsPayload.Enum.PROPERTIES sliding_windows = SlidingWindowsPayload.Enum.PROPERTIES session_windows = SessionWindowsPayload.Enum.PROPERTIES monitoring_info_specs = MonitoringInfoSpecs.Enum monitoring_info_types = MonitoringInfoTypeUrns.Enum monitoring_info_labels = MonitoringInfo.MonitoringInfoLabels protocols = StandardProtocols.Enum runner_protocols = StandardRunnerProtocols.Enum requirements = StandardRequirements.Enum displayData = StandardDisplayData.DisplayData java_class_lookup = ExpansionMethods.Enum.JAVA_CLASS_LOOKUP schematransform_based_expand = ExpansionMethods.Enum.SCHEMA_TRANSFORM decimal = LogicalTypes.Enum.DECIMAL micros_instant = LogicalTypes.Enum.MICROS_INSTANT millis_instant = LogicalTypes.Enum.MILLIS_INSTANT python_callable = LogicalTypes.Enum.PYTHON_CALLABLE fixed_bytes = LogicalTypes.Enum.FIXED_BYTES var_bytes = LogicalTypes.Enum.VAR_BYTES fixed_char = LogicalTypes.Enum.FIXED_CHAR var_char = LogicalTypes.Enum.VAR_CHAR date = LogicalTypes.Enum.DATE ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/portability/python_urns.py0000644000175100017510000000572315206113032023170 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Enumeration of URNs specific to the Python SDK. For internal use only; no backwards-compatibility guarantees.""" PICKLED_CODER = "beam:coder:pickled_python:v1" PICKLED_COMBINE_FN = "beam:combinefn:pickled_python:v1" PICKLED_DOFN = "beam:dofn:pickled_python:v1" PICKLED_DOFN_INFO = "beam:dofn:pickled_python_info:v1" PICKLED_SOURCE = "beam:source:pickled_python:v1" PICKLED_TRANSFORM = "beam:transform:pickled_python:v1" PICKLED_WINDOW_MAPPING_FN = "beam:window_mapping_fn:pickled_python:v1" PICKLED_WINDOWFN = "beam:window_fn:pickled_python:v1" PICKLED_VIEWFN = "beam:view_fn:pickled_python_data:v1" IMPULSE_READ_TRANSFORM = "beam:transform:read_from_impulse_python:v1" GENERIC_COMPOSITE_TRANSFORM = "beam:transform:generic_composite:v1" KEY_WITH_NONE_DOFN = "beam:dofn:python_key_with_none:v1" PACKED_COMBINE_FN = "beam:combinefn:packed_python:v1" # A coder for a tuple. # Components: The coders for the tuple elements, in order. TUPLE_CODER = "beam:coder:tuple:v1" # This allows us to decode TimestampedPrefixed(LengthPrefixed(AnyWindowCoder)). TIMESTAMP_PREFIXED_OPAQUE_WINDOW_CODER = ( "beam:timestamp_prefixed_opaque_window_coder:v1") # Invoke UserFns in process, via direct function calls. # Payload: None. EMBEDDED_PYTHON = "beam:env:embedded_python:v1" # Invoke UserFns in process, but over GRPC channels. # Payload: (optional) Number of worker threads, followed by ',' and the size of # the state cache, as a decimal string, e.g. '2,1000'. EMBEDDED_PYTHON_GRPC = "beam:env:embedded_python_grpc:v1" # Invoke UserFns via a yet-to-be-started loopback external worker. # Payload: None. EMBEDDED_PYTHON_LOOPBACK = "beam:env:embedded_python_loopback:v1" # Instantiate SDK harness via a command line provided in the payload. # This is different than the standard process environment in that it # starts up the SDK harness directly, rather than the bootstrapping # and artifact fetching code. # (Used for testing.) SUBPROCESS_SDK = "beam:env:harness_subprocess_python:v1" # An annotation that indicates combiner packing is OK in all sub-transforms # of this transform. This optimization may result in renamed counters and # PCollection element counts. APPLY_COMBINER_PACKING = "beam:annotation:apply_combiner_packing:v1" ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/portability/utils.py0000644000175100017510000000232615206113032021734 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """For internal use only; no backwards-compatibility guarantees.""" from typing import TYPE_CHECKING from typing import NamedTuple if TYPE_CHECKING: from apache_beam.portability.api import metrics_pb2 PropertiesFromEnumValue = NamedTuple( 'PropertiesFromEnumValue', [ ('urn', str), ('constant', str), ('spec', 'metrics_pb2.MonitoringInfoSpec'), ('label_props', 'metrics_pb2.MonitoringInfoLabelProps'), ]) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/programming_guide_test.py0000644000175100017510000000331715206113032022771 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest import apache_beam as beam from apache_beam import metrics from apache_beam.runners.direct.direct_runner import BundleBasedDirectRunner class ProgrammingGuideTest(unittest.TestCase): def test_metrics_example(self): class MyMetricsDoFn(beam.DoFn): def __init__(self): super().__init__() self.counter = metrics.Metrics.counter("namespace", "counter1") def process(self, element): self.counter.inc() yield element with beam.Pipeline(runner=BundleBasedDirectRunner()) as p: _ = (p | beam.Create([1, 2, 3]) | beam.ParDo(MyMetricsDoFn())) metrics_filter = metrics.MetricsFilter().with_name("counter1") query_result = p.result.metrics().query(metrics_filter) for metric in query_result["counters"]: print(metric) # Not in example but just to confirm that anything is returned assert query_result["counters"] if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/pvalue.py0000644000175100017510000005647515206113032017544 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """PValue, PCollection: one node of a dataflow graph. A node of a dataflow processing graph is a PValue. Currently, there is only one type: PCollection (a potentially very large set of arbitrary values). Once created, a PValue belongs to a pipeline and has an associated transform (of type PTransform), which describes how the value will be produced when the pipeline gets executed. """ # pytype: skip-file import collections import itertools from typing import TYPE_CHECKING from typing import Any from typing import Generic from typing import Iterator from typing import NamedTuple from typing import Optional from typing import Sequence from typing import TypeVar from typing import Union from apache_beam import coders from apache_beam import typehints from apache_beam.internal import pickler from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.portability.api import beam_runner_api_pb2 if TYPE_CHECKING: from apache_beam.pipeline import AppliedPTransform from apache_beam.pipeline import Pipeline from apache_beam.runners.pipeline_context import PipelineContext from apache_beam.transforms import sideinputs from apache_beam.transforms.core import ParDo from apache_beam.transforms.core import Windowing __all__ = [ 'PCollection', 'TaggedOutput', 'AsSideInput', 'AsSingleton', 'AsIter', 'AsList', 'AsDict', 'AsMultiMap', 'EmptySideInput', 'Row', ] T = TypeVar('T') class PValue(object): """Base class for PCollection. Dataflow users should not construct PValue objects directly in their pipelines. A PValue has the following main characteristics: (1) Belongs to a pipeline. Added during object initialization. (2) Has a transform that can compute the value if executed. (3) Has a value which is meaningful if the transform was executed. """ def __init__( self, pipeline: 'Pipeline', tag: Optional[str] = None, element_type: Optional[Union[type, 'typehints.TypeConstraint']] = None, windowing: Optional['Windowing'] = None, is_bounded=True, ): """Initializes a PValue with all arguments hidden behind keyword arguments. Args: pipeline: Pipeline object for this PValue. tag: Tag of this PValue. element_type: The type of this PValue. """ self.pipeline = pipeline self.tag = tag self.element_type = element_type # The AppliedPTransform instance for the application of the PTransform # generating this PValue. The field gets initialized when a transform # gets applied. self.producer: Optional[AppliedPTransform] = None self.is_bounded = is_bounded if windowing: self._windowing = windowing self.requires_deterministic_key_coder = None def __str__(self): return self._str_internal() def __repr__(self): return '<%s at %s>' % (self._str_internal(), hex(id(self))) def _str_internal(self): return "%s[%s.%s]" % ( self.__class__.__name__, self.producer.full_label if self.producer else None, self.tag) def apply(self, *args, **kwargs): """Applies a transform or callable to a PValue. Args: *args: positional arguments. **kwargs: keyword arguments. The method will insert the pvalue as the next argument following an optional first label and a transform/callable object. It will call the pipeline.apply() method with this modified argument list. """ arglist = list(args) arglist.insert(1, self) return self.pipeline.apply(*arglist, **kwargs) def __or__(self, ptransform): return self.pipeline.apply(ptransform, self) class PCollection(PValue, Generic[T]): """A multiple values (potentially huge) container. Dataflow users should not construct PCollection objects directly in their pipelines. """ def __eq__(self, other): if isinstance(other, PCollection): return self.tag == other.tag and self.producer == other.producer def __hash__(self): return hash((self.tag, self.producer)) @property def windowing(self) -> 'Windowing': if not hasattr(self, '_windowing'): assert self.producer is not None and self.producer.transform is not None self._windowing = self.producer.transform.get_windowing( self.producer.inputs) return self._windowing def __reduce_ex__(self, unused_version): # Pickling a PCollection is almost always the wrong thing to do, but we # can't prohibit it as it often gets implicitly picked up (e.g. as part # of a closure). return _InvalidUnpickledPCollection, () @staticmethod def from_(pcoll: PValue, is_bounded: Optional[bool] = None) -> 'PCollection': """Create a PCollection, using another PCollection as a starting point. Transfers relevant attributes. """ if is_bounded is None: is_bounded = pcoll.is_bounded return PCollection(pcoll.pipeline, is_bounded=is_bounded) def to_runner_api( self, context: 'PipelineContext') -> beam_runner_api_pb2.PCollection: return beam_runner_api_pb2.PCollection( unique_name=self._unique_name(), coder_id=context.coder_id_from_element_type( self.element_type, self.requires_deterministic_key_coder), is_bounded=beam_runner_api_pb2.IsBounded.BOUNDED if self.is_bounded else beam_runner_api_pb2.IsBounded.UNBOUNDED, windowing_strategy_id=context.windowing_strategies.get_id( self.windowing)) def _unique_name(self) -> str: if self.producer: return '%d%s.%s' % ( len(self.producer.full_label), self.producer.full_label, self.tag) else: return 'PCollection%s' % id(self) @staticmethod def from_runner_api( proto: beam_runner_api_pb2.PCollection, context: 'PipelineContext') -> 'PCollection': # Producer and tag will be filled in later, the key point is that the same # object is returned for the same pcollection id. # We pass None for the PCollection's Pipeline to avoid a cycle during # deserialization. It will be populated soon after this call, in # Pipeline.from_runner_api(). This brief period is the only time that # PCollection.pipeline is allowed to be None. return PCollection( None, # type: ignore[arg-type] element_type=context.element_type_from_coder_id(proto.coder_id), windowing=context.windowing_strategies.get_by_id( proto.windowing_strategy_id), is_bounded=proto.is_bounded == beam_runner_api_pb2.IsBounded.BOUNDED) class _InvalidUnpickledPCollection(object): pass class PBegin(PValue): """A pipeline begin marker used as input to create/read transforms. The class is used internally to represent inputs to Create and Read transforms. This allows us to have transforms that uniformly take PValue(s) as inputs. """ pass class PDone(PValue): """PDone is the output of a transform that has a trivial result such as Write. """ pass class DoOutputsTuple(object): """An object grouping the multiple outputs of a ParDo or FlatMap transform.""" def __init__( self, pipeline: 'Pipeline', transform: 'ParDo', tags: Sequence[str], main_tag: Optional[str], allow_unknown_tags: Optional[bool] = None, ): self._pipeline = pipeline self._tags = tags self._main_tag = main_tag self._transform = transform self._tagged_output_types = ( transform.get_type_hints().tagged_output_types() if transform else {}) self._allow_unknown_tags = ( not tags if allow_unknown_tags is None else allow_unknown_tags) # The ApplyPTransform instance for the application of the multi FlatMap # generating this value. The field gets initialized when a transform # gets applied. self.producer: Optional[AppliedPTransform] = None # Dictionary of PCollections already associated with tags. self._pcolls: dict[Optional[str], PCollection] = {} def __str__(self): return '<%s>' % self._str_internal() def __repr__(self): return '<%s at %s>' % (self._str_internal(), hex(id(self))) def _str_internal(self): return '%s main_tag=%s tags=%s transform=%s' % ( self.__class__.__name__, self._main_tag, self._tags, self._transform) def __iter__(self) -> Iterator[PCollection]: """Iterates over tags returning for each call a (tag, pcollection) pair.""" if self._main_tag is not None: yield self[self._main_tag] for tag in self._tags: yield self[tag] def __getattr__(self, tag: str) -> PCollection: # Special methods which may be accessed before the object is # fully constructed (e.g. in unpickling). if tag[:2] == tag[-2:] == '__': return object.__getattr__(self, tag) # type: ignore return self[tag] def __getitem__(self, tag: Union[int, str, None]) -> PCollection: # Accept int tags so that we can look at Partition tags with the # same ints that we used in the partition function. # TODO(gildea): Consider requiring string-based tags everywhere. # This will require a partition function that does not return ints. if isinstance(tag, int): tag = str(tag) if tag == self._main_tag: tag = None elif self._tags and tag not in self._tags and not self._allow_unknown_tags: raise ValueError( "Tag '%s' is neither the main tag '%s' " "nor any of the tags %s" % (tag, self._main_tag, self._tags)) # Check if we accessed this tag before. if tag in self._pcolls: return self._pcolls[tag] assert self.producer is not None if tag is not None: self._transform.output_tags.add(tag) is_bounded = all(i.is_bounded for i in self.producer.main_inputs.values()) pcoll = PCollection( self._pipeline, tag=tag, element_type=self._tagged_output_types.get(tag, typehints.Any), is_bounded=is_bounded) # Transfer the producer from the DoOutputsTuple to the resulting # PCollection. pcoll.producer = self.producer.parts[0] # Add this as an output to both the inner ParDo and the outer _MultiParDo # PTransforms. if tag not in self.producer.parts[0].outputs: self.producer.parts[0].add_output(pcoll, tag) self.producer.add_output(pcoll, tag) else: # Main output is output of inner ParDo. pval = self.producer.parts[0].outputs[None] assert isinstance(pval, PCollection), ("DoOutputsTuple should follow a ParDo.") pcoll = pval self._pcolls[tag] = pcoll return pcoll TagType = TypeVar('TagType', bound=str) ValueType = TypeVar('ValueType') class TaggedOutput(Generic[TagType, ValueType]): """An object representing a tagged value. ParDo, Map, and FlatMap transforms can emit values on multiple outputs which are distinguished by string tags. The DoFn will return plain values if it wants to emit on the main output and TaggedOutput objects if it wants to emit a value on a specific tagged output. """ def __init__(self, tag: TagType, value: ValueType) -> None: if not isinstance(tag, str): raise TypeError( 'Attempting to create a TaggedOutput with non-string tag %s' % (tag, )) self.tag = tag self.value = value class AsSideInput(object): """Marker specifying that a PCollection will be used as a side input. When a PCollection is supplied as a side input to a PTransform, it is necessary to indicate how the PCollection should be made available as a PTransform side argument (e.g. in the form of an iterable, mapping, or single value). This class is the superclass of all the various options, and should not be instantiated directly. (See instead AsSingleton, AsIter, etc.) """ def __init__(self, pcoll: PCollection) -> None: from apache_beam.transforms import sideinputs self.pvalue = pcoll self._window_mapping_fn = sideinputs.default_window_mapping_fn( pcoll.windowing.windowfn) def _view_options(self): """Internal options corresponding to specific view. Intended for internal use by runner implementations. Returns: Tuple of options for the given view. """ return { 'window_mapping_fn': self._window_mapping_fn, 'coder': self._windowed_coder(), } @property def element_type(self): return typehints.Any def _windowed_coder(self): return coders.WindowedValueCoder( coders.registry.get_coder( self.pvalue.element_type or self.element_type), self.pvalue.windowing.windowfn.get_window_coder()) # TODO(robertwb): Get rid of _from_runtime_iterable and _view_options # in favor of _side_input_data(). def _side_input_data(self) -> 'SideInputData': view_options = self._view_options() from_runtime_iterable = type(self)._from_runtime_iterable return SideInputData( common_urns.side_inputs.ITERABLE.urn, self._window_mapping_fn, lambda iterable: from_runtime_iterable(iterable, view_options)) def to_runner_api( self, context: 'PipelineContext') -> beam_runner_api_pb2.SideInput: return self._side_input_data().to_runner_api(context) @staticmethod def from_runner_api( proto: beam_runner_api_pb2.SideInput, context: 'PipelineContext') -> '_UnpickledSideInput': return _UnpickledSideInput(SideInputData.from_runner_api(proto, context)) @staticmethod def _from_runtime_iterable(it, options): raise NotImplementedError def requires_keyed_input(self): return False class _UnpickledSideInput(AsSideInput): def __init__(self, side_input_data: 'SideInputData') -> None: self._data = side_input_data self._window_mapping_fn = side_input_data.window_mapping_fn @staticmethod def _from_runtime_iterable(it, options): access_pattern = options['data'].access_pattern if access_pattern == common_urns.side_inputs.ITERABLE.urn: raw_view = it elif access_pattern == common_urns.side_inputs.MULTIMAP.urn: raw_view = collections.defaultdict(list) for k, v in it: raw_view[k].append(v) else: raise ValueError('Unknown access_pattern: %s' % access_pattern) return options['data'].view_fn(raw_view) def _view_options(self): return { 'data': self._data, # For non-fn-api runners. 'window_mapping_fn': self._data.window_mapping_fn, 'coder': self._windowed_coder(), } def _side_input_data(self): return self._data class SideInputData(object): """All of the data about a side input except for the bound PCollection.""" def __init__( self, access_pattern: str, window_mapping_fn: 'sideinputs.WindowMappingFn', view_fn): self.access_pattern = access_pattern self.window_mapping_fn = window_mapping_fn self.view_fn = view_fn def to_runner_api( self, context: 'PipelineContext') -> beam_runner_api_pb2.SideInput: return beam_runner_api_pb2.SideInput( access_pattern=beam_runner_api_pb2.FunctionSpec( urn=self.access_pattern), view_fn=beam_runner_api_pb2.FunctionSpec( urn=python_urns.PICKLED_VIEWFN, payload=pickler.dumps(self.view_fn)), window_mapping_fn=beam_runner_api_pb2.FunctionSpec( urn=python_urns.PICKLED_WINDOW_MAPPING_FN, payload=pickler.dumps(self.window_mapping_fn))) @staticmethod def from_runner_api( proto: beam_runner_api_pb2.SideInput, unused_context: 'PipelineContext') -> 'SideInputData': assert proto.view_fn.urn == python_urns.PICKLED_VIEWFN assert ( proto.window_mapping_fn.urn == python_urns.PICKLED_WINDOW_MAPPING_FN) return SideInputData( proto.access_pattern.urn, pickler.loads(proto.window_mapping_fn.payload), pickler.loads(proto.view_fn.payload)) class AsSingleton(AsSideInput): """Marker specifying that an entire PCollection is to be used as a side input. When a PCollection is supplied as a side input to a PTransform, it is necessary to indicate whether the entire PCollection should be made available as a PTransform side argument (in the form of an iterable), or whether just one value should be pulled from the PCollection and supplied as the side argument (as an ordinary value). Wrapping a PCollection side input argument to a PTransform in this container (e.g., data.apply('label', MyPTransform(), AsSingleton(my_side_input) ) selects the latter behavior. The input PCollection must contain exactly one value per window, unless a default is given, in which case it may be empty. """ _NO_DEFAULT = object() def __init__( self, pcoll: PCollection, default_value: Any = _NO_DEFAULT) -> None: super().__init__(pcoll) self.default_value = default_value def __repr__(self): return 'AsSingleton(%s)' % self.pvalue def _view_options(self): base = super()._view_options() if self.default_value != AsSingleton._NO_DEFAULT: return dict(base, default=self.default_value) return base @staticmethod def _from_runtime_iterable(it, options): head = list(itertools.islice(it, 2)) if not head: return options.get('default', EmptySideInput()) elif len(head) == 1: return head[0] raise ValueError( 'PCollection of size %d with more than one element accessed as a ' 'singleton view. First two elements encountered are "%s", "%s".' % (len(head), str(head[0]), str(head[1]))) @property def element_type(self): return self.pvalue.element_type class AsIter(AsSideInput): """Marker specifying that an entire PCollection is to be used as a side input. When a PCollection is supplied as a side input to a PTransform, it is necessary to indicate whether the entire PCollection should be made available as a PTransform side argument (in the form of an iterable), or whether just one value should be pulled from the PCollection and supplied as the side argument (as an ordinary value). Wrapping a PCollection side input argument to a PTransform in this container (e.g., data.apply('label', MyPTransform(), AsIter(my_side_input) ) selects the former behavor. """ def __repr__(self): return 'AsIter(%s)' % self.pvalue @staticmethod def _from_runtime_iterable(it, options): return it def _side_input_data(self) -> SideInputData: return SideInputData( common_urns.side_inputs.ITERABLE.urn, self._window_mapping_fn, lambda iterable: iterable) @property def element_type(self): return typehints.Iterable[self.pvalue.element_type] class AsList(AsSideInput): """Marker specifying that an entire PCollection is to be used as a side input. Intended for use in side-argument specification---the same places where AsSingleton and AsIter are used, but forces materialization of this PCollection as a list. Args: pcoll: Input pcollection. Returns: An AsList-wrapper around a PCollection whose one element is a list containing all elements in pcoll. """ @staticmethod def _from_runtime_iterable(it, options): return list(it) def _side_input_data(self) -> SideInputData: return SideInputData( common_urns.side_inputs.ITERABLE.urn, self._window_mapping_fn, list) class AsDict(AsSideInput): """Marker specifying a PCollection to be used as an indexable side input. Intended for use in side-argument specification---the same places where AsSingleton and AsIter are used, but returns an interface that allows key lookup. Args: pcoll: Input pcollection. All elements should be key-value pairs (i.e. 2-tuples) with unique keys. Returns: An AsDict-wrapper around a PCollection whose one element is a dict with entries for uniquely-keyed pairs in pcoll. """ @staticmethod def _from_runtime_iterable(it, options): return dict(it) def _side_input_data(self) -> SideInputData: return SideInputData( common_urns.side_inputs.ITERABLE.urn, self._window_mapping_fn, dict) class AsMultiMap(AsSideInput): """Marker specifying a PCollection to be used as an indexable side input. Similar to AsDict, but multiple values may be associated per key, and the keys are fetched lazily rather than all having to fit in memory. Intended for use in side-argument specification---the same places where AsSingleton and AsIter are used, but returns an interface that allows key lookup. """ @staticmethod def _from_runtime_iterable(it, options): # Legacy implementation. result = collections.defaultdict(list) for k, v in it: result[k].append(v) return result def _side_input_data(self) -> SideInputData: return SideInputData( common_urns.side_inputs.MULTIMAP.urn, self._window_mapping_fn, lambda x: x) def requires_keyed_input(self): return True class EmptySideInput(object): """Value indicating when a singleton side input was empty. If a PCollection was furnished as a singleton side input to a PTransform, and that PCollection was empty, then this value is supplied to the DoFn in the place where a value from a non-empty PCollection would have gone. This alerts the DoFn that the side input PCollection was empty. Users may want to check whether side input values are EmptySideInput, but they will very likely never want to create new instances of this class themselves. """ pass class Row(object): """A dynamic schema'd row object. This objects attributes are initialized from the keywords passed into its constructor, e.g. Row(x=3, y=4) will create a Row with two attributes x and y. More importantly, when a Row object is returned from a `Map`, `FlatMap`, or `DoFn` type inference is able to deduce the schema of the resulting PCollection, e.g. pc | beam.Map(lambda x: Row(x=x, y=0.5 * x)) when applied to a PCollection of ints will produce a PCollection with schema `(x=int, y=float)`. Note that in Beam 2.30.0 and later, Row objects are sensitive to field order. So `Row(x=3, y=4)` is not considered equal to `Row(y=4, x=3)`. """ def __init__(self, **kwargs): self.__dict__.update(kwargs) def as_dict(self): return dict(self.__dict__) # For compatibility with named tuples. _asdict = as_dict def __iter__(self): for _, value in self.__dict__.items(): yield value def __repr__(self): return 'Row(%s)' % ', '.join('%s=%r' % kv for kv in self.__dict__.items()) def __hash__(self): return hash(self.__dict__.items()) def __eq__(self, other): if type(self) == type(other): other_dict = other.__dict__ elif type(other) == type(NamedTuple): other_dict = other._asdict() else: return False return ( len(self.__dict__) == len(other_dict) and all(s == o for s, o in zip(self.__dict__.items(), other_dict.items()))) def __reduce__(self): return _make_Row, tuple(self.__dict__.items()) def _make_Row(*items): return Row(**dict(items)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/pvalue_test.py0000644000175100017510000000441715206113032020570 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the PValue and PCollection classes.""" # pytype: skip-file import unittest from apache_beam.pvalue import AsSingleton from apache_beam.pvalue import PValue from apache_beam.pvalue import Row from apache_beam.pvalue import TaggedOutput from apache_beam.testing.test_pipeline import TestPipeline class PValueTest(unittest.TestCase): def test_pvalue_expected_arguments(self): pipeline = TestPipeline() value = PValue(pipeline) self.assertEqual(pipeline, value.pipeline) def test_assingleton_multi_element(self): with self.assertRaisesRegex( ValueError, 'PCollection of size 2 with more than one element accessed as a ' 'singleton view. First two elements encountered are \"1\", \"2\".'): AsSingleton._from_runtime_iterable([1, 2], {}) class TaggedValueTest(unittest.TestCase): def test_passed_tuple_as_tag(self): with self.assertRaisesRegex( TypeError, r'Attempting to create a TaggedOutput with non-string tag \(1, 2, 3\)'): TaggedOutput((1, 2, 3), 'value') class RowTest(unittest.TestCase): def test_row_eq(self): row = Row(a=1, b=2) same = Row(a=1, b=2) self.assertEqual(row, same) def test_trailing_column_row_neq(self): row = Row(a=1, b=2) trail = Row(a=1, b=2, c=3) self.assertNotEqual(row, trail) def test_row_comparison_respects_element_order(self): row = Row(a=1, b=2) different = Row(b=2, a=1) self.assertNotEqual(row, different) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/py.typed0000644000175100017510000000153115206113032017354 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # Marker file for PEP 561. # The apache-beam package uses inline types. ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5644271 apache_beam-2.74.0/apache_beam/runners/0000755000175100017510000000000015206113160017353 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/__init__.py0000644000175100017510000000261315206113032021464 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Runner objects execute a Pipeline. This package defines runners, which are used to execute a pipeline. """ from apache_beam.runners.direct.direct_runner import DirectRunner from apache_beam.runners.direct.test_direct_runner import TestDirectRunner from apache_beam.runners.runner import PipelineRunner from apache_beam.runners.runner import PipelineState from apache_beam.runners.runner import create_runner # isort: off # initialize these last to avoid a circular dependency from apache_beam.runners.dataflow.dataflow_runner import DataflowRunner from apache_beam.runners.dataflow.test_dataflow_runner import TestDataflowRunner ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/common.pxd0000644000175100017510000001462315206113032021364 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # cimport cython from apache_beam.utils.windowed_value cimport WindowedValue from apache_beam.utils.windowed_value cimport WindowedBatch from apache_beam.transforms.cy_dataflow_distribution_counter cimport DataflowDistributionCounter from libc.stdint cimport int64_t cdef type TaggedOutput, TimestampedValue cdef class Receiver(object): cpdef receive(self, WindowedValue windowed_value) cpdef receive_batch(self, WindowedBatch windowed_batch) cpdef flush(self) cdef class MethodWrapper(object): cdef public object args cdef public object defaults cdef public object method_value cdef str method_name cdef bint has_userstate_arguments cdef object state_args_to_replace cdef object timer_args_to_replace cdef object timestamp_arg_name cdef object window_arg_name cdef object key_arg_name cdef object restriction_provider cdef object restriction_provider_arg_name cdef public object watermark_estimator_provider cdef object watermark_estimator_provider_arg_name cdef object dynamic_timer_tag_arg_name cdef bint unbounded_per_element cdef class DoFnSignature(object): cdef public MethodWrapper process_method cdef public MethodWrapper process_batch_method cdef public MethodWrapper start_bundle_method cdef public MethodWrapper finish_bundle_method cdef public MethodWrapper setup_lifecycle_method cdef public MethodWrapper teardown_lifecycle_method cdef public MethodWrapper create_watermark_estimator_method cdef public MethodWrapper initial_restriction_method cdef public MethodWrapper create_tracker_method cdef public MethodWrapper split_method cdef public object batching_configuration cdef public object do_fn cdef public object timer_methods cdef bint _is_stateful_dofn cdef class DoFnInvoker(object): cdef public DoFnSignature signature cdef OutputHandler output_handler cdef object user_state_context cdef public object bundle_finalizer_param cdef object _setup_context_values cdef object _bundle_context_values cpdef invoke_process(self, WindowedValue windowed_value, restriction=*, watermark_estimator_state=*, additional_args=*, additional_kwargs=*) cpdef invoke_start_bundle(self) cpdef invoke_finish_bundle(self) cpdef invoke_split(self, element, restriction) cpdef invoke_initial_restriction(self, element) cpdef invoke_create_tracker(self, restriction) cdef class SimpleInvoker(DoFnInvoker): cdef object process_method cdef object process_batch_method cdef class PerWindowInvoker(DoFnInvoker): cdef list side_inputs cdef DoFnContext context cdef list args_for_process cdef dict kwargs_for_process cdef list placeholders_for_process cdef list args_for_process_batch cdef dict kwargs_for_process_batch cdef list placeholders_for_process_batch cdef bint has_windowed_inputs cdef bint recalculate_window_args cdef bint has_cached_window_args cdef bint has_cached_window_batch_args cdef object process_method cdef object process_batch_method cdef bint is_splittable cdef object threadsafe_restriction_tracker cdef object threadsafe_watermark_estimator cdef WindowedValue current_windowed_value cdef object restriction cdef object watermark_estimator_state cdef object current_window_index cdef object stop_window_index cdef bint is_key_param_required cdef object splitting_lock cdef class DoFnRunner: cdef DoFnContext context cdef object step_name cdef list side_inputs cdef DoFnInvoker do_fn_invoker cdef public object bundle_finalizer_param cdef str transform_id cdef object execution_context cpdef process(self, WindowedValue windowed_value) cdef class OutputHandler(object): @cython.locals(windowed_value=WindowedValue, output_element_count=int64_t) cpdef handle_process_outputs(self, WindowedValue element, results, watermark_estimator=*) @cython.locals(windowed_batch=WindowedBatch, output_element_count=int64_t) cpdef handle_process_batch_outputs(self, WindowedBatch input_batch, results, watermark_estimator=*) cdef class _OutputHandler(OutputHandler): cdef object window_fn cdef Receiver main_receivers cdef object tagged_receivers cdef DataflowDistributionCounter per_element_output_counter cdef object output_batch_converter cdef bint _process_batch_yields_elements cdef bint _process_yields_batches @cython.locals(windowed_value=WindowedValue, windowed_batch=WindowedBatch, output_element_count=int64_t) cpdef handle_process_outputs(self, WindowedValue element, results, watermark_estimator=*) @cython.locals(windowed_value=WindowedValue, windowed_batch=WindowedBatch, output_element_count=int64_t) cpdef handle_process_batch_outputs(self, WindowedBatch input_batch, results, watermark_estimator=*) @cython.locals(windowed_value=WindowedValue) cdef inline WindowedValue _maybe_propagate_windowing_info(self, WindowedValue input_element, result) cdef inline tuple _handle_tagged_output(self, result) cdef inline _write_value_to_tag(self, tag, WindowedValue windowed_value, watermark_estimator) cdef inline _write_batch_to_tag(self, tag, WindowedBatch windowed_batch, watermark_estimator) cdef inline _verify_batch_output(self, result) cdef class DoFnContext(object): cdef object label cdef object state cdef WindowedValue windowed_value cpdef set_element(self, WindowedValue windowed_value) cdef class _ReceiverAdapter(Receiver): cdef object underlying ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/common.py0000644000175100017510000022251115206113032021216 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Worker operations executor. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file # ruff: noqa: UP006 import logging import sys import threading import traceback from enum import Enum from typing import TYPE_CHECKING from typing import Any from typing import Dict from typing import Iterable from typing import List from typing import Mapping from typing import Optional from typing import Tuple from apache_beam.coders import TupleCoder from apache_beam.coders import coders from apache_beam.internal import util from apache_beam.options.value_provider import RuntimeValueProvider from apache_beam.pvalue import TaggedOutput from apache_beam.runners.sdf_utils import NoOpWatermarkEstimatorProvider from apache_beam.runners.sdf_utils import RestrictionTrackerView from apache_beam.runners.sdf_utils import SplitResultPrimary from apache_beam.runners.sdf_utils import SplitResultResidual from apache_beam.runners.sdf_utils import ThreadsafeRestrictionTracker from apache_beam.runners.sdf_utils import ThreadsafeWatermarkEstimator from apache_beam.transforms import DoFn from apache_beam.transforms import core from apache_beam.transforms import userstate from apache_beam.transforms.core import RestrictionProvider from apache_beam.transforms.core import WatermarkEstimatorProvider from apache_beam.transforms.window import GlobalWindow from apache_beam.transforms.window import GlobalWindows from apache_beam.transforms.window import TimestampedValue from apache_beam.transforms.window import WindowFn from apache_beam.typehints.batch import BatchConverter from apache_beam.utils.counters import Counter from apache_beam.utils.counters import CounterName from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.windowed_value import HomogeneousWindowedBatch from apache_beam.utils.windowed_value import WindowedBatch from apache_beam.utils.windowed_value import WindowedValue if TYPE_CHECKING: from apache_beam.io.iobase import RestrictionProgress from apache_beam.iobase import RestrictionTracker from apache_beam.iobase import WatermarkEstimator from apache_beam.runners.worker.bundle_processor import ExecutionContext from apache_beam.transforms import sideinputs from apache_beam.transforms.core import TimerSpec IMPULSE_VALUE_CODER_IMPL = coders.WindowedValueCoder( coders.BytesCoder(), coders.GlobalWindowCoder()).get_impl() ENCODED_IMPULSE_VALUE = IMPULSE_VALUE_CODER_IMPL.encode_nested( GlobalWindows.windowed_value(b'')) _LOGGER = logging.getLogger(__name__) class NameContext(object): """Holds the name information for a step.""" def __init__(self, step_name, transform_id=None): # type: (str, Optional[str]) -> None """Creates a new step NameContext. Args: step_name: The name of the step. """ self.step_name = step_name self.transform_id = transform_id def __eq__(self, other): return self.step_name == other.step_name def __repr__(self): return 'NameContext(%s)' % self.__dict__ def __hash__(self): return hash(self.step_name) def metrics_name(self): """Returns the step name used for metrics reporting.""" return self.step_name def logging_name(self): """Returns the step name used for logging.""" return self.step_name class Receiver(object): """For internal use only; no backwards-compatibility guarantees. An object that consumes a WindowedValue. This class can be efficiently used to pass values between the sdk and worker harnesses. """ def receive(self, windowed_value): # type: (WindowedValue) -> None raise NotImplementedError def receive_batch(self, windowed_batch): # type: (WindowedBatch) -> None raise NotImplementedError def flush(self): raise NotImplementedError class MethodWrapper(object): """For internal use only; no backwards-compatibility guarantees. Represents a method that can be invoked by `DoFnInvoker`.""" def __init__(self, obj_to_invoke, method_name): """ Initiates a ``MethodWrapper``. Args: obj_to_invoke: the object that contains the method. Has to either be a `DoFn` object or a `RestrictionProvider` object. method_name: name of the method as a string. """ if not isinstance(obj_to_invoke, (DoFn, RestrictionProvider, WatermarkEstimatorProvider)): raise ValueError( '\'obj_to_invoke\' has to be either a \'DoFn\' or ' 'a \'RestrictionProvider\'. Received %r instead.' % obj_to_invoke) self.args, self.defaults = core.get_function_arguments(obj_to_invoke, method_name) # TODO(BEAM-5878) support kwonlyargs on Python 3. self.method_value = getattr(obj_to_invoke, method_name) self.method_name = method_name self.has_userstate_arguments = False self.state_args_to_replace = {} # type: Dict[str, core.StateSpec] self.timer_args_to_replace = {} # type: Dict[str, core.TimerSpec] self.timestamp_arg_name = None # type: Optional[str] self.window_arg_name = None # type: Optional[str] self.key_arg_name = None # type: Optional[str] self.restriction_provider = None self.restriction_provider_arg_name = None self.watermark_estimator_provider = None self.watermark_estimator_provider_arg_name = None self.dynamic_timer_tag_arg_name = None if hasattr(self.method_value, 'unbounded_per_element'): self.unbounded_per_element = True else: self.unbounded_per_element = False for kw, v in zip(self.args[-len(self.defaults):], self.defaults): if isinstance(v, core.DoFn.StateParam): self.state_args_to_replace[kw] = v.state_spec self.has_userstate_arguments = True elif isinstance(v, core.DoFn.TimerParam): self.timer_args_to_replace[kw] = v.timer_spec self.has_userstate_arguments = True elif core.DoFn.TimestampParam == v: self.timestamp_arg_name = kw elif core.DoFn.WindowParam == v: self.window_arg_name = kw elif core.DoFn.WindowedValueParam == v: self.window_arg_name = kw elif core.DoFn.KeyParam == v: self.key_arg_name = kw elif isinstance(v, core.DoFn.RestrictionParam): self.restriction_provider = v.restriction_provider or obj_to_invoke self.restriction_provider_arg_name = kw elif isinstance(v, core.DoFn.WatermarkEstimatorParam): self.watermark_estimator_provider = ( v.watermark_estimator_provider or obj_to_invoke) self.watermark_estimator_provider_arg_name = kw elif core.DoFn.DynamicTimerTagParam == v: self.dynamic_timer_tag_arg_name = kw # Create NoOpWatermarkEstimatorProvider if there is no # WatermarkEstimatorParam provided. if self.watermark_estimator_provider is None: self.watermark_estimator_provider = NoOpWatermarkEstimatorProvider() def invoke_timer_callback( self, user_state_context, key, window, timestamp, pane_info, dynamic_timer_tag): # TODO(ccy): support side inputs. kwargs = {} if self.has_userstate_arguments: for kw, state_spec in self.state_args_to_replace.items(): kwargs[kw] = user_state_context.get_state(state_spec, key, window) for kw, timer_spec in self.timer_args_to_replace.items(): kwargs[kw] = user_state_context.get_timer( timer_spec, key, window, timestamp, pane_info) if self.timestamp_arg_name: kwargs[self.timestamp_arg_name] = Timestamp.of(timestamp) if self.window_arg_name: kwargs[self.window_arg_name] = window if self.key_arg_name: kwargs[self.key_arg_name] = key if self.dynamic_timer_tag_arg_name: kwargs[self.dynamic_timer_tag_arg_name] = dynamic_timer_tag if kwargs: return self.method_value(**kwargs) else: return self.method_value() class BatchingPreference(Enum): DO_NOT_CARE = 1 # This operation can operate on batches or element-at-a-time # TODO: Should we also store batching parameters here? (time/size preferences) BATCH_REQUIRED = 2 # This operation can only operate on batches BATCH_FORBIDDEN = 3 # This operation can only work element-at-a-time # Other possibilities: BATCH_PREFERRED (with min batch size specified) @property def supports_batches(self) -> bool: return self in (self.BATCH_REQUIRED, self.DO_NOT_CARE) @property def supports_elements(self) -> bool: return self in (self.BATCH_FORBIDDEN, self.DO_NOT_CARE) @property def requires_batches(self) -> bool: return self == self.BATCH_REQUIRED class DoFnSignature(object): """Represents the signature of a given ``DoFn`` object. Signature of a ``DoFn`` provides a view of the properties of a given ``DoFn``. Among other things, this will give an extensible way for for (1) accessing the structure of the ``DoFn`` including methods and method parameters (2) identifying features that a given ``DoFn`` support, for example, whether a given ``DoFn`` is a Splittable ``DoFn`` ( https://s.apache.org/splittable-do-fn) (3) validating a ``DoFn`` based on the feature set offered by it. """ def __init__(self, do_fn): # type: (core.DoFn) -> None # We add a property here for all methods defined by Beam DoFn features. assert isinstance(do_fn, core.DoFn) self.do_fn = do_fn self.process_method = MethodWrapper(do_fn, 'process') self.process_batch_method = MethodWrapper(do_fn, 'process_batch') self.start_bundle_method = MethodWrapper(do_fn, 'start_bundle') self.finish_bundle_method = MethodWrapper(do_fn, 'finish_bundle') self.setup_lifecycle_method = MethodWrapper(do_fn, 'setup') self.teardown_lifecycle_method = MethodWrapper(do_fn, 'teardown') restriction_provider = self.get_restriction_provider() watermark_estimator_provider = self.get_watermark_estimator_provider() self.create_watermark_estimator_method = ( MethodWrapper( watermark_estimator_provider, 'create_watermark_estimator')) self.initial_restriction_method = ( MethodWrapper(restriction_provider, 'initial_restriction') if restriction_provider else None) self.create_tracker_method = ( MethodWrapper(restriction_provider, 'create_tracker') if restriction_provider else None) self.split_method = ( MethodWrapper(restriction_provider, 'split') if restriction_provider else None) self._validate() # Handle stateful DoFns. self._is_stateful_dofn = userstate.is_stateful_dofn(do_fn) self.timer_methods = {} # type: Dict[TimerSpec, MethodWrapper] if self._is_stateful_dofn: # Populate timer firing methods, keyed by TimerSpec. _, all_timer_specs = userstate.get_dofn_specs(do_fn) for timer_spec in all_timer_specs: method = timer_spec._attached_callback self.timer_methods[timer_spec] = MethodWrapper(do_fn, method.__name__) def get_restriction_provider(self): # type: () -> RestrictionProvider return self.process_method.restriction_provider def get_watermark_estimator_provider(self): # type: () -> WatermarkEstimatorProvider return self.process_method.watermark_estimator_provider def is_unbounded_per_element(self): return self.process_method.unbounded_per_element def _validate(self): # type: () -> None self._validate_process() self._validate_process_batch() self._validate_bundle_method(self.start_bundle_method) self._validate_bundle_method(self.finish_bundle_method) self._validate_stateful_dofn() def _check_duplicate_dofn_params(self, method: MethodWrapper): param_ids = [ d.param_id for d in method.defaults if isinstance(d, core._DoFnParam) ] if len(param_ids) != len(set(param_ids)): raise ValueError( 'DoFn %r has duplicate %s method parameters: %s.' % (self.do_fn, method.method_name, param_ids)) def _validate_process(self): # type: () -> None """Validate that none of the DoFnParameters are repeated in the function """ self._check_duplicate_dofn_params(self.process_method) def _validate_process_batch(self): # type: () -> None self._check_duplicate_dofn_params(self.process_batch_method) for d in self.process_batch_method.defaults: if not isinstance(d, core._DoFnParam): continue # Helpful errors for params which will be supported in the future if d == (core.DoFn.ElementParam): # We currently assume we can just get the typehint from the first # parameter. ElementParam breaks this assumption raise NotImplementedError( f"DoFn {self.do_fn!r} uses unsupported DoFn param ElementParam.") if d in (core.DoFn.KeyParam, core.DoFn.StateParam, core.DoFn.TimerParam): raise NotImplementedError( f"DoFn {self.do_fn!r} has unsupported per-key DoFn param {d}. " "Per-key DoFn params are not yet supported for process_batch " "(https://github.com/apache/beam/issues/21653).") # Fallback to catch anything not explicitly supported if not d in (core.DoFn.WindowParam, core.DoFn.TimestampParam, core.DoFn.PaneInfoParam): raise ValueError( f"DoFn {self.do_fn!r} has unsupported process_batch " f"method parameter {d}") def _validate_bundle_method(self, method_wrapper): """Validate that none of the DoFnParameters are used in the function """ for param in core.DoFn.DoFnProcessParams: if param in method_wrapper.defaults: raise ValueError( 'DoFn.process() method-only parameter %s cannot be used in %s.' % (param, method_wrapper)) def _validate_stateful_dofn(self): # type: () -> None userstate.validate_stateful_dofn(self.do_fn) def is_splittable_dofn(self): # type: () -> bool return self.get_restriction_provider() is not None def get_restriction_coder(self): # type: () -> Optional[TupleCoder] """Get coder for a restriction when processing an SDF. """ if self.is_splittable_dofn(): return TupleCoder([ (self.get_restriction_provider().restriction_coder()), (self.get_watermark_estimator_provider().estimator_state_coder()) ]) else: return None def is_stateful_dofn(self): # type: () -> bool return self._is_stateful_dofn def has_timers(self): # type: () -> bool _, all_timer_specs = userstate.get_dofn_specs(self.do_fn) return bool(all_timer_specs) def has_bundle_finalization(self): for sig in (self.start_bundle_method, self.process_method, self.finish_bundle_method): for d in sig.defaults: try: if d == DoFn.BundleFinalizerParam: return True except Exception: # pylint: disable=broad-except # Default value might be incomparable. pass return False def get_bundle_contexts(self): seen = set() for sig in (self.setup_lifecycle_method, self.start_bundle_method, self.process_method, self.process_batch_method, self.finish_bundle_method, self.teardown_lifecycle_method): for d in sig.defaults: try: if isinstance(d, DoFn.BundleContextParam): if d not in seen: seen.add(d) yield d except Exception: # pylint: disable=broad-except # Default value might be incomparable. pass def get_setup_contexts(self): seen = set() for sig in (self.setup_lifecycle_method, self.start_bundle_method, self.process_method, self.process_batch_method, self.finish_bundle_method, self.teardown_lifecycle_method): for d in sig.defaults: try: if isinstance(d, DoFn.SetupContextParam): if d not in seen: seen.add(d) yield d except Exception: # pylint: disable=broad-except # Default value might be incomparable. pass class DoFnInvoker(object): """An abstraction that can be used to execute DoFn methods. A DoFnInvoker describes a particular way for invoking methods of a DoFn represented by a given DoFnSignature.""" def __init__( self, output_handler, # type: _OutputHandler signature # type: DoFnSignature ): # type: (...) -> None """ Initializes `DoFnInvoker` :param output_handler: an OutputHandler for receiving elements produced by invoking functions of the DoFn. :param signature: a DoFnSignature for the DoFn being invoked """ self.output_handler = output_handler self.signature = signature self.user_state_context = None # type: Optional[userstate.UserStateContext] self.bundle_finalizer_param = None # type: Optional[core._BundleFinalizerParam] @staticmethod def create_invoker( signature, # type: DoFnSignature output_handler, # type: OutputHandler context=None, # type: Optional[DoFnContext] side_inputs=None, # type: Optional[List[sideinputs.SideInputMap]] input_args=None, input_kwargs=None, process_invocation=True, user_state_context=None, # type: Optional[userstate.UserStateContext] bundle_finalizer_param=None # type: Optional[core._BundleFinalizerParam] ): # type: (...) -> DoFnInvoker """ Creates a new DoFnInvoker based on given arguments. Args: output_handler: an OutputHandler for receiving elements produced by invoking functions of the DoFn. signature: a DoFnSignature for the DoFn being invoked. context: Context to be used when invoking the DoFn (deprecated). side_inputs: side inputs to be used when invoking th process method. input_args: arguments to be used when invoking the process method. Some of the arguments given here might be placeholders (for example for side inputs) that get filled before invoking the process method. input_kwargs: keyword arguments to be used when invoking the process method. Some of the keyword arguments given here might be placeholders (for example for side inputs) that get filled before invoking the process method. process_invocation: If True, this function may return an invoker that performs extra optimizations for invoking process() method efficiently. user_state_context: The UserStateContext instance for the current Stateful DoFn. bundle_finalizer_param: The param that passed to a process method, which allows a callback to be registered. """ side_inputs = side_inputs or [] use_per_window_invoker = process_invocation and ( side_inputs or input_args or input_kwargs or signature.process_method.defaults or signature.process_batch_method.defaults or signature.is_stateful_dofn()) if not use_per_window_invoker: return SimpleInvoker(output_handler, signature) else: if context is None: raise TypeError("Must provide context when not using SimpleInvoker") return PerWindowInvoker( output_handler, signature, context, side_inputs, input_args, input_kwargs, user_state_context, bundle_finalizer_param) def invoke_process( self, windowed_value, # type: WindowedValue restriction=None, watermark_estimator_state=None, additional_args=None, additional_kwargs=None): # type: (...) -> Iterable[SplitResultResidual] """Invokes the DoFn.process() function. Args: windowed_value: a WindowedValue object that gives the element for which process() method should be invoked along with the window the element belongs to. restriction: The restriction to use when executing this splittable DoFn. Should only be specified for splittable DoFns. watermark_estimator_state: The watermark estimator state to use when executing this splittable DoFn. Should only be specified for splittable DoFns. additional_args: additional arguments to be passed to the current `DoFn.process()` invocation, usually as side inputs. additional_kwargs: additional keyword arguments to be passed to the current `DoFn.process()` invocation. """ raise NotImplementedError def invoke_process_batch( self, windowed_batch, # type: WindowedBatch additional_args=None, additional_kwargs=None): # type: (...) -> None """Invokes the DoFn.process() function. Args: windowed_batch: a WindowedBatch object that gives a batch of elements for which process_batch() method should be invoked, along with the window each element belongs to. additional_args: additional arguments to be passed to the current `DoFn.process()` invocation, usually as side inputs. additional_kwargs: additional keyword arguments to be passed to the current `DoFn.process()` invocation. """ raise NotImplementedError def invoke_setup(self): # type: () -> None """Invokes the DoFn.setup() method """ self._setup_context_values = { c: c.create_and_enter() for c in self.signature.get_setup_contexts() } self.signature.setup_lifecycle_method.method_value() def invoke_start_bundle(self): # type: () -> None """Invokes the DoFn.start_bundle() method. """ self._bundle_context_values = { c: c.create_and_enter() for c in self.signature.get_bundle_contexts() } self.output_handler.start_bundle_outputs( self.signature.start_bundle_method.method_value()) def invoke_finish_bundle(self): # type: () -> None """Invokes the DoFn.finish_bundle() method. """ self.output_handler.finish_bundle_outputs( self.signature.finish_bundle_method.method_value()) for c in self._bundle_context_values.values(): c[0].__exit__(None, None, None) self._bundle_context_values = None def invoke_teardown(self): # type: () -> None """Invokes the DoFn.teardown() method """ self.signature.teardown_lifecycle_method.method_value() for c in self._setup_context_values.values(): c[0].__exit__(None, None, None) self._setup_context_values = None def invoke_user_timer( self, timer_spec, key, window, timestamp, pane_info, dynamic_timer_tag): # self.output_handler is Optional, but in practice it won't be None here self.output_handler.handle_process_outputs( WindowedValue(None, timestamp, (window, )), self.signature.timer_methods[timer_spec].invoke_timer_callback( self.user_state_context, key, window, timestamp, pane_info, dynamic_timer_tag)) def invoke_create_watermark_estimator(self, estimator_state): return self.signature.create_watermark_estimator_method.method_value( estimator_state) def invoke_split(self, element, restriction): return self.signature.split_method.method_value(element, restriction) def invoke_initial_restriction(self, element): return self.signature.initial_restriction_method.method_value(element) def invoke_create_tracker(self, restriction): return self.signature.create_tracker_method.method_value(restriction) class SimpleInvoker(DoFnInvoker): """An invoker that processes elements ignoring windowing information.""" def __init__( self, output_handler, # type: OutputHandler signature # type: DoFnSignature ): # type: (...) -> None super().__init__(output_handler, signature) self.process_method = signature.process_method.method_value self.process_batch_method = signature.process_batch_method.method_value def invoke_process( self, windowed_value, # type: WindowedValue restriction=None, watermark_estimator_state=None, additional_args=None, additional_kwargs=None): # type: (...) -> Iterable[SplitResultResidual] self.output_handler.handle_process_outputs( windowed_value, self.process_method(windowed_value.value)) return [] def invoke_process_batch( self, windowed_batch, # type: WindowedBatch restriction=None, watermark_estimator_state=None, additional_args=None, additional_kwargs=None): # type: (...) -> None self.output_handler.handle_process_batch_outputs( windowed_batch, self.process_batch_method(windowed_batch.values)) def _get_arg_placeholders( method: MethodWrapper, input_args: Optional[List[Any]], input_kwargs: Optional[Dict[str, any]]): input_args = input_args if input_args else [] input_kwargs = input_kwargs if input_kwargs else {} arg_names = method.args default_arg_values = method.defaults # Create placeholder for element parameter of DoFn.process() method. # Not to be confused with ArgumentPlaceHolder, which may be passed in # input_args and is a placeholder for side-inputs. class ArgPlaceholder(object): def __init__(self, placeholder): self.placeholder = placeholder if all(core.DoFn.ElementParam != arg for arg in default_arg_values): # TODO(https://github.com/apache/beam/issues/19631): Handle cases in which # len(arg_names) == len(default_arg_values). args_to_pick = len(arg_names) - len(default_arg_values) - 1 # Positional argument values for process(), with placeholders for special # values such as the element, timestamp, etc. args_with_placeholders = ([ArgPlaceholder(core.DoFn.ElementParam)] + input_args[:args_to_pick]) else: args_to_pick = len(arg_names) - len(default_arg_values) args_with_placeholders = input_args[:args_to_pick] # Fill the OtherPlaceholders for context, key, window or timestamp remaining_args_iter = iter(input_args[args_to_pick:]) for a, d in zip(arg_names[-len(default_arg_values):], default_arg_values): if core.DoFn.ElementParam == d: args_with_placeholders.append(ArgPlaceholder(d)) elif core.DoFn.KeyParam == d: args_with_placeholders.append(ArgPlaceholder(d)) elif core.DoFn.WindowParam == d: args_with_placeholders.append(ArgPlaceholder(d)) elif core.DoFn.WindowedValueParam == d: args_with_placeholders.append(ArgPlaceholder(d)) elif core.DoFn.TimestampParam == d: args_with_placeholders.append(ArgPlaceholder(d)) elif core.DoFn.PaneInfoParam == d: args_with_placeholders.append(ArgPlaceholder(d)) elif core.DoFn.SideInputParam == d: # If no more args are present then the value must be passed via kwarg try: args_with_placeholders.append(next(remaining_args_iter)) except StopIteration: if a not in input_kwargs: raise ValueError("Value for sideinput %s not provided" % a) elif isinstance(d, core.DoFn.StateParam): args_with_placeholders.append(ArgPlaceholder(d)) elif isinstance(d, core.DoFn.TimerParam): args_with_placeholders.append(ArgPlaceholder(d)) elif isinstance(d, type) and core.DoFn.BundleFinalizerParam == d: args_with_placeholders.append(ArgPlaceholder(d)) elif isinstance(d, core.DoFn.BundleContextParam): args_with_placeholders.append(ArgPlaceholder(d)) elif isinstance(d, core.DoFn.SetupContextParam): args_with_placeholders.append(ArgPlaceholder(d)) else: # If no more args are present then the value must be passed via kwarg try: args_with_placeholders.append(next(remaining_args_iter)) except StopIteration: pass args_with_placeholders.extend(list(remaining_args_iter)) # Stash the list of placeholder positions for performance placeholders = [(i, x.placeholder) for (i, x) in enumerate(args_with_placeholders) if isinstance(x, ArgPlaceholder)] return placeholders, args_with_placeholders, input_kwargs class PerWindowInvoker(DoFnInvoker): """An invoker that processes elements considering windowing information.""" def __init__( self, output_handler, # type: OutputHandler signature, # type: DoFnSignature context, # type: DoFnContext side_inputs, # type: Iterable[sideinputs.SideInputMap] input_args, input_kwargs, user_state_context, # type: Optional[userstate.UserStateContext] bundle_finalizer_param # type: Optional[core._BundleFinalizerParam] ): super().__init__(output_handler, signature) self.side_inputs = side_inputs self.context = context self.process_method = signature.process_method.method_value default_arg_values = signature.process_method.defaults self.has_windowed_inputs = ( not all(si.is_globally_windowed() for si in side_inputs) or any( core.DoFn.WindowParam == arg for arg in signature.process_method.defaults) or any( core.DoFn.WindowParam == arg for arg in signature.process_batch_method.defaults) or signature.is_stateful_dofn()) self.user_state_context = user_state_context self.is_splittable = signature.is_splittable_dofn() self.is_key_param_required = any( core.DoFn.KeyParam == arg for arg in default_arg_values) self.threadsafe_restriction_tracker = None # type: Optional[ThreadsafeRestrictionTracker] self.threadsafe_watermark_estimator = None # type: Optional[ThreadsafeWatermarkEstimator] self.current_windowed_value = None # type: Optional[WindowedValue] self.bundle_finalizer_param = bundle_finalizer_param if self.is_splittable: self.splitting_lock = threading.Lock() self.current_window_index = None self.stop_window_index = None # TODO(https://github.com/apache/beam/issues/28776): Remove caching after # fully rolling out. # If true, always recalculate window args. If false, has_cached_window_args # and has_cached_window_batch_args will be set to true if the corresponding # self.args_for_process,have been updated and should be reused directly. self.recalculate_window_args = ( self.has_windowed_inputs or 'disable_global_windowed_args_caching' in RuntimeValueProvider.experiments) self.has_cached_window_args = False self.has_cached_window_batch_args = False # Try to prepare all the arguments that can just be filled in # without any additional work. in the process function. # Also cache all the placeholders needed in the process function. input_args = list(input_args) ( self.placeholders_for_process, self.args_for_process, self.kwargs_for_process) = _get_arg_placeholders( signature.process_method, input_args, input_kwargs) self.process_batch_method = signature.process_batch_method.method_value ( self.placeholders_for_process_batch, self.args_for_process_batch, self.kwargs_for_process_batch) = _get_arg_placeholders( signature.process_batch_method, input_args, input_kwargs) def invoke_process( self, windowed_value, # type: WindowedValue restriction=None, watermark_estimator_state=None, additional_args=None, additional_kwargs=None): # type: (...) -> Iterable[SplitResultResidual] if not additional_args: additional_args = [] if not additional_kwargs: additional_kwargs = {} self.context.set_element(windowed_value) # Call for the process function for each window if has windowed side inputs # or if the process accesses the window parameter. We can just call it once # otherwise as none of the arguments are changing residuals = [] if self.is_splittable: if restriction is None: # This may be a SDF invoked as an ordinary DoFn on runners that don't # understand SDF. See, e.g. BEAM-11472. # In this case, processing the element is simply processing it against # the entire initial restriction. restriction = self.signature.initial_restriction_method.method_value( windowed_value.value) with self.splitting_lock: self.current_windowed_value = windowed_value self.restriction = restriction self.watermark_estimator_state = watermark_estimator_state try: if self.has_windowed_inputs and len(windowed_value.windows) > 1: for i, w in enumerate(windowed_value.windows): if not self._should_process_window_for_sdf( windowed_value, additional_kwargs, i): break residual = self._invoke_process_per_window( WindowedValue( windowed_value.value, windowed_value.timestamp, (w, )), additional_args, additional_kwargs) if residual: residuals.append(residual) else: if self._should_process_window_for_sdf(windowed_value, additional_kwargs): residual = self._invoke_process_per_window( windowed_value, additional_args, additional_kwargs) if residual: residuals.append(residual) finally: with self.splitting_lock: self.current_windowed_value = None self.restriction = None self.watermark_estimator_state = None self.current_window_index = None self.threadsafe_restriction_tracker = None self.threadsafe_watermark_estimator = None elif self.has_windowed_inputs and len(windowed_value.windows) != 1: for w in windowed_value.windows: self._invoke_process_per_window( WindowedValue( windowed_value.value, windowed_value.timestamp, (w, )), additional_args, additional_kwargs) else: self._invoke_process_per_window( windowed_value, additional_args, additional_kwargs) return residuals def invoke_process_batch( self, windowed_batch, # type: WindowedBatch additional_args=None, additional_kwargs=None): # type: (...) -> None if not additional_args: additional_args = [] if not additional_kwargs: additional_kwargs = {} assert isinstance(windowed_batch, HomogeneousWindowedBatch) if self.has_windowed_inputs and len(windowed_batch.windows) != 1: for w in windowed_batch.windows: self._invoke_process_batch_per_window( HomogeneousWindowedBatch.of( windowed_batch.values, windowed_batch.timestamp, (w, ), windowed_batch.pane_info), additional_args, additional_kwargs) else: self._invoke_process_batch_per_window( windowed_batch, additional_args, additional_kwargs) def _should_process_window_for_sdf( self, windowed_value, # type: WindowedValue additional_kwargs, window_index=None, # type: Optional[int] ): restriction_tracker = self.invoke_create_tracker(self.restriction) watermark_estimator = self.invoke_create_watermark_estimator( self.watermark_estimator_state) with self.splitting_lock: if window_index: self.current_window_index = window_index if window_index == 0: self.stop_window_index = len(windowed_value.windows) if window_index == self.stop_window_index: return False self.threadsafe_restriction_tracker = ThreadsafeRestrictionTracker( restriction_tracker) self.threadsafe_watermark_estimator = ( ThreadsafeWatermarkEstimator(watermark_estimator)) restriction_tracker_param = ( self.signature.process_method.restriction_provider_arg_name) if not restriction_tracker_param: raise ValueError( 'DoFn is splittable but DoFn does not have a ' 'RestrictionTrackerParam defined') additional_kwargs[restriction_tracker_param] = ( RestrictionTrackerView(self.threadsafe_restriction_tracker)) watermark_param = ( self.signature.process_method.watermark_estimator_provider_arg_name) # When the watermark_estimator is a NoOpWatermarkEstimator, the system # will not add watermark_param into the DoFn param list. if watermark_param is not None: additional_kwargs[watermark_param] = self.threadsafe_watermark_estimator return True def _invoke_process_per_window( self, windowed_value, # type: WindowedValue additional_args, additional_kwargs, ): # type: (...) -> Optional[SplitResultResidual] if self.has_cached_window_args: args_for_process, kwargs_for_process = ( self.args_for_process, self.kwargs_for_process) else: if self.has_windowed_inputs: assert len(windowed_value.windows) <= 1 window, = windowed_value.windows else: window = GlobalWindow() side_inputs = [si[window] for si in self.side_inputs] side_inputs.extend(additional_args) args_for_process, kwargs_for_process = util.insert_values_in_args( self.args_for_process, self.kwargs_for_process, side_inputs) if not self.recalculate_window_args: self.args_for_process, self.kwargs_for_process = ( args_for_process, kwargs_for_process) self.has_cached_window_args = True # Extract key in the case of a stateful DoFn. Note that in the case of a # stateful DoFn, we set during __init__ self.has_windowed_inputs to be # True. Therefore, windows will be exploded coming into this method, and # we can rely on the window variable being set above. if self.user_state_context or self.is_key_param_required: try: key, unused_value = windowed_value.value except (TypeError, ValueError): raise ValueError(( 'Input value to a stateful DoFn or KeyParam must be a KV tuple; ' 'instead, got \'%s\'.') % (windowed_value.value, )) for i, p in self.placeholders_for_process: if core.DoFn.ElementParam == p: args_for_process[i] = windowed_value.value elif core.DoFn.KeyParam == p: args_for_process[i] = key elif core.DoFn.WindowParam == p: args_for_process[i] = window elif core.DoFn.WindowedValueParam == p: args_for_process[i] = windowed_value elif core.DoFn.TimestampParam == p: args_for_process[i] = windowed_value.timestamp elif core.DoFn.PaneInfoParam == p: args_for_process[i] = windowed_value.pane_info elif isinstance(p, core.DoFn.StateParam): assert self.user_state_context is not None args_for_process[i] = ( self.user_state_context.get_state(p.state_spec, key, window)) elif isinstance(p, core.DoFn.TimerParam): assert self.user_state_context is not None args_for_process[i] = ( self.user_state_context.get_timer( p.timer_spec, key, window, windowed_value.timestamp, windowed_value.pane_info)) elif core.DoFn.BundleFinalizerParam == p: args_for_process[i] = self.bundle_finalizer_param elif isinstance(p, core.DoFn.BundleContextParam): args_for_process[i] = self._bundle_context_values[p][1] elif isinstance(p, core.DoFn.SetupContextParam): args_for_process[i] = self._setup_context_values[p][1] kwargs_for_process = kwargs_for_process or {} if additional_kwargs: kwargs_for_process.update(additional_kwargs) self.output_handler.handle_process_outputs( windowed_value, self.process_method(*args_for_process, **kwargs_for_process), self.threadsafe_watermark_estimator) if self.is_splittable: assert self.threadsafe_restriction_tracker is not None self.threadsafe_restriction_tracker.check_done() deferred_status = self.threadsafe_restriction_tracker.deferred_status() if deferred_status: deferred_restriction, deferred_timestamp = deferred_status element = windowed_value.value size = self.signature.get_restriction_provider().restriction_size( element, deferred_restriction) if size < 0: raise ValueError('Expected size >= 0 but received %s.' % size) current_watermark = ( self.threadsafe_watermark_estimator.current_watermark()) estimator_state = ( self.threadsafe_watermark_estimator.get_estimator_state()) residual_value = ((element, (deferred_restriction, estimator_state)), size) return SplitResultResidual( residual_value=windowed_value.with_value(residual_value), current_watermark=current_watermark, deferred_timestamp=deferred_timestamp) return None def _invoke_process_batch_per_window( self, windowed_batch: WindowedBatch, additional_args, additional_kwargs, ): # type: (...) -> Optional[SplitResultResidual] if self.has_cached_window_batch_args: args_for_process_batch, kwargs_for_process_batch = ( self.args_for_process_batch, self.kwargs_for_process_batch) else: if self.has_windowed_inputs: assert isinstance(windowed_batch, HomogeneousWindowedBatch) assert len(windowed_batch.windows) <= 1 window, = windowed_batch.windows else: window = GlobalWindow() side_inputs = [si[window] for si in self.side_inputs] side_inputs.extend(additional_args) args_for_process_batch, kwargs_for_process_batch = ( util.insert_values_in_args( self.args_for_process_batch, self.kwargs_for_process_batch, side_inputs, ) ) if not self.recalculate_window_args: self.args_for_process_batch, self.kwargs_for_process_batch = ( args_for_process_batch, kwargs_for_process_batch) self.has_cached_window_batch_args = True for i, p in self.placeholders_for_process_batch: if core.DoFn.ElementParam == p: args_for_process_batch[i] = windowed_batch.values elif core.DoFn.KeyParam == p: raise NotImplementedError( 'https://github.com/apache/beam/issues/21653: Per-key process_batch' ) elif core.DoFn.WindowParam == p: args_for_process_batch[i] = window elif core.DoFn.TimestampParam == p: args_for_process_batch[i] = windowed_batch.timestamp elif core.DoFn.PaneInfoParam == p: assert isinstance(windowed_batch, HomogeneousWindowedBatch) args_for_process_batch[i] = windowed_batch.pane_info elif isinstance(p, core.DoFn.StateParam): raise NotImplementedError( "https://github.com/apache/beam/issues/21653: " "Per-key process_batch") elif isinstance(p, core.DoFn.TimerParam): raise NotImplementedError( "https://github.com/apache/beam/issues/21653: " "Per-key process_batch") elif isinstance(p, core.DoFn.BundleContextParam): args_for_process_batch[i] = self._bundle_context_values[p][1] elif isinstance(p, core.DoFn.SetupContextParam): args_for_process_batch[i] = self._setup_context_values[p][1] kwargs_for_process_batch = kwargs_for_process_batch or {} if additional_kwargs: kwargs_for_process_batch.update(additional_kwargs) self.output_handler.handle_process_batch_outputs( windowed_batch, self.process_batch_method( *args_for_process_batch, **kwargs_for_process_batch), self.threadsafe_watermark_estimator) @staticmethod def _try_split( fraction, window_index, # type: Optional[int] stop_window_index, # type: Optional[int] windowed_value, # type: WindowedValue restriction, watermark_estimator_state, restriction_provider, # type: RestrictionProvider restriction_tracker, # type: RestrictionTracker watermark_estimator, # type: WatermarkEstimator ): # type: (...) -> Optional[Tuple[Iterable[SplitResultPrimary], Iterable[SplitResultResidual], Optional[int]]] """Try to split returning a primaries, residuals and a new stop index. For non-window observing splittable DoFns we split the current restriction and assign the primary and residual to all the windows. For window observing splittable DoFns, we: 1) return a split at a window boundary if the fraction lies outside of the current window. 2) attempt to split the current restriction, if successful then return the primary and residual for the current window and an additional primary and residual for any fully processed and fully unprocessed windows. 3) fall back to returning a split at the window boundary if possible Args: window_index: the current index of the window being processed or None if the splittable DoFn is not window observing. stop_window_index: the current index to stop processing at or None if the splittable DoFn is not window observing. windowed_value: the current windowed value restriction: the initial restriction when processing was started. watermark_estimator_state: the initial watermark estimator state when processing was started. restriction_provider: the DoFn's restriction provider restriction_tracker: the current restriction tracker watermark_estimator: the current watermark estimator Returns: A tuple containing (primaries, residuals, new_stop_index) or None if splitting was not possible. new_stop_index will only be set if the splittable DoFn is window observing otherwise it will be None. """ def compute_whole_window_split(to_index, from_index): restriction_size = restriction_provider.restriction_size( windowed_value, restriction) if restriction_size < 0: raise ValueError( 'Expected size >= 0 but received %s.' % restriction_size) # The primary and residual both share the same value only differing # by the set of windows they are in. value = ((windowed_value.value, (restriction, watermark_estimator_state)), restriction_size) primary_restriction = SplitResultPrimary( primary_value=WindowedValue( value, windowed_value.timestamp, windowed_value.windows[:to_index])) if to_index > 0 else None # Don't report any updated watermarks for the residual since they have # not processed any part of the restriction. residual_restriction = SplitResultResidual( residual_value=WindowedValue( value, windowed_value.timestamp, windowed_value.windows[from_index:stop_window_index]), current_watermark=None, deferred_timestamp=None) if from_index < stop_window_index else None return (primary_restriction, residual_restriction) primary_restrictions = [] residual_restrictions = [] window_observing = window_index is not None # If we are processing each window separately and we aren't on the last # window then compute whether the split lies within the current window # or a future window. if window_observing and window_index != stop_window_index - 1: progress = restriction_tracker.current_progress() if not progress: # Assume no work has been completed for the current window if progress # is unavailable. from apache_beam.io.iobase import RestrictionProgress progress = RestrictionProgress(completed=0, remaining=1) scaled_progress = PerWindowInvoker._scale_progress( progress, window_index, stop_window_index) # Compute the fraction of the remainder relative to the scaled progress. # If the value is greater than or equal to progress.remaining_work then we # should split at the closest window boundary. fraction_of_remainder = scaled_progress.remaining_work * fraction if fraction_of_remainder >= progress.remaining_work: # The fraction is outside of the current window and hence we will # split at the closest window boundary. Favor a split and return the # last window if we would have rounded up to the end of the window # based upon the fraction. new_stop_window_index = min( stop_window_index - 1, window_index + max( 1, int( round(( progress.completed_work + scaled_progress.remaining_work * fraction) / progress.total_work)))) primary, residual = compute_whole_window_split( new_stop_window_index, new_stop_window_index) assert primary is not None assert residual is not None return ([primary], [residual], new_stop_window_index) else: # The fraction is within the current window being processed so compute # the updated fraction based upon the number of windows being processed. new_stop_window_index = window_index + 1 fraction = fraction_of_remainder / progress.remaining_work # Attempt to split below, if we can't then we'll compute a split # using only window boundaries else: # We aren't splitting within multiple windows so we don't change our # stop index. new_stop_window_index = stop_window_index # Temporary workaround for [BEAM-7473]: get current_watermark before # split, in case watermark gets advanced before getting split results. # In worst case, current_watermark is always stale, which is ok. current_watermark = (watermark_estimator.current_watermark()) current_estimator_state = (watermark_estimator.get_estimator_state()) split = restriction_tracker.try_split(fraction) if split: primary, residual = split element = windowed_value.value primary_size = restriction_provider.restriction_size( windowed_value.value, primary) if primary_size < 0: raise ValueError('Expected size >= 0 but received %s.' % primary_size) residual_size = restriction_provider.restriction_size( windowed_value.value, residual) if residual_size < 0: raise ValueError('Expected size >= 0 but received %s.' % residual_size) # We use the watermark estimator state for the original process call # for the primary and the updated watermark estimator state for the # residual for the split. primary_split_value = ((element, (primary, watermark_estimator_state)), primary_size) residual_split_value = ((element, (residual, current_estimator_state)), residual_size) windows = ( windowed_value.windows[window_index], ) if window_observing else windowed_value.windows primary_restrictions.append( SplitResultPrimary( primary_value=WindowedValue( primary_split_value, windowed_value.timestamp, windows))) residual_restrictions.append( SplitResultResidual( residual_value=WindowedValue( residual_split_value, windowed_value.timestamp, windows), current_watermark=current_watermark, deferred_timestamp=None)) if window_observing: assert new_stop_window_index == window_index + 1 primary, residual = compute_whole_window_split( window_index, window_index + 1) if primary: primary_restrictions.append(primary) if residual: residual_restrictions.append(residual) return ( primary_restrictions, residual_restrictions, new_stop_window_index) elif new_stop_window_index and new_stop_window_index != stop_window_index: # If we failed to split but have a new stop index then return a split # at the window boundary. primary, residual = compute_whole_window_split( new_stop_window_index, new_stop_window_index) assert primary is not None assert residual is not None return ([primary], [residual], new_stop_window_index) else: return None def try_split(self, fraction): # type: (...) -> Optional[Tuple[Iterable[SplitResultPrimary], Iterable[SplitResultResidual]]] if not self.is_splittable: return None with self.splitting_lock: if not self.threadsafe_restriction_tracker: return None # Make a local reference to member variables that change references during # processing under lock before attempting to split so we have a consistent # view of all the references. result = PerWindowInvoker._try_split( fraction, self.current_window_index, self.stop_window_index, self.current_windowed_value, self.restriction, self.watermark_estimator_state, self.signature.get_restriction_provider(), self.threadsafe_restriction_tracker, self.threadsafe_watermark_estimator) if not result: return None residuals, primaries, self.stop_window_index = result return (residuals, primaries) @staticmethod def _scale_progress(progress, window_index, stop_window_index): # We scale progress based upon the amount of work we will do for one # window and have it apply for all windows. completed = window_index * progress.total_work + progress.completed_work remaining = ( stop_window_index - (window_index + 1)) * progress.total_work + progress.remaining_work from apache_beam.io.iobase import RestrictionProgress return RestrictionProgress(completed=completed, remaining=remaining) def current_element_progress(self): # type: () -> Optional[RestrictionProgress] if not self.is_splittable: return None with self.splitting_lock: current_window_index = self.current_window_index stop_window_index = self.stop_window_index threadsafe_restriction_tracker = self.threadsafe_restriction_tracker if not threadsafe_restriction_tracker: return None progress = threadsafe_restriction_tracker.current_progress() if not current_window_index or not progress: return progress # stop_window_index should always be set if current_window_index is set, # it is an error otherwise. assert stop_window_index return PerWindowInvoker._scale_progress( progress, current_window_index, stop_window_index) class DoFnRunner: """For internal use only; no backwards-compatibility guarantees. A helper class for executing ParDo operations. """ def __init__( self, fn, # type: core.DoFn args, kwargs, side_inputs, # type: Iterable[sideinputs.SideInputMap] windowing, tagged_receivers, # type: Mapping[Optional[str], Receiver] step_name=None, # type: Optional[str] logging_context=None, state=None, scoped_metrics_container=None, operation_name=None, transform_id=None, user_state_context=None, # type: Optional[userstate.UserStateContext] ): """Initializes a DoFnRunner. Args: fn: user DoFn to invoke args: positional side input arguments (static and placeholder), if any kwargs: keyword side input arguments (static and placeholder), if any side_inputs: list of sideinput.SideInputMaps for deferred side inputs windowing: windowing properties of the output PCollection(s) tagged_receivers: a dict of tag name to Receiver objects step_name: the name of this step logging_context: DEPRECATED [BEAM-4728] state: handle for accessing DoFn state scoped_metrics_container: DEPRECATED operation_name: The system name assigned by the runner for this operation. transform_id: The PTransform Id in the pipeline proto for this DoFn. user_state_context: The UserStateContext instance for the current Stateful DoFn. """ # Need to support multiple iterations. side_inputs = list(side_inputs) self.step_name = step_name self.transform_id = transform_id self.context = DoFnContext(step_name, state=state) self.bundle_finalizer_param = DoFn.BundleFinalizerParam() self.execution_context = None # type: Optional[ExecutionContext] do_fn_signature = DoFnSignature(fn) # Optimize for the common case. main_receivers = tagged_receivers[None] # TODO(https://github.com/apache/beam/issues/18886): Remove if block after # output counter released. if 'outputs_per_element_counter' in RuntimeValueProvider.experiments: # TODO(BEAM-3955): Make step_name and operation_name less confused. output_counter_name = ( CounterName('per-element-output-count', step_name=operation_name)) per_element_output_counter = state._counter_factory.get_counter( output_counter_name, Counter.DATAFLOW_DISTRIBUTION).accumulator else: per_element_output_counter = None output_handler = _OutputHandler( windowing.windowfn, main_receivers, tagged_receivers, per_element_output_counter, getattr(fn, 'output_batch_converter', None), getattr( do_fn_signature.process_method.method_value, '_beam_yields_batches', False), getattr( do_fn_signature.process_batch_method.method_value, '_beam_yields_elements', False), ) if do_fn_signature.is_stateful_dofn() and not user_state_context: raise Exception( 'Requested execution of a stateful DoFn, but no user state context ' 'is available. This likely means that the current runner does not ' 'support the execution of stateful DoFns.') self.do_fn_invoker = DoFnInvoker.create_invoker( do_fn_signature, output_handler, self.context, side_inputs, args, kwargs, user_state_context=user_state_context, bundle_finalizer_param=self.bundle_finalizer_param) def process(self, windowed_value): # type: (WindowedValue) -> Iterable[SplitResultResidual] try: return self.do_fn_invoker.invoke_process(windowed_value) except BaseException as exn: self._reraise_augmented(exn, windowed_value) return [] def _maybe_sample_exception( self, exc_info: Tuple, windowed_value: Optional[WindowedValue]) -> None: if self.execution_context is None: return output_sampler = self.execution_context.output_sampler if output_sampler is None: return output_sampler.sample_exception( windowed_value, exc_info, self.transform_id, self.execution_context.instruction_id) def process_batch(self, windowed_batch): # type: (WindowedBatch) -> None try: self.do_fn_invoker.invoke_process_batch(windowed_batch) except BaseException as exn: self._reraise_augmented(exn) def process_with_sized_restriction(self, windowed_value): # type: (WindowedValue) -> Iterable[SplitResultResidual] (element, (restriction, estimator_state)), _ = windowed_value.value return self.do_fn_invoker.invoke_process( windowed_value.with_value(element), restriction=restriction, watermark_estimator_state=estimator_state) def try_split(self, fraction): # type: (...) -> Optional[Tuple[Iterable[SplitResultPrimary], Iterable[SplitResultResidual]]] assert isinstance(self.do_fn_invoker, PerWindowInvoker) return self.do_fn_invoker.try_split(fraction) def current_element_progress(self): # type: () -> Optional[RestrictionProgress] assert isinstance(self.do_fn_invoker, PerWindowInvoker) return self.do_fn_invoker.current_element_progress() def process_user_timer( self, timer_spec, key, window, timestamp, pane_info, dynamic_timer_tag): try: self.do_fn_invoker.invoke_user_timer( timer_spec, key, window, timestamp, pane_info, dynamic_timer_tag) except BaseException as exn: self._reraise_augmented(exn) def _invoke_bundle_method(self, bundle_method): try: self.context.set_element(None) bundle_method() except BaseException as exn: self._reraise_augmented(exn) def _invoke_lifecycle_method(self, lifecycle_method): try: self.context.set_element(None) lifecycle_method() except BaseException as exn: self._reraise_augmented(exn) def setup(self): # type: () -> None self._invoke_lifecycle_method(self.do_fn_invoker.invoke_setup) def start(self): # type: () -> None self._invoke_bundle_method(self.do_fn_invoker.invoke_start_bundle) def finish(self): # type: () -> None self._invoke_bundle_method(self.do_fn_invoker.invoke_finish_bundle) def teardown(self): # type: () -> None self._invoke_lifecycle_method(self.do_fn_invoker.invoke_teardown) def finalize(self): # type: () -> None self.bundle_finalizer_param.finalize_bundle() def _reraise_augmented(self, exn, windowed_value=None): if getattr(exn, '_tagged_with_step', False) or not self.step_name: raise exn step_annotation = " [while running '%s']" % self.step_name # To emulate exception chaining (not available in Python 2). try: # Attempt to construct the same kind of exception # with an augmented message. new_exn = type(exn)(exn.args[0] + step_annotation, *exn.args[1:]) new_exn._tagged_with_step = True # Could raise attribute error. except: # pylint: disable=bare-except # If anything goes wrong, construct a RuntimeError whose message # records the original exception's type and message. new_exn = RuntimeError( traceback.format_exception_only(type(exn), exn)[-1].strip() + step_annotation) new_exn._tagged_with_step = True exc_info = sys.exc_info() _, _, tb = exc_info new_exn = new_exn.with_traceback(tb) self._maybe_sample_exception(exc_info, windowed_value) _LOGGER.exception(new_exn) raise new_exn class OutputHandler(object): def handle_process_outputs( self, windowed_input_element, results, watermark_estimator=None): # type: (WindowedValue, Iterable[Any], Optional[WatermarkEstimator]) -> None raise NotImplementedError def handle_process_batch_outputs( self, windowed_input_element, results, watermark_estimator=None): # type: (WindowedBatch, Iterable[Any], Optional[WatermarkEstimator]) -> None raise NotImplementedError class _OutputHandler(OutputHandler): """Processes output produced by DoFn method invocations.""" def __init__( self, window_fn, main_receivers, # type: Receiver tagged_receivers, # type: Mapping[Optional[str], Receiver] per_element_output_counter, output_batch_converter, # type: Optional[BatchConverter] process_yields_batches, # type: bool process_batch_yields_elements, # type: bool ): """Initializes ``_OutputHandler``. Args: window_fn: a windowing function (WindowFn). main_receivers: a dict of tag name to Receiver objects. tagged_receivers: main receiver object. per_element_output_counter: per_element_output_counter of one work_item. could be none if experimental flag turn off """ self.window_fn = window_fn self.main_receivers = main_receivers self.tagged_receivers = tagged_receivers if (per_element_output_counter is not None and per_element_output_counter.is_cythonized): self.per_element_output_counter = per_element_output_counter else: self.per_element_output_counter = None self.output_batch_converter = output_batch_converter self._process_yields_batches = process_yields_batches self._process_batch_yields_elements = process_batch_yields_elements def handle_process_outputs( self, windowed_input_element, results, watermark_estimator=None): # type: (WindowedValue, Iterable[Any], Optional[WatermarkEstimator]) -> None """Dispatch the result of process computation to the appropriate receivers. A value wrapped in a TaggedOutput object will be unwrapped and then dispatched to the appropriate indexed output. """ if results is None: results = [] # TODO(https://github.com/apache/beam/issues/20404): Verify that the # results object is a valid iterable type if # performance_runtime_type_check is active, without harming performance output_element_count = 0 for result in results: tag, result = self._handle_tagged_output(result) if not self._process_yields_batches: # process yields elements windowed_value = self._maybe_propagate_windowing_info( windowed_input_element, result) output_element_count += 1 self._write_value_to_tag(tag, windowed_value, watermark_estimator) else: # process yields batches self._verify_batch_output(result) if isinstance(result, WindowedBatch): assert isinstance(result, HomogeneousWindowedBatch) windowed_batch = result if (windowed_input_element is not None and len(windowed_input_element.windows) != 1): windowed_batch.windows *= len(windowed_input_element.windows) else: windowed_batch = ( HomogeneousWindowedBatch.from_batch_and_windowed_value( batch=result, windowed_value=windowed_input_element)) output_element_count += self.output_batch_converter.get_length( windowed_batch.values) self._write_batch_to_tag(tag, windowed_batch, watermark_estimator) # TODO(https://github.com/apache/beam/issues/18886): Remove if block after # output counter released. Only enable per_element_output_counter when # counter cythonized if self.per_element_output_counter is not None: self.per_element_output_counter.add_input(output_element_count) def handle_process_batch_outputs( self, windowed_input_batch, results, watermark_estimator=None): # type: (WindowedBatch, Iterable[Any], Optional[WatermarkEstimator]) -> None """Dispatch the result of process_batch computation to the appropriate receivers. A value wrapped in a TaggedOutput object will be unwrapped and then dispatched to the appropriate indexed output. """ if results is None: results = [] output_element_count = 0 for result in results: tag, result = self._handle_tagged_output(result) if not self._process_batch_yields_elements: # process_batch yields batches assert self.output_batch_converter is not None self._verify_batch_output(result) if isinstance(result, WindowedBatch): assert isinstance(result, HomogeneousWindowedBatch) windowed_batch = result if (windowed_input_batch is not None and len(windowed_input_batch.windows) != 1): windowed_batch.windows *= len(windowed_input_batch.windows) else: windowed_batch = windowed_input_batch.with_values(result) output_element_count += self.output_batch_converter.get_length( windowed_batch.values) self._write_batch_to_tag(tag, windowed_batch, watermark_estimator) else: # process_batch yields elements assert isinstance(windowed_input_batch, HomogeneousWindowedBatch) windowed_value = self._maybe_propagate_windowing_info( windowed_input_batch.as_empty_windowed_value(), result) output_element_count += 1 self._write_value_to_tag(tag, windowed_value, watermark_estimator) # TODO(https://github.com/apache/beam/issues/18886): Remove if block after # output counter released. Only enable per_element_output_counter when # counter cythonized if self.per_element_output_counter is not None: self.per_element_output_counter.add_input(output_element_count) def _maybe_propagate_windowing_info(self, windowed_input_element, result): # type: (WindowedValue, Any) -> WindowedValue if isinstance(result, WindowedValue): windowed_value = result if (windowed_input_element is not None and len(windowed_input_element.windows) != 1): windowed_value.windows *= len(windowed_input_element.windows) return windowed_value elif isinstance(result, TimestampedValue): assign_context = WindowFn.AssignContext(result.timestamp, result.value) windowed_value = WindowedValue( result.value, result.timestamp, self.window_fn.assign(assign_context)) if len(windowed_input_element.windows) != 1: windowed_value.windows *= len(windowed_input_element.windows) return windowed_value else: return windowed_input_element.with_value(result) def _handle_tagged_output(self, result): if isinstance(result, TaggedOutput): tag = result.tag if not isinstance(tag, str): raise TypeError('In %s, tag %s is not a string' % (self, tag)) return tag, result.value return None, result def _write_value_to_tag(self, tag, windowed_value, watermark_estimator): if watermark_estimator is not None: watermark_estimator.observe_timestamp(windowed_value.timestamp) if tag is None: self.main_receivers.receive(windowed_value) else: self.tagged_receivers[tag].receive(windowed_value) def _write_batch_to_tag(self, tag, windowed_batch, watermark_estimator): if watermark_estimator is not None: for timestamp in windowed_batch.timestamps: watermark_estimator.observe_timestamp(timestamp) if tag is None: self.main_receivers.receive_batch(windowed_batch) else: self.tagged_receivers[tag].receive_batch(windowed_batch) def _verify_batch_output(self, result): if isinstance(result, (WindowedValue, TimestampedValue)): raise TypeError( f"Received {type(result).__name__} from DoFn that was " "expected to produce a batch.") def start_bundle_outputs(self, results): """Validate that start_bundle does not output any elements""" if results is None: return raise RuntimeError( 'Start Bundle should not output any elements but got %s' % results) def finish_bundle_outputs(self, results): """Dispatch the result of finish_bundle to the appropriate receivers. A value wrapped in a TaggedOutput object will be unwrapped and then dispatched to the appropriate indexed output. """ if results is None: return for result in results: tag = None if isinstance(result, TaggedOutput): tag = result.tag if not isinstance(tag, str): raise TypeError('In %s, tag %s is not a string' % (self, tag)) result = result.value if isinstance(result, WindowedValue): windowed_value = result else: raise RuntimeError('Finish Bundle should only output WindowedValue ' +\ 'type but got %s' % type(result)) if tag is None: self.main_receivers.receive(windowed_value) else: self.tagged_receivers[tag].receive(windowed_value) class _NoContext(WindowFn.AssignContext): """An uninspectable WindowFn.AssignContext.""" NO_VALUE = object() def __init__(self, value, timestamp=NO_VALUE): self.value = value self._timestamp = timestamp @property def timestamp(self): if self._timestamp is self.NO_VALUE: raise ValueError('No timestamp in this context.') else: return self._timestamp @property def existing_windows(self): raise ValueError('No existing_windows in this context.') class DoFnState(object): """For internal use only; no backwards-compatibility guarantees. Keeps track of state that DoFns want, currently, user counters. """ def __init__(self, counter_factory): self.step_name = '' self._counter_factory = counter_factory def counter_for(self, aggregator): """Looks up the counter for this aggregator, creating one if necessary.""" return self._counter_factory.get_aggregator_counter( self.step_name, aggregator) # TODO(robertwb): Replace core.DoFnContext with this. class DoFnContext(object): """For internal use only; no backwards-compatibility guarantees.""" def __init__(self, label, element=None, state=None): self.label = label self.state = state if element is not None: self.set_element(element) def set_element(self, windowed_value): # type: (Optional[WindowedValue]) -> None self.windowed_value = windowed_value @property def element(self): if self.windowed_value is None: raise AttributeError('element not accessible in this context') else: return self.windowed_value.value @property def timestamp(self): if self.windowed_value is None: raise AttributeError('timestamp not accessible in this context') else: return self.windowed_value.timestamp @property def windows(self): if self.windowed_value is None: raise AttributeError('windows not accessible in this context') else: return self.windowed_value.windows ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/common_test.py0000644000175100017510000005160515206113032022261 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import hamcrest as hc import apache_beam as beam from apache_beam.io.restriction_trackers import OffsetRange from apache_beam.io.restriction_trackers import OffsetRestrictionTracker from apache_beam.io.watermark_estimators import ManualWatermarkEstimator from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.runners.common import DoFnSignature from apache_beam.runners.common import PerWindowInvoker from apache_beam.runners.sdf_utils import SplitResultPrimary from apache_beam.runners.sdf_utils import SplitResultResidual from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.transforms.core import DoFn from apache_beam.transforms.core import RestrictionProvider from apache_beam.transforms.window import IntervalWindow from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.windowed_value import WindowedValue class DoFnSignatureTest(unittest.TestCase): def test_dofn_validate_process_error(self): class MyDoFn(DoFn): def process(self, element, w1=DoFn.WindowParam, w2=DoFn.WindowParam): pass with self.assertRaises(ValueError): DoFnSignature(MyDoFn()) def test_dofn_get_defaults(self): class MyDoFn(DoFn): def process(self, element, w=DoFn.WindowParam): pass signature = DoFnSignature(MyDoFn()) self.assertEqual(signature.process_method.defaults, [DoFn.WindowParam]) @unittest.skip('BEAM-5878') def test_dofn_get_defaults_kwonly(self): class MyDoFn(DoFn): def process(self, element, *, w=DoFn.WindowParam): pass signature = DoFnSignature(MyDoFn()) self.assertEqual(signature.process_method.defaults, [DoFn.WindowParam]) def test_dofn_validate_start_bundle_error(self): class MyDoFn(DoFn): def process(self, element): pass def start_bundle(self, w1=DoFn.WindowParam): pass with self.assertRaises(ValueError): DoFnSignature(MyDoFn()) def test_dofn_validate_finish_bundle_error(self): class MyDoFn(DoFn): def process(self, element): pass def finish_bundle(self, w1=DoFn.WindowParam): pass with self.assertRaises(ValueError): DoFnSignature(MyDoFn()) def test_unbounded_element_process_fn(self): class UnboundedDoFn(DoFn): @DoFn.unbounded_per_element() def process(self, element): pass class BoundedDoFn(DoFn): def process(self, element): pass signature = DoFnSignature(UnboundedDoFn()) self.assertTrue(signature.is_unbounded_per_element()) signature = DoFnSignature(BoundedDoFn()) self.assertFalse(signature.is_unbounded_per_element()) class DoFnProcessTest(unittest.TestCase): # pylint: disable=expression-not-assigned all_records = None def setUp(self): DoFnProcessTest.all_records = [] def record_dofn(self): class RecordDoFn(DoFn): def process(self, element): DoFnProcessTest.all_records.append(element) return RecordDoFn() def test_dofn_process_keyparam(self): class DoFnProcessWithKeyparam(DoFn): def process(self, element, mykey=DoFn.KeyParam): yield "{key}-verify".format(key=mykey) pipeline_options = PipelineOptions() with TestPipeline(options=pipeline_options) as p: test_stream = (TestStream().advance_watermark_to(10).add_elements([1, 2])) ( p | test_stream | beam.Map(lambda x: (x, "some-value")) | "window_into" >> beam.WindowInto( window.FixedWindows(5), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.ParDo(DoFnProcessWithKeyparam()) | beam.ParDo(self.record_dofn())) self.assertEqual(['1-verify', '2-verify'], sorted(DoFnProcessTest.all_records)) def test_dofn_process_keyparam_error_no_key(self): class DoFnProcessWithKeyparam(DoFn): def process(self, element, mykey=DoFn.KeyParam): yield "{key}-verify".format(key=mykey) pipeline_options = PipelineOptions() with self.assertRaises(ValueError),\ TestPipeline(options=pipeline_options) as p: test_stream = (TestStream().advance_watermark_to(10).add_elements([1, 2])) (p | test_stream | beam.ParDo(DoFnProcessWithKeyparam())) def test_pardo_with_unbounded_per_element_dofn(self): class UnboundedDoFn(beam.DoFn): @beam.DoFn.unbounded_per_element() def process(self, element): pass class BoundedDoFn(beam.DoFn): def process(self, element): pass with TestPipeline() as p: source = p | beam.Impulse() unbounded_pcoll = source | beam.ParDo(UnboundedDoFn()) bounded_pcoll = source | beam.ParDo(BoundedDoFn()) self.assertEqual(unbounded_pcoll.is_bounded, False) self.assertEqual(bounded_pcoll.is_bounded, True) class TestOffsetRestrictionProvider(RestrictionProvider): def restriction_size(self, element, restriction): return restriction.size() class PerWindowInvokerSplitTest(unittest.TestCase): def setUp(self): self.window1 = IntervalWindow(0, 10) self.window2 = IntervalWindow(10, 20) self.window3 = IntervalWindow(20, 30) self.windowed_value = WindowedValue( 'a', 57, (self.window1, self.window2, self.window3)) self.restriction = OffsetRange(0, 100) self.watermark_estimator_state = Timestamp(21) self.restriction_provider = TestOffsetRestrictionProvider() self.watermark_estimator = ManualWatermarkEstimator(Timestamp(42)) self.maxDiff = None def create_split_in_window(self, offset_index, windows): return ( SplitResultPrimary( primary_value=WindowedValue((( 'a', (OffsetRange(0, offset_index), self.watermark_estimator_state)), offset_index), 57, windows)), SplitResultResidual( residual_value=WindowedValue((( 'a', ( OffsetRange(offset_index, 100), self.watermark_estimator.get_estimator_state())), 100 - offset_index), 57, windows), current_watermark=self.watermark_estimator.current_watermark(), deferred_timestamp=None)) def create_split_across_windows(self, primary_windows, residual_windows): primary = SplitResultPrimary( primary_value=WindowedValue( (('a', (OffsetRange(0, 100), self.watermark_estimator_state)), 100), 57, primary_windows)) if primary_windows else None residual = SplitResultResidual( residual_value=WindowedValue( (('a', (OffsetRange(0, 100), self.watermark_estimator_state)), 100), 57, residual_windows), current_watermark=None, deferred_timestamp=None) if residual_windows else None return primary, residual def test_non_window_observing_checkpoint(self): # test checkpoint restriction_tracker = OffsetRestrictionTracker(self.restriction) restriction_tracker.try_claim(30) (primaries, residuals, stop_index) = PerWindowInvoker._try_split( 0.0, None, None, self.windowed_value, self.restriction, self.watermark_estimator_state, self.restriction_provider, restriction_tracker, self.watermark_estimator) expected_primary_split, expected_residual_split = ( self.create_split_in_window(31, self.windowed_value.windows)) self.assertEqual([expected_primary_split], primaries) self.assertEqual([expected_residual_split], residuals) # We don't expect the stop index to be set for non window observing splits self.assertIsNone(stop_index) def test_non_window_observing_split(self): restriction_tracker = OffsetRestrictionTracker(self.restriction) restriction_tracker.try_claim(30) (primaries, residuals, stop_index) = PerWindowInvoker._try_split( 0.1, None, None, self.windowed_value, self.restriction, self.watermark_estimator_state, self.restriction_provider, restriction_tracker, self.watermark_estimator) expected_primary_split, expected_residual_split = ( self.create_split_in_window(37, self.windowed_value.windows)) self.assertEqual([expected_primary_split], primaries) self.assertEqual([expected_residual_split], residuals) # We don't expect the stop index to be set for non window observing splits self.assertIsNone(stop_index) def test_non_window_observing_split_when_restriction_is_done(self): restriction_tracker = OffsetRestrictionTracker(self.restriction) restriction_tracker.try_claim(100) self.assertIsNone( PerWindowInvoker._try_split( 0.1, None, None, self.windowed_value, self.restriction, self.watermark_estimator_state, self.restriction_provider, restriction_tracker, self.watermark_estimator)) def test_window_observing_checkpoint_on_first_window(self): restriction_tracker = OffsetRestrictionTracker(self.restriction) restriction_tracker.try_claim(30) (primaries, residuals, stop_index) = PerWindowInvoker._try_split( 0.0, 0, 3, self.windowed_value, self.restriction, self.watermark_estimator_state, self.restriction_provider, restriction_tracker, self.watermark_estimator) expected_primary_split, expected_residual_split = ( self.create_split_in_window(31, (self.window1, ))) _, expected_residual_windows = ( self.create_split_across_windows(None, (self.window2, self.window3,))) hc.assert_that(primaries, hc.contains_inanyorder(expected_primary_split)) hc.assert_that( residuals, hc.contains_inanyorder( expected_residual_split, expected_residual_windows, )) self.assertEqual(stop_index, 1) def test_window_observing_checkpoint_on_first_window_after_prior_split(self): restriction_tracker = OffsetRestrictionTracker(self.restriction) restriction_tracker.try_claim(30) (primaries, residuals, stop_index) = PerWindowInvoker._try_split( 0.0, 0, 2, # stop index < len(windows) representing a prior split had occurred self.windowed_value, self.restriction, self.watermark_estimator_state, self.restriction_provider, restriction_tracker, self.watermark_estimator) expected_primary_split, expected_residual_split = ( self.create_split_in_window(31, (self.window1, ))) _, expected_residual_windows = ( self.create_split_across_windows(None, (self.window2, ))) hc.assert_that(primaries, hc.contains_inanyorder(expected_primary_split)) hc.assert_that( residuals, hc.contains_inanyorder( expected_residual_split, expected_residual_windows, )) self.assertEqual(stop_index, 1) def test_window_observing_split_on_first_window(self): restriction_tracker = OffsetRestrictionTracker(self.restriction) restriction_tracker.try_claim(30) (primaries, residuals, stop_index) = PerWindowInvoker._try_split( 0.2, 0, 3, self.windowed_value, self.restriction, self.watermark_estimator_state, self.restriction_provider, restriction_tracker, self.watermark_estimator) # 20% of 2.7 windows = 20% of 270 offset left = 54 offset # 30 + 54 = 84 split offset expected_primary_split, expected_residual_split = ( self.create_split_in_window(84, (self.window1, ))) _, expected_residual_windows = ( self.create_split_across_windows(None, (self.window2, self.window3, ))) hc.assert_that(primaries, hc.contains_inanyorder(expected_primary_split)) hc.assert_that( residuals, hc.contains_inanyorder( expected_residual_split, expected_residual_windows, )) self.assertEqual(stop_index, 1) def test_window_observing_split_on_middle_window(self): restriction_tracker = OffsetRestrictionTracker(self.restriction) restriction_tracker.try_claim(30) (primaries, residuals, stop_index) = PerWindowInvoker._try_split( 0.2, 1, 3, self.windowed_value, self.restriction, self.watermark_estimator_state, self.restriction_provider, restriction_tracker, self.watermark_estimator) # 20% of 1.7 windows = 20% of 170 offset left = 34 offset # 30 + 34 = 64 split offset expected_primary_split, expected_residual_split = ( self.create_split_in_window(64, (self.window2, ))) expected_primary_windows, expected_residual_windows = ( self.create_split_across_windows((self.window1, ), (self.window3, ))) hc.assert_that( primaries, hc.contains_inanyorder( expected_primary_split, expected_primary_windows, )) hc.assert_that( residuals, hc.contains_inanyorder( expected_residual_split, expected_residual_windows, )) self.assertEqual(stop_index, 2) def test_window_observing_split_on_last_window(self): restriction_tracker = OffsetRestrictionTracker(self.restriction) restriction_tracker.try_claim(30) (primaries, residuals, stop_index) = PerWindowInvoker._try_split( 0.2, 2, 3, self.windowed_value, self.restriction, self.watermark_estimator_state, self.restriction_provider, restriction_tracker, self.watermark_estimator) # 20% of 0.7 windows = 20% of 70 offset left = 14 offset # 30 + 14 = 44 split offset expected_primary_split, expected_residual_split = ( self.create_split_in_window(44, (self.window3, ))) expected_primary_windows, _ = ( self.create_split_across_windows((self.window1, self.window2, ), None)) hc.assert_that( primaries, hc.contains_inanyorder( expected_primary_split, expected_primary_windows, )) hc.assert_that( residuals, hc.contains_inanyorder( expected_residual_split, )) self.assertEqual(stop_index, 3) def test_window_observing_split_on_first_window_fallback(self): restriction_tracker = OffsetRestrictionTracker(self.restriction) restriction_tracker.try_claim(100) # We assume that we can't split this fully claimed restriction self.assertIsNone(restriction_tracker.try_split(0)) (primaries, residuals, stop_index) = PerWindowInvoker._try_split( 0.0, 0, 3, self.windowed_value, self.restriction, self.watermark_estimator_state, self.restriction_provider, restriction_tracker, self.watermark_estimator) expected_primary_windows, expected_residual_windows = ( self.create_split_across_windows( (self.window1, ), (self.window2, self.window3, ))) hc.assert_that( primaries, hc.contains_inanyorder( expected_primary_windows, )) hc.assert_that( residuals, hc.contains_inanyorder( expected_residual_windows, )) self.assertEqual(stop_index, 1) def test_window_observing_split_on_middle_window_fallback(self): restriction_tracker = OffsetRestrictionTracker(self.restriction) restriction_tracker.try_claim(100) # We assume that we can't split this fully claimed restriction self.assertIsNone(restriction_tracker.try_split(0)) (primaries, residuals, stop_index) = PerWindowInvoker._try_split( 0.0, 1, 3, self.windowed_value, self.restriction, self.watermark_estimator_state, self.restriction_provider, restriction_tracker, self.watermark_estimator) expected_primary_windows, expected_residual_windows = ( self.create_split_across_windows( (self.window1, self.window2, ), (self.window3, ))) hc.assert_that( primaries, hc.contains_inanyorder( expected_primary_windows, )) hc.assert_that( residuals, hc.contains_inanyorder( expected_residual_windows, )) self.assertEqual(stop_index, 2) def test_window_observing_split_on_last_window_when_split_not_possible(self): restriction_tracker = OffsetRestrictionTracker(self.restriction) restriction_tracker.try_claim(100) # We assume that we can't split this fully claimed restriction self.assertIsNone(restriction_tracker.try_split(0)) self.assertIsNone( PerWindowInvoker._try_split( 0.0, 2, 3, self.windowed_value, self.restriction, self.watermark_estimator_state, self.restriction_provider, restriction_tracker, self.watermark_estimator)) def test_window_observing_split_on_window_boundary_round_up(self): restriction_tracker = OffsetRestrictionTracker(self.restriction) restriction_tracker.try_claim(30) (primaries, residuals, stop_index) = PerWindowInvoker._try_split( 0.6, 0, 3, self.windowed_value, self.restriction, self.watermark_estimator_state, self.restriction_provider, restriction_tracker, self.watermark_estimator) # 60% of 2.7 windows = 60% of 270 offset left = 162 offset # 30 + 162 = 192 offset --> round to end of window 2 expected_primary_windows, expected_residual_windows = ( self.create_split_across_windows( (self.window1, self.window2, ), (self.window3, ))) hc.assert_that( primaries, hc.contains_inanyorder( expected_primary_windows, )) hc.assert_that( residuals, hc.contains_inanyorder( expected_residual_windows, )) self.assertEqual(stop_index, 2) def test_window_observing_split_on_window_boundary_round_down(self): restriction_tracker = OffsetRestrictionTracker(self.restriction) restriction_tracker.try_claim(30) (primaries, residuals, stop_index) = PerWindowInvoker._try_split( 0.3, 0, 3, self.windowed_value, self.restriction, self.watermark_estimator_state, self.restriction_provider, restriction_tracker, self.watermark_estimator) # 30% of 2.7 windows = 30% of 270 offset left = 81 offset # 30 + 81 = 111 offset --> round to end of window 1 expected_primary_windows, expected_residual_windows = ( self.create_split_across_windows( (self.window1, ), (self.window2, self.window3, ))) hc.assert_that( primaries, hc.contains_inanyorder( expected_primary_windows, )) hc.assert_that( residuals, hc.contains_inanyorder( expected_residual_windows, )) self.assertEqual(stop_index, 1) def test_window_observing_split_on_window_boundary_round_down_on_last_window( self): restriction_tracker = OffsetRestrictionTracker(self.restriction) restriction_tracker.try_claim(30) (primaries, residuals, stop_index) = PerWindowInvoker._try_split( 0.9, 0, 3, self.windowed_value, self.restriction, self.watermark_estimator_state, self.restriction_provider, restriction_tracker, self.watermark_estimator) # 90% of 2.7 windows = 90% of 270 offset left = 243 offset # 30 + 243 = 273 offset --> prefer a split so round to end of window 2 # instead of no split expected_primary_windows, expected_residual_windows = ( self.create_split_across_windows( (self.window1, self.window2, ), (self.window3, ))) hc.assert_that( primaries, hc.contains_inanyorder( expected_primary_windows, )) hc.assert_that( residuals, hc.contains_inanyorder( expected_residual_windows, )) self.assertEqual(stop_index, 2) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5674276 apache_beam-2.74.0/apache_beam/runners/dask/0000755000175100017510000000000015206113160020275 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dask/__init__.py0000644000175100017510000000142015206113032022401 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dask/dask_runner.py0000644000175100017510000002207615206113032023167 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """DaskRunner, executing remote jobs on Dask.distributed. The DaskRunner is a runner implementation that executes a graph of transformations across processes and workers via Dask distributed's scheduler. """ import argparse import collections import dataclasses import typing as t from apache_beam import pvalue from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.pipeline import AppliedPTransform from apache_beam.pipeline import PipelineVisitor from apache_beam.runners.dask.overrides import dask_overrides from apache_beam.runners.dask.transform_evaluator import TRANSLATIONS from apache_beam.runners.dask.transform_evaluator import DaskBagWindowedIterator from apache_beam.runners.dask.transform_evaluator import Flatten from apache_beam.runners.dask.transform_evaluator import NoOp from apache_beam.runners.direct.direct_runner import BundleBasedDirectRunner from apache_beam.runners.runner import PipelineResult from apache_beam.runners.runner import PipelineState from apache_beam.transforms.sideinputs import SideInputMap from apache_beam.utils.interactive_utils import is_in_notebook try: # Added to try to prevent threading related issues, see # https://github.com/pytest-dev/pytest/issues/3216#issuecomment-1502451456 import dask.distributed as ddist except ImportError: ddist = {} class DaskOptions(PipelineOptions): @staticmethod def _parse_timeout(candidate): try: return int(candidate) except (TypeError, ValueError): import dask return dask.config.no_default @staticmethod def _extract_bag_kwargs(dask_options: dict) -> dict: """Parse keyword arguments for `dask.Bag`s; used in graph translation.""" out = {} if npartitions := dask_options.pop('npartitions', None): out['npartitions'] = npartitions if partition_size := dask_options.pop('partition_size', None): out['partition_size'] = partition_size return out @classmethod def _add_argparse_args(cls, parser: argparse.ArgumentParser) -> None: parser.add_argument( '--dask_client_address', dest='address', type=str, default=None, help='Address of a dask Scheduler server. Will default to a ' '`dask.LocalCluster()`.') parser.add_argument( '--dask_connection_timeout', dest='timeout', type=DaskOptions._parse_timeout, help='Timeout duration for initial connection to the scheduler.') parser.add_argument( '--dask_scheduler_file', dest='scheduler_file', type=str, default=None, help='Path to a file with scheduler information if available.') # TODO(alxr): Add options for security. parser.add_argument( '--dask_client_name', dest='name', type=str, default=None, help='Gives the client a name that will be included in logs generated ' 'on the scheduler for matters relating to this client.') parser.add_argument( '--dask_connection_limit', dest='connection_limit', type=int, default=512, help='The number of open comms to maintain at once in the connection ' 'pool.') partitions_parser = parser.add_mutually_exclusive_group() partitions_parser.add_argument( '--dask_npartitions', dest='npartitions', type=int, default=None, help='The desired number of `dask.Bag` partitions. When unspecified, ' 'an educated guess is made.') partitions_parser.add_argument( '--dask_partition_size', dest='partition_size', type=int, default=None, help='The length of each `dask.Bag` partition. When unspecified, ' 'an educated guess is made.') @dataclasses.dataclass class DaskRunnerResult(PipelineResult): client: ddist.Client futures: t.Sequence[ddist.Future] def __post_init__(self): super().__init__(PipelineState.RUNNING) def wait_until_finish(self, duration=None) -> str: try: if duration is not None: # Convert milliseconds to seconds duration /= 1000 for _ in ddist.as_completed(self.futures, timeout=duration, with_results=True): # without gathering results, worker errors are not raised on the client: # https://distributed.dask.org/en/stable/resilience.html#user-code-failures # so we want to gather results to raise errors client-side, but we do # not actually need to use the results here, so we just pass. to gather, # we use the iterative `as_completed(..., with_results=True)`, instead # of aggregate `client.gather`, to minimize memory footprint of results. pass self._state = PipelineState.DONE except: # pylint: disable=broad-except self._state = PipelineState.FAILED raise return self._state def cancel(self) -> str: self._state = PipelineState.CANCELLING self.client.cancel(self.futures) self._state = PipelineState.CANCELLED return self._state def metrics(self): # TODO(alxr): Collect and return metrics... raise NotImplementedError('collecting metrics will come later!') class DaskRunner(BundleBasedDirectRunner): """Executes a pipeline on a Dask distributed client.""" @staticmethod def to_dask_bag_visitor(bag_kwargs=None) -> PipelineVisitor: from dask import bag as db if bag_kwargs is None: bag_kwargs = {} @dataclasses.dataclass class DaskBagVisitor(PipelineVisitor): bags: dict[AppliedPTransform, db.Bag] = dataclasses.field( default_factory=collections.OrderedDict) def visit_transform(self, transform_node: AppliedPTransform) -> None: op_class = TRANSLATIONS.get(transform_node.transform.__class__, NoOp) op = op_class(transform_node, bag_kwargs=bag_kwargs) op_kws = {"input_bag": None, "side_inputs": None} inputs = list(transform_node.inputs) if inputs: bag_inputs = [] for input_value in inputs: if isinstance(input_value, pvalue.PBegin): bag_inputs.append(None) prev_op = input_value.producer if prev_op in self.bags: bag_inputs.append(self.bags[prev_op]) # Input to `Flatten` could be of length 1, e.g. a single-element # tuple: `(pcoll, ) | beam.Flatten()`. If so, we still pass it as # an iterable, because `Flatten.apply` always takes an iterable. if len(bag_inputs) == 1 and not isinstance(op, Flatten): op_kws["input_bag"] = bag_inputs[0] else: op_kws["input_bag"] = bag_inputs side_inputs = list(transform_node.side_inputs) if side_inputs: bag_side_inputs = [] for si in side_inputs: si_asbag = self.bags.get(si.pvalue.producer) bag_side_inputs.append( SideInputMap( type(si), si._view_options(), DaskBagWindowedIterator(si_asbag, si._window_mapping_fn))) op_kws["side_inputs"] = bag_side_inputs self.bags[transform_node] = op.apply(**op_kws) return DaskBagVisitor() @staticmethod def is_fnapi_compatible(): return False def run_pipeline(self, pipeline, options): import dask # TODO(alxmrs): Create interactive notebook support. if is_in_notebook(): raise NotImplementedError('interactive support will come later!') try: import dask.distributed as ddist except ImportError: raise ImportError( 'DaskRunner is not available. Please install apache_beam[dask].') dask_options = options.view_as(DaskOptions).get_all_options( drop_default=True, current_only=True) bag_kwargs = DaskOptions._extract_bag_kwargs(dask_options) client = ddist.Client(**dask_options) pipeline.replace_all(dask_overrides()) dask_visitor = self.to_dask_bag_visitor(bag_kwargs) pipeline.visit(dask_visitor) # The dictionary in this visitor keeps a mapping of every Beam # PTransform to the equivalent Bag operation. This is highly # redundant. Thus, we can get away with computing just the last # value, which should be connected to the full Bag Task Graph. opt_graph = dask.optimize(list(dask_visitor.bags.values())[-1]) futures = client.compute(opt_graph) return DaskRunnerResult(client, futures) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dask/dask_runner_test.py0000644000175100017510000003407715206113032024232 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import datetime import inspect import typing as t import unittest import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing import test_pipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import window try: import dask import dask.distributed as ddist from apache_beam.runners.dask.dask_runner import DaskOptions # pylint: disable=ungrouped-imports from apache_beam.runners.dask.dask_runner import DaskRunner # pylint: disable=ungrouped-imports except (ImportError, ModuleNotFoundError): raise unittest.SkipTest('Dask must be installed to run tests.') class DaskOptionsTest(unittest.TestCase): def test_parses_connection_timeout__defaults_to_none(self): default_options = PipelineOptions([]) default_dask_options = default_options.view_as(DaskOptions) self.assertEqual(None, default_dask_options.timeout) def test_parses_connection_timeout__parses_int(self): conn_options = PipelineOptions('--dask_connection_timeout 12'.split()) dask_conn_options = conn_options.view_as(DaskOptions) self.assertEqual(12, dask_conn_options.timeout) def test_parses_connection_timeout__handles_bad_input(self): err_options = PipelineOptions('--dask_connection_timeout foo'.split()) dask_err_options = err_options.view_as(DaskOptions) self.assertEqual(dask.config.no_default, dask_err_options.timeout) def test_parser_destinations__agree_with_dask_client(self): options = PipelineOptions( '--dask_client_address localhost:8080 --dask_connection_timeout 600 ' '--dask_scheduler_file foobar.cfg --dask_client_name charlie ' '--dask_connection_limit 1024'.split()) dask_options = options.view_as(DaskOptions) # Get the argument names for the constructor. client_args = list(inspect.signature(ddist.Client).parameters) for opt_name in dask_options.get_all_options(drop_default=True).keys(): with self.subTest(f'{opt_name} in dask.distributed.Client constructor'): self.assertIn(opt_name, client_args) def test_parser_extract_bag_kwargs__deletes_dask_kwargs(self): options = PipelineOptions('--dask_npartitions 8'.split()) dask_options = options.view_as(DaskOptions).get_all_options() self.assertIn('npartitions', dask_options) bag_kwargs = DaskOptions._extract_bag_kwargs(dask_options) self.assertNotIn('npartitions', dask_options) self.assertEqual(bag_kwargs, {'npartitions': 8}) def test_parser_extract_bag_kwargs__unconfigured(self): options = PipelineOptions() dask_options = options.view_as(DaskOptions).get_all_options() # It's present as a default option. self.assertIn('npartitions', dask_options) bag_kwargs = DaskOptions._extract_bag_kwargs(dask_options) self.assertNotIn('npartitions', dask_options) self.assertEqual(bag_kwargs, {}) class DaskRunnerRunPipelineTest(unittest.TestCase): """Test class used to introspect the dask runner via a debugger.""" def setUp(self) -> None: self.pipeline = test_pipeline.TestPipeline(runner=DaskRunner()) def test_create(self): with self.pipeline as p: pcoll = p | beam.Create([1]) assert_that(pcoll, equal_to([1])) def test_create_multiple(self): with self.pipeline as p: pcoll = p | beam.Create([1, 2, 3, 4]) assert_that(pcoll, equal_to([1, 2, 3, 4])) def test_create_and_map(self): def double(x): return x * 2 with self.pipeline as p: pcoll = p | beam.Create([1]) | beam.Map(double) assert_that(pcoll, equal_to([2])) def test_create_and_map_multiple(self): def double(x): return x * 2 with self.pipeline as p: pcoll = p | beam.Create([1, 2]) | beam.Map(double) assert_that(pcoll, equal_to([2, 4])) def test_create_and_map_many(self): def double(x): return x * 2 with self.pipeline as p: pcoll = p | beam.Create(list(range(1, 11))) | beam.Map(double) assert_that(pcoll, equal_to(list(range(2, 21, 2)))) def test_create_map_and_groupby(self): def double(x): return x * 2, x with self.pipeline as p: pcoll = p | beam.Create([1]) | beam.Map(double) | beam.GroupByKey() assert_that(pcoll, equal_to([(2, [1])])) def test_create_map_and_groupby_multiple(self): def double(x): return x * 2, x with self.pipeline as p: pcoll = ( p | beam.Create([1, 2, 1, 2, 3]) | beam.Map(double) | beam.GroupByKey()) assert_that(pcoll, equal_to([(2, [1, 1]), (4, [2, 2]), (6, [3])])) def test_map_with_positional_side_input(self): def mult_by(x, y): return x * y with self.pipeline as p: side = p | "side" >> beam.Create([3]) pcoll = ( p | "main" >> beam.Create([1]) | beam.Map(mult_by, beam.pvalue.AsSingleton(side))) assert_that(pcoll, equal_to([3])) def test_map_with_keyword_side_input(self): def mult_by(x, y): return x * y with self.pipeline as p: side = p | "side" >> beam.Create([3]) pcoll = ( p | "main" >> beam.Create([1]) | beam.Map(mult_by, y=beam.pvalue.AsSingleton(side))) assert_that(pcoll, equal_to([3])) def test_pardo_side_inputs(self): def cross_product(elem, sides): for side in sides: yield elem, side with self.pipeline as p: main = p | "main" >> beam.Create(["a", "b", "c"]) side = p | "side" >> beam.Create(["x", "y"]) assert_that( main | beam.FlatMap(cross_product, beam.pvalue.AsList(side)), equal_to([ ("a", "x"), ("b", "x"), ("c", "x"), ("a", "y"), ("b", "y"), ("c", "y"), ]), ) def test_pardo_side_input_dependencies(self): with self.pipeline as p: inputs = [p | beam.Create([None])] for k in range(1, 10): inputs.append( inputs[0] | beam.ParDo( ExpectingSideInputsFn(f"Do{k}"), *[beam.pvalue.AsList(inputs[s]) for s in range(1, k)], )) def test_pardo_side_input_sparse_dependencies(self): with self.pipeline as p: inputs = [] def choose_input(s): return inputs[(389 + s * 5077) % len(inputs)] for k in range(20): num_inputs = int((k * k % 16)**0.5) if num_inputs == 0: inputs.append(p | f"Create{k}" >> beam.Create([f"Create{k}"])) else: inputs.append( choose_input(0) | beam.ParDo( ExpectingSideInputsFn(f"Do{k}"), *[ beam.pvalue.AsList(choose_input(s)) for s in range(1, num_inputs) ], )) @unittest.expectedFailure def test_pardo_windowed_side_inputs(self): with self.pipeline as p: # Now with some windowing. pcoll = ( p | beam.Create(list(range(10))) | beam.Map(lambda t: window.TimestampedValue(t, t))) # Intentionally choosing non-aligned windows to highlight the transition. main = pcoll | "WindowMain" >> beam.WindowInto(window.FixedWindows(5)) side = pcoll | "WindowSide" >> beam.WindowInto(window.FixedWindows(7)) res = main | beam.Map( lambda x, s: (x, sorted(s)), beam.pvalue.AsList(side)) assert_that( res, equal_to([ # The window [0, 5) maps to the window [0, 7). (0, list(range(7))), (1, list(range(7))), (2, list(range(7))), (3, list(range(7))), (4, list(range(7))), # The window [5, 10) maps to the window [7, 14). (5, list(range(7, 10))), (6, list(range(7, 10))), (7, list(range(7, 10))), (8, list(range(7, 10))), (9, list(range(7, 10))), ]), label="windowed", ) def test_flattened_side_input(self, with_transcoding=True): with self.pipeline as p: main = p | "main" >> beam.Create([None]) side1 = p | "side1" >> beam.Create([("a", 1)]) side2 = p | "side2" >> beam.Create([("b", 2)]) if with_transcoding: # Also test non-matching coder types (transcoding required) third_element = [("another_type")] else: third_element = [("b", 3)] side3 = p | "side3" >> beam.Create(third_element) side = (side1, side2) | beam.Flatten() assert_that( main | beam.Map(lambda a, b: (a, b), beam.pvalue.AsDict(side)), equal_to([(None, { "a": 1, "b": 2 })]), label="CheckFlattenAsSideInput", ) assert_that( (side, side3) | "FlattenAfter" >> beam.Flatten(), equal_to([("a", 1), ("b", 2)] + third_element), label="CheckFlattenOfSideInput", ) def test_gbk_side_input(self): with self.pipeline as p: main = p | "main" >> beam.Create([None]) side = p | "side" >> beam.Create([("a", 1)]) | beam.GroupByKey() assert_that( main | beam.Map(lambda a, b: (a, b), beam.pvalue.AsDict(side)), equal_to([(None, { "a": [1] })]), ) def test_multimap_side_input(self): with self.pipeline as p: main = p | "main" >> beam.Create(["a", "b"]) side = p | "side" >> beam.Create([("a", 1), ("b", 2), ("a", 3)]) assert_that( main | beam.Map( lambda k, d: (k, sorted(d[k])), beam.pvalue.AsMultiMap(side)), equal_to([("a", [1, 3]), ("b", [2])]), ) def test_multimap_multiside_input(self): # A test where two transforms in the same stage consume the same PCollection # twice as side input. with self.pipeline as p: main = p | "main" >> beam.Create(["a", "b"]) side = p | "side" >> beam.Create([("a", 1), ("b", 2), ("a", 3)]) assert_that( main | "first map" >> beam.Map( lambda k, d, l: (k, sorted(d[k]), sorted([e[1] for e in l])), beam.pvalue.AsMultiMap(side), beam.pvalue.AsList(side), ) | "second map" >> beam.Map( lambda k, d, l: (k[0], sorted(d[k[0]]), sorted([e[1] for e in l])), beam.pvalue.AsMultiMap(side), beam.pvalue.AsList(side), ), equal_to([("a", [1, 3], [1, 2, 3]), ("b", [2], [1, 2, 3])]), ) def test_multimap_side_input_type_coercion(self): with self.pipeline as p: main = p | "main" >> beam.Create(["a", "b"]) # The type of this side-input is forced to Any (overriding type # inference). Without type coercion to Tuple[Any, Any], the usage of this # side-input in AsMultiMap() below should fail. side = p | "side" >> beam.Create([("a", 1), ("b", 2), ("a", 3)]).with_output_types(t.Any) assert_that( main | beam.Map( lambda k, d: (k, sorted(d[k])), beam.pvalue.AsMultiMap(side)), equal_to([("a", [1, 3]), ("b", [2])]), ) def test_pardo_unfusable_side_inputs__one(self): def cross_product(elem, sides): for side in sides: yield elem, side with self.pipeline as p: pcoll = p | "Create1" >> beam.Create(["a", "b"]) assert_that( pcoll | "FlatMap1" >> beam.FlatMap(cross_product, beam.pvalue.AsList(pcoll)), equal_to([("a", "a"), ("a", "b"), ("b", "a"), ("b", "b")]), label="assert_that1", ) def test_pardo_unfusable_side_inputs__two(self): def cross_product(elem, sides): for side in sides: yield elem, side with self.pipeline as p: pcoll = p | "Create2" >> beam.Create(["a", "b"]) derived = ((pcoll, ) | beam.Flatten() | beam.Map(lambda x: (x, x)) | beam.GroupByKey() | "Unkey" >> beam.Map(lambda kv: kv[0])) assert_that( pcoll | "FlatMap2" >> beam.FlatMap( cross_product, beam.pvalue.AsList(derived)), equal_to([("a", "a"), ("a", "b"), ("b", "a"), ("b", "b")]), label="assert_that2", ) def test_groupby_with_fixed_windows(self): def double(x): return x * 2, x def add_timestamp(pair): delta = datetime.timedelta(seconds=pair[1] * 60) now = (datetime.datetime.now() + delta).timestamp() return window.TimestampedValue(pair, now) with self.pipeline as p: pcoll = ( p | beam.Create([1, 2, 1, 2, 3]) | beam.Map(double) | beam.WindowInto(window.FixedWindows(60)) | beam.Map(add_timestamp) | beam.GroupByKey()) assert_that(pcoll, equal_to([(2, [1, 1]), (4, [2, 2]), (6, [3])])) def test_groupby_string_keys(self): with self.pipeline as p: pcoll = ( p | beam.Create([('a', 1), ('a', 2), ('b', 3), ('b', 4)]) | beam.GroupByKey()) assert_that(pcoll, equal_to([('a', [1, 2]), ('b', [3, 4])])) class ExpectingSideInputsFn(beam.DoFn): def __init__(self, name): self._name = name def default_label(self): return self._name def process(self, element, *side_inputs): if not all(list(s) for s in side_inputs): raise ValueError(f"Missing data in side input {side_inputs}") yield self._name if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dask/overrides.py0000644000175100017510000001310215206113032022644 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import dataclasses import typing as t import apache_beam as beam from apache_beam import typehints from apache_beam.io.iobase import SourceBase from apache_beam.pipeline import AppliedPTransform from apache_beam.pipeline import PTransformOverride from apache_beam.runners.direct.direct_runner import _GroupAlsoByWindowDoFn from apache_beam.transforms import ptransform from apache_beam.transforms.window import GlobalWindows K = t.TypeVar("K") V = t.TypeVar("V") @dataclasses.dataclass class _Create(beam.PTransform): values: tuple[t.Any] def expand(self, input_or_inputs): return beam.pvalue.PCollection.from_(input_or_inputs) def get_windowing(self, inputs: t.Any) -> beam.Windowing: return beam.Windowing(GlobalWindows()) @typehints.with_input_types(K) @typehints.with_output_types(K) class _Reshuffle(beam.PTransform): def expand(self, input_or_inputs): return beam.pvalue.PCollection.from_(input_or_inputs) @dataclasses.dataclass class _Read(beam.PTransform): source: SourceBase def expand(self, input_or_inputs): return beam.pvalue.PCollection.from_(input_or_inputs) @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types(tuple[K, t.Iterable[V]]) class _GroupByKeyOnly(beam.PTransform): def expand(self, input_or_inputs): return beam.pvalue.PCollection.from_(input_or_inputs) def infer_output_type(self, input_type): key_type, value_type = typehints.trivial_inference.key_value_types( input_type ) return typehints.KV[key_type, typehints.Iterable[value_type]] @typehints.with_input_types(tuple[K, t.Iterable[V]]) @typehints.with_output_types(tuple[K, t.Iterable[V]]) class _GroupAlsoByWindow(beam.ParDo): def __init__(self, windowing): super().__init__(_GroupAlsoByWindowDoFn(windowing)) self.windowing = windowing def expand(self, input_or_inputs): return beam.pvalue.PCollection.from_(input_or_inputs) @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types(tuple[K, t.Iterable[V]]) class _GroupByKey(beam.PTransform): def expand(self, input_or_inputs): return ( input_or_inputs | "ReifyWindows" >> beam.ParDo(beam.GroupByKey.ReifyWindows()) | "GroupByKey" >> _GroupByKeyOnly() | "GroupByWindow" >> _GroupAlsoByWindow(input_or_inputs.windowing)) class _Flatten(beam.PTransform): def expand(self, input_or_inputs): if isinstance(input_or_inputs, beam.PCollection): # NOTE(cisaacstern): I needed this to avoid # `TypeError: 'PCollection' object is not iterable` # being raised by `all(...)` call below for single-element flattens, i.e., # `(pcoll, ) | beam.Flatten() | ...` is_bounded = input_or_inputs.is_bounded else: is_bounded = all(pcoll.is_bounded for pcoll in input_or_inputs) return beam.pvalue.PCollection(self.pipeline, is_bounded=is_bounded) def dask_overrides() -> list[PTransformOverride]: class CreateOverride(PTransformOverride): def matches(self, applied_ptransform: AppliedPTransform) -> bool: return applied_ptransform.transform.__class__ == beam.Create def get_replacement_transform_for_applied_ptransform( self, applied_ptransform: AppliedPTransform) -> ptransform.PTransform: return _Create(t.cast(beam.Create, applied_ptransform.transform).values) class ReshuffleOverride(PTransformOverride): def matches(self, applied_ptransform: AppliedPTransform) -> bool: return applied_ptransform.transform.__class__ == beam.Reshuffle def get_replacement_transform_for_applied_ptransform( self, applied_ptransform: AppliedPTransform) -> ptransform.PTransform: return _Reshuffle() class ReadOverride(PTransformOverride): def matches(self, applied_ptransform: AppliedPTransform) -> bool: return applied_ptransform.transform.__class__ == beam.io.Read def get_replacement_transform_for_applied_ptransform( self, applied_ptransform: AppliedPTransform) -> ptransform.PTransform: return _Read(t.cast(beam.io.Read, applied_ptransform.transform).source) class GroupByKeyOverride(PTransformOverride): def matches(self, applied_ptransform: AppliedPTransform) -> bool: return applied_ptransform.transform.__class__ == beam.GroupByKey def get_replacement_transform_for_applied_ptransform( self, applied_ptransform: AppliedPTransform) -> ptransform.PTransform: return _GroupByKey() class FlattenOverride(PTransformOverride): def matches(self, applied_ptransform: AppliedPTransform) -> bool: return applied_ptransform.transform.__class__ == beam.Flatten def get_replacement_transform_for_applied_ptransform( self, applied_ptransform: AppliedPTransform) -> ptransform.PTransform: return _Flatten() return [ CreateOverride(), ReshuffleOverride(), ReadOverride(), GroupByKeyOverride(), FlattenOverride(), ] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dask/transform_evaluator.py0000644000175100017510000002146215206113032024747 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Transform Beam PTransforms into Dask Bag operations. A minimum set of operation substitutions, to adap Beam's PTransform model to Dask Bag functions. """ import abc import dataclasses import logging import math import typing as t from dataclasses import field import dask.bag as db import apache_beam from apache_beam import DoFn from apache_beam import TaggedOutput from apache_beam.pipeline import AppliedPTransform from apache_beam.runners.common import DoFnContext from apache_beam.runners.common import DoFnInvoker from apache_beam.runners.common import DoFnSignature from apache_beam.runners.common import Receiver from apache_beam.runners.common import _OutputHandler from apache_beam.runners.dask.overrides import _Create from apache_beam.runners.dask.overrides import _Flatten from apache_beam.runners.dask.overrides import _GroupByKeyOnly from apache_beam.transforms.sideinputs import SideInputMap from apache_beam.transforms.window import GlobalWindow from apache_beam.transforms.window import TimestampedValue from apache_beam.transforms.window import WindowFn from apache_beam.utils.windowed_value import WindowedValue # Inputs to DaskOps. OpInput = t.Union[db.Bag, t.Sequence[db.Bag], None] OpSide = t.Optional[t.Sequence[SideInputMap]] # Value types for PCollections (possibly Windowed Values). PCollVal = t.Union[WindowedValue, t.Any] _LOGGER = logging.getLogger(__name__) def get_windowed_value(item: t.Any, window_fn: WindowFn) -> WindowedValue: """Wraps a value (item) inside a Window.""" if isinstance(item, TaggedOutput): item = item.value if isinstance(item, WindowedValue): windowed_value = item elif isinstance(item, TimestampedValue): assign_context = WindowFn.AssignContext(item.timestamp, item.value) windowed_value = WindowedValue( item.value, item.timestamp, tuple(window_fn.assign(assign_context))) else: windowed_value = WindowedValue(item, 0, (GlobalWindow(), )) return windowed_value def defenestrate(x): """Extracts the underlying item from a Window.""" if isinstance(x, WindowedValue): return x.value return x @dataclasses.dataclass class DaskBagWindowedIterator: """Iterator for `apache_beam.transforms.sideinputs.SideInputMap`""" bag: db.Bag window_fn: WindowFn def __iter__(self): # FIXME(cisaacstern): list() is likely inefficient, since it presumably # materializes the full result before iterating over it. doing this for # now as a proof-of-concept. can we can generate results incrementally? for result in list(self.bag): yield get_windowed_value(result, self.window_fn) @dataclasses.dataclass class TaggingReceiver(Receiver): """A Receiver that handles tagged `WindowValue`s.""" tag: str values: list[PCollVal] def receive(self, windowed_value: WindowedValue): if self.tag: output = TaggedOutput(self.tag, windowed_value) else: output = windowed_value self.values.append(output) @dataclasses.dataclass class OneReceiver(dict): """A Receiver that tags value via dictionary lookup key.""" values: list[PCollVal] = field(default_factory=list) def __missing__(self, key): if key not in self: self[key] = TaggingReceiver(key, self.values) return self[key] @dataclasses.dataclass class DaskBagOp(abc.ABC): """Abstract Base Class for all Dask-supported Operations. All DaskBagOps must support an `apply()` operation, which invokes the dask bag upon the previous op's input. Attributes applied: The underlying `AppliedPTransform` which holds the code for the target operation. bag_kwargs: (optional) Keyword arguments applied to input bags, usually from the pipeline's `DaskOptions`. """ applied: AppliedPTransform bag_kwargs: dict = dataclasses.field(default_factory=dict) @property def transform(self): return self.applied.transform @abc.abstractmethod def apply(self, input_bag: OpInput, side_inputs: OpSide = None) -> db.Bag: pass class NoOp(DaskBagOp): """An identity on a dask bag: returns the input as-is.""" def apply(self, input_bag: OpInput, side_inputs: OpSide = None) -> db.Bag: return input_bag class Create(DaskBagOp): """The beginning of a Beam pipeline; the input must be `None`.""" def apply(self, input_bag: OpInput, side_inputs: OpSide = None) -> db.Bag: assert input_bag is None, 'Create expects no input!' original_transform = t.cast(_Create, self.transform) items = original_transform.values npartitions = self.bag_kwargs.get('npartitions') partition_size = self.bag_kwargs.get('partition_size') if npartitions and partition_size: raise ValueError( f'Please specify either `dask_npartitions` or ' f'`dask_parition_size` but not both: ' f'{npartitions=}, {partition_size=}.') if not npartitions and not partition_size: # partition_size is inversely related to `npartitions`. # Ideal "chunk sizes" in dask are around 10-100 MBs. # Let's hope ~128 items per partition is around this # memory overhead. default_size = 128 partition_size = max(default_size, math.ceil(math.sqrt(len(items)) / 10)) if partition_size == default_size: _LOGGER.warning( 'The new default partition size is %d, it used to be 1 ' 'in previous DaskRunner versions.' % default_size) return db.from_sequence( items, npartitions=npartitions, partition_size=partition_size) def apply_dofn_to_bundle( items, do_fn_invoker_args, do_fn_invoker_kwargs, tagged_receivers): """Invokes a DoFn within a bundle, implemented as a Dask partition.""" do_fn_invoker = DoFnInvoker.create_invoker( *do_fn_invoker_args, **do_fn_invoker_kwargs) do_fn_invoker.invoke_setup() do_fn_invoker.invoke_start_bundle() for it in items: do_fn_invoker.invoke_process(it) results = [v.value for v in tagged_receivers.values] do_fn_invoker.invoke_finish_bundle() do_fn_invoker.invoke_teardown() return results class ParDo(DaskBagOp): """Apply a pure function in an embarrassingly-parallel way. This consumes a sequence of items and returns a sequence of items. """ def apply(self, input_bag: db.Bag, side_inputs: OpSide = None) -> db.Bag: transform = t.cast(apache_beam.ParDo, self.transform) args, kwargs = transform.raw_side_inputs args = list(args) main_input = next(iter(self.applied.main_inputs.values())) window_fn = main_input.windowing.windowfn if hasattr( main_input, "windowing") else None tagged_receivers = OneReceiver() do_fn_invoker_args = [ DoFnSignature(transform.fn), _OutputHandler( window_fn=window_fn, main_receivers=tagged_receivers[None], tagged_receivers=tagged_receivers, per_element_output_counter=None, output_batch_converter=None, process_yields_batches=False, process_batch_yields_elements=False), ] do_fn_invoker_kwargs = dict( context=DoFnContext(transform.label, state=None), side_inputs=side_inputs, input_args=args, input_kwargs=kwargs, user_state_context=None, bundle_finalizer_param=DoFn.BundleFinalizerParam(), ) return input_bag.map(get_windowed_value, window_fn).map_partitions( apply_dofn_to_bundle, do_fn_invoker_args, do_fn_invoker_kwargs, tagged_receivers, ) class GroupByKey(DaskBagOp): """Group a PCollection into a mapping of keys to elements.""" def apply(self, input_bag: db.Bag, side_inputs: OpSide = None) -> db.Bag: def key(item): return item[0] def value(item): k, v = item return k, [defenestrate(elm[1]) for elm in v] return input_bag.groupby(key).map(value) class Flatten(DaskBagOp): """Produces a flattened bag from a collection of bags.""" def apply( self, input_bag: list[db.Bag], side_inputs: OpSide = None) -> db.Bag: assert isinstance(input_bag, list), 'Must take a sequence of bags!' return db.concat(input_bag) TRANSLATIONS = { _Create: Create, apache_beam.ParDo: ParDo, _GroupByKeyOnly: GroupByKey, _Flatten: Flatten, } ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5804288 apache_beam-2.74.0/apache_beam/runners/dataflow/0000755000175100017510000000000015206113160021154 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/__init__.py0000644000175100017510000000215715206113032023270 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """The DataflowRunner executes pipelines on Google Cloud Dataflow. Anything in this package not imported here is an internal implementation detail with no backwards-compatibility guarantees. """ from apache_beam.runners.dataflow.dataflow_runner import DataflowRunner from apache_beam.runners.dataflow.test_dataflow_runner import TestDataflowRunner ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/dataflow_exercise_metrics_pipeline.py0000644000175100017510000001326115206113032030632 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A word-counting workflow.""" # pytype: skip-file import time from hamcrest.library.number.ordering_comparison import greater_than import apache_beam as beam from apache_beam.metrics import Metrics from apache_beam.testing.metric_result_matchers import DistributionMatcher from apache_beam.testing.metric_result_matchers import MetricResultMatcher SLEEP_TIME_SECS = 1 INPUT = [0, 0, 0, 100] METRIC_NAMESPACE = ( 'apache_beam.runners.dataflow.' 'dataflow_exercise_metrics_pipeline.UserMetricsDoFn') def metric_matchers(): """MetricResult matchers common to all tests.""" # TODO(ajamato): Matcher for the 'metrics' step's ElementCount. # TODO(ajamato): Matcher for the 'metrics' step's MeanByteCount. # TODO(ajamato): Matcher for the start and finish exec times. # TODO(ajamato): Matcher for a gauge metric once implemented in dataflow. matchers = [ # User Counter Metrics. MetricResultMatcher( name='total_values', namespace=METRIC_NAMESPACE, step='metrics', attempted=sum(INPUT), committed=sum(INPUT)), MetricResultMatcher( name='ExecutionTime_StartBundle', step='metrics', attempted=greater_than(0), committed=greater_than(0)), MetricResultMatcher( name='ExecutionTime_ProcessElement', step='metrics', attempted=greater_than(0), committed=greater_than(0)), MetricResultMatcher( name='ExecutionTime_FinishBundle', step='metrics', attempted=greater_than(0), committed=greater_than(0)), MetricResultMatcher( name='distribution_values', namespace=METRIC_NAMESPACE, step='metrics', attempted=DistributionMatcher( sum_value=sum(INPUT), count_value=len(INPUT), min_value=min(INPUT), max_value=max(INPUT)), committed=DistributionMatcher( sum_value=sum(INPUT), count_value=len(INPUT), min_value=min(INPUT), max_value=max(INPUT)), ), # Element count and MeanByteCount for a User ParDo. MetricResultMatcher( name='ElementCount', labels={ 'output_user_name': 'metrics-out0', 'original_name': 'metrics-out0-ElementCount' }, attempted=greater_than(0), committed=greater_than(0)), MetricResultMatcher( name='MeanByteCount', labels={ 'output_user_name': 'metrics-out0', 'original_name': 'metrics-out0-MeanByteCount' }, attempted=greater_than(0), committed=greater_than(0)) ] pcoll_names = [ 'GroupByKey/Reify-out0', 'GroupByKey/Read-out0', 'map_to_common_key-out0', 'GroupByKey/GroupByWindow-out0', 'GroupByKey/Read-out0', 'GroupByKey/Reify-out0' ] for name in pcoll_names: matchers.extend([ MetricResultMatcher( name='ElementCount', labels={ 'output_user_name': name, 'original_name': '%s-ElementCount' % name }, attempted=greater_than(0), committed=greater_than(0)), MetricResultMatcher( name='MeanByteCount', labels={ 'output_user_name': name, 'original_name': '%s-MeanByteCount' % name }, attempted=greater_than(0), committed=greater_than(0)), ]) return matchers class UserMetricsDoFn(beam.DoFn): """Parse each line of input text into words.""" def __init__(self): self.total_metric = Metrics.counter(self.__class__, 'total_values') self.dist_metric = Metrics.distribution( self.__class__, 'distribution_values') # TODO(ajamato): Add a verifier for gauge once it is supported by the SDKs # and runners. self.latest_metric = Metrics.gauge(self.__class__, 'latest_value') def start_bundle(self): time.sleep(SLEEP_TIME_SECS) def process(self, element): """Returns the processed element and increments the metrics.""" elem_int = int(element) self.total_metric.inc(elem_int) self.dist_metric.update(elem_int) self.latest_metric.set(elem_int) time.sleep(SLEEP_TIME_SECS) return [elem_int] def finish_bundle(self): time.sleep(SLEEP_TIME_SECS) def apply_and_run(pipeline): """Given an initialized Pipeline applies transforms and runs it.""" _ = ( pipeline | beam.Create(INPUT) | 'metrics' >> (beam.ParDo(UserMetricsDoFn())) | 'map_to_common_key' >> beam.Map(lambda x: ('key', x)) | beam.GroupByKey() | 'm_out' >> beam.FlatMap( lambda x: [ 1, 2, 3, 4, 5, beam.pvalue.TaggedOutput('once', x), beam.pvalue. TaggedOutput('twice', x), beam.pvalue.TaggedOutput('twice', x) ])) result = pipeline.run() result.wait_until_finish() return result ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/dataflow_exercise_metrics_pipeline_test.py0000644000175100017510000000377515206113032031702 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A word-counting workflow.""" # pytype: skip-file import argparse import unittest import pytest import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.runners.dataflow import dataflow_exercise_metrics_pipeline from apache_beam.testing import metric_result_matchers from apache_beam.testing.test_pipeline import TestPipeline class ExerciseMetricsPipelineTest(unittest.TestCase): def run_pipeline(self, **opts): test_pipeline = TestPipeline(is_integration_test=True) argv = test_pipeline.get_full_options_as_args(**opts) parser = argparse.ArgumentParser() unused_known_args, pipeline_args = parser.parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) p = beam.Pipeline(options=pipeline_options) return dataflow_exercise_metrics_pipeline.apply_and_run(p) @pytest.mark.it_postcommit @unittest.skip('https://github.com/apache/beam/issues/22605') def test_metrics_it(self): result = self.run_pipeline() errors = metric_result_matchers.verify_all( result.metrics().all_metrics(), dataflow_exercise_metrics_pipeline.metric_matchers()) self.assertFalse(errors, str(errors)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/dataflow_exercise_streaming_metrics_pipeline.py0000644000175100017510000000651315206113032032705 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A word-counting workflow.""" # pytype: skip-file import argparse import logging import time import apache_beam as beam from apache_beam.metrics import Metrics from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions SLEEP_TIME_SECS = 1 _LOGGER = logging.getLogger(__name__) class StreamingUserMetricsDoFn(beam.DoFn): """Generates user metrics and outputs same element.""" def __init__(self): self.double_message_counter = Metrics.counter( self.__class__, 'double_msg_counter_name') self.msg_len_dist_metric = Metrics.distribution( self.__class__, 'msg_len_dist_metric_name') def start_bundle(self): time.sleep(SLEEP_TIME_SECS) def process(self, element): """Returns the processed element and increments the metrics.""" text_line = element.strip() self.double_message_counter.inc() self.double_message_counter.inc() self.msg_len_dist_metric.update(len(text_line)) _LOGGER.debug("Done processing returning element array: '%s'", element) return [element] def finish_bundle(self): time.sleep(SLEEP_TIME_SECS) def run(argv=None): """Given an initialized Pipeline applies transforms and runs it.""" parser = argparse.ArgumentParser() parser.add_argument( '--output_topic', required=True, help=( 'Output PubSub topic of the form ' '"projects//topic/".')) group = parser.add_mutually_exclusive_group(required=True) group.add_argument( '--input_topic', help=( 'Input PubSub topic of the form ' '"projects//topics/".')) group.add_argument( '--input_subscription', help=( 'Input PubSub subscription of the form ' '"projects//subscriptions/."')) known_args, pipeline_args = parser.parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = True pipeline_options.view_as(StandardOptions).streaming = True pipeline = beam.Pipeline(options=pipeline_options) _ = ( pipeline | beam.io.ReadFromPubSub(subscription=known_args.input_subscription) | 'generate_metrics' >> (beam.ParDo(StreamingUserMetricsDoFn())) | 'dump_to_pub' >> beam.io.WriteToPubSub(known_args.output_topic)) result = pipeline.run() result.wait_until_finish() return result if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/dataflow_exercise_streaming_metrics_pipeline_test.py0000644000175100017510000001527315206113032033747 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A word-counting workflow.""" # pytype: skip-file import logging import unittest import uuid import pytest from hamcrest.core.core.allof import all_of from apache_beam.io.gcp.tests.pubsub_matcher import PubSubMessageMatcher from apache_beam.runners.dataflow import dataflow_exercise_streaming_metrics_pipeline from apache_beam.runners.runner import PipelineState from apache_beam.testing import metric_result_matchers from apache_beam.testing import test_utils from apache_beam.testing.metric_result_matchers import DistributionMatcher from apache_beam.testing.metric_result_matchers import MetricResultMatcher from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline INPUT_TOPIC = 'exercise_streaming_metrics_topic_input' INPUT_SUB = 'exercise_streaming_metrics_subscription_input' OUTPUT_TOPIC = 'exercise_streaming_metrics_topic_output' OUTPUT_SUB = 'exercise_streaming_metrics_subscription_output' WAIT_UNTIL_FINISH_DURATION = 5 * 60 * 1000 # in milliseconds MESSAGES_TO_PUBLISH = ["message a", "message b b", "message c"] SLEEP_TIME_SECS = 1 _LOGGER = logging.getLogger(__name__) class ExerciseStreamingMetricsPipelineTest(unittest.TestCase): def setUp(self): """Creates all required topics and subs.""" self.test_pipeline = TestPipeline(is_integration_test=True) self.project = self.test_pipeline.get_option('project') self.uuid = str(uuid.uuid4()) # Set up PubSub environment. from google.cloud import pubsub self.pub_client = pubsub.PublisherClient() self.input_topic_name = INPUT_TOPIC + self.uuid self.input_topic = self.pub_client.create_topic( name=self.pub_client.topic_path(self.project, self.input_topic_name)) self.output_topic_name = OUTPUT_TOPIC + self.uuid self.output_topic = self.pub_client.create_topic( name=self.pub_client.topic_path(self.project, self.output_topic_name)) self.sub_client = pubsub.SubscriberClient() self.input_sub_name = INPUT_SUB + self.uuid self.input_sub = self.sub_client.create_subscription( name=self.sub_client.subscription_path( self.project, self.input_sub_name), topic=self.input_topic.name) self.output_sub_name = OUTPUT_SUB + self.uuid self.output_sub = self.sub_client.create_subscription( name=self.sub_client.subscription_path( self.project, self.output_sub_name), topic=self.output_topic.name, ack_deadline_seconds=60) def _inject_words(self, topic, messages): """Inject messages as test data to PubSub.""" _LOGGER.debug('Injecting messages to topic %s', topic.name) for msg in messages: self.pub_client.publish(self.input_topic.name, msg.encode('utf-8')) _LOGGER.debug('Done. Injecting messages to topic %s', topic.name) def tearDown(self): """Delete all created topics and subs.""" test_utils.cleanup_subscriptions( self.sub_client, [self.input_sub, self.output_sub]) test_utils.cleanup_topics( self.pub_client, [self.input_topic, self.output_topic]) def run_pipeline(self): # Waits for messages to appear in output topic. expected_msg = [msg.encode('utf-8') for msg in MESSAGES_TO_PUBLISH] pubsub_msg_verifier = PubSubMessageMatcher( self.project, self.output_sub.name, expected_msg, timeout=600) # Checks that pipeline initializes to RUNNING state. state_verifier = PipelineStateMatcher(PipelineState.RUNNING) extra_opts = { 'wait_until_finish_duration': WAIT_UNTIL_FINISH_DURATION, 'on_success_matcher': all_of(state_verifier, pubsub_msg_verifier), 'experiment': 'beam_fn_api', 'input_subscription': self.input_sub.name, 'output_topic': self.output_topic.name, } argv = self.test_pipeline.get_full_options_as_args(**extra_opts) return dataflow_exercise_streaming_metrics_pipeline.run(argv) # Need not run streaming test in batch mode. @pytest.mark.it_validatesrunner @pytest.mark.no_sickbay_batch @pytest.mark.no_xdist def test_streaming_pipeline_returns_expected_user_metrics_fnapi_it(self): """ Runs streaming Dataflow job and verifies that user metrics are reported correctly. """ self._inject_words(self.input_topic, MESSAGES_TO_PUBLISH) result = self.run_pipeline() METRIC_NAMESPACE = \ ('apache_beam.runners.dataflow.' 'dataflow_exercise_streaming_metrics_pipeline.StreamingUserMetricsDoFn') matchers = [ # System metrics MetricResultMatcher( name='ElementCount', labels={ "output_user_name": "generate_metrics-out0", "original_name": "generate_metrics-out0-ElementCount" }, attempted=len(MESSAGES_TO_PUBLISH), committed=len(MESSAGES_TO_PUBLISH), ), # User Counter Metrics. MetricResultMatcher( name='double_msg_counter_name', namespace=METRIC_NAMESPACE, step='generate_metrics', attempted=len(MESSAGES_TO_PUBLISH) * 2, committed=len(MESSAGES_TO_PUBLISH) * 2), MetricResultMatcher( name='msg_len_dist_metric_name', namespace=METRIC_NAMESPACE, step='generate_metrics', attempted=DistributionMatcher( sum_value=len(''.join(MESSAGES_TO_PUBLISH)), count_value=len(MESSAGES_TO_PUBLISH), min_value=len(MESSAGES_TO_PUBLISH[0]), max_value=len(MESSAGES_TO_PUBLISH[1])), committed=DistributionMatcher( sum_value=len(''.join(MESSAGES_TO_PUBLISH)), count_value=len(MESSAGES_TO_PUBLISH), min_value=len(MESSAGES_TO_PUBLISH[0]), max_value=len(MESSAGES_TO_PUBLISH[1]))), ] metrics = result.metrics().all_metrics() errors = metric_result_matchers.verify_all(metrics, matchers) self.assertFalse(errors, str(errors)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/dataflow_job_service.py0000644000175100017510000000600215206113032025675 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import argparse import logging import sys from apache_beam.portability.api import beam_job_api_pb2 from apache_beam.runners.dataflow import dataflow_runner from apache_beam.runners.portability import local_job_service from apache_beam.runners.portability import local_job_service_main from apache_beam.runners.portability import portable_runner class DataflowBeamJob(local_job_service.BeamJob): """A representation of a single Beam job to be run on the Dataflow runner. """ def _invoke_runner(self): """Actually calls Dataflow and waits for completion. """ runner = dataflow_runner.DataflowRunner() self.result = runner.run_pipeline( None, self.pipeline_options(), self._pipeline_proto) # The result can be None if there is no need to send a request # to the service (e.g. template creation). if not getattr(self.result, 'has_job', None): self.set_state(beam_job_api_pb2.JobState.DONE) return self.result # Prefer this to result.wait_until_finish() to get state updates # and avoid creating an extra thread (which also messes with logging). dataflow_runner.DataflowRunner.poll_for_job_completion( runner, self.result, None, lambda dataflow_state: self.set_state( portable_runner.PipelineResult.pipeline_state_to_runner_api_state( self.result.api_jobstate_to_pipeline_state(dataflow_state)))) return self.result def cancel(self): if not self.is_terminal_state(self.state): self.result.cancel() def run(argv, beam_job_type=DataflowBeamJob): if argv[0] == __file__: argv = argv[1:] parser = argparse.ArgumentParser() parser.add_argument( '-p', '--port', '--job_port', type=int, default=0, help='port on which to serve the job api') parser.add_argument('--staging_dir') options = parser.parse_args(argv) job_servicer = local_job_service.LocalJobServicer( options.staging_dir, beam_job_type=beam_job_type) port = job_servicer.start_grpc_server(options.port) try: local_job_service_main.serve("Listening for beam jobs on port %d." % port) finally: job_servicer.stop() if __name__ == '__main__': logging.basicConfig() logging.getLogger().setLevel(logging.INFO) run(sys.argv) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/dataflow_job_service_test.py0000644000175100017510000000632415206113032026743 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import time import unittest import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.runners.dataflow import dataflow_job_service from apache_beam.runners.portability import local_job_service # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.runners.dataflow.internal import apiclient except ImportError: apiclient = None # type: ignore # pylint: enable=wrong-import-order, wrong-import-position @unittest.skipIf(apiclient is None, 'GCP dependencies are not installed') class DirectPipelineResultTest(unittest.TestCase): def test_dry_run(self): # Not an integration test that actually runs on Dataflow, # but does exercise (most of) the translation and setup code, # as well as the connection. job_servicer = local_job_service.LocalJobServicer( None, beam_job_type=dataflow_job_service.DataflowBeamJob) port = job_servicer.start_grpc_server(0) try: options = PipelineOptions( runner='PortableRunner', job_endpoint=f'localhost:{port}', project='some_project', temp_location='gs://bucket/dir', region='us-central1', dry_run=True, template_location='gs://bucket/dir', ) with beam.Pipeline(options=options) as p: _ = p | beam.Create([1, 2, 3]) | beam.Map(lambda x: x * x) finally: job_servicer.stop() @unittest.skipIf(apiclient is None, 'GCP dependencies are not installed') class DirectPipelineTemplateTest(unittest.TestCase): def test_template(self): job_servicer = local_job_service.LocalJobServicer( None, beam_job_type=dataflow_job_service.DataflowBeamJob) port = job_servicer.start_grpc_server(0) try: template_path = ( 'gs://apache-beam-testing-temp/test/template-{}'.format(time.time())) options = PipelineOptions( runner='PortableRunner', job_endpoint=f'localhost:{port}', project='apache-beam-testing', region='us-central1', staging_location='gs://apache-beam-testing-stg/stg/', temp_location='gs://apache-beam-testing-temp/tmp', template_location=template_path, ) with beam.Pipeline(options=options) as p: _ = p | beam.Create([1, 2, 3]) | beam.Map(lambda x: x * x) finally: job_servicer.stop() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/dataflow_metrics.py0000644000175100017510000003211015206113032025050 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ DataflowRunner implementation of MetricResults. It is in charge of responding to queries of current metrics by going to the dataflow service. """ # pytype: skip-file import argparse import logging import numbers import sys from collections import defaultdict from apache_beam.metrics.cells import DistributionData from apache_beam.metrics.cells import DistributionResult from apache_beam.metrics.execution import MetricKey from apache_beam.metrics.execution import MetricResult from apache_beam.metrics.metric import MetricResults from apache_beam.metrics.metricbase import MetricName from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions _LOGGER = logging.getLogger(__name__) def _get_match(proto, filter_fn): """Finds and returns the first element that matches a query. If no element matches the query, it throws ValueError. If more than one element matches the query, it returns only the first. """ query = [elm for elm in proto if filter_fn(elm)] if len(query) == 0: raise ValueError('Could not find element') elif len(query) > 1: raise ValueError('Too many matches') return query[0] # V1b3 MetricStructuredName keys to accept and copy to the MetricKey labels. STEP_LABEL = 'step' STRUCTURED_NAME_LABELS = set( ['execution_step', 'original_name', 'output_user_name']) class DataflowMetrics(MetricResults): """Implementation of MetricResults class for the Dataflow runner.""" def __init__(self, dataflow_client=None, job_result=None, job_graph=None): """Initialize the Dataflow metrics object. Args: dataflow_client: apiclient.DataflowApplicationClient to interact with the dataflow service. job_result: DataflowPipelineResult with the state and id information of the job. job_graph: apiclient.Job instance to be able to translate between internal step names (e.g. "s2"), and user step names (e.g. "split"). """ super().__init__() self._dataflow_client = dataflow_client self.job_result = job_result self._queried_after_termination = False self._cached_metrics = None self._job_graph = job_graph @staticmethod def _is_counter(metric_result): return isinstance(metric_result.attempted, numbers.Number) @staticmethod def _is_distribution(metric_result): return isinstance(metric_result.attempted, DistributionResult) @staticmethod def _is_string_set(metric_result): return isinstance(metric_result.attempted, set) def _translate_step_name(self, internal_name): """Translate between internal step names (e.g. "s1") and user step names.""" if not self._job_graph: raise ValueError( 'Could not translate the internal step name %r since job graph is ' 'not available.' % internal_name) user_step_name = None if (self._job_graph and internal_name in self._job_graph.proto_pipeline.components.transforms.keys()): # Dataflow Runner v2 with portable job submission uses proto transform map # IDs for step names. Also PTransform.unique_name maps to user step names. # Hence we lookup user step names based on the proto. user_step_name = self._job_graph.proto_pipeline.components.transforms[ internal_name].unique_name else: try: step = _get_match( self._job_graph.proto.steps, lambda x: x.name == internal_name) user_step_name = _get_match( step.properties.additionalProperties, lambda x: x.key == 'user_name').value.string_value except ValueError: pass # Exception is handled below. if not user_step_name: raise ValueError( 'Could not translate the internal step name %r.' % internal_name) return user_step_name def _get_metric_key(self, metric): """Populate the MetricKey object for a queried metric result.""" step = "" name = metric.name.name # Always extract a name labels = {} try: # Try to extract the user step name. # If ValueError is thrown within this try-block, it is because of # one of the following: # 1. Unable to translate the step name. Only happening with improperly # formatted job graph (unlikely), or step name not being the internal # step name (only happens for unstructured-named metrics). # 2. Unable to unpack [step] or [namespace]; which should only happen # for unstructured names. step = _get_match( metric.name.context.additionalProperties, lambda x: x.key == STEP_LABEL).value step = self._translate_step_name(step) except ValueError: pass namespace = "dataflow/v1b3" # Try to extract namespace or add a default. try: namespace = _get_match( metric.name.context.additionalProperties, lambda x: x.key == 'namespace').value except ValueError: pass for kv in metric.name.context.additionalProperties: if kv.key in STRUCTURED_NAME_LABELS: labels[kv.key] = kv.value # Package everything besides namespace and name the labels as well, # including unmodified step names to assist in integration the exact # unmodified values which come from dataflow. return MetricKey(step, MetricName(namespace, name), labels=labels) def _populate_metrics(self, response, result, user_metrics=False): """Move metrics from response to results as MetricResults.""" if user_metrics: metrics = [ metric for metric in response.metrics if metric.name.origin == 'user' ] else: metrics = [ metric for metric in response.metrics if metric.name.origin == 'dataflow/v1b3' ] # Get the tentative/committed versions of every metric together. metrics_by_name = defaultdict(lambda: {}) for metric in metrics: if (metric.name.name.endswith('_MIN') or metric.name.name.endswith('_MAX') or metric.name.name.endswith('_MEAN') or metric.name.name.endswith('_COUNT')): # The Dataflow Service presents distribution metrics in two ways: # One way is as a single distribution object with all its fields, and # another way is as four different scalar metrics labeled as _MIN, # _MAX, _COUNT_, _MEAN. # TODO(pabloem) remove these when distributions are not being broken up # in the service. # The second way is only useful for the UI, and should be ignored. continue is_tentative = [ prop for prop in metric.name.context.additionalProperties if prop.key == 'tentative' and prop.value == 'true' ] tentative_or_committed = 'tentative' if is_tentative else 'committed' metric_key = self._get_metric_key(metric) if metric_key is None: continue metrics_by_name[metric_key][tentative_or_committed] = metric # Now we create the MetricResult elements. for metric_key, metric in metrics_by_name.items(): attempted = self._get_metric_value(metric['tentative']) committed = self._get_metric_value(metric['committed']) result.append( MetricResult(metric_key, attempted=attempted, committed=committed)) def _get_metric_value(self, metric): """Get a metric result object from a MetricUpdate from Dataflow API.""" if metric is None: return None if metric.scalar is not None: return metric.scalar.integer_value elif metric.distribution is not None: dist_count = _get_match( metric.distribution.object_value.properties, lambda x: x.key == 'count').value.integer_value dist_min = _get_match( metric.distribution.object_value.properties, lambda x: x.key == 'min').value.integer_value dist_max = _get_match( metric.distribution.object_value.properties, lambda x: x.key == 'max').value.integer_value dist_sum = _get_match( metric.distribution.object_value.properties, lambda x: x.key == 'sum').value.integer_value if dist_sum is None: # distribution metric is not meant to use on large values, but in case # it is, the value can overflow and become double_value, the correctness # of the value may not be guaranteed. _LOGGER.info( "Distribution metric sum value seems to have " "overflowed integer_value range, the correctness of sum or mean " "value may not be guaranteed: %s" % metric.distribution) dist_sum = int( _get_match( metric.distribution.object_value.properties, lambda x: x.key == 'sum').value.double_value) return DistributionResult( DistributionData(dist_sum, dist_count, dist_min, dist_max)) #TODO(https://github.com/apache/beam/issues/31788) support StringSet after # re-generate apiclient else: return None def _get_metrics_from_dataflow(self, job_id=None): """Return cached metrics or query the dataflow service.""" if not job_id: try: job_id = self.job_result.job_id() except AttributeError: job_id = None if not job_id: raise ValueError('Can not query metrics. Job id is unknown.') if self._cached_metrics: return self._cached_metrics job_metrics = self._dataflow_client.get_job_metrics(job_id) # If we cannot determine that the job has terminated, # then metrics will not change and we can cache them. if self.job_result and self.job_result.is_in_terminal_state(): self._cached_metrics = job_metrics return job_metrics def all_metrics(self, job_id=None): """Return all user and system metrics from the dataflow service.""" metric_results = [] response = self._get_metrics_from_dataflow(job_id=job_id) self._populate_metrics(response, metric_results, user_metrics=True) self._populate_metrics(response, metric_results, user_metrics=False) return metric_results def query(self, filter=None): metric_results = [] response = self._get_metrics_from_dataflow() self._populate_metrics(response, metric_results, user_metrics=True) return { self.COUNTERS: [ elm for elm in metric_results if self.matches(filter, elm.key) and DataflowMetrics._is_counter(elm) ], self.DISTRIBUTIONS: [ elm for elm in metric_results if self.matches(filter, elm.key) and DataflowMetrics._is_distribution(elm) ], # TODO(pabloem): Add Gauge support for dataflow. self.GAUGES: [], self.STRINGSETS: [ elm for elm in metric_results if self.matches(filter, elm.key) and DataflowMetrics._is_string_set(elm) ] } def main(argv): """Print the metric results for the dataflow --job_id and --project. Instead of running an entire pipeline which takes several minutes, use this main method to display MetricResults for a specific --job_id and --project which takes only a few seconds. """ # TODO(https://github.com/apache/beam/issues/19452): The MetricResults do not # show translated step names as the job_graph is not provided to # DataflowMetrics. Import here to avoid adding the dependency for local # running scenarios. try: # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.runners.dataflow.internal import apiclient except ImportError: raise ImportError( 'Google Cloud Dataflow runner not available, ' 'please install apache_beam[gcp]') if argv[0] == __file__: argv = argv[1:] parser = argparse.ArgumentParser() parser.add_argument( '-j', '--job_id', type=str, help='The job id to query metrics for.') parser.add_argument( '-p', '--project', type=str, help='The project name to query metrics for.') flags = parser.parse_args(argv) # Get a Dataflow API client and set its project and job_id in the options. options = PipelineOptions() gcloud_options = options.view_as(GoogleCloudOptions) gcloud_options.project = flags.project dataflow_client = apiclient.DataflowApplicationClient(options) df_metrics = DataflowMetrics(dataflow_client) all_metrics = df_metrics.all_metrics(job_id=flags.job_id) _LOGGER.info( 'Printing all MetricResults for %s in %s', flags.job_id, flags.project) for metric_result in all_metrics: _LOGGER.info(metric_result) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) main(sys.argv) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/dataflow_metrics_test.py0000644000175100017510000005106715206113032026123 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Tests corresponding to the DataflowRunner implementation of MetricsResult, the DataflowMetrics class. """ # pytype: skip-file # pylint: skip-file import types import unittest import mock from apache_beam import DoFn from apache_beam import ParDo from apache_beam.metrics.cells import DistributionData from apache_beam.metrics.cells import DistributionResult from apache_beam.metrics.execution import MetricKey from apache_beam.metrics.execution import MetricResult from apache_beam.metrics.metricbase import MetricName from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.pipeline import Pipeline from apache_beam.runners.dataflow import dataflow_metrics from apache_beam.testing import metric_result_matchers from apache_beam.testing.metric_result_matchers import MetricResultMatcher from apache_beam.transforms import Create from apache_beam.transforms.environments import DockerEnvironment # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.runners.dataflow.internal import apiclient except ImportError: apiclient = None # type: ignore # pylint: enable=wrong-import-order, wrong-import-position class DictToObject(object): """Translate from a dict(list()) structure to an object structure""" def __init__(self, data): for name, value in data.items(): setattr(self, name, self._wrap(value)) def _wrap(self, value): if isinstance(value, (tuple, list, set, frozenset)): return type(value)([self._wrap(v) for v in value]) return DictToObject(value) if isinstance(value, dict) else value class TestDataflowMetrics(unittest.TestCase): # TODO(https://github.com/apache/beam/issues/19258): Write a dump tool to # generate this fake data, or somehow make this easier to maintain. ONLY_COUNTERS_LIST = { "metrics": [ { "name": { "context": { "additionalProperties": [{ "key": "namespace", "value": "__main__.WordExtractingDoFn" }, { "key": "step", "value": "s2" }, { "key": "tentative", "value": "true" }] }, "name": "words", "origin": "user" }, "scalar": { "integer_value": 26185 }, "distribution": None, "updateTime": "2017-03-22T18:47:06.402Z" }, { "name": { "context": { "additionalProperties": [{ "key": "namespace", "value": "__main__.WordExtractingDoFn" }, { "key": "step", "value": "s2" }] }, "name": "words", "origin": "user" }, "scalar": { "integer_value": 26181 }, "distribution": None, "updateTime": "2017-03-22T18:47:06.402Z" }, { "name": { "context": { "additionalProperties": [{ "key": "namespace", "value": "__main__.WordExtractingDoFn" }, { "key": "step", "value": "s2" }, { "key": "tentative", "value": "true" }] }, "name": "empty_lines", "origin": "user" }, "scalar": { "integer_value": 1080 }, "distribution": None, "updateTime": "2017-03-22T18:47:06.402Z" }, { "name": { "context": { "additionalProperties": [{ "key": "namespace", "value": "__main__.WordExtractingDoFn" }, { "key": "step", "value": "s2" }] }, "name": "empty_lines", "origin": "user" }, "scalar": { "integer_value": 1080 }, "distribution": None, "updateTime": "2017-03-22T18:47:06.402Z" }, ] } STRUCTURED_COUNTER_LIST = { "metrics": [ { "name": { "context": { "additionalProperties": [{ "key": "namespace", "value": "__main__.WordExtractingDoFn" }, { "key": "step", "value": "s2" }, { "key": "tentative", "value": "true" }] }, "name": "word_lengths", "origin": "user" }, "scalar": { "integer_value": 109475 }, "distribution": None, "updateTime": "2017-03-22T18:47:06.402Z" }, { "name": { "context": { "additionalProperties": [{ "key": "namespace", "value": "__main__.WordExtractingDoFn" }, { "key": "step", "value": "s2" }] }, "name": "word_lengths", "origin": "user" }, "scalar": { "integer_value": 109475 }, "distribution": None, "updateTime": "2017-03-22T18:47:06.402Z" }, { "name": { "context": { "additionalProperties": [{ "key": "namespace", "value": "__main__.WordExtractingDoFn" }, { "key": "step", "value": "s2" }, { "key": "tentative", "value": "true" }] }, "name": "word_length_dist", "origin": "user" }, "scalar": None, "distribution": { "object_value": { "properties": [ { "key": "min", "value": { "integer_value": 2 } }, { "key": "max", "value": { "integer_value": 16 } }, { "key": "count", "value": { "integer_value": 2 } }, { "key": "mean", "value": { "integer_value": 9 } }, { "key": "sum", "value": { "integer_value": 18 } }, ] } }, "updateTime": "2017-03-22T18:47:06.402Z" }, { "name": { "context": { "additionalProperties": [{ "key": "namespace", "value": "__main__.WordExtractingDoFn" }, { "key": "step", "value": "s2" }] }, "name": "word_length_dist", "origin": "user" }, "scalar": None, "distribution": { "object_value": { "properties": [ { "key": "min", "value": { "integer_value": 2 } }, { "key": "max", "value": { "integer_value": 16 } }, { "key": "count", "value": { "integer_value": 2 } }, { "key": "mean", "value": { "integer_value": 9 } }, { "key": "sum", "value": { "integer_value": 18 } }, ] } }, "updateTime": "2017-03-22T18:47:06.402Z" }, ] } SYSTEM_COUNTERS_LIST = { "metrics": [ # ElementCount { "name": { "context": { "additionalProperties": [ { "key": "original_name", "value": "ToIsmRecordForMultimap-out0-ElementCount" }, # yapf: disable { "key": "output_user_name", "value": "ToIsmRecordForMultimap-out0" } ] }, "name": "ElementCount", "origin": "dataflow/v1b3" }, "scalar": { "integer_value": 42 }, "distribution": None, "updateTime": "2017-03-22T18:47:06.402Z" }, { "name": { "context": { "additionalProperties": [ { "key": "original_name", "value": "ToIsmRecordForMultimap-out0-ElementCount" }, # yapf: disable { "key": "output_user_name", "value": "ToIsmRecordForMultimap-out0" }, { "key": "tentative", "value": "true" } ] }, "name": "ElementCount", "origin": "dataflow/v1b3" }, "scalar": { "integer_value": 42 }, "distribution": None, "updateTime": "2017-03-22T18:47:06.402Z" }, # MeanByteCount { "name": { "context": { "additionalProperties": [ { "key": "original_name", "value": "Read-out0-MeanByteCount" }, { "key": "output_user_name", "value": "GroupByKey/Read-out0" } ] }, "name": "MeanByteCount", "origin": "dataflow/v1b3" }, "scalar": { "integer_value": 31 }, "distribution": None, "updateTime": "2017-03-22T18:47:06.402Z" }, { "name": { "context": { "additionalProperties": [ { "key": "original_name", "value": "Read-out0-MeanByteCount" }, { "key": "output_user_name", "value": "GroupByKey/Read-out0" }, { "key": "tentative", "value": "true" } ] }, "name": "MeanByteCount", "origin": "dataflow/v1b3" }, "scalar": { "integer_value": 31 }, "distribution": None, "updateTime": "2017-03-22T18:47:06.402Z" }, # ExecutionTime { "name": { "context": { "additionalProperties": [ { "key": "step", "value": "write/Write/Write" }, ] }, "name": "ExecutionTime_ProcessElement", "origin": "dataflow/v1b3" }, "scalar": { "integer_value": 1000 }, "distribution": None, "updateTime": "2017-03-22T18:47:06.402Z" }, { "name": { "context": { "additionalProperties": [{ "key": "step", "value": "write/Write/Write" }, { "key": "tentative", "value": "true" }] }, "name": "ExecutionTime_ProcessElement", "origin": "dataflow/v1b3" }, "scalar": { "integer_value": 1000 }, "distribution": None, "updateTime": "2017-03-22T18:47:06.402Z" }, ] } def setup_mock_client_result(self, counter_list=None): mock_client = mock.Mock() mock_query_result = DictToObject(counter_list) mock_client.get_job_metrics.return_value = mock_query_result mock_job_result = mock.Mock() mock_job_result.job_id.return_value = 1 mock_job_result.is_in_terminal_state.return_value = False return mock_client, mock_job_result def test_cache_functions(self): mock_client, mock_job_result = self.setup_mock_client_result( self.STRUCTURED_COUNTER_LIST) dm = dataflow_metrics.DataflowMetrics(mock_client, mock_job_result) # At first creation, we should always query dataflow. self.assertTrue(dm._cached_metrics is None) # Right after querying, we still query again. dm.query() self.assertTrue(dm._cached_metrics is None) # The job has ended. The query should not run again after this. mock_job_result.is_in_terminal_state.return_value = True dm.query() self.assertTrue(dm._cached_metrics) def test_query_structured_metrics(self): mock_client, mock_job_result = self.setup_mock_client_result( self.STRUCTURED_COUNTER_LIST) dm = dataflow_metrics.DataflowMetrics(mock_client, mock_job_result) dm._translate_step_name = types.MethodType(lambda self, x: 'split', dm) query_result = dm.query() expected_counters = [ MetricResult( MetricKey( 'split', MetricName('__main__.WordExtractingDoFn', 'word_lengths'), ), 109475, 109475), ] self.assertEqual(query_result['counters'], expected_counters) expected_distributions = [ MetricResult( MetricKey( 'split', MetricName('__main__.WordExtractingDoFn', 'word_length_dist'), ), DistributionResult(DistributionData(18, 2, 2, 16)), DistributionResult(DistributionData(18, 2, 2, 16))), ] self.assertEqual(query_result['distributions'], expected_distributions) @unittest.skipIf(apiclient is None, 'GCP dependencies are not installed') def test_translate_portable_job_step_name(self): mock_client, mock_job_result = self.setup_mock_client_result( self.ONLY_COUNTERS_LIST) pipeline_options = PipelineOptions([ '--experiments=use_portable_job_submission', '--temp_location=gs://any-location/temp', '--project=dummy_project', ]) pipeline = Pipeline(options=pipeline_options) pipeline | Create([1, 2, 3]) | 'MyTestParDo' >> ParDo(DoFn()) # pylint:disable=expression-not-assigned test_environment = DockerEnvironment(container_image='test_default_image') proto_pipeline, _ = pipeline.to_runner_api( return_context=True, default_environment=test_environment) job = apiclient.Job(pipeline_options, proto_pipeline) dm = dataflow_metrics.DataflowMetrics(mock_client, mock_job_result, job) self.assertEqual( 'MyTestParDo', dm._translate_step_name('ref_AppliedPTransform_MyTestParDo_14')) def test_query_counters(self): mock_client, mock_job_result = self.setup_mock_client_result( self.ONLY_COUNTERS_LIST) dm = dataflow_metrics.DataflowMetrics(mock_client, mock_job_result) dm._translate_step_name = types.MethodType(lambda self, x: 'split', dm) query_result = dm.query() expected_counters = [ MetricResult( MetricKey( 'split', MetricName('__main__.WordExtractingDoFn', 'empty_lines')), 1080, 1080), MetricResult( MetricKey( 'split', MetricName('__main__.WordExtractingDoFn', 'words')), 26181, 26185), ] self.assertEqual( sorted(query_result['counters'], key=lambda x: x.key.metric.name), sorted(expected_counters, key=lambda x: x.key.metric.name)) def test_system_counters_set_labels_and_step_name(self): mock_client, mock_job_result = self.setup_mock_client_result( self.SYSTEM_COUNTERS_LIST) test_object = dataflow_metrics.DataflowMetrics(mock_client, mock_job_result) all_metrics = test_object.all_metrics() matchers = [ MetricResultMatcher( name='ElementCount', labels={ 'original_name': 'ToIsmRecordForMultimap-out0-ElementCount', 'output_user_name': 'ToIsmRecordForMultimap-out0' }, attempted=42, committed=42), MetricResultMatcher( name='MeanByteCount', labels={ 'original_name': 'Read-out0-MeanByteCount', 'output_user_name': 'GroupByKey/Read-out0' }, attempted=31, committed=31), MetricResultMatcher( name='ExecutionTime_ProcessElement', step='write/Write/Write', attempted=1000, committed=1000) ] errors = metric_result_matchers.verify_all(all_metrics, matchers) self.assertFalse(errors, errors) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/dataflow_runner.py0000644000175100017510000010457315206113032024730 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A runner implementation that submits a job for remote execution. The runner will create a JSON description of the job graph and then submit it to the Dataflow Service for remote execution by a worker. """ # pytype: skip-file import logging import os import threading import time import warnings from collections import defaultdict from subprocess import DEVNULL from typing import TYPE_CHECKING from typing import List import apache_beam as beam from apache_beam import coders from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import TestOptions from apache_beam.options.pipeline_options import TypeOptions from apache_beam.options.pipeline_options import WorkerOptions from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.dataflow.internal.clients import dataflow as dataflow_api from apache_beam.runners.pipeline_utils import group_by_key_input_visitor from apache_beam.runners.pipeline_utils import merge_common_environments from apache_beam.runners.pipeline_utils import merge_superset_dep_environments from apache_beam.runners.runner import PipelineResult from apache_beam.runners.runner import PipelineRunner from apache_beam.runners.runner import PipelineState from apache_beam.transforms import environments from apache_beam.typehints import typehints from apache_beam.utils import processes from apache_beam.utils.interactive_utils import is_in_notebook from apache_beam.utils.plugin import BeamPlugin if TYPE_CHECKING: from apache_beam.pipeline import PTransformOverride __all__ = ['DataflowRunner'] _LOGGER = logging.getLogger(__name__) BQ_SOURCE_UW_ERROR = ( 'The Read(BigQuerySource(...)) transform is not supported with newer stack ' 'features (Fn API, Dataflow Runner V2, etc). Please use the transform ' 'apache_beam.io.gcp.bigquery.ReadFromBigQuery instead.') class DataflowRunner(PipelineRunner): """A runner that creates job graphs and submits them for remote execution. Every execution of the run() method will submit an independent job for remote execution that consists of the nodes reachable from the passed-in node argument or entire graph if the node is None. The run() method returns after the service creates the job, and the job status is reported as RUNNING. """ # A list of PTransformOverride objects to be applied before running a pipeline # using DataflowRunner. # Currently this only works for overrides where the input and output types do # not change. # For internal SDK use only. This should not be updated by Beam pipeline # authors. # Imported here to avoid circular dependencies. # TODO: Remove the apache_beam.pipeline dependency in CreatePTransformOverride from apache_beam.runners.dataflow.ptransform_overrides import NativeReadPTransformOverride # These overrides should be applied before the proto representation of the # graph is created. _PTRANSFORM_OVERRIDES = [ NativeReadPTransformOverride(), ] # type: List[PTransformOverride] def __init__(self, cache=None): self._default_environment = None def default_pickle_library_override(self): return 'cloudpickle' def is_fnapi_compatible(self): return False @staticmethod def poll_for_job_completion( runner, result, duration, state_update_callback=None): """Polls for the specified job to finish running (successfully or not). Updates the result with the new job information before returning. Args: runner: DataflowRunner instance to use for polling job state. result: DataflowPipelineResult instance used for job information. duration (int): The time to wait (in milliseconds) for job to finish. If it is set to :data:`None`, it will wait indefinitely until the job is finished. """ if result.state == PipelineState.DONE: return last_message_time = None current_seen_messages = set() last_error_rank = float('-inf') last_error_msg = None last_job_state = None # How long to wait after pipeline failure for the error # message to show up giving the reason for the failure. # It typically takes about 30 seconds. final_countdown_timer_secs = 50.0 sleep_secs = 5.0 # Try to prioritize the user-level traceback, if any. def rank_error(msg): if 'work item was attempted' in msg: return -1 elif 'Traceback' in msg: return 1 return 0 if duration: start_secs = time.time() duration_secs = duration // 1000 job_id = result.job_id() while True: response = runner.dataflow_client.get_job(job_id) # If get() is called very soon after Create() the response may not contain # an initialized 'currentState' field. if response.currentState is not None: if response.currentState != last_job_state: if state_update_callback: state_update_callback(response.currentState) _LOGGER.info('Job %s is in state %s', job_id, response.currentState) last_job_state = response.currentState if str(response.currentState) not in ('JOB_STATE_RUNNING', 'JOB_STATE_PAUSED', 'JOB_STATE_PAUSING'): # Stop checking for new messages on timeout, explanatory # message received, success, or a terminal job state caused # by the user that therefore doesn't require explanation. if (final_countdown_timer_secs <= 0.0 or last_error_msg is not None or str(response.currentState) == 'JOB_STATE_DONE' or str(response.currentState) == 'JOB_STATE_CANCELLED' or str(response.currentState) == 'JOB_STATE_UPDATED' or str(response.currentState) == 'JOB_STATE_DRAINED'): break # Check that job is in a post-preparation state before starting the # final countdown. if (str(response.currentState) not in ('JOB_STATE_PENDING', 'JOB_STATE_QUEUED')): # The job has failed; ensure we see any final error messages. sleep_secs = 1.0 # poll faster during the final countdown final_countdown_timer_secs -= sleep_secs time.sleep(sleep_secs) # Get all messages since beginning of the job run or since last message. page_token = None while True: messages, page_token = runner.dataflow_client.list_messages( job_id, page_token=page_token, start_time=last_message_time) for m in messages: message = '%s: %s: %s' % (m.time, m.messageImportance, m.messageText) if not last_message_time or m.time > last_message_time: last_message_time = m.time current_seen_messages = set() if message in current_seen_messages: # Skip the message if it has already been seen at the current # time. This could be the case since the list_messages API is # queried starting at last_message_time. continue else: current_seen_messages.add(message) # Skip empty messages. if m.messageImportance is None: continue message_importance = str(m.messageImportance) if (message_importance == 'JOB_MESSAGE_DEBUG' or message_importance == 'JOB_MESSAGE_DETAILED'): _LOGGER.debug(message) elif message_importance == 'JOB_MESSAGE_BASIC': _LOGGER.info(message) elif message_importance == 'JOB_MESSAGE_WARNING': _LOGGER.warning(message) elif message_importance == 'JOB_MESSAGE_ERROR': _LOGGER.error(message) if rank_error(m.messageText) >= last_error_rank: last_error_rank = rank_error(m.messageText) last_error_msg = m.messageText else: _LOGGER.info(message) if not page_token: break if duration: passed_secs = time.time() - start_secs if passed_secs > duration_secs: _LOGGER.warning( 'Timing out on waiting for job %s after %d seconds', job_id, passed_secs) break result._job = response runner.last_error_msg = last_error_msg @staticmethod def _only_element(iterable): # type: (Iterable[T]) -> T # noqa: F821 element, = iterable return element @staticmethod def side_input_visitor(deterministic_key_coders=True): # Imported here to avoid circular dependencies. # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.pipeline import PipelineVisitor from apache_beam.transforms.core import ParDo class SideInputVisitor(PipelineVisitor): """Ensures input `PCollection` used as a side inputs has a `KV` type. TODO(BEAM-115): Once Python SDK is compatible with the new Runner API, we could directly replace the coder instead of mutating the element type. """ def visit_transform(self, transform_node): if isinstance(transform_node.transform, ParDo): new_side_inputs = [] for side_input in transform_node.side_inputs: access_pattern = side_input._side_input_data().access_pattern if access_pattern == common_urns.side_inputs.ITERABLE.urn: # TODO(https://github.com/apache/beam/issues/20043): Stop # patching up the access pattern to appease Dataflow when # using the UW and hardcode the output type to be Any since # the Dataflow JSON and pipeline proto can differ in coders # which leads to encoding/decoding issues within the runner. side_input.pvalue.element_type = typehints.Any new_side_input = _DataflowIterableSideInput(side_input) elif access_pattern == common_urns.side_inputs.MULTIMAP.urn: # Ensure the input coder is a KV coder and patch up the # access pattern to appease Dataflow. side_input.pvalue.element_type = typehints.coerce_to_kv_type( side_input.pvalue.element_type, transform_node.full_label) side_input.pvalue.requires_deterministic_key_coder = ( deterministic_key_coders and transform_node.full_label) new_side_input = _DataflowMultimapSideInput(side_input) else: raise ValueError( 'Unsupported access pattern for %r: %r' % (transform_node.full_label, access_pattern)) new_side_inputs.append(new_side_input) transform_node.side_inputs = new_side_inputs transform_node.transform.side_inputs = new_side_inputs return SideInputVisitor() @staticmethod def flatten_input_visitor(): # Imported here to avoid circular dependencies. from apache_beam.pipeline import PipelineVisitor class FlattenInputVisitor(PipelineVisitor): """A visitor that replaces the element type for input ``PCollections``s of a ``Flatten`` transform with that of the output ``PCollection``. """ def visit_transform(self, transform_node): # Imported here to avoid circular dependencies. # pylint: disable=wrong-import-order, wrong-import-position from apache_beam import Flatten if isinstance(transform_node.transform, Flatten): output_pcoll = DataflowRunner._only_element( transform_node.outputs.values()) for input_pcoll in transform_node.inputs: input_pcoll.element_type = output_pcoll.element_type return FlattenInputVisitor() @staticmethod def combinefn_visitor(): # Imported here to avoid circular dependencies. from apache_beam import core from apache_beam.pipeline import PipelineVisitor class CombineFnVisitor(PipelineVisitor): """Checks if `CombineFn` has non-default setup or teardown methods. If yes, raises `ValueError`. """ def visit_transform(self, applied_transform): transform = applied_transform.transform if isinstance(transform, core.ParDo) and isinstance( transform.fn, core.CombineValuesDoFn): if self._overrides_setup_or_teardown(transform.fn.combinefn): raise ValueError( 'CombineFn.setup and CombineFn.teardown are ' 'not supported with non-portable Dataflow ' 'runner. Please use Dataflow Runner V2 instead.') @staticmethod def _overrides_setup_or_teardown(combinefn): # TODO(https://github.com/apache/beam/issues/18716): provide an # implementation for this method return False return CombineFnVisitor() def _adjust_pipeline_for_dataflow_v2(self, pipeline): # Dataflow runner requires a KV type for GBK inputs, hence we enforce that # here. pipeline.visit( group_by_key_input_visitor( not pipeline._options.view_as( TypeOptions).allow_non_deterministic_key_coders)) def run_pipeline(self, pipeline, options, pipeline_proto=None): """Remotely executes entire pipeline or parts reachable from node.""" if _is_runner_v2_disabled(options): raise ValueError( 'Disabling Runner V2 no longer supported ' 'using Beam Python %s.' % beam.version.__version__) # Label goog-dataflow-notebook if job is started from notebook. if is_in_notebook(): notebook_version = ( 'goog-dataflow-notebook=' + beam.version.__version__.replace('.', '_')) if options.view_as(GoogleCloudOptions).labels: options.view_as(GoogleCloudOptions).labels.append(notebook_version) else: options.view_as(GoogleCloudOptions).labels = [notebook_version] # Import here to avoid adding the dependency for local running scenarios. try: # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.runners.dataflow.internal import apiclient except ImportError: raise ImportError( 'Google Cloud Dataflow runner not available, ' 'please install apache_beam[gcp]') _check_and_add_missing_options(options) # Convert all side inputs into a form acceptable to Dataflow. if pipeline: pipeline.visit(self.combinefn_visitor()) pipeline.visit( self.side_input_visitor( deterministic_key_coders=not options.view_as( TypeOptions).allow_non_deterministic_key_coders)) # Performing configured PTransform overrides. Note that this is currently # done before Runner API serialization, since the new proto needs to # contain any added PTransforms. pipeline.replace_all(DataflowRunner._PTRANSFORM_OVERRIDES) # Apply DataflowRunner-specific overrides (e.g., streaming PubSub # optimizations) from apache_beam.runners.dataflow.ptransform_overrides import get_dataflow_transform_overrides dataflow_overrides = get_dataflow_transform_overrides(options) if dataflow_overrides: pipeline.replace_all(dataflow_overrides) if options.view_as(DebugOptions).lookup_experiment('use_legacy_bq_sink'): warnings.warn( "Native sinks no longer implemented; " "ignoring use_legacy_bq_sink.") if pipeline_proto: self.proto_pipeline = pipeline_proto else: if options.view_as(SetupOptions).prebuild_sdk_container_engine: # if prebuild_sdk_container_engine is specified we will build a new sdk # container image with dependencies pre-installed and use that image, # instead of using the inferred default container image. self._default_environment = ( environments.DockerEnvironment.from_options(options)) options.view_as(WorkerOptions).sdk_container_image = ( self._default_environment.container_image) else: artifacts = environments.python_sdk_dependencies(options) if artifacts: _LOGGER.info( "Pipeline has additional dependencies to be installed " "in SDK worker container, consider using the SDK " "container image pre-building workflow to avoid " "repetitive installations. Learn more on " "https://cloud.google.com/dataflow/docs/guides/" "using-custom-containers#prebuild") self._default_environment = ( environments.DockerEnvironment.from_container_image( apiclient.get_container_image_from_options(options), artifacts=artifacts, resource_hints=environments.resource_hints_from_options( options))) # This has to be performed before pipeline proto is constructed to make # sure that the changes are reflected in the portable job submission path. self._adjust_pipeline_for_dataflow_v2(pipeline) # Snapshot the pipeline in a portable proto. self.proto_pipeline, self.proto_context = pipeline.to_runner_api( return_context=True, default_environment=self._default_environment) if any(pcoll.is_bounded == beam_runner_api_pb2.IsBounded.UNBOUNDED for pcoll in self.proto_pipeline.components.pcollections.values()): if (not options.view_as(StandardOptions).streaming and not options.view_as(DebugOptions).lookup_experiment( 'unsafely_attempt_to_process_unbounded_data_in_batch_mode')): _LOGGER.info( 'Automatically inferring streaming mode ' 'due to unbounded PCollections.') options.view_as(StandardOptions).streaming = True if options.view_as(StandardOptions).streaming: _check_and_add_missing_streaming_options(options) # Dataflow can only handle Docker environments. for env_id, env in self.proto_pipeline.components.environments.items(): self.proto_pipeline.components.environments[env_id].CopyFrom( environments.resolve_anyof_environment( env, common_urns.environments.DOCKER.urn)) self.proto_pipeline = merge_common_environments( merge_superset_dep_environments(self.proto_pipeline)) # Optimize the pipeline if it not streaming and the pre_optimize # experiment is set. if not options.view_as(StandardOptions).streaming: pre_optimize = options.view_as(DebugOptions).lookup_experiment( 'pre_optimize', 'default').lower() from apache_beam.runners.portability.fn_api_runner import translations if pre_optimize == 'none': phases = [] elif pre_optimize == 'default' or pre_optimize == 'all': phases = [translations.pack_combiners, translations.sort_stages] else: phases = [] for phase_name in pre_optimize.split(','): # For now, these are all we allow. if phase_name in ('pack_combiners', ): phases.append(getattr(translations, phase_name)) else: raise ValueError( 'Unknown or inapplicable phase for pre_optimize: %s' % phase_name) phases.append(translations.sort_stages) if phases: self.proto_pipeline = translations.optimize_pipeline( self.proto_pipeline, phases=phases, known_runner_urns=frozenset(), partial=True) # Add setup_options for all the BeamPlugin imports setup_options = options.view_as(SetupOptions) plugins = BeamPlugin.get_all_plugin_paths() if setup_options.beam_plugins is not None: plugins = list(set(plugins + setup_options.beam_plugins)) setup_options.beam_plugins = plugins # Elevate "min_cpu_platform" to pipeline option, but using the existing # experiment. debug_options = options.view_as(DebugOptions) worker_options = options.view_as(WorkerOptions) if worker_options.min_cpu_platform: debug_options.add_experiment( 'min_cpu_platform=' + worker_options.min_cpu_platform) self.job = apiclient.Job(options, self.proto_pipeline) test_options = options.view_as(TestOptions) # If it is a dry run, return without submitting the job. if test_options.dry_run: result = PipelineResult(PipelineState.DONE) result.wait_until_finish = lambda duration=None: None result.job = self.job return result # Get a Dataflow API client and set its options self.dataflow_client = apiclient.DataflowApplicationClient( options, self.job.root_staging_location) # Create the job description and send a request to the service. The result # can be None if there is no need to send a request to the service (e.g. # template creation). If a request was sent and failed then the call will # raise an exception. result = DataflowPipelineResult( self.dataflow_client.create_job(self.job), self, options) # TODO(BEAM-4274): Circular import runners-metrics. Requires refactoring. from apache_beam.runners.dataflow.dataflow_metrics import DataflowMetrics self._metrics = DataflowMetrics(self.dataflow_client, result, self.job) result.metric_results = self._metrics return result @staticmethod def _get_coder(typehint, window_coder): """Returns a coder based on a typehint object.""" if window_coder: return coders.WindowedValueCoder( coders.registry.get_coder(typehint), window_coder=window_coder) return coders.registry.get_coder(typehint) def _verify_gbk_coders(self, transform, pcoll): # Infer coder of parent. # # TODO(ccy): make Coder inference and checking less specialized and more # comprehensive. parent = pcoll.producer if parent: coder = parent.transform._infer_output_coder() # pylint: disable=protected-access if not coder: coder = self._get_coder(pcoll.element_type or typehints.Any, None) if not coder.is_kv_coder(): raise ValueError(( 'Coder for the GroupByKey operation "%s" is not a ' 'key-value coder: %s.') % (transform.label, coder)) # TODO(robertwb): Update the coder itself if it changed. coders.registry.verify_deterministic( coder.key_coder(), 'GroupByKey operation "%s"' % transform.label) def get_default_gcp_region(self): """Get a default value for Google Cloud region according to https://cloud.google.com/compute/docs/gcloud-compute/#default-properties. If no default can be found, returns None. """ environment_region = os.environ.get('CLOUDSDK_COMPUTE_REGION') if environment_region: _LOGGER.info( 'Using default GCP region %s from $CLOUDSDK_COMPUTE_REGION', environment_region) return environment_region try: cmd = ['gcloud', 'config', 'get-value', 'compute/region'] raw_output = processes.check_output(cmd, stderr=DEVNULL) formatted_output = raw_output.decode('utf-8').strip() if formatted_output: _LOGGER.info( 'Using default GCP region %s from `%s`', formatted_output, ' '.join(cmd)) return formatted_output except RuntimeError: pass return None class _DataflowSideInput(beam.pvalue.AsSideInput): """Wraps a side input as a dataflow-compatible side input.""" def _view_options(self): return { 'data': self._data, } def _side_input_data(self): return self._data def _add_runner_v2_missing_options(options): debug_options = options.view_as(DebugOptions) debug_options.add_experiment('beam_fn_api') debug_options.add_experiment('use_unified_worker') debug_options.add_experiment('use_runner_v2') debug_options.add_experiment('use_portable_job_submission') def _check_and_add_missing_options(options): # Type: (PipelineOptions) -> None """Validates and adds missing pipeline options depending on options set. :param options: PipelineOptions for this pipeline. """ debug_options = options.view_as(DebugOptions) dataflow_service_options = options.view_as( GoogleCloudOptions).dataflow_service_options or [] # Add use_gbek to dataflow_service_options if gbek is set. if options.view_as(SetupOptions).gbek: if 'use_gbek' not in dataflow_service_options: dataflow_service_options.append('use_gbek') elif 'use_gbek' in dataflow_service_options: raise ValueError( 'Do not set use_gbek directly, pass in the --gbek pipeline option ' 'with a valid secret instead.') _add_runner_v2_missing_options(options) # Ensure that prime is specified as an experiment if specified as a dataflow # service option if 'enable_prime' in dataflow_service_options: debug_options.add_experiment('enable_prime') elif debug_options.lookup_experiment('enable_prime'): dataflow_service_options.append('enable_prime') options.view_as( GoogleCloudOptions).dataflow_service_options = dataflow_service_options sdk_location = options.view_as(SetupOptions).sdk_location if 'dev' in beam.version.__version__ and sdk_location == 'default': raise ValueError( "You are submitting a pipeline with Apache Beam Python SDK " f"{beam.version.__version__}. " "When launching Dataflow jobs with an unreleased (dev) SDK, " "please provide an SDK distribution in the --sdk_location option " "to use a consistent SDK version at " "pipeline submission and runtime. To ignore this error and use " "an SDK preinstalled in the default Dataflow dev runtime environment " "or in a custom container image, use --sdk_location=container.") def _check_and_add_missing_streaming_options(options): # Type: (PipelineOptions) -> None """Validates and adds missing pipeline options depending on options set. Must be called after it has been determined whether we're running in streaming mode. :param options: PipelineOptions for this pipeline. """ # Streaming only supports using runner v2 (aka unified worker). # Runner v2 only supports using streaming engine (aka windmill service) if options.view_as(StandardOptions).streaming: debug_options = options.view_as(DebugOptions) debug_options.add_experiment('enable_streaming_engine') debug_options.add_experiment('enable_windmill_service') def _is_runner_v2_disabled(options): # Type: (PipelineOptions) -> bool """Returns true if runner v2 is disabled.""" debug_options = options.view_as(DebugOptions) return ( debug_options.lookup_experiment('disable_runner_v2') or debug_options.lookup_experiment('disable_runner_v2_until_2023') or debug_options.lookup_experiment('disable_runner_v2_until_v2.50') or debug_options.lookup_experiment('disable_prime_runner_v2')) class _DataflowIterableSideInput(_DataflowSideInput): """Wraps an iterable side input as dataflow-compatible side input.""" def __init__(self, side_input): # pylint: disable=protected-access self.pvalue = side_input.pvalue side_input_data = side_input._side_input_data() assert ( side_input_data.access_pattern == common_urns.side_inputs.ITERABLE.urn) self._data = beam.pvalue.SideInputData( common_urns.side_inputs.ITERABLE.urn, side_input_data.window_mapping_fn, side_input_data.view_fn) class _DataflowMultimapSideInput(_DataflowSideInput): """Wraps a multimap side input as dataflow-compatible side input.""" def __init__(self, side_input): # pylint: disable=protected-access self.pvalue = side_input.pvalue side_input_data = side_input._side_input_data() assert ( side_input_data.access_pattern == common_urns.side_inputs.MULTIMAP.urn) self._data = beam.pvalue.SideInputData( common_urns.side_inputs.MULTIMAP.urn, side_input_data.window_mapping_fn, side_input_data.view_fn) class DataflowPipelineResult(PipelineResult): """Represents the state of a pipeline run on the Dataflow service.""" def __init__(self, job, runner, options=None): """Initialize a new DataflowPipelineResult instance. Args: job: Job message from the Dataflow API. Could be :data:`None` if a job request was not sent to Dataflow service (e.g. template jobs). runner: DataflowRunner instance. """ self._job = job self._runner = runner self._options = options self.metric_results = None def _update_job(self): # We need the job id to be able to update job information. There is no need # to update the job if we are in a known terminal state. if self.has_job and not self.is_in_terminal_state(): self._job = self._runner.dataflow_client.get_job(self.job_id()) def job_id(self): return self._job.id def metrics(self): return self.metric_results def monitoring_infos(self): logging.warning('Monitoring infos not yet supported for Dataflow runner.') return [] @property def has_job(self): return self._job is not None @staticmethod def api_jobstate_to_pipeline_state(api_jobstate): values_enum = dataflow_api.Job.CurrentStateValueValuesEnum # Ordered by the enum values. Values that may be introduced in # future versions of Dataflow API are considered UNRECOGNIZED by this SDK. api_jobstate_map = defaultdict( lambda: PipelineState.UNRECOGNIZED, { values_enum.JOB_STATE_UNKNOWN: PipelineState.UNKNOWN, values_enum.JOB_STATE_STOPPED: PipelineState.STOPPED, values_enum.JOB_STATE_RUNNING: PipelineState.RUNNING, values_enum.JOB_STATE_DONE: PipelineState.DONE, values_enum.JOB_STATE_FAILED: PipelineState.FAILED, values_enum.JOB_STATE_CANCELLED: PipelineState.CANCELLED, values_enum.JOB_STATE_UPDATED: PipelineState.UPDATED, values_enum.JOB_STATE_DRAINING: PipelineState.DRAINING, values_enum.JOB_STATE_DRAINED: PipelineState.DRAINED, values_enum.JOB_STATE_PENDING: PipelineState.PENDING, values_enum.JOB_STATE_CANCELLING: PipelineState.CANCELLING, values_enum.JOB_STATE_RESOURCE_CLEANING_UP: PipelineState. RESOURCE_CLEANING_UP, values_enum.JOB_STATE_PAUSING: PipelineState.PAUSING, values_enum.JOB_STATE_PAUSED: PipelineState.PAUSED, }) return ( api_jobstate_map[api_jobstate] if api_jobstate else PipelineState.UNKNOWN) def _get_job_state(self): return self.api_jobstate_to_pipeline_state(self._job.currentState) @property def state(self): """Return the current state of the remote job. Returns: A PipelineState object. """ if not self.has_job: # https://github.com/apache/beam/blob/8f71dc41b30a978095ca0e0699009e4f4445a618/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py#L867-L870 return PipelineState.DONE self._update_job() return self._get_job_state() def is_in_terminal_state(self): if not self.has_job: return True return PipelineState.is_terminal(self._get_job_state()) def wait_until_finish(self, duration=None): if not self.is_in_terminal_state(): if not self.has_job: raise IOError('Failed to get the Dataflow job id.') gcp_options = self._options.view_as(GoogleCloudOptions) consoleUrl = ( "Console URL: https://console.cloud.google.com/" f"dataflow/jobs/{gcp_options.region}/{self.job_id()}" f"?project={gcp_options.project}") thread = threading.Thread( target=DataflowRunner.poll_for_job_completion, args=(self._runner, self, duration)) # Mark the thread as a daemon thread so a keyboard interrupt on the main # thread will terminate everything. This is also the reason we will not # use thread.join() to wait for the polling thread. thread.daemon = True thread.start() while thread.is_alive(): time.sleep(5.0) # TODO: Merge the termination code in poll_for_job_completion and # is_in_terminal_state. terminated = self.is_in_terminal_state() assert duration or terminated, ( 'Job did not reach to a terminal state after waiting indefinitely. ' '{}'.format(consoleUrl)) if terminated and self.state != PipelineState.DONE: # TODO(BEAM-1290): Consider converting this to an error log based on # theresolution of the issue. _LOGGER.error(consoleUrl) raise DataflowRuntimeException( 'Dataflow pipeline failed. State: %s, Error:\n%s' % (self.state, getattr(self._runner, 'last_error_msg', None)), self) elif PipelineState.is_terminal( self.state) and self.state == PipelineState.FAILED and self._runner: raise DataflowRuntimeException( 'Dataflow pipeline failed. State: %s, Error:\n%s' % (self.state, getattr(self._runner, 'last_error_msg', None)), self) return self.state def cancel(self): if not self.has_job: raise IOError('Failed to get the Dataflow job id.') self._update_job() if self.is_in_terminal_state(): _LOGGER.warning( 'Cancel failed because job %s is already terminated in state %s.', self.job_id(), self.state) else: if not self._runner.dataflow_client.modify_job_state( self.job_id(), 'JOB_STATE_CANCELLED'): cancel_failed_message = ( 'Failed to cancel job %s, please go to the Developers Console to ' 'cancel it manually.') % self.job_id() _LOGGER.error(cancel_failed_message) raise DataflowRuntimeException(cancel_failed_message, self) return self.state def __str__(self): return '<%s %s %s>' % (self.__class__.__name__, self.job_id(), self.state) def __repr__(self): return '<%s %s at %s>' % (self.__class__.__name__, self._job, hex(id(self))) class DataflowRuntimeException(Exception): """Indicates an error has occurred in running this pipeline.""" def __init__(self, msg, result): super().__init__(msg) self.result = result ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/dataflow_runner_test.py0000644000175100017510000007457615206113032026000 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the DataflowRunner class.""" # pytype: skip-file import unittest import mock import apache_beam as beam import apache_beam.transforms as ptransform from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.pipeline import AppliedPTransform from apache_beam.pipeline import Pipeline from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.pvalue import PCollection from apache_beam.runners import DataflowRunner from apache_beam.runners import TestDataflowRunner from apache_beam.runners import create_runner from apache_beam.runners import pipeline_utils from apache_beam.runners.dataflow.dataflow_runner import DataflowPipelineResult from apache_beam.runners.dataflow.dataflow_runner import DataflowRuntimeException from apache_beam.runners.dataflow.dataflow_runner import _check_and_add_missing_options from apache_beam.runners.dataflow.dataflow_runner import _check_and_add_missing_streaming_options from apache_beam.runners.dataflow.internal.clients import dataflow as dataflow_api from apache_beam.runners.internal import names from apache_beam.runners.runner import PipelineState from apache_beam.testing.extra_assertions import ExtraAssertionsMixin from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.transforms import combiners from apache_beam.transforms import environments from apache_beam.typehints import typehints # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.runners.dataflow.internal import apiclient except ImportError: apiclient = None # type: ignore # pylint: enable=wrong-import-order, wrong-import-position # SpecialParDo and SpecialDoFn are used in test_remote_runner_display_data. # Due to https://github.com/apache/beam/issues/19848, these need to be declared # outside of the test method. # TODO: Should not subclass ParDo. Switch to PTransform as soon as # composite transforms support display data. class SpecialParDo(beam.ParDo): def __init__(self, fn, now): super().__init__(fn) self.fn = fn self.now = now # Make this a list to be accessible within closure def display_data(self): return { 'asubcomponent': self.fn, 'a_class': SpecialParDo, 'a_time': self.now } class SpecialDoFn(beam.DoFn): def display_data(self): return {'dofn_value': 42} def process(self): pass @unittest.skipIf(apiclient is None, 'GCP dependencies are not installed') class DataflowRunnerTest(unittest.TestCase, ExtraAssertionsMixin): def setUp(self): self.default_properties = [ '--job_name=test-job', '--project=test-project', '--region=us-central1', '--staging_location=gs://beam/test', '--temp_location=gs://beam/tmp', '--no_auth', '--dry_run=True', '--sdk_location=container' ] @mock.patch('time.sleep', return_value=None) def test_wait_until_finish_unrecognized(self, patched_time_sleep): values_enum = dataflow_api.Job.CurrentStateValueValuesEnum options = PipelineOptions(self.default_properties) class MockDataflowRunner(object): def __init__(self, states): self.dataflow_client = mock.MagicMock() self.job = mock.MagicMock() self.job.id = "test-job-id" self.job.currentState = values_enum.JOB_STATE_UNKNOWN self._states = states self._next_state_index = 0 def get_job_side_effect(*args, **kwargs): self.job.currentState = self._states[self._next_state_index] if self._next_state_index < (len(self._states) - 1): self._next_state_index += 1 return mock.DEFAULT self.dataflow_client.get_job = mock.MagicMock( return_value=self.job, side_effect=get_job_side_effect) self.dataflow_client.list_messages = mock.MagicMock( return_value=([], None)) with self.assertRaisesRegex( AssertionError, r'Job did not reach to a terminal state after waiting indefinitely. ' r'Console URL: ' r'https://console.cloud.google.com/dataflow/jobs/' r'us-central1/test-job-id\?project=test-project'): failed_runner = MockDataflowRunner("some_unrecognized_state") failed_result = DataflowPipelineResult( failed_runner.job, failed_runner, options) failed_result.wait_until_finish() @mock.patch('time.sleep', return_value=None) def test_wait_until_finish(self, patched_time_sleep): values_enum = dataflow_api.Job.CurrentStateValueValuesEnum options = PipelineOptions(self.default_properties) class MockDataflowRunner(object): def __init__(self, states): self.dataflow_client = mock.MagicMock() self.job = mock.MagicMock() self.job.currentState = values_enum.JOB_STATE_UNKNOWN self._states = states self._next_state_index = 0 def get_job_side_effect(*args, **kwargs): self.job.currentState = self._states[self._next_state_index] if self._next_state_index < (len(self._states) - 1): self._next_state_index += 1 return mock.DEFAULT self.dataflow_client.get_job = mock.MagicMock( return_value=self.job, side_effect=get_job_side_effect) self.dataflow_client.list_messages = mock.MagicMock( return_value=([], None)) with self.assertRaisesRegex(DataflowRuntimeException, 'Dataflow pipeline failed. State: FAILED'): failed_runner = MockDataflowRunner([values_enum.JOB_STATE_FAILED]) failed_result = DataflowPipelineResult( failed_runner.job, failed_runner, options) failed_result.wait_until_finish() # check the second call can still triggers the exception with self.assertRaisesRegex(DataflowRuntimeException, 'Dataflow pipeline failed. State: FAILED'): failed_result.wait_until_finish() succeeded_runner = MockDataflowRunner([values_enum.JOB_STATE_DONE]) succeeded_result = DataflowPipelineResult( succeeded_runner.job, succeeded_runner, options) result = succeeded_result.wait_until_finish() self.assertEqual(result, PipelineState.DONE) # Time array has duplicate items, because some logging implementations also # call time. with mock.patch('time.time', mock.MagicMock(side_effect=[1, 1, 2, 2, 3])): duration_succeeded_runner = MockDataflowRunner( [values_enum.JOB_STATE_RUNNING, values_enum.JOB_STATE_DONE]) duration_succeeded_result = DataflowPipelineResult( duration_succeeded_runner.job, duration_succeeded_runner, options) result = duration_succeeded_result.wait_until_finish(5000) self.assertEqual(result, PipelineState.DONE) with mock.patch('time.time', mock.MagicMock(side_effect=[1, 9, 9, 20, 20])): duration_timedout_runner = MockDataflowRunner( [values_enum.JOB_STATE_RUNNING]) duration_timedout_result = DataflowPipelineResult( duration_timedout_runner.job, duration_timedout_runner, options) result = duration_timedout_result.wait_until_finish(5000) self.assertEqual(result, PipelineState.RUNNING) with mock.patch('time.time', mock.MagicMock(side_effect=[1, 9, 9, 20, 20])): duration_timedout_runner = MockDataflowRunner( [values_enum.JOB_STATE_PAUSING]) duration_timedout_result = DataflowPipelineResult( duration_timedout_runner.job, duration_timedout_runner, options) result = duration_timedout_result.wait_until_finish(5000) self.assertEqual(result, PipelineState.PAUSING) with mock.patch('time.time', mock.MagicMock(side_effect=[1, 9, 9, 20, 20])): duration_timedout_runner = MockDataflowRunner( [values_enum.JOB_STATE_PAUSED]) duration_timedout_result = DataflowPipelineResult( duration_timedout_runner.job, duration_timedout_runner, options) result = duration_timedout_result.wait_until_finish(5000) self.assertEqual(result, PipelineState.PAUSED) with mock.patch('time.time', mock.MagicMock(side_effect=[1, 1, 2, 2, 3])): with self.assertRaisesRegex(DataflowRuntimeException, 'Dataflow pipeline failed. State: CANCELLED'): duration_failed_runner = MockDataflowRunner( [values_enum.JOB_STATE_CANCELLED]) duration_failed_result = DataflowPipelineResult( duration_failed_runner.job, duration_failed_runner, options) duration_failed_result.wait_until_finish(5000) @mock.patch('time.sleep', return_value=None) def test_cancel(self, patched_time_sleep): values_enum = dataflow_api.Job.CurrentStateValueValuesEnum options = PipelineOptions( self.default_properties).view_as(GoogleCloudOptions) class MockDataflowRunner(object): def __init__(self, state, cancel_result): self.dataflow_client = mock.MagicMock() self.job = mock.MagicMock() self.job.currentState = state self.dataflow_client.get_job = mock.MagicMock(return_value=self.job) self.dataflow_client.modify_job_state = mock.MagicMock( return_value=cancel_result) self.dataflow_client.list_messages = mock.MagicMock( return_value=([], None)) with self.assertRaisesRegex(DataflowRuntimeException, 'Failed to cancel job'): failed_runner = MockDataflowRunner(values_enum.JOB_STATE_RUNNING, False) failed_result = DataflowPipelineResult( failed_runner.job, failed_runner, options) failed_result.cancel() succeeded_runner = MockDataflowRunner(values_enum.JOB_STATE_RUNNING, True) succeeded_result = DataflowPipelineResult( succeeded_runner.job, succeeded_runner, options) succeeded_result.cancel() terminal_runner = MockDataflowRunner(values_enum.JOB_STATE_DONE, False) terminal_result = DataflowPipelineResult( terminal_runner.job, terminal_runner, options) terminal_result.cancel() def test_api_jobstate_to_pipeline_state(self): values_enum = dataflow_api.Job.CurrentStateValueValuesEnum expected_mappings = [ (values_enum.JOB_STATE_UNKNOWN, PipelineState.UNKNOWN), (values_enum.JOB_STATE_STOPPED, PipelineState.STOPPED), (values_enum.JOB_STATE_RUNNING, PipelineState.RUNNING), (values_enum.JOB_STATE_DONE, PipelineState.DONE), (values_enum.JOB_STATE_FAILED, PipelineState.FAILED), (values_enum.JOB_STATE_CANCELLED, PipelineState.CANCELLED), (values_enum.JOB_STATE_UPDATED, PipelineState.UPDATED), (values_enum.JOB_STATE_DRAINING, PipelineState.DRAINING), (values_enum.JOB_STATE_DRAINED, PipelineState.DRAINED), (values_enum.JOB_STATE_PENDING, PipelineState.PENDING), (values_enum.JOB_STATE_CANCELLING, PipelineState.CANCELLING), ( values_enum.JOB_STATE_RESOURCE_CLEANING_UP, PipelineState.RESOURCE_CLEANING_UP), (values_enum.JOB_STATE_PAUSING, PipelineState.PAUSING), (values_enum.JOB_STATE_PAUSED, PipelineState.PAUSED), ] for api_state, pipeline_state in expected_mappings: self.assertEqual( DataflowPipelineResult.api_jobstate_to_pipeline_state(api_state), pipeline_state) def test_create_runner(self): self.assertTrue(isinstance(create_runner('DataflowRunner'), DataflowRunner)) self.assertTrue( isinstance(create_runner('TestDataflowRunner'), TestDataflowRunner)) @staticmethod def dependency_proto_from_main_session_file(serialized_path): return [ beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=serialized_path, role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2.ArtifactStagingToRolePayload( staged_name=names.PICKLED_MAIN_SESSION_FILE).SerializeToString( )) ] def test_environment_override_translation_legacy_worker_harness_image(self): self.default_properties.append('--experiments=beam_fn_api') self.default_properties.append('--worker_harness_container_image=LEGACY') remote_runner = DataflowRunner() options = PipelineOptions(self.default_properties) options.view_as(DebugOptions).add_experiment( 'disable_logging_submission_environment') with Pipeline(remote_runner, options=options) as p: ( # pylint: disable=expression-not-assigned p | ptransform.Create([1, 2, 3]) | 'Do' >> ptransform.FlatMap(lambda x: [(x, x)]) | ptransform.GroupByKey()) actual = list(remote_runner.proto_pipeline.components.environments.values()) self.assertEqual(len(actual), 1) actual = actual[0] file_path = actual.dependencies[0].type_payload # Dependency payload contains main_session from a transient temp directory # Use actual for expected value. main_session_dep = self.dependency_proto_from_main_session_file(file_path) self.assertEqual( actual, beam_runner_api_pb2.Environment( urn=common_urns.environments.DOCKER.urn, payload=beam_runner_api_pb2.DockerPayload( container_image='LEGACY').SerializeToString(), capabilities=environments.python_sdk_docker_capabilities(), dependencies=environments.python_sdk_dependencies(options=options) + main_session_dep)) def test_environment_override_translation_sdk_container_image(self): self.default_properties.append('--experiments=beam_fn_api') self.default_properties.append('--sdk_container_image=FOO') remote_runner = DataflowRunner() options = PipelineOptions(self.default_properties) options.view_as(DebugOptions).add_experiment( 'disable_logging_submission_environment') with Pipeline(remote_runner, options=options) as p: ( # pylint: disable=expression-not-assigned p | ptransform.Create([1, 2, 3]) | 'Do' >> ptransform.FlatMap(lambda x: [(x, x)]) | ptransform.GroupByKey()) actual = list(remote_runner.proto_pipeline.components.environments.values()) self.assertEqual(len(actual), 1) actual = actual[0] file_path = actual.dependencies[0].type_payload # Dependency payload contains main_session from a transient temp directory # Use actual for expected value. main_session_dep = self.dependency_proto_from_main_session_file(file_path) self.assertEqual( actual, beam_runner_api_pb2.Environment( urn=common_urns.environments.DOCKER.urn, payload=beam_runner_api_pb2.DockerPayload( container_image='FOO').SerializeToString(), capabilities=environments.python_sdk_docker_capabilities(), dependencies=environments.python_sdk_dependencies(options=options) + main_session_dep)) def test_remote_runner_translation(self): remote_runner = DataflowRunner() with Pipeline(remote_runner, options=PipelineOptions(self.default_properties)) as p: ( # pylint: disable=expression-not-assigned p | ptransform.Create([1, 2, 3]) | 'Do' >> ptransform.FlatMap(lambda x: [(x, x)]) | ptransform.GroupByKey()) def test_group_by_key_input_visitor_with_valid_inputs(self): p = TestPipeline() pcoll1 = PCollection(p) pcoll2 = PCollection(p) pcoll3 = PCollection(p) pcoll1.element_type = None pcoll2.element_type = typehints.Any pcoll3.element_type = typehints.KV[typehints.Any, typehints.Any] for pcoll in [pcoll1, pcoll2, pcoll3]: applied = AppliedPTransform( None, beam.GroupByKey(), "label", {'pcoll': pcoll}, None, None) applied.outputs[None] = PCollection(None) pipeline_utils.group_by_key_input_visitor().visit_transform(applied) self.assertEqual( pcoll.element_type, typehints.KV[typehints.Any, typehints.Any]) def test_group_by_key_input_visitor_with_invalid_inputs(self): p = TestPipeline() pcoll1 = PCollection(p) pcoll2 = PCollection(p) pcoll1.element_type = str pcoll2.element_type = typehints.Set err_msg = ( r"Input to 'label' must be compatible with KV\[Any, Any\]. " "Found .*") for pcoll in [pcoll1, pcoll2]: with self.assertRaisesRegex(ValueError, err_msg): pipeline_utils.group_by_key_input_visitor().visit_transform( AppliedPTransform( None, beam.GroupByKey(), "label", {'in': pcoll}, None, None)) def test_group_by_key_input_visitor_for_non_gbk_transforms(self): p = TestPipeline() pcoll = PCollection(p) for transform in [beam.Flatten(), beam.Map(lambda x: x)]: pcoll.element_type = typehints.Any pipeline_utils.group_by_key_input_visitor().visit_transform( AppliedPTransform( None, transform, "label", {'in': pcoll}, None, None)) self.assertEqual(pcoll.element_type, typehints.Any) def test_flatten_input_with_visitor_with_single_input(self): self._test_flatten_input_visitor(typehints.KV[int, int], typehints.Any, 1) def test_flatten_input_with_visitor_with_multiple_inputs(self): self._test_flatten_input_visitor( typehints.KV[int, typehints.Any], typehints.Any, 5) def _test_flatten_input_visitor(self, input_type, output_type, num_inputs): p = TestPipeline() inputs = {} for ix in range(num_inputs): input_pcoll = PCollection(p) input_pcoll.element_type = input_type inputs[str(ix)] = input_pcoll output_pcoll = PCollection(p) output_pcoll.element_type = output_type flatten = AppliedPTransform( None, beam.Flatten(), "label", inputs, None, None) flatten.add_output(output_pcoll, None) DataflowRunner.flatten_input_visitor().visit_transform(flatten) for _ in range(num_inputs): self.assertEqual(inputs['0'].element_type, output_type) def test_gbk_then_flatten_input_visitor(self): p = TestPipeline( runner=DataflowRunner(), options=PipelineOptions(self.default_properties)) none_str_pc = p | 'c1' >> beam.Create({None: 'a'}) none_int_pc = p | 'c2' >> beam.Create({None: 3}) flat = (none_str_pc, none_int_pc) | beam.Flatten() _ = flat | beam.GroupByKey() # This may change if type inference changes, but we assert it here # to make sure the check below is not vacuous. self.assertNotIsInstance(flat.element_type, typehints.TupleConstraint) p.visit(pipeline_utils.group_by_key_input_visitor()) p.visit(DataflowRunner.flatten_input_visitor()) # The dataflow runner requires gbk input to be tuples *and* flatten # inputs to be equal to their outputs. Assert both hold. self.assertIsInstance(flat.element_type, typehints.TupleConstraint) self.assertEqual(flat.element_type, none_str_pc.element_type) self.assertEqual(flat.element_type, none_int_pc.element_type) def test_side_input_visitor(self): p = TestPipeline() pc = p | beam.Create([]) transform = beam.Map( lambda x, y, z: (x, y, z), beam.pvalue.AsSingleton(pc), beam.pvalue.AsMultiMap(pc)) applied_transform = AppliedPTransform( None, transform, "label", {'pc': pc}, None, None) DataflowRunner.side_input_visitor().visit_transform(applied_transform) self.assertEqual(2, len(applied_transform.side_inputs)) self.assertEqual( common_urns.side_inputs.ITERABLE.urn, applied_transform.side_inputs[0]._side_input_data().access_pattern) self.assertEqual( common_urns.side_inputs.MULTIMAP.urn, applied_transform.side_inputs[1]._side_input_data().access_pattern) def test_min_cpu_platform_flag_is_propagated_to_experiments(self): remote_runner = DataflowRunner() self.default_properties.append('--min_cpu_platform=Intel Haswell') with Pipeline(remote_runner, PipelineOptions(self.default_properties)) as p: p | ptransform.Create([1]) # pylint: disable=expression-not-assigned self.assertIn( 'min_cpu_platform=Intel Haswell', remote_runner.job.options.view_as(DebugOptions).experiments) def test_streaming_adds_windmill_experiments(self): remote_runner = DataflowRunner() self.default_properties.append('--streaming') self.default_properties.append('--experiment=some_other_experiment') with Pipeline(remote_runner, PipelineOptions(self.default_properties)) as p: p | ptransform.Create([1]) # pylint: disable=expression-not-assigned experiments_for_job = ( remote_runner.job.options.view_as(DebugOptions).experiments) self.assertIn('enable_streaming_engine', experiments_for_job) self.assertIn('enable_windmill_service', experiments_for_job) self.assertIn('some_other_experiment', experiments_for_job) def test_upload_graph_experiment(self): remote_runner = DataflowRunner() self.default_properties.append('--experiment=upload_graph') with Pipeline(remote_runner, PipelineOptions(self.default_properties)) as p: p | ptransform.Create([1]) # pylint: disable=expression-not-assigned experiments_for_job = ( remote_runner.job.options.view_as(DebugOptions).experiments) self.assertIn('upload_graph', experiments_for_job) def test_use_fastavro_experiment_is_not_added_when_use_avro_is_present(self): remote_runner = DataflowRunner() self.default_properties.append('--experiment=use_avro') with Pipeline(remote_runner, PipelineOptions(self.default_properties)) as p: p | ptransform.Create([1]) # pylint: disable=expression-not-assigned debug_options = remote_runner.job.options.view_as(DebugOptions) self.assertFalse(debug_options.lookup_experiment('use_fastavro', False)) @mock.patch('os.environ.get', return_value=None) @mock.patch('apache_beam.utils.processes.check_output', return_value=b'') def test_get_default_gcp_region_no_default_returns_none( self, patched_environ, patched_processes): runner = DataflowRunner() result = runner.get_default_gcp_region() self.assertIsNone(result) @mock.patch('os.environ.get', return_value='some-region1') @mock.patch('apache_beam.utils.processes.check_output', return_value=b'') def test_get_default_gcp_region_from_environ( self, patched_environ, patched_processes): runner = DataflowRunner() result = runner.get_default_gcp_region() self.assertEqual(result, 'some-region1') @mock.patch('os.environ.get', return_value=None) @mock.patch( 'apache_beam.utils.processes.check_output', return_value=b'some-region2\n') def test_get_default_gcp_region_from_gcloud( self, patched_environ, patched_processes): runner = DataflowRunner() result = runner.get_default_gcp_region() self.assertEqual(result, 'some-region2') @mock.patch('os.environ.get', return_value=None) @mock.patch( 'apache_beam.utils.processes.check_output', side_effect=RuntimeError('Executable gcloud not found')) def test_get_default_gcp_region_ignores_error( self, patched_environ, patched_processes): runner = DataflowRunner() result = runner.get_default_gcp_region() self.assertIsNone(result) @unittest.skip( 'https://github.com/apache/beam/issues/18716: enable once ' 'CombineFnVisitor is fixed') def test_unsupported_combinefn_detection(self): class CombinerWithNonDefaultSetupTeardown(combiners.CountCombineFn): def setup(self, *args, **kwargs): pass def teardown(self, *args, **kwargs): pass runner = DataflowRunner() with self.assertRaisesRegex(ValueError, 'CombineFn.setup and CombineFn.' 'teardown are not supported'): with beam.Pipeline(runner=runner, options=PipelineOptions(self.default_properties)) as p: _ = ( p | beam.Create([1]) | beam.CombineGlobally(CombinerWithNonDefaultSetupTeardown())) try: with beam.Pipeline(runner=runner, options=PipelineOptions(self.default_properties)) as p: _ = ( p | beam.Create([1]) | beam.CombineGlobally( combiners.SingleInputTupleCombineFn( combiners.CountCombineFn(), combiners.CountCombineFn()))) except ValueError: self.fail('ValueError raised unexpectedly') def test_pack_combiners(self): class PackableCombines(beam.PTransform): def annotations(self): return {python_urns.APPLY_COMBINER_PACKING: b''} def expand(self, pcoll): _ = pcoll | 'PackableMin' >> beam.CombineGlobally(min) _ = pcoll | 'PackableMax' >> beam.CombineGlobally(max) runner = DataflowRunner() with beam.Pipeline(runner=runner, options=PipelineOptions(self.default_properties)) as p: _ = p | beam.Create([10, 20, 30]) | PackableCombines() unpacked_minimum_step_name = ( 'PackableCombines/PackableMin/CombinePerKey/Combine') unpacked_maximum_step_name = ( 'PackableCombines/PackableMax/CombinePerKey/Combine') packed_step_name = ( 'PackableCombines/Packed[PackableMin_CombinePerKey, ' 'PackableMax_CombinePerKey]/Pack') transform_names = set( transform.unique_name for transform in runner.proto_pipeline.components.transforms.values()) self.assertNotIn(unpacked_minimum_step_name, transform_names) self.assertNotIn(unpacked_maximum_step_name, transform_names) self.assertIn(packed_step_name, transform_names) def test_batch_is_runner_v2(self): options = PipelineOptions(['--sdk_location=container']) _check_and_add_missing_options(options) for expected in ['beam_fn_api', 'use_unified_worker', 'use_runner_v2', 'use_portable_job_submission']: self.assertTrue( options.view_as(DebugOptions).lookup_experiment(expected, False), expected) def test_streaming_is_runner_v2(self): options = PipelineOptions(['--sdk_location=container', '--streaming']) _check_and_add_missing_options(options) _check_and_add_missing_streaming_options(options) for expected in ['beam_fn_api', 'use_unified_worker', 'use_runner_v2', 'use_portable_job_submission', 'enable_windmill_service', 'enable_streaming_engine']: self.assertTrue( options.view_as(DebugOptions).lookup_experiment(expected, False), expected) def test_dataflow_service_options_enable_prime_sets_runner_v2(self): options = PipelineOptions([ '--sdk_location=container', '--streaming', '--dataflow_service_options=enable_prime' ]) _check_and_add_missing_options(options) for expected in ['beam_fn_api', 'use_unified_worker', 'use_runner_v2', 'use_portable_job_submission']: self.assertTrue( options.view_as(DebugOptions).lookup_experiment(expected, False), expected) options = PipelineOptions([ '--sdk_location=container', '--streaming', '--dataflow_service_options=enable_prime' ]) _check_and_add_missing_options(options) _check_and_add_missing_streaming_options(options) for expected in ['beam_fn_api', 'use_unified_worker', 'use_runner_v2', 'use_portable_job_submission', 'enable_windmill_service', 'enable_streaming_engine']: self.assertTrue( options.view_as(DebugOptions).lookup_experiment(expected, False), expected) @unittest.skipIf(apiclient is None, 'GCP dependencies are not installed') @mock.patch( 'apache_beam.options.pipeline_options.GoogleCloudOptions.validate', lambda *args: []) def test_auto_streaming_with_unbounded(self): options = PipelineOptions([ '--sdk_location=container', '--runner=DataflowRunner', '--dry_run=True', '--temp_location=gs://bucket', '--project=project', '--region=region' ]) with beam.Pipeline(options=options) as p: _ = p | beam.io.ReadFromPubSub('projects/some-project/topics/some-topic') self.assertEqual( p.result.job.proto.type, apiclient.dataflow.Job.TypeValueValuesEnum.JOB_TYPE_STREAMING) @unittest.skipIf(apiclient is None, 'GCP dependencies are not installed') @mock.patch( 'apache_beam.options.pipeline_options.GoogleCloudOptions.validate', lambda *args: []) def test_auto_streaming_no_unbounded(self): options = PipelineOptions([ '--sdk_location=container', '--runner=DataflowRunner', '--dry_run=True', '--temp_location=gs://bucket', '--project=project', '--region=region' ]) with beam.Pipeline(options=options) as p: _ = p | beam.Create([1, 2, 3]) self.assertEqual( p.result.job.proto.type, apiclient.dataflow.Job.TypeValueValuesEnum.JOB_TYPE_BATCH) @unittest.skipIf(apiclient is None, 'GCP dependencies are not installed') @mock.patch( 'apache_beam.options.pipeline_options.GoogleCloudOptions.validate', lambda *args: []) def test_explicit_streaming_no_unbounded(self): options = PipelineOptions([ '--streaming', '--sdk_location=container', '--runner=DataflowRunner', '--dry_run=True', '--temp_location=gs://bucket', '--project=project', '--region=region' ]) with beam.Pipeline(options=options) as p: _ = p | beam.Create([1, 2, 3]) self.assertEqual( p.result.job.proto.type, apiclient.dataflow.Job.TypeValueValuesEnum.JOB_TYPE_STREAMING) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5854294 apache_beam-2.74.0/apache_beam/runners/dataflow/internal/0000755000175100017510000000000015206113160022770 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/internal/__init__.py0000644000175100017510000000142015206113032025074 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/internal/apiclient.py0000644000175100017510000015657015206113032025326 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ For internal use only. No backwards compatibility guarantees. Dataflow client utility functions.""" # pytype: skip-file # To regenerate the client: # pip install google-apitools[cli] # gen_client --discovery_url=cloudbuild.v1 --overwrite \ # --outdir=apache_beam/runners/dataflow/internal/clients/cloudbuild \ # --root_package=. client import ast import codecs from functools import partial import getpass import hashlib import io import json import logging import os import random import string from packaging import version import re import sys import time import traceback import warnings from copy import copy from datetime import datetime from datetime import timezone from apitools.base.py import encoding from apitools.base.py import exceptions from apache_beam import version as beam_version from apache_beam.internal.gcp.auth import get_service_credentials from apache_beam.internal.gcp.json_value import to_json_value from apache_beam.internal.http_client import get_new_http from apache_beam.io.filesystems import FileSystems from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem from apache_beam.io.gcp.gcsio import create_storage_client from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import WorkerOptions from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.dataflow.internal import names from apache_beam.runners.dataflow.internal.clients import dataflow from apache_beam.runners.internal import names as shared_names from apache_beam.runners.pipeline_utils import validate_pipeline_graph from apache_beam.runners.portability.stager import Stager from apache_beam.transforms import DataflowDistributionCounter from apache_beam.transforms import cy_combiners from apache_beam.transforms.display import DisplayData from apache_beam.transforms.environments import is_apache_beam_container from apache_beam.utils import retry from apache_beam.utils import proto_utils # Environment version information. It is passed to the service during a # a job submission and is used by the service to establish what features # are expected by the workers. _LEGACY_ENVIRONMENT_MAJOR_VERSION = '8' _FNAPI_ENVIRONMENT_MAJOR_VERSION = '8' _LOGGER = logging.getLogger(__name__) _PYTHON_VERSIONS_SUPPORTED_BY_DATAFLOW = [ '3.10', '3.11', '3.12', '3.13', '3.14' ] class Environment(object): """Wrapper for a dataflow Environment protobuf.""" def __init__( self, packages, options, environment_version, proto_pipeline_staged_url, proto_pipeline=None, pipeline_proto_hash=None): self.standard_options = options.view_as(StandardOptions) self.google_cloud_options = options.view_as(GoogleCloudOptions) self.worker_options = options.view_as(WorkerOptions) self.debug_options = options.view_as(DebugOptions) self.pipeline_url = proto_pipeline_staged_url self.proto = dataflow.Environment() self.proto.clusterManagerApiService = GoogleCloudOptions.COMPUTE_API_SERVICE self.proto.dataset = '{}/cloud_dataflow'.format( GoogleCloudOptions.BIGQUERY_API_SERVICE) self.proto.tempStoragePrefix = ( self.google_cloud_options.temp_location.replace( 'gs:/', GoogleCloudOptions.STORAGE_API_SERVICE)) if self.worker_options.worker_region: self.proto.workerRegion = self.worker_options.worker_region if self.worker_options.worker_zone: self.proto.workerZone = self.worker_options.worker_zone # User agent information. self.proto.userAgent = dataflow.Environment.UserAgentValue() self.local = 'localhost' in self.google_cloud_options.dataflow_endpoint self._proto_pipeline = proto_pipeline if self.google_cloud_options.service_account_email: self.proto.serviceAccountEmail = ( self.google_cloud_options.service_account_email) if self.google_cloud_options.dataflow_kms_key: self.proto.serviceKmsKeyName = self.google_cloud_options.dataflow_kms_key self.proto.userAgent.additionalProperties.extend([ dataflow.Environment.UserAgentValue.AdditionalProperty( key='name', value=to_json_value(self._get_python_sdk_name())), dataflow.Environment.UserAgentValue.AdditionalProperty( key='version', value=to_json_value(beam_version.__version__)) ]) # Version information. self.proto.version = dataflow.Environment.VersionValue() _verify_interpreter_version_is_supported(options) if self.standard_options.streaming: job_type = 'FNAPI_STREAMING' else: job_type = 'FNAPI_BATCH' self.proto.version.additionalProperties.extend([ dataflow.Environment.VersionValue.AdditionalProperty( key='job_type', value=to_json_value(job_type)), dataflow.Environment.VersionValue.AdditionalProperty( key='major', value=to_json_value(environment_version)) ]) # TODO: Use enumerated type instead of strings for job types. if job_type.startswith('FNAPI_'): self.debug_options.experiments = self.debug_options.experiments or [] debug_options_experiments = self.debug_options.experiments # Add use_multiple_sdk_containers flag if it's not already present. Do not # add the flag if 'no_use_multiple_sdk_containers' is present. # TODO: Cleanup use_multiple_sdk_containers once we deprecate Python SDK # till version 2.4. if ('use_multiple_sdk_containers' not in debug_options_experiments and 'no_use_multiple_sdk_containers' not in debug_options_experiments): debug_options_experiments.append('use_multiple_sdk_containers') # FlexRS if self.google_cloud_options.flexrs_goal == 'COST_OPTIMIZED': self.proto.flexResourceSchedulingGoal = ( dataflow.Environment.FlexResourceSchedulingGoalValueValuesEnum. FLEXRS_COST_OPTIMIZED) elif self.google_cloud_options.flexrs_goal == 'SPEED_OPTIMIZED': self.proto.flexResourceSchedulingGoal = ( dataflow.Environment.FlexResourceSchedulingGoalValueValuesEnum. FLEXRS_SPEED_OPTIMIZED) # Experiments if self.debug_options.experiments: for experiment in self.debug_options.experiments: self.proto.experiments.append(experiment) # Worker pool(s) information. package_descriptors = [] for package in packages: package_descriptors.append( dataflow.Package( location='%s/%s' % ( self.google_cloud_options.staging_location.replace( 'gs:/', GoogleCloudOptions.STORAGE_API_SERVICE), package), name=package)) pool = dataflow.WorkerPool( kind='local' if self.local else 'harness', packages=package_descriptors, taskrunnerSettings=dataflow.TaskRunnerSettings( parallelWorkerSettings=dataflow.WorkerSettings( baseUrl=GoogleCloudOptions.DATAFLOW_ENDPOINT, servicePath=self.google_cloud_options.dataflow_endpoint))) pool.autoscalingSettings = dataflow.AutoscalingSettings() # Set worker pool options received through command line. if self.worker_options.num_workers: pool.numWorkers = self.worker_options.num_workers if self.worker_options.max_num_workers: pool.autoscalingSettings.maxNumWorkers = ( self.worker_options.max_num_workers) if self.worker_options.autoscaling_algorithm: values_enum = dataflow.AutoscalingSettings.AlgorithmValueValuesEnum pool.autoscalingSettings.algorithm = { 'NONE': values_enum.AUTOSCALING_ALGORITHM_NONE, 'THROUGHPUT_BASED': values_enum.AUTOSCALING_ALGORITHM_BASIC, }.get(self.worker_options.autoscaling_algorithm) if self.worker_options.machine_type: pool.machineType = self.worker_options.machine_type if self.worker_options.disk_size_gb: pool.diskSizeGb = self.worker_options.disk_size_gb if self.worker_options.disk_type: pool.diskType = self.worker_options.disk_type if self.worker_options.disk_provisioned_iops is not None: pool.diskProvisionedIops = self.worker_options.disk_provisioned_iops if self.worker_options.disk_provisioned_throughput_mibps is not None: pool.diskProvisionedThroughputMibps = ( self.worker_options.disk_provisioned_throughput_mibps) if self.worker_options.zone: pool.zone = self.worker_options.zone if self.worker_options.network: pool.network = self.worker_options.network if self.worker_options.subnetwork: pool.subnetwork = self.worker_options.subnetwork # Setting worker pool sdk_harness_container_images option for supported # Dataflow workers. environments_to_use = self._get_environments_from_tranforms() # Adding container images for other SDKs that may be needed for # cross-language pipelines. for id, environment in environments_to_use: if environment.urn != common_urns.environments.DOCKER.urn: raise Exception( 'Dataflow can only execute pipeline steps in Docker environments.' ' Received %r.' % environment) environment_payload = proto_utils.parse_Bytes( environment.payload, beam_runner_api_pb2.DockerPayload) container_image_url = environment_payload.container_image container_image = dataflow.SdkHarnessContainerImage() container_image.containerImage = container_image_url container_image.useSingleCorePerContainer = ( common_urns.protocols.MULTI_CORE_BUNDLE_PROCESSING.urn not in environment.capabilities) container_image.environmentId = id for capability in environment.capabilities: container_image.capabilities.append(capability) pool.sdkHarnessContainerImages.append(container_image) if not pool.sdkHarnessContainerImages: pool.workerHarnessContainerImage = ( get_container_image_from_options(options)) elif len(pool.sdkHarnessContainerImages) == 1: # Dataflow expects a value here when there is only one environment. pool.workerHarnessContainerImage = ( pool.sdkHarnessContainerImages[0].containerImage) if self.debug_options.number_of_worker_harness_threads: pool.numThreadsPerWorker = ( self.debug_options.number_of_worker_harness_threads) if self.worker_options.use_public_ips is not None: if self.worker_options.use_public_ips: pool.ipConfiguration = ( dataflow.WorkerPool.IpConfigurationValueValuesEnum.WORKER_IP_PUBLIC) else: pool.ipConfiguration = ( dataflow.WorkerPool.IpConfigurationValueValuesEnum.WORKER_IP_PRIVATE ) if self.standard_options.streaming: # Use separate data disk for streaming. disk = dataflow.Disk() if self.local: disk.diskType = 'local' if self.worker_options.disk_type: disk.diskType = self.worker_options.disk_type pool.dataDisks.append(disk) self.proto.workerPools.append(pool) sdk_pipeline_options = options.get_all_options(retain_unknown_options=True) if sdk_pipeline_options: self.proto.sdkPipelineOptions = ( dataflow.Environment.SdkPipelineOptionsValue()) options_dict = { k: v for k, v in sdk_pipeline_options.items() if v is not None } options_dict["pipelineUrl"] = proto_pipeline_staged_url if pipeline_proto_hash: options_dict["pipelineProtoHash"] = pipeline_proto_hash # Don't pass impersonate_service_account through to the harness. # Though impersonation should start a job, the workers should # not try to modify their credentials. options_dict.pop('impersonate_service_account', None) self.proto.sdkPipelineOptions.additionalProperties.append( dataflow.Environment.SdkPipelineOptionsValue.AdditionalProperty( key='options', value=to_json_value(options_dict))) dd = DisplayData.create_from_options(options) items = [item.get_dict() for item in dd.items] self.proto.sdkPipelineOptions.additionalProperties.append( dataflow.Environment.SdkPipelineOptionsValue.AdditionalProperty( key='display_data', value=to_json_value(items))) if self.google_cloud_options.dataflow_service_options: for option in self.google_cloud_options.dataflow_service_options: self.proto.serviceOptions.append(option) if self.google_cloud_options.enable_hot_key_logging: self.proto.debugOptions = dataflow.DebugOptions(enableHotKeyLogging=True) def _get_environments_from_tranforms(self): if not self._proto_pipeline: return [] environment_ids = set( transform.environment_id for transform in self._proto_pipeline.components.transforms.values() if transform.environment_id) return [(id, self._proto_pipeline.components.environments[id]) for id in environment_ids] def _get_python_sdk_name(self): python_version = '%d.%d' % (sys.version_info[0], sys.version_info[1]) return 'Apache Beam Python %s SDK' % python_version class Job(object): """Wrapper for a dataflow Job protobuf.""" def __str__(self): def encode_shortstrings(input_buffer, errors='strict'): """Encoder (from Unicode) that suppresses long base64 strings.""" original_len = len(input_buffer) if original_len > 150: if self.base64_str_re.match(input_buffer): input_buffer = '' % original_len input_buffer = input_buffer.encode('ascii', errors=errors) else: matched = self.coder_str_re.match(input_buffer) if matched: input_buffer = '%s' % ( matched.group(1), matched.end(2) - matched.start(2)) input_buffer = input_buffer.encode('ascii', errors=errors) return input_buffer, original_len def decode_shortstrings(input_buffer, errors='strict'): """Decoder (to Unicode) that suppresses long base64 strings.""" shortened, length = encode_shortstrings(input_buffer, errors) return str(shortened), length def shortstrings_registerer(encoding_name): if encoding_name == 'shortstrings': return codecs.CodecInfo( name='shortstrings', encode=encode_shortstrings, decode=decode_shortstrings) return None codecs.register(shortstrings_registerer) # Use json "dump string" method to get readable formatting; # further modify it to not output too-long strings, aimed at the # 10,000+ character hex-encoded "serialized_fn" values. return json.dumps( json.loads(encoding.MessageToJson(self.proto)), indent=2, sort_keys=True) @staticmethod def _build_default_job_name(user_name): """Generates a default name for a job. user_name is lowercased, and any characters outside of [-a-z0-9] are removed. If necessary, the user_name is truncated to shorten the job name to 63 characters.""" user_name = re.sub('[^-a-z0-9]', '', user_name.lower()) date_component = datetime.now(timezone.utc).strftime('%m%d%H%M%S-%f') app_user_name = 'beamapp-{}'.format(user_name) # append 8 random alphanumeric characters to avoid collisions. random_component = ''.join( random.choices(string.ascii_lowercase + string.digits, k=8)) job_name = '{}-{}-{}'.format( app_user_name, date_component, random_component) if len(job_name) > 63: job_name = '{}-{}-{}'.format( app_user_name[:-(len(job_name) - 63)], date_component, random_component) return job_name @staticmethod def default_job_name(job_name): if job_name is None: job_name = Job._build_default_job_name(getpass.getuser()) return job_name def __init__(self, options, proto_pipeline): self.options = options validate_pipeline_graph(proto_pipeline) self.proto_pipeline = proto_pipeline self.google_cloud_options = options.view_as(GoogleCloudOptions) if not self.google_cloud_options.job_name: self.google_cloud_options.job_name = self.default_job_name( self.google_cloud_options.job_name) required_google_cloud_options = ['project', 'job_name', 'temp_location'] missing = [ option for option in required_google_cloud_options if not getattr(self.google_cloud_options, option) ] if missing: raise ValueError( 'Missing required configuration parameters: %s' % missing) if not self.google_cloud_options.staging_location: _LOGGER.info( 'Defaulting to the temp_location as staging_location: %s', self.google_cloud_options.temp_location) ( self.google_cloud_options.staging_location ) = self.google_cloud_options.temp_location self.root_staging_location = self.google_cloud_options.staging_location # Make the staging and temp locations job name and time specific. This is # needed to avoid clashes between job submissions using the same staging # area or team members using same job names. This method is not entirely # foolproof since two job submissions with same name can happen at exactly # the same time. However the window is extremely small given that # time.time() has at least microseconds granularity. We add the suffix only # for GCS staging locations where the potential for such clashes is high. if self.google_cloud_options.staging_location.startswith('gs://'): path_suffix = '%s.%f' % (self.google_cloud_options.job_name, time.time()) self.google_cloud_options.staging_location = FileSystems.join( self.google_cloud_options.staging_location, path_suffix) self.google_cloud_options.temp_location = FileSystems.join( self.google_cloud_options.temp_location, path_suffix) self.proto = dataflow.Job(name=self.google_cloud_options.job_name) if self.options.view_as(StandardOptions).streaming: self.proto.type = dataflow.Job.TypeValueValuesEnum.JOB_TYPE_STREAMING else: self.proto.type = dataflow.Job.TypeValueValuesEnum.JOB_TYPE_BATCH if self.google_cloud_options.update: self.proto.replaceJobId = self.job_id_for_name(self.proto.name) if self.google_cloud_options.transform_name_mapping: self.proto.transformNameMapping = ( dataflow.Job.TransformNameMappingValue()) for _, (key, value) in enumerate( self.google_cloud_options.transform_name_mapping.items()): self.proto.transformNameMapping.additionalProperties.append( dataflow.Job.TransformNameMappingValue.AdditionalProperty( key=key, value=value)) if self.google_cloud_options.create_from_snapshot: self.proto.createdFromSnapshotId = ( self.google_cloud_options.create_from_snapshot) # Labels. if self.google_cloud_options.labels: self.proto.labels = dataflow.Job.LabelsValue() labels = self.google_cloud_options.labels if isinstance(labels, str): labels = [labels] elif isinstance(labels, dict): labels = [str(labels)] for label in labels: if '{' in label: label = ast.literal_eval(label) for key, value in label.items(): self.proto.labels.additionalProperties.append( dataflow.Job.LabelsValue.AdditionalProperty( key=key, value=value)) else: parts = label.split('=', 1) key = parts[0] value = parts[1] if len(parts) > 1 else '' self.proto.labels.additionalProperties.append( dataflow.Job.LabelsValue.AdditionalProperty(key=key, value=value)) # Client Request ID self.proto.clientRequestId = '{}-{}'.format( datetime.now(timezone.utc).strftime('%Y%m%d%H%M%S%f'), random.randrange(9000) + 1000) self.base64_str_re = re.compile(r'^[A-Za-z0-9+/]*=*$') self.coder_str_re = re.compile(r'^([A-Za-z]+\$)([A-Za-z0-9+/]*=*)$') def job_id_for_name(self, job_name): return DataflowApplicationClient( self.google_cloud_options).job_id_for_name(job_name) def json(self): return encoding.MessageToJson(self.proto) def __reduce__(self): """Reduce hook for pickling the Job class more easily.""" return (Job, (self.options, )) class DataflowApplicationClient(object): _HASH_CHUNK_SIZE = 1024 * 8 _GCS_CACHE_PREFIX = "artifact_cache" """A Dataflow API client used by application code to create and query jobs.""" def __init__(self, options, root_staging_location=None): """Initializes a Dataflow API client object.""" self.standard_options = options.view_as(StandardOptions) self.google_cloud_options = options.view_as(GoogleCloudOptions) self._enable_caching = self.google_cloud_options.enable_artifact_caching self._enable_bucket_read_metric_counter = \ self.google_cloud_options.enable_bucket_read_metric_counter self._enable_bucket_write_metric_counter =\ self.google_cloud_options.enable_bucket_write_metric_counter self._root_staging_location = ( root_staging_location or self.google_cloud_options.staging_location) self.environment_version = _FNAPI_ENVIRONMENT_MAJOR_VERSION if self.google_cloud_options.no_auth: credentials = None else: credentials = get_service_credentials(options) http_client = get_new_http() self._client = dataflow.DataflowV1b3( url=self.google_cloud_options.dataflow_endpoint, credentials=credentials, get_credentials=(not self.google_cloud_options.no_auth), http=http_client, response_encoding=get_response_encoding()) self._storage_client = create_storage_client( options, not self.google_cloud_options.no_auth) self._sdk_image_overrides = self._get_sdk_image_overrides(options) def _get_sdk_image_overrides(self, pipeline_options): worker_options = pipeline_options.view_as(WorkerOptions) sdk_overrides = worker_options.sdk_harness_container_image_overrides return ( dict(s.split(',', 1) for s in sdk_overrides) if sdk_overrides else {}) @staticmethod def _compute_sha256(file): hasher = hashlib.sha256() with open(file, 'rb') as f: for chunk in iter(partial(f.read, DataflowApplicationClient._HASH_CHUNK_SIZE), b""): hasher.update(chunk) return hasher.hexdigest() def _cached_location(self, sha256): sha_prefix = sha256[0:2] return FileSystems.join( self._root_staging_location, DataflowApplicationClient._GCS_CACHE_PREFIX, sha_prefix, sha256) def _gcs_file_copy(self, from_path, to_path, sha256): if self._enable_caching and sha256: self._cached_gcs_file_copy(from_path, to_path, sha256) else: self._uncached_gcs_file_copy(from_path, to_path) def _cached_gcs_file_copy(self, from_path, to_path, sha256): cached_path = self._cached_location(sha256) if FileSystems.exists(cached_path): _LOGGER.info( 'Skipping upload of %s because it already exists at %s', to_path, cached_path) else: self._uncached_gcs_file_copy(from_path, cached_path) FileSystems.copy( source_file_names=[cached_path], destination_file_names=[to_path]) _LOGGER.info('Copied cached artifact from %s to %s', from_path, to_path) def _uncached_gcs_file_copy(self, from_path, to_path): to_folder, to_name = os.path.split(to_path) total_size = os.path.getsize(from_path) self.stage_file_with_retry( to_folder, to_name, from_path, total_size=total_size) def _stage_resources(self, pipeline, options): google_cloud_options = options.view_as(GoogleCloudOptions) if google_cloud_options.staging_location is None: raise RuntimeError('The --staging_location option must be specified.') if google_cloud_options.temp_location is None: raise RuntimeError('The --temp_location option must be specified.') resources = [] staged_paths = {} staged_hashes = {} for _, env in sorted(pipeline.components.environments.items(), key=lambda kv: kv[0]): for dep in env.dependencies: if dep.type_urn != common_urns.artifact_types.FILE.urn: raise RuntimeError('unsupported artifact type %s' % dep.type_urn) type_payload = beam_runner_api_pb2.ArtifactFilePayload.FromString( dep.type_payload) if dep.role_urn == common_urns.artifact_roles.STAGING_TO.urn: remote_name = ( beam_runner_api_pb2.ArtifactStagingToRolePayload.FromString( dep.role_payload)).staged_name is_staged_role = True else: remote_name = os.path.basename(type_payload.path) is_staged_role = False # compute sha256 even if caching is disabled. # This is used to check the payload integrity along with caching. if not type_payload.sha256: type_payload.sha256 = self._compute_sha256(type_payload.path) if type_payload.sha256 and type_payload.sha256 in staged_hashes: _LOGGER.info( 'Found duplicated artifact sha256: %s (%s)', type_payload.path, type_payload.sha256) remote_name = staged_hashes[type_payload.sha256] if is_staged_role: # We should not be overriding this, as dep.role_payload.staged_name # refers to the desired name on the worker, whereas staged_name # refers to its placement in a distributed filesystem. # TODO(heejong): Clean this up. dep.role_payload = beam_runner_api_pb2.ArtifactStagingToRolePayload( staged_name=remote_name).SerializeToString() elif type_payload.path and type_payload.path in staged_paths: _LOGGER.info( 'Found duplicated artifact path: %s (%s)', type_payload.path, type_payload.sha256) remote_name = staged_paths[type_payload.path] if is_staged_role: # We should not be overriding this, as dep.role_payload.staged_name # refers to the desired name on the worker, whereas staged_name # refers to its placement in a distributed filesystem. # TODO(heejong): Clean this up. dep.role_payload = beam_runner_api_pb2.ArtifactStagingToRolePayload( staged_name=remote_name).SerializeToString() else: resources.append( (type_payload.path, remote_name, type_payload.sha256)) staged_paths[type_payload.path] = remote_name staged_hashes[type_payload.sha256] = remote_name if FileSystems.get_scheme( google_cloud_options.staging_location) == GCSFileSystem.scheme(): dep.type_urn = common_urns.artifact_types.URL.urn dep.type_payload = beam_runner_api_pb2.ArtifactUrlPayload( url=FileSystems.join( google_cloud_options.staging_location, remote_name), sha256=type_payload.sha256).SerializeToString() else: dep.type_payload = beam_runner_api_pb2.ArtifactFilePayload( path=FileSystems.join( google_cloud_options.staging_location, remote_name), sha256=type_payload.sha256).SerializeToString() resource_stager = _LegacyDataflowStager(self) staged_resources = resource_stager.stage_job_resources( resources, staging_location=google_cloud_options.staging_location) return staged_resources def stage_file( self, gcs_or_local_path, file_name, stream, mime_type='application/octet-stream', total_size=None): """Stages a file at a GCS or local path with stream-supplied contents.""" from google.cloud.exceptions import Forbidden from google.cloud.exceptions import NotFound if not gcs_or_local_path.startswith('gs://'): local_path = FileSystems.join(gcs_or_local_path, file_name) _LOGGER.info('Staging file locally to %s', local_path) with open(local_path, 'wb') as f: f.write(stream.read()) return gcs_location = FileSystems.join(gcs_or_local_path, file_name) bucket_name, blob_name = gcs_location[5:].split('/', 1) start_time = time.time() _LOGGER.info('Starting GCS upload to %s...', gcs_location) try: from google.cloud.storage import Blob from google.cloud.storage.fileio import BlobWriter bucket = self._storage_client.get_bucket(bucket_name) blob = bucket.get_blob(blob_name) if not blob: blob = Blob(blob_name, bucket) with BlobWriter(blob) as f: f.write(stream.read()) _LOGGER.info( 'Completed GCS upload to %s in %s seconds.', gcs_location, int(time.time() - start_time)) return except Exception as e: reportable_errors = [ Forbidden, NotFound, ] if type(e) in reportable_errors: raise IOError(( 'Could not upload to GCS path %s: %s. Please verify ' 'that credentials are valid, that the specified path ' 'exists, and that you have write access to it.') % (gcs_or_local_path, e)) raise @retry.with_exponential_backoff( retry_filter=retry.retry_on_server_errors_and_timeout_filter) def stage_file_with_retry( self, gcs_or_local_path, file_name, stream_or_path, mime_type='application/octet-stream', total_size=None): if isinstance(stream_or_path, str): path = stream_or_path with open(path, 'rb') as stream: self.stage_file( gcs_or_local_path, file_name, stream, mime_type, total_size) elif isinstance(stream_or_path, io.IOBase): stream = stream_or_path try: self.stage_file( gcs_or_local_path, file_name, stream, mime_type, total_size) except Exception as exn: if stream.seekable(): # reset cursor for possible retrying stream.seek(0) raise exn else: raise retry.PermanentException( "Skip retrying because we caught exception:" + ''.join(traceback.format_exception_only(exn.__class__, exn)) + ', but the stream is not seekable.') else: raise retry.PermanentException( "Skip retrying because type " + str(type(stream_or_path)) + "stream_or_path is unsupported.") @retry.no_retries # Using no_retries marks this as an integration point. def create_job(self, job): """Creates job description. May stage and/or submit for remote execution.""" self.create_job_description(job) # Stage and submit the job when necessary dataflow_job_file = job.options.view_as(DebugOptions).dataflow_job_file template_location = ( job.options.view_as(GoogleCloudOptions).template_location) if job.options.view_as(DebugOptions).lookup_experiment('upload_graph'): self.stage_file_with_retry( job.options.view_as(GoogleCloudOptions).staging_location, "dataflow_graph.json", io.BytesIO(job.json().encode('utf-8'))) del job.proto.steps[:] job.proto.stepsLocation = FileSystems.join( job.options.view_as(GoogleCloudOptions).staging_location, "dataflow_graph.json") # template file generation should be placed immediately before the # conditional API call. job_location = template_location or dataflow_job_file if job_location: gcs_or_local_path = os.path.dirname(job_location) file_name = os.path.basename(job_location) self.stage_file_with_retry( gcs_or_local_path, file_name, io.BytesIO(job.json().encode('utf-8'))) if not template_location: return self.submit_job_description(job) _LOGGER.info( 'A template was just created at location %s', template_location) return None @staticmethod def _update_container_image_for_dataflow(beam_container_image_url): # By default Dataflow pipelines use containers hosted in Dataflow GCR # instead of Docker Hub. image_suffix = beam_container_image_url.rsplit('/', 1)[1] # trim "RCX" as release candidate tag exists on Docker Hub but not GCR check_rc = image_suffix.lower().split('rc') if len(check_rc) == 2: image_suffix = image_suffix[:-2 - len(check_rc[1])] return names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + '/' + image_suffix @staticmethod def _apply_sdk_environment_overrides( proto_pipeline, sdk_overrides, pipeline_options): # Updates container image URLs for Dataflow. # For a given container image URL # * If a matching override has been provided that will be used. # * For improved performance, External Apache Beam container images that are # not explicitly overridden will be # updated to use GCR copies instead of directly downloading from the # Docker Hub. current_sdk_container_image = get_container_image_from_options( pipeline_options) for environment in proto_pipeline.components.environments.values(): docker_payload = proto_utils.parse_Bytes( environment.payload, beam_runner_api_pb2.DockerPayload) overridden = False new_container_image = docker_payload.container_image for pattern, override in sdk_overrides.items(): new_container_image = re.sub(pattern, override, new_container_image) if new_container_image != docker_payload.container_image: overridden = True # Container of the current (Python) SDK is overridden separately, hence # not updated here. if (is_apache_beam_container(new_container_image) and not overridden and new_container_image != current_sdk_container_image): new_container_image = ( DataflowApplicationClient._update_container_image_for_dataflow( docker_payload.container_image)) if not new_container_image: raise ValueError( 'SDK Docker container image has to be a non-empty string') new_payload = copy(docker_payload) new_payload.container_image = new_container_image environment.payload = new_payload.SerializeToString() def create_job_description(self, job): """Creates a job described by the workflow proto.""" DataflowApplicationClient._apply_sdk_environment_overrides( job.proto_pipeline, self._sdk_image_overrides, job.options) # Stage other resources for the SDK harness resources = self._stage_resources(job.proto_pipeline, job.options) # Stage proto pipeline. serialized_pipeline = job.proto_pipeline.SerializeToString() pipeline_proto_hash = hashlib.sha256(serialized_pipeline).hexdigest() self.stage_file_with_retry( job.google_cloud_options.staging_location, shared_names.STAGED_PIPELINE_FILENAME, io.BytesIO(serialized_pipeline)) job.proto.environment = Environment( proto_pipeline_staged_url=FileSystems.join( job.google_cloud_options.staging_location, shared_names.STAGED_PIPELINE_FILENAME), packages=resources, options=job.options, environment_version=self.environment_version, proto_pipeline=job.proto_pipeline, pipeline_proto_hash=pipeline_proto_hash).proto _LOGGER.debug('JOB: %s', job) @retry.with_exponential_backoff(num_retries=3, initial_delay_secs=3) def get_job_metrics(self, job_id): request = dataflow.DataflowProjectsLocationsJobsGetMetricsRequest() request.jobId = job_id request.location = self.google_cloud_options.region request.projectId = self.google_cloud_options.project try: response = self._client.projects_locations_jobs.GetMetrics(request) except exceptions.BadStatusCodeError as e: _LOGGER.error( 'HTTP status %d. Unable to query metrics', e.response.status) raise return response @retry.with_exponential_backoff(num_retries=3) def submit_job_description(self, job): """Creates and excutes a job request.""" request = dataflow.DataflowProjectsLocationsJobsCreateRequest() request.projectId = self.google_cloud_options.project request.location = self.google_cloud_options.region request.job = job.proto try: response = self._client.projects_locations_jobs.Create(request) except exceptions.BadStatusCodeError as e: _LOGGER.error( 'HTTP status %d trying to create job' ' at dataflow service endpoint %s', e.response.status, self.google_cloud_options.dataflow_endpoint) _LOGGER.fatal('details of server error: %s', e) raise if response.clientRequestId and \ response.clientRequestId != job.proto.clientRequestId: if self.google_cloud_options.update: raise DataflowJobAlreadyExistsError( "The job named %s with id: %s has already been updated into job " "id: %s and cannot be updated again." % (response.name, job.proto.replaceJobId, response.id)) else: raise DataflowJobAlreadyExistsError( 'There is already active job named %s with id: %s. If you want to ' 'submit a second job, try again by setting a different name using ' '--job_name.' % (response.name, response.id)) _LOGGER.info('Create job: %s', response) # The response is a Job proto with the id for the new job. _LOGGER.info('Created job with id: [%s]', response.id) _LOGGER.info('Submitted job: %s', response.id) _LOGGER.info( 'To access the Dataflow monitoring console, please navigate to ' 'https://console.cloud.google.com/dataflow/jobs/%s/%s?project=%s', self.google_cloud_options.region, response.id, self.google_cloud_options.project) return response @retry.with_exponential_backoff() # Using retry defaults from utils/retry.py def modify_job_state(self, job_id, new_state): """Modify the run state of the job. Args: job_id: The id of the job. new_state: A string representing the new desired state. It could be set to either 'JOB_STATE_DONE', 'JOB_STATE_CANCELLED' or 'JOB_STATE_DRAINING'. Returns: True if the job was modified successfully. """ if new_state == 'JOB_STATE_DONE': new_state = dataflow.Job.RequestedStateValueValuesEnum.JOB_STATE_DONE elif new_state == 'JOB_STATE_CANCELLED': new_state = dataflow.Job.RequestedStateValueValuesEnum.JOB_STATE_CANCELLED elif new_state == 'JOB_STATE_DRAINING': new_state = dataflow.Job.RequestedStateValueValuesEnum.JOB_STATE_DRAINING else: # Other states could only be set by the service. return False request = dataflow.DataflowProjectsLocationsJobsUpdateRequest() request.jobId = job_id request.projectId = self.google_cloud_options.project request.location = self.google_cloud_options.region request.job = dataflow.Job(requestedState=new_state) self._client.projects_locations_jobs.Update(request) return True @retry.with_exponential_backoff( retry_filter=retry.retry_on_server_errors_and_notfound_filter) def get_job(self, job_id): """Gets the job status for a submitted job. Args: job_id: A string representing the job_id for the workflow as returned by the create_job() request. Returns: A Job proto. See below for interesting fields. The Job proto returned from a get_job() request contains some interesting fields: currentState: An object representing the current state of the job. The string representation of the object (str() result) has the following possible values: JOB_STATE_UNKNONW, JOB_STATE_STOPPED, JOB_STATE_RUNNING, JOB_STATE_DONE, JOB_STATE_FAILED, JOB_STATE_CANCELLED. createTime: UTC time when the job was created (e.g. '2015-03-10T00:01:53.074Z') currentStateTime: UTC time for the current state of the job. """ request = dataflow.DataflowProjectsLocationsJobsGetRequest() request.jobId = job_id request.projectId = self.google_cloud_options.project request.location = self.google_cloud_options.region response = self._client.projects_locations_jobs.Get(request) return response @retry.with_exponential_backoff( retry_filter=retry.retry_on_server_errors_and_notfound_filter) def list_messages( self, job_id, start_time=None, end_time=None, page_token=None, minimum_importance=None): """List messages associated with the execution of a job. Args: job_id: A string representing the job_id for the workflow as returned by the create_job() request. start_time: If specified, only messages generated after the start time will be returned, otherwise all messages since job started will be returned. The value is a string representing UTC time (e.g., '2015-08-18T21:03:50.644Z') end_time: If specified, only messages generated before the end time will be returned, otherwise all messages up to current time will be returned. The value is a string representing UTC time (e.g., '2015-08-18T21:03:50.644Z') page_token: A string to be used as next page token if the list call returned paginated results. minimum_importance: Filter for messages based on importance. The possible string values in increasing order of importance are: JOB_MESSAGE_DEBUG, JOB_MESSAGE_DETAILED, JOB_MESSAGE_BASIC, JOB_MESSAGE_WARNING, JOB_MESSAGE_ERROR. For example, a filter set on warning will allow only warnings and errors and exclude all others. Returns: A tuple consisting of a list of JobMessage instances and a next page token string. Raises: RuntimeError: if an unexpected value for the message_importance argument is used. The JobMessage objects returned by the call contain the following fields: id: A unique string identifier for the message. time: A string representing the UTC time of the message (e.g., '2015-08-18T21:03:50.644Z') messageImportance: An enumeration value for the message importance. The value if converted to string will have the following possible values: JOB_MESSAGE_DEBUG, JOB_MESSAGE_DETAILED, JOB_MESSAGE_BASIC, JOB_MESSAGE_WARNING, JOB_MESSAGE_ERROR. messageText: A message string. """ request = dataflow.DataflowProjectsLocationsJobsMessagesListRequest( jobId=job_id, location=self.google_cloud_options.region, projectId=self.google_cloud_options.project) if page_token is not None: request.pageToken = page_token if start_time is not None: request.startTime = start_time if end_time is not None: request.endTime = end_time if minimum_importance is not None: if minimum_importance == 'JOB_MESSAGE_DEBUG': request.minimumImportance = ( dataflow.DataflowProjectsLocationsJobsMessagesListRequest. MinimumImportanceValueValuesEnum.JOB_MESSAGE_DEBUG) elif minimum_importance == 'JOB_MESSAGE_DETAILED': request.minimumImportance = ( dataflow.DataflowProjectsLocationsJobsMessagesListRequest. MinimumImportanceValueValuesEnum.JOB_MESSAGE_DETAILED) elif minimum_importance == 'JOB_MESSAGE_BASIC': request.minimumImportance = ( dataflow.DataflowProjectsLocationsJobsMessagesListRequest. MinimumImportanceValueValuesEnum.JOB_MESSAGE_BASIC) elif minimum_importance == 'JOB_MESSAGE_WARNING': request.minimumImportance = ( dataflow.DataflowProjectsLocationsJobsMessagesListRequest. MinimumImportanceValueValuesEnum.JOB_MESSAGE_WARNING) elif minimum_importance == 'JOB_MESSAGE_ERROR': request.minimumImportance = ( dataflow.DataflowProjectsLocationsJobsMessagesListRequest. MinimumImportanceValueValuesEnum.JOB_MESSAGE_ERROR) else: raise RuntimeError( 'Unexpected value for minimum_importance argument: %r' % minimum_importance) response = self._client.projects_locations_jobs_messages.List(request) return response.jobMessages, response.nextPageToken def job_id_for_name(self, job_name): token = None while True: request = dataflow.DataflowProjectsLocationsJobsListRequest( projectId=self.google_cloud_options.project, location=self.google_cloud_options.region, pageToken=token) response = self._client.projects_locations_jobs.List(request) for job in response.jobs: if (job.name == job_name and job.currentState in [dataflow.Job.CurrentStateValueValuesEnum.JOB_STATE_RUNNING, dataflow.Job.CurrentStateValueValuesEnum.JOB_STATE_DRAINING]): return job.id token = response.nextPageToken if token is None: raise ValueError("No running job found with name '%s'" % job_name) class MetricUpdateTranslators(object): """Translators between accumulators and dataflow metric updates.""" @staticmethod def translate_boolean(accumulator, metric_update_proto): metric_update_proto.boolean = accumulator.value @staticmethod def translate_scalar_mean_int(accumulator, metric_update_proto): if accumulator.count: metric_update_proto.integerMean = dataflow.IntegerMean() metric_update_proto.integerMean.sum = to_split_int(accumulator.sum) metric_update_proto.integerMean.count = to_split_int(accumulator.count) else: metric_update_proto.nameAndKind.kind = None @staticmethod def translate_scalar_mean_float(accumulator, metric_update_proto): if accumulator.count: metric_update_proto.floatingPointMean = dataflow.FloatingPointMean() metric_update_proto.floatingPointMean.sum = accumulator.sum metric_update_proto.floatingPointMean.count = to_split_int( accumulator.count) else: metric_update_proto.nameAndKind.kind = None @staticmethod def translate_scalar_counter_int(accumulator, metric_update_proto): metric_update_proto.integer = to_split_int(accumulator.value) @staticmethod def translate_scalar_counter_float(accumulator, metric_update_proto): metric_update_proto.floatingPoint = accumulator.value class _LegacyDataflowStager(Stager): def __init__(self, dataflow_application_client): super().__init__() self._dataflow_application_client = dataflow_application_client def stage_artifact(self, local_path_to_artifact, artifact_name, sha256): self._dataflow_application_client._gcs_file_copy( local_path_to_artifact, artifact_name, sha256) def commit_manifest(self): pass @staticmethod def get_sdk_package_name(): """For internal use only; no backwards-compatibility guarantees. Returns the PyPI package name to be staged to Google Cloud Dataflow. """ return shared_names.BEAM_PACKAGE_NAME class DataflowJobAlreadyExistsError(retry.PermanentException): """A non-retryable exception that a job with the given name already exists.""" # Inherits retry.PermanentException to avoid retry in # DataflowApplicationClient.submit_job_description pass def to_split_int(n): res = dataflow.SplitInt64() res.lowBits = n & 0xffffffff res.highBits = n >> 32 return res # TODO: Used in legacy batch worker. Move under MetricUpdateTranslators # after Runner V2 transition. def translate_distribution(distribution_update, metric_update_proto): """Translate metrics DistributionUpdate to dataflow distribution update. Args: distribution_update: Instance of DistributionData, DistributionInt64Accumulator or DataflowDistributionCounter. metric_update_proto: Used for report metrics. """ dist_update_proto = dataflow.DistributionUpdate() dist_update_proto.min = to_split_int(distribution_update.min) dist_update_proto.max = to_split_int(distribution_update.max) dist_update_proto.count = to_split_int(distribution_update.count) dist_update_proto.sum = to_split_int(distribution_update.sum) # DataflowDistributionCounter needs to translate histogram if isinstance(distribution_update, DataflowDistributionCounter): dist_update_proto.histogram = dataflow.Histogram() distribution_update.translate_to_histogram(dist_update_proto.histogram) metric_update_proto.distribution = dist_update_proto # TODO: Used in legacy batch worker. Delete after Runner V2 transition. def translate_value(value, metric_update_proto): metric_update_proto.integer = to_split_int(value) def _get_container_image_tag(): base_version = version.parse(beam_version.__version__).base_version if base_version != beam_version.__version__: warnings.warn( "A non-standard version of Beam SDK detected: %s. " "Dataflow runner will use container image tag %s. " "This use case is not supported." % (beam_version.__version__, base_version)) return base_version def get_container_image_from_options(pipeline_options): """For internal use only; no backwards-compatibility guarantees. Args: pipeline_options (PipelineOptions): A container for pipeline options. Returns: str: Container image for remote execution. """ worker_options = pipeline_options.view_as(WorkerOptions) if worker_options.sdk_container_image: return worker_options.sdk_container_image # Legacy and runner v2 exist in different repositories. # Set to legacy format, override if runner v2 container_repo = names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY image_name = '{repository}/beam_python{major}.{minor}_sdk'.format( repository=container_repo, major=sys.version_info[0], minor=sys.version_info[1]) image_tag = _get_required_container_version() return image_name + ':' + image_tag def _get_required_container_version(): """For internal use only; no backwards-compatibility guarantees. Returns: str: The tag of worker container images in GCR that corresponds to current version of the SDK. """ if 'dev' in beam_version.__version__: return names.BEAM_DEV_SDK_CONTAINER_TAG else: return _get_container_image_tag() def get_response_encoding(): """Encoding to use to decode HTTP response from Google APIs.""" return 'utf8' def _verify_interpreter_version_is_supported(pipeline_options): if ('%s.%s' % (sys.version_info[0], sys.version_info[1]) in _PYTHON_VERSIONS_SUPPORTED_BY_DATAFLOW): return if 'dev' in beam_version.__version__: return debug_options = pipeline_options.view_as(DebugOptions) if (debug_options.experiments and 'use_unsupported_python_version' in debug_options.experiments): return raise Exception( 'Dataflow runner currently supports Python versions %s, got %s.\n' 'To ignore this requirement and start a job ' 'using an unsupported version of Python interpreter, pass ' '--experiment use_unsupported_python_version pipeline option.' % (_PYTHON_VERSIONS_SUPPORTED_BY_DATAFLOW, sys.version)) # To enable a counter on the service, add it to this dictionary. # This is required for the legacy python dataflow runner, as portability # does not communicate to the service via python code, but instead via a # a runner harness (in C++ or Java). # TODO(https://github.com/apache/beam/issues/19433) : Remove this antipattern, # legacy dataflow python pipelines will break whenever a new cy_combiner type # is used. structured_counter_translations = { cy_combiners.CountCombineFn: ( dataflow.CounterMetadata.KindValueValuesEnum.SUM, MetricUpdateTranslators.translate_scalar_counter_int), cy_combiners.SumInt64Fn: ( dataflow.CounterMetadata.KindValueValuesEnum.SUM, MetricUpdateTranslators.translate_scalar_counter_int), cy_combiners.MinInt64Fn: ( dataflow.CounterMetadata.KindValueValuesEnum.MIN, MetricUpdateTranslators.translate_scalar_counter_int), cy_combiners.MaxInt64Fn: ( dataflow.CounterMetadata.KindValueValuesEnum.MAX, MetricUpdateTranslators.translate_scalar_counter_int), cy_combiners.MeanInt64Fn: ( dataflow.CounterMetadata.KindValueValuesEnum.MEAN, MetricUpdateTranslators.translate_scalar_mean_int), cy_combiners.SumFloatFn: ( dataflow.CounterMetadata.KindValueValuesEnum.SUM, MetricUpdateTranslators.translate_scalar_counter_float), cy_combiners.MinFloatFn: ( dataflow.CounterMetadata.KindValueValuesEnum.MIN, MetricUpdateTranslators.translate_scalar_counter_float), cy_combiners.MaxFloatFn: ( dataflow.CounterMetadata.KindValueValuesEnum.MAX, MetricUpdateTranslators.translate_scalar_counter_float), cy_combiners.MeanFloatFn: ( dataflow.CounterMetadata.KindValueValuesEnum.MEAN, MetricUpdateTranslators.translate_scalar_mean_float), cy_combiners.AllCombineFn: ( dataflow.CounterMetadata.KindValueValuesEnum.AND, MetricUpdateTranslators.translate_boolean), cy_combiners.AnyCombineFn: ( dataflow.CounterMetadata.KindValueValuesEnum.OR, MetricUpdateTranslators.translate_boolean), cy_combiners.DataflowDistributionCounterFn: ( dataflow.CounterMetadata.KindValueValuesEnum.DISTRIBUTION, translate_distribution), cy_combiners.DistributionInt64Fn: ( dataflow.CounterMetadata.KindValueValuesEnum.DISTRIBUTION, translate_distribution), } counter_translations = { cy_combiners.CountCombineFn: ( dataflow.NameAndKind.KindValueValuesEnum.SUM, MetricUpdateTranslators.translate_scalar_counter_int), cy_combiners.SumInt64Fn: ( dataflow.NameAndKind.KindValueValuesEnum.SUM, MetricUpdateTranslators.translate_scalar_counter_int), cy_combiners.MinInt64Fn: ( dataflow.NameAndKind.KindValueValuesEnum.MIN, MetricUpdateTranslators.translate_scalar_counter_int), cy_combiners.MaxInt64Fn: ( dataflow.NameAndKind.KindValueValuesEnum.MAX, MetricUpdateTranslators.translate_scalar_counter_int), cy_combiners.MeanInt64Fn: ( dataflow.NameAndKind.KindValueValuesEnum.MEAN, MetricUpdateTranslators.translate_scalar_mean_int), cy_combiners.SumFloatFn: ( dataflow.NameAndKind.KindValueValuesEnum.SUM, MetricUpdateTranslators.translate_scalar_counter_float), cy_combiners.MinFloatFn: ( dataflow.NameAndKind.KindValueValuesEnum.MIN, MetricUpdateTranslators.translate_scalar_counter_float), cy_combiners.MaxFloatFn: ( dataflow.NameAndKind.KindValueValuesEnum.MAX, MetricUpdateTranslators.translate_scalar_counter_float), cy_combiners.MeanFloatFn: ( dataflow.NameAndKind.KindValueValuesEnum.MEAN, MetricUpdateTranslators.translate_scalar_mean_float), cy_combiners.AllCombineFn: ( dataflow.NameAndKind.KindValueValuesEnum.AND, MetricUpdateTranslators.translate_boolean), cy_combiners.AnyCombineFn: ( dataflow.NameAndKind.KindValueValuesEnum.OR, MetricUpdateTranslators.translate_boolean), cy_combiners.DataflowDistributionCounterFn: ( dataflow.NameAndKind.KindValueValuesEnum.DISTRIBUTION, translate_distribution), cy_combiners.DistributionInt64Fn: ( dataflow.CounterMetadata.KindValueValuesEnum.DISTRIBUTION, translate_distribution), } ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/internal/apiclient_test.py0000644000175100017510000023066415206113032026362 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the apiclient module.""" # pytype: skip-file import hashlib import io import itertools import json import logging import os import sys import time import unittest import mock from apache_beam.io.filesystems import FileSystems from apache_beam.metrics.cells import DistributionData from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.pipeline import Pipeline from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.dataflow.internal import names from apache_beam.runners.dataflow.internal.clients import dataflow from apache_beam.transforms import Create from apache_beam.transforms import DataflowDistributionCounter from apache_beam.transforms import DoFn from apache_beam.transforms import ParDo from apache_beam.transforms.environments import DockerEnvironment from apache_beam.utils import retry # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from apache_beam.runners.dataflow.internal import apiclient except ImportError: apiclient = None # type: ignore # pylint: enable=wrong-import-order, wrong-import-position FAKE_PIPELINE_URL = "gs://invalid-bucket/anywhere" _LOGGER = logging.getLogger(__name__) @unittest.skipIf(apiclient is None, 'GCP dependencies are not installed') class UtilTest(unittest.TestCase): @unittest.skip("Enable once BEAM-1080 is fixed.") def test_create_application_client(self): pipeline_options = PipelineOptions() apiclient.DataflowApplicationClient(pipeline_options) def test_pipeline_url(self): pipeline_options = PipelineOptions([ '--subnetwork', '/regions/MY/subnetworks/SUBNETWORK', '--temp_location', 'gs://any-location/temp' ]) env = apiclient.Environment( [], pipeline_options, '2.0.0', # any environment version FAKE_PIPELINE_URL) recovered_options = None for additionalProperty in env.proto.sdkPipelineOptions.additionalProperties: if additionalProperty.key == 'options': recovered_options = additionalProperty.value break else: self.fail( 'No pipeline options found in %s' % env.proto.sdkPipelineOptions) pipeline_url = None for property in recovered_options.object_value.properties: if property.key == 'pipelineUrl': pipeline_url = property.value break else: self.fail('No pipeline_url found in %s' % recovered_options) self.assertEqual(pipeline_url.string_value, FAKE_PIPELINE_URL) def test_pipeline_proto_hash(self): pipeline_options = PipelineOptions( ['--temp_location', 'gs://any-location/temp']) proto_pipeline = beam_runner_api_pb2.Pipeline() proto_pipeline.components.transforms['dummy'].unique_name = 'dummy' expected_hash = hashlib.sha256( proto_pipeline.SerializeToString()).hexdigest() env = apiclient.Environment([], pipeline_options, '2.0.0', FAKE_PIPELINE_URL, proto_pipeline, pipeline_proto_hash=expected_hash) recovered_options = None for additionalProperty in env.proto.sdkPipelineOptions.additionalProperties: if additionalProperty.key == 'options': recovered_options = additionalProperty.value break else: self.fail('No pipeline options found') pipeline_proto_hash = None for property in recovered_options.object_value.properties: if property.key == 'pipelineProtoHash': pipeline_proto_hash = property.value break else: self.fail('No pipelineProtoHash found') self.assertEqual(pipeline_proto_hash.string_value, expected_hash) def test_set_network(self): pipeline_options = PipelineOptions([ '--network', 'anetworkname', '--temp_location', 'gs://any-location/temp' ]) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual(env.proto.workerPools[0].network, 'anetworkname') def test_set_subnetwork(self): pipeline_options = PipelineOptions([ '--subnetwork', '/regions/MY/subnetworks/SUBNETWORK', '--temp_location', 'gs://any-location/temp' ]) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual( env.proto.workerPools[0].subnetwork, '/regions/MY/subnetworks/SUBNETWORK') def test_flexrs_blank(self): pipeline_options = PipelineOptions( ['--temp_location', 'gs://any-location/temp']) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual(env.proto.flexResourceSchedulingGoal, None) def test_flexrs_cost(self): pipeline_options = PipelineOptions([ '--flexrs_goal', 'COST_OPTIMIZED', '--temp_location', 'gs://any-location/temp' ]) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual( env.proto.flexResourceSchedulingGoal, ( dataflow.Environment.FlexResourceSchedulingGoalValueValuesEnum. FLEXRS_COST_OPTIMIZED)) def test_flexrs_speed(self): pipeline_options = PipelineOptions([ '--flexrs_goal', 'SPEED_OPTIMIZED', '--temp_location', 'gs://any-location/temp' ]) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual( env.proto.flexResourceSchedulingGoal, ( dataflow.Environment.FlexResourceSchedulingGoalValueValuesEnum. FLEXRS_SPEED_OPTIMIZED)) def _verify_sdk_harness_container_images_get_set(self, pipeline_options): pipeline = Pipeline(options=pipeline_options) pipeline | Create([1, 2, 3]) | ParDo(DoFn()) # pylint:disable=expression-not-assigned test_environment = DockerEnvironment(container_image='test_default_image') proto_pipeline, _ = pipeline.to_runner_api( return_context=True, default_environment=test_environment) dummy_env = beam_runner_api_pb2.Environment( urn=common_urns.environments.DOCKER.urn, payload=( beam_runner_api_pb2.DockerPayload( container_image='dummy_image')).SerializeToString()) dummy_env.capabilities.append( common_urns.protocols.MULTI_CORE_BUNDLE_PROCESSING.urn) proto_pipeline.components.environments['dummy_env_id'].CopyFrom(dummy_env) dummy_transform = beam_runner_api_pb2.PTransform( environment_id='dummy_env_id') proto_pipeline.components.transforms['dummy_transform_id'].CopyFrom( dummy_transform) env = apiclient.Environment( [], # packages pipeline_options, '2.0.0', # any environment version FAKE_PIPELINE_URL, proto_pipeline) worker_pool = env.proto.workerPools[0] self.assertEqual(2, len(worker_pool.sdkHarnessContainerImages)) # Only one of the environments is missing MULTI_CORE_BUNDLE_PROCESSING. self.assertEqual( 1, sum( c.useSingleCorePerContainer for c in worker_pool.sdkHarnessContainerImages)) env_and_image = [(item.environmentId, item.containerImage) for item in worker_pool.sdkHarnessContainerImages] self.assertIn(('dummy_env_id', 'dummy_image'), env_and_image) self.assertIn((mock.ANY, 'test_default_image'), env_and_image) def test_sdk_harness_container_images_get_set_runner_v2(self): pipeline_options = PipelineOptions([ '--experiments=use_runner_v2', '--temp_location', 'gs://any-location/temp' ]) self._verify_sdk_harness_container_images_get_set(pipeline_options) def test_sdk_harness_container_images_get_set_prime(self): pipeline_options = PipelineOptions([ '--dataflow_service_options=enable_prime', '--temp_location', 'gs://any-location/temp' ]) self._verify_sdk_harness_container_images_get_set(pipeline_options) def _verify_sdk_harness_container_image_overrides(self, pipeline_options): test_environment = DockerEnvironment( container_image='dummy_container_image') proto_pipeline, _ = Pipeline().to_runner_api( return_context=True, default_environment=test_environment) # Accessing non-public method for testing. apiclient.DataflowApplicationClient._apply_sdk_environment_overrides( proto_pipeline, { '.*dummy.*': 'new_dummy_container_image', '.*notfound.*': 'new_dummy_container_image_2' }, pipeline_options) self.assertIsNotNone(1, len(proto_pipeline.components.environments)) env = list(proto_pipeline.components.environments.values())[0] from apache_beam.utils import proto_utils docker_payload = proto_utils.parse_Bytes( env.payload, beam_runner_api_pb2.DockerPayload) # Container image should be overridden by the given override. self.assertEqual( docker_payload.container_image, 'new_dummy_container_image') def test_sdk_harness_container_image_overrides_runner_v2(self): pipeline_options = PipelineOptions([ '--experiments=use_runner_v2', '--temp_location', 'gs://any-location/temp' ]) self._verify_sdk_harness_container_image_overrides(pipeline_options) def test_sdk_harness_container_image_overrides_prime(self): pipeline_options = PipelineOptions([ '--dataflow_service_options=enable_prime', '--temp_location', 'gs://any-location/temp' ]) self._verify_sdk_harness_container_image_overrides(pipeline_options) def _verify_dataflow_container_image_override(self, pipeline_options): pipeline = Pipeline(options=pipeline_options) pipeline | Create([1, 2, 3]) | ParDo(DoFn()) # pylint:disable=expression-not-assigned dummy_env = DockerEnvironment( container_image='apache/beam_dummy_name:dummy_tag') proto_pipeline, _ = pipeline.to_runner_api( return_context=True, default_environment=dummy_env) # Accessing non-public method for testing. apiclient.DataflowApplicationClient._apply_sdk_environment_overrides( proto_pipeline, {}, pipeline_options) from apache_beam.utils import proto_utils found_override = False for env in proto_pipeline.components.environments.values(): docker_payload = proto_utils.parse_Bytes( env.payload, beam_runner_api_pb2.DockerPayload) if docker_payload.container_image.startswith( names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY): found_override = True self.assertTrue(found_override) def test_dataflow_container_image_override_runner_v2(self): pipeline_options = PipelineOptions([ '--experiments=use_runner_v2', '--temp_location', 'gs://any-location/temp' ]) self._verify_dataflow_container_image_override(pipeline_options) def test_dataflow_container_image_override_prime(self): pipeline_options = PipelineOptions([ '--dataflow_service_options=enable_prime', '--temp_location', 'gs://any-location/temp' ]) self._verify_dataflow_container_image_override(pipeline_options) def _verify_dataflow_container_image_override_rc(self, pipeline_options): pipeline = Pipeline(options=pipeline_options) pipeline | Create([1, 2, 3]) | ParDo(DoFn()) # pylint:disable=expression-not-assigned dummy_env = DockerEnvironment( container_image='apache/beam_dummy_name:2.00.0RC10') proto_pipeline, _ = pipeline.to_runner_api( return_context=True, default_environment=dummy_env) # Accessing non-public method for testing. apiclient.DataflowApplicationClient._apply_sdk_environment_overrides( proto_pipeline, {}, pipeline_options) from apache_beam.utils import proto_utils found_override = False trimed_rc = True for env in proto_pipeline.components.environments.values(): docker_payload = proto_utils.parse_Bytes( env.payload, beam_runner_api_pb2.DockerPayload) if docker_payload.container_image.startswith( names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY): found_override = True if docker_payload.container_image.split(':')[-1] != '2.00.0': trimed_rc = False self.assertTrue(found_override) self.assertTrue(trimed_rc) def test_dataflow_container_image_override_rc(self): pipeline_options = PipelineOptions([ '--experiments=use_runner_v2', '--temp_location', 'gs://any-location/temp' ]) self._verify_dataflow_container_image_override_rc(pipeline_options) def _verify_non_apache_container_not_overridden(self, pipeline_options): pipeline = Pipeline(options=pipeline_options) pipeline | Create([1, 2, 3]) | ParDo(DoFn()) # pylint:disable=expression-not-assigned dummy_env = DockerEnvironment( container_image='other_org/dummy_name:dummy_tag') proto_pipeline, _ = pipeline.to_runner_api( return_context=True, default_environment=dummy_env) # Accessing non-public method for testing. apiclient.DataflowApplicationClient._apply_sdk_environment_overrides( proto_pipeline, {}, pipeline_options) self.assertIsNotNone(2, len(proto_pipeline.components.environments)) from apache_beam.utils import proto_utils found_override = False for env in proto_pipeline.components.environments.values(): docker_payload = proto_utils.parse_Bytes( env.payload, beam_runner_api_pb2.DockerPayload) if docker_payload.container_image.startswith( names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY): found_override = True self.assertFalse(found_override) def test_non_apache_container_not_overridden_runner_v2(self): pipeline_options = PipelineOptions([ '--experiments=use_runner_v2', '--temp_location', 'gs://any-location/temp' ]) self._verify_non_apache_container_not_overridden(pipeline_options) def test_non_apache_container_not_overridden_prime(self): pipeline_options = PipelineOptions([ '--dataflow_service_options=enable_prime', '--temp_location', 'gs://any-location/temp' ]) self._verify_non_apache_container_not_overridden(pipeline_options) def _verify_pipeline_sdk_not_overridden(self, pipeline_options): pipeline = Pipeline(options=pipeline_options) pipeline | Create([1, 2, 3]) | ParDo(DoFn()) # pylint:disable=expression-not-assigned proto_pipeline, _ = pipeline.to_runner_api(return_context=True) dummy_env = DockerEnvironment( container_image='dummy_prefix/dummy_name:dummy_tag') proto_pipeline, _ = pipeline.to_runner_api( return_context=True, default_environment=dummy_env) # Accessing non-public method for testing. apiclient.DataflowApplicationClient._apply_sdk_environment_overrides( proto_pipeline, {}, pipeline_options) self.assertIsNotNone(2, len(proto_pipeline.components.environments)) from apache_beam.utils import proto_utils found_override = False for env in proto_pipeline.components.environments.values(): docker_payload = proto_utils.parse_Bytes( env.payload, beam_runner_api_pb2.DockerPayload) if docker_payload.container_image.startswith( names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY): found_override = True self.assertFalse(found_override) def test_pipeline_sdk_not_overridden_runner_v2(self): pipeline_options = PipelineOptions([ '--experiments=use_runner_v2', '--temp_location', 'gs://any-location/temp', '--sdk_container_image=dummy_prefix/dummy_name:dummy_tag' ]) self._verify_pipeline_sdk_not_overridden(pipeline_options) def test_pipeline_sdk_not_overridden_prime(self): pipeline_options = PipelineOptions([ '--dataflow_service_options=enable_prime', '--temp_location', 'gs://any-location/temp', '--sdk_container_image=dummy_prefix/dummy_name:dummy_tag' ]) self._verify_pipeline_sdk_not_overridden(pipeline_options) def test_invalid_default_job_name(self): # Regexp for job names in dataflow. regexp = '^[a-z]([-a-z0-9]{0,61}[a-z0-9])?$' job_name = apiclient.Job._build_default_job_name('invalid.-_user_n*/ame') self.assertRegex(job_name, regexp) job_name = apiclient.Job._build_default_job_name( 'invalid-extremely-long.username_that_shouldbeshortened_or_is_invalid') self.assertRegex(job_name, regexp) def test_default_job_name(self): job_name = apiclient.Job.default_job_name(None) regexp = 'beamapp-.*-[0-9]{10}-[0-9]{6}-[a-z0-9]{8}$' self.assertRegex(job_name, regexp) def test_split_int(self): number = 12345 split_number = apiclient.to_split_int(number) self.assertEqual((split_number.lowBits, split_number.highBits), (number, 0)) shift_number = number << 32 split_number = apiclient.to_split_int(shift_number) self.assertEqual((split_number.lowBits, split_number.highBits), (0, number)) def test_translate_distribution_using_accumulator(self): metric_update = dataflow.CounterUpdate() accumulator = mock.Mock() accumulator.min = 1 accumulator.max = 15 accumulator.sum = 16 accumulator.count = 2 apiclient.translate_distribution(accumulator, metric_update) self.assertEqual(metric_update.distribution.min.lowBits, accumulator.min) self.assertEqual(metric_update.distribution.max.lowBits, accumulator.max) self.assertEqual(metric_update.distribution.sum.lowBits, accumulator.sum) self.assertEqual( metric_update.distribution.count.lowBits, accumulator.count) def test_translate_distribution_using_distribution_data(self): metric_update = dataflow.CounterUpdate() distribution_update = DistributionData(16, 2, 1, 15) apiclient.translate_distribution(distribution_update, metric_update) self.assertEqual( metric_update.distribution.min.lowBits, distribution_update.min) self.assertEqual( metric_update.distribution.max.lowBits, distribution_update.max) self.assertEqual( metric_update.distribution.sum.lowBits, distribution_update.sum) self.assertEqual( metric_update.distribution.count.lowBits, distribution_update.count) def test_translate_distribution_using_dataflow_distribution_counter(self): counter_update = DataflowDistributionCounter() counter_update.add_input(1) counter_update.add_input(3) metric_proto = dataflow.CounterUpdate() apiclient.translate_distribution(counter_update, metric_proto) histogram = mock.Mock(firstBucketOffset=None, bucketCounts=None) counter_update.translate_to_histogram(histogram) self.assertEqual(metric_proto.distribution.min.lowBits, counter_update.min) self.assertEqual(metric_proto.distribution.max.lowBits, counter_update.max) self.assertEqual(metric_proto.distribution.sum.lowBits, counter_update.sum) self.assertEqual( metric_proto.distribution.count.lowBits, counter_update.count) self.assertEqual( metric_proto.distribution.histogram.bucketCounts, histogram.bucketCounts) self.assertEqual( metric_proto.distribution.histogram.firstBucketOffset, histogram.firstBucketOffset) def test_translate_means(self): metric_update = dataflow.CounterUpdate() accumulator = mock.Mock() accumulator.sum = 16 accumulator.count = 2 apiclient.MetricUpdateTranslators.translate_scalar_mean_int( accumulator, metric_update) self.assertEqual(metric_update.integerMean.sum.lowBits, accumulator.sum) self.assertEqual(metric_update.integerMean.count.lowBits, accumulator.count) accumulator.sum = 16.0 accumulator.count = 2 apiclient.MetricUpdateTranslators.translate_scalar_mean_float( accumulator, metric_update) self.assertEqual(metric_update.floatingPointMean.sum, accumulator.sum) self.assertEqual( metric_update.floatingPointMean.count.lowBits, accumulator.count) def test_translate_means_using_distribution_accumulator(self): # This is the special case for MeanByteCount. # Which is reported over the FnAPI as a beam distribution, # and to the service as a MetricUpdate IntegerMean. metric_update = dataflow.CounterUpdate() accumulator = mock.Mock() accumulator.min = 7 accumulator.max = 9 accumulator.sum = 16 accumulator.count = 2 apiclient.MetricUpdateTranslators.translate_scalar_mean_int( accumulator, metric_update) self.assertEqual(metric_update.integerMean.sum.lowBits, accumulator.sum) self.assertEqual(metric_update.integerMean.count.lowBits, accumulator.count) accumulator.sum = 16.0 accumulator.count = 2 apiclient.MetricUpdateTranslators.translate_scalar_mean_float( accumulator, metric_update) self.assertEqual(metric_update.floatingPointMean.sum, accumulator.sum) self.assertEqual( metric_update.floatingPointMean.count.lowBits, accumulator.count) def test_default_ip_configuration(self): pipeline_options = PipelineOptions( ['--temp_location', 'gs://any-location/temp']) env = apiclient.Environment([], pipeline_options, '2.0.0', FAKE_PIPELINE_URL) self.assertEqual(env.proto.workerPools[0].ipConfiguration, None) def test_public_ip_configuration(self): pipeline_options = PipelineOptions( ['--temp_location', 'gs://any-location/temp', '--use_public_ips']) env = apiclient.Environment([], pipeline_options, '2.0.0', FAKE_PIPELINE_URL) self.assertEqual( env.proto.workerPools[0].ipConfiguration, dataflow.WorkerPool.IpConfigurationValueValuesEnum.WORKER_IP_PUBLIC) def test_private_ip_configuration(self): pipeline_options = PipelineOptions( ['--temp_location', 'gs://any-location/temp', '--no_use_public_ips']) env = apiclient.Environment([], pipeline_options, '2.0.0', FAKE_PIPELINE_URL) self.assertEqual( env.proto.workerPools[0].ipConfiguration, dataflow.WorkerPool.IpConfigurationValueValuesEnum.WORKER_IP_PRIVATE) def test_number_of_worker_harness_threads(self): pipeline_options = PipelineOptions([ '--temp_location', 'gs://any-location/temp', '--number_of_worker_harness_threads', '2' ]) env = apiclient.Environment([], pipeline_options, '2.0.0', FAKE_PIPELINE_URL) self.assertEqual(env.proto.workerPools[0].numThreadsPerWorker, 2) def test_disk_provisioning_options(self): pipeline_options = PipelineOptions([ '--temp_location', 'gs://any-location/temp', '--disk_provisioned_iops', '4000', '--disk_provisioned_throughput_mibps', '200' ]) env = apiclient.Environment([], pipeline_options, '2.0.0', FAKE_PIPELINE_URL) self.assertEqual(env.proto.workerPools[0].diskProvisionedIops, 4000) self.assertEqual( env.proto.workerPools[0].diskProvisionedThroughputMibps, 200) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.' 'beam_version.__version__', '2.2.0') def test_harness_override_absent_with_runner_v2(self): pipeline_options = PipelineOptions([ '--temp_location', 'gs://any-location/temp', '--streaming', '--experiments=use_runner_v2' ]) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) if env.proto.experiments: for experiment in env.proto.experiments: self.assertNotIn('runner_harness_container_image=', experiment) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.' 'beam_version.__version__', '2.2.0') def test_custom_harness_override_present_with_runner_v2(self): pipeline_options = PipelineOptions([ '--temp_location', 'gs://any-location/temp', '--streaming', '--experiments=runner_harness_container_image=fake_image', '--experiments=use_runner_v2', ]) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual( 1, len([ x for x in env.proto.experiments if x.startswith('runner_harness_container_image=') ])) self.assertIn( 'runner_harness_container_image=fake_image', env.proto.experiments) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.' 'beam_version.__version__', '2.2.0.dev') def test_pinned_worker_harness_image_tag_used_in_dev_sdk(self): # streaming, fnapi pipeline. pipeline_options = PipelineOptions( ['--temp_location', 'gs://any-location/temp', '--streaming']) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual( env.proto.workerPools[0].workerHarnessContainerImage, ( names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + '/beam_python%d.%d_sdk:%s' % ( sys.version_info[0], sys.version_info[1], names.BEAM_DEV_SDK_CONTAINER_TAG))) pipeline_options = PipelineOptions( ['--temp_location', 'gs://any-location/temp']) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual( env.proto.workerPools[0].workerHarnessContainerImage, ( names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + '/beam_python%d.%d_sdk:%s' % ( sys.version_info[0], sys.version_info[1], names.BEAM_DEV_SDK_CONTAINER_TAG))) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.' 'beam_version.__version__', '2.2.0') def test_worker_harness_image_tag_matches_released_sdk_version(self): # streaming, fnapi pipeline. pipeline_options = PipelineOptions( ['--temp_location', 'gs://any-location/temp', '--streaming']) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual( env.proto.workerPools[0].workerHarnessContainerImage, ( names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + '/beam_python%d.%d_sdk:2.2.0' % (sys.version_info[0], sys.version_info[1]))) pipeline_options = PipelineOptions( ['--temp_location', 'gs://any-location/temp']) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual( env.proto.workerPools[0].workerHarnessContainerImage, ( names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + '/beam_python%d.%d_sdk:2.2.0' % (sys.version_info[0], sys.version_info[1]))) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.' 'beam_version.__version__', '2.2.0.rc1') def test_worker_harness_image_tag_matches_base_sdk_version_of_an_rc(self): # streaming, fnapi pipeline. pipeline_options = PipelineOptions( ['--temp_location', 'gs://any-location/temp', '--streaming']) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual( env.proto.workerPools[0].workerHarnessContainerImage, ( names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + '/beam_python%d.%d_sdk:2.2.0' % (sys.version_info[0], sys.version_info[1]))) pipeline_options = PipelineOptions( ['--temp_location', 'gs://any-location/temp']) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual( env.proto.workerPools[0].workerHarnessContainerImage, ( names.DATAFLOW_CONTAINER_IMAGE_REPOSITORY + '/beam_python%d.%d_sdk:2.2.0' % (sys.version_info[0], sys.version_info[1]))) def test_worker_harness_override_takes_precedence_over_sdk_defaults(self): # streaming, fnapi pipeline. pipeline_options = PipelineOptions([ '--temp_location', 'gs://any-location/temp', '--streaming', '--sdk_container_image=some:image' ]) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual( env.proto.workerPools[0].workerHarnessContainerImage, 'some:image') # batch, legacy pipeline. pipeline_options = PipelineOptions([ '--temp_location', 'gs://any-location/temp', '--sdk_container_image=some:image' ]) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual( env.proto.workerPools[0].workerHarnessContainerImage, 'some:image') @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.Job.' 'job_id_for_name', return_value='test_id') def test_transform_name_mapping(self, mock_job): pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp', '--update', '--transform_name_mapping', '{\"from\":\"to\"}' ]) job = apiclient.Job(pipeline_options, beam_runner_api_pb2.Pipeline()) self.assertIsNotNone(job.proto.transformNameMapping) def test_created_from_snapshot_id(self): pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp', '--create_from_snapshot', 'test_snapshot_id' ]) job = apiclient.Job(pipeline_options, beam_runner_api_pb2.Pipeline()) self.assertEqual('test_snapshot_id', job.proto.createdFromSnapshotId) def test_labels(self): pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp' ]) job = apiclient.Job(pipeline_options, beam_runner_api_pb2.Pipeline()) self.assertIsNone(job.proto.labels) pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp', '--label', 'key1=value1', '--label', 'key2', '--label', 'key3=value3', '--labels', 'key4=value4', '--labels', 'key5' ]) job = apiclient.Job(pipeline_options, beam_runner_api_pb2.Pipeline()) self.assertEqual(5, len(job.proto.labels.additionalProperties)) self.assertEqual('key1', job.proto.labels.additionalProperties[0].key) self.assertEqual('value1', job.proto.labels.additionalProperties[0].value) self.assertEqual('key2', job.proto.labels.additionalProperties[1].key) self.assertEqual('', job.proto.labels.additionalProperties[1].value) self.assertEqual('key3', job.proto.labels.additionalProperties[2].key) self.assertEqual('value3', job.proto.labels.additionalProperties[2].value) self.assertEqual('key4', job.proto.labels.additionalProperties[3].key) self.assertEqual('value4', job.proto.labels.additionalProperties[3].value) self.assertEqual('key5', job.proto.labels.additionalProperties[4].key) self.assertEqual('', job.proto.labels.additionalProperties[4].value) pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp', '--labels', '{ "name": "wrench", "mass": "1_3kg", "count": "3" }' ]) job = apiclient.Job(pipeline_options, beam_runner_api_pb2.Pipeline()) self.assertEqual(3, len(job.proto.labels.additionalProperties)) self.assertEqual('name', job.proto.labels.additionalProperties[0].key) self.assertEqual('wrench', job.proto.labels.additionalProperties[0].value) self.assertEqual('mass', job.proto.labels.additionalProperties[1].key) self.assertEqual('1_3kg', job.proto.labels.additionalProperties[1].value) self.assertEqual('count', job.proto.labels.additionalProperties[2].key) self.assertEqual('3', job.proto.labels.additionalProperties[2].value) pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp' ]) pipeline_options.view_as(GoogleCloudOptions).labels = { "name": "wrench", "mass": "1_3kg", "count": "3" } job = apiclient.Job(pipeline_options, beam_runner_api_pb2.Pipeline()) self.assertEqual(3, len(job.proto.labels.additionalProperties)) self.assertEqual('name', job.proto.labels.additionalProperties[0].key) self.assertEqual('wrench', job.proto.labels.additionalProperties[0].value) self.assertEqual('mass', job.proto.labels.additionalProperties[1].key) self.assertEqual('1_3kg', job.proto.labels.additionalProperties[1].value) self.assertEqual('count', job.proto.labels.additionalProperties[2].key) self.assertEqual('3', job.proto.labels.additionalProperties[2].value) pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp' ]) pipeline_options.view_as( GoogleCloudOptions ).labels = '{ "name": "wrench", "mass": "1_3kg", "count": "3" }' job = apiclient.Job(pipeline_options, beam_runner_api_pb2.Pipeline()) self.assertEqual(3, len(job.proto.labels.additionalProperties)) self.assertEqual('name', job.proto.labels.additionalProperties[0].key) self.assertEqual('wrench', job.proto.labels.additionalProperties[0].value) self.assertEqual('mass', job.proto.labels.additionalProperties[1].key) self.assertEqual('1_3kg', job.proto.labels.additionalProperties[1].value) self.assertEqual('count', job.proto.labels.additionalProperties[2].key) self.assertEqual('3', job.proto.labels.additionalProperties[2].value) def test_experiment_use_multiple_sdk_containers(self): pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp', '--experiments', 'beam_fn_api' ]) environment = apiclient.Environment([], pipeline_options, 1, FAKE_PIPELINE_URL) self.assertIn('use_multiple_sdk_containers', environment.proto.experiments) pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp', '--experiments', 'beam_fn_api', '--experiments', 'use_multiple_sdk_containers' ]) environment = apiclient.Environment([], pipeline_options, 1, FAKE_PIPELINE_URL) self.assertIn('use_multiple_sdk_containers', environment.proto.experiments) pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp', '--experiments', 'beam_fn_api', '--experiments', 'no_use_multiple_sdk_containers' ]) environment = apiclient.Environment([], pipeline_options, 1, FAKE_PIPELINE_URL) self.assertNotIn( 'use_multiple_sdk_containers', environment.proto.experiments) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.sys.version_info', (3, 9)) def test_get_python_sdk_name(self): pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp', '--experiments', 'beam_fn_api', '--experiments', 'use_multiple_sdk_containers' ]) environment = apiclient.Environment([], pipeline_options, 1, FAKE_PIPELINE_URL) self.assertEqual( 'Apache Beam Python 3.9 SDK', environment._get_python_sdk_name()) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.sys.version_info', (2, 7)) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.' 'beam_version.__version__', '2.2.0') def test_interpreter_version_check_fails_py27(self): pipeline_options = PipelineOptions([]) self.assertRaises( Exception, apiclient._verify_interpreter_version_is_supported, pipeline_options) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.sys.version_info', (3, 0, 0)) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.' 'beam_version.__version__', '2.2.0.dev') def test_interpreter_version_check_passes_on_dev_sdks(self): pipeline_options = PipelineOptions([]) apiclient._verify_interpreter_version_is_supported(pipeline_options) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.' 'beam_version.__version__', '2.2.0') @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.sys.version_info', (3, 0, 0)) def test_interpreter_version_check_passes_with_experiment(self): pipeline_options = PipelineOptions( ["--experiment=use_unsupported_python_version"]) apiclient._verify_interpreter_version_is_supported(pipeline_options) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.sys.version_info', (3, 8, 2)) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.' 'beam_version.__version__', '2.2.0') def test_interpreter_version_check_fails_py38(self): pipeline_options = PipelineOptions([]) self.assertRaises( Exception, apiclient._verify_interpreter_version_is_supported, pipeline_options) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.sys.version_info', (3, 10, 10)) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.' 'beam_version.__version__', '2.2.0') def test_interpreter_version_check_passes_py310(self): pipeline_options = PipelineOptions([]) apiclient._verify_interpreter_version_is_supported(pipeline_options) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.sys.version_info', (3, 15, 0)) @mock.patch( 'apache_beam.runners.dataflow.internal.apiclient.' 'beam_version.__version__', '2.2.0') def test_interpreter_version_check_fails_on_not_yet_supported_version(self): pipeline_options = PipelineOptions([]) self.assertRaises( Exception, apiclient._verify_interpreter_version_is_supported, pipeline_options) def test_get_response_encoding(self): encoding = apiclient.get_response_encoding() assert encoding == 'utf8' def test_graph_is_uploaded(self): pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp', '--experiments', 'beam_fn_api', '--experiments', 'upload_graph' ]) job = apiclient.Job(pipeline_options, beam_runner_api_pb2.Pipeline()) pipeline_options.view_as(GoogleCloudOptions).no_auth = True client = apiclient.DataflowApplicationClient(pipeline_options) with mock.patch.object(client, 'stage_file', side_effect=None): with mock.patch.object(client, 'create_job_description', side_effect=None): with mock.patch.object(client, 'submit_job_description', side_effect=None): client.create_job(job) client.stage_file.assert_called_once_with( mock.ANY, "dataflow_graph.json", mock.ANY, 'application/octet-stream', None) client.create_job_description.assert_called_once() def test_create_job_returns_existing_job(self): pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp', ]) job = apiclient.Job(pipeline_options, beam_runner_api_pb2.Pipeline()) self.assertTrue(job.proto.clientRequestId) # asserts non-empty string pipeline_options.view_as(GoogleCloudOptions).no_auth = True client = apiclient.DataflowApplicationClient(pipeline_options) response = dataflow.Job() # different clientRequestId from `job` response.clientRequestId = "20210821081910123456-1234" response.name = 'test_job_name' response.id = '2021-08-19_21_18_43-9756917246311111021' with mock.patch.object(client._client.projects_locations_jobs, 'Create', side_effect=[response]): with mock.patch.object(client, 'create_job_description', side_effect=None): with self.assertRaises( apiclient.DataflowJobAlreadyExistsError) as context: client.create_job(job) self.assertEqual( str(context.exception), 'There is already active job named %s with id: %s. If you want to ' 'submit a second job, try again by setting a different name using ' '--job_name.' % ('test_job_name', response.id)) def test_update_job_returns_existing_job(self): pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp', '--region', 'us-central1', '--update', ]) replace_job_id = '2021-08-21_00_00_01-6081497447916622336' with mock.patch('apache_beam.runners.dataflow.internal.apiclient.Job.' 'job_id_for_name', return_value=replace_job_id) as job_id_for_name_mock: job = apiclient.Job(pipeline_options, beam_runner_api_pb2.Pipeline()) job_id_for_name_mock.assert_called_once() self.assertTrue(job.proto.clientRequestId) # asserts non-empty string pipeline_options.view_as(GoogleCloudOptions).no_auth = True client = apiclient.DataflowApplicationClient(pipeline_options) response = dataflow.Job() # different clientRequestId from `job` response.clientRequestId = "20210821083254123456-1234" response.name = 'test_job_name' response.id = '2021-08-19_21_29_07-5725551945600207770' with mock.patch.object(client, 'create_job_description', side_effect=None): with mock.patch.object(client._client.projects_locations_jobs, 'Create', side_effect=[response]): with self.assertRaises( apiclient.DataflowJobAlreadyExistsError) as context: client.create_job(job) self.assertEqual( str(context.exception), 'The job named %s with id: %s has already been updated into job ' 'id: %s and cannot be updated again.' % ('test_job_name', replace_job_id, response.id)) def test_template_file_generation_with_upload_graph(self): pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp', '--experiments', 'upload_graph', '--template_location', 'gs://test-location/template' ]) job = apiclient.Job(pipeline_options, beam_runner_api_pb2.Pipeline()) job.proto.steps.append(dataflow.Step(name='test_step_name')) pipeline_options.view_as(GoogleCloudOptions).no_auth = True client = apiclient.DataflowApplicationClient(pipeline_options) with mock.patch.object(client, 'stage_file', side_effect=None): with mock.patch.object(client, 'create_job_description', side_effect=None): with mock.patch.object(client, 'submit_job_description', side_effect=None): client.create_job(job) client.stage_file.assert_has_calls([ mock.call( mock.ANY, 'dataflow_graph.json', mock.ANY, 'application/octet-stream', None), mock.call( mock.ANY, 'template', mock.ANY, 'application/octet-stream', None) ]) client.create_job_description.assert_called_once() # template is generated, but job should not be submitted to the # service. client.submit_job_description.assert_not_called() template_filename = client.stage_file.call_args_list[-1][0][1] self.assertTrue('template' in template_filename) template_content = client.stage_file.call_args_list[-1][0][2].read( ).decode('utf-8') template_obj = json.loads(template_content) self.assertFalse(template_obj.get('steps')) self.assertTrue(template_obj['stepsLocation']) def test_stage_resources(self): pipeline_options = PipelineOptions([ '--temp_location', 'gs://test-location/temp', '--staging_location', 'gs://test-location/staging', '--no_auth' ]) pipeline = beam_runner_api_pb2.Pipeline( components=beam_runner_api_pb2.Components( environments={ 'env1': beam_runner_api_pb2.Environment( dependencies=[ beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload( path='/tmp/foo1').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='foo1').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload( path='/tmp/bar1').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='bar1').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload( path='/tmp/baz').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='baz1').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload( path='/tmp/renamed1', sha256='abcdefg').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='renamed1').SerializeToString()) ]), 'env2': beam_runner_api_pb2.Environment( dependencies=[ beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload( path='/tmp/foo2').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='foo2').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload( path='/tmp/bar2').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='bar2').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload( path='/tmp/baz').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='baz2').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload( path='/tmp/renamed2', sha256='abcdefg').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='renamed2').SerializeToString()) ]) })) client = apiclient.DataflowApplicationClient(pipeline_options) with mock.patch.object(apiclient.DataflowApplicationClient, '_compute_sha256', side_effect=lambda path: 'hash' + path): with mock.patch.object(apiclient._LegacyDataflowStager, 'stage_job_resources') as mock_stager: client._stage_resources(pipeline, pipeline_options) mock_stager.assert_called_once_with( [('/tmp/foo1', 'foo1', 'hash/tmp/foo1'), ('/tmp/bar1', 'bar1', 'hash/tmp/bar1'), ('/tmp/baz', 'baz1', 'hash/tmp/baz'), ('/tmp/renamed1', 'renamed1', 'abcdefg'), ('/tmp/foo2', 'foo2', 'hash/tmp/foo2'), ('/tmp/bar2', 'bar2', 'hash/tmp/bar2')], staging_location='gs://test-location/staging') pipeline_expected = beam_runner_api_pb2.Pipeline( components=beam_runner_api_pb2.Components( environments={ 'env1': beam_runner_api_pb2.Environment( dependencies=[ beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/foo1', sha256='hash/tmp/foo1').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='foo1').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/bar1', sha256='hash/tmp/bar1').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='bar1').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/baz1', sha256='hash/tmp/baz').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='baz1').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/renamed1', sha256='abcdefg').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='renamed1').SerializeToString()) ]), 'env2': beam_runner_api_pb2.Environment( dependencies=[ beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/foo2', sha256='hash/tmp/foo2').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='foo2').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/bar2', sha256='hash/tmp/bar2').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='bar2').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/baz1', sha256='hash/tmp/baz').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='baz1').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/renamed1', sha256='abcdefg').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='renamed1').SerializeToString()) ]) })) self.assertEqual(pipeline, pipeline_expected) def test_set_dataflow_service_option(self): pipeline_options = PipelineOptions([ '--dataflow_service_option', 'whizz=bang', '--temp_location', 'gs://any-location/temp' ]) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual(env.proto.serviceOptions, ['whizz=bang']) def test_enable_hot_key_logging(self): # Tests that the enable_hot_key_logging is not set by default. pipeline_options = PipelineOptions( ['--temp_location', 'gs://any-location/temp']) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertIsNone(env.proto.debugOptions) # Now test that it is set when given. pipeline_options = PipelineOptions([ '--enable_hot_key_logging', '--temp_location', 'gs://any-location/temp' ]) env = apiclient.Environment( [], #packages pipeline_options, '2.0.0', #any environment version FAKE_PIPELINE_URL) self.assertEqual( env.proto.debugOptions, dataflow.DebugOptions(enableHotKeyLogging=True)) def _mock_uncached_copy(self, staging_root, src, sha256, dst_name=None): sha_prefix = sha256[0:2] gcs_cache_path = FileSystems.join( staging_root, apiclient.DataflowApplicationClient._GCS_CACHE_PREFIX, sha_prefix, sha256) if not dst_name: _, dst_name = os.path.split(src) return [ mock.call.gcs_exists(gcs_cache_path), mock.call.gcs_upload(src, gcs_cache_path), mock.call.gcs_gcs_copy( source_file_names=[gcs_cache_path], destination_file_names=[f'gs://test-location/staging/{dst_name}']) ] def _mock_cached_copy(self, staging_root, src, sha256, dst_name=None): uncached = self._mock_uncached_copy(staging_root, src, sha256, dst_name) uncached.pop(1) return uncached def test_stage_artifacts_with_caching(self): pipeline_options = PipelineOptions([ '--temp_location', 'gs://test-location/temp', '--staging_location', 'gs://test-location/staging', '--no_auth', '--enable_artifact_caching' ]) pipeline = beam_runner_api_pb2.Pipeline( components=beam_runner_api_pb2.Components( environments={ 'env1': beam_runner_api_pb2.Environment( dependencies=[ beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload( path='/tmp/foo1', sha256='abcd').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='foo1').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload( path='/tmp/bar1', sha256='defg').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='bar1').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload(path='/tmp/baz', sha256='hijk' ).SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='baz1').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload( path='/tmp/renamed1', sha256='abcdefg').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='renamed1').SerializeToString()) ]), 'env2': beam_runner_api_pb2.Environment( dependencies=[ beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload( path='/tmp/foo2', sha256='lmno').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='foo2').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload( path='/tmp/bar2', sha256='pqrs').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='bar2').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload(path='/tmp/baz', sha256='tuv' ).SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='baz2').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2. ArtifactFilePayload( path='/tmp/renamed2', sha256='abcdefg').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='renamed2').SerializeToString()) ]) })) client = apiclient.DataflowApplicationClient(pipeline_options) staging_root = 'gs://test-location/staging' # every other artifact already exists n = [0] def exists_return_value(*args): n[0] += 1 return n[0] % 2 == 0 with mock.patch.object(FileSystems, 'exists', side_effect=exists_return_value) as mock_gcs_exists: with mock.patch.object(apiclient.DataflowApplicationClient, '_uncached_gcs_file_copy') as mock_gcs_copy: with mock.patch.object(FileSystems, 'copy') as mock_gcs_gcs_copy: manager = mock.Mock() manager.attach_mock(mock_gcs_exists, 'gcs_exists') manager.attach_mock(mock_gcs_copy, 'gcs_upload') manager.attach_mock(mock_gcs_gcs_copy, 'gcs_gcs_copy') client._stage_resources(pipeline, pipeline_options) expected_calls = list( itertools.chain.from_iterable([ self._mock_uncached_copy(staging_root, '/tmp/foo1', 'abcd'), self._mock_cached_copy(staging_root, '/tmp/bar1', 'defg'), self._mock_uncached_copy( staging_root, '/tmp/baz', 'hijk', 'baz1'), self._mock_cached_copy( staging_root, '/tmp/renamed1', 'abcdefg'), self._mock_uncached_copy(staging_root, '/tmp/foo2', 'lmno'), self._mock_cached_copy(staging_root, '/tmp/bar2', 'pqrs'), ])) assert manager.mock_calls == expected_calls pipeline_expected = beam_runner_api_pb2.Pipeline( components=beam_runner_api_pb2.Components( environments={ 'env1': beam_runner_api_pb2.Environment( dependencies=[ beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/foo1', sha256='abcd').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='foo1').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/bar1', sha256='defg').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='bar1').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/baz1', sha256='hijk').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='baz1').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/renamed1', sha256='abcdefg').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='renamed1').SerializeToString()) ]), 'env2': beam_runner_api_pb2.Environment( dependencies=[ beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/foo2', sha256='lmno').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='foo2').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/bar2', sha256='pqrs').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='bar2').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/baz1', sha256='tuv').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='baz1').SerializeToString()), beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='gs://test-location/staging/renamed1', sha256='abcdefg').SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2. ArtifactStagingToRolePayload( staged_name='renamed1').SerializeToString()) ]) })) self.assertEqual(pipeline, pipeline_expected) def test_stage_file_with_retry(self): def effect(self, *args, **kwargs): nonlocal count count += 1 # Fail the first two calls and succeed afterward if count <= 2: raise Exception("This exception is raised for testing purpose.") class Unseekable(io.IOBase): def seekable(self): return False pipeline_options = PipelineOptions([ '--project', 'test_project', '--job_name', 'test_job_name', '--temp_location', 'gs://test-location/temp', ]) pipeline_options.view_as(GoogleCloudOptions).no_auth = True client = apiclient.DataflowApplicationClient(pipeline_options) with mock.patch.object(client, 'stage_file') as mock_stage_file: mock_stage_file.side_effect = effect with mock.patch.object(time, 'sleep') as mock_sleep: with mock.patch("builtins.open", mock.mock_open(read_data="data")) as mock_file_open: count = 0 # calling with a file name client.stage_file_with_retry( "/to", "new_name", "/from/old_name", total_size=4) self.assertEqual(mock_stage_file.call_count, 3) self.assertEqual(mock_sleep.call_count, 2) self.assertEqual(mock_file_open.call_count, 3) count = 0 mock_stage_file.reset_mock() mock_sleep.reset_mock() mock_file_open.reset_mock() # calling with a seekable stream client.stage_file_with_retry( "/to", "new_name", io.BytesIO(b'test'), total_size=4) self.assertEqual(mock_stage_file.call_count, 3) self.assertEqual(mock_sleep.call_count, 2) # no open() is called if a stream is provided mock_file_open.assert_not_called() count = 0 mock_sleep.reset_mock() mock_file_open.reset_mock() mock_stage_file.reset_mock() # calling with an unseekable stream self.assertRaises( retry.PermanentException, client.stage_file_with_retry, "/to", "new_name", Unseekable(), total_size=4) # Unseekable streams are staged once. If staging fails, no retries are # attempted. self.assertEqual(mock_stage_file.call_count, 1) mock_sleep.assert_not_called() mock_file_open.assert_not_called() count = 0 mock_sleep.reset_mock() mock_file_open.reset_mock() mock_stage_file.reset_mock() # calling with something else self.assertRaises( retry.PermanentException, client.stage_file_with_retry, "/to", "new_name", object(), total_size=4) # No staging will be called for wrong arg type mock_stage_file.assert_not_called() mock_sleep.assert_not_called() mock_file_open.assert_not_called() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5854294 apache_beam-2.74.0/apache_beam/runners/dataflow/internal/clients/0000755000175100017510000000000015206113160024431 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/internal/clients/__init__.py0000644000175100017510000000142015206113032026535 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.5874298 apache_beam-2.74.0/apache_beam/runners/dataflow/internal/clients/dataflow/0000755000175100017510000000000015206113160026232 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/internal/clients/dataflow/__init__.py0000644000175100017510000000255615206113032030351 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Common imports for generated dataflow client library.""" # pylint:disable=wildcard-import import pkgutil # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apitools.base.py import * from apache_beam.runners.dataflow.internal.clients.dataflow.dataflow_v1b3_client import * from apache_beam.runners.dataflow.internal.clients.dataflow.dataflow_v1b3_messages import * except ImportError: pass # pylint: enable=wrong-import-order, wrong-import-position __path__ = pkgutil.extend_path(__path__, __name__) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/internal/clients/dataflow/dataflow_v1b3_client.py0000644000175100017510000016343415206113032032607 0ustar00runnerrunner"""Generated client library for dataflow version v1b3.""" # NOTE: This file is autogenerated and should not be edited by hand. from apitools.base.py import base_api from . import dataflow_v1b3_messages as messages class DataflowV1b3(base_api.BaseApiClient): """Generated client library for service dataflow version v1b3.""" MESSAGES_MODULE = messages BASE_URL = 'https://dataflow.googleapis.com/' MTLS_BASE_URL = 'https://dataflow.mtls.googleapis.com/' _PACKAGE = 'dataflow' _SCOPES = ['https://www.googleapis.com/auth/cloud-platform', 'https://www.googleapis.com/auth/compute'] _VERSION = 'v1b3' _CLIENT_ID = 'CLIENT_ID' _CLIENT_SECRET = 'CLIENT_SECRET' _USER_AGENT = 'x_Tw5K8nnjoRAqULM9PFAC2b' _CLIENT_CLASS_NAME = 'DataflowV1b3' _URL_VERSION = 'v1b3' _API_KEY = None def __init__(self, url='', credentials=None, get_credentials=True, http=None, model=None, log_request=False, log_response=False, credentials_args=None, default_global_params=None, additional_http_headers=None, response_encoding=None): """Create a new dataflow handle.""" url = url or self.BASE_URL super(DataflowV1b3, self).__init__( url, credentials=credentials, get_credentials=get_credentials, http=http, model=model, log_request=log_request, log_response=log_response, credentials_args=credentials_args, default_global_params=default_global_params, additional_http_headers=additional_http_headers, response_encoding=response_encoding) self.projects_jobs_debug = self.ProjectsJobsDebugService(self) self.projects_jobs_messages = self.ProjectsJobsMessagesService(self) self.projects_jobs_workItems = self.ProjectsJobsWorkItemsService(self) self.projects_jobs = self.ProjectsJobsService(self) self.projects_locations_flexTemplates = self.ProjectsLocationsFlexTemplatesService(self) self.projects_locations_jobs_debug = self.ProjectsLocationsJobsDebugService(self) self.projects_locations_jobs_messages = self.ProjectsLocationsJobsMessagesService(self) self.projects_locations_jobs_snapshots = self.ProjectsLocationsJobsSnapshotsService(self) self.projects_locations_jobs_stages = self.ProjectsLocationsJobsStagesService(self) self.projects_locations_jobs_workItems = self.ProjectsLocationsJobsWorkItemsService(self) self.projects_locations_jobs = self.ProjectsLocationsJobsService(self) self.projects_locations_snapshots = self.ProjectsLocationsSnapshotsService(self) self.projects_locations_templates = self.ProjectsLocationsTemplatesService(self) self.projects_locations = self.ProjectsLocationsService(self) self.projects_snapshots = self.ProjectsSnapshotsService(self) self.projects_templates = self.ProjectsTemplatesService(self) self.projects = self.ProjectsService(self) class ProjectsJobsDebugService(base_api.BaseApiService): """Service class for the projects_jobs_debug resource.""" _NAME = 'projects_jobs_debug' def __init__(self, client): super(DataflowV1b3.ProjectsJobsDebugService, self).__init__(client) self._upload_configs = { } def GetConfig(self, request, global_params=None): r"""Get encoded debug configuration for component. Not cacheable. Args: request: (DataflowProjectsJobsDebugGetConfigRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (GetDebugConfigResponse) The response message. """ config = self.GetMethodConfig('GetConfig') return self._RunMethod( config, request, global_params=global_params) GetConfig.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.jobs.debug.getConfig', ordered_params=['projectId', 'jobId'], path_params=['jobId', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/jobs/{jobId}/debug/getConfig', request_field='getDebugConfigRequest', request_type_name='DataflowProjectsJobsDebugGetConfigRequest', response_type_name='GetDebugConfigResponse', supports_download=False, ) def SendCapture(self, request, global_params=None): r"""Send encoded debug capture data for component. Args: request: (DataflowProjectsJobsDebugSendCaptureRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (SendDebugCaptureResponse) The response message. """ config = self.GetMethodConfig('SendCapture') return self._RunMethod( config, request, global_params=global_params) SendCapture.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.jobs.debug.sendCapture', ordered_params=['projectId', 'jobId'], path_params=['jobId', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/jobs/{jobId}/debug/sendCapture', request_field='sendDebugCaptureRequest', request_type_name='DataflowProjectsJobsDebugSendCaptureRequest', response_type_name='SendDebugCaptureResponse', supports_download=False, ) class ProjectsJobsMessagesService(base_api.BaseApiService): """Service class for the projects_jobs_messages resource.""" _NAME = 'projects_jobs_messages' def __init__(self, client): super(DataflowV1b3.ProjectsJobsMessagesService, self).__init__(client) self._upload_configs = { } def List(self, request, global_params=None): r"""Request the job status. To request the status of a job, we recommend using `projects.locations.jobs.messages.list` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.jobs.messages.list` is not recommended, as you can only request the status of jobs that are running in `us-central1`. Args: request: (DataflowProjectsJobsMessagesListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ListJobMessagesResponse) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.jobs.messages.list', ordered_params=['projectId', 'jobId'], path_params=['jobId', 'projectId'], query_params=['endTime', 'location', 'minimumImportance', 'pageSize', 'pageToken', 'startTime'], relative_path='v1b3/projects/{projectId}/jobs/{jobId}/messages', request_field='', request_type_name='DataflowProjectsJobsMessagesListRequest', response_type_name='ListJobMessagesResponse', supports_download=False, ) class ProjectsJobsWorkItemsService(base_api.BaseApiService): """Service class for the projects_jobs_workItems resource.""" _NAME = 'projects_jobs_workItems' def __init__(self, client): super(DataflowV1b3.ProjectsJobsWorkItemsService, self).__init__(client) self._upload_configs = { } def Lease(self, request, global_params=None): r"""Leases a dataflow WorkItem to run. Args: request: (DataflowProjectsJobsWorkItemsLeaseRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (LeaseWorkItemResponse) The response message. """ config = self.GetMethodConfig('Lease') return self._RunMethod( config, request, global_params=global_params) Lease.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.jobs.workItems.lease', ordered_params=['projectId', 'jobId'], path_params=['jobId', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/jobs/{jobId}/workItems:lease', request_field='leaseWorkItemRequest', request_type_name='DataflowProjectsJobsWorkItemsLeaseRequest', response_type_name='LeaseWorkItemResponse', supports_download=False, ) def ReportStatus(self, request, global_params=None): r"""Reports the status of dataflow WorkItems leased by a worker. Args: request: (DataflowProjectsJobsWorkItemsReportStatusRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ReportWorkItemStatusResponse) The response message. """ config = self.GetMethodConfig('ReportStatus') return self._RunMethod( config, request, global_params=global_params) ReportStatus.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.jobs.workItems.reportStatus', ordered_params=['projectId', 'jobId'], path_params=['jobId', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/jobs/{jobId}/workItems:reportStatus', request_field='reportWorkItemStatusRequest', request_type_name='DataflowProjectsJobsWorkItemsReportStatusRequest', response_type_name='ReportWorkItemStatusResponse', supports_download=False, ) class ProjectsJobsService(base_api.BaseApiService): """Service class for the projects_jobs resource.""" _NAME = 'projects_jobs' def __init__(self, client): super(DataflowV1b3.ProjectsJobsService, self).__init__(client) self._upload_configs = { } def Aggregated(self, request, global_params=None): r"""List the jobs of a project across all regions. **Note:** This method doesn't support filtering the list of jobs by name. Args: request: (DataflowProjectsJobsAggregatedRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ListJobsResponse) The response message. """ config = self.GetMethodConfig('Aggregated') return self._RunMethod( config, request, global_params=global_params) Aggregated.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.jobs.aggregated', ordered_params=['projectId'], path_params=['projectId'], query_params=['filter', 'location', 'name', 'pageSize', 'pageToken', 'view'], relative_path='v1b3/projects/{projectId}/jobs:aggregated', request_field='', request_type_name='DataflowProjectsJobsAggregatedRequest', response_type_name='ListJobsResponse', supports_download=False, ) def Create(self, request, global_params=None): r"""Creates a Dataflow job. To create a job, we recommend using `projects.locations.jobs.create` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.jobs.create` is not recommended, as your job will always start in `us-central1`. Do not enter confidential information when you supply string values using the API. Args: request: (DataflowProjectsJobsCreateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Job) The response message. """ config = self.GetMethodConfig('Create') return self._RunMethod( config, request, global_params=global_params) Create.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.jobs.create', ordered_params=['projectId'], path_params=['projectId'], query_params=['location', 'replaceJobId', 'view'], relative_path='v1b3/projects/{projectId}/jobs', request_field='job', request_type_name='DataflowProjectsJobsCreateRequest', response_type_name='Job', supports_download=False, ) def Get(self, request, global_params=None): r"""Gets the state of the specified Cloud Dataflow job. To get the state of a job, we recommend using `projects.locations.jobs.get` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.jobs.get` is not recommended, as you can only get the state of jobs that are running in `us-central1`. Args: request: (DataflowProjectsJobsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Job) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params) Get.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.jobs.get', ordered_params=['projectId', 'jobId'], path_params=['jobId', 'projectId'], query_params=['location', 'view'], relative_path='v1b3/projects/{projectId}/jobs/{jobId}', request_field='', request_type_name='DataflowProjectsJobsGetRequest', response_type_name='Job', supports_download=False, ) def GetMetrics(self, request, global_params=None): r"""Request the job status. To request the status of a job, we recommend using `projects.locations.jobs.getMetrics` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.jobs.getMetrics` is not recommended, as you can only request the status of jobs that are running in `us-central1`. Args: request: (DataflowProjectsJobsGetMetricsRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (JobMetrics) The response message. """ config = self.GetMethodConfig('GetMetrics') return self._RunMethod( config, request, global_params=global_params) GetMetrics.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.jobs.getMetrics', ordered_params=['projectId', 'jobId'], path_params=['jobId', 'projectId'], query_params=['location', 'startTime'], relative_path='v1b3/projects/{projectId}/jobs/{jobId}/metrics', request_field='', request_type_name='DataflowProjectsJobsGetMetricsRequest', response_type_name='JobMetrics', supports_download=False, ) def List(self, request, global_params=None): r"""List the jobs of a project. To list the jobs of a project in a region, we recommend using `projects.locations.jobs.list` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). To list the all jobs across all regions, use `projects.jobs.aggregated`. Using `projects.jobs.list` is not recommended, because you can only get the list of jobs that are running in `us-central1`. `projects.locations.jobs.list` and `projects.jobs.list` support filtering the list of jobs by name. Filtering by name isn't supported by `projects.jobs.aggregated`. Args: request: (DataflowProjectsJobsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ListJobsResponse) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.jobs.list', ordered_params=['projectId'], path_params=['projectId'], query_params=['filter', 'location', 'name', 'pageSize', 'pageToken', 'view'], relative_path='v1b3/projects/{projectId}/jobs', request_field='', request_type_name='DataflowProjectsJobsListRequest', response_type_name='ListJobsResponse', supports_download=False, ) def Snapshot(self, request, global_params=None): r"""Snapshot the state of a streaming job. Args: request: (DataflowProjectsJobsSnapshotRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Snapshot) The response message. """ config = self.GetMethodConfig('Snapshot') return self._RunMethod( config, request, global_params=global_params) Snapshot.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.jobs.snapshot', ordered_params=['projectId', 'jobId'], path_params=['jobId', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/jobs/{jobId}:snapshot', request_field='snapshotJobRequest', request_type_name='DataflowProjectsJobsSnapshotRequest', response_type_name='Snapshot', supports_download=False, ) def Update(self, request, global_params=None): r"""Updates the state of an existing Cloud Dataflow job. To update the state of an existing job, we recommend using `projects.locations.jobs.update` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.jobs.update` is not recommended, as you can only update the state of jobs that are running in `us-central1`. Args: request: (DataflowProjectsJobsUpdateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Job) The response message. """ config = self.GetMethodConfig('Update') return self._RunMethod( config, request, global_params=global_params) Update.method_config = lambda: base_api.ApiMethodInfo( http_method='PUT', method_id='dataflow.projects.jobs.update', ordered_params=['projectId', 'jobId'], path_params=['jobId', 'projectId'], query_params=['location', 'updateMask'], relative_path='v1b3/projects/{projectId}/jobs/{jobId}', request_field='job', request_type_name='DataflowProjectsJobsUpdateRequest', response_type_name='Job', supports_download=False, ) class ProjectsLocationsFlexTemplatesService(base_api.BaseApiService): """Service class for the projects_locations_flexTemplates resource.""" _NAME = 'projects_locations_flexTemplates' def __init__(self, client): super(DataflowV1b3.ProjectsLocationsFlexTemplatesService, self).__init__(client) self._upload_configs = { } def Launch(self, request, global_params=None): r"""Launch a job with a FlexTemplate. Args: request: (DataflowProjectsLocationsFlexTemplatesLaunchRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (LaunchFlexTemplateResponse) The response message. """ config = self.GetMethodConfig('Launch') return self._RunMethod( config, request, global_params=global_params) Launch.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.locations.flexTemplates.launch', ordered_params=['projectId', 'location'], path_params=['location', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/locations/{location}/flexTemplates:launch', request_field='launchFlexTemplateRequest', request_type_name='DataflowProjectsLocationsFlexTemplatesLaunchRequest', response_type_name='LaunchFlexTemplateResponse', supports_download=False, ) class ProjectsLocationsJobsDebugService(base_api.BaseApiService): """Service class for the projects_locations_jobs_debug resource.""" _NAME = 'projects_locations_jobs_debug' def __init__(self, client): super(DataflowV1b3.ProjectsLocationsJobsDebugService, self).__init__(client) self._upload_configs = { } def GetConfig(self, request, global_params=None): r"""Get encoded debug configuration for component. Not cacheable. Args: request: (DataflowProjectsLocationsJobsDebugGetConfigRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (GetDebugConfigResponse) The response message. """ config = self.GetMethodConfig('GetConfig') return self._RunMethod( config, request, global_params=global_params) GetConfig.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.locations.jobs.debug.getConfig', ordered_params=['projectId', 'location', 'jobId'], path_params=['jobId', 'location', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs/{jobId}/debug/getConfig', request_field='getDebugConfigRequest', request_type_name='DataflowProjectsLocationsJobsDebugGetConfigRequest', response_type_name='GetDebugConfigResponse', supports_download=False, ) def GetWorkerStacktraces(self, request, global_params=None): r"""Get worker stacktraces from debug capture. Args: request: (DataflowProjectsLocationsJobsDebugGetWorkerStacktracesRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (GetWorkerStacktracesResponse) The response message. """ config = self.GetMethodConfig('GetWorkerStacktraces') return self._RunMethod( config, request, global_params=global_params) GetWorkerStacktraces.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.locations.jobs.debug.getWorkerStacktraces', ordered_params=['projectId', 'location', 'jobId'], path_params=['jobId', 'location', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs/{jobId}/debug/getWorkerStacktraces', request_field='getWorkerStacktracesRequest', request_type_name='DataflowProjectsLocationsJobsDebugGetWorkerStacktracesRequest', response_type_name='GetWorkerStacktracesResponse', supports_download=False, ) def SendCapture(self, request, global_params=None): r"""Send encoded debug capture data for component. Args: request: (DataflowProjectsLocationsJobsDebugSendCaptureRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (SendDebugCaptureResponse) The response message. """ config = self.GetMethodConfig('SendCapture') return self._RunMethod( config, request, global_params=global_params) SendCapture.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.locations.jobs.debug.sendCapture', ordered_params=['projectId', 'location', 'jobId'], path_params=['jobId', 'location', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs/{jobId}/debug/sendCapture', request_field='sendDebugCaptureRequest', request_type_name='DataflowProjectsLocationsJobsDebugSendCaptureRequest', response_type_name='SendDebugCaptureResponse', supports_download=False, ) class ProjectsLocationsJobsMessagesService(base_api.BaseApiService): """Service class for the projects_locations_jobs_messages resource.""" _NAME = 'projects_locations_jobs_messages' def __init__(self, client): super(DataflowV1b3.ProjectsLocationsJobsMessagesService, self).__init__(client) self._upload_configs = { } def List(self, request, global_params=None): r"""Request the job status. To request the status of a job, we recommend using `projects.locations.jobs.messages.list` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.jobs.messages.list` is not recommended, as you can only request the status of jobs that are running in `us-central1`. Args: request: (DataflowProjectsLocationsJobsMessagesListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ListJobMessagesResponse) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.locations.jobs.messages.list', ordered_params=['projectId', 'location', 'jobId'], path_params=['jobId', 'location', 'projectId'], query_params=['endTime', 'minimumImportance', 'pageSize', 'pageToken', 'startTime'], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs/{jobId}/messages', request_field='', request_type_name='DataflowProjectsLocationsJobsMessagesListRequest', response_type_name='ListJobMessagesResponse', supports_download=False, ) class ProjectsLocationsJobsSnapshotsService(base_api.BaseApiService): """Service class for the projects_locations_jobs_snapshots resource.""" _NAME = 'projects_locations_jobs_snapshots' def __init__(self, client): super(DataflowV1b3.ProjectsLocationsJobsSnapshotsService, self).__init__(client) self._upload_configs = { } def List(self, request, global_params=None): r"""Lists snapshots. Args: request: (DataflowProjectsLocationsJobsSnapshotsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ListSnapshotsResponse) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.locations.jobs.snapshots.list', ordered_params=['projectId', 'location', 'jobId'], path_params=['jobId', 'location', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs/{jobId}/snapshots', request_field='', request_type_name='DataflowProjectsLocationsJobsSnapshotsListRequest', response_type_name='ListSnapshotsResponse', supports_download=False, ) class ProjectsLocationsJobsStagesService(base_api.BaseApiService): """Service class for the projects_locations_jobs_stages resource.""" _NAME = 'projects_locations_jobs_stages' def __init__(self, client): super(DataflowV1b3.ProjectsLocationsJobsStagesService, self).__init__(client) self._upload_configs = { } def GetExecutionDetails(self, request, global_params=None): r"""Request detailed information about the execution status of a stage of the job. EXPERIMENTAL. This API is subject to change or removal without notice. Args: request: (DataflowProjectsLocationsJobsStagesGetExecutionDetailsRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (StageExecutionDetails) The response message. """ config = self.GetMethodConfig('GetExecutionDetails') return self._RunMethod( config, request, global_params=global_params) GetExecutionDetails.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.locations.jobs.stages.getExecutionDetails', ordered_params=['projectId', 'location', 'jobId', 'stageId'], path_params=['jobId', 'location', 'projectId', 'stageId'], query_params=['endTime', 'pageSize', 'pageToken', 'startTime'], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs/{jobId}/stages/{stageId}/executionDetails', request_field='', request_type_name='DataflowProjectsLocationsJobsStagesGetExecutionDetailsRequest', response_type_name='StageExecutionDetails', supports_download=False, ) class ProjectsLocationsJobsWorkItemsService(base_api.BaseApiService): """Service class for the projects_locations_jobs_workItems resource.""" _NAME = 'projects_locations_jobs_workItems' def __init__(self, client): super(DataflowV1b3.ProjectsLocationsJobsWorkItemsService, self).__init__(client) self._upload_configs = { } def Lease(self, request, global_params=None): r"""Leases a dataflow WorkItem to run. Args: request: (DataflowProjectsLocationsJobsWorkItemsLeaseRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (LeaseWorkItemResponse) The response message. """ config = self.GetMethodConfig('Lease') return self._RunMethod( config, request, global_params=global_params) Lease.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.locations.jobs.workItems.lease', ordered_params=['projectId', 'location', 'jobId'], path_params=['jobId', 'location', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs/{jobId}/workItems:lease', request_field='leaseWorkItemRequest', request_type_name='DataflowProjectsLocationsJobsWorkItemsLeaseRequest', response_type_name='LeaseWorkItemResponse', supports_download=False, ) def ReportStatus(self, request, global_params=None): r"""Reports the status of dataflow WorkItems leased by a worker. Args: request: (DataflowProjectsLocationsJobsWorkItemsReportStatusRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ReportWorkItemStatusResponse) The response message. """ config = self.GetMethodConfig('ReportStatus') return self._RunMethod( config, request, global_params=global_params) ReportStatus.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.locations.jobs.workItems.reportStatus', ordered_params=['projectId', 'location', 'jobId'], path_params=['jobId', 'location', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs/{jobId}/workItems:reportStatus', request_field='reportWorkItemStatusRequest', request_type_name='DataflowProjectsLocationsJobsWorkItemsReportStatusRequest', response_type_name='ReportWorkItemStatusResponse', supports_download=False, ) class ProjectsLocationsJobsService(base_api.BaseApiService): """Service class for the projects_locations_jobs resource.""" _NAME = 'projects_locations_jobs' def __init__(self, client): super(DataflowV1b3.ProjectsLocationsJobsService, self).__init__(client) self._upload_configs = { } def Create(self, request, global_params=None): r"""Creates a Dataflow job. To create a job, we recommend using `projects.locations.jobs.create` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.jobs.create` is not recommended, as your job will always start in `us-central1`. Do not enter confidential information when you supply string values using the API. Args: request: (DataflowProjectsLocationsJobsCreateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Job) The response message. """ config = self.GetMethodConfig('Create') return self._RunMethod( config, request, global_params=global_params) Create.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.locations.jobs.create', ordered_params=['projectId', 'location'], path_params=['location', 'projectId'], query_params=['replaceJobId', 'view'], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs', request_field='job', request_type_name='DataflowProjectsLocationsJobsCreateRequest', response_type_name='Job', supports_download=False, ) def Get(self, request, global_params=None): r"""Gets the state of the specified Cloud Dataflow job. To get the state of a job, we recommend using `projects.locations.jobs.get` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.jobs.get` is not recommended, as you can only get the state of jobs that are running in `us-central1`. Args: request: (DataflowProjectsLocationsJobsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Job) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params) Get.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.locations.jobs.get', ordered_params=['projectId', 'location', 'jobId'], path_params=['jobId', 'location', 'projectId'], query_params=['view'], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs/{jobId}', request_field='', request_type_name='DataflowProjectsLocationsJobsGetRequest', response_type_name='Job', supports_download=False, ) def GetExecutionDetails(self, request, global_params=None): r"""Request detailed information about the execution status of the job. EXPERIMENTAL. This API is subject to change or removal without notice. Args: request: (DataflowProjectsLocationsJobsGetExecutionDetailsRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (JobExecutionDetails) The response message. """ config = self.GetMethodConfig('GetExecutionDetails') return self._RunMethod( config, request, global_params=global_params) GetExecutionDetails.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.locations.jobs.getExecutionDetails', ordered_params=['projectId', 'location', 'jobId'], path_params=['jobId', 'location', 'projectId'], query_params=['pageSize', 'pageToken'], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs/{jobId}/executionDetails', request_field='', request_type_name='DataflowProjectsLocationsJobsGetExecutionDetailsRequest', response_type_name='JobExecutionDetails', supports_download=False, ) def GetMetrics(self, request, global_params=None): r"""Request the job status. To request the status of a job, we recommend using `projects.locations.jobs.getMetrics` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.jobs.getMetrics` is not recommended, as you can only request the status of jobs that are running in `us-central1`. Args: request: (DataflowProjectsLocationsJobsGetMetricsRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (JobMetrics) The response message. """ config = self.GetMethodConfig('GetMetrics') return self._RunMethod( config, request, global_params=global_params) GetMetrics.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.locations.jobs.getMetrics', ordered_params=['projectId', 'location', 'jobId'], path_params=['jobId', 'location', 'projectId'], query_params=['startTime'], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs/{jobId}/metrics', request_field='', request_type_name='DataflowProjectsLocationsJobsGetMetricsRequest', response_type_name='JobMetrics', supports_download=False, ) def List(self, request, global_params=None): r"""List the jobs of a project. To list the jobs of a project in a region, we recommend using `projects.locations.jobs.list` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). To list the all jobs across all regions, use `projects.jobs.aggregated`. Using `projects.jobs.list` is not recommended, because you can only get the list of jobs that are running in `us-central1`. `projects.locations.jobs.list` and `projects.jobs.list` support filtering the list of jobs by name. Filtering by name isn't supported by `projects.jobs.aggregated`. Args: request: (DataflowProjectsLocationsJobsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ListJobsResponse) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.locations.jobs.list', ordered_params=['projectId', 'location'], path_params=['location', 'projectId'], query_params=['filter', 'name', 'pageSize', 'pageToken', 'view'], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs', request_field='', request_type_name='DataflowProjectsLocationsJobsListRequest', response_type_name='ListJobsResponse', supports_download=False, ) def Snapshot(self, request, global_params=None): r"""Snapshot the state of a streaming job. Args: request: (DataflowProjectsLocationsJobsSnapshotRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Snapshot) The response message. """ config = self.GetMethodConfig('Snapshot') return self._RunMethod( config, request, global_params=global_params) Snapshot.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.locations.jobs.snapshot', ordered_params=['projectId', 'location', 'jobId'], path_params=['jobId', 'location', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs/{jobId}:snapshot', request_field='snapshotJobRequest', request_type_name='DataflowProjectsLocationsJobsSnapshotRequest', response_type_name='Snapshot', supports_download=False, ) def Update(self, request, global_params=None): r"""Updates the state of an existing Cloud Dataflow job. To update the state of an existing job, we recommend using `projects.locations.jobs.update` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.jobs.update` is not recommended, as you can only update the state of jobs that are running in `us-central1`. Args: request: (DataflowProjectsLocationsJobsUpdateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Job) The response message. """ config = self.GetMethodConfig('Update') return self._RunMethod( config, request, global_params=global_params) Update.method_config = lambda: base_api.ApiMethodInfo( http_method='PUT', method_id='dataflow.projects.locations.jobs.update', ordered_params=['projectId', 'location', 'jobId'], path_params=['jobId', 'location', 'projectId'], query_params=['updateMask'], relative_path='v1b3/projects/{projectId}/locations/{location}/jobs/{jobId}', request_field='job', request_type_name='DataflowProjectsLocationsJobsUpdateRequest', response_type_name='Job', supports_download=False, ) class ProjectsLocationsSnapshotsService(base_api.BaseApiService): """Service class for the projects_locations_snapshots resource.""" _NAME = 'projects_locations_snapshots' def __init__(self, client): super(DataflowV1b3.ProjectsLocationsSnapshotsService, self).__init__(client) self._upload_configs = { } def Delete(self, request, global_params=None): r"""Deletes a snapshot. Args: request: (DataflowProjectsLocationsSnapshotsDeleteRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (DeleteSnapshotResponse) The response message. """ config = self.GetMethodConfig('Delete') return self._RunMethod( config, request, global_params=global_params) Delete.method_config = lambda: base_api.ApiMethodInfo( http_method='DELETE', method_id='dataflow.projects.locations.snapshots.delete', ordered_params=['projectId', 'location', 'snapshotId'], path_params=['location', 'projectId', 'snapshotId'], query_params=[], relative_path='v1b3/projects/{projectId}/locations/{location}/snapshots/{snapshotId}', request_field='', request_type_name='DataflowProjectsLocationsSnapshotsDeleteRequest', response_type_name='DeleteSnapshotResponse', supports_download=False, ) def Get(self, request, global_params=None): r"""Gets information about a snapshot. Args: request: (DataflowProjectsLocationsSnapshotsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Snapshot) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params) Get.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.locations.snapshots.get', ordered_params=['projectId', 'location', 'snapshotId'], path_params=['location', 'projectId', 'snapshotId'], query_params=[], relative_path='v1b3/projects/{projectId}/locations/{location}/snapshots/{snapshotId}', request_field='', request_type_name='DataflowProjectsLocationsSnapshotsGetRequest', response_type_name='Snapshot', supports_download=False, ) def List(self, request, global_params=None): r"""Lists snapshots. Args: request: (DataflowProjectsLocationsSnapshotsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ListSnapshotsResponse) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.locations.snapshots.list', ordered_params=['projectId', 'location'], path_params=['location', 'projectId'], query_params=['jobId'], relative_path='v1b3/projects/{projectId}/locations/{location}/snapshots', request_field='', request_type_name='DataflowProjectsLocationsSnapshotsListRequest', response_type_name='ListSnapshotsResponse', supports_download=False, ) class ProjectsLocationsTemplatesService(base_api.BaseApiService): """Service class for the projects_locations_templates resource.""" _NAME = 'projects_locations_templates' def __init__(self, client): super(DataflowV1b3.ProjectsLocationsTemplatesService, self).__init__(client) self._upload_configs = { } def Create(self, request, global_params=None): r"""Creates a Cloud Dataflow job from a template. Do not enter confidential information when you supply string values using the API. To create a job, we recommend using `projects.locations.templates.create` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.templates.create` is not recommended, because your job will always start in `us-central1`. Args: request: (DataflowProjectsLocationsTemplatesCreateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Job) The response message. """ config = self.GetMethodConfig('Create') return self._RunMethod( config, request, global_params=global_params) Create.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.locations.templates.create', ordered_params=['projectId', 'location'], path_params=['location', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/locations/{location}/templates', request_field='createJobFromTemplateRequest', request_type_name='DataflowProjectsLocationsTemplatesCreateRequest', response_type_name='Job', supports_download=False, ) def Get(self, request, global_params=None): r"""Get the template associated with a template. To get the template, we recommend using `projects.locations.templates.get` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.templates.get` is not recommended, because only templates that are running in `us-central1` are retrieved. Args: request: (DataflowProjectsLocationsTemplatesGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (GetTemplateResponse) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params) Get.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.locations.templates.get', ordered_params=['projectId', 'location'], path_params=['location', 'projectId'], query_params=['gcsPath', 'view'], relative_path='v1b3/projects/{projectId}/locations/{location}/templates:get', request_field='', request_type_name='DataflowProjectsLocationsTemplatesGetRequest', response_type_name='GetTemplateResponse', supports_download=False, ) def Launch(self, request, global_params=None): r"""Launches a template. To launch a template, we recommend using `projects.locations.templates.launch` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.templates.launch` is not recommended, because jobs launched from the template will always start in `us-central1`. Args: request: (DataflowProjectsLocationsTemplatesLaunchRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (LaunchTemplateResponse) The response message. """ config = self.GetMethodConfig('Launch') return self._RunMethod( config, request, global_params=global_params) Launch.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.locations.templates.launch', ordered_params=['projectId', 'location'], path_params=['location', 'projectId'], query_params=['dynamicTemplate_gcsPath', 'dynamicTemplate_stagingLocation', 'gcsPath', 'validateOnly'], relative_path='v1b3/projects/{projectId}/locations/{location}/templates:launch', request_field='launchTemplateParameters', request_type_name='DataflowProjectsLocationsTemplatesLaunchRequest', response_type_name='LaunchTemplateResponse', supports_download=False, ) class ProjectsLocationsService(base_api.BaseApiService): """Service class for the projects_locations resource.""" _NAME = 'projects_locations' def __init__(self, client): super(DataflowV1b3.ProjectsLocationsService, self).__init__(client) self._upload_configs = { } def WorkerMessages(self, request, global_params=None): r"""Send a worker_message to the service. Args: request: (DataflowProjectsLocationsWorkerMessagesRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (SendWorkerMessagesResponse) The response message. """ config = self.GetMethodConfig('WorkerMessages') return self._RunMethod( config, request, global_params=global_params) WorkerMessages.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.locations.workerMessages', ordered_params=['projectId', 'location'], path_params=['location', 'projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/locations/{location}/WorkerMessages', request_field='sendWorkerMessagesRequest', request_type_name='DataflowProjectsLocationsWorkerMessagesRequest', response_type_name='SendWorkerMessagesResponse', supports_download=False, ) class ProjectsSnapshotsService(base_api.BaseApiService): """Service class for the projects_snapshots resource.""" _NAME = 'projects_snapshots' def __init__(self, client): super(DataflowV1b3.ProjectsSnapshotsService, self).__init__(client) self._upload_configs = { } def Get(self, request, global_params=None): r"""Gets information about a snapshot. Args: request: (DataflowProjectsSnapshotsGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Snapshot) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params) Get.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.snapshots.get', ordered_params=['projectId', 'snapshotId'], path_params=['projectId', 'snapshotId'], query_params=['location'], relative_path='v1b3/projects/{projectId}/snapshots/{snapshotId}', request_field='', request_type_name='DataflowProjectsSnapshotsGetRequest', response_type_name='Snapshot', supports_download=False, ) def List(self, request, global_params=None): r"""Lists snapshots. Args: request: (DataflowProjectsSnapshotsListRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (ListSnapshotsResponse) The response message. """ config = self.GetMethodConfig('List') return self._RunMethod( config, request, global_params=global_params) List.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.snapshots.list', ordered_params=['projectId'], path_params=['projectId'], query_params=['jobId', 'location'], relative_path='v1b3/projects/{projectId}/snapshots', request_field='', request_type_name='DataflowProjectsSnapshotsListRequest', response_type_name='ListSnapshotsResponse', supports_download=False, ) class ProjectsTemplatesService(base_api.BaseApiService): """Service class for the projects_templates resource.""" _NAME = 'projects_templates' def __init__(self, client): super(DataflowV1b3.ProjectsTemplatesService, self).__init__(client) self._upload_configs = { } def Create(self, request, global_params=None): r"""Creates a Cloud Dataflow job from a template. Do not enter confidential information when you supply string values using the API. To create a job, we recommend using `projects.locations.templates.create` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.templates.create` is not recommended, because your job will always start in `us-central1`. Args: request: (DataflowProjectsTemplatesCreateRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (Job) The response message. """ config = self.GetMethodConfig('Create') return self._RunMethod( config, request, global_params=global_params) Create.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.templates.create', ordered_params=['projectId'], path_params=['projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/templates', request_field='createJobFromTemplateRequest', request_type_name='DataflowProjectsTemplatesCreateRequest', response_type_name='Job', supports_download=False, ) def Get(self, request, global_params=None): r"""Get the template associated with a template. To get the template, we recommend using `projects.locations.templates.get` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.templates.get` is not recommended, because only templates that are running in `us-central1` are retrieved. Args: request: (DataflowProjectsTemplatesGetRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (GetTemplateResponse) The response message. """ config = self.GetMethodConfig('Get') return self._RunMethod( config, request, global_params=global_params) Get.method_config = lambda: base_api.ApiMethodInfo( http_method='GET', method_id='dataflow.projects.templates.get', ordered_params=['projectId'], path_params=['projectId'], query_params=['gcsPath', 'location', 'view'], relative_path='v1b3/projects/{projectId}/templates:get', request_field='', request_type_name='DataflowProjectsTemplatesGetRequest', response_type_name='GetTemplateResponse', supports_download=False, ) def Launch(self, request, global_params=None): r"""Launches a template. To launch a template, we recommend using `projects.locations.templates.launch` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). Using `projects.templates.launch` is not recommended, because jobs launched from the template will always start in `us-central1`. Args: request: (DataflowProjectsTemplatesLaunchRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (LaunchTemplateResponse) The response message. """ config = self.GetMethodConfig('Launch') return self._RunMethod( config, request, global_params=global_params) Launch.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.templates.launch', ordered_params=['projectId'], path_params=['projectId'], query_params=['dynamicTemplate_gcsPath', 'dynamicTemplate_stagingLocation', 'gcsPath', 'location', 'validateOnly'], relative_path='v1b3/projects/{projectId}/templates:launch', request_field='launchTemplateParameters', request_type_name='DataflowProjectsTemplatesLaunchRequest', response_type_name='LaunchTemplateResponse', supports_download=False, ) class ProjectsService(base_api.BaseApiService): """Service class for the projects resource.""" _NAME = 'projects' def __init__(self, client): super(DataflowV1b3.ProjectsService, self).__init__(client) self._upload_configs = { } def DeleteSnapshots(self, request, global_params=None): r"""Deletes a snapshot. Args: request: (DataflowProjectsDeleteSnapshotsRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (DeleteSnapshotResponse) The response message. """ config = self.GetMethodConfig('DeleteSnapshots') return self._RunMethod( config, request, global_params=global_params) DeleteSnapshots.method_config = lambda: base_api.ApiMethodInfo( http_method='DELETE', method_id='dataflow.projects.deleteSnapshots', ordered_params=['projectId'], path_params=['projectId'], query_params=['location', 'snapshotId'], relative_path='v1b3/projects/{projectId}/snapshots', request_field='', request_type_name='DataflowProjectsDeleteSnapshotsRequest', response_type_name='DeleteSnapshotResponse', supports_download=False, ) def WorkerMessages(self, request, global_params=None): r"""Send a worker_message to the service. Args: request: (DataflowProjectsWorkerMessagesRequest) input message global_params: (StandardQueryParameters, default: None) global arguments Returns: (SendWorkerMessagesResponse) The response message. """ config = self.GetMethodConfig('WorkerMessages') return self._RunMethod( config, request, global_params=global_params) WorkerMessages.method_config = lambda: base_api.ApiMethodInfo( http_method='POST', method_id='dataflow.projects.workerMessages', ordered_params=['projectId'], path_params=['projectId'], query_params=[], relative_path='v1b3/projects/{projectId}/WorkerMessages', request_field='sendWorkerMessagesRequest', request_type_name='DataflowProjectsWorkerMessagesRequest', response_type_name='SendWorkerMessagesResponse', supports_download=False, ) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/internal/clients/dataflow/dataflow_v1b3_messages.py0000644000175100017510000117654015206113032033143 0ustar00runnerrunner"""Generated message classes for dataflow version v1b3. Manages Google Cloud Dataflow projects on Google Cloud Platform. """ # NOTE: This file is autogenerated and should not be edited by hand. from apitools.base.protorpclite import messages as _messages from apitools.base.py import encoding from apitools.base.py import extra_types package = 'dataflow' class ApproximateProgress(_messages.Message): r"""Obsolete in favor of ApproximateReportedProgress and ApproximateSplitRequest. Fields: percentComplete: Obsolete. position: Obsolete. remainingTime: Obsolete. """ percentComplete = _messages.FloatField(1, variant=_messages.Variant.FLOAT) position = _messages.MessageField('Position', 2) remainingTime = _messages.StringField(3) class ApproximateReportedProgress(_messages.Message): r"""A progress measurement of a WorkItem by a worker. Fields: consumedParallelism: Total amount of parallelism in the portion of input of this task that has already been consumed and is no longer active. In the first two examples above (see remaining_parallelism), the value should be 29 or 2 respectively. The sum of remaining_parallelism and consumed_parallelism should equal the total amount of parallelism in this work item. If specified, must be finite. fractionConsumed: Completion as fraction of the input consumed, from 0.0 (beginning, nothing consumed), to 1.0 (end of the input, entire input consumed). position: A Position within the work to represent a progress. remainingParallelism: Total amount of parallelism in the input of this task that remains, (i.e. can be delegated to this task and any new tasks via dynamic splitting). Always at least 1 for non-finished work items and 0 for finished. "Amount of parallelism" refers to how many non-empty parts of the input can be read in parallel. This does not necessarily equal number of records. An input that can be read in parallel down to the individual records is called "perfectly splittable". An example of non-perfectly parallelizable input is a block-compressed file format where a block of records has to be read as a whole, but different blocks can be read in parallel. Examples: * If we are processing record #30 (starting at 1) out of 50 in a perfectly splittable 50-record input, this value should be 21 (20 remaining + 1 current). * If we are reading through block 3 in a block-compressed file consisting of 5 blocks, this value should be 3 (since blocks 4 and 5 can be processed in parallel by new tasks via dynamic splitting and the current task remains processing block 3). * If we are reading through the last block in a block- compressed file, or reading or processing the last record in a perfectly splittable input, this value should be 1, because apart from the current task, no additional remainder can be split off. """ consumedParallelism = _messages.MessageField('ReportedParallelism', 1) fractionConsumed = _messages.FloatField(2) position = _messages.MessageField('Position', 3) remainingParallelism = _messages.MessageField('ReportedParallelism', 4) class ApproximateSplitRequest(_messages.Message): r"""A suggestion by the service to the worker to dynamically split the WorkItem. Fields: fractionConsumed: A fraction at which to split the work item, from 0.0 (beginning of the input) to 1.0 (end of the input). fractionOfRemainder: The fraction of the remainder of work to split the work item at, from 0.0 (split at the current position) to 1.0 (end of the input). position: A Position at which to split the work item. """ fractionConsumed = _messages.FloatField(1) fractionOfRemainder = _messages.FloatField(2) position = _messages.MessageField('Position', 3) class AutoscalingEvent(_messages.Message): r"""A structured message reporting an autoscaling decision made by the Dataflow service. Enums: EventTypeValueValuesEnum: The type of autoscaling event to report. Fields: currentNumWorkers: The current number of workers the job has. description: A message describing why the system decided to adjust the current number of workers, why it failed, or why the system decided to not make any changes to the number of workers. eventType: The type of autoscaling event to report. targetNumWorkers: The target number of workers the worker pool wants to resize to use. time: The time this event was emitted to indicate a new target or current num_workers value. workerPool: A short and friendly name for the worker pool this event refers to. """ class EventTypeValueValuesEnum(_messages.Enum): r"""The type of autoscaling event to report. Values: TYPE_UNKNOWN: Default type for the enum. Value should never be returned. TARGET_NUM_WORKERS_CHANGED: The TARGET_NUM_WORKERS_CHANGED type should be used when the target worker pool size has changed at the start of an actuation. An event should always be specified as TARGET_NUM_WORKERS_CHANGED if it reflects a change in the target_num_workers. CURRENT_NUM_WORKERS_CHANGED: The CURRENT_NUM_WORKERS_CHANGED type should be used when actual worker pool size has been changed, but the target_num_workers has not changed. ACTUATION_FAILURE: The ACTUATION_FAILURE type should be used when we want to report an error to the user indicating why the current number of workers in the pool could not be changed. Displayed in the current status and history widgets. NO_CHANGE: Used when we want to report to the user a reason why we are not currently adjusting the number of workers. Should specify both target_num_workers, current_num_workers and a decision_message. """ TYPE_UNKNOWN = 0 TARGET_NUM_WORKERS_CHANGED = 1 CURRENT_NUM_WORKERS_CHANGED = 2 ACTUATION_FAILURE = 3 NO_CHANGE = 4 currentNumWorkers = _messages.IntegerField(1) description = _messages.MessageField('StructuredMessage', 2) eventType = _messages.EnumField('EventTypeValueValuesEnum', 3) targetNumWorkers = _messages.IntegerField(4) time = _messages.StringField(5) workerPool = _messages.StringField(6) class AutoscalingSettings(_messages.Message): r"""Settings for WorkerPool autoscaling. Enums: AlgorithmValueValuesEnum: The algorithm to use for autoscaling. Fields: algorithm: The algorithm to use for autoscaling. maxNumWorkers: The maximum number of workers to cap scaling at. """ class AlgorithmValueValuesEnum(_messages.Enum): r"""The algorithm to use for autoscaling. Values: AUTOSCALING_ALGORITHM_UNKNOWN: The algorithm is unknown, or unspecified. AUTOSCALING_ALGORITHM_NONE: Disable autoscaling. AUTOSCALING_ALGORITHM_BASIC: Increase worker count over time to reduce job execution time. """ AUTOSCALING_ALGORITHM_UNKNOWN = 0 AUTOSCALING_ALGORITHM_NONE = 1 AUTOSCALING_ALGORITHM_BASIC = 2 algorithm = _messages.EnumField('AlgorithmValueValuesEnum', 1) maxNumWorkers = _messages.IntegerField(2, variant=_messages.Variant.INT32) class Base2Exponent(_messages.Message): r"""Exponential buckets where the growth factor between buckets is `2**(2**-scale)`. e.g. for `scale=1` growth factor is `2**(2**(-1))=sqrt(2)`. `n` buckets will have the following boundaries. - 0th: [0, gf) - i in [1, n-1]: [gf^(i), gf^(i+1)) Fields: numberOfBuckets: Must be greater than 0. scale: Must be between -3 and 3. This forces the growth factor of the bucket boundaries to be between `2^(1/8)` and `256`. """ numberOfBuckets = _messages.IntegerField(1, variant=_messages.Variant.INT32) scale = _messages.IntegerField(2, variant=_messages.Variant.INT32) class BigQueryIODetails(_messages.Message): r"""Metadata for a BigQuery connector used by the job. Fields: dataset: Dataset accessed in the connection. projectId: Project accessed in the connection. query: Query used to access data in the connection. table: Table accessed in the connection. """ dataset = _messages.StringField(1) projectId = _messages.StringField(2) query = _messages.StringField(3) table = _messages.StringField(4) class BigTableIODetails(_messages.Message): r"""Metadata for a Cloud Bigtable connector used by the job. Fields: instanceId: InstanceId accessed in the connection. projectId: ProjectId accessed in the connection. tableId: TableId accessed in the connection. """ instanceId = _messages.StringField(1) projectId = _messages.StringField(2) tableId = _messages.StringField(3) class BoundedTrie(_messages.Message): r"""The message type used for encoding metrics of type bounded trie. Fields: bound: The maximum number of elements to store before truncation. root: A compact representation of all the elements in this trie. singleton: A more efficient representation for metrics consisting of a single value. """ bound = _messages.IntegerField(1, variant=_messages.Variant.INT32) root = _messages.MessageField('BoundedTrieNode', 2) singleton = _messages.StringField(3, repeated=True) class BoundedTrieNode(_messages.Message): r"""A single node in a BoundedTrie. Messages: ChildrenValue: Children of this node. Must be empty if truncated is true. Fields: children: Children of this node. Must be empty if truncated is true. truncated: Whether this node has been truncated. A truncated leaf represents possibly many children with the same prefix. """ @encoding.MapUnrecognizedFields('additionalProperties') class ChildrenValue(_messages.Message): r"""Children of this node. Must be empty if truncated is true. Messages: AdditionalProperty: An additional property for a ChildrenValue object. Fields: additionalProperties: Additional properties of type ChildrenValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a ChildrenValue object. Fields: key: Name of the additional property. value: A BoundedTrieNode attribute. """ key = _messages.StringField(1) value = _messages.MessageField('BoundedTrieNode', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) children = _messages.MessageField('ChildrenValue', 1) truncated = _messages.BooleanField(2) class BucketOptions(_messages.Message): r"""`BucketOptions` describes the bucket boundaries used in the histogram. Fields: exponential: Bucket boundaries grow exponentially. linear: Bucket boundaries grow linearly. """ exponential = _messages.MessageField('Base2Exponent', 1) linear = _messages.MessageField('Linear', 2) class CPUTime(_messages.Message): r"""Modeled after information exposed by /proc/stat. Fields: rate: Average CPU utilization rate (% non-idle cpu / second) since previous sample. timestamp: Timestamp of the measurement. totalMs: Total active CPU time across all cores (ie., non-idle) in milliseconds since start-up. """ rate = _messages.FloatField(1) timestamp = _messages.StringField(2) totalMs = _messages.IntegerField(3, variant=_messages.Variant.UINT64) class ComponentSource(_messages.Message): r"""Description of an interstitial value between transforms in an execution stage. Fields: name: Dataflow service generated name for this source. originalTransformOrCollection: User name for the original user transform or collection with which this source is most closely associated. userName: Human-readable name for this transform; may be user or system generated. """ name = _messages.StringField(1) originalTransformOrCollection = _messages.StringField(2) userName = _messages.StringField(3) class ComponentTransform(_messages.Message): r"""Description of a transform executed as part of an execution stage. Fields: name: Dataflow service generated name for this source. originalTransform: User name for the original user transform with which this transform is most closely associated. userName: Human-readable name for this transform; may be user or system generated. """ name = _messages.StringField(1) originalTransform = _messages.StringField(2) userName = _messages.StringField(3) class ComputationTopology(_messages.Message): r"""All configuration data for a particular Computation. Fields: computationId: The ID of the computation. inputs: The inputs to the computation. keyRanges: The key ranges processed by the computation. outputs: The outputs from the computation. stateFamilies: The state family values. systemStageName: The system stage name. """ computationId = _messages.StringField(1) inputs = _messages.MessageField('StreamLocation', 2, repeated=True) keyRanges = _messages.MessageField('KeyRangeLocation', 3, repeated=True) outputs = _messages.MessageField('StreamLocation', 4, repeated=True) stateFamilies = _messages.MessageField('StateFamilyConfig', 5, repeated=True) systemStageName = _messages.StringField(6) class ConcatPosition(_messages.Message): r"""A position that encapsulates an inner position and an index for the inner position. A ConcatPosition can be used by a reader of a source that encapsulates a set of other sources. Fields: index: Index of the inner source. position: Position within the inner source. """ index = _messages.IntegerField(1, variant=_messages.Variant.INT32) position = _messages.MessageField('Position', 2) class ContainerSpec(_messages.Message): r"""Container Spec. Fields: defaultEnvironment: Default runtime environment for the job. image: Name of the docker container image. E.g., gcr.io/project/some-image imageRepositoryCertPath: Cloud Storage path to self-signed certificate of private registry. imageRepositoryPasswordSecretId: Secret Manager secret id for password to authenticate to private registry. imageRepositoryUsernameSecretId: Secret Manager secret id for username to authenticate to private registry. metadata: Metadata describing a template including description and validation rules. sdkInfo: Required. SDK info of the Flex Template. """ defaultEnvironment = _messages.MessageField('FlexTemplateRuntimeEnvironment', 1) image = _messages.StringField(2) imageRepositoryCertPath = _messages.StringField(3) imageRepositoryPasswordSecretId = _messages.StringField(4) imageRepositoryUsernameSecretId = _messages.StringField(5) metadata = _messages.MessageField('TemplateMetadata', 6) sdkInfo = _messages.MessageField('SDKInfo', 7) class CounterMetadata(_messages.Message): r"""CounterMetadata includes all static non-name non-value counter attributes. Enums: KindValueValuesEnum: Counter aggregation kind. StandardUnitsValueValuesEnum: System defined Units, see above enum. Fields: description: Human-readable description of the counter semantics. kind: Counter aggregation kind. otherUnits: A string referring to the unit type. standardUnits: System defined Units, see above enum. """ class KindValueValuesEnum(_messages.Enum): r"""Counter aggregation kind. Values: INVALID: Counter aggregation kind was not set. SUM: Aggregated value is the sum of all contributed values. MAX: Aggregated value is the max of all contributed values. MIN: Aggregated value is the min of all contributed values. MEAN: Aggregated value is the mean of all contributed values. OR: Aggregated value represents the logical 'or' of all contributed values. AND: Aggregated value represents the logical 'and' of all contributed values. SET: Aggregated value is a set of unique contributed values. DISTRIBUTION: Aggregated value captures statistics about a distribution. LATEST_VALUE: Aggregated value tracks the latest value of a variable. """ INVALID = 0 SUM = 1 MAX = 2 MIN = 3 MEAN = 4 OR = 5 AND = 6 SET = 7 DISTRIBUTION = 8 LATEST_VALUE = 9 class StandardUnitsValueValuesEnum(_messages.Enum): r"""System defined Units, see above enum. Values: BYTES: Counter returns a value in bytes. BYTES_PER_SEC: Counter returns a value in bytes per second. MILLISECONDS: Counter returns a value in milliseconds. MICROSECONDS: Counter returns a value in microseconds. NANOSECONDS: Counter returns a value in nanoseconds. TIMESTAMP_MSEC: Counter returns a timestamp in milliseconds. TIMESTAMP_USEC: Counter returns a timestamp in microseconds. TIMESTAMP_NSEC: Counter returns a timestamp in nanoseconds. """ BYTES = 0 BYTES_PER_SEC = 1 MILLISECONDS = 2 MICROSECONDS = 3 NANOSECONDS = 4 TIMESTAMP_MSEC = 5 TIMESTAMP_USEC = 6 TIMESTAMP_NSEC = 7 description = _messages.StringField(1) kind = _messages.EnumField('KindValueValuesEnum', 2) otherUnits = _messages.StringField(3) standardUnits = _messages.EnumField('StandardUnitsValueValuesEnum', 4) class CounterStructuredName(_messages.Message): r"""Identifies a counter within a per-job namespace. Counters whose structured names are the same get merged into a single value for the job. Enums: OriginValueValuesEnum: One of the standard Origins defined above. PortionValueValuesEnum: Portion of this counter, either key or value. Fields: componentStepName: Name of the optimized step being executed by the workers. executionStepName: Name of the stage. An execution step contains multiple component steps. inputIndex: Index of an input collection that's being read from/written to as a side input. The index identifies a step's side inputs starting by 1 (e.g. the first side input has input_index 1, the third has input_index 3). Side inputs are identified by a pair of (original_step_name, input_index). This field helps uniquely identify them. name: Counter name. Not necessarily globally-unique, but unique within the context of the other fields. Required. origin: One of the standard Origins defined above. originNamespace: A string containing a more specific namespace of the counter's origin. originalRequestingStepName: The step name requesting an operation, such as GBK. I.e. the ParDo causing a read/write from shuffle to occur, or a read from side inputs. originalStepName: System generated name of the original step in the user's graph, before optimization. portion: Portion of this counter, either key or value. workerId: ID of a particular worker. """ class OriginValueValuesEnum(_messages.Enum): r"""One of the standard Origins defined above. Values: SYSTEM: Counter was created by the Dataflow system. USER: Counter was created by the user. """ SYSTEM = 0 USER = 1 class PortionValueValuesEnum(_messages.Enum): r"""Portion of this counter, either key or value. Values: ALL: Counter portion has not been set. KEY: Counter reports a key. VALUE: Counter reports a value. """ ALL = 0 KEY = 1 VALUE = 2 componentStepName = _messages.StringField(1) executionStepName = _messages.StringField(2) inputIndex = _messages.IntegerField(3, variant=_messages.Variant.INT32) name = _messages.StringField(4) origin = _messages.EnumField('OriginValueValuesEnum', 5) originNamespace = _messages.StringField(6) originalRequestingStepName = _messages.StringField(7) originalStepName = _messages.StringField(8) portion = _messages.EnumField('PortionValueValuesEnum', 9) workerId = _messages.StringField(10) class CounterStructuredNameAndMetadata(_messages.Message): r"""A single message which encapsulates structured name and metadata for a given counter. Fields: metadata: Metadata associated with a counter name: Structured name of the counter. """ metadata = _messages.MessageField('CounterMetadata', 1) name = _messages.MessageField('CounterStructuredName', 2) class CounterUpdate(_messages.Message): r"""An update to a Counter sent from a worker. Next ID: 17 Fields: boolean: Boolean value for And, Or. boundedTrie: Bounded trie data cumulative: True if this counter is reported as the total cumulative aggregate value accumulated since the worker started working on this WorkItem. By default this is false, indicating that this counter is reported as a delta. distribution: Distribution data floatingPoint: Floating point value for Sum, Max, Min. floatingPointList: List of floating point numbers, for Set. floatingPointMean: Floating point mean aggregation value for Mean. integer: Integer value for Sum, Max, Min. integerGauge: Gauge data integerList: List of integers, for Set. integerMean: Integer mean aggregation value for Mean. internal: Value for internally-defined counters used by the Dataflow service. nameAndKind: Counter name and aggregation type. shortId: The service-generated short identifier for this counter. The short_id -> (name, metadata) mapping is constant for the lifetime of a job. stringList: List of strings, for Set. structuredNameAndMetadata: Counter structured name and metadata. """ boolean = _messages.BooleanField(1) boundedTrie = _messages.MessageField('BoundedTrie', 2) cumulative = _messages.BooleanField(3) distribution = _messages.MessageField('DistributionUpdate', 4) floatingPoint = _messages.FloatField(5) floatingPointList = _messages.MessageField('FloatingPointList', 6) floatingPointMean = _messages.MessageField('FloatingPointMean', 7) integer = _messages.MessageField('SplitInt64', 8) integerGauge = _messages.MessageField('IntegerGauge', 9) integerList = _messages.MessageField('IntegerList', 10) integerMean = _messages.MessageField('IntegerMean', 11) internal = _messages.MessageField('extra_types.JsonValue', 12) nameAndKind = _messages.MessageField('NameAndKind', 13) shortId = _messages.IntegerField(14) stringList = _messages.MessageField('StringList', 15) structuredNameAndMetadata = _messages.MessageField('CounterStructuredNameAndMetadata', 16) class CreateJobFromTemplateRequest(_messages.Message): r"""A request to create a Cloud Dataflow job from a template. Messages: ParametersValue: The runtime parameters to pass to the job. Fields: environment: The runtime environment for the job. gcsPath: Required. A Cloud Storage path to the template from which to create the job. Must be a valid Cloud Storage URL, beginning with `gs://`. jobName: Required. The job name to use for the created job. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) to which to direct the request. parameters: The runtime parameters to pass to the job. """ @encoding.MapUnrecognizedFields('additionalProperties') class ParametersValue(_messages.Message): r"""The runtime parameters to pass to the job. Messages: AdditionalProperty: An additional property for a ParametersValue object. Fields: additionalProperties: Additional properties of type ParametersValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a ParametersValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) environment = _messages.MessageField('RuntimeEnvironment', 1) gcsPath = _messages.StringField(2) jobName = _messages.StringField(3) location = _messages.StringField(4) parameters = _messages.MessageField('ParametersValue', 5) class CustomSourceLocation(_messages.Message): r"""Identifies the location of a custom souce. Fields: stateful: Whether this source is stateful. """ stateful = _messages.BooleanField(1) class DataDiskAssignment(_messages.Message): r"""Data disk assignment for a given VM instance. Fields: dataDisks: Mounted data disks. The order is important a data disk's 0-based index in this list defines which persistent directory the disk is mounted to, for example the list of { "myproject-1014-104817-4c2-harness-0-disk-0" }, { "myproject-1014-104817-4c2-harness-0-disk-1" }. vmInstance: VM instance name the data disks mounted to, for example "myproject-1014-104817-4c2-harness-0". """ dataDisks = _messages.StringField(1, repeated=True) vmInstance = _messages.StringField(2) class DataSamplingConfig(_messages.Message): r"""Configuration options for sampling elements. Enums: BehaviorsValueListEntryValuesEnum: Fields: behaviors: List of given sampling behaviors to enable. For example, specifying behaviors = [ALWAYS_ON] samples in-flight elements but does not sample exceptions. Can be used to specify multiple behaviors like, behaviors = [ALWAYS_ON, EXCEPTIONS] for specifying periodic sampling and exception sampling. If DISABLED is in the list, then sampling will be disabled and ignore the other given behaviors. Ordering does not matter. """ class BehaviorsValueListEntryValuesEnum(_messages.Enum): r"""BehaviorsValueListEntryValuesEnum enum type. Values: DATA_SAMPLING_BEHAVIOR_UNSPECIFIED: If given, has no effect on sampling behavior. Used as an unknown or unset sentinel value. DISABLED: When given, disables element sampling. Has same behavior as not setting the behavior. ALWAYS_ON: When given, enables sampling in-flight from all PCollections. EXCEPTIONS: When given, enables sampling input elements when a user- defined DoFn causes an exception. """ DATA_SAMPLING_BEHAVIOR_UNSPECIFIED = 0 DISABLED = 1 ALWAYS_ON = 2 EXCEPTIONS = 3 behaviors = _messages.EnumField('BehaviorsValueListEntryValuesEnum', 1, repeated=True) class DataSamplingReport(_messages.Message): r"""Contains per-worker telemetry about the data sampling feature. Fields: bytesWrittenDelta: Optional. Delta of bytes written to file from previous report. elementsSampledBytes: Optional. Delta of bytes sampled from previous report. elementsSampledCount: Optional. Delta of number of elements sampled from previous report. exceptionsSampledCount: Optional. Delta of number of samples taken from user code exceptions from previous report. pcollectionsSampledCount: Optional. Delta of number of PCollections sampled from previous report. persistenceErrorsCount: Optional. Delta of errors counts from persisting the samples from previous report. translationErrorsCount: Optional. Delta of errors counts from retrieving, or translating the samples from previous report. """ bytesWrittenDelta = _messages.IntegerField(1) elementsSampledBytes = _messages.IntegerField(2) elementsSampledCount = _messages.IntegerField(3) exceptionsSampledCount = _messages.IntegerField(4) pcollectionsSampledCount = _messages.IntegerField(5) persistenceErrorsCount = _messages.IntegerField(6) translationErrorsCount = _messages.IntegerField(7) class DataflowGaugeValue(_messages.Message): r"""The gauge value of a metric. Fields: measuredTime: The timestamp when the gauge was recorded. value: The value of the gauge. """ measuredTime = _messages.StringField(1) value = _messages.IntegerField(2) class DataflowHistogramValue(_messages.Message): r"""Summary statistics for a population of values. HistogramValue contains a sequence of buckets and gives a count of values that fall into each bucket. Bucket boundares are defined by a formula and bucket widths are either fixed or exponentially increasing. Fields: bucketCounts: Optional. The number of values in each bucket of the histogram, as described in `bucket_options`. `bucket_counts` should contain N values, where N is the number of buckets specified in `bucket_options`. If `bucket_counts` has fewer than N values, the remaining values are assumed to be 0. bucketOptions: Describes the bucket boundaries used in the histogram. count: Number of values recorded in this histogram. outlierStats: Statistics on the values recorded in the histogram that fall out of the bucket boundaries. """ bucketCounts = _messages.IntegerField(1, repeated=True) bucketOptions = _messages.MessageField('BucketOptions', 2) count = _messages.IntegerField(3) outlierStats = _messages.MessageField('OutlierStats', 4) class DataflowProjectsDeleteSnapshotsRequest(_messages.Message): r"""A DataflowProjectsDeleteSnapshotsRequest object. Fields: location: The location that contains this snapshot. projectId: The ID of the Cloud Platform project that the snapshot belongs to. snapshotId: The ID of the snapshot. """ location = _messages.StringField(1) projectId = _messages.StringField(2, required=True) snapshotId = _messages.StringField(3) class DataflowProjectsJobsAggregatedRequest(_messages.Message): r"""A DataflowProjectsJobsAggregatedRequest object. Enums: FilterValueValuesEnum: The kind of filter to use. ViewValueValuesEnum: Deprecated. ListJobs always returns summaries now. Use GetJob for other JobViews. Fields: filter: The kind of filter to use. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains this job. name: Optional. The job name. pageSize: If there are many jobs, limit response to at most this many. The actual number of jobs returned will be the lesser of max_responses and an unspecified server-defined limit. pageToken: Set this to the 'next_page_token' field of a previous response to request additional results in a long list. projectId: The project which owns the jobs. view: Deprecated. ListJobs always returns summaries now. Use GetJob for other JobViews. """ class FilterValueValuesEnum(_messages.Enum): r"""The kind of filter to use. Values: UNKNOWN: The filter isn't specified, or is unknown. This returns all jobs ordered on descending `JobUuid`. ALL: Returns all running jobs first ordered on creation timestamp, then returns all terminated jobs ordered on the termination timestamp. TERMINATED: Filters the jobs that have a terminated state, ordered on the termination timestamp. Example terminated states: `JOB_STATE_STOPPED`, `JOB_STATE_UPDATED`, `JOB_STATE_DRAINED`, etc. ACTIVE: Filters the jobs that are running ordered on the creation timestamp. """ UNKNOWN = 0 ALL = 1 TERMINATED = 2 ACTIVE = 3 class ViewValueValuesEnum(_messages.Enum): r"""Deprecated. ListJobs always returns summaries now. Use GetJob for other JobViews. Values: JOB_VIEW_UNKNOWN: The job view to return isn't specified, or is unknown. Responses will contain at least the `JOB_VIEW_SUMMARY` information, and may contain additional information. JOB_VIEW_SUMMARY: Request summary information only: Project ID, Job ID, job name, job type, job status, start/end time, and Cloud SDK version details. JOB_VIEW_ALL: Request all information available for this job. When the job is in `JOB_STATE_PENDING`, the job has been created but is not yet running, and not all job information is available. For complete job information, wait until the job in is `JOB_STATE_RUNNING`. For more information, see [JobState](https://cloud.google.com/dataflow/docs/ref erence/rest/v1b3/projects.jobs#jobstate). JOB_VIEW_DESCRIPTION: Request summary info and limited job description data for steps, labels and environment. """ JOB_VIEW_UNKNOWN = 0 JOB_VIEW_SUMMARY = 1 JOB_VIEW_ALL = 2 JOB_VIEW_DESCRIPTION = 3 filter = _messages.EnumField('FilterValueValuesEnum', 1) location = _messages.StringField(2) name = _messages.StringField(3) pageSize = _messages.IntegerField(4, variant=_messages.Variant.INT32) pageToken = _messages.StringField(5) projectId = _messages.StringField(6, required=True) view = _messages.EnumField('ViewValueValuesEnum', 7) class DataflowProjectsJobsCreateRequest(_messages.Message): r"""A DataflowProjectsJobsCreateRequest object. Enums: ViewValueValuesEnum: The level of information requested in response. Fields: job: A Job resource to be passed as the request body. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains this job. projectId: The ID of the Cloud Platform project that the job belongs to. replaceJobId: Deprecated. This field is now in the Job message. view: The level of information requested in response. """ class ViewValueValuesEnum(_messages.Enum): r"""The level of information requested in response. Values: JOB_VIEW_UNKNOWN: The job view to return isn't specified, or is unknown. Responses will contain at least the `JOB_VIEW_SUMMARY` information, and may contain additional information. JOB_VIEW_SUMMARY: Request summary information only: Project ID, Job ID, job name, job type, job status, start/end time, and Cloud SDK version details. JOB_VIEW_ALL: Request all information available for this job. When the job is in `JOB_STATE_PENDING`, the job has been created but is not yet running, and not all job information is available. For complete job information, wait until the job in is `JOB_STATE_RUNNING`. For more information, see [JobState](https://cloud.google.com/dataflow/docs/ref erence/rest/v1b3/projects.jobs#jobstate). JOB_VIEW_DESCRIPTION: Request summary info and limited job description data for steps, labels and environment. """ JOB_VIEW_UNKNOWN = 0 JOB_VIEW_SUMMARY = 1 JOB_VIEW_ALL = 2 JOB_VIEW_DESCRIPTION = 3 job = _messages.MessageField('Job', 1) location = _messages.StringField(2) projectId = _messages.StringField(3, required=True) replaceJobId = _messages.StringField(4) view = _messages.EnumField('ViewValueValuesEnum', 5) class DataflowProjectsJobsDebugGetConfigRequest(_messages.Message): r"""A DataflowProjectsJobsDebugGetConfigRequest object. Fields: getDebugConfigRequest: A GetDebugConfigRequest resource to be passed as the request body. jobId: The job id. projectId: The project id. """ getDebugConfigRequest = _messages.MessageField('GetDebugConfigRequest', 1) jobId = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) class DataflowProjectsJobsDebugSendCaptureRequest(_messages.Message): r"""A DataflowProjectsJobsDebugSendCaptureRequest object. Fields: jobId: The job id. projectId: The project id. sendDebugCaptureRequest: A SendDebugCaptureRequest resource to be passed as the request body. """ jobId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) sendDebugCaptureRequest = _messages.MessageField('SendDebugCaptureRequest', 3) class DataflowProjectsJobsGetMetricsRequest(_messages.Message): r"""A DataflowProjectsJobsGetMetricsRequest object. Fields: jobId: The job to get metrics for. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the job specified by job_id. projectId: A project id. startTime: Return only metric data that has changed since this time. Default is to return all information about all metrics for the job. """ jobId = _messages.StringField(1, required=True) location = _messages.StringField(2) projectId = _messages.StringField(3, required=True) startTime = _messages.StringField(4) class DataflowProjectsJobsGetRequest(_messages.Message): r"""A DataflowProjectsJobsGetRequest object. Enums: ViewValueValuesEnum: The level of information requested in response. Fields: jobId: The job ID. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains this job. projectId: The ID of the Cloud Platform project that the job belongs to. view: The level of information requested in response. """ class ViewValueValuesEnum(_messages.Enum): r"""The level of information requested in response. Values: JOB_VIEW_UNKNOWN: The job view to return isn't specified, or is unknown. Responses will contain at least the `JOB_VIEW_SUMMARY` information, and may contain additional information. JOB_VIEW_SUMMARY: Request summary information only: Project ID, Job ID, job name, job type, job status, start/end time, and Cloud SDK version details. JOB_VIEW_ALL: Request all information available for this job. When the job is in `JOB_STATE_PENDING`, the job has been created but is not yet running, and not all job information is available. For complete job information, wait until the job in is `JOB_STATE_RUNNING`. For more information, see [JobState](https://cloud.google.com/dataflow/docs/ref erence/rest/v1b3/projects.jobs#jobstate). JOB_VIEW_DESCRIPTION: Request summary info and limited job description data for steps, labels and environment. """ JOB_VIEW_UNKNOWN = 0 JOB_VIEW_SUMMARY = 1 JOB_VIEW_ALL = 2 JOB_VIEW_DESCRIPTION = 3 jobId = _messages.StringField(1, required=True) location = _messages.StringField(2) projectId = _messages.StringField(3, required=True) view = _messages.EnumField('ViewValueValuesEnum', 4) class DataflowProjectsJobsListRequest(_messages.Message): r"""A DataflowProjectsJobsListRequest object. Enums: FilterValueValuesEnum: The kind of filter to use. ViewValueValuesEnum: Deprecated. ListJobs always returns summaries now. Use GetJob for other JobViews. Fields: filter: The kind of filter to use. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains this job. name: Optional. The job name. pageSize: If there are many jobs, limit response to at most this many. The actual number of jobs returned will be the lesser of max_responses and an unspecified server-defined limit. pageToken: Set this to the 'next_page_token' field of a previous response to request additional results in a long list. projectId: The project which owns the jobs. view: Deprecated. ListJobs always returns summaries now. Use GetJob for other JobViews. """ class FilterValueValuesEnum(_messages.Enum): r"""The kind of filter to use. Values: UNKNOWN: The filter isn't specified, or is unknown. This returns all jobs ordered on descending `JobUuid`. ALL: Returns all running jobs first ordered on creation timestamp, then returns all terminated jobs ordered on the termination timestamp. TERMINATED: Filters the jobs that have a terminated state, ordered on the termination timestamp. Example terminated states: `JOB_STATE_STOPPED`, `JOB_STATE_UPDATED`, `JOB_STATE_DRAINED`, etc. ACTIVE: Filters the jobs that are running ordered on the creation timestamp. """ UNKNOWN = 0 ALL = 1 TERMINATED = 2 ACTIVE = 3 class ViewValueValuesEnum(_messages.Enum): r"""Deprecated. ListJobs always returns summaries now. Use GetJob for other JobViews. Values: JOB_VIEW_UNKNOWN: The job view to return isn't specified, or is unknown. Responses will contain at least the `JOB_VIEW_SUMMARY` information, and may contain additional information. JOB_VIEW_SUMMARY: Request summary information only: Project ID, Job ID, job name, job type, job status, start/end time, and Cloud SDK version details. JOB_VIEW_ALL: Request all information available for this job. When the job is in `JOB_STATE_PENDING`, the job has been created but is not yet running, and not all job information is available. For complete job information, wait until the job in is `JOB_STATE_RUNNING`. For more information, see [JobState](https://cloud.google.com/dataflow/docs/ref erence/rest/v1b3/projects.jobs#jobstate). JOB_VIEW_DESCRIPTION: Request summary info and limited job description data for steps, labels and environment. """ JOB_VIEW_UNKNOWN = 0 JOB_VIEW_SUMMARY = 1 JOB_VIEW_ALL = 2 JOB_VIEW_DESCRIPTION = 3 filter = _messages.EnumField('FilterValueValuesEnum', 1) location = _messages.StringField(2) name = _messages.StringField(3) pageSize = _messages.IntegerField(4, variant=_messages.Variant.INT32) pageToken = _messages.StringField(5) projectId = _messages.StringField(6, required=True) view = _messages.EnumField('ViewValueValuesEnum', 7) class DataflowProjectsJobsMessagesListRequest(_messages.Message): r"""A DataflowProjectsJobsMessagesListRequest object. Enums: MinimumImportanceValueValuesEnum: Filter to only get messages with importance >= level Fields: endTime: Return only messages with timestamps < end_time. The default is now (i.e. return up to the latest messages available). jobId: The job to get messages about. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the job specified by job_id. minimumImportance: Filter to only get messages with importance >= level pageSize: If specified, determines the maximum number of messages to return. If unspecified, the service may choose an appropriate default, or may return an arbitrarily large number of results. pageToken: If supplied, this should be the value of next_page_token returned by an earlier call. This will cause the next page of results to be returned. projectId: A project id. startTime: If specified, return only messages with timestamps >= start_time. The default is the job creation time (i.e. beginning of messages). """ class MinimumImportanceValueValuesEnum(_messages.Enum): r"""Filter to only get messages with importance >= level Values: JOB_MESSAGE_IMPORTANCE_UNKNOWN: The message importance isn't specified, or is unknown. JOB_MESSAGE_DEBUG: The message is at the 'debug' level: typically only useful for software engineers working on the code the job is running. Typically, Dataflow pipeline runners do not display log messages at this level by default. JOB_MESSAGE_DETAILED: The message is at the 'detailed' level: somewhat verbose, but potentially useful to users. Typically, Dataflow pipeline runners do not display log messages at this level by default. These messages are displayed by default in the Dataflow monitoring UI. JOB_MESSAGE_BASIC: The message is at the 'basic' level: useful for keeping track of the execution of a Dataflow pipeline. Typically, Dataflow pipeline runners display log messages at this level by default, and these messages are displayed by default in the Dataflow monitoring UI. JOB_MESSAGE_WARNING: The message is at the 'warning' level: indicating a condition pertaining to a job which may require human intervention. Typically, Dataflow pipeline runners display log messages at this level by default, and these messages are displayed by default in the Dataflow monitoring UI. JOB_MESSAGE_ERROR: The message is at the 'error' level: indicating a condition preventing a job from succeeding. Typically, Dataflow pipeline runners display log messages at this level by default, and these messages are displayed by default in the Dataflow monitoring UI. """ JOB_MESSAGE_IMPORTANCE_UNKNOWN = 0 JOB_MESSAGE_DEBUG = 1 JOB_MESSAGE_DETAILED = 2 JOB_MESSAGE_BASIC = 3 JOB_MESSAGE_WARNING = 4 JOB_MESSAGE_ERROR = 5 endTime = _messages.StringField(1) jobId = _messages.StringField(2, required=True) location = _messages.StringField(3) minimumImportance = _messages.EnumField('MinimumImportanceValueValuesEnum', 4) pageSize = _messages.IntegerField(5, variant=_messages.Variant.INT32) pageToken = _messages.StringField(6) projectId = _messages.StringField(7, required=True) startTime = _messages.StringField(8) class DataflowProjectsJobsSnapshotRequest(_messages.Message): r"""A DataflowProjectsJobsSnapshotRequest object. Fields: jobId: The job to be snapshotted. projectId: The project which owns the job to be snapshotted. snapshotJobRequest: A SnapshotJobRequest resource to be passed as the request body. """ jobId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) snapshotJobRequest = _messages.MessageField('SnapshotJobRequest', 3) class DataflowProjectsJobsUpdateRequest(_messages.Message): r"""A DataflowProjectsJobsUpdateRequest object. Fields: job: A Job resource to be passed as the request body. jobId: The job ID. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains this job. projectId: The ID of the Cloud Platform project that the job belongs to. updateMask: The list of fields to update relative to Job. If empty, only RequestedJobState will be considered for update. If the FieldMask is not empty and RequestedJobState is none/empty, The fields specified in the update mask will be the only ones considered for update. If both RequestedJobState and update_mask are specified, an error will be returned as we cannot update both state and mask. """ job = _messages.MessageField('Job', 1) jobId = _messages.StringField(2, required=True) location = _messages.StringField(3) projectId = _messages.StringField(4, required=True) updateMask = _messages.StringField(5) class DataflowProjectsJobsWorkItemsLeaseRequest(_messages.Message): r"""A DataflowProjectsJobsWorkItemsLeaseRequest object. Fields: jobId: Identifies the workflow job this worker belongs to. leaseWorkItemRequest: A LeaseWorkItemRequest resource to be passed as the request body. projectId: Identifies the project this worker belongs to. """ jobId = _messages.StringField(1, required=True) leaseWorkItemRequest = _messages.MessageField('LeaseWorkItemRequest', 2) projectId = _messages.StringField(3, required=True) class DataflowProjectsJobsWorkItemsReportStatusRequest(_messages.Message): r"""A DataflowProjectsJobsWorkItemsReportStatusRequest object. Fields: jobId: The job which the WorkItem is part of. projectId: The project which owns the WorkItem's job. reportWorkItemStatusRequest: A ReportWorkItemStatusRequest resource to be passed as the request body. """ jobId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) reportWorkItemStatusRequest = _messages.MessageField('ReportWorkItemStatusRequest', 3) class DataflowProjectsLocationsFlexTemplatesLaunchRequest(_messages.Message): r"""A DataflowProjectsLocationsFlexTemplatesLaunchRequest object. Fields: launchFlexTemplateRequest: A LaunchFlexTemplateRequest resource to be passed as the request body. location: Required. The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) to which to direct the request. E.g., us-central1, us-west1. projectId: Required. The ID of the Cloud Platform project that the job belongs to. """ launchFlexTemplateRequest = _messages.MessageField('LaunchFlexTemplateRequest', 1) location = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) class DataflowProjectsLocationsJobsCreateRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsCreateRequest object. Enums: ViewValueValuesEnum: The level of information requested in response. Fields: job: A Job resource to be passed as the request body. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains this job. projectId: The ID of the Cloud Platform project that the job belongs to. replaceJobId: Deprecated. This field is now in the Job message. view: The level of information requested in response. """ class ViewValueValuesEnum(_messages.Enum): r"""The level of information requested in response. Values: JOB_VIEW_UNKNOWN: The job view to return isn't specified, or is unknown. Responses will contain at least the `JOB_VIEW_SUMMARY` information, and may contain additional information. JOB_VIEW_SUMMARY: Request summary information only: Project ID, Job ID, job name, job type, job status, start/end time, and Cloud SDK version details. JOB_VIEW_ALL: Request all information available for this job. When the job is in `JOB_STATE_PENDING`, the job has been created but is not yet running, and not all job information is available. For complete job information, wait until the job in is `JOB_STATE_RUNNING`. For more information, see [JobState](https://cloud.google.com/dataflow/docs/ref erence/rest/v1b3/projects.jobs#jobstate). JOB_VIEW_DESCRIPTION: Request summary info and limited job description data for steps, labels and environment. """ JOB_VIEW_UNKNOWN = 0 JOB_VIEW_SUMMARY = 1 JOB_VIEW_ALL = 2 JOB_VIEW_DESCRIPTION = 3 job = _messages.MessageField('Job', 1) location = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) replaceJobId = _messages.StringField(4) view = _messages.EnumField('ViewValueValuesEnum', 5) class DataflowProjectsLocationsJobsDebugGetConfigRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsDebugGetConfigRequest object. Fields: getDebugConfigRequest: A GetDebugConfigRequest resource to be passed as the request body. jobId: The job id. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the job specified by job_id. projectId: The project id. """ getDebugConfigRequest = _messages.MessageField('GetDebugConfigRequest', 1) jobId = _messages.StringField(2, required=True) location = _messages.StringField(3, required=True) projectId = _messages.StringField(4, required=True) class DataflowProjectsLocationsJobsDebugGetWorkerStacktracesRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsDebugGetWorkerStacktracesRequest object. Fields: getWorkerStacktracesRequest: A GetWorkerStacktracesRequest resource to be passed as the request body. jobId: The job for which to get stacktraces. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the job specified by job_id. projectId: The project id. """ getWorkerStacktracesRequest = _messages.MessageField('GetWorkerStacktracesRequest', 1) jobId = _messages.StringField(2, required=True) location = _messages.StringField(3, required=True) projectId = _messages.StringField(4, required=True) class DataflowProjectsLocationsJobsDebugSendCaptureRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsDebugSendCaptureRequest object. Fields: jobId: The job id. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the job specified by job_id. projectId: The project id. sendDebugCaptureRequest: A SendDebugCaptureRequest resource to be passed as the request body. """ jobId = _messages.StringField(1, required=True) location = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) sendDebugCaptureRequest = _messages.MessageField('SendDebugCaptureRequest', 4) class DataflowProjectsLocationsJobsGetExecutionDetailsRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsGetExecutionDetailsRequest object. Fields: jobId: The job to get execution details for. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the job specified by job_id. pageSize: If specified, determines the maximum number of stages to return. If unspecified, the service may choose an appropriate default, or may return an arbitrarily large number of results. pageToken: If supplied, this should be the value of next_page_token returned by an earlier call. This will cause the next page of results to be returned. projectId: A project id. """ jobId = _messages.StringField(1, required=True) location = _messages.StringField(2, required=True) pageSize = _messages.IntegerField(3, variant=_messages.Variant.INT32) pageToken = _messages.StringField(4) projectId = _messages.StringField(5, required=True) class DataflowProjectsLocationsJobsGetMetricsRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsGetMetricsRequest object. Fields: jobId: The job to get metrics for. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the job specified by job_id. projectId: A project id. startTime: Return only metric data that has changed since this time. Default is to return all information about all metrics for the job. """ jobId = _messages.StringField(1, required=True) location = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) startTime = _messages.StringField(4) class DataflowProjectsLocationsJobsGetRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsGetRequest object. Enums: ViewValueValuesEnum: The level of information requested in response. Fields: jobId: The job ID. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains this job. projectId: The ID of the Cloud Platform project that the job belongs to. view: The level of information requested in response. """ class ViewValueValuesEnum(_messages.Enum): r"""The level of information requested in response. Values: JOB_VIEW_UNKNOWN: The job view to return isn't specified, or is unknown. Responses will contain at least the `JOB_VIEW_SUMMARY` information, and may contain additional information. JOB_VIEW_SUMMARY: Request summary information only: Project ID, Job ID, job name, job type, job status, start/end time, and Cloud SDK version details. JOB_VIEW_ALL: Request all information available for this job. When the job is in `JOB_STATE_PENDING`, the job has been created but is not yet running, and not all job information is available. For complete job information, wait until the job in is `JOB_STATE_RUNNING`. For more information, see [JobState](https://cloud.google.com/dataflow/docs/ref erence/rest/v1b3/projects.jobs#jobstate). JOB_VIEW_DESCRIPTION: Request summary info and limited job description data for steps, labels and environment. """ JOB_VIEW_UNKNOWN = 0 JOB_VIEW_SUMMARY = 1 JOB_VIEW_ALL = 2 JOB_VIEW_DESCRIPTION = 3 jobId = _messages.StringField(1, required=True) location = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) view = _messages.EnumField('ViewValueValuesEnum', 4) class DataflowProjectsLocationsJobsListRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsListRequest object. Enums: FilterValueValuesEnum: The kind of filter to use. ViewValueValuesEnum: Deprecated. ListJobs always returns summaries now. Use GetJob for other JobViews. Fields: filter: The kind of filter to use. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains this job. name: Optional. The job name. pageSize: If there are many jobs, limit response to at most this many. The actual number of jobs returned will be the lesser of max_responses and an unspecified server-defined limit. pageToken: Set this to the 'next_page_token' field of a previous response to request additional results in a long list. projectId: The project which owns the jobs. view: Deprecated. ListJobs always returns summaries now. Use GetJob for other JobViews. """ class FilterValueValuesEnum(_messages.Enum): r"""The kind of filter to use. Values: UNKNOWN: The filter isn't specified, or is unknown. This returns all jobs ordered on descending `JobUuid`. ALL: Returns all running jobs first ordered on creation timestamp, then returns all terminated jobs ordered on the termination timestamp. TERMINATED: Filters the jobs that have a terminated state, ordered on the termination timestamp. Example terminated states: `JOB_STATE_STOPPED`, `JOB_STATE_UPDATED`, `JOB_STATE_DRAINED`, etc. ACTIVE: Filters the jobs that are running ordered on the creation timestamp. """ UNKNOWN = 0 ALL = 1 TERMINATED = 2 ACTIVE = 3 class ViewValueValuesEnum(_messages.Enum): r"""Deprecated. ListJobs always returns summaries now. Use GetJob for other JobViews. Values: JOB_VIEW_UNKNOWN: The job view to return isn't specified, or is unknown. Responses will contain at least the `JOB_VIEW_SUMMARY` information, and may contain additional information. JOB_VIEW_SUMMARY: Request summary information only: Project ID, Job ID, job name, job type, job status, start/end time, and Cloud SDK version details. JOB_VIEW_ALL: Request all information available for this job. When the job is in `JOB_STATE_PENDING`, the job has been created but is not yet running, and not all job information is available. For complete job information, wait until the job in is `JOB_STATE_RUNNING`. For more information, see [JobState](https://cloud.google.com/dataflow/docs/ref erence/rest/v1b3/projects.jobs#jobstate). JOB_VIEW_DESCRIPTION: Request summary info and limited job description data for steps, labels and environment. """ JOB_VIEW_UNKNOWN = 0 JOB_VIEW_SUMMARY = 1 JOB_VIEW_ALL = 2 JOB_VIEW_DESCRIPTION = 3 filter = _messages.EnumField('FilterValueValuesEnum', 1) location = _messages.StringField(2, required=True) name = _messages.StringField(3) pageSize = _messages.IntegerField(4, variant=_messages.Variant.INT32) pageToken = _messages.StringField(5) projectId = _messages.StringField(6, required=True) view = _messages.EnumField('ViewValueValuesEnum', 7) class DataflowProjectsLocationsJobsMessagesListRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsMessagesListRequest object. Enums: MinimumImportanceValueValuesEnum: Filter to only get messages with importance >= level Fields: endTime: Return only messages with timestamps < end_time. The default is now (i.e. return up to the latest messages available). jobId: The job to get messages about. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the job specified by job_id. minimumImportance: Filter to only get messages with importance >= level pageSize: If specified, determines the maximum number of messages to return. If unspecified, the service may choose an appropriate default, or may return an arbitrarily large number of results. pageToken: If supplied, this should be the value of next_page_token returned by an earlier call. This will cause the next page of results to be returned. projectId: A project id. startTime: If specified, return only messages with timestamps >= start_time. The default is the job creation time (i.e. beginning of messages). """ class MinimumImportanceValueValuesEnum(_messages.Enum): r"""Filter to only get messages with importance >= level Values: JOB_MESSAGE_IMPORTANCE_UNKNOWN: The message importance isn't specified, or is unknown. JOB_MESSAGE_DEBUG: The message is at the 'debug' level: typically only useful for software engineers working on the code the job is running. Typically, Dataflow pipeline runners do not display log messages at this level by default. JOB_MESSAGE_DETAILED: The message is at the 'detailed' level: somewhat verbose, but potentially useful to users. Typically, Dataflow pipeline runners do not display log messages at this level by default. These messages are displayed by default in the Dataflow monitoring UI. JOB_MESSAGE_BASIC: The message is at the 'basic' level: useful for keeping track of the execution of a Dataflow pipeline. Typically, Dataflow pipeline runners display log messages at this level by default, and these messages are displayed by default in the Dataflow monitoring UI. JOB_MESSAGE_WARNING: The message is at the 'warning' level: indicating a condition pertaining to a job which may require human intervention. Typically, Dataflow pipeline runners display log messages at this level by default, and these messages are displayed by default in the Dataflow monitoring UI. JOB_MESSAGE_ERROR: The message is at the 'error' level: indicating a condition preventing a job from succeeding. Typically, Dataflow pipeline runners display log messages at this level by default, and these messages are displayed by default in the Dataflow monitoring UI. """ JOB_MESSAGE_IMPORTANCE_UNKNOWN = 0 JOB_MESSAGE_DEBUG = 1 JOB_MESSAGE_DETAILED = 2 JOB_MESSAGE_BASIC = 3 JOB_MESSAGE_WARNING = 4 JOB_MESSAGE_ERROR = 5 endTime = _messages.StringField(1) jobId = _messages.StringField(2, required=True) location = _messages.StringField(3, required=True) minimumImportance = _messages.EnumField('MinimumImportanceValueValuesEnum', 4) pageSize = _messages.IntegerField(5, variant=_messages.Variant.INT32) pageToken = _messages.StringField(6) projectId = _messages.StringField(7, required=True) startTime = _messages.StringField(8) class DataflowProjectsLocationsJobsSnapshotRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsSnapshotRequest object. Fields: jobId: The job to be snapshotted. location: The location that contains this job. projectId: The project which owns the job to be snapshotted. snapshotJobRequest: A SnapshotJobRequest resource to be passed as the request body. """ jobId = _messages.StringField(1, required=True) location = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) snapshotJobRequest = _messages.MessageField('SnapshotJobRequest', 4) class DataflowProjectsLocationsJobsSnapshotsListRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsSnapshotsListRequest object. Fields: jobId: If specified, list snapshots created from this job. location: The location to list snapshots in. projectId: The project ID to list snapshots for. """ jobId = _messages.StringField(1, required=True) location = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) class DataflowProjectsLocationsJobsStagesGetExecutionDetailsRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsStagesGetExecutionDetailsRequest object. Fields: endTime: Upper time bound of work items to include, by start time. jobId: The job to get execution details for. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the job specified by job_id. pageSize: If specified, determines the maximum number of work items to return. If unspecified, the service may choose an appropriate default, or may return an arbitrarily large number of results. pageToken: If supplied, this should be the value of next_page_token returned by an earlier call. This will cause the next page of results to be returned. projectId: A project id. stageId: The stage for which to fetch information. startTime: Lower time bound of work items to include, by start time. """ endTime = _messages.StringField(1) jobId = _messages.StringField(2, required=True) location = _messages.StringField(3, required=True) pageSize = _messages.IntegerField(4, variant=_messages.Variant.INT32) pageToken = _messages.StringField(5) projectId = _messages.StringField(6, required=True) stageId = _messages.StringField(7, required=True) startTime = _messages.StringField(8) class DataflowProjectsLocationsJobsUpdateRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsUpdateRequest object. Fields: job: A Job resource to be passed as the request body. jobId: The job ID. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains this job. projectId: The ID of the Cloud Platform project that the job belongs to. updateMask: The list of fields to update relative to Job. If empty, only RequestedJobState will be considered for update. If the FieldMask is not empty and RequestedJobState is none/empty, The fields specified in the update mask will be the only ones considered for update. If both RequestedJobState and update_mask are specified, an error will be returned as we cannot update both state and mask. """ job = _messages.MessageField('Job', 1) jobId = _messages.StringField(2, required=True) location = _messages.StringField(3, required=True) projectId = _messages.StringField(4, required=True) updateMask = _messages.StringField(5) class DataflowProjectsLocationsJobsWorkItemsLeaseRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsWorkItemsLeaseRequest object. Fields: jobId: Identifies the workflow job this worker belongs to. leaseWorkItemRequest: A LeaseWorkItemRequest resource to be passed as the request body. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the WorkItem's job. projectId: Identifies the project this worker belongs to. """ jobId = _messages.StringField(1, required=True) leaseWorkItemRequest = _messages.MessageField('LeaseWorkItemRequest', 2) location = _messages.StringField(3, required=True) projectId = _messages.StringField(4, required=True) class DataflowProjectsLocationsJobsWorkItemsReportStatusRequest(_messages.Message): r"""A DataflowProjectsLocationsJobsWorkItemsReportStatusRequest object. Fields: jobId: The job which the WorkItem is part of. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the WorkItem's job. projectId: The project which owns the WorkItem's job. reportWorkItemStatusRequest: A ReportWorkItemStatusRequest resource to be passed as the request body. """ jobId = _messages.StringField(1, required=True) location = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) reportWorkItemStatusRequest = _messages.MessageField('ReportWorkItemStatusRequest', 4) class DataflowProjectsLocationsSnapshotsDeleteRequest(_messages.Message): r"""A DataflowProjectsLocationsSnapshotsDeleteRequest object. Fields: location: The location that contains this snapshot. projectId: The ID of the Cloud Platform project that the snapshot belongs to. snapshotId: The ID of the snapshot. """ location = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) snapshotId = _messages.StringField(3, required=True) class DataflowProjectsLocationsSnapshotsGetRequest(_messages.Message): r"""A DataflowProjectsLocationsSnapshotsGetRequest object. Fields: location: The location that contains this snapshot. projectId: The ID of the Cloud Platform project that the snapshot belongs to. snapshotId: The ID of the snapshot. """ location = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) snapshotId = _messages.StringField(3, required=True) class DataflowProjectsLocationsSnapshotsListRequest(_messages.Message): r"""A DataflowProjectsLocationsSnapshotsListRequest object. Fields: jobId: If specified, list snapshots created from this job. location: The location to list snapshots in. projectId: The project ID to list snapshots for. """ jobId = _messages.StringField(1) location = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) class DataflowProjectsLocationsTemplatesCreateRequest(_messages.Message): r"""A DataflowProjectsLocationsTemplatesCreateRequest object. Fields: createJobFromTemplateRequest: A CreateJobFromTemplateRequest resource to be passed as the request body. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) to which to direct the request. projectId: Required. The ID of the Cloud Platform project that the job belongs to. """ createJobFromTemplateRequest = _messages.MessageField('CreateJobFromTemplateRequest', 1) location = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) class DataflowProjectsLocationsTemplatesGetRequest(_messages.Message): r"""A DataflowProjectsLocationsTemplatesGetRequest object. Enums: ViewValueValuesEnum: The view to retrieve. Defaults to METADATA_ONLY. Fields: gcsPath: Required. A Cloud Storage path to the template from which to create the job. Must be valid Cloud Storage URL, beginning with 'gs://'. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) to which to direct the request. projectId: Required. The ID of the Cloud Platform project that the job belongs to. view: The view to retrieve. Defaults to METADATA_ONLY. """ class ViewValueValuesEnum(_messages.Enum): r"""The view to retrieve. Defaults to METADATA_ONLY. Values: METADATA_ONLY: Template view that retrieves only the metadata associated with the template. """ METADATA_ONLY = 0 gcsPath = _messages.StringField(1) location = _messages.StringField(2, required=True) projectId = _messages.StringField(3, required=True) view = _messages.EnumField('ViewValueValuesEnum', 4) class DataflowProjectsLocationsTemplatesLaunchRequest(_messages.Message): r"""A DataflowProjectsLocationsTemplatesLaunchRequest object. Fields: dynamicTemplate_gcsPath: Path to the dynamic template specification file on Cloud Storage. The file must be a JSON serialized `DynamicTemplateFileSpec` object. dynamicTemplate_stagingLocation: Cloud Storage path for staging dependencies. Must be a valid Cloud Storage URL, beginning with `gs://`. gcsPath: A Cloud Storage path to the template to use to create the job. Must be valid Cloud Storage URL, beginning with `gs://`. launchTemplateParameters: A LaunchTemplateParameters resource to be passed as the request body. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) to which to direct the request. projectId: Required. The ID of the Cloud Platform project that the job belongs to. validateOnly: If true, the request is validated but not actually executed. Defaults to false. """ dynamicTemplate_gcsPath = _messages.StringField(1) dynamicTemplate_stagingLocation = _messages.StringField(2) gcsPath = _messages.StringField(3) launchTemplateParameters = _messages.MessageField('LaunchTemplateParameters', 4) location = _messages.StringField(5, required=True) projectId = _messages.StringField(6, required=True) validateOnly = _messages.BooleanField(7) class DataflowProjectsLocationsWorkerMessagesRequest(_messages.Message): r"""A DataflowProjectsLocationsWorkerMessagesRequest object. Fields: location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the job. projectId: The project to send the WorkerMessages to. sendWorkerMessagesRequest: A SendWorkerMessagesRequest resource to be passed as the request body. """ location = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) sendWorkerMessagesRequest = _messages.MessageField('SendWorkerMessagesRequest', 3) class DataflowProjectsSnapshotsGetRequest(_messages.Message): r"""A DataflowProjectsSnapshotsGetRequest object. Fields: location: The location that contains this snapshot. projectId: The ID of the Cloud Platform project that the snapshot belongs to. snapshotId: The ID of the snapshot. """ location = _messages.StringField(1) projectId = _messages.StringField(2, required=True) snapshotId = _messages.StringField(3, required=True) class DataflowProjectsSnapshotsListRequest(_messages.Message): r"""A DataflowProjectsSnapshotsListRequest object. Fields: jobId: If specified, list snapshots created from this job. location: The location to list snapshots in. projectId: The project ID to list snapshots for. """ jobId = _messages.StringField(1) location = _messages.StringField(2) projectId = _messages.StringField(3, required=True) class DataflowProjectsTemplatesCreateRequest(_messages.Message): r"""A DataflowProjectsTemplatesCreateRequest object. Fields: createJobFromTemplateRequest: A CreateJobFromTemplateRequest resource to be passed as the request body. projectId: Required. The ID of the Cloud Platform project that the job belongs to. """ createJobFromTemplateRequest = _messages.MessageField('CreateJobFromTemplateRequest', 1) projectId = _messages.StringField(2, required=True) class DataflowProjectsTemplatesGetRequest(_messages.Message): r"""A DataflowProjectsTemplatesGetRequest object. Enums: ViewValueValuesEnum: The view to retrieve. Defaults to METADATA_ONLY. Fields: gcsPath: Required. A Cloud Storage path to the template from which to create the job. Must be valid Cloud Storage URL, beginning with 'gs://'. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) to which to direct the request. projectId: Required. The ID of the Cloud Platform project that the job belongs to. view: The view to retrieve. Defaults to METADATA_ONLY. """ class ViewValueValuesEnum(_messages.Enum): r"""The view to retrieve. Defaults to METADATA_ONLY. Values: METADATA_ONLY: Template view that retrieves only the metadata associated with the template. """ METADATA_ONLY = 0 gcsPath = _messages.StringField(1) location = _messages.StringField(2) projectId = _messages.StringField(3, required=True) view = _messages.EnumField('ViewValueValuesEnum', 4) class DataflowProjectsTemplatesLaunchRequest(_messages.Message): r"""A DataflowProjectsTemplatesLaunchRequest object. Fields: dynamicTemplate_gcsPath: Path to the dynamic template specification file on Cloud Storage. The file must be a JSON serialized `DynamicTemplateFileSpec` object. dynamicTemplate_stagingLocation: Cloud Storage path for staging dependencies. Must be a valid Cloud Storage URL, beginning with `gs://`. gcsPath: A Cloud Storage path to the template to use to create the job. Must be valid Cloud Storage URL, beginning with `gs://`. launchTemplateParameters: A LaunchTemplateParameters resource to be passed as the request body. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) to which to direct the request. projectId: Required. The ID of the Cloud Platform project that the job belongs to. validateOnly: If true, the request is validated but not actually executed. Defaults to false. """ dynamicTemplate_gcsPath = _messages.StringField(1) dynamicTemplate_stagingLocation = _messages.StringField(2) gcsPath = _messages.StringField(3) launchTemplateParameters = _messages.MessageField('LaunchTemplateParameters', 4) location = _messages.StringField(5) projectId = _messages.StringField(6, required=True) validateOnly = _messages.BooleanField(7) class DataflowProjectsWorkerMessagesRequest(_messages.Message): r"""A DataflowProjectsWorkerMessagesRequest object. Fields: projectId: The project to send the WorkerMessages to. sendWorkerMessagesRequest: A SendWorkerMessagesRequest resource to be passed as the request body. """ projectId = _messages.StringField(1, required=True) sendWorkerMessagesRequest = _messages.MessageField('SendWorkerMessagesRequest', 2) class DatastoreIODetails(_messages.Message): r"""Metadata for a Datastore connector used by the job. Fields: namespace: Namespace used in the connection. projectId: ProjectId accessed in the connection. """ namespace = _messages.StringField(1) projectId = _messages.StringField(2) class DebugOptions(_messages.Message): r"""Describes any options that have an effect on the debugging of pipelines. Fields: dataSampling: Configuration options for sampling elements from a running pipeline. enableHotKeyLogging: Optional. When true, enables the logging of the literal hot key to the user's Cloud Logging. """ dataSampling = _messages.MessageField('DataSamplingConfig', 1) enableHotKeyLogging = _messages.BooleanField(2) class DeleteSnapshotResponse(_messages.Message): r"""Response from deleting a snapshot.""" class DerivedSource(_messages.Message): r"""Specification of one of the bundles produced as a result of splitting a Source (e.g. when executing a SourceSplitRequest, or when splitting an active task using WorkItemStatus.dynamic_source_split), relative to the source being split. Enums: DerivationModeValueValuesEnum: What source to base the produced source on (if any). Fields: derivationMode: What source to base the produced source on (if any). source: Specification of the source. """ class DerivationModeValueValuesEnum(_messages.Enum): r"""What source to base the produced source on (if any). Values: SOURCE_DERIVATION_MODE_UNKNOWN: The source derivation is unknown, or unspecified. SOURCE_DERIVATION_MODE_INDEPENDENT: Produce a completely independent Source with no base. SOURCE_DERIVATION_MODE_CHILD_OF_CURRENT: Produce a Source based on the Source being split. SOURCE_DERIVATION_MODE_SIBLING_OF_CURRENT: Produce a Source based on the base of the Source being split. """ SOURCE_DERIVATION_MODE_UNKNOWN = 0 SOURCE_DERIVATION_MODE_INDEPENDENT = 1 SOURCE_DERIVATION_MODE_CHILD_OF_CURRENT = 2 SOURCE_DERIVATION_MODE_SIBLING_OF_CURRENT = 3 derivationMode = _messages.EnumField('DerivationModeValueValuesEnum', 1) source = _messages.MessageField('Source', 2) class Disk(_messages.Message): r"""Describes the data disk used by a workflow job. Fields: diskType: Disk storage type, as defined by Google Compute Engine. This must be a disk type appropriate to the project and zone in which the workers will run. If unknown or unspecified, the service will attempt to choose a reasonable default. For example, the standard persistent disk type is a resource name typically ending in "pd-standard". If SSD persistent disks are available, the resource name typically ends with "pd-ssd". The actual valid values are defined the Google Compute Engine API, not by the Cloud Dataflow API; consult the Google Compute Engine documentation for more information about determining the set of available disk types for a particular project and zone. Google Compute Engine Disk types are local to a particular project in a particular zone, and so the resource name will typically look something like this: compute.googleapis.com/projects/project-id/zones/zone/diskTypes/pd- standard mountPoint: Directory in a VM where disk is mounted. sizeGb: Size of disk in GB. If zero or unspecified, the service will attempt to choose a reasonable default. """ diskType = _messages.StringField(1) mountPoint = _messages.StringField(2) sizeGb = _messages.IntegerField(3, variant=_messages.Variant.INT32) class DisplayData(_messages.Message): r"""Data provided with a pipeline or transform to provide descriptive info. Fields: boolValue: Contains value if the data is of a boolean type. durationValue: Contains value if the data is of duration type. floatValue: Contains value if the data is of float type. int64Value: Contains value if the data is of int64 type. javaClassValue: Contains value if the data is of java class type. key: The key identifying the display data. This is intended to be used as a label for the display data when viewed in a dax monitoring system. label: An optional label to display in a dax UI for the element. namespace: The namespace for the key. This is usually a class name or programming language namespace (i.e. python module) which defines the display data. This allows a dax monitoring system to specially handle the data and perform custom rendering. shortStrValue: A possible additional shorter value to display. For example a java_class_name_value of com.mypackage.MyDoFn will be stored with MyDoFn as the short_str_value and com.mypackage.MyDoFn as the java_class_name value. short_str_value can be displayed and java_class_name_value will be displayed as a tooltip. strValue: Contains value if the data is of string type. timestampValue: Contains value if the data is of timestamp type. url: An optional full URL. """ boolValue = _messages.BooleanField(1) durationValue = _messages.StringField(2) floatValue = _messages.FloatField(3, variant=_messages.Variant.FLOAT) int64Value = _messages.IntegerField(4) javaClassValue = _messages.StringField(5) key = _messages.StringField(6) label = _messages.StringField(7) namespace = _messages.StringField(8) shortStrValue = _messages.StringField(9) strValue = _messages.StringField(10) timestampValue = _messages.StringField(11) url = _messages.StringField(12) class DistributionUpdate(_messages.Message): r"""A metric value representing a distribution. Fields: count: The count of the number of elements present in the distribution. histogram: (Optional) Histogram of value counts for the distribution. max: The maximum value present in the distribution. min: The minimum value present in the distribution. sum: Use an int64 since we'd prefer the added precision. If overflow is a common problem we can detect it and use an additional int64 or a double. sumOfSquares: Use a double since the sum of squares is likely to overflow int64. """ count = _messages.MessageField('SplitInt64', 1) histogram = _messages.MessageField('Histogram', 2) max = _messages.MessageField('SplitInt64', 3) min = _messages.MessageField('SplitInt64', 4) sum = _messages.MessageField('SplitInt64', 5) sumOfSquares = _messages.FloatField(6) class DynamicSourceSplit(_messages.Message): r"""When a task splits using WorkItemStatus.dynamic_source_split, this message describes the two parts of the split relative to the description of the current task's input. Fields: primary: Primary part (continued to be processed by worker). Specified relative to the previously-current source. Becomes current. residual: Residual part (returned to the pool of work). Specified relative to the previously-current source. """ primary = _messages.MessageField('DerivedSource', 1) residual = _messages.MessageField('DerivedSource', 2) class Environment(_messages.Message): r"""Describes the environment in which a Dataflow Job runs. Enums: FlexResourceSchedulingGoalValueValuesEnum: Optional. Which Flexible Resource Scheduling mode to run in. ShuffleModeValueValuesEnum: Output only. The shuffle mode used for the job. StreamingModeValueValuesEnum: Optional. Specifies the Streaming Engine message processing guarantees. Reduces cost and latency but might result in duplicate messages committed to storage. Designed to run simple mapping streaming ETL jobs at the lowest cost. For example, Change Data Capture (CDC) to BigQuery is a canonical use case. For more information, see [Set the pipeline streaming mode](https://cloud.google.com/dataflow/docs/guides/streaming-modes). Messages: InternalExperimentsValue: Experimental settings. SdkPipelineOptionsValue: The Cloud Dataflow SDK pipeline options specified by the user. These options are passed through the service and are used to recreate the SDK pipeline options on the worker in a language agnostic and platform independent way. UserAgentValue: Optional. A description of the process that generated the request. VersionValue: A structure describing which components and their versions of the service are required in order to run the job. Fields: clusterManagerApiService: The type of cluster manager API to use. If unknown or unspecified, the service will attempt to choose a reasonable default. This should be in the form of the API service name, e.g. "compute.googleapis.com". dataset: Optional. The dataset for the current project where various workflow related tables are stored. The supported resource type is: Google BigQuery: bigquery.googleapis.com/{dataset} debugOptions: Optional. Any debugging options to be supplied to the job. experiments: The list of experiments to enable. This field should be used for SDK related experiments and not for service related experiments. The proper field for service related experiments is service_options. flexResourceSchedulingGoal: Optional. Which Flexible Resource Scheduling mode to run in. internalExperiments: Experimental settings. sdkPipelineOptions: The Cloud Dataflow SDK pipeline options specified by the user. These options are passed through the service and are used to recreate the SDK pipeline options on the worker in a language agnostic and platform independent way. serviceAccountEmail: Optional. Identity to run virtual machines as. Defaults to the default account. serviceKmsKeyName: Optional. If set, contains the Cloud KMS key identifier used to encrypt data at rest, AKA a Customer Managed Encryption Key (CMEK). Format: projects/PROJECT_ID/locations/LOCATION/keyRings/KEY_RING/cryptoKeys/KEY serviceOptions: Optional. The list of service options to enable. This field should be used for service related experiments only. These experiments, when graduating to GA, should be replaced by dedicated fields or become default (i.e. always on). shuffleMode: Output only. The shuffle mode used for the job. streamingMode: Optional. Specifies the Streaming Engine message processing guarantees. Reduces cost and latency but might result in duplicate messages committed to storage. Designed to run simple mapping streaming ETL jobs at the lowest cost. For example, Change Data Capture (CDC) to BigQuery is a canonical use case. For more information, see [Set the pipeline streaming mode](https://cloud.google.com/dataflow/docs/guides/streaming-modes). tempStoragePrefix: The prefix of the resources the system should use for temporary storage. The system will append the suffix "/temp-{JOBNAME} to this resource prefix, where {JOBNAME} is the value of the job_name field. The resulting bucket and object prefix is used as the prefix of the resources used to store temporary data needed during the job execution. NOTE: This will override the value in taskrunner_settings. The supported resource type is: Google Cloud Storage: storage.googleapis.com/{bucket}/{object} bucket.storage.googleapis.com/{object} usePublicIps: Optional. True when any worker pool that uses public IPs is present. useStreamingEngineResourceBasedBilling: Output only. Whether the job uses the Streaming Engine resource-based billing model. userAgent: Optional. A description of the process that generated the request. version: A structure describing which components and their versions of the service are required in order to run the job. workerPools: The worker pools. At least one "harness" worker pool must be specified in order for the job to have workers. workerRegion: Optional. The Compute Engine region (https://cloud.google.com/compute/docs/regions-zones/regions-zones) in which worker processing should occur, e.g. "us-west1". Mutually exclusive with worker_zone. If neither worker_region nor worker_zone is specified, default to the control plane's region. workerZone: Optional. The Compute Engine zone (https://cloud.google.com/compute/docs/regions-zones/regions-zones) in which worker processing should occur, e.g. "us-west1-a". Mutually exclusive with worker_region. If neither worker_region nor worker_zone is specified, a zone in the control plane's region is chosen based on available capacity. """ class FlexResourceSchedulingGoalValueValuesEnum(_messages.Enum): r"""Optional. Which Flexible Resource Scheduling mode to run in. Values: FLEXRS_UNSPECIFIED: Run in the default mode. FLEXRS_SPEED_OPTIMIZED: Optimize for lower execution time. FLEXRS_COST_OPTIMIZED: Optimize for lower cost. """ FLEXRS_UNSPECIFIED = 0 FLEXRS_SPEED_OPTIMIZED = 1 FLEXRS_COST_OPTIMIZED = 2 class ShuffleModeValueValuesEnum(_messages.Enum): r"""Output only. The shuffle mode used for the job. Values: SHUFFLE_MODE_UNSPECIFIED: Shuffle mode information is not available. VM_BASED: Shuffle is done on the worker VMs. SERVICE_BASED: Shuffle is done on the service side. """ SHUFFLE_MODE_UNSPECIFIED = 0 VM_BASED = 1 SERVICE_BASED = 2 class StreamingModeValueValuesEnum(_messages.Enum): r"""Optional. Specifies the Streaming Engine message processing guarantees. Reduces cost and latency but might result in duplicate messages committed to storage. Designed to run simple mapping streaming ETL jobs at the lowest cost. For example, Change Data Capture (CDC) to BigQuery is a canonical use case. For more information, see [Set the pipeline streaming mode](https://cloud.google.com/dataflow/docs/guides/streaming-modes). Values: STREAMING_MODE_UNSPECIFIED: Run in the default mode. STREAMING_MODE_EXACTLY_ONCE: In this mode, message deduplication is performed against persistent state to make sure each message is processed and committed to storage exactly once. STREAMING_MODE_AT_LEAST_ONCE: Message deduplication is not performed. Messages might be processed multiple times, and the results are applied multiple times. Note: Setting this value also enables Streaming Engine and Streaming Engine resource-based billing. """ STREAMING_MODE_UNSPECIFIED = 0 STREAMING_MODE_EXACTLY_ONCE = 1 STREAMING_MODE_AT_LEAST_ONCE = 2 @encoding.MapUnrecognizedFields('additionalProperties') class InternalExperimentsValue(_messages.Message): r"""Experimental settings. Messages: AdditionalProperty: An additional property for a InternalExperimentsValue object. Fields: additionalProperties: Properties of the object. Contains field @type with type URL. """ class AdditionalProperty(_messages.Message): r"""An additional property for a InternalExperimentsValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class SdkPipelineOptionsValue(_messages.Message): r"""The Cloud Dataflow SDK pipeline options specified by the user. These options are passed through the service and are used to recreate the SDK pipeline options on the worker in a language agnostic and platform independent way. Messages: AdditionalProperty: An additional property for a SdkPipelineOptionsValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a SdkPipelineOptionsValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class UserAgentValue(_messages.Message): r"""Optional. A description of the process that generated the request. Messages: AdditionalProperty: An additional property for a UserAgentValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a UserAgentValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class VersionValue(_messages.Message): r"""A structure describing which components and their versions of the service are required in order to run the job. Messages: AdditionalProperty: An additional property for a VersionValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a VersionValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) clusterManagerApiService = _messages.StringField(1) dataset = _messages.StringField(2) debugOptions = _messages.MessageField('DebugOptions', 3) experiments = _messages.StringField(4, repeated=True) flexResourceSchedulingGoal = _messages.EnumField('FlexResourceSchedulingGoalValueValuesEnum', 5) internalExperiments = _messages.MessageField('InternalExperimentsValue', 6) sdkPipelineOptions = _messages.MessageField('SdkPipelineOptionsValue', 7) serviceAccountEmail = _messages.StringField(8) serviceKmsKeyName = _messages.StringField(9) serviceOptions = _messages.StringField(10, repeated=True) shuffleMode = _messages.EnumField('ShuffleModeValueValuesEnum', 11) streamingMode = _messages.EnumField('StreamingModeValueValuesEnum', 12) tempStoragePrefix = _messages.StringField(13) usePublicIps = _messages.BooleanField(14) useStreamingEngineResourceBasedBilling = _messages.BooleanField(15) userAgent = _messages.MessageField('UserAgentValue', 16) version = _messages.MessageField('VersionValue', 17) workerPools = _messages.MessageField('WorkerPool', 18, repeated=True) workerRegion = _messages.StringField(19) workerZone = _messages.StringField(20) class ExecutionStageState(_messages.Message): r"""A message describing the state of a particular execution stage. Enums: ExecutionStageStateValueValuesEnum: Executions stage states allow the same set of values as JobState. Fields: currentStateTime: The time at which the stage transitioned to this state. executionStageName: The name of the execution stage. executionStageState: Executions stage states allow the same set of values as JobState. """ class ExecutionStageStateValueValuesEnum(_messages.Enum): r"""Executions stage states allow the same set of values as JobState. Values: JOB_STATE_UNKNOWN: The job's run state isn't specified. JOB_STATE_STOPPED: `JOB_STATE_STOPPED` indicates that the job has not yet started to run. JOB_STATE_RUNNING: `JOB_STATE_RUNNING` indicates that the job is currently running. JOB_STATE_DONE: `JOB_STATE_DONE` indicates that the job has successfully completed. This is a terminal job state. This state may be set by the Cloud Dataflow service, as a transition from `JOB_STATE_RUNNING`. It may also be set via a Cloud Dataflow `UpdateJob` call, if the job has not yet reached a terminal state. JOB_STATE_FAILED: `JOB_STATE_FAILED` indicates that the job has failed. This is a terminal job state. This state may only be set by the Cloud Dataflow service, and only as a transition from `JOB_STATE_RUNNING`. JOB_STATE_CANCELLED: `JOB_STATE_CANCELLED` indicates that the job has been explicitly cancelled. This is a terminal job state. This state may only be set via a Cloud Dataflow `UpdateJob` call, and only if the job has not yet reached another terminal state. JOB_STATE_UPDATED: `JOB_STATE_UPDATED` indicates that the job was successfully updated, meaning that this job was stopped and another job was started, inheriting state from this one. This is a terminal job state. This state may only be set by the Cloud Dataflow service, and only as a transition from `JOB_STATE_RUNNING`. JOB_STATE_DRAINING: `JOB_STATE_DRAINING` indicates that the job is in the process of draining. A draining job has stopped pulling from its input sources and is processing any data that remains in-flight. This state may be set via a Cloud Dataflow `UpdateJob` call, but only as a transition from `JOB_STATE_RUNNING`. Jobs that are draining may only transition to `JOB_STATE_DRAINED`, `JOB_STATE_CANCELLED`, or `JOB_STATE_FAILED`. JOB_STATE_DRAINED: `JOB_STATE_DRAINED` indicates that the job has been drained. A drained job terminated by stopping pulling from its input sources and processing any data that remained in-flight when draining was requested. This state is a terminal state, may only be set by the Cloud Dataflow service, and only as a transition from `JOB_STATE_DRAINING`. JOB_STATE_PENDING: `JOB_STATE_PENDING` indicates that the job has been created but is not yet running. Jobs that are pending may only transition to `JOB_STATE_RUNNING`, or `JOB_STATE_FAILED`. JOB_STATE_CANCELLING: `JOB_STATE_CANCELLING` indicates that the job has been explicitly cancelled and is in the process of stopping. Jobs that are cancelling may only transition to `JOB_STATE_CANCELLED` or `JOB_STATE_FAILED`. JOB_STATE_QUEUED: `JOB_STATE_QUEUED` indicates that the job has been created but is being delayed until launch. Jobs that are queued may only transition to `JOB_STATE_PENDING` or `JOB_STATE_CANCELLED`. JOB_STATE_RESOURCE_CLEANING_UP: `JOB_STATE_RESOURCE_CLEANING_UP` indicates that the batch job's associated resources are currently being cleaned up after a successful run. Currently, this is an opt-in feature, please reach out to Cloud support team if you are interested. JOB_STATE_PAUSING: `JOB_STATE_PAUSING` is not implemented yet. JOB_STATE_PAUSED: `JOB_STATE_PAUSED` is not implemented yet. """ JOB_STATE_UNKNOWN = 0 JOB_STATE_STOPPED = 1 JOB_STATE_RUNNING = 2 JOB_STATE_DONE = 3 JOB_STATE_FAILED = 4 JOB_STATE_CANCELLED = 5 JOB_STATE_UPDATED = 6 JOB_STATE_DRAINING = 7 JOB_STATE_DRAINED = 8 JOB_STATE_PENDING = 9 JOB_STATE_CANCELLING = 10 JOB_STATE_QUEUED = 11 JOB_STATE_RESOURCE_CLEANING_UP = 12 JOB_STATE_PAUSING = 13 JOB_STATE_PAUSED = 14 currentStateTime = _messages.StringField(1) executionStageName = _messages.StringField(2) executionStageState = _messages.EnumField('ExecutionStageStateValueValuesEnum', 3) class ExecutionStageSummary(_messages.Message): r"""Description of the composing transforms, names/ids, and input/outputs of a stage of execution. Some composing transforms and sources may have been generated by the Dataflow service during execution planning. Enums: KindValueValuesEnum: Type of transform this stage is executing. Fields: componentSource: Collections produced and consumed by component transforms of this stage. componentTransform: Transforms that comprise this execution stage. id: Dataflow service generated id for this stage. inputSource: Input sources for this stage. kind: Type of transform this stage is executing. name: Dataflow service generated name for this stage. outputSource: Output sources for this stage. prerequisiteStage: Other stages that must complete before this stage can run. """ class KindValueValuesEnum(_messages.Enum): r"""Type of transform this stage is executing. Values: UNKNOWN_KIND: Unrecognized transform type. PAR_DO_KIND: ParDo transform. GROUP_BY_KEY_KIND: Group By Key transform. FLATTEN_KIND: Flatten transform. READ_KIND: Read transform. WRITE_KIND: Write transform. CONSTANT_KIND: Constructs from a constant value, such as with Create.of. SINGLETON_KIND: Creates a Singleton view of a collection. SHUFFLE_KIND: Opening or closing a shuffle session, often as part of a GroupByKey. """ UNKNOWN_KIND = 0 PAR_DO_KIND = 1 GROUP_BY_KEY_KIND = 2 FLATTEN_KIND = 3 READ_KIND = 4 WRITE_KIND = 5 CONSTANT_KIND = 6 SINGLETON_KIND = 7 SHUFFLE_KIND = 8 componentSource = _messages.MessageField('ComponentSource', 1, repeated=True) componentTransform = _messages.MessageField('ComponentTransform', 2, repeated=True) id = _messages.StringField(3) inputSource = _messages.MessageField('StageSource', 4, repeated=True) kind = _messages.EnumField('KindValueValuesEnum', 5) name = _messages.StringField(6) outputSource = _messages.MessageField('StageSource', 7, repeated=True) prerequisiteStage = _messages.StringField(8, repeated=True) class FailedLocation(_messages.Message): r"""Indicates which [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) failed to respond to a request for data. Fields: name: The name of the [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that failed to respond. """ name = _messages.StringField(1) class FileIODetails(_messages.Message): r"""Metadata for a File connector used by the job. Fields: filePattern: File Pattern used to access files by the connector. """ filePattern = _messages.StringField(1) class FlattenInstruction(_messages.Message): r"""An instruction that copies its inputs (zero or more) to its (single) output. Fields: inputs: Describes the inputs to the flatten instruction. """ inputs = _messages.MessageField('InstructionInput', 1, repeated=True) class FlexTemplateRuntimeEnvironment(_messages.Message): r"""The environment values to be set at runtime for flex template. Enums: AutoscalingAlgorithmValueValuesEnum: The algorithm to use for autoscaling FlexrsGoalValueValuesEnum: Set FlexRS goal for the job. https://cloud.google.com/dataflow/docs/guides/flexrs IpConfigurationValueValuesEnum: Configuration for VM IPs. StreamingModeValueValuesEnum: Optional. Specifies the Streaming Engine message processing guarantees. Reduces cost and latency but might result in duplicate messages committed to storage. Designed to run simple mapping streaming ETL jobs at the lowest cost. For example, Change Data Capture (CDC) to BigQuery is a canonical use case. For more information, see [Set the pipeline streaming mode](https://cloud.google.com/dataflow/docs/guides/streaming-modes). Messages: AdditionalUserLabelsValue: Additional user labels to be specified for the job. Keys and values must follow the restrictions specified in the [labeling restrictions](https://cloud.google.com/compute/docs/labeling- resources#restrictions) page. An object containing a list of "key": value pairs. Example: { "name": "wrench", "mass": "1kg", "count": "3" }. Fields: additionalExperiments: Additional experiment flags for the job. additionalPipelineOptions: Optional. Additional pipeline option flags for the job. additionalUserLabels: Additional user labels to be specified for the job. Keys and values must follow the restrictions specified in the [labeling restrictions](https://cloud.google.com/compute/docs/labeling- resources#restrictions) page. An object containing a list of "key": value pairs. Example: { "name": "wrench", "mass": "1kg", "count": "3" }. autoscalingAlgorithm: The algorithm to use for autoscaling diskSizeGb: Worker disk size, in gigabytes. dumpHeapOnOom: If true, when processing time is spent almost entirely on garbage collection (GC), saves a heap dump before ending the thread or process. If false, ends the thread or process without saving a heap dump. Does not save a heap dump when the Java Virtual Machine (JVM) has an out of memory error during processing. The location of the heap file is either echoed back to the user, or the user is given the opportunity to download the heap file. enableLauncherVmSerialPortLogging: If true serial port logging will be enabled for the launcher VM. enableStreamingEngine: Whether to enable Streaming Engine for the job. flexrsGoal: Set FlexRS goal for the job. https://cloud.google.com/dataflow/docs/guides/flexrs ipConfiguration: Configuration for VM IPs. kmsKeyName: Name for the Cloud KMS key for the job. Key format is: projects//locations//keyRings//cryptoKeys/ launcherMachineType: The machine type to use for launching the job. If not set, Dataflow will select a default machine type. machineType: The machine type to use for the job. Defaults to the value from the template if not specified. maxWorkers: The maximum number of Google Compute Engine instances to be made available to your pipeline during execution, from 1 to 1000. network: Network to which VMs will be assigned. If empty or unspecified, the service will use the network "default". numWorkers: The initial number of Google Compute Engine instances for the job. saveHeapDumpsToGcsPath: Cloud Storage bucket (directory) to upload heap dumps to. Enabling this field implies that `dump_heap_on_oom` is set to true. sdkContainerImage: Docker registry location of container image to use for the 'worker harness. Default is the container for the version of the SDK. Note this field is only valid for portable pipelines. serviceAccountEmail: The email address of the service account to run the job as. stagingLocation: The Cloud Storage path for staging local files. Must be a valid Cloud Storage URL, beginning with `gs://`. streamingMode: Optional. Specifies the Streaming Engine message processing guarantees. Reduces cost and latency but might result in duplicate messages committed to storage. Designed to run simple mapping streaming ETL jobs at the lowest cost. For example, Change Data Capture (CDC) to BigQuery is a canonical use case. For more information, see [Set the pipeline streaming mode](https://cloud.google.com/dataflow/docs/guides/streaming-modes). subnetwork: Subnetwork to which VMs will be assigned, if desired. You can specify a subnetwork using either a complete URL or an abbreviated path. Expected to be of the form "https://www.googleapis.com/compute/v1/projec ts/HOST_PROJECT_ID/regions/REGION/subnetworks/SUBNETWORK" or "regions/REGION/subnetworks/SUBNETWORK". If the subnetwork is located in a Shared VPC network, you must use the complete URL. tempLocation: The Cloud Storage path to use for temporary files. Must be a valid Cloud Storage URL, beginning with `gs://`. workerRegion: The Compute Engine region (https://cloud.google.com/compute/docs/regions-zones/regions-zones) in which worker processing should occur, e.g. "us-west1". Mutually exclusive with worker_zone. If neither worker_region nor worker_zone is specified, default to the control plane's region. workerZone: The Compute Engine zone (https://cloud.google.com/compute/docs/regions-zones/regions-zones) in which worker processing should occur, e.g. "us-west1-a". Mutually exclusive with worker_region. If neither worker_region nor worker_zone is specified, a zone in the control plane's region is chosen based on available capacity. If both `worker_zone` and `zone` are set, `worker_zone` takes precedence. zone: The Compute Engine [availability zone](https://cloud.google.com/compute/docs/regions-zones/regions-zones) for launching worker instances to run your pipeline. In the future, worker_zone will take precedence. """ class AutoscalingAlgorithmValueValuesEnum(_messages.Enum): r"""The algorithm to use for autoscaling Values: AUTOSCALING_ALGORITHM_UNKNOWN: The algorithm is unknown, or unspecified. AUTOSCALING_ALGORITHM_NONE: Disable autoscaling. AUTOSCALING_ALGORITHM_BASIC: Increase worker count over time to reduce job execution time. """ AUTOSCALING_ALGORITHM_UNKNOWN = 0 AUTOSCALING_ALGORITHM_NONE = 1 AUTOSCALING_ALGORITHM_BASIC = 2 class FlexrsGoalValueValuesEnum(_messages.Enum): r"""Set FlexRS goal for the job. https://cloud.google.com/dataflow/docs/guides/flexrs Values: FLEXRS_UNSPECIFIED: Run in the default mode. FLEXRS_SPEED_OPTIMIZED: Optimize for lower execution time. FLEXRS_COST_OPTIMIZED: Optimize for lower cost. """ FLEXRS_UNSPECIFIED = 0 FLEXRS_SPEED_OPTIMIZED = 1 FLEXRS_COST_OPTIMIZED = 2 class IpConfigurationValueValuesEnum(_messages.Enum): r"""Configuration for VM IPs. Values: WORKER_IP_UNSPECIFIED: The configuration is unknown, or unspecified. WORKER_IP_PUBLIC: Workers should have public IP addresses. WORKER_IP_PRIVATE: Workers should have private IP addresses. """ WORKER_IP_UNSPECIFIED = 0 WORKER_IP_PUBLIC = 1 WORKER_IP_PRIVATE = 2 class StreamingModeValueValuesEnum(_messages.Enum): r"""Optional. Specifies the Streaming Engine message processing guarantees. Reduces cost and latency but might result in duplicate messages committed to storage. Designed to run simple mapping streaming ETL jobs at the lowest cost. For example, Change Data Capture (CDC) to BigQuery is a canonical use case. For more information, see [Set the pipeline streaming mode](https://cloud.google.com/dataflow/docs/guides/streaming-modes). Values: STREAMING_MODE_UNSPECIFIED: Run in the default mode. STREAMING_MODE_EXACTLY_ONCE: In this mode, message deduplication is performed against persistent state to make sure each message is processed and committed to storage exactly once. STREAMING_MODE_AT_LEAST_ONCE: Message deduplication is not performed. Messages might be processed multiple times, and the results are applied multiple times. Note: Setting this value also enables Streaming Engine and Streaming Engine resource-based billing. """ STREAMING_MODE_UNSPECIFIED = 0 STREAMING_MODE_EXACTLY_ONCE = 1 STREAMING_MODE_AT_LEAST_ONCE = 2 @encoding.MapUnrecognizedFields('additionalProperties') class AdditionalUserLabelsValue(_messages.Message): r"""Additional user labels to be specified for the job. Keys and values must follow the restrictions specified in the [labeling restrictions](https://cloud.google.com/compute/docs/labeling- resources#restrictions) page. An object containing a list of "key": value pairs. Example: { "name": "wrench", "mass": "1kg", "count": "3" }. Messages: AdditionalProperty: An additional property for a AdditionalUserLabelsValue object. Fields: additionalProperties: Additional properties of type AdditionalUserLabelsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a AdditionalUserLabelsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) additionalExperiments = _messages.StringField(1, repeated=True) additionalPipelineOptions = _messages.StringField(2, repeated=True) additionalUserLabels = _messages.MessageField('AdditionalUserLabelsValue', 3) autoscalingAlgorithm = _messages.EnumField('AutoscalingAlgorithmValueValuesEnum', 4) diskSizeGb = _messages.IntegerField(5, variant=_messages.Variant.INT32) dumpHeapOnOom = _messages.BooleanField(6) enableLauncherVmSerialPortLogging = _messages.BooleanField(7) enableStreamingEngine = _messages.BooleanField(8) flexrsGoal = _messages.EnumField('FlexrsGoalValueValuesEnum', 9) ipConfiguration = _messages.EnumField('IpConfigurationValueValuesEnum', 10) kmsKeyName = _messages.StringField(11) launcherMachineType = _messages.StringField(12) machineType = _messages.StringField(13) maxWorkers = _messages.IntegerField(14, variant=_messages.Variant.INT32) network = _messages.StringField(15) numWorkers = _messages.IntegerField(16, variant=_messages.Variant.INT32) saveHeapDumpsToGcsPath = _messages.StringField(17) sdkContainerImage = _messages.StringField(18) serviceAccountEmail = _messages.StringField(19) stagingLocation = _messages.StringField(20) streamingMode = _messages.EnumField('StreamingModeValueValuesEnum', 21) subnetwork = _messages.StringField(22) tempLocation = _messages.StringField(23) workerRegion = _messages.StringField(24) workerZone = _messages.StringField(25) zone = _messages.StringField(26) class FloatingPointList(_messages.Message): r"""A metric value representing a list of floating point numbers. Fields: elements: Elements of the list. """ elements = _messages.FloatField(1, repeated=True) class FloatingPointMean(_messages.Message): r"""A representation of a floating point mean metric contribution. Fields: count: The number of values being aggregated. sum: The sum of all values being aggregated. """ count = _messages.MessageField('SplitInt64', 1) sum = _messages.FloatField(2) class GPUUsage(_messages.Message): r"""Information about the GPU usage on the worker. Fields: timestamp: Required. Timestamp of the measurement. utilization: Required. Utilization info about the GPU. """ timestamp = _messages.StringField(1) utilization = _messages.MessageField('GPUUtilization', 2) class GPUUtilization(_messages.Message): r"""Utilization details about the GPU. Fields: rate: Required. GPU utilization rate of any kernel over the last sample period in the range of [0, 1]. """ rate = _messages.FloatField(1) class GetDebugConfigRequest(_messages.Message): r"""Request to get updated debug configuration for component. Fields: componentId: The internal component id for which debug configuration is requested. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the job specified by job_id. workerId: The worker id, i.e., VM hostname. """ componentId = _messages.StringField(1) location = _messages.StringField(2) workerId = _messages.StringField(3) class GetDebugConfigResponse(_messages.Message): r"""Response to a get debug configuration request. Fields: config: The encoded debug configuration for the requested component. """ config = _messages.StringField(1) class GetTemplateResponse(_messages.Message): r"""The response to a GetTemplate request. Enums: TemplateTypeValueValuesEnum: Template Type. Fields: metadata: The template metadata describing the template name, available parameters, etc. runtimeMetadata: Describes the runtime metadata with SDKInfo and available parameters. status: The status of the get template request. Any problems with the request will be indicated in the error_details. templateType: Template Type. """ class TemplateTypeValueValuesEnum(_messages.Enum): r"""Template Type. Values: UNKNOWN: Unknown Template Type. LEGACY: Legacy Template. FLEX: Flex Template. """ UNKNOWN = 0 LEGACY = 1 FLEX = 2 metadata = _messages.MessageField('TemplateMetadata', 1) runtimeMetadata = _messages.MessageField('RuntimeMetadata', 2) status = _messages.MessageField('Status', 3) templateType = _messages.EnumField('TemplateTypeValueValuesEnum', 4) class GetWorkerStacktracesRequest(_messages.Message): r"""Request to get worker stacktraces from debug capture. Fields: endTime: The end time for the stacktrace query. The returned stacktraces will be a recent stack trace at or shortly before this time. workerId: The worker for which to get stacktraces. The returned stacktraces will be for the SDK harness running on this worker. """ endTime = _messages.StringField(1) workerId = _messages.StringField(2) class GetWorkerStacktracesResponse(_messages.Message): r"""Response to get worker stacktraces from debug capture. Fields: sdks: Repeated as unified worker may have multiple SDK processes. """ sdks = _messages.MessageField('Sdk', 1, repeated=True) class Histogram(_messages.Message): r"""Histogram of value counts for a distribution. Buckets have an inclusive lower bound and exclusive upper bound and use "1,2,5 bucketing": The first bucket range is from [0,1) and all subsequent bucket boundaries are powers of ten multiplied by 1, 2, or 5. Thus, bucket boundaries are 0, 1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, ... Negative values are not supported. Fields: bucketCounts: Counts of values in each bucket. For efficiency, prefix and trailing buckets with count = 0 are elided. Buckets can store the full range of values of an unsigned long, with ULLONG_MAX falling into the 59th bucket with range [1e19, 2e19). firstBucketOffset: Starting index of first stored bucket. The non- inclusive upper-bound of the ith bucket is given by: pow(10,(i-first_bucket_offset)/3) * (1,2,5)[(i-first_bucket_offset)%3] """ bucketCounts = _messages.IntegerField(1, repeated=True) firstBucketOffset = _messages.IntegerField(2, variant=_messages.Variant.INT32) class HotKeyDebuggingInfo(_messages.Message): r"""Information useful for debugging a hot key detection. Messages: DetectedHotKeysValue: Debugging information for each detected hot key. Keyed by a hash of the key. Fields: detectedHotKeys: Debugging information for each detected hot key. Keyed by a hash of the key. """ @encoding.MapUnrecognizedFields('additionalProperties') class DetectedHotKeysValue(_messages.Message): r"""Debugging information for each detected hot key. Keyed by a hash of the key. Messages: AdditionalProperty: An additional property for a DetectedHotKeysValue object. Fields: additionalProperties: Additional properties of type DetectedHotKeysValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a DetectedHotKeysValue object. Fields: key: Name of the additional property. value: A HotKeyInfo attribute. """ key = _messages.StringField(1) value = _messages.MessageField('HotKeyInfo', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) detectedHotKeys = _messages.MessageField('DetectedHotKeysValue', 1) class HotKeyDetection(_messages.Message): r"""Proto describing a hot key detected on a given WorkItem. Fields: hotKeyAge: The age of the hot key measured from when it was first detected. systemName: System-defined name of the step containing this hot key. Unique across the workflow. userStepName: User-provided name of the step that contains this hot key. """ hotKeyAge = _messages.StringField(1) systemName = _messages.StringField(2) userStepName = _messages.StringField(3) class HotKeyInfo(_messages.Message): r"""Information about a hot key. Fields: hotKeyAge: The age of the hot key measured from when it was first detected. key: A detected hot key that is causing limited parallelism. This field will be populated only if the following flag is set to true: "-- enable_hot_key_logging". keyTruncated: If true, then the above key is truncated and cannot be deserialized. This occurs if the key above is populated and the key size is >5MB. """ hotKeyAge = _messages.StringField(1) key = _messages.StringField(2) keyTruncated = _messages.BooleanField(3) class InstructionInput(_messages.Message): r"""An input of an instruction, as a reference to an output of a producer instruction. Fields: outputNum: The output index (origin zero) within the producer. producerInstructionIndex: The index (origin zero) of the parallel instruction that produces the output to be consumed by this input. This index is relative to the list of instructions in this input's instruction's containing MapTask. """ outputNum = _messages.IntegerField(1, variant=_messages.Variant.INT32) producerInstructionIndex = _messages.IntegerField(2, variant=_messages.Variant.INT32) class InstructionOutput(_messages.Message): r"""An output of an instruction. Messages: CodecValue: The codec to use to encode data being written via this output. Fields: codec: The codec to use to encode data being written via this output. name: The user-provided name of this output. onlyCountKeyBytes: For system-generated byte and mean byte metrics, certain instructions should only report the key size. onlyCountValueBytes: For system-generated byte and mean byte metrics, certain instructions should only report the value size. originalName: System-defined name for this output in the original workflow graph. Outputs that do not contribute to an original instruction do not set this. systemName: System-defined name of this output. Unique across the workflow. """ @encoding.MapUnrecognizedFields('additionalProperties') class CodecValue(_messages.Message): r"""The codec to use to encode data being written via this output. Messages: AdditionalProperty: An additional property for a CodecValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a CodecValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) codec = _messages.MessageField('CodecValue', 1) name = _messages.StringField(2) onlyCountKeyBytes = _messages.BooleanField(3) onlyCountValueBytes = _messages.BooleanField(4) originalName = _messages.StringField(5) systemName = _messages.StringField(6) class IntegerGauge(_messages.Message): r"""A metric value representing temporal values of a variable. Fields: timestamp: The time at which this value was measured. Measured as msecs from epoch. value: The value of the variable represented by this gauge. """ timestamp = _messages.StringField(1) value = _messages.MessageField('SplitInt64', 2) class IntegerList(_messages.Message): r"""A metric value representing a list of integers. Fields: elements: Elements of the list. """ elements = _messages.MessageField('SplitInt64', 1, repeated=True) class IntegerMean(_messages.Message): r"""A representation of an integer mean metric contribution. Fields: count: The number of values being aggregated. sum: The sum of all values being aggregated. """ count = _messages.MessageField('SplitInt64', 1) sum = _messages.MessageField('SplitInt64', 2) class Job(_messages.Message): r"""Defines a job to be run by the Cloud Dataflow service. Do not enter confidential information when you supply string values using the API. Enums: CurrentStateValueValuesEnum: The current state of the job. Jobs are created in the `JOB_STATE_STOPPED` state unless otherwise specified. A job in the `JOB_STATE_RUNNING` state may asynchronously enter a terminal state. After a job has reached a terminal state, no further state updates may be made. This field might be mutated by the Dataflow service; callers cannot mutate it. RequestedStateValueValuesEnum: The job's requested state. Applies to `UpdateJob` requests. Set `requested_state` with `UpdateJob` requests to switch between the states `JOB_STATE_STOPPED` and `JOB_STATE_RUNNING`. You can also use `UpdateJob` requests to change a job's state from `JOB_STATE_RUNNING` to `JOB_STATE_CANCELLED`, `JOB_STATE_DONE`, or `JOB_STATE_DRAINED`. These states irrevocably terminate the job if it hasn't already reached a terminal state. This field has no effect on `CreateJob` requests. TypeValueValuesEnum: Optional. The type of Dataflow job. Messages: LabelsValue: User-defined labels for this job. The labels map can contain no more than 64 entries. Entries of the labels map are UTF8 strings that comply with the following restrictions: * Keys must conform to regexp: \p{Ll}\p{Lo}{0,62} * Values must conform to regexp: [\p{Ll}\p{Lo}\p{N}_-]{0,63} * Both keys and values are additionally constrained to be <= 128 bytes in size. TransformNameMappingValue: Optional. The map of transform name prefixes of the job to be replaced to the corresponding name prefixes of the new job. Fields: clientRequestId: The client's unique identifier of the job, re-used across retried attempts. If this field is set, the service will ensure its uniqueness. The request to create a job will fail if the service has knowledge of a previously submitted job with the same client's ID and job name. The caller may use this field to ensure idempotence of job creation across retried attempts to create a job. By default, the field is empty and, in that case, the service ignores it. createTime: The timestamp when the job was initially created. Immutable and set by the Cloud Dataflow service. createdFromSnapshotId: If this is specified, the job's initial state is populated from the given snapshot. currentState: The current state of the job. Jobs are created in the `JOB_STATE_STOPPED` state unless otherwise specified. A job in the `JOB_STATE_RUNNING` state may asynchronously enter a terminal state. After a job has reached a terminal state, no further state updates may be made. This field might be mutated by the Dataflow service; callers cannot mutate it. currentStateTime: The timestamp associated with the current state. environment: Optional. The environment for the job. executionInfo: Deprecated. id: The unique ID of this job. This field is set by the Dataflow service when the job is created, and is immutable for the life of the job. jobMetadata: This field is populated by the Dataflow service to support filtering jobs by the metadata values provided here. Populated for ListJobs and all GetJob views SUMMARY and higher. labels: User-defined labels for this job. The labels map can contain no more than 64 entries. Entries of the labels map are UTF8 strings that comply with the following restrictions: * Keys must conform to regexp: \p{Ll}\p{Lo}{0,62} * Values must conform to regexp: [\p{Ll}\p{Lo}\p{N}_-]{0,63} * Both keys and values are additionally constrained to be <= 128 bytes in size. location: Optional. The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains this job. name: Optional. The user-specified Dataflow job name. Only one active job with a given name can exist in a project within one region at any given time. Jobs in different regions can have the same name. If a caller attempts to create a job with the same name as an active job that already exists, the attempt returns the existing job. The name must match the regular expression `[a-z]([-a-z0-9]{0,1022}[a-z0-9])?` pausable: Output only. Indicates whether the job can be paused. pipelineDescription: Preliminary field: The format of this data may change at any time. A description of the user pipeline and stages through which it is executed. Created by Cloud Dataflow service. Only retrieved with JOB_VIEW_DESCRIPTION or JOB_VIEW_ALL. projectId: The ID of the Google Cloud project that the job belongs to. replaceJobId: If this job is an update of an existing job, this field is the job ID of the job it replaced. When sending a `CreateJobRequest`, you can update a job by specifying it here. The job named here is stopped, and its intermediate state is transferred to this job. replacedByJobId: If another job is an update of this job (and thus, this job is in `JOB_STATE_UPDATED`), this field contains the ID of that job. requestedState: The job's requested state. Applies to `UpdateJob` requests. Set `requested_state` with `UpdateJob` requests to switch between the states `JOB_STATE_STOPPED` and `JOB_STATE_RUNNING`. You can also use `UpdateJob` requests to change a job's state from `JOB_STATE_RUNNING` to `JOB_STATE_CANCELLED`, `JOB_STATE_DONE`, or `JOB_STATE_DRAINED`. These states irrevocably terminate the job if it hasn't already reached a terminal state. This field has no effect on `CreateJob` requests. runtimeUpdatableParams: This field may ONLY be modified at runtime using the projects.jobs.update method to adjust job behavior. This field has no effect when specified at job creation. satisfiesPzi: Output only. Reserved for future use. This field is set only in responses from the server; it is ignored if it is set in any requests. satisfiesPzs: Reserved for future use. This field is set only in responses from the server; it is ignored if it is set in any requests. serviceResources: Output only. Resources used by the Dataflow Service to run the job. stageStates: This field may be mutated by the Cloud Dataflow service; callers cannot mutate it. startTime: The timestamp when the job was started (transitioned to JOB_STATE_PENDING). Flexible resource scheduling jobs are started with some delay after job creation, so start_time is unset before start and is updated when the job is started by the Cloud Dataflow service. For other jobs, start_time always equals to create_time and is immutable and set by the Cloud Dataflow service. steps: Exactly one of step or steps_location should be specified. The top- level steps that constitute the entire job. Only retrieved with JOB_VIEW_ALL. stepsLocation: The Cloud Storage location where the steps are stored. tempFiles: A set of files the system should be aware of that are used for temporary storage. These temporary files will be removed on job completion. No duplicates are allowed. No file patterns are supported. The supported files are: Google Cloud Storage: storage.googleapis.com/{bucket}/{object} bucket.storage.googleapis.com/{object} transformNameMapping: Optional. The map of transform name prefixes of the job to be replaced to the corresponding name prefixes of the new job. type: Optional. The type of Dataflow job. """ class CurrentStateValueValuesEnum(_messages.Enum): r"""The current state of the job. Jobs are created in the `JOB_STATE_STOPPED` state unless otherwise specified. A job in the `JOB_STATE_RUNNING` state may asynchronously enter a terminal state. After a job has reached a terminal state, no further state updates may be made. This field might be mutated by the Dataflow service; callers cannot mutate it. Values: JOB_STATE_UNKNOWN: The job's run state isn't specified. JOB_STATE_STOPPED: `JOB_STATE_STOPPED` indicates that the job has not yet started to run. JOB_STATE_RUNNING: `JOB_STATE_RUNNING` indicates that the job is currently running. JOB_STATE_DONE: `JOB_STATE_DONE` indicates that the job has successfully completed. This is a terminal job state. This state may be set by the Cloud Dataflow service, as a transition from `JOB_STATE_RUNNING`. It may also be set via a Cloud Dataflow `UpdateJob` call, if the job has not yet reached a terminal state. JOB_STATE_FAILED: `JOB_STATE_FAILED` indicates that the job has failed. This is a terminal job state. This state may only be set by the Cloud Dataflow service, and only as a transition from `JOB_STATE_RUNNING`. JOB_STATE_CANCELLED: `JOB_STATE_CANCELLED` indicates that the job has been explicitly cancelled. This is a terminal job state. This state may only be set via a Cloud Dataflow `UpdateJob` call, and only if the job has not yet reached another terminal state. JOB_STATE_UPDATED: `JOB_STATE_UPDATED` indicates that the job was successfully updated, meaning that this job was stopped and another job was started, inheriting state from this one. This is a terminal job state. This state may only be set by the Cloud Dataflow service, and only as a transition from `JOB_STATE_RUNNING`. JOB_STATE_DRAINING: `JOB_STATE_DRAINING` indicates that the job is in the process of draining. A draining job has stopped pulling from its input sources and is processing any data that remains in-flight. This state may be set via a Cloud Dataflow `UpdateJob` call, but only as a transition from `JOB_STATE_RUNNING`. Jobs that are draining may only transition to `JOB_STATE_DRAINED`, `JOB_STATE_CANCELLED`, or `JOB_STATE_FAILED`. JOB_STATE_DRAINED: `JOB_STATE_DRAINED` indicates that the job has been drained. A drained job terminated by stopping pulling from its input sources and processing any data that remained in-flight when draining was requested. This state is a terminal state, may only be set by the Cloud Dataflow service, and only as a transition from `JOB_STATE_DRAINING`. JOB_STATE_PENDING: `JOB_STATE_PENDING` indicates that the job has been created but is not yet running. Jobs that are pending may only transition to `JOB_STATE_RUNNING`, or `JOB_STATE_FAILED`. JOB_STATE_CANCELLING: `JOB_STATE_CANCELLING` indicates that the job has been explicitly cancelled and is in the process of stopping. Jobs that are cancelling may only transition to `JOB_STATE_CANCELLED` or `JOB_STATE_FAILED`. JOB_STATE_QUEUED: `JOB_STATE_QUEUED` indicates that the job has been created but is being delayed until launch. Jobs that are queued may only transition to `JOB_STATE_PENDING` or `JOB_STATE_CANCELLED`. JOB_STATE_RESOURCE_CLEANING_UP: `JOB_STATE_RESOURCE_CLEANING_UP` indicates that the batch job's associated resources are currently being cleaned up after a successful run. Currently, this is an opt-in feature, please reach out to Cloud support team if you are interested. JOB_STATE_PAUSING: `JOB_STATE_PAUSING` is not implemented yet. JOB_STATE_PAUSED: `JOB_STATE_PAUSED` is not implemented yet. """ JOB_STATE_UNKNOWN = 0 JOB_STATE_STOPPED = 1 JOB_STATE_RUNNING = 2 JOB_STATE_DONE = 3 JOB_STATE_FAILED = 4 JOB_STATE_CANCELLED = 5 JOB_STATE_UPDATED = 6 JOB_STATE_DRAINING = 7 JOB_STATE_DRAINED = 8 JOB_STATE_PENDING = 9 JOB_STATE_CANCELLING = 10 JOB_STATE_QUEUED = 11 JOB_STATE_RESOURCE_CLEANING_UP = 12 JOB_STATE_PAUSING = 13 JOB_STATE_PAUSED = 14 class RequestedStateValueValuesEnum(_messages.Enum): r"""The job's requested state. Applies to `UpdateJob` requests. Set `requested_state` with `UpdateJob` requests to switch between the states `JOB_STATE_STOPPED` and `JOB_STATE_RUNNING`. You can also use `UpdateJob` requests to change a job's state from `JOB_STATE_RUNNING` to `JOB_STATE_CANCELLED`, `JOB_STATE_DONE`, or `JOB_STATE_DRAINED`. These states irrevocably terminate the job if it hasn't already reached a terminal state. This field has no effect on `CreateJob` requests. Values: JOB_STATE_UNKNOWN: The job's run state isn't specified. JOB_STATE_STOPPED: `JOB_STATE_STOPPED` indicates that the job has not yet started to run. JOB_STATE_RUNNING: `JOB_STATE_RUNNING` indicates that the job is currently running. JOB_STATE_DONE: `JOB_STATE_DONE` indicates that the job has successfully completed. This is a terminal job state. This state may be set by the Cloud Dataflow service, as a transition from `JOB_STATE_RUNNING`. It may also be set via a Cloud Dataflow `UpdateJob` call, if the job has not yet reached a terminal state. JOB_STATE_FAILED: `JOB_STATE_FAILED` indicates that the job has failed. This is a terminal job state. This state may only be set by the Cloud Dataflow service, and only as a transition from `JOB_STATE_RUNNING`. JOB_STATE_CANCELLED: `JOB_STATE_CANCELLED` indicates that the job has been explicitly cancelled. This is a terminal job state. This state may only be set via a Cloud Dataflow `UpdateJob` call, and only if the job has not yet reached another terminal state. JOB_STATE_UPDATED: `JOB_STATE_UPDATED` indicates that the job was successfully updated, meaning that this job was stopped and another job was started, inheriting state from this one. This is a terminal job state. This state may only be set by the Cloud Dataflow service, and only as a transition from `JOB_STATE_RUNNING`. JOB_STATE_DRAINING: `JOB_STATE_DRAINING` indicates that the job is in the process of draining. A draining job has stopped pulling from its input sources and is processing any data that remains in-flight. This state may be set via a Cloud Dataflow `UpdateJob` call, but only as a transition from `JOB_STATE_RUNNING`. Jobs that are draining may only transition to `JOB_STATE_DRAINED`, `JOB_STATE_CANCELLED`, or `JOB_STATE_FAILED`. JOB_STATE_DRAINED: `JOB_STATE_DRAINED` indicates that the job has been drained. A drained job terminated by stopping pulling from its input sources and processing any data that remained in-flight when draining was requested. This state is a terminal state, may only be set by the Cloud Dataflow service, and only as a transition from `JOB_STATE_DRAINING`. JOB_STATE_PENDING: `JOB_STATE_PENDING` indicates that the job has been created but is not yet running. Jobs that are pending may only transition to `JOB_STATE_RUNNING`, or `JOB_STATE_FAILED`. JOB_STATE_CANCELLING: `JOB_STATE_CANCELLING` indicates that the job has been explicitly cancelled and is in the process of stopping. Jobs that are cancelling may only transition to `JOB_STATE_CANCELLED` or `JOB_STATE_FAILED`. JOB_STATE_QUEUED: `JOB_STATE_QUEUED` indicates that the job has been created but is being delayed until launch. Jobs that are queued may only transition to `JOB_STATE_PENDING` or `JOB_STATE_CANCELLED`. JOB_STATE_RESOURCE_CLEANING_UP: `JOB_STATE_RESOURCE_CLEANING_UP` indicates that the batch job's associated resources are currently being cleaned up after a successful run. Currently, this is an opt-in feature, please reach out to Cloud support team if you are interested. JOB_STATE_PAUSING: `JOB_STATE_PAUSING` is not implemented yet. JOB_STATE_PAUSED: `JOB_STATE_PAUSED` is not implemented yet. """ JOB_STATE_UNKNOWN = 0 JOB_STATE_STOPPED = 1 JOB_STATE_RUNNING = 2 JOB_STATE_DONE = 3 JOB_STATE_FAILED = 4 JOB_STATE_CANCELLED = 5 JOB_STATE_UPDATED = 6 JOB_STATE_DRAINING = 7 JOB_STATE_DRAINED = 8 JOB_STATE_PENDING = 9 JOB_STATE_CANCELLING = 10 JOB_STATE_QUEUED = 11 JOB_STATE_RESOURCE_CLEANING_UP = 12 JOB_STATE_PAUSING = 13 JOB_STATE_PAUSED = 14 class TypeValueValuesEnum(_messages.Enum): r"""Optional. The type of Dataflow job. Values: JOB_TYPE_UNKNOWN: The type of the job is unspecified, or unknown. JOB_TYPE_BATCH: A batch job with a well-defined end point: data is read, data is processed, data is written, and the job is done. JOB_TYPE_STREAMING: A continuously streaming job with no end: data is read, processed, and written continuously. """ JOB_TYPE_UNKNOWN = 0 JOB_TYPE_BATCH = 1 JOB_TYPE_STREAMING = 2 @encoding.MapUnrecognizedFields('additionalProperties') class LabelsValue(_messages.Message): r"""User-defined labels for this job. The labels map can contain no more than 64 entries. Entries of the labels map are UTF8 strings that comply with the following restrictions: * Keys must conform to regexp: \p{Ll}\p{Lo}{0,62} * Values must conform to regexp: [\p{Ll}\p{Lo}\p{N}_-]{0,63} * Both keys and values are additionally constrained to be <= 128 bytes in size. Messages: AdditionalProperty: An additional property for a LabelsValue object. Fields: additionalProperties: Additional properties of type LabelsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a LabelsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class TransformNameMappingValue(_messages.Message): r"""Optional. The map of transform name prefixes of the job to be replaced to the corresponding name prefixes of the new job. Messages: AdditionalProperty: An additional property for a TransformNameMappingValue object. Fields: additionalProperties: Additional properties of type TransformNameMappingValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a TransformNameMappingValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) clientRequestId = _messages.StringField(1) createTime = _messages.StringField(2) createdFromSnapshotId = _messages.StringField(3) currentState = _messages.EnumField('CurrentStateValueValuesEnum', 4) currentStateTime = _messages.StringField(5) environment = _messages.MessageField('Environment', 6) executionInfo = _messages.MessageField('JobExecutionInfo', 7) id = _messages.StringField(8) jobMetadata = _messages.MessageField('JobMetadata', 9) labels = _messages.MessageField('LabelsValue', 10) location = _messages.StringField(11) name = _messages.StringField(12) pausable = _messages.BooleanField(13) pipelineDescription = _messages.MessageField('PipelineDescription', 14) projectId = _messages.StringField(15) replaceJobId = _messages.StringField(16) replacedByJobId = _messages.StringField(17) requestedState = _messages.EnumField('RequestedStateValueValuesEnum', 18) runtimeUpdatableParams = _messages.MessageField('RuntimeUpdatableParams', 19) satisfiesPzi = _messages.BooleanField(20) satisfiesPzs = _messages.BooleanField(21) serviceResources = _messages.MessageField('ServiceResources', 22) stageStates = _messages.MessageField('ExecutionStageState', 23, repeated=True) startTime = _messages.StringField(24) steps = _messages.MessageField('Step', 25, repeated=True) stepsLocation = _messages.StringField(26) tempFiles = _messages.StringField(27, repeated=True) transformNameMapping = _messages.MessageField('TransformNameMappingValue', 28) type = _messages.EnumField('TypeValueValuesEnum', 29) class JobExecutionDetails(_messages.Message): r"""Information about the execution of a job. Fields: nextPageToken: If present, this response does not contain all requested tasks. To obtain the next page of results, repeat the request with page_token set to this value. stages: The stages of the job execution. """ nextPageToken = _messages.StringField(1) stages = _messages.MessageField('StageSummary', 2, repeated=True) class JobExecutionInfo(_messages.Message): r"""Additional information about how a Cloud Dataflow job will be executed that isn't contained in the submitted job. Messages: StagesValue: A mapping from each stage to the information about that stage. Fields: stages: A mapping from each stage to the information about that stage. """ @encoding.MapUnrecognizedFields('additionalProperties') class StagesValue(_messages.Message): r"""A mapping from each stage to the information about that stage. Messages: AdditionalProperty: An additional property for a StagesValue object. Fields: additionalProperties: Additional properties of type StagesValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a StagesValue object. Fields: key: Name of the additional property. value: A JobExecutionStageInfo attribute. """ key = _messages.StringField(1) value = _messages.MessageField('JobExecutionStageInfo', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) stages = _messages.MessageField('StagesValue', 1) class JobExecutionStageInfo(_messages.Message): r"""Contains information about how a particular google.dataflow.v1beta3.Step will be executed. Fields: stepName: The steps associated with the execution stage. Note that stages may have several steps, and that a given step might be run by more than one stage. """ stepName = _messages.StringField(1, repeated=True) class JobMessage(_messages.Message): r"""A particular message pertaining to a Dataflow job. Enums: MessageImportanceValueValuesEnum: Importance level of the message. Fields: id: Deprecated. messageImportance: Importance level of the message. messageText: The text of the message. time: The timestamp of the message. """ class MessageImportanceValueValuesEnum(_messages.Enum): r"""Importance level of the message. Values: JOB_MESSAGE_IMPORTANCE_UNKNOWN: The message importance isn't specified, or is unknown. JOB_MESSAGE_DEBUG: The message is at the 'debug' level: typically only useful for software engineers working on the code the job is running. Typically, Dataflow pipeline runners do not display log messages at this level by default. JOB_MESSAGE_DETAILED: The message is at the 'detailed' level: somewhat verbose, but potentially useful to users. Typically, Dataflow pipeline runners do not display log messages at this level by default. These messages are displayed by default in the Dataflow monitoring UI. JOB_MESSAGE_BASIC: The message is at the 'basic' level: useful for keeping track of the execution of a Dataflow pipeline. Typically, Dataflow pipeline runners display log messages at this level by default, and these messages are displayed by default in the Dataflow monitoring UI. JOB_MESSAGE_WARNING: The message is at the 'warning' level: indicating a condition pertaining to a job which may require human intervention. Typically, Dataflow pipeline runners display log messages at this level by default, and these messages are displayed by default in the Dataflow monitoring UI. JOB_MESSAGE_ERROR: The message is at the 'error' level: indicating a condition preventing a job from succeeding. Typically, Dataflow pipeline runners display log messages at this level by default, and these messages are displayed by default in the Dataflow monitoring UI. """ JOB_MESSAGE_IMPORTANCE_UNKNOWN = 0 JOB_MESSAGE_DEBUG = 1 JOB_MESSAGE_DETAILED = 2 JOB_MESSAGE_BASIC = 3 JOB_MESSAGE_WARNING = 4 JOB_MESSAGE_ERROR = 5 id = _messages.StringField(1) messageImportance = _messages.EnumField('MessageImportanceValueValuesEnum', 2) messageText = _messages.StringField(3) time = _messages.StringField(4) class JobMetadata(_messages.Message): r"""Metadata available primarily for filtering jobs. Will be included in the ListJob response and Job SUMMARY view. Messages: UserDisplayPropertiesValue: List of display properties to help UI filter jobs. Fields: bigTableDetails: Identification of a Cloud Bigtable source used in the Dataflow job. bigqueryDetails: Identification of a BigQuery source used in the Dataflow job. datastoreDetails: Identification of a Datastore source used in the Dataflow job. fileDetails: Identification of a File source used in the Dataflow job. pubsubDetails: Identification of a Pub/Sub source used in the Dataflow job. sdkVersion: The SDK version used to run the job. spannerDetails: Identification of a Spanner source used in the Dataflow job. userDisplayProperties: List of display properties to help UI filter jobs. """ @encoding.MapUnrecognizedFields('additionalProperties') class UserDisplayPropertiesValue(_messages.Message): r"""List of display properties to help UI filter jobs. Messages: AdditionalProperty: An additional property for a UserDisplayPropertiesValue object. Fields: additionalProperties: Additional properties of type UserDisplayPropertiesValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a UserDisplayPropertiesValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) bigTableDetails = _messages.MessageField('BigTableIODetails', 1, repeated=True) bigqueryDetails = _messages.MessageField('BigQueryIODetails', 2, repeated=True) datastoreDetails = _messages.MessageField('DatastoreIODetails', 3, repeated=True) fileDetails = _messages.MessageField('FileIODetails', 4, repeated=True) pubsubDetails = _messages.MessageField('PubSubIODetails', 5, repeated=True) sdkVersion = _messages.MessageField('SdkVersion', 6) spannerDetails = _messages.MessageField('SpannerIODetails', 7, repeated=True) userDisplayProperties = _messages.MessageField('UserDisplayPropertiesValue', 8) class JobMetrics(_messages.Message): r"""JobMetrics contains a collection of metrics describing the detailed progress of a Dataflow job. Metrics correspond to user-defined and system- defined metrics in the job. For more information, see [Dataflow job metrics] (https://cloud.google.com/dataflow/docs/guides/using-monitoring-intf). This resource captures only the most recent values of each metric; time-series data can be queried for them (under the same metric names) from Cloud Monitoring. Fields: metricTime: Timestamp as of which metric values are current. metrics: All metrics for this job. """ metricTime = _messages.StringField(1) metrics = _messages.MessageField('MetricUpdate', 2, repeated=True) class KeyRangeDataDiskAssignment(_messages.Message): r"""Data disk assignment information for a specific key-range of a sharded computation. Currently we only support UTF-8 character splits to simplify encoding into JSON. Fields: dataDisk: The name of the data disk where data for this range is stored. This name is local to the Google Cloud Platform project and uniquely identifies the disk within that project, for example "myproject-1014-104817-4c2-harness-0-disk-1". end: The end (exclusive) of the key range. start: The start (inclusive) of the key range. """ dataDisk = _messages.StringField(1) end = _messages.StringField(2) start = _messages.StringField(3) class KeyRangeLocation(_messages.Message): r"""Location information for a specific key-range of a sharded computation. Currently we only support UTF-8 character splits to simplify encoding into JSON. Fields: dataDisk: The name of the data disk where data for this range is stored. This name is local to the Google Cloud Platform project and uniquely identifies the disk within that project, for example "myproject-1014-104817-4c2-harness-0-disk-1". deliveryEndpoint: The physical location of this range assignment to be used for streaming computation cross-worker message delivery. deprecatedPersistentDirectory: DEPRECATED. The location of the persistent state for this range, as a persistent directory in the worker local filesystem. end: The end (exclusive) of the key range. start: The start (inclusive) of the key range. """ dataDisk = _messages.StringField(1) deliveryEndpoint = _messages.StringField(2) deprecatedPersistentDirectory = _messages.StringField(3) end = _messages.StringField(4) start = _messages.StringField(5) class LaunchFlexTemplateParameter(_messages.Message): r"""Launch FlexTemplate Parameter. Messages: LaunchOptionsValue: Launch options for this flex template job. This is a common set of options across languages and templates. This should not be used to pass job parameters. ParametersValue: The parameters for FlexTemplate. Ex. {"num_workers":"5"} TransformNameMappingsValue: Use this to pass transform_name_mappings for streaming update jobs. Ex:{"oldTransformName":"newTransformName",...}' Fields: containerSpec: Spec about the container image to launch. containerSpecGcsPath: Cloud Storage path to a file with json serialized ContainerSpec as content. environment: The runtime environment for the FlexTemplate job jobName: Required. The job name to use for the created job. For update job request, job name should be same as the existing running job. launchOptions: Launch options for this flex template job. This is a common set of options across languages and templates. This should not be used to pass job parameters. parameters: The parameters for FlexTemplate. Ex. {"num_workers":"5"} transformNameMappings: Use this to pass transform_name_mappings for streaming update jobs. Ex:{"oldTransformName":"newTransformName",...}' update: Set this to true if you are sending a request to update a running streaming job. When set, the job name should be the same as the running job. """ @encoding.MapUnrecognizedFields('additionalProperties') class LaunchOptionsValue(_messages.Message): r"""Launch options for this flex template job. This is a common set of options across languages and templates. This should not be used to pass job parameters. Messages: AdditionalProperty: An additional property for a LaunchOptionsValue object. Fields: additionalProperties: Additional properties of type LaunchOptionsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a LaunchOptionsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class ParametersValue(_messages.Message): r"""The parameters for FlexTemplate. Ex. {"num_workers":"5"} Messages: AdditionalProperty: An additional property for a ParametersValue object. Fields: additionalProperties: Additional properties of type ParametersValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a ParametersValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class TransformNameMappingsValue(_messages.Message): r"""Use this to pass transform_name_mappings for streaming update jobs. Ex:{"oldTransformName":"newTransformName",...}' Messages: AdditionalProperty: An additional property for a TransformNameMappingsValue object. Fields: additionalProperties: Additional properties of type TransformNameMappingsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a TransformNameMappingsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) containerSpec = _messages.MessageField('ContainerSpec', 1) containerSpecGcsPath = _messages.StringField(2) environment = _messages.MessageField('FlexTemplateRuntimeEnvironment', 3) jobName = _messages.StringField(4) launchOptions = _messages.MessageField('LaunchOptionsValue', 5) parameters = _messages.MessageField('ParametersValue', 6) transformNameMappings = _messages.MessageField('TransformNameMappingsValue', 7) update = _messages.BooleanField(8) class LaunchFlexTemplateRequest(_messages.Message): r"""A request to launch a Cloud Dataflow job from a FlexTemplate. Fields: launchParameter: Required. Parameter to launch a job form Flex Template. validateOnly: If true, the request is validated but not actually executed. Defaults to false. """ launchParameter = _messages.MessageField('LaunchFlexTemplateParameter', 1) validateOnly = _messages.BooleanField(2) class LaunchFlexTemplateResponse(_messages.Message): r"""Response to the request to launch a job from Flex Template. Fields: job: The job that was launched, if the request was not a dry run and the job was successfully launched. """ job = _messages.MessageField('Job', 1) class LaunchTemplateParameters(_messages.Message): r"""Parameters to provide to the template being launched. Note that the [metadata in the pipeline code] (https://cloud.google.com/dataflow/docs/guides/templates/creating- templates#metadata) determines which runtime parameters are valid. Messages: ParametersValue: The runtime parameters to pass to the job. TransformNameMappingValue: Only applicable when updating a pipeline. Map of transform name prefixes of the job to be replaced to the corresponding name prefixes of the new job. Fields: environment: The runtime environment for the job. jobName: Required. The job name to use for the created job. The name must match the regular expression `[a-z]([-a-z0-9]{0,1022}[a-z0-9])?` parameters: The runtime parameters to pass to the job. transformNameMapping: Only applicable when updating a pipeline. Map of transform name prefixes of the job to be replaced to the corresponding name prefixes of the new job. update: If set, replace the existing pipeline with the name specified by jobName with this pipeline, preserving state. """ @encoding.MapUnrecognizedFields('additionalProperties') class ParametersValue(_messages.Message): r"""The runtime parameters to pass to the job. Messages: AdditionalProperty: An additional property for a ParametersValue object. Fields: additionalProperties: Additional properties of type ParametersValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a ParametersValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class TransformNameMappingValue(_messages.Message): r"""Only applicable when updating a pipeline. Map of transform name prefixes of the job to be replaced to the corresponding name prefixes of the new job. Messages: AdditionalProperty: An additional property for a TransformNameMappingValue object. Fields: additionalProperties: Additional properties of type TransformNameMappingValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a TransformNameMappingValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) environment = _messages.MessageField('RuntimeEnvironment', 1) jobName = _messages.StringField(2) parameters = _messages.MessageField('ParametersValue', 3) transformNameMapping = _messages.MessageField('TransformNameMappingValue', 4) update = _messages.BooleanField(5) class LaunchTemplateResponse(_messages.Message): r"""Response to the request to launch a template. Fields: job: The job that was launched, if the request was not a dry run and the job was successfully launched. """ job = _messages.MessageField('Job', 1) class LeaseWorkItemRequest(_messages.Message): r"""Request to lease WorkItems. Messages: UnifiedWorkerRequestValue: Untranslated bag-of-bytes WorkRequest from UnifiedWorker. Fields: currentWorkerTime: The current timestamp at the worker. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the WorkItem's job. projectNumber: Optional. The project number of the project this worker belongs to. requestedLeaseDuration: The initial lease period. unifiedWorkerRequest: Untranslated bag-of-bytes WorkRequest from UnifiedWorker. workItemTypes: Filter for WorkItem type. workerCapabilities: Worker capabilities. WorkItems might be limited to workers with specific capabilities. workerId: Identifies the worker leasing work -- typically the ID of the virtual machine running the worker. """ @encoding.MapUnrecognizedFields('additionalProperties') class UnifiedWorkerRequestValue(_messages.Message): r"""Untranslated bag-of-bytes WorkRequest from UnifiedWorker. Messages: AdditionalProperty: An additional property for a UnifiedWorkerRequestValue object. Fields: additionalProperties: Properties of the object. Contains field @type with type URL. """ class AdditionalProperty(_messages.Message): r"""An additional property for a UnifiedWorkerRequestValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) currentWorkerTime = _messages.StringField(1) location = _messages.StringField(2) projectNumber = _messages.IntegerField(3) requestedLeaseDuration = _messages.StringField(4) unifiedWorkerRequest = _messages.MessageField('UnifiedWorkerRequestValue', 5) workItemTypes = _messages.StringField(6, repeated=True) workerCapabilities = _messages.StringField(7, repeated=True) workerId = _messages.StringField(8) class LeaseWorkItemResponse(_messages.Message): r"""Response to a request to lease WorkItems. Messages: UnifiedWorkerResponseValue: Untranslated bag-of-bytes WorkResponse for UnifiedWorker. Fields: unifiedWorkerResponse: Untranslated bag-of-bytes WorkResponse for UnifiedWorker. workItems: A list of the leased WorkItems. """ @encoding.MapUnrecognizedFields('additionalProperties') class UnifiedWorkerResponseValue(_messages.Message): r"""Untranslated bag-of-bytes WorkResponse for UnifiedWorker. Messages: AdditionalProperty: An additional property for a UnifiedWorkerResponseValue object. Fields: additionalProperties: Properties of the object. Contains field @type with type URL. """ class AdditionalProperty(_messages.Message): r"""An additional property for a UnifiedWorkerResponseValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) unifiedWorkerResponse = _messages.MessageField('UnifiedWorkerResponseValue', 1) workItems = _messages.MessageField('WorkItem', 2, repeated=True) class Linear(_messages.Message): r"""Linear buckets with the following boundaries for indices in 0 to n-1. - i in [0, n-1]: [start + (i)*width, start + (i+1)*width) Fields: numberOfBuckets: Must be greater than 0. start: Lower bound of the first bucket. width: Distance between bucket boundaries. Must be greater than 0. """ numberOfBuckets = _messages.IntegerField(1, variant=_messages.Variant.INT32) start = _messages.FloatField(2) width = _messages.FloatField(3) class ListJobMessagesResponse(_messages.Message): r"""Response to a request to list job messages. Fields: autoscalingEvents: Autoscaling events in ascending timestamp order. jobMessages: Messages in ascending timestamp order. nextPageToken: The token to obtain the next page of results if there are more. """ autoscalingEvents = _messages.MessageField('AutoscalingEvent', 1, repeated=True) jobMessages = _messages.MessageField('JobMessage', 2, repeated=True) nextPageToken = _messages.StringField(3) class ListJobsResponse(_messages.Message): r"""Response to a request to list Cloud Dataflow jobs in a project. This might be a partial response, depending on the page size in the ListJobsRequest. However, if the project does not have any jobs, an instance of ListJobsResponse is not returned and the requests's response body is empty {}. Fields: failedLocation: Zero or more messages describing the [regional endpoints] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that failed to respond. jobs: A subset of the requested job information. nextPageToken: Set if there may be more results than fit in this response. """ failedLocation = _messages.MessageField('FailedLocation', 1, repeated=True) jobs = _messages.MessageField('Job', 2, repeated=True) nextPageToken = _messages.StringField(3) class ListSnapshotsResponse(_messages.Message): r"""List of snapshots. Fields: snapshots: Returned snapshots. """ snapshots = _messages.MessageField('Snapshot', 1, repeated=True) class MapTask(_messages.Message): r"""MapTask consists of an ordered set of instructions, each of which describes one particular low-level operation for the worker to perform in order to accomplish the MapTask's WorkItem. Each instruction must appear in the list before any instructions which depends on its output. Fields: counterPrefix: Counter prefix that can be used to prefix counters. Not currently used in Dataflow. instructions: The instructions in the MapTask. stageName: System-defined name of the stage containing this MapTask. Unique across the workflow. systemName: System-defined name of this MapTask. Unique across the workflow. """ counterPrefix = _messages.StringField(1) instructions = _messages.MessageField('ParallelInstruction', 2, repeated=True) stageName = _messages.StringField(3) systemName = _messages.StringField(4) class MemInfo(_messages.Message): r"""Information about the memory usage of a worker or a container within a worker. Fields: currentLimitBytes: Instantenous memory limit in bytes. currentOoms: Number of Out of Memory (OOM) events recorded since the previous measurement. currentRssBytes: Instantenous memory (RSS) size in bytes. timestamp: Timestamp of the measurement. totalGbMs: Total memory (RSS) usage since start up in GB * ms. """ currentLimitBytes = _messages.IntegerField(1, variant=_messages.Variant.UINT64) currentOoms = _messages.IntegerField(2) currentRssBytes = _messages.IntegerField(3, variant=_messages.Variant.UINT64) timestamp = _messages.StringField(4) totalGbMs = _messages.IntegerField(5, variant=_messages.Variant.UINT64) class MetricShortId(_messages.Message): r"""The metric short id is returned to the user alongside an offset into ReportWorkItemStatusRequest Fields: metricIndex: The index of the corresponding metric in the ReportWorkItemStatusRequest. Required. shortId: The service-generated short identifier for the metric. """ metricIndex = _messages.IntegerField(1, variant=_messages.Variant.INT32) shortId = _messages.IntegerField(2) class MetricStructuredName(_messages.Message): r"""Identifies a metric, by describing the source which generated the metric. Messages: ContextValue: Zero or more labeled fields which identify the part of the job this metric is associated with, such as the name of a step or collection. For example, built-in counters associated with steps will have context['step'] = . Counters associated with PCollections in the SDK will have context['pcollection'] = . Fields: context: Zero or more labeled fields which identify the part of the job this metric is associated with, such as the name of a step or collection. For example, built-in counters associated with steps will have context['step'] = . Counters associated with PCollections in the SDK will have context['pcollection'] = . name: Worker-defined metric name. origin: Origin (namespace) of metric name. May be blank for user-define metrics; will be "dataflow" for metrics defined by the Dataflow service or SDK. """ @encoding.MapUnrecognizedFields('additionalProperties') class ContextValue(_messages.Message): r"""Zero or more labeled fields which identify the part of the job this metric is associated with, such as the name of a step or collection. For example, built-in counters associated with steps will have context['step'] = . Counters associated with PCollections in the SDK will have context['pcollection'] = . Messages: AdditionalProperty: An additional property for a ContextValue object. Fields: additionalProperties: Additional properties of type ContextValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a ContextValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) context = _messages.MessageField('ContextValue', 1) name = _messages.StringField(2) origin = _messages.StringField(3) class MetricUpdate(_messages.Message): r"""Describes the state of a metric. Fields: boundedTrie: Worker-computed aggregate value for the "Trie" aggregation kind. The only possible value type is a BoundedTrieNode. Introduced this field to avoid breaking older SDKs when Dataflow service starts to populate the `bounded_trie` field. cumulative: True if this metric is reported as the total cumulative aggregate value accumulated since the worker started working on this WorkItem. By default this is false, indicating that this metric is reported as a delta that is not associated with any WorkItem. distribution: A struct value describing properties of a distribution of numeric values. gauge: A struct value describing properties of a Gauge. Metrics of gauge type show the value of a metric across time, and is aggregated based on the newest value. internal: Worker-computed aggregate value for internal use by the Dataflow service. kind: Metric aggregation kind. The possible metric aggregation kinds are "Sum", "Max", "Min", "Mean", "Set", "And", "Or", and "Distribution". The specified aggregation kind is case-insensitive. If omitted, this is not an aggregated value but instead a single metric sample value. meanCount: Worker-computed aggregate value for the "Mean" aggregation kind. This holds the count of the aggregated values and is used in combination with mean_sum above to obtain the actual mean aggregate value. The only possible value type is Long. meanSum: Worker-computed aggregate value for the "Mean" aggregation kind. This holds the sum of the aggregated values and is used in combination with mean_count below to obtain the actual mean aggregate value. The only possible value types are Long and Double. name: Name of the metric. scalar: Worker-computed aggregate value for aggregation kinds "Sum", "Max", "Min", "And", and "Or". The possible value types are Long, Double, and Boolean. set: Worker-computed aggregate value for the "Set" aggregation kind. The only possible value type is a list of Values whose type can be Long, Double, String, or BoundedTrie according to the metric's type. All Values in the list must be of the same type. trie: Worker-computed aggregate value for the "Trie" aggregation kind. The only possible value type is a BoundedTrieNode. updateTime: Timestamp associated with the metric value. Optional when workers are reporting work progress; it will be filled in responses from the metrics API. """ boundedTrie = _messages.MessageField('extra_types.JsonValue', 1) cumulative = _messages.BooleanField(2) distribution = _messages.MessageField('extra_types.JsonValue', 3) gauge = _messages.MessageField('extra_types.JsonValue', 4) internal = _messages.MessageField('extra_types.JsonValue', 5) kind = _messages.StringField(6) meanCount = _messages.MessageField('extra_types.JsonValue', 7) meanSum = _messages.MessageField('extra_types.JsonValue', 8) name = _messages.MessageField('MetricStructuredName', 9) scalar = _messages.MessageField('extra_types.JsonValue', 10) set = _messages.MessageField('extra_types.JsonValue', 11) trie = _messages.MessageField('extra_types.JsonValue', 12) updateTime = _messages.StringField(13) class MetricValue(_messages.Message): r"""The value of a metric along with its name and labels. Messages: MetricLabelsValue: Optional. Set of metric labels for this metric. Fields: metric: Base name for this metric. metricLabels: Optional. Set of metric labels for this metric. valueGauge64: Non-cumulative int64 value of this metric. valueHistogram: Histogram value of this metric. valueInt64: Integer value of this metric. """ @encoding.MapUnrecognizedFields('additionalProperties') class MetricLabelsValue(_messages.Message): r"""Optional. Set of metric labels for this metric. Messages: AdditionalProperty: An additional property for a MetricLabelsValue object. Fields: additionalProperties: Additional properties of type MetricLabelsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a MetricLabelsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) metric = _messages.StringField(1) metricLabels = _messages.MessageField('MetricLabelsValue', 2) valueGauge64 = _messages.MessageField('DataflowGaugeValue', 3) valueHistogram = _messages.MessageField('DataflowHistogramValue', 4) valueInt64 = _messages.IntegerField(5) class MountedDataDisk(_messages.Message): r"""Describes mounted data disk. Fields: dataDisk: The name of the data disk. This name is local to the Google Cloud Platform project and uniquely identifies the disk within that project, for example "myproject-1014-104817-4c2-harness-0-disk-1". """ dataDisk = _messages.StringField(1) class MultiOutputInfo(_messages.Message): r"""Information about an output of a multi-output DoFn. Fields: tag: The id of the tag the user code will emit to this output by; this should correspond to the tag of some SideInputInfo. """ tag = _messages.StringField(1) class NameAndKind(_messages.Message): r"""Basic metadata about a counter. Enums: KindValueValuesEnum: Counter aggregation kind. Fields: kind: Counter aggregation kind. name: Name of the counter. """ class KindValueValuesEnum(_messages.Enum): r"""Counter aggregation kind. Values: INVALID: Counter aggregation kind was not set. SUM: Aggregated value is the sum of all contributed values. MAX: Aggregated value is the max of all contributed values. MIN: Aggregated value is the min of all contributed values. MEAN: Aggregated value is the mean of all contributed values. OR: Aggregated value represents the logical 'or' of all contributed values. AND: Aggregated value represents the logical 'and' of all contributed values. SET: Aggregated value is a set of unique contributed values. DISTRIBUTION: Aggregated value captures statistics about a distribution. LATEST_VALUE: Aggregated value tracks the latest value of a variable. """ INVALID = 0 SUM = 1 MAX = 2 MIN = 3 MEAN = 4 OR = 5 AND = 6 SET = 7 DISTRIBUTION = 8 LATEST_VALUE = 9 kind = _messages.EnumField('KindValueValuesEnum', 1) name = _messages.StringField(2) class OutlierStats(_messages.Message): r"""Statistics for the underflow and overflow bucket. Fields: overflowCount: Number of values that are larger than the upper bound of the largest bucket. overflowMean: Mean of values in the overflow bucket. underflowCount: Number of values that are smaller than the lower bound of the smallest bucket. underflowMean: Mean of values in the undeflow bucket. """ overflowCount = _messages.IntegerField(1) overflowMean = _messages.FloatField(2) underflowCount = _messages.IntegerField(3) underflowMean = _messages.FloatField(4) class Package(_messages.Message): r"""The packages that must be installed in order for a worker to run the steps of the Cloud Dataflow job that will be assigned to its worker pool. This is the mechanism by which the Cloud Dataflow SDK causes code to be loaded onto the workers. For example, the Cloud Dataflow Java SDK might use this to install jars containing the user's code and all of the various dependencies (libraries, data files, etc.) required in order for that code to run. Fields: location: The resource to read the package from. The supported resource type is: Google Cloud Storage: storage.googleapis.com/{bucket} bucket.storage.googleapis.com/ name: The name of the package. """ location = _messages.StringField(1) name = _messages.StringField(2) class ParDoInstruction(_messages.Message): r"""An instruction that does a ParDo operation. Takes one main input and zero or more side inputs, and produces zero or more outputs. Runs user code. Messages: UserFnValue: The user function to invoke. Fields: input: The input. multiOutputInfos: Information about each of the outputs, if user_fn is a MultiDoFn. numOutputs: The number of outputs. sideInputs: Zero or more side inputs. userFn: The user function to invoke. """ @encoding.MapUnrecognizedFields('additionalProperties') class UserFnValue(_messages.Message): r"""The user function to invoke. Messages: AdditionalProperty: An additional property for a UserFnValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a UserFnValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) input = _messages.MessageField('InstructionInput', 1) multiOutputInfos = _messages.MessageField('MultiOutputInfo', 2, repeated=True) numOutputs = _messages.IntegerField(3, variant=_messages.Variant.INT32) sideInputs = _messages.MessageField('SideInputInfo', 4, repeated=True) userFn = _messages.MessageField('UserFnValue', 5) class ParallelInstruction(_messages.Message): r"""Describes a particular operation comprising a MapTask. Fields: flatten: Additional information for Flatten instructions. name: User-provided name of this operation. originalName: System-defined name for the operation in the original workflow graph. outputs: Describes the outputs of the instruction. parDo: Additional information for ParDo instructions. partialGroupByKey: Additional information for PartialGroupByKey instructions. read: Additional information for Read instructions. systemName: System-defined name of this operation. Unique across the workflow. write: Additional information for Write instructions. """ flatten = _messages.MessageField('FlattenInstruction', 1) name = _messages.StringField(2) originalName = _messages.StringField(3) outputs = _messages.MessageField('InstructionOutput', 4, repeated=True) parDo = _messages.MessageField('ParDoInstruction', 5) partialGroupByKey = _messages.MessageField('PartialGroupByKeyInstruction', 6) read = _messages.MessageField('ReadInstruction', 7) systemName = _messages.StringField(8) write = _messages.MessageField('WriteInstruction', 9) class Parameter(_messages.Message): r"""Structured data associated with this message. Fields: key: Key or name for this parameter. value: Value for this parameter. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) class ParameterMetadata(_messages.Message): r"""Metadata for a specific parameter. Enums: ParamTypeValueValuesEnum: Optional. The type of the parameter. Used for selecting input picker. Messages: CustomMetadataValue: Optional. Additional metadata for describing this parameter. Fields: customMetadata: Optional. Additional metadata for describing this parameter. defaultValue: Optional. The default values will pre-populate the parameter with the given value from the proto. If default_value is left empty, the parameter will be populated with a default of the relevant type, e.g. false for a boolean. enumOptions: Optional. The options shown when ENUM ParameterType is specified. groupName: Optional. Specifies a group name for this parameter to be rendered under. Group header text will be rendered exactly as specified in this field. Only considered when parent_name is NOT provided. helpText: Required. The help text to display for the parameter. hiddenUi: Optional. Whether the parameter should be hidden in the UI. isOptional: Optional. Whether the parameter is optional. Defaults to false. label: Required. The label to display for the parameter. name: Required. The name of the parameter. paramType: Optional. The type of the parameter. Used for selecting input picker. parentName: Optional. Specifies the name of the parent parameter. Used in conjunction with 'parent_trigger_values' to make this parameter conditional (will only be rendered conditionally). Should be mappable to a ParameterMetadata.name field. parentTriggerValues: Optional. The value(s) of the 'parent_name' parameter which will trigger this parameter to be shown. If left empty, ANY non- empty value in parent_name will trigger this parameter to be shown. Only considered when this parameter is conditional (when 'parent_name' has been provided). regexes: Optional. Regexes that the parameter must match. """ class ParamTypeValueValuesEnum(_messages.Enum): r"""Optional. The type of the parameter. Used for selecting input picker. Values: DEFAULT: Default input type. TEXT: The parameter specifies generic text input. GCS_READ_BUCKET: The parameter specifies a Cloud Storage Bucket to read from. GCS_WRITE_BUCKET: The parameter specifies a Cloud Storage Bucket to write to. GCS_READ_FILE: The parameter specifies a Cloud Storage file path to read from. GCS_WRITE_FILE: The parameter specifies a Cloud Storage file path to write to. GCS_READ_FOLDER: The parameter specifies a Cloud Storage folder path to read from. GCS_WRITE_FOLDER: The parameter specifies a Cloud Storage folder to write to. PUBSUB_TOPIC: The parameter specifies a Pub/Sub Topic. PUBSUB_SUBSCRIPTION: The parameter specifies a Pub/Sub Subscription. BIGQUERY_TABLE: The parameter specifies a BigQuery table. JAVASCRIPT_UDF_FILE: The parameter specifies a JavaScript UDF in Cloud Storage. SERVICE_ACCOUNT: The parameter specifies a Service Account email. MACHINE_TYPE: The parameter specifies a Machine Type. KMS_KEY_NAME: The parameter specifies a KMS Key name. WORKER_REGION: The parameter specifies a Worker Region. WORKER_ZONE: The parameter specifies a Worker Zone. BOOLEAN: The parameter specifies a boolean input. ENUM: The parameter specifies an enum input. NUMBER: The parameter specifies a number input. KAFKA_TOPIC: Deprecated. Please use KAFKA_READ_TOPIC instead. KAFKA_READ_TOPIC: The parameter specifies the fully-qualified name of an Apache Kafka topic. This can be either a Google Managed Kafka topic or a non-managed Kafka topic. KAFKA_WRITE_TOPIC: The parameter specifies the fully-qualified name of an Apache Kafka topic. This can be an existing Google Managed Kafka topic, the name for a new Google Managed Kafka topic, or an existing non-managed Kafka topic. """ DEFAULT = 0 TEXT = 1 GCS_READ_BUCKET = 2 GCS_WRITE_BUCKET = 3 GCS_READ_FILE = 4 GCS_WRITE_FILE = 5 GCS_READ_FOLDER = 6 GCS_WRITE_FOLDER = 7 PUBSUB_TOPIC = 8 PUBSUB_SUBSCRIPTION = 9 BIGQUERY_TABLE = 10 JAVASCRIPT_UDF_FILE = 11 SERVICE_ACCOUNT = 12 MACHINE_TYPE = 13 KMS_KEY_NAME = 14 WORKER_REGION = 15 WORKER_ZONE = 16 BOOLEAN = 17 ENUM = 18 NUMBER = 19 KAFKA_TOPIC = 20 KAFKA_READ_TOPIC = 21 KAFKA_WRITE_TOPIC = 22 @encoding.MapUnrecognizedFields('additionalProperties') class CustomMetadataValue(_messages.Message): r"""Optional. Additional metadata for describing this parameter. Messages: AdditionalProperty: An additional property for a CustomMetadataValue object. Fields: additionalProperties: Additional properties of type CustomMetadataValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a CustomMetadataValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) customMetadata = _messages.MessageField('CustomMetadataValue', 1) defaultValue = _messages.StringField(2) enumOptions = _messages.MessageField('ParameterMetadataEnumOption', 3, repeated=True) groupName = _messages.StringField(4) helpText = _messages.StringField(5) hiddenUi = _messages.BooleanField(6) isOptional = _messages.BooleanField(7) label = _messages.StringField(8) name = _messages.StringField(9) paramType = _messages.EnumField('ParamTypeValueValuesEnum', 10) parentName = _messages.StringField(11) parentTriggerValues = _messages.StringField(12, repeated=True) regexes = _messages.StringField(13, repeated=True) class ParameterMetadataEnumOption(_messages.Message): r"""ParameterMetadataEnumOption specifies the option shown in the enum form. Fields: description: Optional. The description to display for the enum option. label: Optional. The label to display for the enum option. value: Required. The value of the enum option. """ description = _messages.StringField(1) label = _messages.StringField(2) value = _messages.StringField(3) class PartialGroupByKeyInstruction(_messages.Message): r"""An instruction that does a partial group-by-key. One input and one output. Messages: InputElementCodecValue: The codec to use for interpreting an element in the input PTable. ValueCombiningFnValue: The value combining function to invoke. Fields: input: Describes the input to the partial group-by-key instruction. inputElementCodec: The codec to use for interpreting an element in the input PTable. originalCombineValuesInputStoreName: If this instruction includes a combining function this is the name of the intermediate store between the GBK and the CombineValues. originalCombineValuesStepName: If this instruction includes a combining function, this is the name of the CombineValues instruction lifted into this instruction. sideInputs: Zero or more side inputs. valueCombiningFn: The value combining function to invoke. """ @encoding.MapUnrecognizedFields('additionalProperties') class InputElementCodecValue(_messages.Message): r"""The codec to use for interpreting an element in the input PTable. Messages: AdditionalProperty: An additional property for a InputElementCodecValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a InputElementCodecValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class ValueCombiningFnValue(_messages.Message): r"""The value combining function to invoke. Messages: AdditionalProperty: An additional property for a ValueCombiningFnValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a ValueCombiningFnValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) input = _messages.MessageField('InstructionInput', 1) inputElementCodec = _messages.MessageField('InputElementCodecValue', 2) originalCombineValuesInputStoreName = _messages.StringField(3) originalCombineValuesStepName = _messages.StringField(4) sideInputs = _messages.MessageField('SideInputInfo', 5, repeated=True) valueCombiningFn = _messages.MessageField('ValueCombiningFnValue', 6) class PerStepNamespaceMetrics(_messages.Message): r"""Metrics for a particular unfused step and namespace. A metric is uniquely identified by the `metrics_namespace`, `original_step`, `metric name` and `metric_labels`. Fields: metricValues: Optional. Metrics that are recorded for this namespace and unfused step. metricsNamespace: The namespace of these metrics on the worker. originalStep: The original system name of the unfused step that these metrics are reported from. """ metricValues = _messages.MessageField('MetricValue', 1, repeated=True) metricsNamespace = _messages.StringField(2) originalStep = _messages.StringField(3) class PerWorkerMetrics(_messages.Message): r"""Per worker metrics. Fields: perStepNamespaceMetrics: Optional. Metrics for a particular unfused step and namespace. """ perStepNamespaceMetrics = _messages.MessageField('PerStepNamespaceMetrics', 1, repeated=True) class PipelineDescription(_messages.Message): r"""A descriptive representation of submitted pipeline as well as the executed form. This data is provided by the Dataflow service for ease of visualizing the pipeline and interpreting Dataflow provided metrics. Fields: displayData: Pipeline level display data. executionPipelineStage: Description of each stage of execution of the pipeline. originalPipelineTransform: Description of each transform in the pipeline and collections between them. stepNamesHash: A hash value of the submitted pipeline portable graph step names if exists. """ displayData = _messages.MessageField('DisplayData', 1, repeated=True) executionPipelineStage = _messages.MessageField('ExecutionStageSummary', 2, repeated=True) originalPipelineTransform = _messages.MessageField('TransformSummary', 3, repeated=True) stepNamesHash = _messages.StringField(4) class Point(_messages.Message): r"""A point in the timeseries. Fields: time: The timestamp of the point. value: The value of the point. """ time = _messages.StringField(1) value = _messages.FloatField(2) class Position(_messages.Message): r"""Position defines a position within a collection of data. The value can be either the end position, a key (used with ordered collections), a byte offset, or a record index. Fields: byteOffset: Position is a byte offset. concatPosition: CloudPosition is a concat position. end: Position is past all other positions. Also useful for the end position of an unbounded range. key: Position is a string key, ordered lexicographically. recordIndex: Position is a record index. shufflePosition: CloudPosition is a base64 encoded BatchShufflePosition (with FIXED sharding). """ byteOffset = _messages.IntegerField(1) concatPosition = _messages.MessageField('ConcatPosition', 2) end = _messages.BooleanField(3) key = _messages.StringField(4) recordIndex = _messages.IntegerField(5) shufflePosition = _messages.StringField(6) class ProgressTimeseries(_messages.Message): r"""Information about the progress of some component of job execution. Fields: currentProgress: The current progress of the component, in the range [0,1]. dataPoints: History of progress for the component. Points are sorted by time. """ currentProgress = _messages.FloatField(1) dataPoints = _messages.MessageField('Point', 2, repeated=True) class PubSubIODetails(_messages.Message): r"""Metadata for a Pub/Sub connector used by the job. Fields: subscription: Subscription used in the connection. topic: Topic accessed in the connection. """ subscription = _messages.StringField(1) topic = _messages.StringField(2) class PubsubLocation(_messages.Message): r"""Identifies a pubsub location to use for transferring data into or out of a streaming Dataflow job. Fields: dropLateData: Indicates whether the pipeline allows late-arriving data. dynamicDestinations: If true, then this location represents dynamic topics. idLabel: If set, contains a pubsub label from which to extract record ids. If left empty, record deduplication will be strictly best effort. subscription: A pubsub subscription, in the form of "pubsub.googleapis.com/subscriptions//" timestampLabel: If set, contains a pubsub label from which to extract record timestamps. If left empty, record timestamps will be generated upon arrival. topic: A pubsub topic, in the form of "pubsub.googleapis.com/topics//" trackingSubscription: If set, specifies the pubsub subscription that will be used for tracking custom time timestamps for watermark estimation. withAttributes: If true, then the client has requested to get pubsub attributes. """ dropLateData = _messages.BooleanField(1) dynamicDestinations = _messages.BooleanField(2) idLabel = _messages.StringField(3) subscription = _messages.StringField(4) timestampLabel = _messages.StringField(5) topic = _messages.StringField(6) trackingSubscription = _messages.StringField(7) withAttributes = _messages.BooleanField(8) class PubsubSnapshotMetadata(_messages.Message): r"""Represents a Pubsub snapshot. Fields: expireTime: The expire time of the Pubsub snapshot. snapshotName: The name of the Pubsub snapshot. topicName: The name of the Pubsub topic. """ expireTime = _messages.StringField(1) snapshotName = _messages.StringField(2) topicName = _messages.StringField(3) class ReadInstruction(_messages.Message): r"""An instruction that reads records. Takes no inputs, produces one output. Fields: source: The source to read from. """ source = _messages.MessageField('Source', 1) class ReportWorkItemStatusRequest(_messages.Message): r"""Request to report the status of WorkItems. Messages: UnifiedWorkerRequestValue: Untranslated bag-of-bytes WorkProgressUpdateRequest from UnifiedWorker. Fields: currentWorkerTime: The current timestamp at the worker. location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the WorkItem's job. projectNumber: Optional. The project number of the project which owns the WorkItem's job. unifiedWorkerRequest: Untranslated bag-of-bytes WorkProgressUpdateRequest from UnifiedWorker. workItemStatuses: The order is unimportant, except that the order of the WorkItemServiceState messages in the ReportWorkItemStatusResponse corresponds to the order of WorkItemStatus messages here. workerId: The ID of the worker reporting the WorkItem status. If this does not match the ID of the worker which the Dataflow service believes currently has the lease on the WorkItem, the report will be dropped (with an error response). """ @encoding.MapUnrecognizedFields('additionalProperties') class UnifiedWorkerRequestValue(_messages.Message): r"""Untranslated bag-of-bytes WorkProgressUpdateRequest from UnifiedWorker. Messages: AdditionalProperty: An additional property for a UnifiedWorkerRequestValue object. Fields: additionalProperties: Properties of the object. Contains field @type with type URL. """ class AdditionalProperty(_messages.Message): r"""An additional property for a UnifiedWorkerRequestValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) currentWorkerTime = _messages.StringField(1) location = _messages.StringField(2) projectNumber = _messages.IntegerField(3) unifiedWorkerRequest = _messages.MessageField('UnifiedWorkerRequestValue', 4) workItemStatuses = _messages.MessageField('WorkItemStatus', 5, repeated=True) workerId = _messages.StringField(6) class ReportWorkItemStatusResponse(_messages.Message): r"""Response from a request to report the status of WorkItems. Messages: UnifiedWorkerResponseValue: Untranslated bag-of-bytes WorkProgressUpdateResponse for UnifiedWorker. Fields: unifiedWorkerResponse: Untranslated bag-of-bytes WorkProgressUpdateResponse for UnifiedWorker. workItemServiceStates: A set of messages indicating the service-side state for each WorkItem whose status was reported, in the same order as the WorkItemStatus messages in the ReportWorkItemStatusRequest which resulting in this response. """ @encoding.MapUnrecognizedFields('additionalProperties') class UnifiedWorkerResponseValue(_messages.Message): r"""Untranslated bag-of-bytes WorkProgressUpdateResponse for UnifiedWorker. Messages: AdditionalProperty: An additional property for a UnifiedWorkerResponseValue object. Fields: additionalProperties: Properties of the object. Contains field @type with type URL. """ class AdditionalProperty(_messages.Message): r"""An additional property for a UnifiedWorkerResponseValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) unifiedWorkerResponse = _messages.MessageField('UnifiedWorkerResponseValue', 1) workItemServiceStates = _messages.MessageField('WorkItemServiceState', 2, repeated=True) class ReportedParallelism(_messages.Message): r"""Represents the level of parallelism in a WorkItem's input, reported by the worker. Fields: isInfinite: Specifies whether the parallelism is infinite. If true, "value" is ignored. Infinite parallelism means the service will assume that the work item can always be split into more non-empty work items by dynamic splitting. This is a work-around for lack of support for infinity by the current JSON-based Java RPC stack. value: Specifies the level of parallelism in case it is finite. """ isInfinite = _messages.BooleanField(1) value = _messages.FloatField(2) class ResourceUtilizationReport(_messages.Message): r"""Worker metrics exported from workers. This contains resource utilization metrics accumulated from a variety of sources. For more information, see go/df-resource-signals. Messages: ContainersValue: Per container information. Key: container name. Fields: containers: Per container information. Key: container name. cpuTime: CPU utilization samples. gpuUsage: Optional. GPU usage samples. memoryInfo: Memory utilization samples. """ @encoding.MapUnrecognizedFields('additionalProperties') class ContainersValue(_messages.Message): r"""Per container information. Key: container name. Messages: AdditionalProperty: An additional property for a ContainersValue object. Fields: additionalProperties: Additional properties of type ContainersValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a ContainersValue object. Fields: key: Name of the additional property. value: A ResourceUtilizationReport attribute. """ key = _messages.StringField(1) value = _messages.MessageField('ResourceUtilizationReport', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) containers = _messages.MessageField('ContainersValue', 1) cpuTime = _messages.MessageField('CPUTime', 2, repeated=True) gpuUsage = _messages.MessageField('GPUUsage', 3, repeated=True) memoryInfo = _messages.MessageField('MemInfo', 4, repeated=True) class ResourceUtilizationReportResponse(_messages.Message): r"""Service-side response to WorkerMessage reporting resource utilization. """ class RuntimeEnvironment(_messages.Message): r"""The environment values to set at runtime. Enums: IpConfigurationValueValuesEnum: Optional. Configuration for VM IPs. StreamingModeValueValuesEnum: Optional. Specifies the Streaming Engine message processing guarantees. Reduces cost and latency but might result in duplicate messages committed to storage. Designed to run simple mapping streaming ETL jobs at the lowest cost. For example, Change Data Capture (CDC) to BigQuery is a canonical use case. For more information, see [Set the pipeline streaming mode](https://cloud.google.com/dataflow/docs/guides/streaming-modes). Messages: AdditionalUserLabelsValue: Optional. Additional user labels to be specified for the job. Keys and values should follow the restrictions specified in the [labeling restrictions](https://cloud.google.com/compute/docs/labeling- resources#restrictions) page. An object containing a list of "key": value pairs. Example: { "name": "wrench", "mass": "1kg", "count": "3" }. Fields: additionalExperiments: Optional. Additional experiment flags for the job, specified with the `--experiments` option. additionalPipelineOptions: Optional. Additional pipeline option flags for the job. additionalUserLabels: Optional. Additional user labels to be specified for the job. Keys and values should follow the restrictions specified in the [labeling restrictions](https://cloud.google.com/compute/docs/labeling- resources#restrictions) page. An object containing a list of "key": value pairs. Example: { "name": "wrench", "mass": "1kg", "count": "3" }. bypassTempDirValidation: Optional. Whether to bypass the safety checks for the job's temporary directory. Use with caution. diskSizeGb: Optional. The disk size, in gigabytes, to use on each remote Compute Engine worker instance. enableStreamingEngine: Optional. Whether to enable Streaming Engine for the job. ipConfiguration: Optional. Configuration for VM IPs. kmsKeyName: Optional. Name for the Cloud KMS key for the job. Key format is: projects//locations//keyRings//cryptoKeys/ machineType: Optional. The machine type to use for the job. Defaults to the value from the template if not specified. maxWorkers: Optional. The maximum number of Google Compute Engine instances to be made available to your pipeline during execution, from 1 to 1000. The default value is 1. network: Optional. Network to which VMs will be assigned. If empty or unspecified, the service will use the network "default". numWorkers: Optional. The initial number of Google Compute Engine instances for the job. The default value is 11. serviceAccountEmail: Optional. The email address of the service account to run the job as. streamingMode: Optional. Specifies the Streaming Engine message processing guarantees. Reduces cost and latency but might result in duplicate messages committed to storage. Designed to run simple mapping streaming ETL jobs at the lowest cost. For example, Change Data Capture (CDC) to BigQuery is a canonical use case. For more information, see [Set the pipeline streaming mode](https://cloud.google.com/dataflow/docs/guides/streaming-modes). subnetwork: Optional. Subnetwork to which VMs will be assigned, if desired. You can specify a subnetwork using either a complete URL or an abbreviated path. Expected to be of the form "https://www.googleapis.com /compute/v1/projects/HOST_PROJECT_ID/regions/REGION/subnetworks/SUBNETWO RK" or "regions/REGION/subnetworks/SUBNETWORK". If the subnetwork is located in a Shared VPC network, you must use the complete URL. tempLocation: Required. The Cloud Storage path to use for temporary files. Must be a valid Cloud Storage URL, beginning with `gs://`. workerRegion: Required. The Compute Engine region (https://cloud.google.com/compute/docs/regions-zones/regions-zones) in which worker processing should occur, e.g. "us-west1". Mutually exclusive with worker_zone. If neither worker_region nor worker_zone is specified, default to the control plane's region. workerZone: Optional. The Compute Engine zone (https://cloud.google.com/compute/docs/regions-zones/regions-zones) in which worker processing should occur, e.g. "us-west1-a". Mutually exclusive with worker_region. If neither worker_region nor worker_zone is specified, a zone in the control plane's region is chosen based on available capacity. If both `worker_zone` and `zone` are set, `worker_zone` takes precedence. zone: Optional. The Compute Engine [availability zone](https://cloud.google.com/compute/docs/regions-zones/regions-zones) for launching worker instances to run your pipeline. In the future, worker_zone will take precedence. """ class IpConfigurationValueValuesEnum(_messages.Enum): r"""Optional. Configuration for VM IPs. Values: WORKER_IP_UNSPECIFIED: The configuration is unknown, or unspecified. WORKER_IP_PUBLIC: Workers should have public IP addresses. WORKER_IP_PRIVATE: Workers should have private IP addresses. """ WORKER_IP_UNSPECIFIED = 0 WORKER_IP_PUBLIC = 1 WORKER_IP_PRIVATE = 2 class StreamingModeValueValuesEnum(_messages.Enum): r"""Optional. Specifies the Streaming Engine message processing guarantees. Reduces cost and latency but might result in duplicate messages committed to storage. Designed to run simple mapping streaming ETL jobs at the lowest cost. For example, Change Data Capture (CDC) to BigQuery is a canonical use case. For more information, see [Set the pipeline streaming mode](https://cloud.google.com/dataflow/docs/guides/streaming-modes). Values: STREAMING_MODE_UNSPECIFIED: Run in the default mode. STREAMING_MODE_EXACTLY_ONCE: In this mode, message deduplication is performed against persistent state to make sure each message is processed and committed to storage exactly once. STREAMING_MODE_AT_LEAST_ONCE: Message deduplication is not performed. Messages might be processed multiple times, and the results are applied multiple times. Note: Setting this value also enables Streaming Engine and Streaming Engine resource-based billing. """ STREAMING_MODE_UNSPECIFIED = 0 STREAMING_MODE_EXACTLY_ONCE = 1 STREAMING_MODE_AT_LEAST_ONCE = 2 @encoding.MapUnrecognizedFields('additionalProperties') class AdditionalUserLabelsValue(_messages.Message): r"""Optional. Additional user labels to be specified for the job. Keys and values should follow the restrictions specified in the [labeling restrictions](https://cloud.google.com/compute/docs/labeling- resources#restrictions) page. An object containing a list of "key": value pairs. Example: { "name": "wrench", "mass": "1kg", "count": "3" }. Messages: AdditionalProperty: An additional property for a AdditionalUserLabelsValue object. Fields: additionalProperties: Additional properties of type AdditionalUserLabelsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a AdditionalUserLabelsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) additionalExperiments = _messages.StringField(1, repeated=True) additionalPipelineOptions = _messages.StringField(2, repeated=True) additionalUserLabels = _messages.MessageField('AdditionalUserLabelsValue', 3) bypassTempDirValidation = _messages.BooleanField(4) diskSizeGb = _messages.IntegerField(5, variant=_messages.Variant.INT32) enableStreamingEngine = _messages.BooleanField(6) ipConfiguration = _messages.EnumField('IpConfigurationValueValuesEnum', 7) kmsKeyName = _messages.StringField(8) machineType = _messages.StringField(9) maxWorkers = _messages.IntegerField(10, variant=_messages.Variant.INT32) network = _messages.StringField(11) numWorkers = _messages.IntegerField(12, variant=_messages.Variant.INT32) serviceAccountEmail = _messages.StringField(13) streamingMode = _messages.EnumField('StreamingModeValueValuesEnum', 14) subnetwork = _messages.StringField(15) tempLocation = _messages.StringField(16) workerRegion = _messages.StringField(17) workerZone = _messages.StringField(18) zone = _messages.StringField(19) class RuntimeMetadata(_messages.Message): r"""RuntimeMetadata describing a runtime environment. Fields: parameters: The parameters for the template. sdkInfo: SDK Info for the template. """ parameters = _messages.MessageField('ParameterMetadata', 1, repeated=True) sdkInfo = _messages.MessageField('SDKInfo', 2) class RuntimeUpdatableParams(_messages.Message): r"""Additional job parameters that can only be updated during runtime using the projects.jobs.update method. These fields have no effect when specified during job creation. Fields: acceptableBacklogDuration: Optional. Deprecated: Use `latency_tier` instead. The backlog threshold duration in seconds for autoscaling. Value must be non-negative. autoscalingTier: Optional. Deprecated: Use `latency_tier` instead. The backlog threshold tier for autoscaling. Value must be one of "low- latency", "medium-latency", or "high-latency". latencyTier: Optional. The backlog threshold tier for autoscaling. Value must be one of "low-latency", "medium-latency", or "high-latency". maxNumWorkers: The maximum number of workers to cap autoscaling at. This field is currently only supported for Streaming Engine jobs. minNumWorkers: The minimum number of workers to scale down to. This field is currently only supported for Streaming Engine jobs. workerUtilizationHint: Target worker utilization, compared against the aggregate utilization of the worker pool by autoscaler, to determine upscaling and downscaling when absent other constraints such as backlog. For more information, see [Update an existing pipeline](https://cloud.google.com/dataflow/docs/guides/updating-a- pipeline). """ acceptableBacklogDuration = _messages.StringField(1) autoscalingTier = _messages.StringField(2) latencyTier = _messages.StringField(3) maxNumWorkers = _messages.IntegerField(4, variant=_messages.Variant.INT32) minNumWorkers = _messages.IntegerField(5, variant=_messages.Variant.INT32) workerUtilizationHint = _messages.FloatField(6) class SDKInfo(_messages.Message): r"""SDK Information. Enums: LanguageValueValuesEnum: Required. The SDK Language. Fields: language: Required. The SDK Language. version: Optional. The SDK version. """ class LanguageValueValuesEnum(_messages.Enum): r"""Required. The SDK Language. Values: UNKNOWN: UNKNOWN Language. JAVA: Java. PYTHON: Python. GO: Go. YAML: YAML. """ UNKNOWN = 0 JAVA = 1 PYTHON = 2 GO = 3 YAML = 4 language = _messages.EnumField('LanguageValueValuesEnum', 1) version = _messages.StringField(2) class Sdk(_messages.Message): r"""A structured representation of an SDK. Fields: sdkId: The SDK harness id. stacks: The stacktraces for the processes running on the SDK harness. """ sdkId = _messages.StringField(1) stacks = _messages.MessageField('Stack', 2, repeated=True) class SdkBug(_messages.Message): r"""A bug found in the Dataflow SDK. Enums: SeverityValueValuesEnum: Output only. How severe the SDK bug is. TypeValueValuesEnum: Output only. Describes the impact of this SDK bug. Fields: severity: Output only. How severe the SDK bug is. type: Output only. Describes the impact of this SDK bug. uri: Output only. Link to more information on the bug. """ class SeverityValueValuesEnum(_messages.Enum): r"""Output only. How severe the SDK bug is. Values: SEVERITY_UNSPECIFIED: A bug of unknown severity. NOTICE: A minor bug that that may reduce reliability or performance for some jobs. Impact will be minimal or non-existent for most jobs. WARNING: A bug that has some likelihood of causing performance degradation, data loss, or job failures. SEVERE: A bug with extremely significant impact. Jobs may fail erroneously, performance may be severely degraded, and data loss may be very likely. """ SEVERITY_UNSPECIFIED = 0 NOTICE = 1 WARNING = 2 SEVERE = 3 class TypeValueValuesEnum(_messages.Enum): r"""Output only. Describes the impact of this SDK bug. Values: TYPE_UNSPECIFIED: Unknown issue with this SDK. GENERAL: Catch-all for SDK bugs that don't fit in the below categories. PERFORMANCE: Using this version of the SDK may result in degraded performance. DATALOSS: Using this version of the SDK may cause data loss. """ TYPE_UNSPECIFIED = 0 GENERAL = 1 PERFORMANCE = 2 DATALOSS = 3 severity = _messages.EnumField('SeverityValueValuesEnum', 1) type = _messages.EnumField('TypeValueValuesEnum', 2) uri = _messages.StringField(3) class SdkHarnessContainerImage(_messages.Message): r"""Defines an SDK harness container for executing Dataflow pipelines. Fields: capabilities: The set of capabilities enumerated in the above Environment proto. See also [beam_runner_api.proto](https://github.com/apache/beam/b lob/master/model/pipeline/src/main/proto/org/apache/beam/model/pipeline/ v1/beam_runner_api.proto) containerImage: A docker container image that resides in Google Container Registry. environmentId: Environment ID for the Beam runner API proto Environment that corresponds to the current SDK Harness. useSingleCorePerContainer: If true, recommends the Dataflow service to use only one core per SDK container instance with this image. If false (or unset) recommends using more than one core per SDK container instance with this image for efficiency. Note that Dataflow service may choose to override this property if needed. """ capabilities = _messages.StringField(1, repeated=True) containerImage = _messages.StringField(2) environmentId = _messages.StringField(3) useSingleCorePerContainer = _messages.BooleanField(4) class SdkVersion(_messages.Message): r"""The version of the SDK used to run the job. Enums: SdkSupportStatusValueValuesEnum: The support status for this SDK version. Fields: bugs: Output only. Known bugs found in this SDK version. sdkSupportStatus: The support status for this SDK version. version: The version of the SDK used to run the job. versionDisplayName: A readable string describing the version of the SDK. """ class SdkSupportStatusValueValuesEnum(_messages.Enum): r"""The support status for this SDK version. Values: UNKNOWN: Cloud Dataflow is unaware of this version. SUPPORTED: This is a known version of an SDK, and is supported. STALE: A newer version of the SDK family exists, and an update is recommended. DEPRECATED: This version of the SDK is deprecated and will eventually be unsupported. UNSUPPORTED: Support for this SDK version has ended and it should no longer be used. """ UNKNOWN = 0 SUPPORTED = 1 STALE = 2 DEPRECATED = 3 UNSUPPORTED = 4 bugs = _messages.MessageField('SdkBug', 1, repeated=True) sdkSupportStatus = _messages.EnumField('SdkSupportStatusValueValuesEnum', 2) version = _messages.StringField(3) versionDisplayName = _messages.StringField(4) class SendDebugCaptureRequest(_messages.Message): r"""Request to send encoded debug information. Next ID: 8 Enums: DataFormatValueValuesEnum: Format for the data field above (id=5). Fields: componentId: The internal component id for which debug information is sent. data: The encoded debug information. dataFormat: Format for the data field above (id=5). location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the job specified by job_id. workerId: The worker id, i.e., VM hostname. """ class DataFormatValueValuesEnum(_messages.Enum): r"""Format for the data field above (id=5). Values: DATA_FORMAT_UNSPECIFIED: Format unspecified, parsing is determined based upon page type and legacy encoding. (go/protodosdonts#do-include-an- unspecified-value-in-an-enum) RAW: Raw HTML string. JSON: JSON-encoded string. ZLIB: Websafe encoded zlib-compressed string. BROTLI: Websafe encoded brotli-compressed string. """ DATA_FORMAT_UNSPECIFIED = 0 RAW = 1 JSON = 2 ZLIB = 3 BROTLI = 4 componentId = _messages.StringField(1) data = _messages.StringField(2) dataFormat = _messages.EnumField('DataFormatValueValuesEnum', 3) location = _messages.StringField(4) workerId = _messages.StringField(5) class SendDebugCaptureResponse(_messages.Message): r"""Response to a send capture request. nothing""" class SendWorkerMessagesRequest(_messages.Message): r"""A request for sending worker messages to the service. Fields: location: The [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints) that contains the job. workerMessages: The WorkerMessages to send. """ location = _messages.StringField(1) workerMessages = _messages.MessageField('WorkerMessage', 2, repeated=True) class SendWorkerMessagesResponse(_messages.Message): r"""The response to the worker messages. Fields: workerMessageResponses: The servers response to the worker messages. """ workerMessageResponses = _messages.MessageField('WorkerMessageResponse', 1, repeated=True) class SeqMapTask(_messages.Message): r"""Describes a particular function to invoke. Messages: UserFnValue: The user function to invoke. Fields: inputs: Information about each of the inputs. name: The user-provided name of the SeqDo operation. outputInfos: Information about each of the outputs. stageName: System-defined name of the stage containing the SeqDo operation. Unique across the workflow. systemName: System-defined name of the SeqDo operation. Unique across the workflow. userFn: The user function to invoke. """ @encoding.MapUnrecognizedFields('additionalProperties') class UserFnValue(_messages.Message): r"""The user function to invoke. Messages: AdditionalProperty: An additional property for a UserFnValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a UserFnValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) inputs = _messages.MessageField('SideInputInfo', 1, repeated=True) name = _messages.StringField(2) outputInfos = _messages.MessageField('SeqMapTaskOutputInfo', 3, repeated=True) stageName = _messages.StringField(4) systemName = _messages.StringField(5) userFn = _messages.MessageField('UserFnValue', 6) class SeqMapTaskOutputInfo(_messages.Message): r"""Information about an output of a SeqMapTask. Fields: sink: The sink to write the output value to. tag: The id of the TupleTag the user code will tag the output value by. """ sink = _messages.MessageField('Sink', 1) tag = _messages.StringField(2) class ServiceResources(_messages.Message): r"""Resources used by the Dataflow Service to run the job. Fields: zones: Output only. List of Cloud Zones being used by the Dataflow Service for this job. Example: us-central1-c """ zones = _messages.StringField(1, repeated=True) class ShellTask(_messages.Message): r"""A task which consists of a shell command for the worker to execute. Fields: command: The shell command to run. exitCode: Exit code for the task. """ command = _messages.StringField(1) exitCode = _messages.IntegerField(2, variant=_messages.Variant.INT32) class SideInputInfo(_messages.Message): r"""Information about a side input of a DoFn or an input of a SeqDoFn. Messages: KindValue: How to interpret the source element(s) as a side input value. Fields: kind: How to interpret the source element(s) as a side input value. sources: The source(s) to read element(s) from to get the value of this side input. If more than one source, then the elements are taken from the sources, in the specified order if order matters. At least one source is required. tag: The id of the tag the user code will access this side input by; this should correspond to the tag of some MultiOutputInfo. """ @encoding.MapUnrecognizedFields('additionalProperties') class KindValue(_messages.Message): r"""How to interpret the source element(s) as a side input value. Messages: AdditionalProperty: An additional property for a KindValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a KindValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) kind = _messages.MessageField('KindValue', 1) sources = _messages.MessageField('Source', 2, repeated=True) tag = _messages.StringField(3) class Sink(_messages.Message): r"""A sink that records can be encoded and written to. Messages: CodecValue: The codec to use to encode data written to the sink. SpecValue: The sink to write to, plus its parameters. Fields: codec: The codec to use to encode data written to the sink. spec: The sink to write to, plus its parameters. """ @encoding.MapUnrecognizedFields('additionalProperties') class CodecValue(_messages.Message): r"""The codec to use to encode data written to the sink. Messages: AdditionalProperty: An additional property for a CodecValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a CodecValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class SpecValue(_messages.Message): r"""The sink to write to, plus its parameters. Messages: AdditionalProperty: An additional property for a SpecValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a SpecValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) codec = _messages.MessageField('CodecValue', 1) spec = _messages.MessageField('SpecValue', 2) class Snapshot(_messages.Message): r"""Represents a snapshot of a job. Enums: StateValueValuesEnum: State of the snapshot. Fields: creationTime: The time this snapshot was created. description: User specified description of the snapshot. Maybe empty. diskSizeBytes: The disk byte size of the snapshot. Only available for snapshots in READY state. id: The unique ID of this snapshot. projectId: The project this snapshot belongs to. pubsubMetadata: Pub/Sub snapshot metadata. region: Cloud region where this snapshot lives in, e.g., "us-central1". sourceJobId: The job this snapshot was created from. state: State of the snapshot. ttl: The time after which this snapshot will be automatically deleted. """ class StateValueValuesEnum(_messages.Enum): r"""State of the snapshot. Values: UNKNOWN_SNAPSHOT_STATE: Unknown state. PENDING: Snapshot intent to create has been persisted, snapshotting of state has not yet started. RUNNING: Snapshotting is being performed. READY: Snapshot has been created and is ready to be used. FAILED: Snapshot failed to be created. DELETED: Snapshot has been deleted. """ UNKNOWN_SNAPSHOT_STATE = 0 PENDING = 1 RUNNING = 2 READY = 3 FAILED = 4 DELETED = 5 creationTime = _messages.StringField(1) description = _messages.StringField(2) diskSizeBytes = _messages.IntegerField(3) id = _messages.StringField(4) projectId = _messages.StringField(5) pubsubMetadata = _messages.MessageField('PubsubSnapshotMetadata', 6, repeated=True) region = _messages.StringField(7) sourceJobId = _messages.StringField(8) state = _messages.EnumField('StateValueValuesEnum', 9) ttl = _messages.StringField(10) class SnapshotJobRequest(_messages.Message): r"""Request to create a snapshot of a job. Fields: description: User specified description of the snapshot. Maybe empty. location: The location that contains this job. snapshotSources: If true, perform snapshots for sources which support this. ttl: TTL for the snapshot. """ description = _messages.StringField(1) location = _messages.StringField(2) snapshotSources = _messages.BooleanField(3) ttl = _messages.StringField(4) class Source(_messages.Message): r"""A source that records can be read and decoded from. Messages: BaseSpecsValueListEntry: A BaseSpecsValueListEntry object. CodecValue: The codec to use to decode data read from the source. SpecValue: The source to read from, plus its parameters. Fields: baseSpecs: While splitting, sources may specify the produced bundles as differences against another source, in order to save backend-side memory and allow bigger jobs. For details, see SourceSplitRequest. To support this use case, the full set of parameters of the source is logically obtained by taking the latest explicitly specified value of each parameter in the order: base_specs (later items win), spec (overrides anything in base_specs). codec: The codec to use to decode data read from the source. doesNotNeedSplitting: Setting this value to true hints to the framework that the source doesn't need splitting, and using SourceSplitRequest on it would yield SOURCE_SPLIT_OUTCOME_USE_CURRENT. E.g. a file splitter may set this to true when splitting a single file into a set of byte ranges of appropriate size, and set this to false when splitting a filepattern into individual files. However, for efficiency, a file splitter may decide to produce file subranges directly from the filepattern to avoid a splitting round-trip. See SourceSplitRequest for an overview of the splitting process. This field is meaningful only in the Source objects populated by the user (e.g. when filling in a DerivedSource). Source objects supplied by the framework to the user don't have this field populated. metadata: Optionally, metadata for this source can be supplied right away, avoiding a SourceGetMetadataOperation roundtrip (see SourceOperationRequest). This field is meaningful only in the Source objects populated by the user (e.g. when filling in a DerivedSource). Source objects supplied by the framework to the user don't have this field populated. spec: The source to read from, plus its parameters. """ @encoding.MapUnrecognizedFields('additionalProperties') class BaseSpecsValueListEntry(_messages.Message): r"""A BaseSpecsValueListEntry object. Messages: AdditionalProperty: An additional property for a BaseSpecsValueListEntry object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a BaseSpecsValueListEntry object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class CodecValue(_messages.Message): r"""The codec to use to decode data read from the source. Messages: AdditionalProperty: An additional property for a CodecValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a CodecValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class SpecValue(_messages.Message): r"""The source to read from, plus its parameters. Messages: AdditionalProperty: An additional property for a SpecValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a SpecValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) baseSpecs = _messages.MessageField('BaseSpecsValueListEntry', 1, repeated=True) codec = _messages.MessageField('CodecValue', 2) doesNotNeedSplitting = _messages.BooleanField(3) metadata = _messages.MessageField('SourceMetadata', 4) spec = _messages.MessageField('SpecValue', 5) class SourceFork(_messages.Message): r"""DEPRECATED in favor of DynamicSourceSplit. Fields: primary: DEPRECATED primarySource: DEPRECATED residual: DEPRECATED residualSource: DEPRECATED """ primary = _messages.MessageField('SourceSplitShard', 1) primarySource = _messages.MessageField('DerivedSource', 2) residual = _messages.MessageField('SourceSplitShard', 3) residualSource = _messages.MessageField('DerivedSource', 4) class SourceGetMetadataRequest(_messages.Message): r"""A request to compute the SourceMetadata of a Source. Fields: source: Specification of the source whose metadata should be computed. """ source = _messages.MessageField('Source', 1) class SourceGetMetadataResponse(_messages.Message): r"""The result of a SourceGetMetadataOperation. Fields: metadata: The computed metadata. """ metadata = _messages.MessageField('SourceMetadata', 1) class SourceMetadata(_messages.Message): r"""Metadata about a Source useful for automatically optimizing and tuning the pipeline, etc. Fields: estimatedSizeBytes: An estimate of the total size (in bytes) of the data that would be read from this source. This estimate is in terms of external storage size, before any decompression or other processing done by the reader. infinite: Specifies that the size of this source is known to be infinite (this is a streaming source). producesSortedKeys: Whether this source is known to produce key/value pairs with the (encoded) keys in lexicographically sorted order. """ estimatedSizeBytes = _messages.IntegerField(1) infinite = _messages.BooleanField(2) producesSortedKeys = _messages.BooleanField(3) class SourceOperationRequest(_messages.Message): r"""A work item that represents the different operations that can be performed on a user-defined Source specification. Fields: getMetadata: Information about a request to get metadata about a source. name: User-provided name of the Read instruction for this source. originalName: System-defined name for the Read instruction for this source in the original workflow graph. split: Information about a request to split a source. stageName: System-defined name of the stage containing the source operation. Unique across the workflow. systemName: System-defined name of the Read instruction for this source. Unique across the workflow. """ getMetadata = _messages.MessageField('SourceGetMetadataRequest', 1) name = _messages.StringField(2) originalName = _messages.StringField(3) split = _messages.MessageField('SourceSplitRequest', 4) stageName = _messages.StringField(5) systemName = _messages.StringField(6) class SourceOperationResponse(_messages.Message): r"""The result of a SourceOperationRequest, specified in ReportWorkItemStatusRequest.source_operation when the work item is completed. Fields: getMetadata: A response to a request to get metadata about a source. split: A response to a request to split a source. """ getMetadata = _messages.MessageField('SourceGetMetadataResponse', 1) split = _messages.MessageField('SourceSplitResponse', 2) class SourceSplitOptions(_messages.Message): r"""Hints for splitting a Source into bundles (parts for parallel processing) using SourceSplitRequest. Fields: desiredBundleSizeBytes: The source should be split into a set of bundles where the estimated size of each is approximately this many bytes. desiredShardSizeBytes: DEPRECATED in favor of desired_bundle_size_bytes. """ desiredBundleSizeBytes = _messages.IntegerField(1) desiredShardSizeBytes = _messages.IntegerField(2) class SourceSplitRequest(_messages.Message): r"""Represents the operation to split a high-level Source specification into bundles (parts for parallel processing). At a high level, splitting of a source into bundles happens as follows: SourceSplitRequest is applied to the source. If it returns SOURCE_SPLIT_OUTCOME_USE_CURRENT, no further splitting happens and the source is used "as is". Otherwise, splitting is applied recursively to each produced DerivedSource. As an optimization, for any Source, if its does_not_need_splitting is true, the framework assumes that splitting this source would return SOURCE_SPLIT_OUTCOME_USE_CURRENT, and doesn't initiate a SourceSplitRequest. This applies both to the initial source being split and to bundles produced from it. Fields: options: Hints for tuning the splitting process. source: Specification of the source to be split. """ options = _messages.MessageField('SourceSplitOptions', 1) source = _messages.MessageField('Source', 2) class SourceSplitResponse(_messages.Message): r"""The response to a SourceSplitRequest. Enums: OutcomeValueValuesEnum: Indicates whether splitting happened and produced a list of bundles. If this is USE_CURRENT_SOURCE_AS_IS, the current source should be processed "as is" without splitting. "bundles" is ignored in this case. If this is SPLITTING_HAPPENED, then "bundles" contains a list of bundles into which the source was split. Fields: bundles: If outcome is SPLITTING_HAPPENED, then this is a list of bundles into which the source was split. Otherwise this field is ignored. This list can be empty, which means the source represents an empty input. outcome: Indicates whether splitting happened and produced a list of bundles. If this is USE_CURRENT_SOURCE_AS_IS, the current source should be processed "as is" without splitting. "bundles" is ignored in this case. If this is SPLITTING_HAPPENED, then "bundles" contains a list of bundles into which the source was split. shards: DEPRECATED in favor of bundles. """ class OutcomeValueValuesEnum(_messages.Enum): r"""Indicates whether splitting happened and produced a list of bundles. If this is USE_CURRENT_SOURCE_AS_IS, the current source should be processed "as is" without splitting. "bundles" is ignored in this case. If this is SPLITTING_HAPPENED, then "bundles" contains a list of bundles into which the source was split. Values: SOURCE_SPLIT_OUTCOME_UNKNOWN: The source split outcome is unknown, or unspecified. SOURCE_SPLIT_OUTCOME_USE_CURRENT: The current source should be processed "as is" without splitting. SOURCE_SPLIT_OUTCOME_SPLITTING_HAPPENED: Splitting produced a list of bundles. """ SOURCE_SPLIT_OUTCOME_UNKNOWN = 0 SOURCE_SPLIT_OUTCOME_USE_CURRENT = 1 SOURCE_SPLIT_OUTCOME_SPLITTING_HAPPENED = 2 bundles = _messages.MessageField('DerivedSource', 1, repeated=True) outcome = _messages.EnumField('OutcomeValueValuesEnum', 2) shards = _messages.MessageField('SourceSplitShard', 3, repeated=True) class SourceSplitShard(_messages.Message): r"""DEPRECATED in favor of DerivedSource. Enums: DerivationModeValueValuesEnum: DEPRECATED Fields: derivationMode: DEPRECATED source: DEPRECATED """ class DerivationModeValueValuesEnum(_messages.Enum): r"""DEPRECATED Values: SOURCE_DERIVATION_MODE_UNKNOWN: The source derivation is unknown, or unspecified. SOURCE_DERIVATION_MODE_INDEPENDENT: Produce a completely independent Source with no base. SOURCE_DERIVATION_MODE_CHILD_OF_CURRENT: Produce a Source based on the Source being split. SOURCE_DERIVATION_MODE_SIBLING_OF_CURRENT: Produce a Source based on the base of the Source being split. """ SOURCE_DERIVATION_MODE_UNKNOWN = 0 SOURCE_DERIVATION_MODE_INDEPENDENT = 1 SOURCE_DERIVATION_MODE_CHILD_OF_CURRENT = 2 SOURCE_DERIVATION_MODE_SIBLING_OF_CURRENT = 3 derivationMode = _messages.EnumField('DerivationModeValueValuesEnum', 1) source = _messages.MessageField('Source', 2) class SpannerIODetails(_messages.Message): r"""Metadata for a Spanner connector used by the job. Fields: databaseId: DatabaseId accessed in the connection. instanceId: InstanceId accessed in the connection. projectId: ProjectId accessed in the connection. """ databaseId = _messages.StringField(1) instanceId = _messages.StringField(2) projectId = _messages.StringField(3) class SplitInt64(_messages.Message): r"""A representation of an int64, n, that is immune to precision loss when encoded in JSON. Fields: highBits: The high order bits, including the sign: n >> 32. lowBits: The low order bits: n & 0xffffffff. """ highBits = _messages.IntegerField(1, variant=_messages.Variant.INT32) lowBits = _messages.IntegerField(2, variant=_messages.Variant.UINT32) class Stack(_messages.Message): r"""A structuredstacktrace for a process running on the worker. Fields: stackContent: The raw stack trace. threadCount: With java thread dumps we may get collapsed stacks e.g., N threads in stack "". Instead of having to copy over the same stack trace N times, this int field captures this. threadName: Thread name. For example, "CommitThread-0,10,main" threadState: The state of the thread. For example, "WAITING". timestamp: Timestamp at which the stack was captured. """ stackContent = _messages.StringField(1) threadCount = _messages.IntegerField(2, variant=_messages.Variant.INT32) threadName = _messages.StringField(3) threadState = _messages.StringField(4) timestamp = _messages.StringField(5) class StageExecutionDetails(_messages.Message): r"""Information about the workers and work items within a stage. Fields: nextPageToken: If present, this response does not contain all requested tasks. To obtain the next page of results, repeat the request with page_token set to this value. workers: Workers that have done work on the stage. """ nextPageToken = _messages.StringField(1) workers = _messages.MessageField('WorkerDetails', 2, repeated=True) class StageSource(_messages.Message): r"""Description of an input or output of an execution stage. Fields: name: Dataflow service generated name for this source. originalTransformOrCollection: User name for the original user transform or collection with which this source is most closely associated. sizeBytes: Size of the source, if measurable. userName: Human-readable name for this source; may be user or system generated. """ name = _messages.StringField(1) originalTransformOrCollection = _messages.StringField(2) sizeBytes = _messages.IntegerField(3) userName = _messages.StringField(4) class StageSummary(_messages.Message): r"""Information about a particular execution stage of a job. Enums: StateValueValuesEnum: State of this stage. Fields: endTime: End time of this stage. If the work item is completed, this is the actual end time of the stage. Otherwise, it is the predicted end time. metrics: Metrics for this stage. progress: Progress for this stage. Only applicable to Batch jobs. stageId: ID of this stage startTime: Start time of this stage. state: State of this stage. stragglerSummary: Straggler summary for this stage. """ class StateValueValuesEnum(_messages.Enum): r"""State of this stage. Values: EXECUTION_STATE_UNKNOWN: The component state is unknown or unspecified. EXECUTION_STATE_NOT_STARTED: The component is not yet running. EXECUTION_STATE_RUNNING: The component is currently running. EXECUTION_STATE_SUCCEEDED: The component succeeded. EXECUTION_STATE_FAILED: The component failed. EXECUTION_STATE_CANCELLED: Execution of the component was cancelled. """ EXECUTION_STATE_UNKNOWN = 0 EXECUTION_STATE_NOT_STARTED = 1 EXECUTION_STATE_RUNNING = 2 EXECUTION_STATE_SUCCEEDED = 3 EXECUTION_STATE_FAILED = 4 EXECUTION_STATE_CANCELLED = 5 endTime = _messages.StringField(1) metrics = _messages.MessageField('MetricUpdate', 2, repeated=True) progress = _messages.MessageField('ProgressTimeseries', 3) stageId = _messages.StringField(4) startTime = _messages.StringField(5) state = _messages.EnumField('StateValueValuesEnum', 6) stragglerSummary = _messages.MessageField('StragglerSummary', 7) class StandardQueryParameters(_messages.Message): r"""Query parameters accepted by all methods. Enums: FXgafvValueValuesEnum: V1 error format. AltValueValuesEnum: Data format for response. Fields: f__xgafv: V1 error format. access_token: OAuth access token. alt: Data format for response. callback: JSONP fields: Selector specifying which fields to include in a partial response. key: API key. Your API key identifies your project and provides you with API access, quota, and reports. Required unless you provide an OAuth 2.0 token. oauth_token: OAuth 2.0 token for the current user. prettyPrint: Returns response with indentations and line breaks. quotaUser: Available to use for quota purposes for server-side applications. Can be any arbitrary string assigned to a user, but should not exceed 40 characters. trace: A tracing token of the form "token:" to include in api requests. uploadType: Legacy upload protocol for media (e.g. "media", "multipart"). upload_protocol: Upload protocol for media (e.g. "raw", "multipart"). """ class AltValueValuesEnum(_messages.Enum): r"""Data format for response. Values: json: Responses with Content-Type of application/json media: Media download with context-dependent Content-Type proto: Responses with Content-Type of application/x-protobuf """ json = 0 media = 1 proto = 2 class FXgafvValueValuesEnum(_messages.Enum): r"""V1 error format. Values: _1: v1 error format _2: v2 error format """ _1 = 0 _2 = 1 f__xgafv = _messages.EnumField('FXgafvValueValuesEnum', 1) access_token = _messages.StringField(2) alt = _messages.EnumField('AltValueValuesEnum', 3, default='json') callback = _messages.StringField(4) fields = _messages.StringField(5) key = _messages.StringField(6) oauth_token = _messages.StringField(7) prettyPrint = _messages.BooleanField(8, default=True) quotaUser = _messages.StringField(9) trace = _messages.StringField(10) uploadType = _messages.StringField(11) upload_protocol = _messages.StringField(12) class StateFamilyConfig(_messages.Message): r"""State family configuration. Fields: isRead: If true, this family corresponds to a read operation. stateFamily: The state family value. """ isRead = _messages.BooleanField(1) stateFamily = _messages.StringField(2) class Status(_messages.Message): r"""The `Status` type defines a logical error model that is suitable for different programming environments, including REST APIs and RPC APIs. It is used by [gRPC](https://github.com/grpc). Each `Status` message contains three pieces of data: error code, error message, and error details. You can find out more about this error model and how to work with it in the [API Design Guide](https://cloud.google.com/apis/design/errors). Messages: DetailsValueListEntry: A DetailsValueListEntry object. Fields: code: The status code, which should be an enum value of google.rpc.Code. details: A list of messages that carry the error details. There is a common set of message types for APIs to use. message: A developer-facing error message, which should be in English. Any user-facing error message should be localized and sent in the google.rpc.Status.details field, or localized by the client. """ @encoding.MapUnrecognizedFields('additionalProperties') class DetailsValueListEntry(_messages.Message): r"""A DetailsValueListEntry object. Messages: AdditionalProperty: An additional property for a DetailsValueListEntry object. Fields: additionalProperties: Properties of the object. Contains field @type with type URL. """ class AdditionalProperty(_messages.Message): r"""An additional property for a DetailsValueListEntry object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) code = _messages.IntegerField(1, variant=_messages.Variant.INT32) details = _messages.MessageField('DetailsValueListEntry', 2, repeated=True) message = _messages.StringField(3) class Step(_messages.Message): r"""Defines a particular step within a Cloud Dataflow job. A job consists of multiple steps, each of which performs some specific operation as part of the overall job. Data is typically passed from one step to another as part of the job. **Note:** The properties of this object are not stable and might change. Here's an example of a sequence of steps which together implement a Map-Reduce job: * Read a collection of data from some source, parsing the collection's elements. * Validate the elements. * Apply a user-defined function to map each element to some value and extract an element-specific key value. * Group elements with the same key into a single element with that key, transforming a multiply-keyed collection into a uniquely-keyed collection. * Write the elements out to some data sink. Note that the Cloud Dataflow service may be used to run many different types of jobs, not just Map-Reduce. Messages: PropertiesValue: Named properties associated with the step. Each kind of predefined step has its own required set of properties. Must be provided on Create. Only retrieved with JOB_VIEW_ALL. Fields: kind: The kind of step in the Cloud Dataflow job. name: The name that identifies the step. This must be unique for each step with respect to all other steps in the Cloud Dataflow job. properties: Named properties associated with the step. Each kind of predefined step has its own required set of properties. Must be provided on Create. Only retrieved with JOB_VIEW_ALL. """ @encoding.MapUnrecognizedFields('additionalProperties') class PropertiesValue(_messages.Message): r"""Named properties associated with the step. Each kind of predefined step has its own required set of properties. Must be provided on Create. Only retrieved with JOB_VIEW_ALL. Messages: AdditionalProperty: An additional property for a PropertiesValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a PropertiesValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) kind = _messages.StringField(1) name = _messages.StringField(2) properties = _messages.MessageField('PropertiesValue', 3) class Straggler(_messages.Message): r"""Information for a straggler. Fields: batchStraggler: Batch straggler identification and debugging information. streamingStraggler: Streaming straggler identification and debugging information. """ batchStraggler = _messages.MessageField('StragglerInfo', 1) streamingStraggler = _messages.MessageField('StreamingStragglerInfo', 2) class StragglerDebuggingInfo(_messages.Message): r"""Information useful for debugging a straggler. Each type will provide specialized debugging information relevant for a particular cause. The StragglerDebuggingInfo will be 1:1 mapping to the StragglerCause enum. Fields: hotKey: Hot key debugging details. """ hotKey = _messages.MessageField('HotKeyDebuggingInfo', 1) class StragglerInfo(_messages.Message): r"""Information useful for straggler identification and debugging. Messages: CausesValue: The straggler causes, keyed by the string representation of the StragglerCause enum and contains specialized debugging information for each straggler cause. Fields: causes: The straggler causes, keyed by the string representation of the StragglerCause enum and contains specialized debugging information for each straggler cause. startTime: The time when the work item attempt became a straggler. """ @encoding.MapUnrecognizedFields('additionalProperties') class CausesValue(_messages.Message): r"""The straggler causes, keyed by the string representation of the StragglerCause enum and contains specialized debugging information for each straggler cause. Messages: AdditionalProperty: An additional property for a CausesValue object. Fields: additionalProperties: Additional properties of type CausesValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a CausesValue object. Fields: key: Name of the additional property. value: A StragglerDebuggingInfo attribute. """ key = _messages.StringField(1) value = _messages.MessageField('StragglerDebuggingInfo', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) causes = _messages.MessageField('CausesValue', 1) startTime = _messages.StringField(2) class StragglerSummary(_messages.Message): r"""Summarized straggler identification details. Messages: StragglerCauseCountValue: Aggregated counts of straggler causes, keyed by the string representation of the StragglerCause enum. Fields: recentStragglers: The most recent stragglers. stragglerCauseCount: Aggregated counts of straggler causes, keyed by the string representation of the StragglerCause enum. totalStragglerCount: The total count of stragglers. """ @encoding.MapUnrecognizedFields('additionalProperties') class StragglerCauseCountValue(_messages.Message): r"""Aggregated counts of straggler causes, keyed by the string representation of the StragglerCause enum. Messages: AdditionalProperty: An additional property for a StragglerCauseCountValue object. Fields: additionalProperties: Additional properties of type StragglerCauseCountValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a StragglerCauseCountValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.IntegerField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) recentStragglers = _messages.MessageField('Straggler', 1, repeated=True) stragglerCauseCount = _messages.MessageField('StragglerCauseCountValue', 2) totalStragglerCount = _messages.IntegerField(3) class StreamLocation(_messages.Message): r"""Describes a stream of data, either as input to be processed or as output of a streaming Dataflow job. Fields: customSourceLocation: The stream is a custom source. pubsubLocation: The stream is a pubsub stream. sideInputLocation: The stream is a streaming side input. streamingStageLocation: The stream is part of another computation within the current streaming Dataflow job. """ customSourceLocation = _messages.MessageField('CustomSourceLocation', 1) pubsubLocation = _messages.MessageField('PubsubLocation', 2) sideInputLocation = _messages.MessageField('StreamingSideInputLocation', 3) streamingStageLocation = _messages.MessageField('StreamingStageLocation', 4) class StreamingApplianceSnapshotConfig(_messages.Message): r"""Streaming appliance snapshot configuration. Fields: importStateEndpoint: Indicates which endpoint is used to import appliance state. snapshotId: If set, indicates the snapshot id for the snapshot being performed. """ importStateEndpoint = _messages.StringField(1) snapshotId = _messages.StringField(2) class StreamingComputationConfig(_messages.Message): r"""Configuration information for a single streaming computation. Messages: TransformUserNameToStateFamilyValue: Map from user name of stateful transforms in this stage to their state family. Fields: computationId: Unique identifier for this computation. instructions: Instructions that comprise the computation. stageName: Stage name of this computation. systemName: System defined name for this computation. transformUserNameToStateFamily: Map from user name of stateful transforms in this stage to their state family. """ @encoding.MapUnrecognizedFields('additionalProperties') class TransformUserNameToStateFamilyValue(_messages.Message): r"""Map from user name of stateful transforms in this stage to their state family. Messages: AdditionalProperty: An additional property for a TransformUserNameToStateFamilyValue object. Fields: additionalProperties: Additional properties of type TransformUserNameToStateFamilyValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a TransformUserNameToStateFamilyValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) computationId = _messages.StringField(1) instructions = _messages.MessageField('ParallelInstruction', 2, repeated=True) stageName = _messages.StringField(3) systemName = _messages.StringField(4) transformUserNameToStateFamily = _messages.MessageField('TransformUserNameToStateFamilyValue', 5) class StreamingComputationRanges(_messages.Message): r"""Describes full or partial data disk assignment information of the computation ranges. Fields: computationId: The ID of the computation. rangeAssignments: Data disk assignments for ranges from this computation. """ computationId = _messages.StringField(1) rangeAssignments = _messages.MessageField('KeyRangeDataDiskAssignment', 2, repeated=True) class StreamingComputationTask(_messages.Message): r"""A task which describes what action should be performed for the specified streaming computation ranges. Enums: TaskTypeValueValuesEnum: A type of streaming computation task. Fields: computationRanges: Contains ranges of a streaming computation this task should apply to. dataDisks: Describes the set of data disks this task should apply to. taskType: A type of streaming computation task. """ class TaskTypeValueValuesEnum(_messages.Enum): r"""A type of streaming computation task. Values: STREAMING_COMPUTATION_TASK_UNKNOWN: The streaming computation task is unknown, or unspecified. STREAMING_COMPUTATION_TASK_STOP: Stop processing specified streaming computation range(s). STREAMING_COMPUTATION_TASK_START: Start processing specified streaming computation range(s). """ STREAMING_COMPUTATION_TASK_UNKNOWN = 0 STREAMING_COMPUTATION_TASK_STOP = 1 STREAMING_COMPUTATION_TASK_START = 2 computationRanges = _messages.MessageField('StreamingComputationRanges', 1, repeated=True) dataDisks = _messages.MessageField('MountedDataDisk', 2, repeated=True) taskType = _messages.EnumField('TaskTypeValueValuesEnum', 3) class StreamingConfigTask(_messages.Message): r"""A task that carries configuration information for streaming computations. Messages: UserStepToStateFamilyNameMapValue: Map from user step names to state families. Fields: commitStreamChunkSizeBytes: Chunk size for commit streams from the harness to windmill. getDataStreamChunkSizeBytes: Chunk size for get data streams from the harness to windmill. maxWorkItemCommitBytes: Maximum size for work item commit supported windmill storage layer. operationalLimits: Operational limits for the streaming job. Can be used by the worker to validate outputs sent to the backend. streamingComputationConfigs: Set of computation configuration information. streamingEngineStateTagEncodingVersion: Optional. The state tag encoding format version for streaming engine jobs. userStepToStateFamilyNameMap: Map from user step names to state families. userWorkerRunnerV1Settings: Binary encoded proto to control runtime behavior of the java runner v1 user worker. userWorkerRunnerV2Settings: Binary encoded proto to control runtime behavior of the runner v2 user worker. windmillServiceEndpoint: If present, the worker must use this endpoint to communicate with Windmill Service dispatchers, otherwise the worker must continue to use whatever endpoint it had been using. windmillServicePort: If present, the worker must use this port to communicate with Windmill Service dispatchers. Only applicable when windmill_service_endpoint is specified. """ @encoding.MapUnrecognizedFields('additionalProperties') class UserStepToStateFamilyNameMapValue(_messages.Message): r"""Map from user step names to state families. Messages: AdditionalProperty: An additional property for a UserStepToStateFamilyNameMapValue object. Fields: additionalProperties: Additional properties of type UserStepToStateFamilyNameMapValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a UserStepToStateFamilyNameMapValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) commitStreamChunkSizeBytes = _messages.IntegerField(1) getDataStreamChunkSizeBytes = _messages.IntegerField(2) maxWorkItemCommitBytes = _messages.IntegerField(3) operationalLimits = _messages.MessageField('StreamingOperationalLimits', 4) streamingComputationConfigs = _messages.MessageField('StreamingComputationConfig', 5, repeated=True) streamingEngineStateTagEncodingVersion = _messages.IntegerField(6, variant=_messages.Variant.INT32) userStepToStateFamilyNameMap = _messages.MessageField('UserStepToStateFamilyNameMapValue', 7) userWorkerRunnerV1Settings = _messages.BytesField(8) userWorkerRunnerV2Settings = _messages.BytesField(9) windmillServiceEndpoint = _messages.StringField(10) windmillServicePort = _messages.IntegerField(11) class StreamingOperationalLimits(_messages.Message): r"""Operational limits imposed on streaming jobs by the backend. Fields: maxBagElementBytes: The maximum size for an element in bag state. maxGlobalDataBytes: The maximum size for an element in global data. maxKeyBytes: The maximum size allowed for a key. maxProductionOutputBytes: The maximum size for a single output element. maxSortedListElementBytes: The maximum size for an element in sorted list state. maxSourceStateBytes: The maximum size for a source state update. maxTagBytes: The maximum size for a state tag. maxValueBytes: The maximum size for a value state field. """ maxBagElementBytes = _messages.IntegerField(1) maxGlobalDataBytes = _messages.IntegerField(2) maxKeyBytes = _messages.IntegerField(3) maxProductionOutputBytes = _messages.IntegerField(4) maxSortedListElementBytes = _messages.IntegerField(5) maxSourceStateBytes = _messages.IntegerField(6) maxTagBytes = _messages.IntegerField(7) maxValueBytes = _messages.IntegerField(8) class StreamingScalingReport(_messages.Message): r"""Contains per-user worker telemetry used in streaming autoscaling. Fields: activeBundleCount: A integer attribute. activeThreadCount: Current acive thread count. maximumBundleCount: Maximum bundle count. maximumBytes: Maximum bytes. maximumBytesCount: A integer attribute. maximumThreadCount: Maximum thread count limit. outstandingBundleCount: Current outstanding bundle count. outstandingBytes: Current outstanding bytes. outstandingBytesCount: A integer attribute. """ activeBundleCount = _messages.IntegerField(1, variant=_messages.Variant.INT32) activeThreadCount = _messages.IntegerField(2, variant=_messages.Variant.INT32) maximumBundleCount = _messages.IntegerField(3, variant=_messages.Variant.INT32) maximumBytes = _messages.IntegerField(4) maximumBytesCount = _messages.IntegerField(5, variant=_messages.Variant.INT32) maximumThreadCount = _messages.IntegerField(6, variant=_messages.Variant.INT32) outstandingBundleCount = _messages.IntegerField(7, variant=_messages.Variant.INT32) outstandingBytes = _messages.IntegerField(8) outstandingBytesCount = _messages.IntegerField(9, variant=_messages.Variant.INT32) class StreamingScalingReportResponse(_messages.Message): r"""Contains per-user-worker streaming scaling recommendation from the backend. Fields: maximumThreadCount: Maximum thread count limit; """ maximumThreadCount = _messages.IntegerField(1, variant=_messages.Variant.INT32) class StreamingSetupTask(_messages.Message): r"""A task which initializes part of a streaming Dataflow job. Fields: drain: The user has requested drain. receiveWorkPort: The TCP port on which the worker should listen for messages from other streaming computation workers. snapshotConfig: Configures streaming appliance snapshot. streamingComputationTopology: The global topology of the streaming Dataflow job. workerHarnessPort: The TCP port used by the worker to communicate with the Dataflow worker harness. """ drain = _messages.BooleanField(1) receiveWorkPort = _messages.IntegerField(2, variant=_messages.Variant.INT32) snapshotConfig = _messages.MessageField('StreamingApplianceSnapshotConfig', 3) streamingComputationTopology = _messages.MessageField('TopologyConfig', 4) workerHarnessPort = _messages.IntegerField(5, variant=_messages.Variant.INT32) class StreamingSideInputLocation(_messages.Message): r"""Identifies the location of a streaming side input. Fields: stateFamily: Identifies the state family where this side input is stored. tag: Identifies the particular side input within the streaming Dataflow job. """ stateFamily = _messages.StringField(1) tag = _messages.StringField(2) class StreamingStageLocation(_messages.Message): r"""Identifies the location of a streaming computation stage, for stage-to- stage communication. Fields: streamId: Identifies the particular stream within the streaming Dataflow job. """ streamId = _messages.StringField(1) class StreamingStragglerInfo(_messages.Message): r"""Information useful for streaming straggler identification and debugging. Fields: dataWatermarkLag: The event-time watermark lag at the time of the straggler detection. endTime: End time of this straggler. startTime: Start time of this straggler. systemWatermarkLag: The system watermark lag at the time of the straggler detection. workerName: Name of the worker where the straggler was detected. """ dataWatermarkLag = _messages.StringField(1) endTime = _messages.StringField(2) startTime = _messages.StringField(3) systemWatermarkLag = _messages.StringField(4) workerName = _messages.StringField(5) class StringList(_messages.Message): r"""A metric value representing a list of strings. Fields: elements: Elements of the list. """ elements = _messages.StringField(1, repeated=True) class StructuredMessage(_messages.Message): r"""A rich message format, including a human readable string, a key for identifying the message, and structured data associated with the message for programmatic consumption. Fields: messageKey: Identifier for this message type. Used by external systems to internationalize or personalize message. messageText: Human-readable version of message. parameters: The structured data associated with this message. """ messageKey = _messages.StringField(1) messageText = _messages.StringField(2) parameters = _messages.MessageField('Parameter', 3, repeated=True) class TaskRunnerSettings(_messages.Message): r"""Taskrunner configuration settings. Fields: alsologtostderr: Whether to also send taskrunner log info to stderr. baseTaskDir: The location on the worker for task-specific subdirectories. baseUrl: The base URL for the taskrunner to use when accessing Google Cloud APIs. When workers access Google Cloud APIs, they logically do so via relative URLs. If this field is specified, it supplies the base URL to use for resolving these relative URLs. The normative algorithm used is defined by RFC 1808, "Relative Uniform Resource Locators". If not specified, the default value is "http://www.googleapis.com/" commandlinesFileName: The file to store preprocessing commands in. continueOnException: Whether to continue taskrunner if an exception is hit. dataflowApiVersion: The API version of endpoint, e.g. "v1b3" harnessCommand: The command to launch the worker harness. languageHint: The suggested backend language. logDir: The directory on the VM to store logs. logToSerialconsole: Whether to send taskrunner log info to Google Compute Engine VM serial console. logUploadLocation: Indicates where to put logs. If this is not specified, the logs will not be uploaded. The supported resource type is: Google Cloud Storage: storage.googleapis.com/{bucket}/{object} bucket.storage.googleapis.com/{object} oauthScopes: The OAuth2 scopes to be requested by the taskrunner in order to access the Cloud Dataflow API. parallelWorkerSettings: The settings to pass to the parallel worker harness. streamingWorkerMainClass: The streaming worker main class name. taskGroup: The UNIX group ID on the worker VM to use for tasks launched by taskrunner; e.g. "wheel". taskUser: The UNIX user ID on the worker VM to use for tasks launched by taskrunner; e.g. "root". tempStoragePrefix: The prefix of the resources the taskrunner should use for temporary storage. The supported resource type is: Google Cloud Storage: storage.googleapis.com/{bucket}/{object} bucket.storage.googleapis.com/{object} vmId: The ID string of the VM. workflowFileName: The file to store the workflow in. """ alsologtostderr = _messages.BooleanField(1) baseTaskDir = _messages.StringField(2) baseUrl = _messages.StringField(3) commandlinesFileName = _messages.StringField(4) continueOnException = _messages.BooleanField(5) dataflowApiVersion = _messages.StringField(6) harnessCommand = _messages.StringField(7) languageHint = _messages.StringField(8) logDir = _messages.StringField(9) logToSerialconsole = _messages.BooleanField(10) logUploadLocation = _messages.StringField(11) oauthScopes = _messages.StringField(12, repeated=True) parallelWorkerSettings = _messages.MessageField('WorkerSettings', 13) streamingWorkerMainClass = _messages.StringField(14) taskGroup = _messages.StringField(15) taskUser = _messages.StringField(16) tempStoragePrefix = _messages.StringField(17) vmId = _messages.StringField(18) workflowFileName = _messages.StringField(19) class TemplateMetadata(_messages.Message): r"""Metadata describing a template. Fields: defaultStreamingMode: Optional. Indicates the default streaming mode for a streaming template. Only valid if both supports_at_least_once and supports_exactly_once are true. Possible values: UNSPECIFIED, EXACTLY_ONCE and AT_LEAST_ONCE description: Optional. A description of the template. name: Required. The name of the template. parameters: The parameters for the template. streaming: Optional. Indicates if the template is streaming or not. supportsAtLeastOnce: Optional. Indicates if the streaming template supports at least once mode. supportsExactlyOnce: Optional. Indicates if the streaming template supports exactly once mode. yamlDefinition: Optional. For future use. """ defaultStreamingMode = _messages.StringField(1) description = _messages.StringField(2) name = _messages.StringField(3) parameters = _messages.MessageField('ParameterMetadata', 4, repeated=True) streaming = _messages.BooleanField(5) supportsAtLeastOnce = _messages.BooleanField(6) supportsExactlyOnce = _messages.BooleanField(7) yamlDefinition = _messages.StringField(8) class TopologyConfig(_messages.Message): r"""Global topology of the streaming Dataflow job, including all computations and their sharded locations. Messages: UserStageToComputationNameMapValue: Maps user stage names to stable computation names. Fields: computations: The computations associated with a streaming Dataflow job. dataDiskAssignments: The disks assigned to a streaming Dataflow job. forwardingKeyBits: The size (in bits) of keys that will be assigned to source messages. persistentStateVersion: Version number for persistent state. userStageToComputationNameMap: Maps user stage names to stable computation names. """ @encoding.MapUnrecognizedFields('additionalProperties') class UserStageToComputationNameMapValue(_messages.Message): r"""Maps user stage names to stable computation names. Messages: AdditionalProperty: An additional property for a UserStageToComputationNameMapValue object. Fields: additionalProperties: Additional properties of type UserStageToComputationNameMapValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a UserStageToComputationNameMapValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) computations = _messages.MessageField('ComputationTopology', 1, repeated=True) dataDiskAssignments = _messages.MessageField('DataDiskAssignment', 2, repeated=True) forwardingKeyBits = _messages.IntegerField(3, variant=_messages.Variant.INT32) persistentStateVersion = _messages.IntegerField(4, variant=_messages.Variant.INT32) userStageToComputationNameMap = _messages.MessageField('UserStageToComputationNameMapValue', 5) class TransformSummary(_messages.Message): r"""Description of the type, names/ids, and input/outputs for a transform. Enums: KindValueValuesEnum: Type of transform. Fields: displayData: Transform-specific display data. id: SDK generated id of this transform instance. inputCollectionName: User names for all collection inputs to this transform. kind: Type of transform. name: User provided name for this transform instance. outputCollectionName: User names for all collection outputs to this transform. """ class KindValueValuesEnum(_messages.Enum): r"""Type of transform. Values: UNKNOWN_KIND: Unrecognized transform type. PAR_DO_KIND: ParDo transform. GROUP_BY_KEY_KIND: Group By Key transform. FLATTEN_KIND: Flatten transform. READ_KIND: Read transform. WRITE_KIND: Write transform. CONSTANT_KIND: Constructs from a constant value, such as with Create.of. SINGLETON_KIND: Creates a Singleton view of a collection. SHUFFLE_KIND: Opening or closing a shuffle session, often as part of a GroupByKey. """ UNKNOWN_KIND = 0 PAR_DO_KIND = 1 GROUP_BY_KEY_KIND = 2 FLATTEN_KIND = 3 READ_KIND = 4 WRITE_KIND = 5 CONSTANT_KIND = 6 SINGLETON_KIND = 7 SHUFFLE_KIND = 8 displayData = _messages.MessageField('DisplayData', 1, repeated=True) id = _messages.StringField(2) inputCollectionName = _messages.StringField(3, repeated=True) kind = _messages.EnumField('KindValueValuesEnum', 4) name = _messages.StringField(5) outputCollectionName = _messages.StringField(6, repeated=True) class WorkItem(_messages.Message): r"""WorkItem represents basic information about a WorkItem to be executed in the cloud. Fields: configuration: Work item-specific configuration as an opaque blob. id: Identifies this WorkItem. initialReportIndex: The initial index to use when reporting the status of the WorkItem. jobId: Identifies the workflow job this WorkItem belongs to. leaseExpireTime: Time when the lease on this Work will expire. mapTask: Additional information for MapTask WorkItems. packages: Any required packages that need to be fetched in order to execute this WorkItem. projectId: Identifies the cloud project this WorkItem belongs to. reportStatusInterval: Recommended reporting interval. seqMapTask: Additional information for SeqMapTask WorkItems. shellTask: Additional information for ShellTask WorkItems. sourceOperationTask: Additional information for source operation WorkItems. streamingComputationTask: Additional information for StreamingComputationTask WorkItems. streamingConfigTask: Additional information for StreamingConfigTask WorkItems. streamingSetupTask: Additional information for StreamingSetupTask WorkItems. """ configuration = _messages.StringField(1) id = _messages.IntegerField(2) initialReportIndex = _messages.IntegerField(3) jobId = _messages.StringField(4) leaseExpireTime = _messages.StringField(5) mapTask = _messages.MessageField('MapTask', 6) packages = _messages.MessageField('Package', 7, repeated=True) projectId = _messages.StringField(8) reportStatusInterval = _messages.StringField(9) seqMapTask = _messages.MessageField('SeqMapTask', 10) shellTask = _messages.MessageField('ShellTask', 11) sourceOperationTask = _messages.MessageField('SourceOperationRequest', 12) streamingComputationTask = _messages.MessageField('StreamingComputationTask', 13) streamingConfigTask = _messages.MessageField('StreamingConfigTask', 14) streamingSetupTask = _messages.MessageField('StreamingSetupTask', 15) class WorkItemDetails(_messages.Message): r"""Information about an individual work item execution. Enums: StateValueValuesEnum: State of this work item. Fields: attemptId: Attempt ID of this work item endTime: End time of this work item attempt. If the work item is completed, this is the actual end time of the work item. Otherwise, it is the predicted end time. metrics: Metrics for this work item. progress: Progress of this work item. startTime: Start time of this work item attempt. state: State of this work item. stragglerInfo: Information about straggler detections for this work item. taskId: Name of this work item. """ class StateValueValuesEnum(_messages.Enum): r"""State of this work item. Values: EXECUTION_STATE_UNKNOWN: The component state is unknown or unspecified. EXECUTION_STATE_NOT_STARTED: The component is not yet running. EXECUTION_STATE_RUNNING: The component is currently running. EXECUTION_STATE_SUCCEEDED: The component succeeded. EXECUTION_STATE_FAILED: The component failed. EXECUTION_STATE_CANCELLED: Execution of the component was cancelled. """ EXECUTION_STATE_UNKNOWN = 0 EXECUTION_STATE_NOT_STARTED = 1 EXECUTION_STATE_RUNNING = 2 EXECUTION_STATE_SUCCEEDED = 3 EXECUTION_STATE_FAILED = 4 EXECUTION_STATE_CANCELLED = 5 attemptId = _messages.StringField(1) endTime = _messages.StringField(2) metrics = _messages.MessageField('MetricUpdate', 3, repeated=True) progress = _messages.MessageField('ProgressTimeseries', 4) startTime = _messages.StringField(5) state = _messages.EnumField('StateValueValuesEnum', 6) stragglerInfo = _messages.MessageField('StragglerInfo', 7) taskId = _messages.StringField(8) class WorkItemServiceState(_messages.Message): r"""The Dataflow service's idea of the current state of a WorkItem being processed by a worker. Messages: HarnessDataValue: Other data returned by the service, specific to the particular worker harness. Fields: completeWorkStatus: If set, a request to complete the work item with the given status. This will not be set to OK, unless supported by the specific kind of WorkItem. It can be used for the backend to indicate a WorkItem must terminate, e.g., for aborting work. harnessData: Other data returned by the service, specific to the particular worker harness. hotKeyDetection: A hot key is a symptom of poor data distribution in which there are enough elements mapped to a single key to impact pipeline performance. When present, this field includes metadata associated with any hot key. leaseExpireTime: Time at which the current lease will expire. metricShortId: The short ids that workers should use in subsequent metric updates. Workers should strive to use short ids whenever possible, but it is ok to request the short_id again if a worker lost track of it (e.g. if the worker is recovering from a crash). NOTE: it is possible that the response may have short ids for a subset of the metrics. nextReportIndex: The index value to use for the next report sent by the worker. Note: If the report call fails for whatever reason, the worker should reuse this index for subsequent report attempts. reportStatusInterval: New recommended reporting interval. splitRequest: The progress point in the WorkItem where the Dataflow service suggests that the worker truncate the task. suggestedStopPoint: DEPRECATED in favor of split_request. suggestedStopPosition: Obsolete, always empty. """ @encoding.MapUnrecognizedFields('additionalProperties') class HarnessDataValue(_messages.Message): r"""Other data returned by the service, specific to the particular worker harness. Messages: AdditionalProperty: An additional property for a HarnessDataValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a HarnessDataValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) completeWorkStatus = _messages.MessageField('Status', 1) harnessData = _messages.MessageField('HarnessDataValue', 2) hotKeyDetection = _messages.MessageField('HotKeyDetection', 3) leaseExpireTime = _messages.StringField(4) metricShortId = _messages.MessageField('MetricShortId', 5, repeated=True) nextReportIndex = _messages.IntegerField(6) reportStatusInterval = _messages.StringField(7) splitRequest = _messages.MessageField('ApproximateSplitRequest', 8) suggestedStopPoint = _messages.MessageField('ApproximateProgress', 9) suggestedStopPosition = _messages.MessageField('Position', 10) class WorkItemStatus(_messages.Message): r"""Conveys a worker's progress through the work described by a WorkItem. Fields: completed: True if the WorkItem was completed (successfully or unsuccessfully). counterUpdates: Worker output counters for this WorkItem. dynamicSourceSplit: See documentation of stop_position. errors: Specifies errors which occurred during processing. If errors are provided, and completed = true, then the WorkItem is considered to have failed. metricUpdates: DEPRECATED in favor of counter_updates. progress: DEPRECATED in favor of reported_progress. reportIndex: The report index. When a WorkItem is leased, the lease will contain an initial report index. When a WorkItem's status is reported to the system, the report should be sent with that report index, and the response will contain the index the worker should use for the next report. Reports received with unexpected index values will be rejected by the service. In order to preserve idempotency, the worker should not alter the contents of a report, even if the worker must submit the same report multiple times before getting back a response. The worker should not submit a subsequent report until the response for the previous report had been received from the service. reportedProgress: The worker's progress through this WorkItem. requestedLeaseDuration: Amount of time the worker requests for its lease. sourceFork: DEPRECATED in favor of dynamic_source_split. sourceOperationResponse: If the work item represented a SourceOperationRequest, and the work is completed, contains the result of the operation. stopPosition: A worker may split an active map task in two parts, "primary" and "residual", continuing to process the primary part and returning the residual part into the pool of available work. This event is called a "dynamic split" and is critical to the dynamic work rebalancing feature. The two obtained sub-tasks are called "parts" of the split. The parts, if concatenated, must represent the same input as would be read by the current task if the split did not happen. The exact way in which the original task is decomposed into the two parts is specified either as a position demarcating them (stop_position), or explicitly as two DerivedSources, if this task consumes a user-defined source type (dynamic_source_split). The "current" task is adjusted as a result of the split: after a task with range [A, B) sends a stop_position update at C, its range is considered to be [A, C), e.g.: * Progress should be interpreted relative to the new range, e.g. "75% completed" means "75% of [A, C) completed" * The worker should interpret proposed_stop_position relative to the new range, e.g. "split at 68%" should be interpreted as "split at 68% of [A, C)". * If the worker chooses to split again using stop_position, only stop_positions in [A, C) will be accepted. * Etc. dynamic_source_split has similar semantics: e.g., if a task with source S splits using dynamic_source_split into {P, R} (where P and R must be together equivalent to S), then subsequent progress and proposed_stop_position should be interpreted relative to P, and in a potential subsequent dynamic_source_split into {P', R'}, P' and R' must be together equivalent to P, etc. totalThrottlerWaitTimeSeconds: Total time the worker spent being throttled by external systems. workItemId: Identifies the WorkItem. """ completed = _messages.BooleanField(1) counterUpdates = _messages.MessageField('CounterUpdate', 2, repeated=True) dynamicSourceSplit = _messages.MessageField('DynamicSourceSplit', 3) errors = _messages.MessageField('Status', 4, repeated=True) metricUpdates = _messages.MessageField('MetricUpdate', 5, repeated=True) progress = _messages.MessageField('ApproximateProgress', 6) reportIndex = _messages.IntegerField(7) reportedProgress = _messages.MessageField('ApproximateReportedProgress', 8) requestedLeaseDuration = _messages.StringField(9) sourceFork = _messages.MessageField('SourceFork', 10) sourceOperationResponse = _messages.MessageField('SourceOperationResponse', 11) stopPosition = _messages.MessageField('Position', 12) totalThrottlerWaitTimeSeconds = _messages.FloatField(13) workItemId = _messages.StringField(14) class WorkerDetails(_messages.Message): r"""Information about a worker Fields: workItems: Work items processed by this worker, sorted by time. workerName: Name of this worker """ workItems = _messages.MessageField('WorkItemDetails', 1, repeated=True) workerName = _messages.StringField(2) class WorkerHealthReport(_messages.Message): r"""WorkerHealthReport contains information about the health of a worker. The VM should be identified by the labels attached to the WorkerMessage that this health ping belongs to. Messages: PodsValueListEntry: A PodsValueListEntry object. Fields: msg: Message describing any unusual health reports. pods: The pods running on the worker. See: http://kubernetes.io/v1.1/docs/api-reference/v1/definitions.html#_v1_pod This field is used by the worker to send the status of the indvidual containers running on each worker. reportInterval: The interval at which the worker is sending health reports. The default value of 0 should be interpreted as the field is not being explicitly set by the worker. vmBrokenCode: Code to describe a specific reason, if known, that a VM has reported broken state. vmIsBroken: Whether the VM is in a permanently broken state. Broken VMs should be abandoned or deleted ASAP to avoid assigning or completing any work. vmIsHealthy: Whether the VM is currently healthy. vmStartupTime: The time the VM was booted. """ @encoding.MapUnrecognizedFields('additionalProperties') class PodsValueListEntry(_messages.Message): r"""A PodsValueListEntry object. Messages: AdditionalProperty: An additional property for a PodsValueListEntry object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a PodsValueListEntry object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) msg = _messages.StringField(1) pods = _messages.MessageField('PodsValueListEntry', 2, repeated=True) reportInterval = _messages.StringField(3) vmBrokenCode = _messages.StringField(4) vmIsBroken = _messages.BooleanField(5) vmIsHealthy = _messages.BooleanField(6) vmStartupTime = _messages.StringField(7) class WorkerHealthReportResponse(_messages.Message): r"""WorkerHealthReportResponse contains information returned to the worker in response to a health ping. Fields: reportInterval: A positive value indicates the worker should change its reporting interval to the specified value. The default value of zero means no change in report rate is requested by the server. """ reportInterval = _messages.StringField(1) class WorkerLifecycleEvent(_messages.Message): r"""A report of an event in a worker's lifecycle. The proto contains one event, because the worker is expected to asynchronously send each message immediately after the event. Due to this asynchrony, messages may arrive out of order (or missing), and it is up to the consumer to interpret. The timestamp of the event is in the enclosing WorkerMessage proto. Enums: EventValueValuesEnum: The event being reported. Messages: MetadataValue: Other stats that can accompany an event. E.g. { "downloaded_bytes" : "123456" } Fields: containerStartTime: The start time of this container. All events will report this so that events can be grouped together across container/VM restarts. event: The event being reported. metadata: Other stats that can accompany an event. E.g. { "downloaded_bytes" : "123456" } """ class EventValueValuesEnum(_messages.Enum): r"""The event being reported. Values: UNKNOWN_EVENT: Invalid event. OS_START: The time the VM started. CONTAINER_START: Our container code starts running. Multiple containers could be distinguished with WorkerMessage.labels if desired. NETWORK_UP: The worker has a functional external network connection. STAGING_FILES_DOWNLOAD_START: Started downloading staging files. STAGING_FILES_DOWNLOAD_FINISH: Finished downloading all staging files. SDK_INSTALL_START: For applicable SDKs, started installation of SDK and worker packages. SDK_INSTALL_FINISH: Finished installing SDK. """ UNKNOWN_EVENT = 0 OS_START = 1 CONTAINER_START = 2 NETWORK_UP = 3 STAGING_FILES_DOWNLOAD_START = 4 STAGING_FILES_DOWNLOAD_FINISH = 5 SDK_INSTALL_START = 6 SDK_INSTALL_FINISH = 7 @encoding.MapUnrecognizedFields('additionalProperties') class MetadataValue(_messages.Message): r"""Other stats that can accompany an event. E.g. { "downloaded_bytes" : "123456" } Messages: AdditionalProperty: An additional property for a MetadataValue object. Fields: additionalProperties: Additional properties of type MetadataValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a MetadataValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) containerStartTime = _messages.StringField(1) event = _messages.EnumField('EventValueValuesEnum', 2) metadata = _messages.MessageField('MetadataValue', 3) class WorkerMessage(_messages.Message): r"""WorkerMessage provides information to the backend about a worker. Messages: LabelsValue: Labels are used to group WorkerMessages. For example, a worker_message about a particular container might have the labels: { "JOB_ID": "2015-04-22", "WORKER_ID": "wordcount-vm-2015..." "CONTAINER_TYPE": "worker", "CONTAINER_ID": "ac1234def"} Label tags typically correspond to Label enum values. However, for ease of development other strings can be used as tags. LABEL_UNSPECIFIED should not be used here. Fields: dataSamplingReport: Optional. Contains metrics related to go/dataflow- data-sampling-telemetry. labels: Labels are used to group WorkerMessages. For example, a worker_message about a particular container might have the labels: { "JOB_ID": "2015-04-22", "WORKER_ID": "wordcount-vm-2015..." "CONTAINER_TYPE": "worker", "CONTAINER_ID": "ac1234def"} Label tags typically correspond to Label enum values. However, for ease of development other strings can be used as tags. LABEL_UNSPECIFIED should not be used here. perWorkerMetrics: System defined metrics for this worker. streamingScalingReport: Contains per-user worker telemetry used in streaming autoscaling. time: The timestamp of the worker_message. workerHealthReport: The health of a worker. workerLifecycleEvent: Record of worker lifecycle events. workerMessageCode: A worker message code. workerMetrics: Resource metrics reported by workers. workerShutdownNotice: Shutdown notice by workers. workerThreadScalingReport: Thread scaling information reported by workers. """ @encoding.MapUnrecognizedFields('additionalProperties') class LabelsValue(_messages.Message): r"""Labels are used to group WorkerMessages. For example, a worker_message about a particular container might have the labels: { "JOB_ID": "2015-04-22", "WORKER_ID": "wordcount-vm-2015..." "CONTAINER_TYPE": "worker", "CONTAINER_ID": "ac1234def"} Label tags typically correspond to Label enum values. However, for ease of development other strings can be used as tags. LABEL_UNSPECIFIED should not be used here. Messages: AdditionalProperty: An additional property for a LabelsValue object. Fields: additionalProperties: Additional properties of type LabelsValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a LabelsValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) dataSamplingReport = _messages.MessageField('DataSamplingReport', 1) labels = _messages.MessageField('LabelsValue', 2) perWorkerMetrics = _messages.MessageField('PerWorkerMetrics', 3) streamingScalingReport = _messages.MessageField('StreamingScalingReport', 4) time = _messages.StringField(5) workerHealthReport = _messages.MessageField('WorkerHealthReport', 6) workerLifecycleEvent = _messages.MessageField('WorkerLifecycleEvent', 7) workerMessageCode = _messages.MessageField('WorkerMessageCode', 8) workerMetrics = _messages.MessageField('ResourceUtilizationReport', 9) workerShutdownNotice = _messages.MessageField('WorkerShutdownNotice', 10) workerThreadScalingReport = _messages.MessageField('WorkerThreadScalingReport', 11) class WorkerMessageCode(_messages.Message): r"""A message code is used to report status and error messages to the service. The message codes are intended to be machine readable. The service will take care of translating these into user understandable messages if necessary. Example use cases: 1. Worker processes reporting successful startup. 2. Worker processes reporting specific errors (e.g. package staging failure). Messages: ParametersValue: Parameters contains specific information about the code. This is a struct to allow parameters of different types. Examples: 1. For a "HARNESS_STARTED" message parameters might provide the name of the worker and additional data like timing information. 2. For a "GCS_DOWNLOAD_ERROR" parameters might contain fields listing the Cloud Storage objects being downloaded and fields containing errors. In general complex data structures should be avoided. If a worker needs to send a specific and complicated data structure then please consider defining a new proto and adding it to the data oneof in WorkerMessageResponse. Conventions: Parameters should only be used for information that isn't typically passed as a label. hostname and other worker identifiers should almost always be passed as labels since they will be included on most messages. Fields: code: The code is a string intended for consumption by a machine that identifies the type of message being sent. Examples: 1. "HARNESS_STARTED" might be used to indicate the worker harness has started. 2. "GCS_DOWNLOAD_ERROR" might be used to indicate an error downloading a Cloud Storage file as part of the boot process of one of the worker containers. This is a string and not an enum to make it easy to add new codes without waiting for an API change. parameters: Parameters contains specific information about the code. This is a struct to allow parameters of different types. Examples: 1. For a "HARNESS_STARTED" message parameters might provide the name of the worker and additional data like timing information. 2. For a "GCS_DOWNLOAD_ERROR" parameters might contain fields listing the Cloud Storage objects being downloaded and fields containing errors. In general complex data structures should be avoided. If a worker needs to send a specific and complicated data structure then please consider defining a new proto and adding it to the data oneof in WorkerMessageResponse. Conventions: Parameters should only be used for information that isn't typically passed as a label. hostname and other worker identifiers should almost always be passed as labels since they will be included on most messages. """ @encoding.MapUnrecognizedFields('additionalProperties') class ParametersValue(_messages.Message): r"""Parameters contains specific information about the code. This is a struct to allow parameters of different types. Examples: 1. For a "HARNESS_STARTED" message parameters might provide the name of the worker and additional data like timing information. 2. For a "GCS_DOWNLOAD_ERROR" parameters might contain fields listing the Cloud Storage objects being downloaded and fields containing errors. In general complex data structures should be avoided. If a worker needs to send a specific and complicated data structure then please consider defining a new proto and adding it to the data oneof in WorkerMessageResponse. Conventions: Parameters should only be used for information that isn't typically passed as a label. hostname and other worker identifiers should almost always be passed as labels since they will be included on most messages. Messages: AdditionalProperty: An additional property for a ParametersValue object. Fields: additionalProperties: Properties of the object. """ class AdditionalProperty(_messages.Message): r"""An additional property for a ParametersValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) code = _messages.StringField(1) parameters = _messages.MessageField('ParametersValue', 2) class WorkerMessageResponse(_messages.Message): r"""A worker_message response allows the server to pass information to the sender. Fields: streamingScalingReportResponse: Service's streaming scaling response for workers. workerHealthReportResponse: The service's response to a worker's health report. workerMetricsResponse: Service's response to reporting worker metrics (currently empty). workerShutdownNoticeResponse: Service's response to shutdown notice (currently empty). workerThreadScalingReportResponse: Service's thread scaling recommendation for workers. """ streamingScalingReportResponse = _messages.MessageField('StreamingScalingReportResponse', 1) workerHealthReportResponse = _messages.MessageField('WorkerHealthReportResponse', 2) workerMetricsResponse = _messages.MessageField('ResourceUtilizationReportResponse', 3) workerShutdownNoticeResponse = _messages.MessageField('WorkerShutdownNoticeResponse', 4) workerThreadScalingReportResponse = _messages.MessageField('WorkerThreadScalingReportResponse', 5) class WorkerPool(_messages.Message): r"""Describes one particular pool of Cloud Dataflow workers to be instantiated by the Cloud Dataflow service in order to perform the computations required by a job. Note that a workflow job may use multiple pools, in order to match the various computational requirements of the various stages of the job. Enums: DefaultPackageSetValueValuesEnum: The default package set to install. This allows the service to select a default set of packages which are useful to worker harnesses written in a particular language. IpConfigurationValueValuesEnum: Configuration for VM IPs. TeardownPolicyValueValuesEnum: Sets the policy for determining when to turndown worker pool. Allowed values are: `TEARDOWN_ALWAYS`, `TEARDOWN_ON_SUCCESS`, and `TEARDOWN_NEVER`. `TEARDOWN_ALWAYS` means workers are always torn down regardless of whether the job succeeds. `TEARDOWN_ON_SUCCESS` means workers are torn down if the job succeeds. `TEARDOWN_NEVER` means the workers are never torn down. If the workers are not torn down by the service, they will continue to run and use Google Compute Engine VM resources in the user's project until they are explicitly terminated by the user. Because of this, Google recommends using the `TEARDOWN_ALWAYS` policy except for small, manually supervised test jobs. If unknown or unspecified, the service will attempt to choose a reasonable default. Messages: MetadataValue: Metadata to set on the Google Compute Engine VMs. PoolArgsValue: Extra arguments for this worker pool. Fields: autoscalingSettings: Settings for autoscaling of this WorkerPool. dataDisks: Data disks that are used by a VM in this workflow. defaultPackageSet: The default package set to install. This allows the service to select a default set of packages which are useful to worker harnesses written in a particular language. diskProvisionedIops: Optional. IOPS provisioned for the root disk for VMs. diskProvisionedThroughputMibps: Optional. Throughput provisioned for the root disk for VMs. diskSizeGb: Size of root disk for VMs, in GB. If zero or unspecified, the service will attempt to choose a reasonable default. diskSourceImage: Fully qualified source image for disks. diskType: Type of root disk for VMs. If empty or unspecified, the service will attempt to choose a reasonable default. ipConfiguration: Configuration for VM IPs. kind: The kind of the worker pool; currently only `harness` and `shuffle` are supported. machineType: Machine type (e.g. "n1-standard-1"). If empty or unspecified, the service will attempt to choose a reasonable default. metadata: Metadata to set on the Google Compute Engine VMs. network: Network to which VMs will be assigned. If empty or unspecified, the service will use the network "default". numThreadsPerWorker: The number of threads per worker harness. If empty or unspecified, the service will choose a number of threads (according to the number of cores on the selected machine type for batch, or 1 by convention for streaming). numWorkers: Number of Google Compute Engine workers in this pool needed to execute the job. If zero or unspecified, the service will attempt to choose a reasonable default. onHostMaintenance: The action to take on host maintenance, as defined by the Google Compute Engine API. packages: Packages to be installed on workers. poolArgs: Extra arguments for this worker pool. sdkHarnessContainerImages: Set of SDK harness containers needed to execute this pipeline. This will only be set in the Fn API path. For non-cross- language pipelines this should have only one entry. Cross-language pipelines will have two or more entries. subnetwork: Subnetwork to which VMs will be assigned, if desired. Expected to be of the form "regions/REGION/subnetworks/SUBNETWORK". taskrunnerSettings: Settings passed through to Google Compute Engine workers when using the standard Dataflow task runner. Users should ignore this field. teardownPolicy: Sets the policy for determining when to turndown worker pool. Allowed values are: `TEARDOWN_ALWAYS`, `TEARDOWN_ON_SUCCESS`, and `TEARDOWN_NEVER`. `TEARDOWN_ALWAYS` means workers are always torn down regardless of whether the job succeeds. `TEARDOWN_ON_SUCCESS` means workers are torn down if the job succeeds. `TEARDOWN_NEVER` means the workers are never torn down. If the workers are not torn down by the service, they will continue to run and use Google Compute Engine VM resources in the user's project until they are explicitly terminated by the user. Because of this, Google recommends using the `TEARDOWN_ALWAYS` policy except for small, manually supervised test jobs. If unknown or unspecified, the service will attempt to choose a reasonable default. workerHarnessContainerImage: Required. Docker container image that executes the Cloud Dataflow worker harness, residing in Google Container Registry. Deprecated for the Fn API path. Use sdk_harness_container_images instead. zone: Zone to run the worker pools in. If empty or unspecified, the service will attempt to choose a reasonable default. """ class DefaultPackageSetValueValuesEnum(_messages.Enum): r"""The default package set to install. This allows the service to select a default set of packages which are useful to worker harnesses written in a particular language. Values: DEFAULT_PACKAGE_SET_UNKNOWN: The default set of packages to stage is unknown, or unspecified. DEFAULT_PACKAGE_SET_NONE: Indicates that no packages should be staged at the worker unless explicitly specified by the job. DEFAULT_PACKAGE_SET_JAVA: Stage packages typically useful to workers written in Java. DEFAULT_PACKAGE_SET_PYTHON: Stage packages typically useful to workers written in Python. """ DEFAULT_PACKAGE_SET_UNKNOWN = 0 DEFAULT_PACKAGE_SET_NONE = 1 DEFAULT_PACKAGE_SET_JAVA = 2 DEFAULT_PACKAGE_SET_PYTHON = 3 class IpConfigurationValueValuesEnum(_messages.Enum): r"""Configuration for VM IPs. Values: WORKER_IP_UNSPECIFIED: The configuration is unknown, or unspecified. WORKER_IP_PUBLIC: Workers should have public IP addresses. WORKER_IP_PRIVATE: Workers should have private IP addresses. """ WORKER_IP_UNSPECIFIED = 0 WORKER_IP_PUBLIC = 1 WORKER_IP_PRIVATE = 2 class TeardownPolicyValueValuesEnum(_messages.Enum): r"""Sets the policy for determining when to turndown worker pool. Allowed values are: `TEARDOWN_ALWAYS`, `TEARDOWN_ON_SUCCESS`, and `TEARDOWN_NEVER`. `TEARDOWN_ALWAYS` means workers are always torn down regardless of whether the job succeeds. `TEARDOWN_ON_SUCCESS` means workers are torn down if the job succeeds. `TEARDOWN_NEVER` means the workers are never torn down. If the workers are not torn down by the service, they will continue to run and use Google Compute Engine VM resources in the user's project until they are explicitly terminated by the user. Because of this, Google recommends using the `TEARDOWN_ALWAYS` policy except for small, manually supervised test jobs. If unknown or unspecified, the service will attempt to choose a reasonable default. Values: TEARDOWN_POLICY_UNKNOWN: The teardown policy isn't specified, or is unknown. TEARDOWN_ALWAYS: Always teardown the resource. TEARDOWN_ON_SUCCESS: Teardown the resource on success. This is useful for debugging failures. TEARDOWN_NEVER: Never teardown the resource. This is useful for debugging and development. """ TEARDOWN_POLICY_UNKNOWN = 0 TEARDOWN_ALWAYS = 1 TEARDOWN_ON_SUCCESS = 2 TEARDOWN_NEVER = 3 @encoding.MapUnrecognizedFields('additionalProperties') class MetadataValue(_messages.Message): r"""Metadata to set on the Google Compute Engine VMs. Messages: AdditionalProperty: An additional property for a MetadataValue object. Fields: additionalProperties: Additional properties of type MetadataValue """ class AdditionalProperty(_messages.Message): r"""An additional property for a MetadataValue object. Fields: key: Name of the additional property. value: A string attribute. """ key = _messages.StringField(1) value = _messages.StringField(2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) @encoding.MapUnrecognizedFields('additionalProperties') class PoolArgsValue(_messages.Message): r"""Extra arguments for this worker pool. Messages: AdditionalProperty: An additional property for a PoolArgsValue object. Fields: additionalProperties: Properties of the object. Contains field @type with type URL. """ class AdditionalProperty(_messages.Message): r"""An additional property for a PoolArgsValue object. Fields: key: Name of the additional property. value: A extra_types.JsonValue attribute. """ key = _messages.StringField(1) value = _messages.MessageField('extra_types.JsonValue', 2) additionalProperties = _messages.MessageField('AdditionalProperty', 1, repeated=True) autoscalingSettings = _messages.MessageField('AutoscalingSettings', 1) dataDisks = _messages.MessageField('Disk', 2, repeated=True) defaultPackageSet = _messages.EnumField('DefaultPackageSetValueValuesEnum', 3) diskProvisionedIops = _messages.IntegerField(4) diskProvisionedThroughputMibps = _messages.IntegerField(5) diskSizeGb = _messages.IntegerField(6, variant=_messages.Variant.INT32) diskSourceImage = _messages.StringField(7) diskType = _messages.StringField(8) ipConfiguration = _messages.EnumField('IpConfigurationValueValuesEnum', 9) kind = _messages.StringField(10) machineType = _messages.StringField(11) metadata = _messages.MessageField('MetadataValue', 12) network = _messages.StringField(13) numThreadsPerWorker = _messages.IntegerField(14, variant=_messages.Variant.INT32) numWorkers = _messages.IntegerField(15, variant=_messages.Variant.INT32) onHostMaintenance = _messages.StringField(16) packages = _messages.MessageField('Package', 17, repeated=True) poolArgs = _messages.MessageField('PoolArgsValue', 18) sdkHarnessContainerImages = _messages.MessageField('SdkHarnessContainerImage', 19, repeated=True) subnetwork = _messages.StringField(20) taskrunnerSettings = _messages.MessageField('TaskRunnerSettings', 21) teardownPolicy = _messages.EnumField('TeardownPolicyValueValuesEnum', 22) workerHarnessContainerImage = _messages.StringField(23) zone = _messages.StringField(24) class WorkerSettings(_messages.Message): r"""Provides data to pass through to the worker harness. Fields: baseUrl: The base URL for accessing Google Cloud APIs. When workers access Google Cloud APIs, they logically do so via relative URLs. If this field is specified, it supplies the base URL to use for resolving these relative URLs. The normative algorithm used is defined by RFC 1808, "Relative Uniform Resource Locators". If not specified, the default value is "http://www.googleapis.com/" reportingEnabled: Whether to send work progress updates to the service. servicePath: The Cloud Dataflow service path relative to the root URL, for example, "dataflow/v1b3/projects". shuffleServicePath: The Shuffle service path relative to the root URL, for example, "shuffle/v1beta1". tempStoragePrefix: The prefix of the resources the system should use for temporary storage. The supported resource type is: Google Cloud Storage: storage.googleapis.com/{bucket}/{object} bucket.storage.googleapis.com/{object} workerId: The ID of the worker running this pipeline. """ baseUrl = _messages.StringField(1) reportingEnabled = _messages.BooleanField(2) servicePath = _messages.StringField(3) shuffleServicePath = _messages.StringField(4) tempStoragePrefix = _messages.StringField(5) workerId = _messages.StringField(6) class WorkerShutdownNotice(_messages.Message): r"""Shutdown notification from workers. This is to be sent by the shutdown script of the worker VM so that the backend knows that the VM is being shut down. Fields: reason: The reason for the worker shutdown. Current possible values are: "UNKNOWN": shutdown reason is unknown. "PREEMPTION": shutdown reason is preemption. Other possible reasons may be added in the future. """ reason = _messages.StringField(1) class WorkerShutdownNoticeResponse(_messages.Message): r"""Service-side response to WorkerMessage issuing shutdown notice.""" class WorkerThreadScalingReport(_messages.Message): r"""Contains information about the thread scaling information of a worker. Fields: currentThreadCount: Current number of active threads in a worker. """ currentThreadCount = _messages.IntegerField(1, variant=_messages.Variant.INT32) class WorkerThreadScalingReportResponse(_messages.Message): r"""Contains the thread scaling recommendation for a worker from the backend. Fields: recommendedThreadCount: Recommended number of threads for a worker. """ recommendedThreadCount = _messages.IntegerField(1, variant=_messages.Variant.INT32) class WriteInstruction(_messages.Message): r"""An instruction that writes records. Takes one input, produces no outputs. Fields: input: The input. sink: The sink to write to. """ input = _messages.MessageField('InstructionInput', 1) sink = _messages.MessageField('Sink', 2) encoding.AddCustomJsonFieldMapping( StandardQueryParameters, 'f__xgafv', '$.xgafv') encoding.AddCustomJsonEnumMapping( StandardQueryParameters.FXgafvValueValuesEnum, '_1', '1') encoding.AddCustomJsonEnumMapping( StandardQueryParameters.FXgafvValueValuesEnum, '_2', '2') encoding.AddCustomJsonFieldMapping( DataflowProjectsLocationsTemplatesLaunchRequest, 'dynamicTemplate_gcsPath', 'dynamicTemplate.gcsPath') encoding.AddCustomJsonFieldMapping( DataflowProjectsLocationsTemplatesLaunchRequest, 'dynamicTemplate_stagingLocation', 'dynamicTemplate.stagingLocation') encoding.AddCustomJsonFieldMapping( DataflowProjectsTemplatesLaunchRequest, 'dynamicTemplate_gcsPath', 'dynamicTemplate.gcsPath') encoding.AddCustomJsonFieldMapping( DataflowProjectsTemplatesLaunchRequest, 'dynamicTemplate_stagingLocation', 'dynamicTemplate.stagingLocation') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/internal/clients/dataflow/message_matchers.py0000644000175100017510000001012415206113032032112 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file from hamcrest.core.base_matcher import BaseMatcher IGNORED = object() class MetricStructuredNameMatcher(BaseMatcher): """Matches a MetricStructuredName.""" def __init__(self, name=IGNORED, origin=IGNORED, context=IGNORED): """Creates a MetricsStructuredNameMatcher. Any property not passed in to the constructor will be ignored when matching. Args: name: A string with the metric name. origin: A string with the metric namespace. context: A key:value dictionary that will be matched to the structured name. """ if context != IGNORED and not isinstance(context, dict): raise ValueError('context must be a Python dictionary.') self.name = name self.origin = origin self.context = context def _matches(self, item): if self.name != IGNORED and item.name != self.name: return False if self.origin != IGNORED and item.origin != self.origin: return False if self.context != IGNORED: for key, name in self.context.items(): if key not in item.context: return False if name != IGNORED and item.context[key] != name: return False return True def describe_to(self, description): descriptors = [] if self.name != IGNORED: descriptors.append('name is {}'.format(self.name)) if self.origin != IGNORED: descriptors.append('origin is {}'.format(self.origin)) if self.context != IGNORED: descriptors.append('context is ({})'.format(str(self.context))) item_description = ' and '.join(descriptors) description.append(item_description) class MetricUpdateMatcher(BaseMatcher): """Matches a metrics update protocol buffer.""" def __init__( self, cumulative=IGNORED, name=IGNORED, scalar=IGNORED, kind=IGNORED): """Creates a MetricUpdateMatcher. Any property not passed in to the constructor will be ignored when matching. Args: cumulative: A boolean. name: A MetricStructuredNameMatcher object that matches the name. scalar: An integer with the metric update. kind: A string defining the kind of counter. """ if name != IGNORED and not isinstance(name, MetricStructuredNameMatcher): raise ValueError('name must be a MetricStructuredNameMatcher.') self.cumulative = cumulative self.name = name self.scalar = scalar self.kind = kind def _matches(self, item): if self.cumulative != IGNORED and item.cumulative != self.cumulative: return False if self.name != IGNORED and not self.name._matches(item.name): return False if self.kind != IGNORED and item.kind != self.kind: return False if self.scalar != IGNORED: value_property = [ p for p in item.scalar.object_value.properties if p.key == 'value' ] int_value = value_property[0].value.integer_value if self.scalar != int_value: return False return True def describe_to(self, description): descriptors = [] if self.cumulative != IGNORED: descriptors.append('cumulative is {}'.format(self.cumulative)) if self.name != IGNORED: descriptors.append('name is {}'.format(self.name)) if self.scalar != IGNORED: descriptors.append('scalar is ({})'.format(str(self.scalar))) item_description = ' and '.join(descriptors) description.append(item_description) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/internal/clients/dataflow/message_matchers_test.py0000644000175100017510000000530615206113032033157 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import hamcrest as hc import apache_beam.runners.dataflow.internal.clients.dataflow as dataflow from apache_beam.internal.gcp.json_value import to_json_value from apache_beam.runners.dataflow.internal.clients.dataflow import message_matchers # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apitools.base.py import base_api except ImportError: base_api = None # pylint: enable=wrong-import-order, wrong-import-position @unittest.skipIf(base_api is None, 'GCP dependencies are not installed') class TestMatchers(unittest.TestCase): def test_structured_name_matcher_basic(self): metric_name = dataflow.MetricStructuredName() metric_name.name = 'metric1' metric_name.origin = 'origin2' matcher = message_matchers.MetricStructuredNameMatcher( name='metric1', origin='origin2') hc.assert_that(metric_name, hc.is_(matcher)) with self.assertRaises(AssertionError): matcher = message_matchers.MetricStructuredNameMatcher( name='metric1', origin='origin1') hc.assert_that(metric_name, hc.is_(matcher)) def test_metric_update_basic(self): metric_update = dataflow.MetricUpdate() metric_update.name = dataflow.MetricStructuredName() metric_update.name.name = 'metric1' metric_update.name.origin = 'origin1' metric_update.cumulative = False metric_update.kind = 'sum' metric_update.scalar = to_json_value(1, with_type=True) name_matcher = message_matchers.MetricStructuredNameMatcher( name='metric1', origin='origin1') matcher = message_matchers.MetricUpdateMatcher( name=name_matcher, kind='sum', scalar=1) hc.assert_that(metric_update, hc.is_(matcher)) with self.assertRaises(AssertionError): matcher.kind = 'suma' hc.assert_that(metric_update, hc.is_(matcher)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/internal/names.py0000644000175100017510000000314115206113032024442 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Various names for properties, transforms, etc.""" # All constants are for internal use only; no backwards-compatibility # guarantees. # pytype: skip-file # Referenced by Dataflow legacy worker. from apache_beam.runners.internal.names import PICKLED_MAIN_SESSION_FILE # pylint: disable=unused-import # String constants related to sources framework SOURCE_FORMAT = 'custom_source' SOURCE_TYPE = 'CustomSourcesType' SERIALIZED_SOURCE_KEY = 'serialized_source' # In a released SDK, Python sdk container image is tagged with the SDK version. # Unreleased sdks use container image tag specified below. # Update this tag whenever there is a change that # requires changes to SDK harness container or SDK harness launcher. BEAM_DEV_SDK_CONTAINER_TAG = 'beam-master-20260427' DATAFLOW_CONTAINER_IMAGE_REPOSITORY = 'gcr.io/cloud-dataflow/v1beta3' ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/ptransform_overrides.py0000644000175100017510000001123415206113032026002 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Ptransform overrides for DataflowRunner.""" # pytype: skip-file from apache_beam.options.pipeline_options import StandardOptions from apache_beam.pipeline import PTransformOverride class StreamingPubSubWriteDoFnOverride(PTransformOverride): """Override ParDo(_PubSubWriteDoFn) for streaming mode in DataflowRunner. This override specifically targets the final ParDo step in WriteToPubSub and replaces it with Write(sink) for streaming optimization. """ def matches(self, applied_ptransform): from apache_beam.io.gcp.pubsub import _PubSubWriteDoFn from apache_beam.transforms import ParDo if not isinstance(applied_ptransform.transform, ParDo): return False # Check if this ParDo uses _PubSubWriteDoFn dofn = applied_ptransform.transform.dofn return isinstance(dofn, _PubSubWriteDoFn) def get_replacement_transform_for_applied_ptransform( self, applied_ptransform): from apache_beam.io.iobase import Write # Get the WriteToPubSub transform from the DoFn constructor parameter dofn = applied_ptransform.transform.dofn # The DoFn was initialized with the WriteToPubSub transform # We need to reconstruct the sink from the DoFn's stored properties if hasattr(dofn, 'project') and hasattr(dofn, 'short_topic_name'): from apache_beam.io.gcp.pubsub import _PubSubSink # Create a sink with the same properties as the original topic = f"projects/{dofn.project}/topics/{dofn.short_topic_name}" sink = _PubSubSink( topic=topic, id_label=getattr(dofn, 'id_label', None), timestamp_attribute=getattr(dofn, 'timestamp_attribute', None)) return Write(sink) else: # Fallback: return the original transform if we can't reconstruct it return applied_ptransform.transform def get_dataflow_transform_overrides(pipeline_options): """Returns DataflowRunner-specific transform overrides. Args: pipeline_options: Pipeline options to determine which overrides to apply. Returns: List of PTransformOverride objects for DataflowRunner. """ overrides = [] # Only add streaming-specific overrides when in streaming mode if pipeline_options.view_as(StandardOptions).streaming: # Add PubSub ParDo streaming override that targets only the final step overrides.append(StreamingPubSubWriteDoFnOverride()) return overrides class NativeReadPTransformOverride(PTransformOverride): """A ``PTransformOverride`` for ``Read`` using native sources. The DataflowRunner expects that the Read PTransform using native sources act as a primitive. So this override replaces the Read with a primitive. """ def matches(self, applied_ptransform): # Imported here to avoid circular dependencies. # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.io import Read # Consider the native Read to be a primitive for Dataflow by replacing. return ( isinstance(applied_ptransform.transform, Read) and not getattr(applied_ptransform.transform, 'override', False) and hasattr(applied_ptransform.transform.source, 'format')) def get_replacement_transform(self, ptransform): # Imported here to avoid circular dependencies. # pylint: disable=wrong-import-order, wrong-import-position from apache_beam import pvalue from apache_beam.io import iobase # This is purposely subclassed from the Read transform to take advantage of # the existing windowing, typing, and display data. class Read(iobase.Read): override = True def expand(self, pbegin): return pvalue.PCollection.from_(pbegin) # Use the source's coder type hint as this replacement's output. Otherwise, # the typing information is not properly forwarded to the DataflowRunner # and will choose the incorrect coder for this transform. return Read(ptransform.source).with_output_types( ptransform.source.coder.to_type_hint()) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/template_runner_test.py0000644000175100017510000000562515206113032025777 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for templated pipelines.""" # pytype: skip-file import json import tempfile import unittest import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.pipeline import Pipeline from apache_beam.runners.dataflow.dataflow_runner import DataflowRunner # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.runners.dataflow.internal import apiclient except ImportError: apiclient = None # type: ignore # pylint: enable=wrong-import-order, wrong-import-position @unittest.skipIf(apiclient is None, 'GCP dependencies are not installed') class TemplatingDataflowRunnerTest(unittest.TestCase): """TemplatingDataflow tests.""" def test_full_completion(self): # Create dummy file and close it. Note that we need to do this because # Windows does not allow NamedTemporaryFiles to be reopened elsewhere # before the temporary file is closed. dummy_file = tempfile.NamedTemporaryFile(delete=False) dummy_file_name = dummy_file.name dummy_file.close() dummy_dir = tempfile.mkdtemp() remote_runner = DataflowRunner() options = PipelineOptions([ '--sdk_location=' + dummy_file_name, '--job_name=test-job', '--project=apache-beam-testing', '--region=us-central1', '--staging_location=gs://apache-beam-testing-stg/stg/', '--temp_location=gs://apache-beam-testing-temp/tmp', '--template_location=' + dummy_file_name ]) with Pipeline(remote_runner, options) as pipeline: pipeline | beam.Create([1, 2, 3]) | beam.Map(lambda x: x) # pylint: disable=expression-not-assigned with open(dummy_file_name) as template_file: saved_job_dict = json.load(template_file) self.assertEqual( saved_job_dict['environment']['sdkPipelineOptions']['options'] ['project'], 'apache-beam-testing') self.assertEqual( saved_job_dict['environment']['sdkPipelineOptions']['options'] ['job_name'], 'test-job') if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/dataflow/test_dataflow_runner.py0000644000175100017510000001016615206113032025761 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Wrapper of Beam runners that's built for running and verifying e2e tests.""" # pytype: skip-file import logging import time from apache_beam.internal import pickler from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import TestOptions from apache_beam.runners.dataflow.dataflow_runner import DataflowRunner from apache_beam.runners.runner import PipelineState __all__ = ['TestDataflowRunner'] # Dataflow take up to 10mins for the long tail of starting/stopping worker # pool. WAIT_IN_STATE_TIMEOUT = 10 * 60 _LOGGER = logging.getLogger(__name__) class TestDataflowRunner(DataflowRunner): def run_pipeline(self, pipeline, options): """Execute test pipeline and verify test matcher""" test_options = options.view_as(TestOptions) on_success_matcher = test_options.on_success_matcher wait_duration = test_options.wait_until_finish_duration is_streaming = options.view_as(StandardOptions).streaming # [BEAM-1889] Do not send this to remote workers also, there is no need to # send this option to remote executors. test_options.on_success_matcher = None self.result = super().run_pipeline(pipeline, options) if self.result.has_job: # TODO(markflyhigh)(https://github.com/apache/beam/issues/18254): Use # print since Nose dosen't show logs in some cases. print('Worker logs: %s' % self.build_console_url(options)) _LOGGER.info('Console log: ') _LOGGER.info(self.build_console_url(options)) try: self.wait_until_in_state(PipelineState.RUNNING) if is_streaming and not wait_duration: _LOGGER.warning('Waiting indefinitely for streaming job.') self.result.wait_until_finish(duration=wait_duration) if on_success_matcher: from hamcrest import assert_that as hc_assert_that hc_assert_that(self.result, pickler.loads(on_success_matcher)) finally: if not self.result.is_in_terminal_state(): self.result.cancel() self.wait_until_in_state(PipelineState.CANCELLED) return self.result def build_console_url(self, options): """Build a console url of Dataflow job.""" project = options.view_as(GoogleCloudOptions).project region_id = options.view_as(GoogleCloudOptions).region job_id = self.result.job_id() return ( 'https://console.cloud.google.com/dataflow/jobs/%s/%s?project=%s' % (region_id, job_id, project)) def wait_until_in_state(self, expected_state, timeout=WAIT_IN_STATE_TIMEOUT): """Wait until Dataflow pipeline enters a certain state.""" consoleUrl = ( "Console URL: https://console.cloud.google.com/dataflow/" f"/{self.result.job_id()}?project=") if not self.result.has_job: _LOGGER.error(consoleUrl) raise IOError('Failed to get the Dataflow job id.') start_time = time.time() while time.time() - start_time <= timeout: job_state = self.result.state if self.result.is_in_terminal_state() or job_state == expected_state: return job_state time.sleep(5) _LOGGER.error(consoleUrl) raise RuntimeError( 'Timeout after %d seconds while waiting for job %s ' 'enters expected state %s. Current state is %s.' % (timeout, self.result.job_id(), expected_state, self.result.state)) ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.6014311 apache_beam-2.74.0/apache_beam/runners/direct/0000755000175100017510000000000015206113160020625 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/__init__.py0000644000175100017510000000203115206113032022730 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Inprocess runner executes pipelines locally in a single process. Anything in this package not imported here is an internal implementation detail with no backwards-compatibility guarantees. """ from apache_beam.runners.direct.direct_runner import DirectRunner ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/bundle_factory.py0000644000175100017510000001664315206113032024207 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A factory that creates UncommittedBundles.""" # pytype: skip-file from typing import Iterable from typing import Iterator from typing import Union from typing import cast from apache_beam import pvalue from apache_beam.runners import common from apache_beam.utils.windowed_value import WindowedValue class BundleFactory(object): """For internal use only; no backwards-compatibility guarantees. BundleFactory creates output bundles to be used by transform evaluators. Args: stacked: whether or not to stack the WindowedValues within the bundle in case consecutive ones share the same timestamp and windows. DirectRunnerOptions.direct_runner_use_stacked_bundle controls this option. """ def __init__(self, stacked: bool) -> None: self._stacked = stacked def create_bundle( self, output_pcollection: Union[pvalue.PBegin, pvalue.PCollection]) -> '_Bundle': return _Bundle(output_pcollection, self._stacked) def create_empty_committed_bundle( self, output_pcollection: Union[pvalue.PBegin, pvalue.PCollection]) -> '_Bundle': bundle = self.create_bundle(output_pcollection) bundle.commit(None) return bundle # a bundle represents a unit of work that will be processed by a transform. class _Bundle(common.Receiver): """Part of a PCollection with output elements. Part of a PCollection. Elements are output to a bundle, which will cause them to be executed by PTransform that consume the PCollection this bundle is a part of at a later point. It starts as an uncommitted bundle and can have elements added to it. It needs to be committed to make it immutable before passing it to a downstream ptransform. The stored elements are WindowedValues, which contains timestamp and windows information. Bundle internally optimizes storage by stacking elements with the same timestamp and windows into StackedWindowedValues, and then returns an iterable to restore WindowedValues upon get_elements() call. When this optimization is not desired, it can be avoided by an option when creating bundles, like::: b = Bundle(stacked=False) """ class _StackedWindowedValues(object): """A stack of WindowedValues with the same timestamp and windows. It must be initialized from a single WindowedValue. Example::: s = StackedWindowedValues(windowed_value) if (another_windowed_value.timestamp == s.timestamp and another_windowed_value.windows == s.windows): s.add_value(another_windowed_value.value) windowed_values = [wv for wv in s.windowed_values()] # now windowed_values equals to [windowed_value, another_windowed_value] """ def __init__(self, initial_windowed_value): self._initial_windowed_value = initial_windowed_value self._appended_values = [] @property def timestamp(self): return self._initial_windowed_value.timestamp @property def windows(self): return self._initial_windowed_value.windows @property def pane_info(self): return self._initial_windowed_value.pane_info def add_value(self, value): self._appended_values.append(value) def windowed_values(self) -> Iterator[WindowedValue]: # yield first windowed_value as is, then iterate through # _appended_values to yield WindowedValue on the fly. yield self._initial_windowed_value for v in self._appended_values: yield self._initial_windowed_value.with_value(v) def __init__( self, pcollection: Union[pvalue.PBegin, pvalue.PCollection], stacked: bool = True) -> None: assert isinstance(pcollection, (pvalue.PBegin, pvalue.PCollection)) self._pcollection = pcollection self._elements: list[Union[WindowedValue, _Bundle._StackedWindowedValues]] = [] self._stacked = stacked self._committed = False self._tag = None # optional tag information for this bundle def get_elements_iterable(self, make_copy: bool = False) -> Iterable[WindowedValue]: """Returns iterable elements. Args: make_copy: whether to force returning copy or yielded iterable. Returns: unstacked elements, in the form of iterable if committed and make_copy is not True, or as a list of copied WindowedValues. """ if not self._stacked: # we can safely assume self._elements contains only WindowedValues elements = cast('list[WindowedValue]', self._elements) if self._committed and not make_copy: return elements return list(elements) def iterable_stacked_or_elements(elements): for e in elements: if isinstance(e, _Bundle._StackedWindowedValues): for w in e.windowed_values(): yield w else: yield e if self._committed and not make_copy: return iterable_stacked_or_elements(self._elements) # returns a copy. return [e for e in iterable_stacked_or_elements(self._elements)] def has_elements(self): return len(self._elements) > 0 @property def tag(self): return self._tag @tag.setter def tag(self, value): assert not self._tag self._tag = value @property def pcollection(self): """PCollection that the elements of this UncommittedBundle belong to.""" return self._pcollection def add(self, element): """Outputs an element to this bundle. Args: element: WindowedValue """ assert not self._committed if not self._stacked: self._elements.append(element) return if (self._elements and (isinstance(self._elements[-1], (WindowedValue, _Bundle._StackedWindowedValues))) and self._elements[-1].timestamp == element.timestamp and self._elements[-1].windows == element.windows and self._elements[-1].pane_info == element.pane_info): if isinstance(self._elements[-1], WindowedValue): self._elements[-1] = _Bundle._StackedWindowedValues(self._elements[-1]) self._elements[-1].add_value(element.value) else: self._elements.append(element) def output(self, element): self.add(element) def receive(self, element: WindowedValue) -> None: self.add(element) def commit(self, synchronized_processing_time): """Commits this bundle. Uncommitted bundle will become committed (immutable) after this call. Args: synchronized_processing_time: the synchronized processing time at which this bundle was committed """ assert not self._committed self._committed = True self._elements = tuple(self._elements) self._synchronized_processing_time = synchronized_processing_time ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/clock.py0000644000175100017510000000312315206113032022267 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Clock implementations for real time processing and testing. For internal use only. No backwards compatibility guarantees. """ # pytype: skip-file import time from apache_beam.utils.timestamp import Timestamp class Clock(object): def time(self): """Returns the number of seconds since epoch.""" raise NotImplementedError() def advance_time(self, advance_by): """Advances the clock by a number of seconds.""" raise NotImplementedError() class RealClock(object): def time(self): return time.time() class TestClock(object): """Clock used for Testing""" def __init__(self, current_time=None): self._current_time = current_time if current_time else Timestamp() def time(self): return float(self._current_time) def advance_time(self, advance_by): self._current_time += advance_by ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/consumer_tracking_pipeline_visitor.py0000644000175100017510000000500515206113032030356 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ConsumerTrackingPipelineVisitor, a PipelineVisitor object.""" # pytype: skip-file from apache_beam import pvalue from apache_beam.pipeline import AppliedPTransform from apache_beam.pipeline import PipelineVisitor class ConsumerTrackingPipelineVisitor(PipelineVisitor): """For internal use only; no backwards-compatibility guarantees. Visitor for extracting value-consumer relations from the graph. Tracks the AppliedPTransforms that consume each PValue in the Pipeline. This is used to schedule consuming PTransforms to consume input after the upstream transform has produced and committed output. """ def __init__(self): self.value_to_consumers: dict[pvalue.PValue, set[AppliedPTransform]] = {} self.root_transforms: set[AppliedPTransform] = set() self.step_names: dict[AppliedPTransform, str] = {} self._num_transforms = 0 self._views = set() @property def views(self): """Returns a list of side intputs extracted from the graph. Returns: A list of pvalue.AsSideInput. """ return list(self._views) def visit_transform(self, applied_ptransform: AppliedPTransform) -> None: inputs = list(applied_ptransform.inputs) if inputs: for input_value in inputs: if isinstance(input_value, pvalue.PBegin): self.root_transforms.add(applied_ptransform) if input_value not in self.value_to_consumers: self.value_to_consumers[input_value] = set() self.value_to_consumers[input_value].add(applied_ptransform) else: self.root_transforms.add(applied_ptransform) self.step_names[applied_ptransform] = 's%d' % (self._num_transforms) self._num_transforms += 1 for side_input in applied_ptransform.side_inputs: self._views.add(side_input) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/consumer_tracking_pipeline_visitor_test.py0000644000175100017510000001332715206113032031423 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for consumer_tracking_pipeline_visitor.""" # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam import pvalue from apache_beam.pipeline import Pipeline from apache_beam.pvalue import AsList from apache_beam.runners.direct import DirectRunner from apache_beam.runners.direct.consumer_tracking_pipeline_visitor import ConsumerTrackingPipelineVisitor from apache_beam.transforms import CoGroupByKey from apache_beam.transforms import Create from apache_beam.transforms import DoFn from apache_beam.transforms import FlatMap from apache_beam.transforms import Flatten from apache_beam.transforms import ParDo # Disable frequent lint warning due to pipe operator for chaining transforms. # pylint: disable=expression-not-assigned # pylint: disable=pointless-statement class ConsumerTrackingPipelineVisitorTest(unittest.TestCase): def setUp(self): self.pipeline = Pipeline(DirectRunner()) self.visitor = ConsumerTrackingPipelineVisitor() def test_root_transforms(self): root_read = beam.Impulse() root_flatten = Flatten(pipeline=self.pipeline) pbegin = pvalue.PBegin(self.pipeline) pcoll_read = pbegin | 'read' >> root_read pcoll_read | FlatMap(lambda x: x) [] | 'flatten' >> root_flatten self.pipeline.visit(self.visitor) root_transforms = [t.transform for t in self.visitor.root_transforms] self.assertCountEqual(root_transforms, [root_read, root_flatten]) pbegin_consumers = [ c.transform for c in self.visitor.value_to_consumers[pbegin] ] self.assertCountEqual(pbegin_consumers, [root_read]) self.assertEqual(len(self.visitor.step_names), 3) def test_side_inputs(self): class SplitNumbersFn(DoFn): def process(self, element): if element < 0: yield pvalue.TaggedOutput('tag_negative', element) else: yield element class ProcessNumbersFn(DoFn): def process(self, element, negatives): yield element def _process_numbers(pcoll, negatives): first_output = ( pcoll | 'process numbers step 1' >> ParDo(ProcessNumbersFn(), negatives)) second_output = ( first_output | 'process numbers step 2' >> ParDo(ProcessNumbersFn(), negatives)) output_pc = ((first_output, second_output) | 'flatten results' >> beam.Flatten()) return output_pc root_read = beam.Impulse() result = ( self.pipeline | 'read' >> root_read | ParDo(SplitNumbersFn()).with_outputs('tag_negative', main='positive')) positive, negative = result _process_numbers(positive, AsList(negative)) self.pipeline.visit(self.visitor) root_transforms = [t.transform for t in self.visitor.root_transforms] self.assertEqual(root_transforms, [root_read]) self.assertEqual(len(self.visitor.step_names), 5) self.assertEqual(len(self.visitor.views), 1) self.assertTrue(isinstance(self.visitor.views[0], pvalue.AsList)) def test_co_group_by_key(self): emails = self.pipeline | 'email' >> Create([('joe', 'joe@example.com')]) phones = self.pipeline | 'phone' >> Create([('mary', '111-222-3333')]) {'emails': emails, 'phones': phones} | CoGroupByKey() self.pipeline.visit(self.visitor) root_transforms = [t.transform for t in self.visitor.root_transforms] self.assertEqual(len(root_transforms), 2) self.assertGreater( len(self.visitor.step_names), 3) # 2 creates + expanded CoGBK self.assertEqual(len(self.visitor.views), 0) def test_visitor_not_sorted(self): p = Pipeline() # pylint: disable=expression-not-assigned from apache_beam.testing.test_stream import TestStream p | TestStream().add_elements(['']) | beam.Map(lambda _: _) original_graph = p.to_runner_api(return_context=False) out_of_order_graph = p.to_runner_api(return_context=False) root_id = out_of_order_graph.root_transform_ids[0] root = out_of_order_graph.components.transforms[root_id] tmp = root.subtransforms[0] root.subtransforms[0] = root.subtransforms[1] root.subtransforms[1] = tmp p = beam.Pipeline().from_runner_api( out_of_order_graph, runner='BundleBasedDirectRunner', options=None) v_out_of_order = ConsumerTrackingPipelineVisitor() p.visit(v_out_of_order) p = beam.Pipeline().from_runner_api( original_graph, runner='BundleBasedDirectRunner', options=None) v_original = ConsumerTrackingPipelineVisitor() p.visit(v_original) # Convert to string to assert they are equal. out_of_order_labels = { str(k): [str(t) for t in value_to_consumer] for k, value_to_consumer in v_out_of_order.value_to_consumers.items() } original_labels = { str(k): [str(t) for t in value_to_consumer] for k, value_to_consumer in v_original.value_to_consumers.items() } self.assertDictEqual(out_of_order_labels, original_labels) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/direct_metrics.py0000644000175100017510000001741415206113032024204 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ DirectRunner implementation of MetricResults. It is in charge not only of responding to queries of current metrics, but also of keeping the common state consistent. """ # pytype: skip-file import threading from collections import defaultdict from typing import Any from typing import SupportsInt from apache_beam.metrics.cells import BoundedTrieData from apache_beam.metrics.cells import DistributionData from apache_beam.metrics.cells import GaugeData from apache_beam.metrics.cells import StringSetData from apache_beam.metrics.execution import MetricKey from apache_beam.metrics.execution import MetricResult from apache_beam.metrics.metric import MetricResults class MetricAggregator(object): """For internal use only; no backwards-compatibility guarantees. Base interface for aggregating metric data during pipeline execution.""" def identity_element(self): # type: () -> Any """Returns the identical element of an Aggregation. For the identity element, it must hold that Aggregator.combine(any_element, identity_element) == any_element. """ raise NotImplementedError def combine(self, x, y): # type: (Any, Any) -> Any raise NotImplementedError def result(self, x): # type: (Any) -> Any raise NotImplementedError class CounterAggregator(MetricAggregator): """For internal use only; no backwards-compatibility guarantees. Aggregator for Counter metric data during pipeline execution. Values aggregated should be ``int`` objects. """ @staticmethod def identity_element(): # type: () -> int return 0 def combine(self, x, y): # type: (SupportsInt, SupportsInt) -> int return int(x) + int(y) def result(self, x): # type: (SupportsInt) -> int return int(x) _IDENTITY_HISTOGRAM = object() class HistogramAggregator(MetricAggregator): @staticmethod def identity_element(): return _IDENTITY_HISTOGRAM def combine(self, x, y): if x is _IDENTITY_HISTOGRAM: return y if y is _IDENTITY_HISTOGRAM: return x return x.combine(y) def result(self, x): if x is _IDENTITY_HISTOGRAM: raise TypeError return x.get_result() class GenericAggregator(MetricAggregator): def __init__(self, data_class): self._data_class = data_class def identity_element(self): return self._data_class.identity_element() def combine(self, x, y): return x.combine(y) def result(self, x): return x.get_result() class DirectMetrics(MetricResults): def __init__(self): self._counters = defaultdict(lambda: DirectMetric(CounterAggregator())) self._distributions = defaultdict( lambda: DirectMetric(GenericAggregator(DistributionData))) self._gauges = defaultdict( lambda: DirectMetric(GenericAggregator(GaugeData))) self._string_sets = defaultdict( lambda: DirectMetric(GenericAggregator(StringSetData))) self._bounded_tries = defaultdict( lambda: DirectMetric(GenericAggregator(BoundedTrieData))) self._histograms = defaultdict(lambda: DirectMetric(HistogramAggregator())) def _apply_operation(self, bundle, updates, op): for k, v in updates.counters.items(): op(self._counters[k], bundle, v) for k, v in updates.distributions.items(): op(self._distributions[k], bundle, v) for k, v in updates.gauges.items(): op(self._gauges[k], bundle, v) for k, v in updates.string_sets.items(): op(self._string_sets[k], bundle, v) for k, v in updates.bounded_tries.items(): op(self._bounded_tries[k], bundle, v) for k, v in updates.histograms.items(): op(self._histograms[k], bundle, v) def commit_logical(self, bundle, updates): op = lambda obj, bundle, update: obj.commit_logical(bundle, update) self._apply_operation(bundle, updates, op) def commit_physical(self, bundle, updates): op = lambda obj, bundle, update: obj.commit_physical(bundle, update) self._apply_operation(bundle, updates, op) def update_physical(self, bundle, updates): op = lambda obj, bundle, update: obj.update_physical(bundle, update) self._apply_operation(bundle, updates, op) def query(self, filter=None): counters = [ MetricResult( MetricKey(k.step, k.metric), v.extract_committed(), v.extract_latest_attempted()) for k, v in self._counters.items() if self.matches(filter, k) ] distributions = [ MetricResult( MetricKey(k.step, k.metric), v.extract_committed(), v.extract_latest_attempted()) for k, v in self._distributions.items() if self.matches(filter, k) ] gauges = [ MetricResult( MetricKey(k.step, k.metric), v.extract_committed(), v.extract_latest_attempted()) for k, v in self._gauges.items() if self.matches(filter, k) ] string_sets = [ MetricResult( MetricKey(k.step, k.metric), v.extract_committed(), v.extract_latest_attempted()) for k, v in self._string_sets.items() if self.matches(filter, k) ] bounded_tries = [ MetricResult( MetricKey(k.step, k.metric), v.extract_committed(), v.extract_latest_attempted()) for k, v in self._bounded_tries.items() if self.matches(filter, k) ] histograms = [ MetricResult( MetricKey(k.step, k.metric), v.extract_committed(), v.extract_latest_attempted()) for k, v in self._histograms.items() if self.matches(filter, k) ] return { self.COUNTERS: counters, self.DISTRIBUTIONS: distributions, self.GAUGES: gauges, self.STRINGSETS: string_sets, self.BOUNDED_TRIES: bounded_tries, self.HISTOGRAMS: histograms, } class DirectMetric(object): """ Keeps a consistent state for a single metric. It keeps track of the metric's physical and logical updates. It's thread safe. """ def __init__(self, aggregator): self.aggregator = aggregator self._attempted_lock = threading.Lock() self.finished_attempted = aggregator.identity_element() self.inflight_attempted = {} self._committed_lock = threading.Lock() self.finished_committed = aggregator.identity_element() def commit_logical(self, bundle, update): with self._committed_lock: self.finished_committed = self.aggregator.combine( update, self.finished_committed) def commit_physical(self, bundle, update): with self._attempted_lock: self.inflight_attempted[bundle] = update self.finished_attempted = self.aggregator.combine( update, self.finished_attempted) del self.inflight_attempted[bundle] def update_physical(self, bundle, update): self.inflight_attempted[bundle] = update def extract_committed(self): return self.aggregator.result(self.finished_committed) def extract_latest_attempted(self): res = self.finished_attempted for _, u in self.inflight_attempted.items(): res = self.aggregator.combine(res, u) return self.aggregator.result(res) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/direct_metrics_test.py0000644000175100017510000002234415206113032025241 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import hamcrest as hc from apache_beam.metrics.cells import DistributionData from apache_beam.metrics.cells import DistributionResult from apache_beam.metrics.execution import MetricKey from apache_beam.metrics.execution import MetricResult from apache_beam.metrics.execution import MetricUpdates from apache_beam.metrics.metricbase import MetricName from apache_beam.runners.direct.direct_metrics import DirectMetrics class DirectMetricsTest(unittest.TestCase): name1 = MetricName('namespace1', 'name1') name2 = MetricName('namespace1', 'name2') name3 = MetricName('namespace2', 'name1') bundle1 = object() # For this test, any object can be a bundle bundle2 = object() def test_combiner_functions(self): metrics = DirectMetrics() counter = metrics._counters['anykey'] counter.commit_logical(self.bundle1, 5) self.assertEqual(counter.extract_committed(), 5) with self.assertRaises(TypeError): counter.commit_logical(self.bundle1, None) distribution = metrics._distributions['anykey'] distribution.commit_logical(self.bundle1, DistributionData(4, 1, 4, 4)) self.assertEqual( distribution.extract_committed(), DistributionResult(DistributionData(4, 1, 4, 4))) with self.assertRaises(AttributeError): distribution.commit_logical(self.bundle1, None) def test_commit_logical_no_filter(self): metrics = DirectMetrics() metrics.commit_logical( self.bundle1, MetricUpdates( counters={ MetricKey('step1', self.name1): 5, MetricKey('step1', self.name2): 8 }, distributions={ MetricKey('step1', self.name1): DistributionData(8, 2, 3, 5) })) metrics.commit_logical( self.bundle1, MetricUpdates( counters={ MetricKey('step2', self.name1): 7, MetricKey('step1', self.name2): 4 }, distributions={ MetricKey('step1', self.name1): DistributionData(4, 1, 4, 4) })) results = metrics.query() hc.assert_that( results['counters'], hc.contains_inanyorder( *[ MetricResult(MetricKey('step1', self.name2), 12, 0), MetricResult(MetricKey('step2', self.name1), 7, 0), MetricResult(MetricKey('step1', self.name1), 5, 0) ])) hc.assert_that( results['distributions'], hc.contains_inanyorder( MetricResult( MetricKey('step1', self.name1), DistributionResult(DistributionData(12, 3, 3, 5)), DistributionResult(DistributionData(0, 0, None, None))))) def test_apply_physical_no_filter(self): metrics = DirectMetrics() metrics.update_physical( object(), MetricUpdates( counters={ MetricKey('step1', self.name1): 5, MetricKey('step1', self.name3): 8 })) metrics.update_physical( object(), MetricUpdates( counters={ MetricKey('step2', self.name1): 7, MetricKey('step1', self.name3): 4 })) results = metrics.query() hc.assert_that( results['counters'], hc.contains_inanyorder( *[ MetricResult(MetricKey('step1', self.name1), 0, 5), MetricResult(MetricKey('step1', self.name3), 0, 12), MetricResult(MetricKey('step2', self.name1), 0, 7) ])) metrics.commit_physical(object(), MetricUpdates()) results = metrics.query() hc.assert_that( results['counters'], hc.contains_inanyorder( *[ MetricResult(MetricKey('step1', self.name1), 0, 5), MetricResult(MetricKey('step1', self.name3), 0, 12), MetricResult(MetricKey('step2', self.name1), 0, 7) ])) def test_apply_physical_logical(self): metrics = DirectMetrics() dist_zero = DistributionData(0, 0, None, None) metrics.update_physical( object(), MetricUpdates( counters={ MetricKey('step1', self.name1): 7, MetricKey('step1', self.name2): 5, MetricKey('step2', self.name1): 1 }, distributions={ MetricKey('step1', self.name1): DistributionData(3, 1, 3, 3), MetricKey('step2', self.name3): DistributionData(8, 2, 4, 4) })) results = metrics.query() hc.assert_that( results['counters'], hc.contains_inanyorder( *[ MetricResult(MetricKey('step1', self.name1), 0, 7), MetricResult(MetricKey('step1', self.name2), 0, 5), MetricResult(MetricKey('step2', self.name1), 0, 1) ])) hc.assert_that( results['distributions'], hc.contains_inanyorder( *[ MetricResult( MetricKey('step1', self.name1), DistributionResult(dist_zero), DistributionResult(DistributionData(3, 1, 3, 3))), MetricResult( MetricKey('step2', self.name3), DistributionResult(dist_zero), DistributionResult(DistributionData(8, 2, 4, 4))) ])) metrics.commit_physical( object(), MetricUpdates( counters={ MetricKey('step1', self.name1): -3, MetricKey('step2', self.name1): -5 }, distributions={ MetricKey('step1', self.name1): DistributionData(8, 4, 1, 5), MetricKey('step2', self.name2): DistributionData(8, 8, 1, 1) })) results = metrics.query() hc.assert_that( results['counters'], hc.contains_inanyorder( *[ MetricResult(MetricKey('step1', self.name1), 0, 4), MetricResult(MetricKey('step1', self.name2), 0, 5), MetricResult(MetricKey('step2', self.name1), 0, -4) ])) hc.assert_that( results['distributions'], hc.contains_inanyorder( *[ MetricResult( MetricKey('step1', self.name1), DistributionResult(dist_zero), DistributionResult(DistributionData(11, 5, 1, 5))), MetricResult( MetricKey('step2', self.name3), DistributionResult(dist_zero), DistributionResult(DistributionData(8, 2, 4, 4))), MetricResult( MetricKey('step2', self.name2), DistributionResult(dist_zero), DistributionResult(DistributionData(8, 8, 1, 1))) ])) metrics.commit_logical( object(), MetricUpdates( counters={ MetricKey('step1', self.name1): 3, MetricKey('step1', self.name2): 5, MetricKey('step2', self.name1): -3 }, distributions={ MetricKey('step1', self.name1): DistributionData(11, 5, 1, 5), MetricKey('step2', self.name2): DistributionData(8, 8, 1, 1), MetricKey('step2', self.name3): DistributionData(4, 1, 4, 4) })) results = metrics.query() hc.assert_that( results['counters'], hc.contains_inanyorder( *[ MetricResult(MetricKey('step1', self.name1), 3, 4), MetricResult(MetricKey('step1', self.name2), 5, 5), MetricResult(MetricKey('step2', self.name1), -3, -4) ])) hc.assert_that( results['distributions'], hc.contains_inanyorder( *[ MetricResult( MetricKey('step1', self.name1), DistributionResult(DistributionData(11, 5, 1, 5)), DistributionResult(DistributionData(11, 5, 1, 5))), MetricResult( MetricKey('step2', self.name3), DistributionResult(DistributionData(4, 1, 4, 4)), DistributionResult(DistributionData(8, 2, 4, 4))), MetricResult( MetricKey('step2', self.name2), DistributionResult(DistributionData(8, 8, 1, 1)), DistributionResult(DistributionData(8, 8, 1, 1))) ])) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/direct_runner.py0000644000175100017510000006666115206113032024057 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """DirectRunner, executing on the local machine. The DirectRunner is a runner implementation that executes the entire graph of transformations belonging to a pipeline on the local machine. """ # pytype: skip-file import itertools import logging import typing from google.protobuf import wrappers_pb2 import apache_beam as beam from apache_beam import coders from apache_beam import typehints from apache_beam.internal.util import ArgumentPlaceholder from apache_beam.options.pipeline_options import DirectOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.value_provider import RuntimeValueProvider from apache_beam.pvalue import PCollection from apache_beam.runners.direct.bundle_factory import BundleFactory from apache_beam.runners.direct.clock import RealClock from apache_beam.runners.direct.clock import TestClock from apache_beam.runners.runner import PipelineResult from apache_beam.runners.runner import PipelineRunner from apache_beam.runners.runner import PipelineState from apache_beam.transforms import userstate from apache_beam.transforms.core import CombinePerKey from apache_beam.transforms.core import CombineValuesDoFn from apache_beam.transforms.core import DoFn from apache_beam.transforms.core import ParDo from apache_beam.transforms.ptransform import PTransform from apache_beam.transforms.timeutil import TimeDomain from apache_beam.typehints import trivial_inference from apache_beam.utils.interactive_utils import is_in_ipython __all__ = ['BundleBasedDirectRunner', 'DirectRunner', 'SwitchingDirectRunner'] _LOGGER = logging.getLogger(__name__) class SwitchingDirectRunner(PipelineRunner): """Executes a single pipeline on the local machine. This implementation switches between using the FnApiRunner (which has high throughput for batch jobs) and using the BundleBasedDirectRunner, which supports streaming execution and certain primitives not yet implemented in the FnApiRunner. """ _is_interactive = False def is_fnapi_compatible(self): return BundleBasedDirectRunner.is_fnapi_compatible() def is_interactive(self): self._is_interactive = True def run_pipeline(self, pipeline, options): from apache_beam.io.gcp.pubsub import ReadFromPubSub from apache_beam.io.gcp.pubsub import WriteToPubSub from apache_beam.pipeline import PipelineVisitor from apache_beam.testing.test_stream import TestStream class _FnApiRunnerSupportVisitor(PipelineVisitor): """Visitor determining if a Pipeline can be run on the FnApiRunner.""" def accept(self, pipeline): self.supported_by_fnapi_runner = True pipeline.visit(self) return self.supported_by_fnapi_runner def enter_composite_transform(self, applied_ptransform): # The FnApiRunner does not support streaming execution. if isinstance(applied_ptransform.transform, (ReadFromPubSub, WriteToPubSub)): self.supported_by_fnapi_runner = False def visit_transform(self, applied_ptransform): transform = applied_ptransform.transform # The FnApiRunner does not support streaming execution. if isinstance(transform, TestStream): self.supported_by_fnapi_runner = False if isinstance(transform, beam.ParDo): dofn = transform.dofn # The FnApiRunner does not support execution of CombineFns with # deferred side inputs. if isinstance(dofn, CombineValuesDoFn): args, kwargs = transform.raw_side_inputs args_to_check = itertools.chain(args, kwargs.values()) if any(isinstance(arg, ArgumentPlaceholder) for arg in args_to_check): self.supported_by_fnapi_runner = False if userstate.is_stateful_dofn(dofn): _, timer_specs = userstate.get_dofn_specs(dofn) for timer in timer_specs: if timer.time_domain == TimeDomain.REAL_TIME: self.supported_by_fnapi_runner = False class _PrismRunnerSupportVisitor(PipelineVisitor): """Visitor determining if a Pipeline can be run on the PrismRunner.""" def accept(self, pipeline, is_interactive): all_options = options.get_all_options() self.supported_by_prism_runner = True # TODO(https://github.com/apache/beam/issues/33623): Prism currently # double fires on AfterCount trigger, once appropriately, and once # incorrectly at the end of the window. This if condition could be # more targeted, but for now we'll just ignore all unsafe triggers. if pipeline.allow_unsafe_triggers: self.supported_by_prism_runner = False # TODO(https://github.com/apache/beam/issues/33623): Prism currently # does not support interactive mode elif is_in_ipython() or is_interactive: self.supported_by_prism_runner = False # TODO(https://github.com/apache/beam/issues/33623): Prism currently # does not support the update compat flag elif all_options['update_compatibility_version']: self.supported_by_prism_runner = False else: pipeline.visit(self) # Avoid circular import from apache_beam.pipeline import ExternalTransformFinder if ExternalTransformFinder.contains_external_transforms(pipeline): # TODO(https://github.com/apache/beam/issues/33623): Prism currently # seems to not be able to consistently bring up external transforms. # It does sometimes, but at volume suites start to fail. We will try # to enable this in a future release. self.supported_by_prism_runner = False return self.supported_by_prism_runner def visit_transform(self, applied_ptransform): transform = applied_ptransform.transform # Python SDK assumes the direct runner TestStream implementation is # being used. if isinstance(transform, TestStream): self.supported_by_prism_runner = False if isinstance(transform, beam.ParDo): dofn = transform.dofn # TODO(https://github.com/apache/beam/issues/33623): Prism currently # does not seem to handle DoFns using exception handling very well. # This may be limited just to subprocess DoFns, but more # investigation is needed before making it default if isinstance(dofn, beam.transforms.core._ExceptionHandlingWrapperDoFn): self.supported_by_prism_runner = False # https://github.com/apache/beam/issues/34549 # Remote once we can support local materialization if (hasattr(dofn, 'is_materialize_values_do_fn') and dofn.is_materialize_values_do_fn): self.supported_by_prism_runner = False # It's uncertain if the Prism Runner supports execution of CombineFns # with deferred side inputs. if isinstance(dofn, CombineValuesDoFn): args, kwargs = transform.raw_side_inputs args_to_check = itertools.chain(args, kwargs.values()) if any(isinstance(arg, ArgumentPlaceholder) for arg in args_to_check): self.supported_by_prism_runner = False if userstate.is_stateful_dofn(dofn): # https://github.com/apache/beam/issues/32786 - # Remove once Real time clock is used. state_specs, timer_specs = userstate.get_dofn_specs(dofn) for timer in timer_specs: if timer.time_domain == TimeDomain.REAL_TIME: self.supported_by_prism_runner = False for state in state_specs: if isinstance(state, userstate.CombiningValueStateSpec): self.supported_by_prism_runner = False if isinstance( dofn, beam.transforms.combiners._PartialGroupByKeyCombiningValues): if len(transform.side_inputs) > 0: # Prism doesn't support side input combiners (this is within spec) self.supported_by_prism_runner = False # TODO(https://github.com/apache/beam/issues/33623): Prism seems to # not handle session windows correctly. Examples are: # util_test.py::ReshuffleTest::test_reshuffle_window_fn_preserved # and util_test.py::ReshuffleTest::test_reshuffle_windows_unchanged if isinstance(transform, beam.WindowInto) and isinstance( transform.get_windowing('').windowfn, beam.window.Sessions): self.supported_by_prism_runner = False # Use BundleBasedDirectRunner if other runners are missing needed features. runner = BundleBasedDirectRunner() # Check whether all transforms used in the pipeline are supported by the # PrismRunner if _PrismRunnerSupportVisitor().accept(pipeline, self._is_interactive): _LOGGER.info('Running pipeline with PrismRunner.') from apache_beam.runners.portability import prism_runner runner = prism_runner.PrismRunner() try: pr = runner.run_pipeline(pipeline, options) # This is non-blocking, so if the state is *already* finished, something # probably failed on job submission. if (PipelineState.is_terminal(pr.state) and pr.state != PipelineState.DONE): _LOGGER.info( 'Pipeline failed on PrismRunner, falling back to DirectRunner.') runner = BundleBasedDirectRunner() else: return pr except Exception as e: # If prism fails in Preparing the portable job, then the PortableRunner # code raises an exception. Catch it, log it, and use the Direct runner # instead. _LOGGER.info('Exception with PrismRunner:\n %s\n' % (e)) _LOGGER.info('Falling back to DirectRunner') runner = BundleBasedDirectRunner() # Check whether all transforms used in the pipeline are supported by the # FnApiRunner, and the pipeline was not meant to be run as streaming. if _FnApiRunnerSupportVisitor().accept(pipeline): from apache_beam.portability.api import beam_provision_api_pb2 from apache_beam.runners.portability.fn_api_runner import fn_runner from apache_beam.runners.portability.portable_runner import JobServiceHandle all_options = options.get_all_options() encoded_options = JobServiceHandle.encode_pipeline_options(all_options) provision_info = fn_runner.ExtendedProvisionInfo( beam_provision_api_pb2.ProvisionInfo( pipeline_options=encoded_options)) runner = fn_runner.FnApiRunner(provision_info=provision_info) return runner.run_pipeline(pipeline, options) # Type variables. K = typing.TypeVar('K') V = typing.TypeVar('V') @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types(tuple[K, typing.Iterable[V]]) class _GroupByKeyOnly(PTransform): """A group by key transform, ignoring windows.""" def infer_output_type(self, input_type): key_type, value_type = trivial_inference.key_value_types(input_type) return typehints.KV[key_type, typehints.Iterable[value_type]] def expand(self, pcoll): self._check_pcollection(pcoll) return PCollection.from_(pcoll) @typehints.with_input_types(tuple[K, typing.Iterable[V]]) @typehints.with_output_types(tuple[K, typing.Iterable[V]]) class _GroupAlsoByWindow(ParDo): """The GroupAlsoByWindow transform.""" def __init__(self, windowing): super().__init__(_GroupAlsoByWindowDoFn(windowing)) self.windowing = windowing def expand(self, pcoll): self._check_pcollection(pcoll) return PCollection.from_(pcoll) class _GroupAlsoByWindowDoFn(DoFn): # TODO(robertwb): Support combiner lifting. def __init__(self, windowing): super().__init__() self.windowing = windowing def infer_output_type(self, input_type): key_type, windowed_value_iter_type = trivial_inference.key_value_types( input_type) value_type = windowed_value_iter_type.inner_type.inner_type return typehints.KV[key_type, typehints.Iterable[value_type]] def start_bundle(self): # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.transforms.trigger import create_trigger_driver # pylint: enable=wrong-import-order, wrong-import-position self.driver = create_trigger_driver(self.windowing, True) def process(self, element): k, vs = element return self.driver.process_entire_key(k, vs) @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types(tuple[K, typing.Iterable[V]]) class _StreamingGroupByKeyOnly(_GroupByKeyOnly): """Streaming GroupByKeyOnly placeholder for overriding in DirectRunner.""" urn = "direct_runner:streaming_gbko:v0.1" # These are needed due to apply overloads. def to_runner_api_parameter(self, unused_context): return _StreamingGroupByKeyOnly.urn, None @staticmethod @PTransform.register_urn(urn, None) def from_runner_api_parameter( unused_ptransform, unused_payload, unused_context): return _StreamingGroupByKeyOnly() @typehints.with_input_types(tuple[K, typing.Iterable[V]]) @typehints.with_output_types(tuple[K, typing.Iterable[V]]) class _StreamingGroupAlsoByWindow(_GroupAlsoByWindow): """Streaming GroupAlsoByWindow placeholder for overriding in DirectRunner.""" urn = "direct_runner:streaming_gabw:v0.1" # These are needed due to apply overloads. def to_runner_api_parameter(self, context): return ( _StreamingGroupAlsoByWindow.urn, wrappers_pb2.BytesValue( value=context.windowing_strategies.get_id(self.windowing))) @staticmethod @PTransform.register_urn(urn, wrappers_pb2.BytesValue) def from_runner_api_parameter(unused_ptransform, payload, context): return _StreamingGroupAlsoByWindow( context.windowing_strategies.get_by_id(payload.value)) @typehints.with_input_types(tuple[K, typing.Iterable[V]]) @typehints.with_output_types(tuple[K, typing.Iterable[V]]) class _GroupByKey(PTransform): """The DirectRunner GroupByKey implementation.""" def expand(self, pcoll): # Imported here to avoid circular dependencies. # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.coders import typecoders input_type = pcoll.element_type if input_type is not None: # Initialize type-hints used below to enforce type-checking and to # pass downstream to further PTransforms. key_type, value_type = trivial_inference.key_value_types(input_type) # Enforce the input to a GBK has a KV element type. pcoll.element_type = typehints.typehints.coerce_to_kv_type( pcoll.element_type) typecoders.registry.verify_deterministic( typecoders.registry.get_coder(key_type), 'GroupByKey operation "%s"' % self.label) reify_output_type = typehints.KV[ key_type, typehints.WindowedValue[value_type]] # type: ignore[misc] gbk_input_type = ( typehints.KV[ key_type, typehints.Iterable[typehints.WindowedValue[ # type: ignore[misc] value_type]]]) gbk_output_type = typehints.KV[key_type, typehints.Iterable[value_type]] # pylint: disable=bad-option-value return ( pcoll | 'ReifyWindows' >> ( ParDo(beam.GroupByKey.ReifyWindows()).with_output_types( reify_output_type)) | 'GroupByKey' >> ( _GroupByKeyOnly().with_input_types( reify_output_type).with_output_types(gbk_input_type)) | ( 'GroupByWindow' >> _GroupAlsoByWindow(pcoll.windowing).with_input_types( gbk_input_type).with_output_types(gbk_output_type))) else: # The input_type is None, run the default return ( pcoll | 'ReifyWindows' >> ParDo(beam.GroupByKey.ReifyWindows()) | 'GroupByKey' >> _GroupByKeyOnly() | 'GroupByWindow' >> _GroupAlsoByWindow(pcoll.windowing)) def _get_transform_overrides(pipeline_options): # A list of PTransformOverride objects to be applied before running a pipeline # using DirectRunner. # Currently this only works for overrides where the input and output types do # not change. # For internal use only; no backwards-compatibility guarantees. # Importing following locally to avoid a circular dependency. from apache_beam.pipeline import PTransformOverride from apache_beam.runners.direct.sdf_direct_runner import ProcessKeyedElementsViaKeyedWorkItemsOverride from apache_beam.runners.direct.sdf_direct_runner import SplittableParDoOverride from apache_beam.transforms.combiners import LiftedCombinePerKey class CombinePerKeyOverride(PTransformOverride): def matches(self, applied_ptransform): if isinstance(applied_ptransform.transform, CombinePerKey): return applied_ptransform.inputs[0].windowing.is_default() def get_replacement_transform_for_applied_ptransform( self, applied_ptransform): # TODO: Move imports to top. Pipeline <-> Runner dependency cause problems # with resolving imports when they are at top. # pylint: disable=wrong-import-position try: transform = applied_ptransform.transform return LiftedCombinePerKey( transform.fn, transform.args, transform.kwargs) except NotImplementedError: return transform class StreamingGroupByKeyOverride(PTransformOverride): def matches(self, applied_ptransform): # Note: we match the exact class, since we replace it with a subclass. return applied_ptransform.transform.__class__ == _GroupByKeyOnly def get_replacement_transform_for_applied_ptransform( self, applied_ptransform): # Use specialized streaming implementation. transform = _StreamingGroupByKeyOnly() return transform class StreamingGroupAlsoByWindowOverride(PTransformOverride): def matches(self, applied_ptransform): # Note: we match the exact class, since we replace it with a subclass. transform = applied_ptransform.transform return ( isinstance(applied_ptransform.transform, ParDo) and isinstance(transform.dofn, _GroupAlsoByWindowDoFn) and transform.__class__ != _StreamingGroupAlsoByWindow) def get_replacement_transform_for_applied_ptransform( self, applied_ptransform): # Use specialized streaming implementation. transform = _StreamingGroupAlsoByWindow( applied_ptransform.transform.dofn.windowing) return transform class TestStreamOverride(PTransformOverride): def matches(self, applied_ptransform): from apache_beam.testing.test_stream import TestStream self.applied_ptransform = applied_ptransform return isinstance(applied_ptransform.transform, TestStream) def get_replacement_transform_for_applied_ptransform( self, applied_ptransform): from apache_beam.runners.direct.test_stream_impl import _ExpandableTestStream return _ExpandableTestStream(applied_ptransform.transform) class GroupByKeyPTransformOverride(PTransformOverride): """A ``PTransformOverride`` for ``GroupByKey``. This replaces the Beam implementation as a primitive. """ def matches(self, applied_ptransform): # Imported here to avoid circular dependencies. # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.transforms.core import GroupByKey return isinstance(applied_ptransform.transform, GroupByKey) def get_replacement_transform_for_applied_ptransform( self, applied_ptransform): return _GroupByKey() overrides = [ # This needs to be the first and the last override. Other overrides depend # on the GroupByKey implementation to be composed of _GroupByKeyOnly and # _GroupAlsoByWindow. GroupByKeyPTransformOverride(), SplittableParDoOverride(), ProcessKeyedElementsViaKeyedWorkItemsOverride(), CombinePerKeyOverride(), TestStreamOverride(), ] # Add streaming overrides, if necessary. if pipeline_options.view_as(StandardOptions).streaming: overrides.append(StreamingGroupByKeyOverride()) overrides.append(StreamingGroupAlsoByWindowOverride()) # Add PubSub overrides, if PubSub is available. try: from apache_beam.io.gcp import pubsub as unused_pubsub overrides += _get_pubsub_transform_overrides(pipeline_options) except ImportError: pass # This also needs to be last because other transforms apply GBKs which need to # be translated into a DirectRunner-compatible transform. overrides.append(GroupByKeyPTransformOverride()) return overrides class _DirectReadFromPubSub(PTransform): def __init__(self, source): self._source = source def _infer_output_coder( self, unused_input_type=None, unused_input_coder=None) -> typing.Optional[coders.Coder]: return coders.BytesCoder() def get_windowing(self, unused_inputs): return beam.Windowing(beam.window.GlobalWindows()) def expand(self, pvalue): # This is handled as a native transform. return PCollection(self.pipeline, is_bounded=self._source.is_bounded()) def _get_pubsub_transform_overrides(pipeline_options): from apache_beam.io.gcp import pubsub as beam_pubsub from apache_beam.pipeline import PTransformOverride class ReadFromPubSubOverride(PTransformOverride): def matches(self, applied_ptransform): return isinstance( applied_ptransform.transform, beam_pubsub.ReadFromPubSub) def get_replacement_transform_for_applied_ptransform( self, applied_ptransform): if not pipeline_options.view_as(StandardOptions).streaming: raise Exception( 'PubSub I/O is only available in streaming mode ' '(use the --streaming flag).') return _DirectReadFromPubSub(applied_ptransform.transform._source) # WriteToPubSub no longer needs an override - it works by default for both # batch and streaming return [ReadFromPubSubOverride()] class BundleBasedDirectRunner(PipelineRunner): """Executes a single pipeline on the local machine.""" @staticmethod def is_fnapi_compatible(): return False def run_pipeline(self, pipeline, options): """Execute the entire pipeline and returns an DirectPipelineResult.""" # TODO: Move imports to top. Pipeline <-> Runner dependency cause problems # with resolving imports when they are at top. # pylint: disable=wrong-import-position from apache_beam.pipeline import PipelineVisitor from apache_beam.runners.direct.consumer_tracking_pipeline_visitor import ConsumerTrackingPipelineVisitor from apache_beam.runners.direct.evaluation_context import EvaluationContext from apache_beam.runners.direct.executor import Executor from apache_beam.runners.direct.transform_evaluator import TransformEvaluatorRegistry from apache_beam.testing.test_stream import TestStream from apache_beam.transforms.external import ExternalTransform class VerifyNoCrossLanguageTransforms(PipelineVisitor): """Visitor determining whether a Pipeline uses a TestStream.""" def visit_transform(self, applied_ptransform): if isinstance(applied_ptransform.transform, ExternalTransform): raise RuntimeError( "Streaming Python direct runner " "does not support cross-language pipelines." "Please use other runners such as FlinkRunner, " "DataflowRunner, or PrismRunner.") pipeline.visit(VerifyNoCrossLanguageTransforms()) # If the TestStream I/O is used, use a mock test clock. class TestStreamUsageVisitor(PipelineVisitor): """Visitor determining whether a Pipeline uses a TestStream.""" def __init__(self): self.uses_test_stream = False def visit_transform(self, applied_ptransform): if isinstance(applied_ptransform.transform, TestStream): self.uses_test_stream = True visitor = TestStreamUsageVisitor() pipeline.visit(visitor) clock = TestClock() if visitor.uses_test_stream else RealClock() # Performing configured PTransform overrides. pipeline.replace_all(_get_transform_overrides(options)) _LOGGER.info('Running pipeline with DirectRunner.') self.consumer_tracking_visitor = ConsumerTrackingPipelineVisitor() pipeline.visit(self.consumer_tracking_visitor) evaluation_context = EvaluationContext( options, BundleFactory( stacked=options.view_as( DirectOptions).direct_runner_use_stacked_bundle), self.consumer_tracking_visitor.root_transforms, self.consumer_tracking_visitor.value_to_consumers, self.consumer_tracking_visitor.step_names, self.consumer_tracking_visitor.views, clock) executor = Executor( self.consumer_tracking_visitor.value_to_consumers, TransformEvaluatorRegistry(evaluation_context), evaluation_context) # DirectRunner does not support injecting # PipelineOptions values at runtime RuntimeValueProvider.set_runtime_options({}) # Start the executor. This is a non-blocking call, it will start the # execution in background threads and return. executor.start(self.consumer_tracking_visitor.root_transforms) result = DirectPipelineResult(executor, evaluation_context) return result # Use the SwitchingDirectRunner as the default. DirectRunner = SwitchingDirectRunner class DirectPipelineResult(PipelineResult): """A DirectPipelineResult provides access to info about a pipeline.""" def __init__(self, executor, evaluation_context): super().__init__(PipelineState.RUNNING) self._executor = executor self._evaluation_context = evaluation_context def __del__(self): if self._state == PipelineState.RUNNING: _LOGGER.warning( 'The DirectPipelineResult is being garbage-collected while the ' 'DirectRunner is still running the corresponding pipeline. This may ' 'lead to incomplete execution of the pipeline if the main thread ' 'exits before pipeline completion. Consider using ' 'result.wait_until_finish() to wait for completion of pipeline ' 'execution.') def wait_until_finish(self, duration=None): if not PipelineState.is_terminal(self.state): if duration: raise NotImplementedError( 'DirectRunner does not support duration argument.') try: self._executor.await_completion() self._state = PipelineState.DONE except: # pylint: disable=broad-except self._state = PipelineState.FAILED raise return self._state def aggregated_values(self, aggregator_or_name): return self._evaluation_context.get_aggregator_values(aggregator_or_name) def metrics(self): return self._evaluation_context.metrics() def cancel(self): """Shuts down pipeline workers. For testing use only. Does not properly wait for pipeline workers to shut down. """ self._state = PipelineState.CANCELLING self._executor.shutdown() self._state = PipelineState.CANCELLED ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/direct_runner_test.py0000644000175100017510000002253115206113032025102 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import threading import unittest from collections import defaultdict import hamcrest as hc import apache_beam as beam from apache_beam.metrics.cells import DistributionData from apache_beam.metrics.cells import DistributionResult from apache_beam.metrics.execution import MetricKey from apache_beam.metrics.execution import MetricResult from apache_beam.metrics.metric import Metrics from apache_beam.metrics.metric import MetricsFilter from apache_beam.metrics.metricbase import MetricName from apache_beam.pipeline import Pipeline from apache_beam.runners import DirectRunner from apache_beam.runners import TestDirectRunner from apache_beam.runners import create_runner from apache_beam.runners.direct.evaluation_context import _ExecutionContext from apache_beam.runners.direct.transform_evaluator import _GroupByKeyOnlyEvaluator from apache_beam.runners.direct.transform_evaluator import _TransformEvaluator from apache_beam.testing import test_pipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.utils import shared class DictWrapper(dict): pass def acquire_dict(): return DictWrapper() class DirectPipelineResultTest(unittest.TestCase): def test_waiting_on_result_stops_executor_threads(self): pre_test_threads = set(t.ident for t in threading.enumerate()) for runner in [ 'BundleBasedDirectRunner', 'apache_beam.runners.portability.fn_api_runner.fn_runner.FnApiRunner' ]: pipeline = test_pipeline.TestPipeline(runner=runner) _ = (pipeline | beam.Create([{'foo': 'bar'}])) result = pipeline.run() result.wait_until_finish() post_test_threads = set(t.ident for t in threading.enumerate()) new_threads = post_test_threads - pre_test_threads self.assertEqual(len(new_threads), 0) def test_direct_runner_metrics(self): class MyDoFn(beam.DoFn): def start_bundle(self): count = Metrics.counter(self.__class__, 'bundles') count.inc() def finish_bundle(self): count = Metrics.counter(self.__class__, 'finished_bundles') count.inc() def process(self, element): gauge = Metrics.gauge(self.__class__, 'latest_element') gauge.set(element) count = Metrics.counter(self.__class__, 'elements') count.inc() distro = Metrics.distribution(self.__class__, 'element_dist') distro.update(element) str_set = Metrics.string_set(self.__class__, 'element_str_set') str_set.add(str(element % 4)) Metrics.bounded_trie(self.__class__, 'element_bounded_trie').add( ("a", "b", str(element % 4))) return [element] # TODO(https://github.com/apache/beam/issues/34549): This test relies on # metrics filtering which doesn't work on Prism yet because Prism renames # steps (e.g. "Do" becomes "ref_AppliedPTransform_Do_7"). # https://github.com/apache/beam/blob/5f9cd73b7c9a2f37f83971ace3a399d633201dd1/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py#L1590 p = Pipeline('FnApiRunner') pcoll = ( p | beam.Create([1, 2, 3, 4, 5], reshuffle=False) | 'Do' >> beam.ParDo(MyDoFn())) assert_that(pcoll, equal_to([1, 2, 3, 4, 5])) result = p.run() result.wait_until_finish() metrics = result.metrics().query(MetricsFilter().with_step('Do')) namespace = '{}.{}'.format(MyDoFn.__module__, MyDoFn.__name__) hc.assert_that( metrics['counters'], hc.contains_inanyorder( MetricResult( MetricKey('Do', MetricName(namespace, 'elements')), 5, 5), MetricResult( MetricKey('Do', MetricName(namespace, 'bundles')), 1, 1), MetricResult( MetricKey('Do', MetricName(namespace, 'finished_bundles')), 1, 1))) hc.assert_that( metrics['distributions'], hc.contains_inanyorder( MetricResult( MetricKey('Do', MetricName(namespace, 'element_dist')), DistributionResult(DistributionData(15, 5, 1, 5)), DistributionResult(DistributionData(15, 5, 1, 5))))) gauge_result = metrics['gauges'][0] hc.assert_that( gauge_result.key, hc.equal_to(MetricKey('Do', MetricName(namespace, 'latest_element')))) hc.assert_that(gauge_result.committed.value, hc.equal_to(5)) hc.assert_that(gauge_result.attempted.value, hc.equal_to(5)) str_set_result = metrics['string_sets'][0] hc.assert_that( str_set_result.key, hc.equal_to(MetricKey('Do', MetricName(namespace, 'element_str_set')))) hc.assert_that(len(str_set_result.committed), hc.equal_to(4)) hc.assert_that(len(str_set_result.attempted), hc.equal_to(4)) bounded_trie_results = metrics['bounded_tries'][0] hc.assert_that( bounded_trie_results.key, hc.equal_to( MetricKey('Do', MetricName(namespace, 'element_bounded_trie')))) hc.assert_that(bounded_trie_results.committed.size(), hc.equal_to(4)) hc.assert_that(bounded_trie_results.attempted.size(), hc.equal_to(4)) def test_create_runner(self): self.assertTrue(isinstance(create_runner('DirectRunner'), DirectRunner)) self.assertTrue( isinstance(create_runner('TestDirectRunner'), TestDirectRunner)) class BundleBasedRunnerTest(unittest.TestCase): def test_type_hints(self): with test_pipeline.TestPipeline(runner='BundleBasedDirectRunner') as p: _ = ( p | beam.Create([[]]).with_output_types(beam.typehints.List[int]) | beam.combiners.Count.Globally()) def test_impulse(self): with test_pipeline.TestPipeline(runner='BundleBasedDirectRunner') as p: assert_that(p | beam.Impulse(), equal_to([b''])) class DirectRunnerRetryTests(unittest.TestCase): def test_retry_fork_graph(self): # TODO(https://github.com/apache/beam/issues/18640): The FnApiRunner # currently does not currently support retries. p = beam.Pipeline(runner='BundleBasedDirectRunner') shared_handler = shared.Shared() counts = shared_handler.acquire(acquire_dict) counts['count_b'] = 0 counts['count_c'] = 0 def f_b(x): shared_handler.acquire(acquire_dict)['count_b'] += 1 raise Exception('exception in f_b') def f_c(x): shared_handler.acquire(acquire_dict)['count_c'] += 1 raise Exception('exception in f_c') names = p | 'CreateNodeA' >> beam.Create(['Ann', 'Joe']) fork_b = names | 'SendToB' >> beam.Map(f_b) # pylint: disable=unused-variable fork_c = names | 'SendToC' >> beam.Map(f_c) # pylint: disable=unused-variable with self.assertRaises(Exception): p.run().wait_until_finish() assert counts['count_b'] == counts['count_c'] == 4 def test_no_partial_writeouts(self): class TestTransformEvaluator(_TransformEvaluator): def __init__(self): self._execution_context = _ExecutionContext(None, {}) def start_bundle(self): self.step_context = self._execution_context.get_step_context() def process_element(self, element): k, v = element state = self.step_context.get_keyed_state(k) state.add_state(None, _GroupByKeyOnlyEvaluator.ELEMENTS_TAG, v) # Create instance and add key/value, key/value2 evaluator = TestTransformEvaluator() evaluator.start_bundle() self.assertIsNone(evaluator.step_context.existing_keyed_state.get('key')) self.assertIsNone(evaluator.step_context.partial_keyed_state.get('key')) evaluator.process_element(['key', 'value']) self.assertEqual( evaluator.step_context.existing_keyed_state['key'].state, defaultdict(lambda: defaultdict(list))) self.assertEqual( evaluator.step_context.partial_keyed_state['key'].state, {None: { 'elements': ['value'] }}) evaluator.process_element(['key', 'value2']) self.assertEqual( evaluator.step_context.existing_keyed_state['key'].state, defaultdict(lambda: defaultdict(list))) self.assertEqual( evaluator.step_context.partial_keyed_state['key'].state, {None: { 'elements': ['value', 'value2'] }}) # Simulate an exception (redo key/value) evaluator._execution_context.reset() evaluator.start_bundle() evaluator.process_element(['key', 'value']) self.assertEqual( evaluator.step_context.existing_keyed_state['key'].state, defaultdict(lambda: defaultdict(list))) self.assertEqual( evaluator.step_context.partial_keyed_state['key'].state, {None: { 'elements': ['value'] }}) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/direct_userstate.py0000644000175100017510000002603215206113032024551 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Support for user state in the BundleBasedDirectRunner.""" # pytype: skip-file import copy import itertools from apache_beam.transforms import userstate from apache_beam.transforms.trigger import _ListStateTag from apache_beam.transforms.trigger import _ReadModifyWriteStateTag from apache_beam.transforms.trigger import _SetStateTag class DirectRuntimeState(userstate.RuntimeState): def __init__(self, state_spec, state_tag, current_value_accessor): self._state_spec = state_spec self._state_tag = state_tag self._current_value_accessor = current_value_accessor @staticmethod def for_spec(state_spec, state_tag, current_value_accessor): if isinstance(state_spec, userstate.ReadModifyWriteStateSpec): return ReadModifyWriteRuntimeState( state_spec, state_tag, current_value_accessor) elif isinstance(state_spec, userstate.BagStateSpec): return BagRuntimeState(state_spec, state_tag, current_value_accessor) elif isinstance(state_spec, userstate.CombiningValueStateSpec): return CombiningValueRuntimeState( state_spec, state_tag, current_value_accessor) elif isinstance(state_spec, userstate.SetStateSpec): return SetRuntimeState(state_spec, state_tag, current_value_accessor) else: raise ValueError('Invalid state spec: %s' % state_spec) def _encode(self, value): return self._state_spec.coder.encode(value) def _decode(self, value): return self._state_spec.coder.decode(value) # Sentinel designating an unread value. UNREAD_VALUE = object() class ReadModifyWriteRuntimeState(DirectRuntimeState, userstate.ReadModifyWriteRuntimeState): def __init__(self, state_spec, state_tag, current_value_accessor): super().__init__(state_spec, state_tag, current_value_accessor) self._value = UNREAD_VALUE self._cleared = False self._modified = False def read(self): if self._cleared: return None if self._value is UNREAD_VALUE: self._value = self._current_value_accessor() if not self._value: return None return self._decode(self._value[0]) def write(self, value): self._cleared = False self._modified = True self._value = [self._encode(value)] def clear(self): self._cleared = True self._modified = False self._value = [] def is_cleared(self): return self._cleared def is_modified(self): return self._modified class BagRuntimeState(DirectRuntimeState, userstate.BagRuntimeState): def __init__(self, state_spec, state_tag, current_value_accessor): super().__init__(state_spec, state_tag, current_value_accessor) self._cached_value = UNREAD_VALUE self._cleared = False self._new_values = [] def read(self): if self._cached_value is UNREAD_VALUE: self._cached_value = self._current_value_accessor() if not self._cleared: encoded_values = itertools.chain(self._cached_value, self._new_values) else: encoded_values = self._new_values return (self._decode(v) for v in encoded_values) def add(self, value): self._new_values.append(self._encode(value)) def clear(self): self._cleared = True self._cached_value = [] self._new_values = [] class SetRuntimeState(DirectRuntimeState, userstate.SetRuntimeState): def __init__(self, state_spec, state_tag, current_value_accessor): super().__init__(state_spec, state_tag, current_value_accessor) self._current_accumulator = UNREAD_VALUE self._modified = False def _read_initial_value(self): if self._current_accumulator is UNREAD_VALUE: self._current_accumulator = { self._decode(a) for a in self._current_value_accessor() } def read(self): self._read_initial_value() return self._current_accumulator def add(self, value): self._read_initial_value() self._modified = True self._current_accumulator.add(value) def clear(self): self._current_accumulator = set() self._modified = True def is_modified(self): return self._modified class CombiningValueRuntimeState(DirectRuntimeState, userstate.CombiningValueRuntimeState): """Combining value state interface object passed to user code.""" def __init__(self, state_spec, state_tag, current_value_accessor): super().__init__(state_spec, state_tag, current_value_accessor) self._current_accumulator = UNREAD_VALUE self._modified = False self._combine_fn = copy.deepcopy(state_spec.combine_fn) self._combine_fn.setup() self._finalized = False def _read_initial_value(self): if self._current_accumulator is UNREAD_VALUE: existing_accumulators = list( self._decode(a) for a in self._current_value_accessor()) if existing_accumulators: self._current_accumulator = self._combine_fn.merge_accumulators( existing_accumulators) else: self._current_accumulator = self._combine_fn.create_accumulator() def read(self): self._read_initial_value() return self._combine_fn.extract_output(self._current_accumulator) def add(self, value): self._read_initial_value() self._modified = True self._current_accumulator = self._combine_fn.add_input( self._current_accumulator, value) def clear(self): self._modified = True self._current_accumulator = self._combine_fn.create_accumulator() def finalize(self): if not self._finalized: self._combine_fn.teardown() self._finalized = True class DirectUserStateContext(userstate.UserStateContext): """userstate.UserStateContext for the BundleBasedDirectRunner. The DirectUserStateContext buffers up updates that are to be committed by the TransformEvaluator after running a DoFn. """ def __init__(self, step_context, dofn, key_coder): self.step_context = step_context self.dofn = dofn self.key_coder = key_coder self.all_state_specs, self.all_timer_specs = userstate.get_dofn_specs(dofn) self.state_tags = {} for state_spec in self.all_state_specs: state_key = 'user/%s' % state_spec.name if isinstance(state_spec, userstate.ReadModifyWriteStateSpec): state_tag = _ReadModifyWriteStateTag(state_key) elif isinstance(state_spec, userstate.BagStateSpec): state_tag = _ListStateTag(state_key) elif isinstance(state_spec, userstate.CombiningValueStateSpec): state_tag = _ListStateTag(state_key) elif isinstance(state_spec, userstate.SetStateSpec): state_tag = _SetStateTag(state_key) else: raise ValueError('Invalid state spec: %s' % state_spec) self.state_tags[state_spec] = state_tag self.cached_states = {} self.cached_timers = {} def get_timer( self, timer_spec: userstate.TimerSpec, key, window, timestamp, pane) -> userstate.RuntimeTimer: assert timer_spec in self.all_timer_specs encoded_key = self.key_coder.encode(key) cache_key = (encoded_key, window, timer_spec) if cache_key not in self.cached_timers: self.cached_timers[cache_key] = userstate.RuntimeTimer() return self.cached_timers[cache_key] def get_state(self, state_spec, key, window): assert state_spec in self.all_state_specs encoded_key = self.key_coder.encode(key) cache_key = (encoded_key, window, state_spec) if cache_key not in self.cached_states: state_tag = self.state_tags[state_spec] value_accessor = ( lambda: self._get_underlying_state(state_spec, key, window)) self.cached_states[cache_key] = DirectRuntimeState.for_spec( state_spec, state_tag, value_accessor) return self.cached_states[cache_key] def _get_underlying_state(self, state_spec, key, window): state_tag = self.state_tags[state_spec] encoded_key = self.key_coder.encode(key) return ( self.step_context.get_keyed_state(encoded_key).get_state( window, state_tag)) def commit(self): # Commit state modifications. for cache_key, runtime_state in self.cached_states.items(): encoded_key, window, state_spec = cache_key state = self.step_context.get_keyed_state(encoded_key) state_tag = self.state_tags[state_spec] if isinstance(state_spec, userstate.BagStateSpec): if runtime_state._cleared: state.clear_state(window, state_tag) for new_value in runtime_state._new_values: state.add_state(window, state_tag, new_value) elif isinstance(state_spec, userstate.CombiningValueStateSpec): if runtime_state._modified: state.clear_state(window, state_tag) state.add_state( window, state_tag, state_spec.coder.encode(runtime_state._current_accumulator)) elif isinstance(state_spec, userstate.SetStateSpec): if runtime_state.is_modified(): state.clear_state(window, state_tag) for new_value in runtime_state._current_accumulator: state.add_state( window, state_tag, state_spec.coder.encode(new_value)) elif isinstance(state_spec, userstate.ReadModifyWriteStateSpec): if runtime_state.is_cleared(): state.clear_state(window, state_tag) if runtime_state.is_modified(): state.clear_state(window, state_tag) state.add_state(window, state_tag, runtime_state._value) else: raise ValueError('Invalid state spec: %s' % state_spec) # Commit new timers. for cache_key, runtime_timer in self.cached_timers.items(): encoded_key, window, timer_spec = cache_key state = self.step_context.get_keyed_state(encoded_key) timer_name = 'user/%s' % timer_spec.name for dynamic_timer_tag, timer in runtime_timer._timer_recordings.items(): if timer.cleared: state.clear_timer( window, timer_name, timer_spec.time_domain, dynamic_timer_tag=dynamic_timer_tag) if timer.timestamp: # TODO(ccy): add corresponding watermark holds after the DirectRunner # allows for keyed watermark holds. state.set_timer( window, timer_name, timer_spec.time_domain, timer.timestamp, dynamic_timer_tag=dynamic_timer_tag) def reset(self): for state in self.cached_states.values(): state.finalize() self.cached_states = {} self.cached_timers = {} ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/evaluation_context.py0000644000175100017510000004200315206113032025107 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """EvaluationContext tracks global state, triggers and watermarks.""" # pytype: skip-file import collections import threading from typing import TYPE_CHECKING from typing import Any from typing import Iterable from typing import Optional from typing import Union from apache_beam import pvalue from apache_beam.pipeline import AppliedPTransform from apache_beam.runners.direct.direct_metrics import DirectMetrics from apache_beam.runners.direct.executor import TransformExecutor from apache_beam.runners.direct.watermark_manager import WatermarkManager from apache_beam.transforms import sideinputs from apache_beam.transforms.trigger import InMemoryUnmergedState from apache_beam.utils import counters from apache_beam.utils.timestamp import Timestamp if TYPE_CHECKING: from apache_beam.runners.direct.bundle_factory import BundleFactory from apache_beam.runners.direct.bundle_factory import _Bundle from apache_beam.runners.direct.util import TimerFiring from apache_beam.runners.direct.util import TransformResult from apache_beam.runners.direct.watermark_manager import _TransformWatermarks class _ExecutionContext(object): """Contains the context for the execution of a single PTransform. It holds the watermarks for that transform, as well as keyed states. """ def __init__(self, watermarks: '_TransformWatermarks', keyed_states): self.watermarks = watermarks self.keyed_states = keyed_states self._step_context = None def get_step_context(self): if not self._step_context: self._step_context = DirectStepContext(self.keyed_states) return self._step_context def reset(self): self._step_context = None class _SideInputView(object): def __init__(self, view): self._view = view self.blocked_tasks = collections.deque() self.elements = [] self.value = None self.watermark = None def __repr__(self): elements_string = ( ', '.join(str(elm) for elm in self.elements) if self.elements else '[]') return '_SideInputView(elements=%s)' % elements_string class _SideInputsContainer(object): """An in-process container for side inputs. It provides methods for blocking until a side-input is available and writing to a side input. """ def __init__(self, side_inputs: Iterable['pvalue.AsSideInput']) -> None: self._lock = threading.Lock() self._views: dict[pvalue.AsSideInput, _SideInputView] = {} self._transform_to_side_inputs: collections.defaultdict[ Optional[AppliedPTransform], list[pvalue.AsSideInput]] = collections.defaultdict(list) # this appears unused: self._side_input_to_blocked_tasks = collections.defaultdict(list) # type: ignore for side in side_inputs: self._views[side] = _SideInputView(side) self._transform_to_side_inputs[side.pvalue.producer].append(side) def __repr__(self): views_string = ( ', '.join(str(elm) for elm in self._views.values()) if self._views else '[]') return '_SideInputsContainer(_views=%s)' % views_string def get_value_or_block_until_ready( self, side_input, task: TransformExecutor, block_until: Timestamp) -> Any: """Returns the value of a view whose task is unblocked or blocks its task. It gets the value of a view whose watermark has been updated and surpasses a given value. Args: side_input: ``_UnpickledSideInput`` value. task: ``TransformExecutor`` task waiting on a side input. block_until: Timestamp after which the task gets unblocked. Returns: The ``SideInputMap`` value of a view when the tasks it blocks are unblocked. Otherwise, None. """ with self._lock: view = self._views[side_input] if view.watermark and view.watermark.output_watermark >= block_until: view.value = self._pvalue_to_value(side_input, view.elements) return view.value else: view.blocked_tasks.append((task, block_until)) task.blocked = True def add_values(self, side_input, values): with self._lock: view = self._views[side_input] view.elements.extend(values) def update_watermarks_for_transform_and_unblock_tasks( self, ptransform, watermark) -> list[tuple[TransformExecutor, Timestamp]]: """Updates _SideInputsContainer after a watermark update and unbloks tasks. It traverses the list of side inputs per PTransform and calls _update_watermarks_for_side_input_and_unblock_tasks to unblock tasks. Args: ptransform: Value of a PTransform. watermark: Value of the watermark after an update for a PTransform. Returns: Tasks that get unblocked as a result of the watermark advancing. """ unblocked_tasks = [] for side in self._transform_to_side_inputs[ptransform]: unblocked_tasks.extend( self._update_watermarks_for_side_input_and_unblock_tasks( side, watermark)) return unblocked_tasks def _update_watermarks_for_side_input_and_unblock_tasks( self, side_input, watermark) -> list[tuple[TransformExecutor, Timestamp]]: """Helps update _SideInputsContainer after a watermark update. For each view of the side input, it updates the value of the watermark recorded when the watermark moved and unblocks tasks accordingly. Args: side_input: ``_UnpickledSideInput`` value. watermark: Value of the watermark after an update for a PTransform. Returns: Tasks that get unblocked as a result of the watermark advancing. """ with self._lock: view = self._views[side_input] view.watermark = watermark unblocked_tasks = [] tasks_just_unblocked = [] for task, block_until in view.blocked_tasks: if watermark.output_watermark >= block_until: view.value = self._pvalue_to_value(side_input, view.elements) unblocked_tasks.append(task) tasks_just_unblocked.append((task, block_until)) task.blocked = False for task in tasks_just_unblocked: view.blocked_tasks.remove(task) return unblocked_tasks def _pvalue_to_value(self, side_input, values): """Given a side input, returns the associated value in its requested form. Args: side_input: _UnpickledSideInput object. values: Iterable values associated with the side input. Returns: The side input in its requested form. Raises: ValueError: If values cannot be converted into the requested form. """ return sideinputs.SideInputMap( type(side_input), side_input._view_options(), values) class EvaluationContext(object): """Evaluation context with the global state information of the pipeline. The evaluation context for a specific pipeline being executed by the DirectRunner. Contains state shared within the execution across all transforms. EvaluationContext contains shared state for an execution of the DirectRunner that can be used while evaluating a PTransform. This consists of views into underlying state and watermark implementations, access to read and write side inputs, and constructing counter sets and execution contexts. This includes executing callbacks asynchronously when state changes to the appropriate point (e.g. when a side input is requested and known to be empty). EvaluationContext also handles results by committing finalizing bundles based on the current global state and updating the global state appropriately. This includes updating the per-(step,key) state, updating global watermarks, and executing any callbacks that can be executed. """ def __init__( self, pipeline_options, bundle_factory: 'BundleFactory', root_transforms, value_to_consumers, step_names, views: Iterable[pvalue.AsSideInput], clock): self.pipeline_options = pipeline_options self._bundle_factory = bundle_factory self._root_transforms = root_transforms self._value_to_consumers = value_to_consumers self._step_names = step_names self.views = views self._pcollection_to_views: collections.defaultdict[ pvalue.PValue, list[pvalue.AsSideInput]] = collections.defaultdict(list) for view in views: self._pcollection_to_views[view.pvalue].append(view) self._transform_keyed_states = self._initialize_keyed_states( root_transforms, value_to_consumers) self._side_inputs_container = _SideInputsContainer(views) self._watermark_manager = WatermarkManager( clock, root_transforms, value_to_consumers, self._transform_keyed_states) self._pending_unblocked_tasks: list[tuple[TransformExecutor, Timestamp]] = [] self._counter_factory = counters.CounterFactory() self._metrics = DirectMetrics() self._lock = threading.Lock() self.shutdown_requested = False def _initialize_keyed_states(self, root_transforms, value_to_consumers): """Initialize user state dicts. These dicts track user state per-key, per-transform and per-window. """ transform_keyed_states = {} for transform in root_transforms: transform_keyed_states[transform] = {} for consumers in value_to_consumers.values(): for consumer in consumers: transform_keyed_states[consumer] = {} return transform_keyed_states def metrics(self): # TODO. Should this be made a @property? return self._metrics def is_root_transform(self, applied_ptransform: AppliedPTransform) -> bool: return applied_ptransform in self._root_transforms def handle_result( self, completed_bundle: '_Bundle', completed_timers, result: 'TransformResult'): """Handle the provided result produced after evaluating the input bundle. Handle the provided TransformResult, produced after evaluating the provided committed bundle (potentially None, if the result of a root PTransform). The result is the output of running the transform contained in the TransformResult on the contents of the provided bundle. Args: completed_bundle: the bundle that was processed to produce the result. completed_timers: the timers that were delivered to produce the completed_bundle. result: the ``TransformResult`` of evaluating the input bundle Returns: the committed bundles contained within the handled result. """ with self._lock: committed_bundles, unprocessed_bundles = self._commit_bundles( result.uncommitted_output_bundles, result.unprocessed_bundles) self._metrics.commit_logical( completed_bundle, result.logical_metric_updates) # If the result is for a view, update side inputs container. self._update_side_inputs_container(committed_bundles, result) # Tasks generated from unblocked side inputs as the watermark progresses. tasks = self._watermark_manager.update_watermarks( completed_bundle, result.transform, completed_timers, committed_bundles, unprocessed_bundles, result.keyed_watermark_holds, self._side_inputs_container) self._pending_unblocked_tasks.extend(tasks) if result.counters: for counter in result.counters: merged_counter = self._counter_factory.get_counter( counter.name, counter.combine_fn) merged_counter.accumulator.merge([counter.accumulator]) # Commit partial GBK states existing_keyed_state = self._transform_keyed_states[result.transform] for k, v in result.partial_keyed_state.items(): existing_keyed_state[k] = v return committed_bundles def _update_side_inputs_container( self, committed_bundles: Iterable['_Bundle'], result: 'TransformResult'): """Update the side inputs container if we are outputting into a side input. Look at the result, and if it's outputing into a PCollection that we have registered as a PCollectionView, we add the result to the PCollectionView. """ if (result.uncommitted_output_bundles and result.uncommitted_output_bundles[0].pcollection in self._pcollection_to_views): for view in self._pcollection_to_views[ result.uncommitted_output_bundles[0].pcollection]: for committed_bundle in committed_bundles: # side_input must be materialized. self._side_inputs_container.add_values( view, committed_bundle.get_elements_iterable(make_copy=True)) def get_aggregator_values(self, aggregator_or_name): return self._counter_factory.get_aggregator_values(aggregator_or_name) def schedule_pending_unblocked_tasks(self, executor_service): if self._pending_unblocked_tasks: with self._lock: for task in self._pending_unblocked_tasks: executor_service.submit(task) self._pending_unblocked_tasks = [] def _commit_bundles( self, uncommitted_bundles: Iterable['_Bundle'], unprocessed_bundles: Iterable['_Bundle'] ) -> tuple[tuple['_Bundle', ...], tuple['_Bundle', ...]]: """Commits bundles and returns a immutable set of committed bundles.""" for in_progress_bundle in uncommitted_bundles: producing_applied_ptransform = in_progress_bundle.pcollection.producer watermarks = self._watermark_manager.get_watermarks( producing_applied_ptransform) in_progress_bundle.commit(watermarks.synchronized_processing_output_time) for unprocessed_bundle in unprocessed_bundles: unprocessed_bundle.commit(None) return tuple(uncommitted_bundles), tuple(unprocessed_bundles) def get_execution_context( self, applied_ptransform: AppliedPTransform) -> _ExecutionContext: return _ExecutionContext( self._watermark_manager.get_watermarks(applied_ptransform), self._transform_keyed_states[applied_ptransform]) def create_bundle( self, output_pcollection: Union[pvalue.PBegin, pvalue.PCollection]) -> '_Bundle': """Create an uncommitted bundle for the specified PCollection.""" return self._bundle_factory.create_bundle(output_pcollection) def create_empty_committed_bundle( self, output_pcollection: pvalue.PCollection) -> '_Bundle': """Create empty bundle useful for triggering evaluation.""" return self._bundle_factory.create_empty_committed_bundle( output_pcollection) def extract_all_timers( self) -> tuple[list[tuple[AppliedPTransform, list['TimerFiring']]], bool]: return self._watermark_manager.extract_all_timers() def is_done(self, transform: Optional[AppliedPTransform] = None) -> bool: """Checks completion of a step or the pipeline. Args: transform: AppliedPTransform to check for completion. Returns: True if the step will not produce additional output. If transform is None returns true if all steps are done. """ if transform: return self._is_transform_done(transform) for applied_ptransform in self._step_names: if not self._is_transform_done(applied_ptransform): return False return True def _is_transform_done(self, transform: AppliedPTransform) -> bool: tw = self._watermark_manager.get_watermarks(transform) return tw.output_watermark == WatermarkManager.WATERMARK_POS_INF def get_value_or_block_until_ready(self, side_input, task, block_until): assert isinstance(task, TransformExecutor) return self._side_inputs_container.get_value_or_block_until_ready( side_input, task, block_until) def shutdown(self): self.shutdown_requested = True class DirectUnmergedState(InMemoryUnmergedState): """UnmergedState implementation for the DirectRunner.""" def __init__(self): super().__init__(defensive_copy=False) class DirectStepContext(object): """Context for the currently-executing step.""" def __init__(self, existing_keyed_state): self.existing_keyed_state = existing_keyed_state # In order to avoid partial writes of a bundle, every time # existing_keyed_state is accessed, a copy of the state is made # to be transferred to the bundle state once the bundle is committed. self.partial_keyed_state = {} def get_keyed_state(self, key): if not self.existing_keyed_state.get(key): self.existing_keyed_state[key] = DirectUnmergedState() if not self.partial_keyed_state.get(key): self.partial_keyed_state[key] = self.existing_keyed_state[key].copy() return self.partial_keyed_state[key] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/executor.py0000644000175100017510000006246515206113032023050 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """An executor that schedules and executes applied ptransforms.""" # pytype: skip-file import collections import itertools import logging import queue import threading import traceback from typing import TYPE_CHECKING from typing import Any from typing import Optional from weakref import WeakValueDictionary from apache_beam.metrics.execution import MetricsContainer from apache_beam.runners.worker import statesampler from apache_beam.transforms import sideinputs from apache_beam.utils import counters if TYPE_CHECKING: from apache_beam import pvalue from apache_beam.runners.direct.bundle_factory import _Bundle from apache_beam.runners.direct.evaluation_context import EvaluationContext from apache_beam.runners.direct.transform_evaluator import TransformEvaluatorRegistry _LOGGER = logging.getLogger(__name__) class _ExecutorService(object): """Thread pool for executing tasks in parallel.""" class CallableTask(object): def call(self, state_sampler): pass @property def name(self): return None class _ExecutorServiceWorker(threading.Thread): """Worker thread for executing a single task at a time.""" # Amount to block waiting for getting an item from the queue in seconds. TIMEOUT = 5 def __init__( self, queue: 'queue.Queue[_ExecutorService.CallableTask]', index): super().__init__() self.queue = queue self._index = index self._default_name = 'ExecutorServiceWorker-' + str(index) self._update_name() self.shutdown_requested = False # Stop worker thread when main thread exits. self.daemon = True self.start() def _update_name(self, task=None): if task and task.name: name = task.name else: name = self._default_name self.name = 'Thread: %d, %s (%s)' % ( self._index, name, 'executing' if task else 'idle') def _get_task_or_none(self) -> Optional['_ExecutorService.CallableTask']: try: # Do not block indefinitely, otherwise we may not act for a requested # shutdown. return self.queue.get( timeout=_ExecutorService._ExecutorServiceWorker.TIMEOUT) except queue.Empty: return None def run(self): state_sampler = statesampler.StateSampler('', counters.CounterFactory()) statesampler.set_current_tracker(state_sampler) while not self.shutdown_requested: task = self._get_task_or_none() if task: try: if not self.shutdown_requested: self._update_name(task) task.call(state_sampler) self._update_name() finally: self.queue.task_done() def shutdown(self): self.shutdown_requested = True def __init__(self, num_workers): self.queue: queue.Queue[_ExecutorService.CallableTask] = queue.Queue() self.workers = [ _ExecutorService._ExecutorServiceWorker(self.queue, i) for i in range(num_workers) ] self.shutdown_requested = False def submit(self, task: '_ExecutorService.CallableTask') -> None: assert isinstance(task, _ExecutorService.CallableTask) if not self.shutdown_requested: self.queue.put(task) def await_completion(self): for worker in self.workers: worker.join() def shutdown(self): self.shutdown_requested = True for worker in self.workers: worker.shutdown() # Consume all the remaining items in the queue while not self.queue.empty(): try: self.queue.get_nowait() self.queue.task_done() except queue.Empty: continue # All existing threads will eventually terminate (after they complete their # last task). class _TransformEvaluationState(object): def __init__(self, executor_service, scheduled: set['TransformExecutor']): self.executor_service = executor_service self.scheduled = scheduled def schedule(self, work): self.scheduled.add(work) self.executor_service.submit(work) def complete(self, completed_work): self.scheduled.remove(completed_work) class _ParallelEvaluationState(_TransformEvaluationState): """A TransformEvaluationState with unlimited parallelism. Any TransformExecutor scheduled will be immediately submitted to the ExecutorService. A principal use of this is for evaluators that can generate output bundles only using the input bundle (e.g. ParDo). """ pass class _SerialEvaluationState(_TransformEvaluationState): """A TransformEvaluationState with a single work queue. Any TransformExecutor scheduled will be placed on the work queue. Only one item of work will be submitted to the ExecutorService at any time. A principal use of this is for evaluators that keeps a global state such as _GroupByKeyOnly. """ def __init__(self, executor_service, scheduled): super().__init__(executor_service, scheduled) self.serial_queue = collections.deque() self.currently_evaluating = None self._lock = threading.Lock() def complete(self, completed_work): self._update_currently_evaluating(None, completed_work) super().complete(completed_work) def schedule(self, new_work): self._update_currently_evaluating(new_work, None) def _update_currently_evaluating(self, new_work, completed_work): with self._lock: if new_work: self.serial_queue.append(new_work) if completed_work: assert self.currently_evaluating == completed_work self.currently_evaluating = None if self.serial_queue and not self.currently_evaluating: next_work = self.serial_queue.pop() self.currently_evaluating = next_work super().schedule(next_work) class _TransformExecutorServices(object): """Schedules and completes TransformExecutors. Controls the concurrency as appropriate for the applied transform the executor exists for. """ def __init__(self, executor_service: _ExecutorService) -> None: self._executor_service = executor_service self._scheduled: set[TransformExecutor] = set() self._parallel = _ParallelEvaluationState( self._executor_service, self._scheduled) self._serial_cache: WeakValueDictionary[ Any, _SerialEvaluationState] = WeakValueDictionary() def parallel(self) -> _ParallelEvaluationState: return self._parallel def serial(self, step: Any) -> _SerialEvaluationState: cached = self._serial_cache.get(step) if not cached: cached = _SerialEvaluationState(self._executor_service, self._scheduled) self._serial_cache[step] = cached return cached @property def executors(self) -> frozenset['TransformExecutor']: return frozenset(self._scheduled) class _CompletionCallback(object): """The default completion callback. The default completion callback is used to complete transform evaluations that are triggered due to the arrival of elements from an upstream transform, or for a source transform. """ def __init__( self, evaluation_context: 'EvaluationContext', all_updates, timer_firings=None): self._evaluation_context = evaluation_context self._all_updates = all_updates self._timer_firings = timer_firings or [] def handle_result( self, transform_executor, input_committed_bundle, transform_result): output_committed_bundles = self._evaluation_context.handle_result( input_committed_bundle, self._timer_firings, transform_result) for output_committed_bundle in output_committed_bundles: self._all_updates.offer( _ExecutorServiceParallelExecutor._ExecutorUpdate( transform_executor, committed_bundle=output_committed_bundle)) for unprocessed_bundle in transform_result.unprocessed_bundles: self._all_updates.offer( _ExecutorServiceParallelExecutor._ExecutorUpdate( transform_executor, unprocessed_bundle=unprocessed_bundle)) return output_committed_bundles def handle_exception(self, transform_executor, exception): self._all_updates.offer( _ExecutorServiceParallelExecutor._ExecutorUpdate( transform_executor, exception=exception)) class TransformExecutor(_ExecutorService.CallableTask): """For internal use only; no backwards-compatibility guarantees. TransformExecutor will evaluate a bundle using an applied ptransform. A CallableTask responsible for constructing a TransformEvaluator and evaluating it on some bundle of input, and registering the result using the completion callback. """ _MAX_RETRY_PER_BUNDLE = 4 def __init__( self, transform_evaluator_registry: 'TransformEvaluatorRegistry', evaluation_context: 'EvaluationContext', input_bundle: '_Bundle', fired_timers, applied_ptransform, completion_callback, transform_evaluation_state: _TransformEvaluationState): self._transform_evaluator_registry = transform_evaluator_registry self._evaluation_context = evaluation_context self._input_bundle = input_bundle # For non-empty bundles, store the window of the max EOW. # TODO(mariagh): Move to class _Bundle's inner _StackedWindowedValues self._latest_main_input_window = None if input_bundle.has_elements(): self._latest_main_input_window = input_bundle._elements[0].windows[0] for elem in input_bundle.get_elements_iterable(): if elem.windows[0].end > self._latest_main_input_window.end: self._latest_main_input_window = elem.windows[0] self._fired_timers = fired_timers self._applied_ptransform = applied_ptransform self._completion_callback = completion_callback self._transform_evaluation_state = transform_evaluation_state self._side_input_values: dict[pvalue.AsSideInput, Any] = {} self.blocked = False self._call_count = 0 self._retry_count = 0 self._max_retries_per_bundle = TransformExecutor._MAX_RETRY_PER_BUNDLE def call(self, state_sampler): self._call_count += 1 assert self._call_count <= (1 + len(self._applied_ptransform.side_inputs)) metrics_container = MetricsContainer(self._applied_ptransform.full_label) start_state = state_sampler.scoped_state( self._applied_ptransform.full_label, 'start', metrics_container=metrics_container) process_state = state_sampler.scoped_state( self._applied_ptransform.full_label, 'process', metrics_container=metrics_container) finish_state = state_sampler.scoped_state( self._applied_ptransform.full_label, 'finish', metrics_container=metrics_container) with start_state: # Side input initialization should be accounted for in start_state. for side_input in self._applied_ptransform.side_inputs: # Find the projection of main's window onto the side input's window. window_mapping_fn = side_input._view_options().get( 'window_mapping_fn', sideinputs._global_window_mapping_fn) main_onto_side_window = window_mapping_fn( self._latest_main_input_window) block_until = main_onto_side_window.end if side_input not in self._side_input_values: value = self._evaluation_context.get_value_or_block_until_ready( side_input, self, block_until) if not value: # Monitor task will reschedule this executor once the side input is # available. return self._side_input_values[side_input] = value side_input_values = [ self._side_input_values[side_input] for side_input in self._applied_ptransform.side_inputs ] while self._retry_count < self._max_retries_per_bundle: try: self.attempt_call( metrics_container, side_input_values, start_state, process_state, finish_state) break except Exception as e: self._retry_count += 1 _LOGGER.error( 'Exception at bundle %r, due to an exception.\n %s', self._input_bundle, traceback.format_exc()) if self._retry_count == self._max_retries_per_bundle: _LOGGER.error( 'Giving up after %s attempts.', self._max_retries_per_bundle) self._completion_callback.handle_exception(self, e) self._evaluation_context.metrics().commit_physical( self._input_bundle, metrics_container.get_cumulative()) self._transform_evaluation_state.complete(self) def attempt_call( self, metrics_container, side_input_values, start_state, process_state, finish_state): """Attempts to run a bundle.""" evaluator = self._transform_evaluator_registry.get_evaluator( self._applied_ptransform, self._input_bundle, side_input_values) with start_state: evaluator.start_bundle() with process_state: if self._fired_timers: for timer_firing in self._fired_timers: evaluator.process_timer_wrapper(timer_firing) if self._input_bundle: for value in self._input_bundle.get_elements_iterable(): evaluator.process_element(value) with finish_state: result = evaluator.finish_bundle() result.logical_metric_updates = metrics_container.get_cumulative() self._completion_callback.handle_result(self, self._input_bundle, result) return result class Executor(object): """For internal use only; no backwards-compatibility guarantees.""" def __init__(self, *args, **kwargs): self._executor = _ExecutorServiceParallelExecutor(*args, **kwargs) def start(self, roots): self._executor.start(roots) def await_completion(self): self._executor.await_completion() def shutdown(self): self._executor.request_shutdown() class _ExecutorServiceParallelExecutor(object): """An internal implementation for Executor.""" NUM_WORKERS = 1 def __init__( self, value_to_consumers, transform_evaluator_registry, evaluation_context: 'EvaluationContext'): self.executor_service = _ExecutorService( _ExecutorServiceParallelExecutor.NUM_WORKERS) self.transform_executor_services = _TransformExecutorServices( self.executor_service) self.value_to_consumers = value_to_consumers self.transform_evaluator_registry = transform_evaluator_registry self.evaluation_context = evaluation_context self.all_updates = _ExecutorServiceParallelExecutor._TypedUpdateQueue( _ExecutorServiceParallelExecutor._ExecutorUpdate) self.visible_updates = _ExecutorServiceParallelExecutor._TypedUpdateQueue( _ExecutorServiceParallelExecutor._VisibleExecutorUpdate) self.default_completion_callback = _CompletionCallback( evaluation_context, self.all_updates) def start(self, roots): self.root_nodes = frozenset(roots) self.all_nodes = frozenset( itertools.chain( roots, *itertools.chain(self.value_to_consumers.values()))) self.node_to_pending_bundles = {} for root_node in self.root_nodes: provider = ( self.transform_evaluator_registry.get_root_bundle_provider(root_node)) self.node_to_pending_bundles[root_node] = provider.get_root_bundles() self.executor_service.submit( _ExecutorServiceParallelExecutor._MonitorTask(self)) def await_completion(self): update = self.visible_updates.take() try: if update.exception: raise update.exception finally: self.executor_service.shutdown() self.executor_service.await_completion() def request_shutdown(self): self.executor_service.shutdown() self.executor_service.await_completion() self.evaluation_context.shutdown() def schedule_consumers(self, committed_bundle: '_Bundle') -> None: if committed_bundle.pcollection in self.value_to_consumers: consumers = self.value_to_consumers[committed_bundle.pcollection] for applied_ptransform in consumers: self.schedule_consumption( applied_ptransform, committed_bundle, [], self.default_completion_callback) def schedule_unprocessed_bundle(self, applied_ptransform, unprocessed_bundle): self.node_to_pending_bundles[applied_ptransform].append(unprocessed_bundle) def schedule_consumption( self, consumer_applied_ptransform, committed_bundle: '_Bundle', fired_timers, on_complete): """Schedules evaluation of the given bundle with the transform.""" assert consumer_applied_ptransform assert committed_bundle assert on_complete if self.transform_evaluator_registry.should_execute_serially( consumer_applied_ptransform): transform_executor_service: _TransformEvaluationState = ( self.transform_executor_services.serial(consumer_applied_ptransform)) else: transform_executor_service = self.transform_executor_services.parallel() transform_executor = TransformExecutor( self.transform_evaluator_registry, self.evaluation_context, committed_bundle, fired_timers, consumer_applied_ptransform, on_complete, transform_executor_service) transform_executor_service.schedule(transform_executor) class _TypedUpdateQueue(object): """Type checking update queue with blocking and non-blocking operations.""" def __init__(self, item_type): self._item_type = item_type self._queue = queue.Queue() def poll(self): try: item = self._queue.get_nowait() self._queue.task_done() return item except queue.Empty: return None def take(self): # The implementation of Queue.Queue.get() does not propagate # KeyboardInterrupts when a timeout is not used. We therefore use a # one-second timeout in the following loop to allow KeyboardInterrupts # to be correctly propagated. while True: try: item = self._queue.get(timeout=1) self._queue.task_done() return item except queue.Empty: pass def offer(self, item): assert isinstance(item, self._item_type) self._queue.put_nowait(item) class _ExecutorUpdate(object): """An internal status update on the state of the executor.""" def __init__( self, transform_executor, committed_bundle=None, unprocessed_bundle=None, exception=None): self.transform_executor = transform_executor # Exactly one of them should be not-None assert sum( [bool(committed_bundle), bool(unprocessed_bundle), bool(exception)]) == 1 self.committed_bundle = committed_bundle self.unprocessed_bundle = unprocessed_bundle self.exception = exception class _VisibleExecutorUpdate(object): """An update of interest to the user. Used for awaiting the completion to decide whether to return normally or raise an exception. """ def __init__(self, exception=None): self.finished = exception is not None self.exception = exception class _MonitorTask(_ExecutorService.CallableTask): """MonitorTask continuously runs to ensure that pipeline makes progress.""" def __init__(self, executor: '_ExecutorServiceParallelExecutor') -> None: self._executor = executor @property def name(self): return 'monitor' def call(self, state_sampler): try: update = self._executor.all_updates.poll() while update: if update.committed_bundle: self._executor.schedule_consumers(update.committed_bundle) elif update.unprocessed_bundle: self._executor.schedule_unprocessed_bundle( update.transform_executor._applied_ptransform, update.unprocessed_bundle) else: assert update.exception _LOGGER.warning( 'A task failed with exception: %s', update.exception) self._executor.visible_updates.offer( _ExecutorServiceParallelExecutor._VisibleExecutorUpdate( update.exception)) update = self._executor.all_updates.poll() self._executor.evaluation_context.schedule_pending_unblocked_tasks( self._executor.executor_service) self._add_work_if_necessary(self._fire_timers()) except Exception as e: # pylint: disable=broad-except _LOGGER.error('Monitor task died due to exception.\n %s', e) self._executor.visible_updates.offer( _ExecutorServiceParallelExecutor._VisibleExecutorUpdate(e)) finally: if not self._should_shutdown(): self._executor.executor_service.submit(self) def _should_shutdown(self) -> bool: """Checks whether the pipeline is completed and should be shut down. If there is anything in the queue of tasks to do or if there are any realtime timers set, do not shut down. Otherwise, check if all the transforms' watermarks are complete. If they are not, the pipeline is not progressing (stall detected). Whether the pipeline has stalled or not, the executor should shut down the pipeline. Returns: True only if the pipeline has reached a terminal state and should be shut down. """ if self._is_executing(): # There are some bundles still in progress. return False watermark_manager = self._executor.evaluation_context._watermark_manager _, any_unfired_realtime_timers = watermark_manager.extract_all_timers() if any_unfired_realtime_timers: return False else: if self._executor.evaluation_context.is_done(): self._executor.visible_updates.offer( _ExecutorServiceParallelExecutor._VisibleExecutorUpdate()) else: # Nothing is scheduled for execution, but watermarks incomplete. self._executor.visible_updates.offer( _ExecutorServiceParallelExecutor._VisibleExecutorUpdate(( Exception('Monitor task detected a pipeline stall.'), None, None))) self._executor.executor_service.shutdown() return True def _fire_timers(self): """Schedules triggered consumers if any timers fired. Returns: True if timers fired. """ transform_fired_timers, _ = ( self._executor.evaluation_context.extract_all_timers()) for applied_ptransform, fired_timers in transform_fired_timers: # Use an empty committed bundle. just to trigger. empty_bundle = ( self._executor.evaluation_context.create_empty_committed_bundle( applied_ptransform.inputs[0])) timer_completion_callback = _CompletionCallback( self._executor.evaluation_context, self._executor.all_updates, timer_firings=fired_timers) self._executor.schedule_consumption( applied_ptransform, empty_bundle, fired_timers, timer_completion_callback) return bool(transform_fired_timers) def _is_executing(self) -> bool: """Checks whether the job is still executing. Returns: True if there is at least one non-blocked TransformExecutor active.""" executors = self._executor.transform_executor_services.executors if not executors: # Nothing is executing. return False # Ensure that at least one of those executors is not blocked. for transform_executor in executors: if not transform_executor.blocked: return True return False def _add_work_if_necessary(self, timers_fired): """Adds more work from the roots if pipeline requires more input. If all active TransformExecutors are in a blocked state, add more work from root nodes that may have additional work. This ensures that if a pipeline has elements available from the root nodes it will add those elements when necessary. Args: timers_fired: True if any timers fired prior to this call. """ # If any timers have fired, they will add more work; No need to add more. if timers_fired: return if self._is_executing(): # We have at least one executor that can proceed without adding # additional work. return # All current TransformExecutors are blocked; add more work from any # pending bundles. for applied_ptransform in self._executor.all_nodes: if not self._executor.evaluation_context.is_done(applied_ptransform): pending_bundles = self._executor.node_to_pending_bundles.get( applied_ptransform, []) for bundle in pending_bundles: self._executor.schedule_consumption( applied_ptransform, bundle, [], self._executor.default_completion_callback) self._executor.node_to_pending_bundles[applied_ptransform] = [] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/sdf_direct_runner.py0000644000175100017510000005123015206113032024675 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """This module contains Splittable DoFn logic that is specific to DirectRunner. """ # pytype: skip-file import uuid from threading import Lock from threading import Timer from typing import Any from typing import Iterable from typing import Optional import apache_beam as beam from apache_beam import TimeDomain from apache_beam import pvalue from apache_beam.coders import typecoders from apache_beam.io.iobase import WatermarkEstimator from apache_beam.pipeline import AppliedPTransform from apache_beam.pipeline import PTransformOverride from apache_beam.runners.common import DoFnContext from apache_beam.runners.common import DoFnInvoker from apache_beam.runners.common import DoFnSignature from apache_beam.runners.common import OutputHandler from apache_beam.runners.direct.evaluation_context import DirectStepContext from apache_beam.runners.direct.util import KeyedWorkItem from apache_beam.runners.direct.watermark_manager import WatermarkManager from apache_beam.transforms.core import ParDo from apache_beam.transforms.core import ProcessContinuation from apache_beam.transforms.ptransform import PTransform from apache_beam.transforms.trigger import _ReadModifyWriteStateTag from apache_beam.utils.windowed_value import WindowedValue class SplittableParDoOverride(PTransformOverride): """A transform override for ParDo transformss of SplittableDoFns. Replaces the ParDo transform with a SplittableParDo transform that performs SDF specific logic. """ def matches(self, applied_ptransform): assert isinstance(applied_ptransform, AppliedPTransform) transform = applied_ptransform.transform if isinstance(transform, ParDo): signature = DoFnSignature(transform.fn) return signature.is_splittable_dofn() def get_replacement_transform_for_applied_ptransform( self, applied_ptransform): ptransform = applied_ptransform.transform assert isinstance(ptransform, ParDo) do_fn = ptransform.fn signature = DoFnSignature(do_fn) if signature.is_splittable_dofn(): return SplittableParDo(ptransform) else: return ptransform class SplittableParDo(PTransform): """A transform that processes a PCollection using a Splittable DoFn.""" def __init__(self, ptransform): assert isinstance(ptransform, ParDo) self._ptransform = ptransform def expand(self, pcoll): sdf = self._ptransform.fn signature = DoFnSignature(sdf) restriction_coder = signature.get_restriction_coder() element_coder = typecoders.registry.get_coder(pcoll.element_type) keyed_elements = ( pcoll | 'pair' >> ParDo(PairWithRestrictionFn(sdf)) | 'split' >> ParDo(SplitRestrictionFn(sdf)) | 'explode' >> ParDo(ExplodeWindowsFn()) | 'random' >> ParDo(RandomUniqueKeyFn())) return keyed_elements | ProcessKeyedElements( sdf, element_coder, restriction_coder, pcoll.windowing, self._ptransform.args, self._ptransform.kwargs, self._ptransform.side_inputs) class ElementAndRestriction(object): """A holder for an element, restriction, and watermark estimator state.""" def __init__(self, element, restriction, watermark_estimator_state): self.element = element self.restriction = restriction self.watermark_estimator_state = watermark_estimator_state class PairWithRestrictionFn(beam.DoFn): """A transform that pairs each element with a restriction.""" def __init__(self, do_fn): self._signature = DoFnSignature(do_fn) def start_bundle(self): self._invoker = DoFnInvoker.create_invoker( self._signature, output_handler=_NoneShallPassOutputHandler(), process_invocation=False) def process(self, element, window=beam.DoFn.WindowParam, *args, **kwargs): initial_restriction = self._invoker.invoke_initial_restriction(element) watermark_estimator_state = ( self._signature.process_method.watermark_estimator_provider. initial_estimator_state(element, initial_restriction)) yield ElementAndRestriction( element, initial_restriction, watermark_estimator_state) class SplitRestrictionFn(beam.DoFn): """A transform that perform initial splitting of Splittable DoFn inputs.""" def __init__(self, do_fn): self._do_fn = do_fn def start_bundle(self): signature = DoFnSignature(self._do_fn) self._invoker = DoFnInvoker.create_invoker( signature, output_handler=_NoneShallPassOutputHandler(), process_invocation=False) def process(self, element_and_restriction, *args, **kwargs): element = element_and_restriction.element restriction = element_and_restriction.restriction restriction_parts = self._invoker.invoke_split(element, restriction) for part in restriction_parts: yield ElementAndRestriction( element, part, element_and_restriction.watermark_estimator_state) class ExplodeWindowsFn(beam.DoFn): """A transform that forces the runner to explode windows. This is done to make sure that Splittable DoFn proceses an element for each of the windows that element belongs to. """ def process(self, element, window=beam.DoFn.WindowParam, *args, **kwargs): yield element class RandomUniqueKeyFn(beam.DoFn): """A transform that assigns a unique key to each element.""" def process(self, element, window=beam.DoFn.WindowParam, *args, **kwargs): # We ignore UUID collisions here since they are extremely rare. yield (uuid.uuid4().bytes, element) class ProcessKeyedElements(PTransform): """A primitive transform that performs SplittableDoFn magic. Input to this transform should be a PCollection of keyed ElementAndRestriction objects. """ def __init__( self, sdf, element_coder, restriction_coder, windowing_strategy, ptransform_args, ptransform_kwargs, ptransform_side_inputs): self.sdf = sdf self.element_coder = element_coder self.restriction_coder = restriction_coder self.windowing_strategy = windowing_strategy self.ptransform_args = ptransform_args self.ptransform_kwargs = ptransform_kwargs self.ptransform_side_inputs = ptransform_side_inputs def expand(self, pcoll): return pvalue.PCollection.from_(pcoll) class ProcessKeyedElementsViaKeyedWorkItemsOverride(PTransformOverride): """A transform override for ProcessElements transform.""" def matches(self, applied_ptransform): return isinstance(applied_ptransform.transform, ProcessKeyedElements) def get_replacement_transform_for_applied_ptransform( self, applied_ptransform): return ProcessKeyedElementsViaKeyedWorkItems(applied_ptransform.transform) class ProcessKeyedElementsViaKeyedWorkItems(PTransform): """A transform that processes Splittable DoFn input via KeyedWorkItems.""" def __init__(self, process_keyed_elements_transform): self._process_keyed_elements_transform = process_keyed_elements_transform def expand(self, pcoll): process_elements = ProcessElements(self._process_keyed_elements_transform) process_elements.args = ( self._process_keyed_elements_transform.ptransform_args) process_elements.kwargs = ( self._process_keyed_elements_transform.ptransform_kwargs) process_elements.side_inputs = ( self._process_keyed_elements_transform.ptransform_side_inputs) return pcoll | beam.core.GroupByKey() | process_elements class ProcessElements(PTransform): """A primitive transform for processing keyed elements or KeyedWorkItems. Will be evaluated by `runners.direct.transform_evaluator._ProcessElementsEvaluator`. """ def __init__(self, process_keyed_elements_transform): self._process_keyed_elements_transform = process_keyed_elements_transform self.sdf = self._process_keyed_elements_transform.sdf def expand(self, pcoll): return pvalue.PCollection.from_(pcoll) def new_process_fn(self, sdf): return ProcessFn( sdf, self._process_keyed_elements_transform.ptransform_args, self._process_keyed_elements_transform.ptransform_kwargs) class ProcessFn(beam.DoFn): """A `DoFn` that executes machineary for invoking a Splittable `DoFn`. Input to the `ParDo` step that includes a `ProcessFn` will be a `PCollection` of `ElementAndRestriction` objects. This class is mainly responsible for following. (1) setup environment for properly invoking a Splittable `DoFn`. (2) invoke `process()` method of a Splittable `DoFn`. (3) after the `process()` invocation of the Splittable `DoFn`, determine if a re-invocation of the element is needed. If this is the case, set state and a timer for a re-invocation and hold output watermark till this re-invocation. (4) after the final invocation of a given element clear any previous state set for re-invoking the element and release the output watermark. """ def __init__(self, sdf, args_for_invoker, kwargs_for_invoker): self.sdf = sdf self._element_tag = _ReadModifyWriteStateTag('element') self._restriction_tag = _ReadModifyWriteStateTag('restriction') self._watermark_state_tag = _ReadModifyWriteStateTag( 'watermark_estimator_state') self.watermark_hold_tag = _ReadModifyWriteStateTag('watermark_hold') self._process_element_invoker = None self._output_processor = _OutputHandler() self.sdf_invoker = DoFnInvoker.create_invoker( DoFnSignature(self.sdf), context=DoFnContext('unused_context'), output_handler=self._output_processor, input_args=args_for_invoker, input_kwargs=kwargs_for_invoker) self._step_context = None @property def step_context(self): return self._step_context @step_context.setter def step_context(self, step_context): assert isinstance(step_context, DirectStepContext) self._step_context = step_context def set_process_element_invoker(self, process_element_invoker): assert isinstance(process_element_invoker, SDFProcessElementInvoker) self._process_element_invoker = process_element_invoker def process( self, element, timestamp=beam.DoFn.TimestampParam, window=beam.DoFn.WindowParam, *args, **kwargs): if isinstance(element, KeyedWorkItem): # Must be a timer firing. key = element.encoded_key else: key, values = element values = list(values) assert len(values) == 1 # Value here will either be a WindowedValue or an ElementAndRestriction # object. # TODO: handle key collisions here. assert len(values) == 1, 'Internal error. Processing of splittable ' \ 'DoFn cannot continue since elements did not ' \ 'have unique keys.' value = values[0] if len(values) != 1: raise ValueError('') state = self._step_context.get_keyed_state(key) element_state = state.get_state(window, self._element_tag) # Initially element_state is an empty list. is_seed_call = not element_state if not is_seed_call: element = state.get_state(window, self._element_tag) restriction = state.get_state(window, self._restriction_tag) watermark_estimator_state = state.get_state( window, self._watermark_state_tag) windowed_element = WindowedValue(element, timestamp, [window]) else: # After values iterator is expanded above we should have gotten a list # with a single ElementAndRestriction object. assert isinstance(value, ElementAndRestriction) element_and_restriction = value element = element_and_restriction.element restriction = element_and_restriction.restriction watermark_estimator_state = ( element_and_restriction.watermark_estimator_state) if isinstance(value, WindowedValue): windowed_element = WindowedValue( element, value.timestamp, value.windows) else: windowed_element = WindowedValue(element, timestamp, [window]) assert self._process_element_invoker assert isinstance(self._process_element_invoker, SDFProcessElementInvoker) output_values = self._process_element_invoker.invoke_process_element( self.sdf_invoker, self._output_processor, windowed_element, restriction, watermark_estimator_state, *args, **kwargs) sdf_result = None for output in output_values: if isinstance(output, SDFProcessElementInvoker.Result): # SDFProcessElementInvoker.Result should be the last item yielded. sdf_result = output break yield output assert sdf_result, ('SDFProcessElementInvoker must return a ' 'SDFProcessElementInvoker.Result object as the last ' 'value of a SDF invoke_process_element() invocation.') if not sdf_result.residual_restriction: # All work for current residual and restriction pair is complete. state.clear_state(window, self._element_tag) state.clear_state(window, self._restriction_tag) state.clear_state(window, self._watermark_state_tag) # Releasing output watermark by setting it to positive infinity. state.add_state( window, self.watermark_hold_tag, WatermarkManager.WATERMARK_POS_INF) else: state.add_state(window, self._element_tag, element) state.add_state( window, self._restriction_tag, sdf_result.residual_restriction) state.add_state( window, self._watermark_state_tag, watermark_estimator_state) # Holding output watermark by setting it to negative infinity. state.add_state( window, self.watermark_hold_tag, WatermarkManager.WATERMARK_NEG_INF) # Setting a timer to be reinvoked to continue processing the element. # Currently Python SDK only supports setting timers based on watermark. So # forcing a reinvocation by setting a timer for watermark negative # infinity. # TODO(chamikara): update this by setting a timer for the proper # processing time when Python SDK supports that. state.set_timer( window, '', TimeDomain.WATERMARK, WatermarkManager.WATERMARK_NEG_INF) class SDFProcessElementInvoker(object): """A utility that invokes SDF `process()` method and requests checkpoints. This class is responsible for invoking the `process()` method of a Splittable `DoFn` and making sure that invocation terminated properly. Based on the input configuration, this class may decide to request a checkpoint for a `process()` execution so that runner can process current output and resume the invocation at a later time. More specifically, when initializing a `SDFProcessElementInvoker`, caller may specify the number of output elements or processing time after which a checkpoint should be requested. This class is responsible for properly requesting a checkpoint based on either of these criteria. When the `process()` call of Splittable `DoFn` ends, this class performs validations to make sure that processing ended gracefully and returns a `SDFProcessElementInvoker.Result` that contains information which can be used by the caller to perform another `process()` invocation for the residual. A `process()` invocation may decide to give up processing voluntarily by returning a `ProcessContinuation` object (see documentation of `ProcessContinuation` for more details). So if a 'ProcessContinuation' is produced this class ends the execution and performs steps to finalize the current invocation. """ class Result(object): def __init__( self, residual_restriction=None, process_continuation=None, future_output_watermark=None): """Returned as a result of a `invoke_process_element()` invocation. Args: residual_restriction: a restriction for the unprocessed part of the element. process_continuation: a `ProcessContinuation` if one was returned as the last element of the SDF `process()` invocation. future_output_watermark: output watermark of the results that will be produced when invoking the Splittable `DoFn` for the current element with `residual_restriction`. """ self.residual_restriction = residual_restriction self.process_continuation = process_continuation self.future_output_watermark = future_output_watermark def __init__(self, max_num_outputs, max_duration): self._max_num_outputs = max_num_outputs self._max_duration = max_duration self._checkpoint_lock = Lock() def test_method(self): raise ValueError def invoke_process_element( self, sdf_invoker, output_processor, element, restriction, watermark_estimator_state, *args, **kwargs): """Invokes `process()` method of a Splittable `DoFn` for a given element. Args: sdf_invoker: a `DoFnInvoker` for the Splittable `DoFn`. element: the element to process Returns: a `SDFProcessElementInvoker.Result` object. """ assert isinstance(sdf_invoker, DoFnInvoker) class CheckpointState(object): def __init__(self): self.checkpointed = None self.residual_restriction = None checkpoint_state = CheckpointState() def initiate_checkpoint(): with self._checkpoint_lock: if checkpoint_state.checkpointed: return checkpoint_state.checkpointed = object() split = sdf_invoker.try_split(0) if split: _, checkpoint_state.residual_restriction = split else: # Clear the checkpoint if the split didn't happen. This counters # a very unlikely race condition that the Timer attempted to initiate # a checkpoint before invoke_process set the current element allowing # for another attempt to checkpoint. checkpoint_state.checkpointed = None output_processor.reset() Timer(self._max_duration, initiate_checkpoint).start() sdf_invoker.invoke_process( element, additional_args=args, restriction=restriction, watermark_estimator_state=watermark_estimator_state) assert output_processor.output_iter is not None output_count = 0 # We have to expand and re-yield here to support ending execution for a # given number of output elements as well as to capture the # ProcessContinuation of one was returned. process_continuation = None for output in output_processor.output_iter: # A ProcessContinuation, if returned, should be the last element. assert not process_continuation if isinstance(output, ProcessContinuation): # Taking a checkpoint so that we can determine primary and residual # restrictions. initiate_checkpoint() # A ProcessContinuation should always be the last element produced by # the output iterator. # TODO: support continuing after the specified amount of delay. # Continuing here instead of breaking to enforce that this is the last # element. process_continuation = output continue yield output output_count += 1 if self._max_num_outputs and output_count >= self._max_num_outputs: initiate_checkpoint() result = ( SDFProcessElementInvoker.Result( residual_restriction=checkpoint_state.residual_restriction) if checkpoint_state.residual_restriction else SDFProcessElementInvoker.Result()) yield result class _OutputHandler(OutputHandler): def __init__(self): self.output_iter = None def handle_process_outputs( self, windowed_input_element: WindowedValue, output_iter: Iterable[Any], watermark_estimator: Optional[WatermarkEstimator] = None) -> None: self.output_iter = output_iter def reset(self): self.output_iter = None class _NoneShallPassOutputHandler(OutputHandler): def handle_process_outputs( self, windowed_input_element: WindowedValue, output_iter: Iterable[Any], watermark_estimator: Optional[WatermarkEstimator] = None) -> None: raise RuntimeError() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/sdf_direct_runner_test.py0000644000175100017510000002120115206113032025727 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for SDF implementation for DirectRunner.""" # pytype: skip-file import logging import os import unittest import apache_beam as beam from apache_beam import Create from apache_beam import DoFn from apache_beam.io import filebasedsource_test from apache_beam.io.restriction_trackers import OffsetRange from apache_beam.io.restriction_trackers import OffsetRestrictionTracker from apache_beam.pvalue import AsList from apache_beam.pvalue import AsSingleton from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.core import RestrictionProvider from apache_beam.transforms.trigger import AccumulationMode from apache_beam.transforms.window import SlidingWindows from apache_beam.transforms.window import TimestampedValue class ReadFilesProvider(RestrictionProvider): def initial_restriction(self, element): size = os.path.getsize(element) return OffsetRange(0, size) def create_tracker(self, restriction): return OffsetRestrictionTracker(restriction) def restriction_size(self, element, restriction): return restriction.size() class ReadFiles(DoFn): def __init__(self, resume_count=None): self._resume_count = resume_count def process( self, element, restriction_tracker=DoFn.RestrictionParam(ReadFilesProvider()), *args, **kwargs): file_name = element with open(file_name, 'rb') as file: pos = restriction_tracker.current_restriction().start if restriction_tracker.current_restriction().start > 0: file.seek(restriction_tracker.current_restriction().start - 1) line = file.readline() pos = pos - 1 + len(line) output_count = 0 while restriction_tracker.try_claim(pos): line = file.readline() len_line = len(line) line = line.strip() if not line: break if line is None: break yield line output_count += 1 if self._resume_count and output_count == self._resume_count: restriction_tracker.defer_remainder() break pos += len_line class ExpandStringsProvider(RestrictionProvider): def initial_restriction(self, element): return OffsetRange(0, len(element[0])) def create_tracker(self, restriction): return OffsetRestrictionTracker(restriction) # No initial split performed. def split(self, element, restriction): return [ restriction, ] def restriction_size(self, element, restriction): return restriction.size() class ExpandStrings(DoFn): def __init__(self, record_window=False): self._record_window = record_window def process( self, element, side1, side2, side3, window=beam.DoFn.WindowParam, restriction_tracker=DoFn.RestrictionParam(ExpandStringsProvider()), *args, **kwargs): side = [] side.extend(side1) side.extend(side2) side.extend(side3) side = list(side) for i in range(restriction_tracker.current_restriction().start, restriction_tracker.current_restriction().stop): if restriction_tracker.try_claim(i): if not side: yield ( element[0] + ':' + str(element[1]) + ':' + str(int(window.start)) if self._record_window else element) else: for val in side: ret = ( element[0] + ':' + str(element[1]) + ':' + str(int(window.start)) if self._record_window else element) yield ret + ':' + val else: break class SDFDirectRunnerTest(unittest.TestCase): def setUp(self): super().setUp() # Importing following for DirectRunner SDF implemenation for testing. from apache_beam.runners.direct import transform_evaluator self._old_default_max_num_outputs = ( transform_evaluator._ProcessElementsEvaluator.DEFAULT_MAX_NUM_OUTPUTS) self._default_max_num_outputs = ( transform_evaluator._ProcessElementsEvaluator.DEFAULT_MAX_NUM_OUTPUTS ) = 100 def tearDown(self): from apache_beam.runners.direct import transform_evaluator transform_evaluator._ProcessElementsEvaluator.DEFAULT_MAX_NUM_OUTPUTS = ( self._old_default_max_num_outputs) def run_sdf_read_pipeline( self, num_files, num_records_per_file, resume_count=None): expected_data = [] file_names = [] for _ in range(num_files): new_file_name, new_expected_data = filebasedsource_test.write_data( num_records_per_file) assert len(new_expected_data) == num_records_per_file file_names.append(new_file_name) expected_data.extend(new_expected_data) assert len(expected_data) > 0 with TestPipeline() as p: pc1 = ( p | 'Create1' >> beam.Create(file_names) | 'SDF' >> beam.ParDo(ReadFiles(resume_count))) assert_that(pc1, equal_to(expected_data)) # TODO(chamikara: verify the number of times process method was invoked # using a side output once SDFs supports producing side outputs. def test_sdf_no_checkpoint_single_element(self): self.run_sdf_read_pipeline(1, self._default_max_num_outputs - 1) def test_sdf_one_checkpoint_single_element(self): self.run_sdf_read_pipeline(1, int(self._default_max_num_outputs + 1)) def test_sdf_multiple_checkpoints_single_element(self): self.run_sdf_read_pipeline(1, int(self._default_max_num_outputs * 3)) def test_sdf_no_checkpoint_multiple_element(self): self.run_sdf_read_pipeline(5, int(self._default_max_num_outputs - 1)) def test_sdf_one_checkpoint_multiple_element(self): self.run_sdf_read_pipeline(5, int(self._default_max_num_outputs + 1)) def test_sdf_multiple_checkpoints_multiple_element(self): self.run_sdf_read_pipeline(5, int(self._default_max_num_outputs * 3)) def test_sdf_with_resume_single_element(self): resume_count = self._default_max_num_outputs // 10 # Makes sure that resume_count is not trivial. assert resume_count > 0 self.run_sdf_read_pipeline( 1, self._default_max_num_outputs - 1, resume_count) def test_sdf_with_resume_multiple_elements(self): resume_count = self._default_max_num_outputs // 10 assert resume_count > 0 self.run_sdf_read_pipeline( 5, int(self._default_max_num_outputs - 1), resume_count) def test_sdf_with_windowed_timestamped_input(self): with TestPipeline() as p: result = ( p | beam.Create([1, 3, 5, 10]) | beam.FlatMap( lambda t: [TimestampedValue(('A', t), t), TimestampedValue(('B', t), t)]) | beam.WindowInto( SlidingWindows(10, 5), accumulation_mode=AccumulationMode.DISCARDING) | beam.ParDo(ExpandStrings(record_window=True), [], [], [])) expected_result = [ 'A:1:-5', 'A:1:0', 'A:3:-5', 'A:3:0', 'A:5:0', 'A:5:5', 'A:10:5', 'A:10:10', 'B:1:-5', 'B:1:0', 'B:3:-5', 'B:3:0', 'B:5:0', 'B:5:5', 'B:10:5', 'B:10:10', ] assert_that(result, equal_to(expected_result)) def test_sdf_with_side_inputs(self): with TestPipeline() as p: side1 = p | 'Create1' >> Create(['1', '2']) side2 = p | 'Create2' >> Create(['3', '4']) side3 = p | 'Create3' >> Create(['5']) result = ( p | 'create_main' >> beam.Create(['a', 'b', 'c']) | beam.ParDo( ExpandStrings(), AsList(side1), AsList(side2), AsSingleton(side3))) expected_result = [] for c in ['a', 'b', 'c']: for i in range(5): expected_result.append(c + ':' + str(i + 1)) assert_that(result, equal_to(expected_result)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/test_direct_runner.py0000644000175100017510000000413515206113032025102 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Wrapper of Beam runners that's built for running and verifying e2e tests.""" # pytype: skip-file from apache_beam.internal import pickler from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import TestOptions from apache_beam.runners.direct.direct_runner import DirectRunner from apache_beam.runners.runner import PipelineState __all__ = ['TestDirectRunner'] class TestDirectRunner(DirectRunner): def run_pipeline(self, pipeline, options): """Execute test pipeline and verify test matcher""" test_options = options.view_as(TestOptions) on_success_matcher = test_options.on_success_matcher is_streaming = options.view_as(StandardOptions).streaming # [BEAM-1889] Do not send this to remote workers also, there is no need to # send this option to remote executors. test_options.on_success_matcher = None self.result = super().run_pipeline(pipeline, options) try: if not is_streaming: self.result.wait_until_finish() if on_success_matcher: from hamcrest import assert_that as hc_assert_that hc_assert_that(self.result, pickler.loads(on_success_matcher)) finally: if not PipelineState.is_terminal(self.result.state): self.result.cancel() self.result.wait_until_finish() return self.result ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/test_stream_impl.py0000644000175100017510000003267715206113032024567 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """The TestStream implementation for the DirectRunner The DirectRunner implements TestStream as the _TestStream class which is used to store the events in memory, the _WatermarkController which is used to set the watermark and emit events, and the multiplexer which sends events to the correct tagged PCollection. """ # pytype: skip-file import itertools import logging from queue import Empty as EmptyException from queue import Queue from threading import Thread from typing import Union import grpc from apache_beam import ParDo from apache_beam import coders from apache_beam import pvalue from apache_beam.portability.api import beam_interactive_api_pb2 from apache_beam.portability.api import beam_interactive_api_pb2_grpc from apache_beam.testing.test_stream import ElementEvent from apache_beam.testing.test_stream import ProcessingTimeEvent from apache_beam.testing.test_stream import WatermarkEvent from apache_beam.transforms import PTransform from apache_beam.transforms import core from apache_beam.transforms import window from apache_beam.transforms.window import TimestampedValue from apache_beam.utils import timestamp from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import Timestamp _LOGGER = logging.getLogger(__name__) class _EndOfStream: pass class _WatermarkController(PTransform): """A runner-overridable PTransform Primitive to control the watermark. Expected implementation behavior: - If the instance recieves a WatermarkEvent, it sets its output watermark to the specified value then drops the event. - If the instance receives an ElementEvent, it emits all specified elements to the Global Window with the event time set to the element's timestamp. """ def __init__(self, output_tag): self.output_tag = output_tag def get_windowing(self, _): return core.Windowing(window.GlobalWindows()) def expand(self, pcoll): ret = pvalue.PCollection.from_(pcoll) ret.tag = self.output_tag return ret class _ExpandableTestStream(PTransform): def __init__(self, test_stream): self.test_stream = test_stream def expand(self, pbegin): """Expands the TestStream into the DirectRunner implementation. Takes the TestStream transform and creates a _TestStream -> multiplexer -> _WatermarkController. """ assert isinstance(pbegin, pvalue.PBegin) # If there is only one tag there is no need to add the multiplexer. if len(self.test_stream.output_tags) == 1: return ( pbegin | _TestStream( self.test_stream.output_tags, events=self.test_stream._events, coder=self.test_stream.coder, endpoint=self.test_stream._endpoint) | _WatermarkController(list(self.test_stream.output_tags)[0])) # Multiplex to the correct PCollection based upon the event tag. def mux(event): if event.tag: yield pvalue.TaggedOutput(event.tag, event) else: yield event mux_output = ( pbegin | _TestStream( self.test_stream.output_tags, events=self.test_stream._events, coder=self.test_stream.coder, endpoint=self.test_stream._endpoint) | 'TestStream Multiplexer' >> ParDo(mux).with_outputs()) # Apply a way to control the watermark per output. It is necessary to # have an individual _WatermarkController per PCollection because the # calculation of the input watermark of a transform is based on the event # timestamp of the elements flowing through it. Meaning, it is impossible # to control the output watermarks of the individual PCollections solely # on the event timestamps. outputs = {} for tag in self.test_stream.output_tags: label = '_WatermarkController[{}]'.format(tag) outputs[tag] = (mux_output[tag] | label >> _WatermarkController(tag)) return outputs class _TestStream(PTransform): """Test stream that generates events on an unbounded PCollection of elements. Each event emits elements, advances the watermark or advances the processing time. After all of the specified elements are emitted, ceases to produce output. Expected implementation behavior: - If the instance receives a WatermarkEvent with the WATERMARK_CONTROL_TAG then the instance sets its own watermark hold at the specified value and drops the event. - If the instance receives any other WatermarkEvent or ElementEvent, it passes it to the consumer. """ # This tag is used on WatermarkEvents to control the watermark at the root # TestStream. WATERMARK_CONTROL_TAG = '_TestStream_Watermark' def __init__( self, output_tags, coder=coders.FastPrimitivesCoder(), events=None, endpoint=None): assert coder is not None self.coder = coder self._raw_events = events self._events = self._add_watermark_advancements(output_tags, events) self.output_tags = output_tags self.endpoint = endpoint def _watermark_starts(self, output_tags): """Sentinel values to hold the watermark of outputs to -inf. The output watermarks of the output PCollections (fake unbounded sources) in a TestStream are controlled by watermark holds. This sets the hold of each output PCollection so that the individual holds can be controlled by the given events. """ return [WatermarkEvent(timestamp.MIN_TIMESTAMP, tag) for tag in output_tags] def _watermark_stops(self, output_tags): """Sentinel values to close the watermark of outputs.""" return [WatermarkEvent(timestamp.MAX_TIMESTAMP, tag) for tag in output_tags] def _test_stream_start(self): """Sentinel value to move the watermark hold of the TestStream to +inf. This sets a hold to +inf such that the individual holds of the output PCollections are allowed to modify their individial output watermarks with their holds. This is because the calculation of the output watermark is a min over all input watermarks. """ return [ WatermarkEvent( timestamp.MAX_TIMESTAMP - timestamp.TIME_GRANULARITY, _TestStream.WATERMARK_CONTROL_TAG) ] def _test_stream_stop(self): """Sentinel value to close the watermark of the TestStream.""" return [ WatermarkEvent( timestamp.MAX_TIMESTAMP, _TestStream.WATERMARK_CONTROL_TAG) ] def _test_stream_init(self): """Sentinel value to hold the watermark of the TestStream to -inf. This sets a hold to ensure that the output watermarks of the output PCollections do not advance to +inf before their watermark holds are set. """ return [ WatermarkEvent( timestamp.MIN_TIMESTAMP, _TestStream.WATERMARK_CONTROL_TAG) ] def _set_up(self, output_tags): return ( self._test_stream_init() + self._watermark_starts(output_tags) + self._test_stream_start()) def _tear_down(self, output_tags): return self._watermark_stops(output_tags) + self._test_stream_stop() def _add_watermark_advancements(self, output_tags, events): """Adds watermark advancements to the given events. The following watermark advancements can be done on the runner side. However, it makes the logic on the runner side much more complicated than it needs to be. In order for watermarks to be properly advanced in a TestStream, a specific sequence of watermark holds must be sent: 1. Hold the root watermark at -inf (this prevents the pipeline from immediately returning). 2. Hold the watermarks at the WatermarkControllerss at -inf (this prevents the pipeline from immediately returning). 3. Advance the root watermark to +inf - 1 (this allows the downstream WatermarkControllers to control their watermarks via holds). 4. Advance watermarks as normal. 5. Advance WatermarkController watermarks to +inf 6. Advance root watermark to +inf. """ if not events: return [] return self._set_up(output_tags) + events + self._tear_down(output_tags) def get_windowing(self, unused_inputs): return core.Windowing(window.GlobalWindows()) def expand(self, pcoll): return pvalue.PCollection(pcoll.pipeline, is_bounded=False) def _infer_output_coder(self, input_type=None, input_coder=None): return self.coder @staticmethod def events_from_script(events): """Yields the in-memory events. """ return itertools.chain(events) @staticmethod def _stream_events_from_rpc(endpoint, output_tags, coder, channel, is_alive): """Yields the events received from the given endpoint. This is the producer thread that reads events from the TestStreamService and puts them onto the shared queue. At the end of the stream, an _EndOfStream is placed on the channel to signify a successful end. """ stub_channel = grpc.insecure_channel(endpoint) stub = beam_interactive_api_pb2_grpc.TestStreamServiceStub(stub_channel) # Request the PCollections that we are looking for from the service. event_request = beam_interactive_api_pb2.EventsRequest( output_ids=[str(tag) for tag in output_tags]) event_stream = stub.Events(event_request) try: for e in event_stream: channel.put(_TestStream.test_stream_payload_to_events(e, coder)) if not is_alive(): return except grpc.RpcError as e: # Do not raise an exception in the non-error status codes. These can occur # when the Python interpreter shuts down or when in a notebook environment # when the kernel is interrupted. if e.code() in (grpc.StatusCode.CANCELLED, grpc.StatusCode.UNAVAILABLE): return raise e finally: # Gracefully stop the job if there is an exception. channel.put(_EndOfStream()) @staticmethod def events_from_rpc(endpoint, output_tags, coder, evaluation_context): """Yields the events received from the given endpoint. This method starts a new thread that reads from the TestStreamService and puts the events onto a shared queue. This method then yields all elements from the queue. Unfortunately, this is necessary because the GRPC API does not allow for non-blocking calls when utilizing a streaming RPC. It is officially suggested from the docs to use a producer/consumer pattern to handle streaming RPCs. By doing so, this gives this method control over when to cancel reading from the RPC if the server takes too long to respond. """ # Shared variable with the producer queue. This shuts down the producer if # the consumer exits early. shutdown_requested = False def is_alive(): return not (shutdown_requested or evaluation_context.shutdown_requested) # The shared queue that allows the producer and consumer to communicate. channel: 'Queue[Union[test_stream.Event, _EndOfStream]]' = ( # noqa: F821 Queue()) event_stream = Thread( target=_TestStream._stream_events_from_rpc, args=(endpoint, output_tags, coder, channel, is_alive)) event_stream.setDaemon(True) event_stream.start() # This pumps the shared queue for events until the _EndOfStream sentinel is # reached. If the TestStreamService takes longer than expected, the queue # will timeout and an EmptyException will be raised. This also sets the # shared is_alive sentinel to shut down the producer. while True: try: # Raise an EmptyException if there are no events during the last timeout # period. event = channel.get(timeout=30) if isinstance(event, _EndOfStream): break yield event except EmptyException as e: _LOGGER.warning( 'TestStream timed out waiting for new events from service.' ' Stopping pipeline.') shutdown_requested = True raise e @staticmethod def test_stream_payload_to_events(payload, coder): """Returns a TestStream Python event object from a TestStream event Proto. """ if payload.HasField('element_event'): element_event = payload.element_event elements = [ TimestampedValue( coder.decode(e.encoded_element), Timestamp(micros=e.timestamp)) for e in element_event.elements ] return ElementEvent(timestamped_values=elements, tag=element_event.tag) if payload.HasField('watermark_event'): watermark_event = payload.watermark_event return WatermarkEvent( Timestamp(micros=watermark_event.new_watermark), tag=watermark_event.tag) if payload.HasField('processing_time_event'): processing_time_event = payload.processing_time_event return ProcessingTimeEvent( Duration(micros=processing_time_event.advance_duration)) raise RuntimeError( 'Received a proto without the specified fields: {}'.format(payload)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/transform_evaluator.py0000644000175100017510000012657215206113032025307 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """An evaluator of a specific application of a transform.""" # pytype: skip-file import atexit import collections import logging import random import time from collections import abc from typing import TYPE_CHECKING from typing import Any from apache_beam import coders from apache_beam import io from apache_beam import pvalue from apache_beam.internal import pickler from apache_beam.pipeline import AppliedPTransform from apache_beam.runners import common from apache_beam.runners.common import DoFnRunner from apache_beam.runners.common import DoFnState from apache_beam.runners.direct.direct_runner import _DirectReadFromPubSub from apache_beam.runners.direct.direct_runner import _GroupByKeyOnly from apache_beam.runners.direct.direct_runner import _StreamingGroupAlsoByWindow from apache_beam.runners.direct.direct_runner import _StreamingGroupByKeyOnly from apache_beam.runners.direct.direct_userstate import DirectUserStateContext from apache_beam.runners.direct.sdf_direct_runner import ProcessElements from apache_beam.runners.direct.sdf_direct_runner import ProcessFn from apache_beam.runners.direct.sdf_direct_runner import SDFProcessElementInvoker from apache_beam.runners.direct.test_stream_impl import _TestStream from apache_beam.runners.direct.test_stream_impl import _WatermarkController from apache_beam.runners.direct.util import KeyedWorkItem from apache_beam.runners.direct.util import TransformResult from apache_beam.runners.direct.watermark_manager import WatermarkManager from apache_beam.testing.test_stream import ElementEvent from apache_beam.testing.test_stream import PairWithTiming from apache_beam.testing.test_stream import ProcessingTimeEvent from apache_beam.testing.test_stream import TimingInfo from apache_beam.testing.test_stream import WatermarkEvent from apache_beam.testing.test_stream import WindowedValueHolder from apache_beam.transforms import core from apache_beam.transforms.trigger import InMemoryUnmergedState from apache_beam.transforms.trigger import TimeDomain from apache_beam.transforms.trigger import _CombiningValueStateTag from apache_beam.transforms.trigger import _ListStateTag from apache_beam.transforms.trigger import _ReadModifyWriteStateTag from apache_beam.transforms.trigger import create_trigger_driver from apache_beam.transforms.userstate import get_dofn_specs from apache_beam.transforms.userstate import is_stateful_dofn from apache_beam.transforms.window import GlobalWindows from apache_beam.transforms.window import WindowedValue from apache_beam.typehints.typecheck import TypeCheckError from apache_beam.utils import counters from apache_beam.utils.timestamp import MIN_TIMESTAMP from apache_beam.utils.timestamp import Timestamp if TYPE_CHECKING: from apache_beam.io.gcp.pubsub import PubsubMessage from apache_beam.io.gcp.pubsub import _PubSubSource from apache_beam.runners.direct.evaluation_context import EvaluationContext _LOGGER = logging.getLogger(__name__) class TransformEvaluatorRegistry(object): """For internal use only; no backwards-compatibility guarantees. Creates instances of TransformEvaluator for the application of a transform. """ _test_evaluators_overrides: dict[type[core.PTransform], type['_TransformEvaluator']] = {} def __init__(self, evaluation_context: 'EvaluationContext') -> None: assert evaluation_context self._evaluation_context = evaluation_context self._evaluators: dict[type[core.PTransform], type[_TransformEvaluator]] = { io.Read: _BoundedReadEvaluator, _DirectReadFromPubSub: _PubSubReadEvaluator, core.Flatten: _FlattenEvaluator, core.Impulse: _ImpulseEvaluator, core.ParDo: _ParDoEvaluator, _GroupByKeyOnly: _GroupByKeyOnlyEvaluator, _StreamingGroupByKeyOnly: _StreamingGroupByKeyOnlyEvaluator, _StreamingGroupAlsoByWindow: _StreamingGroupAlsoByWindowEvaluator, _TestStream: _TestStreamEvaluator, ProcessElements: _ProcessElementsEvaluator, _WatermarkController: _WatermarkControllerEvaluator, PairWithTiming: _PairWithTimingEvaluator, } self._evaluators.update(self._test_evaluators_overrides) self._root_bundle_providers = { core.PTransform: DefaultRootBundleProvider, _TestStream: _TestStreamRootBundleProvider, } def get_evaluator( self, applied_ptransform, input_committed_bundle, side_inputs): """Returns a TransformEvaluator suitable for processing given inputs.""" assert applied_ptransform assert bool(applied_ptransform.side_inputs) == bool(side_inputs) # Walk up the class hierarchy to find an evaluable type. This is necessary # for supporting sub-classes of core transforms. for cls in applied_ptransform.transform.__class__.mro(): evaluator = self._evaluators.get(cls) if evaluator: break if not evaluator: raise NotImplementedError( 'Execution of [%s] not implemented in runner %s.' % (type(applied_ptransform.transform), self)) return evaluator( self._evaluation_context, applied_ptransform, input_committed_bundle, side_inputs) def get_root_bundle_provider(self, applied_ptransform): provider_cls = None for cls in applied_ptransform.transform.__class__.mro(): provider_cls = self._root_bundle_providers.get(cls) if provider_cls: break if not provider_cls: raise NotImplementedError( 'Root provider for [%s] not implemented in runner %s' % (type(applied_ptransform.transform), self)) return provider_cls(self._evaluation_context, applied_ptransform) def should_execute_serially(self, applied_ptransform): """Returns True if this applied_ptransform should run one bundle at a time. Some TransformEvaluators use a global state object to keep track of their global execution state. For example evaluator for _GroupByKeyOnly uses this state as an in memory dictionary to buffer keys. Serially executed evaluators will act as syncing point in the graph and execution will not move forward until they receive all of their inputs. Once they receive all of their input, they will release the combined output. Their output may consist of multiple bundles as they may divide their output into pieces before releasing. Args: applied_ptransform: Transform to be used for execution. Returns: True if executor should execute applied_ptransform serially. """ if isinstance( applied_ptransform.transform, (_GroupByKeyOnly, _StreamingGroupByKeyOnly, _StreamingGroupAlsoByWindow)): return True elif (isinstance(applied_ptransform.transform, core.ParDo) and is_stateful_dofn(applied_ptransform.transform.dofn)): return True return False class RootBundleProvider(object): """Provides bundles for the initial execution of a root transform.""" def __init__(self, evaluation_context, applied_ptransform): self._evaluation_context = evaluation_context self._applied_ptransform = applied_ptransform def get_root_bundles(self): raise NotImplementedError class DefaultRootBundleProvider(RootBundleProvider): """Provides an empty bundle by default for root transforms.""" def get_root_bundles(self): input_node = pvalue.PBegin(self._applied_ptransform.transform.pipeline) empty_bundle = ( self._evaluation_context.create_empty_committed_bundle(input_node)) return [empty_bundle] class _TestStreamRootBundleProvider(RootBundleProvider): """Provides an initial bundle for the TestStream evaluator. This bundle is used as the initial state to the TestStream. Each unprocessed bundle emitted from the TestStream afterwards is its state: index into the stream, and the watermark. """ def get_root_bundles(self): test_stream = self._applied_ptransform.transform # If there was an endpoint defined then get the events from the # TestStreamService. if test_stream.endpoint: _TestStreamEvaluator.event_stream = _TestStream.events_from_rpc( test_stream.endpoint, test_stream.output_tags, test_stream.coder, self._evaluation_context) else: _TestStreamEvaluator.event_stream = ( _TestStream.events_from_script(test_stream._events)) bundle = self._evaluation_context.create_bundle( pvalue.PBegin(self._applied_ptransform.transform.pipeline)) bundle.add(GlobalWindows.windowed_value(b'', timestamp=MIN_TIMESTAMP)) bundle.commit(None) return [bundle] class _TransformEvaluator(object): """An evaluator of a specific application of a transform.""" def __init__( self, evaluation_context: 'EvaluationContext', applied_ptransform: AppliedPTransform, input_committed_bundle, side_inputs): self._evaluation_context = evaluation_context self._applied_ptransform = applied_ptransform self._input_committed_bundle = input_committed_bundle self._side_inputs = side_inputs self._expand_outputs() self._execution_context = evaluation_context.get_execution_context( applied_ptransform) self._step_context = self._execution_context.get_step_context() def _expand_outputs(self): outputs = set() for pval in self._applied_ptransform.outputs.values(): if isinstance(pval, pvalue.DoOutputsTuple): pvals = (v for v in pval) else: pvals = (pval, ) for v in pvals: outputs.add(v) self._outputs = frozenset(outputs) def _split_list_into_bundles( self, output_pcollection, elements, max_element_per_bundle, element_size_fn): """Splits elements, an iterable, into multiple output bundles. Args: output_pcollection: PCollection that the elements belong to. elements: elements to be chunked into bundles. max_element_per_bundle: (approximately) the maximum element per bundle. If it is None, only a single bundle will be produced. element_size_fn: Function to return the size of a given element. Returns: List of output uncommitted bundles with at least one bundle. """ bundle = self._evaluation_context.create_bundle(output_pcollection) bundle_size = 0 bundles = [bundle] for element in elements: if max_element_per_bundle and bundle_size >= max_element_per_bundle: bundle = self._evaluation_context.create_bundle(output_pcollection) bundle_size = 0 bundles.append(bundle) bundle.output(element) bundle_size += element_size_fn(element) return bundles def start_bundle(self): """Starts a new bundle.""" pass def process_timer_wrapper(self, timer_firing): """Process timer by clearing and then calling process_timer(). This method is called with any timer firing and clears the delivered timer from the keyed state and then calls process_timer(). The default process_timer() implementation emits a KeyedWorkItem for the particular timer and passes it to process_element(). Evaluator subclasses which desire different timer delivery semantics can override process_timer(). """ state = self._step_context.get_keyed_state(timer_firing.encoded_key) state.clear_timer( timer_firing.window, timer_firing.name, timer_firing.time_domain, dynamic_timer_tag=timer_firing.dynamic_timer_tag) self.process_timer(timer_firing) def process_timer(self, timer_firing): """Default process_timer() impl. generating KeyedWorkItem element.""" self.process_element( GlobalWindows.windowed_value( KeyedWorkItem( timer_firing.encoded_key, timer_firings=[timer_firing]))) def process_element(self, element): """Processes a new element as part of the current bundle.""" raise NotImplementedError('%s do not process elements.' % type(self)) def finish_bundle(self) -> TransformResult: """Finishes the bundle and produces output.""" pass class _BoundedReadEvaluator(_TransformEvaluator): """TransformEvaluator for bounded Read transform.""" # After some benchmarks, 1000 was optimal among {100,1000,10000} MAX_ELEMENT_PER_BUNDLE = 1000 def __init__( self, evaluation_context, applied_ptransform, input_committed_bundle, side_inputs): assert not side_inputs self._source = applied_ptransform.transform.source self._source.pipeline_options = evaluation_context.pipeline_options super().__init__( evaluation_context, applied_ptransform, input_committed_bundle, side_inputs) def finish_bundle(self): assert len(self._outputs) == 1 output_pcollection = list(self._outputs)[0] def _read_values_to_bundles(reader): read_result = [GlobalWindows.windowed_value(e) for e in reader] return self._split_list_into_bundles( output_pcollection, read_result, _BoundedReadEvaluator.MAX_ELEMENT_PER_BUNDLE, lambda _: 1) if isinstance(self._source, io.iobase.BoundedSource): # Getting a RangeTracker for the default range of the source and reading # the full source using that. range_tracker = self._source.get_range_tracker(None, None) reader = self._source.read(range_tracker) bundles = _read_values_to_bundles(reader) else: with self._source.reader() as reader: bundles = _read_values_to_bundles(reader) return TransformResult(self, bundles, [], None, None) class _WatermarkControllerEvaluator(_TransformEvaluator): """TransformEvaluator for the _WatermarkController transform. This is used to enable multiple output watermarks for the TestStream. """ # The state tag used to store the watermark. WATERMARK_TAG = _ReadModifyWriteStateTag( '_WatermarkControllerEvaluator_Watermark_Tag') def __init__( self, evaluation_context, applied_ptransform, input_committed_bundle, side_inputs): assert not side_inputs self.transform = applied_ptransform.transform super().__init__( evaluation_context, applied_ptransform, input_committed_bundle, side_inputs) self._init_state() def _init_state(self): """Gets and sets the initial state. This is used to keep track of the watermark hold between calls. """ transform_states = self._evaluation_context._transform_keyed_states state = transform_states[self._applied_ptransform] if self.WATERMARK_TAG not in state: watermark_state = InMemoryUnmergedState() watermark_state.set_global_state(self.WATERMARK_TAG, MIN_TIMESTAMP) state[self.WATERMARK_TAG] = watermark_state self._state = state[self.WATERMARK_TAG] @property def _watermark(self): return self._state.get_global_state(self.WATERMARK_TAG) @_watermark.setter def _watermark(self, watermark): self._state.set_global_state(self.WATERMARK_TAG, watermark) def start_bundle(self): self.bundles = [] def process_element(self, element): # In order to keep the order of the elements between the script and what # flows through the pipeline the same, emit the elements here. event = element.value if isinstance(event, WatermarkEvent): self._watermark = event.new_watermark elif isinstance(event, ElementEvent): main_output = list(self._outputs)[0] bundle = self._evaluation_context.create_bundle(main_output) for tv in event.timestamped_values: # Unreify the value into the correct window. if isinstance(tv.value, WindowedValueHolder): bundle.output(tv.value.windowed_value) else: bundle.output( GlobalWindows.windowed_value(tv.value, timestamp=tv.timestamp)) self.bundles.append(bundle) def finish_bundle(self): # The watermark hold we set here is the way we allow the TestStream events # to control the output watermark. return TransformResult( self, self.bundles, [], None, {None: self._watermark}) class _PairWithTimingEvaluator(_TransformEvaluator): """TransformEvaluator for the PairWithTiming transform. This transform takes an element as an input and outputs KV(element, `TimingInfo`). Where the `TimingInfo` contains both the processing time timestamp and watermark. """ def __init__( self, evaluation_context, applied_ptransform, input_committed_bundle, side_inputs): assert not side_inputs super().__init__( evaluation_context, applied_ptransform, input_committed_bundle, side_inputs) def start_bundle(self): main_output = list(self._outputs)[0] self.bundle = self._evaluation_context.create_bundle(main_output) watermark_manager = self._evaluation_context._watermark_manager watermarks = watermark_manager.get_watermarks(self._applied_ptransform) output_watermark = watermarks.output_watermark now = Timestamp(seconds=watermark_manager._clock.time()) self.timing_info = TimingInfo(now, output_watermark) def process_element(self, element): result = WindowedValue((element.value, self.timing_info), element.timestamp, element.windows, element.pane_info) self.bundle.output(result) def finish_bundle(self): return TransformResult(self, [self.bundle], [], None, {}) class _TestStreamEvaluator(_TransformEvaluator): """TransformEvaluator for the TestStream transform. This evaluator's responsibility is to retrieve the next event from the _TestStream and either: advance the clock, advance the _TestStream watermark, or pass the event to the _WatermarkController. The _WatermarkController is in charge of emitting the elements to the downstream consumers and setting its own output watermark. """ event_stream = None def __init__( self, evaluation_context, applied_ptransform, input_committed_bundle, side_inputs): assert not side_inputs super().__init__( evaluation_context, applied_ptransform, input_committed_bundle, side_inputs) self.test_stream = applied_ptransform.transform self.is_done = False def start_bundle(self): self.bundles = [] self.watermark = MIN_TIMESTAMP def process_element(self, element): # The watermark of the _TestStream transform itself. self.watermark = element.timestamp # Set up the correct watermark holds in the Watermark controllers and the # TestStream so that the watermarks will not automatically advance to +inf # when elements start streaming. This can happen multiple times in the first # bundle, but the operations are idempotent and adding state to keep track # of this would add unnecessary code complexity. events = [] if self.watermark == MIN_TIMESTAMP: for event in self.test_stream._set_up(self.test_stream.output_tags): events.append(event) # Retrieve the TestStream's event stream and read from it. try: events.append(next(self.event_stream)) except StopIteration: # Advance the watermarks to +inf to cleanly stop the pipeline. self.is_done = True events += ([ e for e in self.test_stream._tear_down(self.test_stream.output_tags) ]) for event in events: # We can either have the _TestStream or the _WatermarkController to emit # the elements. We chose to emit in the _WatermarkController so that the # element is emitted at the correct watermark value. if isinstance(event, (ElementEvent, WatermarkEvent)): # The WATERMARK_CONTROL_TAG is used to hold the _TestStream's # watermark to -inf, then +inf-1, then +inf. This watermark progression # is ultimately used to set up the proper holds to allow the # _WatermarkControllers to control their own output watermarks. if event.tag == _TestStream.WATERMARK_CONTROL_TAG: self.watermark = event.new_watermark else: main_output = list(self._outputs)[0] bundle = self._evaluation_context.create_bundle(main_output) bundle.output(GlobalWindows.windowed_value(event)) self.bundles.append(bundle) elif isinstance(event, ProcessingTimeEvent): self._evaluation_context._watermark_manager._clock.advance_time( event.advance_by) else: raise ValueError('Invalid TestStream event: %s.' % event) def finish_bundle(self): unprocessed_bundles = [] # Continue to send its own state to itself via an unprocessed bundle. This # acts as a heartbeat, where each element will read the next event from the # event stream. if not self.is_done: unprocessed_bundle = self._evaluation_context.create_bundle( pvalue.PBegin(self._applied_ptransform.transform.pipeline)) unprocessed_bundle.add( GlobalWindows.windowed_value(b'', timestamp=self.watermark)) unprocessed_bundles.append(unprocessed_bundle) # Returning the watermark in the dict here is used as a watermark hold. return TransformResult( self, self.bundles, unprocessed_bundles, None, {None: self.watermark}) class _PubSubReadEvaluator(_TransformEvaluator): """TransformEvaluator for PubSub read.""" # A mapping of transform to _PubSubSubscriptionWrapper. # TODO(https://github.com/apache/beam/issues/19751): Prevents garbage # collection of pipeline instances. _subscription_cache: dict[AppliedPTransform, str] = {} def __init__( self, evaluation_context, applied_ptransform, input_committed_bundle, side_inputs): assert not side_inputs super().__init__( evaluation_context, applied_ptransform, input_committed_bundle, side_inputs) self.source: _PubSubSource = self._applied_ptransform.transform._source if self.source.id_label: raise NotImplementedError( 'DirectRunner: id_label is not supported for PubSub reads') sub_project = None if hasattr(self._evaluation_context, 'pipeline_options'): from apache_beam.options.pipeline_options import GoogleCloudOptions sub_project = ( self._evaluation_context.pipeline_options.view_as( GoogleCloudOptions).project) if not sub_project: sub_project = self.source.project self._sub_name = self.get_subscription( self._applied_ptransform, self.source.project, self.source.topic_name, sub_project, self.source.subscription_name) @classmethod def get_subscription( cls, transform, project, short_topic_name, sub_project, short_sub_name): from google.cloud import pubsub if short_sub_name: return pubsub.SubscriberClient.subscription_path(project, short_sub_name) if transform in cls._subscription_cache: return cls._subscription_cache[transform] sub_client = pubsub.SubscriberClient() sub_name = sub_client.subscription_path( sub_project, 'beam_%d_%x' % (int(time.time()), random.randrange(1 << 32))) topic_name = sub_client.topic_path(project, short_topic_name) sub_client.create_subscription(name=sub_name, topic=topic_name) atexit.register(sub_client.delete_subscription, subscription=sub_name) cls._subscription_cache[transform] = sub_name return cls._subscription_cache[transform] def start_bundle(self): pass def process_element(self, element): pass def _read_from_pubsub( self, timestamp_attribute) -> list[tuple[Timestamp, 'PubsubMessage']]: from google.cloud import pubsub from apache_beam.io.gcp.pubsub import PubsubMessage def _get_element(message): parsed_message = PubsubMessage._from_message(message) if (timestamp_attribute and timestamp_attribute in parsed_message.attributes): rfc3339_or_milli = parsed_message.attributes[timestamp_attribute] try: timestamp = Timestamp(micros=int(rfc3339_or_milli) * 1000) except ValueError: try: timestamp = Timestamp.from_rfc3339(rfc3339_or_milli) except ValueError as e: raise ValueError('Bad timestamp value: %s' % e) else: if message.publish_time is None: raise ValueError('No publish time present in message: %s' % message) try: timestamp = Timestamp.from_utc_datetime(message.publish_time) except ValueError as e: raise ValueError('Bad timestamp value for message %s: %s', message, e) return timestamp, parsed_message # Because of the AutoAck, we are not able to reread messages if this # evaluator fails with an exception before emitting a bundle. However, # the DirectRunner currently doesn't retry work items anyway, so the # pipeline would enter an inconsistent state on any error. sub_client = pubsub.SubscriberClient() try: response = sub_client.pull( subscription=self._sub_name, max_messages=10, timeout=30) results = [_get_element(rm.message) for rm in response.received_messages] ack_ids = [rm.ack_id for rm in response.received_messages] if ack_ids: sub_client.acknowledge(subscription=self._sub_name, ack_ids=ack_ids) finally: sub_client.close() return results def finish_bundle(self) -> TransformResult: data = self._read_from_pubsub(self.source.timestamp_attribute) if data: output_pcollection = list(self._outputs)[0] bundle = self._evaluation_context.create_bundle(output_pcollection) # TODO(ccy): Respect the PubSub source's id_label field. for timestamp, message in data: if self.source.with_attributes: element = message else: element = message.data bundle.output( GlobalWindows.windowed_value(element, timestamp=timestamp)) bundles = [bundle] else: bundles = [] assert self._applied_ptransform.transform is not None if self._applied_ptransform.inputs: input_pvalue = self._applied_ptransform.inputs[0] else: input_pvalue = pvalue.PBegin(self._applied_ptransform.transform.pipeline) unprocessed_bundle = self._evaluation_context.create_bundle(input_pvalue) # TODO(udim): Correct value for watermark hold. return TransformResult( self, bundles, [unprocessed_bundle], None, {None: Timestamp.of(time.time())}) class _FlattenEvaluator(_TransformEvaluator): """TransformEvaluator for Flatten transform.""" def __init__( self, evaluation_context, applied_ptransform, input_committed_bundle, side_inputs): assert not side_inputs super().__init__( evaluation_context, applied_ptransform, input_committed_bundle, side_inputs) def start_bundle(self): assert len(self._outputs) == 1 output_pcollection = list(self._outputs)[0] self.bundle = self._evaluation_context.create_bundle(output_pcollection) def process_element(self, element): self.bundle.output(element) def finish_bundle(self): bundles = [self.bundle] return TransformResult(self, bundles, [], None, None) class _ImpulseEvaluator(_TransformEvaluator): """TransformEvaluator for Impulse transform.""" def finish_bundle(self): assert len(self._outputs) == 1 output_pcollection = list(self._outputs)[0] bundle = self._evaluation_context.create_bundle(output_pcollection) bundle.output(GlobalWindows.windowed_value(b'')) return TransformResult(self, [bundle], [], None, None) class _TaggedReceivers(dict): """Received ParDo output and redirect to the associated output bundle.""" def __init__(self, evaluation_context): self._evaluation_context = evaluation_context self._null_receiver = None super().__init__() class NullReceiver(common.Receiver): """Ignores undeclared outputs, default execution mode.""" def receive(self, element: WindowedValue) -> None: pass class _InMemoryReceiver(common.Receiver): """Buffers undeclared outputs to the given dictionary.""" def __init__(self, target, tag): self._target = target self._tag = tag def receive(self, element: WindowedValue) -> None: self._target[self._tag].append(element) def __missing__(self, key): if not self._null_receiver: self._null_receiver = _TaggedReceivers.NullReceiver() return self._null_receiver class _ParDoEvaluator(_TransformEvaluator): """TransformEvaluator for ParDo transform.""" def __init__( self, evaluation_context: 'EvaluationContext', applied_ptransform: AppliedPTransform, input_committed_bundle, side_inputs, perform_dofn_pickle_test=True): super().__init__( evaluation_context, applied_ptransform, input_committed_bundle, side_inputs) # This is a workaround for SDF implementation. SDF implementation adds state # to the SDF that is not picklable. self._perform_dofn_pickle_test = perform_dofn_pickle_test def start_bundle(self): transform = self._applied_ptransform.transform self._tagged_receivers = _TaggedReceivers(self._evaluation_context) for output_tag in self._applied_ptransform.outputs: output_pcollection = pvalue.PCollection(None, tag=output_tag) output_pcollection.producer = self._applied_ptransform self._tagged_receivers[output_tag] = ( self._evaluation_context.create_bundle(output_pcollection)) self._tagged_receivers[output_tag].tag = output_tag self._counter_factory = counters.CounterFactory() # TODO(aaltay): Consider storing the serialized form as an optimization. dofn = ( pickler.roundtrip(transform.dofn) if self._perform_dofn_pickle_test else transform.dofn) args = transform.args if hasattr(transform, 'args') else [] kwargs = transform.kwargs if hasattr(transform, 'kwargs') else {} self.user_state_context = None self.user_timer_map = {} if is_stateful_dofn(dofn): kv_type_hint = self._applied_ptransform.inputs[0].element_type if kv_type_hint and kv_type_hint != Any: coder = coders.registry.get_coder(kv_type_hint) self.key_coder = coder.key_coder() else: self.key_coder = coders.registry.get_coder(Any) self.user_state_context = DirectUserStateContext( self._step_context, dofn, self.key_coder) _, all_timer_specs = get_dofn_specs(dofn) for timer_spec in all_timer_specs: self.user_timer_map['user/%s' % timer_spec.name] = timer_spec self.runner = DoFnRunner( dofn, args, kwargs, self._side_inputs, self._applied_ptransform.inputs[0].windowing, tagged_receivers=self._tagged_receivers, step_name=self._applied_ptransform.full_label, state=DoFnState(self._counter_factory), user_state_context=self.user_state_context) self.runner.setup() self.runner.start() def process_timer(self, timer_firing): if timer_firing.name not in self.user_timer_map: _LOGGER.warning('Unknown timer fired: %s', timer_firing) timer_spec = self.user_timer_map[timer_firing.name] self.runner.process_user_timer( timer_spec, self.key_coder.decode(timer_firing.encoded_key), timer_firing.window, timer_firing.timestamp, # TODO Add paneinfo to timer_firing in DirectRunner None, timer_firing.dynamic_timer_tag) def process_element(self, element): self.runner.process(element) def finish_bundle(self): self.runner.finish() self.runner.teardown() bundles = list(self._tagged_receivers.values()) result_counters = self._counter_factory.get_counters() if self.user_state_context: self.user_state_context.commit() self.user_state_context.reset() return TransformResult(self, bundles, [], result_counters, None) class _GroupByKeyOnlyEvaluator(_TransformEvaluator): """TransformEvaluator for _GroupByKeyOnly transform.""" MAX_ELEMENT_PER_BUNDLE = None ELEMENTS_TAG = _ListStateTag('elements') COMPLETION_TAG = _CombiningValueStateTag('completed', any) def __init__( self, evaluation_context, applied_ptransform, input_committed_bundle, side_inputs): assert not side_inputs super().__init__( evaluation_context, applied_ptransform, input_committed_bundle, side_inputs) def _is_final_bundle(self): return ( self._execution_context.watermarks.input_watermark == WatermarkManager.WATERMARK_POS_INF) def start_bundle(self): self.global_state = self._step_context.get_keyed_state(None) assert len(self._outputs) == 1 self.output_pcollection = list(self._outputs)[0] # The output type of a GroupByKey will be Tuple[Any, Any] or more specific. # TODO(https://github.com/apache/beam/issues/18490): Infer coders earlier. kv_type_hint = ( self._applied_ptransform.outputs[None].element_type or self._applied_ptransform.transform.get_type_hints().input_types[0][0]) self.key_coder = coders.registry.get_coder(kv_type_hint.tuple_types[0]) def process_timer(self, timer_firing): # We do not need to emit a KeyedWorkItem to process_element(). pass def process_element(self, element): assert not self.global_state.get_state( None, _GroupByKeyOnlyEvaluator.COMPLETION_TAG) if (isinstance(element, WindowedValue) and isinstance(element.value, abc.Iterable) and len(element.value) == 2): k, v = element.value encoded_k = self.key_coder.encode(k) state = self._step_context.get_keyed_state(encoded_k) state.add_state(None, _GroupByKeyOnlyEvaluator.ELEMENTS_TAG, v) else: raise TypeCheckError( 'Input to _GroupByKeyOnly must be a PCollection of ' 'windowed key-value pairs. Instead received: %r.' % element) def finish_bundle(self): if self._is_final_bundle(): if self.global_state.get_state(None, _GroupByKeyOnlyEvaluator.COMPLETION_TAG): # Ignore empty bundles after emitting output. (This may happen because # empty bundles do not affect input watermarks.) bundles = [] else: gbk_result = [] # TODO(ccy): perhaps we can clean this up to not use this # internal attribute of the DirectStepContext. for encoded_k in self._step_context.existing_keyed_state: # Ignore global state. if encoded_k is None: continue k = self.key_coder.decode(encoded_k) state = self._step_context.get_keyed_state(encoded_k) vs = state.get_state(None, _GroupByKeyOnlyEvaluator.ELEMENTS_TAG) gbk_result.append(GlobalWindows.windowed_value((k, vs))) def len_element_fn(element): _, v = element.value return len(v) bundles = self._split_list_into_bundles( self.output_pcollection, gbk_result, _GroupByKeyOnlyEvaluator.MAX_ELEMENT_PER_BUNDLE, len_element_fn) self.global_state.add_state( None, _GroupByKeyOnlyEvaluator.COMPLETION_TAG, True) hold = WatermarkManager.WATERMARK_POS_INF else: bundles = [] hold = WatermarkManager.WATERMARK_NEG_INF self.global_state.set_timer( None, '', TimeDomain.WATERMARK, WatermarkManager.WATERMARK_POS_INF) return TransformResult(self, bundles, [], None, {None: hold}) class _StreamingGroupByKeyOnlyEvaluator(_TransformEvaluator): """TransformEvaluator for _StreamingGroupByKeyOnly transform. The _GroupByKeyOnlyEvaluator buffers elements until its input watermark goes to infinity, which is suitable for batch mode execution. During streaming mode execution, we emit each bundle as it comes to the next transform. """ MAX_ELEMENT_PER_BUNDLE = None def __init__( self, evaluation_context, applied_ptransform, input_committed_bundle, side_inputs): assert not side_inputs super().__init__( evaluation_context, applied_ptransform, input_committed_bundle, side_inputs) def start_bundle(self): self.gbk_items = collections.defaultdict(list) assert len(self._outputs) == 1 self.output_pcollection = list(self._outputs)[0] # The input type of a GroupByKey will be Tuple[Any, Any] or more specific. kv_type_hint = self._applied_ptransform.inputs[0].element_type key_type_hint = (kv_type_hint.tuple_types[0] if kv_type_hint else Any) self.key_coder = coders.registry.get_coder(key_type_hint) def process_element(self, element): if (isinstance(element, WindowedValue) and isinstance(element.value, collections.abc.Iterable) and len(element.value) == 2): k, v = element.value self.gbk_items[self.key_coder.encode(k)].append(v) else: raise TypeCheckError( 'Input to _GroupByKeyOnly must be a PCollection of ' 'windowed key-value pairs. Instead received: %r.' % element) def finish_bundle(self): bundles = [] bundle = None for encoded_k, vs in self.gbk_items.items(): if not bundle: bundle = self._evaluation_context.create_bundle(self.output_pcollection) bundles.append(bundle) kwi = KeyedWorkItem(encoded_k, elements=vs) bundle.add(GlobalWindows.windowed_value(kwi)) return TransformResult(self, bundles, [], None, None) class _StreamingGroupAlsoByWindowEvaluator(_TransformEvaluator): """TransformEvaluator for the _StreamingGroupAlsoByWindow transform. This evaluator is only used in streaming mode. In batch mode, the GroupAlsoByWindow operation is evaluated as a normal DoFn, as defined in transforms/core.py. """ def __init__( self, evaluation_context, applied_ptransform, input_committed_bundle, side_inputs): assert not side_inputs super().__init__( evaluation_context, applied_ptransform, input_committed_bundle, side_inputs) def start_bundle(self): assert len(self._outputs) == 1 self.output_pcollection = list(self._outputs)[0] self.driver = create_trigger_driver( self._applied_ptransform.transform.windowing, clock=self._evaluation_context._watermark_manager._clock) self.gabw_items = [] self.keyed_holds = {} # The input type (which is the same as the output type) of a # GroupAlsoByWindow will be Tuple[Any, Iter[Any]] or more specific. kv_type_hint = self._applied_ptransform.outputs[None].element_type key_type_hint = (kv_type_hint.tuple_types[0] if kv_type_hint else Any) self.key_coder = coders.registry.get_coder(key_type_hint) def process_element(self, element): kwi = element.value assert isinstance(kwi, KeyedWorkItem), kwi encoded_k, timer_firings, vs = ( kwi.encoded_key, kwi.timer_firings, kwi.elements) k = self.key_coder.decode(encoded_k) state = self._step_context.get_keyed_state(encoded_k) watermarks = self._evaluation_context._watermark_manager.get_watermarks( self._applied_ptransform) for timer_firing in timer_firings: for wvalue in self.driver.process_timer(timer_firing.window, timer_firing.name, timer_firing.time_domain, timer_firing.timestamp, state, watermarks.input_watermark): self.gabw_items.append(wvalue.with_value((k, wvalue.value))) if vs: for wvalue in self.driver.process_elements(state, vs, watermarks.output_watermark, watermarks.input_watermark): self.gabw_items.append(wvalue.with_value((k, wvalue.value))) self.keyed_holds[encoded_k] = state.get_earliest_hold() def finish_bundle(self): bundles = [] if self.gabw_items: bundle = self._evaluation_context.create_bundle(self.output_pcollection) for item in self.gabw_items: bundle.add(item) bundles.append(bundle) return TransformResult(self, bundles, [], None, self.keyed_holds) class _ProcessElementsEvaluator(_TransformEvaluator): """An evaluator for sdf_direct_runner.ProcessElements transform.""" # Maximum number of elements that will be produced by a Splittable DoFn before # a checkpoint is requested by the runner. DEFAULT_MAX_NUM_OUTPUTS = None # Maximum duration a Splittable DoFn will process an element before a # checkpoint is requested by the runner. DEFAULT_MAX_DURATION = 1 def __init__( self, evaluation_context, applied_ptransform, input_committed_bundle, side_inputs): super().__init__( evaluation_context, applied_ptransform, input_committed_bundle, side_inputs) process_elements_transform = applied_ptransform.transform assert isinstance(process_elements_transform, ProcessElements) # Replacing the do_fn of the transform with a wrapper do_fn that performs # SDF magic. transform = applied_ptransform.transform sdf = transform.sdf self._process_fn = transform.new_process_fn(sdf) transform.dofn = self._process_fn assert isinstance(self._process_fn, ProcessFn) self._process_fn.step_context = self._step_context process_element_invoker = ( SDFProcessElementInvoker( max_num_outputs=self.DEFAULT_MAX_NUM_OUTPUTS, max_duration=self.DEFAULT_MAX_DURATION)) self._process_fn.set_process_element_invoker(process_element_invoker) self._par_do_evaluator = _ParDoEvaluator( evaluation_context, applied_ptransform, input_committed_bundle, side_inputs, perform_dofn_pickle_test=False) self.keyed_holds = {} def start_bundle(self): self._par_do_evaluator.start_bundle() def process_element(self, element): assert isinstance(element, WindowedValue) assert len(element.windows) == 1 window = element.windows[0] if isinstance(element.value, KeyedWorkItem): key = element.value.encoded_key else: # If not a `KeyedWorkItem`, this must be a tuple where key is a randomly # generated key and the value is a `WindowedValue` that contains an # `ElementAndRestriction` object. assert isinstance(element.value, tuple) key = element.value[0] self._par_do_evaluator.process_element(element) state = self._step_context.get_keyed_state(key) self.keyed_holds[key] = state.get_state( window, self._process_fn.watermark_hold_tag) def finish_bundle(self): par_do_result = self._par_do_evaluator.finish_bundle() transform_result = TransformResult( self, par_do_result.uncommitted_output_bundles, par_do_result.unprocessed_bundles, par_do_result.counters, par_do_result.keyed_watermark_holds, par_do_result.undeclared_tag_values) for key, keyed_hold in self.keyed_holds.items(): transform_result.keyed_watermark_holds[key] = keyed_hold return transform_result ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/util.py0000644000175100017510000000636615206113032022165 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utility classes used by the DirectRunner. For internal use only. No backwards compatibility guarantees. """ # pytype: skip-file class TransformResult(object): """Result of evaluating an AppliedPTransform with a TransformEvaluator.""" def __init__( self, transform_evaluator, uncommitted_output_bundles, unprocessed_bundles, counters, keyed_watermark_holds, undeclared_tag_values=None): self.transform = transform_evaluator._applied_ptransform self.uncommitted_output_bundles = uncommitted_output_bundles self.unprocessed_bundles = unprocessed_bundles self.counters = counters # Mapping of key -> earliest hold timestamp or None. Keys should be # strings or None. # # For each key, we receive as its corresponding value the earliest # watermark hold for that key (the key can be None for global state), past # which the output watermark for the currently-executing step will not # advance. If the value is None or utils.timestamp.MAX_TIMESTAMP, the # watermark hold will be removed. self.keyed_watermark_holds = keyed_watermark_holds or {} # Only used when caching (materializing) all values is requested. self.undeclared_tag_values = undeclared_tag_values # Populated by the TransformExecutor. self.logical_metric_updates = None step_context = transform_evaluator._execution_context.get_step_context() self.partial_keyed_state = step_context.partial_keyed_state class TimerFiring(object): """A single instance of a fired timer.""" def __init__( self, encoded_key, window, name, time_domain, timestamp, dynamic_timer_tag=''): self.encoded_key = encoded_key self.window = window self.name = name self.time_domain = time_domain self.timestamp = timestamp self.dynamic_timer_tag = dynamic_timer_tag def __repr__(self): return 'TimerFiring({!r}, {!r}, {}, {}, {})'.format( self.encoded_key, self.name, self.time_domain, self.timestamp, self.dynamic_timer_tag) class KeyedWorkItem(object): """A keyed item that can either be a timer firing or a list of elements.""" def __init__(self, encoded_key, timer_firings=None, elements=None): self.encoded_key = encoded_key self.timer_firings = timer_firings or [] self.elements = elements or [] def __repr__(self): return 'KeyedWorkItem({!r}, {}, {})'.format( self.encoded_key, self.timer_firings, self.elements) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/direct/watermark_manager.py0000644000175100017510000003045515206113032024673 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Manages watermarks of PCollections and AppliedPTransforms.""" # pytype: skip-file import threading from typing import TYPE_CHECKING from typing import Iterable from apache_beam import pipeline from apache_beam import pvalue from apache_beam.pipeline import AppliedPTransform from apache_beam.runners.direct.util import TimerFiring from apache_beam.utils.timestamp import MAX_TIMESTAMP from apache_beam.utils.timestamp import MIN_TIMESTAMP from apache_beam.utils.timestamp import TIME_GRANULARITY from apache_beam.utils.timestamp import Timestamp if TYPE_CHECKING: from apache_beam.runners.direct.bundle_factory import _Bundle class WatermarkManager(object): """For internal use only; no backwards-compatibility guarantees. Tracks and updates watermarks for all AppliedPTransforms.""" WATERMARK_POS_INF = MAX_TIMESTAMP WATERMARK_NEG_INF = MIN_TIMESTAMP def __init__( self, clock, root_transforms, value_to_consumers, transform_keyed_states): self._clock = clock self._root_transforms = root_transforms self._value_to_consumers = value_to_consumers self._transform_keyed_states = transform_keyed_states # AppliedPTransform -> TransformWatermarks self._transform_to_watermarks: dict[AppliedPTransform, _TransformWatermarks] = {} for root_transform in root_transforms: self._transform_to_watermarks[root_transform] = _TransformWatermarks( self._clock, transform_keyed_states[root_transform], root_transform) for consumers in value_to_consumers.values(): for consumer in consumers: self._transform_to_watermarks[consumer] = _TransformWatermarks( self._clock, transform_keyed_states[consumer], consumer) for consumers in value_to_consumers.values(): for consumer in consumers: self._update_input_transform_watermarks(consumer) def _update_input_transform_watermarks( self, applied_ptransform: AppliedPTransform) -> None: assert isinstance(applied_ptransform, pipeline.AppliedPTransform) input_transform_watermarks = [] for input_pvalue in applied_ptransform.inputs: assert input_pvalue.producer or isinstance(input_pvalue, pvalue.PBegin) if input_pvalue.producer: input_transform_watermarks.append( self.get_watermarks(input_pvalue.producer)) self._transform_to_watermarks[ applied_ptransform].update_input_transform_watermarks( input_transform_watermarks) def get_watermarks( self, applied_ptransform: AppliedPTransform) -> '_TransformWatermarks': """Gets the input and output watermarks for an AppliedPTransform. If the applied_ptransform has not processed any elements, return a watermark with minimum value. Args: applied_ptransform: AppliedPTransform to get the watermarks for. Returns: A snapshot (TransformWatermarks) of the input watermark and output watermark for the provided transform. """ # TODO(altay): Composite transforms should have a composite watermark. Until # then they are represented by their last transform. while applied_ptransform.parts: applied_ptransform = applied_ptransform.parts[-1] return self._transform_to_watermarks[applied_ptransform] def update_watermarks( self, completed_committed_bundle: '_Bundle', applied_ptransform: AppliedPTransform, completed_timers, outputs, unprocessed_bundles, keyed_earliest_holds, side_inputs_container): assert isinstance(applied_ptransform, pipeline.AppliedPTransform) self._update_pending( completed_committed_bundle, applied_ptransform, completed_timers, outputs, unprocessed_bundles) tw = self.get_watermarks(applied_ptransform) tw.hold(keyed_earliest_holds) return self._refresh_watermarks(applied_ptransform, side_inputs_container) def _update_pending( self, input_committed_bundle, applied_ptransform: AppliedPTransform, completed_timers, output_committed_bundles: Iterable['_Bundle'], unprocessed_bundles: Iterable['_Bundle']): """Updated list of pending bundles for the given AppliedPTransform.""" # Update pending elements. Filter out empty bundles. They do not impact # watermarks and should not trigger downstream execution. for output in output_committed_bundles: if output.has_elements(): if output.pcollection in self._value_to_consumers: consumers = self._value_to_consumers[output.pcollection] for consumer in consumers: consumer_tw = self._transform_to_watermarks[consumer] consumer_tw.add_pending(output) completed_tw = self._transform_to_watermarks[applied_ptransform] completed_tw.update_timers(completed_timers) for unprocessed_bundle in unprocessed_bundles: completed_tw.add_pending(unprocessed_bundle) assert input_committed_bundle or applied_ptransform in self._root_transforms if input_committed_bundle and input_committed_bundle.has_elements(): completed_tw.remove_pending(input_committed_bundle) def _refresh_watermarks(self, applied_ptransform, side_inputs_container): assert isinstance(applied_ptransform, pipeline.AppliedPTransform) unblocked_tasks = [] tw = self.get_watermarks(applied_ptransform) if tw.refresh(): for pval in applied_ptransform.outputs.values(): if isinstance(pval, pvalue.DoOutputsTuple): pvals = (v for v in pval) else: pvals = (pval, ) for v in pvals: if v in self._value_to_consumers: # If there are downstream consumers consumers = self._value_to_consumers[v] for consumer in consumers: unblocked_tasks.extend( self._refresh_watermarks(consumer, side_inputs_container)) # Notify the side_inputs_container. unblocked_tasks.extend( side_inputs_container. update_watermarks_for_transform_and_unblock_tasks( applied_ptransform, tw)) return unblocked_tasks def extract_all_timers( self) -> tuple[list[tuple[AppliedPTransform, list[TimerFiring]]], bool]: """Extracts fired timers for all transforms and reports if there are any timers set.""" all_timers: list[tuple[AppliedPTransform, list[TimerFiring]]] = [] has_realtime_timer = False for applied_ptransform, tw in self._transform_to_watermarks.items(): fired_timers, had_realtime_timer = tw.extract_transform_timers() if fired_timers: # We should sort the timer firings, so they are fired in order. fired_timers.sort(key=lambda ft: ft.timestamp) all_timers.append((applied_ptransform, fired_timers)) if (had_realtime_timer and tw.output_watermark < WatermarkManager.WATERMARK_POS_INF): has_realtime_timer = True return all_timers, has_realtime_timer class _TransformWatermarks(object): """Tracks input and output watermarks for an AppliedPTransform.""" def __init__(self, clock, keyed_states, transform): self._clock = clock self._keyed_states = keyed_states self._input_transform_watermarks: list[_TransformWatermarks] = [] self._input_watermark = WatermarkManager.WATERMARK_NEG_INF self._output_watermark = WatermarkManager.WATERMARK_NEG_INF self._keyed_earliest_holds = {} # Scheduled bundles targeted for this transform. self._pending: set['_Bundle'] = set() self._fired_timers = set() self._lock = threading.Lock() self._label = str(transform) def update_input_transform_watermarks( self, input_transform_watermarks: list['_TransformWatermarks']) -> None: with self._lock: self._input_transform_watermarks = input_transform_watermarks def update_timers(self, completed_timers): with self._lock: for timer_firing in completed_timers: self._fired_timers.remove(timer_firing) @property def input_watermark(self) -> Timestamp: with self._lock: return self._input_watermark @property def output_watermark(self) -> Timestamp: with self._lock: return self._output_watermark def hold(self, keyed_earliest_holds): with self._lock: for key, hold_value in keyed_earliest_holds.items(): self._keyed_earliest_holds[key] = hold_value if (hold_value is None or hold_value == WatermarkManager.WATERMARK_POS_INF): del self._keyed_earliest_holds[key] def add_pending(self, pending: '_Bundle') -> None: with self._lock: self._pending.add(pending) def remove_pending(self, completed: '_Bundle') -> None: with self._lock: # Ignore repeated removes. This will happen if a transform has a repeated # input. if completed in self._pending: self._pending.remove(completed) def refresh(self) -> bool: """Refresh the watermark for a given transform. This method looks at the watermark coming from all input PTransforms, and the timestamp of the minimum element, as well as any watermark holds. Returns: True if the watermark has advanced, and False if it has not. """ with self._lock: min_pending_timestamp = WatermarkManager.WATERMARK_POS_INF has_pending_elements = False for input_bundle in self._pending: # TODO(ccy): we can have the Bundle class keep track of the minimum # timestamp so we don't have to do an iteration here. for wv in input_bundle.get_elements_iterable(): has_pending_elements = True if wv.timestamp < min_pending_timestamp: min_pending_timestamp = wv.timestamp # If there is a pending element with a certain timestamp, we can at most # advance our watermark to the maximum timestamp less than that # timestamp. pending_holder = WatermarkManager.WATERMARK_POS_INF if has_pending_elements: pending_holder = min_pending_timestamp - TIME_GRANULARITY input_watermarks = [ tw.output_watermark for tw in self._input_transform_watermarks ] input_watermarks.append(WatermarkManager.WATERMARK_POS_INF) producer_watermark = min(input_watermarks) self._input_watermark = max( self._input_watermark, min(pending_holder, producer_watermark)) earliest_hold = WatermarkManager.WATERMARK_POS_INF for hold in self._keyed_earliest_holds.values(): if hold < earliest_hold: earliest_hold = hold new_output_watermark = min(self._input_watermark, earliest_hold) advanced = new_output_watermark > self._output_watermark self._output_watermark = new_output_watermark return advanced @property def synchronized_processing_output_time(self): return self._clock.time() def extract_transform_timers(self) -> tuple[list[TimerFiring], bool]: """Extracts fired timers and reports of any timers set per transform.""" with self._lock: fired_timers = [] has_realtime_timer = False for encoded_key, state in self._keyed_states.items(): timers, had_realtime_timer = state.get_timers( watermark=self._input_watermark, processing_time=self._clock.time()) if had_realtime_timer: has_realtime_timer = True for expired in timers: window, (name, time_domain, timestamp, dynamic_timer_tag) = expired fired_timers.append( TimerFiring( encoded_key, window, name, time_domain, timestamp, dynamic_timer_tag=dynamic_timer_tag)) self._fired_timers.update(fired_timers) return fired_timers, has_realtime_timer ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.6224334 apache_beam-2.74.0/apache_beam/runners/interactive/0000755000175100017510000000000015206113160021670 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/__init__.py0000644000175100017510000000142015206113032023774 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/augmented_pipeline.py0000644000175100017510000001177515206113032026111 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module to augment interactive flavor into the given pipeline. For internal use only; no backward-compatibility guarantees. """ # pytype: skip-file import copy from typing import Optional import apache_beam as beam from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.interactive import background_caching_job from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.caching.cacheable import Cacheable from apache_beam.runners.interactive.caching.read_cache import ReadCache from apache_beam.runners.interactive.caching.write_cache import WriteCache class AugmentedPipeline: """A pipeline with augmented interactive flavor that caches intermediate PCollections defined by the user, reads computed PCollections as source and prunes unnecessary pipeline parts for fast computation. """ def __init__( self, user_pipeline: beam.Pipeline, pcolls: Optional[set[beam.pvalue.PCollection]] = None): """ Initializes a pipelilne for augmenting interactive flavor. Args: user_pipeline: a beam.Pipeline instance defined by the user. pcolls: cacheable pcolls to be computed/retrieved. If the set is empty, all intermediate pcolls assigned to variables are applicable. """ assert not pcolls or all(pcoll.pipeline is user_pipeline for pcoll in pcolls), 'All %s need to belong to %s' % (pcolls, user_pipeline) self._user_pipeline = user_pipeline self._pcolls = pcolls self._cache_manager = ie.current_env().get_cache_manager( self._user_pipeline, create_if_absent=True) if background_caching_job.has_source_to_cache(self._user_pipeline): self._cache_manager = ie.current_env().get_cache_manager( self._user_pipeline) _, self._context = self._user_pipeline.to_runner_api(return_context=True) self._context.component_id_map = copy.copy( self._user_pipeline.component_id_map) self._cacheables = self.cacheables() @property def augmented_pipeline(self) -> beam_runner_api_pb2.Pipeline: return self.augment() # TODO(https://github.com/apache/beam/issues/20526): Support generating a # background recording job that contains unbound source recording transforms # only. @property def background_recording_pipeline(self) -> beam_runner_api_pb2.Pipeline: raise NotImplementedError def cacheables(self) -> dict[beam.pvalue.PCollection, Cacheable]: """Finds all the cacheable intermediate PCollections in the pipeline with their metadata. """ c = {} for watching in ie.current_env().watching(): for key, val in watching: if (isinstance(val, beam.pvalue.PCollection) and val.pipeline is self._user_pipeline and (not self._pcolls or val in self._pcolls)): c[val] = Cacheable( var=key, pcoll=val, version=str(id(val)), producer_version=str(id(val.producer))) return c def augment(self) -> beam_runner_api_pb2.Pipeline: """Augments the pipeline with cache. Always calculates a new result. For a cacheable PCollection, if cache exists, read cache; else, write cache. """ pipeline = self._user_pipeline.to_runner_api() # Find pcolls eligible for reading or writing cache. readcache_pcolls = set() for pcoll, cacheable in self._cacheables.items(): key = repr(cacheable.to_key()) if (self._cache_manager.exists('full', key) and pcoll in ie.current_env().computed_pcollections): readcache_pcolls.add(pcoll) writecache_pcolls = set( self._cacheables.keys()).difference(readcache_pcolls) # Wire in additional transforms to read cache and write cache. for readcache_pcoll in readcache_pcolls: ReadCache( pipeline, self._context, self._cache_manager, self._cacheables[readcache_pcoll]).read_cache() for writecache_pcoll in writecache_pcolls: WriteCache( pipeline, self._context, self._cache_manager, self._cacheables[writecache_pcoll]).write_cache() # TODO(https://github.com/apache/beam/issues/20526): Support streaming, add # pruning logic, and integrate pipeline fragment logic. return pipeline ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/augmented_pipeline_test.py0000644000175100017510000000532515206113032027142 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for augmented_pipeline module.""" # pytest: skip-file import unittest import apache_beam as beam from apache_beam.runners.interactive import augmented_pipeline as ap from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie class CacheableTest(unittest.TestCase): def setUp(self): ie.new_env() def test_find_all_cacheables(self): p = beam.Pipeline() cacheable_pcoll_1 = p | beam.Create([1, 2, 3]) cacheable_pcoll_2 = cacheable_pcoll_1 | beam.Map(lambda x: x * x) ib.watch(locals()) aug_p = ap.AugmentedPipeline(p) cacheables = aug_p.cacheables() self.assertIn(cacheable_pcoll_1, cacheables) self.assertIn(cacheable_pcoll_2, cacheables) def test_ignore_cacheables(self): p = beam.Pipeline() cacheable_pcoll_1 = p | 'cacheable_pcoll_1' >> beam.Create([1, 2, 3]) cacheable_pcoll_2 = p | 'cacheable_pcoll_2' >> beam.Create([4, 5, 6]) ib.watch(locals()) aug_p = ap.AugmentedPipeline(p, (cacheable_pcoll_1, )) cacheables = aug_p.cacheables() self.assertIn(cacheable_pcoll_1, cacheables) self.assertNotIn(cacheable_pcoll_2, cacheables) def test_ignore_pcoll_from_other_pipeline(self): p = beam.Pipeline() p2 = beam.Pipeline() cacheable_from_p2 = p2 | beam.Create([1, 2, 3]) ib.watch(locals()) aug_p = ap.AugmentedPipeline(p) cacheables = aug_p.cacheables() self.assertNotIn(cacheable_from_p2, cacheables) class AugmentTest(unittest.TestCase): def setUp(self): ie.new_env() def test_error_when_pcolls_from_mixed_pipelines(self): p = beam.Pipeline() cacheable_from_p = p | beam.Create([1, 2, 3]) p2 = beam.Pipeline() cacheable_from_p2 = p2 | beam.Create([1, 2, 3]) ib.watch(locals()) self.assertRaises( AssertionError, lambda: ap.AugmentedPipeline(p, (cacheable_from_p, cacheable_from_p2))) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/background_caching_job.py0000644000175100017510000003404715206113032026675 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module to build and run background source recording jobs. For internal use only; no backwards-compatibility guarantees. A background source recording job is a job that records events for all recordable sources of a given pipeline. With Interactive Beam, one such job is started when a pipeline run happens (which produces a main job in contrast to the background source recording job) and meets the following conditions: #. The pipeline contains recordable sources, configured through interactive_beam.options.recordable_sources. #. No such background job is running. #. No such background job has completed successfully and the cached events are still valid (invalidated when recordable sources change in the pipeline). Once started, the background source recording job runs asynchronously until it hits some recording limit configured in interactive_beam.options. Meanwhile, the main job and future main jobs from the pipeline will run using the deterministic replayable recorded events until they are invalidated. """ # pytype: skip-file import logging import threading import time import apache_beam as beam from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import utils from apache_beam.runners.interactive.caching import streaming_cache from apache_beam.runners.runner import PipelineState _LOGGER = logging.getLogger(__name__) class BackgroundCachingJob(object): """A simple abstraction that controls necessary components of a timed and space limited background source recording job. A background source recording job successfully completes source data recording in 2 conditions: #. The job is finite and runs into DONE state; #. The job is infinite but hits an interactive_beam.options configured limit and gets cancelled into CANCELLED/CANCELLING state. In both situations, the background source recording job should be treated as done successfully. """ def __init__(self, pipeline_result, limiters): self._pipeline_result = pipeline_result self._result_lock = threading.RLock() self._condition_checker = threading.Thread( target=self._background_caching_job_condition_checker, daemon=True) # Limiters are checks s.t. if any are triggered then the background caching # job gets cancelled. self._limiters = limiters self._condition_checker.start() def _background_caching_job_condition_checker(self): while True: with self._result_lock: if PipelineState.is_terminal(self._pipeline_result.state): break if self._should_end_condition_checker(): self.cancel() break time.sleep(0.5) def _should_end_condition_checker(self): return any(l.is_triggered() for l in self._limiters) def is_done(self): with self._result_lock: is_terminated = self._pipeline_result.state in ( PipelineState.DONE, PipelineState.CANCELLED) is_triggered = self._should_end_condition_checker() is_cancelling = self._pipeline_result.state is PipelineState.CANCELLING return is_terminated or (is_triggered and is_cancelling) def is_running(self): with self._result_lock: return self._pipeline_result.state is PipelineState.RUNNING def cancel(self): """Cancels this background source recording job. """ with self._result_lock: if not PipelineState.is_terminal(self._pipeline_result.state): try: self._pipeline_result.cancel() except NotImplementedError: # Ignore the cancel invocation if it is never implemented by the # runner. pass @property def state(self): with self._result_lock: return self._pipeline_result.state def attempt_to_run_background_caching_job( runner, user_pipeline, options=None, limiters=None): """Attempts to run a background source recording job for a user-defined pipeline. Returns True if a job was started, False otherwise. The pipeline result is automatically tracked by Interactive Beam in case future cancellation/cleanup is needed. """ if is_background_caching_job_needed(user_pipeline): # Cancel non-terminal jobs if there is any before starting a new one. attempt_to_cancel_background_caching_job(user_pipeline) # Cancel the gRPC server serving the test stream if there is one. attempt_to_stop_test_stream_service(user_pipeline) # TODO(BEAM-8335): refactor background source recording job logic from # pipeline_instrument module to this module and aggregate tests. from apache_beam.runners.interactive import pipeline_instrument as instr runner_pipeline = beam.pipeline.Pipeline.from_runner_api( user_pipeline.to_runner_api(), runner, options) ie.current_env().add_derived_pipeline(user_pipeline, runner_pipeline) background_caching_job_result = beam.pipeline.Pipeline.from_runner_api( instr.build_pipeline_instrument( runner_pipeline).background_caching_pipeline_proto(), runner, options).run() recording_limiters = ( limiters if limiters else ie.current_env().options.capture_control.limiters()) ie.current_env().set_background_caching_job( user_pipeline, BackgroundCachingJob( background_caching_job_result, limiters=recording_limiters)) return True return False def is_background_caching_job_needed(user_pipeline): """Determines if a background source recording job needs to be started. It does several state checks and recording state changes throughout the process. It is not idempotent to simplify the usage. """ job = ie.current_env().get_background_caching_job(user_pipeline) # Checks if the pipeline contains any source that needs to be cached. need_cache = has_source_to_cache(user_pipeline) # If this is True, we can invalidate a previous done/running job if there is # one. cache_changed = is_source_to_cache_changed(user_pipeline) # When recording replay is disabled, cache is always needed for recordable # sources (if any). if need_cache and not ie.current_env().options.enable_recording_replay: from apache_beam.runners.interactive.options import capture_control capture_control.evict_captured_data() return True return ( need_cache and # Checks if it's the first time running a job from the pipeline. ( not job or # Or checks if there is no previous job. # DONE means a previous job has completed successfully and the # cached events might still be valid. not ( job.is_done() or # RUNNING means a previous job has been started and is still # running. job.is_running()) or # Or checks if we can invalidate the previous job. cache_changed)) def is_cache_complete(pipeline_id: str) -> bool: """Returns True if the backgrond cache for the given pipeline is done. """ user_pipeline = ie.current_env().pipeline_id_to_pipeline(pipeline_id) job = ie.current_env().get_background_caching_job(user_pipeline) is_done = job and job.is_done() cache_changed = is_source_to_cache_changed( user_pipeline, update_cached_source_signature=False) # Stop reading from the cache if the background job is done or the underlying # cache signature changed that requires a new background source recording job. return is_done or cache_changed def has_source_to_cache(user_pipeline): """Determines if a user-defined pipeline contains any source that need to be cached. If so, also immediately wrap current cache manager held by current interactive environment into a streaming cache if this has not been done. The wrapping doesn't invalidate existing cache in any way. This can help determining if a background source recording job is needed to write cache for sources and if a test stream service is needed to serve the cache. Throughout the check, if source-to-cache has changed from the last check, it also cleans up the invalidated cache early on. """ # TODO(BEAM-8335): we temporarily only cache replaceable unbounded sources. # Add logic for other cacheable sources here when they are available. has_cache = utils.has_unbounded_sources(user_pipeline) if has_cache: if not isinstance(ie.current_env().get_cache_manager(user_pipeline, create_if_absent=True), streaming_cache.StreamingCache): file_based_cm = ie.current_env().get_cache_manager(user_pipeline) cache_dir = file_based_cm._cache_dir cache_root = ie.current_env().options.cache_root if cache_root: if cache_root.startswith('gs://'): raise ValueError( 'GCS cache paths are not currently supported for ' 'streaming pipelines.') cache_dir = cache_root ie.current_env().set_cache_manager( streaming_cache.StreamingCache( cache_dir, is_cache_complete=is_cache_complete, sample_resolution_sec=1.0, saved_pcoders=file_based_cm._saved_pcoders), user_pipeline) return has_cache def attempt_to_cancel_background_caching_job(user_pipeline): """Attempts to cancel background source recording job for a user-defined pipeline. If no background source recording job needs to be cancelled, NOOP. Otherwise, cancel such job. """ job = ie.current_env().get_background_caching_job(user_pipeline) if job: job.cancel() def attempt_to_stop_test_stream_service(user_pipeline): """Attempts to stop the gRPC server/service serving the test stream. If there is no such server started, NOOP. Otherwise, stop it. """ if is_a_test_stream_service_running(user_pipeline): ie.current_env().evict_test_stream_service_controller(user_pipeline).stop() def is_a_test_stream_service_running(user_pipeline): """Checks to see if there is a gPRC server/service running that serves the test stream to any job started from the given user_pipeline. """ return ie.current_env().get_test_stream_service_controller( user_pipeline) is not None def is_source_to_cache_changed( user_pipeline, update_cached_source_signature=True): """Determines if there is any change in the sources that need to be cached used by the user-defined pipeline. Due to the expensiveness of computations and for the simplicity of usage, this function is not idempotent because Interactive Beam automatically discards previously tracked signature of transforms and tracks the current signature of transforms for the user-defined pipeline if there is any change. When it's True, there is addition/deletion/mutation of source transforms that requires a new background source recording job. """ # By default gets empty set if the user_pipeline is first time seen because # we can treat it as adding transforms. recorded_signature = ie.current_env().get_cached_source_signature( user_pipeline) current_signature = extract_source_to_cache_signature(user_pipeline) is_changed = not current_signature.issubset(recorded_signature) # The computation of extract_unbounded_source_signature is expensive, track on # change by default. if is_changed and update_cached_source_signature: options = ie.current_env().options # No info needed when recording replay is disabled. if options.enable_recording_replay: if not recorded_signature: def sizeof_fmt(num, suffix='B'): for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']: if abs(num) < 1000.0: return "%3.1f%s%s" % (num, unit, suffix) num /= 1000.0 return "%.1f%s%s" % (num, 'Yi', suffix) _LOGGER.info( 'Interactive Beam has detected unbounded sources in your pipeline. ' 'In order to have a deterministic replay, a segment of data will ' 'be recorded from all sources for %s seconds or until a total of ' '%s have been written to disk.', options.recording_duration.total_seconds(), sizeof_fmt(options.recording_size_limit)) else: _LOGGER.info( 'Interactive Beam has detected a new streaming source was ' 'added to the pipeline. In order for the cached streaming ' 'data to start at the same time, all recorded data has been ' 'cleared and a new segment of data will be recorded.') ie.current_env().cleanup(user_pipeline) ie.current_env().set_cached_source_signature( user_pipeline, current_signature) ie.current_env().add_user_pipeline(user_pipeline) return is_changed def extract_source_to_cache_signature(user_pipeline): """Extracts a set of signature for sources that need to be cached in the user-defined pipeline. A signature is a str representation of urn and payload of a source. """ # TODO(BEAM-8335): we temporarily only cache replaceable unbounded sources. # Add logic for other cacheable sources here when they are available. unbounded_sources_as_applied_transforms = utils.unbounded_sources( user_pipeline) unbounded_sources_as_ptransforms = set( map(lambda x: x.transform, unbounded_sources_as_applied_transforms)) _, context = user_pipeline.to_runner_api(return_context=True) signature = set( map( lambda transform: str(transform.to_runner_api(context)), unbounded_sources_as_ptransforms)) return signature ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/background_caching_job_test.py0000644000175100017510000003314415206113032027731 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.runners.interactive.background_caching_job.""" # pytype: skip-file import unittest from unittest.mock import patch import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.pipeline import PipelineVisitor from apache_beam.runners import runner from apache_beam.runners.interactive import background_caching_job as bcj from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import interactive_runner from apache_beam.runners.interactive.caching.streaming_cache import StreamingCache from apache_beam.runners.interactive.testing.mock_ipython import mock_get_ipython from apache_beam.runners.interactive.testing.test_cache_manager import FileRecordsBuilder from apache_beam.testing.test_stream import TestStream from apache_beam.testing.test_stream_service import TestStreamServiceController from apache_beam.transforms.window import TimestampedValue _FOO_PUBSUB_SUB = 'projects/test-project/subscriptions/foo' _BAR_PUBSUB_SUB = 'projects/test-project/subscriptions/bar' _TEST_CACHE_KEY = 'test' def _build_a_test_stream_pipeline(): test_stream = ( TestStream().advance_watermark_to(0).add_elements([ TimestampedValue('a', 1) ]).advance_processing_time(5).advance_watermark_to_infinity()) p = beam.Pipeline(runner=interactive_runner.InteractiveRunner()) events = p | test_stream # pylint: disable=possibly-unused-variable ib.watch(locals()) return p def _build_an_empty_stream_pipeline(): pipeline_options = PipelineOptions(streaming=True) p = beam.Pipeline( interactive_runner.InteractiveRunner(), options=pipeline_options) ib.watch({'pipeline': p}) return p def _setup_test_streaming_cache(pipeline): cache_manager = StreamingCache(cache_dir=None) ie.current_env().set_cache_manager(cache_manager, pipeline) builder = FileRecordsBuilder(tag=_TEST_CACHE_KEY) (builder .advance_watermark(watermark_secs=0) .advance_processing_time(5) .add_element(element='a', event_time_secs=1) .advance_watermark(watermark_secs=100) .advance_processing_time(10)) # yapf: disable cache_manager.write(builder.build(), _TEST_CACHE_KEY) @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') class BackgroundCachingJobTest(unittest.TestCase): def tearDown(self): ie.new_env() # TODO(BEAM-8335): remove the patches when there are appropriate test sources # that meet the boundedness checks. @patch( 'apache_beam.runners.interactive.background_caching_job' '.has_source_to_cache', lambda x: True) # Disable the clean up so that we can keep the test streaming cache. @patch( 'apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.cleanup', lambda x, y: None) def test_background_caching_job_starts_when_none_such_job_exists(self): # Create a fake PipelineResult and PipelineRunner. This is because we want # to test whether the BackgroundCachingJob can be started without having to # rely on a real pipeline run. class FakePipelineResult(beam.runners.runner.PipelineResult): def wait_until_finish(self): return class FakePipelineRunner(beam.runners.PipelineRunner): def run_pipeline(self, pipeline, options): return FakePipelineResult(beam.runners.runner.PipelineState.RUNNING) p = beam.Pipeline( runner=interactive_runner.InteractiveRunner(FakePipelineRunner()), options=PipelineOptions(streaming=True)) # pylint: disable=possibly-unused-variable elems = p | 'Read' >> beam.io.ReadFromPubSub(subscription=_FOO_PUBSUB_SUB) ib.watch(locals()) _setup_test_streaming_cache(p) p.run() self.assertIsNotNone(ie.current_env().get_background_caching_job(p)) expected_cached_source_signature = bcj.extract_source_to_cache_signature(p) # This is to check whether the cached source signature is set correctly # when the background caching job is started. self.assertEqual( expected_cached_source_signature, ie.current_env().get_cached_source_signature(p)) @patch( 'apache_beam.runners.interactive.background_caching_job' '.has_source_to_cache', lambda x: False) def test_background_caching_job_not_start_for_batch_pipeline(self): p = beam.Pipeline() # pylint: disable=expression-not-assigned p | beam.Create([]) p.run() self.assertIsNone(ie.current_env().get_background_caching_job(p)) @patch( 'apache_beam.runners.interactive.background_caching_job' '.has_source_to_cache', lambda x: True) # Disable the clean up so that we can keep the test streaming cache. @patch( 'apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.cleanup', lambda x, y: None) def test_background_caching_job_not_start_when_such_job_exists(self): p = _build_a_test_stream_pipeline() _setup_test_streaming_cache(p) a_running_background_caching_job = bcj.BackgroundCachingJob( runner.PipelineResult(runner.PipelineState.RUNNING), limiters=[]) ie.current_env().set_background_caching_job( p, a_running_background_caching_job) main_job_result = p.run() # No background caching job is started so result is still the running one. self.assertIs( a_running_background_caching_job, ie.current_env().get_background_caching_job(p)) # A new main job is started so result of the main job is set. self.assertIs(main_job_result, ie.current_env().pipeline_result(p)) @patch( 'apache_beam.runners.interactive.background_caching_job' '.has_source_to_cache', lambda x: True) # Disable the clean up so that we can keep the test streaming cache. @patch( 'apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.cleanup', lambda x, y: None) def test_background_caching_job_not_start_when_such_job_is_done(self): p = _build_a_test_stream_pipeline() _setup_test_streaming_cache(p) a_done_background_caching_job = bcj.BackgroundCachingJob( runner.PipelineResult(runner.PipelineState.DONE), limiters=[]) ie.current_env().set_background_caching_job( p, a_done_background_caching_job) main_job_result = p.run() # No background caching job is started so result is still the running one. self.assertIs( a_done_background_caching_job, ie.current_env().get_background_caching_job(p)) # A new main job is started so result of the main job is set. self.assertIs(main_job_result, ie.current_env().pipeline_result(p)) @patch('IPython.get_ipython', new_callable=mock_get_ipython) def test_source_to_cache_changed_when_pipeline_is_first_time_seen(self, cell): with cell: # Cell 1 pipeline = _build_an_empty_stream_pipeline() with cell: # Cell 2 read_foo = pipeline | 'Read' >> beam.io.ReadFromPubSub( subscription=_FOO_PUBSUB_SUB) ib.watch({'read_foo': read_foo}) self.assertTrue(bcj.is_source_to_cache_changed(pipeline)) @patch('IPython.get_ipython', new_callable=mock_get_ipython) def test_source_to_cache_changed_when_new_source_is_added(self, cell): with cell: # Cell 1 pipeline = _build_an_empty_stream_pipeline() read_foo = pipeline | 'Read' >> beam.io.ReadFromPubSub( subscription=_FOO_PUBSUB_SUB) ib.watch({'read_foo': read_foo}) # Sets the signature for current pipeline state. ie.current_env().set_cached_source_signature( pipeline, bcj.extract_source_to_cache_signature(pipeline)) self.assertFalse(bcj.is_cache_complete(str(id(pipeline)))) with cell: # Cell 2 read_bar = pipeline | 'Read' >> beam.io.ReadFromPubSub( subscription=_BAR_PUBSUB_SUB) ib.watch({'read_bar': read_bar}) self.assertTrue(bcj.is_cache_complete(str(id(pipeline)))) self.assertTrue(bcj.is_source_to_cache_changed(pipeline)) @patch('IPython.get_ipython', new_callable=mock_get_ipython) def test_source_to_cache_changed_when_source_is_altered(self, cell): with cell: # Cell 1 pipeline = _build_an_empty_stream_pipeline() transform = beam.io.ReadFromPubSub(subscription=_FOO_PUBSUB_SUB) read_foo = pipeline | 'Read' >> transform ib.watch({'read_foo': read_foo}) # Sets the signature for current pipeline state. ie.current_env().set_cached_source_signature( pipeline, bcj.extract_source_to_cache_signature(pipeline)) with cell: # Cell 2 from apache_beam.io.gcp.pubsub import _PubSubSource # Alter the transform. transform._source = _PubSubSource(subscription=_BAR_PUBSUB_SUB) self.assertTrue(bcj.is_source_to_cache_changed(pipeline)) @patch('IPython.get_ipython', new_callable=mock_get_ipython) def test_source_to_cache_not_changed_for_same_source(self, cell): with cell: # Cell 1 pipeline = _build_an_empty_stream_pipeline() transform = beam.io.ReadFromPubSub(subscription=_FOO_PUBSUB_SUB) with cell: # Cell 2 read_foo_1 = pipeline | 'Read' >> transform ib.watch({'read_foo_1': read_foo_1}) # Sets the signature for current pipeline state. ie.current_env().set_cached_source_signature( pipeline, bcj.extract_source_to_cache_signature(pipeline)) with cell: # Cell 3 # Apply exactly the same transform and the same instance. read_foo_2 = pipeline | 'Read' >> transform ib.watch({'read_foo_2': read_foo_2}) self.assertFalse(bcj.is_source_to_cache_changed(pipeline)) with cell: # Cell 4 # Apply the same transform but represented in a different instance. # The signature representing the urn and payload is still the same, so it # is not treated as a new unbounded source. read_foo_3 = pipeline | 'Read' >> beam.io.ReadFromPubSub( subscription=_FOO_PUBSUB_SUB) ib.watch({'read_foo_3': read_foo_3}) self.assertFalse(bcj.is_source_to_cache_changed(pipeline)) @patch('IPython.get_ipython', new_callable=mock_get_ipython) def test_source_to_cache_not_changed_when_source_is_removed(self, cell): with cell: # Cell 1 pipeline = _build_an_empty_stream_pipeline() foo_transform = beam.io.ReadFromPubSub(subscription=_FOO_PUBSUB_SUB) bar_transform = beam.io.ReadFromPubSub(subscription=_BAR_PUBSUB_SUB) with cell: # Cell 2 read_foo = pipeline | 'Read' >> foo_transform ib.watch({'read_foo': read_foo}) signature_with_only_foo = bcj.extract_source_to_cache_signature(pipeline) with cell: # Cell 3 read_bar = pipeline | 'Read' >> bar_transform ib.watch({'read_bar': read_bar}) self.assertTrue(bcj.is_source_to_cache_changed(pipeline)) signature_with_foo_bar = ie.current_env().get_cached_source_signature( pipeline) self.assertNotEqual(signature_with_only_foo, signature_with_foo_bar) class BarPruneVisitor(PipelineVisitor): def enter_composite_transform(self, transform_node): pruned_parts = list(transform_node.parts) for part in transform_node.parts: if part.transform is bar_transform: pruned_parts.remove(part) transform_node.parts = tuple(pruned_parts) self.visit_transform(transform_node) def visit_transform(self, transform_node): if transform_node.transform is bar_transform: transform_node.parent = None v = BarPruneVisitor() pipeline.visit(v) signature_after_pruning_bar = bcj.extract_source_to_cache_signature( pipeline) self.assertEqual(signature_with_only_foo, signature_after_pruning_bar) self.assertFalse(bcj.is_source_to_cache_changed(pipeline)) def test_determine_a_test_stream_service_running(self): pipeline = _build_an_empty_stream_pipeline() test_stream_service = TestStreamServiceController(reader=None) test_stream_service.start() ie.current_env().set_test_stream_service_controller( pipeline, test_stream_service) self.assertTrue(bcj.is_a_test_stream_service_running(pipeline)) # the test_stream_service will be cleaned up on teardown. def test_stop_a_running_test_stream_service(self): pipeline = _build_an_empty_stream_pipeline() test_stream_service = TestStreamServiceController(reader=None) test_stream_service.start() ie.current_env().set_test_stream_service_controller( pipeline, test_stream_service) bcj.attempt_to_stop_test_stream_service(pipeline) self.assertFalse(bcj.is_a_test_stream_service_running(pipeline)) @patch( 'apache_beam.testing.test_stream_service.TestStreamServiceController' '.stop') def test_noop_when_no_test_stream_service_running(self, _mocked_stop): pipeline = _build_an_empty_stream_pipeline() self.assertFalse(bcj.is_a_test_stream_service_running(pipeline)) bcj.attempt_to_stop_test_stream_service(pipeline) _mocked_stop.assert_not_called() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/cache_manager.py0000644000175100017510000003324515206113032025004 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import base64 import collections import os import tempfile from urllib.parse import quote from urllib.parse import unquote_to_bytes import apache_beam as beam from apache_beam import coders from apache_beam.io import filesystems from apache_beam.io import textio from apache_beam.io import tfrecordio from apache_beam.testing import test_stream from apache_beam.transforms import combiners class CacheManager(object): """Abstract class for caching PCollections. A PCollection cache is identified by labels, which consist of a prefix (either 'full' or 'sample') and a cache_label which is a hash of the PCollection derivation. """ def exists(self, *labels): # type (*str) -> bool """Returns if the PCollection cache exists.""" raise NotImplementedError def is_latest_version(self, version, *labels): # type (str, *str) -> bool """Returns if the given version number is the latest.""" return version == self._latest_version(*labels) def _latest_version(self, *labels): # type (*str) -> str """Returns the latest version number of the PCollection cache.""" raise NotImplementedError def read(self, *labels, **args): # type (*str, Dict[str, Any]) -> Tuple[str, Generator[Any]] """Return the PCollection as a list as well as the version number. Args: *labels: List of labels for PCollection instance. **args: Dict of additional arguments. Currently only 'tail' as a boolean. When tail is True, will wait and read new elements until the cache is complete. Returns: A tuple containing an iterator for the items in the PCollection and the version number. It is possible that the version numbers from read() and_latest_version() are different. This usually means that the cache's been evicted (thus unavailable => read() returns version = -1), but it had reached version n before eviction. """ raise NotImplementedError def write(self, value, *labels): # type (Any, *str) -> None """Writes the value to the given cache. Args: value: An encodable (with corresponding PCoder) value *labels: List of labels for PCollection instance """ raise NotImplementedError def clear(self, *labels): # type (*str) -> Boolean """Clears the cache entry of the given labels and returns True on success. Args: value: An encodable (with corresponding PCoder) value *labels: List of labels for PCollection instance """ raise NotImplementedError def source(self, *labels): # type (*str) -> ptransform.PTransform """Returns a PTransform that reads the PCollection cache.""" raise NotImplementedError def sink(self, labels, is_capture=False): # type (*str, bool) -> ptransform.PTransform """Returns a PTransform that writes the PCollection cache. TODO(BEAM-10514): Make sure labels will not be converted into an arbitrarily long file path: e.g., windows has a 260 path limit. """ raise NotImplementedError def save_pcoder(self, pcoder, *labels): # type (coders.Coder, *str) -> None """Saves pcoder for given PCollection. Correct reading of PCollection from Cache requires PCoder to be known. This method saves desired PCoder for PCollection that will subsequently be used by sink(...), source(...), and, most importantly, read(...) method. The latter must be able to read a PCollection written by Beam using non-Beam IO. Args: pcoder: A PCoder to be used for reading and writing a PCollection. *labels: List of labels for PCollection instance. """ raise NotImplementedError def load_pcoder(self, *labels): # type (*str) -> coders.Coder """Returns previously saved PCoder for reading and writing PCollection.""" raise NotImplementedError def cleanup(self): # type () -> None """Cleans up all the PCollection caches.""" raise NotImplementedError def size(self, *labels: str) -> int: """Returns the size of the PCollection on disk in bytes.""" raise NotImplementedError class FileBasedCacheManager(CacheManager): """Maps PCollections to local temp files for materialization.""" _available_formats = { 'text': ( lambda path: textio.ReadFromText( path, coder=Base64Coder(), compression_type=filesystems. CompressionTypes.BZIP2), lambda path: textio.WriteToText( path, coder=Base64Coder(), compression_type=filesystems. CompressionTypes.BZIP2)), 'tfrecord': (tfrecordio.ReadFromTFRecord, tfrecordio.WriteToTFRecord) } def __init__(self, cache_dir=None, cache_format='text'): if cache_dir: self._cache_dir = cache_dir else: self._cache_dir = tempfile.mkdtemp( prefix='ib-', dir=os.environ.get('TEST_TMPDIR', None)) self._versions = collections.defaultdict(lambda: self._CacheVersion()) self.cache_format = cache_format if cache_format not in self._available_formats: raise ValueError("Unsupported cache format: '%s'." % cache_format) self._reader_class, self._writer_class = self._available_formats[ cache_format] self._default_pcoder = ( SafeFastPrimitivesCoder() if cache_format == 'text' else None) # List of saved pcoders keyed by PCollection path. It is OK to keep this # list in memory because once FileBasedCacheManager object is # destroyed/re-created it loses the access to previously written cache # objects anyways even if cache_dir already exists. In other words, # it is not possible to resume execution of Beam pipeline from the # saved cache if FileBasedCacheManager has been reset. # # However, if we are to implement better cache persistence, one needs # to take care of keeping consistency between the cached PCollection # and its PCoder type. self._saved_pcoders = {} def size(self, *labels): if self.exists(*labels): matched_path = self._match(*labels) # if any matched path has a gs:// prefix, it must be cached on GCS if 'gs://' in matched_path[0]: from apache_beam.io.gcp import gcsio return sum( sum(s for _, s in gcsio.GcsIO().list_files(path)) for path in matched_path) return sum(os.path.getsize(path) for path in matched_path) return 0 def exists(self, *labels): if labels and any(labels[1:]): return bool(self._match(*labels)) return False def _latest_version(self, *labels): timestamp = 0 for path in self._match(*labels): timestamp = max(timestamp, filesystems.FileSystems.last_updated(path)) result = self._versions["-".join(labels)].get_version(timestamp) return result def save_pcoder(self, pcoder, *labels): self._saved_pcoders[self._path(*labels)] = pcoder def load_pcoder(self, *labels): saved_pcoder = self._saved_pcoders.get(self._path(*labels), None) if saved_pcoder is None or isinstance(saved_pcoder, coders.FastPrimitivesCoder): return self._default_pcoder return saved_pcoder def read(self, *labels, **args): # Return an iterator to an empty list if it doesn't exist. if not self.exists(*labels): return iter([]), -1 # Otherwise, return a generator to the cached PCollection. coder = self.load_pcoder('reify', *labels[1:]) source = self.raw_source(*labels)._source range_tracker = source.get_range_tracker(None, None) reader = source.read(range_tracker) version = self._latest_version(*labels) return (coder.decode(b) for b in reader), version def write(self, values, *labels): """Imitates how a WriteCache transform works without running a pipeline. For testing and cache manager development, not for production usage because the write is not sharded and does not use Beam execution model. """ pcoder = coders.registry.get_coder(type(values[0])) # Save the pcoder for the actual labels. self.save_pcoder(pcoder, *labels) self.save_pcoder(pcoder, 'reify', *labels[-1:]) single_shard_labels = [*labels[:-1], '-00000-of-00001'] # Save the pcoder for the labels that imitates the sharded cache file name # suffix. self.save_pcoder(pcoder, *single_shard_labels) # Put a '-%05d-of-%05d' suffix to the cache file. sink = self.raw_sink(single_shard_labels)._sink path = self._path(*labels[:-1]) writer = sink.open_writer(path, labels[-1]) for v in values: writer.write(pcoder.encode(v)) writer.close() def clear(self, *labels): if self.exists(*labels): filesystems.FileSystems.delete(self._match(*labels)) return True return False def source(self, *labels): coder = self.load_pcoder('reify', *labels[1:]) return self.raw_source(*labels) | beam.Map( lambda b: test_stream.WindowedValueHolder(coder.decode(b))) def sink(self, labels, is_capture=False): coder = self.load_pcoder('reify', *labels[1:]) return beam.Map(lambda wvh: coder.encode(wvh.windowed_value) ) | self.raw_sink(labels, is_capture) def raw_sink(self, labels, is_capture=False): return self._writer_class(self._path(*labels)) def raw_source(self, *labels): return self._reader_class(self._glob_path(*labels)) def cleanup(self): if self._cache_dir.startswith('gs://'): from apache_beam.io.gcp import gcsfilesystem from apache_beam.options.pipeline_options import PipelineOptions fs = gcsfilesystem.GCSFileSystem(PipelineOptions()) fs.delete([self._cache_dir + '/full/']) elif filesystems.FileSystems.exists(self._cache_dir): filesystems.FileSystems.delete([self._cache_dir]) self._saved_pcoders = {} def _glob_path(self, *labels): return self._path(*labels) + '*-*-of-*' def _path(self, *labels): return filesystems.FileSystems.join(self._cache_dir, *labels) def _match(self, *labels): match = filesystems.FileSystems.match([self._glob_path(*labels)]) assert len(match) == 1 return [metadata.path for metadata in match[0].metadata_list] class _CacheVersion(object): """This class keeps track of the timestamp and the corresponding version.""" def __init__(self): self.current_version = -1 self.current_timestamp = 0 def get_version(self, timestamp): """Updates version if necessary and returns the version number. Args: timestamp: (int) unix timestamp when the cache is updated. This value is zero if the cache has been evicted or doesn't exist. """ # Do not update timestamp if the cache's been evicted. if timestamp != 0 and timestamp != self.current_timestamp: assert timestamp > self.current_timestamp self.current_version = self.current_version + 1 self.current_timestamp = timestamp return self.current_version class ReadCache(beam.PTransform): """A PTransform that reads the PCollections from the cache.""" def __init__(self, cache_manager, label): self._cache_manager = cache_manager self._label = label def expand(self, pbegin): # pylint: disable=expression-not-assigned return pbegin | 'Read' >> self._cache_manager.source('full', self._label) class WriteCache(beam.PTransform): """A PTransform that writes the PCollections to the cache.""" def __init__( self, cache_manager, label, sample=False, sample_size=0, is_capture=False): self._cache_manager = cache_manager self._label = label self._sample = sample self._sample_size = sample_size self._is_capture = is_capture def expand(self, pcoll): prefix = 'sample' if self._sample else 'full' # We save pcoder that is necessary for proper reading of # cached PCollection. _cache_manager.sink(...) call below # should be using this saved pcoder. self._cache_manager.save_pcoder( beam.coders.WindowedValueCoder( beam.coders.registry.get_coder(pcoll.element_type), pcoll.windowing.windowfn.get_window_coder()), 'reify', self._label) self._cache_manager.save_pcoder( coders.registry.get_coder(pcoll.element_type), prefix, self._label) if self._sample: pcoll |= 'Sample' >> ( combiners.Sample.FixedSizeGlobally(self._sample_size) | beam.FlatMap(lambda sample: sample)) # pylint: disable=expression-not-assigned return pcoll | 'Write' >> self._cache_manager.sink( (prefix, self._label), is_capture=self._is_capture) class SafeFastPrimitivesCoder(coders.Coder): """This class add an quote/unquote step to escape special characters.""" # pylint: disable=bad-option-value def encode(self, value): return quote( coders.coders.FastPrimitivesCoder().encode(value)).encode('utf-8') def decode(self, value): return coders.coders.FastPrimitivesCoder().decode(unquote_to_bytes(value)) class Base64Coder(coders.Coder): """Used to safely encode arbitrary bytes to textio.""" encode = staticmethod(base64.b64encode) decode = staticmethod(base64.b64decode) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/cache_manager_test.py0000644000175100017510000002467415206113032026051 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import time import unittest import apache_beam as beam from apache_beam import coders from apache_beam.runners.interactive import cache_manager as cache from apache_beam.runners.interactive import interactive_beam as ib class FileBasedCacheManagerTest(object): """Unit test for FileBasedCacheManager. Note that this set of tests focuses only the methods that interacts with the LOCAL file system. The idea is that once FileBasedCacheManager works well with the local file system, it should work with any file system with `apache_beam.io.filesystem` interface. Those tests that involve interactions with Beam pipeline (i.e. source(), sink(), ReadCache, and WriteCache) will be tested with InteractiveRunner as a part of integration tests instead. """ cache_format: str = None def setUp(self): self.cache_manager = cache.FileBasedCacheManager( cache_format=self.cache_format) def tearDown(self): self.cache_manager.cleanup() def mock_write_cache(self, values, prefix, cache_label): """Cache the PCollection where cache.WriteCache would write to.""" # Pause for 0.1 sec, because the Jenkins test runs so fast that the file # writes happen at the same timestamp. time.sleep(0.1) labels = [prefix, cache_label] self.cache_manager.write(values, *labels) def test_exists(self): """Test that CacheManager can correctly tell if the cache exists or not.""" prefix = 'full' cache_label = 'some-cache-label' cache_version_one = ['cache', 'version', 'one'] self.assertFalse(self.cache_manager.exists(prefix, cache_label)) self.mock_write_cache(cache_version_one, prefix, cache_label) self.assertTrue(self.cache_manager.exists(prefix, cache_label)) self.cache_manager.cleanup() self.assertFalse(self.cache_manager.exists(prefix, cache_label)) self.mock_write_cache(cache_version_one, prefix, cache_label) self.assertTrue(self.cache_manager.exists(prefix, cache_label)) def test_empty_label_not_exist(self): prefix = 'full' cache_label = 'some-cache-label' cache_version_one = ['cache', 'version', 'one'] self.assertFalse(self.cache_manager.exists(prefix, cache_label)) self.mock_write_cache(cache_version_one, prefix, cache_label) self.assertTrue(self.cache_manager.exists(prefix, cache_label)) # '' shouldn't be treated as a wildcard to match everything. self.assertFalse(self.cache_manager.exists(prefix, '')) def test_size(self): """Test getting the size of some cache label.""" # The Beam API for writing doesn't return the number of bytes that was # written to disk. So this test is only possible when the coder encodes the # bytes that will be written directly to disk, which only the WriteToText # transform does (with respect to the WriteToTFRecord transform). if self.cache_manager.cache_format != 'text': return prefix = 'full' cache_label = 'some-cache-label' # Test that if nothing is written the size is 0. self.assertEqual(self.cache_manager.size(prefix, cache_label), 0) value = 'a' self.mock_write_cache([value], prefix, cache_label) coder = self.cache_manager.load_pcoder(prefix, cache_label) encoded = coder.encode(value) # We encode in a format that escapes newlines. self.assertGreater( self.cache_manager.size(prefix, cache_label), len(encoded)) def test_clear(self): """Test that CacheManager can correctly tell if the cache exists or not.""" prefix = 'full' cache_label = 'some-cache-label' cache_version_one = ['cache', 'version', 'one'] self.assertFalse(self.cache_manager.exists(prefix, cache_label)) self.mock_write_cache(cache_version_one, prefix, cache_label) self.assertTrue(self.cache_manager.exists(prefix, cache_label)) self.assertTrue(self.cache_manager.clear(prefix, cache_label)) self.assertFalse(self.cache_manager.exists(prefix, cache_label)) def test_read_basic(self): """Test the condition where the cache is read once after written once.""" prefix = 'full' cache_label = 'some-cache-label' cache_version_one = ['cache', 'version', 'one'] self.mock_write_cache(cache_version_one, prefix, cache_label) reader, version = self.cache_manager.read(prefix, cache_label) pcoll_list = list(reader) self.assertListEqual(pcoll_list, cache_version_one) self.assertEqual(version, 0) self.assertTrue( self.cache_manager.is_latest_version(version, prefix, cache_label)) def test_read_version_update(self): """Tests if the version is properly updated after the files are updated.""" prefix = 'full' cache_label = 'some-cache-label' cache_version_one = ['cache', 'version', 'one'] cache_version_two = ['cache', 'version', 'two'] self.mock_write_cache(cache_version_one, prefix, cache_label) reader, version = self.cache_manager.read(prefix, cache_label) pcoll_list = list(reader) self.mock_write_cache(cache_version_two, prefix, cache_label) self.assertFalse( self.cache_manager.is_latest_version(version, prefix, cache_label)) reader, version = self.cache_manager.read(prefix, cache_label) pcoll_list = list(reader) self.assertListEqual(pcoll_list, cache_version_two) self.assertEqual(version, 1) self.assertTrue( self.cache_manager.is_latest_version(version, prefix, cache_label)) def test_read_before_write(self): """Test the behavior when read() is called before WriteCache completes.""" prefix = 'full' cache_label = 'some-cache-label' self.assertFalse(self.cache_manager.exists(prefix, cache_label)) reader, version = self.cache_manager.read(prefix, cache_label) pcoll_list = list(reader) self.assertListEqual(pcoll_list, []) self.assertEqual(version, -1) self.assertTrue( self.cache_manager.is_latest_version(version, prefix, cache_label)) def test_read_over_cleanup(self): """Test the behavior of read() over cache cleanup.""" prefix = 'full' cache_label = 'some-cache-label' cache_version_one = ['cache', 'version', 'one'] cache_version_two = ['cache', 'version', 'two'] # The initial write and read. self.mock_write_cache(cache_version_one, prefix, cache_label) reader, version = self.cache_manager.read(prefix, cache_label) pcoll_list = list(reader) # Cache cleanup. self.cache_manager.cleanup() # Check that even if cache is evicted, the latest version stays the same. self.assertTrue( self.cache_manager.is_latest_version(version, prefix, cache_label)) reader, version = self.cache_manager.read(prefix, cache_label) pcoll_list = list(reader) self.assertListEqual(pcoll_list, []) self.assertEqual(version, -1) self.assertFalse( self.cache_manager.is_latest_version(version, prefix, cache_label)) # PCollection brought back to cache. self.mock_write_cache(cache_version_two, prefix, cache_label) self.assertFalse( self.cache_manager.is_latest_version(version, prefix, cache_label)) reader, version = self.cache_manager.read(prefix, cache_label) pcoll_list = list(reader) self.assertListEqual(pcoll_list, cache_version_two) # Check that version continues from the previous value instead of starting # from 0 again. self.assertEqual(version, 1) self.assertTrue( self.cache_manager.is_latest_version(version, prefix, cache_label)) def test_load_saved_pcoder(self): pipeline = beam.Pipeline() pcoll = pipeline | beam.Create([1, 2, 3]) _ = pcoll | cache.WriteCache(self.cache_manager, 'a key') self.assertIs( type(self.cache_manager.load_pcoder('full', 'a key')), type(coders.registry.get_coder(int))) def test_cache_manager_uses_gcs_ib_cache_root(self): """ Checks that FileBasedCacheManager._cache_dir is set to the cache_root set under Interactive Beam for a GCS directory. """ # Set Interactive Beam specified cache dir to cloud storage ib.options.cache_root = 'gs://' cache_manager_with_ib_option = cache.FileBasedCacheManager( cache_dir=ib.options.cache_root) self.assertEqual( ib.options.cache_root, cache_manager_with_ib_option._cache_dir) # Reset Interactive Beam setting ib.options.cache_root = None def test_cache_manager_uses_local_ib_cache_root(self): """ Checks that FileBasedCacheManager._cache_dir is set to the cache_root set under Interactive Beam for a local directory and that the cached values are the same as the values of a cache using default settings. """ prefix = 'full' cache_label = 'some-cache-label' cached_values = [1, 2, 3] self.mock_write_cache(cached_values, prefix, cache_label) reader_one, _ = self.cache_manager.read(prefix, cache_label) pcoll_list_one = list(reader_one) # Set Interactive Beam specified cache dir to local directory ib.options.cache_root = '/tmp/it-test/' cache_manager_with_ib_option = cache.FileBasedCacheManager( cache_dir=ib.options.cache_root) self.assertEqual( ib.options.cache_root, cache_manager_with_ib_option._cache_dir) cache_manager_with_ib_option.write(cached_values, *[prefix, cache_label]) reader_two, _ = self.cache_manager.read(prefix, cache_label) pcoll_list_two = list(reader_two) # Writing to a different directory should not impact the cached values self.assertEqual(pcoll_list_one, pcoll_list_two) # Reset Interactive Beam setting ib.options.cache_root = None class TextFileBasedCacheManagerTest( FileBasedCacheManagerTest, unittest.TestCase, ): cache_format = 'text' class TFRecordBasedCacheManagerTest( FileBasedCacheManagerTest, unittest.TestCase, ): cache_format = 'tfrecord' if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.6304343 apache_beam-2.74.0/apache_beam/runners/interactive/caching/0000755000175100017510000000000015206113160023264 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/caching/__init__.py0000644000175100017510000000142015206113032025370 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/caching/cacheable.py0000644000175100017510000000560315206113032025527 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for dataclasses to hold metadata of cacheable PCollections in the user code scope. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file from dataclasses import dataclass import apache_beam as beam @dataclass class Cacheable: var: str version: str producer_version: str pcoll: beam.pvalue.PCollection def __hash__(self): return hash((self.var, self.version, self.producer_version, self.pcoll)) @staticmethod def from_pcoll( pcoll_name: str, pcoll: beam.pvalue.PCollection) -> 'Cacheable': return Cacheable(pcoll_name, str(id(pcoll)), str(id(pcoll.producer)), pcoll) def to_key(self): return CacheKey( self.var, self.version, self.producer_version, str(id(self.pcoll.pipeline))) @dataclass class CacheKey: """The identifier of a cacheable PCollection in cache. It contains 4 stringified components: var: The obfuscated variable name of the PCollection. version: The id() of the PCollection. producer_version: The id() of the producer of the PCollection. pipeline_id: The id() of the pipeline the PCollection belongs to. """ var: str version: str producer_version: str pipeline_id: str def __post_init__(self): from apache_beam.runners.interactive.utils import obfuscate # Normalize arbitrary variable name to a fixed length hex str. self.var = obfuscate(self.var)[:10] def __hash__(self): return hash( (self.var, self.version, self.producer_version, self.pipeline_id)) @staticmethod def from_str(r: str) -> 'CacheKey': r_split = r.split('-') ck = CacheKey(*r_split) # Avoid double obfuscation. ck.var = r_split[0] return ck @staticmethod def from_pcoll(pcoll_name: str, pcoll: beam.pvalue.PCollection) -> 'CacheKey': return CacheKey( pcoll_name, str(id(pcoll)), str(id(pcoll.producer)), str(id(pcoll.pipeline))) def to_str(self): return '-'.join( [self.var, self.version, self.producer_version, self.pipeline_id]) def __repr__(self): return self.to_str() def __str__(self): return self.to_str() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/caching/expression_cache.py0000644000175100017510000001015515206113032027160 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from typing import * import apache_beam as beam from apache_beam.dataframe import convert from apache_beam.dataframe import expressions class ExpressionCache(object): """Utility class for caching deferred DataFrames expressions. This is cache is currently a light-weight wrapper around the TO_PCOLLECTION_CACHE in the beam.dataframes.convert module and the computed_pcollections in the interactive module. Example:: df : beam.dataframe.DeferredDataFrame = ... ... cache = ExpressionCache() cache.replace_with_cached(df._expr) This will automatically link the instance to the existing caches. After it is created, the cache can then be used to modify an existing deferred dataframe expression tree to replace nodes with computed PCollections. This object can be created and destroyed whenever. This class holds no state and the only side-effect is modifying the given expression. """ def __init__(self, pcollection_cache=None, computed_cache=None): from apache_beam.runners.interactive import interactive_environment as ie self._pcollection_cache = ( convert.TO_PCOLLECTION_CACHE if pcollection_cache is None else pcollection_cache) self._computed_cache = ( ie.current_env().computed_pcollections if computed_cache is None else computed_cache) def replace_with_cached( self, expr: expressions.Expression) -> Dict[str, expressions.Expression]: """Replaces any previously computed expressions with PlaceholderExpressions. This is used to correctly read any expressions that were cached in previous runs. This enables the InteractiveRunner to prune off old calculations from the expression tree. """ replaced_inputs: Dict[str, expressions.Expression] = {} self._replace_with_cached_recur(expr, replaced_inputs) return replaced_inputs def _replace_with_cached_recur( self, expr: expressions.Expression, replaced_inputs: Dict[str, expressions.Expression]) -> None: """Recursive call for `replace_with_cached`. Recurses through the expression tree and replaces any cached inputs with `PlaceholderExpression`s. """ final_inputs = [] for input in expr.args(): pc = self._get_cached(input) # Only read from cache when there is the PCollection has been fully # computed. This is so that no partial results are used. if self._is_computed(pc): # Reuse previously seen cached expressions. This is so that the same # value isn't cached multiple times. if input._id in replaced_inputs: cached = replaced_inputs[input._id] else: cached = expressions.PlaceholderExpression( input.proxy(), self._pcollection_cache[input._id]) replaced_inputs[input._id] = cached final_inputs.append(cached) else: final_inputs.append(input) self._replace_with_cached_recur(input, replaced_inputs) expr._args = tuple(final_inputs) def _get_cached(self, expr: expressions.Expression) -> Optional[beam.PCollection]: """Returns the PCollection associated with the expression.""" return self._pcollection_cache.get(expr._id, None) def _is_computed(self, pc: beam.PCollection) -> bool: """Returns True if the PCollection has been run and computed.""" return pc is not None and pc in self._computed_cache ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/caching/expression_cache_test.py0000644000175100017510000001074115206113032030220 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest import apache_beam as beam from apache_beam.dataframe import expressions from apache_beam.runners.interactive.caching.expression_cache import ExpressionCache class ExpressionCacheTest(unittest.TestCase): def setUp(self): self._pcollection_cache = {} self._computed_cache = set() self._pipeline = beam.Pipeline() self.cache = ExpressionCache(self._pcollection_cache, self._computed_cache) def create_trace(self, expr): trace = [expr] for input in expr.args(): trace += self.create_trace(input) return trace def mock_cache(self, expr): pcoll = beam.PCollection(self._pipeline) self._pcollection_cache[expr._id] = pcoll self._computed_cache.add(pcoll) def assertTraceTypes(self, expr, expected): actual_types = [type(e).__name__ for e in self.create_trace(expr)] expected_types = [e.__name__ for e in expected] self.assertListEqual(actual_types, expected_types) def test_only_replaces_cached(self): in_expr = expressions.ConstantExpression(0) comp_expr = expressions.ComputedExpression('test', lambda x: x, [in_expr]) # Expect that no replacement of expressions is performed. expected_trace = [ expressions.ComputedExpression, expressions.ConstantExpression ] self.assertTraceTypes(comp_expr, expected_trace) self.cache.replace_with_cached(comp_expr) self.assertTraceTypes(comp_expr, expected_trace) # Now "cache" the expression and assert that the cached expression was # replaced with a placeholder. self.mock_cache(in_expr) replaced = self.cache.replace_with_cached(comp_expr) expected_trace = [ expressions.ComputedExpression, expressions.PlaceholderExpression ] self.assertTraceTypes(comp_expr, expected_trace) self.assertIn(in_expr._id, replaced) def test_only_replaces_inputs(self): arg_0_expr = expressions.ConstantExpression(0) ident_val = expressions.ComputedExpression( 'ident', lambda x: x, [arg_0_expr]) arg_1_expr = expressions.ConstantExpression(1) comp_expr = expressions.ComputedExpression( 'add', lambda x, y: x + y, [ident_val, arg_1_expr]) self.mock_cache(ident_val) replaced = self.cache.replace_with_cached(comp_expr) # Assert that ident_val was replaced and that its arguments were removed # from the expression tree. expected_trace = [ expressions.ComputedExpression, expressions.PlaceholderExpression, expressions.ConstantExpression ] self.assertTraceTypes(comp_expr, expected_trace) self.assertIn(ident_val._id, replaced) self.assertNotIn(arg_0_expr, self.create_trace(comp_expr)) def test_only_caches_same_input(self): arg_0_expr = expressions.ConstantExpression(0) ident_val = expressions.ComputedExpression( 'ident', lambda x: x, [arg_0_expr]) comp_expr = expressions.ComputedExpression( 'add', lambda x, y: x + y, [ident_val, arg_0_expr]) self.mock_cache(arg_0_expr) replaced = self.cache.replace_with_cached(comp_expr) # Assert that arg_0_expr, being an input to two computations, was replaced # with the same placeholder expression. expected_trace = [ expressions.ComputedExpression, expressions.ComputedExpression, expressions.PlaceholderExpression, expressions.PlaceholderExpression ] actual_trace = self.create_trace(comp_expr) unique_placeholders = set( t for t in actual_trace if isinstance(t, expressions.PlaceholderExpression)) self.assertTraceTypes(comp_expr, expected_trace) self.assertTrue( all(e == replaced[arg_0_expr._id] for e in unique_placeholders)) self.assertIn(arg_0_expr._id, replaced) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/caching/read_cache.py0000644000175100017510000001423415206113032025676 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module to read cache of computed PCollections. For internal use only; no backward-compatibility guarantees. """ # pytype: skip-file import apache_beam as beam from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.interactive import cache_manager as cache from apache_beam.runners.interactive.caching.cacheable import Cacheable from apache_beam.runners.interactive.caching.reify import unreify_from_cache from apache_beam.runners.pipeline_context import PipelineContext from apache_beam.transforms.ptransform import PTransform class ReadCache: """Class that facilitates reading cache of computed PCollections. """ def __init__( self, pipeline: beam_runner_api_pb2.Pipeline, context: PipelineContext, cache_manager: cache.CacheManager, cacheable: Cacheable): self._pipeline = pipeline self._context = context self._cache_manager = cache_manager self._cacheable = cacheable self._key = repr(cacheable.to_key()) def read_cache(self) -> tuple[str, str]: """Reads cache of the cacheable PCollection and wires the cache into the pipeline proto. Returns the pipeline-scoped ids of the cacheable PCollection and the cache reading output PCollection that replaces it. First, it creates a temporary pipeline instance on top of the existing component_id_map from the self._pipeline's context so that both pipelines share the context and have no conflict component ids. Second, it instantiates a _ReadCacheTransform to build the temporary pipeline with a subgraph under top level transforms that reads the cache of a cacheable PCollection. Third, it copies components of the subgraph from the temporary pipeline to self._pipeline, skipping components that are not in the temporary pipeline but presents in the component_id_map of self._pipeline. Since to_runner_api generates components for all entries in the component_id_map, those component ids from the context shared by self._pipeline need to be ignored. Last, it replaces inputs of all transforms that consume the cacheable PCollection with the output PCollection of the _ReadCacheTransform so that the whole pipeline computes with data from the cache. The pipeline fragment of reading the cacheable PCollection is now disconnected from the rest of the pipeline and can be pruned later. """ template, read_output = self._build_runner_api_template() output_id = self._context.pcollections.get_id(read_output) source_id = self._context.pcollections.get_id(self._cacheable.pcoll) # Copy cache reading subgraph from the template to the pipeline proto. for pcoll_id in template.components.pcollections: if pcoll_id in self._pipeline.components.pcollections: continue self._pipeline.components.pcollections[pcoll_id].CopyFrom( template.components.pcollections[pcoll_id]) for coder_id in template.components.coders: if coder_id in self._pipeline.components.coders: continue self._pipeline.components.coders[coder_id].CopyFrom( template.components.coders[coder_id]) for windowing_strategy_id in template.components.windowing_strategies: if (windowing_strategy_id in self._pipeline.components.windowing_strategies): continue self._pipeline.components.windowing_strategies[ windowing_strategy_id].CopyFrom( template.components.windowing_strategies[windowing_strategy_id]) template_root_transform_id = template.root_transform_ids[0] root_transform_id = self._pipeline.root_transform_ids[0] for transform_id in template.components.transforms: if (transform_id == template_root_transform_id or transform_id in self._pipeline.components.transforms): continue self._pipeline.components.transforms[transform_id].CopyFrom( template.components.transforms[transform_id]) self._pipeline.components.transforms[ root_transform_id].subtransforms.extend( template.components.transforms[template_root_transform_id]. subtransforms) # Replace all the input pcoll of source_id with output pcoll of output_id # from cache reading. for transform in self._pipeline.components.transforms.values(): inputs = transform.inputs if source_id in inputs.values(): keys_need_replacement = set() for key in inputs: if inputs[key] == source_id: keys_need_replacement.add(key) for key in keys_need_replacement: inputs[key] = output_id return source_id, output_id def _build_runner_api_template( self) -> tuple[beam_runner_api_pb2.Pipeline, beam.pvalue.PCollection]: transform = _ReadCacheTransform(self._cache_manager, self._key) tmp_pipeline = beam.Pipeline() tmp_pipeline.component_id_map = self._context.component_id_map read_output = tmp_pipeline | 'source_cache_' >> transform return tmp_pipeline.to_runner_api(), read_output class _ReadCacheTransform(PTransform): """A composite transform encapsulates reading cache of PCollections. """ def __init__(self, cache_manager: cache.CacheManager, key: str): self._cache_manager = cache_manager self._key = key def expand(self, pcoll: beam.pvalue.PCollection) -> beam.pvalue.PCollection: return unreify_from_cache( pipeline=pcoll.pipeline, cache_key=self._key, cache_manager=self._cache_manager) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/caching/read_cache_test.py0000644000175100017510000000665115206113032026741 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for read_cache.""" # pytype: skip-file import unittest from unittest.mock import patch import apache_beam as beam from apache_beam.runners.interactive import augmented_pipeline as ap from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.caching import read_cache from apache_beam.runners.interactive.testing.pipeline_assertion import assert_pipeline_proto_equal from apache_beam.runners.interactive.testing.test_cache_manager import InMemoryCache class ReadCacheTest(unittest.TestCase): def setUp(self): ie.new_env() @patch( 'apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.get_cache_manager') def test_read_cache(self, mocked_get_cache_manager): p = beam.Pipeline() pcoll = p | beam.Create([1, 2, 3]) consumer_transform = beam.Map(lambda x: x * x) _ = pcoll | consumer_transform ib.watch(locals()) # Create the cache in memory. cache_manager = InMemoryCache() mocked_get_cache_manager.return_value = cache_manager aug_p = ap.AugmentedPipeline(p) key = repr(aug_p._cacheables[pcoll].to_key()) cache_manager.write('test', 'full', key) # Capture the applied transform of the consumer_transform. pcoll_id = aug_p._context.pcollections.get_id(pcoll) consumer_transform_id = None pipeline_proto = p.to_runner_api() for (transform_id, transform) in pipeline_proto.components.transforms.items(): if pcoll_id in transform.inputs.values(): consumer_transform_id = transform_id break self.assertIsNotNone(consumer_transform_id) # Read cache on the pipeline proto. _, cache_id = read_cache.ReadCache( pipeline_proto, aug_p._context, aug_p._cache_manager, aug_p._cacheables[pcoll]).read_cache() actual_pipeline = pipeline_proto # Read cache directly on the pipeline instance. transform = read_cache._ReadCacheTransform(aug_p._cache_manager, key) p | 'source_cache_' + key >> transform expected_pipeline = p.to_runner_api() # This rougly checks the equivalence between two protos, not detailed # wiring in sub transforms under top level transforms. assert_pipeline_proto_equal(self, expected_pipeline, actual_pipeline) # Check if the actual_pipeline uses cache as input of the # consumer_transform instead of the original pcoll from source. inputs = actual_pipeline.components.transforms[consumer_transform_id].inputs self.assertIn(cache_id, inputs.values()) self.assertNotIn(pcoll_id, inputs.values()) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/caching/reify.py0000644000175100017510000001001115206113032024743 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for transforms that reifies and unreifies PCollection values with window info. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file from typing import Optional import apache_beam as beam from apache_beam.runners.interactive import cache_manager as cache from apache_beam.testing import test_stream READ_CACHE = 'ReadCache_' WRITE_CACHE = 'WriteCache_' class Reify(beam.DoFn): """Reifies elements with window info into windowed values. Internally used to capture window info with each element into cache for replayability. """ def process(self, e, wv=beam.DoFn.WindowedValueParam): yield test_stream.WindowedValueHolder(wv) class Unreify(beam.DoFn): """Unreifies elements from windowed values. Cached values are elements with window info. This unpacks the elements. """ def process(self, e): # Row coder was used when encoding windowed values. if isinstance(e, beam.Row) and hasattr(e, 'windowed_value'): yield e.windowed_value def reify_to_cache( pcoll: beam.pvalue.PCollection, cache_key: str, cache_manager: cache.CacheManager, reify_label: Optional[str] = None, write_cache_label: Optional[str] = None, is_capture: bool = False) -> beam.pvalue.PValue: """Reifies elements into windowed values and write to cache. Args: pcoll: The PCollection to be cached. cache_key: The key of the cache. cache_manager: The cache manager to manage the cache. reify_label: (optional) A transform label for the Reify transform. write_cache_label: (optional) A transform label for the cache-writing transform. is_capture: Whether the cache is capturing a record of recordable sources. """ if not reify_label: reify_label = '{}{}{}'.format('ReifyBefore_', WRITE_CACHE, cache_key) if not write_cache_label: write_cache_label = '{}{}'.format(WRITE_CACHE, cache_key) return ( pcoll | reify_label >> beam.ParDo(Reify()) | write_cache_label >> cache.WriteCache( cache_manager, cache_key, is_capture=is_capture)) def unreify_from_cache( pipeline: beam.Pipeline, cache_key: str, cache_manager: cache.CacheManager, element_type: Optional[type] = None, source_label: Optional[str] = None, unreify_label: Optional[str] = None) -> beam.pvalue.PCollection: """Reads from cache and unreifies elements from windowed values. pipeline: The pipeline that's reading from the cache. cache_key: The key of the cache. cache_manager: The cache manager to manage the cache. element_type: (optional) The element type of the PCollection's elements. source_label: (optional) A transform label for the cache-reading transform. unreify_label: (optional) A transform label for the Unreify transform. """ if not source_label: source_label = '{}{}'.format(READ_CACHE, cache_key) if not unreify_label: unreify_label = '{}{}{}'.format('UnreifyAfter_', READ_CACHE, cache_key) read_cache = pipeline | source_label >> cache.ReadCache( cache_manager, cache_key) if element_type: # If the PCollection is schema-aware, explicitly sets the output types. return read_cache | unreify_label >> beam.ParDo( Unreify()).with_output_types(element_type) return read_cache | unreify_label >> beam.ParDo(Unreify()) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/caching/streaming_cache.py0000644000175100017510000004710515206113032026757 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import os import shutil import tempfile import time import traceback from collections import OrderedDict # We don't have an explicit pathlib dependency because this code only works with # the interactive target installed which has an indirect dependency on pathlib # through ipython>=5.9.0. from pathlib import Path from google.protobuf.message import DecodeError import apache_beam as beam from apache_beam import coders from apache_beam.portability.api import beam_interactive_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.interactive.cache_manager import CacheManager from apache_beam.runners.interactive.cache_manager import SafeFastPrimitivesCoder from apache_beam.runners.interactive.caching.cacheable import CacheKey from apache_beam.testing.test_stream import OutputFormat from apache_beam.testing.test_stream import ReverseTestStream from apache_beam.utils import timestamp _LOGGER = logging.getLogger(__name__) class StreamingCacheSink(beam.PTransform): """A PTransform that writes TestStreamFile(Header|Records)s to file. This transform takes in an arbitrary element stream and writes the list of TestStream events (as TestStreamFileRecords) to file. When replayed, this will produce the best-effort replay of the original job (e.g. some elements may be produced slightly out of order from the original stream). Note that this PTransform is assumed to be only run on a single machine where the following assumptions are correct: elements come in ordered, no two transforms are writing to the same file. This PTransform is assumed to only run correctly with the DirectRunner. TODO(https://github.com/apache/beam/issues/20002): Generalize this to more source/sink types aside from file based. Also, generalize to cases where there might be multiple workers writing to the same sink. """ def __init__( self, cache_dir, filename, sample_resolution_sec, coder=SafeFastPrimitivesCoder()): self._cache_dir = cache_dir self._filename = filename self._sample_resolution_sec = sample_resolution_sec self._coder = coder self._path = os.path.join(self._cache_dir, self._filename) @property def path(self): """Returns the path the sink leads to.""" return self._path @property def size_in_bytes(self): """Returns the space usage in bytes of the sink.""" try: return os.stat(self._path).st_size except OSError: _LOGGER.debug( 'Failed to calculate cache size for file %s, the file might have not ' 'been created yet. Return 0. %s', self._path, traceback.format_exc()) return 0 def expand(self, pcoll): class StreamingWriteToText(beam.DoFn): """DoFn that performs the writing. Note that the other file writing methods cannot be used in streaming contexts. """ def __init__(self, full_path, coder=SafeFastPrimitivesCoder()): self._full_path = full_path self._coder = coder # Try and make the given path. Path(os.path.dirname(full_path)).mkdir(parents=True, exist_ok=True) def start_bundle(self): # Open the file for 'append-mode' and writing 'bytes'. self._fh = open(self._full_path, 'ab') def finish_bundle(self): self._fh.close() def process(self, e): """Appends the given element to the file. """ self._fh.write(self._coder.encode(e) + b'\n') return ( pcoll | ReverseTestStream( output_tag=self._filename, sample_resolution_sec=self._sample_resolution_sec, output_format=OutputFormat.SERIALIZED_TEST_STREAM_FILE_RECORDS, coder=self._coder) | beam.ParDo( StreamingWriteToText(full_path=self._path, coder=self._coder))) class StreamingCacheSource: """A class that reads and parses TestStreamFile(Header|Reader)s. This source operates in the following way: 1. Wait for up to `timeout_secs` for the file to be available. 2. Read, parse, and emit the entire contents of the file 3. Wait for more events to come or until `is_cache_complete` returns True 4. If there are more events, then go to 2 5. Otherwise, stop emitting. This class is used to read from file and send its to the TestStream via the StreamingCacheManager.Reader. """ def __init__(self, cache_dir, labels, is_cache_complete=None, coder=None): if not coder: coder = SafeFastPrimitivesCoder() if not is_cache_complete: is_cache_complete = lambda _: True self._cache_dir = cache_dir self._coder = coder self._labels = labels self._path = os.path.join(self._cache_dir, *self._labels) self._is_cache_complete = is_cache_complete self._pipeline_id = CacheKey.from_str(labels[-1]).pipeline_id def _wait_until_file_exists(self, timeout_secs=30): """Blocks until the file exists for a maximum of timeout_secs. """ # Wait for up to `timeout_secs` for the file to be available. start = time.time() while not os.path.exists(self._path): time.sleep(1) if time.time() - start > timeout_secs: pcollection_var = CacheKey.from_str(self._labels[-1]).var raise RuntimeError( 'Timed out waiting for cache file for PCollection `{}` to be ' 'available with path {}.'.format(pcollection_var, self._path)) return open(self._path, mode='rb') def _emit_from_file(self, fh, tail): """Emits the TestStreamFile(Header|Record)s from file. This returns a generator to be able to read all lines from the given file. If `tail` is True, then it will wait until the cache is complete to exit. Otherwise, it will read the file only once. """ # Always read at least once to read the whole file. while True: pos = fh.tell() line = fh.readline() # Check if we are at EOF or if we have an incomplete line. if not line or (line and line[-1] != b'\n'[0]): # Read at least the first line to get the header. if not tail and pos != 0: break # Complete reading only when the cache is complete. if self._is_cache_complete(self._pipeline_id): break # Otherwise wait for new data in the file to be written. time.sleep(0.5) fh.seek(pos) else: # The first line at pos = 0 is always the header. Read the line without # the new line. to_decode = line[:-1] if pos == 0: proto_cls = beam_interactive_api_pb2.TestStreamFileHeader else: proto_cls = beam_interactive_api_pb2.TestStreamFileRecord msg = self._try_parse_as(proto_cls, to_decode) if msg: yield msg else: break def _try_parse_as(self, proto_cls, to_decode): try: msg = proto_cls() msg.ParseFromString(self._coder.decode(to_decode)) except DecodeError: _LOGGER.error( 'Could not parse as %s. This can indicate that the cache is ' 'corruputed. Please restart the kernel. ' '\nfile: %s \nmessage: %s', proto_cls, self._path, to_decode) msg = None return msg def read(self, tail): """Reads all TestStreamFile(Header|TestStreamFileRecord)s from file. This returns a generator to be able to read all lines from the given file. If `tail` is True, then it will wait until the cache is complete to exit. Otherwise, it will read the file only once. """ with self._wait_until_file_exists() as f: for e in self._emit_from_file(f, tail): yield e # TODO(victorhc): Add support for cache_dir locations that are on GCS class StreamingCache(CacheManager): """Abstraction that holds the logic for reading and writing to cache. """ def __init__( self, cache_dir, is_cache_complete=None, sample_resolution_sec=0.1, saved_pcoders=None): self._sample_resolution_sec = sample_resolution_sec self._is_cache_complete = is_cache_complete if cache_dir: self._cache_dir = cache_dir else: self._cache_dir = tempfile.mkdtemp( prefix='ib-', dir=os.environ.get('TEST_TMPDIR', None)) # List of saved pcoders keyed by PCollection path. It is OK to keep this # list in memory because once FileBasedCacheManager object is # destroyed/re-created it loses the access to previously written cache # objects anyways even if cache_dir already exists. In other words, # it is not possible to resume execution of Beam pipeline from the # saved cache if FileBasedCacheManager has been reset. # # However, if we are to implement better cache persistence, one needs # to take care of keeping consistency between the cached PCollection # and its PCoder type. self._saved_pcoders = saved_pcoders or {} self._default_pcoder = SafeFastPrimitivesCoder() # The sinks to capture data from capturable sources. # Dict([str, StreamingCacheSink]) self._capture_sinks = {} self._capture_keys = set() def size(self, *labels): if self.exists(*labels): return os.path.getsize(os.path.join(self._cache_dir, *labels)) return 0 @property def capture_size(self): return sum([sink.size_in_bytes for _, sink in self._capture_sinks.items()]) @property def capture_paths(self): return list(self._capture_sinks.keys()) @property def capture_keys(self): return self._capture_keys def exists(self, *labels): if labels and any(labels): path = os.path.join(self._cache_dir, *labels) return os.path.exists(path) return False # TODO(srohde): Modify this to return the correct version. def read(self, *labels, **args): """Returns a generator to read all records from file.""" tail = args.pop('tail', False) # Only immediately return when the file doesn't exist when the user wants a # snapshot of the cache (when tail is false). if not self.exists(*labels) and not tail: return iter([]), -1 reader = StreamingCacheSource( self._cache_dir, labels, self._is_cache_complete, self.load_pcoder(*labels)).read(tail=tail) # Return an empty iterator if there is nothing in the file yet. This can # only happen when tail is False. try: header = next(reader) except StopIteration: return iter([]), -1 return StreamingCache.Reader([header], [reader]).read(), 1 def read_multiple(self, labels, tail=True): """Returns a generator to read all records from file. Does tail until the cache is complete. This is because it is used in the TestStreamServiceController to read from file which is only used during pipeline runtime which needs to block. """ readers = [ StreamingCacheSource( self._cache_dir, l, self._is_cache_complete, self.load_pcoder(*l)).read(tail=tail) for l in labels ] headers = [next(r) for r in readers] return StreamingCache.Reader(headers, readers).read() def write(self, values, *labels): """Writes the given values to cache. """ directory = os.path.join(self._cache_dir, *labels[:-1]) filepath = os.path.join(directory, labels[-1]) if not os.path.exists(directory): os.makedirs(directory) with open(filepath, 'ab') as f: for v in values: if isinstance(v, (beam_interactive_api_pb2.TestStreamFileHeader, beam_interactive_api_pb2.TestStreamFileRecord)): val = v.SerializeToString() else: raise TypeError( 'Values given to streaming cache should be either ' 'TestStreamFileHeader or TestStreamFileRecord.') f.write(self.load_pcoder(*labels).encode(val) + b'\n') def clear(self, *labels): directory = os.path.join(self._cache_dir, *labels[:-1]) filepath = os.path.join(directory, labels[-1]) self._capture_keys.discard(labels[-1]) if os.path.exists(filepath): os.remove(filepath) return True return False def source(self, *labels): """Returns the StreamingCacheManager source. This is beam.Impulse() because unbounded sources will be marked with this and then the PipelineInstrument will replace these with a TestStream. """ return beam.Impulse() def sink(self, labels, is_capture=False): """Returns a StreamingCacheSink to write elements to file. Note that this is assumed to only work in the DirectRunner as the underlying StreamingCacheSink assumes a single machine to have correct element ordering. """ filename = labels[-1] cache_dir = os.path.join(self._cache_dir, *labels[:-1]) sink = StreamingCacheSink( cache_dir, filename, self._sample_resolution_sec, self.load_pcoder(*labels)) if is_capture: self._capture_sinks[sink.path] = sink self._capture_keys.add(filename) return sink def save_pcoder(self, pcoder, *labels): self._saved_pcoders[os.path.join(self._cache_dir, *labels)] = pcoder def load_pcoder(self, *labels): saved_pcoder = self._saved_pcoders.get( os.path.join(self._cache_dir, *labels), None) if saved_pcoder is None or isinstance(saved_pcoder, coders.FastPrimitivesCoder): return self._default_pcoder return saved_pcoder def cleanup(self): if os.path.exists(self._cache_dir): def on_fail_to_cleanup(function, path, excinfo): _LOGGER.warning( 'Failed to clean up temporary files: %s. You may' 'manually delete them if necessary. Error was: %s', path, excinfo) shutil.rmtree(self._cache_dir, onerror=on_fail_to_cleanup) self._saved_pcoders = {} self._capture_sinks = {} self._capture_keys = set() class Reader(object): """Abstraction that reads from PCollection readers. This class is an Abstraction layer over multiple PCollection readers to be used for supplying a TestStream service with events. This class is also responsible for holding the state of the clock, injecting clock advancement events, and watermark advancement events. """ def __init__(self, headers, readers): # This timestamp is used as the monotonic clock to order events in the # replay. self._monotonic_clock = timestamp.Timestamp.of(0) # The PCollection cache readers. self._readers = {} # The file headers that are metadata for that particular PCollection. # The header allows for metadata about an entire stream, so that the data # isn't copied per record. self._headers = {header.tag: header for header in headers} self._readers = OrderedDict( ((h.tag, r) for (h, r) in zip(headers, readers))) # The most recently read timestamp per tag. self._stream_times = { tag: timestamp.Timestamp(seconds=0) for tag in self._headers } def _test_stream_events_before_target(self, target_timestamp): """Reads the next iteration of elements from each stream. Retrieves an element from each stream iff the most recently read timestamp from that stream is less than the target_timestamp. Since the amount of events may not fit into memory, this StreamingCache reads at most one element from each stream at a time. """ records = [] for tag, r in self._readers.items(): # The target_timestamp is the maximum timestamp that was read from the # stream. Some readers may have elements that are less than this. Thus, # we skip all readers that already have elements that are at this # timestamp so that we don't read everything into memory. if self._stream_times[tag] >= target_timestamp: continue try: record = next(r).recorded_event if record.HasField('processing_time_event'): self._stream_times[tag] += timestamp.Duration( micros=record.processing_time_event.advance_duration) records.append((tag, record, self._stream_times[tag])) except StopIteration: pass return records def _merge_sort(self, previous_events, new_events): return sorted( previous_events + new_events, key=lambda x: x[2], reverse=True) def _min_timestamp_of(self, events): return events[-1][2] if events else timestamp.MAX_TIMESTAMP def _event_stream_caught_up_to_target(self, events, target_timestamp): empty_events = not events stream_is_past_target = self._min_timestamp_of(events) > target_timestamp return empty_events or stream_is_past_target def read(self): """Reads records from PCollection readers. """ # The largest timestamp read from the different streams. target_timestamp = timestamp.MAX_TIMESTAMP # The events from last iteration that are past the target timestamp. unsent_events = [] # Emit events until all events have been read. while True: # Read the next set of events. The read events will most likely be # out of order if there are multiple readers. Here we sort them into # a more manageable state. new_events = self._test_stream_events_before_target(target_timestamp) events_to_send = self._merge_sort(unsent_events, new_events) if not events_to_send: break # Get the next largest timestamp in the stream. This is used as the # timestamp for readers to "catch-up" to. This will only read from # readers with a timestamp less than this. target_timestamp = self._min_timestamp_of(events_to_send) # Loop through the elements with the correct timestamp. while not self._event_stream_caught_up_to_target(events_to_send, target_timestamp): # First advance the clock to match the time of the stream. This has # a side-effect of also advancing this cache's clock. tag, r, curr_timestamp = events_to_send.pop() if curr_timestamp > self._monotonic_clock: yield self._advance_processing_time(curr_timestamp) # Then, send either a new element or watermark. if r.HasField('element_event'): r.element_event.tag = tag yield r elif r.HasField('watermark_event'): r.watermark_event.tag = tag yield r unsent_events = events_to_send target_timestamp = self._min_timestamp_of(unsent_events) def _advance_processing_time(self, new_timestamp): """Advances the internal clock and returns an AdvanceProcessingTime event. """ advancy_by = new_timestamp.micros - self._monotonic_clock.micros e = beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=advancy_by)) self._monotonic_clock = new_timestamp return e ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/caching/streaming_cache_test.py0000644000175100017510000005526515206113032030024 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest from apache_beam import coders from apache_beam.options.pipeline_options import StandardOptions from apache_beam.portability.api import beam_interactive_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive.cache_manager import SafeFastPrimitivesCoder from apache_beam.runners.interactive.caching.cacheable import CacheKey from apache_beam.runners.interactive.caching.streaming_cache import StreamingCache from apache_beam.runners.interactive.testing.test_cache_manager import FileRecordsBuilder from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import * from apache_beam.transforms.window import TimestampedValue # Nose automatically detects tests if they match a regex. Here, it mistakens # these protos as tests. For more info see the Nose docs at: # https://nose.readthedocs.io/en/latest/writing_tests.html beam_runner_api_pb2.TestStreamPayload.__test__ = False # type: ignore[attr-defined] beam_interactive_api_pb2.TestStreamFileHeader.__test__ = False # type: ignore[attr-defined] beam_interactive_api_pb2.TestStreamFileRecord.__test__ = False # type: ignore[attr-defined] class StreamingCacheTest(unittest.TestCase): def setUp(self): pass def test_exists(self): cache = StreamingCache(cache_dir=None) self.assertFalse(cache.exists('my_label')) cache.write([beam_interactive_api_pb2.TestStreamFileRecord()], 'my_label') self.assertTrue(cache.exists('my_label')) # '' shouldn't be treated as a wildcard to match everything. self.assertFalse(cache.exists('')) def test_empty(self): CACHED_PCOLLECTION_KEY = repr(CacheKey('arbitrary_key', '', '', '')) cache = StreamingCache(cache_dir=None) self.assertFalse(cache.exists(CACHED_PCOLLECTION_KEY)) cache.write([], CACHED_PCOLLECTION_KEY) reader, _ = cache.read(CACHED_PCOLLECTION_KEY) # Assert that an empty reader returns an empty list. self.assertFalse([e for e in reader]) def test_size(self): cache = StreamingCache(cache_dir=None) cache.write([beam_interactive_api_pb2.TestStreamFileRecord()], 'my_label') coder = cache.load_pcoder('my_label') # Add one because of the new-line character that is also written. size = len( coder.encode( beam_interactive_api_pb2.TestStreamFileRecord().SerializeToString()) ) + 1 self.assertEqual(cache.size('my_label'), size) def test_clear(self): cache = StreamingCache(cache_dir=None) self.assertFalse(cache.exists('my_label')) cache.sink(['my_label'], is_capture=True) cache.write([beam_interactive_api_pb2.TestStreamFileRecord()], 'my_label') self.assertTrue(cache.exists('my_label')) self.assertEqual(cache.capture_keys, set(['my_label'])) self.assertTrue(cache.clear('my_label')) self.assertFalse(cache.exists('my_label')) self.assertFalse(cache.capture_keys) def test_single_reader(self): """ Tests that we expect to see all the correctly emitted TestStreamPayloads. """ CACHED_PCOLLECTION_KEY = repr(CacheKey('arbitrary_key', '', '', '')) values = (FileRecordsBuilder(tag=CACHED_PCOLLECTION_KEY) .add_element(element=0, event_time_secs=0) .advance_processing_time(1) .add_element(element=1, event_time_secs=1) .advance_processing_time(1) .add_element(element=2, event_time_secs=2) .build()) # yapf: disable cache = StreamingCache(cache_dir=None) cache.write(values, CACHED_PCOLLECTION_KEY) reader, _ = cache.read(CACHED_PCOLLECTION_KEY) coder = coders.FastPrimitivesCoder() events = list(reader) # Units here are in microseconds. expected = [ beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode(0), timestamp=0) ], tag=CACHED_PCOLLECTION_KEY)), beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=1 * 10**6)), beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode(1), timestamp=1 * 10**6) ], tag=CACHED_PCOLLECTION_KEY)), beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=1 * 10**6)), beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode(2), timestamp=2 * 10**6) ], tag=CACHED_PCOLLECTION_KEY)), ] self.assertSequenceEqual(events, expected) def test_multiple_readers(self): """Tests that the service advances the clock with multiple outputs. """ CACHED_LETTERS = repr(CacheKey('letters', '', '', '')) CACHED_NUMBERS = repr(CacheKey('numbers', '', '', '')) CACHED_LATE = repr(CacheKey('late', '', '', '')) letters = (FileRecordsBuilder(CACHED_LETTERS) .advance_processing_time(1) .advance_watermark(watermark_secs=0) .add_element(element='a', event_time_secs=0) .advance_processing_time(10) .advance_watermark(watermark_secs=10) .add_element(element='b', event_time_secs=10) .build()) # yapf: disable numbers = (FileRecordsBuilder(CACHED_NUMBERS) .advance_processing_time(2) .add_element(element=1, event_time_secs=0) .advance_processing_time(1) .add_element(element=2, event_time_secs=0) .advance_processing_time(1) .add_element(element=2, event_time_secs=0) .build()) # yapf: disable late = (FileRecordsBuilder(CACHED_LATE) .advance_processing_time(101) .add_element(element='late', event_time_secs=0) .build()) # yapf: disable cache = StreamingCache(cache_dir=None) cache.write(letters, CACHED_LETTERS) cache.write(numbers, CACHED_NUMBERS) cache.write(late, CACHED_LATE) reader = cache.read_multiple([[CACHED_LETTERS], [CACHED_NUMBERS], [CACHED_LATE]]) coder = coders.FastPrimitivesCoder() events = list(reader) # Units here are in microseconds. expected = [ # Advances clock from 0 to 1 beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=1 * 10**6)), beam_runner_api_pb2.TestStreamPayload.Event( watermark_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceWatermark(new_watermark=0, tag=CACHED_LETTERS)), beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('a'), timestamp=0) ], tag=CACHED_LETTERS)), # Advances clock from 1 to 2 beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=1 * 10**6)), beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode(1), timestamp=0) ], tag=CACHED_NUMBERS)), # Advances clock from 2 to 3 beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=1 * 10**6)), beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode(2), timestamp=0) ], tag=CACHED_NUMBERS)), # Advances clock from 3 to 4 beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=1 * 10**6)), beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode(2), timestamp=0) ], tag=CACHED_NUMBERS)), # Advances clock from 4 to 11 beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=7 * 10**6)), beam_runner_api_pb2.TestStreamPayload.Event( watermark_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceWatermark(new_watermark=10 * 10**6, tag=CACHED_LETTERS)), beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('b'), timestamp=10 * 10**6) ], tag=CACHED_LETTERS)), # Advances clock from 11 to 101 beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=90 * 10**6)), beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('late'), timestamp=0) ], tag=CACHED_LATE)), ] self.assertSequenceEqual(events, expected) def test_read_and_write(self): """An integration test between the Sink and Source. This ensures that the sink and source speak the same language in terms of coders, protos, order, and units. """ CACHED_RECORDS = repr(CacheKey('records', '', '', '')) # Units here are in seconds. test_stream = ( TestStream(output_tags=(CACHED_RECORDS)) .advance_watermark_to(0, tag=CACHED_RECORDS) .advance_processing_time(5) .add_elements(['a', 'b', 'c'], tag=CACHED_RECORDS) .advance_watermark_to(10, tag=CACHED_RECORDS) .advance_processing_time(1) .add_elements( [ TimestampedValue('1', 15), TimestampedValue('2', 15), TimestampedValue('3', 15) ], tag=CACHED_RECORDS)) # yapf: disable coder = SafeFastPrimitivesCoder() cache = StreamingCache(cache_dir=None, sample_resolution_sec=1.0) # Assert that there are no capture keys at first. self.assertEqual(cache.capture_keys, set()) options = StandardOptions(streaming=True) with TestPipeline(options=options) as p: records = (p | test_stream)[CACHED_RECORDS] # pylint: disable=expression-not-assigned records | cache.sink([CACHED_RECORDS], is_capture=True) reader, _ = cache.read(CACHED_RECORDS) actual_events = list(reader) # Assert that the capture keys are forwarded correctly. self.assertEqual(cache.capture_keys, set([CACHED_RECORDS])) # Units here are in microseconds. expected_events = [ beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=5 * 10**6)), beam_runner_api_pb2.TestStreamPayload.Event( watermark_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceWatermark(new_watermark=0, tag=CACHED_RECORDS)), beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('a'), timestamp=0), beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('b'), timestamp=0), beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('c'), timestamp=0), ], tag=CACHED_RECORDS)), beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=1 * 10**6)), beam_runner_api_pb2.TestStreamPayload.Event( watermark_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceWatermark(new_watermark=10 * 10**6, tag=CACHED_RECORDS)), beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('1'), timestamp=15 * 10**6), beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('2'), timestamp=15 * 10**6), beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('3'), timestamp=15 * 10**6), ], tag=CACHED_RECORDS)), ] self.assertEqual(actual_events, expected_events) def test_read_and_write_multiple_outputs(self): """An integration test between the Sink and Source with multiple outputs. This tests the funcionatlity that the StreamingCache reads from multiple files and combines them into a single sorted output. """ LETTERS_TAG = repr(CacheKey('letters', '', '', '')) NUMBERS_TAG = repr(CacheKey('numbers', '', '', '')) # Units here are in seconds. test_stream = (TestStream() .advance_watermark_to(0, tag=LETTERS_TAG) .advance_processing_time(5) .add_elements(['a', 'b', 'c'], tag=LETTERS_TAG) .advance_watermark_to(10, tag=NUMBERS_TAG) .advance_processing_time(1) .add_elements( [ TimestampedValue('1', 15), TimestampedValue('2', 15), TimestampedValue('3', 15) ], tag=NUMBERS_TAG)) # yapf: disable cache = StreamingCache(cache_dir=None, sample_resolution_sec=1.0) coder = SafeFastPrimitivesCoder() options = StandardOptions(streaming=True) with TestPipeline(options=options) as p: # pylint: disable=expression-not-assigned events = p | test_stream events[LETTERS_TAG] | 'Letters sink' >> cache.sink([LETTERS_TAG]) events[NUMBERS_TAG] | 'Numbers sink' >> cache.sink([NUMBERS_TAG]) reader = cache.read_multiple([[LETTERS_TAG], [NUMBERS_TAG]]) actual_events = list(reader) # Units here are in microseconds. expected_events = [ beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=5 * 10**6)), beam_runner_api_pb2.TestStreamPayload.Event( watermark_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceWatermark(new_watermark=0, tag=LETTERS_TAG)), beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('a'), timestamp=0), beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('b'), timestamp=0), beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('c'), timestamp=0), ], tag=LETTERS_TAG)), beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=1 * 10**6)), beam_runner_api_pb2.TestStreamPayload.Event( watermark_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceWatermark(new_watermark=10 * 10**6, tag=NUMBERS_TAG)), beam_runner_api_pb2.TestStreamPayload.Event( watermark_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceWatermark(new_watermark=0, tag=LETTERS_TAG)), beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('1'), timestamp=15 * 10**6), beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('2'), timestamp=15 * 10**6), beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode('3'), timestamp=15 * 10**6), ], tag=NUMBERS_TAG)), ] self.assertListEqual(actual_events, expected_events) def test_always_default_coder_for_test_stream_records(self): CACHED_NUMBERS = repr(CacheKey('numbers', '', '', '')) numbers = (FileRecordsBuilder(CACHED_NUMBERS) .advance_processing_time(2) .add_element(element=1, event_time_secs=0) .advance_processing_time(1) .add_element(element=2, event_time_secs=0) .advance_processing_time(1) .add_element(element=2, event_time_secs=0) .build()) # yapf: disable cache = StreamingCache(cache_dir=None) cache.write(numbers, CACHED_NUMBERS) self.assertIs( type(cache.load_pcoder(CACHED_NUMBERS)), type(cache._default_pcoder)) def test_streaming_cache_does_not_write_non_record_or_header_types(self): cache = StreamingCache(cache_dir=None) self.assertRaises(TypeError, cache.write, 'some value', 'a key') def test_streaming_cache_uses_gcs_ib_cache_root(self): """ Checks that StreamingCache._cache_dir is set to the cache_root set under Interactive Beam for a GCS directory. """ # Set Interactive Beam specified cache dir to cloud storage ib.options.cache_root = 'gs://' cache_manager_with_ib_option = StreamingCache( cache_dir=ib.options.cache_root) self.assertEqual( ib.options.cache_root, cache_manager_with_ib_option._cache_dir) # Reset Interactive Beam setting ib.options.cache_root = None def test_streaming_cache_uses_local_ib_cache_root(self): """ Checks that StreamingCache._cache_dir is set to the cache_root set under Interactive Beam for a local directory and that the cached values are the same as the values of a cache using default settings. """ CACHED_PCOLLECTION_KEY = repr(CacheKey('arbitrary_key', '', '', '')) values = (FileRecordsBuilder(CACHED_PCOLLECTION_KEY) .advance_processing_time(1) .advance_watermark(watermark_secs=0) .add_element(element=1, event_time_secs=0) .build()) # yapf: disable local_cache = StreamingCache(cache_dir=None) local_cache.write(values, CACHED_PCOLLECTION_KEY) reader_one, _ = local_cache.read(CACHED_PCOLLECTION_KEY) pcoll_list_one = list(reader_one) # Set Interactive Beam specified cache dir to cloud storage ib.options.cache_root = '/tmp/it-test/' cache_manager_with_ib_option = StreamingCache( cache_dir=ib.options.cache_root) self.assertEqual( ib.options.cache_root, cache_manager_with_ib_option._cache_dir) cache_manager_with_ib_option.write(values, CACHED_PCOLLECTION_KEY) reader_two, _ = cache_manager_with_ib_option.read(CACHED_PCOLLECTION_KEY) pcoll_list_two = list(reader_two) self.assertEqual(pcoll_list_one, pcoll_list_two) # Reset Interactive Beam setting ib.options.cache_root = None if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/caching/write_cache.py0000644000175100017510000001675115206113032026123 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module to write cache for PCollections being computed. For internal use only; no backward-compatibility guarantees. """ # pytype: skip-file import apache_beam as beam from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.interactive import cache_manager as cache from apache_beam.runners.interactive.caching.cacheable import Cacheable from apache_beam.runners.interactive.caching.reify import reify_to_cache from apache_beam.runners.pipeline_context import PipelineContext from apache_beam.transforms.ptransform import PTransform class WriteCache: """Class that facilitates writing cache for PCollections being computed. """ def __init__( self, pipeline: beam_runner_api_pb2.Pipeline, context: PipelineContext, cache_manager: cache.CacheManager, cacheable: Cacheable): self._pipeline = pipeline self._context = context self._cache_manager = cache_manager self._cacheable = cacheable self._key = cacheable.to_key().to_str() def write_cache(self) -> None: """Writes cache for the cacheable PCollection that is being computed. First, it creates a temporary pipeline instance on top of the existing component_id_map from self._pipeline's context so that both pipelines share the context and have no conflict component ids. Second, it creates a _PCollectionPlaceHolder in the temporary pipeline that mimics the attributes of the cacheable PCollection to be written into cache. It also marks all components in the current temporary pipeline as ignorable when later copying components to self._pipeline. Third, it instantiates a _WriteCacheTransform that uses the _PCollectionPlaceHolder as the input. This adds a subgraph under top level transforms that writes the _PCollectionPlaceHolder into cache. Fourth, it copies components of the subgraph from the temporary pipeline to self._pipeline, skipping components that are ignored in the temporary pipeline and components that are not in the temporary pipeline but presents in the component_id_map of self._pipeline. Last, it replaces inputs of all transforms that consume the _PCollectionPlaceHolder with the cacheable PCollection to be written to cache. """ template, write_input_placeholder = self._build_runner_api_template() input_placeholder_id = self._context.pcollections.get_id( write_input_placeholder.placeholder_pcoll) input_id = self._context.pcollections.get_id(self._cacheable.pcoll) # Copy cache writing subgraph from the template to the pipeline proto. for pcoll_id in template.components.pcollections: if (pcoll_id in self._pipeline.components.pcollections or pcoll_id in write_input_placeholder.ignorable_components.pcollections): continue self._pipeline.components.pcollections[pcoll_id].CopyFrom( template.components.pcollections[pcoll_id]) for coder_id in template.components.coders: if (coder_id in self._pipeline.components.coders or coder_id in write_input_placeholder.ignorable_components.coders): continue self._pipeline.components.coders[coder_id].CopyFrom( template.components.coders[coder_id]) for windowing_strategy_id in template.components.windowing_strategies: if (windowing_strategy_id in self._pipeline.components.windowing_strategies or windowing_strategy_id in write_input_placeholder.ignorable_components.windowing_strategies): continue self._pipeline.components.windowing_strategies[ windowing_strategy_id].CopyFrom( template.components.windowing_strategies[windowing_strategy_id]) template_root_transform_id = template.root_transform_ids[0] root_transform_id = self._pipeline.root_transform_ids[0] for transform_id in template.components.transforms: if (transform_id in self._pipeline.components.transforms or transform_id in write_input_placeholder.ignorable_components.transforms): continue self._pipeline.components.transforms[transform_id].CopyFrom( template.components.transforms[transform_id]) for top_level_transform in template.components.transforms[ template_root_transform_id].subtransforms: if (top_level_transform in write_input_placeholder.ignorable_components.transforms): continue self._pipeline.components.transforms[ root_transform_id].subtransforms.append(top_level_transform) # Replace all the input pcoll of input_placeholder_id from cache writing # with cacheable pcoll of input_id. for transform in self._pipeline.components.transforms.values(): inputs = transform.inputs if input_placeholder_id in inputs.values(): keys_need_replacement = set() for key in inputs: if inputs[key] == input_placeholder_id: keys_need_replacement.add(key) for key in keys_need_replacement: inputs[key] = input_id def _build_runner_api_template( self) -> tuple[beam_runner_api_pb2.Pipeline, '_PCollectionPlaceHolder']: pph = _PCollectionPlaceHolder(self._cacheable.pcoll, self._context) transform = _WriteCacheTransform(self._cache_manager, self._key) _ = pph.placeholder_pcoll | 'sink_cache_' + self._key >> transform return pph.placeholder_pcoll.pipeline.to_runner_api(), pph class _WriteCacheTransform(PTransform): """A composite transform encapsulates writing cache for PCollections. """ def __init__(self, cache_manager: cache.CacheManager, key: str): self._cache_manager = cache_manager self._key = key def expand(self, pcoll: beam.pvalue.PCollection) -> beam.pvalue.PValue: return reify_to_cache( pcoll=pcoll, cache_key=self._key, cache_manager=self._cache_manager) class _PCollectionPlaceHolder: """A placeholder as an input to the cache writing transform. """ def __init__(self, pcoll: beam.pvalue.PCollection, context: PipelineContext): tmp_pipeline = beam.Pipeline() tmp_pipeline.component_id_map = context.component_id_map self._input_placeholder = tmp_pipeline | 'CreatePInput' >> beam.Create( [], reshuffle=False) self._input_placeholder.tag = pcoll.tag self._input_placeholder.element_type = pcoll.element_type self._input_placeholder.is_bounded = pcoll.is_bounded self._input_placeholder._windowing = pcoll.windowing self._ignorable_components = tmp_pipeline.to_runner_api().components @property def placeholder_pcoll(self) -> beam.pvalue.PCollection: return self._input_placeholder @property def ignorable_components(self) -> beam_runner_api_pb2.Components: """Subgraph generated by the placeholder that can be ignored in the final pipeline proto. """ return self._ignorable_components ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/caching/write_cache_test.py0000644000175100017510000000573415206113032027161 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for write_cache.""" # pytype: skip-file import unittest from unittest.mock import patch import apache_beam as beam from apache_beam.runners.interactive import augmented_pipeline as ap from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.caching import write_cache from apache_beam.runners.interactive.testing.pipeline_assertion import assert_pipeline_proto_equal from apache_beam.runners.interactive.testing.test_cache_manager import InMemoryCache class WriteCacheTest(unittest.TestCase): def setUp(self): ie.new_env() @patch( 'apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.get_cache_manager') def test_write_cache(self, mocked_get_cache_manager): p = beam.Pipeline() pcoll = p | beam.Create([1, 2, 3]) ib.watch(locals()) cache_manager = InMemoryCache() mocked_get_cache_manager.return_value = cache_manager aug_p = ap.AugmentedPipeline(p) key = repr(aug_p._cacheables[pcoll].to_key()) pipeline_proto = p.to_runner_api() # Write cache on the pipeline proto. write_cache.WriteCache( pipeline_proto, aug_p._context, aug_p._cache_manager, aug_p._cacheables[pcoll]).write_cache() actual_pipeline = pipeline_proto # Write cache directly on the piepline instance. transform = write_cache._WriteCacheTransform(aug_p._cache_manager, key) _ = pcoll | 'sink_cache_' + key >> transform expected_pipeline = p.to_runner_api() assert_pipeline_proto_equal(self, expected_pipeline, actual_pipeline) # Check if the actual_pipeline uses pcoll as an input of a write transform. pcoll_id = aug_p._context.pcollections.get_id(pcoll) write_transform_id = None for transform_id, transform in \ actual_pipeline.components.transforms.items(): if pcoll_id in transform.inputs.values(): write_transform_id = transform_id break self.assertIsNotNone(write_transform_id) self.assertIn( 'sink', actual_pipeline.components.transforms[write_transform_id].unique_name) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.6314342 apache_beam-2.74.0/apache_beam/runners/interactive/dataproc/0000755000175100017510000000000015206113160023465 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/dataproc/__init__.py0000644000175100017510000000142015206113032025571 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/dataproc/dataproc_cluster_manager.py0000644000175100017510000003430615206113032031073 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import re import time from typing import Optional from apache_beam import version as beam_version from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.dataproc.types import ClusterMetadata from apache_beam.runners.interactive.utils import obfuscate from apache_beam.runners.interactive.utils import progress_indicated try: from google.cloud import dataproc_v1 from apache_beam.io.gcp import gcsfilesystem # pylint: disable=ungrouped-imports except ImportError: class UnimportedDataproc: Cluster = None dataproc_v1 = UnimportedDataproc() _LOGGER = logging.getLogger(__name__) # Name of the log file auto-generated by Dataproc. We use it to locate the # startup output of the Flink daemon to retrieve master url and dashboard # information. DATAPROC_STAGING_LOG_NAME = 'dataproc-initialization-script-0_output' # Home dir of os user yarn. YARN_HOME = '/var/lib/hadoop-yarn' # Configures the os user yarn to use gcloud as the docker credHelper. # Also sets some taskmanager configurations for better parallelism. # Finally starts the yarn application: flink cluster in session mode. INIT_ACTION = """#!/bin/bash sudo -u yarn gcloud auth configure-docker --quiet readonly FLINK_INSTALL_DIR='/usr/lib/flink' readonly MASTER_HOSTNAME="$(/usr/share/google/get_metadata_value attributes/dataproc-master)" cat <>${FLINK_INSTALL_DIR}/conf/flink-conf.yaml taskmanager.memory.network.fraction: 0.2 taskmanager.memory.network.min: 64mb taskmanager.memory.network.max: 1gb EOF sed -i \ "s/^taskmanager.network.numberOfBuffers: 2048/taskmanager.network.numberOfBuffers: 8192/" \ ${FLINK_INSTALL_DIR}/conf/flink-conf.yaml if [[ "${HOSTNAME}" == "${MASTER_HOSTNAME}" ]]; then . /usr/bin/flink-yarn-daemon fi """ class DataprocClusterManager: """Self-contained cluster manager that controls the lifecyle of a Dataproc cluster connected by one or more pipelines under Interactive Beam. """ def __init__(self, cluster_metadata: ClusterMetadata) -> None: """Initializes the DataprocClusterManager with properties required to interface with the Dataproc ClusterControllerClient. """ self.cluster_metadata = cluster_metadata # Pipelines whose jobs are executed on the cluster. self.pipelines = set() self._cluster_client = dataproc_v1.ClusterControllerClient( client_options={ 'api_endpoint': \ f'{self.cluster_metadata.region}-dataproc.googleapis.com:443' }) self._fs = gcsfilesystem.GCSFileSystem(PipelineOptions()) self._staging_directory = None cache_dir = ie.current_env().options.cache_root if not cache_dir.startswith('gs://'): error_msg = ( 'ib.options.cache_root needs to be a Cloud Storage ' 'Bucket to cache source recording and PCollections in current ' f'interactive setup, instead \'{cache_dir}\' is assigned.') _LOGGER.error(error_msg) raise ValueError(error_msg) self._cache_root = cache_dir.rstrip('/') def stage_init_action(self) -> str: """Stages the initialization action script to GCS cache root to set up Dataproc clusters. Returns the staged gcs file path. """ # Versionizes the initialization action script. init_action_ver = obfuscate(INIT_ACTION) path = f'{self._cache_root}/dataproc-init-action-{init_action_ver}.sh' if not self._fs.exists(path): with self._fs.create(path) as bwriter: bwriter.write(INIT_ACTION.encode()) return path @progress_indicated def create_cluster(self, cluster: dict) -> None: """Attempts to create a cluster using attributes that were initialized with the DataprocClusterManager instance. Args: cluster: Dictionary representing Dataproc cluster. Read more about the schema for clusters here: https://cloud.google.com/python/docs/reference/dataproc/latest/google.cloud.dataproc_v1.types.Cluster """ if self.cluster_metadata.master_url: return try: self._cluster_client.create_cluster( request={ 'project_id': self.cluster_metadata.project_id, 'region': self.cluster_metadata.region, 'cluster': cluster }) except Exception as e: if e.code == 409: _LOGGER.info( 'Cluster %s already exists. Continuing...', self.cluster_metadata.cluster_name) elif e.code == 403: _LOGGER.error( 'Due to insufficient project permissions, ' 'unable to create cluster: %s', self.cluster_metadata.cluster_name) raise ValueError( 'You cannot create a cluster in project: {}'.format( self.cluster_metadata.project_id)) elif e.code == 501: _LOGGER.error( 'Invalid region provided: %s', self.cluster_metadata.region) raise ValueError( 'Region {} does not exist!'.format(self.cluster_metadata.region)) else: _LOGGER.error( 'Unable to create cluster: %s', self.cluster_metadata.cluster_name) raise e else: _LOGGER.info( 'Cluster created successfully: %s', self.cluster_metadata.cluster_name) self._staging_directory = self.get_staging_location() master_url, dashboard = self.get_master_url_and_dashboard() self.cluster_metadata.master_url = master_url self.cluster_metadata.dashboard = dashboard def create_flink_cluster(self) -> None: """Calls _create_cluster with a configuration that enables FlinkRunner.""" init_action_path = self.stage_init_action() # https://cloud.google.com/php/docs/reference/cloud-dataproc/latest/V1.Cluster cluster = { 'project_id': self.cluster_metadata.project_id, 'cluster_name': self.cluster_metadata.cluster_name, 'config': { 'software_config': { # TODO(https://github.com/apache/beam/issues/21527): Uncomment # these lines when a Dataproc image is released with previously # missing dependencies. # 'image_version': ie.current_env().clusters. # DATAPROC_IMAGE_VERSION, 'optional_components': ['DOCKER', 'FLINK'], 'properties': { # Enforces HOME dir for user yarn. 'yarn:yarn.nodemanager.user-home-dir': YARN_HOME } }, 'initialization_actions': [{ 'executable_file': init_action_path }], 'gce_cluster_config': { 'metadata': { 'flink-start-yarn-session': 'false' }, 'service_account_scopes': [ 'https://www.googleapis.com/auth/cloud-platform' ], 'internal_ip_only': False }, 'master_config': { # There must be 1 and only 1 instance of master. 'num_instances': 1 }, 'worker_config': {}, 'endpoint_config': { 'enable_http_port_access': True } }, 'labels': { 'goog-dataflow-notebook': beam_version.__version__.replace( '.', '_') } } # Additional gce_cluster_config. gce_cluster_config = cluster['config']['gce_cluster_config'] if self.cluster_metadata.subnetwork: gce_cluster_config['subnetwork_uri'] = self.cluster_metadata.subnetwork # Additional InstanceGroupConfig for master and workers. master_config = cluster['config']['master_config'] worker_config = cluster['config']['worker_config'] if self.cluster_metadata.num_workers: worker_config['num_instances'] = self.cluster_metadata.num_workers if self.cluster_metadata.machine_type: master_config['machine_type_uri'] = self.cluster_metadata.machine_type worker_config['machine_type_uri'] = self.cluster_metadata.machine_type self.create_cluster(cluster) def cleanup(self) -> None: """Deletes the cluster that uses the attributes initialized with the DataprocClusterManager instance.""" try: self._cluster_client.delete_cluster( request={ 'project_id': self.cluster_metadata.project_id, 'region': self.cluster_metadata.region, 'cluster_name': self.cluster_metadata.cluster_name, }) self.cleanup_staging_files() except Exception as e: if e.code == 403: _LOGGER.error( 'Due to insufficient project permissions, ' 'unable to clean up the default cluster: %s', self.cluster_metadata.cluster_name) raise ValueError( 'You cannot delete a cluster in project: {}'.format( self.cluster_metadata.project_id)) elif e.code == 404: _LOGGER.error( 'Cluster does not exist: %s', self.cluster_metadata.cluster_name) raise ValueError( 'Cluster was not found: {}'.format( self.cluster_metadata.cluster_name)) else: _LOGGER.error( 'Failed to delete cluster: %s', self.cluster_metadata.cluster_name) raise e def get_cluster_details(self) -> dataproc_v1.Cluster: """Gets the Dataproc_v1 Cluster object for the current cluster manager.""" try: return self._cluster_client.get_cluster( request={ 'project_id': self.cluster_metadata.project_id, 'region': self.cluster_metadata.region, 'cluster_name': self.cluster_metadata.cluster_name }) except Exception as e: if e.code == 403: _LOGGER.error( 'Due to insufficient project permissions, ' 'unable to retrieve information for cluster: %s', self.cluster_metadata.cluster_name) raise ValueError( 'You cannot view clusters in project: {}'.format( self.cluster_metadata.project_id)) elif e.code == 404: _LOGGER.error( 'Cluster does not exist: %s', self.cluster_metadata.cluster_name) raise ValueError( 'Cluster was not found: {}'.format( self.cluster_metadata.cluster_name)) else: _LOGGER.error( 'Failed to get information for cluster: %s', self.cluster_metadata.cluster_name) raise e def wait_for_cluster_to_provision(self) -> None: while self.get_cluster_details().status.state.name == 'CREATING': time.sleep(15) def get_staging_location(self) -> str: """Gets the staging bucket of an existing Dataproc cluster.""" try: self.wait_for_cluster_to_provision() cluster_details = self.get_cluster_details() bucket_name = cluster_details.config.config_bucket gcs_path = 'gs://' + bucket_name + '/google-cloud-dataproc-metainfo/' for file in self._fs._list(gcs_path): if self.cluster_metadata.cluster_name in file.path: # this file path split will look something like: # ['gs://.../google-cloud-dataproc-metainfo/{staging_dir}/', # '-{node-type}/dataproc-initialization-script-0_output'] return file.path.split(self.cluster_metadata.cluster_name)[0] except Exception as e: _LOGGER.error( 'Failed to get %s cluster staging bucket.', self.cluster_metadata.cluster_name) raise e def parse_master_url_and_dashboard(self, line: str) -> tuple[str, str]: """Parses the master_url and YARN application_id of the Flink process from an input line. The line containing both the master_url and application id is always formatted as such: {text} Found Web Interface {master_url} of application '{application_id}'.\\n Truncated example where '...' represents additional text between segments: ... google-dataproc-startup[000]: ... activate-component-flink[0000]: ...org.apache.flink.yarn.YarnClusterDescriptor... [] - Found Web Interface example-master-url:50000 of application 'application_123456789000_0001'. Returns the flink_master_url and dashboard link as a tuple.""" cluster_details = self.get_cluster_details() yarn_endpoint = cluster_details.config.endpoint_config.http_ports[ 'YARN ResourceManager'] segment = line.split('Found Web Interface ')[1].split(' of application ') master_url = segment[0] application_id = re.sub('\'|.\n', '', segment[1]) dashboard = re.sub( '/yarn/', '/gateway/default/yarn/proxy/' + application_id + '/', yarn_endpoint) return master_url, dashboard def get_master_url_and_dashboard(self) -> tuple[Optional[str], Optional[str]]: """Returns the master_url of the current cluster.""" startup_logs = [] for file in self._fs._list(self._staging_directory): if DATAPROC_STAGING_LOG_NAME in file.path: startup_logs.append(file.path) for log in startup_logs: content = self._fs.open(log) for line in content.readlines(): decoded_line = line.decode() if 'Found Web Interface' in decoded_line: return self.parse_master_url_and_dashboard(decoded_line) return None, None def cleanup_staging_files(self) -> None: if self._staging_directory: staging_files = [ file.path for file in self._fs._list(self._staging_directory) ] self._fs.delete(staging_files) if self._cache_root: cache_files = [file.path for file in self._fs._list(self._cache_root)] self._fs.delete(cache_files) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/dataproc/dataproc_cluster_manager_test.py0000644000175100017510000003504115206113032032127 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.runners.interactive.dataproc. dataproc_cluster_manager.""" # pytype: skip-file import unittest from unittest.mock import patch from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import DataprocClusterManager from apache_beam.runners.interactive.dataproc.types import ClusterMetadata try: from google.cloud import dataproc_v1 # pylint: disable=unused-import except ImportError: _dataproc_imported = False else: _dataproc_imported = True class MockProperty: def __init__(self, property, value): object.__setattr__(self, property, value) class MockException(Exception): def __init__(self, code=-1): self.code = code class MockCluster: def __init__(self, config_bucket=None): self.config = MockProperty('config_bucket', config_bucket) self.status = MockProperty('state', MockProperty('name', None)) class MockFileSystem: def _list(self, dir=None): return [ MockProperty( 'path', 'test-path/dataproc-initialization-script-0_output') ] def open(self, dir=None): return MockFileIO('test-line Found Web Interface test-master-url' \ ' of application \'test-app-id\'.\n') class MockFileIO: def __init__(self, contents): self.contents = contents def readlines(self): return [self.contents.encode('utf-8')] @unittest.skipIf(not _dataproc_imported, 'dataproc package was not imported.') class DataprocClusterManagerTest(unittest.TestCase): """Unit test for DataprocClusterManager""" def setUp(self): self.patcher = patch( 'apache_beam.runners.interactive.interactive_environment.current_env') self.m_env = self.patcher.start() self.m_env().clusters = ib.Clusters() self.m_env().options.cache_root = 'gs://fake' def tearDown(self): self.m_env().options.cache_root = None self.patcher.stop() @patch( 'google.cloud.dataproc_v1.ClusterControllerClient.create_cluster', side_effect=MockException(409)) def test_create_cluster_default_already_exists(self, mock_cluster_client): """ Tests that no exception is thrown when a cluster already exists, but is using ie.current_env().clusters.default_cluster_name. """ cluster_metadata = ClusterMetadata( project_id='test-project', region='test-region') cluster_manager = DataprocClusterManager(cluster_metadata) from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import _LOGGER with self.assertLogs(_LOGGER, level='INFO') as context_manager: cluster_manager.create_cluster({}) self.assertTrue('already exists' in context_manager.output[0]) @patch( 'google.cloud.dataproc_v1.ClusterControllerClient.create_cluster', side_effect=MockException(403)) def test_create_cluster_permission_denied(self, mock_cluster_client): """ Tests that an exception is thrown when a user is trying to write to a project while having insufficient permissions. """ cluster_metadata = ClusterMetadata( project_id='test-project', region='test-region') cluster_manager = DataprocClusterManager(cluster_metadata) from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import _LOGGER with self.assertLogs(_LOGGER, level='ERROR') as context_manager: self.assertRaises(ValueError, cluster_manager.create_cluster, {}) self.assertTrue( 'Due to insufficient project permissions' in context_manager.output[0]) @patch( 'google.cloud.dataproc_v1.ClusterControllerClient.create_cluster', side_effect=MockException(501)) def test_create_cluster_region_does_not_exist(self, mock_cluster_client): """ Tests that an exception is thrown when a user specifies a region that does not exist. """ cluster_metadata = ClusterMetadata( project_id='test-project', region='test-region') cluster_manager = DataprocClusterManager(cluster_metadata) from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import _LOGGER with self.assertLogs(_LOGGER, level='ERROR') as context_manager: self.assertRaises(ValueError, cluster_manager.create_cluster, {}) self.assertTrue('Invalid region provided' in context_manager.output[0]) @patch( 'google.cloud.dataproc_v1.ClusterControllerClient.create_cluster', side_effect=MockException()) def test_create_cluster_other_exception(self, mock_cluster_client): """ Tests that an exception is thrown when the exception is not handled by any other case under _create_cluster. """ cluster_metadata = ClusterMetadata( project_id='test-project', region='test-region') cluster_manager = DataprocClusterManager(cluster_metadata) from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import _LOGGER with self.assertLogs(_LOGGER, level='ERROR') as context_manager: self.assertRaises(MockException, cluster_manager.create_cluster, {}) self.assertTrue('Unable to create cluster' in context_manager.output[0]) @patch( 'apache_beam.runners.interactive.dataproc.dataproc_cluster_manager.' 'DataprocClusterManager.cleanup_staging_files', return_value=None) @patch( 'google.cloud.dataproc_v1.ClusterControllerClient.delete_cluster', side_effect=MockException(403)) def test_cleanup_permission_denied(self, mock_cluster_client, mock_cleanup): """ Tests that an exception is thrown when a user is trying to delete a project that they have insufficient permissions for. """ cluster_metadata = ClusterMetadata( project_id='test-project', region='test-region') cluster_manager = DataprocClusterManager(cluster_metadata) from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import _LOGGER with self.assertLogs(_LOGGER, level='ERROR') as context_manager: self.assertRaises(ValueError, cluster_manager.cleanup) self.assertTrue( 'Due to insufficient project permissions' in context_manager.output[0]) @patch( 'apache_beam.runners.interactive.dataproc.dataproc_cluster_manager.' 'DataprocClusterManager.cleanup_staging_files', return_value=None) @patch( 'google.cloud.dataproc_v1.ClusterControllerClient.delete_cluster', side_effect=MockException(404)) def test_cleanup_does_not_exist(self, mock_cluster_client, mock_cleanup): """ Tests that an exception is thrown when cleanup attempts to delete a cluster that does not exist. """ cluster_metadata = ClusterMetadata( project_id='test-project', region='test-region') cluster_manager = DataprocClusterManager(cluster_metadata) from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import _LOGGER with self.assertLogs(_LOGGER, level='ERROR') as context_manager: self.assertRaises(ValueError, cluster_manager.cleanup) self.assertTrue('Cluster does not exist' in context_manager.output[0]) @patch( 'apache_beam.runners.interactive.dataproc.dataproc_cluster_manager.' 'DataprocClusterManager.cleanup_staging_files', return_value=None) @patch( 'google.cloud.dataproc_v1.ClusterControllerClient.delete_cluster', side_effect=MockException()) def test_cleanup_other_exception(self, mock_cluster_client, mock_cleanup): """ Tests that an exception is thrown when the exception is not handled by any other case under cleanup. """ cluster_metadata = ClusterMetadata( project_id='test-project', region='test-region') cluster_manager = DataprocClusterManager(cluster_metadata) from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import _LOGGER with self.assertLogs(_LOGGER, level='ERROR') as context_manager: self.assertRaises(MockException, cluster_manager.cleanup) self.assertTrue('Failed to delete cluster' in context_manager.output[0]) @patch( 'apache_beam.io.gcp.gcsfilesystem.GCSFileSystem._list', return_value=[ MockProperty( 'path', 'gs://test-bucket/google-cloud-dataproc-metainfo' '/test-cluster/item') ]) @patch( 'google.cloud.dataproc_v1.ClusterControllerClient.get_cluster', return_value=MockCluster('test-bucket')) def test_get_staging_location(self, mock_cluster_client, mock_list): """ Test to receive a mock staging location successfully under get_staging_location. """ cluster_metadata = ClusterMetadata( project_id='test-project', region='test-region', cluster_name='test-cluster') cluster_manager = DataprocClusterManager(cluster_metadata) self.assertEqual( cluster_manager.get_staging_location(), 'gs://test-bucket/google-cloud-dataproc-metainfo/') @patch( 'google.cloud.dataproc_v1.ClusterControllerClient.get_cluster', side_effect=MockException()) def test_get_staging_location_exception(self, mock_cluster_client): """ Test to catch when an error is raised inside get_staging_location. """ cluster_metadata = ClusterMetadata( project_id='test-project', region='test-region', cluster_name='test-cluster') cluster_manager = DataprocClusterManager(cluster_metadata) with self.assertRaises(MockException): cluster_manager.get_staging_location() @patch( 'apache_beam.runners.interactive.dataproc.dataproc_cluster_manager.' 'DataprocClusterManager.get_cluster_details', return_value=MockProperty( 'config', MockProperty( 'endpoint_config', MockProperty( 'http_ports', {'YARN ResourceManager': 'test-resource-manager/yarn/'})))) def test_parse_master_url_and_dashboard(self, mock_cluster_details): """ Tests that parse_master_url_and_dashboard properly parses the input string and produces a mock master_url and mock dashboard link. """ cluster_metadata = ClusterMetadata( project_id='test-project', region='test-region') cluster_manager = DataprocClusterManager(cluster_metadata) line = 'test-line Found Web Interface test-master-url' \ ' of application \'test-app-id\'.\n' master_url, dashboard = cluster_manager.parse_master_url_and_dashboard(line) self.assertEqual('test-master-url', master_url) self.assertEqual( 'test-resource-manager/gateway/default/yarn/proxy/test-app-id/', dashboard) @patch( 'google.cloud.dataproc_v1.ClusterControllerClient.get_cluster', side_effect=MockException(403)) def test_get_cluster_details_permission_denied(self, mock_cluster_client): """ Tests that an exception is thrown when a user is trying to get information for a project without sufficient permissions to do so. """ cluster_metadata = ClusterMetadata( project_id='test-project', region='test-region') cluster_manager = DataprocClusterManager(cluster_metadata) from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import _LOGGER with self.assertLogs( _LOGGER, level='ERROR') as context_manager, self.assertRaises(ValueError): cluster_manager.get_cluster_details() self.assertTrue( 'Due to insufficient project permissions' in context_manager.output[0]) @patch( 'google.cloud.dataproc_v1.ClusterControllerClient.get_cluster', side_effect=MockException(404)) def test_get_cluster_details_does_not_exist(self, mock_cluster_client): """ Tests that an exception is thrown when cleanup attempts to get information for a cluster that does not exist. """ cluster_metadata = ClusterMetadata( project_id='test-project', region='test-region') cluster_manager = DataprocClusterManager(cluster_metadata) from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import _LOGGER with self.assertLogs( _LOGGER, level='ERROR') as context_manager, self.assertRaises(ValueError): cluster_manager.get_cluster_details() self.assertTrue('Cluster does not exist' in context_manager.output[0]) @patch( 'google.cloud.dataproc_v1.ClusterControllerClient.get_cluster', side_effect=MockException()) def test_get_cluster_details_other_exception(self, mock_cluster_client): """ Tests that an exception is thrown when the exception is not handled by any other case under get_cluster_details. """ cluster_metadata = ClusterMetadata( project_id='test-project', region='test-region') cluster_manager = DataprocClusterManager(cluster_metadata) from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import _LOGGER with self.assertLogs( _LOGGER, level='ERROR') as context_manager, self.assertRaises(MockException): cluster_manager.get_cluster_details() self.assertTrue( 'Failed to get information for cluster' in context_manager.output[0]) @patch( 'apache_beam.runners.interactive.dataproc.dataproc_cluster_manager.' 'DataprocClusterManager.parse_master_url_and_dashboard', return_value=('test-master-url', 'test-dashboard-link')) def test_get_master_url_and_dashboard(self, mock_parse_method): """ Tests that get_master_url_and_dashboard detect the line containing the unique substring which identifies the location of the master_url and application id of the Flink master. """ cluster_metadata = ClusterMetadata( project_id='test-project', region='test-region') cluster_manager = DataprocClusterManager(cluster_metadata) cluster_manager._fs = MockFileSystem() cluster_metadata._staging_directory = 'test-staging-bucket' master_url, dashboard = cluster_manager.get_master_url_and_dashboard() self.assertEqual(master_url, 'test-master-url') self.assertEqual(dashboard, 'test-dashboard-link') if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/dataproc/types.py0000644000175100017510000000551015206113032025202 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import uuid from dataclasses import dataclass from dataclasses import field from typing import Optional from typing import Union from apache_beam.pipeline import Pipeline def _generate_unique_cluster_name(): return f'interactive-beam-{uuid.uuid4().hex}' @dataclass class ClusterMetadata: """Metadata of a provisioned worker cluster that executes Beam pipelines. Apache Beam supports running Beam pipelines on different runners provisioned in different setups based on the runner and pipeline options associated with each pipeline. To provide similar portability features, Interactive Beam automatically extracts such ClusterMetadata information from pipeline options of a pipeline in the REPL context and provision suitable clusters to execute the pipeline. The lifecyle of the clusters is managed by Interactive Beam and the user doesn not need to interact with it. It's not recommended to build this ClusterMetadata from raw values nor use it to interact with the cluster management logic directly. Interactive Beam now supports:: 1. Runner: FlinkRunner; Setup: on Google Cloud with Flink on Dataproc. """ project_id: Optional[str] = None region: Optional[str] = 'us-central1' cluster_name: Optional[str] = field( default_factory=_generate_unique_cluster_name) # From WorkerOptions. subnetwork: Optional[str] = None num_workers: Optional[int] = None machine_type: Optional[int] = None # Derivative fields do not affect hash or comparison. master_url: Optional[str] = None dashboard: Optional[str] = None def __key(self): return ( self.project_id, self.region, self.cluster_name, self.subnetwork, self.num_workers, self.machine_type) def __hash__(self): return hash(self.__key()) def __eq__(self, other): if isinstance(other, ClusterMetadata): return self.__key() == other.__key() return False def reset_name(self): self.cluster_name = _generate_unique_cluster_name() ClusterIdentifier = Union[str, Pipeline, ClusterMetadata] ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.637435 apache_beam-2.74.0/apache_beam/runners/interactive/display/0000755000175100017510000000000015206113160023335 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/display/__init__.py0000644000175100017510000000142015206113032025441 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/display/display_manager.py0000644000175100017510000001543215206113032027051 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Manages displaying pipeline graph and execution status on the frontend. This module is experimental. No backwards-compatibility guarantees. """ # pytype: skip-file import collections import threading import time from typing import TYPE_CHECKING from apache_beam.runners.interactive.display import interactive_pipeline_graph try: import IPython # pylint: disable=import-error from IPython import get_ipython # pylint: disable=import-error from IPython.display import display as ip_display # pylint: disable=import-error # _display_progress defines how outputs are printed on the frontend. _display_progress = ip_display if not TYPE_CHECKING: def _formatter(string, pp, cycle): # pylint: disable=unused-argument pp.text(string) if get_ipython(): plain = get_ipython().display_formatter.formatters['text/plain'] # pylint: disable=undefined-variable plain.for_type(str, _formatter) except ImportError: IPython = None _display_progress = print class DisplayManager(object): """Manages displaying pipeline graph and execution status on the frontend.""" def __init__( self, pipeline_proto, pipeline_analyzer, cache_manager, pipeline_graph_renderer): """Constructor of DisplayManager. Args: pipeline_proto: (Pipeline proto) pipeline_analyzer: (PipelineAnalyzer) the pipeline analyzer that corresponds to this round of execution. This will provide more detailed informations about the pipeline cache_manager: (interactive_runner.CacheManager) DisplayManager fetches the latest status of pipeline execution by querying cache_manager. pipeline_graph_renderer: (pipeline_graph_renderer.PipelineGraphRenderer) decides how a pipeline graph is rendered. """ # Every parameter except cache_manager is expected to remain constant. self._analyzer = pipeline_analyzer self._cache_manager = cache_manager self._pipeline_graph = interactive_pipeline_graph.InteractivePipelineGraph( pipeline_proto, required_transforms=self._analyzer.tl_required_trans_ids(), referenced_pcollections=self._analyzer.tl_referenced_pcoll_ids(), cached_pcollections=self._analyzer.caches_used()) self._renderer = pipeline_graph_renderer # _text_to_print keeps track of information to be displayed. self._text_to_print = collections.OrderedDict() self._text_to_print['summary'] = ( 'Using %s cached PCollections\nExecuting %s of %s ' 'transforms.') % ( len(self._analyzer.caches_used()), ( len(self._analyzer.tl_required_trans_ids()) - len(self._analyzer.read_cache_ids()) - len(self._analyzer.write_cache_ids())), len( pipeline_proto.components.transforms[ pipeline_proto.root_transform_ids[0]].subtransforms)) self._text_to_print.update( {pcoll_id: "" for pcoll_id in self._analyzer.tl_referenced_pcoll_ids()}) # _pcollection_stats maps pcoll_id to # { 'cache_label': cache_label, version': version, 'sample': pcoll_in_list } self._pcollection_stats = {} for pcoll_id in self._analyzer.tl_referenced_pcoll_ids(): self._pcollection_stats[pcoll_id] = { 'cache_label': self._analyzer.pipeline_info().cache_label(pcoll_id), 'version': -1, 'sample': [] } self._producers = {} for _, transform in pipeline_proto.components.transforms.items(): for pcoll_id in transform.outputs.values(): if pcoll_id not in self._producers or '/' not in transform.unique_name: self._producers[pcoll_id] = transform.unique_name # For periodic update. self._lock = threading.Lock() self._periodic_update = False def update_display(self, force=False): """Updates display on the frontend. Retrieves the latest execution status by querying CacheManager and updates display on the fronend. The assumption is that there is only one pipeline in a cell, because it clears up everything in the cell output every update cycle. Args: force: (bool) whether to force updating when no stats change happens. """ with self._lock: stats_updated = False for pcoll_id, stats in self._pcollection_stats.items(): cache_label = stats['cache_label'] version = stats['version'] if force or not self._cache_manager.is_latest_version( version, 'sample', cache_label): pcoll_list, version = self._cache_manager.read('sample', cache_label) stats['sample'] = list(pcoll_list) stats['version'] = version stats_updated = True if pcoll_id in self._analyzer.tl_referenced_pcoll_ids(): self._text_to_print[pcoll_id] = ( str( '%s produced %s' % ( self._producers[pcoll_id], interactive_pipeline_graph.format_sample(pcoll_list, 5)))) if force or stats_updated: self._pipeline_graph.update_pcollection_stats(self._pcollection_stats) if IPython: from IPython import display display.clear_output(True) rendered_graph = self._renderer.render_pipeline_graph( self._pipeline_graph) display.display(display.HTML(rendered_graph)) _display_progress('Running...') for text in self._text_to_print.values(): if text != "": _display_progress(text) def start_periodic_update(self): """Start a thread that periodically updates the display.""" self.update_display(True) self._periodic_update = True def _updater(): while self._periodic_update: self.update_display() time.sleep(.02) t = threading.Thread(target=_updater) t.daemon = True t.start() def stop_periodic_update(self): """Stop periodically updating the display.""" self.update_display(True) self._periodic_update = False ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/display/interactive_pipeline_graph.py0000644000175100017510000001254215206113032031274 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Helper to render pipeline graph in IPython when running interactively. This module is experimental. No backwards-compatibility guarantees. """ # pytype: skip-file import re from apache_beam.runners.interactive.display import pipeline_graph def nice_str(o): s = repr(o) s = s.replace('"', "'") s = s.replace('\\', '|') s = re.sub(r'[^\x20-\x7F]', ' ', s) assert '"' not in s if len(s) > 35: s = s[:35] + '...' return s def format_sample(contents, count=1000): contents = list(contents) elems = ', '.join([nice_str(o) for o in contents[:count]]) if len(contents) > count: elems += ', ...' assert '"' not in elems return '{%s}' % elems class InteractivePipelineGraph(pipeline_graph.PipelineGraph): """Creates the DOT representation of an interactive pipeline. Thread-safe.""" def __init__( self, pipeline, required_transforms=None, referenced_pcollections=None, cached_pcollections=None): """Constructor of PipelineGraph. Args: pipeline: (Pipeline proto) or (Pipeline) pipeline to be rendered. required_transforms: (list/set of str) ID of top level PTransforms that lead to visible results. referenced_pcollections: (list/set of str) ID of PCollections that are referenced by top level PTransforms executed (i.e. required_transforms) cached_pcollections: (set of str) a set of PCollection IDs of those whose cached results are used in the execution. """ self._required_transforms = required_transforms or set() self._referenced_pcollections = referenced_pcollections or set() self._cached_pcollections = cached_pcollections or set() super().__init__( pipeline=pipeline, default_vertex_attrs={ 'color': 'gray', 'fontcolor': 'gray' }, default_edge_attrs={'color': 'gray'}) transform_updates, pcollection_updates = self._generate_graph_update_dicts() self._update_graph(transform_updates, pcollection_updates) def update_pcollection_stats(self, pcollection_stats): """Updates PCollection stats. Args: pcollection_stats: (dict of dict) maps PCollection IDs to informations. In particular, we only care about the field 'sample' which should be a the PCollection result in as a list. """ edge_dict = {} for pcoll_id, stats in pcollection_stats.items(): attrs = {} pcoll_list = stats['sample'] if pcoll_list: attrs['label'] = format_sample(pcoll_list, 1) attrs['labeltooltip'] = format_sample(pcoll_list, 10) else: attrs['label'] = '?' edge_dict[pcoll_id] = attrs self._update_graph(edge_dict=edge_dict) def _generate_graph_update_dicts(self): """Generate updates specific to interactive pipeline. Returns: vertex_dict: (Dict[str, Dict[str, str]]) maps vertex name to attributes edge_dict: (Dict[str, Dict[str, str]]) maps vertex name to attributes """ transform_dict = {} # maps PTransform IDs to properties pcoll_dict = {} # maps PCollection IDs to properties for transform_id, transform_proto in self._top_level_transforms(): transform_dict[transform_proto.unique_name] = { 'required': transform_id in self._required_transforms } for pcoll_id in transform_proto.outputs.values(): pcoll_dict[pcoll_id] = { 'cached': pcoll_id in self._cached_pcollections, 'referenced': pcoll_id in self._referenced_pcollections } def vertex_properties_to_attributes(vertex): """Converts PCollection properties to DOT vertex attributes.""" attrs = {} if 'leaf' in vertex: attrs['style'] = 'invis' elif vertex.get('required'): attrs['color'] = 'blue' attrs['fontcolor'] = 'blue' else: attrs['color'] = 'grey' return attrs def edge_properties_to_attributes(edge): """Converts PTransform properties to DOT edge attributes.""" attrs = {} if edge.get('cached'): attrs['color'] = 'red' elif edge.get('referenced'): attrs['color'] = 'black' else: attrs['color'] = 'grey' return attrs vertex_dict = {} # maps vertex names to attributes edge_dict = {} # maps edge names to attributes for transform_name, transform_properties in transform_dict.items(): vertex_dict[transform_name] = vertex_properties_to_attributes( transform_properties) for pcoll_id, pcoll_properties in pcoll_dict.items(): edge_dict[pcoll_id] = edge_properties_to_attributes(pcoll_properties) return vertex_dict, edge_dict ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/display/pcoll_visualization.py0000644000175100017510000005166115206113032030010 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module visualizes PCollection data. For internal use only; no backwards-compatibility guarantees. Only works with Python 3.5+. """ # pytype: skip-file import base64 import datetime import html import logging from datetime import timedelta from typing import Optional from dateutil import tz import apache_beam as beam from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.utils import elements_to_df from apache_beam.transforms.window import GlobalWindow from apache_beam.transforms.window import IntervalWindow try: from facets_overview.generic_feature_statistics_generator import \ GenericFeatureStatisticsGenerator # pylint: disable=import-error from IPython import get_ipython # pylint: disable=import-error from IPython.display import HTML # pylint: disable=import-error from IPython.display import Javascript # pylint: disable=import-error from IPython.display import display # pylint: disable=import-error from IPython.display import display_javascript # pylint: disable=import-error from timeloop import Timeloop # pylint: disable=import-error if get_ipython(): _pcoll_visualization_ready = True else: _pcoll_visualization_ready = False except ImportError: _pcoll_visualization_ready = False _LOGGER = logging.getLogger(__name__) _CSS = """ """ _DIVE_SCRIPT_TEMPLATE = """ try {{ document .getElementById("{display_id}") .contentDocument .getElementById("{display_id}") .data = {jsonstr}; }} catch (e) {{ // NOOP when the user has cleared the output from the notebook. }}""" _DIVE_HTML_TEMPLATE = _CSS + """ """ _OVERVIEW_SCRIPT_TEMPLATE = """ try {{ document .getElementById("{display_id}") .contentDocument .getElementById("{display_id}") .protoInput = "{protostr}"; }} catch (e) {{ // NOOP when the user has cleared the output from the notebook. }}""" _OVERVIEW_HTML_TEMPLATE = _CSS + """ """ _DATATABLE_INITIALIZATION_CONFIG = """ bAutoWidth: false, columns: {columns}, destroy: true, responsive: true, columnDefs: [ {{ targets: "_all", className: "dt-left" }}, {{ "targets": 0, "width": "10px", "title": "" }} ]""" _DATAFRAME_SCRIPT_TEMPLATE = """ var dt; if ($.fn.dataTable.isDataTable("#{table_id}")) {{ dt = $("#{table_id}").dataTable(); }} else if ($("#{table_id}_wrapper").length == 0) {{ dt = $("#{table_id}").dataTable({{ """ + _DATATABLE_INITIALIZATION_CONFIG + """ }}); }} else {{ return; }} dt.api() .clear() .rows.add({data_as_rows}) .draw('full-hold');""" _DATAFRAME_PAGINATION_TEMPLATE = _CSS + """
""" _NO_DATA_TEMPLATE = _CSS + """
No data to display.
""" _NO_DATA_REMOVAL_SCRIPT = """ $("#no_data_{id}").remove();""" def visualize( stream, dynamic_plotting_interval=None, include_window_info=False, display_facets=False, element_type=None): """Visualizes the data of a given PCollection. Optionally enables dynamic plotting with interval in seconds if the PCollection is being produced by a running pipeline or the pipeline is streaming indefinitely. The function always returns immediately and is asynchronous when dynamic plotting is on. If dynamic plotting enabled, the visualization is updated continuously until the pipeline producing the PCollection is in an end state. The visualization would be anchored to the notebook cell output area. The function asynchronously returns a handle to the visualization job immediately. The user could manually do:: # In one notebook cell, enable dynamic plotting every 1 second: handle = visualize(pcoll, dynamic_plotting_interval=1) # Visualization anchored to the cell's output area. # In a different cell: handle.stop() # Will stop the dynamic plotting of the above visualization manually. # Otherwise, dynamic plotting ends when pipeline is not running anymore. If dynamic_plotting is not enabled (by default), None is returned. If include_window_info is True, the data will include window information, which consists of the event timestamps, windows, and pane info. If display_facets is True, the facets widgets will be rendered. Otherwise, the facets widgets will not be rendered. The function is experimental. For internal use only; no backwards-compatibility guarantees. """ if not _pcoll_visualization_ready: return None pv = PCollectionVisualization( stream, include_window_info=include_window_info, display_facets=display_facets, element_type=element_type) if ie.current_env().is_in_notebook: pv.display() else: pv.display_plain_text() # We don't want to do dynamic plotting if there is no notebook frontend. return None if dynamic_plotting_interval: # Disables the verbose logging from timeloop. logging.getLogger('timeloop').disabled = True tl = Timeloop() def dynamic_plotting(stream, pv, tl, include_window_info, display_facets): @tl.job(interval=timedelta(seconds=dynamic_plotting_interval)) def continuous_update_display(): # pylint: disable=unused-variable # Always creates a new PCollVisualization instance when the # PCollection materialization is being updated and dynamic # plotting is in-process. # PCollectionVisualization created at this level doesn't need dynamic # plotting interval information when instantiated because it's already # in dynamic plotting logic. updated_pv = PCollectionVisualization( stream, include_window_info=include_window_info, display_facets=display_facets, element_type=element_type) updated_pv.display(updating_pv=pv) # Stop updating the visualizations as soon as the stream will not yield # new elements. if stream.is_done(): try: tl.stop() except RuntimeError: # The job can only be stopped once. Ignore excessive stops. pass tl.start() return tl return dynamic_plotting(stream, pv, tl, include_window_info, display_facets) return None def visualize_computed_pcoll( pcoll_name: str, pcoll: beam.pvalue.PCollection, max_n: int, max_duration_secs: float, dynamic_plotting_interval: Optional[int] = None, include_window_info: bool = False, display_facets: bool = False) -> None: """A simple visualize alternative. When the pcoll_name and pcoll pair identifies a watched and computed PCollection in the current interactive environment without ambiguity, an ElementStream can be built directly from cache. Returns immediately, the visualization is asynchronous, but guaranteed to end in the near future. Args: pcoll_name: the variable name of the PCollection. pcoll: the PCollection to be visualized. max_n: the maximum number of elements to visualize. max_duration_secs: max duration of elements to read in seconds. dynamic_plotting_interval: the interval in seconds between visualization updates if provided; otherwise, no dynamic plotting. include_window_info: whether to include windowing info in the elements. display_facets: whether to display the facets widgets. """ pipeline = ie.current_env().user_pipeline(pcoll.pipeline) rm = ie.current_env().get_recording_manager(pipeline, create_if_absent=True) stream = rm.read( pcoll_name, pcoll, max_n=max_n, max_duration_secs=max_duration_secs) if stream: visualize( stream, dynamic_plotting_interval=dynamic_plotting_interval, include_window_info=include_window_info, display_facets=display_facets, element_type=pcoll.element_type) class PCollectionVisualization(object): """A visualization of a PCollection. The class relies on creating a PipelineInstrument w/o actual instrument to access current interactive environment for materialized PCollection data at the moment of self instantiation through cache. """ def __init__( self, stream, include_window_info=False, display_facets=False, element_type=None): assert _pcoll_visualization_ready, ( 'Dependencies for PCollection visualization are not available. Please ' 'use `pip install apache-beam[interactive]` to install necessary ' 'dependencies and make sure that you are executing code in an ' 'interactive environment such as a Jupyter notebook.') self._stream = stream # Variable name as the title for element value in the rendered data table. self._pcoll_var = stream.var if not self._pcoll_var: self._pcoll_var = 'Value' obfuscated_id = stream.display_id(id(self)) self._dive_display_id = 'facets_dive_{}'.format(obfuscated_id) self._overview_display_id = 'facets_overview_{}'.format(obfuscated_id) self._df_display_id = 'df_{}'.format(obfuscated_id) self._include_window_info = include_window_info self._display_facets = display_facets self._is_datatable_empty = True self._element_type = element_type def display_plain_text(self): """Displays a head sample of the normalized PCollection data. This function is used when the ipython kernel is not connected to a notebook frontend such as when running ipython in terminal or in unit tests. It's a visualization in terminal-like UI, not a function to retrieve data for programmatically usages. """ # Double check if the dependency is ready in case someone mistakenly uses # the function. if _pcoll_visualization_ready: data = self._to_dataframe() # Displays a data-table with at most 25 entries from the head. data_sample = data.head(25) display(data_sample) def display(self, updating_pv=None): """Displays the visualization through IPython. Args: updating_pv: A PCollectionVisualization object. When provided, the display_id of each visualization part will inherit from the initial display of updating_pv and only update that visualization web element instead of creating new ones. The visualization has 3 parts: facets-dive, facets-overview and paginated data table. Each part is assigned an auto-generated unique display id (the uniqueness is guaranteed throughout the lifespan of the PCollection variable). """ # Ensures that dive, overview and table render the same data because the # materialized PCollection data might being updated continuously. data = self._to_dataframe() # Give the numbered column names when visualizing. data.columns = [ self._pcoll_var + '.' + str(column) if isinstance(column, int) else column for column in data.columns ] # String-ify the dictionaries for display because elements of type dict # cannot be ordered. data = data.map(lambda x: str(x) if isinstance(x, dict) else x) if updating_pv: # Only updates when data is not empty. Otherwise, consider it a bad # iteration and noop since there is nothing to be updated. if data.empty: _LOGGER.debug('Skip a visualization update due to empty data.') else: self._display_dataframe(data.copy(deep=True), updating_pv) if self._display_facets: self._display_dive(data.copy(deep=True), updating_pv) self._display_overview(data.copy(deep=True), updating_pv) else: self._display_dataframe(data.copy(deep=True)) if self._display_facets: self._display_dive(data.copy(deep=True)) self._display_overview(data.copy(deep=True)) def _display_dive(self, data, update=None): sprite_size = 32 if len(data.index) > 50000 else 64 format_window_info_in_dataframe(data) jsonstr = data.to_json(orient='records', default_handler=str) if update: script = _DIVE_SCRIPT_TEMPLATE.format( display_id=update._dive_display_id, jsonstr=jsonstr) display_javascript(Javascript(script)) else: html_str = _DIVE_HTML_TEMPLATE.format( display_id=self._dive_display_id, jsonstr=html.escape(jsonstr), sprite_size=sprite_size) display(HTML(html_str)) def _display_overview(self, data, update=None): if (not data.empty and self._include_window_info and all(column in data.columns for column in ('event_time', 'windows', 'pane_info'))): data = data.drop(['event_time', 'windows', 'pane_info'], axis=1) # GFSG expects all column names to be strings. data.columns = data.columns.astype(str) gfsg = GenericFeatureStatisticsGenerator() proto = gfsg.ProtoFromDataFrames([{'name': 'data', 'table': data}]) protostr = base64.b64encode(proto.SerializeToString()).decode('utf-8') if update: script = _OVERVIEW_SCRIPT_TEMPLATE.format( display_id=update._overview_display_id, protostr=protostr) display_javascript(Javascript(script)) else: html_str = _OVERVIEW_HTML_TEMPLATE.format( display_id=self._overview_display_id, protostr=protostr) display(HTML(html_str)) def _display_dataframe(self, data, update=None): table_id = 'table_{}'.format( update._df_display_id if update else self._df_display_id) columns = [{ 'title': '' }] + [{ 'title': str(column) } for column in data.columns] format_window_info_in_dataframe(data) # Convert the dataframe into rows, each row looks like # [column_1_val, column_2_val, ...]. rows = data.map(lambda x: str(x)).to_dict('split')['data'] # Convert each row into dict where keys are column index in the datatable # to be rendered and values are data from the dataframe. Column index 0 is # left out to hold the int index (not part of the data) from dataframe. # Each row becomes: {1: column_1_val, 2: column_2_val, ...}. rows = [{k + 1: v for k, v in enumerate(row)} for row in rows] # Add the dataframe int index (used as default ordering column) to datatable # column index 0 (will be rendered as the first column). # Each row becomes: # {1: column_1_val, 2: column_2_val, ..., 0: int_index_in_dataframe}. for k, row in enumerate(rows): row[0] = k script = _DATAFRAME_SCRIPT_TEMPLATE.format( table_id=table_id, columns=columns, data_as_rows=rows) script_in_jquery_with_datatable = ie._JQUERY_WITH_DATATABLE_TEMPLATE.format( customized_script=script) # Dynamically load data into the existing datatable if not empty. if update and not update._is_datatable_empty: display_javascript(Javascript(script_in_jquery_with_datatable)) else: if data.empty: html_str = _NO_DATA_TEMPLATE.format(id=table_id) else: html_str = _DATAFRAME_PAGINATION_TEMPLATE.format( table_id=table_id, script_in_jquery_with_datatable=script_in_jquery_with_datatable) if update: if not data.empty: # Initialize a datatable to replace the existing no data div. display( Javascript( ie._JQUERY_WITH_DATATABLE_TEMPLATE.format( customized_script=_NO_DATA_REMOVAL_SCRIPT.format( id=table_id)))) display(HTML(html_str), display_id=update._df_display_id) update._is_datatable_empty = False else: display(HTML(html_str), display_id=self._df_display_id) if not data.empty: self._is_datatable_empty = False def _to_dataframe(self): results = list(self._stream.read(tail=False)) return elements_to_df( results, self._include_window_info, element_type=self._element_type) def format_window_info_in_dataframe(data): if 'event_time' in data.columns: data['event_time'] = data['event_time'].apply(event_time_formatter) if 'windows' in data.columns: data['windows'] = data['windows'].apply(windows_formatter) if 'pane_info' in data.columns: data['pane_info'] = data['pane_info'].apply(pane_info_formatter) def event_time_formatter(event_time_us): options = ie.current_env().options to_tz = options.display_timezone try: return ( datetime.datetime.utcfromtimestamp(event_time_us / 1000000).replace( tzinfo=tz.tzutc()).astimezone(to_tz).strftime( options.display_timestamp_format)) except ValueError: if event_time_us < 0: return 'Min Timestamp' return 'Max Timestamp' def windows_formatter(windows): result = [] for w in windows: if isinstance(w, GlobalWindow): result.append(str(w)) elif isinstance(w, IntervalWindow): # First get the duration in terms of hours, minutes, seconds, and # micros. duration = w.end.micros - w.start.micros duration_secs = duration // 1000000 hours, remainder = divmod(duration_secs, 3600) minutes, seconds = divmod(remainder, 60) micros = (duration - duration_secs * 1000000) % 1000000 # Construct the duration string. Try and write the string in such a # way that minimizes the amount of characters written. duration = '' if hours: duration += '{}h '.format(hours) if minutes or (hours and seconds): duration += '{}m '.format(minutes) if seconds: if micros: duration += '{}.{:06}s'.format(seconds, micros) else: duration += '{}s'.format(seconds) start = event_time_formatter(w.start.micros) result.append('{} ({})'.format(start, duration)) return ','.join(result) def pane_info_formatter(pane_info): from apache_beam.utils.windowed_value import PaneInfo from apache_beam.utils.windowed_value import PaneInfoTiming assert isinstance(pane_info, PaneInfo) result = 'Pane {}'.format(pane_info.index) timing_info = '{}{}'.format( 'Final ' if pane_info.is_last else '', PaneInfoTiming.to_string(pane_info.timing).lower().capitalize() if pane_info.timing in (PaneInfoTiming.EARLY, PaneInfoTiming.LATE) else '') if timing_info: result += ': ' + timing_info return result ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/display/pcoll_visualization_test.py0000644000175100017510000002075715206113032031051 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.runners.interactive.display.pcoll_visualization.""" # pytype: skip-file import unittest from unittest.mock import ANY from unittest.mock import PropertyMock from unittest.mock import patch import pytz import apache_beam as beam from apache_beam.runners import runner from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import interactive_runner as ir from apache_beam.runners.interactive.display import pcoll_visualization as pv from apache_beam.runners.interactive.recording_manager import RecordingManager from apache_beam.runners.interactive.testing.mock_ipython import mock_get_ipython from apache_beam.transforms.window import GlobalWindow from apache_beam.transforms.window import IntervalWindow from apache_beam.utils.windowed_value import PaneInfo from apache_beam.utils.windowed_value import PaneInfoTiming try: import timeloop except ImportError: pass @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') class PCollectionVisualizationTest(unittest.TestCase): def setUp(self): ie.new_env() # Allow unit test to run outside of ipython kernel since we don't test the # frontend rendering in unit tests. pv._pcoll_visualization_ready = True # Generally test the logic where notebook is connected to the assumed # ipython kernel by forcefully setting notebook check to True. ie.current_env()._is_in_notebook = True ib.options.display_timezone = pytz.timezone('US/Pacific') self._p = beam.Pipeline(ir.InteractiveRunner()) # pylint: disable=bad-option-value self._pcoll = self._p | 'Create' >> beam.Create(range(5)) ib.watch(self) ie.current_env().track_user_pipelines() recording_manager = RecordingManager(self._p) recording = recording_manager.record([self._pcoll], max_n=5, max_duration=5) self._stream = recording.stream(self._pcoll) def test_pcoll_visualization_generate_unique_display_id(self): pv_1 = pv.PCollectionVisualization(self._stream) pv_2 = pv.PCollectionVisualization(self._stream) self.assertNotEqual(pv_1._dive_display_id, pv_2._dive_display_id) self.assertNotEqual(pv_1._overview_display_id, pv_2._overview_display_id) self.assertNotEqual(pv_1._df_display_id, pv_2._df_display_id) @patch('IPython.get_ipython', new_callable=mock_get_ipython) @patch( 'apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.is_in_notebook', new_callable=PropertyMock) def test_one_shot_visualization_not_return_handle( self, mocked_is_in_notebook, unused): mocked_is_in_notebook.return_value = True self.assertIsNone(pv.visualize(self._stream, display_facets=True)) @patch('IPython.get_ipython', new_callable=mock_get_ipython) @patch( 'apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.is_in_notebook', new_callable=PropertyMock) def test_dynamic_plotting_return_handle(self, mocked_is_in_notebook, unused): mocked_is_in_notebook.return_value = True h = pv.visualize( self._stream, dynamic_plotting_interval=1, display_facets=True) self.assertIsInstance(h, timeloop.Timeloop) h.stop() @patch('IPython.get_ipython', new_callable=mock_get_ipython) @patch( 'apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.is_in_notebook', new_callable=PropertyMock) def test_no_dynamic_plotting_when_not_in_notebook( self, mocked_is_in_notebook, unused): mocked_is_in_notebook.return_value = False h = pv.visualize( self._stream, dynamic_plotting_interval=1, display_facets=True) self.assertIsNone(h) @patch( 'apache_beam.runners.interactive.display.pcoll_visualization' '.PCollectionVisualization._display_dive') @patch( 'apache_beam.runners.interactive.display.pcoll_visualization' '.PCollectionVisualization._display_overview') @patch( 'apache_beam.runners.interactive.display.pcoll_visualization' '.PCollectionVisualization._display_dataframe') def test_dynamic_plotting_updates_same_display( self, mocked_display_dataframe, mocked_display_overview, mocked_display_dive): original_pcollection_visualization = pv.PCollectionVisualization( self._stream, display_facets=True) # Dynamic plotting always creates a new PCollectionVisualization. new_pcollection_visualization = pv.PCollectionVisualization( self._stream, display_facets=True) # The display uses ANY data the moment display is invoked, and updates # web elements with ids fetched from the given updating_pv. new_pcollection_visualization.display( updating_pv=original_pcollection_visualization) mocked_display_dataframe.assert_called_once_with( ANY, original_pcollection_visualization) # Below assertions are still true without newer calls. mocked_display_overview.assert_called_once_with( ANY, original_pcollection_visualization) mocked_display_dive.assert_called_once_with( ANY, original_pcollection_visualization) def test_auto_stop_dynamic_plotting_when_job_is_terminated(self): fake_pipeline_result = runner.PipelineResult(runner.PipelineState.RUNNING) ie.current_env().set_pipeline_result(self._p, fake_pipeline_result) # When job is running, the dynamic plotting will not be stopped. self.assertFalse(ie.current_env().is_terminated(self._p)) fake_pipeline_result = runner.PipelineResult(runner.PipelineState.DONE) ie.current_env().set_pipeline_result(self._p, fake_pipeline_result) # When job is done, the dynamic plotting will be stopped. self.assertTrue(ie.current_env().is_terminated(self._p)) @patch('pandas.DataFrame.head') def test_display_plain_text_when_kernel_has_no_frontend(self, _mocked_head): # Resets the notebook check to False. ie.current_env()._is_in_notebook = False self.assertIsNone(pv.visualize(self._stream, display_facets=True)) _mocked_head.assert_called_once() def test_event_time_formatter(self): # In microseconds: Monday, March 2, 2020 3:14:54 PM GMT-08:00 event_time_us = 1583190894000000 self.assertEqual( '2020-03-02 15:14:54.000000-0800', pv.event_time_formatter(event_time_us)) def test_event_time_formatter_overflow_lower_bound(self): # A relatively small negative event time, which could be valid in Beam but # has no meaning when visualized. event_time_us = -100000000000000000 self.assertEqual('Min Timestamp', pv.event_time_formatter(event_time_us)) def test_event_time_formatter_overflow_upper_bound(self): # A relatively large event time, which exceeds the upper bound of unix time # Year 2038. It could mean infinite future in Beam but has no meaning # when visualized. # The value in test is supposed to be year 10000. event_time_us = 253402300800000000 self.assertEqual('Max Timestamp', pv.event_time_formatter(event_time_us)) def test_windows_formatter_global(self): gw = GlobalWindow() self.assertEqual(str(gw), pv.windows_formatter([gw])) def test_windows_formatter_interval(self): # The unit is second. iw = IntervalWindow(start=1583190894, end=1583200000) self.assertEqual( '2020-03-02 15:14:54.000000-0800 (2h 31m 46s)', pv.windows_formatter([iw])) def test_pane_info_formatter(self): self.assertEqual( 'Pane 0: Final Early', pv.pane_info_formatter( PaneInfo( is_first=False, is_last=True, timing=PaneInfoTiming.EARLY, index=0, nonspeculative_index=0))) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/display/pipeline_graph.py0000644000175100017510000002611115206113032026674 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """For generating Beam pipeline graph in DOT representation. This module is experimental. No backwards-compatibility guarantees. """ # pytype: skip-file import collections import logging import threading from typing import Iterator from typing import Union import apache_beam as beam from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import pipeline_instrument as inst from apache_beam.runners.interactive.display import pipeline_graph_renderer try: import pydot except ImportError: pass # pylint does not understand context # pylint:disable=dangerous-default-value class PipelineGraph(object): """Creates a DOT representing the pipeline. Thread-safe. Runner agnostic.""" def __init__( self, pipeline: Union[beam_runner_api_pb2.Pipeline, beam.Pipeline], default_vertex_attrs={'shape': 'box'}, default_edge_attrs=None, render_option=None): """Constructor of PipelineGraph. Examples: graph = pipeline_graph.PipelineGraph(pipeline_proto) graph.get_dot() or graph = pipeline_graph.PipelineGraph(pipeline) graph.get_dot() Args: pipeline: (Pipeline proto) or (Pipeline) pipeline to be rendered. default_vertex_attrs: (Dict[str, str]) a dict of default vertex attributes default_edge_attrs: (Dict[str, str]) a dict of default edge attributes render_option: (str) this parameter decides how the pipeline graph is rendered. See display.pipeline_graph_renderer for available options. """ self._lock = threading.Lock() self._graph: pydot.Dot = None self._pipeline_instrument = None if isinstance(pipeline, beam.Pipeline): self._pipeline_instrument = inst.PipelineInstrument( pipeline, pipeline._options) # The pre-process links user pipeline to runner pipeline through analysis # but without mutating runner pipeline. self._pipeline_instrument.preprocess() if isinstance(pipeline, beam_runner_api_pb2.Pipeline): self._pipeline_proto = pipeline elif isinstance(pipeline, beam.Pipeline): self._pipeline_proto = pipeline.to_runner_api() else: raise TypeError( 'pipeline should either be a %s or %s, while %s is given' % (beam_runner_api_pb2.Pipeline, beam.Pipeline, type(pipeline))) # A dict from PCollection ID to a list of its consuming Transform IDs self._consumers: collections.defaultdict[ str, list[str]] = collections.defaultdict(list) # A dict from PCollection ID to its producing Transform ID self._producers: dict[str, str] = {} for transform_id, transform_proto in self._top_level_transforms(): for pcoll_id in transform_proto.inputs.values(): self._consumers[pcoll_id].append(transform_id) for pcoll_id in transform_proto.outputs.values(): self._producers[pcoll_id] = transform_id default_vertex_attrs = default_vertex_attrs or {'shape': 'box'} if 'color' not in default_vertex_attrs: default_vertex_attrs['color'] = 'blue' if 'fontcolor' not in default_vertex_attrs: default_vertex_attrs['fontcolor'] = 'blue' vertex_dict, edge_dict = self._generate_graph_dicts() self._construct_graph( vertex_dict, edge_dict, default_vertex_attrs, default_edge_attrs) self._renderer = pipeline_graph_renderer.get_renderer(render_option) def get_dot(self) -> str: return self._get_graph().to_string() def display_graph(self): """Displays the graph generated.""" rendered_graph = self._renderer.render_pipeline_graph(self) if ie.current_env().is_in_notebook: try: from IPython import display display.display(display.HTML(rendered_graph)) except ImportError: # Unlikely to happen when is_in_notebook. logging.warning( 'Failed to import IPython display module when current ' 'environment is in a notebook. Cannot display the ' 'pipeline graph.') def _top_level_transforms( self) -> Iterator[tuple[str, beam_runner_api_pb2.PTransform]]: """Yields all top level PTransforms (subtransforms of the root PTransform). Yields: (str, PTransform proto) ID, proto pair of top level PTransforms. """ transforms = self._pipeline_proto.components.transforms for root_transform_id in self._pipeline_proto.root_transform_ids: root_transform_proto = transforms[root_transform_id] for top_level_transform_id in root_transform_proto.subtransforms: top_level_transform_proto = transforms[top_level_transform_id] yield top_level_transform_id, top_level_transform_proto def _decorate(self, value): """Decorates label-ish values used for rendering in dot language. Escapes special characters in the given str value for dot language. All PTransform unique names are escaped implicitly in this module when building dot representation. Otherwise, special characters will break the graph rendered or cause runtime errors. """ # Replace py str literal `\\` which is `\` in dot with py str literal # `\\\\` which is `\\` in dot so that dot `\\` can be rendered as `\`. Then # replace `"` with `\\"` so that the dot generated will be `\"` and be # rendered as `"`. return '"{}"'.format(value.replace('\\', '\\\\').replace('"', '\\"')) def _generate_graph_dicts(self): """From pipeline_proto and other info, generate the graph. Returns: vertex_dict: (Dict[str, Dict[str, str]]) vertex mapped to attributes. edge_dict: (Dict[(str, str), Dict[str, str]]) vertex pair mapped to the edge's attribute. """ transforms = self._pipeline_proto.components.transforms # A dict from vertex name (i.e. PCollection ID) to its attributes. vertex_dict = collections.defaultdict(dict) # A dict from vertex name pairs defining the edge (i.e. a pair of PTransform # IDs defining the PCollection) to its attributes. edge_dict = collections.defaultdict(dict) self._edge_to_vertex_pairs = collections.defaultdict(list) for _, transform in self._top_level_transforms(): vertex_dict[self._decorate(transform.unique_name)] = {} for pcoll_id in transform.outputs.values(): pcoll_node = None if self._pipeline_instrument: cacheable = self._pipeline_instrument.cacheables.get(pcoll_id) pcoll_node = cacheable.var if cacheable else None # If no PipelineInstrument is available or the PCollection is not # watched. if not pcoll_node: pcoll_node = 'pcoll%s' % (hash(pcoll_id) % 10000) vertex_dict[pcoll_node] = { 'shape': 'circle', 'label': '', # The pcoll node has no name. } # There is PipelineInstrument and the PCollection is watched with an # assigned variable. else: vertex_dict[pcoll_node] = {'shape': 'circle'} if pcoll_id not in self._consumers: self._edge_to_vertex_pairs[pcoll_id].append( (self._decorate(transform.unique_name), pcoll_node)) edge_dict[(self._decorate(transform.unique_name), pcoll_node)] = {} else: for consumer in self._consumers[pcoll_id]: producer_name = self._decorate(transform.unique_name) consumer_name = self._decorate(transforms[consumer].unique_name) self._edge_to_vertex_pairs[pcoll_id].append( (producer_name, pcoll_node)) edge_dict[(producer_name, pcoll_node)] = {} self._edge_to_vertex_pairs[pcoll_id].append( (pcoll_node, consumer_name)) edge_dict[(pcoll_node, consumer_name)] = {} return vertex_dict, edge_dict def _get_graph(self): """Returns pydot.Dot object for the pipeline graph. The purpose of this method is to avoid accessing the graph while it is updated. No one except for this method should be accessing _graph directly. Returns: (pydot.Dot) """ with self._lock: return self._graph def _construct_graph( self, vertex_dict, edge_dict, default_vertex_attrs, default_edge_attrs): """Constructs the pydot.Dot object for the pipeline graph. Args: vertex_dict: (Dict[str, Dict[str, str]]) maps vertex names to attributes edge_dict: (Dict[(str, str), Dict[str, str]]) maps vertex name pairs to attributes default_vertex_attrs: (Dict[str, str]) a dict of attributes default_edge_attrs: (Dict[str, str]) a dict of attributes """ with self._lock: self._graph = pydot.Dot() if default_vertex_attrs: self._graph.set_node_defaults(**default_vertex_attrs) if default_edge_attrs: self._graph.set_edge_defaults(**default_edge_attrs) self._vertex_refs = {} # Maps vertex name to pydot.Node self._edge_refs = {} # Maps vertex name pairs to pydot.Edge for vertex, vertex_attrs in vertex_dict.items(): vertex_ref = pydot.Node(vertex, **vertex_attrs) self._vertex_refs[vertex] = vertex_ref self._graph.add_node(vertex_ref) for edge, edge_attrs in edge_dict.items(): vertex_src = self._vertex_refs[edge[0]] vertex_dst = self._vertex_refs[edge[1]] edge_ref = pydot.Edge(vertex_src, vertex_dst, **edge_attrs) self._edge_refs[edge] = edge_ref self._graph.add_edge(edge_ref) def _update_graph(self, vertex_dict=None, edge_dict=None): """Updates the pydot.Dot object with the given attribute update Args: vertex_dict: (Dict[str, Dict[str, str]]) maps vertex names to attributes edge_dict: This should be Either (Dict[str, Dict[str, str]]) which maps edge names to attributes Or (Dict[(str, str), Dict[str, str]]) which maps vertex pairs to edge attributes """ def set_attrs(ref, attrs): for attr_name, attr_val in attrs.items(): ref.set(attr_name, attr_val) with self._lock: if vertex_dict: for vertex, vertex_attrs in vertex_dict.items(): set_attrs(self._vertex_refs[vertex], vertex_attrs) if edge_dict: for edge, edge_attrs in edge_dict.items(): if isinstance(edge, tuple): set_attrs(self._edge_refs[edge], edge_attrs) else: for vertex_pair in self._edge_to_vertex_pairs[edge]: set_attrs(self._edge_refs[vertex_pair], edge_attrs) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/display/pipeline_graph_renderer.py0000644000175100017510000000712515206113032030566 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """For rendering pipeline graph in HTML-compatible format. This module is experimental. No backwards-compatibility guarantees. """ # pytype: skip-file import abc import os import subprocess from typing import TYPE_CHECKING from typing import Optional from apache_beam.utils.plugin import BeamPlugin if TYPE_CHECKING: from apache_beam.runners.interactive.display.pipeline_graph import PipelineGraph class PipelineGraphRenderer(BeamPlugin, metaclass=abc.ABCMeta): """Abstract class for renderers, who decide how pipeline graphs are rendered. """ @classmethod @abc.abstractmethod def option(cls) -> str: """The corresponding rendering option for the renderer. """ raise NotImplementedError @abc.abstractmethod def render_pipeline_graph(self, pipeline_graph: 'PipelineGraph') -> str: """Renders the pipeline graph in HTML-compatible format. Args: pipeline_graph: (pipeline_graph.PipelineGraph) the graph to be rendererd. Returns: unicode, str or bytes that can be expressed as HTML. """ raise NotImplementedError class MuteRenderer(PipelineGraphRenderer): """Use this renderer to mute the pipeline display. """ @classmethod def option(cls) -> str: return 'mute' def render_pipeline_graph(self, pipeline_graph: 'PipelineGraph') -> str: return '' class TextRenderer(PipelineGraphRenderer): """This renderer simply returns the dot representation in text format. """ @classmethod def option(cls) -> str: return 'text' def render_pipeline_graph(self, pipeline_graph: 'PipelineGraph') -> str: return pipeline_graph.get_dot() class PydotRenderer(PipelineGraphRenderer): """This renderer renders the graph using pydot. It depends on 1. The software Graphviz: https://www.graphviz.org/ 2. The python module pydot: https://pypi.org/project/pydot/ """ @classmethod def option(cls) -> str: return 'graph' def render_pipeline_graph(self, pipeline_graph: 'PipelineGraph') -> str: return pipeline_graph._get_graph().create_svg().decode("utf-8") # pylint: disable=protected-access def get_renderer(option: Optional[str] = None) -> type[PipelineGraphRenderer]: """Get an instance of PipelineGraphRenderer given rendering option. Args: option: (str) the rendering option. Returns: (PipelineGraphRenderer) """ if option is None: if os.name == 'nt': exists = subprocess.call(['where', 'dot.exe']) == 0 else: exists = subprocess.call(['which', 'dot']) == 0 if exists: option = 'graph' else: option = 'text' renderer = [ r for r in PipelineGraphRenderer.get_all_subclasses() if option == r.option() ] if len(renderer) == 0: raise ValueError() elif len(renderer) == 1: return renderer[0]() else: raise ValueError('Found more than one renderer for option: %s', option) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/display/pipeline_graph_test.py0000644000175100017510000001135515206113032027737 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.runners.interactive.display.pipeline_graph.""" # pytype: skip-file import unittest from unittest.mock import patch import apache_beam as beam from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import interactive_runner as ir from apache_beam.runners.interactive.display import pipeline_graph from apache_beam.runners.interactive.testing.mock_ipython import mock_get_ipython # pylint: disable=bad-option-value,unused-variable,possibly-unused-variable # Reason: # Disable pylint for pipelines built for testing. Not all PCollections are # used but they need to be assigned to variables so that we can test how # interactive beam applies the magic around user-defined variables. # The tests need graphviz to work. @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') class PipelineGraphTest(unittest.TestCase): def setUp(self): ie.new_env() def test_decoration(self): p = beam.Pipeline(ir.InteractiveRunner()) # We are examining if literal `"` and trailing literal `\` are decorated # correctly. pcoll = p | '"[1]": "Create\\"' >> beam.Create(range(1000)) ib.watch(locals()) self.assertEqual( ( 'digraph G {\n' 'node [color=blue, fontcolor=blue, shape=box];\n' # The py string literal from `\\\\\\"` is `\\\"` in dot and will be # rendered as `\"` because they are enclosed by `"`. '"\\"[1]\\": \\"Create\\\\\\"";\n' 'pcoll [shape=circle];\n' '"\\"[1]\\": \\"Create\\\\\\"" -> pcoll;\n' '}\n'), pipeline_graph.PipelineGraph(p).get_dot()) def test_get_dot(self): p = beam.Pipeline(ir.InteractiveRunner()) init_pcoll = p | 'Init' >> beam.Create(range(10)) squares = init_pcoll | 'Square' >> beam.Map(lambda x: x * x) cubes = init_pcoll | 'Cube' >> beam.Map(lambda x: x**3) ib.watch(locals()) self.assertEqual(( 'digraph G {\n' 'node [color=blue, fontcolor=blue, shape=box];\n' '"Init";\n' 'init_pcoll [shape=circle];\n' '"Square";\n' 'squares [shape=circle];\n' '"Cube";\n' 'cubes [shape=circle];\n' '"Init" -> init_pcoll;\n' 'init_pcoll -> "Square";\n' 'init_pcoll -> "Cube";\n' '"Square" -> squares;\n' '"Cube" -> cubes;\n' '}\n'), pipeline_graph.PipelineGraph(p).get_dot()) @patch('IPython.get_ipython', new_callable=mock_get_ipython) def test_get_dot_within_notebook(self, cell): # Assume a mocked ipython kernel and notebook frontend have been set up. ie.current_env()._is_in_ipython = True ie.current_env()._is_in_notebook = True with cell: # Cell 1 p = beam.Pipeline(ir.InteractiveRunner()) # Immediately track this local pipeline so that ipython prompts when # applying transforms will be tracked and used for labels. ib.watch(locals()) with cell: # Cell 2 init_pcoll = p | 'Init' >> beam.Create(range(10)) with cell: # Cell 3 squares = init_pcoll | 'Square' >> beam.Map(lambda x: x * x) with cell: # Cell 4 cubes = init_pcoll | 'Cube' >> beam.Map(lambda x: x**3) # Tracks all PCollections defined so far. ib.watch(locals()) self.assertEqual(( 'digraph G {\n' 'node [color=blue, fontcolor=blue, shape=box];\n' '"[2]: Init";\n' 'init_pcoll [shape=circle];\n' '"[3]: Square";\n' 'squares [shape=circle];\n' '"[4]: Cube";\n' 'cubes [shape=circle];\n' '"[2]: Init" -> init_pcoll;\n' 'init_pcoll -> "[3]: Square";\n' 'init_pcoll -> "[4]: Cube";\n' '"[3]: Square" -> squares;\n' '"[4]: Cube" -> cubes;\n' '}\n'), pipeline_graph.PipelineGraph(p).get_dot()) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/interactive_beam.py0000644000175100017510000012533315206113032025550 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module of Interactive Beam features that can be used in notebook. The purpose of the module is to reduce the learning curve of Interactive Beam users, provide a single place for importing and add sugar syntax for all Interactive Beam components. It gives users capability to interact with existing environment/session/context for Interactive Beam and visualize PCollections as bounded dataset. In the meantime, it hides the interactivity implementation from users so that users can focus on developing Beam pipeline without worrying about how hidden states in the interactive session are managed. A convention to import this module: from apache_beam.runners.interactive import interactive_beam as ib Note: If you want backward-compatibility, only invoke interfaces provided by this module in your notebook or application code. """ # pytype: skip-file import logging from collections.abc import Iterable from datetime import timedelta from typing import Any from typing import Optional from typing import Union import pandas as pd import apache_beam as beam from apache_beam.dataframe.frame_base import DeferredBase from apache_beam.options.pipeline_options import FlinkRunnerOptions from apache_beam.pvalue import PCollection from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import DataprocClusterManager from apache_beam.runners.interactive.dataproc.types import ClusterIdentifier from apache_beam.runners.interactive.dataproc.types import ClusterMetadata from apache_beam.runners.interactive.display import pipeline_graph from apache_beam.runners.interactive.display.pcoll_visualization import visualize from apache_beam.runners.interactive.display.pcoll_visualization import visualize_computed_pcoll from apache_beam.runners.interactive.options import interactive_options from apache_beam.runners.interactive.recording_manager import AsyncComputationResult from apache_beam.runners.interactive.utils import deferred_df_to_pcollection from apache_beam.runners.interactive.utils import elements_to_df from apache_beam.runners.interactive.utils import find_pcoll_name from apache_beam.runners.interactive.utils import progress_indicated from apache_beam.runners.runner import PipelineState _LOGGER = logging.getLogger(__name__) class Options(interactive_options.InteractiveOptions): """Options that guide how Interactive Beam works.""" @property def enable_recording_replay(self): """Whether replayable source data recorded should be replayed for multiple PCollection evaluations and pipeline runs as long as the data recorded is still valid.""" return self.capture_control._enable_capture_replay @enable_recording_replay.setter def enable_recording_replay(self, value): """Sets whether source data recorded should be replayed. True - Enables recording of replayable source data so that following PCollection evaluations and pipeline runs always use the same data recorded; False - Disables recording of replayable source data so that following PCollection evaluation and pipeline runs always use new data from sources. """ # This makes sure the log handler is configured correctly in case the # options are configured in an early stage. _ = ie.current_env() if value: _LOGGER.info( 'Record replay is enabled. When a PCollection is evaluated or the ' 'pipeline is executed, existing data recorded from previous ' 'computations will be replayed for consistent results. If no ' 'recorded data is available, new data from recordable sources will ' 'be recorded.') else: _LOGGER.info( 'Record replay is disabled. The next time a PCollection is ' 'evaluated or the pipeline is executed, new data will always be ' 'consumed from sources in the pipeline. You will not have ' 'replayability until re-enabling this option.') self.capture_control._enable_capture_replay = value @property def recordable_sources(self): """Interactive Beam automatically records data from sources in this set. """ return self.capture_control._capturable_sources @property def recording_duration(self): """The data recording of sources ends as soon as the background source recording job has run for this long.""" return self.capture_control._capture_duration @recording_duration.setter def recording_duration(self, value): """Sets the recording duration as a timedelta. The input can be a datetime.timedelta, a possitive integer as seconds or a string representation that is parsable by pandas.to_timedelta. Example:: # Sets the recording duration limit to 10 seconds. ib.options.recording_duration = timedelta(seconds=10) ib.options.recording_duration = 10 ib.options.recording_duration = '10s' # Explicitly control the recordings. ib.recordings.stop(p) ib.recordings.clear(p) ib.recordings.record(p) # The next PCollection evaluation uses fresh data from sources, # and the data recorded will be replayed until another clear. ib.collect(some_pcoll) """ duration = None if isinstance(value, int): assert value > 0, 'Duration must be a positive value.' duration = timedelta(seconds=value) elif isinstance(value, str): duration = pd.to_timedelta(value) else: assert isinstance(value, timedelta), ('The input can only abe a ' 'datetime.timedelta, a possitive integer as seconds, or a string ' 'representation that is parsable by pandas.to_timedelta.') duration = value if self.capture_control._capture_duration.total_seconds( ) != duration.total_seconds(): _ = ie.current_env() _LOGGER.info( 'You have changed recording duration from %s seconds to %s seconds. ' 'To allow new data to be recorded for the updated duration the ' 'next time a PCollection is evaluated or the pipeline is executed, ' 'please invoke ib.recordings.stop, ib.recordings.clear and ' 'ib.recordings.record.', self.capture_control._capture_duration.total_seconds(), duration.total_seconds()) self.capture_control._capture_duration = duration @property def recording_size_limit(self): """The data recording of sources ends as soon as the size (in bytes) of data recorded from recordable sources reaches the limit.""" return self.capture_control._capture_size_limit @recording_size_limit.setter def recording_size_limit(self, value): """Sets the recording size in bytes. Example:: # Sets the recording size limit to 1GB. interactive_beam.options.recording_size_limit = 1e9 """ if self.capture_control._capture_size_limit != value: _ = ie.current_env() _LOGGER.info( 'You have changed recording size limit from %s bytes to %s bytes. To ' 'allow new data to be recorded under the updated size limit the ' 'next time a PCollection is recorded or the pipeline is executed, ' 'please invoke ib.recordings.stop, ib.recordings.clear and ' 'ib.recordings.record.', self.capture_control._capture_size_limit, value) self.capture_control._capture_size_limit = value @property def display_timestamp_format(self): """The format in which timestamps are displayed. Default is '%Y-%m-%d %H:%M:%S.%f%z', e.g. 2020-02-01 15:05:06.000015-08:00. """ return self._display_timestamp_format @display_timestamp_format.setter def display_timestamp_format(self, value): """Sets the format in which timestamps are displayed. Default is '%Y-%m-%d %H:%M:%S.%f%z', e.g. 2020-02-01 15:05:06.000015-08:00. Example:: # Sets the format to not display the timezone or microseconds. interactive_beam.options.display_timestamp_format = %Y-%m-%d %H:%M:%S' """ self._display_timestamp_format = value @property def display_timezone(self): """The timezone in which timestamps are displayed. Defaults to local timezone. """ return self._display_timezone @display_timezone.setter def display_timezone(self, value): """Sets the timezone (datetime.tzinfo) in which timestamps are displayed. Defaults to local timezone. Example:: # Imports the timezone library. from pytz import timezone # Will display all timestamps in the US/Eastern time zone. tz = timezone('US/Eastern') # You can also use dateutil.tz to get a timezone. tz = dateutil.tz.gettz('US/Eastern') interactive_beam.options.display_timezone = tz """ self._display_timezone = value @property def cache_root(self): """The cache directory specified by the user. Defaults to None. """ return self._cache_root @cache_root.setter def cache_root(self, value): """Sets the cache directory. Defaults to None. Example of local directory usage:: interactive_beam.options.cache_root = '/Users/username/my/cache/dir' Example of GCS directory usage:: interactive_beam.options.cache_root = 'gs://my-gcs-bucket/cache/dir' """ _LOGGER.warning( 'Interactive Beam has detected a set value for the cache_root ' 'option. Please note: existing cache managers will not have ' 'their current cache directory changed. The option must be ' 'set in Interactive Beam prior to the initialization of new ' 'pipelines to take effect. To apply changes to new pipelines, ' 'the kernel must be restarted or the pipeline creation codes ' 'must be re-executed. ') self._cache_root = value class Recordings(): """An introspection interface for recordings for pipelines. When a user materializes a PCollection onto disk (eg. ib.show) for a streaming pipeline, a background source recording job is started. This job pulls data from all defined unbounded sources for that PCollection's pipeline. The following methods allow for introspection into that background recording job. """ def describe( self, pipeline: Optional[beam.Pipeline] = None) -> dict[str, Any]: # noqa: F821 """Returns a description of all the recordings for the given pipeline. If no pipeline is given then this returns a dictionary of descriptions for all pipelines. """ # Create the RecordingManager if it doesn't already exist. if pipeline: ie.current_env().get_recording_manager(pipeline, create_if_absent=True) description = ie.current_env().describe_all_recordings() if pipeline: return description[pipeline] return description def clear(self, pipeline: beam.Pipeline) -> bool: """Clears all recordings of the given pipeline. Returns True if cleared.""" description = self.describe(pipeline) if (not PipelineState.is_terminal(description['state']) and description['state'] != PipelineState.STOPPED): _LOGGER.warning( 'Trying to clear a recording with a running pipeline. Did ' 'you forget to call ib.recordings.stop?') return False ie.current_env().cleanup(pipeline) return True def stop(self, pipeline: beam.Pipeline) -> None: """Stops the background source recording of the given pipeline.""" recording_manager = ie.current_env().get_recording_manager( pipeline, create_if_absent=True) recording_manager.cancel() def record(self, pipeline: beam.Pipeline) -> bool: """Starts a background source recording job for the given pipeline. Returns True if the recording job was started. """ description = self.describe(pipeline) if (not PipelineState.is_terminal(description['state']) and description['state'] != PipelineState.STOPPED): _LOGGER.warning( 'Trying to start a recording with a running pipeline. Did ' 'you forget to call ib.recordings.stop?') return False if description['size'] > 0: _LOGGER.warning( 'A recording already exists for this pipeline. To start a ' 'recording, make sure to call ib.recordings.clear first.') return False recording_manager = ie.current_env().get_recording_manager( pipeline, create_if_absent=True) return recording_manager.record_pipeline() class Clusters: """An interface to control clusters implicitly created and managed by the current interactive environment. This class is not needed and should not be used otherwise. Do not use it for clusters a user explicitly manages: e.g., if you have a Flink cluster running somewhere and provides the flink master when running a pipeline with the FlinkRunner, the cluster will not be tracked or managed by Beam. To reuse the same cluster for your pipelines, use the same pipeline options: e.g., a pipeline option with the same flink master if you are using FlinkRunner. This module is experimental. No backwards-compatibility guarantees. Interactive Beam automatically creates/reuses existing worker clusters to execute pipelines when it detects the need from configurations. Currently, the only supported cluster implementation is Flink running on Cloud Dataproc. To configure a pipeline to run on Cloud Dataproc with Flink, set the underlying runner of the InteractiveRunner to FlinkRunner and the pipeline options to indicate where on Cloud the FlinkRunner should be deployed to. An example to enable automatic Dataproc cluster creation/reuse:: options = PipelineOptions([ '--project=my-project', '--region=my-region', '--environment_type=DOCKER']) pipeline = beam.Pipeline(InteractiveRunner( underlying_runner=FlinkRunner()), options=options) Reuse a pipeline options in another pipeline would configure Interactive Beam to reuse the same Dataproc cluster implicitly managed by the current interactive environment. If a flink_master is identified as a known cluster, the corresponding cluster is also resued. Furthermore, if a cluster is explicitly created by using a pipeline as an identifier to a known cluster, the cluster is reused. An example:: # If pipeline runs on a known cluster, below code reuses the cluster # manager without creating a new one. dcm = ib.clusters.create(pipeline) To provision the cluster, use WorkerOptions. Supported configurations are:: 1. subnetwork 2. num_workers 3. machine_type To configure a pipeline to run on an existing FlinkRunner deployed elsewhere, set the flink_master explicitly so no cluster will be created/reused. An example pipeline options to skip automatic Dataproc cluster usage:: options = PipelineOptions([ '--flink_master=some.self.hosted.flink:port', '--environment_type=DOCKER']) To configure a pipeline to run on a local FlinkRunner, explicitly set the default cluster metadata to None: ib.clusters.set_default_cluster(None). """ # Explicitly set the Flink version here to ensure compatibility with 2.2 # Dataproc images: # https://cloud.google.com/dataproc/docs/concepts/versioning/dataproc-release-2.2 # you can manually override this by importing Clusters DATAPROC_FLINK_VERSION = '1.17' # The minimum worker number to create a Dataproc cluster. DATAPROC_MINIMUM_WORKER_NUM = 2 # TODO(https://github.com/apache/beam/issues/21527): Fix the Dataproc image # version after a released image contains all missing dependencies for Flink # to run. # DATAPROC_IMAGE_VERSION = '2.0.XX-debian10' def __init__(self) -> None: self.dataproc_cluster_managers: dict[ClusterMetadata, DataprocClusterManager] = {} self.master_urls: dict[str, ClusterMetadata] = {} self.pipelines: dict[beam.Pipeline, DataprocClusterManager] = {} self.default_cluster_metadata: Optional[ClusterMetadata] = None def create( self, cluster_identifier: ClusterIdentifier) -> DataprocClusterManager: """Creates a Dataproc cluster manager provisioned for the cluster identified. If the cluster is known, returns an existing cluster manager. """ # Try to get some not-None cluster metadata. cluster_metadata = self.cluster_metadata(cluster_identifier) if not cluster_metadata: raise ValueError( 'Unknown cluster identifier: %s. Cannot create or reuse' 'a Dataproc cluster.') if not cluster_metadata.region: _LOGGER.info( 'No region information was detected, defaulting Dataproc cluster ' 'region to: us-central1.') cluster_metadata.region = 'us-central1' elif cluster_metadata.region == 'global': # The global region is unsupported as it will be eventually deprecated. raise ValueError('Clusters in the global region are not supported.') # else use the provided region. if (cluster_metadata.num_workers and cluster_metadata.num_workers < self.DATAPROC_MINIMUM_WORKER_NUM): _LOGGER.info( 'At least %s workers are required for a cluster, defaulting to %s.', self.DATAPROC_MINIMUM_WORKER_NUM, self.DATAPROC_MINIMUM_WORKER_NUM) cluster_metadata.num_workers = self.DATAPROC_MINIMUM_WORKER_NUM known_dcm = self.dataproc_cluster_managers.get(cluster_metadata, None) if known_dcm: return known_dcm dcm = DataprocClusterManager(cluster_metadata) dcm.create_flink_cluster() # ClusterMetadata with derivative fields populated by the dcm. derived_meta = dcm.cluster_metadata self.dataproc_cluster_managers[derived_meta] = dcm self.master_urls[derived_meta.master_url] = derived_meta # Update the default cluster metadata to the one just created. self.set_default_cluster(derived_meta) return dcm def cleanup( self, cluster_identifier: Optional[ClusterIdentifier] = None, force: bool = False) -> None: """Cleans up the cluster associated with the given cluster_identifier. When None cluster_identifier is provided: if force is True, cleans up for all clusters; otherwise, do a dry run and NOOP. If a beam.Pipeline is given as the ClusterIdentifier while multiple pipelines share the same cluster, it only cleans up the association between the pipeline and the cluster identified. If the cluster_identifier is unknown, NOOP. """ if not cluster_identifier: dcm_to_cleanup = set(self.dataproc_cluster_managers.values()) if force: for dcm in dcm_to_cleanup: self._cleanup(dcm) self.default_cluster_metadata = None else: _LOGGER.warning( 'No cluster_identifier provided. If you intend to ' 'clean up all clusters, invoke ib.clusters.cleanup(force=True). ' 'Current clusters are %s.', self.describe()) elif isinstance(cluster_identifier, beam.Pipeline): p = cluster_identifier dcm = self.pipelines.pop(p, None) if dcm: dcm.pipelines.remove(p) p_flink_options = p.options.view_as(FlinkRunnerOptions) p_flink_options.flink_master = '[auto]' p_flink_options.flink_version = None # Only cleans up when there is no pipeline using the cluster. if not dcm.pipelines: self._cleanup(dcm) else: if isinstance(cluster_identifier, str): meta = self.master_urls.get(cluster_identifier, None) else: meta = cluster_identifier dcm = self.dataproc_cluster_managers.get(meta, None) if dcm: self._cleanup(dcm) def describe( self, cluster_identifier: Optional[ClusterIdentifier] = None ) -> Union[ClusterMetadata, list[ClusterMetadata]]: """Describes the ClusterMetadata by a ClusterIdentifier. If no cluster_identifier is given or if the cluster_identifier is unknown, it returns descriptions for all known clusters. Example usage: # Describe the cluster executing work for a pipeline. ib.clusters.describe(pipeline) # Describe the cluster with the flink master url. ib.clusters.describe(master_url) # Describe all existing clusters. ib.clusters.describe() """ if cluster_identifier: meta = self._cluster_metadata(cluster_identifier) if meta in self.dataproc_cluster_managers: return meta return list(self.dataproc_cluster_managers.keys()) def set_default_cluster( self, cluster_identifier: Optional[ClusterIdentifier] = None) -> None: """Temporarily sets the default metadata for creating or reusing a DataprocClusterManager. It is always updated to the most recently created cluster. If no known ClusterMetadata can be identified by the ClusterIdentifer, NOOP. If None is set, next time when Flink is in use, if no cluster is explicitly configured by a pipeline, the job runs locally. """ if cluster_identifier: self.default_cluster_metadata = self.cluster_metadata(cluster_identifier) else: self.default_cluster_metadata = None def cluster_metadata( self, cluster_identifier: Optional[ClusterIdentifier] = None ) -> Optional[ClusterMetadata]: """Fetches the ClusterMetadata by a ClusterIdentifier that could be a URL in string, a Beam pipeline, or an equivalent to a known ClusterMetadata; If the given cluster_identifier is an URL or a pipeline that is unknown to the current environment, the default cluster metadata (could be None) is returned. If the given cluster_identifier is a ClusterMetadata but unknown to the current environment, passes it through (NOOP). """ meta = self._cluster_metadata(cluster_identifier) return meta if meta else self.default_cluster_metadata def _cluster_metadata( self, cluster_identifier: Optional[ClusterIdentifier] = None ) -> Optional[ClusterMetadata]: meta = None if cluster_identifier: if isinstance(cluster_identifier, str): meta = self.master_urls.get(cluster_identifier, None) elif isinstance(cluster_identifier, beam.Pipeline): dcm = self.pipelines.get(cluster_identifier, None) if dcm: meta = dcm.cluster_metadata elif isinstance(cluster_identifier, ClusterMetadata): meta = cluster_identifier if meta in self.dataproc_cluster_managers: meta = self.dataproc_cluster_managers[meta].cluster_metadata elif (meta and self.default_cluster_metadata and meta.cluster_name == self.default_cluster_metadata.cluster_name): _LOGGER.warning( 'Cannot change the configuration of the running cluster %s. ' 'Existing is %s, desired is %s.', self.default_cluster_metadata.cluster_name, self.default_cluster_metadata, meta) meta.reset_name() _LOGGER.warning( 'To avoid conflict, issuing a new cluster name %s ' 'for a new cluster.', meta.cluster_name) else: raise TypeError( 'A cluster_identifier should be Optional[Union[str, ' 'beam.Pipeline, ClusterMetadata], instead %s was given.', type(cluster_identifier)) return meta def _cleanup(self, dcm: DataprocClusterManager) -> None: dcm.cleanup() self.dataproc_cluster_managers.pop(dcm.cluster_metadata, None) self.master_urls.pop(dcm.cluster_metadata.master_url, None) for p in dcm.pipelines: self.pipelines.pop(p, None) if dcm.cluster_metadata == self.default_cluster_metadata: self.default_cluster_metadata = None # Users can set options to guide how Interactive Beam works. # Examples: # ib.options.enable_recording_replay = False/True # ib.options.recording_duration = '1m' # ib.options.recordable_sources.add(SourceClass) # Check the docstrings for detailed usages. options = Options() # Users can introspect into recordings by using the recordings class. # Examples: # p = beam.Pipeline(InteractiveRunner()) # elems = p | beam.Create([1, 2, 3]) # ib.show(elems) # ib.recordings.describe(p) recordings = Recordings() # Users can interact with the clusters used by their environment. # Examples: # ib.clusters.describe(p) # Check the docstrings for detailed usages. clusters = Clusters() def watch(watchable): """Monitors a watchable. This allows Interactive Beam to implicitly pass on the information about the location of your pipeline definition. Current implementation mainly watches for PCollection variables defined in user code. A watchable can be a dictionary of variable metadata such as locals(), a str name of a module, a module object or an instance of a class. The variable can come from any scope even local variables in a method of a class defined in a module. Below are all valid:: watch(__main__) # if import __main__ is already invoked watch('__main__') # does not require invoking import __main__ beforehand watch(self) # inside a class watch(SomeInstance()) # an instance of a class watch(locals()) # inside a function, watching local variables within If you write a Beam pipeline in the __main__ module directly, since the __main__ module is always watched, you don't have to instruct Interactive Beam. If your Beam pipeline is defined in some module other than __main__, such as inside a class function or a unit test, you can watch() the scope. For example:: class Foo(object) def run_pipeline(self): with beam.Pipeline() as p: init_pcoll = p | 'Init Create' >> beam.Create(range(10)) watch(locals()) return init_pcoll init_pcoll = Foo().run_pipeline() Interactive Beam caches init_pcoll for the first run. Then you can use:: show(init_pcoll) To visualize data from init_pcoll once the pipeline is executed. """ ie.current_env().watch(watchable) @progress_indicated def show( *pcolls: Union[dict[Any, PCollection], Iterable[PCollection], PCollection], include_window_info: bool = False, visualize_data: bool = False, n: Union[int, str] = 'inf', duration: Union[int, str] = 'inf'): """Shows given PCollections in an interactive exploratory way if used within a notebook, or prints a heading sampled data if used within an ipython shell. Noop if used in a non-interactive environment. Args: include_window_info: (optional) if True, windowing information of the data will be visualized too. Default is false. visualize_data: (optional) by default, the visualization contains data tables rendering data from given pcolls separately as if they are converted into dataframes. If visualize_data is True, there will be a more dive-in widget and statistically overview widget of the data. Otherwise, those 2 data visualization widgets will not be displayed. n: (optional) max number of elements to visualize. Default 'inf'. duration: (optional) max duration of elements to read in integer seconds or a string duration. Default 'inf'. The given pcolls can be dictionary of PCollections (as values), or iterable of PCollections or plain PCollection values. The user can specify either the max number of elements with `n` to read or the maximum duration of elements to read with `duration`. When a limiter is not supplied, it is assumed to be infinite. By default, the visualization contains data tables rendering data from given pcolls separately as if they are converted into dataframes. If visualize_data is True, there will be a more dive-in widget and statistically overview widget of the data. Otherwise, those 2 data visualization widgets will not be displayed. Ad hoc builds a pipeline fragment including only transforms that are necessary to produce data for given PCollections pcolls, runs the pipeline fragment to compute data for those pcolls and then visualizes the data. The function is always blocking. If used within a notebook, the data visualized might be dynamically updated before the function returns as more and more data could getting processed and emitted when the pipeline fragment is being executed. If used within an ipython shell, there will be no dynamic plotting but a static plotting in the end of pipeline fragment execution. The PCollections given must belong to the same pipeline. For example:: p = beam.Pipeline(InteractiveRunner()) init = p | 'Init' >> beam.Create(range(1000)) square = init | 'Square' >> beam.Map(lambda x: x * x) cube = init | 'Cube' >> beam.Map(lambda x: x ** 3) # Below builds a pipeline fragment from the defined pipeline `p` that # contains only applied transforms of `Init` and `Square`. Then the # interactive runner runs the pipeline fragment implicitly to compute data # represented by PCollection `square` and visualizes it. show(square) # This is equivalent to `show(square)` because `square` depends on `init` # and `init` is included in the pipeline fragment and computed anyway. show(init, square) # Below is similar to running `p.run()`. It computes data for both # PCollection `square` and PCollection `cube`, then visualizes them. show(square, cube) """ flatten_pcolls = [] for pcoll_container in pcolls: if isinstance(pcoll_container, dict): flatten_pcolls.extend(pcoll_container.values()) elif isinstance(pcoll_container, (beam.pvalue.PCollection, DeferredBase)): flatten_pcolls.append(pcoll_container) else: try: flatten_pcolls.extend(iter(pcoll_container)) except TypeError: raise ValueError( 'The given pcoll %s is not a dict, an iterable or a PCollection.' % pcoll_container) # Iterate through the given PCollections and convert any deferred DataFrames # or Series into PCollections. pcolls = set() # The element type is used to help visualize the given PCollection. For the # deferred DataFrame/Series case it is the proxy of the frame. element_types = {} for pcoll in flatten_pcolls: if isinstance(pcoll, DeferredBase): pcoll, element_type = deferred_df_to_pcollection(pcoll) watch({'anonymous_pcollection_{}'.format(id(pcoll)): pcoll}) else: element_type = pcoll.element_type element_types[pcoll] = element_type pcolls.add(pcoll) assert isinstance(pcoll, beam.pvalue.PCollection), ( '{} is not an apache_beam.pvalue.PCollection.'.format(pcoll)) assert len(pcolls) > 0, ( 'Need at least 1 PCollection to show data visualization.') pcoll_pipeline = next(iter(pcolls)).pipeline user_pipeline = ie.current_env().user_pipeline(pcoll_pipeline) # Possibly showing a PCollection defined in a local scope that is not # explicitly watched. Ad hoc watch it though it's a little late. if not user_pipeline: watch({'anonymous_pipeline_{}'.format(id(pcoll_pipeline)): pcoll_pipeline}) user_pipeline = pcoll_pipeline if isinstance(n, str): assert n == 'inf', ( 'Currently only the string \'inf\' is supported. This denotes reading ' 'elements until the recording is stopped via a kernel interrupt.') elif isinstance(n, int): assert n > 0, 'n needs to be positive or the string \'inf\'' if isinstance(duration, int): assert duration > 0, ('duration needs to be positive, a duration string, ' 'or the string \'inf\'') if n == 'inf': n = float('inf') if duration == 'inf': duration = float('inf') previously_computed_pcolls = { pcoll for pcoll in pcolls if pcoll in ie.current_env().computed_pcollections } for pcoll in previously_computed_pcolls: visualize_computed_pcoll( find_pcoll_name(pcoll), pcoll, n, duration, include_window_info=include_window_info, display_facets=visualize_data) pcolls = pcolls - previously_computed_pcolls recording_manager = ie.current_env().get_recording_manager( user_pipeline, create_if_absent=True) recording = recording_manager.record(pcolls, max_n=n, max_duration=duration) # Catch a KeyboardInterrupt to gracefully cancel the recording and # visualizations. try: # If in notebook, static plotting computed pcolls as computation is done. if ie.current_env().is_in_notebook: for stream in recording.computed().values(): visualize( stream, include_window_info=include_window_info, display_facets=visualize_data, element_type=element_types[stream.pcoll]) elif ie.current_env().is_in_ipython: for stream in recording.computed().values(): visualize( stream, include_window_info=include_window_info, element_type=element_types[stream.pcoll]) if recording.is_computed(): return # If in notebook, dynamic plotting as computation goes. if ie.current_env().is_in_notebook: for stream in recording.uncomputed().values(): visualize( stream, dynamic_plotting_interval=1, include_window_info=include_window_info, display_facets=visualize_data, element_type=element_types[stream.pcoll]) # Invoke wait_until_finish to ensure the blocking nature of this API without # relying on the run to be blocking. recording.wait_until_finish() # If just in ipython shell, plotting once when the computation is completed. if ie.current_env().is_in_ipython and not ie.current_env().is_in_notebook: for stream in recording.computed().values(): visualize(stream, include_window_info=include_window_info) except KeyboardInterrupt: if recording: recording.cancel() @progress_indicated def collect( *pcolls, n='inf', duration='inf', include_window_info=False, runner=None, options=None, force_compute=False, force_tuple=False, raw_records=False): """Materializes the elements from a PCollection into a Dataframe. This reads each element from file and reads only the amount that it needs into memory. The user can specify either the max number of elements to read or the maximum duration of elements to read. When a limiter is not supplied, it is assumed to be infinite. Args: pcolls: PCollections to compute. n: (optional) max number of elements to visualize. Default 'inf'. duration: (optional) max duration of elements to read in integer seconds or a string duration. Default 'inf'. include_window_info: (optional) if True, appends the windowing information to each row. Default False. runner: (optional) the runner with which to compute the results options: (optional) any additional pipeline options to use to compute the results force_compute: (optional) if True, forces recomputation rather than using cached PCollections force_tuple: (optional) if True, return a 1-tuple or results rather than the bare results if only one PCollection is computed raw_records: (optional) if True, return a list of collected records without converting to a DataFrame. Default False. For example:: p = beam.Pipeline(InteractiveRunner()) init = p | 'Init' >> beam.Create(range(10)) square = init | 'Square' >> beam.Map(lambda x: x * x) # Run the pipeline and bring the PCollection into memory as a Dataframe. in_memory_square = head(square, n=5) # Run the pipeline and get the raw list of elements. raw_squares = collect(square, n=5, raw_records=True) """ if len(pcolls) == 0: return () def as_pcollection(pcoll_or_df): if isinstance(pcoll_or_df, DeferredBase): # Get the proxy so we can get the output shape of the DataFrame. pcoll, element_type = deferred_df_to_pcollection(pcoll_or_df) watch({'anonymous_pcollection_{}'.format(id(pcoll)): pcoll}) return pcoll, element_type elif isinstance(pcoll_or_df, beam.pvalue.PCollection): return pcoll_or_df, pcoll_or_df.element_type else: raise TypeError(f'{pcoll} is not an apache_beam.pvalue.PCollection.') pcolls_with_element_types = [as_pcollection(p) for p in pcolls] pcolls_to_element_types = dict(pcolls_with_element_types) pcolls = [pcoll for pcoll, _ in pcolls_with_element_types] pipelines = set(pcoll.pipeline for pcoll in pcolls) if len(pipelines) != 1: raise ValueError('All PCollections must belong to the same pipeline.') pipeline, = pipelines if isinstance(n, str): assert n == 'inf', ( 'Currently only the string \'inf\' is supported. This denotes reading ' 'elements until the recording is stopped via a kernel interrupt.') elif isinstance(n, int): assert n > 0, 'n needs to be positive or the string \'inf\'' if isinstance(duration, int): assert duration > 0, ('duration needs to be positive, a duration string, ' 'or the string \'inf\'') if n == 'inf': n = float('inf') if duration == 'inf': duration = float('inf') user_pipeline = ie.current_env().user_pipeline(pipeline) # Possibly collecting a PCollection defined in a local scope that is not # explicitly watched. Ad hoc watch it though it's a little late. if not user_pipeline: watch({'anonymous_pipeline_{}'.format(id(pipeline)): pipeline}) user_pipeline = pipeline recording_manager = ie.current_env().get_recording_manager( user_pipeline, create_if_absent=True) # If already computed, directly read the stream and return. computed = {} for pcoll in pcolls_to_element_types.keys(): if pcoll in ie.current_env().computed_pcollections and not force_compute: pcoll_name = find_pcoll_name(pcoll) computed[pcoll] = list( recording_manager.read(pcoll_name, pcoll, n, duration).read()) uncomputed = set(pcolls) - set(computed.keys()) if uncomputed: recording = recording_manager.record( uncomputed, max_n=n, max_duration=duration, runner=runner, options=options, force_compute=force_compute) try: for pcoll in uncomputed: computed[pcoll] = list(recording.stream(pcoll).read()) except KeyboardInterrupt: recording.cancel() if n == float('inf'): n = None if raw_records: result_tuple = tuple([el.value for el in computed.get(pcoll, [])][:n] for pcoll in pcolls) else: # Collecting DataFrames may have a length > n, so slice again to be sure. # Note that array[:None] returns everything. empty = pd.DataFrame() result_tuple = tuple( elements_to_df( computed.get(pcoll, []), include_window_info=include_window_info, element_type=pcolls_to_element_types[pcoll])[:n] if pcoll in computed else empty for pcoll in pcolls) if len(result_tuple) == 1 and not force_tuple: return result_tuple[0] else: return result_tuple @progress_indicated def compute( *pcolls: Union[dict[Any, PCollection], Iterable[PCollection], PCollection], wait_for_inputs: bool = True, blocking: bool = False, runner=None, options=None, force_compute=False, ) -> Optional[AsyncComputationResult]: """Computes the given PCollections, potentially asynchronously. Args: *pcolls: PCollections to compute. Can be a single PCollection, an iterable of PCollections, or a dictionary with PCollections as values. wait_for_inputs: Whether to wait until the asynchronous dependencies are computed. Setting this to False allows to immediately schedule the computation, but also potentially results in running the same pipeline stages multiple times. blocking: If False, the computation will run in non-blocking fashion. In Colab/IPython environment this mode will also provide the controls for the running pipeline. If True, the computation will block until the pipeline is done. runner: (optional) the runner with which to compute the results. options: (optional) any additional pipeline options to use to compute the results. force_compute: (optional) if True, forces recomputation rather than using cached PCollections. Returns: An AsyncComputationResult object if blocking is False, otherwise None. """ flatten_pcolls = [] for pcoll_container in pcolls: if isinstance(pcoll_container, dict): flatten_pcolls.extend(pcoll_container.values()) elif isinstance(pcoll_container, (beam.pvalue.PCollection, DeferredBase)): flatten_pcolls.append(pcoll_container) else: try: flatten_pcolls.extend(iter(pcoll_container)) except TypeError: raise ValueError( f'The given pcoll {pcoll_container} is not a dict, an iterable or ' 'a PCollection.') pcolls_set = set() for pcoll in flatten_pcolls: if isinstance(pcoll, DeferredBase): pcoll, _ = deferred_df_to_pcollection(pcoll) watch({f'anonymous_pcollection_{id(pcoll)}': pcoll}) assert isinstance( pcoll, beam.pvalue.PCollection ), f'{pcoll} is not an apache_beam.pvalue.PCollection.' pcolls_set.add(pcoll) if not pcolls_set: _LOGGER.info('No PCollections to compute.') return None pcoll_pipeline = next(iter(pcolls_set)).pipeline user_pipeline = ie.current_env().user_pipeline(pcoll_pipeline) if not user_pipeline: watch({f'anonymous_pipeline_{id(pcoll_pipeline)}': pcoll_pipeline}) user_pipeline = pcoll_pipeline for pcoll in pcolls_set: if pcoll.pipeline is not user_pipeline: raise ValueError('All PCollections must belong to the same pipeline.') recording_manager = ie.current_env().get_recording_manager( user_pipeline, create_if_absent=True) return recording_manager.compute_async( pcolls_set, wait_for_inputs=wait_for_inputs, blocking=blocking, runner=runner, options=options, force_compute=force_compute, ) @progress_indicated def show_graph(pipeline): """Shows the current pipeline shape of a given Beam pipeline as a DAG. """ pipeline_graph.PipelineGraph(pipeline).display_graph() def evict_recorded_data(pipeline=None): """Forcefully evicts all recorded replayable data for the given pipeline. If no pipeline is specified, evicts for all user defined pipelines. Once invoked, Interactive Beam will record new data based on the guidance of options the next time it evaluates/visualizes PCollections or runs pipelines. """ from apache_beam.runners.interactive.options import capture_control capture_control.evict_captured_data(pipeline) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/interactive_beam_test.py0000644000175100017510000011603315206113032026604 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.runners.interactive.interactive_beam.""" # pytype: skip-file import dataclasses import importlib import sys import time import unittest from concurrent.futures import TimeoutError from typing import NamedTuple from unittest.mock import ANY from unittest.mock import MagicMock from unittest.mock import call from unittest.mock import patch import apache_beam as beam from apache_beam import dataframe as frames from apache_beam.dataframe.frame_base import DeferredBase from apache_beam.options.pipeline_options import FlinkRunnerOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import interactive_runner as ir from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import DataprocClusterManager from apache_beam.runners.interactive.dataproc.types import ClusterMetadata from apache_beam.runners.interactive.options.capture_limiters import Limiter from apache_beam.runners.interactive.recording_manager import AsyncComputationResult from apache_beam.runners.interactive.testing.mock_env import isolated_env from apache_beam.runners.runner import PipelineState from apache_beam.testing.test_stream import TestStream @dataclasses.dataclass class MockClusterMetadata: master_url = 'mock_url' class Record(NamedTuple): order_id: int product_id: int quantity: int # The module name is also a variable in module. _module_name = 'apache_beam.runners.interactive.interactive_beam_test' def _get_watched_pcollections_with_variable_names(): watched_pcollections = {} for watching in ie.current_env().watching(): for key, val in watching: if hasattr(val, '__class__') and isinstance(val, beam.pvalue.PCollection): watched_pcollections[val] = key return watched_pcollections @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') @isolated_env class InteractiveBeamTest(unittest.TestCase): def setUp(self): self._var_in_class_instance = 'a var in class instance, not directly used' def tearDown(self): ib.options.capture_control.set_limiters_for_test([]) def test_watch_main_by_default(self): test_env = ie.InteractiveEnvironment() # Current Interactive Beam env fetched and the test env are 2 instances. self.assertNotEqual(id(ie.current_env()), id(test_env)) self.assertEqual(ie.current_env().watching(), test_env.watching()) def test_watch_a_module_by_name(self): test_env = ie.InteractiveEnvironment() ib.watch(_module_name) test_env.watch(_module_name) self.assertEqual(ie.current_env().watching(), test_env.watching()) def test_watch_a_module_by_module_object(self): test_env = ie.InteractiveEnvironment() module = importlib.import_module(_module_name) ib.watch(module) test_env.watch(module) self.assertEqual(ie.current_env().watching(), test_env.watching()) def test_watch_locals(self): # test_env serves as local var too. test_env = ie.InteractiveEnvironment() ib.watch(locals()) test_env.watch(locals()) self.assertEqual(ie.current_env().watching(), test_env.watching()) def test_watch_class_instance(self): test_env = ie.InteractiveEnvironment() ib.watch(self) test_env.watch(self) self.assertEqual(ie.current_env().watching(), test_env.watching()) @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_show_always_watch_given_pcolls(self): p = beam.Pipeline(ir.InteractiveRunner()) # pylint: disable=bad-option-value pcoll = p | 'Create' >> beam.Create(range(10)) # The pcoll is not watched since watch(locals()) is not explicitly called. self.assertFalse(pcoll in _get_watched_pcollections_with_variable_names()) # The call of show watches pcoll. ib.watch({'p': p}) ie.current_env().track_user_pipelines() ib.show(pcoll) self.assertTrue(pcoll in _get_watched_pcollections_with_variable_names()) @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_show_mark_pcolls_computed_when_done(self): p = beam.Pipeline(ir.InteractiveRunner()) # pylint: disable=bad-option-value pcoll = p | 'Create' >> beam.Create(range(10)) self.assertFalse(pcoll in ie.current_env().computed_pcollections) # The call of show marks pcoll computed. ib.watch(locals()) ie.current_env().track_user_pipelines() ib.show(pcoll) self.assertTrue(pcoll in ie.current_env().computed_pcollections) @patch(( 'apache_beam.runners.interactive.interactive_beam.' 'visualize_computed_pcoll')) def test_show_handles_dict_of_pcolls(self, mocked_visualize): p = beam.Pipeline(ir.InteractiveRunner()) # pylint: disable=bad-option-value pcoll = p | 'Create' >> beam.Create(range(10)) ib.watch(locals()) ie.current_env().track_user_pipelines() ie.current_env().mark_pcollection_computed([pcoll]) ie.current_env()._is_in_ipython = True ie.current_env()._is_in_notebook = True ib.show({'pcoll': pcoll}) mocked_visualize.assert_called_once() @patch(( 'apache_beam.runners.interactive.interactive_beam.' 'visualize_computed_pcoll')) def test_show_handles_iterable_of_pcolls(self, mocked_visualize): p = beam.Pipeline(ir.InteractiveRunner()) # pylint: disable=bad-option-value pcoll = p | 'Create' >> beam.Create(range(10)) ib.watch(locals()) ie.current_env().track_user_pipelines() ie.current_env().mark_pcollection_computed([pcoll]) ie.current_env()._is_in_ipython = True ie.current_env()._is_in_notebook = True ib.show([pcoll]) mocked_visualize.assert_called_once() @patch('apache_beam.runners.interactive.interactive_beam.visualize') def test_show_handles_deferred_dataframes(self, mocked_visualize): p = beam.Pipeline(ir.InteractiveRunner()) deferred = frames.convert.to_dataframe(p | beam.Create([Record(0, 0, 0)])) ib.watch(locals()) ie.current_env().track_user_pipelines() ie.current_env()._is_in_ipython = True ie.current_env()._is_in_notebook = True ib.show(deferred) mocked_visualize.assert_called_once() @patch(( 'apache_beam.runners.interactive.interactive_beam.' 'visualize_computed_pcoll')) def test_show_noop_when_pcoll_container_is_invalid(self, mocked_visualize): class SomeRandomClass: def __init__(self, pcoll): self._pcoll = pcoll p = beam.Pipeline(ir.InteractiveRunner()) # pylint: disable=bad-option-value pcoll = p | 'Create' >> beam.Create(range(10)) ie.current_env().mark_pcollection_computed([pcoll]) ie.current_env()._is_in_ipython = True ie.current_env()._is_in_notebook = True self.assertRaises(ValueError, ib.show, SomeRandomClass(pcoll)) mocked_visualize.assert_not_called() def test_recordings_describe(self): """Tests that getting the description works.""" # Create the pipelines to test. p1 = beam.Pipeline(ir.InteractiveRunner()) p2 = beam.Pipeline(ir.InteractiveRunner()) ib.watch(locals()) # Get the descriptions. This test is simple as there isn't much logic in the # method. self.assertEqual(ib.recordings.describe(p1)['size'], 0) self.assertEqual(ib.recordings.describe(p2)['size'], 0) all_descriptions = ib.recordings.describe() self.assertEqual(all_descriptions[p1]['size'], 0) self.assertEqual(all_descriptions[p2]['size'], 0) # Ensure that the variable name for the pipeline is set correctly. self.assertEqual(all_descriptions[p1]['pipeline_var'], 'p1') self.assertEqual(all_descriptions[p2]['pipeline_var'], 'p2') def test_recordings_clear(self): """Tests that clearing the pipeline is correctly forwarded.""" # Create a basic pipeline to store something in the cache. p = beam.Pipeline(ir.InteractiveRunner()) elem = p | beam.Create([1]) ib.watch(locals()) ie.current_env().track_user_pipelines() # This records the pipeline so that the cache size is > 0. ib.collect(elem) self.assertGreater(ib.recordings.describe(p)['size'], 0) # After clearing, the cache should be empty. ib.recordings.clear(p) self.assertEqual(ib.recordings.describe(p)['size'], 0) def test_recordings_record(self): """Tests that recording pipeline succeeds.""" # Add the TestStream so that it can be cached. ib.options.recordable_sources.add(TestStream) # Create a pipeline with an arbitrary amonunt of elements. p = beam.Pipeline( ir.InteractiveRunner(), options=PipelineOptions(streaming=True)) # pylint: disable=unused-variable _ = (p | TestStream() .advance_watermark_to(0) .advance_processing_time(1) .add_elements(list(range(10))) .advance_processing_time(1)) # yapf: disable ib.watch(locals()) ie.current_env().track_user_pipelines() # Assert that the pipeline starts in a good state. self.assertEqual(ib.recordings.describe(p)['state'], PipelineState.STOPPED) self.assertEqual(ib.recordings.describe(p)['size'], 0) # Create a lmiter that stops the background caching job when something is # written to cache. This is used to make ensure that the pipeline is # functioning properly and that there are no data races with the test. class SizeLimiter(Limiter): def __init__(self, pipeline): self.pipeline = pipeline self.should_trigger = False def is_triggered(self): return ( ib.recordings.describe(self.pipeline)['size'] > 0 and self.should_trigger) limiter = SizeLimiter(p) ib.options.capture_control.set_limiters_for_test([limiter]) # Assert that a recording can be started only once. self.assertTrue(ib.recordings.record(p)) self.assertFalse(ib.recordings.record(p)) self.assertEqual(ib.recordings.describe(p)['state'], PipelineState.RUNNING) # Wait for the pipeline to start and write something to cache. limiter.should_trigger = True for _ in range(60): if limiter.is_triggered(): break time.sleep(1) self.assertTrue( limiter.is_triggered(), 'Test timed out waiting for limiter to be triggered. This indicates ' 'that the BackgroundCachingJob did not cache anything.') # Assert that a recording can be stopped and can't be started again until # after the cache is cleared. ib.recordings.stop(p) self.assertEqual(ib.recordings.describe(p)['state'], PipelineState.STOPPED) self.assertFalse(ib.recordings.record(p)) ib.recordings.clear(p) self.assertTrue(ib.recordings.record(p)) ib.recordings.stop(p) def test_collect_raw_records_true(self): p = beam.Pipeline(ir.InteractiveRunner()) data = list(range(5)) pcoll = p | 'Create' >> beam.Create(data) ib.watch(locals()) ie.current_env().track_user_pipelines() result = ib.collect(pcoll, raw_records=True) self.assertIsInstance(result, list) self.assertEqual(result, data) result_n = ib.collect(pcoll, n=3, raw_records=True) self.assertIsInstance(result_n, list) self.assertEqual(result_n, data[:3]) def test_collect_raw_records_false(self): p = beam.Pipeline(ir.InteractiveRunner()) data = list(range(5)) pcoll = p | 'Create' >> beam.Create(data) ib.watch(locals()) ie.current_env().track_user_pipelines() result = ib.collect(pcoll) self.assertNotIsInstance(result, list) self.assertTrue( hasattr(result, 'columns'), "Result should have 'columns' attribute") self.assertTrue( hasattr(result, 'values'), "Result should have 'values' attribute") result_n = ib.collect(pcoll, n=3) self.assertNotIsInstance(result_n, list) self.assertTrue( hasattr(result_n, 'columns'), "Result (n=3) should have 'columns' attribute") self.assertTrue( hasattr(result_n, 'values'), "Result (n=3) should have 'values' attribute") def test_collect_raw_records_true_multiple_pcolls(self): p = beam.Pipeline(ir.InteractiveRunner()) data1 = list(range(3)) data2 = [x * x for x in range(3)] pcoll1 = p | 'Create1' >> beam.Create(data1) pcoll2 = p | 'Create2' >> beam.Create(data2) ib.watch(locals()) ie.current_env().track_user_pipelines() result = ib.collect(pcoll1, pcoll2, raw_records=True) self.assertIsInstance(result, tuple) self.assertEqual(len(result), 2) self.assertIsInstance(result[0], list) self.assertEqual(result[0], data1) self.assertIsInstance(result[1], list) self.assertEqual(result[1], data2) def test_collect_raw_records_false_multiple_pcolls(self): p = beam.Pipeline(ir.InteractiveRunner()) data1 = list(range(3)) data2 = [x * x for x in range(3)] pcoll1 = p | 'Create1' >> beam.Create(data1) pcoll2 = p | 'Create2' >> beam.Create(data2) ib.watch(locals()) ie.current_env().track_user_pipelines() result = ib.collect(pcoll1, pcoll2) self.assertIsInstance(result, tuple) self.assertEqual(len(result), 2) self.assertNotIsInstance(result[0], list) self.assertTrue(hasattr(result[0], 'columns')) self.assertNotIsInstance(result[1], list) self.assertTrue(hasattr(result[1], 'columns')) def test_collect_raw_records_true_force_tuple(self): p = beam.Pipeline(ir.InteractiveRunner()) data = list(range(5)) pcoll = p | 'Create' >> beam.Create(data) ib.watch(locals()) ie.current_env().track_user_pipelines() result = ib.collect(pcoll, raw_records=True, force_tuple=True) self.assertIsInstance(result, tuple) self.assertEqual(len(result), 1) self.assertIsInstance(result[0], list) self.assertEqual(result[0], data) @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') @isolated_env class InteractiveBeamClustersTest(unittest.TestCase): def setUp(self): self.current_env.options.cache_root = 'gs://fake' self.clusters = self.current_env.clusters def tearDown(self): self.current_env.options.cache_root = None def test_cluster_metadata_pass_through_metadata(self): cid = ClusterMetadata(project_id='test-project') meta = self.clusters.cluster_metadata(cid) self.assertIs(meta, cid) def test_cluster_metadata_identifies_pipeline(self): cid = beam.Pipeline() known_meta = ClusterMetadata(project_id='test-project') dcm = DataprocClusterManager(known_meta) self.clusters.pipelines[cid] = dcm meta = self.clusters.cluster_metadata(cid) self.assertIs(meta, known_meta) def test_cluster_metadata_identifies_master_url(self): cid = 'test-url' known_meta = ClusterMetadata(project_id='test-project') _ = DataprocClusterManager(known_meta) self.clusters.master_urls[cid] = known_meta meta = self.clusters.cluster_metadata(cid) self.assertIs(meta, known_meta) def test_cluster_metadata_default_value(self): cid_none = None cid_unknown_p = beam.Pipeline() cid_unknown_master_url = 'test-url' default_meta = ClusterMetadata(project_id='test-project') self.clusters.set_default_cluster(default_meta) self.assertIs(default_meta, self.clusters.cluster_metadata(cid_none)) self.assertIs(default_meta, self.clusters.cluster_metadata(cid_unknown_p)) self.assertIs( default_meta, self.clusters.cluster_metadata(cid_unknown_master_url)) def test_create_a_new_cluster(self): meta = ClusterMetadata(project_id='test-project') _ = self.clusters.create(meta) # Derived fields are populated. self.assertTrue(meta.master_url.startswith('test-url')) self.assertEqual(meta.dashboard, 'test-dashboard') # The cluster is known. self.assertIn(meta, self.clusters.dataproc_cluster_managers) self.assertIn(meta.master_url, self.clusters.master_urls) # The default cluster is updated to the created cluster. self.assertIs(meta, self.clusters.default_cluster_metadata) def test_create_but_reuse_a_known_cluster(self): known_meta = ClusterMetadata( project_id='test-project', region='test-region') known_dcm = DataprocClusterManager(known_meta) known_meta.master_url = 'test-url' self.clusters.set_default_cluster(known_meta) self.clusters.dataproc_cluster_managers[known_meta] = known_dcm self.clusters.master_urls[known_meta.master_url] = known_meta # Use an equivalent meta as the identifier to create a cluster. cid_meta = ClusterMetadata( project_id=known_meta.project_id, region=known_meta.region, cluster_name=known_meta.cluster_name) dcm = self.clusters.create(cid_meta) # The known cluster manager is returned. self.assertIs(dcm, known_dcm) # Then use an equivalent master_url as the identifier. cid_master_url = known_meta.master_url dcm = self.clusters.create(cid_master_url) self.assertIs(dcm, known_dcm) def test_cleanup_by_a_pipeline(self): meta = ClusterMetadata(project_id='test-project') dcm = self.clusters.create(meta) # Set up the association between a pipeline and a cluster. # In real code, it's set by the runner the 1st time a pipeline is executed. options = PipelineOptions() options.view_as(FlinkRunnerOptions).flink_master = meta.master_url p = beam.Pipeline(options=options) self.clusters.pipelines[p] = dcm dcm.pipelines.add(p) self.clusters.cleanup(p) # Delete the cluster. self.m_delete_cluster.assert_called_once() # Pipeline association is cleaned up. self.assertNotIn(p, self.clusters.pipelines) self.assertNotIn(p, dcm.pipelines) # The internal option in the pipeline is overwritten. self.assertEqual( p.options.view_as(FlinkRunnerOptions).flink_master, '[auto]') # The original option is unchanged. self.assertEqual( options.view_as(FlinkRunnerOptions).flink_master, meta.master_url) # The cluster is unknown now. self.assertNotIn(meta, self.clusters.dataproc_cluster_managers) self.assertNotIn(meta.master_url, self.clusters.master_urls) # The cleaned up cluster is also the default cluster. Clean the default. self.assertIsNone(self.clusters.default_cluster_metadata) def test_not_cleanup_if_multiple_pipelines_share_a_manager(self): meta = ClusterMetadata(project_id='test-project') dcm = self.clusters.create(meta) options = PipelineOptions() options.view_as(FlinkRunnerOptions).flink_master = meta.master_url options2 = PipelineOptions() options2.view_as(FlinkRunnerOptions).flink_master = meta.master_url p = beam.Pipeline(options=options) p2 = beam.Pipeline(options=options2) self.clusters.pipelines[p] = dcm self.clusters.pipelines[p2] = dcm dcm.pipelines.add(p) dcm.pipelines.add(p2) self.clusters.cleanup(p) # No cluster deleted. self.m_delete_cluster.assert_not_called() # Pipeline association of p is cleaned up. self.assertNotIn(p, self.clusters.pipelines) self.assertNotIn(p, dcm.pipelines) # The internal option in the pipeline is overwritten. self.assertEqual( p.options.view_as(FlinkRunnerOptions).flink_master, '[auto]') # The original option is unchanged. self.assertEqual( options.view_as(FlinkRunnerOptions).flink_master, meta.master_url) # Pipeline association of p2 still presents. self.assertIn(p2, self.clusters.pipelines) self.assertIn(p2, dcm.pipelines) self.assertEqual( p2.options.view_as(FlinkRunnerOptions).flink_master, meta.master_url) self.assertEqual( options2.view_as(FlinkRunnerOptions).flink_master, meta.master_url) # The cluster is still known. self.assertIn(meta, self.clusters.dataproc_cluster_managers) self.assertIn(meta.master_url, self.clusters.master_urls) # The default cluster still presents. self.assertIs(meta, self.clusters.default_cluster_metadata) def test_cleanup_by_a_master_url(self): meta = ClusterMetadata(project_id='test-project') _ = self.clusters.create(meta) self.clusters.cleanup(meta.master_url) self.m_delete_cluster.assert_called_once() self.assertNotIn(meta, self.clusters.dataproc_cluster_managers) self.assertNotIn(meta.master_url, self.clusters.master_urls) self.assertIsNone(self.clusters.default_cluster_metadata) def test_cleanup_by_meta(self): known_meta = ClusterMetadata( project_id='test-project', region='test-region') _ = self.clusters.create(known_meta) meta = ClusterMetadata( project_id=known_meta.project_id, region=known_meta.region, cluster_name=known_meta.cluster_name) self.clusters.cleanup(meta) self.m_delete_cluster.assert_called_once() self.assertNotIn(known_meta, self.clusters.dataproc_cluster_managers) self.assertNotIn(known_meta.master_url, self.clusters.master_urls) self.assertIsNone(self.clusters.default_cluster_metadata) def test_force_cleanup_everything(self): meta = ClusterMetadata(project_id='test-project') meta2 = ClusterMetadata(project_id='test-project-2') _ = self.clusters.create(meta) _ = self.clusters.create(meta2) self.clusters.cleanup(force=True) self.assertEqual(self.m_delete_cluster.call_count, 2) self.assertNotIn(meta, self.clusters.dataproc_cluster_managers) self.assertNotIn(meta2, self.clusters.dataproc_cluster_managers) self.assertIsNone(self.clusters.default_cluster_metadata) def test_cleanup_noop_for_no_cluster_identifier(self): meta = ClusterMetadata(project_id='test-project') _ = self.clusters.create(meta) self.clusters.cleanup() self.m_delete_cluster.assert_not_called() def test_cleanup_noop_unknown_cluster(self): meta = ClusterMetadata(project_id='test-project') dcm = self.clusters.create(meta) p = beam.Pipeline() self.clusters.pipelines[p] = dcm dcm.pipelines.add(p) cid_pipeline = beam.Pipeline() self.clusters.cleanup(cid_pipeline) self.m_delete_cluster.assert_not_called() cid_master_url = 'some-random-url' self.clusters.cleanup(cid_master_url) self.m_delete_cluster.assert_not_called() cid_meta = ClusterMetadata(project_id='random-project') self.clusters.cleanup(cid_meta) self.m_delete_cluster.assert_not_called() self.assertIn(meta, self.clusters.dataproc_cluster_managers) self.assertIn(meta.master_url, self.clusters.master_urls) self.assertIs(meta, self.clusters.default_cluster_metadata) self.assertIn(p, self.clusters.pipelines) self.assertIn(p, dcm.pipelines) def test_describe_everything(self): meta = ClusterMetadata(project_id='test-project') meta2 = ClusterMetadata( project_id='test-project', region='some-other-region') _ = self.clusters.create(meta) _ = self.clusters.create(meta2) meta_list = self.clusters.describe() self.assertEqual([meta, meta2], meta_list) def test_describe_by_cluster_identifier(self): known_meta = ClusterMetadata(project_id='test-project') known_meta2 = ClusterMetadata( project_id='test-project', region='some-other-region') dcm = self.clusters.create(known_meta) dcm2 = self.clusters.create(known_meta2) p = beam.Pipeline() p2 = beam.Pipeline() self.clusters.pipelines[p] = dcm dcm.pipelines.add(p) self.clusters.pipelines[p2] = dcm2 dcm.pipelines.add(p2) cid_pipeline = p meta = self.clusters.describe(cid_pipeline) self.assertIs(meta, known_meta) cid_master_url = known_meta.master_url meta = self.clusters.describe(cid_master_url) self.assertIs(meta, known_meta) cid_meta = ClusterMetadata( project_id=known_meta.project_id, region=known_meta.region, cluster_name=known_meta.cluster_name) meta = self.clusters.describe(cid_meta) self.assertIs(meta, known_meta) def test_describe_everything_when_cluster_identifer_unknown(self): known_meta = ClusterMetadata(project_id='test-project') known_meta2 = ClusterMetadata( project_id='test-project', region='some-other-region') dcm = self.clusters.create(known_meta) dcm2 = self.clusters.create(known_meta2) p = beam.Pipeline() p2 = beam.Pipeline() self.clusters.pipelines[p] = dcm dcm.pipelines.add(p) self.clusters.pipelines[p2] = dcm2 dcm.pipelines.add(p2) cid_pipeline = beam.Pipeline() meta_list = self.clusters.describe(cid_pipeline) self.assertEqual([known_meta, known_meta2], meta_list) cid_master_url = 'some-random-url' meta_list = self.clusters.describe(cid_master_url) self.assertEqual([known_meta, known_meta2], meta_list) cid_meta = ClusterMetadata(project_id='some-random-project') meta_list = self.clusters.describe(cid_meta) self.assertEqual([known_meta, known_meta2], meta_list) def test_default_value_for_invalid_worker_number(self): meta = ClusterMetadata(project_id='test-project', num_workers=1) self.clusters.create(meta) self.assertEqual(meta.num_workers, 2) @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') @isolated_env class InteractiveBeamComputeTest(unittest.TestCase): def setUp(self): self.env = ie.current_env() self.env._is_in_ipython = False # Default to non-IPython def test_compute_blocking(self): p = beam.Pipeline(ir.InteractiveRunner()) data = list(range(10)) pcoll = p | 'Create' >> beam.Create(data) ib.watch(locals()) self.env.track_user_pipelines() result = ib.compute(pcoll, blocking=True) self.assertIsNone(result) # Blocking returns None self.assertTrue(pcoll in self.env.computed_pcollections) collected = ib.collect(pcoll, raw_records=True) self.assertEqual(collected, data) def test_compute_non_blocking(self): p = beam.Pipeline(ir.InteractiveRunner()) data = list(range(5)) pcoll = p | 'Create' >> beam.Create(data) ib.watch(locals()) self.env.track_user_pipelines() async_result = ib.compute(pcoll, blocking=False) self.assertIsInstance(async_result, AsyncComputationResult) pipeline_result = async_result.result(timeout=60) self.assertTrue(async_result.done()) self.assertIsNone(async_result.exception()) self.assertEqual(pipeline_result.state, PipelineState.DONE) self.assertTrue(pcoll in self.env.computed_pcollections) collected = ib.collect(pcoll, raw_records=True) self.assertEqual(collected, data) def test_compute_with_list_input(self): p = beam.Pipeline(ir.InteractiveRunner()) pcoll1 = p | 'Create1' >> beam.Create([1, 2, 3]) pcoll2 = p | 'Create2' >> beam.Create([4, 5, 6]) ib.watch(locals()) self.env.track_user_pipelines() ib.compute([pcoll1, pcoll2], blocking=True) self.assertTrue(pcoll1 in self.env.computed_pcollections) self.assertTrue(pcoll2 in self.env.computed_pcollections) def test_compute_with_dict_input(self): p = beam.Pipeline(ir.InteractiveRunner()) pcoll1 = p | 'Create1' >> beam.Create([1, 2, 3]) pcoll2 = p | 'Create2' >> beam.Create([4, 5, 6]) ib.watch(locals()) self.env.track_user_pipelines() ib.compute({'a': pcoll1, 'b': pcoll2}, blocking=True) self.assertTrue(pcoll1 in self.env.computed_pcollections) self.assertTrue(pcoll2 in self.env.computed_pcollections) def test_compute_empty_input(self): result = ib.compute([], blocking=True) self.assertIsNone(result) result_async = ib.compute([], blocking=False) self.assertIsNone(result_async) def test_compute_force_recompute(self): p = beam.Pipeline(ir.InteractiveRunner()) pcoll = p | 'Create' >> beam.Create([1, 2, 3]) ib.watch(locals()) self.env.track_user_pipelines() ib.compute(pcoll, blocking=True) self.assertTrue(pcoll in self.env.computed_pcollections) # Mock evict_computed_pcollections to check if it's called with patch.object(self.env, 'evict_computed_pcollections') as mock_evict: ib.compute(pcoll, blocking=True, force_compute=True) mock_evict.assert_called_once_with(p) self.assertTrue(pcoll in self.env.computed_pcollections) def test_compute_non_blocking_exception(self): p = beam.Pipeline(ir.InteractiveRunner()) def raise_error(elem): raise ValueError('Test Error') pcoll = p | 'Create' >> beam.Create([1]) | 'Error' >> beam.Map(raise_error) ib.watch(locals()) self.env.track_user_pipelines() async_result = ib.compute(pcoll, blocking=False) self.assertIsInstance(async_result, AsyncComputationResult) with self.assertRaises(ValueError): async_result.result(timeout=60) self.assertTrue(async_result.done()) self.assertIsInstance(async_result.exception(), ValueError) self.assertFalse(pcoll in self.env.computed_pcollections) @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', True) @patch('apache_beam.runners.interactive.recording_manager.display') @patch('apache_beam.runners.interactive.recording_manager.clear_output') @patch('apache_beam.runners.interactive.recording_manager.HTML') @patch('ipywidgets.Button') @patch('ipywidgets.FloatProgress') @patch('ipywidgets.Output') @patch('ipywidgets.HBox') @patch('ipywidgets.VBox') def test_compute_non_blocking_ipython_widgets( self, mock_vbox, mock_hbox, mock_output, mock_progress, mock_button, mock_html, mock_clear_output, mock_display, ): self.env._is_in_ipython = True p = beam.Pipeline(ir.InteractiveRunner()) pcoll = p | 'Create' >> beam.Create(range(3)) ib.watch(locals()) self.env.track_user_pipelines() mock_controls = mock_vbox.return_value mock_html_instance = mock_html.return_value async_result = ib.compute(pcoll, blocking=False) self.assertIsNotNone(async_result) mock_button.assert_called_once_with(description='Cancel') mock_progress.assert_called_once() mock_output.assert_called_once() mock_hbox.assert_called_once() mock_vbox.assert_called_once() mock_html.assert_called_once_with('

Initializing...

') self.assertEqual(mock_display.call_count, 2) mock_display.assert_has_calls([ call(mock_controls, display_id=async_result._display_id), call(mock_html_instance) ]) mock_clear_output.assert_called_once() async_result.result(timeout=60) # Let it finish def test_compute_dependency_wait_true(self): p = beam.Pipeline(ir.InteractiveRunner()) pcoll1 = p | 'Create1' >> beam.Create([1, 2, 3]) pcoll2 = pcoll1 | 'Map' >> beam.Map(lambda x: x * 2) ib.watch(locals()) self.env.track_user_pipelines() rm = self.env.get_recording_manager(p) # Start pcoll1 computation async_res1 = ib.compute(pcoll1, blocking=False) self.assertTrue(self.env.is_pcollection_computing(pcoll1)) # Spy on _wait_for_dependencies with patch.object(rm, '_wait_for_dependencies', wraps=rm._wait_for_dependencies) as spy_wait: async_res2 = ib.compute(pcoll2, blocking=False, wait_for_inputs=True) # Check that wait_for_dependencies was called for pcoll2 spy_wait.assert_called_with({pcoll2}, async_res2) # Let pcoll1 finish async_res1.result(timeout=60) self.assertTrue(pcoll1 in self.env.computed_pcollections) self.assertFalse(self.env.is_pcollection_computing(pcoll1)) # pcoll2 should now run and complete async_res2.result(timeout=60) self.assertTrue(pcoll2 in self.env.computed_pcollections) @patch.object(ie.InteractiveEnvironment, 'is_pcollection_computing') def test_compute_dependency_wait_false(self, mock_is_computing): p = beam.Pipeline(ir.InteractiveRunner()) pcoll1 = p | 'Create1' >> beam.Create([1, 2, 3]) pcoll2 = pcoll1 | 'Map' >> beam.Map(lambda x: x * 2) ib.watch(locals()) self.env.track_user_pipelines() rm = self.env.get_recording_manager(p) # Pretend pcoll1 is computing mock_is_computing.side_effect = lambda pcoll: pcoll is pcoll1 with patch.object(rm, '_execute_pipeline_fragment', wraps=rm._execute_pipeline_fragment) as spy_execute: async_res2 = ib.compute(pcoll2, blocking=False, wait_for_inputs=False) async_res2.result(timeout=60) # Assert that execute was called for pcoll2 without waiting spy_execute.assert_called_with({pcoll2}, async_res2, ANY, ANY) self.assertTrue(pcoll2 in self.env.computed_pcollections) def test_async_computation_result_cancel(self): p = beam.Pipeline(ir.InteractiveRunner()) # A stream that never finishes to test cancellation pcoll = p | beam.Create([1]) | beam.Map(lambda x: time.sleep(100)) ib.watch(locals()) self.env.track_user_pipelines() async_result = ib.compute(pcoll, blocking=False) self.assertIsInstance(async_result, AsyncComputationResult) # Give it a moment to start time.sleep(0.1) # Mock the pipeline result's cancel method mock_pipeline_result = MagicMock() mock_pipeline_result.state = PipelineState.RUNNING async_result.set_pipeline_result(mock_pipeline_result) self.assertTrue(async_result.cancel()) mock_pipeline_result.cancel.assert_called_once() # The future should be cancelled eventually by the runner # This part is hard to test without deeper runner integration with self.assertRaises(TimeoutError): async_result.result(timeout=1) # It should not complete successfully @patch( 'apache_beam.runners.interactive.recording_manager.RecordingManager.' '_execute_pipeline_fragment') def test_compute_multiple_async(self, mock_execute_fragment): p = beam.Pipeline(ir.InteractiveRunner()) pcoll1 = p | 'Create1' >> beam.Create([1, 2, 3]) pcoll2 = p | 'Create2' >> beam.Create([4, 5, 6]) pcoll3 = pcoll1 | 'Map1' >> beam.Map(lambda x: x * 2) ib.watch(locals()) self.env.track_user_pipelines() mock_pipeline_result = MagicMock() mock_pipeline_result.state = PipelineState.DONE mock_execute_fragment.return_value = mock_pipeline_result res1 = ib.compute(pcoll1, blocking=False) res2 = ib.compute(pcoll2, blocking=False) res3 = ib.compute(pcoll3, blocking=False) # Depends on pcoll1 self.assertIsNotNone(res1) self.assertIsNotNone(res2) self.assertIsNotNone(res3) res1.result(timeout=60) res2.result(timeout=60) res3.result(timeout=60) time.sleep(0.1) self.assertTrue( pcoll1 in self.env.computed_pcollections, "pcoll1 not marked computed") self.assertTrue( pcoll2 in self.env.computed_pcollections, "pcoll2 not marked computed") self.assertTrue( pcoll3 in self.env.computed_pcollections, "pcoll3 not marked computed") self.assertEqual(mock_execute_fragment.call_count, 3) @patch( 'apache_beam.runners.interactive.interactive_beam.' 'deferred_df_to_pcollection') def test_compute_input_flattening(self, mock_deferred_to_pcoll): p = beam.Pipeline(ir.InteractiveRunner()) pcoll1 = p | 'C1' >> beam.Create([1]) pcoll2 = p | 'C2' >> beam.Create([2]) pcoll3 = p | 'C3' >> beam.Create([3]) pcoll4 = p | 'C4' >> beam.Create([4]) class MockDeferred(DeferredBase): def __init__(self, pcoll): mock_expr = MagicMock() super().__init__(mock_expr) self._pcoll = pcoll def _get_underlying_pcollection(self): return self._pcoll deferred_pcoll = MockDeferred(pcoll4) mock_deferred_to_pcoll.return_value = (pcoll4, p) ib.watch(locals()) self.env.track_user_pipelines() with patch.object(self.env, 'get_recording_manager') as mock_get_rm: mock_rm = MagicMock() mock_get_rm.return_value = mock_rm ib.compute(pcoll1, [pcoll2], {'a': pcoll3}, deferred_pcoll) expected_pcolls = {pcoll1, pcoll2, pcoll3, pcoll4} mock_rm.compute_async.assert_called_once_with( expected_pcolls, wait_for_inputs=True, blocking=False, runner=None, options=None, force_compute=False) def test_compute_invalid_input_type(self): with self.assertRaisesRegex(ValueError, "not a dict, an iterable or a PCollection"): ib.compute(123) def test_compute_mixed_pipelines(self): p1 = beam.Pipeline(ir.InteractiveRunner()) pcoll1 = p1 | 'C1' >> beam.Create([1]) p2 = beam.Pipeline(ir.InteractiveRunner()) pcoll2 = p2 | 'C2' >> beam.Create([2]) ib.watch(locals()) self.env.track_user_pipelines() with self.assertRaisesRegex( ValueError, "All PCollections must belong to the same pipeline"): ib.compute(pcoll1, pcoll2) @patch( 'apache_beam.runners.interactive.interactive_beam.' 'deferred_df_to_pcollection') @patch.object(ib, 'watch') def test_compute_with_deferred_base(self, mock_watch, mock_deferred_to_pcoll): p = beam.Pipeline(ir.InteractiveRunner()) pcoll = p | 'C1' >> beam.Create([1]) class MockDeferred(DeferredBase): def __init__(self, pcoll): # Provide a dummy expression to satisfy DeferredBase.__init__ mock_expr = MagicMock() super().__init__(mock_expr) self._pcoll = pcoll def _get_underlying_pcollection(self): return self._pcoll deferred = MockDeferred(pcoll) mock_deferred_to_pcoll.return_value = (pcoll, p) with patch.object(self.env, 'get_recording_manager') as mock_get_rm: mock_rm = MagicMock() mock_get_rm.return_value = mock_rm ib.compute(deferred) mock_deferred_to_pcoll.assert_called_once_with(deferred) self.assertEqual(mock_watch.call_count, 2) mock_watch.assert_has_calls([ call({f'anonymous_pcollection_{id(pcoll)}': pcoll}), call({f'anonymous_pipeline_{id(p)}': p}) ], any_order=False) mock_rm.compute_async.assert_called_once_with({pcoll}, wait_for_inputs=True, blocking=False, runner=None, options=None, force_compute=False) def test_compute_new_pipeline(self): p = beam.Pipeline(ir.InteractiveRunner()) pcoll = p | 'Create' >> beam.Create([1]) # NOT calling ib.watch() or track_user_pipelines() with patch.object(self.env, 'get_recording_manager') as mock_get_rm, \ patch.object(ib, 'watch') as mock_watch: mock_rm = MagicMock() mock_get_rm.return_value = mock_rm ib.compute(pcoll) mock_watch.assert_called_with({f'anonymous_pipeline_{id(p)}': p}) mock_get_rm.assert_called_once_with(p, create_if_absent=True) mock_rm.compute_async.assert_called_once() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/interactive_environment.py0000644000175100017510000007616515206113032027220 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module of the current Interactive Beam environment. For internal use only; no backwards-compatibility guarantees. Provides interfaces to interact with existing Interactive Beam environment. External Interactive Beam users please use interactive_beam module in application code or notebook. """ # pytype: skip-file import atexit import importlib import logging import os import tempfile from collections.abc import Iterable from pathlib import PurePath import apache_beam as beam from apache_beam.runners import DataflowRunner from apache_beam.runners import runner from apache_beam.runners.direct import direct_runner from apache_beam.runners.interactive import cache_manager as cache from apache_beam.runners.interactive.messaging.interactive_environment_inspector import InteractiveEnvironmentInspector from apache_beam.runners.interactive.sql.sql_chain import SqlChain from apache_beam.runners.interactive.user_pipeline_tracker import UserPipelineTracker from apache_beam.runners.interactive.utils import assert_bucket_exists from apache_beam.runners.interactive.utils import detect_pipeline_runner from apache_beam.runners.interactive.utils import register_ipython_log_handler from apache_beam.utils.interactive_utils import is_in_ipython from apache_beam.utils.interactive_utils import is_in_notebook # Interactive Beam user flow is data-centric rather than pipeline-centric, so # there is only one global interactive environment instance that manages # implementation that enables interactivity. _interactive_beam_env = None _LOGGER = logging.getLogger(__name__) # By `format(customized_script=xxx)`, the given `customized_script` is # guaranteed to be executed within access to a jquery with datatable plugin # configured which is useful so that any `customized_script` is resilient to # browser refresh. Inside `customized_script`, use `$` as jQuery. _JQUERY_WITH_DATATABLE_TEMPLATE = """ if (typeof window.interactive_beam_jquery == 'undefined') {{ var jqueryScript = document.createElement('script'); jqueryScript.src = 'https://code.jquery.com/jquery-3.4.1.slim.min.js'; jqueryScript.type = 'text/javascript'; jqueryScript.onload = function() {{ var datatableScript = document.createElement('script'); datatableScript.src = 'https://cdn.datatables.net/1.10.20/js/jquery.dataTables.min.js'; datatableScript.type = 'text/javascript'; datatableScript.onload = function() {{ window.interactive_beam_jquery = jQuery.noConflict(true); window.interactive_beam_jquery(document).ready(function($){{ {customized_script} }}); }} document.head.appendChild(datatableScript); }}; document.head.appendChild(jqueryScript); }} else {{ window.interactive_beam_jquery(document).ready(function($){{ {customized_script} }}); }}""" # By `format(hrefs=xxx)`, the given `hrefs` will be imported as HTML imports. # Since HTML import might not be supported by the browser, we check if HTML # import is supported by the browser, if so, import HTMLs else setup # webcomponents and chain the HTML import to the end of onload. _HTML_IMPORT_TEMPLATE = """ var import_html = () => {{ {hrefs}.forEach(href => {{ var link = document.createElement('link'); link.rel = 'import' link.href = href; document.head.appendChild(link); }}); }} if ('import' in document.createElement('link')) {{ import_html(); }} else {{ var webcomponentScript = document.createElement('script'); webcomponentScript.src = 'https://cdnjs.cloudflare.com/ajax/libs/webcomponentsjs/1.3.3/webcomponents-lite.js'; webcomponentScript.type = 'text/javascript'; webcomponentScript.onload = function(){{ import_html(); }}; document.head.appendChild(webcomponentScript); }}""" def current_env(): """Gets current Interactive Beam environment.""" global _interactive_beam_env if not _interactive_beam_env: _interactive_beam_env = InteractiveEnvironment() return _interactive_beam_env def new_env(): """Creates a new Interactive Beam environment to replace current one.""" global _interactive_beam_env if _interactive_beam_env: _interactive_beam_env.cleanup() _interactive_beam_env = None return current_env() class InteractiveEnvironment(object): """An interactive environment with cache and pipeline variable metadata. Interactive Beam will use the watched variable information to determine if a PCollection is assigned to a variable in user pipeline definition. When executing the pipeline, interactivity is applied with implicit cache mechanism for those PCollections if the pipeline is interactive. Users can also visualize and introspect those PCollections in user code since they have handles to the variables. """ def __init__(self): # Registers a cleanup routine when system exits. atexit.register(self.cleanup) # Holds cache managers that manage source recording and intermediate # PCollection cache for each pipeline. Each key is a stringified user # defined pipeline instance's id. self._cache_managers = {} # Holds RecordingManagers keyed by pipeline instance id. self._recording_managers = {} # Holds class instances, module object, string of module names. self._watching_set = set() # Holds variables list of (Dict[str, object]). self._watching_dict_list = [] # Holds results of main jobs as Dict[str, PipelineResult]. # Each key is a pipeline instance defined by the end user. The # InteractiveRunner is responsible for populating this dictionary # implicitly. self._main_pipeline_results = {} # Holds background caching jobs as Dict[str, BackgroundCachingJob]. # Each key is a pipeline instance defined by the end user. The # InteractiveRunner or its enclosing scope is responsible for populating # this dictionary implicitly when a background caching jobs is started. self._background_caching_jobs = {} # Holds TestStreamServiceControllers that controls gRPC servers serving # events as test stream of TestStreamPayload.Event. # Dict[str, TestStreamServiceController]. Each key is a pipeline # instance defined by the end user. The InteractiveRunner or its enclosing # scope is responsible for populating this dictionary implicitly when a new # controller is created to start a new gRPC server. The server stays alive # until a new background caching job is started thus invalidating everything # the gRPC server serves. self._test_stream_service_controllers = {} self._cached_source_signature = {} self._tracked_user_pipelines = UserPipelineTracker() from apache_beam.runners.interactive.interactive_beam import clusters self.clusters = clusters # Tracks the computation completeness of PCollections. PCollections tracked # here don't need to be re-computed when data introspection is needed. self._computed_pcolls = set() self._computing_pcolls = set() # Always watch __main__ module. self.watch('__main__') # Check if [interactive] dependencies are installed. try: import IPython # pylint: disable=unused-import import timeloop # pylint: disable=unused-import from facets_overview.generic_feature_statistics_generator import \ GenericFeatureStatisticsGenerator # pylint: disable=unused-import from google.cloud import dataproc_v1 # pylint: disable=unused-import self._is_interactive_ready = True except ImportError: self._is_interactive_ready = False _LOGGER.warning( 'Dependencies required for Interactive Beam PCollection ' 'visualization are not available, please use: `pip ' 'install apache-beam[interactive]` to install necessary ' 'dependencies to enable all data visualization features.') self._is_in_ipython = is_in_ipython() self._is_in_notebook = is_in_notebook() if not self._is_in_ipython: _LOGGER.warning( 'You cannot use Interactive Beam features when you are ' 'not in an interactive environment such as a Jupyter ' 'notebook or ipython terminal.') if self._is_in_ipython and not self._is_in_notebook: _LOGGER.warning( 'You have limited Interactive Beam features since your ' 'ipython kernel is not connected to any notebook frontend.') if self._is_in_notebook: self.load_jquery_with_datatable() register_ipython_log_handler() # A singleton inspector instance to message information of current # environment to other applications. self._inspector = InteractiveEnvironmentInspector() # A similar singleton inspector except it includes synthetic variables # generated by Interactive Beam. self._inspector_with_synthetic = InteractiveEnvironmentInspector( ignore_synthetic=False) self.sql_chain = {} @property def options(self): """A reference to the global interactive options. Provided to avoid import loop or excessive dynamic import. All internal Interactive Beam modules should access interactive_beam.options through this property. """ from apache_beam.runners.interactive.interactive_beam import options return options @property def is_interactive_ready(self): """If the [interactive] dependencies are installed.""" return self._is_interactive_ready @property def is_in_ipython(self): """If the runtime is within an IPython kernel.""" return self._is_in_ipython @property def is_in_notebook(self): """If the kernel is connected to a notebook frontend. If not, it could be that the user is using kernel in a terminal or a unit test. """ return self._is_in_notebook @property def inspector(self): """Gets the singleton InteractiveEnvironmentInspector to retrieve information consumable by other applications such as a notebook extension.""" return self._inspector @property def inspector_with_synthetic(self): """Gets the singleton InteractiveEnvironmentInspector with additional synthetic variables generated by Interactive Beam. Internally used.""" return self._inspector_with_synthetic def cleanup_pipeline(self, pipeline): from apache_beam.runners.interactive import background_caching_job as bcj bcj.attempt_to_cancel_background_caching_job(pipeline) bcj.attempt_to_stop_test_stream_service(pipeline) cache_manager = self.get_cache_manager(pipeline) # Recording manager performs cache manager cleanup during eviction, so we # don't need to clean it up here. if cache_manager and self.get_recording_manager(pipeline) is None: cache_manager.cleanup() self.clusters.cleanup(pipeline) def cleanup_environment(self): for _, job in self._background_caching_jobs.items(): if job: job.cancel() for _, controller in self._test_stream_service_controllers.items(): if controller: controller.stop() for pipeline_id, cache_manager in self._cache_managers.items(): # Recording manager performs cache manager cleanup during eviction, so # we don't need to clean it up here. if cache_manager and pipeline_id not in self._recording_managers: cache_manager.cleanup() self.clusters.cleanup(force=True) def cleanup(self, pipeline=None): """Cleans up cached states for the given pipeline. Noop if the given pipeline is absent from the environment. Cleans up for all pipelines if no pipeline is specified.""" if pipeline: self.cleanup_pipeline(pipeline) else: self.cleanup_environment() self.evict_recording_manager(pipeline) self.evict_background_caching_job(pipeline) self.evict_test_stream_service_controller(pipeline) self.evict_computed_pcollections(pipeline) self.evict_cached_source_signature(pipeline) self.evict_pipeline_result(pipeline) self.evict_tracked_pipelines(pipeline) def _track_user_pipelines(self, watchable): """Tracks user pipelines from the given watchable.""" pipelines = set() if isinstance(watchable, beam.Pipeline): pipelines.add(watchable) elif isinstance(watchable, dict): for v in watchable.values(): if isinstance(v, beam.Pipeline): pipelines.add(v) elif isinstance(watchable, Iterable): for v in watchable: if isinstance(v, beam.Pipeline): pipelines.add(v) for p in pipelines: self._tracked_user_pipelines.add_user_pipeline(p) _ = self.get_cache_manager(p, create_if_absent=True) _ = self.get_recording_manager(p, create_if_absent=True) def watch(self, watchable): """Watches a watchable. A watchable can be a dictionary of variable metadata such as locals(), a str name of a module, a module object or an instance of a class. The variable can come from any scope even local. Duplicated variable naming doesn't matter since they are different instances. Duplicated variables are also allowed when watching. """ if isinstance(watchable, dict): self._watching_dict_list.append(watchable.items()) else: self._watching_set.add(watchable) self._track_user_pipelines(watchable) def watching(self): """Analyzes and returns a list of pair lists referring to variable names and values from watched scopes. Each entry in the list represents the variable defined within a watched watchable. Currently, each entry holds a list of pairs. The format might change in the future to hold more metadata. Duplicated pairs are allowed. And multiple paris can have the same variable name as the "first" while having different variable values as the "second" since variables in different scopes can have the same name. """ watching = list(self._watching_dict_list) for watchable in self._watching_set: if isinstance(watchable, str): module = importlib.import_module(watchable) watching.append(vars(module).items()) else: watching.append(vars(watchable).items()) return watching def set_cache_manager(self, cache_manager, pipeline): """Sets the cache manager held by current Interactive Environment for the given pipeline.""" if self.get_cache_manager(pipeline) is cache_manager: # NOOP if setting to the same cache_manager. return if self.get_cache_manager(pipeline): # Invoke cleanup routine when a new cache_manager is forcefully set and # current cache_manager is not None. self.cleanup(pipeline) self._cache_managers[str(id(pipeline))] = cache_manager def get_cache_manager(self, pipeline, create_if_absent=False): """Gets the cache manager held by current Interactive Environment for the given pipeline. If the pipeline is absent from the environment while create_if_absent is True, creates and returns a new file based cache manager for the pipeline.""" cache_manager = self._cache_managers.get(str(id(pipeline)), None) pipeline_runner = detect_pipeline_runner(pipeline) if not cache_manager and create_if_absent: cache_root = self.options.cache_root if cache_root: if cache_root.startswith('gs://'): cache_dir = self._get_gcs_cache_dir(pipeline, cache_root) else: cache_dir = tempfile.mkdtemp(dir=cache_root) if not isinstance(pipeline_runner, direct_runner.DirectRunner): _LOGGER.warning( 'A local cache directory has been specified while ' 'not using DirectRunner. It is recommended to cache into a ' 'GCS bucket instead.') else: staging_location = pipeline.options.get_all_options( )['staging_location'] if isinstance(pipeline_runner, DataflowRunner) and staging_location: cache_dir = self._get_gcs_cache_dir(pipeline, staging_location) _LOGGER.info( 'No cache_root detected. ' 'Defaulting to staging_location %s for cache location.', staging_location) else: cache_dir = tempfile.mkdtemp( suffix=str(id(pipeline)), prefix='it-', dir=os.environ.get('TEST_TMPDIR', None)) cache_manager = cache.FileBasedCacheManager(cache_dir) self._cache_managers[str(id(pipeline))] = cache_manager return cache_manager def evict_cache_manager(self, pipeline=None): """Evicts the cache manager held by current Interactive Environment for the given pipeline. Noop if the pipeline is absent from the environment. If no pipeline is specified, evicts for all pipelines.""" self.cleanup(pipeline) if pipeline: return self._cache_managers.pop(str(id(pipeline)), None) self._cache_managers.clear() def set_recording_manager(self, recording_manager, pipeline): """Sets the recording manager for the given pipeline.""" if self.get_recording_manager(pipeline) is recording_manager: # NOOP if setting to the same recording_manager. return self._recording_managers[str(id(pipeline))] = recording_manager def get_recording_manager(self, pipeline, create_if_absent=False): """Gets the recording manager for the given pipeline.""" # Allow initial module loading to be complete and not have a circular # import. from apache_beam.runners.interactive.recording_manager import RecordingManager recording_manager = self._recording_managers.get(str(id(pipeline)), None) if not recording_manager and create_if_absent: # Get the pipeline variable name for the user. This is useful if the user # has multiple pipelines. pipeline_var = '' for w in self.watching(): for var, val in w: if val is pipeline: pipeline_var = var break recording_manager = RecordingManager(pipeline, pipeline_var) self._recording_managers[str(id(pipeline))] = recording_manager return recording_manager def evict_recording_manager(self, pipeline): """Evicts the recording manager for the given pipeline. This stops the background caching job and clears the cache. Noop if the pipeline is absent from the environment. If no pipeline is specified, evicts for all pipelines. """ if not pipeline: for rm in self._recording_managers.values(): rm.cancel() rm.clear() self._recording_managers = {} return recording_manager = self.get_recording_manager(pipeline) if recording_manager: recording_manager.cancel() recording_manager.clear() del self._recording_managers[str(id(pipeline))] def describe_all_recordings(self): """Returns a description of the recording for all watched pipelnes.""" return { self.pipeline_id_to_pipeline(pid): rm.describe() for pid, rm in self._recording_managers.items() } def set_pipeline_result(self, pipeline, result): """Sets the pipeline run result. Adds one if absent. Otherwise, replace.""" assert issubclass(type(pipeline), beam.Pipeline), ( 'pipeline must be an instance of apache_beam.Pipeline or its subclass') assert issubclass(type(result), runner.PipelineResult), ( 'result must be an instance of ' 'apache_beam.runners.runner.PipelineResult or its subclass') self._main_pipeline_results[str(id(pipeline))] = result def evict_pipeline_result(self, pipeline=None): """Evicts the last run result of the given pipeline. Noop if the pipeline is absent from the environment. If no pipeline is specified, evicts for all pipelines.""" if pipeline: return self._main_pipeline_results.pop(str(id(pipeline)), None) self._main_pipeline_results.clear() def pipeline_result(self, pipeline): """Gets the pipeline run result. None if absent.""" return self._main_pipeline_results.get(str(id(pipeline)), None) def set_background_caching_job(self, pipeline, background_caching_job): """Sets the background caching job started from the given pipeline.""" assert issubclass(type(pipeline), beam.Pipeline), ( 'pipeline must be an instance of apache_beam.Pipeline or its subclass') from apache_beam.runners.interactive.background_caching_job import BackgroundCachingJob assert isinstance(background_caching_job, BackgroundCachingJob), ( 'background_caching job must be an instance of BackgroundCachingJob') self._background_caching_jobs[str(id(pipeline))] = background_caching_job def get_background_caching_job(self, pipeline): """Gets the background caching job started from the given pipeline.""" return self._background_caching_jobs.get(str(id(pipeline)), None) def evict_background_caching_job(self, pipeline=None): """Evicts the background caching job started from the given pipeline. Noop if the given pipeline is absent from the environment. If no pipeline is specified, evicts for all pipelines.""" if pipeline: return self._background_caching_jobs.pop(str(id(pipeline)), None) self._background_caching_jobs.clear() def set_test_stream_service_controller(self, pipeline, controller): """Sets the test stream service controller that has started a gRPC server serving the test stream for any job started from the given user defined pipeline. """ self._test_stream_service_controllers[str(id(pipeline))] = controller def get_test_stream_service_controller(self, pipeline): """Gets the test stream service controller that has started a gRPC server serving the test stream for any job started from the given user defined pipeline. """ return self._test_stream_service_controllers.get(str(id(pipeline)), None) def evict_test_stream_service_controller(self, pipeline): """Evicts and pops the test stream service controller that has started a gRPC server serving the test stream for any job started from the given user defined pipeline. Noop if the given pipeline is absent from the environment. If no pipeline is specified, evicts for all pipelines. """ if pipeline: return self._test_stream_service_controllers.pop(str(id(pipeline)), None) self._test_stream_service_controllers.clear() def is_terminated(self, pipeline): """Queries if the most recent job (by executing the given pipeline) state is in a terminal state. True if absent.""" result = self.pipeline_result(pipeline) if result: return runner.PipelineState.is_terminal(result.state) return True def set_cached_source_signature(self, pipeline, signature): self._cached_source_signature[str(id(pipeline))] = signature def get_cached_source_signature(self, pipeline): return self._cached_source_signature.get(str(id(pipeline)), set()) def evict_cached_source_signature(self, pipeline=None): """Evicts the signature generated for each recorded source of the given pipeline. Noop if the given pipeline is absent from the environment. If no pipeline is specified, evicts for all pipelines.""" if pipeline: return self._cached_source_signature.pop(str(id(pipeline)), None) self._cached_source_signature.clear() def track_user_pipelines(self): """Record references to all user defined pipeline instances watched in current environment. Current static global singleton interactive environment holds references to a set of pipeline instances defined by the user in the watched scope. Interactive Beam features could use the references to determine if a given pipeline is defined by user or implicitly created by Beam SDK or runners, then handle them differently. This is invoked every time a PTransform is to be applied if the current code execution is under ipython due to the possibility that any user defined pipeline can be re-evaluated through notebook cell re-execution at any time. Each time this is invoked, it will check if there is a cache manager already created for each user defined pipeline. If not, create one for it. If a pipeline is no longer watched due to re-execution while its PCollections are still in watched scope, the pipeline becomes anonymous but still accessible indirectly through references to its PCollections. This function also clears up internal states for those anonymous pipelines once all their PCollections are anonymous. """ for watching in self.watching(): for _, val in watching: if isinstance(val, beam.pipeline.Pipeline): self._tracked_user_pipelines.add_user_pipeline(val) _ = self.get_cache_manager(val, create_if_absent=True) _ = self.get_recording_manager(val, create_if_absent=True) all_tracked_pipeline_ids = set(self._background_caching_jobs.keys()).union( set(self._test_stream_service_controllers.keys()), set(self._cache_managers.keys()), {str(id(pcoll.pipeline)) for pcoll in self._computed_pcolls}, set(self._cached_source_signature.keys()), set(self._main_pipeline_results.keys())) inspectable_pipelines = self._inspector.inspectable_pipelines for pipeline in all_tracked_pipeline_ids: if pipeline not in inspectable_pipelines: self.cleanup(pipeline) @property def tracked_user_pipelines(self): """Returns the user pipelines in this environment.""" for p in self._tracked_user_pipelines: yield p def user_pipeline(self, derived_pipeline): """Returns the user pipeline for the given derived pipeline.""" return self._tracked_user_pipelines.get_user_pipeline(derived_pipeline) def add_user_pipeline(self, user_pipeline): self._tracked_user_pipelines.add_user_pipeline(user_pipeline) def add_derived_pipeline(self, user_pipeline, derived_pipeline): """Adds the derived pipeline to the parent user pipeline.""" self._tracked_user_pipelines.add_derived_pipeline( user_pipeline, derived_pipeline) def evict_tracked_pipelines(self, user_pipeline): """Evicts the user pipeline and its derived pipelines.""" if user_pipeline: self._tracked_user_pipelines.evict(user_pipeline) else: self._tracked_user_pipelines.clear() def pipeline_id_to_pipeline(self, pid): """Converts a pipeline id to a user pipeline. """ return self._tracked_user_pipelines.get_pipeline(pid) def mark_pcollection_computed(self, pcolls): """Marks computation completeness for the given pcolls. Interactive Beam can use this information to determine if a computation is needed to introspect the data of any given PCollection. """ self._computed_pcolls.update(pcoll for pcoll in pcolls) def evict_computed_pcollections(self, pipeline=None): """Evicts all computed PCollections for the given pipeline. If no pipeline is specified, evicts for all pipelines. """ if pipeline: discarded = set() for pcoll in self._computed_pcolls: if pcoll.pipeline is pipeline: discarded.add(pcoll) self._computed_pcolls -= discarded else: self._computed_pcolls = set() @property def computed_pcollections(self): return self._computed_pcolls def load_jquery_with_datatable(self): """Loads common resources to enable jquery with datatable configured for notebook frontends if necessary. If the resources have been loaded, NOOP. A window.interactive_beam_jquery with datatable plugin configured can be used in following notebook cells once this is invoked. #. There should only be one jQuery imported. #. Datatable needs to be imported after jQuery is loaded. #. Imported jQuery is attached to window named as jquery[version]. #. The window attachment needs to happen at the end of import chain until all jQuery plugins are set. """ try: from IPython.display import Javascript from IPython.display import display_javascript display_javascript( Javascript( _JQUERY_WITH_DATATABLE_TEMPLATE.format(customized_script=''))) except ImportError: pass # NOOP if dependencies are not available. def import_html_to_head(self, html_hrefs): """Imports given external HTMLs (supported through webcomponents) into the head of the document. On load of webcomponentsjs, import given HTMLs. If HTML import is already supported, skip loading webcomponentsjs. No matter how many times an HTML import occurs in the document, only the first occurrence really embeds the external HTML. In a notebook environment, the body of the document is always changing due to cell [re-]execution, deletion and re-ordering. Thus, HTML imports shouldn't be put in the body especially the output areas of notebook cells. """ try: from IPython.display import Javascript from IPython.display import display_javascript display_javascript( Javascript(_HTML_IMPORT_TEMPLATE.format(hrefs=html_hrefs))) except ImportError: pass # NOOP if dependencies are not available. def get_sql_chain(self, pipeline, set_user_pipeline=False): if pipeline not in self.sql_chain: self.sql_chain[pipeline] = SqlChain() chain = self.sql_chain[pipeline] if set_user_pipeline: if chain.user_pipeline and chain.user_pipeline is not pipeline: raise ValueError( 'The beam_sql magic tries to query PCollections from multiple ' 'pipelines: %s and %s', chain.user_pipeline, pipeline) chain.user_pipeline = pipeline return chain def _get_gcs_cache_dir(self, pipeline, cache_dir): cache_dir_path = PurePath(cache_dir) if len(cache_dir_path.parts) < 2: _LOGGER.error( 'GCS bucket cache path "%s" is too short to be valid. See ' 'https://cloud.google.com/storage/docs/naming-buckets for ' 'the expected format.', cache_dir) raise ValueError('cache_root GCS bucket path is invalid.') bucket_name = cache_dir_path.parts[1] assert_bucket_exists(bucket_name) return 'gs://{}/{}'.format('/'.join(cache_dir_path.parts[1:]), id(pipeline)) @property def computing_pcollections(self): return self._computing_pcolls def mark_pcollection_computing(self, pcolls): """Marks the given pcolls as currently being computed.""" self._computing_pcolls.update(pcolls) def unmark_pcollection_computing(self, pcolls): """Removes the given pcolls from the computing set.""" self._computing_pcolls.difference_update(pcolls) def is_pcollection_computing(self, pcoll): """Checks if the given pcollection is currently being computed.""" return pcoll in self._computing_pcolls ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/interactive_environment_test.py0000644000175100017510000003766215206113032030256 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.runners.interactive.interactive_environment.""" # pytype: skip-file import importlib import unittest from unittest.mock import patch import apache_beam as beam from apache_beam.runners import runner from apache_beam.runners.interactive import cache_manager as cache from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.recording_manager import RecordingManager from apache_beam.runners.interactive.sql.sql_chain import SqlNode from apache_beam.runners.interactive.testing.mock_env import isolated_env # The module name is also a variable in module. _module_name = 'apache_beam.runners.interactive.interactive_environment_test' @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') @isolated_env class InteractiveEnvironmentTest(unittest.TestCase): def setUp(self): self._p = beam.Pipeline() self._var_in_class_instance = 'a var in class instance' def assertVariableWatched(self, variable_name, variable_val): self.assertTrue(self._is_variable_watched(variable_name, variable_val)) def assertVariableNotWatched(self, variable_name, variable_val): self.assertFalse(self._is_variable_watched(variable_name, variable_val)) def _is_variable_watched(self, variable_name, variable_val): return any((variable_name, variable_val) in watching for watching in ie.current_env().watching()) def _a_function_with_local_watched(self): local_var_watched = 123 # pylint: disable=possibly-unused-variable ie.current_env().watch(locals()) def _a_function_not_watching_local(self): local_var_not_watched = 456 # pylint: disable=unused-variable def test_watch_main_by_default(self): self.assertTrue('__main__' in ie.current_env()._watching_set) # __main__ module has variable __name__ with value '__main__' self.assertVariableWatched('__name__', '__main__') def test_watch_a_module_by_name(self): self.assertFalse(_module_name in ie.current_env()._watching_set) self.assertVariableNotWatched('_module_name', _module_name) ie.current_env().watch(_module_name) self.assertTrue(_module_name in ie.current_env()._watching_set) self.assertVariableWatched('_module_name', _module_name) def test_watch_a_module_by_module_object(self): module = importlib.import_module(_module_name) self.assertFalse(module in ie.current_env()._watching_set) self.assertVariableNotWatched('_module_name', _module_name) ie.current_env().watch(module) self.assertTrue(module in ie.current_env()._watching_set) self.assertVariableWatched('_module_name', _module_name) def test_watch_locals(self): self.assertVariableNotWatched('local_var_watched', 123) self.assertVariableNotWatched('local_var_not_watched', 456) self._a_function_with_local_watched() self.assertVariableWatched('local_var_watched', 123) self._a_function_not_watching_local() self.assertVariableNotWatched('local_var_not_watched', 456) def test_watch_class_instance(self): self.assertVariableNotWatched( '_var_in_class_instance', self._var_in_class_instance) ie.current_env().watch(self) self.assertVariableWatched( '_var_in_class_instance', self._var_in_class_instance) def test_fail_to_set_pipeline_result_key_not_pipeline(self): class NotPipeline(object): pass with self.assertRaises(AssertionError) as ctx: ie.current_env().set_pipeline_result( NotPipeline(), runner.PipelineResult(runner.PipelineState.RUNNING)) self.assertTrue( 'pipeline must be an instance of apache_beam.Pipeline ' 'or its subclass' in ctx.exception) def test_fail_to_set_pipeline_result_value_not_pipeline_result(self): class NotResult(object): pass with self.assertRaises(AssertionError) as ctx: ie.current_env().set_pipeline_result(self._p, NotResult()) self.assertTrue( 'result must be an instance of ' 'apache_beam.runners.runner.PipelineResult or its ' 'subclass' in ctx.exception) def test_set_pipeline_result_successfully(self): class PipelineSubClass(beam.Pipeline): pass class PipelineResultSubClass(runner.PipelineResult): pass pipeline = PipelineSubClass() pipeline_result = PipelineResultSubClass(runner.PipelineState.RUNNING) ie.current_env().set_pipeline_result(pipeline, pipeline_result) self.assertIs(ie.current_env().pipeline_result(pipeline), pipeline_result) def test_determine_terminal_state(self): for state in (runner.PipelineState.DONE, runner.PipelineState.FAILED, runner.PipelineState.CANCELLED, runner.PipelineState.UPDATED, runner.PipelineState.DRAINED): ie.current_env().set_pipeline_result( self._p, runner.PipelineResult(state)) self.assertTrue(ie.current_env().is_terminated(self._p)) for state in (runner.PipelineState.UNKNOWN, runner.PipelineState.STARTING, runner.PipelineState.STOPPED, runner.PipelineState.RUNNING, runner.PipelineState.DRAINING, runner.PipelineState.PENDING, runner.PipelineState.CANCELLING, runner.PipelineState.UNRECOGNIZED): ie.current_env().set_pipeline_result( self._p, runner.PipelineResult(state)) self.assertFalse(ie.current_env().is_terminated(self._p)) def test_evict_pipeline_result(self): pipeline_result = runner.PipelineResult(runner.PipelineState.DONE) ie.current_env().set_pipeline_result(self._p, pipeline_result) self.assertIs( ie.current_env().evict_pipeline_result(self._p), pipeline_result) self.assertIs(ie.current_env().pipeline_result(self._p), None) def test_pipeline_result_is_none_when_pipeline_absent(self): self.assertIs(ie.current_env().pipeline_result(self._p), None) self.assertIs(ie.current_env().is_terminated(self._p), True) self.assertIs(ie.current_env().evict_pipeline_result(self._p), None) def test_cleanup_registered_when_creating_new_env(self): with patch('atexit.register') as mocked_atexit: _ = ie.InteractiveEnvironment() mocked_atexit.assert_called_once() def test_cleanup_invoked_when_new_env_replace_not_none_env(self): ie._interactive_beam_env = self.current_env with patch('apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.cleanup') as mocked_cleanup: ie.new_env() mocked_cleanup.assert_called_once() def test_cleanup_not_invoked_when_cm_changed_from_none(self): env = ie.InteractiveEnvironment() with patch('apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.cleanup') as mocked_cleanup: dummy_pipeline = 'dummy' self.assertIsNone(env.get_cache_manager(dummy_pipeline)) cache_manager = cache.FileBasedCacheManager() env.set_cache_manager(cache_manager, dummy_pipeline) mocked_cleanup.assert_not_called() self.assertIs(env.get_cache_manager(dummy_pipeline), cache_manager) def test_cleanup_invoked_when_not_none_cm_changed(self): env = ie.InteractiveEnvironment() with patch('apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.cleanup') as mocked_cleanup: dummy_pipeline = 'dummy' env.set_cache_manager(cache.FileBasedCacheManager(), dummy_pipeline) mocked_cleanup.assert_not_called() env.set_cache_manager(cache.FileBasedCacheManager(), dummy_pipeline) mocked_cleanup.assert_called_once() def test_noop_when_cm_is_not_changed(self): cache_manager = cache.FileBasedCacheManager() dummy_pipeline = 'dummy' env = ie.InteractiveEnvironment() with patch('apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.cleanup') as mocked_cleanup: env._cache_managers[str(id(dummy_pipeline))] = cache_manager mocked_cleanup.assert_not_called() env.set_cache_manager(cache_manager, dummy_pipeline) mocked_cleanup.assert_not_called() def test_get_cache_manager_creates_cache_manager_if_absent(self): env = ie.InteractiveEnvironment() dummy_pipeline = beam.Pipeline() self.assertIsNone(env.get_cache_manager(dummy_pipeline)) self.assertIsNotNone( env.get_cache_manager(dummy_pipeline, create_if_absent=True)) def test_track_user_pipeline_cleanup_non_inspectable_pipeline(self): dummy_pipeline_1 = beam.Pipeline() dummy_pipeline_2 = beam.Pipeline() dummy_pipeline_3 = beam.Pipeline() dummy_pipeline_4 = beam.Pipeline() dummy_pcoll = dummy_pipeline_4 | beam.Create([1]) dummy_pipeline_5 = beam.Pipeline() dummy_non_inspectable_pipeline = 'dummy' ie.current_env().watch(locals()) from apache_beam.runners.interactive.background_caching_job import BackgroundCachingJob ie.current_env().set_background_caching_job( dummy_pipeline_1, BackgroundCachingJob( runner.PipelineResult(runner.PipelineState.DONE), limiters=[])) ie.current_env().set_test_stream_service_controller(dummy_pipeline_2, None) ie.current_env().set_cache_manager( cache.FileBasedCacheManager(), dummy_pipeline_3) ie.current_env().mark_pcollection_computed([dummy_pcoll]) ie.current_env().set_cached_source_signature( dummy_non_inspectable_pipeline, None) ie.current_env().set_pipeline_result( dummy_pipeline_5, runner.PipelineResult(runner.PipelineState.RUNNING)) with patch('apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.cleanup') as mocked_cleanup: ie.current_env().track_user_pipelines() mocked_cleanup.assert_called_once() def test_evict_pcollections(self): """Tests the evicton logic in the InteractiveEnvironment.""" # Create two PCollection, one that will be evicted and another that won't. p_to_evict = beam.Pipeline() to_evict = p_to_evict | beam.Create([]) p_not_evicted = beam.Pipeline() not_evicted = p_not_evicted | beam.Create([]) # Mark the PCollections as computed because the eviction logic only works # on computed PCollections. ie.current_env().mark_pcollection_computed([to_evict, not_evicted]) self.assertSetEqual( ie.current_env().computed_pcollections, {to_evict, not_evicted}) # Evict the PCollection and then check that the other PCollection is safe. ie.current_env().evict_computed_pcollections(p_to_evict) self.assertSetEqual(ie.current_env().computed_pcollections, {not_evicted}) def test_set_get_recording_manager(self): p = beam.Pipeline() rm = RecordingManager(p) ie.current_env().set_recording_manager(rm, p) self.assertIs(rm, ie.current_env().get_recording_manager(p)) def test_recording_manager_create_if_absent(self): p = beam.Pipeline() self.assertFalse(ie.current_env().get_recording_manager(p)) self.assertTrue( ie.current_env().get_recording_manager(p, create_if_absent=True)) def test_evict_recording_manager(self): p = beam.Pipeline() self.assertFalse(ie.current_env().get_recording_manager(p)) self.assertTrue( ie.current_env().get_recording_manager(p, create_if_absent=True)) def test_describe_all_recordings(self): self.assertFalse(ie.current_env().describe_all_recordings()) p1 = beam.Pipeline() p2 = beam.Pipeline() ie.current_env().watch(locals()) ie.current_env().track_user_pipelines() rm1 = ie.current_env().get_recording_manager(p1, create_if_absent=True) rm2 = ie.current_env().get_recording_manager(p2, create_if_absent=True) description = ie.current_env().describe_all_recordings() self.assertTrue(description) expected_description = {p1: rm1.describe(), p2: rm2.describe()} self.assertDictEqual(description, expected_description) def test_get_empty_sql_chain(self): env = ie.InteractiveEnvironment() p = beam.Pipeline() chain = env.get_sql_chain(p) self.assertIsNotNone(chain) self.assertEqual(chain.nodes, {}) def test_get_sql_chain_with_nodes(self): env = ie.InteractiveEnvironment() p = beam.Pipeline() chain_with_node = env.get_sql_chain(p).append( SqlNode(output_name='name', source=p, query="query")) chain_got = env.get_sql_chain(p) self.assertIs(chain_with_node, chain_got) def test_get_sql_chain_setting_user_pipeline(self): env = ie.InteractiveEnvironment() p = beam.Pipeline() chain = env.get_sql_chain(p, set_user_pipeline=True) self.assertIs(chain.user_pipeline, p) def test_get_sql_chain_None_when_setting_multiple_user_pipelines(self): env = ie.InteractiveEnvironment() p = beam.Pipeline() chain = env.get_sql_chain(p, set_user_pipeline=True) p2 = beam.Pipeline() # Set the chain for a different pipeline. env.sql_chain[p2] = chain with self.assertRaises(ValueError): env.get_sql_chain(p2, set_user_pipeline=True) @patch( 'apache_beam.runners.interactive.interactive_environment.' 'assert_bucket_exists', return_value=None) def test_get_gcs_cache_dir_valid_path(self, mock_assert_bucket_exists): env = ie.InteractiveEnvironment() p = beam.Pipeline() cache_root = 'gs://test-cache-dir/' actual_cache_dir = env._get_gcs_cache_dir(p, cache_root) expected_cache_dir = 'gs://test-cache-dir/{}'.format(id(p)) self.assertEqual(actual_cache_dir, expected_cache_dir) def test_get_gcs_cache_dir_invalid_path(self): env = ie.InteractiveEnvironment() p = beam.Pipeline() cache_root = 'gs://' with self.assertRaises(ValueError): env._get_gcs_cache_dir(p, cache_root) def test_pcollection_computing_state(self): env = ie.InteractiveEnvironment() p = beam.Pipeline() pcoll1 = p | 'Create1' >> beam.Create([1]) pcoll2 = p | 'Create2' >> beam.Create([2]) self.assertFalse(env.is_pcollection_computing(pcoll1)) self.assertFalse(env.is_pcollection_computing(pcoll2)) self.assertEqual(env.computing_pcollections, set()) env.mark_pcollection_computing({pcoll1}) self.assertTrue(env.is_pcollection_computing(pcoll1)) self.assertFalse(env.is_pcollection_computing(pcoll2)) self.assertEqual(env.computing_pcollections, {pcoll1}) env.mark_pcollection_computing({pcoll2}) self.assertTrue(env.is_pcollection_computing(pcoll1)) self.assertTrue(env.is_pcollection_computing(pcoll2)) self.assertEqual(env.computing_pcollections, {pcoll1, pcoll2}) env.unmark_pcollection_computing({pcoll1}) self.assertFalse(env.is_pcollection_computing(pcoll1)) self.assertTrue(env.is_pcollection_computing(pcoll2)) self.assertEqual(env.computing_pcollections, {pcoll2}) env.unmark_pcollection_computing({pcoll2}) self.assertFalse(env.is_pcollection_computing(pcoll1)) self.assertFalse(env.is_pcollection_computing(pcoll2)) self.assertEqual(env.computing_pcollections, set()) def test_mark_unmark_empty(self): env = ie.InteractiveEnvironment() # Ensure no errors with empty sets env.mark_pcollection_computing(set()) self.assertEqual(env.computing_pcollections, set()) env.unmark_pcollection_computing(set()) self.assertEqual(env.computing_pcollections, set()) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/interactive_runner.py0000644000175100017510000003450615206113032026156 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A runner that allows running of Beam pipelines interactively. This module is experimental. No backwards-compatibility guarantees. """ # pytype: skip-file import logging from typing import Optional import apache_beam as beam from apache_beam import runners from apache_beam.options.pipeline_options import FlinkRunnerOptions from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import WorkerOptions from apache_beam.pipeline import PipelineVisitor from apache_beam.runners.direct import direct_runner from apache_beam.runners.interactive import background_caching_job from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import pipeline_instrument as inst from apache_beam.runners.interactive.dataproc.types import ClusterMetadata from apache_beam.runners.interactive.display import pipeline_graph from apache_beam.runners.interactive.options import capture_control from apache_beam.runners.interactive.utils import to_element_list from apache_beam.runners.interactive.utils import watch_sources from apache_beam.testing.test_stream_service import TestStreamServiceController # size of PCollection samples cached. SAMPLE_SIZE = 8 _LOGGER = logging.getLogger(__name__) class InteractiveRunner(runners.PipelineRunner): """An interactive runner for Beam Python pipelines. Allows interactively building and running Beam Python pipelines. """ def __init__( self, underlying_runner=None, render_option=None, skip_display=True, force_compute=True, blocking=True): """Constructor of InteractiveRunner. Args: underlying_runner: (runner.PipelineRunner) render_option: (str) this parameter decides how the pipeline graph is rendered. See display.pipeline_graph_renderer for available options. skip_display: (bool) whether to skip display operations when running the pipeline. Useful if running large pipelines when display is not needed. force_compute: (bool) whether sequential pipeline runs can use cached data of PCollections computed from the previous runs including show API invocation from interactive_beam module. If True, always run the whole pipeline and compute data for PCollections forcefully. If False, use available data and run minimum pipeline fragment to only compute data not available. blocking: (bool) whether the pipeline run should be blocking or not. """ self._underlying_runner = ( underlying_runner or direct_runner.DirectRunner()) if hasattr(self._underlying_runner, 'is_interactive'): self._underlying_runner.is_interactive() self._render_option = render_option self._in_session = False self._skip_display = skip_display self._force_compute = force_compute self._blocking = blocking def is_fnapi_compatible(self): # TODO(https://github.com/apache/beam/issues/19937): # return self._underlying_runner.is_fnapi_compatible() return False def default_pickle_library_override(self): """Delegates pickler override to the underlying runner.""" if hasattr(self._underlying_runner, 'default_pickle_library_override'): return self._underlying_runner.default_pickle_library_override() return super().default_pickle_library_override() def set_render_option(self, render_option): """Sets the rendering option. Args: render_option: (str) this parameter decides how the pipeline graph is rendered. See display.pipeline_graph_renderer for available options. """ self._render_option = render_option def start_session(self): """Start the session that keeps back-end managers and workers alive. """ if self._in_session: return enter = getattr(self._underlying_runner, '__enter__', None) if enter is not None: _LOGGER.info('Starting session.') self._in_session = True enter() else: _LOGGER.error('Keep alive not supported.') def end_session(self): """End the session that keeps backend managers and workers alive. """ if not self._in_session: return exit = getattr(self._underlying_runner, '__exit__', None) if exit is not None: self._in_session = False _LOGGER.info('Ending session.') exit(None, None, None) def apply(self, transform, pvalueish, options): # TODO(qinyeli, BEAM-646): Remove runner interception of apply. return self._underlying_runner.apply(transform, pvalueish, options) def run_pipeline(self, pipeline, options): if not ie.current_env().options.enable_recording_replay: capture_control.evict_captured_data() if self._force_compute: ie.current_env().evict_computed_pcollections() # Make sure that sources without a user reference are still cached. watch_sources(pipeline) user_pipeline = ie.current_env().user_pipeline(pipeline) from apache_beam.runners.portability.flink_runner import FlinkRunner if isinstance(self._underlying_runner, FlinkRunner): self.configure_for_flink(user_pipeline, options) pipeline_instrument = inst.build_pipeline_instrument(pipeline, options) # The user_pipeline analyzed might be None if the pipeline given has nothing # to be cached and tracing back to the user defined pipeline is impossible. # When it's None, there is no need to cache including the background # caching job and no result to track since no background caching job is # started at all. if user_pipeline: # Should use the underlying runner and run asynchronously. background_caching_job.attempt_to_run_background_caching_job( self._underlying_runner, user_pipeline, options) if (background_caching_job.has_source_to_cache(user_pipeline) and not background_caching_job.is_a_test_stream_service_running( user_pipeline)): streaming_cache_manager = ie.current_env().get_cache_manager( user_pipeline) # Only make the server if it doesn't exist already. if (streaming_cache_manager and not ie.current_env().get_test_stream_service_controller( user_pipeline)): def exception_handler(e): _LOGGER.error(str(e)) return True test_stream_service = TestStreamServiceController( streaming_cache_manager, exception_handler=exception_handler) test_stream_service.start() ie.current_env().set_test_stream_service_controller( user_pipeline, test_stream_service) pipeline_to_execute = beam.pipeline.Pipeline.from_runner_api( pipeline_instrument.instrumented_pipeline_proto(), self._underlying_runner, options) if ie.current_env().get_test_stream_service_controller(user_pipeline): endpoint = ie.current_env().get_test_stream_service_controller( user_pipeline).endpoint # TODO: make the StreamingCacheManager and TestStreamServiceController # constructed when the InteractiveEnvironment is imported. class TestStreamVisitor(PipelineVisitor): def visit_transform(self, transform_node): from apache_beam.testing.test_stream import TestStream if (isinstance(transform_node.transform, TestStream) and not transform_node.transform._events): transform_node.transform._endpoint = endpoint pipeline_to_execute.visit(TestStreamVisitor()) if not self._skip_display: a_pipeline_graph = pipeline_graph.PipelineGraph( pipeline_instrument.original_pipeline_proto, render_option=self._render_option) a_pipeline_graph.display_graph() main_job_result = PipelineResult( pipeline_to_execute.run(), pipeline_instrument) # In addition to this pipeline result setting, redundant result setting from # outer scopes are also recommended since the user_pipeline might not be # available from within this scope. if user_pipeline: ie.current_env().set_pipeline_result(user_pipeline, main_job_result) if self._blocking: main_job_result.wait_until_finish() if main_job_result.state is beam.runners.runner.PipelineState.DONE: # pylint: disable=bad-option-value ie.current_env().mark_pcollection_computed( pipeline_instrument.cached_pcolls) return main_job_result def configure_for_flink( self, user_pipeline: beam.Pipeline, options: PipelineOptions) -> None: """Configures the pipeline options for running a job with Flink. When running with a FlinkRunner, a job server started from an uber jar (locally built or remotely downloaded) hosting the beam_job_api will communicate with the Flink cluster located at the given flink_master in the pipeline options. """ clusters = ie.current_env().clusters if clusters.pipelines.get(user_pipeline, None): # Noop for a known pipeline using a known Dataproc cluster. return flink_master = self._strip_protocol_if_any( options.view_as(FlinkRunnerOptions).flink_master) cluster_metadata = clusters.default_cluster_metadata if flink_master == '[auto]': # Try to create/reuse a cluster when no flink_master is given. project_id = options.view_as(GoogleCloudOptions).project region = options.view_as(GoogleCloudOptions).region or 'us-central1' if project_id: if clusters.default_cluster_metadata: # Reuse the cluster name from default in case of a known cluster. cluster_metadata = ClusterMetadata( project_id=project_id, region=region, cluster_name=clusters.default_cluster_metadata.cluster_name) else: # Generate the metadata with a new unique cluster name. cluster_metadata = ClusterMetadata( project_id=project_id, region=region) # Add additional configurations. self._worker_options_to_cluster_metadata(options, cluster_metadata) # else use the default cluster metadata. elif flink_master in clusters.master_urls: cluster_metadata = clusters.cluster_metadata(flink_master) else: # Noop if a self-hosted Flink is in use. return if not cluster_metadata: return # Not even a default cluster to create/reuse, run Flink locally. dcm = clusters.create(cluster_metadata) # Side effects associated with the user_pipeline. clusters.pipelines[user_pipeline] = dcm dcm.pipelines.add(user_pipeline) self._configure_flink_options( options, clusters.DATAPROC_FLINK_VERSION, dcm.cluster_metadata.master_url) def _strip_protocol_if_any(self, flink_master: Optional[str]): if flink_master: parts = flink_master.split('://') if len(parts) > 1: return parts[1] return flink_master def _worker_options_to_cluster_metadata( self, options: PipelineOptions, cluster_metadata: ClusterMetadata): worker_options = options.view_as(WorkerOptions) if worker_options.subnetwork: cluster_metadata.subnetwork = worker_options.subnetwork if worker_options.num_workers: cluster_metadata.num_workers = worker_options.num_workers if worker_options.machine_type: cluster_metadata.machine_type = worker_options.machine_type def _configure_flink_options( self, options: PipelineOptions, flink_version: str, master_url: str): flink_options = options.view_as(FlinkRunnerOptions) flink_options.flink_version = flink_version # flink_options.flink_job_server_jar will be populated by the # apache_beam.utils.subprocess_server.JavaJarServer.path_to_beam_jar, # do not populate it explicitly. flink_options.flink_master = master_url class PipelineResult(beam.runners.runner.PipelineResult): """Provides access to information about a pipeline.""" def __init__(self, underlying_result, pipeline_instrument): """Constructor of PipelineResult. Args: underlying_result: (PipelineResult) the result returned by the underlying runner running the pipeline. pipeline_instrument: (PipelineInstrument) pipeline instrument describing the pipeline being executed with interactivity applied and related metadata including where the interactivity-backing cache lies. """ super().__init__(underlying_result.state) self._underlying_result = underlying_result self._pipeline_instrument = pipeline_instrument @property def state(self): return self._underlying_result.state def wait_until_finish(self): self._underlying_result.wait_until_finish() def get(self, pcoll, include_window_info=False): """Materializes the PCollection into a list. If include_window_info is True, then returns the elements as WindowedValues. Otherwise, return the element as itself. """ return list(self.read(pcoll, include_window_info)) def read(self, pcoll, include_window_info=False): """Reads the PCollection one element at a time from cache. If include_window_info is True, then returns the elements as WindowedValues. Otherwise, return the element as itself. """ key = self._pipeline_instrument.cache_key(pcoll) cache_manager = ie.current_env().get_cache_manager( self._pipeline_instrument.user_pipeline) if key and cache_manager.exists('full', key): coder = cache_manager.load_pcoder('full', key) reader, _ = cache_manager.read('full', key) return to_element_list(reader, coder, include_window_info) else: raise ValueError('PCollection not available, please run the pipeline.') def cancel(self): self._underlying_result.cancel() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/interactive_runner_test.py0000644000175100017510000006746015206113032027222 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for google3.pipeline.dataflow.python.interactive.interactive_runner. This module is experimental. No backwards-compatibility guarantees. """ # pytype: skip-file import sys import unittest import unittest.mock from typing import NamedTuple import pandas as pd import apache_beam as beam from apache_beam.dataframe.convert import to_dataframe from apache_beam.options.pipeline_options import FlinkRunnerOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import WorkerOptions from apache_beam.runners.direct import direct_runner from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import interactive_runner from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import DataprocClusterManager from apache_beam.runners.interactive.dataproc.types import ClusterMetadata from apache_beam.runners.interactive.testing.mock_env import isolated_env from apache_beam.runners.portability.flink_runner import FlinkRunner from apache_beam.runners.runner import PipelineRunner from apache_beam.testing.test_stream import TestStream from apache_beam.transforms.window import GlobalWindow from apache_beam.transforms.window import IntervalWindow from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.windowed_value import PaneInfo from apache_beam.utils.windowed_value import PaneInfoTiming from apache_beam.utils.windowed_value import WindowedValue def print_with_message(msg): def printer(elem): print(msg, elem) return elem return printer class Record(NamedTuple): name: str age: int height: int @isolated_env class InteractiveRunnerTest(unittest.TestCase): @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_basic(self): p = beam.Pipeline( runner=interactive_runner.InteractiveRunner( direct_runner.DirectRunner())) ib.watch({'p': p}) p.run().wait_until_finish() pc0 = ( p | 'read' >> beam.Create([1, 2, 3]) | 'Print1.1' >> beam.Map(print_with_message('Run1.1'))) pc = pc0 | 'Print1.2' >> beam.Map(print_with_message('Run1.2')) ib.watch(locals()) p.run().wait_until_finish() _ = pc | 'Print2' >> beam.Map(print_with_message('Run2')) p.run().wait_until_finish() _ = pc0 | 'Print3' >> beam.Map(print_with_message('Run3')) p.run().wait_until_finish() @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_wordcount(self): class WordExtractingDoFn(beam.DoFn): def process(self, element): text_line = element.strip() words = text_line.split() return words p = beam.Pipeline( runner=interactive_runner.InteractiveRunner( direct_runner.DirectRunner())) # Count the occurrences of each word. counts = ( p | beam.Create(['to be or not to be that is the question']) | 'split' >> beam.ParDo(WordExtractingDoFn()) | 'pair_with_one' >> beam.Map(lambda x: (x, 1)) | 'group' >> beam.GroupByKey() | 'count' >> beam.Map(lambda wordones: (wordones[0], sum(wordones[1])))) # Watch the local scope for Interactive Beam so that counts will be cached. ib.watch(locals()) # This is normally done in the interactive_utils when a transform is # applied but needs an IPython environment. So we manually run this here. ie.current_env().track_user_pipelines() result = p.run() result.wait_until_finish() actual = list(result.get(counts)) self.assertSetEqual( set(actual), set([ ('or', 1), ('that', 1), ('be', 2), ('is', 1), ('question', 1), ('to', 2), ('the', 1), ('not', 1), ])) # Truncate the precision to millis because the window coder uses millis # as units then gets upcast to micros. end_of_window = (GlobalWindow().max_timestamp().micros // 1000) * 1000 df_counts = ib.collect(counts, include_window_info=True, n=10) df_expected = pd.DataFrame({ 0: [e[0] for e in actual], 1: [e[1] for e in actual], 'event_time': [end_of_window for _ in actual], 'windows': [(GlobalWindow(), ) for _ in actual], 'pane_info': [ PaneInfo(True, True, PaneInfoTiming.ON_TIME, 0, 0) for _ in actual ] }, columns=[ 0, 1, 'event_time', 'windows', 'pane_info' ]) pd.testing.assert_frame_equal(df_expected, df_counts) actual_reified = result.get(counts, include_window_info=True) expected_reified = [ WindowedValue( e, Timestamp(micros=end_of_window), (GlobalWindow(), ), PaneInfo(True, True, PaneInfoTiming.ON_TIME, 0, 0)) for e in actual ] self.assertEqual(actual_reified, expected_reified) def test_streaming_wordcount(self): class WordExtractingDoFn(beam.DoFn): def process(self, element): text_line = element.strip() words = text_line.split() return words # Add the TestStream so that it can be cached. ib.options.recordable_sources.add(TestStream) p = beam.Pipeline( runner=interactive_runner.InteractiveRunner(), options=StandardOptions(streaming=True)) data = ( p | TestStream() .advance_watermark_to(0) .advance_processing_time(1) .add_elements(['to', 'be', 'or', 'not', 'to', 'be']) .advance_watermark_to(20) .advance_processing_time(1) .add_elements(['that', 'is', 'the', 'question']) .advance_watermark_to(30) .advance_processing_time(1) .advance_watermark_to(40) .advance_processing_time(1) .advance_watermark_to(50) .advance_processing_time(1) | beam.WindowInto(beam.window.FixedWindows(10))) # yapf: disable counts = ( data | 'split' >> beam.ParDo(WordExtractingDoFn()) | 'pair_with_one' >> beam.Map(lambda x: (x, 1)) | 'group' >> beam.GroupByKey() | 'count' >> beam.Map(lambda wordones: (wordones[0], sum(wordones[1])))) # Watch the local scope for Interactive Beam so that referenced PCollections # will be cached. ib.watch(locals()) # This is normally done in the interactive_utils when a transform is # applied but needs an IPython environment. So we manually run this here. ie.current_env().track_user_pipelines() # This tests that the data was correctly cached. pane_info = PaneInfo(True, True, PaneInfoTiming.UNKNOWN, 0, 0) expected_data_df = pd.DataFrame([ ('to', 0, [IntervalWindow(0, 10)], pane_info), ('be', 0, [IntervalWindow(0, 10)], pane_info), ('or', 0, [IntervalWindow(0, 10)], pane_info), ('not', 0, [IntervalWindow(0, 10)], pane_info), ('to', 0, [IntervalWindow(0, 10)], pane_info), ('be', 0, [IntervalWindow(0, 10)], pane_info), ('that', 20000000, [IntervalWindow(20, 30)], pane_info), ('is', 20000000, [IntervalWindow(20, 30)], pane_info), ('the', 20000000, [IntervalWindow(20, 30)], pane_info), ('question', 20000000, [IntervalWindow(20, 30)], pane_info) ], columns=[0, 'event_time', 'windows', 'pane_info']) # yapf: disable data_df = ib.collect(data, n=10, include_window_info=True) pd.testing.assert_frame_equal(expected_data_df, data_df) # This tests that the windowing was passed correctly so that all the data # is aggregated also correctly. pane_info = PaneInfo(True, False, PaneInfoTiming.ON_TIME, 0, 0) expected_counts_df = pd.DataFrame([ ('be', 2, 9999999, [IntervalWindow(0, 10)], pane_info), ('not', 1, 9999999, [IntervalWindow(0, 10)], pane_info), ('or', 1, 9999999, [IntervalWindow(0, 10)], pane_info), ('to', 2, 9999999, [IntervalWindow(0, 10)], pane_info), ('is', 1, 29999999, [IntervalWindow(20, 30)], pane_info), ('question', 1, 29999999, [IntervalWindow(20, 30)], pane_info), ('that', 1, 29999999, [IntervalWindow(20, 30)], pane_info), ('the', 1, 29999999, [IntervalWindow(20, 30)], pane_info), ], columns=[0, 1, 'event_time', 'windows', 'pane_info']) # yapf: disable counts_df = ib.collect(counts, n=8, include_window_info=True) # The group by key has no guarantee of order. So we post-process the DF by # sorting so we can test equality. sorted_counts_df = (counts_df .sort_values(['event_time', 0], ascending=True) .reset_index(drop=True)) # yapf: disable pd.testing.assert_frame_equal(expected_counts_df, sorted_counts_df) def test_session(self): class MockPipelineRunner(object): def __init__(self): self._in_session = False def __enter__(self): self._in_session = True def __exit__(self, exc_type, exc_val, exc_tb): self._in_session = False underlying_runner = MockPipelineRunner() runner = interactive_runner.InteractiveRunner(underlying_runner) runner.start_session() self.assertTrue(underlying_runner._in_session) runner.end_session() self.assertFalse(underlying_runner._in_session) @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') def test_mark_pcollection_completed_after_successful_run(self): with self.cell: # Cell 1 p = beam.Pipeline(interactive_runner.InteractiveRunner()) ib.watch({'p': p}) with self.cell: # Cell 2 # pylint: disable=bad-option-value init = p | 'Init' >> beam.Create(range(5)) with self.cell: # Cell 3 square = init | 'Square' >> beam.Map(lambda x: x * x) cube = init | 'Cube' >> beam.Map(lambda x: x**3) ib.watch(locals()) result = p.run() self.assertTrue(init in ie.current_env().computed_pcollections) self.assertEqual({0, 1, 2, 3, 4}, set(result.get(init))) self.assertTrue(square in ie.current_env().computed_pcollections) self.assertEqual({0, 1, 4, 9, 16}, set(result.get(square))) self.assertTrue(cube in ie.current_env().computed_pcollections) self.assertEqual({0, 1, 8, 27, 64}, set(result.get(cube))) @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_dataframes(self): p = beam.Pipeline( runner=interactive_runner.InteractiveRunner( direct_runner.DirectRunner())) data = p | beam.Create( [1, 2, 3]) | beam.Map(lambda x: beam.Row(square=x * x, cube=x * x * x)) df = to_dataframe(data) # Watch the local scope for Interactive Beam so that values will be cached. ib.watch(locals()) # This is normally done in the interactive_utils when a transform is # applied but needs an IPython environment. So we manually run this here. ie.current_env().track_user_pipelines() df_expected = pd.DataFrame({'square': [1, 4, 9], 'cube': [1, 8, 27]}) pd.testing.assert_frame_equal( df_expected, ib.collect(df, n=10).reset_index(drop=True)) @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_dataframes_with_grouped_index(self): p = beam.Pipeline( runner=interactive_runner.InteractiveRunner( direct_runner.DirectRunner())) data = [ Record('a', 20, 170), Record('a', 30, 170), Record('b', 22, 180), Record('c', 18, 150) ] aggregate = lambda df: df.groupby('height').mean(numeric_only=True) deferred_df = aggregate(to_dataframe(p | beam.Create(data))) df_expected = aggregate(pd.DataFrame(data)) # Watch the local scope for Interactive Beam so that values will be cached. ib.watch(locals()) # This is normally done in the interactive_utils when a transform is # applied but needs an IPython environment. So we manually run this here. ie.current_env().track_user_pipelines() pd.testing.assert_frame_equal(df_expected, ib.collect(deferred_df, n=10)) @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_dataframes_with_multi_index(self): p = beam.Pipeline( runner=interactive_runner.InteractiveRunner( direct_runner.DirectRunner())) data = [ Record('a', 20, 170), Record('a', 30, 170), Record('b', 22, 180), Record('c', 18, 150) ] aggregate = lambda df: df.groupby(['name', 'height']).mean() deferred_df = aggregate(to_dataframe(p | beam.Create(data))) df_input = pd.DataFrame(data) df_input.name = df_input.name.astype(pd.StringDtype()) df_expected = aggregate(df_input) # Watch the local scope for Interactive Beam so that values will be cached. ib.watch(locals()) # This is normally done in the interactive_utils when a transform is # applied but needs an IPython environment. So we manually run this here. ie.current_env().track_user_pipelines() pd.testing.assert_frame_equal(df_expected, ib.collect(deferred_df, n=10)) @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_dataframes_with_multi_index_get_result(self): p = beam.Pipeline( runner=interactive_runner.InteractiveRunner( direct_runner.DirectRunner())) data = [ Record('a', 20, 170), Record('a', 30, 170), Record('b', 22, 180), Record('c', 18, 150) ] aggregate = lambda df: df.groupby(['name', 'height']).mean()['age'] deferred_df = aggregate(to_dataframe(p | beam.Create(data))) df_input = pd.DataFrame(data) df_input.name = df_input.name.astype(pd.StringDtype()) df_expected = aggregate(df_input) # Watch the local scope for Interactive Beam so that values will be cached. ib.watch(locals()) # This is normally done in the interactive_utils when a transform is # applied but needs an IPython environment. So we manually run this here. ie.current_env().track_user_pipelines() pd.testing.assert_series_equal(df_expected, ib.collect(deferred_df, n=10)) @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_dataframes_same_cell_twice(self): p = beam.Pipeline( runner=interactive_runner.InteractiveRunner( direct_runner.DirectRunner())) data = p | beam.Create( [1, 2, 3]) | beam.Map(lambda x: beam.Row(square=x * x, cube=x * x * x)) df = to_dataframe(data) # Watch the local scope for Interactive Beam so that values will be cached. ib.watch(locals()) # This is normally done in the interactive_utils when a transform is # applied but needs an IPython environment. So we manually run this here. ie.current_env().track_user_pipelines() df_expected = pd.DataFrame({'square': [1, 4, 9], 'cube': [1, 8, 27]}) pd.testing.assert_series_equal( df_expected['square'], ib.collect(df['square'], n=10).reset_index(drop=True)) pd.testing.assert_series_equal( df_expected['cube'], ib.collect(df['cube'], n=10).reset_index(drop=True)) @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_dataframe_caching(self): # Create a pipeline that exercises the DataFrame API. This will also use # caching in the background. with self.cell: # Cell 1 p = beam.Pipeline(interactive_runner.InteractiveRunner()) ib.watch({'p': p}) with self.cell: # Cell 2 data = p | beam.Create([ 1, 2, 3 ]) | beam.Map(lambda x: beam.Row(square=x * x, cube=x * x * x)) with beam.dataframe.allow_non_parallel_operations(): df = to_dataframe(data).reset_index(drop=True) ib.collect(df) with self.cell: # Cell 3 df['output'] = df['square'] * df['cube'] ib.collect(df) with self.cell: # Cell 4 df['output'] = 0 ib.collect(df) # We use a trace through the graph to perform an isomorphism test. The end # output should look like a linear graph. This indicates that the dataframe # transform was correctly broken into separate pieces to cache. If caching # isn't enabled, all the dataframe computation nodes are connected to a # single shared node. trace = [] # Only look at the top-level transforms for the isomorphism. The test # doesn't care about the transform implementations, just the overall shape. class TopLevelTracer(beam.pipeline.PipelineVisitor): def _find_root_producer(self, node: beam.pipeline.AppliedPTransform): if node is None or not node.full_label: return None parent = self._find_root_producer(node.parent) if parent is None: return node return parent def _add_to_trace(self, node, trace): if '/' not in str(node): if node.inputs: producer = self._find_root_producer(node.inputs[0].producer) producer_name = producer.full_label if producer else '' trace.append((producer_name, node.full_label)) def visit_transform(self, node: beam.pipeline.AppliedPTransform): self._add_to_trace(node, trace) def enter_composite_transform( self, node: beam.pipeline.AppliedPTransform): self._add_to_trace(node, trace) p.visit(TopLevelTracer()) # Do the isomorphism test which states that the topological sort of the # graph yields a linear graph. trace_string = '\n'.join(str(t) for t in trace) prev_producer = '' for producer, consumer in trace: self.assertEqual(producer, prev_producer, trace_string) prev_producer = consumer @staticmethod def only_none_shall_pass(value): if value is None: return b'\0' else: raise RuntimeError("Should be using a more efficient coder.") @unittest.mock.patch.object( beam.coders.coders.FastPrimitivesCoder, 'encode', only_none_shall_pass) def test_defaults_to_efficient_cache(self): p = beam.Pipeline( runner=interactive_runner.InteractiveRunner( direct_runner.DirectRunner())) inputs = [1, 10, 100, 1000, 10000] big = ( p | beam.Create(inputs) | 'Explode' >> beam.FlatMap(lambda n: ("v_%s" % ix for ix in range(n)))) # Watch the local scope for Interactive Beam so that counts will be cached. ib.watch(locals()) # This is normally done in the interactive_utils when a transform is # applied but needs an IPython environment. So we manually run this here. ie.current_env().track_user_pipelines() result = p.run() result.wait_until_finish() self.assertEqual(len(result.get(big)), sum(inputs)) self.assertEqual( len(result.get(big, include_window_info=True)), sum(inputs)) cache_manager = ie.current_env().get_cache_manager( result._pipeline_instrument.user_pipeline) key = result._pipeline_instrument.cache_key(big) size = cache_manager.size('full', key) # Despite (highly redundant) windowing information, the cache is small. self.assertLess(size, sum(inputs)) def test_default_pickle_library_override_delegates(self): mock_underlying = unittest.mock.MagicMock(spec=PipelineRunner) mock_underlying.default_pickle_library_override.return_value = 'cloudpickle' runner = interactive_runner.InteractiveRunner( underlying_runner=mock_underlying) self.assertEqual(runner.default_pickle_library_override(), 'cloudpickle') def test_default_pickle_library_override_fallback(self): mock_underlying = unittest.mock.MagicMock(spec=PipelineRunner) del mock_underlying.default_pickle_library_override runner = interactive_runner.InteractiveRunner( underlying_runner=mock_underlying) # Should fallback to the base class implementation without crashing self.assertIsNone(runner.default_pickle_library_override()) @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') @isolated_env class ConfigForFlinkTest(unittest.TestCase): def setUp(self): self.current_env.options.cache_root = 'gs://fake' def tearDown(self): self.current_env.options.cache_root = None def test_create_a_new_cluster_for_a_new_pipeline(self): clusters = self.current_env.clusters runner = interactive_runner.InteractiveRunner( underlying_runner=FlinkRunner()) options = PipelineOptions(project='test-project', region='test-region') p = beam.Pipeline(runner=runner, options=options) runner.configure_for_flink(p, options) # Fetch the metadata and assert all side effects. meta = clusters.cluster_metadata(p) # The metadata should have all fields populated. self.assertEqual(meta.project_id, 'test-project') self.assertEqual(meta.region, 'test-region') self.assertTrue(meta.cluster_name.startswith('interactive-beam-')) self.assertTrue(meta.master_url.startswith('test-url')) self.assertEqual(meta.dashboard, 'test-dashboard') # The cluster is known now. self.assertIn(meta, clusters.dataproc_cluster_managers) self.assertIn(meta.master_url, clusters.master_urls) self.assertIn(p, clusters.pipelines) # The default cluster is updated to the created cluster. self.assertIs(meta, clusters.default_cluster_metadata) # The pipeline options is tuned for execution on the cluster. flink_options = options.view_as(FlinkRunnerOptions) self.assertEqual(flink_options.flink_master, meta.master_url) self.assertEqual( flink_options.flink_version, clusters.DATAPROC_FLINK_VERSION) def test_reuse_a_cluster_for_a_known_pipeline(self): clusters = self.current_env.clusters runner = interactive_runner.InteractiveRunner( underlying_runner=FlinkRunner()) options = PipelineOptions(project='test-project', region='test-region') p = beam.Pipeline(runner=runner, options=options) meta = ClusterMetadata(project_id='test-project', region='test-region') dcm = DataprocClusterManager(meta) # Configure the clusters so that the pipeline is known. clusters.pipelines[p] = dcm runner.configure_for_flink(p, options) # A known cluster is reused. tuned_meta = clusters.cluster_metadata(p) self.assertIs(tuned_meta, meta) def test_reuse_a_known_cluster_for_unknown_pipeline(self): clusters = self.current_env.clusters runner = interactive_runner.InteractiveRunner( underlying_runner=FlinkRunner()) options = PipelineOptions(project='test-project', region='test-region') p = beam.Pipeline(runner=runner, options=options) meta = ClusterMetadata(project_id='test-project', region='test-region') dcm = DataprocClusterManager(meta) # Configure the clusters so that the cluster is known. clusters.dataproc_cluster_managers[meta] = dcm clusters.set_default_cluster(meta) runner.configure_for_flink(p, options) # A known cluster is reused. tuned_meta = clusters.cluster_metadata(p) self.assertIs(tuned_meta, meta) # The pipeline is known. self.assertIn(p, clusters.pipelines) registered_dcm = clusters.pipelines[p] self.assertIn(p, registered_dcm.pipelines) def test_reuse_default_cluster_if_not_configured(self): clusters = self.current_env.clusters runner = interactive_runner.InteractiveRunner( underlying_runner=FlinkRunner()) options = PipelineOptions() # Pipeline is not configured to run on Cloud. p = beam.Pipeline(runner=runner, options=options) meta = ClusterMetadata(project_id='test-project', region='test-region') meta.master_url = 'test-url' meta.dashboard = 'test-dashboard' dcm = DataprocClusterManager(meta) # Configure the clusters so that a default cluster is known. clusters.dataproc_cluster_managers[meta] = dcm clusters.set_default_cluster(meta) runner.configure_for_flink(p, options) # The default cluster is used. tuned_meta = clusters.cluster_metadata(p) self.assertIs(tuned_meta, clusters.default_cluster_metadata) # The pipeline is known. self.assertIn(p, clusters.pipelines) registered_dcm = clusters.pipelines[p] self.assertIn(p, registered_dcm.pipelines) # The pipeline options is tuned for execution on the cluster. flink_options = options.view_as(FlinkRunnerOptions) self.assertEqual(flink_options.flink_master, tuned_meta.master_url) self.assertEqual( flink_options.flink_version, clusters.DATAPROC_FLINK_VERSION) def test_worker_options_to_cluster_metadata(self): clusters = self.current_env.clusters runner = interactive_runner.InteractiveRunner( underlying_runner=FlinkRunner()) options = PipelineOptions(project='test-project', region='test-region') worker_options = options.view_as(WorkerOptions) worker_options.num_workers = 2 worker_options.subnetwork = 'test-network' worker_options.machine_type = 'test-machine-type' p = beam.Pipeline(runner=runner, options=options) runner.configure_for_flink(p, options) configured_meta = clusters.cluster_metadata(p) self.assertEqual(configured_meta.num_workers, worker_options.num_workers) self.assertEqual(configured_meta.subnetwork, worker_options.subnetwork) self.assertEqual(configured_meta.machine_type, worker_options.machine_type) def test_configure_flink_options(self): clusters = self.current_env.clusters runner = interactive_runner.InteractiveRunner( underlying_runner=FlinkRunner()) options = PipelineOptions(project='test-project', region='test-region') p = beam.Pipeline(runner=runner, options=options) runner.configure_for_flink(p, options) flink_options = options.view_as(FlinkRunnerOptions) self.assertEqual( flink_options.flink_version, clusters.DATAPROC_FLINK_VERSION) self.assertTrue(flink_options.flink_master.startswith('test-url-')) def test_configure_flink_options_with_flink_version_overridden(self): clusters = self.current_env.clusters runner = interactive_runner.InteractiveRunner( underlying_runner=FlinkRunner()) options = PipelineOptions(project='test-project', region='test-region') flink_options = options.view_as(FlinkRunnerOptions) flink_options.flink_version = 'test-version' p = beam.Pipeline(runner=runner, options=options) runner.configure_for_flink(p, options) # The version is overridden to the flink version used by the EMR solution, # currently only 1: Cloud Dataproc. self.assertEqual( flink_options.flink_version, clusters.DATAPROC_FLINK_VERSION) def test_strip_http_protocol_from_flink_master(self): runner = interactive_runner.InteractiveRunner( underlying_runner=FlinkRunner()) stripped = runner._strip_protocol_if_any('https://flink-master') self.assertEqual('flink-master', stripped) def test_no_strip_from_flink_master(self): runner = interactive_runner.InteractiveRunner( underlying_runner=FlinkRunner()) stripped = runner._strip_protocol_if_any('flink-master') self.assertEqual('flink-master', stripped) def test_no_strip_from_non_flink_master(self): runner = interactive_runner.InteractiveRunner( underlying_runner=FlinkRunner()) stripped = runner._strip_protocol_if_any(None) self.assertIsNone(stripped) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.6384351 apache_beam-2.74.0/apache_beam/runners/interactive/messaging/0000755000175100017510000000000015206113160023645 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/messaging/__init__.py0000644000175100017510000000142015206113032025751 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/messaging/interactive_environment_inspector.py0000644000175100017510000002016215206113032033245 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Messaging mechanism to inspect the interactive environment. A singleton instance is accessible from interactive_environment.current_env().inspector. """ # pytype: skip-file import apache_beam as beam from apache_beam.runners.interactive.utils import as_json from apache_beam.runners.interactive.utils import obfuscate class InteractiveEnvironmentInspector(object): """Inspector that converts information of the current interactive environment including pipelines and pcollections into JSON data suitable for messaging with applications within/outside the Python kernel. The usage is always that the application side reads the inspectables or list_inspectables first then communicates back to the kernel and get_val for usage on the kernel side. """ def __init__(self, ignore_synthetic=True): self._inspectables = {} self._anonymous = {} self._inspectable_pipelines = set() self._ignore_synthetic = ignore_synthetic self._clusters = {} @property def inspectables(self): """Lists pipelines and pcollections assigned to variables as inspectables. """ self._inspectables = inspect(self._ignore_synthetic) return self._inspectables @property def inspectable_pipelines(self): """Returns a dictionary of all inspectable pipelines. The keys are stringified id of pipeline instances. This includes user defined pipeline assigned to variables and anonymous pipelines with inspectable PCollections. If a user defined pipeline is not within the returned dict, it can be considered out of scope, and all resources and memory states related to it should be released. """ _ = self.list_inspectables() return self._inspectable_pipelines @as_json def list_inspectables(self): """Lists inspectables in JSON format. When listing, pcollections are organized by the pipeline they belong to. If a pipeline is no longer assigned to a variable but its pcollections assigned to variables are still in scope, the pipeline will be given a name as 'anonymous_pipeline[id:$inMemoryId]'. The listing doesn't contain object values of the pipelines or pcollections. The obfuscated identifier can be used to trace back to those values in the kernel. The listing includes anonymous pipelines that are not assigned to variables but still containing inspectable PCollections. """ listing = {} pipelines = inspect_pipelines() for pipeline, name in pipelines.items(): metadata = meta(name, pipeline) listing[obfuscate(metadata)] = {'metadata': metadata, 'pcolls': {}} for identifier, inspectable in self.inspectables.items(): if inspectable['metadata']['type'] == 'pcollection': pipeline = inspectable['value'].pipeline if pipeline not in list(pipelines.keys()): pipeline_name = synthesize_pipeline_name(pipeline) pipelines[pipeline] = pipeline_name pipeline_metadata = meta(pipeline_name, pipeline) pipeline_identifier = obfuscate(pipeline_metadata) self._anonymous[pipeline_identifier] = { 'metadata': pipeline_metadata, 'value': pipeline } listing[pipeline_identifier] = { 'metadata': pipeline_metadata, 'pcolls': { identifier: inspectable['metadata'] } } else: pipeline_identifier = obfuscate(meta(pipelines[pipeline], pipeline)) listing[pipeline_identifier]['pcolls'][identifier] = inspectable[ 'metadata'] self._inspectable_pipelines = dict( (str(id(pipeline)), pipeline) for pipeline in pipelines) return listing def get_val(self, identifier): """Retrieves the in memory object itself by identifier. The retrieved object could be a pipeline or a pcollection. If the identifier is not recognized, return None. The identifier can refer to an anonymous pipeline and the object will still be retrieved. """ inspectable = self._inspectables.get(identifier, None) if inspectable: return inspectable['value'] inspectable = self._anonymous.get(identifier, None) if inspectable: return inspectable['value'] return None def get_pcoll_data(self, identifier, include_window_info=False): """Retrieves the json formatted PCollection data. If no PCollection value can be retieved from the given identifier, an empty json string will be returned. """ value = self.get_val(identifier) if isinstance(value, beam.pvalue.PCollection): from apache_beam.runners.interactive import interactive_beam as ib dataframe = ib.collect(value, include_window_info=include_window_info) return dataframe.to_json(orient='table') return {} @as_json def list_clusters(self): """Retrieves information for all clusters as a json. The json object maps a unique obfuscated identifier of a cluster to the corresponding cluster_name, project, region, master_url, dashboard, and pipelines. Furthermore, copies the mapping to self._clusters. """ from apache_beam.runners.interactive import interactive_environment as ie clusters = ie.current_env().clusters all_cluster_data = {} for meta, dcm in clusters.dataproc_cluster_managers.items(): all_cluster_data[obfuscate(meta)] = { 'cluster_name': meta.cluster_name, 'project': meta.project_id, 'region': meta.region, 'master_url': meta.master_url, 'dashboard': meta.dashboard, 'pipelines': [str(id(p)) for p in dcm.pipelines] } self._clusters = all_cluster_data return all_cluster_data def get_cluster_master_url(self, identifier: str) -> str: """Returns the master_url corresponding to the obfuscated identifier.""" return self._clusters[identifier]['master_url'] # Guaranteed to exist. def inspect(ignore_synthetic=True): """Inspects current interactive environment to track metadata and values of pipelines and pcollections. Each pipeline and pcollections tracked is given a unique identifier. """ from apache_beam.runners.interactive import interactive_environment as ie inspectables = {} for watching in ie.current_env().watching(): for name, value in watching: # Ignore synthetic vars created by Interactive Beam itself. if ignore_synthetic and name.startswith('synthetic_var_'): continue metadata = meta(name, value) identifier = obfuscate(metadata) if isinstance(value, (beam.pipeline.Pipeline, beam.pvalue.PCollection)): inspectables[identifier] = {'metadata': metadata, 'value': value} return inspectables def inspect_pipelines(): """Inspects current interactive environment to track all pipelines assigned to variables. The keys are pipeline objects and values are pipeline names. """ from apache_beam.runners.interactive import interactive_environment as ie pipelines = {} for watching in ie.current_env().watching(): for name, value in watching: if isinstance(value, beam.pipeline.Pipeline): pipelines[value] = name return pipelines def meta(name, val): """Generates meta data for the given name and value.""" return { 'name': name, 'inMemoryId': id(val), 'type': type(val).__name__.lower() } def synthesize_pipeline_name(val): """Synthesizes a pipeline name for the given pipeline object.""" return 'anonymous_pipeline[id:{}]'.format(id(val)) ././@PaxHeader0000000000000000000000000000020600000000000010213 xustar00112 path=apache_beam-2.74.0/apache_beam/runners/interactive/messaging/interactive_environment_inspector_test.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/messaging/interactive_environment_inspector_test.0000644000175100017510000002052415206113032033735 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for interactive_environment_inspector.""" # pytype: skip-file import json import sys import unittest import apache_beam as beam import apache_beam.runners.interactive.messaging.interactive_environment_inspector as inspector from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import interactive_runner as ir from apache_beam.runners.interactive.dataproc.types import ClusterMetadata from apache_beam.runners.interactive.testing.mock_env import isolated_env from apache_beam.runners.interactive.utils import obfuscate @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') @unittest.skipIf( sys.version_info < (3, 7), 'The tests require at least Python 3.7 to work.') @isolated_env class InteractiveEnvironmentInspectorTest(unittest.TestCase): def test_inspect(self): with self.cell: # Cell 1 pipeline = beam.Pipeline(ir.InteractiveRunner()) # Early watch the pipeline so that cell re-execution can be handled. ib.watch({'pipeline': pipeline}) # pylint: disable=bad-option-value pcoll = pipeline | 'Create' >> beam.Create(range(10)) with self.cell: # Cell 2 # Re-executes the line that created the pcoll causing the original # pcoll no longer inspectable. # pylint: disable=bad-option-value pcoll = pipeline | 'Create' >> beam.Create(range(10)) ib.watch(locals()) pipeline_metadata = inspector.meta('pipeline', pipeline) pcoll_metadata = inspector.meta('pcoll', pcoll) expected_inspectables = { obfuscate(pipeline_metadata): { 'metadata': pipeline_metadata, 'value': pipeline }, obfuscate(pcoll_metadata): { 'metadata': pcoll_metadata, 'value': pcoll } } for inspectable in inspector.inspect().items(): self.assertTrue(inspectable in expected_inspectables.items()) def test_inspect_pipelines(self): with self.cell: # Cell 1 pipeline_1 = beam.Pipeline(ir.InteractiveRunner()) pipeline_2 = beam.Pipeline(ir.InteractiveRunner()) with self.cell: # Cell 2 # Re-executes the line that created pipeline_1 causing the original # pipeline_1 no longer inspectable. pipeline_1 = beam.Pipeline(ir.InteractiveRunner()) ib.watch(locals()) expected_inspectable_pipelines = { pipeline_1: 'pipeline_1', pipeline_2: 'pipeline_2' } for inspectable_pipeline in inspector.inspect_pipelines().items(): self.assertTrue( inspectable_pipeline in expected_inspectable_pipelines.items()) pipeline_1_metadata = inspector.meta('pipeline_1', pipeline_1) pipeline_2_metadata = inspector.meta('pipeline_2', pipeline_2) expected_inspectables = { obfuscate(pipeline_2_metadata): { 'metadata': pipeline_2_metadata, 'value': pipeline_2 }, obfuscate(pipeline_1_metadata): { 'metadata': pipeline_1_metadata, 'value': pipeline_1 } } for inspectable in inspector.inspect().items(): self.assertTrue(inspectable in expected_inspectables.items()) def test_list_inspectables(self): with self.cell: # Cell 1 pipeline = beam.Pipeline(ir.InteractiveRunner()) # pylint: disable=bad-option-value pcoll_1 = pipeline | 'Create' >> beam.Create(range(10)) pcoll_2 = pcoll_1 | 'Square' >> beam.Map(lambda x: x * x) with self.cell: # Cell 2 # Re-executes the line that created pipeline causing the original # pipeline become an anonymous pipeline that is still inspectable because # its pcoll_1 and pcoll_2 are still inspectable. pipeline = beam.Pipeline(ir.InteractiveRunner()) ib.watch(locals()) anonymous_pipeline_name = inspector.synthesize_pipeline_name( pcoll_1.pipeline) anonymous_pipeline_metadata = inspector.meta( anonymous_pipeline_name, pcoll_1.pipeline) pipeline_metadata = inspector.meta('pipeline', pipeline) pcoll_1_metadata = inspector.meta('pcoll_1', pcoll_1) pcoll_2_metadata = inspector.meta('pcoll_2', pcoll_2) expected_inspectable_list = { obfuscate(pipeline_metadata): { 'metadata': pipeline_metadata, 'pcolls': {} }, obfuscate(anonymous_pipeline_metadata): { 'metadata': anonymous_pipeline_metadata, 'pcolls': { obfuscate(pcoll_1_metadata): pcoll_1_metadata, obfuscate(pcoll_2_metadata): pcoll_2_metadata } } } ins = inspector.InteractiveEnvironmentInspector() actual_listings = ins.list_inspectables() self.assertEqual(actual_listings, json.dumps(expected_inspectable_list)) def test_get_val(self): with self.cell: # Cell 1 pipeline = beam.Pipeline(ir.InteractiveRunner()) # pylint: disable=bad-option-value pcoll = pipeline | 'Create' >> beam.Create(range(10)) with self.cell: # Cell 2 # Re-executes the line that created pipeline causing the original # pipeline become an anonymous pipeline that is still inspectable because # its pcoll is still inspectable. pipeline = beam.Pipeline(ir.InteractiveRunner()) ib.watch(locals()) ins = inspector.InteractiveEnvironmentInspector() _ = ins.list_inspectables() pipeline_identifier = obfuscate(inspector.meta('pipeline', pipeline)) self.assertIs(ins.get_val(pipeline_identifier), pipeline) pcoll_identifier = obfuscate(inspector.meta('pcoll', pcoll)) self.assertIs(ins.get_val(pcoll_identifier), pcoll) anonymous_pipeline_name = inspector.synthesize_pipeline_name(pcoll.pipeline) anonymous_pipeline_identifier = obfuscate( inspector.meta(anonymous_pipeline_name, pcoll.pipeline)) self.assertIs(ins.get_val(anonymous_pipeline_identifier), pcoll.pipeline) def test_get_pcoll_data(self): pipeline = beam.Pipeline(ir.InteractiveRunner()) # pylint: disable=bad-option-value pcoll = pipeline | 'Create' >> beam.Create(list(range(10))) counts = pcoll | beam.combiners.Count.PerElement() ib.watch(locals()) ie.current_env().track_user_pipelines() counts_identifier = obfuscate(inspector.meta('counts', counts)) ins = inspector.InteractiveEnvironmentInspector() _ = ins.list_inspectables() actual_counts_pcoll_data = ins.get_pcoll_data(counts_identifier) expected_counts_pcoll_data = ib.collect( counts, n=10).to_json(orient='table') self.assertEqual(actual_counts_pcoll_data, expected_counts_pcoll_data) actual_counts_with_window_info = ins.get_pcoll_data(counts_identifier, True) expected_counts_with_window_info = ib.collect( counts, include_window_info=True).to_json(orient='table') self.assertEqual( actual_counts_with_window_info, expected_counts_with_window_info) def test_list_clusters(self): self.current_env.options.cache_root = 'gs://fake' meta = ClusterMetadata(project_id='project') dcm = self.current_env.clusters.create(meta) p = beam.Pipeline() dcm.pipelines.add(p) self.current_env.clusters.pipelines[p] = dcm cluster_id = obfuscate(meta) self.assertEqual({ cluster_id: { 'cluster_name': meta.cluster_name, 'project': meta.project_id, 'region': meta.region, 'master_url': meta.master_url, 'dashboard': meta.dashboard, 'pipelines': [str(id(p)) for p in dcm.pipelines] } }, json.loads(self.current_env.inspector.list_clusters())) self.current_env.options.cache_root = None if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/non_interactive_runner_test.py0000644000175100017510000002442615206113032030067 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for interactive utilities without explicitly using InteractiveRunner. """ # pytype: skip-file import importlib import sys import unittest from collections import defaultdict from typing import NamedTuple import pandas as pd import apache_beam as beam from apache_beam.dataframe.convert import to_dataframe from apache_beam.runners.direct import direct_runner from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive.testing.mock_env import isolated_env from apache_beam.transforms.window import GlobalWindow from apache_beam.utils.windowed_value import PaneInfo from apache_beam.utils.windowed_value import PaneInfoTiming def print_with_message(msg): def printer(elem): print(msg, elem) return elem return printer class Record(NamedTuple): name: str age: int height: int _side_effects = defaultdict(int) def cause_side_effect(elem): mod = importlib.import_module(__name__) mod._side_effects[elem] += 1 return elem def count_side_effects(elem): mod = importlib.import_module(__name__) return mod._side_effects[elem] def clear_side_effect(): mod = importlib.import_module(__name__) mod._side_effects.clear() @isolated_env class NonInteractiveRunnerTest(unittest.TestCase): @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_basic(self): clear_side_effect() # This test relies on the pipeline cache being populated. Prism doesn't # consistently populate this cache, forcing FnApiRunner p = beam.Pipeline('FnApiRunner') # Initial collection runs the pipeline. pcoll1 = p | beam.Create(['a', 'b', 'c']) | beam.Map(cause_side_effect) collected1 = ib.collect(pcoll1) self.assertEqual(set(collected1[0]), set(['a', 'b', 'c'])) self.assertEqual(count_side_effects('a'), 1) # Collecting the PCollection again uses the cache. collected1again = ib.collect(pcoll1) self.assertEqual(set(collected1again[0]), set(['a', 'b', 'c'])) self.assertEqual(count_side_effects('a'), 1) # Using the PCollection uses the cache. pcoll2 = pcoll1 | beam.Map(str.upper) collected2 = ib.collect(pcoll2) self.assertEqual(set(collected2[0]), set(['A', 'B', 'C'])) self.assertEqual(count_side_effects('a'), 1) # Force re-computation. collected2 = ib.collect(pcoll2, force_compute=True) self.assertEqual(set(collected2[0]), set(['A', 'B', 'C'])) self.assertEqual(count_side_effects('a'), 2) @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_multiple_collect(self): clear_side_effect() p = beam.Pipeline(direct_runner.DirectRunner()) # Initial collection runs the pipeline. pcollA = p | 'A' >> beam.Create(['a']) | 'As' >> beam.Map(cause_side_effect) pcollB = p | 'B' >> beam.Create(['b']) | 'Bs' >> beam.Map(cause_side_effect) collectedA, collectedB = ib.collect(pcollA, pcollB) self.assertEqual(set(collectedA[0]), set(['a'])) self.assertEqual(set(collectedB[0]), set(['b'])) self.assertEqual(count_side_effects('a'), 1) self.assertEqual(count_side_effects('b'), 1) # Collecting the PCollection again uses the cache. collectedA, collectedB = ib.collect(pcollA, pcollB) self.assertEqual(set(collectedA[0]), set(['a'])) self.assertEqual(set(collectedB[0]), set(['b'])) self.assertEqual(count_side_effects('a'), 1) self.assertEqual(count_side_effects('b'), 1) # Using the PCollection uses the cache. pcollAA = pcollA | beam.Map( lambda x: 2 * x) | 'AAs' >> beam.Map(cause_side_effect) collectedA, collectedB, collectedAA = ib.collect(pcollA, pcollB, pcollAA) self.assertEqual(set(collectedA[0]), set(['a'])) self.assertEqual(set(collectedB[0]), set(['b'])) self.assertEqual(set(collectedAA[0]), set(['aa'])) self.assertEqual(count_side_effects('a'), 1) self.assertEqual(count_side_effects('b'), 1) self.assertEqual(count_side_effects('aa'), 1) # Duplicates are only computed once. pcollBB = pcollB | beam.Map( lambda x: 2 * x) | 'BBs' >> beam.Map(cause_side_effect) collectedAA, collectedAAagain, collectedBB, collectedBBagain = ib.collect( pcollAA, pcollAA, pcollBB, pcollBB) self.assertEqual(set(collectedAA[0]), set(['aa'])) self.assertEqual(set(collectedAAagain[0]), set(['aa'])) self.assertEqual(set(collectedBB[0]), set(['bb'])) self.assertEqual(set(collectedBBagain[0]), set(['bb'])) self.assertEqual(count_side_effects('aa'), 1) self.assertEqual(count_side_effects('bb'), 1) @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_wordcount(self): class WordExtractingDoFn(beam.DoFn): def process(self, element): text_line = element.strip() words = text_line.split() return words p = beam.Pipeline(runner=direct_runner.DirectRunner()) # Count the occurrences of each word. counts = ( p | beam.Create(['to be or not to be that is the question']) | 'split' >> beam.ParDo(WordExtractingDoFn()) | 'pair_with_one' >> beam.Map(lambda x: (x, 1)) | 'group' >> beam.GroupByKey() | 'count' >> beam.Map(lambda wordones: (wordones[0], sum(wordones[1])))) actual = ib.collect(counts) self.assertSetEqual( set(zip(actual[0], actual[1])), set([ ('or', 1), ('that', 1), ('be', 2), ('is', 1), ('question', 1), ('to', 2), ('the', 1), ('not', 1), ])) # Truncate the precision to millis because the window coder uses millis # as units then gets upcast to micros. end_of_window = (GlobalWindow().max_timestamp().micros // 1000) * 1000 df_counts = ib.collect(counts, include_window_info=True, n=10) df_expected = pd.DataFrame({ 0: list(actual[0]), 1: list(actual[1]), 'event_time': [end_of_window] * len(actual), 'windows': [[GlobalWindow()]] * len(actual), 'pane_info': [PaneInfo(True, True, PaneInfoTiming.ON_TIME, 0, 0)] * len(actual) }, ) pd.testing.assert_frame_equal(df_expected, df_counts) @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_dataframes(self): p = beam.Pipeline(runner=direct_runner.DirectRunner()) data = p | beam.Create( [1, 2, 3]) | beam.Map(lambda x: beam.Row(square=x * x, cube=x * x * x)) df = to_dataframe(data) df_expected = pd.DataFrame({'square': [1, 4, 9], 'cube': [1, 8, 27]}) pd.testing.assert_frame_equal( df_expected, ib.collect(df, n=10).reset_index(drop=True)) @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_dataframes_with_grouped_index(self): p = beam.Pipeline(runner=direct_runner.DirectRunner()) data = [ Record('a', 20, 170), Record('a', 30, 170), Record('b', 22, 180), Record('c', 18, 150) ] aggregate = lambda df: df.groupby('height').mean(numeric_only=True) deferred_df = aggregate(to_dataframe(p | beam.Create(data))) df_expected = aggregate(pd.DataFrame(data)) pd.testing.assert_frame_equal(df_expected, ib.collect(deferred_df, n=10)) @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_dataframes_with_multi_index(self): p = beam.Pipeline(runner=direct_runner.DirectRunner()) data = [ Record('a', 20, 170), Record('a', 30, 170), Record('b', 22, 180), Record('c', 18, 150) ] aggregate = lambda df: df.groupby(['name', 'height']).mean() deferred_df = aggregate(to_dataframe(p | beam.Create(data))) df_input = pd.DataFrame(data) df_input.name = df_input.name.astype(pd.StringDtype()) df_expected = aggregate(df_input) pd.testing.assert_frame_equal(df_expected, ib.collect(deferred_df, n=10)) @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_dataframes_same_cell_twice(self): p = beam.Pipeline(runner=direct_runner.DirectRunner()) data = p | beam.Create( [1, 2, 3]) | beam.Map(lambda x: beam.Row(square=x * x, cube=x * x * x)) df = to_dataframe(data) df_expected = pd.DataFrame({'square': [1, 4, 9], 'cube': [1, 8, 27]}) pd.testing.assert_series_equal( df_expected['square'], ib.collect(df['square'], n=10).reset_index(drop=True)) pd.testing.assert_series_equal( df_expected['cube'], ib.collect(df['cube'], n=10).reset_index(drop=True)) @unittest.skipIf(sys.platform == "win32", "[BEAM-10627]") def test_new_runner_and_options(self): class MyRunner(beam.runners.PipelineRunner): run_count = 0 @classmethod def run_pipeline(cls, pipeline, options): assert options._all_options['my_option'] == 123 cls.run_count += 1 return direct_runner.DirectRunner().run_pipeline(pipeline, options) clear_side_effect() p = beam.Pipeline(direct_runner.DirectRunner()) # Initial collection runs the pipeline. pcoll1 = p | beam.Create(['a', 'b', 'c']) | beam.Map(cause_side_effect) collected1 = ib.collect(pcoll1) self.assertEqual(set(collected1[0]), set(['a', 'b', 'c'])) self.assertEqual(count_side_effects('a'), 1) # Using the PCollection uses the cache with a different runner and options. pcoll2 = pcoll1 | beam.Map(str.upper) collected2 = ib.collect( pcoll2, runner=MyRunner(), options=beam.options.pipeline_options.PipelineOptions(my_option=123)) self.assertEqual(set(collected2[0]), set(['A', 'B', 'C'])) self.assertEqual(count_side_effects('a'), 1) self.assertEqual(MyRunner.run_count, 1) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.6444356 apache_beam-2.74.0/apache_beam/runners/interactive/options/0000755000175100017510000000000015206113160023363 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/options/__init__.py0000644000175100017510000000142015206113032025467 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/options/capture_control.py0000644000175100017510000000531515206113032027142 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module to control how Interactive Beam captures data from sources for deterministic replayable PCollection evaluation and pipeline runs. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file import logging from datetime import timedelta from apache_beam.io.gcp.pubsub import ReadFromPubSub from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.options import capture_limiters _LOGGER = logging.getLogger(__name__) class CaptureControl(object): """Options and their utilities that controls how Interactive Beam captures deterministic replayable data from sources.""" def __init__(self): self._enable_capture_replay = True self._capturable_sources = { ReadFromPubSub, } # yapf: disable self._capture_duration = timedelta(seconds=60) self._capture_size_limit = 1e9 self._test_limiters = None def limiters(self) -> list['capture_limiters.Limiter']: # noqa: F821 if self._test_limiters: return self._test_limiters return [ capture_limiters.SizeLimiter(self._capture_size_limit), capture_limiters.DurationLimiter(self._capture_duration) ] def set_limiters_for_test( self, limiters: list['capture_limiters.Limiter']) -> None: # noqa: F821 self._test_limiters = limiters def evict_captured_data(pipeline=None): """Evicts all deterministic replayable data that have been captured by Interactive Beam for the given pipeline. If no pipeline is specified, evicts for all user defined pipelines. In future PCollection evaluation/visualization and pipeline runs, Interactive Beam will capture fresh data.""" if ie.current_env().options.enable_recording_replay: _LOGGER.info( 'You have requested Interactive Beam to evict all recorded ' 'data that could be deterministically replayed among multiple ' 'pipeline runs.') ie.current_env().cleanup(pipeline) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/options/capture_control_test.py0000644000175100017510000001605115206113032030200 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.runners.interactive.options.capture_control.""" # pytype: skip-file import unittest from unittest.mock import patch import apache_beam as beam from apache_beam import coders from apache_beam.portability.api import beam_interactive_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners import runner from apache_beam.runners.interactive import background_caching_job as bcj from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import interactive_runner from apache_beam.runners.interactive.caching.streaming_cache import StreamingCache from apache_beam.runners.interactive.options import capture_control from apache_beam.runners.interactive.options import capture_limiters from apache_beam.testing.test_stream_service import TestStreamServiceController def _build_an_empty_streaming_pipeline(): from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions pipeline_options = PipelineOptions() pipeline_options.view_as(StandardOptions).streaming = True p = beam.Pipeline( interactive_runner.InteractiveRunner(), options=pipeline_options) ib.watch({'pipeline': p}) return p def _fake_a_running_test_stream_service(pipeline): class FakeReader: def read_multiple(self): yield 1 test_stream_service = TestStreamServiceController(FakeReader()) test_stream_service.start() ie.current_env().set_test_stream_service_controller( pipeline, test_stream_service) @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') class CaptureControlTest(unittest.TestCase): def setUp(self): ie.new_env() @patch( 'apache_beam.runners.interactive.background_caching_job' '.BackgroundCachingJob.cancel') @patch( 'apache_beam.testing.test_stream_service.TestStreamServiceController' '.stop') def test_capture_control_evict_captured_data( self, mocked_test_stream_service_stop, mocked_background_caching_job_cancel): p = _build_an_empty_streaming_pipeline() ie.current_env().track_user_pipelines() self.assertFalse(ie.current_env().tracked_user_pipelines == set()) background_caching_job = bcj.BackgroundCachingJob( runner.PipelineResult(runner.PipelineState.RUNNING), limiters=[]) ie.current_env().set_background_caching_job(p, background_caching_job) _fake_a_running_test_stream_service(p) # Fake the canceling state of the main job. background_caching_job._pipeline_result = runner.PipelineResult( runner.PipelineState.CANCELLING) self.assertIsNotNone(ie.current_env().get_test_stream_service_controller(p)) ie.current_env().set_cached_source_signature(p, 'a signature') ie.current_env().mark_pcollection_computed(['fake_pcoll']) capture_control.evict_captured_data() mocked_background_caching_job_cancel.assert_called() mocked_test_stream_service_stop.assert_called_once() # Neither timer nor capture size limit is reached, thus, the cancelling # main job's background caching job is not considered as done. self.assertFalse(background_caching_job.is_done()) self.assertIsNone(ie.current_env().get_test_stream_service_controller(p)) self.assertTrue(ie.current_env().computed_pcollections == set()) self.assertTrue(ie.current_env().get_cached_source_signature(p) == set()) def test_capture_size_limit_not_reached_when_no_cache(self): self.assertEqual(len(ie.current_env()._cache_managers), 0) limiter = capture_limiters.SizeLimiter(1) self.assertFalse(limiter.is_triggered()) def test_capture_size_limit_not_reached_when_no_file(self): cache = StreamingCache(cache_dir=None) self.assertFalse(cache.exists('my_label')) ie.current_env().set_cache_manager(cache, 'dummy pipeline') limiter = capture_limiters.SizeLimiter(1) self.assertFalse(limiter.is_triggered()) def test_capture_size_limit_not_reached_when_file_size_under_limit(self): ib.options.capture_size_limit = 100 cache = StreamingCache(cache_dir=None) # Build a sink object to track the label as a capture in the test. cache.sink(['my_label'], is_capture=True) cache.write([beam_interactive_api_pb2.TestStreamFileRecord()], 'my_label') self.assertTrue(cache.exists('my_label')) ie.current_env().set_cache_manager(cache, 'dummy pipeline') limiter = capture_limiters.SizeLimiter(ib.options.capture_size_limit) self.assertFalse(limiter.is_triggered()) def test_capture_size_limit_reached_when_file_size_above_limit(self): ib.options.capture_size_limit = 1 cache = StreamingCache(cache_dir=None) cache.sink(['my_label'], is_capture=True) cache.write([ beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload. TimestampedElement( encoded_element=coders.FastPrimitivesCoder().encode( 'a'), timestamp=0) ]))) ], 'my_label') self.assertTrue(cache.exists('my_label')) p = _build_an_empty_streaming_pipeline() ie.current_env().set_cache_manager(cache, p) limiter = capture_limiters.SizeLimiter(1) self.assertTrue(limiter.is_triggered()) def test_timer_terminates_capture_size_checker(self): p = _build_an_empty_streaming_pipeline() class FakeLimiter(capture_limiters.Limiter): def __init__(self): self.trigger = False def is_triggered(self): return self.trigger limiter = FakeLimiter() background_caching_job = bcj.BackgroundCachingJob( runner.PipelineResult(runner.PipelineState.CANCELLING), limiters=[limiter]) ie.current_env().set_background_caching_job(p, background_caching_job) self.assertFalse(background_caching_job.is_done()) limiter.trigger = True self.assertTrue(background_caching_job.is_done()) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/options/capture_limiters.py0000644000175100017510000001172615206113032027315 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module to condition how Interactive Beam stops capturing data. For internal use only; no backwards-compatibility guarantees. """ import datetime import threading from typing import Any import pandas as pd from apache_beam.portability.api import beam_interactive_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.utils.windowed_value import WindowedValue class Limiter: """Limits an aspect of the caching layer.""" def is_triggered(self) -> bool: """Returns True if the limiter has triggered, and caching should stop.""" raise NotImplementedError class ElementLimiter(Limiter): """A `Limiter` that limits reading from cache based on some property of an element. """ def update(self, e: Any) -> None: # noqa: F821 """Update the internal state based on some property of an element. This is executed on every element that is read from cache. """ raise NotImplementedError class SizeLimiter(Limiter): """Limits the cache size to a specified byte limit.""" def __init__(self, size_limit: int): self._size_limit = size_limit def is_triggered(self): total_capture_size = 0 ie.current_env().track_user_pipelines() for user_pipeline in ie.current_env().tracked_user_pipelines: cache_manager = ie.current_env().get_cache_manager(user_pipeline) if hasattr(cache_manager, 'capture_size'): total_capture_size += cache_manager.capture_size return total_capture_size >= self._size_limit class DurationLimiter(Limiter): """Limits the duration of the capture.""" def __init__( self, duration_limit: datetime.timedelta # noqa: F821 ): self._duration_limit = duration_limit self._timer = threading.Timer(duration_limit.total_seconds(), self._trigger) self._timer.daemon = True self._triggered = False self._timer.start() def _trigger(self): self._triggered = True def is_triggered(self): return self._triggered class CountLimiter(ElementLimiter): """Limits by counting the number of elements seen.""" def __init__(self, max_count): self._max_count = max_count self._count = 0 def update(self, e): # A TestStreamFileRecord can contain many elements at once. If e is a file # record, then count the number of elements in the bundle. if isinstance(e, beam_interactive_api_pb2.TestStreamFileRecord): if not e.recorded_event.element_event: return self._count += len(e.recorded_event.element_event.elements) # Otherwise, count everything else but the header of the file since it is # not an element. elif not isinstance(e, beam_interactive_api_pb2.TestStreamFileHeader): # When elements are DataFrames, we want the output to be constrained by # how many rows we have read, not how many DataFrames we have read. if isinstance(e, WindowedValue) and isinstance(e.value, pd.DataFrame): self._count += len(e.value) else: self._count += 1 def is_triggered(self): return self._count >= self._max_count class ProcessingTimeLimiter(ElementLimiter): """Limits by how long the ProcessingTime passed in the element stream. Reads all elements from the timespan [start, start + duration). This measures the duration from the first element in the stream. Each subsequent element has a delta "advance_duration" that moves the internal clock forward. This triggers when the duration from the internal clock and the start exceeds the given duration. """ def __init__(self, max_duration_secs): """Initialize the ProcessingTimeLimiter.""" self._max_duration_us = max_duration_secs * 1e6 self._start_us = 0 self._cur_time_us = 0 def update(self, e): # Only look at TestStreamFileRecords which hold the processing time. if not isinstance(e, beam_runner_api_pb2.TestStreamPayload.Event): return if not e.HasField('processing_time_event'): return if self._start_us == 0: self._start_us = e.processing_time_event.advance_duration self._cur_time_us += e.processing_time_event.advance_duration def is_triggered(self): return self._cur_time_us - self._start_us >= self._max_duration_us ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/options/capture_limiters_test.py0000644000175100017510000000437315206113032030354 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest import pandas as pd from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.interactive.options.capture_limiters import CountLimiter from apache_beam.runners.interactive.options.capture_limiters import ProcessingTimeLimiter from apache_beam.utils.windowed_value import WindowedValue class CaptureLimitersTest(unittest.TestCase): def test_count_limiter(self): limiter = CountLimiter(5) for e in range(4): limiter.update(e) self.assertFalse(limiter.is_triggered()) limiter.update(4) self.assertTrue(limiter.is_triggered()) def test_count_limiter_with_dataframes(self): limiter = CountLimiter(5) # Test that empty dataframes don't count. for _ in range(10): df = WindowedValue(pd.DataFrame(), 0, []) limiter.update(df) self.assertFalse(limiter.is_triggered()) df = WindowedValue(pd.DataFrame({'col': list(range(10))}), 0, []) limiter.update(df) self.assertTrue(limiter.is_triggered()) def test_processing_time_limiter(self): limiter = ProcessingTimeLimiter(max_duration_secs=2) e = beam_runner_api_pb2.TestStreamPayload.Event() e.processing_time_event.advance_duration = int(1 * 1e6) limiter.update(e) self.assertFalse(limiter.is_triggered()) e = beam_runner_api_pb2.TestStreamPayload.Event() e.processing_time_event.advance_duration = int(2 * 1e6) limiter.update(e) self.assertTrue(limiter.is_triggered()) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/options/interactive_options.py0000644000175100017510000000324015206113032030022 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module to expose options that control how Interactive Beam works. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file from dateutil import tz from apache_beam.runners.interactive.options import capture_control class InteractiveOptions(object): """An intermediate facade to query and configure options that guide how Interactive Beam works.""" def __init__(self): self._capture_control = capture_control.CaptureControl() self._display_timestamp_format = '%Y-%m-%d %H:%M:%S.%f%z' self._display_timezone = tz.tzlocal() self._cache_root = None def __repr__(self): options_str = '\n'.join( '{} = {}'.format(k, getattr(self, k)) for k in dir(self) if k[0] != '_' and k != 'capture_control') return 'interactive_beam.options:\n{}'.format(options_str) @property def capture_control(self): return self._capture_control ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/pipeline_fragment.py0000644000175100017510000002714415206113032025740 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module to build pipeline fragment that produces given PCollections. For internal use only; no backwards-compatibility guarantees. """ import apache_beam as beam from apache_beam.pipeline import AppliedPTransform from apache_beam.pipeline import PipelineVisitor from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import pipeline_instrument as instr from apache_beam.testing.test_stream import TestStream class PipelineFragment(object): """A fragment of a pipeline definition. A pipeline fragment is built from the original pipeline definition to include only PTransforms that are necessary to produce the given PCollections. """ def __init__(self, pcolls, options=None, runner=None): """Constructor of PipelineFragment. Args: pcolls: (List[PCollection]) a list of PCollections to build pipeline fragment for. options: (PipelineOptions) the pipeline options for the implicit pipeline run. runner: (Runner) the pipeline runner for the implicit pipeline run. """ assert len(pcolls) > 0, ( 'Need at least 1 PCollection as the target data to build a pipeline ' 'fragment that produces it.') for pcoll in pcolls: assert isinstance(pcoll, beam.pvalue.PCollection), ( '{} is not an apache_beam.pvalue.PCollection.'.format(pcoll)) # No modification to self._user_pipeline is allowed. self._user_pipeline = pcolls[0].pipeline # These are user PCollections. Do not use them to deduce anything that # will be executed by any runner. Instead, use # `self._runner_pcolls_to_user_pcolls.keys()` to get copied PCollections. self._pcolls = set(pcolls) for pcoll in self._pcolls: assert pcoll.pipeline is self._user_pipeline, ( '{} belongs to a different user pipeline than other PCollections ' 'given and cannot be used to build a pipeline fragment that produces ' 'the given PCollections.'.format(pcoll)) self._options = options self._runner = runner # A copied pipeline instance for modification without changing the user # pipeline instance held by the end user. This instance can be processed # into a pipeline fragment that later run by the underlying runner. self._runner_pipeline = self._build_runner_pipeline() _, self._context = self._runner_pipeline.to_runner_api(return_context=True) self._runner_pcoll_to_id = instr.pcoll_to_pcoll_id( self._runner_pipeline, self._context) # Correlate components in the runner pipeline to components in the user # pipeline. The target pcolls are the pcolls given and defined in the user # pipeline. self._id_to_target_pcoll = self._calculate_target_pcoll_ids() self._label_to_user_transform = self._calculate_user_transform_labels() # Below will give us the 1:1 correlation between # PCollections/AppliedPTransforms from the copied runner pipeline and # PCollections/AppliedPTransforms from the user pipeline. # (Dict[PCollection, PCollection]) ( self._runner_pcolls_to_user_pcolls, # (Dict[AppliedPTransform, AppliedPTransform]) self._runner_transforms_to_user_transforms ) = self._build_correlation_between_pipelines( self._runner_pcoll_to_id, self._id_to_target_pcoll, self._label_to_user_transform) # Below are operated on the runner pipeline. (self._necessary_transforms, self._necessary_pcollections) = self._mark_necessary_transforms_and_pcolls( self._runner_pcolls_to_user_pcolls) self._runner_pipeline = self._prune_runner_pipeline_to_fragment( self._runner_pipeline, self._necessary_transforms) def deduce_fragment(self): """Deduce the pipeline fragment as an apache_beam.Pipeline instance.""" fragment = beam.pipeline.Pipeline.from_runner_api( self._runner_pipeline.to_runner_api(), self._runner or self._runner_pipeline.runner, self._options) ie.current_env().add_derived_pipeline(self._runner_pipeline, fragment) return fragment def run(self, display_pipeline_graph=False, use_cache=True, blocking=False): """Shorthand to run the pipeline fragment.""" fragment = self.deduce_fragment() from apache_beam.runners.interactive.interactive_runner import InteractiveRunner try: if isinstance(self._runner_pipeline.runner, InteractiveRunner): preserved_skip_display = self._runner_pipeline.runner._skip_display preserved_force_compute = self._runner_pipeline.runner._force_compute preserved_blocking = self._runner_pipeline.runner._blocking self._runner_pipeline.runner._skip_display = not display_pipeline_graph self._runner_pipeline.runner._force_compute = not use_cache self._runner_pipeline.runner._blocking = blocking return fragment.run() else: pipeline_instrument = instr.build_pipeline_instrument( fragment, self._runner_pipeline._options) pipeline_instrument_proto = ( pipeline_instrument.instrumented_pipeline_proto()) if any(pcoll.is_bounded == beam_runner_api_pb2.IsBounded.UNBOUNDED for pcoll in pipeline_instrument_proto.components.pcollections.values()): raise RuntimeError( 'Please specify InteractiveRunner when creating ' 'the Beam pipeline to use this function ' 'on unbouded PCollections.') result = beam.pipeline.Pipeline.from_runner_api( pipeline_instrument_proto, fragment.runner, fragment._options).run() result.wait_until_finish() ie.current_env().mark_pcollection_computed( pipeline_instrument.cached_pcolls) return result finally: if isinstance(self._runner_pipeline.runner, InteractiveRunner): self._runner_pipeline.runner._skip_display = preserved_skip_display self._runner_pipeline.runner._force_compute = preserved_force_compute self._runner_pipeline.runner._blocking = preserved_blocking def _build_runner_pipeline(self): runner_pipeline = beam.pipeline.Pipeline.from_runner_api( self._user_pipeline.to_runner_api(), self._user_pipeline.runner, self._options) ie.current_env().add_derived_pipeline(self._user_pipeline, runner_pipeline) return runner_pipeline def _calculate_target_pcoll_ids(self): pcoll_id_to_target_pcoll = {} for pcoll in self._pcolls: pcoll_id_to_target_pcoll[self._runner_pcoll_to_id.get(str(pcoll), '')] = pcoll return pcoll_id_to_target_pcoll def _calculate_user_transform_labels(self): label_to_user_transform = {} class UserTransformVisitor(PipelineVisitor): def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): if transform_node is not None: label_to_user_transform[transform_node.full_label] = transform_node v = UserTransformVisitor() self._runner_pipeline.visit(v) return label_to_user_transform def _build_correlation_between_pipelines( self, runner_pcoll_to_id, id_to_target_pcoll, label_to_user_transform): runner_pcolls_to_user_pcolls = {} runner_transforms_to_user_transforms = {} class CorrelationVisitor(PipelineVisitor): def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): self._process_transform(transform_node) for in_pcoll in transform_node.inputs: self._process_pcoll(in_pcoll) for out_pcoll in transform_node.outputs.values(): self._process_pcoll(out_pcoll) def _process_pcoll(self, pcoll): pcoll_id = runner_pcoll_to_id.get(str(pcoll), '') if pcoll_id in id_to_target_pcoll: runner_pcolls_to_user_pcolls[pcoll] = (id_to_target_pcoll[pcoll_id]) def _process_transform(self, transform_node): if transform_node.full_label in label_to_user_transform: runner_transforms_to_user_transforms[transform_node] = ( label_to_user_transform[transform_node.full_label]) v = CorrelationVisitor() self._runner_pipeline.visit(v) return runner_pcolls_to_user_pcolls, runner_transforms_to_user_transforms def _mark_necessary_transforms_and_pcolls(self, runner_pcolls_to_user_pcolls): necessary_transforms = set() all_inputs = set() updated_all_inputs = set(runner_pcolls_to_user_pcolls.keys()) # Do this until no more new PCollection is recorded. while len(updated_all_inputs) != len(all_inputs): all_inputs = set(updated_all_inputs) for pcoll in all_inputs: producer = pcoll.producer while producer: if producer in necessary_transforms: break # Mark the AppliedPTransform as necessary. necessary_transforms.add(producer) # Also mark composites that are not the root transform. If the root # transform is added, then all transforms are incorrectly marked as # necessary. If composites are not handled, then there will be # orphaned PCollections. if producer.parent is not None: necessary_transforms.update(producer.parts) # This will recursively add all the PCollections in this composite. for part in producer.parts: updated_all_inputs.update(part.outputs.values()) # Record all necessary input and side input PCollections. updated_all_inputs.update(producer.inputs) # pylint: disable=bad-option-value side_input_pvalues = set( map(lambda side_input: side_input.pvalue, producer.side_inputs)) updated_all_inputs.update(side_input_pvalues) # Go to its parent AppliedPTransform. producer = producer.parent return necessary_transforms, all_inputs def _prune_runner_pipeline_to_fragment( self, runner_pipeline, necessary_transforms): class PruneVisitor(PipelineVisitor): def enter_composite_transform(self, transform_node): if should_skip_pruning(transform_node): return pruned_parts = list(transform_node.parts) for part in transform_node.parts: if part not in necessary_transforms: pruned_parts.remove(part) transform_node.parts = tuple(pruned_parts) self.visit_transform(transform_node) def visit_transform(self, transform_node): if transform_node not in necessary_transforms: transform_node.parent = None v = PruneVisitor() runner_pipeline.visit(v) return runner_pipeline def should_skip_pruning(transform: AppliedPTransform): return ( isinstance(transform.transform, TestStream) or '_DataFrame_' in transform.full_label) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/pipeline_fragment_test.py0000644000175100017510000001442215206113032026772 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.runners.interactive.pipeline_fragment.""" import unittest from unittest.mock import patch import apache_beam as beam from apache_beam.options.pipeline_options import StandardOptions from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import interactive_runner as ir from apache_beam.runners.interactive import pipeline_fragment as pf from apache_beam.runners.interactive.testing.mock_ipython import mock_get_ipython from apache_beam.runners.interactive.testing.pipeline_assertion import assert_pipeline_equal from apache_beam.runners.interactive.testing.pipeline_assertion import assert_pipeline_proto_equal from apache_beam.testing.test_stream import TestStream @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') class PipelineFragmentTest(unittest.TestCase): def setUp(self): ie.new_env() # Assume a notebook frontend is connected to the mocked ipython kernel. ie.current_env()._is_in_ipython = True ie.current_env()._is_in_notebook = True @patch('IPython.get_ipython', new_callable=mock_get_ipython) def test_build_pipeline_fragment(self, cell): with cell: # Cell 1 p = beam.Pipeline(ir.InteractiveRunner()) p_expected = beam.Pipeline(ir.InteractiveRunner()) # Watch local scope now to allow interactive beam to track the pipelines. ib.watch(locals()) with cell: # Cell 2 # pylint: disable=bad-option-value init = p | 'Init' >> beam.Create(range(10)) init_expected = p_expected | 'Init' >> beam.Create(range(10)) with cell: # Cell 3 square = init | 'Square' >> beam.Map(lambda x: x * x) _ = init | 'Cube' >> beam.Map(lambda x: x**3) _ = init_expected | 'Square' >> beam.Map(lambda x: x * x) # Watch every PCollection has been defined so far in local scope. ib.watch(locals()) fragment = pf.PipelineFragment([square]).deduce_fragment() assert_pipeline_equal(self, p_expected, fragment) @patch('IPython.get_ipython', new_callable=mock_get_ipython) def test_user_pipeline_intact_after_deducing_pipeline_fragment(self, cell): with cell: # Cell 1 p = beam.Pipeline(ir.InteractiveRunner()) # Watch the pipeline `p` immediately without calling locals(). ib.watch({'p': p}) with cell: # Cell 2 # pylint: disable=bad-option-value init = p | 'Init' >> beam.Create(range(10)) with cell: # Cell 3 square = init | 'Square' >> beam.Map(lambda x: x * x) with cell: # Cell 4 cube = init | 'Cube' >> beam.Map(lambda x: x**3) # Watch every PCollection has been defined so far in local scope without # calling locals(). ib.watch({'init': init, 'square': square, 'cube': cube}) user_pipeline_proto_before_deducing_fragment = p.to_runner_api( return_context=False) _ = pf.PipelineFragment([square]).deduce_fragment() user_pipeline_proto_after_deducing_fragment = p.to_runner_api( return_context=False) assert_pipeline_proto_equal( self, user_pipeline_proto_before_deducing_fragment, user_pipeline_proto_after_deducing_fragment) @patch('IPython.get_ipython', new_callable=mock_get_ipython) def test_pipeline_fragment_produces_correct_data(self, cell): with cell: # Cell 1 p = beam.Pipeline(ir.InteractiveRunner()) ib.watch({'p': p}) with cell: # Cell 2 # pylint: disable=bad-option-value init = p | 'Init' >> beam.Create(range(5)) with cell: # Cell 3 square = init | 'Square' >> beam.Map(lambda x: x * x) _ = init | 'Cube' >> beam.Map(lambda x: x**3) ib.watch(locals()) result = pf.PipelineFragment([square]).run() self.assertEqual([0, 1, 4, 9, 16], list(result.get(square))) def test_fragment_does_not_prune_teststream(self): """Tests that the fragment does not prune the TestStream composite parts. """ options = StandardOptions(streaming=True) p = beam.Pipeline(ir.InteractiveRunner(), options) test_stream = p | TestStream(output_tags=['a', 'b']) # pylint: disable=unused-variable a = test_stream['a'] | 'a' >> beam.Map(lambda _: _) b = test_stream['b'] | 'b' >> beam.Map(lambda _: _) fragment = pf.PipelineFragment([b]).deduce_fragment() # If the fragment does prune the TestStreawm composite parts, then the # resulting graph is invalid and the following call will raise an exception. fragment.to_runner_api() @patch('IPython.get_ipython', new_callable=mock_get_ipython) def test_pipeline_composites(self, cell): """Tests that composites are supported. """ with cell: # Cell 1 p = beam.Pipeline(ir.InteractiveRunner()) ib.watch({'p': p}) with cell: # Cell 2 # pylint: disable=bad-option-value init = p | 'Init' >> beam.Create(range(5)) with cell: # Cell 3 # Have a composite within a composite to test that all transforms under a # composite are added. @beam.ptransform_fn def Bar(pcoll): return pcoll | beam.Map(lambda n: 2 * n) @beam.ptransform_fn def Foo(pcoll): p1 = pcoll | beam.Map(lambda n: 3 * n) p2 = pcoll | beam.Map(str) bar = p1 | Bar() return {'pc1': p1, 'pc2': p2, 'bar': bar} res = init | Foo() ib.watch(res) pc = res['bar'] result = pf.PipelineFragment([pc]).run() self.assertEqual([0, 6, 12, 18, 24], list(result.get(pc))) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/pipeline_instrument.py0000644000175100017510000007573615206113032026357 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module to instrument interactivity to the given pipeline. For internal use only; no backwards-compatibility guarantees. This module accesses current interactive environment and analyzes given pipeline to transform original pipeline into a one-shot pipeline with interactivity. """ # pytype: skip-file import logging import apache_beam as beam from apache_beam.pipeline import PipelineVisitor from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.interactive import background_caching_job from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import pipeline_fragment as pf from apache_beam.runners.interactive import utils from apache_beam.runners.interactive.caching.cacheable import Cacheable from apache_beam.runners.interactive.caching.cacheable import CacheKey from apache_beam.runners.interactive.caching.reify import WRITE_CACHE from apache_beam.runners.interactive.caching.reify import reify_to_cache from apache_beam.runners.interactive.caching.reify import unreify_from_cache from apache_beam.testing import test_stream _LOGGER = logging.getLogger(__name__) class PipelineInstrument(object): """A pipeline instrument for pipeline to be executed by interactive runner. This module should never depend on underlying runner that interactive runner delegates. It instruments the original instance of pipeline directly by appending or replacing transforms with help of cache. It provides interfaces to recover states of original pipeline. It's the interactive runner's responsibility to coordinate supported underlying runners to run the pipeline instrumented and recover the original pipeline states if needed. """ def __init__(self, pipeline, options=None): self._pipeline = pipeline self._user_pipeline = ie.current_env().user_pipeline(pipeline) if not self._user_pipeline: self._user_pipeline = pipeline self._cache_manager = ie.current_env().get_cache_manager( self._user_pipeline, create_if_absent=True) # Check if the user defined pipeline contains any source to cache. # If so, during the check, the cache manager is converted into a # streaming cache manager, thus re-assign. if background_caching_job.has_source_to_cache(self._user_pipeline): self._cache_manager = ie.current_env().get_cache_manager( self._user_pipeline) self._background_caching_pipeline = beam.pipeline.Pipeline.from_runner_api( pipeline.to_runner_api(), pipeline.runner, options) ie.current_env().add_derived_pipeline( self._pipeline, self._background_caching_pipeline) # Snapshot of original pipeline information. (self._original_pipeline_proto, context) = self._pipeline.to_runner_api(return_context=True) # All compute-once-against-original-pipeline fields. self._unbounded_sources = utils.unbounded_sources( self._background_caching_pipeline) self._pcoll_to_pcoll_id = pcoll_to_pcoll_id(self._pipeline, context) # A Dict[str, Cacheable] from a PCollection id to a Cacheable that belongs # to the analyzed pipeline. self._cacheables = self.find_cacheables() # A dict from cache key to PCollection that is read from cache. # If exists, caller should reuse the PCollection read. If not, caller # should create new transform and track the PCollection read from cache. # (Dict[str, AppliedPTransform]). self._cached_pcoll_read = {} # A dict from PCollections in the runner pipeline instance to their # corresponding PCollections in the user pipeline instance. Populated # after preprocess(). self._runner_pcoll_to_user_pcoll = {} self._pruned_pipeline_proto = None # Refers target pcolls output by instrumented write cache transforms, used # by pruning logic as supplemental targets to build pipeline fragment up # from. self._extended_targets = set() # Refers pcolls used as inputs but got replaced by outputs of read cache # transforms instrumented, used by pruning logic as targets no longer need # to be produced during pipeline runs. self._ignored_targets = set() # Set of PCollections that are written to cache. self.cached_pcolls = set() def instrumented_pipeline_proto(self): """Always returns a new instance of portable instrumented proto.""" targets = set(self._runner_pcoll_to_user_pcoll.keys()) targets.update(self._extended_targets) targets = targets.difference(self._ignored_targets) if len(targets) > 0: # Prunes upstream transforms that don't contribute to the targets the # instrumented pipeline run cares. return pf.PipelineFragment( list(targets)).deduce_fragment().to_runner_api() return self._pipeline.to_runner_api() def _required_components( self, pipeline_proto, required_transforms_ids, visited, follow_outputs=False, follow_inputs=False): """Returns the components and subcomponents of the given transforms. This method returns required components such as transforms and PCollections related to the given transforms and to all of their subtransforms. This method accomplishes this recursively. """ if not required_transforms_ids: return ({}, {}) transforms = pipeline_proto.components.transforms pcollections = pipeline_proto.components.pcollections # Cache the transforms that will be copied into the new pipeline proto. required_transforms = {k: transforms[k] for k in required_transforms_ids} # Cache all the output PCollections of the transforms. pcollection_ids = [ pc for t in required_transforms.values() for pc in t.outputs.values() ] required_pcollections = { pc_id: pcollections[pc_id] for pc_id in pcollection_ids } subtransforms = {} subpcollections = {} # Recursively go through all the subtransforms and add their components. for transform_id, transform in required_transforms.items(): if transform_id in pipeline_proto.root_transform_ids: continue (t, pc) = self._required_components( pipeline_proto, transform.subtransforms, visited, follow_outputs=False, follow_inputs=False) subtransforms.update(t) subpcollections.update(pc) if follow_outputs: outputs = [ pc_id for t in required_transforms.values() for pc_id in t.outputs.values() ] visited_copy = visited.copy() consuming_transforms = { t_id: t for t_id, t in transforms.items() if set(outputs).intersection(set(t.inputs.values())) } consuming_transforms = set(consuming_transforms.keys()) visited.update(consuming_transforms) consuming_transforms = consuming_transforms - visited_copy (t, pc) = self._required_components( pipeline_proto, list(consuming_transforms), visited, follow_outputs, follow_inputs) subtransforms.update(t) subpcollections.update(pc) if follow_inputs: inputs = [ pc_id for t in required_transforms.values() for pc_id in t.inputs.values() ] producing_transforms = { t_id: t for t_id, t in transforms.items() if set(inputs).intersection(set(t.outputs.values())) } (t, pc) = self._required_components( pipeline_proto, list(producing_transforms.keys()), visited, follow_outputs, follow_inputs) subtransforms.update(t) subpcollections.update(pc) # Now we got all the components and their subcomponents, so return the # complete collection. required_transforms.update(subtransforms) required_pcollections.update(subpcollections) return (required_transforms, required_pcollections) def prune_subgraph_for(self, pipeline, required_transform_ids): # Create the pipeline_proto to read all the components from. It will later # create a new pipeline proto from the cut out components. pipeline_proto, context = pipeline.to_runner_api(return_context=True) # Get all the root transforms. The caching transforms will be subtransforms # of one of these roots. roots = [root for root in pipeline_proto.root_transform_ids] (t, p) = self._required_components( pipeline_proto, roots + required_transform_ids, set(), follow_outputs=True, follow_inputs=True) def set_proto_map(proto_map, new_value): proto_map.clear() for key, value in new_value.items(): proto_map[key].CopyFrom(value) # Copy the transforms into the new pipeline. pipeline_to_execute = beam_runner_api_pb2.Pipeline() pipeline_to_execute.root_transform_ids[:] = roots set_proto_map(pipeline_to_execute.components.transforms, t) set_proto_map(pipeline_to_execute.components.pcollections, p) set_proto_map( pipeline_to_execute.components.coders, context.to_runner_api().coders) set_proto_map( pipeline_to_execute.components.windowing_strategies, context.to_runner_api().windowing_strategies) # Cut out all subtransforms in the root that aren't the required transforms. for root_id in roots: root = pipeline_to_execute.components.transforms[root_id] root.subtransforms[:] = [ transform_id for transform_id in root.subtransforms if transform_id in pipeline_to_execute.components.transforms ] return pipeline_to_execute def background_caching_pipeline_proto(self): """Returns the background caching pipeline. This method creates a background caching pipeline by: adding writes to cache from each unbounded source (done in the instrument method), and cutting out all components (transform, PCollections, coders, windowing strategies) that are not the unbounded sources or writes to cache (or subtransforms thereof). """ # Create the pipeline_proto to read all the components from. It will later # create a new pipeline proto from the cut out components. pipeline_proto, context = self._background_caching_pipeline.to_runner_api( return_context=True) # Get all the sources we want to cache. sources = utils.unbounded_sources(self._background_caching_pipeline) # Get all the root transforms. The caching transforms will be subtransforms # of one of these roots. roots = [root for root in pipeline_proto.root_transform_ids] # Get the transform IDs of the caching transforms. These caching operations # are added to the _background_caching_pipeline in the instrument() method. # It's added there so that multiple calls to this method won't add multiple # caching operations (idempotent). transforms = pipeline_proto.components.transforms caching_transform_ids = [ t_id for root in roots for t_id in transforms[root].subtransforms if WRITE_CACHE in t_id ] # Get the IDs of the unbounded sources. required_transform_labels = [src.full_label for src in sources] unbounded_source_ids = [ k for k, v in transforms.items() if v.unique_name in required_transform_labels ] # The required transforms are the transforms that we want to cut out of # the pipeline_proto and insert into a new pipeline to return. required_transform_ids = ( roots + caching_transform_ids + unbounded_source_ids) (t, p) = self._required_components( pipeline_proto, required_transform_ids, set()) def set_proto_map(proto_map, new_value): proto_map.clear() for key, value in new_value.items(): proto_map[key].CopyFrom(value) # Copy the transforms into the new pipeline. pipeline_to_execute = beam_runner_api_pb2.Pipeline() pipeline_to_execute.root_transform_ids[:] = roots set_proto_map(pipeline_to_execute.components.transforms, t) set_proto_map(pipeline_to_execute.components.pcollections, p) set_proto_map( pipeline_to_execute.components.coders, context.to_runner_api().coders) set_proto_map( pipeline_to_execute.components.windowing_strategies, context.to_runner_api().windowing_strategies) # Cut out all subtransforms in the root that aren't the required transforms. for root_id in roots: root = pipeline_to_execute.components.transforms[root_id] root.subtransforms[:] = [ transform_id for transform_id in root.subtransforms if transform_id in pipeline_to_execute.components.transforms ] return pipeline_to_execute @property def cacheables(self) -> dict[str, Cacheable]: """Returns the Cacheables by PCollection ids. If you're already working with user defined pipelines and PCollections, do not build a PipelineInstrument just to get the cacheables. Instead, use apache_beam.runners.interactive.utils.cacheables. """ return self._cacheables @property def has_unbounded_sources(self): """Returns whether the pipeline has any recordable sources. """ return len(self._unbounded_sources) > 0 @property def original_pipeline_proto(self): """Returns a snapshot of the pipeline proto before instrumentation.""" return self._original_pipeline_proto @property def user_pipeline(self): """Returns a reference to the pipeline instance defined by the user. If a pipeline has no cacheable PCollection and the user pipeline cannot be found, return None indicating there is nothing to be cached in the user pipeline. The pipeline given for instrumenting and mutated in this class is not necessarily the pipeline instance defined by the user. From the watched scopes, this class figures out what the user pipeline instance is. This metadata can be used for tracking pipeline results. """ return self._user_pipeline @property def runner_pcoll_to_user_pcoll(self): """Returns cacheable PCollections correlated from instances in the runner pipeline to instances in the user pipeline.""" return self._runner_pcoll_to_user_pcoll def find_cacheables(self) -> dict[str, Cacheable]: """Finds PCollections that need to be cached for analyzed pipeline. There might be multiple pipelines defined and watched, this will only find cacheables belong to the analyzed pipeline. """ result = {} cacheables = utils.cacheables() for _, cacheable in cacheables.items(): if cacheable.pcoll.pipeline is not self._user_pipeline: # Ignore all cacheables from other pipelines. continue pcoll_id = self.pcoll_id(cacheable.pcoll) if not pcoll_id: _LOGGER.debug( 'Unable to retrieve PCollection id for %s. Ignored.', cacheable.pcoll) continue result[self.pcoll_id(cacheable.pcoll)] = cacheable return result def instrument(self): """Instruments original pipeline with cache. For cacheable output PCollection, if cache for the key doesn't exist, do _write_cache(); for cacheable input PCollection, if cache for the key exists, do _read_cache(). No instrument in any other situation. Modifies: self._pipeline """ cacheable_inputs = set() all_inputs = set() all_outputs = set() unbounded_source_pcolls = set() class InstrumentVisitor(PipelineVisitor): """Visitor utilizes cache to instrument the pipeline.""" def __init__(self, pin): self._pin = pin def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): if isinstance(transform_node.transform, tuple(ie.current_env().options.recordable_sources)): unbounded_source_pcolls.update(transform_node.outputs.values()) cacheable_inputs.update(self._pin._cacheable_inputs(transform_node)) ins, outs = self._pin._all_inputs_outputs(transform_node) all_inputs.update(ins) all_outputs.update(outs) v = InstrumentVisitor(self) self._pipeline.visit(v) # Every output PCollection that is never used as an input PCollection is # considered as a side effect of the pipeline run and should be included. self._extended_targets.update(all_outputs.difference(all_inputs)) # Add the unbounded source PCollections to the cacheable inputs. This allows # for the caching of unbounded sources without a variable reference. cacheable_inputs.update(unbounded_source_pcolls) # Create ReadCache transforms. for cacheable_input in cacheable_inputs: self._read_cache( self._pipeline, cacheable_input, cacheable_input in unbounded_source_pcolls) # Replace/wire inputs w/ cached PCollections from ReadCache transforms. self._replace_with_cached_inputs(self._pipeline) # Write cache for all cacheables. for _, cacheable in self._cacheables.items(): self._write_cache( self._pipeline, cacheable.pcoll, ignore_unbounded_reads=True) # Instrument the background caching pipeline if we can. if self.has_unbounded_sources: for source in self._unbounded_sources: self._write_cache( self._background_caching_pipeline, source.outputs[None], output_as_extended_target=False, is_capture=True) class TestStreamVisitor(PipelineVisitor): def __init__(self): self.test_stream = None def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): if (self.test_stream is None and isinstance(transform_node.transform, test_stream.TestStream)): self.test_stream = transform_node.full_label v = TestStreamVisitor() self._pipeline.visit(v) pipeline_proto = self._pipeline.to_runner_api(return_context=False) test_stream_id = '' for t_id, t in pipeline_proto.components.transforms.items(): if t.unique_name == v.test_stream: test_stream_id = t_id break self._pruned_pipeline_proto = self.prune_subgraph_for( self._pipeline, [test_stream_id]) pruned_pipeline = beam.Pipeline.from_runner_api( proto=self._pruned_pipeline_proto, runner=self._pipeline.runner, options=self._pipeline._options) ie.current_env().add_derived_pipeline(self._pipeline, pruned_pipeline) self._pipeline = pruned_pipeline def preprocess(self): """Pre-processes the pipeline. Since the pipeline instance in the class might not be the same instance defined in the user code, the pre-process will figure out the relationship of cacheable PCollections between these 2 instances by replacing 'pcoll' fields in the cacheable dictionary with ones from the running instance. """ class PreprocessVisitor(PipelineVisitor): def __init__(self, pin): self._pin = pin def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): for in_pcoll in transform_node.inputs: self._process(in_pcoll) for out_pcoll in transform_node.outputs.values(): self._process(out_pcoll) def _process(self, pcoll): pcoll_id = self._pin._pcoll_to_pcoll_id.get(str(pcoll), '') if pcoll_id in self._pin._cacheables: pcoll_id = self._pin.pcoll_id(pcoll) user_pcoll = self._pin._cacheables[pcoll_id].pcoll if (pcoll_id in self._pin._cacheables and user_pcoll != pcoll): self._pin._runner_pcoll_to_user_pcoll[pcoll] = user_pcoll self._pin._cacheables[pcoll_id].pcoll = pcoll v = PreprocessVisitor(self) self._pipeline.visit(v) def _write_cache( self, pipeline, pcoll, output_as_extended_target=True, ignore_unbounded_reads=False, is_capture=False): """Caches a cacheable PCollection. For the given PCollection, by appending sub transform part that materialize the PCollection through sink into cache implementation. The cache write is not immediate. It happens when the runner runs the transformed pipeline and thus not usable for this run as intended. This function always writes the cache for the given PCollection as long as the PCollection belongs to the pipeline being instrumented and the keyed cache is absent. Modifies: pipeline """ # Makes sure the pcoll belongs to the pipeline being instrumented. if pcoll.pipeline is not pipeline: return # Ignore the unbounded reads from recordable sources as these will be pruned # out using the PipelineFragment later on. if ignore_unbounded_reads: ignore = False producer = pcoll.producer while producer: if isinstance(producer.transform, tuple(ie.current_env().options.recordable_sources)): ignore = True break producer = producer.parent if ignore: self._ignored_targets.add(pcoll) return # The keyed cache is always valid within this instrumentation. key = self.cache_key(pcoll) # Only need to write when the cache with expected key doesn't exist. if not self._cache_manager.exists('full', key): self.cached_pcolls.add(self.runner_pcoll_to_user_pcoll.get(pcoll, pcoll)) # Read the windowing information and cache it along with the element. This # caches the arguments to a WindowedValue object because Python has logic # that detects if a DoFn returns a WindowedValue. When it detecs one, it # puts the element into the correct window then emits the value to # downstream transforms. extended_target = reify_to_cache( pcoll=pcoll, cache_key=key, cache_manager=self._cache_manager, is_capture=is_capture) if output_as_extended_target: self._extended_targets.add(extended_target) def _read_cache(self, pipeline, pcoll, is_unbounded_source_output): """Reads a cached pvalue. A noop will cause the pipeline to execute the transform as it is and cache nothing from this transform for next run. Modifies: pipeline """ # Makes sure the pcoll belongs to the pipeline being instrumented. if pcoll.pipeline is not pipeline: return # The keyed cache is always valid within this instrumentation. key = self.cache_key(pcoll) # Can only read from cache when the cache with expected key exists and its # computation has been completed. is_cached = self._cache_manager.exists('full', key) is_computed = ( pcoll in self._runner_pcoll_to_user_pcoll and self._runner_pcoll_to_user_pcoll[pcoll] in ie.current_env().computed_pcollections) if ((is_cached and is_computed) or is_unbounded_source_output): if key not in self._cached_pcoll_read: # Mutates the pipeline with cache read transform attached # to root of the pipeline. # To put the cached value into the correct window, simply return a # WindowedValue constructed from the element. pcoll_from_cache = unreify_from_cache( pipeline=pipeline, cache_key=key, cache_manager=self._cache_manager) self._cached_pcoll_read[key] = pcoll_from_cache # else: NOOP when cache doesn't exist, just compute the original graph. def _replace_with_cached_inputs(self, pipeline): """Replace PCollection inputs in the pipeline with cache if possible. For any input PCollection, find out whether there is valid cache. If so, replace the input of the AppliedPTransform with output of the AppliedPtransform that sources pvalue from the cache. If there is no valid cache, noop. """ # Find all cached unbounded PCollections. # If the pipeline has unbounded sources, then we want to force all cache # reads to go through the TestStream (even if they are bounded sources). if self.has_unbounded_sources: class CacheableUnboundedPCollectionVisitor(PipelineVisitor): def __init__(self, pin): self._pin = pin self.unbounded_pcolls = set() def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): if transform_node.outputs: for output_pcoll in transform_node.outputs.values(): key = self._pin.cache_key(output_pcoll) if key in self._pin._cached_pcoll_read: self.unbounded_pcolls.add(key) if transform_node.inputs: for input_pcoll in transform_node.inputs: key = self._pin.cache_key(input_pcoll) if key in self._pin._cached_pcoll_read: self.unbounded_pcolls.add(key) v = CacheableUnboundedPCollectionVisitor(self) pipeline.visit(v) # The set of keys from the cached unbounded PCollections will be used as # the output tags for the TestStream. This is to remember what cache-key # is associated with which PCollection. output_tags = v.unbounded_pcolls # Take the PCollections that will be read from the TestStream and insert # them back into the dictionary of cached PCollections. The next step will # replace the downstream consumer of the non-cached PCollections with # these PCollections. if output_tags: output_pcolls = pipeline | test_stream.TestStream( output_tags=output_tags, coder=self._cache_manager._default_pcoder) for tag, pcoll in output_pcolls.items(): self._cached_pcoll_read[tag] = pcoll class ReadCacheWireVisitor(PipelineVisitor): """Visitor wires cache read as inputs to replace corresponding original input PCollections in pipeline. """ def __init__(self, pin): """Initializes with a PipelineInstrument.""" self._pin = pin def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): if transform_node.inputs: main_inputs = dict(transform_node.main_inputs) for tag, input_pcoll in main_inputs.items(): key = self._pin.cache_key(input_pcoll) # Replace the input pcollection with the cached pcollection (if it # has been cached). if key in self._pin._cached_pcoll_read: # Ignore this pcoll in the final pruned instrumented pipeline. self._pin._ignored_targets.add(input_pcoll) main_inputs[tag] = self._pin._cached_pcoll_read[key] # Update the transform with its new inputs. transform_node.main_inputs = main_inputs v = ReadCacheWireVisitor(self) pipeline.visit(v) def _cacheable_inputs(self, transform): inputs = set() for in_pcoll in transform.inputs: if self.pcoll_id(in_pcoll) in self._cacheables: inputs.add(in_pcoll) return inputs def _all_inputs_outputs(self, transform): inputs = set() outputs = set() for in_pcoll in transform.inputs: inputs.add(in_pcoll) for _, out_pcoll in transform.outputs.items(): outputs.add(out_pcoll) return inputs, outputs def pcoll_id(self, pcoll): """Gets the PCollection id of the given pcoll. Returns '' if not found. """ return self._pcoll_to_pcoll_id.get(str(pcoll), '') def cache_key(self, pcoll): """Gets the identifier of a cacheable PCollection in cache. If the pcoll is not a cacheable, return ''. This is only needed in pipeline instrument when the origin of given pcoll is unknown (whether it's from the user pipeline or a runner pipeline). If a pcoll is from the user pipeline, always use CacheKey.from_pcoll to build the key. The key is what the pcoll would use as identifier if it's materialized in cache. It doesn't mean that there would definitely be such cache already. Also, the pcoll can come from the original user defined pipeline object or an equivalent pcoll from a transformed copy of the original pipeline. """ cacheable = self._cacheables.get(self.pcoll_id(pcoll), None) if cacheable: if cacheable.pcoll in self.runner_pcoll_to_user_pcoll: user_pcoll = self.runner_pcoll_to_user_pcoll[cacheable.pcoll] else: user_pcoll = cacheable.pcoll return CacheKey.from_pcoll(cacheable.var, user_pcoll).to_str() return '' def build_pipeline_instrument(pipeline, options=None): """Creates PipelineInstrument for a pipeline and its options with cache. Throughout the process, the returned PipelineInstrument snapshots the given pipeline and then mutates the pipeline. It's invoked by interactive components such as the InteractiveRunner and the given pipeline should be implicitly created runner pipelines instead of pipeline instances defined by the user. This is the shorthand for doing 3 steps: 1) compute once for metadata of the given runner pipeline and everything watched from user pipelines; 2) associate info between the runner pipeline and its corresponding user pipeline, eliminate data from other user pipelines if there are any; 3) mutate the runner pipeline to apply interactivity. """ pi = PipelineInstrument(pipeline, options) pi.preprocess() pi.instrument() # Instruments the pipeline only once. return pi def pcoll_to_pcoll_id(pipeline, original_context): """Returns a dict mapping PCollections string to PCollection IDs. Using a PipelineVisitor to iterate over every node in the pipeline, records the mapping from PCollections to PCollections IDs. This mapping will be used to query cached PCollections. Returns: (dict from str to str) a dict mapping str(pcoll) to pcoll_id. """ class PCollVisitor(PipelineVisitor): """"A visitor that records input and output values to be replaced. Input and output values that should be updated are recorded in maps input_replacements and output_replacements respectively. We cannot update input and output values while visiting since that results in validation errors. """ def __init__(self): self.pcoll_to_pcoll_id = {} def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): for pcoll in transform_node.outputs.values(): self.pcoll_to_pcoll_id[str(pcoll)] = ( original_context.pcollections.get_id(pcoll)) v = PCollVisitor() pipeline.visit(v) return v.pcoll_to_pcoll_id ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/pipeline_instrument_test.py0000644000175100017510000010365015206113032027401 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.runners.interactive.pipeline_instrument.""" # pytype: skip-file import unittest import apache_beam as beam from apache_beam import coders from apache_beam.pipeline import PipelineVisitor from apache_beam.runners.interactive import cache_manager as cache from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import interactive_runner from apache_beam.runners.interactive import pipeline_instrument as instr from apache_beam.runners.interactive import utils from apache_beam.runners.interactive.caching.cacheable import Cacheable from apache_beam.runners.interactive.caching.cacheable import CacheKey from apache_beam.runners.interactive.caching.streaming_cache import StreamingCache from apache_beam.runners.interactive.testing.pipeline_assertion import assert_pipeline_equal from apache_beam.runners.interactive.testing.pipeline_assertion import assert_pipeline_proto_contain_top_level_transform from apache_beam.runners.interactive.testing.pipeline_assertion import assert_pipeline_proto_equal from apache_beam.runners.interactive.testing.pipeline_assertion import \ assert_pipeline_proto_not_contain_top_level_transform from apache_beam.runners.interactive.testing.test_cache_manager import InMemoryCache from apache_beam.testing.test_stream import TestStream class PipelineInstrumentTest(unittest.TestCase): def setUp(self): ie.new_env() def cache_key_of(self, name, pcoll): return CacheKey.from_pcoll(name, pcoll).to_str() def test_pcoll_to_pcoll_id(self): p = beam.Pipeline(interactive_runner.InteractiveRunner()) ie.current_env().set_cache_manager(InMemoryCache(), p) # pylint: disable=bad-option-value init_pcoll = p | 'Init Create' >> beam.Impulse() _, ctx = p.to_runner_api(return_context=True) self.assertEqual( instr.pcoll_to_pcoll_id(p, ctx), {str(init_pcoll): 'ref_PCollection_PCollection_1'}) def test_pcoll_id_with_user_pipeline(self): p_id_user = beam.Pipeline(interactive_runner.InteractiveRunner()) ie.current_env().set_cache_manager(InMemoryCache(), p_id_user) init_pcoll = p_id_user | 'Init Create' >> beam.Create([1, 2, 3]) instrumentation = instr.build_pipeline_instrument(p_id_user) self.assertEqual( instrumentation.pcoll_id(init_pcoll), 'ref_PCollection_PCollection_8') def test_pcoll_id_with_runner_pipeline(self): p_id_runner = beam.Pipeline(interactive_runner.InteractiveRunner()) ie.current_env().set_cache_manager(InMemoryCache(), p_id_runner) # pylint: disable=possibly-unused-variable init_pcoll = p_id_runner | 'Init Create' >> beam.Create([1, 2, 3]) ib.watch(locals()) # It's normal that when executing, the pipeline object is a different # but equivalent instance from what user has built. The pipeline instrument # should be able to identify if the original instance has changed in an # interactive env while mutating the other instance for execution. The # version map can be used to figure out what the PCollection instances are # in the original instance and if the evaluation has changed since last # execution. p2_id_runner = beam.Pipeline(interactive_runner.InteractiveRunner()) # pylint: disable=bad-option-value init_pcoll_2 = p2_id_runner | 'Init Create' >> beam.Create(range(10)) ie.current_env().add_derived_pipeline(p_id_runner, p2_id_runner) instrumentation = instr.build_pipeline_instrument(p2_id_runner) # The cache_key should use id(init_pcoll) as prefix even when # init_pcoll_2 is supplied as long as the version map is given. self.assertEqual( instrumentation.pcoll_id(init_pcoll_2), 'ref_PCollection_PCollection_8') def test_cache_key(self): p = beam.Pipeline(interactive_runner.InteractiveRunner()) ie.current_env().set_cache_manager(InMemoryCache(), p) # pylint: disable=bad-option-value init_pcoll = p | 'Init Create' >> beam.Create(range(10)) squares = init_pcoll | 'Square' >> beam.Map(lambda x: x * x) cubes = init_pcoll | 'Cube' >> beam.Map(lambda x: x**3) # Watch the local variables, i.e., the Beam pipeline defined. ib.watch(locals()) pipeline_instrument = instr.build_pipeline_instrument(p) self.assertEqual( pipeline_instrument.cache_key(init_pcoll), self.cache_key_of('init_pcoll', init_pcoll)) self.assertEqual( pipeline_instrument.cache_key(squares), self.cache_key_of('squares', squares)) self.assertEqual( pipeline_instrument.cache_key(cubes), self.cache_key_of('cubes', cubes)) def test_cacheables(self): p_cacheables = beam.Pipeline(interactive_runner.InteractiveRunner()) ie.current_env().set_cache_manager(InMemoryCache(), p_cacheables) # pylint: disable=bad-option-value init_pcoll = p_cacheables | 'Init Create' >> beam.Create(range(10)) squares = init_pcoll | 'Square' >> beam.Map(lambda x: x * x) cubes = init_pcoll | 'Cube' >> beam.Map(lambda x: x**3) ib.watch(locals()) pipeline_instrument = instr.build_pipeline_instrument(p_cacheables) self.assertEqual( pipeline_instrument._cacheables, { pipeline_instrument.pcoll_id(init_pcoll): Cacheable( var='init_pcoll', version=str(id(init_pcoll)), producer_version=str(id(init_pcoll.producer)), pcoll=init_pcoll), pipeline_instrument.pcoll_id(squares): Cacheable( var='squares', version=str(id(squares)), producer_version=str(id(squares.producer)), pcoll=squares), pipeline_instrument.pcoll_id(cubes): Cacheable( var='cubes', version=str(id(cubes)), producer_version=str(id(cubes.producer)), pcoll=cubes) }) def test_background_caching_pipeline_proto(self): p = beam.Pipeline(interactive_runner.InteractiveRunner()) ie.current_env().set_cache_manager(StreamingCache(cache_dir=None), p) # Test that the two ReadFromPubSub are correctly cut out. a = p | 'ReadUnboundedSourceA' >> beam.io.ReadFromPubSub( subscription='projects/fake-project/subscriptions/fake_sub') b = p | 'ReadUnboundedSourceB' >> beam.io.ReadFromPubSub( subscription='projects/fake-project/subscriptions/fake_sub') # Add some extra PTransform afterwards to make sure that only the unbounded # sources remain. c = (a, b) | beam.Flatten() _ = c | beam.Map(lambda x: x) ib.watch(locals()) instrumenter = instr.build_pipeline_instrument(p) actual_pipeline = instrumenter.background_caching_pipeline_proto() # Now recreate the expected pipeline, which should only have the unbounded # sources. p = beam.Pipeline(interactive_runner.InteractiveRunner()) ie.current_env().set_cache_manager(StreamingCache(cache_dir=None), p) a = p | 'ReadUnboundedSourceA' >> beam.io.ReadFromPubSub( subscription='projects/fake-project/subscriptions/fake_sub') _ = ( a | 'reify a' >> beam.Map(lambda _: _) | 'a' >> cache.WriteCache(ie.current_env().get_cache_manager(p), '')) b = p | 'ReadUnboundedSourceB' >> beam.io.ReadFromPubSub( subscription='projects/fake-project/subscriptions/fake_sub') _ = ( b | 'reify b' >> beam.Map(lambda _: _) | 'b' >> cache.WriteCache(ie.current_env().get_cache_manager(p), '')) expected_pipeline = p.to_runner_api(return_context=False) assert_pipeline_proto_equal(self, expected_pipeline, actual_pipeline) def _example_pipeline(self, watch=True, bounded=True): p_example = beam.Pipeline(interactive_runner.InteractiveRunner()) ie.current_env().set_cache_manager(InMemoryCache(), p_example) # pylint: disable=bad-option-value if bounded: source = beam.Create(range(10)) else: source = beam.io.ReadFromPubSub( subscription='projects/fake-project/subscriptions/fake_sub') init_pcoll = p_example | 'Init Source' >> source second_pcoll = init_pcoll | 'Second' >> beam.Map(lambda x: x * x) if watch: ib.watch(locals()) return (p_example, init_pcoll, second_pcoll) def _mock_write_cache(self, pipeline, values, cache_key): """Cache the PCollection where cache.WriteCache would write to.""" labels = ['full', cache_key] # Usually, the pcoder will be inferred from `pcoll.element_type` pcoder = coders.registry.get_coder(object) cache_manager = ie.current_env().get_cache_manager(pipeline) cache_manager.save_pcoder(pcoder, *labels) cache_manager.write(values, *labels) def test_instrument_example_pipeline_to_write_cache(self): # Original instance defined by user code has all variables handlers. p_origin, init_pcoll, second_pcoll = self._example_pipeline() # Copied instance when execution has no user defined variables. p_copy, _, _ = self._example_pipeline(watch=False) ie.current_env().add_derived_pipeline(p_origin, p_copy) # Instrument the copied pipeline. pipeline_instrument = instr.build_pipeline_instrument(p_copy) # Manually instrument original pipeline with expected pipeline transforms. init_pcoll_cache_key = pipeline_instrument.cache_key(init_pcoll) _ = ( init_pcoll | 'reify init' >> beam.Map(lambda _: _) | '_WriteCache_' + init_pcoll_cache_key >> cache.WriteCache( ie.current_env().get_cache_manager(p_origin), init_pcoll_cache_key)) second_pcoll_cache_key = pipeline_instrument.cache_key(second_pcoll) _ = ( second_pcoll | 'reify second' >> beam.Map(lambda _: _) | '_WriteCache_' + second_pcoll_cache_key >> cache.WriteCache( ie.current_env().get_cache_manager(p_origin), second_pcoll_cache_key)) # The 2 pipelines should be the same now. assert_pipeline_equal(self, p_copy, p_origin) def test_instrument_example_pipeline_to_read_cache(self): p_origin, init_pcoll, second_pcoll = self._example_pipeline() p_copy, _, _ = self._example_pipeline(False) # Mock as if cacheable PCollections are cached. init_pcoll_cache_key = self.cache_key_of('init_pcoll', init_pcoll) self._mock_write_cache(p_origin, [b'1', b'2', b'3'], init_pcoll_cache_key) second_pcoll_cache_key = self.cache_key_of('second_pcoll', second_pcoll) self._mock_write_cache(p_origin, [b'1', b'4', b'9'], second_pcoll_cache_key) # Mark the completeness of PCollections from the original(user) pipeline. ie.current_env().mark_pcollection_computed((init_pcoll, second_pcoll)) ie.current_env().add_derived_pipeline(p_origin, p_copy) instr.build_pipeline_instrument(p_copy) cached_init_pcoll = ( p_origin | '_ReadCache_' + init_pcoll_cache_key >> cache.ReadCache( ie.current_env().get_cache_manager(p_origin), init_pcoll_cache_key) | 'unreify' >> beam.Map(lambda _: _)) # second_pcoll is never used as input and there is no need to read cache. class TestReadCacheWireVisitor(PipelineVisitor): """Replace init_pcoll with cached_init_pcoll for all occuring inputs.""" def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): if transform_node.inputs: main_inputs = dict(transform_node.main_inputs) for tag, main_input in main_inputs.items(): if main_input == init_pcoll: main_inputs[tag] = cached_init_pcoll transform_node.main_inputs = main_inputs v = TestReadCacheWireVisitor() p_origin.visit(v) assert_pipeline_equal(self, p_origin, p_copy) def test_find_out_correct_user_pipeline(self): # This is the user pipeline instance we care in the watched scope. user_pipeline, _, _ = self._example_pipeline() # This is a new runner pipeline instance with the same pipeline graph to # what the user_pipeline represents. runner_pipeline = beam.pipeline.Pipeline.from_runner_api( user_pipeline.to_runner_api(), user_pipeline.runner, options=None) ie.current_env().add_derived_pipeline(user_pipeline, runner_pipeline) # This is a totally irrelevant user pipeline in the watched scope. irrelevant_user_pipeline = beam.Pipeline( interactive_runner.InteractiveRunner()) ib.watch({'irrelevant_user_pipeline': irrelevant_user_pipeline}) # Build instrument from the runner pipeline. pipeline_instrument = instr.build_pipeline_instrument(runner_pipeline) self.assertIs(pipeline_instrument.user_pipeline, user_pipeline) def test_instrument_example_unbounded_pipeline_to_read_cache(self): """Tests that the instrumenter works for a single unbounded source. """ # Create the pipeline that will be instrumented. p_original = beam.Pipeline(interactive_runner.InteractiveRunner()) ie.current_env().set_cache_manager( StreamingCache(cache_dir=None), p_original) source_1 = p_original | 'source1' >> beam.io.ReadFromPubSub( subscription='projects/fake-project/subscriptions/fake_sub') # pylint: disable=possibly-unused-variable pcoll_1 = source_1 | 'square1' >> beam.Map(lambda x: x * x) # Mock as if cacheable PCollections are cached. ib.watch(locals()) # This should be noop. utils.watch_sources(p_original) for name, pcoll in locals().items(): if not isinstance(pcoll, beam.pvalue.PCollection): continue cache_key = self.cache_key_of(name, pcoll) self._mock_write_cache(p_original, [], cache_key) # Instrument the original pipeline to create the pipeline the user will see. instrumenter = instr.build_pipeline_instrument(p_original) actual_pipeline = beam.Pipeline.from_runner_api( proto=instrumenter.instrumented_pipeline_proto(), runner=interactive_runner.InteractiveRunner(), options=None) # Now, build the expected pipeline which replaces the unbounded source with # a TestStream. source_1_cache_key = self.cache_key_of('source_1', source_1) p_expected = beam.Pipeline() test_stream = (p_expected | TestStream(output_tags=[source_1_cache_key])) # pylint: disable=expression-not-assigned test_stream[source_1_cache_key] | 'square1' >> beam.Map(lambda x: x * x) # Test that the TestStream is outputting to the correct PCollection. class TestStreamVisitor(PipelineVisitor): def __init__(self): self.output_tags = set() def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): transform = transform_node.transform if isinstance(transform, TestStream): self.output_tags = transform.output_tags v = TestStreamVisitor() actual_pipeline.visit(v) expected_output_tags = set([source_1_cache_key]) actual_output_tags = v.output_tags self.assertSetEqual(expected_output_tags, actual_output_tags) # Test that the pipeline is as expected. assert_pipeline_proto_equal( self, p_expected.to_runner_api(), instrumenter.instrumented_pipeline_proto()) def test_able_to_cache_intermediate_unbounded_source_pcollection(self): """Tests being able to cache an intermediate source PCollection. In the following pipeline, the source doesn't have a reference and so is not automatically cached in the watch() command. This tests that this case is taken care of. """ # Create the pipeline that will be instrumented. from apache_beam.options.pipeline_options import StandardOptions options = StandardOptions(streaming=True) streaming_cache_manager = StreamingCache(cache_dir=None) p_original_cache_source = beam.Pipeline( interactive_runner.InteractiveRunner(), options) ie.current_env().set_cache_manager( streaming_cache_manager, p_original_cache_source) # pylint: disable=possibly-unused-variable source_1 = ( p_original_cache_source | 'source1' >> beam.io.ReadFromPubSub( subscription='projects/fake-project/subscriptions/fake_sub') | beam.Map(lambda e: e)) # Watch but do not cache the PCollections. ib.watch(locals()) # Make sure that sources without a user reference are still cached. utils.watch_sources(p_original_cache_source) intermediate_source_pcoll = None for watching in ie.current_env().watching(): watching = list(watching) for var, watchable in watching: if 'synthetic' in var: intermediate_source_pcoll = watchable break # Instrument the original pipeline to create the pipeline the user will see. p_copy = beam.Pipeline.from_runner_api( p_original_cache_source.to_runner_api(), runner=interactive_runner.InteractiveRunner(), options=options) ie.current_env().add_derived_pipeline(p_original_cache_source, p_copy) instrumenter = instr.build_pipeline_instrument(p_copy) actual_pipeline = beam.Pipeline.from_runner_api( proto=instrumenter.instrumented_pipeline_proto(), runner=interactive_runner.InteractiveRunner(), options=options) ie.current_env().add_derived_pipeline( p_original_cache_source, actual_pipeline) # Now, build the expected pipeline which replaces the unbounded source with # a TestStream. intermediate_source_pcoll_cache_key = \ self.cache_key_of('synthetic_var_' + str(id(intermediate_source_pcoll)), intermediate_source_pcoll) p_expected = beam.Pipeline() ie.current_env().set_cache_manager(streaming_cache_manager, p_expected) test_stream = ( p_expected | TestStream(output_tags=[intermediate_source_pcoll_cache_key])) # pylint: disable=expression-not-assigned ( test_stream[intermediate_source_pcoll_cache_key] | 'square1' >> beam.Map(lambda e: e) | 'reify' >> beam.Map(lambda _: _) | cache.WriteCache( ie.current_env().get_cache_manager(p_expected), 'unused')) # Test that the TestStream is outputting to the correct PCollection. class TestStreamVisitor(PipelineVisitor): def __init__(self): self.output_tags = set() def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): transform = transform_node.transform if isinstance(transform, TestStream): self.output_tags = transform.output_tags v = TestStreamVisitor() actual_pipeline.visit(v) expected_output_tags = set([intermediate_source_pcoll_cache_key]) actual_output_tags = v.output_tags self.assertSetEqual(expected_output_tags, actual_output_tags) # Test that the pipeline is as expected. assert_pipeline_proto_equal( self, p_expected.to_runner_api(), instrumenter.instrumented_pipeline_proto()) def test_instrument_mixed_streaming_batch(self): """Tests caching for both batch and streaming sources in the same pipeline. This ensures that cached bounded and unbounded sources are read from the TestStream. """ # Create the pipeline that will be instrumented. from apache_beam.options.pipeline_options import StandardOptions options = StandardOptions(streaming=True) p_original = beam.Pipeline(interactive_runner.InteractiveRunner(), options) streaming_cache_manager = StreamingCache(cache_dir=None) ie.current_env().set_cache_manager(streaming_cache_manager, p_original) source_1 = p_original | 'source1' >> beam.io.ReadFromPubSub( subscription='projects/fake-project/subscriptions/fake_sub') source_2 = p_original | 'source2' >> beam.Create([1, 2, 3, 4, 5]) # pylint: disable=possibly-unused-variable pcoll_1 = ((source_1, source_2) | beam.Flatten() | 'square1' >> beam.Map(lambda x: x * x)) # Watch but do not cache the PCollections. ib.watch(locals()) # This should be noop. utils.watch_sources(p_original) self._mock_write_cache( p_original, [], self.cache_key_of('source_2', source_2)) ie.current_env().mark_pcollection_computed([source_2]) # Instrument the original pipeline to create the pipeline the user will see. p_copy = beam.Pipeline.from_runner_api( p_original.to_runner_api(), runner=interactive_runner.InteractiveRunner(), options=options) ie.current_env().add_derived_pipeline(p_original, p_copy) instrumenter = instr.build_pipeline_instrument(p_copy) actual_pipeline = beam.Pipeline.from_runner_api( proto=instrumenter.instrumented_pipeline_proto(), runner=interactive_runner.InteractiveRunner(), options=options) # Now, build the expected pipeline which replaces the unbounded source with # a TestStream. source_1_cache_key = self.cache_key_of('source_1', source_1) source_2_cache_key = self.cache_key_of('source_2', source_2) p_expected = beam.Pipeline() ie.current_env().set_cache_manager(streaming_cache_manager, p_expected) test_stream = ( p_expected | TestStream(output_tags=[source_1_cache_key, source_2_cache_key])) # pylint: disable=expression-not-assigned (( test_stream[self.cache_key_of('source_1', source_1)], test_stream[self.cache_key_of('source_2', source_2)]) | beam.Flatten() | 'square1' >> beam.Map(lambda x: x * x) | 'reify' >> beam.Map(lambda _: _) | cache.WriteCache( ie.current_env().get_cache_manager(p_expected), 'unused')) # Test that the TestStream is outputting to the correct PCollection. class TestStreamVisitor(PipelineVisitor): def __init__(self): self.output_tags = set() def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): transform = transform_node.transform if isinstance(transform, TestStream): self.output_tags = transform.output_tags v = TestStreamVisitor() actual_pipeline.visit(v) expected_output_tags = set([source_1_cache_key, source_2_cache_key]) actual_output_tags = v.output_tags self.assertSetEqual(expected_output_tags, actual_output_tags) # Test that the pipeline is as expected. assert_pipeline_proto_equal( self, p_expected.to_runner_api(), instrumenter.instrumented_pipeline_proto()) def test_instrument_example_unbounded_pipeline_direct_from_source(self): """Tests that the it caches PCollections from a source. """ # Create the pipeline that will be instrumented. from apache_beam.options.pipeline_options import StandardOptions options = StandardOptions(streaming=True) p_original_direct_source = beam.Pipeline( interactive_runner.InteractiveRunner(), options) ie.current_env().set_cache_manager( StreamingCache(cache_dir=None), p_original_direct_source) source_1 = p_original_direct_source | 'source1' >> beam.io.ReadFromPubSub( subscription='projects/fake-project/subscriptions/fake_sub') # pylint: disable=possibly-unused-variable p_expected = beam.Pipeline() # pylint: disable=unused-variable test_stream = ( p_expected | TestStream(output_tags=[self.cache_key_of('source_1', source_1)])) # Watch but do not cache the PCollections. ib.watch(locals()) # This should be noop. utils.watch_sources(p_original_direct_source) # Instrument the original pipeline to create the pipeline the user will see. p_copy = beam.Pipeline.from_runner_api( p_original_direct_source.to_runner_api(), runner=interactive_runner.InteractiveRunner(), options=options) ie.current_env().add_derived_pipeline(p_original_direct_source, p_copy) instrumenter = instr.build_pipeline_instrument(p_copy) actual_pipeline = beam.Pipeline.from_runner_api( proto=instrumenter.instrumented_pipeline_proto(), runner=interactive_runner.InteractiveRunner(), options=options) ie.current_env().add_derived_pipeline( p_original_direct_source, actual_pipeline) # Now, build the expected pipeline which replaces the unbounded source with # a TestStream. source_1_cache_key = self.cache_key_of('source_1', source_1) # Test that the TestStream is outputting to the correct PCollection. class TestStreamVisitor(PipelineVisitor): def __init__(self): self.output_tags = set() def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): transform = transform_node.transform if isinstance(transform, TestStream): self.output_tags = transform.output_tags v = TestStreamVisitor() actual_pipeline.visit(v) expected_output_tags = set([source_1_cache_key]) actual_output_tags = v.output_tags self.assertSetEqual(expected_output_tags, actual_output_tags) # Test that the pipeline is as expected. assert_pipeline_proto_equal( self, p_expected.to_runner_api(), instrumenter.instrumented_pipeline_proto()) def test_instrument_example_unbounded_pipeline_to_read_cache_not_cached(self): """Tests that the instrumenter works when the PCollection is not cached. """ # Create the pipeline that will be instrumented. from apache_beam.options.pipeline_options import StandardOptions options = StandardOptions(streaming=True) p_original_read_cache = beam.Pipeline( interactive_runner.InteractiveRunner(), options) ie.current_env().set_cache_manager( StreamingCache(cache_dir=None), p_original_read_cache) source_1 = p_original_read_cache | 'source1' >> beam.io.ReadFromPubSub( subscription='projects/fake-project/subscriptions/fake_sub') # pylint: disable=possibly-unused-variable pcoll_1 = source_1 | 'square1' >> beam.Map(lambda x: x * x) # Watch but do not cache the PCollections. ib.watch(locals()) # This should be noop. utils.watch_sources(p_original_read_cache) # Instrument the original pipeline to create the pipeline the user will see. p_copy = beam.Pipeline.from_runner_api( p_original_read_cache.to_runner_api(), runner=interactive_runner.InteractiveRunner(), options=options) ie.current_env().add_derived_pipeline(p_original_read_cache, p_copy) instrumenter = instr.build_pipeline_instrument(p_copy) actual_pipeline = beam.Pipeline.from_runner_api( proto=instrumenter.instrumented_pipeline_proto(), runner=interactive_runner.InteractiveRunner(), options=options) # Now, build the expected pipeline which replaces the unbounded source with # a TestStream. source_1_cache_key = self.cache_key_of('source_1', source_1) p_expected = beam.Pipeline() ie.current_env().set_cache_manager( StreamingCache(cache_dir=None), p_expected) test_stream = (p_expected | TestStream(output_tags=[source_1_cache_key])) # pylint: disable=expression-not-assigned ( test_stream[source_1_cache_key] | 'square1' >> beam.Map(lambda x: x * x) | 'reify' >> beam.Map(lambda _: _) | cache.WriteCache( ie.current_env().get_cache_manager(p_expected), 'unused')) # Test that the TestStream is outputting to the correct PCollection. class TestStreamVisitor(PipelineVisitor): def __init__(self): self.output_tags = set() def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): transform = transform_node.transform if isinstance(transform, TestStream): self.output_tags = transform.output_tags v = TestStreamVisitor() actual_pipeline.visit(v) expected_output_tags = set([source_1_cache_key]) actual_output_tags = v.output_tags self.assertSetEqual(expected_output_tags, actual_output_tags) # Test that the pipeline is as expected. assert_pipeline_proto_equal( self, p_expected.to_runner_api(), instrumenter.instrumented_pipeline_proto()) def test_instrument_example_unbounded_pipeline_to_multiple_read_cache(self): """Tests that the instrumenter works for multiple unbounded sources. """ # Create the pipeline that will be instrumented. p_original = beam.Pipeline(interactive_runner.InteractiveRunner()) ie.current_env().set_cache_manager( StreamingCache(cache_dir=None), p_original) source_1 = p_original | 'source1' >> beam.io.ReadFromPubSub( subscription='projects/fake-project/subscriptions/fake_sub') source_2 = p_original | 'source2' >> beam.io.ReadFromPubSub( subscription='projects/fake-project/subscriptions/fake_sub') # pylint: disable=possibly-unused-variable pcoll_1 = source_1 | 'square1' >> beam.Map(lambda x: x * x) # pylint: disable=possibly-unused-variable pcoll_2 = source_2 | 'square2' >> beam.Map(lambda x: x * x) # Mock as if cacheable PCollections are cached. ib.watch(locals()) # This should be noop. utils.watch_sources(p_original) for name, pcoll in locals().items(): if not isinstance(pcoll, beam.pvalue.PCollection): continue cache_key = self.cache_key_of(name, pcoll) self._mock_write_cache(p_original, [], cache_key) # Instrument the original pipeline to create the pipeline the user will see. instrumenter = instr.build_pipeline_instrument(p_original) actual_pipeline = beam.Pipeline.from_runner_api( proto=instrumenter.instrumented_pipeline_proto(), runner=interactive_runner.InteractiveRunner(), options=None) # Now, build the expected pipeline which replaces the unbounded source with # a TestStream. source_1_cache_key = self.cache_key_of('source_1', source_1) source_2_cache_key = self.cache_key_of('source_2', source_2) p_expected = beam.Pipeline() test_stream = ( p_expected | TestStream( output_tags=[ self.cache_key_of('source_1', source_1), self.cache_key_of('source_2', source_2) ])) # pylint: disable=expression-not-assigned test_stream[source_1_cache_key] | 'square1' >> beam.Map(lambda x: x * x) # pylint: disable=expression-not-assigned test_stream[source_2_cache_key] | 'square2' >> beam.Map(lambda x: x * x) # Test that the TestStream is outputting to the correct PCollection. class TestStreamVisitor(PipelineVisitor): def __init__(self): self.output_tags = set() def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): transform = transform_node.transform if isinstance(transform, TestStream): self.output_tags = transform.output_tags v = TestStreamVisitor() actual_pipeline.visit(v) expected_output_tags = set([source_1_cache_key, source_2_cache_key]) actual_output_tags = v.output_tags self.assertSetEqual(expected_output_tags, actual_output_tags) # Test that the pipeline is as expected. assert_pipeline_proto_equal( self, p_expected.to_runner_api(), instrumenter.instrumented_pipeline_proto()) def test_pipeline_pruned_when_input_pcoll_is_cached(self): user_pipeline, init_pcoll, _ = self._example_pipeline() runner_pipeline = beam.Pipeline.from_runner_api( user_pipeline.to_runner_api(), user_pipeline.runner, None) ie.current_env().add_derived_pipeline(user_pipeline, runner_pipeline) # Mock as if init_pcoll is cached. init_pcoll_cache_key = self.cache_key_of('init_pcoll', init_pcoll) self._mock_write_cache( user_pipeline, [b'1', b'2', b'3'], init_pcoll_cache_key) ie.current_env().mark_pcollection_computed([init_pcoll]) # Build an instrument from the runner pipeline. pipeline_instrument = instr.build_pipeline_instrument(runner_pipeline) pruned_proto = pipeline_instrument.instrumented_pipeline_proto() # Skip the prune step for comparison, it should contain the sub-graph that # produces init_pcoll but not useful anymore. full_proto = pipeline_instrument._pipeline.to_runner_api() self.assertEqual( len( pruned_proto.components.transforms[ 'ref_AppliedPTransform_AppliedPTransform_1'].subtransforms), 5) assert_pipeline_proto_not_contain_top_level_transform( self, pruned_proto, 'Init Source') self.assertEqual( len( full_proto.components.transforms[ 'ref_AppliedPTransform_AppliedPTransform_1'].subtransforms), 6) assert_pipeline_proto_contain_top_level_transform( self, full_proto, 'Init-Source') def test_side_effect_pcoll_is_included(self): pipeline_with_side_effect = beam.Pipeline( interactive_runner.InteractiveRunner()) ie.current_env().set_cache_manager( InMemoryCache(), pipeline_with_side_effect) # Deliberately not assign the result to a variable to make it a # "side effect" transform. Note we never watch anything from # the pipeline defined locally either. # pylint: disable=bad-option-value,expression-not-assigned pipeline_with_side_effect | 'Init Create' >> beam.Create(range(10)) pipeline_instrument = instr.build_pipeline_instrument( pipeline_with_side_effect) self.assertTrue(pipeline_instrument._extended_targets) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/recording_manager.py0000644000175100017510000010061515206113032025711 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import collections import logging import os import threading import time import uuid import warnings from concurrent.futures import Future from concurrent.futures import ThreadPoolExecutor from typing import Any from typing import Optional from typing import Union import pandas as pd import apache_beam as beam from apache_beam.dataframe.frame_base import DeferredBase from apache_beam.options import pipeline_options from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners import runner from apache_beam.runners.interactive import background_caching_job as bcj from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import interactive_runner as ir from apache_beam.runners.interactive import pipeline_fragment as pf from apache_beam.runners.interactive import utils from apache_beam.runners.interactive.caching.cacheable import CacheKey from apache_beam.runners.interactive.display.pipeline_graph import PipelineGraph from apache_beam.runners.interactive.options import capture_control from apache_beam.runners.runner import PipelineState _LOGGER = logging.getLogger(__name__) try: import ipywidgets as widgets from IPython.display import HTML from IPython.display import clear_output from IPython.display import display IS_IPYTHON = True except ImportError: IS_IPYTHON = False class AsyncComputationResult: """Represents the result of an asynchronous computation.""" def __init__( self, future: Future, pcolls: set[beam.pvalue.PCollection], user_pipeline: beam.Pipeline, recording_manager: 'RecordingManager', ): self._future = future self._pcolls = pcolls self._user_pipeline = user_pipeline self._env = ie.current_env() self._recording_manager = recording_manager self._pipeline_result: Optional[beam.runners.runner.PipelineResult] = None self._display_id = str(uuid.uuid4()) self._output_widget = widgets.Output() if IS_IPYTHON else None self._cancel_button = ( widgets.Button(description='Cancel') if IS_IPYTHON else None) self._progress_bar = ( widgets.FloatProgress( value=0.0, min=0.0, max=1.0, description='Running:', bar_style='info', ) if IS_IPYTHON else None) self._cancel_requested = False if IS_IPYTHON: self._cancel_button.on_click(self._cancel_clicked) controls = widgets.VBox([ widgets.HBox([self._cancel_button, self._progress_bar]), self._output_widget, ]) display(controls, display_id=self._display_id) self.update_display('Initializing...') self._future.add_done_callback(self._on_done) def _cancel_clicked(self, b): self._cancel_requested = True self._cancel_button.disabled = True self.update_display('Cancel requested...') self.cancel() def update_display(self, msg: str, progress: Optional[float] = None): if not IS_IPYTHON: print(f'AsyncCompute: {msg}') return with self._output_widget: clear_output(wait=True) display(HTML(f'

{msg}

')) if progress is not None: self._progress_bar.value = progress if self.done(): self._cancel_button.disabled = True if self.exception(): self._progress_bar.bar_style = 'danger' self._progress_bar.description = 'Failed' elif self._future.cancelled(): self._progress_bar.bar_style = 'warning' self._progress_bar.description = 'Cancelled' else: self._progress_bar.bar_style = 'success' self._progress_bar.description = 'Done' elif self._cancel_requested: self._cancel_button.disabled = True self._progress_bar.description = 'Cancelling...' else: self._cancel_button.disabled = False def set_pipeline_result( self, pipeline_result: beam.runners.runner.PipelineResult): self._pipeline_result = pipeline_result if self._cancel_requested: self.cancel() def result(self, timeout=None): return self._future.result(timeout=timeout) def done(self): return self._future.done() def exception(self, timeout=None): try: return self._future.exception(timeout=timeout) except TimeoutError: return None def _on_done(self, future: Future): self._env.unmark_pcollection_computing(self._pcolls) self._recording_manager._async_computations.pop(self._display_id, None) if future.cancelled(): self.update_display('Computation Cancelled.', 1.0) return exc = future.exception() if exc: self.update_display(f'Error: {exc}', 1.0) _LOGGER.error('Asynchronous computation failed: %s', exc, exc_info=exc) else: self.update_display('Computation Finished Successfully.', 1.0) res = future.result() if res and res.state == PipelineState.DONE: self._env.mark_pcollection_computed(self._pcolls) else: _LOGGER.warning( 'Async computation finished but state is not DONE: %s', res.state if res else 'Unknown') def cancel(self): if self._future.done(): self.update_display('Cannot cancel: Computation already finished.') return False self._cancel_requested = True self._cancel_button.disabled = True self.update_display('Attempting to cancel...') if self._pipeline_result: try: # Check pipeline state before cancelling current_state = self._pipeline_result.state if PipelineState.is_terminal(current_state): self.update_display( 'Cannot cancel: Pipeline already in terminal state' f' {current_state}.') return False self._pipeline_result.cancel() self.update_display('Cancel signal sent to pipeline.') # The future will be cancelled by the runner if successful return True except Exception as e: self.update_display('Error sending cancel signal: %s', e) _LOGGER.warning('Error during pipeline cancel(): %s', e, exc_info=e) # Still try to cancel the future as a fallback return self._future.cancel() else: self.update_display('Pipeline not yet fully started, cancelling future.') return self._future.cancel() def __repr__(self): return ( f'") class ElementStream: """A stream of elements from a given PCollection.""" def __init__( self, pcoll: beam.pvalue.PCollection, var: str, cache_key: str, max_n: int, max_duration_secs: float): self._pcoll = pcoll self._cache_key = cache_key self._pipeline = ie.current_env().user_pipeline(pcoll.pipeline) self._var = var self._n = max_n self._duration_secs = max_duration_secs # A small state variable that when True, indicates that no more new elements # will be yielded if read() is called again. self._done = False @property def var(self) -> str: """Returns the variable named that defined this PCollection.""" return self._var @property def pcoll(self) -> beam.pvalue.PCollection: """Returns the PCollection that supplies this stream with data.""" return self._pcoll @property def cache_key(self) -> str: """Returns the cache key for this stream.""" return self._cache_key def display_id(self, suffix: str) -> str: """Returns a unique id able to be displayed in a web browser.""" return utils.obfuscate(self._cache_key, suffix) def is_computed(self) -> bool: # noqa: F821 """Returns True if no more elements will be recorded.""" return self._pcoll in ie.current_env().computed_pcollections def is_done(self) -> bool: # noqa: F821 """Returns True if no more new elements will be yielded.""" return self._done def read(self, tail: bool = True) -> Any: """Reads the elements currently recorded.""" # Get the cache manager and wait until the file exists. cache_manager = ie.current_env().get_cache_manager(self._pipeline) # Retrieve the coder for the particular PCollection which will be used to # decode elements read from cache. coder = cache_manager.load_pcoder('full', self._cache_key) # Read the elements from the cache. # Import limiters here to prevent a circular import. from apache_beam.runners.interactive.options.capture_limiters import CountLimiter from apache_beam.runners.interactive.options.capture_limiters import ProcessingTimeLimiter reader, _ = cache_manager.read('full', self._cache_key, tail=tail) # Because a single TestStreamFileRecord can yield multiple elements, we # limit the count again here in the to_element_list call. # # There are two ways of exiting this loop either a limiter was triggered or # all elements from the cache were read. In the latter situation, it may be # the case that the pipeline was still running. Thus, another invocation of # `read` will yield new elements. count_limiter = CountLimiter(self._n) time_limiter = ProcessingTimeLimiter(self._duration_secs) limiters = (count_limiter, time_limiter) for e in utils.to_element_list(reader, coder, include_window_info=True, n=self._n, include_time_events=True): # From the to_element_list we either get TestStreamPayload.Events if # include_time_events or decoded elements from the reader. Make sure we # only count the decoded elements to break early. if isinstance(e, beam_runner_api_pb2.TestStreamPayload.Event): time_limiter.update(e) else: count_limiter.update(e) yield e if any(l.is_triggered() for l in limiters): break # A limiter being triggered means that we have fulfilled the user's request. # This implies that reading from the cache again won't yield any new # elements. WLOG, this applies to the user pipeline being terminated. if any(l.is_triggered() for l in limiters) or ie.current_env().is_terminated(self._pipeline): self._done = True class Recording: """A group of PCollections from a given pipeline run.""" def __init__( self, user_pipeline: beam.Pipeline, pcolls: list[beam.pvalue.PCollection], # noqa: F821 result: 'beam.runner.PipelineResult', max_n: int, max_duration_secs: float, ): self._user_pipeline = user_pipeline self._result = result self._result_lock = threading.Lock() self._pcolls = pcolls pcoll_var = lambda pcoll: { v: k for k, v in utils.pcoll_by_name().items() }.get(pcoll, None) self._streams = { pcoll: ElementStream( pcoll, pcoll_var(pcoll), CacheKey.from_pcoll(pcoll_var(pcoll), pcoll).to_str(), max_n, max_duration_secs) for pcoll in pcolls } self._start = time.time() self._duration_secs = max_duration_secs self._set_computed = bcj.is_cache_complete(str(id(user_pipeline))) # Run a separate thread for marking the PCollections done. This is because # the pipeline run may be asynchronous. self._mark_computed = threading.Thread(target=self._mark_all_computed) self._mark_computed.daemon = True self._mark_computed.start() def _mark_all_computed(self) -> None: """Marks all the PCollections upon a successful pipeline run.""" if not self._result: return while not PipelineState.is_terminal(self._result.state): with self._result_lock: bcj = ie.current_env().get_background_caching_job(self._user_pipeline) if bcj and bcj.is_done(): self._result.wait_until_finish() elif time.time() - self._start >= self._duration_secs: self._result.cancel() self._result.wait_until_finish() elif all(s.is_done() for s in self._streams.values()): self._result.cancel() self._result.wait_until_finish() time.sleep(0.1) # Mark the PCollection as computed so that Interactive Beam wouldn't need to # re-compute. if self._result.state is PipelineState.DONE and self._set_computed: ie.current_env().mark_pcollection_computed(self._pcolls) def is_computed(self) -> bool: """Returns True if all PCollections are computed.""" return all(s.is_computed() for s in self._streams.values()) def stream(self, pcoll: beam.pvalue.PCollection) -> ElementStream: """Returns an ElementStream for a given PCollection.""" return self._streams[pcoll] def computed(self) -> None: """Returns all computed ElementStreams.""" return {p: s for p, s in self._streams.items() if s.is_computed()} def uncomputed(self) -> None: """Returns all uncomputed ElementStreams.""" return {p: s for p, s in self._streams.items() if not s.is_computed()} def cancel(self) -> None: """Cancels the recording.""" with self._result_lock: self._result.cancel() def wait_until_finish(self) -> None: """Waits until the pipeline is done and returns the final state. This also marks any PCollections as computed right away if the pipeline is successful. """ if not self._result: return beam.runners.runner.PipelineState.DONE self._mark_computed.join() return self._result.state def describe(self) -> dict[str, int]: """Returns a dictionary describing the cache and recording.""" cache_manager = ie.current_env().get_cache_manager(self._user_pipeline) size = sum( cache_manager.size('full', s.cache_key) for s in self._streams.values()) return {'size': size, 'duration': self._duration_secs} class RecordingManager: """Manages recordings of PCollections for a given pipeline.""" def __init__( self, user_pipeline: beam.Pipeline, pipeline_var: str = None, test_limiters: list['Limiter'] = None) -> None: # noqa: F821 self.user_pipeline: beam.Pipeline = user_pipeline self.pipeline_var: str = pipeline_var if pipeline_var else '' self._recordings: set[Recording] = set() self._start_time_sec: float = 0 self._test_limiters = test_limiters if test_limiters else [] self._executor = ThreadPoolExecutor(max_workers=os.cpu_count()) self._env = ie.current_env() self._async_computations: dict[str, AsyncComputationResult] = {} self._pipeline_graph = None def _execute_pipeline_fragment( self, pcolls_to_compute: set[beam.pvalue.PCollection], async_result: Optional['AsyncComputationResult'] = None, runner: runner.PipelineRunner = None, options: pipeline_options.PipelineOptions = None, ) -> beam.runners.runner.PipelineResult: """Synchronously executes a pipeline fragment for the given PCollections.""" merged_options = pipeline_options.PipelineOptions(**{ **self.user_pipeline.options.get_all_options( drop_default=True, retain_unknown_options=True ), **( options.get_all_options( drop_default=True, retain_unknown_options=True ) if options else {} ), }) fragment = pf.PipelineFragment( list(pcolls_to_compute), merged_options, runner=runner) if async_result: async_result.update_display('Building pipeline fragment...', 0.1) pipeline_to_run = fragment.deduce_fragment() if async_result: async_result.update_display('"Pipeline running, awaiting finish..."', 0.2) pipeline_result = pipeline_to_run.run() if async_result: async_result.set_pipeline_result(pipeline_result) pipeline_result.wait_until_finish() return pipeline_result def _run_async_computation( self, pcolls_to_compute: set[beam.pvalue.PCollection], async_result: 'AsyncComputationResult', wait_for_inputs: bool, runner: runner.PipelineRunner = None, options: pipeline_options.PipelineOptions = None, ): """The function to be run in the thread pool for async computation.""" try: if wait_for_inputs: if not self._wait_for_dependencies(pcolls_to_compute, async_result): raise RuntimeError('Dependency computation failed or was cancelled.') _LOGGER.info( 'Starting asynchronous computation for %d PCollections.', len(pcolls_to_compute)) pipeline_result = self._execute_pipeline_fragment( pcolls_to_compute, async_result, runner, options) # if pipeline_result.state == PipelineState.DONE: # self._env.mark_pcollection_computed(pcolls_to_compute) # _LOGGER.info( # 'Asynchronous computation finished successfully for' # f' {len(pcolls_to_compute)} PCollections.' # ) # else: # _LOGGER.error( # 'Asynchronous computation failed for' # f' {len(pcolls_to_compute)} PCollections. State:' # f' {pipeline_result.state}' # ) return pipeline_result except Exception as e: _LOGGER.exception('Exception during asynchronous computation: %s', e) raise # finally: # self._env.unmark_pcollection_computing(pcolls_to_compute) def _watch(self, pcolls: list[beam.pvalue.PCollection]) -> None: """Watch any pcollections not being watched. This allows for the underlying caching layer to identify the PCollection as something to be cached. """ watched_pcollections = set() watched_dataframes = set() for watching in ie.current_env().watching(): for _, val in watching: if isinstance(val, beam.pvalue.PCollection): watched_pcollections.add(val) elif isinstance(val, DeferredBase): watched_dataframes.add(val) # Convert them one-by-one to generate a unique label for each. This allows # caching at a more fine-grained granularity. # # TODO(https://github.com/apache/beam/issues/20929): investigate the mixing # pcollections in multiple pipelines error when using the default label. for df in watched_dataframes: pcoll, _ = utils.deferred_df_to_pcollection(df) watched_pcollections.add(pcoll) for pcoll in pcolls: if pcoll not in watched_pcollections: ie.current_env().watch( {'anonymous_pcollection_{}'.format(id(pcoll)): pcoll}) def _clear(self) -> None: """Clears the recording of all non-source PCollections.""" cache_manager = ie.current_env().get_cache_manager(self.user_pipeline) # Only clear the PCollections that aren't being populated from the # BackgroundCachingJob. computed = ie.current_env().computed_pcollections cacheables = [ c for c in utils.cacheables().values() if c.pcoll.pipeline is self.user_pipeline and c.pcoll not in computed ] all_cached = set(str(c.to_key()) for c in cacheables) source_pcolls = getattr(cache_manager, 'capture_keys', set()) to_clear = all_cached - source_pcolls self._clear_pcolls(cache_manager, set(to_clear)) def _clear_pcolls(self, cache_manager, pcolls): for pc in pcolls: cache_manager.clear('full', pc) def clear(self) -> None: """Clears all cached PCollections for this RecordingManager.""" cache_manager = ie.current_env().get_cache_manager(self.user_pipeline) if cache_manager: cache_manager.cleanup() def cancel(self: None) -> None: """Cancels the current background recording job.""" bcj.attempt_to_cancel_background_caching_job(self.user_pipeline) for r in self._recordings: r.wait_until_finish() self._recordings = set() # The recordings rely on a reference to the BCJ to correctly finish. So we # evict the BCJ after they complete. ie.current_env().evict_background_caching_job(self.user_pipeline) def describe(self) -> dict[str, int]: """Returns a dictionary describing the cache and recording.""" cache_manager = ie.current_env().get_cache_manager(self.user_pipeline) capture_size = getattr(cache_manager, 'capture_size', 0) descriptions = [r.describe() for r in self._recordings] size = sum(d['size'] for d in descriptions) + capture_size start = self._start_time_sec bcj = ie.current_env().get_background_caching_job(self.user_pipeline) if bcj: state = bcj.state else: state = PipelineState.STOPPED return { 'size': size, 'start': start, 'state': state, 'pipeline_var': self.pipeline_var } def record_pipeline(self) -> bool: """Starts a background caching job for this RecordingManager's pipeline.""" runner = self.user_pipeline.runner if isinstance(runner, ir.InteractiveRunner): runner = runner._underlying_runner if hasattr(runner, 'is_interactive'): runner.is_interactive() # Make sure that sources without a user reference are still cached. ie.current_env().add_user_pipeline(self.user_pipeline) utils.watch_sources(self.user_pipeline) # Attempt to run background caching job to record any sources. warnings.filterwarnings( 'ignore', 'options is deprecated since First stable release. References to ' '.options will not be supported', category=DeprecationWarning) if bcj.attempt_to_run_background_caching_job( runner, self.user_pipeline, options=self.user_pipeline.options, limiters=self._test_limiters): self._start_time_sec = time.time() return True return False def compute_async( self, pcolls: set[beam.pvalue.PCollection], wait_for_inputs: bool = True, blocking: bool = False, runner: runner.PipelineRunner = None, options: pipeline_options.PipelineOptions = None, force_compute: bool = False, ) -> Optional[AsyncComputationResult]: """Computes the given PCollections, potentially asynchronously.""" if force_compute: self._env.evict_computed_pcollections(self.user_pipeline) computed_pcolls = { pcoll for pcoll in pcolls if pcoll in self._env.computed_pcollections } computing_pcolls = { pcoll for pcoll in pcolls if self._env.is_pcollection_computing(pcoll) } pcolls_to_compute = pcolls - computed_pcolls - computing_pcolls if not pcolls_to_compute: _LOGGER.info( 'All requested PCollections are already computed or are being' ' computed.') return None self._watch(list(pcolls_to_compute)) self.record_pipeline() if blocking: self._env.mark_pcollection_computing(pcolls_to_compute) try: if wait_for_inputs: if not self._wait_for_dependencies(pcolls_to_compute): raise RuntimeError( 'Dependency computation failed or was cancelled.') pipeline_result = self._execute_pipeline_fragment( pcolls_to_compute, None, runner, options) if pipeline_result.state == PipelineState.DONE: self._env.mark_pcollection_computed(pcolls_to_compute) else: _LOGGER.error( 'Blocking computation failed. State: %s', pipeline_result.state) raise RuntimeError( 'Blocking computation failed. State: %s', pipeline_result.state) finally: self._env.unmark_pcollection_computing(pcolls_to_compute) return None else: # Asynchronous future = Future() async_result = AsyncComputationResult( future, pcolls_to_compute, self.user_pipeline, self) self._async_computations[async_result._display_id] = async_result self._env.mark_pcollection_computing(pcolls_to_compute) def task(): try: result = self._run_async_computation( pcolls_to_compute, async_result, wait_for_inputs, runner, options) future.set_result(result) except Exception as e: if not future.cancelled(): future.set_exception(e) self._executor.submit(task) return async_result def _get_pipeline_graph(self): """Lazily initializes and returns the PipelineGraph.""" if self._pipeline_graph is None: try: # Try to create the graph. self._pipeline_graph = PipelineGraph(self.user_pipeline) except (ImportError, NameError, AttributeError): # If pydot is missing, PipelineGraph() might crash. _LOGGER.warning( "Could not create PipelineGraph (pydot missing?). " \ "Async features disabled." ) self._pipeline_graph = None return self._pipeline_graph def _get_pcoll_id_map(self): """Creates a map from PCollection object to its ID in the proto.""" pcoll_to_id = {} graph = self._get_pipeline_graph() if graph and graph._pipeline_instrument: pcoll_to_id = graph._pipeline_instrument._pcoll_to_pcoll_id return {v: k for k, v in pcoll_to_id.items()} def _get_all_dependencies( self, pcolls: set[beam.pvalue.PCollection]) -> set[beam.pvalue.PCollection]: """Gets all upstream PCollection dependencies for the given set of PCollections.""" graph = self._get_pipeline_graph() if not graph: return set() analyzer = graph._pipeline_instrument if not analyzer: return set() pcoll_to_id = analyzer._pcoll_to_pcoll_id target_pcoll_ids = { pcoll_to_id.get(str(pcoll)) for pcoll in pcolls if str(pcoll) in pcoll_to_id } if not target_pcoll_ids: return set() # Build a map from PCollection ID to the actual PCollection object id_to_pcoll_obj = {} for _, inspectable in self._env.inspector.inspectables.items(): value = inspectable['value'] if isinstance(value, beam.pvalue.PCollection): pcoll_id = pcoll_to_id.get(str(value)) if pcoll_id: id_to_pcoll_obj[pcoll_id] = value dependencies = set() queue = collections.deque(target_pcoll_ids) visited_pcoll_ids = set(target_pcoll_ids) producers = graph._producers transforms = graph._pipeline_proto.components.transforms while queue: pcoll_id = queue.popleft() if pcoll_id not in producers: continue producer_id = producers[pcoll_id] transform_proto = transforms.get(producer_id) if not transform_proto: continue for input_pcoll_id in transform_proto.inputs.values(): if input_pcoll_id not in visited_pcoll_ids: visited_pcoll_ids.add(input_pcoll_id) queue.append(input_pcoll_id) dep_obj = id_to_pcoll_obj.get(input_pcoll_id) if dep_obj and dep_obj not in pcolls: dependencies.add(dep_obj) return dependencies def _wait_for_dependencies( self, pcolls: set[beam.pvalue.PCollection], async_result: Optional[AsyncComputationResult] = None, ) -> bool: """Waits for any dependencies of the given PCollections that are currently being computed.""" dependencies = self._get_all_dependencies(pcolls) computing_deps: dict[beam.pvalue.PCollection, AsyncComputationResult] = {} for dep in dependencies: if self._env.is_pcollection_computing(dep): for comp in self._async_computations.values(): if dep in comp._pcolls: computing_deps[dep] = comp break if not computing_deps: return True if async_result: async_result.update_display( 'Waiting for %d dependencies to finish...', len(computing_deps)) _LOGGER.info( 'Waiting for %d dependencies: %s', len(computing_deps), computing_deps.keys()) futures_to_wait = list( set(comp._future for comp in computing_deps.values())) try: for i, future in enumerate(futures_to_wait): if async_result: async_result.update_display( f'Waiting for dependency {i + 1}/{len(futures_to_wait)}...', progress=0.05 + 0.05 * (i / len(futures_to_wait)), ) future.result() if async_result: async_result.update_display('Dependencies finished.', progress=0.1) _LOGGER.info('Dependencies finished successfully.') return True except Exception as e: if async_result: async_result.update_display(f'Dependency failed: {e}') _LOGGER.error('Dependency computation failed: %s', e, exc_info=e) return False def record( self, pcolls: list[beam.pvalue.PCollection], *, max_n: int, max_duration: Union[int, str], runner: runner.PipelineRunner = None, options: pipeline_options.PipelineOptions = None, force_compute: bool = False) -> Recording: # noqa: F821 """Records the given PCollections.""" if not ie.current_env().options.enable_recording_replay: capture_control.evict_captured_data() if force_compute: ie.current_env().evict_computed_pcollections() # Assert that all PCollection come from the same user_pipeline. for pcoll in pcolls: assert pcoll.pipeline is self.user_pipeline, ( '{} belongs to a different user-defined pipeline ({}) than that of' ' other PCollections ({}).'.format( pcoll, pcoll.pipeline, self.user_pipeline)) if isinstance(max_duration, str) and max_duration != 'inf': max_duration_secs = pd.to_timedelta(max_duration).total_seconds() else: max_duration_secs = max_duration # Make sure that all PCollections to be shown are watched. If a PCollection # has not been watched, make up a variable name for that PCollection and # watch it. No validation is needed here because the watch logic can handle # arbitrary variables. self._watch(pcolls) self.record_pipeline() # Get the subset of computed PCollections. These do not to be recomputed. computed_pcolls = set( pcoll for pcoll in pcolls if pcoll in ie.current_env().computed_pcollections) # Start a pipeline fragment to start computing the PCollections. uncomputed_pcolls = set(pcolls).difference(computed_pcolls) if uncomputed_pcolls: if not self._wait_for_dependencies(uncomputed_pcolls): raise RuntimeError( 'Cannot record because a dependency failed to compute' ' asynchronously.') self._clear() merged_options = pipeline_options.PipelineOptions( **{ **self.user_pipeline.options.get_all_options( drop_default=True, retain_unknown_options=True), **options.get_all_options( drop_default=True, retain_unknown_options=True) }) if options else self.user_pipeline.options cache_path = ie.current_env().options.cache_root is_remote_run = cache_path and ie.current_env( ).options.cache_root.startswith('gs://') pf.PipelineFragment( list(uncomputed_pcolls), merged_options, runner=runner).run(blocking=is_remote_run) result = ie.current_env().pipeline_result(self.user_pipeline) else: result = None recording = Recording( self.user_pipeline, pcolls, result, max_n, max_duration_secs) self._recordings.add(recording) return recording def read( self, pcoll_name: str, pcoll: beam.pvalue.PValue, max_n: int, max_duration_secs: float) -> Union[None, ElementStream]: # noqa: F821 """Reads an ElementStream of a computed PCollection. Returns None if an error occurs. The caller is responsible of validating if the given pcoll_name and pcoll can identify a watched and computed PCollection without ambiguity in the notebook. """ try: cache_key = CacheKey.from_pcoll(pcoll_name, pcoll).to_str() return ElementStream( pcoll, pcoll_name, cache_key, max_n, max_duration_secs) except (KeyboardInterrupt, SystemExit): raise except Exception as e: # Caller should handle all validations. Here to avoid redundant # validations, simply log errors if caller fails to do so. _LOGGER.error(str(e)) return None ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/recording_manager_test.py0000644000175100017510000012013615206113032026750 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import time import unittest from concurrent.futures import Future from unittest.mock import MagicMock from unittest.mock import call from unittest.mock import patch import apache_beam as beam from apache_beam import coders from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.portability.api import beam_interactive_api_pb2 from apache_beam.runners.interactive import background_caching_job as bcj from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.caching.cacheable import CacheKey from apache_beam.runners.interactive.interactive_runner import InteractiveRunner from apache_beam.runners.interactive.options.capture_limiters import Limiter from apache_beam.runners.interactive.recording_manager import _LOGGER from apache_beam.runners.interactive.recording_manager import AsyncComputationResult from apache_beam.runners.interactive.recording_manager import ElementStream from apache_beam.runners.interactive.recording_manager import Recording from apache_beam.runners.interactive.recording_manager import RecordingManager from apache_beam.runners.interactive.testing.test_cache_manager import FileRecordsBuilder from apache_beam.runners.interactive.testing.test_cache_manager import InMemoryCache from apache_beam.runners.runner import PipelineState from apache_beam.testing.test_stream import TestStream from apache_beam.testing.test_stream import WindowedValueHolder from apache_beam.transforms.window import GlobalWindow from apache_beam.utils.timestamp import MIN_TIMESTAMP from apache_beam.utils.windowed_value import WindowedValue @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') class AsyncComputationResultTest(unittest.TestCase): def setUp(self): self.mock_future = MagicMock(spec=Future) self.pcolls = {MagicMock(spec=beam.pvalue.PCollection)} self.user_pipeline = MagicMock(spec=beam.Pipeline) self.recording_manager = MagicMock(spec=RecordingManager) self.recording_manager._async_computations = {} self.env = ie.InteractiveEnvironment() patch.object(ie, 'current_env', return_value=self.env).start() self.mock_button = patch('ipywidgets.Button', autospec=True).start() self.mock_float_progress = patch( 'ipywidgets.FloatProgress', autospec=True).start() self.mock_output = patch('ipywidgets.Output', autospec=True).start() self.mock_hbox = patch('ipywidgets.HBox', autospec=True).start() self.mock_vbox = patch('ipywidgets.VBox', autospec=True).start() self.mock_display = patch( 'apache_beam.runners.interactive.recording_manager.display', autospec=True).start() self.mock_clear_output = patch( 'apache_beam.runners.interactive.recording_manager.clear_output', autospec=True).start() self.mock_html = patch( 'apache_beam.runners.interactive.recording_manager.HTML', autospec=True).start() self.addCleanup(patch.stopall) @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', False) def test_async_result_init_non_ipython(self): async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) self.assertIsNotNone(async_res) self.mock_future.add_done_callback.assert_called_once() self.assertIsNone(async_res._cancel_button) def test_on_done_success(self): async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) mock_pipeline_result = MagicMock() mock_pipeline_result.state = PipelineState.DONE self.mock_future.result.return_value = mock_pipeline_result self.mock_future.exception.return_value = None self.mock_future.cancelled.return_value = False async_res._display_id = 'test_id' self.recording_manager._async_computations['test_id'] = async_res with patch.object( self.env, 'unmark_pcollection_computing' ) as mock_unmark, patch.object( self.env, 'mark_pcollection_computed' ) as mock_mark_computed, patch.object( async_res, 'update_display' ) as mock_update: async_res._on_done(self.mock_future) mock_unmark.assert_called_once_with(self.pcolls) mock_mark_computed.assert_called_once_with(self.pcolls) self.assertNotIn('test_id', self.recording_manager._async_computations) mock_update.assert_called_with('Computation Finished Successfully.', 1.0) @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', False) def test_on_done_failure(self): async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) test_exception = ValueError('Test') self.mock_future.exception.return_value = test_exception self.mock_future.cancelled.return_value = False with patch.object( self.env, 'unmark_pcollection_computing' ) as mock_unmark, patch.object( self.env, 'mark_pcollection_computed' ) as mock_mark_computed: async_res._on_done(self.mock_future) mock_unmark.assert_called_once_with(self.pcolls) mock_mark_computed.assert_not_called() @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', False) def test_on_done_cancelled(self): async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) self.mock_future.cancelled.return_value = True with patch.object(self.env, 'unmark_pcollection_computing') as mock_unmark: async_res._on_done(self.mock_future) mock_unmark.assert_called_once_with(self.pcolls) @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', True) def test_cancel(self): async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) mock_pipeline_result = MagicMock() mock_pipeline_result.state = PipelineState.RUNNING async_res.set_pipeline_result(mock_pipeline_result) self.mock_future.done.return_value = False self.assertTrue(async_res.cancel()) mock_pipeline_result.cancel.assert_called_once() self.assertTrue(async_res._cancel_requested) self.assertTrue(async_res._cancel_button.disabled) @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', False) def test_cancel_already_done(self): async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) self.mock_future.done.return_value = True self.assertFalse(async_res.cancel()) @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', True) @patch('apache_beam.runners.interactive.recording_manager.display') @patch('ipywidgets.Button') @patch('ipywidgets.FloatProgress') @patch('ipywidgets.Output') @patch('ipywidgets.HBox') @patch('ipywidgets.VBox') def test_async_result_init_ipython( self, mock_vbox, mock_hbox, mock_output, mock_progress, mock_button, mock_display, ): mock_btn_instance = mock_button.return_value async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) self.assertIsNotNone(async_res) mock_button.assert_called_once_with(description='Cancel') mock_progress.assert_called_once() mock_output.assert_called_once() mock_hbox.assert_called_once() mock_vbox.assert_called_once() mock_display.assert_called() mock_btn_instance.on_click.assert_called_once_with( async_res._cancel_clicked) self.mock_future.add_done_callback.assert_called_once() @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', True) @patch( 'apache_beam.runners.interactive.recording_manager.display', MagicMock()) @patch('ipywidgets.Button', MagicMock()) @patch('ipywidgets.FloatProgress', MagicMock()) @patch('ipywidgets.Output', MagicMock()) def test_cancel_clicked(self): async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) with patch.object(async_res, 'cancel') as mock_cancel, patch.object( async_res, 'update_display' ) as mock_update: async_res._cancel_clicked(None) self.assertTrue(async_res._cancel_requested) self.assertTrue(async_res._cancel_button.disabled) mock_update.assert_called_once_with('Cancel requested...') mock_cancel.assert_called_once() @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', False) def test_update_display_non_ipython(self): async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) with patch('builtins.print') as mock_print: async_res.update_display('Test Message') mock_print.assert_called_once_with('AsyncCompute: Test Message') @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', True) def test_update_display_ipython(self): mock_prog_instance = self.mock_float_progress.return_value mock_btn_instance = self.mock_button.return_value async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) update_call_count = 1 self.assertEqual(self.mock_clear_output.call_count, update_call_count) # State: Running self.mock_future.done.return_value = False async_res._cancel_requested = False async_res.update_display('Running Test', 0.5) update_call_count += 1 self.mock_display.assert_called() self.assertEqual(self.mock_clear_output.call_count, update_call_count) self.assertEqual(mock_prog_instance.value, 0.5) self.assertFalse(mock_btn_instance.disabled) self.mock_html.assert_called_with('

Running Test

') # State: Done Success self.mock_future.done.return_value = True self.mock_future.exception.return_value = None self.mock_future.cancelled.return_value = False async_res.update_display('Done') update_call_count += 1 self.assertEqual(self.mock_clear_output.call_count, update_call_count) self.assertTrue(mock_btn_instance.disabled) self.assertEqual(mock_prog_instance.bar_style, 'success') self.assertEqual(mock_prog_instance.description, 'Done') # State: Done Failed self.mock_future.exception.return_value = Exception() async_res.update_display('Failed') update_call_count += 1 self.assertEqual(self.mock_clear_output.call_count, update_call_count) self.assertEqual(mock_prog_instance.bar_style, 'danger') self.assertEqual(mock_prog_instance.description, 'Failed') # State: Done Cancelled self.mock_future.exception.return_value = None self.mock_future.cancelled.return_value = True async_res.update_display('Cancelled') update_call_count += 1 self.assertEqual(self.mock_clear_output.call_count, update_call_count) self.assertEqual(mock_prog_instance.bar_style, 'warning') self.assertEqual(mock_prog_instance.description, 'Cancelled') # State: Cancelling self.mock_future.done.return_value = False async_res._cancel_requested = True async_res.update_display('Cancelling') update_call_count += 1 self.assertEqual(self.mock_clear_output.call_count, update_call_count) self.assertTrue(mock_btn_instance.disabled) self.assertEqual(mock_prog_instance.description, 'Cancelling...') @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', False) def test_set_pipeline_result_cancel_requested(self): async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) async_res._cancel_requested = True mock_pipeline_result = MagicMock() with patch.object(async_res, 'cancel') as mock_cancel: async_res.set_pipeline_result(mock_pipeline_result) self.assertIs(async_res._pipeline_result, mock_pipeline_result) mock_cancel.assert_called_once() @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', False) def test_exception_timeout(self): async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) self.mock_future.exception.side_effect = TimeoutError self.assertIsNone(async_res.exception(timeout=0.1)) @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', False) @patch.object(_LOGGER, 'warning') def test_on_done_not_done_state(self, mock_logger_warning): async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) mock_pipeline_result = MagicMock() mock_pipeline_result.state = PipelineState.FAILED self.mock_future.result.return_value = mock_pipeline_result self.mock_future.exception.return_value = None self.mock_future.cancelled.return_value = False with patch.object(self.env, 'mark_pcollection_computed') as mock_mark_computed: async_res._on_done(self.mock_future) mock_mark_computed.assert_not_called() mock_logger_warning.assert_called_once() @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', True) def test_cancel_no_pipeline_result(self): async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) self.mock_future.done.return_value = False self.mock_future.cancel.return_value = True with patch.object(async_res, 'update_display') as mock_update: self.assertTrue(async_res.cancel()) mock_update.assert_any_call( 'Pipeline not yet fully started, cancelling future.') self.mock_future.cancel.assert_called_once() @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', True) def test_cancel_pipeline_terminal_state(self): async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) self.mock_future.done.return_value = False mock_pipeline_result = MagicMock() mock_pipeline_result.state = PipelineState.DONE async_res.set_pipeline_result(mock_pipeline_result) with patch.object(async_res, 'update_display') as mock_update: self.assertFalse(async_res.cancel()) mock_update.assert_any_call( 'Cannot cancel: Pipeline already in terminal state DONE.') mock_pipeline_result.cancel.assert_not_called() @patch('apache_beam.runners.interactive.recording_manager.IS_IPYTHON', True) @patch.object(_LOGGER, 'warning') @patch.object(AsyncComputationResult, 'update_display') def test_cancel_pipeline_exception( self, mock_update_display, mock_logger_warning): async_res = AsyncComputationResult( self.mock_future, self.pcolls, self.user_pipeline, self.recording_manager, ) self.mock_future.done.return_value = False mock_pipeline_result = MagicMock() mock_pipeline_result.state = PipelineState.RUNNING test_exception = RuntimeError('Cancel Failed') mock_pipeline_result.cancel.side_effect = test_exception async_res.set_pipeline_result(mock_pipeline_result) self.mock_future.cancel.return_value = False self.assertFalse(async_res.cancel()) expected_calls = [ call('Initializing...'), # From __init__ call('Attempting to cancel...'), # From cancel() start call('Error sending cancel signal: %s', test_exception) # From except block ] mock_update_display.assert_has_calls(expected_calls, any_order=False) mock_logger_warning.assert_called_once() self.mock_future.cancel.assert_called_once() class MockPipelineResult(beam.runners.runner.PipelineResult): """Mock class for controlling a PipelineResult.""" def __init__(self): self._state = PipelineState.RUNNING def wait_until_finish(self): pass def set_state(self, state): self._state = state @property def state(self): return self._state def cancel(self): self._state = PipelineState.CANCELLED class ElementStreamTest(unittest.TestCase): def setUp(self): self.cache = InMemoryCache() self.p = beam.Pipeline() self.pcoll = self.p | beam.Create([]) self.cache_key = str(CacheKey('pcoll', '', '', '')) # Create a MockPipelineResult to control the state of a fake run of the # pipeline. self.mock_result = MockPipelineResult() ie.current_env().add_user_pipeline(self.p) ie.current_env().set_pipeline_result(self.p, self.mock_result) ie.current_env().set_cache_manager(self.cache, self.p) def test_read(self): """Test reading and if a stream is done no more elements are returned.""" self.mock_result.set_state(PipelineState.DONE) self.cache.write(['expected'], 'full', self.cache_key) self.cache.save_pcoder(None, 'full', self.cache_key) stream = ElementStream( self.pcoll, '', self.cache_key, max_n=1, max_duration_secs=1) self.assertFalse(stream.is_done()) self.assertEqual(list(stream.read())[0], 'expected') self.assertTrue(stream.is_done()) def test_done_if_terminated(self): """Test that terminating the job sets the stream as done.""" self.cache.write(['expected'], 'full', self.cache_key) self.cache.save_pcoder(None, 'full', self.cache_key) stream = ElementStream( self.pcoll, '', self.cache_key, max_n=100, max_duration_secs=10) self.assertFalse(stream.is_done()) self.assertEqual(list(stream.read(tail=False))[0], 'expected') # The limiters were not reached, so the stream is not done yet. self.assertFalse(stream.is_done()) self.mock_result.set_state(PipelineState.DONE) self.assertEqual(list(stream.read(tail=False))[0], 'expected') # The underlying pipeline is terminated, so the stream won't yield new # elements. self.assertTrue(stream.is_done()) def test_read_n(self): """Test that the stream only reads 'n' elements.""" self.mock_result.set_state(PipelineState.DONE) self.cache.write(list(range(5)), 'full', self.cache_key) self.cache.save_pcoder(None, 'full', self.cache_key) stream = ElementStream( self.pcoll, '', self.cache_key, max_n=1, max_duration_secs=1) self.assertEqual(list(stream.read()), [0]) self.assertTrue(stream.is_done()) stream = ElementStream( self.pcoll, '', self.cache_key, max_n=2, max_duration_secs=1) self.assertEqual(list(stream.read()), [0, 1]) self.assertTrue(stream.is_done()) stream = ElementStream( self.pcoll, '', self.cache_key, max_n=5, max_duration_secs=1) self.assertEqual(list(stream.read()), list(range(5))) self.assertTrue(stream.is_done()) # Test that if the user asks for more than in the cache it still returns. stream = ElementStream( self.pcoll, '', self.cache_key, max_n=10, max_duration_secs=1) self.assertEqual(list(stream.read()), list(range(5))) self.assertTrue(stream.is_done()) def test_read_duration(self): """Test that the stream only reads a 'duration' of elements.""" def as_windowed_value(element): return WindowedValueHolder(WindowedValue(element, 0, [])) values = (FileRecordsBuilder(tag=self.cache_key) .advance_processing_time(1) .add_element(element=as_windowed_value(0), event_time_secs=0) .advance_processing_time(1) .add_element(element=as_windowed_value(1), event_time_secs=1) .advance_processing_time(1) .add_element(element=as_windowed_value(2), event_time_secs=3) .advance_processing_time(1) .add_element(element=as_windowed_value(3), event_time_secs=4) .advance_processing_time(1) .add_element(element=as_windowed_value(4), event_time_secs=5) .build()) # yapf: disable values = [ v.recorded_event for v in values if isinstance(v, beam_interactive_api_pb2.TestStreamFileRecord) ] self.mock_result.set_state(PipelineState.DONE) self.cache.write(values, 'full', self.cache_key) self.cache.save_pcoder(coders.FastPrimitivesCoder(), 'full', self.cache_key) # The following tests a progression of reading different durations from the # cache. stream = ElementStream( self.pcoll, '', self.cache_key, max_n=100, max_duration_secs=1) self.assertSequenceEqual([e.value for e in stream.read()], [0]) stream = ElementStream( self.pcoll, '', self.cache_key, max_n=100, max_duration_secs=2) self.assertSequenceEqual([e.value for e in stream.read()], [0, 1]) stream = ElementStream( self.pcoll, '', self.cache_key, max_n=100, max_duration_secs=10) self.assertSequenceEqual([e.value for e in stream.read()], [0, 1, 2, 3, 4]) class RecordingTest(unittest.TestCase): def test_computed(self): """Tests that a PCollection is marked as computed only in a complete state. Because the background caching job is now long-lived, repeated runs of a PipelineFragment may yield different results for the same PCollection. """ p = beam.Pipeline(InteractiveRunner()) elems = p | beam.Create([0, 1, 2]) ib.watch(locals()) # Create a MockPipelineResult to control the state of a fake run of the # pipeline. mock_result = MockPipelineResult() ie.current_env().track_user_pipelines() ie.current_env().set_pipeline_result(p, mock_result) # Create a mock BackgroundCachingJob that will control whether to set the # PCollections as computed or not. bcj_mock_result = MockPipelineResult() background_caching_job = bcj.BackgroundCachingJob(bcj_mock_result, []) # Create a recording. recording = Recording( p, [elems], mock_result, max_n=10, max_duration_secs=60) # The background caching job and the recording isn't done yet so there may # be more elements to be recorded. self.assertFalse(recording.is_computed()) self.assertFalse(recording.computed()) self.assertTrue(recording.uncomputed()) # The recording is finished but the background caching job is not. There # may still be more elements to record, or the intermediate PCollection may # have stopped caching in an incomplete state, e.g. before a window could # fire. mock_result.set_state(PipelineState.DONE) recording.wait_until_finish() self.assertFalse(recording.is_computed()) self.assertFalse(recording.computed()) self.assertTrue(recording.uncomputed()) # The background caching job finished before we started a recording which # is a sure signal that there will be no more elements. bcj_mock_result.set_state(PipelineState.DONE) ie.current_env().set_background_caching_job(p, background_caching_job) recording = Recording( p, [elems], mock_result, max_n=10, max_duration_secs=60) recording.wait_until_finish() # There are no more elements and the recording finished, meaning that the # intermediate PCollections are in a complete state. They can now be marked # as computed. self.assertTrue(recording.is_computed()) self.assertTrue(recording.computed()) self.assertFalse(recording.uncomputed()) def test_describe(self): p = beam.Pipeline(InteractiveRunner()) numbers = p | 'numbers' >> beam.Create([0, 1, 2]) letters = p | 'letters' >> beam.Create(['a', 'b', 'c']) ib.watch(locals()) # Create a MockPipelineResult to control the state of a fake run of the # pipeline. mock_result = MockPipelineResult() ie.current_env().track_user_pipelines() ie.current_env().set_pipeline_result(p, mock_result) cache_manager = InMemoryCache() ie.current_env().set_cache_manager(cache_manager, p) # Create a recording with an arbitrary start time. recording = Recording( p, [numbers, letters], mock_result, max_n=10, max_duration_secs=60) # Get the cache key of the stream and write something to cache. This is # so that a pipeline doesn't have to run in the test. numbers_stream = recording.stream(numbers) cache_manager.write([0, 1, 2], 'full', numbers_stream.cache_key) cache_manager.save_pcoder(None, 'full', numbers_stream.cache_key) letters_stream = recording.stream(letters) cache_manager.write(['a', 'b', 'c'], 'full', letters_stream.cache_key) cache_manager.save_pcoder(None, 'full', letters_stream.cache_key) # Get the description. description = recording.describe() size = description['size'] self.assertEqual( size, cache_manager.size('full', numbers_stream.cache_key) + cache_manager.size('full', letters_stream.cache_key)) @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') class RecordingManagerTest(unittest.TestCase): def test_basic_execution(self): """A basic pipeline to be used as a smoke test.""" # Create the pipeline that will emit 0, 1, 2. p = beam.Pipeline(InteractiveRunner()) numbers = p | 'numbers' >> beam.Create([0, 1, 2]) letters = p | 'letters' >> beam.Create(['a', 'b', 'c']) # Watch the pipeline and PCollections. This is normally done in a notebook # environment automatically, but we have to do it manually here. ib.watch(locals()) ie.current_env().track_user_pipelines() # Create the recording objects. By calling `record` a new PipelineFragment # is started to compute the given PCollections and cache to disk. rm = RecordingManager(p) numbers_recording = rm.record([numbers], max_n=3, max_duration=500) numbers_stream = numbers_recording.stream(numbers) numbers_recording.wait_until_finish() # Once the pipeline fragment completes, we can read from the stream and know # that all elements were written to cache. elems = list(numbers_stream.read()) expected_elems = [ WindowedValue(i, MIN_TIMESTAMP, (GlobalWindow(), )) for i in range(3) ] self.assertListEqual(elems, expected_elems) # Make an extra recording and test the description. letters_recording = rm.record([letters], max_n=3, max_duration=500) letters_recording.wait_until_finish() self.assertEqual( rm.describe()['size'], numbers_recording.describe()['size'] + letters_recording.describe()['size']) rm.cancel() def test_duration_parsing(self): p = beam.Pipeline(InteractiveRunner()) elems = p | beam.Create([0, 1, 2]) # Watch the pipeline and PCollections. This is normally done in a notebook # environment automatically, but we have to do it manually here. ib.watch(locals()) ie.current_env().track_user_pipelines() # Create the recording objects. rm = RecordingManager(p) recording = rm.record([elems], max_n=3, max_duration='500s') recording.wait_until_finish() # Assert that the duration was parsed correctly to integer seconds. self.assertEqual(recording.describe()['duration'], 500) def test_cancel_stops_recording(self): # Add the TestStream so that it can be cached. ib.options.recordable_sources.add(TestStream) p = beam.Pipeline( InteractiveRunner(), options=PipelineOptions(streaming=True)) elems = ( p | TestStream().advance_watermark_to(0).advance_processing_time( 1).add_elements(list(range(10))).advance_processing_time(1)) squares = elems | beam.Map(lambda x: x**2) # Watch the local scope for Interactive Beam so that referenced PCollections # will be cached. ib.watch(locals()) # This is normally done in the interactive_utils when a transform is # applied but needs an IPython environment. So we manually run this here. ie.current_env().track_user_pipelines() class SemaphoreLimiter(Limiter): def __init__(self): self.triggered = False def is_triggered(self): return self.triggered # Get the recording then the BackgroundCachingJob. semaphore_limiter = SemaphoreLimiter() rm = RecordingManager(p, test_limiters=[semaphore_limiter]) rm.record([squares], max_n=10, max_duration=500) # The BackgroundCachingJob is still waiting for more elements, so it isn't # done yet. bcj = ie.current_env().get_background_caching_job(p) self.assertFalse(bcj.is_done()) # Assert that something was read and that the BackgroundCachingJob was # sucessfully stopped. # self.assertTrue(list(recording.stream(squares).read())) semaphore_limiter.triggered = True rm.cancel() self.assertTrue(bcj.is_done()) def test_recording_manager_clears_cache(self): """Tests that the RecordingManager clears the cache before recording. A job may have incomplete PCollections when the job terminates. Clearing the cache ensures that correct results are computed every run. """ # Add the TestStream so that it can be cached. ib.options.recordable_sources.add(TestStream) p = beam.Pipeline( InteractiveRunner(), options=PipelineOptions(streaming=True)) elems = ( p | TestStream().advance_watermark_to(0).advance_processing_time( 1).add_elements(list(range(10))).advance_processing_time(1)) squares = elems | beam.Map(lambda x: x**2) # Watch the local scope for Interactive Beam so that referenced PCollections # will be cached. ib.watch(locals()) # This is normally done in the interactive_utils when a transform is # applied but needs an IPython environment. So we manually run this here. ie.current_env().track_user_pipelines() # Do the first recording to get the timestamp of the first time the fragment # was run. rm = RecordingManager(p) # Set up a mock for the Cache's clear function which will be used to clear # uncomputed PCollections. rm._clear_pcolls = MagicMock() rm.record([squares], max_n=1, max_duration=500) rm.cancel() # Assert that the cache cleared the PCollection. rm._clear_pcolls.assert_any_call( unittest.mock.ANY, # elems is unbounded source populated by the background job, thus not # cleared. {CacheKey.from_pcoll('squares', squares).to_str()}) def test_clear(self): p1 = beam.Pipeline(InteractiveRunner()) elems_1 = p1 | 'elems 1' >> beam.Create([0, 1, 2]) ib.watch(locals()) ie.current_env().track_user_pipelines() recording_manager = RecordingManager(p1) recording = recording_manager.record([elems_1], max_n=3, max_duration=500) recording.wait_until_finish() record_describe = recording_manager.describe() self.assertGreater(record_describe['size'], 0) recording_manager.clear() self.assertEqual(recording_manager.describe()['size'], 0) def test_clear_specific_pipeline(self): """Tests that clear can empty the cache for a specific pipeline.""" # Create two pipelines so we can check that clearing the cache won't clear # all defined pipelines. p1 = beam.Pipeline(InteractiveRunner()) elems_1 = p1 | 'elems 1' >> beam.Create([0, 1, 2]) p2 = beam.Pipeline(InteractiveRunner()) elems_2 = p2 | 'elems 2' >> beam.Create([0, 1, 2]) # Watch the pipeline and PCollections. This is normally done in a notebook # environment automatically, but we have to do it manually here. ib.watch(locals()) ie.current_env().track_user_pipelines() # Create the recording objects. By calling `record` a new PipelineFragment # is started to compute the given PCollections and cache to disk. rm_1 = RecordingManager(p1) recording = rm_1.record([elems_1], max_n=3, max_duration=500) recording.wait_until_finish() rm_2 = RecordingManager(p2) recording = rm_2.record([elems_2], max_n=3, max_duration=500) recording.wait_until_finish() # Assert that clearing only one recording clears that recording. if rm_1.describe()['state'] == PipelineState.STOPPED \ and rm_2.describe()['state'] == PipelineState.STOPPED: self.assertGreater(rm_1.describe()['size'], 0) self.assertGreater(rm_2.describe()['size'], 0) rm_1.clear() self.assertEqual(rm_1.describe()['size'], 0) self.assertGreater(rm_2.describe()['size'], 0) rm_2.clear() self.assertEqual(rm_2.describe()['size'], 0) def test_record_pipeline(self): # Add the TestStream so that it can be cached. ib.options.recordable_sources.add(TestStream) p = beam.Pipeline( InteractiveRunner(), options=PipelineOptions(streaming=True)) # pylint: disable=unused-variable _ = (p | TestStream() .advance_watermark_to(0) .advance_processing_time(1) .add_elements(list(range(10))) .advance_processing_time(1)) # yapf: disable # Watch the local scope for Interactive Beam so that referenced PCollections # will be cached. ib.watch(locals()) # This is normally done in the interactive_utils when a transform is # applied but needs an IPython environment. So we manually run this here. ie.current_env().track_user_pipelines() # Create a lmiter that stops the background caching job when something is # written to cache. This is used to make ensure that the pipeline is # functioning properly and that there are no data races with the test. class SizeLimiter(Limiter): def __init__(self, p): self.pipeline = p self._rm = None def set_recording_manager(self, rm): self._rm = rm def is_triggered(self): return self._rm.describe()['size'] > 0 if self._rm else False # Do the first recording to get the timestamp of the first time the fragment # was run. size_limiter = SizeLimiter(p) rm = RecordingManager(p, test_limiters=[size_limiter]) size_limiter.set_recording_manager(rm) self.assertEqual(rm.describe()['state'], PipelineState.STOPPED) self.assertTrue(rm.record_pipeline()) # A recording is in progress, no need to start another one. self.assertFalse(rm.record_pipeline()) for _ in range(60): if rm.describe()['state'] == PipelineState.CANCELLED: break time.sleep(1) self.assertTrue( rm.describe()['state'] == PipelineState.CANCELLED, 'Test timed out waiting for pipeline to be cancelled. This indicates ' 'that the BackgroundCachingJob did not cache anything.') @patch( 'apache_beam.runners.interactive.recording_manager.' 'RecordingManager._clear_pcolls', return_value=None) @patch( 'apache_beam.runners.interactive.pipeline_fragment.' 'PipelineFragment.run', return_value=None) def test_record_detects_remote_runner( self, mock_pipeline_fragment, mock_clear_pcolls): """Tests that a remote runner is detected, resulting in the PipelineFragment instance to have blocking enabled.""" # Create the pipeline that will emit 0, 1, 2. p = beam.Pipeline(InteractiveRunner()) numbers = p | 'numbers' >> beam.Create([0, 1, 2]) # Set the cache directory for Interactive Beam to be in a GCS bucket. ib.options.cache_root = 'gs://test-bucket/' # Create the recording objects. By calling `record` a new PipelineFragment # is started to compute the given PCollections and cache to disk. rm = RecordingManager(p) # Run record() and check if the PipelineFragment.run had blocking set to # True due to the GCS cache_root value. rm.record([numbers], max_n=3, max_duration=500) mock_pipeline_fragment.assert_called_with(blocking=True) # Reset cache_root value. ib.options.cache_root = None def test_compute_async_blocking(self): p = beam.Pipeline(InteractiveRunner()) pcoll = p | beam.Create([1, 2, 3]) ib.watch(locals()) ie.current_env().track_user_pipelines() rm = RecordingManager(p) with patch.object(rm, '_execute_pipeline_fragment') as mock_execute: mock_result = MagicMock() mock_result.state = PipelineState.DONE mock_execute.return_value = mock_result res = rm.compute_async({pcoll}, blocking=True) self.assertIsNone(res) mock_execute.assert_called_once() self.assertTrue(pcoll in ie.current_env().computed_pcollections) @patch( 'apache_beam.runners.interactive.recording_manager.AsyncComputationResult' ) @patch( 'apache_beam.runners.interactive.recording_manager.ThreadPoolExecutor.' 'submit') def test_compute_async_non_blocking(self, mock_submit, mock_async_result_cls): p = beam.Pipeline(InteractiveRunner()) pcoll = p | beam.Create([1, 2, 3]) ib.watch(locals()) ie.current_env().track_user_pipelines() rm = RecordingManager(p) mock_async_res_instance = mock_async_result_cls.return_value # Capture the task task_submitted = None def capture_task(task): nonlocal task_submitted task_submitted = task # Return a mock future return MagicMock() mock_submit.side_effect = capture_task with patch.object( rm, '_wait_for_dependencies', return_value=True ), patch.object( rm, '_execute_pipeline_fragment' ) as _, patch.object( ie.current_env(), 'mark_pcollection_computing', wraps=ie.current_env().mark_pcollection_computing, ) as wrapped_mark: res = rm.compute_async({pcoll}, blocking=False) wrapped_mark.assert_called_once_with({pcoll}) # Run the task to trigger the marks self.assertIs(res, mock_async_res_instance) mock_submit.assert_called_once() self.assertIsNotNone(task_submitted) with patch.object( rm, '_wait_for_dependencies', return_value=True ), patch.object( rm, '_execute_pipeline_fragment' ) as _: task_submitted() self.assertTrue(pcoll in ie.current_env().computing_pcollections) def test_get_all_dependencies(self): p = beam.Pipeline(InteractiveRunner()) p1 = p | 'C1' >> beam.Create([1]) p2 = p | 'C2' >> beam.Create([2]) p3 = p1 | 'M1' >> beam.Map(lambda x: x) p4 = (p2, p3) | 'F1' >> beam.Flatten() p5 = p3 | 'M2' >> beam.Map(lambda x: x) ib.watch(locals()) ie.current_env().track_user_pipelines() rm = RecordingManager(p) rm.record_pipeline() # Analyze pipeline self.assertEqual(rm._get_all_dependencies({p1}), set()) self.assertEqual(rm._get_all_dependencies({p3}), {p1}) self.assertEqual(rm._get_all_dependencies({p4}), {p1, p2, p3}) self.assertEqual(rm._get_all_dependencies({p5}), {p1, p3}) self.assertEqual(rm._get_all_dependencies({p4, p5}), {p1, p2, p3}) @patch( 'apache_beam.runners.interactive.recording_manager.AsyncComputationResult' ) def test_wait_for_dependencies(self, mock_async_result_cls): p = beam.Pipeline(InteractiveRunner()) p1 = p | 'C1' >> beam.Create([1]) p2 = p1 | 'M1' >> beam.Map(lambda x: x) ib.watch(locals()) ie.current_env().track_user_pipelines() rm = RecordingManager(p) rm.record_pipeline() # Scenario 1: No dependencies computing self.assertTrue(rm._wait_for_dependencies({p2})) # Scenario 2: Dependency is computing mock_future = MagicMock(spec=Future) mock_async_res = MagicMock(spec=AsyncComputationResult) mock_async_res._future = mock_future mock_async_res._pcolls = {p1} rm._async_computations['dep_id'] = mock_async_res ie.current_env().mark_pcollection_computing({p1}) self.assertTrue(rm._wait_for_dependencies({p2})) mock_future.result.assert_called_once() ie.current_env().unmark_pcollection_computing({p1}) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.6494362 apache_beam-2.74.0/apache_beam/runners/interactive/sql/0000755000175100017510000000000015206113160022467 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/sql/__init__.py0000644000175100017510000000142015206113032024573 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/sql/beam_sql_magics.py0000644000175100017510000004302615206113032026152 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module of beam_sql cell magic that executes a Beam SQL. Only works within an IPython kernel. """ import argparse import importlib import keyword import logging import traceback from typing import Optional from typing import Union from IPython.core.magic import Magics from IPython.core.magic import line_cell_magic from IPython.core.magic import magics_class import apache_beam as beam from apache_beam.pvalue import PValue from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.background_caching_job import has_source_to_cache from apache_beam.runners.interactive.caching.cacheable import CacheKey from apache_beam.runners.interactive.caching.reify import reify_to_cache from apache_beam.runners.interactive.caching.reify import unreify_from_cache from apache_beam.runners.interactive.display.pcoll_visualization import visualize_computed_pcoll from apache_beam.runners.interactive.sql.sql_chain import SqlChain from apache_beam.runners.interactive.sql.sql_chain import SqlNode from apache_beam.runners.interactive.sql.utils import DataflowOptionsForm from apache_beam.runners.interactive.sql.utils import find_pcolls from apache_beam.runners.interactive.sql.utils import pformat_namedtuple from apache_beam.runners.interactive.sql.utils import register_coder_for_schema from apache_beam.runners.interactive.sql.utils import replace_single_pcoll_token from apache_beam.runners.interactive.utils import create_var_in_main from apache_beam.runners.interactive.utils import obfuscate from apache_beam.runners.interactive.utils import pcoll_by_name from apache_beam.runners.interactive.utils import progress_indicated from apache_beam.testing import test_stream from apache_beam.testing.test_stream_service import TestStreamServiceController from apache_beam.transforms.sql import SqlTransform from apache_beam.typehints.native_type_compatibility import match_is_named_tuple _LOGGER = logging.getLogger(__name__) _EXAMPLE_USAGE = """beam_sql magic to execute Beam SQL in notebooks --------------------------------------------------------- %%beam_sql [-o OUTPUT_NAME] [-v] [-r RUNNER] query --------------------------------------------------------- Or --------------------------------------------------------- %%%%beam_sql [-o OUTPUT_NAME] [-v] [-r RUNNER] query-line#1 query-line#2 ... query-line#N --------------------------------------------------------- """ _NOT_SUPPORTED_MSG = """The query was valid and successfully applied. But beam_sql failed to execute the query: %s Runner used by beam_sql was %s. Some Beam features might have not been supported by the Python SDK and runner combination. Please check the runner output for more details about the failed items. In the meantime, you may check: https://beam.apache.org/documentation/runners/capability-matrix/ to choose a runner other than the InteractiveRunner and explicitly apply SqlTransform to build Beam pipelines in a non-interactive manner. """ _SUPPORTED_RUNNERS = ['DirectRunner', 'DataflowRunner'] class BeamSqlParser: """A parser to parse beam_sql inputs.""" def __init__(self): self._parser = argparse.ArgumentParser(usage=_EXAMPLE_USAGE) self._parser.add_argument( '-o', '--output-name', dest='output_name', help=( 'The output variable name of the magic, usually a PCollection. ' 'Auto-generated if omitted.')) self._parser.add_argument( '-v', '--verbose', action='store_true', help='Display more details about the magic execution.') self._parser.add_argument( '-r', '--runner', dest='runner', help=( 'The runner to run the query. Supported runners are %s. If not ' 'provided, DirectRunner is used and results can be inspected ' 'locally.' % _SUPPORTED_RUNNERS)) self._parser.add_argument( 'query', type=str, nargs='*', help=( 'The Beam SQL query to execute. ' 'Syntax: https://beam.apache.org/documentation/dsls/sql/calcite/' 'query-syntax/. ' 'Please make sure that there is no conflict between your variable ' 'names and the SQL keywords, such as "SELECT", "FROM", "WHERE" and ' 'etc.')) def parse(self, args: list[str]) -> Optional[argparse.Namespace]: """Parses a list of string inputs. The parsed namespace contains these attributes: output_name: Optional[str], the output variable name. verbose: bool, whether to display more details of the magic execution. query: Optional[List[str]], the beam SQL query to execute. Returns: The parsed args or None if fail to parse. """ try: return self._parser.parse_args(args) except KeyboardInterrupt: raise except: # pylint: disable=bare-except # -h or --help results in SystemExit 0. Do not raise. return None def print_help(self) -> None: self._parser.print_help() def on_error(error_msg, *args): """Logs the error and the usage example.""" _LOGGER.error(error_msg, *args) BeamSqlParser().print_help() @magics_class class BeamSqlMagics(Magics): def __init__(self, shell): super().__init__(shell) # Eagerly initializes the environment. _ = ie.current_env() self._parser = BeamSqlParser() @line_cell_magic def beam_sql(self, line: str, cell: Optional[str] = None) -> Optional[PValue]: """The beam_sql line/cell magic that executes a Beam SQL. Args: line: the string on the same line after the beam_sql magic. cell: everything else in the same notebook cell as a string. If None, beam_sql is used as line magic. Otherwise, cell magic. Returns None if running into an error or waiting for user input (running on a selected runner remotely), otherwise a PValue as if a SqlTransform is applied. """ input_str = line if cell: input_str += ' ' + cell parsed = self._parser.parse(input_str.strip().split()) if not parsed: # Failed to parse inputs, let the parser handle the exit. return output_name = parsed.output_name verbose = parsed.verbose query = parsed.query runner = parsed.runner if output_name and not output_name.isidentifier() or keyword.iskeyword( output_name): on_error( 'The output_name "%s" is not a valid identifier. Please supply a ' 'valid identifier that is not a Python keyword.', line) return if not query: on_error('Please supply the SQL query to be executed.') return if runner and runner not in _SUPPORTED_RUNNERS: on_error( 'Runner "%s" is not supported. Supported runners are %s.', runner, _SUPPORTED_RUNNERS) return query = ' '.join(query) found = find_pcolls(query, pcoll_by_name(), verbose=verbose) schemas = set() main_session = importlib.import_module('__main__') for _, pcoll in found.items(): if not match_is_named_tuple(pcoll.element_type): on_error( 'PCollection %s of type %s is not a NamedTuple. See ' 'https://beam.apache.org/documentation/programming-guide/#schemas ' 'for more details.', pcoll, pcoll.element_type) return register_coder_for_schema(pcoll.element_type, verbose=verbose) # Only care about schemas defined by the user in the main module. if hasattr(main_session, pcoll.element_type.__name__): schemas.add(pcoll.element_type) if runner in ('DirectRunner', None): collect_data_for_local_run(query, found) output_name, output, chain = apply_sql(query, output_name, found) chain.current.schemas = schemas cache_output(output_name, output) return output output_name, current_node, chain = apply_sql( query, output_name, found, False) current_node.schemas = schemas # TODO(BEAM-10708): Move the options setup and result handling to a # separate module when more runners are supported. if runner == 'DataflowRunner': _ = chain.to_pipeline() _ = DataflowOptionsForm( output_name, pcoll_by_name()[output_name], verbose).display_for_input() return None else: raise ValueError('Unsupported runner %s.', runner) @progress_indicated def collect_data_for_local_run(query: str, found: dict[str, beam.PCollection]): from apache_beam.runners.interactive import interactive_beam as ib for name, pcoll in found.items(): try: _ = ib.collect(pcoll) except (KeyboardInterrupt, SystemExit): raise except: # pylint: disable=bare-except _LOGGER.error( 'Cannot collect data for PCollection %s. Please make sure the ' 'PCollections queried in the sql "%s" are all from a single ' 'pipeline using an InteractiveRunner. Make sure there is no ' 'ambiguity, for example, same named PCollections from multiple ' 'pipelines or notebook re-executions.', name, query) raise @progress_indicated def apply_sql( query: str, output_name: Optional[str], found: dict[str, beam.PCollection], run: bool = True) -> tuple[str, Union[PValue, SqlNode], SqlChain]: """Applies a SqlTransform with the given sql and queried PCollections. Args: query: The SQL query executed in the magic. output_name: (optional) The output variable name in __main__ module. found: The PCollections with variable names found to be used in the query. run: Whether to prepare the SQL pipeline for a local run or not. Returns: A tuple of values. First str value is the output variable name in __main__ module, auto-generated if not provided. Second value: if run, it's a PValue; otherwise, a SqlNode tracks the SQL without applying it or executing it. Third value: SqlChain is a chain of SqlNodes that have been applied. """ output_name = _generate_output_name(output_name, query, found) query, sql_source, chain = _build_query_components( query, found, output_name, run) if run: try: output = sql_source | SqlTransform(query) # Declare a variable with the output_name and output value in the # __main__ module so that the user can use the output smoothly. output_name, output = create_var_in_main(output_name, output) _LOGGER.info( "The output PCollection variable is %s with element_type %s", output_name, pformat_namedtuple(output.element_type)) return output_name, output, chain except (KeyboardInterrupt, SystemExit): raise except: # pylint: disable=bare-except on_error('Error when applying the Beam SQL: %s', traceback.format_exc()) raise else: return output_name, chain.current, chain def pcolls_from_streaming_cache( user_pipeline: beam.Pipeline, query_pipeline: beam.Pipeline, name_to_pcoll: dict[str, beam.PCollection]) -> dict[str, beam.PCollection]: """Reads PCollection cache through the TestStream. Args: user_pipeline: The beam.Pipeline object defined by the user in the notebook. query_pipeline: The beam.Pipeline object built by the magic to execute the SQL query. name_to_pcoll: PCollections with variable names used in the SQL query. Returns: A Dict[str, beam.PCollection], where each PCollection is tagged with their PCollection variable names, read from the cache. When the user_pipeline has unbounded sources, we force all cache reads to go through the TestStream even if they are bounded sources. """ def exception_handler(e): _LOGGER.error(str(e)) return True cache_manager = ie.current_env().get_cache_manager( user_pipeline, create_if_absent=True) test_stream_service = ie.current_env().get_test_stream_service_controller( user_pipeline) if not test_stream_service: test_stream_service = TestStreamServiceController( cache_manager, exception_handler=exception_handler) test_stream_service.start() ie.current_env().set_test_stream_service_controller( user_pipeline, test_stream_service) tag_to_name = {} for name, pcoll in name_to_pcoll.items(): key = CacheKey.from_pcoll(name, pcoll).to_str() tag_to_name[key] = name output_pcolls = query_pipeline | test_stream.TestStream( output_tags=set(tag_to_name.keys()), coder=cache_manager._default_pcoder, endpoint=test_stream_service.endpoint) sql_source = {} for tag, output in output_pcolls.items(): name = tag_to_name[tag] # Must mark the element_type to avoid introducing pickled Python coder # to the Java expansion service. output.element_type = name_to_pcoll[name].element_type sql_source[name] = output return sql_source def _generate_output_name( output_name: Optional[str], query: str, found: dict[str, beam.PCollection]) -> str: """Generates a unique output name if None is provided. Otherwise, returns the given output name directly. The generated output name is sql_output_{uuid} where uuid is an obfuscated value from the query and PCollections found to be used in the query. """ if not output_name: execution_id = obfuscate(query, found)[:12] output_name = 'sql_output_' + execution_id return output_name def _build_query_components( query: str, found: dict[str, beam.PCollection], output_name: str, run: bool = True ) -> tuple[str, Union[dict[str, beam.PCollection], beam.PCollection, beam.Pipeline], SqlChain]: """Builds necessary components needed to apply the SqlTransform. Args: query: The SQL query to be executed by the magic. found: The PCollections with variable names found to be used by the query. output_name: The output variable name in __main__ module. run: Whether to prepare components for a local run or not. Returns: The processed query to be executed by the magic; a source to apply the SqlTransform to: a dictionary of tagged PCollections, or a single PCollection, or the pipeline to execute the query; the chain of applied beam_sql magics this one belongs to. """ if found: user_pipeline = ie.current_env().user_pipeline( next(iter(found.values())).pipeline) sql_pipeline = beam.Pipeline(options=user_pipeline._options) ie.current_env().add_derived_pipeline(user_pipeline, sql_pipeline) sql_source = {} if run: if has_source_to_cache(user_pipeline): sql_source = pcolls_from_streaming_cache( user_pipeline, sql_pipeline, found) else: cache_manager = ie.current_env().get_cache_manager( user_pipeline, create_if_absent=True) for pcoll_name, pcoll in found.items(): cache_key = CacheKey.from_pcoll(pcoll_name, pcoll).to_str() sql_source[pcoll_name] = unreify_from_cache( pipeline=sql_pipeline, cache_key=cache_key, cache_manager=cache_manager, element_type=pcoll.element_type) else: sql_source = found if len(sql_source) == 1: query = replace_single_pcoll_token(query, next(iter(sql_source.keys()))) sql_source = next(iter(sql_source.values())) node = SqlNode( output_name=output_name, source=set(found.keys()), query=query) chain = ie.current_env().get_sql_chain( user_pipeline, set_user_pipeline=True).append(node) else: # does not query any existing PCollection sql_source = beam.Pipeline() ie.current_env().add_user_pipeline(sql_source) # The node should be the root node of the chain created below. node = SqlNode(output_name=output_name, source=sql_source, query=query) chain = ie.current_env().get_sql_chain(sql_source).append(node) return query, sql_source, chain @progress_indicated def cache_output(output_name: str, output: PValue) -> None: user_pipeline = ie.current_env().user_pipeline(output.pipeline) if user_pipeline: cache_manager = ie.current_env().get_cache_manager( user_pipeline, create_if_absent=True) else: _LOGGER.warning( 'Something is wrong with %s. Cannot introspect its data.', output) return key = CacheKey.from_pcoll(output_name, output).to_str() _ = reify_to_cache(pcoll=output, cache_key=key, cache_manager=cache_manager) try: output.pipeline.run().wait_until_finish() except (KeyboardInterrupt, SystemExit): raise except: # pylint: disable=bare-except _LOGGER.warning( _NOT_SUPPORTED_MSG, traceback.format_exc(), output.pipeline.runner) return ie.current_env().mark_pcollection_computed([output]) visualize_computed_pcoll( output_name, output, max_n=float('inf'), max_duration_secs=float('inf')) def load_ipython_extension(ipython): """Marks this module as an IPython extension. To load this magic in an IPython environment, execute: %load_ext apache_beam.runners.interactive.sql.beam_sql_magics. """ ipython.register_magics(BeamSqlMagics) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/sql/beam_sql_magics_test.py0000644000175100017510000001506615206113032027214 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for beam_sql_magics module.""" # pytype: skip-file import unittest from unittest.mock import patch import pytest import apache_beam as beam from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.cache_manager import FileBasedCacheManager from apache_beam.runners.interactive.caching.cacheable import CacheKey try: from apache_beam.runners.interactive.sql.beam_sql_magics import BeamSqlParser from apache_beam.runners.interactive.sql.beam_sql_magics import _build_query_components from apache_beam.runners.interactive.sql.beam_sql_magics import _generate_output_name from apache_beam.runners.interactive.sql.beam_sql_magics import cache_output except (ImportError, NameError): pass # The test is to be skipped because [interactive] dep not installed. @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') @pytest.mark.skipif( not ie.current_env().is_interactive_ready, reason='[interactive] dependency is not installed.') class BeamSqlMagicsTest(unittest.TestCase): def test_generate_output_name_when_not_provided(self): output_name = None self.assertTrue( _generate_output_name(output_name, '', {}).startswith('sql_output_')) def test_use_given_output_name_when_provided(self): output_name = 'output' self.assertEqual(_generate_output_name(output_name, '', {}), output_name) def test_build_query_components_when_no_pcoll_queried(self): query = """SELECT CAST(1 AS INT) AS `id`, CAST('foo' AS VARCHAR) AS `str`, CAST(3.14 AS DOUBLE) AS `flt`""" processed_query, sql_source, chain = _build_query_components( query, {}, 'output') self.assertEqual(processed_query, query) self.assertIsInstance(sql_source, beam.Pipeline) self.assertIsInstance(chain.current.source, beam.Pipeline) self.assertEqual('output', chain.current.output_name) self.assertEqual(query, chain.current.query) def test_build_query_components_when_single_pcoll_queried(self): p = beam.Pipeline() target = p | beam.Create([1, 2, 3]) ib.watch(locals()) query = 'SELECT * FROM target where a=1' found = {'target': target} with patch('apache_beam.runners.interactive.sql.beam_sql_magics.' 'unreify_from_cache', lambda pipeline, cache_key, cache_manager, element_type: target): processed_query, sql_source, chain = _build_query_components( query, found, 'output') expected_query = 'SELECT * FROM PCOLLECTION where a=1' self.assertEqual(expected_query, processed_query) self.assertIsInstance(sql_source, beam.PCollection) self.assertIn('target', chain.current.source) self.assertEqual(expected_query, chain.current.query) self.assertEqual('output', chain.current.output_name) def test_build_query_components_when_multiple_pcolls_queried(self): p = beam.Pipeline() pcoll_1 = p | 'Create 1' >> beam.Create([1, 2, 3]) pcoll_2 = p | 'Create 2' >> beam.Create([4, 5, 6]) ib.watch(locals()) query = 'SELECT * FROM pcoll_1 JOIN pcoll_2 USING (a)' found = {'pcoll_1': pcoll_1, 'pcoll_2': pcoll_2} with patch( 'apache_beam.runners.interactive.sql.beam_sql_magics.' 'unreify_from_cache', lambda pipeline, cache_key, cache_manager, element_type: pcoll_1): processed_query, sql_source, chain = _build_query_components( query, found, 'output') self.assertEqual(processed_query, query) self.assertIsInstance(sql_source, dict) self.assertIn('pcoll_1', sql_source) self.assertIn('pcoll_2', sql_source) self.assertIn('pcoll_1', chain.current.source) self.assertIn('pcoll_2', chain.current.source) self.assertEqual(query, chain.current.query) self.assertEqual('output', chain.current.output_name) def test_build_query_components_when_unbounded_pcolls_queried(self): p = beam.Pipeline() pcoll = p | beam.io.ReadFromPubSub( subscription='projects/fake-project/subscriptions/fake_sub') ib.watch(locals()) query = 'SELECT * FROM pcoll' found = {'pcoll': pcoll} with patch('apache_beam.runners.interactive.sql.beam_sql_magics.' 'pcolls_from_streaming_cache', lambda a, b, c: found): _, sql_source, chain = _build_query_components(query, found, 'output') self.assertIs(sql_source, pcoll) self.assertIn('pcoll', chain.current.source) self.assertEqual('SELECT * FROM PCOLLECTION', chain.current.query) self.assertEqual('output', chain.current.output_name) def test_cache_output(self): p_cache_output = beam.Pipeline() pcoll_co = p_cache_output | 'Create Source' >> beam.Create([1, 2, 3]) cache_manager = FileBasedCacheManager() ie.current_env().set_cache_manager(cache_manager, p_cache_output) ib.watch(locals()) with patch('apache_beam.runners.interactive.display.pcoll_visualization.' 'visualize_computed_pcoll', lambda a, b: None): cache_output('pcoll_co', pcoll_co) self.assertIn(pcoll_co, ie.current_env().computed_pcollections) self.assertTrue( cache_manager.exists( 'full', CacheKey.from_pcoll('pcoll_co', pcoll_co).to_str())) def test_parser_with_all_inputs(self): parsed = BeamSqlParser().parse( '-o output_name -v SELECT * FROM abc'.split()) self.assertTrue(parsed.verbose) self.assertEqual('output_name', parsed.output_name) self.assertEqual('SELECT * FROM abc', ' '.join(parsed.query)) def test_parser_with_no_input(self): parsed = BeamSqlParser().parse([]) self.assertFalse(parsed.verbose) self.assertIsNone(parsed.output_name) self.assertFalse(parsed.query) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/sql/sql_chain.py0000644000175100017510000002121615206113032025002 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module for tracking a chain of beam_sql magics applied. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file import importlib import logging from dataclasses import dataclass from typing import Any from typing import Optional from typing import Union import apache_beam as beam from apache_beam.internal import pickler from apache_beam.runners.interactive.sql.utils import register_coder_for_schema from apache_beam.runners.interactive.utils import create_var_in_main from apache_beam.runners.interactive.utils import pcoll_by_name from apache_beam.runners.interactive.utils import progress_indicated from apache_beam.transforms.sql import SqlTransform from apache_beam.utils.interactive_utils import is_in_ipython _LOGGER = logging.getLogger(__name__) @dataclass class SqlNode: """Each SqlNode represents a beam_sql magic applied. Attributes: output_name: the watched unique name of the beam_sql output. Can be used as an identifier. source: the inputs consumed by this node. Can be a pipeline or a set of PCollections represented by their variable names watched. When it's a pipeline, the node computes from raw values in the query, so the output can be consumed by any SqlNode in any SqlChain. query: the SQL query applied by this node. schemas: the schemas (NamedTuple classes) used by this node. evaluated: the pipelines this node has been evaluated for. next: the next SqlNode applied chronologically. execution_count: the execution count if in an IPython env. """ output_name: str source: Union[beam.Pipeline, set[str]] query: str schemas: set[Any] = None evaluated: set[beam.Pipeline] = None next: Optional['SqlNode'] = None execution_count: int = 0 def __post_init__(self): if not self.schemas: self.schemas = set() if not self.evaluated: self.evaluated = set() if is_in_ipython(): from IPython import get_ipython self.execution_count = get_ipython().execution_count def __hash__(self): return hash( (self.output_name, self.source, self.query, self.execution_count)) def to_pipeline(self, pipeline: Optional[beam.Pipeline]) -> beam.Pipeline: """Converts the chain into an executable pipeline.""" if pipeline not in self.evaluated: # The whole chain should form a single pipeline. source = self.source if isinstance(self.source, beam.Pipeline): if pipeline: # use the known pipeline source = pipeline else: # use the source pipeline pipeline = self.source else: name_to_pcoll = pcoll_by_name() if len(self.source) == 1: source = name_to_pcoll.get(next(iter(self.source))) else: source = {s: name_to_pcoll.get(s) for s in self.source} if isinstance(source, beam.Pipeline): output = source | 'beam_sql_{}_{}'.format( self.output_name, self.execution_count) >> SqlTransform(self.query) else: output = source | 'schema_loaded_beam_sql_{}_{}'.format( self.output_name, self.execution_count ) >> SchemaLoadedSqlTransform( self.output_name, self.query, self.schemas, self.execution_count) _ = create_var_in_main(self.output_name, output) self.evaluated.add(pipeline) if self.next: return self.next.to_pipeline(pipeline) else: return pipeline class SchemaLoadedSqlTransform(beam.PTransform): """PTransform that loads schema before executing SQL. When submitting a pipeline to remote runner for execution, schemas defined in the main module are not available without save_main_session. However, save_main_session might fail when there is anything unpicklable. This DoFn makes sure only the schemas needed are pickled locally and restored later on workers. """ def __init__(self, output_name, query, schemas, execution_count): self.output_name = output_name self.query = query self.schemas = schemas self.execution_count = execution_count # TODO(BEAM-8123): clean up this attribute or the whole wrapper PTransform. # Dill does not preserve everything. On the other hand, save_main_session # is not stable. Until cloudpickle replaces dill in Beam, we work around # it by explicitly pickling annotations and load schemas in remote main # sessions. self.schema_annotations = [s.__annotations__ for s in self.schemas] class _SqlTransformDoFn(beam.DoFn): """The DoFn yields all its input without any transform but a setup to configure the main session.""" def __init__(self, schemas, annotations): self.pickled_schemas = [pickler.dumps(s) for s in schemas] self.pickled_annotations = [pickler.dumps(a) for a in annotations] def setup(self): main_session = importlib.import_module('__main__') for pickled_schema, pickled_annotation in zip( self.pickled_schemas, self.pickled_annotations): schema = pickler.loads(pickled_schema) schema.__annotations__ = pickler.loads(pickled_annotation) if not hasattr(main_session, schema.__name__) or not hasattr( getattr(main_session, schema.__name__), '__annotations__'): # Restore the schema in the main session on the [remote] worker. setattr(main_session, schema.__name__, schema) register_coder_for_schema(schema) def process(self, e): yield e def expand(self, source): """Applies the SQL transform. If a PCollection uses a schema defined in the main session, use the additional DoFn to restore it on the worker.""" if isinstance(source, dict): schema_loaded = { tag: pcoll | 'load_schemas_{}_tag_{}_{}'.format( self.output_name, tag, self.execution_count) >> beam.ParDo( self._SqlTransformDoFn(self.schemas, self.schema_annotations)) if pcoll.element_type in self.schemas else pcoll for tag, pcoll in source.items() } elif isinstance(source, beam.pvalue.PCollection): schema_loaded = source | 'load_schemas_{}_{}'.format( self.output_name, self.execution_count) >> beam.ParDo( self._SqlTransformDoFn(self.schemas, self.schema_annotations) ) if source.element_type in self.schemas else source else: raise ValueError( '{} should be either a single PCollection or a dict of named ' 'PCollections.'.format(source)) return schema_loaded | 'beam_sql_{}_{}'.format( self.output_name, self.execution_count) >> SqlTransform(self.query) @dataclass class SqlChain: """A chain of SqlNodes. Attributes: nodes: all nodes by their output_names. root: the first SqlNode applied chronologically. current: the last node applied. user_pipeline: the user defined pipeline this chain originates from. If None, the whole chain just computes from raw values in queries. Otherwise, at least some of the nodes in chain has queried against PCollections. """ nodes: dict[str, SqlNode] = None root: Optional[SqlNode] = None current: Optional[SqlNode] = None user_pipeline: Optional[beam.Pipeline] = None def __post_init__(self): if not self.nodes: self.nodes = {} @progress_indicated def to_pipeline(self) -> beam.Pipeline: """Converts the chain into a beam pipeline.""" pipeline_to_execute = self.root.to_pipeline(self.user_pipeline) # The pipeline definitely contains external transform: SqlTransform. pipeline_to_execute.contains_external_transforms = True return pipeline_to_execute def append(self, node: SqlNode) -> 'SqlChain': """Appends a node to the chain.""" if self.current: self.current.next = node else: self.root = node self.current = node self.nodes[node.output_name] = node return self def get(self, output_name: str) -> Optional[SqlNode]: """Gets a node from the chain based on the given output_name.""" return self.nodes.get(output_name, None) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/sql/sql_chain_test.py0000644000175100017510000001006215206113032026036 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for sql_chain module.""" # pytype: skip-file import unittest from unittest.mock import patch import pytest import apache_beam as beam from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.sql.sql_chain import SqlChain from apache_beam.runners.interactive.sql.sql_chain import SqlNode from apache_beam.runners.interactive.testing.mock_ipython import mock_get_ipython class SqlChainTest(unittest.TestCase): def test_init(self): chain = SqlChain() self.assertEqual({}, chain.nodes) self.assertIsNone(chain.root) self.assertIsNone(chain.current) self.assertIsNone(chain.user_pipeline) def test_append_first_node(self): node = SqlNode(output_name='first', source='a', query='q1') chain = SqlChain().append(node) self.assertIs(node, chain.get(node.output_name)) self.assertIs(node, chain.root) self.assertIs(node, chain.current) def test_append_non_root_node(self): chain = SqlChain().append( SqlNode(output_name='root', source='root', query='q1')) self.assertIsNone(chain.root.next) node = SqlNode(output_name='next_node', source='root', query='q2') chain.append(node) self.assertIs(node, chain.root.next) self.assertIs(node, chain.get(node.output_name)) @patch( 'apache_beam.runners.interactive.sql.sql_chain.SchemaLoadedSqlTransform.' '__rrshift__') def test_to_pipeline_only_evaluate_once_per_pipeline_and_node( self, mocked_sql_transform): p = beam.Pipeline() ie.current_env().watch({'p': p}) pcoll_1 = p | 'create pcoll_1' >> beam.Create([1, 2, 3]) pcoll_2 = p | 'create pcoll_2' >> beam.Create([4, 5, 6]) ie.current_env().watch({'pcoll_1': pcoll_1, 'pcoll_2': pcoll_2}) node = SqlNode( output_name='root', source={'pcoll_1', 'pcoll_2'}, query='q1') chain = SqlChain(user_pipeline=p).append(node) _ = chain.to_pipeline() mocked_sql_transform.assert_called_once() _ = chain.to_pipeline() mocked_sql_transform.assert_called_once() @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') @pytest.mark.skipif( not ie.current_env().is_interactive_ready, reason='[interactive] dependency is not installed.') @patch( 'apache_beam.runners.interactive.sql.sql_chain.SchemaLoadedSqlTransform.' '__rrshift__') def test_nodes_with_same_outputs(self, mocked_sql_transform): p = beam.Pipeline() ie.current_env().watch({'p_nodes_with_same_output': p}) pcoll = p | 'create pcoll' >> beam.Create([1, 2, 3]) ie.current_env().watch({'pcoll': pcoll}) chain = SqlChain(user_pipeline=p) output_name = 'output' with patch('IPython.get_ipython', new_callable=mock_get_ipython) as cell: with cell: node_cell_1 = SqlNode(output_name, source='pcoll', query='q1') chain.append(node_cell_1) _ = chain.to_pipeline() mocked_sql_transform.assert_called_with( 'schema_loaded_beam_sql_output_1') with cell: node_cell_2 = SqlNode(output_name, source='pcoll', query='q2') chain.append(node_cell_2) _ = chain.to_pipeline() mocked_sql_transform.assert_called_with( 'schema_loaded_beam_sql_output_2') if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/sql/utils.py0000644000175100017510000003605415206113032024207 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module of utilities for SQL magics. For internal use only; no backward-compatibility guarantees. """ # pytype: skip-file import logging import os import tempfile from dataclasses import dataclass from typing import Any from typing import Callable from typing import NamedTuple from typing import Optional from typing import Union import apache_beam as beam from apache_beam.io import WriteToText from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import WorkerOptions from apache_beam.runners.interactive.utils import create_var_in_main from apache_beam.runners.interactive.utils import progress_indicated from apache_beam.runners.runner import create_runner from apache_beam.typehints.native_type_compatibility import match_is_named_tuple from apache_beam.utils.interactive_utils import is_in_ipython _LOGGER = logging.getLogger(__name__) def register_coder_for_schema( schema: NamedTuple, verbose: bool = False) -> None: """Registers a RowCoder for the given schema if hasn't. Notifies the user of what code has been implicitly executed. """ assert match_is_named_tuple(schema), ( 'Schema %s is not a typing.NamedTuple.' % schema) coder = beam.coders.registry.get_coder(schema) if not isinstance(coder, beam.coders.RowCoder): if verbose: _LOGGER.warning( 'Schema %s has not been registered to use a RowCoder. ' 'Automatically registering it by running: ' 'beam.coders.registry.register_coder(%s, ' 'beam.coders.RowCoder)', schema.__name__, schema.__name__) beam.coders.registry.register_coder(schema, beam.coders.RowCoder) def find_pcolls( sql: str, pcolls: dict[str, beam.PCollection], verbose: bool = False) -> dict[str, beam.PCollection]: """Finds all PCollections used in the given sql query. It does a simple word by word match and calls ib.collect for each PCollection found. """ found = {} for word in sql.split(): if word in pcolls: found[word] = pcolls[word] if found: if verbose: _LOGGER.info('Found PCollections used in the magic: %s.', found) _LOGGER.info('Collecting data...') return found def replace_single_pcoll_token(sql: str, pcoll_name: str) -> str: """Replaces the pcoll_name used in the sql with 'PCOLLECTION'. For sql query using only a single PCollection, the PCollection needs to be referred to as 'PCOLLECTION' instead of its variable/tag name. """ words = sql.split() token_locations = [] i = 0 for word in words: if word.lower() == 'from': token_locations.append(i + 1) i += 2 continue i += 1 for token_location in token_locations: if token_location < len(words) and words[token_location] == pcoll_name: words[token_location] = 'PCOLLECTION' return ' '.join(words) def pformat_namedtuple(schema: NamedTuple) -> str: return '{}({})'.format( schema.__name__, ', '.join([ '{}: {}'.format(k, repr(v)) for k, v in schema.__annotations__.items() ])) def pformat_dict(raw_input: dict[str, Any]) -> str: return '{{\n{}\n}}'.format( ',\n'.join(['{}: {}'.format(k, v) for k, v in raw_input.items()])) @dataclass class OptionsEntry: """An entry of PipelineOptions that can be visualized through ipywidgets to take inputs in IPython notebooks interactively. Attributes: label: The value of the Label widget. help: The help message of the entry, usually the same to the help in PipelineOptions. cls: The PipelineOptions class/subclass the options belong to. arg_builder: Builds the argument/option. If it's a str, this entry assigns the input ipywidget's value directly to the argument. If it's a Dict, use the corresponding Callable to assign the input value to each argument. If Callable is None, fallback to assign the input value directly. This allows building multiple similar PipelineOptions arguments from a single input, such as staging_location and temp_location in GoogleCloudOptions. default: The default value of the entry, None if absent. """ label: str help: str cls: type[PipelineOptions] arg_builder: Union[str, dict[str, Optional[Callable]]] default: Optional[str] = None def __post_init__(self): # The attribute holds an ipywidget, currently only supports Text. # The str value can be accessed by self.input.value. self.input = None class OptionsForm: """A form visualized to take inputs from users in IPython Notebooks and generate PipelineOptions to run pipelines. """ def __init__(self): # The current Python SDK incorrectly parses unparsable pipeline options # Here we ignore all flags for the interactive beam_sql magic # since the beam_sql magic does not use flags self.options = PipelineOptions(flags={}) self.entries = [] def add(self, entry: OptionsEntry) -> 'OptionsForm': """Adds an OptionsEntry to the form. """ self.entries.append(entry) return self def to_options(self) -> PipelineOptions: """Builds the PipelineOptions based on user inputs. Can only be invoked after display_for_input. """ for entry in self.entries: assert entry.input, ( 'to_options invoked before display_for_input. ' 'Wrong usage.') view = self.options.view_as(entry.cls) if isinstance(entry.arg_builder, str): setattr(view, entry.arg_builder, entry.input.value) else: for arg, builder in entry.arg_builder.items(): if builder: setattr(view, arg, builder(entry.input.value)) else: setattr(view, arg, entry.input.value) self.additional_options() return self.options def additional_options(self): """Alters the self.options with additional config.""" pass def display_for_input(self) -> 'OptionsForm': """Displays the widgets to take user inputs.""" from IPython.display import display from ipywidgets import GridBox from ipywidgets import Label from ipywidgets import Layout from ipywidgets import Text widgets = [] for entry in self.entries: text_label = Label(value=entry.label) text_input = entry.input if entry.input else Text( value=entry.default if entry.default else '') text_help = Label(value=entry.help) entry.input = text_input widgets.append(text_label) widgets.append(text_input) widgets.append(text_help) grid = GridBox(widgets, layout=Layout(grid_template_columns='1fr 2fr 6fr')) display(grid) self.display_actions() return self def display_actions(self): """Displays actionable widgets to utilize the options, run pipelines and etc.""" pass class DataflowOptionsForm(OptionsForm): """A form to take inputs from users in IPython Notebooks to build PipelineOptions to run pipelines on Dataflow. Only contains minimum fields needed. """ @staticmethod def _build_default_project() -> str: """Builds a default project id.""" try: # pylint: disable=c-extension-no-member import google.auth return google.auth.default()[1] except (KeyboardInterrupt, SystemExit): raise except Exception as e: _LOGGER.warning('There is some issue with your gcloud auth: %s', e) return 'your-project-id' @staticmethod def _build_req_file_from_pkgs(pkgs) -> Optional[str]: """Builds a requirements file that contains all additional PYPI packages needed.""" if pkgs: deps = pkgs.split(',') req_file = os.path.join( tempfile.mkdtemp(prefix='beam-sql-dataflow-'), 'req.txt') with open(req_file, 'a') as f: for dep in deps: f.write(dep.strip() + '\n') return req_file return None def __init__( self, output_name: str, output_pcoll: beam.PCollection, verbose: bool = False): """Inits the OptionsForm for setting up Dataflow jobs.""" super().__init__() self.p = output_pcoll.pipeline self.output_name = output_name self.output_pcoll = output_pcoll self.verbose = verbose self.notice_shown = False self.add( OptionsEntry( label='Project Id', help='Name of the Cloud project owning the Dataflow job.', cls=GoogleCloudOptions, arg_builder='project', default=DataflowOptionsForm._build_default_project()) ).add( OptionsEntry( label='Region', help='The Google Compute Engine region for creating Dataflow job.', cls=GoogleCloudOptions, arg_builder='region', default='us-central1') ).add( OptionsEntry( label='GCS Bucket', help=( 'GCS path to stage code packages needed by workers and save ' 'temporary workflow jobs.'), cls=GoogleCloudOptions, arg_builder={ 'staging_location': lambda x: x + '/staging', 'temp_location': lambda x: x + '/temp' }, default='gs://YOUR_GCS_BUCKET_HERE') ).add( OptionsEntry( label='Additional Packages', help=( 'PYPI packages installed, comma-separated. If None, leave ' 'this field empty.'), cls=SetupOptions, arg_builder={ 'requirements_file': lambda x: DataflowOptionsForm. _build_req_file_from_pkgs(x) }, default='')) def additional_options(self): # Use the latest Java SDK by default. sdk_overrides = self.options.view_as( WorkerOptions).sdk_harness_container_image_overrides override = '.*java.*,apache/beam_java11_sdk:latest' if sdk_overrides and override not in sdk_overrides: sdk_overrides.append(override) else: self.options.view_as( WorkerOptions).sdk_harness_container_image_overrides = [override] def display_actions(self): from IPython.display import HTML from IPython.display import display from ipywidgets import Button from ipywidgets import GridBox from ipywidgets import Layout from ipywidgets import Output options_output_area = Output() run_output_area = Output() run_btn = Button( description='Run on Dataflow', button_style='success', tooltip=( 'Submit to Dataflow for execution with the configured options. The ' 'output PCollection\'s data will be written to the GCS bucket you ' 'configure.')) show_options_btn = Button( description='Show Options', button_style='info', tooltip='Show current pipeline options configured.') def _run_on_dataflow(btn): with run_output_area: run_output_area.clear_output() @progress_indicated def _inner(): options = self.to_options() # Caches the output_pcoll to a GCS bucket. try: execution_count = 0 if is_in_ipython(): from IPython import get_ipython execution_count = get_ipython().execution_count output_location = '{}/{}'.format( options.view_as(GoogleCloudOptions).staging_location, self.output_name) _ = self.output_pcoll | 'WriteOuput{}_{}ToGCS'.format( self.output_name, execution_count) >> WriteToText(output_location) _LOGGER.info( 'Data of output PCollection %s will be written to %s', self.output_name, output_location) except (KeyboardInterrupt, SystemExit): raise except: # pylint: disable=bare-except # The transform has been added before, noop. pass if self.verbose: _LOGGER.info( 'Running the pipeline on Dataflow with pipeline options %s.', pformat_dict(options.display_data())) result = create_runner('DataflowRunner').run_pipeline(self.p, options) cloud_options = options.view_as(GoogleCloudOptions) url = ( 'https://console.cloud.google.com/dataflow/jobs/%s/%s?project=%s' % (cloud_options.region, result.job_id(), cloud_options.project)) display( HTML( 'Click here for the details ' 'of your Dataflow job.' % url)) result_name = 'result_{}'.format(self.output_name) create_var_in_main(result_name, result) if self.verbose: _LOGGER.info( 'The pipeline result of the run can be accessed from variable ' '%s. The current status is %s.', result_name, result) try: btn.disabled = True _inner() finally: btn.disabled = False run_btn.on_click(_run_on_dataflow) def _show_options(btn): with options_output_area: options_output_area.clear_output() options = self.to_options() options_name = 'options_{}'.format(self.output_name) create_var_in_main(options_name, options) _LOGGER.info( 'The pipeline options configured is: %s.', pformat_dict(options.display_data())) show_options_btn.on_click(_show_options) grid = GridBox([run_btn, show_options_btn], layout=Layout(grid_template_columns='repeat(2, 200px)')) display(grid) # Implicitly initializes the options variable before 1st time showing # options. options_name_inited, _ = create_var_in_main('options_{}'.format( self.output_name), self.to_options()) if not self.notice_shown: _LOGGER.info( 'The pipeline options can be configured through variable %s. You ' 'may also add additional options or sink transforms such as write ' 'to BigQuery in other notebook cells. Come back to click "Run on ' 'Dataflow" button once you complete additional configurations. ' 'Optionally, you can chain more beam_sql magics with DataflowRunner ' 'and click "Run on Dataflow" in their outputs.', options_name_inited) self.notice_shown = True display(options_output_area) display(run_output_area) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/sql/utils_test.py0000644000175100017510000001122315206113032025235 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for utils module.""" # pytype: skip-file import unittest from typing import NamedTuple from typing import Optional from typing import Union from unittest.mock import patch import pytest import apache_beam as beam from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.sql.utils import DataflowOptionsForm from apache_beam.runners.interactive.sql.utils import find_pcolls from apache_beam.runners.interactive.sql.utils import pformat_dict from apache_beam.runners.interactive.sql.utils import pformat_namedtuple from apache_beam.runners.interactive.sql.utils import register_coder_for_schema from apache_beam.runners.interactive.sql.utils import replace_single_pcoll_token class ANamedTuple(NamedTuple): a: int b: str class OptionalUnionType(NamedTuple): unnamed: Optional[Union[int, str]] class UtilsTest(unittest.TestCase): def test_register_coder_for_schema(self): self.assertNotIsInstance( beam.coders.registry.get_coder(ANamedTuple), beam.coders.RowCoder) register_coder_for_schema(ANamedTuple) self.assertIsInstance( beam.coders.registry.get_coder(ANamedTuple), beam.coders.RowCoder) def test_find_pcolls(self): with patch('apache_beam.runners.interactive.interactive_beam.collect', lambda _: None): found = find_pcolls( """SELECT * FROM pcoll_1 JOIN pcoll_2 USING (common_column)""", { 'pcoll_1': None, 'pcoll_2': None }) self.assertIn('pcoll_1', found) self.assertIn('pcoll_2', found) def test_replace_single_pcoll_token(self): sql = 'SELECT * FROM abc WHERE a=1 AND b=2' replaced_sql = replace_single_pcoll_token(sql, 'wow') self.assertEqual(replaced_sql, sql) replaced_sql = replace_single_pcoll_token(sql, 'abc') self.assertEqual( replaced_sql, 'SELECT * FROM PCOLLECTION WHERE a=1 AND b=2') def test_pformat_namedtuple(self): actual = pformat_namedtuple(ANamedTuple) self.assertEqual("ANamedTuple(a: , b: )", actual) def test_pformat_namedtuple_with_unnamed_fields(self): actual = pformat_namedtuple(OptionalUnionType) # Parameters of an Union type can be in any order. possible_expected = ( 'OptionalUnionType(unnamed: typing.Union[int, str, NoneType])', 'OptionalUnionType(unnamed: typing.Union[str, int, NoneType])', 'OptionalUnionType(unnamed: int | str | None)') self.assertIn(actual, possible_expected) def test_pformat_dict(self): actual = pformat_dict({'a': 1, 'b': '2'}) self.assertEqual('{\na: 1,\nb: 2\n}', actual) @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') @pytest.mark.skipif( not ie.current_env().is_interactive_ready, reason='[interactive] dependency is not installed.') class OptionsFormTest(unittest.TestCase): def test_dataflow_options_form(self): p = beam.Pipeline() pcoll = p | beam.Create([1, 2, 3]) with patch('google.auth') as ga: ga.default = lambda: ['', 'default_project_id'] df_form = DataflowOptionsForm('pcoll', pcoll) df_form.display_for_input() df_form.entries[2].input.value = 'gs://test-bucket' df_form.entries[3].input.value = 'a-pkg' options = df_form.to_options() cloud_options = options.view_as(GoogleCloudOptions) self.assertEqual(cloud_options.project, 'default_project_id') self.assertEqual(cloud_options.region, 'us-central1') self.assertEqual( cloud_options.staging_location, 'gs://test-bucket/staging') self.assertEqual(cloud_options.temp_location, 'gs://test-bucket/temp') self.assertIsNotNone(options.view_as(SetupOptions).requirements_file) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.6514363 apache_beam-2.74.0/apache_beam/runners/interactive/testing/0000755000175100017510000000000015206113160023345 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/testing/__init__.py0000644000175100017510000000142015206113032025451 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.6524365 apache_beam-2.74.0/apache_beam/runners/interactive/testing/integration/0000755000175100017510000000000015206113160025670 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/testing/integration/__init__.py0000644000175100017510000000142015206113032027774 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/testing/integration/notebook_executor.py0000644000175100017510000001306615206113032032004 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module to execute jupyter notebooks and gather the output into renderable HTML files.""" # pytype: skip-file import html import os import shutil import subprocess from html.parser import HTMLParser from apache_beam.runners.interactive.utils import obfuscate try: import nbformat from jupyter_client.kernelspec import KernelSpecManager from nbconvert.preprocessors import ExecutePreprocessor _interactive_integration_ready = True except ImportError: _interactive_integration_ready = False class NotebookExecutor(object): """Executor that reads notebooks, executes it and gathers outputs into static HTML pages that can be served.""" def __init__(self, path: str) -> None: assert _interactive_integration_ready, ( '[interactive_test] dependency is not installed.') assert os.path.exists(path), '{} does not exist.'.format(path) self._paths = [] if os.path.isdir(path): for root, _, files in os.walk(path): for filename in files: if filename.endswith('.ipynb'): self._paths.append(os.path.join(root, filename)) elif path.endswith('.ipynb'): self._paths.append(path) assert len( self._paths) > 0, ('No notebooks to be executed under{}'.format(path)) self._dir = os.path.dirname(self._paths[0]) self._output_html_dir = os.path.join(self._dir, 'output') self.cleanup() self._output_html_paths = {} self._notebook_path_to_execution_id = {} kernel_specs = KernelSpecManager().get_all_specs() if 'test' not in kernel_specs: # Install a test ipython kernel in current runtime environment. If this # errors out, it means the test env is broken and should fail the test. process = subprocess.run( ['python', '-m', 'ipykernel', 'install', '--user', '--name', 'test'], check=True) process.check_returncode() def cleanup(self): """Cleans up the output folder.""" _cleanup(self._output_html_dir) def execute(self): """Executes all notebooks found in the scoped path and gathers their outputs into HTML pages stored in the output folder.""" for path in self._paths: with open(path, 'r') as nb_f: nb = nbformat.read(nb_f, as_version=4) ep = ExecutePreprocessor( timeout=-1, allow_errors=True, kernel_name='test') ep.preprocess(nb, {'metadata': {'path': os.path.dirname(path)}}) execution_id = obfuscate(path) output_html_path = os.path.join( self._output_html_dir, execution_id + '.html') with open(output_html_path, 'a+') as sink: sink.write('\n') sink.write('\n') sink.write('\n') sink.write('\n') for cell in nb['cells']: if cell['cell_type'] == 'code': for output in cell['outputs']: _extract_html(output, sink) sink.write('\n') sink.write('\n') self._output_html_paths[execution_id] = output_html_path self._notebook_path_to_execution_id[path] = execution_id @property def output_html_paths(self): """Mapping from execution ids to output html page paths. An execution/test id is an obfuscated value from the executed notebook path. It identifies the input notebook, the output html, the screenshot of the output html, and the golden screenshot for comparison. """ return self._output_html_paths @property def output_html_dir(self): """The directory's path to all the output html pages generated.""" return self._output_html_dir @property def notebook_path_to_execution_id(self): """Mapping from input notebook paths to their obfuscated execution ids.""" return self._notebook_path_to_execution_id def _cleanup(output_dir): """Cleans up the given output_dir.""" if os.path.exists(output_dir): shutil.rmtree(output_dir) os.makedirs(output_dir) def _extract_html(output, sink): """Extracts html elements from the output of an executed notebook node and writes them into a file sink.""" if output['output_type'] == 'display_data': data = output['data'] if 'application/javascript' in data: sink.write('\n') if 'text/html' in data: parser = IFrameParser() parser.feed(data['text/html']) if parser.srcdocs: sink.write(parser.srcdocs) else: sink.write(data['text/html']) class IFrameParser(HTMLParser): """A parser to extract iframe content from given HTML.""" def __init__(self): self._srcdocs = [] super().__init__() def handle_starttag(self, tag, attrs): if tag == 'iframe': for attr in attrs: if 'srcdoc' in attr: self._srcdocs.append(html.unescape(attr[1])) @property def srcdocs(self): return '\n'.join(self._srcdocs) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/testing/integration/screen_diff.py0000644000175100017510000002202215206113032030505 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module to conduct screen diff based notebook integration tests.""" # pytype: skip-file import os import platform import threading import unittest from http.server import HTTPServer from http.server import SimpleHTTPRequestHandler import pytest from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.testing.integration import notebook_executor try: import chromedriver_binary # pylint: disable=unused-import from needle.cases import NeedleTestCase from needle.driver import NeedleChrome from selenium.webdriver.chrome.options import Options from selenium.webdriver.common.by import By from selenium.webdriver.support import expected_conditions from selenium.webdriver.support.ui import WebDriverWait _interactive_integration_ready = ( notebook_executor._interactive_integration_ready) except ImportError: _interactive_integration_ready = False # Web elements will be rendered differently on different platforms. List all # supported platforms with goldens here. _SUPPORTED_PLATFORMS = ['Darwin', 'Linux'] class ScreenDiffIntegrationTestEnvironment(object): """A test environment to conduct screen diff integration tests for notebooks. """ def __init__( self, test_notebook_path: str, golden_dir: str, cleanup: bool = True) -> None: assert _interactive_integration_ready, ( '[interactive_test] dependency is not installed.') assert os.path.exists(golden_dir), '{} does not exist.'.format(golden_dir) assert os.path.isdir(golden_dir), '{} is not a directory.'.format( golden_dir) self._golden_dir = golden_dir self._notebook_executor = notebook_executor.NotebookExecutor( test_notebook_path) self._cleanup = cleanup self._test_urls = {} self._server = None def __enter__(self): self._notebook_executor.execute() self._server = HTTPServer(('', 0), SimpleHTTPRequestHandler) def start_serving(server): server.serve_forever() threading.Thread( target=start_serving, args=[self._server], daemon=True).start() for test_id, output_path in\ self._notebook_executor.output_html_paths.items(): self._test_urls[test_id] = self.base_url + output_path return self def __exit__(self, exc_type, exc_value, traceback): if self._notebook_executor and self._cleanup: self._notebook_executor.cleanup() if self._server: def stop_serving(server): server.shutdown() threading.Thread( target=stop_serving, args=[self._server], daemon=True).start() @property def base_url(self): """The base url where the locally started server serving HTMLs generated by notebook executions.""" assert self._server, 'Server has not started.' host_n_port = self._server.server_address return 'http://{}:{}/'.format(host_n_port[0], host_n_port[1]) @property def test_urls(self): """Mapping from test_id/execution_id to urls serving the output HTML pages generated by the corresponding notebook executions.""" return self._test_urls @property def notebook_path_to_test_id(self): """Mapping from input notebook paths to their obfuscated execution/test ids. """ return self._notebook_executor.notebook_path_to_execution_id def should_skip(): """Whether a screen diff test should be skipped.""" return not ( platform.system() in _SUPPORTED_PLATFORMS and ie.current_env().is_interactive_ready and _interactive_integration_ready) if should_skip(): @unittest.skip( reason='[interactive] and [interactive_test] deps are both required.') @pytest.mark.skip( reason='[interactive] and [interactive_test] deps are both required.') class BaseTestCase(unittest.TestCase): """A skipped base test case if interactive_test dependency is not installed. """ pass else: class BaseTestCase(NeedleTestCase): """A base test case to execute screen diff integration tests.""" # Whether the browser should be headless. _headless = True def __init__(self, *args, **kwargs): """Initializes a test. Some kwargs that could be configured: #. golden_dir=. A directory path pointing to all the golden screenshots as baselines for comparison. #. test_notebook_dir=. A path pointing to a directory of notebook files in ipynb format. #. headless=. Whether the browser should be headless when executing the tests. #. golden_size=<(int, int)>. The size of the screenshot to take and compare. #. cleanup=. Whether to clean up the output directory. Should always be True in automated test environment. When debugging, turn it False to manually check the output for difference. #. threshold=. An image difference threshold, when the image pixel distance is bigger than the value, the test will fail. """ golden_root = kwargs.pop( 'golden_dir', 'apache_beam/runners/interactive/testing/integration/goldens') self._golden_dir = os.path.join(golden_root, platform.system()) self._test_notebook_dir = kwargs.pop( 'test_notebook_dir', 'apache_beam/runners/interactive/testing/integration/test_notebooks') BaseTestCase._headless = kwargs.pop('headless', True) self._test_env = None self._viewport_width, self._viewport_height = kwargs.pop( 'golden_size', (1024, 10000)) self._cleanup = kwargs.pop('cleanup', True) self._threshold = kwargs.pop('threshold', 5000) self.baseline_directory = os.path.join(os.getcwd(), self._golden_dir) self.output_directory = os.path.join( os.getcwd(), self._test_notebook_dir, 'output') super().__init__(*args, **kwargs) @classmethod def get_web_driver(cls): chrome_options = Options() if cls._headless: chrome_options.add_argument('--headless') chrome_options.add_argument('--no-sandbox') chrome_options.add_argument('--disable-dev-shm-usage') chrome_options.add_argument('--force-color-profile=srgb') return NeedleChrome(options=chrome_options) def setUp(self): self.set_viewport_size(self._viewport_width, self._viewport_height) def run(self, result=None): with ScreenDiffIntegrationTestEnvironment(self._test_notebook_dir, self._golden_dir, self._cleanup) as test_env: self._test_env = test_env super().run(result) def explicit_wait(self): """Wait for common elements to be visible.""" WebDriverWait(self.driver, 5).until( expected_conditions.visibility_of_element_located( (By.TAG_NAME, 'facets-overview'))) WebDriverWait(self.driver, 5).until( expected_conditions.visibility_of_element_located( (By.TAG_NAME, 'facets-dive'))) def assert_all(self): """Asserts screenshots for all notebooks in the test_notebook_path.""" for test_id, test_url in self._test_env.test_urls.items(): self.driver.get(test_url) self.explicit_wait() self.assertScreenshot('body', test_id, self._threshold) def assert_single(self, test_id): """Asserts the screenshot for a single test. The given test id will be the name of the golden screenshot.""" test_url = self._test_env.test_urls.get(test_id, None) assert test_url, '{} is not a valid test id.'.format(test_id) self.driver.get(test_url) self.explicit_wait() self.assertScreenshot('body', test_id, self._threshold) def assert_notebook(self, notebook_name): """Asserts the screenshot for a single notebook. The notebook with the given notebook_name under test_notebook_dir will be executed and asserted. """ if not notebook_name.endswith('.ipynb'): notebook_name += '.ipynb' notebook_path = os.path.join(self._test_notebook_dir, notebook_name) test_id = self._test_env.notebook_path_to_test_id.get(notebook_path, None) assert test_id, 'Cannot find notebook with name {}.'.format(notebook_name) self.assert_single(test_id) # This file contains no tests. Below lines are purely for passing lint. if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/testing/mock_env.py0000644000175100017510000000710315206113032025517 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module of mocks to isolated the test environment for each Interactive Beam test. """ import unittest import uuid from unittest.mock import patch from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import DataprocClusterManager from apache_beam.runners.interactive.interactive_environment import InteractiveEnvironment from apache_beam.runners.interactive.testing.mock_ipython import mock_get_ipython def isolated_env(cls: type[unittest.TestCase]): """A class decorator for unittest.TestCase to set up an isolated test environment for Interactive Beam.""" class IsolatedInteractiveEnvironmentTest(cls): def setUp(self): self.env_patchers = [] interactive_path = 'apache_beam.runners.interactive' if ie.current_env().is_interactive_ready: # Group lines of code into a notebook/IPython cell: # with self.cell: # arbitrary python code # ... # arbitrary python code self.ipython_patcher = patch( 'IPython.get_ipython', new_callable=mock_get_ipython) self.cell = self.ipython_patcher.start() self.env_patchers.append(self.ipython_patcher) # self.current_env IS interactive_environment.current_env() in tests. self.ie_patcher = patch( f'{interactive_path}.interactive_environment.current_env') self.m_current_env = self.ie_patcher.start() self.current_env = InteractiveEnvironment() self.m_current_env.return_value = self.current_env self.env_patchers.append(self.ie_patcher) # Patches dataproc cluster creation and deletion. self.create_cluster_patcher = patch.object( DataprocClusterManager, 'create_flink_cluster', mock_create_cluster) self.alt_create_cluster = self.create_cluster_patcher.start() self.env_patchers.append(self.create_cluster_patcher) self.delete_cluster_patcher = patch.object( DataprocClusterManager, 'cleanup') self.m_delete_cluster = self.delete_cluster_patcher.start() self.env_patchers.append(self.delete_cluster_patcher) super().setUp() def tearDown(self): super().tearDown() # Explicitly calls the cleanup instead of letting it be called at exit # when the mocks and isolated env instance are out of scope. self.current_env.cleanup() for patcher in reversed(self.env_patchers): patcher.stop() return IsolatedInteractiveEnvironmentTest def mock_create_cluster(self): """Mocks a cluster creation and populates derived fields.""" self.cluster_metadata.master_url = 'test-url-' + uuid.uuid4().hex self.cluster_metadata.dashboard = 'test-dashboard' # This file contains no tests. Below lines are purely for passing lint. if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/testing/mock_ipython.py0000644000175100017510000000431615206113032026424 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # def mock_get_ipython(): """Mock an ipython environment w/o setting up real ipython kernel. Each entering of get_ipython() invocation will have the prompt increased by one. Grouping arbitrary python code into separate cells using `with` clause. Examples:: # Usage, before each test function, prepend: @patch('IPython.get_ipython', new_callable=mock_get_ipython) # In the test function's signature, add an argument for the patch, e.g.: def some_test(self, cell): # Group lines of code into a cell using the argument: with cell: # arbitrary python code # ... # arbitrary python code # Next cell with prompt increased by one: with cell: # Auto-incremental # arbitrary python code # ... # arbitrary python code """ class MockedGetIpython(object): def __init__(self): self._execution_count = 0 # Mock as if the kernel is connected to a notebook frontend. self.config = {'IPKernelApp': 'mock'} def __call__(self): return self @property def execution_count(self): """Execution count always starts from 1 and is constant within a cell.""" return self._execution_count def __enter__(self): """Marks entering of a cell/prompt.""" self._execution_count = self._execution_count + 1 def __exit__(self, exc_type, exc_value, traceback): """Marks exiting of a cell/prompt.""" pass return MockedGetIpython() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/testing/pipeline_assertion.py0000644000175100017510000001151715206113032027616 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module to verify implicit cache transforms applied by Interactive Beam. For internal use only; no backwards-compatibility guarantees. This utility should only be used by Interactive Beam tests. For example, it can be used to verify if the implicit cache transforms are applied as expected when running a pipeline with the InteractiveRunner. It can also be used to verify if a pipeline fragment has pruned unnecessary transforms. It shouldn't be used to verify equivalence between pipelines if the code to be tested depends on or mutates user code within transforms in pipelines. """ def assert_pipeline_equal(test_case, expected_pipeline, actual_pipeline): """Asserts the equivalence between two given apache_beam.Pipeline instances. Args: test_case: (unittest.TestCase) the unittest testcase where it asserts. expected_pipeline: (Pipeline) the pipeline instance expected. actual_pipeline: (Pipeline) the actual pipeline instance to be asserted. """ expected_pipeline_proto = expected_pipeline.to_runner_api( use_fake_coders=True) actual_pipeline_proto = actual_pipeline.to_runner_api(use_fake_coders=True) assert_pipeline_proto_equal( test_case, expected_pipeline_proto, actual_pipeline_proto) def assert_pipeline_proto_equal( test_case, expected_pipeline_proto, actual_pipeline_proto): """Asserts the equivalence between two pipeline proto representations.""" components1 = expected_pipeline_proto.components components2 = actual_pipeline_proto.components test_case.assertEqual( len(components1.transforms), len(components2.transforms)) test_case.assertEqual( len(components1.pcollections), len(components2.pcollections)) # TODO(BEAM-7926): Update tests and make below 2 assertions assertEqual. test_case.assertLessEqual( len(components1.windowing_strategies), len(components2.windowing_strategies)) test_case.assertLessEqual(len(components1.coders), len(components2.coders)) _assert_transform_equal( test_case, actual_pipeline_proto, actual_pipeline_proto.root_transform_ids[0], expected_pipeline_proto, expected_pipeline_proto.root_transform_ids[0]) def assert_pipeline_proto_contain_top_level_transform( test_case, pipeline_proto, transform_label): """Asserts the top level transforms contain a transform with the given transform label.""" _assert_pipeline_proto_contains_top_level_transform( test_case, pipeline_proto, transform_label, True) def assert_pipeline_proto_not_contain_top_level_transform( test_case, pipeline_proto, transform_label): """Asserts the top level transforms do not contain a transform with the given transform label.""" _assert_pipeline_proto_contains_top_level_transform( test_case, pipeline_proto, transform_label, False) def _assert_pipeline_proto_contains_top_level_transform( test_case, pipeline_proto, transform_label, contain): top_level_transform_labels = pipeline_proto.components.transforms[ pipeline_proto.root_transform_ids[0]].subtransforms test_case.assertEqual( contain, any( transform_label in top_level_transform_label for top_level_transform_label in top_level_transform_labels)) def _assert_transform_equal( test_case, expected_pipeline_proto, expected_transform_id, actual_pipeline_proto, actual_transform_id): """Asserts the equivalence between transforms from two given pipelines. """ transform_proto1 = expected_pipeline_proto.components.transforms[ expected_transform_id] transform_proto2 = actual_pipeline_proto.components.transforms[ actual_transform_id] test_case.assertEqual(transform_proto1.spec.urn, transform_proto2.spec.urn) # Skipping payload checking because PTransforms of the same functionality # could generate different payloads. test_case.assertEqual( len(transform_proto1.subtransforms), len(transform_proto2.subtransforms)) test_case.assertSetEqual( set(transform_proto1.inputs), set(transform_proto2.inputs)) test_case.assertSetEqual( set(transform_proto1.outputs), set(transform_proto2.outputs)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/testing/test_cache_manager.py0000644000175100017510000001046415206113032027516 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import collections import itertools import sys import apache_beam as beam from apache_beam import coders from apache_beam.portability.api import beam_interactive_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.interactive.cache_manager import CacheManager from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import Timestamp class InMemoryCache(CacheManager): """A cache that stores all PCollections in an in-memory map. This is only used for checking the pipeline shape. This can't be used for running the pipeline isn't shared between the SDK and the Runner. """ def __init__(self): self._cached = {} self._pcoders = {} def exists(self, *labels): return self._key(*labels) in self._cached def _latest_version(self, *labels): return True def read(self, *labels, **args): if not self.exists(*labels): return itertools.chain([]), -1 return itertools.chain(self._cached[self._key(*labels)]), None def write(self, value, *labels): if not self.exists(*labels): self._cached[self._key(*labels)] = [] self._cached[self._key(*labels)] += value def save_pcoder(self, pcoder, *labels): self._pcoders[self._key(*labels)] = pcoder def load_pcoder(self, *labels): return self._pcoders[self._key(*labels)] def cleanup(self): self._cached = collections.defaultdict(list) self._pcoders = {} def clear(self, *label): # Noop because in-memory. pass def source(self, *labels): vals = self._cached[self._key(*labels)] return beam.Create(vals) def sink(self, labels, is_capture=False): return beam.Map(lambda _: _) def size(self, *labels): if self.exists(*labels): return sys.getsizeof(self._cached[self._key(*labels)]) return 0 def _key(self, *labels): return '/'.join([l for l in labels]) class NoopSink(beam.PTransform): def expand(self, pcoll): return pcoll | beam.Map(lambda x: x) class FileRecordsBuilder(object): def __init__(self, tag=None): self._header = beam_interactive_api_pb2.TestStreamFileHeader(tag=tag) self._records = [] self._coder = coders.FastPrimitivesCoder() def add_element(self, element, event_time_secs): element_payload = beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=self._coder.encode(element), timestamp=Timestamp.of(event_time_secs).micros) record = beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event. AddElements(elements=[element_payload]))) self._records.append(record) return self def advance_watermark(self, watermark_secs): record = beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload.Event( watermark_event=beam_runner_api_pb2.TestStreamPayload. Event.AdvanceWatermark( new_watermark=Timestamp.of(watermark_secs).micros))) self._records.append(record) return self def advance_processing_time(self, delta_secs): record = beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime( advance_duration=Duration.of(delta_secs).micros))) self._records.append(record) return self def build(self): return [self._header] + self._records ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/user_pipeline_tracker.py0000644000175100017510000001314715206113032026624 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Class that tracks derived/pipeline fragments from user pipelines. For internal use only; no backwards-compatibility guarantees. In the InteractiveRunner the design is to keep the user pipeline unchanged, create a copy of the user pipeline, and modify the copy. When the derived pipeline runs, there should only be per-user pipeline state. This makes sure that derived pipelines can link back to the parent user pipeline. """ import shutil from typing import Iterator from typing import Optional import apache_beam as beam # type: ignore class UserPipelineTracker: """Tracks user pipelines from derived pipelines. This data structure is similar to a disjoint set data structure. A derived pipeline can only have one parent user pipeline. A user pipeline can have many derived pipelines. """ def __init__(self): self._user_pipelines: dict[beam.Pipeline, list[beam.Pipeline]] = {} self._derived_pipelines: dict[beam.Pipeline] = {} self._pid_to_pipelines: dict[beam.Pipeline] = {} def __iter__(self) -> Iterator[beam.Pipeline]: """Iterates through all the user pipelines.""" for p in self._user_pipelines: yield p def _key(self, pipeline: beam.Pipeline) -> str: return str(id(pipeline)) def evict(self, pipeline: beam.Pipeline) -> None: """Evicts the pipeline. Removes the given pipeline and derived pipelines if a user pipeline. Otherwise, removes the given derived pipeline. """ user_pipeline = self.get_user_pipeline(pipeline) if user_pipeline: for d in self._user_pipelines[user_pipeline]: del self._derived_pipelines[d] del self._user_pipelines[user_pipeline] elif pipeline in self._derived_pipelines: del self._derived_pipelines[pipeline] def clear(self) -> None: """Clears the tracker of all user and derived pipelines.""" # Remove all local_tempdir of created pipelines. for p in self._pid_to_pipelines.values(): shutil.rmtree(p.local_tempdir, ignore_errors=True) self._user_pipelines.clear() self._derived_pipelines.clear() self._pid_to_pipelines.clear() def get_pipeline(self, pid: str) -> Optional[beam.Pipeline]: """Returns the pipeline corresponding to the given pipeline id.""" return self._pid_to_pipelines.get(pid, None) def add_user_pipeline(self, p: beam.Pipeline) -> beam.Pipeline: """Adds a user pipeline with an empty set of derived pipelines.""" self._memoize_pipieline(p) # Create a new node for the user pipeline if it doesn't exist already. user_pipeline = self.get_user_pipeline(p) if not user_pipeline: user_pipeline = p self._user_pipelines[p] = [] return user_pipeline def _memoize_pipieline(self, p: beam.Pipeline) -> None: """Memoizes the pid of the pipeline to the pipeline object.""" pid = self._key(p) if pid not in self._pid_to_pipelines: self._pid_to_pipelines[pid] = p def add_derived_pipeline( self, maybe_user_pipeline: beam.Pipeline, derived_pipeline: beam.Pipeline) -> None: """Adds a derived pipeline with the user pipeline. If the `maybe_user_pipeline` is a user pipeline, then the derived pipeline will be added to its set. Otherwise, the derived pipeline will be added to the user pipeline of the `maybe_user_pipeline`. By doing the above one can do: p = beam.Pipeline() derived1 = beam.Pipeline() derived2 = beam.Pipeline() ut = UserPipelineTracker() ut.add_derived_pipeline(p, derived1) ut.add_derived_pipeline(derived1, derived2) # Returns p. ut.get_user_pipeline(derived2) """ self._memoize_pipieline(maybe_user_pipeline) self._memoize_pipieline(derived_pipeline) # Cannot add a derived pipeline twice. assert derived_pipeline not in self._derived_pipelines # Get the "true" user pipeline. This allows for the user to derive a # pipeline from another derived pipeline, use both as arguments, and this # method will still get the correct user pipeline. user = self.add_user_pipeline(maybe_user_pipeline) # Map the derived pipeline to the user pipeline. self._derived_pipelines[derived_pipeline] = user self._user_pipelines[user].append(derived_pipeline) def get_user_pipeline(self, p: beam.Pipeline) -> Optional[beam.Pipeline]: """Returns the user pipeline of the given pipeline. If the given pipeline has no user pipeline, i.e. not added to this tracker, then this returns None. If the given pipeline is a user pipeline then this returns the same pipeline. If the given pipeline is a derived pipeline then this returns the user pipeline. """ # If `p` is a user pipeline then return it. if p in self._user_pipelines: return p # If `p` exists then return its user pipeline. if p in self._derived_pipelines: return self._derived_pipelines[p] # Otherwise, `p` is not in this tracker. return None ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/user_pipeline_tracker_test.py0000644000175100017510000001370715206113032027665 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest import apache_beam as beam from apache_beam.runners.interactive.user_pipeline_tracker import UserPipelineTracker class UserPipelineTrackerTest(unittest.TestCase): def test_getting_unknown_pid_returns_none(self): ut = UserPipelineTracker() p = beam.Pipeline() self.assertIsNone(ut.get_pipeline(str(id(p)))) def test_getting_unknown_pipeline_returns_none(self): ut = UserPipelineTracker() p = beam.Pipeline() self.assertIsNone(ut.get_user_pipeline(p)) def test_no_parent_returns_none(self): ut = UserPipelineTracker() user = beam.Pipeline() derived = beam.Pipeline() orphan = beam.Pipeline() ut.add_derived_pipeline(user, derived) self.assertIsNone(ut.get_user_pipeline(orphan)) def test_get_user_pipeline_is_same(self): ut = UserPipelineTracker() p = beam.Pipeline() ut.add_user_pipeline(p) self.assertIs(ut.get_user_pipeline(p), p) def test_can_add_derived(self): ut = UserPipelineTracker() user = beam.Pipeline() derived = beam.Pipeline() ut.add_derived_pipeline(user, derived) self.assertIs(ut.get_user_pipeline(derived), user) def test_can_add_multiple_derived(self): """Tests that there can be many user pipelines with many derived pipelines. """ ut = UserPipelineTracker() # Add the first set of user and derived pipelines. user1 = beam.Pipeline() derived11 = beam.Pipeline() derived12 = beam.Pipeline() ut.add_derived_pipeline(user1, derived11) ut.add_derived_pipeline(user1, derived12) # Add the second set of user and derived pipelines. user2 = beam.Pipeline() derived21 = beam.Pipeline() derived22 = beam.Pipeline() ut.add_derived_pipeline(user2, derived21) ut.add_derived_pipeline(user2, derived22) # Assert that the user pipelines are correct. self.assertIs(ut.get_user_pipeline(derived11), user1) self.assertIs(ut.get_user_pipeline(derived12), user1) self.assertIs(ut.get_user_pipeline(derived21), user2) self.assertIs(ut.get_user_pipeline(derived22), user2) def test_cannot_have_multiple_parents(self): ut = UserPipelineTracker() user1 = beam.Pipeline() user2 = beam.Pipeline() derived = beam.Pipeline() ut.add_derived_pipeline(user1, derived) with self.assertRaises(AssertionError): ut.add_derived_pipeline(user2, derived) self.assertIs(ut.get_user_pipeline(derived), user1) def test_adding_derived_with_derived_gets_user_pipeline(self): """Tests that one can correctly add a derived pipeline from a derived pipeline and still get the correct user pipeline. """ ut = UserPipelineTracker() user = beam.Pipeline() derived1 = beam.Pipeline() derived2 = beam.Pipeline() # Add the first derived pipeline to the user pipelne. ut.add_derived_pipeline(user, derived1) # Add the second derived pipeline to the first derived pipeline. This should # get the user pipeline of the first and add the second to it. ut.add_derived_pipeline(derived1, derived2) # Asserts that both derived pipelines are under the same user pipeline. self.assertIs(ut.get_user_pipeline(derived1), user) self.assertIs(ut.get_user_pipeline(derived2), user) def test_can_get_pipeline_from_id(self): """Tests the pid -> pipeline memoization.""" ut = UserPipelineTracker() user = beam.Pipeline() derived = beam.Pipeline() ut.add_user_pipeline(user) ut.add_derived_pipeline(user, derived) self.assertIs(ut.get_pipeline(str(id(user))), user) self.assertIs(ut.get_pipeline(str(id(derived))), derived) def test_clear(self): ut = UserPipelineTracker() user = beam.Pipeline() derived = beam.Pipeline() ut.add_derived_pipeline(user, derived) self.assertIs(ut.get_user_pipeline(derived), user) ut.clear() self.assertIsNone(ut.get_user_pipeline(user)) self.assertIsNone(ut.get_user_pipeline(derived)) def test_can_iterate(self): ut = UserPipelineTracker() user1 = beam.Pipeline() derived11 = beam.Pipeline() derived12 = beam.Pipeline() ut.add_derived_pipeline(user1, derived11) ut.add_derived_pipeline(user1, derived12) user2 = beam.Pipeline() derived21 = beam.Pipeline() derived22 = beam.Pipeline() ut.add_derived_pipeline(user2, derived21) ut.add_derived_pipeline(user2, derived22) user_pipelines = set(p for p in ut) self.assertSetEqual(set([user1, user2]), user_pipelines) def test_can_evict_user_pipeline(self): ut = UserPipelineTracker() user1 = beam.Pipeline() derived11 = beam.Pipeline() derived12 = beam.Pipeline() ut.add_derived_pipeline(user1, derived11) ut.add_derived_pipeline(user1, derived12) user2 = beam.Pipeline() derived21 = beam.Pipeline() derived22 = beam.Pipeline() ut.add_derived_pipeline(user2, derived21) ut.add_derived_pipeline(user2, derived22) ut.evict(user1) self.assertIsNone(ut.get_user_pipeline(user1)) self.assertIsNone(ut.get_user_pipeline(derived11)) self.assertIsNone(ut.get_user_pipeline(derived12)) self.assertIs(user2, ut.get_user_pipeline(derived21)) self.assertIs(user2, ut.get_user_pipeline(derived22)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/utils.py0000644000175100017510000004166715206113032023416 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utilities to be used in Interactive Beam. """ import functools import hashlib import importlib import json import logging from typing import Any from typing import Callable from typing import Iterator from typing import Union import pandas as pd import apache_beam as beam from apache_beam.coders import Coder from apache_beam.dataframe.convert import to_pcollection from apache_beam.dataframe.frame_base import DeferredBase from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.pipeline import Pipeline from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.interactive.caching.cacheable import Cacheable from apache_beam.runners.interactive.caching.cacheable import CacheKey from apache_beam.runners.interactive.caching.expression_cache import ExpressionCache from apache_beam.testing.test_stream import WindowedValueHolder from apache_beam.typehints.schemas import named_fields_from_element_type from apache_beam.utils.windowed_value import WindowedValue _LOGGER = logging.getLogger(__name__) # Add line breaks to the IPythonLogHandler's HTML output. _INTERACTIVE_LOG_STYLE = """ """ def to_element_list( reader: Iterator[Union[beam_runner_api_pb2.TestStreamPayload.Event, WindowedValueHolder]], coder: Coder, include_window_info: bool, n: int = None, include_time_events: bool = False, ) -> list[WindowedValue]: """Returns an iterator that properly decodes the elements from the reader. """ # Defining a generator like this makes it easier to limit the count of # elements read. Otherwise, the count limit would need to be duplicated. def elements(): for e in reader: if isinstance(e, beam_runner_api_pb2.TestStreamPayload.Event): if (e.HasField('watermark_event') or e.HasField('processing_time_event')): if include_time_events: yield e else: for tv in e.element_event.elements: decoded = coder.decode(tv.encoded_element) yield ( decoded.windowed_value if include_window_info else decoded.windowed_value.value) elif isinstance(e, WindowedValueHolder): yield ( e.windowed_value if include_window_info else e.windowed_value.value) elif isinstance(e, WindowedValue): yield (e if include_window_info else e.value) else: yield e # Because we can yield multiple elements from a single TestStreamFileRecord, # we have to limit the count here to ensure that `n` is fulfilled. count = 0 for e in elements(): if n and count >= n: break yield e if not isinstance(e, beam_runner_api_pb2.TestStreamPayload.Event): count += 1 def elements_to_df( elements: list[WindowedValue], include_window_info: bool = False, element_type: Any = None) -> 'DataFrame': # noqa: F821 """Parses the given elements into a Dataframe. If the elements are a list of WindowedValues, then it will break out the elements into their own DataFrame and return it. If include_window_info is True, then it will concatenate the windowing information onto the elements DataFrame. """ try: columns_names = [ name for name, _ in named_fields_from_element_type(element_type) ] except TypeError: columns_names = None rows = [] windowed_info = [] for e in elements: rows.append(e.value) if include_window_info: windowed_info.append([e.timestamp.micros, e.windows, e.pane_info]) using_dataframes = isinstance(element_type, pd.DataFrame) using_series = isinstance(element_type, pd.Series) if using_dataframes or using_series: rows_df = pd.concat(rows) else: rows_df = pd.DataFrame(rows, columns=columns_names) if include_window_info and not using_series: windowed_info_df = pd.DataFrame( windowed_info, columns=['event_time', 'windows', 'pane_info']) final_df = pd.concat([rows_df, windowed_info_df], axis=1) else: final_df = rows_df return final_df def register_ipython_log_handler() -> None: """Adds the IPython handler to a dummy parent logger (named 'apache_beam.runners.interactive') of all interactive modules' loggers so that if is_in_notebook, logging displays the logs as HTML in frontends. """ # apache_beam.runners.interactive is not a module, thus this "root" logger is # a dummy one created to hold the IPython log handler. When children loggers # have propagate as True (by default) and logging level as NOTSET (by default, # so the "root" logger's logging level takes effect), the IPython log handler # will be triggered at the "root"'s own logging level. And if a child logger # sets its logging level, it can take control back. interactive_root_logger = logging.getLogger('apache_beam.runners.interactive') if any(isinstance(h, IPythonLogHandler) for h in interactive_root_logger.handlers): return interactive_root_logger.setLevel(logging.INFO) interactive_root_logger.addHandler(IPythonLogHandler()) # Disable the propagation so that logs emitted from interactive modules should # only be handled by loggers and handlers defined within interactive packages. interactive_root_logger.propagate = False class IPythonLogHandler(logging.Handler): """A logging handler to display logs as HTML in IPython backed frontends.""" # TODO(BEAM-7923): Switch to Google hosted CDN once # https://code.google.com/archive/p/google-ajax-apis/issues/637 is resolved. log_template = """
{msg}
""" logging_to_alert_level_map = { logging.CRITICAL: 'danger', logging.ERROR: 'danger', logging.WARNING: 'warning', logging.INFO: 'info', logging.DEBUG: 'dark', logging.NOTSET: 'light' } def emit(self, record): try: from html import escape from IPython.display import HTML from IPython.display import display display(HTML(_INTERACTIVE_LOG_STYLE)) display( HTML( self.log_template.format( level=self.logging_to_alert_level_map[record.levelno], msg=escape(record.msg % record.args)))) except ImportError: pass # NOOP when dependencies are not available. def obfuscate(*inputs: Any) -> str: """Obfuscates any inputs into a hexadecimal string.""" str_inputs = [str(input) for input in inputs] merged_inputs = '_'.join(str_inputs) return hashlib.md5(merged_inputs.encode('utf-8')).hexdigest() class ProgressIndicator(object): """An indicator visualizing code execution in progress.""" # TODO(BEAM-7923): Switch to Google hosted CDN once # https://code.google.com/archive/p/google-ajax-apis/issues/637 is resolved. spinner_template = """
{text}
""" spinner_removal_template = """ $("#{id}").remove();""" def __init__(self, enter_text: str, exit_text: str) -> None: self._id = 'progress_indicator_{}'.format(obfuscate(id(self))) self._enter_text = enter_text self._exit_text = exit_text def __enter__(self): try: from IPython.display import HTML from IPython.display import display from apache_beam.runners.interactive import interactive_environment as ie if ie.current_env().is_in_notebook: display( HTML( self.spinner_template.format( id=self._id, text=self._enter_text))) else: display(self._enter_text) except ImportError as e: _LOGGER.error( 'Please use interactive Beam features in an IPython' 'or notebook environment: %s' % e) def __exit__(self, exc_type, exc_value, traceback): try: from IPython.display import Javascript from IPython.display import display from IPython.display import display_javascript from apache_beam.runners.interactive import interactive_environment as ie if ie.current_env().is_in_notebook: script = self.spinner_removal_template.format(id=self._id) display_javascript( Javascript( ie._JQUERY_WITH_DATATABLE_TEMPLATE.format( customized_script=script))) else: display(self._exit_text) except ImportError as e: _LOGGER.error( 'Please use interactive Beam features in an IPython' 'or notebook environment: %s' % e) def progress_indicated(func: Callable[..., Any]) -> Callable[..., Any]: """A decorator using a unique progress indicator as a context manager to execute the given function within.""" @functools.wraps(func) def run_within_progress_indicator(*args, **kwargs): with ProgressIndicator(f'Processing... {func.__name__}', 'Done.'): return func(*args, **kwargs) return run_within_progress_indicator def as_json(func: Callable[..., Any]) -> Callable[..., str]: """A decorator convert python objects returned by callables to json string. The decorated function should always return an object parsable by json.dumps. If the object is not parsable, the str() of original object is returned instead. """ def return_as_json(*args, **kwargs): try: return_value = func(*args, **kwargs) return json.dumps(return_value) except TypeError: return str(return_value) return return_as_json def deferred_df_to_pcollection(df): assert isinstance(df, DeferredBase), '{} is not a DeferredBase'.format(df) # The proxy is used to output a DataFrame with the correct columns. # # TODO(https://github.com/apache/beam/issues/20577): Once type hints are # implemented for pandas, use those instead of the proxy. cache = ExpressionCache() cache.replace_with_cached(df._expr) proxy = df._expr.proxy() return to_pcollection(df, yield_elements='pandas', label=str(df._expr)), proxy def pcoll_by_name() -> dict[str, beam.PCollection]: """Finds all PCollections by their variable names defined in the notebook.""" from apache_beam.runners.interactive import interactive_environment as ie inspectables = ie.current_env().inspector_with_synthetic.inspectables pcolls = {} for _, inspectable in inspectables.items(): metadata = inspectable['metadata'] if metadata['type'] == 'pcollection': pcolls[metadata['name']] = inspectable['value'] return pcolls def find_pcoll_name(pcoll: beam.PCollection) -> str: """Finds the variable name of a PCollection defined by the user. Returns None if not assigned to any variable. """ from apache_beam.runners.interactive import interactive_environment as ie inspectables = ie.current_env().inspector.inspectables for _, inspectable in inspectables.items(): if inspectable['value'] is pcoll: return inspectable['metadata']['name'] return None def cacheables() -> dict[CacheKey, Cacheable]: """Finds all Cacheables with their CacheKeys.""" from apache_beam.runners.interactive import interactive_environment as ie inspectables = ie.current_env().inspector_with_synthetic.inspectables cacheables = {} for _, inspectable in inspectables.items(): metadata = inspectable['metadata'] if metadata['type'] == 'pcollection': cacheable = Cacheable.from_pcoll(metadata['name'], inspectable['value']) cacheables[cacheable.to_key()] = cacheable return cacheables def watch_sources(pipeline): """Watches the unbounded sources in the pipeline. Sources can output to a PCollection without a user variable reference. In this case the source is not cached. We still want to cache the data so we synthetically create a variable to the intermediate PCollection. """ from apache_beam.pipeline import PipelineVisitor from apache_beam.runners.interactive import interactive_environment as ie retrieved_user_pipeline = ie.current_env().user_pipeline(pipeline) pcoll_to_name = {v: k for k, v in pcoll_by_name().items()} class CacheableUnboundedPCollectionVisitor(PipelineVisitor): def __init__(self): self.unbounded_pcolls = set() def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): if isinstance(transform_node.transform, tuple(ie.current_env().options.recordable_sources)): for pcoll in transform_node.outputs.values(): # Only generate a synthetic var when it's not already watched. For # example, the user could have assigned the unbounded source output # to a variable, watching it again with a different variable name # creates ambiguity. if pcoll not in pcoll_to_name: ie.current_env().watch({'synthetic_var_' + str(id(pcoll)): pcoll}) retrieved_user_pipeline.visit(CacheableUnboundedPCollectionVisitor()) def has_unbounded_sources(pipeline): """Checks if a given pipeline has recordable sources.""" return len(unbounded_sources(pipeline)) > 0 def unbounded_sources(pipeline): """Returns a pipeline's recordable sources.""" from apache_beam.pipeline import PipelineVisitor from apache_beam.runners.interactive import interactive_environment as ie class CheckUnboundednessVisitor(PipelineVisitor): """Visitor checks if there are any unbounded read sources in the Pipeline. Visitor visits all nodes and checks if it is an instance of recordable sources. """ def __init__(self): self.unbounded_sources = [] def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): if isinstance(transform_node.transform, tuple(ie.current_env().options.recordable_sources)): self.unbounded_sources.append(transform_node) v = CheckUnboundednessVisitor() pipeline.visit(v) return v.unbounded_sources def create_var_in_main(name: str, value: Any, watch: bool = True) -> tuple[str, Any]: """Declares a variable in the main module. Args: name: the variable name in the main module. value: the value of the variable. watch: whether to watch it in the interactive environment. Returns: A 2-entry tuple of the variable name and value. """ setattr(importlib.import_module('__main__'), name, value) if watch: from apache_beam.runners.interactive import interactive_environment as ie ie.current_env().watch({name: value}) return name, value def assert_bucket_exists(bucket_name: str) -> None: """Asserts whether the specified GCS bucket with the name bucket_name exists. Logs an error and raises a ValueError if the bucket does not exist. Logs a warning if the bucket cannot be verified to exist. """ try: from google.cloud.exceptions import ClientError from google.cloud.exceptions import NotFound from apache_beam.io.gcp.gcsio import create_storage_client storage_client = create_storage_client(PipelineOptions()) storage_client.get_bucket(bucket_name) except ClientError as e: if isinstance(e, NotFound): _LOGGER.error('%s bucket does not exist!', bucket_name) raise ValueError('Invalid GCS bucket provided!') else: _LOGGER.warning( 'ClientError - unable to verify whether bucket %s exists', bucket_name) except ImportError: _LOGGER.warning( 'ImportError - unable to verify whether bucket %s exists', bucket_name) def detect_pipeline_runner(pipeline): if isinstance(pipeline, Pipeline): from apache_beam.runners.interactive.interactive_runner import InteractiveRunner if isinstance(pipeline.runner, InteractiveRunner): pipeline_runner = pipeline.runner._underlying_runner else: pipeline_runner = pipeline.runner else: pipeline_runner = None return pipeline_runner ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/interactive/utils_test.py0000644000175100017510000003606315206113032024447 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import importlib import json import logging import tempfile import unittest from typing import NamedTuple from unittest.mock import PropertyMock from unittest.mock import patch import numpy as np import pandas as pd import pytest import apache_beam as beam from apache_beam import coders from apache_beam.dataframe.convert import to_dataframe from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.interactive import interactive_beam as ib from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive import utils from apache_beam.runners.interactive.caching.cacheable import Cacheable from apache_beam.runners.interactive.interactive_runner import InteractiveRunner from apache_beam.runners.interactive.testing.mock_ipython import mock_get_ipython from apache_beam.runners.interactive.testing.test_cache_manager import InMemoryCache from apache_beam.runners.portability.flink_runner import FlinkRunner from apache_beam.testing.test_stream import WindowedValueHolder from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.windowed_value import WindowedValue # Protect against environments where apitools library is not available. try: from google.cloud.exceptions import BadRequest from google.cloud.exceptions import NotFound except ImportError: _http_error_imported = False else: _http_error_imported = True class MockStorageClient(): def __init__(self): pass def get_bucket(self, path): if path == 'test-bucket-not-found': raise NotFound('Bucket not found') elif path == 'test-bucket-not-verified': raise BadRequest('Request faulty') class Record(NamedTuple): order_id: int product_id: int quantity: int def windowed_value(e): from apache_beam.transforms.window import GlobalWindow return WindowedValue(e, 1, [GlobalWindow()]) class ParseToDataframeTest(unittest.TestCase): def test_parse_windowedvalue(self): """Tests that WindowedValues are supported but not present. """ els = [windowed_value(('a', 2)), windowed_value(('b', 3))] actual_df = utils.elements_to_df(els, include_window_info=False) expected_df = pd.DataFrame([['a', 2], ['b', 3]], columns=[0, 1]) # check_like so that ordering of indices doesn't matter. pd.testing.assert_frame_equal(actual_df, expected_df, check_like=True) def test_parse_windowedvalue_with_window_info(self): """Tests that WindowedValues are supported and have their own columns. """ els = [windowed_value(('a', 2)), windowed_value(('b', 3))] actual_df = utils.elements_to_df(els, include_window_info=True) expected_df = pd.DataFrame( [['a', 2, int(1e6), els[0].windows, els[0].pane_info], ['b', 3, int(1e6), els[1].windows, els[1].pane_info]], columns=[0, 1, 'event_time', 'windows', 'pane_info']) # check_like so that ordering of indices doesn't matter. pd.testing.assert_frame_equal(actual_df, expected_df, check_like=True) def test_parse_windowedvalue_with_dicts(self): """Tests that dicts play well with WindowedValues. """ els = [ windowed_value({ 'b': 2, 'd': 4 }), windowed_value({ 'a': 1, 'b': 2, 'c': 3 }) ] actual_df = utils.elements_to_df(els, include_window_info=True) expected_df = pd.DataFrame( [[np.nan, 2, np.nan, 4, int(1e6), els[0].windows, els[0].pane_info], [1, 2, 3, np.nan, int(1e6), els[1].windows, els[1].pane_info]], columns=['a', 'b', 'c', 'd', 'event_time', 'windows', 'pane_info']) # check_like so that ordering of indices doesn't matter. pd.testing.assert_frame_equal(actual_df, expected_df, check_like=True) def test_parse_dataframes(self): """Tests that it correctly parses a DataFrame. """ deferred = to_dataframe(beam.Pipeline() | beam.Create([Record(0, 0, 0)])) els = [windowed_value(pd.DataFrame(Record(n, 0, 0))) for n in range(10)] actual_df = utils.elements_to_df( els, element_type=deferred._expr.proxy()).reset_index(drop=True) expected_df = pd.concat([e.value for e in els], ignore_index=True) pd.testing.assert_frame_equal(actual_df, expected_df) def test_parse_series(self): """Tests that it correctly parses a Pandas Series. """ deferred = to_dataframe(beam.Pipeline() | beam.Create([Record(0, 0, 0)]))['order_id'] els = [windowed_value(pd.Series([n])) for n in range(10)] actual_df = utils.elements_to_df( els, element_type=deferred._expr.proxy()).reset_index(drop=True) expected_df = pd.concat([e.value for e in els], ignore_index=True) pd.testing.assert_series_equal(actual_df, expected_df) class ToElementListTest(unittest.TestCase): def test_test_stream_payload_events(self): """Tests that the to_element_list can limit the count in a single bundle.""" coder = coders.FastPrimitivesCoder() def reader(): element_payload = [ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=coder.encode( WindowedValueHolder(WindowedValue(e, 0, []))), timestamp=Timestamp.of(0).micros) for e in range(10) ] event = beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event.AddElements( elements=element_payload)) yield event # The reader creates 10 elements in a single TestStreamPayload but we limit # the number of elements read to 5 here. This tests that the to_element_list # can limit the number of elements in a single bundle. elements = utils.to_element_list( reader(), coder, include_window_info=False, n=5) self.assertSequenceEqual(list(elements), list(range(5))) def test_element_limit_count(self): """Tests that the to_element_list can limit the count.""" elements = utils.to_element_list( iter(range(10)), None, include_window_info=False, n=5) self.assertSequenceEqual(list(elements), list(range(5))) @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') class IPythonLogHandlerTest(unittest.TestCase): def setUp(self): utils.register_ipython_log_handler() self._interactive_root_logger = logging.getLogger( 'apache_beam.runners.interactive') def test_ipython_log_handler_not_double_registered(self): utils.register_ipython_log_handler() ipython_log_handlers = list( filter( lambda x: isinstance(x, utils.IPythonLogHandler), [handler for handler in self._interactive_root_logger.handlers])) self.assertEqual(1, len(ipython_log_handlers)) @patch('apache_beam.runners.interactive.utils.IPythonLogHandler.emit') def test_default_logging_level_is_info(self, mock_emit): # By default the logging level of loggers and log handlers are NOTSET. Also, # the propagation is default to true for all loggers. In this scenario, all # loggings from child loggers will be propagated to the interactive "root" # logger which is set to INFO level that gets handled by the sole log # handler IPythonLogHandler which is set to NOTSET. The effect will be # everything >= info level will be logged through IPython.display to # all frontends connected to current kernel. dummy_logger = logging.getLogger('apache_beam.runners.interactive.dummy1') dummy_logger.info('info') mock_emit.assert_called_once() dummy_logger.debug('debug') # Emit is not called, so it's still called once. mock_emit.assert_called_once() @patch('apache_beam.runners.interactive.utils.IPythonLogHandler.emit') def test_child_module_logger_can_override_logging_level(self, mock_emit): # When a child logger's logging level is configured to something that is not # NOTSET, it takes back the logging control from the interactive "root" # logger by not propagating anything. dummy_logger = logging.getLogger('apache_beam.runners.interactive.dummy2') dummy_logger.setLevel(logging.DEBUG) mock_emit.assert_not_called() dummy_logger.debug('debug') # Because the dummy child logger is configured to log at DEBUG level, it # now propagates DEBUG loggings to the interactive "root" logger. mock_emit.assert_called_once() # When the dummy child logger is configured to log at CRITICAL level, it # will only propagate CRITICAL loggings to the interactive "root" logger. dummy_logger.setLevel(logging.CRITICAL) # Error loggings will not be handled now. dummy_logger.error('error') # Emit is not called, so it's still called once. mock_emit.assert_called_once() @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') @pytest.mark.skipif( not ie.current_env().is_interactive_ready, reason='[interactive] dependency is not installed.') class ProgressIndicatorTest(unittest.TestCase): def setUp(self): self.gcs_patcher = patch( 'apache_beam.io.gcp.gcsfilesystem.GCSFileSystem.delete') self.gcs_patcher.start() ie.new_env() @patch('IPython.get_ipython', new_callable=mock_get_ipython) @patch( 'apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.is_in_notebook', new_callable=PropertyMock) def test_progress_in_plain_text_when_not_in_notebook( self, mocked_is_in_notebook, unused): mocked_is_in_notebook.return_value = False with patch('IPython.display.display') as mocked_display: @utils.progress_indicated def progress_indicated_dummy(): mocked_display.assert_any_call('Processing... progress_indicated_dummy') progress_indicated_dummy() mocked_display.assert_any_call('Done.') @patch('IPython.get_ipython', new_callable=mock_get_ipython) @patch( 'apache_beam.runners.interactive.interactive_environment' '.InteractiveEnvironment.is_in_notebook', new_callable=PropertyMock) def test_progress_in_HTML_JS_when_in_notebook( self, mocked_is_in_notebook, unused): mocked_is_in_notebook.return_value = True with patch('IPython.display.HTML') as mocked_html,\ patch('IPython.display.Javascript') as mocked_js: with utils.ProgressIndicator('enter', 'exit'): mocked_html.assert_called() mocked_js.assert_called() def tearDown(self): self.gcs_patcher.stop() @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') class MessagingUtilTest(unittest.TestCase): SAMPLE_DATA = {'a': [1, 2, 3], 'b': 4, 'c': '5', 'd': {'e': 'f'}} def setUp(self): self.gcs_patcher = patch( 'apache_beam.io.gcp.gcsfilesystem.GCSFileSystem.delete') self.gcs_patcher.start() ie.new_env() def test_as_json_decorator(self): @utils.as_json def dummy(): return MessagingUtilTest.SAMPLE_DATA # As of Python 3.6, for the CPython implementation of Python, # dictionaries remember the order of items inserted. self.assertEqual(json.loads(dummy()), MessagingUtilTest.SAMPLE_DATA) def tearDown(self): self.gcs_patcher.stop() class GeneralUtilTest(unittest.TestCase): def test_pcoll_by_name(self): p = beam.Pipeline() pcoll = p | beam.Create([1]) ib.watch({'p': p, 'pcoll': pcoll}) name_to_pcoll = utils.pcoll_by_name() self.assertIn('pcoll', name_to_pcoll) def test_cacheables(self): p2 = beam.Pipeline() pcoll2 = p2 | beam.Create([2]) ib.watch({'p2': p2, 'pcoll2': pcoll2}) cacheables = utils.cacheables() cacheable_key = Cacheable.from_pcoll('pcoll2', pcoll2).to_key() self.assertIn(cacheable_key, cacheables) def test_has_unbounded_source(self): p = beam.Pipeline() ie.current_env().set_cache_manager(InMemoryCache(), p) _ = p | 'ReadUnboundedSource' >> beam.io.ReadFromPubSub( subscription='projects/fake-project/subscriptions/fake_sub') self.assertTrue(utils.has_unbounded_sources(p)) def test_not_has_unbounded_source(self): p = beam.Pipeline() ie.current_env().set_cache_manager(InMemoryCache(), p) with tempfile.NamedTemporaryFile(delete=False) as f: f.write(b'test') _ = p | 'ReadBoundedSource' >> beam.io.ReadFromText(f.name) self.assertFalse(utils.has_unbounded_sources(p)) def test_find_pcoll_name(self): p = beam.Pipeline() pcoll = p | beam.Create([1, 2, 3]) ib.watch({ 'p_test_find_pcoll_name': p, 'pcoll_test_find_pcoll_name': pcoll, }) self.assertEqual('pcoll_test_find_pcoll_name', utils.find_pcoll_name(pcoll)) def test_create_var_in_main(self): name = 'test_create_var_in_main' value = Record(0, 0, 0) _ = utils.create_var_in_main(name, value) main_session = importlib.import_module('__main__') self.assertIs(getattr(main_session, name, None), value) @patch('google.cloud.storage.Client', return_value=MockStorageClient()) @unittest.skipIf(not _http_error_imported, 'http errors are not imported.') class GCSUtilsTest(unittest.TestCase): @patch('google.cloud.storage.Client.get_bucket') def test_assert_bucket_exists_not_found(self, mock_response, mock_client): with self.assertRaises(ValueError): utils.assert_bucket_exists('test-bucket-not-found') @patch('google.cloud.storage.Client.get_bucket') def test_assert_bucket_exists_not_verified(self, mock_response, mock_client): from apache_beam.runners.interactive.utils import _LOGGER with self.assertLogs(_LOGGER, level='WARNING'): utils.assert_bucket_exists('test-bucket-not-verified') @patch('google.cloud.storage.Client.get_bucket') def test_assert_bucket_exists_found(self, mock_response, mock_client): utils.assert_bucket_exists('') class PipelineUtilTest(unittest.TestCase): def test_detect_pipeline_underlying_runner(self): p = beam.Pipeline(InteractiveRunner(underlying_runner=FlinkRunner())) pipeline_runner = utils.detect_pipeline_runner(p) self.assertTrue(isinstance(pipeline_runner, FlinkRunner)) def test_detect_pipeline_no_underlying_runner(self): p = beam.Pipeline(InteractiveRunner()) pipeline_runner = utils.detect_pipeline_runner(p) from apache_beam.runners.direct.direct_runner import DirectRunner self.assertTrue(isinstance(pipeline_runner, DirectRunner)) def test_detect_pipeline_no_runner(self): pipeline_runner = utils.detect_pipeline_runner(None) self.assertEqual(pipeline_runner, None) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.6554368 apache_beam-2.74.0/apache_beam/runners/internal/0000755000175100017510000000000015206113160021167 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/internal/__init__.py0000644000175100017510000000142015206113032023273 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/internal/names.py0000644000175100017510000000262315206113032022645 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Various names shared by more than one runner.""" # All constants are for internal use only; no backwards-compatibility # guarantees. # Current value is hardcoded in Dataflow internal infrastructure; # please don't change without a review from Dataflow maintainers. STAGED_SDK_SOURCES_FILENAME = 'dataflow_python_sdk.tar' PICKLED_MAIN_SESSION_FILE = 'pickled_main_session' STAGED_PIPELINE_FILENAME = "pipeline.pb" STAGED_PIPELINE_URL_METADATA_FIELD = 'pipeline_url' # Package names for different distributions BEAM_PACKAGE_NAME = 'apache-beam' # SDK identifiers for different distributions BEAM_SDK_NAME = 'Apache Beam SDK for Python' ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.656437 apache_beam-2.74.0/apache_beam/runners/job/0000755000175100017510000000000015206113160020125 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/job/__init__.py0000644000175100017510000000142015206113032022231 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/job/manager.py0000644000175100017510000000347015206113032022113 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A object to control to the Job API Co-Process """ # pytype: skip-file import logging import subprocess import time import grpc from apache_beam.portability.api import beam_job_api_pb2_grpc _LOGGER = logging.getLogger(__name__) class DockerRPCManager(object): """A native co-process to start a contianer that speaks the JobApi """ def __init__(self, run_command=None): # TODO(BEAM-2431): Change this to a docker container from a command. self.process = subprocess.Popen([ 'python', '-m', 'apache_beam.runners.experimental.python_rpc_direct.server' ]) self.channel = grpc.insecure_channel('localhost:50051') self.service = beam_job_api_pb2_grpc.JobServiceStub(self.channel) # Sleep for 2 seconds for process to start completely # This is just for the co-process and would be removed # once we migrate to docker. time.sleep(2) def __del__(self): """Terminate the co-process when the manager is GC'ed """ _LOGGER.info('Shutting the co-process') self.process.terminate() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/job/utils.py0000644000175100017510000000246315206113032021642 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utility functions for efficiently processing with the job API """ # pytype: skip-file import json import logging from google.protobuf import json_format from google.protobuf import struct_pb2 def dict_to_struct(dict_obj: dict) -> struct_pb2.Struct: try: return json_format.ParseDict(dict_obj, struct_pb2.Struct()) except json_format.ParseError: logging.error('Failed to parse dict %s', dict_obj) raise def struct_to_dict(struct_obj: struct_pb2.Struct) -> dict: return json.loads(json_format.MessageToJson(struct_obj)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/pipeline_context.py0000644000175100017510000002756615206113032023314 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utility class for serializing pipelines via the runner API. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file # mypy: disallow-untyped-defs from typing import Any from typing import Generic from typing import Iterable from typing import Mapping from typing import Optional from typing import TypeVar from typing import Union from google.protobuf import message from typing_extensions import Protocol from apache_beam import coders from apache_beam import pipeline from apache_beam import pvalue from apache_beam.coders.coder_impl import IterableStateReader from apache_beam.coders.coder_impl import IterableStateWriter from apache_beam.internal import pickler from apache_beam.pipeline import ComponentIdMap from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.transforms import core from apache_beam.transforms import environments from apache_beam.transforms.resources import merge_resource_hints from apache_beam.typehints import native_type_compatibility PortableObjectT = TypeVar('PortableObjectT', bound='PortableObject') class PortableObject(Protocol): def to_runner_api(self, __context: 'PipelineContext') -> Any: pass @classmethod def from_runner_api(cls, __proto: Any, __context: 'PipelineContext') -> Any: pass class _PipelineContextMap(Generic[PortableObjectT]): """This is a bi-directional map between objects and ids. Under the hood it encodes and decodes these objects into runner API representations. """ def __init__( self, context: 'PipelineContext', obj_type: type[PortableObjectT], namespace: str, proto_map: Optional[Mapping[str, message.Message]] = None) -> None: self._pipeline_context = context self._obj_type = obj_type self._namespace = namespace self._obj_to_id: dict[Any, str] = {} self._id_to_obj: dict[str, Any] = {} self._id_to_proto = dict(proto_map) if proto_map else {} def populate_map(self, proto_map: Mapping[str, message.Message]) -> None: for id, proto in self._id_to_proto.items(): proto_map[id].CopyFrom(proto) def get_id(self, obj: PortableObjectT, label: Optional[str] = None) -> str: if obj not in self._obj_to_id: id = self._pipeline_context.component_id_map.get_or_assign( obj, self._obj_type, label) self._id_to_obj[id] = obj self._obj_to_id[obj] = id self._id_to_proto[id] = obj.to_runner_api(self._pipeline_context) return self._obj_to_id[obj] def get_proto( self, obj: PortableObjectT, label: Optional[str] = None) -> message.Message: return self._id_to_proto[self.get_id(obj, label)] def get_by_id(self, id: str) -> PortableObjectT: if id not in self._id_to_obj: self._id_to_obj[id] = self._obj_type.from_runner_api( self._id_to_proto[id], self._pipeline_context) return self._id_to_obj[id] def get_by_proto( self, maybe_new_proto: message.Message, label: Optional[str] = None, deduplicate: bool = False) -> str: # TODO: this method may not be safe for arbitrary protos due to # xlang concerns, hence limiting usage to the only current use-case it has. # See: https://github.com/apache/beam/pull/14390#discussion_r616062377 assert isinstance(maybe_new_proto, beam_runner_api_pb2.Environment) obj = self._obj_type.from_runner_api( maybe_new_proto, self._pipeline_context) if deduplicate: if obj in self._obj_to_id: return self._obj_to_id[obj] for id, proto in self._id_to_proto.items(): if proto == maybe_new_proto: return id return self.put_proto( self._pipeline_context.component_id_map.get_or_assign( obj=obj, obj_type=self._obj_type, label=label), maybe_new_proto) def get_id_to_proto_map(self) -> dict[str, message.Message]: return self._id_to_proto def get_proto_from_id(self, id: str) -> message.Message: return self.get_id_to_proto_map()[id] def put_proto( self, id: str, proto: message.Message, ignore_duplicates: bool = False) -> str: if not ignore_duplicates and id in self._id_to_proto: raise ValueError("Id '%s' is already taken." % id) elif (ignore_duplicates and id in self._id_to_proto and self._id_to_proto[id] != proto): raise ValueError( 'Cannot insert different protos %r and %r with the same ID %r', self._id_to_proto[id], proto, id) self._id_to_proto[id] = proto return id def __getitem__(self, id: str) -> Any: return self.get_by_id(id) def __contains__(self, id: str) -> bool: return id in self._id_to_proto class PipelineContext(object): """For internal use only; no backwards-compatibility guarantees. Used for accessing and constructing the referenced objects of a Pipeline. """ def __init__( self, proto: Optional[Union[beam_runner_api_pb2.Components, beam_fn_api_pb2.ProcessBundleDescriptor]] = None, component_id_map: Optional[pipeline.ComponentIdMap] = None, default_environment: Optional[environments.Environment] = None, use_fake_coders: bool = False, iterable_state_read: Optional[IterableStateReader] = None, iterable_state_write: Optional[IterableStateWriter] = None, namespace: str = 'ref', requirements: Iterable[str] = (), ) -> None: if isinstance(proto, beam_fn_api_pb2.ProcessBundleDescriptor): proto = beam_runner_api_pb2.Components( coders=dict(proto.coders.items()), windowing_strategies=dict(proto.windowing_strategies.items()), environments=dict(proto.environments.items())) self.component_id_map = component_id_map or ComponentIdMap(namespace) assert self.component_id_map.namespace == namespace # TODO(https://github.com/apache/beam/issues/20827) Initialize # component_id_map with objects from proto. self.transforms = _PipelineContextMap( self, pipeline.AppliedPTransform, namespace, proto.transforms if proto is not None else None) self.pcollections = _PipelineContextMap( self, pvalue.PCollection, namespace, proto.pcollections if proto is not None else None) self.coders = _PipelineContextMap( self, coders.Coder, namespace, proto.coders if proto is not None else None) self.windowing_strategies = _PipelineContextMap( self, core.Windowing, namespace, proto.windowing_strategies if proto is not None else None) self.environments = _PipelineContextMap( self, environments.Environment, namespace, proto.environments if proto is not None else None) if default_environment is None: default_environment = environments.DefaultEnvironment() self._default_environment_id: str = self.environments.get_id( default_environment, label='default_environment') self.use_fake_coders = use_fake_coders self.deterministic_coder_map: Mapping[coders.Coder, coders.Coder] = {} self.iterable_state_read = iterable_state_read self.iterable_state_write = iterable_state_write self._requirements = set(requirements) self.enable_best_effort_deterministic_pickling = False self.enable_stable_code_identifier_pickling = False def add_requirement(self, requirement: str) -> None: self._requirements.add(requirement) def requirements(self) -> frozenset[str]: return frozenset(self._requirements) # If fake coders are requested, return a pickled version of the element type # rather than an actual coder. The element type is required for some runners, # as well as performing a round-trip through protos. # TODO(https://github.com/apache/beam/issues/18490): Remove once this is no # longer needed. def coder_id_from_element_type( self, element_type: Any, requires_deterministic_key_coder: Optional[str] = None) -> str: if self.use_fake_coders: return pickler.dumps(element_type).decode('ascii') else: coder = coders.registry.get_coder(element_type) if requires_deterministic_key_coder: coder = coders.TupleCoder([ self.deterministic_coder( coder.key_coder(), requires_deterministic_key_coder), coder.value_coder() ]) return self.coders.get_id(coder) def deterministic_coder(self, coder: coders.Coder, msg: str) -> coders.Coder: if coder not in self.deterministic_coder_map: self.deterministic_coder_map[coder] = coder.as_deterministic_coder(msg) # type: ignore return self.deterministic_coder_map[coder] def element_type_from_coder_id(self, coder_id: str) -> Any: if self.use_fake_coders or coder_id not in self.coders: return pickler.loads(coder_id) else: return native_type_compatibility.convert_to_beam_type( self.coders[coder_id].to_type_hint()) @staticmethod def from_runner_api( proto: beam_runner_api_pb2.Components) -> 'PipelineContext': return PipelineContext(proto) def to_runner_api(self) -> beam_runner_api_pb2.Components: context_proto = beam_runner_api_pb2.Components() self.transforms.populate_map(context_proto.transforms) self.pcollections.populate_map(context_proto.pcollections) self.coders.populate_map(context_proto.coders) self.windowing_strategies.populate_map(context_proto.windowing_strategies) self.environments.populate_map(context_proto.environments) return context_proto def default_environment_id(self) -> str: return self._default_environment_id def get_environment_id_for_resource_hints( self, hints: dict[str, bytes]) -> str: """Returns an environment id that has necessary resource hints.""" if not hints: return self.default_environment_id() def get_or_create_environment_with_resource_hints( template_env_id: str, resource_hints: dict[str, bytes], ) -> str: """Creates an environment that has necessary hints and returns its id.""" template_env = self.environments.get_proto_from_id(template_env_id) cloned_env = beam_runner_api_pb2.Environment() # (TODO https://github.com/apache/beam/issues/25615) # Remove the suppress warning for type once mypy is updated to 0.941 or # higher. # mypy 0.790 throws the warning below but 0.941 doesn't. # error: Argument 1 to "CopyFrom" of "Message" has incompatible type # "Message"; expected "Environment" [arg-type] # Here, Environment is a subclass of Message but mypy still # throws an error. cloned_env.CopyFrom(template_env) cloned_env.resource_hints.clear() cloned_env.resource_hints.update(resource_hints) return self.environments.get_by_proto( cloned_env, label='environment_with_resource_hints', deduplicate=True) default_env_id = self.default_environment_id() env_hints = self.environments.get_by_id(default_env_id).resource_hints() hints = merge_resource_hints(outer_hints=env_hints, inner_hints=hints) maybe_new_env_id = get_or_create_environment_with_resource_hints( default_env_id, hints) return maybe_new_env_id ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/pipeline_context_test.py0000644000175100017510000000721715206113032024342 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the windowing classes.""" # pytype: skip-file import unittest from apache_beam import coders from apache_beam.runners import pipeline_context from apache_beam.transforms import environments class PipelineContextTest(unittest.TestCase): def test_deduplication(self): context = pipeline_context.PipelineContext() bytes_coder_ref = context.coders.get_id(coders.BytesCoder()) bytes_coder_ref2 = context.coders.get_id(coders.BytesCoder()) self.assertEqual(bytes_coder_ref, bytes_coder_ref2) def test_deduplication_by_proto(self): context = pipeline_context.PipelineContext() env_proto = environments.SubprocessSDKEnvironment( command_string="foo").to_runner_api(None) env_ref_1 = context.environments.get_by_proto(env_proto) env_ref_2 = context.environments.get_by_proto(env_proto, deduplicate=True) self.assertEqual(env_ref_1, env_ref_2) def test_equal_environments_are_deduplicated_when_fetched_by_obj_or_proto( self): context = pipeline_context.PipelineContext() env = environments.SubprocessSDKEnvironment(command_string="foo") env_proto = env.to_runner_api(None) id_from_proto = context.environments.get_by_proto(env_proto) id_from_obj = context.environments.get_id(env) self.assertEqual(id_from_obj, id_from_proto) self.assertEqual( context.environments.get_by_id(id_from_obj).command_string, "foo") env = environments.SubprocessSDKEnvironment(command_string="bar") env_proto = env.to_runner_api(None) id_from_obj = context.environments.get_id(env) id_from_proto = context.environments.get_by_proto( env_proto, deduplicate=True) self.assertEqual(id_from_obj, id_from_proto) self.assertEqual( context.environments.get_by_id(id_from_obj).command_string, "bar") def test_serialization(self): context = pipeline_context.PipelineContext() float_coder_ref = context.coders.get_id(coders.FloatCoder()) bytes_coder_ref = context.coders.get_id(coders.BytesCoder()) proto = context.to_runner_api() context2 = pipeline_context.PipelineContext.from_runner_api(proto) self.assertEqual( coders.FloatCoder(), context2.coders.get_by_id(float_coder_ref)) self.assertEqual( coders.BytesCoder(), context2.coders.get_by_id(bytes_coder_ref)) def test_common_id_assignment(self): context = pipeline_context.PipelineContext() float_coder_ref = context.coders.get_id(coders.FloatCoder()) bytes_coder_ref = context.coders.get_id(coders.BytesCoder()) context2 = pipeline_context.PipelineContext( component_id_map=context.component_id_map) bytes_coder_ref2 = context2.coders.get_id(coders.BytesCoder()) float_coder_ref2 = context2.coders.get_id(coders.FloatCoder()) self.assertEqual(bytes_coder_ref, bytes_coder_ref2) self.assertEqual(float_coder_ref, float_coder_ref2) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/pipeline_utils.py0000644000175100017510000002402515206113032022753 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Pipeline manipulation utilities useful for many runners. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-fileimport collections import collections import copy from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.transforms import environments from apache_beam.typehints import typehints def group_by_key_input_visitor(deterministic_key_coders=True): # Importing here to avoid a circular dependency # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.pipeline import PipelineVisitor from apache_beam.transforms.core import GroupByKey class GroupByKeyInputVisitor(PipelineVisitor): """A visitor that replaces `Any` element type for input `PCollection` of a `GroupByKey` with a `KV` type. TODO(BEAM-115): Once Python SDK is compatible with the new Runner API, we could directly replace the coder instead of mutating the element type. """ def __init__(self, deterministic_key_coders=True): self.deterministic_key_coders = deterministic_key_coders def enter_composite_transform(self, transform_node): self.visit_transform(transform_node) def visit_transform(self, transform_node): if isinstance(transform_node.transform, GroupByKey): pcoll = transform_node.inputs[0] pcoll.element_type = typehints.coerce_to_kv_type( pcoll.element_type, transform_node.full_label) pcoll.requires_deterministic_key_coder = ( self.deterministic_key_coders and transform_node.full_label) key_type, value_type = pcoll.element_type.tuple_types if transform_node.outputs: key = next(iter(transform_node.outputs.keys())) transform_node.outputs[key].element_type = typehints.KV[ key_type, typehints.Iterable[value_type]] transform_node.outputs[key].requires_deterministic_key_coder = ( self.deterministic_key_coders and transform_node.full_label) return GroupByKeyInputVisitor(deterministic_key_coders) def validate_pipeline_graph(pipeline_proto): """Ensures this is a correctly constructed Beam pipeline. """ def get_coder(pcoll_id): return pipeline_proto.components.coders[ pipeline_proto.components.pcollections[pcoll_id].coder_id] def validate_transform(transform_id): transform_proto = pipeline_proto.components.transforms[transform_id] # Currently the only validation we perform is that GBK operations have # their coders set properly. if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn: if len(transform_proto.inputs) != 1: raise ValueError("Unexpected number of inputs: %s" % transform_proto) if len(transform_proto.outputs) != 1: raise ValueError("Unexpected number of outputs: %s" % transform_proto) input_coder = get_coder(next(iter(transform_proto.inputs.values()))) output_coder = get_coder(next(iter(transform_proto.outputs.values()))) if input_coder.spec.urn != common_urns.coders.KV.urn: raise ValueError( "Bad coder for input of %s: %s" % (transform_id, input_coder)) if output_coder.spec.urn != common_urns.coders.KV.urn: raise ValueError( "Bad coder for output of %s: %s" % (transform_id, output_coder)) output_values_coder = pipeline_proto.components.coders[ output_coder.component_coder_ids[1]] if (input_coder.component_coder_ids[0] != output_coder.component_coder_ids[0] or output_values_coder.spec.urn != common_urns.coders.ITERABLE.urn or output_values_coder.component_coder_ids[0] != input_coder.component_coder_ids[1]): raise ValueError( "Incompatible input coder %s and output coder %s for transform %s" % (transform_id, input_coder, output_coder)) elif transform_proto.spec.urn == common_urns.primitives.ASSIGN_WINDOWS.urn: if not transform_proto.inputs: raise ValueError("Missing input for transform: %s" % transform_proto) elif transform_proto.spec.urn == common_urns.primitives.PAR_DO.urn: if not transform_proto.inputs: raise ValueError("Missing input for transform: %s" % transform_proto) for t in transform_proto.subtransforms: validate_transform(t) for t in pipeline_proto.root_transform_ids: validate_transform(t) def _dep_key(dep): if dep.type_urn == common_urns.artifact_types.FILE.urn: payload = beam_runner_api_pb2.ArtifactFilePayload.FromString( dep.type_payload) if payload.sha256: type_info = 'sha256', payload.sha256 else: type_info = 'path', payload.path elif dep.type_urn == common_urns.artifact_types.URL.urn: payload = beam_runner_api_pb2.ArtifactUrlPayload.FromString( dep.type_payload) if payload.sha256: type_info = 'sha256', payload.sha256 else: type_info = 'url', payload.url else: type_info = dep.type_urn, dep.type_payload return type_info, dep.role_urn, dep.role_payload def _expanded_dep_keys(dep): if (dep.type_urn == common_urns.artifact_types.FILE.urn and dep.role_urn == common_urns.artifact_roles.STAGING_TO.urn): payload = beam_runner_api_pb2.ArtifactFilePayload.FromString( dep.type_payload) role = beam_runner_api_pb2.ArtifactStagingToRolePayload.FromString( dep.role_payload) if role.staged_name == 'submission_environment_dependencies.txt': return elif role.staged_name == 'requirements.txt': with open(payload.path) as fin: for line in fin: yield 'requirements.txt', line.strip() return yield _dep_key(dep) def _base_env_key(env, include_deps=True): return ( env.urn, env.payload, tuple(sorted(env.capabilities)), tuple(sorted(env.resource_hints.items())), tuple(sorted(_dep_key(dep) for dep in env.dependencies)) if include_deps else None) def _env_key(env): return tuple( sorted( _base_env_key(e) for e in environments.expand_anyof_environments(env))) def merge_common_environments(pipeline_proto, inplace=False): canonical_environments = collections.defaultdict(list) for env_id, env in pipeline_proto.components.environments.items(): canonical_environments[_env_key(env)].append(env_id) if len(canonical_environments) == len(pipeline_proto.components.environments): # All environments are already sufficiently distinct. return pipeline_proto environment_remappings = { e: es[0] for es in canonical_environments.values() for e in es } return update_environments(pipeline_proto, environment_remappings, inplace) def merge_superset_dep_environments(pipeline_proto): """Merges all environemnts A and B where A and B are equivalent except that A has a superset of the dependencies of B. """ docker_envs = {} for env_id, env in pipeline_proto.components.environments.items(): docker_env = environments.resolve_anyof_environment( env, common_urns.environments.DOCKER.urn) if docker_env.urn == common_urns.environments.DOCKER.urn: docker_envs[env_id] = docker_env has_base_and_dep = collections.defaultdict(set) env_scores = { env_id: (len(env.dependencies), env_id) for (env_id, env) in docker_envs.items() } for env_id, env in docker_envs.items(): base_key = _base_env_key(env, include_deps=False) has_base_and_dep[base_key, None].add(env_id) for dep in env.dependencies: for dep_key in _expanded_dep_keys(dep): has_base_and_dep[base_key, dep_key].add(env_id) environment_remappings = {} for env_id, env in docker_envs.items(): base_key = _base_env_key(env, include_deps=False) # This is the set of all environments that have at least all of env's deps. candidates = set.intersection( has_base_and_dep[base_key, None], *[ has_base_and_dep[base_key, dep_key] for dep in env.dependencies for dep_key in _expanded_dep_keys(dep) ]) # Choose the maximal one. best = max(candidates, key=env_scores.get) if best != env_id: environment_remappings[env_id] = best return update_environments(pipeline_proto, environment_remappings) def update_environments(pipeline_proto, environment_remappings, inplace=False): if not environment_remappings: return pipeline_proto if not inplace: pipeline_proto = copy.copy(pipeline_proto) for t in pipeline_proto.components.transforms.values(): if t.environment_id not in pipeline_proto.components.environments: # TODO(https://github.com/apache/beam/issues/30876): Remove this # workaround. continue if t.environment_id and t.environment_id in environment_remappings: t.environment_id = environment_remappings[t.environment_id] for w in pipeline_proto.components.windowing_strategies.values(): if w.environment_id not in pipeline_proto.components.environments: # TODO(https://github.com/apache/beam/issues/30876): Remove this # workaround. continue if w.environment_id and w.environment_id in environment_remappings: w.environment_id = environment_remappings[w.environment_id] for e in set(environment_remappings.keys()) - set( environment_remappings.values()): del pipeline_proto.components.environments[e] return pipeline_proto ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/pipeline_utils_test.py0000644000175100017510000001773315206113032024022 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import hashlib import os import random import tempfile import unittest import apache_beam as beam from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.pipeline_utils import merge_common_environments from apache_beam.runners.pipeline_utils import merge_superset_dep_environments from apache_beam.runners.portability.expansion_service_test import FibTransform class PipelineUtilitiesTest(unittest.TestCase): def test_equal_environments_merged(self): pipeline_proto = merge_common_environments( beam_runner_api_pb2.Pipeline( components=beam_runner_api_pb2.Components( environments={ 'a1': beam_runner_api_pb2.Environment(urn='A'), 'a2': beam_runner_api_pb2.Environment(urn='A'), 'b1': beam_runner_api_pb2.Environment( urn='B', payload=b'x'), 'b2': beam_runner_api_pb2.Environment( urn='B', payload=b'x'), 'b3': beam_runner_api_pb2.Environment( urn='B', payload=b'y'), }, transforms={ 't1': beam_runner_api_pb2.PTransform( unique_name='t1', environment_id='a1'), 't2': beam_runner_api_pb2.PTransform( unique_name='t2', environment_id='a2'), }, windowing_strategies={ 'w1': beam_runner_api_pb2.WindowingStrategy( environment_id='b1'), 'w2': beam_runner_api_pb2.WindowingStrategy( environment_id='b2'), }))) self.assertEqual(len(pipeline_proto.components.environments), 3) self.assertTrue(('a1' in pipeline_proto.components.environments) ^ ('a2' in pipeline_proto.components.environments)) self.assertTrue(('b1' in pipeline_proto.components.environments) ^ ('b2' in pipeline_proto.components.environments)) self.assertEqual( len( set( t.environment_id for t in pipeline_proto.components.transforms.values())), 1) self.assertEqual( len( set( w.environment_id for w in pipeline_proto.components.windowing_strategies.values())), 1) def _make_dep(self, path): hasher = hashlib.sha256() if os.path.exists(path): with open(path, 'rb') as fin: hasher.update(fin.read()) else: # A fake file, identified only by its path. hasher.update(path.encode('utf-8')) return beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2.ArtifactFilePayload( path=path, sha256=hasher.hexdigest()).SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2.ArtifactStagingToRolePayload( staged_name=os.path.basename(path)).SerializeToString()) def _docker_env(self, id, deps=()): return beam_runner_api_pb2.Environment( urn=common_urns.environments.DOCKER.urn, payload=id.encode('utf8'), dependencies=[self._make_dep(path) for path in deps], ) def test_subset_deps_environments_merged(self): environments = { 'A': self._docker_env('A'), 'Ax': self._docker_env('A', ['x']), 'Ay': self._docker_env('A', ['y']), 'Axy': self._docker_env('A', ['x', 'y']), 'Bx': self._docker_env('B', ['x']), 'Bxy': self._docker_env('B', ['x', 'y']), 'Byz': self._docker_env('B', ['y', 'z']), } transforms = { env_id: beam_runner_api_pb2.PTransform( unique_name=env_id, environment_id=env_id) for env_id in environments.keys() } pipeline_proto = merge_superset_dep_environments( beam_runner_api_pb2.Pipeline( components=beam_runner_api_pb2.Components( environments=environments, transforms=transforms))) # These can all be merged into the same environment. self.assertEqual( pipeline_proto.components.transforms['A'].environment_id, 'Axy') self.assertEqual( pipeline_proto.components.transforms['Ax'].environment_id, 'Axy') self.assertEqual( pipeline_proto.components.transforms['Ay'].environment_id, 'Axy') self.assertEqual( pipeline_proto.components.transforms['Axy'].environment_id, 'Axy') # Despite having the same dependencies, these must be merged into their own. self.assertEqual( pipeline_proto.components.transforms['Bx'].environment_id, 'Bxy') self.assertEqual( pipeline_proto.components.transforms['Bxy'].environment_id, 'Bxy') # This is not a subset of any, must be left alone. self.assertEqual( pipeline_proto.components.transforms['Byz'].environment_id, 'Byz') def test_subset_deps_environments_merged_with_requirements_txt(self): with tempfile.TemporaryDirectory() as tmpdir: def make_file(basename, content): subdir = tempfile.mkdtemp(dir=tmpdir) path = os.path.join(subdir, basename) with open(path, 'w') as fout: fout.write(content) return path def make_py_deps(*pkgs): return [ make_file('requirements.txt', '\n'.join(pkgs)), make_file( 'submission_environment_dependencies.txt', str( random.random())), ] + [make_file(pkg, pkg) for pkg in pkgs] environments = { 'A': self._docker_env('A'), 'Ax': self._docker_env('A', make_py_deps('x')), 'Ay': self._docker_env('A', make_py_deps('y')), 'Axy': self._docker_env('A', make_py_deps('x', 'y')), } transforms = { env_id: beam_runner_api_pb2.PTransform( unique_name=env_id, environment_id=env_id) for env_id in environments.keys() } pipeline_proto = merge_superset_dep_environments( beam_runner_api_pb2.Pipeline( components=beam_runner_api_pb2.Components( environments=environments, transforms=transforms))) # These can all be merged into the same environment. self.assertEqual( pipeline_proto.components.transforms['A'].environment_id, 'Axy') self.assertEqual( pipeline_proto.components.transforms['Ax'].environment_id, 'Axy') self.assertEqual( pipeline_proto.components.transforms['Ay'].environment_id, 'Axy') self.assertEqual( pipeline_proto.components.transforms['Axy'].environment_id, 'Axy') def test_external_merged(self): p = beam.Pipeline() # This transform recursively creates several external environments. _ = p | FibTransform(4) pipeline_proto = p.to_runner_api() # All our external environments are equal and consolidated. # We also have a placeholder "default" environment that has not been # resolved do anything concrete yet. envs = pipeline_proto.components.environments self.assertEqual( len(envs), 2, f'should be 2 environments, instead got: {envs}') if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.676439 apache_beam-2.74.0/apache_beam/runners/portability/0000755000175100017510000000000015206113160021715 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/__init__.py0000644000175100017510000000153315206113032024026 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """This runner is experimental; no backwards-compatibility guarantees.""" ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/abstract_job_service.py0000644000175100017510000003171615206113032026452 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import copy import itertools import json import logging import shutil import tempfile import uuid import zipfile from concurrent import futures from typing import BinaryIO from typing import Iterator from typing import Optional from typing import Union import grpc from google.protobuf import json_format from google.protobuf import struct_pb2 from google.protobuf import timestamp_pb2 from apache_beam.portability.api import beam_artifact_api_pb2_grpc from apache_beam.portability.api import beam_job_api_pb2 from apache_beam.portability.api import beam_job_api_pb2_grpc from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.portability.api import endpoints_pb2 from apache_beam.runners.portability import artifact_service from apache_beam.utils.timestamp import Timestamp _LOGGER = logging.getLogger(__name__) StateEvent = tuple[int, Union[timestamp_pb2.Timestamp, Timestamp]] def make_state_event(state, timestamp): if isinstance(timestamp, Timestamp): proto_timestamp = timestamp.to_proto() elif isinstance(timestamp, timestamp_pb2.Timestamp): proto_timestamp = timestamp else: raise ValueError( "Expected apache_beam.utils.timestamp.Timestamp, " "or google.protobuf.timestamp_pb2.Timestamp. " "Got %s" % type(timestamp)) return beam_job_api_pb2.JobStateEvent(state=state, timestamp=proto_timestamp) class AbstractJobServiceServicer(beam_job_api_pb2_grpc.JobServiceServicer): """Manages one or more pipelines, possibly concurrently. Experimental: No backward compatibility guaranteed. Servicer for the Beam Job API. """ def __init__(self): self._jobs: dict[str, AbstractBeamJob] = {} def create_beam_job( self, preparation_id, # stype: str job_name: str, pipeline: beam_runner_api_pb2.Pipeline, options: struct_pb2.Struct) -> 'AbstractBeamJob': """Returns an instance of AbstractBeamJob specific to this servicer.""" raise NotImplementedError(type(self)) def Prepare( self, request: beam_job_api_pb2.PrepareJobRequest, context=None, timeout=None) -> beam_job_api_pb2.PrepareJobResponse: _LOGGER.debug('Got Prepare request.') preparation_id = '%s-%s' % (request.job_name, uuid.uuid4()) self._jobs[preparation_id] = self.create_beam_job( preparation_id, request.job_name, request.pipeline, request.pipeline_options) self._jobs[preparation_id].prepare() _LOGGER.debug("Prepared job '%s' as '%s'", request.job_name, preparation_id) return beam_job_api_pb2.PrepareJobResponse( preparation_id=preparation_id, artifact_staging_endpoint=self._jobs[preparation_id]. artifact_staging_endpoint(), staging_session_token=preparation_id) def Run( self, request: beam_job_api_pb2.RunJobRequest, context=None, timeout=None) -> beam_job_api_pb2.RunJobResponse: # For now, just use the preparation id as the job id. job_id = request.preparation_id _LOGGER.info("Running job '%s'", job_id) self._jobs[job_id].run() return beam_job_api_pb2.RunJobResponse(job_id=job_id) def GetJobs( self, request: beam_job_api_pb2.GetJobsRequest, context=None, timeout=None) -> beam_job_api_pb2.GetJobsResponse: return beam_job_api_pb2.GetJobsResponse( job_info=[job.to_runner_api() for job in self._jobs.values()]) def GetState( self, request: beam_job_api_pb2.GetJobStateRequest, context=None) -> beam_job_api_pb2.JobStateEvent: return make_state_event(*self._jobs[request.job_id].get_state()) def GetPipeline( self, request: beam_job_api_pb2.GetJobPipelineRequest, context=None, timeout=None) -> beam_job_api_pb2.GetJobPipelineResponse: return beam_job_api_pb2.GetJobPipelineResponse( pipeline=self._jobs[request.job_id].get_pipeline()) def Cancel( self, request: beam_job_api_pb2.CancelJobRequest, context=None, timeout=None) -> beam_job_api_pb2.CancelJobResponse: self._jobs[request.job_id].cancel() return beam_job_api_pb2.CancelJobResponse( state=self._jobs[request.job_id].get_state()[0]) def GetStateStream(self, request, context=None, timeout=None) -> Iterator[beam_job_api_pb2.JobStateEvent]: """Yields state transitions since the stream started. """ if request.job_id not in self._jobs: raise LookupError("Job {} does not exist".format(request.job_id)) job = self._jobs[request.job_id] for state, timestamp in job.get_state_stream(): yield make_state_event(state, timestamp) def GetMessageStream( self, request, context=None, timeout=None) -> Iterator[beam_job_api_pb2.JobMessagesResponse]: """Yields messages since the stream started. """ if request.job_id not in self._jobs: raise LookupError("Job {} does not exist".format(request.job_id)) job = self._jobs[request.job_id] for msg in job.get_message_stream(): if isinstance(msg, tuple): resp = beam_job_api_pb2.JobMessagesResponse( state_response=make_state_event(*msg)) else: resp = beam_job_api_pb2.JobMessagesResponse(message_response=msg) yield resp def DescribePipelineOptions( self, request, context=None, timeout=None) -> beam_job_api_pb2.DescribePipelineOptionsResponse: return beam_job_api_pb2.DescribePipelineOptionsResponse() class AbstractBeamJob(object): """Abstract baseclass for managing a single Beam job.""" def __init__( self, job_id: str, job_name: str, pipeline: beam_runner_api_pb2.Pipeline, options: struct_pb2.Struct): self._job_id = job_id self._job_name = job_name self._pipeline_proto = pipeline self._pipeline_options = options self._state_history = [(beam_job_api_pb2.JobState.STOPPED, Timestamp.now())] def prepare(self) -> None: """Called immediately after this class is instantiated""" raise NotImplementedError(self) def run(self) -> None: raise NotImplementedError(self) def cancel(self) -> Optional['beam_job_api_pb2.JobState.Enum']: raise NotImplementedError(self) def artifact_staging_endpoint( self) -> Optional[endpoints_pb2.ApiServiceDescriptor]: raise NotImplementedError(self) def get_state_stream(self) -> Iterator[StateEvent]: raise NotImplementedError(self) def get_message_stream( self ) -> Iterator[Union[StateEvent, Optional[beam_job_api_pb2.JobMessage]]]: raise NotImplementedError(self) @property def state(self): """Get the latest state enum.""" return self.get_state()[0] def get_state(self): """Get a tuple of the latest state and its timestamp.""" # this is safe: initial state is set in __init__ return self._state_history[-1] def set_state(self, new_state): """Set the latest state as an int enum and update the state history. :param new_state: int latest state enum :return: Timestamp or None the new timestamp if the state has not changed, else None """ if new_state != self._state_history[-1][0]: timestamp = Timestamp.now() self._state_history.append((new_state, timestamp)) return timestamp else: return None def with_state_history(self, state_stream): """Utility to prepend recorded state history to an active state stream""" return itertools.chain(self._state_history[:], state_stream) def get_pipeline(self) -> beam_runner_api_pb2.Pipeline: return self._pipeline_proto @staticmethod def is_terminal_state(state): from apache_beam.runners.portability import portable_runner return state in portable_runner.TERMINAL_STATES def to_runner_api(self) -> beam_job_api_pb2.JobInfo: return beam_job_api_pb2.JobInfo( job_id=self._job_id, job_name=self._job_name, pipeline_options=self._pipeline_options, state=self.state) class JarArtifactManager(object): def __init__(self, jar_path, root): self._root = root self._zipfile_handle = zipfile.ZipFile(jar_path, 'a') def close(self): self._zipfile_handle.close() def file_writer(self, path: str) -> tuple[BinaryIO, str]: """Given a relative path, returns an open handle that can be written to and an reference that can later be used to read this file.""" full_path = '%s/%s' % (self._root, path) return self._zipfile_handle.open( full_path, 'w', force_zip64=True), 'classpath://%s' % full_path def zipfile_handle(self): return self._zipfile_handle class UberJarBeamJob(AbstractBeamJob): """Abstract baseclass for creating a Beam job. The resulting job will be packaged and run in an executable uber jar.""" # These must agree with those defined in PortablePipelineJarUtils.java. PIPELINE_FOLDER = 'BEAM-PIPELINE' PIPELINE_MANIFEST = PIPELINE_FOLDER + '/pipeline-manifest.json' # We only stage a single pipeline in the jar. PIPELINE_NAME = 'pipeline' PIPELINE_PATH = '/'.join([PIPELINE_FOLDER, PIPELINE_NAME, "pipeline.json"]) PIPELINE_OPTIONS_PATH = '/'.join( [PIPELINE_FOLDER, PIPELINE_NAME, 'pipeline-options.json']) ARTIFACT_FOLDER = '/'.join([PIPELINE_FOLDER, PIPELINE_NAME, 'artifacts']) def __init__( self, executable_jar, job_id, job_name, pipeline, options, artifact_port=0): super().__init__(job_id, job_name, pipeline, options) self._executable_jar = executable_jar self._jar_uploaded = False self._artifact_port = artifact_port def prepare(self): # Copy the executable jar, injecting the pipeline and options as resources. with tempfile.NamedTemporaryFile(suffix='.jar') as tout: self._jar = tout.name shutil.copy(self._executable_jar, self._jar) self._start_artifact_service(self._jar, self._artifact_port) def _start_artifact_service(self, jar, requested_port): self._artifact_manager = JarArtifactManager(self._jar, self.ARTIFACT_FOLDER) self._artifact_staging_service = artifact_service.ArtifactStagingService( self._artifact_manager.file_writer) self._artifact_staging_service.register_job( self._job_id, { env_id: env.dependencies for (env_id, env) in self._pipeline_proto.components.environments.items() }) options = [("grpc.http2.max_pings_without_data", 0), ("grpc.http2.max_ping_strikes", 0)] self._artifact_staging_server = grpc.server( futures.ThreadPoolExecutor(), options=options) port = self._artifact_staging_server.add_insecure_port( '[::]:%s' % requested_port) beam_artifact_api_pb2_grpc.add_ArtifactStagingServiceServicer_to_server( self._artifact_staging_service, self._artifact_staging_server) self._artifact_staging_endpoint = endpoints_pb2.ApiServiceDescriptor( url='localhost:%d' % port) self._artifact_staging_server.start() _LOGGER.info('Artifact server started on port %s', port) return port def _stop_artifact_service(self): self._artifact_staging_server.stop(1) # Update dependencies to point to staged files. pipeline = copy.copy(self._pipeline_proto) if any(env.dependencies for env in pipeline.components.environments.values()): for env_id, deps in self._artifact_staging_service.resolved_deps( self._job_id).items(): # Slice assignment not supported for repeated fields. env = self._pipeline_proto.components.environments[env_id] del env.dependencies[:] env.dependencies.extend(deps) # Copy the pipeline definition and metadata into the jar. z = self._artifact_manager.zipfile_handle() with z.open(self.PIPELINE_PATH, 'w') as fout: fout.write( json_format.MessageToJson(self._pipeline_proto).encode('utf-8')) with z.open(self.PIPELINE_OPTIONS_PATH, 'w') as fout: fout.write( json_format.MessageToJson(self._pipeline_options).encode('utf-8')) with z.open(self.PIPELINE_MANIFEST, 'w') as fout: fout.write( json.dumps({ 'defaultJobName': self.PIPELINE_NAME }).encode('utf-8')) # Closes the jar file. self._artifact_manager.close() def artifact_staging_endpoint(self): return self._artifact_staging_endpoint ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/artifact_service.py0000644000175100017510000002726315206113032025614 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Implementation of an Artifact{Staging,Retrieval}Service. The staging service here can be backed by any beam filesystem. """ # pytype: skip-file import concurrent.futures import hashlib import os import queue import sys import tempfile import threading from io import BytesIO from typing import Any from typing import BinaryIO # pylint: disable=unused-import from typing import Callable from typing import MutableMapping from typing import Optional from urllib.request import urlopen import grpc from apache_beam.io import filesystems from apache_beam.io.filesystems import CompressionTypes from apache_beam.portability import common_urns from apache_beam.portability.api import beam_artifact_api_pb2 from apache_beam.portability.api import beam_artifact_api_pb2_grpc from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.utils import proto_utils class ArtifactRetrievalService( beam_artifact_api_pb2_grpc.ArtifactRetrievalServiceServicer): _DEFAULT_CHUNK_SIZE = 2 << 20 def __init__( self, file_reader: Callable[[str], BinaryIO], chunk_size=None, ): self._file_reader = file_reader self._chunk_size = chunk_size or self._DEFAULT_CHUNK_SIZE def ResolveArtifacts(self, request, context=None): return beam_artifact_api_pb2.ResolveArtifactsResponse( replacements=request.artifacts) def GetArtifact(self, request, context=None): if request.artifact.type_urn == common_urns.artifact_types.FILE.urn: payload = proto_utils.parse_Bytes( request.artifact.type_payload, beam_runner_api_pb2.ArtifactFilePayload) read_handle = self._file_reader(payload.path) elif request.artifact.type_urn == common_urns.artifact_types.URL.urn: payload = proto_utils.parse_Bytes( request.artifact.type_payload, beam_runner_api_pb2.ArtifactUrlPayload) read_handle = urlopen(payload.url) elif request.artifact.type_urn == common_urns.artifact_types.EMBEDDED.urn: payload = proto_utils.parse_Bytes( request.artifact.type_payload, beam_runner_api_pb2.EmbeddedFilePayload) read_handle = BytesIO(payload.data) else: raise NotImplementedError(request.artifact.type_urn) with read_handle as fin: while True: chunk = fin.read(self._chunk_size) if not chunk: break yield beam_artifact_api_pb2.GetArtifactResponse(data=chunk) class ArtifactStagingService( beam_artifact_api_pb2_grpc.ArtifactStagingServiceServicer): def __init__( self, file_writer: Callable[[str, Optional[str]], tuple[BinaryIO, str]], ): self._lock = threading.Lock() self._jobs_to_stage: dict[ str, tuple[dict[Any, list[beam_runner_api_pb2.ArtifactInformation]], threading.Event]] = {} self._file_writer = file_writer def register_job( self, staging_token: str, dependency_sets: MutableMapping[ Any, list[beam_runner_api_pb2.ArtifactInformation]]): if staging_token in self._jobs_to_stage: raise ValueError('Already staging %s' % staging_token) with self._lock: self._jobs_to_stage[staging_token] = ( dict(dependency_sets), threading.Event()) def resolved_deps(self, staging_token, timeout=None): with self._lock: dependency_sets, event = self._jobs_to_stage[staging_token] try: if not event.wait(timeout): raise concurrent.futures.TimeoutError() return dependency_sets finally: with self._lock: del self._jobs_to_stage[staging_token] def ReverseArtifactRetrievalService(self, responses, context=None): staging_token = next(responses).staging_token with self._lock: try: dependency_sets, event = self._jobs_to_stage[staging_token] except KeyError: if context: context.set_code(grpc.StatusCode.NOT_FOUND) context.set_details('No such staging token: %r' % staging_token) raise requests = _QueueIter() class ForwardingRetrievalService(object): def ResolveArtifactss(self, request): requests.put( beam_artifact_api_pb2.ArtifactRequestWrapper( resolve_artifact=request)) return next(responses).resolve_artifact_response def GetArtifact(self, request): requests.put( beam_artifact_api_pb2.ArtifactRequestWrapper(get_artifact=request)) while True: response = next(responses) yield response.get_artifact_response if response.is_last: break def resolve(): try: for key, dependencies in dependency_sets.items(): dependency_sets[key] = list( resolve_as_files( ForwardingRetrievalService(), lambda name: self._file_writer( os.path.join(staging_token, name)), dependencies)) requests.done() except: # pylint: disable=bare-except requests.abort() raise finally: event.set() t = threading.Thread(target=resolve) t.daemon = True t.start() return requests def resolve_as_files(retrieval_service, file_writer, dependencies): """Translates a set of dependencies into file-based dependencies.""" # Resolve until nothing changes. This ensures that they can be fetched. resolution = retrieval_service.ResolveArtifactss( beam_artifact_api_pb2.ResolveArtifactsRequest( artifacts=dependencies, # Anything fetchable will do. # TODO(robertwb): Take advantage of shared filesystems, urls. preferred_urns=[], )) dependencies = resolution.replacements # Fetch each of the dependencies, using file_writer to store them as # file-based artifacts. # TODO(robertwb): Consider parallelizing the actual writes. for dep in dependencies: if dep.role_urn == common_urns.artifact_roles.STAGING_TO.urn: base_name = os.path.basename( proto_utils.parse_Bytes( dep.role_payload, beam_runner_api_pb2.ArtifactStagingToRolePayload).staged_name) else: base_name = None unique_name = '-'.join( filter( None, [hashlib.sha256(dep.SerializeToString()).hexdigest(), base_name])) file_handle, path = file_writer(unique_name) with file_handle as fout: for chunk in retrieval_service.GetArtifact( beam_artifact_api_pb2.GetArtifactRequest(artifact=dep)): fout.write(chunk.data) yield beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2.ArtifactFilePayload( path=path).SerializeToString(), role_urn=dep.role_urn, role_payload=dep.role_payload) def offer_artifacts( artifact_staging_service, artifact_retrieval_service, staging_token): """Offers a set of artifacts to an artifact staging service, via the ReverseArtifactRetrievalService API. The given artifact_retrieval_service should be able to resolve/get all artifacts relevant to this job. """ responses = _QueueIter() responses.put( beam_artifact_api_pb2.ArtifactResponseWrapper( staging_token=staging_token)) requests = artifact_staging_service.ReverseArtifactRetrievalService(responses) try: for request in requests: if request.HasField('resolve_artifact'): responses.put( beam_artifact_api_pb2.ArtifactResponseWrapper( resolve_artifact_response=artifact_retrieval_service. ResolveArtifacts(request.resolve_artifact))) elif request.HasField('get_artifact'): for chunk in artifact_retrieval_service.GetArtifact( request.get_artifact): responses.put( beam_artifact_api_pb2.ArtifactResponseWrapper( get_artifact_response=chunk)) responses.put( beam_artifact_api_pb2.ArtifactResponseWrapper( get_artifact_response=beam_artifact_api_pb2.GetArtifactResponse( data=b''), is_last=True)) responses.done() except: # pylint: disable=bare-except responses.abort() raise class BeamFilesystemHandler(object): def __init__(self, root): self._root = root def file_reader(self, path): return filesystems.FileSystems.open( path, compression_type=CompressionTypes.UNCOMPRESSED) def file_writer(self, name=None): full_path = filesystems.FileSystems.join(self._root, name) return filesystems.FileSystems.create( full_path, compression_type=CompressionTypes.UNCOMPRESSED), full_path def resolve_artifacts(artifacts, service, dest_dir): if not artifacts: return artifacts else: return [ maybe_store_artifact(artifact, service, dest_dir) for artifact in service.ResolveArtifacts( beam_artifact_api_pb2.ResolveArtifactsRequest( artifacts=artifacts)).replacements ] def maybe_store_artifact(artifact, service, dest_dir): if artifact.type_urn in (common_urns.artifact_types.URL.urn, common_urns.artifact_types.EMBEDDED.urn): return artifact elif artifact.type_urn == common_urns.artifact_types.FILE.urn: payload = beam_runner_api_pb2.ArtifactFilePayload.FromString( artifact.type_payload) # pylint: disable=condition-evals-to-constant if os.path.exists( payload.path) and payload.sha256 and payload.sha256 == sha256( payload.path) and False: return artifact else: return store_artifact(artifact, service, dest_dir) else: return store_artifact(artifact, service, dest_dir) def store_artifact(artifact, service, dest_dir): hasher = hashlib.sha256() with tempfile.NamedTemporaryFile(dir=dest_dir, delete=False) as fout: for block in service.GetArtifact( beam_artifact_api_pb2.GetArtifactRequest(artifact=artifact)): hasher.update(block.data) fout.write(block.data) return beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2.ArtifactFilePayload( path=fout.name, sha256=hasher.hexdigest()).SerializeToString(), role_urn=artifact.role_urn, role_payload=artifact.role_payload) def sha256(path): hasher = hashlib.sha256() with open(path, 'rb') as fin: for block in iter(lambda: fin.read(4 << 20), b''): hasher.update(block) return hasher.hexdigest() class _QueueIter(object): _END = object() def __init__(self): self._queue = queue.Queue() def put(self, item): self._queue.put(item) def done(self): self._queue.put(self._END) self._queue.put(StopIteration) def abort(self, exn=None): if exn is None: exn = sys.exc_info()[1] self._queue.put(self._END) self._queue.put(exn) def __iter__(self): return self def __next__(self): item = self._queue.get() if item is self._END: raise self._queue.get() else: return item ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/artifact_service_test.py0000644000175100017510000001613715206113032026651 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test cases for :module:`artifact_service_client`.""" # pytype: skip-file import contextlib import io import threading import unittest from urllib.parse import quote from apache_beam.portability import common_urns from apache_beam.portability.api import beam_artifact_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.portability import artifact_service from apache_beam.utils import proto_utils class InMemoryFileManager(object): def __init__(self, contents=()): self._contents = dict(contents) def get(self, path): return self._contents[path] def file_reader(self, path): return io.BytesIO(self._contents[path]) def file_writer(self, name): path = 'prefix:' + name @contextlib.contextmanager def writable(): buffer = io.BytesIO() yield buffer buffer.seek(0) self._contents[path] = buffer.read() return writable(), path class ArtifactServiceTest(unittest.TestCase): def file_artifact(self, path): return beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2.ArtifactFilePayload( path=path).SerializeToString()) def embedded_artifact(self, data, name=None): return beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.EMBEDDED.urn, type_payload=beam_runner_api_pb2.EmbeddedFilePayload( data=data).SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn if name else None, role_payload=beam_runner_api_pb2.ArtifactStagingToRolePayload( staged_name=name).SerializeToString() if name else None) def test_file_retrieval(self): file_manager = InMemoryFileManager({ 'path/to/a': b'a', 'path/to/b': b'b' * 37 }) retrieval_service = artifact_service.ArtifactRetrievalService( file_manager.file_reader, chunk_size=10) dep_a = self.file_artifact('path/to/a') self.assertEqual( retrieval_service.ResolveArtifacts( beam_artifact_api_pb2.ResolveArtifactsRequest(artifacts=[dep_a])), beam_artifact_api_pb2.ResolveArtifactsResponse(replacements=[dep_a])) self.assertEqual( list( retrieval_service.GetArtifact( beam_artifact_api_pb2.GetArtifactRequest(artifact=dep_a))), [beam_artifact_api_pb2.GetArtifactResponse(data=b'a')]) self.assertEqual( list( retrieval_service.GetArtifact( beam_artifact_api_pb2.GetArtifactRequest( artifact=self.file_artifact('path/to/b')))), [ beam_artifact_api_pb2.GetArtifactResponse(data=b'b' * 10), beam_artifact_api_pb2.GetArtifactResponse(data=b'b' * 10), beam_artifact_api_pb2.GetArtifactResponse(data=b'b' * 10), beam_artifact_api_pb2.GetArtifactResponse(data=b'b' * 7) ]) def test_embedded_retrieval(self): retrieval_service = artifact_service.ArtifactRetrievalService(None) embedded_dep = self.embedded_artifact(b'some_data') self.assertEqual( list( retrieval_service.GetArtifact( beam_artifact_api_pb2.GetArtifactRequest( artifact=embedded_dep))), [beam_artifact_api_pb2.GetArtifactResponse(data=b'some_data')]) def test_url_retrieval(self): retrieval_service = artifact_service.ArtifactRetrievalService(None) url_dep = beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.URL.urn, type_payload=beam_runner_api_pb2.ArtifactUrlPayload( url='file:' + quote(__file__)).SerializeToString()) content = b''.join([ r.data for r in retrieval_service.GetArtifact( beam_artifact_api_pb2.GetArtifactRequest(artifact=url_dep)) ]) with open(__file__, 'rb') as fin: self.assertEqual(content, fin.read()) def test_push_artifacts(self): unresolved = beam_runner_api_pb2.ArtifactInformation(type_urn='unresolved') resolved_a = self.embedded_artifact(data=b'a', name='a.txt') resolved_b = self.embedded_artifact(data=b'bb', name='b.txt') dep_big = self.embedded_artifact(data=b'big ' * 100, name='big.txt') class TestArtifacts(object): def ResolveArtifacts(self, request): replacements = [] for artifact in request.artifacts: if artifact.type_urn == 'unresolved': replacements += [resolved_a, resolved_b] else: replacements.append(artifact) return beam_artifact_api_pb2.ResolveArtifactsResponse( replacements=replacements) def GetArtifact(self, request): if request.artifact.type_urn == common_urns.artifact_types.EMBEDDED.urn: content = proto_utils.parse_Bytes( request.artifact.type_payload, beam_runner_api_pb2.EmbeddedFilePayload).data for k in range(0, len(content), 13): yield beam_artifact_api_pb2.GetArtifactResponse( data=content[k:k + 13]) else: raise NotImplementedError file_manager = InMemoryFileManager() server = artifact_service.ArtifactStagingService(file_manager.file_writer) server.register_job('staging_token', {'env': [unresolved, dep_big]}) # "Push" artifacts as if from a client. t = threading.Thread( target=lambda: artifact_service.offer_artifacts( server, TestArtifacts(), 'staging_token')) t.daemon = True t.start() resolved_deps = server.resolved_deps('staging_token', timeout=5)['env'] expected = { 'a.txt': b'a', 'b.txt': b'bb', 'big.txt': b'big ' * 100, } for dep in resolved_deps: self.assertEqual(dep.type_urn, common_urns.artifact_types.FILE.urn) self.assertEqual(dep.role_urn, common_urns.artifact_roles.STAGING_TO.urn) type_payload = proto_utils.parse_Bytes( dep.type_payload, beam_runner_api_pb2.ArtifactFilePayload) role_payload = proto_utils.parse_Bytes( dep.role_payload, beam_runner_api_pb2.ArtifactStagingToRolePayload) self.assertTrue( type_payload.path.endswith(role_payload.staged_name), type_payload.path) self.assertEqual( file_manager.get(type_payload.path), expected.pop(role_payload.staged_name)) self.assertEqual(expected, {}) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/expansion_service.py0000644000175100017510000001310615206113032026012 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A PipelineExpansion service. """ # pytype: skip-file import copy import traceback from apache_beam import pipeline as beam_pipeline from apache_beam.options import pipeline_options from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.portability.api import beam_expansion_api_pb2 from apache_beam.portability.api import beam_expansion_api_pb2_grpc from apache_beam.runners import pipeline_context from apache_beam.runners.portability import artifact_service from apache_beam.runners.portability.artifact_service import BeamFilesystemHandler from apache_beam.transforms import environments from apache_beam.transforms import external from apache_beam.transforms import ptransform class ExpansionServiceServicer( beam_expansion_api_pb2_grpc.ExpansionServiceServicer): def __init__(self, options=None, loopback_address=None): self._options = options or beam_pipeline.PipelineOptions( flags=[], environment_type=python_urns.EMBEDDED_PYTHON, sdk_location='container') default_environment = (environments.Environment.from_options(self._options)) if loopback_address: loopback_environment = environments.Environment.from_options( beam_pipeline.PipelineOptions( environment_type=common_urns.environments.EXTERNAL.urn, environment_config=loopback_address)) default_environment = environments.AnyOfEnvironment( [default_environment, loopback_environment]) self._default_environment = default_environment def Expand(self, request, context=None): try: options = copy.deepcopy(self._options) options = pipeline_options.PipelineOptions.from_runner_api( request.pipeline_options, options) pipeline = beam_pipeline.Pipeline(options=options) def with_pipeline(component, pcoll_id=None): component.pipeline = pipeline if pcoll_id: component.producer, component.tag = producers[pcoll_id] # We need the lookup to resolve back to this id. context.pcollections._obj_to_id[component] = pcoll_id return component context = pipeline_context.PipelineContext( request.components, default_environment=self._default_environment, namespace=request.namespace, requirements=request.requirements) producers = { pcoll_id: (context.transforms.get_by_id(t_id), pcoll_tag) for t_id, t_proto in request.components.transforms.items() for pcoll_tag, pcoll_id in t_proto.outputs.items() } transform = with_pipeline( ptransform.PTransform.from_runner_api(request.transform, context)) if len(request.output_coder_requests) == 1: output_coder = { k: context.element_type_from_coder_id(v) for k, v in request.output_coder_requests.items() } transform = transform.with_output_types(list(output_coder.values())[0]) elif len(request.output_coder_requests) > 1: raise ValueError( 'type annotation for multiple outputs is not allowed yet: %s' % request.output_coder_requests) inputs = transform._pvaluish_from_dict({ tag: with_pipeline( context.pcollections.get_by_id(pcoll_id), pcoll_id) for tag, pcoll_id in request.transform.inputs.items() }) if not inputs: inputs = pipeline with external.ExternalTransform.outer_namespace(request.namespace): result = pipeline.apply( transform, inputs, request.transform.unique_name) expanded_transform = pipeline._root_transform().parts[-1] # TODO(BEAM-1833): Use named outputs internally. if isinstance(result, dict): expanded_transform.outputs = result pipeline_proto = pipeline.to_runner_api(context=context) # TODO(BEAM-1833): Use named inputs internally. expanded_transform_id = context.transforms.get_id(expanded_transform) expanded_transform_proto = pipeline_proto.components.transforms.pop( expanded_transform_id) expanded_transform_proto.inputs.clear() expanded_transform_proto.inputs.update(request.transform.inputs) for transform_id in pipeline_proto.root_transform_ids: del pipeline_proto.components.transforms[transform_id] return beam_expansion_api_pb2.ExpansionResponse( components=pipeline_proto.components, transform=expanded_transform_proto, requirements=pipeline_proto.requirements) except Exception: # pylint: disable=broad-except return beam_expansion_api_pb2.ExpansionResponse( error=traceback.format_exc()) def artifact_service(self): """Returns a service to retrieve artifacts for use in a job.""" return artifact_service.ArtifactRetrievalService( BeamFilesystemHandler(None).file_reader) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/expansion_service_main.py0000644000175100017510000000711015206113032027014 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import argparse import logging import signal import sys import grpc from apache_beam.internal import pickler from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.portability.api import beam_artifact_api_pb2_grpc from apache_beam.portability.api import beam_expansion_api_pb2_grpc from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.runners.portability import artifact_service from apache_beam.runners.portability import expansion_service from apache_beam.runners.worker import worker_pool_main from apache_beam.transforms import fully_qualified_named_transform from apache_beam.utils import thread_pool_executor _LOGGER = logging.getLogger(__name__) def main(argv): parser = argparse.ArgumentParser() parser.add_argument( '-p', '--port', type=int, help='port on which to serve the job api') parser.add_argument('--fully_qualified_name_glob', default=None) parser.add_argument('--serve_loopback_worker', action='store_true') known_args, pipeline_args = parser.parse_known_args(argv) pipeline_options = PipelineOptions( pipeline_args + ["--experiments=beam_fn_api", "--sdk_location=container"]) # Set this before any pipeline construction occurs. # See https://github.com/apache/beam/issues/21615 pickler.set_library(pipeline_options.view_as(SetupOptions).pickle_library) with fully_qualified_named_transform.FullyQualifiedNamedTransform.with_filter( known_args.fully_qualified_name_glob): address = '0.0.0.0:{}'.format(known_args.port) server = grpc.server(thread_pool_executor.shared_unbounded_instance()) if known_args.serve_loopback_worker: beam_fn_api_pb2_grpc.add_BeamFnExternalWorkerPoolServicer_to_server( worker_pool_main.BeamFnExternalWorkerPoolServicer(), server) loopback_address = address else: loopback_address = None beam_expansion_api_pb2_grpc.add_ExpansionServiceServicer_to_server( expansion_service.ExpansionServiceServicer( pipeline_options, loopback_address=loopback_address), server) beam_artifact_api_pb2_grpc.add_ArtifactRetrievalServiceServicer_to_server( artifact_service.ArtifactRetrievalService( artifact_service.BeamFilesystemHandler(None).file_reader), server) server.add_insecure_port(address) server.start() _LOGGER.info('Listening for expansion requests at %d', known_args.port) def cleanup(unused_signum, unused_frame): _LOGGER.info('Shutting down expansion service.') server.stop(None) signal.signal(signal.SIGTERM, cleanup) signal.signal(signal.SIGINT, cleanup) # blocking main thread forever. signal.pause() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) main(sys.argv) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/expansion_service_test.py0000644000175100017510000003320715206113032027055 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import argparse import logging import pickle import signal import sys import typing import grpc import apache_beam as beam import apache_beam.transforms.combiners as combine from apache_beam.coders import RowCoder from apache_beam.pipeline import PipelineOptions from apache_beam.portability.api import beam_artifact_api_pb2_grpc from apache_beam.portability.api import beam_expansion_api_pb2_grpc from apache_beam.portability.api import external_transforms_pb2 from apache_beam.runners.portability import artifact_service from apache_beam.runners.portability import expansion_service from apache_beam.runners.portability.stager import Stager from apache_beam.transforms import fully_qualified_named_transform from apache_beam.transforms import ptransform from apache_beam.transforms.environments import PyPIArtifactRegistry from apache_beam.transforms.external import ImplicitSchemaPayloadBuilder from apache_beam.utils import thread_pool_executor # This script provides an expansion service and example ptransforms for running # external transform test cases. See external_test.py for details. _LOGGER = logging.getLogger(__name__) TEST_PREFIX_URN = "beam:transforms:xlang:test:prefix" TEST_MULTI_URN = "beam:transforms:xlang:test:multi" TEST_GBK_URN = "beam:transforms:xlang:test:gbk" TEST_CGBK_URN = "beam:transforms:xlang:test:cgbk" TEST_COMGL_URN = "beam:transforms:xlang:test:comgl" TEST_COMPK_URN = "beam:transforms:xlang:test:compk" TEST_FLATTEN_URN = "beam:transforms:xlang:test:flatten" TEST_PARTITION_URN = "beam:transforms:xlang:test:partition" TEST_PYTHON_BS4_URN = "beam:transforms:xlang:test:python_bs4" # A transform that does not produce an output. TEST_NO_OUTPUT_URN = "beam:transforms:xlang:test:nooutput" @ptransform.PTransform.register_urn('beam:transforms:xlang:count', None) class CountPerElementTransform(ptransform.PTransform): def expand(self, pcoll): return pcoll | combine.Count.PerElement() def to_runner_api_parameter(self, unused_context): return 'beam:transforms:xlang:count', None @staticmethod def from_runner_api_parameter( unused_ptransform, unused_parameter, unused_context): return CountPerElementTransform() @ptransform.PTransform.register_urn( 'beam:transforms:xlang:filter_less_than_eq', bytes) class FilterLessThanTransform(ptransform.PTransform): def __init__(self, payload): self._payload = payload def expand(self, pcoll): return ( pcoll | beam.Filter( lambda elem, target: elem <= target, int(ord(self._payload[0])))) def to_runner_api_parameter(self, unused_context): return ( 'beam:transforms:xlang:filter_less_than', self._payload.encode('utf8')) @staticmethod def from_runner_api_parameter(unused_ptransform, payload, unused_context): return FilterLessThanTransform(payload.decode('utf8')) @ptransform.PTransform.register_urn(TEST_PREFIX_URN, None) @beam.typehints.with_output_types(str) class PrefixTransform(ptransform.PTransform): def __init__(self, payload): self._payload = payload def expand(self, pcoll): return pcoll | 'TestLabel' >> beam.Map( lambda x: '{}{}'.format(self._payload, x)) def to_runner_api_parameter(self, unused_context): return TEST_PREFIX_URN, ImplicitSchemaPayloadBuilder( {'data': self._payload}).payload() @staticmethod def from_runner_api_parameter(unused_ptransform, payload, unused_context): return PrefixTransform(parse_string_payload(payload)['data']) @ptransform.PTransform.register_urn(TEST_MULTI_URN, None) class MutltiTransform(ptransform.PTransform): def expand(self, pcolls): return { 'main': (pcolls['main1'], pcolls['main2']) | beam.Flatten() | beam.Map(lambda x, s: x + s, beam.pvalue.AsSingleton( pcolls['side'])).with_output_types(str), 'side': pcolls['side'] | beam.Map(lambda x: x + x).with_output_types(str), } def to_runner_api_parameter(self, unused_context): return TEST_MULTI_URN, None @staticmethod def from_runner_api_parameter( unused_ptransform, unused_parameter, unused_context): return MutltiTransform() @ptransform.PTransform.register_urn(TEST_GBK_URN, None) class GBKTransform(ptransform.PTransform): def expand(self, pcoll): return pcoll | 'TestLabel' >> beam.GroupByKey() def to_runner_api_parameter(self, unused_context): return TEST_GBK_URN, None @staticmethod def from_runner_api_parameter( unused_ptransform, unused_parameter, unused_context): return GBKTransform() @ptransform.PTransform.register_urn(TEST_CGBK_URN, None) class CoGBKTransform(ptransform.PTransform): class ConcatFn(beam.DoFn): def process(self, element): (k, v) = element return [(k, v['col1'] + v['col2'])] def expand(self, pcoll): return pcoll \ | beam.CoGroupByKey() \ | beam.ParDo(self.ConcatFn()).with_output_types( tuple[int, typing.Iterable[str]]) def to_runner_api_parameter(self, unused_context): return TEST_CGBK_URN, None @staticmethod def from_runner_api_parameter( unused_ptransform, unused_parameter, unused_context): return CoGBKTransform() @ptransform.PTransform.register_urn(TEST_COMGL_URN, None) class CombineGloballyTransform(ptransform.PTransform): def expand(self, pcoll): return pcoll \ | beam.CombineGlobally(sum).with_output_types(int) def to_runner_api_parameter(self, unused_context): return TEST_COMGL_URN, None @staticmethod def from_runner_api_parameter( unused_ptransform, unused_parameter, unused_context): return CombineGloballyTransform() @ptransform.PTransform.register_urn(TEST_COMPK_URN, None) class CombinePerKeyTransform(ptransform.PTransform): def expand(self, pcoll): output = pcoll \ | beam.CombinePerKey(sum) # TODO: Use `with_output_types` instead of explicitly # assigning to `.element_type` after fixing BEAM-12872 output.element_type = beam.typehints.Tuple[str, int] return output def to_runner_api_parameter(self, unused_context): return TEST_COMPK_URN, None @staticmethod def from_runner_api_parameter( unused_ptransform, unused_parameter, unused_context): return CombinePerKeyTransform() @ptransform.PTransform.register_urn(TEST_FLATTEN_URN, None) class FlattenTransform(ptransform.PTransform): def expand(self, pcoll): return pcoll.values() | beam.Flatten().with_output_types(int) def to_runner_api_parameter(self, unused_context): return TEST_FLATTEN_URN, None @staticmethod def from_runner_api_parameter( unused_ptransform, unused_parameter, unused_context): return FlattenTransform() @ptransform.PTransform.register_urn(TEST_PARTITION_URN, None) class PartitionTransform(ptransform.PTransform): def expand(self, pcoll): col1, col2 = pcoll | beam.Partition( lambda elem, n: 0 if elem % 2 == 0 else 1, 2) typed_col1 = col1 | beam.Map(lambda x: x).with_output_types(int) typed_col2 = col2 | beam.Map(lambda x: x).with_output_types(int) return {'0': typed_col1, '1': typed_col2} def to_runner_api_parameter(self, unused_context): return TEST_PARTITION_URN, None @staticmethod def from_runner_api_parameter( unused_ptransform, unused_parameter, unused_context): return PartitionTransform() class ExtractHtmlTitleDoFn(beam.DoFn): def process(self, element): from bs4 import BeautifulSoup soup = BeautifulSoup(element, 'html.parser') return [soup.title.string] @ptransform.PTransform.register_urn(TEST_PYTHON_BS4_URN, None) class ExtractHtmlTitleTransform(ptransform.PTransform): def expand(self, pcoll): return pcoll | beam.ParDo(ExtractHtmlTitleDoFn()).with_output_types(str) def to_runner_api_parameter(self, unused_context): return TEST_PYTHON_BS4_URN, None @staticmethod def from_runner_api_parameter( unused_ptransform, unused_parameter, unused_context): return ExtractHtmlTitleTransform() @ptransform.PTransform.register_urn('payload', bytes) class PayloadTransform(ptransform.PTransform): def __init__(self, payload): self._payload = payload def expand(self, pcoll): return pcoll | beam.Map(lambda x, s: x + s, self._payload) def to_runner_api_parameter(self, unused_context): return b'payload', self._payload.encode('ascii') @staticmethod def from_runner_api_parameter(unused_ptransform, payload, unused_context): return PayloadTransform(payload.decode('ascii')) @ptransform.PTransform.register_urn('map_to_union_types', None) class MapToUnionTypesTransform(ptransform.PTransform): class CustomDoFn(beam.DoFn): def process(self, element): if element == 1: return ['1'] elif element == 2: return [2] else: return [3.0] def expand(self, pcoll): return pcoll | beam.ParDo(self.CustomDoFn()) def to_runner_api_parameter(self, unused_context): return b'map_to_union_types', None @staticmethod def from_runner_api_parameter( unused_ptransform, unused_payload, unused_context): return MapToUnionTypesTransform() @ptransform.PTransform.register_urn('fib', bytes) class FibTransform(ptransform.PTransform): def __init__(self, level): self._level = level def expand(self, p): if self._level <= 2: return p | beam.Create([1]) else: a = p | 'A' >> beam.ExternalTransform( 'fib', str(self._level - 1).encode('ascii'), expansion_service.ExpansionServiceServicer()) b = p | 'B' >> beam.ExternalTransform( 'fib', str(self._level - 2).encode('ascii'), expansion_service.ExpansionServiceServicer()) return ((a, b) | beam.Flatten() | beam.CombineGlobally(sum).without_defaults()) def to_runner_api_parameter(self, unused_context): return 'fib', str(self._level).encode('ascii') @staticmethod def from_runner_api_parameter(unused_ptransform, level, unused_context): return FibTransform(int(level.decode('ascii'))) @ptransform.PTransform.register_urn(TEST_NO_OUTPUT_URN, None) class NoOutputTransform(ptransform.PTransform): def expand(self, pcoll): def log_val(val): logging.debug('Got value: %r', val) # Logging without returning anything _ = (pcoll | 'TestLabel' >> beam.ParDo(log_val)) def to_runner_api_parameter(self, unused_context): return TEST_NO_OUTPUT_URN, None @staticmethod def from_runner_api_parameter(unused_ptransform, payload, unused_context): return NoOutputTransform(parse_string_payload(payload)['data']) def parse_string_payload(input_byte): payload = external_transforms_pb2.ExternalConfigurationPayload() payload.ParseFromString(input_byte) return RowCoder(payload.schema).decode(payload.payload)._asdict() def create_test_sklearn_model(file_name): from sklearn import svm x = [[0, 0], [1, 1]] y = [0, 1] model = svm.SVC() model.fit(x, y) with open(file_name, 'wb') as file: pickle.dump(model, file) def update_sklearn_model_dependency(env): model_file = "/tmp/sklearn_test_model" staged_name = "sklearn_model" create_test_sklearn_model(model_file) env._artifacts.append( Stager._create_file_stage_to_artifact(model_file, staged_name)) server = None def cleanup(unused_signum, unused_frame): _LOGGER.info('Shutting down expansion service.') server.stop(None) def main(unused_argv): # TODO: use the regular expansion service (expansion_service_main) instead of # this custom service for testing. PyPIArtifactRegistry.register_artifact('beautifulsoup4', '>=4.9,<5.0') parser = argparse.ArgumentParser() parser.add_argument( '-p', '--port', type=int, help='port on which to serve the job api') parser.add_argument('--fully_qualified_name_glob', default=None) options = parser.parse_args() global server with fully_qualified_named_transform.FullyQualifiedNamedTransform.with_filter( options.fully_qualified_name_glob): server = grpc.server(thread_pool_executor.shared_unbounded_instance()) expansion_servicer = expansion_service.ExpansionServiceServicer( PipelineOptions([ "--experiments", "beam_fn_api", "--sdk_location", "container", "--pickle_library", "cloudpickle" ])) update_sklearn_model_dependency(expansion_servicer._default_environment) beam_expansion_api_pb2_grpc.add_ExpansionServiceServicer_to_server( expansion_servicer, server) beam_artifact_api_pb2_grpc.add_ArtifactRetrievalServiceServicer_to_server( artifact_service.ArtifactRetrievalService( artifact_service.BeamFilesystemHandler(None).file_reader), server) server.add_insecure_port('localhost:{}'.format(options.port)) server.start() _LOGGER.info('Listening for expansion requests at %d', options.port) signal.signal(signal.SIGTERM, cleanup) signal.signal(signal.SIGINT, cleanup) # blocking main thread forever. signal.pause() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) main(sys.argv) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/flink_runner.py0000644000175100017510000001123115206113032024757 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A runner for executing portable pipelines on Flink.""" # pytype: skip-file import logging import os import re import urllib from apache_beam.options import pipeline_options from apache_beam.runners.portability import flink_uber_jar_job_server from apache_beam.runners.portability import job_server from apache_beam.runners.portability import portable_runner MAGIC_HOST_NAMES = ['[local]', '[auto]'] _LOGGER = logging.getLogger(__name__) class FlinkRunner(portable_runner.PortableRunner): """A runner for launching jobs on Flink, automatically starting a local flink master if one is not given. """ # Inherits run_portable_pipeline from PortableRunner. def default_environment(self, options): portable_options = options.view_as(pipeline_options.PortableOptions) flink_options = options.view_as(pipeline_options.FlinkRunnerOptions) if (flink_options.flink_master in MAGIC_HOST_NAMES and not portable_options.environment_type and not portable_options.output_executable_path): portable_options.environment_type = 'LOOPBACK' return super().default_environment(options) def default_job_server(self, options): flink_options = options.view_as(pipeline_options.FlinkRunnerOptions) flink_master = self.add_http_scheme(flink_options.flink_master) flink_options.flink_master = flink_master if (flink_options.flink_submit_uber_jar and flink_master not in MAGIC_HOST_NAMES): # This has to be changed [auto], otherwise we will attempt to submit a # the pipeline remotely on the Flink JobMaster which will _fail_. # DO NOT CHANGE the following line, unless you have tested this. flink_options.flink_master = '[auto]' return flink_uber_jar_job_server.FlinkUberJarJobServer( flink_master, options) else: return job_server.StopOnExitJobServer(FlinkJarJobServer(options)) def create_job_service_handle(self, job_service, options): return portable_runner.JobServiceHandle( job_service, options, retain_unknown_options=options.view_as( pipeline_options.FlinkRunnerOptions).flink_submit_uber_jar) @staticmethod def add_http_scheme(flink_master): """Adds a http protocol scheme if none provided.""" flink_master = flink_master.strip() if not flink_master in MAGIC_HOST_NAMES and \ not re.search('^http[s]?://', flink_master): _LOGGER.info( 'Adding HTTP protocol scheme to flink_master parameter: ' 'http://%s', flink_master) flink_master = 'http://' + flink_master return flink_master class FlinkJarJobServer(job_server.JavaJarJobServer): def __init__(self, options): super().__init__(options) options = options.view_as(pipeline_options.FlinkRunnerOptions) self._jar = options.flink_job_server_jar self._master_url = options.flink_master self._flink_version = options.flink_version def path_to_jar(self): if self._jar: if not os.path.exists(self._jar): url = urllib.parse.urlparse(self._jar) if not url.scheme: raise ValueError( 'Unable to parse jar URL "%s". If using a full URL, make sure ' 'the scheme is specified. If using a local file path, make sure ' 'the file exists; you may have to first build the job server ' 'using `./gradlew runners:flink:%s:job-server:shadowJar`.' % (self._jar, self._flink_version)) return self._jar else: return self.path_to_beam_jar( ':runners:flink:%s:job-server:shadowJar' % self._flink_version) def java_arguments( self, job_port, artifact_port, expansion_port, artifacts_dir): return [ '--flink-master', self._master_url, '--artifacts-dir', artifacts_dir, '--job-port', job_port, '--artifact-port', artifact_port, '--expansion-port', expansion_port ] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/flink_runner_test.py0000644000175100017510000003720615206113032026030 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import argparse import logging import platform import shlex import typing import unittest from os import linesep from os import path from os.path import exists from shutil import rmtree from tempfile import mkdtemp import pytest import apache_beam as beam from apache_beam import Impulse from apache_beam import Map from apache_beam.io.external.generate_sequence import GenerateSequence from apache_beam.io.kafka import ReadFromKafka from apache_beam.io.kafka import WriteToKafka from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import FlinkRunnerOptions from apache_beam.options.pipeline_options import PortableOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.runners.portability import job_server from apache_beam.runners.portability import portable_runner from apache_beam.runners.portability import portable_runner_test from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.sql import SqlTransform # Run as # # pytest flink_runner_test.py[::TestClass::test_case] \ # --test-pipeline-options="--environment_type=LOOPBACK" _LOGGER = logging.getLogger(__name__) Row = typing.NamedTuple("Row", [("col1", int), ("col2", str)]) beam.coders.registry.register_coder(Row, beam.coders.RowCoder) class FlinkRunnerTest(portable_runner_test.PortableRunnerTest): _use_grpc = True _use_subprocesses = True conf_dir = None expansion_port = None flink_job_server_jar = None def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.environment_type = None self.environment_config = None self.enable_commit = False def setUp(self): self.enable_commit = False @pytest.fixture(autouse=True) def parse_options(self, request): if not request.config.option.test_pipeline_options: raise unittest.SkipTest( 'Skipping because --test-pipeline-options is not specified.') test_pipeline_options = request.config.option.test_pipeline_options parser = argparse.ArgumentParser(add_help=True) parser.add_argument( '--flink_job_server_jar', help='Job server jar to submit jobs.', action='store') parser.add_argument( '--environment_type', default='LOOPBACK', choices=['DOCKER', 'PROCESS', 'LOOPBACK'], help='Set the environment type for running user code. DOCKER runs ' 'user code in a container. PROCESS runs user code in ' 'automatically started processes. LOOPBACK runs user code on ' 'the same process that originally submitted the job.') parser.add_argument( '--environment_option', '--environment_options', dest='environment_options', action='append', default=None, help=( 'Environment configuration for running the user code. ' 'Recognized options depend on --environment_type.\n ' 'For DOCKER: docker_container_image (optional)\n ' 'For PROCESS: process_command (required), process_variables ' '(optional, comma-separated)\n ' 'For EXTERNAL: external_service_address (required)')) known_args, unknown_args = parser.parse_known_args( shlex.split(test_pipeline_options)) if unknown_args: _LOGGER.warning('Discarding unrecognized arguments %s' % unknown_args) self.set_flink_job_server_jar( known_args.flink_job_server_jar or job_server.JavaJarJobServer.path_to_beam_jar(( ':runners:flink:%s:job-server:shadowJar' % FlinkRunnerOptions.PUBLISHED_FLINK_VERSIONS[-1]))) self.environment_type = known_args.environment_type self.environment_options = known_args.environment_options @classmethod def tearDownClass(cls): if cls.conf_dir and exists(cls.conf_dir): _LOGGER.info("removing conf dir: %s" % cls.conf_dir) rmtree(cls.conf_dir) super().tearDownClass() @classmethod def _create_conf_dir(cls): """Create (and save a static reference to) a "conf dir", used to provide metrics configs and verify metrics output It gets cleaned up when the suite is done executing""" if hasattr(cls, 'conf_dir'): cls.conf_dir = mkdtemp(prefix='flinktest-conf') # path for a FileReporter to write metrics to cls.test_metrics_path = path.join(cls.conf_dir, 'test-metrics.txt') # path to write Flink configuration to # Flink 1.x conf: conf_path = path.join(cls.conf_dir, 'flink-conf.yaml') file_reporter = 'org.apache.beam.runners.flink.metrics.FileReporter' with open(conf_path, 'w') as f: f.write( linesep.join([ 'metrics.reporters: file', 'metrics.reporter.file.class: %s' % file_reporter, 'metrics.reporter.file.path: %s' % cls.test_metrics_path, 'metrics.scope.operator: ', ])) # Flink 2.x conf: conf_path_2 = path.join(cls.conf_dir, 'config.yaml') with open(conf_path_2, 'w') as f: f.write( '''metrics: reporters: file reporter: file: class: %s path: %s scope: operator: ''' % (file_reporter, cls.test_metrics_path)) @classmethod def _subprocess_command(cls, job_port, expansion_port): # will be cleaned up at the end of this method, and recreated and used by # the job server tmp_dir = mkdtemp(prefix='flinktest') cls._create_conf_dir() cls.expansion_port = expansion_port platform_specific_opts = [] if platform.system() == 'Linux': # UseContainerSupport is supported in Linux and turned on by default platform_specific_opts.append('-XX:-UseContainerSupport') try: return ['java'] + platform_specific_opts + [ '--add-opens=java.base/java.lang=ALL-UNNAMED', '--add-opens=java.base/java.nio=ALL-UNNAMED', '--add-opens=java.base/java.util=ALL-UNNAMED', '-Dslf4j.provider=org.slf4j.simple.SimpleServiceProvider', '-Dorg.slf4j.simpleLogger.log.org.apache.flink.metrics=error', '-Dorg.slf4j.simpleLogger.log.org.apache.flink.runtime=error', '-Dorg.slf4j.simpleLogger.log.org.apache.flink.streaming=error', '-jar', cls.flink_job_server_jar, '--flink-master', '[local]', '--flink-conf-dir', cls.conf_dir, '--artifacts-dir', tmp_dir, '--job-port', str(job_port), '--artifact-port', '0', '--expansion-port', str(expansion_port), ] finally: rmtree(tmp_dir) @classmethod def get_runner(cls): return portable_runner.PortableRunner() @classmethod def get_expansion_service(cls): # TODO Move expansion address resides into PipelineOptions return 'localhost:%s' % cls.expansion_port @classmethod def set_flink_job_server_jar(cls, flink_job_server_jar): cls.flink_job_server_jar = flink_job_server_jar def create_options(self): options = super().create_options() options.view_as(DebugOptions).experiments = ['beam_fn_api'] options._all_options['parallelism'] = 2 options.view_as(PortableOptions).environment_type = self.environment_type options.view_as( PortableOptions).environment_options = self.environment_options if self.enable_commit: options.view_as(StandardOptions).streaming = True options._all_options['checkpointing_interval'] = 3000 options._all_options['shutdown_sources_after_idle_ms'] = 60000 options._all_options['number_of_execution_retries'] = 1 return options def test_batch_rebatch_pardos(self): if self.environment_type == 'DOCKER': pytest.skip( "Skipping test_batch_rebatch_pardos for DOCKER environment in " "FlinkRunnerTest as warnings are not propagated back when Python SDK " "is run in Docker.") super().test_batch_rebatch_pardos() # Can't read host files from within docker, read a "local" file there. def test_read(self): print('name:', __name__) with self.create_pipeline() as p: lines = p | beam.io.ReadFromText('/etc/profile') assert_that(lines, lambda lines: len(lines) > 0) def test_no_subtransform_composite(self): raise unittest.SkipTest("BEAM-4781") def test_external_transform(self): if self.environment_type == 'PROCESS': self.skipTest( "Skipping external transform test in PROCESS mode due to " "https://github.com/apache/beam/issues/19461 (Flink expansion " "service incompatibility).") with self.create_pipeline() as p: res = ( p | GenerateSequence( start=1, stop=10, expansion_service=self.get_expansion_service())) assert_that(res, equal_to([i for i in range(1, 10)])) def test_expand_kafka_read(self): if self.environment_type == 'PROCESS': self.skipTest( "Skipping Kafka read test in PROCESS mode due to " "https://github.com/apache/beam/issues/19461 (Flink expansion " "service incompatibility).") # We expect to fail here because we do not have a Kafka cluster handy. # Nevertheless, we check that the transform is expanded by the # ExpansionService and that the pipeline fails during execution. with self.assertRaises(Exception) as ctx: self.enable_commit = True with self.create_pipeline() as p: # pylint: disable=expression-not-assigned ( p | ReadFromKafka( consumer_config={ 'bootstrap.servers': 'notvalid1:7777, notvalid2:3531', 'group.id': 'any_group' }, topics=['topic1', 'topic2'], key_deserializer='org.apache.kafka.' 'common.serialization.' 'ByteArrayDeserializer', value_deserializer='org.apache.kafka.' 'common.serialization.' 'LongDeserializer', commit_offset_in_finalize=True, timestamp_policy=ReadFromKafka.create_time_policy, redistribute=False, redistribute_num_keys=0, allow_duplicates=False, expansion_service=self.get_expansion_service())) self.assertTrue( 'No resolvable bootstrap urls given in bootstrap.servers' in str(ctx.exception), 'Expected to fail due to invalid bootstrap.servers, but ' 'failed due to:\n%s' % str(ctx.exception)) def test_expand_kafka_write(self): if self.environment_type == 'PROCESS': self.skipTest( "Skipping Kafka write test in PROCESS mode due to " "https://github.com/apache/beam/issues/19461 (Flink expansion " "service incompatibility).") # We just test the expansion but do not execute. # pylint: disable=expression-not-assigned ( self.create_pipeline() | Impulse() | Map(lambda input: (1, input)) | WriteToKafka( producer_config={ 'bootstrap.servers': 'localhost:9092, notvalid2:3531' }, topic='topic1', key_serializer='org.apache.kafka.' 'common.serialization.' 'LongSerializer', value_serializer='org.apache.kafka.' 'common.serialization.' 'ByteArraySerializer', expansion_service=self.get_expansion_service())) def test_sql(self): if self.environment_type == 'PROCESS': self.skipTest( "Skipping SQL test in PROCESS mode due to " "https://github.com/apache/beam/issues/19461 (Flink expansion " "service incompatibility).") with self.create_pipeline() as p: output = ( p | 'Create' >> beam.Create([Row(x, str(x)) for x in range(5)]) | 'Sql' >> SqlTransform( """SELECT col1, col2 || '*' || col2 as col2, power(col1, 2) as col3 FROM PCOLLECTION """, expansion_service=self.get_expansion_service())) assert_that( output, equal_to([(x, '{x}*{x}'.format(x=x), x * x) for x in range(5)])) def test_flattened_side_input(self): # Blocked on support for transcoding # https://github.com/apache/beam/issues/19365 super().test_flattened_side_input(with_transcoding=False) def test_flatten_and_gbk(self): # Blocked on support for transcoding # https://github.com/apache/beam/issues/19365 # Also blocked on support of flatten and groupby sharing the same input # https://github.com/apache/beam/issues/34647 raise unittest.SkipTest("https://github.com/apache/beam/issues/34647") def test_metrics(self): super().test_metrics(check_gauge=False, check_bounded_trie=False) def test_sdf_with_watermark_tracking(self): raise unittest.SkipTest("BEAM-2939") def test_callbacks_with_exception(self): raise unittest.SkipTest("https://github.com/apache/beam/issues/19526") def test_register_finalizations(self): raise unittest.SkipTest("https://github.com/apache/beam/issues/19526") def test_custom_merging_window(self): raise unittest.SkipTest("https://github.com/apache/beam/issues/20641") def test_custom_window_type(self): raise unittest.SkipTest("https://github.com/apache/beam/issues/20641") # Inherits all other tests. class FlinkRunnerTestOptimized(FlinkRunnerTest): # TODO: Remove these tests after resolving # https://github.com/apache/beam/issues/19422 and enabling # PortableRunnerOptimized def create_options(self): options = super().create_options() options.view_as(DebugOptions).experiments = [ 'pre_optimize=all' ] + options.view_as(DebugOptions).experiments return options def test_external_transform(self): raise unittest.SkipTest("https://github.com/apache/beam/issues/19461") def test_expand_kafka_read(self): raise unittest.SkipTest("https://github.com/apache/beam/issues/19461") def test_expand_kafka_write(self): raise unittest.SkipTest("https://github.com/apache/beam/issues/19461") def test_sql(self): raise unittest.SkipTest("https://github.com/apache/beam/issues/19461") def test_pack_combiners(self): # Stages produced by translations.pack_combiners are fused # by translations.greedily_fuse, which prevent the stages # from being detecting using counters by the test. self._test_pack_combiners(assert_using_counter_names=False) class FlinkRunnerTestStreaming(FlinkRunnerTest): def create_options(self): options = super().create_options() options.view_as(StandardOptions).streaming = True return options def test_callbacks_with_exception(self): self.enable_commit = True super().test_callbacks_with_exception() def test_register_finalizations(self): self.enable_commit = True super().test_register_finalizations() if __name__ == '__main__': # Run the tests. logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/flink_uber_jar_job_server.py0000644000175100017510000002156615206113032027473 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A job server submitting portable pipelines as uber jars to Flink.""" # pytype: skip-file import logging import os import tempfile import time import urllib import requests from google.protobuf import json_format from apache_beam.options import pipeline_options from apache_beam.portability.api import beam_job_api_pb2 from apache_beam.runners.portability import abstract_job_service from apache_beam.runners.portability import job_server _LOGGER = logging.getLogger(__name__) class FlinkUberJarJobServer(abstract_job_service.AbstractJobServiceServicer): """A Job server which submits a self-contained Jar to a Flink cluster. The jar contains the Beam pipeline definition, dependencies, and the pipeline artifacts. """ def __init__(self, master_url, options): super().__init__() self._master_url = master_url self._executable_jar = ( options.view_as( pipeline_options.FlinkRunnerOptions).flink_job_server_jar) self._user_agent = options.view_as(pipeline_options.SetupOptions).user_agent self._artifact_port = ( options.view_as(pipeline_options.JobServerOptions).artifact_port) self._temp_dir = tempfile.mkdtemp(prefix='apache-beam-flink') def start(self): return self def stop(self): pass def executable_jar(self): if self._executable_jar: if not os.path.exists(self._executable_jar): parsed = urllib.parse.urlparse(self._executable_jar) if not parsed.scheme: try: flink_version = self.flink_version() except Exception: flink_version = '$FLINK_VERSION' raise ValueError( 'Unable to parse jar URL "%s". If using a full URL, make sure ' 'the scheme is specified. If using a local file path, make sure ' 'the file exists; you may have to first build the job server ' 'using `./gradlew runners:flink:%s:job-server:shadowJar`.' % (self._executable_jar, flink_version)) url = self._executable_jar else: url = job_server.JavaJarJobServer.path_to_beam_jar( ':runners:flink:%s:job-server:shadowJar' % self.flink_version()) return job_server.JavaJarJobServer.local_jar( url, user_agent=self._user_agent) def flink_version(self): full_version = requests.get( '%s/v1/config' % self._master_url, timeout=60).json()['flink-version'] # Only return up to minor version. return '.'.join(full_version.split('.')[:2]) def create_beam_job(self, job_id, job_name, pipeline, options): return FlinkBeamJob( self._master_url, self.executable_jar(), job_id, job_name, pipeline, options, artifact_port=self._artifact_port) def GetJobMetrics(self, request, context=None): if request.job_id not in self._jobs: raise LookupError("Job {} does not exist".format(request.job_id)) metrics_text = self._jobs[request.job_id].get_metrics() response = beam_job_api_pb2.GetJobMetricsResponse() json_format.Parse(metrics_text, response) return response class FlinkBeamJob(abstract_job_service.UberJarBeamJob): """Runs a single Beam job on Flink by staging all contents into a Jar and uploading it via the Flink Rest API.""" def __init__( self, master_url, executable_jar, job_id, job_name, pipeline, options, artifact_port=0): super().__init__( executable_jar, job_id, job_name, pipeline, options, artifact_port=artifact_port) self._master_url = master_url def request(self, method, path, expected_status=200, **kwargs): url = '%s/%s' % (self._master_url, path) response = method(url, **kwargs) if response.status_code != expected_status: raise RuntimeError( "Request to %s failed with status %d: %s" % (url, response.status_code, response.text)) if response.text: return response.json() def get(self, path, **kwargs): return self.request(requests.get, path, **kwargs) def post(self, path, **kwargs): return self.request(requests.post, path, **kwargs) def delete(self, path, **kwargs): return self.request(requests.delete, path, **kwargs) def run(self): self._stop_artifact_service() # Upload the jar and start the job. with open(self._jar, 'rb') as jar_file: self._flink_jar_id = self.post( 'v1/jars/upload', files={'jarfile': ('beam.jar', jar_file)})['filename'].split('/')[-1] self._jar_uploaded = True self._flink_job_id = self.post( 'v1/jars/%s/run' % self._flink_jar_id, json={ 'entryClass': 'org.apache.beam.runners.flink.FlinkPipelineRunner' })['jobid'] os.unlink(self._jar) _LOGGER.info('Started Flink job as %s' % self._flink_job_id) def cancel(self): self.post('v1/%s/stop' % self._flink_job_id, expected_status=202) self.delete_jar() def delete_jar(self): if self._jar_uploaded: self._jar_uploaded = False try: self.delete('v1/jars/%s' % self._flink_jar_id) except Exception: _LOGGER.info( 'Error deleting jar %s' % self._flink_jar_id, exc_info=True) def _get_state(self): """Query flink to get the current state. :return: tuple of int and Timestamp or None timestamp will be None if the state has not changed since the last query. """ # For just getting the status, execution-result seems cheaper. flink_status = self.get('v1/jobs/%s/execution-result' % self._flink_job_id)['status']['id'] if flink_status == 'COMPLETED': flink_status = self.get('v1/jobs/%s' % self._flink_job_id)['state'] beam_state = { 'CREATED': beam_job_api_pb2.JobState.STARTING, 'RUNNING': beam_job_api_pb2.JobState.RUNNING, 'FAILING': beam_job_api_pb2.JobState.RUNNING, 'FAILED': beam_job_api_pb2.JobState.FAILED, 'CANCELLING': beam_job_api_pb2.JobState.CANCELLING, 'CANCELED': beam_job_api_pb2.JobState.CANCELLED, 'FINISHED': beam_job_api_pb2.JobState.DONE, 'RESTARTING': beam_job_api_pb2.JobState.RUNNING, 'SUSPENDED': beam_job_api_pb2.JobState.RUNNING, 'RECONCILING': beam_job_api_pb2.JobState.RUNNING, 'IN_PROGRESS': beam_job_api_pb2.JobState.RUNNING, 'COMPLETED': beam_job_api_pb2.JobState.DONE, }.get(flink_status, beam_job_api_pb2.JobState.UNSPECIFIED) if self.is_terminal_state(beam_state): self.delete_jar() # update the state history if it has changed return beam_state, self.set_state(beam_state) def get_state(self): state, timestamp = self._get_state() if timestamp is None: # state has not changed since it was last checked: use previous timestamp return super().get_state() else: return state, timestamp def get_state_stream(self): def _state_iter(): sleep_secs = 1.0 while True: yield self.get_state() sleep_secs = min(60, sleep_secs * 1.2) time.sleep(sleep_secs) for state, timestamp in self.with_state_history(_state_iter()): yield state, timestamp if self.is_terminal_state(state): break def get_message_stream(self): for state, timestamp in self.get_state_stream(): if self.is_terminal_state(state): response = self.get('v1/jobs/%s/exceptions' % self._flink_job_id) for ix, exc in enumerate(response['all-exceptions']): yield beam_job_api_pb2.JobMessage( message_id='message%d' % ix, time=str(exc['timestamp']), importance=beam_job_api_pb2.JobMessage.MessageImportance. JOB_MESSAGE_ERROR, message_text=exc['exception']) yield state, timestamp break else: yield state, timestamp def get_metrics(self): accumulators = self.get('v1/jobs/%s/accumulators' % self._flink_job_id)['user-task-accumulators'] for accumulator in accumulators: if accumulator['name'] == '__metricscontainers': return accumulator['value'] raise LookupError( "Found no metrics container for job {}".format(self._flink_job_id)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/flink_uber_jar_job_server_test.py0000644000175100017510000002110515206113032030517 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import contextlib import logging import os import tempfile import unittest import zipfile import requests_mock from apache_beam.options import pipeline_options from apache_beam.portability.api import beam_job_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.portability import flink_runner from apache_beam.runners.portability import flink_uber_jar_job_server from apache_beam.runners.portability.local_job_service_test import TestJobServicePlan @contextlib.contextmanager def temp_name(*args, **kwargs): with tempfile.NamedTemporaryFile(*args, **kwargs) as t: name = t.name try: yield name finally: if os.path.exists(name): os.unlink(name) class FlinkUberJarJobServerTest(unittest.TestCase): @requests_mock.mock() def test_flink_version(self, http_mock): http_mock.get('http://flink/v1/config', json={'flink-version': '3.1.4.1'}) job_server = flink_uber_jar_job_server.FlinkUberJarJobServer( 'http://flink', pipeline_options.FlinkRunnerOptions()) self.assertEqual(job_server.flink_version(), "3.1") @requests_mock.mock() def test_get_job_metrics(self, http_mock): response = { "user-task-accumulators": [{ "name": "__metricscontainers", "type": "MetricsAccumulator", "value": "{\"metrics\": {\"attempted\": [{\"urn\": " "\"metric_urn\", \"type\": \"beam:metrics:sum_int64:v1\", " "\"payload\": \"AA==\", \"labels\": " "{\"PTRANSFORM\": \"ptransform_id\"}}]}}" }] } http_mock.get( 'http://flink/v1/jobs/flink_job_id/accumulators', json=response) options = pipeline_options.FlinkRunnerOptions() job_server = flink_uber_jar_job_server.FlinkUberJarJobServer( 'http://flink', options) job = flink_uber_jar_job_server.FlinkBeamJob( 'http://flink', None, 'job_id', 'job_name', None, options) job._flink_job_id = 'flink_job_id' job_server._jobs['job_id'] = job request = beam_job_api_pb2.GetJobMetricsRequest(job_id='job_id') expected = beam_job_api_pb2.GetJobMetricsResponse( metrics=beam_job_api_pb2.MetricResults( attempted=[{ "urn": "metric_urn", "type": "beam:metrics:sum_int64:v1", "payload": b'\000', "labels": { "PTRANSFORM": "ptransform_id" } }])) actual = job_server.GetJobMetrics(request) self.assertEqual(actual, expected) @requests_mock.mock() def test_end_to_end(self, http_mock): with temp_name(suffix='fake.jar') as fake_jar: # Create the jar file with some trivial contents. with zipfile.ZipFile(fake_jar, 'w') as zip: with zip.open('FakeClass.class', 'w') as fout: fout.write(b'[original_contents]') options = pipeline_options.FlinkRunnerOptions() options.flink_job_server_jar = fake_jar job_server = flink_uber_jar_job_server.FlinkUberJarJobServer( 'http://flink', options) plan = TestJobServicePlan(job_server) # Prepare the job. prepare_response = plan.prepare(beam_runner_api_pb2.Pipeline()) plan.stage( beam_runner_api_pb2.Pipeline(), prepare_response.artifact_staging_endpoint.url, prepare_response.staging_session_token) # Now actually run the job. http_mock.post( 'http://flink/v1/jars/upload', json={'filename': '/path/to/jar/nonce'}) http_mock.post( 'http://flink/v1/jars/nonce/run', json={'jobid': 'some_job_id'}) _, message_stream, state_stream = plan.run( prepare_response.preparation_id) # Check the status until the job is "done" and get all error messages. http_mock.get( 'http://flink/v1/jobs/some_job_id/execution-result', [{ 'json': { 'status': { 'id': 'IN_PROGRESS' } } }, { 'json': { 'status': { 'id': 'IN_PROGRESS' } } }, { 'json': { 'status': { 'id': 'COMPLETED' } } }]) http_mock.get( 'http://flink/v1/jobs/some_job_id', json={'state': 'FINISHED'}) http_mock.delete('http://flink/v1/jars/nonce') self.assertEqual([s.state for s in state_stream], [ beam_job_api_pb2.JobState.STOPPED, beam_job_api_pb2.JobState.RUNNING, beam_job_api_pb2.JobState.RUNNING, beam_job_api_pb2.JobState.DONE ]) http_mock.get( 'http://flink/v1/jobs/some_job_id/exceptions', json={'all-exceptions': [{ 'exception': 'exc_text', 'timestamp': 0 }]}) def get_item(x): if x.HasField('message_response'): return x.message_response else: return x.state_response.state self.assertEqual([get_item(m) for m in message_stream], [ beam_job_api_pb2.JobState.STOPPED, beam_job_api_pb2.JobState.RUNNING, beam_job_api_pb2.JobMessage( message_id='message0', time='0', importance=beam_job_api_pb2.JobMessage. MessageImportance.JOB_MESSAGE_ERROR, message_text='exc_text'), beam_job_api_pb2.JobState.DONE, ]) def test_retain_unknown_options(self): original_options = pipeline_options.PipelineOptions( ['--unknown_option_foo=some_value']) flink_options = original_options.view_as( pipeline_options.FlinkRunnerOptions) flink_options.flink_submit_uber_jar = True flink_options.flink_master = 'http://host:port' runner = flink_runner.FlinkRunner() job_service_handle = runner.create_job_service(original_options) options_proto = job_service_handle.get_pipeline_options() self.assertEqual( options_proto['beam:option:unknown_option_foo:v1'], 'some_value') @requests_mock.mock() def test_bad_url_flink_version(self, http_mock): http_mock.get('http://flink/v1/config', json={'flink-version': '1.2.3.4'}) options = pipeline_options.FlinkRunnerOptions() options.flink_job_server_jar = "bad url" job_server = flink_uber_jar_job_server.FlinkUberJarJobServer( 'http://flink', options) with self.assertRaises(ValueError) as context: job_server.executable_jar() self.assertEqual( 'Unable to parse jar URL "bad url". If using a full URL, make sure ' 'the scheme is specified. If using a local file path, make sure ' 'the file exists; you may have to first build the job server ' 'using `./gradlew runners:flink:1.2:job-server:shadowJar`.', str(context.exception)) def test_bad_url_placeholder_version(self): options = pipeline_options.FlinkRunnerOptions() options.flink_job_server_jar = "bad url" job_server = flink_uber_jar_job_server.FlinkUberJarJobServer( 'http://example.com/bad', options) with self.assertRaises(ValueError) as context: job_server.executable_jar() self.assertEqual( 'Unable to parse jar URL "bad url". If using a full URL, make sure ' 'the scheme is specified. If using a local file path, make sure ' 'the file exists; you may have to first build the job server ' 'using `./gradlew runners:flink:$FLINK_VERSION:job-server:shadowJar`.', str(context.exception)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.6814396 apache_beam-2.74.0/apache_beam/runners/portability/fn_api_runner/0000755000175100017510000000000015206113160024542 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/fn_api_runner/__init__.py0000644000175100017510000000154115206113032026652 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from apache_beam.runners.portability.fn_api_runner.fn_runner import FnApiRunner ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/fn_api_runner/execution.py0000644000175100017510000014513615206113032027127 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Set of utilities for execution of a pipeline by the FnApiRunner.""" # mypy: disallow-untyped-defs import collections import copy import itertools import logging import struct import uuid import weakref from typing import TYPE_CHECKING from typing import Any from typing import Callable from typing import Generic from typing import Iterable from typing import Iterator from typing import MutableMapping from typing import Optional from typing import Sequence from typing import TypeVar from typing import Union from typing_extensions import Protocol from apache_beam import coders from apache_beam.coders.coder_impl import CoderImpl from apache_beam.coders.coder_impl import WindowedValueCoderImpl from apache_beam.coders.coder_impl import create_InputStream from apache_beam.coders.coder_impl import create_OutputStream from apache_beam.coders.coders import WindowedValueCoder from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.portability.api import endpoints_pb2 from apache_beam.runners import pipeline_context from apache_beam.runners.common import ENCODED_IMPULSE_VALUE from apache_beam.runners.direct.clock import RealClock from apache_beam.runners.direct.clock import TestClock from apache_beam.runners.portability.fn_api_runner import translations from apache_beam.runners.portability.fn_api_runner.translations import DataInput from apache_beam.runners.portability.fn_api_runner.translations import DataOutput from apache_beam.runners.portability.fn_api_runner.translations import OutputTimers from apache_beam.runners.portability.fn_api_runner.translations import Stage from apache_beam.runners.portability.fn_api_runner.translations import create_buffer_id from apache_beam.runners.portability.fn_api_runner.translations import only_element from apache_beam.runners.portability.fn_api_runner.translations import split_buffer_id from apache_beam.runners.portability.fn_api_runner.translations import unique_name from apache_beam.runners.portability.fn_api_runner.watermark_manager import WatermarkManager from apache_beam.runners.worker import bundle_processor from apache_beam.transforms import core from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.transforms.window import BoundedWindow from apache_beam.transforms.window import GlobalWindow from apache_beam.transforms.window import GlobalWindows from apache_beam.utils import proto_utils from apache_beam.utils import windowed_value from apache_beam.utils.timestamp import MAX_TIMESTAMP from apache_beam.utils.timestamp import Timestamp if TYPE_CHECKING: from apache_beam.runners.portability.fn_api_runner import worker_handlers from apache_beam.runners.portability.fn_api_runner.translations import DataSideInput _LOGGER = logging.getLogger(__name__) SAFE_WINDOW_FNS = set( window.WindowFn._known_urns.keys()) - {python_urns.PICKLED_WINDOWFN} class Buffer(Protocol): def __iter__(self) -> Iterator[bytes]: pass def append(self, item: bytes) -> None: pass def extend(self, other: 'Buffer') -> None: pass class PartitionableBuffer(Buffer, Protocol): def copy(self) -> 'PartitionableBuffer': pass def partition(self, n: int) -> list[list[bytes]]: pass @property def cleared(self) -> bool: pass def clear(self) -> None: pass def reset(self) -> None: pass class ListBuffer: """Used to support parititioning of a list.""" def __init__(self, coder_impl: Optional[CoderImpl]) -> None: self._coder_impl = coder_impl or CoderImpl() self._inputs: list[bytes] = [] self._grouped_output: Optional[list[list[bytes]]] = None self.cleared = False def copy(self) -> 'ListBuffer': new = ListBuffer(self._coder_impl) new._inputs = [v for v in self._inputs] return new def extend(self, extra: 'Buffer') -> None: if self.cleared: raise RuntimeError('Trying to append to a cleared ListBuffer.') if self._grouped_output: raise RuntimeError('ListBuffer append after read.') assert isinstance(extra, ListBuffer) self._inputs.extend(extra._inputs) def append(self, element: bytes) -> None: if self.cleared: raise RuntimeError('Trying to append to a cleared ListBuffer.') if self._grouped_output: raise RuntimeError('ListBuffer append after read.') self._inputs.append(element) def partition(self, n: int) -> list[list[bytes]]: if self.cleared: raise RuntimeError('Trying to partition a cleared ListBuffer.') if len(self._inputs) >= n or len(self._inputs) == 0: return [self._inputs[k::n] for k in range(n)] else: if not self._grouped_output: output_stream_list = [create_OutputStream() for _ in range(n)] idx = 0 for input in self._inputs: input_stream = create_InputStream(input) while input_stream.size() > 0: decoded_value = self._coder_impl.decode_from_stream( input_stream, True) self._coder_impl.encode_to_stream( decoded_value, output_stream_list[idx], True) idx = (idx + 1) % n self._grouped_output = [[output_stream.get()] for output_stream in output_stream_list] return self._grouped_output def __iter__(self) -> Iterator[bytes]: if self.cleared: raise RuntimeError('Trying to iterate through a cleared ListBuffer.') return iter(self._inputs) def clear(self) -> None: self.cleared = True self._inputs = [] self._grouped_output = None def reset(self) -> None: """Resets a cleared buffer for reuse.""" if not self.cleared: raise RuntimeError('Trying to reset a non-cleared ListBuffer.') self.cleared = False class GroupingBuffer(object): """Used to accumulate groupded (shuffled) results.""" def __init__( self, pre_grouped_coder: coders.Coder, post_grouped_coder: coders.Coder, windowing: core.Windowing) -> None: self._key_coder = pre_grouped_coder.key_coder() self._pre_grouped_coder = pre_grouped_coder self._post_grouped_coder = post_grouped_coder self._table: collections.defaultdict[bytes, list[Any]] = collections.defaultdict( list) self._windowing = windowing self._grouped_output: Optional[list[list[bytes]]] = None def copy(self) -> 'GroupingBuffer': # This is a silly temporary optimization. This class must be removed once # full support for streaming is added (i.e. once we use trigger_manager for # data grouping instead of GroupingBuffer). return self def append(self, elements_data: bytes) -> None: if self._grouped_output: raise RuntimeError('Grouping table append after read.') input_stream = create_InputStream(elements_data) coder_impl = self._pre_grouped_coder.get_impl() key_coder_impl = self._key_coder.get_impl() # TODO(robertwb): We could optimize this even more by using a # window-dropping coder for the data plane. is_trivial_windowing = self._windowing.is_default() while input_stream.size() > 0: windowed_key_value = coder_impl.decode_from_stream(input_stream, True) key, value = windowed_key_value.value self._table[key_coder_impl.encode(key)].append( value if is_trivial_windowing else windowed_key_value. with_value(value)) def extend(self, input_buffer: Buffer) -> None: if isinstance(input_buffer, ListBuffer): # TODO(pabloem): GroupingBuffer will be removed once shuffling is done # via state. Remove this workaround along with that. return assert isinstance(input_buffer, GroupingBuffer), \ 'Input was not GroupingBuffer: %s' % input_buffer for key, values in input_buffer._table.items(): self._table[key].extend(values) def partition(self, n: int) -> list[list[bytes]]: """ It is used to partition _GroupingBuffer to N parts. Once it is partitioned, it would not be re-partitioned with diff N. Re-partition is not supported now. """ if not self._grouped_output: if self._windowing.is_default(): globally_window = GlobalWindows.windowed_value( None, timestamp=GlobalWindow().max_timestamp(), pane_info=windowed_value.PaneInfo( is_first=True, is_last=True, timing=windowed_value.PaneInfoTiming.ON_TIME, index=0, nonspeculative_index=0)).with_value windowed_key_values = lambda key, values: [ globally_window((key, values)) ] else: # TODO(pabloem, BEAM-7514): Trigger driver needs access to the clock # note that this only comes through if windowing is default - but what # about having multiple firings on the global window. # May need to revise. trigger_driver = trigger.create_trigger_driver(self._windowing, True) windowed_key_values = trigger_driver.process_entire_key coder_impl = self._post_grouped_coder.get_impl() key_coder_impl = self._key_coder.get_impl() self._grouped_output = [[] for _ in range(n)] output_stream_list = [create_OutputStream() for _ in range(n)] for idx, (encoded_key, windowed_values) in enumerate(self._table.items()): key = key_coder_impl.decode(encoded_key) for wkvs in windowed_key_values(key, windowed_values): coder_impl.encode_to_stream(wkvs, output_stream_list[idx % n], True) for ix, output_stream in enumerate(output_stream_list): self._grouped_output[ix] = [output_stream.get()] self._table.clear() return self._grouped_output def __iter__(self) -> Iterator[bytes]: """ Since partition() returns a list of lists, add this __iter__ to return a list to simplify code when we need to iterate through ALL elements of _GroupingBuffer. """ return itertools.chain(*self.partition(1)) # these should never be accessed, but they allow this class to meet the # PartionableBuffer protocol cleared = False def clear(self) -> None: pass def reset(self) -> None: pass class WindowGroupingBuffer(object): """Used to partition windowed side inputs.""" def __init__( self, access_pattern: beam_runner_api_pb2.FunctionSpec, coder: WindowedValueCoder) -> None: # Here's where we would use a different type of partitioning # (e.g. also by key) for a different access pattern. if access_pattern.urn == common_urns.side_inputs.ITERABLE.urn: self._kv_extractor = lambda value: ('', value) self._key_coder: coders.Coder = coders.SingletonCoder('') self._value_coder = coder.wrapped_value_coder elif access_pattern.urn == common_urns.side_inputs.MULTIMAP.urn: self._kv_extractor = lambda value: value self._key_coder = coder.wrapped_value_coder.key_coder() self._value_coder = (coder.wrapped_value_coder.value_coder()) else: raise ValueError("Unknown access pattern: '%s'" % access_pattern.urn) self._windowed_value_coder = coder self._window_coder = coder.window_coder self._values_by_window: collections.defaultdict[ tuple[str, BoundedWindow], list[Any]] = collections.defaultdict(list) def append(self, elements_data: bytes) -> None: input_stream = create_InputStream(elements_data) while input_stream.size() > 0: windowed_val_coder_impl: WindowedValueCoderImpl = ( self._windowed_value_coder.get_impl()) windowed_value = windowed_val_coder_impl.decode_from_stream( input_stream, True) key, value = self._kv_extractor(windowed_value.value) for window in windowed_value.windows: self._values_by_window[key, window].append(value) def encoded_items(self) -> Iterator[tuple[bytes, bytes, bytes, int]]: value_coder_impl = self._value_coder.get_impl() key_coder_impl = self._key_coder.get_impl() for (key, window), values in self._values_by_window.items(): encoded_window = self._window_coder.encode(window) encoded_key = key_coder_impl.encode_nested(key) output_stream = create_OutputStream() for value in values: value_coder_impl.encode_to_stream(value, output_stream, True) yield encoded_key, encoded_window, output_stream.get(), len(values) class GenericNonMergingWindowFn(window.NonMergingWindowFn): URN = 'internal-generic-non-merging' def __init__(self, coder: coders.Coder) -> None: self._coder = coder def assign( self, assign_context: window.WindowFn.AssignContext) -> Iterable[BoundedWindow]: raise NotImplementedError() def get_window_coder(self) -> coders.Coder: return self._coder @staticmethod @window.urns.RunnerApiFn.register_urn(URN, bytes) def from_runner_api_parameter( window_coder_id: bytes, context: Any) -> 'GenericNonMergingWindowFn': return GenericNonMergingWindowFn( context.coders[window_coder_id.decode('utf-8')]) QUEUE_KEY_TYPE = TypeVar('QUEUE_KEY_TYPE') class _ProcessingQueueManager(object): """Manages the queues for ProcessBundle inputs. There are three queues: - ready_inputs(_ProcessingQueueManager.KeyedQueue). This queue contains input data that is ready to be processed. These are data such as timers past their trigger time, and data to be processed. The ready_inputs_queue contains tuples of (stage_name, inputs), where inputs are dictionaries mapping PCollection name to data buffers. - watermark_pending_inputs(_ProcessingQueueManager.KeyedQueue). This queue contains input data that is not yet ready to be processed, and is blocked on the watermark advancing. ((stage_name, watermark), inputs), where the watermark is the watermark at which the inputs should be scheduled, and inputs are dictionaries mapping PCollection name to data buffers. - time_pending_inputs(_ProcessingQueueManager.KeyedQueue). This queue contains input data that is not yet ready to be processed, and is blocked on time advancing. ((stage_name, time), inputs), where the time is the real time point at which the inputs should be scheduled, and inputs are dictionaries mapping PCollection name to data buffers. """ class KeyedQueue(Generic[QUEUE_KEY_TYPE]): def __init__(self) -> None: self._q: collections.deque[tuple[QUEUE_KEY_TYPE, DataInput]] = collections.deque() self._keyed_elements: MutableMapping[QUEUE_KEY_TYPE, tuple[QUEUE_KEY_TYPE, DataInput]] = {} def enque(self, elm: tuple[QUEUE_KEY_TYPE, DataInput]) -> None: key = elm[0] incoming_inputs: DataInput = elm[1] if not incoming_inputs: return if key in self._keyed_elements: existing_inputs = self._keyed_elements[key][1] for pcoll in incoming_inputs.data: if incoming_inputs.data[pcoll] and existing_inputs.data.get(pcoll): existing_inputs.data[pcoll].extend(incoming_inputs.data[pcoll]) elif incoming_inputs.data[pcoll]: existing_inputs.data[pcoll] = incoming_inputs.data[pcoll] for timer_family in (incoming_inputs.timers or []): if (incoming_inputs.timers[timer_family] and existing_inputs.timers.get(timer_family)): existing_inputs.timers[timer_family].extend( incoming_inputs.timers[timer_family]) elif incoming_inputs.timers[timer_family]: existing_inputs.timers[timer_family] = incoming_inputs.timers[ timer_family] else: self._keyed_elements[key] = elm self._q.appendleft(elm) def deque(self) -> tuple[QUEUE_KEY_TYPE, DataInput]: elm = self._q.pop() key = elm[0] del self._keyed_elements[key] return elm def __len__(self) -> int: return len(self._q) def __repr__(self) -> str: return '<%s at 0x%x>' % (str(self), id(self)) def __str__(self) -> str: return '<%s len: %s %s>' % ( self.__class__.__name__, len(self), list(self._q)) def __init__(self) -> None: # For time-pending and watermark-pending inputs, the key type is # STAGE+TIMESTAMP, while for the ready inputs, the key type is only STAGE. self.time_pending_inputs = _ProcessingQueueManager.KeyedQueue[tuple[ str, Timestamp]]() self.watermark_pending_inputs = _ProcessingQueueManager.KeyedQueue[tuple[ str, Timestamp]]() self.ready_inputs = _ProcessingQueueManager.KeyedQueue[str]() def __str__(self) -> str: return '_ProcessingQueueManager(%s)' % self.__dict__ class GenericMergingWindowFn(window.WindowFn): URN = 'internal-generic-merging' TO_SDK_TRANSFORM = 'read' FROM_SDK_TRANSFORM = 'write' _HANDLES: dict[str, 'GenericMergingWindowFn'] = {} def __init__( self, execution_context: 'FnApiRunnerExecutionContext', windowing_strategy_proto: beam_runner_api_pb2.WindowingStrategy) -> None: self._worker_handler: Optional[worker_handlers.WorkerHandler] = None self._handle_id = handle_id = uuid.uuid4().hex self._HANDLES[handle_id] = self # ExecutionContexts are expensive, we don't want to keep them in the # static dictionary forever. Instead we hold a weakref and pop self # out of the dict once this context goes away. self._execution_context_ref_obj = weakref.ref( execution_context, lambda _: self._HANDLES.pop(handle_id, None)) self._windowing_strategy_proto = windowing_strategy_proto self._counter = 0 # Lazily created in make_process_bundle_descriptor() self._process_bundle_descriptor = None self._bundle_processor_id: str = '' self.windowed_input_coder_impl: Optional[CoderImpl] = None self.windowed_output_coder_impl: Optional[CoderImpl] = None def _execution_context_ref(self) -> 'FnApiRunnerExecutionContext': result = self._execution_context_ref_obj() assert result is not None return result def payload(self) -> bytes: return self._handle_id.encode('utf-8') @staticmethod @window.urns.RunnerApiFn.register_urn(URN, bytes) def from_runner_api_parameter( handle_id: bytes, unused_context: Any) -> 'GenericMergingWindowFn': return GenericMergingWindowFn._HANDLES[handle_id.decode('utf-8')] def assign( self, assign_context: window.WindowFn.AssignContext ) -> Iterable[window.BoundedWindow]: raise NotImplementedError() def merge(self, merge_context: window.WindowFn.MergeContext) -> None: worker_handler = self.worker_handle() assert self.windowed_input_coder_impl is not None assert self.windowed_output_coder_impl is not None process_bundle_id = self.uid('process') to_worker = worker_handler.data_conn.output_stream( process_bundle_id, self.TO_SDK_TRANSFORM) to_worker.write( self.windowed_input_coder_impl.encode_nested( window.GlobalWindows.windowed_value((b'', merge_context.windows)))) to_worker.close() process_bundle_req = beam_fn_api_pb2.InstructionRequest( instruction_id=process_bundle_id, process_bundle=beam_fn_api_pb2.ProcessBundleRequest( process_bundle_descriptor_id=self._bundle_processor_id)) result_future = worker_handler.control_conn.push(process_bundle_req) for output in worker_handler.data_conn.input_elements( process_bundle_id, [self.FROM_SDK_TRANSFORM], abort_callback=lambda: bool(result_future.is_done() and result_future. get().error)): if isinstance(output, beam_fn_api_pb2.Elements.Data): windowed_result = self.windowed_output_coder_impl.decode_nested( output.data) for merge_result, originals in windowed_result.value[1][1]: merge_context.merge(originals, merge_result) else: raise RuntimeError("Unexpected data: %s" % output) result = result_future.get() if result.error: raise RuntimeError(result.error) # The result was "returned" via the merge callbacks on merge_context above. def get_window_coder(self) -> coders.Coder: return self._execution_context_ref().pipeline_context.coders[ self._windowing_strategy_proto.window_coder_id] def worker_handle(self) -> 'worker_handlers.WorkerHandler': if self._worker_handler is None: worker_handler_manager = self._execution_context_ref( ).worker_handler_manager self._worker_handler = worker_handler_manager.get_worker_handlers( self._windowing_strategy_proto.environment_id, 1)[0] process_bundle_decriptor = self.make_process_bundle_descriptor( self._worker_handler.data_api_service_descriptor(), self._worker_handler.state_api_service_descriptor()) worker_handler_manager.register_process_bundle_descriptor( process_bundle_decriptor) return self._worker_handler def make_process_bundle_descriptor( self, data_api_service_descriptor: Optional[endpoints_pb2.ApiServiceDescriptor], state_api_service_descriptor: Optional[endpoints_pb2.ApiServiceDescriptor] ) -> beam_fn_api_pb2.ProcessBundleDescriptor: """Creates a ProcessBundleDescriptor for invoking the WindowFn's merge operation. """ def make_channel_payload(coder_id: str) -> bytes: data_spec = beam_fn_api_pb2.RemoteGrpcPort(coder_id=coder_id) if data_api_service_descriptor: data_spec.api_service_descriptor.url = (data_api_service_descriptor.url) return data_spec.SerializeToString() pipeline_context = self._execution_context_ref().pipeline_context global_windowing_strategy_id = self.uid('global_windowing_strategy') global_windowing_strategy_proto = core.Windowing( window.GlobalWindows()).to_runner_api(pipeline_context) coders = dict(pipeline_context.coders.get_id_to_proto_map()) def make_coder(urn: str, *components: str) -> str: coder_proto = beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec(urn=urn), component_coder_ids=components) coder_id = self.uid('coder') coders[coder_id] = coder_proto pipeline_context.coders.put_proto(coder_id, coder_proto) return coder_id bytes_coder_id = make_coder(common_urns.coders.BYTES.urn) window_coder_id = self._windowing_strategy_proto.window_coder_id global_window_coder_id = make_coder(common_urns.coders.GLOBAL_WINDOW.urn) iter_window_coder_id = make_coder( common_urns.coders.ITERABLE.urn, window_coder_id) input_coder_id = make_coder( common_urns.coders.KV.urn, bytes_coder_id, iter_window_coder_id) output_coder_id = make_coder( common_urns.coders.KV.urn, bytes_coder_id, make_coder( common_urns.coders.KV.urn, iter_window_coder_id, make_coder( common_urns.coders.ITERABLE.urn, make_coder( common_urns.coders.KV.urn, window_coder_id, iter_window_coder_id)))) windowed_input_coder_id = make_coder( common_urns.coders.WINDOWED_VALUE.urn, input_coder_id, global_window_coder_id) windowed_output_coder_id = make_coder( common_urns.coders.WINDOWED_VALUE.urn, output_coder_id, global_window_coder_id) self.windowed_input_coder_impl = pipeline_context.coders[ windowed_input_coder_id].get_impl() self.windowed_output_coder_impl = pipeline_context.coders[ windowed_output_coder_id].get_impl() self._bundle_processor_id = self.uid('merge_windows') return beam_fn_api_pb2.ProcessBundleDescriptor( id=self._bundle_processor_id, transforms={ self.TO_SDK_TRANSFORM: beam_runner_api_pb2.PTransform( unique_name='MergeWindows/Read', spec=beam_runner_api_pb2.FunctionSpec( urn=bundle_processor.DATA_INPUT_URN, payload=make_channel_payload(windowed_input_coder_id)), outputs={'input': 'input'}), 'Merge': beam_runner_api_pb2.PTransform( unique_name='MergeWindows/Merge', spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.primitives.MERGE_WINDOWS.urn, payload=self._windowing_strategy_proto.window_fn. SerializeToString()), inputs={'input': 'input'}, outputs={'output': 'output'}), self.FROM_SDK_TRANSFORM: beam_runner_api_pb2.PTransform( unique_name='MergeWindows/Write', spec=beam_runner_api_pb2.FunctionSpec( urn=bundle_processor.DATA_OUTPUT_URN, payload=make_channel_payload(windowed_output_coder_id)), inputs={'output': 'output'}), }, pcollections={ 'input': beam_runner_api_pb2.PCollection( unique_name='input', windowing_strategy_id=global_windowing_strategy_id, coder_id=input_coder_id), 'output': beam_runner_api_pb2.PCollection( unique_name='output', windowing_strategy_id=global_windowing_strategy_id, coder_id=output_coder_id), }, coders=coders, windowing_strategies={ global_windowing_strategy_id: global_windowing_strategy_proto, }, environments=dict( self._execution_context_ref().pipeline_components.environments. items()), state_api_service_descriptor=state_api_service_descriptor, timer_api_service_descriptor=data_api_service_descriptor) def uid(self, name: str = '') -> str: self._counter += 1 return '%s_%s_%s' % (self._handle_id, name, self._counter) class FnApiRunnerExecutionContext(object): """ :var pcoll_buffers: (dict): Mapping of PCollection IDs to list that functions as buffer for the ``beam.PCollection``. """ def __init__( self, stages: list[translations.Stage], worker_handler_manager: 'worker_handlers.WorkerHandlerManager', pipeline_components: beam_runner_api_pb2.Components, safe_coders: translations.SafeCoderMapping, data_channel_coders: dict[str, str], num_workers: int, uses_teststream: bool = False, split_managers: Sequence[tuple[str, Callable[[int], Iterable[float]]]] = () ) -> None: """ :param worker_handler_manager: This class manages the set of worker handlers, and the communication with state / control APIs. :param pipeline_components: (beam_runner_api_pb2.Components) :param safe_coders: A map from Coder ID to Safe Coder ID. :param data_channel_coders: A map from PCollection ID to the ID of the Coder for that PCollection. """ self.stages = {s.name: s for s in stages} self.side_input_descriptors_by_stage = ( self._build_data_side_inputs_map(stages)) self.pcoll_buffers: MutableMapping[bytes, PartitionableBuffer] = {} self.timer_buffers: MutableMapping[bytes, ListBuffer] = {} self.worker_handler_manager = worker_handler_manager self.pipeline_components = pipeline_components self.safe_coders = safe_coders self.data_channel_coders = data_channel_coders self.num_workers = num_workers self.split_managers = split_managers # TODO(pabloem): Move Clock classes out of DirectRunner and into FnApiRnr self.clock: Union[TestClock, RealClock] = ( TestClock() if uses_teststream else RealClock()) self.queues = _ProcessingQueueManager() # The following set of dictionaries hold information mapping relationships # between various pipeline elements. self.input_transform_to_buffer_id: MutableMapping[str, bytes] = {} self.pcollection_to_producer_transform: MutableMapping[Union[str, bytes], Optional[str]] = {} # Map of buffer_id to its consumers. A consumer is the pair of # Stage name + Ptransform name that consume that buffer. self.buffer_id_to_consumer_pairs: dict[bytes, set[tuple[str, str]]] = {} self._compute_pipeline_dictionaries() self.watermark_manager = WatermarkManager(stages) # from apache_beam.runners.portability.fn_api_runner import \ # visualization_tools # visualization_tools.show_watermark_manager(self.watermark_manager) self.pipeline_context = pipeline_context.PipelineContext( self.pipeline_components, iterable_state_write=self._iterable_state_write) self._last_uid = -1 self.safe_windowing_strategies = { id: self._make_safe_windowing_strategy(id) for id in self.pipeline_components.windowing_strategies.keys() } self._stage_managers: dict[str, BundleContextManager] = {} def bundle_manager_for( self, stage: Stage, num_workers: Optional[int] = None) -> 'BundleContextManager': if stage.name not in self._stage_managers: self._stage_managers[stage.name] = BundleContextManager( self, stage, num_workers or self.num_workers, split_managers=self.split_managers) return self._stage_managers[stage.name] def _compute_pipeline_dictionaries(self) -> None: for s in self.stages.values(): for t in s.transforms: buffer_id = t.spec.payload if t.spec.urn == bundle_processor.DATA_INPUT_URN: self.input_transform_to_buffer_id[t.unique_name] = buffer_id if t.spec.payload == translations.IMPULSE_BUFFER: # Impulse data is not produced by any PTransform. self.pcollection_to_producer_transform[ translations.IMPULSE_BUFFER] = None else: assert t.spec.payload != translations.IMPULSE_BUFFER _, input_pcoll = split_buffer_id(buffer_id) # Adding PCollections that may not have a producer. # This is necessary, for example, for the case where we pass an # empty list of PCollections into a Flatten transform. if input_pcoll not in self.pcollection_to_producer_transform: self.pcollection_to_producer_transform[input_pcoll] = None if buffer_id not in self.buffer_id_to_consumer_pairs: self.buffer_id_to_consumer_pairs[buffer_id] = set() if (s.name, t.unique_name ) not in self.buffer_id_to_consumer_pairs[buffer_id]: self.buffer_id_to_consumer_pairs[buffer_id].add( (s.name, t.unique_name)) elif t.spec.urn == bundle_processor.DATA_OUTPUT_URN: _, output_pcoll = split_buffer_id(buffer_id) self.pcollection_to_producer_transform[output_pcoll] = t.unique_name elif t.spec.urn in translations.PAR_DO_URNS: pass def setup(self) -> None: """This sets up the pipeline to begin running. 1. This function enqueues all initial pipeline bundles to be executed. 2. It also updates payload fields on DATA_INPUT and DATA_OUTPUT operations to the Data API endpoints that are live. """ for stage in self.stages.values(): self._enqueue_stage_initial_inputs(stage) def _enqueue_stage_initial_inputs(self, stage: Stage) -> None: """Sets up IMPULSE inputs for a stage, and the data GRPC API endpoint.""" data_input: MutableMapping[str, PartitionableBuffer] = {} ready_to_schedule = True for transform in stage.transforms: if (transform.spec.urn in {bundle_processor.DATA_INPUT_URN, bundle_processor.DATA_OUTPUT_URN}): if transform.spec.urn == bundle_processor.DATA_INPUT_URN: coder_id = self.data_channel_coders[only_element( transform.outputs.values())] coder = self.pipeline_context.coders[self.safe_coders.get( coder_id, coder_id)] if transform.spec.payload == translations.IMPULSE_BUFFER: data_input[transform.unique_name] = ListBuffer(coder.get_impl()) data_input[transform.unique_name].append(ENCODED_IMPULSE_VALUE) else: # If this is not an IMPULSE input, then it is not part of the # initial inputs of a pipeline, and we'll ignore it. pass else: coder_id = self.data_channel_coders[only_element( transform.inputs.values())] # For every DATA_INPUT or DATA_OUTPUT operation, we need to replace the # payload with the GRPC configuration for the Data channel. bundle_manager = self.bundle_manager_for(stage) data_spec = beam_fn_api_pb2.RemoteGrpcPort(coder_id=coder_id) data_api_service_descriptor = ( bundle_manager.data_api_service_descriptor()) if data_api_service_descriptor: data_spec.api_service_descriptor.url = ( data_api_service_descriptor.url) transform.spec.payload = data_spec.SerializeToString() elif transform.spec.urn in translations.PAR_DO_URNS: payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.ParDoPayload) if payload.side_inputs: # If the stage needs side inputs, then it's not ready to be # executed. ready_to_schedule = False if data_input and ready_to_schedule: # We push the data inputs, along with the name of the consuming stage. _LOGGER.debug('Scheduling bundle in stage for execution: %s', stage.name) self.queues.ready_inputs.enque((stage.name, DataInput(data_input, {}))) elif data_input and not ready_to_schedule: _LOGGER.debug( 'Enqueuing stage pending watermark. Stage name: %s', stage.name) self.queues.watermark_pending_inputs.enque( ((stage.name, MAX_TIMESTAMP), DataInput(data_input, {}))) @staticmethod def _build_data_side_inputs_map( stages: Iterable[translations.Stage] ) -> MutableMapping[str, 'DataSideInput']: """Builds an index mapping stages to side input descriptors. A side input descriptor is a map of side input IDs to side input access patterns for all of the outputs of a stage that will be consumed as a side input. """ transform_consumers: collections.defaultdict[ str, list[beam_runner_api_pb2.PTransform]] = collections.defaultdict(list) stage_consumers: collections.defaultdict[ str, list[translations.Stage]] = collections.defaultdict(list) def get_all_side_inputs() -> set[str]: all_side_inputs: set[str] = set() for stage in stages: for transform in stage.transforms: for input in transform.inputs.values(): transform_consumers[input].append(transform) stage_consumers[input].append(stage) for si in stage.side_inputs(): all_side_inputs.add(si) return all_side_inputs all_side_inputs = frozenset(get_all_side_inputs()) data_side_inputs_by_producing_stage: dict[str, DataSideInput] = {} producing_stages_by_pcoll = {} for s in stages: data_side_inputs_by_producing_stage[s.name] = {} for transform in s.transforms: for o in transform.outputs.values(): if o in s.side_inputs(): continue if o in producing_stages_by_pcoll: continue producing_stages_by_pcoll[o] = s for side_pc in all_side_inputs: for consuming_transform in transform_consumers[side_pc]: if consuming_transform.spec.urn not in translations.PAR_DO_URNS: continue producing_stage = producing_stages_by_pcoll[side_pc] payload = proto_utils.parse_Bytes( consuming_transform.spec.payload, beam_runner_api_pb2.ParDoPayload) for si_tag in payload.side_inputs: if consuming_transform.inputs[si_tag] == side_pc: side_input_id = (consuming_transform.unique_name, si_tag) data_side_inputs_by_producing_stage[ producing_stage.name][side_input_id] = ( translations.create_buffer_id(side_pc), payload.side_inputs[si_tag].access_pattern) return data_side_inputs_by_producing_stage def _make_safe_windowing_strategy(self, id: str) -> str: windowing_strategy_proto = self.pipeline_components.windowing_strategies[id] if windowing_strategy_proto.window_fn.urn in SAFE_WINDOW_FNS: return id else: safe_id = id + '_safe' while safe_id in self.pipeline_components.windowing_strategies: safe_id += '_' safe_proto = copy.copy(windowing_strategy_proto) if (windowing_strategy_proto.merge_status == beam_runner_api_pb2.MergeStatus.NON_MERGING): safe_proto.window_fn.urn = GenericNonMergingWindowFn.URN safe_proto.window_fn.payload = ( windowing_strategy_proto.window_coder_id.encode('utf-8')) elif (windowing_strategy_proto.merge_status == beam_runner_api_pb2.MergeStatus.NEEDS_MERGE): window_fn = GenericMergingWindowFn(self, windowing_strategy_proto) safe_proto.window_fn.urn = GenericMergingWindowFn.URN safe_proto.window_fn.payload = window_fn.payload() else: raise NotImplementedError( 'Unsupported merging strategy: %s' % windowing_strategy_proto.merge_status) self.pipeline_context.windowing_strategies.put_proto(safe_id, safe_proto) return safe_id @property def state_servicer(self) -> 'worker_handlers.StateServicer': # TODO(BEAM-9625): Ensure FnApiRunnerExecutionContext owns StateServicer return self.worker_handler_manager.state_servicer def next_uid(self) -> str: self._last_uid += 1 return str(self._last_uid) def _iterable_state_write( self, values: Iterable, element_coder_impl: CoderImpl) -> bytes: token = unique_name(None, 'iter').encode('ascii') out = create_OutputStream() for element in values: element_coder_impl.encode_to_stream(element, out, True) self.worker_handler_manager.state_servicer.append_raw( beam_fn_api_pb2.StateKey( runner=beam_fn_api_pb2.StateKey.Runner(key=token)), out.get()) return token def commit_side_inputs_to_state( self, data_side_input: 'DataSideInput', ) -> None: for (consuming_transform_id, tag), (buffer_id, func_spec) in data_side_input.items(): _, pcoll_id = split_buffer_id(buffer_id) value_coder = self.pipeline_context.coders[self.safe_coders[ self.data_channel_coders[pcoll_id]]] elements_by_window = WindowGroupingBuffer(func_spec, value_coder) if buffer_id not in self.pcoll_buffers: self.pcoll_buffers[buffer_id] = ListBuffer( coder_impl=value_coder.get_impl()) for element_data in self.pcoll_buffers[buffer_id]: elements_by_window.append(element_data) if func_spec.urn == common_urns.side_inputs.ITERABLE.urn: for _, window, elements_data, _ in elements_by_window.encoded_items(): state_key = beam_fn_api_pb2.StateKey( iterable_side_input=beam_fn_api_pb2.StateKey.IterableSideInput( transform_id=consuming_transform_id, side_input_id=tag, window=window)) self.state_servicer.append_raw(state_key, elements_data) elif func_spec.urn == common_urns.side_inputs.MULTIMAP.urn: # TODO(robertwb): Consider computing these lazily on demand rather than # anticipating all potentail state requests which will be more cpu and # memory efficient for large side inputs. for (key, window, elements_data, elements_count) in elements_by_window.encoded_items(): state_key = beam_fn_api_pb2.StateKey( multimap_side_input=beam_fn_api_pb2.StateKey.MultimapSideInput( transform_id=consuming_transform_id, side_input_id=tag, window=window, key=key)) self.state_servicer.append_raw(state_key, elements_data) key_iter_state_key = beam_fn_api_pb2.StateKey( multimap_keys_side_input=beam_fn_api_pb2.StateKey. MultimapKeysSideInput( transform_id=consuming_transform_id, side_input_id=tag, window=window)) self.state_servicer.append_raw(key_iter_state_key, key) kv_iter_state_key = beam_fn_api_pb2.StateKey( multimap_keys_values_side_input=beam_fn_api_pb2.StateKey. MultimapKeysValuesSideInput( transform_id=consuming_transform_id, side_input_id=tag, window=window)) self.state_servicer.append_raw( kv_iter_state_key, # KV> encoding. key + struct.pack('>i', elements_count) + elements_data) else: raise ValueError("Unknown access pattern: '%s'" % func_spec.urn) class BundleContextManager(object): def __init__( self, execution_context: FnApiRunnerExecutionContext, stage: translations.Stage, num_workers: int, split_managers: Sequence[tuple[str, Callable[[int], Iterable[float]]]], ) -> None: self.execution_context = execution_context self.stage = stage self.bundle_uid = self.execution_context.next_uid() self.num_workers = num_workers self.split_managers = split_managers # Properties that are lazily initialized self._process_bundle_descriptor: Optional[ beam_fn_api_pb2.ProcessBundleDescriptor] = None self._worker_handlers: Optional[list[worker_handlers.WorkerHandler]] = None # a mapping of {(transform_id, timer_family_id): timer_coder_id}. The map # is built after self._process_bundle_descriptor is initialized. # This field can be used to tell whether current bundle has timers. self._timer_coder_ids: Optional[dict[tuple[str, str], str]] = None # A mapping from transform_name to Buffer ID self.stage_data_outputs: DataOutput = {} # A mapping of {(transform_id, timer_family_id) : buffer_id} self.stage_timer_outputs: OutputTimers = {} self._compute_expected_outputs() def _compute_expected_outputs(self) -> None: for transform in self.stage.transforms: if transform.spec.urn == bundle_processor.DATA_OUTPUT_URN: buffer_id = transform.spec.payload self.stage_data_outputs[transform.unique_name] = buffer_id elif transform.spec.urn in translations.PAR_DO_URNS: payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.ParDoPayload) for timer_family_id in payload.timer_family_specs.keys(): time_domain = payload.timer_family_specs[timer_family_id].time_domain self.stage_timer_outputs[(transform.unique_name, timer_family_id)] = ( create_buffer_id(timer_family_id, 'timers'), time_domain) @property def worker_handlers(self) -> list['worker_handlers.WorkerHandler']: if self._worker_handlers is None: self._worker_handlers = ( self.execution_context.worker_handler_manager.get_worker_handlers( self.stage.environment, self.num_workers)) return self._worker_handlers def data_api_service_descriptor( self) -> Optional[endpoints_pb2.ApiServiceDescriptor]: # All worker_handlers share the same grpc server, so we can read grpc server # info from any worker_handler and read from the first worker_handler. return self.worker_handlers[0].data_api_service_descriptor() def state_api_service_descriptor( self) -> Optional[endpoints_pb2.ApiServiceDescriptor]: # All worker_handlers share the same grpc server, so we can read grpc server # info from any worker_handler and read from the first worker_handler. return self.worker_handlers[0].state_api_service_descriptor() @property def process_bundle_descriptor( self) -> beam_fn_api_pb2.ProcessBundleDescriptor: if self._process_bundle_descriptor is None: self._process_bundle_descriptor = self._build_process_bundle_descriptor() self._timer_coder_ids = self._build_timer_coders_id_map() return self._process_bundle_descriptor def _build_process_bundle_descriptor( self) -> beam_fn_api_pb2.ProcessBundleDescriptor: # Cannot be invoked until *after* _extract_endpoints is called. # Always populate the timer_api_service_descriptor. return beam_fn_api_pb2.ProcessBundleDescriptor( id=self.bundle_uid, transforms={ transform.unique_name: transform for transform in self.stage.transforms }, pcollections=dict( self.execution_context.pipeline_components.pcollections.items()), coders=dict(self.execution_context.pipeline_components.coders.items()), windowing_strategies=dict( self.execution_context.pipeline_components.windowing_strategies. items()), environments=dict( self.execution_context.pipeline_components.environments.items()), state_api_service_descriptor=self.state_api_service_descriptor(), timer_api_service_descriptor=self.data_api_service_descriptor()) def get_input_coder_impl(self, transform_id: str) -> CoderImpl: coder_id = beam_fn_api_pb2.RemoteGrpcPort.FromString( self.process_bundle_descriptor.transforms[transform_id].spec.payload ).coder_id assert coder_id return self.get_coder_impl(coder_id) def _build_timer_coders_id_map(self) -> dict[tuple[str, str], str]: assert self._process_bundle_descriptor is not None timer_coder_ids = {} for transform_id, transform_proto in (self._process_bundle_descriptor .transforms.items()): if transform_proto.spec.urn == common_urns.primitives.PAR_DO.urn: pardo_payload = proto_utils.parse_Bytes( transform_proto.spec.payload, beam_runner_api_pb2.ParDoPayload) for id, timer_family_spec in pardo_payload.timer_family_specs.items(): timer_coder_ids[(transform_id, id)] = ( timer_family_spec.timer_family_coder_id) return timer_coder_ids def get_coder_impl(self, coder_id: str) -> CoderImpl: if coder_id in self.execution_context.safe_coders: return self.execution_context.pipeline_context.coders[ self.execution_context.safe_coders[coder_id]].get_impl() else: return self.execution_context.pipeline_context.coders[coder_id].get_impl() def get_timer_coder_impl( self, transform_id: str, timer_family_id: str) -> CoderImpl: assert self._timer_coder_ids is not None return self.get_coder_impl( self._timer_coder_ids[(transform_id, timer_family_id)]) def get_buffer( self, buffer_id: bytes, transform_id: str) -> PartitionableBuffer: """Returns the buffer for a given (operation_type, PCollection ID). For grouping-typed operations, we produce a ``GroupingBuffer``. For others, we produce a ``ListBuffer``. """ kind, name = split_buffer_id(buffer_id) if kind == 'materialize': if buffer_id not in self.execution_context.pcoll_buffers: self.execution_context.pcoll_buffers[buffer_id] = ListBuffer( coder_impl=self.get_input_coder_impl(transform_id)) return self.execution_context.pcoll_buffers[buffer_id] # For timer buffer, name = timer_family_id elif kind == 'timers': if buffer_id not in self.execution_context.timer_buffers: timer_coder_impl = self.get_timer_coder_impl(transform_id, name) self.execution_context.timer_buffers[buffer_id] = ListBuffer( timer_coder_impl) return self.execution_context.timer_buffers[buffer_id] elif kind == 'group': # This is a grouping write, create a grouping buffer if needed. if buffer_id not in self.execution_context.pcoll_buffers: original_gbk_transform = name transform_proto = self.execution_context.pipeline_components.transforms[ original_gbk_transform] input_pcoll = only_element(list(transform_proto.inputs.values())) output_pcoll = only_element(list(transform_proto.outputs.values())) pre_gbk_coder = self.execution_context.pipeline_context.coders[ self.execution_context.safe_coders[ self.execution_context.data_channel_coders[input_pcoll]]] post_gbk_coder = self.execution_context.pipeline_context.coders[ self.execution_context.safe_coders[ self.execution_context.data_channel_coders[output_pcoll]]] windowing_strategy = ( self.execution_context.pipeline_context.windowing_strategies[ self.execution_context.safe_windowing_strategies[ self.execution_context.pipeline_components. pcollections[input_pcoll].windowing_strategy_id]]) self.execution_context.pcoll_buffers[buffer_id] = GroupingBuffer( pre_gbk_coder, post_gbk_coder, windowing_strategy) else: # These should be the only two identifiers we produce for now, # but special side input writes may go here. raise NotImplementedError(buffer_id) return self.execution_context.pcoll_buffers[buffer_id] def input_for(self, transform_id: str, input_id: str) -> str: """Returns the name of the transform producing the given PCollection.""" input_pcoll = self.process_bundle_descriptor.transforms[ transform_id].inputs[input_id] for read_id, proto in self.process_bundle_descriptor.transforms.items(): # The GrpcRead is followed by the SDF/Process. if (proto.spec.urn == bundle_processor.DATA_INPUT_URN and input_pcoll in proto.outputs.values()): return read_id # The GrpcRead is followed by the SDF/Truncate -> SDF/Process. if (proto.spec.urn == common_urns.sdf_components.TRUNCATE_SIZED_RESTRICTION.urn and input_pcoll in proto.outputs.values()): read_input = list( self.process_bundle_descriptor.transforms[read_id].inputs.values() )[0] for (grpc_read, transform_proto) in self.process_bundle_descriptor.transforms.items(): # pylint: disable=line-too-long if (transform_proto.spec.urn == bundle_processor.DATA_INPUT_URN and read_input in transform_proto.outputs.values()): return grpc_read raise RuntimeError('No IO transform feeds %s' % transform_id) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/fn_api_runner/fn_runner.py0000644000175100017510000020537015206113032027115 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A PipelineRunner using the SDK harness. """ # pytype: skip-file # mypy: check-untyped-defs import contextlib import copy import itertools import json import logging import multiprocessing import os import subprocess import sys import threading import time from typing import Callable from typing import Iterable from typing import Iterator from typing import Mapping from typing import MutableMapping from typing import Optional from typing import TypeVar from typing import Union from apache_beam.coders.coder_impl import create_OutputStream from apache_beam.metrics import metric from apache_beam.metrics import monitoring_infos from apache_beam.metrics.execution import MetricResult from apache_beam.metrics.execution import MetricsContainer from apache_beam.metrics.metricbase import MetricName from apache_beam.metrics.monitoring_infos import consolidate as consolidate_monitoring_infos from apache_beam.options import pipeline_options from apache_beam.options.value_provider import RuntimeValueProvider from apache_beam.pipeline import Pipeline from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_provision_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.portability.api import metrics_pb2 from apache_beam.runners import runner from apache_beam.runners.pipeline_utils import group_by_key_input_visitor from apache_beam.runners.pipeline_utils import merge_common_environments from apache_beam.runners.pipeline_utils import merge_superset_dep_environments from apache_beam.runners.pipeline_utils import validate_pipeline_graph from apache_beam.runners.portability import portable_metrics from apache_beam.runners.portability.fn_api_runner import execution from apache_beam.runners.portability.fn_api_runner import translations from apache_beam.runners.portability.fn_api_runner.execution import ListBuffer from apache_beam.runners.portability.fn_api_runner.translations import BundleProcessResult from apache_beam.runners.portability.fn_api_runner.translations import DataInput from apache_beam.runners.portability.fn_api_runner.translations import DataOutput from apache_beam.runners.portability.fn_api_runner.translations import OutputTimerData from apache_beam.runners.portability.fn_api_runner.translations import OutputTimers from apache_beam.runners.portability.fn_api_runner.translations import create_buffer_id from apache_beam.runners.portability.fn_api_runner.translations import only_element from apache_beam.runners.portability.fn_api_runner.worker_handlers import WorkerHandler from apache_beam.runners.portability.fn_api_runner.worker_handlers import WorkerHandlerManager from apache_beam.runners.worker import bundle_processor from apache_beam.transforms import environments from apache_beam.utils import proto_utils from apache_beam.utils import thread_pool_executor from apache_beam.utils import timestamp from apache_beam.utils.profiler import Profile _LOGGER = logging.getLogger(__name__) _BUNDLE_LOGGER = logging.getLogger(__name__ + ".run_bundle") T = TypeVar('T') # This module is experimental. No backwards-compatibility guarantees. class FnApiRunner(runner.PipelineRunner): NUM_FUSED_STAGES_COUNTER = "__num_fused_stages" def __init__( self, default_environment: Optional[environments.Environment] = None, bundle_repeat: int = 0, use_state_iterables: bool = False, provision_info: Optional['ExtendedProvisionInfo'] = None, progress_request_frequency: Optional[float] = None, is_drain: bool = False) -> None: """Creates a new Fn API Runner. Args: default_environment: the default environment to use for UserFns. bundle_repeat: replay every bundle this many extra times, for profiling and debugging use_state_iterables: Intentionally split gbk iterables over state API (for testing) provision_info: provisioning info to make available to workers, or None progress_request_frequency: The frequency (in seconds) that the runner waits before requesting progress from the SDK. is_drain: identify whether expand the sdf graph in the drain mode. """ super().__init__() self._default_environment = ( default_environment or environments.EmbeddedPythonEnvironment.default()) self._bundle_repeat = bundle_repeat self._num_workers = 1 self._progress_frequency = progress_request_frequency self._profiler_factory: Optional[Callable[..., Profile]] = None self._use_state_iterables = use_state_iterables self._is_drain = is_drain self._provision_info = provision_info or ExtendedProvisionInfo( beam_provision_api_pb2.ProvisionInfo( retrieval_token='unused-retrieval-token')) @staticmethod def supported_requirements() -> tuple[str, ...]: return ( common_urns.requirements.REQUIRES_STATEFUL_PROCESSING.urn, common_urns.requirements.REQUIRES_BUNDLE_FINALIZATION.urn, common_urns.requirements.REQUIRES_SPLITTABLE_DOFN.urn, ) def run_pipeline( self, pipeline: Pipeline, options: pipeline_options.PipelineOptions) -> 'RunnerResult': RuntimeValueProvider.set_runtime_options({}) # Setup "beam_fn_api" experiment options if lacked. experiments = ( options.view_as(pipeline_options.DebugOptions).experiments or []) if not 'beam_fn_api' in experiments: experiments.append('beam_fn_api') options.view_as(pipeline_options.DebugOptions).experiments = experiments # This is sometimes needed if type checking is disabled # to enforce that the inputs (and outputs) of GroupByKey operations # are known to be KVs. pipeline.visit( group_by_key_input_visitor( not options.view_as(pipeline_options.TypeOptions). allow_non_deterministic_key_coders)) self._bundle_repeat = self._bundle_repeat or options.view_as( pipeline_options.DirectOptions).direct_runner_bundle_repeat pipeline_direct_num_workers = options.view_as( pipeline_options.DirectOptions).direct_num_workers if pipeline_direct_num_workers == 0: self._num_workers = multiprocessing.cpu_count() else: self._num_workers = pipeline_direct_num_workers or self._num_workers # set direct workers running mode if it is defined with pipeline options. running_mode = \ options.view_as(pipeline_options.DirectOptions).direct_running_mode if running_mode == 'multi_threading': self._default_environment = ( environments.EmbeddedPythonGrpcEnvironment.default()) elif running_mode == 'multi_processing': command_string = '%s -m apache_beam.runners.worker.sdk_worker_main' \ % sys.executable self._default_environment = ( environments.SubprocessSDKEnvironment.from_command_string( command_string=command_string)) if running_mode == 'in_memory' and self._num_workers != 1: _LOGGER.warning( 'If direct_num_workers is not equal to 1, direct_running_mode ' 'should be `multi_processing` or `multi_threading` instead of ' '`in_memory` in order for it to have the desired worker parallelism ' 'effect. direct_num_workers: %d ; running_mode: %s', self._num_workers, running_mode) self._profiler_factory = Profile.factory_from_options( options.view_as(pipeline_options.ProfilingOptions)) self._latest_run_result = self.run_via_runner_api( pipeline.to_runner_api(default_environment=self._default_environment), options) return self._latest_run_result def run_via_runner_api( self, pipeline_proto: beam_runner_api_pb2.Pipeline, options: pipeline_options.PipelineOptions) -> 'RunnerResult': validate_pipeline_graph(pipeline_proto) self._validate_requirements(pipeline_proto) self._check_requirements(pipeline_proto) direct_options = options.view_as(pipeline_options.DirectOptions) test_splits = ( json.loads(direct_options.direct_test_splits) if isinstance( direct_options.direct_test_splits, str) else direct_options.direct_test_splits) self._split_managers = _split_managers_from_context + [ (stage, create_test_split_manager(**data)) for (stage, data) in test_splits.items() ] if direct_options.direct_embed_docker_python: pipeline_proto = self.embed_default_docker_image(pipeline_proto) pipeline_proto = merge_common_environments( self.resolve_any_environments( merge_superset_dep_environments(pipeline_proto))) stage_context, stages = self.create_stages(pipeline_proto) return self.run_stages(stage_context, stages) def embed_default_docker_image(self, pipeline_proto): """Updates the pipeline proto to execute transforms that would normally be executed in the default docker image for this SDK to execute inline via the "embedded" environment. """ # Context is unused for these types. embedded_env = environments.EmbeddedPythonEnvironment.default( ).to_runner_api(None) # type: ignore[arg-type] docker_env = environments.DockerEnvironment.from_container_image( environments.DockerEnvironment.default_docker_image()).to_runner_api( None) # type: ignore[arg-type] # We'd rather deal with the complexity of any environments here rather # than resolve them first so we can get optimal substitution in case # docker is not high in the preferred environment type list. def is_this_python_docker_env(env): return any( e == docker_env for e in environments.expand_anyof_environments(env)) python_docker_environments = set( env_id for (env_id, env) in pipeline_proto.components.environments.items() if is_this_python_docker_env(env)) if not python_docker_environments: # No matching docker environments. return pipeline_proto for env_id, env in pipeline_proto.components.environments.items(): if env.urn == embedded_env.urn: embedded_env_id = env_id break else: # No existing embedded environment. Create one. embedded_env_id = "python_embedded_env" while embedded_env_id in pipeline_proto.components.environments: embedded_env_id += '_' pipeline_proto.components.environments[embedded_env_id].CopyFrom( embedded_env) for transform in pipeline_proto.components.transforms.values(): if transform.environment_id in python_docker_environments: transform.environment_id = embedded_env_id return pipeline_proto def resolve_any_environments(self, pipeline_proto): for env_id, env in pipeline_proto.components.environments.items(): pipeline_proto.components.environments[env_id].CopyFrom( environments.resolve_anyof_environment( env, python_urns.EMBEDDED_PYTHON, common_urns.environments.EXTERNAL.urn, common_urns.environments.DOCKER.urn)) return pipeline_proto @contextlib.contextmanager def maybe_profile(self) -> Iterator[None]: if self._profiler_factory: try: profile_id = 'direct-' + subprocess.check_output([ 'git', 'rev-parse', '--abbrev-ref', 'HEAD' ]).decode(errors='ignore').strip() except subprocess.CalledProcessError: profile_id = 'direct-unknown' profiler: Optional[Profile] = self._profiler_factory( profile_id, time_prefix='') else: profiler = None if profiler: with profiler: yield if not self._bundle_repeat: _LOGGER.warning( 'The --direct_runner_bundle_repeat option is not set; ' 'a significant portion of the profile may be one-time overhead.') path = profiler.profile_output print('CPU Profile written to %s' % path) try: import gprof2dot # pylint: disable=unused-import if not subprocess.call([sys.executable, '-m', 'gprof2dot', '-f', 'pstats', path, '-o', path + '.dot']): if not subprocess.call( ['dot', '-Tsvg', '-o', path + '.svg', path + '.dot']): print( 'CPU Profile rendering at file://%s.svg' % os.path.abspath(path)) except ImportError: # pylint: disable=superfluous-parens print('Please install gprof2dot and dot for profile renderings.') else: # Empty context. yield def _validate_requirements( self, pipeline_proto: beam_runner_api_pb2.Pipeline) -> None: """As a test runner, validate requirements were set correctly.""" expected_requirements = set() def add_requirements(transform_id: str) -> None: transform = pipeline_proto.components.transforms[transform_id] if transform.spec.urn in translations.PAR_DO_URNS: payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.ParDoPayload) if payload.requests_finalization: expected_requirements.add( common_urns.requirements.REQUIRES_BUNDLE_FINALIZATION.urn) if (payload.state_specs or payload.timer_family_specs): expected_requirements.add( common_urns.requirements.REQUIRES_STATEFUL_PROCESSING.urn) if payload.requires_stable_input: expected_requirements.add( common_urns.requirements.REQUIRES_STABLE_INPUT.urn) if payload.requires_time_sorted_input: expected_requirements.add( common_urns.requirements.REQUIRES_TIME_SORTED_INPUT.urn) if payload.restriction_coder_id: expected_requirements.add( common_urns.requirements.REQUIRES_SPLITTABLE_DOFN.urn) else: for sub in transform.subtransforms: add_requirements(sub) for root in pipeline_proto.root_transform_ids: add_requirements(root) if not expected_requirements.issubset(pipeline_proto.requirements): raise ValueError( 'Missing requirement declaration: %s' % (expected_requirements - set(pipeline_proto.requirements))) def _check_requirements( self, pipeline_proto: beam_runner_api_pb2.Pipeline) -> None: """Check that this runner can satisfy all pipeline requirements.""" supported_requirements = set(self.supported_requirements()) for requirement in pipeline_proto.requirements: if requirement not in supported_requirements: raise ValueError( 'Unable to run pipeline with requirement: %s' % requirement) for transform in pipeline_proto.components.transforms.values(): if transform.spec.urn == common_urns.primitives.TEST_STREAM.urn: raise NotImplementedError(transform.spec.urn) elif transform.spec.urn in translations.PAR_DO_URNS: payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.ParDoPayload) for timer in payload.timer_family_specs.values(): if timer.time_domain not in ( beam_runner_api_pb2.TimeDomain.EVENT_TIME, beam_runner_api_pb2.TimeDomain.PROCESSING_TIME): raise NotImplementedError(timer.time_domain) def create_stages( self, pipeline_proto: beam_runner_api_pb2.Pipeline ) -> tuple[translations.TransformContext, list[translations.Stage]]: return translations.create_and_optimize_stages( copy.deepcopy(pipeline_proto), phases=[ translations.annotate_downstream_side_inputs, translations.fix_side_input_pcoll_coders, translations.pack_combiners, translations.lift_combiners, translations.expand_sdf, translations.expand_gbk, translations.sink_flattens, translations.greedily_fuse, translations.read_to_impulse, translations.impulse_to_input, translations.sort_stages, translations.add_impulse_to_dangling_transforms, translations.setup_timer_mapping, translations.populate_data_channel_coders, ], known_runner_urns=frozenset([ common_urns.primitives.FLATTEN.urn, common_urns.primitives.GROUP_BY_KEY.urn, ]), use_state_iterables=self._use_state_iterables, is_drain=self._is_drain) def run_stages( self, stage_context: translations.TransformContext, stages: list[translations.Stage]) -> 'RunnerResult': """Run a list of topologically-sorted stages in batch mode. Args: stage_context (translations.TransformContext) stages (list[fn_api_runner.translations.Stage]) """ worker_handler_manager = WorkerHandlerManager( stage_context.components.environments, self._provision_info) pipeline_metrics = MetricsContainer('') pipeline_metrics.get_counter( MetricName( str(type(self)), self.NUM_FUSED_STAGES_COUNTER, urn='internal:' + self.NUM_FUSED_STAGES_COUNTER)).update( len(stages)) monitoring_infos_by_stage: MutableMapping[ str, Iterable['metrics_pb2.MonitoringInfo']] = {} runner_execution_context = execution.FnApiRunnerExecutionContext( stages, worker_handler_manager, stage_context.components, stage_context.safe_coders, stage_context.data_channel_coders, self._num_workers, split_managers=self._split_managers) try: with self.maybe_profile(): # Initialize Runner context: # - Pipeline dictionaries, initial inputs and pipeline triggers # - Replace Data API endpoints in protobufs. runner_execution_context.setup() bundle_counter = 0 # Start executing all ready bundles. while len(runner_execution_context.queues.ready_inputs) > 0: _LOGGER.debug( "Remaining ready bundles: %s\n" "\tWatermark pending bundles: %s\n" "\tTime pending bundles: %s", len(runner_execution_context.queues.ready_inputs), len(runner_execution_context.queues.watermark_pending_inputs), len(runner_execution_context.queues.time_pending_inputs)) consuming_stage_name, bundle_input = ( runner_execution_context.queues.ready_inputs.deque()) stage = runner_execution_context.stages[consuming_stage_name] bundle_context_manager = runner_execution_context.bundle_manager_for( stage) _BUNDLE_LOGGER.debug( 'Running bundle for stage %s\n\tExpected outputs: %s timers: %s', bundle_context_manager.stage.name, bundle_context_manager.stage_data_outputs, bundle_context_manager.stage_timer_outputs) assert consuming_stage_name == bundle_context_manager.stage.name bundle_counter += 1 bundle_results = self._execute_bundle( runner_execution_context, bundle_context_manager, bundle_input) if consuming_stage_name in monitoring_infos_by_stage: monitoring_infos_by_stage[ consuming_stage_name] = consolidate_monitoring_infos( itertools.chain( bundle_results.process_bundle.monitoring_infos, monitoring_infos_by_stage[consuming_stage_name])) else: assert isinstance( bundle_results.process_bundle.monitoring_infos, Iterable) monitoring_infos_by_stage[consuming_stage_name] = \ bundle_results.process_bundle.monitoring_infos # Within monitoring_infos_by_stage we also keep monitoring information # for the whole pipeline, which we key under ''. if '' not in monitoring_infos_by_stage: monitoring_infos_by_stage[''] = list( pipeline_metrics.to_runner_api_monitoring_infos('').values()) else: monitoring_infos_by_stage[''] = consolidate_monitoring_infos( itertools.chain( pipeline_metrics.to_runner_api_monitoring_infos( '').values(), monitoring_infos_by_stage[''])) # We only compute new ready bundles whenever we run out of current # ready bundles, but we could do it after every new bundle and # it should work either way. if len(runner_execution_context.queues.ready_inputs) == 0: self._schedule_ready_bundles(runner_execution_context) assert len(runner_execution_context.queues.ready_inputs) == 0, ( 'A total of %d ready bundles did not execute.' % len(runner_execution_context.queues.ready_inputs)) assert len( runner_execution_context.queues.watermark_pending_inputs) == 0, \ ('A total of %d watermark-pending bundles did not execute.' % len(runner_execution_context.queues.watermark_pending_inputs)) assert len(runner_execution_context.queues.time_pending_inputs) == 0, ( 'A total of %d time-pending bundles did not execute.' % len(runner_execution_context.queues.time_pending_inputs)) finally: worker_handler_manager.close_all() return RunnerResult(runner.PipelineState.DONE, monitoring_infos_by_stage) def _schedule_ready_bundles( self, runner_execution_context: execution.FnApiRunnerExecutionContext): to_add_watermarks = [] while len(runner_execution_context.queues.watermark_pending_inputs) > 0: (stage_name, bundle_watermark), data_input = ( runner_execution_context.queues.watermark_pending_inputs.deque()) current_watermark = ( runner_execution_context.watermark_manager.get_stage_node( stage_name).input_watermark()) if current_watermark >= bundle_watermark: _BUNDLE_LOGGER.debug( 'Watermark: %s. Enqueuing bundle scheduled for (%s) for stage %s', current_watermark, bundle_watermark, stage_name) _LOGGER.debug( 'Stage info:\n\t:%s\n', runner_execution_context.watermark_manager.get_stage_node( stage_name)) runner_execution_context.queues.ready_inputs.enque( (stage_name, data_input)) else: _BUNDLE_LOGGER.debug( 'Unable to add bundle for stage %s\n' '\tStage input watermark: %s\n' '\tBundle schedule watermark: %s', stage_name, current_watermark, bundle_watermark) _LOGGER.debug( 'Stage info:\n\t:%s\n', runner_execution_context.watermark_manager.get_stage_node( stage_name)) to_add_watermarks.append(((stage_name, bundle_watermark), data_input)) for elm in to_add_watermarks: runner_execution_context.queues.watermark_pending_inputs.enque(elm) to_add_real_time = [] while len(runner_execution_context.queues.time_pending_inputs) > 0: current_time = runner_execution_context.clock.time() (stage_name, work_timestamp), data_input = ( runner_execution_context.queues.time_pending_inputs.deque()) if current_time >= work_timestamp: _LOGGER.debug( 'Time: %s. Enqueuing bundle scheduled for (%s) for stage %s', current_time, work_timestamp, stage_name) runner_execution_context.queues.ready_inputs.enque( (stage_name, data_input)) else: _LOGGER.debug( 'Unable to add bundle for stage %s\n' '\tCurrent time: %s\n' '\tBundle schedule time: %s\n', stage_name, current_time, work_timestamp) to_add_real_time.append(((stage_name, work_timestamp), data_input)) for elm in to_add_real_time: runner_execution_context.queues.time_pending_inputs.enque(elm) def _run_bundle_multiple_times_for_testing( self, runner_execution_context: execution.FnApiRunnerExecutionContext, bundle_manager: 'BundleManager', data_input: MutableMapping[str, execution.PartitionableBuffer], data_output: DataOutput, fired_timers: Mapping[translations.TimerFamilyId, execution.PartitionableBuffer], expected_output_timers: OutputTimers, ) -> None: """ If bundle_repeat > 0, replay every bundle for profiling and debugging. """ # all workers share state, so use any worker_handler. for _ in range(self._bundle_repeat): try: runner_execution_context.state_servicer.checkpoint() bundle_manager.process_bundle( data_input, data_output, fired_timers, expected_output_timers, dry_run=True) finally: runner_execution_context.state_servicer.restore() @staticmethod def _collect_written_timers( bundle_context_manager: execution.BundleContextManager ) -> tuple[dict[translations.TimerFamilyId, timestamp.Timestamp], OutputTimerData]: """Review output buffers, and collect written timers. This function reviews a stage that has just been run. The stage will have written timers to its output buffers. The function then takes the timers, and adds them to the `newly_set_timers` dictionary, and the timer_watermark_data dictionary. The function then returns the following two elements in a tuple: - timer_watermark_data: A dictionary mapping timer family to upcoming timestamp to fire. - newly_set_timers: A dictionary mapping timer family to timer buffers to be passed to the SDK upon firing. """ timer_watermark_data = {} newly_set_timers: OutputTimerData = {} for (transform_id, timer_family_id) in bundle_context_manager.stage.timers: written_timers = bundle_context_manager.get_buffer( create_buffer_id(timer_family_id, kind='timers'), transform_id) assert isinstance(written_timers, ListBuffer) timer_coder_impl = bundle_context_manager.get_timer_coder_impl( transform_id, timer_family_id) if not written_timers.cleared: timers_by_key_tag_and_window = {} for elements_timers in written_timers: for decoded_timer in timer_coder_impl.decode_all(elements_timers): key_tag_win = ( decoded_timer.user_key, decoded_timer.dynamic_timer_tag, decoded_timer.windows[0]) if not decoded_timer.clear_bit: timers_by_key_tag_and_window[key_tag_win] = decoded_timer elif (decoded_timer.clear_bit and key_tag_win in timers_by_key_tag_and_window): del timers_by_key_tag_and_window[key_tag_win] out = create_OutputStream() for decoded_timer in timers_by_key_tag_and_window.values(): # Only add not cleared timer to fired timers. if not decoded_timer.clear_bit: timer_coder_impl.encode_to_stream(decoded_timer, out, True) if (transform_id, timer_family_id) not in timer_watermark_data: timer_watermark_data[(transform_id, timer_family_id)] = timestamp.MAX_TIMESTAMP timer_watermark_data[(transform_id, timer_family_id)] = min( timer_watermark_data[(transform_id, timer_family_id)], decoded_timer.hold_timestamp) if (transform_id, timer_family_id) not in timer_watermark_data: continue newly_set_timers[(transform_id, timer_family_id)] = ( ListBuffer(coder_impl=timer_coder_impl), timer_watermark_data[(transform_id, timer_family_id)]) newly_set_timers[(transform_id, timer_family_id)][0].append(out.get()) written_timers.clear() return timer_watermark_data, newly_set_timers def _add_sdk_delayed_applications_to_deferred_inputs( self, bundle_context_manager: execution.BundleContextManager, bundle_result: beam_fn_api_pb2.InstructionResponse, deferred_inputs: MutableMapping[str, execution.PartitionableBuffer] ) -> set[str]: """Returns a set of PCollection IDs of PColls having delayed applications. This transform inspects the bundle_context_manager, and bundle_result objects, and adds all deferred inputs to the deferred_inputs object. """ pcolls_with_delayed_apps = set() for delayed_application in bundle_result.process_bundle.residual_roots: producer_name = bundle_context_manager.input_for( delayed_application.application.transform_id, delayed_application.application.input_id) if producer_name not in deferred_inputs: deferred_inputs[producer_name] = ListBuffer( coder_impl=bundle_context_manager.get_input_coder_impl( producer_name)) deferred_inputs[producer_name].append( delayed_application.application.element) transform = bundle_context_manager.process_bundle_descriptor.transforms[ producer_name] # We take the output with tag 'out' from the producer transform. The # producer transform is a GRPC read, and it has a single output. pcolls_with_delayed_apps.add(only_element(transform.outputs.values())) return pcolls_with_delayed_apps def _add_residuals_and_channel_splits_to_deferred_inputs( self, splits: list[beam_fn_api_pb2.ProcessBundleSplitResponse], bundle_context_manager: execution.BundleContextManager, last_sent: MutableMapping[str, execution.PartitionableBuffer], deferred_inputs: MutableMapping[str, execution.PartitionableBuffer] ) -> tuple[set[str], set[str]]: """Returns a two sets representing PCollections with watermark holds. The first set represents PCollections with delayed root applications. The second set represents PTransforms with channel splits. """ pcolls_with_delayed_apps = set() transforms_with_channel_splits = set() prev_stops: dict[str, int] = {} for split in splits: for delayed_application in split.residual_roots: producer_name = bundle_context_manager.input_for( delayed_application.application.transform_id, delayed_application.application.input_id) if producer_name not in deferred_inputs: deferred_inputs[producer_name] = ListBuffer( coder_impl=bundle_context_manager.get_input_coder_impl( producer_name)) deferred_inputs[producer_name].append( delayed_application.application.element) # We take the output with tag 'out' from the producer transform. The # producer transform is a GRPC read, and it has a single output. pcolls_with_delayed_apps.add( bundle_context_manager.process_bundle_descriptor. transforms[producer_name].outputs['out']) for channel_split in split.channel_splits: coder_impl = bundle_context_manager.get_input_coder_impl( channel_split.transform_id) # TODO(SDF): This requires deterministic ordering of buffer iteration. # TODO(SDF): The return split is in terms of indices. Ideally, # a runner could map these back to actual positions to effectively # describe the two "halves" of the now-split range. Even if we have # to buffer each element we send (or at the very least a bit of # metadata, like position, about each of them) this should be doable # if they're already in memory and we are bounding the buffer size # (e.g. to 10mb plus whatever is eagerly read from the SDK). In the # case of non-split-points, we can either immediately replay the # "non-split-position" elements or record them as we do the other # delayed applications. # Decode and recode to split the encoded buffer by element index. all_elements = list( coder_impl.decode_all( b''.join(last_sent[channel_split.transform_id]))) residual_elements = all_elements[ channel_split.first_residual_element:prev_stops. get(channel_split.transform_id, len(all_elements)) + 1] if residual_elements: transform = ( bundle_context_manager.process_bundle_descriptor.transforms[ channel_split.transform_id]) assert transform.spec.urn == bundle_processor.DATA_INPUT_URN transforms_with_channel_splits.add(transform.unique_name) if channel_split.transform_id not in deferred_inputs: coder_impl = bundle_context_manager.get_input_coder_impl( channel_split.transform_id) deferred_inputs[channel_split.transform_id] = ListBuffer( coder_impl=coder_impl) deferred_inputs[channel_split.transform_id].append( coder_impl.encode_all(residual_elements)) prev_stops[ channel_split.transform_id] = channel_split.last_primary_element return pcolls_with_delayed_apps, transforms_with_channel_splits def _execute_bundle( self, runner_execution_context: execution.FnApiRunnerExecutionContext, bundle_context_manager: execution.BundleContextManager, bundle_input: DataInput) -> beam_fn_api_pb2.InstructionResponse: """Execute a bundle end-to-end. Args: runner_execution_context (execution.FnApiRunnerExecutionContext): An object containing execution information for the pipeline. bundle_context_manager (execution.BundleContextManager): A description of the stage to execute, and its context. bundle_input: The set of buffers to input into this bundle """ worker_handler_manager = runner_execution_context.worker_handler_manager # TODO(pabloem): Should move this to be done once per stage worker_handler_manager.register_process_bundle_descriptor( bundle_context_manager.process_bundle_descriptor) # We create the bundle manager here, as it can be reused for bundles of # the same stage, but it may have to be created by-bundle later on. bundle_manager = self._get_bundle_manager(bundle_context_manager) last_result, deferred_inputs, newly_set_timers, watermark_updates = ( self._run_bundle( runner_execution_context, bundle_context_manager, bundle_input, bundle_context_manager.stage_data_outputs, bundle_context_manager.stage_timer_outputs, bundle_manager)) for pc_name, watermark in watermark_updates.items(): _BUNDLE_LOGGER.debug('Update: %s %s', pc_name, watermark) runner_execution_context.watermark_manager.set_pcoll_watermark( pc_name, watermark) if deferred_inputs: assert (runner_execution_context.watermark_manager.get_stage_node( bundle_context_manager.stage.name).output_watermark() < timestamp.MAX_TIMESTAMP), ( 'wrong timestamp for %s. ' % runner_execution_context.watermark_manager.get_stage_node( bundle_context_manager.stage.name)) runner_execution_context.queues.ready_inputs.enque( (bundle_context_manager.stage.name, DataInput(deferred_inputs, {}))) self._enqueue_set_timers( runner_execution_context, bundle_context_manager, newly_set_timers, bundle_input) if not deferred_inputs and not newly_set_timers: # Marking all inputs as having been fully produced. for _, output_pc in bundle_context_manager.stage_data_outputs.items(): _, update_output_pc = translations.split_buffer_id(output_pc) runner_execution_context.watermark_manager.get_pcoll_node( update_output_pc).set_produced_watermark(timestamp.MAX_TIMESTAMP) # Store the required downstream side inputs into state so it is accessible # for the worker when it runs bundles that consume this stage's output. data_side_input = ( runner_execution_context.side_input_descriptors_by_stage.get( bundle_context_manager.stage.name, {})) runner_execution_context.commit_side_inputs_to_state(data_side_input) buffers_to_clean = set() known_consumers = set() for transform_id, buffer_id in ( bundle_context_manager.stage_data_outputs.items()): for (consuming_stage_name, consuming_transform) in \ runner_execution_context.buffer_id_to_consumer_pairs.get(buffer_id, []): buffer = runner_execution_context.pcoll_buffers.get(buffer_id, None) if (buffer_id in runner_execution_context.pcoll_buffers and buffer_id not in buffers_to_clean): buffers_to_clean.add(buffer_id) elif buffer and buffer_id in buffers_to_clean: # If the buffer_id has already been added to buffers_to_clean, this # means that the buffer is being consumed by two separate stages, # so we create a copy of the buffer for every new stage. runner_execution_context.pcoll_buffers[buffer_id] = buffer.copy() buffer = runner_execution_context.pcoll_buffers[buffer_id] # empty buffer. Add it to the pcoll_buffer to avoid element # duplication. if buffer_id not in runner_execution_context.pcoll_buffers: buffer = bundle_context_manager.get_buffer(buffer_id, transform_id) runner_execution_context.pcoll_buffers[buffer_id] = buffer buffers_to_clean.add(buffer_id) # If the buffer has already been added to be consumed by # (stage, transform), then we don't need to add it again. This case # can happen whenever we flatten the same PCollection with itself. if (consuming_stage_name, consuming_transform, buffer_id) in known_consumers: continue else: known_consumers.add( (consuming_stage_name, consuming_transform, buffer_id)) # We enqueue all of the pending output buffers to be scheduled at the # MAX_TIMESTAMP for the downstream stage. runner_execution_context.queues.watermark_pending_inputs.enque( ((consuming_stage_name, timestamp.MAX_TIMESTAMP), DataInput({consuming_transform: buffer}, {}))) # type: ignore for bid in buffers_to_clean: if bid in runner_execution_context.pcoll_buffers: del runner_execution_context.pcoll_buffers[bid] return last_result def _enqueue_set_timers( self, runner_execution_context: execution.FnApiRunnerExecutionContext, bundle_context_manager: execution.BundleContextManager, fired_timers: translations.OutputTimerData, previous_bundle_input: DataInput): empty_data_input: MutableMapping[str, execution.PartitionableBuffer] = { k: ListBuffer(None) for k in previous_bundle_input.data.keys() } current_time = runner_execution_context.clock.time() current_watermark = \ runner_execution_context.watermark_manager.get_stage_node( bundle_context_manager.stage.name).input_watermark() for unique_timer_family in fired_timers: timer_data, target_timestamp = fired_timers[unique_timer_family] _, time_domain = bundle_context_manager.stage_timer_outputs[ unique_timer_family] if time_domain == beam_runner_api_pb2.TimeDomain.PROCESSING_TIME: if current_time >= target_timestamp: runner_execution_context.queues.ready_inputs.enque(( bundle_context_manager.stage.name, DataInput(empty_data_input, {unique_timer_family: timer_data}))) else: runner_execution_context.queues.time_pending_inputs.enque( ((bundle_context_manager.stage.name, target_timestamp), DataInput(empty_data_input, {unique_timer_family: timer_data}))) else: assert time_domain == beam_runner_api_pb2.TimeDomain.EVENT_TIME if current_watermark >= target_timestamp: runner_execution_context.queues.ready_inputs.enque(( bundle_context_manager.stage.name, DataInput(empty_data_input, {unique_timer_family: timer_data}))) else: runner_execution_context.queues.watermark_pending_inputs.enque( ((bundle_context_manager.stage.name, target_timestamp), DataInput(empty_data_input, {unique_timer_family: timer_data}))) def _get_bundle_manager( self, bundle_context_manager: execution.BundleContextManager): # TODO(pabloem): Consider moving this function to the BundleContextManager cache_token_generator = FnApiRunner.get_cache_token_generator(static=False) if bundle_context_manager.num_workers == 1: # Avoid thread/processor pools for increased performance and debugability. bundle_manager_type: Union[type[BundleManager], type[ParallelBundleManager]] = BundleManager elif bundle_context_manager.stage.is_stateful(): # State is keyed, and a single key cannot be processed concurrently. # Alternatively, we could arrange to partition work by key. # TODO(pabloem): Arrange to partition work by key. bundle_manager_type = BundleManager else: bundle_manager_type = ParallelBundleManager return bundle_manager_type( bundle_context_manager, self._progress_frequency, cache_token_generator=cache_token_generator) @staticmethod def _build_watermark_updates( runner_execution_context: execution.FnApiRunnerExecutionContext, stage_inputs: Iterable[str], expected_timers: Iterable[translations.TimerFamilyId], pcolls_with_da: set[str], transforms_w_splits: set[str], watermarks_by_transform_and_timer_family: dict[translations.TimerFamilyId, timestamp.Timestamp] ) -> dict[Union[str, translations.TimerFamilyId], timestamp.Timestamp]: """Builds a dictionary of PCollection (or TimerFamilyId) to timestamp. Args: stage_inputs: represent the set of expected input PCollections for a stage. These do not include timers. expected_timers: represent the set of TimerFamilyIds that the stage can expect to receive as inputs. pcolls_with_da: represent the set of stage input PCollections that had delayed applications. transforms_w_splits: represent the set of transforms in the stage that had input splits. watermarks_by_transform_and_timer_family: represent the set of watermark holds to be added for each timer family. """ updates: dict[Union[str, translations.TimerFamilyId], timestamp.Timestamp] = {} def get_pcoll_id(transform_id): buffer_id = runner_execution_context.input_transform_to_buffer_id[ transform_id] # For IMPULSE-reading transforms, we use the transform name as buffer id. if buffer_id == translations.IMPULSE_BUFFER: pcollection_id = transform_id else: _, pcollection_id = translations.split_buffer_id(buffer_id) return pcollection_id # Any PCollections that have deferred applications should have their # watermark held back. for pcoll in pcolls_with_da: updates[pcoll] = timestamp.MIN_TIMESTAMP # Also any transforms with splits should have their input PCollection's # watermark held back. for tr in transforms_w_splits: pcoll_id = get_pcoll_id(tr) updates[pcoll_id] = timestamp.MIN_TIMESTAMP # For all expected stage timers, we have two possible outcomes: # 1) If the stage set a firing time for the timer, then we hold the # watermark at that time # 2) If the stage did not set a firing time for the timer, then we # advance the watermark for that timer to MAX_TIMESTAMP. for timer_pcoll_id in expected_timers: updates[timer_pcoll_id] = watermarks_by_transform_and_timer_family.get( timer_pcoll_id, timestamp.MAX_TIMESTAMP) # For any PCollection in the set of stage inputs, if its watermark was not # held back (i.e. there weren't splits in its consumer PTransform, and there # weren't delayed applications of the PCollection's elements), then the # watermark should be advanced to MAX_TIMESTAMP. for transform_id in stage_inputs: pcoll_id = get_pcoll_id(transform_id) if pcoll_id not in updates: updates[pcoll_id] = timestamp.MAX_TIMESTAMP return updates def _run_bundle( self, runner_execution_context: execution.FnApiRunnerExecutionContext, bundle_context_manager: execution.BundleContextManager, bundle_input: DataInput, data_output: DataOutput, expected_timer_output: OutputTimers, bundle_manager: 'BundleManager' ) -> tuple[beam_fn_api_pb2.InstructionResponse, dict[str, execution.PartitionableBuffer], OutputTimerData, dict[Union[str, translations.TimerFamilyId], timestamp.Timestamp]]: """Execute a bundle, and return a result object, and deferred inputs.""" data_input = bundle_input.data input_timers = bundle_input.timers self._run_bundle_multiple_times_for_testing( runner_execution_context, bundle_manager, data_input, data_output, input_timers, expected_timer_output) result, splits = bundle_manager.process_bundle( data_input, data_output, input_timers, expected_timer_output) # Now we collect all the deferred inputs remaining from bundle execution. # Deferred inputs can be: # - timers # - SDK-initiated deferred applications of root elements # - Runner-initiated deferred applications of root elements deferred_inputs: dict[str, execution.PartitionableBuffer] = {} watermarks_by_transform_and_timer_family, newly_set_timers = ( self._collect_written_timers(bundle_context_manager)) sdk_pcolls_with_da = self._add_sdk_delayed_applications_to_deferred_inputs( bundle_context_manager, result, deferred_inputs) runner_pcolls_with_da, transforms_with_channel_splits = ( self._add_residuals_and_channel_splits_to_deferred_inputs( splits, bundle_context_manager, data_input, deferred_inputs)) watermark_updates = self._build_watermark_updates( runner_execution_context, data_input.keys(), expected_timer_output.keys(), runner_pcolls_with_da.union(sdk_pcolls_with_da), transforms_with_channel_splits, watermarks_by_transform_and_timer_family) # After collecting deferred inputs, we 'pad' the structure with empty # buffers for other expected inputs. if deferred_inputs: # The worker will be waiting on these inputs as well. for other_input in data_input: if other_input not in deferred_inputs: deferred_inputs[other_input] = ListBuffer( coder_impl=bundle_context_manager.get_input_coder_impl( other_input)) return result, deferred_inputs, newly_set_timers, watermark_updates @staticmethod def get_cache_token_generator( static: bool = True ) -> Iterator[beam_fn_api_pb2.ProcessBundleRequest.CacheToken]: """A generator for cache tokens. :arg static If True, generator always returns the same cache token If False, generator returns a new cache token each time :return A generator which returns a cache token on next(generator) """ def generate_token( identifier: int) -> beam_fn_api_pb2.ProcessBundleRequest.CacheToken: return beam_fn_api_pb2.ProcessBundleRequest.CacheToken( user_state=beam_fn_api_pb2.ProcessBundleRequest.CacheToken.UserState( ), token="cache_token_{}".format(identifier).encode("utf-8")) class StaticGenerator(object): def __init__(self) -> None: self._token = generate_token(1) def __iter__(self) -> 'StaticGenerator': # pylint: disable=non-iterator-returned return self def __next__(self) -> beam_fn_api_pb2.ProcessBundleRequest.CacheToken: return self._token class DynamicGenerator(object): def __init__(self) -> None: self._counter = 0 self._lock = threading.Lock() def __iter__(self) -> 'DynamicGenerator': # pylint: disable=non-iterator-returned return self def __next__(self) -> beam_fn_api_pb2.ProcessBundleRequest.CacheToken: with self._lock: self._counter += 1 return generate_token(self._counter) if static: return StaticGenerator() else: return DynamicGenerator() class ExtendedProvisionInfo(object): def __init__( self, provision_info: Optional[beam_provision_api_pb2.ProvisionInfo] = None, artifact_staging_dir: Optional[str] = None, job_name: str = '', ) -> None: self.provision_info = ( provision_info or beam_provision_api_pb2.ProvisionInfo()) self.artifact_staging_dir = artifact_staging_dir self.job_name = job_name def for_environment(self, env) -> 'ExtendedProvisionInfo': if env.dependencies: provision_info_with_deps = copy.deepcopy(self.provision_info) provision_info_with_deps.dependencies.extend(env.dependencies) return ExtendedProvisionInfo( provision_info_with_deps, self.artifact_staging_dir, self.job_name) else: return self _split_managers_from_context = [] @contextlib.contextmanager def split_manager(stage_name, split_manager): """Registers a split manager to control the flow of elements to a given stage. Used for testing. A split manager should be a coroutine yielding desired split fractions, receiving the corresponding split results. Currently, only one input is supported. """ try: _split_managers_from_context.append((stage_name, split_manager)) yield finally: _split_managers_from_context.pop() def create_test_split_manager(timings, fractions): # In general split managers can "wait" until events are observed in a # pipeline DoFn, but this is non-trivial to do cross-process, so we provide # a simple wall-clock-based splitter for out-of-process specification # via a pipeline option (e.g. for other languages). first = True def time_based_split_manager(unused_num_elements): nonlocal first if not first: return first = False start = time.time() for timing, fraction in zip(timings, fractions): now = time.time() if start + timing > now: time.sleep(start + timing - now) yield fraction return time_based_split_manager class BundleManager(object): """Manages the execution of a bundle from the runner-side. This class receives a bundle descriptor, and performs the following tasks: - Registration of the bundle with the worker. - Splitting of the bundle - Setting up any other bundle requirements (e.g. side inputs). - Submitting the bundle to worker for execution - Passing bundle input data to the worker - Collecting bundle output data from the worker - Finalizing the bundle. """ _uid_counter = 0 _lock = threading.Lock() def __init__( self, bundle_context_manager: execution.BundleContextManager, progress_frequency: Optional[float] = None, cache_token_generator=FnApiRunner.get_cache_token_generator(), split_managers=() ) -> None: """Set up a bundle manager. Args: progress_frequency """ self.bundle_context_manager: execution.BundleContextManager = ( bundle_context_manager) self._progress_frequency = progress_frequency self._worker_handler: Optional[WorkerHandler] = None self._cache_token_generator = cache_token_generator self.split_managers = split_managers def _send_input_to_worker( self, process_bundle_id: str, read_transform_id: str, byte_streams) -> None: assert self._worker_handler is not None data_out = self._worker_handler.data_conn.output_stream( process_bundle_id, read_transform_id) for byte_stream in (byte_streams or []): data_out.write(byte_stream) data_out.close() def _send_timers_to_worker( self, process_bundle_id, transform_id, timer_family_id, timers) -> None: assert self._worker_handler is not None timer_out = self._worker_handler.data_conn.output_timer_stream( process_bundle_id, transform_id, timer_family_id) for timer in timers: timer_out.write(timer) timer_out.close() def _select_split_manager(self) -> Optional[Callable[[int], Iterable[float]]]: """Returns the split manager, if any, associated with this stage.""" unique_names = set( t.unique_name for t in self.bundle_context_manager. process_bundle_descriptor.transforms.values()) for stage_name, candidate in reversed( self.bundle_context_manager.split_managers): if (stage_name in unique_names or (stage_name + '/Process') in unique_names): return candidate return None def _generate_splits_for_testing( self, split_manager, inputs: Mapping[str, execution.PartitionableBuffer], process_bundle_id) -> list[beam_fn_api_pb2.ProcessBundleSplitResponse]: split_results: list[beam_fn_api_pb2.ProcessBundleSplitResponse] = [] read_transform_id, buffer_data = only_element(inputs.items()) byte_stream = b''.join(buffer_data or []) num_elements = len( list( self.bundle_context_manager.get_input_coder_impl( read_transform_id).decode_all(byte_stream))) # Start the split manager in case it wants to set any breakpoints. split_manager_generator = split_manager(num_elements) try: split_fraction = next(split_manager_generator) done = False except StopIteration: split_fraction = None done = True # Send all the data. self._send_input_to_worker( process_bundle_id, read_transform_id, [byte_stream]) assert self._worker_handler is not None # Execute the requested splits. while not done: if split_fraction is None: split_result = None else: split_request = beam_fn_api_pb2.InstructionRequest( process_bundle_split=beam_fn_api_pb2.ProcessBundleSplitRequest( instruction_id=process_bundle_id, desired_splits={ read_transform_id: beam_fn_api_pb2. ProcessBundleSplitRequest.DesiredSplit( fraction_of_remainder=split_fraction, estimated_input_elements=num_elements) })) logging.info("Requesting split %s", split_request) split_response: beam_fn_api_pb2.InstructionResponse = ( self._worker_handler.control_conn.push(split_request).get()) for t in (0.05, 0.1, 0.2): if ('Unknown process bundle' in split_response.error or split_response.process_bundle_split == beam_fn_api_pb2.ProcessBundleSplitResponse()): time.sleep(t) split_response = self._worker_handler.control_conn.push( split_request).get() logging.info('Got split response %s', split_response) if ('Unknown process bundle' in split_response.error or split_response.process_bundle_split == beam_fn_api_pb2.ProcessBundleSplitResponse()): # It may have finished too fast. split_result = None elif split_response.error: raise RuntimeError(split_response.error) else: split_result = split_response.process_bundle_split split_results.append(split_result) try: split_fraction = split_manager_generator.send(split_result) except StopIteration: break return split_results def process_bundle( self, inputs: Mapping[str, execution.PartitionableBuffer], expected_outputs: DataOutput, fired_timers: Mapping[translations.TimerFamilyId, execution.PartitionableBuffer], expected_output_timers: OutputTimers, dry_run: bool = False, ) -> BundleProcessResult: # Unique id for the instruction processing this bundle. with BundleManager._lock: BundleManager._uid_counter += 1 process_bundle_id = 'bundle_%s' % BundleManager._uid_counter self._worker_handler = self.bundle_context_manager.worker_handlers[ BundleManager._uid_counter % len(self.bundle_context_manager.worker_handlers)] split_manager = self._select_split_manager() if split_manager: logging.info('Using split manager %s', split_manager) if not split_manager: # Send timers. for transform_id, timer_family_id in expected_output_timers.keys(): self._send_timers_to_worker( process_bundle_id, transform_id, timer_family_id, fired_timers.get((transform_id, timer_family_id), [])) # If there is no split_manager, write all input data to the channel. for transform_id, elements in inputs.items(): self._send_input_to_worker(process_bundle_id, transform_id, elements) # Actually start the bundle. process_bundle_req = beam_fn_api_pb2.InstructionRequest( instruction_id=process_bundle_id, process_bundle=beam_fn_api_pb2.ProcessBundleRequest( process_bundle_descriptor_id=self.bundle_context_manager. process_bundle_descriptor.id, cache_tokens=[next(self._cache_token_generator)])) result_future = self._worker_handler.control_conn.push(process_bundle_req) split_results: list[beam_fn_api_pb2.ProcessBundleSplitResponse] = [] with ProgressRequester(self._worker_handler, process_bundle_id, self._progress_frequency): if split_manager: split_results = self._generate_splits_for_testing( split_manager, inputs, process_bundle_id) expect_reads: list[Union[str, tuple[str, str]]] = list(expected_outputs.keys()) expect_reads.extend(list(expected_output_timers.keys())) # Gather all output data. for output in self._worker_handler.data_conn.input_elements( process_bundle_id, expect_reads, abort_callback=lambda: (result_future.is_done() and bool(result_future.get().error))): if isinstance(output, beam_fn_api_pb2.Elements.Timers) and not dry_run: with BundleManager._lock: timer_buffer = self.bundle_context_manager.get_buffer( expected_output_timers[( output.transform_id, output.timer_family_id)][0], output.transform_id) if timer_buffer.cleared: timer_buffer.reset() timer_buffer.append(output.timers) if isinstance(output, beam_fn_api_pb2.Elements.Data) and not dry_run: with BundleManager._lock: self.bundle_context_manager.get_buffer( expected_outputs[output.transform_id], output.transform_id).append(output.data) result: beam_fn_api_pb2.InstructionResponse = result_future.get() if result.error: raise RuntimeError(result.error) if result.process_bundle.requires_finalization: finalize_request = beam_fn_api_pb2.InstructionRequest( finalize_bundle=beam_fn_api_pb2.FinalizeBundleRequest( instruction_id=process_bundle_id)) finalize_response = self._worker_handler.control_conn.push( finalize_request).get() if finalize_response.error: raise RuntimeError(finalize_response.error) return result, split_results class ParallelBundleManager(BundleManager): def __init__( self, bundle_context_manager: execution.BundleContextManager, progress_frequency: Optional[float] = None, cache_token_generator=None, **kwargs) -> None: super().__init__( bundle_context_manager, progress_frequency, cache_token_generator=cache_token_generator) self._num_workers = bundle_context_manager.num_workers def process_bundle( self, inputs: Mapping[str, execution.PartitionableBuffer], expected_outputs: DataOutput, fired_timers: Mapping[translations.TimerFamilyId, execution.PartitionableBuffer], expected_output_timers: OutputTimers, dry_run: bool = False, ) -> BundleProcessResult: part_inputs: list[dict[str, list[bytes]]] = [{} for _ in range(self._num_workers)] # Timers are only executed on the first worker # TODO(BEAM-9741): Split timers to multiple workers timer_inputs = [ fired_timers if i == 0 else {} for i in range(self._num_workers) ] for name, input in inputs.items(): for ix, part in enumerate(input.partition(self._num_workers)): part_inputs[ix][name] = part merged_result: Optional[beam_fn_api_pb2.InstructionResponse] = None split_result_list: list[beam_fn_api_pb2.ProcessBundleSplitResponse] = [] def execute(part_map_input_timers) -> BundleProcessResult: part_map, input_timers = part_map_input_timers bundle_manager = BundleManager( self.bundle_context_manager, self._progress_frequency, cache_token_generator=self._cache_token_generator) return bundle_manager.process_bundle( part_map, expected_outputs, input_timers, expected_output_timers, dry_run) with thread_pool_executor.shared_unbounded_instance() as executor: for result, split_result in executor.map(execute, zip(part_inputs, # pylint: disable=bad-option-value timer_inputs)): split_result_list += split_result if merged_result is None: merged_result = result else: merged_result = beam_fn_api_pb2.InstructionResponse( process_bundle=beam_fn_api_pb2.ProcessBundleResponse( monitoring_infos=monitoring_infos.consolidate( itertools.chain( result.process_bundle.monitoring_infos, merged_result.process_bundle.monitoring_infos))), error=result.error or merged_result.error) assert merged_result is not None return merged_result, split_result_list class ProgressRequester(threading.Thread): """ Thread that asks SDK Worker for progress reports with a certain frequency. A callback can be passed to call with progress updates. """ def __init__( self, worker_handler: WorkerHandler, instruction_id, frequency, callback=None) -> None: super().__init__() self._worker_handler = worker_handler self._instruction_id = instruction_id self._frequency = frequency self._done = False self._latest_progress: Optional[ beam_fn_api_pb2.ProcessBundleProgressResponse] = None self._callback = callback self.daemon = True def __enter__(self): if self._frequency: self.start() def __exit__(self, *unused_exc_info): if self._frequency: self.stop() def run(self): while not self._done: try: progress_result = self._worker_handler.control_conn.push( beam_fn_api_pb2.InstructionRequest( process_bundle_progress=beam_fn_api_pb2. ProcessBundleProgressRequest( instruction_id=self._instruction_id))).get() self._latest_progress = progress_result.process_bundle_progress if self._callback: self._callback(self._latest_progress) except Exception as exn: _LOGGER.error("Bad progress: %s", exn) time.sleep(self._frequency) def stop(self): self._done = True class FnApiMetrics(metric.MetricResults): def __init__(self, step_monitoring_infos, user_metrics_only=True): """Used for querying metrics from the PipelineResult object. step_monitoring_infos: Per step metrics specified as MonitoringInfos. user_metrics_only: If true, includes user metrics only. """ self._counters = {} self._distributions = {} self._gauges = {} self._string_sets = {} self._bounded_tries = {} self._user_metrics_only = user_metrics_only self._monitoring_infos = step_monitoring_infos for smi in step_monitoring_infos.values(): counters, distributions, gauges, string_sets, bounded_tries = ( portable_metrics.from_monitoring_infos(smi, user_metrics_only)) self._counters.update(counters) self._distributions.update(distributions) self._gauges.update(gauges) self._string_sets.update(string_sets) self._bounded_tries.update(bounded_tries) def query(self, filter=None): counters = [ MetricResult(k, v, v) for k, v in self._counters.items() if self.matches(filter, k) ] distributions = [ MetricResult(k, v, v) for k, v in self._distributions.items() if self.matches(filter, k) ] gauges = [ MetricResult(k, v, v) for k, v in self._gauges.items() if self.matches(filter, k) ] string_sets = [ MetricResult(k, v, v) for k, v in self._string_sets.items() if self.matches(filter, k) ] bounded_tries = [ MetricResult(k, v, v) for k, v in self._bounded_tries.items() if self.matches(filter, k) ] return { self.COUNTERS: counters, self.DISTRIBUTIONS: distributions, self.GAUGES: gauges, self.STRINGSETS: string_sets, self.BOUNDED_TRIES: bounded_tries, } def monitoring_infos(self) -> list[metrics_pb2.MonitoringInfo]: return [ item for sublist in self._monitoring_infos.values() for item in sublist ] class RunnerResult(runner.PipelineResult): def __init__(self, state, monitoring_infos_by_stage): super().__init__(state) self._monitoring_infos_by_stage = monitoring_infos_by_stage self._metrics = None self._monitoring_metrics = None def wait_until_finish(self, duration=None): return self._state def metrics(self): """Returns a queryable object including user metrics only.""" if self._metrics is None: self._metrics = FnApiMetrics( self._monitoring_infos_by_stage, user_metrics_only=True) return self._metrics def monitoring_metrics(self): """Returns a queryable object including all metrics.""" if self._monitoring_metrics is None: self._monitoring_metrics = FnApiMetrics( self._monitoring_infos_by_stage, user_metrics_only=False) return self._monitoring_metrics def monitoring_infos(self): for ms in self._monitoring_infos_by_stage.values(): yield from ms ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py0000644000175100017510000027610415206113032030157 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import collections import gc import logging import os import random import re import shutil import tempfile import threading import time import traceback import typing import unittest import uuid from typing import Any from typing import Iterator from typing import no_type_check import hamcrest # pylint: disable=ungrouped-imports import numpy as np import pytest from hamcrest.core.matcher import Matcher from hamcrest.core.string_description import StringDescription from tenacity import retry from tenacity import stop_after_attempt import apache_beam as beam from apache_beam.coders import coders from apache_beam.coders.coders import StrUtf8Coder from apache_beam.io import restriction_trackers from apache_beam.io.watermark_estimators import ManualWatermarkEstimator from apache_beam.metrics import monitoring_infos from apache_beam.metrics.execution import MetricKey from apache_beam.metrics.metricbase import MetricName from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import DirectOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.value_provider import RuntimeValueProvider from apache_beam.portability import python_urns from apache_beam.runners.portability import fn_api_runner from apache_beam.runners.portability.fn_api_runner import fn_runner from apache_beam.runners.sdf_utils import RestrictionTrackerView from apache_beam.runners.worker import data_plane from apache_beam.runners.worker import statesampler from apache_beam.runners.worker.operations import InefficientExecutionWarning from apache_beam.testing.synthetic_pipeline import SyntheticSDFAsSource from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.testing.util import has_at_least_one from apache_beam.tools import utils from apache_beam.transforms import environments from apache_beam.transforms import trigger from apache_beam.transforms import userstate from apache_beam.transforms import window from apache_beam.transforms.periodicsequence import PeriodicImpulse from apache_beam.utils import timestamp from apache_beam.utils import windowed_value if statesampler.FAST_SAMPLER: DEFAULT_SAMPLING_PERIOD_MS = statesampler.DEFAULT_SAMPLING_PERIOD_MS else: DEFAULT_SAMPLING_PERIOD_MS = 0 _LOGGER = logging.getLogger(__name__) def _matcher_or_equal_to(value_or_matcher): """Pass-thru for matchers, and wraps value inputs in an equal_to matcher.""" if value_or_matcher is None: return None if isinstance(value_or_matcher, Matcher): return value_or_matcher return hamcrest.equal_to(value_or_matcher) def has_urn_and_labels(mi, urn, labels): """Returns true if it the monitoring_info contains the labels and urn.""" def contains_labels(mi, labels): # Check all the labels and their values exist in the monitoring_info return all(item in mi.labels.items() for item in labels.items()) return contains_labels(mi, labels) and mi.urn == urn class FnApiRunnerTest(unittest.TestCase): def create_pipeline(self, is_drain=False): return beam.Pipeline(runner=fn_api_runner.FnApiRunner(is_drain=is_drain)) def test_assert_that(self): # TODO: figure out a way for fn_api_runner to parse and raise the # underlying exception. with self.assertRaisesRegex(Exception, 'Failed assert'): with self.create_pipeline() as p: assert_that(p | beam.Create(['a', 'b']), equal_to(['a'])) def test_create(self): with self.create_pipeline() as p: assert_that(p | beam.Create(['a', 'b']), equal_to(['a', 'b'])) def test_pardo(self): with self.create_pipeline() as p: res = ( p | beam.Create(['a', 'bc']) | beam.Map(lambda e: e * 2) | beam.Map(lambda e: e + 'x')) assert_that(res, equal_to(['aax', 'bcbcx'])) def test_batch_pardo(self): with self.create_pipeline() as p: res = ( p | beam.Create(np.array([1, 2, 3], dtype=np.int64)).with_output_types( np.int64) | beam.ParDo(ArrayMultiplyDoFn()) | beam.Map(lambda x: x * 3)) assert_that(res, equal_to([6, 12, 18])) def test_batch_pardo_override_type_inference(self): class ArrayMultiplyDoFnOverride(beam.DoFn): def process_batch(self, batch, *unused_args, **unused_kwargs) -> Iterator[np.ndarray]: assert isinstance(batch, np.ndarray) yield batch * 2 # infer_output_type must be defined (when there's no process method), # otherwise we don't know the input type is the same as output type. def infer_output_type(self, input_type): return input_type def get_input_batch_type(self, input_element_type): from apache_beam.typehints.batch import NumpyArray return NumpyArray[input_element_type] def get_output_batch_type(self, input_element_type): return self.get_input_batch_type(input_element_type) with self.create_pipeline() as p: res = ( p | beam.Create(np.array([1, 2, 3], dtype=np.int64)).with_output_types( np.int64) | beam.ParDo(ArrayMultiplyDoFnOverride()) | beam.Map(lambda x: x * 3)) assert_that(res, equal_to([6, 12, 18])) @unittest.skip('https://github.com/apache/beam/issues/23944') def test_batch_pardo_trigger_flush(self): try: utils.check_compiled('apache_beam.coders.coder_impl') except RuntimeError: self.skipTest( 'https://github.com/apache/beam/issues/21643: FnRunnerTest with ' 'non-trivial inputs flakes in non-cython environments') with self.create_pipeline() as p: res = ( p # Pass more than GeneralPurposeConsumerSet.MAX_BATCH_SIZE elements # here to make sure we exercise the batch size limit. | beam.Create(np.array(range(5000), dtype=np.int64)).with_output_types(np.int64) | beam.ParDo(ArrayMultiplyDoFn()) | beam.Map(lambda x: x * 3)) assert_that(res, equal_to([i * 2 * 3 for i in range(5000)])) def test_batch_rebatch_pardos(self): # Should raise a warning about the rebatching that mentions: # - The consuming DoFn # - The output batch type of the producer # - The input batch type of the consumer with self.assertWarnsRegex(InefficientExecutionWarning, (r'ListPlusOneDoFn.*NumpyArray.*List\[\]')): with self.create_pipeline() as p: res = ( p | beam.Create(np.array([1, 2, 3], dtype=np.int64)).with_output_types(np.int64) | beam.ParDo(ArrayMultiplyDoFn()) | beam.ParDo(ListPlusOneDoFn()) | beam.Map(lambda x: x * 3)) assert_that(res, equal_to([9, 15, 21])) def test_batch_pardo_fusion_break(self): class NormalizeDoFn(beam.DoFn): @no_type_check def process_batch( self, batch: np.ndarray, mean: np.float64, ) -> Iterator[np.ndarray]: assert isinstance(batch, np.ndarray) yield batch - mean # infer_output_type must be defined (when there's no process method), # otherwise we don't know the input type is the same as output type. def infer_output_type(self, input_type): return np.float64 with self.create_pipeline() as p: pc = ( p | beam.Create(np.array([1, 2, 3], dtype=np.int64)).with_output_types( np.int64) | beam.ParDo(ArrayMultiplyDoFn())) res = ( pc | beam.ParDo( NormalizeDoFn(), mean=beam.pvalue.AsSingleton( pc | beam.CombineGlobally(beam.combiners.MeanCombineFn())))) assert_that(res, equal_to([-2, 0, 2])) def test_batch_pardo_dofn_params(self): class ConsumeParamsDoFn(beam.DoFn): @no_type_check def process_batch( self, batch: np.ndarray, ts=beam.DoFn.TimestampParam, pane_info=beam.DoFn.PaneInfoParam, ) -> Iterator[np.ndarray]: assert isinstance(batch, np.ndarray) assert isinstance(ts, timestamp.Timestamp) assert isinstance(pane_info, windowed_value.PaneInfo) yield batch * ts.seconds() # infer_output_type must be defined (when there's no process method), # otherwise we don't know the input type is the same as output type. def infer_output_type(self, input_type): return input_type with self.create_pipeline() as p: res = ( p | beam.Create(np.array(range(10), dtype=np.int64)).with_output_types( np.int64) | beam.Map(lambda t: window.TimestampedValue(t, int(t % 2))). with_output_types(np.int64) | beam.ParDo(ConsumeParamsDoFn())) assert_that(res, equal_to([0, 1, 0, 3, 0, 5, 0, 7, 0, 9])) def test_batch_pardo_window_param(self): class PerWindowDoFn(beam.DoFn): @no_type_check def process_batch( self, batch: np.ndarray, window=beam.DoFn.WindowParam, ) -> Iterator[np.ndarray]: yield batch * window.start.seconds() # infer_output_type must be defined (when there's no process method), # otherwise we don't know the input type is the same as output type. def infer_output_type(self, input_type): return input_type with self.create_pipeline() as p: res = ( p | beam.Create(np.array(range(10), dtype=np.int64)).with_output_types( np.int64) | beam.Map(lambda t: window.TimestampedValue(t, int(t))). with_output_types(np.int64) | beam.WindowInto(window.FixedWindows(5)) | beam.ParDo(PerWindowDoFn())) assert_that(res, equal_to([0, 0, 0, 0, 0, 25, 30, 35, 40, 45])) def test_batch_pardo_overlapping_windows(self): class PerWindowDoFn(beam.DoFn): @no_type_check def process_batch(self, batch: np.ndarray, window=beam.DoFn.WindowParam) -> Iterator[np.ndarray]: yield batch * window.start.seconds() # infer_output_type must be defined (when there's no process method), # otherwise we don't know the input type is the same as output type. def infer_output_type(self, input_type): return input_type with self.create_pipeline() as p: res = ( p | beam.Create(np.array(range(10), dtype=np.int64)).with_output_types( np.int64) | beam.Map(lambda t: window.TimestampedValue(t, int(t))). with_output_types(np.int64) | beam.WindowInto(window.SlidingWindows(size=5, period=3)) | beam.ParDo(PerWindowDoFn())) assert_that( res, equal_to([ 0 * -3, 1 * -3, # [-3, 2) 0 * 0, 1 * 0, 2 * 0, 3 * 0, 4 * 0, # [ 0, 5) 3 * 3, 4 * 3, 5 * 3, 6 * 3, 7 * 3, # [ 3, 8) 6 * 6, 7 * 6, 8 * 6, 9 * 6, # [ 6, 11) 9 * 9 # [ 9, 14) ])) def test_batch_to_element_pardo(self): class ArraySumDoFn(beam.DoFn): @beam.DoFn.yields_elements def process_batch(self, batch: np.ndarray, *unused_args, **unused_kwargs) -> Iterator[np.int64]: yield batch.sum() def infer_output_type(self, input_type): assert input_type == np.int64 return np.int64 with self.create_pipeline() as p: res = ( p | beam.Create(np.array(range(100), dtype=np.int64)).with_output_types( np.int64) | beam.ParDo(ArrayMultiplyDoFn()) | beam.ParDo(ArraySumDoFn()) | beam.CombineGlobally(sum)) assert_that(res, equal_to([99 * 50 * 2])) def test_element_to_batch_pardo(self): class ArrayProduceDoFn(beam.DoFn): @beam.DoFn.yields_batches def process(self, element: np.int64, *unused_args, **unused_kwargs) -> Iterator[np.ndarray]: yield np.array([element] * int(element)) # infer_output_type must be defined (when there's no process method), # otherwise we don't know the input type is the same as output type. def infer_output_type(self, input_type): return np.int64 with self.create_pipeline() as p: res = ( p | beam.Create(np.array([1, 2, 3], dtype=np.int64)).with_output_types( np.int64) | beam.ParDo(ArrayProduceDoFn()) | beam.ParDo(ArrayMultiplyDoFn()) | beam.Map(lambda x: x * 3)) assert_that(res, equal_to([6, 12, 12, 18, 18, 18])) @unittest.skip('https://github.com/apache/beam/issues/23944') def test_pardo_large_input(self): try: utils.check_compiled('apache_beam.coders.coder_impl') except RuntimeError: self.skipTest( 'https://github.com/apache/beam/issues/21643: FnRunnerTest with ' 'non-trivial inputs flakes in non-cython environments') with self.create_pipeline() as p: res = ( p | beam.Create(np.array(range(5000), dtype=np.int64)).with_output_types(np.int64) | beam.Map(lambda e: e * 2) | beam.Map(lambda e: e + 3)) assert_that(res, equal_to([(i * 2) + 3 for i in range(5000)])) def test_pardo_side_outputs(self): def tee(elem, *tags): for tag in tags: if tag in elem: yield beam.pvalue.TaggedOutput(tag, elem) with self.create_pipeline() as p: xy = ( p | 'Create' >> beam.Create(['x', 'y', 'xy']) | beam.FlatMap(tee, 'x', 'y').with_outputs()) assert_that(xy.x, equal_to(['x', 'xy']), label='x') assert_that(xy.y, equal_to(['y', 'xy']), label='y') def test_pardo_side_and_main_outputs(self): def even_odd(elem): yield elem yield beam.pvalue.TaggedOutput('odd' if elem % 2 else 'even', elem) with self.create_pipeline() as p: ints = p | beam.Create([1, 2, 3]) named = ints | 'named' >> beam.FlatMap(even_odd).with_outputs( 'even', 'odd', main='all') assert_that(named.all, equal_to([1, 2, 3]), label='named.all') assert_that(named.even, equal_to([2]), label='named.even') assert_that(named.odd, equal_to([1, 3]), label='named.odd') unnamed = ints | 'unnamed' >> beam.FlatMap(even_odd).with_outputs() unnamed[None] | beam.Map(id) # pylint: disable=expression-not-assigned assert_that(unnamed[None], equal_to([1, 2, 3]), label='unnamed.all') assert_that(unnamed.even, equal_to([2]), label='unnamed.even') assert_that(unnamed.odd, equal_to([1, 3]), label='unnamed.odd') def test_pardo_side_inputs(self): def cross_product(elem, sides): for side in sides: yield elem, side with self.create_pipeline() as p: main = p | 'main' >> beam.Create(['a', 'b', 'c']) side = p | 'side' >> beam.Create(['x', 'y']) assert_that( main | beam.FlatMap(cross_product, beam.pvalue.AsList(side)), equal_to([('a', 'x'), ('b', 'x'), ('c', 'x'), ('a', 'y'), ('b', 'y'), ('c', 'y')])) def test_pardo_side_input_dependencies(self): ## # The issue that this test surfaces is that whenever a PCollection is # consumed as main input by several stages, we have a bug. # # The bug is: A stage assumes that it has run if its upstream PCollection # has watermark=MAX. If multiple stages depend on a single PCollection, then # the first stage that runs it will set its watermar to MAX, and other # stages will think they've run. # # How to resolve? # Option1: to make sure that a PCollection's watermark only advances with # its consumption by all consumers? # - I tested this and it didn't seem fruitful/obvious. (YET) # # Option2: Change execution schema: A PCollection's watermark represents # its *production* watermark, not its *consumption* watermark.(?) with self.create_pipeline() as p: inputs = [p | beam.Create([None])] for k in range(1, 10): inputs.append( inputs[0] | beam.ParDo( ExpectingSideInputsFn(f'Do{k}'), *[beam.pvalue.AsList(inputs[s]) for s in range(1, k)])) def test_flatmap_numpy_array(self): with self.create_pipeline() as p: pc = ( p | beam.Create([np.array(range(10))]) | beam.FlatMap(lambda arr: arr)) assert_that(pc, equal_to([np.int64(i) for i in range(10)])) @unittest.skip('https://github.com/apache/beam/issues/21228') def test_pardo_side_input_sparse_dependencies(self): with self.create_pipeline() as p: inputs = [] def choose_input(s): return inputs[(389 + s * 5077) % len(inputs)] for k in range(20): num_inputs = int((k * k % 16)**0.5) if num_inputs == 0: inputs.append(p | f'Create{k}' >> beam.Create([f'Create{k}'])) else: inputs.append( choose_input(0) | beam.ParDo( ExpectingSideInputsFn(f'Do{k}'), *[ beam.pvalue.AsList(choose_input(s)) for s in range(1, num_inputs) ])) def test_pardo_windowed_side_inputs(self): with self.create_pipeline() as p: # Now with some windowing. pcoll = p | beam.Create(list( range(10))) | beam.Map(lambda t: window.TimestampedValue(t, t)) # Intentionally choosing non-aligned windows to highlight the transition. main = pcoll | 'WindowMain' >> beam.WindowInto(window.FixedWindows(5)) side = pcoll | 'WindowSide' >> beam.WindowInto(window.FixedWindows(7)) res = main | beam.Map( lambda x, s: (x, sorted(s)), beam.pvalue.AsList(side)) assert_that( res, equal_to([ # The window [0, 5) maps to the window [0, 7). (0, list(range(7))), (1, list(range(7))), (2, list(range(7))), (3, list(range(7))), (4, list(range(7))), # The window [5, 10) maps to the window [7, 14). (5, list(range(7, 10))), (6, list(range(7, 10))), (7, list(range(7, 10))), (8, list(range(7, 10))), (9, list(range(7, 10))) ]), label='windowed') def test_flattened_side_input(self, with_transcoding=True): with self.create_pipeline() as p: main = p | 'main' >> beam.Create([None]) side1 = p | 'side1' >> beam.Create([('a', 1)]) side2 = p | 'side2' >> beam.Create([('b', 2)]) if with_transcoding: # Also test non-matching coder types (transcoding required) third_element = [('another_type')] else: third_element = [('b', 3)] side3 = p | 'side3' >> beam.Create(third_element) side = (side1, side2) | beam.Flatten() assert_that( main | beam.Map(lambda a, b: (a, b), beam.pvalue.AsDict(side)), equal_to([(None, { 'a': 1, 'b': 2 })]), label='CheckFlattenAsSideInput') assert_that((side, side3) | 'FlattenAfter' >> beam.Flatten(), equal_to([('a', 1), ('b', 2)] + third_element), label='CheckFlattenOfSideInput') def test_flatten_and_gbk(self, with_transcoding=True): with self.create_pipeline() as p: side1 = p | 'side1' >> beam.Create([('a', 1)]) if with_transcoding: # Also test non-matching coder types (transcoding required) second_element = [('another_type')] else: second_element = [('b', 2)] side2 = p | 'side2' >> beam.Create(second_element) flatten_out = (side1, side2) | beam.Flatten() gbk_out = side1 | beam.GroupByKey() assert_that(flatten_out, equal_to([('a', 1)] + second_element)) assert_that(gbk_out, equal_to([('a', [1])])) def test_gbk_side_input(self): with self.create_pipeline() as p: main = p | 'main' >> beam.Create([None]) side = p | 'side' >> beam.Create([('a', 1)]) | beam.GroupByKey() assert_that( main | beam.Map(lambda a, b: (a, b), beam.pvalue.AsDict(side)), equal_to([(None, { 'a': [1] })])) def test_multimap_side_input(self): with self.create_pipeline() as p: main = p | 'main' >> beam.Create(['a', 'b']) side = p | 'side' >> beam.Create([('a', 1), ('b', 2), ('a', 3)]) assert_that( main | beam.Map( lambda k, d: (k, sorted(d[k])), beam.pvalue.AsMultiMap(side)), equal_to([('a', [1, 3]), ('b', [2])])) def test_multimap_multiside_input(self): # A test where two transforms in the same stage consume the same PCollection # twice as side input. with self.create_pipeline() as p: main = p | 'main' >> beam.Create(['a', 'b']) side = p | 'side' >> beam.Create([('a', 1), ('b', 2), ('a', 3)]) assert_that( main | 'first map' >> beam.Map( lambda k, d, l: (k, sorted(d[k]), sorted([e[1] for e in l])), beam.pvalue.AsMultiMap(side), beam.pvalue.AsList(side)) | 'second map' >> beam.Map( lambda k, d, l: (k[0], sorted(d[k[0]]), sorted([e[1] for e in l])), beam.pvalue.AsMultiMap(side), beam.pvalue.AsList(side)), equal_to([('a', [1, 3], [1, 2, 3]), ('b', [2], [1, 2, 3])])) def test_multimap_side_input_type_coercion(self): with self.create_pipeline() as p: main = p | 'main' >> beam.Create(['a', 'b']) # The type of this side-input is forced to Any (overriding type # inference). Without type coercion to Tuple[Any, Any], the usage of this # side-input in AsMultiMap() below should fail. side = ( p | 'side' >> beam.Create([('a', 1), ('b', 2), ('a', 3)]).with_output_types(typing.Any)) assert_that( main | beam.Map( lambda k, d: (k, sorted(d[k])), beam.pvalue.AsMultiMap(side)), equal_to([('a', [1, 3]), ('b', [2])])) def test_pardo_unfusable_side_inputs(self): def cross_product(elem, sides): for side in sides: yield elem, side with self.create_pipeline() as p: pcoll = p | beam.Create(['a', 'b']) assert_that( pcoll | beam.FlatMap(cross_product, beam.pvalue.AsList(pcoll)), equal_to([('a', 'a'), ('a', 'b'), ('b', 'a'), ('b', 'b')])) def test_pardo_unfusable_side_inputs_with_separation(self): def cross_product(elem, sides): for side in sides: yield elem, side with self.create_pipeline() as p: pcoll = p | beam.Create(['a', 'b']) derived = ((pcoll, ) | beam.Flatten() | beam.Map(lambda x: (x, x)) | beam.GroupByKey() | 'Unkey' >> beam.Map(lambda kv: kv[0])) assert_that( pcoll | beam.FlatMap(cross_product, beam.pvalue.AsList(derived)), equal_to([('a', 'a'), ('a', 'b'), ('b', 'a'), ('b', 'b')])) def test_pardo_state_only(self): index_state_spec = userstate.CombiningValueStateSpec('index', sum) value_and_index_state_spec = userstate.ReadModifyWriteStateSpec( 'value:index', StrUtf8Coder()) # TODO(ccy): State isn't detected with Map/FlatMap. class AddIndex(beam.DoFn): def process( self, kv, index=beam.DoFn.StateParam(index_state_spec), value_and_index=beam.DoFn.StateParam(value_and_index_state_spec)): k, v = kv index.add(1) value_and_index.write('%s:%s' % (v, index.read())) yield k, v, index.read(), value_and_index.read() inputs = [('A', 'a')] * 2 + [('B', 'b')] * 3 expected = [('A', 'a', 1, 'a:1'), ('A', 'a', 2, 'a:2'), ('B', 'b', 1, 'b:1'), ('B', 'b', 2, 'b:2'), ('B', 'b', 3, 'b:3')] with self.create_pipeline() as p: assert_that( p | beam.Create(inputs) | beam.ParDo(AddIndex()), equal_to(expected)) @unittest.skip('TestStream not yet supported') def test_teststream_pardo_timers(self): timer_spec = userstate.TimerSpec('timer', userstate.TimeDomain.WATERMARK) class TimerDoFn(beam.DoFn): def process(self, element, timer=beam.DoFn.TimerParam(timer_spec)): unused_key, ts = element timer.set(ts) timer.set(2 * ts) @userstate.on_timer(timer_spec) def process_timer(self): yield 'fired' ts = ( TestStream().add_elements([('k1', 10)]) # Set timer for 20 .advance_watermark_to(100).add_elements([('k2', 100) ]) # Set timer for 200 .advance_watermark_to(1000)) with self.create_pipeline() as p: _ = ( p | ts | beam.ParDo(TimerDoFn()) | beam.Map(lambda x, ts=beam.DoFn.TimestampParam: (x, ts))) #expected = [('fired', ts) for ts in (20, 200)] #assert_that(actual, equal_to(expected)) def test_pardo_timers(self): timer_spec = userstate.TimerSpec('timer', userstate.TimeDomain.WATERMARK) state_spec = userstate.CombiningValueStateSpec('num_called', sum) class TimerDoFn(beam.DoFn): def process(self, element, timer=beam.DoFn.TimerParam(timer_spec)): unused_key, ts = element timer.set(ts) timer.set(2 * ts) @userstate.on_timer(timer_spec) def process_timer( self, ts=beam.DoFn.TimestampParam, timer=beam.DoFn.TimerParam(timer_spec), state=beam.DoFn.StateParam(state_spec)): if state.read() == 0: state.add(1) timer.set(timestamp.Timestamp(micros=2 * ts.micros)) yield 'fired' with self.create_pipeline() as p: actual = ( p | beam.Create([('k1', 10), ('k2', 100)]) | beam.ParDo(TimerDoFn()) | beam.Map(lambda x, ts=beam.DoFn.TimestampParam: (x, ts))) expected = [('fired', ts) for ts in (20, 200, 40, 400)] assert_that(actual, equal_to(expected)) def test_pardo_timers_clear(self): timer_spec = userstate.TimerSpec('timer', userstate.TimeDomain.WATERMARK) clear_timer_spec = userstate.TimerSpec( 'clear_timer', userstate.TimeDomain.WATERMARK) class TimerDoFn(beam.DoFn): def process( self, element, timer=beam.DoFn.TimerParam(timer_spec), clear_timer=beam.DoFn.TimerParam(clear_timer_spec)): unused_key, ts = element timer.set(ts) timer.set(2 * ts) clear_timer.set(ts) clear_timer.clear() @userstate.on_timer(timer_spec) def process_timer(self): yield 'fired' @userstate.on_timer(clear_timer_spec) def process_clear_timer(self): yield 'should not fire' with self.create_pipeline() as p: actual = ( p | beam.Create([('k1', 10), ('k2', 100)]) | beam.ParDo(TimerDoFn()) | beam.Map(lambda x, ts=beam.DoFn.TimestampParam: (x, ts))) expected = [('fired', ts) for ts in (20, 200)] assert_that(actual, equal_to(expected)) def _run_pardo_et_timer_test( self, n, timer_delay, reset_count=True, clear_timer=True, expected=None): class EventTimeTimerDoFn(beam.DoFn): COUNT = userstate.ReadModifyWriteStateSpec( 'count', coders.VarInt32Coder()) # event-time timer TIMER = userstate.TimerSpec('timer', userstate.TimeDomain.WATERMARK) def __init__(self): self._n = n self._timer_delay = timer_delay self._reset_count = reset_count self._clear_timer = clear_timer def process( self, element_pair, t=beam.DoFn.TimestampParam, count=beam.DoFn.StateParam(COUNT), timer=beam.DoFn.TimerParam(TIMER)): local_count = count.read() or 0 local_count += 1 _LOGGER.info( "get element %s, count=%d", str(element_pair[1]), local_count) if local_count == 1: _LOGGER.info("set timer to %s", str(t + self._timer_delay)) timer.set(t + self._timer_delay) if local_count == self._n: if self._reset_count: _LOGGER.info("reset count") local_count = 0 # don't need the timer now if self._clear_timer: _LOGGER.info("clear timer") timer.clear() count.write(local_count) @userstate.on_timer(TIMER) def timer_callback(self, t=beam.DoFn.TimestampParam): _LOGGER.error("Timer should not fire here") _LOGGER.info("timer callback start (timestamp=%s)", str(t)) yield "fired" with self.create_pipeline() as p: actual = ( p | PeriodicImpulse( start_timestamp=timestamp.Timestamp.now(), stop_timestamp=timestamp.Timestamp.now() + 14, fire_interval=1) | beam.WithKeys(0) | beam.ParDo(EventTimeTimerDoFn())) assert_that(actual, equal_to(expected)) def test_pardo_et_timer_with_no_firing(self): if type(self).__name__ in {'FnApiRunnerTest', 'FnApiRunnerTestWithGrpc', 'FnApiRunnerTestWithGrpcAndMultiWorkers', 'FnApiRunnerTestWithDisabledCaching', 'FnApiRunnerTestWithMultiWorkers', 'FnApiRunnerTestWithBundleRepeat', 'FnApiRunnerTestWithBundleRepeatAndMultiWorkers', 'SparkRunnerTest'}: raise unittest.SkipTest("https://github.com/apache/beam/issues/35168") # The timer will not fire. It is initially set to T + 10, but then it is # cleared at T + 4 (count == 5), and reset to T + 5 + 10 # (count is reset every 5 seconds). self._run_pardo_et_timer_test(5, 10, True, True, []) def test_pardo_et_timer_with_no_reset(self): if type(self).__name__ in {'FnApiRunnerTest', 'FnApiRunnerTestWithGrpc', 'FnApiRunnerTestWithGrpcAndMultiWorkers', 'FnApiRunnerTestWithDisabledCaching', 'FnApiRunnerTestWithMultiWorkers', 'FnApiRunnerTestWithBundleRepeat', 'FnApiRunnerTestWithBundleRepeatAndMultiWorkers', 'SparkRunnerTest'}: raise unittest.SkipTest("https://github.com/apache/beam/issues/35168") # The timer will not fire. It is initially set to T + 10, and then it is # cleared at T + 4 and never set again (count is not reset). self._run_pardo_et_timer_test(5, 10, False, True, []) def test_pardo_et_timer_with_no_reset_and_no_clear(self): if type(self).__name__ in {'FnApiRunnerTest', 'FnApiRunnerTestWithGrpc', 'FnApiRunnerTestWithGrpcAndMultiWorkers', 'FnApiRunnerTestWithDisabledCaching', 'FnApiRunnerTestWithMultiWorkers', 'FnApiRunnerTestWithBundleRepeat', 'FnApiRunnerTestWithBundleRepeatAndMultiWorkers', 'SparkRunnerTest'}: raise unittest.SkipTest("https://github.com/apache/beam/issues/35168") # The timer will fire at T + 10. After the timer is set, it is never # cleared or set again. self._run_pardo_et_timer_test(5, 10, False, False, ["fired"]) def test_pardo_state_timers(self): self._run_pardo_state_timers(windowed=False) def test_pardo_state_timers_non_standard_coder(self): self._run_pardo_state_timers(windowed=False, key_type=Any) def test_windowed_pardo_state_timers(self): self._run_pardo_state_timers(windowed=True) def _run_pardo_state_timers(self, windowed, key_type=None): """ :param windowed: If True, uses an interval window, otherwise a global window :param key_type: Allows to override the inferred key type. This is useful to test the use of non-standard coders, e.g. Python's FastPrimitivesCoder. """ state_spec = userstate.BagStateSpec('state', beam.coders.StrUtf8Coder()) timer_spec = userstate.TimerSpec('timer', userstate.TimeDomain.WATERMARK) elements = list('abcdefgh') key = 'key' buffer_size = 3 class BufferDoFn(beam.DoFn): def process( self, kv, ts=beam.DoFn.TimestampParam, timer=beam.DoFn.TimerParam(timer_spec), state=beam.DoFn.StateParam(state_spec)): _, element = kv state.add(element) buffer = state.read() # For real use, we'd keep track of this size separately. if len(list(buffer)) >= 3: state.clear() yield buffer else: # Set the timer to fire within it's window. timer.set(ts + (1 - timestamp.Duration(micros=1000))) @userstate.on_timer(timer_spec) def process_timer(self, state=beam.DoFn.StateParam(state_spec)): buffer = state.read() state.clear() yield buffer def is_buffered_correctly(actual): # Pickling self in the closure for asserts gives errors (only on jenkins). self = FnApiRunnerTest('__init__') # Acutal should be a grouping of the inputs into batches of size # at most buffer_size, but the actual batching is nondeterministic # based on ordering and trigger firing timing. self.assertEqual( sorted(sum((list(b) for b in actual), [])), elements, actual) self.assertEqual( max(len(list(buffer)) for buffer in actual), buffer_size, actual) if windowed: # Elements were assigned to windows based on their parity. # Assert that each grouping consists of elements belonging to the # same window to ensure states and timers were properly partitioned. for b in actual: parity = set(ord(e) % 2 for e in b) self.assertEqual(1, len(parity), b) with self.create_pipeline() as p: actual = ( p | beam.Create(elements) # Send even and odd elements to different windows. | beam.Map(lambda e: window.TimestampedValue(e, ord(e) % 2)) | beam.WindowInto( window.FixedWindows(1) if windowed else window.GlobalWindows()) | beam.Map(lambda x: (key, x)).with_output_types( tuple[key_type if key_type else type(key), Any]) | beam.ParDo(BufferDoFn())) assert_that(actual, is_buffered_correctly) def test_pardo_dynamic_timer(self): class DynamicTimerDoFn(beam.DoFn): dynamic_timer_spec = userstate.TimerSpec( 'dynamic_timer', userstate.TimeDomain.WATERMARK) def process( self, element, dynamic_timer=beam.DoFn.TimerParam(dynamic_timer_spec)): dynamic_timer.set(element[1], dynamic_timer_tag=element[0]) @userstate.on_timer(dynamic_timer_spec) def dynamic_timer_callback( self, tag=beam.DoFn.DynamicTimerTagParam, timestamp=beam.DoFn.TimestampParam): yield (tag, timestamp) with self.create_pipeline() as p: actual = ( p | beam.Create([('key1', 10), ('key2', 20), ('key3', 30)]) | beam.ParDo(DynamicTimerDoFn())) assert_that(actual, equal_to([('key1', 10), ('key2', 20), ('key3', 30)])) def test_sdf(self): class ExpandingStringsDoFn(beam.DoFn): def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam( ExpandStringsProvider())): assert isinstance(restriction_tracker, RestrictionTrackerView) cur = restriction_tracker.current_restriction().start while restriction_tracker.try_claim(cur): yield element[cur] cur += 1 with self.create_pipeline() as p: data = ['abc', 'defghijklmno', 'pqrstuv', 'wxyz'] actual = (p | beam.Create(data) | beam.ParDo(ExpandingStringsDoFn())) assert_that(actual, equal_to(list(''.join(data)))) def test_sdf_with_dofn_as_restriction_provider(self): class ExpandingStringsDoFn(beam.DoFn, ExpandStringsProvider): def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam()): assert isinstance(restriction_tracker, RestrictionTrackerView) cur = restriction_tracker.current_restriction().start while restriction_tracker.try_claim(cur): yield element[cur] cur += 1 with self.create_pipeline() as p: data = ['abc', 'defghijklmno', 'pqrstuv', 'wxyz'] actual = (p | beam.Create(data) | beam.ParDo(ExpandingStringsDoFn())) assert_that(actual, equal_to(list(''.join(data)))) def test_sdf_with_check_done_failed(self): class ExpandingStringsDoFn(beam.DoFn): def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam( ExpandStringsProvider())): assert isinstance(restriction_tracker, RestrictionTrackerView) cur = restriction_tracker.current_restriction().start while restriction_tracker.try_claim(cur): yield element[cur] cur += 1 return with self.assertRaises(Exception): with self.create_pipeline() as p: data = ['abc', 'defghijklmno', 'pqrstuv', 'wxyz'] _ = (p | beam.Create(data) | beam.ParDo(ExpandingStringsDoFn())) def test_sdf_with_watermark_tracking(self): class ExpandingStringsDoFn(beam.DoFn): def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam( ExpandStringsProvider()), watermark_estimator=beam.DoFn.WatermarkEstimatorParam( ManualWatermarkEstimator.default_provider())): cur = restriction_tracker.current_restriction().start while restriction_tracker.try_claim(cur): watermark_estimator.set_watermark(timestamp.Timestamp(cur)) assert ( watermark_estimator.current_watermark() == timestamp.Timestamp( cur)) yield element[cur] if cur % 2 == 1: restriction_tracker.defer_remainder(timestamp.Duration(micros=5)) return cur += 1 with self.create_pipeline() as p: data = ['abc', 'defghijklmno', 'pqrstuv', 'wxyz'] actual = (p | beam.Create(data) | beam.ParDo(ExpandingStringsDoFn())) assert_that(actual, equal_to(list(''.join(data)))) def test_sdf_with_dofn_as_watermark_estimator(self): class ExpandingStringsDoFn(beam.DoFn, beam.WatermarkEstimatorProvider): def initial_estimator_state(self, element, restriction): return None def create_watermark_estimator(self, state): return beam.io.watermark_estimators.ManualWatermarkEstimator(state) def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam( ExpandStringsProvider()), watermark_estimator=beam.DoFn.WatermarkEstimatorParam( ManualWatermarkEstimator.default_provider())): cur = restriction_tracker.current_restriction().start while restriction_tracker.try_claim(cur): watermark_estimator.set_watermark(timestamp.Timestamp(cur)) assert ( watermark_estimator.current_watermark() == timestamp.Timestamp( cur)) yield element[cur] if cur % 2 == 1: restriction_tracker.defer_remainder(timestamp.Duration(micros=5)) return cur += 1 with self.create_pipeline() as p: data = ['abc', 'defghijklmno', 'pqrstuv', 'wxyz'] actual = (p | beam.Create(data) | beam.ParDo(ExpandingStringsDoFn())) assert_that(actual, equal_to(list(''.join(data)))) def run_sdf_initiated_checkpointing(self, is_drain=False): counter = beam.metrics.Metrics.counter('ns', 'my_counter') class ExpandStringsDoFn(beam.DoFn): def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam( ExpandStringsProvider())): assert isinstance(restriction_tracker, RestrictionTrackerView) cur = restriction_tracker.current_restriction().start while restriction_tracker.try_claim(cur): counter.inc() yield element[cur] if cur % 2 == 1: restriction_tracker.defer_remainder() return cur += 1 with self.create_pipeline(is_drain=is_drain) as p: data = ['abc', 'defghijklmno', 'pqrstuv', 'wxyz'] actual = (p | beam.Create(data) | beam.ParDo(ExpandStringsDoFn())) assert_that(actual, equal_to(list(''.join(data)))) if isinstance(p.runner, fn_api_runner.FnApiRunner): res = p.runner._latest_run_result counters = res.metrics().query( beam.metrics.MetricsFilter().with_name('my_counter'))['counters'] self.assertEqual(1, len(counters)) self.assertEqual(counters[0].committed, len(''.join(data))) def test_sdf_with_sdf_initiated_checkpointing(self): self.run_sdf_initiated_checkpointing(is_drain=False) def test_draining_sdf_with_sdf_initiated_checkpointing(self): self.run_sdf_initiated_checkpointing(is_drain=True) def test_sdf_default_truncate_when_bounded(self): class SimpleSDF(beam.DoFn): def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam( OffsetRangeProvider(use_bounded_offset_range=True))): assert isinstance(restriction_tracker, RestrictionTrackerView) cur = restriction_tracker.current_restriction().start while restriction_tracker.try_claim(cur): yield cur cur += 1 with self.create_pipeline(is_drain=True) as p: actual = (p | beam.Create([10]) | beam.ParDo(SimpleSDF())) assert_that(actual, equal_to(range(10))) def test_sdf_default_truncate_when_unbounded(self): class SimpleSDF(beam.DoFn): def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam( OffsetRangeProvider(use_bounded_offset_range=False))): assert isinstance(restriction_tracker, RestrictionTrackerView) cur = restriction_tracker.current_restriction().start while restriction_tracker.try_claim(cur): yield cur cur += 1 with self.create_pipeline(is_drain=True) as p: actual = (p | beam.Create([10]) | beam.ParDo(SimpleSDF())) assert_that(actual, equal_to([])) def test_sdf_with_truncate(self): class SimpleSDF(beam.DoFn): def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam( OffsetRangeProviderWithTruncate())): assert isinstance(restriction_tracker, RestrictionTrackerView) cur = restriction_tracker.current_restriction().start while restriction_tracker.try_claim(cur): yield cur cur += 1 with self.create_pipeline(is_drain=True) as p: actual = (p | beam.Create([10]) | beam.ParDo(SimpleSDF())) assert_that(actual, equal_to(range(5))) def test_group_by_key(self): with self.create_pipeline() as p: res = ( p | beam.Create([('a', 1), ('a', 2), ('b', 3)]) | beam.GroupByKey() | beam.Map(lambda k_vs: (k_vs[0], sorted(k_vs[1])))) assert_that(res, equal_to([('a', [1, 2]), ('b', [3])])) # Runners may special case the Reshuffle transform urn. def test_reshuffle(self): with self.create_pipeline() as p: assert_that( p | beam.Create([1, 2, 3]) | beam.Reshuffle(), equal_to([1, 2, 3])) def test_reshuffle_after_custom_window(self): with self.create_pipeline() as p: assert_that( p | beam.Create([12, 2, 1]) | beam.Map(lambda t: window.TimestampedValue(t, t)) | beam.WindowInto(beam.transforms.window.FixedWindows(2)) | beam.Reshuffle(), equal_to([12, 2, 1])) def test_flatten(self, with_transcoding=True): with self.create_pipeline() as p: if with_transcoding: # Additional element which does not match with the first type additional = [ord('d')] else: additional = ['d'] res = ( p | 'a' >> beam.Create(['a']), p | 'bc' >> beam.Create(['b', 'c']), p | 'd' >> beam.Create(additional)) | beam.Flatten() assert_that(res, equal_to(['a', 'b', 'c'] + additional)) def test_flatten_same_pcollections(self, with_transcoding=True): with self.create_pipeline() as p: pc = p | beam.Create(['a', 'b']) assert_that((pc, pc, pc) | beam.Flatten(), equal_to(['a', 'b'] * 3)) def test_combine_per_key(self): with self.create_pipeline() as p: res = ( p | beam.Create([('a', 1), ('a', 2), ('b', 3)]) | beam.CombinePerKey(beam.combiners.MeanCombineFn())) assert_that(res, equal_to([('a', 1.5), ('b', 3.0)])) def test_windowed_combine_per_key(self): with self.create_pipeline() as p: input = ( p | beam.Create([12, 2, 1]) | beam.Map(lambda t: window.TimestampedValue(('k', t), t))) fixed = input | 'Fixed' >> ( beam.WindowInto(beam.transforms.window.FixedWindows(10)) | beam.CombinePerKey(beam.combiners.MeanCombineFn())) assert_that(fixed, equal_to([('k', 1.5), ('k', 12)])) sliding = input | 'Sliding' >> ( beam.WindowInto(beam.transforms.window.SlidingWindows(20, 10)) | beam.CombinePerKey(beam.combiners.MeanCombineFn())) assert_that(sliding, equal_to([('k', 1.5), ('k', 5.0), ('k', 12)])) sessions = input | 'Sessions' >> ( beam.WindowInto(beam.transforms.window.Sessions(5)) | beam.CombinePerKey(beam.combiners.MeanCombineFn())) assert_that(sessions, equal_to([('k', 1.5), ('k', 12)])) def test_read(self): # Can't use NamedTemporaryFile as a context # due to https://bugs.python.org/issue14243 temp_file = tempfile.NamedTemporaryFile(delete=False) try: temp_file.write(b'a\nb\nc') temp_file.close() with self.create_pipeline() as p: assert_that( p | beam.io.ReadFromText(temp_file.name), equal_to(['a', 'b', 'c'])) finally: os.unlink(temp_file.name) def test_sliding_windows(self): data = [(timestamp.Timestamp(i), i) for i in range(1, 10)] with self.create_pipeline() as p: ret = ( p | PeriodicImpulse(data=data, fire_interval=1) | beam.WithKeys(0) | beam.WindowInto(beam.transforms.window.SlidingWindows(6, 3)) | beam.GroupByKey()) assert_that( ret, equal_to([(0, [1, 2]), (0, [1, 2, 3, 4, 5]), (0, [3, 4, 5, 6, 7, 8]), (0, [6, 7, 8, 9]), (0, [9])], lambda x, y: x[0] == y[0] and sorted(x[1]) == sorted(y[1]))) def test_windowing(self): with self.create_pipeline() as p: res = ( p | beam.Create([1, 2, 100, 101, 102]) | beam.Map(lambda t: window.TimestampedValue(('k', t), t)) | beam.WindowInto(beam.transforms.window.Sessions(10)) | beam.GroupByKey() | beam.Map(lambda k_vs1: (k_vs1[0], sorted(k_vs1[1])))) assert_that(res, equal_to([('k', [1, 2]), ('k', [100, 101, 102])])) def test_custom_merging_window(self): with self.create_pipeline() as p: res = ( p | beam.Create([1, 2, 100, 101, 102]) | beam.Map(lambda t: window.TimestampedValue(('k', t), t)) | beam.WindowInto(CustomMergingWindowFn()) | beam.GroupByKey() | beam.Map(lambda k_vs1: (k_vs1[0], sorted(k_vs1[1])))) assert_that( res, equal_to([('k', [1]), ('k', [101]), ('k', [2, 100, 102])])) gc.collect() from apache_beam.runners.portability.fn_api_runner.execution import GenericMergingWindowFn self.assertEqual(GenericMergingWindowFn._HANDLES, {}) def test_custom_window_type(self): with self.create_pipeline() as p: res = ( p | beam.Create([1, 2, 100, 101, 102]) | beam.Map(lambda t: window.TimestampedValue(('k', t), t)) | beam.WindowInto(EvenOddWindows()) | beam.GroupByKey() | beam.Map(lambda k_vs1: (k_vs1[0], sorted(k_vs1[1])))) assert_that( res, equal_to([('k', [1]), ('k', [2]), ('k', [101]), ('k', [100, 102])])) @unittest.skip('BEAM-9119: test is flaky') def test_large_elements(self): with self.create_pipeline() as p: big = ( p | beam.Create(['a', 'a', 'b']) | beam.Map(lambda x: (x, x * data_plane._DEFAULT_SIZE_FLUSH_THRESHOLD))) side_input_res = ( big | beam.Map( lambda x, side: (x[0], side.count(x[0])), beam.pvalue.AsList(big | beam.Map(lambda x: x[0])))) assert_that( side_input_res, equal_to([('a', 2), ('a', 2), ('b', 1)]), label='side') gbk_res = (big | beam.GroupByKey() | beam.Map(lambda x: x[0])) assert_that(gbk_res, equal_to(['a', 'b']), label='gbk') def test_error_message_includes_stage(self): with self.assertRaises(BaseException) as e_cm: with self.create_pipeline() as p: def raise_error(x): raise RuntimeError( 'This error is expected and does not indicate a test failure.') # pylint: disable=expression-not-assigned ( p | beam.Create(['a', 'b']) | 'StageA' >> beam.Map(lambda x: x) | 'StageB' >> beam.Map(lambda x: x) | 'StageC' >> beam.Map(raise_error) | 'StageD' >> beam.Map(lambda x: x)) message = e_cm.exception.args[0] self.assertIn('StageC', message) self.assertNotIn('StageB', message) def test_error_traceback_includes_user_code(self): def first(x): return second(x) def second(x): return third(x) def third(x): raise ValueError( 'This error is expected and does not indicate a test failure.') try: with self.create_pipeline() as p: p | beam.Create([0]) | beam.Map(first) # pylint: disable=expression-not-assigned except Exception: # pylint: disable=broad-except message = traceback.format_exc() else: raise AssertionError('expected exception not raised') self.assertIn('first', message) self.assertIn('second', message) self.assertIn('third', message) def test_no_subtransform_composite(self): class First(beam.PTransform): def expand(self, pcolls): return pcolls[0] with self.create_pipeline() as p: pcoll_a = p | 'a' >> beam.Create(['a']) pcoll_b = p | 'b' >> beam.Create(['b']) assert_that((pcoll_a, pcoll_b) | First(), equal_to(['a'])) def test_metrics(self, check_gauge=True, check_bounded_trie=False): p = self.create_pipeline() counter = beam.metrics.Metrics.counter('ns', 'counter') distribution = beam.metrics.Metrics.distribution('ns', 'distribution') gauge = beam.metrics.Metrics.gauge('ns', 'gauge') string_set = beam.metrics.Metrics.string_set('ns', 'string_set') bounded_trie = beam.metrics.Metrics.bounded_trie('ns', 'bounded_trie') elements = ['a', 'zzz'] pcoll = p | beam.Create(elements) # pylint: disable=expression-not-assigned pcoll | 'count1' >> beam.FlatMap(lambda x: counter.inc()) pcoll | 'count2' >> beam.FlatMap(lambda x: counter.inc(len(x))) pcoll | 'dist' >> beam.FlatMap(lambda x: distribution.update(len(x))) pcoll | 'gauge' >> beam.FlatMap(lambda x: gauge.set(3)) pcoll | 'string_set' >> beam.FlatMap(lambda x: string_set.add(x)) pcoll | 'bounded_trie' >> beam.FlatMap(lambda x: bounded_trie.add(tuple(x))) res = p.run() res.wait_until_finish() t1, t2 = res.metrics().query(beam.metrics.MetricsFilter() .with_name('counter'))['counters'] self.assertEqual(t1.committed + t2.committed, 6) dist, = res.metrics().query(beam.metrics.MetricsFilter() .with_name('distribution'))['distributions'] self.assertEqual( dist.committed.data, beam.metrics.cells.DistributionData(4, 2, 1, 3)) self.assertEqual(dist.committed.mean, 2.0) if check_gauge: gaug, = res.metrics().query(beam.metrics.MetricsFilter() .with_name('gauge'))['gauges'] self.assertEqual(gaug.committed.value, 3) str_set, = res.metrics().query(beam.metrics.MetricsFilter() .with_name('string_set'))['string_sets'] self.assertEqual(str_set.committed, set(elements)) if check_bounded_trie: bounded_trie, = res.metrics().query(beam.metrics.MetricsFilter() .with_name('bounded_trie'))['bounded_tries'] self.assertEqual(bounded_trie.committed.size(), 2) for element in elements: self.assertTrue( bounded_trie.committed.contains(tuple(element)), element) def test_callbacks_with_exception(self): elements_list = ['1', '2'] def raise_expetion(): raise Exception('raise exception when calling callback') class FinalizebleDoFnWithException(beam.DoFn): def process( self, element, bundle_finalizer=beam.DoFn.BundleFinalizerParam): bundle_finalizer.register(raise_expetion) yield element with self.create_pipeline() as p: res = ( p | beam.Create(elements_list) | beam.ParDo(FinalizebleDoFnWithException())) assert_that(res, equal_to(['1', '2'])) def test_register_finalizations(self): event_recorder = EventRecorder(tempfile.gettempdir()) class FinalizableSplittableDoFn(beam.DoFn): def process( self, element, bundle_finalizer=beam.DoFn.BundleFinalizerParam, restriction_tracker=beam.DoFn.RestrictionParam( OffsetRangeProvider( use_bounded_offset_range=True, checkpoint_only=True))): # We use SDF to enforce finalization call happens by using # self-initiated checkpoint. if 'finalized' in event_recorder.events(): restriction_tracker.try_claim( restriction_tracker.current_restriction().start) yield element restriction_tracker.try_claim(element) return if restriction_tracker.try_claim( restriction_tracker.current_restriction().start): bundle_finalizer.register(lambda: event_recorder.record('finalized')) # We sleep here instead of setting a resume time since the resume time # doesn't need to be honored. time.sleep(1) restriction_tracker.defer_remainder() with self.create_pipeline() as p: max_retries = 100 res = ( p | beam.Create([max_retries]) | beam.ParDo(FinalizableSplittableDoFn())) assert_that(res, equal_to([max_retries])) event_recorder.cleanup() def test_sdf_synthetic_source(self): common_attrs = { 'key_size': 1, 'value_size': 1, 'initial_splitting_num_bundles': 2, 'initial_splitting_desired_bundle_size': 2, 'sleep_per_input_record_sec': 0, 'initial_splitting': 'const' } num_source_description = 5 min_num_record = 10 max_num_record = 20 # pylint: disable=unused-variable source_descriptions = ([ dict({'num_records': random.randint(min_num_record, max_num_record)}, **common_attrs) for i in range(0, num_source_description) ]) total_num_records = 0 for source in source_descriptions: total_num_records += source['num_records'] with self.create_pipeline() as p: res = ( p | beam.Create(source_descriptions) | beam.ParDo(SyntheticSDFAsSource()) | beam.combiners.Count.Globally()) assert_that(res, equal_to([total_num_records])) def test_create_value_provider_pipeline_option(self): # Verify that the runner can execute a pipeline when there are value # provider pipeline options # pylint: disable=unused-variable class FooOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument( "--foo", help='a value provider argument', default="bar") RuntimeValueProvider.set_runtime_options({}) with self.create_pipeline() as p: assert_that(p | beam.Create(['a', 'b']), equal_to(['a', 'b'])) def _test_pack_combiners(self, assert_using_counter_names): counter = beam.metrics.Metrics.counter('ns', 'num_values') def min_with_counter(values): counter.inc() return min(values) def max_with_counter(values): counter.inc() return max(values) class PackableCombines(beam.PTransform): def annotations(self): return {python_urns.APPLY_COMBINER_PACKING: b''} def expand(self, pcoll): assert_that( pcoll | 'PackableMin' >> beam.CombineGlobally(min_with_counter), equal_to([10]), label='AssertMin') assert_that( pcoll | 'PackableMax' >> beam.CombineGlobally(max_with_counter), equal_to([30]), label='AssertMax') with self.create_pipeline() as p: _ = p | beam.Create([10, 20, 30]) | PackableCombines() res = p.result packed_step_name_regex = ( r'.*Packed.*PackableMin.*CombinePerKey.*PackableMax.*CombinePerKey.*' + 'Pack.*') counters = res.metrics().query(beam.metrics.MetricsFilter())['counters'] step_names = set(m.key.step for m in counters if m.key.step) pipeline_options = p._options if assert_using_counter_names: if pipeline_options.view_as(StandardOptions).streaming: self.assertFalse( any(re.match(packed_step_name_regex, s) for s in step_names)) else: self.assertTrue( any(re.match(packed_step_name_regex, s) for s in step_names)) def test_pack_combiners(self): self._test_pack_combiners(assert_using_counter_names=True) def test_group_by_key_with_empty_pcoll_elements(self): with self.create_pipeline() as p: res = ( p | beam.Create([('test_key', 'test_value')]) | beam.Filter(lambda x: False) | beam.GroupByKey()) assert_that(res, equal_to([])) def test_first_pane(self): with self.create_pipeline() as p: res = ( p | beam.Create([1, 2]) | beam.WithKeys(0) | beam.WindowInto( window.GlobalWindows(), trigger=trigger.Repeatedly(trigger.AfterCount(1)), accumulation_mode=trigger.AccumulationMode.ACCUMULATING, allowed_lateness=0, ) | beam.GroupByKey() | beam.Values()) has_at_least_one(res, lambda e, t, w, p: p.is_first) has_at_least_one(res, lambda e, t, w, p: p.index == 0) # These tests are kept in a separate group so that they are # not ran in the FnApiRunnerTestWithBundleRepeat which repeats # bundle processing. This breaks the byte sampling metrics as # it makes the probability of sampling far too small # upon repeating bundle processing due to unncessarily incrementing # the sampling counter. class FnApiRunnerMetricsTest(unittest.TestCase): def assert_has_counter( self, mon_infos, urn, labels, value=None, ge_value=None): # TODO(ajamato): Consider adding a matcher framework found = 0 matches = [] for mi in mon_infos: if has_urn_and_labels(mi, urn, labels): extracted_value = monitoring_infos.extract_counter_value(mi) if ge_value is not None: if extracted_value >= ge_value: found = found + 1 elif value is not None: if extracted_value == value: found = found + 1 else: found = found + 1 ge_value_str = {'ge_value': ge_value} if ge_value else '' value_str = {'value': value} if value else '' self.assertEqual( 1, found, "Found (%s, %s) Expected only 1 monitoring_info for %s." % ( found, matches, (urn, labels, value_str, ge_value_str), )) def assert_has_distribution( self, mon_infos, urn, labels, sum=None, count=None, min=None, max=None): # TODO(ajamato): Consider adding a matcher framework sum = _matcher_or_equal_to(sum) count = _matcher_or_equal_to(count) min = _matcher_or_equal_to(min) max = _matcher_or_equal_to(max) found = 0 description = StringDescription() for mi in mon_infos: if has_urn_and_labels(mi, urn, labels): (extracted_count, extracted_sum, extracted_min, extracted_max) = monitoring_infos.extract_distribution(mi) increment = 1 if sum is not None: description.append_text(' sum: ') sum.describe_to(description) if not sum.matches(extracted_sum): increment = 0 if count is not None: description.append_text(' count: ') count.describe_to(description) if not count.matches(extracted_count): increment = 0 if min is not None: description.append_text(' min: ') min.describe_to(description) if not min.matches(extracted_min): increment = 0 if max is not None: description.append_text(' max: ') max.describe_to(description) if not max.matches(extracted_max): increment = 0 found += increment self.assertEqual( 1, found, "Found (%s) Expected only 1 monitoring_info for %s." % ( found, (urn, labels, str(description)), )) def create_pipeline(self): return beam.Pipeline(runner=fn_api_runner.FnApiRunner()) def test_element_count_metrics(self): class GenerateTwoOutputs(beam.DoFn): def process(self, element): yield str(element) + '1' yield beam.pvalue.TaggedOutput('SecondOutput', str(element) + '2') yield beam.pvalue.TaggedOutput('SecondOutput', str(element) + '2') yield beam.pvalue.TaggedOutput('ThirdOutput', str(element) + '3') class PassThrough(beam.DoFn): def process(self, element): yield element p = self.create_pipeline() # Produce enough elements to make sure byte sampling occurs. num_source_elems = 100 pcoll = p | beam.Create(['a%d' % i for i in range(num_source_elems)], reshuffle=False) # pylint: disable=expression-not-assigned pardo = ( 'StepThatDoesTwoOutputs' >> beam.ParDo( GenerateTwoOutputs()).with_outputs( 'SecondOutput', 'ThirdOutput', main='FirstAndMainOutput')) # Actually feed pcollection to pardo second_output, third_output, first_output = (pcoll | pardo) # consume some of elements merged = ((first_output, second_output, third_output) | beam.Flatten()) merged | ('PassThrough') >> beam.ParDo(PassThrough()) second_output | ('PassThrough2') >> beam.ParDo(PassThrough()) res = p.run() res.wait_until_finish() result_metrics = res.monitoring_metrics() counters = result_metrics.monitoring_infos() # All element count and byte count metrics must have a PCOLLECTION_LABEL. self.assertFalse([ x for x in counters if x.urn in [ monitoring_infos.ELEMENT_COUNT_URN, monitoring_infos.SAMPLED_BYTE_SIZE_URN ] and monitoring_infos.PCOLLECTION_LABEL not in x.labels ]) try: labels = { monitoring_infos.PCOLLECTION_LABEL: 'ref_PCollection_PCollection_1' } self.assert_has_counter( counters, monitoring_infos.ELEMENT_COUNT_URN, labels, 1) # Create output. labels = { monitoring_infos.PCOLLECTION_LABEL: 'ref_PCollection_PCollection_3' } self.assert_has_counter( counters, monitoring_infos.ELEMENT_COUNT_URN, labels, num_source_elems) self.assert_has_distribution( counters, monitoring_infos.SAMPLED_BYTE_SIZE_URN, labels, min=hamcrest.greater_than(0), max=hamcrest.greater_than(0), sum=hamcrest.greater_than(0), count=hamcrest.greater_than(0)) # GenerateTwoOutputs, main output. labels = { monitoring_infos.PCOLLECTION_LABEL: 'ref_PCollection_PCollection_4' } self.assert_has_counter( counters, monitoring_infos.ELEMENT_COUNT_URN, labels, num_source_elems) self.assert_has_distribution( counters, monitoring_infos.SAMPLED_BYTE_SIZE_URN, labels, min=hamcrest.greater_than(0), max=hamcrest.greater_than(0), sum=hamcrest.greater_than(0), count=hamcrest.greater_than(0)) # GenerateTwoOutputs, "SecondOutput" output. labels = { monitoring_infos.PCOLLECTION_LABEL: 'ref_PCollection_PCollection_5' } self.assert_has_counter( counters, monitoring_infos.ELEMENT_COUNT_URN, labels, 2 * num_source_elems) self.assert_has_distribution( counters, monitoring_infos.SAMPLED_BYTE_SIZE_URN, labels, min=hamcrest.greater_than(0), max=hamcrest.greater_than(0), sum=hamcrest.greater_than(0), count=hamcrest.greater_than(0)) # GenerateTwoOutputs, "ThirdOutput" output. labels = { monitoring_infos.PCOLLECTION_LABEL: 'ref_PCollection_PCollection_6' } self.assert_has_counter( counters, monitoring_infos.ELEMENT_COUNT_URN, labels, num_source_elems) self.assert_has_distribution( counters, monitoring_infos.SAMPLED_BYTE_SIZE_URN, labels, min=hamcrest.greater_than(0), max=hamcrest.greater_than(0), sum=hamcrest.greater_than(0), count=hamcrest.greater_than(0)) # Skipping other pcollections due to non-deterministic naming for multiple # outputs. # Flatten/Read, main output. labels = { monitoring_infos.PCOLLECTION_LABEL: 'ref_PCollection_PCollection_7' } self.assert_has_counter( counters, monitoring_infos.ELEMENT_COUNT_URN, labels, 4 * num_source_elems) self.assert_has_distribution( counters, monitoring_infos.SAMPLED_BYTE_SIZE_URN, labels, min=hamcrest.greater_than(0), max=hamcrest.greater_than(0), sum=hamcrest.greater_than(0), count=hamcrest.greater_than(0)) # PassThrough, main output labels = { monitoring_infos.PCOLLECTION_LABEL: 'ref_PCollection_PCollection_8' } self.assert_has_counter( counters, monitoring_infos.ELEMENT_COUNT_URN, labels, 4 * num_source_elems) self.assert_has_distribution( counters, monitoring_infos.SAMPLED_BYTE_SIZE_URN, labels, min=hamcrest.greater_than(0), max=hamcrest.greater_than(0), sum=hamcrest.greater_than(0), count=hamcrest.greater_than(0)) # PassThrough2, main output labels = { monitoring_infos.PCOLLECTION_LABEL: 'ref_PCollection_PCollection_9' } self.assert_has_counter( counters, monitoring_infos.ELEMENT_COUNT_URN, labels, num_source_elems) self.assert_has_distribution( counters, monitoring_infos.SAMPLED_BYTE_SIZE_URN, labels, min=hamcrest.greater_than(0), max=hamcrest.greater_than(0), sum=hamcrest.greater_than(0), count=hamcrest.greater_than(0)) except: print(res._monitoring_infos_by_stage) raise def test_non_user_metrics(self): p = self.create_pipeline() pcoll = p | beam.Create(['a', 'zzz']) # pylint: disable=expression-not-assigned pcoll | 'MyStep' >> beam.FlatMap(lambda x: None) res = p.run() res.wait_until_finish() result_metrics = res.monitoring_metrics() all_metrics_via_montoring_infos = result_metrics.query() def assert_counter_exists(metrics, namespace, name, step): found = 0 metric_key = MetricKey(step, MetricName(namespace, name)) for m in metrics['counters']: if m.key == metric_key: found = found + 1 self.assertEqual( 1, found, "Did not find exactly 1 metric for %s." % metric_key) urns = [ monitoring_infos.START_BUNDLE_MSECS_URN, monitoring_infos.PROCESS_BUNDLE_MSECS_URN, monitoring_infos.FINISH_BUNDLE_MSECS_URN, monitoring_infos.TOTAL_MSECS_URN, ] for urn in urns: split = urn.split(':') namespace = split[0] name = ':'.join(split[1:]) assert_counter_exists( all_metrics_via_montoring_infos, namespace, name, step='Create/Impulse') assert_counter_exists( all_metrics_via_montoring_infos, namespace, name, step='MyStep') # Due to somewhat non-deterministic nature of state sampling and sleep, # this test is flaky when state duration is low. # Since increasing state duration significantly would also slow down # the test suite, we are retrying twice on failure as a mitigation. @retry(reraise=True, stop=stop_after_attempt(3)) def test_progress_metrics(self): p = self.create_pipeline() _ = ( p | beam.Create([0, 0, 0, 5e-3 * DEFAULT_SAMPLING_PERIOD_MS], reshuffle=False) | beam.Map(time.sleep) | beam.Map(lambda x: ('key', x)) | beam.GroupByKey() | 'm_out' >> beam.FlatMap( lambda x: [ 1, 2, 3, 4, 5, beam.pvalue.TaggedOutput('once', x), beam.pvalue. TaggedOutput('twice', x), beam.pvalue.TaggedOutput('twice', x) ])) res = p.run() res.wait_until_finish() def has_mi_for_ptransform(mon_infos, ptransform): for mi in mon_infos: if ptransform in mi.labels[monitoring_infos.PTRANSFORM_LABEL]: return True return False try: # Test the new MonitoringInfo monitoring format. self.assertEqual(3, len(res._monitoring_infos_by_stage)) pregbk_mis, postgbk_mis = [ mi for stage, mi in res._monitoring_infos_by_stage.items() if stage] if not has_mi_for_ptransform(pregbk_mis, 'Create/Map(decode)'): # The monitoring infos above are actually unordered. Swap. pregbk_mis, postgbk_mis = postgbk_mis, pregbk_mis # pregbk monitoring infos labels = { monitoring_infos.PCOLLECTION_LABEL: 'ref_PCollection_PCollection_3' } self.assert_has_counter( pregbk_mis, monitoring_infos.ELEMENT_COUNT_URN, labels, value=4) self.assert_has_distribution( pregbk_mis, monitoring_infos.SAMPLED_BYTE_SIZE_URN, labels) labels = { monitoring_infos.PCOLLECTION_LABEL: 'ref_PCollection_PCollection_4' } self.assert_has_counter( pregbk_mis, monitoring_infos.ELEMENT_COUNT_URN, labels, value=4) self.assert_has_distribution( pregbk_mis, monitoring_infos.SAMPLED_BYTE_SIZE_URN, labels) labels = {monitoring_infos.PTRANSFORM_LABEL: 'Map(sleep)'} self.assert_has_counter( pregbk_mis, monitoring_infos.TOTAL_MSECS_URN, labels, ge_value=4 * DEFAULT_SAMPLING_PERIOD_MS) # postgbk monitoring infos labels = { monitoring_infos.PCOLLECTION_LABEL: 'ref_PCollection_PCollection_6' } self.assert_has_counter( postgbk_mis, monitoring_infos.ELEMENT_COUNT_URN, labels, value=1) self.assert_has_distribution( postgbk_mis, monitoring_infos.SAMPLED_BYTE_SIZE_URN, labels) labels = { monitoring_infos.PCOLLECTION_LABEL: 'ref_PCollection_PCollection_7' } self.assert_has_counter( postgbk_mis, monitoring_infos.ELEMENT_COUNT_URN, labels, value=5) self.assert_has_distribution( postgbk_mis, monitoring_infos.SAMPLED_BYTE_SIZE_URN, labels) except: print(res._monitoring_infos_by_stage) raise class FnApiRunnerTestWithGrpc(FnApiRunnerTest): def create_pipeline(self, is_drain=False): return beam.Pipeline( runner=fn_api_runner.FnApiRunner( default_environment=environments.EmbeddedPythonGrpcEnvironment. default(), is_drain=is_drain)) class FnApiRunnerTestWithDisabledCaching(FnApiRunnerTest): def create_pipeline(self, is_drain=False): return beam.Pipeline( runner=fn_api_runner.FnApiRunner( default_environment=environments.EmbeddedPythonGrpcEnvironment( state_cache_size=0, data_buffer_time_limit_ms=0, capabilities=environments.python_sdk_capabilities(), artifacts=()), is_drain=is_drain)) class FnApiRunnerTestWithMultiWorkers(FnApiRunnerTest): def create_pipeline(self, is_drain=False): pipeline_options = PipelineOptions(direct_num_workers=2) p = beam.Pipeline( runner=fn_api_runner.FnApiRunner(is_drain=is_drain), options=pipeline_options) #TODO(https://github.com/apache/beam/issues/19936): Fix these tests. p._options.view_as(DebugOptions).experiments.remove('beam_fn_api') return p def test_metrics(self): raise unittest.SkipTest("This test is for a single worker only.") def test_sdf_with_sdf_initiated_checkpointing(self): raise unittest.SkipTest("This test is for a single worker only.") def test_draining_sdf_with_sdf_initiated_checkpointing(self): raise unittest.SkipTest("This test is for a single worker only.") def test_sdf_with_watermark_tracking(self): raise unittest.SkipTest("This test is for a single worker only.") def test_sdf_with_dofn_as_watermark_estimator(self): raise unittest.SkipTest("This test is for a single worker only.") def test_register_finalizations(self): raise unittest.SkipTest("This test is for a single worker only.") def test_sliding_windows(self): raise unittest.SkipTest("This test is for a single worker only.") class FnApiRunnerTestWithGrpcAndMultiWorkers(FnApiRunnerTest): def create_pipeline(self, is_drain=False): pipeline_options = PipelineOptions( direct_num_workers=2, direct_running_mode='multi_threading') p = beam.Pipeline( runner=fn_api_runner.FnApiRunner(is_drain=is_drain), options=pipeline_options) #TODO(https://github.com/apache/beam/issues/19936): Fix these tests. p._options.view_as(DebugOptions).experiments.remove('beam_fn_api') return p def test_metrics(self): raise unittest.SkipTest("This test is for a single worker only.") def test_sdf_with_sdf_initiated_checkpointing(self): raise unittest.SkipTest("This test is for a single worker only.") def test_draining_sdf_with_sdf_initiated_checkpointing(self): raise unittest.SkipTest("This test is for a single worker only.") def test_sdf_with_watermark_tracking(self): raise unittest.SkipTest("This test is for a single worker only.") def test_sdf_with_dofn_as_watermark_estimator(self): raise unittest.SkipTest("This test is for a single worker only.") def test_register_finalizations(self): raise unittest.SkipTest("This test is for a single worker only.") def test_sliding_windows(self): raise unittest.SkipTest("This test is for a single worker only.") class FnApiRunnerTestWithBundleRepeat(FnApiRunnerTest): def create_pipeline(self, is_drain=False): return beam.Pipeline( runner=fn_api_runner.FnApiRunner(bundle_repeat=3, is_drain=is_drain)) def test_register_finalizations(self): raise unittest.SkipTest("TODO: Avoid bundle finalizations on repeat.") class FnApiRunnerTestWithBundleRepeatAndMultiWorkers(FnApiRunnerTest): def create_pipeline(self, is_drain=False): pipeline_options = PipelineOptions(direct_num_workers=2) p = beam.Pipeline( runner=fn_api_runner.FnApiRunner(bundle_repeat=3, is_drain=is_drain), options=pipeline_options) #TODO(https://github.com/apache/beam/issues/19936): Fix these tests. p._options.view_as(DebugOptions).experiments.remove('beam_fn_api') return p def test_register_finalizations(self): raise unittest.SkipTest("TODO: Avoid bundle finalizations on repeat.") def test_metrics(self): raise unittest.SkipTest("This test is for a single worker only.") def test_sdf_with_sdf_initiated_checkpointing(self): raise unittest.SkipTest("This test is for a single worker only.") def test_draining_sdf_with_sdf_initiated_checkpointing(self): raise unittest.SkipTest("This test is for a single worker only.") def test_sdf_with_watermark_tracking(self): raise unittest.SkipTest("This test is for a single worker only.") def test_sdf_with_dofn_as_watermark_estimator(self): raise unittest.SkipTest("This test is for a single worker only.") def test_sliding_windows(self): raise unittest.SkipTest("This test is for a single worker only.") class FnApiRunnerSplitTest(unittest.TestCase): def create_pipeline(self, is_drain=False): # Must be GRPC so we can send data and split requests concurrent # to the bundle process request. return beam.Pipeline( runner=fn_api_runner.FnApiRunner( default_environment=environments.EmbeddedPythonGrpcEnvironment. default(), is_drain=is_drain)) def test_checkpoint(self): # This split manager will get re-invoked on each smaller split, # so N times for N elements. element_counter = ElementCounter() def split_manager(num_elements): # Send at least one element so it can make forward progress. element_counter.reset() breakpoint = element_counter.set_breakpoint(1) # Cede control back to the runner so data can be sent. yield breakpoint.wait() # Split as close to current as possible. split_result = yield 0.0 # Verify we split at exactly the first element. self.verify_channel_split(split_result, 0, 1) # Continue processing. breakpoint.clear() self.run_split_pipeline(split_manager, list('abc'), element_counter) def test_split_half(self): total_num_elements = 25 seen_bundle_sizes = [] element_counter = ElementCounter() def split_manager(num_elements): seen_bundle_sizes.append(num_elements) if num_elements == total_num_elements: element_counter.reset() breakpoint = element_counter.set_breakpoint(5) yield breakpoint.wait() # Split the remainder (20, then 10, elements) in half. split1 = yield 0.5 self.verify_channel_split(split1, 14, 15) # remainder is 15 to end split2 = yield 0.5 self.verify_channel_split(split2, 9, 10) # remainder is 10 to end breakpoint.clear() self.run_split_pipeline( split_manager, range(total_num_elements), element_counter) self.assertEqual([25, 15], seen_bundle_sizes) def run_split_pipeline(self, split_manager, elements, element_counter=None): with fn_runner.split_manager('Identity', split_manager): with self.create_pipeline() as p: res = ( p | beam.Create(elements) | beam.Reshuffle() | 'Identity' >> beam.Map(lambda x: x) | beam.Map(lambda x: element_counter.increment() or x)) assert_that(res, equal_to(elements)) def run_sdf_checkpoint(self, is_drain=False): element_counter = ElementCounter() def split_manager(num_elements): if num_elements > 0: element_counter.reset() breakpoint = element_counter.set_breakpoint(1) yield breakpoint.wait() yield 0 breakpoint.clear() # Everything should be perfectly split. elements = [2, 3] expected_groups = [[(2, 0)], [(2, 1)], [(3, 0)], [(3, 1)], [(3, 2)]] self.run_sdf_split_pipeline( split_manager, elements, element_counter, expected_groups, is_drain=is_drain) def run_sdf_split_half(self, is_drain=False): element_counter = ElementCounter() is_first_bundle = True def split_manager(num_elements): nonlocal is_first_bundle if is_first_bundle and num_elements > 0: is_first_bundle = False breakpoint = element_counter.set_breakpoint(1) yield breakpoint.wait() split1 = yield 0.5 split2 = yield 0.5 split3 = yield 0.5 self.verify_channel_split(split1, 0, 1) self.verify_channel_split(split2, -1, 1) self.verify_channel_split(split3, -1, 1) breakpoint.clear() elements = [4, 4] expected_groups = [[(4, 0)], [(4, 1)], [(4, 2), (4, 3)], [(4, 0), (4, 1), (4, 2), (4, 3)]] self.run_sdf_split_pipeline( split_manager, elements, element_counter, expected_groups, is_drain=is_drain) def run_split_crazy_sdf(self, seed=None, is_drain=False): if seed is None: seed = random.randrange(1 << 20) r = random.Random(seed) element_counter = ElementCounter() def split_manager(num_elements): if num_elements > 0: element_counter.reset() wait_for = r.randrange(num_elements) breakpoint = element_counter.set_breakpoint(wait_for) yield breakpoint.wait() yield r.random() yield r.random() breakpoint.clear() try: elements = [r.randrange(5, 10) for _ in range(5)] self.run_sdf_split_pipeline( split_manager, elements, element_counter, is_drain=is_drain) except Exception: _LOGGER.error('test_split_crazy_sdf.seed = %s', seed) raise def test_nosplit_sdf(self): def split_manager(num_elements): yield elements = [1, 2, 3] expected_groups = [[(e, k) for k in range(e)] for e in elements] self.run_sdf_split_pipeline( split_manager, elements, ElementCounter(), expected_groups) def test_checkpoint_sdf(self): self.run_sdf_checkpoint(is_drain=False) def test_checkpoint_draining_sdf(self): self.run_sdf_checkpoint(is_drain=True) def test_split_half_sdf(self): self.run_sdf_split_half(is_drain=False) def test_split_half_draining_sdf(self): self.run_sdf_split_half(is_drain=True) def test_split_crazy_sdf(self, seed=None): self.run_split_crazy_sdf(seed=seed, is_drain=False) def test_split_crazy_draining_sdf(self, seed=None): self.run_split_crazy_sdf(seed=seed, is_drain=True) def run_sdf_split_pipeline( self, split_manager, elements, element_counter, expected_groups=None, is_drain=False): # Define an SDF that for each input x produces [(x, k) for k in range(x)]. class EnumerateProvider(beam.transforms.core.RestrictionProvider): def initial_restriction(self, element): return restriction_trackers.OffsetRange(0, element) def create_tracker(self, restriction): return restriction_trackers.OffsetRestrictionTracker(restriction) def split(self, element, restriction): # Don't do any initial splitting to simplify test. return [restriction] def restriction_size(self, element, restriction): return restriction.size() def is_bounded(self): return True class EnumerateSdf(beam.DoFn): def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam(EnumerateProvider())): to_emit = [] cur = restriction_tracker.current_restriction().start while restriction_tracker.try_claim(cur): to_emit.append((element, cur)) element_counter.increment() cur += 1 # Emitting in batches for tighter testing. yield to_emit expected = [(e, k) for e in elements for k in range(e)] with fn_runner.split_manager('SDF', split_manager): with self.create_pipeline(is_drain=is_drain) as p: grouped = ( p | beam.Create(elements, reshuffle=False) | 'SDF' >> beam.ParDo(EnumerateSdf())) flat = grouped | beam.FlatMap(lambda x: x) assert_that(flat, equal_to(expected)) if expected_groups: assert_that(grouped, equal_to(expected_groups), label='CheckGrouped') def test_time_based_split_manager(self): elements = [str(x) for x in range(100)] class BundleCountingDoFn(beam.DoFn): def process(self, element): time.sleep(0.005) yield element def finish_bundle(self): yield window.GlobalWindows.windowed_value('endOfBundle') with self.create_pipeline() as p: p._options.view_as(DirectOptions).direct_test_splits = { 'SplitMarker': { 'timings': [0, .05], 'fractions': [0.5, 0.5] } } assert_that( p | beam.Create(elements) | 'SplitMarker' >> beam.ParDo(BundleCountingDoFn()), # We split the first bundle twice (once at 50%, and again at 50% of # what was left). All returned split remainders get processed # (together) in a (single) subsequent bundle. equal_to(elements + ['endOfBundle'] * 2)) def verify_channel_split(self, split_result, last_primary, first_residual): self.assertEqual(1, len(split_result.channel_splits), split_result) channel_split, = split_result.channel_splits self.assertEqual(last_primary, channel_split.last_primary_element) self.assertEqual(first_residual, channel_split.first_residual_element) # There should be a primary and residual application for each element # not covered above. self.assertEqual( first_residual - last_primary - 1, len(split_result.primary_roots), split_result.primary_roots) self.assertEqual( first_residual - last_primary - 1, len(split_result.residual_roots), split_result.residual_roots) class ElementCounter(object): """Used to wait until a certain number of elements are seen.""" def __init__(self): self._cv = threading.Condition() self.reset() def reset(self): with self._cv: self._breakpoints = collections.defaultdict(list) self._count = 0 def increment(self): with self._cv: self._count += 1 self._cv.notify_all() breakpoints = list(self._breakpoints[self._count]) for breakpoint in breakpoints: breakpoint.wait() def set_breakpoint(self, value): with self._cv: event = threading.Event() self._breakpoints[value].append(event) class Breakpoint(object): @staticmethod def wait(timeout=10): with self._cv: start = time.time() while self._count < value: elapsed = time.time() - start if elapsed > timeout: raise RuntimeError('Timed out waiting for %s' % value) self._cv.wait(timeout - elapsed) @staticmethod def clear(): event.set() return Breakpoint() def __reduce__(self): # Ensure we get the same element back through a pickling round-trip. name = uuid.uuid4().hex _pickled_element_counters[name] = self return _unpickle_element_counter, (name, ) _pickled_element_counters: dict[str, ElementCounter] = {} def _unpickle_element_counter(name): return _pickled_element_counters[name] class EventRecorder(object): """Used to be registered as a callback in bundle finalization. The reason why records are written into a tmp file is, the in-memory dataset cannot keep callback records when passing into one DoFn. """ def __init__(self, tmp_dir): self.tmp_dir = os.path.join(tmp_dir, uuid.uuid4().hex) os.mkdir(self.tmp_dir) def record(self, content): file_path = os.path.join(self.tmp_dir, uuid.uuid4().hex + '.txt') with open(file_path, 'w') as f: f.write(content) def events(self): content = [] record_files = [ f for f in os.listdir(self.tmp_dir) if os.path.isfile(os.path.join(self.tmp_dir, f)) ] for file in record_files: with open(os.path.join(self.tmp_dir, file), 'r') as f: content.append(f.read()) return sorted(content) def cleanup(self): shutil.rmtree(self.tmp_dir) class ExpandStringsProvider(beam.transforms.core.RestrictionProvider): """A RestrictionProvider that used for sdf related tests.""" def initial_restriction(self, element): return restriction_trackers.OffsetRange(0, len(element)) def create_tracker(self, restriction): return restriction_trackers.OffsetRestrictionTracker(restriction) def split(self, element, restriction): desired_bundle_size = restriction.size() // 2 return restriction.split(desired_bundle_size) def restriction_size(self, element, restriction): return restriction.size() class UnboundedOffsetRestrictionTracker( restriction_trackers.OffsetRestrictionTracker): def is_bounded(self): return False class OffsetRangeProvider(beam.transforms.core.RestrictionProvider): def __init__(self, use_bounded_offset_range, checkpoint_only=False): self.use_bounded_offset_range = use_bounded_offset_range self.checkpoint_only = checkpoint_only def initial_restriction(self, element): return restriction_trackers.OffsetRange(0, element) def create_tracker(self, restriction): if self.checkpoint_only: class CheckpointOnlyOffsetRestrictionTracker( restriction_trackers.OffsetRestrictionTracker): def try_split(self, unused_fraction_of_remainder): return super().try_split(0.0) return CheckpointOnlyOffsetRestrictionTracker(restriction) if self.use_bounded_offset_range: return restriction_trackers.OffsetRestrictionTracker(restriction) return UnboundedOffsetRestrictionTracker(restriction) def split(self, element, restriction): return [restriction] def restriction_size(self, element, restriction): return restriction.size() class OffsetRangeProviderWithTruncate(OffsetRangeProvider): def __init__(self): super().__init__(True) def truncate(self, element, restriction): return restriction_trackers.OffsetRange( restriction.start, restriction.stop // 2) class FnApiBasedLullLoggingTest(unittest.TestCase): def create_pipeline(self): return beam.Pipeline( runner=fn_api_runner.FnApiRunner( default_environment=environments.EmbeddedPythonGrpcEnvironment. default(), progress_request_frequency=0.5)) class StateBackedTestElementType(object): live_element_count = 0 def __init__(self, num_elements, unused): self.num_elements = num_elements StateBackedTestElementType.live_element_count += 1 # Due to using state backed iterable, we expect there is a few instances # alive at any given time. if StateBackedTestElementType.live_element_count > 5: raise RuntimeError('Too many live instances.') def __del__(self): StateBackedTestElementType.live_element_count -= 1 def __reduce__(self): return (self.__class__, (self.num_elements, 'x' * self.num_elements)) @pytest.mark.it_validatesrunner class FnApiBasedStateBackedCoderTest(unittest.TestCase): def create_pipeline(self): return beam.Pipeline( runner=fn_api_runner.FnApiRunner(use_state_iterables=True)) def test_gbk_many_values(self): with self.create_pipeline() as p: # The number of integers could be a knob to test against # different runners' default settings on page size. VALUES_PER_ELEMENT = 300 NUM_OF_ELEMENTS = 200 r = ( p | beam.Create([None]) | beam.FlatMap( lambda x: ((1, StateBackedTestElementType(VALUES_PER_ELEMENT, _)) for _ in range(NUM_OF_ELEMENTS))) | beam.GroupByKey() | beam.MapTuple(lambda _, vs: sum(e.num_elements for e in vs))) assert_that(r, equal_to([VALUES_PER_ELEMENT * NUM_OF_ELEMENTS])) # TODO(robertwb): Why does pickling break when this is inlined? class CustomMergingWindowFn(window.WindowFn): def assign(self, assign_context): return [ window.IntervalWindow( assign_context.timestamp, assign_context.timestamp + 1000) ] def merge(self, merge_context): evens = [w for w in merge_context.windows if w.start % 2 == 0] if evens: merge_context.merge( evens, window.IntervalWindow( min(w.start for w in evens), max(w.end for w in evens))) def get_window_coder(self): return coders.IntervalWindowCoder() class ColoredFixedWindow(window.BoundedWindow): def __init__(self, end, color): super().__init__(end) self.color = color def __hash__(self): return hash((self.end, self.color)) def __eq__(self, other): return ( type(self) == type(other) and self.end == other.end and self.color == other.color) class ColoredFixedWindowCoder(beam.coders.Coder): kv_coder = beam.coders.TupleCoder( [beam.coders.TimestampCoder(), beam.coders.StrUtf8Coder()]) def encode(self, colored_window): return self.kv_coder.encode((colored_window.end, colored_window.color)) def decode(self, encoded_window): return ColoredFixedWindow(*self.kv_coder.decode(encoded_window)) def is_deterministic(self): return True class EvenOddWindows(window.NonMergingWindowFn): def assign(self, context): timestamp = context.timestamp return [ ColoredFixedWindow( timestamp - timestamp % 10 + 10, 'red' if timestamp.micros // 1000000 % 2 else 'black') ] def get_window_coder(self): return ColoredFixedWindowCoder() class ExpectingSideInputsFn(beam.DoFn): def __init__(self, name): self._name = name def default_label(self): return self._name def process(self, element, *side_inputs): logging.info('Running %s (side inputs: %s)', self._name, side_inputs) if not all(list(s) for s in side_inputs): raise ValueError(f'Missing data in side input {side_inputs}') yield self._name class ArrayMultiplyDoFn(beam.DoFn): def process_batch(self, batch: np.ndarray, *unused_args, **unused_kwargs) -> Iterator[np.ndarray]: assert isinstance(batch, np.ndarray) # GeneralPurposeConsumerSet should limit batches to MAX_BATCH_SIZE (4096) # elements assert np.size(batch, axis=0) <= 4096 yield batch * 2 # infer_output_type must be defined (when there's no process method), # otherwise we don't know the input type is the same as output type. def infer_output_type(self, input_type): return input_type class ListPlusOneDoFn(beam.DoFn): def process_batch(self, batch: list[np.int64], *unused_args, **unused_kwargs) -> Iterator[list[np.int64]]: assert isinstance(batch, list) yield [element + 1 for element in batch] # infer_output_type must be defined (when there's no process method), # otherwise we don't know the input type is the same as output type. def infer_output_type(self, input_type): return input_type if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/fn_api_runner/translations.py0000644000175100017510000024407315206113032027645 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Pipeline transformations for the FnApiRunner. """ # pytype: skip-file # mypy: check-untyped-defs import collections import copy import functools import itertools import logging import operator from builtins import object from typing import TYPE_CHECKING from typing import Any from typing import Callable from typing import Collection from typing import Container from typing import DefaultDict from typing import Iterable from typing import Iterator from typing import MutableMapping from typing import NamedTuple from typing import Optional from typing import TypeVar from typing import Union from apache_beam import coders from apache_beam.internal import pickler from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.worker import bundle_processor from apache_beam.transforms import combiners from apache_beam.transforms import core from apache_beam.utils import proto_utils from apache_beam.utils import timestamp if TYPE_CHECKING: from apache_beam.runners.portability.fn_api_runner.execution import PartitionableBuffer T = TypeVar('T') # This module is experimental. No backwards-compatibility guarantees. _LOGGER = logging.getLogger(__name__) KNOWN_COMPOSITES = frozenset([ common_urns.primitives.GROUP_BY_KEY.urn, common_urns.composites.COMBINE_PER_KEY.urn, common_urns.primitives.PAR_DO.urn, # After SDF expansion. ]) COMBINE_URNS = frozenset([ common_urns.composites.COMBINE_PER_KEY.urn, ]) PAR_DO_URNS = frozenset([ common_urns.primitives.PAR_DO.urn, common_urns.sdf_components.PAIR_WITH_RESTRICTION.urn, common_urns.sdf_components.SPLIT_AND_SIZE_RESTRICTIONS.urn, common_urns.sdf_components.PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS.urn, common_urns.sdf_components.TRUNCATE_SIZED_RESTRICTION.urn, ]) IMPULSE_BUFFER = b'impulse' # TimerFamilyId is identified by transform name + timer family # TODO(pabloem): Rename this type to express this name is unique per pipeline. TimerFamilyId = tuple[str, str] BufferId = bytes # SideInputId is identified by a consumer ParDo + tag. SideInputId = tuple[str, str] SideInputAccessPattern = beam_runner_api_pb2.FunctionSpec # A map from a PCollection coder ID to a Safe Coder ID # A safe coder is a coder that can be used on the runner-side of the FnApi. # A safe coder receives a byte string, and returns a type that can be # understood by the runner when deserializing. SafeCoderMapping = dict[str, str] # DataSideInput maps SideInputIds to a tuple of the encoded bytes of the side # input content, and a payload specification regarding the type of side input # (MultiMap / Iterable). DataSideInput = dict[SideInputId, tuple[bytes, SideInputAccessPattern]] DataOutput = dict[str, BufferId] # A map of [Transform ID, Timer Family ID] to [Buffer ID, Time Domain for timer] # The time domain comes from beam_runner_api_pb2.TimeDomain. It may be # EVENT_TIME or PROCESSING_TIME. OutputTimers = MutableMapping[TimerFamilyId, tuple[BufferId, Any]] # A map of [Transform ID, Timer Family ID] to [Buffer CONTENTS, Timestamp] OutputTimerData = MutableMapping[TimerFamilyId, tuple['PartitionableBuffer', timestamp.Timestamp]] BundleProcessResult = tuple[beam_fn_api_pb2.InstructionResponse, list[beam_fn_api_pb2.ProcessBundleSplitResponse]] # TODO(pabloem): Change tha name to a more representative one class DataInput(NamedTuple): data: MutableMapping[str, 'PartitionableBuffer'] timers: MutableMapping[TimerFamilyId, 'PartitionableBuffer'] class Stage(object): """A set of Transforms that can be sent to the worker for processing.""" def __init__( self, name, # type: str transforms, # type: List[beam_runner_api_pb2.PTransform] downstream_side_inputs=None, # type: Optional[FrozenSet[str]] must_follow=frozenset(), # type: FrozenSet[Stage] parent=None, # type: Optional[str] environment=None, # type: Optional[str] forced_root=False): self.name = name self.transforms = transforms self.downstream_side_inputs = downstream_side_inputs self.must_follow = must_follow self.timers = set() # type: Set[TimerFamilyId] self.parent = parent if environment is None: environment = functools.reduce( self._merge_environments, (self._extract_environment(t) for t in transforms)) self.environment = environment self.forced_root = forced_root def __repr__(self): must_follow = ', '.join(prev.name for prev in self.must_follow) if self.downstream_side_inputs is None: downstream_side_inputs = '' else: downstream_side_inputs = ', '.join( str(si) for si in self.downstream_side_inputs) return "%s\n %s\n must follow: %s\n downstream_side_inputs: %s" % ( self.name, '\n'.join([ "%s:%s" % (transform.unique_name, transform.spec.urn) for transform in self.transforms ]), must_follow, downstream_side_inputs) @staticmethod def _extract_environment(transform): # type: (beam_runner_api_pb2.PTransform) -> Optional[str] environment = transform.environment_id return environment if environment else None @staticmethod def _merge_environments(env1, env2): # type: (Optional[str], Optional[str]) -> Optional[str] if env1 is None: return env2 elif env2 is None: return env1 else: if env1 != env2: raise ValueError( "Incompatible environments: '%s' != '%s'" % (str(env1).replace('\n', ' '), str(env2).replace('\n', ' '))) return env1 def can_fuse(self, consumer, context): # type: (Stage, TransformContext) -> bool try: self._merge_environments(self.environment, consumer.environment) except ValueError: return False def no_overlap(a, b): return not a or not b or not a.intersection(b) return ( not consumer.forced_root and not self in consumer.must_follow and self.is_all_sdk_urns(context) and consumer.is_all_sdk_urns(context) and no_overlap(self.downstream_side_inputs, consumer.side_inputs())) def fuse(self, other, context): # type: (Stage, TransformContext) -> Stage return Stage( "(%s)+(%s)" % (self.name, other.name), self.transforms + other.transforms, union(self.downstream_side_inputs, other.downstream_side_inputs), union(self.must_follow, other.must_follow), environment=self._merge_environments( self.environment, other.environment), parent=_parent_for_fused_stages([self, other], context), forced_root=self.forced_root or other.forced_root) def is_runner_urn(self, context): # type: (TransformContext) -> bool return any( transform.spec.urn in context.known_runner_urns for transform in self.transforms) def is_all_sdk_urns(self, context): def is_sdk_transform(transform): # Execute multi-input flattens in the runner. if transform.spec.urn == common_urns.primitives.FLATTEN.urn and len( transform.inputs) > 1: return False else: return transform.spec.urn not in context.runner_only_urns return all(is_sdk_transform(transform) for transform in self.transforms) def is_stateful(self): for transform in self.transforms: if transform.spec.urn in PAR_DO_URNS: payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.ParDoPayload) if payload.state_specs or payload.timer_family_specs: return True return False def side_inputs(self): # type: () -> Iterator[str] for transform in self.transforms: yield from side_inputs(transform).values() def has_as_main_input(self, pcoll): for transform in self.transforms: if transform.spec.urn in PAR_DO_URNS: payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.ParDoPayload) local_side_inputs = payload.side_inputs else: local_side_inputs = {} for local_id, pipeline_id in transform.inputs.items(): if pcoll == pipeline_id and local_id not in local_side_inputs: return True def deduplicate_read(self): # type: () -> None seen_pcolls = set() # type: Set[str] new_transforms = [] for transform in self.transforms: if transform.spec.urn == bundle_processor.DATA_INPUT_URN: pcoll = only_element(list(transform.outputs.items()))[1] if pcoll in seen_pcolls: continue seen_pcolls.add(pcoll) new_transforms.append(transform) self.transforms = new_transforms def executable_stage_transform( self, known_runner_urns, # type: FrozenSet[str] all_consumers, components # type: beam_runner_api_pb2.Components ): # type: (...) -> beam_runner_api_pb2.PTransform if (len(self.transforms) == 1 and self.transforms[0].spec.urn in known_runner_urns): result = copy.copy(self.transforms[0]) del result.subtransforms[:] return result else: all_inputs = set( pcoll for t in self.transforms for pcoll in t.inputs.values()) all_outputs = set( pcoll for t in self.transforms for pcoll in t.outputs.values()) internal_transforms = set(id(t) for t in self.transforms) external_outputs = [ pcoll for pcoll in all_outputs if all_consumers[pcoll] - internal_transforms ] stage_components = beam_runner_api_pb2.Components() stage_components.CopyFrom(components) # Only keep the PCollections referenced in this stage. stage_components.pcollections.clear() for pcoll_id in all_inputs.union(all_outputs): stage_components.pcollections[pcoll_id].CopyFrom( components.pcollections[pcoll_id]) # Only keep the transforms in this stage. # Also gather up payload data as we iterate over the transforms. stage_components.transforms.clear() main_inputs = set() # type: Set[str] side_inputs = [] user_states = [] timers = [] for ix, transform in enumerate(self.transforms): transform_id = 'transform_%d' % ix if transform.spec.urn == common_urns.primitives.PAR_DO.urn: payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.ParDoPayload) for tag in payload.side_inputs.keys(): side_inputs.append( beam_runner_api_pb2.ExecutableStagePayload.SideInputId( transform_id=transform_id, local_name=tag)) for tag in payload.state_specs.keys(): user_states.append( beam_runner_api_pb2.ExecutableStagePayload.UserStateId( transform_id=transform_id, local_name=tag)) for tag in payload.timer_family_specs.keys(): timers.append( beam_runner_api_pb2.ExecutableStagePayload.TimerId( transform_id=transform_id, local_name=tag)) main_inputs.update( pcoll_id for tag, pcoll_id in transform.inputs.items() if tag not in payload.side_inputs) else: main_inputs.update(transform.inputs.values()) stage_components.transforms[transform_id].CopyFrom(transform) main_input_id = only_element(main_inputs - all_outputs) named_inputs = dict({ '%s:%s' % (side.transform_id, side.local_name): stage_components. transforms[side.transform_id].inputs[side.local_name] for side in side_inputs }, main_input=main_input_id) # at this point we should have resolved an environment, as the key of # components.environments cannot be None. assert self.environment is not None exec_payload = beam_runner_api_pb2.ExecutableStagePayload( environment=components.environments[self.environment], input=main_input_id, outputs=external_outputs, transforms=stage_components.transforms.keys(), components=stage_components, side_inputs=side_inputs, user_states=user_states, timers=timers) return beam_runner_api_pb2.PTransform( unique_name=unique_name(None, self.name), spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.executable_stage, payload=exec_payload.SerializeToString()), inputs=named_inputs, outputs={ 'output_%d' % ix: pcoll for ix, pcoll in enumerate(external_outputs) }, ) def memoize_on_instance(f): missing = object() def wrapper(self, *args): try: cache = getattr(self, '_cache_%s' % f.__name__) except AttributeError: cache = {} setattr(self, '_cache_%s' % f.__name__, cache) result = cache.get(args, missing) if result is missing: result = cache[args] = f(self, *args) return result return wrapper class TransformContext(object): _COMMON_CODER_URNS = set( value.urn for (key, value) in common_urns.coders.__dict__.items() if not key.startswith('_') # Length prefix Rows rather than re-coding them. ) - set([common_urns.coders.ROW.urn]) _REQUIRED_CODER_URNS = set([ common_urns.coders.WINDOWED_VALUE.urn, # For impulse. common_urns.coders.BYTES.urn, common_urns.coders.GLOBAL_WINDOW.urn, # For GBK. common_urns.coders.KV.urn, common_urns.coders.ITERABLE.urn, # For SDF. common_urns.coders.DOUBLE.urn, # For timers. common_urns.coders.TIMER.urn, # For everything else. common_urns.coders.LENGTH_PREFIX.urn, common_urns.coders.CUSTOM_WINDOW.urn, ]) def __init__( self, components, # type: beam_runner_api_pb2.Components known_runner_urns, # type: FrozenSet[str] use_state_iterables=False, is_drain=False): self.components = components self.known_runner_urns = known_runner_urns self.runner_only_urns = known_runner_urns - frozenset( [common_urns.primitives.FLATTEN.urn]) self._known_coder_urns = set.union( # Those which are required. self._REQUIRED_CODER_URNS, # Those common coders which are understood by many environments. self._COMMON_CODER_URNS.intersection( *( set(env.capabilities) for env in self.components.environments.values()))) self.use_state_iterables = use_state_iterables self.is_drain = is_drain # ok to pass None for context because BytesCoder has no components coder_proto = coders.BytesCoder().to_runner_api( None) # type: ignore[arg-type] self.bytes_coder_id = self.add_or_get_coder_id(coder_proto, 'bytes_coder') self.safe_coders: SafeCoderMapping = { self.bytes_coder_id: self.bytes_coder_id } # A map of PCollection ID to Coder ID. self.data_channel_coders = {} # type: Dict[str, str] def add_or_get_coder_id( self, coder_proto, # type: beam_runner_api_pb2.Coder coder_prefix='coder'): # type: (...) -> str for coder_id, coder in self.components.coders.items(): if coder == coder_proto: return coder_id new_coder_id = unique_name(self.components.coders, coder_prefix) self.components.coders[new_coder_id].CopyFrom(coder_proto) return new_coder_id def add_data_channel_coder(self, pcoll_id, is_input=False): pcoll = self.components.pcollections[pcoll_id] proto = beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.coders.WINDOWED_VALUE.urn), component_coder_ids=[ pcoll.coder_id, self.components.windowing_strategies[ pcoll.windowing_strategy_id].window_coder_id ]) windowed_coder_id = self.add_or_get_coder_id( proto, pcoll.coder_id + '_windowed') if is_input and self.use_state_iterables: windowed_coder_id = self.with_state_iterables(windowed_coder_id) self.data_channel_coders[pcoll_id] = self.maybe_length_prefixed_coder( windowed_coder_id) @memoize_on_instance def with_state_iterables(self, coder_id): # type: (str) -> str coder = self.components.coders[coder_id] if coder.spec.urn == common_urns.coders.ITERABLE.urn: new_coder_id = unique_name( self.components.coders, coder_id + '_state_backed') new_coder = self.components.coders[new_coder_id] new_coder.CopyFrom(coder) new_coder.spec.urn = common_urns.coders.STATE_BACKED_ITERABLE.urn new_coder.spec.payload = b'1' new_coder.component_coder_ids[0] = self.with_state_iterables( coder.component_coder_ids[0]) return new_coder_id else: new_component_ids = [ self.with_state_iterables(c) for c in coder.component_coder_ids ] if new_component_ids == coder.component_coder_ids: return coder_id else: new_coder_id = unique_name( self.components.coders, coder_id + '_state_backed') self.components.coders[new_coder_id].CopyFrom( beam_runner_api_pb2.Coder( spec=coder.spec, component_coder_ids=new_component_ids)) return new_coder_id @memoize_on_instance def maybe_length_prefixed_coder(self, coder_id): # type: (str) -> str if coder_id in self.safe_coders: return coder_id (maybe_length_prefixed_id, safe_id) = self.maybe_length_prefixed_and_safe_coder(coder_id) self.safe_coders[maybe_length_prefixed_id] = safe_id return maybe_length_prefixed_id @memoize_on_instance def maybe_length_prefixed_and_safe_coder(self, coder_id): # type: (str) -> Tuple[str, str] coder = self.components.coders[coder_id] if coder.spec.urn == common_urns.coders.LENGTH_PREFIX.urn: # If the coder is already length prefixed, we can use it as is, and # have the runner treat it as opaque bytes. return coder_id, self.bytes_coder_id elif (coder.spec.urn == common_urns.coders.WINDOWED_VALUE.urn and self.components.coders[coder.component_coder_ids[1]].spec.urn not in self._known_coder_urns): # A WindowedValue coder with an unknown window type. # This needs to be encoded in such a way that we still have access to its # timestmap. lp_elem_coder = self.maybe_length_prefixed_coder( coder.component_coder_ids[0]) tp_window_coder = self.timestamped_prefixed_window_coder( coder.component_coder_ids[1]) new_coder_id = unique_name( self.components.coders, coder_id + '_timestamp_prefixed') self.components.coders[new_coder_id].CopyFrom( beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.coders.WINDOWED_VALUE.urn), component_coder_ids=[lp_elem_coder, tp_window_coder])) safe_coder_id = unique_name( self.components.coders, coder_id + '_timestamp_prefixed_opaque') self.components.coders[safe_coder_id].CopyFrom( beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.coders.WINDOWED_VALUE.urn), component_coder_ids=[ self.safe_coders[lp_elem_coder], self.safe_coders[tp_window_coder] ])) return new_coder_id, safe_coder_id elif coder.spec.urn in self._known_coder_urns: # A known coder type, but its components may still need to be length # prefixed. new_component_ids = [ self.maybe_length_prefixed_coder(c) for c in coder.component_coder_ids ] if new_component_ids == coder.component_coder_ids: new_coder_id = coder_id else: new_coder_id = unique_name( self.components.coders, coder_id + '_length_prefixed') self.components.coders[new_coder_id].CopyFrom( beam_runner_api_pb2.Coder( spec=coder.spec, component_coder_ids=new_component_ids)) safe_component_ids = [self.safe_coders[c] for c in new_component_ids] if safe_component_ids == coder.component_coder_ids: safe_coder_id = coder_id else: safe_coder_id = unique_name(self.components.coders, coder_id + '_safe') self.components.coders[safe_coder_id].CopyFrom( beam_runner_api_pb2.Coder( spec=coder.spec, component_coder_ids=safe_component_ids)) return new_coder_id, safe_coder_id else: # A completely unkown coder. Wrap the entire thing in a length prefix. new_coder_id = unique_name( self.components.coders, coder_id + '_length_prefixed') self.components.coders[new_coder_id].CopyFrom( beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.coders.LENGTH_PREFIX.urn), component_coder_ids=[coder_id])) return new_coder_id, self.bytes_coder_id @memoize_on_instance def timestamped_prefixed_window_coder(self, coder_id): length_prefixed = self.maybe_length_prefixed_coder(coder_id) new_coder_id = unique_name( self.components.coders, coder_id + '_timestamp_prefixed') self.components.coders[new_coder_id].CopyFrom( beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.coders.CUSTOM_WINDOW.urn), component_coder_ids=[length_prefixed])) safe_coder_id = unique_name( self.components.coders, coder_id + '_timestamp_prefixed_opaque') self.components.coders[safe_coder_id].CopyFrom( beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec( urn=python_urns.TIMESTAMP_PREFIXED_OPAQUE_WINDOW_CODER))) self.safe_coders[new_coder_id] = safe_coder_id return new_coder_id def length_prefix_pcoll_coders(self, pcoll_id): # type: (str) -> None self.components.pcollections[pcoll_id].coder_id = ( self.maybe_length_prefixed_coder( self.components.pcollections[pcoll_id].coder_id)) @memoize_on_instance def parents_map(self): return { child: parent for (parent, transform) in self.components.transforms.items() for child in transform.subtransforms } def leaf_transform_stages( root_ids, # type: Iterable[str] components, # type: beam_runner_api_pb2.Components parent=None, # type: Optional[str] known_composites=KNOWN_COMPOSITES # type: FrozenSet[str] ): # type: (...) -> Iterator[Stage] for root_id in root_ids: root = components.transforms[root_id] if root.spec.urn in known_composites: yield Stage(root_id, [root], parent=parent) elif not root.subtransforms: # Make sure its outputs are not a subset of its inputs. if set(root.outputs.values()) - set(root.inputs.values()): yield Stage(root_id, [root], parent=parent) else: for stage in leaf_transform_stages(root.subtransforms, components, root_id, known_composites): yield stage def pipeline_from_stages( pipeline_proto, # type: beam_runner_api_pb2.Pipeline stages, # type: Iterable[Stage] known_runner_urns, # type: FrozenSet[str] partial # type: bool ): # type: (...) -> beam_runner_api_pb2.Pipeline # In case it was a generator that mutates components as it # produces outputs (as is the case with most transformations). stages = list(stages) new_proto = beam_runner_api_pb2.Pipeline() new_proto.CopyFrom(pipeline_proto) components = new_proto.components components.transforms.clear() components.pcollections.clear() # order preserving but still has fast contains checking roots = {} # type: Dict[str, Any] parents = { child: parent for parent, proto in pipeline_proto.components.transforms.items() for child in proto.subtransforms } def copy_output_pcollections(transform): for pcoll_id in transform.outputs.values(): components.pcollections[pcoll_id].CopyFrom( pipeline_proto.components.pcollections[pcoll_id]) def add_parent(child, parent): if parent is None: if child not in roots: roots[child] = None else: if (parent not in components.transforms and parent in pipeline_proto.components.transforms): components.transforms[parent].CopyFrom( pipeline_proto.components.transforms[parent]) copy_output_pcollections(components.transforms[parent]) del components.transforms[parent].subtransforms[:] # Ensure that child is the last item in the parent's subtransforms. # If the stages were previously sorted into topological order using # sort_stages, this ensures that the parent transforms are also # added in topological order. if child in components.transforms[parent].subtransforms: components.transforms[parent].subtransforms.remove(child) components.transforms[parent].subtransforms.append(child) add_parent(parent, parents.get(parent)) def copy_subtransforms(transform): for subtransform_id in transform.subtransforms: if subtransform_id not in pipeline_proto.components.transforms: raise RuntimeError( 'Could not find subtransform to copy: ' + subtransform_id) subtransform = pipeline_proto.components.transforms[subtransform_id] components.transforms[subtransform_id].CopyFrom(subtransform) copy_output_pcollections(components.transforms[subtransform_id]) copy_subtransforms(subtransform) all_consumers = collections.defaultdict( set) # type: DefaultDict[str, Set[int]] for stage in stages: for transform in stage.transforms: for pcoll in transform.inputs.values(): all_consumers[pcoll].add(id(transform)) for stage in stages: if partial: transform = only_element(stage.transforms) copy_subtransforms(transform) else: transform = stage.executable_stage_transform( known_runner_urns, all_consumers, pipeline_proto.components) transform_id = unique_name(components.transforms, stage.name) components.transforms[transform_id].CopyFrom(transform) copy_output_pcollections(transform) add_parent(transform_id, stage.parent) del new_proto.root_transform_ids[:] new_proto.root_transform_ids.extend(roots.keys()) return new_proto def create_and_optimize_stages( pipeline_proto, # type: beam_runner_api_pb2.Pipeline phases, known_runner_urns, # type: FrozenSet[str] use_state_iterables=False, is_drain=False): # type: (...) -> Tuple[TransformContext, List[Stage]] """Create a set of stages given a pipeline proto, and set of optimizations. Args: pipeline_proto (beam_runner_api_pb2.Pipeline): A pipeline defined by a user. phases (callable): Each phase identifies a specific transformation to be applied to the pipeline graph. Existing phases are defined in this file, and receive a list of stages, and a pipeline context. Some available transformations are ``lift_combiners``, ``expand_sdf``, ``expand_gbk``, etc. Returns: A tuple with a pipeline context, and a list of stages (i.e. an optimized graph). """ pipeline_context = TransformContext( pipeline_proto.components, known_runner_urns, use_state_iterables=use_state_iterables, is_drain=is_drain) # Initial set of stages are singleton leaf transforms. stages = list( leaf_transform_stages( pipeline_proto.root_transform_ids, pipeline_proto.components, known_composites=union(known_runner_urns, KNOWN_COMPOSITES))) # Apply each phase in order. for phase in phases: stages = list(phase(stages, pipeline_context)) # Return the (possibly mutated) context and ordered set of stages. return pipeline_context, stages def optimize_pipeline( pipeline_proto, # type: beam_runner_api_pb2.Pipeline phases, known_runner_urns, # type: FrozenSet[str] partial=False, **kwargs): unused_context, stages = create_and_optimize_stages( pipeline_proto, phases, known_runner_urns, **kwargs) return pipeline_from_stages( pipeline_proto, stages, known_runner_urns, partial) # Optimization stages. def standard_optimize_phases(): """Returns the basic set of phases, to be passed to optimize_pipeline, that result in a pipeline consisting only of fused stages (with urn beam:runner:executable_stage:v1) and, of course, those designated as known runner urns, in topological order. """ return [ annotate_downstream_side_inputs, annotate_stateful_dofns_as_roots, fix_side_input_pcoll_coders, pack_combiners, lift_combiners, expand_sdf, fix_flatten_coders, # sink_flattens, greedily_fuse, read_to_impulse, extract_impulse_stages, remove_data_plane_ops, sort_stages, ] def annotate_downstream_side_inputs(stages, pipeline_context): # type: (Iterable[Stage], TransformContext) -> Iterable[Stage] """Annotate each stage with fusion-prohibiting information. Each stage is annotated with the (transitive) set of pcollections that depend on this stage that are also used later in the pipeline as a side input. While theoretically this could result in O(n^2) annotations, the size of each set is bounded by the number of side inputs (typically much smaller than the number of total nodes) and the number of *distinct* side-input sets is also generally small (and shared due to the use of union defined above). This representation is also amenable to simple recomputation on fusion. """ consumers = collections.defaultdict( list) # type: DefaultDict[str, List[Stage]] def get_all_side_inputs(): # type: () -> Set[str] all_side_inputs = set() # type: Set[str] for stage in stages: for transform in stage.transforms: for input in transform.inputs.values(): consumers[input].append(stage) for si in stage.side_inputs(): all_side_inputs.add(si) return all_side_inputs all_side_inputs = frozenset(get_all_side_inputs()) downstream_side_inputs_by_stage = {} # type: Dict[Stage, FrozenSet[str]] def compute_downstream_side_inputs(stage): # type: (Stage) -> FrozenSet[str] if stage not in downstream_side_inputs_by_stage: downstream_side_inputs = frozenset() # type: FrozenSet[str] for transform in stage.transforms: for output in transform.outputs.values(): if output in all_side_inputs: downstream_side_inputs = union( downstream_side_inputs, frozenset([output])) for consumer in consumers[output]: downstream_side_inputs = union( downstream_side_inputs, compute_downstream_side_inputs(consumer)) downstream_side_inputs_by_stage[stage] = downstream_side_inputs return downstream_side_inputs_by_stage[stage] for stage in stages: stage.downstream_side_inputs = compute_downstream_side_inputs(stage) return stages def annotate_stateful_dofns_as_roots(stages, pipeline_context): # type: (Iterable[Stage], TransformContext) -> Iterable[Stage] for stage in stages: for transform in stage.transforms: if transform.spec.urn == common_urns.primitives.PAR_DO.urn: pardo_payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.ParDoPayload) if pardo_payload.state_specs or pardo_payload.timer_family_specs: stage.forced_root = True yield stage def fix_side_input_pcoll_coders(stages, pipeline_context): # type: (Iterable[Stage], TransformContext) -> Iterable[Stage] """Length prefix side input PCollection coders. """ for stage in stages: for si in stage.side_inputs(): pipeline_context.length_prefix_pcoll_coders(si) return stages def _group_stages_by_key(stages, get_stage_key): grouped_stages = collections.defaultdict(list) stages_with_none_key = [] for stage in stages: stage_key = get_stage_key(stage) if stage_key is None: stages_with_none_key.append(stage) else: grouped_stages[stage_key].append(stage) return (grouped_stages, stages_with_none_key) def _group_stages_with_limit(stages, get_limit): # type: (Iterable[Stage], Callable[[str], int]) -> Iterable[Collection[Stage]] stages_with_limit = [(stage, get_limit(stage.name)) for stage in stages] group: list[Stage] = [] group_limit = 0 for stage, limit in sorted(stages_with_limit, key=operator.itemgetter(1)): if limit < 1: raise Exception( 'expected get_limit to return an integer >= 1, ' 'instead got: %d for stage: %s' % (limit, stage)) if not group: group_limit = limit assert len(group) < group_limit group.append(stage) if len(group) >= group_limit: yield group group = [] if group: yield group def _remap_input_pcolls(transform, pcoll_id_remap): for input_key in list(transform.inputs.keys()): if transform.inputs[input_key] in pcoll_id_remap: transform.inputs[input_key] = pcoll_id_remap[transform.inputs[input_key]] def _make_pack_name(names): """Return the packed Transform or Stage name. The output name will contain the input names' common prefix, the infix '/Packed', and the input names' suffixes in square brackets. For example, if the input names are 'a/b/c1/d1' and 'a/b/c2/d2, then the output name is 'a/b/Packed[c1_d1, c2_d2]'. """ assert names tokens_in_names = [name.split('/') for name in names] common_prefix_tokens = [] # Find the longest common prefix of tokens. while True: first_token_in_names = set() for tokens in tokens_in_names: if not tokens: break first_token_in_names.add(tokens[0]) if len(first_token_in_names) != 1: break common_prefix_tokens.append(next(iter(first_token_in_names))) for tokens in tokens_in_names: tokens.pop(0) common_prefix_tokens.append('Packed') common_prefix = '/'.join(common_prefix_tokens) suffixes = ['_'.join(tokens) for tokens in tokens_in_names] return '%s[%s]' % (common_prefix, ', '.join(suffixes)) def _eliminate_common_key_with_none(stages, context, can_pack=lambda s: True): # type: (Iterable[Stage], TransformContext, Callable[[str], Union[bool, int]]) -> Iterable[Stage] """Runs common subexpression elimination for sibling KeyWithNone stages. If multiple KeyWithNone stages share a common input, then all but one stages will be eliminated along with their output PCollections. Transforms that originally read input from the output PCollection of the eliminated KeyWithNone stages will be remapped to read input from the output PCollection of the remaining KeyWithNone stage. """ # Partition stages by whether they are eligible for common KeyWithNone # elimination, and group eligible KeyWithNone stages by parent and # environment. def get_stage_key(stage): if len(stage.transforms) == 1 and can_pack(stage.name): transform = only_transform(stage.transforms) if (transform.spec.urn == common_urns.primitives.PAR_DO.urn and len(transform.inputs) == 1 and len(transform.outputs) == 1): pardo_payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.ParDoPayload) if pardo_payload.do_fn.urn == python_urns.KEY_WITH_NONE_DOFN: return (only_element(transform.inputs.values()), stage.environment) return None grouped_eligible_stages, ineligible_stages = _group_stages_by_key( stages, get_stage_key) # Eliminate stages and build the PCollection remapping dictionary. pcoll_id_remap = {} remaining_stages = [] for sibling_stages in grouped_eligible_stages.values(): if len(sibling_stages) > 1: output_pcoll_ids = [ only_element(stage.transforms[0].outputs.values()) for stage in sibling_stages ] parent = _parent_for_fused_stages(sibling_stages, context) for to_delete_pcoll_id in output_pcoll_ids[1:]: pcoll_id_remap[to_delete_pcoll_id] = output_pcoll_ids[0] del context.components.pcollections[to_delete_pcoll_id] sibling_stages[0].parent = parent sibling_stages[0].name = _make_pack_name( stage.name for stage in sibling_stages) only_transform( sibling_stages[0].transforms).unique_name = _make_pack_name( only_transform(stage.transforms).unique_name for stage in sibling_stages) remaining_stages.append(sibling_stages[0]) # Remap all transforms in components. for transform in context.components.transforms.values(): _remap_input_pcolls(transform, pcoll_id_remap) # Yield stages while remapping input PCollections if needed. stages_to_yield = itertools.chain(ineligible_stages, remaining_stages) for stage in stages_to_yield: transform = only_transform(stage.transforms) _remap_input_pcolls(transform, pcoll_id_remap) yield stage _DEFAULT_PACK_COMBINERS_LIMIT = 128 def pack_per_key_combiners(stages, context, can_pack=lambda s: True): # type: (Iterable[Stage], TransformContext, Callable[[str], Union[bool, int]]) -> Iterator[Stage] """Packs sibling CombinePerKey stages into a single CombinePerKey. If CombinePerKey stages have a common input, one input each, and one output each, pack the stages into a single stage that runs all CombinePerKeys and outputs resulting tuples to a new PCollection. A subsequent stage unpacks tuples from this PCollection and sends them to the original output PCollections. """ class _UnpackFn(core.DoFn): """A DoFn that unpacks a packed to multiple tagged outputs. Example: tags = (T1, T2, ...) input = (K, (V1, V2, ...)) output = TaggedOutput(T1, (K, V1)), TaggedOutput(T2, (K, V1)), ... """ def __init__(self, tags): self._tags = tags def process(self, element): key, values = element return [ core.pvalue.TaggedOutput(tag, (key, value)) for tag, value in zip(self._tags, values) ] def _get_fallback_coder_id(): return context.add_or_get_coder_id( # passing None works here because there are no component coders coders.registry.get_coder(object).to_runner_api(None)) # type: ignore[arg-type] def _get_component_coder_id_from_kv_coder(coder, index): assert index < 2 if coder.spec.urn == common_urns.coders.KV.urn and len( coder.component_coder_ids) == 2: return coder.component_coder_ids[index] return _get_fallback_coder_id() def _get_key_coder_id_from_kv_coder(coder): return _get_component_coder_id_from_kv_coder(coder, 0) def _get_value_coder_id_from_kv_coder(coder): return _get_component_coder_id_from_kv_coder(coder, 1) def _try_fuse_stages(a, b): if a.can_fuse(b, context): return a.fuse(b, context) else: raise ValueError def _get_limit(stage_name): result = can_pack(stage_name) if result is True: return _DEFAULT_PACK_COMBINERS_LIMIT else: return int(result) # Partition stages by whether they are eligible for CombinePerKey packing # and group eligible CombinePerKey stages by parent and environment. def get_stage_key(stage): if (len(stage.transforms) == 1 and can_pack(stage.name) and stage.environment is not None and python_urns.PACKED_COMBINE_FN in context.components.environments[stage.environment].capabilities): transform = only_transform(stage.transforms) if (transform.spec.urn == common_urns.composites.COMBINE_PER_KEY.urn and len(transform.inputs) == 1 and len(transform.outputs) == 1): combine_payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.CombinePayload) if combine_payload.combine_fn.urn == python_urns.PICKLED_COMBINE_FN: return (only_element(transform.inputs.values()), stage.environment) return None grouped_eligible_stages, ineligible_stages = _group_stages_by_key( stages, get_stage_key) for stage in ineligible_stages: yield stage grouped_packable_stages = [ (stage_key, subgrouped_stages) for stage_key, grouped_stages in grouped_eligible_stages.items() for subgrouped_stages in _group_stages_with_limit( grouped_stages, _get_limit) ] for stage_key, packable_stages in grouped_packable_stages: input_pcoll_id, _ = stage_key try: if not len(packable_stages) > 1: raise ValueError('Only one stage in this group: Skipping stage packing') # Fused stage is used as template and is not yielded. fused_stage = functools.reduce(_try_fuse_stages, packable_stages) except ValueError: # Skip packing stages in this group. # Yield the stages unmodified, and then continue to the next group. for stage in packable_stages: yield stage continue transforms = [only_transform(stage.transforms) for stage in packable_stages] combine_payloads = [ proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.CombinePayload) for transform in transforms ] output_pcoll_ids = [ only_element(transform.outputs.values()) for transform in transforms ] # Build accumulator coder for (acc1, acc2, ...) accumulator_coder_ids = [ combine_payload.accumulator_coder_id for combine_payload in combine_payloads ] tuple_accumulator_coder_id = context.add_or_get_coder_id( beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec(urn=python_urns.TUPLE_CODER), component_coder_ids=accumulator_coder_ids)) # Build packed output coder for (key, (out1, out2, ...)) input_kv_coder_id = context.components.pcollections[input_pcoll_id].coder_id key_coder_id = _get_key_coder_id_from_kv_coder( context.components.coders[input_kv_coder_id]) output_kv_coder_ids = [ context.components.pcollections[output_pcoll_id].coder_id for output_pcoll_id in output_pcoll_ids ] output_value_coder_ids = [ _get_value_coder_id_from_kv_coder( context.components.coders[output_kv_coder_id]) for output_kv_coder_id in output_kv_coder_ids ] pack_output_value_coder = beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec(urn=python_urns.TUPLE_CODER), component_coder_ids=output_value_coder_ids) pack_output_value_coder_id = context.add_or_get_coder_id( pack_output_value_coder) pack_output_kv_coder = beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec(urn=common_urns.coders.KV.urn), component_coder_ids=[key_coder_id, pack_output_value_coder_id]) pack_output_kv_coder_id = context.add_or_get_coder_id(pack_output_kv_coder) pack_stage_name = _make_pack_name([stage.name for stage in packable_stages]) pack_transform_name = _make_pack_name([ only_transform(stage.transforms).unique_name for stage in packable_stages ]) pack_pcoll_id = unique_name(context.components.pcollections, 'pcollection') input_pcoll = context.components.pcollections[input_pcoll_id] context.components.pcollections[pack_pcoll_id].CopyFrom( beam_runner_api_pb2.PCollection( unique_name=pack_transform_name + '/Pack.out', coder_id=pack_output_kv_coder_id, windowing_strategy_id=input_pcoll.windowing_strategy_id, is_bounded=input_pcoll.is_bounded)) # Set up Pack stage. # TODO(https://github.com/apache/beam/issues/19737): classes that inherit # from RunnerApiFn are expected to accept a PipelineContext for # from_runner_api/to_runner_api. Determine how to accomodate this. pack_combine_fn = combiners.SingleInputTupleCombineFn( *[ core.CombineFn.from_runner_api(combine_payload.combine_fn, context) # type: ignore[arg-type] for combine_payload in combine_payloads ]).to_runner_api(context) # type: ignore[arg-type] pack_transform = beam_runner_api_pb2.PTransform( unique_name=pack_transform_name + '/Pack', spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.composites.COMBINE_PER_KEY.urn, payload=beam_runner_api_pb2.CombinePayload( combine_fn=pack_combine_fn, accumulator_coder_id=tuple_accumulator_coder_id). SerializeToString()), inputs={'in': input_pcoll_id}, # 'None' single output key follows convention for CombinePerKey. outputs={'None': pack_pcoll_id}, environment_id=fused_stage.environment) pack_stage = Stage( pack_stage_name + '/Pack', [pack_transform], downstream_side_inputs=fused_stage.downstream_side_inputs, must_follow=fused_stage.must_follow, parent=fused_stage.parent, environment=fused_stage.environment) yield pack_stage # Set up Unpack stage tags = [str(i) for i in range(len(output_pcoll_ids))] pickled_do_fn_data = pickler.dumps((_UnpackFn(tags), (), {}, [], None)) unpack_transform = beam_runner_api_pb2.PTransform( unique_name=pack_transform_name + '/Unpack', spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.primitives.PAR_DO.urn, payload=beam_runner_api_pb2.ParDoPayload( do_fn=beam_runner_api_pb2.FunctionSpec( urn=python_urns.PICKLED_DOFN_INFO, payload=pickled_do_fn_data)).SerializeToString()), inputs={'in': pack_pcoll_id}, outputs=dict(zip(tags, output_pcoll_ids)), environment_id=fused_stage.environment) unpack_stage = Stage( pack_stage_name + '/Unpack', [unpack_transform], downstream_side_inputs=fused_stage.downstream_side_inputs, must_follow=fused_stage.must_follow, parent=fused_stage.parent, environment=fused_stage.environment) yield unpack_stage def pack_combiners(stages, context, can_pack=None): # type: (Iterable[Stage], TransformContext, Optional[Callable[[str], Union[bool, int]]]) -> Iterator[Stage] if can_pack is None: can_pack_names = {} # type: Dict[str, Union[bool, int]] parents = context.parents_map() def can_pack_fn(name: str) -> Union[bool, int]: if name in can_pack_names: return can_pack_names[name] else: transform = context.components.transforms[name] if python_urns.APPLY_COMBINER_PACKING in transform.annotations: try: result = int( transform.annotations[python_urns.APPLY_COMBINER_PACKING]) except ValueError: result = True elif name in parents: result = can_pack_fn(parents[name]) else: result = False can_pack_names[name] = result return result can_pack = can_pack_fn yield from pack_per_key_combiners( _eliminate_common_key_with_none(stages, context, can_pack), context, can_pack) def lift_combiners(stages, context): # type: (List[Stage], TransformContext) -> Iterator[Stage] """Expands CombinePerKey into pre- and post-grouping stages. ... -> CombinePerKey -> ... becomes ... -> PreCombine -> GBK -> MergeAccumulators -> ExtractOutput -> ... """ def is_compatible_with_combiner_lifting(trigger): '''Returns whether this trigger is compatible with combiner lifting. Certain triggers, such as those that fire after a certain number of elements, need to observe every element, and as such are incompatible with combiner lifting (which may aggregate several elements into one before they reach the triggering code after shuffle). ''' if trigger is None: return True elif trigger.WhichOneof('trigger') in ( 'default', 'always', 'never', 'after_processing_time', 'after_synchronized_processing_time'): return True elif trigger.HasField('element_count'): return trigger.element_count.element_count == 1 elif trigger.HasField('after_end_of_window'): return is_compatible_with_combiner_lifting( trigger.after_end_of_window.early_firings ) and is_compatible_with_combiner_lifting( trigger.after_end_of_window.late_firings) elif trigger.HasField('after_any'): return all( is_compatible_with_combiner_lifting(t) for t in trigger.after_any.subtriggers) elif trigger.HasField('repeat'): return is_compatible_with_combiner_lifting(trigger.repeat.subtrigger) else: return False def can_lift(combine_per_key_transform): windowing = context.components.windowing_strategies[ context.components.pcollections[only_element( list(combine_per_key_transform.inputs.values()) )].windowing_strategy_id] return is_compatible_with_combiner_lifting(windowing.trigger) def make_stage(base_stage, transform): # type: (Stage, beam_runner_api_pb2.PTransform) -> Stage return Stage( transform.unique_name, [transform], downstream_side_inputs=base_stage.downstream_side_inputs, must_follow=base_stage.must_follow, parent=base_stage.name, environment=base_stage.environment) def lifted_stages(stage): transform = stage.transforms[0] combine_payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.CombinePayload) input_pcoll = context.components.pcollections[only_element( list(transform.inputs.values()))] output_pcoll = context.components.pcollections[only_element( list(transform.outputs.values()))] element_coder_id = input_pcoll.coder_id element_coder = context.components.coders[element_coder_id] key_coder_id, _ = element_coder.component_coder_ids accumulator_coder_id = combine_payload.accumulator_coder_id key_accumulator_coder = beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec(urn=common_urns.coders.KV.urn), component_coder_ids=[key_coder_id, accumulator_coder_id]) key_accumulator_coder_id = context.add_or_get_coder_id( key_accumulator_coder) accumulator_iter_coder = beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.coders.ITERABLE.urn), component_coder_ids=[accumulator_coder_id]) accumulator_iter_coder_id = context.add_or_get_coder_id( accumulator_iter_coder) key_accumulator_iter_coder = beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec(urn=common_urns.coders.KV.urn), component_coder_ids=[key_coder_id, accumulator_iter_coder_id]) key_accumulator_iter_coder_id = context.add_or_get_coder_id( key_accumulator_iter_coder) precombined_pcoll_id = unique_name( context.components.pcollections, 'pcollection') context.components.pcollections[precombined_pcoll_id].CopyFrom( beam_runner_api_pb2.PCollection( unique_name=transform.unique_name + '/Precombine.out', coder_id=key_accumulator_coder_id, windowing_strategy_id=input_pcoll.windowing_strategy_id, is_bounded=input_pcoll.is_bounded)) grouped_pcoll_id = unique_name( context.components.pcollections, 'pcollection') context.components.pcollections[grouped_pcoll_id].CopyFrom( beam_runner_api_pb2.PCollection( unique_name=transform.unique_name + '/Group.out', coder_id=key_accumulator_iter_coder_id, windowing_strategy_id=output_pcoll.windowing_strategy_id, is_bounded=output_pcoll.is_bounded)) merged_pcoll_id = unique_name( context.components.pcollections, 'pcollection') context.components.pcollections[merged_pcoll_id].CopyFrom( beam_runner_api_pb2.PCollection( unique_name=transform.unique_name + '/Merge.out', coder_id=key_accumulator_coder_id, windowing_strategy_id=output_pcoll.windowing_strategy_id, is_bounded=output_pcoll.is_bounded)) yield make_stage( stage, beam_runner_api_pb2.PTransform( unique_name=transform.unique_name + '/Precombine', spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.combine_components.COMBINE_PER_KEY_PRECOMBINE. urn, payload=transform.spec.payload), inputs=transform.inputs, outputs={'out': precombined_pcoll_id}, annotations=transform.annotations, environment_id=transform.environment_id)) yield make_stage( stage, beam_runner_api_pb2.PTransform( unique_name=transform.unique_name + '/Group', spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.primitives.GROUP_BY_KEY.urn), inputs={'in': precombined_pcoll_id}, annotations=transform.annotations, outputs={'out': grouped_pcoll_id})) yield make_stage( stage, beam_runner_api_pb2.PTransform( unique_name=transform.unique_name + '/Merge', spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.combine_components. COMBINE_PER_KEY_MERGE_ACCUMULATORS.urn, payload=transform.spec.payload), inputs={'in': grouped_pcoll_id}, outputs={'out': merged_pcoll_id}, annotations=transform.annotations, environment_id=transform.environment_id)) yield make_stage( stage, beam_runner_api_pb2.PTransform( unique_name=transform.unique_name + '/ExtractOutputs', spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.combine_components. COMBINE_PER_KEY_EXTRACT_OUTPUTS.urn, payload=transform.spec.payload), inputs={'in': merged_pcoll_id}, outputs=transform.outputs, annotations=transform.annotations, environment_id=transform.environment_id)) def unlifted_stages(stage): transform = stage.transforms[0] for sub in transform.subtransforms: yield make_stage(stage, context.components.transforms[sub]) for stage in stages: transform = only_transform(stage.transforms) if transform.spec.urn == common_urns.composites.COMBINE_PER_KEY.urn: expansion = lifted_stages if can_lift(transform) else unlifted_stages for substage in expansion(stage): yield substage else: yield stage def _lowest_common_ancestor(a, b, parents): # type: (str, str, Dict[str, str]) -> Optional[str] '''Returns the name of the lowest common ancestor of the two named stages. The map of stage names to their parents' stage names should be provided in parents. Note that stages are considered to be ancestors of themselves. ''' assert a != b def get_ancestors(name): ancestor = name while ancestor is not None: yield ancestor ancestor = parents.get(ancestor) a_ancestors = set(get_ancestors(a)) for b_ancestor in get_ancestors(b): if b_ancestor in a_ancestors: return b_ancestor return None def _parent_for_fused_stages(stages, context): # type: (Iterable[Stage], TransformContext) -> Optional[str] '''Returns the name of the new parent for the fused stages. The new parent is the lowest common ancestor of the fused stages that is not contained in the set of stages to be fused. The provided context is used to compute ancestors of stages. ''' parents = context.parents_map() # If any of the input stages were produced by fusion or an optimizer phase, # or had its parent modified by an optimizer phase, its parent will not be # be reflected in the PipelineContext yet, so we need to add it to the # parents map. for stage in stages: parents[stage.name] = stage.parent def reduce_fn(a, b): # type: (Optional[str], Optional[str]) -> Optional[str] if a is None or b is None: return None return _lowest_common_ancestor(a, b, parents) stage_names = [stage.name for stage in stages] # type: List[Optional[str]] result = functools.reduce(reduce_fn, stage_names) if result in stage_names: result = parents.get(result) return result def expand_sdf(stages, context): # type: (Iterable[Stage], TransformContext) -> Iterator[Stage] """Transforms splitable DoFns into pair+split+read.""" for stage in stages: transform = only_transform(stage.transforms) if transform.spec.urn == common_urns.primitives.PAR_DO.urn: pardo_payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.ParDoPayload) if pardo_payload.restriction_coder_id: def copy_like(protos, original, suffix='_copy', **kwargs): if isinstance(original, str): key = original original = protos[original] else: key = 'component' new_id = unique_name(protos, key + suffix) protos[new_id].CopyFrom(original) proto = protos[new_id] for name, value in kwargs.items(): if isinstance(value, dict): getattr(proto, name).clear() getattr(proto, name).update(value) elif isinstance(value, list): del getattr(proto, name)[:] getattr(proto, name).extend(value) elif name == 'urn': proto.spec.urn = value elif name == 'payload': proto.spec.payload = value else: setattr(proto, name, value) if 'unique_name' not in kwargs and hasattr(proto, 'unique_name'): proto.unique_name = unique_name( {p.unique_name for p in protos.values()}, original.unique_name + suffix) return new_id def make_stage(base_stage, transform_id, extra_must_follow=()): # type: (Stage, str, Iterable[Stage]) -> Stage transform = context.components.transforms[transform_id] return Stage( transform.unique_name, [transform], base_stage.downstream_side_inputs, union(base_stage.must_follow, frozenset(extra_must_follow)), parent=base_stage.name, environment=base_stage.environment) main_input_tag = only_element( tag for tag in transform.inputs.keys() if tag not in pardo_payload.side_inputs) main_input_id = transform.inputs[main_input_tag] element_coder_id = context.components.pcollections[ main_input_id].coder_id # Tuple[element, restriction] paired_coder_id = context.add_or_get_coder_id( beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.coders.KV.urn), component_coder_ids=[ element_coder_id, pardo_payload.restriction_coder_id ])) # Tuple[Tuple[element, restriction], double] sized_coder_id = context.add_or_get_coder_id( beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.coders.KV.urn), component_coder_ids=[ paired_coder_id, context.add_or_get_coder_id( # context can be None here only because FloatCoder does # not have components coders.FloatCoder().to_runner_api(None), # type: ignore 'doubles_coder') ])) paired_pcoll_id = copy_like( context.components.pcollections, main_input_id, '_paired', coder_id=paired_coder_id) pair_transform_id = copy_like( context.components.transforms, transform, unique_name=transform.unique_name + '/PairWithRestriction', urn=common_urns.sdf_components.PAIR_WITH_RESTRICTION.urn, outputs={'out': paired_pcoll_id}) split_pcoll_id = copy_like( context.components.pcollections, main_input_id, '_split', coder_id=sized_coder_id) split_transform_id = copy_like( context.components.transforms, transform, unique_name=transform.unique_name + '/SplitAndSizeRestriction', urn=common_urns.sdf_components.SPLIT_AND_SIZE_RESTRICTIONS.urn, inputs=dict(transform.inputs, **{main_input_tag: paired_pcoll_id}), outputs={'out': split_pcoll_id}) reshuffle_stage = None if common_urns.composites.RESHUFFLE.urn in context.known_runner_urns: reshuffle_pcoll_id = copy_like( context.components.pcollections, main_input_id, '_reshuffle', coder_id=sized_coder_id) reshuffle_transform_id = copy_like( context.components.transforms, transform, unique_name=transform.unique_name + '/Reshuffle', urn=common_urns.composites.RESHUFFLE.urn, payload=b'', inputs=dict(transform.inputs, **{main_input_tag: split_pcoll_id}), outputs={'out': reshuffle_pcoll_id}) reshuffle_stage = make_stage(stage, reshuffle_transform_id) else: reshuffle_pcoll_id = split_pcoll_id reshuffle_transform_id = None if context.is_drain: truncate_pcoll_id = copy_like( context.components.pcollections, main_input_id, '_truncate_restriction', coder_id=sized_coder_id) # Lengthprefix the truncate output. context.length_prefix_pcoll_coders(truncate_pcoll_id) truncate_transform_id = copy_like( context.components.transforms, transform, unique_name=transform.unique_name + '/TruncateAndSizeRestriction', urn=common_urns.sdf_components.TRUNCATE_SIZED_RESTRICTION.urn, inputs=dict( transform.inputs, **{main_input_tag: reshuffle_pcoll_id}), outputs={'out': truncate_pcoll_id}) process_transform_id = copy_like( context.components.transforms, transform, unique_name=transform.unique_name + '/Process', urn=common_urns.sdf_components. PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS.urn, inputs=dict( transform.inputs, **{main_input_tag: truncate_pcoll_id})) else: process_transform_id = copy_like( context.components.transforms, transform, unique_name=transform.unique_name + '/Process', urn=common_urns.sdf_components. PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS.urn, inputs=dict( transform.inputs, **{main_input_tag: reshuffle_pcoll_id})) yield make_stage(stage, pair_transform_id) split_stage = make_stage(stage, split_transform_id) yield split_stage if reshuffle_stage: yield reshuffle_stage if context.is_drain: yield make_stage( stage, truncate_transform_id, extra_must_follow=[split_stage]) yield make_stage(stage, process_transform_id) else: yield make_stage( stage, process_transform_id, extra_must_follow=[split_stage]) else: yield stage else: yield stage def expand_gbk(stages, pipeline_context): # type: (Iterable[Stage], TransformContext) -> Iterator[Stage] """Transforms each GBK into a write followed by a read.""" for stage in stages: transform = only_transform(stage.transforms) if transform.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn: for pcoll_id in transform.inputs.values(): pipeline_context.length_prefix_pcoll_coders(pcoll_id) for pcoll_id in transform.outputs.values(): pipeline_context.length_prefix_pcoll_coders(pcoll_id) # This is used later to correlate the read and write. transform_id = stage.name if transform != pipeline_context.components.transforms.get(transform_id): transform_id = unique_name( pipeline_context.components.transforms, stage.name) pipeline_context.components.transforms[transform_id].CopyFrom(transform) grouping_buffer = create_buffer_id(transform_id, kind='group') gbk_write = Stage( transform.unique_name + '/Write', [ beam_runner_api_pb2.PTransform( unique_name=transform.unique_name + '/Write', inputs=transform.inputs, spec=beam_runner_api_pb2.FunctionSpec( urn=bundle_processor.DATA_OUTPUT_URN, payload=grouping_buffer)) ], downstream_side_inputs=frozenset(), must_follow=stage.must_follow) yield gbk_write yield Stage( transform.unique_name + '/Read', [ beam_runner_api_pb2.PTransform( unique_name=transform.unique_name + '/Read', outputs=transform.outputs, spec=beam_runner_api_pb2.FunctionSpec( urn=bundle_processor.DATA_INPUT_URN, payload=grouping_buffer)) ], downstream_side_inputs=stage.downstream_side_inputs, must_follow=union(frozenset([gbk_write]), stage.must_follow)) else: yield stage def fix_flatten_coders( stages, pipeline_context, identity_urn=bundle_processor.IDENTITY_DOFN_URN): # type: (Iterable[Stage], TransformContext, str) -> Iterator[Stage] """Ensures that the inputs of Flatten have the same coders as the output. """ pcollections = pipeline_context.components.pcollections for stage in stages: transform = only_element(stage.transforms) if transform.spec.urn == common_urns.primitives.FLATTEN.urn: output_pcoll_id = only_element(transform.outputs.values()) output_coder_id = pcollections[output_pcoll_id].coder_id for local_in, pcoll_in in list(transform.inputs.items()): if pcollections[pcoll_in].coder_id != output_coder_id: # Flatten requires that all its inputs be materialized with the # same coder as its output. Add stages to transcode flatten # inputs that use different coders. transcoded_pcollection = unique_name( pcollections, transform.unique_name + '/Transcode/' + local_in + '/out') transcode_name = unique_name( pipeline_context.components.transforms, transform.unique_name + '/Transcode/' + local_in) yield Stage( transcode_name, [ beam_runner_api_pb2.PTransform( unique_name=transcode_name, inputs={local_in: pcoll_in}, outputs={'out': transcoded_pcollection}, spec=beam_runner_api_pb2.FunctionSpec(urn=identity_urn), environment_id=transform.environment_id) ], downstream_side_inputs=frozenset(), must_follow=stage.must_follow) pcollections[transcoded_pcollection].CopyFrom(pcollections[pcoll_in]) pcollections[transcoded_pcollection].unique_name = ( transcoded_pcollection) pcollections[transcoded_pcollection].coder_id = output_coder_id transform.inputs[local_in] = transcoded_pcollection yield stage def sink_flattens(stages, pipeline_context): # type: (Iterable[Stage], TransformContext) -> Iterator[Stage] """Sink flattens and remove them from the graph. A flatten that cannot be sunk/fused away becomes multiple writes (to the same logical sink) followed by a read. """ # TODO(robertwb): Actually attempt to sink rather than always materialize. # TODO(robertwb): Possibly fuse multi-input flattens into one of the stages. for stage in fix_flatten_coders(stages, pipeline_context, common_urns.primitives.FLATTEN.urn): transform = only_element(stage.transforms) if (transform.spec.urn == common_urns.primitives.FLATTEN.urn and len(transform.inputs) > 1): # This is used later to correlate the read and writes. buffer_id = create_buffer_id(transform.unique_name) flatten_writes = [] # type: List[Stage] for local_in, pcoll_in in transform.inputs.items(): flatten_write = Stage( transform.unique_name + '/Write/' + local_in, [ beam_runner_api_pb2.PTransform( unique_name=transform.unique_name + '/Write/' + local_in, inputs={local_in: pcoll_in}, spec=beam_runner_api_pb2.FunctionSpec( urn=bundle_processor.DATA_OUTPUT_URN, payload=buffer_id), environment_id=transform.environment_id) ], downstream_side_inputs=frozenset(), must_follow=stage.must_follow) flatten_writes.append(flatten_write) yield flatten_write yield Stage( transform.unique_name + '/Read', [ beam_runner_api_pb2.PTransform( unique_name=transform.unique_name + '/Read', outputs=transform.outputs, spec=beam_runner_api_pb2.FunctionSpec( urn=bundle_processor.DATA_INPUT_URN, payload=buffer_id), environment_id=transform.environment_id) ], downstream_side_inputs=stage.downstream_side_inputs, must_follow=union(frozenset(flatten_writes), stage.must_follow)) else: yield stage def greedily_fuse(stages, pipeline_context): # type: (Iterable[Stage], TransformContext) -> FrozenSet[Stage] """Places transforms sharing an edge in the same stage, whenever possible. """ producers_by_pcoll = {} # type: Dict[str, Stage] consumers_by_pcoll = collections.defaultdict( list) # type: DefaultDict[str, List[Stage]] # Used to always reference the correct stage as the producer and # consumer maps are not updated when stages are fused away. replacements = {} # type: Dict[Stage, Stage] def replacement(s): old_ss = [] while s in replacements: old_ss.append(s) s = replacements[s] for old_s in old_ss[:-1]: replacements[old_s] = s return s def fuse(producer, consumer): fused = producer.fuse(consumer, pipeline_context) replacements[producer] = fused replacements[consumer] = fused # First record the producers and consumers of each PCollection. for stage in stages: for transform in stage.transforms: for input in transform.inputs.values(): consumers_by_pcoll[input].append(stage) for output in transform.outputs.values(): producers_by_pcoll[output] = stage # Now try to fuse away all pcollections. for pcoll, producer in producers_by_pcoll.items(): write_pcoll = None for consumer in consumers_by_pcoll[pcoll]: producer = replacement(producer) consumer = replacement(consumer) # Update consumer.must_follow set, as it's used in can_fuse. consumer.must_follow = frozenset( replacement(s) for s in consumer.must_follow) if producer.can_fuse(consumer, pipeline_context): fuse(producer, consumer) else: # If we can't fuse, do a read + write. pipeline_context.length_prefix_pcoll_coders(pcoll) buffer_id = create_buffer_id(pcoll) if write_pcoll is None: write_pcoll = Stage( pcoll + '/Write', [ beam_runner_api_pb2.PTransform( unique_name=pcoll + '/Write', inputs={'in': pcoll}, spec=beam_runner_api_pb2.FunctionSpec( urn=bundle_processor.DATA_OUTPUT_URN, payload=buffer_id)) ], downstream_side_inputs=producer.downstream_side_inputs) fuse(producer, write_pcoll) if consumer.has_as_main_input(pcoll): read_pcoll = Stage( pcoll + '/Read', [ beam_runner_api_pb2.PTransform( unique_name=pcoll + '/Read', outputs={'out': pcoll}, spec=beam_runner_api_pb2.FunctionSpec( urn=bundle_processor.DATA_INPUT_URN, payload=buffer_id)) ], downstream_side_inputs=consumer.downstream_side_inputs, must_follow=frozenset([write_pcoll])) fuse(read_pcoll, consumer) else: consumer.must_follow = union( consumer.must_follow, frozenset([write_pcoll])) # Everything that was originally a stage or a replacement, but wasn't # replaced, should be in the final graph. final_stages = frozenset(stages).union(list(replacements.values()))\ .difference(list(replacements)) for stage in final_stages: # Update all references to their final values before throwing # the replacement data away. stage.must_follow = frozenset(replacement(s) for s in stage.must_follow) # Two reads of the same stage may have been fused. This is unneeded. stage.deduplicate_read() return final_stages def read_to_impulse(stages, pipeline_context): # type: (Iterable[Stage], TransformContext) -> Iterator[Stage] """Translates Read operations into Impulse operations.""" for stage in stages: # First map Reads, if any, to Impulse + triggered read op. for transform in list(stage.transforms): if transform.spec.urn == common_urns.deprecated_primitives.READ.urn: read_pc = only_element(transform.outputs.values()) read_pc_proto = pipeline_context.components.pcollections[read_pc] impulse_pc = unique_name( pipeline_context.components.pcollections, 'Impulse') pipeline_context.components.pcollections[impulse_pc].CopyFrom( beam_runner_api_pb2.PCollection( unique_name=impulse_pc, coder_id=pipeline_context.bytes_coder_id, windowing_strategy_id=read_pc_proto.windowing_strategy_id, is_bounded=read_pc_proto.is_bounded)) stage.transforms.remove(transform) # TODO(robertwb): If this goes multi-process before fn-api # read is default, expand into split + reshuffle + read. stage.transforms.append( beam_runner_api_pb2.PTransform( unique_name=transform.unique_name + '/Impulse', spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.primitives.IMPULSE.urn), outputs={'out': impulse_pc})) stage.transforms.append( beam_runner_api_pb2.PTransform( unique_name=transform.unique_name, spec=beam_runner_api_pb2.FunctionSpec( urn=python_urns.IMPULSE_READ_TRANSFORM, payload=transform.spec.payload), inputs={'in': impulse_pc}, outputs={'out': read_pc})) yield stage def impulse_to_input(stages, pipeline_context): # type: (Iterable[Stage], TransformContext) -> Iterator[Stage] """Translates Impulse operations into GRPC reads.""" for stage in stages: for transform in list(stage.transforms): if transform.spec.urn == common_urns.primitives.IMPULSE.urn: stage.transforms.remove(transform) stage.transforms.append( beam_runner_api_pb2.PTransform( unique_name=transform.unique_name, spec=beam_runner_api_pb2.FunctionSpec( urn=bundle_processor.DATA_INPUT_URN, payload=IMPULSE_BUFFER), outputs=transform.outputs)) yield stage def extract_impulse_stages(stages, pipeline_context): # type: (Iterable[Stage], TransformContext) -> Iterator[Stage] """Splits fused Impulse operations into their own stage.""" for stage in stages: for transform in list(stage.transforms): if transform.spec.urn == common_urns.primitives.IMPULSE.urn: stage.transforms.remove(transform) yield Stage( transform.unique_name, transforms=[transform], downstream_side_inputs=stage.downstream_side_inputs, must_follow=stage.must_follow, parent=stage.parent) if stage.transforms: yield stage def remove_data_plane_ops(stages, pipeline_context): # type: (Iterable[Stage], TransformContext) -> Iterator[Stage] for stage in stages: for transform in list(stage.transforms): if transform.spec.urn in (bundle_processor.DATA_INPUT_URN, bundle_processor.DATA_OUTPUT_URN): stage.transforms.remove(transform) if stage.transforms: yield stage def setup_timer_mapping(stages, pipeline_context): # type: (Iterable[Stage], TransformContext) -> Iterator[Stage] """Set up a mapping of {transform_id: [timer_ids]} for each stage. """ for stage in stages: for transform in stage.transforms: if transform.spec.urn in PAR_DO_URNS: payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.ParDoPayload) for timer_family_id in payload.timer_family_specs.keys(): stage.timers.add((transform.unique_name, timer_family_id)) yield stage def sort_stages(stages, pipeline_context): # type: (Iterable[Stage], TransformContext) -> List[Stage] """Order stages suitable for sequential execution. """ all_stages = set(stages) seen = set() # type: Set[Stage] ordered = [] producers = { pcoll: stage for stage in all_stages for t in stage.transforms for pcoll in t.outputs.values() } def process(stage): if stage not in seen: seen.add(stage) if stage not in all_stages: return for prev in stage.must_follow: process(prev) stage_outputs = set( pcoll for transform in stage.transforms for pcoll in transform.outputs.values()) for transform in stage.transforms: for pcoll in transform.inputs.values(): if pcoll not in stage_outputs: process(producers[pcoll]) ordered.append(stage) for stage in stages: process(stage) return ordered def populate_data_channel_coders(stages, pipeline_context): # type: (Iterable[Stage], TransformContext) -> Iterable[Stage] """Populate coders for GRPC input and output ports.""" for stage in stages: for transform in stage.transforms: if transform.spec.urn in (bundle_processor.DATA_INPUT_URN, bundle_processor.DATA_OUTPUT_URN): if transform.spec.urn == bundle_processor.DATA_INPUT_URN: sdk_pcoll_id = only_element(transform.outputs.values()) else: sdk_pcoll_id = only_element(transform.inputs.values()) pipeline_context.add_data_channel_coder( sdk_pcoll_id, transform.spec.urn == bundle_processor.DATA_INPUT_URN) return stages def add_impulse_to_dangling_transforms(stages, pipeline_context): # type: (Iterable[Stage], TransformContext) -> Iterable[Stage] """Populate coders for GRPC input and output ports.""" for stage in stages: for transform in stage.transforms: if len(transform.inputs ) == 0 and transform.spec.urn != bundle_processor.DATA_INPUT_URN: # We look through the various stages in the DAG, and transforms. If we # see a transform that has no inputs whatsoever. impulse_pc = unique_name( pipeline_context.components.pcollections, 'Impulse') output_pcoll = pipeline_context.components.pcollections[next( iter(transform.outputs.values()))] pipeline_context.components.pcollections[impulse_pc].CopyFrom( beam_runner_api_pb2.PCollection( unique_name=impulse_pc, coder_id=pipeline_context.bytes_coder_id, windowing_strategy_id=output_pcoll.windowing_strategy_id, is_bounded=output_pcoll.is_bounded)) transform.inputs['in'] = impulse_pc stage.transforms.append( beam_runner_api_pb2.PTransform( unique_name=transform.unique_name, spec=beam_runner_api_pb2.FunctionSpec( urn=bundle_processor.DATA_INPUT_URN, payload=IMPULSE_BUFFER), outputs={'out': impulse_pc})) yield stage def union(a, b): # Minimize the number of distinct sets. if not a or a == b: return b elif not b: return a else: return frozenset.union(a, b) _global_counter = 0 def side_inputs(transform): result = {} if transform.spec.urn in PAR_DO_URNS: payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.ParDoPayload) for side_input in payload.side_inputs: result[side_input] = transform.inputs[side_input] return result def unique_name(existing, prefix): # type: (Optional[Container[str]], str) -> str if existing is None: global _global_counter _global_counter += 1 return '%s_%d' % (prefix, _global_counter) elif prefix in existing: counter = 0 while True: counter += 1 prefix_counter = prefix + "_%s" % counter if prefix_counter not in existing: return prefix_counter else: return prefix def only_element(iterable): # type: (Iterable[T]) -> T element, = iterable return element def only_transform(transforms): # type: (List[beam_runner_api_pb2.PTransform]) -> beam_runner_api_pb2.PTransform assert len(transforms) == 1 return transforms[0] def create_buffer_id(name, kind='materialize'): # type: (str, str) -> bytes return ('%s:%s' % (kind, name)).encode('utf-8') def split_buffer_id(buffer_id): # type: (bytes) -> Tuple[str, str] """A buffer id is "kind:pcollection_id". Split into (kind, pcoll_id). """ kind, pcoll_id = buffer_id.decode('utf-8').split(':', 1) return kind, pcoll_id ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/fn_api_runner/translations_test.py0000644000175100017510000004531315206113032030700 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import unittest import pytest import apache_beam as beam from apache_beam import runners from apache_beam.options import pipeline_options from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.runners.portability.fn_api_runner import translations from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import combiners from apache_beam.transforms import core from apache_beam.transforms import environments from apache_beam.transforms.core import Create class TranslationsTest(unittest.TestCase): def test_eliminate_common_key_with_void(self): class MultipleKeyWithNone(beam.PTransform): def expand(self, pcoll): _ = pcoll | 'key-with-none-a' >> beam.ParDo(core._KeyWithNone()) _ = pcoll | 'key-with-none-b' >> beam.ParDo(core._KeyWithNone()) _ = pcoll | 'key-with-none-c' >> beam.ParDo(core._KeyWithNone()) pipeline = beam.Pipeline() _ = pipeline | beam.Create( [1, 2, 3]) | 'multiple-key-with-none' >> MultipleKeyWithNone() pipeline_proto = pipeline.to_runner_api() _, stages = translations.create_and_optimize_stages( pipeline_proto, [translations._eliminate_common_key_with_none], known_runner_urns=frozenset()) key_with_none_stages = [ stage for stage in stages if 'key-with-none' in stage.name ] self.assertEqual(len(key_with_none_stages), 1) self.assertIn('multiple-key-with-none', key_with_none_stages[0].parent) def test_pack_combiners(self): class MultipleCombines(beam.PTransform): def annotations(self): return {python_urns.APPLY_COMBINER_PACKING: b''} def expand(self, pcoll): _ = pcoll | 'mean-perkey' >> combiners.Mean.PerKey() _ = pcoll | 'count-perkey' >> combiners.Count.PerKey() _ = pcoll | 'largest-perkey' >> core.CombinePerKey(combiners.Largest(1)) pipeline = beam.Pipeline() vals = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6] _ = pipeline | Create([('a', x) for x in vals ]) | 'multiple-combines' >> MultipleCombines() environment = environments.DockerEnvironment.from_options( pipeline_options.PortableOptions(sdk_location='container')) pipeline_proto = pipeline.to_runner_api(default_environment=environment) _, stages = translations.create_and_optimize_stages( pipeline_proto, [translations.pack_combiners], known_runner_urns=frozenset()) combine_per_key_stages = [] for stage in stages: for transform in stage.transforms: if transform.spec.urn == common_urns.composites.COMBINE_PER_KEY.urn: combine_per_key_stages.append(stage) self.assertEqual(len(combine_per_key_stages), 1) self.assertIn('Packed', combine_per_key_stages[0].name) self.assertIn('Packed', combine_per_key_stages[0].transforms[0].unique_name) self.assertIn('multiple-combines', combine_per_key_stages[0].parent) self.assertNotIn('-perkey', combine_per_key_stages[0].parent) def test_pack_combiners_with_missing_environment_capability(self): class MultipleCombines(beam.PTransform): def annotations(self): return {python_urns.APPLY_COMBINER_PACKING: b''} def expand(self, pcoll): _ = pcoll | 'mean-perkey' >> combiners.Mean.PerKey() _ = pcoll | 'count-perkey' >> combiners.Count.PerKey() _ = pcoll | 'largest-perkey' >> core.CombinePerKey(combiners.Largest(1)) pipeline = beam.Pipeline() vals = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6] _ = pipeline | Create([('a', x) for x in vals]) | MultipleCombines() environment = environments.DockerEnvironment(capabilities=()) pipeline_proto = pipeline.to_runner_api(default_environment=environment) _, stages = translations.create_and_optimize_stages( pipeline_proto, [translations.pack_combiners], known_runner_urns=frozenset()) combine_per_key_stages = [] for stage in stages: for transform in stage.transforms: if transform.spec.urn == common_urns.composites.COMBINE_PER_KEY.urn: combine_per_key_stages.append(stage) # Combiner packing should be skipped because the environment is missing # the beam:combinefn:packed_python:v1 capability. self.assertEqual(len(combine_per_key_stages), 3) for combine_per_key_stage in combine_per_key_stages: self.assertNotIn('Packed', combine_per_key_stage.name) self.assertNotIn( 'Packed', combine_per_key_stage.transforms[0].unique_name) def test_pack_global_combiners(self): class MultipleCombines(beam.PTransform): def annotations(self): return {python_urns.APPLY_COMBINER_PACKING: b''} def expand(self, pcoll): _ = pcoll | 'mean-globally' >> combiners.Mean.Globally() _ = pcoll | 'count-globally' >> combiners.Count.Globally() _ = pcoll | 'largest-globally' >> core.CombineGlobally( combiners.Largest(1)) pipeline = beam.Pipeline() vals = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6] _ = pipeline | Create(vals) | 'multiple-combines' >> MultipleCombines() environment = environments.DockerEnvironment.from_options( pipeline_options.PortableOptions(sdk_location='container')) pipeline_proto = pipeline.to_runner_api(default_environment=environment) _, stages = translations.create_and_optimize_stages( pipeline_proto, [ translations.pack_combiners, ], known_runner_urns=frozenset()) key_with_void_stages = [ stage for stage in stages if 'KeyWithVoid' in stage.name ] self.assertEqual(len(key_with_void_stages), 1) self.assertIn('multiple-combines', key_with_void_stages[0].parent) self.assertNotIn('-globally', key_with_void_stages[0].parent) combine_per_key_stages = [] for stage in stages: for transform in stage.transforms: if transform.spec.urn == common_urns.composites.COMBINE_PER_KEY.urn: combine_per_key_stages.append(stage) self.assertEqual(len(combine_per_key_stages), 1) self.assertIn('Packed', combine_per_key_stages[0].name) self.assertIn('Packed', combine_per_key_stages[0].transforms[0].unique_name) self.assertIn('multiple-combines', combine_per_key_stages[0].parent) self.assertNotIn('-globally', combine_per_key_stages[0].parent) def test_optimize_empty_pipeline(self): pipeline = beam.Pipeline() pipeline_proto = pipeline.to_runner_api() optimized_pipeline_proto = translations.optimize_pipeline( pipeline_proto, [], known_runner_urns=frozenset(), partial=True) # Tests that Pipeline.from_runner_api() does not throw an exception. runner = runners.DirectRunner() beam.Pipeline.from_runner_api( optimized_pipeline_proto, runner, pipeline_options.PipelineOptions()) def test_optimize_single_combine_globally(self): class SingleCombine(beam.PTransform): def annotations(self): return {python_urns.APPLY_COMBINER_PACKING: b''} def expand(self, pcoll): _ = pcoll | combiners.Count.Globally() pipeline = beam.Pipeline() vals = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6] _ = pipeline | Create(vals) | SingleCombine() pipeline_proto = pipeline.to_runner_api() optimized_pipeline_proto = translations.optimize_pipeline( pipeline_proto, [ translations.pack_combiners, ], known_runner_urns=frozenset(), partial=True) # Tests that Pipeline.from_runner_api() does not throw an exception. runner = runners.DirectRunner() beam.Pipeline.from_runner_api( optimized_pipeline_proto, runner, pipeline_options.PipelineOptions()) def test_optimize_multiple_combine_globally(self): class MultipleCombines(beam.PTransform): def annotations(self): return {python_urns.APPLY_COMBINER_PACKING: b''} def expand(self, pcoll): _ = pcoll | 'mean-globally' >> combiners.Mean.Globally() _ = pcoll | 'count-globally' >> combiners.Count.Globally() _ = pcoll | 'largest-globally' >> core.CombineGlobally( combiners.Largest(1)) pipeline = beam.Pipeline() vals = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6] _ = pipeline | Create(vals) | MultipleCombines() pipeline_proto = pipeline.to_runner_api() optimized_pipeline_proto = translations.optimize_pipeline( pipeline_proto, [ translations.pack_combiners, ], known_runner_urns=frozenset(), partial=True) # Tests that Pipeline.from_runner_api() does not throw an exception. runner = runners.DirectRunner() beam.Pipeline.from_runner_api( optimized_pipeline_proto, runner, pipeline_options.PipelineOptions()) def test_pipeline_from_sorted_stages_is_toplogically_ordered(self): pipeline = beam.Pipeline() side = pipeline | 'side' >> Create([3, 4]) class CreateAndMultiplyBySide(beam.PTransform): def expand(self, pcoll): return ( pcoll | 'main' >> Create([1, 2]) | 'compute' >> beam.FlatMap( lambda x, s: [x * y for y in s], beam.pvalue.AsIter(side))) _ = pipeline | 'create-and-multiply-by-side' >> CreateAndMultiplyBySide() pipeline_proto = pipeline.to_runner_api() optimized_pipeline_proto = translations.optimize_pipeline( pipeline_proto, [ (lambda stages, _: reversed(list(stages))), translations.sort_stages, ], known_runner_urns=frozenset(), partial=True) def assert_is_topologically_sorted(transform_id, visited_pcolls): transform = optimized_pipeline_proto.components.transforms[transform_id] self.assertTrue(set(transform.inputs.values()).issubset(visited_pcolls)) visited_pcolls.update(transform.outputs.values()) for subtransform in transform.subtransforms: assert_is_topologically_sorted(subtransform, visited_pcolls) self.assertEqual(len(optimized_pipeline_proto.root_transform_ids), 1) assert_is_topologically_sorted( optimized_pipeline_proto.root_transform_ids[0], set()) @pytest.mark.it_validatesrunner def test_run_packable_combine_per_key(self): class MultipleCombines(beam.PTransform): def annotations(self): return {python_urns.APPLY_COMBINER_PACKING: b''} def expand(self, pcoll): # These CombinePerKey stages will be packed if and only if # translations.pack_combiners is enabled in the TestPipeline runner. assert_that( pcoll | 'min-perkey' >> core.CombinePerKey(min), equal_to([('a', -1)]), label='assert-min-perkey') assert_that( pcoll | 'count-perkey' >> combiners.Count.PerKey(), equal_to([('a', 10)]), label='assert-count-perkey') assert_that( pcoll | 'largest-perkey' >> combiners.Top.LargestPerKey(2), equal_to([('a', [9, 6])]), label='assert-largest-perkey') with TestPipeline() as pipeline: vals = [6, 3, 1, -1, 9, 1, 5, 2, 0, 6] _ = ( pipeline | Create([('a', x) for x in vals]) | 'multiple-combines' >> MultipleCombines()) @pytest.mark.it_validatesrunner def test_run_packable_combine_globally(self): class MultipleCombines(beam.PTransform): def annotations(self): return {python_urns.APPLY_COMBINER_PACKING: b''} def expand(self, pcoll): # These CombineGlobally stages will be packed if and only if # translations.eliminate_common_key_with_void and # translations.pack_combiners are enabled in the TestPipeline runner. assert_that( pcoll | 'min-globally' >> core.CombineGlobally(min), equal_to([-1]), label='assert-min-globally') assert_that( pcoll | 'count-globally' >> combiners.Count.Globally(), equal_to([10]), label='assert-count-globally') assert_that( pcoll | 'largest-globally' >> combiners.Top.Largest(2), equal_to([[9, 6]]), label='assert-largest-globally') with TestPipeline() as pipeline: vals = [6, 3, 1, -1, 9, 1, 5, 2, 0, 6] _ = pipeline | Create(vals) | 'multiple-combines' >> MultipleCombines() @pytest.mark.it_validatesrunner def test_run_packable_combine_limit(self): class MultipleLargeCombines(beam.PTransform): def annotations(self): # Limit to at most 2 combiners per packed combiner. return {python_urns.APPLY_COMBINER_PACKING: b'2'} def expand(self, pcoll): assert_that( pcoll | 'min-1-globally' >> core.CombineGlobally(min), equal_to([-1]), label='assert-min-1-globally') assert_that( pcoll | 'min-2-globally' >> core.CombineGlobally(min), equal_to([-1]), label='assert-min-2-globally') assert_that( pcoll | 'min-3-globally' >> core.CombineGlobally(min), equal_to([-1]), label='assert-min-3-globally') class MultipleSmallCombines(beam.PTransform): def annotations(self): # Limit to at most 4 combiners per packed combiner. return {python_urns.APPLY_COMBINER_PACKING: b'4'} def expand(self, pcoll): assert_that( pcoll | 'min-4-globally' >> core.CombineGlobally(min), equal_to([-1]), label='assert-min-4-globally') assert_that( pcoll | 'min-5-globally' >> core.CombineGlobally(min), equal_to([-1]), label='assert-min-5-globally') with TestPipeline() as pipeline: vals = [6, 3, 1, -1, 9, 1, 5, 2, 0, 6] pcoll = pipeline | Create(vals) _ = pcoll | 'multiple-large-combines' >> MultipleLargeCombines() _ = pcoll | 'multiple-small-combines' >> MultipleSmallCombines() proto = pipeline.to_runner_api( default_environment=environments.EmbeddedPythonEnvironment( capabilities=environments.python_sdk_capabilities())) optimized = translations.optimize_pipeline( proto, phases=[translations.pack_combiners], known_runner_urns=frozenset(), partial=True) optimized_stage_names = [ t.unique_name for t in optimized.components.transforms.values() ] self.assertIn( 'multiple-large-combines/Packed[min-1-globally_CombinePerKey, ' 'min-2-globally_CombinePerKey]/Pack', optimized_stage_names) self.assertIn( 'Packed[multiple-large-combines_min-3-globally_CombinePerKey, ' 'multiple-small-combines_min-4-globally_CombinePerKey]/Pack', optimized_stage_names) self.assertIn( 'multiple-small-combines/min-5-globally/CombinePerKey', optimized_stage_names) self.assertNotIn( 'multiple-large-combines/min-1-globally/CombinePerKey', optimized_stage_names) self.assertNotIn( 'multiple-large-combines/min-2-globally/CombinePerKey', optimized_stage_names) self.assertNotIn( 'multiple-large-combines/min-3-globally/CombinePerKey', optimized_stage_names) self.assertNotIn( 'multiple-small-combines/min-4-globally/CombinePerKey', optimized_stage_names) def test_combineperkey_annotation_propagation(self): """ Test that the CPK component transforms inherit annotations from the source CPK """ class MyCombinePerKey(beam.CombinePerKey): def annotations(self): return {"my_annotation": b""} with TestPipeline() as pipeline: _ = pipeline | beam.Create([(1, 2)]) | MyCombinePerKey(min) # Verify the annotations are propagated to the split up # CPK transforms proto = pipeline.to_runner_api( default_environment=environments.EmbeddedPythonEnvironment( capabilities=environments.python_sdk_capabilities())) optimized = translations.optimize_pipeline( proto, phases=[translations.lift_combiners], known_runner_urns=frozenset(), partial=True) for transform_id in ['MyCombinePerKey(min)/Precombine', 'MyCombinePerKey(min)/Group', 'MyCombinePerKey(min)/Merge', 'MyCombinePerKey(min)/ExtractOutputs']: assert ( "my_annotation" in optimized.components.transforms[transform_id].annotations) def test_conditionally_packed_combiners(self): class RecursiveCombine(beam.PTransform): def __init__(self, labels): self._labels = labels def expand(self, pcoll): base = pcoll | 'Sum' >> beam.CombineGlobally(sum) if self._labels: rest = pcoll | self._labels[0] >> RecursiveCombine(self._labels[1:]) return (base, rest) | beam.Flatten() else: return base def annotations(self): if len(self._labels) == 2: return {python_urns.APPLY_COMBINER_PACKING: b''} else: return {} # Verify the results are as expected. with TestPipeline() as pipeline: result = pipeline | beam.Create([1, 2, 3]) | RecursiveCombine('ABCD') assert_that(result, equal_to([6, 6, 6, 6, 6])) # Verify the optimization is as expected. proto = pipeline.to_runner_api( default_environment=environments.EmbeddedPythonEnvironment( capabilities=environments.python_sdk_capabilities())) optimized = translations.optimize_pipeline( proto, phases=[translations.pack_combiners], known_runner_urns=frozenset(), partial=True) optimized_stage_names = sorted( t.unique_name for t in optimized.components.transforms.values()) self.assertIn('RecursiveCombine/Sum/CombinePerKey', optimized_stage_names) self.assertIn('RecursiveCombine/A/Sum/CombinePerKey', optimized_stage_names) self.assertNotIn( 'RecursiveCombine/A/B/Sum/CombinePerKey', optimized_stage_names) self.assertIn( 'RecursiveCombine/A/B/Packed[Sum_CombinePerKey, ' 'C_Sum_CombinePerKey, C_D_Sum_CombinePerKey]/Pack', optimized_stage_names) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/fn_api_runner/trigger_manager.py0000644000175100017510000004510015206113032030247 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import collections import logging import typing from collections import defaultdict from apache_beam import typehints from apache_beam.coders import PickleCoder from apache_beam.coders import StrUtf8Coder from apache_beam.coders import TupleCoder from apache_beam.coders import VarIntCoder from apache_beam.coders.coders import IntervalWindowCoder from apache_beam.transforms import DoFn from apache_beam.transforms.core import Windowing from apache_beam.transforms.timeutil import TimeDomain from apache_beam.transforms.trigger import AccumulationMode from apache_beam.transforms.trigger import TriggerContext from apache_beam.transforms.trigger import _CombiningValueStateTag from apache_beam.transforms.trigger import _StateTag from apache_beam.transforms.userstate import AccumulatingRuntimeState from apache_beam.transforms.userstate import BagRuntimeState from apache_beam.transforms.userstate import BagStateSpec from apache_beam.transforms.userstate import CombiningValueStateSpec from apache_beam.transforms.userstate import RuntimeTimer from apache_beam.transforms.userstate import SetRuntimeState from apache_beam.transforms.userstate import SetStateSpec from apache_beam.transforms.userstate import TimerSpec from apache_beam.transforms.userstate import on_timer from apache_beam.transforms.window import BoundedWindow from apache_beam.transforms.window import GlobalWindow from apache_beam.transforms.window import TimestampedValue from apache_beam.transforms.window import WindowFn from apache_beam.typehints import TypeCheckError from apache_beam.utils import windowed_value from apache_beam.utils.timestamp import MIN_TIMESTAMP from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.windowed_value import WindowedValue _LOGGER = logging.getLogger(__name__) _LOGGER.setLevel(logging.DEBUG) K = typing.TypeVar('K') class _ReifyWindows(DoFn): """Receives KV pairs, and wraps the values into WindowedValues.""" def process( self, element, window=DoFn.WindowParam, timestamp=DoFn.TimestampParam): try: k, v = element except TypeError: raise TypeCheckError( 'Input to GroupByKey must be a PCollection with ' 'elements compatible with KV[A, B]') yield (k, windowed_value.WindowedValue(v, timestamp, [window])) class _GroupBundlesByKey(DoFn): def start_bundle(self): self.keys = defaultdict(list) def process(self, element): key, windowed_value = element self.keys[key].append(windowed_value) def finish_bundle(self): for k, vals in self.keys.items(): yield windowed_value.WindowedValue((k, vals), MIN_TIMESTAMP, [GlobalWindow()]) def read_watermark(watermark_state): try: return watermark_state.read() except ValueError: watermark_state.add(MIN_TIMESTAMP) return watermark_state.read() class TriggerMergeContext(WindowFn.MergeContext): def __init__( self, all_windows, context: 'FnRunnerStatefulTriggerContext', windowing): super().__init__(all_windows) self.trigger_context = context self.windowing = windowing self.merged_away: dict[BoundedWindow, BoundedWindow] = {} def merge(self, to_be_merged, merge_result): _LOGGER.debug("Merging %s into %s", to_be_merged, merge_result) self.trigger_context.merge_windows(to_be_merged, merge_result) for window in to_be_merged: if window != merge_result: self.merged_away[window] = merge_result # Clear state associated with PaneInfo since it is # not preserved across merges. self.trigger_context.for_window(window).clear_state(None) self.windowing.triggerfn.on_merge( to_be_merged, merge_result, self.trigger_context.for_window(merge_result)) @typehints.with_input_types( tuple[K, typing.Iterable[windowed_value.WindowedValue]]) @typehints.with_output_types( tuple[K, typing.Iterable[windowed_value.WindowedValue]]) class GeneralTriggerManagerDoFn(DoFn): """A trigger manager that supports all windowing / triggering cases. This implements a DoFn that manages triggering in a per-key basis. All elements for a single key are processed together. Per-key state holds data related to all windows. """ # TODO(BEAM-12026) Add support for Global and custom window fns. KNOWN_WINDOWS = SetStateSpec('known_windows', IntervalWindowCoder()) FINISHED_WINDOWS = SetStateSpec('finished_windows', IntervalWindowCoder()) LAST_KNOWN_TIME = CombiningValueStateSpec('last_known_time', combine_fn=max) LAST_KNOWN_WATERMARK = CombiningValueStateSpec( 'last_known_watermark', combine_fn=max) # TODO(pabloem) What's the coder for the elements/keys here? WINDOW_ELEMENT_PAIRS = BagStateSpec( 'all_elements', TupleCoder([IntervalWindowCoder(), PickleCoder()])) WINDOW_TAG_VALUES = BagStateSpec( 'per_window_per_tag_value_state', TupleCoder([IntervalWindowCoder(), StrUtf8Coder(), VarIntCoder()])) PROCESSING_TIME_TIMER = TimerSpec( 'processing_time_timer', TimeDomain.REAL_TIME) WATERMARK_TIMER = TimerSpec('watermark_timer', TimeDomain.WATERMARK) def __init__(self, windowing: Windowing): self.windowing = windowing # Only session windows are merging. Other windows are non-merging. self.merging_windows = self.windowing.windowfn.is_merging() def process( self, element: tuple[K, typing.Iterable[windowed_value.WindowedValue]], all_elements: BagRuntimeState = DoFn.StateParam(WINDOW_ELEMENT_PAIRS), # type: ignore latest_processing_time: AccumulatingRuntimeState = DoFn.StateParam(LAST_KNOWN_TIME), # type: ignore latest_watermark: AccumulatingRuntimeState = DoFn.StateParam( # type: ignore LAST_KNOWN_WATERMARK), window_tag_values: BagRuntimeState = DoFn.StateParam(WINDOW_TAG_VALUES), # type: ignore windows_state: SetRuntimeState = DoFn.StateParam(KNOWN_WINDOWS), # type: ignore finished_windows_state: SetRuntimeState = DoFn.StateParam( # type: ignore FINISHED_WINDOWS), processing_time_timer=DoFn.TimerParam(PROCESSING_TIME_TIMER), watermark_timer=DoFn.TimerParam(WATERMARK_TIMER), *args, **kwargs): context = FnRunnerStatefulTriggerContext( processing_time_timer=processing_time_timer, watermark_timer=watermark_timer, latest_processing_time=latest_processing_time, latest_watermark=latest_watermark, all_elements_state=all_elements, window_tag_values=window_tag_values, finished_windows_state=finished_windows_state) key, windowed_values = element watermark = read_watermark(latest_watermark) windows_to_elements = collections.defaultdict(list) for wv in windowed_values: for window in wv.windows: # ignore expired windows if watermark > window.end + self.windowing.allowed_lateness: continue if window in finished_windows_state.read(): continue windows_to_elements[window].append( TimestampedValue(wv.value, wv.timestamp)) # Processing merging of windows if self.merging_windows: old_windows = set(windows_state.read()) all_windows = old_windows.union(list(windows_to_elements)) if all_windows != old_windows: merge_context = TriggerMergeContext( all_windows, context, self.windowing) self.windowing.windowfn.merge(merge_context) merged_windows_to_elements = collections.defaultdict(list) for window, values in windows_to_elements.items(): while window in merge_context.merged_away: window = merge_context.merged_away[window] merged_windows_to_elements[window].extend(values) windows_to_elements = merged_windows_to_elements for w in windows_to_elements: windows_state.add(w) # Done processing merging of windows seen_windows = set() for w in windows_to_elements: window_context = context.for_window(w) seen_windows.add(w) for value_w_timestamp in windows_to_elements[w]: _LOGGER.debug(value_w_timestamp) all_elements.add((w, value_w_timestamp)) self.windowing.triggerfn.on_element(windowed_values, w, window_context) return self._fire_eligible_windows( key, TimeDomain.WATERMARK, watermark, None, context, seen_windows) def _fire_eligible_windows( self, key: K, time_domain, timestamp: Timestamp, timer_tag: typing.Optional[str], context: 'FnRunnerStatefulTriggerContext', windows_of_interest: typing.Optional[set[BoundedWindow]] = None): windows_to_elements = context.windows_to_elements_map() context.all_elements_state.clear() fired_windows = set() _LOGGER.debug( '%s - tag %s - timestamp %s', time_domain, timer_tag, timestamp) for w, elems in windows_to_elements.items(): if windows_of_interest is not None and w not in windows_of_interest: # windows_of_interest=None means that we care about all windows. # If we care only about some windows, and this window is not one of # them, then we do not intend to fire this window. continue window_context = context.for_window(w) if self.windowing.triggerfn.should_fire(time_domain, timestamp, w, window_context): finished = self.windowing.triggerfn.on_fire( timestamp, w, window_context) _LOGGER.debug('Firing on window %s. Finished: %s', w, finished) fired_windows.add(w) if finished: context.finished_windows_state.add(w) # TODO(pabloem): Format the output: e.g. pane info elems = [WindowedValue(e.value, e.timestamp, (w, )) for e in elems] yield (key, elems) finished_windows: set[BoundedWindow] = set( context.finished_windows_state.read()) # Add elements that were not fired back into state. for w, elems in windows_to_elements.items(): for e in elems: if (w in finished_windows or (w in fired_windows and self.windowing.accumulation_mode == AccumulationMode.DISCARDING)): continue context.all_elements_state.add((w, e)) @on_timer(PROCESSING_TIME_TIMER) def processing_time_trigger( self, key=DoFn.KeyParam, timer_tag=DoFn.DynamicTimerTagParam, timestamp=DoFn.TimestampParam, latest_processing_time=DoFn.StateParam(LAST_KNOWN_TIME), all_elements=DoFn.StateParam(WINDOW_ELEMENT_PAIRS), processing_time_timer=DoFn.TimerParam(PROCESSING_TIME_TIMER), window_tag_values: BagRuntimeState = DoFn.StateParam(WINDOW_TAG_VALUES), # type: ignore finished_windows_state: SetRuntimeState = DoFn.StateParam( # type: ignore FINISHED_WINDOWS), watermark_timer=DoFn.TimerParam(WATERMARK_TIMER)): context = FnRunnerStatefulTriggerContext( processing_time_timer=processing_time_timer, watermark_timer=watermark_timer, latest_processing_time=latest_processing_time, latest_watermark=None, all_elements_state=all_elements, window_tag_values=window_tag_values, finished_windows_state=finished_windows_state) result = self._fire_eligible_windows( key, TimeDomain.REAL_TIME, timestamp, timer_tag, context) latest_processing_time.add(timestamp) return result @on_timer(WATERMARK_TIMER) def watermark_trigger( self, key=DoFn.KeyParam, timer_tag=DoFn.DynamicTimerTagParam, timestamp=DoFn.TimestampParam, latest_watermark=DoFn.StateParam(LAST_KNOWN_WATERMARK), all_elements=DoFn.StateParam(WINDOW_ELEMENT_PAIRS), processing_time_timer=DoFn.TimerParam(PROCESSING_TIME_TIMER), window_tag_values: BagRuntimeState = DoFn.StateParam(WINDOW_TAG_VALUES), # type: ignore finished_windows_state: SetRuntimeState = DoFn.StateParam( # type: ignore FINISHED_WINDOWS), watermark_timer=DoFn.TimerParam(WATERMARK_TIMER)): context = FnRunnerStatefulTriggerContext( processing_time_timer=processing_time_timer, watermark_timer=watermark_timer, latest_processing_time=None, latest_watermark=latest_watermark, all_elements_state=all_elements, window_tag_values=window_tag_values, finished_windows_state=finished_windows_state) result = self._fire_eligible_windows( key, TimeDomain.WATERMARK, timestamp, timer_tag, context) latest_watermark.add(timestamp) return result class FnRunnerStatefulTriggerContext(TriggerContext): def __init__( self, processing_time_timer: RuntimeTimer, watermark_timer: RuntimeTimer, latest_processing_time: typing.Optional[AccumulatingRuntimeState], latest_watermark: typing.Optional[AccumulatingRuntimeState], all_elements_state: BagRuntimeState, window_tag_values: BagRuntimeState, finished_windows_state: SetRuntimeState): self.timers = { TimeDomain.REAL_TIME: processing_time_timer, TimeDomain.WATERMARK: watermark_timer } self.current_times = { TimeDomain.REAL_TIME: latest_processing_time, TimeDomain.WATERMARK: latest_watermark } self.all_elements_state = all_elements_state self.window_tag_values = window_tag_values self.finished_windows_state = finished_windows_state def windows_to_elements_map( self) -> dict[BoundedWindow, list[windowed_value.WindowedValue]]: window_element_pairs: typing.Iterable[ tuple[BoundedWindow, windowed_value.WindowedValue]] = self.all_elements_state.read() result: dict[BoundedWindow, list[windowed_value.WindowedValue]] = {} for w, e in window_element_pairs: if w not in result: result[w] = [] result[w].append(e) return result def for_window(self, window): return PerWindowTriggerContext(window, self) def get_current_time(self): return self.current_times[TimeDomain.REAL_TIME].read() def set_timer(self, name, time_domain, timestamp): _LOGGER.debug('Setting timer (%s, %s) at %s', time_domain, name, timestamp) self.timers[time_domain].set(timestamp, dynamic_timer_tag=name) def clear_timer(self, name, time_domain): _LOGGER.debug('Clearing timer (%s, %s)', time_domain, name) self.timers[time_domain].clear(dynamic_timer_tag=name) def merge_windows(self, to_be_merged, merge_result): all_triplets = list(self.window_tag_values.read()) # Collect all the triplets for the window we are merging away, and tag them # with the new window (merge_result). merging_away_triplets = [(merge_result, state_tag, state) for (window, state_tag, state) in all_triplets if window in to_be_merged] # Collect all of the other triplets, and joining them with the newly-tagged # set of triplets. resulting_triplets = [(window, state_tag, state) for (window, state_tag, state) in all_triplets if window not in to_be_merged] + merging_away_triplets self.window_tag_values.clear() for t in resulting_triplets: self.window_tag_values.add(t) # Merge also element-window pairs all_elements = self.all_elements_state.read() resulting_elements = [ (merge_result if e[0] in to_be_merged else e[0], e[1]) for e in all_elements ] self.all_elements_state.clear() for e in resulting_elements: self.all_elements_state.add(e) def add_state(self, tag, value): # State can only be kept in per-window context, so this is not implemented. raise NotImplementedError('unimplemented') def get_state(self, tag): # State can only be kept in per-window context, so this is not implemented. raise NotImplementedError('unimplemented') def clear_state(self, tag): # State can only be kept in per-window context, so this is not implemented. raise NotImplementedError('unimplemented') class PerWindowTriggerContext(TriggerContext): def __init__(self, window, parent: FnRunnerStatefulTriggerContext): self.window = window self.parent = parent def get_current_time(self): return self.parent.get_current_time() def set_timer(self, name, time_domain, timestamp): self.parent.set_timer(name, time_domain, timestamp) def clear_timer(self, name, time_domain): _LOGGER.debug('Clearing timer (%s, %s)', time_domain, name) self.parent.clear_timer(name, time_domain) def add_state(self, tag: _StateTag, value): assert isinstance(tag, _CombiningValueStateTag) # Used to count: # 1) number of elements in a window ('count') # 2) number of triggers matched individually ('index') # 3) whether the watermark has passed end of window ('is_late') self.parent.window_tag_values.add((self.window, tag.tag, value)) def get_state(self, tag: _StateTag): assert isinstance(tag, _CombiningValueStateTag) # Used to count: # 1) number of elements in a window ('count') # 2) number of triggers matched individually ('index') # 3) whether the watermark has passed end of window ('is_late') all_triplets = self.parent.window_tag_values.read() relevant_triplets = [(window, state_tag, state) for (window, state_tag, state) in all_triplets if window == self.window and state_tag == tag.tag] return tag.combine_fn.apply(relevant_triplets) def clear_state(self, tag: _StateTag): if tag is None: matches = lambda x: True else: matches = lambda x: x == tag.tag all_triplets = self.parent.window_tag_values.read() remaining_triplets = [(window, state_tag, state) for (window, state_tag, state) in all_triplets if not (window == self.window and matches(state_tag))] _LOGGER.debug('Tag: %s | Remaining triplets: %s', tag, remaining_triplets) self.parent.window_tag_values.clear() for t in remaining_triplets: self.parent.window_tag_values.add(t) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/fn_api_runner/trigger_manager_test.py0000644000175100017510000002261615206113032031315 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest from apache_beam import Map from apache_beam import WindowInto from apache_beam.runners.portability.fn_api_runner import trigger_manager from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.core import ParDo from apache_beam.transforms.core import Windowing from apache_beam.transforms.trigger import AccumulationMode from apache_beam.transforms.trigger import AfterCount from apache_beam.transforms.trigger import AfterWatermark from apache_beam.transforms.trigger import Repeatedly from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import IntervalWindow from apache_beam.transforms.window import Sessions from apache_beam.transforms.window import SlidingWindows from apache_beam.transforms.window import TimestampedValue from apache_beam.utils.timestamp import MAX_TIMESTAMP class TriggerManagerTest(unittest.TestCase): def test_with_trigger_window_that_finish(self): def tsv(key, value, ts): return TimestampedValue((key, value), timestamp=ts) # yapf: disable test_stream = ( TestStream() .advance_watermark_to(0) .add_elements([tsv('k1', 1, 0), tsv('k1', 2, 0)]) .add_elements([tsv('k1', 3, 0)]) .advance_watermark_to(2) .add_elements([tsv('k1', 6, 0)]) # This value is discarded. .advance_watermark_to_infinity()) # yapf: enable # Fixed, one-second windows with DefaultTrigger (after watermark) windowing = Windowing( FixedWindows(1), triggerfn=AfterWatermark(), allowed_lateness=0, accumulation_mode=AccumulationMode.DISCARDING) with TestPipeline() as p: result = ( p | test_stream | WindowInto(windowing.windowfn) | ParDo(trigger_manager._ReifyWindows()) | ParDo(trigger_manager._GroupBundlesByKey()) | ParDo(trigger_manager.GeneralTriggerManagerDoFn(windowing)) | Map( lambda elm: (elm[0], elm[1][0].windows[0], [v.value for v in elm[1]]))) assert_that( result, equal_to([ ('k1', IntervalWindow(0, 1), [1, 2, 3]), # On the watermark ])) def test_fixed_windows_simple_watermark(self): def tsv(key, value, ts): return TimestampedValue((key, value), timestamp=ts) # yapf: disable test_stream = ( TestStream() .advance_watermark_to(0) .add_elements([tsv('k1', 1, 0), tsv('k2', 1, 0), tsv('k1', 2, 0), tsv('k2', 2, 0)]) .add_elements([tsv('k1', 3, 0), tsv('k2', 3, 0)]) .add_elements([tsv('k1', 4, 1), tsv('k2', 4, 1)]) .add_elements([tsv('k1', 5, 1), tsv('k2', 5, 1)]) .advance_watermark_to(1) .add_elements([tsv('k1', 6, 0)]) .advance_watermark_to_infinity()) # yapf: enable # Fixed, one-second windows with DefaultTrigger (after watermark) windowing = Windowing( FixedWindows(1), allowed_lateness=MAX_TIMESTAMP.seconds()) with TestPipeline() as p: result = ( p | test_stream | WindowInto(windowing.windowfn) | ParDo(trigger_manager._ReifyWindows()) | ParDo(trigger_manager._GroupBundlesByKey()) | ParDo(trigger_manager.GeneralTriggerManagerDoFn(windowing)) | Map( lambda elm: (elm[0], elm[1][0].windows[0], [v.value for v in elm[1]]))) assert_that( result, equal_to([ ('k1', IntervalWindow(0, 1), [1, 2, 3]), # On the watermark ('k2', IntervalWindow(0, 1), [1, 2, 3]), # On the watermark ('k1', IntervalWindow(1, 2), [4, 5]), # On the watermark ('k2', IntervalWindow(1, 2), [4, 5]), # On the watermark ('k1', IntervalWindow(0, 1), [6]), # After the watermark ])) def test_sliding_windows_simple_watermark(self): # yapf: disable test_stream = ( TestStream() .advance_watermark_to(0) .add_elements([('k1', 1), ('k2', 1), ('k1', 1), ('k2', 1)]) .add_elements([('k1', 1), ('k2', 1)]) .advance_watermark_to(1) .add_elements([('k1', 2), ('k2', 2)]) .add_elements([('k1', 2), ('k2', 2)]) .advance_watermark_to(2) .add_elements([('k1', 3), ('k2', 3)]) .add_elements([('k1', 3), ('k2', 3)]) .advance_watermark_to_infinity()) # yapf: enable # Fixed, one-second windows with DefaultTrigger (after watermark) windowing = Windowing(SlidingWindows(2, 1)) with TestPipeline() as p: result = ( p | test_stream | WindowInto(windowing.windowfn) | ParDo(trigger_manager._ReifyWindows()) | ParDo(trigger_manager._GroupBundlesByKey()) | ParDo(trigger_manager.GeneralTriggerManagerDoFn(windowing)) | Map( lambda elm: (elm[0], elm[1][0].windows[0], [v.value for v in elm[1]]))) assert_that( result, equal_to([ ('k1', IntervalWindow(-1, 1), [1, 1, 1]), ('k2', IntervalWindow(-1, 1), [1, 1, 1]), ('k1', IntervalWindow(0, 2), [1, 1, 1, 2, 2]), ('k2', IntervalWindow(0, 2), [1, 1, 1, 2, 2]), ('k1', IntervalWindow(1, 3), [2, 2, 3, 3]), ('k2', IntervalWindow(1, 3), [2, 2, 3, 3]), ('k1', IntervalWindow(2, 4), [3, 3]), ('k2', IntervalWindow(2, 4), [3, 3]), ])) def test_fixed_after_count_accumulating(self): # yapf: disable test_stream = ( TestStream() .advance_watermark_to(0) .add_elements([('k1', 1), ('k1', 1), ('k2', 1), ('k2', 1)]) .add_elements([('k1', 1), ('k1', 1)]) .advance_watermark_to(2) .add_elements([('k1', 2), ('k2', 2)]) # This values are discarded. .advance_watermark_to_infinity()) # yapf: enable # Fixed, one-second windows with DefaultTrigger (after watermark) windowing = Windowing( FixedWindows(2), triggerfn=Repeatedly(AfterCount(2)), accumulation_mode=AccumulationMode.ACCUMULATING) with TestPipeline() as p: result = ( p | test_stream | WindowInto(windowing.windowfn) | ParDo(trigger_manager._ReifyWindows()) | ParDo(trigger_manager._GroupBundlesByKey()) | ParDo(trigger_manager.GeneralTriggerManagerDoFn(windowing)) | Map( lambda elm: (elm[0], elm[1][0].windows[0], [v.value for v in elm[1]]))) assert_that( result, equal_to([ ('k1', IntervalWindow(0, 2), [1, 1]), ('k2', IntervalWindow(0, 2), [1, 1]), ('k1', IntervalWindow(0, 2), [1, 1, 1, 1]), ])) def test_sessions_and_complex_trigger_accumulating(self): def tsv(key, value, ts): return TimestampedValue((key, value), timestamp=ts) # yapf: disable test_stream = ( TestStream() .advance_watermark_to(0) .add_elements([tsv('k1', 1, 1), tsv('k1', 2, 15), tsv('k1', 3, 7), tsv('k1', 4, 30)]) .advance_watermark_to(50) .add_elements([tsv('k1', -3, 1), tsv('k1', -2, 2),]) .add_elements([tsv('k1', -1, 21)]) .advance_watermark_to_infinity()) # yapf: enable # Fixed, one-second windows with DefaultTrigger (after watermark) windowing = Windowing( Sessions(10), triggerfn=AfterWatermark(early=AfterCount(2), late=AfterCount(1)), accumulation_mode=AccumulationMode.ACCUMULATING, allowed_lateness=MAX_TIMESTAMP.seconds()) with TestPipeline() as p: result = ( p | test_stream | WindowInto(windowing.windowfn) | ParDo(trigger_manager._ReifyWindows()) | ParDo(trigger_manager._GroupBundlesByKey()) | ParDo(trigger_manager.GeneralTriggerManagerDoFn(windowing)) | Map( lambda elm: (elm[0], elm[1][0].windows[0], set(v.value for v in elm[1])))) assert_that( result, equal_to([ ('k1', IntervalWindow(1, 25), {1, 2, 3}), # early ('k1', IntervalWindow(1, 25), {1, 2, 3}), # on time ('k1', IntervalWindow(30, 40), {4}), # on time ('k1', IntervalWindow(1, 25), {1, 2, 3, -3, -2}), # late ('k1', IntervalWindow(1, 40), {1, 2, 3, 4, -3, -2, -1}), # late ])) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/fn_api_runner/visualization_tools.py0000644000175100017510000001050515206113032031234 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Set of utilities to visualize a pipeline to be executed by FnApiRunner.""" from apache_beam.runners.portability.fn_api_runner.translations import Stage from apache_beam.runners.portability.fn_api_runner.watermark_manager import WatermarkManager from apache_beam.utils import timestamp def show_stage(stage: Stage): try: import graphviz except ImportError: import warnings warnings.warn('Unable to draw pipeline. graphviz library missing.') return g = graphviz.Digraph() seen_pcollections = set() for t in stage.transforms: g.node(t.unique_name, shape='box') for i in t.inputs.values(): assert isinstance(i, str) if i not in seen_pcollections: g.node(i) seen_pcollections.add(i) g.edge(i, t.unique_name) for o in t.outputs.values(): assert isinstance(o, str) if o not in seen_pcollections: g.node(o) seen_pcollections.add(o) g.edge(t.unique_name, o) g.render('stage_graph', format='png') def show_watermark_manager(watermark_manager: WatermarkManager, filename=None): try: import graphviz except ImportError: import warnings warnings.warn('Unable to draw pipeline. graphviz library missing.') return g = graphviz.Digraph() def pcoll_node_name(pcoll_node: WatermarkManager.PCollectionNode): if isinstance(pcoll_node.name, tuple): return 'PCOLL_%s_%s' % pcoll_node.name else: return 'PCOLL_%s' % pcoll_node.name def add_node(name, shape=None, color=None, label=None): if name not in seen_nodes: seen_nodes.add(name) g.node( name, shape=shape, fillcolor=color, style='filled', label=name + (label or '')) def add_links(link_from=None, link_to=None, edge_style="solid"): if link_from and link_to: if (link_to, link_from, edge_style) not in seen_links: g.edge(link_from, link_to, style=edge_style) seen_links.add((link_to, link_from, edge_style)) seen_nodes: set[str] = set() seen_links: set[tuple[str, str]] = set() for node in watermark_manager._stages_by_name.values(): name = 'STAGE_%s...%s' % (node.name[:30], node.name[-30:]) add_node(name, 'box') for pcnode in watermark_manager._pcollections_by_name.values(): assert isinstance(pcnode, WatermarkManager.PCollectionNode) name = pcoll_node_name(pcnode) if pcnode.watermark() == timestamp.MIN_TIMESTAMP: color = 'aquamarine' elif pcnode.watermark() == timestamp.MAX_TIMESTAMP: color = 'aquamarine4' else: color = 'aquamarine2' add_node( name, color=color, label='\n%s\nprod: %s' % (pcnode.watermark(), pcnode._produced_watermark)) for node in watermark_manager._stages_by_name.values(): stage = 'STAGE_%s...%s' % (node.name[:30], node.name[-30:]) for pcoll in node.inputs: input_name = pcoll_node_name(pcoll) # Main inputs have a BOLD edge. add_links(link_from=input_name, link_to=stage, edge_style="bold") for pcoll in node.side_inputs: # Side inputs have a dashed edge. input_name = pcoll_node_name(pcoll) add_links(link_from=input_name, link_to=stage, edge_style="dashed") for pcnode in watermark_manager._pcollections_by_name.values(): assert isinstance(pcnode, WatermarkManager.PCollectionNode) pcoll_name = pcoll_node_name(pcnode) for producer in pcnode.producers: prod_name = 'STAGE_%s...%s' % (producer.name[:30], producer.name[-30:]) add_links(link_from=prod_name, link_to=pcoll_name) g.render(filename or 'pipeline_graph', format='png') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/fn_api_runner/watermark_manager.py0000644000175100017510000002017215206113032030603 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utilities for managing watermarks for a pipeline execution by FnApiRunner.""" from __future__ import absolute_import from typing import Union from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.portability.fn_api_runner import translations from apache_beam.runners.portability.fn_api_runner.translations import split_buffer_id from apache_beam.runners.worker import bundle_processor from apache_beam.utils import proto_utils from apache_beam.utils import timestamp class WatermarkManager(object): """Manages the watermarks of a pipeline's stages. It works by constructing an internal graph representation of the pipeline, and keeping track of dependencies.""" class PCollectionNode(object): def __init__(self, name): self.name = name self._watermark = timestamp.MIN_TIMESTAMP self.producers: set[WatermarkManager.StageNode] = set() self.consumers = 0 self._fully_consumed_by = 0 self._produced_watermark = timestamp.MIN_TIMESTAMP def __str__(self): return 'PCollectionNode' % ( self.name, list(self.producers)) def set_watermark(self, wm: timestamp.Timestamp): self._fully_consumed_by += 1 if self._fully_consumed_by >= self.consumers: self._watermark = min(self.upstream_watermark(), wm) def set_produced_watermark(self, wm: timestamp.Timestamp): # TODO(pabloem): Consider case where there are various producers self._produced_watermark = wm def upstream_watermark(self): if len(self.producers) == 1: return self._produced_watermark elif self.producers: return min(p.output_watermark() for p in self.producers) else: return timestamp.MAX_TIMESTAMP def watermark(self): return self._watermark class StageNode(object): def __init__(self, name): # We keep separate inputs and side inputs because side inputs # should hold back a stage's input watermark, to hold back execution # for that stage; but they should not be considered when calculating # the output watermark of the stage, because only the main input # can actually advance that watermark. self.name = name self.inputs: set[WatermarkManager.PCollectionNode] = set() self.side_inputs: set[WatermarkManager.PCollectionNode] = set() self.outputs: set[WatermarkManager.PCollectionNode] = set() def __str__(self): return 'StageNode None: self._pcollections_by_name: dict[Union[str, translations.TimerFamilyId], WatermarkManager.PCollectionNode] = {} self._stages_by_name: dict[str, WatermarkManager.StageNode] = {} def add_pcollection( pcname: str, snode: WatermarkManager.StageNode) -> WatermarkManager.PCollectionNode: if pcname not in self._pcollections_by_name: self._pcollections_by_name[pcname] = WatermarkManager.PCollectionNode( pcname) pcnode = self._pcollections_by_name[pcname] pcnode.consumers += 1 assert isinstance(pcnode, WatermarkManager.PCollectionNode) snode.inputs.add(pcnode) return pcnode for s in stages: stage_name = s.name stage_node = WatermarkManager.StageNode(stage_name) self._stages_by_name[stage_name] = stage_node # 1. Get stage inputs, create nodes for them, add to _stages_by_name, # and add as inputs to stage node. for transform in s.transforms: if transform.spec.urn == bundle_processor.DATA_INPUT_URN: buffer_id = transform.spec.payload if buffer_id == translations.IMPULSE_BUFFER: pcoll_name = transform.unique_name add_pcollection(pcoll_name, stage_node) continue else: _, pcoll_name = split_buffer_id(buffer_id) add_pcollection(pcoll_name, stage_node) # 2. Get stage timers, and add them as inputs to the stage. for transform in s.transforms: if transform.spec.urn in translations.PAR_DO_URNS: payload = proto_utils.parse_Bytes( transform.spec.payload, beam_runner_api_pb2.ParDoPayload) for timer_family_id in payload.timer_family_specs.keys(): timer_pcoll_name = (transform.unique_name, timer_family_id) self._pcollections_by_name[ timer_pcoll_name] = WatermarkManager.PCollectionNode( timer_pcoll_name) timer_pcoll_node = self._pcollections_by_name[timer_pcoll_name] assert isinstance( timer_pcoll_node, WatermarkManager.PCollectionNode) stage_node.inputs.add(timer_pcoll_node) timer_pcoll_node.consumers += 1 # 3. Get stage outputs, create nodes for them, add to # _pcollections_by_name, and add stage as their producer for transform in s.transforms: if transform.spec.urn == bundle_processor.DATA_OUTPUT_URN: buffer_id = transform.spec.payload _, pcoll_name = split_buffer_id(buffer_id) if pcoll_name not in self._pcollections_by_name: self._pcollections_by_name[ pcoll_name] = WatermarkManager.PCollectionNode(pcoll_name) pcoll_node = self._pcollections_by_name[pcoll_name] assert isinstance(pcoll_node, WatermarkManager.PCollectionNode) pcoll_node.producers.add(stage_node) stage_node.outputs.add(pcoll_node) # 4. Get stage side inputs, create nodes for them, add to # _pcollections_by_name, and add them as side inputs of the stage. for pcoll_name in s.side_inputs(): if pcoll_name not in self._pcollections_by_name: self._pcollections_by_name[ pcoll_name] = WatermarkManager.PCollectionNode(pcoll_name) pcoll_node = self._pcollections_by_name[pcoll_name] assert isinstance(pcoll_node, WatermarkManager.PCollectionNode) stage_node.side_inputs.add(pcoll_node) self._verify(stages) def _verify(self, stages: list[translations.Stage]): for s in stages: if len(self._stages_by_name[s.name].inputs) == 0: from apache_beam.runners.portability.fn_api_runner import visualization_tools visualization_tools.show_stage(s) raise ValueError( 'Stage %s has no main inputs. ' 'At least one main input is necessary.' % s.name) def get_stage_node(self, name: str) -> StageNode: # noqa: F821 return self._stages_by_name[name] def get_pcoll_node(self, name: str) -> PCollectionNode: # noqa: F821 return self._pcollections_by_name[name] def set_pcoll_watermark(self, name, watermark): element = self._pcollections_by_name[name] element.set_watermark(watermark) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/fn_api_runner/worker_handlers.py0000644000175100017510000013233615206113032030313 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Code for communicating with the Workers.""" # mypy: disallow-untyped-defs import collections import contextlib import copy import logging import os import queue import subprocess import sys import threading import time from typing import TYPE_CHECKING from typing import Any from typing import BinaryIO # pylint: disable=unused-import from typing import Callable from typing import DefaultDict from typing import Iterable from typing import Iterator from typing import Mapping from typing import Optional from typing import TypeVar from typing import Union from typing import cast from typing import overload import grpc from sortedcontainers import SortedSet from apache_beam import coders from apache_beam.io import filesystems from apache_beam.io.filesystems import CompressionTypes from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.portability.api import beam_artifact_api_pb2_grpc from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.portability.api import beam_provision_api_pb2 from apache_beam.portability.api import beam_provision_api_pb2_grpc from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.portability.api import endpoints_pb2 from apache_beam.runners.portability import artifact_service from apache_beam.runners.portability.fn_api_runner.execution import Buffer from apache_beam.runners.worker import data_plane from apache_beam.runners.worker import sdk_worker from apache_beam.runners.worker.channel_factory import GRPCChannelFactory from apache_beam.runners.worker.log_handler import LOGENTRY_TO_LOG_LEVEL_MAP from apache_beam.runners.worker.sdk_worker import _Future from apache_beam.runners.worker.statecache import StateCache from apache_beam.utils import proto_utils from apache_beam.utils import thread_pool_executor from apache_beam.utils.interactive_utils import is_in_notebook from apache_beam.utils.sentinel import Sentinel if TYPE_CHECKING: from google.protobuf import message from grpc import ServicerContext from apache_beam.runners.portability.fn_api_runner.fn_runner import \ ExtendedProvisionInfo # pylint: disable=ungrouped-imports # State caching is enabled in the fn_api_runner for testing, except for one # test which runs without state caching (FnApiRunnerTestWithDisabledCaching). # The cache is disabled in production for other runners. STATE_CACHE_SIZE_MB = 100 MB_TO_BYTES = 1 << 20 # Time-based flush is enabled in the fn_api_runner by default. DATA_BUFFER_TIME_LIMIT_MS = 1000 FNAPI_RUNNER_CAPABILITIES = frozenset([ common_urns.runner_protocols.MULTIMAP_KEYS_VALUES_SIDE_INPUT.urn, ]) _LOGGER = logging.getLogger(__name__) T = TypeVar('T') ConstructorFn = Callable[[ Union['message.Message', bytes], 'sdk_worker.StateHandler', 'ExtendedProvisionInfo', 'GrpcServer' ], 'WorkerHandler'] class ControlConnection(object): _uid_counter = 0 _lock = threading.Lock() def __init__(self): # type: () -> None self._push_queue = queue.Queue( ) # type: queue.Queue[Union[beam_fn_api_pb2.InstructionRequest, Sentinel]] self._input = None # type: Optional[Iterable[beam_fn_api_pb2.InstructionResponse]] self._futures_by_id = {} # type: Dict[str, ControlFuture] self._read_thread = threading.Thread( name='beam_control_read', target=self._read) self._state = BeamFnControlServicer.UNSTARTED_STATE def _read(self): # type: () -> None assert self._input is not None for data in self._input: self._futures_by_id.pop(data.instruction_id).set(data) @overload def push(self, req): # type: (Sentinel) -> None pass @overload def push(self, req): # type: (beam_fn_api_pb2.InstructionRequest) -> ControlFuture pass def push(self, req): # type: (Union[Sentinel, beam_fn_api_pb2.InstructionRequest]) -> Optional[ControlFuture] if req is BeamFnControlServicer._DONE_MARKER: self._push_queue.put(req) return None if not req.instruction_id: with ControlConnection._lock: ControlConnection._uid_counter += 1 req.instruction_id = 'control_%s' % ControlConnection._uid_counter future = ControlFuture(req.instruction_id) self._futures_by_id[req.instruction_id] = future self._push_queue.put(req) return future def get_req(self): # type: () -> Union[Sentinel, beam_fn_api_pb2.InstructionRequest] return self._push_queue.get() def set_input(self, input): # type: (Iterable[beam_fn_api_pb2.InstructionResponse]) -> None with ControlConnection._lock: if self._input: raise RuntimeError('input is already set.') self._input = input self._read_thread.start() self._state = BeamFnControlServicer.STARTED_STATE def close(self): # type: () -> None with ControlConnection._lock: if self._state == BeamFnControlServicer.STARTED_STATE: self.push(BeamFnControlServicer._DONE_MARKER) self._read_thread.join() self._state = BeamFnControlServicer.DONE_STATE def abort(self, exn): # type: (Exception) -> None for future in self._futures_by_id.values(): future.abort(exn) class BeamFnControlServicer(beam_fn_api_pb2_grpc.BeamFnControlServicer): """Implementation of BeamFnControlServicer for clients.""" UNSTARTED_STATE = 'unstarted' STARTED_STATE = 'started' DONE_STATE = 'done' _DONE_MARKER = Sentinel.sentinel def __init__( self, worker_manager, # type: WorkerHandlerManager ): # type: (...) -> None self._worker_manager = worker_manager self._lock = threading.Lock() self._uid_counter = 0 self._state = self.UNSTARTED_STATE # following self._req_* variables are used for debugging purpose, data is # added only when self._log_req is True. self._req_sent = collections.defaultdict(int) # type: DefaultDict[str, int] self._log_req = logging.getLogger().getEffectiveLevel() <= logging.DEBUG self._connections_by_worker_id = collections.defaultdict( ControlConnection) # type: DefaultDict[str, ControlConnection] def get_conn_by_worker_id(self, worker_id): # type: (str) -> ControlConnection with self._lock: return self._connections_by_worker_id[worker_id] def Control( self, iterator, # type: Iterable[beam_fn_api_pb2.InstructionResponse] context # type: ServicerContext ): # type: (...) -> Iterator[beam_fn_api_pb2.InstructionRequest] with self._lock: if self._state == self.DONE_STATE: return else: self._state = self.STARTED_STATE worker_id = dict(context.invocation_metadata()).get('worker_id') if not worker_id: raise RuntimeError( 'All workers communicate through gRPC should have ' 'worker_id. Received None.') control_conn = self.get_conn_by_worker_id(worker_id) control_conn.set_input(iterator) while True: to_push = control_conn.get_req() if to_push is self._DONE_MARKER: return yield to_push if self._log_req: self._req_sent[to_push.instruction_id] += 1 def done(self): # type: () -> None self._state = self.DONE_STATE _LOGGER.debug( 'Runner: Requests sent by runner: %s', [(str(req), cnt) for req, cnt in self._req_sent.items()]) def GetProcessBundleDescriptor(self, id, context=None): # type: (beam_fn_api_pb2.GetProcessBundleDescriptorRequest, Any) -> beam_fn_api_pb2.ProcessBundleDescriptor return self._worker_manager.get_process_bundle_descriptor(id) class WorkerHandler(object): """worker_handler for a worker. It provides utilities to start / stop the worker, provision any resources for it, as well as provide descriptors for the data, state and logging APIs for it. """ _registered_environments = {} # type: Dict[str, Tuple[ConstructorFn, type]] _worker_id_counter = -1 _lock = threading.Lock() control_conn = None # type: ControlConnection data_conn = None # type: data_plane._GrpcDataChannel def __init__( self, control_handler, # type: Any data_plane_handler, # type: Any state, # type: sdk_worker.StateHandler provision_info # type: ExtendedProvisionInfo ): # type: (...) -> None """Initialize a WorkerHandler. Args: control_handler: data_plane_handler (data_plane.DataChannel): state: provision_info: """ self.control_handler = control_handler self.data_plane_handler = data_plane_handler self.state = state self.provision_info = provision_info with WorkerHandler._lock: WorkerHandler._worker_id_counter += 1 self.worker_id = 'worker_%s' % WorkerHandler._worker_id_counter def close(self): # type: () -> None self.stop_worker() def start_worker(self): # type: () -> None raise NotImplementedError def stop_worker(self): # type: () -> None raise NotImplementedError def control_api_service_descriptor(self): # type: () -> endpoints_pb2.ApiServiceDescriptor raise NotImplementedError def artifact_api_service_descriptor(self): # type: () -> endpoints_pb2.ApiServiceDescriptor raise NotImplementedError def data_api_service_descriptor(self): # type: () -> Optional[endpoints_pb2.ApiServiceDescriptor] raise NotImplementedError def state_api_service_descriptor(self): # type: () -> Optional[endpoints_pb2.ApiServiceDescriptor] raise NotImplementedError def logging_api_service_descriptor(self): # type: () -> Optional[endpoints_pb2.ApiServiceDescriptor] raise NotImplementedError @classmethod def register_environment( cls, urn, # type: str payload_type # type: Optional[Type[T]] ): # type: (...) -> Callable[[Type[WorkerHandler]], Callable[[T, sdk_worker.StateHandler, ExtendedProvisionInfo, GrpcServer], WorkerHandler]] def wrapper(constructor): # type: (Callable) -> Callable cls._registered_environments[urn] = constructor, payload_type # type: ignore[assignment] return constructor return wrapper @classmethod def create( cls, environment, # type: beam_runner_api_pb2.Environment state, # type: sdk_worker.StateHandler provision_info, # type: ExtendedProvisionInfo grpc_server # type: GrpcServer ): # type: (...) -> WorkerHandler constructor, payload_type = cls._registered_environments[environment.urn] return constructor( proto_utils.parse_Bytes(environment.payload, payload_type), state, provision_info, grpc_server) # This takes a WorkerHandlerManager instead of GrpcServer, so it is not # compatible with WorkerHandler.register_environment. There is a special case # in WorkerHandlerManager.get_worker_handlers() that allows it to work. @WorkerHandler.register_environment(python_urns.EMBEDDED_PYTHON, None) class EmbeddedWorkerHandler(WorkerHandler): """An in-memory worker_handler for fn API control, state and data planes.""" def __init__( self, unused_payload, # type: None state, # type: sdk_worker.StateHandler provision_info, # type: ExtendedProvisionInfo worker_manager, # type: WorkerHandlerManager ): # type: (...) -> None super().__init__( self, data_plane.InMemoryDataChannel(), state, provision_info) self.control_conn = self # type: ignore # need Protocol to describe this self.data_conn = self.data_plane_handler state_cache = StateCache(STATE_CACHE_SIZE_MB * MB_TO_BYTES) self.bundle_processor_cache = sdk_worker.BundleProcessorCache( FNAPI_RUNNER_CAPABILITIES, SingletonStateHandlerFactory( sdk_worker.GlobalCachingStateHandler(state_cache, state)), data_plane.InMemoryDataChannelFactory( self.data_plane_handler.inverse()), worker_manager._process_bundle_descriptors) self.worker = sdk_worker.SdkWorker(self.bundle_processor_cache) self._uid_counter = 0 def push(self, request): # type: (beam_fn_api_pb2.InstructionRequest) -> ControlFuture if not request.instruction_id: self._uid_counter += 1 request.instruction_id = 'control_%s' % self._uid_counter response = self.worker.do_instruction(request) return ControlFuture(request.instruction_id, response) def start_worker(self): # type: () -> None pass def stop_worker(self): # type: () -> None self.bundle_processor_cache.shutdown() def done(self): # type: () -> None pass def data_api_service_descriptor(self): # type: () -> endpoints_pb2.ApiServiceDescriptor # A fake endpoint is needed for properly constructing timer info map in # bundle_processor for fnapi_runner. return endpoints_pb2.ApiServiceDescriptor(url='fake') def state_api_service_descriptor(self): # type: () -> None return None def logging_api_service_descriptor(self): # type: () -> None return None class BasicLoggingService(beam_fn_api_pb2_grpc.BeamFnLoggingServicer): def Logging(self, log_messages, context=None): # type: (Iterable[beam_fn_api_pb2.LogEntry.List], Any) -> Iterator[beam_fn_api_pb2.LogControl] yield beam_fn_api_pb2.LogControl() for log_message in log_messages: for log in log_message.log_entries: logging.log(LOGENTRY_TO_LOG_LEVEL_MAP[log.severity], str(log)) class BasicProvisionService(beam_provision_api_pb2_grpc.ProvisionServiceServicer ): def __init__(self, base_info, worker_manager): # type: (beam_provision_api_pb2.ProvisionInfo, WorkerHandlerManager) -> None self._base_info = base_info self._worker_manager = worker_manager def GetProvisionInfo(self, request, context=None): # type: (Any, Optional[ServicerContext]) -> beam_provision_api_pb2.GetProvisionInfoResponse if context: worker_id = dict(context.invocation_metadata())['worker_id'] worker = self._worker_manager.get_worker(worker_id) info = copy.copy(worker.provision_info.provision_info) info.logging_endpoint.CopyFrom(worker.logging_api_service_descriptor()) info.artifact_endpoint.CopyFrom(worker.artifact_api_service_descriptor()) info.control_endpoint.CopyFrom(worker.control_api_service_descriptor()) else: info = self._base_info info.runner_capabilities[:] = FNAPI_RUNNER_CAPABILITIES return beam_provision_api_pb2.GetProvisionInfoResponse(info=info) class GrpcServer(object): _DEFAULT_SHUTDOWN_TIMEOUT_SECS = 5 def __init__( self, state, # type: StateServicer provision_info, # type: Optional[ExtendedProvisionInfo] worker_manager, # type: WorkerHandlerManager ): # type: (...) -> None # Options to have no limits (-1) on the size of the messages # received or sent over the data plane. The actual buffer size # is controlled in a layer above. Also, options to keep the server alive # when too many pings are received. options = [ ("grpc.max_receive_message_length", -1), ("grpc.max_send_message_length", -1), ("grpc.http2.max_pings_without_data", 0), ("grpc.http2.max_ping_strikes", 0), # match `grpc.keepalive_time_ms` defined in the client # (channel_factory.py) ("grpc.http2.min_ping_interval_without_data_ms", 20_000), ] self.state = state self.provision_info = provision_info self.control_server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=options) self.control_port = self.control_server.add_insecure_port('[::]:0') self.control_address = 'localhost:%s' % self.control_port self.data_server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=options) self.data_port = self.data_server.add_insecure_port('[::]:0') self.state_server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=options) self.state_port = self.state_server.add_insecure_port('[::]:0') self.control_handler = BeamFnControlServicer(worker_manager) beam_fn_api_pb2_grpc.add_BeamFnControlServicer_to_server( self.control_handler, self.control_server) # If we have provision info, serve these off the control port as well. if self.provision_info: if self.provision_info.provision_info: beam_provision_api_pb2_grpc.add_ProvisionServiceServicer_to_server( BasicProvisionService( self.provision_info.provision_info, worker_manager), self.control_server) def open_uncompressed(f): # type: (str) -> BinaryIO return filesystems.FileSystems.open( f, compression_type=CompressionTypes.UNCOMPRESSED) beam_artifact_api_pb2_grpc.add_ArtifactRetrievalServiceServicer_to_server( artifact_service.ArtifactRetrievalService( file_reader=open_uncompressed), self.control_server) self.data_plane_handler = data_plane.BeamFnDataServicer( DATA_BUFFER_TIME_LIMIT_MS) beam_fn_api_pb2_grpc.add_BeamFnDataServicer_to_server( self.data_plane_handler, self.data_server) beam_fn_api_pb2_grpc.add_BeamFnStateServicer_to_server( GrpcStateServicer(state), self.state_server) self.logging_server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=options) self.logging_port = self.logging_server.add_insecure_port('[::]:0') beam_fn_api_pb2_grpc.add_BeamFnLoggingServicer_to_server( BasicLoggingService(), self.logging_server) _LOGGER.info('starting control server on port %s', self.control_port) _LOGGER.info('starting data server on port %s', self.data_port) _LOGGER.info('starting state server on port %s', self.state_port) _LOGGER.info('starting logging server on port %s', self.logging_port) self.logging_server.start() self.state_server.start() self.data_server.start() self.control_server.start() def close(self): # type: () -> None self.control_handler.done() to_wait = [ self.control_server.stop(self._DEFAULT_SHUTDOWN_TIMEOUT_SECS), self.data_server.stop(self._DEFAULT_SHUTDOWN_TIMEOUT_SECS), self.state_server.stop(self._DEFAULT_SHUTDOWN_TIMEOUT_SECS), self.logging_server.stop(self._DEFAULT_SHUTDOWN_TIMEOUT_SECS) ] for w in to_wait: w.wait() class GrpcWorkerHandler(WorkerHandler): """An grpc based worker_handler for fn API control, state and data planes.""" def __init__( self, state, # type: StateServicer provision_info, # type: ExtendedProvisionInfo grpc_server # type: GrpcServer ): # type: (...) -> None self._grpc_server = grpc_server super().__init__( self._grpc_server.control_handler, self._grpc_server.data_plane_handler, state, provision_info) self.state = state self.control_address = self.port_from_worker(self._grpc_server.control_port) self.control_conn = self._grpc_server.control_handler.get_conn_by_worker_id( self.worker_id) self.data_conn = self._grpc_server.data_plane_handler.get_conn_by_worker_id( self.worker_id) def control_api_service_descriptor(self): # type: () -> endpoints_pb2.ApiServiceDescriptor return endpoints_pb2.ApiServiceDescriptor( url=self.port_from_worker(self._grpc_server.control_port)) def artifact_api_service_descriptor(self): # type: () -> endpoints_pb2.ApiServiceDescriptor return endpoints_pb2.ApiServiceDescriptor( url=self.port_from_worker(self._grpc_server.control_port)) def data_api_service_descriptor(self): # type: () -> endpoints_pb2.ApiServiceDescriptor return endpoints_pb2.ApiServiceDescriptor( url=self.port_from_worker(self._grpc_server.data_port)) def state_api_service_descriptor(self): # type: () -> endpoints_pb2.ApiServiceDescriptor return endpoints_pb2.ApiServiceDescriptor( url=self.port_from_worker(self._grpc_server.state_port)) def logging_api_service_descriptor(self): # type: () -> endpoints_pb2.ApiServiceDescriptor return endpoints_pb2.ApiServiceDescriptor( url=self.port_from_worker(self._grpc_server.logging_port)) def close(self): # type: () -> None self.control_conn.close() self.data_conn.close() super().close() def port_from_worker(self, port): # type: (int) -> str return '%s:%s' % (self.host_from_worker(), port) def host_from_worker(self): # type: () -> str return 'localhost' @WorkerHandler.register_environment( common_urns.environments.EXTERNAL.urn, beam_runner_api_pb2.ExternalPayload) class ExternalWorkerHandler(GrpcWorkerHandler): def __init__( self, external_payload, # type: beam_runner_api_pb2.ExternalPayload state, # type: StateServicer provision_info, # type: ExtendedProvisionInfo grpc_server # type: GrpcServer ): # type: (...) -> None super().__init__(state, provision_info, grpc_server) self._external_payload = external_payload def start_worker(self): # type: () -> None _LOGGER.info("Requesting worker at %s", self._external_payload.endpoint.url) stub = beam_fn_api_pb2_grpc.BeamFnExternalWorkerPoolStub( GRPCChannelFactory.insecure_channel( self._external_payload.endpoint.url)) _LOGGER.info('self.control_address: %s' % self.control_address) control_descriptor = endpoints_pb2.ApiServiceDescriptor( url=self.control_address) response = stub.StartWorker( beam_fn_api_pb2.StartWorkerRequest( worker_id=self.worker_id, control_endpoint=control_descriptor, artifact_endpoint=control_descriptor, provision_endpoint=control_descriptor, logging_endpoint=self.logging_api_service_descriptor(), params=self._external_payload.params)) if response.error: raise RuntimeError("Error starting worker: %s" % response.error) def stop_worker(self): # type: () -> None pass def host_from_worker(self): # type: () -> str # TODO(https://github.com/apache/beam/issues/19947): Reconcile across # platforms. if sys.platform in ['win32', 'darwin']: return 'localhost' import socket return socket.getfqdn() @WorkerHandler.register_environment(python_urns.EMBEDDED_PYTHON_GRPC, bytes) class EmbeddedGrpcWorkerHandler(GrpcWorkerHandler): def __init__( self, payload, # type: bytes state, # type: StateServicer provision_info, # type: ExtendedProvisionInfo grpc_server # type: GrpcServer ): # type: (...) -> None super().__init__(state, provision_info, grpc_server) from apache_beam.transforms.environments import EmbeddedPythonGrpcEnvironment config = EmbeddedPythonGrpcEnvironment.parse_config(payload.decode('utf-8')) self._state_cache_size = ( config.get('state_cache_size') or STATE_CACHE_SIZE_MB) << 20 self._data_buffer_time_limit_ms = \ config.get('data_buffer_time_limit_ms') or DATA_BUFFER_TIME_LIMIT_MS def start_worker(self): # type: () -> None self.worker = sdk_worker.SdkHarness( self.control_address, state_cache_size=self._state_cache_size, data_buffer_time_limit_ms=self._data_buffer_time_limit_ms, worker_id=self.worker_id, runner_capabilities=FNAPI_RUNNER_CAPABILITIES) self.worker_thread = threading.Thread( name='run_worker', target=self.worker.run) self.worker_thread.daemon = True self.worker_thread.start() def stop_worker(self): # type: () -> None self.worker_thread.join() # The subprocesses module is not threadsafe on Python 2.7. Use this lock to # prevent concurrent calls to Popen(). SUBPROCESS_LOCK = threading.Lock() @WorkerHandler.register_environment(python_urns.SUBPROCESS_SDK, bytes) class SubprocessSdkWorkerHandler(GrpcWorkerHandler): def __init__( self, worker_command_line, # type: bytes state, # type: StateServicer provision_info, # type: ExtendedProvisionInfo grpc_server # type: GrpcServer ): # type: (...) -> None super().__init__(state, provision_info, grpc_server) self._worker_command_line = worker_command_line def start_worker(self): # type: () -> None from apache_beam.runners.portability import local_job_service self.worker = local_job_service.SubprocessSdkWorker( self._worker_command_line, self.control_address, self.provision_info, self.worker_id) self.worker_thread = threading.Thread( name='run_worker', target=self.worker.run) self.worker_thread.start() def stop_worker(self): # type: () -> None self.worker_thread.join() @WorkerHandler.register_environment( common_urns.environments.DOCKER.urn, beam_runner_api_pb2.DockerPayload) class DockerSdkWorkerHandler(GrpcWorkerHandler): def __init__( self, payload, # type: beam_runner_api_pb2.DockerPayload state, # type: StateServicer provision_info, # type: ExtendedProvisionInfo grpc_server # type: GrpcServer ): # type: (...) -> None super().__init__(state, provision_info, grpc_server) self._container_image = payload.container_image self._container_id = None # type: Optional[bytes] def host_from_worker(self): # type: () -> str if sys.platform == 'darwin': # See https://docs.docker.com/docker-for-mac/networking/ return 'host.docker.internal' if sys.platform == 'linux' and is_in_notebook(): import socket # Gets ipv4 address of current host. Note the host is not guaranteed to # be localhost because the python SDK could be running within a container. return socket.gethostbyname(socket.getfqdn()) return super().host_from_worker() def start_worker(self): # type: () -> None credential_options = [] try: # This is the public facing API, skip if it is not available. # (If this succeeds but the imports below fail, better to actually raise # an error below rather than silently fail.) # pylint: disable=unused-import import google.auth except ImportError: pass else: from google.auth import _cloud_sdk from google.auth import environment_vars gcloud_cred_file = os.environ.get( environment_vars.CREDENTIALS, _cloud_sdk.get_application_default_credentials_path()) if os.path.exists(gcloud_cred_file): docker_cred_file = '/docker_cred_file.json' credential_options.extend([ '--mount', f'type=bind,source={gcloud_cred_file},target={docker_cred_file}', '--env', f'{environment_vars.CREDENTIALS}={docker_cred_file}' ]) with SUBPROCESS_LOCK: try: _LOGGER.info('Attempting to pull image %s', self._container_image) subprocess.check_call(['docker', 'pull', self._container_image]) except Exception: _LOGGER.info( 'Unable to pull image %s, defaulting to local image if it exists' % self._container_image) self._container_id = subprocess.check_output([ 'docker', 'run', '-d', '--network=host', ] + credential_options + [ self._container_image, '--id=%s' % self.worker_id, '--logging_endpoint=%s' % self.logging_api_service_descriptor().url, '--control_endpoint=%s' % self.control_address, '--artifact_endpoint=%s' % self.control_address, '--provision_endpoint=%s' % self.control_address, ]).strip() assert self._container_id is not None while True: status = subprocess.check_output([ 'docker', 'inspect', '-f', '{{.State.Status}}', self._container_id ]).strip() _LOGGER.info( 'Waiting for docker to start up. Current status is %s' % status.decode('utf-8')) if status == b'running': _LOGGER.info( 'Docker container is running. container_id = %s, ' 'worker_id = %s', self._container_id, self.worker_id) break elif status in (b'dead', b'exited'): subprocess.call(['docker', 'container', 'logs', self._container_id]) raise RuntimeError( 'SDK failed to start. Final status is %s' % status.decode('utf-8')) time.sleep(1) self._done = False t = threading.Thread(target=self.watch_container) t.daemon = True t.start() def watch_container(self): # type: () -> None while not self._done: assert self._container_id is not None status = subprocess.check_output( ['docker', 'inspect', '-f', '{{.State.Status}}', self._container_id]).strip() if status != b'running': if not self._done: logs = subprocess.check_output([ 'docker', 'container', 'logs', '--tail', '10', self._container_id ], stderr=subprocess.STDOUT) _LOGGER.info(logs) self.control_conn.abort( RuntimeError( 'SDK exited unexpectedly. ' 'Final status is %s. Final log line is %s' % ( status.decode('utf-8'), logs.decode('utf-8').strip().rsplit('\n', maxsplit=1)[-1]))) time.sleep(5) def stop_worker(self): # type: () -> None self._done = True if self._container_id: with SUBPROCESS_LOCK: subprocess.call(['docker', 'kill', self._container_id]) class WorkerHandlerManager(object): """ Manages creation of ``WorkerHandler``s. Caches ``WorkerHandler``s based on environment id. """ def __init__( self, environments, # type: Mapping[str, beam_runner_api_pb2.Environment] job_provision_info # type: ExtendedProvisionInfo ): # type: (...) -> None self._environments = environments self._job_provision_info = job_provision_info self._cached_handlers = collections.defaultdict( list) # type: DefaultDict[str, List[WorkerHandler]] self._workers_by_id = {} # type: Dict[str, WorkerHandler] self.state_servicer = StateServicer() self._grpc_server = None # type: Optional[GrpcServer] self._process_bundle_descriptors = { } # type: Dict[str, beam_fn_api_pb2.ProcessBundleDescriptor] def register_process_bundle_descriptor(self, process_bundle_descriptor): # type: (beam_fn_api_pb2.ProcessBundleDescriptor) -> None self._process_bundle_descriptors[ process_bundle_descriptor.id] = process_bundle_descriptor def get_process_bundle_descriptor(self, request): # type: (beam_fn_api_pb2.GetProcessBundleDescriptorRequest) -> beam_fn_api_pb2.ProcessBundleDescriptor return self._process_bundle_descriptors[ request.process_bundle_descriptor_id] def get_worker_handlers( self, environment_id, # type: Optional[str] num_workers # type: int ): # type: (...) -> List[WorkerHandler] if environment_id is None: # Any environment will do, pick one arbitrarily. environment_id = next(iter(self._environments.keys())) environment = self._environments[environment_id] if environment.urn == common_urns.environments.ANYOF.urn: payload = beam_runner_api_pb2.AnyOfEnvironmentPayload.FromString( environment.payload) env_rankings = { python_urns.EMBEDDED_PYTHON: 10, common_urns.environments.EXTERNAL.urn: 5, common_urns.environments.DOCKER.urn: 1, } environment = sorted( payload.environments, key=lambda env: env_rankings.get(env.urn, -1))[-1] # assume all environments except EMBEDDED_PYTHON use gRPC. if environment.urn == python_urns.EMBEDDED_PYTHON: # special case for EmbeddedWorkerHandler: there's no need for a gRPC # server, but we need to pass self instead. Cast to make the type check # on WorkerHandler.create() think we have a GrpcServer instance. grpc_server = cast(GrpcServer, self) elif self._grpc_server is None: self._grpc_server = GrpcServer( self.state_servicer, self._job_provision_info, self) grpc_server = self._grpc_server else: grpc_server = self._grpc_server worker_handler_list = self._cached_handlers[environment_id] if len(worker_handler_list) < num_workers: for _ in range(len(worker_handler_list), num_workers): worker_handler = WorkerHandler.create( environment, self.state_servicer, self._job_provision_info.for_environment(environment), grpc_server) _LOGGER.debug( "Created Worker handler %s for environment %s (%s, %r)", worker_handler, environment_id, environment.urn, environment.payload) self._cached_handlers[environment_id].append(worker_handler) self._workers_by_id[worker_handler.worker_id] = worker_handler worker_handler.start_worker() return self._cached_handlers[environment_id][:num_workers] def close_all(self): # type: () -> None for worker_handler_list in self._cached_handlers.values(): for worker_handler in set(worker_handler_list): try: worker_handler.close() except Exception: _LOGGER.error( "Error closing worker_handler %s" % worker_handler, exc_info=True) self._cached_handlers = collections.defaultdict(list) self._workers_by_id = {} if self._grpc_server is not None: self._grpc_server.close() self._grpc_server = None def get_worker(self, worker_id): # type: (str) -> WorkerHandler return self._workers_by_id[worker_id] class StateServicer(beam_fn_api_pb2_grpc.BeamFnStateServicer, sdk_worker.StateHandler): _SUPPORTED_STATE_TYPES = frozenset([ 'runner', 'multimap_side_input', 'multimap_keys_side_input', 'multimap_keys_values_side_input', 'iterable_side_input', 'bag_user_state', 'multimap_user_state', 'ordered_list_user_state' ]) class CopyOnWriteState(object): def __init__(self, underlying): # type: (DefaultDict[bytes, Buffer]) -> None self._underlying = underlying self._overlay = {} # type: Dict[bytes, Buffer] def __getitem__(self, key): # type: (bytes) -> Buffer if key in self._overlay: return self._overlay[key] else: return StateServicer.CopyOnWriteList( self._underlying, self._overlay, key) def __delitem__(self, key): # type: (bytes) -> None self._overlay[key] = [] def commit(self): # type: () -> DefaultDict[bytes, Buffer] self._underlying.update(self._overlay) return self._underlying class CopyOnWriteList(object): def __init__( self, underlying, # type: DefaultDict[bytes, Buffer] overlay, # type: Dict[bytes, Buffer] key # type: bytes ): # type: (...) -> None self._underlying = underlying self._overlay = overlay self._key = key def __iter__(self): # type: () -> Iterator[bytes] if self._key in self._overlay: return iter(self._overlay[self._key]) else: return iter(self._underlying[self._key]) def append(self, item): # type: (bytes) -> None if self._key not in self._overlay: self._overlay[self._key] = list(self._underlying[self._key]) self._overlay[self._key].append(item) def extend(self, other: Buffer) -> None: raise NotImplementedError() StateType = Union[CopyOnWriteState, collections.defaultdict[bytes, Buffer]] def __init__(self): # type: () -> None self._lock = threading.Lock() self._state = collections.defaultdict(list) # type: StateServicer.StateType self._checkpoint = None # type: Optional[StateServicer.StateType] self._use_continuation_tokens = False self._continuations = {} # type: Dict[bytes, Tuple[bytes, ...]] self._ordered_list_keys = collections.defaultdict( SortedSet) # type: DefaultDict[bytes, SortedSet] def checkpoint(self): # type: () -> None assert self._checkpoint is None and not \ isinstance(self._state, StateServicer.CopyOnWriteState) self._checkpoint = self._state self._state = StateServicer.CopyOnWriteState(self._state) def commit(self): # type: () -> None assert isinstance(self._state, StateServicer.CopyOnWriteState) and \ isinstance(self._checkpoint, StateServicer.CopyOnWriteState) self._state.commit() self._state = self._checkpoint.commit() self._checkpoint = None def restore(self): # type: () -> None assert self._checkpoint is not None self._state = self._checkpoint self._checkpoint = None @contextlib.contextmanager def process_instruction_id(self, unused_instruction_id): # type: (Any) -> Iterator yield def _get_one_interval_key(self, state_key, start): # type: (beam_fn_api_pb2.StateKey, int) -> bytes state_key_copy = beam_fn_api_pb2.StateKey() state_key_copy.CopyFrom(state_key) state_key_copy.ordered_list_user_state.range.start = start state_key_copy.ordered_list_user_state.range.end = start + 1 return self._to_key(state_key_copy) def get_raw( self, state_key, # type: beam_fn_api_pb2.StateKey continuation_token=None # type: Optional[bytes] ): # type: (...) -> Tuple[bytes, Optional[bytes]] if state_key.WhichOneof('type') not in self._SUPPORTED_STATE_TYPES: raise NotImplementedError( 'Unknown state type: ' + state_key.WhichOneof('type')) # type: ignore[operator] with self._lock: if not continuation_token: # Compute full_state only when no continuation token is provided. # If there is continuation token, full_state is already in # continuation cache. No need to recompute. full_state = [] # type: List[bytes] if state_key.WhichOneof('type') == 'ordered_list_user_state': maybe_start = state_key.ordered_list_user_state.range.start maybe_end = state_key.ordered_list_user_state.range.end persistent_state_key = beam_fn_api_pb2.StateKey() persistent_state_key.CopyFrom(state_key) persistent_state_key.ordered_list_user_state.ClearField("range") available_keys = self._ordered_list_keys[self._to_key( persistent_state_key)] for i in available_keys.irange(maybe_start, maybe_end, inclusive=(True, False)): entries = self._state[self._get_one_interval_key( persistent_state_key, i)] full_state.extend(entries) else: full_state.extend(self._state[self._to_key(state_key)]) if self._use_continuation_tokens: # The token is "nonce:index". if not continuation_token: token_base = b'token_%x' % len(self._continuations) self._continuations[token_base] = tuple(full_state) return b'', b'%s:0' % token_base else: token_base, index = continuation_token.split(b':') ix = int(index) full_state_cont = self._continuations[token_base] if ix == len(full_state_cont): return b'', None else: return full_state_cont[ix], b'%s:%d' % (token_base, ix + 1) else: assert not continuation_token return b''.join(full_state), None def append_raw( self, state_key, # type: beam_fn_api_pb2.StateKey data # type: bytes ): # type: (...) -> _Future with self._lock: if state_key.WhichOneof('type') == 'ordered_list_user_state': coder = coders.TupleCoder([ coders.VarIntCoder(), coders.coders.LengthPrefixCoder(coders.BytesCoder()) ]).get_impl() for key, value in coder.decode_all(data): self._state[self._get_one_interval_key(state_key, key)].append( coder.encode((key, value))) self._ordered_list_keys[self._to_key(state_key)].add(key) else: self._state[self._to_key(state_key)].append(data) return _Future.done() def clear(self, state_key): # type: (beam_fn_api_pb2.StateKey) -> _Future with self._lock: try: if state_key.WhichOneof('type') == 'ordered_list_user_state': start = state_key.ordered_list_user_state.range.start end = state_key.ordered_list_user_state.range.end persistent_state_key = beam_fn_api_pb2.StateKey() persistent_state_key.CopyFrom(state_key) persistent_state_key.ordered_list_user_state.ClearField("range") available_keys = self._ordered_list_keys[self._to_key( persistent_state_key)] for i in list(available_keys.irange(start, end, inclusive=(True, False))): del self._state[self._get_one_interval_key(persistent_state_key, i)] available_keys.remove(i) else: del self._state[self._to_key(state_key)] except KeyError: # This may happen with the caching layer across bundles. Caching may # skip this storage layer for a blocking_get(key) request. Without # the caching, the state for a key would be initialized via the # defaultdict that _state uses. pass return _Future.done() def done(self): # type: () -> None pass @staticmethod def _to_key(state_key): # type: (beam_fn_api_pb2.StateKey) -> bytes return state_key.SerializeToString() class GrpcStateServicer(beam_fn_api_pb2_grpc.BeamFnStateServicer): def __init__(self, state): # type: (StateServicer) -> None self._state = state def State( self, request_stream, # type: Iterable[beam_fn_api_pb2.StateRequest] context=None # type: Any ): # type: (...) -> Iterator[beam_fn_api_pb2.StateResponse] # Note that this eagerly mutates state, assuming any failures are fatal. # Thus it is safe to ignore instruction_id. for request in request_stream: try: request_type = request.WhichOneof('request') if request_type == 'get': data, continuation_token = self._state.get_raw( request.state_key, request.get.continuation_token) yield beam_fn_api_pb2.StateResponse( id=request.id, get=beam_fn_api_pb2.StateGetResponse( data=data, continuation_token=continuation_token)) elif request_type == 'append': self._state.append_raw(request.state_key, request.append.data) yield beam_fn_api_pb2.StateResponse( id=request.id, append=beam_fn_api_pb2.StateAppendResponse()) elif request_type == 'clear': self._state.clear(request.state_key) yield beam_fn_api_pb2.StateResponse( id=request.id, clear=beam_fn_api_pb2.StateClearResponse()) else: raise NotImplementedError('Unknown state request: %s' % request_type) except Exception as exn: yield beam_fn_api_pb2.StateResponse(id=request.id, error=str(exn)) class SingletonStateHandlerFactory(sdk_worker.StateHandlerFactory): """A singleton cache for a StateServicer.""" def __init__(self, state_handler): # type: (sdk_worker.CachingStateHandler) -> None self._state_handler = state_handler def create_state_handler(self, api_service_descriptor): # type: (endpoints_pb2.ApiServiceDescriptor) -> sdk_worker.CachingStateHandler """Returns the singleton state handler.""" return self._state_handler def close(self): # type: () -> None """Does nothing.""" pass class ControlFuture(object): def __init__( self, instruction_id, # type: str response=None # type: Optional[beam_fn_api_pb2.InstructionResponse] ): # type: (...) -> None self.instruction_id = instruction_id self._response = response if response is None: self._condition = threading.Condition() self._exception = None # type: Optional[Exception] def is_done(self): # type: () -> bool return self._response is not None def set(self, response): # type: (beam_fn_api_pb2.InstructionResponse) -> None with self._condition: self._response = response self._condition.notify_all() def get(self, timeout=None): # type: (Optional[float]) -> beam_fn_api_pb2.InstructionResponse if not self._response and not self._exception: with self._condition: if not self._response and not self._exception: self._condition.wait(timeout) if self._exception: raise self._exception else: assert self._response is not None return self._response def abort(self, exception): # type: (Exception) -> None with self._condition: self._exception = exception self._condition.notify_all() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/fn_api_runner/worker_handlers_test.py0000644000175100017510000000400715206113032031343 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import unittest from apache_beam.portability.api import beam_provision_api_pb2 from apache_beam.runners.portability.fn_api_runner.fn_runner import ExtendedProvisionInfo from apache_beam.runners.portability.fn_api_runner.worker_handlers import WorkerHandlerManager from apache_beam.transforms import environments _LOGGER = logging.getLogger(__name__) class WorkerHandlerManagerTest(unittest.TestCase): def test_close_all(self): inprocess_env = environments.EmbeddedPythonEnvironment( capabilities=environments.python_sdk_capabilities(), artifacts=(), resource_hints={}).to_runner_api(None) envs = { 'inprocess': inprocess_env, } prov_info = ExtendedProvisionInfo( beam_provision_api_pb2.ProvisionInfo( retrieval_token='unused-retrieval-token')) manager = WorkerHandlerManager(envs, job_provision_info=prov_info) first_workers = manager.get_worker_handlers('inprocess', 1) manager.close_all() second_workers = manager.get_worker_handlers('inprocess', 1) assert len(first_workers) == len(second_workers) == 1 assert first_workers != second_workers if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/job_server.py0000644000175100017510000001350615206113032024432 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import atexit import logging import shutil import signal import tempfile import threading import grpc from apache_beam.options import pipeline_options from apache_beam.portability.api import beam_job_api_pb2_grpc from apache_beam.runners.portability import local_job_service from apache_beam.utils import subprocess_server from apache_beam.version import __version__ as beam_version class JobServer(object): def start(self): """Starts this JobServer, returning a grpc service to which to submit jobs. """ raise NotImplementedError(type(self)) def stop(self): """Stops this job server.""" raise NotImplementedError(type(self)) class ExternalJobServer(JobServer): def __init__(self, endpoint, timeout=None): self._endpoint = endpoint self._timeout = timeout def start(self) -> beam_job_api_pb2_grpc.JobServiceStub: channel = grpc.insecure_channel(self._endpoint) grpc.channel_ready_future(channel).result(timeout=self._timeout) return beam_job_api_pb2_grpc.JobServiceStub(channel) def stop(self): pass class EmbeddedJobServer(JobServer): def start(self) -> 'local_job_service.LocalJobServicer': return local_job_service.LocalJobServicer() def stop(self): pass class StopOnExitJobServer(JobServer): """Wraps a JobServer such that its stop will automatically be called on exit. """ def __init__(self, job_server): self._lock = threading.Lock() self._job_server = job_server self._started = False # save original signal handler self._original_sigint_handler = signal.getsignal(signal.SIGINT) def _sigint_handler(self, sig, frame): self.stop() if callable(self._original_sigint_handler): # call original signal handler to handle sigint gracefully self._original_sigint_handler(sig, frame) def start(self): with self._lock: if not self._started: self._endpoint = self._job_server.start() self._started = True atexit.register(self.stop) signal.signal(signal.SIGINT, self._sigint_handler) return self._endpoint def stop(self): with self._lock: if self._started: try: self._job_server.stop() finally: self._started = False # Unregister the atexit handler to prevent duplicate # registrations when the server is restarted/reused. atexit.unregister(self.stop) class SubprocessJobServer(JobServer): """An abstract base class for JobServers run as an external process.""" def __init__(self): self._local_temp_root = None self._server = None self._log_filter = None def subprocess_cmd_and_endpoint(self): raise NotImplementedError(type(self)) def start(self): if self._server is None: self._local_temp_root = tempfile.mkdtemp(prefix='beam-temp') cmd, endpoint = self.subprocess_cmd_and_endpoint() port = int(endpoint.split(':')[-1]) logger = logging.getLogger(f"{self.__class__.__name__}") if self._log_filter is not None: logger.addFilter(self._log_filter) self._server = subprocess_server.SubprocessServer( beam_job_api_pb2_grpc.JobServiceStub, cmd, port=port, logger=logger) return self._server.start() def stop(self): if self._local_temp_root: shutil.rmtree(self._local_temp_root) self._local_temp_root = None return self._server.stop() def local_temp_dir(self, **kwargs): return tempfile.mkdtemp(dir=self._local_temp_root, **kwargs) class JavaJarJobServer(SubprocessJobServer): def __init__(self, options): super().__init__() options = options.view_as(pipeline_options.JobServerOptions) self._job_port = options.job_port self._artifact_port = options.artifact_port self._expansion_port = options.expansion_port self._artifacts_dir = options.artifacts_dir self._java_launcher = options.job_server_java_launcher self._jvm_properties = options.job_server_jvm_properties self._jar_cache_dir = options.jar_cache_dir def java_arguments( self, job_port, artifact_port, expansion_port, artifacts_dir): raise NotImplementedError(type(self)) def path_to_jar(self): raise NotImplementedError(type(self)) @staticmethod def path_to_beam_jar(gradle_target, artifact_id=None): return subprocess_server.JavaJarServer.path_to_beam_jar( gradle_target, artifact_id=artifact_id) @staticmethod def local_jar(url, jar_cache_dir=None, user_agent=None): return subprocess_server.JavaJarServer.local_jar( url, jar_cache_dir, user_agent) def subprocess_cmd_and_endpoint(self): jar_path = self.local_jar(self.path_to_jar(), self._jar_cache_dir) artifacts_dir = ( self._artifacts_dir if self._artifacts_dir else self.local_temp_dir( prefix='artifacts')) job_port, = subprocess_server.pick_port(self._job_port) subprocess_cmd = [self._java_launcher, '-jar'] + self._jvm_properties + [ jar_path ] + list( self.java_arguments( job_port, self._artifact_port, self._expansion_port, artifacts_dir)) return (subprocess_cmd, 'localhost:%s' % job_port) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/job_server_test.py0000644000175100017510000000501515206113032025465 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import unittest from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.runners.portability.job_server import JavaJarJobServer class JavaJarJobServerStub(JavaJarJobServer): def java_arguments( self, job_port, artifact_port, expansion_port, artifacts_dir): return [ '--artifacts-dir', artifacts_dir, '--job-port', job_port, '--artifact-port', artifact_port, '--expansion-port', expansion_port ] def path_to_jar(self): return '/path/to/jar' @staticmethod def local_jar(url, jar_cache_dir=None): logging.debug("url({%s}), jar_cache_dir({%s})", url, jar_cache_dir) return url class JavaJarJobServerTest(unittest.TestCase): def test_subprocess_cmd_and_endpoint(self): pipeline_options = PipelineOptions([ '--job_port=8099', '--artifact_port=8098', '--expansion_port=8097', '--artifacts_dir=/path/to/artifacts/', '--job_server_java_launcher=/path/to/java', '--job_server_jvm_properties=-Dsome.property=value' ]) job_server = JavaJarJobServerStub(pipeline_options) subprocess_cmd, endpoint = job_server.subprocess_cmd_and_endpoint() self.assertEqual( subprocess_cmd, [ '/path/to/java', '-jar', '-Dsome.property=value', '/path/to/jar', '--artifacts-dir', '/path/to/artifacts/', '--job-port', 8099, '--artifact-port', 8098, '--expansion-port', 8097 ]) self.assertEqual(endpoint, 'localhost:8099') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/local_job_service.py0000644000175100017510000004273415206113032025743 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import concurrent.futures import itertools import logging import os import queue import shutil import subprocess import tempfile import threading import time import traceback from typing import Any from typing import Mapping from typing import Optional import grpc from google.protobuf import json_format from google.protobuf import struct_pb2 from google.protobuf import text_format from apache_beam import pipeline from apache_beam.metrics import monitoring_infos from apache_beam.options import pipeline_options from apache_beam.portability.api import beam_artifact_api_pb2_grpc from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.portability.api import beam_job_api_pb2 from apache_beam.portability.api import beam_job_api_pb2_grpc from apache_beam.portability.api import beam_provision_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.portability.api import endpoints_pb2 from apache_beam.runners.job import utils as job_utils from apache_beam.runners.portability import abstract_job_service from apache_beam.runners.portability import artifact_service from apache_beam.runners.portability import portable_runner from apache_beam.runners.portability.fn_api_runner import fn_runner from apache_beam.runners.portability.fn_api_runner import worker_handlers from apache_beam.runners.worker.log_handler import LOGENTRY_TO_LOG_LEVEL_MAP from apache_beam.transforms import environments from apache_beam.utils import thread_pool_executor _LOGGER = logging.getLogger(__name__) def _iter_queue(q): while True: yield q.get(block=True) class LocalJobServicer(abstract_job_service.AbstractJobServiceServicer): """Manages one or more pipelines, possibly concurrently. Experimental: No backward compatibility guaranteed. Servicer for the Beam Job API. This JobService uses a basic local implementation of runner to run the job. This JobService is not capable of managing job on remote clusters. By default, this JobService executes the job in process but still uses GRPC to communicate pipeline and worker state. It can also be configured to use inline calls rather than GRPC (for speed) or launch completely separate subprocesses for the runner and worker(s). """ def __init__(self, staging_dir=None, beam_job_type=None): super().__init__() self._cleanup_staging_dir = staging_dir is None self._staging_dir = staging_dir or tempfile.mkdtemp() self._artifact_service = artifact_service.ArtifactStagingService( artifact_service.BeamFilesystemHandler(self._staging_dir).file_writer) self._artifact_staging_endpoint: Optional[ endpoints_pb2.ApiServiceDescriptor] = None self._beam_job_type = beam_job_type or BeamJob def create_beam_job( self, preparation_id, # stype: str job_name: str, pipeline: beam_runner_api_pb2.Pipeline, options: struct_pb2.Struct) -> 'BeamJob': self._artifact_service.register_job( staging_token=preparation_id, dependency_sets=_extract_dependency_sets( pipeline.components.environments)) provision_info = fn_runner.ExtendedProvisionInfo( beam_provision_api_pb2.ProvisionInfo(pipeline_options=options), self._staging_dir, job_name=job_name) return self._beam_job_type( preparation_id, pipeline, options, provision_info, self._artifact_staging_endpoint, self._artifact_service) def get_bind_address(self): """Return the address used to open the port on the gRPC server. This is often, but not always the same as the service address. For example, to make the service accessible to external machines, override this to return '[::]' and override `get_service_address()` to return a publicly accessible host name. """ return self.get_service_address() def get_service_address(self): """Return the host name at which this server will be accessible. In particular, this is provided to the client upon connection as the artifact staging endpoint. """ return 'localhost' def start_grpc_server(self, port=0): options = [("grpc.max_receive_message_length", -1), ("grpc.max_send_message_length", -1), ("grpc.http2.max_pings_without_data", 0), ("grpc.http2.max_ping_strikes", 0)] self._server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=options) port = self._server.add_insecure_port( '%s:%d' % (self.get_bind_address(), port)) beam_job_api_pb2_grpc.add_JobServiceServicer_to_server(self, self._server) beam_artifact_api_pb2_grpc.add_ArtifactStagingServiceServicer_to_server( self._artifact_service, self._server) hostname = self.get_service_address() self._artifact_staging_endpoint = endpoints_pb2.ApiServiceDescriptor( url='%s:%d' % (hostname, port)) self._server.start() _LOGGER.info('Grpc server started at %s on port %d' % (hostname, port)) return port def stop(self, timeout=1): self._server.stop(timeout) if os.path.exists(self._staging_dir) and self._cleanup_staging_dir: shutil.rmtree(self._staging_dir, ignore_errors=True) def GetJobMetrics(self, request, context=None): if request.job_id not in self._jobs: raise LookupError("Job {} does not exist".format(request.job_id)) result = self._jobs[request.job_id].result if result is None: monitoring_info_list = [] else: monitoring_info_list = result.monitoring_infos() # Filter out system metrics user_monitoring_info_list = [ x for x in monitoring_info_list if monitoring_infos.is_user_monitoring_info(x) ] return beam_job_api_pb2.GetJobMetricsResponse( metrics=beam_job_api_pb2.MetricResults( committed=user_monitoring_info_list)) class SubprocessSdkWorker(object): """Manages a SDK worker implemented as a subprocess communicating over grpc. """ def __init__( self, worker_command_line: bytes, control_address, provision_info, worker_id=None): # worker_command_line is of bytes type received from grpc. It was encoded in # apache_beam.transforms.environments.SubprocessSDKEnvironment earlier. # decode it back as subprocess.Popen does not support bytes args in win32. self._worker_command_line = worker_command_line.decode('utf-8') self._control_address = control_address self._provision_info = provision_info self._worker_id = worker_id def run(self): options = [("grpc.http2.max_pings_without_data", 0), ("grpc.http2.max_ping_strikes", 0)] logging_server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=options) logging_port = logging_server.add_insecure_port('[::]:0') logging_server.start() logging_servicer = BeamFnLoggingServicer() beam_fn_api_pb2_grpc.add_BeamFnLoggingServicer_to_server( logging_servicer, logging_server) logging_descriptor = text_format.MessageToString( endpoints_pb2.ApiServiceDescriptor(url='localhost:%s' % logging_port)) control_descriptor = text_format.MessageToString( endpoints_pb2.ApiServiceDescriptor(url=self._control_address)) pipeline_options = json_format.MessageToJson( self._provision_info.provision_info.pipeline_options) env_dict = dict( os.environ, CONTROL_API_SERVICE_DESCRIPTOR=control_descriptor, LOGGING_API_SERVICE_DESCRIPTOR=logging_descriptor, PIPELINE_OPTIONS=pipeline_options) # only add worker_id when it is set. if self._worker_id: env_dict['WORKER_ID'] = self._worker_id with worker_handlers.SUBPROCESS_LOCK: p = subprocess.Popen(self._worker_command_line, shell=True, env=env_dict) try: p.wait() if p.returncode: raise RuntimeError( 'Worker subprocess exited with return code %s' % p.returncode) finally: if p.poll() is None: p.kill() logging_server.stop(0) class BeamJob(abstract_job_service.AbstractBeamJob): """This class handles running and managing a single pipeline. The current state of the pipeline is available as self.state. """ def __init__( self, job_id: str, pipeline, options, provision_info: fn_runner.ExtendedProvisionInfo, artifact_staging_endpoint: Optional[endpoints_pb2.ApiServiceDescriptor], artifact_service: artifact_service.ArtifactStagingService, ): super().__init__(job_id, provision_info.job_name, pipeline, options) self._provision_info = provision_info self._artifact_staging_endpoint = artifact_staging_endpoint self._artifact_service = artifact_service self._state_queues: list[queue.Queue] = [] self._log_queues = JobLogQueues() self.daemon = True self.result = None def pipeline_options(self): def from_urn(key): assert key.startswith('beam:option:') assert key.endswith(':v1') return key[12:-3] return pipeline_options.PipelineOptions( **{ from_urn(key): value for (key, value ) in job_utils.struct_to_dict(self._pipeline_options).items() }) def set_state(self, new_state): """Set the latest state as an int enum and notify consumers""" timestamp = super().set_state(new_state) if timestamp is not None: # Inform consumers of the new state. for queue in self._state_queues: queue.put((new_state, timestamp)) def prepare(self): pass def artifact_staging_endpoint(self): return self._artifact_staging_endpoint def run(self): self.set_state(beam_job_api_pb2.JobState.STARTING) self._run_thread = threading.Thread(target=self._run_job) self._run_thread.start() def _run_job(self): with JobLogHandler(self._log_queues) as log_handler: self._update_dependencies() pipeline.Pipeline.merge_compatible_environments(self._pipeline_proto) try: start = time.time() self.result = self._invoke_runner() self.result.wait_until_finish() _LOGGER.info( 'Completed job in %s seconds with state %s.', time.time() - start, self.result.state) self.set_state( portable_runner.PipelineResult.pipeline_state_to_runner_api_state( self.result.state)) except: # pylint: disable=bare-except self._log_queues.put( beam_job_api_pb2.JobMessage( message_id=log_handler._next_id(), time=time.strftime('%Y-%m-%d %H:%M:%S.'), importance=beam_job_api_pb2.JobMessage.JOB_MESSAGE_ERROR, message_text=traceback.format_exc())) _LOGGER.exception('Error running pipeline.') self.set_state(beam_job_api_pb2.JobState.FAILED) def _invoke_runner(self): self.set_state(beam_job_api_pb2.JobState.RUNNING) return fn_runner.FnApiRunner( provision_info=self._provision_info).run_via_runner_api( self._pipeline_proto, self.pipeline_options()) def _update_dependencies(self): try: _update_dependency_sets( self._pipeline_proto.components.environments, self._artifact_service.resolved_deps(self._job_id, timeout=0)) self._provision_info.provision_info.ClearField('retrieval_token') except concurrent.futures.TimeoutError: # TODO(https://github.com/apache/beam/issues/20267): Require this once # all SDKs support it. pass def cancel(self): if not self.is_terminal_state(self.state): self.set_state(beam_job_api_pb2.JobState.CANCELLING) # TODO(robertwb): Actually cancel... self.set_state(beam_job_api_pb2.JobState.CANCELLED) def get_state_stream(self): # Register for any new state changes. state_queue = queue.Queue() self._state_queues.append(state_queue) for state, timestamp in self.with_state_history(_iter_queue(state_queue)): yield state, timestamp if self.is_terminal_state(state): break def get_message_stream(self): # Register for any new messages. log_queue = queue.Queue() self._log_queues.append(log_queue) self._state_queues.append(log_queue) for msg in itertools.chain(self._log_queues.cache(), self.with_state_history(_iter_queue(log_queue))): if isinstance(msg, tuple): assert len(msg) == 2 and isinstance(msg[0], int) current_state = msg[0] yield msg if self.is_terminal_state(current_state): break else: yield msg class BeamFnLoggingServicer(beam_fn_api_pb2_grpc.BeamFnLoggingServicer): def Logging(self, log_bundles, context=None): for log_bundle in log_bundles: for log_entry in log_bundle.log_entries: _LOGGER.log( LOGENTRY_TO_LOG_LEVEL_MAP[log_entry.severity], 'Worker: %s', str(log_entry).replace('\n', ' ')) return iter([]) class JobLogQueues(object): def __init__(self): self._queues: list[queue.Queue] = [] self._cache = [] self._cache_size = 10 self._lock = threading.Lock() def cache(self): with self._lock: return list(self._cache) def append(self, queue): with self._lock: self._queues.append(queue) def put(self, msg): with self._lock: if len(self._cache) < self._cache_size: self._cache.append(msg) else: min_level = min(m.importance for m in self._cache) if msg.importance >= min_level: self._cache.append(msg) for ix, m in enumerate(self._cache): if m.importance == min_level: del self._cache[ix] break for queue in self._queues: queue.put(msg) class JobLogHandler(logging.Handler): """Captures logs to be returned via the Beam Job API. Enabled via the with statement.""" # Mapping from logging levels to LogEntry levels. LOG_LEVEL_MAP = { logging.FATAL: beam_job_api_pb2.JobMessage.JOB_MESSAGE_ERROR, logging.CRITICAL: beam_job_api_pb2.JobMessage.JOB_MESSAGE_ERROR, logging.ERROR: beam_job_api_pb2.JobMessage.JOB_MESSAGE_ERROR, logging.WARNING: beam_job_api_pb2.JobMessage.JOB_MESSAGE_WARNING, logging.INFO: beam_job_api_pb2.JobMessage.JOB_MESSAGE_BASIC, logging.DEBUG: beam_job_api_pb2.JobMessage.JOB_MESSAGE_DEBUG, } def __init__(self, log_queues): super().__init__() self._last_id = 0 self._logged_thread = None self._log_queues = log_queues def __enter__(self): # Remember the current thread to demultiplex the logs of concurrently # running pipelines (as Python log handlers are global). self._logged_thread = threading.current_thread() logging.getLogger().addHandler(self) return self def __exit__(self, *args): self._logged_thread = None self.close() def _next_id(self): self._last_id += 1 return str(self._last_id) def emit(self, record): if self._logged_thread is threading.current_thread(): msg = beam_job_api_pb2.JobMessage( message_id=self._next_id(), time=time.strftime( '%Y-%m-%d %H:%M:%S.', time.localtime(record.created)), importance=self.LOG_LEVEL_MAP[record.levelno], message_text=self.format(record)) # Inform all message consumers. self._log_queues.put(msg) def _extract_dependency_sets( envs: Mapping[str, beam_runner_api_pb2.Environment] ) -> Mapping[Any, list[beam_runner_api_pb2.ArtifactInformation]]: """Expands the set of environments into a mapping of (opaque) keys to dependency sets. This is not 1:1 in the case of AnyOf environments. The values can then be resolved and the mapping passed back to _update_dependency_sets to update the dependencies in the original protos. """ def dependencies_iter(): for env_id, env in envs.items(): for ix, sub_env in enumerate(environments.expand_anyof_environments(env)): yield (env_id, ix), sub_env.dependencies return dict(dependencies_iter()) def _update_dependency_sets( envs: Mapping[str, beam_runner_api_pb2.Environment], resolved_deps: Mapping[Any, list[beam_runner_api_pb2.ArtifactInformation]]): """Takes the mapping of beam Environments (originally passed to `_extract_dependency_sets`) and a set of (key-wise) updated dependencies, and updates the original environment protos to contain the updated dependencies. """ for env_id, env in envs.items(): new_envs = [] for ix, sub_env in enumerate(environments.expand_anyof_environments(env)): # Slice assignment not supported for repeated fields. del sub_env.dependencies[:] sub_env.dependencies.extend(resolved_deps[env_id, ix]) new_envs.append(sub_env) if len(new_envs) == 1: envs[env_id].CopyFrom(new_envs[0]) else: envs[env_id].CopyFrom( environments.AnyOfEnvironment.create_proto(new_envs)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/local_job_service_main.py0000644000175100017510000001317715206113032026746 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file """Starts a service for running portable beam pipelines. The basic usage is simply python -m apache_beam.runners.portability.local_job_service_main Many other options are also supported, such as starting in the background or passing in a lockfile to ensure that only one copy of the service is running at a time. Pass --help to see them all. """ import argparse import logging import os import pathlib import signal import subprocess import sys import time from apache_beam.runners.portability import local_job_service _LOGGER = logging.getLogger(__name__) def run(argv): if argv[0] == __file__: argv = argv[1:] parser = argparse.ArgumentParser() parser.add_argument( '-p', '--port', '--job_port', type=int, default=0, help='port on which to serve the job api') parser.add_argument('--staging_dir') parser.add_argument( '--pid_file', help='File in which to store the process id of the server.') parser.add_argument( '--port_file', help='File in which to store the port of the server.') parser.add_argument( '--background', action='store_true', help='Start the server up as a background process.' ' Will fail if pid_file already exists, unless --stop is also specified.') parser.add_argument( '--stderr_file', help='Where to write stderr (if not specified, merged with stdout).') parser.add_argument( '--stdout_file', help='Where to write stdout for background job service.') parser.add_argument( '--stop', action='store_true', help='Stop the existing process, if any, specified in pid_file.' ' Will not start up a new service unless --background is specified.') options = parser.parse_args(argv) if options.stop: if not options.pid_file: raise RuntimeError('--pid_file must be specified with --stop') if os.path.exists(options.pid_file): with open(options.pid_file) as fin: pid = int(fin.read()) print('Killing process at', pid) try: os.kill(pid, signal.SIGTERM) except Exception: print('Process', pid, 'already killed.') os.unlink(options.pid_file) else: print('Process id file', options.pid_file, 'already removed.') if not options.background: return if options.background: if not options.pid_file: raise RuntimeError('--pid_file must be specified with --start') if options.stop: argv.remove('--stop') argv.remove('--background') if not options.port_file: options.port_file = os.path.splitext(options.pid_file)[0] + '.port' argv.append('--port_file') argv.append(options.port_file) if not options.stdout_file: raise RuntimeError('--stdout_file must be specified with --background') os.makedirs(pathlib.PurePath(options.stdout_file).parent, exist_ok=True) stdout_dest = open(options.stdout_file, mode='w') if options.stderr_file: os.makedirs(pathlib.PurePath(options.stderr_file).parent, exist_ok=True) stderr_dest = open(options.stderr_file, mode='w') else: stderr_dest = subprocess.STDOUT subprocess.Popen([ sys.executable, '-m', 'apache_beam.runners.portability.local_job_service_main' ] + argv, stderr=stderr_dest, stdout=stdout_dest) print('Waiting for server to start up...') while not os.path.exists(options.port_file): time.sleep(.1) with open(options.port_file) as fin: port = fin.read() print('Server started at port', port) return if options.pid_file: print('Writing process id to', options.pid_file) os.makedirs(pathlib.PurePath(options.pid_file).parent, exist_ok=True) fd = os.open(options.pid_file, os.O_CREAT | os.O_EXCL | os.O_RDWR) with os.fdopen(fd, 'w') as fout: fout.write(str(os.getpid())) try: job_servicer = local_job_service.LocalJobServicer(options.staging_dir) port = job_servicer.start_grpc_server(options.port) try: if options.port_file: print('Writing port to', options.port_file) os.makedirs(pathlib.PurePath(options.port_file).parent, exist_ok=True) with open(options.port_file + '.tmp', 'w') as fout: fout.write(str(port)) os.rename(options.port_file + '.tmp', options.port_file) serve("Listening for beam jobs on port %d." % port) finally: job_servicer.stop() finally: if options.pid_file and os.path.exists(options.pid_file): os.unlink(options.pid_file) if options.port_file and os.path.exists(options.port_file): os.unlink(options.port_file) def serve(msg): logging_delay = 30 while True: _LOGGER.info(msg) time.sleep(logging_delay) logging_delay *= 1.25 if __name__ == '__main__': signal.signal(signal.SIGTERM, lambda *args: sys.exit(0)) logging.basicConfig() logging.getLogger().setLevel(logging.INFO) run(sys.argv) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/local_job_service_test.py0000644000175100017510000001100115206113032026761 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import unittest import mock from apache_beam.options import pipeline_options from apache_beam.portability.api import beam_job_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.portability import local_job_service from apache_beam.runners.portability.portable_runner import JobServiceHandle from apache_beam.runners.portability.portable_runner import PortableRunner class TestJobServicePlan(JobServiceHandle): def __init__(self, job_service): self.job_service = job_service self.options = None self.timeout = None self.artifact_endpoint = None def get_pipeline_options(self): return None class LocalJobServerTest(unittest.TestCase): def test_end_to_end(self): job_service = local_job_service.LocalJobServicer() job_service.start_grpc_server() plan = TestJobServicePlan(job_service) _, message_stream, state_stream = plan.submit( beam_runner_api_pb2.Pipeline()) state_results = list(state_stream) message_results = list(message_stream) expected_states = [ beam_job_api_pb2.JobState.STOPPED, beam_job_api_pb2.JobState.STARTING, beam_job_api_pb2.JobState.RUNNING, beam_job_api_pb2.JobState.DONE, ] self.assertEqual([s.state for s in state_results], expected_states) self.assertEqual([ s.state_response.state for s in message_results if s.HasField('state_response') ], expected_states) def test_error_messages_after_pipeline_failure(self): job_service = local_job_service.LocalJobServicer() job_service.start_grpc_server() plan = TestJobServicePlan(job_service) job_id, message_stream, state_stream = plan.submit( beam_runner_api_pb2.Pipeline(requirements=['unsupported_requirement'])) message_results = list(message_stream) state_results = list(state_stream) expected_states = [ beam_job_api_pb2.JobState.STOPPED, beam_job_api_pb2.JobState.STARTING, beam_job_api_pb2.JobState.RUNNING, beam_job_api_pb2.JobState.FAILED, ] self.assertEqual([s.state for s in state_results], expected_states) self.assertTrue( any( 'unsupported_requirement' in m.message_response.message_text for m in message_results), message_results) # Assert we still see the error message if we fetch error messages after # the job has completed. messages_again = list( plan.job_service.GetMessageStream( beam_job_api_pb2.JobMessagesRequest(job_id=job_id))) self.assertTrue( any( 'unsupported_requirement' in m.message_response.message_text for m in message_results), messages_again) def test_artifact_service_override(self): job_service = local_job_service.LocalJobServicer() port = job_service.start_grpc_server() test_artifact_endpoint = 'testartifactendpoint:4242' options = pipeline_options.PipelineOptions([ '--job_endpoint', 'localhost:%d' % port, '--artifact_endpoint', test_artifact_endpoint, ]) runner = PortableRunner() job_service_handle = runner.create_job_service(options) with mock.patch.object(job_service_handle, 'stage') as mocked_stage: job_service_handle.submit(beam_runner_api_pb2.Pipeline()) mocked_stage.assert_called_once_with( mock.ANY, test_artifact_endpoint, mock.ANY) # Confirm the artifact_endpoint is in the options protobuf options_proto = job_service_handle.get_pipeline_options() self.assertEqual( options_proto['beam:option:artifact_endpoint:v1'], test_artifact_endpoint) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/portable_metrics.py0000644000175100017510000000511615206113032025626 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging from apache_beam.metrics import monitoring_infos from apache_beam.metrics.execution import MetricKey from apache_beam.metrics.metric import MetricName _LOGGER = logging.getLogger(__name__) def from_monitoring_infos(monitoring_info_list, user_metrics_only=False): """Groups MonitoringInfo objects into counters, distributions, gauges and string sets Args: monitoring_info_list: An iterable of MonitoringInfo objects. user_metrics_only: If true, includes user metrics only. Returns: A tuple containing three dictionaries: counters, distributions, gauges and string set, respectively. Each dictionary contains (MetricKey, metric result) pairs. """ counters = {} distributions = {} gauges = {} string_sets = {} bounded_tries = {} for mi in monitoring_info_list: if (user_metrics_only and not monitoring_infos.is_user_monitoring_info(mi)): continue try: key = _create_metric_key(mi) except ValueError as e: _LOGGER.debug(str(e)) continue metric_result = (monitoring_infos.extract_metric_result_map_value(mi)) if monitoring_infos.is_counter(mi): counters[key] = metric_result elif monitoring_infos.is_distribution(mi): distributions[key] = metric_result elif monitoring_infos.is_gauge(mi): gauges[key] = metric_result elif monitoring_infos.is_string_set(mi): string_sets[key] = metric_result elif monitoring_infos.is_bounded_trie(mi): bounded_tries[key] = metric_result return counters, distributions, gauges, string_sets, bounded_tries def _create_metric_key(monitoring_info): step_name = monitoring_infos.get_step_name(monitoring_info) namespace, name = monitoring_infos.parse_namespace_and_name(monitoring_info) return MetricKey(step_name, MetricName(namespace, name)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/portable_runner.py0000644000175100017510000005627615206113032025506 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # mypy: check-untyped-defs import atexit import copy import functools import itertools import logging import threading import time from typing import Any from typing import Iterator from typing import Optional import grpc from google.protobuf import struct_pb2 from apache_beam.metrics import metric from apache_beam.metrics.execution import MetricResult from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import PortableOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.value_provider import ValueProvider from apache_beam.pipeline import Pipeline from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.portability.api import beam_artifact_api_pb2_grpc from apache_beam.portability.api import beam_job_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners import runner from apache_beam.runners.job import utils as job_utils from apache_beam.runners.portability import artifact_service from apache_beam.runners.portability import job_server from apache_beam.runners.portability import portable_metrics from apache_beam.runners.portability.fn_api_runner.fn_runner import translations from apache_beam.runners.worker import sdk_worker_main from apache_beam.runners.worker import worker_pool_main from apache_beam.transforms import environments __all__ = ['PortableRunner'] MESSAGE_LOG_LEVELS = { beam_job_api_pb2.JobMessage.MESSAGE_IMPORTANCE_UNSPECIFIED: logging.INFO, beam_job_api_pb2.JobMessage.JOB_MESSAGE_DEBUG: logging.DEBUG, beam_job_api_pb2.JobMessage.JOB_MESSAGE_DETAILED: logging.DEBUG, beam_job_api_pb2.JobMessage.JOB_MESSAGE_BASIC: logging.INFO, beam_job_api_pb2.JobMessage.JOB_MESSAGE_WARNING: logging.WARNING, beam_job_api_pb2.JobMessage.JOB_MESSAGE_ERROR: logging.ERROR, } TERMINAL_STATES = [ beam_job_api_pb2.JobState.DONE, beam_job_api_pb2.JobState.DRAINED, beam_job_api_pb2.JobState.FAILED, beam_job_api_pb2.JobState.CANCELLED, ] _LOGGER = logging.getLogger(__name__) class JobServiceHandle(object): """ Encapsulates the interactions necessary to submit a pipeline to a job service. The base set of interactions consists of 3 steps: - prepare - stage - run """ def __init__(self, job_service, options, retain_unknown_options=False): self.job_service = job_service self.options = options self.timeout = options.view_as(PortableOptions).job_server_timeout self.artifact_endpoint = options.view_as(PortableOptions).artifact_endpoint self._retain_unknown_options = retain_unknown_options def submit( self, proto_pipeline: beam_runner_api_pb2.Pipeline ) -> tuple[str, Iterator[beam_job_api_pb2.JobStateEvent], Iterator[beam_job_api_pb2.JobMessagesResponse]]: """ Submit and run the pipeline defined by `proto_pipeline`. """ prepare_response = self.prepare(proto_pipeline) artifact_endpoint = ( self.artifact_endpoint or prepare_response.artifact_staging_endpoint.url) self.stage( proto_pipeline, artifact_endpoint, prepare_response.staging_session_token) return self.run(prepare_response.preparation_id) def get_pipeline_options(self) -> struct_pb2.Struct: """ Get `self.options` as a protobuf Struct """ # fetch runner options from job service # retries in case the channel is not ready def send_options_request(max_retries=5): num_retries = 0 while True: try: # This reports channel is READY but connections may fail # Seems to be only an issue on Mac with port forwardings return self.job_service.DescribePipelineOptions( beam_job_api_pb2.DescribePipelineOptionsRequest(), timeout=self.timeout) except grpc.FutureTimeoutError: # no retry for timeout errors raise except grpc.RpcError as e: num_retries += 1 if num_retries > max_retries: raise e time.sleep(1) options_response = send_options_request() def add_runner_options(parser): for option in options_response.options: try: # no default values - we don't want runner options # added unless they were specified by the user add_arg_args = {'action': 'store'} if option.description is not None: # Prevent bare %'s in help strings in 3.14+ add_arg_args['help'] = option.description.replace("%", "%%") if option.type == beam_job_api_pb2.PipelineOptionType.BOOLEAN: add_arg_args['action'] = 'store_true' \ if option.default_value != 'true' else 'store_false' elif option.type == beam_job_api_pb2.PipelineOptionType.INTEGER: add_arg_args['type'] = int # type: ignore elif option.type == beam_job_api_pb2.PipelineOptionType.ARRAY: add_arg_args['action'] = 'append' parser.add_argument("--%s" % option.name, **add_arg_args) except Exception as e: # ignore runner options that are already present # only in this case is duplicate not treated as error if 'conflicting option string' not in str(e): raise _LOGGER.debug("Runner option '%s' was already added" % option.name) all_options = self.options.get_all_options( add_extra_args_fn=add_runner_options, retain_unknown_options=self._retain_unknown_options) return self.encode_pipeline_options(all_options) @staticmethod def encode_pipeline_options( all_options: dict[str, Any]) -> 'struct_pb2.Struct': def convert_pipeline_option_value(v): # convert int values: BEAM-5509 if type(v) == int: return str(v) elif isinstance(v, ValueProvider): return convert_pipeline_option_value( v.get()) if v.is_accessible() else None return v # TODO: Define URNs for options. p_options = { 'beam:option:' + k + ':v1': convert_pipeline_option_value(v) for k, v in all_options.items() if v is not None } return job_utils.dict_to_struct(p_options) def prepare( self, proto_pipeline: beam_runner_api_pb2.Pipeline ) -> beam_job_api_pb2.PrepareJobResponse: """Prepare the job on the job service""" return self.job_service.Prepare( beam_job_api_pb2.PrepareJobRequest( job_name='job', pipeline=proto_pipeline, pipeline_options=self.get_pipeline_options()), timeout=self.timeout) def stage( self, proto_pipeline: beam_runner_api_pb2.Pipeline, artifact_staging_endpoint, staging_session_token) -> None: """Stage artifacts""" if artifact_staging_endpoint: artifact_service.offer_artifacts( beam_artifact_api_pb2_grpc.ArtifactStagingServiceStub( channel=grpc.insecure_channel(artifact_staging_endpoint)), artifact_service.ArtifactRetrievalService( artifact_service.BeamFilesystemHandler(None).file_reader), staging_session_token) def run( self, preparation_id: str ) -> tuple[str, Iterator[beam_job_api_pb2.JobStateEvent], Iterator[beam_job_api_pb2.JobMessagesResponse]]: """Run the job""" try: state_stream = self.job_service.GetStateStream( beam_job_api_pb2.GetJobStateRequest(job_id=preparation_id), timeout=self.timeout) # If there's an error, we don't always get it until we try to read. # Fortunately, there's always an immediate current state published. state_stream = itertools.chain([next(state_stream)], state_stream) message_stream = self.job_service.GetMessageStream( beam_job_api_pb2.JobMessagesRequest(job_id=preparation_id), timeout=self.timeout) except Exception: # TODO(https://github.com/apache/beam/issues/19284): Unify preparation_id # and job_id for all runners. state_stream = message_stream = None # Run the job and wait for a result, we don't set a timeout here because # it may take a long time for a job to complete and streaming # jobs currently never return a response. run_response = self.job_service.Run( beam_job_api_pb2.RunJobRequest(preparation_id=preparation_id)) if state_stream is None: state_stream = self.job_service.GetStateStream( beam_job_api_pb2.GetJobStateRequest(job_id=run_response.job_id)) message_stream = self.job_service.GetMessageStream( beam_job_api_pb2.JobMessagesRequest(job_id=run_response.job_id)) return run_response.job_id, message_stream, state_stream class PortableRunner(runner.PipelineRunner): """ Experimental: No backward compatibility guaranteed. A BeamRunner that executes Python pipelines via the Beam Job API. This runner is a stub and does not run the actual job. This runner schedules the job on a job service. The responsibility of running and managing the job lies with the job service used. """ def __init__(self): self._dockerized_job_server: Optional[job_server.JobServer] = None @staticmethod def _create_environment(options: PipelineOptions) -> environments.Environment: return environments.Environment.from_options( options.view_as(PortableOptions)) def default_job_server(self, options): raise NotImplementedError( 'You must specify a --job_endpoint when using --runner=PortableRunner. ' 'Alternatively, you may specify which portable runner you intend to ' 'use, such as --runner=FlinkRunner or --runner=SparkRunner.') def create_job_service_handle(self, job_service, options) -> JobServiceHandle: return JobServiceHandle(job_service, options) def create_job_service(self, options: PipelineOptions) -> JobServiceHandle: """ Start the job service and return a `JobServiceHandle` """ job_endpoint = options.view_as(PortableOptions).job_endpoint if job_endpoint: if job_endpoint == 'embed': server: job_server.JobServer = job_server.EmbeddedJobServer() else: job_server_timeout = options.view_as(PortableOptions).job_server_timeout server = job_server.ExternalJobServer(job_endpoint, job_server_timeout) else: server = self.default_job_server(options) return self.create_job_service_handle(server.start(), options) @staticmethod def get_proto_pipeline( pipeline: Pipeline, options: PipelineOptions) -> beam_runner_api_pb2.Pipeline: proto_pipeline = pipeline.to_runner_api( default_environment=environments.Environment.from_options( options.view_as(PortableOptions))) return PortableRunner._optimize_pipeline(proto_pipeline, options) @staticmethod def _optimize_pipeline( proto_pipeline: beam_runner_api_pb2.Pipeline, options: PipelineOptions) -> beam_runner_api_pb2.Pipeline: # TODO: https://github.com/apache/beam/issues/19493 # Eventually remove the 'pre_optimize' option alltogether and only perform # the equivalent of the 'default' case below (minus the 'lift_combiners' # part). pre_optimize = options.view_as(DebugOptions).lookup_experiment( 'pre_optimize', 'default').lower() if (not options.view_as(StandardOptions).streaming and pre_optimize != 'none'): if pre_optimize == 'default': phases = [ # TODO: https://github.com/apache/beam/issues/18584 # https://github.com/apache/beam/issues/18586 # Eventually remove the 'lift_combiners' phase from 'default'. translations.pack_combiners, translations.lift_combiners, translations.sort_stages ] partial = True elif pre_optimize == 'all': phases = translations.standard_optimize_phases() partial = False elif pre_optimize == 'all_except_fusion': # TODO(https://github.com/apache/beam/issues/19422): Delete this branch # after PortableRunner supports beam:runner:executable_stage:v1. phases = translations.standard_optimize_phases() phases.remove(translations.greedily_fuse) partial = True else: phases = [] for phase_name in pre_optimize.split(','): # For now, these are all we allow. if phase_name in ('pack_combiners', 'lift_combiners', 'expand_sdf'): phases.append(getattr(translations, phase_name)) else: raise ValueError( 'Unknown or inapplicable phase for pre_optimize: %s' % phase_name) phases.append(translations.sort_stages) partial = True # All (known) portable runners (ie Flink and Spark) support these URNs. known_urns = frozenset([ common_urns.composites.RESHUFFLE.urn, common_urns.primitives.IMPULSE.urn, common_urns.primitives.FLATTEN.urn, common_urns.primitives.GROUP_BY_KEY.urn ]) proto_pipeline = translations.optimize_pipeline( proto_pipeline, phases=phases, known_runner_urns=known_urns, partial=partial) return proto_pipeline def run_portable_pipeline( self, pipeline: beam_runner_api_pb2.Pipeline, options: PipelineOptions) -> runner.PipelineResult: portable_options = options.view_as(PortableOptions) # Do not set a Runner. Otherwise this can cause problems in Java's # PipelineOptions, i.e. ClassNotFoundException, if the corresponding Runner # does not exist in the Java SDK. In portability, the entry point is clearly # defined via the JobService. portable_options.view_as(StandardOptions).runner = None cleanup_callbacks = self.start_and_replace_loopback_environments( pipeline, options) optimized_pipeline = self._optimize_pipeline(pipeline, options) job_service_handle = self.create_job_service(options) job_id, message_stream, state_stream = job_service_handle.submit( optimized_pipeline) result = PipelineResult( job_service_handle.job_service, job_id, message_stream, state_stream, cleanup_callbacks) if cleanup_callbacks: # Register an exit handler to ensure cleanup on exit. atexit.register(functools.partial(result._cleanup, on_exit=True)) _LOGGER.info( 'Environment "%s" has started a component necessary for the ' 'execution. Be sure to run the pipeline using\n' ' with Pipeline() as p:\n' ' p.apply(..)\n' 'This ensures that the pipeline finishes before this program exits.', portable_options.environment_type) return result @staticmethod def start_and_replace_loopback_environments(pipeline, options): portable_options = copy.deepcopy(options.view_as(PortableOptions)) experiments = options.view_as(DebugOptions).experiments or [] cleanup_callbacks = [] for env in pipeline.components.environments.values(): if env.urn == python_urns.EMBEDDED_PYTHON_LOOPBACK: # Start a worker and change the environment to point to that worker. use_loopback_process_worker = options.view_as( DebugOptions).lookup_experiment( 'use_loopback_process_worker', False) portable_options.environment_type = 'EXTERNAL' portable_options.environment_config, server = ( worker_pool_main.BeamFnExternalWorkerPoolServicer.start( state_cache_size= sdk_worker_main._get_state_cache_size_bytes( options=options), data_buffer_time_limit_ms= sdk_worker_main._get_data_buffer_time_limit_ms(experiments), use_process=use_loopback_process_worker)) external_env = environments.ExternalEnvironment.from_options( portable_options).to_runner_api(None) # type: ignore env.urn = external_env.urn env.payload = external_env.payload cleanup_callbacks.append(functools.partial(server.stop, 1)) return cleanup_callbacks class PortableMetrics(metric.MetricResults): def __init__(self, job_metrics_response): metrics = job_metrics_response.metrics self.attempted = portable_metrics.from_monitoring_infos(metrics.attempted) self.committed = portable_metrics.from_monitoring_infos(metrics.committed) @staticmethod def _combine(committed, attempted, filter): all_keys = set(committed.keys()) | set(attempted.keys()) return [ MetricResult(key, committed.get(key), attempted.get(key)) for key in all_keys if metric.MetricResults.matches(filter, key) ] def query(self, filter=None): counters, distributions, gauges, stringsets, bounded_tries = [ self._combine(x, y, filter) for x, y in zip(self.committed, self.attempted) ] return { self.COUNTERS: counters, self.DISTRIBUTIONS: distributions, self.GAUGES: gauges, self.STRINGSETS: stringsets, self.BOUNDED_TRIES: bounded_tries, } class PipelineResult(runner.PipelineResult): def __init__( self, job_service, job_id, message_stream, state_stream, cleanup_callbacks=()): super().__init__(beam_job_api_pb2.JobState.UNSPECIFIED) self._job_service = job_service self._job_id = job_id self._messages = [] self._message_stream = message_stream self._state_stream = state_stream self._cleanup_callbacks = cleanup_callbacks self._metrics = None self._runtime_exception = None def cancel(self) -> None: try: self._job_service.Cancel( beam_job_api_pb2.CancelJobRequest(job_id=self._job_id)) finally: self._cleanup() @property def state(self): runner_api_state = self._job_service.GetState( beam_job_api_pb2.GetJobStateRequest(job_id=self._job_id)).state self._state = self.runner_api_state_to_pipeline_state(runner_api_state) return self._state @staticmethod def runner_api_state_to_pipeline_state(runner_api_state): return getattr( runner.PipelineState, beam_job_api_pb2.JobState.Enum.Name(runner_api_state)) @staticmethod def pipeline_state_to_runner_api_state(pipeline_state): if pipeline_state == runner.PipelineState.PENDING: return beam_job_api_pb2.JobState.STARTING else: try: return beam_job_api_pb2.JobState.Enum.Value(pipeline_state) except ValueError: return beam_job_api_pb2.JobState.UNSPECIFIED def metrics(self): if not self._metrics: job_metrics_response = self._job_service.GetJobMetrics( beam_job_api_pb2.GetJobMetricsRequest(job_id=self._job_id)) self._metrics = PortableMetrics(job_metrics_response) return self._metrics def _last_error_message(self) -> str: # Filter only messages with the "message_response" and error messages. messages = [ m.message_response for m in self._messages if m.HasField('message_response') ] error_messages = [ m for m in messages if m.importance == beam_job_api_pb2.JobMessage.JOB_MESSAGE_ERROR ] if error_messages: return error_messages[-1].message_text else: return 'unknown error' def wait_until_finish(self, duration=None): """ :param duration: The maximum time in milliseconds to wait for the result of the execution. If None or zero, will wait until the pipeline finishes. :return: The result of the pipeline, i.e. PipelineResult. """ last_error_text = None def read_messages() -> None: nonlocal last_error_text previous_state = -1 for message in self._message_stream: if message.HasField('message_response'): mr = message.message_response logging.log(MESSAGE_LOG_LEVELS[mr.importance], "%s", mr.message_text) if mr.importance == beam_job_api_pb2.JobMessage.JOB_MESSAGE_ERROR: last_error_text = mr.message_text else: current_state = message.state_response.state if current_state != previous_state: _LOGGER.info( "Job state changed to %s", self.runner_api_state_to_pipeline_state(current_state)) previous_state = current_state self._messages.append(message) message_thread = threading.Thread( target=read_messages, name='wait_until_finish_read') message_thread.daemon = True message_thread.start() if duration: state_thread = threading.Thread( target=functools.partial(self._observe_state, message_thread), name='wait_until_finish_state_observer') state_thread.daemon = True state_thread.start() start_time = time.time() duration_secs = duration / 1000 while (time.time() - start_time < duration_secs and state_thread.is_alive()): time.sleep(1) else: self._observe_state(message_thread) if self._runtime_exception: raise self._runtime_exception from apache_beam.runners.runner import PipelineState if self._state == PipelineState.FAILED: raise RuntimeError(last_error_text or "Pipeline failed.") return self._state def _observe_state(self, message_thread): try: for state_response in self._state_stream: self._state = self.runner_api_state_to_pipeline_state( state_response.state) if state_response.state in TERMINAL_STATES: # Wait for any last messages. message_thread.join(10) break if self._state != runner.PipelineState.DONE: self._runtime_exception = RuntimeError( 'Pipeline %s failed in state %s: %s' % (self._job_id, self._state, self._last_error_message())) except Exception as e: self._runtime_exception = e finally: self._cleanup() def _cleanup(self, on_exit: bool = False) -> None: if on_exit and self._cleanup_callbacks: _LOGGER.info( 'Running cleanup on exit. If your pipeline should continue running, ' 'be sure to use the following syntax:\n' ' with Pipeline() as p:\n' ' p.apply(..)\n' 'This ensures that the pipeline finishes before this program exits.') callback_exceptions = [] for callback in self._cleanup_callbacks: try: callback() except Exception as e: callback_exceptions.append(e) self._cleanup_callbacks = () if callback_exceptions: formatted_exceptions = ''.join( [f"\n\t{repr(e)}" for e in callback_exceptions]) raise RuntimeError('Errors: {}'.format(formatted_exceptions)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/portable_runner_test.py0000644000175100017510000004767315206113032026546 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import inspect import logging import socket import subprocess import sys import tempfile import time import unittest from unittest.mock import MagicMock import grpc import apache_beam as beam from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import DirectOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import PortableOptions from apache_beam.portability import python_urns from apache_beam.portability.api import beam_job_api_pb2 from apache_beam.portability.api import beam_job_api_pb2_grpc from apache_beam.runners.portability import portable_runner from apache_beam.runners.portability.fn_api_runner import fn_runner_test from apache_beam.runners.portability.local_job_service import LocalJobServicer from apache_beam.runners.portability.portable_runner import PortableRunner from apache_beam.runners.worker import worker_pool_main from apache_beam.runners.worker.channel_factory import GRPCChannelFactory from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import environments from apache_beam.transforms import userstate _LOGGER = logging.getLogger(__name__) class PortableRunnerTest(fn_runner_test.FnApiRunnerTest): # Controls job service interaction, not sdk harness interaction. _use_subprocesses = False @classmethod def _pick_unused_port(cls): return cls._pick_unused_ports(num_ports=1)[0] @staticmethod def _pick_unused_ports(num_ports): """Not perfect, but we have to provide a port to the subprocess.""" # TODO(robertwb): Consider letting the subprocess communicate a choice of # port back. sockets = [] ports = [] for _ in range(0, num_ports): s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sockets.append(s) s.bind(('localhost', 0)) _, port = s.getsockname() ports.append(port) try: return ports finally: for s in sockets: s.close() @classmethod def _start_local_runner_subprocess_job_service(cls): cls._maybe_kill_subprocess() # TODO(robertwb): Consider letting the subprocess pick one and # communicate it back... # pylint: disable=unbalanced-tuple-unpacking job_port, expansion_port = cls._pick_unused_ports(num_ports=2) _LOGGER.info('Starting server on port %d.', job_port) cls._subprocess = subprocess.Popen( cls._subprocess_command(job_port, expansion_port)) address = 'localhost:%d' % job_port job_service = beam_job_api_pb2_grpc.JobServiceStub( GRPCChannelFactory.insecure_channel(address)) _LOGGER.info('Waiting for server to be ready...') start = time.time() # Long timeout based previously flaky test. See issue #22115. timeout = 300 while True: time.sleep(0.1) if cls._subprocess.poll() is not None: raise RuntimeError( 'Subprocess terminated unexpectedly with exit code %d.' % cls._subprocess.returncode) elif time.time() - start > timeout: raise RuntimeError( 'Pipeline timed out waiting for job service subprocess.') else: try: job_service.GetState( beam_job_api_pb2.GetJobStateRequest(job_id='[fake]')) break except grpc.RpcError as exn: if exn.code() != grpc.StatusCode.UNAVAILABLE: # We were able to contact the service for our fake state request. break _LOGGER.info('Server ready.') return address @classmethod def _get_job_endpoint(cls): if '_job_endpoint' not in cls.__dict__: cls._job_endpoint = cls._create_job_endpoint() return cls._job_endpoint @classmethod def _create_job_endpoint(cls): if cls._use_subprocesses: return cls._start_local_runner_subprocess_job_service() else: cls._servicer = LocalJobServicer() return 'localhost:%d' % cls._servicer.start_grpc_server() @classmethod def get_runner(cls): return portable_runner.PortableRunner() @classmethod def tearDownClass(cls): cls._maybe_kill_subprocess() @classmethod def _maybe_kill_subprocess(cls): if hasattr(cls, '_subprocess') and cls._subprocess.poll() is None: cls._subprocess.kill() time.sleep(0.1) def create_options(self): def get_pipeline_name(): for _, _, _, method_name, _, _ in inspect.stack(): if method_name.find('test') != -1: return method_name return 'unknown_test' # Set the job name for better debugging. options = PipelineOptions.from_dictionary( {'job_name': get_pipeline_name() + '_' + str(time.time())}) options.view_as(PortableOptions).job_endpoint = self._get_job_endpoint() # Override the default environment type for testing. options.view_as(PortableOptions).environment_type = ( python_urns.EMBEDDED_PYTHON) # Enable caching (disabled by default) options.view_as(DebugOptions).add_experiment('state_cache_size=100') # Enable time-based data buffer (disabled by default) options.view_as(DebugOptions).add_experiment( 'data_buffer_time_limit_ms=1000') return options def create_pipeline(self, is_drain=False): return beam.Pipeline(self.get_runner(), self.create_options()) def test_pardo_state_with_custom_key_coder(self): """Tests that state requests work correctly when the key coder is an SDK-specific coder, i.e. non standard coder. This is additionally enforced by Java's ProcessBundleDescriptorsTest and by Flink's ExecutableStageDoFnOperator which detects invalid encoding by checking for the correct key group of the encoded key.""" index_state_spec = userstate.CombiningValueStateSpec('index', sum) # Test params # Ensure decent amount of elements to serve all partitions n = 200 duplicates = 1 split = n // (duplicates + 1) inputs = [(i % split, str(i % split)) for i in range(0, n)] # Use a DoFn which has to use FastPrimitivesCoder because the type cannot # be inferred class Input(beam.DoFn): def process(self, impulse): for i in inputs: yield i class AddIndex(beam.DoFn): def process(self, kv, index=beam.DoFn.StateParam(index_state_spec)): k, v = kv index.add(1) yield k, v, index.read() expected = [(i % split, str(i % split), i // split + 1) for i in range(0, n)] with self.create_pipeline() as p: assert_that( p | beam.Impulse() | beam.ParDo(Input()) | beam.ParDo(AddIndex()), equal_to(expected)) # Inherits all other tests from fn_api_runner_test.FnApiRunnerTest def test_sdf_default_truncate_when_bounded(self): raise unittest.SkipTest("Portable runners don't support drain yet.") def test_sdf_default_truncate_when_unbounded(self): raise unittest.SkipTest("Portable runners don't support drain yet.") def test_sdf_with_truncate(self): raise unittest.SkipTest("Portable runners don't support drain yet.") def test_draining_sdf_with_sdf_initiated_checkpointing(self): raise unittest.SkipTest("Portable runners don't support drain yet.") def test_pardo_et_timer_with_no_firing(self): if type(self).__name__ in { 'PortableRunnerTest', 'PortableRunnerTestWithSubprocesses', 'PortableRunnerTestWithSubprocessesAndMultiWorkers', 'PortableRunnerTestWithExternalEnv', 'PortableRunnerTestWithLocalDocker', 'PortableRunnerOptimizedWithoutFusion', 'SparkRunnerTest' }: raise unittest.SkipTest("https://github.com/apache/beam/issues/35168") super().test_pardo_et_timer_with_no_firing() def test_pardo_et_timer_with_no_reset(self): if type(self).__name__ in { 'PortableRunnerTest', 'PortableRunnerTestWithSubprocesses', 'PortableRunnerTestWithSubprocessesAndMultiWorkers', 'PortableRunnerTestWithExternalEnv', 'PortableRunnerTestWithLocalDocker', 'PortableRunnerOptimizedWithoutFusion', 'SparkRunnerTest' }: raise unittest.SkipTest("https://github.com/apache/beam/issues/35168") super().test_pardo_et_timer_with_no_reset() @unittest.skip("https://github.com/apache/beam/issues/19422") class PortableRunnerOptimized(PortableRunnerTest): def create_options(self): options = super().create_options() options.view_as(DebugOptions).add_experiment('pre_optimize=all') options.view_as(DebugOptions).add_experiment('state_cache_size=100') options.view_as(DebugOptions).add_experiment( 'data_buffer_time_limit_ms=1000') return options # TODO(https://github.com/apache/beam/issues/19422): Delete this test after # PortableRunner supports beam:runner:executable_stage:v1. class PortableRunnerOptimizedWithoutFusion(PortableRunnerTest): def create_options(self): options = super().create_options() options.view_as(DebugOptions).add_experiment( 'pre_optimize=all_except_fusion') options.view_as(DebugOptions).add_experiment('state_cache_size=100') options.view_as(DebugOptions).add_experiment( 'data_buffer_time_limit_ms=1000') return options class PortableRunnerTestWithExternalEnv(PortableRunnerTest): @classmethod def setUpClass(cls): cls._worker_address, cls._worker_server = ( worker_pool_main.BeamFnExternalWorkerPoolServicer.start( state_cache_size=100, data_buffer_time_limit_ms=1000)) @classmethod def tearDownClass(cls): cls._worker_server.stop(1) def create_options(self): options = super().create_options() options.view_as(PortableOptions).environment_type = 'EXTERNAL' options.view_as(PortableOptions).environment_config = self._worker_address return options class PortableRunnerTestWithSubprocesses(PortableRunnerTest): _use_subprocesses = True def create_options(self): options = super().create_options() options.view_as(PortableOptions).environment_type = ( python_urns.SUBPROCESS_SDK) options.view_as(PortableOptions).environment_config = ( b'%s -m apache_beam.runners.worker.sdk_worker_main' % sys.executable.encode('ascii')).decode('utf-8') # Enable caching (disabled by default) options.view_as(DebugOptions).add_experiment('state_cache_size=100') # Enable time-based data buffer (disabled by default) options.view_as(DebugOptions).add_experiment( 'data_buffer_time_limit_ms=1000') return options @classmethod def _subprocess_command(cls, job_port, _): return [ sys.executable, '-m', 'apache_beam.runners.portability.local_job_service_main', '-p', str(job_port), ] def test_batch_rebatch_pardos(self): raise unittest.SkipTest( "Portable runners with subprocess can't make " "assertions about warnings raised on the worker.") class PortableRunnerTestWithSubprocessesAndMultiWorkers( PortableRunnerTestWithSubprocesses): _use_subprocesses = True def create_options(self): options = super() \ .create_options() options.view_as(DirectOptions).direct_num_workers = 2 return options class PortableRunnerInternalTest(unittest.TestCase): def setUp(self) -> None: self.tmp_dir = tempfile.TemporaryDirectory() self.actual_mkdtemp = tempfile.mkdtemp tempfile.mkdtemp = MagicMock(return_value=self.tmp_dir.name) def tearDown(self) -> None: tempfile.mkdtemp = self.actual_mkdtemp self.tmp_dir.cleanup() def test__create_default_environment(self): docker_image = environments.DockerEnvironment.default_docker_image() self.assertEqual( PortableRunner._create_environment( options=PipelineOptions.from_dictionary( {'sdk_location': 'container'})), environments.DockerEnvironment( container_image=docker_image, artifacts=environments.python_sdk_dependencies(PipelineOptions()))) def test__create_docker_environment(self): docker_image = 'py-docker' self.assertEqual( PortableRunner._create_environment( PipelineOptions.from_dictionary({ 'environment_type': 'DOCKER', 'environment_config': docker_image, 'sdk_location': 'container', })), environments.DockerEnvironment( container_image=docker_image, artifacts=environments.python_sdk_dependencies(PipelineOptions()))) def test__create_process_environment(self): self.assertEqual( PortableRunner._create_environment( PipelineOptions.from_dictionary({ 'environment_type': "PROCESS", 'environment_config': '{"os": "linux", "arch": "amd64", ' '"command": "run.sh", ' '"env":{"k1": "v1"} }', 'sdk_location': 'container', })), environments.ProcessEnvironment( 'run.sh', os='linux', arch='amd64', env={'k1': 'v1'}, artifacts=environments.python_sdk_dependencies(PipelineOptions()))) self.assertEqual( PortableRunner._create_environment( PipelineOptions.from_dictionary({ 'environment_type': 'PROCESS', 'environment_config': '{"command": "run.sh"}', 'sdk_location': 'container', })), environments.ProcessEnvironment( 'run.sh', artifacts=environments.python_sdk_dependencies(PipelineOptions()))) def test__create_external_environment(self): self.assertEqual( PortableRunner._create_environment( PipelineOptions.from_dictionary({ 'environment_type': "EXTERNAL", 'environment_config': 'localhost:50000', 'sdk_location': 'container', })), environments.ExternalEnvironment( 'localhost:50000', artifacts=environments.python_sdk_dependencies(PipelineOptions()))) raw_config = ' {"url":"localhost:50000", "params":{"k1":"v1"}} ' for env_config in (raw_config, raw_config.lstrip(), raw_config.strip()): self.assertEqual( PortableRunner._create_environment( PipelineOptions.from_dictionary({ 'environment_type': "EXTERNAL", 'environment_config': env_config, 'sdk_location': 'container', })), environments.ExternalEnvironment( 'localhost:50000', params={"k1": "v1"}, artifacts=environments.python_sdk_dependencies( PipelineOptions()))) with self.assertRaises(ValueError): PortableRunner._create_environment( PipelineOptions.from_dictionary({ 'environment_type': "EXTERNAL", 'environment_config': '{invalid}', 'sdk_location': 'container', })) with self.assertRaises(ValueError) as ctx: PortableRunner._create_environment( PipelineOptions.from_dictionary({ 'environment_type': "EXTERNAL", 'environment_config': '{"params":{"k1":"v1"}}', 'sdk_location': 'container', })) self.assertIn( 'External environment endpoint must be set.', ctx.exception.args) def hasDockerImage(): image = environments.DockerEnvironment.default_docker_image() try: check_image = subprocess.check_output( "docker images -q %s" % image, shell=True) return check_image != '' except Exception: return False @unittest.skipIf( not hasDockerImage(), "docker not installed or " "no docker image") class PortableRunnerTestWithLocalDocker(PortableRunnerTest): def create_options(self): options = super().create_options() options.view_as(PortableOptions).job_endpoint = 'embed' return options class PortableRunnerOptimizationTest(unittest.TestCase): """Tests for PortableRunner._optimize_pipeline.""" @staticmethod def _transform_urns(proto, options): optimized = PortableRunner._optimize_pipeline(proto, options) return { t.spec.urn for t in optimized.components.transforms.values() if t.spec.urn } def test_custom_optimize_expand_sdf(self): """Verify that expand_sdf can be requested explicitly. See https://github.com/apache/beam/issues/24422. """ from apache_beam.io import restriction_trackers from apache_beam.portability import common_urns class ExpandStringsProvider(beam.transforms.core.RestrictionProvider): def initial_restriction(self, element): return restriction_trackers.OffsetRange(0, len(element)) def create_tracker(self, restriction): return restriction_trackers.OffsetRestrictionTracker(restriction) def restriction_size(self, element, restriction): return restriction.size() class ExpandingStringsDoFn(beam.DoFn): def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam( ExpandStringsProvider())): cur = restriction_tracker.current_restriction().start while restriction_tracker.try_claim(cur): yield element[cur] cur += 1 p = beam.Pipeline() _ = (p | beam.Create(['abc']) | beam.ParDo(ExpandingStringsDoFn())) proto = p.to_runner_api() transform_urns = self._transform_urns( proto, PipelineOptions(['--experiments=pre_optimize=expand_sdf'])) self.assertIn( common_urns.sdf_components.PAIR_WITH_RESTRICTION.urn, transform_urns) self.assertIn( common_urns.sdf_components.SPLIT_AND_SIZE_RESTRICTIONS.urn, transform_urns) self.assertIn( common_urns.sdf_components.PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS.urn, transform_urns) def test_custom_optimize_expands_bounded_read(self): """Verify that iobase.Read(BoundedSource) expands with explicit expand_sdf. This is the end-to-end scenario from https://github.com/apache/beam/issues/24422: Read transforms like ReadFromParquet use SDFs internally. With explicit expand_sdf, these are expanded before reaching the Spark job server as a single ParDo, executing on one partition with no parallelization. """ from apache_beam.io import iobase from apache_beam.portability import common_urns class _FakeBoundedSource(iobase.BoundedSource): def get_range_tracker(self, start_position, stop_position): return None def read(self, range_tracker): return iter([]) def estimate_size(self): return 0 p = beam.Pipeline() _ = p | beam.io.Read(_FakeBoundedSource()) proto = p.to_runner_api() transform_urns = self._transform_urns( proto, PipelineOptions(['--experiments=pre_optimize=expand_sdf'])) # The SDFBoundedSourceReader DoFn should have been expanded into # SDF component stages. self.assertIn( common_urns.sdf_components.PAIR_WITH_RESTRICTION.urn, transform_urns) self.assertIn( common_urns.sdf_components.SPLIT_AND_SIZE_RESTRICTIONS.urn, transform_urns) self.assertIn( common_urns.sdf_components.PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS.urn, transform_urns) # Reshuffle should be present to enable parallelization. self.assertIn(common_urns.composites.RESHUFFLE.urn, transform_urns) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/prism_runner.py0000644000175100017510000004547615206113032025030 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A runner for executing portable pipelines on Apache Beam Prism.""" # this will make using the list parameterized generic happy # on python 3.8 so we aren't revisiting this code after we # sunset it from __future__ import annotations import datetime import hashlib import json import logging import os import platform import re import shutil import stat import subprocess import typing import urllib import zipfile from urllib.error import URLError from urllib.request import urlopen from apache_beam.io.filesystems import FileSystems from apache_beam.options import pipeline_options from apache_beam.runners.portability import job_server from apache_beam.runners.portability import portable_runner from apache_beam.transforms import environments from apache_beam.utils import shared from apache_beam.utils import subprocess_server from apache_beam.version import __version__ as beam_version # pytype: skip-file # Prefix for constructing a download URL GITHUB_DOWNLOAD_PREFIX = 'https://github.com/apache/beam/releases/download/' # Prefix for constructing a release URL, so we can derive a download URL GITHUB_TAG_PREFIX = 'https://github.com/apache/beam/releases/tag/' _LOGGER = logging.getLogger(__name__) class PrismRunner(portable_runner.PortableRunner): """A runner for launching jobs on Prism, automatically downloading and starting a Prism instance if needed. """ shared_handle = shared.Shared() def default_environment( self, options: pipeline_options.PipelineOptions) -> environments.Environment: portable_options = options.view_as(pipeline_options.PortableOptions) if (not portable_options.environment_type and not portable_options.output_executable_path): portable_options.environment_type = 'LOOPBACK' return super().default_environment(options) def default_job_server(self, options): debug_options = options.view_as(pipeline_options.DebugOptions) get_job_server = lambda: job_server.StopOnExitJobServer( PrismJobServer(options)) if debug_options.lookup_experiment("disable_prism_server_singleton"): return get_job_server() return PrismRunner.shared_handle.acquire(get_job_server) def create_job_service_handle(self, job_service, options): return portable_runner.JobServiceHandle( job_service, options, retain_unknown_options=True) def _md5sum(filename, block_size=8192) -> str: md5 = hashlib.md5() with open(filename, 'rb') as f: while True: data = f.read(block_size) if not data: break md5.update(data) return md5.hexdigest() def _rename_if_different(src, dst): assert (os.path.isfile(src)) if os.path.isfile(dst): if _md5sum(src) != _md5sum(dst): # Remove existing binary to prevent exception on Windows during # os.rename. # See: https://docs.python.org/3/library/os.html#os.rename os.remove(dst) os.rename(src, dst) else: _LOGGER.info( 'Found %s and %s with the same md5. Skipping overwrite.' % (src, dst)) os.remove(src) else: os.rename(src, dst) class PrismRunnerLogFilter(logging.Filter): COMMON_FIELDS = set(["level", "source", "msg", "time"]) def filter(self, record): if record.funcName == 'log_stdout': try: message = record.getMessage() json_record = json.loads(message) level_str = json_record["level"] # Example level with offset: 'ERROR+2' if "+" in level_str or "-" in level_str: match = re.match(r"([A-Z]+)([+-]\d+)", level_str) if match: base, offset = match.groups() base_level = getattr(logging, base, logging.INFO) record.levelno = base_level + int(offset) else: record.levelno = getattr(logging, level_str, logging.INFO) else: record.levelno = getattr(logging, level_str, logging.INFO) record.levelname = logging.getLevelName(record.levelno) if "source" in json_record: record.funcName = json_record["source"]["function"] record.pathname = json_record["source"]["file"] record.filename = os.path.basename(record.pathname) record.lineno = json_record["source"]["line"] record.created = datetime.datetime.fromisoformat( json_record["time"]).timestamp() extras = { k: v for k, v in json_record.items() if k not in PrismRunnerLogFilter.COMMON_FIELDS } if json_record["msg"] == "log from SDK worker": # TODO: Use location and time inside the nested message to set record record.name = "SdkWorker" + "@" + json_record["worker"]["ID"] record.msg = json_record["sdk"]["msg"] else: record.name = "PrismRunner" record.msg = ( f"{json_record['msg']} " f"({', '.join(f'{k}={v!r}' for k, v in extras.items())})") except (json.JSONDecodeError, KeyError, ValueError, TypeError, AttributeError): # The log parsing/filtering is best-effort. pass return True # Always return True to allow the record to pass. class PrismJobServer(job_server.SubprocessJobServer): BIN_CACHE = os.path.expanduser("~/.apache_beam/cache/prism/bin") def __init__(self, options): super().__init__() prism_options = options.view_as(pipeline_options.PrismRunnerOptions) # Options flow: # If the path is set, always download and unzip the provided path, # even if a binary is cached. self._path = prism_options.prism_location # Which version to use when constructing the prism download url. if prism_options.prism_beam_version_override: self._version = prism_options.prism_beam_version_override else: self._version = 'v' + beam_version job_options = options.view_as(pipeline_options.JobServerOptions) self._job_port = job_options.job_port self._log_level = prism_options.prism_log_level self._log_kind = prism_options.prism_log_kind # override console to json with log filter enabled if self._log_kind == "console": self._log_kind = "json" self._log_filter = PrismRunnerLogFilter() # the method is only kept for testing and backward compatibility @classmethod def local_bin( cls, url: str, bin_cache: str = '', ignore_cache: bool = False) -> str: url, ignore_cache = cls._download_to_local_path(url, bin_cache, ignore_cache) return cls._prepare_executable(url, bin_cache, ignore_cache) # the method is only kept for testing and backward compatibility def path_to_binary(self) -> str: return self._resolve_source_path() # the method is only kept for testing and backward compatibility def construct_download_url(self, root_tag: str, sys: str, mach: str) -> str: return self._construct_download_url(self._version, root_tag, sys, mach) @staticmethod def _prepare_executable( url: str, bin_cache: str, ignore_cache: bool = True) -> str: """ Given a path to a local artifact (zip or binary), makes it an executable binary file. Returns the path to the final, executable binary. """ assert (os.path.isfile(url)) if zipfile.is_zipfile(url): target = os.path.splitext(os.path.basename(url))[0] target_url = os.path.join(bin_cache, target) if not ignore_cache and os.path.exists(target_url): _LOGGER.info( 'Using cached prism binary from %s for %s' % (target_url, url)) else: # Only unzip the zip file if the url is a zip file and ignore_cache is # True (cache disabled) _LOGGER.info("Unzipping prism from %s to %s" % (url, target_url)) z = zipfile.ZipFile(url) bin_cache_tmp = os.path.join(bin_cache, 'tmp') if not os.path.exists(bin_cache_tmp): os.makedirs(bin_cache_tmp) target_tmp_url = z.extract(target, path=bin_cache_tmp) _rename_if_different(target_tmp_url, target_url) else: target_url = url _LOGGER.info("Prism binary path resolved to: %s", target_url) # Make sure the binary is executable. try: st = os.stat(target_url) os.chmod(target_url, st.st_mode | stat.S_IEXEC) except PermissionError: _LOGGER.warning( 'Could not change permissions of prism binary; invoking may fail if ' + 'current process does not have exec permissions on binary.') return target_url @staticmethod def _download_to_local_path( url: str, bin_cache: str = '', ignore_cache: bool = False) -> tuple[str, bool]: """ Ensures the artifact is on local disk, downloading it if necessary. Returns the path to the local (potentially cached) artifact. """ # ignore_cache sets whether we should always be downloading and unzipping # the file or not, to avoid staleness issues. if bin_cache == '': bin_cache = PrismJobServer.BIN_CACHE if os.path.exists(url): _LOGGER.info('Using local prism binary/zip from %s' % url) cached_file = url else: cached_file = os.path.join(bin_cache, os.path.basename(url)) if os.path.exists(cached_file) and not ignore_cache: _LOGGER.info( 'Using cached prism binary/zip from %s for %s' % (cached_file, url)) else: _LOGGER.info('Downloading prism from %s' % url) if not os.path.exists(bin_cache): os.makedirs(bin_cache) try: try: url_read = FileSystems.open(url) except ValueError: url_read = urlopen(url) with open(cached_file + '.tmp', 'wb') as zip_write: shutil.copyfileobj(url_read, zip_write, length=1 << 20) _rename_if_different(cached_file + '.tmp', cached_file) except URLError as e: raise RuntimeError( 'Unable to fetch remote prism binary at %s: %s' % (url, e)) # If we download a new prism, then we should always use it but not # the cached one. ignore_cache = True return cached_file, ignore_cache @staticmethod def _construct_download_url( version: str, root_tag: str, sys: str, mach: str) -> str: """Construct the prism download URL with the appropriate release tag. This maps operating systems and machine architectures to the compatible and canonical names used by the Go build targets. platform.system() provides compatible listings, so we need to filter out the unsupported versions.""" opsys = sys.lower() if opsys not in ['linux', 'windows', 'darwin']: raise ValueError( 'Operating System "%s" unsupported for constructing a Prism release ' 'binary URL.' % (opsys)) # platform.machine() will vary by system, but many names are compatible. arch = mach.lower() if arch in ['amd64', 'x86_64', 'x86-64', 'x64']: arch = 'amd64' if arch in ['arm64', 'aarch64_be', 'aarch64', 'armv8b', 'armv8l']: arch = 'arm64' if arch not in ['amd64', 'arm64']: raise ValueError( 'Machine architecture "%s" unsupported for constructing a Prism ' 'release binary URL.' % (opsys)) # Some special handling is needed when creating url for release candidates. # For example, v2.66.0rc2 should have the following url # https://github.com/apache/beam/releases/download/v2.66.0-RC2/apache_beam-v2.66.0-prism-xxx-yyy.zip if 'rc' in version: version = version.split('rc')[0] if 'rc' in root_tag: root_tag = '-RC'.join(root_tag.split('rc')) return ( GITHUB_DOWNLOAD_PREFIX + f"{root_tag}/apache_beam-{version}-prism-{opsys}-{arch}.zip") @staticmethod def _resolve_from_location_override(path, version) -> str: """Handles the case where --prism_location is explicitly set.""" # The path is overridden, check various cases. if os.path.exists(path): # The path is local and exists, use directly. return path try: if FileSystems.exists(path): # The path is in one of the supported filesystems. return path except ValueError: # If there is a value error raised by Filesystems, try to resolve # the path with the following steps. pass # Check if the path is a URL. url = urllib.parse.urlparse(path) if not url.scheme: raise ValueError( 'Unable to parse binary URL "%s". If using a full URL, make ' 'sure the scheme is specified. If using a local file xpath, ' 'make sure the file exists; you may have to first build prism ' 'using `go build `.' % (path)) # We have a URL, see if we need to construct a valid file name. if path.startswith(GITHUB_DOWNLOAD_PREFIX): # If this URL starts with the download prefix, let it through. return path # The only other valid option is a github release page. if not path.startswith(GITHUB_TAG_PREFIX): raise ValueError( 'Provided --prism_location URL is not an Apache Beam Github ' 'Release page URL or download URL: %s' % (path)) # Get the root tag for this URL root_tag = os.path.basename(os.path.normpath(path)) return PrismJobServer._construct_download_url( version, root_tag, platform.system(), platform.machine()) @staticmethod def _install_from_source(version): """Builds and installs Prism from a Go source package. It first tries the local module, then falls back to @latest. """ # This is a development version! Assume Go is installed. # Set the install directory to the cache location. envdict = {**os.environ, "GOBIN": PrismJobServer.BIN_CACHE} PRISMPKG = "github.com/apache/beam/sdks/v2/go/cmd/prism" _LOGGER.info( 'Installing prism from local source into "%s".', PrismJobServer.BIN_CACHE) process = subprocess.run(["go", "install", PRISMPKG], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=envdict, check=False) if process.returncode == 0: # Successfully installed return '%s/prism' % (PrismJobServer.BIN_CACHE) # We failed to build for some reason. output = process.stdout.decode("utf-8") if ("not in a module" not in output) and ("no required module provides" not in output): # This branch handles two classes of failures: # 1. Go isn't installed, so it needs to be installed by the Beam SDK # developer. # 2. Go is installed, and they are building in a local version of Prism, # but there was a compile error that the developer should address. # Either way, the @latest fallback either would fail, or hide the error, # so fail now. _LOGGER.info(output) raise ValueError( 'Unable to install a local of Prism: "%s";\n' 'Likely Go is not installed, or a local change to Prism did not ' 'compile.\nPlease install Go (see https://go.dev/doc/install) to ' 'enable automatic local builds.\n' 'Alternatively provide a binary with the --prism_location flag.' '\nCaptured output:\n %s' % (version, output)) # Go is installed and claims we're not in a Go module that has access to # the Prism package. # Fallback to using the @latest version of prism, which works everywhere. _LOGGER.info( 'Installing prism from "%s@latest" into "%s".', PRISMPKG, PrismJobServer.BIN_CACHE) process = subprocess.run(["go", "install", PRISMPKG + "@latest"], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=envdict, check=False) if process.returncode == 0: return '%s/prism' % (PrismJobServer.BIN_CACHE) output = process.stdout.decode("utf-8") raise ValueError( 'We were unable to execute the subprocess "%s" to automatically ' 'build prism.\nAlternatively provide an alternate binary with the ' '--prism_location flag.' '\nCaptured output:\n %s' % (process.args, output)) def _resolve_source_path(self) -> str: """Resolves and returns the source for the Prism binary. The resolution follows this order: 1. A user-provided location (local path, GCS, or URL). 2. A pre-built binary from GitHub for a release version. 3. Build from local Go source for a development version. """ if self._path: return self._resolve_from_location_override(self._path, self._version) if '.dev' not in self._version: # Not a development version, so construct the production download URL return self._construct_download_url( self._version, self._version, platform.system(), platform.machine()) return self._install_from_source(self._version) def _get_executable_path(self) -> str: """Orchestrates the process of getting a ready-to-use Prism binary.""" source = self._resolve_source_path() if source == "%s/prism" % (self.BIN_CACHE): # source is from go installation, so it is already a local binary return self._prepare_executable(source, self.BIN_CACHE, True) # Always re-download/extract if a custom path was provided to avoid # staleness ignore_cache = self._path is not None local_path, ignore_cache = self._download_to_local_path(source, self.BIN_CACHE, ignore_cache) return self._prepare_executable(local_path, self.BIN_CACHE, ignore_cache) def subprocess_cmd_and_endpoint(self) -> tuple[list[typing.Any], str]: bin_path = self._get_executable_path() job_port, = subprocess_server.pick_port(self._job_port) subprocess_cmd = [bin_path] + self.prism_arguments(job_port) return (subprocess_cmd, f"localhost:{job_port}") def prism_arguments(self, job_port) -> list[typing.Any]: return [ '--job_port', job_port, '--log_level', self._log_level, '--log_kind', self._log_kind, # Go does not support "-flag x" format for boolean flags. '--serve_http=false', ] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/prism_runner_test.py0000644000175100017510000004644215206113032026061 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import argparse import logging import os.path import queue import shlex import threading import time import typing import unittest import zipfile from os import linesep from os.path import exists from shutil import rmtree from tempfile import mkdtemp from unittest import mock import pytest from parameterized import parameterized import apache_beam as beam from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import PortableOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import TypeOptions from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.runners.portability import portable_runner_test from apache_beam.runners.portability import prism_runner from apache_beam.runners.worker import worker_pool_main from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.utils import shared # Run as # # pytest prism_runner_test.py[::TestClass::test_case] \ # --test-pipeline-options="--environment_type=LOOPBACK" _LOGGER = logging.getLogger(__name__) Row = typing.NamedTuple("Row", [("col1", int), ("col2", str)]) beam.coders.registry.register_coder(Row, beam.coders.RowCoder) class PrismRunnerTest(portable_runner_test.PortableRunnerTest): _use_grpc = True _use_subprocesses = True conf_dir = None expansion_port = None def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.environment_type = None self.environment_config = None self.enable_commit = False self.streaming = False self.allow_unsafe_triggers = False def setUp(self): self.enable_commit = False @pytest.fixture(autouse=True) def parse_options(self, request): if not request.config.option.test_pipeline_options: raise unittest.SkipTest( 'Skipping because --test-pipeline-options is not specified.') test_pipeline_options = request.config.option.test_pipeline_options parser = argparse.ArgumentParser(add_help=True) parser.add_argument( '--prism_bin', help='Prism binary to submit jobs.', action='store') parser.add_argument( '--environment_type', default='LOOPBACK', choices=['DOCKER', 'PROCESS', 'LOOPBACK'], help='Set the environment type for running user code. DOCKER runs ' 'user code in a container. PROCESS runs user code in ' 'automatically started processes. LOOPBACK runs user code on ' 'the same process that originally submitted the job.') parser.add_argument( '--environment_option', '--environment_options', dest='environment_options', action='append', default=None, help=( 'Environment configuration for running the user code. ' 'Recognized options depend on --environment_type.\n ' 'For DOCKER: docker_container_image (optional)\n ' 'For PROCESS: process_command (required), process_variables ' '(optional, comma-separated)\n ' 'For EXTERNAL: external_service_address (required)')) known_args, unknown_args = parser.parse_known_args( shlex.split(test_pipeline_options)) if unknown_args: _LOGGER.warning('Discarding unrecognized arguments %s' % unknown_args) self.set_prism_bin(known_args.prism_bin) self.environment_type = known_args.environment_type self.environment_options = known_args.environment_options @classmethod def tearDownClass(cls): if cls.conf_dir and exists(cls.conf_dir): _LOGGER.info("removing conf dir: %s" % cls.conf_dir) rmtree(cls.conf_dir) super().tearDownClass() @classmethod def _create_conf_dir(cls): """Create (and save a static reference to) a "conf dir", used to provide metrics configs and verify metrics output It gets cleaned up when the suite is done executing""" if hasattr(cls, 'conf_dir'): cls.conf_dir = mkdtemp(prefix='prismtest-conf') # path for a FileReporter to write metrics to cls.test_metrics_path = os.path.join(cls.conf_dir, 'test-metrics.txt') # path to write Prism configuration to conf_path = os.path.join(cls.conf_dir, 'prism-conf.yaml') file_reporter = 'org.apache.beam.runners.prism.metrics.FileReporter' with open(conf_path, 'w') as f: f.write( linesep.join([ 'metrics.reporters: file', 'metrics.reporter.file.class: %s' % file_reporter, 'metrics.reporter.file.path: %s' % cls.test_metrics_path, 'metrics.scope.operator: ', ])) @classmethod def _subprocess_command(cls, job_port, expansion_port): # will be cleaned up at the end of this method, and recreated and used by # the job server tmp_dir = mkdtemp(prefix='prismtest') cls._create_conf_dir() cls.expansion_port = expansion_port try: return [ cls.prism_bin, '--job_port', str(job_port), ] finally: rmtree(tmp_dir) @classmethod def get_expansion_service(cls): # TODO Move expansion address resides into PipelineOptions return 'localhost:%s' % cls.expansion_port @classmethod def set_prism_bin(cls, prism_bin): cls.prism_bin = prism_bin def create_options(self): options = super().create_options() options.view_as(DebugOptions).experiments = ['beam_fn_api'] options.view_as(DebugOptions).experiments = [ 'pre_optimize=default' ] + options.view_as(DebugOptions).experiments options.view_as(PortableOptions).environment_type = self.environment_type options.view_as( PortableOptions).environment_options = self.environment_options options.view_as(StandardOptions).streaming = self.streaming options.view_as( TypeOptions).allow_unsafe_triggers = self.allow_unsafe_triggers return options # Can't read host files from within docker, read a "local" file there. def test_read(self): print('name:', __name__) with self.create_pipeline() as p: lines = p | beam.io.ReadFromText('/etc/profile') assert_that(lines, lambda lines: len(lines) > 0) def test_create_transform(self): with self.create_pipeline() as p: out = (p | beam.Create([1])) assert_that(out, equal_to([1])) with self.create_pipeline() as p: out = (p | beam.Create([1, 2], reshuffle=False)) assert_that(out, equal_to([1, 2])) with self.create_pipeline() as p: out = (p | beam.Create([1, 2], reshuffle=True)) assert_that(out, equal_to([1, 2])) def test_external_transform(self): raise unittest.SkipTest("Requires an expansion service to execute.") def test_expand_kafka_read(self): raise unittest.SkipTest("Requires an expansion service to execute.") def test_expand_kafka_write(self): raise unittest.SkipTest("Requires an expansion service to execute.") def test_sql(self): raise unittest.SkipTest("Requires an expansion service to execute.") # The following tests require additional implementation in Prism. def test_custom_merging_window(self): raise unittest.SkipTest( "Requires Prism to support Custom Window " + "Coders, and Merging Custom Windows. " + "https://github.com/apache/beam/issues/31921") def test_custom_window_type(self): raise unittest.SkipTest( "Requires Prism to support Custom Window Coders." + " https://github.com/apache/beam/issues/31921") def test_metrics(self): super().test_metrics(check_bounded_trie=False) def construct_timestamped(k, t): return window.TimestampedValue((k, t), t) def format_result(k, vs): return ('%s-%s' % (k, len(list(vs))), set(vs)) def test_after_count_trigger_batch(self): self.allow_unsafe_triggers = True with self.create_pipeline() as p: result = ( p | beam.Create([1, 2, 3, 4, 5, 10, 11]) | beam.FlatMap(lambda t: [('A', t), ('B', t + 5)]) #A1, A2, A3, A4, A5, A10, A11, B6, B7, B8, B9, B10, B15, B16 | beam.MapTuple(PrismRunnerTest.construct_timestamped) | beam.WindowInto( window.FixedWindows(10), trigger=trigger.AfterCount(3), accumulation_mode=trigger.AccumulationMode.DISCARDING, ) | beam.GroupByKey() | beam.MapTuple(PrismRunnerTest.format_result)) assert_that( result, equal_to( list([ ('A-5', {1, 2, 3, 4, 5}), ('A-2', {10, 11}), ('B-4', {6, 7, 8, 9}), ('B-3', {10, 15, 16}), ]))) def test_after_count_trigger_streaming(self): self.allow_unsafe_triggers = True self.streaming = True with self.create_pipeline() as p: result = ( p | beam.Create([1, 2, 3, 4, 5, 10, 11]) | beam.FlatMap(lambda t: [('A', t), ('B', t + 5)]) #A1, A2, A3, A4, A5, A10, A11, B6, B7, B8, B9, B10, B15, B16 | beam.MapTuple(PrismRunnerTest.construct_timestamped) | beam.WindowInto( window.FixedWindows(10), trigger=trigger.AfterCount(3), accumulation_mode=trigger.AccumulationMode.DISCARDING, ) | beam.GroupByKey() | beam.MapTuple(PrismRunnerTest.format_result)) assert_that( result, equal_to( list([ ('A-3', {1, 2, 3}), ('A-2', {4, 5}), ('A-2', {10, 11}), ('B-3', {6, 7, 8}), ('B-1', {9}), ('B-3', {10, 15, 16}), ]))) class PrismJobServerTest(unittest.TestCase): def setUp(self) -> None: self.local_dir = mkdtemp() self.cache_dir = os.path.join(self.local_dir, "cache") os.mkdir(self.cache_dir) self.job_server = prism_runner.PrismJobServer(options=PortableOptions()) self.local_bin_path = os.path.join(self.local_dir, "my_prism_bin") self.local_zip_path = os.path.join(self.local_dir, "my_prism_bin.zip") self.cache_bin_path = os.path.join(self.cache_dir, "my_prism_bin") self.cache_zip_path = os.path.join(self.cache_dir, "my_prism_bin.zip") self.remote_zip_path = "https://github.com/apache/beam/releases/download/fake_ver/my_prism_bin.zip" # pylint: disable=line-too-long def tearDown(self) -> None: rmtree(self.local_dir) pass def _make_local_bin(self, fn=None): fn = fn or self.local_bin_path with open(fn, 'wb'): pass def _make_local_zip(self, fn=None): fn = fn or self.local_zip_path with zipfile.ZipFile(fn, 'w', zipfile.ZIP_DEFLATED): pass def _make_cache_bin(self, fn=None): fn = fn or self.cache_bin_path with open(fn, 'wb'): pass def _make_cache_zip(self, fn=None): fn = fn or self.cache_zip_path with zipfile.ZipFile(fn, 'w', zipfile.ZIP_DEFLATED): pass def _extract_side_effect(self, fn, path=None): if path is None: return fn full_path = os.path.join(str(path), fn) if path.startswith(self.cache_dir): self._make_cache_bin(full_path) else: self._make_local_bin(full_path) return full_path @parameterized.expand([[True, True], [True, False], [False, True], [False, False]]) def test_with_unknown_path(self, custom_bin_cache, ignore_cache): self.assertRaises( FileNotFoundError, lambda: self.job_server.local_bin( "/path/unknown", bin_cache=self.cache_dir if custom_bin_cache else '', ignore_cache=ignore_cache)) @parameterized.expand([ [True, True, True], [True, True, False], [True, False, True], [True, False, False], [False, True, True], [False, True, False], [False, False, True], [False, False, False], ]) def test_with_local_binary_and_zip( self, has_cache_bin, has_cache_zip, ignore_cache): self._make_local_bin() self._make_local_zip() if has_cache_bin: self._make_cache_bin() if has_cache_zip: self._make_cache_zip() with mock.patch('zipfile.is_zipfile') as mock_is_zipfile: with mock.patch('zipfile.ZipFile') as mock_zipfile_init: mock_zipfile = mock.MagicMock() mock_zipfile.extract = mock.Mock(side_effect=self._extract_side_effect) mock_zipfile_init.return_value = mock_zipfile # path is set to local binary # always use local binary even if we have a cached copy, no unzipping mock_is_zipfile.return_value = False self.assertEqual( self.job_server.local_bin( self.local_bin_path, self.cache_dir, ignore_cache), self.local_bin_path) mock_zipfile_init.assert_not_called() # path is set to local zip # use local zip and unzip only if cache binary not available or # ignore_cache is true mock_is_zipfile.return_value = True self.assertEqual( self.job_server.local_bin( self.local_zip_path, self.cache_dir, ignore_cache), self.cache_bin_path) if has_cache_bin and not ignore_cache: # if cache is enabled and binary is in cache, we wont't unzip mock_zipfile_init.assert_not_called() else: mock_zipfile_init.assert_called_once() mock_zipfile_init.reset_mock() @parameterized.expand([ [True, True, True], [True, True, False], [True, False, True], [True, False, False], [False, True, True], [False, True, False], [False, False, True], [False, False, False], ]) def test_with_remote_path(self, has_cache_bin, has_cache_zip, ignore_cache): if has_cache_bin: self._make_cache_bin() if has_cache_zip: self._make_cache_zip() with mock.patch( 'apache_beam.runners.portability.prism_runner.urlopen') as mock_urlopen: mock_response = mock.MagicMock() if has_cache_zip: with open(self.cache_zip_path, 'rb') as f: mock_response.read.side_effect = [f.read(), b''] else: mock_response.read.return_value = b'' mock_urlopen.return_value = mock_response with mock.patch('zipfile.is_zipfile') as mock_is_zipfile: with mock.patch('zipfile.ZipFile') as mock_zipfile_init: mock_zipfile = mock.MagicMock() mock_zipfile.extract = mock.Mock( side_effect=self._extract_side_effect) mock_zipfile_init.return_value = mock_zipfile mock_is_zipfile.return_value = True self.assertEqual( self.job_server.local_bin( self.remote_zip_path, self.cache_dir, ignore_cache=ignore_cache), self.cache_bin_path) if has_cache_zip and not ignore_cache: # if cache is enabled and zip is in cache, we wont't download mock_urlopen.assert_not_called() else: mock_urlopen.assert_called_once() if has_cache_bin and has_cache_zip and not ignore_cache: # if cache is enabled and both binary and zip are in cache, we # wont't unzip mock_zipfile_init.assert_not_called() else: mock_zipfile_init.assert_called_once() class PrismRunnerSingletonTest(unittest.TestCase): @parameterized.expand([True, False]) def test_singleton(self, enable_singleton): if enable_singleton: options = DebugOptions() # prism singleton is enabled by default else: options = DebugOptions(["--experiment=disable_prism_server_singleton"]) runner = prism_runner.PrismRunner() with mock.patch( 'apache_beam.runners.portability.prism_runner.PrismJobServer' ) as mock_prism_server: # Reset the class-level singleton for every fresh run prism_runner.PrismRunner.shared_handle = shared.Shared() runner = prism_runner.PrismRunner() runner.default_job_server(options) mock_prism_server.assert_called_once() mock_prism_server.reset_mock() runner = prism_runner.PrismRunner() runner.default_job_server(options) if enable_singleton: # If singleton is enabled, we won't try to create a new server for the # second run. mock_prism_server.assert_not_called() else: mock_prism_server.assert_called_once() def test_loopback_worker_daemon_thread_accumulation(self): """Verifies that in LOOPBACK mode, the external worker pool servicer properly tracks active thread-based SdkHarness workers and cleanly shuts them down in StopWorker via sentinel messages. This prevents background daemon threads from accumulating across sequential pipeline executions and leaking resources. """ servicer = worker_pool_main.BeamFnExternalWorkerPoolServicer( use_process=False, state_cache_size=0, data_buffer_time_limit_ms=0) active_workers = [] mock_responses = queue.Queue() def mock_run(self_worker): active_workers.append(self_worker) mock_responses.get() active_workers.remove(self_worker) def wait_for_workers(expected_count, timeout=5.0): start = time.time() while time.time() - start < timeout: if len(active_workers) == expected_count: return time.sleep(0.01) self.assertEqual(len(active_workers), expected_count) with mock.patch( 'apache_beam.runners.worker.sdk_worker.SdkHarness') as mock_harness: mock_harness.return_value._responses = mock_responses mock_harness.return_value.run = lambda: mock_run(mock_harness) # Simulate starting Worker 1 for Pipeline 1 req1 = beam_fn_api_pb2.StartWorkerRequest(worker_id="worker_1") req1.control_endpoint.url = "localhost:12345" servicer.StartWorker(req1, None) wait_for_workers(1) # Simulate stopping Worker 1 at the end of Pipeline 1 stop_req1 = beam_fn_api_pb2.StopWorkerRequest(worker_id="worker_1") servicer.StopWorker(stop_req1, None) # Verify the fix: StopWorker successfully tells the thread harness to shut down, # completely resolving the thread leak! wait_for_workers(0) # Simulate starting Worker 2 for Pipeline 2 req2 = beam_fn_api_pb2.StartWorkerRequest(worker_id="worker_2") req2.control_endpoint.url = "localhost:12345" servicer.StartWorker(req2, None) wait_for_workers(1) # Clean up the second worker servicer.StopWorker( beam_fn_api_pb2.StopWorkerRequest(worker_id="worker_2"), None) wait_for_workers(0) if __name__ == '__main__': # Run the tests. logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/requirements_cache_it_test.py0000644000175100017510000000557715206113032027704 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A pipeline to verify the installation of packages specified in the requirements.txt. A requirements text is created during runtime with package specified in _PACKAGE_IN_REQUIREMENTS_FILE. """ import argparse import logging import os import shutil import tempfile from importlib.metadata import PackageNotFoundError from importlib.metadata import distribution import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions _PACKAGE_IN_REQUIREMENTS_FILE = ['matplotlib', 'seaborn'] def verify_packages_from_requirements_file_are_installed(unused_element): _PACKAGE_NOT_IN_REQUIREMENTS_FILE = ['torch'] packages_to_test = _PACKAGE_IN_REQUIREMENTS_FILE + ( _PACKAGE_NOT_IN_REQUIREMENTS_FILE) for package_name in packages_to_test: try: output = distribution(package_name) except PackageNotFoundError as e: # pylint: disable=unused-variable output = None if package_name in _PACKAGE_IN_REQUIREMENTS_FILE: assert output is not None, ('Please check if package %s is specified' ' in requirements file' % package_name) if package_name in _PACKAGE_NOT_IN_REQUIREMENTS_FILE: assert output is None def run(argv=None, save_main_session=True): parser = argparse.ArgumentParser() _, pipeline_args = parser.parse_known_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session temp_dir = tempfile.mkdtemp() requirements_text_path = os.path.join(temp_dir, 'requirements.txt') with open(requirements_text_path, 'w') as f: f.write('\n'.join(_PACKAGE_IN_REQUIREMENTS_FILE)) pipeline_options.view_as( SetupOptions).requirements_file = requirements_text_path with beam.Pipeline(options=pipeline_options) as p: ( # pylint: disable=expression-not-assigned p | beam.Create([None]) | beam.Map(verify_packages_from_requirements_file_are_installed)) shutil.rmtree(temp_dir) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/sdk_container_builder.py0000644000175100017510000003523315206113032026624 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """SdkContainerImageBuilder builds the portable SDK container with dependencies. It copies the right boot dependencies, namely: apache beam sdk, python packages from requirements.txt, python packages from extra_packages.txt, workflow tarball, into the latest public python sdk container image, and run the dependencies installation in advance with the boot program in setup only mode to build the new image. """ import json import logging import os import posixpath import shutil import subprocess import sys import tarfile import tempfile import time import uuid from google.protobuf.json_format import MessageToJson from apache_beam import version as beam_version from apache_beam.internal.gcp.auth import get_service_credentials from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions # pylint: disable=unused-import from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import WorkerOptions from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.portability.stager import Stager from apache_beam.utils import plugin ARTIFACTS_CONTAINER_DIR = '/opt/apache/beam/artifacts' ARTIFACTS_MANIFEST_FILE = 'artifacts_info.json' SDK_CONTAINER_ENTRYPOINT = '/opt/apache/beam/boot' DOCKERFILE_TEMPLATE = ( """FROM {base_image} RUN mkdir -p {workdir} COPY ./* {workdir}/ RUN {entrypoint} --setup_only --artifacts {workdir}/{manifest_file} """) SOURCE_FOLDER = 'source' _LOGGER = logging.getLogger(__name__) class SdkContainerImageBuilder(plugin.BeamPlugin): def __init__(self, options): self._options = options self._docker_registry_push_url = self._options.view_as( SetupOptions).docker_registry_push_url version = ( beam_version.__version__ if 'dev' not in beam_version.__version__ else 'latest') self._base_image = ( self._options.view_as(WorkerOptions).sdk_container_image or 'apache/beam_python%s.%s_sdk:%s' % (sys.version_info[0], sys.version_info[1], version)) self._temp_src_dir = None def _build(self): container_image_tag = str(uuid.uuid4()) container_image_name = posixpath.join( self._docker_registry_push_url or '', 'beam_python_prebuilt_sdk:%s' % container_image_tag) with tempfile.TemporaryDirectory() as temp_folder: self._temp_src_dir = temp_folder self._prepare_dependencies() self._invoke_docker_build_and_push(container_image_name) return container_image_name def _prepare_dependencies(self): with tempfile.TemporaryDirectory() as tmp: artifacts = Stager.create_job_resources( self._options, tmp, log_submission_env_dependencies=False) resources = Stager.extract_staging_tuple_iter(artifacts) # make a copy of the staged artifacts into the temp source folder. file_names = [] for path, name, _ in resources: shutil.copyfile(path, os.path.join(self._temp_src_dir, name)) file_names.append(name) with open(os.path.join(self._temp_src_dir, 'Dockerfile'), 'w') as file: file.write( DOCKERFILE_TEMPLATE.format( base_image=self._base_image, workdir=ARTIFACTS_CONTAINER_DIR, manifest_file=ARTIFACTS_MANIFEST_FILE, entrypoint=SDK_CONTAINER_ENTRYPOINT)) self._generate_artifacts_manifests_json_file( file_names, self._temp_src_dir) def _invoke_docker_build_and_push(self, container_image_name): raise NotImplementedError @classmethod def _builder_key(cls) -> str: return f'{cls.__module__}.{cls.__name__}' @staticmethod def _generate_artifacts_manifests_json_file(file_names, temp_dir): infos = [] for name in file_names: info = beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.StandardArtifacts.Types.FILE.urn, type_payload=beam_runner_api_pb2.ArtifactFilePayload( path=name).SerializeToString(), ) infos.append(json.dumps(MessageToJson(info))) with open(os.path.join(temp_dir, ARTIFACTS_MANIFEST_FILE), 'w') as file: file.write('[\n' + ',\n'.join(infos) + '\n]') @classmethod def build_container_image(cls, pipeline_options: PipelineOptions) -> str: setup_options = pipeline_options.view_as(SetupOptions) container_build_engine = setup_options.prebuild_sdk_container_engine builder_cls = cls._get_subclass_by_key(container_build_engine) builder = builder_cls(pipeline_options) return builder._build() @classmethod def _get_subclass_by_key(cls, key: str) -> type['SdkContainerImageBuilder']: available_builders = [ subclass for subclass in cls.get_all_subclasses() if subclass._builder_key() == key ] if not available_builders: available_builder_keys = [ subclass._builder_key() for subclass in cls.get_all_subclasses() ] raise ValueError( f'Cannot find SDK builder type {key} in ' f'{available_builder_keys}') elif len(available_builders) > 1: raise ValueError(f'Found multiple builders under key {key}') return available_builders[0] class _SdkContainerImageLocalBuilder(SdkContainerImageBuilder): """SdkContainerLocalBuilder builds the sdk container image with local docker.""" @classmethod def _builder_key(cls): return 'local_docker' def _invoke_docker_build_and_push(self, container_image_name): try: _LOGGER.info("Building sdk container, this may take a few minutes...") now = time.time() subprocess.run(['docker', 'build', '.', '-t', container_image_name], check=True, cwd=self._temp_src_dir) except subprocess.CalledProcessError as err: raise RuntimeError( 'Failed to build sdk container with local docker, ' 'stderr:\n %s.' % err.stderr) else: _LOGGER.info( "Successfully built %s in %.2f seconds" % (container_image_name, time.time() - now)) if self._docker_registry_push_url: _LOGGER.info("Pushing prebuilt sdk container...") try: subprocess.run(['docker', 'push', container_image_name], check=True) except subprocess.CalledProcessError as err: raise RuntimeError( 'Failed to push prebuilt sdk container %s, stderr: \n%s' % (container_image_name, err.stderr)) _LOGGER.info( "Successfully pushed %s in %.2f seconds" % (container_image_name, time.time() - now)) else: _LOGGER.info( "no --docker_registry_push_url option is specified in pipeline " "options, specify it if the new image is intended to be " "pushed to a registry.") class _SdkContainerImageCloudBuilder(SdkContainerImageBuilder): """SdkContainerLocalBuilder builds the sdk container image with google cloud build.""" def __init__(self, options): super().__init__(options) self._google_cloud_options = options.view_as(GoogleCloudOptions) self._cloud_build_machine_type = self._get_cloud_build_machine_type_enum( options.view_as(SetupOptions).cloud_build_machine_type) if self._google_cloud_options.no_auth: credentials = None else: credentials = get_service_credentials(options) from apache_beam.io.gcp.gcsio import create_storage_client self._storage_client = create_storage_client( options, not self._google_cloud_options.no_auth) from google.cloud.devtools.cloudbuild_v1.services import cloud_build self._cloudbuild_client = cloud_build.CloudBuildClient( credentials=credentials) if not self._docker_registry_push_url: self._docker_registry_push_url = ( 'gcr.io/%s/prebuilt_beam_sdk' % self._google_cloud_options.project) @classmethod def _builder_key(cls): return 'cloud_build' def _invoke_docker_build_and_push(self, container_image_name): from google.cloud.devtools.cloudbuild_v1 import types as cloud_build_types project_id = self._google_cloud_options.project temp_location = self._google_cloud_options.temp_location # google cloud build service expects all the build source file to be # compressed into a tarball. tarball_path = os.path.join(self._temp_src_dir, '%s.tgz' % SOURCE_FOLDER) self._make_tarfile(tarball_path, self._temp_src_dir) _LOGGER.info( "Compressed source files for building sdk container at %s" % tarball_path) container_image_tag = container_image_name.split(':')[-1] gcs_location = os.path.join( temp_location, '%s-%s.tgz' % (SOURCE_FOLDER, container_image_tag)) self._upload_to_gcs(tarball_path, gcs_location) build = cloud_build_types.Build() if self._cloud_build_machine_type: build.options = cloud_build_types.BuildOptions() build.options.machineType = self._cloud_build_machine_type build.steps = [] step = cloud_build_types.BuildStep() step.name = 'quay.io/buildah/stable:latest' step.entrypoint = 'sh' step.args = [ '-c', # The --storage-driver=vfs option is used to run buildah in a # rootless environment. ( 'buildah bud --storage-driver=vfs -t {0} . && ' 'buildah push --storage-driver=vfs {0} docker://{0}' ).format(container_image_name), ] step.dir = SOURCE_FOLDER build.steps.append(step) source = cloud_build_types.Source() storage_source = cloud_build_types.StorageSource() gcs_bucket, gcs_object = self._get_gcs_bucket_and_name(gcs_location) storage_source.bucket = os.path.join(gcs_bucket) storage_source.object_ = gcs_object source.storage_source = storage_source build.source = source # TODO(zyichi): make timeout configurable build.timeout = '7200s' now = time.time() # operation = client.create_build(project_id=project_id, build=build) request = cloud_build_types.CreateBuildRequest( project_id=project_id, build=build) build = self._cloudbuild_client.create_build(request) build_id, log_url = self._get_cloud_build_id_and_log_url(build.metadata) _LOGGER.info( 'Building sdk container with Google Cloud Build, this may ' 'take a few minutes, you may check build log at %s' % log_url) # block until build finish, if build fails exception will be raised and # stops the job submission. response = self._cloudbuild_client.get_build( request=cloud_build_types.GetBuildRequest( id=build_id, project_id=project_id)) while response.status in [cloud_build_types.Build.Status.QUEUED, cloud_build_types.Build.Status.PENDING, cloud_build_types.Build.Status.WORKING]: time.sleep(10) response = self._cloudbuild_client.get_build( cloud_build_types.GetBuildRequest(id=build_id, project_id=project_id)) if response.status != cloud_build_types.Build.Status.SUCCESS: raise RuntimeError( 'Failed to build python sdk container image on google cloud build, ' 'please check build log for error.') _LOGGER.info( "Python SDK container pre-build finished in %.2f seconds" % (time.time() - now)) _LOGGER.info( "Python SDK container built and pushed as %s." % container_image_name) def _upload_to_gcs(self, local_file_path, gcs_location): bucket_name, blob_name = self._get_gcs_bucket_and_name(gcs_location) _LOGGER.info('Starting GCS upload to %s...', gcs_location) from google.cloud import storage from google.cloud.exceptions import Forbidden from google.cloud.exceptions import NotFound try: bucket = self._storage_client.get_bucket(bucket_name) blob = bucket.get_blob(blob_name) if not blob: blob = storage.Blob(name=blob_name, bucket=bucket) blob.upload_from_filename(local_file_path) except Exception as e: if isinstance(e, (Forbidden, NotFound)): raise IOError(( 'Could not upload to GCS path %s: %s. Please verify ' 'that credentials are valid and that you have write ' 'access to the specified path.') % (gcs_location, e.message)) raise _LOGGER.info('Completed GCS upload to %s.', gcs_location) def _get_cloud_build_id_and_log_url(self, metadata): # google-cloud-build 3.35+ if getattr(metadata, 'build', None): build = metadata.build return (build.id, build.log_url) # Fallback for older clients that use additionalProperties. id = None log_url = None additional_props = getattr(metadata, 'additionalProperties', None) if additional_props: for item in additional_props: if item.key == 'build': for field in item.value.object_value.properties: if field.key == 'logUrl': log_url = field.value.string_value if field.key == 'id': id = field.value.string_value return id, log_url @staticmethod def _get_gcs_bucket_and_name(gcs_location): return gcs_location[5:].split('/', 1) @staticmethod def _make_tarfile(output_filename, source_dir): with tarfile.open(output_filename, "w:gz") as tar: tar.add(source_dir, arcname=SOURCE_FOLDER) @staticmethod def _get_cloud_build_machine_type_enum(machine_type: str): from google.cloud.devtools.cloudbuild_v1 import types as cloud_build_types if not machine_type: return None mappings = { 'n1-highcpu-8': cloud_build_types.BuildOptions.MachineType.N1_HIGHCPU_8, 'n1-highcpu-32': cloud_build_types.BuildOptions.MachineType. N1_HIGHCPU_32, 'e2-highcpu-8': cloud_build_types.BuildOptions.MachineType.E2_HIGHCPU_8, 'e2-highcpu-32': cloud_build_types.BuildOptions.MachineType. E2_HIGHCPU_32 } if machine_type.lower() in mappings: return mappings[machine_type.lower()] else: raise ValueError( 'Unknown Cloud Build Machine Type option, please specify one of ' '[n1-highcpu-8, n1-highcpu-32, e2-highcpu-8, e2-highcpu-32].') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/sdk_container_builder_test.py0000644000175100017510000001132015206113032027652 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the sdk container builder module.""" # pytype: skip-file import gc import logging import unittest import unittest.mock from apache_beam.options import pipeline_options from apache_beam.runners.portability import sdk_container_builder class SdkContainerBuilderTest(unittest.TestCase): def tearDown(self): # Ensures SdkContainerImageBuilder subclasses are cleared gc.collect() def test_can_find_local_builder(self): local_builder = ( sdk_container_builder.SdkContainerImageBuilder._get_subclass_by_key( 'local_docker')) self.assertEqual( local_builder, sdk_container_builder._SdkContainerImageLocalBuilder) def test_can_find_cloud_builder(self): local_builder = ( sdk_container_builder.SdkContainerImageBuilder._get_subclass_by_key( 'cloud_build')) self.assertEqual( local_builder, sdk_container_builder._SdkContainerImageCloudBuilder) def test_missing_builder_key_throws_value_error(self): with self.assertRaises(ValueError): sdk_container_builder.SdkContainerImageBuilder._get_subclass_by_key( 'missing-id') def test_multiple_matchings_keys_throws_value_error(self): # pylint: disable=unused-variable class _PluginSdkBuilder(sdk_container_builder.SdkContainerImageBuilder): @classmethod def _builder_key(cls): return 'test-id' class _PluginSdkBuilder2(sdk_container_builder.SdkContainerImageBuilder): @classmethod def _builder_key(cls): return 'test-id' # pylint: enable=unused-variable with self.assertRaises(ValueError): sdk_container_builder.SdkContainerImageBuilder._get_subclass_by_key( 'test-id') def test_can_find_new_subclass(self): class _PluginSdkBuilder(sdk_container_builder.SdkContainerImageBuilder): pass expected_key = f'{_PluginSdkBuilder.__module__}._PluginSdkBuilder' local_builder = ( sdk_container_builder.SdkContainerImageBuilder._get_subclass_by_key( expected_key)) self.assertEqual(local_builder, _PluginSdkBuilder) @unittest.mock.patch( 'apache_beam.runners.portability.sdk_container_builder._SdkContainerImageLocalBuilder' # pylint: disable=line-too-long ) @unittest.mock.patch.object( sdk_container_builder.SdkContainerImageBuilder, "_get_subclass_by_key") def test_build_container_image_locates_subclass_invokes_build( self, mock_get_subclass, mocked_local_builder): mock_get_subclass.return_value = mocked_local_builder options = pipeline_options.PipelineOptions() sdk_container_builder.SdkContainerImageBuilder.build_container_image( options) mocked_local_builder.assert_called_once_with(options) mocked_local_builder.return_value._build.assert_called_once_with() def test_container_image_name_uses_forward_slashes(self): """Verify container image names use forward slashes as URI separators.""" options = pipeline_options.PipelineOptions([ '--docker_registry_push_url=europe-west1-docker.pkg.dev/project-id'\ '/repo-name', ]) builder = sdk_container_builder._SdkContainerImageLocalBuilder(options) # Mock the file and docker operations with unittest.mock.patch( 'apache_beam.runners.portability.sdk_container_builder.tempfile.' \ 'TemporaryDirectory' ): with unittest.mock.patch.object(builder, '_prepare_dependencies'): with unittest.mock.patch.object(builder, '_invoke_docker_build_and_push'): container_image_name = builder._build() expected_prefix = 'europe-west1-docker.pkg.dev/project-id/repo-name/' \ 'beam_python_prebuilt_sdk:' self.assertTrue( container_image_name.startswith(expected_prefix), f'Expected image name to start with {expected_prefix},'\ f' got: {container_image_name}' ) if __name__ == '__main__': # Run the tests. logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/spark_java_job_server_test.py0000644000175100017510000000520715206113032027671 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import unittest from apache_beam.options import pipeline_options from apache_beam.runners.portability.spark_runner import SparkRunner class SparkTestPipelineOptions(pipeline_options.PipelineOptions): def view_as(self, cls): # Ensure only SparkRunnerOptions and JobServerOptions are used when calling # default_job_server. If other options classes are needed, the cache key # must include them to prevent incorrect hits. assert ( cls is pipeline_options.SparkRunnerOptions or cls is pipeline_options.JobServerOptions) return super().view_as(cls) class SparkJavaJobServerTest(unittest.TestCase): def test_job_server_cache(self): # Multiple SparkRunner instances may be created, so we need to make sure we # cache job servers across runner instances. # Most pipeline-specific options, such as sdk_worker_parallelism, don't # affect job server configuration, so it is ok to ignore them for caching. job_server1 = SparkRunner().default_job_server( SparkTestPipelineOptions(['--sdk_worker_parallelism=1'])) job_server2 = SparkRunner().default_job_server( SparkTestPipelineOptions(['--sdk_worker_parallelism=2'])) self.assertIs(job_server2, job_server1) # JobServerOptions and SparkRunnerOptions do affect job server # configuration, so using different pipeline options gives us a different # job server. job_server3 = SparkRunner().default_job_server( SparkTestPipelineOptions(['--job_port=1234'])) self.assertIsNot(job_server3, job_server1) job_server4 = SparkRunner().default_job_server( SparkTestPipelineOptions(['--spark_master_url=spark://localhost:5678'])) self.assertIsNot(job_server4, job_server1) self.assertIsNot(job_server4, job_server3) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/spark_runner.py0000644000175100017510000001106715206113032025003 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A runner for executing portable pipelines on Spark.""" # pytype: skip-file import os import re import urllib from apache_beam.options import pipeline_options from apache_beam.runners.portability import job_server from apache_beam.runners.portability import portable_runner from apache_beam.runners.portability import spark_uber_jar_job_server # https://spark.apache.org/docs/latest/submitting-applications.html#master-urls LOCAL_MASTER_PATTERN = r'^local(\[.+\])?$' # Since Java job servers are heavyweight external processes, cache them. # This applies only to SparkJarJobServer, not SparkUberJarJobServer. JOB_SERVER_CACHE = {} class SparkRunner(portable_runner.PortableRunner): """A runner for launching jobs on Spark, automatically starting a local spark master if one is not given. """ # Inherits run_portable_pipeline from PortableRunner. def default_environment(self, options): spark_options = options.view_as(pipeline_options.SparkRunnerOptions) portable_options = options.view_as(pipeline_options.PortableOptions) if (re.match(LOCAL_MASTER_PATTERN, spark_options.spark_master_url) and not portable_options.environment_type and not portable_options.output_executable_path): portable_options.environment_type = 'LOOPBACK' return super().default_environment(options) def default_job_server(self, options): spark_options = options.view_as(pipeline_options.SparkRunnerOptions) if spark_options.spark_submit_uber_jar: if not spark_options.spark_rest_url: raise ValueError('Option spark_rest_url must be set.') return spark_uber_jar_job_server.SparkUberJarJobServer( spark_options.spark_rest_url, options) # Use Java job server by default. # Only SparkRunnerOptions and JobServerOptions affect job server # configuration, so concat those as the cache key. job_server_options = options.view_as(pipeline_options.JobServerOptions) options_str = str(spark_options) + str(job_server_options) if not options_str in JOB_SERVER_CACHE: JOB_SERVER_CACHE[options_str] = job_server.StopOnExitJobServer( SparkJarJobServer(options)) return JOB_SERVER_CACHE[options_str] def create_job_service_handle(self, job_service, options): return portable_runner.JobServiceHandle( job_service, options, retain_unknown_options=options.view_as( pipeline_options.SparkRunnerOptions).spark_submit_uber_jar) class SparkJarJobServer(job_server.JavaJarJobServer): def __init__(self, options): super().__init__(options) options = options.view_as(pipeline_options.SparkRunnerOptions) self._jar = options.spark_job_server_jar self._master_url = options.spark_master_url self._spark_version = options.spark_version def path_to_jar(self): if self._jar: if not os.path.exists(self._jar): url = urllib.parse.urlparse(self._jar) if not url.scheme: raise ValueError( 'Unable to parse jar URL "%s". If using a full URL, make sure ' 'the scheme is specified. If using a local file path, make sure ' 'the file exists; you may have to first build the job server ' 'using `./gradlew runners:spark:3:job-server:shadowJar`.' % self._jar) return self._jar else: if self._spark_version == '2': raise ValueError('Support for Spark 2 was dropped.') return self.path_to_beam_jar(':runners:spark:3:job-server:shadowJar') def java_arguments( self, job_port, artifact_port, expansion_port, artifacts_dir): return [ '--spark-master-url', self._master_url, '--artifacts-dir', artifacts_dir, '--job-port', job_port, '--artifact-port', artifact_port, '--expansion-port', expansion_port ] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/spark_runner_test.py0000644000175100017510000001631315206113032026041 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import argparse import logging import shlex import unittest from shutil import rmtree from tempfile import mkdtemp import pytest from apache_beam.options.pipeline_options import PortableOptions from apache_beam.runners.portability import job_server from apache_beam.runners.portability import portable_runner from apache_beam.runners.portability import portable_runner_test from apache_beam.utils import subprocess_server # Run as # # pytest spark_runner_test.py[::TestClass::test_case] \ # --test-pipeline-options="--environment_type=LOOPBACK" _LOGGER = logging.getLogger(__name__) class SparkRunnerTest(portable_runner_test.PortableRunnerTest): _use_grpc = True _use_subprocesses = True expansion_port = None spark_job_server_jar = None @pytest.fixture(autouse=True) def parse_options(self, request): if not request.config.option.test_pipeline_options: raise unittest.SkipTest( 'Skipping because --test-pipeline-options is not specified.') test_pipeline_options = request.config.option.test_pipeline_options parser = argparse.ArgumentParser(add_help=True) parser.add_argument( '--spark_job_server_jar', help='Job server jar to submit jobs.', action='store') parser.add_argument( '--environment_type', default='LOOPBACK', choices=['DOCKER', 'PROCESS', 'LOOPBACK'], help='Set the environment type for running user code. DOCKER runs ' 'user code in a container. PROCESS runs user code in ' 'automatically started processes. LOOPBACK runs user code on ' 'the same process that originally submitted the job.') parser.add_argument( '--environment_option', '--environment_options', dest='environment_options', action='append', default=None, help=( 'Environment configuration for running the user code. ' 'Recognized options depend on --environment_type.\n ' 'For DOCKER: docker_container_image (optional)\n ' 'For PROCESS: process_command (required), process_variables ' '(optional, comma-separated)\n ' 'For EXTERNAL: external_service_address (required)')) known_args, unknown_args = parser.parse_known_args( shlex.split(test_pipeline_options)) if unknown_args: _LOGGER.warning('Discarding unrecognized arguments %s' % unknown_args) self.set_spark_job_server_jar( known_args.spark_job_server_jar or job_server.JavaJarJobServer.path_to_beam_jar( ':runners:spark:3:job-server:shadowJar')) self.environment_type = known_args.environment_type self.environment_options = known_args.environment_options @classmethod def _subprocess_command(cls, job_port, expansion_port): # will be cleaned up at the end of this method, and recreated and used by # the job server tmp_dir = mkdtemp(prefix='sparktest') cls.expansion_port = expansion_port try: return [ subprocess_server.JavaHelper.get_java(), '-Dbeam.spark.test.reuseSparkContext=true', '-jar', cls.spark_job_server_jar, '--spark-master-url', 'local', '--artifacts-dir', tmp_dir, '--job-port', str(job_port), '--artifact-port', '0', '--expansion-port', str(expansion_port), ] finally: rmtree(tmp_dir) @classmethod def get_runner(cls): return portable_runner.PortableRunner() @classmethod def get_expansion_service(cls): # TODO Move expansion address resides into PipelineOptions return 'localhost:%s' % cls.expansion_port @classmethod def set_spark_job_server_jar(cls, spark_job_server_jar): cls.spark_job_server_jar = spark_job_server_jar def create_options(self): options = super().create_options() options.view_as(PortableOptions).environment_type = self.environment_type options.view_as( PortableOptions).environment_options = self.environment_options return options def test_metrics(self): # Skip until Spark runner supports metrics. raise unittest.SkipTest("https://github.com/apache/beam/issues/19496") def test_sdf(self): # Skip until Spark runner supports SDF. raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") def test_sdf_with_watermark_tracking(self): # Skip until Spark runner supports SDF. raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") def test_sdf_with_sdf_initiated_checkpointing(self): # Skip until Spark runner supports SDF. raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") def test_sdf_synthetic_source(self): # Skip until Spark runner supports SDF. raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") def test_callbacks_with_exception(self): # Skip until Spark runner supports bundle finalization. raise unittest.SkipTest("https://github.com/apache/beam/issues/19517") def test_register_finalizations(self): # Skip until Spark runner supports bundle finalization. raise unittest.SkipTest("https://github.com/apache/beam/issues/19517") def test_sdf_with_dofn_as_watermark_estimator(self): # Skip until Spark runner supports SDF and self-checkpoint. raise unittest.SkipTest("https://github.com/apache/beam/issues/19468") def test_pardo_dynamic_timer(self): raise unittest.SkipTest("https://github.com/apache/beam/issues/20179") def test_flattened_side_input(self): # Blocked on support for transcoding # https://github.com/apache/beam/issues/19504 super().test_flattened_side_input(with_transcoding=False) def test_flatten_and_gbk(self): # Blocked on support for transcoding # https://github.com/apache/beam/issues/19504 # Also blocked on support of flatten and groupby sharing the same input # https://github.com/apache/beam/issues/34647 raise unittest.SkipTest("https://github.com/apache/beam/issues/34647") def test_custom_merging_window(self): raise unittest.SkipTest("https://github.com/apache/beam/issues/20641") def test_custom_window_type(self): raise unittest.SkipTest("https://github.com/apache/beam/issues/20641") def test_sliding_windows(self): raise unittest.SkipTest("https://github.com/apache/beam/issues/35429") # Inherits all other tests from PortableRunnerTest. if __name__ == '__main__': # Run the tests. logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/spark_uber_jar_job_server.py0000644000175100017510000002374215206113032027506 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A job server submitting portable pipelines as uber jars to Spark.""" # pytype: skip-file import itertools import logging import os import tempfile import time import urllib import zipfile import requests from apache_beam.options import pipeline_options from apache_beam.portability.api import beam_job_api_pb2 from apache_beam.runners.portability import abstract_job_service from apache_beam.runners.portability import job_server from apache_beam.utils.timestamp import Timestamp _LOGGER = logging.getLogger(__name__) class SparkUberJarJobServer(abstract_job_service.AbstractJobServiceServicer): """A Job server which submits a self-contained Jar to a Spark cluster. The jar contains the Beam pipeline definition, dependencies, and the pipeline artifacts. """ def __init__(self, rest_url, options): super().__init__() self._rest_url = rest_url self._artifact_port = ( options.view_as(pipeline_options.JobServerOptions).artifact_port) self._temp_dir = tempfile.mkdtemp(prefix='apache-beam-spark') spark_options = options.view_as(pipeline_options.SparkRunnerOptions) self._executable_jar = spark_options.spark_job_server_jar self._spark_version = spark_options.spark_version self._user_agent = options.view_as(pipeline_options.SetupOptions).user_agent def start(self): return self def stop(self): pass def executable_jar(self): if self._executable_jar: if not os.path.exists(self._executable_jar): parsed = urllib.parse.urlparse(self._executable_jar) if not parsed.scheme: raise ValueError( 'Unable to parse jar URL "%s". If using a full URL, make sure ' 'the scheme is specified. If using a local file path, make sure ' 'the file exists; you may have to first build the job server ' 'using `./gradlew runners:spark:3:job-server:shadowJar`.' % self._executable_jar) url = self._executable_jar else: if self._spark_version == '2': raise ValueError('Support for Spark 2 was dropped.') else: url = job_server.JavaJarJobServer.path_to_beam_jar( ':runners:spark:3:job-server:shadowJar') return job_server.JavaJarJobServer.local_jar( url, user_agent=self._user_agent) def create_beam_job(self, job_id, job_name, pipeline, options): return SparkBeamJob( self._rest_url, self.executable_jar(), job_id, job_name, pipeline, options, artifact_port=self._artifact_port) class SparkBeamJob(abstract_job_service.UberJarBeamJob): """Runs a single Beam job on Spark by staging all contents into a Jar and uploading it via the Spark Rest API. Note that the Spark Rest API is not enabled by default. It must be enabled by setting the configuration property spark.master.rest.enabled to true.""" def __init__( self, rest_url, executable_jar, job_id, job_name, pipeline, options, artifact_port=0): super().__init__( executable_jar, job_id, job_name, pipeline, options, artifact_port=artifact_port) self._rest_url = rest_url # Message history is a superset of state history. self._message_history = self._state_history[:] def request(self, method, path, expected_status=200, **kwargs): url = '%s/%s' % (self._rest_url, path) response = method(url, **kwargs) if response.status_code != expected_status: raise RuntimeError( "Request to %s failed with status %d: %s" % (url, response.status_code, response.text)) if response.text: return response.json() def get(self, path, **kwargs): return self.request(requests.get, path, **kwargs) def post(self, path, **kwargs): return self.request(requests.post, path, **kwargs) def delete(self, path, **kwargs): return self.request(requests.delete, path, **kwargs) def _get_server_spark_version(self): # Spark REST API doesn't seem to offer a dedicated endpoint for getting the # version, but it does include the version in all responses, even errors. return self.get('', expected_status=400)['serverSparkVersion'] def _get_client_spark_version_from_properties(self, jar): """Parse Spark version from spark-version-info.properties file in the jar. https://github.com/apache/spark/blob/dddfeca175bdce5294debe00d4a993daef92ca60/build/spark-build-info#L30 """ with zipfile.ZipFile(jar, 'a', compression=zipfile.ZIP_DEFLATED) as z: with z.open('spark-version-info.properties') as fin: for line in fin.read().decode('utf-8').splitlines(): split = list(map(lambda s: s.strip(), line.split('='))) if len(split) == 2 and split[0] == 'version' and split[1] != '': return split[1] raise ValueError( 'Property "version" not found in spark-version-info.properties.') def _get_client_spark_version(self, jar): try: return self._get_client_spark_version_from_properties(jar) except Exception as e: _LOGGER.debug(e) server_version = self._get_server_spark_version() _LOGGER.warning( 'Unable to parse Spark version from ' 'spark-version-info.properties. Defaulting to %s' % server_version) return server_version def _create_submission_request(self, jar, job_name): jar_url = "file:%s" % jar return { "action": "CreateSubmissionRequest", "appArgs": [], "appResource": jar_url, "clientSparkVersion": self._get_client_spark_version(jar), "environmentVariables": {}, "mainClass": "org.apache.beam.runners.spark.SparkPipelineRunner", "sparkProperties": { "spark.jars": jar_url, "spark.app.name": job_name, "spark.submit.deployMode": "cluster", } } def run(self): self._stop_artifact_service() # Upload the jar and start the job. self._spark_submission_id = self.post( 'v1/submissions/create', json=self._create_submission_request(self._jar, self._job_name))['submissionId'] _LOGGER.info('Submitted Spark job with ID %s' % self._spark_submission_id) def cancel(self): self.post('v1/submissions/kill/%s' % self._spark_submission_id) @staticmethod def _get_beam_state(spark_response): return { 'SUBMITTED': beam_job_api_pb2.JobState.STARTING, 'RUNNING': beam_job_api_pb2.JobState.RUNNING, 'FINISHED': beam_job_api_pb2.JobState.DONE, 'RELAUNCHING': beam_job_api_pb2.JobState.RUNNING, 'UNKNOWN': beam_job_api_pb2.JobState.UNSPECIFIED, 'KILLED': beam_job_api_pb2.JobState.CANCELLED, 'FAILED': beam_job_api_pb2.JobState.FAILED, 'ERROR': beam_job_api_pb2.JobState.FAILED, }.get(spark_response['driverState'], beam_job_api_pb2.JobState.UNSPECIFIED) def _get_spark_status(self): return self.get('v1/submissions/status/%s' % self._spark_submission_id) def get_state(self): response = self._get_spark_status() state = self._get_beam_state(response) timestamp = self.set_state(state) if timestamp is None: # State has not changed since last check. Use previous timestamp. return super().get_state() else: return state, timestamp def _with_message_history(self, message_stream): return itertools.chain(self._message_history[:], message_stream) def _get_message_iter(self): """Returns an iterator of messages from the Spark server. Note that while message history is de-duped, this function's returned iterator may contain duplicate values.""" sleep_secs = 1.0 message_ix = 0 while True: response = self._get_spark_status() state = self._get_beam_state(response) timestamp = Timestamp.now() message = None if 'message' in response: importance = ( beam_job_api_pb2.JobMessage.MessageImportance.JOB_MESSAGE_ERROR if state == beam_job_api_pb2.JobState.FAILED else beam_job_api_pb2.JobMessage.MessageImportance.JOB_MESSAGE_BASIC) message = beam_job_api_pb2.JobMessage( message_id='message%d' % message_ix, time=str(int(timestamp)), importance=importance, message_text=response['message']) yield message message_ix += 1 # TODO(https://github.com/apache/beam/issues/20019) In the event of a # failure, query additional info from Spark master and/or workers. check_timestamp = self.set_state(state) if check_timestamp is not None: if message: self._message_history.append(message) self._message_history.append((state, check_timestamp)) yield state, timestamp sleep_secs = min(60, sleep_secs * 1.2) time.sleep(sleep_secs) def get_state_stream(self): for msg in self._with_message_history(self._get_message_iter()): if isinstance(msg, tuple): state, timestamp = msg yield state, timestamp if self.is_terminal_state(state): break def get_message_stream(self): for msg in self._with_message_history(self._get_message_iter()): yield msg if isinstance(msg, tuple): state, _ = msg if self.is_terminal_state(state): break ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/spark_uber_jar_job_server_test.py0000644000175100017510000002054215206113032030540 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import contextlib import logging import os import tempfile import unittest import zipfile import freezegun import requests_mock from apache_beam.options import pipeline_options from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SparkRunnerOptions from apache_beam.portability.api import beam_job_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.portability import spark_runner from apache_beam.runners.portability import spark_uber_jar_job_server from apache_beam.runners.portability.local_job_service_test import TestJobServicePlan @contextlib.contextmanager def temp_name(*args, **kwargs): with tempfile.NamedTemporaryFile(*args, **kwargs) as t: name = t.name try: yield name finally: if os.path.exists(name): os.unlink(name) def spark_job(): return spark_uber_jar_job_server.SparkBeamJob( 'http://host:6066', '', '', '', '', '', pipeline_options.SparkRunnerOptions()) class SparkUberJarJobServerTest(unittest.TestCase): @requests_mock.mock() def test_get_server_spark_version(self, http_mock): http_mock.get( 'http://host:6066', json={ "action": "ErrorResponse", "message": "Missing protocol version. Please submit requests " "through http://[host]:[port]/v1/submissions/...", "serverSparkVersion": "1.2.3" }, status_code=400) self.assertEqual(spark_job()._get_server_spark_version(), "1.2.3") def test_get_client_spark_version_from_properties(self): with temp_name(suffix='fake.jar') as fake_jar: with zipfile.ZipFile(fake_jar, 'w') as zip: with zip.open('spark-version-info.properties', 'w') as fout: fout.write(b'version=4.5.6') self.assertEqual( spark_job()._get_client_spark_version_from_properties(fake_jar), "4.5.6") def test_get_client_spark_version_from_properties_no_properties_file(self): with self.assertRaises(KeyError): with temp_name(suffix='fake.jar') as fake_jar: with zipfile.ZipFile(fake_jar, 'w') as zip: # Write some other file to the jar. with zip.open('FakeClass.class', 'w') as fout: fout.write(b'[original_contents]') spark_job()._get_client_spark_version_from_properties(fake_jar) def test_get_client_spark_version_from_properties_missing_version(self): with self.assertRaises(ValueError): with temp_name(suffix='fake.jar') as fake_jar: with zipfile.ZipFile(fake_jar, 'w') as zip: with zip.open('spark-version-info.properties', 'w') as fout: fout.write(b'version=') spark_job()._get_client_spark_version_from_properties(fake_jar) @requests_mock.mock() @freezegun.freeze_time("1970-01-01") def test_end_to_end(self, http_mock): submission_id = "submission-id" worker_host_port = "workerhost:12345" worker_id = "worker-id" server_spark_version = "1.2.3" def spark_submission_status_response(state): return { 'json': { "action": "SubmissionStatusResponse", "driverState": state, "serverSparkVersion": server_spark_version, "submissionId": submission_id, "success": "true", "workerHostPort": worker_host_port, "workerId": worker_id } } with temp_name(suffix='fake.jar') as fake_jar: with zipfile.ZipFile(fake_jar, 'w') as zip: with zip.open('spark-version-info.properties', 'w') as fout: fout.write(b'version=4.5.6') options = pipeline_options.SparkRunnerOptions() options.spark_job_server_jar = fake_jar job_server = spark_uber_jar_job_server.SparkUberJarJobServer( 'http://host:6066', options) # Prepare the job. plan = TestJobServicePlan(job_server) # Prepare the job. prepare_response = plan.prepare(beam_runner_api_pb2.Pipeline()) # pylint: disable=assignment-from-no-return retrieval_token = plan.stage( beam_runner_api_pb2.Pipeline(), prepare_response.artifact_staging_endpoint.url, prepare_response.staging_session_token) # Now actually run the job. http_mock.post( 'http://host:6066/v1/submissions/create', json={ "action": "CreateSubmissionResponse", "message": "Driver successfully submitted as submission-id", "serverSparkVersion": "1.2.3", "submissionId": "submission-id", "success": "true" }) job_server.Run( beam_job_api_pb2.RunJobRequest( preparation_id=prepare_response.preparation_id, retrieval_token=retrieval_token)) # Check the status until the job is "done" and get all error messages. http_mock.get( 'http://host:6066/v1/submissions/status/submission-id', [ spark_submission_status_response('RUNNING'), spark_submission_status_response('RUNNING'), { 'json': { "action": "SubmissionStatusResponse", "driverState": "ERROR", "message": "oops", "serverSparkVersion": "1.2.3", "submissionId": submission_id, "success": "true", "workerHostPort": worker_host_port, "workerId": worker_id } } ]) state_stream = job_server.GetStateStream( beam_job_api_pb2.GetJobStateRequest( job_id=prepare_response.preparation_id)) self.assertEqual([s.state for s in state_stream], [ beam_job_api_pb2.JobState.STOPPED, beam_job_api_pb2.JobState.RUNNING, beam_job_api_pb2.JobState.RUNNING, beam_job_api_pb2.JobState.FAILED ]) message_stream = job_server.GetMessageStream( beam_job_api_pb2.JobMessagesRequest( job_id=prepare_response.preparation_id)) def get_item(x): if x.HasField('message_response'): return x.message_response else: return x.state_response.state self.assertEqual([get_item(m) for m in message_stream], [ beam_job_api_pb2.JobState.STOPPED, beam_job_api_pb2.JobState.RUNNING, beam_job_api_pb2.JobMessage( message_id='message0', time='0', importance=beam_job_api_pb2.JobMessage. MessageImportance.JOB_MESSAGE_ERROR, message_text="oops"), beam_job_api_pb2.JobState.FAILED, ]) def test_retain_unknown_options(self): original_options = PipelineOptions(['--unknown_option_foo=some_value']) spark_options = original_options.view_as(SparkRunnerOptions) spark_options.spark_submit_uber_jar = True spark_options.spark_rest_url = 'spark://localhost:6066' runner = spark_runner.SparkRunner() job_service_handle = runner.create_job_service(original_options) options_proto = job_service_handle.get_pipeline_options() self.assertEqual( options_proto['beam:option:unknown_option_foo:v1'], 'some_value') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/stager.py0000644000175100017510000011512315206113032023555 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Support for installing custom code and required dependencies. Workflows, with the exception of very simple ones, are organized in multiple modules and packages. Typically, these modules and packages have dependencies on other standard libraries. Beam relies on the Python setuptools package to handle these scenarios. For further details please read: https://pythonhosted.org/an_example_pypi_project/setuptools.html When a runner tries to run a pipeline it will check for a --requirements_file and a --setup_file option. If --setup_file is present then it is assumed that the folder containing the file specified by the option has the typical layout required by setuptools and it will run 'python setup.py sdist' to produce a source distribution. The resulting tarball (a .tar or .tar.gz file) will be staged at the staging location specified as job option. When a worker starts it will check for the presence of this file and will run 'easy_install tarball' to install the package in the worker. If --requirements_file is present then the file specified by the option will be staged in the staging location. When a worker starts it will check for the presence of this file and will run 'pip install -r requirements.txt'. A requirements file can be easily generated by running 'pip freeze -r requirements.txt'. The reason a runner does not run this automatically is because quite often only a small fraction of the dependencies present in a requirements.txt file are actually needed for remote execution and therefore a one-time manual trimming is desirable. TODO(silviuc): Should we allow several setup packages? TODO(silviuc): We should allow customizing the exact command for setup build. """ # pytype: skip-file import glob import hashlib import logging import os import shutil import subprocess import sys import tempfile from importlib.metadata import distribution from typing import Callable from typing import Optional from urllib.parse import urlparse from packaging import version from apache_beam.internal import pickler from apache_beam.internal.http_client import get_new_http from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystems import FileSystems from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import PipelineOptions # pylint: disable=unused-import from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import WorkerOptions from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.internal import names from apache_beam.utils import processes from apache_beam.utils import retry # All constants are for internal use only; no backwards-compatibility # guarantees. # Standard file names used for staging files. WORKFLOW_TARBALL_FILE = 'workflow.tar.gz' REQUIREMENTS_FILE = 'requirements.txt' EXTRA_PACKAGES_FILE = 'extra_packages.txt' # Filename that stores the submission environment dependencies. SUBMISSION_ENV_DEPENDENCIES_FILE = 'submission_environment_dependencies.txt' # One of the choices for user to use for requirements cache during staging SKIP_REQUIREMENTS_CACHE = 'skip' _LOGGER = logging.getLogger(__name__) def retry_on_non_zero_exit(exception): if (isinstance(exception, processes.CalledProcessError) and exception.returncode != 0): return True class Stager(object): """Abstract Stager identifies and copies the appropriate artifacts to the staging location. Implementation of this stager has to implement :func:`stage_artifact` and :func:`commit_manifest`. """ _DEFAULT_CHUNK_SIZE = 2 << 20 def stage_artifact( self, local_path_to_artifact: str, artifact_name: str, sha256: str) -> None: """ Stages the artifact to Stager._staging_location and adds artifact_name to the manifest of artifacts that have been staged.""" raise NotImplementedError def commit_manifest(self): """Commits manifest.""" raise NotImplementedError @staticmethod def _create_file_stage_to_artifact(local_path, staged_name): return beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2.ArtifactFilePayload( path=local_path).SerializeToString(), role_urn=common_urns.artifact_roles.STAGING_TO.urn, role_payload=beam_runner_api_pb2.ArtifactStagingToRolePayload( staged_name=staged_name).SerializeToString()) @staticmethod def _create_file_pip_requirements_artifact(local_path): return beam_runner_api_pb2.ArtifactInformation( type_urn=common_urns.artifact_types.FILE.urn, type_payload=beam_runner_api_pb2.ArtifactFilePayload( path=local_path).SerializeToString(), role_urn=common_urns.artifact_roles.PIP_REQUIREMENTS_FILE.urn) @staticmethod def extract_staging_tuple_iter( artifacts: list[beam_runner_api_pb2.ArtifactInformation]): for artifact in artifacts: if artifact.type_urn == common_urns.artifact_types.FILE.urn: file_payload = beam_runner_api_pb2.ArtifactFilePayload() file_payload.ParseFromString(artifact.type_payload) src = file_payload.path sha256 = file_payload.sha256 if artifact.role_urn == common_urns.artifact_roles.STAGING_TO.urn: role_payload = beam_runner_api_pb2.ArtifactStagingToRolePayload() role_payload.ParseFromString(artifact.role_payload) dst = role_payload.staged_name elif (artifact.role_urn == common_urns.artifact_roles.PIP_REQUIREMENTS_FILE.urn): dst = hashlib.sha256(artifact.SerializeToString()).hexdigest() else: raise RuntimeError("unknown role type: %s" % artifact.role_urn) yield (src, dst, sha256) else: raise RuntimeError("unknown artifact type: %s" % artifact.type_urn) @staticmethod def create_job_resources( options: PipelineOptions, temp_dir: str, build_setup_args: Optional[list[str]] = None, pypi_requirements: Optional[list[str]] = None, populate_requirements_cache: Optional[Callable[[str, str, bool], None]] = None, skip_prestaged_dependencies: Optional[bool] = False, log_submission_env_dependencies: Optional[bool] = True, ): """For internal use only; no backwards-compatibility guarantees. Creates (if needed) a list of job resources. Args: options: Command line options. More specifically the function will expect requirements_file, setup_file, and save_main_session options to be present. temp_dir: Temporary folder where the resource building can happen. If None then a unique temp directory will be created. Used only for testing. build_setup_args: A list of command line arguments used to build a setup package. Used only if options.setup_file is not None. Used only for testing. pypi_requirements: A list of PyPI requirements used to cache source packages. populate_requirements_cache: Callable for populating the requirements cache. Used only for testing. skip_prestaged_dependencies: Skip staging dependencies that can be added into SDK containers during prebuilding. log_submission_env_dependencies: (Optional) param to stage and log submission environment dependencies. Defaults to True. Returns: A list of ArtifactInformation to be used for staging resources. Raises: RuntimeError: If files specified are not found or error encountered while trying to create the resources (e.g., build a setup package). """ resources: list[beam_runner_api_pb2.ArtifactInformation] = [] setup_options = options.view_as(SetupOptions) use_beam_default_container = options.view_as( WorkerOptions).sdk_container_image is None pickler.set_library(setup_options.pickle_library) # We can skip boot dependencies: apache beam sdk, python packages from # requirements.txt, python packages from extra_packages and workflow tarball # if we know we are using a dependency pre-installed sdk container image. if not skip_prestaged_dependencies: requirements_cache_path = ( os.path.join(tempfile.gettempdir(), 'beam-requirements-cache') if (setup_options.requirements_cache is None) else setup_options.requirements_cache) if (setup_options.requirements_cache != SKIP_REQUIREMENTS_CACHE and not os.path.exists(requirements_cache_path)): os.makedirs(requirements_cache_path, exist_ok=True) # Track packages to stage for this specific run. packages_to_stage = set() # Stage a requirements file if present. if setup_options.requirements_file is not None: if not os.path.isfile(setup_options.requirements_file): raise RuntimeError( 'The file %s cannot be found. It was specified in the ' '--requirements_file command line option.' % setup_options.requirements_file) extra_packages, thinned_requirements_file = ( Stager._extract_local_packages(setup_options.requirements_file)) if extra_packages: setup_options.extra_packages = ( setup_options.extra_packages or []) + extra_packages resources.append( Stager._create_file_stage_to_artifact( thinned_requirements_file, REQUIREMENTS_FILE)) # Populate cache with packages from the requirement file option and # stage the files in the cache. if not use_beam_default_container: _LOGGER.warning( 'When using a custom container image, prefer installing' ' additional PyPI dependencies directly into the image,' ' instead of specifying them via runtime options, ' 'such as --requirements_file. ') if setup_options.requirements_cache != SKIP_REQUIREMENTS_CACHE: populate_cache_callable = ( populate_requirements_cache if populate_requirements_cache else Stager._populate_requirements_cache) downloaded_packages = populate_cache_callable( setup_options.requirements_file, requirements_cache_path, setup_options.requirements_cache_only_sources) if downloaded_packages: packages_to_stage.update(downloaded_packages) if pypi_requirements: tf = tempfile.NamedTemporaryFile(mode='w', delete=False) tf.writelines(pypi_requirements) tf.close() resources.append(Stager._create_file_pip_requirements_artifact(tf.name)) # Populate cache with packages from PyPI requirements and stage # the files in the cache. if setup_options.requirements_cache != SKIP_REQUIREMENTS_CACHE: populate_cache_callable = ( populate_requirements_cache if populate_requirements_cache else Stager._populate_requirements_cache) downloaded_packages = populate_cache_callable( tf.name, requirements_cache_path, setup_options.requirements_cache_only_sources) if downloaded_packages: packages_to_stage.update(downloaded_packages) if (setup_options.requirements_cache != SKIP_REQUIREMENTS_CACHE) and ( setup_options.requirements_file is not None or pypi_requirements): for pkg in packages_to_stage: pkg_path = os.path.join(requirements_cache_path, pkg) if os.path.exists(pkg_path): resources.append( Stager._create_file_stage_to_artifact(pkg_path, pkg)) # Handle a setup file if present. # We will build the setup package locally and then copy it to the staging # location because the staging location is a remote path and the file # cannot be created directly there. if setup_options.setup_file is not None: if not os.path.isfile(setup_options.setup_file): raise RuntimeError( 'The file %s cannot be found. It was specified in the ' '--setup_file command line option.' % setup_options.setup_file) if os.path.basename(setup_options.setup_file) not in ('setup.py', 'pyproject.toml'): raise RuntimeError( 'The --setup_file option expects the full path to a file named ' 'setup.py or pyproject.toml instead of %s' % setup_options.setup_file) tarball_file = Stager._build_setup_package( setup_options.setup_file, temp_dir, build_setup_args) resources.append( Stager._create_file_stage_to_artifact( tarball_file, WORKFLOW_TARBALL_FILE)) if setup_options.files_to_stage is not None: for file in setup_options.files_to_stage: resources.append( Stager._create_file_stage_to_artifact( file, os.path.basename(file))) # Handle extra local packages that should be staged. if setup_options.extra_packages is not None: resources.extend( Stager._create_extra_packages( setup_options.extra_packages, temp_dir=temp_dir)) if hasattr(setup_options, 'sdk_location'): sdk_location = setup_options.sdk_location if Stager._is_remote_path(sdk_location): try: resources.extend( Stager._create_beam_sdk( sdk_remote_location=setup_options.sdk_location, temp_dir=temp_dir, )) except: raise RuntimeError( 'The --sdk_location option was used with an unsupported ' 'type of location: %s' % sdk_location) elif sdk_location == 'default': # Use default location for a runner. pass elif sdk_location == 'container': # Used in the past to indicate that SDK should be used from container # image instead of being staged. # Equivalent to 'default' now, leaving for backwards compatibility. pass else: if os.path.isdir(setup_options.sdk_location): sdk_path = os.path.join( setup_options.sdk_location, names.STAGED_SDK_SOURCES_FILENAME) else: sdk_path = setup_options.sdk_location if os.path.isfile(sdk_path): _LOGGER.info('Copying Beam SDK "%s" to staging location.', sdk_path) resources.append( Stager._create_file_stage_to_artifact( sdk_path, Stager._desired_sdk_filename_in_staging_location( setup_options.sdk_location))) else: if setup_options.sdk_location == 'default': raise RuntimeError( 'Cannot find default Beam SDK tar file "%s"' % sdk_path) elif not setup_options.sdk_location: _LOGGER.info( 'Beam SDK will not be staged since --sdk_location ' 'is empty.') else: raise RuntimeError( 'The file "%s" cannot be found. Its location was specified ' 'by the --sdk_location command-line option.' % sdk_path) # The following artifacts are not processed by python sdk container boot # sequence in a setup mode and hence should not be skipped even if a # prebuilt sdk container image is used. # TODO(heejong): remove jar_packages experimental flag when cross-language # dependency management is implemented for all runners. # Handle jar packages that should be staged for Java SDK Harness. jar_packages = options.view_as(DebugOptions).lookup_experiment( 'jar_packages') if jar_packages is not None: resources.extend( Stager._create_jar_packages( jar_packages.split(','), temp_dir=temp_dir)) # Pickle the main session if requested. # We will create the pickled main session locally and then copy it to the # staging location because the staging location is a remote path and the # file cannot be created directly there. if setup_options.save_main_session: pickled_session_file = os.path.join( temp_dir, names.PICKLED_MAIN_SESSION_FILE) pickler.dump_session(pickled_session_file) if os.path.exists(pickled_session_file): resources.append( Stager._create_file_stage_to_artifact( pickled_session_file, names.PICKLED_MAIN_SESSION_FILE)) # stage the submission environment dependencies, if enabled. if (log_submission_env_dependencies and not options.view_as(DebugOptions).lookup_experiment( 'disable_logging_submission_environment')): resources.extend( Stager._create_stage_submission_env_dependencies(temp_dir)) return resources def stage_job_resources( self, resources: list[tuple[str, str, str]], staging_location: Optional[str] = None): """For internal use only; no backwards-compatibility guarantees. Stages job resources to staging_location. Args: resources: A list of tuples of local file paths and file names (no paths) to be used for staging resources. staging_location: Location to stage the file. Returns: A list of file names (no paths) for the resources staged. All the files are assumed to be staged at staging_location. Raises: RuntimeError: If files specified are not found or error encountered while trying to create the resources (e.g., build a setup package). """ # Make sure that all required options are specified. if staging_location is None: raise RuntimeError('The staging_location must be specified.') staged_resources = [] for file_path, staged_path, sha256 in resources: self.stage_artifact( file_path, FileSystems.join(staging_location, staged_path), sha256) staged_resources.append(staged_path) return staged_resources def create_and_stage_job_resources( self, options: PipelineOptions, build_setup_args: Optional[list[str]] = None, temp_dir: Optional[str] = None, pypi_requirements: Optional[list[str]] = None, populate_requirements_cache: Optional[Callable[[str, str, bool], None]] = None, staging_location: Optional[str] = None): """For internal use only; no backwards-compatibility guarantees. Creates (if needed) and stages job resources to staging_location. Args: options: Command line options. More specifically the function will expect requirements_file, setup_file, and save_main_session options to be present. build_setup_args: A list of command line arguments used to build a setup package. Used only if options.setup_file is not None. Used only for testing. temp_dir: Temporary folder where the resource building can happen. If None then a unique temp directory will be created. Used only for testing. pypi_requirements: A list of PyPI requirements used to cache source packages. populate_requirements_cache: Callable for populating the requirements cache. Used only for testing. staging_location: Location to stage the file. Returns: A tuple of: 1) retrieval token 2) A list of file names (no paths) for the resources staged. All the files are assumed to be staged at staging_location Raises: RuntimeError: If files specified are not found or error encountered while trying to create the resources (e.g., build a setup package). """ temp_dir = temp_dir or tempfile.mkdtemp() resources = self.create_job_resources( options, temp_dir, build_setup_args, pypi_requirements=pypi_requirements, populate_requirements_cache=populate_requirements_cache) staged_resources = self.stage_job_resources( list(Stager.extract_staging_tuple_iter(resources)), staging_location) # Delete all temp files created while staging job resources. shutil.rmtree(temp_dir) retrieval_token = self.commit_manifest() return retrieval_token, staged_resources @staticmethod @retry.with_exponential_backoff(num_retries=4) def _download_file(from_url, to_path): """Downloads a file over http/https from a url or copy it from a remote path to local path.""" if from_url.startswith('http://') or from_url.startswith('https://'): # TODO(silviuc): We should cache downloads so we do not do it for every # job. try: # We check if the file is actually there because wget returns a file # even for a 404 response (file will contain the contents of the 404 # response). response, content = get_new_http().request(from_url) if int(response['status']) >= 400: raise RuntimeError( 'Artifact not found at %s (response: %s)' % (from_url, response)) with open(to_path, 'wb') as f: f.write(content) except Exception: _LOGGER.info('Failed to download Artifact from %s', from_url) raise else: try: read_handle = FileSystems.open( from_url, compression_type=CompressionTypes.UNCOMPRESSED) with read_handle as fin: with open(to_path, 'wb') as f: while True: chunk = fin.read(Stager._DEFAULT_CHUNK_SIZE) if not chunk: break f.write(chunk) _LOGGER.info('Copied remote file from %s to %s.', from_url, to_path) return except Exception as e: _LOGGER.info( 'Failed to download file from %s via apache_beam.io.filesystems.' 'Trying to copy directly. %s', from_url, repr(e)) if not os.path.isdir(os.path.dirname(to_path)): _LOGGER.info( 'Created folder (since we have not done yet, and any errors ' 'will follow): %s ', os.path.dirname(to_path)) os.mkdir(os.path.dirname(to_path)) shutil.copyfile(from_url, to_path) @staticmethod def _is_remote_path(path): return path.find('://') != -1 @staticmethod def _create_jar_packages( jar_packages, temp_dir) -> list[beam_runner_api_pb2.ArtifactInformation]: """Creates a list of local jar packages for Java SDK Harness. :param jar_packages: Ordered list of local paths to jar packages to be staged. Only packages on localfile system and GCS are supported. :param temp_dir: Temporary folder where the resource building can happen. :return: A list of tuples of local file paths and file names (no paths) for the resource staged. All the files are assumed to be staged in staging_location. :raises: RuntimeError: If files specified are not found or do not have expected name patterns. """ resources: list[beam_runner_api_pb2.ArtifactInformation] = [] staging_temp_dir = tempfile.mkdtemp(dir=temp_dir) local_packages: list[str] = [] for package in jar_packages: if not os.path.basename(package).endswith('.jar'): raise RuntimeError( 'The --experiment=\'jar_packages=\' option expects a full path ' 'ending with ".jar" instead of %s' % package) if not os.path.isfile(package): if Stager._is_remote_path(package): # Download remote package. _LOGGER.info( 'Downloading jar package: %s locally before staging', package) _, last_component = FileSystems.split(package) local_file_path = FileSystems.join(staging_temp_dir, last_component) Stager._download_file(package, local_file_path) else: raise RuntimeError( 'The file %s cannot be found. It was specified in the ' '--experiment=\'jar_packages=\' command line option.' % package) else: local_packages.append(package) local_packages.extend([ FileSystems.join(staging_temp_dir, f) for f in os.listdir(staging_temp_dir) ]) for package in local_packages: basename = os.path.basename(package) resources.append(Stager._create_file_stage_to_artifact(package, basename)) return resources @staticmethod def _create_extra_packages( extra_packages, temp_dir) -> list[beam_runner_api_pb2.ArtifactInformation]: """Creates a list of local extra packages. Args: extra_packages: Ordered list of local paths to extra packages to be staged. Only packages on localfile system and GCS are supported. temp_dir: Temporary folder where the resource building can happen. Caller is responsible for cleaning up this folder after this function returns. Returns: A list of ArtifactInformation of local file paths and file names (no paths) for the resources staged. All the files are assumed to be staged in staging_location. Raises: RuntimeError: If files specified are not found or do not have expected name patterns. """ resources: list[beam_runner_api_pb2.ArtifactInformation] = [] staging_temp_dir = tempfile.mkdtemp(dir=temp_dir) local_packages: list[str] = [] for package in extra_packages: if not (os.path.basename(package).endswith('.tar') or os.path.basename(package).endswith('.tar.gz') or os.path.basename(package).endswith('.whl') or os.path.basename(package).endswith('.zip')): raise RuntimeError( 'The --extra_package option expects a full path ending with ' '".tar", ".tar.gz", ".whl" or ".zip" instead of %s' % package) if os.path.basename(package).endswith('.whl'): _LOGGER.warning( 'The .whl package "%s" provided in --extra_package ' 'must be binary-compatible with the worker runtime environment.' % package) if not os.path.isfile(package): if Stager._is_remote_path(package): # Download remote package. _LOGGER.info( 'Downloading extra package: %s locally before staging', package) _, last_component = FileSystems.split(package) local_file_path = FileSystems.join(staging_temp_dir, last_component) Stager._download_file(package, local_file_path) else: raise RuntimeError( 'The file %s cannot be found. It was specified in the ' '--extra_packages command line option.' % package) else: local_packages.append(package) local_packages.extend([ FileSystems.join(staging_temp_dir, f) for f in os.listdir(staging_temp_dir) ]) for package in local_packages: basename = os.path.basename(package) resources.append(Stager._create_file_stage_to_artifact(package, basename)) # Create a file containing the list of extra packages and stage it. # The file is important so that in the worker the packages are installed # exactly in the order specified. This approach will avoid extra PyPI # requests. For example if package A depends on package B and package A # is installed first then the installer will try to satisfy the # dependency on B by downloading the package from PyPI. If package B is # installed first this is avoided. with open(os.path.join(temp_dir, EXTRA_PACKAGES_FILE), 'wt') as f: for package in local_packages: f.write('%s\n' % os.path.basename(package)) # Note that the caller of this function is responsible for deleting the # temporary folder where all temp files are created, including this one. resources.append( Stager._create_file_stage_to_artifact( os.path.join(temp_dir, EXTRA_PACKAGES_FILE), EXTRA_PACKAGES_FILE)) return resources @staticmethod def _get_python_executable(): # Allow overriding the python executable to use for downloading and # installing dependencies, otherwise use the python executable for # the current process. python_bin = os.environ.get('BEAM_PYTHON') or sys.executable if not python_bin: raise ValueError('Could not find Python executable.') return python_bin @staticmethod def _remove_dependency_from_requirements( requirements_file: str, dependency_to_remove: str, temp_directory_path): """Function to remove dependencies from a given requirements file.""" # read all the dependency names with open(requirements_file, 'r') as f: lines = f.readlines() tmp_requirements_filename = os.path.join( temp_directory_path, 'tmp_requirements.txt') with open(tmp_requirements_filename, 'w') as tf: for i in range(len(lines)): if not lines[i].startswith(dependency_to_remove): tf.write(lines[i]) return tmp_requirements_filename @staticmethod def _extract_local_packages(requirements_file): local_deps = [] pypi_deps = [] with open(requirements_file, 'r') as fin: for line in fin: dep = line.strip() if os.path.exists(dep): local_deps.append(dep) else: pypi_deps.append(dep) if local_deps: with tempfile.NamedTemporaryFile(suffix='-requirements.txt', delete=False) as fout: fout.write('\n'.join(pypi_deps).encode('utf-8')) return local_deps, fout.name else: return [], requirements_file @staticmethod def _get_platform_for_default_sdk_container(): """ Get the platform for apache beam SDK container based on Pip version. Note: pip is still expected to download compatible wheel of a package with platform tag manylinux1 if the package on PyPI doesn't have (manylinux2014) or (manylinux2010) wheels. Reference: https://www.python.org/dev/peps/pep-0599/#id21 """ # TODO(anandinguva): When https://github.com/pypa/pip/issues/10760 is # addressed, download wheel based on glibc version in Beam's Python # Base image pip_version = distribution('pip').version if version.parse(pip_version) >= version.parse('19.3'): # pip can only recognize manylinux2014_x86_64 wheels # from version 19.3. return 'manylinux2014_x86_64' else: return 'manylinux2010_x86_64' @staticmethod @retry.with_exponential_backoff( num_retries=4, retry_filter=retry_on_non_zero_exit) def _populate_requirements_cache( requirements_file, cache_dir, populate_cache_with_sdists=False): # The 'pip download' command will not download again if it finds the # tarball with the proper version already present. # It will get the packages downloaded in the order they are presented in # the requirements file and will download package dependencies. # The apache-beam dependency is excluded from requirements cache population # because we stage the SDK separately. with tempfile.TemporaryDirectory() as temp_directory: tmp_requirements_filepath = Stager._remove_dependency_from_requirements( requirements_file=requirements_file, dependency_to_remove='apache-beam', temp_directory_path=temp_directory) # Download to a temporary directory first, then copy to cache. # This allows us to track exactly which packages are needed for this # requirements file. download_dir = tempfile.mkdtemp(dir=temp_directory) cmd_args = [ Stager._get_python_executable(), '-m', 'pip', 'download', '--dest', download_dir, '--find-links', cache_dir, '-r', tmp_requirements_filepath, '--exists-action', 'i', '--no-deps' ] if populate_cache_with_sdists: cmd_args.extend(['--no-binary', ':all:']) else: language_implementation_tag = 'cp' abi_suffix = 'm' if sys.version_info < (3, 8) else '' abi_tag = 'cp%d%d%s' % ( sys.version_info[0], sys.version_info[1], abi_suffix) platform_tag = Stager._get_platform_for_default_sdk_container() cmd_args.extend([ '--implementation', language_implementation_tag, '--abi', abi_tag, '--platform', platform_tag ]) _LOGGER.info('Executing command: %s', cmd_args) processes.check_output(cmd_args, stderr=processes.STDOUT) # Get list of downloaded packages and copy them to the cache downloaded_packages = set() for pkg_file in os.listdir(download_dir): downloaded_packages.add(pkg_file) src_path = os.path.join(download_dir, pkg_file) dest_path = os.path.join(cache_dir, pkg_file) # Only copy if not already in cache if not os.path.exists(dest_path): shutil.copy2(src_path, dest_path) return downloaded_packages @staticmethod def _build_setup_package( setup_file: str, temp_dir: str, build_setup_args: Optional[list[str]] = None) -> str: saved_current_directory = os.getcwd() try: os.chdir(os.path.dirname(setup_file)) if build_setup_args is None: # if build is installed in the user env, use it to # build the sdist else fallback to legacy # setup.py sdist call for setup.py file. try: build_setup_args = [ Stager._get_python_executable(), '-m', 'build', '--no-isolation', # Otherwise, we need internet access to PyPI. '--sdist', '--outdir', temp_dir, os.path.dirname(setup_file), ] _LOGGER.info('Executing command: %s', build_setup_args) processes.check_output(build_setup_args) except RuntimeError: if setup_file.endswith('setup.py'): build_setup_args = [ Stager._get_python_executable(), os.path.basename(setup_file), 'sdist', '--dist-dir', temp_dir ] _LOGGER.info('Executing command: %s', build_setup_args) processes.check_output(build_setup_args) else: # If it's pyproject.toml and `python -m build` failed, # there's no direct legacy fallback. raise RuntimeError( f"Failed to build package from '{setup_file}' using . " f"'python -m build'. Please ensure that the 'build' module " f"is installed and your project's build configuration is valid." ) output_files = glob.glob(os.path.join(temp_dir, '*.tar.gz')) if not output_files: raise RuntimeError( 'File %s not found.' % os.path.join(temp_dir, '*.tar.gz')) return output_files[0] finally: os.chdir(saved_current_directory) @staticmethod def _desired_sdk_filename_in_staging_location(sdk_location) -> str: """Returns the name that SDK file should have in the staging location. Args: sdk_location: Full path to SDK file. """ if sdk_location.endswith('.whl'): _, wheel_filename = FileSystems.split(sdk_location) if wheel_filename.startswith('apache_beam'): return wheel_filename else: raise RuntimeError('Unrecognized SDK wheel file: %s' % sdk_location) else: return names.STAGED_SDK_SOURCES_FILENAME @staticmethod def _create_beam_sdk( sdk_remote_location, temp_dir) -> list[beam_runner_api_pb2.ArtifactInformation]: """Creates a Beam SDK file with the appropriate version. Args: sdk_remote_location: A URL from which the file can be downloaded or a remote file location. The SDK file can be a tarball or a wheel. temp_dir: path to temporary location where the file should be downloaded. Returns: A list of ArtifactInformation of local files path and SDK files that will be staged to the staging location. Raises: RuntimeError: if staging was not successful. """ sdk_remote_parsed = urlparse(sdk_remote_location) sdk_remote_filename = os.path.basename(sdk_remote_parsed.path) local_download_file = os.path.join(temp_dir, sdk_remote_filename) Stager._download_file(sdk_remote_location, local_download_file) staged_name = Stager._desired_sdk_filename_in_staging_location( local_download_file) _LOGGER.info('Staging Beam SDK from %s', sdk_remote_location) return [ Stager._create_file_stage_to_artifact(local_download_file, staged_name) ] @staticmethod def _create_stage_submission_env_dependencies(temp_dir): """Create and stage a file with list of dependencies installed in the submission environment. This list can be used at runtime to compare against the dependencies in the runtime environment. This allows runners to warn users about any potential dependency mismatches and help debug issues related to environment mismatches. Args: temp_dir: path to temporary location where the file should be downloaded. Returns: A list of ArtifactInformation of local file path that will be staged to the staging location. """ try: local_dependency_file_path = os.path.join( temp_dir, SUBMISSION_ENV_DEPENDENCIES_FILE) dependencies = subprocess.check_output( [sys.executable, '-m', 'pip', 'freeze']) local_python_path = f"Python Path: {sys.executable}\n" with open(local_dependency_file_path, 'w') as f: f.write(local_python_path + str(dependencies)) return [ Stager._create_file_stage_to_artifact( local_dependency_file_path, SUBMISSION_ENV_DEPENDENCIES_FILE), ] except Exception as e: _LOGGER.warning( "Couldn't stage a list of installed dependencies in " "submission environment. Got exception: %s", e) return [] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/portability/stager_test.py0000644000175100017510000011756515206113032024630 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the stager module.""" # pytype: skip-file import io import logging import os import shutil import sys import tempfile import unittest import mock import pytest from apache_beam.internal import pickler from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystems import FileSystems from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.runners.internal import names from apache_beam.runners.portability import stager _LOGGER = logging.getLogger(__name__) try: import setuptools except ImportError: setuptools = None class StagerTest(unittest.TestCase): def setUp(self): self._temp_dir = None self.stager = TestStager() def tearDown(self): if self._temp_dir: shutil.rmtree(self._temp_dir) self.stager = None # [https://github.com/apache/beam/issues/21457] set pickler to dill by # default. pickler.set_library(pickler.DEFAULT_PICKLE_LIB) def make_temp_dir(self): if self._temp_dir is None: self._temp_dir = tempfile.mkdtemp() return tempfile.mkdtemp(dir=self._temp_dir) def update_options(self, options): setup_options = options.view_as(SetupOptions) setup_options.sdk_location = '' def create_temp_file(self, path, contents): with open(path, 'w') as f: f.write(contents) return f.name # We can not rely on actual remote file systems paths hence making # '/tmp/remote/' a new remote path. def is_remote_path(self, path): return path.startswith('/tmp/remote/') remote_copied_files: list[str] = [] def file_copy(self, from_path, to_path): if self.is_remote_path(from_path): self.remote_copied_files.append(from_path) _, from_name = os.path.split(from_path) if os.path.isdir(to_path): to_path = os.path.join(to_path, from_name) self.create_temp_file(to_path, 'nothing') _LOGGER.info('Fake copied remote file: %s to %s', from_path, to_path) elif self.is_remote_path(to_path): _LOGGER.info('Faking upload_file(%s, %s)', from_path, to_path) else: shutil.copyfile(from_path, to_path) def populate_requirements_cache( self, requirements_file, cache_dir, populate_cache_with_sdists=False): _ = requirements_file self.create_temp_file(os.path.join(cache_dir, 'abc.txt'), 'nothing') self.create_temp_file(os.path.join(cache_dir, 'def.txt'), 'nothing') # Return the list of packages that were "downloaded" for this requirements return {'abc.txt', 'def.txt'} @mock.patch('apache_beam.runners.portability.stager.open') @mock.patch('apache_beam.runners.portability.stager.get_new_http') def test_download_file_https(self, mock_new_http, mock_open): from_url = 'https://from_url' to_path = '/tmp/http_file/' mock_new_http.return_value.request.return_value = ({ 'status': 200 }, 'file_content') self.stager._download_file(from_url, to_path) assert mock_open.mock_calls == [ mock.call('/tmp/http_file/', 'wb'), mock.call().__enter__(), mock.call().__enter__().write('file_content'), mock.call().__exit__(None, None, None) ] @mock.patch('apache_beam.runners.portability.stager.open') @mock.patch('apache_beam.runners.portability.stager.get_new_http') @mock.patch.object(FileSystems, 'open') def test_download_file_non_http(self, mock_fs_open, mock_new_http, mock_open): from_url = 'gs://bucket/from_url' to_path = '/tmp/file/' mock_fs_open.return_value = io.BytesIO(b"file_content") self.stager._download_file(from_url, to_path) assert not mock_new_http.called mock_fs_open.assert_called_with( from_url, compression_type=CompressionTypes.UNCOMPRESSED) assert mock_open.mock_calls == [ mock.call('/tmp/file/', 'wb'), mock.call().__enter__(), mock.call().__enter__().write(b'file_content'), mock.call().__exit__(None, None, None) ] @mock.patch('apache_beam.runners.portability.stager.os.mkdir') @mock.patch('apache_beam.runners.portability.stager.shutil.copyfile') @mock.patch('apache_beam.runners.portability.stager.get_new_http') def test_download_file_unrecognized( self, mock_new_http, mock_copyfile, mock_mkdir): from_url = '/tmp/from_file' to_path = '/tmp/to_file/' with mock.patch('apache_beam.runners.portability.stager.os.path.isdir', return_value=True): self.stager._download_file(from_url, to_path) assert not mock_new_http.called mock_copyfile.assert_called_with(from_url, to_path) with mock.patch('apache_beam.runners.portability.stager.os.path.isdir', return_value=False): self.stager._download_file(from_url, to_path) assert mock_mkdir.called def test_no_staging_location(self): with self.assertRaises(RuntimeError) as cm: self.stager.stage_job_resources([], staging_location=None) self.assertEqual( 'The staging_location must be specified.', cm.exception.args[0]) def test_no_main_session(self): staging_dir = self.make_temp_dir() options = PipelineOptions() options.view_as(SetupOptions).save_main_session = False self.update_options(options) self.assertEqual([stager.SUBMISSION_ENV_DEPENDENCIES_FILE], self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1]) # xdist adds unpicklable modules to the main session. @pytest.mark.no_xdist @pytest.mark.uses_dill @unittest.skipIf( sys.platform == "win32" and sys.version_info < (3, 8), 'https://github.com/apache/beam/issues/20659: pytest on Windows pulls ' 'in a zipimporter, unpicklable before py3.8') def test_with_main_session(self): pytest.importorskip("dill") staging_dir = self.make_temp_dir() options = PipelineOptions() options.view_as(SetupOptions).save_main_session = True options.view_as(SetupOptions).pickle_library = pickler.USE_DILL self.update_options(options) self.assertEqual([ names.PICKLED_MAIN_SESSION_FILE, stager.SUBMISSION_ENV_DEPENDENCIES_FILE ], self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1]) self.assertTrue( os.path.isfile( os.path.join(staging_dir, names.PICKLED_MAIN_SESSION_FILE))) # (https://github.com/apache/beam/issues/21457): Remove the decorator once # cloudpickle is default pickle library @pytest.mark.no_xdist def test_main_session_staged_when_using_cloudpickle(self): staging_dir = self.make_temp_dir() options = PipelineOptions() options.view_as(SetupOptions).save_main_session = True # even if the save main session is on, no pickle file for main # session is saved when pickle_library==cloudpickle. options.view_as(SetupOptions).pickle_library = pickler.USE_CLOUDPICKLE self.update_options(options) self.assertEqual([ names.PICKLED_MAIN_SESSION_FILE, stager.SUBMISSION_ENV_DEPENDENCIES_FILE ], self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1]) def test_default_resources(self): staging_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) self.assertEqual([stager.SUBMISSION_ENV_DEPENDENCIES_FILE], self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1]) def test_no_submission_env_staging(self): staging_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) resources = self.stager.create_job_resources( options, staging_dir, log_submission_env_dependencies=False) self.assertEqual([], resources) def test_with_requirements_file(self): staging_dir = self.make_temp_dir() requirements_cache_dir = self.make_temp_dir() source_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).requirements_cache = requirements_cache_dir options.view_as(SetupOptions).requirements_file = os.path.join( source_dir, stager.REQUIREMENTS_FILE) self.create_temp_file( os.path.join(source_dir, stager.REQUIREMENTS_FILE), 'nothing') self.assertEqual( sorted([ stager.REQUIREMENTS_FILE, 'abc.txt', 'def.txt', stager.SUBMISSION_ENV_DEPENDENCIES_FILE ]), sorted( self.stager.create_and_stage_job_resources( options, populate_requirements_cache=self.populate_requirements_cache, staging_location=staging_dir)[1])) self.assertTrue( os.path.isfile(os.path.join(staging_dir, stager.REQUIREMENTS_FILE))) def test_with_pypi_requirements(self): staging_dir = self.make_temp_dir() requirements_cache_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).requirements_cache = requirements_cache_dir resources = self.stager.create_and_stage_job_resources( options, pypi_requirements=['nothing>=1.0,<2.0'], populate_requirements_cache=self.populate_requirements_cache, staging_location=staging_dir)[1] self.assertEqual(4, len(resources)) self.assertTrue({'abc.txt', 'def.txt'} <= set(resources)) generated_requirements = ( set(resources) - {'abc.txt', 'def.txt', stager.SUBMISSION_ENV_DEPENDENCIES_FILE}).pop() with open(os.path.join(staging_dir, generated_requirements)) as f: data = f.read() self.assertEqual('nothing>=1.0,<2.0', data) self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'abc.txt'))) self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'def.txt'))) def test_requirements_file_not_present(self): staging_dir = self.make_temp_dir() with self.assertRaises(RuntimeError) as cm: options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).requirements_file = 'nosuchfile' self.stager.create_and_stage_job_resources( options, populate_requirements_cache=self.populate_requirements_cache, staging_location=staging_dir) self.assertEqual( cm.exception.args[0], 'The file %s cannot be found. It was specified in the ' '--requirements_file command line option.' % 'nosuchfile') def test_with_requirements_file_and_cache(self): staging_dir = self.make_temp_dir() source_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).requirements_file = os.path.join( source_dir, stager.REQUIREMENTS_FILE) options.view_as(SetupOptions).requirements_cache = self.make_temp_dir() self.create_temp_file( os.path.join(source_dir, stager.REQUIREMENTS_FILE), 'nothing') self.assertEqual( sorted([ stager.REQUIREMENTS_FILE, 'abc.txt', 'def.txt', stager.SUBMISSION_ENV_DEPENDENCIES_FILE ]), sorted( self.stager.create_and_stage_job_resources( options, populate_requirements_cache=self.populate_requirements_cache, staging_location=staging_dir)[1])) self.assertTrue( os.path.isfile(os.path.join(staging_dir, stager.REQUIREMENTS_FILE))) self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'abc.txt'))) self.assertTrue(os.path.isfile(os.path.join(staging_dir, 'def.txt'))) def test_requirements_cache_not_populated_when_cache_disabled(self): staging_dir = self.make_temp_dir() source_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).requirements_file = os.path.join( source_dir, stager.REQUIREMENTS_FILE) options.view_as( SetupOptions).requirements_cache = stager.SKIP_REQUIREMENTS_CACHE self.create_temp_file( os.path.join(source_dir, stager.REQUIREMENTS_FILE), 'nothing') with mock.patch( 'apache_beam.runners.portability.stager_test.StagerTest.' 'populate_requirements_cache') as (populate_requirements_cache): resources = self.stager.create_and_stage_job_resources( options, populate_requirements_cache=self.populate_requirements_cache, staging_location=staging_dir)[1] assert not populate_requirements_cache.called self.assertEqual( [stager.REQUIREMENTS_FILE, stager.SUBMISSION_ENV_DEPENDENCIES_FILE], resources) self.assertTrue(not os.path.isfile(os.path.join(staging_dir, 'abc.txt'))) self.assertTrue(not os.path.isfile(os.path.join(staging_dir, 'def.txt'))) def test_with_pypi_requirements_skipping_cache(self): staging_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) options.view_as( SetupOptions).requirements_cache = stager.SKIP_REQUIREMENTS_CACHE resources = self.stager.create_and_stage_job_resources( options, pypi_requirements=['nothing>=1.0,<2.0'], populate_requirements_cache=self.populate_requirements_cache, staging_location=staging_dir)[1] with open(os.path.join(staging_dir, resources[0])) as f: data = f.read() self.assertEqual('nothing>=1.0,<2.0', data) self.assertTrue(not os.path.isfile(os.path.join(staging_dir, 'abc.txt'))) self.assertTrue(not os.path.isfile(os.path.join(staging_dir, 'def.txt'))) def test_setup_file_not_present(self): staging_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).setup_file = 'nosuchfile' with self.assertRaises(RuntimeError) as cm: self.stager.create_and_stage_job_resources( options, staging_location=staging_dir) self.assertEqual( cm.exception.args[0], 'The file %s cannot be found. It was specified in the ' '--setup_file command line option.' % 'nosuchfile') def test_setup_file_not_named_setup_dot_py(self): staging_dir = self.make_temp_dir() source_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).setup_file = ( os.path.join(source_dir, 'xyz-setup.py')) self.create_temp_file(os.path.join(source_dir, 'xyz-setup.py'), 'notused') with self.assertRaises(RuntimeError) as cm: self.stager.create_and_stage_job_resources( options, staging_location=staging_dir) self.assertTrue( cm.exception.args[0].startswith( 'The --setup_file option expects the full path to a file named ' 'setup.py or pyproject.toml instead of ')) def test_setup_file_supplies_unexpected_filename(self): staging_dir = self.make_temp_dir() source_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).setup_file = ( os.path.join(source_dir, 'xyz-pyproject.toml')) self.create_temp_file( os.path.join(source_dir, 'xyz-pyproject.toml'), 'notused') with self.assertRaises(RuntimeError) as cm: self.stager.create_and_stage_job_resources( options, staging_location=staging_dir) self.assertTrue( cm.exception.args[0].startswith( 'The --setup_file option expects the full path to a file named ' 'setup.py or pyproject.toml instead of ')) @unittest.skipIf(setuptools is None, "setuptools not available in this env") def test_setup_file_as_setup_dot_py_file(self): source_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).setup_file = ( os.path.join(source_dir, 'setup.py')) self.create_temp_file( os.path.join(source_dir, 'setup.py'), "from setuptools import setup; setup(name='my_package')") temp_dir = tempfile.mkdtemp() resources = self.stager.create_job_resources( options=options, temp_dir=temp_dir) file_path = list(self.stager.extract_staging_tuple_iter(resources))[0][0] self.assertEqual( os.path.join(temp_dir, 'my_package-0.0.0.tar.gz'), file_path) def test_sdk_location_default(self): staging_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).sdk_location = 'default' _, staged_resources = self.stager.create_and_stage_job_resources( options, temp_dir=self.make_temp_dir(), staging_location=staging_dir) self.assertEqual([stager.SUBMISSION_ENV_DEPENDENCIES_FILE], staged_resources) def test_sdk_location_local_directory(self): staging_dir = self.make_temp_dir() sdk_location = self.make_temp_dir() self.create_temp_file( os.path.join(sdk_location, names.STAGED_SDK_SOURCES_FILENAME), 'Package content.') options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).sdk_location = sdk_location self.assertEqual([ names.STAGED_SDK_SOURCES_FILENAME, stager.SUBMISSION_ENV_DEPENDENCIES_FILE ], self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1]) tarball_path = os.path.join(staging_dir, names.STAGED_SDK_SOURCES_FILENAME) with open(tarball_path) as f: self.assertEqual(f.read(), 'Package content.') def test_files_to_stage(self): staging_dir = self.make_temp_dir() source_dir = self.make_temp_dir() foo_ca = os.path.join(source_dir, 'foo.ca') self.create_temp_file(foo_ca, 'ca content') test_txt = os.path.join(source_dir, 'test.txt') self.create_temp_file(test_txt, 'test content') files_to_stage = [foo_ca, test_txt] options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).files_to_stage = files_to_stage self.assertEqual( ['foo.ca', 'test.txt', stager.SUBMISSION_ENV_DEPENDENCIES_FILE], self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1]) foo = os.path.join(staging_dir, 'foo.ca') with open(foo) as f: self.assertEqual(f.read(), 'ca content') txt = os.path.join(staging_dir, 'test.txt') with open(txt) as f: self.assertEqual(f.read(), 'test content') def test_sdk_location_local_source_file(self): staging_dir = self.make_temp_dir() sdk_directory = self.make_temp_dir() sdk_filename = 'apache-beam-3.0.0.tar.gz' sdk_location = os.path.join(sdk_directory, sdk_filename) self.create_temp_file(sdk_location, 'Package content.') options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).sdk_location = sdk_location self.assertEqual([ names.STAGED_SDK_SOURCES_FILENAME, stager.SUBMISSION_ENV_DEPENDENCIES_FILE ], self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1]) tarball_path = os.path.join(staging_dir, names.STAGED_SDK_SOURCES_FILENAME) with open(tarball_path) as f: self.assertEqual(f.read(), 'Package content.') def test_sdk_location_local_wheel_file(self): staging_dir = self.make_temp_dir() sdk_directory = self.make_temp_dir() sdk_filename = 'apache_beam-1.0.0-cp27-cp27mu-manylinux1_x86_64.whl' sdk_location = os.path.join(sdk_directory, sdk_filename) self.create_temp_file(sdk_location, 'Package content.') options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).sdk_location = sdk_location self.assertEqual([sdk_filename, stager.SUBMISSION_ENV_DEPENDENCIES_FILE], self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1]) tarball_path = os.path.join(staging_dir, sdk_filename) with open(tarball_path) as f: self.assertEqual(f.read(), 'Package content.') def test_sdk_location_local_directory_not_present(self): staging_dir = self.make_temp_dir() sdk_location = 'nosuchdir' with self.assertRaises(RuntimeError) as cm: options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).sdk_location = sdk_location self.stager.create_and_stage_job_resources( options, staging_location=staging_dir) self.assertEqual( 'The file "%s" cannot be found. Its ' 'location was specified by the --sdk_location command-line option.' % sdk_location, cm.exception.args[0]) @mock.patch( 'apache_beam.runners.portability.stager_test.TestStager.stage_artifact') @mock.patch( 'apache_beam.runners.portability.stager_test.stager.Stager._download_file' ) def test_sdk_location_remote_source_file(self, *unused_mocks): staging_dir = self.make_temp_dir() sdk_location = 'gs://my-gcs-bucket/tarball.tar.gz' options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).sdk_location = sdk_location self.assertEqual([ names.STAGED_SDK_SOURCES_FILENAME, stager.SUBMISSION_ENV_DEPENDENCIES_FILE ], self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1]) def test_sdk_location_remote_wheel_file(self, *unused_mocks): staging_dir = self.make_temp_dir() sdk_filename = 'apache_beam-1.0.0-cp27-cp27mu-manylinux1_x86_64.whl' sdk_location = 'https://storage.googleapis.com/my-gcs-bucket/' + \ sdk_filename options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).sdk_location = sdk_location def file_download(_, to_path): with open(to_path, 'w') as f: f.write('Package content.') return to_path with mock.patch('apache_beam.runners.portability.stager_test' '.stager.Stager._download_file', staticmethod(file_download)): self.assertEqual([sdk_filename, stager.SUBMISSION_ENV_DEPENDENCIES_FILE], self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1]) wheel_file_path = os.path.join(staging_dir, sdk_filename) with open(wheel_file_path) as f: self.assertEqual(f.read(), 'Package content.') def test_sdk_location_http(self): staging_dir = self.make_temp_dir() sdk_location = 'http://storage.googleapis.com/my-gcs-bucket/tarball.tar.gz' options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).sdk_location = sdk_location def file_download(_, to_path): with open(to_path, 'w') as f: f.write('Package content.') return to_path with mock.patch('apache_beam.runners.portability.stager_test' '.stager.Stager._download_file', staticmethod(file_download)): self.assertEqual([ names.STAGED_SDK_SOURCES_FILENAME, stager.SUBMISSION_ENV_DEPENDENCIES_FILE ], self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1]) tarball_path = os.path.join(staging_dir, names.STAGED_SDK_SOURCES_FILENAME) with open(tarball_path) as f: self.assertEqual(f.read(), 'Package content.') def test_with_extra_packages(self): staging_dir = self.make_temp_dir() source_dir = self.make_temp_dir() self.create_temp_file(os.path.join(source_dir, 'abc.tar.gz'), 'nothing') self.create_temp_file(os.path.join(source_dir, 'xyz.tar.gz'), 'nothing') self.create_temp_file(os.path.join(source_dir, 'xyz2.tar'), 'nothing') self.create_temp_file(os.path.join(source_dir, 'whl.whl'), 'nothing') self.create_temp_file( os.path.join(source_dir, stager.EXTRA_PACKAGES_FILE), 'nothing') options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).extra_packages = [ os.path.join(source_dir, 'abc.tar.gz'), os.path.join(source_dir, 'xyz.tar.gz'), os.path.join(source_dir, 'xyz2.tar'), os.path.join(source_dir, 'whl.whl'), '/tmp/remote/remote_file.tar.gz' ] self.remote_copied_files = [] with mock.patch('apache_beam.runners.portability.stager_test' '.stager.Stager._download_file', staticmethod(self.file_copy)): with mock.patch('apache_beam.runners.portability.stager_test' '.stager.Stager._is_remote_path', staticmethod(self.is_remote_path)): self.assertEqual([ 'abc.tar.gz', 'xyz.tar.gz', 'xyz2.tar', 'whl.whl', 'remote_file.tar.gz', stager.EXTRA_PACKAGES_FILE, stager.SUBMISSION_ENV_DEPENDENCIES_FILE ], self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1]) with open(os.path.join(staging_dir, stager.EXTRA_PACKAGES_FILE)) as f: self.assertEqual([ 'abc.tar.gz\n', 'xyz.tar.gz\n', 'xyz2.tar\n', 'whl.whl\n', 'remote_file.tar.gz\n' ], f.readlines()) self.assertEqual(['/tmp/remote/remote_file.tar.gz'], self.remote_copied_files) def test_with_extra_packages_missing_files(self): staging_dir = self.make_temp_dir() with self.assertRaises(RuntimeError) as cm: options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).extra_packages = ['nosuchfile.tar.gz'] self.stager.create_and_stage_job_resources( options, staging_location=staging_dir) self.assertEqual( cm.exception.args[0], 'The file %s cannot be found. It was specified in the ' '--extra_packages command line option.' % 'nosuchfile.tar.gz') def test_with_extra_packages_invalid_file_name(self): staging_dir = self.make_temp_dir() source_dir = self.make_temp_dir() self.create_temp_file(os.path.join(source_dir, 'abc.tgz'), 'nothing') with self.assertRaises(RuntimeError) as cm: options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).extra_packages = [ os.path.join(source_dir, 'abc.tgz') ] self.stager.create_and_stage_job_resources( options, staging_location=staging_dir) self.assertEqual( cm.exception.args[0], 'The --extra_package option expects a full path ending with ' '".tar", ".tar.gz", ".whl" or ".zip" ' 'instead of %s' % os.path.join(source_dir, 'abc.tgz')) def test_with_jar_packages_missing_files(self): staging_dir = self.make_temp_dir() with self.assertRaises(RuntimeError) as cm: options = PipelineOptions() self.update_options(options) options.view_as(DebugOptions).experiments = [ 'jar_packages=nosuchfile.jar' ] self.stager.create_and_stage_job_resources( options, staging_location=staging_dir) self.assertEqual( cm.exception.args[0], 'The file %s cannot be found. It was specified in the ' '--experiment=\'jar_packages=\' command line option.' % 'nosuchfile.jar') def test_with_jar_packages_invalid_file_name(self): staging_dir = self.make_temp_dir() source_dir = self.make_temp_dir() self.create_temp_file(os.path.join(source_dir, 'abc.tgz'), 'nothing') with self.assertRaises(RuntimeError) as cm: options = PipelineOptions() self.update_options(options) options.view_as(DebugOptions).experiments = [ 'jar_packages=' + os.path.join(source_dir, 'abc.tgz') ] self.stager.create_and_stage_job_resources( options, staging_location=staging_dir) self.assertEqual( cm.exception.args[0], 'The --experiment=\'jar_packages=\' option expects a full path ending ' 'with ".jar" instead of %s' % os.path.join(source_dir, 'abc.tgz')) def test_with_jar_packages(self): staging_dir = self.make_temp_dir() source_dir = self.make_temp_dir() self.create_temp_file(os.path.join(source_dir, 'abc.jar'), 'nothing') self.create_temp_file(os.path.join(source_dir, 'xyz.jar'), 'nothing') self.create_temp_file(os.path.join(source_dir, 'ijk.jar'), 'nothing') options = PipelineOptions() self.update_options(options) options.view_as(DebugOptions).experiments = [ 'jar_packages=%s,%s,%s,%s' % ( os.path.join(source_dir, 'abc.jar'), os.path.join(source_dir, 'xyz.jar'), os.path.join(source_dir, 'ijk.jar'), '/tmp/remote/remote.jar') ] self.remote_copied_files = [] with mock.patch('apache_beam.runners.portability.stager_test' '.stager.Stager._download_file', staticmethod(self.file_copy)): with mock.patch('apache_beam.runners.portability.stager_test' '.stager.Stager._is_remote_path', staticmethod(self.is_remote_path)): self.assertEqual([ 'abc.jar', 'xyz.jar', 'ijk.jar', 'remote.jar', stager.SUBMISSION_ENV_DEPENDENCIES_FILE ], self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1]) self.assertEqual(['/tmp/remote/remote.jar'], self.remote_copied_files) def test_remove_dependency_from_requirements(self): requirements_cache_dir = self.make_temp_dir() requirements = ['apache_beam\n', 'avro-python3\n', 'fastavro\n', 'numpy\n'] with open(os.path.join(requirements_cache_dir, 'abc.txt'), 'w') as f: for i in range(len(requirements)): f.write(requirements[i]) tmp_req_filename = self.stager._remove_dependency_from_requirements( requirements_file=os.path.join(requirements_cache_dir, 'abc.txt'), dependency_to_remove='apache_beam', temp_directory_path=requirements_cache_dir) with open(tmp_req_filename, 'r') as tf: lines = tf.readlines() self.assertEqual(['avro-python3\n', 'fastavro\n', 'numpy\n'], sorted(lines)) tmp_req_filename = self.stager._remove_dependency_from_requirements( requirements_file=os.path.join(requirements_cache_dir, 'abc.txt'), dependency_to_remove='fastavro', temp_directory_path=requirements_cache_dir) with open(tmp_req_filename, 'r') as tf: lines = tf.readlines() self.assertEqual(['apache_beam\n', 'avro-python3\n', 'numpy\n'], sorted(lines)) def _populate_requitements_cache_fake( self, requirements_file, temp_dir, populate_cache_with_sdists): packages = set() if not populate_cache_with_sdists: self.create_temp_file(os.path.join(temp_dir, 'nothing.whl'), 'Fake whl') packages.add('nothing.whl') self.create_temp_file( os.path.join(temp_dir, 'nothing.tar.gz'), 'Fake tarball') packages.add('nothing.tar.gz') return packages # requirements cache will popultated with bdist/whl if present # else source would be downloaded. def test_populate_requirements_cache_with_bdist(self): staging_dir = self.make_temp_dir() requirements_cache_dir = self.make_temp_dir() source_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).requirements_cache = requirements_cache_dir options.view_as(SetupOptions).requirements_file = os.path.join( source_dir, stager.REQUIREMENTS_FILE) self.create_temp_file( os.path.join(source_dir, stager.REQUIREMENTS_FILE), 'nothing') # for default container image, the sdk_container_image option would be none with mock.patch('apache_beam.runners.portability.stager_test' '.stager.Stager._populate_requirements_cache', staticmethod(self._populate_requitements_cache_fake)): options.view_as(SetupOptions).requirements_cache_only_sources = False resources = self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1] for f in resources: if (f != stager.REQUIREMENTS_FILE and f != stager.SUBMISSION_ENV_DEPENDENCIES_FILE): self.assertTrue(('.tar.gz' in f) or ('.whl' in f)) # requirements cache will populated only with sdists/sources def test_populate_requirements_cache_with_sdist(self): staging_dir = self.make_temp_dir() requirements_cache_dir = self.make_temp_dir() source_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).requirements_cache = requirements_cache_dir options.view_as(SetupOptions).requirements_file = os.path.join( source_dir, stager.REQUIREMENTS_FILE) self.create_temp_file( os.path.join(source_dir, stager.REQUIREMENTS_FILE), 'nothing') with mock.patch('apache_beam.runners.portability.stager_test' '.stager.Stager._populate_requirements_cache', staticmethod(self._populate_requitements_cache_fake)): options.view_as(SetupOptions).requirements_cache_only_sources = True resources = self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1] for f in resources: if (f != stager.REQUIREMENTS_FILE and f != stager.SUBMISSION_ENV_DEPENDENCIES_FILE): self.assertTrue('.tar.gz' in f) self.assertTrue('.whl' not in f) def test_populate_requirements_cache_with_local_files(self): staging_dir = self.make_temp_dir() requirements_cache_dir = self.make_temp_dir() source_dir = self.make_temp_dir() pkg_dir = self.make_temp_dir() options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).requirements_cache = requirements_cache_dir options.view_as(SetupOptions).requirements_file = os.path.join( source_dir, stager.REQUIREMENTS_FILE) local_package = os.path.join(pkg_dir, 'local_package.tar.gz') self.create_temp_file(local_package, 'local-package-content') self.create_temp_file( os.path.join(source_dir, stager.REQUIREMENTS_FILE), '\n'.join(['fake_pypi', local_package])) with mock.patch('apache_beam.runners.portability.stager_test' '.stager.Stager._populate_requirements_cache', staticmethod(self._populate_requitements_cache_fake)): options.view_as(SetupOptions).requirements_cache_only_sources = True resources = self.stager.create_and_stage_job_resources( options, staging_location=staging_dir)[1] self.assertEqual( sorted([ stager.REQUIREMENTS_FILE, stager.EXTRA_PACKAGES_FILE, 'nothing.tar.gz', 'local_package.tar.gz', stager.SUBMISSION_ENV_DEPENDENCIES_FILE ]), sorted(resources)) with open(os.path.join(staging_dir, stager.REQUIREMENTS_FILE)) as fin: requirements_contents = fin.read() self.assertIn('fake_pypi', requirements_contents) self.assertNotIn('local_package', requirements_contents) with open(os.path.join(staging_dir, stager.EXTRA_PACKAGES_FILE)) as fin: extra_packages_contents = fin.read() self.assertNotIn('fake_pypi', extra_packages_contents) self.assertIn('local_package', extra_packages_contents) def test_only_required_packages_staged_from_cache(self): """Test that only packages needed for current requirements are staged. This test verifies the fix for the issue where the entire cache directory was being staged, even when some packages weren't needed for the current workflow. """ staging_dir = self.make_temp_dir() requirements_cache_dir = self.make_temp_dir() source_dir = self.make_temp_dir() # Pre-populate cache with packages from a "previous run" self.create_temp_file( os.path.join(requirements_cache_dir, 'old_package.whl'), 'old package') self.create_temp_file( os.path.join(requirements_cache_dir, 'another_old.tar.gz'), 'another') options = PipelineOptions() self.update_options(options) options.view_as(SetupOptions).requirements_cache = requirements_cache_dir options.view_as(SetupOptions).requirements_file = os.path.join( source_dir, stager.REQUIREMENTS_FILE) self.create_temp_file( os.path.join(source_dir, stager.REQUIREMENTS_FILE), 'new_package') def populate_cache_for_new_requirements( requirements_file, cache_dir, populate_cache_with_sdists=False): # Simulate downloading only the packages needed for new requirements self.create_temp_file( os.path.join(cache_dir, 'new_package.whl'), 'new package content') return {'new_package.whl'} resources = self.stager.create_and_stage_job_resources( options, populate_requirements_cache=populate_cache_for_new_requirements, staging_location=staging_dir)[1] # Verify only new_package.whl is staged, not old packages self.assertIn('new_package.whl', resources) self.assertNotIn('old_package.whl', resources) self.assertNotIn('another_old.tar.gz', resources) # Verify the file exists in staging self.assertTrue( os.path.isfile(os.path.join(staging_dir, 'new_package.whl'))) # Verify old packages are NOT in staging self.assertFalse( os.path.isfile(os.path.join(staging_dir, 'old_package.whl'))) self.assertFalse( os.path.isfile(os.path.join(staging_dir, 'another_old.tar.gz'))) def test_populate_requirements_cache_uses_find_links(self): """Test that _populate_requirements_cache uses --find-links to reuse cache. This test verifies that pip download is called with --find-links pointing to the cache directory, so packages already in cache are reused instead of being re-downloaded from PyPI. """ requirements_cache_dir = self.make_temp_dir() source_dir = self.make_temp_dir() # Create a requirements file requirements_file = os.path.join(source_dir, 'requirements.txt') self.create_temp_file(requirements_file, 'some_package==1.0.0') captured_cmd_args = [] def mock_check_output(cmd_args, **kwargs): captured_cmd_args.extend(cmd_args) return b'' with mock.patch( 'apache_beam.runners.portability.stager.processes.check_output', side_effect=mock_check_output): stager.Stager._populate_requirements_cache( requirements_file, requirements_cache_dir) # Verify --find-links is in the command with the cache directory self.assertIn('--find-links', captured_cmd_args) find_links_index = captured_cmd_args.index('--find-links') self.assertEqual( requirements_cache_dir, captured_cmd_args[find_links_index + 1]) class TestStager(stager.Stager): def stage_artifact(self, local_path_to_artifact, artifact_name, sha256): _LOGGER.info( 'File copy from %s to %s.', local_path_to_artifact, artifact_name) shutil.copyfile(local_path_to_artifact, artifact_name) def commit_manifest(self): pass if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/render.py0000644000175100017510000004741415206113032021214 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A portable "runner" that renders a beam graph. This runner can either render the graph to a (set of) output path(s), as designated by (possibly repeated) --render_output, or serve the pipeline as an interactive graph, if --render_port is set. In Python, this runner can be passed directly at pipeline construction, e.g.:: with beam.Pipeline(runner=beam.runners.render.RenderRunner(), options=...) For other languages, start this service a by running:: python -m apache_beam.runners.render --job_port=PORT ... and then run your pipeline with the PortableRunner setting the job endpoint to `localhost:PORT`. If any `--render_output=path.ext` flags are passed, each submitted job will get written to the given output (overwriting any previously existing file). If `--render_port` is set to a non-negative value, a local http server will be started which allows for interactive exploration of the pipeline graph. As an alternative to starting a job server, a single pipeline can be rendered by passing a pipeline proto file to `--pipeline_proto`. For example python -m apache_beam.runners.render \\ --pipeline_proto gs:///pipeline.pb \\ --render_output=/tmp/pipeline.svg Requires the graphviz dot executable to be available in the path. """ import argparse import base64 import collections import http.server import json import logging import os import re import subprocess import sys import tempfile import threading import time import urllib.parse from google.protobuf import json_format from google.protobuf import text_format from apache_beam.options import pipeline_options from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners import runner from apache_beam.runners.portability import local_job_service from apache_beam.runners.portability import local_job_service_main from apache_beam.runners.portability.fn_api_runner import translations try: from apache_beam.io.gcp import gcsio except ImportError: gcsio = None # type: ignore _LOGGER = logging.getLogger(__name__) # From the Beam site, circa November 2022. DEFAULT_EDGE_STYLE = 'color="#ff570b"' DEFAULT_TRANSFORM_STYLE = ( 'shape=rect style="rounded, filled" color="#ff570b" fillcolor="#fff6dd"') DEFAULT_HIGHLIGHT_STYLE = ( 'shape=rect style="rounded, filled" color="#ff570b" fillcolor="#ffdb97"') class RenderOptions(pipeline_options.PipelineOptions): """Rendering options.""" @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--render_port', type=int, default=-1, help='The port at which to serve the graph. ' 'If 0, an unused port will be chosen. ' 'If -1, the server will not be started.') parser.add_argument( '--render_output', action='append', help='A path or paths to which to write rendered output. ' 'The output type will be deduced from the file extension.') parser.add_argument( '--render_leaf_composite_nodes', action='append', help='A set of regular expressions for transform names that should ' 'not be expanded. For example, one could pass "\bRead.*" to indicate ' 'the inner structure of read nodes should not be expanded. ' 'If not given, defaults to the top-level nodes if interactively ' 'serving the graph and expanding all nodes otherwise.') parser.add_argument( '--render_edge_attributes', default='', help='Graphviz attributes to add to all edges.') parser.add_argument( '--render_node_attributes', default='', help='Graphviz attributes to add to all nodes.') parser.add_argument( '--render_highlight_attributes', default='', help='Graphviz attributes to add to all highlighted nodes.') parser.add_argument( '--log_proto', default=False, action='store_true', help='Set to also log input pipeline proto to stdout.') return parser class PipelineRenderer: def __init__(self, pipeline, options): self.pipeline = pipeline self.options = options # Drill down into any uninteresting, top-level transforms that contain # the whole pipeline (often added by the SDK). roots = self.pipeline.root_transform_ids while len(roots) == 1: root = self.pipeline.components.transforms[roots[0]] if not root.subtransforms: break roots = root.subtransforms self.roots = roots # Figure out at what point to stop rendering composite internals. if options.render_leaf_composite_nodes: is_leaf = lambda transform_id: any( re.match( pattern, self.pipeline.components.transforms[transform_id]. unique_name) for patterns in options.render_leaf_composite_nodes for pattern in patterns.split(',')) self.leaf_composites = set() def mark_leaves(transform_ids): for transform_id in transform_ids: if is_leaf(transform_id): self.leaf_composites.add(transform_id) else: mark_leaves( self.pipeline.components.transforms[transform_id].subtransforms) mark_leaves(self.roots) elif options.render_port >= 0: # Start interactive with no unfolding. self.leaf_composites = set(self.roots) else: # For non-interactive, expand fully. self.leaf_composites = set() # Useful for attempting graph layout consistency. self.latest_positions = {} self.highlighted = [] def update(self, toggle=None): if toggle: transform_id = toggle[0] self.highlighted = [transform_id] if transform_id in self.leaf_composites: transform = self.pipeline.components.transforms[transform_id] if transform.subtransforms: self.leaf_composites.remove(transform_id) for subtransform in transform.subtransforms: self.leaf_composites.add(subtransform) if transform_id in self.latest_positions: self.latest_positions[subtransform] = self.latest_positions[ transform_id] else: self.leaf_composites.add(transform_id) def style(self, transform_id): base = ' '.join( [DEFAULT_TRANSFORM_STYLE, self.options.render_node_attributes]) if transform_id in self.highlighted: return ' '.join([ base, DEFAULT_HIGHLIGHT_STYLE, self.options.render_highlight_attributes ]) else: return base def to_dot(self): return '\n'.join(self.to_dot_iter()) def to_dot_iter(self): yield 'digraph G {' # Defer drawing any edges until the end lest we declare nodes too early. edges_out = [] for transform_id in self.roots: yield from self.transform_to_dot( transform_id, self.pcoll_leaf_consumers(), edges_out) yield from edges_out yield '}' def transform_to_dot(self, transform_id, pcoll_leaf_consumers, edges_out): transform = self.pipeline.components.transforms[transform_id] if self.is_leaf(transform_id): yield self.transform_node(transform_id) transform_inputs = set(transform.inputs.values()) for name, output in transform.outputs.items(): # For outputs that are also inputs, it's ambiguous whether they are # consumed as the outputs of this transform, or of the upstream # transform. Render the latter. if output in transform_inputs: continue output_label = name if len(transform.outputs) > 1 else '' for consumer, is_side_input in pcoll_leaf_consumers[output]: # Can't yield this here as the consumer might not be in this cluster. edge_style = 'dashed' if is_side_input else 'solid' edge_attributes = ' '.join([ f'label="{output_label}" style={edge_style}', DEFAULT_EDGE_STYLE, self.options.render_edge_attributes ]) edges_out.append( f'"{transform_id}" -> "{consumer}" [{edge_attributes}]') else: yield f'subgraph "cluster_{transform_id}" {{' yield self.transform_attributes(transform_id) for subtransform in transform.subtransforms: yield from self.transform_to_dot( subtransform, pcoll_leaf_consumers, edges_out) yield '}' def transform_node(self, transform_id): return f'"{transform_id}" [{self.transform_attributes(transform_id)}]' def transform_attributes(self, transform_id): transform = self.pipeline.components.transforms[transform_id] local_name = transform.unique_name.split('/')[-1] if transform_id in self.latest_positions: pos_str = f'pos="{self.latest_positions[transform_id]}"' else: pos_str = '' return ( f'label="{local_name}" {self.style(transform_id)} ' f'URL="javascript:click(\'{transform_id}\')" {pos_str}') def pcoll_leaf_consumers_iter(self, transform_id): transform = self.pipeline.components.transforms[transform_id] transform_inputs = set(transform.inputs.values()) side_inputs = set(translations.side_inputs(transform).values()) if self.is_leaf(transform_id): for pcoll in transform.inputs.values(): yield pcoll, (transform_id, pcoll in side_inputs) for subtransform in transform.subtransforms: for pcoll, (consumer, annotation) in self.pcoll_leaf_consumers_iter(subtransform): if self.is_leaf(transform_id): if pcoll not in transform_inputs: yield pcoll, (transform_id, annotation) else: yield pcoll, (consumer, annotation) def pcoll_leaf_consumers(self): result = collections.defaultdict(list) for transform_id in self.roots: for pcoll, consumer_info in self.pcoll_leaf_consumers_iter(transform_id): result[pcoll].append(consumer_info) return result def is_leaf(self, transform_id): return ( transform_id in self.leaf_composites or not self.pipeline.components.transforms[transform_id].subtransforms) def info(self): if len(self.highlighted) != 1: return '' transform_id = self.highlighted[0] return f'
{self.pipeline.components.transforms[transform_id]}
' def layout_dot(self): layout = subprocess.run(['dot', '-Tdot'], input=self.to_dot().encode('utf-8'), capture_output=True, check=True).stdout # Try to capture the positions for layout consistency. json_out = json.loads( subprocess.run(['dot', '-n2', '-Kneato', '-Tjson'], input=layout, capture_output=True, check=True).stdout) for box in json_out['objects']: name = box.get('name', None) if name in self.pipeline.components.transforms: if 'pos' in box: self.latest_positions[name] = box['pos'] elif 'bb' in box: x0, y0, x1, y1 = [float(r) for r in box['bb'].split(',')] self.latest_positions[name] = f'{(x0+x1)/2},{(y0+y1)/2}' return layout def page_callback_data(self, layout): svg = subprocess.run(['dot', '-Kneato', '-n2', '-Tsvg'], input=layout, capture_output=True, check=True).stdout cmapx = subprocess.run(['dot', '-Kneato', '-n2', '-Tcmapx'], input=layout, capture_output=True, check=True).stdout return { 'src': 'data:image/svg+xml;base64,' + base64.b64encode(svg).decode('utf-8'), 'cmapx': cmapx.decode('utf-8'), 'info': self.info(), } def render_data(self): _LOGGER.info("Re-rendering pipeline...") layout = self.layout_dot() if self.options.render_output: for path in self.options.render_output: format = os.path.splitext(path)[-1][1:] result = subprocess.run( ['dot', '-Kneato', '-n2', '-T' + format, '-o', path], input=layout, check=False) if result.returncode: _LOGGER.error( "Failed render pipeline as %r: exit %s", path, result.returncode) else: _LOGGER.info("Rendered pipeline as %r", path) return self.page_callback_data(layout) def render_json(self): return json.dumps(self.render_data()) def page(self): data = self.render_data() src = data['src'] cmapx = data['cmapx'] return """ """ + f""" Click on a composite transform to expand.
{cmapx}
""" class RenderRunner(runner.PipelineRunner): # TODO(robertwb): Consider making this a runner wrapper, where live status # (such as counters, stage completion status, or possibly even PCollection # samples) queryable and/or displayed. This could evolve into a full Beam # UI. def run_pipeline(self, pipeline_object, options): return self.run_portable_pipeline(pipeline_object.to_runner_api(), options) def run_portable_pipeline(self, pipeline_proto, options): render_options = options.view_as(RenderOptions) if render_options.render_port < 0 and not render_options.render_output: raise ValueError( 'At least one of --render_port or --render_output must be provided.') if render_options.log_proto: _LOGGER.info(pipeline_proto) renderer = PipelineRenderer(pipeline_proto, render_options) try: subprocess.run(['dot', '-V'], capture_output=True, check=True) except FileNotFoundError as exn: # If dot is not available, we can at least output the raw .dot files. dot_files = [ output for output in render_options.render_output if output.endswith('.dot') ] for output in dot_files: with open(output, 'w') as fout: fout.write(renderer.to_dot()) _LOGGER.info("Wrote pipeline as %s", output) non_dot_files = set(render_options.render_output) - set(dot_files) if non_dot_files: raise RuntimeError( "Graphviz dot executable not available " f"for rendering non-dot output files {non_dot_files}") from exn elif render_options.render_port >= 0: raise RuntimeError( "Graphviz dot executable not available for serving") from exn return RenderPipelineResult(None) renderer.page() if render_options.render_port >= 0: # TODO: If this gets more complex, we could consider taking on a # framework like Flask as a dependency. class RequestHandler(http.server.BaseHTTPRequestHandler): def do_GET(self): parts = urllib.parse.urlparse(self.path) args = urllib.parse.parse_qs(parts.query) renderer.update(**args) if parts.path == '/': response = renderer.page() elif parts.path == '/render': response = renderer.render_json() else: self.send_response(400) return self.send_response(200) self.send_header("Content-type", "text/html") self.end_headers() self.wfile.write(response.encode('utf-8')) server = http.server.HTTPServer(('localhost', render_options.render_port), RequestHandler) server_thread = threading.Thread(target=server.serve_forever, daemon=True) server_thread.start() print('Serving at http://%s:%s' % server.server_address) return RenderPipelineResult(server) else: return RenderPipelineResult(None) class RenderPipelineResult(runner.PipelineResult): def __init__(self, server): super().__init__(runner.PipelineState.RUNNING) self.server = server def wait_until_finish(self, duration=None): if self.server: time.sleep(duration or 1e8) self.server.shutdown() self._state = runner.PipelineState.DONE def monitoring_infos(self): return [] def run(argv): if argv[0] == __file__: argv = argv[1:] parser = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument( '--job_port', type=int, default=0, help='port on which to serve the job api') parser.add_argument( '--pipeline_proto', help='file containing the beam pipeline definition') RenderOptions._add_argparse_args(parser) options = parser.parse_args(argv) if options.pipeline_proto: if not options.render_output and options.render_port < 0: options.render_port = 0 render_one(options) if options.render_output: return run_server(options) def render_one(options): if options.pipeline_proto == '-': content = sys.stdin.buffer.read() if content[0] == b'{': ext = '.json' else: try: content.decode('utf-8') ext = '.textproto' except UnicodeDecodeError: ext = '.pb' else: if options.pipeline_proto.startswith('gs://'): if gcsio is None: raise ImportError('GCS not available; please install apache_beam[gcp]') open_fn = gcsio.GcsIO().open else: open_fn = open with open_fn(options.pipeline_proto, 'rb') as fin: content = fin.read() ext = os.path.splitext(options.pipeline_proto)[-1] if ext == '.textproto': pipeline_proto = text_format.Parse(content, beam_runner_api_pb2.Pipeline()) elif ext == '.json': pipeline_proto = json_format.Parse(content, beam_runner_api_pb2.Pipeline()) else: pipeline_proto = beam_runner_api_pb2.Pipeline() pipeline_proto.ParseFromString(content) RenderRunner().run_portable_pipeline( pipeline_proto, pipeline_options.PipelineOptions(**vars(options))) def run_server(options): class RenderBeamJob(local_job_service.BeamJob): def _invoke_runner(self): return RenderRunner().run_portable_pipeline( self._pipeline_proto, pipeline_options.PipelineOptions(**vars(options))) with tempfile.TemporaryDirectory() as staging_dir: job_servicer = local_job_service.LocalJobServicer( staging_dir, beam_job_type=RenderBeamJob) port = job_servicer.start_grpc_server(options.job_port) try: local_job_service_main.serve("Listening for beam jobs on port %d." % port) finally: job_servicer.stop() if __name__ == '__main__': logging.basicConfig() logging.getLogger().setLevel(logging.INFO) run(sys.argv) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/render_test.py0000644000175100017510000001135315206113032022244 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import os import argparse import logging import subprocess import unittest import tempfile import pytest import apache_beam as beam from apache_beam.runners import render default_options = render.RenderOptions._add_argparse_args( argparse.ArgumentParser()).parse_args([]) class RenderRunnerTest(unittest.TestCase): def test_basic_graph(self): p = beam.Pipeline() _ = ( p | beam.Impulse() | beam.Map(lambda _: 2) | 'CustomName' >> beam.Map(lambda x: x * x)) dot = render.PipelineRenderer(p.to_runner_api(), default_options).to_dot() self.assertIn('digraph', dot) self.assertIn('CustomName', dot) self.assertEqual(dot.count('->'), 2) def test_render_config_validation(self): p = beam.Pipeline() _ = ( p | beam.Impulse() | beam.Map(lambda _: 2) | 'CustomName' >> beam.Map(lambda x: x * x)) pipeline_proto = p.to_runner_api() with pytest.raises(ValueError): render.RenderRunner().run_portable_pipeline( pipeline_proto, render.RenderOptions()) def test_side_input(self): p = beam.Pipeline() pcoll = p | beam.Impulse() | beam.FlatMap(lambda x: [1, 2, 3]) dot = render.PipelineRenderer(p.to_runner_api(), default_options).to_dot() self.assertEqual(dot.count('->'), 1) self.assertNotIn('dashed', dot) _ = pcoll | beam.Map( lambda x, side: x * side, side=beam.pvalue.AsList(pcoll)) dot = render.PipelineRenderer(p.to_runner_api(), default_options).to_dot() self.assertEqual(dot.count('->'), 3) self.assertIn('dashed', dot) def test_composite_collapse(self): p = beam.Pipeline() _ = p | beam.Create([1, 2, 3]) | beam.Map(lambda x: x * x) pipeline_proto = p.to_runner_api() renderer = render.PipelineRenderer(pipeline_proto, default_options) self.assertEqual(renderer.to_dot().count('->'), 8) create_transform_id, = [ id for (id, transform) in pipeline_proto.components.transforms.items() if transform.unique_name == 'Create'] renderer.update(toggle=[create_transform_id]) self.assertEqual(renderer.to_dot().count('->'), 1) class DotRequiringRenderingTest(unittest.TestCase): @classmethod def setUpClass(cls): try: subprocess.run(['dot', '-V'], capture_output=True, check=True) except FileNotFoundError: cls._dot_installed = False else: cls._dot_installed = True def setUp(self) -> None: if not self._dot_installed: # type: ignore[attr-defined] self.skipTest('dot executable not installed') def test_run_portable_pipeline(self): p = beam.Pipeline() _ = ( p | beam.Impulse() | beam.Map(lambda _: 2) | 'CustomName' >> beam.Map(lambda x: x * x)) pipeline_proto = p.to_runner_api() with tempfile.TemporaryDirectory() as tmpdir: svg_path = os.path.join(tmpdir, "my_output.svg") render.RenderRunner().run_portable_pipeline( pipeline_proto, render.RenderOptions(render_output=[svg_path])) assert os.path.exists(svg_path) def test_dot_well_formed(self): p = beam.Pipeline() _ = p | beam.Create([1, 2, 3]) | beam.Map(lambda x: x * x) pipeline_proto = p.to_runner_api() renderer = render.PipelineRenderer(pipeline_proto, default_options) # Doesn't actually look at the output, but ensures dot executes correctly. renderer.render_data() create_transform_id, = [ id for (id, transform) in pipeline_proto.components.transforms.items() if transform.unique_name == 'Create'] renderer.update(toggle=[create_transform_id]) renderer.render_data() def test_leaf_composite_filter(self): p = beam.Pipeline() _ = p | beam.Create([1, 2, 3]) | beam.Map(lambda x: x * x) dot = render.PipelineRenderer( p.to_runner_api(), render.RenderOptions(['--render_leaf_composite_nodes=Create' ])).to_dot() self.assertEqual(dot.count('->'), 1) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/runner.py0000644000175100017510000002734415206113032021246 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """PipelineRunner, an abstract base runner object.""" # pytype: skip-file import importlib import logging from typing import TYPE_CHECKING from typing import Iterable from typing import Optional from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import PortableOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import TypeOptions from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.pipeline_utils import group_by_key_input_visitor from apache_beam.transforms import environments if TYPE_CHECKING: from apache_beam import PTransform from apache_beam import pvalue from apache_beam.pipeline import Pipeline __all__ = ['PipelineRunner', 'PipelineState', 'PipelineResult'] _RUNNER_MAP = { path.rsplit('.', maxsplit=1)[-1].lower(): path for path in StandardOptions.ALL_KNOWN_RUNNERS } # Allow this alias, but don't make public. _RUNNER_MAP['pythonrpcdirectrunner'] = ( 'apache_beam.runners.experimental' '.python_rpc_direct.python_rpc_direct_runner.PythonRPCDirectRunner') _LOGGER = logging.getLogger(__name__) def create_runner(runner_name: str) -> 'PipelineRunner': """For internal use only; no backwards-compatibility guarantees. Creates a runner instance from a runner class name. Args: runner_name: Name of the pipeline runner. Possible values are listed in _RUNNER_MAP above. Returns: A runner object. Raises: RuntimeError: if an invalid runner name is used. """ # Get the qualified runner name by using the lower case runner name. If that # fails try appending the name with 'runner' and check if it matches. # If that also fails, use the given runner name as is. runner_name = _RUNNER_MAP.get( runner_name.lower(), _RUNNER_MAP.get(runner_name.lower() + 'runner', runner_name)) if '.' in runner_name: module, runner = runner_name.rsplit('.', 1) try: return getattr(importlib.import_module(module), runner)() except ImportError: if 'dataflow' in runner_name.lower(): raise ImportError( 'Google Cloud Dataflow runner not available, ' 'please install apache_beam[gcp]') elif 'interactive' in runner_name.lower(): raise ImportError( 'Interactive runner not available, ' 'please install apache_beam[interactive]') else: raise else: raise ValueError( 'Unexpected pipeline runner: %s. Valid values are %s ' 'or the fully qualified name of a PipelineRunner subclass.' % (runner_name, ', '.join(StandardOptions.KNOWN_RUNNER_NAMES))) class PipelineRunner(object): """A runner of a pipeline object. The base runner provides a run() method for visiting every node in the pipeline's DAG and executing the transforms computing the PValue in the node. A custom runner will typically provide implementations for some of the transform methods (ParDo, GroupByKey, Create, etc.). It may also provide a new implementation for clear_pvalue(), which is used to wipe out materialized values in order to reduce footprint. """ def run( self, transform: 'PTransform', options: Optional[PipelineOptions] = None) -> 'PipelineResult': """Run the given transform or callable with this runner. Blocks until the pipeline is complete. See also `PipelineRunner.run_async`. """ result = self.run_async(transform, options) result.wait_until_finish() return result def run_async( self, transform: 'PTransform', options: Optional[PipelineOptions] = None) -> 'PipelineResult': """Run the given transform or callable with this runner. May return immediately, executing the pipeline in the background. The returned result object can be queried for progress, and `wait_until_finish` may be called to block until completion. """ # Imported here to avoid circular dependencies. # pylint: disable=wrong-import-order, wrong-import-position from apache_beam import PTransform from apache_beam.pipeline import Pipeline from apache_beam.pvalue import PBegin p = Pipeline(runner=self, options=options) if isinstance(transform, PTransform): p | transform else: transform(PBegin(p)) return p.run() def run_portable_pipeline( self, pipeline: beam_runner_api_pb2.Pipeline, options: PipelineOptions) -> 'PipelineResult': """Execute the entire pipeline. Runners should override this method. """ raise NotImplementedError def default_environment( self, options: PipelineOptions) -> environments.Environment: """Returns the default environment that should be used for this runner. Runners may override this method to provide alternative environments. """ return environments.Environment.from_options( options.view_as(PortableOptions)) def run_pipeline( self, pipeline: 'Pipeline', options: PipelineOptions) -> 'PipelineResult': """Execute the entire pipeline or the sub-DAG reachable from a node. """ pipeline.visit( group_by_key_input_visitor( not options.view_as(TypeOptions).allow_non_deterministic_key_coders) ) # TODO: https://github.com/apache/beam/issues/19168 # portable runner specific default if options.view_as(SetupOptions).sdk_location == 'default': options.view_as(SetupOptions).sdk_location = 'container' return self.run_portable_pipeline( pipeline.to_runner_api( default_environment=self.default_environment(options)), options) def apply( self, transform: 'PTransform', input: Optional['pvalue.PValue'], options: PipelineOptions): # TODO(robertwb): Remove indirection once internal references are fixed. return self.apply_PTransform(transform, input, options) def apply_PTransform(self, transform, input, options): # TODO(robertwb): Remove indirection once internal references are fixed. return transform.expand(input) def is_fnapi_compatible(self): """Whether to enable the beam_fn_api experiment by default.""" return True def check_requirements( self, pipeline_proto: beam_runner_api_pb2.Pipeline, supported_requirements: Iterable[str]): """Check that this runner can satisfy all pipeline requirements.""" # Imported here to avoid circular dependencies. # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.runners.portability.fn_api_runner import translations supported_requirements = set(supported_requirements) for requirement in pipeline_proto.requirements: if requirement not in supported_requirements: raise ValueError( 'Unable to run pipeline with requirement: %s' % requirement) for transform in pipeline_proto.components.transforms.values(): if transform.spec.urn == common_urns.primitives.TEST_STREAM.urn: if common_urns.primitives.TEST_STREAM.urn not in supported_requirements: raise NotImplementedError(transform.spec.urn) elif transform.spec.urn in translations.PAR_DO_URNS: payload = beam_runner_api_pb2.ParDoPayload.FromString( transform.spec.payload) for timer in payload.timer_family_specs.values(): if timer.time_domain not in ( beam_runner_api_pb2.TimeDomain.EVENT_TIME, beam_runner_api_pb2.TimeDomain.PROCESSING_TIME): raise NotImplementedError(timer.time_domain) def default_pickle_library_override(self): """Default pickle library, can be overridden by runner implementation.""" return None # FIXME: replace with PipelineState(str, enum.Enum) class PipelineState(object): """State of the Pipeline, as returned by :attr:`PipelineResult.state`. This is meant to be the union of all the states any runner can put a pipeline in. Currently, it represents the values of the dataflow API JobState enum. """ UNKNOWN = 'UNKNOWN' # not specified by a runner, or unknown to a runner. STARTING = 'STARTING' # not yet started STOPPED = 'STOPPED' # paused or not yet started RUNNING = 'RUNNING' # currently running DONE = 'DONE' # successfully completed (terminal state) FAILED = 'FAILED' # failed (terminal state) CANCELLED = 'CANCELLED' # explicitly cancelled (terminal state) UPDATED = 'UPDATED' # replaced by another job (terminal state) DRAINING = 'DRAINING' # still processing, no longer reading data DRAINED = 'DRAINED' # draining completed (terminal state) PENDING = 'PENDING' # the job has been created but is not yet running. CANCELLING = 'CANCELLING' # job has been explicitly cancelled and is # in the process of stopping RESOURCE_CLEANING_UP = 'RESOURCE_CLEANING_UP' # job's resources are being # cleaned up PAUSING = 'PAUSING' # job is in the process of pausing PAUSED = 'PAUSED' # job has been paused UNRECOGNIZED = 'UNRECOGNIZED' # the job state reported by a runner cannot be # interpreted by the SDK. @classmethod def is_terminal(cls, state): return state in [ cls.DONE, cls.FAILED, cls.CANCELLED, cls.UPDATED, cls.DRAINED ] class PipelineResult(object): """A :class:`PipelineResult` provides access to info about a pipeline.""" def __init__(self, state): self._state = state @property def state(self): """Return the current state of the pipeline execution.""" return self._state def wait_until_finish(self, duration=None): # pylint: disable=unused-argument """Waits until the pipeline finishes and returns the final status. Args: duration (int): The time to wait (in milliseconds) for job to finish. If it is set to :data:`None`, it will wait indefinitely until the job is finished. Raises: IOError: If there is a persistent problem getting job information. NotImplementedError: If the runner does not support this operation. Returns: The final state of the pipeline, or :data:`None` on timeout. """ if not PipelineState.is_terminal(self._state): raise NotImplementedError() def cancel(self): """Cancels the pipeline execution. Raises: IOError: If there is a persistent problem getting job information. NotImplementedError: If the runner does not support this operation. Returns: The final state of the pipeline. """ raise NotImplementedError() def metrics(self): """Returns :class:`~apache_beam.metrics.metric.MetricResults` object to query metrics from the runner. Raises: NotImplementedError: If the runner does not support this operation. """ raise NotImplementedError() # pylint: disable=unused-argument def aggregated_values(self, aggregator_or_name): """Return a dict of step names to values of the Aggregator.""" _LOGGER.warning( '%s does not implement aggregated_values', self.__class__.__name__) return {} ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/runner_test.py0000644000175100017510000000606315206113032022300 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the PipelineRunner and DirectRunner classes. Note that PipelineRunner and DirectRunner functionality is tested in all the other unit tests. In this file we choose to test only aspects related to caching and clearing values that are not tested elsewhere. """ # pytype: skip-file import unittest import apache_beam as beam from apache_beam.metrics.metric import Metrics from apache_beam.runners import DirectRunner from apache_beam.runners import create_runner from apache_beam.runners.portability.fn_api_runner import FnApiRunner class RunnerTest(unittest.TestCase): default_properties = [ '--dataflow_endpoint=ignored', '--job_name=test-job', '--project=test-project', '--staging_location=ignored', '--temp_location=/dev/null', '--no_auth' ] def test_create_runner(self): self.assertTrue(isinstance(create_runner('DirectRunner'), DirectRunner)) self.assertRaises(ValueError, create_runner, 'xyz') def test_create_runner_shorthand(self): self.assertTrue(isinstance(create_runner('DiReCtRuNnEr'), DirectRunner)) self.assertTrue(isinstance(create_runner('directrunner'), DirectRunner)) self.assertTrue(isinstance(create_runner('direct'), DirectRunner)) self.assertTrue(isinstance(create_runner('DiReCt'), DirectRunner)) self.assertTrue(isinstance(create_runner('Direct'), DirectRunner)) def test_run_api(self): my_metric = Metrics.counter('namespace', 'my_metric') runner = FnApiRunner() result = runner.run( beam.Create([1, 10, 100]) | beam.Map(lambda x: my_metric.inc(x))) result.wait_until_finish() # Use counters to assert the pipeline actually ran. my_metric_value = result.metrics().query()['counters'][0].committed self.assertEqual(my_metric_value, 111) def test_run_api_with_callable(self): my_metric = Metrics.counter('namespace', 'my_metric') def fn(start): return ( start | beam.Create([1, 10, 100]) | beam.Map(lambda x: my_metric.inc(x))) runner = FnApiRunner() result = runner.run(fn) result.wait_until_finish() # Use counters to assert the pipeline actually ran. my_metric_value = result.metrics().query()['counters'][0].committed self.assertEqual(my_metric_value, 111) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/sdf_utils.py0000644000175100017510000002166215206113032021726 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file """Common utility class to help SDK harness to execute an SDF. """ import logging import threading from typing import TYPE_CHECKING from typing import Any from typing import NamedTuple from typing import Optional from typing import Union from apache_beam.transforms.core import WatermarkEstimatorProvider from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.windowed_value import WindowedValue if TYPE_CHECKING: from apache_beam.io.iobase import RestrictionProgress from apache_beam.io.iobase import RestrictionTracker from apache_beam.io.iobase import WatermarkEstimator _LOGGER = logging.getLogger(__name__) SplitResultPrimary = NamedTuple( 'SplitResultPrimary', [('primary_value', WindowedValue)]) SplitResultResidual = NamedTuple( 'SplitResultResidual', [('residual_value', WindowedValue), ('current_watermark', Timestamp), ('deferred_timestamp', Optional[Duration])]) class ThreadsafeRestrictionTracker(object): """A thread-safe wrapper which wraps a `RestrictionTracker`. This wrapper guarantees synchronization of modifying restrictions across multi-thread. """ def __init__(self, restriction_tracker: 'RestrictionTracker') -> None: from apache_beam.io.iobase import RestrictionTracker if not isinstance(restriction_tracker, RestrictionTracker): raise ValueError( 'Initialize ThreadsafeRestrictionTracker requires' 'RestrictionTracker.') self._restriction_tracker = restriction_tracker # Records an absolute timestamp when defer_remainder is called. self._timestamp = None self._lock = threading.RLock() self._deferred_residual = None self._deferred_timestamp: Optional[Union[Timestamp, Duration]] = None def current_restriction(self): with self._lock: return self._restriction_tracker.current_restriction() def try_claim(self, position): with self._lock: return self._restriction_tracker.try_claim(position) def defer_remainder(self, deferred_time=None): """Performs self-checkpoint on current processing restriction with an expected resuming time. Self-checkpoint could happen during processing elements. When executing an DoFn.process(), you may want to stop processing an element and resuming later if current element has been processed quit a long time or you also want to have some outputs from other elements. ``defer_remainder()`` can be called on per element if needed. Args: deferred_time: A relative ``Duration`` that indicates the ideal time gap between now and resuming, or an absolute ``Timestamp`` for resuming execution time. If the time_delay is None, the deferred work will be executed as soon as possible. """ # Record current time for calculating deferred_time later. with self._lock: self._timestamp = Timestamp.now() if deferred_time and not isinstance(deferred_time, (Duration, Timestamp)): raise ValueError( 'The timestamp of deter_remainder() should be a ' 'Duration or a Timestamp, or None.') self._deferred_timestamp = deferred_time checkpoint = self.try_split(0) if checkpoint: _, self._deferred_residual = checkpoint def check_done(self): with self._lock: return self._restriction_tracker.check_done() def current_progress(self) -> 'RestrictionProgress': with self._lock: return self._restriction_tracker.current_progress() def try_split(self, fraction_of_remainder): with self._lock: return self._restriction_tracker.try_split(fraction_of_remainder) def deferred_status(self) -> Optional[tuple[Any, Duration]]: """Returns deferred work which is produced by ``defer_remainder()``. When there is a self-checkpoint performed, the system needs to fulfill the DelayedBundleApplication with deferred_work for a ProcessBundleResponse. The system calls this API to get deferred_residual with watermark together to help the runner to schedule a future work. Returns: (deferred_residual, time_delay) if having any residual, else None. """ if self._deferred_residual: # If _deferred_timestamp is None, create Duration(0). if not self._deferred_timestamp: self._deferred_timestamp = Duration() # If an absolute timestamp is provided, calculate the delta between # the absoluted time and the time deferred_status() is called. elif isinstance(self._deferred_timestamp, Timestamp): self._deferred_timestamp = (self._deferred_timestamp - Timestamp.now()) # If a Duration is provided, the deferred time should be: # provided duration - the spent time since the defer_remainder() is # called. elif isinstance(self._deferred_timestamp, Duration): self._deferred_timestamp -= (Timestamp.now() - self._timestamp) return self._deferred_residual, self._deferred_timestamp return None def is_bounded(self): return self._restriction_tracker.is_bounded() class RestrictionTrackerView(object): """A DoFn view of thread-safe RestrictionTracker. The RestrictionTrackerView wraps a ThreadsafeRestrictionTracker and only exposes APIs that will be called by a ``DoFn.process()``. During execution time, the RestrictionTrackerView will be fed into the ``DoFn.process`` as a restriction_tracker. """ def __init__( self, threadsafe_restriction_tracker: ThreadsafeRestrictionTracker) -> None: if not isinstance(threadsafe_restriction_tracker, ThreadsafeRestrictionTracker): raise ValueError( 'Initialize RestrictionTrackerView requires ' 'ThreadsafeRestrictionTracker.') self._threadsafe_restriction_tracker = threadsafe_restriction_tracker def current_restriction(self): return self._threadsafe_restriction_tracker.current_restriction() def try_claim(self, position): return self._threadsafe_restriction_tracker.try_claim(position) def defer_remainder(self, deferred_time=None): self._threadsafe_restriction_tracker.defer_remainder(deferred_time) def is_bounded(self): self._threadsafe_restriction_tracker.is_bounded() class ThreadsafeWatermarkEstimator(object): """A threadsafe wrapper which wraps a WatermarkEstimator with locking mechanism to guarantee multi-thread safety. """ def __init__(self, watermark_estimator: 'WatermarkEstimator') -> None: from apache_beam.io.iobase import WatermarkEstimator if not isinstance(watermark_estimator, WatermarkEstimator): raise ValueError('Initializing Threadsafe requires a WatermarkEstimator') self._watermark_estimator = watermark_estimator self._lock = threading.Lock() def __getattr__(self, attr): if hasattr(self._watermark_estimator, attr): def method_wrapper(*args, **kw): with self._lock: return getattr(self._watermark_estimator, attr)(*args, **kw) return method_wrapper raise AttributeError(attr) def get_estimator_state(self): with self._lock: return self._watermark_estimator.get_estimator_state() def current_watermark(self) -> Timestamp: with self._lock: return self._watermark_estimator.current_watermark() def observe_timestamp(self, timestamp: Timestamp) -> None: if not isinstance(timestamp, Timestamp): raise ValueError( 'Input of observe_timestamp should be a Timestamp ' 'object') with self._lock: self._watermark_estimator.observe_timestamp(timestamp) class NoOpWatermarkEstimatorProvider(WatermarkEstimatorProvider): """A WatermarkEstimatorProvider which creates NoOpWatermarkEstimator for the framework. """ def initial_estimator_state(self, element, restriction): return None def create_watermark_estimator(self, estimator_state): from apache_beam.io.iobase import WatermarkEstimator class _NoOpWatermarkEstimator(WatermarkEstimator): """A No-op WatermarkEstimator which is provided for the framework if there is no custom one. """ def observe_timestamp(self, timestamp): pass def current_watermark(self): return None def get_estimator_state(self): return None return _NoOpWatermarkEstimator() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/sdf_utils_test.py0000644000175100017510000001336615206113032022767 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for classes in sdf_utils.py.""" # pytype: skip-file import time import unittest from apache_beam.io.concat_source_test import RangeSource from apache_beam.io.restriction_trackers import OffsetRange from apache_beam.io.restriction_trackers import OffsetRestrictionTracker from apache_beam.io.watermark_estimators import ManualWatermarkEstimator from apache_beam.runners.sdf_utils import RestrictionTrackerView from apache_beam.runners.sdf_utils import ThreadsafeRestrictionTracker from apache_beam.runners.sdf_utils import ThreadsafeWatermarkEstimator from apache_beam.utils import timestamp class ThreadsafeRestrictionTrackerTest(unittest.TestCase): def test_initialization(self): with self.assertRaises(ValueError): ThreadsafeRestrictionTracker(RangeSource(0, 1)) def test_defer_remainder_with_wrong_time_type(self): threadsafe_tracker = ThreadsafeRestrictionTracker( OffsetRestrictionTracker(OffsetRange(0, 10))) with self.assertRaises(ValueError): threadsafe_tracker.defer_remainder(10) def test_self_checkpoint_immediately(self): restriction_tracker = OffsetRestrictionTracker(OffsetRange(0, 10)) threadsafe_tracker = ThreadsafeRestrictionTracker(restriction_tracker) threadsafe_tracker.defer_remainder() deferred_residual, deferred_time = threadsafe_tracker.deferred_status() expected_residual = OffsetRange(0, 10) self.assertEqual(deferred_residual, expected_residual) self.assertTrue(isinstance(deferred_time, timestamp.Duration)) self.assertEqual(deferred_time, 0) def test_self_checkpoint_with_relative_time(self): threadsafe_tracker = ThreadsafeRestrictionTracker( OffsetRestrictionTracker(OffsetRange(0, 10))) threadsafe_tracker.defer_remainder(timestamp.Duration(100)) time.sleep(2) _, deferred_time = threadsafe_tracker.deferred_status() self.assertTrue(isinstance(deferred_time, timestamp.Duration)) # The expectation = 100 - 2 - some_delta self.assertTrue(deferred_time <= 98) def test_self_checkpoint_with_absolute_time(self): threadsafe_tracker = ThreadsafeRestrictionTracker( OffsetRestrictionTracker(OffsetRange(0, 10))) now = timestamp.Timestamp.now() schedule_time = now + timestamp.Duration(100) self.assertTrue(isinstance(schedule_time, timestamp.Timestamp)) threadsafe_tracker.defer_remainder(schedule_time) time.sleep(2) _, deferred_time = threadsafe_tracker.deferred_status() self.assertTrue(isinstance(deferred_time, timestamp.Duration)) # The expectation = # schedule_time - the time when deferred_status is called - some_delta self.assertTrue(deferred_time <= 98) class RestrictionTrackerViewTest(unittest.TestCase): def test_initialization(self): with self.assertRaises(ValueError): RestrictionTrackerView(OffsetRestrictionTracker(OffsetRange(0, 10))) def test_api_expose(self): threadsafe_tracker = ThreadsafeRestrictionTracker( OffsetRestrictionTracker(OffsetRange(0, 10))) tracker_view = RestrictionTrackerView(threadsafe_tracker) current_restriction = tracker_view.current_restriction() self.assertEqual(current_restriction, OffsetRange(0, 10)) self.assertTrue(tracker_view.try_claim(0)) tracker_view.defer_remainder() deferred_remainder, deferred_watermark = ( threadsafe_tracker.deferred_status()) self.assertEqual(deferred_remainder, OffsetRange(1, 10)) self.assertEqual(deferred_watermark, timestamp.Duration()) def test_non_expose_apis(self): threadsafe_tracker = ThreadsafeRestrictionTracker( OffsetRestrictionTracker(OffsetRange(0, 10))) tracker_view = RestrictionTrackerView(threadsafe_tracker) with self.assertRaises(AttributeError): tracker_view.check_done() with self.assertRaises(AttributeError): tracker_view.current_progress() with self.assertRaises(AttributeError): tracker_view.try_split() with self.assertRaises(AttributeError): tracker_view.deferred_status() class ThreadsafeWatermarkEstimatorTest(unittest.TestCase): def test_initialization(self): with self.assertRaises(ValueError): ThreadsafeWatermarkEstimator(None) def test_get_estimator_state(self): estimator = ThreadsafeWatermarkEstimator(ManualWatermarkEstimator(None)) self.assertIsNone(estimator.get_estimator_state()) estimator.set_watermark(timestamp.Timestamp(10)) self.assertEqual(estimator.get_estimator_state(), timestamp.Timestamp(10)) def test_track_timestamp(self): estimator = ThreadsafeWatermarkEstimator(ManualWatermarkEstimator(None)) estimator.observe_timestamp(timestamp.Timestamp(10)) self.assertIsNone(estimator.current_watermark()) estimator.set_watermark(timestamp.Timestamp(20)) self.assertEqual(estimator.current_watermark(), timestamp.Timestamp(20)) def test_non_exsited_attr(self): estimator = ThreadsafeWatermarkEstimator(ManualWatermarkEstimator(None)) with self.assertRaises(AttributeError): estimator.non_existed_call() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.6814396 apache_beam-2.74.0/apache_beam/runners/test/0000755000175100017510000000000015206113160020332 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/test/__init__.py0000644000175100017510000000245015206113032022442 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test runner objects that's only for end-to-end tests. This package defines runners, which are used to execute test pipeline and verify results. """ # Protect against environments where dataflow runner is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.runners.dataflow.test_dataflow_runner import TestDataflowRunner from apache_beam.runners.direct.test_direct_runner import TestDirectRunner except ImportError: pass # pylint: enable=wrong-import-order, wrong-import-position ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/trivial_runner.py0000644000175100017510000004550315206113032022775 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import collections import logging from typing import Any from typing import Iterable from typing import Iterator from typing import TypeVar from apache_beam import coders from apache_beam.coders.coder_impl import create_InputStream from apache_beam.coders.coder_impl import create_OutputStream from apache_beam.portability import common_urns from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners import common from apache_beam.runners import pipeline_context from apache_beam.runners import runner from apache_beam.runners.portability.fn_api_runner import translations from apache_beam.runners.portability.fn_api_runner import worker_handlers from apache_beam.runners.worker import bundle_processor from apache_beam.transforms import core from apache_beam.transforms import trigger from apache_beam.utils import windowed_value T = TypeVar("T") _LOGGER = logging.getLogger(__name__) class TrivialRunner(runner.PipelineRunner): """A bare-bones batch Python pipeline runner illistrating how to use the RunnerAPI and FnAPI to execute pipelines. Note that this runner is primarily for pedagogical purposes and is missing several features in order to keep it as simple as possible. Where possible pointers are provided which this should serve as a useful starting point. """ def run_portable_pipeline(self, pipeline, options): # First ensure we are able to run this pipeline. # Specifically, that it does not depend on requirements that were # added since this runner was developed. self.check_requirements(pipeline, self.supported_requirements()) # Now we optimize the pipeline, notably performing pipeline fusion, # to turn it into a DAG where all the operations are one of # Impulse, Flatten, GroupByKey, or beam:runner:executable_stage. optimized_pipeline = translations.optimize_pipeline( pipeline, phases=translations.standard_optimize_phases(), known_runner_urns=frozenset([ common_urns.primitives.IMPULSE.urn, common_urns.primitives.FLATTEN.urn, common_urns.primitives.GROUP_BY_KEY.urn ]), # This boolean indicates we want fused executable_stages. partial=False) # standard_optimize_phases() has a final step giving the stages in # topological order, so now we can just walk over them and execute # them. (This are not quite so simple if we were attempting to execute # a streaming pipeline, but this is a trivial runner that only supports # batch...) execution_state = ExecutionState(optimized_pipeline) for transform_id in optimized_pipeline.root_transform_ids: self.execute_transform(transform_id, execution_state) # A more sophisticated runner may perform the execution in the background # and return a PipelineResult that can be used to monitor/cancel the # concurrent execution. return runner.PipelineResult(runner.PipelineState.DONE) def execute_transform(self, transform_id, execution_state): """Execute a single transform.""" transform_proto = execution_state.optimized_pipeline.components.transforms[ transform_id] _LOGGER.info( "Executing stage %s %s", transform_id, transform_proto.unique_name) if not is_primitive_transform(transform_proto): # A composite is simply executed by executing its parts. for sub_transform in transform_proto.subtransforms: self.execute_transform(sub_transform, execution_state) elif transform_proto.spec.urn == common_urns.primitives.IMPULSE.urn: # An impulse has no inputs and produces a single output (which happens # to be an empty byte string in the global window). execution_state.set_pcollection_contents( only_element(transform_proto.outputs.values()), [common.ENCODED_IMPULSE_VALUE]) elif transform_proto.spec.urn == common_urns.primitives.FLATTEN.urn: # The output of a flatten is simply the union of its inputs. output_pcoll_id = only_element(transform_proto.outputs.values()) execution_state.set_pcollection_contents( output_pcoll_id, sum([ execution_state.get_pcollection_contents(pc) for pc in transform_proto.inputs.values() ], [])) elif transform_proto.spec.urn == common_urns.executable_stage: # This is a collection of user DoFns. self.execute_executable_stage(transform_proto, execution_state) elif transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn: # Execute the grouping operation. self.group_by_key_and_window( only_element(transform_proto.inputs.values()), only_element(transform_proto.outputs.values()), execution_state) else: raise RuntimeError( f"Unsupported transform {transform_id}" " of type {transform_proto.spec.urn}") def execute_executable_stage(self, transform_proto, execution_state): # Stage here is like a mini pipeline, with PTransforms, PCollections, etc. # inside of it. stage = beam_runner_api_pb2.ExecutableStagePayload.FromString( transform_proto.spec.payload) if stage.side_inputs: # To support these we would need to make the side input PCollections # available over the state API before processing this bundle. raise NotImplementedError() # This is the set of transforms that were fused together. stage_transforms = { id: stage.components.transforms[id] for id in stage.transforms } # The executable stage has bare PCollections as its inputs and outputs. # # We need an operation to feed data into the bundle (runner -> SDK). # This is done by attaching special transform that reads from the data # channel. input_transform = execution_state.new_id('stage_input') input_pcoll = stage.input stage_transforms[input_transform] = beam_runner_api_pb2.PTransform( # Read data, encoded with the given coder, from the data channel. spec=beam_runner_api_pb2.FunctionSpec( urn=bundle_processor.DATA_INPUT_URN, payload=beam_fn_api_pb2.RemoteGrpcPort( # If we were using a cross-process data channel, we would also # need to set the address of the data channel itself here. coder_id=execution_state.windowed_coder_id( stage.input)).SerializeToString()), # Wire its "output" to the required PCollection. outputs={'out': input_pcoll}) # Also add operations to consume data produced by the bundle and writes # them to the data channel (SDK -> runner). output_ops_to_pcoll = {} for output_pcoll in stage.outputs: output_transform = execution_state.new_id('stage_output') stage_transforms[output_transform] = beam_runner_api_pb2.PTransform( # Data will be written here, with the given coder. spec=beam_runner_api_pb2.FunctionSpec( urn=bundle_processor.DATA_OUTPUT_URN, payload=beam_fn_api_pb2.RemoteGrpcPort( # Again, the grpc address itself is implicit. coder_id=execution_state.windowed_coder_id( output_pcoll)).SerializeToString()), # This operation takes as input the bundle's output pcollection. inputs={'input': output_pcoll}) output_ops_to_pcoll[output_transform] = output_pcoll # Now we can create a description of what it means to process a bundle # of this type. When processing the bundle, we simply refer to this # descriptor by id. (In our case we only do so once, but in a more general # runner one may want to process the same "type" of bundle of many distinct # partitions of the input, especially in streaming where one may have # hundreds of concurrently processed bundles per key.) process_bundle_descriptor = beam_fn_api_pb2.ProcessBundleDescriptor( id=execution_state.new_id('descriptor'), transforms=stage_transforms, pcollections=stage.components.pcollections, coders=execution_state.optimized_pipeline.components.coders, windowing_strategies=stage.components.windowing_strategies, environments=stage.components.environments, # Were timers and state supported, their endpoints would be listed here. ) execution_state.register_process_bundle_descriptor( process_bundle_descriptor) # Now we are ready to actually execute the bundle. process_bundle_id = execution_state.new_id('bundle') # First, push the all input data onto the data channel. # Timers would be sent over this channel as well. # In a real runner we could do this after (or concurrently with) starting # the bundle to avoid having to hold all the data to process in memory, # but for simplicity our bundle invocation (below) is synchronous so the # data must be available for processing right away. to_worker = execution_state.worker_handler.data_conn.output_stream( process_bundle_id, input_transform) for encoded_data in execution_state.get_pcollection_contents(input_pcoll): to_worker.write(encoded_data) to_worker.close() # Now we send a process bundle request over the control plane. process_bundle_request = beam_fn_api_pb2.InstructionRequest( instruction_id=process_bundle_id, process_bundle=beam_fn_api_pb2.ProcessBundleRequest( process_bundle_descriptor_id=process_bundle_descriptor.id)) result_future = execution_state.worker_handler.control_conn.push( process_bundle_request) # Read the results off the data channel. # Note that if there are multiple outputs, we may get them in any order, # possibly interleaved. for output in execution_state.worker_handler.data_conn.input_elements( process_bundle_id, list(output_ops_to_pcoll.keys())): if isinstance(output, beam_fn_api_pb2.Elements.Data): # Adds the output to the appropriate PCollection. execution_state.set_pcollection_contents( output_ops_to_pcoll[output.transform_id], [output.data]) else: # E.g. timers to set. raise RuntimeError("Unexpected data type: %s" % output) # Ensure the operation completed successfully. # This result contains things like metrics and continuation tokens as well. result = result_future.get() if result.error: raise RuntimeError(result.error) if result.process_bundle.residual_roots: # We would need to re-schedule execution of this bundle with this data. raise NotImplementedError('SDF continuation') if result.process_bundle.requires_finalization: # We would need to invoke the finalization callback, on a best effort # basis, *after* the outputs are durably committed. raise NotImplementedError('finalization') if result.process_bundle.elements.data: # These should be processed just like outputs from the data channel. raise NotImplementedError('control-channel data') if result.process_bundle.elements.timers: # These should be processed just like outputs from the data channel. raise NotImplementedError('timers') def group_by_key_and_window(self, input_pcoll, output_pcoll, execution_state): """Groups the elements of input_pcoll, placing their output in output_pcoll. """ # Note that we are using the designated coders to decode and deeply inspect # the elements. This could be useful for other operations as well (e.g. # splitting the returned bundles of encoded elements into smaller chunks). # One issue that we're sweeping under the rug here is that for a truly # portable multi-language runner we may not be able to instanciate # every coder in Python. This can be overcome by wrapping unknown coders as # length-prefixing (or otherwise delimiting) variants. # Decode the input elements to get at their individual keys and values. input_coder = execution_state.windowed_coder(input_pcoll) key_coder = input_coder.key_coder() input_elements = [] for encoded_elements in execution_state.get_pcollection_contents( input_pcoll): for element in decode_all(encoded_elements, input_coder): input_elements.append(element) # Now perform the actual grouping. components = execution_state.optimized_pipeline.components windowing = components.windowing_strategies[ components.pcollections[input_pcoll].windowing_strategy_id] if (windowing.merge_status == beam_runner_api_pb2.MergeStatus.Enum.NON_MERGING and windowing.output_time == beam_runner_api_pb2.OutputTime.Enum.END_OF_WINDOW): # This is the "easy" case, show how to do it by hand. # Note that we're grouping by encoded key, and also by the window. grouped = collections.defaultdict(list) for element in input_elements: for window in element.windows: key, value = element.value grouped[window, key_coder.encode(key)].append(value) output_elements = [ windowed_value.WindowedValue( (key_coder.decode(encoded_key), values), window.end, [window], trigger.BatchGlobalTriggerDriver.ONLY_FIRING) for ((window, encoded_key), values) in grouped.items() ] else: # This handles generic merging and triggering. trigger_driver = trigger.create_trigger_driver( execution_state.windowing_strategy(input_pcoll), True) grouped_by_key = collections.defaultdict(list) for element in input_elements: key, value = element.value grouped_by_key[key_coder.encode(key)].append(element.with_value(value)) output_elements = [] for encoded_key, windowed_values in grouped_by_key.items(): for grouping in trigger_driver.process_entire_key( key_coder.decode(encoded_key), windowed_values): output_elements.append(grouping) # Store the grouped values in the output PCollection. output_coder = execution_state.windowed_coder(output_pcoll) execution_state.set_pcollection_contents( output_pcoll, [encode_all(output_elements, output_coder)]) def supported_requirements(self) -> Iterable[str]: # Nothing non-trivial is supported. return [] class ExecutionState: """A helper class holding various values and context during execution.""" def __init__(self, optimized_pipeline): self.optimized_pipeline = optimized_pipeline self._pcollections_to_encoded_chunks = {} self._counter = 0 self._process_bundle_descriptors = {} # This emulates a connection to an SDK worker (e.g. its data, control, # etc. channels). # There are other variants available as well (e.g. GRPC, Docker, ...) self.worker_handler = worker_handlers.EmbeddedWorkerHandler( None, state=worker_handlers.StateServicer(), provision_info=None, worker_manager=self) self._windowed_coders = {} # Populate all windowed coders before creating _pipeline_context. for pcoll_id in self.optimized_pipeline.components.pcollections.keys(): self.windowed_coder_id(pcoll_id) self._pipeline_context = pipeline_context.PipelineContext( optimized_pipeline.components) def register_process_bundle_descriptor( self, process_bundle_descriptor: beam_fn_api_pb2.ProcessBundleDescriptor): self._process_bundle_descriptors[ process_bundle_descriptor.id] = process_bundle_descriptor def get_pcollection_contents(self, pcoll_id: str) -> list[bytes]: return self._pcollections_to_encoded_chunks[pcoll_id] def set_pcollection_contents(self, pcoll_id: str, chunks: list[bytes]): self._pcollections_to_encoded_chunks[pcoll_id] = chunks def new_id(self, prefix='') -> str: self._counter += 1 return f'runner_{prefix}_{self._counter}' def windowed_coder(self, pcollection_id: str) -> coders.Coder: return self._pipeline_context.coders.get_by_id( self.windowed_coder_id(pcollection_id)) def windowing_strategy(self, pcollection_id: str) -> core.Windowing: return self._pipeline_context.windowing_strategies.get_by_id( self.optimized_pipeline.components.pcollections[pcollection_id]. windowing_strategy_id) def windowed_coder_id(self, pcollection_id: str) -> str: pcoll = self.optimized_pipeline.components.pcollections[pcollection_id] windowing = self.optimized_pipeline.components.windowing_strategies[ pcoll.windowing_strategy_id] return self._windowed_coder_id_from( pcoll.coder_id, windowing.window_coder_id) def _windowed_coder_id_from(self, coder_id: str, window_coder_id: str) -> str: if (coder_id, window_coder_id) not in self._windowed_coders: windowed_coder_id = self.new_id('windowed_coder') self._windowed_coders[coder_id, window_coder_id] = windowed_coder_id self.optimized_pipeline.components.coders[windowed_coder_id].CopyFrom( beam_runner_api_pb2.Coder( spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.coders.WINDOWED_VALUE.urn), component_coder_ids=[coder_id, window_coder_id])) return self._windowed_coders[coder_id, window_coder_id] def is_primitive_transform(transform: beam_runner_api_pb2.PTransform) -> bool: # pylint: disable=consider-using-ternary return ( # Only non-primitive (aka composite) transforms have subtransforms. not transform.subtransforms # If it has outputs but all of the outputs are also inputs to this # transform, there is nothing to compute. and not transform.outputs or bool(set(transform.outputs.values()) - set(transform.inputs.values()))) def only_element(iterable: Iterable[T]) -> T: element, = iterable return element def decode_all(encoded_elements: bytes, coder: coders.Coder) -> Iterator[Any]: coder_impl = coder.get_impl() input_stream = create_InputStream(encoded_elements) while input_stream.size() > 0: yield coder_impl.decode_from_stream(input_stream, True) def encode_all(elements: Iterator[T], coder: coders.Coder) -> bytes: coder_impl = coder.get_impl() output_stream = create_OutputStream() for element in elements: coder_impl.encode_to_stream(element, output_stream, True) return output_stream.get() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/trivial_runner_test.py0000644000175100017510000000543715206113032024036 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest import apache_beam as beam from apache_beam.runners.trivial_runner import TrivialRunner from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class TrivialRunnerTest(unittest.TestCase): def test_trivial(self): # The most trivial pipeline, to ensure at least something is working. # (Notably avoids the non-trivial complexity within assert_that.) with beam.Pipeline(runner=TrivialRunner()) as p: _ = p | beam.Impulse() def test_assert_that(self): # If this fails, the other tests may be vacuous. with self.assertRaisesRegex(Exception, 'Failed assert'): with beam.Pipeline(runner=TrivialRunner()) as p: assert_that(p | beam.Impulse(), equal_to(['a'])) def test_impulse(self): with beam.Pipeline(runner=TrivialRunner()) as p: assert_that(p | beam.Impulse(), equal_to([b''])) def test_create(self): with beam.Pipeline(runner=TrivialRunner()) as p: assert_that(p | beam.Create(['a', 'b']), equal_to(['a', 'b'])) def test_flatten(self): with beam.Pipeline(runner=TrivialRunner()) as p: # Disabling reshuffle here (and elsewhere) just to give a simpler graph # for debugging. ab = p | 'AB' >> beam.Create(['a', 'b'], reshuffle=False) c = p | 'C' >> beam.Create(['c'], reshuffle=False) assert_that((ab, c, c) | beam.Flatten(), equal_to(['a', 'b', 'c', 'c'])) def test_map(self): with beam.Pipeline(runner=TrivialRunner()) as p: assert_that( p | beam.Create(['a', 'b'], reshuffle=False) | beam.Map(str.upper), equal_to(['A', 'B'])) def test_gbk(self): with beam.Pipeline(runner=TrivialRunner()) as p: result = ( p | beam.Create([('a', 1), ('b', 2), ('b', 3)], reshuffle=False) | beam.GroupByKey() | beam.MapTuple(lambda k, vs: (k, sorted(vs)))) assert_that(result, equal_to([('a', [1]), ('b', [2, 3])])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.6924407 apache_beam-2.74.0/apache_beam/runners/worker/0000755000175100017510000000000015206113160020664 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/__init__.py0000644000175100017510000000152515206113032022776 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """For internal use only; no backwards-compatibility guarantees.""" ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/bundle_processor.py0000644000175100017510000024261115206113032024612 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """SDK harness for executing Python Fns via the Fn API.""" # pytype: skip-file from __future__ import annotations import base64 import bisect import collections import concurrent.futures import copy import heapq import itertools import json import logging import random import threading from dataclasses import dataclass from dataclasses import field from itertools import chain from typing import TYPE_CHECKING from typing import Any from typing import Callable from typing import Container from typing import Iterable from typing import Iterator from typing import Mapping from typing import MutableMapping from typing import Optional from typing import TypeVar from typing import Union from typing import cast from google.protobuf import duration_pb2 from google.protobuf import timestamp_pb2 from sortedcontainers import SortedDict from sortedcontainers import SortedList import apache_beam as beam from apache_beam import coders from apache_beam.coders import WindowedValueCoder from apache_beam.coders import coder_impl from apache_beam.internal import pickler from apache_beam.io import iobase from apache_beam.metrics import monitoring_infos from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners import common from apache_beam.runners import pipeline_context from apache_beam.runners.worker import data_sampler from apache_beam.runners.worker import operation_specs from apache_beam.runners.worker import operations from apache_beam.runners.worker import statesampler from apache_beam.runners.worker.worker_status import thread_dump from apache_beam.transforms import TimeDomain from apache_beam.transforms import core from apache_beam.transforms import environments from apache_beam.transforms import sideinputs from apache_beam.transforms import userstate from apache_beam.transforms import window from apache_beam.utils import counters from apache_beam.utils import proto_utils from apache_beam.utils import timestamp from apache_beam.utils.windowed_value import WindowedValue if TYPE_CHECKING: from google.protobuf import message # pylint: disable=ungrouped-imports from apache_beam import pvalue from apache_beam.portability.api import metrics_pb2 from apache_beam.runners.sdf_utils import SplitResultPrimary from apache_beam.runners.sdf_utils import SplitResultResidual from apache_beam.runners.worker import data_plane from apache_beam.runners.worker import sdk_worker from apache_beam.transforms.core import Windowing from apache_beam.transforms.window import BoundedWindow from apache_beam.utils import windowed_value T = TypeVar('T') ConstructorFn = Callable[[ 'BeamTransformFactory', Any, beam_runner_api_pb2.PTransform, Union['message.Message', bytes], dict[str, list[operations.Operation]] ], operations.Operation] OperationT = TypeVar('OperationT', bound=operations.Operation) FnApiUserRuntimeStateTypes = Union['ReadModifyWriteRuntimeState', 'CombiningValueRuntimeState', 'SynchronousSetRuntimeState', 'SynchronousBagRuntimeState', 'SynchronousOrderedListRuntimeState'] DATA_INPUT_URN = 'beam:runner:source:v1' DATA_OUTPUT_URN = 'beam:runner:sink:v1' SYNTHETIC_DATA_SAMPLING_URN = 'beam:internal:sampling:v1' IDENTITY_DOFN_URN = 'beam:dofn:identity:0.1' # TODO(vikasrk): Fix this once runner sends appropriate common_urns. OLD_DATAFLOW_RUNNER_HARNESS_PARDO_URN = 'beam:dofn:javasdk:0.1' OLD_DATAFLOW_RUNNER_HARNESS_READ_URN = 'beam:source:java:0.1' URNS_NEEDING_PCOLLECTIONS = set([ monitoring_infos.ELEMENT_COUNT_URN, monitoring_infos.SAMPLED_BYTE_SIZE_URN ]) _LOGGER = logging.getLogger(__name__) class RunnerIOOperation(operations.Operation): """Common baseclass for runner harness IO operations.""" def __init__( self, name_context: common.NameContext, step_name: Any, consumers: Mapping[Any, Iterable[operations.Operation]], counter_factory: counters.CounterFactory, state_sampler: statesampler.StateSampler, windowed_coder: coders.Coder, transform_id: str, data_channel: data_plane.DataChannel) -> None: super().__init__(name_context, None, counter_factory, state_sampler) self.windowed_coder = windowed_coder self.windowed_coder_impl = windowed_coder.get_impl() # transform_id represents the consumer for the bytes in the data plane for a # DataInputOperation or a producer of these bytes for a DataOutputOperation. self.transform_id = transform_id self.data_channel = data_channel for _, consumer_ops in consumers.items(): for consumer in consumer_ops: self.add_receiver(consumer, 0) class DataOutputOperation(RunnerIOOperation): """A sink-like operation that gathers outputs to be sent back to the runner. """ def set_output_stream( self, output_stream: data_plane.ClosableOutputStream) -> None: self.output_stream = output_stream def process(self, windowed_value: windowed_value.WindowedValue) -> None: self.windowed_coder_impl.encode_to_stream( windowed_value, self.output_stream, True) self.output_stream.maybe_flush() def finish(self) -> None: super().finish() self.output_stream.close() class DataInputOperation(RunnerIOOperation): """A source-like operation that gathers input from the runner.""" def __init__( self, operation_name: common.NameContext, step_name, consumers: Mapping[Any, list[operations.Operation]], counter_factory: counters.CounterFactory, state_sampler: statesampler.StateSampler, windowed_coder: coders.Coder, transform_id, data_channel: data_plane.GrpcClientDataChannel) -> None: super().__init__( operation_name, step_name, consumers, counter_factory, state_sampler, windowed_coder, transform_id=transform_id, data_channel=data_channel) self.consumer = next(iter(consumers.values())) self.splitting_lock = threading.Lock() self.index = -1 self.stop = float('inf') self.started = False def setup(self, data_sampler=None): super().setup(data_sampler) # We must do this manually as we don't have a spec or spec.output_coders. self.receivers = [ operations.ConsumerSet.create( counter_factory=self.counter_factory, step_name=self.name_context.step_name, output_index=0, consumers=self.consumer, coder=self.windowed_coder, producer_type_hints=self._get_runtime_performance_hints(), producer_batch_converter=self.get_output_batch_converter()) ] def start(self) -> None: super().start() with self.splitting_lock: self.started = True def process(self, windowed_value: windowed_value.WindowedValue) -> None: self.output(windowed_value) def process_encoded(self, encoded_windowed_values: bytes) -> None: input_stream = coder_impl.create_InputStream(encoded_windowed_values) while input_stream.size() > 0: with self.splitting_lock: if self.index == self.stop - 1: return self.index += 1 try: decoded_value = self.windowed_coder_impl.decode_from_stream( input_stream, True) except Exception as exn: coder = str(self.windowed_coder) step = self.name_context.step_name raise ValueError( f"Error decoding input stream with coder {coder} in step {step}" ) from exn self.output(decoded_value) def monitoring_infos( self, transform_id: str, tag_to_pcollection_id: dict[str, str] ) -> dict[frozenset, metrics_pb2.MonitoringInfo]: all_monitoring_infos = super().monitoring_infos( transform_id, tag_to_pcollection_id) read_progress_info = monitoring_infos.int64_counter( monitoring_infos.DATA_CHANNEL_READ_INDEX, self.index, ptransform=transform_id) all_monitoring_infos[monitoring_infos.to_key( read_progress_info)] = read_progress_info return all_monitoring_infos # TODO(https://github.com/apache/beam/issues/19737): typing not compatible # with super type def try_split( # type: ignore[override] self, fraction_of_remainder, total_buffer_size, allowed_split_points ) -> Optional[ tuple[ int, Iterable[operations.SdfSplitResultsPrimary], Iterable[operations.SdfSplitResultsResidual], int]]: with self.splitting_lock: if not self.started: return None if self.index == -1: # We are "finished" with the (non-existent) previous element. current_element_progress = 1.0 else: current_element_progress_object = ( self.receivers[0].current_element_progress()) if current_element_progress_object is None: current_element_progress = 0.5 else: current_element_progress = ( current_element_progress_object.fraction_completed) # Now figure out where to split. split = self._compute_split( self.index, current_element_progress, self.stop, fraction_of_remainder, total_buffer_size, allowed_split_points, self.receivers[0].try_split) if split: self.stop = split[-1] return split @staticmethod def _compute_split( index, current_element_progress, stop, fraction_of_remainder, total_buffer_size, allowed_split_points=(), try_split=lambda fraction: None): def is_valid_split_point(index): return not allowed_split_points or index in allowed_split_points if total_buffer_size < index + 1: total_buffer_size = index + 1 elif total_buffer_size > stop: total_buffer_size = stop # The units here (except for keep_of_element_remainder) are all in # terms of number of (possibly fractional) elements. remainder = total_buffer_size - index - current_element_progress keep = remainder * fraction_of_remainder if current_element_progress < 1: keep_of_element_remainder = keep / (1 - current_element_progress) # If it's less than what's left of the current element, # try splitting at the current element. if (keep_of_element_remainder < 1 and is_valid_split_point(index) and is_valid_split_point(index + 1)): split: Optional[tuple[ Iterable[operations.SdfSplitResultsPrimary], Iterable[operations.SdfSplitResultsResidual]]] = try_split( keep_of_element_remainder) if split: element_primaries, element_residuals = split return index - 1, element_primaries, element_residuals, index + 1 # Otherwise, split at the closest element boundary. # pylint: disable=bad-option-value stop_index = index + max(1, int(round(current_element_progress + keep))) if allowed_split_points and stop_index not in allowed_split_points: # Choose the closest allowed split point. allowed_split_points = sorted(allowed_split_points) closest = bisect.bisect(allowed_split_points, stop_index) if closest == 0: stop_index = allowed_split_points[0] elif closest == len(allowed_split_points): stop_index = allowed_split_points[-1] else: prev = allowed_split_points[closest - 1] next = allowed_split_points[closest] if index < prev and stop_index - prev < next - stop_index: stop_index = prev else: stop_index = next if index < stop_index < stop: return stop_index - 1, [], [], stop_index else: return None def finish(self) -> None: super().finish() with self.splitting_lock: self.index += 1 self.started = False def reset(self) -> None: with self.splitting_lock: self.index = -1 self.stop = float('inf') super().reset() class _StateBackedIterable(object): def __init__( self, state_handler: sdk_worker.CachingStateHandler, state_key: beam_fn_api_pb2.StateKey, coder_or_impl: Union[coders.Coder, coder_impl.CoderImpl], ) -> None: self._state_handler = state_handler self._state_key = state_key if isinstance(coder_or_impl, coders.Coder): self._coder_impl = coder_or_impl.get_impl() else: self._coder_impl = coder_or_impl def __iter__(self) -> Iterator[Any]: return iter( self._state_handler.blocking_get(self._state_key, self._coder_impl)) def __reduce__(self): return list, (list(self), ) coder_impl.FastPrimitivesCoderImpl.register_iterable_like_type( _StateBackedIterable) class StateBackedSideInputMap(object): _BULK_READ_LIMIT = 100 _BULK_READ_FULLY = "fully" _BULK_READ_PARTIALLY = "partially" def __init__( self, state_handler: sdk_worker.CachingStateHandler, transform_id: str, tag: Optional[str], side_input_data: pvalue.SideInputData, coder: WindowedValueCoder, use_bulk_read: bool = False, ) -> None: self._state_handler = state_handler self._transform_id = transform_id self._tag = tag self._side_input_data = side_input_data self._element_coder = coder.wrapped_value_coder self._target_window_coder = coder.window_coder # TODO(robertwb): Limit the cache size. self._cache: dict[BoundedWindow, Any] = {} self._use_bulk_read = use_bulk_read def __getitem__(self, window): target_window = self._side_input_data.window_mapping_fn(window) if target_window not in self._cache: state_handler = self._state_handler access_pattern = self._side_input_data.access_pattern if access_pattern == common_urns.side_inputs.ITERABLE.urn: state_key = beam_fn_api_pb2.StateKey( iterable_side_input=beam_fn_api_pb2.StateKey.IterableSideInput( transform_id=self._transform_id, side_input_id=self._tag, window=self._target_window_coder.encode(target_window))) raw_view = _StateBackedIterable( state_handler, state_key, self._element_coder) elif access_pattern == common_urns.side_inputs.MULTIMAP.urn: state_key = beam_fn_api_pb2.StateKey( multimap_side_input=beam_fn_api_pb2.StateKey.MultimapSideInput( transform_id=self._transform_id, side_input_id=self._tag, window=self._target_window_coder.encode(target_window), key=b'')) kv_iter_state_key = beam_fn_api_pb2.StateKey( multimap_keys_values_side_input=beam_fn_api_pb2.StateKey. MultimapKeysValuesSideInput( transform_id=self._transform_id, side_input_id=self._tag, window=self._target_window_coder.encode(target_window))) cache = {} key_coder = self._element_coder.key_coder() key_coder_impl = key_coder.get_impl() value_coder = self._element_coder.value_coder() use_bulk_read = self._use_bulk_read class MultiMap(object): _bulk_read = None _lock = threading.Lock() def __getitem__(self, key): if use_bulk_read: if self._bulk_read is None: with self._lock: if self._bulk_read is None: try: # Attempt to bulk read the key-values over the iterable # protocol which, if supported, can be much more efficient # than point lookups if it fits into memory. for ix, (k, vs) in enumerate(_StateBackedIterable( state_handler, kv_iter_state_key, coders.TupleCoder( (key_coder, coders.IterableCoder(value_coder))))): cache[k] = vs if ix > StateBackedSideInputMap._BULK_READ_LIMIT: self._bulk_read = ( StateBackedSideInputMap._BULK_READ_PARTIALLY) break else: # We reached the end of the iteration without breaking. self._bulk_read = ( StateBackedSideInputMap._BULK_READ_FULLY) except Exception: _LOGGER.error( "Iterable access of map side inputs unsupported.", exc_info=True) self._bulk_read = ( StateBackedSideInputMap._BULK_READ_PARTIALLY) if (self._bulk_read == StateBackedSideInputMap._BULK_READ_FULLY): return cache.get(key, []) if key not in cache: keyed_state_key = beam_fn_api_pb2.StateKey() keyed_state_key.CopyFrom(state_key) keyed_state_key.multimap_side_input.key = ( key_coder_impl.encode_nested(key)) cache[key] = _StateBackedIterable( state_handler, keyed_state_key, value_coder) return cache[key] def __reduce__(self): # TODO(robertwb): Figure out how to support this. raise TypeError(common_urns.side_inputs.MULTIMAP.urn) raw_view = MultiMap() else: raise ValueError("Unknown access pattern: '%s'" % access_pattern) self._cache[target_window] = self._side_input_data.view_fn(raw_view) return self._cache[target_window] def is_globally_windowed(self) -> bool: return ( self._side_input_data.window_mapping_fn == sideinputs._global_window_mapping_fn) def reset(self) -> None: # TODO(BEAM-5428): Cross-bundle caching respecting cache tokens. self._cache = {} class ReadModifyWriteRuntimeState(userstate.ReadModifyWriteRuntimeState): def __init__(self, underlying_bag_state): self._underlying_bag_state = underlying_bag_state def read(self) -> Any: values = list(self._underlying_bag_state.read()) if not values: return None return values[0] def write(self, value: Any) -> None: self.clear() self._underlying_bag_state.add(value) def clear(self) -> None: self._underlying_bag_state.clear() def commit(self) -> None: self._underlying_bag_state.commit() class CombiningValueRuntimeState(userstate.CombiningValueRuntimeState): def __init__( self, underlying_bag_state: userstate.AccumulatingRuntimeState, combinefn: core.CombineFn) -> None: self._combinefn = combinefn self._combinefn.setup() self._underlying_bag_state = underlying_bag_state self._finalized = False def _read_accumulator(self, rewrite=True): merged_accumulator = self._combinefn.merge_accumulators( self._underlying_bag_state.read()) if rewrite: self._underlying_bag_state.clear() self._underlying_bag_state.add(merged_accumulator) return merged_accumulator def read(self) -> Iterable[Any]: return self._combinefn.extract_output(self._read_accumulator()) def add(self, value: Any) -> None: # Prefer blind writes, but don't let them grow unboundedly. # This should be tuned to be much lower, but for now exercise # both paths well. if random.random() < 0.5: accumulator = self._read_accumulator(False) self._underlying_bag_state.clear() else: accumulator = self._combinefn.create_accumulator() self._underlying_bag_state.add( self._combinefn.add_input(accumulator, value)) def clear(self) -> None: self._underlying_bag_state.clear() def commit(self): self._underlying_bag_state.commit() def finalize(self): if not self._finalized: self._combinefn.teardown() self._finalized = True class _ConcatIterable(object): """An iterable that is the concatination of two iterables. Unlike itertools.chain, this allows reiteration. """ def __init__(self, first: Iterable[Any], second: Iterable[Any]) -> None: self.first = first self.second = second def __iter__(self) -> Iterator[Any]: for elem in self.first: yield elem for elem in self.second: yield elem coder_impl.FastPrimitivesCoderImpl.register_iterable_like_type(_ConcatIterable) class SynchronousBagRuntimeState(userstate.BagRuntimeState): def __init__( self, state_handler: sdk_worker.CachingStateHandler, state_key: beam_fn_api_pb2.StateKey, value_coder: coders.Coder) -> None: self._state_handler = state_handler self._state_key = state_key self._value_coder = value_coder self._cleared = False self._added_elements: list[Any] = [] def read(self) -> Iterable[Any]: return _ConcatIterable([] if self._cleared else cast( 'Iterable[Any]', _StateBackedIterable( self._state_handler, self._state_key, self._value_coder)), self._added_elements) def add(self, value: Any) -> None: self._added_elements.append(value) def clear(self) -> None: self._cleared = True self._added_elements = [] def commit(self) -> None: to_await = None if self._cleared: to_await = self._state_handler.clear(self._state_key) if self._added_elements: to_await = self._state_handler.extend( self._state_key, self._value_coder.get_impl(), self._added_elements) if to_await: # To commit, we need to wait on the last state request future to complete. to_await.get() class SynchronousSetRuntimeState(userstate.SetRuntimeState): def __init__( self, state_handler: sdk_worker.CachingStateHandler, state_key: beam_fn_api_pb2.StateKey, value_coder: coders.Coder) -> None: self._state_handler = state_handler self._state_key = state_key self._value_coder = value_coder self._cleared = False self._added_elements: set[Any] = set() def _compact_data(self, rewrite=True): accumulator = set( _ConcatIterable( set() if self._cleared else _StateBackedIterable( self._state_handler, self._state_key, self._value_coder), self._added_elements)) if rewrite and accumulator: self._state_handler.clear(self._state_key) self._state_handler.extend( self._state_key, self._value_coder.get_impl(), accumulator) # Since everthing is already committed so we can safely reinitialize # added_elements here. self._added_elements = set() return accumulator def read(self) -> set[Any]: return self._compact_data(rewrite=False) def add(self, value: Any) -> None: if self._cleared: # This is a good time explicitly clear. self._state_handler.clear(self._state_key) self._cleared = False self._added_elements.add(value) if random.random() > 0.5: self._compact_data() def clear(self) -> None: self._cleared = True self._added_elements = set() def commit(self) -> None: to_await = None if self._cleared: to_await = self._state_handler.clear(self._state_key) if self._added_elements: to_await = self._state_handler.extend( self._state_key, self._value_coder.get_impl(), self._added_elements) if to_await: # To commit, we need to wait on the last state request future to complete. to_await.get() class RangeSet: """For Internal Use only. A simple range set for ranges of [x,y).""" def __init__(self) -> None: # The start points and end points are stored separately in order. self._sorted_starts = SortedList() self._sorted_ends = SortedList() def add(self, start: int, end: int) -> None: if start >= end: return # ranges[:min_idx] and ranges[max_idx:] is unaffected by this insertion # the first range whose end point >= the start of the new range min_idx = self._sorted_ends.bisect_left(start) # the first range whose start point > the end point of the new range max_idx = self._sorted_starts.bisect_right(end) if min_idx >= len(self._sorted_starts) or max_idx <= 0: # the new range is beyond any current ranges new_start = start new_end = end else: # the new range overlaps with ranges[min_idx:max_idx] new_start = min(start, self._sorted_starts[min_idx]) new_end = max(end, self._sorted_ends[max_idx - 1]) del self._sorted_starts[min_idx:max_idx] del self._sorted_ends[min_idx:max_idx] self._sorted_starts.add(new_start) self._sorted_ends.add(new_end) def __contains__(self, key: int) -> bool: idx = self._sorted_starts.bisect_left(key) return (idx < len(self._sorted_starts) and self._sorted_starts[idx] == key) or (idx > 0 and self._sorted_ends[idx - 1] > key) def __len__(self) -> int: assert len(self._sorted_starts) == len(self._sorted_ends) return len(self._sorted_starts) def __iter__(self) -> Iterator[tuple[int, int]]: return zip(self._sorted_starts, self._sorted_ends) def __str__(self) -> str: return str(list(zip(self._sorted_starts, self._sorted_ends))) class SynchronousOrderedListRuntimeState(userstate.OrderedListRuntimeState): RANGE_MIN = -(1 << 63) RANGE_MAX = (1 << 63) - 1 TIMESTAMP_RANGE_MIN = timestamp.Timestamp(micros=RANGE_MIN) TIMESTAMP_RANGE_MAX = timestamp.Timestamp(micros=RANGE_MAX) def __init__( self, state_handler: sdk_worker.CachingStateHandler, state_key: beam_fn_api_pb2.StateKey, value_coder: coders.Coder) -> None: self._state_handler = state_handler self._state_key = state_key self._elem_coder = beam.coders.TupleCoder( [coders.VarIntCoder(), coders.coders.LengthPrefixCoder(value_coder)]) self._cleared = False self._pending_adds = SortedDict() self._pending_removes = RangeSet() def add(self, elem: tuple[timestamp.Timestamp, Any]) -> None: assert len(elem) == 2 key_ts, value = elem key = key_ts.micros if key >= self.RANGE_MAX or key < self.RANGE_MIN: raise ValueError("key value %d is out of range" % key) self._pending_adds.setdefault(key, []).append(value) def read(self) -> Iterable[tuple[timestamp.Timestamp, Any]]: return self.read_range(self.TIMESTAMP_RANGE_MIN, self.TIMESTAMP_RANGE_MAX) def read_range( self, min_timestamp: timestamp.Timestamp, limit_timestamp: timestamp.Timestamp ) -> Iterable[tuple[timestamp.Timestamp, Any]]: # convert timestamp to int, as sort keys are stored as int internally. min_key = min_timestamp.micros limit_key = limit_timestamp.micros keys_to_add = self._pending_adds.irange( min_key, limit_key, inclusive=(True, False)) # use list interpretation here to construct the actual list # of iterators of the selected range. local_items = chain.from_iterable([ itertools.islice( zip(itertools.cycle([ k, ]), self._pending_adds[k]), len(self._pending_adds[k])) for k in keys_to_add ]) if not self._cleared: range_query_state_key = beam_fn_api_pb2.StateKey() range_query_state_key.CopyFrom(self._state_key) range_query_state_key.ordered_list_user_state.range.start = min_key range_query_state_key.ordered_list_user_state.range.end = limit_key # make a deep copy here because there could be other operations occur in # the middle of an iteration and change pending_removes pending_removes_snapshot = copy.deepcopy(self._pending_removes) persistent_items = filter( lambda kv: kv[0] not in pending_removes_snapshot, _StateBackedIterable( self._state_handler, range_query_state_key, self._elem_coder)) return map( lambda x: (timestamp.Timestamp(micros=x[0]), x[1]), heapq.merge(persistent_items, local_items)) return map(lambda x: (timestamp.Timestamp(micros=x[0]), x[1]), local_items) def clear(self) -> None: self._cleared = True self._pending_adds = SortedDict() self._pending_removes = RangeSet() self._pending_removes.add(self.RANGE_MIN, self.RANGE_MAX) def clear_range( self, min_timestamp: timestamp.Timestamp, limit_timestamp: timestamp.Timestamp) -> None: min_key = min_timestamp.micros limit_key = limit_timestamp.micros # materialize the keys to remove before the actual removal keys_to_remove = list( self._pending_adds.irange(min_key, limit_key, inclusive=(True, False))) for k in keys_to_remove: del self._pending_adds[k] if not self._cleared: self._pending_removes.add(min_key, limit_key) def commit(self) -> None: futures = [] if self._pending_removes: for start, end in self._pending_removes: range_query_state_key = beam_fn_api_pb2.StateKey() range_query_state_key.CopyFrom(self._state_key) range_query_state_key.ordered_list_user_state.range.start = start range_query_state_key.ordered_list_user_state.range.end = end futures.append(self._state_handler.clear(range_query_state_key)) self._pending_removes = RangeSet() if self._pending_adds: items_to_add = [] for k in self._pending_adds: items_to_add.extend(zip(itertools.cycle([ k, ]), self._pending_adds[k])) futures.append( self._state_handler.extend( self._state_key, self._elem_coder.get_impl(), items_to_add)) self._pending_adds = SortedDict() if len(futures): # To commit, we need to wait on every state request futures to complete. for to_await in futures: to_await.get() self._cleared = False class OutputTimer(userstate.BaseTimer): def __init__( self, key, window: BoundedWindow, timestamp: timestamp.Timestamp, paneinfo: windowed_value.PaneInfo, time_domain: str, timer_family_id: str, timer_coder_impl: coder_impl.TimerCoderImpl, output_stream: data_plane.ClosableOutputStream): self._key = key self._window = window self._input_timestamp = timestamp self._paneinfo = paneinfo self._time_domain = time_domain self._timer_family_id = timer_family_id self._output_stream = output_stream self._timer_coder_impl = timer_coder_impl def set(self, ts: timestamp.TimestampTypes, dynamic_timer_tag='') -> None: ts = timestamp.Timestamp.of(ts) timer = userstate.Timer( user_key=self._key, dynamic_timer_tag=dynamic_timer_tag, windows=(self._window, ), clear_bit=False, fire_timestamp=ts, hold_timestamp=ts if TimeDomain.is_event_time(self._time_domain) else self._input_timestamp, paneinfo=self._paneinfo) self._timer_coder_impl.encode_to_stream(timer, self._output_stream, True) self._output_stream.maybe_flush() def clear(self, dynamic_timer_tag='') -> None: timer = userstate.Timer( user_key=self._key, dynamic_timer_tag=dynamic_timer_tag, windows=(self._window, ), clear_bit=True, fire_timestamp=None, hold_timestamp=None, paneinfo=None) self._timer_coder_impl.encode_to_stream(timer, self._output_stream, True) self._output_stream.maybe_flush() class TimerInfo(object): """A data class to store information related to a timer.""" def __init__(self, timer_coder_impl, output_stream=None): self.timer_coder_impl = timer_coder_impl self.output_stream = output_stream class FnApiUserStateContext(userstate.UserStateContext): """Interface for state and timers from SDK to Fn API servicer of state..""" def __init__( self, state_handler: sdk_worker.CachingStateHandler, transform_id: str, key_coder: coders.Coder, window_coder: coders.Coder, ) -> None: """Initialize a ``FnApiUserStateContext``. Args: state_handler: A StateServicer object. transform_id: The name of the PTransform that this context is associated. key_coder: Coder for the key type. window_coder: Coder for the window type. """ self._state_handler = state_handler self._transform_id = transform_id self._key_coder = key_coder self._window_coder = window_coder # A mapping of {timer_family_id: TimerInfo} self._timers_info: dict[str, TimerInfo] = {} self._all_states: dict[tuple, FnApiUserRuntimeStateTypes] = {} def add_timer_info(self, timer_family_id: str, timer_info: TimerInfo) -> None: self._timers_info[timer_family_id] = timer_info def get_timer( self, timer_spec: userstate.TimerSpec, key, window, timestamp, pane) -> OutputTimer: assert self._timers_info[timer_spec.name].output_stream is not None timer_coder_impl = self._timers_info[timer_spec.name].timer_coder_impl output_stream = self._timers_info[timer_spec.name].output_stream return OutputTimer( key, window, timestamp, pane, timer_spec.time_domain, timer_spec.name, timer_coder_impl, output_stream) def get_state(self, *args: Any) -> FnApiUserRuntimeStateTypes: state_handle = self._all_states.get(args) if state_handle is None: state_handle = self._all_states[args] = self._create_state(*args) return state_handle def _create_state( self, state_spec: userstate.StateSpec, key, window: BoundedWindow) -> FnApiUserRuntimeStateTypes: if isinstance(state_spec, (userstate.BagStateSpec, userstate.CombiningValueStateSpec, userstate.ReadModifyWriteStateSpec)): bag_state = SynchronousBagRuntimeState( self._state_handler, state_key=beam_fn_api_pb2.StateKey( bag_user_state=beam_fn_api_pb2.StateKey.BagUserState( transform_id=self._transform_id, user_state_id=state_spec.name, window=self._window_coder.encode(window), # State keys are expected in nested encoding format key=self._key_coder.encode_nested(key))), value_coder=state_spec.coder) if isinstance(state_spec, userstate.BagStateSpec): return bag_state elif isinstance(state_spec, userstate.ReadModifyWriteStateSpec): return ReadModifyWriteRuntimeState(bag_state) else: return CombiningValueRuntimeState( bag_state, copy.deepcopy(state_spec.combine_fn)) elif isinstance(state_spec, userstate.SetStateSpec): return SynchronousSetRuntimeState( self._state_handler, state_key=beam_fn_api_pb2.StateKey( bag_user_state=beam_fn_api_pb2.StateKey.BagUserState( transform_id=self._transform_id, user_state_id=state_spec.name, window=self._window_coder.encode(window), # State keys are expected in nested encoding format key=self._key_coder.encode_nested(key))), value_coder=state_spec.coder) elif isinstance(state_spec, userstate.OrderedListStateSpec): return SynchronousOrderedListRuntimeState( self._state_handler, state_key=beam_fn_api_pb2.StateKey( ordered_list_user_state=beam_fn_api_pb2.StateKey. OrderedListUserState( transform_id=self._transform_id, user_state_id=state_spec.name, window=self._window_coder.encode(window), key=self._key_coder.encode_nested(key))), value_coder=state_spec.coder) else: raise NotImplementedError(state_spec) def commit(self) -> None: for state in self._all_states.values(): state.commit() def reset(self) -> None: for state in self._all_states.values(): state.finalize() self._all_states = {} def memoize(func): cache = {} missing = object() def wrapper(*args): result = cache.get(args, missing) if result is missing: result = cache[args] = func(*args) return result return wrapper def only_element(iterable: Iterable[T]) -> T: element, = iterable return element def _environments_compatible(submission: str, runtime: str) -> bool: if submission == runtime: return True if 'rc' in submission and runtime in submission: # TODO(https://github.com/apache/beam/issues/28084): Loosen # the check for RCs until RC containers install the matching version. return True return False def _verify_descriptor_created_in_a_compatible_env( process_bundle_descriptor: beam_fn_api_pb2.ProcessBundleDescriptor) -> None: runtime_sdk = environments.sdk_base_version_capability() for t in process_bundle_descriptor.transforms.values(): env = process_bundle_descriptor.environments[t.environment_id] for c in env.capabilities: if (c.startswith(environments.SDK_VERSION_CAPABILITY_PREFIX) and not _environments_compatible(c, runtime_sdk)): raise RuntimeError( "Pipeline construction environment and pipeline runtime " "environment are not compatible. If you use a custom " "container image, check that the Python interpreter minor version " "and the Apache Beam version in your image match the versions " "used at pipeline construction time. " f"Submission environment: {c}. " f"Runtime environment: {runtime_sdk}.") # TODO: Consider warning on mismatches in versions of installed packages. class BundleProcessor(object): """ A class for processing bundles of elements. """ def __init__( self, runner_capabilities: frozenset[str], process_bundle_descriptor: beam_fn_api_pb2.ProcessBundleDescriptor, state_handler: sdk_worker.CachingStateHandler, data_channel_factory: data_plane.DataChannelFactory, data_sampler: Optional[data_sampler.DataSampler] = None, ) -> None: """Initialize a bundle processor. Args: runner_capabilities (``FrozenSet[str]``): The set of capabilities of the runner with which we will be interacting process_bundle_descriptor (``beam_fn_api_pb2.ProcessBundleDescriptor``): a description of the stage that this ``BundleProcessor``is to execute. state_handler (CachingStateHandler). data_channel_factory (``data_plane.DataChannelFactory``). """ self.runner_capabilities = runner_capabilities self.process_bundle_descriptor = process_bundle_descriptor self.state_handler = state_handler self.data_channel_factory = data_channel_factory self.data_sampler = data_sampler self.current_instruction_id: Optional[str] = None # Represents whether the SDK is consuming received data. self.consuming_received_data = False _verify_descriptor_created_in_a_compatible_env(process_bundle_descriptor) # There is no guarantee that the runner only set # timer_api_service_descriptor when having timers. So this field cannot be # used as an indicator of timers. if self.process_bundle_descriptor.timer_api_service_descriptor.url: self.timer_data_channel = ( data_channel_factory.create_data_channel_from_url( self.process_bundle_descriptor.timer_api_service_descriptor.url)) else: self.timer_data_channel = None # A mapping of # {(transform_id, timer_family_id): TimerInfo} # The mapping is empty when there is no timer_family_specs in the # ProcessBundleDescriptor. self.timers_info: dict[tuple[str, str], TimerInfo] = {} # TODO(robertwb): Figure out the correct prefix to use for output counters # from StateSampler. self.counter_factory = counters.CounterFactory() self.state_sampler = statesampler.StateSampler( 'fnapi-step-%s' % self.process_bundle_descriptor.id, self.counter_factory) with concurrent.futures.ThreadPoolExecutor( max_workers=1, thread_name_prefix='ExecutionTreeCreator') as executor: future = executor.submit( self.create_execution_tree, self.process_bundle_descriptor) try: self.ops = future.result(timeout=3600) except concurrent.futures.TimeoutError: # In rare cases, unpickling a DoFn might get permanently stuck, # for example when unpickling involves importing a module and # a subprocess is launched during the import operation. _LOGGER.error( 'Timed out while reconstructing a pipeline fragment for: %s.\n' 'This is likely a transient error. The SDK harness ' 'will self-terminate, and the runner can retry the operation. ' 'If the error is frequent, check whether the stuckness happens ' 'while deserializing (unpickling) a dependency of your pipeline ' 'in the stacktrace below: \n%s\n', self.process_bundle_descriptor.id, thread_dump('ExecutionTreeCreator')) # Raising an exception here doesn't interrupt the left-over thread. # Out of caution, terminate the SDK harness process. from apache_beam.runners.worker.sdk_worker_main import terminate_sdk_harness terminate_sdk_harness() for op in reversed(self.ops.values()): op.setup(self.data_sampler) self.splitting_lock = threading.Lock() def create_execution_tree( self, descriptor: beam_fn_api_pb2.ProcessBundleDescriptor ) -> collections.OrderedDict[str, operations.DoOperation]: transform_factory = BeamTransformFactory( self.runner_capabilities, descriptor, self.data_channel_factory, self.counter_factory, self.state_sampler, self.state_handler, self.data_sampler, ) self.timers_info = transform_factory.extract_timers_info() def is_side_input(transform_proto, tag): if transform_proto.spec.urn == common_urns.primitives.PAR_DO.urn: return tag in proto_utils.parse_Bytes( transform_proto.spec.payload, beam_runner_api_pb2.ParDoPayload).side_inputs pcoll_consumers: collections.defaultdict[ str, list[str]] = collections.defaultdict(list) for transform_id, transform_proto in descriptor.transforms.items(): for tag, pcoll_id in transform_proto.inputs.items(): if not is_side_input(transform_proto, tag): pcoll_consumers[pcoll_id].append(transform_id) @memoize def get_operation(transform_id: str) -> operations.Operation: transform_consumers = { tag: [get_operation(op) for op in pcoll_consumers[pcoll_id]] for tag, pcoll_id in descriptor.transforms[transform_id].outputs.items() } # Initialize transform-specific state in the Data Sampler. if self.data_sampler: self.data_sampler.initialize_samplers( transform_id, descriptor, transform_factory.get_coder) return transform_factory.create_operation( transform_id, transform_consumers) # Operations must be started (hence returned) in order. @memoize def topological_height(transform_id: str) -> int: return 1 + max([0] + [ topological_height(consumer) for pcoll in descriptor.transforms[transform_id].outputs.values() for consumer in pcoll_consumers[pcoll] ]) return collections.OrderedDict([( transform_id, cast(operations.DoOperation, get_operation(transform_id))) for transform_id in sorted( descriptor.transforms, key=topological_height, reverse=True)]) def reset(self) -> None: self.counter_factory.reset() self.state_sampler.reset() # Side input caches. for op in self.ops.values(): op.reset() def process_bundle( self, instruction_id: str ) -> tuple[list[beam_fn_api_pb2.DelayedBundleApplication], bool]: expected_input_ops: list[DataInputOperation] = [] for op in self.ops.values(): if isinstance(op, DataOutputOperation): # TODO(robertwb): Is there a better way to pass the instruction id to # the operation? op.set_output_stream( op.data_channel.output_stream(instruction_id, op.transform_id)) elif isinstance(op, DataInputOperation): # We must wait until we receive "end of stream" for each of these ops. expected_input_ops.append(op) try: execution_context = ExecutionContext(instruction_id=instruction_id) self.current_instruction_id = instruction_id self.state_sampler.start() # Start all operations. for op in reversed(self.ops.values()): _LOGGER.debug('start %s', op) op.execution_context = execution_context op.start() # Each data_channel is mapped to a list of expected inputs which includes # both data input and timer input. The data input is identied by # transform_id. The data input is identified by # (transform_id, timer_family_id). data_channels: collections.defaultdict[ data_plane.DataChannel, list[Union[str, tuple[str, str]]]] = collections.defaultdict(list) # Add expected data inputs for each data channel. input_op_by_transform_id = {} for input_op in expected_input_ops: data_channels[input_op.data_channel].append(input_op.transform_id) input_op_by_transform_id[input_op.transform_id] = input_op # Update timer_data channel with expected timer inputs. if self.timer_data_channel: data_channels[self.timer_data_channel].extend( list(self.timers_info.keys())) # Set up timer output stream for DoOperation. for ((transform_id, timer_family_id), timer_info) in self.timers_info.items(): output_stream = self.timer_data_channel.output_timer_stream( instruction_id, transform_id, timer_family_id) timer_info.output_stream = output_stream self.ops[transform_id].add_timer_info(timer_family_id, timer_info) # Process data and timer inputs # We are currently not consuming received data. self.consuming_received_data = False for data_channel, expected_inputs in data_channels.items(): for element in data_channel.input_elements(instruction_id, expected_inputs): # Since we have received a set of elements and are consuming it. self.consuming_received_data = True if isinstance(element, beam_fn_api_pb2.Elements.Timers): timer_coder_impl = ( self.timers_info[( element.transform_id, element.timer_family_id)].timer_coder_impl) for timer_data in timer_coder_impl.decode_all(element.timers): self.ops[element.transform_id].process_timer( element.timer_family_id, timer_data) elif isinstance(element, beam_fn_api_pb2.Elements.Data): input_op_by_transform_id[element.transform_id].process_encoded( element.data) # We are done consuming the set of elements. self.consuming_received_data = False # Finish all operations. for op in self.ops.values(): _LOGGER.debug('finish %s', op) op.finish() # Close every timer output stream for timer_info in self.timers_info.values(): assert timer_info.output_stream is not None timer_info.output_stream.close() return ([ self.delayed_bundle_application(op, residual) for op, residual in execution_context.delayed_applications ], self.requires_finalization()) finally: self.consuming_received_data = False # Ensure any in-flight split attempts complete. with self.splitting_lock: self.current_instruction_id = None self.state_sampler.stop_if_still_running() def finalize_bundle(self) -> beam_fn_api_pb2.FinalizeBundleResponse: for op in self.ops.values(): op.finalize_bundle() return beam_fn_api_pb2.FinalizeBundleResponse() def requires_finalization(self) -> bool: return any(op.needs_finalization() for op in self.ops.values()) def try_split( self, bundle_split_request: beam_fn_api_pb2.ProcessBundleSplitRequest ) -> beam_fn_api_pb2.ProcessBundleSplitResponse: split_response = beam_fn_api_pb2.ProcessBundleSplitResponse() with self.splitting_lock: if bundle_split_request.instruction_id != self.current_instruction_id: # This may be a delayed split for a former bundle, see BEAM-12475. return split_response for op in self.ops.values(): if isinstance(op, DataInputOperation): desired_split = bundle_split_request.desired_splits.get( op.transform_id) if desired_split: split = op.try_split( desired_split.fraction_of_remainder, desired_split.estimated_input_elements, desired_split.allowed_split_points) if split: ( primary_end, element_primaries, element_residuals, residual_start, ) = split for element_primary in element_primaries: split_response.primary_roots.add().CopyFrom( self.bundle_application(*element_primary)) for element_residual in element_residuals: split_response.residual_roots.add().CopyFrom( self.delayed_bundle_application(*element_residual)) split_response.channel_splits.extend([ beam_fn_api_pb2.ProcessBundleSplitResponse.ChannelSplit( transform_id=op.transform_id, last_primary_element=primary_end, first_residual_element=residual_start) ]) return split_response def delayed_bundle_application( self, op: operations.DoOperation, deferred_remainder: SplitResultResidual ) -> beam_fn_api_pb2.DelayedBundleApplication: assert op.input_info is not None # TODO(SDF): For non-root nodes, need main_input_coder + residual_coder. (element_and_restriction, current_watermark, deferred_timestamp) = ( deferred_remainder) if deferred_timestamp: assert isinstance(deferred_timestamp, timestamp.Duration) proto_deferred_watermark: Optional[ duration_pb2.Duration] = proto_utils.from_micros( duration_pb2.Duration, deferred_timestamp.micros) else: proto_deferred_watermark = None return beam_fn_api_pb2.DelayedBundleApplication( requested_time_delay=proto_deferred_watermark, application=self.construct_bundle_application( op.input_info, current_watermark, element_and_restriction)) def bundle_application( self, op: operations.DoOperation, primary: SplitResultPrimary) -> beam_fn_api_pb2.BundleApplication: assert op.input_info is not None return self.construct_bundle_application( op.input_info, None, primary.primary_value) def construct_bundle_application( self, op_input_info: operations.OpInputInfo, output_watermark: Optional[timestamp.Timestamp], element) -> beam_fn_api_pb2.BundleApplication: transform_id, main_input_tag, main_input_coder, outputs = op_input_info if output_watermark: proto_output_watermark = proto_utils.from_micros( timestamp_pb2.Timestamp, output_watermark.micros) output_watermarks: Optional[dict[str, timestamp_pb2.Timestamp]] = { output: proto_output_watermark for output in outputs } else: output_watermarks = None return beam_fn_api_pb2.BundleApplication( transform_id=transform_id, input_id=main_input_tag, output_watermarks=output_watermarks, element=main_input_coder.get_impl().encode_nested(element)) def monitoring_infos(self) -> list[metrics_pb2.MonitoringInfo]: """Returns the list of MonitoringInfos collected processing this bundle.""" # Construct a new dict first to remove duplicates. all_monitoring_infos_dict = {} for transform_id, op in self.ops.items(): tag_to_pcollection_id = self.process_bundle_descriptor.transforms[ transform_id].outputs all_monitoring_infos_dict.update( op.monitoring_infos(transform_id, dict(tag_to_pcollection_id))) return list(all_monitoring_infos_dict.values()) def shutdown(self) -> None: for op in self.ops.values(): op.teardown() @dataclass class ExecutionContext: # Any splits to be processed later. delayed_applications: list[tuple[operations.DoOperation, common.SplitResultResidual]] = field( default_factory=list) # The exception sampler for the currently executing PTransform. output_sampler: Optional[data_sampler.OutputSampler] = None # The current instruction being executed. instruction_id: Optional[str] = None class BeamTransformFactory(object): """Factory for turning transform_protos into executable operations.""" def __init__( self, runner_capabilities: frozenset[str], descriptor: beam_fn_api_pb2.ProcessBundleDescriptor, data_channel_factory: data_plane.DataChannelFactory, counter_factory: counters.CounterFactory, state_sampler: statesampler.StateSampler, state_handler: sdk_worker.CachingStateHandler, data_sampler: Optional[data_sampler.DataSampler], ): self.runner_capabilities = runner_capabilities self.descriptor = descriptor self.data_channel_factory = data_channel_factory self.counter_factory = counter_factory self.state_sampler = state_sampler self.state_handler = state_handler self.context = pipeline_context.PipelineContext( descriptor, iterable_state_read=lambda token, element_coder_impl: _StateBackedIterable( state_handler, beam_fn_api_pb2.StateKey( runner=beam_fn_api_pb2.StateKey.Runner(key=token)), element_coder_impl)) self.data_sampler = data_sampler _known_urns: dict[str, tuple[ConstructorFn, Union[type[message.Message], type[bytes], None]]] = {} @classmethod def register_urn( cls, urn: str, parameter_type: Optional[type[T]] ) -> Callable[[ Callable[[ BeamTransformFactory, str, beam_runner_api_pb2.PTransform, T, dict[str, list[operations.Operation]] ], operations.Operation] ], Callable[[ BeamTransformFactory, str, beam_runner_api_pb2.PTransform, T, dict[str, list[operations.Operation]] ], operations.Operation]]: def wrapper(func): cls._known_urns[urn] = func, parameter_type return func return wrapper def create_operation( self, transform_id: str, consumers: dict[str, list[operations.Operation]]) -> operations.Operation: transform_proto = self.descriptor.transforms[transform_id] if not transform_proto.unique_name: _LOGGER.debug("No unique name set for transform %s" % transform_id) transform_proto.unique_name = transform_id creator, parameter_type = self._known_urns[transform_proto.spec.urn] payload = proto_utils.parse_Bytes( transform_proto.spec.payload, parameter_type) return creator(self, transform_id, transform_proto, payload, consumers) def extract_timers_info(self) -> dict[tuple[str, str], TimerInfo]: timers_info = {} for transform_id, transform_proto in self.descriptor.transforms.items(): if transform_proto.spec.urn == common_urns.primitives.PAR_DO.urn: pardo_payload = proto_utils.parse_Bytes( transform_proto.spec.payload, beam_runner_api_pb2.ParDoPayload) for (timer_family_id, timer_family_spec) in pardo_payload.timer_family_specs.items(): timer_coder_impl = self.get_coder( timer_family_spec.timer_family_coder_id).get_impl() # The output_stream should be updated when processing a bundle. timers_info[(transform_id, timer_family_id)] = TimerInfo( timer_coder_impl=timer_coder_impl) return timers_info def get_coder(self, coder_id: str) -> coders.Coder: if coder_id not in self.descriptor.coders: raise KeyError("No such coder: %s" % coder_id) coder_proto = self.descriptor.coders[coder_id] if coder_proto.spec.urn: return self.context.coders.get_by_id(coder_id) else: # No URN, assume cloud object encoding json bytes. return operation_specs.get_coder_from_spec( json.loads(coder_proto.spec.payload.decode('utf-8'))) def get_windowed_coder(self, pcoll_id: str) -> WindowedValueCoder: coder = self.get_coder(self.descriptor.pcollections[pcoll_id].coder_id) # TODO(robertwb): Remove this condition once all runners are consistent. if not isinstance(coder, WindowedValueCoder): windowing_strategy = self.descriptor.windowing_strategies[ self.descriptor.pcollections[pcoll_id].windowing_strategy_id] return WindowedValueCoder( coder, self.get_coder(windowing_strategy.window_coder_id)) else: return coder def get_output_coders( self, transform_proto: beam_runner_api_pb2.PTransform ) -> dict[str, coders.Coder]: return { tag: self.get_windowed_coder(pcoll_id) for tag, pcoll_id in transform_proto.outputs.items() } def get_only_output_coder( self, transform_proto: beam_runner_api_pb2.PTransform) -> coders.Coder: return only_element(self.get_output_coders(transform_proto).values()) def get_input_coders( self, transform_proto: beam_runner_api_pb2.PTransform ) -> dict[str, coders.WindowedValueCoder]: return { tag: self.get_windowed_coder(pcoll_id) for tag, pcoll_id in transform_proto.inputs.items() } def get_only_input_coder( self, transform_proto: beam_runner_api_pb2.PTransform) -> coders.Coder: return only_element(list(self.get_input_coders(transform_proto).values())) def get_input_windowing( self, transform_proto: beam_runner_api_pb2.PTransform) -> Windowing: pcoll_id = only_element(transform_proto.inputs.values()) windowing_strategy_id = self.descriptor.pcollections[ pcoll_id].windowing_strategy_id return self.context.windowing_strategies.get_by_id(windowing_strategy_id) # TODO(robertwb): Update all operations to take these in the constructor. @staticmethod def augment_oldstyle_op( op: OperationT, step_name: str, consumers: Mapping[str, Iterable[operations.Operation]], tag_list: Optional[list[str]] = None) -> OperationT: op.step_name = step_name for tag, op_consumers in consumers.items(): for consumer in op_consumers: op.add_receiver(consumer, tag_list.index(tag) if tag_list else 0) return op @BeamTransformFactory.register_urn( DATA_INPUT_URN, beam_fn_api_pb2.RemoteGrpcPort) def create_source_runner( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, grpc_port: beam_fn_api_pb2.RemoteGrpcPort, consumers: dict[str, list[operations.Operation]]) -> DataInputOperation: output_coder = factory.get_coder(grpc_port.coder_id) return DataInputOperation( common.NameContext(transform_proto.unique_name, transform_id), transform_proto.unique_name, consumers, factory.counter_factory, factory.state_sampler, output_coder, transform_id=transform_id, data_channel=factory.data_channel_factory.create_data_channel(grpc_port)) @BeamTransformFactory.register_urn( DATA_OUTPUT_URN, beam_fn_api_pb2.RemoteGrpcPort) def create_sink_runner( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, grpc_port: beam_fn_api_pb2.RemoteGrpcPort, consumers: dict[str, list[operations.Operation]]) -> DataOutputOperation: output_coder = factory.get_coder(grpc_port.coder_id) return DataOutputOperation( common.NameContext(transform_proto.unique_name, transform_id), transform_proto.unique_name, consumers, factory.counter_factory, factory.state_sampler, output_coder, transform_id=transform_id, data_channel=factory.data_channel_factory.create_data_channel(grpc_port)) @BeamTransformFactory.register_urn(OLD_DATAFLOW_RUNNER_HARNESS_READ_URN, None) def create_source_java( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, parameter, consumers: dict[str, list[operations.Operation]]) -> operations.ReadOperation: # The Dataflow runner harness strips the base64 encoding. source = pickler.loads(base64.b64encode(parameter)) spec = operation_specs.WorkerRead( iobase.SourceBundle(1.0, source, None, None), [factory.get_only_output_coder(transform_proto)]) return factory.augment_oldstyle_op( operations.ReadOperation( common.NameContext(transform_proto.unique_name, transform_id), spec, factory.counter_factory, factory.state_sampler), transform_proto.unique_name, consumers) @BeamTransformFactory.register_urn( common_urns.deprecated_primitives.READ.urn, beam_runner_api_pb2.ReadPayload) def create_deprecated_read( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, parameter: beam_runner_api_pb2.ReadPayload, consumers: dict[str, list[operations.Operation]]) -> operations.ReadOperation: source = iobase.BoundedSource.from_runner_api( parameter.source, factory.context) spec = operation_specs.WorkerRead( iobase.SourceBundle(1.0, source, None, None), [WindowedValueCoder(source.default_output_coder())]) return factory.augment_oldstyle_op( operations.ReadOperation( common.NameContext(transform_proto.unique_name, transform_id), spec, factory.counter_factory, factory.state_sampler), transform_proto.unique_name, consumers) @BeamTransformFactory.register_urn( python_urns.IMPULSE_READ_TRANSFORM, beam_runner_api_pb2.ReadPayload) def create_read_from_impulse_python( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, parameter: beam_runner_api_pb2.ReadPayload, consumers: dict[str, list[operations.Operation]] ) -> operations.ImpulseReadOperation: return operations.ImpulseReadOperation( common.NameContext(transform_proto.unique_name, transform_id), factory.counter_factory, factory.state_sampler, consumers, iobase.BoundedSource.from_runner_api(parameter.source, factory.context), factory.get_only_output_coder(transform_proto)) @BeamTransformFactory.register_urn(OLD_DATAFLOW_RUNNER_HARNESS_PARDO_URN, None) def create_dofn_javasdk( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, serialized_fn, consumers: dict[str, list[operations.Operation]]): return _create_pardo_operation( factory, transform_id, transform_proto, consumers, serialized_fn) @BeamTransformFactory.register_urn( common_urns.sdf_components.PAIR_WITH_RESTRICTION.urn, beam_runner_api_pb2.ParDoPayload) def create_pair_with_restriction(*args): class PairWithRestriction(beam.DoFn): def __init__(self, fn, restriction_provider, watermark_estimator_provider): self.restriction_provider = restriction_provider self.watermark_estimator_provider = watermark_estimator_provider def process(self, element, *args, **kwargs): # TODO(SDF): Do we want to allow mutation of the element? # (E.g. it could be nice to shift bulky description to the portion # that can be distributed.) initial_restriction = self.restriction_provider.initial_restriction( element) initial_estimator_state = ( self.watermark_estimator_provider.initial_estimator_state( element, initial_restriction)) yield (element, (initial_restriction, initial_estimator_state)) return _create_sdf_operation(PairWithRestriction, *args) @BeamTransformFactory.register_urn( common_urns.sdf_components.SPLIT_AND_SIZE_RESTRICTIONS.urn, beam_runner_api_pb2.ParDoPayload) def create_split_and_size_restrictions(*args): class SplitAndSizeRestrictions(beam.DoFn): def __init__(self, fn, restriction_provider, watermark_estimator_provider): self.restriction_provider = restriction_provider self.watermark_estimator_provider = watermark_estimator_provider def process(self, element_restriction, *args, **kwargs): element, (restriction, _) = element_restriction for part, size in self.restriction_provider.split_and_size( element, restriction): if size < 0: raise ValueError('Expected size >= 0 but received %s.' % size) estimator_state = ( self.watermark_estimator_provider.initial_estimator_state( element, part)) yield ((element, (part, estimator_state)), size) return _create_sdf_operation(SplitAndSizeRestrictions, *args) @BeamTransformFactory.register_urn( common_urns.sdf_components.TRUNCATE_SIZED_RESTRICTION.urn, beam_runner_api_pb2.ParDoPayload) def create_truncate_sized_restriction(*args): class TruncateAndSizeRestriction(beam.DoFn): def __init__(self, fn, restriction_provider, watermark_estimator_provider): self.restriction_provider = restriction_provider def process(self, element_restriction, *args, **kwargs): ((element, (restriction, estimator_state)), _) = element_restriction truncated_restriction = self.restriction_provider.truncate( element, restriction) if truncated_restriction: truncated_restriction_size = ( self.restriction_provider.restriction_size( element, truncated_restriction)) if truncated_restriction_size < 0: raise ValueError( 'Expected size >= 0 but received %s.' % truncated_restriction_size) yield ((element, (truncated_restriction, estimator_state)), truncated_restriction_size) return _create_sdf_operation( TruncateAndSizeRestriction, *args, operation_cls=operations.SdfTruncateSizedRestrictions) @BeamTransformFactory.register_urn( common_urns.sdf_components.PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS.urn, beam_runner_api_pb2.ParDoPayload) def create_process_sized_elements_and_restrictions( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, parameter: beam_runner_api_pb2.ParDoPayload, consumers: dict[str, list[operations.Operation]]): return _create_pardo_operation( factory, transform_id, transform_proto, consumers, core.DoFnInfo.from_runner_api(parameter.do_fn, factory.context).serialized_dofn_data(), parameter, operation_cls=operations.SdfProcessSizedElements) def _create_sdf_operation( proxy_dofn, factory, transform_id, transform_proto, parameter, consumers, operation_cls=operations.DoOperation): dofn_data = pickler.loads(parameter.do_fn.payload) dofn = dofn_data[0] restriction_provider = common.DoFnSignature(dofn).get_restriction_provider() watermark_estimator_provider = ( common.DoFnSignature(dofn).get_watermark_estimator_provider()) serialized_fn = pickler.dumps( (proxy_dofn(dofn, restriction_provider, watermark_estimator_provider), ) + dofn_data[1:]) return _create_pardo_operation( factory, transform_id, transform_proto, consumers, serialized_fn, parameter, operation_cls=operation_cls) @BeamTransformFactory.register_urn( common_urns.primitives.PAR_DO.urn, beam_runner_api_pb2.ParDoPayload) def create_par_do( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, parameter: beam_runner_api_pb2.ParDoPayload, consumers: dict[str, list[operations.Operation]]) -> operations.DoOperation: return _create_pardo_operation( factory, transform_id, transform_proto, consumers, core.DoFnInfo.from_runner_api(parameter.do_fn, factory.context).serialized_dofn_data(), parameter) def _create_pardo_operation( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, consumers, serialized_fn, pardo_proto: Optional[beam_runner_api_pb2.ParDoPayload] = None, operation_cls=operations.DoOperation): if pardo_proto and pardo_proto.side_inputs: input_tags_to_coders = factory.get_input_coders(transform_proto) tagged_side_inputs = [ (tag, beam.pvalue.SideInputData.from_runner_api(si, factory.context)) for tag, si in pardo_proto.side_inputs.items() ] tagged_side_inputs.sort( key=lambda tag_si: sideinputs.get_sideinput_index(tag_si[0])) side_input_maps = [ StateBackedSideInputMap( factory.state_handler, transform_id, tag, si, input_tags_to_coders[tag], use_bulk_read=( common_urns.runner_protocols.MULTIMAP_KEYS_VALUES_SIDE_INPUT.urn in factory.runner_capabilities)) for (tag, si) in tagged_side_inputs ] else: side_input_maps = [] output_tags = list(transform_proto.outputs.keys()) dofn_data = pickler.loads(serialized_fn) if not dofn_data[-1]: # Windowing not set. if pardo_proto: other_input_tags: Container[str] = set.union( set(pardo_proto.side_inputs), set(pardo_proto.timer_family_specs)) else: other_input_tags = () pcoll_id, = [pcoll for tag, pcoll in transform_proto.inputs.items() if tag not in other_input_tags] windowing = factory.context.windowing_strategies.get_by_id( factory.descriptor.pcollections[pcoll_id].windowing_strategy_id) serialized_fn = pickler.dumps(dofn_data[:-1] + (windowing, )) if pardo_proto and (pardo_proto.timer_family_specs or pardo_proto.state_specs or pardo_proto.restriction_coder_id): found_input_coder = None for tag, pcoll_id in transform_proto.inputs.items(): if tag in pardo_proto.side_inputs: pass else: # Must be the main input assert found_input_coder is None main_input_tag = tag found_input_coder = factory.get_windowed_coder(pcoll_id) assert found_input_coder is not None main_input_coder = found_input_coder if pardo_proto.timer_family_specs or pardo_proto.state_specs: user_state_context: Optional[ FnApiUserStateContext] = FnApiUserStateContext( factory.state_handler, transform_id, main_input_coder.key_coder(), main_input_coder.window_coder) else: user_state_context = None else: user_state_context = None output_coders = factory.get_output_coders(transform_proto) spec = operation_specs.WorkerDoFn( serialized_fn=serialized_fn, output_tags=output_tags, input=None, side_inputs=None, # Fn API uses proto definitions and the Fn State API output_coders=[output_coders[tag] for tag in output_tags]) result = factory.augment_oldstyle_op( operation_cls( common.NameContext(transform_proto.unique_name, transform_id), spec, factory.counter_factory, factory.state_sampler, side_input_maps, user_state_context), transform_proto.unique_name, consumers, output_tags) if pardo_proto and pardo_proto.restriction_coder_id: result.input_info = operations.OpInputInfo( transform_id, main_input_tag, main_input_coder, transform_proto.outputs.keys()) return result def _create_simple_pardo_operation( factory: BeamTransformFactory, transform_id, transform_proto, consumers, dofn: beam.DoFn, ): serialized_fn = pickler.dumps((dofn, (), {}, [], None)) return _create_pardo_operation( factory, transform_id, transform_proto, consumers, serialized_fn) @BeamTransformFactory.register_urn( common_urns.primitives.ASSIGN_WINDOWS.urn, beam_runner_api_pb2.WindowingStrategy) def create_assign_windows( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, parameter: beam_runner_api_pb2.WindowingStrategy, consumers: dict[str, list[operations.Operation]]): class WindowIntoDoFn(beam.DoFn): def __init__(self, windowing): self.windowing = windowing def process( self, element, timestamp=beam.DoFn.TimestampParam, window=beam.DoFn.WindowParam): new_windows = self.windowing.windowfn.assign( WindowFn.AssignContext(timestamp, element=element, window=window)) yield WindowedValue(element, timestamp, new_windows) from apache_beam.transforms.core import Windowing from apache_beam.transforms.window import WindowFn windowing = Windowing.from_runner_api(parameter, factory.context) return _create_simple_pardo_operation( factory, transform_id, transform_proto, consumers, WindowIntoDoFn(windowing)) @BeamTransformFactory.register_urn(IDENTITY_DOFN_URN, None) def create_identity_dofn( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, parameter, consumers: dict[str, list[operations.Operation]] ) -> operations.FlattenOperation: return factory.augment_oldstyle_op( operations.FlattenOperation( common.NameContext(transform_proto.unique_name, transform_id), operation_specs.WorkerFlatten( None, [factory.get_only_output_coder(transform_proto)]), factory.counter_factory, factory.state_sampler), transform_proto.unique_name, consumers) @BeamTransformFactory.register_urn( common_urns.combine_components.COMBINE_PER_KEY_PRECOMBINE.urn, beam_runner_api_pb2.CombinePayload) def create_combine_per_key_precombine( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, payload: beam_runner_api_pb2.CombinePayload, consumers: dict[str, list[operations.Operation]]) -> operations.PGBKCVOperation: serialized_combine_fn = pickler.dumps(( beam.CombineFn.from_runner_api(payload.combine_fn, factory.context), [], {})) return factory.augment_oldstyle_op( operations.PGBKCVOperation( common.NameContext(transform_proto.unique_name, transform_id), operation_specs.WorkerPartialGroupByKey( serialized_combine_fn, None, [factory.get_only_output_coder(transform_proto)]), factory.counter_factory, factory.state_sampler, factory.get_input_windowing(transform_proto)), transform_proto.unique_name, consumers) @BeamTransformFactory.register_urn( common_urns.combine_components.COMBINE_PER_KEY_MERGE_ACCUMULATORS.urn, beam_runner_api_pb2.CombinePayload) def create_combbine_per_key_merge_accumulators( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, payload: beam_runner_api_pb2.CombinePayload, consumers: dict[str, list[operations.Operation]]): return _create_combine_phase_operation( factory, transform_id, transform_proto, payload, consumers, 'merge') @BeamTransformFactory.register_urn( common_urns.combine_components.COMBINE_PER_KEY_EXTRACT_OUTPUTS.urn, beam_runner_api_pb2.CombinePayload) def create_combine_per_key_extract_outputs( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, payload: beam_runner_api_pb2.CombinePayload, consumers: dict[str, list[operations.Operation]]): return _create_combine_phase_operation( factory, transform_id, transform_proto, payload, consumers, 'extract') @BeamTransformFactory.register_urn( common_urns.combine_components.COMBINE_PER_KEY_CONVERT_TO_ACCUMULATORS.urn, beam_runner_api_pb2.CombinePayload) def create_combine_per_key_convert_to_accumulators( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, payload: beam_runner_api_pb2.CombinePayload, consumers: dict[str, list[operations.Operation]]): return _create_combine_phase_operation( factory, transform_id, transform_proto, payload, consumers, 'convert') @BeamTransformFactory.register_urn( common_urns.combine_components.COMBINE_GROUPED_VALUES.urn, beam_runner_api_pb2.CombinePayload) def create_combine_grouped_values( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, payload: beam_runner_api_pb2.CombinePayload, consumers: dict[str, list[operations.Operation]]): return _create_combine_phase_operation( factory, transform_id, transform_proto, payload, consumers, 'all') def _create_combine_phase_operation( factory, transform_id, transform_proto, payload, consumers, phase) -> operations.CombineOperation: serialized_combine_fn = pickler.dumps(( beam.CombineFn.from_runner_api(payload.combine_fn, factory.context), [], {})) return factory.augment_oldstyle_op( operations.CombineOperation( common.NameContext(transform_proto.unique_name, transform_id), operation_specs.WorkerCombineFn( serialized_combine_fn, phase, None, [factory.get_only_output_coder(transform_proto)]), factory.counter_factory, factory.state_sampler), transform_proto.unique_name, consumers) @BeamTransformFactory.register_urn(common_urns.primitives.FLATTEN.urn, None) def create_flatten( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, payload, consumers: dict[str, list[operations.Operation]] ) -> operations.FlattenOperation: return factory.augment_oldstyle_op( operations.FlattenOperation( common.NameContext(transform_proto.unique_name, transform_id), operation_specs.WorkerFlatten( None, [factory.get_only_output_coder(transform_proto)]), factory.counter_factory, factory.state_sampler), transform_proto.unique_name, consumers) @BeamTransformFactory.register_urn( common_urns.primitives.MAP_WINDOWS.urn, beam_runner_api_pb2.FunctionSpec) def create_map_windows( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, mapping_fn_spec: beam_runner_api_pb2.FunctionSpec, consumers: dict[str, list[operations.Operation]]): assert mapping_fn_spec.urn == python_urns.PICKLED_WINDOW_MAPPING_FN window_mapping_fn = pickler.loads(mapping_fn_spec.payload) class MapWindows(beam.DoFn): def process(self, element): key, window = element return [(key, window_mapping_fn(window))] return _create_simple_pardo_operation( factory, transform_id, transform_proto, consumers, MapWindows()) @BeamTransformFactory.register_urn( common_urns.primitives.MERGE_WINDOWS.urn, beam_runner_api_pb2.FunctionSpec) def create_merge_windows( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, mapping_fn_spec: beam_runner_api_pb2.FunctionSpec, consumers: dict[str, list[operations.Operation]]): assert mapping_fn_spec.urn == python_urns.PICKLED_WINDOWFN window_fn = pickler.loads(mapping_fn_spec.payload) class MergeWindows(beam.DoFn): def process(self, element): nonce, windows = element original_windows: set[window.BoundedWindow] = set(windows) merged_windows: MutableMapping[ window.BoundedWindow, set[window.BoundedWindow]] = collections.defaultdict( set) # noqa: F821 class RecordingMergeContext(window.WindowFn.MergeContext): def merge( self, to_be_merged: Iterable[window.BoundedWindow], merge_result: window.BoundedWindow, ): originals = merged_windows[merge_result] for w in to_be_merged: if w in original_windows: originals.add(w) original_windows.remove(w) else: originals.update(merged_windows.pop(w)) window_fn.merge(RecordingMergeContext(windows)) yield nonce, (original_windows, merged_windows.items()) return _create_simple_pardo_operation( factory, transform_id, transform_proto, consumers, MergeWindows()) @BeamTransformFactory.register_urn(common_urns.primitives.TO_STRING.urn, None) def create_to_string_fn( factory: BeamTransformFactory, transform_id: str, transform_proto: beam_runner_api_pb2.PTransform, mapping_fn_spec: beam_runner_api_pb2.FunctionSpec, consumers: dict[str, list[operations.Operation]]): class ToString(beam.DoFn): def process(self, element): key, value = element return [(key, str(value))] return _create_simple_pardo_operation( factory, transform_id, transform_proto, consumers, ToString()) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/bundle_processor_test.py0000644000175100017510000006652015206113032025654 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for bundle processing.""" # pytype: skip-file import random import unittest import apache_beam as beam from apache_beam.coders import StrUtf8Coder from apache_beam.coders.coders import FastPrimitivesCoder from apache_beam.portability import common_urns from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.runners import common from apache_beam.runners.portability.fn_api_runner.worker_handlers import StateServicer from apache_beam.runners.worker import bundle_processor from apache_beam.runners.worker import operations from apache_beam.runners.worker.bundle_processor import BeamTransformFactory from apache_beam.runners.worker.bundle_processor import BundleProcessor from apache_beam.runners.worker.bundle_processor import DataInputOperation from apache_beam.runners.worker.bundle_processor import FnApiUserStateContext from apache_beam.runners.worker.bundle_processor import SynchronousOrderedListRuntimeState from apache_beam.runners.worker.bundle_processor import TimerInfo from apache_beam.runners.worker.data_plane import SizeBasedBufferingClosableOutputStream from apache_beam.runners.worker.data_sampler import DataSampler from apache_beam.runners.worker.sdk_worker import GlobalCachingStateHandler from apache_beam.runners.worker.statecache import StateCache from apache_beam.transforms import userstate from apache_beam.transforms.window import GlobalWindow from apache_beam.utils import timestamp from apache_beam.utils.windowed_value import WindowedValue class FnApiUserStateContextTest(unittest.TestCase): def testOutputTimerTimestamp(self): class Coder(object): """Dummy coder to capture the timer result befor encoding.""" def encode_to_stream(self, timer, *args, **kwargs): self.timer = timer coder = Coder() ctx = FnApiUserStateContext(None, 'transform_id', None, None) ctx.add_timer_info( 'ts-event-timer', TimerInfo(coder, SizeBasedBufferingClosableOutputStream())) ctx.add_timer_info( 'ts-proc-timer', TimerInfo(coder, SizeBasedBufferingClosableOutputStream())) timer_spec1 = userstate.TimerSpec( 'event-timer', userstate.TimeDomain.WATERMARK) timer_spec2 = userstate.TimerSpec( 'proc-timer', userstate.TimeDomain.REAL_TIME) # Set event time timer event_timer = ctx.get_timer(timer_spec1, 'key', GlobalWindow, 23, None) event_timer.set(42) # Output timestamp should be equal to the fire timestamp self.assertEqual(coder.timer.hold_timestamp, 42) # Set processing time timer proc_timer = ctx.get_timer(timer_spec2, 'key', GlobalWindow, 23, None) proc_timer.set(42) # Output timestamp should be equal to the input timestamp self.assertEqual(coder.timer.hold_timestamp, 23) class SplitTest(unittest.TestCase): def split( self, index, current_element_progress, fraction_of_remainder, buffer_size, allowed=(), sdf=False): return DataInputOperation._compute_split( index, current_element_progress, float('inf'), fraction_of_remainder, buffer_size, allowed_split_points=allowed, try_split=lambda frac: element_split(frac, 0)[1:3] if sdf else None) def sdf_split(self, *args, **kwargs): return self.split(*args, sdf=True, **kwargs) def test_simple_split(self): # Split as close to the beginning as possible. self.assertEqual(self.split(0, 0, 0, 16), simple_split(1)) # The closest split is at 4, even when just above or below it. self.assertEqual(self.split(0, 0, 0.24, 16), simple_split(4)) self.assertEqual(self.split(0, 0, 0.25, 16), simple_split(4)) self.assertEqual(self.split(0, 0, 0.26, 16), simple_split(4)) # Split the *remainder* in half. self.assertEqual(self.split(0, 0, 0.5, 16), simple_split(8)) self.assertEqual(self.split(2, 0, 0.5, 16), simple_split(9)) self.assertEqual(self.split(6, 0, 0.5, 16), simple_split(11)) def test_split_with_element_progress(self): # Progress into the active element influences where the split of the # remainder falls. self.assertEqual(self.split(0, 0.5, 0.25, 4), simple_split(1)) self.assertEqual(self.split(0, 0.9, 0.25, 4), simple_split(2)) self.assertEqual(self.split(1, 0.0, 0.25, 4), simple_split(2)) self.assertEqual(self.split(1, 0.1, 0.25, 4), simple_split(2)) def test_split_with_element_allowed_splits(self): # The desired split point is at 4. self.assertEqual( self.split(0, 0, 0.25, 16, allowed=(2, 3, 4, 5)), simple_split(4)) # If we can't split at 4, choose the closest possible split point. self.assertEqual( self.split(0, 0, 0.25, 16, allowed=(2, 3, 5)), simple_split(5)) self.assertEqual( self.split(0, 0, 0.25, 16, allowed=(2, 3, 6)), simple_split(3)) # Also test the case where all possible split points lie above or below # the desired split point. self.assertEqual( self.split(0, 0, 0.25, 16, allowed=(5, 6, 7)), simple_split(5)) self.assertEqual( self.split(0, 0, 0.25, 16, allowed=(1, 2, 3)), simple_split(3)) # We have progressed beyond all possible split points, so can't split. self.assertEqual(self.split(5, 0, 0.25, 16, allowed=(1, 2, 3)), None) def test_sdf_split(self): # Split between future elements at element boundaries. self.assertEqual(self.sdf_split(0, 0, 0.51, 4), simple_split(2)) self.assertEqual(self.sdf_split(0, 0, 0.49, 4), simple_split(2)) self.assertEqual(self.sdf_split(0, 0, 0.26, 4), simple_split(1)) self.assertEqual(self.sdf_split(0, 0, 0.25, 4), simple_split(1)) # If the split falls inside the first, splittable element, split there. self.assertEqual( self.sdf_split(0, 0, 0.20, 4), (-1, ['Primary(0.8)'], ['Residual(0.2)'], 1)) # The choice of split depends on the progress into the first element. self.assertEqual( self.sdf_split(0, 0, .125, 4), (-1, ['Primary(0.5)'], ['Residual(0.5)'], 1)) # Here we are far enough into the first element that splitting at 0.2 of the # remainder falls outside the first element. self.assertEqual(self.sdf_split(0, .5, 0.2, 4), simple_split(1)) # Verify the above logic when we are partially through the stream. self.assertEqual(self.sdf_split(2, 0, 0.6, 4), simple_split(3)) self.assertEqual(self.sdf_split(2, 0.9, 0.6, 4), simple_split(4)) self.assertEqual( self.sdf_split(2, 0.5, 0.2, 4), (1, ['Primary(0.6)'], ['Residual(0.4)'], 3)) def test_sdf_split_with_allowed_splits(self): # This is where we would like to split, when all split points are available. self.assertEqual( self.sdf_split(2, 0, 0.2, 5, allowed=(1, 2, 3, 4, 5)), (1, ['Primary(0.6)'], ['Residual(0.4)'], 3)) # We can't split element at index 2, because 3 is not a split point. self.assertEqual( self.sdf_split(2, 0, 0.2, 5, allowed=(1, 2, 4, 5)), simple_split(4)) # We can't even split element at index 4 as above, because 4 is also not a # split point. self.assertEqual( self.sdf_split(2, 0, 0.2, 5, allowed=(1, 2, 5)), simple_split(5)) # We can't split element at index 2, because 2 is not a split point. self.assertEqual( self.sdf_split(2, 0, 0.2, 5, allowed=(1, 3, 4, 5)), simple_split(3)) def simple_split(first_residual_index): return first_residual_index - 1, [], [], first_residual_index def element_split(frac, index): return ( index - 1, ['Primary(%0.1f)' % frac], ['Residual(%0.1f)' % (1 - frac)], index + 1) class TestOperation(operations.Operation): """Test operation that forwards its payload to consumers.""" class Spec: def __init__(self, transform_proto): self.output_coders = [ FastPrimitivesCoder() for _ in transform_proto.outputs ] def __init__( self, transform_proto, name_context, counter_factory, state_sampler, consumers, payload, ): super().__init__( name_context, self.Spec(transform_proto), counter_factory, state_sampler) self.payload = payload for _, consumer_ops in consumers.items(): for consumer in consumer_ops: self.add_receiver(consumer, 0) def start(self): super().start() # Not using windowing logic, so just using simple defaults here. if self.payload: self.process( WindowedValue(self.payload, timestamp=0, windows=[GlobalWindow()])) def process(self, windowed_value): self.output(windowed_value) @BeamTransformFactory.register_urn('beam:internal:testop:v1', bytes) def create_test_op(factory, transform_id, transform_proto, payload, consumers): return TestOperation( transform_proto, common.NameContext(transform_proto.unique_name, transform_id), factory.counter_factory, factory.state_sampler, consumers, payload) @BeamTransformFactory.register_urn('beam:internal:testexn:v1', bytes) def create_exception_dofn( factory, transform_id, transform_proto, payload, consumers): """Returns a test DoFn that raises the given exception.""" class RaiseException(beam.DoFn): def __init__(self, msg): self.msg = msg.decode() def process(self, _): raise RuntimeError(self.msg) return bundle_processor._create_simple_pardo_operation( factory, transform_id, transform_proto, consumers, RaiseException(payload)) class DataSamplingTest(unittest.TestCase): def test_disabled_by_default(self): """Test that not providing the sampler does not enable Data Sampling. Note that data sampling is enabled by providing the sampler to the processor. """ descriptor = beam_fn_api_pb2.ProcessBundleDescriptor() descriptor.pcollections['a'].unique_name = 'a' _ = BundleProcessor(set(), descriptor, None, None) self.assertEqual(len(descriptor.transforms), 0) def test_can_sample(self): """Test that elements are sampled. This is a small integration test with the BundleProcessor and the DataSampler. It ensures that samples are taken from in-flight elements. These elements are then finally queried. """ data_sampler = DataSampler(sample_every_sec=0.1) descriptor = beam_fn_api_pb2.ProcessBundleDescriptor() # Create the PCollection to sample from. PCOLLECTION_ID = 'pc' CODER_ID = 'c' descriptor.pcollections[PCOLLECTION_ID].unique_name = PCOLLECTION_ID descriptor.pcollections[PCOLLECTION_ID].coder_id = CODER_ID descriptor.coders[ CODER_ID].spec.urn = common_urns.StandardCoders.Enum.BYTES.urn # Add a simple transform to inject an element into the data sampler. This # doesn't use the FnApi, so this uses a simple operation to forward its # payload to consumers. TRANSFORM_ID = 'test_transform' test_transform = descriptor.transforms[TRANSFORM_ID] test_transform.outputs['None'] = PCOLLECTION_ID test_transform.spec.urn = 'beam:internal:testop:v1' test_transform.spec.payload = b'hello, world!' try: # Create and process a fake bundle. The instruction id doesn't matter # here. processor = BundleProcessor( set(), descriptor, None, None, data_sampler=data_sampler) processor.process_bundle('instruction_id') samples = data_sampler.wait_for_samples([PCOLLECTION_ID]) expected = beam_fn_api_pb2.SampleDataResponse( element_samples={ PCOLLECTION_ID: beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=b'\rhello, world!') ]) }) self.assertEqual(samples, expected) finally: data_sampler.stop() def test_can_sample_exceptions(self): """Test that exceptions are sampled.""" data_sampler = DataSampler(sample_every_sec=0.1) descriptor = beam_fn_api_pb2.ProcessBundleDescriptor() # Boiler plate for the DoFn. WINDOWING_ID = 'window' WINDOW_CODER_ID = 'cw' window = descriptor.windowing_strategies[WINDOWING_ID] window.window_fn.urn = common_urns.global_windows.urn window.window_coder_id = WINDOW_CODER_ID window.trigger.default.SetInParent() window_coder = descriptor.coders[WINDOW_CODER_ID] window_coder.spec.urn = common_urns.StandardCoders.Enum.GLOBAL_WINDOW.urn # Input collection to the exception raising DoFn. INPUT_PCOLLECTION_ID = 'pc-in' INPUT_CODER_ID = 'c-in' descriptor.pcollections[ INPUT_PCOLLECTION_ID].unique_name = INPUT_PCOLLECTION_ID descriptor.pcollections[INPUT_PCOLLECTION_ID].coder_id = INPUT_CODER_ID descriptor.pcollections[ INPUT_PCOLLECTION_ID].windowing_strategy_id = WINDOWING_ID descriptor.coders[ INPUT_CODER_ID].spec.urn = common_urns.StandardCoders.Enum.BYTES.urn # Output collection to the exception raising DoFn. Because the transform # "failed" to process the input element, we do NOT expect to see a sample in # this PCollection. OUTPUT_PCOLLECTION_ID = 'pc-out' OUTPUT_CODER_ID = 'c-out' descriptor.pcollections[ OUTPUT_PCOLLECTION_ID].unique_name = OUTPUT_PCOLLECTION_ID descriptor.pcollections[OUTPUT_PCOLLECTION_ID].coder_id = OUTPUT_CODER_ID descriptor.pcollections[ OUTPUT_PCOLLECTION_ID].windowing_strategy_id = WINDOWING_ID descriptor.coders[ OUTPUT_CODER_ID].spec.urn = common_urns.StandardCoders.Enum.BYTES.urn # Add a simple transform to inject an element into the data sampler. This # doesn't use the FnApi, so this uses a simple operation to forward its # payload to consumers. TEST_OP_TRANSFORM_ID = 'test_op' test_transform = descriptor.transforms[TEST_OP_TRANSFORM_ID] test_transform.outputs['None'] = INPUT_PCOLLECTION_ID test_transform.spec.urn = 'beam:internal:testop:v1' test_transform.spec.payload = b'hello, world!' # Add the DoFn to create an exception to sample from. TEST_EXCEPTION_TRANSFORM_ID = 'test_transform' test_transform = descriptor.transforms[TEST_EXCEPTION_TRANSFORM_ID] test_transform.inputs['0'] = INPUT_PCOLLECTION_ID test_transform.outputs['None'] = OUTPUT_PCOLLECTION_ID test_transform.spec.urn = 'beam:internal:testexn:v1' test_transform.spec.payload = b'expected exception' try: # Create and process a fake bundle. The instruction id doesn't matter # here. processor = BundleProcessor( set(), descriptor, None, None, data_sampler=data_sampler) with self.assertRaisesRegex(RuntimeError, 'expected exception'): processor.process_bundle('instruction_id') # NOTE: The expected sample comes from the input PCollection. This is very # important because there can be coder issues if the sample is put in the # wrong PCollection. samples = data_sampler.wait_for_samples([INPUT_PCOLLECTION_ID]) self.assertEqual(len(samples.element_samples), 1) element = samples.element_samples[INPUT_PCOLLECTION_ID].elements[0] self.assertEqual(element.element, b'\rhello, world!') self.assertTrue(element.HasField('exception')) exception = element.exception self.assertEqual(exception.instruction_id, 'instruction_id') self.assertEqual(exception.transform_id, TEST_EXCEPTION_TRANSFORM_ID) self.assertRegex( exception.error, 'Traceback(\n|.)*RuntimeError: expected exception') finally: data_sampler.stop() class EnvironmentCompatibilityTest(unittest.TestCase): def test_rc_environments_are_compatible_with_released_images(self): # TODO(https://github.com/apache/beam/issues/28084): remove when # resolved. self.assertTrue( bundle_processor._environments_compatible( "beam:version:sdk_base:apache/beam_python3.5_sdk:2.1.0rc1", "beam:version:sdk_base:apache/beam_python3.5_sdk:2.1.0")) def test_user_modified_sdks_need_to_be_installed_in_runtime_env(self): self.assertFalse( bundle_processor._environments_compatible( "beam:version:sdk_base:apache/beam_python3.5_sdk:2.1.0-custom", "beam:version:sdk_base:apache/beam_python3.5_sdk:2.1.0")) self.assertTrue( bundle_processor._environments_compatible( "beam:version:sdk_base:apache/beam_python3.5_sdk:2.1.0-custom", "beam:version:sdk_base:apache/beam_python3.5_sdk:2.1.0-custom")) class OrderedListStateTest(unittest.TestCase): class NoStateCache(StateCache): def __init__(self): super().__init__(max_weight=0) @staticmethod def _create_state(window=b"my_window", key=b"my_key", coder=StrUtf8Coder()): state_handler = GlobalCachingStateHandler( OrderedListStateTest.NoStateCache(), StateServicer()) state_key = beam_fn_api_pb2.StateKey( ordered_list_user_state=beam_fn_api_pb2.StateKey.OrderedListUserState( window=window, key=key)) return SynchronousOrderedListRuntimeState(state_handler, state_key, coder) def setUp(self): self.state = self._create_state() def test_read_range(self): T0 = timestamp.Timestamp.of(0) T1 = timestamp.Timestamp.of(1) T2 = timestamp.Timestamp.of(2) T3 = timestamp.Timestamp.of(3) T4 = timestamp.Timestamp.of(4) T5 = timestamp.Timestamp.of(5) T9 = timestamp.Timestamp.of(9) A1, B1, A4 = [(T1, "a1"), (T1, "b1"), (T4, "a4")] self.assertEqual([], list(self.state.read_range(T0, T5))) self.state.add(A1) self.assertEqual([A1], list(self.state.read_range(T0, T5))) self.state.add(B1) self.assertEqual([A1, B1], list(self.state.read_range(T0, T5))) self.state.add(A4) self.assertEqual([A1, B1, A4], list(self.state.read_range(T0, T5))) self.assertEqual([], list(self.state.read_range(T0, T1))) self.assertEqual([], list(self.state.read_range(T5, T9))) self.assertEqual([A1, B1], list(self.state.read_range(T1, T2))) self.assertEqual([], list(self.state.read_range(T2, T3))) self.assertEqual([], list(self.state.read_range(T2, T4))) self.assertEqual([A4], list(self.state.read_range(T4, T5))) def test_read(self): T1 = timestamp.Timestamp.of(1) T4 = timestamp.Timestamp.of(4) A1, B1, A4 = [(T1, "a1"), (T1, "b1"), (T4, "a4")] self.assertEqual([], list(self.state.read())) self.state.add(A1) self.assertEqual([A1], list(self.state.read())) self.state.add(A1) self.assertEqual([A1, A1], list(self.state.read())) self.state.add(B1) self.assertEqual([A1, A1, B1], list(self.state.read())) self.state.add(A4) self.assertEqual([A1, A1, B1, A4], list(self.state.read())) def test_clear_range(self): T0 = timestamp.Timestamp.of(0) T1 = timestamp.Timestamp.of(1) T2 = timestamp.Timestamp.of(2) T3 = timestamp.Timestamp.of(3) T4 = timestamp.Timestamp.of(4) T5 = timestamp.Timestamp.of(5) A1, B1, A4, A5 = [(T1, "a1"), (T1, "b1"), (T4, "a4"), (T5, "a5")] self.state.clear_range(T0, T1) self.assertEqual([], list(self.state.read())) self.state.add(A1) self.state.add(B1) self.state.add(A4) self.state.add(A5) self.assertEqual([A1, B1, A4, A5], list(self.state.read())) self.state.clear_range(T0, T1) self.assertEqual([A1, B1, A4, A5], list(self.state.read())) self.state.clear_range(T1, T2) self.assertEqual([A4, A5], list(self.state.read())) # no side effect on clearing the same range twice self.state.clear_range(T1, T2) self.assertEqual([A4, A5], list(self.state.read())) self.state.clear_range(T3, T4) self.assertEqual([A4, A5], list(self.state.read())) self.state.clear_range(T3, T5) self.assertEqual([A5], list(self.state.read())) def test_add_and_clear_range_after_commit(self): T1 = timestamp.Timestamp.of(1) T4 = timestamp.Timestamp.of(4) T5 = timestamp.Timestamp.of(5) T6 = timestamp.Timestamp.of(6) A1, B1, C1, A4, A5, A6 = [(T1, "a1"), (T1, "b1"), (T1, "c1"), (T4, "a4"), (T5, "a5"), (T6, "a6")] self.state.add(A1) self.state.add(B1) self.state.add(A4) self.state.add(A5) self.state.clear_range(T4, T5) self.assertEqual([A1, B1, A5], list(self.state.read())) self.state.commit() self.assertEqual(len(self.state._pending_adds), 0) self.assertEqual(len(self.state._pending_removes), 0) self.assertEqual([A1, B1, A5], list(self.state.read())) self.state.add(C1) self.state.add(A6) self.assertEqual([A1, B1, C1, A5, A6], list(self.state.read())) self.state.clear_range(T5, T6) self.assertEqual([A1, B1, C1, A6], list(self.state.read())) self.state.commit() self.assertEqual(len(self.state._pending_adds), 0) self.assertEqual(len(self.state._pending_removes), 0) self.assertEqual([A1, B1, C1, A6], list(self.state.read())) def test_clear(self): T1 = timestamp.Timestamp.of(1) T4 = timestamp.Timestamp.of(4) T5 = timestamp.Timestamp.of(5) T9 = timestamp.Timestamp.of(9) A1, B1, C1, A4, A5, B5 = [(T1, "a1"), (T1, "b1"), (T1, "c1"), (T4, "a4"), (T5, "a5"), (T5, "b5")] self.state.add(A1) self.state.add(B1) self.state.add(A4) self.state.add(A5) self.state.clear_range(T4, T5) self.assertEqual([A1, B1, A5], list(self.state.read())) self.state.commit() self.state.add(C1) self.state.clear_range(T5, T9) self.assertEqual([A1, B1, C1], list(self.state.read())) self.state.clear() self.assertEqual(len(self.state._pending_adds), 0) self.assertEqual(len(self.state._pending_removes), 1) self.state.add(B5) self.assertEqual([B5], list(self.state.read())) self.state.commit() self.assertEqual(len(self.state._pending_adds), 0) self.assertEqual(len(self.state._pending_removes), 0) self.assertEqual([B5], list(self.state.read())) def test_multiple_iterators(self): T1 = timestamp.Timestamp.of(1) T3 = timestamp.Timestamp.of(3) T9 = timestamp.Timestamp.of(9) A1, B1, A3, B3 = [(T1, "a1"), (T1, "b1"), (T3, "a3"), (T3, "b3")] self.state.add(A1) self.state.add(A3) self.state.commit() iter_before_b1 = iter(self.state.read()) self.assertEqual(A1, next(iter_before_b1)) self.state.add(B1) self.assertEqual(A3, next(iter_before_b1)) self.assertRaises(StopIteration, lambda: next(iter_before_b1)) self.state.add(B3) iter_before_clear_range = iter(self.state.read()) self.assertEqual(A1, next(iter_before_clear_range)) self.state.clear_range(T3, T9) self.assertEqual(B1, next(iter_before_clear_range)) self.assertEqual(A3, next(iter_before_clear_range)) self.assertEqual(B3, next(iter_before_clear_range)) self.assertRaises(StopIteration, lambda: next(iter_before_clear_range)) self.assertEqual([A1, B1], list(self.state.read())) iter_before_clear = iter(self.state.read()) self.assertEqual(A1, next(iter_before_clear)) self.state.clear() self.assertEqual(B1, next(iter_before_clear)) self.assertRaises(StopIteration, lambda: next(iter_before_clear)) self.assertEqual([], list(self.state.read())) def fuzz_test_helper(self, seed=0, lower=0, upper=20): class NaiveState: def __init__(self): self._data = [[] for i in range((upper - lower + 1))] self._logs = [] def add(self, elem): k, v = elem k = k.micros self._data[k - lower].append(v) self._logs.append("add(%d, %s)" % (k, v)) def clear_range(self, lo, hi): lo = lo.micros hi = hi.micros for i in range(lo, hi): self._data[i - lower] = [] self._logs.append("clear_range(%d, %d)" % (lo, hi)) def clear(self): for i in range(len(self._data)): self._data[i] = [] self._logs.append("clear()") def read(self): self._logs.append("read()") for i in range(len(self._data)): for v in self._data[i]: yield (timestamp.Timestamp(micros=(i + lower)), v) random.seed(seed) state = self._create_state() bench_state = NaiveState() steps = random.randint(20, 50) for i in range(steps): op = random.randint(1, 100) if 1 <= op < 70: num = random.randint(lower, upper) state.add((timestamp.Timestamp(micros=num), "a%d" % num)) bench_state.add((timestamp.Timestamp(micros=num), "a%d" % num)) elif 70 <= op < 95: num1 = random.randint(lower, upper) num2 = random.randint(lower, upper) min_time = timestamp.Timestamp(micros=min(num1, num2)) max_time = timestamp.Timestamp(micros=max(num1, num2)) state.clear_range(min_time, max_time) bench_state.clear_range(min_time, max_time) elif op >= 95: state.clear() bench_state.clear() op = random.randint(1, 10) if 1 <= op <= 9: pass else: state.commit() a = list(bench_state.read()) b = list(state.read()) self.assertEqual( a, b, "Mismatch occurred on seed=%d, step=%d, logs=%s" % (seed, i, ';'.join(bench_state._logs))) def test_fuzz(self): for _ in range(1000): seed = random.randint(0, 0xffffffffffffffff) try: self.fuzz_test_helper(seed=seed) except Exception as e: raise RuntimeError("Exception occurred on seed=%d: %s" % (seed, e)) def test_min_max(self): T_MIN = timestamp.Timestamp(micros=(-(1 << 63))) T_MAX_MINUS_ONE = timestamp.Timestamp(micros=((1 << 63) - 2)) T_MAX = timestamp.Timestamp(micros=((1 << 63) - 1)) T0 = timestamp.Timestamp(micros=0) INT64_MIN, INT64_MAX_MINUS_ONE, INT64_MAX = [(T_MIN, "min"), (T_MAX_MINUS_ONE, "max"), (T_MAX, "err")] self.state.add(INT64_MIN) self.state.add(INT64_MAX_MINUS_ONE) self.assertRaises(ValueError, lambda: self.state.add(INT64_MAX)) self.assertEqual([INT64_MIN, INT64_MAX_MINUS_ONE], list(self.state.read())) self.assertEqual([INT64_MIN], list(self.state.read_range(T_MIN, T0))) self.assertEqual([INT64_MAX_MINUS_ONE], list(self.state.read_range(T0, T_MAX))) def test_continuation_token(self): T1 = timestamp.Timestamp.of(1) T2 = timestamp.Timestamp.of(2) T7 = timestamp.Timestamp.of(7) T8 = timestamp.Timestamp.of(8) A1, A2, A7, B7, A8 = [(T1, "a1"), (T2, "a2"), (T7, "a7"), (T7, "b7"), (T8, "a8")] self.state._state_handler._underlying._use_continuation_tokens = True self.assertEqual([], list(self.state.read_range(T1, T8))) self.state.add(A1) self.state.add(A2) self.state.add(A7) self.state.add(B7) self.state.add(A8) self.assertEqual([A2, A7, B7], list(self.state.read_range(T2, T8))) self.state.commit() self.assertEqual([A2, A7, B7], list(self.state.read_range(T2, T8))) self.assertEqual([A1, A2, A7, B7, A8], list(self.state.read())) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/channel_factory.py0000644000175100017510000000353115206113032024375 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Factory to create grpc channel.""" # pytype: skip-file import grpc class GRPCChannelFactory(grpc.StreamStreamClientInterceptor): DEFAULT_OPTIONS = [ # Setting keepalive_time_ms is needed for other options to work. ("grpc.keepalive_time_ms", 20_000), # Default: 20s. Increasing to 5 min. ("grpc.keepalive_timeout_ms", 300_000), # Default: 2, set to 0 to allow unlimited pings without data ("grpc.http2.max_pings_without_data", 0), # Default: False, set to True to allow keepalive pings when no calls ("grpc.keepalive_permit_without_calls", True), ] def __init__(self): pass @staticmethod def insecure_channel(target, options=None): if options is None: options = [] return grpc.insecure_channel( target, options=options + GRPCChannelFactory.DEFAULT_OPTIONS) @staticmethod def secure_channel(target, credentials, options=None): if options is None: options = [] return grpc.secure_channel( target, credentials, options=options + GRPCChannelFactory.DEFAULT_OPTIONS) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/crossplatform_time.h0000644000175100017510000000457515206113032024762 0ustar00runnerrunner/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ #ifndef BEAM_CROSSPLATFORM_TIME_H #define BEAM_CROSSPLATFORM_TIME_H #include #ifdef _WIN32 #include /** * Alternative to POSIX clock_gettime that may be run on Windows platform. The clk_id parameter is * ignored and function act as for CLOCK_MONOTONIC. On Windows XP or later Performance Counter is used. * On older platforms, where there's no Performance Counter, SystemTime will be used as a failover. */ int clock_gettime(int clk_id, struct timespec *tv) { static LARGE_INTEGER counterFrequency = {0}; static BOOL performanceCounterAvailable = TRUE; LARGE_INTEGER counterValue = {0}; //initialization if (0 == counterFrequency.QuadPart) { if (0 == QueryPerformanceFrequency(&counterFrequency)) { performanceCounterAvailable = FALSE; counterFrequency.QuadPart = 10000000; // failover SystemTime is provided in 100-nanosecond intervals } } if (performanceCounterAvailable) { QueryPerformanceCounter(&counterValue); } else { FILETIME filetime = {0}; GetSystemTimeAsFileTime(&filetime); counterValue.QuadPart = filetime.dwHighDateTime; counterValue.QuadPart <<= 32; counterValue.QuadPart |= filetime.dwLowDateTime; } tv->tv_sec = counterValue.QuadPart / counterFrequency.QuadPart; #pragma warning( suppress : 4244 ) // nanoseconds may not exceed billion, therefore it's safe to cast tv->tv_nsec = ((counterValue.QuadPart % counterFrequency.QuadPart) * 1000000000) / counterFrequency.QuadPart; return 0; } #endif #endif //BEAM_CROSSPLATFORM_TIME_H ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/crossplatform_unistd.h0000644000175100017510000000264015206113032025321 0ustar00runnerrunner/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ #ifndef BEAM_CROSSPLATFORM_UNISTD_H #define BEAM_CROSSPLATFORM_UNISTD_H #ifdef _WIN32 #include /** * Alternative to POSIX usleep that may be run on Windows platform. */ void usleep(__int64 usec) { HANDLE timer; LARGE_INTEGER timerDueTime; // Timer is measuring in 100 of nanoseconds, negative means relative timerDueTime.QuadPart = -(10 * usec); timer = CreateWaitableTimer(NULL, TRUE, NULL); SetWaitableTimer(timer, &timerDueTime, 0, NULL, NULL, 0); WaitForSingleObject(timer, INFINITE); CloseHandle(timer); } #else #include "unistd.h" #endif #endif //BEAM_CROSSPLATFORM_UNISTD_H ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/data_plane.py0000644000175100017510000007641615206113032023342 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Implementation of ``DataChannel``s to communicate across the data plane.""" # pytype: skip-file # mypy: disallow-untyped-defs import abc import collections import json import logging import queue import threading import time from typing import TYPE_CHECKING from typing import Any from typing import Callable from typing import Collection from typing import DefaultDict from typing import Dict from typing import Iterable from typing import Iterator from typing import List from typing import Mapping from typing import Optional from typing import Set from typing import Tuple from typing import Union import grpc from apache_beam.coders import coder_impl from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.runners.worker.channel_factory import GRPCChannelFactory from apache_beam.runners.worker.worker_id_interceptor import WorkerIdInterceptor from apache_beam.utils.byte_limited_queue import ByteLimitedQueue if TYPE_CHECKING: import apache_beam.coders.slow_stream OutputStream = apache_beam.coders.slow_stream.OutputStream DataOrTimers = Union[beam_fn_api_pb2.Elements.Data, beam_fn_api_pb2.Elements.Timers] else: OutputStream = type(coder_impl.create_OutputStream()) _LOGGER = logging.getLogger(__name__) _DEFAULT_SIZE_FLUSH_THRESHOLD = 10 << 20 # 10MB _DEFAULT_TIME_FLUSH_THRESHOLD_MS = 0 # disable time-based flush by default _FLUSH_MAX_SIZE = (2 << 30) - 100 # 2GB less some overhead, protobuf/grpc limit # Keep a set of completed instructions to discard late received data. The set # can have up to _MAX_CLEANED_INSTRUCTIONS items. See _GrpcDataChannel. _MAX_CLEANED_INSTRUCTIONS = 10000 # retry on transient UNAVAILABLE grpc error from data channels. _GRPC_SERVICE_CONFIG = json.dumps({ "methodConfig": [{ "name": [{ "service": "org.apache.beam.model.fn_execution.v1.BeamFnData" }], "retryPolicy": { "maxAttempts": 5, "initialBackoff": "0.1s", "maxBackoff": "5s", "backoffMultiplier": 2, "retryableStatusCodes": ["UNAVAILABLE"], }, }] }) class ClosableOutputStream(OutputStream): """A Outputstream for use with CoderImpls that has a close() method.""" def __init__( self, close_callback=None # type: Optional[Callable[[bytes], None]] ): # type: (...) -> None super().__init__() self._close_callback = close_callback def close(self): # type: () -> None if self._close_callback: self._close_callback(self.get()) def maybe_flush(self): # type: () -> None pass def flush(self): # type: () -> None pass @staticmethod def create( close_callback, # type: Optional[Callable[[bytes], None]] flush_callback, # type: Optional[Callable[[bytes], None]] data_buffer_time_limit_ms # type: int ): # type: (...) -> ClosableOutputStream if data_buffer_time_limit_ms > 0: return TimeBasedBufferingClosableOutputStream( close_callback, flush_callback=flush_callback, time_flush_threshold_ms=data_buffer_time_limit_ms) else: return SizeBasedBufferingClosableOutputStream( close_callback, flush_callback=flush_callback) class SizeBasedBufferingClosableOutputStream(ClosableOutputStream): """A size-based buffering OutputStream.""" _large_flush_last_observed_timestamp = 0.0 def __init__( self, close_callback=None, # type: Optional[Callable[[bytes], None]] flush_callback=None, # type: Optional[Callable[[bytes], None]] size_flush_threshold=_DEFAULT_SIZE_FLUSH_THRESHOLD, # type: int large_buffer_warn_threshold_bytes=512 << 20 # type: int ): super().__init__(close_callback) self._flush_callback = flush_callback self._size_flush_threshold = size_flush_threshold self._large_buffer_warn_threshold_bytes = large_buffer_warn_threshold_bytes # This must be called explicitly to avoid flushing partial elements. def maybe_flush(self): # type: () -> None if self.size() > self._size_flush_threshold: self.flush() def flush(self): # type: () -> None if self._flush_callback: size = self.size() if (self._large_buffer_warn_threshold_bytes and size > self._large_buffer_warn_threshold_bytes): if size > _FLUSH_MAX_SIZE: raise ValueError( f'Buffer size {size} exceeds GRPC limit {_FLUSH_MAX_SIZE}. ' 'This is likely due to a single element that is too large. ' 'To resolve, prefer multiple small elements over single large ' 'elements in PCollections. If needed, store large blobs in ' 'external storage systems, and use PCollections to pass their ' 'metadata, or use a custom coder that reduces the element\'s ' 'size.') if self._large_flush_last_observed_timestamp + 600 < time.time(): self._large_flush_last_observed_timestamp = time.time() _LOGGER.warning( 'Data output stream buffer size %s exceeds %s bytes. ' 'This is likely due to a large element in a PCollection. ' 'Large elements increase pipeline RAM requirements and ' 'can cause runtime errors. ' 'Prefer multiple small elements over single large elements ' 'in PCollections. If needed, store large blobs in external ' 'storage systems, and use PCollections to pass their metadata, ' 'or use a custom coder that reduces the element\'s size.', size, self._large_buffer_warn_threshold_bytes) self._flush_callback(self.get()) self._clear() class TimeBasedBufferingClosableOutputStream( SizeBasedBufferingClosableOutputStream): """A buffering OutputStream with both time-based and size-based.""" _periodic_flusher = None # type: Optional[PeriodicThread] def __init__( self, close_callback=None, # type: Optional[Callable[[bytes], None]] flush_callback=None, # type: Optional[Callable[[bytes], None]] size_flush_threshold=_DEFAULT_SIZE_FLUSH_THRESHOLD, # type: int time_flush_threshold_ms=_DEFAULT_TIME_FLUSH_THRESHOLD_MS # type: int ): # type: (...) -> None super().__init__(close_callback, flush_callback, size_flush_threshold) assert time_flush_threshold_ms > 0 self._time_flush_threshold_ms = time_flush_threshold_ms self._flush_lock = threading.Lock() self._schedule_lock = threading.Lock() self._closed = False self._schedule_periodic_flush() def flush(self): # type: () -> None with self._flush_lock: super().flush() def close(self): # type: () -> None with self._schedule_lock: self._closed = True if self._periodic_flusher: self._periodic_flusher.cancel() self._periodic_flusher = None super().close() def _schedule_periodic_flush(self): # type: () -> None def _flush(): # type: () -> None with self._schedule_lock: if not self._closed: self.flush() self._periodic_flusher = PeriodicThread( self._time_flush_threshold_ms / 1000.0, _flush) self._periodic_flusher.daemon = True self._periodic_flusher.start() class PeriodicThread(threading.Thread): """Call a function periodically with the specified number of seconds""" def __init__( self, interval, # type: float function, # type: Callable args=None, # type: Optional[Iterable] kwargs=None # type: Optional[Mapping[str, Any]] ): # type: (...) -> None threading.Thread.__init__(self) self._interval = interval self._function = function self._args = args if args is not None else [] self._kwargs = kwargs if kwargs is not None else {} self._finished = threading.Event() def run(self): # type: () -> None next_call = time.time() + self._interval while not self._finished.wait(next_call - time.time()): next_call = next_call + self._interval self._function(*self._args, **self._kwargs) def cancel(self): # type: () -> None """Stop the thread if it hasn't finished yet.""" self._finished.set() class DataChannel(metaclass=abc.ABCMeta): """Represents a channel for reading and writing data over the data plane. Read data and timer from this channel with the input_elements method:: for elements_data in data_channel.input_elements( instruction_id, transform_ids, timers): [process elements_data] Write data to this channel using the output_stream method:: out1 = data_channel.output_stream(instruction_id, transform_id) out1.write(...) out1.close() Write timer to this channel using the output_timer_stream method:: out1 = data_channel.output_timer_stream(instruction_id, transform_id, timer_family_id) out1.write(...) out1.close() When all data/timer for all instructions is written, close the channel:: data_channel.close() """ @abc.abstractmethod def input_elements( self, instruction_id, # type: str expected_inputs, # type: Collection[Union[str, Tuple[str, str]]] abort_callback=None # type: Optional[Callable[[], bool]] ): # type: (...) -> Iterator[DataOrTimers] """Returns an iterable of all Element.Data and Element.Timers bundles for instruction_id. This iterable terminates only once the full set of data has been recieved for each of the expected transforms. It may block waiting for more data. Args: instruction_id: which instruction the results must belong to expected_inputs: which transforms to wait on for completion abort_callback: a callback to invoke if blocking returning whether to abort before consuming all the data """ raise NotImplementedError(type(self)) @abc.abstractmethod def output_stream( self, instruction_id, # type: str transform_id # type: str ): # type: (...) -> ClosableOutputStream """Returns an output stream writing elements to transform_id. Args: instruction_id: which instruction this stream belongs to transform_id: the transform_id of the returned stream """ raise NotImplementedError(type(self)) @abc.abstractmethod def output_timer_stream( self, instruction_id, # type: str transform_id, # type: str timer_family_id # type: str ): # type: (...) -> ClosableOutputStream """Returns an output stream written timers to transform_id. Args: instruction_id: which instruction this stream belongs to transform_id: the transform_id of the returned stream timer_family_id: the timer family of the written timer """ raise NotImplementedError(type(self)) @abc.abstractmethod def close(self): # type: () -> None """Closes this channel, indicating that all data has been written. Data can continue to be read. If this channel is shared by many instructions, should only be called on worker shutdown. """ raise NotImplementedError(type(self)) class InMemoryDataChannel(DataChannel): """An in-memory implementation of a DataChannel. This channel is two-sided. What is written to one side is read by the other. The inverse() method returns the other side of a instance. """ def __init__(self, inverse=None, data_buffer_time_limit_ms=0): # type: (Optional[InMemoryDataChannel], int) -> None self._inputs = [] # type: List[DataOrTimers] self._data_buffer_time_limit_ms = data_buffer_time_limit_ms self._inverse = inverse or InMemoryDataChannel( self, data_buffer_time_limit_ms=data_buffer_time_limit_ms) def inverse(self): # type: () -> InMemoryDataChannel return self._inverse def input_elements( self, instruction_id, # type: str unused_expected_inputs, # type: Any abort_callback=None # type: Optional[Callable[[], bool]] ): # type: (...) -> Iterator[DataOrTimers] other_inputs = [] for element in self._inputs: if element.instruction_id == instruction_id: if isinstance(element, beam_fn_api_pb2.Elements.Timers): if not element.is_last: yield element if isinstance(element, beam_fn_api_pb2.Elements.Data): if element.data or element.is_last: yield element else: other_inputs.append(element) self._inputs = other_inputs def output_timer_stream( self, instruction_id, # type: str transform_id, # type: str timer_family_id # type: str ): # type: (...) -> ClosableOutputStream def add_to_inverse_output(timer): # type: (bytes) -> None if timer: self._inverse._inputs.append( beam_fn_api_pb2.Elements.Timers( instruction_id=instruction_id, transform_id=transform_id, timer_family_id=timer_family_id, timers=timer, is_last=False)) def close_stream(timer): # type: (bytes) -> None add_to_inverse_output(timer) self._inverse._inputs.append( beam_fn_api_pb2.Elements.Timers( instruction_id=instruction_id, transform_id=transform_id, timer_family_id='', is_last=True)) return ClosableOutputStream.create( add_to_inverse_output, close_stream, self._data_buffer_time_limit_ms) def output_stream(self, instruction_id, transform_id): # type: (str, str) -> ClosableOutputStream def add_to_inverse_output(data): # type: (bytes) -> None self._inverse._inputs.append( # pylint: disable=protected-access beam_fn_api_pb2.Elements.Data( instruction_id=instruction_id, transform_id=transform_id, data=data)) return ClosableOutputStream.create( add_to_inverse_output, add_to_inverse_output, self._data_buffer_time_limit_ms) def close(self): # type: () -> None pass class _GrpcDataChannel(DataChannel): """Base class for implementing a BeamFnData-based DataChannel.""" _WRITES_FINISHED = beam_fn_api_pb2.Elements.Data() def __init__(self, data_buffer_time_limit_ms=0): # type: (int) -> None self._data_buffer_time_limit_ms = data_buffer_time_limit_ms self._to_send = ByteLimitedQueue( maxsize=10000, maxbytes=100 << 20) # type: ByteLimitedQueue[DataOrTimers] self._received = collections.defaultdict( lambda: ByteLimitedQueue(maxsize=5, maxbytes=100 << 20) ) # type: DefaultDict[str, ByteLimitedQueue[DataOrTimers]] # Keep a cache of completed instructions. Data for completed instructions # must be discarded. See input_elements() and _clean_receiving_queue(). # OrderedDict is used as FIFO set with the value being always `True`. self._cleaned_instruction_ids = collections.OrderedDict( ) # type: collections.OrderedDict[str, bool] self._receive_lock = threading.Lock() self._reads_finished = threading.Event() self._closed = False self._exception = None # type: Optional[Exception] def close(self): # type: () -> None self._to_send.put(self._WRITES_FINISHED, 0) self._closed = True def wait(self, timeout=None): # type: (Optional[int]) -> None self._reads_finished.wait(timeout) def _receiving_queue(self, instruction_id): # type: (str) -> Optional[ByteLimitedQueue[DataOrTimers]] """ Gets or creates queue for a instruction_id. Or, returns None if the instruction_id is already cleaned up. This is best-effort as we track a limited number of cleaned-up instructions. """ with self._receive_lock: if instruction_id in self._cleaned_instruction_ids: return None return self._received[instruction_id] def _clean_receiving_queue(self, instruction_id): # type: (str) -> None """ Removes the queue and adds the instruction_id to the cleaned-up list. The instruction_id cannot be reused for new queue. """ with self._receive_lock: # Per-instruction read queue may or may not be created yet when # we mark an instruction as 'cleaned up' when creating # a bundle processor failed, e.g. due to a flake in DoFn.setup(). # We want to mark an instruction as cleaned up regardless. self._received.pop(instruction_id, None) self._cleaned_instruction_ids[instruction_id] = True while len(self._cleaned_instruction_ids) > _MAX_CLEANED_INSTRUCTIONS: self._cleaned_instruction_ids.popitem(last=False) def input_elements( self, instruction_id, # type: str expected_inputs, # type: Collection[Union[str, Tuple[str, str]]] abort_callback=None # type: Optional[Callable[[], bool]] ): # type: (...) -> Iterator[DataOrTimers] """ Generator to retrieve elements for an instruction_id input_elements should be called only once for an instruction_id Args: instruction_id(str): instruction_id for which data is read expected_inputs(collection): expected inputs, include both data and timer. """ received = self._receiving_queue(instruction_id) if received is None: raise RuntimeError('Instruction cleaned up already %s' % instruction_id) done_inputs = set() # type: Set[Union[str, Tuple[str, str]]] abort_callback = abort_callback or (lambda: False) log_interval_sec = 5 * 60 try: start_time = time.time() next_waiting_log_time = start_time + log_interval_sec while len(done_inputs) < len(expected_inputs): try: element = received.get(timeout=1) except queue.Empty: if self._closed: raise RuntimeError('Channel closed prematurely.') if abort_callback(): return if self._exception: raise self._exception from None current_time = time.time() if next_waiting_log_time <= current_time: # If at the same time another instruction is waiting on input queue # to become available, it is a sign of inefficiency in data plane. _LOGGER.info( 'Detected input queue delay longer than %s seconds. ' 'Waiting to receive elements in input queue ' 'for instruction: %s for %.2f seconds.', log_interval_sec, instruction_id, current_time - start_time) next_waiting_log_time = current_time + log_interval_sec else: start_time = time.time() next_waiting_log_time = start_time + log_interval_sec if isinstance(element, beam_fn_api_pb2.Elements.Timers): if element.is_last: done_inputs.add((element.transform_id, element.timer_family_id)) else: yield element elif isinstance(element, beam_fn_api_pb2.Elements.Data): if element.is_last: done_inputs.add(element.transform_id) else: assert element.transform_id not in done_inputs yield element else: raise ValueError('Unexpected input element type %s' % type(element)) finally: # Instruction_ids are not reusable so Clean queue once we are done with # an instruction_id self._clean_receiving_queue(instruction_id) def output_stream(self, instruction_id, transform_id): # type: (str, str) -> ClosableOutputStream def add_to_send_queue(data): # type: (bytes) -> None if data: elem = beam_fn_api_pb2.Elements.Data( instruction_id=instruction_id, transform_id=transform_id, data=data) self._to_send.put(elem, self._get_element_size_bytes(elem)) def close_callback(data): # type: (bytes) -> None add_to_send_queue(data) # End of stream marker. elem = beam_fn_api_pb2.Elements.Data( instruction_id=instruction_id, transform_id=transform_id, is_last=True) self._to_send.put(elem, self._get_element_size_bytes(elem)) return ClosableOutputStream.create( close_callback, add_to_send_queue, self._data_buffer_time_limit_ms) def output_timer_stream( self, instruction_id, # type: str transform_id, # type: str timer_family_id # type: str ): # type: (...) -> ClosableOutputStream def add_to_send_queue(timer): # type: (bytes) -> None if timer: elem = beam_fn_api_pb2.Elements.Timers( instruction_id=instruction_id, transform_id=transform_id, timer_family_id=timer_family_id, timers=timer, is_last=False) self._to_send.put(elem, self._get_element_size_bytes(elem)) def close_callback(timer): # type: (bytes) -> None add_to_send_queue(timer) elem = beam_fn_api_pb2.Elements.Timers( instruction_id=instruction_id, transform_id=transform_id, timer_family_id=timer_family_id, is_last=True) self._to_send.put(elem, self._get_element_size_bytes(elem)) return ClosableOutputStream.create( close_callback, add_to_send_queue, self._data_buffer_time_limit_ms) def _write_outputs(self): # type: () -> Iterator[beam_fn_api_pb2.Elements] stream_done = False while not stream_done: streams = [self._to_send.get()] try: # Coalesce up to 100 other items. total_size_bytes = streams[0].ByteSize() while (total_size_bytes < _DEFAULT_SIZE_FLUSH_THRESHOLD and len(streams) <= 100): data_or_timer = self._to_send.get_nowait() total_size_bytes += data_or_timer.ByteSize() streams.append(data_or_timer) except queue.Empty: pass if streams[-1] is self._WRITES_FINISHED: stream_done = True streams.pop() if streams: data_stream = [] timer_stream = [] for stream in streams: if isinstance(stream, beam_fn_api_pb2.Elements.Timers): timer_stream.append(stream) elif isinstance(stream, beam_fn_api_pb2.Elements.Data): data_stream.append(stream) else: raise ValueError('Unexpected output element type %s' % type(stream)) yield beam_fn_api_pb2.Elements(data=data_stream, timers=timer_stream) def _get_element_size_bytes(self, element): # type: (Union[beam_fn_api_pb2.Elements.Data, beam_fn_api_pb2.Elements.Timers]) -> int if isinstance(element, beam_fn_api_pb2.Elements.Data): return len(element.data) elif isinstance(element, beam_fn_api_pb2.Elements.Timers): return len(element.timers) else: return 0 def _read_inputs(self, elements_iterator): # type: (Iterable[beam_fn_api_pb2.Elements]) -> None next_discard_log_time = 0 # type: float def _put_queue(instruction_id, element): # type: (str, Union[beam_fn_api_pb2.Elements.Data, beam_fn_api_pb2.Elements.Timers]) -> None """ Puts element to the queue of the instruction_id, or discards it if the instruction_id is already cleaned up. """ nonlocal next_discard_log_time start_time = time.time() next_waiting_log_time = start_time + 300 while True: input_queue = self._receiving_queue(instruction_id) if input_queue is None: current_time = time.time() if next_discard_log_time <= current_time: # Log every 10 seconds across all _put_queue calls _LOGGER.info( 'Discard inputs for cleaned up instruction: %s', instruction_id) next_discard_log_time = current_time + 10 return try: input_queue.put( element, self._get_element_size_bytes(element), timeout=1) return except queue.Full: current_time = time.time() if next_waiting_log_time <= current_time: # Log every 5 mins in each _put_queue call _LOGGER.info( 'Waiting on input queue of instruction: %s for %.2f seconds', instruction_id, current_time - start_time) next_waiting_log_time = current_time + 300 try: for elements in elements_iterator: for timer in elements.timers: _put_queue(timer.instruction_id, timer) for data in elements.data: _put_queue(data.instruction_id, data) except Exception as e: if not self._closed: _LOGGER.exception('Failed to read inputs in the data plane.') self._exception = e raise finally: self._closed = True self._reads_finished.set() def set_inputs(self, elements_iterator): # type: (Iterable[beam_fn_api_pb2.Elements]) -> None reader = threading.Thread( target=lambda: self._read_inputs(elements_iterator), name='read_grpc_client_inputs') reader.daemon = True reader.start() class GrpcClientDataChannel(_GrpcDataChannel): """A DataChannel wrapping the client side of a BeamFnData connection.""" def __init__( self, data_stub, # type: beam_fn_api_pb2_grpc.BeamFnDataStub data_buffer_time_limit_ms=0 # type: int ): # type: (...) -> None super().__init__(data_buffer_time_limit_ms) self.set_inputs(data_stub.Data(self._write_outputs())) class BeamFnDataServicer(beam_fn_api_pb2_grpc.BeamFnDataServicer): """Implementation of BeamFnDataServicer for any number of clients""" def __init__( self, data_buffer_time_limit_ms=0 # type: int ): self._lock = threading.Lock() self._connections_by_worker_id = collections.defaultdict( lambda: _GrpcDataChannel(data_buffer_time_limit_ms) ) # type: DefaultDict[str, _GrpcDataChannel] def get_conn_by_worker_id(self, worker_id): # type: (str) -> _GrpcDataChannel with self._lock: return self._connections_by_worker_id[worker_id] def Data( self, elements_iterator, # type: Iterable[beam_fn_api_pb2.Elements] context # type: Any ): # type: (...) -> Iterator[beam_fn_api_pb2.Elements] worker_id = dict(context.invocation_metadata())['worker_id'] data_conn = self.get_conn_by_worker_id(worker_id) data_conn.set_inputs(elements_iterator) for elements in data_conn._write_outputs(): yield elements class DataChannelFactory(metaclass=abc.ABCMeta): """An abstract factory for creating ``DataChannel``.""" @abc.abstractmethod def create_data_channel(self, remote_grpc_port): # type: (beam_fn_api_pb2.RemoteGrpcPort) -> GrpcClientDataChannel """Returns a ``DataChannel`` from the given RemoteGrpcPort.""" raise NotImplementedError(type(self)) @abc.abstractmethod def create_data_channel_from_url(self, url): # type: (str) -> Optional[GrpcClientDataChannel] """Returns a ``DataChannel`` from the given url.""" raise NotImplementedError(type(self)) @abc.abstractmethod def close(self): # type: () -> None """Close all channels that this factory owns.""" raise NotImplementedError(type(self)) def cleanup(self, instruction_id): # type: (str) -> None """Clean up resources for a given instruction.""" pass class GrpcClientDataChannelFactory(DataChannelFactory): """A factory for ``GrpcClientDataChannel``. Caches the created channels by ``data descriptor url``. """ def __init__( self, credentials=None, # type: Any worker_id=None, # type: Optional[str] data_buffer_time_limit_ms=0 # type: int ): # type: (...) -> None self._data_channel_cache = {} # type: Dict[str, GrpcClientDataChannel] self._lock = threading.Lock() self._credentials = None self._worker_id = worker_id self._data_buffer_time_limit_ms = data_buffer_time_limit_ms if credentials is not None: _LOGGER.info('Using secure channel creds.') self._credentials = credentials def create_data_channel_from_url(self, url): # type: (str) -> Optional[GrpcClientDataChannel] if not url: return None if url not in self._data_channel_cache: with self._lock: if url not in self._data_channel_cache: _LOGGER.info('Creating client data channel for %s', url) # Options to have no limits (-1) on the size of the messages # received or sent over the data plane. The actual buffer size # is controlled in a layer above. channel_options = [("grpc.max_receive_message_length", -1), ("grpc.max_send_message_length", -1), ("grpc.service_config", _GRPC_SERVICE_CONFIG)] grpc_channel = None if self._credentials is None: grpc_channel = GRPCChannelFactory.insecure_channel( url, options=channel_options) else: grpc_channel = GRPCChannelFactory.secure_channel( url, self._credentials, options=channel_options) _LOGGER.info('Data channel established.') # Add workerId to the grpc channel grpc_channel = grpc.intercept_channel( grpc_channel, WorkerIdInterceptor(self._worker_id)) self._data_channel_cache[url] = GrpcClientDataChannel( beam_fn_api_pb2_grpc.BeamFnDataStub(grpc_channel), self._data_buffer_time_limit_ms) return self._data_channel_cache[url] def create_data_channel(self, remote_grpc_port): # type: (beam_fn_api_pb2.RemoteGrpcPort) -> GrpcClientDataChannel url = remote_grpc_port.api_service_descriptor.url # TODO(https://github.com/apache/beam/issues/19737): this can return None # if url is falsey, but this seems incorrect, as code that calls this # method seems to always expect non-Optional values. return self.create_data_channel_from_url(url) # type: ignore[return-value] def close(self): # type: () -> None _LOGGER.info('Closing all cached grpc data channels.') for channel in list(self._data_channel_cache.values()): channel.close() self._data_channel_cache.clear() def cleanup(self, instruction_id): # type: (str) -> None for channel in list(self._data_channel_cache.values()): channel._clean_receiving_queue(instruction_id) class InMemoryDataChannelFactory(DataChannelFactory): """A singleton factory for ``InMemoryDataChannel``.""" def __init__(self, in_memory_data_channel): # type: (GrpcClientDataChannel) -> None self._in_memory_data_channel = in_memory_data_channel def create_data_channel(self, unused_remote_grpc_port): # type: (beam_fn_api_pb2.RemoteGrpcPort) -> GrpcClientDataChannel return self._in_memory_data_channel def create_data_channel_from_url(self, url): # type: (Any) -> GrpcClientDataChannel return self._in_memory_data_channel def close(self): # type: () -> None pass ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/data_plane_test.py0000644000175100017510000001217515206113032024371 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.runners.worker.data_plane.""" # pytype: skip-file import itertools import logging import time import unittest import grpc from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.runners.worker import data_plane from apache_beam.runners.worker.worker_id_interceptor import WorkerIdInterceptor from apache_beam.utils import thread_pool_executor class DataChannelTest(unittest.TestCase): def test_grpc_data_channel(self): self._grpc_data_channel_test() def test_time_based_flush_grpc_data_channel(self): self._grpc_data_channel_test(True) def _grpc_data_channel_test(self, time_based_flush=False): if time_based_flush: data_servicer = data_plane.BeamFnDataServicer( data_buffer_time_limit_ms=100) else: data_servicer = data_plane.BeamFnDataServicer() worker_id = 'worker_0' data_channel_service = \ data_servicer.get_conn_by_worker_id(worker_id) server = grpc.server(thread_pool_executor.shared_unbounded_instance()) beam_fn_api_pb2_grpc.add_BeamFnDataServicer_to_server(data_servicer, server) test_port = server.add_insecure_port('[::]:0') server.start() grpc_channel = grpc.insecure_channel('localhost:%s' % test_port) # Add workerId to the grpc channel grpc_channel = grpc.intercept_channel( grpc_channel, WorkerIdInterceptor(worker_id)) data_channel_stub = beam_fn_api_pb2_grpc.BeamFnDataStub(grpc_channel) if time_based_flush: data_channel_client = data_plane.GrpcClientDataChannel( data_channel_stub, data_buffer_time_limit_ms=100) else: data_channel_client = data_plane.GrpcClientDataChannel(data_channel_stub) try: self._data_channel_test( data_channel_service, data_channel_client, time_based_flush) finally: data_channel_client.close() data_channel_service.close() data_channel_client.wait() data_channel_service.wait() def test_in_memory_data_channel(self): channel = data_plane.InMemoryDataChannel() self._data_channel_test(channel, channel.inverse()) def _data_channel_test(self, server, client, time_based_flush=False): self._data_channel_test_one_direction(server, client, time_based_flush) self._data_channel_test_one_direction(client, server, time_based_flush) def _data_channel_test_one_direction( self, from_channel, to_channel, time_based_flush): transform_1 = '1' transform_2 = '2' # Single write. stream01 = from_channel.output_stream('0', transform_1) stream01.write(b'abc') if not time_based_flush: stream01.close() self.assertEqual( list( itertools.islice(to_channel.input_elements('0', [transform_1]), 1)), [ beam_fn_api_pb2.Elements.Data( instruction_id='0', transform_id=transform_1, data=b'abc') ]) # Multiple interleaved writes to multiple instructions. stream11 = from_channel.output_stream('1', transform_1) stream11.write(b'abc') stream21 = from_channel.output_stream('2', transform_1) stream21.write(b'def') if not time_based_flush: stream11.close() self.assertEqual( list( itertools.islice(to_channel.input_elements('1', [transform_1]), 1)), [ beam_fn_api_pb2.Elements.Data( instruction_id='1', transform_id=transform_1, data=b'abc') ]) if time_based_flush: # Wait to ensure stream21 is flushed before stream22. # Because the flush callback is invoked periodically starting from when a # stream is constructed, there is no guarantee that one stream's callback # is called before the other. time.sleep(0.1) else: stream21.close() stream22 = from_channel.output_stream('2', transform_2) stream22.write(b'ghi') if not time_based_flush: stream22.close() self.assertEqual( list( itertools.islice( to_channel.input_elements('2', [transform_1, transform_2]), 2)), [ beam_fn_api_pb2.Elements.Data( instruction_id='2', transform_id=transform_1, data=b'def'), beam_fn_api_pb2.Elements.Data( instruction_id='2', transform_id=transform_2, data=b'ghi') ]) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/data_sampler.py0000644000175100017510000003167715206113032023706 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Functionaliry for sampling elements during bundle execution.""" # pytype: skip-file from __future__ import annotations import collections import logging import threading import time import traceback from dataclasses import dataclass from threading import Timer from typing import Any from typing import Iterable from typing import Optional from typing import Union from apache_beam.coders.coder_impl import CoderImpl from apache_beam.coders.coder_impl import WindowedValueCoderImpl from apache_beam.coders.coders import Coder from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.utils.windowed_value import WindowedValue _LOGGER = logging.getLogger(__name__) class SampleTimer: """Periodic timer for sampling elements.""" def __init__(self, timeout_secs: float, sampler: OutputSampler) -> None: self._target_timeout_secs = timeout_secs self._timeout_secs = min(timeout_secs, 0.5) if timeout_secs > 0 else 0.0 self._timer = Timer(self._timeout_secs, self.sample) self._sampler = sampler self._sample_duration_secs = 0.0 def reset(self) -> None: # For the first 30 seconds, sample every 0.5 seconds. After that, sample at # the normal rate. if self._sample_duration_secs >= 30.0: self._timeout_secs = self._target_timeout_secs self._sample_duration_secs += self._timeout_secs self._timer.cancel() self._timer = Timer(self._timeout_secs, self.sample) self._timer.start() def stop(self) -> None: self._timer.cancel() def sample(self) -> None: self._sampler.sample() self.reset() @dataclass class ExceptionMetadata: # The repr-ified Exception. msg: str # The transform where the exception occurred. transform_id: str # The instruction when the exception occurred. instruction_id: str @dataclass class ElementSampler: """Record class to hold sampled elements. This class is used as an optimization to quickly sample elements. This is a shared reference between the Operation and the OutputSampler. """ # Is true iff the `el` has been set with a sample. has_element: bool = False # The sampled element. Note that `None` is a valid element and cannot be uesd # as a sentintel to check if there is a sample. Use the `has_element` flag to # check for this case. el: Any = None class OutputSampler: """Represents a way to sample an output of a PTransform. This is configurable to only keep `max_samples` (see constructor) sampled elements in memory. Samples are taken every `sample_every_sec`. """ def __init__( self, coder: Coder, max_samples: int = 10, sample_every_sec: float = 5) -> None: self._samples: collections.deque[Any] = collections.deque( maxlen=max_samples) self._samples_lock: threading.Lock = threading.Lock() self._coder_impl: CoderImpl = coder.get_impl() self._sample_timer = SampleTimer(sample_every_sec, self) self.element_sampler = ElementSampler() self.element_sampler.has_element = False self._exceptions: collections.deque[tuple[ Any, ExceptionMetadata]] = collections.deque(maxlen=max_samples) # For testing, it's easier to disable the Timer and manually sample. if sample_every_sec > 0: self._sample_timer.reset() def stop(self) -> None: """Stops sampling.""" self._sample_timer.stop() def remove_windowed_value(self, el: Union[WindowedValue, Any]) -> Any: """Retrieves the value from the WindowedValue. The Python SDK passes elements as WindowedValues, which may not match the coder for that particular PCollection. """ while isinstance(el, WindowedValue): el = el.value return el def flush(self, clear: bool = True) -> list[beam_fn_api_pb2.SampledElement]: """Returns all samples and optionally clears buffer if clear is True.""" with self._samples_lock: # TODO(rohdesamuel): There can duplicates between the exceptions and # samples. This happens when the OutputSampler samples during an # exception. The fix is to create a OutputSampler per process bundle. # Until then use a set to keep track of the elements. seen = set(id(el) for el, _ in self._exceptions) if isinstance(self._coder_impl, WindowedValueCoderImpl): exceptions = [s for s in self._exceptions] samples = [s for s in self._samples if id(s) not in seen] else: exceptions = [(self.remove_windowed_value(a), b) for a, b in self._exceptions] samples = [ self.remove_windowed_value(s) for s in self._samples if id(s) not in seen ] # Encode in the nested context b/c this ensures that the SDK can decode # the bytes with the ToStringFn. if clear: self._samples.clear() self._exceptions.clear() ret = [] try: ret = [ beam_fn_api_pb2.SampledElement( element=self._coder_impl.encode_nested(s), ) for s in samples ] except Exception as e: # pylint: disable=broad-except _LOGGER.warning('Could not encode sampled values: %s' % e) try: ret.extend( beam_fn_api_pb2.SampledElement( element=self._coder_impl.encode_nested(s), exception=beam_fn_api_pb2.SampledElement.Exception( instruction_id=exn.instruction_id, transform_id=exn.transform_id, error=exn.msg)) for s, exn in exceptions) except Exception as e: # pylint: disable=broad-except _LOGGER.warning('Could not encode sampled exception values: %s' % e) return ret def sample(self) -> None: """Samples the given element to an internal buffer.""" with self._samples_lock: if self.element_sampler.has_element: self._samples.append(self.element_sampler.el) self.element_sampler.has_element = False def sample_exception( self, el: Any, exc_info: Any, transform_id: str, instruction_id: str) -> None: """Adds the given exception to the samples.""" with self._samples_lock: err_string = ''.join(traceback.format_exception(*exc_info)) self._exceptions.append( (el, ExceptionMetadata(err_string, transform_id, instruction_id))) class DataSampler: """A class for querying any samples generated during execution. This class is meant to be a singleton with regard to a particular `sdk_worker.SdkHarness`. When creating the operators, individual `OutputSampler`s are created from `DataSampler.initialize_samplers`. This allows for multi-threaded sampling of a PCollection across the SdkHarness. Samples generated during execution can then be sampled with the `samples` method. This filters samples from the given pcollection ids. """ def __init__( self, max_samples: int = 10, sample_every_sec: float = 30, sample_only_exceptions: bool = False, clock=None) -> None: # Key is PCollection id. Is guarded by the _samplers_lock. self._samplers: dict[str, OutputSampler] = {} # Bundles are processed in parallel, so new samplers may be added when the # runner queries for samples. self._samplers_lock: threading.Lock = threading.Lock() self._max_samples = max_samples self._sample_every_sec = 0.0 if sample_only_exceptions else sample_every_sec self._samplers_by_output: dict[str, list[OutputSampler]] = {} self._clock = clock _ENABLE_DATA_SAMPLING = 'enable_data_sampling' _ENABLE_ALWAYS_ON_EXCEPTION_SAMPLING = 'enable_always_on_exception_sampling' _DISABLE_ALWAYS_ON_EXCEPTION_SAMPLING = 'disable_always_on_exception_sampling' @staticmethod def create(sdk_pipeline_options: PipelineOptions, **kwargs): experiments = sdk_pipeline_options.view_as(DebugOptions).experiments or [] # When true, enables only the sampling of exceptions. always_on_exception_sampling = ( DataSampler._ENABLE_ALWAYS_ON_EXCEPTION_SAMPLING in experiments and DataSampler._DISABLE_ALWAYS_ON_EXCEPTION_SAMPLING not in experiments) # When true, enables the sampling of all PCollections and exceptions. enable_data_sampling = DataSampler._ENABLE_DATA_SAMPLING in experiments if enable_data_sampling or always_on_exception_sampling: sample_only_exceptions = ( always_on_exception_sampling and not enable_data_sampling) return DataSampler( sample_only_exceptions=sample_only_exceptions, **kwargs) else: return None def stop(self) -> None: """Stops all sampling, does not clear samplers in case there are outstanding samples. """ with self._samplers_lock: for sampler in self._samplers.values(): sampler.stop() def sampler_for_output(self, transform_id: str, output_index: int) -> Optional[OutputSampler]: """Returns the OutputSampler for the given output.""" try: with self._samplers_lock: outputs = self._samplers_by_output[transform_id] return outputs[output_index] except KeyError: _LOGGER.warning( f'Out-of-bounds access for transform "{transform_id}" ' + 'and output "{output_index}" OutputSampler. This may ' + 'indicate that the transform was improperly ' + 'initialized with the DataSampler.') return None def initialize_samplers( self, transform_id: str, descriptor: beam_fn_api_pb2.ProcessBundleDescriptor, coder_factory) -> list[OutputSampler]: """Creates the OutputSamplers for the given PTransform. This initializes the samplers only once per PCollection Id. Note that an OutputSampler is created per PCollection and an ElementSampler is created per OutputSampler. This means that multiple ProcessBundles can and will share the same ElementSampler for a given PCollection. """ transform_proto = descriptor.transforms[transform_id] with self._samplers_lock: if transform_id in self._samplers_by_output: return self._samplers_by_output[transform_id] # Initialize the samplers. for pcoll_id in transform_proto.outputs.values(): # Only initialize new PCollections. if pcoll_id in self._samplers: continue # Create the sampler with the corresponding coder. coder_id = descriptor.pcollections[pcoll_id].coder_id coder = coder_factory(coder_id) sampler = OutputSampler( coder, self._max_samples, self._sample_every_sec) self._samplers[pcoll_id] = sampler # Next update the lookup table for ElementSamplers for a given PTransform. # Operations look up the ElementSampler for an output based on the index # of the tag in the PTransform's outputs. The following code intializes # the array with ElementSamplers in the correct indices. outputs = transform_proto.outputs samplers = [self._samplers[pcoll_id] for pcoll_id in outputs.values()] self._samplers_by_output[transform_id] = samplers return samplers def samples( self, pcollection_ids: Optional[Iterable[str]] = None ) -> beam_fn_api_pb2.SampleDataResponse: """Returns samples filtered PCollection ids. All samples from the given PCollections are returned. Empty lists are wildcards. """ ret = beam_fn_api_pb2.SampleDataResponse() with self._samplers_lock: samplers = self._samplers.copy() for pcoll_id in samplers: if pcollection_ids and pcoll_id not in pcollection_ids: continue samples = samplers[pcoll_id].flush() if samples: ret.element_samples[pcoll_id].elements.extend(samples) return ret def wait_for_samples( self, pcollection_ids: list[str]) -> beam_fn_api_pb2.SampleDataResponse: """Waits for samples to exist for the given PCollections (only testing).""" now = time.time() end = now + 30 samples = beam_fn_api_pb2.SampleDataResponse() while now < end: time.sleep(0.1) now = time.time() samples.MergeFrom(self.samples(pcollection_ids)) if not samples: continue has_all = all( pcoll_id in samples.element_samples for pcoll_id in pcollection_ids) if has_all: break return samples ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/data_sampler_test.py0000644000175100017510000005736115206113032024743 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import sys import time import traceback import unittest from typing import Any from typing import Optional from apache_beam.coders import FastPrimitivesCoder from apache_beam.coders import WindowedValueCoder from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.runners.worker.data_sampler import DataSampler from apache_beam.runners.worker.data_sampler import OutputSampler from apache_beam.transforms.window import GlobalWindow from apache_beam.utils.windowed_value import WindowedValue MAIN_TRANSFORM_ID = 'transform' MAIN_PCOLLECTION_ID = 'pcoll' PRIMITIVES_CODER = FastPrimitivesCoder() class DataSamplerTest(unittest.TestCase): def make_test_descriptor( self, outputs: Optional[list[str]] = None, transforms: Optional[list[str]] = None ) -> beam_fn_api_pb2.ProcessBundleDescriptor: outputs = outputs or [MAIN_PCOLLECTION_ID] transforms = transforms or [MAIN_TRANSFORM_ID] descriptor = beam_fn_api_pb2.ProcessBundleDescriptor() for transform_id in transforms: transform = descriptor.transforms[transform_id] for output in outputs: transform.outputs[output] = output return descriptor def setUp(self): self.data_sampler = DataSampler.create( PipelineOptions(experiments=[DataSampler._ENABLE_DATA_SAMPLING]), sample_every_sec=0.1) def tearDown(self): self.data_sampler.stop() def primitives_coder_factory(self, _): return PRIMITIVES_CODER def gen_sample( self, data_sampler: DataSampler, element: Any, output_index: int, transform_id: str = MAIN_TRANSFORM_ID): """Generates a sample for the given transform's output.""" element_sampler = self.data_sampler.sampler_for_output( transform_id, output_index).element_sampler element_sampler.el = element element_sampler.has_element = True def test_single_output(self): """Simple test for a single sample.""" descriptor = self.make_test_descriptor() self.data_sampler.initialize_samplers( MAIN_TRANSFORM_ID, descriptor, self.primitives_coder_factory) self.gen_sample(self.data_sampler, 'a', output_index=0) expected_sample = beam_fn_api_pb2.SampleDataResponse( element_samples={ MAIN_PCOLLECTION_ID: beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('a')) ]) }) samples = self.data_sampler.wait_for_samples([MAIN_PCOLLECTION_ID]) self.assertEqual(samples, expected_sample) def test_not_initialized(self): """Tests that transforms fail gracefully if not properly initialized.""" with self.assertLogs() as cm: self.data_sampler.sampler_for_output(MAIN_TRANSFORM_ID, 0) self.assertRegex(cm.output[0], 'Out-of-bounds access.*') def map_outputs_to_indices( self, outputs, descriptor, transform_id=MAIN_TRANSFORM_ID): tag_list = list(descriptor.transforms[transform_id].outputs) return {output: tag_list.index(output) for output in outputs} def test_sampler_mapping(self): """Tests that the ElementSamplers are created for the correct output.""" # Initialize the DataSampler with the following outputs. The order here may # get shuffled when inserting into the descriptor. pcollection_ids = ['o0', 'o1', 'o2'] descriptor = self.make_test_descriptor(outputs=pcollection_ids) samplers = self.data_sampler.initialize_samplers( MAIN_TRANSFORM_ID, descriptor, self.primitives_coder_factory) # Create a map from the PCollection id to the index into the transform # output. This mirrors what happens when operators are created. The index of # an output is where in the PTransform.outputs it is located (when the map # is converted to a list). outputs = self.map_outputs_to_indices(pcollection_ids, descriptor) # Assert that the mapping is correct, i.e. that we can go from the # PCollection id -> output index and that this is the same as the created # samplers. index = outputs['o0'] self.assertEqual( self.data_sampler.sampler_for_output(MAIN_TRANSFORM_ID, index).element_sampler, samplers[index].element_sampler) index = outputs['o1'] self.assertEqual( self.data_sampler.sampler_for_output(MAIN_TRANSFORM_ID, index).element_sampler, samplers[index].element_sampler) index = outputs['o2'] self.assertEqual( self.data_sampler.sampler_for_output(MAIN_TRANSFORM_ID, index).element_sampler, samplers[index].element_sampler) def test_multiple_outputs(self): """Tests that multiple PCollections have their own sampler.""" pcollection_ids = ['o0', 'o1', 'o2'] descriptor = self.make_test_descriptor(outputs=pcollection_ids) outputs = self.map_outputs_to_indices(pcollection_ids, descriptor) self.data_sampler.initialize_samplers( MAIN_TRANSFORM_ID, descriptor, self.primitives_coder_factory) self.gen_sample(self.data_sampler, 'a', output_index=outputs['o0']) self.gen_sample(self.data_sampler, 'b', output_index=outputs['o1']) self.gen_sample(self.data_sampler, 'c', output_index=outputs['o2']) samples = self.data_sampler.wait_for_samples(['o0', 'o1', 'o2']) expected_samples = beam_fn_api_pb2.SampleDataResponse( element_samples={ 'o0': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('a')) ]), 'o1': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('b')) ]), 'o2': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('c')) ]), }) self.assertEqual(samples, expected_samples) def test_multiple_transforms(self): """Test that multiple transforms with the same PCollections can be sampled. """ # Initialize two transform both with the same two outputs. pcollection_ids = ['o0', 'o1'] descriptor = self.make_test_descriptor( outputs=pcollection_ids, transforms=['t0', 't1']) t0_outputs = self.map_outputs_to_indices( pcollection_ids, descriptor, transform_id='t0') t1_outputs = self.map_outputs_to_indices( pcollection_ids, descriptor, transform_id='t1') self.data_sampler.initialize_samplers( 't0', descriptor, self.primitives_coder_factory) self.data_sampler.initialize_samplers( 't1', descriptor, self.primitives_coder_factory) # The OutputSampler is on a different thread so we don't test the same # PCollections to ensure that no data race occurs. self.gen_sample( self.data_sampler, 'a', output_index=t0_outputs['o0'], transform_id='t0') self.gen_sample( self.data_sampler, 'd', output_index=t1_outputs['o1'], transform_id='t1') expected_samples = beam_fn_api_pb2.SampleDataResponse( element_samples={ 'o0': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('a')) ]), 'o1': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('d')) ]), }) samples = self.data_sampler.wait_for_samples(['o0', 'o1']) self.assertEqual(samples, expected_samples) self.gen_sample( self.data_sampler, 'b', output_index=t0_outputs['o1'], transform_id='t0') self.gen_sample( self.data_sampler, 'c', output_index=t1_outputs['o0'], transform_id='t1') expected_samples = beam_fn_api_pb2.SampleDataResponse( element_samples={ 'o0': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('c')) ]), 'o1': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('b')) ]), }) samples = self.data_sampler.wait_for_samples(['o0', 'o1']) self.assertEqual(samples, expected_samples) def test_sample_filters_single_pcollection_ids(self): """Tests the samples can be filtered based on a single pcollection id.""" pcollection_ids = ['o0', 'o1', 'o2'] descriptor = self.make_test_descriptor(outputs=pcollection_ids) outputs = self.map_outputs_to_indices(pcollection_ids, descriptor) self.data_sampler.initialize_samplers( MAIN_TRANSFORM_ID, descriptor, self.primitives_coder_factory) self.gen_sample(self.data_sampler, 'a', output_index=outputs['o0']) self.gen_sample(self.data_sampler, 'b', output_index=outputs['o1']) self.gen_sample(self.data_sampler, 'c', output_index=outputs['o2']) samples = self.data_sampler.wait_for_samples(['o0']) expected_samples = beam_fn_api_pb2.SampleDataResponse( element_samples={ 'o0': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('a')) ]), }) self.assertEqual(samples, expected_samples) samples = self.data_sampler.wait_for_samples(['o1']) expected_samples = beam_fn_api_pb2.SampleDataResponse( element_samples={ 'o1': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('b')) ]), }) self.assertEqual(samples, expected_samples) samples = self.data_sampler.wait_for_samples(['o2']) expected_samples = beam_fn_api_pb2.SampleDataResponse( element_samples={ 'o2': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('c')) ]), }) self.assertEqual(samples, expected_samples) def test_sample_filters_multiple_pcollection_ids(self): """Tests the samples can be filtered based on a multiple pcollection ids.""" pcollection_ids = ['o0', 'o1', 'o2'] descriptor = self.make_test_descriptor(outputs=pcollection_ids) outputs = self.map_outputs_to_indices(pcollection_ids, descriptor) self.data_sampler.initialize_samplers( MAIN_TRANSFORM_ID, descriptor, self.primitives_coder_factory) self.gen_sample(self.data_sampler, 'a', output_index=outputs['o0']) self.gen_sample(self.data_sampler, 'b', output_index=outputs['o1']) self.gen_sample(self.data_sampler, 'c', output_index=outputs['o2']) samples = self.data_sampler.wait_for_samples(['o0', 'o2']) expected_samples = beam_fn_api_pb2.SampleDataResponse( element_samples={ 'o0': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('a')) ]), 'o2': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('c')) ]), }) self.assertEqual(samples, expected_samples) def test_can_sample_exceptions(self): """Tests that exceptions sampled can be queried by the DataSampler.""" descriptor = self.make_test_descriptor() self.data_sampler.initialize_samplers( MAIN_TRANSFORM_ID, descriptor, self.primitives_coder_factory) sampler = self.data_sampler.sampler_for_output(MAIN_TRANSFORM_ID, 0) exc_info = None try: raise Exception('test') except Exception: exc_info = sys.exc_info() sampler.sample_exception('a', exc_info, MAIN_TRANSFORM_ID, 'instid') samples = self.data_sampler.wait_for_samples([MAIN_PCOLLECTION_ID]) self.assertGreater(len(samples.element_samples), 0) def test_create_experiments(self): """Tests that the experiments correctly make the DataSampler.""" enable_exception_exp = DataSampler._ENABLE_ALWAYS_ON_EXCEPTION_SAMPLING disable_exception_exp = DataSampler._DISABLE_ALWAYS_ON_EXCEPTION_SAMPLING enable_sampling_exp = DataSampler._ENABLE_DATA_SAMPLING self.assertIsNone(DataSampler.create(PipelineOptions())) exp = [disable_exception_exp] self.assertIsNone(DataSampler.create(PipelineOptions(experiments=exp))) exp = [enable_exception_exp, disable_exception_exp] self.assertIsNone(DataSampler.create(PipelineOptions(experiments=exp))) exp = [enable_exception_exp] self.assertIsNotNone(DataSampler.create(PipelineOptions(experiments=exp))) exp = [enable_sampling_exp] self.assertIsNotNone(DataSampler.create(PipelineOptions(experiments=exp))) exp = [enable_sampling_exp, enable_exception_exp, disable_exception_exp] self.assertIsNotNone(DataSampler.create(PipelineOptions(experiments=exp))) def test_samples_all_with_both_experiments(self): """Tests that the using both sampling experiments samples everything.""" self.data_sampler = DataSampler.create( PipelineOptions( experiments=[ DataSampler._ENABLE_DATA_SAMPLING, DataSampler._ENABLE_ALWAYS_ON_EXCEPTION_SAMPLING ]), sample_every_sec=0.1) # Create a descriptor with one transform with two outputs, 'a' and 'b'. descriptor = self.make_test_descriptor(outputs=['a', 'b']) self.data_sampler.initialize_samplers( MAIN_TRANSFORM_ID, descriptor, self.primitives_coder_factory) # Get the samples for the two outputs. # N.B. the order of the samplers is not guaranteed due to Protobuf not # guaranteeing map iteration order. first_sampler = self.data_sampler.sampler_for_output(MAIN_TRANSFORM_ID, 0) second_sampler = self.data_sampler.sampler_for_output(MAIN_TRANSFORM_ID, 1) # Sample an exception for the output 'a', this will show up in the final # samples response. exc_info = None try: raise Exception('test') except Exception: exc_info = sys.exc_info() first_sampler.sample_exception( 'first', exc_info, MAIN_TRANSFORM_ID, 'instid') # Sample a normal element for the output 'b', this will not show up in the # final samples response. second_sampler.element_sampler.el = 'second' second_sampler.element_sampler.has_element = True samples = self.data_sampler.wait_for_samples(['a', 'b']) self.assertEqual(len(samples.element_samples), 2) sample_elements = list( s.elements[0] for s in samples.element_samples.values()) num_exceptions = sum( 1 for element in sample_elements if element.HasField('exception')) self.assertEqual( num_exceptions, 1, "Only one of the samples should have an exception, found: {}".format( sample_elements)) def test_only_sample_exceptions(self): """Tests that the exception sampling experiment only samples exceptions.""" self.data_sampler = DataSampler.create( PipelineOptions( experiments=[DataSampler._ENABLE_ALWAYS_ON_EXCEPTION_SAMPLING]), sample_every_sec=0.1) # Create a descriptor with one transform with two outputs, 'a' and 'b'. descriptor = self.make_test_descriptor(outputs=['a', 'b']) self.data_sampler.initialize_samplers( MAIN_TRANSFORM_ID, descriptor, self.primitives_coder_factory) # Get the samples for the two outputs. # N.B. the order of the samplers is not guaranteed due to Protobuf not # guaranteeing map iteration order. first_sampler = self.data_sampler.sampler_for_output(MAIN_TRANSFORM_ID, 0) second_sampler = self.data_sampler.sampler_for_output(MAIN_TRANSFORM_ID, 1) # Sample an exception for the output 'a', this will show up in the final # samples response. exc_info = None try: raise Exception('test') except Exception: exc_info = sys.exc_info() first_sampler.sample_exception( 'first', exc_info, MAIN_TRANSFORM_ID, 'instid') # Sample a normal element for the output 'b', this will not show up in the # final samples response. second_sampler.element_sampler.el = 'second' second_sampler.element_sampler.has_element = True samples = self.data_sampler.wait_for_samples([]) self.assertEqual(len(samples.element_samples), 1) value = list(samples.element_samples.values())[0] self.assertIsNotNone(value.elements[0].exception) class OutputSamplerTest(unittest.TestCase): def tearDown(self): self.sampler.stop() def wait_for_samples(self, output_sampler: OutputSampler, expected_num: int): """Waits for the expected number of samples for the given sampler.""" now = time.time() end = now + 30 while now < end: time.sleep(0.1) now = time.time() samples = output_sampler.flush(clear=False) if not samples: continue if len(samples) == expected_num: return samples self.assertLess(now, end, 'Timed out waiting for samples') def test_can_sample(self): """Tests that the underlying timer can sample.""" self.sampler = OutputSampler(PRIMITIVES_CODER, sample_every_sec=0.05) element_sampler = self.sampler.element_sampler element_sampler.el = 'a' element_sampler.has_element = True self.wait_for_samples(self.sampler, expected_num=1) self.assertEqual( self.sampler.flush(), [ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('a')) ]) def test_acts_like_circular_buffer(self): """Tests that the buffer overwrites old samples.""" self.sampler = OutputSampler( PRIMITIVES_CODER, max_samples=2, sample_every_sec=0) element_sampler = self.sampler.element_sampler for i in range(10): element_sampler.el = i element_sampler.has_element = True self.sampler.sample() self.assertEqual( self.sampler.flush(), [ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested(i)) for i in (8, 9) ]) def test_samples_multiple_times(self): """Tests that the underlying timer repeats.""" self.sampler = OutputSampler( PRIMITIVES_CODER, max_samples=10, sample_every_sec=0.05) # Always samples the first ten. for i in range(10): self.sampler.element_sampler.el = i self.sampler.element_sampler.has_element = True self.wait_for_samples(self.sampler, i + 1) self.assertEqual( self.sampler.flush(), [ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested(i)) for i in range(10) ]) def test_can_sample_windowed_value(self): """Tests that values with WindowedValueCoders are sampled wholesale.""" coder = WindowedValueCoder(FastPrimitivesCoder()) value = WindowedValue('Hello, World!', 0, [GlobalWindow()]) self.sampler = OutputSampler(coder, sample_every_sec=0) element_sampler = self.sampler.element_sampler element_sampler.el = value element_sampler.has_element = True self.sampler.sample() self.assertEqual( self.sampler.flush(), [beam_fn_api_pb2.SampledElement(element=coder.encode_nested(value))]) def test_can_sample_non_windowed_value(self): """Tests that windowed values with WindowedValueCoders sample only the value. This is important because the Python SDK wraps all values in a WindowedValue even if the coder is not a WindowedValueCoder. In this case, the value must be retrieved from the WindowedValue to match the correct coder. """ value = WindowedValue('Hello, World!', 0, [GlobalWindow()]) self.sampler = OutputSampler(PRIMITIVES_CODER, sample_every_sec=0) element_sampler = self.sampler.element_sampler element_sampler.el = value element_sampler.has_element = True self.sampler.sample() self.assertEqual( self.sampler.flush(), [ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('Hello, World!')) ]) def test_can_sample_exceptions(self): """Tests that exceptions are sampled.""" val = WindowedValue('Hello, World!', 0, [GlobalWindow()]) exc_info = None try: raise Exception('test') except Exception: exc_info = sys.exc_info() err_string = ''.join(traceback.format_exception(*exc_info)) self.sampler = OutputSampler(PRIMITIVES_CODER, sample_every_sec=0) self.sampler.sample_exception( el=val, exc_info=exc_info, transform_id='tid', instruction_id='instid') self.assertEqual( self.sampler.flush(), [ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('Hello, World!'), exception=beam_fn_api_pb2.SampledElement.Exception( instruction_id='instid', transform_id='tid', error=err_string)) ]) def test_can_sample_multiple_exceptions(self): """Tests that multiple exceptions in the same PCollection are sampled.""" exc_info = None try: raise Exception('test') except Exception: exc_info = sys.exc_info() err_string = ''.join(traceback.format_exception(*exc_info)) self.sampler = OutputSampler(PRIMITIVES_CODER, sample_every_sec=0) self.sampler.sample_exception( el='a', exc_info=exc_info, transform_id='tid', instruction_id='instid') self.sampler.sample_exception( el='b', exc_info=exc_info, transform_id='tid', instruction_id='instid') self.assertEqual( self.sampler.flush(), [ beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('a'), exception=beam_fn_api_pb2.SampledElement.Exception( instruction_id='instid', transform_id='tid', error=err_string)), beam_fn_api_pb2.SampledElement( element=PRIMITIVES_CODER.encode_nested('b'), exception=beam_fn_api_pb2.SampledElement.Exception( instruction_id='instid', transform_id='tid', error=err_string)), ]) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/log_handler.py0000644000175100017510000002070215206113032023513 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Beam fn API log handler.""" # pytype: skip-file # mypy: disallow-untyped-defs import logging import math import queue import sys import threading import time import traceback from typing import Iterable from typing import Iterator from typing import Union from typing import cast import grpc from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.portability.api import endpoints_pb2 from apache_beam.runners.worker import statesampler from apache_beam.runners.worker.channel_factory import GRPCChannelFactory from apache_beam.runners.worker.worker_id_interceptor import WorkerIdInterceptor from apache_beam.utils.sentinel import Sentinel # Mapping from logging levels to LogEntry levels. LOG_LEVEL_TO_LOGENTRY_MAP = { logging.FATAL: beam_fn_api_pb2.LogEntry.Severity.CRITICAL, logging.ERROR: beam_fn_api_pb2.LogEntry.Severity.ERROR, logging.WARNING: beam_fn_api_pb2.LogEntry.Severity.WARN, logging.INFO: beam_fn_api_pb2.LogEntry.Severity.INFO, logging.DEBUG: beam_fn_api_pb2.LogEntry.Severity.DEBUG, logging.NOTSET: beam_fn_api_pb2.LogEntry.Severity.UNSPECIFIED, -float('inf'): beam_fn_api_pb2.LogEntry.Severity.DEBUG, } # Mapping from LogEntry levels to logging levels LOGENTRY_TO_LOG_LEVEL_MAP = { beam_fn_api_pb2.LogEntry.Severity.CRITICAL: logging.CRITICAL, beam_fn_api_pb2.LogEntry.Severity.ERROR: logging.ERROR, beam_fn_api_pb2.LogEntry.Severity.WARN: logging.WARNING, beam_fn_api_pb2.LogEntry.Severity.NOTICE: logging.INFO + 1, beam_fn_api_pb2.LogEntry.Severity.INFO: logging.INFO, beam_fn_api_pb2.LogEntry.Severity.DEBUG: logging.DEBUG, beam_fn_api_pb2.LogEntry.Severity.TRACE: logging.DEBUG - 1, beam_fn_api_pb2.LogEntry.Severity.UNSPECIFIED: logging.NOTSET, } class FnApiLogRecordHandler(logging.Handler): """A handler that writes log records to the fn API.""" # Maximum number of log entries in a single stream request. _MAX_BATCH_SIZE = 1000 # Used to indicate the end of stream. _FINISHED = Sentinel.sentinel # Size of the queue used to buffer messages. Once full, messages will be # dropped. If the average log size is 1KB this may use up to 10MB of memory. _QUEUE_SIZE = 10000 def __init__( self, log_service_descriptor: endpoints_pb2.ApiServiceDescriptor) -> None: super().__init__() self._alive = True self._dropped_logs = 0 self._log_entry_queue: queue.Queue[Union[beam_fn_api_pb2.LogEntry, Sentinel]] = queue.Queue( maxsize=self._QUEUE_SIZE) ch = GRPCChannelFactory.insecure_channel(log_service_descriptor.url) # Make sure the channel is ready to avoid [BEAM-4649] grpc.channel_ready_future(ch).result(timeout=60) self._log_channel = grpc.intercept_channel(ch, WorkerIdInterceptor()) self._reader = threading.Thread( target=lambda: self._read_log_control_messages(), name='read_log_control_messages') self._reader.daemon = True self._reader.start() def connect(self) -> Iterable: if hasattr(self, '_logging_stub'): del self._logging_stub # type: ignore[has-type] self._logging_stub = beam_fn_api_pb2_grpc.BeamFnLoggingStub( self._log_channel) return self._logging_stub.Logging(self._write_log_entries()) def map_log_level( self, level: int) -> beam_fn_api_pb2.LogEntry.Severity.Enum.ValueType: try: return LOG_LEVEL_TO_LOGENTRY_MAP[level] except KeyError: return max( beam_level for python_level, beam_level in LOG_LEVEL_TO_LOGENTRY_MAP.items() if python_level <= level) def emit(self, record: logging.LogRecord) -> None: log_entry = beam_fn_api_pb2.LogEntry() log_entry.severity = self.map_log_level(record.levelno) try: log_entry.message = self.format(record) except Exception: # record.msg could be an arbitrary object, convert it to a string first. log_entry.message = ( "Failed to format '%s' with args '%s' during logging." % (str(record.msg), record.args)) log_entry.thread = record.threadName # type: ignore[assignment] log_entry.log_location = '%s:%s' % ( record.pathname or record.module, record.lineno or record.funcName) (fraction, seconds) = math.modf(record.created) nanoseconds = 1e9 * fraction log_entry.timestamp.seconds = int(seconds) log_entry.timestamp.nanos = int(nanoseconds) if record.exc_info: log_entry.trace = ''.join(traceback.format_exception(*record.exc_info)) instruction_id = statesampler.get_current_instruction_id() if instruction_id: log_entry.instruction_id = instruction_id tracker = statesampler.get_current_tracker() if tracker: current_state = tracker.current_state() if (current_state and current_state.name_context and current_state.name_context.transform_id): log_entry.transform_id = current_state.name_context.transform_id try: self._log_entry_queue.put(log_entry, block=False) except queue.Full: self._dropped_logs += 1 def close(self) -> None: """Flush out all existing log entries and unregister this handler.""" try: self._alive = False # Acquiring the handler lock ensures ``emit`` is not run until the lock is # released. self.acquire() self._log_entry_queue.put(self._FINISHED, timeout=5) # wait on server to close. self._reader.join() self.release() # Unregister this handler. super().close() except Exception: # Log rather than raising exceptions, to avoid clobbering # underlying errors that may have caused this to close # prematurely. logging.error("Error closing the logging channel.", exc_info=True) def _write_log_entries(self) -> Iterator[beam_fn_api_pb2.LogEntry.List]: done = False while not done: log_entries = [self._log_entry_queue.get()] try: for _ in range(self._MAX_BATCH_SIZE): log_entries.append(self._log_entry_queue.get_nowait()) except queue.Empty: pass if log_entries[-1] is self._FINISHED: done = True log_entries.pop() if log_entries: # typing: log_entries was initialized as List[Union[..., Sentinel]], # but now that we've popped the sentinel out (above) we can safely cast yield beam_fn_api_pb2.LogEntry.List( log_entries=cast(list[beam_fn_api_pb2.LogEntry], log_entries)) def _read_log_control_messages(self) -> None: # Only reconnect when we are alive. # We can drop some logs in the unlikely event of logging connection # dropped(not closed) during termination when we still have logs to be sent. # This case is unlikely and the chance of reconnection and successful # transmission of logs is also very less as the process is terminating. # I choose not to handle this case to avoid un-necessary code complexity. alive = True # Force at least one connection attempt. while alive: # Loop for reconnection. log_control_iterator = self.connect() if self._dropped_logs > 0: logging.warning( "Dropped %d logs while logging client disconnected", self._dropped_logs) self._dropped_logs = 0 try: for _ in log_control_iterator: # Loop for consuming messages from server. # TODO(vikasrk): Handle control messages. pass # iterator is closed return except Exception as ex: print( "Logging client failed: {}... resetting".format(ex), file=sys.stderr) # Wait a bit before trying a reconnect time.sleep(0.5) # 0.5 seconds alive = self._alive ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/log_handler_test.py0000644000175100017510000002642215206113032024557 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import re import unittest import grpc import apache_beam as beam from apache_beam.coders.coders import FastPrimitivesCoder from apache_beam.portability import common_urns from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.portability.api import endpoints_pb2 from apache_beam.runners import common from apache_beam.runners.common import NameContext from apache_beam.runners.worker import bundle_processor from apache_beam.runners.worker import log_handler from apache_beam.runners.worker import operations from apache_beam.runners.worker import statesampler from apache_beam.runners.worker.bundle_processor import BeamTransformFactory from apache_beam.runners.worker.bundle_processor import BundleProcessor from apache_beam.transforms.window import GlobalWindow from apache_beam.utils import thread_pool_executor from apache_beam.utils.windowed_value import WindowedValue _LOGGER = logging.getLogger(__name__) @BeamTransformFactory.register_urn('beam:internal:testexn:v1', bytes) def create_exception_dofn( factory, transform_id, transform_proto, payload, consumers): """Returns a test DoFn that raises the given exception.""" class RaiseException(beam.DoFn): def __init__(self, msg): self.msg = msg.decode() def process(self, _): raise RuntimeError(self.msg) return bundle_processor._create_simple_pardo_operation( factory, transform_id, transform_proto, consumers, RaiseException(payload)) class TestOperation(operations.Operation): """Test operation that forwards its payload to consumers.""" class Spec: def __init__(self, transform_proto): self.output_coders = [ FastPrimitivesCoder() for _ in transform_proto.outputs ] def __init__( self, transform_proto, name_context, counter_factory, state_sampler, consumers, payload, ): super().__init__( name_context, self.Spec(transform_proto), counter_factory, state_sampler) self.payload = payload for _, consumer_ops in consumers.items(): for consumer in consumer_ops: self.add_receiver(consumer, 0) def start(self): super().start() # Not using windowing logic, so just using simple defaults here. if self.payload: self.process( WindowedValue(self.payload, timestamp=0, windows=[GlobalWindow()])) def process(self, windowed_value): self.output(windowed_value) @BeamTransformFactory.register_urn('beam:internal:testop:v1', bytes) def create_test_op(factory, transform_id, transform_proto, payload, consumers): return TestOperation( transform_proto, common.NameContext(transform_proto.unique_name, transform_id), factory.counter_factory, factory.state_sampler, consumers, payload) class BeamFnLoggingServicer(beam_fn_api_pb2_grpc.BeamFnLoggingServicer): def __init__(self): self.log_records_received = [] def Logging(self, request_iterator, context): for log_record in request_iterator: self.log_records_received.append(log_record) yield beam_fn_api_pb2.LogControl() class FnApiLogRecordHandlerTest(unittest.TestCase): def setUp(self): self.test_logging_service = BeamFnLoggingServicer() self.server = grpc.server(thread_pool_executor.shared_unbounded_instance()) beam_fn_api_pb2_grpc.add_BeamFnLoggingServicer_to_server( self.test_logging_service, self.server) self.test_port = self.server.add_insecure_port('[::]:0') self.server.start() self.logging_service_descriptor = endpoints_pb2.ApiServiceDescriptor() self.logging_service_descriptor.url = 'localhost:%s' % self.test_port self.fn_log_handler = log_handler.FnApiLogRecordHandler( self.logging_service_descriptor) logging.getLogger().setLevel(logging.INFO) logging.getLogger().addHandler(self.fn_log_handler) def tearDown(self): # wait upto 5 seconds. self.server.stop(5) def _verify_fn_log_handler(self, num_log_entries): msg = 'Testing fn logging' _LOGGER.debug('Debug Message 1') for idx in range(num_log_entries): _LOGGER.info('%s: %s', msg, idx) _LOGGER.debug('Debug Message 2') # Wait for logs to be sent to server. self.fn_log_handler.close() num_received_log_entries = 0 for outer in self.test_logging_service.log_records_received: for log_entry in outer.log_entries: self.assertEqual( beam_fn_api_pb2.LogEntry.Severity.INFO, log_entry.severity) self.assertEqual( '%s: %s' % (msg, num_received_log_entries), log_entry.message) self.assertTrue( re.match(r'.*log_handler_test.py:\d+', log_entry.log_location), log_entry.log_location) self.assertGreater(log_entry.timestamp.seconds, 0) self.assertGreaterEqual(log_entry.timestamp.nanos, 0) num_received_log_entries += 1 self.assertEqual(num_received_log_entries, num_log_entries) def assertContains(self, haystack, needle): self.assertTrue( needle in haystack, 'Expected %r to contain %r.' % (haystack, needle)) def test_exc_info(self): try: raise ValueError('some message') except ValueError: _LOGGER.error('some error', exc_info=True) self.fn_log_handler.close() log_entry = self.test_logging_service.log_records_received[0].log_entries[0] self.assertContains(log_entry.message, 'some error') self.assertContains(log_entry.trace, 'some message') self.assertContains(log_entry.trace, 'log_handler_test.py') def test_format_bad_message(self): # We specifically emit to the handler directly since we don't want to emit # to all handlers in general since we know that this record will raise an # exception during formatting. self.fn_log_handler.emit( logging.LogRecord( 'name', logging.ERROR, 'pathname', 777, 'TestLog %d', (None, ), exc_info=None)) self.fn_log_handler.close() log_entry = self.test_logging_service.log_records_received[0].log_entries[0] self.assertContains( log_entry.message, "Failed to format 'TestLog %d' with args '(None,)' during logging.") def test_context(self): try: with statesampler.instruction_id('A'): tracker = statesampler.for_test() with tracker.scoped_state(NameContext('name', 'tid'), 'stage'): _LOGGER.info('message a') with statesampler.instruction_id('B'): _LOGGER.info('message b') _LOGGER.info('message c') self.fn_log_handler.close() a, b, c = sum( [list(logs.log_entries) for logs in self.test_logging_service.log_records_received], []) self.assertEqual(a.instruction_id, 'A') self.assertEqual(b.instruction_id, 'B') self.assertEqual(c.instruction_id, '') self.assertEqual(a.transform_id, 'tid') self.assertEqual(b.transform_id, '') self.assertEqual(c.transform_id, '') finally: statesampler.set_current_tracker(None) def test_extracts_transform_id_during_exceptions(self): """Tests that transform ids are captured during user code exceptions.""" descriptor = beam_fn_api_pb2.ProcessBundleDescriptor() # Boiler plate for the DoFn. WINDOWING_ID = 'window' WINDOW_CODER_ID = 'cw' window = descriptor.windowing_strategies[WINDOWING_ID] window.window_fn.urn = common_urns.global_windows.urn window.window_coder_id = WINDOW_CODER_ID window.trigger.default.SetInParent() window_coder = descriptor.coders[WINDOW_CODER_ID] window_coder.spec.urn = common_urns.StandardCoders.Enum.GLOBAL_WINDOW.urn # Input collection to the exception raising DoFn. INPUT_PCOLLECTION_ID = 'pc-in' INPUT_CODER_ID = 'c-in' descriptor.pcollections[ INPUT_PCOLLECTION_ID].unique_name = INPUT_PCOLLECTION_ID descriptor.pcollections[INPUT_PCOLLECTION_ID].coder_id = INPUT_CODER_ID descriptor.pcollections[ INPUT_PCOLLECTION_ID].windowing_strategy_id = WINDOWING_ID descriptor.coders[ INPUT_CODER_ID].spec.urn = common_urns.StandardCoders.Enum.BYTES.urn # Output collection to the exception raising DoFn. OUTPUT_PCOLLECTION_ID = 'pc-out' OUTPUT_CODER_ID = 'c-out' descriptor.pcollections[ OUTPUT_PCOLLECTION_ID].unique_name = OUTPUT_PCOLLECTION_ID descriptor.pcollections[OUTPUT_PCOLLECTION_ID].coder_id = OUTPUT_CODER_ID descriptor.pcollections[ OUTPUT_PCOLLECTION_ID].windowing_strategy_id = WINDOWING_ID descriptor.coders[ OUTPUT_CODER_ID].spec.urn = common_urns.StandardCoders.Enum.BYTES.urn # Add a simple transform to inject an element into the fake pipeline. TEST_OP_TRANSFORM_ID = 'test_op' test_transform = descriptor.transforms[TEST_OP_TRANSFORM_ID] test_transform.outputs['None'] = INPUT_PCOLLECTION_ID test_transform.spec.urn = 'beam:internal:testop:v1' test_transform.spec.payload = b'hello, world!' # Add the DoFn to create an exception. TEST_EXCEPTION_TRANSFORM_ID = 'test_transform' test_transform = descriptor.transforms[TEST_EXCEPTION_TRANSFORM_ID] test_transform.inputs['0'] = INPUT_PCOLLECTION_ID test_transform.outputs['None'] = OUTPUT_PCOLLECTION_ID test_transform.spec.urn = 'beam:internal:testexn:v1' test_transform.spec.payload = b'expected exception' # Create and process a fake bundle. The instruction id doesn't matter # here. processor = BundleProcessor(set(), descriptor, None, None) with self.assertRaisesRegex(RuntimeError, 'expected exception'): processor.process_bundle('instruction_id') self.fn_log_handler.close() logs = [ log for logs in self.test_logging_service.log_records_received for log in logs.log_entries ] actual_log = logs[0] self.assertEqual( actual_log.severity, beam_fn_api_pb2.LogEntry.Severity.ERROR) self.assertTrue('expected exception' in actual_log.message) self.assertEqual(actual_log.transform_id, 'test_transform') # Test cases. data = { 'one_batch': log_handler.FnApiLogRecordHandler._MAX_BATCH_SIZE - 47, 'exact_multiple': log_handler.FnApiLogRecordHandler._MAX_BATCH_SIZE, 'multi_batch': log_handler.FnApiLogRecordHandler._MAX_BATCH_SIZE * 3 + 47 } def _create_test(name, num_logs): setattr( FnApiLogRecordHandlerTest, 'test_%s' % name, lambda self: self._verify_fn_log_handler(num_logs)) for test_name, num_logs_entries in data.items(): _create_test(test_name, num_logs_entries) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/logger.py0000644000175100017510000001515115206113032022516 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Python worker logging.""" # pytype: skip-file # mypy: disallow-untyped-defs # ruff: noqa: UP006 import contextlib import json import logging import threading import traceback from typing import Any from typing import Dict from typing import Iterator from typing import List from apache_beam.runners.worker import statesampler # Per-thread worker information. This is used only for logging to set # context information that changes while work items get executed: # work_item_id, step_name, stage_name. class _PerThreadWorkerData(threading.local): def __init__(self) -> None: super().__init__() # in the list, as going up and down all the way to zero incurs several # reallocations. self.stack: List[Dict[str, Any]] = [] def get_data(self) -> Dict[str, Any]: all_data = {} for datum in self.stack: all_data.update(datum) return all_data per_thread_worker_data = _PerThreadWorkerData() @contextlib.contextmanager def PerThreadLoggingContext(**kwargs: Any) -> Iterator[None]: """A context manager to add per thread attributes.""" stack = per_thread_worker_data.stack stack.append(kwargs) try: yield finally: stack.pop() class JsonLogFormatter(logging.Formatter): """A JSON formatter class as expected by the logging standard module.""" def __init__(self, job_id: str, worker_id: str) -> None: super().__init__() self.job_id = job_id self.worker_id = worker_id def format(self, record: logging.LogRecord) -> str: """Returns a JSON string based on a LogRecord instance. Args: record: A LogRecord instance. See below for details. Returns: A JSON string representing the record. A LogRecord instance has the following attributes and is used for formatting the final message. Attributes: created: A double representing the timestamp for record creation (e.g., 1438365207.624597). Note that the number contains also msecs and microsecs information. Part of this is also available in the 'msecs' attribute. msecs: A double representing the msecs part of the record creation (e.g., 624.5970726013184). msg: Logging message containing formatting instructions or an arbitrary object. This is the first argument of a log call. args: A tuple containing the positional arguments for the logging call. levelname: A string. Possible values are: INFO, WARNING, ERROR, etc. exc_info: None or a 3-tuple with exception information as it is returned by a call to sys.exc_info(). name: Logger's name. Most logging is done using the default root logger and therefore the name will be 'root'. filename: Basename of the file where logging occurred. funcName: Name of the function where logging occurred. process: The PID of the process running the worker. thread: An id for the thread where the record was logged. This is not a real TID (the one provided by OS) but rather the id (address) of a Python thread object. Nevertheless having this value can allow to filter log statement from only one specific thread. """ output: Dict[str, Any] = {} output['timestamp'] = { 'seconds': int(record.created), 'nanos': int(record.msecs * 1000000) } # ERROR. INFO, DEBUG log levels translate into the same for severity # property. WARNING becomes WARN. output['severity'] = ( record.levelname if record.levelname != 'WARNING' else 'WARN') # msg could be an arbitrary object, convert it to a string first. record_msg = str(record.msg) # Prepare the actual message using the message formatting string and the # positional arguments as they have been used in the log call. if record.args: try: output['message'] = record_msg % record.args except (TypeError, ValueError): output['message'] = '%s with args (%s)' % (record_msg, record.args) else: output['message'] = record_msg # The thread ID is logged as a combination of the process ID and thread ID # since workers can run in multiple processes. output['thread'] = '%s:%s' % (record.process, record.thread) # job ID and worker ID. These do not change during the lifetime of a worker. output['job'] = self.job_id output['worker'] = self.worker_id # Stage, step and work item ID come from thread local storage since they # change with every new work item leased for execution. If there is no # work item ID then we make sure the step is undefined too. data = per_thread_worker_data.get_data() if 'work_item_id' in data: output['work'] = data['work_item_id'] tracker = statesampler.get_current_tracker() if tracker: output['stage'] = tracker.stage_name if tracker.current_state() and tracker.current_state().name_context: output['step'] = tracker.current_state().name_context.logging_name() # All logging happens using the root logger. We will add the basename of the # file and the function name where the logging happened to make it easier # to identify who generated the record. output['logger'] = '%s:%s:%s' % ( record.name, record.filename, record.funcName) # Add exception information if any is available. if record.exc_info: output['exception'] = ''.join( traceback.format_exception(*record.exc_info)) return json.dumps(output) def initialize( job_id: str, worker_id: str, log_path: str, log_level: int = logging.INFO) -> None: """Initialize root logger so that we log JSON to a file and text to stdout.""" file_handler = logging.FileHandler(log_path) file_handler.setFormatter(JsonLogFormatter(job_id, worker_id)) logging.getLogger().addHandler(file_handler) # Default level is set to INFO to avoid logging various DEBUG level log calls # sprinkled throughout the code. logging.getLogger().setLevel(log_level) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/logger_test.py0000644000175100017510000001763315206113032023564 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for worker logging utilities.""" # pytype: skip-file import json import logging import sys import threading import unittest from apache_beam.runners.worker import logger from apache_beam.runners.worker import statesampler from apache_beam.utils.counters import CounterFactory class PerThreadLoggingContextTest(unittest.TestCase): def thread_check_attribute(self, name): self.assertFalse(name in logger.per_thread_worker_data.get_data()) with logger.PerThreadLoggingContext(**{name: 'thread-value'}): self.assertEqual( logger.per_thread_worker_data.get_data()[name], 'thread-value') self.assertFalse(name in logger.per_thread_worker_data.get_data()) def test_per_thread_attribute(self): self.assertFalse('xyz' in logger.per_thread_worker_data.get_data()) with logger.PerThreadLoggingContext(xyz='value'): self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value') thread = threading.Thread( target=self.thread_check_attribute, args=('xyz', )) thread.start() thread.join() self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value') self.assertFalse('xyz' in logger.per_thread_worker_data.get_data()) def test_set_when_undefined(self): self.assertFalse('xyz' in logger.per_thread_worker_data.get_data()) with logger.PerThreadLoggingContext(xyz='value'): self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value') self.assertFalse('xyz' in logger.per_thread_worker_data.get_data()) def test_set_when_already_defined(self): self.assertFalse('xyz' in logger.per_thread_worker_data.get_data()) with logger.PerThreadLoggingContext(xyz='value'): self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value') with logger.PerThreadLoggingContext(xyz='value2'): self.assertEqual( logger.per_thread_worker_data.get_data()['xyz'], 'value2') self.assertEqual(logger.per_thread_worker_data.get_data()['xyz'], 'value') self.assertFalse('xyz' in logger.per_thread_worker_data.get_data()) class JsonLogFormatterTest(unittest.TestCase): SAMPLE_RECORD = { 'created': 123456.789, 'msecs': 789.654321, 'msg': '%s:%d:%.2f', 'args': ('xyz', 4, 3.14), 'levelname': 'WARNING', 'process': 'pid', 'thread': 'tid', 'name': 'name', 'filename': 'file', 'funcName': 'func', 'exc_info': None } SAMPLE_OUTPUT = { 'timestamp': { 'seconds': 123456, 'nanos': 789654321 }, 'severity': 'WARN', 'message': 'xyz:4:3.14', 'thread': 'pid:tid', 'job': 'jobid', 'worker': 'workerid', 'logger': 'name:file:func' } def create_log_record(self, **kwargs): class Record(object): def __init__(self, **kwargs): for k, v in kwargs.items(): setattr(self, k, v) return Record(**kwargs) def test_basic_record(self): formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid') record = self.create_log_record(**self.SAMPLE_RECORD) self.assertEqual(json.loads(formatter.format(record)), self.SAMPLE_OUTPUT) def execute_multiple_cases(self, test_cases): record = self.SAMPLE_RECORD output = self.SAMPLE_OUTPUT formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid') for case in test_cases: record['msg'] = case['msg'] record['args'] = case['args'] output['message'] = case['expected'] self.assertEqual( json.loads(formatter.format(self.create_log_record(**record))), output) def test_record_with_format_character(self): test_cases = [ { 'msg': '%A', 'args': (), 'expected': '%A' }, { 'msg': '%s', 'args': (), 'expected': '%s' }, { 'msg': '%A%s', 'args': ('xy'), 'expected': '%A%s with args (xy)' }, { 'msg': '%s%s', 'args': (1), 'expected': '%s%s with args (1)' }, ] self.execute_multiple_cases(test_cases) def test_record_with_arbitrary_messages(self): test_cases = [ { 'msg': ImportError('abc'), 'args': (), 'expected': 'abc' }, { 'msg': TypeError('abc %s'), 'args': ('def'), 'expected': 'abc def' }, ] self.execute_multiple_cases(test_cases) def test_record_with_per_thread_info(self): self.maxDiff = None tracker = statesampler.StateSampler('stage', CounterFactory()) statesampler.set_current_tracker(tracker) formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid') with logger.PerThreadLoggingContext(work_item_id='workitem'): with tracker.scoped_state('step', 'process'): record = self.create_log_record(**self.SAMPLE_RECORD) log_output = json.loads(formatter.format(record)) expected_output = dict(self.SAMPLE_OUTPUT) expected_output.update({ 'work': 'workitem', 'stage': 'stage', 'step': 'step' }) self.assertEqual(log_output, expected_output) statesampler.set_current_tracker(None) def test_nested_with_per_thread_info(self): self.maxDiff = None tracker = statesampler.StateSampler('stage', CounterFactory()) statesampler.set_current_tracker(tracker) formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid') with logger.PerThreadLoggingContext(work_item_id='workitem'): with tracker.scoped_state('step1', 'process'): record = self.create_log_record(**self.SAMPLE_RECORD) log_output1 = json.loads(formatter.format(record)) with tracker.scoped_state('step2', 'process'): record = self.create_log_record(**self.SAMPLE_RECORD) log_output2 = json.loads(formatter.format(record)) record = self.create_log_record(**self.SAMPLE_RECORD) log_output3 = json.loads(formatter.format(record)) statesampler.set_current_tracker(None) record = self.create_log_record(**self.SAMPLE_RECORD) log_output4 = json.loads(formatter.format(record)) self.assertEqual( log_output1, dict(self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step1')) self.assertEqual( log_output2, dict(self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step2')) self.assertEqual( log_output3, dict(self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step1')) self.assertEqual(log_output4, self.SAMPLE_OUTPUT) def test_exception_record(self): formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid') try: raise ValueError('Something') except ValueError: attribs = dict(self.SAMPLE_RECORD) attribs.update({'exc_info': sys.exc_info()}) record = self.create_log_record(**attribs) log_output = json.loads(formatter.format(record)) # Check if exception type, its message, and stack trace information are in. exn_output = log_output.pop('exception') self.assertNotEqual(exn_output.find('ValueError: Something'), -1) self.assertNotEqual(exn_output.find('logger_test.py'), -1) self.assertEqual(log_output, self.SAMPLE_OUTPUT) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/opcounters.pxd0000644000175100017510000000464315206113032023607 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # cimport cython cimport libc.stdint from apache_beam.utils.counters cimport Counter from apache_beam.runners.worker cimport statesampler_fast cdef class TransformIOCounter(object): cdef readonly object _counter_factory cdef readonly object _state_sampler cdef Counter bytes_read_counter cdef statesampler_fast.ScopedState scoped_state cdef object _latest_step cpdef update_current_step(self) cpdef add_bytes_read(self, libc.stdint.int64_t n) cpdef __enter__(self) cpdef __exit__(self, exc_type, exc_value, traceback) cdef class NoOpTransformIOCounter(TransformIOCounter): pass cdef class SideInputReadCounter(TransformIOCounter): cdef readonly object declaring_step cdef readonly object input_index cdef class SumAccumulator(object): cdef libc.stdint.int64_t _value cpdef update(self, libc.stdint.int64_t value) cpdef libc.stdint.int64_t value(self) cdef class OperationCounters(object): cdef public _counter_factory cdef public Counter element_counter cdef public Counter mean_byte_counter cdef public coder_impl cdef public SumAccumulator active_accumulator cdef public object current_size cdef public libc.stdint.int64_t _sample_counter cdef public libc.stdint.int64_t _next_sample cdef public object output_type_constraints cdef public object producer_batch_converter cpdef update_from(self, windowed_value) cpdef update_from_batch(self, windowed_batch) cdef inline do_sample(self, windowed_value) cpdef update_collect(self) cpdef type_check(self, value) cdef libc.stdint.int64_t _compute_next_sample(self, libc.stdint.int64_t i) cdef inline bint _should_sample(self) cpdef bint should_sample(self) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/opcounters.py0000644000175100017510000003221615206113032023441 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Counters collect the progress of the Worker for reporting to the service.""" # pytype: skip-file # ruff: noqa: UP006 import math import random import sys from typing import TYPE_CHECKING from typing import Any from typing import Optional from apache_beam.typehints import TypeCheckError from apache_beam.typehints.decorators import _check_instance_type from apache_beam.utils import counters from apache_beam.utils import windowed_value from apache_beam.utils.counters import Counter from apache_beam.utils.counters import CounterName if TYPE_CHECKING: from apache_beam.runners.worker.statesampler import StateSampler from apache_beam.typehints.batch import BatchConverter class TransformIOCounter(object): """Class to track time and bytes consumed while reading from IO. Subclasses should be able to track consumption of IO across steps in the same stage - for instance, if a Shuffle or Side Input iterable is passed down to a next step. Some examples of IO can be side inputs, shuffle, or streaming state. """ def __init__(self, counter_factory, state_sampler): """Create a new IO read counter. Args: counter_factory: A counters.CounterFactory to create byte counters. state_sampler: A statesampler.StateSampler to transition into read states. """ self._counter_factory = counter_factory self._state_sampler = state_sampler self._latest_step = None self.bytes_read_counter = None self.scoped_state = None def update_current_step(self): """Update the current running step. Due to the fusion optimization, user code may choose to emit the data structure that holds side inputs (Iterable, Dict, or others). This call updates the current step, to attribute the data consumption to the step that is responsible for actual consumption. CounterName uses the io_target field for information pertinent to the consumption of IO. """ current_state = self._state_sampler.current_state() current_step_name = current_state.name.step_name if current_step_name != self._latest_step: self._latest_step = current_step_name self._update_counters_for_requesting_step(current_step_name) def _update_counters_for_requesting_step(self, step_name): pass def add_bytes_read(self, count): if count > 0 and self.bytes_read_counter: self.bytes_read_counter.update(count) def __enter__(self): self.scoped_state.__enter__() def __exit__(self, exception_type, exception_value, traceback): self.scoped_state.__exit__(exception_type, exception_value, traceback) class NoOpTransformIOCounter(TransformIOCounter): """All operations for IO tracking are no-ops.""" def __init__(self): super().__init__(None, None) def update_current_step(self): pass def __enter__(self): pass def __exit__(self, exception_type, exception_value, traceback): pass def add_bytes_read(self, count): pass class SideInputReadCounter(TransformIOCounter): """Tracks time and bytes consumed while reading from side inputs. This class is designed to track consumption of side inputs across fused steps. We represent a side input as a declaring step, and an input index. The declaring step is the step that originally receives the side input for consumption, and the input index in which the declaring step receives the side input that we want to identify. Note that the declaring step originally receives the side input, but it may not be the only step that spends time reading from this side input. """ def __init__( self, counter_factory, state_sampler, # type: StateSampler declaring_step, input_index): """Create a side input read counter. Args: counter_factory: A counters.CounterFactory to create byte counters. state_sampler: A statesampler.StateSampler to transition into read states. declaring_step: A string with the step name of the step that directly receives the side input initially. input_index: The index of the side input in the list of inputs of the declaring step. The side input is uniquely identified by (declaring_step, input_index); where declaring_step is the step that receives the PCollectionView as a side input, and input_index is the index of the PCollectionView within the list of inputs. """ super().__init__(counter_factory, state_sampler) self.declaring_step = declaring_step self.input_index = input_index # Side inputs are set up within the start state of the first receiving # step. We check the current state to create the internal counters. self.update_current_step() def _update_counters_for_requesting_step(self, step_name): side_input_id = counters.side_input_id(step_name, self.input_index) self.scoped_state = self._state_sampler.scoped_state( self.declaring_step, 'read-sideinput', io_target=side_input_id) self.bytes_read_counter = self._counter_factory.get_counter( CounterName( 'read-sideinput-byte-count', step_name=self.declaring_step, io_target=side_input_id), Counter.SUM) class SumAccumulator(object): """Accumulator for collecting byte counts.""" def __init__(self): self._value = 0 def update(self, value): self._value += value def value(self): return self._value class OperationCounters(object): """The set of basic counters to attach to an Operation.""" def __init__( self, counter_factory, step_name, # type: str coder, index, suffix='out', producer_type_hints=None, producer_batch_converter=None, # type: Optional[BatchConverter] ): self._counter_factory = counter_factory self.element_counter = counter_factory.get_counter( '%s-%s%s-ElementCount' % (step_name, suffix, index), Counter.SUM) self.mean_byte_counter = counter_factory.get_counter( '%s-%s%s-MeanByteCount' % (step_name, suffix, index), Counter.BEAM_DISTRIBUTION) self.coder_impl = coder.get_impl() if coder else None self.active_accumulator = None # type: Optional[SumAccumulator] self.current_size = None # type: Optional[int] self._sample_counter = 0 self._next_sample = 0 self.output_type_constraints = producer_type_hints or {} self.producer_batch_converter = producer_batch_converter def update_from(self, windowed_value): # type: (windowed_value.WindowedValue) -> None """Add one value to this counter.""" if self._should_sample(): self.do_sample(windowed_value) def update_from_batch(self, windowed_batch): # type: (windowed_value.WindowedBatch) -> None assert self.producer_batch_converter is not None assert isinstance(windowed_batch, windowed_value.HomogeneousWindowedBatch) batch_length = self.producer_batch_converter.get_length( windowed_batch.values) self.element_counter.update(batch_length) mean_element_size = self.producer_batch_converter.estimate_byte_size( windowed_batch.values) / batch_length self.mean_byte_counter.update_n(mean_element_size, batch_length) def _observable_callback(self, inner_coder_impl, accumulator): def _observable_callback_inner(value, is_encoded=False): # TODO(ccy): If this stream is large, sample it as well. # To do this, we'll need to compute the average size of elements # in this stream to add the *total* size of this stream to accumulator. # We'll also want make sure we sample at least some of this stream # (as self.should_sample() may be sampling very sparsely by now). if is_encoded: size = len(value) accumulator.update(size) else: accumulator.update(inner_coder_impl.estimate_size(value)) return _observable_callback_inner def type_check(self, value): # type: (Any, bool) -> None for transform_label, type_constraint_tuple in ( self.output_type_constraints.items()): parameter_name, constraint = type_constraint_tuple try: _check_instance_type(constraint, value, parameter_name, verbose=True) except TypeCheckError as e: # TODO: Remove the 'ParDo' prefix for the label name (BEAM-10710) if not transform_label.startswith('ParDo'): transform_label = 'ParDo(%s)' % transform_label error_msg = ( 'Runtime type violation detected within %s: ' '%s' % (transform_label, e)) _, _, traceback = sys.exc_info() raise TypeCheckError(error_msg).with_traceback(traceback) def do_sample(self, windowed_value): # type: (windowed_value.WindowedValue) -> None self.type_check(windowed_value.value) size, observables = ( self.coder_impl.get_estimated_size_and_observables(windowed_value)) # type: ignore[union-attr] if not observables: self.current_size = size else: self.active_accumulator = SumAccumulator() self.active_accumulator.update(size) for observable, inner_coder_impl in observables: observable.register_observer( self._observable_callback( inner_coder_impl, self.active_accumulator)) def update_collect(self): """Collects the accumulated size estimates. Now that the element has been processed, we ask our accumulator for the total and store the result in a counter. """ self.element_counter.update(1) if self.current_size is not None: self.mean_byte_counter.update(self.current_size) self.current_size = None elif self.active_accumulator is not None: self.mean_byte_counter.update(self.active_accumulator.value()) self.active_accumulator = None def _compute_next_sample(self, i): # https://en.wikipedia.org/wiki/Reservoir_sampling#Fast_Approximation gap = math.log(1.0 - random.random()) / math.log(1.0 - (10.0 / i)) return i + math.floor(gap) def _should_sample(self): """Determines whether to sample the next element. Size calculation can be expensive, so we don't do it for each element. Because we need only an estimate of average size, we sample. We always sample the first 10 elements, then the sampling rate is approximately 10/N. After reading N elements, of the next N, we will sample approximately 10*ln(2) (about 7) elements. This algorithm samples at the same rate as Reservoir Sampling, but it never throws away early results. (Because we keep only a running accumulation, storage is not a problem, so there is no need to discard earlier calculations.) Because we accumulate and do not replace, our statistics are biased toward early data. If the data are distributed uniformly, this is not a problem. If the data change over time (i.e., the element size tends to grow or shrink over time), our estimate will show the bias. We could correct this by giving weight N to each sample, since each sample is a stand-in for the N/(10*ln(2)) samples around it, which is proportional to N. Since we do not expect biased data, for efficiency we omit the extra multiplication. We could reduce the early-data bias by putting a lower bound on the sampling rate. Computing random.randint(1, self._sample_counter) for each element is too slow, so when the sample size is big enough (we estimate 30 is big enough), we estimate the size of the gap after each sample. This estimation allows us to call random much less often. Returns: True if it is time to compute another element's size. """ if self.coder_impl is None: return False self._sample_counter += 1 if self._next_sample == 0: if random.randint(1, self._sample_counter) <= 10: if self._sample_counter > 30: self._next_sample = self._compute_next_sample(self._sample_counter) return True return False elif self._sample_counter >= self._next_sample: self._next_sample = self._compute_next_sample(self._sample_counter) return True return False def should_sample(self): # We create this separate method because the above "_should_sample()" method # is marked as inline in Cython and thus can't be exposed to Python code. return self._should_sample() def restart_sampling(self): self._sample_counter = 0 def __str__(self): return '<%s [%s]>' % ( self.__class__.__name__, ', '.join([str(x) for x in self.__iter__()])) def __repr__(self): return '<%s %s at %s>' % ( self.__class__.__name__, [x for x in self.__iter__()], hex(id(self))) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/opcounters_test.py0000644000175100017510000001774415206113032024511 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import math import random import unittest from apache_beam import coders from apache_beam import typehints from apache_beam.runners.worker import opcounters from apache_beam.runners.worker import statesampler from apache_beam.runners.worker.opcounters import OperationCounters from apache_beam.transforms.window import GlobalWindows from apache_beam.utils import counters from apache_beam.utils.counters import CounterFactory # Classes to test that we can handle a variety of objects. # These have to be at top level so the pickler can find them. class OldClassThatDoesNotImplementLen(object): def __init__(self): pass class ObjectThatDoesNotImplementLen(object): def __init__(self): pass class TransformIoCounterTest(unittest.TestCase): def test_basic_counters(self): counter_factory = CounterFactory() sampler = statesampler.StateSampler('stage1', counter_factory) sampler.start() with sampler.scoped_state('step1', 'stateA'): counter = opcounters.SideInputReadCounter( counter_factory, sampler, declaring_step='step1', input_index=1) with sampler.scoped_state('step2', 'stateB'): with counter: counter.add_bytes_read(10) counter.update_current_step() sampler.stop() sampler.commit_counters() actual_counter_names = {c.name for c in counter_factory.get_counters()} expected_counter_names = set([ # Counter names for STEP 1 counters.CounterName( 'read-sideinput-msecs', stage_name='stage1', step_name='step1', io_target=counters.side_input_id('step1', 1)), counters.CounterName( 'read-sideinput-byte-count', step_name='step1', io_target=counters.side_input_id('step1', 1)), # Counter names for STEP 2 counters.CounterName( 'read-sideinput-msecs', stage_name='stage1', step_name='step1', io_target=counters.side_input_id('step2', 1)), counters.CounterName( 'read-sideinput-byte-count', step_name='step1', io_target=counters.side_input_id('step2', 1)), ]) self.assertTrue(actual_counter_names.issuperset(expected_counter_names)) class OperationCountersTest(unittest.TestCase): def verify_counters(self, opcounts, expected_elements, expected_size=None): self.assertEqual(expected_elements, opcounts.element_counter.value()) if expected_size is not None: if math.isnan(expected_size): self.assertTrue(math.isnan(opcounts.mean_byte_counter.value()[0])) else: self.assertEqual(expected_size, opcounts.mean_byte_counter.value()[0]) def test_update_int(self): opcounts = OperationCounters( CounterFactory(), 'some-name', coders.PickleCoder(), 0) self.verify_counters(opcounts, 0) opcounts.update_from(GlobalWindows.windowed_value(1)) opcounts.update_collect() self.verify_counters(opcounts, 1) def test_update_str(self): coder = coders.PickleCoder() opcounts = OperationCounters(CounterFactory(), 'some-name', coder, 0) self.verify_counters(opcounts, 0, float('nan')) value = GlobalWindows.windowed_value('abcde') opcounts.update_from(value) opcounts.update_collect() estimated_size = coder.estimate_size(value) self.verify_counters(opcounts, 1, estimated_size) def test_update_old_object(self): coder = coders.PickleCoder() opcounts = OperationCounters(CounterFactory(), 'some-name', coder, 0) self.verify_counters(opcounts, 0, float('nan')) obj = OldClassThatDoesNotImplementLen() value = GlobalWindows.windowed_value(obj) opcounts.update_from(value) opcounts.update_collect() estimated_size = coder.estimate_size(value) self.verify_counters(opcounts, 1, estimated_size) def test_update_new_object(self): coder = coders.PickleCoder() opcounts = OperationCounters(CounterFactory(), 'some-name', coder, 0) self.verify_counters(opcounts, 0, float('nan')) obj = ObjectThatDoesNotImplementLen() value = GlobalWindows.windowed_value(obj) opcounts.update_from(value) opcounts.update_collect() estimated_size = coder.estimate_size(value) self.verify_counters(opcounts, 1, estimated_size) def test_update_multiple(self): coder = coders.PickleCoder() total_size = 0 opcounts = OperationCounters(CounterFactory(), 'some-name', coder, 0) self.verify_counters(opcounts, 0, float('nan')) value = GlobalWindows.windowed_value('abcde') opcounts.update_from(value) opcounts.update_collect() total_size += coder.estimate_size(value) value = GlobalWindows.windowed_value('defghij') opcounts.update_from(value) opcounts.update_collect() total_size += coder.estimate_size(value) self.verify_counters(opcounts, 2, (float(total_size) / 2)) value = GlobalWindows.windowed_value('klmnop') opcounts.update_from(value) opcounts.update_collect() total_size += coder.estimate_size(value) self.verify_counters(opcounts, 3, (float(total_size) / 3)) def test_update_batch(self): coder = coders.FastPrimitivesCoder() opcounts = OperationCounters( CounterFactory(), 'some-name', coder, 0, producer_batch_converter=typehints.batch.BatchConverter.from_typehints( element_type=typehints.Any, batch_type=typehints.List[typehints.Any])) size_per_element = coder.estimate_size(50) self.verify_counters(opcounts, 0, float('nan')) opcounts.update_from_batch(GlobalWindows.windowed_batch(list(range(100)))) self.verify_counters(opcounts, 100, size_per_element) opcounts.update_from_batch( GlobalWindows.windowed_batch(list(range(100, 200)))) self.verify_counters(opcounts, 200, size_per_element) def test_should_sample(self): # Order of magnitude more buckets than highest constant in code under test. buckets = [0] * 300 # The seed is arbitrary and exists just to ensure this test is robust. # If you don't like this seed, try your own; the test should still pass. random.seed(1720) # Do enough runs that the expected hits even in the last buckets # is big enough to expect some statistical smoothing. total_runs = 10 * len(buckets) # Fill the buckets. for _ in range(total_runs): opcounts = OperationCounters( CounterFactory(), 'some-name', coders.PickleCoder(), 0) for i in range(len(buckets)): if opcounts.should_sample(): buckets[i] += 1 # Look at the buckets to see if they are likely. for i in range(10): self.assertEqual(total_runs, buckets[i]) for i in range(10, len(buckets)): self.assertTrue( buckets[i] > 7 * total_runs / i, 'i=%d, buckets[i]=%d, expected=%d, ratio=%f' % ( i, buckets[i], 10 * total_runs / i, buckets[i] / (10.0 * total_runs / i))) self.assertTrue( buckets[i] < 14 * total_runs / i, 'i=%d, buckets[i]=%d, expected=%d, ratio=%f' % ( i, buckets[i], 10 * total_runs / i, buckets[i] / (10.0 * total_runs / i))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/operation_specs.py0000644000175100017510000003461315206113032024440 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Worker utilities for representing MapTasks. Each MapTask represents a sequence of ParallelInstruction(s): read from a source, write to a sink, parallel do, etc. """ # pytype: skip-file import collections from apache_beam import coders def build_worker_instruction(*args): """Create an object representing a ParallelInstruction protobuf. This will be a collections.namedtuple with a custom __str__ method. Alas, this wrapper is not known to pylint, which thinks it creates constants. You may have to put a disable=invalid-name pylint annotation on any use of this, depending on your names. Args: *args: first argument is the name of the type to create. Should start with "Worker". Second arguments is alist of the attributes of this object. Returns: A new class, a subclass of tuple, that represents the protobuf. """ tuple_class = collections.namedtuple(*args) tuple_class.__str__ = worker_object_to_string tuple_class.__repr__ = worker_object_to_string return tuple_class def worker_printable_fields(workerproto): """Returns the interesting fields of a Worker* object.""" return [ '%s=%s' % (name, value) # _asdict is the only way and cannot subclass this generated class # pylint: disable=protected-access for name, value in workerproto._asdict().items() # want to output value 0 but not None nor [] if (value or value == 0) and name not in ( 'coder', 'coders', 'output_coders', 'elements', 'combine_fn', 'serialized_fn', 'window_fn', 'append_trailing_newlines', 'strip_trailing_newlines', 'compression_type', 'context', 'start_shuffle_position', 'end_shuffle_position', 'shuffle_reader_config', 'shuffle_writer_config') ] def worker_object_to_string(worker_object): """Returns a string compactly representing a Worker* object.""" return '%s(%s)' % ( worker_object.__class__.__name__, ', '.join(worker_printable_fields(worker_object))) # All the following Worker* definitions will have these lint problems: # pylint: disable=invalid-name # pylint: disable=pointless-string-statement WorkerRead = build_worker_instruction('WorkerRead', ['source', 'output_coders']) """Worker details needed to read from a source. Attributes: source: a source object. output_coders: 1-tuple of the coder for the output. """ WorkerSideInputSource = build_worker_instruction( 'WorkerSideInputSource', ['source', 'tag']) """Worker details needed to read from a side input source. Attributes: source: a source object. tag: string tag for this side input. """ WorkerGroupingShuffleRead = build_worker_instruction( 'WorkerGroupingShuffleRead', [ 'start_shuffle_position', 'end_shuffle_position', 'shuffle_reader_config', 'coder', 'output_coders' ]) """Worker details needed to read from a grouping shuffle source. Attributes: start_shuffle_position: An opaque string to be passed to the shuffle source to indicate where to start reading. end_shuffle_position: An opaque string to be passed to the shuffle source to indicate where to stop reading. shuffle_reader_config: An opaque string used to initialize the shuffle reader. Contains things like connection endpoints for the shuffle server appliance and various options. coder: The KV coder used to decode shuffle entries. output_coders: 1-tuple of the coder for the output. """ WorkerUngroupedShuffleRead = build_worker_instruction( 'WorkerUngroupedShuffleRead', [ 'start_shuffle_position', 'end_shuffle_position', 'shuffle_reader_config', 'coder', 'output_coders' ]) """Worker details needed to read from an ungrouped shuffle source. Attributes: start_shuffle_position: An opaque string to be passed to the shuffle source to indicate where to start reading. end_shuffle_position: An opaque string to be passed to the shuffle source to indicate where to stop reading. shuffle_reader_config: An opaque string used to initialize the shuffle reader. Contains things like connection endpoints for the shuffle server appliance and various options. coder: The value coder used to decode shuffle entries. """ WorkerWrite = build_worker_instruction( 'WorkerWrite', ['sink', 'input', 'output_coders']) """Worker details needed to write to a sink. Attributes: sink: a sink object. input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). output_coders: 1-tuple, coder to use to estimate bytes written. """ WorkerInMemoryWrite = build_worker_instruction( 'WorkerInMemoryWrite', ['output_buffer', 'write_windowed_values', 'input', 'output_coders']) """Worker details needed to write to a in-memory sink. Used only for unit testing. It makes worker tests less cluttered with code like "write to a file and then check file contents". Attributes: output_buffer: list to which output elements will be appended write_windowed_values: whether to record the entire WindowedValue outputs, or just the raw (unwindowed) value input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). output_coders: 1-tuple, coder to use to estimate bytes written. """ WorkerShuffleWrite = build_worker_instruction( 'WorkerShuffleWrite', ['shuffle_kind', 'shuffle_writer_config', 'input', 'output_coders']) """Worker details needed to write to a shuffle sink. Attributes: shuffle_kind: A string describing the shuffle kind. This can control the way the worker interacts with the shuffle sink. The possible values are: 'ungrouped', 'group_keys', and 'group_keys_and_sort_values'. shuffle_writer_config: An opaque string used to initialize the shuffle write. Contains things like connection endpoints for the shuffle server appliance and various options. input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). output_coders: 1-tuple of the coder for input elements. If the shuffle_kind is grouping, this is expected to be a KV coder. """ WorkerDoFn = build_worker_instruction( 'WorkerDoFn', ['serialized_fn', 'output_tags', 'input', 'side_inputs', 'output_coders']) """Worker details needed to run a DoFn. Attributes: serialized_fn: A serialized DoFn object to be run for each input element. output_tags: The string tags used to identify the outputs of a ParDo operation. The tag is present even if the ParDo has just one output (e.g., ['out']. output_coders: array of coders, one for each output. input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). side_inputs: A list of Worker...Read instances describing sources to be used for getting values. The types supported right now are WorkerInMemoryRead and WorkerTextRead. """ WorkerReifyTimestampAndWindows = build_worker_instruction( 'WorkerReifyTimestampAndWindows', ['output_tags', 'input', 'output_coders']) """Worker details needed to run a WindowInto. Attributes: output_tags: The string tags used to identify the outputs of a ParDo operation. The tag is present even if the ParDo has just one output (e.g., ['out']. output_coders: array of coders, one for each output. input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). """ WorkerMergeWindows = build_worker_instruction( 'WorkerMergeWindows', [ 'window_fn', 'combine_fn', 'phase', 'output_tags', 'input', 'coders', 'context', 'output_coders' ]) """Worker details needed to run a MergeWindows (aka. GroupAlsoByWindows). Attributes: window_fn: A serialized Windowing object representing the windowing strategy. combine_fn: A serialized CombineFn object to be used after executing the GroupAlsoByWindows operation. May be None if not a combining operation. phase: Possible values are 'all', 'add', 'merge', and 'extract'. A runner optimizer may split the user combiner in 3 separate phases (ADD, MERGE, and EXTRACT), on separate VMs, as it sees fit. The phase attribute dictates which DoFn is actually running in the worker. May be None if not a combining operation. output_tags: The string tags used to identify the outputs of a ParDo operation. The tag is present even if the ParDo has just one output (e.g., ['out']. output_coders: array of coders, one for each output. input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). coders: A 2-tuple of coders (key, value) to encode shuffle entries. context: The ExecutionContext object for the current work item. """ WorkerCombineFn = build_worker_instruction( 'WorkerCombineFn', ['serialized_fn', 'phase', 'input', 'output_coders']) """Worker details needed to run a CombineFn. Attributes: serialized_fn: A serialized CombineFn object to be used. phase: Possible values are 'all', 'add', 'merge', and 'extract'. A runner optimizer may split the user combiner in 3 separate phases (ADD, MERGE, and EXTRACT), on separate VMs, as it sees fit. The phase attribute dictates which DoFn is actually running in the worker. input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). output_coders: 1-tuple of the coder for the output. """ WorkerPartialGroupByKey = build_worker_instruction( 'WorkerPartialGroupByKey', ['combine_fn', 'input', 'output_coders']) """Worker details needed to run a partial group-by-key. Attributes: combine_fn: A serialized CombineFn object to be used. input: A (producer index, output index) tuple representing the ParallelInstruction operation whose output feeds into this operation. The output index is 0 except for multi-output operations (like ParDo). output_coders: 1-tuple of the coder for the output. """ WorkerFlatten = build_worker_instruction( 'WorkerFlatten', ['inputs', 'output_coders']) """Worker details needed to run a Flatten. Attributes: inputs: A list of tuples, each (producer index, output index), representing the ParallelInstruction operations whose output feeds into this operation. The output index is 0 unless the input is from a multi-output operation (such as ParDo). output_coders: 1-tuple of the coder for the output. """ def get_coder_from_spec(coder_spec): """Return a coder instance from a coder spec. Args: coder_spec: A dict where the value of the '@type' key is a pickled instance of a Coder instance. Returns: A coder instance (has encode/decode methods). """ assert coder_spec is not None # Ignore the wrappers in these encodings. ignored_wrappers = ( 'com.google.cloud.dataflow.sdk.util.TimerOrElement$TimerOrElementCoder') if coder_spec['@type'] in ignored_wrappers: assert len(coder_spec['component_encodings']) == 1 coder_spec = coder_spec['component_encodings'][0] return get_coder_from_spec(coder_spec) # Handle a few well known types of coders. if coder_spec['@type'] == 'kind:pair': assert len(coder_spec['component_encodings']) == 2 component_coders = [ get_coder_from_spec(c) for c in coder_spec['component_encodings'] ] return coders.TupleCoder(component_coders) elif coder_spec['@type'] == 'kind:stream': assert len(coder_spec['component_encodings']) == 1 return coders.IterableCoder( get_coder_from_spec(coder_spec['component_encodings'][0])) elif coder_spec['@type'] == 'kind:windowed_value': assert len(coder_spec['component_encodings']) == 2 value_coder, window_coder = [ get_coder_from_spec(c) for c in coder_spec['component_encodings']] return coders.coders.WindowedValueCoder( value_coder, window_coder=window_coder) elif coder_spec['@type'] == 'kind:interval_window': assert ( 'component_encodings' not in coder_spec or not coder_spec['component_encodings']) return coders.coders.IntervalWindowCoder() elif coder_spec['@type'] == 'kind:global_window': assert ( 'component_encodings' not in coder_spec or not coder_spec['component_encodings']) return coders.coders.GlobalWindowCoder() elif coder_spec['@type'] == 'kind:varint': assert ( 'component_encodings' not in coder_spec or len(coder_spec['component_encodings'] == 0)) return coders.coders.VarIntCoder() elif coder_spec['@type'] == 'kind:length_prefix': assert len(coder_spec['component_encodings']) == 1 return coders.coders.LengthPrefixCoder( get_coder_from_spec(coder_spec['component_encodings'][0])) elif coder_spec['@type'] == 'kind:bytes': assert ( 'component_encodings' not in coder_spec or len(coder_spec['component_encodings'] == 0)) return coders.BytesCoder() # We pass coders in the form "$" to make the job # description JSON more readable. return coders.coders.deserialize_coder(coder_spec['@type'].encode('ascii')) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/operations.pxd0000644000175100017510000001104515206113032023563 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # cimport cython from apache_beam.runners.common cimport DoFnRunner from apache_beam.runners.common cimport Receiver from apache_beam.runners.worker cimport opcounters from apache_beam.utils.windowed_value cimport WindowedValue from apache_beam.utils.windowed_value cimport WindowedBatch #from libcpp.string cimport string cdef WindowedValue _globally_windowed_value cdef type _global_window_type cdef class ConsumerSet(Receiver): cdef list consumers cdef readonly opcounters.OperationCounters opcounter cdef public step_name cdef public output_index cdef public coder cdef public object output_sampler cdef public object element_sampler cdef public object execution_context cpdef update_counters_start(self, WindowedValue windowed_value) cpdef update_counters_finish(self) cpdef update_counters_batch(self, WindowedBatch windowed_batch) cdef class SingletonElementConsumerSet(ConsumerSet): cdef Operation consumer cpdef receive(self, WindowedValue windowed_value) cpdef receive_batch(self, WindowedBatch windowed_batch) cpdef flush(self) cdef class GeneralPurposeConsumerSet(ConsumerSet): cdef list element_consumers cdef list passthrough_batch_consumers cdef dict other_batch_consumers cdef bint has_batch_consumers cdef list _batched_elements cdef object producer_batch_converter cpdef receive(self, WindowedValue windowed_value) cpdef receive_batch(self, WindowedBatch windowed_batch) cpdef flush(self) cdef class Operation(object): cdef readonly name_context cdef readonly operation_name cdef readonly spec cdef object consumers cdef readonly counter_factory cdef public metrics_container cdef public execution_context # Public for access by Fn harness operations. # TODO(robertwb): Cythonize FnHarness. cdef public list receivers cdef readonly bint debug_logging_enabled # For legacy workers. cdef bint setup_done cdef public step_name # initialized lazily cdef readonly object state_sampler cdef readonly object scoped_start_state cdef readonly object scoped_process_state cdef readonly object scoped_finish_state cdef readonly object data_sampler cpdef start(self) cpdef process(self, WindowedValue windowed_value) cpdef finish(self) cpdef teardown(self) cpdef output(self, WindowedValue windowed_value, int output_index=*) cpdef execution_time_monitoring_infos(self, transform_id) cpdef user_monitoring_infos(self, transform_id) cpdef pcollection_count_monitoring_infos(self, tag_to_pcollection_id) cpdef monitoring_infos(self, transform_id, tag_to_pcollection_id) cdef class ReadOperation(Operation): @cython.locals(windowed_value=WindowedValue) cpdef start(self) cdef class ImpulseReadOperation(Operation): cdef object source @cython.locals(windowed_value=WindowedValue) cpdef process(self, WindowedValue impulse) cdef class DoOperation(Operation): cdef DoFnRunner dofn_runner cdef object tagged_receivers cdef object side_input_maps cdef object user_state_context cdef public dict timer_inputs cdef dict timer_specs cdef public object input_info cdef object fn cdef object scoped_timer_processing_state cdef class SdfProcessSizedElements(DoOperation): cdef object lock cdef object element_start_output_bytes cdef class SdfTruncateSizedRestrictions(DoOperation): pass cdef class CombineOperation(Operation): cdef object phased_combine_fn cdef class PGBKCVOperation(Operation): cdef public object combine_fn cdef public object combine_fn_add_input cdef public object combine_fn_compact cdef public bint is_default_windowing cdef public object timestamp_combiner cdef dict table cdef long max_keys cdef long key_count cpdef add_key_value(self, wkey, value, timestamp) cpdef output_key(self, wkey, value, timestamp) cdef class FlattenOperation(Operation): pass ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/operations.py0000644000175100017510000015672315206113032023435 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Worker operations executor.""" # pytype: skip-file # pylint: disable=super-with-arguments # ruff: noqa: UP006 import collections import logging import threading import warnings from typing import TYPE_CHECKING from typing import Any from typing import DefaultDict from typing import Dict from typing import FrozenSet from typing import Iterable from typing import Iterator from typing import List from typing import Mapping from typing import NamedTuple from typing import Optional from typing import Tuple from apache_beam import coders from apache_beam.internal import pickler from apache_beam.io import iobase from apache_beam.metrics import monitoring_infos from apache_beam.metrics.cells import DistributionData from apache_beam.metrics.execution import MetricsContainer from apache_beam.portability.api import metrics_pb2 from apache_beam.runners import common from apache_beam.runners.common import Receiver from apache_beam.runners.worker import opcounters from apache_beam.runners.worker import operation_specs from apache_beam.runners.worker import sideinputs from apache_beam.runners.worker.data_sampler import DataSampler from apache_beam.transforms import combiners from apache_beam.transforms import core from apache_beam.transforms import sideinputs as apache_sideinputs from apache_beam.transforms import userstate from apache_beam.transforms import window from apache_beam.transforms.combiners import PhasedCombineFnExecutor from apache_beam.transforms.combiners import curry_combine_fn from apache_beam.transforms.window import GlobalWindows from apache_beam.typehints.batch import BatchConverter from apache_beam.utils.windowed_value import WindowedBatch from apache_beam.utils.windowed_value import WindowedValue if TYPE_CHECKING: from apache_beam.runners.sdf_utils import SplitResultPrimary from apache_beam.runners.sdf_utils import SplitResultResidual from apache_beam.runners.worker.bundle_processor import ExecutionContext from apache_beam.runners.worker.data_sampler import OutputSampler from apache_beam.runners.worker.statesampler import StateSampler from apache_beam.transforms.userstate import TimerSpec # Allow some "pure mode" declarations. try: import cython except ImportError: class FakeCython(object): compiled = False globals()['cython'] = FakeCython() _globally_windowed_value = GlobalWindows.windowed_value(None) _global_window_type = type(_globally_windowed_value.windows[0]) _LOGGER = logging.getLogger(__name__) SdfSplitResultsPrimary = Tuple['DoOperation', 'SplitResultPrimary'] SdfSplitResultsResidual = Tuple['DoOperation', 'SplitResultResidual'] # TODO(BEAM-9324) Remove these workarounds once upgraded to Cython 3 def _cast_to_operation(value): if cython.compiled: return cython.cast(Operation, value) else: return value # TODO(BEAM-9324) Remove these workarounds once upgraded to Cython 3 def _cast_to_receiver(value): if cython.compiled: return cython.cast(Receiver, value) else: return value class ConsumerSet(Receiver): """A ConsumerSet represents a graph edge between two Operation nodes. The ConsumerSet object collects information from the output of the Operation at one end of its edge and the input of the Operation at the other edge. ConsumerSet are attached to the outputting Operation. """ @staticmethod def create( counter_factory, step_name, # type: str output_index, consumers, # type: List[Operation] coder, producer_type_hints, producer_batch_converter, # type: Optional[BatchConverter] output_sampler=None, # type: Optional[OutputSampler] ): # type: (...) -> ConsumerSet if len(consumers) == 1: consumer = consumers[0] consumer_batch_preference = consumer.get_batching_preference() consumer_batch_converter = consumer.get_input_batch_converter() if (not consumer_batch_preference.supports_batches and producer_batch_converter is None and consumer_batch_converter is None): return SingletonElementConsumerSet( counter_factory, step_name, output_index, consumer, coder, producer_type_hints, output_sampler) return GeneralPurposeConsumerSet( counter_factory, step_name, output_index, coder, producer_type_hints, consumers, producer_batch_converter, output_sampler) def __init__( self, counter_factory, step_name, # type: str output_index, consumers, coder, producer_type_hints, producer_batch_converter, output_sampler): self.opcounter = opcounters.OperationCounters( counter_factory, step_name, coder, output_index, producer_type_hints=producer_type_hints, producer_batch_converter=producer_batch_converter) # Used in repr. self.step_name = step_name self.output_index = output_index self.coder = coder self.consumers = consumers self.output_sampler = output_sampler self.element_sampler = ( output_sampler.element_sampler if output_sampler else None) self.execution_context = None # type: Optional[ExecutionContext] def try_split(self, fraction_of_remainder): # type: (...) -> Optional[Any] # TODO(SDF): Consider supporting splitting each consumer individually. # This would never come up in the existing SDF expansion, but might # be useful to support fused SDF nodes. # This would require dedicated delivery of the split results to each # of the consumers separately. return None def current_element_progress(self): # type: () -> Optional[iobase.RestrictionProgress] """Returns the progress of the current element. This progress should be an instance of apache_beam.io.iobase.RestrictionProgress, or None if progress is unknown. """ # TODO(SDF): Could implement this as a weighted average, if it becomes # useful to split on. return None def update_counters_start(self, windowed_value): # type: (WindowedValue) -> None self.opcounter.update_from(windowed_value) if self.execution_context is not None: self.execution_context.output_sampler = self.output_sampler # The following code is optimized by inlining a function call. Because this # is called for every element, a function call is too expensive (order of # 100s of nanoseconds). Furthermore, a lock was purposefully not used # between here and the DataSampler as an additional operation. The tradeoff # is that some samples might be dropped, but it is better than the # alternative which is double sampling the same element. if self.element_sampler is not None: if not self.element_sampler.has_element: self.element_sampler.el = windowed_value self.element_sampler.has_element = True def update_counters_finish(self): # type: () -> None self.opcounter.update_collect() def update_counters_batch(self, windowed_batch): # type: (WindowedBatch) -> None self.opcounter.update_from_batch(windowed_batch) def __repr__(self): return '%s[%s.out%s, coder=%s, len(consumers)=%s]' % ( self.__class__.__name__, self.step_name, self.output_index, self.coder, len(self.consumers)) class SingletonElementConsumerSet(ConsumerSet): """ConsumerSet representing a single consumer that can only process elements (not batches).""" def __init__( self, counter_factory, step_name, output_index, consumer, # type: Operation coder, producer_type_hints, output_sampler): super().__init__( counter_factory, step_name, output_index, [consumer], coder, producer_type_hints, None, output_sampler) self.consumer = consumer def receive(self, windowed_value): # type: (WindowedValue) -> None self.update_counters_start(windowed_value) self.consumer.process(windowed_value) self.update_counters_finish() def receive_batch(self, windowed_batch): raise AssertionError( "SingletonElementConsumerSet.receive_batch is not implemented") def flush(self): # SingletonElementConsumerSet has no buffer to flush pass def try_split(self, fraction_of_remainder): # type: (...) -> Optional[Any] return self.consumer.try_split(fraction_of_remainder) def current_element_progress(self): return self.consumer.current_element_progress() class GeneralPurposeConsumerSet(ConsumerSet): """ConsumerSet implementation that handles all combinations of possible edges. """ MAX_BATCH_SIZE = 4096 def __init__( self, counter_factory, step_name, # type: str output_index, coder, producer_type_hints, consumers, # type: List[Operation] producer_batch_converter, output_sampler): super().__init__( counter_factory, step_name, output_index, consumers, coder, producer_type_hints, producer_batch_converter, output_sampler) self.producer_batch_converter = producer_batch_converter # Partition consumers into three groups: # - consumers that will be passed elements # - consumers that will be passed batches (where their input batch type # matches the output of the producer) # - consumers that will be passed converted batches self.element_consumers: List[Operation] = [] self.passthrough_batch_consumers: List[Operation] = [] other_batch_consumers: DefaultDict[ BatchConverter, List[Operation]] = collections.defaultdict(lambda: []) for consumer in consumers: if not consumer.get_batching_preference().supports_batches: self.element_consumers.append(consumer) elif (consumer.get_input_batch_converter() == self.producer_batch_converter): self.passthrough_batch_consumers.append(consumer) else: # Batch consumer with a mismatched batch type if consumer.get_batching_preference().supports_elements: # Pass it elements if we can self.element_consumers.append(consumer) else: # As a last resort, explode and rebatch consumer_batch_converter = consumer.get_input_batch_converter() # This consumer supports batches, it must have a batch converter assert consumer_batch_converter is not None other_batch_consumers[consumer_batch_converter].append(consumer) self.other_batch_consumers: Dict[BatchConverter, List[Operation]] = dict( other_batch_consumers) self.has_batch_consumers = ( self.passthrough_batch_consumers or self.other_batch_consumers) self._batched_elements: List[Any] = [] def receive(self, windowed_value): # type: (WindowedValue) -> None self.update_counters_start(windowed_value) for consumer in self.element_consumers: _cast_to_operation(consumer).process(windowed_value) # TODO: Do this branching when contstructing ConsumerSet if self.has_batch_consumers: self._batched_elements.append(windowed_value) if len(self._batched_elements) >= self.MAX_BATCH_SIZE: self.flush() # TODO(https://github.com/apache/beam/issues/21655): Properly estimate # sizes in the batch-consumer only case, this undercounts large iterables self.update_counters_finish() def receive_batch(self, windowed_batch): if self.element_consumers: for wv in windowed_batch.as_windowed_values( self.producer_batch_converter.explode_batch): for consumer in self.element_consumers: _cast_to_operation(consumer).process(wv) for consumer in self.passthrough_batch_consumers: _cast_to_operation(consumer).process_batch(windowed_batch) for (consumer_batch_converter, consumers) in self.other_batch_consumers.items(): # Explode and rebatch into the new batch type (ouch!) # TODO: Register direct conversions for equivalent batch types for consumer in consumers: warnings.warn( f"Input to operation {consumer} must be rebatched from type " f"{self.producer_batch_converter.batch_type!r} to " f"{consumer_batch_converter.batch_type!r}.\n" "This is very inefficient, consider re-structuring your pipeline " "or adding a DoFn to directly convert between these types.", InefficientExecutionWarning) _cast_to_operation(consumer).process_batch( windowed_batch.with_values( consumer_batch_converter.produce_batch( self.producer_batch_converter.explode_batch( windowed_batch.values)))) self.update_counters_batch(windowed_batch) def flush(self): if not self.has_batch_consumers or not self._batched_elements: return for batch_converter, consumers in self.other_batch_consumers.items(): for windowed_batch in WindowedBatch.from_windowed_values( self._batched_elements, produce_fn=batch_converter.produce_batch): for consumer in consumers: _cast_to_operation(consumer).process_batch(windowed_batch) for consumer in self.passthrough_batch_consumers: for windowed_batch in WindowedBatch.from_windowed_values( self._batched_elements, produce_fn=self.producer_batch_converter.produce_batch): _cast_to_operation(consumer).process_batch(windowed_batch) self._batched_elements = [] class Operation(object): """An operation representing the live version of a work item specification. An operation can have one or more outputs and for each output it can have one or more receiver operations that will take that as input. """ def __init__( self, name_context, # type: common.NameContext spec, counter_factory, state_sampler # type: StateSampler ): """Initializes a worker operation instance. Args: name_context: A NameContext instance, with the name information for this operation. spec: A operation_specs.Worker* instance. counter_factory: The CounterFactory to use for our counters. state_sampler: The StateSampler for the current operation. """ assert isinstance(name_context, common.NameContext) self.name_context = name_context self.spec = spec self.counter_factory = counter_factory self.execution_context = None # type: Optional[ExecutionContext] self.consumers = collections.defaultdict( list) # type: DefaultDict[int, List[Operation]] # These are overwritten in the legacy harness. self.metrics_container = MetricsContainer(self.name_context.metrics_name()) self.state_sampler = state_sampler self.scoped_start_state = self.state_sampler.scoped_state( self.name_context, 'start', metrics_container=self.metrics_container) self.scoped_process_state = self.state_sampler.scoped_state( self.name_context, 'process', metrics_container=self.metrics_container) self.scoped_finish_state = self.state_sampler.scoped_state( self.name_context, 'finish', metrics_container=self.metrics_container) # TODO(ccy): the '-abort' state can be added when the abort is supported in # Operations. self.receivers = [] # type: List[ConsumerSet] # Legacy workers cannot call setup() until after setting additional state # on the operation. self.setup_done = False self.step_name = None # type: Optional[str] self.data_sampler: Optional[DataSampler] = None def setup(self, data_sampler=None): # type: (Optional[DataSampler]) -> None """Set up operation. This must be called before any other methods of the operation.""" with self.scoped_start_state: self.data_sampler = data_sampler self.debug_logging_enabled = logging.getLogger().isEnabledFor( logging.DEBUG) transform_id = self.name_context.transform_id # Everything except WorkerSideInputSource, which is not a # top-level operation, should have output_coders #TODO(pabloem): Define better what step name is used here. if getattr(self.spec, 'output_coders', None): def get_output_sampler(output_num): if data_sampler is None: return None return data_sampler.sampler_for_output(transform_id, output_num) self.receivers = [ ConsumerSet.create( self.counter_factory, self.name_context.logging_name(), i, self.consumers[i], coder, self._get_runtime_performance_hints(), self.get_output_batch_converter(), get_output_sampler(i)) for i, coder in enumerate(self.spec.output_coders) ] self.setup_done = True def start(self): # type: () -> None """Start operation.""" if not self.setup_done: # For legacy workers. self.setup(self.data_sampler) # The ExecutionContext is per instruction and so cannot be set at # initialization time. if self.data_sampler is not None: for receiver in self.receivers: receiver.execution_context = self.execution_context def get_batching_preference(self): # By default operations don't support batching, require Receiver to unbatch return common.BatchingPreference.BATCH_FORBIDDEN def get_input_batch_converter(self) -> Optional[BatchConverter]: """Returns a batch type converter if this operation can accept a batch, otherwise None.""" return None def get_output_batch_converter(self) -> Optional[BatchConverter]: """Returns a batch type converter if this operation can produce a batch, otherwise None.""" return None def process(self, o): # type: (WindowedValue) -> None """Process element in operation.""" pass def process_batch(self, batch: WindowedBatch): pass def finalize_bundle(self): # type: () -> None pass def needs_finalization(self): return False def try_split(self, fraction_of_remainder): # type: (...) -> Optional[Any] return None def current_element_progress(self): return None def finish(self): # type: () -> None """Finish operation.""" for receiver in self.receivers: _cast_to_receiver(receiver).flush() def teardown(self): # type: () -> None """Tear down operation. No other methods of this operation should be called after this.""" pass def reset(self): # type: () -> None self.metrics_container.reset() def output(self, windowed_value, output_index=0): # type: (WindowedValue, int) -> None _cast_to_receiver(self.receivers[output_index]).receive(windowed_value) def add_receiver(self, operation, output_index=0): # type: (Operation, int) -> None """Adds a receiver operation for the specified output.""" self.consumers[output_index].append(operation) def monitoring_infos(self, transform_id, tag_to_pcollection_id): # type: (str, Dict[str, str]) -> Dict[FrozenSet, metrics_pb2.MonitoringInfo] """Returns the list of MonitoringInfos collected by this operation.""" all_monitoring_infos = self.execution_time_monitoring_infos(transform_id) all_monitoring_infos.update( self.pcollection_count_monitoring_infos(tag_to_pcollection_id)) all_monitoring_infos.update(self.user_monitoring_infos(transform_id)) return all_monitoring_infos def pcollection_count_monitoring_infos(self, tag_to_pcollection_id): # type: (Dict[str, str]) -> Dict[FrozenSet, metrics_pb2.MonitoringInfo] """Returns the element count MonitoringInfo collected by this operation.""" # Skip producing monitoring infos if there is more then one receiver # since there is no way to provide a mapping from tag to pcollection id # within Operation. if len(self.receivers) != 1 or len(tag_to_pcollection_id) != 1: return {} all_monitoring_infos = {} pcollection_id = next(iter(tag_to_pcollection_id.values())) receiver = self.receivers[0] elem_count_mi = monitoring_infos.int64_counter( monitoring_infos.ELEMENT_COUNT_URN, receiver.opcounter.element_counter.value(), pcollection=pcollection_id, ) (unused_mean, sum, count, min, max) = ( receiver.opcounter.mean_byte_counter.value()) sampled_byte_count = monitoring_infos.int64_distribution( monitoring_infos.SAMPLED_BYTE_SIZE_URN, DistributionData(sum, count, min, max), pcollection=pcollection_id, ) all_monitoring_infos[monitoring_infos.to_key(elem_count_mi)] = elem_count_mi all_monitoring_infos[monitoring_infos.to_key( sampled_byte_count)] = sampled_byte_count return all_monitoring_infos def user_monitoring_infos(self, transform_id): # type: (str) -> Dict[FrozenSet, metrics_pb2.MonitoringInfo] """Returns the user MonitoringInfos collected by this operation.""" return self.metrics_container.to_runner_api_monitoring_infos(transform_id) def execution_time_monitoring_infos(self, transform_id): # type: (str) -> Dict[FrozenSet, metrics_pb2.MonitoringInfo] total_time_spent_msecs = ( self.scoped_start_state.sampled_msecs_int() + self.scoped_process_state.sampled_msecs_int() + self.scoped_finish_state.sampled_msecs_int()) mis = [ monitoring_infos.int64_counter( monitoring_infos.START_BUNDLE_MSECS_URN, self.scoped_start_state.sampled_msecs_int(), ptransform=transform_id), monitoring_infos.int64_counter( monitoring_infos.PROCESS_BUNDLE_MSECS_URN, self.scoped_process_state.sampled_msecs_int(), ptransform=transform_id), monitoring_infos.int64_counter( monitoring_infos.FINISH_BUNDLE_MSECS_URN, self.scoped_finish_state.sampled_msecs_int(), ptransform=transform_id), monitoring_infos.int64_counter( monitoring_infos.TOTAL_MSECS_URN, total_time_spent_msecs, ptransform=transform_id), ] return {monitoring_infos.to_key(mi): mi for mi in mis} def __str__(self): """Generates a useful string for this object. Compactly displays interesting fields. In particular, pickled fields are not displayed. Note that we collapse the fields of the contained Worker* object into this object, since there is a 1-1 mapping between Operation and operation_specs.Worker*. Returns: Compact string representing this object. """ return self.str_internal() def str_internal(self, is_recursive=False): """Internal helper for __str__ that supports recursion. When recursing on receivers, keep the output short. Args: is_recursive: whether to omit some details, particularly receivers. Returns: Compact string representing this object. """ printable_name = self.__class__.__name__ if hasattr(self, 'step_name'): printable_name += ' %s' % self.name_context.logging_name() if is_recursive: # If we have a step name, stop here, no more detail needed. return '<%s>' % printable_name if self.spec is None: printable_fields = [] else: printable_fields = operation_specs.worker_printable_fields(self.spec) if not is_recursive and getattr(self, 'receivers', []): printable_fields.append( 'receivers=[%s]' % ', '.join([str(receiver) for receiver in self.receivers])) return '<%s %s>' % (printable_name, ', '.join(printable_fields)) def _get_runtime_performance_hints(self): # type: () -> Optional[Dict[Optional[str], Tuple[Optional[str], Any]]] """Returns any type hints required for performance runtime type-checking.""" return None class ReadOperation(Operation): def start(self): with self.scoped_start_state: super(ReadOperation, self).start() range_tracker = self.spec.source.source.get_range_tracker( self.spec.source.start_position, self.spec.source.stop_position) for value in self.spec.source.source.read(range_tracker): if isinstance(value, WindowedValue): windowed_value = value else: windowed_value = _globally_windowed_value.with_value(value) self.output(windowed_value) class ImpulseReadOperation(Operation): def __init__( self, name_context, # type: common.NameContext counter_factory, state_sampler, # type: StateSampler consumers, # type: Mapping[Any, List[Operation]] source, # type: iobase.BoundedSource output_coder): super(ImpulseReadOperation, self).__init__(name_context, None, counter_factory, state_sampler) self.source = source self.receivers = [ ConsumerSet.create( self.counter_factory, self.name_context.step_name, 0, next(iter(consumers.values())), output_coder, self._get_runtime_performance_hints(), self.get_output_batch_converter()) ] def process(self, unused_impulse): # type: (WindowedValue) -> None with self.scoped_process_state: range_tracker = self.source.get_range_tracker(None, None) for value in self.source.read(range_tracker): if isinstance(value, WindowedValue): windowed_value = value else: windowed_value = _globally_windowed_value.with_value(value) self.output(windowed_value) class InMemoryWriteOperation(Operation): """A write operation that will write to an in-memory sink.""" def process(self, o): # type: (WindowedValue) -> None with self.scoped_process_state: if self.debug_logging_enabled: _LOGGER.debug('Processing [%s] in %s', o, self) self.spec.output_buffer.append( o if self.spec.write_windowed_values else o.value) class _TaggedReceivers(dict): def __init__(self, counter_factory, step_name): self._counter_factory = counter_factory self._step_name = step_name def __missing__(self, tag): self[tag] = receiver = ConsumerSet.create( self._counter_factory, self._step_name, tag, [], None, None, None) return receiver def total_output_bytes(self): # type: () -> int total = 0 for receiver in self.values(): elements = receiver.opcounter.element_counter.value() if elements > 0: mean = (receiver.opcounter.mean_byte_counter.value())[0] total += elements * mean return total OpInputInfo = NamedTuple( 'OpInputInfo', [ ('transform_id', str), ('main_input_tag', str), ('main_input_coder', coders.WindowedValueCoder), ('outputs', Iterable[str]), ]) class DoOperation(Operation): """A Do operation that will execute a custom DoFn for each input element.""" def __init__( self, name, # type: common.NameContext spec, # operation_specs.WorkerDoFn # need to fix this type counter_factory, sampler, side_input_maps=None, user_state_context=None, ): super(DoOperation, self).__init__(name, spec, counter_factory, sampler) self.side_input_maps = side_input_maps self.user_state_context = user_state_context self.tagged_receivers = None # type: Optional[_TaggedReceivers] # A mapping of timer tags to the input "PCollections" they come in on. self.input_info = None # type: Optional[OpInputInfo] self.scoped_timer_processing_state = self.state_sampler.scoped_state( self.name_context, 'process-timers', metrics_container=self.metrics_container) # See fn_data in dataflow_runner.py # TODO: Store all the items from spec? self.fn, _, _, _, _ = (pickler.loads(self.spec.serialized_fn)) def _read_side_inputs(self, tags_and_types): # type: (...) -> Iterator[apache_sideinputs.SideInputMap] """Generator reading side inputs in the order prescribed by tags_and_types. Args: tags_and_types: List of tuples (tag, type). Each side input has a string tag that is specified in the worker instruction. The type is actually a boolean which is True for singleton input (read just first value) and False for collection input (read all values). Yields: With each iteration it yields the result of reading an entire side source either in singleton or collection mode according to the tags_and_types argument. """ # Only call this on the old path where side_input_maps was not # provided directly. assert self.side_input_maps is None # We will read the side inputs in the order prescribed by the # tags_and_types argument because this is exactly the order needed to # replace the ArgumentPlaceholder objects in the args/kwargs of the DoFn # getting the side inputs. # # Note that for each tag there could be several read operations in the # specification. This can happen for instance if the source has been # sharded into several files. for i, (side_tag, view_class, view_options) in enumerate(tags_and_types): sources = [] # Using the side_tag in the lambda below will trigger a pylint warning. # However in this case it is fine because the lambda is used right away # while the variable has the value assigned by the current iteration of # the for loop. # pylint: disable=cell-var-from-loop for si in filter(lambda o: o.tag == side_tag, self.spec.side_inputs): if not isinstance(si, operation_specs.WorkerSideInputSource): raise NotImplementedError('Unknown side input type: %r' % si) sources.append(si.source) si_counter = opcounters.SideInputReadCounter( self.counter_factory, self.state_sampler, declaring_step=self.name_context.step_name, # Inputs are 1-indexed, so we add 1 to i in the side input id input_index=i + 1) element_counter = opcounters.OperationCounters( self.counter_factory, self.name_context.step_name, view_options['coder'], i, suffix='side-input') iterator_fn = sideinputs.get_iterator_fn_for_sources( sources, read_counter=si_counter, element_counter=element_counter) yield apache_sideinputs.SideInputMap( view_class, view_options, sideinputs.EmulatedIterable(iterator_fn)) def setup(self, data_sampler=None): # type: (Optional[DataSampler]) -> None with self.scoped_start_state: super(DoOperation, self).setup(data_sampler) # See fn_data in dataflow_runner.py fn, args, kwargs, tags_and_types, window_fn = ( pickler.loads(self.spec.serialized_fn)) state = common.DoFnState(self.counter_factory) state.step_name = self.name_context.logging_name() # Tag to output index map used to dispatch the output values emitted # by the DoFn function to the appropriate receivers. The main output is # either the only output or the output tagged with 'None' and is # associated with its corresponding index. self.tagged_receivers = _TaggedReceivers( self.counter_factory, self.name_context.logging_name()) if len(self.spec.output_tags) == 1: self.tagged_receivers[None] = self.receivers[0] self.tagged_receivers[self.spec.output_tags[0]] = self.receivers[0] else: for index, tag in enumerate(self.spec.output_tags): self.tagged_receivers[tag] = self.receivers[index] if tag == 'None': self.tagged_receivers[None] = self.receivers[index] if self.user_state_context: self.timer_specs = { spec.name: spec for spec in userstate.get_dofn_specs(fn)[1] } # type: Dict[str, TimerSpec] if self.side_input_maps is None: if tags_and_types: self.side_input_maps = list(self._read_side_inputs(tags_and_types)) else: self.side_input_maps = [] self.dofn_runner = common.DoFnRunner( fn, args, kwargs, self.side_input_maps, window_fn, tagged_receivers=self.tagged_receivers, step_name=self.name_context.logging_name(), state=state, user_state_context=self.user_state_context, transform_id=self.name_context.transform_id, operation_name=self.name_context.metrics_name()) self.dofn_runner.setup() def start(self): # type: () -> None with self.scoped_start_state: super(DoOperation, self).start() self.dofn_runner.execution_context = self.execution_context self.dofn_runner.start() def get_batching_preference(self): if self.fn._process_batch_defined: if self.fn._process_defined: return common.BatchingPreference.DO_NOT_CARE else: return common.BatchingPreference.BATCH_REQUIRED else: return common.BatchingPreference.BATCH_FORBIDDEN def get_input_batch_converter(self) -> Optional[BatchConverter]: return getattr(self.fn, 'input_batch_converter', None) def get_output_batch_converter(self) -> Optional[BatchConverter]: return getattr(self.fn, 'output_batch_converter', None) def process(self, o): # type: (WindowedValue) -> None with self.scoped_process_state: delayed_applications = self.dofn_runner.process(o) if delayed_applications: assert self.execution_context is not None for delayed_application in delayed_applications: self.execution_context.delayed_applications.append( (self, delayed_application)) def process_batch(self, windowed_batch: WindowedBatch) -> None: self.dofn_runner.process_batch(windowed_batch) def finalize_bundle(self): # type: () -> None self.dofn_runner.finalize() def needs_finalization(self): # type: () -> bool return self.dofn_runner.bundle_finalizer_param.has_callbacks() def add_timer_info(self, timer_family_id, timer_info): self.user_state_context.add_timer_info(timer_family_id, timer_info) def process_timer(self, tag, timer_data): with self.scoped_timer_processing_state: timer_spec = self.timer_specs[tag] self.dofn_runner.process_user_timer( timer_spec, timer_data.user_key, timer_data.windows[0], timer_data.fire_timestamp, timer_data.paneinfo, timer_data.dynamic_timer_tag) def finish(self): # type: () -> None super(DoOperation, self).finish() with self.scoped_finish_state: self.dofn_runner.finish() if self.user_state_context: self.user_state_context.commit() def teardown(self): # type: () -> None with self.scoped_finish_state: self.dofn_runner.teardown() if self.user_state_context: self.user_state_context.reset() def reset(self): # type: () -> None super(DoOperation, self).reset() for side_input_map in self.side_input_maps: side_input_map.reset() if self.user_state_context: self.user_state_context.reset() self.dofn_runner.bundle_finalizer_param.reset() def pcollection_count_monitoring_infos(self, tag_to_pcollection_id): # type: (Dict[str, str]) -> Dict[FrozenSet, metrics_pb2.MonitoringInfo] """Returns the element count MonitoringInfo collected by this operation.""" infos = super( DoOperation, self).pcollection_count_monitoring_infos(tag_to_pcollection_id) if self.tagged_receivers: for tag, receiver in self.tagged_receivers.items(): if str(tag) not in tag_to_pcollection_id: continue pcollection_id = tag_to_pcollection_id[str(tag)] mi = monitoring_infos.int64_counter( monitoring_infos.ELEMENT_COUNT_URN, receiver.opcounter.element_counter.value(), pcollection=pcollection_id) infos[monitoring_infos.to_key(mi)] = mi (unused_mean, sum, count, min, max) = ( receiver.opcounter.mean_byte_counter.value()) sampled_byte_count = monitoring_infos.int64_distribution( monitoring_infos.SAMPLED_BYTE_SIZE_URN, DistributionData(sum, count, min, max), pcollection=pcollection_id) infos[monitoring_infos.to_key(sampled_byte_count)] = sampled_byte_count return infos def _get_runtime_performance_hints(self): fns = pickler.loads(self.spec.serialized_fn) if fns and hasattr(fns[0], '_runtime_output_constraints'): return fns[0]._runtime_output_constraints return {} class SdfTruncateSizedRestrictions(DoOperation): def __init__(self, *args, **kwargs): super(SdfTruncateSizedRestrictions, self).__init__(*args, **kwargs) def current_element_progress(self): # type: () -> Optional[iobase.RestrictionProgress] return self.receivers[0].current_element_progress() def try_split( self, fraction_of_remainder ): # type: (...) -> Optional[Tuple[Iterable[SdfSplitResultsPrimary], Iterable[SdfSplitResultsResidual]]] return self.receivers[0].try_split(fraction_of_remainder) class SdfProcessSizedElements(DoOperation): def __init__(self, *args, **kwargs): super(SdfProcessSizedElements, self).__init__(*args, **kwargs) self.lock = threading.RLock() self.element_start_output_bytes = None # type: Optional[int] def process(self, o): # type: (WindowedValue) -> None assert self.tagged_receivers is not None with self.scoped_process_state: try: with self.lock: self.element_start_output_bytes = \ self.tagged_receivers.total_output_bytes() for receiver in self.tagged_receivers.values(): receiver.opcounter.restart_sampling() # Actually processing the element can be expensive; do it without # the lock. delayed_applications = self.dofn_runner.process_with_sized_restriction( o) if delayed_applications: assert self.execution_context is not None for delayed_application in delayed_applications: self.execution_context.delayed_applications.append( (self, delayed_application)) finally: with self.lock: self.element_start_output_bytes = None def try_split(self, fraction_of_remainder): # type: (...) -> Optional[Tuple[Iterable[SdfSplitResultsPrimary], Iterable[SdfSplitResultsResidual]]] split = self.dofn_runner.try_split(fraction_of_remainder) if split: primaries, residuals = split return [(self, primary) for primary in primaries ], [(self, residual) for residual in residuals] return None def current_element_progress(self): # type: () -> Optional[iobase.RestrictionProgress] with self.lock: if self.element_start_output_bytes is not None: progress = self.dofn_runner.current_element_progress() if progress is not None: assert self.tagged_receivers is not None return progress.with_completed( self.tagged_receivers.total_output_bytes() - self.element_start_output_bytes) return None def monitoring_infos(self, transform_id, tag_to_pcollection_id): # type: (str, Dict[str, str]) -> Dict[FrozenSet, metrics_pb2.MonitoringInfo] progress_coder = coders.IterableCoder(coders.FloatCoder()) with self.lock: infos = super(SdfProcessSizedElements, self).monitoring_infos(transform_id, tag_to_pcollection_id) current_element_progress = self.current_element_progress() if current_element_progress: if current_element_progress.completed_work: completed = current_element_progress.completed_work remaining = current_element_progress.remaining_work else: completed = current_element_progress.fraction_completed remaining = current_element_progress.fraction_remaining assert completed is not None assert remaining is not None completed_mi = metrics_pb2.MonitoringInfo( urn=monitoring_infos.WORK_COMPLETED_URN, type=monitoring_infos.PROGRESS_TYPE, labels=monitoring_infos.create_labels(ptransform=transform_id), payload=progress_coder.encode([completed])) remaining_mi = metrics_pb2.MonitoringInfo( urn=monitoring_infos.WORK_REMAINING_URN, type=monitoring_infos.PROGRESS_TYPE, labels=monitoring_infos.create_labels(ptransform=transform_id), payload=progress_coder.encode([remaining])) infos[monitoring_infos.to_key(completed_mi)] = completed_mi infos[monitoring_infos.to_key(remaining_mi)] = remaining_mi return infos class CombineOperation(Operation): """A Combine operation executing a CombineFn for each input element.""" def __init__(self, name_context, spec, counter_factory, state_sampler): super(CombineOperation, self).__init__(name_context, spec, counter_factory, state_sampler) # Combiners do not accept deferred side-inputs (the ignored fourth argument) # and therefore the code to handle the extra args/kwargs is simpler than for # the DoFn's of ParDo. fn, args, kwargs = pickler.loads(self.spec.serialized_fn)[:3] self.phased_combine_fn = ( PhasedCombineFnExecutor(self.spec.phase, fn, args, kwargs)) def setup(self, data_sampler=None): # type: (Optional[DataSampler]) -> None with self.scoped_start_state: _LOGGER.debug('Setup called for %s', self) super(CombineOperation, self).setup(data_sampler) self.phased_combine_fn.combine_fn.setup() def process(self, o): # type: (WindowedValue) -> None with self.scoped_process_state: if self.debug_logging_enabled: _LOGGER.debug('Processing [%s] in %s', o, self) key, values = o.value self.output(o.with_value((key, self.phased_combine_fn.apply(values)))) def finish(self): # type: () -> None _LOGGER.debug('Finishing %s', self) super(CombineOperation, self).finish() def teardown(self): # type: () -> None with self.scoped_finish_state: _LOGGER.debug('Teardown called for %s', self) super(CombineOperation, self).teardown() self.phased_combine_fn.combine_fn.teardown() def create_pgbk_op(step_name, spec, counter_factory, state_sampler): if spec.combine_fn: return PGBKCVOperation(step_name, spec, counter_factory, state_sampler) else: return PGBKOperation(step_name, spec, counter_factory, state_sampler) class PGBKOperation(Operation): """Partial group-by-key operation. This takes (windowed) input (key, value) tuples and outputs (key, [value]) tuples, performing a best effort group-by-key for values in this bundle, memory permitting. """ def __init__(self, name_context, spec, counter_factory, state_sampler): super(PGBKOperation, self).__init__(name_context, spec, counter_factory, state_sampler) assert not self.spec.combine_fn self.table = collections.defaultdict(list) self.size = 0 # TODO(robertwb) Make this configurable. self.max_size = 10 * 1000 def process(self, o): # type: (WindowedValue) -> None with self.scoped_process_state: # TODO(robertwb): Structural (hashable) values. key = o.value[0], tuple(o.windows) self.table[key].append(o) self.size += 1 if self.size > self.max_size: self.flush(9 * self.max_size // 10) def finish(self): # type: () -> None self.flush(0) super().finish() def flush(self, target): # type: (int) -> None limit = self.size - target for ix, (kw, vs) in enumerate(list(self.table.items())): if ix >= limit: break del self.table[kw] key, windows = kw output_value = [v.value[1] for v in vs] windowed_value = WindowedValue((key, output_value), vs[0].timestamp, windows) self.output(windowed_value) class PGBKCVOperation(Operation): """Partial group-by-key operation. This operation handles grouped values with a combine function applied. """ def __init__( self, name_context, spec, counter_factory, state_sampler, windowing=None): super(PGBKCVOperation, self).__init__(name_context, spec, counter_factory, state_sampler) # Combiners do not accept deferred side-inputs (the ignored fourth # argument) and therefore the code to handle the extra args/kwargs is # simpler than for the DoFn's of ParDo. fn, args, kwargs = pickler.loads(self.spec.combine_fn)[:3] self.combine_fn = curry_combine_fn(fn, args, kwargs) self.combine_fn_add_input = self.combine_fn.add_input if self.combine_fn.compact.__func__ is core.CombineFn.compact: self.combine_fn_compact = None else: self.combine_fn_compact = self.combine_fn.compact if windowing: self.is_default_windowing = windowing.is_default() tsc_type = windowing.timestamp_combiner self.timestamp_combiner = ( None if tsc_type == window.TimestampCombiner.OUTPUT_AT_EOW else window.TimestampCombiner.get_impl(tsc_type, windowing.windowfn)) else: self.is_default_windowing = False # unknown self.timestamp_combiner = None # Optimization for the (known tiny accumulator, often wide keyspace) # combine functions. # TODO(b/36567833): Bound by in-memory size rather than key count. self.max_keys = ( 1000 * 1000 if isinstance(fn, (combiners.CountCombineFn, combiners.MeanCombineFn)) or # TODO(b/36597732): Replace this 'or' part by adding the 'cy' optimized # combiners to the short list above. ( isinstance(fn, core.CallableWrapperCombineFn) and fn._fn in (min, max, sum)) else 100 * 1000) # pylint: disable=protected-access self.key_count = 0 self.table = {} def setup(self, data_sampler=None): # type: (Optional[DataSampler]) -> None with self.scoped_start_state: _LOGGER.debug('Setup called for %s', self) super(PGBKCVOperation, self).setup(data_sampler) self.combine_fn.setup() def process(self, wkv): # type: (WindowedValue) -> None with self.scoped_process_state: key, value = wkv.value # pylint: disable=unidiomatic-typecheck # Optimization for the global window case. if self.is_default_windowing: self.add_key_value(key, value, None) else: for window in wkv.windows: self.add_key_value((window, key), value, wkv.timestamp if self.timestamp_combiner else None) def add_key_value(self, wkey, value, timestamp): entry = self.table.get(wkey, None) if entry is None: if self.key_count >= self.max_keys: target = self.key_count * 9 // 10 old_wkeys = [] # TODO(robertwb): Use an LRU cache? for old_wkey, old_wvalue in self.table.items(): old_wkeys.append(old_wkey) # Can't mutate while iterating. self.output_key(old_wkey, old_wvalue[0], old_wvalue[1]) self.key_count -= 1 if self.key_count <= target: break for old_wkey in reversed(old_wkeys): del self.table[old_wkey] self.key_count += 1 # We save the accumulator as a one element list so we can efficiently # mutate when new values are added without searching the cache again. entry = self.table[wkey] = [ self.combine_fn.create_accumulator(), timestamp ] entry[0] = self.combine_fn_add_input(entry[0], value) if not self.is_default_windowing and self.timestamp_combiner: entry[1] = self.timestamp_combiner.combine(entry[1], timestamp) def finish(self): # type: () -> None for wkey, value in self.table.items(): self.output_key(wkey, value[0], value[1]) self.table = {} self.key_count = 0 def teardown(self): # type: () -> None with self.scoped_finish_state: _LOGGER.debug('Teardown called for %s', self) super(PGBKCVOperation, self).teardown() self.combine_fn.teardown() def output_key(self, wkey, accumulator, timestamp): if self.combine_fn_compact is None: value = accumulator else: value = self.combine_fn_compact(accumulator) if self.is_default_windowing: self.output(_globally_windowed_value.with_value((wkey, value))) else: window, key = wkey if self.timestamp_combiner is None: timestamp = window.max_timestamp() self.output(WindowedValue((key, value), timestamp, (window, ))) class FlattenOperation(Operation): """Flatten operation. Receives one or more producer operations, outputs just one list with all the items. """ def process(self, o): # type: (WindowedValue) -> None with self.scoped_process_state: if self.debug_logging_enabled: _LOGGER.debug('Processing [%s] in %s', o, self) self.output(o) def create_operation( name_context, spec, counter_factory, step_name=None, state_sampler=None, test_shuffle_source=None, test_shuffle_sink=None, is_streaming=False): # type: (...) -> Operation """Create Operation object for given operation specification.""" # TODO(pabloem): Document arguments to this function call. if not isinstance(name_context, common.NameContext): name_context = common.NameContext(step_name=name_context) if isinstance(spec, operation_specs.WorkerRead): if isinstance(spec.source, iobase.SourceBundle): op = ReadOperation( name_context, spec, counter_factory, state_sampler) # type: Operation else: from dataflow_worker.native_operations import NativeReadOperation op = NativeReadOperation( name_context, spec, counter_factory, state_sampler) elif isinstance(spec, operation_specs.WorkerWrite): from dataflow_worker.native_operations import NativeWriteOperation op = NativeWriteOperation( name_context, spec, counter_factory, state_sampler) elif isinstance(spec, operation_specs.WorkerCombineFn): op = CombineOperation(name_context, spec, counter_factory, state_sampler) elif isinstance(spec, operation_specs.WorkerPartialGroupByKey): op = create_pgbk_op(name_context, spec, counter_factory, state_sampler) elif isinstance(spec, operation_specs.WorkerDoFn): op = DoOperation(name_context, spec, counter_factory, state_sampler) elif isinstance(spec, operation_specs.WorkerGroupingShuffleRead): from dataflow_worker.shuffle_operations import GroupedShuffleReadOperation op = GroupedShuffleReadOperation( name_context, spec, counter_factory, state_sampler, shuffle_source=test_shuffle_source) elif isinstance(spec, operation_specs.WorkerUngroupedShuffleRead): from dataflow_worker.shuffle_operations import UngroupedShuffleReadOperation op = UngroupedShuffleReadOperation( name_context, spec, counter_factory, state_sampler, shuffle_source=test_shuffle_source) elif isinstance(spec, operation_specs.WorkerInMemoryWrite): op = InMemoryWriteOperation( name_context, spec, counter_factory, state_sampler) elif isinstance(spec, operation_specs.WorkerShuffleWrite): from dataflow_worker.shuffle_operations import ShuffleWriteOperation op = ShuffleWriteOperation( name_context, spec, counter_factory, state_sampler, shuffle_sink=test_shuffle_sink) elif isinstance(spec, operation_specs.WorkerFlatten): op = FlattenOperation(name_context, spec, counter_factory, state_sampler) elif isinstance(spec, operation_specs.WorkerMergeWindows): from dataflow_worker.shuffle_operations import BatchGroupAlsoByWindowsOperation from dataflow_worker.shuffle_operations import StreamingGroupAlsoByWindowsOperation if is_streaming: op = StreamingGroupAlsoByWindowsOperation( name_context, spec, counter_factory, state_sampler) else: op = BatchGroupAlsoByWindowsOperation( name_context, spec, counter_factory, state_sampler) elif isinstance(spec, operation_specs.WorkerReifyTimestampAndWindows): from dataflow_worker.shuffle_operations import ReifyTimestampAndWindowsOperation op = ReifyTimestampAndWindowsOperation( name_context, spec, counter_factory, state_sampler) else: raise TypeError( 'Expected an instance of operation_specs.Worker* class ' 'instead of %s' % (spec, )) return op class SimpleMapTaskExecutor(object): """An executor for map tasks. Stores progress of the read operation that is the first operation of a map task. """ def __init__( self, map_task, counter_factory, state_sampler, test_shuffle_source=None, test_shuffle_sink=None): """Initializes SimpleMapTaskExecutor. Args: map_task: The map task we are to run. The maptask contains a list of operations, and aligned lists for step_names, original_names, system_names of pipeline steps. counter_factory: The CounterFactory instance for the work item. state_sampler: The StateSampler tracking the execution step. test_shuffle_source: Used during tests for dependency injection into shuffle read operation objects. test_shuffle_sink: Used during tests for dependency injection into shuffle write operation objects. """ self._map_task = map_task self._counter_factory = counter_factory self._ops = [] # type: List[Operation] self._state_sampler = state_sampler self._test_shuffle_source = test_shuffle_source self._test_shuffle_sink = test_shuffle_sink def operations(self): # type: () -> List[Operation] return self._ops[:] def execute(self): # type: () -> None """Executes all the operation_specs.Worker* instructions in a map task. We update the map_task with the execution status, expressed as counters. Raises: RuntimeError: if we find more than on read instruction in task spec. TypeError: if the spec parameter is not an instance of the recognized operation_specs.Worker* classes. """ # operations is a list of operation_specs.Worker* instances. # The order of the elements is important because the inputs use # list indexes as references. for name_context, spec in zip(self._map_task.name_contexts, self._map_task.operations): # This is used for logging and assigning names to counters. op = create_operation( name_context, spec, self._counter_factory, None, self._state_sampler, test_shuffle_source=self._test_shuffle_source, test_shuffle_sink=self._test_shuffle_sink) self._ops.append(op) # Add receiver operations to the appropriate producers. if hasattr(op.spec, 'input'): producer, output_index = op.spec.input self._ops[producer].add_receiver(op, output_index) # Flatten has 'inputs', not 'input' if hasattr(op.spec, 'inputs'): for producer, output_index in op.spec.inputs: self._ops[producer].add_receiver(op, output_index) for ix, op in reversed(list(enumerate(self._ops))): _LOGGER.debug('Starting op %d %s', ix, op) op.start() for op in self._ops: op.finish() class InefficientExecutionWarning(RuntimeWarning): """warning to indicate an inefficiency in a Beam pipeline.""" # Don't ignore InefficientExecutionWarning, but only log them once warnings.simplefilter('once', InefficientExecutionWarning) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/sdk_worker.py0000644000175100017510000015275615206113032023426 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """SDK harness for executing Python Fns via the Fn API.""" # pytype: skip-file # mypy: disallow-untyped-defs import abc import collections import contextlib import functools import json import logging import queue import sys import threading import time import traceback from concurrent import futures from typing import TYPE_CHECKING from typing import Any from typing import Callable from typing import DefaultDict from typing import Generic from typing import Iterable from typing import Iterator from typing import MutableMapping from typing import Optional from typing import TypeVar from typing import Union import grpc from apache_beam.coders import coder_impl from apache_beam.metrics import monitoring_infos from apache_beam.metrics.execution import MetricsEnvironment from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.portability.api import metrics_pb2 from apache_beam.runners.worker import bundle_processor from apache_beam.runners.worker import data_plane from apache_beam.runners.worker import statesampler from apache_beam.runners.worker.channel_factory import GRPCChannelFactory from apache_beam.runners.worker.data_plane import PeriodicThread from apache_beam.runners.worker.statecache import CacheAware from apache_beam.runners.worker.statecache import StateCache from apache_beam.runners.worker.worker_id_interceptor import WorkerIdInterceptor from apache_beam.runners.worker.worker_status import FnApiWorkerStatusHandler from apache_beam.utils import thread_pool_executor from apache_beam.utils.sentinel import Sentinel from apache_beam.version import __version__ as beam_version if TYPE_CHECKING: pass T = TypeVar('T') _KT = TypeVar('_KT') _VT = TypeVar('_VT') _LOGGER = logging.getLogger(__name__) DEFAULT_BUNDLE_PROCESSOR_CACHE_SHUTDOWN_THRESHOLD_S = 60 # The number of ProcessBundleRequest instruction ids the BundleProcessorCache # will remember for not running instructions. MAX_KNOWN_NOT_RUNNING_INSTRUCTIONS = 1000 # The number of ProcessBundleRequest instruction ids that BundleProcessorCache # will remember for failed instructions. MAX_FAILED_INSTRUCTIONS = 10000 # retry on transient UNAVAILABLE grpc error from state channels. _GRPC_SERVICE_CONFIG = json.dumps({ "methodConfig": [{ "name": [{ "service": "org.apache.beam.model.fn_execution.v1.BeamFnState" }], "retryPolicy": { "maxAttempts": 5, "initialBackoff": "0.1s", "maxBackoff": "5s", "backoffMultiplier": 2, "retryableStatusCodes": ["UNAVAILABLE"], }, }] }) class ShortIdCache(object): """ Cache for MonitoringInfo "short ids" """ def __init__(self): # type: () -> None self._lock = threading.Lock() self._last_short_id = 0 self._info_key_to_short_id = {} # type: Dict[FrozenSet, str] self._short_id_to_info = {} # type: Dict[str, metrics_pb2.MonitoringInfo] def get_short_id(self, monitoring_info): # type: (metrics_pb2.MonitoringInfo) -> str """ Returns the assigned shortId for a given MonitoringInfo, assigns one if not assigned already. """ key = monitoring_infos.to_key(monitoring_info) with self._lock: try: return self._info_key_to_short_id[key] except KeyError: self._last_short_id += 1 # Convert to a hex string (and drop the '0x') for some compression shortId = hex(self._last_short_id)[2:] payload_cleared = metrics_pb2.MonitoringInfo() payload_cleared.CopyFrom(monitoring_info) payload_cleared.ClearField('payload') self._info_key_to_short_id[key] = shortId self._short_id_to_info[shortId] = payload_cleared return shortId def get_infos(self, short_ids): #type: (Iterable[str]) -> Dict[str, metrics_pb2.MonitoringInfo] """ Gets the base MonitoringInfo (with payload cleared) for each short ID. Throws KeyError if an unassigned short ID is encountered. """ return { short_id: self._short_id_to_info[short_id] for short_id in short_ids } SHORT_ID_CACHE = ShortIdCache() class SdkHarness(object): REQUEST_METHOD_PREFIX = '_request_' def __init__( self, control_address, # type: str credentials=None, # type: Optional[grpc.ChannelCredentials] worker_id=None, # type: Optional[str] # Caching is disabled by default state_cache_size=0, # type: int # time-based data buffering is disabled by default data_buffer_time_limit_ms=0, # type: int profiler_factory=None, # type: Optional[Callable[..., Profile]] status_address=None, # type: Optional[str] # Heap dump through status api is disabled by default enable_heap_dump=False, # type: bool data_sampler=None, # type: Optional[data_sampler.DataSampler] # Unrecoverable SDK harness initialization error (if any) # that should be reported to the runner when proocessing the first bundle. deferred_exception=None, # type: Optional[Exception] runner_capabilities=frozenset(), # type: FrozenSet[str] element_processing_timeout_minutes=None, # type: Optional[int] ): # type: (...) -> None self._alive = True self._worker_index = 0 self._worker_id = worker_id self._state_cache = StateCache(state_cache_size) self._deferred_exception = deferred_exception options = [('grpc.max_receive_message_length', -1), ('grpc.max_send_message_length', -1)] if credentials is None: _LOGGER.info('Creating insecure control channel for %s.', control_address) self._control_channel = GRPCChannelFactory.insecure_channel( control_address, options=options) else: _LOGGER.info('Creating secure control channel for %s.', control_address) self._control_channel = GRPCChannelFactory.secure_channel( control_address, credentials, options=options) grpc.channel_ready_future(self._control_channel).result(timeout=60) _LOGGER.info('Control channel established.') self._control_channel = grpc.intercept_channel( self._control_channel, WorkerIdInterceptor(self._worker_id)) self._data_channel_factory = data_plane.GrpcClientDataChannelFactory( credentials, self._worker_id, data_buffer_time_limit_ms) self._state_handler_factory = GrpcStateHandlerFactory( state_cache=self._state_cache, credentials=credentials, worker_id=self._worker_id) self._profiler_factory = profiler_factory self.data_sampler = data_sampler self.runner_capabilities = runner_capabilities self._element_processing_timeout_minutes = ( element_processing_timeout_minutes) def default_factory(id): # type: (str) -> beam_fn_api_pb2.ProcessBundleDescriptor return self._control_stub.GetProcessBundleDescriptor( beam_fn_api_pb2.GetProcessBundleDescriptorRequest( process_bundle_descriptor_id=id)) self._fns = KeyedDefaultDict(default_factory) # BundleProcessor cache across all workers. self._bundle_processor_cache = BundleProcessorCache( self.runner_capabilities, state_handler_factory=self._state_handler_factory, data_channel_factory=self._data_channel_factory, fns=self._fns, data_sampler=self.data_sampler, ) self._status_handler = None # type: Optional[FnApiWorkerStatusHandler] if status_address: try: self._status_handler = FnApiWorkerStatusHandler( status_address, self._bundle_processor_cache, self._state_cache, enable_heap_dump, element_processing_timeout_minutes=self. _element_processing_timeout_minutes) except Exception: traceback_string = traceback.format_exc() _LOGGER.warning( 'Error creating worker status request handler, ' 'skipping status report. Trace back: %s' % traceback_string) # TODO(BEAM-8998) use common # thread_pool_executor.shared_unbounded_instance() to process bundle # progress once dataflow runner's excessive progress polling is removed. self._report_progress_executor = futures.ThreadPoolExecutor(max_workers=1) self._worker_thread_pool = thread_pool_executor.shared_unbounded_instance() self._responses = queue.Queue( ) # type: queue.Queue[Union[beam_fn_api_pb2.InstructionResponse, Sentinel]] _LOGGER.info('Initializing SDKHarness with unbounded number of workers.') def run(self): # type: () -> None self._control_stub = beam_fn_api_pb2_grpc.BeamFnControlStub( self._control_channel) no_more_work = Sentinel.sentinel def get_responses(): # type: () -> Iterator[beam_fn_api_pb2.InstructionResponse] while True: response = self._responses.get() if response is no_more_work: return yield response self._alive = True try: for work_request in self._control_stub.Control(get_responses()): _LOGGER.debug('Got work %s', work_request.instruction_id) request_type = work_request.WhichOneof('request') if request_type is None: raise RuntimeError( "Cannot interpret a request received over control channel. " "This is not expected. " "Verify that SDK was not accidentally downgraded at runtime. " f"SDK version: {beam_version}, " f"instruction id: {work_request.instruction_id}, " f"raw request: {str(work_request.SerializeToString())}") # Name spacing the request method with 'request_'. The called method # will be like self.request_register(request) getattr(self, SdkHarness.REQUEST_METHOD_PREFIX + request_type)( work_request) finally: self._alive = False if self.data_sampler: self.data_sampler.stop() _LOGGER.info('No more requests from control plane') _LOGGER.info('SDK Harness waiting for in-flight requests to complete') # Wait until existing requests are processed. self._worker_thread_pool.shutdown() # get_responses may be blocked on responses.get(), but we need to return # control to its caller. self._responses.put(no_more_work) # Stop all the workers and clean all the associated resources self._data_channel_factory.close() self._state_handler_factory.close() self._bundle_processor_cache.shutdown() if self._status_handler: self._status_handler.close() _LOGGER.info('Done consuming work.') def _execute( self, task, # type: Callable[[], beam_fn_api_pb2.InstructionResponse] request # type: beam_fn_api_pb2.InstructionRequest ): # type: (...) -> None with statesampler.instruction_id(request.instruction_id): try: response = task() except: # pylint: disable=bare-except traceback_string = traceback.format_exc() print(traceback_string, file=sys.stderr) _LOGGER.error( 'Error processing instruction %s. Original traceback is\n%s\n', request.instruction_id, traceback_string) response = beam_fn_api_pb2.InstructionResponse( instruction_id=request.instruction_id, error=traceback_string) self._responses.put(response) def _request_register(self, request): # type: (beam_fn_api_pb2.InstructionRequest) -> None # registration request is handled synchronously self._execute(lambda: self.create_worker().do_instruction(request), request) def _request_process_bundle(self, request): # type: (beam_fn_api_pb2.InstructionRequest) -> None if self._deferred_exception: raise self._deferred_exception self._bundle_processor_cache.activate(request.instruction_id) self._request_execute(request) def _request_process_bundle_split(self, request): # type: (beam_fn_api_pb2.InstructionRequest) -> None self._request_process_bundle_action(request) def _request_process_bundle_progress(self, request): # type: (beam_fn_api_pb2.InstructionRequest) -> None self._request_process_bundle_action(request) def _request_process_bundle_action(self, request): # type: (beam_fn_api_pb2.InstructionRequest) -> None def task(): # type: () -> None self._execute( lambda: self.create_worker().do_instruction(request), request) self._report_progress_executor.submit(task) def _request_finalize_bundle(self, request): # type: (beam_fn_api_pb2.InstructionRequest) -> None self._request_execute(request) def _request_harness_monitoring_infos(self, request): # type: (beam_fn_api_pb2.InstructionRequest) -> None process_wide_monitoring_infos = MetricsEnvironment.process_wide_container( ).to_runner_api_monitoring_infos(None).values() self._execute( lambda: beam_fn_api_pb2.InstructionResponse( instruction_id=request.instruction_id, harness_monitoring_infos=( beam_fn_api_pb2.HarnessMonitoringInfosResponse( monitoring_data={ SHORT_ID_CACHE.get_short_id(info): info.payload for info in process_wide_monitoring_infos }))), request) def _request_monitoring_infos(self, request): # type: (beam_fn_api_pb2.InstructionRequest) -> None self._execute( lambda: beam_fn_api_pb2.InstructionResponse( instruction_id=request.instruction_id, monitoring_infos= beam_fn_api_pb2.MonitoringInfosMetadataResponse( monitoring_info=SHORT_ID_CACHE.get_infos( request.monitoring_infos.monitoring_info_id))), request) def _request_execute(self, request): # type: (beam_fn_api_pb2.InstructionRequest) -> None def task(): # type: () -> None self._execute( lambda: self.create_worker().do_instruction(request), request) self._worker_thread_pool.submit(task) _LOGGER.debug( "Currently using %s threads." % len(self._worker_thread_pool._workers)) def _request_sample_data(self, request): # type: (beam_fn_api_pb2.InstructionRequest) -> None def get_samples(request): # type: (beam_fn_api_pb2.InstructionRequest) -> beam_fn_api_pb2.InstructionResponse samples = beam_fn_api_pb2.SampleDataResponse() if self.data_sampler is not None: samples = self.data_sampler.samples(request.sample_data.pcollection_ids) return beam_fn_api_pb2.InstructionResponse( instruction_id=request.instruction_id, sample_data=samples) self._execute(lambda: get_samples(request), request) def create_worker(self): # type: () -> SdkWorker return SdkWorker( self._bundle_processor_cache, profiler_factory=self._profiler_factory) class BundleProcessorCache(object): """A cache for ``BundleProcessor``s. ``BundleProcessor`` objects are cached by the id of their ``beam_fn_api_pb2.ProcessBundleDescriptor``. Attributes: fns (dict): A dictionary that maps bundle descriptor IDs to instances of ``beam_fn_api_pb2.ProcessBundleDescriptor``. state_handler_factory (``StateHandlerFactory``): Used to create state handlers to be used by a ``bundle_processor.BundleProcessor`` during processing. data_channel_factory (``data_plane.DataChannelFactory``) active_bundle_processors (dict): A dictionary, indexed by instruction IDs, containing ``bundle_processor.BundleProcessor`` objects that are currently active processing the corresponding instruction. cached_bundle_processors (dict): A dictionary, indexed by bundle processor id, of cached ``bundle_processor.BundleProcessor`` that are not currently performing processing. """ periodic_shutdown = None # type: Optional[PeriodicThread] def __init__( self, runner_capabilities, # type: FrozenSet[str] state_handler_factory, # type: StateHandlerFactory data_channel_factory, # type: data_plane.DataChannelFactory fns, # type: MutableMapping[str, beam_fn_api_pb2.ProcessBundleDescriptor] data_sampler=None, # type: Optional[data_sampler.DataSampler] ): # type: (...) -> None self.runner_capabilities = runner_capabilities self.fns = fns self.state_handler_factory = state_handler_factory self.data_channel_factory = data_channel_factory self.known_not_running_instruction_ids = collections.OrderedDict( ) # type: collections.OrderedDict[str, bool] self.failed_instruction_ids = collections.OrderedDict( ) # type: collections.OrderedDict[str, Exception] self.active_bundle_processors = { } # type: Dict[str, Tuple[str, bundle_processor.BundleProcessor]] self.processors_being_created = { } # type: Dict[str, Tuple[str, threading.Thread, float]] self.cached_bundle_processors = collections.defaultdict( list) # type: DefaultDict[str, List[bundle_processor.BundleProcessor]] self.last_access_times = collections.defaultdict( float) # type: DefaultDict[str, float] self._schedule_periodic_shutdown() self._lock = threading.Lock() self.data_sampler = data_sampler def register(self, bundle_descriptor): # type: (beam_fn_api_pb2.ProcessBundleDescriptor) -> None """Register a ``beam_fn_api_pb2.ProcessBundleDescriptor`` by its id.""" self.fns[bundle_descriptor.id] = bundle_descriptor def activate(self, instruction_id): # type: (str) -> None """Makes the ``instruction_id`` known to the bundle processor. Allows ``lookup`` to return ``None``. Necessary if ``lookup`` can occur before ``get``. """ with self._lock: self.known_not_running_instruction_ids[instruction_id] = True def get(self, instruction_id, bundle_descriptor_id): # type: (str, str) -> bundle_processor.BundleProcessor """ Return the requested ``BundleProcessor``, creating it if necessary. Moves the ``BundleProcessor`` from the inactive to the active cache. """ with self._lock: try: # pop() is threadsafe processor = self.cached_bundle_processors[bundle_descriptor_id].pop() self.active_bundle_processors[ instruction_id] = bundle_descriptor_id, processor try: del self.known_not_running_instruction_ids[instruction_id] except KeyError: # The instruction may have not been pre-registered before execution # since activate() may have never been invoked pass return processor except IndexError: self.processors_being_created[instruction_id] = ( bundle_descriptor_id, threading.current_thread(), time.time()) # Make sure we instantiate the processor while not holding the lock. # Reduce risks of concurrent modifications of the same protos # captured in bundle descriptor when the same bundle descriptor is used # in different instructions. pbd = beam_fn_api_pb2.ProcessBundleDescriptor() pbd.MergeFrom(self.fns[bundle_descriptor_id]) processor = bundle_processor.BundleProcessor( self.runner_capabilities, pbd, self.state_handler_factory.create_state_handler( pbd.state_api_service_descriptor), self.data_channel_factory, self.data_sampler) with self._lock: self.active_bundle_processors[ instruction_id] = bundle_descriptor_id, processor del self.processors_being_created[instruction_id] try: del self.known_not_running_instruction_ids[instruction_id] except KeyError: # The instruction may have not been pre-registered before execution # since activate() may have never been invoked pass return processor def lookup(self, instruction_id): # type: (str) -> Optional[bundle_processor.BundleProcessor] """ Return the requested ``BundleProcessor`` from the cache. Will return ``None`` if the BundleProcessor is known but not yet ready. Will raise an error if the ``instruction_id`` is not known or has been discarded. """ with self._lock: if instruction_id in self.failed_instruction_ids: e = self.failed_instruction_ids[instruction_id] raise RuntimeError( 'Bundle processing associated with %s has failed. ' 'Check prior failing response and attached exception for details.' % instruction_id) from e processor = self.active_bundle_processors.get( instruction_id, (None, None))[-1] if processor: return processor if instruction_id in self.known_not_running_instruction_ids: return None raise RuntimeError('Unknown process bundle id %s.' % instruction_id) def discard(self, instruction_id, exception): # type: (str, Exception) -> None """ Marks the instruction id as failed shutting down the ``BundleProcessor``. """ processor = None with self._lock: self.failed_instruction_ids[instruction_id] = exception while len(self.failed_instruction_ids) > MAX_FAILED_INSTRUCTIONS: self.failed_instruction_ids.popitem(last=False) if instruction_id in self.active_bundle_processors: processor = self.active_bundle_processors.pop(instruction_id)[1] # Perform the shutdown while not holding the lock. if processor: processor.shutdown() self.data_channel_factory.cleanup(instruction_id) def release(self, instruction_id): # type: (str) -> None """ Release the requested ``BundleProcessor``. Resets the ``BundleProcessor`` and moves it from the active to the inactive cache. """ with self._lock: self.known_not_running_instruction_ids[instruction_id] = True while len(self.known_not_running_instruction_ids ) > MAX_KNOWN_NOT_RUNNING_INSTRUCTIONS: self.known_not_running_instruction_ids.popitem(last=False) descriptor_id, processor = ( self.active_bundle_processors.pop(instruction_id)) # Make sure that we reset the processor while not holding the lock. processor.reset() with self._lock: self.last_access_times[descriptor_id] = time.time() self.cached_bundle_processors[descriptor_id].append(processor) def shutdown(self): # type: () -> None """ Shutdown all ``BundleProcessor``s in the cache. """ if self.periodic_shutdown: self.periodic_shutdown.cancel() self.periodic_shutdown.join() self.periodic_shutdown = None for instruction_id in list(self.active_bundle_processors.keys()): self.discard(instruction_id, RuntimeError('Shutdown invoked')) for cached_bundle_processors in self.cached_bundle_processors.values(): BundleProcessorCache._shutdown_cached_bundle_processors( cached_bundle_processors) def _schedule_periodic_shutdown(self): # type: () -> None def shutdown_inactive_bundle_processors(): # type: () -> None inactive_descriptor_ids = [] inactive_time = time.time( ) - DEFAULT_BUNDLE_PROCESSOR_CACHE_SHUTDOWN_THRESHOLD_S with self._lock: for descriptor_id, last_access_time in self.last_access_times.items(): if (inactive_time > last_access_time): inactive_descriptor_ids.append(descriptor_id) # Shutdown can be expensive, keep out of lock for descriptor_id in inactive_descriptor_ids: BundleProcessorCache._shutdown_cached_bundle_processors( self.cached_bundle_processors[descriptor_id]) self.periodic_shutdown = PeriodicThread( DEFAULT_BUNDLE_PROCESSOR_CACHE_SHUTDOWN_THRESHOLD_S, shutdown_inactive_bundle_processors) self.periodic_shutdown.daemon = True self.periodic_shutdown.start() @staticmethod def _shutdown_cached_bundle_processors(cached_bundle_processors): # type: (List[bundle_processor.BundleProcessor]) -> None try: while True: # pop() is threadsafe bundle_processor = cached_bundle_processors.pop() bundle_processor.shutdown() except IndexError: pass class SdkWorker(object): def __init__( self, bundle_processor_cache, # type: BundleProcessorCache profiler_factory=None, # type: Optional[Callable[..., Profile]] ): # type: (...) -> None self.bundle_processor_cache = bundle_processor_cache self.profiler_factory = profiler_factory def do_instruction(self, request): # type: (beam_fn_api_pb2.InstructionRequest) -> beam_fn_api_pb2.InstructionResponse request_type = request.WhichOneof('request') if request_type: # E.g. if register is set, this will call self.register(request.register)) return getattr(self, request_type)( getattr(request, request_type), request.instruction_id) else: raise NotImplementedError def register( self, request, # type: beam_fn_api_pb2.RegisterRequest instruction_id # type: str ): # type: (...) -> beam_fn_api_pb2.InstructionResponse """Registers a set of ``beam_fn_api_pb2.ProcessBundleDescriptor``s. This set of ``beam_fn_api_pb2.ProcessBundleDescriptor`` come as part of a ``beam_fn_api_pb2.RegisterRequest``, which the runner sends to the SDK worker before starting processing to register stages. """ for process_bundle_descriptor in request.process_bundle_descriptor: self.bundle_processor_cache.register(process_bundle_descriptor) return beam_fn_api_pb2.InstructionResponse( instruction_id=instruction_id, register=beam_fn_api_pb2.RegisterResponse()) def process_bundle( self, request, # type: beam_fn_api_pb2.ProcessBundleRequest instruction_id # type: str ): # type: (...) -> beam_fn_api_pb2.InstructionResponse try: bundle_processor = self.bundle_processor_cache.get( instruction_id, request.process_bundle_descriptor_id) with bundle_processor.state_handler.process_instruction_id( instruction_id, request.cache_tokens): with self.maybe_profile(instruction_id): delayed_applications, requests_finalization = ( bundle_processor.process_bundle(instruction_id)) monitoring_infos = bundle_processor.monitoring_infos() response = beam_fn_api_pb2.InstructionResponse( instruction_id=instruction_id, process_bundle=beam_fn_api_pb2.ProcessBundleResponse( residual_roots=delayed_applications, monitoring_infos=monitoring_infos, monitoring_data={ SHORT_ID_CACHE.get_short_id(info): info.payload for info in monitoring_infos }, requires_finalization=requests_finalization)) # Don't release here if finalize is needed. if not requests_finalization: self.bundle_processor_cache.release(instruction_id) return response except Exception as e: # pylint: disable=bare-except # Don't re-use bundle processors on failure. self.bundle_processor_cache.discard(instruction_id, e) raise def process_bundle_split( self, request, # type: beam_fn_api_pb2.ProcessBundleSplitRequest instruction_id # type: str ): # type: (...) -> beam_fn_api_pb2.InstructionResponse try: processor = self.bundle_processor_cache.lookup(request.instruction_id) except RuntimeError: return beam_fn_api_pb2.InstructionResponse( instruction_id=instruction_id, error=traceback.format_exc()) # Return an empty response if we aren't running. This can happen # if the ProcessBundleRequest has not started or already finished. process_bundle_split = ( processor.try_split(request) if processor else beam_fn_api_pb2.ProcessBundleSplitResponse()) return beam_fn_api_pb2.InstructionResponse( instruction_id=instruction_id, process_bundle_split=process_bundle_split) def process_bundle_progress( self, request, # type: beam_fn_api_pb2.ProcessBundleProgressRequest instruction_id # type: str ): # type: (...) -> beam_fn_api_pb2.InstructionResponse try: processor = self.bundle_processor_cache.lookup(request.instruction_id) except RuntimeError: return beam_fn_api_pb2.InstructionResponse( instruction_id=instruction_id, error=traceback.format_exc()) if processor: monitoring_infos = processor.monitoring_infos() consuming_received_data = processor.consuming_received_data else: # Return an empty response if we aren't running. This can happen # if the ProcessBundleRequest has not started or already finished. monitoring_infos = [] consuming_received_data = False return beam_fn_api_pb2.InstructionResponse( instruction_id=instruction_id, process_bundle_progress=beam_fn_api_pb2.ProcessBundleProgressResponse( monitoring_infos=monitoring_infos, monitoring_data={ SHORT_ID_CACHE.get_short_id(info): info.payload for info in monitoring_infos }, consuming_received_data=consuming_received_data)) def finalize_bundle( self, request, # type: beam_fn_api_pb2.FinalizeBundleRequest instruction_id # type: str ): # type: (...) -> beam_fn_api_pb2.InstructionResponse try: processor = self.bundle_processor_cache.lookup(request.instruction_id) except RuntimeError: return beam_fn_api_pb2.InstructionResponse( instruction_id=instruction_id, error=traceback.format_exc()) if processor: try: finalize_response = processor.finalize_bundle() self.bundle_processor_cache.release(request.instruction_id) return beam_fn_api_pb2.InstructionResponse( instruction_id=instruction_id, finalize_bundle=finalize_response) except Exception as e: self.bundle_processor_cache.discard(request.instruction_id, e) raise # We can reach this state if there was an erroneous request to finalize # the bundle while it is being initialized or has already been finalized # and released. raise RuntimeError( 'Bundle is not in a finalizable state for %s' % instruction_id) @contextlib.contextmanager def maybe_profile(self, instruction_id): # type: (str) -> Iterator[None] if self.profiler_factory: profiler = self.profiler_factory(instruction_id) if profiler: with profiler: yield else: yield else: yield class StateHandler(metaclass=abc.ABCMeta): """An abstract object representing a ``StateHandler``.""" @abc.abstractmethod def get_raw( self, state_key, # type: beam_fn_api_pb2.StateKey continuation_token=None # type: Optional[bytes] ): # type: (...) -> Tuple[bytes, Optional[bytes]] """Gets the contents of state for the given state key. State is associated to a state key, AND an instruction_id, which is set when calling process_instruction_id. Returns a tuple with the contents in state, and an optional continuation token, which is used to page the API. """ raise NotImplementedError(type(self)) @abc.abstractmethod def append_raw( self, state_key, # type: beam_fn_api_pb2.StateKey data # type: bytes ): # type: (...) -> _Future """Append the input data into the state key. Returns a future that allows one to wait for the completion of the call. State is associated to a state key, AND an instruction_id, which is set when calling process_instruction_id. """ raise NotImplementedError(type(self)) @abc.abstractmethod def clear(self, state_key): # type: (beam_fn_api_pb2.StateKey) -> _Future """Clears the contents of a cell for the input state key. Returns a future that allows one to wait for the completion of the call. State is associated to a state key, AND an instruction_id, which is set when calling process_instruction_id. """ raise NotImplementedError(type(self)) @abc.abstractmethod @contextlib.contextmanager def process_instruction_id(self, bundle_id): # type: (str) -> Iterator[None] """Switch the context of the state handler to a specific instruction. This must be called before performing any write or read operations on the existing state. """ raise NotImplementedError(type(self)) @abc.abstractmethod def done(self): # type: () -> None """Mark the state handler as done, and potentially delete all context.""" raise NotImplementedError(type(self)) class StateHandlerFactory(metaclass=abc.ABCMeta): """An abstract factory for creating ``DataChannel``.""" @abc.abstractmethod def create_state_handler(self, api_service_descriptor): # type: (endpoints_pb2.ApiServiceDescriptor) -> CachingStateHandler """Returns a ``StateHandler`` from the given ApiServiceDescriptor.""" raise NotImplementedError(type(self)) @abc.abstractmethod def close(self): # type: () -> None """Close all channels that this factory owns.""" raise NotImplementedError(type(self)) class GrpcStateHandlerFactory(StateHandlerFactory): """A factory for ``GrpcStateHandler``. Caches the created channels by ``state descriptor url``. """ def __init__(self, state_cache, credentials=None, worker_id=None): # type: (StateCache, Optional[grpc.ChannelCredentials], Optional[str]) -> None self._state_handler_cache = {} # type: Dict[str, CachingStateHandler] self._lock = threading.Lock() self._throwing_state_handler = ThrowingStateHandler() self._credentials = credentials self._state_cache = state_cache self._worker_id = worker_id def create_state_handler(self, api_service_descriptor): # type: (endpoints_pb2.ApiServiceDescriptor) -> CachingStateHandler if not api_service_descriptor: return self._throwing_state_handler url = api_service_descriptor.url if url not in self._state_handler_cache: with self._lock: if url not in self._state_handler_cache: # Options to have no limits (-1) on the size of the messages # received or sent over the data plane. The actual buffer size is # controlled in a layer above. options = [('grpc.max_receive_message_length', -1), ('grpc.max_send_message_length', -1), ('grpc.service_config', _GRPC_SERVICE_CONFIG)] if self._credentials is None: _LOGGER.info('Creating insecure state channel for %s.', url) grpc_channel = GRPCChannelFactory.insecure_channel( url, options=options) else: _LOGGER.info('Creating secure state channel for %s.', url) grpc_channel = GRPCChannelFactory.secure_channel( url, self._credentials, options=options) _LOGGER.info('State channel established.') # Add workerId to the grpc channel grpc_channel = grpc.intercept_channel( grpc_channel, WorkerIdInterceptor(self._worker_id)) self._state_handler_cache[url] = GlobalCachingStateHandler( self._state_cache, GrpcStateHandler( beam_fn_api_pb2_grpc.BeamFnStateStub(grpc_channel))) return self._state_handler_cache[url] def close(self): # type: () -> None _LOGGER.info('Closing all cached gRPC state handlers.') for _, state_handler in self._state_handler_cache.items(): state_handler.done() self._state_handler_cache.clear() self._state_cache.invalidate_all() class CachingStateHandler(metaclass=abc.ABCMeta): @abc.abstractmethod @contextlib.contextmanager def process_instruction_id(self, bundle_id, cache_tokens): # type: (str, Iterable[beam_fn_api_pb2.ProcessBundleRequest.CacheToken]) -> Iterator[None] raise NotImplementedError(type(self)) @abc.abstractmethod def blocking_get( self, state_key, # type: beam_fn_api_pb2.StateKey coder, # type: coder_impl.CoderImpl ): # type: (...) -> Iterable[Any] raise NotImplementedError(type(self)) @abc.abstractmethod def extend( self, state_key, # type: beam_fn_api_pb2.StateKey coder, # type: coder_impl.CoderImpl elements, # type: Iterable[Any] ): # type: (...) -> _Future raise NotImplementedError(type(self)) @abc.abstractmethod def clear(self, state_key): # type: (beam_fn_api_pb2.StateKey) -> _Future raise NotImplementedError(type(self)) @abc.abstractmethod def done(self): # type: () -> None raise NotImplementedError(type(self)) class ThrowingStateHandler(CachingStateHandler): """A caching state handler that errors on any requests.""" @contextlib.contextmanager def process_instruction_id(self, bundle_id, cache_tokens): # type: (str, Iterable[beam_fn_api_pb2.ProcessBundleRequest.CacheToken]) -> Iterator[None] raise RuntimeError( 'Unable to handle state requests for ProcessBundleDescriptor ' 'for bundle id %s.' % bundle_id) def blocking_get( self, state_key, # type: beam_fn_api_pb2.StateKey coder, # type: coder_impl.CoderImpl ): # type: (...) -> Iterable[Any] raise RuntimeError( 'Unable to handle state requests for ProcessBundleDescriptor without ' 'state ApiServiceDescriptor for state key %s.' % state_key) def extend( self, state_key, # type: beam_fn_api_pb2.StateKey coder, # type: coder_impl.CoderImpl elements, # type: Iterable[Any] ): # type: (...) -> _Future raise RuntimeError( 'Unable to handle state requests for ProcessBundleDescriptor without ' 'state ApiServiceDescriptor for state key %s.' % state_key) def clear(self, state_key): # type: (beam_fn_api_pb2.StateKey) -> _Future raise RuntimeError( 'Unable to handle state requests for ProcessBundleDescriptor without ' 'state ApiServiceDescriptor for state key %s.' % state_key) def done(self): # type: () -> None raise RuntimeError( 'Unable to handle state requests for ProcessBundleDescriptor.') class GrpcStateHandler(StateHandler): _DONE = Sentinel.sentinel def __init__(self, state_stub): # type: (beam_fn_api_pb2_grpc.BeamFnStateStub) -> None self._lock = threading.Lock() self._state_stub = state_stub self._requests = queue.Queue( ) # type: queue.Queue[Union[beam_fn_api_pb2.StateRequest, Sentinel]] self._responses_by_id = {} # type: Dict[str, _Future] self._last_id = 0 self._exception = None # type: Optional[Exception] self._context = threading.local() self.start() @contextlib.contextmanager def process_instruction_id(self, bundle_id): # type: (str) -> Iterator[None] if getattr(self._context, 'process_instruction_id', None) is not None: raise RuntimeError( 'Already bound to %r' % self._context.process_instruction_id) self._context.process_instruction_id = bundle_id try: yield finally: self._context.process_instruction_id = None def start(self): # type: () -> None self._done = False def request_iter(): # type: () -> Iterator[beam_fn_api_pb2.StateRequest] while True: request = self._requests.get() if request is self._DONE or self._done: break yield request responses = self._state_stub.State(request_iter()) def pull_responses(): # type: () -> None try: for response in responses: # Popping an item from a dictionary is atomic in cPython future = self._responses_by_id.pop(response.id) future.set(response) if self._done: break except Exception as e: self._exception = e raise reader = threading.Thread(target=pull_responses, name='read_state') reader.daemon = True reader.start() def done(self): # type: () -> None self._done = True self._requests.put(self._DONE) def get_raw( self, state_key, # type: beam_fn_api_pb2.StateKey continuation_token=None # type: Optional[bytes] ): # type: (...) -> Tuple[bytes, Optional[bytes]] response = self._blocking_request( beam_fn_api_pb2.StateRequest( state_key=state_key, get=beam_fn_api_pb2.StateGetRequest( continuation_token=continuation_token))) return response.get.data, response.get.continuation_token def append_raw( self, state_key, # type: Optional[beam_fn_api_pb2.StateKey] data # type: bytes ): # type: (...) -> _Future return self._request( beam_fn_api_pb2.StateRequest( state_key=state_key, append=beam_fn_api_pb2.StateAppendRequest(data=data))) def clear(self, state_key): # type: (Optional[beam_fn_api_pb2.StateKey]) -> _Future return self._request( beam_fn_api_pb2.StateRequest( state_key=state_key, clear=beam_fn_api_pb2.StateClearRequest())) def _request(self, request): # type: (beam_fn_api_pb2.StateRequest) -> _Future[beam_fn_api_pb2.StateResponse] request.id = self._next_id() request.instruction_id = self._context.process_instruction_id # Adding a new item to a dictionary is atomic in cPython self._responses_by_id[request.id] = future = _Future[ beam_fn_api_pb2.StateResponse]() # Request queue is thread-safe self._requests.put(request) return future def _blocking_request(self, request): # type: (beam_fn_api_pb2.StateRequest) -> beam_fn_api_pb2.StateResponse req_future = self._request(request) while not req_future.wait(timeout=1): if self._exception: raise self._exception elif self._done: raise RuntimeError() response = req_future.get() if response.error: raise RuntimeError(response.error) else: return response def _next_id(self): # type: () -> str with self._lock: # Use a lock here because this GrpcStateHandler is shared across all # requests which have the same process bundle descriptor. State requests # can concurrently access this section if a Runner uses threads / workers # (aka "parallelism") to send data to this SdkHarness and its workers. self._last_id += 1 request_id = self._last_id return str(request_id) class GlobalCachingStateHandler(CachingStateHandler): """ A State handler which retrieves and caches state. If caching is activated, caches across bundles using a supplied cache token. If activated but no cache token is supplied, caching is done at the bundle level. """ def __init__( self, global_state_cache, # type: StateCache underlying_state # type: StateHandler ): # type: (...) -> None self._underlying = underlying_state self._state_cache = global_state_cache self._context = threading.local() # state retrieval time statistics self._retrieval_time = 0.0 self._get_raw_called = 0 self._warn_interval = 60.0 @contextlib.contextmanager def process_instruction_id(self, bundle_id, cache_tokens): # type: (str, Iterable[beam_fn_api_pb2.ProcessBundleRequest.CacheToken]) -> Iterator[None] if getattr(self._context, 'user_state_cache_token', None) is not None: raise RuntimeError( 'Cache tokens already set to %s' % self._context.user_state_cache_token) self._context.side_input_cache_tokens = {} user_state_cache_token = None for cache_token_struct in cache_tokens: if cache_token_struct.HasField("user_state"): # There should only be one user state token present assert not user_state_cache_token user_state_cache_token = cache_token_struct.token elif cache_token_struct.HasField("side_input"): self._context.side_input_cache_tokens[ cache_token_struct.side_input.transform_id, cache_token_struct.side_input. side_input_id] = cache_token_struct.token # TODO: Consider a two-level cache to avoid extra logic and locking # for items cached at the bundle level. self._context.bundle_cache_token = bundle_id try: self._context.user_state_cache_token = user_state_cache_token with self._underlying.process_instruction_id(bundle_id): yield finally: self._context.side_input_cache_tokens = {} self._context.user_state_cache_token = None self._context.bundle_cache_token = None def blocking_get( self, state_key, # type: beam_fn_api_pb2.StateKey coder, # type: coder_impl.CoderImpl ): # type: (...) -> Iterable[Any] cache_token = self._get_cache_token(state_key) if not cache_token: # Cache disabled / no cache token. Can't do a lookup/store in the cache. # Fall back to lazily materializing the state, one element at a time. return self._lazy_iterator(state_key, coder) # Cache lookup cache_state_key = self._convert_to_cache_key(state_key) return self._state_cache.get( (cache_state_key, cache_token), lambda key: self._partially_cached_iterable(state_key, coder)) def extend( self, state_key, # type: beam_fn_api_pb2.StateKey coder, # type: coder_impl.CoderImpl elements, # type: Iterable[Any] ): # type: (...) -> _Future cache_token = self._get_cache_token(state_key) if cache_token: # Update the cache if the value is already present and # can be updated. cache_key = self._convert_to_cache_key(state_key) cached_value = self._state_cache.peek((cache_key, cache_token)) if isinstance(cached_value, list): # The state is fully cached and can be extended # Materialize provided iterable to ensure reproducible iterations, # here and when writing to the state handler below. elements = list(elements) cached_value.extend(elements) # Re-insert into the cache the updated value so the updated size is # reflected. self._state_cache.put((cache_key, cache_token), cached_value) # Write to state handler futures = [] out = coder_impl.create_OutputStream() for element in elements: coder.encode_to_stream(element, out, True) if out.size() > data_plane._DEFAULT_SIZE_FLUSH_THRESHOLD: futures.append(self._underlying.append_raw(state_key, out.get())) out = coder_impl.create_OutputStream() if out.size(): futures.append(self._underlying.append_raw(state_key, out.get())) return _DeferredCall( lambda *results: beam_fn_api_pb2.StateResponse( error='\n'.join( result.error for result in results if result and result.error), append=beam_fn_api_pb2.StateAppendResponse()), *futures) def clear(self, state_key): # type: (beam_fn_api_pb2.StateKey) -> _Future cache_token = self._get_cache_token(state_key) if cache_token: cache_key = self._convert_to_cache_key(state_key) self._state_cache.put((cache_key, cache_token), []) return self._underlying.clear(state_key) def done(self): # type: () -> None self._underlying.done() def _lazy_iterator( self, state_key, # type: beam_fn_api_pb2.StateKey coder, # type: coder_impl.CoderImpl continuation_token=None # type: Optional[bytes] ): # type: (...) -> Iterator[Any] """Materializes the state lazily, one element at a time. :return A generator which returns the next element if advanced. """ while True: input_stream, continuation_token = self._get_raw( state_key, continuation_token) while input_stream.size() > 0: yield coder.decode_from_stream(input_stream, True) if not continuation_token: break def _get_raw( self, state_key, # type: beam_fn_api_pb2.StateKey continuation_token # type: Optional[bytes] ): # type: (...) -> Tuple[coder_impl.create_InputStream, Optional[bytes]] """Call underlying get_raw with performance statistics and detection.""" start_time = time.time() data, continuation_token = ( self._underlying.get_raw(state_key, continuation_token)) input_stream = coder_impl.create_InputStream(data) self._retrieval_time += time.time() - start_time self._get_raw_called += 1 if self._retrieval_time > self._warn_interval: _LOGGER.warning( "Retrieving state %d times costed %.0f seconds. It may be due to " "insufficient state cache size and/or frequent direct access of " "states.\nConsider adding '--max_cache_memory_usage_mb' pipeline " "option to increase state cache size or switch to materialized " "(pvalue.AsList) side input if applicable." % (self._get_raw_called, self._retrieval_time)) # reset counts self._retrieval_time = 0.0 self._get_raw_called = 0 self._warn_interval *= 2 return input_stream, continuation_token def _get_cache_token(self, state_key): # type: (beam_fn_api_pb2.StateKey) -> Optional[bytes] if not self._state_cache.is_cache_enabled(): return None elif state_key.HasField('bag_user_state'): if self._context.user_state_cache_token: return self._context.user_state_cache_token else: return self._context.bundle_cache_token elif state_key.WhichOneof('type').endswith('_side_input'): # type: ignore[union-attr] side_input = getattr(state_key, state_key.WhichOneof('type')) # type: ignore[arg-type] return self._context.side_input_cache_tokens.get( (side_input.transform_id, side_input.side_input_id), self._context.bundle_cache_token) return None def _partially_cached_iterable( self, state_key, # type: beam_fn_api_pb2.StateKey coder # type: coder_impl.CoderImpl ): # type: (...) -> Iterable[Any] """Materialized the first page of data, concatenated with a lazy iterable of the rest, if any. """ input_stream, continuation_token = self._get_raw(state_key, None) head = [] while input_stream.size() > 0: head.append(coder.decode_from_stream(input_stream, True)) if not continuation_token: return head else: return self.ContinuationIterable( head, functools.partial( self._lazy_iterator, state_key, coder, continuation_token)) class ContinuationIterable(Generic[T], CacheAware): def __init__(self, head, continue_iterator_fn): # type: (Iterable[T], Callable[[], Iterable[T]]) -> None self.head = head self.continue_iterator_fn = continue_iterator_fn def __iter__(self): # type: () -> Iterator[T] for item in self.head: yield item for item in self.continue_iterator_fn(): yield item def get_referents_for_cache(self): # type: () -> List[Any] # Only capture the size of the elements and not the # continuation iterator since it references objects # we don't want to include in the cache measurement. return [self.head] @staticmethod def _convert_to_cache_key(state_key): # type: (beam_fn_api_pb2.StateKey) -> bytes return state_key.SerializeToString() class _Future(Generic[T]): """A simple future object to implement blocking requests. """ def __init__(self): # type: () -> None self._event = threading.Event() def wait(self, timeout=None): # type: (Optional[float]) -> bool return self._event.wait(timeout) def get(self, timeout=None): # type: (Optional[float]) -> T if self.wait(timeout): return self._value else: raise LookupError() def set(self, value): # type: (T) -> _Future[T] self._value = value self._event.set() return self @classmethod def done(cls): # type: () -> _Future[None] if not hasattr(cls, 'DONE'): done_future = _Future[None]() done_future.set(None) cls.DONE = done_future # type: ignore[attr-defined] return cls.DONE # type: ignore[attr-defined] class _DeferredCall(_Future[T]): def __init__(self, func, *args): # type: (Callable[..., Any], *Any) -> None self._func = func self._args = [ arg if isinstance(arg, _Future) else _Future().set(arg) for arg in args ] def wait(self, timeout=None): # type: (Optional[float]) -> bool return all(arg.wait(timeout) for arg in self._args) def get(self, timeout=None): # type: (Optional[float]) -> T # List comprehension, not generator: *(gen) causes CPython to build the # argument tuple incrementally via _PyTuple_Resize, which asserts # Py_REFCNT(v)==1. A GC cycle between yields can increment that refcount, # raising SystemError (Objects/tupleobject.c:927). See # https://github.com/python/cpython/issues/127058 (fixed in 3.14.0a3+: # https://github.com/python/cpython/commit/5a23994). *[list] allocates the # tuple once at its final size, avoiding the resize entirely. return self._func(*[arg.get(timeout) for arg in self._args]) def set(self, value): # type: (T) -> _Future[T] raise NotImplementedError() class KeyedDefaultDict(collections.defaultdict[_KT, _VT]): if TYPE_CHECKING: # we promise to only use a subset of what DefaultDict can do def __init__(self, default_factory): # type: (Callable[[_KT], _VT]) -> None pass def __missing__(self, key): # type: (_KT) -> _VT # typing: default_factory takes an arg, but the base class does not self[key] = self.default_factory(key) # type: ignore # pylint: disable=E1137 return self[key] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/sdk_worker_main.py0000644000175100017510000003646215206113032024425 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """SDK Fn Harness entry point.""" # pytype: skip-file import importlib import json import logging import os import re import signal import sys import time import traceback from google.protobuf import text_format from apache_beam.internal import pickler from apache_beam.io import filesystems from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import ProfilingOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import WorkerOptions from apache_beam.options.value_provider import RuntimeValueProvider from apache_beam.portability.api import endpoints_pb2 from apache_beam.runners.internal import names from apache_beam.runners.worker.data_sampler import DataSampler from apache_beam.runners.worker.log_handler import FnApiLogRecordHandler from apache_beam.runners.worker.sdk_worker import SdkHarness from apache_beam.utils import profiler _LOGGER = logging.getLogger(__name__) _ENABLE_GOOGLE_CLOUD_PROFILER = 'enable_google_cloud_profiler' _FN_LOG_HANDLER = None def _import_beam_plugins(plugins): for plugin in plugins: try: importlib.import_module(plugin) _LOGGER.debug('Imported beam-plugin %s', plugin) except ImportError: try: _LOGGER.debug(( "Looks like %s is not a module. " "Trying to import it assuming it's a class"), plugin) module, _ = plugin.rsplit('.', 1) importlib.import_module(module) _LOGGER.debug('Imported %s for beam-plugin %s', module, plugin) except ImportError as exc: _LOGGER.warning('Failed to import beam-plugin %s', plugin, exc_info=exc) def create_harness(environment, dry_run=False): """Creates SDK Fn Harness.""" # Bootstrap log level to capture startup events until pipeline options are # parsed and the actual log level is set. logging.getLogger().setLevel(logging.INFO) deferred_exception = None if 'LOGGING_API_SERVICE_DESCRIPTOR' in environment: try: logging_service_descriptor = endpoints_pb2.ApiServiceDescriptor() text_format.Merge( environment['LOGGING_API_SERVICE_DESCRIPTOR'], logging_service_descriptor) # Send all logs to the runner. fn_log_handler = FnApiLogRecordHandler(logging_service_descriptor) logging.getLogger().addHandler(fn_log_handler) _LOGGER.info('Logging handler created.') except Exception: _LOGGER.error( "Failed to set up logging handler, continuing without.", exc_info=True) fn_log_handler = None else: fn_log_handler = None options_json = environment.get('PIPELINE_OPTIONS') # We check if options are stored in the file. if 'PIPELINE_OPTIONS_FILE' in environment: options_file = environment['PIPELINE_OPTIONS_FILE'] try: with open(options_file, 'r') as f: options_json = f.read() _LOGGER.info('Load pipeline options from file: %s', options_file) except Exception: _LOGGER.error( 'Failed to load pipeline options from file: %s', options_file, exc_info=True) raise pipeline_options_dict = _load_pipeline_options(options_json) default_log_level = _get_log_level_from_options_dict(pipeline_options_dict) logging.getLogger().setLevel(default_log_level) _set_log_level_overrides(pipeline_options_dict) # These are used for dataflow templates. RuntimeValueProvider.set_runtime_options(pipeline_options_dict) sdk_pipeline_options = PipelineOptions.from_dictionary(pipeline_options_dict) filesystems.FileSystems.set_options(sdk_pipeline_options) pickle_library = sdk_pipeline_options.view_as(SetupOptions).pickle_library pickler.set_library(pickle_library) semi_persistent_directory = environment.get('SEMI_PERSISTENT_DIRECTORY', None) runner_capabilities = frozenset( environment.get('RUNNER_CAPABILITIES', '').split()) _LOGGER.info('semi_persistent_directory: %s', semi_persistent_directory) _worker_id = environment.get('WORKER_ID', None) try: _load_main_session(semi_persistent_directory) except LoadMainSessionException: exception_details = traceback.format_exc() _LOGGER.error( 'Could not load main session: %s', exception_details, exc_info=True) raise except Exception: # pylint: disable=broad-except summary = ( "Could not load main session. Inspect which external dependencies " "are used in the main module of your pipeline. Verify that " "corresponding packages are installed in the pipeline runtime " "environment and their installed versions match the versions used in " "pipeline submission environment. For more information, see: https://" "beam.apache.org/documentation/sdks/python-pipeline-dependencies/") _LOGGER.error(summary, exc_info=True) exception_details = traceback.format_exc() deferred_exception = LoadMainSessionException( f"{summary} {exception_details}") _LOGGER.info( 'Pipeline_options: %s', sdk_pipeline_options.get_all_options(drop_default=True)) control_service_descriptor = endpoints_pb2.ApiServiceDescriptor() status_service_descriptor = endpoints_pb2.ApiServiceDescriptor() text_format.Merge( environment['CONTROL_API_SERVICE_DESCRIPTOR'], control_service_descriptor) if 'STATUS_API_SERVICE_DESCRIPTOR' in environment: text_format.Merge( environment['STATUS_API_SERVICE_DESCRIPTOR'], status_service_descriptor) # TODO(robertwb): Support authentication. assert not control_service_descriptor.HasField('authentication') experiments = sdk_pipeline_options.view_as(DebugOptions).experiments or [] enable_heap_dump = 'enable_heap_dump' in experiments beam_plugins = sdk_pipeline_options.view_as(SetupOptions).beam_plugins or [] _import_beam_plugins(beam_plugins) if dry_run: return data_sampler = DataSampler.create(sdk_pipeline_options) sdk_harness = SdkHarness( control_address=control_service_descriptor.url, status_address=status_service_descriptor.url, worker_id=_worker_id, state_cache_size=_get_state_cache_size_bytes( options=sdk_pipeline_options), data_buffer_time_limit_ms=_get_data_buffer_time_limit_ms(experiments), profiler_factory=profiler.Profile.factory_from_options( sdk_pipeline_options.view_as(ProfilingOptions)), enable_heap_dump=enable_heap_dump, data_sampler=data_sampler, deferred_exception=deferred_exception, runner_capabilities=runner_capabilities, element_processing_timeout_minutes=sdk_pipeline_options.view_as( WorkerOptions).element_processing_timeout_minutes) return fn_log_handler, sdk_harness, sdk_pipeline_options def _start_profiler(gcp_profiler_service_name, gcp_profiler_service_version): try: import googlecloudprofiler if gcp_profiler_service_name and gcp_profiler_service_version: googlecloudprofiler.start( service=gcp_profiler_service_name, service_version=gcp_profiler_service_version, verbose=1) _LOGGER.info('Turning on Google Cloud Profiler.') else: raise RuntimeError('Unable to find the job id or job name from envvar.') except Exception as e: # pylint: disable=broad-except _LOGGER.warning( 'Unable to start google cloud profiler due to error: %s. For how to ' 'enable Cloud Profiler with Dataflow see ' 'https://cloud.google.com/dataflow/docs/guides/profiling-a-pipeline.' 'For troubleshooting tips with Cloud Profiler see ' 'https://cloud.google.com/profiler/docs/troubleshooting.' % e) def _get_gcp_profiler_name_if_enabled(sdk_pipeline_options): gcp_profiler_service_name = sdk_pipeline_options.view_as( GoogleCloudOptions).get_cloud_profiler_service_name() return gcp_profiler_service_name def main(unused_argv): """Main entry point for SDK Fn Harness.""" (fn_log_handler, sdk_harness, sdk_pipeline_options) = create_harness(os.environ) global _FN_LOG_HANDLER if fn_log_handler: _FN_LOG_HANDLER = fn_log_handler gcp_profiler_name = _get_gcp_profiler_name_if_enabled(sdk_pipeline_options) if gcp_profiler_name: _start_profiler(gcp_profiler_name, os.environ["JOB_ID"]) try: _LOGGER.info('Python sdk harness starting.') sdk_harness.run() _LOGGER.info('Python sdk harness exiting.') except: # pylint: disable=broad-except _LOGGER.critical('Python sdk harness failed: ', exc_info=True) raise finally: if fn_log_handler: fn_log_handler.close() def terminate_sdk_harness(): """Flushes the FnApiLogRecordHandler if it exists.""" _LOGGER.error('The SDK harness will be terminated in 5 seconds.') time.sleep(5) if _FN_LOG_HANDLER: _FN_LOG_HANDLER.close() os.kill(os.getpid(), signal.SIGINT) # Delay further control flow in the caller until process is terminated. time.sleep(60) # Try to force-terminate if still running. os.kill(os.getpid(), signal.SIGKILL) def _load_pipeline_options(options_json): """Deserialize the pipeline options from a JSON string into a dictionary.""" if options_json is None: return {} options = json.loads(options_json) # Check the options field first for backward compatibility. if 'options' in options: return options.get('options') else: # Remove extra urn part from the key. portable_option_regex = r'^beam:option:(?P.*):v1$' return { re.match(portable_option_regex, k).group('key') if re.match( portable_option_regex, k) else k: v for k, v in options.items() } def _parse_pipeline_options(options_json): """Parses the pipeline options from a JSON string into a PipelineOptions object.""" return PipelineOptions.from_dictionary(_load_pipeline_options(options_json)) def _get_state_cache_size_bytes(options): """Return the maximum size of state cache in bytes. Returns: an int indicating the maximum number of bytes to cache. """ max_cache_memory_usage_mb = options.view_as( WorkerOptions).max_cache_memory_usage_mb # to maintain backward compatibility experiments = options.view_as(DebugOptions).experiments or [] for experiment in experiments: # There should only be 1 match so returning from the loop if re.match(r'state_cache_size=', experiment): _LOGGER.warning( '--experiments=state_cache_size=X is deprecated and will be removed ' 'in future releases.' 'Please use --max_cache_memory_usage_mb=X to set the cache size for ' 'user state API and side inputs.') return int( re.match(r'state_cache_size=(?P.*)', experiment).group('state_cache_size')) << 20 return max_cache_memory_usage_mb << 20 def _get_data_buffer_time_limit_ms(experiments): """Defines the time limt of the outbound data buffering. Note: data_buffer_time_limit_ms is an experimental flag and might not be available in future releases. Returns: an int indicating the time limit in milliseconds of the outbound data buffering. Default is 0 (disabled) """ for experiment in experiments: # There should only be 1 match so returning from the loop if re.match(r'data_buffer_time_limit_ms=', experiment): return int( re.match( r'data_buffer_time_limit_ms=(?P.*)', experiment).group('data_buffer_time_limit_ms')) return 0 def _get_log_level_from_options_dict(options_dict: dict) -> int: """Get log level from options dict's entry `default_sdk_harness_log_level`. If not specified, default log level is logging.INFO. """ dict_level = options_dict.get('default_sdk_harness_log_level', 'INFO') log_level = dict_level if log_level.isdecimal(): log_level = int(log_level) else: # labeled log level log_level = getattr(logging, log_level, None) if not isinstance(log_level, int): # unknown log level. _LOGGER.error("Unknown log level %s. Use default value INFO.", dict_level) log_level = logging.INFO return log_level def _set_log_level_overrides(options_dict: dict) -> None: """Set module log level overrides from options dict's entry `sdk_harness_log_level_overrides`. """ parsed_overrides = options_dict.get('sdk_harness_log_level_overrides', None) if not isinstance(parsed_overrides, dict): if parsed_overrides is not None: _LOGGER.error( "Unable to parse sdk_harness_log_level_overrides: %s", parsed_overrides) return for module_name, log_level in parsed_overrides.items(): try: logging.getLogger(module_name).setLevel(log_level) except Exception as e: # Never crash the worker when exception occurs during log level setting # but logging the error. _LOGGER.error( "Error occurred when setting log level for %s: %s", module_name, e) class LoadMainSessionException(Exception): """ Used to crash this worker if a main session file failed to load. """ pass def _load_main_session(semi_persistent_directory): """Loads a pickled main session from the path specified.""" if pickler.is_currently_dill(): warn_msg = ' Functions defined in __main__ (interactive session) may fail.' err_msg = ' Functions defined in __main__ (interactive session) will ' \ 'almost certainly fail.' elif pickler.is_currently_cloudpickle(): warn_msg = ' User registered objects (e.g. schema, logical type) through' \ 'registeries may not be effective' err_msg = '' if semi_persistent_directory: session_file = os.path.join( semi_persistent_directory, 'staged', names.PICKLED_MAIN_SESSION_FILE) if os.path.isfile(session_file): # If the expected session file is present but empty, it's likely that # the user code run by this worker will likely crash at runtime. # This can happen if the worker fails to download the main session. # Raise a fatal error and crash this worker, forcing a restart. if os.path.getsize(session_file) == 0: if pickler.is_currently_dill(): # Potenitally transient error, unclear if still happening. raise LoadMainSessionException( 'Session file found, but empty: %s.%s' % (session_file, err_msg)) else: _LOGGER.warning('Empty session file: %s.%s', warn_msg, session_file) else: pickler.load_session(session_file) else: _LOGGER.warning('No session file found: %s.%s', warn_msg, session_file) else: _LOGGER.warning('No semi_persistent_directory found: %s', warn_msg) if __name__ == '__main__': main(sys.argv) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/sdk_worker_main_test.py0000644000175100017510000002364515206113032025463 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.runners.worker.sdk_worker_main.""" # pytype: skip-file import io import logging import os import unittest from hamcrest import all_of from hamcrest import assert_that from hamcrest import has_entry from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.value_provider import RuntimeValueProvider from apache_beam.runners.worker import sdk_worker_main from apache_beam.runners.worker import worker_status from apache_beam.utils.plugin import BeamPlugin class SdkWorkerMainTest(unittest.TestCase): # Used for testing newly added flags. class MockOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--eam:option:m_option:v', help='mock option') parser.add_argument('--eam:option:m_option:v1', help='mock option') parser.add_argument('--beam:option:m_option:v', help='mock option') parser.add_argument('--m_flag', action='store_true', help='mock flag') parser.add_argument('--m_option', help='mock option') parser.add_argument( '--m_m_option', action='append', help='mock multi option') def test_status_server(self): # Wrapping the method to see if it appears in threadump def wrapped_method_for_test(): threaddump = worker_status.thread_dump() self.assertRegex(threaddump, '.*wrapped_method_for_test.*') wrapped_method_for_test() def test_parse_pipeline_options(self): assert_that( sdk_worker_main._parse_pipeline_options( '{"options": {' + '"m_option": "/tmp/requirements.txt", ' + '"m_m_option":["beam_fn_api"]' + '}}').get_all_options(), all_of( has_entry('m_m_option', ['beam_fn_api']), has_entry('m_option', '/tmp/requirements.txt'))) assert_that( sdk_worker_main._parse_pipeline_options( '{"beam:option:m_option:v1": "/tmp/requirements.txt", ' + '"beam:option:m_m_option:v1":["beam_fn_api"]}').get_all_options(), all_of( has_entry('m_m_option', ['beam_fn_api']), has_entry('m_option', '/tmp/requirements.txt'))) assert_that( sdk_worker_main._parse_pipeline_options( '{"options": {"beam:option:m_option:v":"mock_val"}}'). get_all_options(), has_entry('beam:option:m_option:v', 'mock_val')) assert_that( sdk_worker_main._parse_pipeline_options( '{"options": {"eam:option:m_option:v1":"mock_val"}}'). get_all_options(), has_entry('eam:option:m_option:v1', 'mock_val')) assert_that( sdk_worker_main._parse_pipeline_options( '{"options": {"eam:option:m_option:v":"mock_val"}}'). get_all_options(), has_entry('eam:option:m_option:v', 'mock_val')) def test_runtime_values(self): test_runtime_provider = RuntimeValueProvider('test_param', int, None) sdk_worker_main.create_harness({ 'CONTROL_API_SERVICE_DESCRIPTOR': '', 'PIPELINE_OPTIONS': '{"test_param": 37}', }, dry_run=True) self.assertTrue(test_runtime_provider.is_accessible()) self.assertEqual(test_runtime_provider.get(), 37) def test_create_sdk_harness_log_handler_received_log(self): # tests that the log handler created in create_harness() does not miss # logs emitted from create_harness() itself. logstream = io.StringIO() class InMemoryHandler(logging.StreamHandler): def __init__(self, *unused): super().__init__(stream=logstream) with unittest.mock.patch( 'apache_beam.runners.worker.sdk_worker_main.FnApiLogRecordHandler', InMemoryHandler): sdk_worker_main.create_harness({ 'LOGGING_API_SERVICE_DESCRIPTOR': '', 'CONTROL_API_SERVICE_DESCRIPTOR': '', 'PIPELINE_OPTIONS': '{"default_sdk_harness_log_level":"INVALID",' '"sdk_harness_log_level_overrides":[]}', }, dry_run=True) logstream.seek(0) logs = logstream.read() self.assertIn('Unknown log level', logs) self.assertIn('Unable to parse sdk_harness_log_level_overrides', logs) def test_import_beam_plugins(self): sdk_worker_main._import_beam_plugins(BeamPlugin.get_all_plugin_paths()) @staticmethod def _overrides_case_to_option_dict(case): """ Return logging level overrides from command line strings via PipelineOption. """ options_list = [] for c in case: options_list += ['--sdk_harness_log_level_overrides', c] options = PipelineOptions(options_list) return options.get_all_options() def test__get_log_level_from_options_dict(self): test_cases = [ {}, { 'default_sdk_harness_log_level': 'DEBUG' }, { 'default_sdk_harness_log_level': '30' }, { 'default_sdk_harness_log_level': 'INVALID_ENTRY' }, ] expected_results = [logging.INFO, logging.DEBUG, 30, logging.INFO] for case, expected in zip(test_cases, expected_results): self.assertEqual( sdk_worker_main._get_log_level_from_options_dict(case), expected) def test__set_log_level_overrides(self): test_cases = [ ([], {}), # not provided, as a smoke test ( # single overrides ['{"fake_module_1a.b":"DEBUG","fake_module_1c.d":"INFO"}'], { "fake_module_1a.b": logging.DEBUG, "fake_module_1a.b.f": logging.DEBUG, "fake_module_1c.d": logging.INFO }), ( # multiple overrides, the last takes precedence [ '{"fake_module_2a.b":"DEBUG"}', '{"fake_module_2c.d":"WARNING","fake_module_2c.d.e":15}', '{"fake_module_2c.d":"ERROR"}' ], { "fake_module_2a.b": logging.DEBUG, "fake_module_2a.b.f": logging.DEBUG, "fake_module_2c.d": logging.ERROR, "fake_module_2c.d.e": 15, "fake_module_2c.d.f": logging.ERROR }) ] for case, expected in test_cases: overrides = self._overrides_case_to_option_dict(case) sdk_worker_main._set_log_level_overrides(overrides) for name, level in expected.items(): self.assertEqual(logging.getLogger(name).getEffectiveLevel(), level) def test__set_log_level_overrides_error(self): test_cases = [ (['{"json.value.is.not.level": ["ERROR"]}'], "Error occurred when setting log level"), (['{"invalid.level":"INVALID"}'], "Error occurred when setting log level"), ] for case, expected in test_cases: overrides = self._overrides_case_to_option_dict(case) with self.assertLogs('apache_beam.runners.worker.sdk_worker_main', level='ERROR') as cm: sdk_worker_main._set_log_level_overrides(overrides) self.assertIn(expected, cm.output[0]) def test_gcp_profiler_uses_provided_service_name_when_specified(self): options = PipelineOptions( ['--dataflow_service_options=enable_google_cloud_profiler=sample']) gcp_profiler_name = sdk_worker_main._get_gcp_profiler_name_if_enabled( options) sdk_worker_main._start_profiler = unittest.mock.MagicMock() sdk_worker_main._start_profiler(gcp_profiler_name, "version") sdk_worker_main._start_profiler.assert_called_with("sample", "version") @unittest.mock.patch.dict(os.environ, {"JOB_NAME": "sample_job"}, clear=True) def test_gcp_profiler_uses_job_name_when_service_name_not_specified(self): options = PipelineOptions( ['--dataflow_service_options=enable_google_cloud_profiler']) gcp_profiler_name = sdk_worker_main._get_gcp_profiler_name_if_enabled( options) sdk_worker_main._start_profiler = unittest.mock.MagicMock() sdk_worker_main._start_profiler(gcp_profiler_name, "version") sdk_worker_main._start_profiler.assert_called_with("sample_job", "version") @unittest.mock.patch.dict(os.environ, {"JOB_NAME": "sample_job"}, clear=True) def test_gcp_profiler_uses_job_name_when_enabled_as_experiment(self): options = PipelineOptions(['--experiment=enable_google_cloud_profiler']) gcp_profiler_name = sdk_worker_main._get_gcp_profiler_name_if_enabled( options) sdk_worker_main._start_profiler = unittest.mock.MagicMock() sdk_worker_main._start_profiler(gcp_profiler_name, "version") sdk_worker_main._start_profiler.assert_called_with("sample_job", "version") @unittest.mock.patch.dict(os.environ, {"JOB_NAME": "sample_job"}, clear=True) def test_pipeline_option_max_cache_memory_usage_mb(self): options = PipelineOptions(flags=['--max_cache_memory_usage_mb=50']) cache_size = sdk_worker_main._get_state_cache_size_bytes(options) self.assertEqual(cache_size, 50 << 20) @unittest.mock.patch.dict(os.environ, {"JOB_NAME": "sample_job"}, clear=True) def test_pipeline_option_max_cache_memory_usage_mb_with_experiments(self): options = PipelineOptions(flags=['--experiments=state_cache_size=50']) cache_size = sdk_worker_main._get_state_cache_size_bytes(options) self.assertEqual(cache_size, 50 << 20) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/sdk_worker_test.py0000644000175100017510000007405415206113032024457 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.runners.worker.sdk_worker.""" # pytype: skip-file import contextlib import logging import unittest from collections import namedtuple import grpc import hamcrest as hc import mock from apache_beam.coders import FastPrimitivesCoder from apache_beam.coders import VarIntCoder from apache_beam.internal.metrics.metric import Metrics as InternalMetrics from apache_beam.metrics import monitoring_infos from apache_beam.metrics.execution import MetricsEnvironment from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.portability.api import metrics_pb2 from apache_beam.runners.worker import data_plane from apache_beam.runners.worker import sdk_worker from apache_beam.runners.worker import statecache from apache_beam.runners.worker.sdk_worker import BundleProcessorCache from apache_beam.runners.worker.sdk_worker import GlobalCachingStateHandler from apache_beam.runners.worker.sdk_worker import SdkWorker from apache_beam.utils import thread_pool_executor _LOGGER = logging.getLogger(__name__) class BeamFnControlServicer(beam_fn_api_pb2_grpc.BeamFnControlServicer): def __init__(self, requests, raise_errors=True): self.requests = requests self.instruction_ids = set(r.instruction_id for r in requests) self.responses = {} self.raise_errors = raise_errors def Control(self, response_iterator, context): for request in self.requests: _LOGGER.info("Sending request %s", request) yield request for response in response_iterator: _LOGGER.info("Got response %s", response) if response.instruction_id != -1: assert response.instruction_id in self.instruction_ids assert response.instruction_id not in self.responses self.responses[response.instruction_id] = response if self.raise_errors and response.error: raise RuntimeError(response.error) elif len(self.responses) == len(self.requests): _LOGGER.info("All %s instructions finished.", len(self.requests)) return raise RuntimeError( "Missing responses: %s" % (self.instruction_ids - set(self.responses.keys()))) class SdkWorkerTest(unittest.TestCase): def _get_process_bundles(self, prefix, size): return [ beam_fn_api_pb2.ProcessBundleDescriptor( id=str(str(prefix) + "-" + str(ix)), transforms={ str(ix): beam_runner_api_pb2.PTransform(unique_name=str(ix)) }) for ix in range(size) ] def _check_fn_registration_multi_request(self, *args): """Check the function registration calls to the sdk_harness. Args: tuple of request_count, number of process_bundles per request and workers counts to process the request. """ for (request_count, process_bundles_per_request) in args: requests = [] process_bundle_descriptors = [] for i in range(request_count): pbd = self._get_process_bundles(i, process_bundles_per_request) process_bundle_descriptors.extend(pbd) requests.append( beam_fn_api_pb2.InstructionRequest( instruction_id=str(i), register=beam_fn_api_pb2.RegisterRequest( process_bundle_descriptor=process_bundle_descriptors))) test_controller = BeamFnControlServicer(requests) server = grpc.server(thread_pool_executor.shared_unbounded_instance()) beam_fn_api_pb2_grpc.add_BeamFnControlServicer_to_server( test_controller, server) test_port = server.add_insecure_port("[::]:0") server.start() harness = sdk_worker.SdkHarness( "localhost:%s" % test_port, state_cache_size=100) harness.run() self.assertEqual( harness._bundle_processor_cache.fns, {item.id: item for item in process_bundle_descriptors}) def test_fn_registration(self): self._check_fn_registration_multi_request((1, 4), (4, 4)) def test_inactive_bundle_processor_returns_empty_progress_response(self): bundle_processor = mock.MagicMock() data_channel_factory = mock.create_autospec( data_plane.GrpcClientDataChannelFactory) bundle_processor_cache = BundleProcessorCache( None, None, data_channel_factory, {}) bundle_processor_cache.activate('instruction_id') worker = SdkWorker(bundle_processor_cache) split_request = beam_fn_api_pb2.InstructionRequest( instruction_id='progress_instruction_id', process_bundle_progress=beam_fn_api_pb2.ProcessBundleProgressRequest( instruction_id='instruction_id')) self.assertEqual( worker.do_instruction(split_request), beam_fn_api_pb2.InstructionResponse( instruction_id='progress_instruction_id', process_bundle_progress=beam_fn_api_pb2. ProcessBundleProgressResponse(consuming_received_data=False))) # Add a mock bundle processor as if it was running before it's released bundle_processor_cache.active_bundle_processors['instruction_id'] = ( 'descriptor_id', bundle_processor) bundle_processor_cache.release('instruction_id') self.assertEqual( worker.do_instruction(split_request), beam_fn_api_pb2.InstructionResponse( instruction_id='progress_instruction_id', process_bundle_progress=beam_fn_api_pb2. ProcessBundleProgressResponse(consuming_received_data=False))) def test_failed_bundle_processor_returns_failed_progress_response(self): bundle_processor = mock.MagicMock() data_channel_factory = mock.create_autospec( data_plane.GrpcClientDataChannelFactory) bundle_processor_cache = BundleProcessorCache( None, None, data_channel_factory, {}) bundle_processor_cache.activate('instruction_id') worker = SdkWorker(bundle_processor_cache) # Add a mock bundle processor as if it was running before it's discarded bundle_processor_cache.active_bundle_processors['instruction_id'] = ( 'descriptor_id', bundle_processor) bundle_processor_cache.discard( 'instruction_id', RuntimeError('test message')) split_request = beam_fn_api_pb2.InstructionRequest( instruction_id='progress_instruction_id', process_bundle_progress=beam_fn_api_pb2.ProcessBundleProgressRequest( instruction_id='instruction_id')) hc.assert_that( worker.do_instruction(split_request).error, hc.contains_string( 'Bundle processing associated with instruction_id has failed')) hc.assert_that( worker.do_instruction(split_request).error, hc.contains_string('test message')) def test_inactive_bundle_processor_returns_empty_split_response(self): bundle_processor = mock.MagicMock() data_channel_factory = mock.create_autospec( data_plane.GrpcClientDataChannelFactory) bundle_processor_cache = BundleProcessorCache( None, None, data_channel_factory, {}) bundle_processor_cache.activate('instruction_id') worker = SdkWorker(bundle_processor_cache) split_request = beam_fn_api_pb2.InstructionRequest( instruction_id='split_instruction_id', process_bundle_split=beam_fn_api_pb2.ProcessBundleSplitRequest( instruction_id='instruction_id')) self.assertEqual( worker.do_instruction(split_request), beam_fn_api_pb2.InstructionResponse( instruction_id='split_instruction_id', process_bundle_split=beam_fn_api_pb2.ProcessBundleSplitResponse())) # Add a mock bundle processor as if it was running before it's released bundle_processor_cache.active_bundle_processors['instruction_id'] = ( 'descriptor_id', bundle_processor) bundle_processor_cache.release('instruction_id') self.assertEqual( worker.do_instruction(split_request), beam_fn_api_pb2.InstructionResponse( instruction_id='split_instruction_id', process_bundle_split=beam_fn_api_pb2.ProcessBundleSplitResponse())) def get_responses(self, instruction_requests, data_sampler=None): """Evaluates and returns {id: InstructionResponse} for the requests.""" test_controller = BeamFnControlServicer(instruction_requests) server = grpc.server(thread_pool_executor.shared_unbounded_instance()) beam_fn_api_pb2_grpc.add_BeamFnControlServicer_to_server( test_controller, server) test_port = server.add_insecure_port("[::]:0") server.start() harness = sdk_worker.SdkHarness( "localhost:%s" % test_port, state_cache_size=100, data_sampler=data_sampler) harness.run() return test_controller.responses def test_harness_monitoring_infos_and_metadata(self): # Clear the process wide metric container. MetricsEnvironment.process_wide_container().metrics = {} # Create a process_wide metric. urn = 'my.custom.urn' labels = {'key': 'value'} InternalMetrics.counter(urn=urn, labels=labels, process_wide=True).inc(10) harness_monitoring_infos_request = beam_fn_api_pb2.InstructionRequest( instruction_id="monitoring_infos", harness_monitoring_infos=beam_fn_api_pb2.HarnessMonitoringInfosRequest( )) responses = self.get_responses([harness_monitoring_infos_request]) expected_monitoring_info = monitoring_infos.int64_counter( urn, 10, labels=labels) monitoring_data = ( responses['monitoring_infos'].harness_monitoring_infos.monitoring_data) # Request the full MonitoringInfo metadata for the returned short_ids. short_ids = list(monitoring_data.keys()) monitoring_infos_metadata_request = beam_fn_api_pb2.InstructionRequest( instruction_id="monitoring_infos_metadata", monitoring_infos=beam_fn_api_pb2.MonitoringInfosMetadataRequest( monitoring_info_id=short_ids)) responses = self.get_responses([monitoring_infos_metadata_request]) # Request the full MonitoringInfo metadata to be returned now. expected_monitoring_info.ClearField("payload") # Verify that one of the returned monitoring infos is our expected # monitoring info. short_id_to_mi = ( responses['monitoring_infos_metadata'].monitoring_infos.monitoring_info) found = False for mi in short_id_to_mi.values(): # Clear the timestamp before comparing mi.ClearField("start_time") if mi == expected_monitoring_info: found = True self.assertTrue(found, str(responses['monitoring_infos_metadata'])) def test_failed_bundle_processor_returns_failed_split_response(self): bundle_processor = mock.MagicMock() data_channel_factory = mock.create_autospec( data_plane.GrpcClientDataChannelFactory) bundle_processor_cache = BundleProcessorCache( None, None, data_channel_factory, {}) bundle_processor_cache.activate('instruction_id') worker = SdkWorker(bundle_processor_cache) # Add a mock bundle processor as if it was running before it's discarded bundle_processor_cache.active_bundle_processors['instruction_id'] = ( 'descriptor_id', bundle_processor) bundle_processor_cache.discard( 'instruction_id', RuntimeError('test message')) split_request = beam_fn_api_pb2.InstructionRequest( instruction_id='split_instruction_id', process_bundle_split=beam_fn_api_pb2.ProcessBundleSplitRequest( instruction_id='instruction_id')) hc.assert_that( worker.do_instruction(split_request).error, hc.contains_string( 'Bundle processing associated with instruction_id has failed')) hc.assert_that( worker.do_instruction(split_request).error, hc.contains_string('test message')) def test_data_sampling_response(self): # Create a data sampler with some fake sampled data. This data will be seen # in the sample response. coder = FastPrimitivesCoder() class FakeDataSampler: def samples(self, pcollection_ids): return beam_fn_api_pb2.SampleDataResponse( element_samples={ 'pcoll_id_1': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=coder.encode_nested('a')) ]), 'pcoll_id_2': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=coder.encode_nested('b')) ]) }) def stop(self): pass data_sampler = FakeDataSampler() # Create and send the fake reponse. The SdkHarness should query the # DataSampler and fill out the sample response. sample_request = beam_fn_api_pb2.InstructionRequest( instruction_id='sample_request', sample_data=beam_fn_api_pb2.SampleDataRequest( pcollection_ids=['pcoll_id_1', 'pcoll_id_2'])) responses = self.get_responses([sample_request], data_sampler) self.assertEqual(len(responses), 1) # Get and assert the correct response. response = responses['sample_request'] expected_response = beam_fn_api_pb2.InstructionResponse( instruction_id='sample_request', sample_data=beam_fn_api_pb2.SampleDataResponse( element_samples={ 'pcoll_id_1': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=coder.encode_nested('a')) ]), 'pcoll_id_2': beam_fn_api_pb2.SampleDataResponse.ElementList( elements=[ beam_fn_api_pb2.SampledElement( element=coder.encode_nested('b')) ]) })) self.assertEqual(response, expected_response) def test_bundle_processor_creation_failure_cleans_up_grpc_data_channel(self): data_channel_factory = data_plane.GrpcClientDataChannelFactory() channel = data_channel_factory.create_data_channel_from_url('some_url') state_handler_factory = mock.create_autospec( sdk_worker.GrpcStateHandlerFactory) bundle_processor_cache = BundleProcessorCache( frozenset(), state_handler_factory, data_channel_factory, {}) if bundle_processor_cache.periodic_shutdown: bundle_processor_cache.periodic_shutdown.cancel() bundle_processor_cache.get = mock.MagicMock( side_effect=RuntimeError('test error')) worker = SdkWorker(bundle_processor_cache) instruction_id = 'instruction_id' request = beam_fn_api_pb2.ProcessBundleRequest( process_bundle_descriptor_id='descriptor_id') with self.assertRaises(RuntimeError): worker.process_bundle(request, instruction_id) self.assertIn(instruction_id, channel._cleaned_instruction_ids) class CachingStateHandlerTest(unittest.TestCase): def test_caching(self): coder = VarIntCoder() coder_impl = coder.get_impl() class FakeUnderlyingState(object): """Simply returns an incremented counter as the state "value." """ def set_counter(self, n): self._counter = n def get_raw(self, *args): self._counter += 1 return coder.encode(self._counter), None @contextlib.contextmanager def process_instruction_id(self, bundle_id): yield underlying_state = FakeUnderlyingState() state_cache = statecache.StateCache(100 << 20) caching_state_hander = GlobalCachingStateHandler( state_cache, underlying_state) state1 = beam_fn_api_pb2.StateKey( bag_user_state=beam_fn_api_pb2.StateKey.BagUserState( user_state_id='state1')) state2 = beam_fn_api_pb2.StateKey( bag_user_state=beam_fn_api_pb2.StateKey.BagUserState( user_state_id='state2')) side1 = beam_fn_api_pb2.StateKey( multimap_side_input=beam_fn_api_pb2.StateKey.MultimapSideInput( transform_id='transform', side_input_id='side1')) side2 = beam_fn_api_pb2.StateKey( iterable_side_input=beam_fn_api_pb2.StateKey.IterableSideInput( transform_id='transform', side_input_id='side2')) state_token1 = beam_fn_api_pb2.ProcessBundleRequest.CacheToken( token=b'state_token1', user_state=beam_fn_api_pb2.ProcessBundleRequest.CacheToken.UserState()) state_token2 = beam_fn_api_pb2.ProcessBundleRequest.CacheToken( token=b'state_token2', user_state=beam_fn_api_pb2.ProcessBundleRequest.CacheToken.UserState()) side1_token1 = beam_fn_api_pb2.ProcessBundleRequest.CacheToken( token=b'side1_token1', side_input=beam_fn_api_pb2.ProcessBundleRequest.CacheToken.SideInput( transform_id='transform', side_input_id='side1')) side1_token2 = beam_fn_api_pb2.ProcessBundleRequest.CacheToken( token=b'side1_token2', side_input=beam_fn_api_pb2.ProcessBundleRequest.CacheToken.SideInput( transform_id='transform', side_input_id='side1')) def get_as_list(key): return list(caching_state_hander.blocking_get(key, coder_impl)) underlying_state.set_counter(100) with caching_state_hander.process_instruction_id('bundle1', []): self.assertEqual(get_as_list(state1), [101]) # uncached self.assertEqual(get_as_list(state2), [102]) # uncached self.assertEqual(get_as_list(state1), [101]) # cached on bundle self.assertEqual(get_as_list(side1), [103]) # uncached self.assertEqual(get_as_list(side2), [104]) # uncached underlying_state.set_counter(200) with caching_state_hander.process_instruction_id( 'bundle2', [state_token1, side1_token1]): self.assertEqual(get_as_list(state1), [201]) # uncached self.assertEqual(get_as_list(state2), [202]) # uncached self.assertEqual(get_as_list(state1), [201]) # cached on state token1 self.assertEqual(get_as_list(side1), [203]) # uncached self.assertEqual(get_as_list(side1), [203]) # cached on side1_token1 self.assertEqual(get_as_list(side2), [204]) # uncached self.assertEqual(get_as_list(side2), [204]) # cached on bundle underlying_state.set_counter(300) with caching_state_hander.process_instruction_id( 'bundle3', [state_token1, side1_token1]): self.assertEqual(get_as_list(state1), [201]) # cached on state token1 self.assertEqual(get_as_list(state2), [202]) # cached on state token1 self.assertEqual(get_as_list(state1), [201]) # cached on state token1 self.assertEqual(get_as_list(side1), [203]) # cached on side1_token1 self.assertEqual(get_as_list(side1), [203]) # cached on side1_token1 self.assertEqual(get_as_list(side2), [301]) # uncached self.assertEqual(get_as_list(side2), [301]) # cached on bundle underlying_state.set_counter(400) with caching_state_hander.process_instruction_id( 'bundle4', [state_token2, side1_token1]): self.assertEqual(get_as_list(state1), [401]) # uncached self.assertEqual(get_as_list(state2), [402]) # uncached self.assertEqual(get_as_list(state1), [401]) # cached on state token2 self.assertEqual(get_as_list(side1), [203]) # cached on side1_token1 self.assertEqual(get_as_list(side1), [203]) # cached on side1_token1 self.assertEqual(get_as_list(side2), [403]) # uncached self.assertEqual(get_as_list(side2), [403]) # cached on bundle underlying_state.set_counter(500) with caching_state_hander.process_instruction_id( 'bundle5', [state_token2, side1_token2]): self.assertEqual(get_as_list(state1), [401]) # cached on state token2 self.assertEqual(get_as_list(state2), [402]) # cached on state token2 self.assertEqual(get_as_list(state1), [401]) # cached on state token2 self.assertEqual(get_as_list(side1), [501]) # uncached self.assertEqual(get_as_list(side1), [501]) # cached on side1_token2 self.assertEqual(get_as_list(side2), [502]) # uncached self.assertEqual(get_as_list(side2), [502]) # cached on bundle class UnderlyingStateHandler(object): """Simply returns an incremented counter as the state "value." """ def __init__(self): self._encoded_values = [] self._continuations = False def set_value(self, value, coder): self._encoded_values = [coder.encode(value)] def set_values(self, values, coder): self._encoded_values = [coder.encode(value) for value in values] def set_continuations(self, continuations): self._continuations = continuations def get_raw(self, _state_key, continuation_token=None): if self._continuations and len(self._encoded_values) > 0: if not continuation_token: continuation_token = '0' idx = int(continuation_token) next_token = str(idx + 1) if idx + 1 < len(self._encoded_values) else None return self._encoded_values[idx], next_token else: return b''.join(self._encoded_values), None def append_raw(self, _key, bytes): self._encoded_values.append(bytes) def clear(self, *args): self._encoded_values = [] @contextlib.contextmanager def process_instruction_id(self, bundle_id): yield def test_append_clear_with_preexisting_state(self): state = beam_fn_api_pb2.StateKey( bag_user_state=beam_fn_api_pb2.StateKey.BagUserState( user_state_id='state1')) cache_token = beam_fn_api_pb2.ProcessBundleRequest.CacheToken( token=b'state_token1', user_state=beam_fn_api_pb2.ProcessBundleRequest.CacheToken.UserState()) coder = VarIntCoder() underlying_state_handler = self.UnderlyingStateHandler() state_cache = statecache.StateCache(100 << 20) handler = GlobalCachingStateHandler(state_cache, underlying_state_handler) def get(): return handler.blocking_get(state, coder.get_impl()) def append(iterable): handler.extend(state, coder.get_impl(), iterable) def clear(): handler.clear(state) # Initialize state underlying_state_handler.set_value(42, coder) with handler.process_instruction_id('bundle', [cache_token]): # Append without reading beforehand append([43]) self.assertEqual(get(), [42, 43]) clear() self.assertEqual(get(), []) append([44, 45]) self.assertEqual(get(), [44, 45]) append((46, 47)) self.assertEqual(get(), [44, 45, 46, 47]) clear() append(range(1000)) self.assertEqual(get(), list(range(1000))) def test_continuation_token(self): underlying_state_handler = self.UnderlyingStateHandler() state_cache = statecache.StateCache(100 << 20) handler = GlobalCachingStateHandler(state_cache, underlying_state_handler) coder = VarIntCoder() state = beam_fn_api_pb2.StateKey( bag_user_state=beam_fn_api_pb2.StateKey.BagUserState( user_state_id='state1')) cache_token = beam_fn_api_pb2.ProcessBundleRequest.CacheToken( token=b'state_token1', user_state=beam_fn_api_pb2.ProcessBundleRequest.CacheToken.UserState()) def get(materialize=True): result = handler.blocking_get(state, coder.get_impl()) return list(result) if materialize else result def get_type(): return type(get(materialize=False)) def append(*values): handler.extend(state, coder.get_impl(), values) def clear(): handler.clear(state) underlying_state_handler.set_continuations(True) underlying_state_handler.set_values([45, 46, 47], coder) with handler.process_instruction_id('bundle', [cache_token]): self.assertEqual( get_type(), GlobalCachingStateHandler.ContinuationIterable) self.assertEqual(get(), [45, 46, 47]) append(48, 49) self.assertEqual( get_type(), GlobalCachingStateHandler.ContinuationIterable) self.assertEqual(get(), [45, 46, 47, 48, 49]) clear() self.assertEqual(get_type(), list) self.assertEqual(get(), []) append(1) self.assertEqual(get(), [1]) append(2, 3) self.assertEqual(get(), [1, 2, 3]) clear() for i in range(1000): append(i) self.assertEqual(get_type(), list) self.assertEqual(get(), [i for i in range(1000)]) class ShortIdCacheTest(unittest.TestCase): def testShortIdAssignment(self): TestCase = namedtuple('TestCase', ['expected_short_id', 'info']) test_cases = [ TestCase(*args) for args in [ ( "1", metrics_pb2.MonitoringInfo( urn="beam:metric:user:distribution_int64:v1", type="beam:metrics:distribution_int64:v1")), ( "2", metrics_pb2.MonitoringInfo( urn="beam:metric:element_count:v1", type="beam:metrics:sum_int64:v1")), ( "3", metrics_pb2.MonitoringInfo( urn="beam:metric:ptransform_progress:completed:v1", type="beam:metrics:progress:v1")), ( "4", metrics_pb2.MonitoringInfo( urn="beam:metric:user:distribution_double:v1", type="beam:metrics:distribution_double:v1")), ( "5", metrics_pb2.MonitoringInfo( urn="TestingSentinelUrn", type="TestingSentinelType")), ( "6", metrics_pb2.MonitoringInfo( urn= "beam:metric:pardo_execution_time:finish_bundle_msecs:v1", type="beam:metrics:sum_int64:v1")), # This case and the next one validates that different labels # with the same urn are in fact assigned different short ids. ( "7", metrics_pb2.MonitoringInfo( urn="beam:metric:user:sum_int64:v1", type="beam:metrics:sum_int64:v1", labels={ "PTRANSFORM": "myT", "NAMESPACE": "harness", "NAME": "metricNumber7" })), ( "8", metrics_pb2.MonitoringInfo( urn="beam:metric:user:sum_int64:v1", type="beam:metrics:sum_int64:v1", labels={ "PTRANSFORM": "myT", "NAMESPACE": "harness", "NAME": "metricNumber8" })), ( "9", metrics_pb2.MonitoringInfo( urn="beam:metric:user:top_n_double:v1", type="beam:metrics:top_n_double:v1", labels={ "PTRANSFORM": "myT", "NAMESPACE": "harness", "NAME": "metricNumber7" })), ( "a", metrics_pb2.MonitoringInfo( urn="beam:metric:element_count:v1", type="beam:metrics:sum_int64:v1", labels={"PCOLLECTION": "myPCol"})), # validate payload is ignored for shortId assignment ( "3", metrics_pb2.MonitoringInfo( urn="beam:metric:ptransform_progress:completed:v1", type="beam:metrics:progress:v1", payload=b"this is ignored!")) ] ] cache = sdk_worker.ShortIdCache() for case in test_cases: self.assertEqual( case.expected_short_id, cache.get_short_id(case.info), "Got incorrect short id for monitoring info:\n%s" % case.info) # Retrieve all of the monitoring infos by short id, and verify that the # metadata (everything but the payload) matches the originals actual_recovered_infos = cache.get_infos( case.expected_short_id for case in test_cases).values() for recoveredInfo, case in zip(actual_recovered_infos, test_cases): self.assertEqual( monitoringInfoMetadata(case.info), monitoringInfoMetadata(recoveredInfo)) # Retrieve short ids one more time in reverse for case in reversed(test_cases): self.assertEqual( case.expected_short_id, cache.get_short_id(case.info), "Got incorrect short id on second retrieval for monitoring info:\n%s" % case.info) class DeferredCallTest(unittest.TestCase): """Tests for _DeferredCall.get().""" def test_get_single_arg(self): f = sdk_worker._Future().set(42) call = sdk_worker._DeferredCall(lambda x: x, f) self.assertEqual(call.get(), 42) def test_get_multiple_args(self): futures = [sdk_worker._Future().set(i) for i in range(5)] call = sdk_worker._DeferredCall(lambda *args: sum(args), *futures) self.assertEqual(call.get(), sum(range(5))) def test_get_non_future_args_are_wrapped(self): # __init__ wraps non-Future values in _Future().set(v); get() must work. call = sdk_worker._DeferredCall(lambda x, y: x * y, 3, 7) self.assertEqual(call.get(), 21) def test_get_mixed_future_and_value_args(self): a = sdk_worker._Future().set(10) call = sdk_worker._DeferredCall(lambda x, y: x + y, a, 5) self.assertEqual(call.get(), 15) def test_get_zero_args(self): call = sdk_worker._DeferredCall(lambda: 99) self.assertEqual(call.get(), 99) def test_get_preserves_return_value_type(self): f = sdk_worker._Future().set({'key': 'val'}) call = sdk_worker._DeferredCall(lambda d: d, f) self.assertEqual(call.get(), {'key': 'val'}) def monitoringInfoMetadata(info): return { descriptor.name: value for descriptor, value in info.ListFields() if not descriptor.name == "payload" } if __name__ == "__main__": logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/sideinputs.py0000644000175100017510000001672215206113032023433 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utilities for handling side inputs.""" # pytype: skip-file import logging import queue import threading import traceback from collections import abc from apache_beam.coders import observable from apache_beam.io import iobase from apache_beam.runners.worker import opcounters from apache_beam.transforms import window from apache_beam.utils.sentinel import Sentinel # Maximum number of reader threads for reading side input sources, per side # input. MAX_SOURCE_READER_THREADS = 15 # Number of slots for elements in side input element queue. Note that this # value is intentionally smaller than MAX_SOURCE_READER_THREADS so as to reduce # memory pressure of holding potentially-large elements in memory. Note that # the number of pending elements in memory is equal to the sum of # MAX_SOURCE_READER_THREADS and ELEMENT_QUEUE_SIZE. ELEMENT_QUEUE_SIZE = 10 # Special element value sentinel for signaling reader state. READER_THREAD_IS_DONE_SENTINEL = Sentinel.sentinel # Used to efficiently window the values of non-windowed side inputs. _globally_windowed = window.GlobalWindows.windowed_value(None).with_value _LOGGER = logging.getLogger(__name__) class PrefetchingSourceSetIterable(object): """Value iterator that reads concurrently from a set of sources.""" def __init__( self, sources, max_reader_threads=MAX_SOURCE_READER_THREADS, read_counter=None, element_counter=None): self.sources = sources self.num_reader_threads = min(max_reader_threads, len(self.sources)) # Queue for sources that are to be read. self.sources_queue = queue.Queue() for source in sources: self.sources_queue.put(source) # Queue for elements that have been read. self.element_queue = queue.Queue(ELEMENT_QUEUE_SIZE) # Queue for exceptions encountered in reader threads; to be rethrown. self.reader_exceptions = queue.Queue() # Whether we have already iterated; this iterable can only be used once. self.already_iterated = False # Whether an error was encountered in any source reader. self.has_errored = False self.read_counter = read_counter or opcounters.NoOpTransformIOCounter() self.element_counter = element_counter self.reader_threads = [] self._start_reader_threads() def add_byte_counter(self, reader): """Adds byte counter observer to a side input reader. Args: reader: A reader that should inherit from ObservableMixin to have bytes tracked. """ def update_bytes_read(record_size, is_record_size=False, **kwargs): # Let the reader report block size. if is_record_size: self.read_counter.add_bytes_read(record_size) if isinstance(reader, observable.ObservableMixin): reader.register_observer(update_bytes_read) def _start_reader_threads(self): for _ in range(0, self.num_reader_threads): t = threading.Thread(target=self._reader_thread) t.daemon = True t.start() self.reader_threads.append(t) def _reader_thread(self): # pylint: disable=too-many-nested-blocks try: while True: try: source = self.sources_queue.get_nowait() if isinstance(source, iobase.BoundedSource): for value in source.read(source.get_range_tracker(None, None)): if self.has_errored: # If any reader has errored, just return. return if isinstance(value, window.WindowedValue): self.element_queue.put(value) else: self.element_queue.put(_globally_windowed(value)) else: # Native dataflow source. with source.reader() as reader: # The tracking of time spend reading and bytes read from side # inputs is kept behind an experiment flag to test performance # impact. self.add_byte_counter(reader) returns_windowed_values = reader.returns_windowed_values for value in reader: if self.has_errored: # If any reader has errored, just return. return if returns_windowed_values: self.element_queue.put(value) else: self.element_queue.put(_globally_windowed(value)) except queue.Empty: return except Exception as e: # pylint: disable=broad-except _LOGGER.error( 'Encountered exception in PrefetchingSourceSetIterable ' 'reader thread: %s', traceback.format_exc()) self.reader_exceptions.put(e) self.has_errored = True finally: self.element_queue.put(READER_THREAD_IS_DONE_SENTINEL) def __iter__(self): # pylint: disable=too-many-nested-blocks if self.already_iterated: raise RuntimeError( 'Can only iterate once over PrefetchingSourceSetIterable instance.') self.already_iterated = True # The invariants during execution are: # 1) A worker thread always posts the sentinel as the last thing it does # before exiting. # 2) We always wait for all sentinels and then join all threads. num_readers_finished = 0 try: while True: try: with self.read_counter: element = self.element_queue.get() if element is READER_THREAD_IS_DONE_SENTINEL: num_readers_finished += 1 if num_readers_finished == self.num_reader_threads: return else: if self.element_counter: self.element_counter.update_from(element) yield element self.element_counter.update_collect() else: yield element finally: if self.has_errored: raise self.reader_exceptions.get() except GeneratorExit: self.has_errored = True raise finally: while num_readers_finished < self.num_reader_threads: element = self.element_queue.get() if element is READER_THREAD_IS_DONE_SENTINEL: num_readers_finished += 1 for t in self.reader_threads: t.join() def get_iterator_fn_for_sources( sources, max_reader_threads=MAX_SOURCE_READER_THREADS, read_counter=None, element_counter=None): """Returns callable that returns iterator over elements for given sources.""" def _inner(): return iter( PrefetchingSourceSetIterable( sources, max_reader_threads=max_reader_threads, read_counter=read_counter, element_counter=element_counter)) return _inner class EmulatedIterable(abc.Iterable): """Emulates an iterable for a side input.""" def __init__(self, iterator_fn): self.iterator_fn = iterator_fn def __iter__(self): return self.iterator_fn() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/sideinputs_test.py0000644000175100017510000001320215206113032024460 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for side input utilities.""" # pytype: skip-file import logging import time import unittest import mock from apache_beam.coders import observable from apache_beam.runners.worker import sideinputs def strip_windows(iterator): return [wv.value for wv in iterator] class FakeSource(object): def __init__(self, items, notify_observers=False): self.items = items self._should_notify_observers = notify_observers def reader(self): return FakeSourceReader(self.items, self._should_notify_observers) class FakeSourceReader(observable.ObservableMixin): def __init__(self, items, notify_observers=False): super().__init__() self.items = items self.entered = False self.exited = False self._should_notify_observers = notify_observers def __iter__(self): if self._should_notify_observers: self.notify_observers(len(self.items), is_record_size=True) for item in self.items: yield item def __enter__(self): self.entered = True return self def __exit__(self, exception_type, exception_value, traceback): self.exited = True @property def returns_windowed_values(self): return False class PrefetchingSourceIteratorTest(unittest.TestCase): def test_single_source_iterator_fn(self): sources = [ FakeSource([0, 1, 2, 3, 4, 5]), ] iterator_fn = sideinputs.get_iterator_fn_for_sources( sources, max_reader_threads=2) assert list(strip_windows(iterator_fn())) == list(range(6)) def test_bytes_read_are_reported(self): mock_read_counter = mock.MagicMock() source_records = ['a', 'b', 'c', 'd'] sources = [ FakeSource(source_records, notify_observers=True), ] iterator_fn = sideinputs.get_iterator_fn_for_sources( sources, max_reader_threads=3, read_counter=mock_read_counter) assert list(strip_windows(iterator_fn())) == source_records mock_read_counter.add_bytes_read.assert_called_with(4) def test_multiple_sources_iterator_fn(self): sources = [ FakeSource([0]), FakeSource([1, 2, 3, 4, 5]), FakeSource([]), FakeSource([6, 7, 8, 9, 10]), ] iterator_fn = sideinputs.get_iterator_fn_for_sources( sources, max_reader_threads=3) assert sorted(strip_windows(iterator_fn())) == list(range(11)) def test_multiple_sources_single_reader_iterator_fn(self): sources = [ FakeSource([0]), FakeSource([1, 2, 3, 4, 5]), FakeSource([]), FakeSource([6, 7, 8, 9, 10]), ] iterator_fn = sideinputs.get_iterator_fn_for_sources( sources, max_reader_threads=1) assert list(strip_windows(iterator_fn())) == list(range(11)) def test_source_iterator_single_source_exception(self): class MyException(Exception): pass def exception_generator(): yield 0 raise MyException('I am an exception!') sources = [ FakeSource(exception_generator()), ] iterator_fn = sideinputs.get_iterator_fn_for_sources(sources) seen = set() with self.assertRaises(MyException): for value in iterator_fn(): seen.add(value.value) self.assertEqual(sorted(seen), [0]) def test_source_iterator_fn_exception(self): class MyException(Exception): pass def exception_generator(): yield 0 time.sleep(0.1) raise MyException('I am an exception!') def perpetual_generator(value): while True: yield value time.sleep(0.1) sources = [ FakeSource(perpetual_generator(1)), FakeSource(perpetual_generator(2)), FakeSource(perpetual_generator(3)), FakeSource(perpetual_generator(4)), FakeSource(exception_generator()), ] iterator_fn = sideinputs.get_iterator_fn_for_sources(sources) seen = set() with self.assertRaises(MyException): for value in iterator_fn(): seen.add(value.value) self.assertEqual(sorted(seen), list(range(5))) class EmulatedCollectionsTest(unittest.TestCase): def test_emulated_iterable(self): def _iterable_fn(): for i in range(10): yield i iterable = sideinputs.EmulatedIterable(_iterable_fn) # Check that multiple iterations are supported. for _ in range(0, 5): for i, j in enumerate(iterable): self.assertEqual(i, j) def test_large_iterable_values(self): # Here, we create a large collection that would be too big for memory- # constained test environments, but should be under the memory limit if # materialized one at a time. def _iterable_fn(): for i in range(10): yield ('%d' % i) * (200 * 1024 * 1024) iterable = sideinputs.EmulatedIterable(_iterable_fn) # Check that multiple iterations are supported. for _ in range(0, 3): for i, j in enumerate(iterable): self.assertEqual(('%d' % i) * (200 * 1024 * 1024), j) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/statecache.py0000644000175100017510000002777615206113032023363 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A module for caching state reads/writes in Beam applications.""" # pytype: skip-file # mypy: disallow-untyped-defs import collections import gc import logging import sys import threading import time import types import weakref from typing import Any from typing import Callable from typing import Union import objsize _LOGGER = logging.getLogger(__name__) _DEFAULT_WEIGHT = 8 _TYPES_TO_NOT_MEASURE = ( # Do not measure shared types type, types.ModuleType, types.FrameType, types.BuiltinFunctionType, # Do not measure lambdas as they typically share lots of state types.FunctionType, types.LambdaType, # Do not measure weak references as they will be deleted and not counted *weakref.ProxyTypes, weakref.ReferenceType) class WeightedValue(object): """Value type that stores corresponding weight. :arg value The value to be stored. :arg weight The associated weight of the value. If unspecified, the objects size will be used. """ def __init__(self, value: Any, weight: int) -> None: self._value = value if weight <= 0: raise ValueError( 'Expected weight to be > 0 for %s but received %d' % (value, weight)) self._weight = weight def weight(self) -> int: return self._weight def value(self) -> Any: return self._value class CacheAware(object): """Allows cache users to override what objects are measured.""" def __init__(self) -> None: pass def get_referents_for_cache(self) -> list[Any]: """Returns the list of objects accounted during cache measurement.""" raise NotImplementedError() def _safe_isinstance(obj: Any, type: Union[type, tuple[type, ...]]) -> bool: """ Return whether an object is an instance of a class or of a subclass thereof. See `isinstance()` for more information. Returns false on `isinstance()` failure. For example applying `isinstance()` on `weakref.proxy` objects attempts to dereference the proxy objects, which may yield an exception. See https://github.com/apache/beam/issues/23389 for additional details. """ try: return isinstance(obj, type) except Exception: return False def _size_func(obj: Any) -> int: """ Returns the size of the object or a default size if an error occurred during sizing. """ try: return sys.getsizeof(obj) except Exception as e: current_time = time.time() # Limit outputting this log so we don't spam the logs on these # occurrences. if _size_func.last_log_time + 300 < current_time: # type: ignore _LOGGER.warning( 'Failed to size %s of type %s. Note that this may ' 'impact cache sizing such that the cache is over ' 'utilized which may lead to out of memory errors.', obj, type(obj), exc_info=e) _size_func.last_log_time = current_time # type: ignore # Use an arbitrary default size that would account for some of the object # overhead. return _DEFAULT_WEIGHT _size_func.last_log_time = 0 # type: ignore def _get_referents_func(*objs: list[Any]) -> list[Any]: """Returns the list of objects accounted during cache measurement. Users can inherit CacheAware to override which referents should be used when measuring the deep size of the object. The default is to use gc.get_referents(*objs). """ rval = [] for obj in objs: if _safe_isinstance(obj, CacheAware): rval.extend(obj.get_referents_for_cache()) # type: ignore else: rval.extend(gc.get_referents(obj)) return rval def _filter_func(o: Any) -> bool: """ Filter out specific types from being measured. Note that we do want to measure the cost of weak references as they will only stay in scope as long as other code references them and will effectively be garbage collected as soon as there isn't a strong reference anymore. Note that we cannot use the default filter function due to isinstance raising an error on weakref.proxy types. See https://github.com/liran-funaro/objsize/issues/6 for additional details. """ return not _safe_isinstance(o, _TYPES_TO_NOT_MEASURE) def get_deep_size(*objs: Any) -> int: """Calculates the deep size of all the arguments in bytes.""" return objsize.get_deep_size( *objs, get_size_func=_size_func, get_referents_func=_get_referents_func, filter_func=_filter_func) class _LoadingValue(WeightedValue): """Allows concurrent users of the cache to wait for a value to be loaded.""" def __init__(self) -> None: super().__init__(None, 1) self._wait_event = threading.Event() def load(self, key: Any, loading_fn: Callable[[Any], Any]) -> None: try: self._value = loading_fn(key) except Exception as err: self._error = err finally: self._wait_event.set() def value(self) -> Any: self._wait_event.wait() err = getattr(self, "_error", None) if err: raise err return self._value class StateCache(object): """LRU cache for Beam state access, scoped by state key and cache_token. Assumes a bag state implementation. For a given key, caches a value and allows to a) peek at the cache (peek), returns the value for the provided key or None if it doesn't exist. Will never block. b) read from the cache (get), returns the value for the provided key or loads it using the supplied function. Multiple calls for the same key will block until the value is loaded. c) write to the cache (put), store the provided value overwriting any previous result d) invalidate a cached element (invalidate) removes the value from the cache for the provided key e) invalidate all cached elements (invalidate_all) The operations on the cache are thread-safe for use by multiple workers. :arg max_weight The maximum weight of entries to store in the cache in bytes. """ def __init__(self, max_weight: int) -> None: _LOGGER.info('Creating state cache with size %s', max_weight) self._max_weight = max_weight self._current_weight = 0 self._cache: collections.OrderedDict[ Any, WeightedValue] = collections.OrderedDict() self._hit_count = 0 self._miss_count = 0 self._evict_count = 0 self._load_time_ns = 0 self._load_count = 0 self._lock = threading.RLock() def peek(self, key: Any) -> Any: assert self.is_cache_enabled() with self._lock: value = self._cache.get(key, None) if value is None or _safe_isinstance(value, _LoadingValue): self._miss_count += 1 return None self._cache.move_to_end(key) self._hit_count += 1 return value.value() def get(self, key: Any, loading_fn: Callable[[Any], Any]) -> Any: assert self.is_cache_enabled() and callable(loading_fn) self._lock.acquire() value = self._cache.get(key, None) # Return the already cached value if value is not None: self._cache.move_to_end(key) self._hit_count += 1 self._lock.release() return value.value() # Load the value since it isn't in the cache. self._miss_count += 1 loading_value = _LoadingValue() self._cache[key] = loading_value self._current_weight += loading_value.weight() # Ensure that we unlock the lock while loading to allow for parallel gets self._lock.release() start_time_ns = time.time_ns() loading_value.load(key, loading_fn) elapsed_time_ns = time.time_ns() - start_time_ns try: value = loading_value.value() except Exception as err: # If loading failed then delete the value from the cache allowing for # the next lookup to possibly succeed. with self._lock: self._load_count += 1 self._load_time_ns += elapsed_time_ns # Don't remove values that have already been replaced with a different # value by a put/invalidate that occurred concurrently with the load. # The put/invalidate will have been responsible for updating the # cache weight appropriately already. old_value = self._cache.get(key, None) if old_value is not loading_value: raise err self._current_weight -= loading_value.weight() del self._cache[key] raise err # Replace the value in the cache with a weighted value now that the # loading has completed successfully. weight = get_deep_size(value) if weight <= 0: _LOGGER.warning( 'Expected object size to be >= 0 for %s but received %d.', value, weight) weight = 8 value = WeightedValue(value, weight) with self._lock: self._load_count += 1 self._load_time_ns += elapsed_time_ns # Don't replace values that have already been replaced with a different # value by a put/invalidate that occurred concurrently with the load. # The put/invalidate will have been responsible for updating the # cache weight appropriately already. old_value = self._cache.get(key, None) if old_value is not loading_value: return value.value() self._current_weight -= loading_value.weight() self._cache[key] = value self._current_weight += value.weight() while self._current_weight > self._max_weight: (_, weighted_value) = self._cache.popitem(last=False) self._current_weight -= weighted_value.weight() self._evict_count += 1 return value.value() def put(self, key: Any, value: Any) -> None: assert self.is_cache_enabled() if not _safe_isinstance(value, WeightedValue): weight = get_deep_size(value) if weight <= 0: _LOGGER.warning( 'Expected object size to be >= 0 for %s but received %d.', value, weight) weight = _DEFAULT_WEIGHT value = WeightedValue(value, weight) with self._lock: old_value = self._cache.pop(key, None) if old_value is not None: self._current_weight -= old_value.weight() self._cache[key] = value self._current_weight += value.weight() while self._current_weight > self._max_weight: (_, weighted_value) = self._cache.popitem(last=False) self._current_weight -= weighted_value.weight() self._evict_count += 1 def invalidate(self, key: Any) -> None: assert self.is_cache_enabled() with self._lock: weighted_value = self._cache.pop(key, None) if weighted_value is not None: self._current_weight -= weighted_value.weight() def invalidate_all(self) -> None: with self._lock: self._cache.clear() self._current_weight = 0 def describe_stats(self) -> str: with self._lock: request_count = self._hit_count + self._miss_count if request_count > 0: hit_ratio = 100.0 * self._hit_count / request_count else: hit_ratio = 100.0 return ( 'used/max %d/%d MB, hit %.2f%%, lookups %d, ' 'avg load time %.0f ns, loads %d, evictions %d') % ( self._current_weight >> 20, self._max_weight >> 20, hit_ratio, request_count, self._load_time_ns / self._load_count if self._load_count > 0 else 0, self._load_count, self._evict_count) def is_cache_enabled(self) -> bool: return self._max_weight > 0 def size(self) -> int: with self._lock: return len(self._cache) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/statecache_test.py0000644000175100017510000003403715206113032024406 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for state caching.""" # pytype: skip-file import logging import re import sys import threading import time import unittest import weakref import objsize from hamcrest import assert_that from hamcrest import contains_string from apache_beam.runners.worker.statecache import CacheAware from apache_beam.runners.worker.statecache import StateCache from apache_beam.runners.worker.statecache import WeightedValue from apache_beam.runners.worker.statecache import _LoadingValue from apache_beam.runners.worker.statecache import get_deep_size class StateCacheTest(unittest.TestCase): def test_weakref(self): test_value = WeightedValue('test', 10 << 20) class WeightedValueRef(): def __init__(self): self.ref = weakref.ref(test_value) cache = StateCache(5 << 20) wait_event = threading.Event() o = WeightedValueRef() cache.put('deep ref', o) # Ensure that the contents of the internal weak ref isn't sized self.assertIsNotNone(cache.peek('deep ref')) self.assertEqual( cache.describe_stats(), 'used/max 0/5 MB, hit 100.00%, lookups 1, avg load time 0 ns, loads 0, ' 'evictions 0') cache.invalidate_all() # Ensure that putting in a weakref doesn't fail regardless of whether # it is alive or not o_ref = weakref.ref(o, lambda value: wait_event.set()) cache.put('not deleted ref', o_ref) del o wait_event.wait() cache.put('deleted', o_ref) def test_weakref_proxy(self): test_value = WeightedValue('test', 10 << 20) class WeightedValueRef(): def __init__(self): self.ref = weakref.ref(test_value) cache = StateCache(5 << 20) wait_event = threading.Event() o = WeightedValueRef() cache.put('deep ref', o) # Ensure that the contents of the internal weak ref isn't sized self.assertIsNotNone(cache.peek('deep ref')) self.assertEqual( cache.describe_stats(), 'used/max 0/5 MB, hit 100.00%, lookups 1, avg load time 0 ns, loads 0, ' 'evictions 0') cache.invalidate_all() # Ensure that putting in a weakref doesn't fail regardless of whether # it is alive or not o_ref = weakref.proxy(o, lambda value: wait_event.set()) cache.put('not deleted', o_ref) del o wait_event.wait() cache.put('deleted', o_ref) def test_size_of_fails(self): class BadSizeOf(object): def __sizeof__(self): raise RuntimeError("TestRuntimeError") cache = StateCache(5 << 20) with self.assertLogs('apache_beam.runners.worker.statecache', level='WARNING') as context: cache.put('key', BadSizeOf()) self.assertEqual(1, len(context.output)) self.assertTrue('Failed to size' in context.output[0]) # Test that we don't spam the logs cache.put('key', BadSizeOf()) self.assertEqual(1, len(context.output)) def test_empty_cache_peek(self): cache = StateCache(5 << 20) self.assertEqual(cache.peek("key"), None) self.assertEqual( cache.describe_stats(), ( 'used/max 0/5 MB, hit 0.00%, lookups 1, ' 'avg load time 0 ns, loads 0, evictions 0')) def test_put_peek(self): cache = StateCache(5 << 20) cache.put("key", WeightedValue("value", 1 << 20)) self.assertEqual(cache.size(), 1) self.assertEqual(cache.peek("key"), "value") self.assertEqual(cache.peek("key2"), None) self.assertEqual( cache.describe_stats(), ( 'used/max 1/5 MB, hit 50.00%, lookups 2, ' 'avg load time 0 ns, loads 0, evictions 0')) def test_default_sized_put(self): cache = StateCache(5 << 20) cache.put("key", bytearray(1 << 20)) cache.put("key2", bytearray(1 << 20)) cache.put("key3", bytearray(1 << 20)) self.assertEqual(cache.peek("key3"), bytearray(1 << 20)) cache.put("key4", bytearray(1 << 20)) cache.put("key5", bytearray(1 << 20)) # note that each byte array instance takes slightly over 1 MB which is why # these 5 byte arrays can't all be stored in the cache causing a single # eviction self.assertEqual( cache.describe_stats(), ( 'used/max 4/5 MB, hit 100.00%, lookups 1, ' 'avg load time 0 ns, loads 0, evictions 1')) def test_max_size(self): cache = StateCache(2 << 20) cache.put("key", WeightedValue("value", 1 << 20)) cache.put("key2", WeightedValue("value2", 1 << 20)) self.assertEqual(cache.size(), 2) cache.put("key3", WeightedValue("value3", 1 << 20)) self.assertEqual(cache.size(), 2) self.assertEqual( cache.describe_stats(), ( 'used/max 2/2 MB, hit 100.00%, lookups 0, ' 'avg load time 0 ns, loads 0, evictions 1')) def test_invalidate_all(self): cache = StateCache(5 << 20) cache.put("key", WeightedValue("value", 1 << 20)) cache.put("key2", WeightedValue("value2", 1 << 20)) self.assertEqual(cache.size(), 2) cache.invalidate_all() self.assertEqual(cache.size(), 0) self.assertEqual(cache.peek("key"), None) self.assertEqual(cache.peek("key2"), None) self.assertEqual( cache.describe_stats(), ( 'used/max 0/5 MB, hit 0.00%, lookups 2, ' 'avg load time 0 ns, loads 0, evictions 0')) def test_lru(self): cache = StateCache(5 << 20) cache.put("key", WeightedValue("value", 1 << 20)) cache.put("key2", WeightedValue("value2", 1 << 20)) cache.put("key3", WeightedValue("value0", 1 << 20)) cache.put("key3", WeightedValue("value3", 1 << 20)) cache.put("key4", WeightedValue("value4", 1 << 20)) cache.put("key5", WeightedValue("value0", 1 << 20)) cache.put("key5", WeightedValue(["value5"], 1 << 20)) self.assertEqual(cache.size(), 5) self.assertEqual(cache.peek("key"), "value") self.assertEqual(cache.peek("key2"), "value2") self.assertEqual(cache.peek("key3"), "value3") self.assertEqual(cache.peek("key4"), "value4") self.assertEqual(cache.peek("key5"), ["value5"]) # insert another key to trigger cache eviction cache.put("key6", WeightedValue("value6", 1 << 20)) self.assertEqual(cache.size(), 5) # least recently used key should be gone ("key") self.assertEqual(cache.peek("key"), None) # trigger a read on "key2" cache.peek("key2") # insert another key to trigger cache eviction cache.put("key7", WeightedValue("value7", 1 << 20)) self.assertEqual(cache.size(), 5) # least recently used key should be gone ("key3") self.assertEqual(cache.peek("key3"), None) # insert another key to trigger cache eviction cache.put("key8", WeightedValue("put", 1 << 20)) self.assertEqual(cache.size(), 5) # insert another key to trigger cache eviction cache.put("key9", WeightedValue("value8", 1 << 20)) self.assertEqual(cache.size(), 5) # least recently used key should be gone ("key4") self.assertEqual(cache.peek("key4"), None) # make "key5" used by writing to it cache.put("key5", WeightedValue("val", 1 << 20)) # least recently used key should be gone ("key6") self.assertEqual(cache.peek("key6"), None) self.assertEqual( cache.describe_stats(), ( 'used/max 5/5 MB, hit 60.00%, lookups 10, ' 'avg load time 0 ns, loads 0, evictions 5')) def test_get(self): def check_key(key): self.assertEqual(key, "key") time.sleep(0.5) return "value" def raise_exception(key): time.sleep(0.5) raise Exception("TestException") cache = StateCache(5 << 20) self.assertEqual("value", cache.get("key", check_key)) with cache._lock: self.assertFalse(isinstance(cache._cache["key"], _LoadingValue)) self.assertEqual("value", cache.peek("key")) cache.invalidate_all() with self.assertRaisesRegex(Exception, "TestException"): cache.get("key", raise_exception) # The cache should not have the value after the failing load causing # check_key to load the value. self.assertEqual("value", cache.get("key", check_key)) with cache._lock: self.assertFalse(isinstance(cache._cache["key"], _LoadingValue)) self.assertEqual("value", cache.peek("key")) assert_that(cache.describe_stats(), contains_string(", loads 3,")) load_time_ns = re.search( ", avg load time (.+) ns,", cache.describe_stats()).group(1) # Load time should be larger then the sleep time and less than 2x sleep time self.assertGreater(int(load_time_ns), 0.5 * 1_000_000_000) self.assertLess(int(load_time_ns), 1_000_000_000) def test_concurrent_get_waits(self): event = threading.Semaphore(0) threads_running = threading.Barrier(3) def wait_for_event(key): with cache._lock: self.assertTrue(isinstance(cache._cache["key"], _LoadingValue)) event.release() return "value" cache = StateCache(5 << 20) def load_key(output): threads_running.wait() output["value"] = cache.get("key", wait_for_event) output["time"] = time.time_ns() t1_output = {} t1 = threading.Thread( target=load_key, args=(t1_output, )) t1.start() t2_output = {} t2 = threading.Thread( target=load_key, args=(t2_output, )) t2.start() # Wait for both threads to start threads_running.wait() # Record the time and wait for the load to start current_time_ns = time.time_ns() event.acquire() t1.join() t2.join() # Ensure that only one thread did the loading and not both by checking that # the semaphore was only released once self.assertFalse(event.acquire(blocking=False)) # Ensure that the load time is greater than the set time ensuring that # both loads had to wait for the event self.assertLessEqual(current_time_ns, t1_output["time"]) self.assertLessEqual(current_time_ns, t2_output["time"]) self.assertEqual("value", t1_output["value"]) self.assertEqual("value", t2_output["value"]) self.assertEqual("value", cache.peek("key")) def test_concurrent_get_superseded_by_put(self): load_happening = threading.Event() finish_loading = threading.Event() def wait_for_event(key): load_happening.set() finish_loading.wait() return "value" cache = StateCache(5 << 20) def load_key(output): output["value"] = cache.get("key", wait_for_event) t1_output = {} t1 = threading.Thread( target=load_key, args=(t1_output, )) t1.start() # Wait for the load to start, update the key, and then let the load finish load_happening.wait() cache.put("key", "value2") finish_loading.set() t1.join() # Ensure that the original value is loaded and returned and not the # updated value self.assertEqual("value", t1_output["value"]) # Ensure that the updated value supersedes the loaded value. self.assertEqual("value2", cache.peek("key")) def test_is_cached_enabled(self): cache = StateCache(1 << 20) self.assertEqual(cache.is_cache_enabled(), True) self.assertEqual( cache.describe_stats(), ( 'used/max 0/1 MB, hit 100.00%, lookups 0, ' 'avg load time 0 ns, loads 0, evictions 0')) cache = StateCache(0) self.assertEqual(cache.is_cache_enabled(), False) self.assertEqual( cache.describe_stats(), ( 'used/max 0/0 MB, hit 100.00%, lookups 0, ' 'avg load time 0 ns, loads 0, evictions 0')) def test_get_referents_for_cache(self): class GetReferentsForCache(CacheAware): def __init__(self): self.measure_me = bytearray(1 << 20) self.ignore_me = bytearray(2 << 20) def get_referents_for_cache(self): return [self.measure_me] cache = StateCache(5 << 20) cache.put("key", GetReferentsForCache()) self.assertEqual( cache.describe_stats(), ( 'used/max 1/5 MB, hit 100.00%, lookups 0, ' 'avg load time 0 ns, loads 0, evictions 0')) def test_get_deep_size_builtin_objects(self): """ `statecache.get_deep_copy` should work same with objsize unless the `objs` has `CacheAware` or a filtered object. They should return the same size for built-in objects. """ primitive_test_objects = [ 1, # int 2.0, # float 1 + 1j, # complex True, # bool 'hello,world', # str b'\00\01\02', # bytes ] collection_test_objects = [ [3, 4, 5], # list (6, 7), # tuple {'a', 'b', 'c'}, # set { 'k': 8, 'l': 9 }, # dict ] for obj in primitive_test_objects: self.assertEqual( get_deep_size(obj), objsize.get_deep_size(obj), f'different size for obj: `{obj}`, type: {type(obj)}') self.assertEqual( get_deep_size(obj), sys.getsizeof(obj), f'different size for obj: `{obj}`, type: {type(obj)}') for obj in collection_test_objects: self.assertEqual( get_deep_size(obj), objsize.get_deep_size(obj), f'different size for obj: `{obj}`, type: {type(obj)}') def test_current_weight_between_get_and_put(self): value = 1234567 get_cache = StateCache(100) get_cache.get("key", lambda k: value) put_cache = StateCache(100) put_cache.put("key", value) self.assertEqual(get_cache._current_weight, put_cache._current_weight) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/statesampler.py0000644000175100017510000001222015206113032023735 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import contextlib import threading from typing import TYPE_CHECKING from typing import NamedTuple from typing import Optional from typing import Union from apache_beam.runners import common from apache_beam.utils.counters import Counter from apache_beam.utils.counters import CounterFactory from apache_beam.utils.counters import CounterName try: from apache_beam.runners.worker import statesampler_fast as statesampler_impl # type: ignore FAST_SAMPLER = True except ImportError: from apache_beam.runners.worker import statesampler_slow as statesampler_impl FAST_SAMPLER = False if TYPE_CHECKING: from apache_beam.metrics.execution import MetricsContainer _STATE_SAMPLERS = threading.local() def set_current_tracker(tracker): _STATE_SAMPLERS.tracker = tracker def get_current_tracker(): try: return _STATE_SAMPLERS.tracker except AttributeError: return None _INSTRUCTION_IDS = threading.local() def get_current_instruction_id(): try: return _INSTRUCTION_IDS.instruction_id except AttributeError: return None @contextlib.contextmanager def instruction_id(id): try: _INSTRUCTION_IDS.instruction_id = id yield finally: _INSTRUCTION_IDS.instruction_id = None def for_test(): set_current_tracker(StateSampler('test', CounterFactory())) return get_current_tracker() StateSamplerInfo = NamedTuple( 'StateSamplerInfo', [('state_name', CounterName), ('transition_count', int), ('time_since_transition', int), ('tracked_thread', Optional[threading.Thread])]) # Default period for sampling current state of pipeline execution. DEFAULT_SAMPLING_PERIOD_MS = 200 class StateSampler(statesampler_impl.StateSampler): def __init__( self, prefix: str, counter_factory, sampling_period_ms=DEFAULT_SAMPLING_PERIOD_MS): self._prefix = prefix self._counter_factory = counter_factory self._states_by_name: dict[CounterName, statesampler_impl.ScopedState] = {} self.sampling_period_ms = sampling_period_ms self.tracked_thread: Optional[threading.Thread] = None self.finished = False self.started = False super().__init__(sampling_period_ms) @property def stage_name(self) -> str: return self._prefix def stop(self) -> None: set_current_tracker(None) super().stop() def stop_if_still_running(self) -> None: if self.started and not self.finished: self.stop() def start(self) -> None: self.tracked_thread = threading.current_thread() set_current_tracker(self) super().start() self.started = True def get_info(self) -> StateSamplerInfo: """Returns StateSamplerInfo with transition statistics.""" return StateSamplerInfo( self.current_state().name, self.state_transition_count, self.time_since_transition, self.tracked_thread) def scoped_state( self, name_context: Union[str, 'common.NameContext'], state_name: str, io_target=None, metrics_container: Optional['MetricsContainer'] = None ) -> statesampler_impl.ScopedState: """Returns a ScopedState object associated to a Step and a State. Args: name_context: common.NameContext. It is the step name information. state_name: str. It is the state name (e.g. process / start / finish). io_target: metrics_container: MetricsContainer. The step's metrics container. Returns: A ScopedState that keeps the execution context and is able to switch it for the execution thread. """ if not isinstance(name_context, common.NameContext): name_context = common.NameContext(name_context) counter_name = CounterName( state_name + '-msecs', stage_name=self._prefix, step_name=name_context.metrics_name(), io_target=io_target) if counter_name in self._states_by_name: return self._states_by_name[counter_name] else: output_counter = self._counter_factory.get_counter( counter_name, Counter.SUM) self._states_by_name[counter_name] = super()._scoped_state( counter_name, name_context, output_counter, metrics_container) return self._states_by_name[counter_name] def commit_counters(self) -> None: """Updates output counters with latest state statistics.""" for state in self._states_by_name.values(): state_msecs = int(1e-6 * state.nsecs) state.counter.update(state_msecs - state.counter.value()) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/statesampler_fast.pxd0000644000175100017510000000432715206113032025126 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # cimport cython from apache_beam.metrics.execution cimport MetricsContainer from cpython cimport pythread from libc.stdint cimport int32_t, int64_t cdef class StateSampler(object): """Tracks time spent in states during pipeline execution.""" cdef int _sampling_period_ms cdef int _sampling_period_ms_start cdef double _sampling_period_ratio cdef list scoped_states_by_index cdef public bint started cdef public bint finished cdef object sampling_thread # This lock guards members that are shared between threads, specificaly # finished, scoped_states_by_index, and the nsecs field of each state therein. cdef pythread.PyThread_type_lock lock cdef public int64_t state_transition_count cdef public int64_t time_since_transition cdef int32_t current_state_index cpdef ScopedState current_state(self) cdef inline ScopedState current_state_c(self) cpdef _scoped_state( self, counter_name, name_context, output_counter, metrics_container) cdef class ScopedState(object): """Context manager class managing transitions for a given sampler state.""" cdef readonly StateSampler sampler cdef readonly int32_t state_index cdef readonly object counter cdef readonly object name cdef readonly object name_context cdef readonly int64_t _nsecs cdef int32_t old_state_index cdef readonly MetricsContainer metrics_container cpdef __enter__(self) cpdef __exit__(self, unused_exc_type, unused_exc_value, unused_traceback) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/statesampler_fast.pyx0000644000175100017510000002277715206113032025164 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """State sampler for tracking time spent in execution steps. The state sampler profiles the time spent in each step of a pipeline. Operations (defined in executor.py) which are executed as part of a MapTask are instrumented with context managers provided by StateSampler.scoped_state(). These context managers change the internal state of the StateSampler during each relevant Operation's .start(), .process() and .finish() methods. State is sampled by a raw C thread, not holding the Python Global Interpreter Lock, which queries the StateSampler's internal state at a defined sampling frequency. In a common example, a ReadOperation during its .start() method reads an element and calls a DoOperation's .process() method, which can call a WriteOperation's .process() method. Each element processed causes the current state to transition between these states of different Operations. Each time the sampling thread queries the current state, the time spent since the previous sample is attributed to that state and accumulated. Over time, this allows a granular runtime profile to be produced. """ import threading from apache_beam.utils.counters import CounterName from apache_beam.metrics.execution cimport MetricsContainer cimport cython from cpython cimport pythread from libc cimport math from libc.stdint cimport int32_t, int64_t cdef extern from "Python.h": # This typically requires the GIL, but we synchronize the list modifications # we use this on via our own lock. cdef void* PyList_GET_ITEM(list, Py_ssize_t index) nogil cdef extern from "crossplatform_unistd.h" nogil: void usleep(int) cdef extern from "crossplatform_time.h" nogil: struct timespec: long tv_sec # seconds long tv_nsec # nanoseconds int clock_gettime(int clock_id, timespec *result) cdef inline int64_t get_nsec_time() noexcept nogil: """Get current time as microseconds since Unix epoch.""" cdef timespec current_time # First argument value of 0 corresponds to CLOCK_REALTIME. clock_gettime(0, ¤t_time) return ( ( current_time.tv_sec) * 1000000000 + # second to nanoseconds current_time.tv_nsec) cdef class StateSampler(object): """Tracks time spent in states during pipeline execution.""" def __init__(self, sampling_period_ms, sampling_period_ms_start=None, sampling_period_ratio=1.2): self._sampling_period_ms = sampling_period_ms # Slowly ramp up to avoid excessive waiting for short stages, as well # as more precise information in that case. self._sampling_period_ms_start = ( sampling_period_ms_start or max(1, sampling_period_ms // 100)) self._sampling_period_ratio = sampling_period_ratio self.started = False self.finished = False self.lock = pythread.PyThread_allocate_lock() self.current_state_index = 0 self.time_since_transition = 0 self.state_transition_count = 0 unknown_state = ScopedState(self, CounterName('unknown'), None, self.current_state_index, None, None) pythread.PyThread_acquire_lock(self.lock, pythread.WAIT_LOCK) self.scoped_states_by_index = [unknown_state] pythread.PyThread_release_lock(self.lock) # Assert that the compiler correctly aligned the current_state field. This # is necessary for reads and writes to this variable to be atomic across # threads without additional synchronization. # States are referenced via an index rather than, say, a pointer because # of better support for 32-bit atomic reads and writes. assert ( &self.current_state_index) % sizeof(int32_t) == 0, ( 'Address of StateSampler.current_state_index is not word-aligned.') def __dealloc__(self): pythread.PyThread_free_lock(self.lock) def run(self): cdef int64_t last_nsecs = get_nsec_time() cdef int64_t elapsed_nsecs cdef int64_t latest_transition_count = self.state_transition_count cdef int64_t sampling_period_us = self._sampling_period_ms_start * 1000 with nogil: while True: usleep(sampling_period_us) sampling_period_us = math.fmin( sampling_period_us * self._sampling_period_ratio, self._sampling_period_ms * 1000) pythread.PyThread_acquire_lock(self.lock, pythread.WAIT_LOCK) try: if self.finished: break elapsed_nsecs = get_nsec_time() - last_nsecs # Take an address as we can't create a reference to the scope # without the GIL. nsecs_ptr = &(PyList_GET_ITEM( self.scoped_states_by_index, self.current_state_index))._nsecs nsecs_ptr[0] += elapsed_nsecs if latest_transition_count != self.state_transition_count: self.time_since_transition = 0 latest_transition_count = self.state_transition_count self.time_since_transition += elapsed_nsecs last_nsecs += elapsed_nsecs finally: pythread.PyThread_release_lock(self.lock) def start(self): assert not self.started self.sampling_thread = threading.Thread(target=self.run) self.sampling_thread.start() def stop(self): assert not self.finished pythread.PyThread_acquire_lock(self.lock, pythread.WAIT_LOCK) self.finished = True pythread.PyThread_release_lock(self.lock) # May have to wait up to sampling_period_ms, but the platform-independent # pythread doesn't support conditions. self.sampling_thread.join() def reset(self): for state in self.scoped_states_by_index: (state)._nsecs = 0 self.started = self.finished = False cpdef ScopedState current_state(self): return self.current_state_c() cdef inline ScopedState current_state_c(self): # Faster than cpdef due to self always being a Python subclass. return self.scoped_states_by_index[self.current_state_index] cpdef _scoped_state(self, counter_name, name_context, output_counter, metrics_container): """Returns a context manager managing transitions for a given state. Args: counter_name: A CounterName object with information about the execution state. output_counter: A Beam Counter to which msecs are committed for reporting. metrics_container: A MetricsContainer for the current step. Returns: A ScopedState for the set of step-state-io_target. """ new_state_index = len(self.scoped_states_by_index) scoped_state = ScopedState(self, counter_name, name_context, new_state_index, output_counter, metrics_container) # Both scoped_states_by_index and scoped_state.nsecs are accessed # by the sampling thread; initialize them under the lock. pythread.PyThread_acquire_lock(self.lock, pythread.WAIT_LOCK) self.scoped_states_by_index.append(scoped_state) scoped_state._nsecs = 0 pythread.PyThread_release_lock(self.lock) return scoped_state def update_metric(self, typed_metric_name, value): # Each of these is a cdef lookup. metrics_container = self.current_state_c().metrics_container if metrics_container is not None: metrics_container.get_metric_cell(typed_metric_name).update(value) cdef class ScopedState(object): """Context manager class managing transitions for a given sampler state.""" def __init__(self, sampler, name, step_name_context, state_index, counter, metrics_container): self.sampler = sampler self.name = name self.name_context = step_name_context self.state_index = state_index self.counter = counter self.metrics_container = metrics_container @property def nsecs(self): return self._nsecs def sampled_seconds(self): return 1e-9 * self.nsecs def sampled_msecs_int(self): return int(1e-6 * self.nsecs) def __repr__(self): return "ScopedState[%s, %s]" % (self.name, self.nsecs) cpdef __enter__(self): self.old_state_index = self.sampler.current_state_index pythread.PyThread_acquire_lock(self.sampler.lock, pythread.WAIT_LOCK) self.sampler.current_state_index = self.state_index self.sampler.state_transition_count += 1 pythread.PyThread_release_lock(self.sampler.lock) cpdef __exit__(self, unused_exc_type, unused_exc_value, unused_traceback): pythread.PyThread_acquire_lock(self.sampler.lock, pythread.WAIT_LOCK) self.sampler.current_state_index = self.old_state_index self.sampler.state_transition_count += 1 pythread.PyThread_release_lock(self.sampler.lock) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/statesampler_slow.py0000644000175100017510000000612215206113032025005 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file from typing import Optional from apache_beam.runners import common from apache_beam.utils import counters class StateSampler(object): def __init__(self, sampling_period_ms): self._state_stack = [ ScopedState(self, counters.CounterName('unknown'), None) ] self.state_transition_count = 0 self.time_since_transition = 0 def current_state(self) -> 'ScopedState': """Returns the current execution state. This operation is not thread safe, and should only be called from the execution thread.""" return self._state_stack[-1] def _scoped_state( self, counter_name: counters.CounterName, name_context: 'common.NameContext', output_counter, metrics_container=None) -> 'ScopedState': assert isinstance(name_context, common.NameContext) return ScopedState( self, counter_name, name_context, output_counter, metrics_container) def update_metric(self, typed_metric_name, value): metrics_container = self.current_state().metrics_container if metrics_container is not None: metrics_container.get_metric_cell(typed_metric_name).update(value) def _enter_state(self, state: 'ScopedState') -> None: self.state_transition_count += 1 self._state_stack.append(state) def _exit_state(self) -> None: self.state_transition_count += 1 self._state_stack.pop() def start(self) -> None: # Sampling not yet supported. Only state tracking at the moment. pass def stop(self) -> None: pass def reset(self) -> None: pass class ScopedState(object): def __init__( self, sampler: StateSampler, name: counters.CounterName, step_name_context: Optional['common.NameContext'], counter=None, metrics_container=None): self.state_sampler = sampler self.name = name self.name_context = step_name_context self.counter = counter self.nsecs = 0 self.metrics_container = metrics_container def sampled_seconds(self) -> float: return 1e-9 * self.nsecs def sampled_msecs_int(self) -> int: return int(1e-6 * self.nsecs) def __repr__(self): return "ScopedState[%s, %s]" % (self.name, self.nsecs) def __enter__(self): self.state_sampler._enter_state(self) def __exit__(self, exc_type, exc_value, traceback): self.state_sampler._exit_state() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/statesampler_test.py0000644000175100017510000002560215206113032025004 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for state sampler.""" # pytype: skip-file import logging import time import unittest from unittest import mock from unittest.mock import Mock from unittest.mock import patch from tenacity import retry from tenacity import stop_after_attempt from apache_beam.internal import pickler from apache_beam.runners import common from apache_beam.runners.worker import operation_specs from apache_beam.runners.worker import operations from apache_beam.runners.worker import statesampler from apache_beam.transforms import core from apache_beam.transforms import userstate from apache_beam.transforms.core import GlobalWindows from apache_beam.transforms.core import Windowing from apache_beam.transforms.window import GlobalWindow from apache_beam.utils.counters import CounterFactory from apache_beam.utils.counters import CounterName from apache_beam.utils.windowed_value import PaneInfo _LOGGER = logging.getLogger(__name__) class TimerDoFn(core.DoFn): TIMER_SPEC = userstate.TimerSpec('timer', userstate.TimeDomain.WATERMARK) def __init__(self, sleep_duration_s=0): self._sleep_duration_s = sleep_duration_s @userstate.on_timer(TIMER_SPEC) def on_timer_f(self): if self._sleep_duration_s: time.sleep(self._sleep_duration_s) class ExceptionTimerDoFn(core.DoFn): """A DoFn that raises an exception when its timer fires.""" TIMER_SPEC = userstate.TimerSpec('ts-timer', userstate.TimeDomain.WATERMARK) def __init__(self, sleep_duration_s=0): self._sleep_duration_s = sleep_duration_s @userstate.on_timer(TIMER_SPEC) def on_timer_f(self): if self._sleep_duration_s: time.sleep(self._sleep_duration_s) raise RuntimeError("Test exception from timer") class StateSamplerTest(unittest.TestCase): # Due to somewhat non-deterministic nature of state sampling and sleep, # this test is flaky when state duration is low. # Since increasing state duration significantly would also slow down # the test suite, we are retrying twice on failure as a mitigation. @retry(reraise=True, stop=stop_after_attempt(3)) def test_basic_sampler(self): # Set up state sampler. counter_factory = CounterFactory() sampler = statesampler.StateSampler( 'basic', counter_factory, sampling_period_ms=1) # Duration of the fastest state. Total test duration is 6 times longer. state_duration_ms = 1000 margin_of_error = 0.25 # Run basic workload transitioning between 3 states. sampler.start() with sampler.scoped_state('step1', 'statea'): time.sleep(state_duration_ms / 1000) self.assertEqual( sampler.current_state().name, CounterName('statea-msecs', step_name='step1', stage_name='basic')) with sampler.scoped_state('step1', 'stateb'): time.sleep(state_duration_ms / 1000) self.assertEqual( sampler.current_state().name, CounterName('stateb-msecs', step_name='step1', stage_name='basic')) with sampler.scoped_state('step1', 'statec'): time.sleep(3 * state_duration_ms / 1000) self.assertEqual( sampler.current_state().name, CounterName( 'statec-msecs', step_name='step1', stage_name='basic')) time.sleep(state_duration_ms / 1000) sampler.stop() sampler.commit_counters() if not statesampler.FAST_SAMPLER: # The slow sampler does not implement sampling, so we won't test it. return # Test that sampled state timings are close to their expected values. # yapf: disable expected_counter_values = { CounterName('statea-msecs', step_name='step1', stage_name='basic'): state_duration_ms, CounterName('stateb-msecs', step_name='step1', stage_name='basic'): 2 * state_duration_ms, CounterName('statec-msecs', step_name='step1', stage_name='basic'): 3 * state_duration_ms, } # yapf: enable for counter in counter_factory.get_counters(): self.assertIn(counter.name, expected_counter_values) expected_value = expected_counter_values[counter.name] actual_value = counter.value() deviation = float(abs(actual_value - expected_value)) / expected_value _LOGGER.info('Sampling deviation from expectation: %f', deviation) self.assertGreater(actual_value, expected_value * (1.0 - margin_of_error)) self.assertLess(actual_value, expected_value * (1.0 + margin_of_error)) # TODO: This test is flaky when it is run under load. A better solution # would be to change the test structure to not depend on specific timings. @retry(reraise=True, stop=stop_after_attempt(3)) def test_sampler_transition_overhead(self): # Set up state sampler. counter_factory = CounterFactory() sampler = statesampler.StateSampler( 'overhead-', counter_factory, sampling_period_ms=10) # Run basic workload transitioning between 3 states. state_a = sampler.scoped_state('step1', 'statea') state_b = sampler.scoped_state('step1', 'stateb') state_c = sampler.scoped_state('step1', 'statec') start_time = time.time() sampler.start() for _ in range(100000): with state_a: with state_b: for _ in range(10): with state_c: pass sampler.stop() elapsed_time = time.time() - start_time state_transition_count = sampler.get_info().transition_count overhead_us = 1000000.0 * elapsed_time / state_transition_count _LOGGER.info('Overhead per transition: %fus', overhead_us) # Conservative upper bound on overhead in microseconds (we expect this to # take 0.17us when compiled in opt mode or 0.48 us when compiled with in # debug mode). self.assertLess(overhead_us, 20.0) @retry(reraise=True, stop=stop_after_attempt(3)) # Patch the problematic function to return the correct timer spec @patch('apache_beam.transforms.userstate.get_dofn_specs') def test_do_operation_process_timer(self, mock_get_dofn_specs): fn = TimerDoFn() mock_get_dofn_specs.return_value = ([], [fn.TIMER_SPEC]) if not statesampler.FAST_SAMPLER: self.skipTest('DoOperation test requires FAST_SAMPLER') state_duration_ms = 200 margin_of_error = 0.75 counter_factory = CounterFactory() sampler = statesampler.StateSampler( 'test_do_op', counter_factory, sampling_period_ms=1) fn_for_spec = TimerDoFn(sleep_duration_s=state_duration_ms / 1000.0) spec = operation_specs.WorkerDoFn( serialized_fn=pickler.dumps( (fn_for_spec, [], {}, [], Windowing(GlobalWindows()))), output_tags=[], input=None, side_inputs=[], output_coders=[]) mock_user_state_context = mock.MagicMock() op = operations.DoOperation( common.NameContext('step1'), spec, counter_factory, sampler, user_state_context=mock_user_state_context) op.setup() timer_data = Mock() timer_data.user_key = None timer_data.windows = [GlobalWindow()] timer_data.fire_timestamp = 0 timer_data.paneinfo = PaneInfo( is_first=False, is_last=False, timing=0, index=0, nonspeculative_index=0) timer_data.dynamic_timer_tag = '' sampler.start() op.process_timer('ts-timer', timer_data=timer_data) sampler.stop() sampler.commit_counters() expected_name = CounterName( 'process-timers-msecs', step_name='step1', stage_name='test_do_op') found_counter = None for counter in counter_factory.get_counters(): if counter.name == expected_name: found_counter = counter break self.assertIsNotNone( found_counter, f"Expected counter '{expected_name}' to be created.") actual_value = found_counter.value() logging.info("Actual value %d", actual_value) self.assertGreater( actual_value, state_duration_ms * (1.0 - margin_of_error)) @retry(reraise=True, stop=stop_after_attempt(3)) @patch('apache_beam.runners.worker.operations.userstate.get_dofn_specs') def test_do_operation_process_timer_with_exception(self, mock_get_dofn_specs): fn = ExceptionTimerDoFn() mock_get_dofn_specs.return_value = ([], [fn.TIMER_SPEC]) if not statesampler.FAST_SAMPLER: self.skipTest('DoOperation test requires FAST_SAMPLER') state_duration_ms = 200 margin_of_error = 0.50 counter_factory = CounterFactory() sampler = statesampler.StateSampler( 'test_do_op_exception', counter_factory, sampling_period_ms=1) fn_for_spec = ExceptionTimerDoFn( sleep_duration_s=state_duration_ms / 1000.0) spec = operation_specs.WorkerDoFn( serialized_fn=pickler.dumps( (fn_for_spec, [], {}, [], Windowing(GlobalWindows()))), output_tags=[], input=None, side_inputs=[], output_coders=[]) mock_user_state_context = mock.MagicMock() op = operations.DoOperation( common.NameContext('step1'), spec, counter_factory, sampler, user_state_context=mock_user_state_context) op.setup() timer_data = Mock() timer_data.user_key = None timer_data.windows = [GlobalWindow()] timer_data.fire_timestamp = 0 timer_data.paneinfo = PaneInfo( is_first=False, is_last=False, timing=0, index=0, nonspeculative_index=0) timer_data.dynamic_timer_tag = '' sampler.start() # Assert that the expected exception is raised with self.assertRaises(RuntimeError): op.process_timer('ts-ts-timer', timer_data=timer_data) sampler.stop() sampler.commit_counters() expected_name = CounterName( 'process-timers-msecs', step_name='step1', stage_name='test_do_op_exception') found_counter = None for counter in counter_factory.get_counters(): if counter.name == expected_name: found_counter = counter break self.assertIsNotNone( found_counter, f"Expected counter '{expected_name}' to be created.") actual_value = found_counter.value() self.assertGreater( actual_value, state_duration_ms * (1.0 - margin_of_error)) _LOGGER.info("Exception test finished successfully.") if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/worker_id_interceptor.py0000644000175100017510000000541615206113032025645 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Client Interceptor to inject worker_id""" # pytype: skip-file import collections import os from typing import Optional import grpc class _ClientCallDetails(collections.namedtuple( '_ClientCallDetails', ('method', 'timeout', 'metadata', 'credentials')), grpc.ClientCallDetails): pass class WorkerIdInterceptor(grpc.UnaryUnaryClientInterceptor, grpc.StreamStreamClientInterceptor): # TODO: (BEAM-3904) Removed defaulting to UUID when worker_id is not present # and throw exception in worker_id_interceptor.py after we have rolled out # the corresponding container changes. # Unique worker Id for this worker. _worker_id = os.environ.get('WORKER_ID') def __init__(self, worker_id: Optional[str] = None) -> None: if worker_id: self._worker_id = worker_id def intercept_unary_unary(self, continuation, client_call_details, request): return self._intercept(continuation, client_call_details, request) def intercept_unary_stream(self, continuation, client_call_details, request): return self._intercept(continuation, client_call_details, request) def intercept_stream_unary(self, continuation, client_call_details, request): return self._intercept(continuation, client_call_details, request) def intercept_stream_stream( self, continuation, client_call_details, request_iterator): return self._intercept(continuation, client_call_details, request_iterator) def _intercept(self, continuation, client_call_details, request): metadata = [] if client_call_details.metadata is not None: metadata = list(client_call_details.metadata) if 'worker_id' in metadata: raise RuntimeError('Header metadata already has a worker_id.') metadata.append(('worker_id', self._worker_id)) new_client_details = _ClientCallDetails( client_call_details.method, client_call_details.timeout, metadata, client_call_details.credentials) return continuation(new_client_details, request) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/worker_id_interceptor_test.py0000644000175100017510000000447215206113032026705 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for WorkerIdInterceptor""" # pytype: skip-file import collections import logging import unittest import grpc from apache_beam.runners.worker.worker_id_interceptor import WorkerIdInterceptor class _ClientCallDetails(collections.namedtuple( '_ClientCallDetails', ('method', 'timeout', 'metadata', 'credentials')), grpc.ClientCallDetails): pass class WorkerIdInterceptorTest(unittest.TestCase): def test_worker_id_insertion(self): worker_id_key = 'worker_id' headers_holder = {} def continuation(client_details, request_iterator): headers_holder.update( {worker_id_key: dict(client_details.metadata).get(worker_id_key)}) WorkerIdInterceptor._worker_id = 'my_worker_id' WorkerIdInterceptor().intercept_stream_stream( continuation, _ClientCallDetails(None, None, None, None), []) self.assertEqual( headers_holder[worker_id_key], 'my_worker_id', 'worker_id_key not set') def test_failure_when_worker_id_exists(self): worker_id_key = 'worker_id' headers_holder = {} def continuation(client_details, request_iterator): headers_holder.update( {worker_id_key: dict(client_details.metadata).get(worker_id_key)}) WorkerIdInterceptor._worker_id = 'my_worker_id' with self.assertRaises(RuntimeError): WorkerIdInterceptor().intercept_stream_stream( continuation, _ClientCallDetails(None, None, {'worker_id': '1'}, None), []) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/worker_pool_main.py0000644000175100017510000002022715206113032024605 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Worker pool entry point. The worker pool exposes an RPC service that is used with EXTERNAL environment to start and stop the SDK workers. The worker pool uses child processes for parallelism; threads are subject to the GIL and not sufficient. This entry point is used by the Python SDK container in worker pool mode. """ # pytype: skip-file import argparse import atexit import logging import subprocess import sys import threading import time import traceback from typing import Optional import grpc from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.runners.worker import sdk_worker from apache_beam.utils import thread_pool_executor from apache_beam.utils.sentinel import Sentinel _LOGGER = logging.getLogger(__name__) def kill_process_gracefully(proc, timeout=10): """ Kill a worker process gracefully by sending a SIGTERM and waiting for it to finish. A SIGKILL will be sent if the process has not finished after ``timeout`` seconds. """ def _kill(): proc.terminate() try: proc.wait(timeout=timeout) except subprocess.TimeoutExpired: _LOGGER.warning('Worker process did not respond, killing it.') proc.kill() proc.wait() # Avoid zombies kill_thread = threading.Thread(target=_kill) kill_thread.start() kill_thread.join() class BeamFnExternalWorkerPoolServicer( beam_fn_api_pb2_grpc.BeamFnExternalWorkerPoolServicer): def __init__( self, use_process=False, container_executable: Optional[str] = None, state_cache_size=0, data_buffer_time_limit_ms=0): self._use_process = use_process self._container_executable = container_executable self._state_cache_size = state_cache_size self._data_buffer_time_limit_ms = data_buffer_time_limit_ms self._worker_processes: dict[str, subprocess.Popen] = {} self._worker_threads: dict[str, sdk_worker.SdkHarness] = {} @classmethod def start( cls, use_process=False, port=0, state_cache_size=0, data_buffer_time_limit_ms=-1, container_executable: Optional[str] = None) -> tuple[str, grpc.Server]: options = [("grpc.http2.max_pings_without_data", 0), ("grpc.http2.max_ping_strikes", 0)] worker_server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=options) worker_address = 'localhost:%s' % worker_server.add_insecure_port( '[::]:%s' % port) worker_pool = cls( use_process=use_process, container_executable=container_executable, state_cache_size=state_cache_size, data_buffer_time_limit_ms=data_buffer_time_limit_ms) beam_fn_api_pb2_grpc.add_BeamFnExternalWorkerPoolServicer_to_server( worker_pool, worker_server) worker_server.start() _LOGGER.info('Listening for workers at %s', worker_address) # Register to kill the subprocesses on exit. def kill_worker_processes(): for worker_process in worker_pool._worker_processes.values(): kill_process_gracefully(worker_process) atexit.register(kill_worker_processes) return worker_address, worker_server def StartWorker( self, start_worker_request: beam_fn_api_pb2.StartWorkerRequest, unused_context) -> beam_fn_api_pb2.StartWorkerResponse: try: if self._use_process: command = [ 'python', '-c', 'from apache_beam.runners.worker.sdk_worker ' 'import SdkHarness; ' 'SdkHarness(' '"%s",' 'worker_id="%s",' 'state_cache_size=%d,' 'data_buffer_time_limit_ms=%d' ')' '.run()' % ( start_worker_request.control_endpoint.url, start_worker_request.worker_id, self._state_cache_size, self._data_buffer_time_limit_ms) ] if self._container_executable: # command as per container spec # the executable is responsible to handle concurrency # for artifact retrieval and other side effects command = [ self._container_executable, '--id=%s' % start_worker_request.worker_id, '--logging_endpoint=%s' % start_worker_request.logging_endpoint.url, '--artifact_endpoint=%s' % start_worker_request.artifact_endpoint.url, '--provision_endpoint=%s' % start_worker_request.provision_endpoint.url, '--control_endpoint=%s' % start_worker_request.control_endpoint.url, ] _LOGGER.warning("Starting worker with command %s" % command) worker_process = subprocess.Popen( command, stdout=subprocess.PIPE, close_fds=True) self._worker_processes[start_worker_request.worker_id] = worker_process else: worker = sdk_worker.SdkHarness( start_worker_request.control_endpoint.url, worker_id=start_worker_request.worker_id, state_cache_size=self._state_cache_size, data_buffer_time_limit_ms=self._data_buffer_time_limit_ms) self._worker_threads[start_worker_request.worker_id] = worker worker_thread = threading.Thread( name='run_worker_%s' % start_worker_request.worker_id, target=worker.run) worker_thread.daemon = True worker_thread.start() return beam_fn_api_pb2.StartWorkerResponse() except Exception: return beam_fn_api_pb2.StartWorkerResponse(error=traceback.format_exc()) def StopWorker( self, stop_worker_request: beam_fn_api_pb2.StopWorkerRequest, unused_context) -> beam_fn_api_pb2.StopWorkerResponse: # applicable for process mode to ensure process cleanup # thread based workers terminate automatically worker_process = self._worker_processes.pop( stop_worker_request.worker_id, None) if worker_process: _LOGGER.info("Stopping worker %s" % stop_worker_request.worker_id) kill_process_gracefully(worker_process) # applicable for thread mode to ensure thread cleanup by # unblocking the harness request stream. worker_thread_harness = self._worker_threads.pop( stop_worker_request.worker_id, None) if worker_thread_harness: _LOGGER.info("Stopping thread worker %s" % stop_worker_request.worker_id) worker_thread_harness._responses.put(Sentinel.sentinel) return beam_fn_api_pb2.StopWorkerResponse() def main(argv=None): """Entry point for worker pool service for external environments.""" parser = argparse.ArgumentParser() parser.add_argument( '--container_executable', type=str, default=None, help='Executable that implements the Beam SDK ' 'container contract.') parser.add_argument( '--service_port', type=int, required=True, dest='port', help='Bind port for the worker pool service.') args, _ = parser.parse_known_args(argv) address, server = (BeamFnExternalWorkerPoolServicer.start(use_process=True, **vars(args))) logging.getLogger().setLevel(logging.INFO) _LOGGER.info( 'Started worker pool servicer at port: %s with executable: %s', address, args.container_executable) try: while True: time.sleep(60 * 60 * 24) except KeyboardInterrupt: server.stop(0) if __name__ == '__main__': main(sys.argv) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/worker_status.py0000644000175100017510000003327515206113032024162 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Worker status api handler for reporting SDK harness debug info.""" import gc import logging import queue import sys import threading import time import traceback from collections import defaultdict import grpc from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.runners.worker.channel_factory import GRPCChannelFactory from apache_beam.runners.worker.statecache import StateCache from apache_beam.runners.worker.worker_id_interceptor import WorkerIdInterceptor from apache_beam.utils.sentinel import Sentinel try: from guppy import hpy except ImportError: hpy = None _LOGGER = logging.getLogger(__name__) # This SDK harness will (by default), log a "lull" in processing if it sees no # transitions in over 5 minutes. # 5 minutes * 60 seconds * 1000 millis * 1000 micros * 1000 nanoseconds DEFAULT_LOG_LULL_TIMEOUT_NS = 5 * 60 * 1000 * 1000 * 1000 # Full thread dump is performed at most every 20 minutes. LOG_LULL_FULL_THREAD_DUMP_INTERVAL_S = 20 * 60 # Full thread dump is performed if the lull is more than 20 minutes. LOG_LULL_FULL_THREAD_DUMP_LULL_S = 20 * 60 def _current_frames(): # Work around https://github.com/python/cpython/issues/106883 if (sys.version_info.minor == 11 and sys.version_info.major == 3 and gc.isenabled()): gc.disable() frames = sys._current_frames() # pylint: disable=protected-access gc.enable() return frames else: return sys._current_frames() # pylint: disable=protected-access def thread_dump(thread_prefix=None): """Get a thread dump for the current SDK harness. Args: thread_prefix: (str) An optional prefix to filter threads by. """ # deduplicate threads with same stack trace stack_traces = defaultdict(list) frames = _current_frames() threads_to_dump = threading.enumerate() if thread_prefix: threads_to_dump = [ t for t in threads_to_dump if t.name.startswith(thread_prefix) ] for t in threads_to_dump: try: stack_trace = ''.join(traceback.format_stack(frames[t.ident])) except KeyError: # the thread may have been destroyed already while enumerating, in such # case, skip to next thread. continue thread_ident_name = (t.ident, t.name) stack_traces[stack_trace].append(thread_ident_name) all_traces = ['=' * 10 + ' THREAD DUMP ' + '=' * 10] for stack, identity in stack_traces.items(): ident, name = identity[0] trace = '--- Thread #%s name: %s %s---\n' % ( ident, name, 'and other %d threads' % (len(identity) - 1) if len(identity) > 1 else '') if len(identity) > 1: trace += 'threads: %s\n' % identity trace += stack all_traces.append(trace) all_traces.append('=' * 30) return '\n'.join(all_traces) def heap_dump(): """Get a heap dump for the current SDK worker harness. """ banner = '=' * 10 + ' HEAP DUMP ' + '=' * 10 + '\n' if not hpy: heap = 'Unable to import guppy, the heap dump will be skipped.\n' else: heap = '%s\n' % hpy().heap() ending = '=' * 30 return banner + heap + ending def _state_cache_stats(state_cache: StateCache) -> str: """Gather state cache statistics.""" cache_stats = ['=' * 10 + ' CACHE STATS ' + '=' * 10] if not state_cache.is_cache_enabled(): cache_stats.append("Cache disabled") else: cache_stats.append(state_cache.describe_stats()) return '\n'.join(cache_stats) def _active_processing_bundles_state(bundle_processor_cache): """Gather information about the currently in-processing active bundles. The result only keeps the longest lasting 10 bundles to avoid excessive spamming. """ active_bundles = ['=' * 10 + ' ACTIVE PROCESSING BUNDLES ' + '=' * 10] if (not bundle_processor_cache.active_bundle_processors and not bundle_processor_cache.processors_being_created): active_bundles.append("No active processing bundles.") else: cache = [] for instruction in list( bundle_processor_cache.active_bundle_processors.keys()): processor = bundle_processor_cache.lookup(instruction) if processor: info = processor.state_sampler.get_info() cache.append(( instruction, processor.process_bundle_descriptor.id, info.tracked_thread, info.time_since_transition)) # reverse sort active bundle by time since last transition, keep top 10. cache.sort(key=lambda x: x[-1], reverse=True) for s in cache[:10]: state = '--- instruction %s ---\n' % s[0] state += 'ProcessBundleDescriptorId: %s\n' % s[1] state += "tracked thread: %s\n" % s[2] state += "time since transition: %.2f seconds\n" % (s[3] / 1e9) active_bundles.append(state) if bundle_processor_cache.processors_being_created: active_bundles.append("Processors being created:\n") current_time = time.time() for instruction, (bundle_id, thread, creation_time) in ( bundle_processor_cache.processors_being_created.items()): state = '--- instruction %s ---\n' % instruction state += 'ProcessBundleDescriptorId: %s\n' % bundle_id state += "tracked thread: %s\n" % thread state += "time since creation started: %.2f seconds\n" % ( current_time - creation_time) active_bundles.append(state) active_bundles.append('=' * 30) return '\n'.join(active_bundles) DONE = Sentinel.sentinel class FnApiWorkerStatusHandler(object): """FnApiWorkerStatusHandler handles worker status request from Runner. """ def __init__( self, status_address, bundle_processor_cache=None, state_cache=None, enable_heap_dump=False, worker_id=None, log_lull_timeout_ns=DEFAULT_LOG_LULL_TIMEOUT_NS, element_processing_timeout_minutes=None): """Initialize FnApiWorkerStatusHandler. Args: status_address: The URL Runner uses to host the WorkerStatus server. bundle_processor_cache: The BundleProcessor cache dict from sdk worker. state_cache: The StateCache form sdk worker. """ self._alive = True self._bundle_processor_cache = bundle_processor_cache self._state_cache = state_cache ch = GRPCChannelFactory.insecure_channel(status_address) grpc.channel_ready_future(ch).result(timeout=60) self._status_channel = grpc.intercept_channel( ch, WorkerIdInterceptor(worker_id)) self._status_stub = beam_fn_api_pb2_grpc.BeamFnWorkerStatusStub( self._status_channel) self._responses = queue.Queue() self.log_lull_timeout_ns = log_lull_timeout_ns if (element_processing_timeout_minutes and element_processing_timeout_minutes > 0): self._element_processing_timeout_ns = ( element_processing_timeout_minutes * 60 * 1e9) else: self._element_processing_timeout_ns = None self._last_full_thread_dump_secs = 0.0 self._last_lull_logged_secs = 0.0 self._server = threading.Thread( target=lambda: self._serve(), name='fn_api_status_handler') self._server.daemon = True self._enable_heap_dump = enable_heap_dump self._server.start() self._lull_logger = threading.Thread( target=lambda: self._log_lull_in_bundle_processor( self._bundle_processor_cache), name='lull_operation_logger') self._lull_logger.daemon = True self._lull_logger.start() def _get_responses(self): while True: response = self._responses.get() if response is DONE: self._alive = False return yield response def _serve(self): while self._alive: for request in self._status_stub.WorkerStatus(self._get_responses()): try: self._responses.put( beam_fn_api_pb2.WorkerStatusResponse( id=request.id, status_info=self.generate_status_response())) except Exception: traceback_string = traceback.format_exc() self._responses.put( beam_fn_api_pb2.WorkerStatusResponse( id=request.id, error="Exception encountered while generating " "status page: %s" % traceback_string)) def generate_status_response(self): all_status_sections = [] if self._state_cache: all_status_sections.append(_state_cache_stats(self._state_cache)) if self._bundle_processor_cache: all_status_sections.append( _active_processing_bundles_state(self._bundle_processor_cache)) all_status_sections.append(thread_dump()) if self._enable_heap_dump: all_status_sections.append(heap_dump()) return '\n'.join(all_status_sections) def close(self): self._responses.put(DONE, timeout=5) def _log_lull_in_bundle_processor(self, bundle_processor_cache): while True: time.sleep(2 * 60) if not bundle_processor_cache: continue for instruction in list( bundle_processor_cache.active_bundle_processors.keys()): processor = bundle_processor_cache.lookup(instruction) if processor: info = processor.state_sampler.get_info() self._log_lull_sampler_info(info, instruction) for instruction, (bundle_id, thread, creation_time) in list( bundle_processor_cache.processors_being_created.items()): self._log_lull_in_creating_bundle_descriptor( instruction, bundle_id, thread, creation_time) def _log_lull_in_creating_bundle_descriptor( self, instruction, bundle_id, thread, creation_time): time_since_creation_ns = (time.time() - creation_time) * 1e9 if (self._element_processing_timeout_ns and time_since_creation_ns > self._element_processing_timeout_ns): stack_trace = self._get_stack_trace(thread) _LOGGER.error(( 'Creation of bundle processor for instruction %s (bundle %s) ' 'has exceeded the specified timeout of %.2f minutes. ' 'This might indicate stuckness in DoFn.setup() or in DoFn creation. ' 'SDK harness will be terminated.\n' 'Current Traceback:\n%s'), instruction, bundle_id, self._element_processing_timeout_ns / 1e9 / 60, stack_trace) from apache_beam.runners.worker.sdk_worker_main import terminate_sdk_harness terminate_sdk_harness() if (time_since_creation_ns > self.log_lull_timeout_ns and self._passed_lull_timeout_since_last_log()): stack_trace = self._get_stack_trace(thread) _LOGGER.warning(( 'Bundle processor for instruction %s (bundle %s) ' 'has been creating for at least %.2f seconds.\n' 'This might indicate slowness in DoFn.setup() or in DoFn creation. ' 'Current Traceback:\n%s'), instruction, bundle_id, time_since_creation_ns / 1e9, stack_trace) def _log_lull_sampler_info(self, sampler_info, instruction): if (not sampler_info or not sampler_info.time_since_transition): return log_lull = ( sampler_info.time_since_transition > self.log_lull_timeout_ns and self._passed_lull_timeout_since_last_log()) timeout_exceeded = ( self._element_processing_timeout_ns and sampler_info.time_since_transition > self._element_processing_timeout_ns) if not (log_lull or timeout_exceeded): return lull_seconds = sampler_info.time_since_transition / 1e9 step_name = sampler_info.state_name.step_name state_name = sampler_info.state_name.name if step_name and state_name: step_name_log = ( ' for PTransform{name=%s, state=%s}' % (step_name, state_name)) else: step_name_log = '' stack_trace = self._get_stack_trace(sampler_info.tracked_thread) if timeout_exceeded: _LOGGER.error( ( 'Processing of an element in bundle %s%s has exceeded the ' 'specified timeout of %.2f minutes. SDK harness will be ' 'terminated.\n' 'Current Traceback:\n%s'), instruction, step_name_log, self._element_processing_timeout_ns / 1e9 / 60, stack_trace, ) from apache_beam.runners.worker.sdk_worker_main import terminate_sdk_harness terminate_sdk_harness() if log_lull: _LOGGER.warning( ( 'Operation ongoing in bundle %s%s for at least %.2f seconds' ' without outputting or completing.\n' 'Current Traceback:\n%s'), instruction, step_name_log, lull_seconds, stack_trace, ) def _get_stack_trace(self, thread): if thread: thread_frame = _current_frames().get(thread.ident) return '\n'.join( traceback.format_stack(thread_frame)) if thread_frame else '' else: return '-NOT AVAILABLE-' def _passed_lull_timeout_since_last_log(self) -> bool: if (time.time() - self._last_lull_logged_secs > self.log_lull_timeout_ns / 1e9): self._last_lull_logged_secs = time.time() return True else: return False ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/runners/worker/worker_status_test.py0000644000175100017510000002275515206113032025222 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import threading import time import unittest import grpc import mock from apache_beam.portability.api import beam_fn_api_pb2 from apache_beam.portability.api import beam_fn_api_pb2_grpc from apache_beam.runners.worker import statesampler from apache_beam.runners.worker.worker_status import FnApiWorkerStatusHandler from apache_beam.runners.worker.worker_status import heap_dump from apache_beam.utils import thread_pool_executor from apache_beam.utils.counters import CounterName class BeamFnStatusServicer(beam_fn_api_pb2_grpc.BeamFnWorkerStatusServicer): def __init__(self, num_request): self.finished = threading.Condition() self.num_request = num_request self.response_received = [] def WorkerStatus(self, response_iterator, context): for i in range(self.num_request): yield beam_fn_api_pb2.WorkerStatusRequest(id=str(i)) for response in response_iterator: self.finished.acquire() self.response_received.append(response) if len(self.response_received) == self.num_request: self.finished.notify_all() self.finished.release() class FnApiWorkerStatusHandlerTest(unittest.TestCase): def setUp(self): self.num_request = 3 self.test_status_service = BeamFnStatusServicer(self.num_request) self.server = grpc.server(thread_pool_executor.shared_unbounded_instance()) beam_fn_api_pb2_grpc.add_BeamFnWorkerStatusServicer_to_server( self.test_status_service, self.server) self.test_port = self.server.add_insecure_port('[::]:0') self.server.start() self.url = 'localhost:%s' % self.test_port self.fn_status_handler = FnApiWorkerStatusHandler( self.url, element_processing_timeout_minutes=10) def tearDown(self): self.fn_status_handler.close() self.server.stop(5) def test_send_status_response(self): self.test_status_service.finished.acquire() while len(self.test_status_service.response_received) < self.num_request: self.test_status_service.finished.wait(1) self.test_status_service.finished.release() for response in self.test_status_service.response_received: self.assertIsNotNone(response.status_info) @mock.patch( 'apache_beam.runners.worker.worker_status' '.FnApiWorkerStatusHandler.generate_status_response') def test_generate_error(self, mock_method): mock_method.side_effect = RuntimeError('error') self.test_status_service.finished.acquire() while len(self.test_status_service.response_received) < self.num_request: self.test_status_service.finished.wait(1) self.test_status_service.finished.release() for response in self.test_status_service.response_received: self.assertIsNotNone(response.error) def test_log_lull_in_bundle_processor(self): def get_state_sampler_info_for_lull(lull_duration_s): return "bundle-id", statesampler.StateSamplerInfo( CounterName('progress-msecs', 'stage_name', 'step_name'), 1, lull_duration_s * 1e9, threading.current_thread()) now = time.time() with mock.patch('logging.Logger.warning') as warn_mock: with mock.patch( 'apache_beam.runners.worker.sdk_worker_main.terminate_sdk_harness' ) as flush_mock: with mock.patch('time.time') as time_mock: time_mock.return_value = now bundle_id, sampler_info = get_state_sampler_info_for_lull(21 * 60) self.fn_status_handler._log_lull_sampler_info(sampler_info, bundle_id) bundle_id_template = warn_mock.call_args[0][1] step_name_template = warn_mock.call_args[0][2] processing_template = warn_mock.call_args[0][3] traceback = warn_mock.call_args = warn_mock.call_args[0][4] self.assertIn('bundle-id', bundle_id_template) self.assertIn('step_name', step_name_template) self.assertEqual(21 * 60, processing_template) self.assertIn('test_log_lull_in_bundle_processor', traceback) flush_mock.assert_called_once() with mock.patch('time.time') as time_mock: time_mock.return_value = now + 6 * 60 # 6 minutes bundle_id, sampler_info = get_state_sampler_info_for_lull(21 * 60) self.fn_status_handler._log_lull_sampler_info(sampler_info, bundle_id) self.assertEqual(flush_mock.call_count, 2) with mock.patch('time.time') as time_mock: time_mock.return_value = now + 21 * 60 # 21 minutes bundle_id, sampler_info = get_state_sampler_info_for_lull(10 * 60) self.fn_status_handler._log_lull_sampler_info(sampler_info, bundle_id) self.assertEqual(flush_mock.call_count, 2) with mock.patch('time.time') as time_mock: time_mock.return_value = now + 42 * 60 # 42 minutes bundle_id, sampler_info = get_state_sampler_info_for_lull(11 * 60) self.fn_status_handler._log_lull_sampler_info(sampler_info, bundle_id) self.assertEqual(flush_mock.call_count, 3) def test_lull_logs_emitted_when_creating_bundle_processor_takes_time(self): instruction_id = "instruction-1" bundle_id = "bundle-1" thread = threading.current_thread() now = time.time() creation_time = now with ( mock.patch('logging.Logger.warning') as warn_mock, mock.patch('logging.Logger.error') as error_mock, mock.patch('time.time') as time_mock, mock.patch( 'apache_beam.runners.worker.sdk_worker_main.terminate_sdk_harness', ) as terminate_mock): # Set time to be past the lull timeout time_mock.return_value = ( now + self.fn_status_handler.log_lull_timeout_ns / 1e9 + 1) self.fn_status_handler._log_lull_in_creating_bundle_descriptor( instruction_id, bundle_id, thread, creation_time) warn_mock.assert_called_once() args, _ = warn_mock.call_args self.assertIn( 'Bundle processor for instruction %s (bundle %s) has been ' 'creating for at least %.2f seconds', args[0]) # Set time to be past the element processing timeout time_mock.return_value = ( now + self.fn_status_handler._element_processing_timeout_ns / 1e9 + 1) self.fn_status_handler._log_lull_in_creating_bundle_descriptor( instruction_id, bundle_id, thread, creation_time) error_mock.assert_called_once() args, _ = error_mock.call_args self.assertIn( 'Creation of bundle processor for instruction %s (bundle %s) ' 'has exceeded the specified timeout', args[0]) terminate_mock.assert_called_once() def test_lull_logs_emitted_when_processing_a_bundle_takes_time(self): instruction_id = "instruction-1" now = time.time() thread = threading.current_thread() with ( mock.patch('logging.Logger.warning') as warn_mock, mock.patch('logging.Logger.error') as error_mock, mock.patch('time.time') as time_mock, mock.patch( 'apache_beam.runners.worker.sdk_worker_main.terminate_sdk_harness', ) as terminate_mock): time_mock.return_value = now + 1 # Set time to be past the lull timeout sampler_info = statesampler.StateSamplerInfo( state_name=CounterName('test_counter', 'test_stage', 'test_step'), transition_count=1, # Set time to be past the lull timeout time_since_transition=( self.fn_status_handler.log_lull_timeout_ns + 1), tracked_thread=thread) self.fn_status_handler._log_lull_sampler_info( sampler_info, instruction_id) warn_mock.assert_called_once() args, _ = warn_mock.call_args self.assertIn( 'Operation ongoing in bundle %s%s for at least %.2f seconds', args[0]) time_mock.return_value = now + 2 sampler_info = statesampler.StateSamplerInfo( state_name=CounterName('test_counter', 'test_stage', 'test_step'), transition_count=1, # Set time to be past the element processing timeout time_since_transition=( self.fn_status_handler._element_processing_timeout_ns + 1), tracked_thread=thread) self.fn_status_handler._log_lull_sampler_info( sampler_info, instruction_id) error_mock.assert_called_once() args, _ = error_mock.call_args self.assertIn( 'Processing of an element in bundle %s%s has exceeded the ' 'specified timeout of %.2f minutes', args[0]) terminate_mock.assert_called_once() class HeapDumpTest(unittest.TestCase): @mock.patch('apache_beam.runners.worker.worker_status.hpy', None) def test_skip_heap_dump(self): result = '%s' % heap_dump() self.assertTrue( 'Unable to import guppy, the heap dump will be skipped' in result) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.6994414 apache_beam-2.74.0/apache_beam/testing/0000755000175100017510000000000015206113160017334 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/__init__.py0000644000175100017510000000142015206113032021440 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.7014418 apache_beam-2.74.0/apache_beam/testing/analyzers/0000755000175100017510000000000015206113160021344 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/analyzers/__init__.py0000644000175100017510000000155315206113032023457 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Performance alert tooling for Apache Beam. No backwards compatibility guarantees. """ ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/analyzers/constants.py0000644000175100017510000000522315206113032023732 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """The file defines global variables for performance analysis.""" _BQ_PROJECT_NAME = 'apache-beam-testing' _BQ_DATASET = 'beam_perf_storage' _UNIQUE_ID = 'test_id' _ISSUE_CREATION_TIMESTAMP_LABEL = 'issue_timestamp' _CHANGE_POINT_TIMESTAMP_LABEL = 'change_point_timestamp' _CHANGE_POINT_LABEL = 'change_point' _TEST_NAME = 'test_name' _METRIC_NAME = 'metric_name' _ISSUE_NUMBER = 'issue_number' _ISSUE_URL = 'issue_url' # number of results to display on the issue description # from change point index in both directions. _NUM_RESULTS_TO_DISPLAY_ON_ISSUE_DESCRIPTION = 10 _NUM_DATA_POINTS_TO_RUN_CHANGE_POINT_ANALYSIS = 100 # Variables used for finding duplicate change points. _DEFAULT_MIN_RUNS_BETWEEN_CHANGE_POINTS = 3 _DEFAULT_NUM_RUMS_IN_CHANGE_POINT_WINDOW = 14 _DEFAULT_MEDIAN_ABS_DEVIATION_THRESHOLD = 2 _PERF_TEST_KEYS = { 'test_description', 'metrics_dataset', 'metrics_table', 'project', 'metric_name' } _SCHEMA = [{ 'name': _UNIQUE_ID, 'field_type': 'STRING', 'mode': 'REQUIRED' }, { 'name': _ISSUE_CREATION_TIMESTAMP_LABEL, 'field_type': 'TIMESTAMP', 'mode': 'REQUIRED' }, { 'name': _CHANGE_POINT_TIMESTAMP_LABEL, 'field_type': 'TIMESTAMP', 'mode': 'REQUIRED' }, { 'name': _CHANGE_POINT_LABEL, 'field_type': 'FLOAT64', 'mode': 'REQUIRED' }, { 'name': _METRIC_NAME, 'field_type': 'STRING', 'mode': 'REQUIRED' }, { 'name': _TEST_NAME, 'field_type': 'STRING', 'mode': 'REQUIRED' }, { 'name': _ISSUE_NUMBER, 'field_type': 'INT64', 'mode': 'REQUIRED' }, { 'name': _ISSUE_URL, 'field_type': 'STRING', 'mode': 'REQUIRED' }] _ANOMALY_MARKER = ' <---- Anomaly' _EDGE_SEGMENT_SIZE = 3 ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/analyzers/github_issues_utils.py0000644000175100017510000001713015206113032026013 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import json import logging import os from typing import Optional import requests from apache_beam.testing.analyzers import constants from apache_beam.testing.analyzers.perf_analysis_utils import MetricContainer from apache_beam.testing.analyzers.perf_analysis_utils import TestConfigContainer try: _GITHUB_TOKEN: Optional[str] = os.environ['GITHUB_TOKEN'] except KeyError as e: _GITHUB_TOKEN = None logging.warning( 'A Github Personal Access token is required ' 'to create Github Issues.') _GITHUB_REPO_OWNER = os.environ.get('REPO_OWNER', 'apache') _GITHUB_REPO_NAME = os.environ.get('REPO_NAME', 'beam') # Adding GitHub Rest API version to the header to maintain version stability. # For more information, please look at # https://github.blog/2022-11-28-to-infinity-and-beyond-enabling-the-future-of-githubs-rest-api-with-api-versioning/ # pylint: disable=line-too-long _HEADERS = { "Authorization": 'token {}'.format(_GITHUB_TOKEN), "Accept": "application/vnd.github+json", "X-GitHub-Api-Version": "2022-11-28" } _ISSUE_TITLE_TEMPLATE = """ Performance Regression or Improvement: {}:{} """ _ISSUE_DESCRIPTION_TEMPLATE = """ Performance change found in the test: `{}` for the metric: `{}`. For more information on how to triage the alerts, please look at `Triage performance alert issues` section of the [README](https://github.com/apache/beam/tree/master/sdks/python/apache_beam/testing/analyzers/README.md#triage-performance-alert-issues). """ _METRIC_INFO_TEMPLATE = "timestamp: {}, metric_value: {}" _AWAITING_TRIAGE_LABEL = 'awaiting triage' _PERF_ALERT_LABEL = 'perf-alert' _REQUEST_TIMEOUT_SECS = 60 def create_issue( title: str, description: str, labels: Optional[list[str]] = None, ) -> tuple[int, str]: """ Create an issue with title, description with a label. Args: title: GitHub issue title. description: GitHub issue description. labels: Labels used to tag the GitHub issue. Returns: Tuple containing GitHub issue number and issue URL. """ url = "https://api.github.com/repos/{}/{}/issues".format( _GITHUB_REPO_OWNER, _GITHUB_REPO_NAME) data = { 'owner': _GITHUB_REPO_OWNER, 'repo': _GITHUB_REPO_NAME, 'title': title, 'body': description, 'labels': [_AWAITING_TRIAGE_LABEL, _PERF_ALERT_LABEL] } if labels: data['labels'].extend(labels) # type: ignore response = requests.post( url=url, data=json.dumps(data), headers=_HEADERS, timeout=_REQUEST_TIMEOUT_SECS).json() return response['number'], response['html_url'] def comment_on_issue(issue_number: int, comment_description: str) -> tuple[bool, str]: """ This method looks for an issue with provided issue_number. If an open issue is found, comment on the open issue with provided description else do nothing. Args: issue_number: A GitHub issue number. comment_description: If an issue with issue_number is open, then comment on the issue with the using comment_description. Returns: tuple[bool, Optional[str]] indicating if a comment was added to issue, and the comment URL. """ url = 'https://api.github.com/repos/{}/{}/issues/{}'.format( _GITHUB_REPO_OWNER, _GITHUB_REPO_NAME, issue_number) open_issue_response = requests.get( url, json.dumps({ 'owner': _GITHUB_REPO_OWNER, 'repo': _GITHUB_REPO_NAME, 'issue_number': issue_number }, default=str), headers=_HEADERS, timeout=_REQUEST_TIMEOUT_SECS).json() if open_issue_response['state'] == 'open': data = { 'owner': _GITHUB_REPO_OWNER, 'repo': _GITHUB_REPO_NAME, 'body': comment_description, issue_number: issue_number, } response = requests.post( open_issue_response['comments_url'], json.dumps(data), headers=_HEADERS, timeout=_REQUEST_TIMEOUT_SECS) return True, response.json()['html_url'] return False, '' def add_awaiting_triage_label(issue_number: int): url = 'https://api.github.com/repos/{}/{}/issues/{}/labels'.format( _GITHUB_REPO_OWNER, _GITHUB_REPO_NAME, issue_number) requests.post( url, json.dumps({'labels': [_AWAITING_TRIAGE_LABEL]}), headers=_HEADERS, timeout=_REQUEST_TIMEOUT_SECS) def get_issue_description( test_config_container: TestConfigContainer, metric_container: MetricContainer, change_point_index: int, max_results_to_display: int = 5, ) -> str: """ Args: test_config_container: TestConfigContainer containing test metadata. metric_container: MetricContainer containing metric data. change_point_index: Index of the change point in the metric data. max_results_to_display: Max number of results to display from the change point index, in both directions of the change point index. Returns: str: Description used to fill the GitHub issues description. """ # TODO: Add mean and median before and after the changepoint index. description = [] description.append( _ISSUE_DESCRIPTION_TEMPLATE.format( test_config_container.test_id, test_config_container.metric_name)) if test_config_container.test_name: description.append(("`test_name:` " + f'{test_config_container.test_name}')) if test_config_container.test_description: description.append( ("`Test description:` " + f'{test_config_container.test_description}')) description.append('```') runs_to_display = [] max_timestamp_index = min( change_point_index + max_results_to_display, len(metric_container.values) - 1) min_timestamp_index = max(0, change_point_index - max_results_to_display) # run in reverse to display the most recent runs first. for i in reversed(range(min_timestamp_index, max_timestamp_index + 1)): row_template = _METRIC_INFO_TEMPLATE.format( metric_container.timestamps[i].ctime(), format(metric_container.values[i], '.2f')) if i == change_point_index: row_template += constants._ANOMALY_MARKER runs_to_display.append(row_template) description.append(os.linesep.join(runs_to_display)) description.append('```') return (2 * os.linesep).join(description) def report_change_point_on_issues( title: str, description: str, labels: Optional[list[str]] = None, existing_issue_number: Optional[int] = None, ) -> tuple[int, str]: """ Comments the description on the existing issue (if provided and still open), or creates a new issue. """ if existing_issue_number is not None: commented_on_issue, issue_url = comment_on_issue( issue_number=existing_issue_number, comment_description=description ) if commented_on_issue: add_awaiting_triage_label(issue_number=existing_issue_number) return existing_issue_number, issue_url return create_issue(title=title, description=description, labels=labels) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/analyzers/load_test_perf_analysis.py0000644000175100017510000000751715206113032026623 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import argparse import logging from apache_beam.testing.analyzers import constants from apache_beam.testing.analyzers import perf_analysis from apache_beam.testing.analyzers import perf_analysis_utils from apache_beam.testing.analyzers.perf_analysis_utils import MetricContainer from apache_beam.testing.analyzers.perf_analysis_utils import TestConfigContainer try: from google.cloud import bigquery except ImportError: bigquery = None # type: ignore class LoadTestMetricsFetcher(perf_analysis_utils.MetricsFetcher): """ Metrics fetcher used to get metric data from a BigQuery table. The metrics are fetched and returned as a dataclass containing lists of timestamps and metric_values. """ def fetch_metric_data( self, *, test_config: TestConfigContainer) -> MetricContainer: if test_config.test_name: test_name, pipeline_name = test_config.test_name.split(',') else: raise Exception("test_name not provided in config.") query = f""" SELECT timestamp, metric.value FROM {test_config.project}.{test_config.metrics_dataset}.{test_config.metrics_table} CROSS JOIN UNNEST(metrics) AS metric WHERE test_name = "{test_name}" AND pipeline_name = "{pipeline_name}" AND metric.name = "{test_config.metric_name}" ORDER BY timestamp DESC LIMIT {constants._NUM_DATA_POINTS_TO_RUN_CHANGE_POINT_ANALYSIS} """ logging.debug("Running query: %s" % query) if bigquery is None: raise ImportError('Bigquery dependencies are not installed.') client = bigquery.Client() query_job = client.query(query=query) metric_data = query_job.result().to_dataframe() if metric_data.empty: logging.error( "No results returned from BigQuery. Please check the query.") return MetricContainer( values=metric_data['value'].tolist(), timestamps=metric_data['timestamp'].tolist(), ) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) load_test_metrics_fetcher = LoadTestMetricsFetcher() parser = argparse.ArgumentParser() parser.add_argument( '--config_file_path', required=True, type=str, help='Path to the config file that contains data to run the Change Point ' 'Analysis.The default file will used will be ' 'apache_beam/testing/analyzers/tests.config.yml. ' 'If you would like to use the Change Point Analysis for finding ' 'performance regression in the tests, ' 'please provide an .yml file in the same structure as the above ' 'mentioned file. ') parser.add_argument( '--save_alert_metadata', action='store_true', default=False, help='Save perf alert/ GH Issue metadata to BigQuery table.') known_args, unknown_args = parser.parse_known_args() if unknown_args: logging.warning('Discarding unknown arguments : %s ' % unknown_args) perf_analysis.run( big_query_metrics_fetcher=load_test_metrics_fetcher, config_file_path=known_args.config_file_path, # Set this to true while running in production. save_alert_metadata=known_args.save_alert_metadata) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/analyzers/perf_analysis.py0000644000175100017510000002722215206113032024560 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # This script is used to run Change Point Analysis using a config file. # config file holds the parameters required to fetch data, and to run the # change point analysis. Change Point Analysis is used to find Performance # regressions for benchmark/load/performance test. import argparse import logging import uuid from datetime import datetime from datetime import timezone from typing import Any import pandas as pd from apache_beam.testing.analyzers import constants from apache_beam.testing.analyzers.perf_analysis_utils import BigQueryMetricsFetcher from apache_beam.testing.analyzers.perf_analysis_utils import ChangePointConfig from apache_beam.testing.analyzers.perf_analysis_utils import GitHubIssueMetaData from apache_beam.testing.analyzers.perf_analysis_utils import MetricsFetcher from apache_beam.testing.analyzers.perf_analysis_utils import TestConfigContainer from apache_beam.testing.analyzers.perf_analysis_utils import create_performance_alert from apache_beam.testing.analyzers.perf_analysis_utils import find_latest_change_point_index from apache_beam.testing.analyzers.perf_analysis_utils import get_existing_issues_data from apache_beam.testing.analyzers.perf_analysis_utils import is_change_point_in_valid_window from apache_beam.testing.analyzers.perf_analysis_utils import is_sibling_change_point from apache_beam.testing.analyzers.perf_analysis_utils import publish_issue_metadata_to_big_query from apache_beam.testing.analyzers.perf_analysis_utils import read_test_config def get_test_config_container( params: dict[str, Any], test_id: str, metric_name: str, ) -> TestConfigContainer: """ Args: params: Dict containing parameters to run change point analysis. Returns: TestConfigContainer object containing test config parameters. """ return TestConfigContainer( project=params['project'], metrics_dataset=params['metrics_dataset'], metrics_table=params['metrics_table'], metric_name=metric_name, test_id=test_id, test_description=params['test_description'], test_name=params.get('test_name', None), labels=params.get('labels', None), ) def get_change_point_config( params: dict[str, Any], ) -> ChangePointConfig: """ Args: params: Dict containing parameters to run change point analysis. Returns: ChangePointConfig object containing change point analysis parameters. """ return ChangePointConfig( min_runs_between_change_points=params.get( 'min_runs_between_change_points', constants._DEFAULT_MIN_RUNS_BETWEEN_CHANGE_POINTS), num_runs_in_change_point_window=params.get( 'num_runs_in_change_point_window', constants._DEFAULT_NUM_RUMS_IN_CHANGE_POINT_WINDOW), median_abs_deviation_threshold=params.get( 'median_abs_deviation_threshold', constants._DEFAULT_MEDIAN_ABS_DEVIATION_THRESHOLD)) def run_change_point_analysis( test_config_container: TestConfigContainer, big_query_metrics_fetcher: MetricsFetcher, change_point_config: ChangePointConfig = ChangePointConfig(), save_alert_metadata: bool = False, ): """ Args: test_config_container: TestConfigContainer containing test metadata for fetching data and running change point analysis. big_query_metrics_fetcher: BigQuery metrics fetcher used to fetch data for change point analysis. change_point_config: ChangePointConfig containing parameters to run change point analysis. save_alert_metadata: bool indicating if issue metadata should be published to BigQuery table. Returns: bool indicating if a change point is observed and alerted on GitHub. """ logging.info( "Running change point analysis for test ID :%s on metric: % s" % (test_config_container.test_id, test_config_container.metric_name)) # test_name will be used to query a single test from # multiple tests in a single BQ table. Right now, the default # assumption is that all the test have an individual BQ table # but this might not be case for other tests(such as IO tests where # a single BQ tables stores all the data) test_name = test_config_container.test_name min_runs_between_change_points = ( change_point_config.min_runs_between_change_points) num_runs_in_change_point_window = ( change_point_config.num_runs_in_change_point_window) metric_container = big_query_metrics_fetcher.fetch_metric_data( test_config=test_config_container) metric_container.sort_by_timestamp() metric_values = metric_container.values timestamps = metric_container.timestamps change_point_index = find_latest_change_point_index( metric_values=metric_values, median_abs_deviation_threshold=change_point_config. median_abs_deviation_threshold) if not change_point_index: logging.info( "Change point is not detected for the test ID %s" % test_config_container.test_id) return False # since timestamps are ordered in ascending order and # num_runs_in_change_point_window refers to the latest runs, # latest_change_point_run can help determine if the change point # index is recent wrt num_runs_in_change_point_window latest_change_point_run = len(timestamps) - 1 - change_point_index if not is_change_point_in_valid_window(num_runs_in_change_point_window, latest_change_point_run): logging.info( 'Performance regression/improvement found for the test ID: %s. ' 'on metric %s. Since the change point run %s ' 'lies outside the num_runs_in_change_point_window distance: %s, ' 'alert is not raised.' % ( test_config_container.test_id, test_config_container.metric_name, latest_change_point_run + 1, num_runs_in_change_point_window)) return False is_valid_change_point = True last_reported_issue_number = None # create a unique table name for each test and metric combination. # for beam load tests, metric_name and metric table are enough to # create a unique table name. For templates/IO tests, add `test_name`. issue_metadata_table_name = ( f'{test_config_container.metrics_table}_{test_config_container.metric_name}' # pylint: disable=line-too-long ) if test_config_container.test_name: issue_metadata_table_name = ( f'{issue_metadata_table_name}_{test_config_container.test_name}') existing_issue_data = get_existing_issues_data( table_name=issue_metadata_table_name) if existing_issue_data is not None: existing_issue_timestamps = existing_issue_data[ constants._CHANGE_POINT_TIMESTAMP_LABEL].tolist() last_reported_issue_number = existing_issue_data[ constants._ISSUE_NUMBER].tolist()[0] if not isinstance(last_reported_issue_number, int): # convert numpy.int64 to int last_reported_issue_number = last_reported_issue_number.item() is_valid_change_point = is_sibling_change_point( previous_change_point_timestamps=existing_issue_timestamps, change_point_index=change_point_index, timestamps=timestamps, min_runs_between_change_points=min_runs_between_change_points, test_id=test_config_container.test_id) # for testing purposes, we don't want to create an issue even if there is # a valid change point. This is useful when we want to test the change point # analysis logic without creating an issue. if is_valid_change_point and save_alert_metadata: issue_number, issue_url = create_performance_alert( test_config_container=test_config_container, metric_container=metric_container, change_point_index=change_point_index, existing_issue_number=last_reported_issue_number, ) issue_metadata = GitHubIssueMetaData( issue_timestamp=pd.Timestamp( datetime.now().replace(tzinfo=timezone.utc)), # BQ doesn't allow '.' in table name test_id=test_config_container.test_id.replace('.', '_'), test_name=test_name or uuid.uuid4().hex, metric_name=test_config_container.metric_name, change_point=metric_values[change_point_index], issue_number=issue_number, issue_url=issue_url, change_point_timestamp=timestamps[change_point_index], ) publish_issue_metadata_to_big_query( issue_metadata=issue_metadata, table_name=issue_metadata_table_name, project=test_config_container.project, ) return is_valid_change_point def run( *, config_file_path: str, big_query_metrics_fetcher: MetricsFetcher = BigQueryMetricsFetcher(), save_alert_metadata: bool = False, ) -> None: """ run is the entry point to run change point analysis on test metric data, which is read from config file, and if there is a performance regression/improvement observed for a test, an alert will filed with GitHub Issues. If config_file_path is None, then the run method will use default config file to read the required perf test parameters. Please take a look at the README for more information on the parameters defined in the config file. """ tests_config: dict[str, dict[str, Any]] = read_test_config(config_file_path) for test_id, params in tests_config.items(): # single test config can have multiple metrics so we need to # iterate over all the metrics and run change point analysis # for each metric. metric_names = params['metric_name'] if isinstance(metric_names, str): metric_names = [metric_names] for metric_name in metric_names: test_config_container = get_test_config_container( params=params, test_id=test_id, metric_name=metric_name) change_point_config = get_change_point_config(params) run_change_point_analysis( test_config_container=test_config_container, big_query_metrics_fetcher=big_query_metrics_fetcher, change_point_config=change_point_config, save_alert_metadata=save_alert_metadata) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) parser = argparse.ArgumentParser() parser.add_argument( '--config_file_path', required=True, type=str, help='Path to the config file that contains data to run the Change Point ' 'Analysis.The default file will used will be ' 'apache_beam/testing/analyzers/tests.config.yml. ' 'If you would like to use the Change Point Analysis for finding ' 'performance regression in the tests, ' 'please provide an .yml file in the same structure as the above ' 'mentioned file. ') parser.add_argument( '--save_alert_metadata', action='store_true', help='Save perf alert/ GH Issue metadata to BigQuery table.') known_args, unknown_args = parser.parse_known_args() if unknown_args: logging.warning('Discarding unknown arguments : %s ' % unknown_args) run( config_file_path=known_args.config_file_path, # Set this to true while running in production. save_alert_metadata=known_args.save_alert_metadata # pylint: disable=line-too-long ) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/analyzers/perf_analysis_test.py0000644000175100017510000002444215206113032025620 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import os import re import unittest import mock import numpy as np import pandas as pd # pylint: disable=ungrouped-imports try: import apache_beam.testing.analyzers.perf_analysis as analysis from apache_beam.io.filesystems import FileSystems from apache_beam.testing.analyzers import constants from apache_beam.testing.analyzers import github_issues_utils from apache_beam.testing.analyzers.perf_analysis_utils import BigQueryMetricsFetcher from apache_beam.testing.analyzers.perf_analysis_utils import MetricContainer from apache_beam.testing.analyzers.perf_analysis_utils import TestConfigContainer from apache_beam.testing.analyzers.perf_analysis_utils import e_divisive from apache_beam.testing.analyzers.perf_analysis_utils import filter_change_points_by_median_threshold from apache_beam.testing.analyzers.perf_analysis_utils import find_change_points from apache_beam.testing.analyzers.perf_analysis_utils import find_latest_change_point_index from apache_beam.testing.analyzers.perf_analysis_utils import is_change_point_in_valid_window from apache_beam.testing.analyzers.perf_analysis_utils import is_sibling_change_point from apache_beam.testing.analyzers.perf_analysis_utils import validate_config from apache_beam.testing.load_tests import load_test_metrics_utils except ImportError as e: raise unittest.SkipTest('Missing dependencies to run perf analysis tests.') # mock methods. def get_fake_data_with_no_change_point(*args, **kwargs): num_samples = 20 metric_values = [1] * num_samples timestamps = [pd.Timestamp(i) for i in range(num_samples)] return MetricContainer(metric_values, timestamps) def get_fake_data_with_change_point(*args, **kwargs): # change point will be at index 13. num_samples = 20 metric_values = [0] * 12 + [3] + [4] * 7 timestamps = [pd.Timestamp(i) for i in range(num_samples)] return MetricContainer(metric_values, timestamps) def get_existing_issue_data(**kwargs): # change point found at index 13. So passing 13 in the # existing issue data in mock method. return pd.DataFrame([{ constants._CHANGE_POINT_TIMESTAMP_LABEL: pd.Timestamp(13), constants._ISSUE_NUMBER: np.array([0]) }]) class TestChangePointAnalysis(unittest.TestCase): def setUp(self) -> None: self.single_change_point_series = [0] * 10 + [1] * 10 self.multiple_change_point_series = self.single_change_point_series + [ 2 ] * 20 self.timestamps = [pd.Timestamp(i) for i in range(5)] self.params = { 'test_description': 'fake_description', 'metrics_dataset': 'fake_dataset', 'metrics_table': 'fake_table', 'project': 'fake_project', 'metric_name': 'fake_metric_name' } self.test_id = 'fake_id' def test_edivisive_means(self): change_point_indexes = e_divisive(self.single_change_point_series) self.assertEqual(change_point_indexes, [10]) change_point_indexes = e_divisive(self.multiple_change_point_series) self.assertEqual(sorted(change_point_indexes), [10, 20]) def test_is_changepoint_in_valid_window(self): changepoint_to_recent_run_window = 19 change_point_index = 14 is_valid = is_change_point_in_valid_window( changepoint_to_recent_run_window, change_point_index) self.assertEqual(is_valid, True) def test_change_point_outside_inspection_window_is_not_a_valid_alert(self): changepoint_to_recent_run_window = 12 change_point_index = 14 is_valid = is_change_point_in_valid_window( changepoint_to_recent_run_window, change_point_index) self.assertEqual(is_valid, False) def test_validate_config(self): test_keys = { 'test_description', 'metrics_dataset', 'metrics_table', 'project', 'metric_name' } self.assertEqual(test_keys, constants._PERF_TEST_KEYS) self.assertTrue(validate_config(test_keys)) def test_duplicate_change_point(self): change_point_index = 2 min_runs_between_change_points = 1 is_alert = is_sibling_change_point( previous_change_point_timestamps=[self.timestamps[0]], timestamps=self.timestamps, change_point_index=change_point_index, min_runs_between_change_points=min_runs_between_change_points, test_id=self.test_id) self.assertTrue(is_alert) def test_duplicate_change_points_are_not_valid_alerts(self): change_point_index = 2 min_runs_between_change_points = 1 is_alert = is_sibling_change_point( previous_change_point_timestamps=[self.timestamps[3]], timestamps=self.timestamps, change_point_index=change_point_index, min_runs_between_change_points=min_runs_between_change_points, test_id=self.test_id) self.assertFalse(is_alert) is_alert = is_sibling_change_point( previous_change_point_timestamps=[ self.timestamps[0], self.timestamps[3] ], timestamps=self.timestamps, change_point_index=change_point_index, min_runs_between_change_points=min_runs_between_change_points, test_id=self.test_id) self.assertFalse(is_alert) @mock.patch.object( BigQueryMetricsFetcher, 'fetch_metric_data', get_fake_data_with_no_change_point) def test_no_alerts_when_no_change_points(self): test_config_container = analysis.get_test_config_container( params=self.params, test_id=self.test_id, metric_name=self.params['metric_name']) is_alert = analysis.run_change_point_analysis( test_config_container=test_config_container, big_query_metrics_fetcher=BigQueryMetricsFetcher()) self.assertFalse(is_alert) @mock.patch.object( BigQueryMetricsFetcher, 'fetch_metric_data', get_fake_data_with_change_point) @mock.patch( 'apache_beam.testing.analyzers.perf_analysis.get_existing_issues_data', return_value=None) @mock.patch( 'apache_beam.testing.analyzers.perf_analysis.' 'publish_issue_metadata_to_big_query', return_value=None) @mock.patch( 'apache_beam.testing.analyzers.perf_analysis' '.create_performance_alert', return_value=(0, '')) def test_alert_on_data_with_change_point(self, *args): test_config_container = analysis.get_test_config_container( params=self.params, test_id=self.test_id, metric_name=self.params['metric_name']) is_alert = analysis.run_change_point_analysis( test_config_container=test_config_container, big_query_metrics_fetcher=BigQueryMetricsFetcher()) self.assertTrue(is_alert) @mock.patch.object( BigQueryMetricsFetcher, 'fetch_metric_data', get_fake_data_with_change_point) @mock.patch( 'apache_beam.testing.analyzers.perf_analysis.get_existing_issues_data', get_existing_issue_data) @mock.patch( 'apache_beam.testing.analyzers.perf_analysis.' 'publish_issue_metadata_to_big_query', return_value=None) @mock.patch( 'apache_beam.testing.analyzers.perf_analysis.create_performance_alert', return_value=(0, '')) def test_alert_on_data_with_reported_change_point(self, *args): test_config_container = analysis.get_test_config_container( params=self.params, test_id=self.test_id, metric_name=self.params['metric_name']) is_alert = analysis.run_change_point_analysis( test_config_container=test_config_container, big_query_metrics_fetcher=BigQueryMetricsFetcher()) self.assertFalse(is_alert) def test_change_point_has_anomaly_marker_in_gh_description(self): metric_container = get_fake_data_with_change_point() metric_values = metric_container.values change_point_index = find_latest_change_point_index(metric_values) test_config_container = TestConfigContainer( project=self.params['project'], metrics_dataset=self.params['metrics_dataset'], metrics_table=self.params['metrics_table'], metric_name=self.params['metric_name'], test_id=self.test_id, test_description=self.params['test_description'], test_name=self.params.get('test_name', None), labels=self.params.get('labels', None), ) description = github_issues_utils.get_issue_description( test_config_container=test_config_container, metric_container=metric_container, change_point_index=change_point_index, max_results_to_display=( constants._NUM_RESULTS_TO_DISPLAY_ON_ISSUE_DESCRIPTION)) runs_info = next(( line for line in description.split(2 * os.linesep) if re.match(r'timestamp: .*, metric_value: .*', line.strip())), '') pattern = (r'timestamp: .+ (\d{4}), metric_value: (\d+.\d+) <---- Anomaly') match = re.search(pattern, runs_info) self.assertTrue(match) def test_change_point_on_noisy_data(self): def read_csv(path): with FileSystems.open(path) as fp: return pd.read_csv(fp) metric_data = read_csv( 'gs://apache-beam-ml/testing/inputs/test_data_with_noise.csv') metric_values = metric_data[load_test_metrics_utils.VALUE_LABEL].tolist() change_points = find_change_points(metric_values) self.assertEqual(change_points[0], 20) # filter the noise. valid_points = filter_change_points_by_median_threshold( metric_values, change_points) self.assertEqual(len(valid_points), 0) def test_change_point_on_edge_segment(self): data = [1] * 50 + [100] self.assertEqual(find_latest_change_point_index(data), None) if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) os.environ['GITHUB_TOKEN'] = 'fake_token' unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/analyzers/perf_analysis_utils.py0000644000175100017510000003300415206113032025773 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import abc import logging from dataclasses import asdict from dataclasses import dataclass from statistics import median from typing import Optional from typing import Union import pandas as pd import yaml from google.api_core import exceptions from signal_processing_algorithms.energy_statistics.energy_statistics import e_divisive from apache_beam.testing.analyzers import constants from apache_beam.testing.load_tests import load_test_metrics_utils from apache_beam.testing.load_tests.load_test_metrics_utils import BigQueryMetricsPublisher # pylint: disable=ungrouped-imports try: from google.cloud import bigquery except ImportError: bigquery = None # type: ignore @dataclass(frozen=True) class GitHubIssueMetaData: """ This class holds metadata that needs to be published to the BigQuery when a GitHub issue is created on a performance alert. """ issue_timestamp: pd.Timestamp change_point_timestamp: pd.Timestamp test_name: str metric_name: str issue_number: int issue_url: str test_id: str change_point: float @dataclass class ChangePointConfig: """ This class holds the change point configuration parameters. """ min_runs_between_change_points: int = ( constants._DEFAULT_MIN_RUNS_BETWEEN_CHANGE_POINTS) num_runs_in_change_point_window: int = ( constants._DEFAULT_NUM_RUMS_IN_CHANGE_POINT_WINDOW) median_abs_deviation_threshold: int = ( constants._DEFAULT_MEDIAN_ABS_DEVIATION_THRESHOLD) @dataclass class TestConfigContainer: metric_name: str project: str metrics_dataset: str metrics_table: str test_id: str # unique id for each test config. test_description: str test_name: Optional[str] = None labels: Optional[list[str]] = None @dataclass class MetricContainer: """ This class holds the metric values and timestamps for a given metric. Args: metric_values: List of metric values. timestamps: List of pandas timestamps corresponding to the metric values. """ values: list[Union[int, float]] timestamps: list[pd.Timestamp] def sort_by_timestamp(self, in_place=True): """ Sorts the metric values and timestamps in ascending order wrt timestamps. Args: in_place: If True, sort the metric values and timestamps in place. """ timestamps, values = zip(*sorted(zip(self.timestamps, self.values))) if not in_place: return MetricContainer(values=values, timestamps=timestamps) self.timestamps, self.values = zip(*sorted( zip(self.timestamps, self.values))) def is_change_point_in_valid_window( num_runs_in_change_point_window: int, latest_change_point_run: int) -> bool: return num_runs_in_change_point_window > latest_change_point_run def get_existing_issues_data(table_name: str) -> Optional[pd.DataFrame]: """ Finds the most recent GitHub issue created for the test_name. If no table found with name=test_name, return (None, None) else return latest created issue_number along with """ query = f""" SELECT * FROM {constants._BQ_PROJECT_NAME}.{constants._BQ_DATASET}.{table_name} ORDER BY {constants._ISSUE_CREATION_TIMESTAMP_LABEL} DESC LIMIT 10 """ try: if bigquery is None: raise ImportError('Bigquery dependencies are not installed.') client = bigquery.Client() query_job = client.query(query=query) existing_issue_data = query_job.result().to_dataframe() except exceptions.NotFound: # If no table found, that means this is first performance regression # on the current test+metric. return None return existing_issue_data def is_sibling_change_point( previous_change_point_timestamps: list[pd.Timestamp], change_point_index: int, timestamps: list[pd.Timestamp], min_runs_between_change_points: int, test_id: str, ) -> bool: """ Sibling change points are the change points that are close to each other. Search the previous_change_point_timestamps with current observed change point sibling window and determine if it is a duplicate change point or not. timestamps are expected to be in ascending order. Return False if the current observed change point is a duplicate of already reported change points else return True. """ sibling_change_point_min_timestamp = timestamps[max( 0, change_point_index - min_runs_between_change_points)] sibling_change_point_max_timestamp = timestamps[min( change_point_index + min_runs_between_change_points, len(timestamps) - 1)] # Search a list of previous change point timestamps and compare it with # current change point timestamp. We do this in case, if a current change # point is already reported in the past. for previous_change_point_timestamp in previous_change_point_timestamps: if (sibling_change_point_min_timestamp <= previous_change_point_timestamp <= sibling_change_point_max_timestamp): logging.info( 'Performance regression/improvement found for the test ID: %s. ' 'Since the change point timestamp %s ' 'lies within the sibling change point window: %s, ' 'alert is not raised.' % ( test_id, previous_change_point_timestamp.strftime('%Y-%m-%d %H:%M:%S'), ( sibling_change_point_min_timestamp.strftime( '%Y-%m-%d %H:%M:%S'), sibling_change_point_max_timestamp.strftime( '%Y-%m-%d %H:%M:%S')))) return False return True def read_test_config(config_file_path: str) -> dict: """ Reads the config file in which the data required to run the change point analysis is specified. """ with open(config_file_path, 'r') as stream: config = yaml.safe_load(stream) return config def validate_config(keys): return constants._PERF_TEST_KEYS.issubset(keys) def find_change_points(metric_values: list[Union[float, int]]): return e_divisive(metric_values) def find_latest_change_point_index( metric_values: list[Union[float, int]], median_abs_deviation_threshold: int = 2): """ Args: metric_values: Metric values used to run change point analysis. Returns: int: Right most change point index observed on metric_values. """ change_points_indices = find_change_points(metric_values) # reduce noise in the change point analysis by filtering out # the change points that are not significant enough. change_points_indices = filter_change_points_by_median_threshold( metric_values, change_points_indices, threshold=0.1, median_abs_deviation_threshold=median_abs_deviation_threshold) # Consider the latest change point. if not change_points_indices: return None change_points_indices.sort() # Remove the change points that are at the edges of the data. # https://github.com/apache/beam/issues/28757 # Remove this workaround once we have a good solution to deal # with the edge change points. change_point_index = change_points_indices[-1] if is_edge_change_point(change_point_index, len(metric_values), constants._EDGE_SEGMENT_SIZE): logging.info( 'The change point %s is located at the edge of the data with an edge ' 'segment size of %s. This change point will be ignored for now, ' 'awaiting additional data. Should the change point persist after ' 'gathering more data, an alert will be raised.' % (change_point_index, constants._EDGE_SEGMENT_SIZE)) return None return change_point_index def publish_issue_metadata_to_big_query( issue_metadata, table_name, project=constants._BQ_PROJECT_NAME, ): """ Published issue_metadata to BigQuery with table name. """ bq_metrics_publisher = BigQueryMetricsPublisher( project_name=project, dataset=constants._BQ_DATASET, table=table_name, bq_schema=constants._SCHEMA) bq_metrics_publisher.publish([asdict(issue_metadata)]) logging.info( 'GitHub metadata is published to Big Query Dataset %s' ', table %s' % (constants._BQ_DATASET, table_name)) def create_performance_alert( test_config_container: TestConfigContainer, metric_container: MetricContainer, change_point_index: int, existing_issue_number: Optional[int], ) -> tuple[int, str]: """ Creates performance alert on GitHub issues and returns GitHub issue number and issue URL. """ # avoid circular imports # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.testing.analyzers import github_issues_utils description = github_issues_utils.get_issue_description( test_config_container=test_config_container, metric_container=metric_container, change_point_index=change_point_index, max_results_to_display=( constants._NUM_RESULTS_TO_DISPLAY_ON_ISSUE_DESCRIPTION)) issue_number, issue_url = github_issues_utils.report_change_point_on_issues( title=github_issues_utils._ISSUE_TITLE_TEMPLATE.format( test_config_container.test_id, test_config_container.metric_name ), description=description, labels=test_config_container.labels, existing_issue_number=existing_issue_number) logging.info( 'Performance regression/improvement is alerted on issue #%s. Link ' ': %s' % (issue_number, issue_url)) return issue_number, issue_url def filter_change_points_by_median_threshold( data: list[Union[int, float]], change_points: list[int], threshold: float = 0.05, median_abs_deviation_threshold: int = 2, ): """ Reduces the number of change points by filtering out the ones that are not significant enough based on the relative median threshold. Default value of threshold is 0.05. """ valid_change_points = [] for idx in change_points: if idx == 0 or idx == len(data): continue left_segment = data[:idx] right_segment = data[idx:] left_value = median(left_segment) right_value = median(right_segment) # MAD (Median Absolute Deviation) is a robust measure of variability. # A low MAD indicates that the data points are tightly clustered around the # median, while a high MAD suggests greater spread. left_mad = median([abs(x - left_value) for x in left_segment]) right_mad = median([abs(x - right_value) for x in right_segment]) # The change is considered significant only if the absolute difference # between the medians of the two segments is greater than a threshold # times the median absolute deviation (MAD) of the respective segments. # This approach helps to filter out changes that are not statistically # significant, making the detection more robust to noise and outliers. if abs(right_value - left_value) > (median_abs_deviation_threshold * (left_mad + right_mad)): valid_change_points.append(idx) continue return valid_change_points def is_edge_change_point( change_point_index, data_size, edge_segment_size=constants._EDGE_SEGMENT_SIZE): """ Removes the change points that are at the edges of the data. Args: change_point_index: Index of the change point. data_size: Size of the data. edge_segment_size: Size of the edge segment. """ return change_point_index > data_size - edge_segment_size class MetricsFetcher(metaclass=abc.ABCMeta): @abc.abstractmethod def fetch_metric_data( self, *, test_config: TestConfigContainer) -> MetricContainer: """ Define SQL query and fetch the timestamp values and metric values from BigQuery tables. """ raise NotImplementedError class BigQueryMetricsFetcher(MetricsFetcher): def fetch_metric_data( self, *, test_config: TestConfigContainer) -> MetricContainer: """ Args: test_config: TestConfigContainer containing metadata required to fetch metric data from BigQuery. Returns: MetricContainer containing metric values and timestamps. """ project = test_config.project metrics_dataset = test_config.metrics_dataset metrics_table = test_config.metrics_table metric_name = test_config.metric_name query = f""" SELECT * FROM {project}.{metrics_dataset}.{metrics_table} WHERE CONTAINS_SUBSTR(({load_test_metrics_utils.METRICS_TYPE_LABEL}), '{metric_name}') ORDER BY {load_test_metrics_utils.SUBMIT_TIMESTAMP_LABEL} DESC LIMIT {constants._NUM_DATA_POINTS_TO_RUN_CHANGE_POINT_ANALYSIS} """ if bigquery is None: raise ImportError('Bigquery dependencies are not installed.') client = bigquery.Client() query_job = client.query(query=query) metric_data = query_job.result().to_dataframe() # metric_data.sort_values( # by=[load_test_metrics_utils.SUBMIT_TIMESTAMP_LABEL], inplace=True) return MetricContainer( values=metric_data[load_test_metrics_utils.VALUE_LABEL].tolist(), timestamps=metric_data[ load_test_metrics_utils.SUBMIT_TIMESTAMP_LABEL].tolist()) ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.7024417 apache_beam-2.74.0/apache_beam/testing/benchmarks/0000755000175100017510000000000015206113160021451 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/__init__.py0000644000175100017510000000142015206113032023555 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.7034419 apache_beam-2.74.0/apache_beam/testing/benchmarks/chicago_taxi/0000755000175100017510000000000015206113160024073 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/chicago_taxi/__init__.py0000644000175100017510000000142015206113032026177 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/chicago_taxi/preprocess.py0000644000175100017510000002142015206113032026627 0ustar00runnerrunner# Copyright 2019 Google LLC. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Preprocessor applying tf.transform to the chicago_taxi data.""" # pytype: skip-file import argparse import os import tensorflow as tf import tensorflow_transform as transform import tensorflow_transform.beam as tft_beam from tensorflow_transform.coders import example_proto_coder from tensorflow_transform.tf_metadata import dataset_metadata, schema_utils import apache_beam as beam from apache_beam.io.gcp.bigquery import ReadFromBigQuery from apache_beam.metrics.metric import MetricsFilter from apache_beam.testing.load_tests.load_test_metrics_utils import ( MeasureTime, MetricsReader) from trainer import taxi def _fill_in_missing(x): """Replace missing values in a SparseTensor. Fills in missing values of `x` with '' or 0, and converts to a dense tensor. Args: x: A `SparseTensor` of rank 2. Its dense shape should have size at most 1 in the second dimension. Returns: A rank 1 tensor where missing values of `x` have been filled in. """ default_value = '' if x.dtype == tf.string else 0 return tf.squeeze( tf.sparse.to_dense( tf.SparseTensor(x.indices, x.values, [x.dense_shape[0], 1]), default_value), axis=1) def transform_data( input_handle, outfile_prefix, working_dir, schema_file, transform_dir=None, max_rows=None, pipeline_args=None, publish_to_bq=False, project=None, metrics_table=None, metrics_dataset=None): """The main tf.transform method which analyzes and transforms data. Args: input_handle: BigQuery table name to process specified as DATASET.TABLE or path to csv file with input data. outfile_prefix: Filename prefix for emitted transformed examples working_dir: Directory in which transformed examples and transform function will be emitted. schema_file: An file path that contains a text-serialized TensorFlow metadata schema of the input data. transform_dir: Directory in which the transform output is located. If provided, this will load the transform_fn from disk instead of computing it over the data. Hint: this is useful for transforming eval data. max_rows: Number of rows to query from BigQuery pipeline_args: additional DataflowRunner or DirectRunner args passed to the beam pipeline. """ def preprocessing_fn(inputs): """tf.transform's callback function for preprocessing inputs. Args: inputs: map from feature keys to raw not-yet-transformed features. Returns: Map from string feature key to transformed feature operations. """ outputs = {} for key in taxi.DENSE_FLOAT_FEATURE_KEYS: # Preserve this feature as a dense float, setting nan's to the mean. outputs[taxi.transformed_name(key)] = transform.scale_to_z_score( _fill_in_missing(inputs[key])) for key in taxi.VOCAB_FEATURE_KEYS: # Build a vocabulary for this feature. outputs[taxi.transformed_name( key)] = transform.compute_and_apply_vocabulary( _fill_in_missing(inputs[key]), top_k=taxi.VOCAB_SIZE, num_oov_buckets=taxi.OOV_SIZE) for key in taxi.BUCKET_FEATURE_KEYS: outputs[taxi.transformed_name(key)] = transform.bucketize( _fill_in_missing(inputs[key]), taxi.FEATURE_BUCKET_COUNT) for key in taxi.CATEGORICAL_FEATURE_KEYS: outputs[taxi.transformed_name(key)] = _fill_in_missing(inputs[key]) # Was this passenger a big tipper? taxi_fare = _fill_in_missing(inputs[taxi.FARE_KEY]) tips = _fill_in_missing(inputs[taxi.LABEL_KEY]) outputs[taxi.transformed_name(taxi.LABEL_KEY)] = tf.where( tf.is_nan(taxi_fare), tf.cast(tf.zeros_like(taxi_fare), tf.int64), # Test if the tip was > 20% of the fare. tf.cast( tf.greater(tips, tf.multiply(taxi_fare, tf.constant(0.2))), tf.int64)) return outputs namespace = metrics_table metrics_monitor = None if publish_to_bq: metrics_monitor = MetricsReader( publish_to_bq=publish_to_bq, project_name=project, bq_table=metrics_table, bq_dataset=metrics_dataset, namespace=namespace, filters=MetricsFilter().with_namespace(namespace)) schema = taxi.read_schema(schema_file) raw_feature_spec = taxi.get_raw_feature_spec(schema) raw_schema = schema_utils.schema_from_feature_spec(raw_feature_spec) raw_data_metadata = dataset_metadata.DatasetMetadata(raw_schema) pipeline = beam.Pipeline(argv=pipeline_args) with tft_beam.Context(temp_dir=working_dir): query = taxi.make_sql(input_handle, max_rows, for_eval=False) raw_data = ( pipeline | 'ReadBigQuery' >> ReadFromBigQuery( query=query, project=project, use_standard_sql=True) | 'Measure time: start' >> beam.ParDo(MeasureTime(namespace))) decode_transform = beam.Map( taxi.clean_raw_data_dict, raw_feature_spec=raw_feature_spec) if transform_dir is None: decoded_data = raw_data | 'DecodeForAnalyze' >> decode_transform transform_fn = ((decoded_data, raw_data_metadata) | ('Analyze' >> tft_beam.AnalyzeDataset(preprocessing_fn))) _ = ( transform_fn | ('WriteTransformFn' >> tft_beam.WriteTransformFn(working_dir))) else: transform_fn = pipeline | tft_beam.ReadTransformFn(transform_dir) # Shuffling the data before materialization will improve Training # effectiveness downstream. Here we shuffle the raw_data (as opposed to # decoded data) since it has a compact representation. shuffled_data = raw_data | 'RandomizeData' >> beam.transforms.Reshuffle() decoded_data = shuffled_data | 'DecodeForTransform' >> decode_transform (transformed_data, transformed_metadata) = (((decoded_data, raw_data_metadata), transform_fn) | 'Transform' >> tft_beam.TransformDataset()) coder = example_proto_coder.ExampleProtoCoder(transformed_metadata.schema) _ = ( transformed_data | 'SerializeExamples' >> beam.Map(coder.encode) | 'Measure time: end' >> beam.ParDo(MeasureTime(namespace)) | 'WriteExamples' >> beam.io.WriteToTFRecord( os.path.join(working_dir, outfile_prefix), file_name_suffix='.gz')) result = pipeline.run() result.wait_until_finish() if metrics_monitor: metrics_monitor.publish_metrics(result) def main(): tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO) parser = argparse.ArgumentParser() parser.add_argument( '--input', help=('Input BigQuery table to process specified as: ' 'DATASET.TABLE')) parser.add_argument( '--schema_file', help='File holding the schema for the input data') parser.add_argument( '--output_dir', help=( 'Directory in which transformed examples and function ' 'will be emitted.')) parser.add_argument( '--outfile_prefix', help='Filename prefix for emitted transformed examples') parser.add_argument( '--transform_dir', required=False, default=None, help='Directory in which the transform output is located') parser.add_argument( '--max_rows', help='Number of rows to query from BigQuery', default=None, type=int) parser.add_argument( '--publish_to_big_query', help='Whether to publish to BQ', default=None, type=bool) parser.add_argument( '--metrics_dataset', help='BQ dataset', default=None, type=str) parser.add_argument( '--metrics_table', help='BQ table', default=None, type=str) parser.add_argument( '--metric_reporting_project', help='BQ table project', default=None, type=str) known_args, pipeline_args = parser.parse_known_args() transform_data( input_handle=known_args.input, outfile_prefix=known_args.outfile_prefix, working_dir=known_args.output_dir, schema_file=known_args.schema_file, transform_dir=known_args.transform_dir, max_rows=known_args.max_rows, pipeline_args=pipeline_args, publish_to_bq=known_args.publish_to_big_query, metrics_dataset=known_args.metrics_dataset, metrics_table=known_args.metrics_table, project=known_args.metric_reporting_project) if __name__ == '__main__': main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/chicago_taxi/process_tfma.py0000644000175100017510000001417115206113032027134 0ustar00runnerrunner# Copyright 2019 Google LLC. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Runs a batch job for performing Tensorflow Model Analysis.""" # pytype: skip-file import argparse import tensorflow as tf import tensorflow_model_analysis as tfma from tensorflow_model_analysis.evaluators import evaluator import apache_beam as beam from apache_beam.io.gcp.bigquery import ReadFromBigQuery from apache_beam.metrics.metric import MetricsFilter from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime from apache_beam.testing.load_tests.load_test_metrics_utils import MetricsReader from trainer import taxi def process_tfma( schema_file, big_query_table=None, eval_model_dir=None, max_eval_rows=None, pipeline_args=None, publish_to_bq=False, project=None, metrics_table=None, metrics_dataset=None): """Runs a batch job to evaluate the eval_model against the given input. Args: schema_file: A file containing a text-serialized Schema that describes the eval data. big_query_table: A BigQuery table name specified as DATASET.TABLE which should be the input for evaluation. This can only be set if input_csv is None. eval_model_dir: A directory where the eval model is located. max_eval_rows: Number of rows to query from BigQuery. pipeline_args: additional DataflowRunner or DirectRunner args passed to the beam pipeline. publish_to_bq: project: metrics_dataset: metrics_table: Raises: ValueError: if input_csv and big_query_table are not specified correctly. """ if big_query_table is None: raise ValueError('--big_query_table should be provided.') slice_spec = [ tfma.slicer.SingleSliceSpec(), tfma.slicer.SingleSliceSpec(columns=['trip_start_hour']) ] metrics_namespace = metrics_table schema = taxi.read_schema(schema_file) eval_shared_model = tfma.default_eval_shared_model( eval_saved_model_path=eval_model_dir, add_metrics_callbacks=[ tfma.post_export_metrics.calibration_plot_and_prediction_histogram(), tfma.post_export_metrics.auc_plots() ]) metrics_monitor = None if publish_to_bq: metrics_monitor = MetricsReader( publish_to_bq=publish_to_bq, project_name=project, bq_table=metrics_table, bq_dataset=metrics_dataset, namespace=metrics_namespace, filters=MetricsFilter().with_namespace(metrics_namespace)) pipeline = beam.Pipeline(argv=pipeline_args) query = taxi.make_sql(big_query_table, max_eval_rows, for_eval=True) raw_feature_spec = taxi.get_raw_feature_spec(schema) raw_data = ( pipeline | 'ReadBigQuery' >> ReadFromBigQuery( query=query, project=project, use_standard_sql=True) | 'Measure time: Start' >> beam.ParDo(MeasureTime(metrics_namespace)) | 'CleanData' >> beam.Map(lambda x: (taxi.clean_raw_data_dict(x, raw_feature_spec)))) # Examples must be in clean tf-example format. coder = taxi.make_proto_coder(schema) # Prepare arguments for Extract, Evaluate and Write steps extractors = tfma.default_extractors( eval_shared_model=eval_shared_model, slice_spec=slice_spec, desired_batch_size=None, materialize=False) evaluators = tfma.default_evaluators( eval_shared_model=eval_shared_model, desired_batch_size=None, num_bootstrap_samples=1) _ = ( raw_data | 'ToSerializedTFExample' >> beam.Map(coder.encode) | 'Extract Results' >> tfma.InputsToExtracts() | 'Extract and evaluate' >> tfma.ExtractAndEvaluate( extractors=extractors, evaluators=evaluators) | 'Map Evaluations to PCollection' >> MapEvalToPCollection() | 'Measure time: End' >> beam.ParDo(MeasureTime(metrics_namespace))) result = pipeline.run() result.wait_until_finish() if metrics_monitor: metrics_monitor.publish_metrics(result) @beam.ptransform_fn @beam.typehints.with_input_types(evaluator.Evaluation) @beam.typehints.with_output_types(beam.typehints.Any) def MapEvalToPCollection( # pylint: disable=invalid-name evaluation): return evaluation['metrics'] def main(): tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO) parser = argparse.ArgumentParser() parser.add_argument( '--eval_model_dir', help='Input path to the model which will be evaluated.') parser.add_argument( '--big_query_table', help='BigQuery path to input examples which will be evaluated.') parser.add_argument( '--max_eval_rows', help='Maximum number of rows to evaluate on.', default=None, type=int) parser.add_argument( '--schema_file', help='File holding the schema for the input data') parser.add_argument( '--publish_to_big_query', help='Whether to publish to BQ', default=None, type=bool) parser.add_argument( '--metrics_dataset', help='BQ dataset', default=None, type=str) parser.add_argument( '--metrics_table', help='BQ table for storing metrics', default=None, type=str) parser.add_argument( '--metric_reporting_project', help='BQ table project', default=None, type=str) known_args, pipeline_args = parser.parse_known_args() process_tfma( big_query_table=known_args.big_query_table, eval_model_dir=known_args.eval_model_dir, max_eval_rows=known_args.max_eval_rows, schema_file=known_args.schema_file, pipeline_args=pipeline_args, publish_to_bq=known_args.publish_to_big_query, metrics_table=known_args.metrics_table, metrics_dataset=known_args.metrics_dataset, project=known_args.metric_reporting_project) if __name__ == '__main__': main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/chicago_taxi/setup.py0000644000175100017510000000246215206113032025607 0ustar00runnerrunner# Copyright 2019 Google LLC. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Setup dependencies for local and cloud deployment.""" # pytype: skip-file import setuptools TF_VERSION = '1.14.0' if __name__ == '__main__': setuptools.setup( name='beam_chicago_taxi', version='0.13.0', packages=setuptools.find_packages(), install_requires=[ 'jupyter>=1.0,<2', 'notebook>=5.7.8,<5.8', 'tensorflow==' + TF_VERSION, 'tensorflow-data-validation>=0.13.1,<0.14', 'tensorflow-metadata>=0.13.0,<0.14', 'tensorflow-model-analysis>=0.13.2,<0.14', 'tensorflow-serving-api>=1.13.0,<1.14', 'tensorflow-transform>=0.13.0,<0.14', ], python_requires='>=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,<4', ) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/chicago_taxi/tfdv_analyze_and_validate.py0000644000175100017510000001572415206113032031635 0ustar00runnerrunner# Copyright 2019 Google LLC. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Compute stats, infer schema, and validate stats for chicago taxi example.""" # pytype: skip-file import argparse import numpy as np import tensorflow as tf from tensorflow.python.lib.io import file_io import tensorflow_data_validation as tfdv from tensorflow_metadata.proto.v0 import statistics_pb2 import apache_beam as beam from apache_beam.io.gcp.bigquery import ReadFromBigQuery from apache_beam.metrics.metric import MetricsFilter from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime from apache_beam.testing.load_tests.load_test_metrics_utils import MetricsReader from google.protobuf import text_format from trainer import taxi def infer_schema(stats_path, schema_path): """Infers a schema from stats in stats_path. Args: stats_path: Location of the stats used to infer the schema. schema_path: Location where the inferred schema is materialized. """ print('Infering schema from statistics.') schema = tfdv.infer_schema( tfdv.load_statistics(stats_path), infer_feature_shape=False) print(text_format.MessageToString(schema)) print('Writing schema to output path.') file_io.write_string_to_file(schema_path, text_format.MessageToString(schema)) def validate_stats(stats_path, schema_path, anomalies_path): """Validates the statistics against the schema and materializes anomalies. Args: stats_path: Location of the stats used to infer the schema. schema_path: Location of the schema to be used for validation. anomalies_path: Location where the detected anomalies are materialized. """ print('Validating schema against the computed statistics.') schema = taxi.read_schema(schema_path) stats = tfdv.load_statistics(stats_path) anomalies = tfdv.validate_statistics(stats, schema) print('Detected following anomalies:') print(text_format.MessageToString(anomalies)) print('Writing anomalies to anomalies path.') file_io.write_string_to_file( anomalies_path, text_format.MessageToString(anomalies)) def compute_stats( input_handle, stats_path, max_rows=None, for_eval=False, pipeline_args=None, publish_to_bq=None, metrics_dataset=None, metrics_table=None, project=None): """Computes statistics on the input data. Args: input_handle: BigQuery table name to process specified as DATASET.TABLE or path to csv file with input data. stats_path: Directory in which stats are materialized. max_rows: Number of rows to query from BigQuery for_eval: Query for eval set rows from BigQuery pipeline_args: additional DataflowRunner or DirectRunner args passed to the beam pipeline. """ namespace = metrics_table pipeline = beam.Pipeline(argv=pipeline_args) metrics_monitor = None if publish_to_bq: metrics_monitor = MetricsReader( publish_to_bq=publish_to_bq, project_name=project, bq_table=metrics_table, bq_dataset=metrics_dataset, namespace=namespace, filters=MetricsFilter().with_namespace(namespace), ) query = taxi.make_sql( table_name=input_handle, max_rows=max_rows, for_eval=for_eval) raw_data = ( pipeline | 'ReadBigQuery' >> ReadFromBigQuery( query=query, project=project, use_standard_sql=True) | 'Measure time: Start' >> beam.ParDo(MeasureTime(namespace)) | 'ConvertToTFDVInput' >> beam.Map( lambda x: {key: np.asarray([x[key]]) for key in x if x[key] is not None})) _ = ( raw_data | 'GenerateStatistics' >> tfdv.GenerateStatistics() | 'Measure time: End' >> beam.ParDo(MeasureTime(namespace)) | 'WriteStatsOutput' >> beam.io.WriteToTFRecord( stats_path, shard_name_template='', coder=beam.coders.ProtoCoder( statistics_pb2.DatasetFeatureStatisticsList))) result = pipeline.run() result.wait_until_finish() if metrics_monitor: metrics_monitor.publish_metrics(result) def main(): tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO) parser = argparse.ArgumentParser() parser.add_argument( '--input', help=( 'Input BigQuery table to process specified as: ' 'DATASET.TABLE or path to csv file with input data.')) parser.add_argument( '--stats_path', help='Location for the computed stats to be materialized.') parser.add_argument( '--for_eval', help='Query for eval set rows from BigQuery', action='store_true') parser.add_argument( '--max_rows', help='Number of rows to query from BigQuery', default=None, type=int) parser.add_argument( '--schema_path', help='Location for the computed schema is located.', default=None, type=str) parser.add_argument( '--infer_schema', help='If specified, also infers a schema based on the computed stats.', action='store_true') parser.add_argument( '--validate_stats', help='If specified, also validates the stats against the schema.', action='store_true') parser.add_argument( '--anomalies_path', help='Location for detected anomalies are materialized.', default=None, type=str) parser.add_argument( '--publish_to_big_query', help='Whether to publish to BQ', default=None, type=bool) parser.add_argument( '--metrics_dataset', help='BQ dataset', default=None, type=str) parser.add_argument( '--metrics_table', help='BQ table', default=None, type=str) parser.add_argument( '--metric_reporting_project', help='BQ table project', default=None, type=str) known_args, pipeline_args = parser.parse_known_args() compute_stats( input_handle=known_args.input, stats_path=known_args.stats_path, max_rows=known_args.max_rows, for_eval=known_args.for_eval, pipeline_args=pipeline_args, publish_to_bq=known_args.publish_to_big_query, metrics_dataset=known_args.metrics_dataset, metrics_table=known_args.metrics_table, project=known_args.metric_reporting_project) print('Stats computation done.') if known_args.infer_schema: infer_schema( stats_path=known_args.stats_path, schema_path=known_args.schema_path) if known_args.validate_stats: validate_stats( stats_path=known_args.stats_path, schema_path=known_args.schema_path, anomalies_path=known_args.anomalies_path) if __name__ == '__main__': main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.704442 apache_beam-2.74.0/apache_beam/testing/benchmarks/chicago_taxi/trainer/0000755000175100017510000000000015206113160025537 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/chicago_taxi/trainer/__init__.py0000644000175100017510000000142015206113032027643 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/chicago_taxi/trainer/model.py0000644000175100017510000001377215206113032027221 0ustar00runnerrunner# Copyright 2019 Google LLC. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Defines the model used to predict who will tip in the Chicago Taxi demo.""" # pytype: skip-file import tensorflow as tf import tensorflow_model_analysis as tfma from tensorflow import estimator as tf_estimator from trainer import taxi def build_estimator(tf_transform_output, config, hidden_units=None): """Build an estimator for predicting the tipping behavior of taxi riders. Args: tf_transform_output: A TFTransformOutput. config: tf.contrib.learn.RunConfig defining the runtime environment for the estimator (including model_dir). hidden_units: [int], the layer sizes of the DNN (input layer first) Returns: Resulting DNNLinearCombinedClassifier. """ transformed_feature_spec = ( tf_transform_output.transformed_feature_spec().copy()) transformed_feature_spec.pop(taxi.transformed_name(taxi.LABEL_KEY)) real_valued_columns = [ tf.feature_column.numeric_column(key, shape=()) for key in taxi.transformed_names(taxi.DENSE_FLOAT_FEATURE_KEYS) ] categorical_columns = [ tf.feature_column.categorical_column_with_identity( key, num_buckets=taxi.VOCAB_SIZE + taxi.OOV_SIZE, default_value=0) for key in taxi.transformed_names(taxi.VOCAB_FEATURE_KEYS) ] categorical_columns += [ tf.feature_column.categorical_column_with_identity( key, num_buckets=taxi.FEATURE_BUCKET_COUNT, default_value=0) for key in taxi.transformed_names(taxi.BUCKET_FEATURE_KEYS) ] categorical_columns += [ tf.feature_column.categorical_column_with_identity( key, num_buckets=num_buckets, default_value=0) for key, num_buckets in zip( taxi.transformed_names(taxi.CATEGORICAL_FEATURE_KEYS), # taxi.MAX_CATEGORICAL_FEATURE_VALUES) ] return tf_estimator.DNNLinearCombinedClassifier( config=config, linear_feature_columns=categorical_columns, dnn_feature_columns=real_valued_columns, dnn_hidden_units=hidden_units or [100, 70, 50, 25]) def example_serving_receiver_fn(tf_transform_output, schema): """Build the serving in inputs. Args: tf_transform_output: A TFTransformOutput. schema: the schema of the input data. Returns: Tensorflow graph which parses examples, applying tf-transform to them. """ raw_feature_spec = taxi.get_raw_feature_spec(schema) raw_feature_spec.pop(taxi.LABEL_KEY) raw_input_fn = tf_estimator.export.build_parsing_serving_input_receiver_fn( raw_feature_spec, default_batch_size=None) serving_input_receiver = raw_input_fn() transformed_features = tf_transform_output.transform_raw_features( serving_input_receiver.features) return tf_estimator.export.ServingInputReceiver( transformed_features, serving_input_receiver.receiver_tensors) def eval_input_receiver_fn(tf_transform_output, schema): """Build everything needed for the tf-model-analysis to run the model. Args: tf_transform_output: A TFTransformOutput. schema: the schema of the input data. Returns: EvalInputReceiver function, which contains: - Tensorflow graph which parses raw untransformed features, applies the tf-transform preprocessing operators. - Set of raw, untransformed features. - Label against which predictions will be compared. """ # Notice that the inputs are raw features, not transformed features here. raw_feature_spec = taxi.get_raw_feature_spec(schema) serialized_tf_example = tf.placeholder( dtype=tf.string, shape=[None], name='input_example_tensor') # Add a parse_example operator to the tensorflow graph, which will parse # raw, untransformed, tf examples. features = tf.parse_example(serialized_tf_example, raw_feature_spec) # Now that we have our raw examples, process them through the tf-transform # function computed during the preprocessing step. transformed_features = tf_transform_output.transform_raw_features(features) # The key name MUST be 'examples'. receiver_tensors = {'examples': serialized_tf_example} # NOTE: Model is driven by transformed features (since training works on the # materialized output of TFT, but slicing will happen on raw features. features.update(transformed_features) return tfma.export.EvalInputReceiver( features=features, receiver_tensors=receiver_tensors, labels=transformed_features[taxi.transformed_name(taxi.LABEL_KEY)]) def _gzip_reader_fn(): """Small utility returning a record reader that can read gzip'ed files.""" return tf.TFRecordReader( options=tf.python_io.TFRecordOptions( compression_type=tf.python_io.TFRecordCompressionType.GZIP)) def input_fn(filenames, tf_transform_output, batch_size=200): """Generates features and labels for training or evaluation. Args: filenames: [str] list of CSV files to read data from. tf_transform_output: A TFTransformOutput. batch_size: int First dimension size of the Tensors returned by input_fn Returns: A (features, indices) tuple where features is a dictionary of Tensors, and indices is a single Tensor of label indices. """ transformed_feature_spec = ( tf_transform_output.transformed_feature_spec().copy()) transformed_features = tf.contrib.learn.io.read_batch_features( filenames, batch_size, transformed_feature_spec, reader=_gzip_reader_fn) # We pop the label because we do not want to use it as a feature while we're # training. return transformed_features, transformed_features.pop( taxi.transformed_name(taxi.LABEL_KEY)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/chicago_taxi/trainer/task.py0000644000175100017510000001236215206113032027055 0ustar00runnerrunner# Copyright 2019 Google LLC. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Trainer for the chicago_taxi demo.""" # pytype: skip-file import argparse import os import tensorflow as tf import tensorflow_model_analysis as tfma import tensorflow_transform as tft from tensorflow import estimator as tf_estimator from trainer import model from trainer import taxi SERVING_MODEL_DIR = 'serving_model_dir' EVAL_MODEL_DIR = 'eval_model_dir' TRAIN_BATCH_SIZE = 40 EVAL_BATCH_SIZE = 40 # Number of nodes in the first layer of the DNN FIRST_DNN_LAYER_SIZE = 100 NUM_DNN_LAYERS = 4 DNN_DECAY_FACTOR = 0.7 def train_and_maybe_evaluate(hparams): """Run the training and evaluate using the high level API. Args: hparams: Holds hyperparameters used to train the model as name/value pairs. Returns: The estimator that was used for training (and maybe eval) """ schema = taxi.read_schema(hparams.schema_file) tf_transform_output = tft.TFTransformOutput(hparams.tf_transform_dir) train_input = lambda: model.input_fn( hparams.train_files, tf_transform_output, batch_size=TRAIN_BATCH_SIZE) eval_input = lambda: model.input_fn( hparams.eval_files, tf_transform_output, batch_size=EVAL_BATCH_SIZE) train_spec = tf_estimator.TrainSpec( train_input, max_steps=hparams.train_steps) serving_receiver_fn = lambda: model.example_serving_receiver_fn( tf_transform_output, schema) exporter = tf_estimator.FinalExporter('chicago-taxi', serving_receiver_fn) eval_spec = tf_estimator.EvalSpec( eval_input, steps=hparams.eval_steps, exporters=[exporter], name='chicago-taxi-eval') run_config = tf_estimator.RunConfig( save_checkpoints_steps=999, keep_checkpoint_max=1) serving_model_dir = os.path.join(hparams.output_dir, SERVING_MODEL_DIR) run_config = run_config.replace(model_dir=serving_model_dir) estimator = model.build_estimator( tf_transform_output, # Construct layers sizes with exponetial decay hidden_units=[ max(2, int(FIRST_DNN_LAYER_SIZE * DNN_DECAY_FACTOR**i)) for i in range(NUM_DNN_LAYERS) ], config=run_config) tf_estimator.train_and_evaluate(estimator, train_spec, eval_spec) return estimator def run_experiment(hparams): """Train the model then export it for tf.model_analysis evaluation. Args: hparams: Holds hyperparameters used to train the model as name/value pairs. """ estimator = train_and_maybe_evaluate(hparams) schema = taxi.read_schema(hparams.schema_file) tf_transform_output = tft.TFTransformOutput(hparams.tf_transform_dir) # Save a model for tfma eval eval_model_dir = os.path.join(hparams.output_dir, EVAL_MODEL_DIR) receiver_fn = lambda: model.eval_input_receiver_fn( tf_transform_output, schema) tfma.export.export_eval_savedmodel( estimator=estimator, export_dir_base=eval_model_dir, eval_input_receiver_fn=receiver_fn) def main(): parser = argparse.ArgumentParser() # Input Arguments parser.add_argument( '--train-files', help='GCS or local paths to training data', nargs='+', required=True) parser.add_argument( '--tf-transform-dir', help='Tf-transform directory with model from preprocessing step', required=True) parser.add_argument( '--output-dir', help="""\ Directory under which which the serving model (under /serving_model_dir)\ and the tf-mode-analysis model (under /eval_model_dir) will be written\ """, required=True) parser.add_argument( '--eval-files', help='GCS or local paths to evaluation data', nargs='+', required=True) # Training arguments parser.add_argument( '--job-dir', help='GCS location to write checkpoints and export models', required=True) # Argument to turn on all logging parser.add_argument( '--verbosity', choices=['DEBUG', 'ERROR', 'FATAL', 'INFO', 'WARN'], default='INFO', ) # Experiment arguments parser.add_argument( '--train-steps', help='Count of steps to run the training job for', required=True, type=int) parser.add_argument( '--eval-steps', help='Number of steps to run evalution for at each checkpoint', default=100, type=int) parser.add_argument( '--schema-file', help='File holding the schema for the input data') args = parser.parse_args() # Set python level verbosity tf.compat.v1.logging.set_verbosity(args.verbosity) # Set C++ Graph Execution level verbosity os.environ['TF_CPP_MIN_LOG_LEVEL'] = str( getattr(tf.compat.v1.logging, args.verbosity) / 10) # Run the training job hparams = tf.contrib.training.HParams(**args.__dict__) run_experiment(hparams) if __name__ == '__main__': main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/chicago_taxi/trainer/taxi.py0000644000175100017510000001311315206113032027053 0ustar00runnerrunner# Copyright 2019 Google LLC. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Utility and schema methods for the chicago_taxi sample.""" # pytype: skip-file from tensorflow_transform import coders as tft_coders from tensorflow_transform.tf_metadata import schema_utils from google.protobuf import text_format from tensorflow.python.lib.io import file_io from tensorflow_metadata.proto.v0 import schema_pb2 # Categorical features are assumed to each have a maximum value in the dataset. MAX_CATEGORICAL_FEATURE_VALUES = [24, 31, 12] CATEGORICAL_FEATURE_KEYS = [ 'trip_start_hour', 'trip_start_day', 'trip_start_month', 'pickup_census_tract', 'dropoff_census_tract', 'pickup_community_area', 'dropoff_community_area' ] DENSE_FLOAT_FEATURE_KEYS = ['trip_miles', 'fare', 'trip_seconds'] # Number of buckets used by tf.transform for encoding each feature. FEATURE_BUCKET_COUNT = 10 BUCKET_FEATURE_KEYS = [ 'pickup_latitude', 'pickup_longitude', 'dropoff_latitude', 'dropoff_longitude' ] # Number of vocabulary terms used for encoding VOCAB_FEATURES by tf.transform VOCAB_SIZE = 1000 # Count of out-of-vocab buckets in which unrecognized VOCAB_FEATURES are hashed. OOV_SIZE = 10 VOCAB_FEATURE_KEYS = [ 'payment_type', 'company', ] LABEL_KEY = 'tips' FARE_KEY = 'fare' CSV_COLUMN_NAMES = [ 'pickup_community_area', 'fare', 'trip_start_month', 'trip_start_hour', 'trip_start_day', 'trip_start_timestamp', 'pickup_latitude', 'pickup_longitude', 'dropoff_latitude', 'dropoff_longitude', 'trip_miles', 'pickup_census_tract', 'dropoff_census_tract', 'payment_type', 'company', 'trip_seconds', 'dropoff_community_area', 'tips', ] def transformed_name(key): return key + '_xf' def transformed_names(keys): return [transformed_name(key) for key in keys] # Tf.Transform considers these features as "raw" def get_raw_feature_spec(schema): return schema_utils.schema_as_feature_spec(schema).feature_spec def make_proto_coder(schema): raw_feature_spec = get_raw_feature_spec(schema) raw_schema = schema_utils.schema_from_feature_spec(raw_feature_spec) return tft_coders.ExampleProtoCoder(raw_schema) def make_csv_coder(schema): """Return a coder for tf.transform to read csv files.""" raw_feature_spec = get_raw_feature_spec(schema) parsing_schema = schema_utils.schema_from_feature_spec(raw_feature_spec) return tft_coders.CsvCoder(CSV_COLUMN_NAMES, parsing_schema) def clean_raw_data_dict(input_dict, raw_feature_spec): """Clean raw data dict.""" output_dict = {} for key in raw_feature_spec: if key not in input_dict or not input_dict[key]: output_dict[key] = [] else: output_dict[key] = [input_dict[key]] return output_dict def make_sql(table_name, max_rows=None, for_eval=False): """Creates the sql command for pulling data from BigQuery. Args: table_name: BigQuery table name max_rows: if set, limits the number of rows pulled from BigQuery for_eval: True if this is for evaluation, false otherwise Returns: sql command as string """ if for_eval: # 1/3 of the dataset used for eval where_clause = 'WHERE MOD(FARM_FINGERPRINT(unique_key), 3) = 0 ' \ 'AND pickup_latitude is not null AND pickup_longitude ' \ 'is not null AND dropoff_latitude is not null ' \ 'AND dropoff_longitude is not null' else: # 2/3 of the dataset used for training where_clause = 'WHERE MOD(FARM_FINGERPRINT(unique_key), 3) > 0 ' \ 'AND pickup_latitude is not null AND pickup_longitude ' \ 'is not null AND dropoff_latitude is not null ' \ 'AND dropoff_longitude is not null' limit_clause = '' if max_rows: limit_clause = 'LIMIT {max_rows}'.format(max_rows=max_rows) return """ SELECT CAST(pickup_community_area AS string) AS pickup_community_area, CAST(dropoff_community_area AS string) AS dropoff_community_area, CAST(pickup_census_tract AS string) AS pickup_census_tract, CAST(dropoff_census_tract AS string) AS dropoff_census_tract, fare, EXTRACT(MONTH FROM trip_start_timestamp) AS trip_start_month, EXTRACT(HOUR FROM trip_start_timestamp) AS trip_start_hour, EXTRACT(DAYOFWEEK FROM trip_start_timestamp) AS trip_start_day, UNIX_SECONDS(trip_start_timestamp) AS trip_start_timestamp, pickup_latitude, pickup_longitude, dropoff_latitude, dropoff_longitude, trip_miles, payment_type, company, trip_seconds, tips FROM `{table_name}` {where_clause} {limit_clause} """.format( table_name=table_name, where_clause=where_clause, limit_clause=limit_clause) def read_schema(path): """Reads a schema from the provided location. Args: path: The location of the file holding a serialized Schema proto. Returns: An instance of Schema or None if the input argument is None """ result = schema_pb2.Schema() contents = file_io.read_file_to_string(path) text_format.Parse(contents, result) return result ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.7054422 apache_beam-2.74.0/apache_beam/testing/benchmarks/cloudml/0000755000175100017510000000000015206113160023110 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/cloudml/__init__.py0000644000175100017510000000142015206113032025214 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/cloudml/cloudml_benchmark_constants_lib.py0000644000175100017510000000210115206113032032045 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A common file for CloudML benchmarks. This file contains constants for pipeline paths, dependency locations and test data paths. """ INPUT_CRITEO_SMALL = 'train10.tsv' INPUT_CRITEO_SMALL_100MB = '100mb/train.txt' INPUT_CRITEO_10GB = '10gb/train.txt' FREQUENCY_THRESHOLD = '5' ENABLE_SHUFFLE = True ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/cloudml/cloudml_benchmark_test.py0000644000175100017510000001113115206113032030165 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import os import time import unittest import uuid import pytest try: import apache_beam.testing.benchmarks.cloudml.cloudml_benchmark_constants_lib as lib from apache_beam.testing.benchmarks.cloudml.pipelines import workflow from apache_beam.testing.load_tests.load_test_metrics_utils import InfluxDBMetricsPublisherOptions from apache_beam.testing.load_tests.load_test_metrics_utils import MetricsReader from apache_beam.testing.test_pipeline import TestPipeline except ImportError: # pylint: disable=bare-except raise unittest.SkipTest('Dependencies are not installed') _INPUT_GCS_BUCKET_ROOT = 'gs://apache-beam-ml/datasets/cloudml/criteo' _OUTPUT_GCS_BUCKET_ROOT = 'gs://temp-storage-for-end-to-end-tests/tft/' def _publish_metrics(pipeline, metric_value, metrics_table, metric_name): influx_options = InfluxDBMetricsPublisherOptions( metrics_table, pipeline.get_option('influx_db_name'), pipeline.get_option('influx_hostname'), os.getenv('INFLUXDB_USER'), os.getenv('INFLUXDB_USER_PASSWORD'), ) metric_reader = MetricsReader( project_name=pipeline.get_option('project'), bq_table=metrics_table, bq_dataset=pipeline.get_option('metrics_dataset'), publish_to_bq=True, influxdb_options=influx_options, ) metric_reader.publish_values([( metric_name, metric_value, )]) @pytest.mark.uses_tft class CloudMLTFTBenchmarkTest(unittest.TestCase): def test_cloudml_benchmark_criteo_small(self): test_pipeline = TestPipeline(is_integration_test=True) extra_opts = {} extra_opts['input'] = os.path.join( _INPUT_GCS_BUCKET_ROOT, lib.INPUT_CRITEO_SMALL) extra_opts['benchmark_type'] = 'tft' extra_opts['classifier'] = 'criteo' extra_opts['frequency_threshold'] = 0 extra_opts['output'] = os.path.join( _OUTPUT_GCS_BUCKET_ROOT, uuid.uuid4().hex) start_time = time.time() workflow.run(test_pipeline.get_full_options_as_args(**extra_opts)) end_time = time.time() metrics_table = 'cloudml_benchmark_criteo_small' _publish_metrics( pipeline=test_pipeline, metric_value=end_time - start_time, metrics_table=metrics_table, metric_name='runtime_sec') def test_cloudml_benchmark_cirteo_no_shuffle_10GB(self): test_pipeline = TestPipeline(is_integration_test=True) extra_opts = {} extra_opts['input'] = os.path.join( _INPUT_GCS_BUCKET_ROOT, lib.INPUT_CRITEO_10GB) extra_opts['benchmark_type'] = 'tft' extra_opts['classifier'] = 'criteo' extra_opts['frequency_threshold'] = 0 extra_opts['output'] = os.path.join( _OUTPUT_GCS_BUCKET_ROOT, uuid.uuid4().hex) extra_opts['shuffle'] = False start_time = time.time() workflow.run(test_pipeline.get_full_options_as_args(**extra_opts)) end_time = time.time() metrics_table = 'cloudml_benchmark_cirteo_no_shuffle_10GB' _publish_metrics( pipeline=test_pipeline, metric_value=end_time - start_time, metrics_table=metrics_table, metric_name='runtime_sec') def test_cloudml_benchmark_criteo_10GB(self): test_pipeline = TestPipeline(is_integration_test=True) extra_opts = {} extra_opts['input'] = os.path.join( _INPUT_GCS_BUCKET_ROOT, lib.INPUT_CRITEO_10GB) extra_opts['benchmark_type'] = 'tft' extra_opts['classifier'] = 'criteo' extra_opts['frequency_threshold'] = 0 extra_opts['output'] = os.path.join( _OUTPUT_GCS_BUCKET_ROOT, uuid.uuid4().hex) start_time = time.time() workflow.run(test_pipeline.get_full_options_as_args(**extra_opts)) end_time = time.time() metrics_table = 'cloudml_benchmark_criteo_10GB' _publish_metrics( pipeline=test_pipeline, metric_value=end_time - start_time, metrics_table=metrics_table, metric_name='runtime_sec') if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.706442 apache_beam-2.74.0/apache_beam/testing/benchmarks/cloudml/criteo_tft/0000755000175100017510000000000015206113160025252 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/cloudml/criteo_tft/__init__.py0000644000175100017510000000142015206113032027356 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/cloudml/criteo_tft/criteo.py0000644000175100017510000001212715206113032027112 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Schema and transform definition for the Criteo dataset.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function import tensorflow as tf import tensorflow_transform as tft def _get_raw_categorical_column_name(column_idx): return 'categorical-feature-{}'.format(column_idx) def get_transformed_categorical_column_name(column_name_or_id): if isinstance(column_name_or_id, bytes): # assume the input is column name column_name = column_name_or_id else: # assume the input is column id column_name = _get_raw_categorical_column_name(column_name_or_id) return column_name + '_id' _INTEGER_COLUMN_NAMES = [ 'int-feature-{}'.format(column_idx) for column_idx in range(1, 14) ] _CATEGORICAL_COLUMN_NAMES = [ _get_raw_categorical_column_name(column_idx) for column_idx in range(14, 40) ] DEFAULT_DELIMITER = '\t' # Number of buckets for integer columns. _NUM_BUCKETS = 10 # Schema annotations aren't supported in this build. tft.common.IS_ANNOTATIONS_PB_AVAILABLE = False def make_ordered_column_names(include_label=True): """Returns the column names in the dataset in the order as they appear. Args: include_label: Indicates whether the label feature should be included. Returns: A list of column names in the dataset. """ result = ['clicked'] if include_label else [] for name in _INTEGER_COLUMN_NAMES: result.append(name) for name in _CATEGORICAL_COLUMN_NAMES: result.append(name) return result def make_legacy_input_feature_spec(include_label=True): """Input schema definition. Args: include_label: Indicates whether the label feature should be included. Returns: A `Schema` object. """ result = {} if include_label: result['clicked'] = tf.io.FixedLenFeature(shape=[], dtype=tf.int64) for name in _INTEGER_COLUMN_NAMES: result[name] = tf.io.FixedLenFeature( shape=[], dtype=tf.int64, default_value=-1) for name in _CATEGORICAL_COLUMN_NAMES: result[name] = tf.io.FixedLenFeature( shape=[], dtype=tf.string, default_value='') return result def make_input_feature_spec(include_label=True): """Input schema definition. Args: include_label: Indicates whether the label feature should be included. Returns: A `Schema` object. """ result = {} if include_label: result['clicked'] = tf.io.FixedLenFeature(shape=[], dtype=tf.int64) for name in _INTEGER_COLUMN_NAMES: result[name] = tf.io.VarLenFeature(dtype=tf.int64) for name in _CATEGORICAL_COLUMN_NAMES: result[name] = tf.io.VarLenFeature(dtype=tf.string) return result def make_preprocessing_fn(frequency_threshold): """Creates a preprocessing function for criteo. Args: frequency_threshold: The frequency_threshold used when generating vocabularies for the categorical features. Returns: A preprocessing function. """ def preprocessing_fn(inputs): """User defined preprocessing function for criteo columns. Args: inputs: dictionary of input `tensorflow_transform.Column`. Returns: A dictionary of `tensorflow_transform.Column` representing the transformed columns. """ result = {'clicked': inputs['clicked']} for name in _INTEGER_COLUMN_NAMES: feature = inputs[name] # TODO(https://github.com/apache/beam/issues/24902): # Replace this boilerplate with a helper function. # This is a SparseTensor because it is optional. Here we fill in a # default value when it is missing. feature = tft.sparse_tensor_to_dense_with_shape( feature, [None, 1], default_value=-1) # Reshaping from a batch of vectors of size 1 to a batch of scalars and # adding a bucketized version. feature = tf.squeeze(feature, axis=1) result[name] = feature result[name + '_bucketized'] = tft.bucketize(feature, _NUM_BUCKETS) for name in _CATEGORICAL_COLUMN_NAMES: feature = inputs[name] # Similar to for integer columns, but use '' as default. feature = tft.sparse_tensor_to_dense_with_shape( feature, [None, 1], default_value='') feature = tf.squeeze(feature, axis=1) result[get_transformed_categorical_column_name( name)] = tft.compute_and_apply_vocabulary( feature, frequency_threshold=frequency_threshold) return result return preprocessing_fn ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.706442 apache_beam-2.74.0/apache_beam/testing/benchmarks/cloudml/pipelines/0000755000175100017510000000000015206113160025100 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/cloudml/pipelines/__init__.py0000644000175100017510000000142015206113032027204 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/cloudml/pipelines/workflow.py0000644000175100017510000001735115206113032027331 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import argparse import logging import os import tensorflow_transform as tft import tensorflow_transform.beam as tft_beam from tensorflow_transform import coders from tensorflow_transform.tf_metadata import dataset_metadata from tensorflow_transform.tf_metadata import schema_utils from tfx_bsl.public import tfxio import apache_beam as beam from apache_beam.testing.benchmarks.cloudml.criteo_tft import criteo # Name of the column for the synthetic version of the benchmark. _SYNTHETIC_COLUMN = 'x' class _RecordBatchToPyDict(beam.PTransform): """Converts PCollections of pa.RecordBatch to python dicts.""" def __init__(self, input_feature_spec): self._input_feature_spec = input_feature_spec def expand(self, pcoll): def format_values(instance): return { k: v.squeeze(0).tolist() if v is not None else self._input_feature_spec[k].default_value for k, v in instance.items() } return ( pcoll | 'RecordBatchToDicts' >> beam.FlatMap(lambda x: x.to_pandas().to_dict(orient='records')) | 'FormatPyDictValues' >> beam.Map(format_values)) def _synthetic_preprocessing_fn(inputs): return { _SYNTHETIC_COLUMN: tft.compute_and_apply_vocabulary( inputs[_SYNTHETIC_COLUMN], # Execute more codepaths but do no frequency filtration. frequency_threshold=1, # Execute more codepaths but do no top filtration. top_k=2**31 - 1, # Execute more codepaths num_oov_buckets=10) } class _PredictionHistogramFn(beam.DoFn): def __init__(self): # Beam Metrics API for Distributions only works with integers but # predictions are floating point numbers. We thus store a "quantized" # distribution of the prediction with sufficient granularity and for ease # of human interpretation (eg as a percentage for logistic regression). self._prediction_distribution = beam.metrics.Metrics.distribution( self.__class__, 'int(scores[0]*100)') def process(self, element): self._prediction_distribution.update(int(element['scores'][0] * 100)) def setup_pipeline(p, args): if args.classifier == 'criteo': input_feature_spec = criteo.make_input_feature_spec() input_schema = schema_utils.schema_from_feature_spec(input_feature_spec) input_tfxio = tfxio.BeamRecordCsvTFXIO( physical_format='text', column_names=criteo.make_ordered_column_names(), schema=input_schema, delimiter=criteo.DEFAULT_DELIMITER, telemetry_descriptors=['CriteoCloudMLBenchmark']) preprocessing_fn = criteo.make_preprocessing_fn(args.frequency_threshold) else: assert False, 'Unknown args classifier <{}>'.format(args.classifier) input_data = p | 'ReadFromText' >> beam.io.textio.ReadFromText( args.input, coder=beam.coders.BytesCoder()) if args.benchmark_type == 'tft': logging.info('TFT benchmark') # Setting TFXIO output format only for Criteo benchmarks to make sure that # both codepaths are covered. output_record_batches = args.classifier == 'criteo' # pylint: disable=expression-not-assigned input_metadata = dataset_metadata.DatasetMetadata(schema=input_schema) ( input_metadata | 'WriteInputMetadata' >> tft_beam.WriteMetadata( os.path.join(args.output, 'raw_metadata'), pipeline=p)) with tft_beam.Context(temp_dir=os.path.join(args.output, 'tmp'), use_deep_copy_optimization=True): decoded_input_data = ( input_data | 'DecodeForAnalyze' >> input_tfxio.BeamSource()) transform_fn = ((decoded_input_data, input_tfxio.TensorAdapterConfig()) | 'Analyze' >> tft_beam.AnalyzeDataset(preprocessing_fn)) if args.shuffle: # Shuffle the data before any decoding (more compact representation). input_data |= 'Shuffle' >> beam.transforms.Reshuffle() # pylint: disable=no-value-for-parameter decoded_input_data = ( input_data | 'DecodeForTransform' >> input_tfxio.BeamSource()) (dataset, metadata) = ((decoded_input_data, input_tfxio.TensorAdapterConfig()), transform_fn) | 'Transform' >> tft_beam.TransformDataset( output_record_batches=output_record_batches) if output_record_batches: def record_batch_to_examples(batch, unary_passthrough_features): """Encodes transformed data as tf.Examples.""" # Ignore unary pass-through features. del unary_passthrough_features # From beam: "imports, functions and other variables defined in the # global context of your __main__ file of your Dataflow pipeline are, by # default, not available in the worker execution environment, and such # references will cause a NameError, unless the --save_main_session # pipeline option is set to True. Please see # https://cloud.google.com/dataflow/faq#how-do-i-handle-nameerrors ." from tfx_bsl.coders.example_coder import RecordBatchToExamples return RecordBatchToExamples(batch) encode_ptransform = beam.FlatMapTuple(record_batch_to_examples) else: example_coder = coders.ExampleProtoCoder(metadata.schema) encode_ptransform = beam.Map(example_coder.encode) # TODO: Use WriteDataset instead when it becomes available. ( dataset | 'Encode' >> encode_ptransform | 'Write' >> beam.io.WriteToTFRecord( os.path.join(args.output, 'features_train'), file_name_suffix='.tfrecord.gz')) # transform_fn | beam.Map(print) transform_fn | 'WriteTransformFn' >> tft_beam.WriteTransformFn(args.output) # TODO: Remember to eventually also save the statistics. else: logging.fatal('Unknown benchmark type: %s', args.benchmark_type) def parse_known_args(argv): """Parses args for this workflow.""" parser = argparse.ArgumentParser() parser.add_argument( '--input', dest='input', required=True, help='Input path for input files.') parser.add_argument( '--output', dest='output', required=True, help='Output path for output files.') parser.add_argument( '--classifier', dest='classifier', required=True, help='Name of classifier to use.') parser.add_argument( '--frequency_threshold', dest='frequency_threshold', default=5, # TODO: Align default with TFT (ie 0). help='Threshold for minimum number of unique values for a category.') parser.add_argument( '--shuffle', action='store_false', dest='shuffle', default=True, help='Skips shuffling the data.') parser.add_argument( '--benchmark_type', dest='benchmark_type', required=True, help='Type of benchmark to run.') return parser.parse_known_args(argv) def run(argv=None): """Main entry point; defines and runs the pipeline.""" known_args, pipeline_args = parse_known_args(argv) with beam.Pipeline(argv=pipeline_args) as p: setup_pipeline(p, known_args) if __name__ == '__main__': run() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.7084424 apache_beam-2.74.0/apache_beam/testing/benchmarks/inference/0000755000175100017510000000000015206113160023407 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/inference/__init__.py0000644000175100017510000000142015206113032025513 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000020600000000000010213 xustar00112 path=apache_beam-2.74.0/apache_beam/testing/benchmarks/inference/pytorch_image_classification_benchmarks.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/inference/pytorch_image_classification_benchmarks.0000644000175100017510000000532515206113032033515 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging from torchvision import models from apache_beam.examples.inference import pytorch_image_classification from apache_beam.testing.load_tests.dataflow_cost_benchmark import DataflowCostBenchmark _PERF_TEST_MODELS = ['resnet50', 'resnet101', 'resnet152'] _PRETRAINED_MODEL_MODULE = 'torchvision.models' class PytorchVisionBenchmarkTest(DataflowCostBenchmark): def __init__(self): # TODO (https://github.com/apache/beam/issues/23008) # make get_namespace() method in RunInference static self.metrics_namespace = 'BeamML_PyTorch' super().__init__( metrics_namespace=self.metrics_namespace, pcollection='PyTorchRunInference/BeamML_RunInference_Postprocess-0.out0' ) def test(self): pretrained_model_name = self.pipeline.get_option('pretrained_model_name') if not pretrained_model_name: raise RuntimeError( 'Please provide a pretrained torch model name.' ' Model name must be from the module torchvision.models') if pretrained_model_name == _PERF_TEST_MODELS[0]: model_class = models.resnet50 elif pretrained_model_name == _PERF_TEST_MODELS[1]: model_class = models.resnet101 elif pretrained_model_name == _PERF_TEST_MODELS[2]: model_class = models.resnet152 else: raise NotImplementedError # model_params are same for all the models. But this may change if we add # different models. model_params = {'num_classes': 1000, 'pretrained': False} extra_opts = {} extra_opts['input'] = self.pipeline.get_option('input_file') device = self.pipeline.get_option('device') self.result = pytorch_image_classification.run( self.pipeline.get_full_options_as_args(**extra_opts), model_class=model_class, model_params=model_params, test_pipeline=self.pipeline, device=device) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) PytorchVisionBenchmarkTest().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/inference/pytorch_language_modeling_benchmarks.py0000644000175100017510000000320115206113032033361 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging from apache_beam.examples.inference import pytorch_language_modeling from apache_beam.testing.load_tests.dataflow_cost_benchmark import DataflowCostBenchmark class PytorchLanguageModelingBenchmarkTest(DataflowCostBenchmark): def __init__(self): # TODO (https://github.com/apache/beam/issues/23008): # make get_namespace() method in RunInference static self.metrics_namespace = 'BeamML_PyTorch' super().__init__(metrics_namespace=self.metrics_namespace) def test(self): extra_opts = {} extra_opts['input'] = self.pipeline.get_option('input_file') self.result = pytorch_language_modeling.run( self.pipeline.get_full_options_as_args(**extra_opts), test_pipeline=self.pipeline) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) PytorchLanguageModelingBenchmarkTest().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/inference/pytorch_sentiment_benchmarks.py0000644000175100017510000000304015206113032031727 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging from apache_beam.examples.inference import pytorch_sentiment from apache_beam.testing.load_tests.dataflow_cost_benchmark import DataflowCostBenchmark class PytorchSentimentBenchmarkTest(DataflowCostBenchmark): def __init__(self): self.metrics_namespace = 'BeamML_PyTorch' super().__init__( metrics_namespace=self.metrics_namespace, pcollection='PostProcess.out0') def test(self): extra_opts = {} extra_opts['input'] = self.pipeline.get_option('input_file') self.result = pytorch_sentiment.run( self.pipeline.get_full_options_as_args(**extra_opts), test_pipeline=self.pipeline) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) PytorchSentimentBenchmarkTest().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/inference/table_row_inference_benchmark.py0000644000175100017510000001014215206113032031763 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Benchmark test for table row inference pipeline. This benchmark measures the performance of RunInference with continuous table row inputs, including throughput, latency, and cost metrics. """ import logging from apache_beam.examples.inference import table_row_inference from apache_beam.options.pipeline_options import DebugOptions from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import WorkerOptions from apache_beam.testing.load_tests.dataflow_cost_benchmark import DataflowCostBenchmark from apache_beam.testing.load_tests.load_test import LoadTestOptions class TableRowInferenceOptions( LoadTestOptions, StandardOptions, GoogleCloudOptions, WorkerOptions, DebugOptions, SetupOptions, ): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--mode', default='batch') parser.add_argument('--input_subscription') parser.add_argument('--input_file') parser.add_argument('--output_table') parser.add_argument('--model_path') parser.add_argument('--feature_columns') parser.add_argument('--window_size_sec', type=int, default=60) parser.add_argument('--trigger_interval_sec', type=int, default=30) parser.add_argument('--input_expand_factor', type=int, default=1) class TableRowInferenceBenchmarkTest(DataflowCostBenchmark): """Benchmark for continuous table row inference with RunInference. This benchmark measures: - Mean Inference Batch Size: Average batch size for inference - Mean Inference Batch Latency: Average time per batch inference - Mean Load Model Latency: Time to load the model - Throughput: Elements processed per second - Cost: Estimated cost on Dataflow """ options_class = TableRowInferenceOptions def __init__(self): self.metrics_namespace = 'BeamML_TableInference' super().__init__( metrics_namespace=self.metrics_namespace, is_streaming=False, pcollection='RunInference/BeamML_RunInference_Postprocess-0.out0') self.is_streaming = ((self.pipeline.get_option('mode') or 'batch') == 'streaming') if self.is_streaming: self.subscription = self.pipeline.get_option('input_subscription') def test(self): """Execute the table row inference pipeline for benchmarking.""" extra_opts = {} mode = self.pipeline.get_option('mode') or 'batch' extra_opts['mode'] = mode if mode == 'streaming': extra_opts['input_subscription'] = self.pipeline.get_option( 'input_subscription') extra_opts['window_size_sec'] = int( self.pipeline.get_option('window_size_sec') or 60) extra_opts['trigger_interval_sec'] = int( self.pipeline.get_option('trigger_interval_sec') or 30) else: extra_opts['input_file'] = self.pipeline.get_option('input_file') for opt in ['output_table', 'model_path', 'feature_columns']: val = self.pipeline.get_option(opt) if val: extra_opts[opt] = val self.result = table_row_inference.run( self.pipeline.get_full_options_as_args(**extra_opts), test_pipeline=self.pipeline) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) TableRowInferenceBenchmarkTest().run() ././@PaxHeader0000000000000000000000000000021500000000000010213 xustar00119 path=apache_beam-2.74.0/apache_beam/testing/benchmarks/inference/tensorflow_mnist_classification_cost_benchmark.py 22 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/inference/tensorflow_mnist_classification_cost_ben0000644000175100017510000000316415206113032033677 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging from apache_beam.examples.inference import tensorflow_mnist_classification from apache_beam.testing.load_tests.dataflow_cost_benchmark import DataflowCostBenchmark class TensorflowMNISTClassificationCostBenchmark(DataflowCostBenchmark): def __init__(self): super().__init__(pcollection='PostProcessOutputs.out0') def test(self): extra_opts = {} extra_opts['input'] = self.pipeline.get_option('input_file') extra_opts['output'] = self.pipeline.get_option('output_file') extra_opts['model_path'] = self.pipeline.get_option('model') self.result = tensorflow_mnist_classification.run( self.pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) TensorflowMNISTClassificationCostBenchmark().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/inference/vllm_gemma_benchmarks.py0000644000175100017510000000311115206113032030270 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging from apache_beam.examples.inference import vllm_gemma_batch from apache_beam.testing.load_tests.dataflow_cost_benchmark import DataflowCostBenchmark class VllmGemmaBenchmarkTest(DataflowCostBenchmark): def __init__(self): self.metrics_namespace = "BeamML_vLLM" super().__init__( metrics_namespace=self.metrics_namespace, pcollection="FormatForBQ.out0", ) def test(self): # The perf-test framework passes --input_file, # but the pipeline expects --input. extra_opts = {"input": self.pipeline.get_option("input_file")} self.result = vllm_gemma_batch.run( self.pipeline.get_full_options_as_args(**extra_opts), test_pipeline=self.pipeline) if __name__ == "__main__": logging.basicConfig(level=logging.INFO) VllmGemmaBenchmarkTest().run() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.7104425 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/0000755000175100017510000000000015206113160023116 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/__init__.py0000644000175100017510000000142015206113032025222 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.7114427 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/models/0000755000175100017510000000000015206113160024401 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/models/__init__.py0000644000175100017510000000142015206113032026505 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/models/auction_bid.py0000644000175100017510000000365315206113032027240 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Result of WinningBid transform.""" from apache_beam.coders import coder_impl from apache_beam.coders.coders import FastCoder from apache_beam.testing.benchmarks.nexmark.models import nexmark_model class AuctionBidCoder(FastCoder): def to_type_hint(self): return AuctionBid def _create_impl(self): return AuctionBidCoderImpl() def is_deterministic(self): return True class AuctionBid(object): CODER = AuctionBidCoder() def __init__(self, auction, bid): self.auction = auction self.bid = bid def __repr__(self): return nexmark_model.model_to_json(self) class AuctionBidCoderImpl(coder_impl.StreamCoderImpl): _auction_coder_impl = nexmark_model.AuctionCoderImpl() _bid_coder_Impl = nexmark_model.BidCoderImpl() def encode_to_stream(self, value, stream, nested): self._auction_coder_impl.encode_to_stream(value.auction, stream, True) self._bid_coder_Impl.encode_to_stream(value.bid, stream, True) def decode_from_stream(self, stream, nested): auction = self._auction_coder_impl.decode_from_stream(stream, True) bid = self._bid_coder_Impl.decode_from_stream(stream, True) return AuctionBid(auction, bid) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/models/field_name.py0000644000175100017510000000237115206113032027037 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Field names for de-serializing json representation of Models """ class FieldNames: ID = 'id' NAME = 'name' EMAIL_ADDRESS = 'emailAddress' CREDIT_CARD = 'creditCard' CITY = "city" STATE = "state" DATE_TIME = "dateTime" EXTRA = "extra" ITEM_NAME = "itemName" DESCRIPTION = "description" INITIAL_BID = "initialBid" RESERVE = "reserve" EXPIRES = "expires" SELLER = "seller" CATEGORY = "category" AUCTION = "auction" BIDDER = "bidder" PRICE = "price" ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/models/nexmark_model.py0000644000175100017510000001762015206113032027604 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Nexmark model. The nexmark suite is a series of queries (streaming pipelines) performed on a simulation of auction events. The model includes the three roles that generate events: - The person who starts and auction or makes a bid (Person). - The auction item (Auction). - The bid on an item for auction (Bid). """ import json from apache_beam.coders import coder_impl from apache_beam.coders.coders import FastCoder from apache_beam.coders.coders import StrUtf8Coder from apache_beam.utils.timestamp import Timestamp def model_to_json(model): return json.dumps(construct_json_dict(model), separators=(",", ":")) def construct_json_dict(model): return {k: unnest_to_json(v) for k, v in model.__dict__.items()} def unnest_to_json(cand): if isinstance(cand, Timestamp): return cand.micros // 1000 elif isinstance(cand, (Auction, Bid, Person)): return construct_json_dict(cand) else: return cand class PersonCoder(FastCoder): def to_type_hint(self): return Person def _create_impl(self): return PersonCoderImpl() def is_deterministic(self): return True class Person(object): "Author of an auction or a bid." CODER = PersonCoder() def __init__( self, id, name, email, credit_card, city, state, date_time, extra=None): self.id = id self.name = name self.email_address = email # key self.credit_card = credit_card self.city = city self.state = state self.date_time = date_time self.extra = extra def __repr__(self): return model_to_json(self) class AuctionCoder(FastCoder): def to_type_hint(self): return Auction def _create_impl(self): return AuctionCoderImpl() def is_deterministic(self): return True class Auction(object): "Item for auction." CODER = AuctionCoder() def __init__( self, id, item_name, description, initial_bid, reserve_price, date_time, expires, seller, category, extra=None): self.id = id self.item_name = item_name # key self.description = description self.initial_bid = initial_bid self.reserve = reserve_price self.date_time = date_time self.expires = expires self.seller = seller self.category = category self.extra = extra def __repr__(self): return model_to_json(self) class BidCoder(FastCoder): def to_type_hint(self): return Bid def _create_impl(self): return BidCoderImpl() def is_deterministic(self): return True class Bid(object): "A bid for an item for auction." CODER = BidCoder() def __init__(self, auction, bidder, price, date_time, extra=None): self.auction = auction # key self.bidder = bidder self.price = price self.date_time = date_time self.extra = extra def __repr__(self): return model_to_json(self) class AuctionCoderImpl(coder_impl.StreamCoderImpl): _int_coder_impl = coder_impl.VarIntCoderImpl() _str_coder_impl = StrUtf8Coder().get_impl() _time_coder_impl = coder_impl.TimestampCoderImpl() def encode_to_stream(self, value, stream, nested): self._int_coder_impl.encode_to_stream(value.id, stream, True) self._str_coder_impl.encode_to_stream(value.item_name, stream, True) self._str_coder_impl.encode_to_stream(value.description, stream, True) self._int_coder_impl.encode_to_stream(value.initial_bid, stream, True) self._int_coder_impl.encode_to_stream(value.reserve, stream, True) self._time_coder_impl.encode_to_stream(value.date_time, stream, True) self._time_coder_impl.encode_to_stream(value.expires, stream, True) self._int_coder_impl.encode_to_stream(value.seller, stream, True) self._int_coder_impl.encode_to_stream(value.category, stream, True) self._str_coder_impl.encode_to_stream(value.extra, stream, True) def decode_from_stream(self, stream, nested): id = self._int_coder_impl.decode_from_stream(stream, True) item_name = self._str_coder_impl.decode_from_stream(stream, True) description = self._str_coder_impl.decode_from_stream(stream, True) initial_bid = self._int_coder_impl.decode_from_stream(stream, True) reserve = self._int_coder_impl.decode_from_stream(stream, True) date_time = self._time_coder_impl.decode_from_stream(stream, True) expires = self._time_coder_impl.decode_from_stream(stream, True) seller = self._int_coder_impl.decode_from_stream(stream, True) category = self._int_coder_impl.decode_from_stream(stream, True) extra = self._str_coder_impl.decode_from_stream(stream, True) return Auction( id, item_name, description, initial_bid, reserve, date_time, expires, seller, category, extra) class BidCoderImpl(coder_impl.StreamCoderImpl): _int_coder_impl = coder_impl.VarIntCoderImpl() _str_coder_impl = StrUtf8Coder().get_impl() _time_coder_impl = coder_impl.TimestampCoderImpl() def encode_to_stream(self, value, stream, nested): self._int_coder_impl.encode_to_stream(value.auction, stream, True) self._int_coder_impl.encode_to_stream(value.bidder, stream, True) self._int_coder_impl.encode_to_stream(value.price, stream, True) self._time_coder_impl.encode_to_stream(value.date_time, stream, True) self._str_coder_impl.encode_to_stream(value.extra, stream, True) def decode_from_stream(self, stream, nested): auction = self._int_coder_impl.decode_from_stream(stream, True) bidder = self._int_coder_impl.decode_from_stream(stream, True) price = self._int_coder_impl.decode_from_stream(stream, True) date_time = self._time_coder_impl.decode_from_stream(stream, True) extra = self._str_coder_impl.decode_from_stream(stream, True) return Bid(auction, bidder, price, date_time, extra) class PersonCoderImpl(coder_impl.StreamCoderImpl): _int_coder_impl = coder_impl.VarIntCoderImpl() _str_coder_impl = StrUtf8Coder().get_impl() _time_coder_impl = coder_impl.TimestampCoderImpl() def encode_to_stream(self, value, stream, nested): self._int_coder_impl.encode_to_stream(value.id, stream, True) self._str_coder_impl.encode_to_stream(value.name, stream, True) self._str_coder_impl.encode_to_stream(value.email_address, stream, True) self._str_coder_impl.encode_to_stream(value.credit_card, stream, True) self._str_coder_impl.encode_to_stream(value.city, stream, True) self._str_coder_impl.encode_to_stream(value.state, stream, True) self._time_coder_impl.encode_to_stream(value.date_time, stream, True) self._str_coder_impl.encode_to_stream(value.extra, stream, True) def decode_from_stream(self, stream, nested): id = self._int_coder_impl.decode_from_stream(stream, True) name = self._str_coder_impl.decode_from_stream(stream, True) email = self._str_coder_impl.decode_from_stream(stream, True) credit_card = self._str_coder_impl.decode_from_stream(stream, True) city = self._str_coder_impl.decode_from_stream(stream, True) state = self._str_coder_impl.decode_from_stream(stream, True) date_time = self._time_coder_impl.decode_from_stream(stream, True) extra = self._str_coder_impl.decode_from_stream(stream, True) return Person(id, name, email, credit_card, city, state, date_time, extra) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/monitor.py0000644000175100017510000000426115206113032025160 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from time import time import apache_beam as beam from apache_beam.metrics import Metrics class Monitor(object): """ A monitor of elements with support for later retrieving their metrics monitor objects contains a doFn to record metrics Args: namespace: the namespace all metrics within this Monitor uses name_prefix: a prefix for this Monitor's metrics' names, intended to be unique in per-monitor basis in pipeline """ def __init__(self, namespace: str, name_prefix: str) -> None: self.namespace = namespace self.name_prefix = name_prefix self.doFn = MonitorDoFn(namespace, name_prefix) class MonitorDoFn(beam.DoFn): def __init__(self, namespace, prefix): self.element_count = Metrics.counter( namespace, prefix + MonitorSuffix.ELEMENT_COUNTER) self.event_time = Metrics.distribution( namespace, prefix + MonitorSuffix.EVENT_TIME) self.event_timestamp = Metrics.distribution( namespace, prefix + MonitorSuffix.EVENT_TIMESTAMP) def start_bundle(self): self.event_time.update(int(time() * 1000)) def process(self, element, timestamp=beam.DoFn.TimestampParam): self.element_count.inc() self.event_timestamp.update(timestamp) yield element def finish_bundle(self): self.event_time.update(int(time() * 1000)) class MonitorSuffix: ELEMENT_COUNTER = '.elements' EVENT_TIMESTAMP = '.event_timestamp' EVENT_TIME = '.event_time' ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/nexmark_launcher.py0000644000175100017510000004362715206113032027030 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Nexmark launcher. The Nexmark suite is a series of queries (streaming pipelines) performed on a simulation of auction events. The launcher orchestrates the generation and parsing of streaming events and the running of queries. Model - Person: Author of an auction or a bid. - Auction: Item under auction. - Bid: A bid for an item under auction. Events - Create Person - Create Auction - Create Bid Queries - Query0: Pass through (send and receive auction events). Usage - DirectRunner python nexmark_launcher.py \ --query/q \ --project \ --loglevel=DEBUG (optional) \ --wait_until_finish_duration \ --streaming - DataflowRunner python nexmark_launcher.py \ --query/q \ --project \ --region \ --loglevel=DEBUG (optional) \ --wait_until_finish_duration \ --streaming \ --sdk_location \ --staging_location=gs://... \ --temp_location=gs:// """ # pytype: skip-file import argparse import json import logging import os import time import uuid import requests from requests.auth import HTTPBasicAuth import apache_beam as beam from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import TypeOptions from apache_beam.runners import PipelineState from apache_beam.testing.benchmarks.nexmark import nexmark_util from apache_beam.testing.benchmarks.nexmark.monitor import Monitor from apache_beam.testing.benchmarks.nexmark.monitor import MonitorSuffix from apache_beam.testing.benchmarks.nexmark.nexmark_perf import NexmarkPerf from apache_beam.testing.benchmarks.nexmark.queries import query0 from apache_beam.testing.benchmarks.nexmark.queries import query1 from apache_beam.testing.benchmarks.nexmark.queries import query2 from apache_beam.testing.benchmarks.nexmark.queries import query3 from apache_beam.testing.benchmarks.nexmark.queries import query4 from apache_beam.testing.benchmarks.nexmark.queries import query5 from apache_beam.testing.benchmarks.nexmark.queries import query6 from apache_beam.testing.benchmarks.nexmark.queries import query7 from apache_beam.testing.benchmarks.nexmark.queries import query8 from apache_beam.testing.benchmarks.nexmark.queries import query9 from apache_beam.testing.benchmarks.nexmark.queries import query10 from apache_beam.testing.benchmarks.nexmark.queries import query11 from apache_beam.testing.benchmarks.nexmark.queries import query12 from apache_beam.transforms import window class NexmarkLauncher(object): # how long after some result is seen and no activity seen do we cancel job DONE_DELAY = 5 * 60 # delay in seconds between sample perf data PERF_DELAY = 20 # delay before cancelling the job when pipeline appears to be stuck TERMINATE_DELAY = 1 * 60 * 60 # delay before warning when pipeline appears to be stuck WARNING_DELAY = 10 * 60 def __init__(self): self.parse_args() self.manage_resources = self.args.manage_resources self.uuid = str(uuid.uuid4()) if self.manage_resources else '' self.topic_name = ( self.args.topic_name + self.uuid if self.args.topic_name else None) self.subscription_name = ( self.args.subscription_name + self.uuid if self.args.subscription_name else None) self.pubsub_mode = self.args.pubsub_mode if self.manage_resources: from google.cloud import pubsub self.cleanup() publish_client = pubsub.Client(project=self.project) topic = publish_client.topic(self.topic_name) logging.info('creating topic %s', self.topic_name) topic.create() sub = topic.subscription(self.subscription_name) logging.info('creating sub %s', self.topic_name) sub.create() self.export_influxdb = self.args.export_summary_to_influx_db if self.export_influxdb: self.influx_database = self.args.influx_database self.influx_host = self.args.influx_host self.influx_base = self.args.base_influx_measurement self.influx_retention = self.args.influx_retention_policy def parse_args(self): parser = argparse.ArgumentParser() parser.add_argument( '--query', '-q', type=int, action='append', required=True, choices=[i for i in range(13)], help='Query to run') parser.add_argument( '--subscription_name', type=str, help='Pub/Sub subscription to read from') parser.add_argument( '--topic_name', type=str, help='Pub/Sub topic to read from') parser.add_argument( '--loglevel', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], default='INFO', help='Set logging level to debug') parser.add_argument( '--input', type=str, help='Path to the data file containing nexmark events.') parser.add_argument( '--num_events', type=int, default=100000, help='number of events expected to process') parser.add_argument( '--manage_resources', default=False, action='store_true', help='If true, manage the creation and cleanup of topics and ' 'subscriptions.') parser.add_argument( '--pubsub_mode', type=str, default='SUBSCRIBE_ONLY', choices=['PUBLISH_ONLY', 'SUBSCRIBE_ONLY', 'COMBINED'], help='Pubsub mode used in the pipeline.') parser.add_argument( '--export_summary_to_influx_db', default=False, action='store_true', help='If set store results in influxdb') parser.add_argument( '--influx_database', type=str, default='beam_test_metrics', help='Influx database name') parser.add_argument( '--influx_host', type=str, default='http://localhost:8086', help='Influx database url') parser.add_argument( '--base_influx_measurement', type=str, default='nexmark', help='Prefix to influx measurement') parser.add_argument( '--influx_retention_policy', type=str, default='forever', help='Retention policy for stored results') self.args, self.pipeline_args = parser.parse_known_args() logging.basicConfig( level=getattr(logging, self.args.loglevel, None), format='(%(threadName)-10s) %(message)s') self.pipeline_options = PipelineOptions(self.pipeline_args) logging.debug('args, pipeline_args: %s, %s', self.args, self.pipeline_args) # Usage with Dataflow requires a project to be supplied. self.project = self.pipeline_options.view_as(GoogleCloudOptions).project self.streaming = self.pipeline_options.view_as(StandardOptions).streaming self.pipeline_options.view_as(TypeOptions).allow_unsafe_triggers = True if self.streaming: if self.args.subscription_name is None or self.project is None: raise ValueError( 'argument --subscription_name and --project ' + 'are required when running in streaming mode') else: if self.args.input is None: raise ValueError( 'argument --input is required when running in batch mode') # We use the save_main_session option because one or more DoFn's in this # workflow rely on global context (e.g., a module imported at module level). self.pipeline_options.view_as(SetupOptions).save_main_session = True def generate_events(self): from google.cloud import pubsub publish_client = pubsub.Client(project=self.project) topic = publish_client.topic(self.topic_name) logging.info('Generating auction events to topic %s', topic.name) if self.args.input.startswith('gs://'): from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem fs = GCSFileSystem(self.pipeline_options) with fs.open(self.args.input) as infile: for line in infile: topic.publish(line) else: with open(self.args.input) as infile: for line in infile: topic.publish(line) logging.info('Finished event generation.') def read_from_file(self): return ( self.pipeline | 'reading_from_file' >> beam.io.ReadFromText(self.args.input) | 'deserialization' >> beam.ParDo(nexmark_util.ParseJsonEventFn()) | 'timestamping' >> beam.Map(lambda e: window.TimestampedValue(e, e.date_time))) def read_from_pubsub(self): # Read from PubSub into a PCollection. if self.subscription_name: raw_events = self.pipeline | 'ReadPubSub_sub' >> beam.io.ReadFromPubSub( subscription=self.subscription_name, with_attributes=True, timestamp_attribute='timestamp') else: raw_events = self.pipeline | 'ReadPubSub_topic' >> beam.io.ReadFromPubSub( topic=self.topic_name, with_attributes=True, timestamp_attribute='timestamp') events = ( raw_events | 'pubsub_unwrap' >> beam.Map(lambda m: m.data) | 'deserialization' >> beam.ParDo(nexmark_util.ParseJsonEventFn())) return events def run_query( self, query_num, query, query_args, pipeline_options, query_errors): try: self.pipeline = beam.Pipeline(options=self.pipeline_options) nexmark_util.setup_coder() event_monitor = Monitor('.events', 'event') result_monitor = Monitor('.results', 'result') if self.streaming: if self.pubsub_mode != 'SUBSCRIBE_ONLY': self.generate_events() if self.pubsub_mode == 'PUBLISH_ONLY': return events = self.read_from_pubsub() else: events = self.read_from_file() events = events | 'event_monitor' >> beam.ParDo(event_monitor.doFn) output = query.load(events, query_args, pipeline_options) output | 'result_monitor' >> beam.ParDo(result_monitor.doFn) # pylint: disable=expression-not-assigned result = self.pipeline.run() if not self.streaming: result.wait_until_finish() perf = self.monitor(result, event_monitor, result_monitor) self.log_performance(perf) if self.export_influxdb: self.publish_performance_influxdb(query_num, perf) except Exception as exc: query_errors.append(str(exc)) raise def monitor(self, job, event_monitor, result_monitor): """ keep monitoring the performance and progress of running job and cancel the job if the job is stuck or seems to have finished running Returns: the final performance if it is measured """ logging.info('starting to monitor the job') last_active_ms = -1 perf = None cancel_job = False waiting_for_shutdown = False while True: now = int(time.time() * 1000) # current time in ms logging.debug('now is %d', now) curr_perf = NexmarkLauncher.get_performance( job, event_monitor, result_monitor) if perf is None or curr_perf.has_progress(perf): last_active_ms = now # only judge if the job should be cancelled if it is streaming job and # has not been shut down already if self.streaming and not waiting_for_shutdown: quiet_duration = (now - last_active_ms) // 1000 if (curr_perf.event_count >= self.args.num_events and curr_perf.result_count >= 0 and quiet_duration > self.DONE_DELAY): # we think the job is finished if expected input count has been seen # and no new results have been produced for a while logging.info('streaming query appears to have finished executing') waiting_for_shutdown = True cancel_job = True elif quiet_duration > self.TERMINATE_DELAY: logging.error( 'streaming query have been stuck for %d seconds', quiet_duration) logging.error('canceling streaming job') waiting_for_shutdown = True cancel_job = True elif quiet_duration > self.WARNING_DELAY: logging.warning( 'streaming query have been stuck for %d seconds', quiet_duration) if cancel_job: job.cancel() perf = curr_perf stopped = PipelineState.is_terminal(job.state) if stopped: break if not waiting_for_shutdown: if last_active_ms == now: logging.info('activity seen, new performance data extracted') else: logging.info('no activity seen') else: logging.info('waiting for shutdown') time.sleep(self.PERF_DELAY) return perf @staticmethod def log_performance(perf: NexmarkPerf) -> None: logging.info( 'input event count: %d, output event count: %d' % (perf.event_count, perf.result_count)) logging.info( 'query run took %.1f seconds and processed %.1f events per second' % (perf.runtime_sec, perf.event_per_sec)) def publish_performance_influxdb(self, query_num, perf): processingMode = "streaming" if self.streaming else "batch" measurement = "%s_%d_python_%s" % ( self.influx_base, query_num, processingMode) tags = {'runner': self.pipeline_options.view_as(StandardOptions).runner} mt = ','.join([measurement] + [k + "=" + v for k, v in tags.items()]) fields = { 'numResults': "%di" % (perf.result_count), 'runtimeMs': "%di" % (perf.runtime_sec * 1000), } ts = int(time.time()) payload = '\n'.join( ["%s %s=%s %d" % (mt, k, v, ts) for k, v in fields.items()]) url = '%s/write' % (self.influx_host) query_str = { 'db': self.influx_database, 'rp': self.influx_retention, 'precision': 's', } user = os.getenv('INFLUXDB_USER') password = os.getenv('INFLUXDB_USER_PASSWORD') auth = HTTPBasicAuth(user, password) try: response = requests.post( url, params=query_str, data=payload, auth=auth, timeout=60) except requests.exceptions.RequestException as e: logging.warning('Failed to publish metrics to InfluxDB: ' + str(e)) else: if response.status_code != 204: content = json.loads(response.content) logging.warning( 'Failed to publish metrics to InfluxDB. Received status code %s ' 'with an error message: %s' % (response.status_code, content['error'])) @staticmethod def get_performance(result, event_monitor, result_monitor): event_count = nexmark_util.get_counter_metric( result, event_monitor.namespace, event_monitor.name_prefix + MonitorSuffix.ELEMENT_COUNTER) event_start = nexmark_util.get_start_time_metric( result, event_monitor.namespace, event_monitor.name_prefix + MonitorSuffix.EVENT_TIME) event_end = nexmark_util.get_end_time_metric( result, event_monitor.namespace, event_monitor.name_prefix + MonitorSuffix.EVENT_TIME) result_count = nexmark_util.get_counter_metric( result, result_monitor.namespace, result_monitor.name_prefix + MonitorSuffix.ELEMENT_COUNTER) result_end = nexmark_util.get_end_time_metric( result, result_monitor.namespace, result_monitor.name_prefix + MonitorSuffix.EVENT_TIME) perf = NexmarkPerf() perf.event_count = event_count perf.result_count = result_count effective_end = max(event_end, result_end) if effective_end >= 0 and event_start >= 0: perf.runtime_sec = (effective_end - event_start) / 1000 if event_count >= 0 and perf.runtime_sec > 0: perf.event_per_sec = event_count / perf.runtime_sec return perf def cleanup(self): if self.manage_resources: from google.cloud import pubsub publish_client = pubsub.Client(project=self.project) topic = publish_client.topic(self.topic_name) if topic.exists(): logging.info('deleting topic %s', self.topic_name) topic.delete() sub = topic.subscription(self.subscription_name) if sub.exists(): logging.info('deleting sub %s', self.topic_name) sub.delete() def run(self): queries = { 0: query0, 1: query1, 2: query2, 3: query3, 4: query4, 5: query5, 6: query6, 7: query7, 8: query8, 9: query9, 10: query10, 11: query11, 12: query12 } # TODO(mariagh): Move to a config file. query_args = { 'auction_skip': 123, 'window_size_sec': 10, 'window_period_sec': 5, 'fanout': 5, 'num_max_workers': 5, 'max_log_events': 100000, 'occasional_delay_sec': 3, 'max_auction_waiting_time': 600 } query_errors = [] for i in self.args.query: logging.info('Running query %d', i) self.run_query( i, queries[i], query_args, self.pipeline_options, query_errors=query_errors) if query_errors: logging.error('Query failed with %s', ', '.join(query_errors)) else: logging.info('Queries run: %s', self.args.query) if __name__ == '__main__': launcher = NexmarkLauncher() launcher.run() launcher.cleanup() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/nexmark_perf.py0000644000175100017510000000357015206113032026154 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ performance summary for a run of nexmark query """ class NexmarkPerf(object): def __init__( self, runtime_sec=None, event_count=None, event_per_sec=None, result_count=None): # the time difference between first seen input and the last seen result self.runtime_sec = runtime_sec if runtime_sec else -1.0 # the count of input event seen self.event_count = event_count if event_count else -1 # number of event processed per runtime_sec self.event_per_sec = event_per_sec if event_per_sec else -1.0 # number of result produced self.result_count = result_count if result_count else -1 def has_progress(self, previous_perf: 'NexmarkPerf') -> bool: """ Args: previous_perf: a NexmarkPerf object to be compared to self Returns: True if there are observed pipeline activity between self and other NexmarkPerf values """ if (self.runtime_sec != previous_perf.runtime_sec or self.event_count != previous_perf.event_count or self.result_count != previous_perf.result_count): return True return False ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/nexmark_util.py0000644000175100017510000002205415206113032026173 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utilities for the Nexmark suite. The Nexmark suite is a series of queries (streaming pipelines) performed on a simulation of auction events. This util includes: - A Command class used to terminate the streaming jobs launched in nexmark_launcher.py by the DirectRunner. - A ParseEventFn DoFn to parse events received from PubSub. Usage: To run a process for a certain duration, define in the code: command = Command(process_to_terminate, args) command.run(timeout=duration) """ # pytype: skip-file import json import logging import threading import apache_beam as beam from apache_beam.metrics import MetricsFilter from apache_beam.runners.runner import PipelineResult # pylint: disable=unused-import from apache_beam.testing.benchmarks.nexmark.models import auction_bid from apache_beam.testing.benchmarks.nexmark.models import nexmark_model from apache_beam.testing.benchmarks.nexmark.models.field_name import FieldNames from apache_beam.transforms import window from apache_beam.utils.timestamp import Timestamp _LOGGER = logging.getLogger(__name__) class Command(object): def __init__(self, cmd, args): self.cmd = cmd self.args = args def run(self, timeout): def thread_target(): logging.debug( 'Starting thread for %d seconds: %s', timeout, self.cmd.__name__) self.cmd(*self.args) _LOGGER.info( '%d seconds elapsed. Thread (%s) finished.', timeout, self.cmd.__name__) thread = threading.Thread(target=thread_target, name='Thread-timeout') thread.daemon = True thread.start() thread.join(timeout) def setup_coder(): beam.coders.registry.register_coder( nexmark_model.Auction, nexmark_model.AuctionCoder) beam.coders.registry.register_coder( nexmark_model.Person, nexmark_model.PersonCoder) beam.coders.registry.register_coder(nexmark_model.Bid, nexmark_model.BidCoder) beam.coders.registry.register_coder( auction_bid.AuctionBid, auction_bid.AuctionBidCoder) class ParseEventFn(beam.DoFn): """ Original parser for parsing raw events info into a Python objects. Each event line has the following format: person: ,name,email,credit_card,city, \ state,timestamp,extra auction: ,item_name, description,initial_bid, \ reserve_price,timestamp,expires,seller,category,extra bid: ,bidder,price,timestamp,extra For example: 'p12345,maria,maria@maria.com,1234-5678-9012-3456, \ sunnyvale,CA,1528098831536' 'a12345,car67,2012 hyundai elantra,15000,20000, \ 1528098831536,20180630,maria,vehicle' 'b12345,maria,20000,1528098831536' """ def process(self, elem): model_dict = { 'p': nexmark_model.Person, 'a': nexmark_model.Auction, 'b': nexmark_model.Bid, } row = elem.split(',') model = model_dict.get(elem[0]) if not model: raise ValueError('Invalid event: %s.' % row) event = model(*row) logging.debug('Parsed event: %s', event) yield event class ParseJsonEventFn(beam.DoFn): """Parses the raw event info into a Python objects. Each event line has the following format: person: {id,name,email,credit_card,city, \ state,timestamp,extra} auction: {id,item_name, description,initial_bid, \ reserve_price,timestamp,expires,seller,category,extra} bid: {auction,bidder,price,timestamp,extra} For example: {"id":1000,"name":"Peter Jones","emailAddress":"nhd@xcat.com",\ "creditCard":"7241 7320 9143 4888","city":"Portland","state":"WY",\ "dateTime":1528098831026,\"extra":"WN_HS_bnpVQ\\[["} {"id":1000,"itemName":"wkx mgee","description":"eszpqxtdxrvwmmywkmogoahf",\ "initialBid":28873,"reserve":29448,"dateTime":1528098831036,\ "expires":1528098840451,"seller":1000,"category":13,"extra":"zcuupiz"} {"auction":1000,"bidder":1001,"price":32530001,"dateTime":1528098831066,\ "extra":"fdiysaV^]NLVsbolvyqwgticfdrwdyiyofWPYTOuwogvszlxjrcNOORM"} """ def process(self, elem): json_dict = json.loads(elem) if type(json_dict[FieldNames.DATE_TIME]) is dict: json_dict[FieldNames.DATE_TIME] = json_dict[ FieldNames.DATE_TIME]['millis'] if FieldNames.NAME in json_dict: yield nexmark_model.Person( json_dict[FieldNames.ID], json_dict[FieldNames.NAME], json_dict[FieldNames.EMAIL_ADDRESS], json_dict[FieldNames.CREDIT_CARD], json_dict[FieldNames.CITY], json_dict[FieldNames.STATE], millis_to_timestamp(json_dict[FieldNames.DATE_TIME]), json_dict[FieldNames.EXTRA]) elif FieldNames.ITEM_NAME in json_dict: if type(json_dict[FieldNames.EXPIRES]) is dict: json_dict[FieldNames.EXPIRES] = json_dict[FieldNames.EXPIRES]['millis'] yield nexmark_model.Auction( json_dict[FieldNames.ID], json_dict[FieldNames.ITEM_NAME], json_dict[FieldNames.DESCRIPTION], json_dict[FieldNames.INITIAL_BID], json_dict[FieldNames.RESERVE], millis_to_timestamp(json_dict[FieldNames.DATE_TIME]), millis_to_timestamp(json_dict[FieldNames.EXPIRES]), json_dict[FieldNames.SELLER], json_dict[FieldNames.CATEGORY], json_dict[FieldNames.EXTRA]) elif FieldNames.AUCTION in json_dict: yield nexmark_model.Bid( json_dict[FieldNames.AUCTION], json_dict[FieldNames.BIDDER], json_dict[FieldNames.PRICE], millis_to_timestamp(json_dict[FieldNames.DATE_TIME]), json_dict[FieldNames.EXTRA]) else: raise ValueError('Invalid event: %s.' % str(json_dict)) class CountAndLog(beam.PTransform): def expand(self, pcoll): return ( pcoll | 'window' >> beam.WindowInto(window.GlobalWindows()) | "Count" >> beam.combiners.Count.Globally() | "Log" >> beam.Map(log_count_info)) def log_count_info(count): logging.info('Query resulted in %d results', count) return count def display(elm): logging.debug(elm) return elm def millis_to_timestamp(millis: int) -> Timestamp: micro_second = millis * 1000 return Timestamp(micros=micro_second) def get_counter_metric( result: PipelineResult, namespace: str, name: str) -> int: """ get specific counter metric from pipeline result Args: result: the PipelineResult which metrics are read from namespace: a string representing the namespace of wanted metric name: a string representing the name of the wanted metric Returns: the result of the wanted metric if it exist, else -1 """ metrics = result.metrics().query( MetricsFilter().with_namespace(namespace).with_name(name)) counters = metrics['counters'] if len(counters) > 1: raise RuntimeError( '%d instead of one metric result matches name: %s in namespace %s' % (len(counters), name, namespace)) return counters[0].result if len(counters) > 0 else -1 def get_start_time_metric( result: PipelineResult, namespace: str, name: str) -> int: """ get the start time out of all times recorded by the specified distribution metric Args: result: the PipelineResult which metrics are read from namespace: a string representing the namespace of wanted metric name: a string representing the name of the wanted metric Returns: the smallest time in the metric or -1 if it doesn't exist """ distributions = result.metrics().query( MetricsFilter().with_namespace(namespace).with_name( name))['distributions'] min_list = list(map(lambda m: m.result.min, distributions)) return min(min_list) if len(min_list) > 0 else -1 def get_end_time_metric( result: PipelineResult, namespace: str, name: str) -> int: """ get the end time out of all times recorded by the specified distribution metric Args: result: the PipelineResult which metrics are read from namespace: a string representing the namespace of wanted metric name: a string representing the name of the wanted metric Returns: the largest time in the metric or -1 if it doesn't exist """ distributions = result.metrics().query( MetricsFilter().with_namespace(namespace).with_name( name))['distributions'] max_list = list(map(lambda m: m.result.max, distributions)) return max(max_list) if len(max_list) > 0 else -1 ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.7154431 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/0000755000175100017510000000000015206113160024573 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/__init__.py0000644000175100017510000000142015206113032026677 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/nexmark_query_util.py0000644000175100017510000000445415206113032031101 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utilities for working with NEXmark data stream.""" import apache_beam as beam from apache_beam.testing.benchmarks.nexmark.models import nexmark_model AUCTION_TAG = 'auctions' BID_TAG = 'bids' PERSON_TAG = 'person' class ResultNames: SELLER = 'seller' PRICE = 'price' NAME = 'name' CITY = 'city' STATE = 'state' AUCTION_ID = 'auction_id' ID = 'id' RESERVE = 'reserve' CATEGORY = 'category' IS_LAST = 'is_last' BIDDER_ID = 'bidder_id' BID_COUNT = 'bid_count' NUM = 'num' def is_bid(event): return isinstance(event, nexmark_model.Bid) def is_auction(event): return isinstance(event, nexmark_model.Auction) def is_person(event): return isinstance(event, nexmark_model.Person) def auction_or_bid(event): return isinstance(event, (nexmark_model.Auction, nexmark_model.Bid)) class JustBids(beam.PTransform): def expand(self, pcoll): return pcoll | "IsBid" >> beam.Filter(is_bid) class JustAuctions(beam.PTransform): def expand(self, pcoll): return pcoll | "IsAuction" >> beam.Filter(is_auction) class JustPerson(beam.PTransform): def expand(self, pcoll): return pcoll | "IsPerson" >> beam.Filter(is_person) class AuctionByIdFn(beam.DoFn): def process(self, element): yield element.id, element class BidByAuctionIdFn(beam.DoFn): def process(self, element): yield element.auction, element class PersonByIdFn(beam.DoFn): def process(self, element): yield element.id, element class AuctionBySellerFn(beam.DoFn): def process(self, element): yield element.seller, element ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/query0.py0000644000175100017510000000272015206113032026371 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Nexmark Query 0: Pass through (send and receive auction events). The Nexmark suite is a series of queries (streaming pipelines) performed on a simulation of auction events. This query is a pass through that simply parses the events generated by the launcher. It serves as a test to verify the infrastructure. """ # pytype: skip-file import apache_beam as beam class RoundTripFn(beam.DoFn): def process(self, element): coder = element.CODER byte_value = coder.encode(element) recon = coder.decode(byte_value) yield recon def load(events, metadata=None, pipeline_options=None): return ( events | 'serialization_and_deserialization' >> beam.ParDo(RoundTripFn())) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/query1.py0000644000175100017510000000303615206113032026373 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Nexmark Query 1: Convert bid prices from dollars to euros. The Nexmark suite is a series of queries (streaming pipelines) performed on a simulation of auction events. This query converts bid prices from dollars to euros. It illustrates a simple map. """ # pytype: skip-file import apache_beam as beam from apache_beam.testing.benchmarks.nexmark.models import nexmark_model from apache_beam.testing.benchmarks.nexmark.queries import nexmark_query_util USD_TO_EURO = 0.89 def load(events, metadata=None, pipeline_options=None): return ( events | nexmark_query_util.JustBids() | 'ConvertToEuro' >> beam.Map( lambda bid: nexmark_model.Bid( bid.auction, bid.bidder, bid.price * USD_TO_EURO, bid.date_time, bid.extra))) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/query10.py0000644000175100017510000001257015206113032026456 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Query 10, 'Log to sharded files' (Not in original suite.) Every window_size_sec, save all events from the last period into 2*max_workers log files. """ import apache_beam as beam from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.utils.timestamp import Duration NUM_SHARD_PER_WORKER = 5 LATE_BATCHING_PERIOD = 10 output_path = None max_num_workers = 5 num_log_shards = NUM_SHARD_PER_WORKER * max_num_workers class OutputFile(object): def __init__(self, max_timestamp, shard, index, timing, filename): self.max_timestamp = max_timestamp self.shard = shard self.index = index self.timing = timing self.filename = filename def open_writable_gcs_file(options, filename): # TODO: [https://github.com/apache/beam/issues/20670] it seems that beam team # has not yet decided about this method and it is left blank and # unspecified. pass def output_file_for(window, shard, pane): """ Returns: an OutputFile object constructed with pane, window and shard. """ filename = '%s/LOG-%s-%s-%03d-%s' % ( output_path, window.max_timestamp(), shard, pane.index, pane.timing) if output_path else None return OutputFile( window.max_timestamp(), shard, pane.index, pane.timing, filename) def index_path_for(window): """ Returns: path to the index file containing all shard names or None if no output_path is set """ if output_path: return '%s/INDEX-%s' % (output_path, window.max_timestamp()) else: return None def load(events, metadata=None, pipeline_options=None): return ( events | 'query10_shard_events' >> beam.ParDo(ShardEventsDoFn()) # trigger fires when each sub-triger (executed in order) fires # repeatedly 1. after at least maxLogEvents in pane # 2. or finally when watermark pass the end of window # Repeatedly 1. after at least maxLogEvents in pane # 2. or processing time pass the first element in pane + delay | 'query10_fix_window' >> beam.WindowInto( window.FixedWindows(metadata.get('window_size_sec')), trigger=trigger.AfterEach( trigger.OrFinally( trigger.Repeatedly( trigger.AfterCount(metadata.get('max_log_events'))), trigger.AfterWatermark()), trigger.Repeatedly( trigger.AfterAny( trigger.AfterCount(metadata.get('max_log_events')), trigger.AfterProcessingTime(LATE_BATCHING_PERIOD)))), accumulation_mode=trigger.AccumulationMode.DISCARDING, # Use a 1 day allowed lateness so that any forgotten hold will stall # the pipeline for that period and be very noticeable. allowed_lateness=Duration.of(1 * 24 * 60 * 60)) | 'query10_gbk' >> beam.GroupByKey() | 'query10_write_event' >> beam.ParDo(WriteEventDoFn(), pipeline_options) | 'query10_window_log_files' >> beam.WindowInto( window.FixedWindows(metadata.get('window_size_sec')), accumulation_mode=trigger.AccumulationMode.DISCARDING, allowed_lateness=Duration.of(1 * 24 * 60 * 60)) | 'query10_gbk_2' >> beam.GroupByKey() | 'query10_write_index' >> beam.ParDo(WriteIndexDoFn(), pipeline_options)) class ShardEventsDoFn(beam.DoFn): def process(self, element): shard_number = abs(hash(element) % num_log_shards) shard = 'shard-%05d-of-%05d' % (shard_number, num_log_shards) yield shard, element class WriteEventDoFn(beam.DoFn): def process( self, element, pipeline_options, window=beam.DoFn.WindowParam, pane_info=beam.DoFn.PaneInfoParam): shard = element[0] options = pipeline_options.view_as(GoogleCloudOptions) output_file = output_file_for(window, shard, pane_info) if output_file.filename: # not do anything because open_writable_gcs_file does not do anything open_writable_gcs_file(options, output_file.filename) for event in element[1]: # pylint: disable=unused-variable # write to file pass yield None, output_file class WriteIndexDoFn(beam.DoFn): def process(self, element, pipeline_options, window=beam.DoFn.WindowParam): options = pipeline_options.view_as(GoogleCloudOptions) filename = index_path_for(window) if filename: # not do anything because open_writable_gcs_file does not do anything open_writable_gcs_file(options, filename) for output_file in element[1]: # pylint: disable=unused-variable # write to file pass ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/query11.py0000644000175100017510000000427315206113032026460 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Query 11, How many bids did a user make in each session he was active? (Not in original suite.) Group bids by the same user into sessions with window_size_sec max gap. However limit the session to at most max_log_events. Emit the number of bids per session. """ import apache_beam as beam from apache_beam.testing.benchmarks.nexmark.queries import nexmark_query_util from apache_beam.testing.benchmarks.nexmark.queries.nexmark_query_util import ResultNames from apache_beam.transforms import trigger from apache_beam.transforms import window def load(events, metadata=None, pipeline_options=None): return ( events # filter to get only bids and then extract bidder id | nexmark_query_util.JustBids() | 'query11_extract_bidder' >> beam.Map(lambda bid: bid.bidder) # window auction and key by auctions' seller | 'query11_session_window' >> beam.WindowInto( window.Sessions(metadata.get('window_size_sec')), trigger=trigger.AfterWatermark( early=trigger.AfterCount(metadata.get('max_log_events'))), accumulation_mode=trigger.AccumulationMode.DISCARDING, allowed_lateness=metadata.get('occasional_delay_sec') // 2) # count per bidder | beam.combiners.Count.PerElement() | beam.Map( lambda bidder_count: { ResultNames.BIDDER_ID: bidder_count[0], ResultNames.BID_COUNT: bidder_count[1] })) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/query12.py0000644000175100017510000000376715206113032026470 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Query 12, How many bids does a user make within a fixed processing time limit (Not in original suite.) Group bids by the same user into processing time windows of window_size_sec. Emit the count of bids per window. """ import apache_beam as beam from apache_beam.testing.benchmarks.nexmark.queries import nexmark_query_util from apache_beam.testing.benchmarks.nexmark.queries.nexmark_query_util import ResultNames from apache_beam.transforms import trigger from apache_beam.transforms import window def load(events, metadata=None, pipeline_options=None): return ( events | nexmark_query_util.JustBids() | 'query12_extract_bidder' >> beam.Map(lambda bid: bid.bidder) # windowing with processing time trigger, currently not supported in batch | beam.WindowInto( window.GlobalWindows(), trigger=trigger.Repeatedly( trigger.AfterProcessingTime(metadata.get('window_size_sec'))), accumulation_mode=trigger.AccumulationMode.DISCARDING, allowed_lateness=0) | 'query12_bid_count' >> beam.combiners.Count.PerElement() | 'query12_output' >> beam.Map( lambda t: { ResultNames.BIDDER_ID: t[0], ResultNames.BID_COUNT: t[1] })) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/query2.py0000644000175100017510000000307115206113032026373 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Query 2: Find bids with specific auction ids and show their bid price This query selects Bids that have a particular auctiuon id, and output their auction id with bid price. It illustrates a simple filter. """ # pytype: skip-file import apache_beam as beam from apache_beam.testing.benchmarks.nexmark.queries import nexmark_query_util from apache_beam.testing.benchmarks.nexmark.queries.nexmark_query_util import ResultNames def load(events, metadata=None, pipeline_options=None): return ( events | nexmark_query_util.JustBids() | 'filter_by_skip' >> beam.Filter(lambda bid: bid.auction % metadata.get('auction_skip') == 0) | 'project' >> beam.Map( lambda bid: { ResultNames.AUCTION_ID: bid.auction, ResultNames.PRICE: bid.price })) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/query3.py0000644000175100017510000001404315206113032026375 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Query 3, 'Local Item Suggestion'. Who is selling in OR, ID or CA in category 10, and for what auction ids? In CQL syntax:: SELECT Istream(P.name, P.city, P.state, A.id) FROM Auction A [ROWS UNBOUNDED], Person P [ROWS UNBOUNDED] WHERE A.seller = P.id AND (P.state = `OR' OR P.state = `ID' OR P.state = `CA') AND A.category = 10; We'll implement this query to allow 'new auction' events to come before the 'new person' events for the auction seller. Those auctions will be stored until the matching person is seen. Then all subsequent auctions for a person will use the stored person record. """ import logging import typing import apache_beam as beam from apache_beam.testing.benchmarks.nexmark.models import nexmark_model from apache_beam.testing.benchmarks.nexmark.queries import nexmark_query_util from apache_beam.testing.benchmarks.nexmark.queries.nexmark_query_util import ResultNames from apache_beam.transforms import trigger from apache_beam.transforms import userstate from apache_beam.transforms import window from apache_beam.transforms.userstate import on_timer def load(events, metadata=None, pipeline_options=None): num_events_in_pane = 30 windowed_events = ( events | beam.WindowInto( window.GlobalWindows(), trigger=trigger.Repeatedly(trigger.AfterCount(num_events_in_pane)), accumulation_mode=trigger.AccumulationMode.DISCARDING)) auction_by_seller_id = ( windowed_events | nexmark_query_util.JustAuctions() | 'query3_filter_category' >> beam.Filter(lambda auc: auc.category == 10) | 'query3_key_by_seller' >> beam.ParDo( nexmark_query_util.AuctionBySellerFn())) person_by_id = ( windowed_events | nexmark_query_util.JustPerson() | 'query3_filter_region' >> beam.Filter(lambda person: person.state in ['OR', 'ID', 'CA']) | 'query3_key_by_person_id' >> beam.ParDo( nexmark_query_util.PersonByIdFn())) return ({ nexmark_query_util.AUCTION_TAG: auction_by_seller_id, nexmark_query_util.PERSON_TAG: person_by_id, } | beam.CoGroupByKey() | 'query3_join' >> beam.ParDo( JoinFn(metadata.get('max_auction_waiting_time'))) | 'query3_output' >> beam.Map( lambda t: { ResultNames.NAME: t[1].name, ResultNames.CITY: t[1].city, ResultNames.STATE: t[1].state, ResultNames.AUCTION_ID: t[0].id })) class JoinFn(beam.DoFn): """ Join auctions and person by person id and emit their product one pair at a time. We know a person may submit any number of auctions. Thus new person event must have the person record stored in persistent state in order to match future auctions by that person. However we know that each auction is associated with at most one person, so only need to store auction records in persistent state until we have seen the corresponding person record. And of course may have already seen that record. """ AUCTIONS = 'auctions_state' PERSON = 'person_state' PERSON_EXPIRING = 'person_state_expiring' auction_spec = userstate.BagStateSpec(AUCTIONS, nexmark_model.Auction.CODER) person_spec = userstate.ReadModifyWriteStateSpec( PERSON, nexmark_model.Person.CODER) person_timer_spec = userstate.TimerSpec( PERSON_EXPIRING, userstate.TimeDomain.WATERMARK) def __init__(self, max_auction_wait_time): self.max_auction_wait_time = max_auction_wait_time def process( # type: ignore self, element: tuple[ str, dict[str, typing.Union[list[nexmark_model.Auction], list[nexmark_model.Person]]]], auction_state=beam.DoFn.StateParam(auction_spec), person_state=beam.DoFn.StateParam(person_spec), person_timer=beam.DoFn.TimerParam(person_timer_spec)): # extract group with tags from element tuple _, group = element existing_person = person_state.read() if existing_person: # the person exists in person_state for this person id for auction in group[nexmark_query_util.AUCTION_TAG]: yield auction, existing_person return new_person = None for person in group[nexmark_query_util.PERSON_TAG]: if not new_person: new_person = person else: logging.error( 'two new person wtih same key: %s and %s' % (person, new_person)) continue # read all pending auctions for this person id, output and flush it pending_auctions = auction_state.read() if pending_auctions: for pending_auction in pending_auctions: yield pending_auction, new_person auction_state.clear() # output new auction for this person id for auction in group[nexmark_query_util.AUCTION_TAG]: yield auction, new_person # remember person for max_auction_wait_time seconds for future auctions person_state.write(new_person) person_timer.set(new_person.date_time + self.max_auction_wait_time) # we are done if we have seen a new person if new_person: return # remember auction until we see person for auction in group[nexmark_query_util.AUCTION_TAG]: auction_state.add(auction) @on_timer(person_timer_spec) def expiry(self, person_state=beam.DoFn.StateParam(person_spec)): person_state.clear() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/query4.py0000644000175100017510000000630415206113032026377 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Query 4, 'Average Price for a Category'. Select the average of the wining bid prices for all closed auctions in each category. In CQL syntax:: SELECT Istream(AVG(Q.final)) FROM Category C, (SELECT Rstream(MAX(B.price) AS final, A.category) FROM Auction A [ROWS UNBOUNDED], Bid B [ROWS UNBOUNDED] WHERE A.id=B.auction AND B.datetime < A.expires AND A.expires < CURRENT_TIME GROUP BY A.id, A.category) Q WHERE Q.category = C.id GROUP BY C.id; For extra spiciness our implementation differs slightly from the above: * We select both the average winning price and the category. * We don't bother joining with a static category table, since it's contents are never used. * We only consider bids which are above the auction's reserve price. * We accept the highest-price, earliest valid bid as the winner. * We calculate the averages oven a sliding window of size window_size_sec and period window_period_sec. """ import apache_beam as beam from apache_beam.testing.benchmarks.nexmark.queries import nexmark_query_util from apache_beam.testing.benchmarks.nexmark.queries import winning_bids from apache_beam.testing.benchmarks.nexmark.queries.nexmark_query_util import ResultNames from apache_beam.transforms import window def load(events, metadata=None, pipeline_options=None): # find winning bids for each closed auction all_winning_bids = ( events | beam.Filter(nexmark_query_util.auction_or_bid) | winning_bids.WinningBids()) return ( all_winning_bids # key winning bids by auction category | beam.Map(lambda auc_bid: (auc_bid.auction.category, auc_bid.bid.price)) # re-window for sliding average | beam.WindowInto( window.SlidingWindows( metadata.get('window_size_sec'), metadata.get('window_period_sec'))) # average for each category | beam.CombinePerKey(beam.combiners.MeanCombineFn()) # TODO(leiyiz): fanout with sliding window produces duplicated results, # uncomment after it is fixed # [https://github.com/apache/beam/issues/20528] # .with_hot_key_fanout(metadata.get('fanout')) # produce output | beam.ParDo(ProjectToCategoryPriceFn())) class ProjectToCategoryPriceFn(beam.DoFn): def process(self, element, pane_info=beam.DoFn.PaneInfoParam): yield { ResultNames.CATEGORY: element[0], ResultNames.PRICE: element[1], ResultNames.IS_LAST: pane_info.is_last } ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/query5.py0000644000175100017510000000662115206113032026402 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Query 5, 'Hot Items'. Which auctions have seen the most bids in the last hour (updated every minute). In CQL syntax:: SELECT Rstream(auction) FROM (SELECT B1.auction, count(*) AS num FROM Bid [RANGE 60 MINUTE SLIDE 1 MINUTE] B1 GROUP BY B1.auction) WHERE num >= ALL (SELECT count(*) FROM Bid [RANGE 60 MINUTE SLIDE 1 MINUTE] B2 GROUP BY B2.auction); To make things a bit more dynamic and easier to test we use much shorter windows, and we'll also preserve the bid counts. """ import apache_beam as beam from apache_beam.testing.benchmarks.nexmark.queries import nexmark_query_util from apache_beam.testing.benchmarks.nexmark.queries.nexmark_query_util import ResultNames from apache_beam.transforms import window def load(events, metadata=None, pipeline_options=None): return ( events | nexmark_query_util.JustBids() | 'query5_sliding_window' >> beam.WindowInto( window.SlidingWindows( metadata.get('window_size_sec'), metadata.get('window_period_sec'))) # project out only the auction id for each bid | 'extract_bid_auction' >> beam.Map(lambda bid: bid.auction) | 'bid_count_per_auction' >> beam.combiners.Count.PerElement() | 'bid_max_count' >> beam.CombineGlobally( MostBidCombineFn()).without_defaults() # TODO(leiyiz): fanout with sliding window produces duplicated results, # uncomment after it is fixed # [https://github.com/apache/beam/issues/20528] # .with_fanout(metadata.get('fanout')) | beam.FlatMap( lambda auc_count: [{ ResultNames.AUCTION_ID: auction, ResultNames.NUM: auc_count[1] } for auction in auc_count[0]])) class MostBidCombineFn(beam.CombineFn): """ combiner function to find auctions with most bid counts """ def create_accumulator(self): return [], 0 def add_input(self, accumulator, element): accu_list, accu_count = accumulator auction, count = element if accu_count < count: return [auction], count elif accu_count > count: return accu_list, accu_count else: accu_list_new = accu_list.copy() accu_list_new.append(auction) return accu_list_new, accu_count def merge_accumulators(self, accumulators): max_list = [] max_count = 0 for (accu_list, count) in accumulators: if count == max_count: max_list = max_list + accu_list elif count < max_count: continue else: max_list = accu_list max_count = count return max_list, max_count def extract_output(self, accumulator): return accumulator ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/query6.py0000644000175100017510000000647715206113032026414 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Query 6, 'Average Selling Price by Seller'. Select the average selling price over the last 10 closed auctions by the same seller. In CQL syntax:: SELECT Istream(AVG(Q.final), Q.seller) FROM (SELECT Rstream(MAX(B.price) AS final, A.seller) FROM Auction A [ROWS UNBOUNDED], Bid B [ROWS UNBOUNDED] WHERE A.id=B.auction AND B.datetime < A.expires AND A.expires < CURRENT_TIME GROUP BY A.id, A.seller) [PARTITION BY A.seller ROWS 10] Q GROUP BY Q.seller; """ import apache_beam as beam from apache_beam.testing.benchmarks.nexmark.queries import nexmark_query_util from apache_beam.testing.benchmarks.nexmark.queries import winning_bids from apache_beam.testing.benchmarks.nexmark.queries.nexmark_query_util import ResultNames from apache_beam.transforms import trigger from apache_beam.transforms import window def load(events, metadata=None, pipeline_options=None): # find winning bids for each closed auction return ( events # find winning bids | beam.Filter(nexmark_query_util.auction_or_bid) | winning_bids.WinningBids() # (auction_bids -> (aution.seller, bid) | beam.Map(lambda auc_bid: (auc_bid.auction.seller, auc_bid.bid)) # calculate and output mean as data arrives | beam.WindowInto( window.GlobalWindows(), trigger=trigger.Repeatedly(trigger.AfterCount(1)), accumulation_mode=trigger.AccumulationMode.ACCUMULATING, allowed_lateness=0) | beam.CombinePerKey(MovingMeanSellingPriceFn(10)) | beam.Map(lambda t: { ResultNames.SELLER: t[0], ResultNames.PRICE: t[1] })) class MovingMeanSellingPriceFn(beam.CombineFn): """ Combiner to keep track of up to max_num_bids of the most recent wining bids and calculate their average selling price. """ def __init__(self, max_num_bids): self.max_num_bids = max_num_bids def create_accumulator(self): return [] def add_input(self, accumulator, element): accumulator.append(element) new_accu = sorted(accumulator, key=lambda bid: (-bid.date_time, -bid.price)) if len(new_accu) > self.max_num_bids: del new_accu[self.max_num_bids] return new_accu def merge_accumulators(self, accumulators): new_accu = [] for accumulator in accumulators: new_accu += accumulator new_accu.sort(key=lambda bid: (bid.date_time, bid.price)) return new_accu[-self.max_num_bids:] def extract_output(self, accumulator): if len(accumulator) == 0: return 0 sum_price = sum(bid.price for bid in accumulator) return int(sum_price / len(accumulator)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/query7.py0000644000175100017510000000414115206113032026377 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Query 7, 'Highest Bid'. Select the bids with the highest bid price in the last minute. In CQL syntax:: SELECT Rstream(B.auction, B.price, B.bidder) FROM Bid [RANGE 1 MINUTE SLIDE 1 MINUTE] B WHERE B.price = (SELECT MAX(B1.price) FROM BID [RANGE 1 MINUTE SLIDE 1 MINUTE] B1); We will use a shorter window to help make testing easier. We'll also implement this using a side-input in order to exercise that functionality. (A combiner, as used in Query 5, is a more efficient approach.). """ import apache_beam as beam from apache_beam.testing.benchmarks.nexmark.queries import nexmark_query_util from apache_beam.transforms import window def load(events, metadata=None, pipeline_options=None): # window bids into fixed window sliding_bids = ( events | nexmark_query_util.JustBids() | beam.WindowInto(window.FixedWindows(metadata.get('window_size_sec')))) # find the largest price in all bids per window max_prices = ( sliding_bids | beam.Map(lambda bid: bid.price) | beam.CombineGlobally(max).without_defaults()) return ( sliding_bids | 'select_bids' >> beam.ParDo( SelectMaxBidFn(), beam.pvalue.AsSingleton(max_prices))) class SelectMaxBidFn(beam.DoFn): def process(self, element, max_bid_price): if element.price == max_bid_price: yield element ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/query8.py0000644000175100017510000000542215206113032026403 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Query 8, 'Monitor New Users'. Select people who have entered the system and created auctions in the last 12 hours, updated every 12 hours. In CQL syntax:: SELECT Rstream(P.id, P.name, A.reserve) FROM Person [RANGE 12 HOUR] P, Auction [RANGE 12 HOUR] A WHERE P.id = A.seller; To make things a bit more dynamic and easier to test we'll use a much shorter window. """ import apache_beam as beam from apache_beam.testing.benchmarks.nexmark.queries import nexmark_query_util from apache_beam.testing.benchmarks.nexmark.queries.nexmark_query_util import ResultNames from apache_beam.transforms import window def load(events, metadata=None, pipeline_options=None): # window person and key by persons' id persons_by_id = ( events | nexmark_query_util.JustPerson() | 'query8_window_person' >> beam.WindowInto( window.FixedWindows(metadata.get('window_size_sec'))) | 'query8_person_by_id' >> beam.ParDo(nexmark_query_util.PersonByIdFn())) # window auction and key by auctions' seller auctions_by_seller = ( events | nexmark_query_util.JustAuctions() | 'query8_window_auction' >> beam.WindowInto( window.FixedWindows(metadata.get('window_size_sec'))) | 'query8_auction_by_seller' >> beam.ParDo( nexmark_query_util.AuctionBySellerFn())) return ({ nexmark_query_util.PERSON_TAG: persons_by_id, nexmark_query_util.AUCTION_TAG: auctions_by_seller } | beam.CoGroupByKey() | 'query8_join' >> beam.ParDo(JoinPersonAuctionFn())) class JoinPersonAuctionFn(beam.DoFn): def process(self, element): _, group = element persons = group[nexmark_query_util.PERSON_TAG] person = persons[0] if persons else None if person is None: # do nothing if this seller id is not a new person in this window return for auction in group[nexmark_query_util.AUCTION_TAG]: yield { ResultNames.ID: person.id, ResultNames.NAME: person.name, ResultNames.RESERVE: auction.reserve } ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/query9.py0000644000175100017510000000235215206113032026403 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Query 9: Winning-bids: extract the most recent of the highest bids See winning_bids.py for detailed documentation """ import apache_beam as beam from apache_beam.testing.benchmarks.nexmark.queries import nexmark_query_util from apache_beam.testing.benchmarks.nexmark.queries import winning_bids def load(events, metadata=None, pipeline_options=None): return ( events | beam.Filter(nexmark_query_util.auction_or_bid) | winning_bids.WinningBids()) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/nexmark/queries/winning_bids.py0000644000175100017510000001572415206113032027626 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ A transform to find winning bids for each closed auction. In pseudo CQL syntax: SELECT Rstream(A.*, B.auction, B.bidder, MAX(B.price), B.dateTime) FROM Auction A [ROWS UNBOUNDED], Bid B [ROWS UNBOUNDED] WHERE A.id = B.auction AND B.datetime < A.expires AND A.expires < CURRENT_TIME GROUP BY A.id We will also check that the winning bid is above the auction reserve. Note that we ignore the auction opening bid value since it has no impact on which bid eventually wins, if any. Our implementation will use a custom windowing function in order to bring bids and auctions together without requiring global state. """ import apache_beam as beam from apache_beam.coders import coder_impl from apache_beam.coders.coders import FastCoder from apache_beam.testing.benchmarks.nexmark.models import auction_bid from apache_beam.testing.benchmarks.nexmark.models import nexmark_model from apache_beam.testing.benchmarks.nexmark.queries import nexmark_query_util from apache_beam.transforms.window import IntervalWindow from apache_beam.transforms.window import WindowFn from apache_beam.utils.timestamp import Duration class AuctionOrBidWindow(IntervalWindow): """Windows for open auctions and bids.""" def __init__(self, start, end, auction_id, is_auction_window): super().__init__(start, end) self.auction = auction_id self.is_auction_window = is_auction_window @staticmethod def for_auction(timestamp, auction): return AuctionOrBidWindow(timestamp, auction.expires, auction.id, True) @staticmethod def for_bid(expected_duration_micro, timestamp, bid): return AuctionOrBidWindow( timestamp, timestamp + Duration(micros=expected_duration_micro * 2), bid.auction, False) def is_auction_window_fn(self): return self.is_auction_window def __str__(self): return ( 'AuctionOrBidWindow{start:%s; end:%s; auction:%d; isAuctionWindow:%s}' % (self.start, self.end, self.auction, self.is_auction_window)) class AuctionOrBidWindowCoder(FastCoder): def _create_impl(self): return AuctionOrBidWindowCoderImpl() def is_deterministic(self): return True class AuctionOrBidWindowCoderImpl(coder_impl.StreamCoderImpl): _super_coder_impl = coder_impl.IntervalWindowCoderImpl() _id_coder_impl = coder_impl.VarIntCoderImpl() _bool_coder_impl = coder_impl.BooleanCoderImpl() def encode_to_stream(self, value, stream, nested): self._super_coder_impl.encode_to_stream(value, stream, True) self._id_coder_impl.encode_to_stream(value.auction, stream, True) self._bool_coder_impl.encode_to_stream( value.is_auction_window, stream, True) def decode_from_stream(self, stream, nested): super_window = self._super_coder_impl.decode_from_stream(stream, True) auction = self._id_coder_impl.decode_from_stream(stream, True) is_auction = self._bool_coder_impl.decode_from_stream(stream, True) return AuctionOrBidWindow( super_window.start, super_window.end, auction, is_auction) class AuctionOrBidWindowFn(WindowFn): def __init__(self, expected_duration_micro): self.expected_duration = expected_duration_micro def assign(self, assign_context): event = assign_context.element if isinstance(event, nexmark_model.Auction): return [AuctionOrBidWindow.for_auction(assign_context.timestamp, event)] elif isinstance(event, nexmark_model.Bid): return [ AuctionOrBidWindow.for_bid( self.expected_duration, assign_context.timestamp, event) ] else: raise ValueError( '%s can only assign windows to auctions and bids, but received %s' % (self.__class__.__name__, event)) def merge(self, merge_context): auction_id_to_auction_window = {} auction_id_to_bid_window = {} for window in merge_context.windows: if window.is_auction_window_fn(): auction_id_to_auction_window[window.auction] = window else: if window.auction not in auction_id_to_bid_window: auction_id_to_bid_window[window.auction] = [] auction_id_to_bid_window[window.auction].append(window) for auction, auction_window in auction_id_to_auction_window.items(): bid_window_list = auction_id_to_bid_window.get(auction) if bid_window_list is not None: to_merge = [] for bid_window in bid_window_list: if bid_window.start < auction_window.end: to_merge.append(bid_window) if len(to_merge) > 0: to_merge.append(auction_window) merge_context.merge(to_merge, auction_window) def get_window_coder(self): return AuctionOrBidWindowCoder() def get_transformed_output_time(self, window, input_timestamp): return window.max_timestamp() class JoinAuctionBidFn(beam.DoFn): @staticmethod def higher_bid(bid, other): if bid.price > other.price: return True elif bid.price < other.price: return False else: return bid.date_time < other.date_time def process(self, element): _, group = element auctions = group[nexmark_query_util.AUCTION_TAG] auction = auctions[0] if auctions else None if auction is None: return best_bid = None for bid in group[nexmark_query_util.BID_TAG]: if bid.price < auction.reserve: continue if best_bid is None or JoinAuctionBidFn.higher_bid(bid, best_bid): best_bid = bid if best_bid: yield auction_bid.AuctionBid(auction, best_bid) class WinningBids(beam.PTransform): def __init__(self): #TODO: change this to be calculated by event generation expected_duration = 16667000 self.auction_or_bid_windowFn = AuctionOrBidWindowFn(expected_duration) def expand(self, pcoll): events = pcoll | beam.WindowInto(self.auction_or_bid_windowFn) auction_by_id = ( events | nexmark_query_util.JustAuctions() | 'auction_by_id' >> beam.ParDo(nexmark_query_util.AuctionByIdFn())) bids_by_auction_id = ( events | nexmark_query_util.JustBids() | 'bid_by_auction' >> beam.ParDo(nexmark_query_util.BidByAuctionIdFn())) return ({ nexmark_query_util.AUCTION_TAG: auction_by_id, nexmark_query_util.BID_TAG: bids_by_auction_id } | beam.CoGroupByKey() | beam.ParDo(JoinAuctionBidFn())) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/sort_and_batch_benchmark.py0000644000175100017510000005303715206113032027015 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Benchmark: BatchElements vs SortAndBatchElements on real Beam pipelines. Compares two batching strategies for variable-length inference workloads by running the actual Beam transforms under DirectRunner: - Baseline (BatchElements): fixed-count batching by setting ``min_batch_size == max_batch_size``. - Stateless (SortAndBatchElements): sorts elements by size within each runner bundle, then splits batches using ``max_batch_weight``. The benchmark materializes per-batch summaries through a temporary Beam sink and analyzes them after the pipeline completes. This keeps the benchmark on the normal Beam execution path rather than relying on InteractiveRunner-specific result materialization or local side effects. Bundle boundaries are runner-defined. As a result, these measurements are meant to compare the actual DirectRunner behavior of the two transforms rather than a synthetic, user-configurable bundle model. Padding ratio:: padding_ratio = sum(max_len_in_batch * batch_size) / sum(actual_lengths) Lower is better. 1.0 = no padding waste. Methodology: - N=20 independent trials per condition (3 warmup trials excluded). - Same input corpus (seed=42) for A/B comparison. - DirectRunner with in-memory execution and one worker for reproducibility. - Percentile method: linear interpolation between adjacent ranks (equivalent to numpy.percentile with method='linear'). For N=20 trials: P50 interpolates ranks 10-11 (0-indexed 9-10), P95 interpolates ranks 19-20 (0-indexed 18-19), P99 interpolates near rank 20 (0-indexed 18.81). - Reports median [IQR] and P95 for each metric. - Inference model: latency = batch_size * (max_seq_len / 50)^1.5 ms (simulates downstream transformer-like scaling). Run:: python3 -m apache_beam.testing.benchmarks.sort_and_batch_benchmark """ import glob import json import math import os import random import statistics import tempfile import time from collections.abc import Sequence from typing import Any import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.transforms import util # --------------------------------------------------------------------------- # Data generators # --------------------------------------------------------------------------- def generate_highly_skewed_data( num_elements: int, min_length: int = 1, max_length: int = 500, seed: int = 42) -> list[str]: """Pareto(alpha=1.2) -- most short, few very long.""" random.seed(seed) data = [] for _ in range(num_elements): length = int(random.paretovariate(1.2) * min_length) length = min(max(length, min_length), max_length) data.append('x' * length) return data def generate_lognormal_data( num_elements: int, mean_length: int = 50, std_factor: float = 0.8, min_length: int = 1, max_length: int = 500, seed: int = 42) -> list[str]: """Log-normal -- moderate skew, typical NLP.""" random.seed(seed) mu = math.log(mean_length) sigma = std_factor data = [] for _ in range(num_elements): length = int(random.lognormvariate(mu, sigma)) length = min(max(length, min_length), max_length) data.append('x' * length) return data def generate_bimodal_data( num_elements: int, mode1_mean: int = 20, mode2_mean: int = 200, mode1_ratio: float = 0.7, min_length: int = 1, max_length: int = 500, seed: int = 42) -> list[str]: """Bimodal -- two distinct length groups.""" random.seed(seed) data = [] for _ in range(num_elements): if random.random() < mode1_ratio: length = int(random.gauss(mode1_mean, mode1_mean * 0.3)) else: length = int(random.gauss(mode2_mean, mode2_mean * 0.3)) length = min(max(length, min_length), max_length) data.append('x' * length) return data def generate_low_variance_data( num_elements: int, mean_length: int = 100, cv: float = 0.1, min_length: int = 1, max_length: int = 500, seed: int = 42) -> list[str]: """Low-variance control (CV=10%).""" random.seed(seed) std = mean_length * cv data = [] for _ in range(num_elements): length = int(random.gauss(mean_length, std)) length = min(max(length, min_length), max_length) data.append('x' * length) return data # --------------------------------------------------------------------------- # Real Beam batching # --------------------------------------------------------------------------- def _direct_runner_options() -> PipelineOptions: return PipelineOptions([ '--runner=DirectRunner', '--direct_running_mode=in_memory', '--direct_num_workers=1', ]) def _batch_to_json(batch: list[str]) -> str: lengths = [len(element) for element in batch] return json.dumps({ 'batch_size': len(batch), 'actual_total_length': sum(lengths), 'max_len': max(lengths) if lengths else 0, }) def _read_batch_summaries(output_prefix: str) -> list[dict[str, int]]: summaries = [] for path in sorted(glob.glob(f'{output_prefix}*')): if path.endswith('.crc'): continue with open(path, encoding='utf-8') as handle: for line in handle: line = line.strip() if line: summaries.append(json.loads(line)) return summaries def _run_batching_pipeline( strategy: str, data: list[str], max_batch_size: int, max_batch_weight: int) -> tuple[list[dict[str, int]], float]: """Runs one Beam pipeline and returns batch summaries plus runtime.""" with tempfile.TemporaryDirectory(prefix='beam_batch_benchmark_') as temp_dir: output_prefix = os.path.join(temp_dir, strategy) pipeline = beam.Pipeline(options=_direct_runner_options()) batched = pipeline | 'CreateInput' >> beam.Create(data, reshuffle=False) if strategy == 'baseline': batched = batched | 'BatchElements' >> util.BatchElements( min_batch_size=max_batch_size, max_batch_size=max_batch_size) elif strategy == 'stateless': batched = batched | 'SortAndBatchElements' >> util.SortAndBatchElements( min_batch_size=1, max_batch_size=max_batch_size, max_batch_weight=max_batch_weight) else: raise ValueError(f'Unknown strategy: {strategy}') _ = ( batched | 'SerializeBatchSummary' >> beam.Map(_batch_to_json) | 'WriteBatchSummary' >> beam.io.WriteToText(output_prefix)) start = time.perf_counter() result = pipeline.run() result.wait_until_finish() runtime_ms = (time.perf_counter() - start) * 1000 return _read_batch_summaries(output_prefix), runtime_ms # --------------------------------------------------------------------------- # Simulated inference # --------------------------------------------------------------------------- def simulate_inference_latency( batch_size: int, max_len: int, base_latency_ms: float = 1.0) -> float: """Simulate downstream inference: O(batch_size * seq_len^1.5).""" if not batch_size or not max_len: return 0.0 return base_latency_ms * batch_size * (max_len / 50)**1.5 # --------------------------------------------------------------------------- # Stats helpers # --------------------------------------------------------------------------- def percentile(data: Sequence[float], p: float) -> float: """Percentile via linear interpolation between adjacent ranks. Equivalent to numpy.percentile(data, p, method='linear'). For N=20: P50 interpolates ranks 10-11, P95 ranks 19-20, P99 near rank 20 (fractional index 18.81). """ if not data: return 0.0 s = sorted(data) k = (len(s) - 1) * p / 100 f = int(k) c = min(f + 1, len(s) - 1) return s[f] + (k - f) * (s[c] - s[f]) def compute_padding_stats( batch_summaries: list[dict[str, int]]) -> dict[str, Any]: """Padding-efficiency statistics for materialized batch summaries.""" total_actual = sum(s['actual_total_length'] for s in batch_summaries) total_padded = sum(s['max_len'] * s['batch_size'] for s in batch_summaries) batch_sizes = [s['batch_size'] for s in batch_summaries if s['batch_size']] max_lengths = [s['max_len'] for s in batch_summaries if s['batch_size']] efficiency = total_actual / total_padded if total_padded else 0.0 padding_ratio = total_padded / total_actual if total_actual else float('inf') return { 'efficiency': efficiency, 'padding_ratio': padding_ratio, 'num_batches': len(batch_summaries), 'avg_batch_size': statistics.mean(batch_sizes) if batch_sizes else 0, 'total_actual_length': total_actual, 'total_padded_length': total_padded, 'padding_overhead': total_padded - total_actual, 'batch_size_p50': percentile(batch_sizes, 50) if batch_sizes else 0, 'batch_size_p95': percentile(batch_sizes, 95) if batch_sizes else 0, 'batch_size_max': max(batch_sizes) if batch_sizes else 0, 'max_len_p50': percentile(max_lengths, 50) if max_lengths else 0, 'max_len_p95': percentile(max_lengths, 95) if max_lengths else 0, } # --------------------------------------------------------------------------- # Invariant validation # --------------------------------------------------------------------------- def validate_invariants( data: list[str], baseline_summaries: list[dict[str, int]], stateless_summaries: list[dict[str, int]]) -> dict[str, Any]: """Validate element/token counts and batch-size equality.""" n = len(data) b_n = sum(s['batch_size'] for s in baseline_summaries) s_n = sum(s['batch_size'] for s in stateless_summaries) tok = sum(len(s) for s in data) b_tok = sum(s['actual_total_length'] for s in baseline_summaries) s_tok = sum(s['actual_total_length'] for s in stateless_summaries) return { 'input_elements': n, 'baseline_elements': b_n, 'stateless_elements': s_n, 'elements_match': n == b_n == s_n, 'input_tokens': tok, 'baseline_tokens': b_tok, 'stateless_tokens': s_tok, 'tokens_match': tok == b_tok == s_tok, 'baseline_num_batches': len(baseline_summaries), 'stateless_num_batches': len(stateless_summaries), } # --------------------------------------------------------------------------- # Performance benchmark (N=20 trials) # --------------------------------------------------------------------------- def run_performance_benchmark( data: list[str], max_batch_size: int, max_batch_weight: int, num_trials: int = 20, warmup_trials: int = 3 ) -> tuple[ dict[str, Any], dict[str, Any], list[dict[str, int]], list[dict[str, int]], ]: """Run N=20 trials for baseline and stateless.""" total_tokens = sum(len(s) for s in data) baseline_trials = [] stateless_trials = [] baseline_sample_summaries = [] stateless_sample_summaries = [] for trial_idx in range(warmup_trials + num_trials): is_warmup = trial_idx < warmup_trials trial_results = {} if trial_idx % 2 == 0: trial_order = ('baseline', 'stateless') else: trial_order = ('stateless', 'baseline') for strategy in trial_order: summaries, runtime_ms = _run_batching_pipeline( strategy, data, max_batch_size, max_batch_weight) batch_latencies = [ simulate_inference_latency(s['batch_size'], s['max_len']) for s in summaries ] trial_results[strategy] = { 'runtime_ms': runtime_ms, 'inference_ms': sum(batch_latencies), 'e2e_ms': runtime_ms + sum(batch_latencies), 'batch_latencies': batch_latencies, 'num_batches': len(summaries), 'summaries': summaries, } if not is_warmup: baseline_trials.append(trial_results['baseline']) stateless_trials.append(trial_results['stateless']) if not baseline_sample_summaries: baseline_sample_summaries = trial_results['baseline']['summaries'] if not stateless_sample_summaries: stateless_sample_summaries = trial_results['stateless']['summaries'] def _stats(trials): e2e = [t['e2e_ms'] for t in trials] tput = [total_tokens / (t['e2e_ms'] / 1000) for t in trials] runtime = [t['runtime_ms'] for t in trials] all_lat = [l for t in trials for l in t['batch_latencies']] return { 'e2e_median': percentile(e2e, 50), 'e2e_p25': percentile(e2e, 25), 'e2e_p75': percentile(e2e, 75), 'e2e_p95': percentile(e2e, 95), 'tput_median': percentile(tput, 50), 'tput_p25': percentile(tput, 25), 'tput_p75': percentile(tput, 75), 'tput_p95': percentile(tput, 95), 'runtime_median': percentile(runtime, 50), 'runtime_p25': percentile(runtime, 25), 'runtime_p75': percentile(runtime, 75), 'runtime_p95': percentile(runtime, 95), 'batch_lat_p50': percentile(all_lat, 50), 'batch_lat_p95': percentile(all_lat, 95), 'batch_lat_p99': percentile(all_lat, 99), 'inf_p95': percentile(all_lat, 95), 'num_trials': len(trials), 'num_batches': trials[0]['num_batches'] if trials else 0, } return ( _stats(baseline_trials), _stats(stateless_trials), baseline_sample_summaries, stateless_sample_summaries, ) # --------------------------------------------------------------------------- # Single benchmark run # --------------------------------------------------------------------------- def run_benchmark( num_elements: int = 10000, min_length: int = 1, max_length: int = 500, max_batch_size: int = 32, max_batch_weight: int = 2000, distribution: str = 'pareto', seed: int = 42) -> dict[str, Any]: """Run baseline vs stateless comparison.""" generators = { 'pareto': lambda: generate_highly_skewed_data( num_elements, min_length, max_length, seed), 'lognormal': lambda: generate_lognormal_data( num_elements, 50, 0.8, min_length, max_length, seed), 'bimodal': lambda: generate_bimodal_data( num_elements, 20, 200, 0.7, min_length, max_length, seed), 'low_variance': lambda: generate_low_variance_data( num_elements, 100, 0.1, min_length, max_length, seed), } if distribution not in generators: raise ValueError(f"Unknown distribution: {distribution}") data = generators[distribution]() lengths = [len(s) for s in data] baseline_perf, stateless_perf, baseline_summaries, stateless_summaries = ( run_performance_benchmark(data, max_batch_size, max_batch_weight)) baseline_pad = compute_padding_stats(baseline_summaries) stateless_pad = compute_padding_stats(stateless_summaries) baseline_pad.update(baseline_perf) stateless_pad.update(stateless_perf) validation = validate_invariants( data, baseline_summaries, stateless_summaries) return { 'config': { 'num_elements': num_elements, 'max_batch_size': max_batch_size, 'max_batch_weight': max_batch_weight, 'distribution': distribution, 'runner': 'DirectRunner', }, 'data_stats': { 'min': min(lengths), 'max': max(lengths), 'mean': statistics.mean(lengths), 'median': statistics.median(lengths), 'std': statistics.stdev(lengths), }, 'baseline': baseline_pad, 'stateless': stateless_pad, 'validation': validation, } # --------------------------------------------------------------------------- # Printing # --------------------------------------------------------------------------- def _fmt_iqr(median, p25, p75, unit=''): return f"{median:.1f} [{p25:.1f}-{p75:.1f}]{unit}" def print_results(results: dict[str, Any]) -> None: cfg = results['config'] ds = results['data_stats'] bl = results['baseline'] st = results['stateless'] val = results['validation'] print("=" * 80) print( f"Distribution: {cfg['distribution']} | " f"N={cfg['num_elements']} | " f"runner={cfg['runner']} | " f"max_batch_size={cfg['max_batch_size']} | " f"max_batch_weight={cfg['max_batch_weight']}") print( f"Input lengths: min={ds['min']} max={ds['max']} " f"mean={ds['mean']:.1f} median={ds['median']:.0f} std={ds['std']:.1f}") print("-" * 80) def _arm(label, s): print(f"\n {label}:") print(f" Num batches: {s['num_batches']}") print(f" Padding ratio: {s['padding_ratio']:.2f}x") print(" ") print(" Throughput (Ktok/s):") med = s['tput_median'] / 1000 p25 = s['tput_p25'] / 1000 p75 = s['tput_p75'] / 1000 print(f" Median [IQR]: {med:.1f}" f" [{p25:.1f}-{p75:.1f}]") print(f" P95: {s['tput_p95']/1000:.1f}") print(" ") print(" E2E latency (ms):") print( f" Median [IQR]: {s['e2e_median']:.1f}" f" [{s['e2e_p25']:.1f}-{s['e2e_p75']:.1f}]") print(f" P95: {s['e2e_p95']:.1f}") print(" ") print(" Pipeline runtime (ms):") print( f" Median [IQR]:" f" {s['runtime_median']:.2f}" f" [{s['runtime_p25']:.2f}" f"-{s['runtime_p75']:.2f}]") print(f" P95: {s['runtime_p95']:.2f}") print(" ") print(" Batch latency (ms):") print(f" P50: {s['batch_lat_p50']:.1f}") print(f" P95: {s['batch_lat_p95']:.1f}") print(f" P99: {s['batch_lat_p99']:.1f}") _arm("Baseline (BatchElements)", bl) _arm("Stateless (SortAndBatchElements w/ weight-based splitting)", st) # Explicit arrows so direction is unambiguous. # down arrow = value decreased (good for latency/padding) # up arrow = value increased (good for throughput) def _delta_lower(base, new): """For metrics where lower is better (latency, padding).""" if base == 0: return 'N/A' pct = (base - new) / base * 100 arrow = '\u2193' if pct > 0 else '\u2191' return f"{arrow}{abs(pct):.1f}%" def _delta_higher(base, new): """For metrics where higher is better (throughput).""" if base == 0: return 'N/A' pct = (new - base) / base * 100 arrow = '\u2191' if pct > 0 else '\u2193' return f"{arrow}{abs(pct):.1f}%" print(f"\n {'_' * 76}") print(" DELTA (Baseline -> Stateless):") def _line(label, bv, sv, delta_fn, fmt='.1f', unit=''): d = delta_fn(bv, sv) print(f" {label}: {bv:{fmt}}{unit}" f" -> {sv:{fmt}}{unit} ({d})") bl_tmed = bl['tput_median'] / 1000 st_tmed = st['tput_median'] / 1000 bl_tp95 = bl['tput_p95'] / 1000 st_tp95 = st['tput_p95'] / 1000 _line( 'Padding ratio ', bl['padding_ratio'], st['padding_ratio'], _delta_lower, fmt='.2f', unit='x') _line('Throughput median', bl_tmed, st_tmed, _delta_higher, unit=' Ktok/s') _line('Throughput p95 ', bl_tp95, st_tp95, _delta_higher, unit=' Ktok/s') _line( 'E2E latency med ', bl['e2e_median'], st['e2e_median'], _delta_lower, unit=' ms') _line( 'E2E latency p95 ', bl['e2e_p95'], st['e2e_p95'], _delta_lower, unit=' ms') _line( 'Pipeline runtime ', bl['runtime_median'], st['runtime_median'], _delta_lower, unit=' ms') _line( 'Batch lat p95 ', bl['batch_lat_p95'], st['batch_lat_p95'], _delta_lower, unit=' ms') _line( 'Batch lat p99 ', bl['batch_lat_p99'], st['batch_lat_p99'], _delta_lower, unit=' ms') # Invariants e_ok = "Y" if val['elements_match'] else "X" t_ok = "Y" if val['tokens_match'] else "X" b_nb = val['baseline_num_batches'] s_nb = val['stateless_num_batches'] print( f"\n Invariants: elements {e_ok} tokens {t_ok}" f" (baseline {b_nb} -> stateless {s_nb}" f" batches)") print("=" * 80) # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main(): print("=" * 80) print("BASELINE (BatchElements) vs STATELESS (SortAndBatchElements)") print("=" * 80) print() print("Experiment design:") print(" A = Baseline : BatchElements with min=max=32") print(" B = Stateless : SortAndBatchElements with max_batch_weight=2000") print(" (sort within runner bundle, then split by weight)") print() print("Implementation notes:") print(" - Runs beam.Create(...) pipelines on DirectRunner") print(" - Materializes per-batch summaries through a temporary text sink") print(" - Uses runner-defined bundle boundaries rather than a synthetic") print(" bundle_size knob") print() print("Methodology:") print(" - N=20 trials, 3 warmup excluded") print(" - DirectRunner, in_memory mode, single worker") print(" - Percentiles: linear interpolation (= numpy default)") print(" - Same seed=42 for both arms") print(" - Inference model: latency = batch_size * (max_seq_len/50)^1.5 ms") print() dist = 'pareto' print(f"\nRunning: {dist}...") r = run_benchmark( num_elements=10000, max_batch_size=32, max_batch_weight=2000, distribution=dist, seed=42) print_results(r) if __name__ == '__main__': main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.7194436 apache_beam-2.74.0/apache_beam/testing/benchmarks/wordcount/0000755000175100017510000000000015206113160023475 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/wordcount/__init__.py0000644000175100017510000000142015206113032025601 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/benchmarks/wordcount/wordcount.py0000644000175100017510000000261015206113032026070 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging from apache_beam.examples import wordcount from apache_beam.testing.load_tests.dataflow_cost_benchmark import DataflowCostBenchmark class WordcountCostBenchmark(DataflowCostBenchmark): def __init__(self): super().__init__(pcollection='Format.out0') def test(self): extra_opts = {} extra_opts['output'] = self.pipeline.get_option('output_file') self.result = wordcount.run( self.pipeline.get_full_options_as_args(**extra_opts), save_main_session=False) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) WordcountCostBenchmark().run() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.7204437 apache_beam-2.74.0/apache_beam/testing/data/0000755000175100017510000000000015206113160020245 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/data/trigger_transcripts.yaml0000644000175100017510000002076415206113032025237 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # name: fixed_default window_fn: FixedWindows(10) # Python names/syntax, unless otherwise noted. trigger_fn: Default # Same. Empty () may be omitted. transcript: # Ordered list of events. - input: [1, 2, 3, 10, 11] # The elements are the timestamps. - watermark: 25 - expect: # Every expected output from the last action. - {window: [0, 9], values: [1, 2, 3], index: 0} - {window: [10, 19], values: [10, 11]} # Partial match on attributes OK. --- name: fixed_default_late_data window_fn: FixedWindows(10) trigger_fn: Default timestamp_combiner: OUTPUT_AT_EOW allowed_lateness: 100 transcript: - input: [1, 2, 3, 10, 11, 25] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2, 3], timestamp: 9, final: false} - {window: [10, 19], values: [10, 11], timestamp: 19} - {window: [20, 29], values: [25], timestamp: 29, late: false} - input: [7] - expect: - {window: [0, 9], values: [1, 2, 3, 7], timestamp: 9, late: true} --- name: fixed_drop_late_data_after_allowed_lateness window_fn: FixedWindows(10) trigger_fn: AfterWatermark(early=AfterCount(3), late=AfterCount(1)) timestamp_combiner: OUTPUT_AT_EOW allowed_lateness: 20 accumulation_mode: accumulating transcript: - input: [1, 2, 10, 11, 80, 81] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2], timestamp: 9} - {window: [10, 19], values: [10, 11], timestamp: 19} - {window: [80, 89], values: [80, 81], timestamp: 89, late: false} - input: [7, 8] # no output - input: [17, 18] # no output - input: [82] - expect: - {window: [80, 89], values: [80, 81, 82], timestamp: 89, late: true} --- name: timestamp_combiner_earliest window_fn: FixedWindows(10) trigger_fn: Default timestamp_combiner: OUTPUT_AT_EARLIEST transcript: - input: [1, 2, 3, 10, 11, 25] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2, 3], timestamp: 1} - {window: [10, 19], values: [10, 11], timestamp: 10} - {window: [20, 29], values: [25], timestamp: 25, late: false} --- name: timestamp_combiner_latest window_fn: FixedWindows(10) trigger_fn: Default timestamp_combiner: OUTPUT_AT_LATEST transcript: - input: [1, 2, 3, 10, 11, 25] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2, 3], timestamp: 3} - {window: [10, 19], values: [10, 11], timestamp: 11} - {window: [20, 29], values: [25], timestamp: 25, late: false} --- # Test that custom timestamping is not invoked. name: timestamp_combiner_custom_timestamping_eow window_fn: CustomTimestampingFixedWindowsWindowFn(10) trigger_fn: Default timestamp_combiner: OUTPUT_AT_EOW transcript: - input: [1, 2, 3, 10, 11, 25] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2, 3], timestamp: 9} - {window: [10, 19], values: [10, 11], timestamp: 19} - {window: [20, 29], values: [25], timestamp: 29, late: false} --- # Test that custom timestamping is not invoked. name: timestamp_combiner_custom_timestamping_earliest window_fn: CustomTimestampingFixedWindowsWindowFn(10) trigger_fn: Default timestamp_combiner: OUTPUT_AT_EARLIEST transcript: - input: [1, 2, 3, 10, 11, 25] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2, 3], timestamp: 1} - {window: [10, 19], values: [10, 11], timestamp: 10} - {window: [20, 29], values: [25], timestamp: 25, late: false} --- # Test that custom timestamping is in fact invoked. name: timestamp_combiner_custom_timestamping_earliest broken_on: - SwitchingDirectRunner # unsupported OUTPUT_AT_EARLIEST_TRANSFORMED window_fn: CustomTimestampingFixedWindowsWindowFn(10) trigger_fn: Default timestamp_combiner: OUTPUT_AT_EARLIEST_TRANSFORMED transcript: - input: [1, 2, 3, 10, 11, 25] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2, 3], timestamp: 101} - {window: [10, 19], values: [10, 11], timestamp: 110} - {window: [20, 29], values: [25], timestamp: 125, late: false} --- name: early_late_sessions broken_on: # Watermark regresses, causing what should be late data to not be late. - SwitchingDirectRunner window_fn: Sessions(10) trigger_fn: AfterWatermark(early=AfterCount(2), late=AfterCount(3)) allowed_lateness: 100 timestamp_combiner: OUTPUT_AT_EOW transcript: - input: [1, 2] - expect: - {window: [1, 11], values: [1, 2], timestamp: 11, early: true, index: 0} - input: [3] # no output - input: [4] - expect: - {window: [1, 13], values: [1, 2, 3, 4], timestamp: 13, early: true, index: 0} - input: [5] - watermark: 100 - expect: - {window: [1, 14], values:[1, 2, 3, 4, 5], timestamp: 14, index: 0, nonspeculative_index: 0} - input: [1] - input: [3, 4] - expect: - {window: [1, 14], values: [1, 1, 2, 3, 3, 4, 4, 5], timestamp: 14, final: false, index: 1, nonspeculative_index: 1} --- name: discarding_early_fixed window_fn: FixedWindows(10) trigger_fn: AfterWatermark(early=AfterCount(2)) timestamp_combiner: OUTPUT_AT_EOW accumulation_mode: discarding transcript: - input: [1, 2] - expect: - {window: [0, 9], values: [1, 2], timestamp: 9, early: true, index: 0} - input: [4] # no output - input: [14] # no output - input: [5] - expect: - {window: [0, 9], values: [4, 5], timestamp: 9, early: true, index: 1} - input: [18] - expect: - {window: [10, 19], values: [14, 18], timestamp: 19, early: true, index: 0} - input: [6] - watermark: 100 - expect: - {window: [0, 9], values: [6], timestamp: 9, early: false, late: false, final: true, index: 2, nonspeculative_index: 0} - {window: [10, 19], values: [], timestamp: 19, early: false, late: false, final: true, index: 1, nonspeculative_index: 0} --- name: garbage_collection broken_on: - SwitchingDirectRunner # claims pipeline stall window_fn: FixedWindows(10) trigger_fn: AfterCount(2) timestamp_combiner: OUTPUT_AT_EOW allowed_lateness: 10 accumulation_mode: discarding transcript: - input: [1, 2, 3, 10, 11, 25] - expect: - {window: [0, 9], timestamp: 9} - {window: [10, 19], timestamp: 19} - state: present: [[20, 29]] absent: [[0, 9]] tombstone: [[10, 19]] --- name: known_late_data_watermark broken_on: - SwitchingDirectRunner # bad timestamp window_fn: FixedWindows(10) trigger_fn: Default timestamp_combiner: OUTPUT_AT_EARLIEST transcript: - watermark: 5 - input: [2, 3, 7, 8] - watermark: 11 - expect: - {window: [0, 9], values: [2, 3, 7, 8], timestamp: 7} --- name: known_late_data_no_watermark_hold_possible broken_on: - SwitchingDirectRunner # bad timestamp window_fn: FixedWindows(10) trigger_fn: Default timestamp_combiner: OUTPUT_AT_EARLIEST transcript: - watermark: 8 - input: [2, 3, 7] - watermark: 11 - expect: - {window: [0, 9], values: [2, 3, 7], timestamp: 9} # These next examples test that bad/incomplete transcripts are rejected. --- name: bad_output error: Unmatched output windowfn: FixedWindows(10) transcript: - input: [1, 2, 3] - expect: - {window: [0, 9], values: [1, 2, 3]} # bad - watermark: 100 --- name: bad_expected_values error: Unmatched output window_fn: FixedWindows(10) transcript: - input: [1, 2, 3] - watermark: 100 - expect: - {window: [0, 9], values: [1, 2]} # bad values --- name: bad_expected_window error: Unmatched output window_fn: FixedWindows(10) transcript: - input: [1, 2, 3] - watermark: 100 - expect: - {window: [0, 19], values: [1, 2, 3]} # bad window --- name: missing_output error: Unexpected output window_fn: FixedWindows(10) transcript: - input: [1, 2, 3] - watermark: 100 # missing output - watermark: 200 --- name: missing_output_at_end error: Unexpected output window_fn: FixedWindows(10) transcript: - input: [1, 2, 3] - watermark: 100 # missing output ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/datatype_inference.py0000644000175100017510000000721515206113032023542 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import array from collections import OrderedDict import numpy as np from fastavro import parse_schema from apache_beam.typehints import trivial_inference from apache_beam.typehints import typehints try: import pyarrow as pa except ImportError: pa = None def infer_element_type(elements): """For internal use only; no backwards-compatibility guarantees. Infer a Beam type for a list of elements. Args: elements (List[Any]): A list of elements for which the type should be inferred. Returns: A Beam type encompassing all elements. """ element_type = typehints.Union[[ trivial_inference.instance_to_type(e) for e in elements ]] return element_type def infer_typehints_schema(data): """For internal use only; no backwards-compatibility guarantees. Infer Beam types for tabular data. Args: data (List[dict]): A list of dictionaries representing rows in a table. Returns: An OrderedDict mapping column names to Beam types. """ column_data = OrderedDict() for row in data: for key, value in row.items(): column_data.setdefault(key, []).append(value) column_types = OrderedDict([(key, infer_element_type(values)) for key, values in column_data.items()]) return column_types def infer_avro_schema(data): """For internal use only; no backwards-compatibility guarantees. Infer avro schema for tabular data. Args: data (List[dict]): A list of dictionaries representing rows in a table. Returns: An avro schema object. """ _typehint_to_avro_type = { type(None): "null", int: "int", float: "double", str: "string", bytes: "bytes", np.ndarray: "bytes", array.array: "bytes", } def typehint_to_avro_type(value): if isinstance(value, typehints.UnionConstraint): return sorted( typehint_to_avro_type(union_type) for union_type in value.union_types) else: return _typehint_to_avro_type[value] column_types = infer_typehints_schema(data) avro_fields = [{ "name": str(key), "type": typehint_to_avro_type(value) } for key, value in column_types.items()] schema_dict = { "namespace": "example.avro", "name": "User", "type": "record", "fields": avro_fields } return parse_schema(schema_dict) def infer_pyarrow_schema(data): """For internal use only; no backwards-compatibility guarantees. Infer PyArrow schema for tabular data. Args: data (List[dict]): A list of dictionaries representing rows in a table. Returns: A PyArrow schema object. """ column_data = OrderedDict() for row in data: for key, value in row.items(): column_data.setdefault(key, []).append(value) column_types = OrderedDict([(key, pa.array(value).type) for key, value in column_data.items()]) return pa.schema(list(column_types.items())) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/datatype_inference_test.py0000644000175100017510000001402415206113032024575 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import unittest from collections import OrderedDict import numpy as np from parameterized import parameterized from apache_beam.testing import datatype_inference from apache_beam.typehints import typehints try: import pyarrow as pa except ImportError: pa = None TEST_DATA = [ { "name": "empty", "data": [], "type_schema": OrderedDict([]), "pyarrow_schema": pa.schema([]) if pa is not None else None, "avro_schema": { "namespace": "example.avro", "name": "User", "type": "record", "fields": [], }, }, { "name": "main", "data": [ OrderedDict([ ("a", 1), ("b", 0.12345), ("c", "Hello World!!"), ("d", np.array([1, 2, 3], dtype=np.int64)), ("e", b"some bytes"), ]), OrderedDict([ ("a", -5), ("b", 1234.567), ("e", b"more bytes"), ]), OrderedDict([ ("a", 100000), ("c", "XoXoX"), ("d", np.array([4, 5, 6], dtype=np.int64)), ("e", b""), ]), ], "type_schema": OrderedDict([ ("a", int), ("b", float), ("c", str), ("d", np.ndarray), ("e", bytes), ]), "pyarrow_schema": pa.schema([ ("a", pa.int64()), ("b", pa.float64()), ("c", pa.string()), ("d", pa.list_(pa.int64())), ("e", pa.binary()), ]) if pa is not None else None, "avro_schema": { "namespace": "example.avro", "name": "User", "type": "record", "fields": [ { "name": "a", "type": "int" }, { "name": "b", "type": "double" }, { "name": "c", "type": "string" }, { "name": "d", "type": "bytes" }, { "name": "e", "type": "bytes" }, ], }, }, ] def nullify_data_and_schemas(test_data): """Add a row with all columns set to None and adjust the schemas accordingly. """ def nullify_avro_schema(schema): """Add a 'null' type to every field.""" schema = schema.copy() new_fields = [] for field in schema["fields"]: if isinstance(field["type"], str): new_fields.append({ "name": field["name"], "type": sorted([field["type"], "null"]) }) else: new_fields.append({ "name": field["name"], "type": sorted(field["type"] + ["null"]) }) schema["fields"] = new_fields return schema def get_collumns_in_order(test_data): """Get a list of columns while trying to maintain original order. .. note:: Columns which do not apear until later rows are added to the end, even if they preceed some columns which have already been added. """ _seen = set() columns = [ c for test_case in test_data for row in test_case["data"] for c in row if c not in _seen and not _seen.add(c) ] return columns nullified_test_data = [] columns = get_collumns_in_order(test_data) for test_case in test_data: if not test_case["data"]: continue test_case = test_case.copy() test_case["name"] = test_case["name"] + "_nullified" test_case["data"] = test_case["data"] + [ OrderedDict([(c, None) for c in columns]) ] test_case["type_schema"] = OrderedDict([ (k, typehints.Union[v, type(None)]) for k, v in test_case["type_schema"].items() ]) test_case["avro_schema"] = nullify_avro_schema(test_case["avro_schema"]) nullified_test_data.append(test_case) return nullified_test_data TEST_DATA += nullify_data_and_schemas(TEST_DATA) class DatatypeInferenceTest(unittest.TestCase): @parameterized.expand([(d["name"], d["data"], d["type_schema"]) for d in TEST_DATA]) def test_infer_typehints_schema(self, _, data, schema): typehints_schema = datatype_inference.infer_typehints_schema(data) self.assertEqual(typehints_schema, schema) @parameterized.expand([(d["name"], d["data"], d["pyarrow_schema"]) for d in TEST_DATA]) @unittest.skipIf(pa is None, "PyArrow is not installed") def test_infer_pyarrow_schema(self, _, data, schema): pyarrow_schema = datatype_inference.infer_pyarrow_schema(data) self.assertEqual(pyarrow_schema, schema) @parameterized.expand([(d["name"], d["data"], d["avro_schema"]) for d in TEST_DATA]) def test_infer_fastavro_schema(self, _, data, schema): from fastavro import parse_schema schema = parse_schema(schema) avro_schema = datatype_inference.infer_avro_schema(data) fields1 = avro_schema.pop("fields") fields2 = schema.pop("fields") self.assertDictEqual(avro_schema, schema) for field1, field2 in zip(fields1, fields2): self.assertDictEqual(field1, field2) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/extra_assertions.py0000644000175100017510000000351115206113032023301 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import numpy as np class ExtraAssertionsMixin(object): def assertUnhashableCountEqual(self, data1, data2): """Assert that two containers have the same items, with special treatment for numpy arrays. """ try: self.assertCountEqual(data1, data2) except (TypeError, ValueError): data1 = [self._to_hashable(d) for d in data1] data2 = [self._to_hashable(d) for d in data2] self.assertCountEqual(data1, data2) def _to_hashable(self, element): try: hash(element) return element except TypeError: pass if isinstance(element, list): return tuple(self._to_hashable(e) for e in element) if isinstance(element, dict): hashable_elements = [] for key, value in sorted(element.items(), key=lambda t: hash(t[0])): hashable_elements.append((key, self._to_hashable(value))) return tuple(hashable_elements) if isinstance(element, np.ndarray): return element.tobytes() raise AssertionError("Encountered unhashable element: {}.".format(element)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/extra_assertions_test.py0000644000175100017510000000420315206113032024337 0ustar00runnerrunner# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import unittest import numpy as np from apache_beam.testing.extra_assertions import ExtraAssertionsMixin class ExtraAssertionsMixinTest(ExtraAssertionsMixin, unittest.TestCase): def test_assert_array_count_equal_strings(self): data1 = ["±♠Ωℑ", "hello", "world"] data2 = ["hello", "±♠Ωℑ", "world"] self.assertUnhashableCountEqual(data1, data2) def test_assert_array_count_equal_mixed(self): data1 = [ { 'a': 1, 123: 1.234 }, ['d', 1], "±♠Ωℑ", np.zeros((3, 6)), (1, 2, 3, 'b'), 'def', 100, 'abc', ('a', 'b', 'c'), None, ] data2 = [ { 123: 1.234, 'a': 1 }, ('a', 'b', 'c'), ['d', 1], None, 'abc', 'def', "±♠Ωℑ", 100, (1, 2, 3, 'b'), np.zeros((3, 6)), ] self.assertUnhashableCountEqual(data1, data2) self.assertUnhashableCountEqual(data1 * 2, data2 * 2) def test_assert_not_equal(self): data1 = [{'a': 123, 'b': 321}, [1, 2, 3]] data2 = [{'a': 123, 'c': 321}, [1, 2, 3]] with self.assertRaises(AssertionError): self.assertUnhashableCountEqual(data1, data2) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/fast_test_utils.pxd0000644000175100017510000000237115206113032023266 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # cimport libc.stdint cdef class LCGenerator(object): cdef libc.stdint.uint64_t _a cdef libc.stdint.uint64_t _c cdef libc.stdint.int32_t _bits cdef libc.stdint.uint64_t _mask cdef libc.stdint.uint64_t _seed cpdef void seed(self, libc.stdint.uint64_t seed) cpdef void seed_jdk(self, libc.stdint.uint64_t seed) cpdef libc.stdint.int32_t next_int(self) cpdef libc.stdint.uint32_t next_uint(self) cpdef bytes randbytes(self, int length) cpdef double random_sample(self)././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/fast_test_utils.pyx0000644000175100017510000000412115206113032023306 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # cimport libc.stdint cimport libc.stdlib cimport libc.string from cpython.bytes cimport PyBytes_FromStringAndSize cdef class LCGenerator(object): def __cinit__(self): self._a = 0x5DEECE66Dull self._c = 0xBull self._bits = 48 self._mask = (1ull << self._bits) - 1 self.seed(0) cpdef void seed(self, libc.stdint.uint64_t seed): self._seed = (seed * self._a + self._c) & self._mask cpdef void seed_jdk(self, libc.stdint.uint64_t seed): self._seed = (seed ^ self._a) & self._mask cpdef libc.stdint.int32_t next_int(self): self.seed(self._seed) return (self._seed >> (self._bits - 32)) cpdef libc.stdint.uint32_t next_uint(self): self.seed(self._seed) return (self._seed >> (self._bits - 32)) cpdef bytes randbytes(self, int length): cdef libc.stdint.int32_t ints = (length + 3) // 4 cdef char* data = libc.stdlib.malloc(ints * 4) cdef libc.stdint.uint32_t value cdef libc.stdint.int32_t i for i in range(0, ints, 1): value = self.next_uint() libc.string.memcpy((data + i * 4), (& value), 4) retval = PyBytes_FromStringAndSize(data, length) libc.stdlib.free(data) return retval cpdef double random_sample(self): return (self.next_uint() >> 8) / (1 << 24) ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.7244442 apache_beam-2.74.0/apache_beam/testing/load_tests/0000755000175100017510000000000015206113160021475 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/load_tests/__init__.py0000644000175100017510000000142015206113032023601 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/load_tests/co_group_by_key_test.py0000644000175100017510000001126715206113032026272 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ This is CoGroupByKey load test with Synthetic Source. Besides of the standard input options there are additional options: * project (optional) - the gcp project in case of saving metrics in Big Query (in case of Dataflow Runner it is required to specify project of runner), * publish_to_big_query - if metrics should be published in big query, * metrics_namespace (optional) - name of BigQuery dataset where metrics will be stored, * metrics_table (optional) - name of BigQuery table where metrics will be stored, * input_options - options for Synthetic Sources, * co_input_options - options for Synthetic Sources, * iterations - number of reiterations over per-key-grouped values to perform (default: 1). Example test run: python -m apache_beam.testing.load_tests.co_group_by_key_test \ --test-pipeline-options=" --project=big-query-project --region=... --publish_to_big_query=true --metrics_dataset=python_load_tests --metrics_table=co_gbk --iterations=1 --input_options='{ \"num_records\": 1000, \"key_size\": 5, \"value_size\": 15 }' --co_input_options='{ \"num_records\": 1000, \"key_size\": 5, \"value_size\": 15 }'" or: ./gradlew -PloadTest.args=" --publish_to_big_query=true --project=... --region=... --metrics_dataset=python_load_tests --metrics_table=co_gbk --iterations=1 --input_options='{ \"num_records\": 1, \"key_size\": 1, \"value_size\": 1}' --co_input_options='{ \"num_records\": 1, \"key_size\": 1, \"value_size\": 1}' --runner=DirectRunner" \ -PloadTest.mainClass=apache_beam.testing.load_tests.co_group_by_key_test \ -Prunner=DirectRunner :sdks:python:apache_beam:testing:load_tests:run """ # pytype: skip-file import json import logging import apache_beam as beam from apache_beam.testing import synthetic_pipeline from apache_beam.testing.load_tests.load_test import LoadTest from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime class CoGroupByKeyTest(LoadTest): INPUT_TAG = 'pc1' CO_INPUT_TAG = 'pc2' def __init__(self): super().__init__() self.co_input_options = json.loads( self.pipeline.get_option('co_input_options')) self.iterations = self.get_option_or_default('iterations', 1) class _UngroupAndReiterate(beam.DoFn): def __init__(self, input_tag, co_input_tag): self.input_tag = input_tag self.co_input_tag = co_input_tag def process(self, element, iterations): values = element[1] inputs = values.get(self.input_tag) co_inputs = values.get(self.co_input_tag) for i in range(iterations): for value in inputs: if i == iterations - 1: yield value for value in co_inputs: if i == iterations - 1: yield value def test(self): pc1 = ( self.pipeline | 'Read ' + self.INPUT_TAG >> beam.io.Read( synthetic_pipeline.SyntheticSource( self.parse_synthetic_source_options())) | 'Measure time: Start pc1' >> beam.ParDo( MeasureTime(self.metrics_namespace))) pc2 = ( self.pipeline | 'Read ' + self.CO_INPUT_TAG >> beam.io.Read( synthetic_pipeline.SyntheticSource( self.parse_synthetic_source_options(self.co_input_options))) | 'Measure time: Start pc2' >> beam.ParDo( MeasureTime(self.metrics_namespace))) # pylint: disable=expression-not-assigned ({ self.INPUT_TAG: pc1, self.CO_INPUT_TAG: pc2 } | 'CoGroupByKey ' >> beam.CoGroupByKey() | 'Consume Joined Collections' >> beam.ParDo( self._UngroupAndReiterate(self.INPUT_TAG, self.CO_INPUT_TAG), self.iterations) | 'Measure time: End' >> beam.ParDo(MeasureTime(self.metrics_namespace))) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) CoGroupByKeyTest().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/load_tests/combine_test.py0000644000175100017510000001045415206113032024524 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ This is Combine load test with Synthetic Source. Besides of the standard input options there are additional options: * fanout (optional) - number of GBK operations to run in parallel * top_count - an arguments passed to the Top combiner. * project (optional) - the gcp project in case of saving metrics in Big Query (in case of Dataflow Runner it is required to specify project of runner), * publish_to_big_query - if metrics should be published in big query, * metrics_namespace (optional) - name of BigQuery dataset where metrics will be stored, * metrics_table (optional) - name of BigQuery table where metrics will be stored, * input_options - options for Synthetic Sources. Example test run: python -m apache_beam.testing.load_tests.combine_test \ --test-pipeline-options=" --project=big-query-project --region=... --publish_to_big_query=true --metrics_dataset=python_load_tests --metrics_table=combine --fanout=1 --top_count=1000 --input_options='{ \"num_records\": 300, \"key_size\": 5, \"value_size\": 15 }'" or: ./gradlew -PloadTest.args=" --publish_to_big_query=true --project=... --region=... --metrics_dataset=python_load_tests --metrics_table=combine --top_count=1000 --fanout=1 --input_options='{ \"num_records\": 1, \"key_size\": 1, \"value_size\": 1}' --runner=DirectRunner" \ -PloadTest.mainClass=apache_beam.testing.load_tests.combine_test \ -Prunner=DirectRunner :sdks:python:apache_beam:testing:load_tests:run """ # pytype: skip-file import logging import sys import apache_beam as beam from apache_beam.testing.load_tests.load_test import LoadTest from apache_beam.testing.load_tests.load_test_metrics_utils import AssignTimestamps from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime from apache_beam.testing.synthetic_pipeline import StatefulLoadGenerator from apache_beam.testing.synthetic_pipeline import SyntheticSource from apache_beam.transforms.combiners import window class CombineTest(LoadTest): def __init__(self): super().__init__() self.fanout = self.get_option_or_default('fanout', 1) try: self.top_count = int(self.pipeline.get_option('top_count')) except (TypeError, ValueError): logging.error( 'You should set \"--top_count\" option to use TOP ' 'combiners') sys.exit(1) class _GetElement(beam.DoFn): def process(self, element): yield element def test(self): if self.get_option_or_default('use_stateful_load_generator', False): source = ( self.pipeline | 'LoadGenerator' >> StatefulLoadGenerator(self.input_options) | beam.ParDo(AssignTimestamps()) | beam.WindowInto(window.FixedWindows(20))) else: source = ( self.pipeline | 'Read synthetic' >> beam.io.Read( SyntheticSource(self.parse_synthetic_source_options()))) pc = ( source | 'Measure time: Start' >> beam.ParDo( MeasureTime(self.metrics_namespace))) for branch in range(self.fanout): ( # pylint: disable=expression-not-assigned pc | 'Combine with Top %i' % branch >> beam.CombineGlobally( beam.combiners.TopCombineFn(self.top_count)).without_defaults() | 'Consume %i' % branch >> beam.ParDo(self._GetElement()) | 'Measure time: End %i' % branch >> beam.ParDo( MeasureTime(self.metrics_namespace))) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) CombineTest().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/load_tests/dataflow_cost_benchmark.py0000644000175100017510000002743715206113032026725 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # pytype: skip-file import logging import re import time from datetime import datetime from typing import Any from typing import Optional from google.cloud import monitoring_v3 from google.protobuf.duration_pb2 import Duration import apache_beam.testing.load_tests.dataflow_cost_consts as costs from apache_beam.metrics.execution import MetricResult from apache_beam.runners.dataflow.dataflow_runner import DataflowPipelineResult from apache_beam.runners.dataflow.internal.apiclient import DataflowApplicationClient from apache_beam.runners.runner import PipelineState from apache_beam.testing.load_tests.load_test import LoadTest class DataflowCostBenchmark(LoadTest): """Base class for Dataflow performance tests which export metrics to external databases: BigQuery or/and InfluxDB. Calculates the expected cost for running the job on Dataflow in region us-central1. Refer to :class:`~apache_beam.testing.load_tests.LoadTestOptions` for more information on the required pipeline options. If using InfluxDB with Basic HTTP authentication enabled, provide the following environment options: `INFLUXDB_USER` and `INFLUXDB_USER_PASSWORD`. If the hardware configuration for the job includes use of a GPU, please specify the version in use with the Accelerator enumeration. This is used to calculate the cost of the job later, as different accelerators have different billing rates per hour of use. """ WORKER_START_PATTERN = re.compile( r'All workers have finished the startup processes and ' r'began to receive work requests') WORKER_STOP_PATTERN = re.compile(r'Stopping worker pool') def __init__( self, metrics_namespace: Optional[str] = None, is_streaming: bool = False, gpu: Optional[costs.Accelerator] = None, pcollection: str = 'ProcessOutput.out0', subscription: Optional[str] = None): self.is_streaming = is_streaming self.gpu = gpu self.pcollection = pcollection self.subscription = subscription super().__init__(metrics_namespace=metrics_namespace) self.dataflow_client = DataflowApplicationClient( self.pipeline.get_pipeline_options()) self.monitoring_client = monitoring_v3.MetricServiceClient() def run(self) -> None: try: self.test() if not hasattr(self, 'result'): self.result = self.pipeline.run() state = self.result.wait_until_finish(duration=self.timeout_ms) assert state != PipelineState.FAILED logging.info( 'Pipeline complete, sleeping for 4 minutes to allow resource ' 'metrics to populate.') time.sleep(240) self.extra_metrics = self._retrieve_cost_metrics(self.result) additional_metrics = self._get_additional_metrics(self.result) self.extra_metrics.update(additional_metrics) logging.info(self.extra_metrics) self._metrics_monitor.publish_metrics(self.result, self.extra_metrics) except Exception: raise finally: self.cleanup() def _retrieve_cost_metrics(self, result: DataflowPipelineResult) -> dict[str, Any]: """Calculates estimated cost based on pipeline resource usage.""" job_id = result.job_id() metrics = result.metrics().all_metrics(job_id) metrics_dict = self._process_metrics_list(metrics) cost = 0.0 if self.is_streaming: cost += metrics_dict.get( "TotalVcpuTime", 0.0) / 3600 * costs.VCPU_PER_HR_STREAMING cost += metrics_dict.get( "TotalMemoryUsage", 0.0) / 1000 / 3600 * costs.MEM_PER_GB_HR_STREAMING cost += metrics_dict.get( "TotalStreamingDataProcessed", 0.0) * costs.SHUFFLE_PER_GB_STREAMING else: cost += metrics_dict.get( "TotalVcpuTime", 0.0) / 3600 * costs.VCPU_PER_HR_BATCH cost += metrics_dict.get( "TotalMemoryUsage", 0.0) / 1000 / 3600 * costs.MEM_PER_GB_HR_BATCH cost += metrics_dict.get( "TotalStreamingDataProcessed", 0.0) * costs.SHUFFLE_PER_GB_BATCH if self.gpu: rate = costs.ACCELERATOR_TO_COST[self.gpu] cost += metrics_dict.get("TotalGpuTime", 0.0) / 3600 * rate cost += metrics_dict.get("TotalPdUsage", 0.0) / 3600 * costs.PD_PER_GB_HR cost += metrics_dict.get( "TotalSsdUsage", 0.0) / 3600 * costs.PD_SSD_PER_GB_HR metrics_dict["EstimatedCost"] = cost return metrics_dict def _process_metrics_list(self, metrics: list[MetricResult]) -> dict[str, Any]: system_metrics = {} for entry in metrics: metric_key = entry.key metric = metric_key.metric if metric_key.step == '' and metric.namespace == 'dataflow/v1b3': system_metrics[metric.name] = entry.committed or 0.0 return system_metrics def _get_worker_time_interval( self, job_id: str) -> tuple[Optional[str], Optional[str]]: """Extracts worker start and stop times from job messages.""" start_time, end_time = None, None page_token = None all_messages = [] last_message_time = None while True: messages, page_token = self.dataflow_client.list_messages( job_id=job_id, start_time=None, end_time=None, page_token=page_token, minimum_importance='JOB_MESSAGE_DEBUG') for message in messages: text = message.messageText if getattr(message, 'time', None): last_message_time = message.time if text: all_messages.append(text) if self.WORKER_START_PATTERN.search(text): start_time = message.time logging.info('Matched WORKER_START_PATTERN: %r', text) if self.WORKER_STOP_PATTERN.search(text): end_time = message.time logging.info('Matched WORKER_STOP_PATTERN: %r', text) if not page_token or (start_time and end_time): break if start_time and not end_time and self.is_streaming and last_message_time: end_time = last_message_time logging.info( 'Using last job message time as end_time for streaming job: %s', end_time) if not start_time or not end_time: logging.warning( 'Could not determine worker time interval. ' 'start_time=%s, end_time=%s, total messages=%d', start_time, end_time, len(all_messages)) return start_time, end_time def _get_throughput_metrics( self, project: str, job_id: str, start_time: str, end_time: str, pcollection_name: Optional[str] = None, ) -> dict[str, float]: """Query Cloud Monitoring for per-PCollection throughput.""" name = ( pcollection_name if pcollection_name is not None else self.pcollection) interval = monitoring_v3.TimeInterval( start_time=start_time, end_time=end_time) aggregation = monitoring_v3.Aggregation( alignment_period=Duration(seconds=60), per_series_aligner=monitoring_v3.Aggregation.Aligner.ALIGN_MEAN) requests = { "Bytes": monitoring_v3.ListTimeSeriesRequest( name=f"projects/{project}", filter=f'metric.type=' f'"dataflow.googleapis.com/job/estimated_byte_count" ' f'AND metric.labels.job_id=' f'"{job_id}" AND metric.labels.pcollection="{name}"', interval=interval, aggregation=aggregation), "Elements": monitoring_v3.ListTimeSeriesRequest( name=f"projects/{project}", filter=f'metric.type="dataflow.googleapis.com/job/element_count" ' f'AND metric.labels.job_id="{job_id}" ' f'AND metric.labels.pcollection="{name}"', interval=interval, aggregation=aggregation) } metrics = {} for key, req in requests.items(): time_series = self.monitoring_client.list_time_series(request=req) values = [ point.value.double_value for series in time_series for point in series.points ] metrics[f"AvgThroughput{key}"] = sum(values) / len( values) if values else 0.0 return metrics def _get_streaming_throughput_metrics( self, project: str, start_time: str, end_time: str) -> dict[str, float]: if not self.subscription: return {'AvgThroughputBytes': 0.0, 'AvgThroughputElements': 0.0} sub_parts = self.subscription.split('/') subscription_id = sub_parts[-1] if sub_parts else self.subscription interval = monitoring_v3.TimeInterval( start_time=start_time, end_time=end_time) aggregation = monitoring_v3.Aggregation( alignment_period=Duration(seconds=60), per_series_aligner=monitoring_v3.Aggregation.Aligner.ALIGN_RATE) requests = { "Bytes": monitoring_v3.ListTimeSeriesRequest( name=f"projects/{project}", filter=f'metric.type=' f'"pubsub.googleapis.com/subscription/byte_cost" ' f'AND resource.labels.subscription_id="{subscription_id}"', interval=interval, aggregation=aggregation), "Elements": monitoring_v3.ListTimeSeriesRequest( name=f"projects/{project}", filter=f'metric.type=' f'"pubsub.googleapis.com/subscription/sent_message_count" ' f'AND resource.labels.subscription_id="{subscription_id}"', interval=interval, aggregation=aggregation), } metrics = {} for key, req in requests.items(): time_series = list(self.monitoring_client.list_time_series(request=req)) values = [ point.value.double_value for series in time_series for point in series.points ] avg_rate = sum(values) / len(values) if values else 0.0 metrics[f"AvgThroughput{key}"] = avg_rate return metrics def _get_job_runtime(self, start_time: str, end_time: str) -> float: """Calculates the job runtime duration in seconds.""" start_dt = datetime.fromisoformat(start_time[:-1]) end_dt = datetime.fromisoformat(end_time[:-1]) return (end_dt - start_dt).total_seconds() def _get_additional_metrics(self, result: DataflowPipelineResult) -> dict[str, Any]: job_id = result.job_id() job = self.dataflow_client.get_job(job_id) project = job.projectId start_time, end_time = self._get_worker_time_interval(job_id) if not start_time or not end_time: logging.warning('Could not find valid worker start/end times.') return {} runtime_seconds = self._get_job_runtime(start_time, end_time) if self.is_streaming: throughput_metrics = self._get_streaming_throughput_metrics( project, start_time, end_time) else: throughput_metrics = self._get_throughput_metrics( project, job_id, start_time, end_time) if (throughput_metrics.get('AvgThroughputBytes', 0) == 0 and throughput_metrics.get('AvgThroughputElements', 0) == 0): logging.warning( 'No throughput data for PCollection "%s". Check Dataflow job %s ' 'graph for actual PCollection names (Runner v2 may use different ' 'naming).', self.pcollection, job_id) return { **throughput_metrics, "JobRuntimeSeconds": runtime_seconds, } ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/load_tests/dataflow_cost_consts.py0000644000175100017510000000356115206113032026274 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # These values are Dataflow costs for running jobs in us-central1. # The cost values are found at https://cloud.google.com/dataflow/pricing from enum import Enum VCPU_PER_HR_BATCH = 0.056 VCPU_PER_HR_STREAMING = 0.069 MEM_PER_GB_HR_BATCH = 0.003557 MEM_PER_GB_HR_STREAMING = 0.0035557 PD_PER_GB_HR = 0.000054 PD_SSD_PER_GB_HR = 0.000298 SHUFFLE_PER_GB_BATCH = 0.011 SHUFFLE_PER_GB_STREAMING = 0.018 # GPU Resource Pricing P100_PER_GPU_PER_HOUR = 1.752 V100_PER_GPU_PER_HOUR = 2.976 T4_PER_GPU_PER_HOUR = 0.42 P4_PER_GPU_PER_HOUR = 0.72 L4_PER_GPU_PER_HOUR = 0.672 A100_40GB_PER_GPU_PER_HOUR = 3.72 A100_80GB_PER_GPU_PER_HOUR = 4.7137 class Accelerator(Enum): P100 = 1 V100 = 2 T4 = 3 P4 = 4 L4 = 5 A100_40GB = 6 A100_80GB = 7 ACCELERATOR_TO_COST: dict[Accelerator, float] = { Accelerator.P100: P100_PER_GPU_PER_HOUR, Accelerator.V100: V100_PER_GPU_PER_HOUR, Accelerator.T4: T4_PER_GPU_PER_HOUR, Accelerator.P4: P4_PER_GPU_PER_HOUR, Accelerator.L4: L4_PER_GPU_PER_HOUR, Accelerator.A100_40GB: A100_40GB_PER_GPU_PER_HOUR, Accelerator.A100_80GB: A100_80GB_PER_GPU_PER_HOUR, } ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/load_tests/group_by_key_test.py0000644000175100017510000001003015206113032025574 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ This is GroupByKey load test with Synthetic Source. Besides of the standard input options there are additional options: * fanout (optional) - number of GBK operations to run in parallel * iterations (optional) - number of reiteraations over per-key-grouped values to perform * project (optional) - the gcp project in case of saving metrics in Big Query (in case of Dataflow Runner it is required to specify project of runner), * publish_to_big_query - if metrics should be published in big query, * metrics_namespace (optional) - name of BigQuery dataset where metrics will be stored, * metrics_table (optional) - name of BigQuery table where metrics will be stored, * input_options - options for Synthetic Sources. Example test run: python -m apache_beam.testing.load_tests.group_by_key_test \ --test-pipeline-options=" --project=big-query-project --region=... --publish_to_big_query=true --metrics_dataset=python_load_tests --metrics_table=gbk --fanout=1 --iterations=1 --input_options='{ \"num_records\": 300, \"key_size\": 5, \"value_size\": 15 }'" or: ./gradlew -PloadTest.args=" --publish_to_big_query=true --project=... --region=... --metrics_dataset=python_load_tests --metrics_table=gbk --fanout=1 --iterations=1 --input_options='{ \"num_records\": 1, \"key_size\": 1, \"value_size\": 1}' --runner=DirectRunner" \ -PloadTest.mainClass=apache_beam.testing.load_tests.group_by_key_test \ -Prunner=DirectRunner :sdks:python:apache_beam:testing:load_tests:run """ # pytype: skip-file import logging import apache_beam as beam from apache_beam.testing.load_tests.load_test import LoadTest from apache_beam.testing.load_tests.load_test_metrics_utils import AssignTimestamps from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureLatency from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime from apache_beam.testing.synthetic_pipeline import SyntheticSource class GroupByKeyTest(LoadTest): def __init__(self): super().__init__() self.fanout = self.get_option_or_default('fanout', 1) self.iterations = self.get_option_or_default('iterations', 1) @staticmethod def ungroup_and_reiterate(element, iterations): key, value = element for i in range(iterations): for v in value: if i == iterations - 1: return key, v def test(self): pc = ( self.pipeline | beam.io.Read(SyntheticSource(self.parse_synthetic_source_options())) | 'Measure time: Start' >> beam.ParDo( MeasureTime(self.metrics_namespace)) | 'Assign timestamps' >> beam.ParDo(AssignTimestamps())) for branch in range(self.fanout): ( # pylint: disable=expression-not-assigned pc | 'GroupByKey %i' % branch >> beam.GroupByKey() | 'Ungroup %i' % branch >> beam.Map(self.ungroup_and_reiterate, self.iterations) | 'Measure latency %i' % branch >> beam.ParDo( MeasureLatency(self.metrics_namespace)) | 'Measure time: End %i' % branch >> beam.ParDo( MeasureTime(self.metrics_namespace))) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) GroupByKeyTest().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/load_tests/load_test.py0000644000175100017510000001606315206113032024031 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # pytype: skip-file import argparse import json import logging import os import sys from apache_beam.metrics import MetricsFilter from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.runners.runner import PipelineState from apache_beam.testing.load_tests.load_test_metrics_utils import InfluxDBMetricsPublisherOptions from apache_beam.testing.load_tests.load_test_metrics_utils import MetricsReader from apache_beam.testing.test_pipeline import TestPipeline class LoadTestOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--publish_to_big_query', type=cls._str_to_boolean, help='Publishes pipeline metrics to BigQuery table.') parser.add_argument( '--metrics_dataset', help='A BigQuery dataset where metrics should be' 'written.') parser.add_argument( '--metrics_table', help='A BigQuery table where metrics should be ' 'written.') parser.add_argument( '--influx_measurement', help='An InfluxDB measurement where metrics should be published to. ' 'Measurement can be thought of as a SQL table. If empty, reporting to ' 'InfluxDB will be disabled.') parser.add_argument( '--influx_db_name', help='InfluxDB database name. If empty, reporting to InfluxDB will be ' 'disabled.') parser.add_argument( '--influx_hostname', help='Hostname to connect to InfluxDB. Defaults to ' '"http://localhost:8086".', default='http://localhost:8086') parser.add_argument( '--input_options', type=json.loads, help='Input specification of SyntheticSource.') parser.add_argument( '--timeout_ms', type=int, default=0, help='Waiting time for the completion of the pipeline in milliseconds.' 'Defaults to waiting forever.') @staticmethod def _str_to_boolean(value): try: return bool(['false', 'true'].index(value.lower())) except ValueError: raise argparse.ArgumentTypeError( '"true" or "false" expected, got "{}" ' 'instead.'.format(value)) class LoadTest(object): """Base class for all integration and performance tests which export metrics to external databases: BigQuery or/and InfluxDB. Refer to :class:`~apache_beam.testing.load_tests.LoadTestOptions` for more information on the required pipeline options. If using InfluxDB with Basic HTTP authentication enabled, provide the following environment options: `INFLUXDB_USER` and `INFLUXDB_USER_PASSWORD`. """ def __init__(self, metrics_namespace=None): options_class = getattr(self.__class__, 'options_class', None) if options_class is not None: options_list = TestPipeline.get_options_list() options = options_class(options_list) self.pipeline = TestPipeline( options=options, is_integration_test=True, blocking=False) else: self.pipeline = TestPipeline(is_integration_test=True, blocking=False) assert not self.pipeline.blocking options = self.pipeline.get_pipeline_options().view_as(LoadTestOptions) self.timeout_ms = options.timeout_ms self.input_options = options.input_options self.extra_metrics = {} if metrics_namespace: self.metrics_namespace = metrics_namespace else: self.metrics_namespace = options.metrics_table \ if options.metrics_table else 'default' publish_to_bq = options.publish_to_big_query if publish_to_bq is None: logging.info( 'Missing --publish_to_big_query option. Metrics will not ' 'be published to BigQuery.') if options.input_options is None: logging.error('--input_options argument is required.') sys.exit(1) gcloud_options = self.pipeline.get_pipeline_options().view_as( GoogleCloudOptions) self.project_id = gcloud_options.project self._metrics_monitor = MetricsReader( publish_to_bq=publish_to_bq, project_name=self.project_id, bq_table=options.metrics_table, bq_dataset=options.metrics_dataset, namespace=self.metrics_namespace, influxdb_options=InfluxDBMetricsPublisherOptions( options.influx_measurement, options.influx_db_name, options.influx_hostname, os.getenv('INFLUXDB_USER'), os.getenv('INFLUXDB_USER_PASSWORD')), # Apply filter to prevent system metrics from being published filters=MetricsFilter().with_namespace(self.metrics_namespace)) def test(self): """An abstract method where the pipeline definition should be put.""" pass def cleanup(self): """An abstract method that executes after the test method.""" pass def run(self): try: self.test() if not hasattr(self, 'result'): self.result = self.pipeline.run() # Defaults to waiting forever, unless timeout_ms has been set state = self.result.wait_until_finish(duration=self.timeout_ms) assert state != PipelineState.FAILED self._metrics_monitor.publish_metrics(self.result, self.extra_metrics) finally: self.cleanup() def parse_synthetic_source_options(self, options=None): if not options: options = self.input_options return { 'numRecords': options.get('num_records'), 'keySizeBytes': options.get('key_size'), 'valueSizeBytes': options.get('value_size'), 'hotKeyFraction': options.get('hot_key_fraction', 0), 'numHotKeys': options.get('num_hot_keys', 0), 'bundleSizeDistribution': { 'type': options.get('bundle_size_distribution_type', 'const'), 'param': options.get('bundle_size_distribution_param', 0) }, 'forceNumInitialBundles': options.get('force_initial_num_bundles', 0), 'algorithm': options.get('algorithm', None) } def get_option_or_default(self, opt_name, default=0): """Returns a testing option or a default value if it was not provided. The returned value is cast to the type of the default value. """ option = self.pipeline.get_option( opt_name, bool_option=type(default) == bool) if option is None: return default try: return type(default)(option) except: raise ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/load_tests/load_test_metrics_utils.py0000644000175100017510000005274115206113032027002 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Utility functions used for integrating Metrics API into load tests pipelines. Metrics are send to BigQuery in following format: test_id | submit_timestamp | metric_type | value The 'test_id' is common for all metrics for one run. Currently it is possible to have following metrics types: * runtime * total_bytes_count """ # pytype: skip-file import json import logging import time import uuid from typing import Any from typing import Mapping from typing import Optional from typing import Union import requests from requests.auth import HTTPBasicAuth import apache_beam as beam from apache_beam.metrics import Metrics from apache_beam.metrics.metric import MetricResults from apache_beam.metrics.metric import MetricsFilter from apache_beam.runners.dataflow.dataflow_runner import DataflowPipelineResult from apache_beam.runners.runner import PipelineResult from apache_beam.transforms.window import TimestampedValue from apache_beam.utils.timestamp import Timestamp try: from google.cloud import bigquery from google.cloud.bigquery.schema import SchemaField from google.cloud.exceptions import NotFound except ImportError: bigquery = None # type: ignore SchemaField = None # type: ignore NotFound = None # type: ignore RUNTIME_METRIC = 'runtime' COUNTER_LABEL = 'total_bytes_count' ID_LABEL = 'test_id' SUBMIT_TIMESTAMP_LABEL = 'timestamp' METRICS_TYPE_LABEL = 'metric' VALUE_LABEL = 'value' JOB_ID_LABEL = 'job_id' SCHEMA = [{ 'name': ID_LABEL, 'field_type': 'STRING', 'mode': 'REQUIRED' }, { 'name': SUBMIT_TIMESTAMP_LABEL, 'field_type': 'TIMESTAMP', 'mode': 'REQUIRED' }, { 'name': METRICS_TYPE_LABEL, 'field_type': 'STRING', 'mode': 'REQUIRED' }, { 'name': VALUE_LABEL, 'field_type': 'FLOAT', 'mode': 'REQUIRED' }, { 'name': JOB_ID_LABEL, 'field_type': 'STRING', 'mode': 'NULLABLE' }] _LOGGER = logging.getLogger(__name__) def parse_step(step_name): """Replaces white spaces and removes 'Step:' label Args: step_name(str): step name passed in metric ParDo Returns: lower case step name without namespace and step label """ prefix = 'step' step_name = step_name.lower().replace(' ', '_') step_name = ( step_name[len(prefix):] if prefix and step_name.startswith(prefix) else step_name) return step_name.strip(':_') def split_metrics_by_namespace_and_name(metrics, namespace, name): """Splits metrics list namespace and name. Args: metrics: list of metrics from pipeline result namespace(str): filter metrics by namespace name(str): filter metrics by name Returns: two lists - one of metrics which are matching filters and second of not matching """ matching_metrics = [] not_matching_metrics = [] for dist in metrics: if dist.key.metric.namespace == namespace\ and dist.key.metric.name == name: matching_metrics.append(dist) else: not_matching_metrics.append(dist) return matching_metrics, not_matching_metrics def get_generic_distributions(generic_dists, metric_id): """Creates flatten list of distributions per its value type. A generic distribution is the one which is not processed but saved in the most raw version. Args: generic_dists: list of distributions to be saved metric_id(uuid): id of the current test run Returns: list of dictionaries made from :class:`DistributionMetric` """ return sum(( get_all_distributions_by_type(dist, metric_id) for dist in generic_dists), []) def get_all_distributions_by_type(dist, metric_id): """Creates new list of objects with type of each distribution metric value. Args: dist(object): DistributionMetric object to be parsed metric_id(uuid): id of the current test run Returns: list of :class:`DistributionMetric` objects """ submit_timestamp = time.time() dist_types = ['count', 'max', 'min', 'sum', 'mean'] distribution_dicts = [] for dist_type in dist_types: try: distribution_dicts.append( get_distribution_dict(dist_type, submit_timestamp, dist, metric_id)) except ValueError: # Ignore metrics with 'None' values. continue return distribution_dicts def get_distribution_dict(metric_type, submit_timestamp, dist, metric_id): """Function creates :class:`DistributionMetric` Args: metric_type(str): type of value from distribution metric which will be saved (ex. max, min, mean, sum) submit_timestamp: timestamp when metric is saved dist(object) distribution object from pipeline result metric_id(uuid): id of the current test run Returns: dictionary prepared for saving according to schema """ return DistributionMetric(dist, submit_timestamp, metric_id, metric_type).as_dict() class MetricsReader(object): """ A :class:`MetricsReader` retrieves metrics from pipeline result, prepares it for publishers and setup publishers. """ def __init__( self, project_name=None, bq_table=None, bq_dataset=None, publish_to_bq=False, influxdb_options: Optional['InfluxDBMetricsPublisherOptions'] = None, namespace=None, filters=None): """Initializes :class:`MetricsReader` . Args: project_name (str): project with BigQuery where metrics will be saved bq_table (str): BigQuery table where metrics will be saved bq_dataset (str): BigQuery dataset where metrics will be saved namespace (str): Namespace of the metrics filters: MetricFilter to query only filtered metrics """ self._namespace = namespace self.publishers: list[MetricsPublisher] = [] # publish to console output self.publishers.append(ConsoleMetricsPublisher()) bq_check = project_name and bq_table and bq_dataset and publish_to_bq if bq_check: # publish to BigQuery bq_publisher = BigQueryMetricsPublisher( project_name, bq_table, bq_dataset) self.publishers.append(bq_publisher) if influxdb_options and influxdb_options.validate(): # publish to InfluxDB self.publishers.append(InfluxDBMetricsPublisher(influxdb_options)) else: _LOGGER.info( 'Missing InfluxDB options. Metrics will not be published to ' 'InfluxDB') self.filters = filters def get_counter_metric(self, result: PipelineResult, name: str) -> int: """ Return the current value for a long counter, or -1 if can't be retrieved. Note this uses only attempted metrics because some runners don't support committed metrics. """ filters = MetricsFilter().with_namespace(self._namespace).with_name(name) counters = result.metrics().query(filters)[MetricResults.COUNTERS] num_results = len(counters) if num_results > 1: raise ValueError( f"More than one metric result matches name: {name} in namespace "\ f"{self._namespace}. Metric results count: {num_results}") elif num_results == 0: return -1 else: return counters[0].attempted def publish_metrics( self, result: PipelineResult, extra_metrics: Optional[dict] = None): """Publish metrics from pipeline result to registered publishers.""" metric_id = uuid.uuid4().hex metrics = result.metrics().query(self.filters) # Metrics from pipeline result are stored in map with keys: 'gauges', # 'distributions' and 'counters'. # Under each key there is list of objects of each metric type. It is # required to prepare metrics for publishing purposes. Expected is to have # a list of dictionaries matching the schema. insert_dicts = self._prepare_all_metrics(metrics, metric_id) insert_dicts += self._prepare_extra_metrics(metric_id, extra_metrics) # Add job id for dataflow jobs for easier debugging. job_id = None if isinstance(result, DataflowPipelineResult): job_id = result.job_id() self._add_job_id_to_metrics(insert_dicts, job_id) if len(insert_dicts) > 0: for publisher in self.publishers: publisher.publish(insert_dicts) def _add_job_id_to_metrics(self, metrics: list[dict[str, Any]], job_id) -> list[dict[str, Any]]: for metric in metrics: metric[JOB_ID_LABEL] = job_id return metrics def _prepare_extra_metrics( self, metric_id: str, extra_metrics: Optional[dict] = None): ts = time.time() if not extra_metrics: extra_metrics = {} return [ Metric(ts, metric_id, v, label=k).as_dict() for k, v in extra_metrics.items() ] def publish_values(self, labeled_values): """The method to publish simple labeled values. Args: labeled_values (List[Tuple(str, int)]): list of (label, value) """ metric_dicts = [ Metric(time.time(), uuid.uuid4().hex, value, label=label).as_dict() for label, value in labeled_values ] for publisher in self.publishers: publisher.publish(metric_dicts) def _prepare_all_metrics(self, metrics, metric_id): insert_rows = self._get_counters(metrics['counters'], metric_id) insert_rows += self._get_distributions(metrics['distributions'], metric_id) return insert_rows def _get_counters(self, counters, metric_id): submit_timestamp = time.time() return [ CounterMetric(counter, submit_timestamp, metric_id).as_dict() for counter in counters ] def _get_distributions(self, distributions, metric_id): rows = [] matching_namsespace, not_matching_namespace = \ split_metrics_by_namespace_and_name(distributions, self._namespace, RUNTIME_METRIC) if len(matching_namsespace) > 0: runtime_metric = RuntimeMetric(matching_namsespace, metric_id) rows.append(runtime_metric.as_dict()) if len(not_matching_namespace) > 0: rows += get_generic_distributions(not_matching_namespace, metric_id) return rows class Metric(object): """Metric base class in ready-to-save format.""" def __init__( self, submit_timestamp, metric_id, value, metric=None, label=None): """Initializes :class:`Metric` Args: metric (object): object of metric result submit_timestamp (float): date-time of saving metric to database metric_id (uuid): unique id to identify test run value: value of metric label: custom metric name to be saved in database """ self.submit_timestamp = submit_timestamp self.metric_id = metric_id self.label = label or metric.key.metric.namespace + \ '_' + parse_step(metric.key.step) + \ '_' + metric.key.metric.name self.value = value def as_dict(self): return { SUBMIT_TIMESTAMP_LABEL: self.submit_timestamp, ID_LABEL: self.metric_id, VALUE_LABEL: self.value, METRICS_TYPE_LABEL: self.label } class CounterMetric(Metric): """The Counter Metric in ready-to-publish format. Args: counter_metric (object): counter metric object from MetricResult submit_timestamp (float): date-time of saving metric to database metric_id (uuid): unique id to identify test run """ def __init__(self, counter_metric, submit_timestamp, metric_id): value = counter_metric.result super().__init__(submit_timestamp, metric_id, value, counter_metric) class DistributionMetric(Metric): """The Distribution Metric in ready-to-publish format. Args: dist_metric (object): distribution metric object from MetricResult submit_timestamp (float): date-time of saving metric to database metric_id (uuid): unique id to identify test run """ def __init__(self, dist_metric, submit_timestamp, metric_id, metric_type): custom_label = dist_metric.key.metric.namespace + \ '_' + parse_step(dist_metric.key.step) + \ '_' + metric_type + \ '_' + dist_metric.key.metric.name value = getattr(dist_metric.result, metric_type) if value is None: msg = '%s: the result is expected to be an integer, ' \ 'not None.' % custom_label _LOGGER.debug(msg) raise ValueError(msg) super() \ .__init__(submit_timestamp, metric_id, value, dist_metric, custom_label) class RuntimeMetric(Metric): """The Distribution Metric in ready-to-publish format. Args: runtime_list: list of distributions metrics from MetricResult with runtime name metric_id(uuid): unique id to identify test run """ def __init__(self, runtime_list, metric_id): value = self._prepare_runtime_metrics(runtime_list) submit_timestamp = time.time() # Label does not include step name, because it is one value calculated # out of many steps label = runtime_list[0].key.metric.namespace + \ '_' + RUNTIME_METRIC super().__init__(submit_timestamp, metric_id, value, None, label) def _prepare_runtime_metrics(self, distributions): min_values = [] max_values = [] for dist in distributions: min_values.append(dist.result.min) max_values.append(dist.result.max) # finding real start min_value = min(min_values) # finding real end max_value = max(max_values) runtime_in_s = float(max_value - min_value) return runtime_in_s class MetricsPublisher: """Base class for metrics publishers.""" def publish(self, results): raise NotImplementedError class ConsoleMetricsPublisher(MetricsPublisher): """A :class:`ConsoleMetricsPublisher` publishes collected metrics to console output.""" def publish(self, results): if len(results) > 0: log = "Load test results for test: %s and timestamp: %s:" \ % (results[0][ID_LABEL], results[0][SUBMIT_TIMESTAMP_LABEL]) _LOGGER.info(log) for result in results: log = "Metric: %s Value: %d" \ % (result[METRICS_TYPE_LABEL], result[VALUE_LABEL]) _LOGGER.info(log) else: _LOGGER.info("No test results were collected.") class BigQueryMetricsPublisher(MetricsPublisher): """A :class:`BigQueryMetricsPublisher` publishes collected metrics to BigQuery output.""" def __init__(self, project_name, table, dataset, bq_schema=None): if not bq_schema: bq_schema = SCHEMA self.bq = BigQueryClient(project_name, table, dataset, bq_schema) def publish(self, results): outputs = self.bq.save(results) if len(outputs) > 0: for output in outputs: if output['errors']: _LOGGER.error(output) raise ValueError( 'Unable save rows in BigQuery: {}'.format(output['errors'])) class BigQueryClient(object): """A :class:`BigQueryClient` publishes collected metrics to BigQuery output.""" def __init__(self, project_name, table, dataset, bq_schema=None): self.schema = bq_schema self._namespace = table self._client = bigquery.Client(project=project_name) self._schema_names = self._get_schema_names() schema = self._prepare_schema() self._get_or_create_table(schema, dataset) def _get_schema_names(self): return [schema['name'] for schema in self.schema] def _prepare_schema(self): return [SchemaField(**row) for row in self.schema] def _get_or_create_table(self, bq_schemas, dataset): if self._namespace == '': raise ValueError('Namespace cannot be empty.') dataset = self._get_dataset(dataset) table_ref = dataset.table(self._namespace) try: self._bq_table = self._client.get_table(table_ref) except NotFound: table = bigquery.Table(table_ref, schema=bq_schemas) self._bq_table = self._client.create_table(table) def _update_schema(self): table_schema = self._bq_table.schema if self.schema and len(table_schema) != self.schema: self._bq_table.schema = self._prepare_schema() self._bq_table = self._client.update_table(self._bq_table, ["schema"]) def _get_dataset(self, dataset_name): bq_dataset_ref = self._client.dataset(dataset_name) try: bq_dataset = self._client.get_dataset(bq_dataset_ref) except NotFound: raise ValueError( 'Dataset {} does not exist in your project. ' 'You have to create table first.'.format(dataset_name)) return bq_dataset def save(self, results): # update schema if needed self._update_schema() return self._client.insert_rows(self._bq_table, results) class InfluxDBMetricsPublisherOptions(object): def __init__( self, measurement: str, db_name: str, hostname: str, user: Optional[str] = None, password: Optional[str] = None): self.measurement = measurement self.db_name = db_name self.hostname = hostname self.user = user self.password = password def validate(self) -> bool: return bool(self.measurement) and bool(self.db_name) def http_auth_enabled(self) -> bool: return self.user is not None and self.password is not None class InfluxDBMetricsPublisher(MetricsPublisher): """Publishes collected metrics to InfluxDB database.""" def __init__(self, options: InfluxDBMetricsPublisherOptions): self.options = options def publish( self, results: list[Mapping[str, Union[float, str, int]]]) -> None: url = '{}/write'.format(self.options.hostname) payload = self._build_payload(results) query_str = {'db': self.options.db_name, 'precision': 's'} auth = HTTPBasicAuth(self.options.user, self.options.password) if \ self.options.http_auth_enabled() else None try: response = requests.post( url, params=query_str, data=payload, auth=auth, timeout=60) except requests.exceptions.RequestException as e: _LOGGER.warning('Failed to publish metrics to InfluxDB: ' + str(e)) else: if response.status_code != 204: content = json.loads(response.content) _LOGGER.warning( 'Failed to publish metrics to InfluxDB. Received status code %s ' 'with an error message: %s' % (response.status_code, content['error'])) def _build_payload( self, results: list[Mapping[str, Union[float, str, int]]]) -> str: def build_kv(mapping, key): return '{}={}'.format(key, mapping[key]) points = [] for result in results: comma_separated = [ self.options.measurement, build_kv(result, METRICS_TYPE_LABEL), build_kv(result, ID_LABEL), ] point = ','.join(comma_separated) + ' ' + build_kv(result, VALUE_LABEL) \ + ' ' + str(int(result[SUBMIT_TIMESTAMP_LABEL])) points.append(point) return '\n'.join(points) class MeasureTime(beam.DoFn): """A distribution metric prepared to be added to pipeline as ParDo to measure runtime.""" def __init__(self, namespace): """Initializes :class:`MeasureTime`. namespace(str): namespace of metric """ self.namespace = namespace self.runtime = Metrics.distribution(self.namespace, RUNTIME_METRIC) def start_bundle(self): self.runtime.update(time.time()) def finish_bundle(self): self.runtime.update(time.time()) def process(self, element): yield element class MeasureBytes(beam.DoFn): """Metric to measure how many bytes was observed in pipeline.""" LABEL = 'total_bytes' def __init__(self, namespace, extractor=None): """Initializes :class:`MeasureBytes`. Args: namespace(str): metric namespace extractor: function to extract elements to be count """ self.namespace = namespace self.counter = Metrics.counter(self.namespace, self.LABEL) self.extractor = extractor if extractor else lambda x: (yield x) def process(self, element, *args): for value in self.extractor(element, *args): self.counter.inc(len(value)) yield element class CountMessages(beam.DoFn): LABEL = 'total_messages' def __init__(self, namespace): self.namespace = namespace self.counter = Metrics.counter(self.namespace, self.LABEL) def process(self, element): self.counter.inc(1) yield element class MeasureLatency(beam.DoFn): """A distribution metric which captures the latency based on the timestamps of the processed elements.""" LABEL = 'latency' def __init__(self, namespace): """Initializes :class:`MeasureLatency`. namespace(str): namespace of metric """ self.namespace = namespace self.latency_ms = Metrics.distribution(self.namespace, self.LABEL) self.time_fn = time.time def process(self, element, timestamp=beam.DoFn.TimestampParam): self.latency_ms.update( int(self.time_fn() * 1000) - (timestamp.micros // 1000)) yield element class AssignTimestamps(beam.DoFn): """DoFn to assigned timestamps to elements.""" def __init__(self): # Avoid having to use save_main_session self.time_fn = time.time self.timestamp_val_fn = TimestampedValue self.timestamp_fn = Timestamp def process(self, element): yield self.timestamp_val_fn( element, self.timestamp_fn(micros=int(self.time_fn() * 1000000))) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996186.0 apache_beam-2.74.0/apache_beam/testing/load_tests/microbenchmarks_test.py0000644000175100017510000000644515206113032026264 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ This is a load test that runs a set of basic microbenchmarks for the Python SDK and the DirectRunner. This test does not need any additional options passed to run, besides the dataset information. Example test run: python -m apache_beam.testing.load_tests.microbenchmarks_test \ --test-pipeline-options=" --project=big-query-project --input_options='{}' --region=... --publish_to_big_query=true --metrics_dataset=python_load_tests --metrics_table=microbenchmarks" or: ./gradlew -PloadTest.args=" --publish_to_big_query=true --project=... --region=... --input_options='{}' --metrics_dataset=python_load_tests --metrics_table=microbenchmarks --runner=DirectRunner" \ -PloadTest.mainClass=apache_beam.testing.load_tests.microbenchmarks_test \ -Prunner=DirectRunner :sdks:python:apache_beam:testing:load_tests:run """ # pytype: skip-file import logging import time from apache_beam.testing.load_tests.load_test import LoadTest from apache_beam.tools import fn_api_runner_microbenchmark from apache_beam.tools import teststream_microbenchmark from apache_beam.transforms.util import _BatchSizeEstimator class MicroBenchmarksLoadTest(LoadTest): def __init__(self): super().__init__() def test(self): self.extra_metrics.update(self._run_fn_api_runner_microbenchmark()) self.extra_metrics.update(self._run_teststream_microbenchmark()) def _run_teststream_microbenchmark(self): start = time.perf_counter() result = teststream_microbenchmark.run_benchmark(verbose=False) sizes = list(result[0].values())[0] costs = list(result[1].values())[0] a, b = _BatchSizeEstimator.linear_regression_no_numpy(sizes, costs) return { 'teststream_microbenchmark_runtime_sec': time.perf_counter() - start, 'teststream_microbenchmark_fixed_cost_ms': a * 1000, 'teststream_microbenchmark_per_element_cost_ms': b * 1000, } def _run_fn_api_runner_microbenchmark(self): start = time.perf_counter() result = fn_api_runner_microbenchmark.run_benchmark(verbose=False) sizes = list(result[0].values())[0] costs = list(result[1].values())[0] a, b = _BatchSizeEstimator.linear_regression_no_numpy(sizes, costs) return { 'fn_api_runner_microbenchmark_runtime_sec': time.perf_counter() - start, 'fn_api_runner_microbenchmark_fixed_cost_ms': a * 1000, 'fn_api_runner_microbenchmark_per_element_cost_ms': b * 1000, } if __name__ == '__main__': logging.basicConfig(level=logging.INFO) MicroBenchmarksLoadTest().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/load_tests/pardo_test.py0000644000175100017510000001410115206113033024207 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ This is ParDo load test with Synthetic Source. Besides of the standard input options there are additional options: * iterations - number of subsequent ParDo transforms to be performed, * number_of_counters - number of counter metrics to be created for one ParDo transform, * number_of_counter_operations - number of operations on counters to be performed in one ParDo, * project (optional) - the gcp project in case of saving metrics in Big Query (in case of Dataflow Runner it is required to specify project of runner), * publish_to_big_query - if metrics should be published in big query, * metrics_dataset (optional) - name of BigQuery dataset where metrics will be stored, * metrics_table (optional) - name of BigQuery table where metrics will be stored, * input_options - options for Synthetic Sources. * stateful - When true, this will use a stateful DoFn * state_cache - When true, this will enable the Python state cache Example test run: python -m apache_beam.testing.load_tests.pardo_test \ --test-pipeline-options=" --iterations=1 --number_of_counters=1 --number_of_counter_operations=1 --project=big-query-project --region=... --publish_to_big_query=true --metrics_dataset=python_load_tests --metrics_table=pardo --input_options='{ \"num_records\": 300, \"key_size\": 5, \"value_size\": 15 }'" or: ./gradlew -PloadTest.args=" --publish_to_big_query=true --project=... --region=... --metrics_dataset=python_load_tests --metrics_table=pardo --input_options='{ \"num_records\": 1, \"key_size\": 1, \"value_size\": 1}' --runner=DirectRunner" \ -PloadTest.mainClass=apache_beam.testing.load_tests.pardo_test \ -Prunner=DirectRunner :sdks:python:apache_beam:testing:load_tests:run """ # pytype: skip-file import logging import apache_beam as beam from apache_beam.metrics import Metrics from apache_beam.options.pipeline_options import DebugOptions from apache_beam.testing.load_tests.load_test import LoadTest from apache_beam.testing.load_tests.load_test_metrics_utils import AssignTimestamps from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureLatency from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime from apache_beam.testing.synthetic_pipeline import StatefulLoadGenerator from apache_beam.testing.synthetic_pipeline import SyntheticSource from apache_beam.transforms import userstate class ParDoTest(LoadTest): def __init__(self): super().__init__() self.iterations = self.get_option_or_default('iterations') self.number_of_counters = self.get_option_or_default( 'number_of_counters', 1) self.number_of_operations = self.get_option_or_default( 'number_of_counter_operations', 1) self.stateful = self.get_option_or_default('stateful', False) if self.get_option_or_default('state_cache', False): self.pipeline.options.view_as(DebugOptions).add_experiment( 'state_cache_size=1000') def test(self): class BaseCounterOperation(beam.DoFn): def __init__(self, number_of_counters, number_of_operations): self.number_of_operations = number_of_operations self.counters = [] for i in range(number_of_counters): self.counters.append( Metrics.counter('do-not-publish', 'name-{}'.format(i))) class StatefulCounterOperation(BaseCounterOperation): state_param = beam.DoFn.StateParam( userstate.CombiningValueStateSpec( 'count', beam.coders.IterableCoder(beam.coders.VarIntCoder()), sum)) if self.stateful else None def process(self, element, state=state_param): for _ in range(self.number_of_operations): for counter in self.counters: counter.inc() if state: state.add(1) yield element class CounterOperation(BaseCounterOperation): def process(self, element): for _ in range(self.number_of_operations): for counter in self.counters: counter.inc() yield element if self.get_option_or_default('use_stateful_load_generator', False): pc = ( self.pipeline | 'LoadGenerator' >> StatefulLoadGenerator(self.input_options) | 'Measure time: Start' >> beam.ParDo( MeasureTime(self.metrics_namespace)) | 'Assign timestamps' >> beam.ParDo(AssignTimestamps())) for i in range(self.iterations): pc |= 'Step: %d' % i >> beam.ParDo( StatefulCounterOperation( self.number_of_counters, self.number_of_operations)) pc |= 'Measure latency' >> beam.ParDo( MeasureLatency(self.metrics_namespace)) else: pc = ( self.pipeline | 'Read synthetic' >> beam.io.Read( SyntheticSource(self.parse_synthetic_source_options())) | 'Measure time: Start' >> beam.ParDo( MeasureTime(self.metrics_namespace))) for i in range(self.iterations): pc |= 'Step: %d' % i >> beam.ParDo( CounterOperation( self.number_of_counters, self.number_of_operations)) # pylint: disable=expression-not-assigned pc | 'Measure time: End' >> beam.ParDo(MeasureTime(self.metrics_namespace)) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) ParDoTest().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/load_tests/sideinput_test.py0000644000175100017510000001706415206113033025121 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Load test for operations involving side inputs. The purpose of this test is to measure the cost of materialization and accessing side inputs. The test uses synthetic source which can be parametrized to generate records with various sizes of keys and values, impose delays in the pipeline and simulate other performance challenges. This test can accept the following parameters: * side_input_type (str) - Required. Specifies how the side input will be materialized in ParDo operation. Choose from (dict, iter, list). * window_count (int) - The number of fixed sized windows to subdivide the side input into. By default, a global window will be used. * access_percentage (int) - Specifies the percentage of elements in the side input to be accessed. By default, all elements will be accessed. Example test run: python -m apache_beam.testing.load_tests.sideinput_test \ --test-pipeline-options=" --side_input_type=iter --input_options='{ \"num_records\": 300, \"key_size\": 5, \"value_size\": 15 }'" or: ./gradlew -PloadTest.args=" --side_input_type=iter --input_options='{ \"num_records\": 300, \"key_size\": 5, \"value_size\": 15}'" \ -PloadTest.mainClass=apache_beam.testing.load_tests.sideinput_test \ -Prunner=DirectRunner :sdks:python:apache_beam:testing:load_tests:run """ # pytype: skip-file import logging from typing import Any from typing import Iterable from typing import Union import apache_beam as beam from apache_beam.testing.load_tests.load_test import LoadTest from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime from apache_beam.testing.synthetic_pipeline import SyntheticSDFAsSource from apache_beam.transforms import window class SideInputTest(LoadTest): SIDE_INPUT_TYPES = { 'iter': beam.pvalue.AsIter, 'list': beam.pvalue.AsList, 'dict': beam.pvalue.AsDict, } SDF_INITIAL_ELEMENTS = 1000 def __init__(self): super().__init__() self.windows = self.get_option_or_default('window_count', default=1) self.access_percentage = self.get_option_or_default( 'access_percentage', default=100) if self.access_percentage < 0 or self.access_percentage > 100: raise ValueError( 'access_percentage: Invalid value. Should be in range ' 'from 0 to 100, got {} instead'.format(self.access_percentage)) self.elements_per_window = self.input_options['num_records'] // self.windows self.side_input_type = self.pipeline.get_option('side_input_type') if self.side_input_type is None: raise ValueError( 'side_input_type is required. Please provide one of ' 'these: {}'.format(list(self.SIDE_INPUT_TYPES.keys()))) def materialize_as(self): try: return self.SIDE_INPUT_TYPES[self.side_input_type] except KeyError: raise ValueError( 'Unknown side input type. Please provide one of ' 'these: {}'.format(list(self.SIDE_INPUT_TYPES.keys()))) def test(self): class SequenceSideInputTestDoFn(beam.DoFn): """Iterate over first n side_input elements.""" def __init__(self, first_n: int): self._first_n = first_n def process( self, element: Any, side_input: Iterable[tuple[bytes, bytes]]) -> None: i = 0 it = iter(side_input) while i < self._first_n: i += 1 try: # No-op. We only make sure that the element is accessed. next(it) except StopIteration: break class MappingSideInputTestDoFn(beam.DoFn): """Iterates over first n keys in the dictionary and checks the value.""" def __init__(self, first_n: int): self._first_n = first_n def process( self, element: Any, dict_side_input: dict[bytes, bytes]) -> None: i = 0 for key in dict_side_input: if i == self._first_n: break # No-op. We only make sure that the element is accessed. dict_side_input[key] i += 1 class AssignTimestamps(beam.DoFn): """Produces timestamped values. Timestamps are equal to the value of the element.""" def __init__(self): # Avoid having to use save_main_session self.window = window def process(self, element: int) -> Iterable[window.TimestampedValue]: yield self.window.TimestampedValue(element, element) class GetSyntheticSDFOptions(beam.DoFn): def __init__( self, elements_per_record: int, key_size: int, value_size: int): self.elements_per_record = elements_per_record self.key_size = key_size self.value_size = value_size def process(self, element: Any) -> Iterable[dict[str, Union[int, str]]]: yield { 'num_records': self.elements_per_record, 'key_size': self.key_size, 'value_size': self.value_size, 'initial_splitting_num_bundles': 0, 'initial_splitting_desired_bundle_size': 0, 'sleep_per_input_record_sec': 0, 'initial_splitting': 'const' } main_input = self.pipeline | 'Create' >> beam.Create(range(self.windows)) initial_elements = self.SDF_INITIAL_ELEMENTS if self.windows > 1: main_input = ( main_input | 'Assign timestamps' >> beam.ParDo(AssignTimestamps()) | 'Apply windows' >> beam.WindowInto(window.FixedWindows(1))) side_input = main_input initial_elements = self.windows else: side_input = self.pipeline | 'Side input: create' >> beam.Create( range(initial_elements)) side_input = ( side_input | 'Get synthetic SDF options' >> beam.ParDo( GetSyntheticSDFOptions( self.input_options['num_records'] // initial_elements, self.input_options['key_size'], self.input_options['value_size'])) | 'Generate input' >> beam.ParDo(SyntheticSDFAsSource())) main_input |= 'Collect start time metrics' >> beam.ParDo( MeasureTime(self.metrics_namespace)) side_input_type = self.materialize_as() elements_to_access = self.elements_per_window * \ self.access_percentage // 100 logging.info( '%s out of %s total elements in each window will be accessed.', elements_to_access, self.elements_per_window) if side_input_type is beam.pvalue.AsDict: dofn = MappingSideInputTestDoFn(elements_to_access) else: dofn = SequenceSideInputTestDoFn(elements_to_access) _ = ( main_input | beam.ParDo(dofn, side_input_type(side_input)) | 'Collect end time metrics' >> beam.ParDo( MeasureTime(self.metrics_namespace))) if __name__ == '__main__': logging.basicConfig(level=logging.INFO) SideInputTest().run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/metric_result_matchers.py0000644000175100017510000001543415206113033024463 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """MetricResult matchers for validating metrics in PipelineResults. example usage: :: result = my_pipeline.run() all_metrics = result.metrics().all_metrics() matchers = [ MetricResultMatcher( namespace='myNamespace', name='myName', step='myStep', labels={ 'pcollection': 'myCollection', 'myCustomKey': 'myCustomValue' }, attempted=42, committed=42 ) ] errors = metric_result_matchers.verify_all(all_metrics, matchers) self.assertFalse(errors, errors) """ # pytype: skip-file from hamcrest import equal_to from hamcrest.core import string_description from hamcrest.core.base_matcher import BaseMatcher from hamcrest.core.matcher import Matcher from apache_beam.metrics.cells import DistributionResult def _matcher_or_equal_to(value_or_matcher): """Pass-thru for matchers, and wraps value inputs in an equal_to matcher.""" if value_or_matcher is None: return None if isinstance(value_or_matcher, Matcher): return value_or_matcher return equal_to(value_or_matcher) class MetricResultMatcher(BaseMatcher): """A PyHamcrest matcher that validates counter MetricResults.""" def __init__( self, namespace=None, name=None, step=None, labels=None, attempted=None, committed=None, sum_value=None, count_value=None, min_value=None, max_value=None, ): self.namespace = _matcher_or_equal_to(namespace) self.name = _matcher_or_equal_to(name) self.step = _matcher_or_equal_to(step) self.attempted = _matcher_or_equal_to(attempted) self.committed = _matcher_or_equal_to(committed) labels = labels or {} self.label_matchers = {} for (k, v) in labels.items(): self.label_matchers[_matcher_or_equal_to(k)] = _matcher_or_equal_to(v) def _matches(self, metric_result): if self.namespace is not None and not self.namespace.matches( metric_result.key.metric.namespace): return False if self.name and not self.name.matches(metric_result.key.metric.name): return False if self.step and not self.step.matches(metric_result.key.step): return False if (self.attempted is not None and not self.attempted.matches(metric_result.attempted)): return False if (self.committed is not None and not self.committed.matches(metric_result.committed)): return False for (k_matcher, v_matcher) in self.label_matchers.items(): matched_keys = [ key for key in metric_result.key.labels.keys() if k_matcher.matches(key) ] matched_key = matched_keys[0] if matched_keys else None if not matched_key: return False label_value = metric_result.key.labels[matched_key] if not v_matcher.matches(label_value): return False return True def describe_to(self, description): if self.namespace: description.append_text(' namespace: ') self.namespace.describe_to(description) if self.name: description.append_text(' name: ') self.name.describe_to(description) if self.step: description.append_text(' step: ') self.step.describe_to(description) for (k_matcher, v_matcher) in self.label_matchers.items(): description.append_text(' (label_key: ') k_matcher.describe_to(description) description.append_text(' label_value: ') v_matcher.describe_to(description) description.append_text('). ') if self.attempted is not None: description.append_text(' attempted: ') self.attempted.describe_to(description) if self.committed is not None: description.append_text(' committed: ') self.committed.describe_to(description) def describe_mismatch(self, metric_result, mismatch_description): mismatch_description.append_text("was").append_value(metric_result) class DistributionMatcher(BaseMatcher): """A PyHamcrest matcher that validates counter distributions.""" def __init__( self, sum_value=None, count_value=None, min_value=None, max_value=None): self.sum_value = _matcher_or_equal_to(sum_value) self.count_value = _matcher_or_equal_to(count_value) self.min_value = _matcher_or_equal_to(min_value) self.max_value = _matcher_or_equal_to(max_value) def _matches(self, distribution_result): if not isinstance(distribution_result, DistributionResult): return False if self.sum_value and not self.sum_value.matches(distribution_result.sum): return False if self.count_value and not self.count_value.matches( distribution_result.count): return False if self.min_value and not self.min_value.matches(distribution_result.min): return False if self.max_value and not self.max_value.matches(distribution_result.max): return False return True def describe_to(self, description): if self.sum_value: description.append_text(' sum_value: ') self.sum_value.describe_to(description) if self.count_value: description.append_text(' count_value: ') self.count_value.describe_to(description) if self.min_value: description.append_text(' min_value: ') self.min_value.describe_to(description) if self.max_value: description.append_text(' max_value: ') self.max_value.describe_to(description) def describe_mismatch(self, distribution_result, mismatch_description): mismatch_description.append_text('was').append_value(distribution_result) def verify_all(all_metrics, matchers): """Verified that every matcher matches a metric result in all_metrics.""" errors = [] matched_metrics = [] for matcher in matchers: matched_metrics = [mr for mr in all_metrics if matcher.matches(mr)] if not matched_metrics: errors.append( 'Unable to match metrics for matcher %s' % (string_description.tostring(matcher))) if errors: errors.append( '\nActual MetricResults:\n' + '\n'.join([str(mr) for mr in all_metrics])) return ''.join(errors) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/metric_result_matchers_test.py0000644000175100017510000002101015206113033025505 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the metric_result_matchers.""" # pytype: skip-file import unittest from hamcrest import anything from hamcrest import assert_that as hc_assert_that from hamcrest import equal_to from hamcrest.core.core.isnot import is_not from hamcrest.library.number.ordering_comparison import greater_than from hamcrest.library.text.isequal_ignoring_case import equal_to_ignoring_case from apache_beam.metrics.cells import DistributionData from apache_beam.metrics.cells import DistributionResult from apache_beam.metrics.execution import MetricKey from apache_beam.metrics.execution import MetricResult from apache_beam.metrics.metricbase import MetricName from apache_beam.testing.metric_result_matchers import DistributionMatcher from apache_beam.testing.metric_result_matchers import MetricResultMatcher EVERYTHING_DISTRIBUTION = { 'namespace': 'myNamespace', 'name': 'myName', 'step': 'myStep', 'attempted': { 'distribution': { 'sum': 12, 'count': 5, 'min': 0, 'max': 6, } }, 'committed': { 'distribution': { 'sum': 12, 'count': 5, 'min': 0, 'max': 6, } }, 'labels': { 'pcollection': 'myCollection', 'myCustomKey': 'myCustomValue' } } EVERYTHING_COUNTER = { 'namespace': 'myNamespace', 'name': 'myName', 'step': 'myStep', 'attempted': { 'counter': 42 }, 'committed': { 'counter': 42 }, 'labels': { 'pcollection': 'myCollection', 'myCustomKey': 'myCustomValue' } } def _create_metric_result(data_dict): step = data_dict['step'] if 'step' in data_dict else '' labels = data_dict['labels'] if 'labels' in data_dict else {} values = {} for key in ['attempted', 'committed']: if key in data_dict: if 'counter' in data_dict[key]: values[key] = data_dict[key]['counter'] elif 'distribution' in data_dict[key]: distribution = data_dict[key]['distribution'] values[key] = DistributionResult( DistributionData( distribution['sum'], distribution['count'], distribution['min'], distribution['max'], )) attempted = values['attempted'] if 'attempted' in values else None committed = values['committed'] if 'committed' in values else None metric_name = MetricName(data_dict['namespace'], data_dict['name']) metric_key = MetricKey(step, metric_name, labels) return MetricResult(metric_key, committed, attempted) class MetricResultMatchersTest(unittest.TestCase): def test_matches_all_for_counter(self): metric_result = _create_metric_result(EVERYTHING_COUNTER) matcher = MetricResultMatcher( namespace='myNamespace', name='myName', step='myStep', labels={ 'pcollection': 'myCollection', 'myCustomKey': 'myCustomValue' }, attempted=42, committed=42) hc_assert_that(metric_result, matcher) def test_matches_none_for_counter(self): metric_result = _create_metric_result(EVERYTHING_COUNTER) matcher = MetricResultMatcher( namespace=is_not(equal_to('invalidNamespace')), name=is_not(equal_to('invalidName')), step=is_not(equal_to('invalidStep')), labels={ is_not(equal_to('invalidPcollection')): anything(), is_not(equal_to('invalidCustomKey')): is_not( equal_to('invalidCustomValue')) }, attempted=is_not(equal_to(1000)), committed=is_not(equal_to(1000))) hc_assert_that(metric_result, matcher) def test_matches_all_for_distribution(self): metric_result = _create_metric_result(EVERYTHING_DISTRIBUTION) matcher = MetricResultMatcher( namespace='myNamespace', name='myName', step='myStep', labels={ 'pcollection': 'myCollection', 'myCustomKey': 'myCustomValue' }, committed=DistributionMatcher( sum_value=12, count_value=5, min_value=0, max_value=6), attempted=DistributionMatcher( sum_value=12, count_value=5, min_value=0, max_value=6), ) hc_assert_that(metric_result, matcher) def test_matches_none_for_distribution(self): metric_result = _create_metric_result(EVERYTHING_DISTRIBUTION) matcher = MetricResultMatcher( namespace=is_not(equal_to('invalidNamespace')), name=is_not(equal_to('invalidName')), step=is_not(equal_to('invalidStep')), labels={ is_not(equal_to('invalidPcollection')): anything(), is_not(equal_to('invalidCustomKey')): is_not( equal_to('invalidCustomValue')) }, committed=is_not( DistributionMatcher( sum_value=120, count_value=50, min_value=100, max_value=60)), attempted=is_not( DistributionMatcher( sum_value=120, count_value=50, min_value=100, max_value=60)), ) hc_assert_that(metric_result, matcher) def test_matches_key_but_not_value(self): metric_result = _create_metric_result(EVERYTHING_COUNTER) matcher = is_not( MetricResultMatcher(labels={'pcollection': 'invalidCollection'})) hc_assert_that(metric_result, matcher) def test_matches_counter_with_custom_matchers(self): metric_result = _create_metric_result(EVERYTHING_COUNTER) matcher = is_not( MetricResultMatcher( namespace=equal_to_ignoring_case('MYNAMESPACE'), name=equal_to_ignoring_case('MYNAME'), step=equal_to_ignoring_case('MYSTEP'), labels={ equal_to_ignoring_case('PCOLLECTION'): equal_to_ignoring_case( 'MYCUSTOMVALUE'), 'myCustomKey': equal_to_ignoring_case('MYCUSTOMVALUE') }, committed=greater_than(0), attempted=greater_than(0))) hc_assert_that(metric_result, matcher) def test_matches_distribution_with_custom_matchers(self): metric_result = _create_metric_result(EVERYTHING_DISTRIBUTION) matcher = is_not( MetricResultMatcher( namespace=equal_to_ignoring_case('MYNAMESPACE'), name=equal_to_ignoring_case('MYNAME'), step=equal_to_ignoring_case('MYSTEP'), labels={ equal_to_ignoring_case('PCOLLECTION'): equal_to_ignoring_case( 'MYCUSTOMVALUE'), 'myCustomKey': equal_to_ignoring_case('MYCUSTOMVALUE') }, committed=is_not( DistributionMatcher( sum_value=greater_than(-1), count_value=greater_than(-1), min_value=greater_than(-1), max_value=greater_than(-1))), attempted=is_not( DistributionMatcher( sum_value=greater_than(-1), count_value=greater_than(-1), min_value=greater_than(-1), max_value=greater_than(-1))), )) hc_assert_that(metric_result, matcher) def test_counter_does_not_match_distribution_and_doesnt_crash(self): metric_result = _create_metric_result(EVERYTHING_COUNTER) matcher = is_not( MetricResultMatcher( committed=DistributionMatcher( sum_value=120, count_value=50, min_value=100, max_value=60), attempted=DistributionMatcher( sum_value=120, count_value=50, min_value=100, max_value=60), )) hc_assert_that(metric_result, matcher) def test_distribution_does_not_match_counter_and_doesnt_crash(self): metric_result = _create_metric_result(EVERYTHING_DISTRIBUTION) matcher = is_not(MetricResultMatcher(attempted=42, committed=42)) hc_assert_that(metric_result, matcher) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/pipeline_verifiers.py0000644000175100017510000001244315206113033023574 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """End-to-end test result verifiers A set of verifiers that are used in end-to-end tests to verify state/output of test pipeline job. Customized verifier should extend `hamcrest.core.base_matcher.BaseMatcher` and override _matches. """ # pytype: skip-file import logging import time from hamcrest.core.base_matcher import BaseMatcher from apache_beam.io.filesystems import FileSystems from apache_beam.runners.runner import PipelineState from apache_beam.testing import test_utils as utils from apache_beam.utils import retry __all__ = [ 'PipelineStateMatcher', 'FileChecksumMatcher', 'retry_on_io_error_and_server_error', ] try: from google.api_core.exceptions import GoogleAPICallError except ImportError: GoogleAPICallError = None # type: ignore MAX_RETRIES = 4 _LOGGER = logging.getLogger(__name__) class PipelineStateMatcher(BaseMatcher): """Matcher that verify pipeline job terminated in expected state Matcher compares the actual pipeline terminate state with expected. By default, `PipelineState.DONE` is used as expected state. """ def __init__(self, expected_state=PipelineState.DONE): self.expected_state = expected_state def _matches(self, pipeline_result): return pipeline_result.state == self.expected_state def describe_to(self, description): description \ .append_text("Test pipeline expected terminated in state: ") \ .append_text(self.expected_state) def describe_mismatch(self, pipeline_result, mismatch_description): mismatch_description \ .append_text("Test pipeline job terminated in state: ") \ .append_text(pipeline_result.state) def retry_on_io_error_and_server_error(exception): """Filter allowing retries on file I/O errors and service error.""" return isinstance(exception, IOError) or \ (GoogleAPICallError is not None and isinstance(exception, GoogleAPICallError)) # pylint: disable=line-too-long class FileChecksumMatcher(BaseMatcher): """Matcher that verifies file(s) content by comparing file checksum. Use apache_beam.io.filebasedsink to fetch file(s) from given path. File checksum is a hash string computed from content of file(s). """ def __init__(self, file_path, expected_checksum, sleep_secs=None): """Initialize a FileChecksumMatcher object Args: file_path : A string that is the full path of output file. This path can contain globs. expected_checksum : A hash string that is computed from expected result. sleep_secs : Number of seconds to wait before verification start. Extra time are given to make sure output files are ready on FS. """ if sleep_secs is not None: if isinstance(sleep_secs, int): self.sleep_secs = sleep_secs else: raise ValueError( 'Sleep seconds, if received, must be int. ' 'But received: %r, %s' % (sleep_secs, type(sleep_secs))) else: self.sleep_secs = None self.file_path = file_path self.expected_checksum = expected_checksum @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry_on_io_error_and_server_error) def _read_with_retry(self): """Read path with retry if I/O failed""" read_lines = [] match_result = FileSystems.match([self.file_path])[0] matched_path = [f.path for f in match_result.metadata_list] if not matched_path: raise IOError('No such file or directory: %s' % self.file_path) _LOGGER.info( 'Find %d files in %s: \n%s', len(matched_path), self.file_path, '\n'.join(matched_path)) for path in matched_path: with FileSystems.open(path, 'r') as f: for line in f: read_lines.append(line) return read_lines def _matches(self, _): if self.sleep_secs: # Wait to have output file ready on FS _LOGGER.info('Wait %d seconds...', self.sleep_secs) time.sleep(self.sleep_secs) # Read from given file(s) path read_lines = self._read_with_retry() # Compute checksum self.checksum = utils.compute_hash(read_lines) _LOGGER.info( 'Read from given path %s, %d lines, checksum: %s.', self.file_path, len(read_lines), self.checksum) return self.checksum == self.expected_checksum def describe_to(self, description): description \ .append_text("Expected checksum is ") \ .append_text(self.expected_checksum) def describe_mismatch(self, pipeline_result, mismatch_description): mismatch_description \ .append_text("Actual checksum is ") \ .append_text(self.checksum) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/pipeline_verifiers_test.py0000644000175100017510000001350115206113033024627 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the test pipeline verifiers""" # pytype: skip-file import logging import os import tempfile import unittest from hamcrest import assert_that as hc_assert_that from mock import Mock from mock import patch from apache_beam.io.localfilesystem import LocalFileSystem from apache_beam.runners.runner import PipelineResult from apache_beam.runners.runner import PipelineState from apache_beam.testing import pipeline_verifiers as verifiers from apache_beam.testing.test_utils import patch_retry try: # pylint: disable=wrong-import-order, wrong-import-position # pylint: disable=ungrouped-imports from google.api_core.exceptions import GoogleAPICallError from google.api_core.exceptions import NotFound from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem except ImportError: GoogleAPICallError = None # type: ignore NotFound = None # type: ignore GCSFileSystem = None # type: ignore class PipelineVerifiersTest(unittest.TestCase): def setUp(self): self._mock_result = Mock() patch_retry(self, verifiers) def test_pipeline_state_matcher_success(self): """Test PipelineStateMatcher successes when using default expected state and job actually finished in DONE """ pipeline_result = PipelineResult(PipelineState.DONE) hc_assert_that(pipeline_result, verifiers.PipelineStateMatcher()) def test_pipeline_state_matcher_given_state(self): """Test PipelineStateMatcher successes when matches given state""" pipeline_result = PipelineResult(PipelineState.FAILED) hc_assert_that( pipeline_result, verifiers.PipelineStateMatcher(PipelineState.FAILED)) def test_pipeline_state_matcher_fails(self): """Test PipelineStateMatcher fails when using default expected state and job actually finished in CANCELLED/DRAINED/FAILED/UNKNOWN """ failed_state = [ PipelineState.CANCELLED, PipelineState.DRAINED, PipelineState.FAILED, PipelineState.UNKNOWN ] for state in failed_state: pipeline_result = PipelineResult(state) with self.assertRaises(AssertionError): hc_assert_that(pipeline_result, verifiers.PipelineStateMatcher()) test_cases = [ { 'content': 'Test FileChecksumMatcher with single file', 'num_files': 1, 'expected_checksum': 'ebe16840cc1d0b4fe1cf71743e9d772fa31683b8' }, { 'content': 'Test FileChecksumMatcher with multiple files', 'num_files': 3, 'expected_checksum': '58b3d3636de3891ac61afb8ace3b5025c3c37d44' }, { 'content': '', 'num_files': 1, 'expected_checksum': 'da39a3ee5e6b4b0d3255bfef95601890afd80709' }, ] def create_temp_file(self, content, directory=None): with tempfile.NamedTemporaryFile(delete=False, dir=directory) as f: f.write(content.encode('utf-8')) return f.name def test_file_checksum_matcher_success(self): for case in self.test_cases: temp_dir = tempfile.mkdtemp() for _ in range(case['num_files']): self.create_temp_file(case['content'], temp_dir) matcher = verifiers.FileChecksumMatcher( os.path.join(temp_dir, '*'), case['expected_checksum']) hc_assert_that(self._mock_result, matcher) @patch.object(LocalFileSystem, 'match') def test_file_checksum_matcher_read_failed(self, mock_match): mock_match.side_effect = IOError('No file found.') matcher = verifiers.FileChecksumMatcher( os.path.join('dummy', 'path'), Mock()) with self.assertRaises(IOError): hc_assert_that(self._mock_result, matcher) self.assertTrue(mock_match.called) self.assertEqual(verifiers.MAX_RETRIES + 1, mock_match.call_count) @patch.object(GCSFileSystem, 'match') @unittest.skipIf( GoogleAPICallError is None, 'GCP dependencies are not installed') def test_file_checksum_matcher_service_error(self, mock_match): mock_match.side_effect = NotFound('Not Found') matcher = verifiers.FileChecksumMatcher('gs://dummy/path', Mock()) with self.assertRaises(NotFound): hc_assert_that(self._mock_result, matcher) self.assertTrue(mock_match.called) self.assertEqual(verifiers.MAX_RETRIES + 1, mock_match.call_count) def test_file_checksum_matchcer_invalid_sleep_time(self): with self.assertRaises(ValueError) as cm: verifiers.FileChecksumMatcher( 'file_path', 'expected_checksum', 'invalid_sleep_time') self.assertEqual( cm.exception.args[0], 'Sleep seconds, if received, must be int. ' 'But received: \'invalid_sleep_time\', ' '{}'.format(str)) @patch('time.sleep', return_value=None) def test_file_checksum_matcher_sleep_before_verify(self, mocked_sleep): temp_dir = tempfile.mkdtemp() case = self.test_cases[0] self.create_temp_file(case['content'], temp_dir) matcher = verifiers.FileChecksumMatcher( os.path.join(temp_dir, '*'), case['expected_checksum'], 10) hc_assert_that(self._mock_result, matcher) self.assertTrue(mocked_sleep.called) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/synthetic_pipeline.py0000644000175100017510000010427015206113033023610 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A set of utilities to write pipelines for performance tests. This module offers a way to create pipelines using synthetic sources and steps. Exact shape of the pipeline and the behaviour of sources and steps can be controlled through arguments. Please see function 'parse_args()' for more details about the arguments. Shape of the pipeline is primarily controlled through two arguments. Argument 'steps' can be used to define a list of steps as a JSON string. Argument 'barrier' describes how these steps are separated from each other. Argument 'barrier' can be use to build a pipeline as a series of steps or a tree of steps with a fanin or a fanout of size 2. Other arguments describe what gets generated by synthetic sources that produce data for the pipeline. """ # pytype: skip-file import argparse import json import logging import math import os import time from random import Random from typing import Optional import apache_beam as beam from apache_beam import pvalue from apache_beam import typehints from apache_beam.io import WriteToText from apache_beam.io import iobase from apache_beam.io import range_trackers from apache_beam.io import restriction_trackers from apache_beam.io.restriction_trackers import OffsetRange from apache_beam.io.restriction_trackers import OffsetRestrictionTracker from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.transforms import userstate from apache_beam.transforms.core import RestrictionProvider try: import numpy as np except ImportError: np = None try: from .fast_test_utils import LCGenerator except ImportError: # Cythonized lib not available from .test_utils import LCGenerator class _Random(Random): """A subclass of `random.Random` from the Python Standard Library that provides a method returning random bytes of arbitrary length. """ # `numpy.random.RandomState` does not provide `random()` method, we keep this # for compatibility reasons. random_sample = Random.random def get_generator(seed: Optional[int] = None, algorithm: Optional[str] = None): if algorithm is None or algorithm == 'builtin': return _Random(seed) elif algorithm == 'lcg': generator = LCGenerator() if seed is not None: generator.seed(seed) return generator else: raise ValueError( 'Unknown algorithm %s. Supported algorithms are "builtin" or "lcg".', algorithm) def parse_byte_size(s): suffixes = 'BKMGTP' if s[-1] in suffixes: return int(float(s[:-1]) * 1024**suffixes.index(s[-1])) return int(s) def div_round_up(a, b): """Return ceil(a/b).""" return int(math.ceil(float(a) / b)) def rotate_key(element): """Returns a new key-value pair of the same size but with a different key.""" (key, value) = element return key[-1:] + key[:-1], value def initial_splitting_zipf( start_position, stop_position, desired_num_bundles, distribution_parameter, num_total_records=None): """Split the given range (defined by start_position, stop_position) into desired_num_bundles using zipf with the given distribution_parameter. """ if not num_total_records: num_total_records = stop_position - start_position samples = np.random.zipf(distribution_parameter, desired_num_bundles) total = sum(samples) relative_bundle_sizes = [(float(sample) / total) for sample in samples] bundle_ranges = [] start = start_position index = 0 while start < stop_position: if index == desired_num_bundles - 1: bundle_ranges.append((start, stop_position)) break stop = start + int(num_total_records * relative_bundle_sizes[index]) bundle_ranges.append((start, stop)) start = stop index += 1 return bundle_ranges class SyntheticStep(beam.DoFn): """A DoFn of which behavior can be controlled through prespecified parameters. """ def __init__( self, per_element_delay_sec=0, per_bundle_delay_sec=0, output_records_per_input_record=1, output_filter_ratio=0): if per_element_delay_sec and per_element_delay_sec < 1e-3: raise ValueError( 'Per element sleep time must be at least 1e-3. ' 'Received: %r', per_element_delay_sec) self._per_element_delay_sec = per_element_delay_sec self._per_bundle_delay_sec = per_bundle_delay_sec self._output_records_per_input_record = output_records_per_input_record self._output_filter_ratio = output_filter_ratio def start_bundle(self): self._start_time = time.time() def finish_bundle(self): # The target is for the enclosing stage to take as close to as possible # the given number of seconds, so we only sleep enough to make up for # overheads not incurred elsewhere. to_sleep = self._per_bundle_delay_sec - (time.time() - self._start_time) # Ignoring sub-millisecond sleep times. if to_sleep >= 1e-3: time.sleep(to_sleep) def process(self, element): if self._per_element_delay_sec >= 1e-3: time.sleep(self._per_element_delay_sec) filter_element = False if self._output_filter_ratio > 0: if np.random.random() < self._output_filter_ratio: filter_element = True if not filter_element: for _ in range(self._output_records_per_input_record): yield element class NonLiquidShardingOffsetRangeTracker(OffsetRestrictionTracker): """An OffsetRangeTracker that doesn't allow splitting. """ def try_split(self, split_offset): pass # Don't split. def checkpoint(self): pass # Don't split. class SyntheticSDFStepRestrictionProvider(RestrictionProvider): """A `RestrictionProvider` for SyntheticSDFStep. An initial_restriction and split that operate on num_records and ignores source description (element). Splits into initial_splitting_num_bundles. Returns size_estimate_override as restriction size, if set. Otherwise uses element size. If initial_splitting_uneven_chunks, produces uneven chunks. """ def __init__( self, num_records, initial_splitting_num_bundles, initial_splitting_uneven_chunks, disable_liquid_sharding, size_estimate_override): self._num_records = num_records self._initial_splitting_num_bundles = initial_splitting_num_bundles self._initial_splitting_uneven_chunks = initial_splitting_uneven_chunks self._disable_liquid_sharding = disable_liquid_sharding self._size_estimate_override = size_estimate_override def initial_restriction(self, element): return OffsetRange(0, self._num_records) def create_tracker(self, restriction): if self._disable_liquid_sharding: return NonLiquidShardingOffsetRangeTracker(restriction) else: return OffsetRestrictionTracker(restriction) def split(self, element, restriction): elems = restriction.size() if (self._initial_splitting_uneven_chunks and self._initial_splitting_num_bundles > 1 and elems > 1): bundle_ranges = initial_splitting_zipf( restriction.start, restriction.stop, self._initial_splitting_num_bundles, 3.0) for start, stop in bundle_ranges: yield OffsetRange(start, stop) else: offsets_per_split = max(1, (elems // self._initial_splitting_num_bundles)) for split in restriction.split(offsets_per_split, offsets_per_split // 2): yield split def restriction_size(self, element, restriction): if self._size_estimate_override is not None: return self._size_estimate_override element_size = len(element) if isinstance(element, str) else 1 return restriction.size() * element_size def get_synthetic_sdf_step( per_element_delay_sec=0, per_bundle_delay_sec=0, output_records_per_input_record=1, output_filter_ratio=0, initial_splitting_num_bundles=8, initial_splitting_uneven_chunks=False, disable_liquid_sharding=False, size_estimate_override=None, ): """A function which returns a SyntheticSDFStep with given parameters. """ class SyntheticSDFStep(beam.DoFn): """A SplittableDoFn of which behavior can be controlled through prespecified parameters. """ def __init__( self, per_element_delay_sec_arg, per_bundle_delay_sec_arg, output_filter_ratio_arg, output_records_per_input_record_arg): if per_element_delay_sec_arg: per_element_delay_sec_arg = ( per_element_delay_sec_arg // output_records_per_input_record_arg) if per_element_delay_sec_arg < 1e-3: raise ValueError( 'Per element sleep time must be at least 1e-3 after being ' 'divided among output elements.') self._per_element_delay_sec = per_element_delay_sec_arg self._per_bundle_delay_sec = per_bundle_delay_sec_arg self._output_filter_ratio = output_filter_ratio_arg def start_bundle(self): self._start_time = time.time() def finish_bundle(self): # The target is for the enclosing stage to take as close to as possible # the given number of seconds, so we only sleep enough to make up for # overheads not incurred elsewhere. to_sleep = self._per_bundle_delay_sec - (time.time() - self._start_time) # Ignoring sub-millisecond sleep times. if to_sleep >= 1e-3: time.sleep(to_sleep) def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam( SyntheticSDFStepRestrictionProvider( output_records_per_input_record, initial_splitting_num_bundles, initial_splitting_uneven_chunks, disable_liquid_sharding, size_estimate_override))): filter_element = False if self._output_filter_ratio > 0: if np.random.random() < self._output_filter_ratio: filter_element = True current_restriction = restriction_tracker.current_restriction() for cur in range(current_restriction.start, current_restriction.stop): if not restriction_tracker.try_claim(cur): return if self._per_element_delay_sec: time.sleep(self._per_element_delay_sec) if not filter_element: yield element cur += 1 return SyntheticSDFStep( per_element_delay_sec, per_bundle_delay_sec, output_filter_ratio, output_records_per_input_record) class SyntheticSource(iobase.BoundedSource): """A custom source of a specified size. """ def __init__(self, input_spec): """Initiates a synthetic source. Args: input_spec: Input specification of the source. See corresponding option in function 'parse_args()' below for more details. Raises: ValueError: if input parameters are invalid. """ def maybe_parse_byte_size(s): return parse_byte_size(s) if isinstance(s, str) else int(s) self._num_records = input_spec['numRecords'] self._key_size = maybe_parse_byte_size(input_spec.get('keySizeBytes', 1)) self._hot_key_fraction = input_spec.get('hotKeyFraction', 0) self._num_hot_keys = input_spec.get('numHotKeys', 0) self._value_size = maybe_parse_byte_size( input_spec.get('valueSizeBytes', 1)) self._total_size = self.element_size * self._num_records self._initial_splitting = ( input_spec['bundleSizeDistribution']['type'] if 'bundleSizeDistribution' in input_spec else 'const') if self._initial_splitting != 'const' and self._initial_splitting != 'zipf': raise ValueError( 'Only const and zipf distributions are supported for determining ' 'sizes of bundles produced by initial splitting. Received: %s', self._initial_splitting) self._initial_splitting_num_bundles = ( input_spec['forceNumInitialBundles'] if 'forceNumInitialBundles' in input_spec else 0) if self._initial_splitting == 'zipf': self._initial_splitting_distribution_parameter = ( input_spec['bundleSizeDistribution']['param']) if self._initial_splitting_distribution_parameter < 1: raise ValueError( 'Parameter for a Zipf distribution must be larger than 1. ' 'Received %r.', self._initial_splitting_distribution_parameter) else: self._initial_splitting_distribution_parameter = 0 self._dynamic_splitting = ( 'none' if ( 'splitPointFrequencyRecords' in input_spec and input_spec['splitPointFrequencyRecords'] == 0) else 'perfect') if 'delayDistribution' in input_spec: if input_spec['delayDistribution']['type'] != 'const': raise ValueError( 'SyntheticSource currently only supports delay ' 'distributions of type \'const\'. Received %s.', input_spec['delayDistribution']['type']) self._sleep_per_input_record_sec = ( float(input_spec['delayDistribution']['const']) / 1000) if (self._sleep_per_input_record_sec and self._sleep_per_input_record_sec < 1e-3): raise ValueError( 'Sleep time per input record must be at least 1e-3.' ' Received: %r', self._sleep_per_input_record_sec) else: self._sleep_per_input_record_sec = 0 # algorithm of the generator self.gen_algo = input_spec.get('algorithm', None) if self.gen_algo not in (None, 'builtin', 'lcg'): raise ValueError( 'Unknown algorithm for input_spec: %s. Supported ' 'algorithms are "builtin" and "lcg".', self.gen_algo) @property def element_size(self): return self._key_size + self._value_size def estimate_size(self): return self._total_size def split(self, desired_bundle_size, start_position=0, stop_position=None): # Performs initial splitting of SyntheticSource. # # Exact sizes and distribution of initial splits generated here depends on # the input specification of the SyntheticSource. if stop_position is None: stop_position = self._num_records if self._initial_splitting == 'zipf': desired_num_bundles = self._initial_splitting_num_bundles or math.ceil( float(self.estimate_size()) / desired_bundle_size) bundle_ranges = initial_splitting_zipf( start_position, stop_position, desired_num_bundles, self._initial_splitting_distribution_parameter, self._num_records) else: if self._initial_splitting_num_bundles: bundle_size_in_elements = max( 1, int(self._num_records / self._initial_splitting_num_bundles)) else: bundle_size_in_elements = ( max( div_round_up(desired_bundle_size, self.element_size), int(math.floor(math.sqrt(self._num_records))))) bundle_ranges = [] for start in range(start_position, stop_position, bundle_size_in_elements): stop = min(start + bundle_size_in_elements, stop_position) bundle_ranges.append((start, stop)) for start, stop in bundle_ranges: yield iobase.SourceBundle(stop - start, self, start, stop) def get_range_tracker(self, start_position, stop_position): if start_position is None: start_position = 0 if stop_position is None: stop_position = self._num_records tracker = range_trackers.OffsetRangeTracker(start_position, stop_position) if self._dynamic_splitting == 'none': tracker = range_trackers.UnsplittableRangeTracker(tracker) return tracker def _gen_kv_pair(self, generator, index): generator.seed(index) rand = generator.random_sample() # Determines whether to generate hot key or not. if rand < self._hot_key_fraction: # Generate hot key. # An integer is randomly selected from the range [0, numHotKeys-1] # with equal probability. generator_hot = get_generator( seed=index % self._num_hot_keys, algorithm=self.gen_algo) bytes_ = generator_hot.randbytes(self._key_size), generator.randbytes( self._value_size) else: bytes_ = generator.randbytes(self.element_size) bytes_ = bytes_[:self._key_size], bytes_[self._key_size:] return bytes_ def read(self, range_tracker): index = range_tracker.start_position() generator = get_generator(algorithm=self.gen_algo) while range_tracker.try_claim(index): time.sleep(self._sleep_per_input_record_sec) yield self._gen_kv_pair(generator, index) index += 1 def default_output_coder(self): return beam.coders.TupleCoder( [beam.coders.BytesCoder(), beam.coders.BytesCoder()]) class SyntheticSDFSourceRestrictionProvider(RestrictionProvider): """A `RestrictionProvider` for SyntheticSDFAsSource. In initial_restriction(element) and split(element), element means source description. A typical element is like: { 'key_size': 1, 'value_size': 1, 'initial_splitting_num_bundles': 8, 'initial_splitting_desired_bundle_size': 2, 'sleep_per_input_record_sec': 0, 'initial_splitting' : 'const' } """ def initial_restriction(self, element): return OffsetRange(0, element['num_records']) def create_tracker(self, restriction): return restriction_trackers.OffsetRestrictionTracker(restriction) def split(self, element, restriction): bundle_ranges = [] start_position = restriction.start stop_position = restriction.stop element_size = element['key_size'] + element['value_size'] estimate_size = element_size * element['num_records'] if element['initial_splitting'] == 'zipf': desired_num_bundles = ( element['initial_splitting_num_bundles'] or div_round_up( estimate_size, element['initial_splitting_desired_bundle_size'])) samples = np.random.zipf( element['initial_splitting_distribution_parameter'], desired_num_bundles) total = sum(samples) relative_bundle_sizes = [(float(sample) / total) for sample in samples] start = start_position index = 0 while start < stop_position: if index == desired_num_bundles - 1: bundle_ranges.append(OffsetRange(start, stop_position)) break stop = start + int( element['num_records'] * relative_bundle_sizes[index]) bundle_ranges.append(OffsetRange(start, stop)) start = stop index += 1 else: if element['initial_splitting_num_bundles']: bundle_size_in_elements = max( 1, int( element['num_records'] / element['initial_splitting_num_bundles'])) else: bundle_size_in_elements = ( max( div_round_up( element['initial_splitting_desired_bundle_size'], element_size), int(math.floor(math.sqrt(element['num_records']))))) for start in range(start_position, stop_position, bundle_size_in_elements): stop = min(start + bundle_size_in_elements, stop_position) bundle_ranges.append(OffsetRange(start, stop)) return bundle_ranges def restriction_size(self, element, restriction): return (element['key_size'] + element['value_size']) * restriction.size() class SyntheticSDFAsSource(beam.DoFn): """A SDF that generates records like a source. This SDF accepts a PCollection of record-based source description. A typical description is like: { 'key_size': 1, 'value_size': 1, 'initial_splitting_num_bundles': 8, 'initial_splitting_desired_bundle_size': 2, 'sleep_per_input_record_sec': 0, 'initial_splitting' : 'const' } A simple pipeline taking this SDF as a source is like: p | beam.Create([description1, description2,...]) | beam.ParDo(SyntheticSDFAsSource()) NOTE: The SDF.process() will have different param content between defining a DoFn and runtime. When defining an SDF.process, the restriction_tracker should be a `RestrictionProvider`. During runtime, the DoFnRunner.process_with_sized_restriction() will feed a 'RestrictionTracker' based on a restriction to SDF.process(). """ def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam( SyntheticSDFSourceRestrictionProvider())): cur = restriction_tracker.current_restriction().start while restriction_tracker.try_claim(cur): r = get_generator(algorithm=element.get('algorithm', None), seed=cur) time.sleep(element['sleep_per_input_record_sec']) yield ( r.randbytes(element['key_size']), r.randbytes(element['value_size'])) cur += 1 class ShuffleBarrier(beam.PTransform): def expand(self, pc): return ( pc | beam.Map(rotate_key) | beam.GroupByKey() | 'Ungroup' >> beam.FlatMap(lambda elm: [(elm[0], v) for v in elm[1]])) class SideInputBarrier(beam.PTransform): def expand(self, pc): return ( pc | beam.Map(rotate_key) | beam.Map( lambda elem, ignored: elem, beam.pvalue.AsIter(pc | beam.FlatMap(lambda elem: None)))) def merge_using_gbk(name, pc1, pc2): """Merges two given PCollections using a CoGroupByKey.""" pc1_with_key = pc1 | (name + 'AttachKey1') >> beam.Map(lambda x: (x, x)) pc2_with_key = pc2 | (name + 'AttachKey2') >> beam.Map(lambda x: (x, x)) grouped = ({ 'pc1': pc1_with_key, 'pc2': pc2_with_key } | (name + 'Group') >> beam.CoGroupByKey()) return ( grouped | (name + 'DeDup') >> beam.Map(lambda elm: elm[0]) ) # Ignoring values def merge_using_side_input(name, pc1, pc2): """Merges two given PCollections using side inputs.""" def join_fn(val, _): # Ignoring side input return val return pc1 | name >> beam.core.Map(join_fn, beam.pvalue.AsIter(pc2)) def expand_using_gbk(name, pc): """Expands a given PCollection into two copies using GroupByKey.""" ret = [] ret.append((pc | ('%s.a' % name) >> ShuffleBarrier())) ret.append((pc | ('%s.b' % name) >> ShuffleBarrier())) return ret def expand_using_second_output(name, pc): """Expands a given PCollection into two copies using side outputs.""" class ExpandFn(beam.DoFn): def process(self, element): yield beam.pvalue.TaggedOutput('second_out', element) yield element pc1, pc2 = (pc | name >> beam.ParDo( ExpandFn()).with_outputs('second_out', main='main_out')) return [pc1, pc2] def _parse_steps(json_str): """Converts the JSON step description into Python objects. See property 'steps' for more details about the JSON step description. Args: json_str: a JSON string that describes the steps. Returns: Information about steps as a list of dictionaries. Each dictionary may have following properties. (1) per_element_delay - amount of delay for each element in seconds. (2) per_bundle_delay - minimum amount of delay for a given step in seconds. (3) output_records_per_input_record - number of output elements generated for each input element to a step. (4) output_filter_ratio - the probability at which a step may filter out a given element by not producing any output for that element. (5) splittable - if the step should be splittable. (6) initial_splitting_num_bundles - number of bundles initial split if step is splittable. (7) initial_splitting_uneven_chunks - if the bundles should be unevenly-sized (8) disable_liquid_sharding - if liquid sharding should be disabled (9) size_estimate_override - the size estimate or None to use default """ all_steps = [] json_data = json.loads(json_str) for val in json_data: steps = {} steps['per_element_delay'] = ((float(val['per_element_delay_msec']) / 1000) if 'per_element_delay_msec' in val else 0) steps['per_bundle_delay'] = ( float(val['per_bundle_delay_sec']) if 'per_bundle_delay_sec' in val else 0) steps['output_records_per_input_record'] = ( int(val['output_records_per_input_record']) if 'output_records_per_input_record' in val else 1) steps['output_filter_ratio'] = ( float(val['output_filter_ratio']) if 'output_filter_ratio' in val else 0) steps['splittable'] = ( bool(val['splittable']) if 'splittable' in val else False) steps['initial_splitting_num_bundles'] = ( int(val['initial_splitting_num_bundles']) if 'initial_splitting_num_bundles' in val else 8) steps['initial_splitting_uneven_chunks'] = ( bool(val['initial_splitting_uneven_chunks']) if 'initial_splitting_uneven_chunks' in val else False) steps['disable_liquid_sharding'] = ( bool(val['disable_liquid_sharding']) if 'disable_liquid_sharding' in val else False) steps['size_estimate_override'] = ( int(val['size_estimate_override']) if 'size_estimate_override' in val else None) all_steps.append(steps) return all_steps def parse_args(args): """Parses a given set of arguments. Args: args: set of arguments to be passed. Returns: a tuple where first item gives the set of arguments defined and parsed within this method and second item gives the set of unknown arguments. """ parser = argparse.ArgumentParser() parser.add_argument( '--steps', dest='steps', type=_parse_steps, help='A JSON string that gives a list where each entry of the list is ' 'configuration information for a step. Configuration for each step ' 'consists of ' '(1) A float "per_bundle_delay_sec" (in seconds). Defaults to 0.' '(2) A float "per_element_delay_msec" (in milli seconds). ' ' Defaults to 0.' '(3) An integer "output_records_per_input_record". Defaults to 1.' '(4) A float "output_filter_ratio" in the range [0, 1] . ' ' Defaults to 0.' '(5) A bool "splittable" that defaults to false.' '(6) An integer "initial_splitting_num_bundles". Defaults to 8.') parser.add_argument( '--input', dest='input', type=json.loads, help='A JSON string that describes the properties of the SyntheticSource ' 'used by the pipeline. Configuration is similar to Java ' 'SyntheticBoundedInput.' 'Currently supports following properties. ' '(1) An integer "numRecords". ' '(2) An integer "keySize". ' '(3) An integer "valueSize". ' '(4) A tuple "bundleSizeDistribution" with following values. ' ' A string "type". Allowed values are "const" and "zipf". ' ' An float "param". Only used if "type"=="zipf". Must be ' ' larger than 1. ' '(5) An integer "forceNumInitialBundles". ' '(6) An integer "splitPointFrequencyRecords". ' '(7) A tuple "delayDistribution" with following values. ' ' A string "type". Only allowed value is "const". ' ' An integer "const". ' '(8) A string "algorithm". Allowed values are "builtin" for Python ' ' builtin random generator, and "lcg" for the linear congruential ' ' generator equivalent to Java (java.util.Random).') parser.add_argument( '--barrier', dest='barrier', default='shuffle', choices=[ 'shuffle', 'side-input', 'expand-gbk', 'expand-second-output', 'merge-gbk', 'merge-side-input' ], help='Whether to use shuffle as the barrier ' '(as opposed to side inputs).') parser.add_argument( '--output', dest='output', default='', help='Destination to write output.') return parser.parse_known_args(args) def run(argv=None, save_main_session=True): """Runs the workflow.""" known_args, pipeline_args = parse_args(argv) pipeline_options = PipelineOptions(pipeline_args) pipeline_options.view_as(SetupOptions).save_main_session = save_main_session input_info = known_args.input with TestPipeline(options=pipeline_options) as p: source = SyntheticSource(input_info) # pylint: disable=expression-not-assigned barrier = known_args.barrier pc_list = [] num_roots = 2**(len(known_args.steps) - 1) if ( barrier == 'merge-gbk' or barrier == 'merge-side-input') else 1 for read_no in range(num_roots): pc_list.append((p | ('Read %d' % read_no) >> beam.io.Read(source))) for step_no, steps in enumerate(known_args.steps): if step_no != 0: new_pc_list = [] for pc_no, pc in enumerate(pc_list): if barrier == 'shuffle': new_pc_list.append( (pc | ('shuffle %d.%d' % (step_no, pc_no)) >> ShuffleBarrier())) elif barrier == 'side-input': new_pc_list.append(( pc | ('side-input %d.%d' % (step_no, pc_no)) >> SideInputBarrier())) elif barrier == 'expand-gbk': new_pc_list.extend( expand_using_gbk(('expand-gbk %d.%d' % (step_no, pc_no)), pc)) elif barrier == 'expand-second-output': new_pc_list.extend( expand_using_second_output( ('expand-second-output %d.%d' % (step_no, pc_no)), pc)) elif barrier == 'merge-gbk': if pc_no % 2 == 0: new_pc_list.append( merge_using_gbk(('merge-gbk %d.%d' % (step_no, pc_no)), pc, pc_list[pc_no + 1])) else: continue elif barrier == 'merge-side-input': if pc_no % 2 == 0: new_pc_list.append( merge_using_side_input( ('merge-side-input %d.%d' % (step_no, pc_no)), pc, pc_list[pc_no + 1])) else: continue pc_list = new_pc_list new_pc_list = [] for pc_no, pc in enumerate(pc_list): if steps['splittable']: step = get_synthetic_sdf_step( per_element_delay_sec=steps['per_element_delay'], per_bundle_delay_sec=steps['per_bundle_delay'], output_records_per_input_record=steps[ 'output_records_per_input_record'], output_filter_ratio=steps['output_filter_ratio'], initial_splitting_num_bundles=steps[ 'initial_splitting_num_bundles'], initial_splitting_uneven_chunks=steps[ 'initial_splitting_uneven_chunks'], disable_liquid_sharding=steps['disable_liquid_sharding'], size_estimate_override=steps['size_estimate_override']) else: step = SyntheticStep( per_element_delay_sec=steps['per_element_delay'], per_bundle_delay_sec=steps['per_bundle_delay'], output_records_per_input_record=steps[ 'output_records_per_input_record'], output_filter_ratio=steps['output_filter_ratio']) new_pc = pc | 'SyntheticStep %d.%d' % (step_no, pc_no) >> beam.ParDo(step) new_pc_list.append(new_pc) pc_list = new_pc_list if known_args.output: # If an output location is provided we format and write output. if len(pc_list) == 1: ( pc_list[0] | 'FormatOutput' >> beam.Map(lambda elm: (elm[0] + elm[1])) | 'WriteOutput' >> WriteToText(known_args.output)) logging.info('Pipeline run completed.') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) run() class StatefulLoadGenerator(beam.PTransform): """A PTransform for generating random data using Timers API.""" def __init__(self, input_options, num_keys=100): self.num_records = input_options['num_records'] self.key_size = input_options['key_size'] self.value_size = input_options['value_size'] self.num_keys = num_keys @typehints.with_output_types(tuple[bytes, bytes]) class GenerateKeys(beam.DoFn): def __init__(self, num_keys, key_size): self.num_keys = num_keys self.key_size = key_size def process(self, impulse): for _ in range(self.num_keys): key = os.urandom(self.key_size) yield key, b'' class GenerateLoad(beam.DoFn): state_spec = userstate.CombiningValueStateSpec( 'bundles_remaining', combine_fn=sum) timer_spec = userstate.TimerSpec('timer', userstate.TimeDomain.WATERMARK) def __init__(self, num_records_per_key, value_size, bundle_size=1000): self.num_records_per_key = num_records_per_key self.payload = os.urandom(value_size) self.bundle_size = bundle_size self.time_fn = time.time def process( self, _element, records_remaining=beam.DoFn.StateParam(state_spec), timer=beam.DoFn.TimerParam(timer_spec)): records_remaining.add(self.num_records_per_key) timer.set(0) @userstate.on_timer(timer_spec) def process_timer( self, key=beam.DoFn.KeyParam, records_remaining=beam.DoFn.StateParam(state_spec), timer=beam.DoFn.TimerParam(timer_spec)): cur_bundle_size = min(self.bundle_size, records_remaining.read()) for _ in range(cur_bundle_size): records_remaining.add(-1) yield key, self.payload if records_remaining.read() > 0: timer.set(0) def expand(self, pbegin): assert isinstance(pbegin, pvalue.PBegin), ( 'Input to transform must be a PBegin but found %s' % pbegin) return ( pbegin | 'Impulse' >> beam.Impulse() | 'GenerateKeys' >> beam.ParDo( StatefulLoadGenerator.GenerateKeys(self.num_keys, self.key_size)) | 'GenerateLoad' >> beam.ParDo( StatefulLoadGenerator.GenerateLoad( self.num_records // self.num_keys, self.value_size))) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/synthetic_pipeline_test.py0000644000175100017510000002233115206113033024644 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.testing.synthetic_pipeline.""" # pytype: skip-file import glob import json import logging import tempfile import time import unittest import apache_beam as beam from apache_beam.io import source_test_utils from apache_beam.io.restriction_trackers import OffsetRange from apache_beam.testing import synthetic_pipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to try: import numpy # pylint: disable=unused-import except ImportError: NP_INSTALLED = False else: NP_INSTALLED = True def input_spec( num_records, key_size, value_size, bundle_size_distribution_type='const', bundle_size_distribution_param=0, force_initial_num_bundles=0): return { 'numRecords': num_records, 'keySizeBytes': key_size, 'valueSizeBytes': value_size, 'bundleSizeDistribution': { 'type': bundle_size_distribution_type, 'param': bundle_size_distribution_param }, 'forceNumInitialBundles': force_initial_num_bundles, } @unittest.skipIf( not NP_INSTALLED, 'Synthetic source dependencies are not installed') class SyntheticPipelineTest(unittest.TestCase): # pylint: disable=expression-not-assigned def test_synthetic_step_multiplies_output_elements_count(self): with beam.Pipeline() as p: pcoll = p | beam.Create(list(range(10))) | beam.ParDo( synthetic_pipeline.SyntheticStep(0, 0, 10)) assert_that(pcoll | beam.combiners.Count.Globally(), equal_to([100])) def test_minimal_runtime_with_synthetic_step_delay(self): start = time.time() with beam.Pipeline() as p: p | beam.Create(list(range(10))) | beam.ParDo( synthetic_pipeline.SyntheticStep(0, 0.5, 10)) elapsed = time.time() - start self.assertGreaterEqual(elapsed, 0.5, elapsed) def test_synthetic_sdf_step_multiplies_output_elements_count(self): with beam.Pipeline() as p: pcoll = p | beam.Create(list(range(10))) | beam.ParDo( synthetic_pipeline.get_synthetic_sdf_step(0, 0, 10)) assert_that(pcoll | beam.combiners.Count.Globally(), equal_to([100])) def test_minimal_runtime_with_synthetic_sdf_step_bundle_delay(self): start = time.time() with beam.Pipeline() as p: p | beam.Create(list(range(10))) | beam.ParDo( synthetic_pipeline.get_synthetic_sdf_step(0, 0.5, 10)) elapsed = time.time() - start self.assertGreaterEqual(elapsed, 0.5, elapsed) def test_synthetic_step_split_provider(self): provider = synthetic_pipeline.SyntheticSDFStepRestrictionProvider( 5, 2, False, False, None) self.assertEqual( list(provider.split('ab', OffsetRange(2, 15))), [OffsetRange(2, 8), OffsetRange(8, 15)]) self.assertEqual( list(provider.split('ab', OffsetRange(0, 8))), [OffsetRange(0, 4), OffsetRange(4, 8)]) self.assertEqual(list(provider.split('ab', OffsetRange(0, 0))), []) self.assertEqual( list(provider.split('ab', OffsetRange(2, 3))), [OffsetRange(2, 3)]) provider = synthetic_pipeline.SyntheticSDFStepRestrictionProvider( 10, 1, False, False, None) self.assertEqual( list(provider.split('ab', OffsetRange(1, 10))), [OffsetRange(1, 10)]) self.assertEqual(provider.restriction_size('ab', OffsetRange(1, 10)), 9 * 2) provider = synthetic_pipeline.SyntheticSDFStepRestrictionProvider( 10, 3, False, False, None) self.assertEqual( list(provider.split('ab', OffsetRange(1, 10))), [OffsetRange(1, 4), OffsetRange(4, 7), OffsetRange(7, 10)]) self.assertEqual(provider.initial_restriction('a'), OffsetRange(0, 10)) provider = synthetic_pipeline.SyntheticSDFStepRestrictionProvider( 10, 3, False, False, 45) self.assertEqual(provider.restriction_size('ab', OffsetRange(1, 3)), 45) tracker = provider.create_tracker(OffsetRange(1, 6)) tracker.try_claim(1) # Claim to allow splitting. self.assertEqual( tracker.try_split(.5), (OffsetRange(1, 3), OffsetRange(3, 6))) def verify_random_splits(self, provider, restriction, bundles): ranges = list(provider.split('ab', restriction)) prior_stop = restriction.start for r in ranges: self.assertEqual(r.start, prior_stop) prior_stop = r.stop self.assertEqual(prior_stop, restriction.stop) self.assertEqual(len(ranges), bundles) def testSyntheticStepSplitProviderUnevenChunks(self): bundles = 4 provider = synthetic_pipeline.SyntheticSDFStepRestrictionProvider( 5, bundles, True, False, None) self.verify_random_splits(provider, OffsetRange(4, 10), bundles) self.verify_random_splits(provider, OffsetRange(4, 4), 0) self.verify_random_splits(provider, OffsetRange(0, 1), 1) self.verify_random_splits(provider, OffsetRange(0, bundles - 2), bundles) def test_synthetic_step_split_provider_no_liquid_sharding(self): # Verify Liquid Sharding Works provider = synthetic_pipeline.SyntheticSDFStepRestrictionProvider( 5, 5, True, False, None) tracker = provider.create_tracker(OffsetRange(1, 6)) tracker.try_claim(2) self.assertEqual( tracker.try_split(.5), (OffsetRange(1, 4), OffsetRange(4, 6))) # Verify No Liquid Sharding provider = synthetic_pipeline.SyntheticSDFStepRestrictionProvider( 5, 5, True, True, None) tracker = provider.create_tracker(OffsetRange(1, 6)) tracker.try_claim(2) self.assertEqual(tracker.try_split(3), None) def test_synthetic_source(self): def assert_size(element, expected_size): assert len(element) == expected_size with beam.Pipeline() as p: pcoll = ( p | beam.io.Read( synthetic_pipeline.SyntheticSource(input_spec(300, 5, 15)))) (pcoll | beam.Map(lambda elm: elm[0]) | 'key' >> beam.Map(assert_size, 5)) ( pcoll | beam.Map(lambda elm: elm[1]) | 'value' >> beam.Map(assert_size, 15)) assert_that(pcoll | beam.combiners.Count.Globally(), equal_to([300])) def test_synthetic_source_split_even(self): source = synthetic_pipeline.SyntheticSource( input_spec(1000, 1, 1, 'const', 0)) splits = source.split(100) sources_info = [(split.source, split.start_position, split.stop_position) for split in splits] self.assertEqual(20, len(sources_info)) source_test_utils.assert_sources_equal_reference_source( (source, None, None), sources_info) def test_synthetic_source_split_uneven(self): source = synthetic_pipeline.SyntheticSource( input_spec(1000, 1, 1, 'zipf', 3, 10)) splits = source.split(100) sources_info = [(split.source, split.start_position, split.stop_position) for split in splits] self.assertEqual(10, len(sources_info)) source_test_utils.assert_sources_equal_reference_source( (source, None, None), sources_info) def test_split_at_fraction(self): source = synthetic_pipeline.SyntheticSource(input_spec(10, 1, 1)) source_test_utils.assert_split_at_fraction_exhaustive(source) source_test_utils.assert_split_at_fraction_fails(source, 5, 0.3) source_test_utils.assert_split_at_fraction_succeeds_and_consistent( source, 1, 0.3) def run_pipeline(self, barrier, writes_output=True): steps = [{ 'per_element_delay': 1 }, { 'per_element_delay': 1, 'splittable': True }] args = [ '--barrier=%s' % barrier, '--runner=DirectRunner', '--steps=%s' % json.dumps(steps), '--input=%s' % json.dumps(input_spec(10, 1, 1)) ] if writes_output: output_location = tempfile.NamedTemporaryFile().name args.append('--output=%s' % output_location) synthetic_pipeline.run(args, save_main_session=False) # Verify output if writes_output: read_output = [] for file_name in glob.glob(output_location + '*'): with open(file_name, 'rb') as f: read_output.extend(f.read().splitlines()) self.assertEqual(10, len(read_output)) def test_pipeline_shuffle(self): self.run_pipeline('shuffle') def test_pipeline_side_input(self): self.run_pipeline('side-input') def test_pipeline_expand_gbk(self): self.run_pipeline('expand-gbk', False) def test_pipeline_expand_side_output(self): self.run_pipeline('expand-second-output', False) def test_pipeline_merge_gbk(self): self.run_pipeline('merge-gbk') def test_pipeline_merge_side_input(self): self.run_pipeline('merge-side-input') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/test_pipeline.py0000644000175100017510000002062415206113033022555 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test Pipeline, a wrapper of Pipeline for test purpose""" # pytype: skip-file import argparse import shlex import sys from unittest import SkipTest from apache_beam.internal import pickler from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.pipeline import Pipeline from apache_beam.runners.runner import PipelineState __all__ = [ 'TestPipeline', ] class TestPipeline(Pipeline): """:class:`TestPipeline` class is used inside of Beam tests that can be configured to run against pipeline runner. It has a functionality to parse arguments from command line and build pipeline options for tests who runs against a pipeline runner and utilizes resources of the pipeline runner. Those test functions are recommended to be tagged by ``@pytest.mark.it_validatesrunner`` annotation. In order to configure the test with customized pipeline options from command line, system argument ``--test-pipeline-options`` can be used to obtains a list of pipeline options. If no options specified, default value will be used. For example, use following command line to execute all ValidatesRunner tests:: pytest -m it_validatesrunner \\ --test-pipeline-options="--runner=DirectRunner \\ --job_name=myJobName \\ --num_workers=1" For example, use assert_that for test validation:: with TestPipeline() as pipeline: pcoll = ... assert_that(pcoll, equal_to(...)) """ # Command line options read in by pytest. # If this is not None, will use as default value for --test-pipeline-options. pytest_test_pipeline_options = None def __init__( self, runner=None, options=None, argv=None, is_integration_test=False, blocking=True, additional_pipeline_args=None, display_data=None, timeout=None): """Initialize a pipeline object for test. Args: runner (~apache_beam.runners.runner.PipelineRunner): An object of type :class:`~apache_beam.runners.runner.PipelineRunner` that will be used to execute the pipeline. For registered runners, the runner name can be specified, otherwise a runner object must be supplied. options (~apache_beam.options.pipeline_options.PipelineOptions): A configured :class:`~apache_beam.options.pipeline_options.PipelineOptions` object containing arguments that should be used for running the pipeline job. argv (List[str]): A list of arguments (such as :data:`sys.argv`) to be used for building a :class:`~apache_beam.options.pipeline_options.PipelineOptions` object. This will only be used if argument **options** is :data:`None`. is_integration_test (bool): :data:`True` if the test is an integration test, :data:`False` otherwise. blocking (bool): Run method will wait until pipeline execution is completed. additional_pipeline_args (List[str]): additional pipeline arguments to be included when construction the pipeline options object. display_data (Dict[str, Any]): a dictionary of static data associated with this pipeline that can be displayed when it runs. timeout (int, optional): Milliseconds to wait for the pipeline to finish. If the timeout is reached, an AssertionError is raised. Raises: ValueError: if either the runner or options argument is not of the expected type. """ self.is_integration_test = is_integration_test self.not_use_test_runner_api = False additional_pipeline_args = additional_pipeline_args or [] self.options_list = ( self._parse_test_option_args(argv) + additional_pipeline_args) self.blocking = blocking self.timeout = timeout if options is None: options = PipelineOptions(self.options_list) super().__init__(runner, options, display_data=display_data) def run(self, test_runner_api=True): result = super().run( test_runner_api=( False if self.not_use_test_runner_api else test_runner_api)) if self.blocking: state = result.wait_until_finish(duration=self.timeout) assert state in (PipelineState.DONE, PipelineState.CANCELLED), \ "Pipeline execution failed." return result def get_pipeline_options(self): return self._options def _parse_test_option_args(self, argv): """Parse value of command line argument: --test-pipeline-options to get pipeline options. Args: argv: An iterable of command line arguments to be used. If not specified then sys.argv will be used as input for parsing arguments. Returns: An argument list of options that can be parsed by argparser or directly build a pipeline option. """ parser = argparse.ArgumentParser() parser.add_argument( '--test-pipeline-options', type=str, action='store', help='only run tests providing service options') parser.add_argument( '--not-use-test-runner-api', action='store_true', default=False, help='whether not to use test-runner-api') known, unused_argv = parser.parse_known_args(argv) test_pipeline_options = known.test_pipeline_options or \ TestPipeline.pytest_test_pipeline_options if self.is_integration_test and not test_pipeline_options: # Skip integration test when argument '--test-pipeline-options' is not # specified since nose calls integration tests when runs unit test by # 'setup.py test'. raise SkipTest( 'IT is skipped because --test-pipeline-options ' 'is not specified') self.not_use_test_runner_api = known.not_use_test_runner_api return shlex.split(test_pipeline_options) \ if test_pipeline_options else [] @classmethod def get_options_list(cls, argv=None): parser = argparse.ArgumentParser() parser.add_argument( '--test-pipeline-options', type=str, action='store', help='Pipeline options for the test') known, _ = parser.parse_known_args(argv if argv is not None else sys.argv) opts = known.test_pipeline_options or getattr( cls, 'pytest_test_pipeline_options', None) return shlex.split(opts) if opts else [] def get_full_options_as_args(self, **extra_opts): """Get full pipeline options as an argument list. Append extra pipeline options to existing option list if provided. Test verifier (if contains in extra options) should be pickled before appending, and will be unpickled later in the TestRunner. """ options = list(self.options_list) for k, v in extra_opts.items(): if not v: continue elif isinstance(v, bool) and v: options.append('--%s' % k) elif 'matcher' in k: options.append('--%s=%s' % (k, pickler.dumps(v).decode())) else: options.append('--%s=%s' % (k, v)) return options def get_option(self, opt_name, bool_option=False): """Get a pipeline option value by name Args: opt_name: The name of the pipeline option. Returns: None if option is not found in existing option list which is generated by parsing value of argument `test-pipeline-options`. """ parser = argparse.ArgumentParser() opt_name = opt_name[:2] if opt_name[:2] == '--' else opt_name # Option name should start with '--' when it's used for parsing. if bool_option: parser.add_argument('--' + opt_name, action='store_true') else: parser.add_argument('--' + opt_name, type=str, action='store') known, _ = parser.parse_known_args(self.options_list) return getattr(known, opt_name) if hasattr(known, opt_name) else None ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/test_pipeline_test.py0000644000175100017510000001102115206113033023603 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit test for the TestPipeline class""" # pytype: skip-file import logging import unittest import mock from hamcrest.core.assert_that import assert_that as hc_assert_that from hamcrest.core.base_matcher import BaseMatcher from apache_beam.internal import pickler from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline # A simple matcher that is ued for testing extra options appending. class SimpleMatcher(BaseMatcher): def _matches(self, item): return True class TestPipelineTest(unittest.TestCase): TEST_CASE = { 'options': ['--test-pipeline-options', '--job=mockJob --male --age=1'], 'expected_list': ['--job=mockJob', '--male', '--age=1'], 'expected_dict': { 'job': 'mockJob', 'male': True, 'age': 1 } } # Used for testing pipeline option creation. class TestParsingOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument('--job', action='store', help='mock job') parser.add_argument('--male', action='store_true', help='mock gender') parser.add_argument('--age', action='store', type=int, help='mock age') def test_option_args_parsing(self): test_pipeline = TestPipeline(argv=self.TEST_CASE['options']) self.assertListEqual( sorted(test_pipeline.get_full_options_as_args()), sorted(self.TEST_CASE['expected_list'])) def test_empty_option_args_parsing(self): test_pipeline = TestPipeline() self.assertListEqual([], test_pipeline.get_full_options_as_args()) def test_create_test_pipeline_options(self): test_pipeline = TestPipeline(argv=self.TEST_CASE['options']) test_options = PipelineOptions(test_pipeline.get_full_options_as_args()) self.assertLessEqual( self.TEST_CASE['expected_dict'].items(), test_options.get_all_options().items()) EXTRA_OPT_CASES = [{ 'options': { 'name': 'Mark' }, 'expected': ['--name=Mark'] }, { 'options': { 'student': True }, 'expected': ['--student'] }, { 'options': { 'student': False }, 'expected': [] }, { 'options': { 'name': 'Mark', 'student': True }, 'expected': ['--name=Mark', '--student'] }] def test_append_extra_options(self): test_pipeline = TestPipeline() for case in self.EXTRA_OPT_CASES: opt_list = test_pipeline.get_full_options_as_args(**case['options']) self.assertListEqual(sorted(opt_list), sorted(case['expected'])) def test_append_verifier_in_extra_opt(self): extra_opt = {'matcher': SimpleMatcher()} opt_list = TestPipeline().get_full_options_as_args(**extra_opt) _, value = opt_list[0].split('=', 1) matcher = pickler.loads(value) self.assertTrue(isinstance(matcher, BaseMatcher)) hc_assert_that(None, matcher) def test_get_option(self): name, value = ('job', 'mockJob') test_pipeline = TestPipeline() test_pipeline.options_list = ['--%s=%s' % (name, value)] self.assertEqual(test_pipeline.get_option(name), value) def test_skip_IT(self): with TestPipeline(is_integration_test=True) as _: # Note that this will never be reached since it should be skipped above. pass self.fail() @mock.patch('apache_beam.testing.test_pipeline.Pipeline.run', autospec=True) def test_not_use_test_runner_api(self, mock_run): with TestPipeline(argv=['--not-use-test-runner-api'], blocking=False) as test_pipeline: pass mock_run.assert_called_once_with(test_pipeline, test_runner_api=False) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/test_stream.py0000644000175100017510000006431215206113033022245 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Provides TestStream for verifying streaming runner semantics. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file from abc import ABCMeta from abc import abstractmethod from enum import Enum from functools import total_ordering import apache_beam as beam from apache_beam import coders from apache_beam import pvalue from apache_beam.portability import common_urns from apache_beam.portability.api import beam_interactive_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.portability.api import endpoints_pb2 from apache_beam.transforms import PTransform from apache_beam.transforms import core from apache_beam.transforms import window from apache_beam.transforms.timeutil import TimeDomain from apache_beam.transforms.userstate import TimerSpec from apache_beam.transforms.userstate import on_timer from apache_beam.transforms.window import TimestampedValue from apache_beam.utils import timestamp from apache_beam.utils.timestamp import MIN_TIMESTAMP from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.windowed_value import WindowedValue __all__ = [ 'Event', 'ElementEvent', 'WatermarkEvent', 'ProcessingTimeEvent', 'TestStream', ] @total_ordering class Event(metaclass=ABCMeta): # type: ignore[misc] """Test stream event to be emitted during execution of a TestStream.""" @abstractmethod def __eq__(self, other): raise NotImplementedError @abstractmethod def __hash__(self): raise NotImplementedError @abstractmethod def __lt__(self, other): raise NotImplementedError @abstractmethod def to_runner_api(self, element_coder): raise NotImplementedError @staticmethod def from_runner_api(proto, element_coder): if proto.HasField('element_event'): event = proto.element_event tag = None if event.tag == 'None' else event.tag return ElementEvent([ TimestampedValue( element_coder.decode(tv.encoded_element), Timestamp(micros=1000 * tv.timestamp)) for tv in proto.element_event.elements ], tag=tag) # yapf: disable elif proto.HasField('watermark_event'): event = proto.watermark_event tag = None if event.tag == 'None' else event.tag return WatermarkEvent( Timestamp(micros=1000 * proto.watermark_event.new_watermark), tag=tag) elif proto.HasField('processing_time_event'): return ProcessingTimeEvent( timestamp.Duration( micros=1000 * proto.processing_time_event.advance_duration)) else: raise ValueError( 'Unknown TestStream Event type: %s' % proto.WhichOneof('event')) class ElementEvent(Event): """Element-producing test stream event.""" def __init__(self, timestamped_values, tag=None): self.timestamped_values = timestamped_values self.tag = tag def __eq__(self, other): if not isinstance(other, ElementEvent): return False return ( self.timestamped_values == other.timestamped_values and self.tag == other.tag) def __hash__(self): return hash(self.timestamped_values) def __lt__(self, other): if not isinstance(other, ElementEvent): raise TypeError return self.timestamped_values < other.timestamped_values def to_runner_api(self, element_coder): tag = 'None' if self.tag is None else self.tag return beam_runner_api_pb2.TestStreamPayload.Event( element_event=beam_runner_api_pb2.TestStreamPayload.Event.AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=element_coder.encode(tv.value), timestamp=tv.timestamp.micros // 1000) for tv in self.timestamped_values ], tag=tag)) def __repr__(self): return 'ElementEvent: <{}, {}>'.format([(e.value, e.timestamp) for e in self.timestamped_values], self.tag) class WatermarkEvent(Event): """Watermark-advancing test stream event.""" def __init__(self, new_watermark, tag=None): self.new_watermark = Timestamp.of(new_watermark) self.tag = tag def __eq__(self, other): if not isinstance(other, WatermarkEvent): return False return self.new_watermark == other.new_watermark and self.tag == other.tag def __hash__(self): return hash(str(self.new_watermark) + str(self.tag)) def __lt__(self, other): if not isinstance(other, WatermarkEvent): raise TypeError return self.new_watermark < other.new_watermark def to_runner_api(self, unused_element_coder): tag = 'None' if self.tag is None else self.tag # Assert that no precision is lost. assert self.new_watermark.micros % 1000 == 0 return beam_runner_api_pb2.TestStreamPayload.Event( watermark_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceWatermark( new_watermark=self.new_watermark.micros // 1000, tag=tag)) def __repr__(self): return 'WatermarkEvent: <{}, {}>'.format(self.new_watermark, self.tag) class ProcessingTimeEvent(Event): """Processing time-advancing test stream event.""" def __init__(self, advance_by): self.advance_by = Duration.of(advance_by) def __eq__(self, other): if not isinstance(other, ProcessingTimeEvent): return False return self.advance_by == other.advance_by def __hash__(self): return hash(self.advance_by) def __lt__(self, other): if not isinstance(other, ProcessingTimeEvent): raise TypeError return self.advance_by < other.advance_by def to_runner_api(self, unused_element_coder): assert self.advance_by.micros % 1000 == 0 return beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=beam_runner_api_pb2.TestStreamPayload.Event. AdvanceProcessingTime(advance_duration=self.advance_by.micros // 1000)) def __repr__(self): return 'ProcessingTimeEvent: <{}>'.format(self.advance_by) class WindowedValueHolderMeta(type): """A metaclass that overrides the isinstance check for WindowedValueHolder. Python does a quick test for exact match. If an instance is exactly of type WindowedValueHolder, the overridden isinstance check is omitted. The override is needed because WindowedValueHolder elements encoded then decoded become Row elements. """ def __instancecheck__(cls, other): """Checks if a beam.Row typed instance is a WindowedValueHolder. """ return ( isinstance(other, beam.Row) and hasattr(other, 'windowed_value') and hasattr(other, 'urn') and isinstance(other.windowed_value, WindowedValue) and other.urn == common_urns.coders.ROW.urn) class WindowedValueHolder(beam.Row, metaclass=WindowedValueHolderMeta): """A class that holds a WindowedValue. This is a special class that can be used by the runner that implements the TestStream as a signal that the underlying value should be unreified to the specified window. """ # Register WindowedValueHolder to always use RowCoder. coders.registry.register_coder(WindowedValueHolderMeta, coders.RowCoder) def __init__(self, windowed_value): assert isinstance(windowed_value, WindowedValue), ( 'WindowedValueHolder can only hold %s type. Instead, %s is given.') % ( WindowedValue, windowed_value) super().__init__( **{ 'windowed_value': windowed_value, 'urn': common_urns.coders.ROW.urn }) @classmethod def from_row(cls, row): """Converts a beam.Row typed instance to WindowedValueHolder. """ if isinstance(row, WindowedValueHolder): return WindowedValueHolder(row.windowed_value) assert isinstance(row, beam.Row), 'The given row %s must be a %s type' % ( row, beam.Row) assert hasattr(row, 'windowed_value'), ( 'The given %s must have a windowed_value attribute.') % row assert isinstance(row.windowed_value, WindowedValue), ( 'The windowed_value attribute of %s must be a %s type') % ( row, WindowedValue) class TestStream(PTransform): """Test stream that generates events on an unbounded PCollection of elements. Each event emits elements, advances the watermark or advances the processing time. After all of the specified elements are emitted, ceases to produce output. Applying the PTransform will return a single PCollection if only the default output or only one output tag has been used. Otherwise a dictionary of output names to PCollections will be returned. """ def __init__( self, coder=coders.FastPrimitivesCoder(), events=None, output_tags=None, endpoint=None): """ Args: coder: (apache_beam.Coder) the coder to encode/decode elements. events: (List[Event]) a list of instructions for the TestStream to execute. If specified, the events tags must exist in the output_tags. output_tags: (List[str]) Initial set of outputs. If no event references an output tag, no output will be produced for that tag. endpoint: (str) a URL locating a TestStreamService. """ super().__init__() assert coder is not None self.coder = coder self.watermarks = {None: timestamp.MIN_TIMESTAMP} self.output_tags = set(output_tags) if output_tags else set() self._events = [] if events is None else list(events) self._endpoint = endpoint event_tags = set( e.tag for e in self._events if isinstance(e, (WatermarkEvent, ElementEvent))) assert event_tags.issubset(self.output_tags), \ '{} is not a subset of {}'.format(event_tags, output_tags) assert not (self._events and self._endpoint), \ 'Only either events or an endpoint can be given at once.' def get_windowing(self, unused_inputs): return core.Windowing(window.GlobalWindows()) def _infer_output_coder(self, input_type=None, input_coder=None): return self.coder def expand(self, pbegin): assert isinstance(pbegin, pvalue.PBegin) self.pipeline = pbegin.pipeline if not self.output_tags: self.output_tags = {None} # For backwards compatibility return a single PCollection. if self.output_tags == {None}: return pvalue.PCollection( self.pipeline, is_bounded=False, tag=list(self.output_tags)[0]) return { tag: pvalue.PCollection(self.pipeline, is_bounded=False, tag=tag) for tag in self.output_tags } def _add(self, event): if isinstance(event, ElementEvent): for tv in event.timestamped_values: assert tv.timestamp < timestamp.MAX_TIMESTAMP, ( 'Element timestamp must be before timestamp.MAX_TIMESTAMP.') elif isinstance(event, WatermarkEvent): if event.tag not in self.watermarks: self.watermarks[event.tag] = timestamp.MIN_TIMESTAMP assert event.new_watermark > self.watermarks[event.tag], ( 'Watermark must strictly-monotonically advance.') self.watermarks[event.tag] = event.new_watermark elif isinstance(event, ProcessingTimeEvent): assert event.advance_by > 0, ( 'Must advance processing time by positive amount.') else: raise ValueError('Unknown event: %s' % event) self._events.append(event) def add_elements(self, elements, tag=None, event_timestamp=None): """Add elements to the TestStream. Elements added to the TestStream will be produced during pipeline execution. These elements can be TimestampedValue, WindowedValue or raw unwrapped elements that are serializable using the TestStream's specified Coder. When a TimestampedValue or a WindowedValue element is used, the timestamp of the TimestampedValue or WindowedValue will be the timestamp of the produced element; otherwise, the current watermark timestamp will be used for that element. The windows of a given WindowedValue are ignored by the TestStream. """ self.output_tags.add(tag) timestamped_values = [] if tag not in self.watermarks: self.watermarks[tag] = timestamp.MIN_TIMESTAMP for element in elements: if isinstance(element, TimestampedValue): timestamped_values.append(element) elif isinstance(element, WindowedValue): # Drop windows for elements in test stream. timestamped_values.append( TimestampedValue(element.value, element.timestamp)) else: # Add elements with timestamp equal to current watermark. if event_timestamp is None: event_timestamp = self.watermarks[tag] timestamped_values.append(TimestampedValue(element, event_timestamp)) self._add(ElementEvent(timestamped_values, tag)) return self def advance_watermark_to(self, new_watermark, tag=None): """Advance the watermark to a given Unix timestamp. The Unix timestamp value used must be later than the previous watermark value and should be given as an int, float or utils.timestamp.Timestamp object. """ self.output_tags.add(tag) self._add(WatermarkEvent(new_watermark, tag)) return self def advance_watermark_to_infinity(self, tag=None): """Advance the watermark to the end of time, completing this TestStream.""" self.advance_watermark_to(timestamp.MAX_TIMESTAMP, tag) return self def advance_processing_time(self, advance_by): """Advance the current processing time by a given duration in seconds. The duration must be a positive second duration and should be given as an int, float or utils.timestamp.Duration object. """ self._add(ProcessingTimeEvent(advance_by)) return self def to_runner_api_parameter(self, context): # Sort the output tags so that the order is deterministic and we are able # to test equality on a roundtrip through the to/from proto apis. return ( common_urns.primitives.TEST_STREAM.urn, beam_runner_api_pb2.TestStreamPayload( coder_id=context.coders.get_id(self.coder), events=[e.to_runner_api(self.coder) for e in self._events], endpoint=endpoints_pb2.ApiServiceDescriptor(url=self._endpoint))) @staticmethod @PTransform.register_urn( common_urns.primitives.TEST_STREAM.urn, beam_runner_api_pb2.TestStreamPayload) def from_runner_api_parameter(ptransform, payload, context): coder = context.coders.get_by_id(payload.coder_id) output_tags = set( None if k == 'None' else k for k in ptransform.outputs.keys()) return TestStream( coder=coder, events=[Event.from_runner_api(e, coder) for e in payload.events], output_tags=output_tags, endpoint=payload.endpoint.url) class TimingInfo(object): def __init__(self, processing_time, watermark): self._processing_time = Timestamp.of(processing_time) self._watermark = Timestamp.of(watermark) @property def processing_time(self): return self._processing_time @property def watermark(self): return self._watermark def __repr__(self): return '({}, {})'.format(self.processing_time, self.watermark) class PairWithTiming(PTransform): """Pairs the input element with timing information. Input: element; output: KV(element, timing information) Where timing information := (processing time, watermark) This is used in the ReverseTestStream implementation to replay watermark advancements. """ URN = "beam:transform:pair_with_timing:v1" def expand(self, pcoll): return pvalue.PCollection.from_(pcoll) class OutputFormat(Enum): TEST_STREAM_EVENTS = 1 TEST_STREAM_FILE_RECORDS = 2 SERIALIZED_TEST_STREAM_FILE_RECORDS = 3 class ReverseTestStream(PTransform): """A Transform that can create TestStream events from a stream of elements. This currently assumes that this the pipeline being run on a single machine and elements come in order and are outputted in the same order that they came in. """ def __init__( self, sample_resolution_sec, output_tag, coder=None, output_format=None): self._sample_resolution_sec = sample_resolution_sec self._output_tag = output_tag self._output_format = output_format if output_format \ else OutputFormat.TEST_STREAM_EVENTS self._coder = coder if coder else beam.coders.FastPrimitivesCoder() def expand(self, pcoll): ret = ( pcoll | beam.WindowInto(beam.window.GlobalWindows()) # First get the initial timing information. This will be used to start # the periodic timers which will generate processing time and watermark # advancements every `sample_resolution_sec`. | 'initial timing' >> PairWithTiming() # Next, map every element to the same key so that only a single timer is # started for this given ReverseTestStream. | 'first key' >> beam.Map(lambda x: (0, x)) # Next, pass-through each element which will be paired with its timing # info in the next step. Also, start the periodic timers. We use timers # in this situation to capture watermark advancements that occur when # there are no elements being produced upstream. | beam.ParDo( _TimingEventGenerator( output_tag=self._output_tag, sample_resolution_sec=self._sample_resolution_sec)) # Next, retrieve the timing information for watermark events that were # generated in the previous step. This is because elements generated # through the timers don't have their timing information yet. | 'timing info for watermarks' >> PairWithTiming() # Re-key to the same key to keep global state. | 'second key' >> beam.Map(lambda x: (0, x)) # Format the events properly. | beam.ParDo(_TestStreamFormatter(self._coder, self._output_format))) if self._output_format == OutputFormat.SERIALIZED_TEST_STREAM_FILE_RECORDS: def serializer(e): return e.SerializeToString() ret = ret | 'serializer' >> beam.Map(serializer) return ret class _TimingEventGenerator(beam.DoFn): """Generates ProcessingTimeEvents and WatermarkEvents at a regular cadence. The runner keeps the state of the clock (which may be faked) and the watermarks, which are inaccessible to SDKs. This DoFn generates ProcessingTimeEvents and WatermarkEvents at a specified sampling rate to capture any clock or watermark advancements between elements. """ # Used to return the initial timing information. EXECUTE_ONCE_STATE = beam.transforms.userstate.BagStateSpec( name='execute_once_state', coder=beam.coders.FastPrimitivesCoder()) # A processing time timer in an infinite loop that generates the events that # will be paired with the TimingInfo from the runner. TIMING_SAMPLER = TimerSpec('timing_sampler', TimeDomain.REAL_TIME) def __init__(self, output_tag, sample_resolution_sec=0.1): self._output_tag = output_tag self._sample_resolution_sec = sample_resolution_sec @on_timer(TIMING_SAMPLER) def on_timing_sampler( self, timestamp=beam.DoFn.TimestampParam, window=beam.DoFn.WindowParam, timing_sampler=beam.DoFn.TimerParam(TIMING_SAMPLER)): """Yields an unbounded stream of ProcessingTimeEvents and WatermarkEvents. The returned events will be paired with the TimingInfo. This loop's only purpose is to generate these events even when there are no elements. """ next_sample_time = (timestamp.micros * 1e-6) + self._sample_resolution_sec timing_sampler.set(next_sample_time) # Generate two events, the delta since the last sample and a place-holder # WatermarkEvent. This is a placeholder because we can't otherwise add the # watermark from the runner to the event. yield ProcessingTimeEvent(self._sample_resolution_sec) yield WatermarkEvent(MIN_TIMESTAMP) def process( self, e, timestamp=beam.DoFn.TimestampParam, window=beam.DoFn.WindowParam, timing_sampler=beam.DoFn.TimerParam(TIMING_SAMPLER), execute_once_state=beam.DoFn.StateParam(EXECUTE_ONCE_STATE)): _, (element, timing_info) = e # Only set the timers once and only send the header once. first_time = next(execute_once_state.read(), True) if first_time: # Generate the initial timing events. execute_once_state.add(False) now_sec = timing_info.processing_time.micros * 1e-6 timing_sampler.set(now_sec + self._sample_resolution_sec) # Here we capture the initial time offset and initial watermark. This is # where we emit the TestStreamFileHeader. yield beam_interactive_api_pb2.TestStreamFileHeader(tag=self._output_tag) yield ProcessingTimeEvent( Duration(micros=timing_info.processing_time.micros)) yield WatermarkEvent(MIN_TIMESTAMP) yield element class _TestStreamFormatter(beam.DoFn): """Formats the events to the specified output format. """ # In order to generate the processing time deltas, we need to keep track of # the previous clock time we got from the runner. PREV_SAMPLE_TIME_STATE = beam.transforms.userstate.BagStateSpec( name='prev_sample_time_state', coder=beam.coders.FastPrimitivesCoder()) def __init__(self, coder, output_format): self._coder = coder self._output_format = output_format def start_bundle(self): self.elements = [] self.timing_events = [] self.header = None def finish_bundle(self): """Outputs all the buffered elements. """ if self._output_format == OutputFormat.TEST_STREAM_EVENTS: return self._output_as_events() return self._output_as_records() def process( self, e, timestamp=beam.DoFn.TimestampParam, prev_sample_time_state=beam.DoFn.StateParam(PREV_SAMPLE_TIME_STATE)): """Buffers elements until the end of the bundle. This buffers elements instead of emitting them immediately to keep elements that come in the same bundle to be outputted in the same bundle. """ _, (element, timing_info) = e if isinstance(element, beam_interactive_api_pb2.TestStreamFileHeader): self.header = element elif isinstance(element, WatermarkEvent): # WatermarkEvents come in with a watermark of MIN_TIMESTAMP. Fill in the # correct watermark from the runner here. element.new_watermark = timing_info.watermark.micros if element not in self.timing_events: self.timing_events.append(element) elif isinstance(element, ProcessingTimeEvent): # Because the runner holds the clock, calculate the processing time delta # here. The TestStream may have faked out the clock, and thus the # delta calculated in the SDK with time.time() will be wrong. prev_sample = next(prev_sample_time_state.read(), Timestamp()) prev_sample_time_state.clear() prev_sample_time_state.add(timing_info.processing_time) advance_by = timing_info.processing_time - prev_sample element.advance_by = advance_by self.timing_events.append(element) else: self.elements.append(TimestampedValue(element, timestamp)) def _output_as_events(self): """Outputs buffered elements as TestStream events. """ if self.timing_events: yield WindowedValue( self.timing_events, timestamp=0, windows=[beam.window.GlobalWindow()]) if self.elements: yield WindowedValue([ElementEvent(self.elements)], timestamp=0, windows=[beam.window.GlobalWindow()]) def _output_as_records(self): """Outputs buffered elements as TestStreamFileRecords. """ if self.header: yield WindowedValue( self.header, timestamp=0, windows=[beam.window.GlobalWindow()]) if self.timing_events: timing_events = self._timing_events_to_records(self.timing_events) for r in timing_events: yield WindowedValue( r, timestamp=0, windows=[beam.window.GlobalWindow()]) if self.elements: elements = self._elements_to_record(self.elements) yield WindowedValue( elements, timestamp=0, windows=[beam.window.GlobalWindow()]) def _timing_events_to_records(self, timing_events): """Returns given timing_events as TestStreamFileRecords. """ records = [] for e in self.timing_events: if isinstance(e, ProcessingTimeEvent): processing_time_event = beam_runner_api_pb2.\ TestStreamPayload.Event.AdvanceProcessingTime( advance_duration=e.advance_by.micros) records.append( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload.Event( processing_time_event=processing_time_event))) elif isinstance(e, WatermarkEvent): watermark_event = beam_runner_api_pb2.\ TestStreamPayload.Event.AdvanceWatermark( new_watermark=int(e.new_watermark)) records.append( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload.Event( watermark_event=watermark_event))) return records def _elements_to_record(self, elements): """Returns elements as TestStreamFileRecords. """ elements = [] for tv in self.elements: element_timestamp = tv.timestamp.micros element = beam_runner_api_pb2.TestStreamPayload.TimestampedElement( encoded_element=self._coder.encode(tv.value), timestamp=element_timestamp) elements.append(element) element_event = beam_runner_api_pb2.TestStreamPayload.Event.AddElements( elements=elements) return beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload.Event( element_event=element_event)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/test_stream_it_test.py0000644000175100017510000002152215206113033023774 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration tests for the test_stream module.""" # pytype: skip-file import unittest from functools import wraps import pytest import apache_beam as beam from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.testing.util import equal_to_per_window from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import TimestampedValue from apache_beam.utils import timestamp from apache_beam.utils.timestamp import Timestamp def supported(runners): if not isinstance(runners, list): runners = [runners] def inner(fn): @wraps(fn) def wrapped(self): if self.runner_name not in runners: self.skipTest( 'The "{}", does not support the TestStream transform. ' 'Supported runners: {}'.format(self.runner_name, runners)) else: return fn(self) return wrapped return inner class TestStreamIntegrationTests(unittest.TestCase): @classmethod def setUpClass(cls): cls.test_pipeline = TestPipeline(is_integration_test=True) cls.args = cls.test_pipeline.get_full_options_as_args() cls.runner_name = type(cls.test_pipeline.runner).__name__ cls.project = cls.test_pipeline.get_option('project') @supported(['DirectRunner', 'SwitchingDirectRunner']) @pytest.mark.it_postcommit def test_basic_execution(self): test_stream = ( TestStream().advance_watermark_to(10).add_elements([ 'a', 'b', 'c' ]).advance_watermark_to(20).add_elements(['d']).add_elements([ 'e' ]).advance_processing_time(10).advance_watermark_to(300).add_elements([ TimestampedValue('late', 12) ]).add_elements([TimestampedValue('last', 310) ]).advance_watermark_to_infinity()) class RecordFn(beam.DoFn): def process( self, element=beam.DoFn.ElementParam, timestamp=beam.DoFn.TimestampParam): yield (element, timestamp) with beam.Pipeline(argv=self.args) as p: my_record_fn = RecordFn() records = p | test_stream | beam.ParDo(my_record_fn) assert_that( records, equal_to([ ('a', timestamp.Timestamp(10)), ('b', timestamp.Timestamp(10)), ('c', timestamp.Timestamp(10)), ('d', timestamp.Timestamp(20)), ('e', timestamp.Timestamp(20)), ('late', timestamp.Timestamp(12)), ('last', timestamp.Timestamp(310)), ])) @supported(['DirectRunner', 'SwitchingDirectRunner']) @pytest.mark.it_postcommit def test_multiple_outputs(self): """Tests that the TestStream supports emitting to multiple PCollections.""" letters_elements = [ TimestampedValue('a', 6), TimestampedValue('b', 7), TimestampedValue('c', 8), ] numbers_elements = [ TimestampedValue('1', 11), TimestampedValue('2', 12), TimestampedValue('3', 13), ] test_stream = ( TestStream().advance_watermark_to(5, tag='letters').add_elements( letters_elements, tag='letters').advance_watermark_to(10, tag='numbers').add_elements( numbers_elements, tag='numbers')) class RecordFn(beam.DoFn): def process( self, element=beam.DoFn.ElementParam, timestamp=beam.DoFn.TimestampParam): yield (element, timestamp) options = StandardOptions(streaming=True) p = TestPipeline(is_integration_test=True, options=options) main = p | test_stream letters = main['letters'] | 'record letters' >> beam.ParDo(RecordFn()) numbers = main['numbers'] | 'record numbers' >> beam.ParDo(RecordFn()) assert_that( letters, equal_to([('a', Timestamp(6)), ('b', Timestamp(7)), ('c', Timestamp(8))]), label='assert letters') assert_that( numbers, equal_to([('1', Timestamp(11)), ('2', Timestamp(12)), ('3', Timestamp(13))]), label='assert numbers') p.run() @supported(['DirectRunner', 'SwitchingDirectRunner']) @pytest.mark.it_postcommit def test_multiple_outputs_with_watermark_advancement(self): """Tests that the TestStream can independently control output watermarks.""" # Purposely set the watermark of numbers to 20 then letters to 5 to test # that the watermark advancement is per PCollection. # # This creates two PCollections, (a, b, c) and (1, 2, 3). These will be # emitted at different times so that they will have different windows. The # watermark advancement is checked by checking their windows. If the # watermark does not advance, then the windows will be [-inf, -inf). If the # windows do not advance separately, then the PCollections will both # windowed in [15, 30). letters_elements = [ TimestampedValue('a', 6), TimestampedValue('b', 7), TimestampedValue('c', 8), ] numbers_elements = [ TimestampedValue('1', 21), TimestampedValue('2', 22), TimestampedValue('3', 23), ] test_stream = ( TestStream().advance_watermark_to( 0, tag='letters').advance_watermark_to( 0, tag='numbers').advance_watermark_to( 20, tag='numbers').advance_watermark_to( 5, tag='letters').add_elements( letters_elements, tag='letters').advance_watermark_to( 10, tag='letters').add_elements( numbers_elements, tag='numbers').advance_watermark_to( 30, tag='numbers')) options = StandardOptions(streaming=True) p = TestPipeline(is_integration_test=True, options=options) main = p | test_stream # Use an AfterWatermark trigger with an early firing to test that the # watermark is advancing properly and that the element is being emitted in # the correct window. letters = ( main['letters'] | 'letter windows' >> beam.WindowInto( FixedWindows(15), trigger=trigger.AfterWatermark(early=trigger.AfterCount(1)), accumulation_mode=trigger.AccumulationMode.DISCARDING) | 'letter with key' >> beam.Map(lambda x: ('k', x)) | 'letter gbk' >> beam.GroupByKey()) numbers = ( main['numbers'] | 'number windows' >> beam.WindowInto( FixedWindows(15), trigger=trigger.AfterWatermark(early=trigger.AfterCount(1)), accumulation_mode=trigger.AccumulationMode.DISCARDING) | 'number with key' >> beam.Map(lambda x: ('k', x)) | 'number gbk' >> beam.GroupByKey()) # The letters were emitted when the watermark was at 5, thus we expect to # see the elements in the [0, 15) window. We used an early trigger to make # sure that the ON_TIME empty pane was also emitted with a TestStream. # This pane has no data because of the early trigger causes the elements to # fire before the end of the window and because the accumulation mode # discards any data after the trigger fired. expected_letters = { window.IntervalWindow(0, 15): [ ('k', ['a', 'b', 'c']), ('k', []), ], } # Same here, except the numbers were emitted at watermark = 20, thus they # are in the [15, 30) window. expected_numbers = { window.IntervalWindow(15, 30): [ ('k', ['1', '2', '3']), ('k', []), ], } assert_that( letters, equal_to_per_window(expected_letters), label='letters assert per window') assert_that( numbers, equal_to_per_window(expected_numbers), label='numbers assert per window') p.run() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/test_stream_service.py0000644000175100017510000000612515206113033023763 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file from concurrent.futures import ThreadPoolExecutor import grpc from apache_beam.portability.api import beam_interactive_api_pb2_grpc class TestStreamServiceController( beam_interactive_api_pb2_grpc.TestStreamServiceServicer): """A server that streams TestStreamPayload.Events from a single EventRequest. This server is used as a way for TestStreams to receive events from file. """ def __init__(self, reader, endpoint=None, exception_handler=None): self._server = grpc.server(ThreadPoolExecutor(max_workers=10)) self._server_started = False self._server_stopped = False if endpoint: self.endpoint = endpoint self._server.add_insecure_port(self.endpoint) else: port = self._server.add_insecure_port('localhost:0') self.endpoint = 'localhost:{}'.format(port) beam_interactive_api_pb2_grpc.add_TestStreamServiceServicer_to_server( self, self._server) self._reader = reader self._exception_handler = exception_handler if not self._exception_handler: self._exception_handler = lambda _: False def start(self): # A server can only be started if never started and never stopped before. if self._server_started or self._server_stopped: return self._server_started = True self._server.start() def stop(self): # A server can only be stopped if already started and never stopped before. if not self._server_started or self._server_stopped: return self._server_started = False self._server_stopped = True self._server.stop(0) # This was introduced in grpcio 1.24 and might be gone in the future. Keep # this check in case the runtime is on a older, current or future grpcio. if hasattr(self._server, 'wait_for_termination'): self._server.wait_for_termination() def Events(self, request, context): """Streams back all of the events from the streaming cache.""" # TODO(srohde): Once we get rid of the CacheManager, get rid of this 'full' # label. tags = [None if tag == 'None' else tag for tag in request.output_ids] try: reader = self._reader.read_multiple([('full', tag) for tag in tags]) while True: e = next(reader) yield e except StopIteration: pass except Exception as e: if not self._exception_handler(e): raise e ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/test_stream_service_test.py0000644000175100017510000001417315206113033025024 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest from unittest.mock import patch import grpc from apache_beam.portability.api import beam_interactive_api_pb2 from apache_beam.portability.api import beam_interactive_api_pb2_grpc from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.testing.test_stream_service import TestStreamServiceController # Nose automatically detects tests if they match a regex. Here, it mistakens # these protos as tests. For more info see the Nose docs at: # https://nose.readthedocs.io/en/latest/writing_tests.html beam_runner_api_pb2.TestStreamPayload.__test__ = False beam_interactive_api_pb2.TestStreamFileHeader.__test__ = False beam_interactive_api_pb2.TestStreamFileRecord.__test__ = False class EventsReader: def __init__(self, expected_key): self._expected_key = expected_key def read_multiple(self, keys): if keys != self._expected_key: raise ValueError( 'Expected key ({}) is not argument({})'.format( self._expected_key, keys)) for i in range(10): e = beam_runner_api_pb2.TestStreamPayload.Event() e.element_event.elements.append( beam_runner_api_pb2.TestStreamPayload.TimestampedElement(timestamp=i)) yield e EXPECTED_KEY = 'key' EXPECTED_KEYS = [EXPECTED_KEY] class TestStreamServiceTest(unittest.TestCase): def setUp(self): self.controller = TestStreamServiceController( EventsReader(expected_key=[('full', EXPECTED_KEY)])) self.controller.start() channel = grpc.insecure_channel(self.controller.endpoint) self.stub = beam_interactive_api_pb2_grpc.TestStreamServiceStub(channel) def tearDown(self): self.controller.stop() def test_normal_run(self): r = self.stub.Events( beam_interactive_api_pb2.EventsRequest(output_ids=EXPECTED_KEYS)) events = [e for e in r] expected_events = [ e for e in EventsReader( expected_key=[EXPECTED_KEYS]).read_multiple([EXPECTED_KEYS]) ] self.assertEqual(events, expected_events) def test_multiple_sessions(self): resp_a = self.stub.Events( beam_interactive_api_pb2.EventsRequest(output_ids=EXPECTED_KEYS)) resp_b = self.stub.Events( beam_interactive_api_pb2.EventsRequest(output_ids=EXPECTED_KEYS)) events_a = [] events_b = [] done = False while not done: a_is_done = False b_is_done = False try: events_a.append(next(resp_a)) except StopIteration: a_is_done = True try: events_b.append(next(resp_b)) except StopIteration: b_is_done = True done = a_is_done and b_is_done expected_events = [ e for e in EventsReader( expected_key=[EXPECTED_KEYS]).read_multiple([EXPECTED_KEYS]) ] self.assertEqual(events_a, expected_events) self.assertEqual(events_b, expected_events) class TestStreamServiceStartStopTest(unittest.TestCase): # Weak internal use needs to be explicitly imported. from grpc import _server def setUp(self): self.controller = TestStreamServiceController( EventsReader(expected_key=[('full', EXPECTED_KEY)])) self.assertFalse(self.controller._server_started) self.assertFalse(self.controller._server_stopped) def tearDown(self): self.controller.stop() def test_start_when_never_started(self): with patch.object(self._server._Server, 'start', wraps=self.controller._server.start) as mock_start: self.controller.start() mock_start.assert_called_once() self.assertTrue(self.controller._server_started) self.assertFalse(self.controller._server_stopped) def test_start_noop_when_already_started(self): with patch.object(self._server._Server, 'start', wraps=self.controller._server.start) as mock_start: self.controller.start() mock_start.assert_called_once() self.controller.start() mock_start.assert_called_once() def test_start_noop_when_already_stopped(self): with patch.object(self._server._Server, 'start', wraps=self.controller._server.start) as mock_start: self.controller.start() self.controller.stop() mock_start.assert_called_once() self.controller.start() mock_start.assert_called_once() def test_stop_noop_when_not_started(self): with patch.object(self._server._Server, 'stop', wraps=self.controller._server.stop) as mock_stop: self.controller.stop() mock_stop.assert_not_called() def test_stop_when_already_started(self): with patch.object(self._server._Server, 'stop', wraps=self.controller._server.stop) as mock_stop: self.controller.start() mock_stop.assert_not_called() self.controller.stop() mock_stop.assert_called_once() self.assertFalse(self.controller._server_started) self.assertTrue(self.controller._server_stopped) def test_stop_noop_when_already_stopped(self): with patch.object(self._server._Server, 'stop', wraps=self.controller._server.stop) as mock_stop: self.controller.start() self.controller.stop() mock_stop.assert_called_once() self.controller.stop() mock_stop.assert_called_once() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/test_stream_test.py0000644000175100017510000012010115206113033023271 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the test_stream module.""" # pytype: skip-file import unittest import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import TypeOptions from apache_beam.portability import common_urns from apache_beam.portability.api import beam_interactive_api_pb2 from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import ElementEvent from apache_beam.testing.test_stream import OutputFormat from apache_beam.testing.test_stream import ProcessingTimeEvent from apache_beam.testing.test_stream import ReverseTestStream from apache_beam.testing.test_stream import TestStream from apache_beam.testing.test_stream import WatermarkEvent from apache_beam.testing.test_stream import WindowedValueHolder from apache_beam.testing.test_stream_service import TestStreamServiceController from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.testing.util import equal_to_per_window from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import TimestampedValue from apache_beam.utils import timestamp from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.windowed_value import PaneInfo from apache_beam.utils.windowed_value import PaneInfoTiming from apache_beam.utils.windowed_value import WindowedValue class TestStreamTest(unittest.TestCase): def test_basic_test_stream(self): test_stream = (TestStream() .advance_watermark_to(0) .add_elements([ 'a', WindowedValue('b', 3, []), TimestampedValue('c', 6)]) .advance_processing_time(10) .advance_watermark_to(8) .add_elements(['d']) .advance_watermark_to_infinity()) # yapf: disable self.assertEqual( test_stream._events, [ WatermarkEvent(0), ElementEvent([ TimestampedValue('a', 0), TimestampedValue('b', 3), TimestampedValue('c', 6), ]), ProcessingTimeEvent(10), WatermarkEvent(8), ElementEvent([ TimestampedValue('d', 8), ]), WatermarkEvent(timestamp.MAX_TIMESTAMP), ]) def test_test_stream_errors(self): with self.assertRaises( AssertionError, msg=('Watermark must strictly-monotonically advance.')): _ = (TestStream().advance_watermark_to(5).advance_watermark_to(4)) with self.assertRaises( AssertionError, msg=('Must advance processing time by positive amount.')): _ = (TestStream().advance_processing_time(-1)) with self.assertRaises( AssertionError, msg=('Element timestamp must be before timestamp.MAX_TIMESTAMP.')): _ = ( TestStream().add_elements( [TimestampedValue('a', timestamp.MAX_TIMESTAMP)])) def test_basic_execution(self): test_stream = (TestStream() .advance_watermark_to(10) .add_elements(['a', 'b', 'c']) .advance_watermark_to(20) .add_elements(['d']) .add_elements(['e']) .advance_processing_time(10) .advance_watermark_to(300) .add_elements([TimestampedValue('late', 12)]) .add_elements([TimestampedValue('last', 310)]) .advance_watermark_to_infinity()) # yapf: disable class RecordFn(beam.DoFn): def process( self, element=beam.DoFn.ElementParam, timestamp=beam.DoFn.TimestampParam): yield (element, timestamp) options = PipelineOptions() options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: my_record_fn = RecordFn() records = p | test_stream | beam.ParDo(my_record_fn) assert_that( records, equal_to([ ('a', timestamp.Timestamp(10)), ('b', timestamp.Timestamp(10)), ('c', timestamp.Timestamp(10)), ('d', timestamp.Timestamp(20)), ('e', timestamp.Timestamp(20)), ('late', timestamp.Timestamp(12)), ('last', timestamp.Timestamp(310)), ])) def test_multiple_outputs(self): """Tests that the TestStream supports emitting to multiple PCollections.""" letters_elements = [ TimestampedValue('a', 6), TimestampedValue('b', 7), TimestampedValue('c', 8), ] numbers_elements = [ TimestampedValue('1', 11), TimestampedValue('2', 12), TimestampedValue('3', 13), ] test_stream = (TestStream() .advance_watermark_to(5, tag='letters') .add_elements(letters_elements, tag='letters') .advance_watermark_to(10, tag='numbers') .add_elements(numbers_elements, tag='numbers')) # yapf: disable class RecordFn(beam.DoFn): def process( self, element=beam.DoFn.ElementParam, timestamp=beam.DoFn.TimestampParam): yield (element, timestamp) options = StandardOptions(streaming=True) p = TestPipeline(options=options) main = p | test_stream letters = main['letters'] | 'record letters' >> beam.ParDo(RecordFn()) numbers = main['numbers'] | 'record numbers' >> beam.ParDo(RecordFn()) assert_that( letters, equal_to([('a', Timestamp(6)), ('b', Timestamp(7)), ('c', Timestamp(8))]), label='assert letters') assert_that( numbers, equal_to([('1', Timestamp(11)), ('2', Timestamp(12)), ('3', Timestamp(13))]), label='assert numbers') p.run() def test_multiple_outputs_with_watermark_advancement(self): """Tests that the TestStream can independently control output watermarks.""" # Purposely set the watermark of numbers to 20 then letters to 5 to test # that the watermark advancement is per PCollection. # # This creates two PCollections, (a, b, c) and (1, 2, 3). These will be # emitted at different times so that they will have different windows. The # watermark advancement is checked by checking their windows. If the # watermark does not advance, then the windows will be [-inf, -inf). If the # windows do not advance separately, then the PCollections will both # windowed in [15, 30). letters_elements = [ TimestampedValue('a', 6), TimestampedValue('b', 7), TimestampedValue('c', 8), ] numbers_elements = [ TimestampedValue('1', 21), TimestampedValue('2', 22), TimestampedValue('3', 23), ] test_stream = (TestStream() .advance_watermark_to(0, tag='letters') .advance_watermark_to(0, tag='numbers') .advance_watermark_to(20, tag='numbers') .advance_watermark_to(5, tag='letters') .add_elements(letters_elements, tag='letters') .advance_watermark_to(10, tag='letters') .add_elements(numbers_elements, tag='numbers') .advance_watermark_to(30, tag='numbers')) # yapf: disable options = StandardOptions(streaming=True) p = TestPipeline(options=options) main = p | test_stream # Use an AfterWatermark trigger with an early firing to test that the # watermark is advancing properly and that the element is being emitted in # the correct window. letters = ( main['letters'] | 'letter windows' >> beam.WindowInto( FixedWindows(15), trigger=trigger.AfterWatermark(early=trigger.AfterCount(1)), accumulation_mode=trigger.AccumulationMode.DISCARDING) | 'letter with key' >> beam.Map(lambda x: ('k', x)) | 'letter gbk' >> beam.GroupByKey()) numbers = ( main['numbers'] | 'number windows' >> beam.WindowInto( FixedWindows(15), trigger=trigger.AfterWatermark(early=trigger.AfterCount(1)), accumulation_mode=trigger.AccumulationMode.DISCARDING) | 'number with key' >> beam.Map(lambda x: ('k', x)) | 'number gbk' >> beam.GroupByKey()) # The letters were emitted when the watermark was at 5, thus we expect to # see the elements in the [0, 15) window. We used an early trigger to make # sure that the ON_TIME empty pane was also emitted with a TestStream. # This pane has no data because of the early trigger causes the elements to # fire before the end of the window and because the accumulation mode # discards any data after the trigger fired. expected_letters = { window.IntervalWindow(0, 15): [ ('k', ['a', 'b', 'c']), ('k', []), ], } # Same here, except the numbers were emitted at watermark = 20, thus they # are in the [15, 30) window. expected_numbers = { window.IntervalWindow(15, 30): [ ('k', ['1', '2', '3']), ('k', []), ], } assert_that( letters, equal_to_per_window(expected_letters), label='letters assert per window') assert_that( numbers, equal_to_per_window(expected_numbers), label='numbers assert per window') p.run() def test_dicts_not_interpreted_as_windowed_values(self): test_stream = (TestStream() .advance_processing_time(10) .advance_watermark_to(10) .add_elements([{'a': 0, 'b': 1, 'c': 2}]) .advance_watermark_to_infinity()) # yapf: disable class RecordFn(beam.DoFn): def process( self, element=beam.DoFn.ElementParam, timestamp=beam.DoFn.TimestampParam): yield (element, timestamp) options = PipelineOptions() options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: my_record_fn = RecordFn() records = p | test_stream | beam.ParDo(my_record_fn) assert_that( records, equal_to([ ({ 'a': 0, 'b': 1, 'c': 2 }, timestamp.Timestamp(10)), ])) def test_windowed_values_interpreted_correctly(self): windowed_value = WindowedValueHolder( WindowedValue( 'a', Timestamp(5), [beam.window.IntervalWindow(5, 10)], PaneInfo(True, True, PaneInfoTiming.ON_TIME, 0, 0))) test_stream = (TestStream() .advance_processing_time(10) .advance_watermark_to(10) .add_elements([windowed_value]) .advance_watermark_to_infinity()) # yapf: disable class RecordFn(beam.DoFn): def process( self, element=beam.DoFn.ElementParam, timestamp=beam.DoFn.TimestampParam, window=beam.DoFn.WindowParam): yield (element, timestamp, window) options = PipelineOptions() options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: my_record_fn = RecordFn() records = p | test_stream | beam.ParDo(my_record_fn) assert_that( records, equal_to([ ('a', timestamp.Timestamp(5), beam.window.IntervalWindow(5, 10)), ])) def test_instance_check_windowed_value_holder(self): windowed_value = WindowedValue( 'a', Timestamp(5), [beam.window.IntervalWindow(5, 10)], PaneInfo(True, True, PaneInfoTiming.ON_TIME, 0, 0)) self.assertTrue( isinstance(WindowedValueHolder(windowed_value), WindowedValueHolder)) self.assertTrue( isinstance( beam.Row( windowed_value=windowed_value, urn=common_urns.coders.ROW.urn), WindowedValueHolder)) self.assertFalse( isinstance( beam.Row(windowed_value=windowed_value), WindowedValueHolder)) self.assertFalse(isinstance(windowed_value, WindowedValueHolder)) self.assertFalse( isinstance(beam.Row(x=windowed_value), WindowedValueHolder)) self.assertFalse( isinstance(beam.Row(windowed_value=1), WindowedValueHolder)) def test_gbk_execution_no_triggers(self): test_stream = ( TestStream().advance_watermark_to(10).add_elements([ 'a', 'b', 'c' ]).advance_watermark_to(20).add_elements(['d']).add_elements([ 'e' ]).advance_processing_time(10).advance_watermark_to(300).add_elements([ TimestampedValue('late', 12) ]).add_elements([TimestampedValue('last', 310) ]).advance_watermark_to_infinity()) options = PipelineOptions() options.view_as(StandardOptions).streaming = True p = TestPipeline(options=options) records = ( p | test_stream | beam.WindowInto(FixedWindows(15), allowed_lateness=300) | beam.Map(lambda x: ('k', x)) | beam.GroupByKey()) # TODO(https://github.com/apache/beam/issues/18441): timestamp assignment # for elements from a GBK should respect the TimestampCombiner. The test # below should also verify the timestamps of the outputted elements once # this is implemented. # assert per window expected_window_to_elements = { window.IntervalWindow(0, 15): [ ('k', ['a', 'b', 'c']), ('k', ['late']), ], window.IntervalWindow(15, 30): [ ('k', ['d', 'e']), ], window.IntervalWindow(300, 315): [ ('k', ['last']), ], } assert_that( records, equal_to_per_window(expected_window_to_elements), label='assert per window') p.run() def test_gbk_execution_after_watermark_trigger(self): test_stream = (TestStream() .advance_watermark_to(10) .add_elements([TimestampedValue('a', 11)]) .advance_watermark_to(20) .add_elements([TimestampedValue('b', 21)]) .advance_watermark_to_infinity()) # yapf: disable options = PipelineOptions() options.view_as(StandardOptions).streaming = True p = TestPipeline(options=options) records = ( p # pylint: disable=unused-variable | test_stream | beam.WindowInto( FixedWindows(15), trigger=trigger.AfterWatermark(early=trigger.AfterCount(1)), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.Map(lambda x: ('k', x)) | beam.GroupByKey()) # TODO(https://github.com/apache/beam/issues/18441): timestamp assignment # for elements from a GBK should respect the TimestampCombiner. The test # below should also verify the timestamps of the outputted elements once # this is implemented. # assert per window expected_window_to_elements = { window.IntervalWindow(0, 15): [('k', ['a']), ('k', [])], window.IntervalWindow(15, 30): [('k', ['b']), ('k', [])], } assert_that( records, equal_to_per_window(expected_window_to_elements), label='assert per window') p.run() def test_gbk_execution_after_processing_trigger_fired(self): """Advance TestClock to (X + delta) and see the pipeline does finish.""" # TODO(mariagh): Add test_gbk_execution_after_processing_trigger_unfired # Advance TestClock to (X + delta) and see the pipeline does finish # Possibly to the framework trigger_transcripts.yaml test_stream = (TestStream() .advance_watermark_to(10) .add_elements(['a']) .advance_processing_time(5.1) .advance_watermark_to_infinity()) # yapf: disable options = PipelineOptions() options.view_as(StandardOptions).streaming = True options.view_as(TypeOptions).allow_unsafe_triggers = True p = TestPipeline(options=options) records = ( p | test_stream | beam.WindowInto( beam.window.FixedWindows(15), trigger=trigger.AfterProcessingTime(5), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.Map(lambda x: ('k', x)) | beam.GroupByKey()) # TODO(https://github.com/apache/beam/issues/18441): timestamp assignment # for elements from a GBK should respect the TimestampCombiner. The test # below should also verify the timestamps of the outputted elements once # this is implemented. expected_window_to_elements = { window.IntervalWindow(0, 15): [('k', ['a'])], } assert_that( records, equal_to_per_window(expected_window_to_elements), label='assert per window') p.run() def test_basic_execution_batch_sideinputs(self): options = PipelineOptions() options.view_as(StandardOptions).streaming = True p = TestPipeline(options=options) main_stream = (p | 'main TestStream' >> TestStream() .advance_watermark_to(10) .add_elements(['e']) .advance_watermark_to_infinity()) # yapf: disable side = ( p | beam.Create([2, 1, 4]) | beam.Map(lambda t: window.TimestampedValue(t, t))) class RecordFn(beam.DoFn): def process( self, elm=beam.DoFn.ElementParam, ts=beam.DoFn.TimestampParam, side=beam.DoFn.SideInputParam): yield (elm, ts, sorted(side)) records = ( main_stream # pylint: disable=unused-variable | beam.ParDo(RecordFn(), beam.pvalue.AsList(side))) assert_that(records, equal_to([('e', Timestamp(10), [1, 2, 4])])) p.run() def test_basic_execution_sideinputs(self): options = PipelineOptions() options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: test_stream = (p | TestStream() .advance_watermark_to(0, tag='side') .advance_watermark_to(10, tag='main') .add_elements(['e'], tag='main') .add_elements([window.TimestampedValue(2, 2)], tag='side') .add_elements([window.TimestampedValue(1, 1)], tag='side') .add_elements([window.TimestampedValue(7, 7)], tag='side') .add_elements([window.TimestampedValue(4, 4)], tag='side') ) # yapf: disable main_stream = test_stream['main'] side_stream = test_stream['side'] class RecordFn(beam.DoFn): def process( self, elm=beam.DoFn.ElementParam, ts=beam.DoFn.TimestampParam, side=beam.DoFn.SideInputParam): yield (elm, ts, side) records = ( main_stream # pylint: disable=unused-variable | beam.ParDo(RecordFn(), beam.pvalue.AsList(side_stream))) assert_that(records, equal_to([('e', Timestamp(10), [2, 1, 7, 4])])) def test_basic_execution_batch_sideinputs_fixed_windows(self): options = PipelineOptions() options.view_as(StandardOptions).streaming = True p = TestPipeline(options=options) main_stream = ( p | 'main TestStream' >> TestStream().advance_watermark_to(2).add_elements( ['a']).advance_watermark_to(4).add_elements( ['b']).advance_watermark_to_infinity() | 'main window' >> beam.WindowInto(window.FixedWindows(1))) side = ( p | beam.Create([2, 1, 4]) | beam.Map(lambda t: window.TimestampedValue(t, t)) | beam.WindowInto(window.FixedWindows(2))) class RecordFn(beam.DoFn): def process( self, elm=beam.DoFn.ElementParam, ts=beam.DoFn.TimestampParam, side=beam.DoFn.SideInputParam): yield (elm, ts, side) records = ( main_stream # pylint: disable=unused-variable | beam.ParDo(RecordFn(), beam.pvalue.AsList(side))) # assert per window expected_window_to_elements = { window.IntervalWindow(2, 3): [('a', Timestamp(2), [2])], window.IntervalWindow(4, 5): [('b', Timestamp(4), [4])] } assert_that( records, equal_to_per_window(expected_window_to_elements), label='assert per window') p.run() def test_basic_execution_sideinputs_fixed_windows(self): options = PipelineOptions() options.view_as(StandardOptions).streaming = True p = TestPipeline(options=options) test_stream = (p | TestStream() .advance_watermark_to(12, tag='side') .add_elements([window.TimestampedValue('s1', 10)], tag='side') .advance_watermark_to(20, tag='side') .add_elements([window.TimestampedValue('s2', 20)], tag='side') .advance_watermark_to(9, tag='main') .add_elements(['a1', 'a2', 'a3', 'a4'], tag='main') .add_elements(['b'], tag='main') .advance_watermark_to(18, tag='main') .add_elements('c', tag='main') ) # yapf: disable main_stream = ( test_stream['main'] | 'main windowInto' >> beam.WindowInto(window.FixedWindows(1))) side_stream = ( test_stream['side'] | 'side windowInto' >> beam.WindowInto(window.FixedWindows(3))) class RecordFn(beam.DoFn): def process( self, elm=beam.DoFn.ElementParam, ts=beam.DoFn.TimestampParam, side=beam.DoFn.SideInputParam): yield (elm, ts, side) records = ( main_stream # pylint: disable=unused-variable | beam.ParDo(RecordFn(), beam.pvalue.AsList(side_stream))) # assert per window expected_window_to_elements = { window.IntervalWindow(9, 10): [ ('a1', Timestamp(9), ['s1']), ('a2', Timestamp(9), ['s1']), ('a3', Timestamp(9), ['s1']), ('a4', Timestamp(9), ['s1']), ('b', Timestamp(9), ['s1']) ], window.IntervalWindow(18, 19): [('c', Timestamp(18), ['s2'])], } assert_that( records, equal_to_per_window(expected_window_to_elements), label='assert per window') p.run() def test_roundtrip_proto(self): test_stream = (TestStream() .advance_processing_time(1) .advance_watermark_to(2) .add_elements([1, 2, 3])) # yapf: disable p = TestPipeline(options=StandardOptions(streaming=True)) p | test_stream pipeline_proto, context = p.to_runner_api(return_context=True) for t in pipeline_proto.components.transforms.values(): if t.spec.urn == common_urns.primitives.TEST_STREAM.urn: test_stream_proto = t self.assertTrue(test_stream_proto) roundtrip_test_stream = TestStream().from_runner_api( test_stream_proto, context) self.assertListEqual(test_stream._events, roundtrip_test_stream._events) self.assertSetEqual( test_stream.output_tags, roundtrip_test_stream.output_tags) self.assertEqual(test_stream.coder, roundtrip_test_stream.coder) def test_roundtrip_proto_multi(self): test_stream = (TestStream() .advance_processing_time(1) .advance_watermark_to(2, tag='a') .advance_watermark_to(3, tag='b') .add_elements([1, 2, 3], tag='a') .add_elements([4, 5, 6], tag='b')) # yapf: disable options = StandardOptions(streaming=True) p = TestPipeline(options=options) p | test_stream pipeline_proto, context = p.to_runner_api(return_context=True) for t in pipeline_proto.components.transforms.values(): if t.spec.urn == common_urns.primitives.TEST_STREAM.urn: test_stream_proto = t self.assertTrue(test_stream_proto) roundtrip_test_stream = TestStream().from_runner_api( test_stream_proto, context) self.assertListEqual(test_stream._events, roundtrip_test_stream._events) self.assertSetEqual( test_stream.output_tags, roundtrip_test_stream.output_tags) self.assertEqual(test_stream.coder, roundtrip_test_stream.coder) def test_basic_execution_with_service(self): """Tests that the TestStream can correctly read from an RPC service. """ coder = beam.coders.FastPrimitivesCoder() test_stream_events = (TestStream(coder=coder) .advance_watermark_to(10000) .add_elements(['a', 'b', 'c']) .advance_watermark_to(20000) .add_elements(['d']) .add_elements(['e']) .advance_processing_time(10) .advance_watermark_to(300000) .add_elements([TimestampedValue('late', 12000)]) .add_elements([TimestampedValue('last', 310000)]) .advance_watermark_to_infinity())._events # yapf: disable test_stream_proto_events = [ e.to_runner_api(coder) for e in test_stream_events ] class InMemoryEventReader: def read_multiple(self, unused_keys): for e in test_stream_proto_events: yield e service = TestStreamServiceController(reader=InMemoryEventReader()) service.start() test_stream = TestStream(coder=coder, endpoint=service.endpoint) class RecordFn(beam.DoFn): def process( self, element=beam.DoFn.ElementParam, timestamp=beam.DoFn.TimestampParam): yield (element, timestamp) options = StandardOptions(streaming=True) p = TestPipeline(options=options) my_record_fn = RecordFn() records = p | test_stream | beam.ParDo(my_record_fn) assert_that( records, equal_to([ ('a', timestamp.Timestamp(10)), ('b', timestamp.Timestamp(10)), ('c', timestamp.Timestamp(10)), ('d', timestamp.Timestamp(20)), ('e', timestamp.Timestamp(20)), ('late', timestamp.Timestamp(12)), ('last', timestamp.Timestamp(310)), ])) p.run() class ReverseTestStreamTest(unittest.TestCase): def test_basic_execution(self): test_stream = (TestStream() .advance_watermark_to(0) .advance_processing_time(5) .add_elements(['a', 'b', 'c']) .advance_watermark_to(2) .advance_processing_time(1) .advance_watermark_to(4) .advance_processing_time(1) .advance_watermark_to(6) .advance_processing_time(1) .advance_watermark_to(8) .advance_processing_time(1) .advance_watermark_to(10) .advance_processing_time(1) .add_elements([TimestampedValue('1', 15), TimestampedValue('2', 15), TimestampedValue('3', 15)])) # yapf: disable options = StandardOptions(streaming=True) p = TestPipeline(options=options) records = ( p | test_stream | ReverseTestStream(sample_resolution_sec=1, output_tag=None)) assert_that( records, equal_to_per_window({ beam.window.GlobalWindow(): [ [ProcessingTimeEvent(5), WatermarkEvent(0)], [ ElementEvent([ TimestampedValue('a', 0), TimestampedValue('b', 0), TimestampedValue('c', 0) ]) ], [ProcessingTimeEvent(1), WatermarkEvent(2000000)], [ProcessingTimeEvent(1), WatermarkEvent(4000000)], [ProcessingTimeEvent(1), WatermarkEvent(6000000)], [ProcessingTimeEvent(1), WatermarkEvent(8000000)], [ProcessingTimeEvent(1), WatermarkEvent(10000000)], [ ElementEvent([ TimestampedValue('1', 15), TimestampedValue('2', 15), TimestampedValue('3', 15) ]) ], ], })) p.run() def test_windowing(self): test_stream = (TestStream() .advance_watermark_to(0) .add_elements(['a', 'b', 'c']) .advance_processing_time(1) .advance_processing_time(1) .advance_processing_time(1) .advance_processing_time(1) .advance_processing_time(1) .advance_watermark_to(5) .add_elements(['1', '2', '3']) .advance_processing_time(1) .advance_watermark_to(6) .advance_processing_time(1) .advance_watermark_to(7) .advance_processing_time(1) .advance_watermark_to(8) .advance_processing_time(1) .advance_watermark_to(9) .advance_processing_time(1) .advance_watermark_to(10) .advance_processing_time(1) .advance_watermark_to(11) .advance_processing_time(1) .advance_watermark_to(12) .advance_processing_time(1) .advance_watermark_to(13) .advance_processing_time(1) .advance_watermark_to(14) .advance_processing_time(1) .advance_watermark_to(15) .advance_processing_time(1) ) # yapf: disable options = StandardOptions(streaming=True) p = TestPipeline(options=options) records = ( p | test_stream | 'letter windows' >> beam.WindowInto( FixedWindows(5), accumulation_mode=trigger.AccumulationMode.DISCARDING) | 'letter with key' >> beam.Map(lambda x: ('k', x)) | 'letter gbk' >> beam.GroupByKey() | ReverseTestStream(sample_resolution_sec=1, output_tag=None)) assert_that( records, equal_to_per_window({ beam.window.GlobalWindow(): [ [ProcessingTimeEvent(5), WatermarkEvent(4999998)], [ ElementEvent( [TimestampedValue(('k', ['a', 'b', 'c']), 4.999999)]) ], [ProcessingTimeEvent(1), WatermarkEvent(5000000)], [ProcessingTimeEvent(1), WatermarkEvent(6000000)], [ProcessingTimeEvent(1), WatermarkEvent(7000000)], [ProcessingTimeEvent(1), WatermarkEvent(8000000)], [ProcessingTimeEvent(1), WatermarkEvent(9000000)], [ ElementEvent( [TimestampedValue(('k', ['1', '2', '3']), 9.999999)]) ], [ProcessingTimeEvent(1), WatermarkEvent(10000000)], [ProcessingTimeEvent(1), WatermarkEvent(11000000)], [ProcessingTimeEvent(1), WatermarkEvent(12000000)], [ProcessingTimeEvent(1), WatermarkEvent(13000000)], [ProcessingTimeEvent(1), WatermarkEvent(14000000)], [ProcessingTimeEvent(1), WatermarkEvent(15000000)], ], })) p.run() def test_basic_execution_in_records_format(self): test_stream = (TestStream() .advance_watermark_to(0) .advance_processing_time(5) .add_elements(['a', 'b', 'c']) .advance_watermark_to(2) .advance_processing_time(1) .advance_watermark_to(4) .advance_processing_time(1) .advance_watermark_to(6) .advance_processing_time(1) .advance_watermark_to(8) .advance_processing_time(1) .advance_watermark_to(10) .advance_processing_time(1) .add_elements([TimestampedValue('1', 15), TimestampedValue('2', 15), TimestampedValue('3', 15)])) # yapf: disable options = StandardOptions(streaming=True) p = TestPipeline(options=options) coder = beam.coders.FastPrimitivesCoder() records = ( p | test_stream | ReverseTestStream( sample_resolution_sec=1, coder=coder, output_format=OutputFormat.TEST_STREAM_FILE_RECORDS, output_tag=None) | 'stringify' >> beam.Map(str)) assert_that( records, equal_to_per_window({ beam.window.GlobalWindow(): [ str(beam_interactive_api_pb2.TestStreamFileHeader()), str( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload. Event( processing_time_event=beam_runner_api_pb2. TestStreamPayload.Event.AdvanceProcessingTime( advance_duration=5000000)))), str( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload. Event( watermark_event=beam_runner_api_pb2. TestStreamPayload.Event.AdvanceWatermark( new_watermark=0)))), str( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload. Event( element_event=beam_runner_api_pb2.TestStreamPayload. Event.AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload. TimestampedElement( encoded_element=coder.encode('a'), timestamp=0), beam_runner_api_pb2.TestStreamPayload. TimestampedElement( encoded_element=coder.encode('b'), timestamp=0), beam_runner_api_pb2.TestStreamPayload. TimestampedElement( encoded_element=coder.encode('c'), timestamp=0), ])))), str( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload. Event( watermark_event=beam_runner_api_pb2. TestStreamPayload.Event.AdvanceWatermark( new_watermark=2000000)))), str( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload. Event( processing_time_event=beam_runner_api_pb2. TestStreamPayload.Event.AdvanceProcessingTime( advance_duration=1000000)))), str( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload. Event( watermark_event=beam_runner_api_pb2. TestStreamPayload.Event.AdvanceWatermark( new_watermark=4000000)))), str( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload. Event( processing_time_event=beam_runner_api_pb2. TestStreamPayload.Event.AdvanceProcessingTime( advance_duration=1000000)))), str( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload. Event( watermark_event=beam_runner_api_pb2. TestStreamPayload.Event.AdvanceWatermark( new_watermark=6000000)))), str( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload. Event( processing_time_event=beam_runner_api_pb2. TestStreamPayload.Event.AdvanceProcessingTime( advance_duration=1000000)))), str( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload. Event( watermark_event=beam_runner_api_pb2. TestStreamPayload.Event.AdvanceWatermark( new_watermark=8000000)))), str( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload. Event( processing_time_event=beam_runner_api_pb2. TestStreamPayload.Event.AdvanceProcessingTime( advance_duration=1000000)))), str( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload. Event( watermark_event=beam_runner_api_pb2. TestStreamPayload.Event.AdvanceWatermark( new_watermark=10000000)))), str( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload. Event( processing_time_event=beam_runner_api_pb2. TestStreamPayload.Event.AdvanceProcessingTime( advance_duration=1000000)))), str( beam_interactive_api_pb2.TestStreamFileRecord( recorded_event=beam_runner_api_pb2.TestStreamPayload. Event( element_event=beam_runner_api_pb2.TestStreamPayload. Event.AddElements( elements=[ beam_runner_api_pb2.TestStreamPayload. TimestampedElement( encoded_element=coder.encode('1'), timestamp=15000000), beam_runner_api_pb2.TestStreamPayload. TimestampedElement( encoded_element=coder.encode('2'), timestamp=15000000), beam_runner_api_pb2.TestStreamPayload. TimestampedElement( encoded_element=coder.encode('3'), timestamp=15000000), ])))), ], })) p.run() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/test_utils.py0000644000175100017510000002063615206113033022113 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utility methods for testing For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file import hashlib import importlib import os import shutil import struct import tempfile import numpy as np from apache_beam.io.filesystems import FileSystems from apache_beam.utils import retry DEFAULT_HASHING_ALG = 'sha1' class TempDir(object): """Context Manager to create and clean-up a temporary directory.""" def __init__(self): self._tempdir = tempfile.mkdtemp() def __enter__(self): return self def __exit__(self, *args): if os.path.exists(self._tempdir): shutil.rmtree(self._tempdir) def get_path(self): """Returns the path to the temporary directory.""" return self._tempdir def create_temp_file(self, suffix='', lines=None): """Creates a temporary file in the temporary directory. Args: suffix (str): The filename suffix of the temporary file (e.g. '.txt') lines (List[str]): A list of lines that will be written to the temporary file. Returns: The name of the temporary file created. """ with tempfile.NamedTemporaryFile(delete=False, dir=self._tempdir, suffix=suffix) as f: if lines: for line in lines: f.write(line) return f.name def compute_hash(content, hashing_alg=DEFAULT_HASHING_ALG): """Compute a hash value of a list of objects by hashing their string representations.""" content = [ str(x).encode('utf-8') if not isinstance(x, bytes) else x for x in content ] content.sort() m = hashlib.new(hashing_alg) for elem in content: m.update(elem) return m.hexdigest() def patch_retry(testcase, module): """A function to patch retry module to use mock clock and logger. Clock and logger that defined in retry decorator will be replaced in test in order to skip sleep phase when retry happens. Args: testcase: An instance of unittest.TestCase that calls this function to patch retry module. module: The module that uses retry and need to be replaced with mock clock and logger in test. """ # Import mock here to avoid execution time errors for other utilities from mock import Mock from mock import patch real_retry_with_exponential_backoff = retry.with_exponential_backoff def patched_retry_with_exponential_backoff(**kwargs): """A patch for retry decorator to use a mock dummy clock and logger.""" kwargs.update(logger=Mock(), clock=Mock()) return real_retry_with_exponential_backoff(**kwargs) patch.object( retry, 'with_exponential_backoff', side_effect=patched_retry_with_exponential_backoff).start() # Reload module after patching. importlib.reload(module) def remove_patches(): patch.stopall() # Reload module again after removing patch. importlib.reload(module) testcase.addCleanup(remove_patches) @retry.with_exponential_backoff( num_retries=3, retry_filter=retry.retry_on_beam_io_error_filter) def delete_files(file_paths): """A function to clean up files or directories using ``FileSystems``. Glob is supported in file path and directories will be deleted recursively. Args: file_paths: A list of strings contains file paths or directories. """ if len(file_paths) == 0: raise RuntimeError('Clean up failed. Invalid file path: %s.' % file_paths) FileSystems.delete(file_paths) def cleanup_subscriptions(sub_client, subs): """Cleanup PubSub subscriptions if exist.""" for sub in subs: sub_client.delete_subscription(subscription=sub.name) def cleanup_topics(pub_client, topics): """Cleanup PubSub topics if exist.""" for topic in topics: pub_client.delete_topic(topic=topic.name) class PullResponseMessage(object): """Data representing a pull request response. Utility class for ``create_pull_response``. """ def __init__( self, data, attributes=None, publish_time_secs=None, publish_time_nanos=None, ack_id=None): self.data = data self.attributes = attributes self.publish_time_secs = publish_time_secs self.publish_time_nanos = publish_time_nanos self.ack_id = ack_id def create_pull_response(responses): """Create an instance of ``google.cloud.pubsub.types.ReceivedMessage``. Used to simulate the response from pubsub.SubscriberClient().pull(). Args: responses: list of ``PullResponseMessage`` Returns: An instance of ``google.cloud.pubsub.types.PullResponse`` populated with responses. """ from google.cloud import pubsub from google.protobuf import timestamp_pb2 res = pubsub.types.PullResponse() for response in responses: received_message = pubsub.types.ReceivedMessage() message = received_message.message message.data = response.data if response.attributes is not None: for k, v in response.attributes.items(): message.attributes[k] = v publish_time = timestamp_pb2.Timestamp() if response.publish_time_secs is not None: publish_time.seconds = response.publish_time_secs if response.publish_time_nanos is not None: publish_time.nanos = response.publish_time_nanos message.publish_time = publish_time if response.ack_id is not None: received_message.ack_id = response.ack_id res.received_messages.append(received_message) return res def create_file(path, contents): """Create a file to use as input to test pipelines""" with FileSystems.create(path) as f: f.write(str.encode(contents, 'utf-8')) return path def read_files_from_pattern(file_pattern): """Reads the files that match a pattern""" metadata_list = FileSystems.match([file_pattern])[0].metadata_list output = [] for metadata in metadata_list: with FileSystems.open(metadata.path) as f: output.append(f.read().decode('utf-8').strip()) return '\n'.join(output) class LCGenerator: """A pure Python implementation of linear congruential generator.""" def __init__(self, a=0x5DEECE66D, c=0xB, bits=48): self._a = a self._c = c self._bits = bits self._mask = (1 << self._bits) - 1 self.seed(0) def seed(self, seed): """Set the seed for the generator.""" # This implementation is made different than the JDK equivalent # self._seed = (seed ^ self._a) & self._mask # to avoid that close seed gives close first pseudo random number. self._seed = (seed * self._a + self._c) & self._mask def seed_jdk(self, seed): """ Set the seed in the way equivalent to java.util.Random.seed. This method provide a JDK equivalent random generator. Note that the JDK implementation has a known issue that is close seed generating close first pseudo-random number. """ self._seed = (seed ^ self._a) & self._mask def next_int(self): """ Get the next integer value in the range of np.int32 ([0, 2^31-1]). Equivalent to java.util.Random.nextInt. """ self.seed(self._seed) return np.int32(np.int64(self._seed) >> (self._bits - 32)) def next_uint(self): """ Get the next unsigned integer value in the range of np.uint32. Equivalent to java.util.Random.next(32). """ self.seed(self._seed) return self._seed >> (self._bits - 32) def randbytes(self, length): """ Get random bytes of given length. Equivalent to java.util.Random.nextBytes. """ ints = (length + 3) // 4 return struct.pack("<%iI" % ints, *(self.next_uint() for _ in range(ints)))[:length] def random_sample(self): """ Get random sample between [0.0, 1.0), single precision float value. Equivalent to java.util.Random.nextFloat. """ return float(self.next_uint() >> 8) / (1 << 24) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/test_utils_test.py0000644000175100017510000001065415206113033023151 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unittest for testing utilities,""" # pytype: skip-file import logging import os import tempfile import unittest import mock from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystems import FileSystems from apache_beam.testing import test_utils as utils class TestUtilsTest(unittest.TestCase): def setUp(self): utils.patch_retry(self, utils) self.tmpdir = tempfile.mkdtemp() def test_delete_files_succeeds(self): path = os.path.join(self.tmpdir, 'f1') with open(path, 'a') as f: f.write('test') assert FileSystems.exists(path) utils.delete_files([path]) assert not FileSystems.exists(path) def test_delete_files_fails_with_io_error(self): path = os.path.join(self.tmpdir, 'f2') with self.assertRaises(BeamIOError) as error: utils.delete_files([path]) self.assertTrue( error.exception.args[0].startswith('Delete operation failed')) self.assertEqual(list(error.exception.exception_details.keys()), [path]) def test_delete_files_fails_with_invalid_arg(self): with self.assertRaises(RuntimeError): utils.delete_files([]) def test_temp_dir_removes_files(self): with utils.TempDir() as tempdir: dir_path = tempdir.get_path() file_path = tempdir.create_temp_file() self.assertTrue(os.path.exists(dir_path)) self.assertTrue(os.path.exists(file_path)) self.assertFalse(os.path.exists(dir_path)) self.assertFalse(os.path.exists(file_path)) def test_temp_file_field_correct(self): with utils.TempDir() as tempdir: filename = tempdir.create_temp_file( suffix='.txt', lines=[b'line1\n', b'line2\n', b'line3\n']) self.assertTrue(filename.endswith('.txt')) with open(filename, 'rb') as f: self.assertEqual(f.readline(), b'line1\n') self.assertEqual(f.readline(), b'line2\n') self.assertEqual(f.readline(), b'line3\n') def test_cleanup_subscriptions(self): sub_client = mock.Mock() sub = mock.Mock() sub.name = 'test_sub' utils.cleanup_subscriptions(sub_client, [sub]) sub_client.delete_subscription.assert_called_with(subscription=sub.name) def test_cleanup_topics(self): pub_client = mock.Mock() topic = mock.Mock() topic.name = 'test_topic' utils.cleanup_topics(pub_client, [topic]) pub_client.delete_topic.assert_called_with(topic=topic.name) class LCGeneratorTest(unittest.TestCase): Generator = utils.LCGenerator def test_generator_seed_results(self): generator = self.Generator() generator.seed(0) self.assertEqual(generator.next_int(), 4232237) generator.seed(1) self.assertEqual(generator.next_int(), -1151252339) self.assertEqual(generator.next_uint(), 3745583449) self.assertAlmostEqual(generator.random_sample(), 0.375548, delta=1e-6) self.assertEqual(generator.randbytes(10), b'\xa6\x8fW\xcb\xb1\xa88]dP') def test_generator_seed_jdk_results(self): generator = self.Generator() generator.seed_jdk(0) self.assertEqual(generator.next_int(), -1155484576) generator.seed_jdk(1) # the first next_int after seed_jdk(1) is close to seed_jdk(0) self.assertEqual(generator.next_int(), -1155869325) self.assertEqual(generator.next_uint(), 431529176) self.assertAlmostEqual(generator.random_sample(), 0.410081, delta=1e-6) self.assertEqual(generator.randbytes(10), b'\x92\xf9Mh\xfc\xcc,5\xf0\xb8') try: # pylint: disable=wrong-import-position from .fast_test_utils import LCGenerator as FastLCGenerator class FastLCGeneratorTest(LCGeneratorTest): Generator = FastLCGenerator except ImportError: pass if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/util.py0000644000175100017510000003627515206113033020677 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utilities for testing Beam pipelines.""" # pytype: skip-file import collections import glob import io import tempfile from typing import Any from typing import Iterable from typing import NamedTuple from apache_beam import pvalue from apache_beam.transforms import window from apache_beam.transforms.core import Create from apache_beam.transforms.core import DoFn from apache_beam.transforms.core import Filter from apache_beam.transforms.core import Map from apache_beam.transforms.core import ParDo from apache_beam.transforms.core import WindowInto from apache_beam.transforms.ptransform import PTransform from apache_beam.transforms.ptransform import ptransform_fn from apache_beam.transforms.util import CoGroupByKey from apache_beam.utils.windowed_value import PANE_INFO_UNKNOWN from apache_beam.utils.windowed_value import PaneInfo __all__ = [ 'assert_that', 'equal_to', 'equal_to_per_window', 'has_at_least_one', 'is_empty', 'is_not_empty', 'matches_all', # open_shards is internal and has no backwards compatibility guarantees. 'open_shards', 'row_namedtuple_equals_fn', 'TestWindowedValue', ] class BeamAssertException(Exception): """Exception raised by matcher classes used by assert_that transform.""" pass # Used for reifying timestamps and windows for assert_that matchers. class TestWindowedValue(NamedTuple): value: Any timestamp: Any windows: list pane_info: PaneInfo = PANE_INFO_UNKNOWN def contains_in_any_order(iterable): """Creates an object that matches another iterable if they both have the same count of items. Arguments: iterable: An iterable of hashable objects. """ class InAnyOrder(object): def __init__(self, iterable): self._counter = collections.Counter(iterable) def __eq__(self, other): return self._counter == collections.Counter(other) def __hash__(self): return hash(self._counter) def __repr__(self): return "InAnyOrder(%s)" % self._counter return InAnyOrder(iterable) class _EqualToPerWindowMatcher(object): def __init__(self, expected_window_to_elements): self._expected_window_to_elements = expected_window_to_elements def __call__(self, value): # Short-hand. _expected = self._expected_window_to_elements # Match the given windowed value to an expected window. Fails if the window # doesn't exist or the element wasn't found in the window. def match(windowed_value): actual = windowed_value.value window_key = windowed_value.windows[0] try: _expected[window_key] except KeyError: raise BeamAssertException( 'Failed assert: window {} not found in any expected ' \ 'windows {}'.format(window_key, list(_expected.keys())))\ # Remove any matched elements from the window. This is used later on to # assert that all elements in the window were matched with actual # elements. try: _expected[window_key].remove(actual) except ValueError: raise BeamAssertException( 'Failed assert: element {} not found in window ' \ '{}:{}'.format(actual, window_key, _expected[window_key]))\ # Run the matcher for each window and value pair. Fails if the # windowed_value is not a TestWindowedValue. for windowed_value in value: if not isinstance(windowed_value, TestWindowedValue): raise BeamAssertException( 'Failed assert: Received element {} is not of type ' \ 'TestWindowedValue. Did you forget to set reify_windows=True ' \ 'on the assertion?'.format(windowed_value)) match(windowed_value) # Finally, some elements may not have been matched. Assert that we removed # all the elements that we received from the expected list. If the list is # non-empty, then there are unmatched elements. for win in _expected: if _expected[win]: raise BeamAssertException( 'Failed assert: unmatched elements {} in window {}'.format( _expected[win], win)) def equal_to_per_window(expected_window_to_elements): """Matcher used by assert_that to check to assert expected windows. The 'assert_that' statement must have reify_windows=True. This assertion works when elements are emitted and are finally checked at the end of the window. Arguments: expected_window_to_elements: A dictionary where the keys are the windows to check and the values are the elements associated with each window. """ return _EqualToPerWindowMatcher(expected_window_to_elements) # Note that equal_to checks if expected and actual are permutations of each # other. However, only permutations of the top level are checked. Therefore # [1,2] and [2,1] are considered equal and [[1,2]] and [[2,1]] are not. def equal_to(expected, equals_fn=None): def _equal(actual, equals_fn=equals_fn): expected_list = list(expected) # Try to compare actual and expected by sorting. This fails with a # TypeError in Python 3 if different types are present in the same # collection. It can also raise false negatives for types that don't have # a deterministic sort order, like pyarrow Tables as of 0.14.1 if not equals_fn: equals_fn = row_namedtuple_equals_fn try: sorted_expected = sorted(expected) sorted_actual = sorted(actual) if sorted_expected == sorted_actual: return except TypeError: pass # Slower method, used in two cases: # 1) If sorted expected != actual, use this method to verify the inequality. # This ensures we don't raise any false negatives for types that don't # have a deterministic sort order. # 2) As a fallback if we encounter a TypeError in python 3. this method # works on collections that have different types. unexpected = [] for element in actual: found = False for i, v in enumerate(expected_list): if equals_fn(v, element): found = True expected_list.pop(i) break if not found: unexpected.append(element) if unexpected or expected_list: msg = 'Failed assert: %r == %r' % (expected, actual) if unexpected: msg = msg + ', unexpected elements %r' % unexpected if expected_list: msg = msg + ', missing elements %r' % expected_list raise BeamAssertException(msg) return _equal def row_namedtuple_equals_fn(expected, actual, fallback_equals_fn=None): """ equals_fn which can be used by equal_to which treats Rows and NamedTuples as equivalent types. This can be useful since Beam converts Rows to NamedTuples when they are sent across portability layers, so a Row may be converted to a NamedTuple automatically by Beam. """ if fallback_equals_fn is None: fallback_equals_fn = lambda e, a: e == a if type(expected) is not pvalue.Row and not _is_named_tuple(expected): return fallback_equals_fn(expected, actual) if type(actual) is not pvalue.Row and not _is_named_tuple(actual): return fallback_equals_fn(expected, actual) expected_dict = expected._asdict() actual_dict = actual._asdict() if len(expected_dict) != len(actual_dict): return False for k, v in expected_dict.items(): if k not in actual_dict: return False if not row_namedtuple_equals_fn(v, actual_dict[k]): return False return True def matches_all(expected): """Matcher used by assert_that to check a set of matchers. Args: expected: A list of elements or hamcrest matchers to be used to match the elements of a single PCollection. """ def _matches(actual): from hamcrest.core import assert_that as hamcrest_assert from hamcrest.library.collection import contains_inanyorder expected_list = list(expected) hamcrest_assert(actual, contains_inanyorder(*expected_list)) return _matches def is_empty(): def _empty(actual): actual = list(actual) if actual: raise BeamAssertException('Failed assert: [] == %r' % actual) return _empty def is_not_empty(): """ This is test method which makes sure that the pcol is not empty and it has some data in it. :return: """ def _not_empty(actual): actual = list(actual) if not actual: raise BeamAssertException('Failed assert: pcol is empty') return _not_empty def assert_that( actual, matcher, label='assert_that', reify_windows=False, use_global_window=True): """A PTransform that checks a PCollection has an expected value. Note that assert_that should be used only for testing pipelines since the check relies on materializing the entire PCollection being checked. Args: actual: A PCollection. matcher: A matcher function taking as argument the actual value of a materialized PCollection. The matcher validates this actual value against expectations and raises BeamAssertException if they are not met. label: Optional string label. This is needed in case several assert_that transforms are introduced in the same pipeline. reify_windows: If True, matcher is passed a list of TestWindowedValue. use_global_window: If False, matcher is passed a dictionary of (k, v) = (window, elements in the window). Returns: Ignored. """ assert isinstance(actual, pvalue.PCollection), ( '%s is not a supported type for Beam assert' % type(actual)) pipeline = actual.pipeline if getattr(actual.pipeline, 'result', None): # The pipeline was already run. The user most likely called assert_that # after the pipeleline context. raise RuntimeError( 'assert_that must be used within a beam.Pipeline context. ' + 'Prior to Beam 2.60.0, asserts outside of the context of a pipeline ' + 'were silently ignored, starting with Beam 2.60.0 this is no longer ' + 'allowed. To fix, move your assert_that call into your pipeline ' + 'context so that it is added before the pipeline is run. For more ' + 'information, see https://github.com/apache/beam/pull/30771') # Usually, the uniqueness of the label is left to the pipeline # writer to guarantee. Since we're in a testing context, we'll # just automatically append a number to the label if it's # already in use, as tests don't typically have to worry about # long-term update compatibility stability of stage names. if label in pipeline.applied_labels: label_idx = 2 while f"{label}_{label_idx}" in pipeline.applied_labels: label_idx += 1 label = f"{label}_{label_idx}" if isinstance(matcher, _EqualToPerWindowMatcher): reify_windows = True use_global_window = True class ReifyTimestampWindow(DoFn): def process( self, element, timestamp=DoFn.TimestampParam, window=DoFn.WindowParam, pane_info=DoFn.PaneInfoParam): # This returns TestWindowedValue instead of # beam.utils.windowed_value.WindowedValue because ParDo will extract # the timestamp and window out of the latter. return [TestWindowedValue(element, timestamp, [window], pane_info)] class AddWindow(DoFn): def process(self, element, window=DoFn.WindowParam): yield element, window class AssertThat(PTransform): def expand(self, pcoll): if reify_windows: pcoll = pcoll | ParDo(ReifyTimestampWindow()) keyed_singleton = pcoll.pipeline | Create([(None, None)]) keyed_singleton.is_bounded = True if use_global_window: pcoll = pcoll | WindowInto(window.GlobalWindows()) keyed_actual = pcoll | 'ToVoidKey' >> Map(lambda v: (None, v)) keyed_actual.is_bounded = True # This is a CoGroupByKey so that the matcher always runs, even if the # PCollection is empty. plain_actual = ((keyed_singleton, keyed_actual) | 'Group' >> CoGroupByKey() | 'Unkey' >> Map(lambda k_values: list(k_values[1][1]))) if not use_global_window: plain_actual = plain_actual | 'AddWindow' >> ParDo(AddWindow()) return plain_actual | 'Match' >> Map(matcher) def default_label(self): return label return actual | AssertThat() @ptransform_fn def AssertThat(pcoll, *args, **kwargs): """Like assert_that, but as an applicable PTransform.""" return assert_that(pcoll, *args, **kwargs) def has_at_least_one(input, criterion, label="has_at_least_one"): pipeline = input.pipeline # similar to assert_that, we choose a label if it already exists. if label in pipeline.applied_labels: label_idx = 2 while f"{label}_{label_idx}" in pipeline.applied_labels: label_idx += 1 label = f"{label}_{label_idx}" def _apply_criterion( e=DoFn.ElementParam, t=DoFn.TimestampParam, w=DoFn.WindowParam, p=DoFn.PaneInfoParam): if criterion(e, t, w, p): return e, t, w, p def _not_empty(actual): actual = list(actual) if not actual: raise BeamAssertException('Failed assert: nothing matches the criterion') result = input | label >> Map(_apply_criterion) | label + "_filter" >> Filter( lambda e: e is not None) assert_that(result, _not_empty) def open_shards(glob_pattern, mode='rt', encoding='utf-8'): """Returns a composite file of all shards matching the given glob pattern. Args: glob_pattern (str): Pattern used to match files which should be opened. mode (str): Specify the mode in which the file should be opened. For available modes, check io.open() documentation. encoding (str): Name of the encoding used to decode or encode the file. This should only be used in text mode. Returns: A stream with the contents of the opened files. """ if 'b' in mode: encoding = None with tempfile.NamedTemporaryFile(delete=False) as out_file: for shard in glob.glob(glob_pattern): with open(shard, 'rb') as in_file: out_file.write(in_file.read()) concatenated_file_name = out_file.name return io.open(concatenated_file_name, mode, encoding=encoding) def _sort_lists(result): if isinstance(result, list): return sorted(result) elif isinstance(result, tuple): return tuple(_sort_lists(e) for e in result) elif isinstance(result, dict): return {k: _sort_lists(v) for k, v in result.items()} elif isinstance(result, Iterable) and not isinstance(result, str): return sorted(result) else: return result def _is_named_tuple(obj) -> bool: return ( isinstance(obj, tuple) and hasattr(obj, '_asdict') and hasattr(obj, '_fields')) # A utility transform that recursively sorts lists for easier testing. SortLists = Map(_sort_lists) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/testing/util_test.py0000644000175100017510000002673315206113033021734 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for testing utilities.""" # pytype: skip-file import unittest from typing import NamedTuple import apache_beam as beam from apache_beam import Create from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import TestWindowedValue from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.testing.util import equal_to_per_window from apache_beam.testing.util import is_empty from apache_beam.testing.util import is_not_empty from apache_beam.testing.util import row_namedtuple_equals_fn from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import GlobalWindow from apache_beam.transforms.window import IntervalWindow from apache_beam.utils.timestamp import MIN_TIMESTAMP class UtilTest(unittest.TestCase): def test_assert_that_passes(self): with TestPipeline() as p: assert_that(p | Create([1, 2, 3]), equal_to([1, 2, 3])) def test_assert_that_passes_order_does_not_matter(self): with TestPipeline() as p: assert_that(p | Create([1, 2, 3]), equal_to([2, 1, 3])) def test_assert_that_passes_order_does_not_matter_with_negatives(self): with TestPipeline() as p: assert_that(p | Create([1, -2, 3]), equal_to([-2, 1, 3])) def test_assert_that_passes_empty_equal_to(self): with TestPipeline() as p: assert_that(p | Create([]), equal_to([])) def test_assert_that_passes_empty_is_empty(self): with TestPipeline() as p: assert_that(p | Create([]), is_empty()) def test_assert_that_fails(self): with self.assertRaises(Exception): with TestPipeline() as p: assert_that(p | Create([1, 10, 100]), equal_to([1, 2, 3])) def test_assert_missing(self): with self.assertRaisesRegex(Exception, r".*missing elements \['c'\]"): with TestPipeline() as p: assert_that(p | Create(['a', 'b']), equal_to(['a', 'b', 'c'])) def test_assert_unexpected(self): with self.assertRaisesRegex(Exception, r".*unexpected elements \['c', 'd'\]|" r"unexpected elements \['d', 'c'\]"): with TestPipeline() as p: assert_that(p | Create(['a', 'b', 'c', 'd']), equal_to(['a', 'b'])) def test_assert_missing_and_unexpected(self): with self.assertRaisesRegex(Exception, r".*unexpected elements \[" r"'c'\].*missing elements" r" \['d'\]"): with TestPipeline() as p: assert_that(p | Create(['a', 'b', 'c']), equal_to(['a', 'b', 'd'])) def test_assert_with_custom_comparator(self): with TestPipeline() as p: assert_that( p | Create([1, 2, 3]), equal_to(['1', '2', '3'], equals_fn=lambda e, a: int(e) == int(a))) def test_reified_value_passes(self): expected = [ TestWindowedValue(v, MIN_TIMESTAMP, [GlobalWindow()]) for v in [1, 2, 3] ] with TestPipeline() as p: assert_that(p | Create([2, 3, 1]), equal_to(expected), reify_windows=True) def test_reified_value_assert_fail_unmatched_value(self): expected = [ TestWindowedValue(v + 1, MIN_TIMESTAMP, [GlobalWindow()]) for v in [1, 2, 3] ] with self.assertRaises(Exception): with TestPipeline() as p: assert_that( p | Create([2, 3, 1]), equal_to(expected), reify_windows=True) def test_reified_value_assert_fail_unmatched_timestamp(self): expected = [TestWindowedValue(v, 1, [GlobalWindow()]) for v in [1, 2, 3]] with self.assertRaises(Exception): with TestPipeline() as p: assert_that( p | Create([2, 3, 1]), equal_to(expected), reify_windows=True) def test_reified_value_assert_fail_unmatched_window(self): expected = [ TestWindowedValue(v, MIN_TIMESTAMP, [IntervalWindow(0, 1)]) for v in [1, 2, 3] ] with self.assertRaises(Exception): with TestPipeline() as p: assert_that( p | Create([2, 3, 1]), equal_to(expected), reify_windows=True) def test_assert_that_fails_on_empty_input(self): with self.assertRaises(Exception): with TestPipeline() as p: assert_that(p | Create([]), equal_to([1, 2, 3])) def test_assert_that_fails_on_empty_expected(self): with self.assertRaises(Exception): with TestPipeline() as p: assert_that(p | Create([1, 2, 3]), is_empty()) def test_assert_that_passes_is_not_empty(self): with TestPipeline() as p: assert_that(p | Create([1, 2, 3]), is_not_empty()) def test_assert_that_fails_on_is_not_empty_expected(self): with self.assertRaisesRegex(Exception, "pcol is empty"): with TestPipeline() as p: assert_that(p | Create([]), is_not_empty()) def test_equal_to_per_window_passes(self): start = int(MIN_TIMESTAMP.micros // 1e6) - 5 end = start + 20 expected = { window.IntervalWindow(start, end): [('k', [1])], } with TestPipeline(options=StandardOptions(streaming=True)) as p: assert_that(( p | Create([1]) | beam.WindowInto( FixedWindows(20), trigger=trigger.AfterWatermark(), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.Map(lambda x: ('k', x)) | beam.GroupByKey()), equal_to_per_window(expected), reify_windows=True) def test_equal_to_per_window_fail_unmatched_window(self): with self.assertRaisesRegex(Exception, "not found in any expected"): expected = { window.IntervalWindow(50, 100): [('k', [1])], } with TestPipeline(options=StandardOptions(streaming=True)) as p: assert_that(( p | Create([1]) | beam.WindowInto( FixedWindows(20), trigger=trigger.AfterWatermark(), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.Map(lambda x: ('k', x)) | beam.GroupByKey()), equal_to_per_window(expected), reify_windows=True) def test_runtimeerror_outside_of_context(self): with beam.Pipeline() as p: outputs = (p | beam.Create([1, 2, 3]) | beam.Map(lambda x: x + 1)) with self.assertRaises(RuntimeError): assert_that(outputs, equal_to([2, 3, 4])) def test_multiple_assert_that_labels(self): with beam.Pipeline() as p: outputs = (p | beam.Create([1, 2, 3]) | beam.Map(lambda x: x + 1)) assert_that(outputs, equal_to([2, 3, 4])) assert_that(outputs, equal_to([2, 3, 4])) assert_that(outputs, equal_to([2, 3, 4])) def test_equal_to_per_window_fail_unmatched_element(self): with self.assertRaisesRegex(Exception, "unmatched elements"): start = int(MIN_TIMESTAMP.micros // 1e6) - 5 end = start + 20 expected = { window.IntervalWindow(start, end): [('k', [1]), ('k', [2])], } with TestPipeline(options=StandardOptions(streaming=True)) as p: assert_that(( p | Create([1]) | beam.WindowInto( FixedWindows(20), trigger=trigger.AfterWatermark(), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.Map(lambda x: ('k', x)) | beam.GroupByKey()), equal_to_per_window(expected), reify_windows=True) def test_equal_to_per_window_succeeds_no_reify_windows(self): start = int(MIN_TIMESTAMP.micros // 1e6) - 5 end = start + 20 expected = { window.IntervalWindow(start, end): [('k', [1])], } with TestPipeline(options=StandardOptions(streaming=True)) as p: assert_that(( p | Create([1]) | beam.WindowInto( FixedWindows(20), trigger=trigger.AfterWatermark(), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.Map(lambda x: ('k', x)) | beam.GroupByKey()), equal_to_per_window(expected)) def test_equal_to_per_window_fail_unexpected_element(self): with self.assertRaisesRegex(Exception, "not found in window"): start = int(MIN_TIMESTAMP.micros // 1e6) - 5 end = start + 20 expected = { window.IntervalWindow(start, end): [('k', [1])], } with TestPipeline(options=StandardOptions(streaming=True)) as p: assert_that(( p | Create([1, 2]) | beam.WindowInto( FixedWindows(20), trigger=trigger.AfterWatermark(), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.Map(lambda x: ('k', x)) | beam.GroupByKey()), equal_to_per_window(expected), reify_windows=True) def test_row_namedtuple_equals(self): class RowTuple(NamedTuple): a: str b: int self.assertTrue( row_namedtuple_equals_fn( beam.Row(a='123', b=456), beam.Row(a='123', b=456))) self.assertTrue( row_namedtuple_equals_fn( beam.Row(a='123', b=456), RowTuple(a='123', b=456))) self.assertTrue( row_namedtuple_equals_fn( RowTuple(a='123', b=456), RowTuple(a='123', b=456))) self.assertTrue( row_namedtuple_equals_fn( RowTuple(a='123', b=456), beam.Row(a='123', b=456))) self.assertTrue(row_namedtuple_equals_fn('foo', 'foo')) self.assertFalse( row_namedtuple_equals_fn( beam.Row(a='123', b=456), beam.Row(a='123', b=4567))) self.assertFalse( row_namedtuple_equals_fn( beam.Row(a='123', b=456), beam.Row(a='123', b=456, c='a'))) self.assertFalse( row_namedtuple_equals_fn( beam.Row(a='123', b=456), RowTuple(a='123', b=4567))) self.assertFalse( row_namedtuple_equals_fn( beam.Row(a='123', b=456, c='foo'), RowTuple(a='123', b=4567))) self.assertFalse( row_namedtuple_equals_fn(beam.Row(a='123'), RowTuple(a='123', b=4567))) self.assertFalse(row_namedtuple_equals_fn(beam.Row(a='123'), '123')) self.assertFalse(row_namedtuple_equals_fn('123', RowTuple(a='123', b=4567))) class NestedNamedTuple(NamedTuple): a: str b: RowTuple self.assertTrue( row_namedtuple_equals_fn( beam.Row(a='foo', b=beam.Row(a='123', b=456)), NestedNamedTuple(a='foo', b=RowTuple(a='123', b=456)))) self.assertTrue( row_namedtuple_equals_fn( beam.Row(a='foo', b=beam.Row(a='123', b=456)), beam.Row(a='foo', b=RowTuple(a='123', b=456)))) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.7274444 apache_beam-2.74.0/apache_beam/tools/0000755000175100017510000000000015206113160017017 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/tools/__init__.py0000644000175100017510000000157415206113033021136 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A package containing internal tools. For internal use only; no backwards-compatibility guarantees. """ ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/tools/coders_microbenchmark.py0000644000175100017510000002627415206113033023726 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A microbenchmark for measuring performance of coders. This runs a sequence of encode-decode operations on random inputs to collect performance of various coders. To evaluate coders performance we approximate the behavior how the coders are used in PCollections: we encode and decode a list of elements. An element can be a string, a list of integers, a windowed value, or any other object we want a coder to process. Run as: python -m apache_beam.tools.coders_microbenchmark """ # pytype: skip-file import argparse import logging import random import re import string import sys import apache_beam as beam from apache_beam.coders import coder_impl from apache_beam.coders import coders from apache_beam.coders import coders_test_common from apache_beam.coders import proto2_coder_test_messages_pb2 as test_message from apache_beam.coders import row_coder from apache_beam.coders import typecoders from apache_beam.tools import utils from apache_beam.transforms import window from apache_beam.typehints import trivial_inference from apache_beam.typehints.pandas_type_compatibility import DataFrameBatchConverterDropIndex from apache_beam.utils import windowed_value def coder_benchmark_factory(coder, generate_fn): """Creates a benchmark that encodes and decodes a list of elements. Args: coder: coder to use to encode an element. generate_fn: a callable that generates an element. """ class CoderBenchmark(object): def __init__(self, num_elements_per_benchmark): self._coder = coders.IterableCoder(coder) self._list = [generate_fn() for _ in range(num_elements_per_benchmark)] def __call__(self): # Calling coder operations on a single element at a time may incur # unrelevant overhead. To compensate, we use a list elements. _ = self._coder.decode(self._coder.encode(self._list)) CoderBenchmark.__name__ = "%s, %s" % (generate_fn.__name__, str(coder)) return CoderBenchmark def batch_row_coder_benchmark_factory(generate_fn, use_batch): """Creates a benchmark that encodes and decodes a list of elements. Args: coder: coder to use to encode an element. generate_fn: a callable that generates an element. """ class CoderBenchmark(object): def __init__(self, num_elements_per_benchmark): self._use_batch = use_batch row_instance = generate_fn() row_type = trivial_inference.instance_to_type(row_instance) self._row_coder = get_row_coder(row_instance) self._batch_converter = DataFrameBatchConverterDropIndex(row_type) self._seq_coder = coders.IterableCoder(self._row_coder) self._data = self._batch_converter.produce_batch( [generate_fn() for _ in range(num_elements_per_benchmark)]) def __call__(self): if self._use_batch: impl = self._row_coder.get_impl() columnar = { col: self._data[col].to_numpy() for col in self._data.columns } output_stream = coder_impl.create_OutputStream() impl.encode_batch_to_stream(columnar, output_stream) impl.decode_batch_from_stream( columnar, coder_impl.create_InputStream(output_stream.get())) else: # Calling coder operations on a single element at a time may incur # unrelevant overhead. To compensate, we use a list elements. self._batch_converter.produce_batch( self._seq_coder.decode( self._seq_coder.encode( self._batch_converter.explode_batch(self._data)))) CoderBenchmark.__name__ = "%s, BatchRowCoder%s" % ( generate_fn.__name__, use_batch) return CoderBenchmark def small_int(): return random.randint(0, 127) def large_int(): return random.randint(sys.maxsize >> 2, sys.maxsize) def random_string(length): return ''.join( random.choice(string.ascii_letters + string.digits) for _ in range(length)) def small_string(): return random_string(4) def large_string(): return random_string(100) def list_int(size): return [small_int() for _ in range(size)] def dict_int_int(size): return {i: i for i in list_int(size)} def small_list(): return list_int(10) def large_list(): # Bool is the last item in FastPrimitiveCoders before pickle. return [bool(k) for k in list_int(1000)] def small_tuple(): # Benchmark a common case of 2-element tuples. return tuple(list_int(2)) def large_tuple(): return tuple(large_list()) def small_dict(): return {i: i for i in small_list()} def large_dict(): return {i: i for i in large_list()} def large_iterable(): yield 'a' * coders.coder_impl.SequenceCoderImpl._DEFAULT_BUFFER_SIZE for k in range(1000): yield k def random_message_with_map(size): message = test_message.MessageWithMap() keys = list_int(size) random.shuffle(keys) for key in keys: message.field1[str(key)].field1 = small_string() return message def small_message_with_map(): return random_message_with_map(5) def large_message_with_map(): return random_message_with_map(20) def globally_windowed_value(): return windowed_value.WindowedValue( value=small_int(), timestamp=12345678, windows=(window.GlobalWindow(), )) def random_windowed_value(num_windows): return windowed_value.WindowedValue( value=small_int(), timestamp=12345678, windows=tuple( window.IntervalWindow(i * 10, i * 10 + small_int()) for i in range(num_windows))) def wv_with_one_window(): return random_windowed_value(num_windows=1) def wv_with_multiple_windows(): return random_windowed_value(num_windows=32) def tiny_row(): return beam.Row(int_value=1) def large_row(): return beam.Row(**{f'int_{ix}': ix for ix in range(20)}) def nullable_row(): return beam.Row(**{f'int_{ix}': ix if ix % 2 else None for ix in range(20)}) def diverse_row(): return beam.Row( int_value=1, float_value=3.14159, str_value='beam', row_value=beam.Row(int_value=2, float_value=2.718281828)) def get_row_coder(row_instance): coder = typecoders.registry.get_coder( trivial_inference.instance_to_type(row_instance)) assert isinstance(coder, row_coder.RowCoder) return coder def row_coder_benchmark_factory(generate_fn): return coder_benchmark_factory(get_row_coder(generate_fn()), generate_fn) def importable_named_tuple(): return [coders_test_common.MyTypedNamedTuple('a', i) for i in range(1000)] def run_coder_benchmarks( num_runs, input_size, seed, verbose, filter_regex='.*'): random.seed(seed) # TODO(https://github.com/apache/beam/issues/18788): Pick coders using type # hints, for example: # tuple_coder = typecoders.registry.get_coder(typing.Tuple[int, ...]) benchmarks = [ coder_benchmark_factory(coders.FastPrimitivesCoder(), small_int), coder_benchmark_factory(coders.FastPrimitivesCoder(), large_int), coder_benchmark_factory(coders.FastPrimitivesCoder(), small_string), coder_benchmark_factory(coders.FastPrimitivesCoder(), large_string), coder_benchmark_factory(coders.FastPrimitivesCoder(), small_list), coder_benchmark_factory( coders.IterableCoder(coders.FastPrimitivesCoder()), small_list), coder_benchmark_factory(coders.FastPrimitivesCoder(), large_list), coder_benchmark_factory( coders.IterableCoder(coders.FastPrimitivesCoder()), large_list), coder_benchmark_factory( coders.IterableCoder(coders.FastPrimitivesCoder()), large_iterable), coder_benchmark_factory(coders.FastPrimitivesCoder(), small_tuple), coder_benchmark_factory(coders.FastPrimitivesCoder(), large_tuple), coder_benchmark_factory(coders.FastPrimitivesCoder(), small_dict), coder_benchmark_factory(coders.FastPrimitivesCoder(), large_dict), coder_benchmark_factory( coders.ProtoCoder(test_message.MessageWithMap), small_message_with_map), coder_benchmark_factory( coders.ProtoCoder(test_message.MessageWithMap), large_message_with_map), coder_benchmark_factory( coders.DeterministicProtoCoder(test_message.MessageWithMap), small_message_with_map), coder_benchmark_factory( coders.DeterministicProtoCoder(test_message.MessageWithMap), large_message_with_map), coder_benchmark_factory( coders.WindowedValueCoder(coders.FastPrimitivesCoder()), wv_with_one_window), coder_benchmark_factory( coders.WindowedValueCoder( coders.FastPrimitivesCoder(), coders.IntervalWindowCoder()), wv_with_multiple_windows), coder_benchmark_factory( coders.WindowedValueCoder( coders.FastPrimitivesCoder(), coders.GlobalWindowCoder()), globally_windowed_value), coder_benchmark_factory( coders.LengthPrefixCoder(coders.FastPrimitivesCoder()), small_int), row_coder_benchmark_factory(tiny_row), row_coder_benchmark_factory(large_row), row_coder_benchmark_factory(nullable_row), row_coder_benchmark_factory(diverse_row), batch_row_coder_benchmark_factory(tiny_row, False), batch_row_coder_benchmark_factory(tiny_row, True), batch_row_coder_benchmark_factory(large_row, False), batch_row_coder_benchmark_factory(large_row, True), batch_row_coder_benchmark_factory(nullable_row, False), batch_row_coder_benchmark_factory(nullable_row, True), batch_row_coder_benchmark_factory(diverse_row, False), batch_row_coder_benchmark_factory(diverse_row, True), coder_benchmark_factory( coders.IterableCoder( coders.FastPrimitivesCoder().as_deterministic_coder( step_label="step")), importable_named_tuple), ] suite = [ utils.BenchmarkConfig(b, input_size, num_runs) for b in benchmarks if re.search(filter_regex, b.__name__, flags=re.I) ] utils.run_benchmarks(suite, verbose=verbose) if __name__ == "__main__": logging.basicConfig() parser = argparse.ArgumentParser() parser.add_argument('--filter', default='.*') parser.add_argument('--num_runs', default=20, type=int) parser.add_argument('--num_elements_per_benchmark', default=1000, type=int) parser.add_argument('--seed', default=42, type=int) options = parser.parse_args() utils.check_compiled("apache_beam.coders.coder_impl") num_runs = 20 num_elements_per_benchmark = 1000 seed = 42 # Fix the seed for better consistency run_coder_benchmarks( options.num_runs, options.num_elements_per_benchmark, options.seed, verbose=True, filter_regex=options.filter) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/tools/distribution_counter_microbenchmark.py0000644000175100017510000000453215206113033026716 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A microbenchmark for measuring DistributionAccumulator performance This runs a sequence of distribution.update for random input value to calculate average update time per input. A typical update operation should run into 0.6 microseconds Run as python -m apache_beam.tools.distribution_counter_microbenchmark """ # pytype: skip-file import logging import random import sys import time from apache_beam.tools import utils def generate_input_values(num_input, lower_bound, upper_bound): values = [] # pylint: disable=unused-variable for i in range(num_input): values.append(random.randint(lower_bound, upper_bound)) return values def run_benchmark(num_runs=100, num_input=10000, seed=time.time()): total_time = 0 random.seed(seed) lower_bound = 0 upper_bound = sys.maxsize inputs = generate_input_values(num_input, lower_bound, upper_bound) from apache_beam.transforms import DataflowDistributionCounter print("Number of runs:", num_runs) print("Input size:", num_input) print("Input sequence from %d to %d" % (lower_bound, upper_bound)) print("Random seed:", seed) for i in range(num_runs): counter = DataflowDistributionCounter() start = time.time() counter.add_inputs_for_test(inputs) time_cost = time.time() - start print("Run %d: Total time cost %g sec" % (i + 1, time_cost)) total_time += time_cost / num_input print("Per element update time cost:", total_time / num_runs) if __name__ == '__main__': logging.basicConfig() utils.check_compiled( 'apache_beam.transforms.cy_dataflow_distribution_counter') run_benchmark() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/tools/fn_api_runner_microbenchmark.py0000644000175100017510000001165715206113033025273 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A microbenchmark for measuring changes in the critical path of FnApiRunner. This microbenchmark attempts to measure the overhead of the main data paths for the FnApiRunner. Specifically state, timers, and shuffling of data. This runs a series of N parallel pipelines with M parallel stages each. Each stage does the following: 1) Put all the PCollection elements in state 2) Set a timer for the future 3) When the timer fires, change the key and output all the elements downstream This executes the same codepaths that are run on the Fn API (and Dataflow) workers, but is generally easier to run (locally) and more stable.. Run as python -m apache_beam.tools.fn_api_runner_microbenchmark The main metric to work with for this benchmark is Fixed Cost. This represents the fixed cost of ovehead for the data path of the FnApiRunner. Initial results were: run 1 of 10, per element time cost: 3.6778 sec run 2 of 10, per element time cost: 0.053498 sec run 3 of 10, per element time cost: 0.0299434 sec run 4 of 10, per element time cost: 0.0211154 sec run 5 of 10, per element time cost: 0.0170031 sec run 6 of 10, per element time cost: 0.0150809 sec run 7 of 10, per element time cost: 0.013218 sec run 8 of 10, per element time cost: 0.0119685 sec run 9 of 10, per element time cost: 0.0107382 sec run 10 of 10, per element time cost: 0.0103208 sec Fixed cost 4.537164939085642 Per-element 0.005474923321695039 R^2 0.95189 """ # pytype: skip-file import argparse import logging import random import apache_beam as beam from apache_beam.coders import VarIntCoder from apache_beam.runners.portability.fn_api_runner import FnApiRunner from apache_beam.tools import utils from apache_beam.transforms.timeutil import TimeDomain from apache_beam.transforms.userstate import SetStateSpec from apache_beam.transforms.userstate import TimerSpec from apache_beam.transforms.userstate import on_timer from apache_beam.typehints import typehints NUM_PARALLEL_STAGES = 7 NUM_SERIAL_STAGES = 7 class BagInStateOutputAfterTimer(beam.DoFn): SET_STATE = SetStateSpec('buffer', VarIntCoder()) EMIT_TIMER = TimerSpec('emit_timer', TimeDomain.WATERMARK) def process( self, element, set_state=beam.DoFn.StateParam(SET_STATE), emit_timer=beam.DoFn.TimerParam(EMIT_TIMER)): _, values = element for v in values: set_state.add(v) emit_timer.set(1) @on_timer(EMIT_TIMER) def emit_values(self, set_state=beam.DoFn.StateParam(SET_STATE)): values = set_state.read() return [(random.randint(0, 1000), v) for v in values] def _build_serial_stages( pipeline, num_serial_stages, num_elements, stage_count): pc = ( pipeline | ('start_stage%s' % stage_count) >> beam.Create( [(random.randint(0, 1000), i) for i in range(num_elements)]) | ('gbk_start_stage%s' % stage_count) >> beam.GroupByKey()) for i in range(num_serial_stages): pc = ( pc | ('stage%s_map%s' % (stage_count, i)) >> beam.ParDo( BagInStateOutputAfterTimer()).with_output_types( typehints.KV[int, int]) | ('stage%s_gbk%s' % (stage_count, i)) >> beam.GroupByKey()) return pc def run_single_pipeline(size): def _pipeline_runner(): with beam.Pipeline(runner=FnApiRunner()) as p: for i in range(NUM_PARALLEL_STAGES): _build_serial_stages(p, NUM_SERIAL_STAGES, size, i) return _pipeline_runner def run_benchmark( starting_point=1, num_runs=10, num_elements_step=100, verbose=True): suite = [ utils.LinearRegressionBenchmarkConfig( run_single_pipeline, starting_point, num_elements_step, num_runs) ] return utils.run_benchmarks(suite, verbose=verbose) if __name__ == '__main__': logging.basicConfig() utils.check_compiled('apache_beam.runners.common') parser = argparse.ArgumentParser() parser.add_argument('--num_runs', default=10, type=int) parser.add_argument('--starting_point', default=1, type=int) parser.add_argument('--increment', default=100, type=int) parser.add_argument('--verbose', default=True, type=bool) options = parser.parse_args() run_benchmark( options.starting_point, options.num_runs, options.increment, options.verbose) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/tools/map_fn_microbenchmark.py0000644000175100017510000000473715206113033023707 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A microbenchmark for measuring changes in overhead for critical code paths. This runs a sequence of trivial Maps over a variable number of inputs to estimate the per-element processing time. It can be useful to run this benchmark before and after a proposed set of changes. A typical per-element cost should be 1-2 microseconds. This executes the same codepaths that are run on the Fn API (and Dataflow) workers, but is generally easier to run (locally) and more stable. It does not, on the other hand, excercise any non-trivial amount of IO (e.g. shuffle). Run as python -m apache_beam.tools.map_fn_microbenchmark """ # pytype: skip-file import logging import time from scipy import stats import apache_beam as beam from apache_beam.tools import utils def run_benchmark(num_maps=100, num_runs=10, num_elements_step=1000): timings = {} for run in range(num_runs): num_elements = num_elements_step * run + 1 start = time.time() with beam.Pipeline() as p: pc = p | beam.Create(list(range(num_elements))) for ix in range(num_maps): pc = pc | 'Map%d' % ix >> beam.FlatMap(lambda x: (None, )) timings[num_elements] = time.time() - start print( "%6d element%s %g sec" % ( num_elements, " " if num_elements == 1 else "s", timings[num_elements])) print() # pylint: disable=unused-variable gradient, intercept, r_value, p_value, std_err = stats.linregress( *list(zip(*list(timings.items())))) print("Fixed cost ", intercept) print("Per-element ", gradient / num_maps) print("R^2 ", r_value**2) if __name__ == '__main__': logging.basicConfig() utils.check_compiled('apache_beam.runners.common') run_benchmark() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/tools/microbenchmarks_test.py0000644000175100017510000000371615206113033023605 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for microbenchmarks code.""" # pytype: skip-file import unittest from importlib.metadata import PackageNotFoundError from importlib.metadata import distribution from apache_beam.tools import coders_microbenchmark from apache_beam.tools import utils class MicrobenchmarksTest(unittest.TestCase): def test_coders_microbenchmark(self): # Right now, we don't evaluate performance impact, only check that # microbenchmark code can successfully run. coders_microbenchmark.run_coder_benchmarks( num_runs=1, input_size=10, seed=1, verbose=False) def is_cython_installed(self): try: distribution('cython') return True except PackageNotFoundError: return False def test_check_compiled(self): if self.is_cython_installed(): utils.check_compiled('apache_beam.runners.worker.opcounters') # Unfortunately, if cython is not installed, that doesn't mean we # can rule out compiled modules (e.g. if Beam was installed from a wheel). # Technically the other way around could be true as well, e.g. if # Cython was installed after Beam, but this is rarer and more easily # remedied. if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/tools/runtime_type_check_microbenchmark.py0000644000175100017510000001065215206113033026321 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A microbenchmark for measuring the overhead of runtime_type_check vs performance_runtime_type_check vs a pipeline with no runtime type check. This runs a sequence of trivial DoFn's over a set of inputs to simulate a real-world pipeline that processes lots of data. Run as python -m apache_beam.tools.runtime_type_check_microbenchmark """ # pytype: skip-file import logging from collections import defaultdict from time import time from typing import Iterable from typing import Union import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.tools import utils @beam.typehints.with_input_types(tuple[int, ...]) class SimpleInput(beam.DoFn): def process(self, element, *args, **kwargs): yield element @beam.typehints.with_output_types(tuple[int, ...]) class SimpleOutput(beam.DoFn): def process(self, element, *args, **kwargs): yield element @beam.typehints.with_input_types( tuple[int, str, tuple[float, ...], Iterable[int], Union[str, int]]) class NestedInput(beam.DoFn): def process(self, element, *args, **kwargs): yield element @beam.typehints.with_output_types( tuple[int, str, tuple[float, ...], Iterable[int], Union[str, int]]) class NestedOutput(beam.DoFn): def process(self, element, *args, **kwargs): yield element def run_benchmark( num_dofns=100, num_runs=10, num_elements_step=2000, num_for_averaging=4): options_map = { 'No Type Check': PipelineOptions(), 'Runtime Type Check': PipelineOptions(runtime_type_check=True), 'Performance Runtime Type Check': PipelineOptions( performance_runtime_type_check=True) } for run in range(num_runs): num_elements = num_elements_step * run + 1 simple_elements = [ tuple(i for i in range(200)) for _ in range(num_elements) ] nested_elements = [( 1, '2', tuple(float(i) for i in range(100)), [i for i in range(100)], '5') for _ in range(num_elements)] timings = defaultdict(lambda: defaultdict(lambda: defaultdict(float))) # Do each run num_for_averaging times to get an average with reduced noise for _ in range(num_for_averaging): for option_name, options in options_map.items(): # Run a Pipeline using DoFn's with simple typehints start = time() with beam.Pipeline(options=options) as p: pc = p | beam.Create(simple_elements) for ix in range(num_dofns): pc = ( pc | 'SimpleOutput %i' % ix >> beam.ParDo(SimpleOutput()) | 'SimpleInput %i' % ix >> beam.ParDo(SimpleInput())) timings[num_elements]['Simple Types'][option_name] += time() - start # Run a pipeline using DoFn's with nested typehints start = time() with beam.Pipeline(options=options) as p: pc = p | beam.Create(nested_elements) for ix in range(num_dofns): pc = ( pc | 'NestedOutput %i' % ix >> beam.ParDo(NestedOutput()) | 'NestedInput %i' % ix >> beam.ParDo(NestedInput())) timings[num_elements]['Nested Types'][option_name] += time() - start for num_elements, element_type_map in timings.items(): print("%d Element%s" % (num_elements, " " if num_elements == 1 else "s")) for element_type, option_name_map in element_type_map.items(): print("-- %s" % element_type) for option_name, time_elapsed in option_name_map.items(): print( "---- %.2f sec (%s)" % (time_elapsed / num_for_averaging, option_name)) print('\n') if __name__ == '__main__': logging.basicConfig() utils.check_compiled('apache_beam.runners.common') run_benchmark() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/tools/sideinput_microbenchmark.py0000644000175100017510000000533015206113033024441 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A microbenchmark for measuring performance of Side Inputs iterable. Runs side input iterable that fetches from multiple FakeSources, and measures the time to fetch all elements from all sources. Run as python -m apache_beam.tools.sideinput_microbenchmark """ # pytype: skip-file import logging import time from apache_beam.runners.worker import opcounters from apache_beam.runners.worker import sideinputs from apache_beam.runners.worker import statesampler from apache_beam.runners.worker.sideinputs_test import FakeSource from apache_beam.tools import utils from apache_beam.utils.counters import CounterFactory def long_generator(value, elements): for _ in range(elements): yield value def run_benchmark(num_runs=50, input_per_source=4000, num_sources=4): print("Number of runs:", num_runs) print("Input size:", num_sources * input_per_source) print("Sources:", num_sources) times = [] for _ in range(num_runs): counter_factory = CounterFactory() state_sampler = statesampler.StateSampler('basic', counter_factory) state_sampler.start() with state_sampler.scoped_state('step1', 'state'): si_counter = opcounters.SideInputReadCounter( counter_factory, state_sampler, 'step1', 1) si_counter = opcounters.NoOpTransformIOCounter() sources = [ FakeSource(long_generator(i, input_per_source)) for i in range(num_sources) ] iterator_fn = sideinputs.get_iterator_fn_for_sources( sources, read_counter=si_counter) start = time.time() list(iterator_fn()) time_cost = time.time() - start times.append(time_cost) state_sampler.stop() print("Runtimes:", times) avg_runtime = sum(times) / len(times) print("Average runtime:", avg_runtime) print("Time per element:", avg_runtime / (input_per_source * num_sources)) if __name__ == '__main__': logging.basicConfig() utils.check_compiled('apache_beam.runners.worker.opcounters') run_benchmark() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/tools/teststream_microbenchmark.py0000644000175100017510000001012615206113033024627 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A microbenchmark for measuring changes in the performance of TestStream running locally. This microbenchmark attempts to measure the overhead of the main data paths for the TestStream. Specifically new elements, watermark changes and processing time advances. This runs a series of N parallel pipelines with M parallel stages each. Each stage does the following: 1) Put all the PCollection elements in a window 2) Wait until the watermark advances past the end of the window. 3) When the watermark passes, change the key and output all the elements 4) Go back to #1 until all elements in the stream have been consumed. This executes the same codepaths that are run on the Fn API (and Dataflow) workers, but is generally easier to run (locally) and more stable. Run as python -m apache_beam.tools.teststream_microbenchmark """ # pytype: skip-file import argparse import itertools import logging import random import apache_beam as beam from apache_beam import WindowInto from apache_beam.runners import DirectRunner from apache_beam.testing.test_stream import TestStream from apache_beam.tools import utils from apache_beam.transforms.window import FixedWindows from apache_beam.typehints import typehints NUM_PARALLEL_STAGES = 7 NUM_SERIAL_STAGES = 6 class RekeyElements(beam.DoFn): def process(self, element): _, values = element return [(random.randint(0, 1000), v) for v in values] def _build_serial_stages(input_pc, num_serial_stages, stage_count): pc = (input_pc | ('gbk_start_stage%s' % stage_count) >> beam.GroupByKey()) for i in range(num_serial_stages): pc = ( pc | ('stage%s_map%s' % (stage_count, i)) >> beam.ParDo( RekeyElements()).with_output_types(typehints.KV[int, int]) | ('stage%s_gbk%s' % (stage_count, i)) >> beam.GroupByKey()) return pc def run_single_pipeline(size): def _pipeline_runner(): with beam.Pipeline(runner=DirectRunner()) as p: ts = TestStream().advance_watermark_to(0) all_elements = iter(range(size)) watermark = 0 while True: next_batch = list(itertools.islice(all_elements, 100)) if not next_batch: break ts = ts.add_elements([(i, random.randint(0, 1000)) for i in next_batch]) watermark = watermark + 100 ts = ts.advance_watermark_to(watermark) ts = ts.advance_watermark_to_infinity() input_pc = p | ts | WindowInto(FixedWindows(100)) for i in range(NUM_PARALLEL_STAGES): _build_serial_stages(input_pc, NUM_SERIAL_STAGES, i) return _pipeline_runner def run_benchmark( starting_point=1, num_runs=10, num_elements_step=300, verbose=True): suite = [ utils.LinearRegressionBenchmarkConfig( run_single_pipeline, starting_point, num_elements_step, num_runs) ] return utils.run_benchmarks(suite, verbose=verbose) if __name__ == '__main__': logging.basicConfig() utils.check_compiled('apache_beam.runners.common') parser = argparse.ArgumentParser() parser.add_argument('--num_runs', default=10, type=int) parser.add_argument('--starting_point', default=1, type=int) parser.add_argument('--increment', default=300, type=int) parser.add_argument('--verbose', default=True, type=bool) options = parser.parse_args() run_benchmark( options.starting_point, options.num_runs, options.increment, options.verbose) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/tools/utils.py0000644000175100017510000001456115206113033020537 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utility functions for all microbenchmarks.""" # pytype: skip-file import collections import gc import importlib import os import time from typing import Callable from typing import NamedTuple import numpy BenchmarkFn = Callable[[], None] BenchmarkFactoryFn = Callable[[int], BenchmarkFn] def check_compiled(module): """Check whether given module has been compiled. Args: module: string, module name """ check_module = importlib.import_module(module) ext = os.path.splitext(check_module.__file__)[-1] if ext in ('.py', '.pyc'): raise RuntimeError( "Profiling uncompiled code.\n" "To compile beam, run " "'pip install Cython; python setup.py build_ext --inplace'") class BenchmarkConfig(NamedTuple): """ Attributes: benchmark: a callable that takes an int argument - benchmark size, and returns a callable. A returned callable must run the code being benchmarked on an input of specified size. For example, one can implement a benchmark as: class MyBenchmark(object): def __init__(self, size): [do necessary initialization] def __call__(self): [run the code in question] size: int, a size of the input. Aggregated per-element metrics are counted based on the size of the input. num_runs: int, number of times to run each benchmark. """ benchmark: BenchmarkFactoryFn size: int num_runs: int def __str__(self): return "%s, %s element(s)" % ( getattr(self.benchmark, '__name__', str(self.benchmark)), str(self.size)) class LinearRegressionBenchmarkConfig(NamedTuple): """ Attributes: benchmark: a callable that takes an int argument - benchmark size, and returns a callable. A returned callable must run the code being benchmarked on an input of specified size. For example, one can implement a benchmark as: class MyBenchmark(object): def __init__(self, size): [do necessary initialization] def __call__(self): [run the code in question] starting_point: int, an initial size of the input. Regression results are calculated based on the input. increment: int, the rate of growth of the input for each run of the benchmark. num_runs: int, number of times to run each benchmark. """ benchmark: Callable[[int], BenchmarkFn] starting_point: int increment: int num_runs: int def __str__(self): return "%s, %s element(s) at start, %s growth per run" % ( getattr(self.benchmark, '__name__', str(self.benchmark)), str(self.starting_point), str(self.increment)) def run_benchmarks(benchmark_suite, verbose=True): """Runs benchmarks, and collects execution times. A simple instrumentation to run a callable several times, collect and print its execution times. Args: benchmark_suite: A list of BenchmarkConfig. verbose: bool, whether to print benchmark results to stdout. Returns: A dictionary of the form string -> list of floats. Keys of the dictionary are benchmark names, values are execution times in seconds for each run. """ def run(benchmark: BenchmarkFactoryFn, size: int): # Contain each run of a benchmark inside a function so that any temporary # objects can be garbage-collected after the run. benchmark_instance_callable = benchmark(size) start = time.time() _ = benchmark_instance_callable() return time.time() - start cost_series = collections.defaultdict(list) size_series = collections.defaultdict(list) for benchmark_config in benchmark_suite: name = str(benchmark_config) num_runs = benchmark_config.num_runs if isinstance(benchmark_config, LinearRegressionBenchmarkConfig): size = benchmark_config.starting_point step = benchmark_config.increment else: assert isinstance(benchmark_config, BenchmarkConfig) size = benchmark_config.size step = 0 for run_id in range(num_runs): # Do a proactive GC before each run to minimize side-effects of different # runs. gc.collect() time_cost = run(benchmark_config.benchmark, size) # Appending size and time cost to perform linear regression cost_series[name].append(time_cost) size_series[name].append(size) if verbose: per_element_cost = time_cost / size print( "%s: run %d of %d, per element time cost: %g sec" % (name, run_id + 1, num_runs, per_element_cost)) # Incrementing the size of the benchmark run by the step size size += step if verbose: print("") if verbose: pad_length = max([len(str(bc)) for bc in benchmark_suite]) for benchmark_config in benchmark_suite: name = str(benchmark_config) if isinstance(benchmark_config, LinearRegressionBenchmarkConfig): from scipy import stats print() # pylint: disable=unused-variable gradient, intercept, r_value, p_value, std_err = stats.linregress( size_series[name], cost_series[name]) print("Fixed cost ", intercept) print("Per-element ", gradient) print("R^2 ", r_value**2) else: assert isinstance(benchmark_config, BenchmarkConfig) per_element_median_cost = ( numpy.median(cost_series[name]) / benchmark_config.size) std = numpy.std(cost_series[name]) / benchmark_config.size print( "%s: p. element median time cost: %g sec, relative std: %.2f%%" % ( name.ljust(pad_length, " "), per_element_median_cost, std * 100 / per_element_median_cost)) return size_series, cost_series ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.753447 apache_beam-2.74.0/apache_beam/transforms/0000755000175100017510000000000015206113160020055 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/__init__.py0000644000175100017510000000271415206113033022171 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """PTransform and descendants.""" # pylint: disable=wildcard-import from apache_beam.transforms import combiners from apache_beam.transforms.core import * from apache_beam.transforms.external import * from apache_beam.transforms.managed import * from apache_beam.transforms.ptransform import * from apache_beam.transforms.stats import * from apache_beam.transforms.timeutil import TimeDomain from apache_beam.transforms.util import * # No backwards compatibility guarantees. try: from apache_beam.transforms.cy_dataflow_distribution_counter import DataflowDistributionCounter except ImportError: from apache_beam.transforms.py_dataflow_distribution_counter import DataflowDistributionCounter ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/async_dofn.py0000644000175100017510000005163315206113033022561 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from __future__ import absolute_import import asyncio import inspect import logging import random import threading import uuid from collections.abc import AsyncIterable from collections.abc import Iterable from concurrent.futures import ThreadPoolExecutor from math import floor from threading import RLock from time import sleep from time import time from types import GeneratorType from typing import Optional import apache_beam as beam from apache_beam import TimeDomain from apache_beam.coders import coders from apache_beam.transforms.userstate import BagStateSpec from apache_beam.transforms.userstate import ReadModifyWriteStateSpec from apache_beam.transforms.userstate import TimerSpec from apache_beam.transforms.userstate import on_timer from apache_beam.utils.shared import Shared from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import Timestamp # A wrapper around a dofn that processes that dofn in an asynchronous manner. class AsyncWrapper(beam.DoFn): """Class that wraps a dofn and converts it from one which process elements synchronously to one which processes them asynchronously. For synchronous dofns the default settings mean that many (100s) of elements will be processed in parallel and that processing an element will block all other work on that key. In addition runners are optimized for latencies less than a few seconds and longer operations can result in high retry rates. Async should be considered when the default parallelism is not correct and/or items are expected to take longer than a few seconds to process. """ TIMER = TimerSpec('timer', TimeDomain.REAL_TIME) TIMER_SET = ReadModifyWriteStateSpec('timer_set', coders.BooleanCoder()) TO_PROCESS = BagStateSpec( 'to_process', coders.TupleCoder( [coders.FastPrimitivesCoder(), coders.FastPrimitivesCoder()])) # The below items are one per dofn (not instance) so are maps of UUID to # value. _event_loop: Optional[asyncio.AbstractEventLoop] = None _event_loop_thread: Optional[threading.Thread] = None _loop_started: Optional[threading.Event] = None _processing_elements = {} _items_in_buffer = {} _pool = {} # Must be reentrant lock so that callbacks can either be called by a pool # thread OR the thread submitting the callback which must already hold this # lock. _lock = RLock() _verbose_logging = False def __init__( self, sync_fn, parallelism=1, callback_frequency=5, max_items_to_buffer=None, timeout=1, max_wait_time=0.5, id_fn=None, use_asyncio=False, ): """Wraps the sync_fn to create an asynchronous version. Args: sync_fn: The dofn to wrap. Must take (K, V) as input. parallelism: The maximum number of elements to process in parallel per worker for this dofn. By default this is set to 1 as the most common case for async dofns are heavy CPU or GPU dofns where the dofn requires the a signficant fraction of the CPU/GPU. callback_frequency: The frequency with which the runner will check for elements to commit. A short callback frequency will mean items are commited shortly after processing but cause additional work for the worker. A large callback frequency will result in slower commits but less busy work. The default of 5s will result in a maximum added latency of 5s while requiring relatively few resources. If your messages take significantly longer than 5s to process it is recommended to raise this. max_items_to_buffer: We should ideally buffer enough to always be busy but not so much that the worker ooms. By default will be 2x the parallelism which should be good for most pipelines. timeout: The maximum amount of time an item should try to be scheduled locally before it goes in the queue of waiting work. max_wait_time: The maximum amount of sleep time while attempting to schedule an item. Used in testing to ensure timeouts are met. id_fn: A function that returns a hashable object from an element. This will be used to track items instead of the element's default hash. use_asyncio: If true, use asyncio and coroutines to process items. If false, use ThreadPoolExecutor. Use asyncio when the work being done is not CPU intensive and heavily waits on network or IO which can benefit from higher parallelism. """ self._sync_fn = sync_fn self._uuid = uuid.uuid4().hex self._parallelism = parallelism self._timeout = timeout self._max_wait_time = max_wait_time self._timer_frequency = callback_frequency self._id_fn = id_fn or (lambda x: x) self._use_asyncio = use_asyncio if max_items_to_buffer is None: self._max_items_to_buffer = max(parallelism * 2, 10) else: self._max_items_to_buffer = max_items_to_buffer AsyncWrapper._processing_elements[self._uuid] = {} AsyncWrapper._items_in_buffer[self._uuid] = 0 self.max_wait_time = max_wait_time self._shared_handle = Shared() @staticmethod def initialize_pool(parallelism): return lambda: ThreadPoolExecutor(max_workers=parallelism) @staticmethod def _run_event_loop(): """Sets up and runs the asyncio event loop in a background thread.""" loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) AsyncWrapper._event_loop = loop AsyncWrapper._loop_started.set() loop.run_forever() loop.close() @staticmethod def reset_state(): event_loop_thread_to_join = None with AsyncWrapper._lock: if AsyncWrapper._event_loop: AsyncWrapper._event_loop.call_soon_threadsafe( AsyncWrapper._event_loop.stop) if AsyncWrapper._event_loop_thread: event_loop_thread_to_join = AsyncWrapper._event_loop_thread AsyncWrapper._event_loop = None AsyncWrapper._event_loop_thread = None if AsyncWrapper._loop_started is not None: AsyncWrapper._loop_started.clear() pools = list(AsyncWrapper._pool.values()) # We must join the asyncio event loop thread outside of the lock block. # If joined inside the lock, the waiting thread holds the lock while blocking, # preventing active coroutines' done callbacks from acquiring the lock on the # event loop thread, resulting in a deadlock. if event_loop_thread_to_join: event_loop_thread_to_join.join() # We must acquire and shut down the thread pools outside of the lock block. # If shutdown(wait=True) is called inside the lock, the caller blocks holding # the lock, preventing active worker threads from acquiring the lock to run # their done callbacks, resulting in a deadlock. pools_to_shutdown = [ pool.acquire(AsyncWrapper.initialize_pool(1)) for pool in pools ] for pool in pools_to_shutdown: pool.shutdown(wait=True, cancel_futures=True) with AsyncWrapper._lock: AsyncWrapper._pool = {} AsyncWrapper._processing_elements = {} AsyncWrapper._items_in_buffer = {} def setup(self): """Forwards to the wrapped dofn's setup method.""" self._sync_fn.setup() with AsyncWrapper._lock: if self._use_asyncio and AsyncWrapper._event_loop_thread is None: AsyncWrapper._loop_started = threading.Event() AsyncWrapper._event_loop_thread = threading.Thread( target=AsyncWrapper._run_event_loop, daemon=True) AsyncWrapper._event_loop_thread.start() AsyncWrapper._loop_started.wait() if not self._uuid in AsyncWrapper._pool: AsyncWrapper._pool[self._uuid] = Shared() AsyncWrapper._processing_elements[self._uuid] = {} AsyncWrapper._items_in_buffer[self._uuid] = 0 def teardown(self): """Forwards to the wrapped dofn's teardown method.""" self._sync_fn.teardown() def sync_fn_process(self, element, *args, **kwargs): """Makes the call to the wrapped dofn's start_bundle, process methods. It will then combine the results into a single generator. Args: element: The element to process. *args: Any additional arguments to pass to the wrapped dofn's process method. Will be the same args that the async wrapper is called with. **kwargs: Any additional keyword arguments to pass to the wrapped dofn's process method. Will be the same kwargs that the async wrapper is called with. Returns: A generator of elements produced by the input element. """ self._sync_fn.start_bundle() process_result = self._sync_fn.process(element, *args, **kwargs) bundle_result = self._sync_fn.finish_bundle() # both process and finish bundle may or may not return generators. We want # to combine whatever results have been returned into a single generator. If # they are single elements then wrap them in lists so that we can combine # them. if not process_result: process_result = [] elif not isinstance(process_result, GeneratorType): process_result = [process_result] if not bundle_result: bundle_result = [] elif not isinstance(bundle_result, GeneratorType): bundle_result = [bundle_result] to_return = [] for x in process_result: to_return.append(x) for x in bundle_result: to_return.append(x) return to_return async def async_fn_process(self, element, *args, **kwargs): """Makes the call to the wrapped dofn's start_bundle, process and finish_bundle methods for asynchronous DoFns. Args: element: The element to process. *args: Any additional arguments to pass to the wrapped dofn's process method. **kwargs: Any additional keyword arguments to pass to the wrapped dofn's process method. Returns: A list of elements produced by the input element. """ async def _collect(result): if result is None: return [] if inspect.isawaitable(result): result = await result if isinstance(result, AsyncIterable): return [item async for item in result] if isinstance(result, (GeneratorType, Iterable)) and not isinstance(result, (str, bytes)): return list(result) return [result] self._sync_fn.start_bundle() process_result = await _collect( self._sync_fn.process(element, *args, **kwargs)) bundle_result = await _collect(self._sync_fn.finish_bundle()) return process_result + bundle_result def decrement_items_in_buffer(self, future): with AsyncWrapper._lock: if self._uuid in AsyncWrapper._items_in_buffer: AsyncWrapper._items_in_buffer[self._uuid] -= 1 def schedule_if_room(self, element, ignore_buffer=False, *args, **kwargs): """Schedules an item to be processed asynchronously if there is room. Args: element: The element to process. ignore_buffer: If true will ignore the buffer limit and schedule the item regardless of the buffer size. Used when an item needs to skip to the front such as retries. *args: arguments that the wrapped dofn requires. **kwargs: keyword arguments that the wrapped dofn requires. Returns: True if the item was scheduled False otherwise. """ with AsyncWrapper._lock: element_id = self._id_fn(element[1]) if element_id in AsyncWrapper._processing_elements[self._uuid]: logging.info('item %s already in processing elements', element) return True if self.accepting_items() or ignore_buffer: if self._use_asyncio: result = asyncio.run_coroutine_threadsafe( self.async_fn_process(element, *args, **kwargs), AsyncWrapper._event_loop, ) else: result = AsyncWrapper._pool[self._uuid].acquire( AsyncWrapper.initialize_pool(self._parallelism)).submit( lambda: self.sync_fn_process(element, *args, **kwargs), ) result.add_done_callback(self.decrement_items_in_buffer) AsyncWrapper._processing_elements[self._uuid][element_id] = ( element, result) AsyncWrapper._items_in_buffer[self._uuid] += 1 return True else: return False # Add an item to the processing pool. Add the future returned by that item to # processing_elements_. def schedule_item(self, element, ignore_buffer=False, *args, **kwargs): """Schedules an item to be processed asynchronously. If the queue is full will block until room opens up. After calling AsyncWrapper will hold a future pointing to the result of this processing Args: element: The element to process. ignore_buffer: If true will ignore the buffer limit and schedule the item regardless of the buffer size. Used when an item needs to skip to the front such as retries. *args: arguments that the wrapped dofn requires. **kwargs: keyword arguments that the wrapped dofn requires. """ done = False sleep_time = 0.01 total_sleep = 0 while not done and total_sleep < self._timeout: done = self.schedule_if_room(element, ignore_buffer, *args, **kwargs) if not done: sleep_time = min(self.max_wait_time, sleep_time * 2) if self._verbose_logging or total_sleep > 10: logging.info( 'buffer is full for item %s, %s waiting %s seconds. Have waited' ' for %s seconds.', element, AsyncWrapper._items_in_buffer[self._uuid], sleep_time, total_sleep, ) total_sleep += sleep_time sleep(sleep_time) def next_time_to_fire(self, key): random.seed(key) return ( floor((time() + self._timer_frequency) / self._timer_frequency) * self._timer_frequency) + ( random.random() * self._timer_frequency) def accepting_items(self): with AsyncWrapper._lock: return ( AsyncWrapper._items_in_buffer[self._uuid] < self._max_items_to_buffer) def is_empty(self): with AsyncWrapper._lock: return AsyncWrapper._items_in_buffer[self._uuid] == 0 # Add the incoming element to a pool of elements to process asynchronously. def process( self, element, timer=beam.DoFn.TimerParam(TIMER), to_process=beam.DoFn.StateParam(TO_PROCESS), *args, **kwargs): """Add the elements to the list of items to be processed asynchronously. Performs additional bookkeeping to maintain exactly once and set timers to commit item after it has finished processing. Args: element: The element to process. timer: Callback timer that will commit elements. to_process: State that keeps track of queued items for exactly once. *args: arguments that the wrapped dofn requires. **kwargs: keyword arguments that the wrapped dofn requires. Returns: An empty list. The elements will be output asynchronously. """ self.schedule_item(element) to_process.add(element) # Set a timer to fire on the next round increment of timer_frequency_. Note # we do this so that each messages timer doesn't get overwritten by the # next. time_to_fire = self.next_time_to_fire(element[0]) timer.set(time_to_fire) # Don't output any elements. This will be done in commit_finished_items. return [] # Synchronises local state (processing_elements_) with SE state (to_process). # Then outputs all finished elements. Finally, sets a timer to fire on the # next round increment of timer_frequency_. def commit_finished_items( self, to_process=beam.DoFn.StateParam(TO_PROCESS), timer=beam.DoFn.TimerParam(TIMER), ): """Commits finished items and synchronizes local state with runner state. Note timer firings are per key while local state contains messages for all keys. Only messages for the given key will be output/cleaned up. Args: to_process: State that keeps track of queued messagees for this key. timer: Timer that initiated this commit and can be reset if not all items have finished.. Returns: A list of elements that have finished processing for this key. """ # For all elements that are in processing state: # If the element is done processing, delete it from all state and yield the # output. # If the element is not yet done, print it. If the element is not in # local state, schedule it for processing. items_finished = 0 items_not_yet_finished = 0 items_rescheduled = 0 items_cancelled = 0 items_in_processing_state = 0 items_in_se_state = 0 to_process_local = list(to_process.read()) key = None to_reschedule = [] if to_process_local: key = str(to_process_local[0][0]) else: logging.error( 'no elements in state during timer callback. Timer should not have' ' been set.') if self._verbose_logging: logging.info('procesing timer for key: %s', key) # processing state is per key so we expect this state to only contain a # given key. Skip items in processing_elements which are for a different # key. with AsyncWrapper._lock: processing_elements = AsyncWrapper._processing_elements[self._uuid] to_process_local_ids = {self._id_fn(e[1]) for e in to_process_local} to_remove_ids = [] for element_id, (element, future) in processing_elements.items(): if element[0] == key and element_id not in to_process_local_ids: items_cancelled += 1 future.cancel() to_remove_ids.append(element_id) logging.info( 'cancelling item %s which is no longer in processing state', element) for element_id in to_remove_ids: processing_elements.pop(element_id) # For all elements which have finished processing output their result. to_return = [] finished_items = [] for x in to_process_local: items_in_se_state += 1 x_id = self._id_fn(x[1]) if x_id in processing_elements: _, future = processing_elements[x_id] if future.done(): to_return.append(future.result()) finished_items.append(x) processing_elements.pop(x_id) items_finished += 1 else: items_not_yet_finished += 1 else: logging.info( 'item %s found in processing state but not local state,' ' scheduling now', x) to_reschedule.append(x) items_rescheduled += 1 # Reschedule the items not under a lock for x in to_reschedule: self.schedule_item(x, ignore_buffer=False) # Update processing state to remove elements we've finished to_process.clear() for x in to_process_local: if x not in finished_items: items_in_processing_state += 1 to_process.add(x) logging.info('items finished %d', items_finished) logging.info('items not yet finished %d', items_not_yet_finished) logging.info('items rescheduled %d', items_rescheduled) logging.info('items cancelled %d', items_cancelled) logging.info('items in processing state %d', items_in_processing_state) logging.info( 'items in buffer %d', AsyncWrapper._items_in_buffer[self._uuid]) # If there are items not yet finished then set a timer to fire in the # future. self._next_time_to_fire = Timestamp.now() + Duration(seconds=5) if items_in_processing_state > 0: time_to_fire = self.next_time_to_fire(key) timer.set(time_to_fire) # Each result is a list. We want to combine them into a single # list of all elements we wish to output. merged_return = [] for x in to_return: merged_return.extend(x) return merged_return @on_timer(TIMER) def timer_callback( self, to_process=beam.DoFn.StateParam(TO_PROCESS), timer=beam.DoFn.TimerParam(TIMER), ): """Helper method to commit finished items in response to timer firing. Args: to_process: State that keeps track of queued items for exactly once. timer: Timer that initiated this commit and can be reset if not all items have finished. Returns: A generator of elements that have finished processing for this key. """ return self.commit_finished_items(to_process, timer) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/async_dofn_test.py0000644000175100017510000004412315206113033023614 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import multiprocessing import random import time import unittest from concurrent.futures import ThreadPoolExecutor from threading import Lock from parameterized import parameterized_class import apache_beam as beam import apache_beam.transforms.async_dofn as async_lib class BasicDofn(beam.DoFn): def __init__(self, sleep_time=0): self.processed = 0 self.sleep_time = sleep_time self.lock = Lock() def process(self, element): logging.info('start processing element %s', element) time.sleep(self.sleep_time) self.lock.acquire() self.processed += 1 self.lock.release() logging.info('finish processing element %s', element) yield element def getProcessed(self): self.lock.acquire() result = self.processed self.lock.release() return result # Used for testing multiple elements produced by a single input element. class MultiElementDoFn(beam.DoFn): def finish_bundle(self): yield ('key', 'bundle end') def process(self, element): yield element yield element class FakeBagState: def __init__(self, items): self.items = items # Normally SE would have a lock on the BT row protecting this from multiple # updates. Here without SE we must lock ourselves. self.lock = Lock() def add(self, item): with self.lock: self.items.append(item) def clear(self): with self.lock: self.items = [] def read(self): with self.lock: return self.items.copy() class FakeTimer: def __init__(self, time): self.time = time def set(self, time): self.time = time @parameterized_class([ { "use_asyncio": True }, { "use_asyncio": False }, ]) class AsyncTest(unittest.TestCase): def setUp(self): super().setUp() async_lib.AsyncWrapper.reset_state() def wait_for_empty(self, async_dofn, timeout=10): count = 0 increment = 1 while not async_dofn.is_empty(): time.sleep(1) count += increment if count > timeout: raise TimeoutError('Timed out waiting for async dofn to be empty') # Usually we are waiting for an item to be ready to output not just no # longer in the queue. These are not the same atomic operation. Sleep a bit # to make sure the item is also ready to output. time.sleep(1) def check_output(self, result, expected_output): output = [] for x in result: output.append(x) # Ordering is not guaranteed so sort both lists before comparing. output.sort() expected_output.sort() self.assertEqual(output, expected_output) def check_items_in_buffer(self, async_dofn, expected_count): self.assertEqual( async_lib.AsyncWrapper._items_in_buffer[async_dofn._uuid], expected_count, ) def test_custom_id_fn(self): class CustomIdObject: def __init__(self, element_id, value): self.element_id = element_id self.value = value def __hash__(self): return hash(self.element_id) def __eq__(self, other): return self.element_id == other.element_id dofn = BasicDofn() async_dofn = async_lib.AsyncWrapper( dofn, id_fn=lambda x: x.element_id, use_asyncio=self.use_asyncio) async_dofn.setup() fake_bag_state = FakeBagState([]) fake_timer = FakeTimer(0) msg1 = ('key1', CustomIdObject(1, 'a')) msg2 = ('key1', CustomIdObject(1, 'b')) result = async_dofn.process( msg1, to_process=fake_bag_state, timer=fake_timer) self.assertEqual(result, []) # The second message should be a no-op as it has the same id. result = async_dofn.process( msg2, to_process=fake_bag_state, timer=fake_timer) self.assertEqual(result, []) self.wait_for_empty(async_dofn) result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) self.check_output(result, [('key1', msg1[1])]) self.assertEqual(fake_bag_state.items, []) def test_basic(self): # Setup an async dofn and send a message in to process. dofn = BasicDofn() async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio) async_dofn.setup() fake_bag_state = FakeBagState([]) fake_timer = FakeTimer(0) msg = ('key1', 1) result = async_dofn.process( msg, to_process=fake_bag_state, timer=fake_timer) # nothing should be produced because this is async. self.assertEqual(result, []) # there should be an element to process in the bag state. self.assertEqual(fake_bag_state.items, [msg]) # The timer should be set. self.assertNotEqual(fake_timer.time, 0) # Give time for the async element to process. self.wait_for_empty(async_dofn) # Fire the (fake) timer and check that the input message has produced output # and the message has been removed from bag state. result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) self.check_output(result, [msg]) self.assertTrue(dofn.getProcessed() == 1) self.assertEqual(fake_bag_state.items, []) def test_multi_key(self): # Send in two messages with different keys. dofn = BasicDofn() async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio) async_dofn.setup() fake_bag_state_key1 = FakeBagState([]) fake_bag_state_key2 = FakeBagState([]) fake_timer = FakeTimer(0) msg1 = ('key1', 1) msg2 = ('key2', 2) async_dofn.process(msg1, to_process=fake_bag_state_key1, timer=fake_timer) async_dofn.process(msg2, to_process=fake_bag_state_key2, timer=fake_timer) self.wait_for_empty(async_dofn) # The result should depend on which key the timer fires for. We are firing # 'key2' first so expect its message to be produced and msg1 to remain in # state. result = async_dofn.commit_finished_items(fake_bag_state_key2, fake_timer) self.check_output(result, [msg2]) self.assertEqual(fake_bag_state_key1.items, [msg1]) self.assertEqual(fake_bag_state_key2.items, []) # Now fire 'key1' and expect its message to be produced. result = async_dofn.commit_finished_items(fake_bag_state_key1, fake_timer) self.check_output(result, [msg1]) self.assertEqual(fake_bag_state_key1.items, []) self.assertEqual(fake_bag_state_key2.items, []) def test_long_item(self): # Test that everything still works with a long running time for the dofn. dofn = BasicDofn(sleep_time=5) async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio) async_dofn.setup() fake_bag_state = FakeBagState([]) fake_timer = FakeTimer(0) msg = ('key1', 1) async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer) result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) self.check_output(result, []) self.assertTrue(dofn.getProcessed() == 0) self.assertEqual(fake_bag_state.items, [msg]) self.wait_for_empty(async_dofn, 20) result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) self.check_output(result, [msg]) self.assertTrue(dofn.getProcessed() == 1) self.assertEqual(fake_bag_state.items, []) def test_lost_item(self): # Setup an element in the bag state that's not in processing state. # The async dofn should reschedule this element. dofn = BasicDofn() async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio) async_dofn.setup() fake_timer = FakeTimer(0) msg = ('key1', 1) fake_bag_state = FakeBagState([msg]) result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) self.check_output(result, []) # The element should be rescheduled and still processed. self.wait_for_empty(async_dofn) result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) self.check_output(result, [msg]) def test_cancelled_item(self): # Test that an item gets removed for processing and does not get output when # it is not present in the bag state. Either this item moved or a commit # failed making the local state and bag state inconsistent. dofn = BasicDofn() async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio) async_dofn.setup() msg = ('key1', 1) msg2 = ('key1', 2) fake_timer = FakeTimer(0) fake_bag_state = FakeBagState([]) async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer) async_dofn.process(msg2, to_process=fake_bag_state, timer=fake_timer) self.wait_for_empty(async_dofn) fake_bag_state.clear() fake_bag_state.add(msg2) result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) self.check_output(result, [msg2]) self.assertEqual(fake_bag_state.items, []) def test_multi_element_dofn(self): # Test that async works when a dofn produces multiple elements in process # and finish_bundle. dofn = MultiElementDoFn() async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio) async_dofn.setup() fake_bag_state = FakeBagState([]) fake_timer = FakeTimer(0) msg = ('key1', 1) async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer) self.wait_for_empty(async_dofn) result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) self.check_output(result, [msg, msg, ('key', 'bundle end')]) self.assertEqual(fake_bag_state.items, []) def test_duplicates(self): # Test that async will produce a single output when a given input is sent # multiple times. dofn = BasicDofn(5) async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio) async_dofn.setup() fake_bag_state = FakeBagState([]) fake_timer = FakeTimer(0) msg = ('key1', 1) async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer) # SE will only deliver the message again if the commit failed and so bag # state should still be empty. fake_bag_state.clear() async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer) self.assertEqual(fake_bag_state.items, [msg]) self.wait_for_empty(async_dofn) result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) self.check_output(result, [msg]) self.assertEqual(fake_bag_state.items, []) def test_slow_duplicates(self): # Test that async will produce a single output when a given input is sent # multiple times. dofn = BasicDofn(5) async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio) async_dofn.setup() fake_bag_state = FakeBagState([]) fake_timer = FakeTimer(0) msg = ('key1', 1) async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer) time.sleep(10) # SE will only deliver the message again if the commit failed and so bag # state should still be empty. fake_bag_state.clear() result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) self.check_output(result, []) self.assertEqual(fake_bag_state.items, []) async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer) self.assertEqual(fake_bag_state.items, [msg]) self.wait_for_empty(async_dofn) result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) self.check_output(result, [msg]) self.assertEqual(fake_bag_state.items, []) def test_buffer_count(self): # Test that the buffer count is correctly incremented when adding items. dofn = BasicDofn(5) async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio) async_dofn.setup() msg = ('key1', 1) fake_timer = FakeTimer(0) fake_bag_state = FakeBagState([]) async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer) self.check_items_in_buffer(async_dofn, 1) self.wait_for_empty(async_dofn) self.check_items_in_buffer(async_dofn, 0) # Commiting already finished items should not change the count. async_dofn.commit_finished_items(fake_bag_state, fake_timer) self.check_items_in_buffer(async_dofn, 0) def test_buffer_stops_accepting_items(self): # Test that the buffer stops accepting items when it is full. dofn = BasicDofn(5) async_dofn = async_lib.AsyncWrapper( dofn, parallelism=1, max_items_to_buffer=5, use_asyncio=self.use_asyncio) async_dofn.setup() fake_timer = FakeTimer(0) fake_bag_state = FakeBagState([]) # Send more elements than the buffer can hold. pool = ThreadPoolExecutor(max_workers=10) expected_output = [] for i in range(0, 10): def add_item(i): item = ('key', i) expected_output.append(item) pool.submit( lambda: async_dofn.process( item, to_process=fake_bag_state, timer=fake_timer)) add_item(i) # Assert that the buffer stops accepting items at its limit. time.sleep(1) self.assertEqual(async_dofn._max_items_to_buffer, 5) self.check_items_in_buffer(async_dofn, 5) # Wait for all buffered items to finish. self.wait_for_empty(async_dofn, 100) # This will commit buffered items and add new items which didn't fit in the # buffer. result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) # Wait for the new buffered items to finish. self.wait_for_empty(async_dofn, 100) result.extend(async_dofn.commit_finished_items(fake_bag_state, fake_timer)) self.check_output(result, expected_output) self.check_items_in_buffer(async_dofn, 0) def test_buffer_with_cancellation(self): dofn = BasicDofn(3) async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio) async_dofn.setup() msg = ('key1', 1) msg2 = ('key1', 2) fake_timer = FakeTimer(0) fake_bag_state = FakeBagState([]) async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer) async_dofn.process(msg2, to_process=fake_bag_state, timer=fake_timer) self.check_items_in_buffer(async_dofn, 2) self.assertEqual(fake_bag_state.items, [msg, msg2]) # Remove one item from the bag state to simulate the user worker being out # of sync with SE state. The item should be cancelled and also removed from # the buffer. fake_bag_state.clear() fake_bag_state.add(msg2) result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) # Note buffer may still be 2 as msg1 can start processing before we get to # cancel it but also not be finished yet. self.check_output(result, []) self.assertEqual(fake_bag_state.items, [msg2]) self.wait_for_empty(async_dofn) result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) self.check_items_in_buffer(async_dofn, 0) self.check_output(result, [msg2]) self.assertEqual(fake_bag_state.items, []) def test_load_correctness(self): # Test AsyncDofn over heavy load. dofn = BasicDofn(1) max_sleep = 10 async_dofn = async_lib.AsyncWrapper( dofn, max_wait_time=max_sleep, use_asyncio=self.use_asyncio) async_dofn.setup() bag_states = {} timers = {} expected_outputs = {} pool = ThreadPoolExecutor(max_workers=10) futures = [] for i in range(0, 10): bag_states['key' + str(i)] = FakeBagState([]) timers['key' + str(i)] = FakeTimer(0) expected_outputs['key' + str(i)] = [] for i in range(0, 100): def add_item(i): item = ('key' + str(random.randint(0, 9)), i) expected_outputs[item[0]].append(item) futures.append( pool.submit( lambda: async_dofn.process( item, to_process=bag_states[item[0]], timer=timers[item[0]]) )) add_item(i) time.sleep(random.random()) # Run for a while. Should be enough to start all items but not finish them # all. time.sleep(random.randint(30, 50)) done = False results = [[] for _ in range(0, 10)] while not done: done = True for i in range(0, 10): results[i].extend( async_dofn.commit_finished_items( bag_states['key' + str(i)], timers['key' + str(i)])) if not bag_states['key' + str(i)].items: self.check_output(results[i], expected_outputs['key' + str(i)]) else: done = False time.sleep(random.randint(10, 30)) for i in range(0, 10): self.check_output(results[i], expected_outputs['key' + str(i)]) self.assertEqual(bag_states['key' + str(i)].items, []) @staticmethod def _run_reset_state_concurrent_teardown(use_asyncio): dofn = BasicDofn(sleep_time=0.5) async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=use_asyncio) async_dofn.setup() fake_bag_state = FakeBagState([]) fake_timer = FakeTimer(0) # Start processing an item. This starts a worker thread/coroutine sleeping for 0.5s. async_dofn.process(('key1', 1), to_process=fake_bag_state, timer=fake_timer) time.sleep(0.05) # Verify that calling reset_state() while background tasks are actively running # completes cleanly without causing lock-ordering deadlocks. async_lib.AsyncWrapper.reset_state() def test_reset_state_concurrent_teardown(self): # Verify concurrent teardown safety in a separate process to prevent any potential # regressions from freezing the main pytest process at exit. p = multiprocessing.Process( target=AsyncTest._run_reset_state_concurrent_teardown, args=(self.use_asyncio, )) p.start() p.join(timeout=10.0) if p.is_alive(): p.terminate() p.join() self.fail( "reset_state() deadlocked/hung waiting for active threads/tasks to finish" ) else: self.assertEqual(p.exitcode, 0) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/batch_dofn_test.py0000644000175100017510000002430315206113033023556 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """UnitTests for Batched DoFn (process_batch) API.""" # pytype: skip-file import unittest from collections.abc import Iterator from typing import no_type_check from parameterized import parameterized_class import apache_beam as beam class ElementDoFn(beam.DoFn): def process(self, element: int, *args, **kwargs) -> Iterator[float]: yield element / 2 class BatchDoFn(beam.DoFn): def process_batch(self, batch: list[int], *args, **kwargs) -> Iterator[list[float]]: yield [element / 2 for element in batch] class NoReturnAnnotation(beam.DoFn): def process_batch(self, batch: list[int], *args, **kwargs): yield [element * 2 for element in batch] class OverrideTypeInference(beam.DoFn): def process_batch(self, batch, *args, **kwargs): yield [element * 2 for element in batch] def get_input_batch_type(self, input_element_type): return list[input_element_type] def get_output_batch_type(self, input_element_type): return list[input_element_type] class EitherDoFn(beam.DoFn): def process(self, element: int, *args, **kwargs) -> Iterator[float]: yield element / 2 def process_batch(self, batch: list[int], *args, **kwargs) -> Iterator[list[float]]: yield [element / 2 for element in batch] class ElementToBatchDoFn(beam.DoFn): @beam.DoFn.yields_batches def process(self, element: int, *args, **kwargs) -> Iterator[list[int]]: yield [element] * element def infer_output_type(self, input_element_type): return input_element_type class BatchToElementDoFn(beam.DoFn): @beam.DoFn.yields_elements def process_batch(self, batch: list[int], *args, **kwargs) -> Iterator[tuple[int, int]]: yield (sum(batch), len(batch)) def get_test_class_name(cls, num, params_dict): return "%s_%s" % (cls.__name__, params_dict['dofn'].__class__.__name__) @parameterized_class([ { "dofn": ElementDoFn(), "input_element_type": int, "expected_process_defined": True, "expected_process_batch_defined": False, "expected_input_batch_type": None, "expected_output_batch_type": None }, { "dofn": BatchDoFn(), "input_element_type": int, "expected_process_defined": False, "expected_process_batch_defined": True, "expected_input_batch_type": beam.typehints.List[int], "expected_output_batch_type": beam.typehints.List[float] }, { "dofn": NoReturnAnnotation(), "input_element_type": int, "expected_process_defined": False, "expected_process_batch_defined": True, "expected_input_batch_type": beam.typehints.List[int], "expected_output_batch_type": beam.typehints.List[int] }, { "dofn": OverrideTypeInference(), "input_element_type": int, "expected_process_defined": False, "expected_process_batch_defined": True, "expected_input_batch_type": beam.typehints.List[int], "expected_output_batch_type": beam.typehints.List[int] }, { "dofn": EitherDoFn(), "input_element_type": int, "expected_process_defined": True, "expected_process_batch_defined": True, "expected_input_batch_type": beam.typehints.List[int], "expected_output_batch_type": beam.typehints.List[float] }, { "dofn": ElementToBatchDoFn(), "input_element_type": int, "expected_process_defined": True, "expected_process_batch_defined": False, "expected_input_batch_type": None, "expected_output_batch_type": beam.typehints.List[int] }, { "dofn": BatchToElementDoFn(), "input_element_type": int, "expected_process_defined": False, "expected_process_batch_defined": True, "expected_input_batch_type": beam.typehints.List[int], "expected_output_batch_type": None, }, ], class_name_func=get_test_class_name) class BatchDoFnParameterizedTest(unittest.TestCase): def test_process_defined(self): self.assertEqual(self.dofn._process_defined, self.expected_process_defined) def test_process_batch_defined(self): self.assertEqual( self.dofn._process_batch_defined, self.expected_process_batch_defined) def test_get_input_batch_type(self): self.assertEqual( self.dofn._get_input_batch_type_normalized(self.input_element_type), self.expected_input_batch_type) def test_get_output_batch_type(self): self.assertEqual( self.dofn._get_output_batch_type_normalized(self.input_element_type), self.expected_output_batch_type) def test_can_yield_batches(self): expected = self.expected_output_batch_type is not None self.assertEqual(self.dofn._can_yield_batches, expected) class NoInputAnnotation(beam.DoFn): def process_batch(self, batch, *args, **kwargs): yield [element * 2 for element in batch] class MismatchedBatchProducingDoFn(beam.DoFn): """A DoFn that produces batches from both process and process_batch, with mismatched return types (one yields floats, the other ints). Should yield a construction time error when applied.""" @beam.DoFn.yields_batches def process(self, element: int, *args, **kwargs) -> Iterator[list[int]]: yield [element] def process_batch(self, batch: list[int], *args, **kwargs) -> Iterator[list[float]]: yield [element / 2 for element in batch] class MismatchedElementProducingDoFn(beam.DoFn): """A DoFn that produces elements from both process and process_batch, with mismatched return types (one yields floats, the other ints). Should yield a construction time error when applied.""" def process(self, element: int, *args, **kwargs) -> Iterator[float]: yield element / 2 @beam.DoFn.yields_elements def process_batch(self, batch: list[int], *args, **kwargs) -> Iterator[int]: yield batch[0] class NoElementOutputAnnotation(beam.DoFn): def process_batch(self, batch: list[int], *args, **kwargs) -> Iterator[list[int]]: yield [element * 2 for element in batch] class BatchDoFnTest(unittest.TestCase): def test_map_pardo(self): # verify batch dofn accessors work well with beam.Map generated DoFn # checking this in parameterized test causes a circular reference issue dofn = beam.Map(lambda x: x * 2).dofn self.assertTrue(dofn._process_defined) self.assertFalse(dofn._process_batch_defined) self.assertEqual(dofn._get_input_batch_type_normalized(int), None) self.assertEqual(dofn._get_output_batch_type_normalized(int), None) def test_no_input_annotation_raises(self): p = beam.Pipeline() pc = p | beam.Create([1, 2, 3]) with self.assertRaisesRegex(TypeError, r'NoInputAnnotation.process_batch'): _ = pc | beam.ParDo(NoInputAnnotation()) def test_unsupported_dofn_param_raises(self): class BadParam(beam.DoFn): @no_type_check def process_batch(self, batch: list[int], key=beam.DoFn.KeyParam): yield batch * key p = beam.Pipeline() pc = p | beam.Create([1, 2, 3]) with self.assertRaisesRegex(NotImplementedError, r'BadParam.*KeyParam'): _ = pc | beam.ParDo(BadParam()) def test_mismatched_batch_producer_raises(self): p = beam.Pipeline() pc = p | beam.Create([1, 2, 3]) # Note (?ms) makes this a multiline regex, where . matches newlines. # See (?aiLmsux) at # https://docs.python.org/3.4/library/re.html#regular-expression-syntax with self.assertRaisesRegex( TypeError, (r'(?ms)MismatchedBatchProducingDoFn.*' r'process: List\[\].*process_batch: ' r'List\[\]')): _ = pc | beam.ParDo(MismatchedBatchProducingDoFn()) def test_mismatched_element_producer_raises(self): p = beam.Pipeline() pc = p | beam.Create([1, 2, 3]) # Note (?ms) makes this a multiline regex, where . matches newlines. # See (?aiLmsux) at # https://docs.python.org/3.4/library/re.html#regular-expression-syntax with self.assertRaisesRegex( TypeError, r'(?ms)MismatchedElementProducingDoFn.*process:.*process_batch:'): _ = pc | beam.ParDo(MismatchedElementProducingDoFn()) def test_cant_infer_batchconverter_input_raises(self): p = beam.Pipeline() pc = p | beam.Create(['a', 'b', 'c']) with self.assertRaisesRegex( TypeError, # Error should mention "input", and the name of the DoFn r'input.*BatchDoFn.*'): _ = pc | beam.ParDo(BatchDoFn()) def test_cant_infer_batchconverter_output_raises(self): p = beam.Pipeline() pc = p | beam.Create([1, 2, 3]) with self.assertRaisesRegex( TypeError, # Error should mention "output", the name of the DoFn, and suggest # overriding DoFn.infer_output_type r'output.*NoElementOutputAnnotation.*DoFn\.infer_output_type'): _ = pc | beam.ParDo(NoElementOutputAnnotation()) def test_element_to_batch_dofn_typehint(self): # Verify that element to batch DoFn sets the correct typehint on the output # PCollection. p = beam.Pipeline() pc = (p | beam.Create([1, 2, 3]) | beam.ParDo(ElementToBatchDoFn())) self.assertEqual(pc.element_type, int) def test_batch_to_element_dofn_typehint(self): # Verify that batch to element DoFn sets the correct typehint on the output # PCollection. p = beam.Pipeline() pc = (p | beam.Create([1, 2, 3]) | beam.ParDo(BatchToElementDoFn())) self.assertEqual(pc.element_type, beam.typehints.Tuple[int, int]) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/combinefn_lifecycle_pipeline.py0000644000175100017510000001461615206113033026302 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import math import apache_beam as beam from apache_beam.options.pipeline_options import TypeOptions from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import combiners from apache_beam.transforms import trigger from apache_beam.transforms import userstate from apache_beam.transforms import window from apache_beam.typehints import with_input_types from apache_beam.typehints import with_output_types @with_input_types(int) @with_output_types(int) class CallSequenceEnforcingCombineFn(beam.CombineFn): instances: set['CallSequenceEnforcingCombineFn'] = set() def __init__(self): super().__init__() self._setup_called = False self._teardown_called = False def setup(self, *args, **kwargs): assert not self._setup_called, 'setup should not be called twice' assert not self._teardown_called, 'setup should be called before teardown' # Keep track of instances so that we can check if teardown is called # properly after pipeline execution. self.instances.add(self) self._setup_called = True def create_accumulator(self, *args, **kwargs): assert self._setup_called, 'setup should have been called' assert not self._teardown_called, 'teardown should not have been called' return 0 def add_input(self, mutable_accumulator, element, *args, **kwargs): assert self._setup_called, 'setup should have been called' assert not self._teardown_called, 'teardown should not have been called' mutable_accumulator += element return mutable_accumulator def add_inputs(self, mutable_accumulator, elements, *args, **kwargs): return self.add_input(mutable_accumulator, sum(elements)) def merge_accumulators(self, accumulators, *args, **kwargs): assert self._setup_called, 'setup should have been called' assert not self._teardown_called, 'teardown should not have been called' return sum(accumulators) def extract_output(self, accumulator, *args, **kwargs): assert self._setup_called, 'setup should have been called' assert not self._teardown_called, 'teardown should not have been called' return accumulator def teardown(self, *args, **kwargs): assert self._setup_called, 'setup should have been called' assert not self._teardown_called, 'teardown should not be called twice' self._teardown_called = True @with_input_types(tuple[None, str]) @with_output_types(tuple[int, str]) class IndexAssigningDoFn(beam.DoFn): state_param = beam.DoFn.StateParam( userstate.CombiningValueStateSpec( 'index', beam.coders.VarIntCoder(), CallSequenceEnforcingCombineFn())) def process(self, element, state=state_param): _, value = element current_index = state.read() yield current_index, value state.add(1) def run_combine(pipeline, input_elements=5, lift_combiners=True): # Calculate the expected result, which is the sum of an arithmetic sequence. # By default, this is equal to: 0 + 1 + 2 + 3 + 4 = 10 expected_result = input_elements * (input_elements - 1) / 2 # Enable runtime type checking in order to cover TypeCheckCombineFn by # the test. pipeline.get_pipeline_options().view_as(TypeOptions).runtime_type_check = True pipeline.get_pipeline_options().view_as( TypeOptions).allow_unsafe_triggers = True with pipeline as p: pcoll = p | 'Start' >> beam.Create(range(input_elements)) # Certain triggers, such as AfterCount, are incompatible with combiner # lifting. We can use that fact to prevent combiners from being lifted. if not lift_combiners: pcoll |= beam.WindowInto( window.GlobalWindows(), trigger=trigger.AfterCount(input_elements), accumulation_mode=trigger.AccumulationMode.DISCARDING) # Pass an additional 'None' in order to cover _CurriedFn by the test. pcoll |= 'Do' >> beam.CombineGlobally( combiners.SingleInputTupleCombineFn( CallSequenceEnforcingCombineFn(), CallSequenceEnforcingCombineFn()), None).with_fanout(fanout=1) assert_that(pcoll, equal_to([(expected_result, expected_result)])) def run_combine_uncopyable_attr( pipeline, input_elements=5, lift_combiners=True): # Calculate the expected result, which is the sum of an arithmetic sequence. # By default, this is equal to: 0 + 1 + 2 + 3 + 4 = 10 expected_result = input_elements * (input_elements - 1) / 2 # Enable runtime type checking in order to cover TypeCheckCombineFn by # the test. pipeline.get_pipeline_options().view_as(TypeOptions).runtime_type_check = True pipeline.get_pipeline_options().view_as( TypeOptions).allow_unsafe_triggers = True with pipeline as p: pcoll = p | 'Start' >> beam.Create(range(input_elements)) # Certain triggers, such as AfterCount, are incompatible with combiner # lifting. We can use that fact to prevent combiners from being lifted. if not lift_combiners: pcoll |= beam.WindowInto( window.GlobalWindows(), trigger=trigger.AfterCount(input_elements), accumulation_mode=trigger.AccumulationMode.DISCARDING) combine_fn = CallSequenceEnforcingCombineFn() # Modules are not deep copyable. Ensure fanout falls back to pickling for # copying combine_fn. combine_fn.module_attribute = math pcoll |= 'Do' >> beam.CombineGlobally(combine_fn).with_fanout(fanout=1) assert_that(pcoll, equal_to([expected_result])) def run_pardo(pipeline, input_elements=10): with pipeline as p: _ = ( p | 'Start' >> beam.Create(('Hello' for _ in range(input_elements))) | 'KeyWithNone' >> beam.Map(lambda elem: (None, elem)) | 'Do' >> beam.ParDo(IndexAssigningDoFn())) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/combinefn_lifecycle_test.py0000644000175100017510000000740215206113033025447 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ValidatesRunner tests for CombineFn lifecycle and bundle methods.""" # pytype: skip-file import unittest import pytest from parameterized import parameterized_class from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.runners.direct import direct_runner from apache_beam.runners.portability import fn_api_runner from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.transforms.combinefn_lifecycle_pipeline import CallSequenceEnforcingCombineFn from apache_beam.transforms.combinefn_lifecycle_pipeline import run_combine from apache_beam.transforms.combinefn_lifecycle_pipeline import run_combine_uncopyable_attr from apache_beam.transforms.combinefn_lifecycle_pipeline import run_pardo @pytest.mark.it_validatesrunner class CombineFnLifecycleTest(unittest.TestCase): def setUp(self): self.pipeline = TestPipeline(is_integration_test=True) def test_combine(self): run_combine(self.pipeline) def test_non_liftable_combine(self): run_combine(self.pipeline, lift_combiners=False) def test_combining_value_state(self): if ('DataflowRunner' in self.pipeline.get_pipeline_options().view_as( StandardOptions).runner): self.skipTest('https://github.com/apache/beam/issues/20722') run_pardo(self.pipeline) @parameterized_class([ {'runner': direct_runner.BundleBasedDirectRunner, 'pickler': 'dill'}, {'runner': direct_runner.BundleBasedDirectRunner, 'pickler': 'cloudpickle'}, {'runner': fn_api_runner.FnApiRunner, 'pickler': 'dill'}, {'runner': fn_api_runner.FnApiRunner, 'pickler': 'cloudpickle'}, ]) # yapf: disable @pytest.mark.uses_dill class LocalCombineFnLifecycleTest(unittest.TestCase): def setUp(self): if self.pickler == 'dill': pytest.importorskip("dill") def tearDown(self): CallSequenceEnforcingCombineFn.instances.clear() def test_combine(self): test_options = PipelineOptions(flags=[f"--pickle_library={self.pickler}"]) run_combine(TestPipeline(runner=self.runner(), options=test_options)) self._assert_teardown_called() def test_combine_deepcopy_fails(self): test_options = PipelineOptions(flags=[f"--pickle_library={self.pickler}"]) run_combine_uncopyable_attr( TestPipeline(runner=self.runner(), options=test_options)) self._assert_teardown_called() def test_non_liftable_combine(self): test_options = PipelineOptions(flags=['--allow_unsafe_triggers']) run_combine( TestPipeline(runner=self.runner(), options=test_options), lift_combiners=False) self._assert_teardown_called() def test_combining_value_state(self): run_pardo(TestPipeline(runner=self.runner())) self._assert_teardown_called() def _assert_teardown_called(self): """Ensures that teardown has been invoked for all CombineFns.""" for instance in CallSequenceEnforcingCombineFn.instances: self.assertTrue(instance._teardown_called) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/combiners.py0000644000175100017510000010570215206113033022414 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A library of basic combiner PTransform subclasses.""" # pytype: skip-file import copy import heapq import itertools import operator import random from collections.abc import Iterable from typing import Any from typing import TypeVar from typing import Union import numpy as np from apache_beam import typehints from apache_beam.transforms import core from apache_beam.transforms import cy_combiners from apache_beam.transforms import ptransform from apache_beam.transforms import window from apache_beam.transforms.display import DisplayDataItem from apache_beam.typehints import with_input_types from apache_beam.typehints import with_output_types from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.windowed_value import WindowedValue __all__ = [ 'Count', 'Mean', 'Sample', 'Top', 'ToDict', 'ToList', 'ToSet', 'Latest', 'CountCombineFn', 'MeanCombineFn', 'SampleCombineFn', 'TopCombineFn', 'ToDictCombineFn', 'ToListCombineFn', 'ToSetCombineFn', 'LatestCombineFn', ] # Type variables T = TypeVar('T') K = TypeVar('K') V = TypeVar('V') TimestampType = Union[int, float, Timestamp, Duration] class CombinerWithoutDefaults(ptransform.PTransform): """Super class to inherit without_defaults to built-in Combiners.""" def __init__(self, has_defaults=True): super().__init__() self.has_defaults = has_defaults def with_defaults(self, has_defaults=True): new = copy.copy(self) new.has_defaults = has_defaults return new def without_defaults(self): return self.with_defaults(False) class Mean(object): """Combiners for computing arithmetic means of elements.""" class Globally(CombinerWithoutDefaults): """combiners.Mean.Globally computes the arithmetic mean of the elements.""" def expand(self, pcoll): if self.has_defaults: return pcoll | core.CombineGlobally(MeanCombineFn()) else: return pcoll | core.CombineGlobally(MeanCombineFn()).without_defaults() class PerKey(ptransform.PTransform): """combiners.Mean.PerKey finds the means of the values for each key.""" def expand(self, pcoll): return pcoll | core.CombinePerKey(MeanCombineFn()) # TODO(laolu): This type signature is overly restrictive. This should be # more general. @with_input_types(Union[float, int, np.int64, np.float64]) @with_output_types(float) class MeanCombineFn(core.CombineFn): """CombineFn for computing an arithmetic mean.""" def create_accumulator(self): return (0, 0) def add_input(self, sum_count, element): (sum_, count) = sum_count return sum_ + element, count + 1 def merge_accumulators(self, accumulators): sums, counts = zip(*accumulators) return sum(sums), sum(counts) def extract_output(self, sum_count): (sum_, count) = sum_count if count == 0: return float('NaN') return sum_ / float(count) def for_input_type(self, input_type): if input_type is int: return cy_combiners.MeanInt64Fn() elif input_type is float: return cy_combiners.MeanFloatFn() return self class Count(object): """Combiners for counting elements.""" @with_input_types(T) @with_output_types(int) class Globally(CombinerWithoutDefaults): """combiners.Count.Globally counts the total number of elements.""" def expand(self, pcoll): if self.has_defaults: return pcoll | core.CombineGlobally(CountCombineFn()) else: return pcoll | core.CombineGlobally(CountCombineFn()).without_defaults() @with_input_types(tuple[K, V]) @with_output_types(tuple[K, int]) class PerKey(ptransform.PTransform): """combiners.Count.PerKey counts how many elements each unique key has.""" def expand(self, pcoll): return pcoll | core.CombinePerKey(CountCombineFn()) @with_input_types(T) @with_output_types(tuple[T, int]) class PerElement(ptransform.PTransform): """combiners.Count.PerElement counts how many times each element occurs.""" def expand(self, pcoll): paired_with_void_type = typehints.Tuple[pcoll.element_type, Any] output_type = typehints.KV[pcoll.element_type, int] return ( pcoll | ( '%s:PairWithVoid' % self.label >> core.Map( lambda x: (x, None)).with_output_types(paired_with_void_type)) | core.CombinePerKey(CountCombineFn()).with_output_types(output_type)) @with_input_types(Any) @with_output_types(int) class CountCombineFn(core.CombineFn): """CombineFn for computing PCollection size.""" def create_accumulator(self): return 0 def add_input(self, accumulator, element): return accumulator + 1 def add_inputs(self, accumulator, elements): return accumulator + len(list(elements)) def merge_accumulators(self, accumulators): return sum(accumulators) def extract_output(self, accumulator): return accumulator class Top(object): """Combiners for obtaining extremal elements.""" # pylint: disable=no-self-argument @with_input_types(T) @with_output_types(list[T]) class Of(CombinerWithoutDefaults): """Returns the n greatest elements in the PCollection. This transform will retrieve the n greatest elements in the PCollection to which it is applied, where "greatest" is determined by a function supplied as the `key` or `reverse` arguments. """ def __init__(self, n, key=None, reverse=False): """Creates a global Top operation. The arguments 'key' and 'reverse' may be passed as keyword arguments, and have the same meaning as for Python's sort functions. Args: n: number of elements to extract from pcoll. key: (optional) a mapping of elements to a comparable key, similar to the key argument of Python's sorting methods. reverse: (optional) whether to order things smallest to largest, rather than largest to smallest """ super().__init__() self._n = n self._key = key self._reverse = reverse def default_label(self): return 'Top(%d)' % self._n def expand(self, pcoll): if pcoll.windowing.is_default(): # This is a more efficient global algorithm. top_per_bundle = pcoll | core.ParDo( _TopPerBundle(self._n, self._key, self._reverse)) # If pcoll is empty, we can't guarantee that top_per_bundle # won't be empty, so inject at least one empty accumulator # so that downstream is guaranteed to produce non-empty output. empty_bundle = ( pcoll.pipeline | core.Create([(None, [])]).with_output_types( top_per_bundle.element_type)) return ((top_per_bundle, empty_bundle) | core.Flatten() | core.GroupByKey() | core.ParDo( _MergeTopPerBundle(self._n, self._key, self._reverse))) else: if self.has_defaults: return pcoll | core.CombineGlobally( TopCombineFn(self._n, self._key, self._reverse)) else: return pcoll | core.CombineGlobally( TopCombineFn(self._n, self._key, self._reverse)).without_defaults() @with_input_types(tuple[K, V]) @with_output_types(tuple[K, list[V]]) class PerKey(ptransform.PTransform): """Identifies the N greatest elements associated with each key. This transform will produce a PCollection mapping unique keys in the input PCollection to the n greatest elements with which they are associated, where "greatest" is determined by a function supplied as the `key` or `reverse` arguments. """ def __init__(self, n, key=None, reverse=False): """Creates a per-key Top operation. The arguments 'key' and 'reverse' may be passed as keyword arguments, and have the same meaning as for Python's sort functions. Args: n: number of elements to extract from pcoll. key: (optional) a mapping of elements to a comparable key, similar to the key argument of Python's sorting methods. reverse: (optional) whether to order things smallest to largest, rather than largest to smallest """ self._n = n self._key = key self._reverse = reverse def default_label(self): return 'TopPerKey(%d)' % self._n def expand(self, pcoll): """Expands the transform. Raises TypeCheckError: If the output type of the input PCollection is not compatible with tuple[A, B]. Args: pcoll: PCollection to process Returns: the PCollection containing the result. """ return pcoll | core.CombinePerKey( TopCombineFn(self._n, self._key, self._reverse)) @staticmethod @ptransform.ptransform_fn def Largest(pcoll, n, has_defaults=True, key=None): """Obtain a list of the greatest N elements in a PCollection.""" if has_defaults: return pcoll | Top.Of(n, key) else: return pcoll | Top.Of(n, key).without_defaults() @staticmethod @ptransform.ptransform_fn def Smallest(pcoll, n, has_defaults=True, key=None): """Obtain a list of the least N elements in a PCollection.""" if has_defaults: return pcoll | Top.Of(n, key, reverse=True) else: return pcoll | Top.Of(n, key, reverse=True).without_defaults() @staticmethod @ptransform.ptransform_fn def LargestPerKey(pcoll, n, key=None): """Identifies the N greatest elements associated with each key.""" return pcoll | Top.PerKey(n, key) @staticmethod @ptransform.ptransform_fn def SmallestPerKey(pcoll, n, *, key=None): """Identifies the N least elements associated with each key.""" return pcoll | Top.PerKey(n, key, reverse=True) @with_input_types(T) @with_output_types(tuple[None, list[T]]) class _TopPerBundle(core.DoFn): def __init__(self, n, key, reverse): self._n = n self._compare = operator.gt if reverse else None self._key = key def start_bundle(self): self._heap = [] def process(self, element): if self._compare or self._key: element = cy_combiners.ComparableValue(element, self._compare, self._key) if len(self._heap) < self._n: heapq.heappush(self._heap, element) else: heapq.heappushpop(self._heap, element) def finish_bundle(self): # Though sorting here results in more total work, this allows us to # skip most elements in the reducer. # Essentially, given s map bundles, we are trading about O(sn) compares in # the (single) reducer for O(sn log n) compares across all mappers. self._heap.sort() # Unwrap to avoid serialization via pickle. if self._compare or self._key: yield window.GlobalWindows.windowed_value( (None, [wrapper.value for wrapper in self._heap])) else: yield window.GlobalWindows.windowed_value((None, self._heap)) @with_input_types(tuple[None, Iterable[list[T]]]) @with_output_types(list[T]) class _MergeTopPerBundle(core.DoFn): def __init__(self, n, key, reverse): self._n = n self._compare = operator.gt if reverse else None self._key = key def process(self, key_and_bundles): _, bundles = key_and_bundles def push(hp, e): if len(hp) < self._n: heapq.heappush(hp, e) return False elif e < hp[0]: # Because _TopPerBundle returns sorted lists, all other elements # will also be smaller. return True else: heapq.heappushpop(hp, e) return False if self._compare or self._key: heapc: list[cy_combiners.ComparableValue] = [] for bundle in bundles: if not heapc: heapc = [ cy_combiners.ComparableValue(element, self._compare, self._key) for element in bundle ] continue # TODO(https://github.com/apache/beam/issues/21205): Remove this # workaround once legacy dataflow correctly handles coders with # combiner packing and/or is deprecated. if not isinstance(bundle, list): bundle = list(bundle) for element in reversed(bundle): if push(heapc, cy_combiners.ComparableValue(element, self._compare, self._key)): break heapc.sort() yield [wrapper.value for wrapper in reversed(heapc)] else: heap = [] for bundle in bundles: # TODO(https://github.com/apache/beam/issues/21205): Remove this # workaround once legacy dataflow correctly handles coders with # combiner packing and/or is deprecated. if not isinstance(bundle, list): bundle = list(bundle) if not heap: heap = bundle continue for element in reversed(bundle): if push(heap, element): break heap.sort() yield heap[::-1] @with_input_types(T) @with_output_types(list[T]) class TopCombineFn(core.CombineFn): """CombineFn doing the combining for all of the Top transforms. This CombineFn uses a `key` or `reverse` operator to rank the elements. Args: key: (optional) a mapping of elements to a comparable key, similar to the key argument of Python's sorting methods. reverse: (optional) whether to order things smallest to largest, rather than largest to smallest """ def __init__(self, n, key=None, reverse=False): self._n = n self._compare = operator.gt if reverse else operator.lt self._key = key def _hydrated_heap(self, heap): if heap: first = heap[0] if isinstance(first, cy_combiners.ComparableValue): if first.requires_hydration: for comparable in heap: assert comparable.requires_hydration comparable.hydrate(self._compare, self._key) assert not comparable.requires_hydration return heap else: return heap else: return [ cy_combiners.ComparableValue(element, self._compare, self._key) for element in heap ] else: return heap def display_data(self): return { 'n': self._n, 'compare': DisplayDataItem( self._compare.__name__ if hasattr(self._compare, '__name__') else self._compare.__class__.__name__).drop_if_none() } # The accumulator type is a tuple # (bool, Union[list[T], list[ComparableValue[T]]) # where the boolean indicates whether the second slot contains a List of T # (False) or List of ComparableValue[T] (True). In either case, the List # maintains heap invariance. When the contents of the List are # ComparableValue[T] they either all 'requires_hydration' or none do. # This accumulator representation allows us to minimize the data encoding # overheads. Creation of ComparableValues is elided for performance reasons # when there is no need for complicated comparison functions. def create_accumulator(self, *args, **kwargs): return (False, []) def add_input(self, accumulator, element, *args, **kwargs): # Caching to avoid paying the price of variadic expansion of args / kwargs # when it's not needed (for the 'if' case below). holds_comparables, heap = accumulator if self._compare is not operator.lt or self._key: heap = self._hydrated_heap(heap) holds_comparables = True else: assert not holds_comparables comparable = ( cy_combiners.ComparableValue(element, self._compare, self._key) if holds_comparables else element) if len(heap) < self._n: heapq.heappush(heap, comparable) else: heapq.heappushpop(heap, comparable) return (holds_comparables, heap) def merge_accumulators(self, accumulators, *args, **kwargs): result_heap = None holds_comparables = None for accumulator in accumulators: holds_comparables, heap = accumulator if self._compare is not operator.lt or self._key: heap = self._hydrated_heap(heap) holds_comparables = True else: assert not holds_comparables if result_heap is None: result_heap = heap else: for comparable in heap: _, result_heap = self.add_input( (holds_comparables, result_heap), comparable.value if holds_comparables else comparable) assert result_heap is not None and holds_comparables is not None return (holds_comparables, result_heap) def compact(self, accumulator, *args, **kwargs): holds_comparables, heap = accumulator # Unwrap to avoid serialization via pickle. if holds_comparables: return (False, [comparable.value for comparable in heap]) else: return accumulator def extract_output(self, accumulator, *args, **kwargs): holds_comparables, heap = accumulator if self._compare is not operator.lt or self._key: if not holds_comparables: heap = self._hydrated_heap(heap) holds_comparables = True else: assert not holds_comparables assert len(heap) <= self._n heap.sort(reverse=True) return [ comparable.value if holds_comparables else comparable for comparable in heap ] class Largest(TopCombineFn): def default_label(self): return 'Largest(%s)' % self._n class Smallest(TopCombineFn): def __init__(self, n): super().__init__(n, reverse=True) def default_label(self): return 'Smallest(%s)' % self._n class Sample(object): """Combiners for sampling n elements without replacement.""" # pylint: disable=no-self-argument @with_input_types(T) @with_output_types(list[T]) class FixedSizeGlobally(CombinerWithoutDefaults): """Sample n elements from the input PCollection without replacement.""" def __init__(self, n): super().__init__() self._n = n def expand(self, pcoll): if self.has_defaults: return pcoll | core.CombineGlobally(SampleCombineFn(self._n)) else: return pcoll | core.CombineGlobally(SampleCombineFn( self._n)).without_defaults() def display_data(self): return {'n': self._n} def default_label(self): return 'FixedSizeGlobally(%d)' % self._n @with_input_types(tuple[K, V]) @with_output_types(tuple[K, list[V]]) class FixedSizePerKey(ptransform.PTransform): """Sample n elements associated with each key without replacement.""" def __init__(self, n): self._n = n def expand(self, pcoll): return pcoll | core.CombinePerKey(SampleCombineFn(self._n)) def display_data(self): return {'n': self._n} def default_label(self): return 'FixedSizePerKey(%d)' % self._n @with_input_types(T) @with_output_types(list[T]) class SampleCombineFn(core.CombineFn): """CombineFn for all Sample transforms.""" def __init__(self, n): super().__init__() # Most of this combiner's work is done by a TopCombineFn. We could just # subclass TopCombineFn to make this class, but since sampling is not # really a kind of Top operation, we use a TopCombineFn instance as a # helper instead. self._top_combiner = TopCombineFn(n) def setup(self): self._top_combiner.setup() def create_accumulator(self): return self._top_combiner.create_accumulator() def add_input(self, heap, element): # Before passing elements to the Top combiner, we pair them with random # numbers. The elements with the n largest random number "keys" will be # selected for the output. return self._top_combiner.add_input(heap, (random.random(), element)) def merge_accumulators(self, heaps): return self._top_combiner.merge_accumulators(heaps) def compact(self, heap): return self._top_combiner.compact(heap) def extract_output(self, heap): # Here we strip off the random number keys we added in add_input. return [e for _, e in self._top_combiner.extract_output(heap)] def teardown(self): self._top_combiner.teardown() class _TupleCombineFnBase(core.CombineFn): def __init__(self, *combiners, merge_accumulators_batch_size=None): self._combiners = [core.CombineFn.maybe_from_callable(c) for c in combiners] self._named_combiners = combiners # If the `merge_accumulators_batch_size` value is not specified, we chose a # bounded default that is inversely proportional to the number of # accumulators in merged tuples. num_combiners = max(1, len(combiners)) self._merge_accumulators_batch_size = ( merge_accumulators_batch_size or max(10, 1000 // num_combiners)) def display_data(self): combiners = [ c.__name__ if hasattr(c, '__name__') else c.__class__.__name__ for c in self._named_combiners ] return { 'combiners': str(combiners), 'merge_accumulators_batch_size': self._merge_accumulators_batch_size } def setup(self, *args, **kwargs): for c in self._combiners: c.setup(*args, **kwargs) def create_accumulator(self, *args, **kwargs): return [c.create_accumulator(*args, **kwargs) for c in self._combiners] def merge_accumulators(self, accumulators, *args, **kwargs): # Make sure that `accumulators` is an iterator (so that the position is # remembered). accumulators = iter(accumulators) result = next(accumulators) while True: # Load accumulators into memory and merge in batches to decrease peak # memory usage. accumulators_batch = [result] + list( itertools.islice(accumulators, self._merge_accumulators_batch_size)) if len(accumulators_batch) == 1: break result = [ c.merge_accumulators(a, *args, **kwargs) for c, a in zip(self._combiners, zip(*accumulators_batch)) ] return result def compact(self, accumulator, *args, **kwargs): return [ c.compact(a, *args, **kwargs) for c, a in zip(self._combiners, accumulator) ] def extract_output(self, accumulator, *args, **kwargs): return tuple( c.extract_output(a, *args, **kwargs) for c, a in zip(self._combiners, accumulator)) def teardown(self, *args, **kwargs): for c in reversed(self._combiners): c.teardown(*args, **kwargs) class TupleCombineFn(_TupleCombineFnBase): """A combiner for combining tuples via a tuple of combiners. Takes as input a tuple of N CombineFns and combines N-tuples by combining the k-th element of each tuple with the k-th CombineFn, outputting a new N-tuple of combined values. """ def add_input(self, accumulator, element, *args, **kwargs): return [ c.add_input(a, e, *args, **kwargs) for c, a, e in zip(self._combiners, accumulator, element) ] def with_common_input(self): return SingleInputTupleCombineFn(*self._combiners) class SingleInputTupleCombineFn(_TupleCombineFnBase): """A combiner for combining a single value via a tuple of combiners. Takes as input a tuple of N CombineFns and combines elements by applying each CombineFn to each input, producing an N-tuple of the outputs corresponding to each of the N CombineFn's outputs. """ def add_input(self, accumulator, element, *args, **kwargs): return [ c.add_input(a, element, *args, **kwargs) for c, a in zip(self._combiners, accumulator) ] @with_input_types(T) @with_output_types(list[T]) class ToList(CombinerWithoutDefaults): """A global CombineFn that condenses a PCollection into a single list.""" def expand(self, pcoll): if self.has_defaults: return pcoll | self.label >> core.CombineGlobally(ToListCombineFn()) else: return pcoll | self.label >> core.CombineGlobally( ToListCombineFn()).without_defaults() @with_input_types(T) @with_output_types(list[T]) class ToListCombineFn(core.CombineFn): """CombineFn for to_list.""" def create_accumulator(self): return [] def add_input(self, accumulator, element): accumulator.append(element) return accumulator def merge_accumulators(self, accumulators): return sum(accumulators, []) def extract_output(self, accumulator): return accumulator @with_input_types(T) @with_output_types(T) class ConcatListCombineFn(core.CombineFn): """CombineFn for concatenating lists together.""" def create_accumulator(self): return [] def add_input(self, accumulator, element): return accumulator + element def merge_accumulators(self, accumulators): return sum(accumulators, []) def extract_output(self, accumulator): return accumulator @with_input_types(tuple[K, V]) @with_output_types(dict[K, V]) class ToDict(CombinerWithoutDefaults): """A global CombineFn that condenses a PCollection into a single dict. PCollections should consist of 2-tuples, notionally (key, value) pairs. If multiple values are associated with the same key, only one of the values will be present in the resulting dict. """ def expand(self, pcoll): if self.has_defaults: return pcoll | self.label >> core.CombineGlobally(ToDictCombineFn()) else: return pcoll | self.label >> core.CombineGlobally( ToDictCombineFn()).without_defaults() @with_input_types(tuple[K, V]) @with_output_types(dict[K, V]) class ToDictCombineFn(core.CombineFn): """CombineFn for to_dict.""" def create_accumulator(self): return {} def add_input(self, accumulator, element): key, value = element accumulator[key] = value return accumulator def merge_accumulators(self, accumulators): result = {} for a in accumulators: result.update(a) return result def extract_output(self, accumulator): return accumulator @with_input_types(T) @with_output_types(set[T]) class ToSet(CombinerWithoutDefaults): """A global CombineFn that condenses a PCollection into a set.""" def expand(self, pcoll): if self.has_defaults: return pcoll | self.label >> core.CombineGlobally(ToSetCombineFn()) else: return pcoll | self.label >> core.CombineGlobally( ToSetCombineFn()).without_defaults() @with_input_types(T) @with_output_types(set[T]) class ToSetCombineFn(core.CombineFn): """CombineFn for ToSet.""" def create_accumulator(self): return set() def add_input(self, accumulator, element): accumulator.add(element) return accumulator def merge_accumulators(self, accumulators): return set.union(*accumulators) def extract_output(self, accumulator): return accumulator class _CurriedFn(core.CombineFn): """Wrapped CombineFn with extra arguments.""" def __init__(self, fn, args, kwargs): self.fn = fn self.args = args self.kwargs = kwargs def setup(self): self.fn.setup(*self.args, **self.kwargs) def create_accumulator(self): return self.fn.create_accumulator(*self.args, **self.kwargs) def add_input(self, accumulator, element): return self.fn.add_input(accumulator, element, *self.args, **self.kwargs) def merge_accumulators(self, accumulators): return self.fn.merge_accumulators(accumulators, *self.args, **self.kwargs) def compact(self, accumulator): return self.fn.compact(accumulator, *self.args, **self.kwargs) def extract_output(self, accumulator): return self.fn.extract_output(accumulator, *self.args, **self.kwargs) def teardown(self): self.fn.teardown(*self.args, **self.kwargs) def apply(self, elements): return self.fn.apply(elements, *self.args, **self.kwargs) def curry_combine_fn(fn, args, kwargs): if not args and not kwargs: return fn else: return _CurriedFn(fn, args, kwargs) class PhasedCombineFnExecutor(object): """Executor for phases of combine operations.""" def __init__(self, phase, fn, args, kwargs): self.combine_fn = curry_combine_fn(fn, args, kwargs) if phase == 'all': self.apply = self.full_combine elif phase == 'add': self.apply = self.add_only elif phase == 'merge': self.apply = self.merge_only elif phase == 'extract': self.apply = self.extract_only elif phase == 'convert': self.apply = self.convert_to_accumulator else: raise ValueError('Unexpected phase: %s' % phase) def full_combine(self, elements): return self.combine_fn.apply(elements) def add_only(self, elements): return self.combine_fn.add_inputs( self.combine_fn.create_accumulator(), elements) def merge_only(self, accumulators): return self.combine_fn.merge_accumulators(accumulators) def extract_only(self, accumulator): return self.combine_fn.extract_output(accumulator) def convert_to_accumulator(self, element): return self.combine_fn.add_input( self.combine_fn.create_accumulator(), element) class Latest(object): """Combiners for computing the latest element""" @with_input_types(T) @with_output_types(T) class Globally(CombinerWithoutDefaults): """Compute the element with the latest timestamp from a PCollection.""" @staticmethod def add_timestamp(element, timestamp=core.DoFn.TimestampParam): return [(element, timestamp)] def expand(self, pcoll): if self.has_defaults: return ( pcoll | core.ParDo(self.add_timestamp).with_output_types( tuple[T, TimestampType]) | core.CombineGlobally(LatestCombineFn())) else: return ( pcoll | core.ParDo(self.add_timestamp).with_output_types( tuple[T, TimestampType]) | core.CombineGlobally(LatestCombineFn()).without_defaults()) @with_input_types(tuple[K, V]) @with_output_types(tuple[K, V]) class PerKey(ptransform.PTransform): """Compute elements with the latest timestamp for each key from a keyed PCollection""" @staticmethod def add_timestamp(element, timestamp=core.DoFn.TimestampParam): key, value = element return [(key, (value, timestamp))] def expand(self, pcoll): return ( pcoll | core.ParDo(self.add_timestamp).with_output_types( tuple[K, tuple[T, TimestampType]]) | core.CombinePerKey(LatestCombineFn())) @with_input_types(tuple[T, TimestampType]) @with_output_types(T) class LatestCombineFn(core.CombineFn): """CombineFn to get the element with the latest timestamp from a PCollection.""" def create_accumulator(self): return (None, window.MIN_TIMESTAMP) def add_input(self, accumulator, element): if accumulator[1] > element[1]: return accumulator else: return element def merge_accumulators(self, accumulators): result = self.create_accumulator() for accumulator in accumulators: result = self.add_input(result, accumulator) return result def extract_output(self, accumulator): return accumulator[0] class LiftedCombinePerKey(core.PTransform): """An implementation of CombinePerKey that does mapper-side pre-combining. This shouldn't generally be used directly except for use-cases where a runner doesn't support CombinePerKey. This implementation manually implements a CombinePerKey using ParDos, as opposed to runner implementations which may use a more efficient implementation. """ def __init__(self, combine_fn, args, kwargs): side_inputs = _pack_side_inputs(args, kwargs) self._side_inputs: dict = side_inputs if not isinstance(combine_fn, core.CombineFn): combine_fn = core.CombineFn.from_callable(combine_fn) self._combine_fn = combine_fn def expand(self, pcoll): return ( pcoll | core.ParDo( _PartialGroupByKeyCombiningValues(self._combine_fn), **self._side_inputs) | core.GroupByKey() | core.ParDo(_FinishCombine(self._combine_fn), **self._side_inputs)) def _pack_side_inputs(side_input_args, side_input_kwargs): if len(side_input_args) >= 10: # If we have more than 10 side inputs, we can't use the # _side_input_arg_{i} as our keys since they won't sort # correctly. Just punt for now, more than 10 args probably # doesn't happen often. raise NotImplementedError side_inputs = {} for i, si in enumerate(side_input_args): side_inputs[f'_side_input_arg_{i}'] = si for k, v in side_input_kwargs.items(): side_inputs[k] = v return side_inputs def _unpack_side_inputs(side_inputs): side_input_args = [] side_input_kwargs = {} for k, v in sorted(side_inputs.items(), key=lambda x: x[0]): if k.startswith('_side_input_arg_'): side_input_args.append(v) else: side_input_kwargs[k] = v return side_input_args, side_input_kwargs class _PartialGroupByKeyCombiningValues(core.DoFn): """Aggregates values into a per-key-window cache. As bundles are in-memory-sized, we don't bother flushing until the very end. """ def __init__(self, combine_fn): self._combine_fn = combine_fn self.side_input_args = [] self.side_input_kwargs = {} def setup(self): self._combine_fn.setup() def start_bundle(self): self._cache = dict() self._cached_windowed_side_inputs = {} def process(self, element, window=core.DoFn.WindowParam, **side_inputs): k, vi = element side_input_args, side_input_kwargs = _unpack_side_inputs(side_inputs) if (k, window) not in self._cache: self._cache[(k, window)] = self._combine_fn.create_accumulator( *side_input_args, **side_input_kwargs) self._cache[k, window] = self._combine_fn.add_input( self._cache[k, window], vi, *side_input_args, **side_input_kwargs) self._cached_windowed_side_inputs[window] = ( side_input_args, side_input_kwargs) return [] # to prevent DoFn-no-iterator warning def finish_bundle(self): for (k, w), va in self._cache.items(): # We compact the accumulator since a GBK (which necessitates encoding) # will follow. side_input_args, side_input_kwargs = ( self._cached_windowed_side_inputs[w]) yield WindowedValue(( k, self._combine_fn.compact(va, *side_input_args, **side_input_kwargs)), w.end, (w, )) def teardown(self): self._combine_fn.teardown() class _FinishCombine(core.DoFn): """Merges partially combined results. """ def __init__(self, combine_fn): self._combine_fn = combine_fn def setup(self): self._combine_fn.setup() def process(self, element, window=core.DoFn.WindowParam, **side_inputs): k, vs = element side_input_args, side_input_kwargs = _unpack_side_inputs(side_inputs) return [( k, self._combine_fn.extract_output( self._combine_fn.merge_accumulators( vs, *side_input_args, **side_input_kwargs), *side_input_args, **side_input_kwargs))] def teardown(self): try: self._combine_fn.teardown() except AttributeError: pass ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/combiners_test.py0000644000175100017510000013544315206113033023460 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for our libraries of combine PTransforms.""" # pytype: skip-file import itertools import json import os import random import tempfile import time import unittest from pathlib import Path import hamcrest as hc import pytest import apache_beam as beam import apache_beam.transforms.combiners as combine from apache_beam import pvalue from apache_beam.metrics import Metrics from apache_beam.metrics import MetricsFilter from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.testing.util import equal_to_per_window from apache_beam.transforms import WindowInto from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.transforms.core import CombineGlobally from apache_beam.transforms.core import Create from apache_beam.transforms.core import Map from apache_beam.transforms.display import DisplayData from apache_beam.transforms.display_test import DisplayDataItemMatcher from apache_beam.transforms.periodicsequence import PeriodicImpulse from apache_beam.transforms.ptransform import PTransform from apache_beam.transforms.trigger import AfterAll from apache_beam.transforms.trigger import AfterCount from apache_beam.transforms.trigger import AfterWatermark from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import GlobalWindows from apache_beam.transforms.window import TimestampCombiner from apache_beam.transforms.window import TimestampedValue from apache_beam.typehints import TypeCheckError from apache_beam.utils.timestamp import Timestamp class SortedConcatWithCounters(beam.CombineFn): """CombineFn for incrementing three different counters: counter, distribution, gauge, at the same time concatenating words.""" def __init__(self): beam.CombineFn.__init__(self) self.word_counter = Metrics.counter(self.__class__, 'word_counter') self.word_lengths_counter = Metrics.counter(self.__class__, 'word_lengths') self.word_lengths_dist = Metrics.distribution( self.__class__, 'word_len_dist') self.last_word_len = Metrics.gauge(self.__class__, 'last_word_len') def create_accumulator(self): return '' def add_input(self, acc, element): self.word_counter.inc(1) self.word_lengths_counter.inc(len(element)) self.word_lengths_dist.update(len(element)) self.last_word_len.set(len(element)) return acc + element def merge_accumulators(self, accs): return ''.join(accs) def extract_output(self, acc): # The sorted acc became a list of characters # and has to be converted back to a string using join. return ''.join(sorted(acc)) class CombineTest(unittest.TestCase): def test_builtin_combines(self): with TestPipeline() as pipeline: vals = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6] mean = sum(vals) / float(len(vals)) size = len(vals) timestamp = 0 # First for global combines. pcoll = pipeline | 'start' >> Create(vals) result_mean = pcoll | 'mean' >> combine.Mean.Globally() result_count = pcoll | 'count' >> combine.Count.Globally() assert_that(result_mean, equal_to([mean]), label='assert:mean') assert_that(result_count, equal_to([size]), label='assert:size') # Now for global combines without default timestamped = pcoll | Map(lambda x: TimestampedValue(x, timestamp)) windowed = timestamped | 'window' >> WindowInto(FixedWindows(60)) result_windowed_mean = ( windowed | 'mean-wo-defaults' >> combine.Mean.Globally().without_defaults()) assert_that( result_windowed_mean, equal_to([mean]), label='assert:mean-wo-defaults') result_windowed_count = ( windowed | 'count-wo-defaults' >> combine.Count.Globally().without_defaults()) assert_that( result_windowed_count, equal_to([size]), label='assert:count-wo-defaults') # Again for per-key combines. pcoll = pipeline | 'start-perkey' >> Create([('a', x) for x in vals]) result_key_mean = pcoll | 'mean-perkey' >> combine.Mean.PerKey() result_key_count = pcoll | 'count-perkey' >> combine.Count.PerKey() assert_that(result_key_mean, equal_to([('a', mean)]), label='key:mean') assert_that(result_key_count, equal_to([('a', size)]), label='key:size') def test_top(self): with TestPipeline() as pipeline: timestamp = 0 # First for global combines. pcoll = pipeline | 'start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6]) result_top = pcoll | 'top' >> combine.Top.Largest(5) result_bot = pcoll | 'bot' >> combine.Top.Smallest(4) assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='assert:top') assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='assert:bot') # Now for global combines without default timestamped = pcoll | Map(lambda x: TimestampedValue(x, timestamp)) windowed = timestamped | 'window' >> WindowInto(FixedWindows(60)) result_windowed_top = windowed | 'top-wo-defaults' >> combine.Top.Largest( 5, has_defaults=False) result_windowed_bot = ( windowed | 'bot-wo-defaults' >> combine.Top.Smallest(4, has_defaults=False)) assert_that( result_windowed_top, equal_to([[9, 6, 6, 5, 3]]), label='assert:top-wo-defaults') assert_that( result_windowed_bot, equal_to([[0, 1, 1, 1]]), label='assert:bot-wo-defaults') # Again for per-key combines. pcoll = pipeline | 'start-perkey' >> Create( [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]]) result_key_top = pcoll | 'top-perkey' >> combine.Top.LargestPerKey(5) result_key_bot = pcoll | 'bot-perkey' >> combine.Top.SmallestPerKey(4) assert_that( result_key_top, equal_to([('a', [9, 6, 6, 5, 3])]), label='key:top') assert_that( result_key_bot, equal_to([('a', [0, 1, 1, 1])]), label='key:bot') def test_empty_global_top(self): with TestPipeline() as p: assert_that(p | beam.Create([]) | combine.Top.Largest(10), equal_to([[]])) def test_sharded_top(self): elements = list(range(100)) random.shuffle(elements) with TestPipeline() as pipeline: shards = [ pipeline | 'Shard%s' % shard >> beam.Create(elements[shard::7]) for shard in range(7) ] assert_that( shards | beam.Flatten() | combine.Top.Largest(10), equal_to([[99, 98, 97, 96, 95, 94, 93, 92, 91, 90]])) def test_top_key(self): self.assertEqual(['aa', 'bbb', 'c', 'dddd'] | combine.Top.Of(3, key=len), [['dddd', 'bbb', 'aa']]) self.assertEqual(['aa', 'bbb', 'c', 'dddd'] | combine.Top.Of(3, key=len, reverse=True), [['c', 'aa', 'bbb']]) self.assertEqual(['xc', 'zb', 'yd', 'wa'] | combine.Top.Largest(3, key=lambda x: x[-1]), [['yd', 'xc', 'zb']]) self.assertEqual(['xc', 'zb', 'yd', 'wa'] | combine.Top.Smallest(3, key=lambda x: x[-1]), [['wa', 'zb', 'xc']]) self.assertEqual([('a', x) for x in [1, 2, 3, 4, 1, 1]] | combine.Top.LargestPerKey(3, key=lambda x: -x), [('a', [1, 1, 1])]) self.assertEqual([('a', x) for x in [1, 2, 3, 4, 1, 1]] | combine.Top.SmallestPerKey(3, key=lambda x: -x), [('a', [4, 3, 2])]) def test_sharded_top_combine_fn(self): def test_combine_fn(combine_fn, shards, expected): accumulators = [ combine_fn.add_inputs(combine_fn.create_accumulator(), shard) for shard in shards ] final_accumulator = combine_fn.merge_accumulators(accumulators) self.assertEqual(combine_fn.extract_output(final_accumulator), expected) test_combine_fn(combine.TopCombineFn(3), [range(10), range(10)], [9, 9, 8]) test_combine_fn( combine.TopCombineFn(5), [range(1000), range(100), range(1001)], [1000, 999, 999, 998, 998]) def test_combine_per_key_top_display_data(self): def individual_test_per_key_dd(combineFn): transform = beam.CombinePerKey(combineFn) dd = DisplayData.create_from(transform) expected_items = [ DisplayDataItemMatcher('combine_fn', combineFn.__class__), DisplayDataItemMatcher('n', combineFn._n), DisplayDataItemMatcher('compare', combineFn._compare.__name__) ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) individual_test_per_key_dd(combine.Largest(5)) individual_test_per_key_dd(combine.Smallest(3)) individual_test_per_key_dd(combine.TopCombineFn(8)) individual_test_per_key_dd(combine.Largest(5)) def test_combine_sample_display_data(self): def individual_test_per_key_dd(sampleFn, n): trs = [sampleFn(n)] for transform in trs: dd = DisplayData.create_from(transform) hc.assert_that( dd.items, hc.contains_inanyorder(DisplayDataItemMatcher('n', transform._n))) individual_test_per_key_dd(combine.Sample.FixedSizePerKey, 5) individual_test_per_key_dd(combine.Sample.FixedSizeGlobally, 5) def test_combine_globally_display_data(self): transform = beam.CombineGlobally(combine.Smallest(5)) dd = DisplayData.create_from(transform) expected_items = [ DisplayDataItemMatcher('combine_fn', combine.Smallest), DisplayDataItemMatcher('n', 5), DisplayDataItemMatcher('compare', 'gt') ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_basic_combiners_display_data(self): transform = beam.CombineGlobally( combine.TupleCombineFn(max, combine.MeanCombineFn(), sum)) dd = DisplayData.create_from(transform) expected_items = [ DisplayDataItemMatcher('combine_fn', combine.TupleCombineFn), DisplayDataItemMatcher('combiners', "['max', 'MeanCombineFn', 'sum']"), DisplayDataItemMatcher('merge_accumulators_batch_size', 333), ] hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) def test_top_shorthands(self): with TestPipeline() as pipeline: pcoll = pipeline | 'start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6]) result_top = pcoll | 'top' >> beam.CombineGlobally(combine.Largest(5)) result_bot = pcoll | 'bot' >> beam.CombineGlobally(combine.Smallest(4)) assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='assert:top') assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='assert:bot') pcoll = pipeline | 'start-perkey' >> Create( [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]]) result_ktop = pcoll | 'top-perkey' >> beam.CombinePerKey( combine.Largest(5)) result_kbot = pcoll | 'bot-perkey' >> beam.CombinePerKey( combine.Smallest(4)) assert_that(result_ktop, equal_to([('a', [9, 6, 6, 5, 3])]), label='ktop') assert_that(result_kbot, equal_to([('a', [0, 1, 1, 1])]), label='kbot') def test_top_no_compact(self): class TopCombineFnNoCompact(combine.TopCombineFn): def compact(self, accumulator): return accumulator with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6]) result_top = pcoll | 'Top' >> beam.CombineGlobally( TopCombineFnNoCompact(5, key=lambda x: x)) result_bot = pcoll | 'Bot' >> beam.CombineGlobally( TopCombineFnNoCompact(4, reverse=True)) assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='Assert:Top') assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='Assert:Bot') pcoll = pipeline | 'Start-Perkey' >> Create( [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]]) result_ktop = pcoll | 'Top-PerKey' >> beam.CombinePerKey( TopCombineFnNoCompact(5, key=lambda x: x)) result_kbot = pcoll | 'Bot-PerKey' >> beam.CombinePerKey( TopCombineFnNoCompact(4, reverse=True)) assert_that(result_ktop, equal_to([('a', [9, 6, 6, 5, 3])]), label='KTop') assert_that(result_kbot, equal_to([('a', [0, 1, 1, 1])]), label='KBot') def test_global_sample(self): def is_good_sample(actual): assert len(actual) == 1 assert sorted(actual[0]) in [[1, 1, 2], [1, 2, 2]], actual with TestPipeline() as pipeline: timestamp = 0 pcoll = pipeline | 'start' >> Create([1, 1, 2, 2]) # Now for global combines without default timestamped = pcoll | Map(lambda x: TimestampedValue(x, timestamp)) windowed = timestamped | 'window' >> WindowInto(FixedWindows(60)) for ix in range(9): assert_that( pcoll | 'sample-%d' % ix >> combine.Sample.FixedSizeGlobally(3), is_good_sample, label='check-%d' % ix) result_windowed = ( windowed | 'sample-wo-defaults-%d' % ix >> combine.Sample.FixedSizeGlobally(3).without_defaults()) assert_that( result_windowed, is_good_sample, label='check-wo-defaults-%d' % ix) def test_per_key_sample(self): with TestPipeline() as pipeline: pcoll = pipeline | 'start-perkey' >> Create( sum(([(i, 1), (i, 1), (i, 2), (i, 2)] for i in range(9)), [])) result = pcoll | 'sample' >> combine.Sample.FixedSizePerKey(3) def matcher(): def match(actual): for _, samples in actual: equal_to([3])([len(samples)]) num_ones = sum(1 for x in samples if x == 1) num_twos = sum(1 for x in samples if x == 2) equal_to([1, 2])([num_ones, num_twos]) return match assert_that(result, matcher()) def test_tuple_combine_fn(self): with TestPipeline() as p: result = ( p | Create([('a', 100, 0.0), ('b', 10, -1), ('c', 1, 100)]) | beam.CombineGlobally( combine.TupleCombineFn(max, combine.MeanCombineFn(), sum)).without_defaults()) assert_that(result, equal_to([('c', 111.0 / 3, 99.0)])) def test_tuple_combine_fn_without_defaults(self): with TestPipeline() as p: result = ( p | Create([1, 1, 2, 3]) | beam.CombineGlobally( combine.TupleCombineFn( min, combine.MeanCombineFn(), max).with_common_input()).without_defaults()) assert_that(result, equal_to([(1, 7.0 / 4, 3)])) def test_empty_tuple_combine_fn(self): with TestPipeline() as p: result = ( p | Create([(), (), ()]) | beam.CombineGlobally(combine.TupleCombineFn())) assert_that(result, equal_to([()])) def test_tuple_combine_fn_batched_merge(self): num_combine_fns = 10 max_num_accumulators_in_memory = 30 # Maximum number of accumulator tuples in memory - 1 for the merge result. merge_accumulators_batch_size = ( max_num_accumulators_in_memory // num_combine_fns - 1) num_accumulator_tuples_to_merge = 20 class CountedAccumulator: count = 0 oom = False def __init__(self): if CountedAccumulator.count > max_num_accumulators_in_memory: CountedAccumulator.oom = True else: CountedAccumulator.count += 1 class CountedAccumulatorCombineFn(beam.CombineFn): def create_accumulator(self): return CountedAccumulator() def merge_accumulators(self, accumulators): CountedAccumulator.count += 1 for _ in accumulators: CountedAccumulator.count -= 1 combine_fn = combine.TupleCombineFn( *[CountedAccumulatorCombineFn() for _ in range(num_combine_fns)], merge_accumulators_batch_size=merge_accumulators_batch_size) combine_fn.merge_accumulators( combine_fn.create_accumulator() for _ in range(num_accumulator_tuples_to_merge)) assert not CountedAccumulator.oom def test_to_list_and_to_dict1(self): with TestPipeline() as pipeline: the_list = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6] timestamp = 0 pcoll = pipeline | 'start' >> Create(the_list) result = pcoll | 'to list' >> combine.ToList() # Now for global combines without default timestamped = pcoll | Map(lambda x: TimestampedValue(x, timestamp)) windowed = timestamped | 'window' >> WindowInto(FixedWindows(60)) result_windowed = ( windowed | 'to list wo defaults' >> combine.ToList().without_defaults()) def matcher(expected): def match(actual): equal_to(expected[0])(actual[0]) return match assert_that(result, matcher([the_list])) assert_that( result_windowed, matcher([the_list]), label='to-list-wo-defaults') def test_to_list_and_to_dict2(self): with TestPipeline() as pipeline: pairs = [(1, 2), (3, 4), (5, 6)] timestamp = 0 pcoll = pipeline | 'start-pairs' >> Create(pairs) result = pcoll | 'to dict' >> combine.ToDict() # Now for global combines without default timestamped = pcoll | Map(lambda x: TimestampedValue(x, timestamp)) windowed = timestamped | 'window' >> WindowInto(FixedWindows(60)) result_windowed = ( windowed | 'to dict wo defaults' >> combine.ToDict().without_defaults()) def matcher(): def match(actual): equal_to([1])([len(actual)]) equal_to(pairs)(actual[0].items()) return match assert_that(result, matcher()) assert_that(result_windowed, matcher(), label='to-dict-wo-defaults') def test_to_set(self): pipeline = TestPipeline() the_list = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6] timestamp = 0 pcoll = pipeline | 'start' >> Create(the_list) result = pcoll | 'to set' >> combine.ToSet() # Now for global combines without default timestamped = pcoll | Map(lambda x: TimestampedValue(x, timestamp)) windowed = timestamped | 'window' >> WindowInto(FixedWindows(60)) result_windowed = ( windowed | 'to set wo defaults' >> combine.ToSet().without_defaults()) def matcher(expected): def match(actual): equal_to(expected[0])(actual[0]) return match assert_that(result, matcher(set(the_list))) assert_that( result_windowed, matcher(set(the_list)), label='to-set-wo-defaults') def test_combine_globally_with_default(self): with TestPipeline() as p: assert_that(p | Create([]) | CombineGlobally(sum), equal_to([0])) def test_combine_globally_without_default(self): with TestPipeline() as p: result = p | Create([]) | CombineGlobally(sum).without_defaults() assert_that(result, equal_to([])) def test_combine_globally_with_default_side_input(self): class SideInputCombine(PTransform): def expand(self, pcoll): side = pcoll | CombineGlobally(sum).as_singleton_view() main = pcoll.pipeline | Create([None]) return main | Map(lambda _, s: s, side) with TestPipeline() as p: result1 = p | 'i1' >> Create([]) | 'c1' >> SideInputCombine() result2 = p | 'i2' >> Create([1, 2, 3, 4]) | 'c2' >> SideInputCombine() assert_that(result1, equal_to([0]), label='r1') assert_that(result2, equal_to([10]), label='r2') def test_hot_key_fanout(self): with TestPipeline() as p: result = ( p | beam.Create(itertools.product(['hot', 'cold'], range(10))) | beam.CombinePerKey(combine.MeanCombineFn()).with_hot_key_fanout( lambda key: (key == 'hot') * 5)) assert_that(result, equal_to([('hot', 4.5), ('cold', 4.5)])) def test_hot_key_fanout_sharded(self): # Lots of elements with the same key with varying/no fanout. with TestPipeline() as p: elements = [(None, e) for e in range(1000)] random.shuffle(elements) shards = [ p | "Shard%s" % shard >> beam.Create(elements[shard::20]) for shard in range(20) ] result = ( shards | beam.Flatten() | beam.CombinePerKey(combine.MeanCombineFn()).with_hot_key_fanout( lambda key: random.randrange(0, 5))) assert_that(result, equal_to([(None, 499.5)])) def test_global_fanout(self): with TestPipeline() as p: result = ( p | beam.Create(range(100)) | beam.CombineGlobally(combine.MeanCombineFn()).with_fanout(11)) assert_that(result, equal_to([49.5])) def test_combining_with_accumulation_mode_and_fanout(self): # PCollection will contain elements from 1 to 5. elements = [i for i in range(1, 6)] ts = TestStream().advance_watermark_to(0) for i in elements: ts.add_elements([i]) ts.advance_watermark_to_infinity() options = PipelineOptions() options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: result = ( p | ts | beam.WindowInto( GlobalWindows(), accumulation_mode=trigger.AccumulationMode.ACCUMULATING, trigger=AfterWatermark(early=AfterAll(AfterCount(1)))) | beam.CombineGlobally(sum).without_defaults().with_fanout(2)) def has_expected_values(actual): from hamcrest.core import assert_that as hamcrest_assert from hamcrest.library.collection import contains from hamcrest.library.collection import only_contains ordered = sorted(actual) # Early firings. hamcrest_assert(ordered[:4], contains(1, 3, 6, 10)) # Different runners have different number of 15s, but there should # be at least one 15. hamcrest_assert(ordered[4:], only_contains(15)) assert_that(result, has_expected_values) def test_combining_with_sliding_windows_and_fanout(self): options = PipelineOptions() options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: def has_expected_values(actual): from hamcrest.core import assert_that as hamcrest_assert from hamcrest.library.collection import only_contains ordered = sorted(actual) hamcrest_assert( ordered, only_contains([0, 1, 2, 3], [0, 1, 2, 3, 5, 6, 7, 8], [5, 6, 7, 8])) result = ( p | beam.Create([ window.TimestampedValue(0, Timestamp(seconds=1666707510)), window.TimestampedValue(1, Timestamp(seconds=1666707511)), window.TimestampedValue(2, Timestamp(seconds=1666707512)), window.TimestampedValue(3, Timestamp(seconds=1666707513)), window.TimestampedValue(5, Timestamp(seconds=1666707515)), window.TimestampedValue(6, Timestamp(seconds=1666707516)), window.TimestampedValue(7, Timestamp(seconds=1666707517)), window.TimestampedValue(8, Timestamp(seconds=1666707518)) ]) | beam.WindowInto(window.SlidingWindows(10, 5)) | beam.CombineGlobally(beam.combiners.ToListCombineFn()). without_defaults().with_fanout(7)) assert_that(result, has_expected_values) def test_combining_with_session_windows_and_fanout(self): options = PipelineOptions() options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: def has_expected_values(actual): from hamcrest.core import assert_that as hamcrest_assert from hamcrest.library.collection import only_contains ordered = sorted(actual) hamcrest_assert(ordered, only_contains([0, 1, 2, 3], [5, 6, 7, 8])) result = ( p | beam.Create([ window.TimestampedValue(0, Timestamp(seconds=1666707510)), window.TimestampedValue(1, Timestamp(seconds=1666707511)), window.TimestampedValue(2, Timestamp(seconds=1666707512)), window.TimestampedValue(3, Timestamp(seconds=1666707513)), window.TimestampedValue(5, Timestamp(seconds=1666707515)), window.TimestampedValue(6, Timestamp(seconds=1666707516)), window.TimestampedValue(7, Timestamp(seconds=1666707517)), window.TimestampedValue(8, Timestamp(seconds=1666707518)) ]) | beam.WindowInto(window.Sessions(2)) | beam.CombineGlobally(beam.combiners.ToListCombineFn()). without_defaults().with_fanout(7)) assert_that(result, has_expected_values) def test_MeanCombineFn_combine(self): with TestPipeline() as p: input = ( p | beam.Create([('a', 1), ('a', 1), ('a', 4), ('b', 1), ('b', 13)])) # The mean of all values regardless of key. global_mean = ( input | beam.Values() | beam.CombineGlobally(combine.MeanCombineFn())) # The (key, mean) pairs for all keys. mean_per_key = (input | beam.CombinePerKey(combine.MeanCombineFn())) expected_mean_per_key = [('a', 2), ('b', 7)] assert_that(global_mean, equal_to([4]), label='global mean') assert_that( mean_per_key, equal_to(expected_mean_per_key), label='mean per key') def test_MeanCombineFn_combine_empty(self): # For each element in a PCollection, if it is float('NaN'), then emits # a string 'NaN', otherwise emits str(element). with TestPipeline() as p: input = (p | beam.Create([])) # Compute the mean of all values in the PCollection, # then format the mean. Since the Pcollection is empty, # the mean is float('NaN'), and is formatted to be a string 'NaN'. global_mean = ( input | beam.Values() | beam.CombineGlobally(combine.MeanCombineFn()) | beam.Map(str)) mean_per_key = (input | beam.CombinePerKey(combine.MeanCombineFn())) # We can't compare one float('NaN') with another float('NaN'), # but we can compare one 'nan' string with another string. assert_that(global_mean, equal_to(['nan']), label='global mean') assert_that(mean_per_key, equal_to([]), label='mean per key') def test_sessions_combine(self): with TestPipeline() as p: input = ( p | beam.Create([('c', 1), ('c', 9), ('c', 12), ('d', 2), ('d', 4)]) | beam.MapTuple(lambda k, v: window.TimestampedValue((k, v), v)) | beam.WindowInto(window.Sessions(4))) global_sum = ( input | beam.Values() | beam.CombineGlobally(sum).without_defaults()) sum_per_key = input | beam.CombinePerKey(sum) # The first window has 3 elements: ('c', 1), ('d', 2), ('d', 4). # The second window has 2 elements: ('c', 9), ('c', 12). assert_that(global_sum, equal_to([7, 21]), label='global sum') assert_that( sum_per_key, equal_to([('c', 1), ('c', 21), ('d', 6)]), label='sum per key') def test_fixed_windows_combine(self): with TestPipeline() as p: input = ( p | beam.Create([('c', 1), ('c', 2), ('c', 10), ('d', 5), ('d', 8), ('d', 9)]) | beam.MapTuple(lambda k, v: window.TimestampedValue((k, v), v)) | beam.WindowInto(window.FixedWindows(4))) global_sum = ( input | beam.Values() | beam.CombineGlobally(sum).without_defaults()) sum_per_key = input | beam.CombinePerKey(sum) # The first window has 2 elements: ('c', 1), ('c', 2). # The second window has 1 elements: ('d', 5). # The third window has 3 elements: ('c', 10), ('d', 8), ('d', 9). assert_that(global_sum, equal_to([3, 5, 27]), label='global sum') assert_that( sum_per_key, equal_to([('c', 3), ('c', 10), ('d', 5), ('d', 17)]), label='sum per key') # Test that three different kinds of metrics work with a customized # SortedConcatWithCounters CombineFn. def test_custormized_counters_in_combine_fn(self): p = TestPipeline() input = ( p | beam.Create([('key1', 'a'), ('key1', 'ab'), ('key1', 'abc'), ('key2', 'uvxy'), ('key2', 'uvxyz')])) # The result of concatenating all values regardless of key. global_concat = ( input | beam.Values() | beam.CombineGlobally(SortedConcatWithCounters())) # The (key, concatenated_string) pairs for all keys. concat_per_key = (input | beam.CombinePerKey(SortedConcatWithCounters())) # Verify the concatenated strings are correct. expected_concat_per_key = [('key1', 'aaabbc'), ('key2', 'uuvvxxyyz')] assert_that( global_concat, equal_to(['aaabbcuuvvxxyyz']), label='global concat') assert_that( concat_per_key, equal_to(expected_concat_per_key), label='concat per key') result = p.run() result.wait_until_finish() # Verify the values of metrics are correct. word_counter_filter = MetricsFilter().with_name('word_counter') query_result = result.metrics().query(word_counter_filter) if query_result['counters']: word_counter = query_result['counters'][0] self.assertEqual(word_counter.result, 5) word_lengths_filter = MetricsFilter().with_name('word_lengths') query_result = result.metrics().query(word_lengths_filter) if query_result['counters']: word_lengths = query_result['counters'][0] self.assertEqual(word_lengths.result, 15) word_len_dist_filter = MetricsFilter().with_name('word_len_dist') query_result = result.metrics().query(word_len_dist_filter) if query_result['distributions']: word_len_dist = query_result['distributions'][0] self.assertEqual(word_len_dist.result.mean, 3) last_word_len_filter = MetricsFilter().with_name('last_word_len') query_result = result.metrics().query(last_word_len_filter) if query_result['gauges']: last_word_len = query_result['gauges'][0] self.assertIn(last_word_len.result.value, [1, 2, 3, 4, 5]) # Test that three different kinds of metrics work with the customized # SortedConcatWithCounters CombineFn when the PCollection is empty. def test_custormized_counters_in_combine_fn_empty(self): p = TestPipeline() input = p | beam.Create([]) # The result of concatenating all values regardless of key. global_concat = ( input | beam.Values() | beam.CombineGlobally(SortedConcatWithCounters())) # The (key, concatenated_string) pairs for all keys. concat_per_key = (input | beam.CombinePerKey(SortedConcatWithCounters())) # Verify the concatenated strings are correct. assert_that(global_concat, equal_to(['']), label='global concat') assert_that(concat_per_key, equal_to([]), label='concat per key') result = p.run() result.wait_until_finish() # Verify the values of metrics are correct. word_counter_filter = MetricsFilter().with_name('word_counter') query_result = result.metrics().query(word_counter_filter) if query_result['counters']: word_counter = query_result['counters'][0] self.assertEqual(word_counter.result, 0) word_lengths_filter = MetricsFilter().with_name('word_lengths') query_result = result.metrics().query(word_lengths_filter) if query_result['counters']: word_lengths = query_result['counters'][0] self.assertEqual(word_lengths.result, 0) word_len_dist_filter = MetricsFilter().with_name('word_len_dist') query_result = result.metrics().query(word_len_dist_filter) if query_result['distributions']: word_len_dist = query_result['distributions'][0] self.assertEqual(word_len_dist.result.count, 0) last_word_len_filter = MetricsFilter().with_name('last_word_len') query_result = result.metrics().query(last_word_len_filter) # No element has ever been recorded. self.assertFalse(query_result['gauges']) class LatestTest(unittest.TestCase): def test_globally(self): l = [ window.TimestampedValue(3, 100), window.TimestampedValue(1, 200), window.TimestampedValue(2, 300) ] with TestPipeline() as p: # Map(lambda x: x) PTransform is added after Create here, because when # a PCollection of TimestampedValues is created with Create PTransform, # the timestamps are not assigned to it. Adding a Map forces the # PCollection to go through a DoFn so that the PCollection consists of # the elements with timestamps assigned to them instead of a PCollection # of TimestampedValue(element, timestamp). pcoll = p | Create(l) | Map(lambda x: x) latest = pcoll | combine.Latest.Globally() assert_that(latest, equal_to([2])) # Now for global combines without default windowed = pcoll | 'window' >> WindowInto(FixedWindows(180)) result_windowed = ( windowed | 'latest wo defaults' >> combine.Latest.Globally().without_defaults()) assert_that(result_windowed, equal_to([3, 2]), label='latest-wo-defaults') def test_globally_empty(self): l = [] with TestPipeline() as p: pc = p | Create(l) | Map(lambda x: x) latest = pc | combine.Latest.Globally() assert_that(latest, equal_to([None])) def test_per_key(self): l = [ window.TimestampedValue(('a', 1), 300), window.TimestampedValue(('b', 3), 100), window.TimestampedValue(('a', 2), 200) ] with TestPipeline() as p: pc = p | Create(l) | Map(lambda x: x) latest = pc | combine.Latest.PerKey() assert_that(latest, equal_to([('a', 1), ('b', 3)])) def test_per_key_empty(self): l = [] with TestPipeline() as p: pc = p | Create(l) | Map(lambda x: x) latest = pc | combine.Latest.PerKey() assert_that(latest, equal_to([])) class LatestCombineFnTest(unittest.TestCase): def setUp(self): self.fn = combine.LatestCombineFn() def test_create_accumulator(self): accumulator = self.fn.create_accumulator() self.assertEqual(accumulator, (None, window.MIN_TIMESTAMP)) def test_add_input(self): accumulator = self.fn.create_accumulator() element = (1, 100) new_accumulator = self.fn.add_input(accumulator, element) self.assertEqual(new_accumulator, (1, 100)) def test_merge_accumulators(self): accumulators = [(2, 400), (5, 100), (9, 200)] merged_accumulator = self.fn.merge_accumulators(accumulators) self.assertEqual(merged_accumulator, (2, 400)) def test_extract_output(self): accumulator = (1, 100) output = self.fn.extract_output(accumulator) self.assertEqual(output, 1) def test_with_input_types_decorator_violation(self): l_int = [1, 2, 3] l_dict = [{'a': 3}, {'g': 5}, {'r': 8}] l_3_tuple = [(12, 31, 41), (12, 34, 34), (84, 92, 74)] with self.assertRaises(TypeCheckError): with TestPipeline() as p: pc = p | Create(l_int) _ = pc | beam.CombineGlobally(self.fn) with self.assertRaises(TypeCheckError): with TestPipeline() as p: pc = p | Create(l_dict) _ = pc | beam.CombineGlobally(self.fn) with self.assertRaises(TypeCheckError): with TestPipeline() as p: pc = p | Create(l_3_tuple) _ = pc | beam.CombineGlobally(self.fn) @pytest.mark.it_validatesrunner class CombineValuesTest(unittest.TestCase): def test_gbk_immediately_followed_by_combine(self): def merge(vals): return "".join(vals) with TestPipeline() as p: result = ( p \ | Create([("key1", "foo"), ("key2", "bar"), ("key1", "foo")], reshuffle=False) \ | beam.GroupByKey() \ | beam.CombineValues(merge) \ | beam.MapTuple(lambda k, v: '{}: {}'.format(k, v))) assert_that(result, equal_to(['key1: foofoo', 'key2: bar'])) # # Test cases for streaming. # @pytest.mark.it_validatesrunner class TimestampCombinerTest(unittest.TestCase): def test_combiner_earliest(self): """Test TimestampCombiner with EARLIEST.""" options = PipelineOptions(streaming=True) with TestPipeline(options=options) as p: result = ( p | TestStream().add_elements([window.TimestampedValue( ('k', 100), 2)]).add_elements( [window.TimestampedValue( ('k', 400), 7)]).advance_watermark_to_infinity() | beam.WindowInto( window.FixedWindows(10), timestamp_combiner=TimestampCombiner.OUTPUT_AT_EARLIEST) | beam.CombinePerKey(sum)) records = ( result | beam.Map(lambda e, ts=beam.DoFn.TimestampParam: (e, ts))) # All the KV pairs are applied GBK using EARLIEST timestamp for the same # key. expected_window_to_elements = { window.IntervalWindow(0, 10): [ (('k', 500), Timestamp(2)), ], } assert_that( records, equal_to_per_window(expected_window_to_elements), use_global_window=False, label='assert per window') def test_combiner_latest(self): """Test TimestampCombiner with LATEST.""" options = PipelineOptions(streaming=True) with TestPipeline(options=options) as p: result = ( p | TestStream().add_elements([window.TimestampedValue( ('k', 100), 2)]).add_elements( [window.TimestampedValue( ('k', 400), 7)]).advance_watermark_to_infinity() | beam.WindowInto( window.FixedWindows(10), timestamp_combiner=TimestampCombiner.OUTPUT_AT_LATEST) | beam.CombinePerKey(sum)) records = ( result | beam.Map(lambda e, ts=beam.DoFn.TimestampParam: (e, ts))) # All the KV pairs are applied GBK using LATEST timestamp for # the same key. expected_window_to_elements = { window.IntervalWindow(0, 10): [ (('k', 500), Timestamp(7)), ], } assert_that( records, equal_to_per_window(expected_window_to_elements), use_global_window=False, label='assert per window') class CombineGloballyTest(unittest.TestCase): def test_combine_globally_for_unbounded_source_with_default(self): # this error is logged since the below combination is ill-defined. with self.assertLogs() as captured_logs: with TestPipeline() as p: _ = ( p | PeriodicImpulse( start_timestamp=time.time(), stop_timestamp=time.time() + 4, fire_interval=1, apply_windowing=False, ) | beam.Map(lambda x: ('c', 1)) | beam.WindowInto( window.GlobalWindows(), trigger=trigger.Repeatedly(trigger.AfterCount(2)), accumulation_mode=trigger.AccumulationMode.DISCARDING, ) | beam.combiners.Count.Globally()) self.assertIn('unbounded collections', '\n'.join(captured_logs.output)) def test_combine_globally_for_unbounded_source_without_defaults(self): # this is the supported case with TestPipeline() as p: _ = ( p | PeriodicImpulse( start_timestamp=time.time(), stop_timestamp=time.time() + 4, fire_interval=1, apply_windowing=False, ) | beam.Map(lambda x: 1) | beam.WindowInto( window.GlobalWindows(), trigger=trigger.Repeatedly(trigger.AfterCount(2)), accumulation_mode=trigger.AccumulationMode.DISCARDING, ) | beam.CombineGlobally(sum).without_defaults()) def get_common_items(sets, excluded_chars=""): # set.intersection() takes multiple sets as separete arguments. # We unpack the `sets` list into multiple arguments with the * operator. # The combine transform might give us an empty list of `sets`, # so we use a list with an empty set as a default value. common = set.intersection(*(sets or [set()])) return common.difference(excluded_chars) class CombinerWithSideInputs(unittest.TestCase): def test_cpk_with_side_input(self): test_cases = [(get_common_items, True), (beam.CombineFn.from_callable(get_common_items), True), (get_common_items, False), (beam.CombineFn.from_callable(get_common_items), False)] for combiner, with_kwarg in test_cases: self._check_combineperkey_with_side_input(combiner, with_kwarg) self._check_combineglobally_with_side_input(combiner, with_kwarg) def _check_combineperkey_with_side_input(self, combiner, with_kwarg): with beam.Pipeline() as pipeline: pc = (pipeline | beam.Create(['🍅'])) if with_kwarg: cpk = beam.CombinePerKey( combiner, excluded_chars=beam.pvalue.AsSingleton(pc)) else: cpk = beam.CombinePerKey(combiner, beam.pvalue.AsSingleton(pc)) common_items = ( pipeline | 'Create produce' >> beam.Create([ {'🍓', '🥕', '🍌', '🍅', '🌶️'}, {'🍇', '🥕', '🥝', '🍅', '🥔'}, {'🍉', '🥕', '🍆', '🍅', '🍍'}, {'🥑', '🥕', '🌽', '🍅', '🥥'}, ]) | beam.WithKeys(lambda x: None) | cpk) assert_that(common_items, equal_to([(None, {'🥕'})])) def _check_combineglobally_with_side_input(self, combiner, with_kwarg): with beam.Pipeline() as pipeline: pc = (pipeline | beam.Create(['🍅'])) if with_kwarg: cpk = beam.CombineGlobally( combiner, excluded_chars=beam.pvalue.AsSingleton(pc)) else: cpk = beam.CombineGlobally(combiner, beam.pvalue.AsSingleton(pc)) common_items = ( pipeline | 'Create produce' >> beam.Create([ {'🍓', '🥕', '🍌', '🍅', '🌶️'}, {'🍇', '🥕', '🥝', '🍅', '🥔'}, {'🍉', '🥕', '🍆', '🍅', '🍍'}, {'🥑', '🥕', '🌽', '🍅', '🥥'}, ]) | cpk) assert_that(common_items, equal_to([{'🥕'}])) def test_combinefn_methods_with_side_input(self): # Test that the expected combinefn methods are called with the # expected arguments when using side inputs in CombinePerKey. with tempfile.TemporaryDirectory() as tmp_dirname: fname = str(Path(tmp_dirname) / "combinefn_calls.json") with open(fname, "w") as f: json.dump({}, f) def set_in_json(key, values): current_json = {} if os.path.exists(fname): with open(fname, "r") as f: current_json = json.load(f) current_json[key] = values with open(fname, "w") as f: json.dump(current_json, f) class MyCombiner(beam.CombineFn): def create_accumulator(self, *args, **kwargs): set_in_json("create_accumulator_args", args) set_in_json("create_accumulator_kwargs", kwargs) return args, kwargs def add_input(self, accumulator, input, *args, **kwargs): set_in_json("add_input_args", args) set_in_json("add_input_kwargs", kwargs) return accumulator def merge_accumulators(self, accumulators, *args, **kwargs): set_in_json("merge_accumulators_args", args) set_in_json("merge_accumulators_kwargs", kwargs) return args, kwargs def compact(self, accumulator, *args, **kwargs): set_in_json("compact_args", args) set_in_json("compact_kwargs", kwargs) return accumulator def extract_output(self, accumulator, *args, **kwargs): set_in_json("extract_output_args", args) set_in_json("extract_output_kwargs", kwargs) return accumulator with beam.Pipeline() as p: static_pos_arg = 0 deferred_pos_arg = beam.pvalue.AsSingleton( p | "CreateDeferredSideInput" >> beam.Create([1])) static_kwarg = 2 deferred_kwarg = beam.pvalue.AsSingleton( p | "CreateDeferredSideInputKwarg" >> beam.Create([3])) res = ( p | "CreateInputs" >> beam.Create([(None, None)]) | beam.CombinePerKey( MyCombiner(), static_pos_arg, deferred_pos_arg, static_kwarg=static_kwarg, deferred_kwarg=deferred_kwarg)) assert_that( res, equal_to([ (None, ((0, 1), { 'static_kwarg': 2, 'deferred_kwarg': 3 })) ])) # Check that the combinefn was called with the expected arguments with open(fname, "r") as f: data = json.load(f) expected_args = [0, 1] expected_kwargs = {"static_kwarg": 2, "deferred_kwarg": 3} method_names = [ "create_accumulator", "compact", "add_input", "merge_accumulators", "extract_output" ] for key in method_names: print(f"Checking {key}") self.assertEqual(data[key + "_args"], expected_args) self.assertEqual(data[key + "_kwargs"], expected_kwargs) def test_cpk_with_windows(self): # With global window side input with TestPipeline() as p: def sum_with_floor(vals, min_value=0): vals_sum = sum(vals) if vals_sum < min_value: vals_sum += min_value return vals_sum res = ( p | "CreateInputs" >> beam.Create([1, 2, 100, 101, 102]) | beam.Map(lambda x: window.TimestampedValue(('k', x), x)) | beam.WindowInto(FixedWindows(99)) | beam.CombinePerKey( sum_with_floor, min_value=pvalue.AsSingleton(p | beam.Create([100])))) assert_that(res, equal_to([('k', 103), ('k', 303)])) # with matching window side input with TestPipeline() as p: min_value = ( p | "CreateMinValue" >> beam.Create([ window.TimestampedValue(50, 5), window.TimestampedValue(1000, 100) ]) | "WindowSideInputs" >> beam.WindowInto(FixedWindows(99))) res = ( p | "CreateInputs" >> beam.Create([1, 2, 100, 101, 102]) | beam.Map(lambda x: window.TimestampedValue(('k', x), x)) | beam.WindowInto(FixedWindows(99)) | beam.CombinePerKey( sum_with_floor, min_value=pvalue.AsSingleton(min_value))) assert_that(res, equal_to([('k', 53), ('k', 1303)])) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/core.py0000644000175100017510000047760715206113033021403 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Core PTransform subclasses, such as FlatMap, GroupByKey, and Map.""" # pytype: skip-file import concurrent.futures import copy import inspect import logging import random import sys import time import traceback import types import typing from collections import defaultdict from functools import wraps from itertools import dropwhile from apache_beam import coders from apache_beam import pvalue from apache_beam import typehints from apache_beam.coders import typecoders from apache_beam.internal import pickler from apache_beam.internal import util from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import TypeOptions from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.transforms import ptransform from apache_beam.transforms import userstate from apache_beam.transforms.display import DisplayDataItem from apache_beam.transforms.display import HasDisplayData from apache_beam.transforms.ptransform import PTransform from apache_beam.transforms.ptransform import PTransformWithSideInputs from apache_beam.transforms.sideinputs import SIDE_INPUT_PREFIX from apache_beam.transforms.sideinputs import get_sideinput_index from apache_beam.transforms.userstate import StateSpec from apache_beam.transforms.userstate import TimerSpec from apache_beam.transforms.window import GlobalWindows from apache_beam.transforms.window import SlidingWindows from apache_beam.transforms.window import TimestampCombiner from apache_beam.transforms.window import TimestampedValue from apache_beam.transforms.window import WindowedValue from apache_beam.transforms.window import WindowFn from apache_beam.typehints import row_type from apache_beam.typehints import trivial_inference from apache_beam.typehints.batch import BatchConverter from apache_beam.typehints.decorators import TypeCheckError from apache_beam.typehints.decorators import WithTypeHints from apache_beam.typehints.decorators import get_signature from apache_beam.typehints.decorators import get_type_hints from apache_beam.typehints.decorators import with_input_types from apache_beam.typehints.decorators import with_output_types from apache_beam.typehints.trivial_inference import element_type from apache_beam.typehints.typehints import TypeConstraint from apache_beam.typehints.typehints import is_consistent_with from apache_beam.typehints.typehints import visit_inner_types from apache_beam.utils import urns from apache_beam.utils.timestamp import Duration if typing.TYPE_CHECKING: from google.protobuf import message # pylint: disable=ungrouped-imports from apache_beam.io import iobase from apache_beam.pipeline import Pipeline from apache_beam.runners.pipeline_context import PipelineContext from apache_beam.transforms import create_source from apache_beam.transforms.trigger import AccumulationMode from apache_beam.transforms.trigger import DefaultTrigger from apache_beam.transforms.trigger import TriggerFn __all__ = [ 'DoFn', 'CombineFn', 'PartitionFn', 'ParDo', 'FlatMap', 'FlatMapTuple', 'Map', 'MapTuple', 'Filter', 'CombineGlobally', 'CombinePerKey', 'CombineValues', 'GroupBy', 'GroupByKey', 'Select', 'Partition', 'Windowing', 'WindowInto', 'Flatten', 'FlattenWith', 'Create', 'Impulse', 'RestrictionProvider', 'WatermarkEstimatorProvider', ] # Type variables T = typing.TypeVar('T') K = typing.TypeVar('K') V = typing.TypeVar('V') _LOGGER = logging.getLogger(__name__) class DoFnContext(object): """A context available to all methods of DoFn instance.""" pass class DoFnProcessContext(DoFnContext): """A processing context passed to DoFn process() during execution. Most importantly, a DoFn.process method will access context.element to get the element it is supposed to process. Attributes: label: label of the ParDo whose element is being processed. element: element being processed (in process method only; always None in start_bundle and finish_bundle) timestamp: timestamp of the element (in process method only; always None in start_bundle and finish_bundle) windows: windows of the element (in process method only; always None in start_bundle and finish_bundle) state: a DoFnState object, which holds the runner's internal state for this element. Not used by the pipeline code. """ def __init__(self, label, element=None, state=None): """Initialize a processing context object with an element and state. The element represents one value from a PCollection that will be accessed by a DoFn object during pipeline execution, and state is an arbitrary object where counters and other pipeline state information can be passed in. DoFnProcessContext objects are also used as inputs to PartitionFn instances. Args: label: label of the PCollection whose element is being processed. element: element of a PCollection being processed using this context. state: a DoFnState object with state to be passed in to the DoFn object. """ self.label = label self.state = state if element is not None: self.set_element(element) def set_element(self, windowed_value): if windowed_value is None: # Not currently processing an element. if hasattr(self, 'element'): del self.element del self.timestamp del self.windows else: self.element = windowed_value.value self.timestamp = windowed_value.timestamp self.windows = windowed_value.windows class ProcessContinuation(object): """An object that may be produced as the last element of a process method invocation. If produced, indicates that there is more work to be done for the current input element. """ def __init__(self, resume_delay=0): """Initializes a ProcessContinuation object. Args: resume_delay: indicates the minimum time, in seconds, that should elapse before re-invoking process() method for resuming the invocation of the current element. """ self.resume_delay = resume_delay @staticmethod def resume(resume_delay=0): """A convenient method that produces a ``ProcessContinuation``. Args: resume_delay: delay after which processing current element should be resumed. Returns: a ``ProcessContinuation`` for signalling the runner that current input element has not been fully processed and should be resumed later. """ return ProcessContinuation(resume_delay=resume_delay) class RestrictionProvider(object): """Provides methods for generating and manipulating restrictions. This class should be implemented to support Splittable ``DoFn`` in Python SDK. See https://s.apache.org/splittable-do-fn for more details about Splittable ``DoFn``. To denote a ``DoFn`` class to be Splittable ``DoFn``, ``DoFn.process()`` method of that class should have exactly one parameter whose default value is an instance of ``RestrictionParam``. This ``RestrictionParam`` can either be constructed with an explicit ``RestrictionProvider``, or, if no ``RestrictionProvider`` is provided, the ``DoFn`` itself must be a ``RestrictionProvider``. The provided ``RestrictionProvider`` instance must provide suitable overrides for the following methods: * create_tracker() * initial_restriction() * restriction_size() Optionally, ``RestrictionProvider`` may override default implementations of following methods: * restriction_coder() * split() * split_and_size() * truncate() ** Pausing and resuming processing of an element ** As the last element produced by the iterator returned by the ``DoFn.process()`` method, a Splittable ``DoFn`` may return an object of type ``ProcessContinuation``. If restriction_tracker.defer_remander is called in the ```DoFn.process()``, it means that runner should later re-invoke ``DoFn.process()`` method to resume processing the current element and the manner in which the re-invocation should be performed. ** Updating output watermark ** ``DoFn.process()`` method of Splittable ``DoFn``s could contain a parameter with default value ``DoFn.WatermarkReporterParam``. If specified this asks the runner to provide a function that can be used to give the runner a (best-effort) lower bound about the timestamps of future output associated with the current element processed by the ``DoFn``. If the ``DoFn`` has multiple outputs, the watermark applies to all of them. Provided function must be invoked with a single parameter of type ``Timestamp`` or as an integer that gives the watermark in number of seconds. """ def create_tracker(self, restriction): # type: (...) -> iobase.RestrictionTracker """Produces a new ``RestrictionTracker`` for the given restriction. This API is required to be implemented. Args: restriction: an object that defines a restriction as identified by a Splittable ``DoFn`` that utilizes the current ``RestrictionProvider``. For example, a tuple that gives a range of positions for a Splittable ``DoFn`` that reads files based on byte positions. Returns: an object of type ``RestrictionTracker``. """ raise NotImplementedError def initial_restriction(self, element): """Produces an initial restriction for the given element. This API is required to be implemented. """ raise NotImplementedError def split(self, element, restriction): """Splits the given element and restriction initially. This method enables runners to perform bulk splitting initially allowing for a rapid increase in parallelism. Note that initial split is a different concept from the split during element processing time. Please refer to ``iobase.RestrictionTracker.try_split`` for details about splitting when the current element and restriction are actively being processed. Returns an iterator of restrictions. The total set of elements produced by reading input element for each of the returned restrictions should be the same as the total set of elements produced by reading the input element for the input restriction. This API is optional if ``split_and_size`` has been implemented. If this method is not override, there is no initial splitting happening on each restriction. """ yield restriction def restriction_coder(self): """Returns a ``Coder`` for restrictions. Returned``Coder`` will be used for the restrictions produced by the current ``RestrictionProvider``. Returns: an object of type ``Coder``. """ return coders.registry.get_coder(object) def restriction_size(self, element, restriction): """Returns the size of a restriction with respect to the given element. By default, asks a newly-created restriction tracker for the default size of the restriction. The return value must be non-negative. Must be thread safe. Will be invoked concurrently during bundle processing due to runner initiated splitting and progress estimation. This API is required to be implemented. """ raise NotImplementedError def split_and_size(self, element, restriction): """Like split, but also does sizing, returning (restriction, size) pairs. For each pair, size must be non-negative. This API is optional if ``split`` and ``restriction_size`` have been implemented. """ for part in self.split(element, restriction): yield part, self.restriction_size(element, part) def truncate(self, element, restriction): """Truncates the provided restriction into a restriction representing a finite amount of work when the pipeline is `draining for additional details about drain.>_`. # pylint: disable=line-too-long By default, if the restriction is bounded then the restriction will be returned otherwise None will be returned. This API is optional and should only be implemented if more granularity is required. Return a truncated finite restriction if further processing is required otherwise return None to represent that no further processing of this restriction is required. The default behavior when a pipeline is being drained is that bounded restrictions process entirely while unbounded restrictions process till a checkpoint is possible. """ restriction_tracker = self.create_tracker(restriction) if restriction_tracker.is_bounded(): return restriction def get_function_arguments(obj, func): # type: (...) -> typing.Tuple[typing.List[str], typing.List[typing.Any]] """Return the function arguments based on the name provided. If they have a _inspect_function attached to the class then use that otherwise default to the modified version of python inspect library. Returns: Same as get_function_args_defaults. """ func_name = '_inspect_%s' % func if hasattr(obj, func_name): f = getattr(obj, func_name) return f() f = getattr(obj, func) return get_function_args_defaults(f) def get_function_args_defaults(f): # type: (...) -> typing.Tuple[typing.List[str], typing.List[typing.Any]] """Returns the function arguments of a given function. Returns: (args: List[str], defaults: List[Any]). The first list names the arguments of the method and the second one has the values of the default arguments. This is similar to ``inspect.getfullargspec()``'s results, except it doesn't include bound arguments and may follow function wrappers. """ signature = get_signature(f) parameter = inspect.Parameter # TODO(BEAM-5878) support kwonlyargs on Python 3. _SUPPORTED_ARG_TYPES = [ parameter.POSITIONAL_ONLY, parameter.POSITIONAL_OR_KEYWORD ] args = [ name for name, p in signature.parameters.items() if p.kind in _SUPPORTED_ARG_TYPES ] defaults = [ p.default for p in signature.parameters.values() if p.kind in _SUPPORTED_ARG_TYPES and p.default is not p.empty ] return args, defaults class WatermarkEstimatorProvider(object): """Provides methods for generating WatermarkEstimator. This class should be implemented if wanting to providing output_watermark information within an SDF. In order to make an SDF.process() access to the typical WatermarkEstimator, the SDF author should have an argument whose default value is a DoFn.WatermarkEstimatorParam instance. This DoFn.WatermarkEstimatorParam can either be constructed with an explicit WatermarkEstimatorProvider, or, if no WatermarkEstimatorProvider is provided, the DoFn itself must be a WatermarkEstimatorProvider. """ def initial_estimator_state(self, element, restriction): """Returns the initial state of the WatermarkEstimator with given element and restriction. This function is called by the system. """ raise NotImplementedError def create_watermark_estimator(self, estimator_state): """Create a new WatermarkEstimator based on the state. The state is typically useful when resuming processing an element. """ raise NotImplementedError def estimator_state_coder(self): return coders.registry.get_coder(object) class _DoFnParam(object): """DoFn parameter.""" def __init__(self, param_id): self.param_id = param_id def __eq__(self, other): if type(self) == type(other): return self.param_id == other.param_id return False def __hash__(self): return hash(self.param_id) def __repr__(self): return self.param_id class _RestrictionDoFnParam(_DoFnParam): """Restriction Provider DoFn parameter.""" def __init__(self, restriction_provider=None): # type: (typing.Optional[RestrictionProvider]) -> None if (restriction_provider is not None and not isinstance(restriction_provider, RestrictionProvider)): raise ValueError( 'DoFn.RestrictionParam expected RestrictionProvider object.') self.restriction_provider = restriction_provider self.param_id = ( 'RestrictionParam(%s)' % restriction_provider.__class__.__name__) class _StateDoFnParam(_DoFnParam): """State DoFn parameter.""" def __init__(self, state_spec): # type: (StateSpec) -> None if not isinstance(state_spec, StateSpec): raise ValueError("DoFn.StateParam expected StateSpec object.") self.state_spec = state_spec self.param_id = 'StateParam(%s)' % state_spec.name class _TimerDoFnParam(_DoFnParam): """Timer DoFn parameter.""" def __init__(self, timer_spec): # type: (TimerSpec) -> None if not isinstance(timer_spec, TimerSpec): raise ValueError("DoFn.TimerParam expected TimerSpec object.") self.timer_spec = timer_spec self.param_id = 'TimerParam(%s)' % timer_spec.name class _BundleFinalizerParam(_DoFnParam): """Bundle Finalization DoFn parameter.""" def __init__(self): self._callbacks = [] self.param_id = "FinalizeBundle" def register(self, callback): self._callbacks.append(callback) # Log errors when calling callback to make sure all callbacks get called # though there are errors. And errors should not fail pipeline. def finalize_bundle(self): for callback in self._callbacks: try: callback() except Exception as e: _LOGGER.warning("Got exception from finalization call: %s", e) def has_callbacks(self): # type: () -> bool return len(self._callbacks) > 0 def reset(self): # type: () -> None del self._callbacks[:] class _WatermarkEstimatorParam(_DoFnParam): """WatermarkEstimator DoFn parameter.""" def __init__( self, watermark_estimator_provider: typing. Optional[WatermarkEstimatorProvider] = None): if (watermark_estimator_provider is not None and not isinstance( watermark_estimator_provider, WatermarkEstimatorProvider)): raise ValueError( 'DoFn.WatermarkEstimatorParam expected' 'WatermarkEstimatorProvider object.') self.watermark_estimator_provider = watermark_estimator_provider self.param_id = 'WatermarkEstimatorProvider' class _ContextParam(_DoFnParam): def __init__( self, context_manager_constructor, args=(), kwargs=None, *, name=None): class_name = self.__class__.__name__.strip('_') if (not callable(context_manager_constructor) or (hasattr(context_manager_constructor, '__enter__') and len(inspect.signature( context_manager_constructor.__enter__).parameters) == 0)): # Context managers constructed with @contextlib.contextmanager can only # be used once, and in addition cannot be pickled because they invoke # the function on __init__ rather than at __enter__. # In addition, other common context managers such as # tempfile.TemporaryDirectory perform side-effecting actions in __init__ # rather than in __enter__. raise TypeError( "A context manager constructor (not a fully constructed context " "manager) must be passed to avoid issues with one-shot managers. " "For example, " "write {class_name}(tempfile.TemporaryDirectory, args=(...)) " "rather than {class_name}(tempfile.TemporaryDirectory(...))") super().__init__(f'{class_name}_{name or id(self)}') self.context_manager_constructor = context_manager_constructor self.args = args self.kwargs = kwargs or {} def create_and_enter(self): cm = self.context_manager_constructor(*self.args, **self.kwargs) return cm, cm.__enter__() class _BundleContextParam(_ContextParam): """Allows one to use a context manager to manage bundle-scoped parameters. The context will be entered at the start of each bundle and exited at the end, equivalent to the `start_bundle` and `finish_bundle` methods on a DoFn. The object returned from `__enter__`, if any, will be substituted for this parameter in invocations. Multiple context manager parameters may be specified which will all be evaluated (in an unspecified order). This can be especially useful for setting up shared context in transforms like `Map`, `FlatMap`, and `Filter` where one does not have start_bundle and finish_bundle methods. """ class _SetupContextParam(_ContextParam): """Allows one to use a context manager to manage DoFn-scoped parameters. The context will be entered before the DoFn is used and exited when it is discarded, equivalent to the `setup` and `teardown` methods of a DoFn. (Note, like `teardown`, exiting is best effort, as workers may be killed before all DoFns are torn down.) The object returned from `__enter__`, if any, will be substituted for this parameter in invocations. Multiple context manager parameters may be specified which will all be evaluated (in an unspecified order). This can be useful for setting up shared resources like persistent connections to external services for transforms like `Map`, `FlatMap`, and `Filter` where one does not have setup and teardown methods. """ class DoFn(WithTypeHints, HasDisplayData, urns.RunnerApiFn): """A function object used by a transform with custom processing. The ParDo transform is such a transform. The ParDo.apply method will take an object of type DoFn and apply it to all elements of a PCollection object. In order to have concrete DoFn objects one has to subclass from DoFn and define the desired behavior (start_bundle/finish_bundle and process) or wrap a callable object using the CallableWrapperDoFn class. """ # Parameters that can be used in the .process() method. ElementParam = _DoFnParam('ElementParam') SideInputParam = _DoFnParam('SideInputParam') TimestampParam = _DoFnParam('TimestampParam') WindowParam = _DoFnParam('WindowParam') WindowedValueParam = _DoFnParam('WindowedValueParam') PaneInfoParam = _DoFnParam('PaneInfoParam') WatermarkEstimatorParam = _WatermarkEstimatorParam BundleFinalizerParam = _BundleFinalizerParam KeyParam = _DoFnParam('KeyParam') BundleContextParam = _BundleContextParam SetupContextParam = _SetupContextParam # Parameters to access state and timers. Not restricted to use only in the # .process() method. Usage: DoFn.StateParam(state_spec), # DoFn.TimerParam(timer_spec), DoFn.TimestampParam, DoFn.WindowParam, # DoFn.KeyParam StateParam = _StateDoFnParam TimerParam = _TimerDoFnParam DynamicTimerTagParam = _DoFnParam('DynamicTimerTagParam') DoFnProcessParams = [ ElementParam, SideInputParam, TimestampParam, WindowParam, WindowedValueParam, WatermarkEstimatorParam, PaneInfoParam, BundleFinalizerParam, KeyParam, StateParam, TimerParam, BundleContextParam, SetupContextParam, ] RestrictionParam = _RestrictionDoFnParam @staticmethod def from_callable(fn): return CallableWrapperDoFn(fn) @staticmethod def unbounded_per_element(): """A decorator on process fn specifying that the fn performs an unbounded amount of work per input element.""" def wrapper(process_fn): process_fn.unbounded_per_element = True return process_fn return wrapper @staticmethod def yields_elements(fn): """A decorator to apply to ``process_batch`` indicating it yields elements. By default ``process_batch`` is assumed to both consume and produce "batches", which are collections of multiple logical Beam elements. This decorator indicates that ``process_batch`` **produces** individual elements at a time. ``process_batch`` is always expected to consume batches. """ if not fn.__name__ in ('process', 'process_batch'): raise TypeError( "@yields_elements must be applied to a process or " f"process_batch method, got {fn!r}.") fn._beam_yields_elements = True return fn @staticmethod def yields_batches(fn): """A decorator to apply to ``process`` indicating it yields batches. By default ``process`` is assumed to both consume and produce individual elements at a time. This decorator indicates that ``process`` **produces** "batches", which are collections of multiple logical Beam elements. """ if not fn.__name__ in ('process', 'process_batch'): raise TypeError( "@yields_elements must be applied to a process or " f"process_batch method, got {fn!r}.") fn._beam_yields_batches = True return fn def default_label(self): return self.__class__.__name__ def process(self, element, *args, **kwargs): """Method to use for processing elements. This is invoked by ``DoFnRunner`` for each element of a input ``PCollection``. The following parameters can be used as default values on ``process`` arguments to indicate that a DoFn accepts the corresponding parameters. For example, a DoFn might accept the element and its timestamp with the following signature:: def process(element=DoFn.ElementParam, timestamp=DoFn.TimestampParam): ... The full set of parameters is: - ``DoFn.ElementParam``: element to be processed, should not be mutated. - ``DoFn.SideInputParam``: a side input that may be used when processing. - ``DoFn.TimestampParam``: timestamp of the input element. - ``DoFn.WindowParam``: ``Window`` the input element belongs to. - ``DoFn.TimerParam``: a ``userstate.RuntimeTimer`` object defined by the spec of the parameter. - ``DoFn.StateParam``: a ``userstate.RuntimeState`` object defined by the spec of the parameter. - ``DoFn.KeyParam``: key associated with the element. - ``DoFn.RestrictionParam``: an ``iobase.RestrictionTracker`` will be provided here to allow treatment as a Splittable ``DoFn``. The restriction tracker will be derived from the restriction provider in the parameter. - ``DoFn.WatermarkEstimatorParam``: a function that can be used to track output watermark of Splittable ``DoFn`` implementations. - ``DoFn.BundleContextParam``: allows a shared context manager to be used per bundle - ``DoFn.SetupContextParam``: allows a shared context manager to be used per DoFn Args: element: The element to be processed *args: side inputs **kwargs: other keyword arguments. Returns: An Iterable of output elements or None. """ raise NotImplementedError def process_batch(self, batch, *args, **kwargs): raise NotImplementedError def setup(self): """Called to prepare an instance for processing bundles of elements. This is a good place to initialize transient in-memory resources, such as network connections. The resources can then be disposed in ``DoFn.teardown``. """ pass def start_bundle(self): """Called before a bundle of elements is processed on a worker. Elements to be processed are split into bundles and distributed to workers. Before a worker calls process() on the first element of its bundle, it calls this method. """ pass def finish_bundle(self): """Called after a bundle of elements is processed on a worker. """ pass def teardown(self): """Called to use to clean up this instance before it is discarded. A runner will do its best to call this method on any given instance to prevent leaks of transient resources, however, there may be situations where this is impossible (e.g. process crash, hardware failure, etc.) or unnecessary (e.g. the pipeline is shutting down and the process is about to be killed anyway, so all transient resources will be released automatically by the OS). In these cases, the call may not happen. It will also not be retried, because in such situations the DoFn instance no longer exists, so there's no instance to retry it on. Thus, all work that depends on input elements, and all externally important side effects, must be performed in ``DoFn.process`` or ``DoFn.finish_bundle``. """ pass def get_function_arguments(self, func): return get_function_arguments(self, func) def default_type_hints(self): process_type_hints = typehints.decorators.IOTypeHints.from_callable( self.process) or typehints.decorators.IOTypeHints.empty() if self._process_yields_batches: # process() produces batches, don't use it's output typehint process_type_hints = process_type_hints.with_output_types_from( typehints.decorators.IOTypeHints.empty()) if self._process_batch_yields_elements: # process_batch() produces elements, *do* use it's output typehint # First access the typehint process_batch_type_hints = typehints.decorators.IOTypeHints.from_callable( self.process_batch) or typehints.decorators.IOTypeHints.empty() # Then we deconflict with the typehint from process, if it exists if (process_batch_type_hints.output_types != typehints.decorators.IOTypeHints.empty().output_types): if (process_type_hints.output_types != typehints.decorators.IOTypeHints.empty().output_types and process_batch_type_hints.output_types != process_type_hints.output_types): raise TypeError( f"DoFn {self!r} yields element from both process and " "process_batch, but they have mismatched output typehints:\n" f" process: {process_type_hints.output_types}\n" f" process_batch: {process_batch_type_hints.output_types}") process_type_hints = process_type_hints.with_output_types_from( process_batch_type_hints) try: process_type_hints = process_type_hints.strip_iterable() except ValueError as e: raise ValueError('Return value not iterable: %s: %s' % (self, e)) process_type_hints = process_type_hints.extract_tagged_outputs() # Prefer class decorator type hints for backwards compatibility. return get_type_hints(self.__class__).with_defaults(process_type_hints) # TODO(sourabhbajaj): Do we want to remove the responsibility of these from # the DoFn or maybe the runner def infer_output_type(self, input_type): # TODO(https://github.com/apache/beam/issues/19824): Side inputs types. return trivial_inference.element_type( _strip_output_annotations( trivial_inference.infer_return_type(self.process, [input_type]))) @property def _process_defined(self) -> bool: # Check if this DoFn's process method has been overridden # Note that we retrieve the __func__ attribute, if it exists, to get the # underlying function from the bound method. # If __func__ doesn't exist, self.process was likely overridden with a free # function, as in CallableWrapperDoFn. return getattr(self.process, '__func__', self.process) != DoFn.process @property def _process_batch_defined(self) -> bool: # Check if this DoFn's process_batch method has been overridden # Note that we retrieve the __func__ attribute, if it exists, to get the # underlying function from the bound method. # If __func__ doesn't exist, self.process_batch was likely overridden with # a free function. return getattr( self.process_batch, '__func__', self.process_batch) != DoFn.process_batch @property def _can_yield_batches(self) -> bool: return ((self._process_defined and self._process_yields_batches) or ( self._process_batch_defined and not self._process_batch_yields_elements)) @property def _process_yields_batches(self) -> bool: return getattr(self.process, '_beam_yields_batches', False) @property def _process_batch_yields_elements(self) -> bool: return getattr(self.process_batch, '_beam_yields_elements', False) def get_input_batch_type( self, input_element_type ) -> typing.Optional[typing.Union[TypeConstraint, type]]: """Determine the batch type expected as input to process_batch. The default implementation of ``get_input_batch_type`` simply observes the input typehint for the first parameter of ``process_batch``. A Batched DoFn may override this method if a dynamic approach is required. Args: input_element_type: The **element type** of the input PCollection this DoFn is being applied to. Returns: ``None`` if this DoFn cannot accept batches, else a Beam typehint or a native Python typehint. """ if not self._process_batch_defined: return None input_type = list( inspect.signature(self.process_batch).parameters.values())[0].annotation if input_type == inspect.Signature.empty: # TODO(https://github.com/apache/beam/issues/21652): Consider supporting # an alternative (dynamic?) approach for declaring input type raise TypeError( f"Either {self.__class__.__name__}.process_batch() must have a type " f"annotation on its first parameter, or {self.__class__.__name__} " "must override get_input_batch_type.") return input_type def _get_input_batch_type_normalized(self, input_element_type): return typehints.native_type_compatibility.convert_to_beam_type( self.get_input_batch_type(input_element_type)) def _get_output_batch_type_normalized(self, input_element_type): return typehints.native_type_compatibility.convert_to_beam_type( self.get_output_batch_type(input_element_type)) @staticmethod def _get_element_type_from_return_annotation(method, input_type): return_type = inspect.signature(method).return_annotation if return_type == inspect.Signature.empty: # output type not annotated, try to infer it return_type = trivial_inference.infer_return_type(method, [input_type]) return_type = typehints.native_type_compatibility.convert_to_beam_type( return_type) if isinstance(return_type, typehints.typehints.IterableTypeConstraint): return return_type.inner_type elif isinstance(return_type, typehints.typehints.IteratorTypeConstraint): return return_type.yielded_type else: raise TypeError( "Expected Iterator in return type annotation for " f"{method!r}, did you mean Iterator[{return_type}]? Note Beam DoFn " "process and process_batch methods are expected to produce " "generators - they should 'yield' rather than 'return'.") def get_output_batch_type( self, input_element_type ) -> typing.Optional[typing.Union[TypeConstraint, type]]: """Determine the batch type produced by this DoFn's ``process_batch`` implementation and/or its ``process`` implementation with ``@yields_batch``. The default implementation of this method observes the return type annotations on ``process_batch`` and/or ``process``. A Batched DoFn may override this method if a dynamic approach is required. Args: input_element_type: The **element type** of the input PCollection this DoFn is being applied to. Returns: ``None`` if this DoFn will never yield batches, else a Beam typehint or a native Python typehint. """ output_batch_type = None if self._process_defined and self._process_yields_batches: output_batch_type = self._get_element_type_from_return_annotation( self.process, input_element_type) if self._process_batch_defined and not self._process_batch_yields_elements: process_batch_type = self._get_element_type_from_return_annotation( self.process_batch, self._get_input_batch_type_normalized(input_element_type)) # TODO: Consider requiring an inheritance relationship rather than # equality if (output_batch_type is not None and (not process_batch_type == output_batch_type)): raise TypeError( f"DoFn {self!r} yields batches from both process and " "process_batch, but they produce different types:\n" f" process: {output_batch_type}\n" f" process_batch: {process_batch_type!r}") output_batch_type = process_batch_type return output_batch_type def _process_argspec_fn(self): """Returns the Python callable that will eventually be invoked. This should ideally be the user-level function that is called with the main and (if any) side inputs, and is used to relate the type hint parameters with the input parameters (e.g., by argument name). """ return self.process urns.RunnerApiFn.register_pickle_urn(python_urns.PICKLED_DOFN) class CallableWrapperDoFn(DoFn): """For internal use only; no backwards-compatibility guarantees. A DoFn (function) object wrapping a callable object. The purpose of this class is to conveniently wrap simple functions and use them in transforms. """ def __init__(self, fn, fullargspec=None): """Initializes a CallableWrapperDoFn object wrapping a callable. Args: fn: A callable object. Raises: TypeError: if fn parameter is not a callable type. """ if not callable(fn): raise TypeError('Expected a callable object instead of: %r' % fn) self._fn = fn self._fullargspec = fullargspec if isinstance( fn, (types.BuiltinFunctionType, types.MethodType, types.FunctionType)): self.process = fn else: # For cases such as set / list where fn is callable but not a function self.process = lambda element: fn(element) super().__init__() def display_data(self): # If the callable has a name, then it's likely a function, and # we show its name. # Otherwise, it might be an instance of a callable class. We # show its class. display_data_value = ( self._fn.__name__ if hasattr(self._fn, '__name__') else self._fn.__class__) return { 'fn': DisplayDataItem(display_data_value, label='Transform Function') } def __repr__(self): return 'CallableWrapperDoFn(%s)' % self._fn def default_type_hints(self): fn_type_hints = typehints.decorators.IOTypeHints.from_callable(self._fn) type_hints = get_type_hints(self._fn).with_defaults(fn_type_hints) # The fn's output type should be iterable. Strip off the outer # container type due to the 'flatten' portion of FlatMap/ParDo. try: type_hints = type_hints.strip_iterable() except ValueError as e: raise TypeCheckError( 'Return value not iterable: %s: %s' % (self.display_data()['fn'].value, e)) type_hints = type_hints.extract_tagged_outputs() return type_hints def infer_output_type(self, input_type): return trivial_inference.element_type( _strip_output_annotations( trivial_inference.infer_return_type(self._fn, [input_type]))) def _process_argspec_fn(self): return getattr(self._fn, '_argspec_fn', self._fn) def _inspect_process(self): if self._fullargspec: return self._fullargspec else: return get_function_args_defaults(self._process_argspec_fn()) class CombineFn(WithTypeHints, HasDisplayData, urns.RunnerApiFn): """A function object used by a Combine transform with custom processing. A CombineFn specifies how multiple values in all or part of a PCollection can be merged into a single value---essentially providing the same kind of information as the arguments to the Python "reduce" builtin (except for the input argument, which is an instance of CombineFnProcessContext). The combining process proceeds as follows: 1. Input values are partitioned into one or more batches. 2. For each batch, the setup method is invoked. 3. For each batch, the create_accumulator method is invoked to create a fresh initial "accumulator" value representing the combination of zero values. 4. For each input value in the batch, the add_input method is invoked to combine more values with the accumulator for that batch. 5. The merge_accumulators method is invoked to combine accumulators from separate batches into a single combined output accumulator value, once all of the accumulators have had all the input value in their batches added to them. This operation is invoked repeatedly, until there is only one accumulator value left. 6. The extract_output operation is invoked on the final accumulator to get the output value. 7. The teardown method is invoked. Note: If this **CombineFn** is used with a transform that has defaults, **apply** will be called with an empty list at expansion time to get the default value. """ def default_label(self): return self.__class__.__name__ def setup(self, *args, **kwargs): """Called to prepare an instance for combining. This method can be useful if there is some state that needs to be loaded before executing any of the other methods. The resources can then be disposed of in ``CombineFn.teardown``. If you are using Dataflow, you need to enable Dataflow Runner V2 before using this feature. Args: *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ pass def create_accumulator(self, *args, **kwargs): """Return a fresh, empty accumulator for the combine operation. Args: *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ raise NotImplementedError(str(self)) def add_input(self, mutable_accumulator, element, *args, **kwargs): """Return result of folding element into accumulator. CombineFn implementors must override add_input. Args: mutable_accumulator: the current accumulator, may be modified and returned for efficiency element: the element to add, should not be mutated *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ raise NotImplementedError(str(self)) def add_inputs(self, mutable_accumulator, elements, *args, **kwargs): """Returns the result of folding each element in elements into accumulator. This is provided in case the implementation affords more efficient bulk addition of elements. The default implementation simply loops over the inputs invoking add_input for each one. Args: mutable_accumulator: the current accumulator, may be modified and returned for efficiency elements: the elements to add, should not be mutated *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ for element in elements: mutable_accumulator =\ self.add_input(mutable_accumulator, element, *args, **kwargs) return mutable_accumulator def merge_accumulators(self, accumulators, *args, **kwargs): """Returns the result of merging several accumulators to a single accumulator value. Args: accumulators: the accumulators to merge. Only the first accumulator may be modified and returned for efficiency; the other accumulators should not be mutated, because they may be shared with other code and mutating them could lead to incorrect results or data corruption. *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ raise NotImplementedError(str(self)) def compact(self, accumulator, *args, **kwargs): """Optionally returns a more compact representation of the accumulator. This is called before an accumulator is sent across the wire, and can be useful in cases where values are buffered or otherwise lazily kept unprocessed when added to the accumulator. Should return an equivalent, though possibly modified, accumulator. By default returns the accumulator unmodified. Args: accumulator: the current accumulator *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ return accumulator def extract_output(self, accumulator, *args, **kwargs): """Return result of converting accumulator into the output value. Args: accumulator: the final accumulator value computed by this CombineFn for the entire input key or PCollection. Can be modified for efficiency. *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ raise NotImplementedError(str(self)) def teardown(self, *args, **kwargs): """Called to clean up an instance before it is discarded. If you are using Dataflow, you need to enable Dataflow Runner V2 before using this feature. Args: *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ pass def apply(self, elements, *args, **kwargs): """Returns result of applying this CombineFn to the input values. Args: elements: the set of values to combine. *args: Additional arguments and side inputs. **kwargs: Additional arguments and side inputs. """ return self.extract_output( self.add_inputs( self.create_accumulator(*args, **kwargs), elements, *args, **kwargs), *args, **kwargs) def for_input_type(self, input_type): """Returns a specialized implementation of self, if it exists. Otherwise, returns self. Args: input_type: the type of input elements. """ return self @staticmethod def from_callable(fn): return CallableWrapperCombineFn(fn) @staticmethod def maybe_from_callable(fn, has_side_inputs=True): # type: (typing.Union[CombineFn, typing.Callable], bool) -> CombineFn if isinstance(fn, CombineFn): return fn elif callable(fn) and not has_side_inputs: return NoSideInputsCallableWrapperCombineFn(fn) elif callable(fn): return CallableWrapperCombineFn(fn) else: raise TypeError('Expected a CombineFn or callable, got %r' % fn) def get_accumulator_coder(self): return coders.registry.get_coder(object) urns.RunnerApiFn.register_pickle_urn(python_urns.PICKLED_COMBINE_FN) class _ReiterableChain(object): """Like itertools.chain, but allowing re-iteration.""" def __init__(self, iterables): self.iterables = iterables def __iter__(self): for iterable in self.iterables: for item in iterable: yield item def __bool__(self): for iterable in self.iterables: for _ in iterable: return True return False class CallableWrapperCombineFn(CombineFn): """For internal use only; no backwards-compatibility guarantees. A CombineFn (function) object wrapping a callable object. The purpose of this class is to conveniently wrap simple functions and use them in Combine transforms. """ _DEFAULT_BUFFER_SIZE = 10 def __init__(self, fn, buffer_size=_DEFAULT_BUFFER_SIZE): """Initializes a CallableFn object wrapping a callable. Args: fn: A callable object that reduces elements of an iterable to a single value (like the builtins sum and max). This callable must be capable of receiving the kind of values it generates as output in its input, and for best results, its operation must be commutative and associative. Raises: TypeError: if fn parameter is not a callable type. """ if not callable(fn): raise TypeError('Expected a callable object instead of: %r' % fn) super().__init__() self._fn = fn self._buffer_size = buffer_size def display_data(self): return {'fn_dd': self._fn} def __repr__(self): return "%s(%s)" % (self.__class__.__name__, self._fn) def create_accumulator(self, *args, **kwargs): return [] def add_input(self, accumulator, element, *args, **kwargs): accumulator.append(element) if len(accumulator) > self._buffer_size: accumulator = [self._fn(accumulator, *args, **kwargs)] return accumulator def add_inputs(self, accumulator, elements, *args, **kwargs): accumulator.extend(elements) if len(accumulator) > self._buffer_size: accumulator = [self._fn(accumulator, *args, **kwargs)] return accumulator def merge_accumulators(self, accumulators, *args, **kwargs): return [self._fn(_ReiterableChain(accumulators), *args, **kwargs)] def compact(self, accumulator, *args, **kwargs): if len(accumulator) <= 1: return accumulator else: return [self._fn(accumulator, *args, **kwargs)] def extract_output(self, accumulator, *args, **kwargs): return self._fn(accumulator, *args, **kwargs) def default_type_hints(self): fn_type_hints = typehints.decorators.IOTypeHints.from_callable(self._fn) type_hints = get_type_hints(self._fn).with_defaults(fn_type_hints) if type_hints.input_types is None: return type_hints else: # fn(Iterable[V]) -> V becomes CombineFn(V) -> V input_args, input_kwargs = type_hints.input_types if not input_args: if len(input_kwargs) == 1: input_args, input_kwargs = tuple(input_kwargs.values()), {} else: raise TypeError('Combiner input type must be specified positionally.') if not is_consistent_with(input_args[0], typehints.Iterable[typehints.Any]): raise TypeCheckError( 'All functions for a Combine PTransform must accept a ' 'single argument compatible with: Iterable[Any]. ' 'Instead a function with input type: %s was received.' % input_args[0]) input_args = (element_type(input_args[0]), ) + input_args[1:] # TODO(robertwb): Assert output type is consistent with input type? return type_hints.with_input_types(*input_args, **input_kwargs) def infer_output_type(self, input_type): return _strip_output_annotations( trivial_inference.infer_return_type(self._fn, [input_type])) def for_input_type(self, input_type): # Avoid circular imports. from apache_beam.transforms import cy_combiners if self._fn is any: return cy_combiners.AnyCombineFn() elif self._fn is all: return cy_combiners.AllCombineFn() else: known_types = { (sum, int): cy_combiners.SumInt64Fn(), (min, int): cy_combiners.MinInt64Fn(), (max, int): cy_combiners.MaxInt64Fn(), (sum, float): cy_combiners.SumFloatFn(), (min, float): cy_combiners.MinFloatFn(), (max, float): cy_combiners.MaxFloatFn(), } return known_types.get((self._fn, input_type), self) class NoSideInputsCallableWrapperCombineFn(CallableWrapperCombineFn): """For internal use only; no backwards-compatibility guarantees. A CombineFn (function) object wrapping a callable object with no side inputs. This is identical to its parent, but avoids accepting and passing *args and **kwargs for efficiency as they are known to be empty. """ def create_accumulator(self): return [] def add_input(self, accumulator, element): accumulator.append(element) if len(accumulator) > self._buffer_size: accumulator = [self._fn(accumulator)] return accumulator def add_inputs(self, accumulator, elements): accumulator.extend(elements) if len(accumulator) > self._buffer_size: accumulator = [self._fn(accumulator)] return accumulator def merge_accumulators(self, accumulators): return [self._fn(_ReiterableChain(accumulators))] def compact(self, accumulator): if len(accumulator) <= 1: return accumulator else: return [self._fn(accumulator)] def extract_output(self, accumulator): return self._fn(accumulator) class PartitionFn(WithTypeHints): """A function object used by a Partition transform. A PartitionFn specifies how individual values in a PCollection will be placed into separate partitions, indexed by an integer. """ def default_label(self): return self.__class__.__name__ def partition_for(self, element, num_partitions, *args, **kwargs): # type: ignore[empty-body] # type: (T, int, *typing.Any, **typing.Any) -> int """Specify which partition will receive this element. Args: element: An element of the input PCollection. num_partitions: Number of partitions, i.e., output PCollections. *args: optional parameters and side inputs. **kwargs: optional parameters and side inputs. Returns: An integer in [0, num_partitions). """ pass class CallableWrapperPartitionFn(PartitionFn): """For internal use only; no backwards-compatibility guarantees. A PartitionFn object wrapping a callable object. Instances of this class wrap simple functions for use in Partition operations. """ def __init__(self, fn): """Initializes a PartitionFn object wrapping a callable. Args: fn: A callable object, which should accept the following arguments: element - element to assign to a partition. num_partitions - number of output partitions. and may accept additional arguments and side inputs. Raises: TypeError: if fn is not a callable type. """ if not callable(fn): raise TypeError('Expected a callable object instead of: %r' % fn) self._fn = fn def partition_for(self, element, num_partitions, *args, **kwargs): # type: (T, int, *typing.Any, **typing.Any) -> int return self._fn(element, num_partitions, *args, **kwargs) def _get_function_body_without_inners(func): source_lines = inspect.getsourcelines(func)[0] source_lines = dropwhile(lambda x: x.strip().startswith("@"), source_lines) first_def_line = next(source_lines).strip() if first_def_line.startswith("def "): last_def_line_without_comment = first_def_line.split("#")[0] \ .split("\"\"\"")[0] while not last_def_line_without_comment.strip().endswith(":"): last_def_line_without_comment = next(source_lines).split("#")[0] \ .split("\"\"\"")[0] first_line = next(source_lines) indentation = len(first_line) - len(first_line.lstrip()) final_lines = [first_line[indentation:]] skip_inner_def = False if first_line[indentation:].startswith("def "): skip_inner_def = True for line in source_lines: line_indentation = len(line) - len(line.lstrip()) if line[indentation:].startswith("def "): skip_inner_def = True continue if skip_inner_def and line_indentation == indentation: skip_inner_def = False if skip_inner_def and line_indentation > indentation: continue final_lines.append(line[indentation:]) return "".join(final_lines) else: return first_def_line.rsplit(":")[-1].strip() def _check_fn_use_yield_and_return(fn): if isinstance(fn, types.BuiltinFunctionType): return False try: source_code = _get_function_body_without_inners(fn) has_yield = False has_return = False has_return_none = False for line in source_code.split("\n"): lstripped_line = line.lstrip() if lstripped_line.startswith("yield ") or lstripped_line.startswith( "yield("): has_yield = True elif lstripped_line.rstrip() == "return": # Return is likely used to exit the function - ok to use with 'yield'. pass elif lstripped_line.startswith("return ") or lstripped_line.startswith( "return("): if lstripped_line.rstrip() == "return None" or lstripped_line.rstrip( ) == "return(None)": has_return_none = True has_return = True if has_yield and has_return: return True if has_return_none: _LOGGER.warning( "Process method returned None (element won't be emitted): %s." " Check if intended.", fn.__self__.__class__) return False except Exception as e: _LOGGER.debug(str(e)) return False class ParDo(PTransformWithSideInputs): """A :class:`ParDo` transform. Processes an input :class:`~apache_beam.pvalue.PCollection` by applying a :class:`DoFn` to each element and returning the accumulated results into an output :class:`~apache_beam.pvalue.PCollection`. The type of the elements is not fixed as long as the :class:`DoFn` can deal with it. In reality the type is restrained to some extent because the elements sometimes must be persisted to external storage. See the :meth:`.expand()` method comments for a detailed description of all possible arguments. Note that the :class:`DoFn` must return an iterable for each element of the input :class:`~apache_beam.pvalue.PCollection`. An easy way to do this is to use the ``yield`` keyword in the process method. Args: pcoll (~apache_beam.pvalue.PCollection): a :class:`~apache_beam.pvalue.PCollection` to be processed. fn (`typing.Union[DoFn, typing.Callable]`): a :class:`DoFn` object to be applied to each element of **pcoll** argument, or a Callable. *args: positional arguments passed to the :class:`DoFn` object. **kwargs: keyword arguments passed to the :class:`DoFn` object. Note that the positional and keyword arguments will be processed in order to detect :class:`~apache_beam.pvalue.PCollection` s that will be computed as side inputs to the transform. During pipeline execution whenever the :class:`DoFn` object gets executed (its :meth:`DoFn.process()` method gets called) the :class:`~apache_beam.pvalue.PCollection` arguments will be replaced by values from the :class:`~apache_beam.pvalue.PCollection` in the exact positions where they appear in the argument lists. """ def __init__(self, fn, *args, **kwargs): super().__init__(fn, *args, **kwargs) # TODO(robertwb): Change all uses of the dofn attribute to use fn instead. self.dofn = self.fn self.output_tags = set() # type: typing.Set[str] if not isinstance(self.fn, DoFn): raise TypeError('ParDo must be called with a DoFn instance.') # DoFn.process cannot allow both return and yield if _check_fn_use_yield_and_return(self.fn.process): _LOGGER.warning( 'Using yield and return in the process method ' 'of %s can lead to unexpected behavior, see:' 'https://github.com/apache/beam/issues/22969.', self.fn.__class__) # Validate the DoFn by creating a DoFnSignature from apache_beam.runners.common import DoFnSignature self._signature = DoFnSignature(self.fn) def with_exception_handling( self, main_tag='good', dead_letter_tag='bad', *, exc_class=Exception, partial=False, use_subprocess=False, threshold=1, threshold_windowing=None, timeout=None, error_handler=None, on_failure_callback: typing.Optional[typing.Callable[ [Exception, typing.Any], None]] = None, allow_unsafe_userstate_in_process=False): """Automatically provides a dead letter output for saving bad inputs. This can allow a pipeline to continue successfully rather than fail or continuously throw errors on retry when bad elements are encountered. This returns a tagged output with two PCollections, the first being the results of successfully processing the input PCollection, and the second being the set of bad records (those which threw exceptions during processing) along with information about the errors raised. For example, one would write:: good, bad = inputs | Map(maybe_erroring_fn).with_exception_handling() and `good` will be a PCollection of mapped records and `bad` will contain tuples of the form `(input, error_string`) for each input that raised an exception. Args: main_tag: tag to be used for the main (good) output of the DoFn, useful to avoid possible conflicts if this DoFn already produces multiple outputs. Optional, defaults to 'good'. dead_letter_tag: tag to be used for the bad inputs, useful to avoid possible conflicts if this DoFn already produces multiple outputs. Optional, defaults to 'bad'. exc_class: An exception class, or tuple of exception classes, to catch. Optional, defaults to 'Exception'. partial: Whether to emit outputs for an element as they're produced (which could result in partial outputs for a ParDo or FlatMap that throws an error part way through execution) or buffer all outputs until successful processing of the entire element. Optional, defaults to False. use_subprocess: Whether to execute the DoFn logic in a subprocess. This allows one to recover from errors that can crash the calling process (e.g. from an underlying C/C++ library causing a segfault), but is slower as elements and results must cross a process boundary. Note that this starts up a long-running process that is used to handle all the elements (until hard failure, which should be rare) rather than a new process per element, so the overhead should be minimal (and can be amortized if there's any per-process or per-bundle initialization that needs to be done). Optional, defaults to False. threshold: An upper bound on the ratio of inputs that can be bad before aborting the entire pipeline. Optional, defaults to 1.0 (meaning up to 100% of inputs can be bad and the pipeline will still succeed). threshold_windowing: Event-time windowing to use for threshold. Optional, defaults to the windowing of the input. timeout: If the element has not finished processing in timeout seconds, raise a TimeoutError. Defaults to None, meaning no time limit. error_handler: An ErrorHandler that should be used to consume the bad records, rather than returning the good and bad records as a tuple. on_failure_callback: If an element fails or times out, on_failure_callback will be invoked. It will receive the exception and the element being processed in as args. In case of a timeout, the exception will be of type `TimeoutError`. Be careful with this callback - if you set a timeout, it will not apply to the callback, and if the callback fails it will not be retried. allow_unsafe_userstate_in_process: If False, user state will not be permitted in the DoFn's process method. This is disabled by default because user state is potentially unsafe with exception handling since it can be successfully stored or cleared even if the associated element fails and is routed to a dead letter queue. Semantics around state in this kind of failure scenario are not well defined and are subject to change. """ args, kwargs = self.raw_side_inputs return self.label >> _ExceptionHandlingWrapper( self.fn, args, kwargs, main_tag, dead_letter_tag, exc_class, partial, use_subprocess, threshold, threshold_windowing, timeout, error_handler, on_failure_callback, allow_unsafe_userstate_in_process, self.get_resource_hints(), self.get_type_hints()) def with_error_handler(self, error_handler, **exception_handling_kwargs): """An alias for `with_exception_handling(error_handler=error_handler, ...)` This is provided to fit the general ErrorHandler conventions. """ if error_handler is None: return self else: return self.with_exception_handling( error_handler=error_handler, **exception_handling_kwargs) def default_type_hints(self): return self.fn.get_type_hints() def infer_output_type(self, input_type): return self.fn.infer_output_type(input_type) def infer_batch_converters(self, input_element_type): # TODO: Test this code (in batch_dofn_test) if self.fn._process_batch_defined: input_batch_type = self.fn._get_input_batch_type_normalized( input_element_type) if input_batch_type is None: raise TypeError( "process_batch method on {self.fn!r} does not have " "an input type annoation") try: # Generate a batch converter to convert between the input type and the # (batch) input type of process_batch self.fn.input_batch_converter = BatchConverter.from_typehints( element_type=input_element_type, batch_type=input_batch_type) except TypeError as e: raise TypeError( "Failed to find a BatchConverter for the input types of DoFn " f"{self.fn!r} (element_type={input_element_type!r}, " f"batch_type={input_batch_type!r}).") from e else: self.fn.input_batch_converter = None if self.fn._can_yield_batches: output_batch_type = self.fn._get_output_batch_type_normalized( input_element_type) if output_batch_type is None: # TODO: Mention process method in this error raise TypeError( f"process_batch method on {self.fn!r} does not have " "a return type annoation") # Generate a batch converter to convert between the output type and the # (batch) output type of process_batch output_element_type = self.infer_output_type(input_element_type) try: self.fn.output_batch_converter = BatchConverter.from_typehints( element_type=output_element_type, batch_type=output_batch_type) except TypeError as e: raise TypeError( "Failed to find a BatchConverter for the *output* types of DoFn " f"{self.fn!r} (element_type={output_element_type!r}, " f"batch_type={output_batch_type!r}). Maybe you need to override " "DoFn.infer_output_type to set the output element type?") from e else: self.fn.output_batch_converter = None def make_fn(self, fn, has_side_inputs): if isinstance(fn, DoFn): return fn return CallableWrapperDoFn(fn) def _process_argspec_fn(self): return self.fn._process_argspec_fn() def display_data(self): return { 'fn': DisplayDataItem(self.fn.__class__, label='Transform Function'), 'fn_dd': self.fn } def expand(self, pcoll): # In the case of a stateful DoFn, warn if the key coder is not # deterministic. if self._signature.is_stateful_dofn(): kv_type_hint = pcoll.element_type if kv_type_hint and kv_type_hint != typehints.Any: coder = coders.registry.get_coder(kv_type_hint) if not coder.is_kv_coder(): raise ValueError( 'Input elements to the transform %s with stateful DoFn must be ' 'key-value pairs.' % self) key_coder = coder.key_coder() else: key_coder = coders.registry.get_coder(typehints.Any) if not key_coder.is_deterministic(): _LOGGER.warning( 'Key coder %s for transform %s with stateful DoFn may not ' 'be deterministic. This may cause incorrect behavior for complex ' 'key types. Consider adding an input type hint for this transform.', key_coder, self) if self._signature.is_unbounded_per_element(): is_bounded = False else: is_bounded = pcoll.is_bounded self.infer_batch_converters(pcoll.element_type) return pvalue.PCollection.from_(pcoll, is_bounded=is_bounded) def with_outputs(self, *tags, main=None, allow_unknown_tags=None): """Returns a tagged tuple allowing access to the outputs of a :class:`ParDo`. The resulting object supports access to the :class:`~apache_beam.pvalue.PCollection` associated with a tag (e.g. ``o.tag``, ``o[tag]``) and iterating over the available tags (e.g. ``for tag in o: ...``). Args: *tags: if non-empty, list of valid tags. If a list of valid tags is given, it will be an error to use an undeclared tag later in the pipeline. **main_kw: dictionary empty or with one key ``'main'`` defining the tag to be used for the main output (which will not have a tag associated with it). Returns: ~apache_beam.pvalue.DoOutputsTuple: An object of type :class:`~apache_beam.pvalue.DoOutputsTuple` that bundles together all the outputs of a :class:`ParDo` transform and allows accessing the individual :class:`~apache_beam.pvalue.PCollection` s for each output using an ``object.tag`` syntax. Raises: TypeError: if the **self** object is not a :class:`~apache_beam.pvalue.PCollection` that is the result of a :class:`ParDo` transform. ValueError: if **main_kw** contains any key other than ``'main'``. """ if main in tags: raise ValueError( 'Main output tag %r must be different from side output tags %r.' % (main, tags)) type_hints = self.get_type_hints() declared_tags = set(type_hints.tagged_output_types().keys()) requested_tags = set(tags) unknown = requested_tags - declared_tags if unknown and declared_tags: # Only warn if type hints exist logging.warning( "Tags %s requested in with_outputs() but not declared " "in type hints. Declared tags: %s", unknown, declared_tags) return _MultiParDo(self, tags, main, allow_unknown_tags) def _do_fn_info(self): return DoFnInfo.create(self.fn, self.args, self.kwargs) def _get_key_and_window_coder(self, named_inputs): if named_inputs is None or not self._signature.is_stateful_dofn(): return None, None main_input = list(set(named_inputs.keys()) - set(self.side_inputs))[0] input_pcoll = named_inputs[main_input] kv_type_hint = input_pcoll.element_type if kv_type_hint and kv_type_hint != typehints.Any: coder = coders.registry.get_coder(kv_type_hint) if not coder.is_kv_coder(): raise ValueError( 'Input elements to the transform %s with stateful DoFn must be ' 'key-value pairs.' % self) key_coder = coder.key_coder() else: key_coder = coders.registry.get_coder(typehints.Any) window_coder = input_pcoll.windowing.windowfn.get_window_coder() return key_coder, window_coder # typing: PTransform base class does not accept extra_kwargs def to_runner_api_parameter(self, context, **extra_kwargs): # type: ignore[override] # type: (PipelineContext, **typing.Any) -> typing.Tuple[str, message.Message] assert isinstance(self, ParDo), \ "expected instance of ParDo, but got %s" % self.__class__ state_specs, timer_specs = userstate.get_dofn_specs(self.fn) if state_specs or timer_specs: context.add_requirement( common_urns.requirements.REQUIRES_STATEFUL_PROCESSING.urn) from apache_beam.runners.common import DoFnSignature sig = DoFnSignature(self.fn) is_splittable = sig.is_splittable_dofn() if is_splittable: restriction_coder = sig.get_restriction_coder() # restriction_coder will never be None when is_splittable is True assert restriction_coder is not None restriction_coder_id = context.coders.get_id( restriction_coder) # type: typing.Optional[str] context.add_requirement( common_urns.requirements.REQUIRES_SPLITTABLE_DOFN.urn) else: restriction_coder_id = None has_bundle_finalization = sig.has_bundle_finalization() if has_bundle_finalization: context.add_requirement( common_urns.requirements.REQUIRES_BUNDLE_FINALIZATION.urn) # Get key_coder and window_coder for main_input. key_coder, window_coder = self._get_key_and_window_coder( extra_kwargs.get('named_inputs', None)) return ( common_urns.primitives.PAR_DO.urn, beam_runner_api_pb2.ParDoPayload( do_fn=self._do_fn_info().to_runner_api(context), requests_finalization=has_bundle_finalization, restriction_coder_id=restriction_coder_id, state_specs={ spec.name: spec.to_runner_api(context) for spec in state_specs }, timer_family_specs={ spec.name: spec.to_runner_api(context, key_coder, window_coder) for spec in timer_specs }, # It'd be nice to name these according to their actual # names/positions in the orignal argument list, but such a # transformation is currently irreversible given how # remove_objects_from_args and insert_values_in_args # are currently implemented. side_inputs={ (SIDE_INPUT_PREFIX + '%s') % ix: si.to_runner_api(context) for ix, si in enumerate(self.side_inputs) })) @staticmethod @PTransform.register_urn( common_urns.primitives.PAR_DO.urn, beam_runner_api_pb2.ParDoPayload) def from_runner_api_parameter(unused_ptransform, pardo_payload, context): fn, args, kwargs, si_tags_and_types, windowing = pickler.loads( DoFnInfo.from_runner_api( pardo_payload.do_fn, context).serialized_dofn_data()) if si_tags_and_types: raise NotImplementedError('explicit side input data') elif windowing: raise NotImplementedError('explicit windowing') result = ParDo(fn, *args, **kwargs) # This is an ordered list stored as a dict (see the comments in # to_runner_api_parameter above). indexed_side_inputs = [( get_sideinput_index(tag), pvalue.AsSideInput.from_runner_api(si, context)) for tag, si in pardo_payload.side_inputs.items()] result.side_inputs = [si for _, si in sorted(indexed_side_inputs)] return result def runner_api_requires_keyed_input(self): return userstate.is_stateful_dofn(self.fn) def get_restriction_coder(self): """Returns `restriction coder if `DoFn` of this `ParDo` is a SDF. Returns `None` otherwise. """ from apache_beam.runners.common import DoFnSignature return DoFnSignature(self.fn).get_restriction_coder() def _add_type_constraint_from_consumer(self, full_label, input_type_hints): if not hasattr(self.fn, '_runtime_output_constraints'): self.fn._runtime_output_constraints = {} self.fn._runtime_output_constraints[full_label] = input_type_hints class _MultiParDo(PTransform): def __init__(self, do_transform, tags, main_tag, allow_unknown_tags=None): super().__init__(do_transform.label) self._do_transform = do_transform self._tags = tags self._main_tag = main_tag self._allow_unknown_tags = allow_unknown_tags def expand(self, pcoll): _ = pcoll | self._do_transform return pvalue.DoOutputsTuple( pcoll.pipeline, self._do_transform, self._tags, self._main_tag, self._allow_unknown_tags) def with_exception_handling(self, main_tag=None, **kwargs): if main_tag is None: main_tag = self._main_tag or 'good' named = self._do_transform.with_exception_handling( main_tag=main_tag, **kwargs) # named is _NamedPTransform wrapping _ExceptionHandlingWrapper named.transform._extra_tags = self._tags return named class DoFnInfo(object): """This class represents the state in the ParDoPayload's function spec, which is the actual DoFn together with some data required for invoking it. """ @staticmethod def register_stateless_dofn(urn): def wrapper(cls): StatelessDoFnInfo.REGISTERED_DOFNS[urn] = cls cls._stateless_dofn_urn = urn return cls return wrapper @classmethod def create(cls, fn, args, kwargs): if hasattr(fn, '_stateless_dofn_urn'): assert not args and not kwargs return StatelessDoFnInfo(fn._stateless_dofn_urn) else: return PickledDoFnInfo(cls._pickled_do_fn_info(fn, args, kwargs)) @staticmethod def from_runner_api(spec, unused_context): if spec.urn == python_urns.PICKLED_DOFN_INFO: return PickledDoFnInfo(spec.payload) elif spec.urn in StatelessDoFnInfo.REGISTERED_DOFNS: return StatelessDoFnInfo(spec.urn) else: raise ValueError('Unexpected DoFn type: %s' % spec.urn) @staticmethod def _pickled_do_fn_info(fn, args, kwargs): # This can be cleaned up once all runners move to portability. return pickler.dumps((fn, args, kwargs, None, None)) def serialized_dofn_data(self): raise NotImplementedError(type(self)) class PickledDoFnInfo(DoFnInfo): def __init__(self, serialized_data): self._serialized_data = serialized_data def serialized_dofn_data(self): return self._serialized_data def to_runner_api(self, unused_context): return beam_runner_api_pb2.FunctionSpec( urn=python_urns.PICKLED_DOFN_INFO, payload=self._serialized_data) class StatelessDoFnInfo(DoFnInfo): REGISTERED_DOFNS = {} # type: typing.Dict[str, typing.Type[DoFn]] def __init__(self, urn): # type: (str) -> None assert urn in self.REGISTERED_DOFNS self._urn = urn def serialized_dofn_data(self): return self._pickled_do_fn_info(self.REGISTERED_DOFNS[self._urn](), (), {}) def to_runner_api(self, unused_context): return beam_runner_api_pb2.FunctionSpec(urn=self._urn) def identity(x): return x def FlatMap(fn=identity, *args, **kwargs): # pylint: disable=invalid-name """:func:`FlatMap` is like :class:`ParDo` except it takes a callable to specify the transformation. The callable must return an iterable for each element of the input :class:`~apache_beam.pvalue.PCollection`. The elements of these iterables will be flattened into the output :class:`~apache_beam.pvalue.PCollection`. If no callable is given, then all elements of the input PCollection must already be iterables themselves and will be flattened into the output PCollection. Args: fn (callable): a callable object. *args: positional arguments passed to the transform callable. **kwargs: keyword arguments passed to the transform callable. Returns: ~apache_beam.pvalue.PCollection: A :class:`~apache_beam.pvalue.PCollection` containing the :func:`FlatMap` outputs. Raises: TypeError: If the **fn** passed as argument is not a callable. Typical error is to pass a :class:`DoFn` instance which is supported only for :class:`ParDo`. """ label = 'FlatMap(%s)' % ptransform.label_from_callable(fn) if not callable(fn): raise TypeError( 'FlatMap can be used only with callable objects. ' 'Received %r instead.' % (fn)) pardo = ParDo(CallableWrapperDoFn(fn), *args, **kwargs) pardo.label = label return pardo def Map(fn, *args, **kwargs): # pylint: disable=invalid-name """:func:`Map` is like :func:`FlatMap` except its callable returns only a single element. Args: fn (callable): a callable object. *args: positional arguments passed to the transform callable. **kwargs: keyword arguments passed to the transform callable. Returns: ~apache_beam.pvalue.PCollection: A :class:`~apache_beam.pvalue.PCollection` containing the :func:`Map` outputs. Raises: TypeError: If the **fn** passed as argument is not a callable. Typical error is to pass a :class:`DoFn` instance which is supported only for :class:`ParDo`. """ if not callable(fn): raise TypeError( 'Map can be used only with callable objects. ' 'Received %r instead.' % (fn)) from apache_beam.transforms.util import fn_takes_side_inputs if fn_takes_side_inputs(fn): wrapper = lambda x, *args, **kwargs: [fn(x, *args, **kwargs)] else: wrapper = lambda x: [fn(x)] label = 'Map(%s)' % ptransform.label_from_callable(fn) # TODO. What about callable classes? if hasattr(fn, '__name__'): wrapper.__name__ = fn.__name__ # Proxy the type-hint information from the original function to this new # wrapped function. type_hints = get_type_hints(fn).with_defaults( typehints.decorators.IOTypeHints.from_callable(fn)) if type_hints.input_types is not None: wrapper = with_input_types( *type_hints.input_types[0], **type_hints.input_types[1])( wrapper) output_hint = type_hints.simple_output_type(label) if output_hint: tagged = { k: typehints.Iterable[v] for k, v in type_hints.tagged_output_types().items() } wrapper = with_output_types( typehints.Iterable[_strip_output_annotations( output_hint, strip_tagged_output=False)], **tagged)( wrapper) # pylint: disable=protected-access wrapper._argspec_fn = fn # pylint: enable=protected-access pardo = FlatMap(wrapper, *args, **kwargs) pardo.label = label return pardo def MapTuple(fn, *args, **kwargs): # pylint: disable=invalid-name r""":func:`MapTuple` is like :func:`Map` but expects tuple inputs and flattens them into multiple input arguments. In other words "SwapKV" >> beam.Map(lambda kv: (kv[1], kv[0])) is equivalent to "SwapKV" >> beam.MapTuple(lambda k, v: (v, k)) This can be useful when processing a PCollection of tuples (e.g. key-value pairs). Args: fn (callable): a callable object. *args: positional arguments passed to the transform callable. **kwargs: keyword arguments passed to the transform callable. Returns: ~apache_beam.pvalue.PCollection: A :class:`~apache_beam.pvalue.PCollection` containing the :func:`MapTuple` outputs. Raises: TypeError: If the **fn** passed as argument is not a callable. Typical error is to pass a :class:`DoFn` instance which is supported only for :class:`ParDo`. """ if not callable(fn): raise TypeError( 'MapTuple can be used only with callable objects. ' 'Received %r instead.' % (fn)) label = 'MapTuple(%s)' % ptransform.label_from_callable(fn) arg_names, defaults = get_function_args_defaults(fn) num_defaults = len(defaults) if num_defaults < len(args) + len(kwargs): raise TypeError('Side inputs must have defaults for MapTuple.') if defaults or args or kwargs: wrapper = lambda x, *args, **kwargs: [fn(*(tuple(x) + args), **kwargs)] else: wrapper = lambda x: [fn(*x)] # Proxy the type-hint information from the original function to this new # wrapped function. type_hints = get_type_hints(fn).with_defaults( typehints.decorators.IOTypeHints.from_callable(fn)) if type_hints.input_types is not None: # TODO(BEAM-14052): ignore input hints, as we do not have enough # information to infer the input type hint of the wrapper function. pass output_hint = type_hints.simple_output_type(label) if output_hint: tagged = { k: typehints.Iterable[v] for k, v in type_hints.tagged_output_types().items() } wrapper = with_output_types( typehints.Iterable[_strip_output_annotations( output_hint, strip_tagged_output=False)], **tagged)( wrapper) # Replace the first (args) component. modified_arg_names = ['tuple_element'] + arg_names[-num_defaults:] modified_argspec = (modified_arg_names, defaults) pardo = ParDo( CallableWrapperDoFn(wrapper, fullargspec=modified_argspec), *args, **kwargs) pardo.label = label return pardo def FlatMapTuple(fn, *args, **kwargs): # pylint: disable=invalid-name r""":func:`FlatMapTuple` is like :func:`FlatMap` but expects tuple inputs and flattens them into multiple input arguments. In other words beam.FlatMap(lambda start_end: range(start_end[0], start_end[1])) is equivalent to beam.FlatMapTuple(lambda start, end: range(start, end)) This can be useful when processing a PCollection of tuples (e.g. key-value pairs). Args: fn (callable): a callable object. *args: positional arguments passed to the transform callable. **kwargs: keyword arguments passed to the transform callable. Returns: ~apache_beam.pvalue.PCollection: A :class:`~apache_beam.pvalue.PCollection` containing the :func:`FlatMapTuple` outputs. Raises: TypeError: If the **fn** passed as argument is not a callable. Typical error is to pass a :class:`DoFn` instance which is supported only for :class:`ParDo`. """ if not callable(fn): raise TypeError( 'FlatMapTuple can be used only with callable objects. ' 'Received %r instead.' % (fn)) label = 'FlatMapTuple(%s)' % ptransform.label_from_callable(fn) arg_names, defaults = get_function_args_defaults(fn) num_defaults = len(defaults) if num_defaults < len(args) + len(kwargs): raise TypeError('Side inputs must have defaults for FlatMapTuple.') if defaults or args or kwargs: wrapper = lambda x, *args, **kwargs: fn(*(tuple(x) + args), **kwargs) else: wrapper = lambda x: fn(*tuple(x)) # Proxy the type-hint information from the original function to this new # wrapped function. type_hints = get_type_hints(fn).with_defaults( typehints.decorators.IOTypeHints.from_callable(fn)) if type_hints.input_types is not None: # TODO(BEAM-14052): ignore input hints, as we do not have enough # information to infer the input type hint of the wrapper function. pass output_hint = type_hints.simple_output_type(label) if output_hint: wrapper = with_output_types( _strip_output_annotations(output_hint, strip_tagged_output=False), **type_hints.tagged_output_types())( wrapper) # Replace the first (args) component. modified_arg_names = ['tuple_element'] + arg_names[-num_defaults:] modified_argspec = (modified_arg_names, defaults) pardo = ParDo( CallableWrapperDoFn(wrapper, fullargspec=modified_argspec), *args, **kwargs) pardo.label = label return pardo # Element format emitted on the dead-letter tag by `with_exception_handling()`: # (exception_class, repr(exception), formatted_traceback_lines). # The first slot is the bare metaclass `type` rather than `Type[BaseException]` # because the runtime value is `type(exn)`, a class object whose only # universally-correct hint is `type`. Beam has no parametric `Type[...]` # constraint that would yield a non-pickle coder anyway, so narrowing here # would over-promise without changing the wire format. ExceptionInfo = typehints.Tuple[type, str, typehints.List[str]] class DeadLetter: """Type hint for a dead-letter element: (original_element, ExceptionInfo). Use as ``DeadLetter[T]`` in ``with_output_types(...)`` etc.""" def __class_getitem__(cls, element_type): return typehints.Tuple[element_type, ExceptionInfo] class _ExceptionHandlingWrapper(ptransform.PTransform): """Implementation of ParDo.with_exception_handling.""" def __init__( self, fn, args, kwargs, main_tag, dead_letter_tag, exc_class, partial, use_subprocess, threshold, threshold_windowing, timeout, error_handler, on_failure_callback, allow_unsafe_userstate_in_process, resource_hints, pardo_type_hints=None): if partial and use_subprocess: raise ValueError('partial and use_subprocess are mutually incompatible.') self._fn = fn self._args = args self._kwargs = kwargs self._main_tag = main_tag self._dead_letter_tag = dead_letter_tag self._exc_class = exc_class self._partial = partial self._use_subprocess = use_subprocess self._threshold = threshold self._threshold_windowing = threshold_windowing self._timeout = timeout self._error_handler = error_handler self._on_failure_callback = on_failure_callback self._allow_unsafe_userstate_in_process = allow_unsafe_userstate_in_process self._resource_hints = resource_hints self._pardo_type_hints = pardo_type_hints self._extra_tags = None def with_outputs(self, *tags, main=None): self._extra_tags = tags if main is not None: self._main_tag = main return self def _build_pardo(self, pcoll): """Build the inner ParDo with the exception-handling wrapper DoFn.""" if self._allow_unsafe_userstate_in_process: if self._use_subprocess or self._timeout: # TODO(https://github.com/apache/beam/issues/35976): Implement this raise Exception( 'allow_unsafe_userstate_in_process is incompatible with ' + 'exception handling done with subprocesses or timeouts. If you ' + 'need this feature, comment in ' + 'https://github.com/apache/beam/issues/35976') if self._use_subprocess: wrapped_fn = _SubprocessDoFn(self._fn, timeout=self._timeout) elif self._timeout: wrapped_fn = _TimeoutDoFn(self._fn, timeout=self._timeout) else: wrapped_fn = self._fn pardo = ParDo( _ExceptionHandlingWrapperDoFn( wrapped_fn, self._dead_letter_tag, self._exc_class, self._partial, self._on_failure_callback, self._allow_unsafe_userstate_in_process, ), *self._args, **self._kwargs, ) pardo.get_resource_hints().update(self._resource_hints) return pardo def _post_process_result(self, pcoll, result): """Apply threshold checking and error handler logic to the result.""" if self._threshold < 1.0: class MaybeWindow(ptransform.PTransform): @staticmethod def expand(pcoll): if self._threshold_windowing: return pcoll | WindowInto(self._threshold_windowing) else: return pcoll # Map(lambda) produces a label formatted like this, but it cannot be # changed without breaking update compat. Here, we pin to the transform # name used in the 2.68 release to avoid breaking changes when the line # number changes. Context: https://github.com/apache/beam/pull/36381 input_count_view = pcoll | 'CountTotal' >> ( MaybeWindow() | "Map()" >> Map(lambda _: 1) | CombineGlobally(sum).as_singleton_view()) bad_count_pcoll = result[self._dead_letter_tag] | 'CountBad' >> ( MaybeWindow() | "Map()" >> Map(lambda _: 1) | CombineGlobally(sum).without_defaults()) def check_threshold(bad, total, threshold, window=DoFn.WindowParam): if bad > total * threshold: raise ValueError( 'The number of failing elements within the window %r ' 'exceeded threshold: %s / %s = %s > %s' % (window, bad, total, bad / total, threshold)) _ = bad_count_pcoll | Map( check_threshold, input_count_view, self._threshold) if self._error_handler: self._error_handler.add_error_pcollection(result[self._dead_letter_tag]) if self._extra_tags is not None: return result return result[self._main_tag] else: return result def expand_2_72_0(self, pcoll): """Pre-2.73.0 behavior: manual element_type override, no with_output_types. """ pardo = self._build_pardo(pcoll) result = pcoll | pardo.with_outputs( self._dead_letter_tag, main=self._main_tag, allow_unknown_tags=True) #TODO(BEAM-18957): Fix when type inference supports tagged outputs. result[self._main_tag].element_type = self._fn.infer_output_type( pcoll.element_type) return self._post_process_result(pcoll, result) def expand(self, pcoll): if pcoll.pipeline.options.is_compat_version_prior_to("2.73.0"): return self.expand_2_72_0(pcoll) pardo = self._build_pardo(pcoll) if (self._pardo_type_hints and self._pardo_type_hints._has_output_types()): main_output_type = self._pardo_type_hints.simple_output_type(self.label) tagged_type_hints = dict(self._pardo_type_hints.tagged_output_types()) else: main_output_type = self._fn.infer_output_type(pcoll.element_type) tagged_type_hints = dict(self._fn.get_type_hints().tagged_output_types()) tagged_type_hints[self._dead_letter_tag] = DeadLetter[pcoll.element_type] pardo = pardo.with_output_types(main_output_type, **tagged_type_hints) all_tags = tuple(set(self._extra_tags or ()) | {self._dead_letter_tag}) result = pcoll | pardo.with_outputs( *all_tags, main=self._main_tag, allow_unknown_tags=True) return self._post_process_result(pcoll, result) class _ExceptionHandlingWrapperDoFn(DoFn): def __init__( self, fn, dead_letter_tag, exc_class, partial, on_failure_callback, allow_unsafe_userstate_in_process): self._fn = fn self._dead_letter_tag = dead_letter_tag self._exc_class = exc_class self._partial = partial self._on_failure_callback = on_failure_callback # Wrap process and expose any top level state params so that process can # handle state and timers. if allow_unsafe_userstate_in_process: @wraps(self._fn.process) def process_wrapper(self, *args, **kwargs): return self.exception_handling_wrapper_do_fn_custom_process( *args, **kwargs) self.process = types.MethodType(process_wrapper, self) else: self.process = self.exception_handling_wrapper_do_fn_custom_process process_sig = inspect.signature(self._fn.process) for name, param in process_sig.parameters.items(): if isinstance(param.default, (DoFn.StateParam, DoFn.TimerParam)): logging.warning( 'State or timer parameter {} detected in process method of ' + '{}. State and timers are unsupported when using ' + 'with_exception_handling and may lead to errors. To enable ' + 'state and timers with limited consistency guarantees, pass ' + 'in the allow_unsafe_userstate_in_process parameters to the ' + 'with_exception_handling method.', name, self.fn) def __getattribute__(self, name): if (name.startswith('__') or name in self.__dict__ or name in _ExceptionHandlingWrapperDoFn.__dict__): return object.__getattribute__(self, name) else: return getattr(self._fn, name) def exception_handling_wrapper_do_fn_custom_process(self, *args, **kwargs): try: result = self._fn.process(*args, **kwargs) if not self._partial: # Don't emit any results until we know there will be no errors. result = list(result) yield from result except self._exc_class as exn: if self._on_failure_callback is not None: try: self._on_failure_callback(exn, args[0]) except Exception as e: logging.warning( 'on_failure_callback failed with error: %s', e, exc_info=True) yield pvalue.TaggedOutput( self._dead_letter_tag, ( args[0], ( type(exn), repr(exn), traceback.format_exception(*sys.exc_info())))) # Idea adapted from https://github.com/tosun-si/asgarde. # TODO(robertwb): Consider how this could fit into the public API. # TODO(robertwb): Generalize to all PValue types. class _PValueWithErrors(object): """This wraps a PCollection such that transforms can be chained in a linear manner while still accumulating any errors.""" def __init__(self, pcoll, exception_handling_args, upstream_errors=()): self._pcoll = pcoll self._exception_handling_args = exception_handling_args self._upstream_errors = upstream_errors @property def pipeline(self): return self._pcoll.pipeline @property def element_type(self): return self._pcoll.element_type @element_type.setter def element_type(self, value): self._pcoll.element_type = value def main_output_tag(self): return self._exception_handling_args.get('main_tag', 'good') def error_output_tag(self): return self._exception_handling_args.get('dead_letter_tag', 'bad') def __or__(self, transform): return self.apply(transform) def apply(self, transform): if hasattr(transform, 'with_exception_handling'): result = self._pcoll | transform.with_exception_handling( **self._exception_handling_args) if result[self.main_output_tag()].element_type == typehints.Any: result[ self.main_output_tag()].element_type = transform.infer_output_type( self._pcoll.element_type) # TODO(BEAM-18957): Add support for tagged type hints. result[self.error_output_tag()].element_type = typehints.Any return _PValueWithErrors( result[self.main_output_tag()], self._exception_handling_args, self._upstream_errors + (result[self.error_output_tag()], )) else: return _PValueWithErrors( self._pcoll | transform, self._exception_handling_args, self._upstream_errors) def accumulated_errors(self): if len(self._upstream_errors) == 1: return self._upstream_errors[0] else: return self._upstream_errors | Flatten() def as_result(self, error_post_processing=None): return { self.main_output_tag(): self._pcoll, self.error_output_tag(): self.accumulated_errors() if error_post_processing is None else self.accumulated_errors() | error_post_processing, } class _MaybePValueWithErrors(object): """This is like _PValueWithErrors, but only wraps values if exception_handling_args is non-trivial. It is useful for handling error-catching and non-error-catching code in a uniform manner. """ def __init__(self, pvalue, exception_handling_args=None): if isinstance(pvalue, _PValueWithErrors): assert exception_handling_args is None self._pvalue = pvalue elif exception_handling_args is None: self._pvalue = pvalue else: self._pvalue = _PValueWithErrors(pvalue, exception_handling_args) @property def pipeline(self): return self._pvalue.pipeline @property def element_type(self): return self._pvalue.element_type @element_type.setter def element_type(self, value): self._pvalue.element_type = value def __or__(self, transform): return self.apply(transform) def apply(self, transform): return _MaybePValueWithErrors(self._pvalue | transform) def as_result(self, error_post_processing=None): if isinstance(self._pvalue, _PValueWithErrors): return self._pvalue.as_result(error_post_processing) else: return self._pvalue class _SubprocessDoFn(DoFn): """Process method run in a subprocess, turning hard crashes into exceptions. """ def __init__(self, fn, timeout=None): self._fn = fn self._serialized_fn = pickler.dumps(fn) self._timeout = timeout def __getattribute__(self, name): if (name.startswith('__') or name in self.__dict__ or name in type(self).__dict__): return object.__getattribute__(self, name) else: return getattr(self._fn, name) def setup(self): self._pool = None def start_bundle(self): # The pool is initialized lazily, including calls to setup and start_bundle. # This allows us to continue processing elements after a crash. pass def process(self, *args, **kwargs): return self._call_remote(self._remote_process, *args, **kwargs) def finish_bundle(self): self._call_remote(self._remote_finish_bundle) def teardown(self): self._call_remote(self._remote_teardown) self._terminate_pool() def _call_remote(self, method, *args, **kwargs): if self._pool is None: self._pool = concurrent.futures.ProcessPoolExecutor(1) self._pool.submit(self._remote_init, self._serialized_fn).result() try: return self._pool.submit(method, *args, **kwargs).result( self._timeout if method == self._remote_process else None) except (concurrent.futures.process.BrokenProcessPool, TimeoutError, concurrent.futures._base.TimeoutError): self._terminate_pool() raise def _terminate_pool(self): """Forcibly terminate the pool, not leaving any live subprocesses.""" pool = self._pool self._pool = None processes = list(pool._processes.values()) pool.shutdown(wait=False) for p in processes: if p.is_alive(): p.kill() time.sleep(1) for p in processes: if p.is_alive(): p.terminate() # These are classmethods to avoid picking the state of self. # They should only be called in an isolated process, so there's no concern # about sharing state or thread safety. @classmethod def _remote_init(cls, serialized_fn): cls._serialized_fn = serialized_fn cls._fn = None cls._started = False @classmethod def _remote_process(cls, *args, **kwargs): if cls._fn is None: cls._fn = pickler.loads(cls._serialized_fn) cls._fn.setup() if not cls._started: cls._fn.start_bundle() cls._started = True result = cls._fn.process(*args, **kwargs) if result: # Don't return generator objects. result = list(result) return result @classmethod def _remote_finish_bundle(cls): if cls._started: cls._started = False if cls._fn.finish_bundle(): # This is because we restart and re-initialize the pool if it crashed. raise RuntimeError( "Returning elements from _SubprocessDoFn.finish_bundle not safe.") @classmethod def _remote_teardown(cls): if cls._fn: cls._fn.teardown() cls._fn = None class _TimeoutDoFn(DoFn): """Process method run in a separate thread allowing timeouts. """ def __init__(self, fn, timeout=None): self._fn = fn self._timeout = timeout self._pool = None def __getattribute__(self, name): if (name.startswith('__') or name in self.__dict__ or name in type(self).__dict__): return object.__getattribute__(self, name) else: return getattr(self._fn, name) def process(self, *args, **kwargs): if self._pool is None: self._pool = concurrent.futures.ThreadPoolExecutor(10) # Import here to avoid circular dependency from apache_beam.runners.worker.statesampler import get_current_tracker from apache_beam.runners.worker.statesampler import set_current_tracker # State sampler/tracker is stored as a thread local variable, and is used # when incrementing counter metrics. dispatching_thread_state_sampler = get_current_tracker() def wrapped_process(): """Makes the dispatching thread local state sampler available to child thread""" set_current_tracker(dispatching_thread_state_sampler) return list(self._fn.process(*args, **kwargs)) # Ensure we iterate over the entire output list in the given amount of time. try: return self._pool.submit(wrapped_process).result(self._timeout) except TimeoutError: self._pool.shutdown(wait=False) self._pool = None raise def teardown(self): try: self._fn.teardown() finally: if self._pool is not None: self._pool.shutdown(wait=False) self._pool = None def Filter(fn, *args, **kwargs): # pylint: disable=invalid-name """:func:`Filter` is a :func:`FlatMap` with its callable filtering out elements. Filter accepts a function that keeps elements that return True, and filters out the remaining elements. Args: fn (``Callable[..., bool]``): a callable object. First argument will be an element. *args: positional arguments passed to the transform callable. **kwargs: keyword arguments passed to the transform callable. Returns: ~apache_beam.pvalue.PCollection: A :class:`~apache_beam.pvalue.PCollection` containing the :func:`Filter` outputs. Raises: TypeError: If the **fn** passed as argument is not a callable. Typical error is to pass a :class:`DoFn` instance which is supported only for :class:`ParDo`. """ if not callable(fn): raise TypeError( 'Filter can be used only with callable objects. ' 'Received %r instead.' % (fn)) wrapper = lambda x, *args, **kwargs: [x] if fn(x, *args, **kwargs) else [] label = 'Filter(%s)' % ptransform.label_from_callable(fn) # TODO: What about callable classes? if hasattr(fn, '__name__'): wrapper.__name__ = fn.__name__ # Get type hints from this instance or the callable. Do not use output type # hints from the callable (which should be bool if set). fn_type_hints = typehints.decorators.IOTypeHints.from_callable(fn) if fn_type_hints is not None: fn_type_hints = fn_type_hints.with_output_types() type_hints = get_type_hints(fn).with_defaults(fn_type_hints) # Proxy the type-hint information from the function being wrapped, setting the # output type to be the same as the input type. def has_simple_input_type(th): return ( th.input_types is not None and len(th.input_types[0]) == 1 and not th.input_types[1]) def simple_input_type(th): return th.input_types[0][0] if has_simple_input_type(th) else None if type_hints.input_types is not None and simple_input_type( type_hints) is not typehints.Any: wrapper = with_input_types( *type_hints.input_types[0], **type_hints.input_types[1])( wrapper) output_hint = type_hints.simple_output_type(label) if (output_hint is None and get_type_hints(wrapper).input_types and get_type_hints(wrapper).input_types[0]): output_hint = get_type_hints(wrapper).input_types[0][0] if output_hint: wrapper = with_output_types( typehints.Iterable[_strip_output_annotations(output_hint)])( wrapper) # pylint: disable=protected-access wrapper._argspec_fn = fn # pylint: enable=protected-access pardo = FlatMap(wrapper, *args, **kwargs) pardo.label = label return pardo def _combine_payload(combine_fn, context): return beam_runner_api_pb2.CombinePayload( combine_fn=combine_fn.to_runner_api(context), accumulator_coder_id=context.coders.get_id( combine_fn.get_accumulator_coder())) class CombineGlobally(PTransform): """A :class:`CombineGlobally` transform. Reduces a :class:`~apache_beam.pvalue.PCollection` to a single value by progressively applying a :class:`CombineFn` to portions of the :class:`~apache_beam.pvalue.PCollection` (and to intermediate values created thereby). See documentation in :class:`CombineFn` for details on the specifics on how :class:`CombineFn` s are applied. Args: pcoll (~apache_beam.pvalue.PCollection): a :class:`~apache_beam.pvalue.PCollection` to be reduced into a single value. fn (callable): a :class:`CombineFn` object that will be called to progressively reduce the :class:`~apache_beam.pvalue.PCollection` into single values, or a callable suitable for wrapping by :class:`~apache_beam.transforms.core.CallableWrapperCombineFn`. *args: positional arguments passed to the :class:`CombineFn` object. **kwargs: keyword arguments passed to the :class:`CombineFn` object. Raises: TypeError: If the output type of the input :class:`~apache_beam.pvalue.PCollection` is not compatible with ``Iterable[A]``. Returns: ~apache_beam.pvalue.PCollection: A single-element :class:`~apache_beam.pvalue.PCollection` containing the main output of the :class:`CombineGlobally` transform. Note that the positional and keyword arguments will be processed in order to detect :class:`~apache_beam.pvalue.PValue` s that will be computed as side inputs to the transform. During pipeline execution whenever the :class:`CombineFn` object gets executed (i.e. any of the :class:`CombineFn` methods get called), the :class:`~apache_beam.pvalue.PValue` arguments will be replaced by their actual value in the exact position where they appear in the argument lists. """ has_defaults = True as_view = False fanout = None # type: typing.Optional[int] def __init__(self, fn, *args, **kwargs): if not (isinstance(fn, CombineFn) or callable(fn)): raise TypeError( 'CombineGlobally can be used only with combineFn objects. ' 'Received %r instead.' % (fn)) super().__init__() self.fn = fn self.args = args self.kwargs = kwargs def display_data(self): return { 'combine_fn': DisplayDataItem( self.fn.__class__, label='Combine Function'), 'combine_fn_dd': self.fn, } def default_label(self): if self.fanout is None: return '%s(%s)' % ( self.__class__.__name__, ptransform.label_from_callable(self.fn)) else: return '%s(%s, fanout=%s)' % ( self.__class__.__name__, ptransform.label_from_callable(self.fn), self.fanout) def _clone(self, **extra_attributes): clone = copy.copy(self) clone.__dict__.update(extra_attributes) return clone def with_fanout(self, fanout): return self._clone(fanout=fanout) def with_defaults(self, has_defaults=True): return self._clone(has_defaults=has_defaults) def without_defaults(self): return self.with_defaults(False) def as_singleton_view(self): return self._clone(as_view=True) def expand(self, pcoll): def add_input_types(transform): type_hints = self.get_type_hints() if type_hints.input_types: return transform.with_input_types(type_hints.input_types[0][0]) return transform combine_fn = CombineFn.maybe_from_callable( self.fn, has_side_inputs=self.args or self.kwargs) combine_per_key = CombinePerKey(combine_fn, *self.args, **self.kwargs) if self.fanout: combine_per_key = combine_per_key.with_hot_key_fanout(self.fanout) combined = ( pcoll | 'KeyWithVoid' >> add_input_types( ParDo(_KeyWithNone()).with_output_types( typehints.KV[None, pcoll.element_type])) | 'CombinePerKey' >> combine_per_key | 'UnKey' >> Map(lambda k_v: k_v[1])) if not self.has_defaults and not self.as_view: return combined elif self.as_view: if self.has_defaults: try: combine_fn.setup(*self.args, **self.kwargs) # This is called in the main program, but cannot be avoided # in the as_view case as it must be available to all windows. default_value = combine_fn.apply([], *self.args, **self.kwargs) finally: combine_fn.teardown(*self.args, **self.kwargs) else: default_value = pvalue.AsSingleton._NO_DEFAULT return pvalue.AsSingleton(combined, default_value=default_value) else: if pcoll.windowing.windowfn != GlobalWindows(): raise ValueError( "Default values are not yet supported in CombineGlobally() if the " "output PCollection is not windowed by GlobalWindows. " "Instead, use CombineGlobally().without_defaults() to output " "an empty PCollection if the input PCollection is empty, " "or CombineGlobally().as_singleton_view() to get the default " "output of the CombineFn if the input PCollection is empty.") # log the error for this ill-defined streaming case now if not pcoll.is_bounded and not pcoll.windowing.is_default(): _LOGGER.error( "When combining elements in unbounded collections with " "the non-default windowing strategy, you must explicitly " "specify how to define the combined result of an empty window. " "Please use CombineGlobally().without_defaults() to output " "an empty PCollection if the input PCollection is empty.") def typed(transform): # TODO(robertwb): We should infer this. if combined.element_type: return transform.with_output_types(combined.element_type) return transform # Capture in closure (avoiding capturing self). args, kwargs = self.args, self.kwargs def inject_default(_, combined): if combined: if len(combined) > 1: _LOGGER.error( "Multiple combined values unexpectedly provided" " for a global combine: %s", combined) assert len(combined) == 1 return combined[0] else: try: combine_fn.setup(*args, **kwargs) default = combine_fn.apply([], *args, **kwargs) finally: combine_fn.teardown(*args, **kwargs) return default return ( pcoll.pipeline | 'DoOnce' >> Create([None]) | 'InjectDefault' >> typed( Map(inject_default, pvalue.AsList(combined)))) @staticmethod @PTransform.register_urn( common_urns.composites.COMBINE_GLOBALLY.urn, beam_runner_api_pb2.CombinePayload) def from_runner_api_parameter(unused_ptransform, combine_payload, context): return CombineGlobally( CombineFn.from_runner_api(combine_payload.combine_fn, context)) @DoFnInfo.register_stateless_dofn(python_urns.KEY_WITH_NONE_DOFN) class _KeyWithNone(DoFn): def process(self, v): yield None, v class CombinePerKey(PTransformWithSideInputs): """A per-key Combine transform. Identifies sets of values associated with the same key in the input PCollection, then applies a CombineFn to condense those sets to single values. See documentation in CombineFn for details on the specifics on how CombineFns are applied. Args: pcoll: input pcollection. fn: instance of CombineFn to apply to all values under the same key in pcoll, or a callable whose signature is ``f(iterable, *args, **kwargs)`` (e.g., sum, max). *args: arguments and side inputs, passed directly to the CombineFn. **kwargs: arguments and side inputs, passed directly to the CombineFn. Returns: A PObject holding the result of the combine operation. """ def __new__(cls, *args, **kwargs): def has_side_inputs(): return ( any(isinstance(arg, pvalue.AsSideInput) for arg in args) or any(isinstance(arg, pvalue.AsSideInput) for arg in kwargs.values())) if has_side_inputs(): # If the CombineFn has deferred side inputs, the python SDK # doesn't implement it. # Use a ParDo-based CombinePerKey instead. from apache_beam.transforms.combiners import LiftedCombinePerKey combine_fn, *args = args return LiftedCombinePerKey(combine_fn, args, kwargs) return super(CombinePerKey, cls).__new__(cls) def with_hot_key_fanout(self, fanout): """A per-key combine operation like self but with two levels of aggregation. If a given key is produced by too many upstream bundles, the final reduction can become a bottleneck despite partial combining being lifted pre-GroupByKey. In these cases it can be helpful to perform intermediate partial aggregations in parallel and then re-group to peform a final (per-key) combine. This is also useful for high-volume keys in streaming where combiners are not generally lifted for latency reasons. Note that a fanout greater than 1 requires the data to be sent through two GroupByKeys, and a high fanout can also result in more shuffle data due to less per-bundle combining. Setting the fanout for a key at 1 or less places values on the "cold key" path that skip the intermediate level of aggregation. Args: fanout: either None, for no fanout, an int, for a constant-degree fanout, or a callable mapping keys to a key-specific degree of fanout. Returns: A per-key combining PTransform with the specified fanout. """ from apache_beam.transforms.combiners import curry_combine_fn if fanout is None: return self else: return _CombinePerKeyWithHotKeyFanout( curry_combine_fn(self.fn, self.args, self.kwargs), fanout) def display_data(self): return { 'combine_fn': DisplayDataItem( self.fn.__class__, label='Combine Function'), 'combine_fn_dd': self.fn } def make_fn(self, fn, has_side_inputs): self._fn_label = ptransform.label_from_callable(fn) return CombineFn.maybe_from_callable(fn, has_side_inputs) def default_label(self): return '%s(%s)' % (self.__class__.__name__, self._fn_label) def _process_argspec_fn(self): return lambda element, *args, **kwargs: None def expand(self, pcoll): # When using gbek, don't allow overriding default implementation gbek_option = (pcoll.pipeline._options.view_as(SetupOptions).gbek) self._using_gbek = (gbek_option is not None and len(gbek_option) > 0) args, kwargs = util.insert_values_in_args( self.args, self.kwargs, self.side_inputs) return pcoll | GroupByKey() | 'Combine' >> CombineValues( self.fn, *args, **kwargs) def default_type_hints(self): result = self.fn.get_type_hints() k = typehints.TypeVariable('K') if result.input_types: args, kwargs = result.input_types args = (typehints.Tuple[k, args[0]], ) + args[1:] result = result.with_input_types(*args, **kwargs) else: result = result.with_input_types(typehints.Tuple[k, typehints.Any]) if result.output_types: main_output_type = result.simple_output_type('') result = result.with_output_types(typehints.Tuple[k, main_output_type]) else: result = result.with_output_types(typehints.Tuple[k, typehints.Any]) return result def to_runner_api_parameter( self, context, # type: PipelineContext ): # type: (...) -> tuple[str, typing.Optional[typing.Union[message.Message, bytes, str]]] if getattr(self, '_using_gbek', False): return super().to_runner_api_parameter(context) if self.args or self.kwargs: from apache_beam.transforms.combiners import curry_combine_fn combine_fn = curry_combine_fn(self.fn, self.args, self.kwargs) else: combine_fn = self.fn return ( common_urns.composites.COMBINE_PER_KEY.urn, _combine_payload(combine_fn, context)) @staticmethod @PTransform.register_urn( common_urns.composites.COMBINE_PER_KEY.urn, beam_runner_api_pb2.CombinePayload) def from_runner_api_parameter(unused_ptransform, combine_payload, context): return CombinePerKey( CombineFn.from_runner_api(combine_payload.combine_fn, context)) def runner_api_requires_keyed_input(self): return True # TODO(robertwb): Rename to CombineGroupedValues? class CombineValues(PTransformWithSideInputs): def make_fn(self, fn, has_side_inputs): return CombineFn.maybe_from_callable(fn, has_side_inputs) def expand(self, pcoll): args, kwargs = util.insert_values_in_args( self.args, self.kwargs, self.side_inputs) input_type = pcoll.element_type key_type = None if input_type is not None: key_type, _ = input_type.tuple_types runtime_type_check = ( pcoll.pipeline._options.view_as(TypeOptions).runtime_type_check) return pcoll | ParDo( CombineValuesDoFn(key_type, self.fn, runtime_type_check), *args, **kwargs) def to_runner_api_parameter(self, context): if self.args or self.kwargs: from apache_beam.transforms.combiners import curry_combine_fn combine_fn = curry_combine_fn(self.fn, self.args, self.kwargs) else: combine_fn = self.fn return ( common_urns.combine_components.COMBINE_GROUPED_VALUES.urn, _combine_payload(combine_fn, context)) @staticmethod @PTransform.register_urn( common_urns.combine_components.COMBINE_GROUPED_VALUES.urn, beam_runner_api_pb2.CombinePayload) def from_runner_api_parameter(unused_ptransform, combine_payload, context): return CombineValues( CombineFn.from_runner_api(combine_payload.combine_fn, context)) class CombineValuesDoFn(DoFn): """DoFn for performing per-key Combine transforms.""" def __init__( self, input_pcoll_type, combinefn, # type: CombineFn runtime_type_check, # type: bool ): super().__init__() self.combinefn = combinefn self.runtime_type_check = runtime_type_check def setup(self): self.combinefn.setup() def process(self, element, *args, **kwargs): # Expected elements input to this DoFn are 2-tuples of the form # (key, iter), with iter an iterable of all the values associated with key # in the input PCollection. if self.runtime_type_check: # Apply the combiner in a single operation rather than artificially # breaking it up so that output type violations manifest as TypeCheck # errors rather than type errors. return [(element[0], self.combinefn.apply(element[1], *args, **kwargs))] # Add the elements into three accumulators (for testing of merge). elements = list(element[1]) accumulators = [] for k in range(3): if len(elements) <= k: break accumulators.append( self.combinefn.add_inputs( self.combinefn.create_accumulator(*args, **kwargs), elements[k::3], *args, **kwargs)) # Merge the accumulators. accumulator = self.combinefn.merge_accumulators( accumulators, *args, **kwargs) # Convert accumulator to the final result. return [( element[0], self.combinefn.extract_output(accumulator, *args, **kwargs))] def teardown(self): self.combinefn.teardown() def default_type_hints(self): hints = self.combinefn.get_type_hints() if hints.input_types: K = typehints.TypeVariable('K') args, kwargs = hints.input_types args = (typehints.Tuple[K, typehints.Iterable[args[0]]], ) + args[1:] hints = hints.with_input_types(*args, **kwargs) else: K = typehints.Any if hints.output_types: main_output_type = hints.simple_output_type('') hints = hints.with_output_types(typehints.Tuple[K, main_output_type]) return hints class _CombinePerKeyWithHotKeyFanout(PTransform): def __init__( self, combine_fn, # type: CombineFn fanout, # type: typing.Union[int, typing.Callable[[typing.Any], int]] ): # type: (...) -> None self._combine_fn = combine_fn self._fanout_fn = ((lambda key: fanout) if isinstance(fanout, int) else fanout) def default_label(self): return '%s(%s, fanout=%s)' % ( self.__class__.__name__, ptransform.label_from_callable(self._combine_fn), ptransform.label_from_callable(self._fanout_fn)) def expand(self, pcoll): from apache_beam.transforms.trigger import AccumulationMode combine_fn = self._combine_fn fanout_fn = self._fanout_fn use_direct_windowing = isinstance(pcoll.windowing.windowfn, SlidingWindows) class SplitHotCold(DoFn): def start_bundle(self): # Spreading a hot key across all possible sub-keys for all bundles # would defeat the goal of not overwhelming downstream reducers # (as well as making less efficient use of PGBK combining tables). # Instead, each bundle independently makes a consistent choice about # which "shard" of a key to send its intermediate results. self._nonce = int(random.getrandbits(31)) def process(self, element): key, value = element fanout = fanout_fn(key) if fanout <= 1: # Boolean indicates this is not an accumulator. yield (key, (False, value)) # cold else: yield pvalue.TaggedOutput('hot', ((self._nonce % fanout, key), value)) class PreCombineFn(CombineFn): def __init__(self): # Deepcopy of the combine_fn to avoid sharing state between lifted # stages when using cloudpickle. try: self._combine_fn_copy = copy.deepcopy(combine_fn) except Exception: self._combine_fn_copy = pickler.roundtrip(combine_fn) self.setup = self._combine_fn_copy.setup self.create_accumulator = self._combine_fn_copy.create_accumulator self.add_input = self._combine_fn_copy.add_input self.merge_accumulators = self._combine_fn_copy.merge_accumulators self.compact = self._combine_fn_copy.compact self.teardown = self._combine_fn_copy.teardown @staticmethod def extract_output(accumulator): # Boolean indicates this is an accumulator. return (True, accumulator) class PostCombineFn(CombineFn): def __init__(self): # Deepcopy of the combine_fn to avoid sharing state between lifted # stages when using cloudpickle. try: self._combine_fn_copy = copy.deepcopy(combine_fn) except Exception: self._combine_fn_copy = pickler.roundtrip(combine_fn) self.setup = self._combine_fn_copy.setup self.create_accumulator = self._combine_fn_copy.create_accumulator self.merge_accumulators = self._combine_fn_copy.merge_accumulators self.compact = self._combine_fn_copy.compact self.extract_output = self._combine_fn_copy.extract_output self.teardown = self._combine_fn_copy.teardown def add_input(self, accumulator, element): is_accumulator, value = element if is_accumulator: return self._combine_fn_copy.merge_accumulators([accumulator, value]) else: return self._combine_fn_copy.add_input(accumulator, value) def StripNonce(nonce_key_value): (_, key), value = nonce_key_value return key, value cold, hot = pcoll | ParDo(SplitHotCold()).with_outputs('hot', main='cold') cold.element_type = typehints.Any # No multi-output type hints. if use_direct_windowing: # For SlidingWindows, swap windowing strategy metadata directly on the # PCollection without re-assigning windows. This mirrors Java's # setWindowingStrategyInternal(). Using WindowInto would call # windowfn.assign() which re-evaluates window assignments from # timestamps — with SlidingWindows, this causes accumulators to leak # into adjacent overlapping windows. if pcoll.windowing.accumulation_mode == AccumulationMode.ACCUMULATING: discarding_windowing = copy.copy(pcoll.windowing) discarding_windowing.accumulation_mode = AccumulationMode.DISCARDING hot._windowing = discarding_windowing precombined_hot = (hot | CombinePerKey(PreCombineFn()) | Map(StripNonce)) precombined_hot._windowing = pcoll.windowing else: precombined_hot = ( hot # Avoid double counting that may happen with stacked accumulating # mode. | 'WindowIntoDiscarding' >> WindowInto( pcoll.windowing, accumulation_mode=AccumulationMode.DISCARDING) | CombinePerKey(PreCombineFn()) | Map(StripNonce) | 'WindowIntoOriginal' >> WindowInto(pcoll.windowing)) return ((cold, precombined_hot) | Flatten() | CombinePerKey(PostCombineFn())) @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types(tuple[K, typing.Iterable[V]]) class GroupByKey(PTransform): """A group by key transform. Processes an input PCollection consisting of key/value pairs represented as a tuple pair. The result is a PCollection where values having a common key are grouped together. For example (a, 1), (b, 2), (a, 3) will result into (a, [1, 3]), (b, [2]). The implementation here is used only when run on the local direct runner. """ def __init__(self, label=None): self._replaced_by_gbek = False self._inside_gbek = False super().__init__(label) class ReifyWindows(DoFn): def process( self, element, window=DoFn.WindowParam, timestamp=DoFn.TimestampParam): try: k, v = element except TypeError: raise TypeCheckError( 'Input to GroupByKey must be a PCollection with ' 'elements compatible with KV[A, B]') return [(k, WindowedValue(v, timestamp, [window]))] def infer_output_type(self, input_type): key_type, value_type = trivial_inference.key_value_types(input_type) return typehints.KV[ key_type, typehints.WindowedValue[value_type]] # type: ignore[misc] def get_windowing(self, inputs): # Switch to the continuation trigger associated with the current trigger. windowing = inputs[0].windowing triggerfn = windowing.triggerfn.get_continuation_trigger() return Windowing( windowfn=windowing.windowfn, triggerfn=triggerfn, accumulation_mode=windowing.accumulation_mode, timestamp_combiner=windowing.timestamp_combiner, allowed_lateness=windowing.allowed_lateness, environment_id=windowing.environment_id) def expand(self, pcoll): replace_with_gbek_secret = ( pcoll.pipeline._options.view_as(SetupOptions).gbek) if replace_with_gbek_secret is not None and not self._inside_gbek: self._replaced_by_gbek = True from apache_beam.transforms.util import GroupByEncryptedKey from apache_beam.transforms.util import Secret secret = Secret.parse_secret_option(replace_with_gbek_secret) return (pcoll | "Group by encrypted key" >> GroupByEncryptedKey(secret)) from apache_beam.transforms.trigger import DataLossReason from apache_beam.transforms.trigger import DefaultTrigger windowing = pcoll.windowing trigger = windowing.triggerfn if not pcoll.is_bounded and isinstance( windowing.windowfn, GlobalWindows) and isinstance(trigger, DefaultTrigger): if pcoll.pipeline.allow_unsafe_triggers: # TODO(BEAM-9487) Change comment for Beam 2.33 _LOGGER.warning( '%s: PCollection passed to GroupByKey is unbounded, has a global ' 'window, and uses a default trigger. This is being allowed ' 'because --allow_unsafe_triggers is set, but it may prevent ' 'data from making it through the pipeline.', self.label) else: raise ValueError( 'GroupByKey cannot be applied to an unbounded ' + 'PCollection with global windowing and a default trigger') unsafe_reason = trigger.may_lose_data(windowing) if unsafe_reason != DataLossReason.NO_POTENTIAL_LOSS: reason_msg = str(unsafe_reason).replace('DataLossReason.', '') if pcoll.pipeline.allow_unsafe_triggers: _LOGGER.warning( '%s: Unsafe trigger `%s` detected (reason: %s). This is ' 'being allowed because --allow_unsafe_triggers is set. This could ' 'lead to missing or incomplete groups.', self.label, trigger, reason_msg) else: msg = '{}: Unsafe trigger: `{}` may lose data. '.format( self.label, trigger) msg += 'Reason: {}. '.format(reason_msg) msg += 'This can be overriden with the --allow_unsafe_triggers flag.' raise ValueError(msg) return pvalue.PCollection.from_(pcoll) def infer_output_type(self, input_type): key_type, value_type = (typehints.typehints.coerce_to_kv_type( input_type).tuple_types) return typehints.KV[key_type, typehints.Iterable[value_type]] def to_runner_api_parameter(self, unused_context): # type: (PipelineContext) -> tuple[str, typing.Optional[typing.Union[message.Message, bytes, str]]] # if we're containing a GroupByEncryptedKey, don't allow runners to # recognize this transform as a GBEK so that it doesn't get replaced. if self._replaced_by_gbek: return super().to_runner_api_parameter(unused_context) return common_urns.primitives.GROUP_BY_KEY.urn, None @staticmethod @PTransform.register_urn(common_urns.primitives.GROUP_BY_KEY.urn, None) def from_runner_api_parameter( unused_ptransform, unused_payload, unused_context): return GroupByKey() def runner_api_requires_keyed_input(self): return True def _expr_to_callable(expr, pos): if isinstance(expr, str): return lambda x: getattr(x, expr) elif callable(expr): return expr else: raise TypeError( 'Field expression %r at %s must be a callable or a string.' % (expr, pos)) class GroupBy(PTransform): """Groups a PCollection by one or more expressions, used to derive the key. `GroupBy(expr)` is roughly equivalent to beam.Map(lambda v: (expr(v), v)) | beam.GroupByKey() but provides several conveniences, e.g. * Several arguments may be provided, as positional or keyword arguments, resulting in a tuple-like key. For example `GroupBy(a=expr1, b=expr2)` groups by a key with attributes `a` and `b` computed by applying `expr1` and `expr2` to each element. * Strings can be used as a shorthand for accessing an attribute, e.g. `GroupBy('some_field')` is equivalent to `GroupBy(lambda v: getattr(v, 'some_field'))`. The GroupBy operation can be made into an aggregating operation by invoking its `aggregate_field` method. """ def __init__( self, *fields, # type: typing.Union[str, typing.Callable] **kwargs # type: typing.Union[str, typing.Callable] ): if len(fields) == 1 and not kwargs: self._force_tuple_keys = False name = fields[0] if isinstance(fields[0], str) else 'key' key_fields = [(name, _expr_to_callable(fields[0], 0))] else: self._force_tuple_keys = True key_fields = [] for ix, field in enumerate(fields): name = field if isinstance(field, str) else 'key%d' % ix key_fields.append((name, _expr_to_callable(field, ix))) for name, expr in kwargs.items(): key_fields.append((name, _expr_to_callable(expr, name))) self._key_fields = key_fields field_names = tuple(name for name, _ in key_fields) self._key_type = lambda *values: _dynamic_named_tuple('Key', field_names)( *values) def aggregate_field( self, field, # type: typing.Union[str, typing.Callable] combine_fn, # type: typing.Union[typing.Callable, CombineFn] dest, # type: str ): """Returns a grouping operation that also aggregates grouped values. Args: field: indicates the field to be aggregated combine_fn: indicates the aggregation function to be used dest: indicates the name that will be used for the aggregate in the output May be called repeatedly to aggregate multiple fields, e.g. GroupBy('key') .aggregate_field('some_attr', sum, 'sum_attr') .aggregate_field(lambda v: ..., MeanCombineFn, 'mean') """ return _GroupAndAggregate(self, ()).aggregate_field(field, combine_fn, dest) def force_tuple_keys(self, value=True): """Forces the keys to always be tuple-like, even if there is only a single expression. """ res = copy.copy(self) res._force_tuple_keys = value return res def _key_func(self): if not self._force_tuple_keys and len(self._key_fields) == 1: return self._key_fields[0][1] else: key_type = self._key_type key_exprs = [expr for _, expr in self._key_fields] return lambda element: key_type(*(expr(element) for expr in key_exprs)) def _key_type_hint(self, input_type): if not self._force_tuple_keys and len(self._key_fields) == 1: expr = self._key_fields[0][1] return trivial_inference.infer_return_type(expr, [input_type]) else: return row_type.RowTypeConstraint.from_fields([ (name, trivial_inference.infer_return_type(expr, [input_type])) for (name, expr) in self._key_fields ]) def default_label(self): return 'GroupBy(%s)' % ', '.join(name for name, _ in self._key_fields) def expand(self, pcoll): input_type = pcoll.element_type or typing.Any # Map(lambda) produces a label formatted like this, but it cannot be # changed without breaking update compat. Here, we pin to the transform # name used in the 2.68 release to avoid breaking changes when the line # number changes. Context: https://github.com/apache/beam/pull/36381 return ( pcoll | "Map()" >> Map(lambda x: (self._key_func()(x), x)).with_output_types( typehints.Tuple[self._key_type_hint(input_type), input_type]) | GroupByKey()) _dynamic_named_tuple_cache = { } # type: typing.Dict[typing.Tuple[str, typing.Tuple[str, ...]], typing.Type[tuple]] def _dynamic_named_tuple(type_name, field_names): # type: (str, typing.Tuple[str, ...]) -> typing.Type[tuple] cache_key = (type_name, field_names) result = _dynamic_named_tuple_cache.get(cache_key) if result is None: import collections result = _dynamic_named_tuple_cache[cache_key] = collections.namedtuple( type_name, field_names) # typing: can't override a method. also, self type is unknown and can't # be cast to tuple result.__reduce__ = lambda self: ( # type: ignore[method-assign] _unpickle_dynamic_named_tuple, (type_name, field_names, tuple(self))) # type: ignore[arg-type] return result def _unpickle_dynamic_named_tuple(type_name, field_names, values): # type: (str, typing.Tuple[str, ...], typing.Iterable[typing.Any]) -> tuple return _dynamic_named_tuple(type_name, field_names)(*values) class _GroupAndAggregate(PTransform): def __init__(self, grouping, aggregations): self._grouping = grouping self._aggregations = aggregations def aggregate_field( self, field, # type: typing.Union[str, typing.Callable] combine_fn, # type: typing.Union[typing.Callable, CombineFn] dest, # type: str ): field = _expr_to_callable(field, 0) return _GroupAndAggregate( self._grouping, list(self._aggregations) + [(field, combine_fn, dest)]) def expand(self, pcoll): from apache_beam.transforms.combiners import TupleCombineFn key_func = self._grouping.force_tuple_keys(True)._key_func() value_exprs = [expr for expr, _, __ in self._aggregations] value_func = lambda element: [expr(element) for expr in value_exprs] result_fields = tuple(name for name, _ in self._grouping._key_fields) + tuple( dest for _, __, dest in self._aggregations) key_type_hint = self._grouping.force_tuple_keys(True)._key_type_hint( pcoll.element_type) # Map(lambda) produces a label formatted like this, but it cannot be # changed without breaking update compat. Here, we pin to the transform # name used in the 2.68 release to avoid breaking changes when the line # number changes. Context: https://github.com/apache/beam/pull/36381 return ( pcoll | "Map()" >> Map(lambda x: (key_func(x), value_func(x))).with_output_types( typehints.Tuple[key_type_hint, typing.Any]) | CombinePerKey( TupleCombineFn( *[combine_fn for _, combine_fn, __ in self._aggregations])) | "MapTuple()" >> MapTuple( lambda key, value: _dynamic_named_tuple('Result', result_fields) (*(key + value)))) class Select(PTransform): """Converts the elements of a PCollection into a schema'd PCollection of Rows. `Select(...)` is roughly equivalent to `Map(lambda x: Row(...))` where each argument (which may be a string or callable) of `ToRow` is applied to `x`. For example, pcoll | beam.Select('a', b=lambda x: foo(x)) is the same as pcoll | 'label' >> beam.Map(lambda x: beam.Row(a=x.a, b=foo(x))) """ def __init__( self, *args, # type: typing.Union[str, typing.Callable] **kwargs # type: typing.Union[str, typing.Callable] ): self._fields = [( expr if isinstance(expr, str) else 'arg%02d' % ix, _expr_to_callable(expr, ix)) for (ix, expr) in enumerate(args) ] + [(name, _expr_to_callable(expr, name)) for (name, expr) in kwargs.items()] self._exception_handling_args = None def with_exception_handling(self, **kwargs): self._exception_handling_args = kwargs return self def default_label(self): return 'ToRows(%s)' % ', '.join(name for name, _ in self._fields) def expand(self, pcoll): # Map(lambda) produces a label formatted like this, but it cannot be # changed without breaking update compat. Here, we pin to the transform # name used in the 2.68 release to avoid breaking changes when the line # number changes. Context: https://github.com/apache/beam/pull/36381 return ( _MaybePValueWithErrors(pcoll, self._exception_handling_args) | "Map()" >> Map( lambda x: pvalue.Row( **{ name: expr(x) for name, expr in self._fields }))).as_result() def infer_output_type(self, input_type): def extract_return_type(expr): expr_hints = get_type_hints(expr) if (expr_hints and expr_hints.has_simple_output_type() and expr_hints.simple_output_type(None) != typehints.Any): return expr_hints.simple_output_type(None) else: return trivial_inference.infer_return_type(expr, [input_type]) return row_type.RowTypeConstraint.from_fields([ (name, extract_return_type(expr)) for (name, expr) in self._fields ]) class Partition(PTransformWithSideInputs): """Split a PCollection into several partitions. Uses the specified PartitionFn to separate an input PCollection into the specified number of sub-PCollections. When apply()d, a Partition() PTransform requires the following: Args: partitionfn: a PartitionFn, or a callable with the signature described in CallableWrapperPartitionFn. n: number of output partitions. The result of this PTransform is a simple list of the output PCollections representing each of n partitions, in order. """ class ApplyPartitionFnFn(DoFn): """A DoFn that applies a PartitionFn.""" def process(self, element, partitionfn, n, *args, **kwargs): partition = partitionfn.partition_for(element, n, *args, **kwargs) import numbers if isinstance(partition, bool) or not isinstance(partition, numbers.Integral): raise ValueError( f"PartitionFn yielded a '{type(partition).__name__}' " "when it should only yield integers") if not 0 <= int(partition) < n: raise ValueError( 'PartitionFn specified out-of-bounds partition index: ' '%d not in [0, %d)' % (partition, n)) # Each input is directed into the output that corresponds to the # selected partition. yield pvalue.TaggedOutput(str(partition), element) def make_fn(self, fn, has_side_inputs): return fn if isinstance(fn, PartitionFn) else CallableWrapperPartitionFn(fn) def expand(self, pcoll): n = int(self.args[0]) args, kwargs = util.insert_values_in_args( self.args, self.kwargs, self.side_inputs) return pcoll | ParDo(self.ApplyPartitionFnFn(), self.fn, *args, ** kwargs).with_outputs(*[str(t) for t in range(n)]) class Windowing(object): def __init__( self, windowfn, # type: WindowFn triggerfn=None, # type: typing.Optional[TriggerFn] accumulation_mode=None, # type: typing.Optional[beam_runner_api_pb2.AccumulationMode.Enum.ValueType] timestamp_combiner=None, # type: typing.Optional[beam_runner_api_pb2.OutputTime.Enum.ValueType] allowed_lateness=0, # type: typing.Union[int, float] environment_id=None, # type: typing.Optional[str] ): """Class representing the window strategy. Args: windowfn: Window assign function. triggerfn: Trigger function. accumulation_mode: a AccumulationMode, controls what to do with data when a trigger fires multiple times. timestamp_combiner: a TimestampCombiner, determines how output timestamps of grouping operations are assigned. allowed_lateness: Maximum delay in seconds after end of window allowed for any late data to be processed without being discarded directly. environment_id: Environment where the current window_fn should be applied in. """ global AccumulationMode, DefaultTrigger # pylint: disable=global-variable-not-assigned # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.transforms.trigger import AccumulationMode from apache_beam.transforms.trigger import DefaultTrigger # pylint: enable=wrong-import-order, wrong-import-position if triggerfn is None: triggerfn = DefaultTrigger() if accumulation_mode is None: if triggerfn == DefaultTrigger(): accumulation_mode = AccumulationMode.DISCARDING else: raise ValueError( 'accumulation_mode must be provided for non-trivial triggers') if not windowfn.get_window_coder().is_deterministic(): raise ValueError( 'window fn (%s) does not have a determanistic coder (%s)' % (windowfn, windowfn.get_window_coder())) self.windowfn = windowfn self.triggerfn = triggerfn self.accumulation_mode = accumulation_mode self.allowed_lateness = Duration.of(allowed_lateness) self.environment_id = environment_id self.timestamp_combiner = ( timestamp_combiner or TimestampCombiner.OUTPUT_AT_EOW) self._is_default = ( self.windowfn == GlobalWindows() and self.triggerfn == DefaultTrigger() and self.accumulation_mode == AccumulationMode.DISCARDING and self.timestamp_combiner == TimestampCombiner.OUTPUT_AT_EOW and self.allowed_lateness == 0) def __repr__(self): return "Windowing(%s, %s, %s, %s, %s)" % ( self.windowfn, self.triggerfn, self.accumulation_mode, self.timestamp_combiner, self.environment_id) def __eq__(self, other): if type(self) == type(other): if self._is_default and other._is_default: return True return ( self.windowfn == other.windowfn and self.triggerfn == other.triggerfn and self.accumulation_mode == other.accumulation_mode and self.timestamp_combiner == other.timestamp_combiner and self.allowed_lateness == other.allowed_lateness and self.environment_id == self.environment_id) return False def __hash__(self): return hash(( self.windowfn, self.triggerfn, self.accumulation_mode, self.allowed_lateness, self.timestamp_combiner, self.environment_id)) def is_default(self): return self._is_default def to_runner_api(self, context): # type: (PipelineContext) -> beam_runner_api_pb2.WindowingStrategy environment_id = self.environment_id or context.default_environment_id() return beam_runner_api_pb2.WindowingStrategy( window_fn=self.windowfn.to_runner_api(context), # TODO(robertwb): Prohibit implicit multi-level merging. merge_status=( beam_runner_api_pb2.MergeStatus.NEEDS_MERGE if self.windowfn.is_merging() else beam_runner_api_pb2.MergeStatus.NON_MERGING), window_coder_id=context.coders.get_id(self.windowfn.get_window_coder()), trigger=self.triggerfn.to_runner_api(context), accumulation_mode=self.accumulation_mode, output_time=self.timestamp_combiner, # TODO(robertwb): Support EMIT_IF_NONEMPTY closing_behavior=beam_runner_api_pb2.ClosingBehavior.EMIT_ALWAYS, on_time_behavior=beam_runner_api_pb2.OnTimeBehavior.FIRE_ALWAYS, allowed_lateness=self.allowed_lateness.micros // 1000, environment_id=environment_id) @staticmethod def from_runner_api(proto, context): # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.transforms.trigger import TriggerFn return Windowing( windowfn=WindowFn.from_runner_api(proto.window_fn, context), triggerfn=TriggerFn.from_runner_api(proto.trigger, context), accumulation_mode=proto.accumulation_mode, timestamp_combiner=proto.output_time, allowed_lateness=Duration(micros=proto.allowed_lateness * 1000), environment_id=None) @typehints.with_input_types(T) @typehints.with_output_types(T) class WindowInto(ParDo): """A window transform assigning windows to each element of a PCollection. Transforms an input PCollection by applying a windowing function to each element. Each transformed element in the result will be a WindowedValue element with the same input value and timestamp, with its new set of windows determined by the windowing function. """ class WindowIntoFn(DoFn): """A DoFn that applies a WindowInto operation.""" def __init__(self, windowing): # type: (Windowing) -> None self.windowing = windowing def process( self, element, timestamp=DoFn.TimestampParam, window=DoFn.WindowParam): context = WindowFn.AssignContext( timestamp, element=element, window=window) new_windows = self.windowing.windowfn.assign(context) yield WindowedValue(element, context.timestamp, new_windows) def infer_output_type(self, input_type): return input_type def __init__( self, windowfn, # type: typing.Union[Windowing, WindowFn] trigger=None, # type: typing.Optional[TriggerFn] accumulation_mode=None, timestamp_combiner=None, allowed_lateness=0): """Initializes a WindowInto transform. Args: windowfn (Windowing, WindowFn): Function to be used for windowing. trigger: (optional) Trigger used for windowing, or None for default. accumulation_mode: (optional) Accumulation mode used for windowing, required for non-trivial triggers. timestamp_combiner: (optional) Timestamp combniner used for windowing, or None for default. """ if isinstance(windowfn, Windowing): # Overlay windowing with kwargs. windowing = windowfn windowfn = windowing.windowfn # Use windowing to fill in defaults for the extra arguments. trigger = trigger or windowing.triggerfn accumulation_mode = accumulation_mode or windowing.accumulation_mode timestamp_combiner = timestamp_combiner or windowing.timestamp_combiner self.windowing = Windowing( windowfn, trigger, accumulation_mode, timestamp_combiner, allowed_lateness) super().__init__(self.WindowIntoFn(self.windowing)) def get_windowing(self, unused_inputs): # type: (typing.Any) -> Windowing return self.windowing def infer_output_type(self, input_type): return input_type def expand(self, pcoll): input_type = pcoll.element_type if input_type is not None: output_type = input_type self.with_input_types(input_type) self.with_output_types(output_type) return super().expand(pcoll) # typing: PTransform base class does not accept extra_kwargs def to_runner_api_parameter(self, context, **extra_kwargs): # type: ignore[override] # type: (PipelineContext, **typing.Any) -> typing.Tuple[str, message.Message] return ( common_urns.primitives.ASSIGN_WINDOWS.urn, self.windowing.to_runner_api(context)) @staticmethod def from_runner_api_parameter(unused_ptransform, proto, context): windowing = Windowing.from_runner_api(proto, context) return WindowInto( windowing.windowfn, trigger=windowing.triggerfn, accumulation_mode=windowing.accumulation_mode, timestamp_combiner=windowing.timestamp_combiner) PTransform.register_urn( common_urns.primitives.ASSIGN_WINDOWS.urn, # TODO(robertwb): Update WindowIntoPayload to include the full strategy. # (Right now only WindowFn is used, but we need this to reconstitute the # WindowInto transform, and in the future will need it at runtime to # support meta-data driven triggers.) # TODO(robertwb): Use a reference rather than embedding? beam_runner_api_pb2.WindowingStrategy, WindowInto.from_runner_api_parameter) # Python's pickling is broken for nested classes. WindowIntoFn = WindowInto.WindowIntoFn class Flatten(PTransform): """Merges several PCollections into a single PCollection. Copies all elements in 0 or more PCollections into a single output PCollection. If there are no input PCollections, the resulting PCollection will be empty (but see also kwargs below). Args: **kwargs: Accepts a single named argument "pipeline", which specifies the pipeline that "owns" this PTransform. Ordinarily Flatten can obtain this information from one of the input PCollections, but if there are none (or if there's a chance there may be none), this argument is the only way to provide pipeline information and should be considered mandatory. """ def __init__(self, **kwargs): super().__init__() self.pipeline = kwargs.pop( 'pipeline', None) # type: typing.Optional[Pipeline] if kwargs: raise ValueError('Unexpected keyword arguments: %s' % list(kwargs)) def _extract_input_pvalues(self, pvalueish): try: pvalueish = tuple(pvalueish) except TypeError: raise ValueError( 'Input to Flatten must be an iterable. ' 'Got a value of type %s instead.' % type(pvalueish)) # Spot check to see if any of the items are iterables of PCollections # and raise an error if so. This is always a user-error for idx, item in enumerate(pvalueish): if isinstance(item, (list, tuple)) and any( isinstance(sub_item, pvalue.PCollection) for sub_item in item): raise TypeError( 'Inputs to Flatten cannot include an iterable of PCollections. ' f'(input at index {idx}: "{item}")') return pvalueish, pvalueish def expand(self, pcolls): windowing = self.get_windowing(pcolls) for pcoll in pcolls: self._check_pcollection(pcoll) if pcoll.windowing != windowing: _LOGGER.warning( 'All input pcollections must have the same window. Windowing for ' 'flatten set to %s, windowing of pcoll %s set to %s', windowing, pcoll, pcoll.windowing) is_bounded = all(pcoll.is_bounded for pcoll in pcolls) return pvalue.PCollection(self.pipeline, is_bounded=is_bounded) def infer_output_type(self, input_type): return input_type def to_runner_api_parameter(self, context): # type: (PipelineContext) -> typing.Tuple[str, None] return common_urns.primitives.FLATTEN.urn, None @staticmethod def from_runner_api_parameter( unused_ptransform, unused_parameter, unused_context): return Flatten() PTransform.register_urn( common_urns.primitives.FLATTEN.urn, None, Flatten.from_runner_api_parameter) class FlattenWith(PTransform): """A PTransform that flattens its input with other PCollections. This is equivalent to creating a tuple containing both the input and the other PCollection(s), but has the advantage that it can be more easily used inline. Root PTransforms can be passed as well as PCollections, in which case their outputs will be flattened. """ def __init__(self, *others): self._others = others def expand(self, pcoll): pcolls = [pcoll] for other in self._others: if isinstance(other, pvalue.PCollection): pcolls.append(other) elif isinstance(other, PTransform): pcolls.append(pcoll.pipeline | other) else: raise TypeError( 'FlattenWith only takes other PCollections and PTransforms, ' f'got {other}') return tuple(pcolls) | Flatten() class Create(PTransform): """A transform that creates a PCollection from an iterable.""" def __init__(self, values, reshuffle=True): """Initializes a Create transform. Args: values: An object of values for the PCollection """ super().__init__() if isinstance(values, (str, bytes)): raise TypeError( 'PTransform Create: Refusing to treat string as ' 'an iterable. (string=%r)' % values) elif isinstance(values, dict): values = values.items() self.values = tuple(values) self.reshuffle = reshuffle self._coder = typecoders.registry.get_coder(self.get_output_type()) def __getstate__(self): serialized_values = [self._coder.encode(v) for v in self.values] return serialized_values, self.reshuffle, self._coder def __setstate__(self, state): serialized_values, self.reshuffle, self._coder = state self.values = [self._coder.decode(v) for v in serialized_values] def to_runner_api_parameter(self, context): # type: (PipelineContext) -> typing.Tuple[str, bytes] # Required as this is identified by type in PTransformOverrides. # TODO(https://github.com/apache/beam/issues/18713): Use an actual URN # here. return self.to_runner_api_pickled(context) def infer_output_type(self, unused_input_type): if not self.values: return typehints.Any # No field data - just use default Union. if not hasattr(self.values[0], 'as_dict'): return typehints.Union[[ trivial_inference.instance_to_type(v) for v in self.values ]] first_fields = self.values[0].as_dict().keys() # Save field types for each field field_types_by_field = defaultdict(set) for row in self.values: row_dict = row.as_dict() for field in first_fields: field_types_by_field[field].add( trivial_inference.instance_to_type(row_dict.get(field))) # Determine the appropriate type for each field final_fields = [] for field in first_fields: field_types = field_types_by_field[field] non_none_types = {t for t in field_types if t is not type(None)} if len(non_none_types) > 1: final_type = typehints.Union[tuple(non_none_types)] elif len(non_none_types) == 1 and len(field_types) == 1: final_type = non_none_types.pop() elif len(non_none_types) == 1 and len(field_types) == 2: final_type = typehints.Optional[non_none_types.pop()] else: raise TypeError("No types found for field %s", field) final_fields.append((field, final_type)) return row_type.RowTypeConstraint.from_fields(final_fields) def get_output_type(self): return ( self.get_type_hints().simple_output_type(self.label) or self.infer_output_type(None)) def expand(self, pbegin): assert isinstance(pbegin, pvalue.PBegin) serialized_values = [self._coder.encode(v) for v in self.values] reshuffle = self.reshuffle # Avoid the "redistributing" reshuffle for 0 and 1 element Creates. # These special cases are often used in building up more complex # transforms (e.g. Write). class MaybeReshuffle(PTransform): def expand(self, pcoll): if len(serialized_values) > 1 and reshuffle: from apache_beam.transforms.util import Reshuffle return pcoll | Reshuffle() else: return pcoll # Map(lambda) produces a label formatted like this, but it cannot be # changed without breaking update compat. Here, we pin to the transform # name used in the 2.68 release to avoid breaking changes when the line # number changes. Context: https://github.com/apache/beam/pull/36381 return ( pbegin | Impulse() | "FlatMap()" >> FlatMap(lambda _: serialized_values).with_output_types(bytes) | MaybeReshuffle().with_output_types(bytes) | Map(self._coder.decode).with_output_types(self.get_output_type())) def as_read(self): from apache_beam.io import iobase source = self._create_source_from_iterable(self.values, self._coder) return iobase.Read(source).with_output_types(self.get_output_type()) def get_windowing(self, unused_inputs): # type: (typing.Any) -> Windowing return Windowing(GlobalWindows()) @staticmethod def _create_source_from_iterable(values, coder): return Create._create_source(list(map(coder.encode, values)), coder) @staticmethod def _create_source(serialized_values, coder): # type: (typing.Any, typing.Any) -> create_source._CreateSource from apache_beam.transforms.create_source import _CreateSource return _CreateSource(serialized_values, coder) @typehints.with_output_types(bytes) class Impulse(PTransform): """Impulse primitive.""" def expand(self, pbegin): if not isinstance(pbegin, pvalue.PBegin): raise TypeError( 'Input to Impulse transform must be a PBegin but found %s' % pbegin) return pvalue.PCollection(pbegin.pipeline, element_type=bytes) def get_windowing(self, inputs): # type: (typing.Any) -> Windowing return Windowing(GlobalWindows()) def infer_output_type(self, unused_input_type): return bytes def to_runner_api_parameter(self, unused_context): # type: (PipelineContext) -> typing.Tuple[str, None] return common_urns.primitives.IMPULSE.urn, None @staticmethod @PTransform.register_urn(common_urns.primitives.IMPULSE.urn, None) def from_runner_api_parameter( unused_ptransform, unused_parameter, unused_context): return Impulse() def _strip_output_annotations(type_hint, strip_tagged_output=True): # TODO(robertwb): These should be parameterized types that the # type inferencer understands. # Then we can replace them with the correct element types instead of # using Any. Refer to typehints.WindowedValue when doing this. annotations = [TimestampedValue, WindowedValue] if strip_tagged_output: annotations.append(pvalue.TaggedOutput) annotations = tuple(annotations) contains_annotation = False def visitor(t, unused_args): if t in annotations or (hasattr(t, '__name__') and t.__name__ == TimestampedValue.__name__): raise StopIteration try: visit_inner_types(type_hint, visitor, []) except StopIteration: contains_annotation = True return typehints.Any if contains_annotation else type_hint ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/core_it_test.py0000644000175100017510000001467715206113033023130 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration tests for cross-language transform expansion.""" # pytype: skip-file import sys import unittest from datetime import datetime import pytest import apache_beam as beam from apache_beam.options.pipeline_options import SetupOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.util import GcpSecret from apache_beam.transforms.util import Secret try: from google.cloud import secretmanager except ImportError: secretmanager = None # type: ignore[assignment] try: from google.cloud import kms except ImportError: kms = None # type: ignore[assignment] class GbekIT(unittest.TestCase): @classmethod def setUpClass(cls): if secretmanager is not None: cls.project_id = 'apache-beam-testing' py_version = f'_py{sys.version_info.major}{sys.version_info.minor}' secret_postfix = datetime.now().strftime('%m%d_%H%M%S') + py_version cls.secret_id = 'gbekit_secret_tests_' + secret_postfix cls.client = secretmanager.SecretManagerServiceClient() cls.project_path = f'projects/{cls.project_id}' cls.secret_path = f'{cls.project_path}/secrets/{cls.secret_id}' try: cls.client.get_secret(request={'name': cls.secret_path}) except Exception: cls.client.create_secret( request={ 'parent': cls.project_path, 'secret_id': cls.secret_id, 'secret': { 'replication': { 'automatic': {} } } }) cls.client.add_secret_version( request={ 'parent': cls.secret_path, 'payload': { 'data': Secret.generate_secret_bytes() } }) version_name = f'{cls.secret_path}/versions/latest' cls.gcp_secret = GcpSecret(version_name) cls.secret_option = f'type:GcpSecret;version_name:{version_name}' if kms is not None: cls.kms_client = kms.KeyManagementServiceClient() cls.location_id = 'global' py_version = f'_py{sys.version_info.major}{sys.version_info.minor}' secret_postfix = datetime.now().strftime('%m%d_%H%M%S') + py_version cls.key_ring_id = 'gbekit_key_ring_tests' cls.key_ring_path = cls.kms_client.key_ring_path( cls.project_id, cls.location_id, cls.key_ring_id) try: cls.kms_client.get_key_ring(request={'name': cls.key_ring_path}) except Exception: parent = f'projects/{cls.project_id}/locations/{cls.location_id}' cls.kms_client.create_key_ring( request={ 'parent': parent, 'key_ring_id': cls.key_ring_id, }) cls.key_id = 'gbekit_key_tests' cls.key_path = cls.kms_client.crypto_key_path( cls.project_id, cls.location_id, cls.key_ring_id, cls.key_id) try: cls.kms_client.get_crypto_key(request={'name': cls.key_path}) except Exception: cls.kms_client.create_crypto_key( request={ 'parent': cls.key_ring_path, 'crypto_key_id': cls.key_id, 'crypto_key': { 'purpose': kms.CryptoKey.CryptoKeyPurpose.ENCRYPT_DECRYPT } }) cls.hsm_secret_option = ( f'type:GcpHsmGeneratedSecret;project_id:{cls.project_id};' f'location_id:{cls.location_id};key_ring_id:{cls.key_ring_id};' f'key_id:{cls.key_id};job_name:{secret_postfix}') @classmethod def tearDownClass(cls): if secretmanager is not None: cls.client.delete_secret(request={'name': cls.secret_path}) @pytest.mark.it_postcommit @unittest.skipIf(secretmanager is None, 'GCP dependencies are not installed') def test_gbk_with_gbek_it(self): pipeline = TestPipeline(is_integration_test=True) pipeline.options.view_as(SetupOptions).gbek = self.secret_option pcoll_1 = pipeline | 'Start 1' >> beam.Create([('a', 1), ('a', 2), ('b', 3), ('c', 4)]) result = (pcoll_1) | beam.GroupByKey() sorted_result = result | beam.Map(lambda x: (x[0], sorted(x[1]))) assert_that( sorted_result, equal_to([('a', ([1, 2])), ('b', ([3])), ('c', ([4]))])) pipeline.run().wait_until_finish() @pytest.mark.it_postcommit @unittest.skipIf(secretmanager is None, 'GCP dependencies are not installed') @unittest.skipIf(kms is None, 'GCP dependencies are not installed') def test_gbk_with_gbek_hsm_it(self): pipeline = TestPipeline(is_integration_test=True) pipeline.options.view_as(SetupOptions).gbek = self.hsm_secret_option pcoll_1 = pipeline | 'Start 1' >> beam.Create([('a', 1), ('a', 2), ('b', 3), ('c', 4)]) result = (pcoll_1) | beam.GroupByKey() sorted_result = result | beam.Map(lambda x: (x[0], sorted(x[1]))) assert_that( sorted_result, equal_to([('a', ([1, 2])), ('b', ([3])), ('c', ([4]))])) pipeline.run().wait_until_finish() @pytest.mark.it_postcommit @unittest.skipIf(secretmanager is None, 'GCP dependencies are not installed') def test_combineValues_with_gbek_it(self): pipeline = TestPipeline(is_integration_test=True) pipeline.options.view_as(SetupOptions).gbek = self.secret_option pcoll_1 = pipeline | 'Start 1' >> beam.Create([('a', 1), ('a', 2), ('b', 3), ('c', 4)]) result = (pcoll_1) | beam.CombinePerKey(sum) assert_that(result, equal_to([('a', 3), ('b', 3), ('c', 4)])) pipeline.run().wait_until_finish() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/core_test.py0000644000175100017510000007035415206113033022426 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the core python file.""" # pytype: skip-file import logging import os import tempfile import unittest from typing import Iterable from typing import Literal from typing import TypeVar import pytest import apache_beam as beam from apache_beam.coders import coders from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.core import DeadLetter from apache_beam.transforms.resources import ResourceHint from apache_beam.transforms.userstate import BagStateSpec from apache_beam.transforms.userstate import ReadModifyWriteStateSpec from apache_beam.transforms.userstate import TimerSpec from apache_beam.transforms.userstate import on_timer from apache_beam.transforms.window import FixedWindows from apache_beam.typehints import TypeCheckError from apache_beam.typehints import row_type from apache_beam.typehints import typehints RETURN_NONE_PARTIAL_WARNING = "Process method returned None" class TestDoFn0(beam.DoFn): """Returning without a value is allowed""" def process(self, element): if not element: return yield element class TestDoFn1(beam.DoFn): def process(self, element): yield element class TestDoFn2(beam.DoFn): def process(self, element): def inner_func(x): yield x return inner_func(element) class TestDoFn3(beam.DoFn): """mixing return and yield is not allowed""" def process(self, element): if not element: return -1 yield element class TestDoFn4(beam.DoFn): """test the variable name containing return""" def process(self, element): my_return = element yield my_return class TestDoFn5(beam.DoFn): """test the variable name containing yield""" def process(self, element): my_yield = element return my_yield class TestDoFn6(beam.DoFn): """test the variable name containing return""" def process(self, element): return_test = element yield return_test class TestDoFn7(beam.DoFn): """test the variable name containing yield""" def process(self, element): yield_test = element return yield_test class TestDoFn8(beam.DoFn): """test the code containing yield and yield from""" def process(self, element): if not element: yield from [1, 2, 3] else: yield element class TestDoFn9(beam.DoFn): def process(self, element): if len(element) > 3: raise ValueError('Not allowed to have long elements') yield element class TestDoFn10(beam.DoFn): """test process returning None explicitly""" def process(self, element): return None class TestDoFn11(beam.DoFn): """test process returning None (no return and no yield)""" def process(self, element): pass class TestDoFn12(beam.DoFn): """test process returning None in a filter pattern""" def process(self, element): if element == 0: return return element class TestDoFnStateful(beam.DoFn): STATE_SPEC = ReadModifyWriteStateSpec('num_elements', coders.VarIntCoder()) """test process with a stateful dofn""" def process(self, element, state=beam.DoFn.StateParam(STATE_SPEC)): if len(element[1]) > 3: raise ValueError('Not allowed to have long elements') current_value = state.read() or 1 state.write(current_value + 1) yield current_value class TestDoFnWithTimer(beam.DoFn): ALL_ELEMENTS = BagStateSpec('buffer', coders.VarIntCoder()) TIMER = TimerSpec('timer', beam.TimeDomain.WATERMARK) """test process with a stateful dofn""" def process( self, element, t=beam.DoFn.TimestampParam, state=beam.DoFn.StateParam(ALL_ELEMENTS), timer=beam.DoFn.TimerParam(TIMER)): if element[1] > 3: raise ValueError('Not allowed to have large numbers') state.add(element[1]) timer.set(t) return [] @on_timer(TIMER) def expiry_callback(self, state=beam.DoFn.StateParam(ALL_ELEMENTS)): unique_elements = list(state.read()) state.clear() return unique_elements class CreateTest(unittest.TestCase): @pytest.fixture(autouse=True) def inject_fixtures(self, caplog): self._caplog = caplog def test_dofn_with_yield_and_return(self): warning_text = 'Using yield and return' with self._caplog.at_level(logging.WARNING): assert beam.ParDo(sum) assert beam.ParDo(TestDoFn0()) assert beam.ParDo(TestDoFn1()) assert beam.ParDo(TestDoFn2()) assert beam.ParDo(TestDoFn4()) assert beam.ParDo(TestDoFn5()) assert beam.ParDo(TestDoFn6()) assert beam.ParDo(TestDoFn7()) assert beam.ParDo(TestDoFn8()) assert warning_text not in self._caplog.text with self._caplog.at_level(logging.WARNING): beam.ParDo(TestDoFn3()) assert warning_text in self._caplog.text def test_dofn_with_explicit_return_none(self): with self._caplog.at_level(logging.WARNING): beam.ParDo(TestDoFn10()) assert RETURN_NONE_PARTIAL_WARNING in self._caplog.text assert str(TestDoFn10) in self._caplog.text def test_dofn_with_implicit_return_none_missing_return_and_yield(self): with self._caplog.at_level(logging.WARNING): beam.ParDo(TestDoFn11()) assert RETURN_NONE_PARTIAL_WARNING not in self._caplog.text def test_dofn_with_implicit_return_none_and_value(self): with self._caplog.at_level(logging.WARNING): beam.ParDo(TestDoFn12()) assert RETURN_NONE_PARTIAL_WARNING not in self._caplog.text class PartitionTest(unittest.TestCase): def test_partition_with_bools(self): with pytest.raises( (ValueError, RuntimeError), match= r"PartitionFn yielded a '([^']*)' when it should only yield integers"): # Check for RuntimeError too since the portable runner casts # all exceptions to RuntimeError invalid_inputs = [True, 1.2, 'string', None] for input_value in invalid_inputs: with beam.testing.test_pipeline.TestPipeline() as p: _ = ( p | beam.Create([input_value]) | beam.Partition(lambda x, _: x, 2)) def test_partition_with_numpy_integers(self): # Test that numpy integer types are correctly accepted by the # ApplyPartitionFnFn class import numpy as np # Create an instance of the ApplyPartitionFnFn class apply_partition_fn = beam.Partition.ApplyPartitionFnFn() # Define a simple partition function class SimplePartitionFn(beam.PartitionFn): def partition_for(self, element, num_partitions): return element % num_partitions partition_fn = SimplePartitionFn() # Test with numpy.int32 # This should not raise an exception outputs = list(apply_partition_fn.process(np.int32(1), partition_fn, 3)) self.assertEqual(len(outputs), 1) self.assertEqual(outputs[0].tag, '1') # 1 % 3 = 1 # Test with numpy.int64 # This should not raise an exception outputs = list(apply_partition_fn.process(np.int64(2), partition_fn, 3)) self.assertEqual(len(outputs), 1) self.assertEqual(outputs[0].tag, '2') # 2 % 3 = 2 def test_partition_fn_returning_numpy_integers(self): # Test that partition functions can return numpy integer types import numpy as np # Create an instance of the ApplyPartitionFnFn class apply_partition_fn = beam.Partition.ApplyPartitionFnFn() # Define partition functions that return numpy integer types class Int32PartitionFn(beam.PartitionFn): def partition_for(self, element, num_partitions): return np.int32(element % num_partitions) class Int64PartitionFn(beam.PartitionFn): def partition_for(self, element, num_partitions): return np.int64(element % num_partitions) # Test with partition function returning numpy.int32 # This should not raise an exception outputs = list(apply_partition_fn.process(1, Int32PartitionFn(), 3)) self.assertEqual(len(outputs), 1) self.assertEqual(outputs[0].tag, '1') # 1 % 3 = 1 # Test with partition function returning numpy.int64 # This should not raise an exception outputs = list(apply_partition_fn.process(2, Int64PartitionFn(), 3)) self.assertEqual(len(outputs), 1) self.assertEqual(outputs[0].tag, '2') # 2 % 3 = 2 def test_partition_boundedness(self): def partition_fn(val, num_partitions): return val % num_partitions class UnboundedDoFn(beam.DoFn): @beam.DoFn.unbounded_per_element() def process(self, element): yield element with beam.testing.test_pipeline.TestPipeline() as p: source = p | beam.Create([1, 2, 3, 4, 5]) p1, p2, p3 = source | "bounded" >> beam.Partition(partition_fn, 3) self.assertEqual(source.is_bounded, True) self.assertEqual(p1.is_bounded, True) self.assertEqual(p2.is_bounded, True) self.assertEqual(p3.is_bounded, True) unbounded = source | beam.ParDo(UnboundedDoFn()) p4, p5, p6 = unbounded | "unbounded" >> beam.Partition(partition_fn, 3) self.assertEqual(unbounded.is_bounded, False) self.assertEqual(p4.is_bounded, False) self.assertEqual(p5.is_bounded, False) self.assertEqual(p6.is_bounded, False) class FlattenTest(unittest.TestCase): def test_flatten_identical_windows(self): with beam.testing.test_pipeline.TestPipeline() as p: source1 = p | "c1" >> beam.Create( [1, 2, 3, 4, 5]) | "w1" >> beam.WindowInto(FixedWindows(100)) source2 = p | "c2" >> beam.Create([6, 7, 8]) | "w2" >> beam.WindowInto( FixedWindows(100)) source3 = p | "c3" >> beam.Create([9, 10]) | "w3" >> beam.WindowInto( FixedWindows(100)) out = (source1, source2, source3) | "flatten" >> beam.Flatten() assert_that(out, equal_to([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])) def test_flatten_no_windows(self): with beam.testing.test_pipeline.TestPipeline() as p: source1 = p | "c1" >> beam.Create([1, 2, 3, 4, 5]) source2 = p | "c2" >> beam.Create([6, 7, 8]) source3 = p | "c3" >> beam.Create([9, 10]) out = (source1, source2, source3) | "flatten" >> beam.Flatten() assert_that(out, equal_to([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])) def test_flatten_mismatched_windows(self): with beam.testing.test_pipeline.TestPipeline() as p: source1 = p | "c1" >> beam.Create( [1, 2, 3, 4, 5]) | "w1" >> beam.WindowInto(FixedWindows(25)) source2 = p | "c2" >> beam.Create([6, 7, 8]) | "w2" >> beam.WindowInto( FixedWindows(100)) source3 = p | "c3" >> beam.Create([9, 10]) | "w3" >> beam.WindowInto( FixedWindows(100)) _ = (source1, source2, source3) | "flatten" >> beam.Flatten() class ExceptionHandlingTest(unittest.TestCase): def test_routes_failures(self): with beam.Pipeline() as pipeline: good, bad = ( pipeline | beam.Create(['abc', 'long_word', 'foo', 'bar', 'foobar']) | beam.ParDo(TestDoFn9()).with_exception_handling() ) bad_elements = bad | beam.Keys() assert_that(good, equal_to(['abc', 'foo', 'bar']), 'good') assert_that(bad_elements, equal_to(['long_word', 'foobar']), 'bad') def test_handles_callbacks(self): with tempfile.TemporaryDirectory() as tmp_dirname: tmp_path = os.path.join(tmp_dirname, 'tmp_filename') file_contents = 'random content' def failure_callback(e, el): if type(e) is not ValueError: raise Exception(f'Failed to pass in correct exception, received {e}') if el != 'foobar': raise Exception(f'Failed to pass in correct element, received {el}') f = open(tmp_path, "a") logging.warning(tmp_path) f.write(file_contents) f.close() with beam.Pipeline() as pipeline: good, bad = ( pipeline | beam.Create(['abc', 'bcd', 'foo', 'bar', 'foobar']) | beam.ParDo(TestDoFn9()).with_exception_handling( on_failure_callback=failure_callback) ) bad_elements = bad | beam.Keys() assert_that(good, equal_to(['abc', 'bcd', 'foo', 'bar']), 'good') assert_that(bad_elements, equal_to(['foobar']), 'bad') with open(tmp_path) as f: s = f.read() self.assertEqual(s, file_contents) def test_handles_no_callback_triggered(self): with tempfile.TemporaryDirectory() as tmp_dirname: tmp_path = os.path.join(tmp_dirname, 'tmp_filename') file_contents = 'random content' def failure_callback(e, el): f = open(tmp_path, "a") logging.warning(tmp_path) f.write(file_contents) f.close() with beam.Pipeline() as pipeline: good, bad = ( pipeline | beam.Create(['abc', 'bcd', 'foo', 'bar']) | beam.ParDo(TestDoFn9()).with_exception_handling( on_failure_callback=failure_callback) ) bad_elements = bad | beam.Keys() assert_that(good, equal_to(['abc', 'bcd', 'foo', 'bar']), 'good') assert_that(bad_elements, equal_to([]), 'bad') self.assertFalse(os.path.isfile(tmp_path)) def test_stateful_exception_handling(self): with beam.Pipeline() as pipeline: good, bad = ( pipeline | beam.Create([(1, 'abc'), (1, 'long_word'), (1, 'foo'), (1, 'bar'), (1, 'foobar')]) | beam.ParDo(TestDoFnStateful()).with_exception_handling( allow_unsafe_userstate_in_process=True) ) bad_elements = bad | beam.Keys() assert_that(good, equal_to([1, 2, 3]), 'good') assert_that( bad_elements, equal_to([(1, 'long_word'), (1, 'foobar')]), 'bad') def test_timer_exception_handling(self): with beam.Pipeline() as pipeline: good, bad = ( pipeline | beam.Create([(1, 0), (1, 1), (1, 2), (1, 5), (1, 10)]) | beam.ParDo(TestDoFnWithTimer()).with_exception_handling( allow_unsafe_userstate_in_process=True) ) bad_elements = bad | beam.Keys() assert_that(good, equal_to([0, 1, 2]), 'good') assert_that(bad_elements, equal_to([(1, 5), (1, 10)]), 'bad') def test_tags_with_exception_handling_then_resource_hint(self): class TagHint(ResourceHint): urn = 'beam:resources:tags:v1' ResourceHint.register_resource_hint('tags', TagHint) with beam.Pipeline() as pipeline: ok, unused_errors = ( pipeline | beam.Create([1]) | beam.Map(lambda x: x) .with_exception_handling() .with_resource_hints(tags='test_tag') ) pd = ok.producer.transform self.assertIsInstance(pd, beam.transforms.core.ParDo) while hasattr(pd.fn, 'fn'): pd = pd.fn self.assertEqual( pd.get_resource_hints(), {'beam:resources:tags:v1': b'test_tag'}, ) def test_tags_with_exception_handling_timeout_then_resource_hint(self): class TagHint(ResourceHint): urn = 'beam:resources:tags:v1' ResourceHint.register_resource_hint('tags', TagHint) with beam.Pipeline() as pipeline: ok, unused_errors = ( pipeline | beam.Create([1]) | beam.Map(lambda x: x) .with_exception_handling(timeout=1) .with_resource_hints(tags='test_tag') ) pd = ok.producer.transform self.assertIsInstance(pd, beam.transforms.core.ParDo) while hasattr(pd.fn, 'fn'): pd = pd.fn self.assertEqual( pd.get_resource_hints(), {'beam:resources:tags:v1': b'test_tag'}, ) def test_tags_with_resource_hint_then_exception_handling(self): class TagHint(ResourceHint): urn = 'beam:resources:tags:v1' ResourceHint.register_resource_hint('tags', TagHint) with beam.Pipeline() as pipeline: ok, unused_errors = ( pipeline | beam.Create([1]) | beam.Map(lambda x: x) .with_resource_hints(tags='test_tag') .with_exception_handling() ) pd = ok.producer.transform self.assertIsInstance(pd, beam.transforms.core.ParDo) while hasattr(pd.fn, 'fn'): pd = pd.fn self.assertEqual( pd.get_resource_hints(), {'beam:resources:tags:v1': b'test_tag'}, ) def test_tags_with_resource_hint_then_exception_handling_timeout(self): class TagHint(ResourceHint): urn = 'beam:resources:tags:v1' ResourceHint.register_resource_hint('tags', TagHint) with beam.Pipeline() as pipeline: ok, unused_errors = ( pipeline | beam.Create([1]) | beam.Map(lambda x: x) .with_resource_hints(tags='test_tag') .with_exception_handling(timeout=1) ) pd = ok.producer.transform self.assertIsInstance(pd, beam.transforms.core.ParDo) while hasattr(pd.fn, 'fn'): pd = pd.fn self.assertEqual( pd.get_resource_hints(), {'beam:resources:tags:v1': b'test_tag'}, ) class ExceptionHandlingWithOutputsTest(unittest.TestCase): """Tests for combining with_exception_handling() and with_outputs().""" def _create_dofn_with_tagged_outputs(self): """A DoFn that yields tagged outputs and can raise on even numbers.""" class DoWithFailures(beam.DoFn): def process( self, element: int ) -> Iterable[int | beam.pvalue.TaggedOutput[Literal['threes'], int] | beam.pvalue.TaggedOutput[Literal['fives'], str]]: if element % 2 == 0: raise ValueError(f'Even numbers not allowed {element}') if element % 3 == 0: yield beam.pvalue.TaggedOutput('threes', element) # type: ignore[misc] elif element % 5 == 0: yield beam.pvalue.TaggedOutput('fives', str(element)) # type: ignore[misc] else: yield element return DoWithFailures() def test_with_exception_handling_then_with_outputs(self): """Direction 1: .with_exception_handling().with_outputs()""" with beam.Pipeline() as p: results = ( p | beam.Create([1, 2, 3, 4, 5, 6, 7]) | beam.ParDo(self._create_dofn_with_tagged_outputs()). with_exception_handling().with_outputs( 'threes', 'fives', main='main')) assert_that(results.main, equal_to([1, 7]), 'main') assert_that(results.threes, equal_to([3]), 'threes') assert_that(results.fives, equal_to(['5']), 'fives') bad_elements = results.bad | beam.Keys() assert_that(bad_elements, equal_to([2, 4, 6]), 'bad') # Verify type hints from annotations are propagated self.assertEqual(results.main.element_type, int) self.assertEqual(results.threes.element_type, int) self.assertEqual(results.fives.element_type, str) self.assertEqual(results.bad.element_type, DeadLetter[int]) def test_with_outputs_then_with_exception_handling(self): """Direction 2: .with_outputs().with_exception_handling()""" with beam.Pipeline() as p: results = ( p | beam.Create([1, 2, 3, 4, 5, 6, 7]) | beam.ParDo(self._create_dofn_with_tagged_outputs()).with_outputs( 'threes', 'fives', main='main').with_exception_handling()) assert_that(results.main, equal_to([1, 7]), 'main') assert_that(results.threes, equal_to([3]), 'threes') assert_that(results.fives, equal_to(['5']), 'fives') bad_elements = results.bad | beam.Keys() assert_that(bad_elements, equal_to([2, 4, 6]), 'bad') # Verify type hints from annotations are propagated self.assertEqual(results.main.element_type, int) self.assertEqual(results.threes.element_type, int) self.assertEqual(results.fives.element_type, str) self.assertEqual(results.bad.element_type, DeadLetter[int]) def test_with_outputs_then_with_exception_handling_custom_dead_letter_tag( self): """Direction 2 with custom dead_letter_tag.""" with beam.Pipeline() as p: results = ( p | beam.Create([1, 2, 3]) | beam.ParDo(self._create_dofn_with_tagged_outputs()).with_outputs( 'threes', main='main').with_exception_handling(dead_letter_tag='errors')) assert_that(results.main, equal_to([1]), 'main') assert_that(results.threes, equal_to([3]), 'threes') bad_elements = results.errors | beam.Keys() assert_that(bad_elements, equal_to([2]), 'errors') self.assertEqual(results.threes.element_type, int) self.assertEqual(results.errors.element_type, DeadLetter[int]) def test_with_exception_handling_then_with_outputs_custom_dead_letter_tag( self): """Direction 1 with custom dead_letter_tag.""" with beam.Pipeline() as p: results = ( p | beam.Create([1, 2, 3]) | beam.ParDo( self._create_dofn_with_tagged_outputs()).with_exception_handling( dead_letter_tag='errors').with_outputs('threes', main='main')) assert_that(results.main, equal_to([1]), 'main') assert_that(results.threes, equal_to([3]), 'threes') bad_elements = results.errors | beam.Keys() assert_that(bad_elements, equal_to([2]), 'errors') self.assertEqual(results.threes.element_type, int) self.assertEqual(results.errors.element_type, DeadLetter[int]) def test_exception_handling_no_with_outputs_backward_compat(self): """Without with_outputs(), behavior is unchanged.""" with beam.Pipeline() as p: good, bad = ( p | beam.Create([1, 2, 7]) | beam.ParDo(self._create_dofn_with_tagged_outputs()) .with_exception_handling()) assert_that(good, equal_to([1, 7]), 'good') bad_elements = bad | beam.Keys() assert_that(bad_elements, equal_to([2]), 'bad') def test_exception_handling_compat_version_uses_old_behavior(self): """With compat version < 2.73.0, old expand path is used.""" options = PipelineOptions(update_compatibility_version="2.72.0") with beam.Pipeline(options=options) as p: good, bad = ( p | beam.Create([1, 2, 7]) | beam.ParDo(self._create_dofn_with_tagged_outputs()) .with_exception_handling()) assert_that(good, equal_to([1, 7]), 'good') bad_elements = bad | beam.Keys() assert_that(bad_elements, equal_to([2]), 'bad') def test_exception_handling_compat_version_element_type_set_manually(self): """With compat version < 2.73.0, element_type is set via manual override (the old behavior) rather than via with_output_types.""" options = PipelineOptions(update_compatibility_version="2.72.0") with beam.Pipeline(options=options) as p: results = ( p | beam.Create([1, 2, 3]) | beam.ParDo(self._create_dofn_with_tagged_outputs()). with_exception_handling().with_outputs('threes', main='main')) # In old path, dead letter type is Any (no with_output_types call) self.assertEqual(results.bad.element_type, typehints.Any) # Tagged outputs still get types from DoFn Literal annotations # (via DoOutputsTuple.__getitem__ reading tagged_output_types) self.assertEqual(results.threes.element_type, int) # Main output type should still be inferred via manual override assert_that(results.main, equal_to([1]), 'main') def test_with_outputs_then_exception_handling_with_map(self): """with_outputs().with_exception_handling() also works on Map.""" with beam.Pipeline() as p: results = ( p | beam.Create([1, 2, 3, 4, 5]) | beam.Map(lambda x: x if x % 2 != 0 else 1 / 0).with_outputs( main='main').with_exception_handling()) assert_that(results.main, equal_to([1, 3, 5]), 'main') bad_elements = results.bad | beam.Keys() assert_that(bad_elements, equal_to([2, 4]), 'bad') def test_with_output_types_chained_on_pardo(self): """When type hints are chained on the ParDo (not annotations on the DoFn), tagged output types should still be propagated through with_exception_handling().with_outputs().""" class DoWithFailuresNoAnnotations(beam.DoFn): def process(self, element): if element % 2 == 0: raise ValueError(f'Even numbers not allowed {element}') if element % 3 == 0: yield beam.pvalue.TaggedOutput('threes', element) else: yield element with beam.Pipeline() as p: results = ( p | beam.Create([1, 2, 3, 7]) | beam.ParDo(DoWithFailuresNoAnnotations()).with_output_types( int, threes=int).with_exception_handling().with_outputs( 'threes', main='main')) assert_that(results.main, equal_to([1, 7]), 'main') assert_that(results.threes, equal_to([3]), 'threes') bad_elements = results.bad | beam.Keys() assert_that(bad_elements, equal_to([2]), 'bad') self.assertEqual(results.main.element_type, int) self.assertEqual(results.threes.element_type, int) def test_with_outputs_and_error_handler(self): """with_outputs() + error_handler should return DoOutputsTuple, not a bare PCollection.""" from apache_beam.transforms.error_handling import ErrorHandler with beam.Pipeline() as p: with ErrorHandler(beam.Map(lambda x: x)) as handler: results = ( p | beam.Create([1, 2, 3, 4, 5, 6, 7]) | beam.ParDo(self._create_dofn_with_tagged_outputs()).with_outputs( 'threes', 'fives', main='main').with_exception_handling(error_handler=handler)) assert_that(results.main, equal_to([1, 7]), 'main') assert_that(results.threes, equal_to([3]), 'threes') assert_that(results.fives, equal_to(['5']), 'fives') def test_callablewrapper_typehint(): T = TypeVar("T") def identity(x: T) -> T: return x dofn = beam.core.CallableWrapperDoFn(identity) assert dofn.get_type_hints().strip_iterable()[1][0][0] == typehints.Any class FlatMapTest(unittest.TestCase): def test_default(self): with beam.Pipeline() as pipeline: letters = ( pipeline | beam.Create(['abc', 'def'], reshuffle=False) | beam.FlatMap()) assert_that(letters, equal_to(['a', 'b', 'c', 'd', 'e', 'f'])) def test_default_identity_function_with_typehint(self): with beam.Pipeline() as pipeline: letters = ( pipeline | beam.Create([["abc"]], reshuffle=False) | beam.FlatMap() | beam.Map(lambda s: s.upper()).with_input_types(str)) assert_that(letters, equal_to(["ABC"])) def test_typecheck_with_default(self): with pytest.raises(TypeCheckError): with beam.Pipeline() as pipeline: _ = ( pipeline | beam.Create([[1, 2, 3]], reshuffle=False) | beam.FlatMap() | beam.Map(lambda s: s.upper()).with_input_types(str)) class CreateInferOutputSchemaTest(unittest.TestCase): def test_multiple_types_for_field(self): output_type = beam.Create([beam.Row(a=1), beam.Row(a='foo')]).infer_output_type(None) self.assertEqual( output_type, row_type.RowTypeConstraint.from_fields([ ('a', typehints.Union[int, str]) ])) def test_single_type_for_field(self): output_type = beam.Create([beam.Row(a=1), beam.Row(a=2)]).infer_output_type(None) self.assertEqual( output_type, row_type.RowTypeConstraint.from_fields([('a', int)])) def test_optional_type_for_field(self): output_type = beam.Create([beam.Row(a=1), beam.Row(a=None)]).infer_output_type(None) self.assertEqual( output_type, row_type.RowTypeConstraint.from_fields([('a', typehints.Optional[int]) ])) def test_none_type_for_field_raises_error(self): with self.assertRaisesRegex(TypeError, "('No types found for field %s', 'a')"): beam.Create([beam.Row(a=None), beam.Row(a=None)]).infer_output_type(None) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/create_source.py0000644000175100017510000000651715206113033023262 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file from apache_beam.io import iobase from apache_beam.transforms.core import Create class _CreateSource(iobase.BoundedSource): """Internal source that is used by Create()""" def __init__(self, serialized_values, coder): self._coder = coder self._serialized_values = [] self._total_size = 0 self._serialized_values = serialized_values self._total_size = sum(map(len, self._serialized_values)) def read(self, range_tracker): start_position = range_tracker.start_position() current_position = start_position def split_points_unclaimed(stop_position): if current_position >= stop_position: return 0 return stop_position - current_position - 1 range_tracker.set_split_points_unclaimed_callback(split_points_unclaimed) element_iter = iter(self._serialized_values[start_position:]) for i in range(start_position, range_tracker.stop_position()): if not range_tracker.try_claim(i): return current_position = i yield self._coder.decode(next(element_iter)) def split(self, desired_bundle_size, start_position=None, stop_position=None): if len(self._serialized_values) < 2: yield iobase.SourceBundle( weight=0, source=self, start_position=0, stop_position=len(self._serialized_values)) else: if start_position is None: start_position = 0 if stop_position is None: stop_position = len(self._serialized_values) avg_size_per_value = self._total_size // len(self._serialized_values) num_values_per_split = max( int(desired_bundle_size // avg_size_per_value), 1) start = start_position while start < stop_position: end = min(start + num_values_per_split, stop_position) remaining = stop_position - end # Avoid having a too small bundle at the end. if remaining < (num_values_per_split // 4): end = stop_position sub_source = Create._create_source( self._serialized_values[start:end], self._coder) yield iobase.SourceBundle( weight=(end - start), source=sub_source, start_position=0, stop_position=(end - start)) start = end def get_range_tracker(self, start_position, stop_position): if start_position is None: start_position = 0 if stop_position is None: stop_position = len(self._serialized_values) from apache_beam import io return io.OffsetRangeTracker(start_position, stop_position) def estimate_size(self): return self._total_size ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/create_test.py0000644000175100017510000001532015206113033022731 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the Create and _CreateSource classes.""" # pytype: skip-file import logging import unittest from apache_beam import Create from apache_beam import coders from apache_beam.coders import FastPrimitivesCoder from apache_beam.internal import pickler from apache_beam.io import source_test_utils from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class CreateTest(unittest.TestCase): def setUp(self): self.coder = FastPrimitivesCoder() def test_create_transform(self): with TestPipeline() as p: assert_that(p | 'Empty' >> Create([]), equal_to([]), label='empty') assert_that(p | 'One' >> Create([None]), equal_to([None]), label='one') assert_that(p | Create(list(range(10))), equal_to(list(range(10)))) def test_create_source_read(self): self.check_read([], self.coder) self.check_read([1], self.coder) # multiple values. self.check_read(list(range(10)), self.coder) def check_read(self, values, coder): source = Create._create_source_from_iterable(values, coder) read_values = source_test_utils.read_from_source(source) self.assertEqual(sorted(values), sorted(read_values)) def test_create_source_read_with_initial_splits(self): self.check_read_with_initial_splits([], self.coder, num_splits=2) self.check_read_with_initial_splits([1], self.coder, num_splits=2) values = list(range(8)) # multiple values with a single split. self.check_read_with_initial_splits(values, self.coder, num_splits=1) # multiple values with a single split with a large desired bundle size self.check_read_with_initial_splits(values, self.coder, num_splits=0.5) # multiple values with many splits. self.check_read_with_initial_splits(values, self.coder, num_splits=3) # multiple values with uneven sized splits. self.check_read_with_initial_splits(values, self.coder, num_splits=4) # multiple values with num splits equal to num values. self.check_read_with_initial_splits( values, self.coder, num_splits=len(values)) # multiple values with num splits greater than to num values. self.check_read_with_initial_splits(values, self.coder, num_splits=30) def check_read_with_initial_splits(self, values, coder, num_splits): """A test that splits the given source into `num_splits` and verifies that the data read from original source is equal to the union of the data read from the split sources. """ source = Create._create_source_from_iterable(values, coder) desired_bundle_size = source._total_size // num_splits splits = source.split(desired_bundle_size) splits_info = [(split.source, split.start_position, split.stop_position) for split in splits] source_test_utils.assert_sources_equal_reference_source( (source, None, None), splits_info) def test_create_source_read_reentrant(self): source = Create._create_source_from_iterable(range(9), self.coder) source_test_utils.assert_reentrant_reads_succeed((source, None, None)) def test_create_source_read_reentrant_with_initial_splits(self): source = Create._create_source_from_iterable(range(24), self.coder) for split in source.split(desired_bundle_size=5): source_test_utils.assert_reentrant_reads_succeed( (split.source, split.start_position, split.stop_position)) def test_create_source_dynamic_splitting(self): # 2 values source = Create._create_source_from_iterable(range(2), self.coder) source_test_utils.assert_split_at_fraction_exhaustive(source) # Multiple values. source = Create._create_source_from_iterable(range(11), self.coder) source_test_utils.assert_split_at_fraction_exhaustive( source, perform_multi_threaded_test=True) def test_create_source_progress(self): num_values = 10 source = Create._create_source_from_iterable(range(num_values), self.coder) splits = [split for split in source.split(desired_bundle_size=100)] assert len(splits) == 1 fraction_consumed_report = [] split_points_report = [] range_tracker = splits[0].source.get_range_tracker( splits[0].start_position, splits[0].stop_position) for _ in splits[0].source.read(range_tracker): fraction_consumed_report.append(range_tracker.fraction_consumed()) split_points_report.append(range_tracker.split_points()) self.assertEqual([float(i) / num_values for i in range(num_values)], fraction_consumed_report) expected_split_points_report = [((i - 1), num_values - (i - 1)) for i in range(1, num_values + 1)] self.assertEqual(expected_split_points_report, split_points_report) def test_create_uses_coder_for_pickling(self): coders.registry.register_coder(_Unpicklable, _UnpicklableCoder) create = Create([_Unpicklable(1), _Unpicklable(2), _Unpicklable(3)]) unpickled_create = pickler.loads(pickler.dumps(create)) self.assertEqual( sorted(create.values, key=lambda v: v.value), sorted(unpickled_create.values, key=lambda v: v.value)) with self.assertRaises(NotImplementedError): # As there is no special coder for Union types, this will fall back to # FastPrimitivesCoder, which in turn falls back to pickling. create_mixed_types = Create([_Unpicklable(1), 2]) pickler.dumps(create_mixed_types) class _Unpicklable(object): def __init__(self, value): self.value = value def __eq__(self, other): return self.value == other.value def __getstate__(self): raise NotImplementedError() def __setstate__(self, state): raise NotImplementedError() class _UnpicklableCoder(coders.Coder): def encode(self, value): return str(value.value).encode() def decode(self, encoded): return _Unpicklable(int(encoded.decode())) def to_type_hint(self): return _Unpicklable if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/cy_combiners.pxd0000644000175100017510000000677315206113033023262 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # cimport cython from libc.stdint cimport int64_t, INT64_MIN, INT64_MAX cdef double _NEG_INF, _POS_INF, _NAN cdef class CountAccumulator(object): cdef readonly int64_t value cpdef add_input(self, unused_element) @cython.locals(accumulator=CountAccumulator) cpdef merge(self, accumulators) cdef class SumInt64Accumulator(object): cdef readonly int64_t value cpdef add_input(self, int64_t element) @cython.locals(accumulator=SumInt64Accumulator) cpdef merge(self, accumulators) cdef class MinInt64Accumulator(object): cdef readonly int64_t value cpdef add_input(self, int64_t element) @cython.locals(accumulator=MinInt64Accumulator) cpdef merge(self, accumulators) cdef class MaxInt64Accumulator(object): cdef readonly int64_t value cpdef add_input(self, int64_t element) @cython.locals(accumulator=MaxInt64Accumulator) cpdef merge(self, accumulators) cdef class MeanInt64Accumulator(object): cdef readonly int64_t sum cdef readonly int64_t count cpdef add_input(self, int64_t element) @cython.locals(accumulator=MeanInt64Accumulator) cpdef merge(self, accumulators) cdef class DistributionInt64Accumulator(object): cdef readonly int64_t sum cdef readonly int64_t count cdef readonly int64_t min cdef readonly int64_t max cpdef add_input(self, int64_t element) @cython.locals(accumulator=DistributionInt64Accumulator) cpdef merge(self, accumulators) cdef class SumDoubleAccumulator(object): cdef readonly double value cpdef add_input(self, double element) @cython.locals(accumulator=SumDoubleAccumulator) cpdef merge(self, accumulators) cdef class MinDoubleAccumulator(object): cdef readonly double value cpdef add_input(self, double element) @cython.locals(accumulator=MinDoubleAccumulator) cpdef merge(self, accumulators) cdef class MaxDoubleAccumulator(object): cdef readonly double value cpdef add_input(self, double element) @cython.locals(accumulator=MaxDoubleAccumulator) cpdef merge(self, accumulators) cdef class MeanDoubleAccumulator(object): cdef readonly double sum cdef readonly int64_t count cpdef add_input(self, double element) @cython.locals(accumulator=MeanDoubleAccumulator) cpdef merge(self, accumulators) cdef class AllAccumulator(object): cdef readonly bint value cpdef add_input(self, bint element) @cython.locals(accumulator=AllAccumulator) cpdef merge(self, accumulators) cdef class AnyAccumulator(object): cdef readonly bint value cpdef add_input(self, bint element) @cython.locals(accumulator=AnyAccumulator) cpdef merge(self, accumulators) cdef class ComparableValue(object): cdef readonly object value, _less_than_fn, _comparable_value cdef readonly bint requires_hydration cpdef hydrate(self, object less_than_fn, object key_fn) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/cy_combiners.py0000644000175100017510000002571115206113033023110 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A library of basic cythonized CombineFn subclasses. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file # ruff: noqa: UP006 import operator from apache_beam.transforms import core try: from apache_beam.transforms.cy_dataflow_distribution_counter import DataflowDistributionCounter except ImportError: from apache_beam.transforms.py_dataflow_distribution_counter import DataflowDistributionCounter class AccumulatorCombineFn(core.CombineFn): # singleton? def create_accumulator(self): return self._accumulator_type() @staticmethod def add_input(accumulator, element): accumulator.add_input(element) return accumulator def merge_accumulators(self, accumulators): accumulator = self._accumulator_type() accumulator.merge(accumulators) return accumulator @staticmethod def extract_output(accumulator): return accumulator.extract_output() def __eq__(self, other): return ( isinstance(other, AccumulatorCombineFn) and self._accumulator_type is other._accumulator_type) def __hash__(self): return hash(self._accumulator_type) _63 = 63 # Avoid large literals in C source code. globals()['INT64_MAX'] = 2**_63 - 1 globals()['INT64_MIN'] = -2**_63 class CountAccumulator(object): def __init__(self): self.value = 0 def add_input(self, unused_element): self.value += 1 def add_input_n(self, unused_element, n): self.value += n def merge(self, accumulators): for accumulator in accumulators: self.value += accumulator.value def extract_output(self): return self.value class SumInt64Accumulator(object): def __init__(self): self.value = 0 def add_input(self, element): global INT64_MAX, INT64_MIN # pylint: disable=global-variable-not-assigned element = int(element) if not INT64_MIN <= element <= INT64_MAX: raise OverflowError(element) self.value += element def add_input_n(self, element, n): global INT64_MAX, INT64_MIN # pylint: disable=global-variable-not-assigned element = int(element) if not INT64_MIN <= element <= INT64_MAX: raise OverflowError(element) self.value += element * n def merge(self, accumulators): for accumulator in accumulators: self.value += accumulator.value def extract_output(self): if not INT64_MIN <= self.value <= INT64_MAX: self.value %= 2**64 if self.value >= INT64_MAX: self.value -= 2**64 return self.value class MinInt64Accumulator(object): def __init__(self): self.value = INT64_MAX def add_input(self, element): element = int(element) if not INT64_MIN <= element <= INT64_MAX: raise OverflowError(element) if element < self.value: self.value = element def add_input_n(self, element, unused_n): self.add_input(element) def merge(self, accumulators): for accumulator in accumulators: if accumulator.value < self.value: self.value = accumulator.value def extract_output(self): return self.value class MaxInt64Accumulator(object): def __init__(self): self.value = INT64_MIN def add_input(self, element): element = int(element) if not INT64_MIN <= element <= INT64_MAX: raise OverflowError(element) if element > self.value: self.value = element def add_input_n(self, element, unused_n): self.add_input(element) def merge(self, accumulators): for accumulator in accumulators: if accumulator.value > self.value: self.value = accumulator.value def extract_output(self): return self.value class MeanInt64Accumulator(object): def __init__(self): self.sum = 0 self.count = 0 def add_input(self, element): element = int(element) if not INT64_MIN <= element <= INT64_MAX: raise OverflowError(element) self.sum += element self.count += 1 def add_input_n(self, element, n): element = int(element) if not INT64_MIN <= element <= INT64_MAX: raise OverflowError(element) self.sum += element * n self.count += n def merge(self, accumulators): for accumulator in accumulators: self.sum += accumulator.sum self.count += accumulator.count def extract_output(self): if not INT64_MIN <= self.sum <= INT64_MAX: self.sum %= 2**64 if self.sum >= INT64_MAX: self.sum -= 2**64 return self.sum // self.count if self.count else _NAN class DistributionInt64Accumulator(object): def __init__(self): self.sum = 0 self.count = 0 self.min = INT64_MAX self.max = INT64_MIN def add_input(self, element): element = int(element) if not INT64_MIN <= element <= INT64_MAX: raise OverflowError(element) self.sum += element self.count += 1 self.min = min(self.min, element) self.max = max(self.max, element) def add_input_n(self, element, n): element = int(element) if not INT64_MIN <= element <= INT64_MAX: raise OverflowError(element) self.sum += element * n self.count += n self.min = min(self.min, element) self.max = max(self.max, element) def merge(self, accumulators): for accumulator in accumulators: self.sum += accumulator.sum self.count += accumulator.count self.min = min(self.min, accumulator.min) self.max = max(self.max, accumulator.max) def extract_output(self): if not INT64_MIN <= self.sum <= INT64_MAX: self.sum %= 2**64 if self.sum >= INT64_MAX: self.sum -= 2**64 mean = self.sum // self.count if self.count else _NAN return mean, self.sum, self.count, self.min, self.max class CountCombineFn(AccumulatorCombineFn): _accumulator_type = CountAccumulator class SumInt64Fn(AccumulatorCombineFn): _accumulator_type = SumInt64Accumulator class MinInt64Fn(AccumulatorCombineFn): _accumulator_type = MinInt64Accumulator class MaxInt64Fn(AccumulatorCombineFn): _accumulator_type = MaxInt64Accumulator class MeanInt64Fn(AccumulatorCombineFn): _accumulator_type = MeanInt64Accumulator class DistributionInt64Fn(AccumulatorCombineFn): _accumulator_type = DistributionInt64Accumulator _POS_INF = float('inf') _NEG_INF = float('-inf') _NAN = float('nan') class SumDoubleAccumulator(object): def __init__(self): self.value = 0 def add_input(self, element): element = float(element) self.value += element def merge(self, accumulators): for accumulator in accumulators: self.value += accumulator.value def extract_output(self): return self.value class MinDoubleAccumulator(object): def __init__(self): self.value = _POS_INF def add_input(self, element): element = float(element) if element < self.value: self.value = element def merge(self, accumulators): for accumulator in accumulators: if accumulator.value < self.value: self.value = accumulator.value def extract_output(self): return self.value class MaxDoubleAccumulator(object): def __init__(self): self.value = _NEG_INF def add_input(self, element): element = float(element) if element > self.value: self.value = element def merge(self, accumulators): for accumulator in accumulators: if accumulator.value > self.value: self.value = accumulator.value def extract_output(self): return self.value class MeanDoubleAccumulator(object): def __init__(self): self.sum = 0 self.count = 0 def add_input(self, element): element = float(element) self.sum += element self.count += 1 def merge(self, accumulators): for accumulator in accumulators: self.sum += accumulator.sum self.count += accumulator.count def extract_output(self): return self.sum // self.count if self.count else _NAN class SumFloatFn(AccumulatorCombineFn): _accumulator_type = SumDoubleAccumulator class MinFloatFn(AccumulatorCombineFn): _accumulator_type = MinDoubleAccumulator class MaxFloatFn(AccumulatorCombineFn): _accumulator_type = MaxDoubleAccumulator class MeanFloatFn(AccumulatorCombineFn): _accumulator_type = MeanDoubleAccumulator class AllAccumulator(object): def __init__(self): self.value = True def add_input(self, element): self.value &= not not element def merge(self, accumulators): for accumulator in accumulators: self.value &= accumulator.value def extract_output(self): return self.value class AnyAccumulator(object): def __init__(self): self.value = False def add_input(self, element): self.value |= not not element def merge(self, accumulators): for accumulator in accumulators: self.value |= accumulator.value def extract_output(self): return self.value class AnyCombineFn(AccumulatorCombineFn): _accumulator_type = AnyAccumulator class AllCombineFn(AccumulatorCombineFn): _accumulator_type = AllAccumulator class DataflowDistributionCounterFn(AccumulatorCombineFn): """A subclass of cy_combiners.AccumulatorCombineFn. Make DataflowDistributionCounter able to report to Dataflow service via CounterFactory. When cythonized DataflowDistributinoCounter available, make CounterFn combine with cythonized module, otherwise, combine with python version. """ _accumulator_type = DataflowDistributionCounter class ComparableValue(object): """A way to allow comparing elements in a rich fashion.""" __slots__ = ( 'value', '_less_than_fn', '_comparable_value', 'requires_hydration') def __init__(self, value, less_than_fn, key_fn, _requires_hydration=False): self.value = value self.hydrate(less_than_fn, key_fn) self.requires_hydration = _requires_hydration def hydrate(self, less_than_fn, key_fn): self._less_than_fn = less_than_fn if less_than_fn else operator.lt self._comparable_value = key_fn(self.value) if key_fn else self.value self.requires_hydration = False def __lt__(self, other): assert not self.requires_hydration assert self._less_than_fn is other._less_than_fn return self._less_than_fn(self._comparable_value, other._comparable_value) def __repr__(self): return 'ComparableValue[%s]' % str(self.value) def __reduce__(self): # Since we can't pickle the Compare and Key Fn we pass None and we signify # that this object _requires_hydration. return ComparableValue, (self.value, None, None, True) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/cy_dataflow_distribution_counter.pxd0000644000175100017510000000332015206113033027421 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ For internal use only. No backwards compatibility guarantees.""" cimport cython from libc.stdint cimport int64_t # 3 buckets for every power of ten -> 1, 2, 5 cdef enum: BUCKET_PER_TEN = 3 # Assume the max input is max(int64_t), then the possible max bucket size is 59 cdef enum: MAX_BUCKET_SIZE = 59 cdef class DataflowDistributionCounter(object): cdef public int64_t min cdef public int64_t max cdef public int64_t count cdef public int64_t sum cdef int64_t* buckets cdef public bint is_cythonized cpdef bint add_input(self, int64_t element) except -1 cpdef bint add_input_n(self, int64_t element, int64_t n) except -1 cdef int64_t _fast_calculate_bucket_index(self, int64_t element) cpdef void translate_to_histogram(self, histogram) cpdef bint add_inputs_for_test(self, elements) except -1 cpdef int64_t calculate_bucket_index(self, int64_t element) cpdef tuple extract_output(self) cpdef merge(self, accumulators) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/cy_dataflow_distribution_counter.pyx0000644000175100017510000001245715206113033027461 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ For internal use only. No backwards compatibility guarantees.""" cimport cython from libc.stdint cimport int64_t, INT64_MAX from libc.stdlib cimport calloc, free cdef unsigned long long* POWER_TEN = [10e-1, 10e0, 10e1, 10e2, 10e3, 10e4, 10e5, 10e6, 10e7, 10e8, 10e9, 10e10, 10e11, 10e12, 10e13, 10e14, 10e15, 10e16, 10e17, 10e18] cdef int64_t get_log10_round_to_floor(int64_t element): cdef int power = 0 while element >= POWER_TEN[power]: power += 1 return power - 1 cdef class DataflowDistributionCounter(object): """Distribution Counter: Contains value distribution statistics and methods for incrementing. Currently using special bucketing strategy suitable for Dataflow Attributes: min: minimum value of all inputs. max: maximum value of all inputs. count: total count of all inputs. sum: sum of all inputs. buckets: histogram buckets of value counts for a distribution(1,2,5 bucketing). Max bucket_index is 58( sys.maxint as input). is_cythonized: mark whether DataflowDistributionCounter cythonized. """ def __init__(self): self.min = INT64_MAX self.max = 0 self.count = 0 self.sum = 0 self.buckets = calloc(MAX_BUCKET_SIZE, sizeof(int64_t)) self.is_cythonized = True def __dealloc__(self): """free allocated memory""" free(self.buckets) cpdef bint add_input(self, int64_t element) except -1: if element < 0: raise ValueError('Distribution counters support only non-negative value') self.min = min(self.min, element) self.max = max(self.max, element) self.count += 1 self.sum += element cdef int64_t bucket_index = self._fast_calculate_bucket_index(element) self.buckets[bucket_index] += 1 cpdef bint add_input_n(self, int64_t element, int64_t n) except -1: if element < 0: raise ValueError('Distribution counters support only non-negative value') self.min = min(self.min, element) self.max = max(self.max, element) self.count += n self.sum += element*n cdef int64_t bucket_index = self._fast_calculate_bucket_index(element) self.buckets[bucket_index] += n cdef int64_t _fast_calculate_bucket_index(self, int64_t element): """Calculate the bucket index for the given element. Declare calculate_bucket_index as cdef in order to improve performance, since cpdef will have significant overhead. """ if element == 0: return 0 cdef int64_t log10_floor = get_log10_round_to_floor(element) cdef int64_t power_of_ten = POWER_TEN[log10_floor] cdef int64_t bucket_offset = 0 if element < power_of_ten * 2: bucket_offset = 0 elif element < power_of_ten * 5: bucket_offset = 1 else: bucket_offset = 2 return 1 + log10_floor * BUCKET_PER_TEN + bucket_offset cpdef void translate_to_histogram(self, histogram): """Translate buckets into Histogram. Args: histogram: apache_beam.runners.dataflow.internal.clents.dataflow.Histogram Ideally, only call this function when reporting counter to dataflow service. """ cdef int first_bucket_offset = 0 cdef int last_bucket_offset = 0 cdef int index = 0 for index in range(0, MAX_BUCKET_SIZE): if self.buckets[index] != 0: first_bucket_offset = index break for index in range(MAX_BUCKET_SIZE - 1, -1, -1): if self.buckets[index] != 0: last_bucket_offset = index break histogram.firstBucketOffset = first_bucket_offset histogram.bucketCounts = [] for index in range(first_bucket_offset, last_bucket_offset + 1): histogram.bucketCounts.append(self.buckets[index]) cpdef bint add_inputs_for_test(self, elements) except -1: """Used for performance microbenchmark. During runtime, add_input will be called through c-call, so we want to have the same calling routine when running microbenchmark as application runtime. Directly calling cpdef from def will cause significant overhead. """ for element in elements: self.add_input(element) cpdef int64_t calculate_bucket_index(self, int64_t element): """Used for unit tests. cdef calculate_bucket_index cannot be called directly from def. """ return self._fast_calculate_bucket_index(element) cpdef tuple extract_output(self): mean = self.sum // self.count if self.count else float('nan') return mean, self.sum, self.count, self.min, self.max cpdef merge(self, accumulators): raise NotImplementedError() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/dataflow_distribution_counter_test.py0000644000175100017510000000552215206113033027630 0ustar00runnerrunner# the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for DataflowDistributionCounter When Cython is available, unit tests will test on cythonized module, otherwise, test on pure python module """ # pytype: skip-file import unittest from mock import Mock from apache_beam.transforms import DataflowDistributionCounter INT64_MAX = 2**63 - 1 class DataflowDistributionAccumulatorTest(unittest.TestCase): def test_calculate_bucket_index_with_input_0(self): counter = DataflowDistributionCounter() index = counter.calculate_bucket_index(0) self.assertEqual(index, 0) def test_calculate_bucket_index_within_max_long(self): counter = DataflowDistributionCounter() bucket = 1 power_of_ten = 1 while power_of_ten <= INT64_MAX: for multiplier in [1, 2, 5]: value = multiplier * power_of_ten actual_bucket = counter.calculate_bucket_index(value - 1) self.assertEqual(actual_bucket, bucket - 1) bucket += 1 power_of_ten *= 10 def test_add_input(self): counter = DataflowDistributionCounter() expected_buckets = [1, 3, 0, 0, 0, 0, 0, 0, 1, 1] expected_sum = 1510 expected_first_bucket_index = 1 expected_count = 6 expected_min = 1 expected_max = 1000 for element in [1, 500, 2, 3, 1000, 4]: counter.add_input(element) histogram = Mock(firstBucketOffset=None, bucketCounts=None) counter.translate_to_histogram(histogram) self.assertEqual(counter.sum, expected_sum) self.assertEqual(counter.count, expected_count) self.assertEqual(counter.min, expected_min) self.assertEqual(counter.max, expected_max) self.assertEqual(histogram.firstBucketOffset, expected_first_bucket_index) self.assertEqual(histogram.bucketCounts, expected_buckets) def test_translate_to_histogram_with_input_0(self): counter = DataflowDistributionCounter() counter.add_input(0) histogram = Mock(firstBucketOffset=None, bucketCounts=None) counter.translate_to_histogram(histogram) self.assertEqual(histogram.firstBucketOffset, 0) self.assertEqual(histogram.bucketCounts, [1]) def test_translate_to_histogram_with_max_input(self): counter = DataflowDistributionCounter() counter.add_input(INT64_MAX) histogram = Mock(firstBucketOffset=None, bucketCounts=None) counter.translate_to_histogram(histogram) self.assertEqual(histogram.firstBucketOffset, 57) self.assertEqual(histogram.bucketCounts, [1]) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/deduplicate.py0000644000175100017510000001201415206113033022707 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file """a collection of ptransforms for deduplicating elements.""" import typing from apache_beam import typehints from apache_beam.coders.coders import BooleanCoder from apache_beam.transforms import core from apache_beam.transforms import ptransform from apache_beam.transforms import userstate from apache_beam.transforms.timeutil import TimeDomain from apache_beam.utils import timestamp __all__ = [ 'Deduplicate', 'DeduplicatePerKey', ] K = typing.TypeVar('K') V = typing.TypeVar('V') @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types(tuple[K, V]) class DeduplicatePerKey(ptransform.PTransform): """ A PTransform which deduplicates pair over a time domain and threshold. Values in different windows will NOT be considered duplicates of each other. Deduplication is guaranteed with respect of time domain and duration. Time durations are required so as to avoid unbounded memory and/or storage requirements within a runner and care might need to be used to ensure that the deduplication time limit is long enough to remove duplicates but short enough to not cause performance problems within a runner. Each runner may provide an optimized implementation of their choice using the deduplication time domain and threshold specified. Does not preserve any order the input PCollection might have had. """ def __init__(self, processing_time_duration=None, event_time_duration=None): if processing_time_duration is None and event_time_duration is None: raise ValueError( 'DeduplicatePerKey requires at lease provide either' 'processing_time_duration or event_time_duration.') self.processing_time_duration = processing_time_duration self.event_time_duration = event_time_duration def _create_deduplicate_fn(self): processing_timer_spec = userstate.TimerSpec( 'processing_timer', TimeDomain.REAL_TIME) event_timer_spec = userstate.TimerSpec('event_timer', TimeDomain.WATERMARK) state_spec = userstate.BagStateSpec('seen', BooleanCoder()) processing_time_duration = self.processing_time_duration event_time_duration = self.event_time_duration class DeduplicationFn(core.DoFn): def process( self, kv, ts=core.DoFn.TimestampParam, seen_state=core.DoFn.StateParam(state_spec), processing_timer=core.DoFn.TimerParam(processing_timer_spec), event_timer=core.DoFn.TimerParam(event_timer_spec)): if True in seen_state.read(): return if processing_time_duration is not None: processing_timer.set( timestamp.Timestamp.now() + processing_time_duration) if event_time_duration is not None: event_timer.set(ts + event_time_duration) seen_state.add(True) yield kv @userstate.on_timer(processing_timer_spec) def process_processing_timer( self, seen_state=core.DoFn.StateParam(state_spec)): seen_state.clear() @userstate.on_timer(event_timer_spec) def process_event_timer( self, seen_state=core.DoFn.StateParam(state_spec)): seen_state.clear() return DeduplicationFn() def expand(self, pcoll): return ( pcoll | 'DeduplicateFn' >> core.ParDo(self._create_deduplicate_fn())) class Deduplicate(ptransform.PTransform): """Similar to DeduplicatePerKey, the Deduplicate transform takes any arbitrary value as input and uses value as key to deduplicate among certain amount of time duration. """ def __init__(self, processing_time_duration=None, event_time_duration=None): if processing_time_duration is None and event_time_duration is None: raise ValueError( 'Deduplicate requires at least providing either ' 'processing_time_duration or event_time_duration.') self.processing_time_duration = processing_time_duration self.event_time_duration = event_time_duration def expand(self, pcoll): return ( pcoll | 'Use Value as Key' >> core.Map(lambda x: (x, None)) | 'DeduplicatePerKey' >> DeduplicatePerKey( processing_time_duration=self.processing_time_duration, event_time_duration=self.event_time_duration) | 'Output Value' >> core.Map(lambda kv: kv[0])) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/deduplicate_test.py0000644000175100017510000001520215206113033023750 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file """Unit tests for deduplicate transform by using TestStream.""" import unittest import pytest import apache_beam as beam from apache_beam.coders import coders from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.testing.util import equal_to_per_window from apache_beam.transforms import deduplicate from apache_beam.transforms import window from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import Timestamp # TestStream is only supported in streaming pipeline. The Deduplicate transform # also requires Timer support. Sickbaying this testsuite until dataflow runner # supports both TestStream and user timer. @pytest.mark.no_sickbay_batch @pytest.mark.no_sickbay_streaming @pytest.mark.it_validatesrunner class DeduplicateTest(unittest.TestCase): def __init__(self, *args, **kwargs): self.runner = None self.options = None super().__init__(*args, **kwargs) def set_runner(self, runner): self.runner = runner def set_options(self, options): self.options = options def create_pipeline(self): if self.runner and self.options: return TestPipeline(runner=self.runner, options=self.options) elif self.runner: return TestPipeline(runner=self.runner) elif self.options: return TestPipeline(options=self.options) else: return TestPipeline() def test_deduplication_in_different_windows(self): with self.create_pipeline() as p: test_stream = ( TestStream( coder=coders.StrUtf8Coder()).advance_watermark_to(0).add_elements( [ window.TimestampedValue('k1', 0), window.TimestampedValue('k2', 10), window.TimestampedValue('k3', 20), window.TimestampedValue('k1', 30), window.TimestampedValue('k2', 40), window.TimestampedValue('k3', 50), window.TimestampedValue('k4', 60), window.TimestampedValue('k5', 70), window.TimestampedValue('k6', 80) ]).advance_watermark_to_infinity()) res = ( p | test_stream | beam.WindowInto(window.FixedWindows(30)) | deduplicate.Deduplicate(processing_time_duration=10 * 60) | beam.Map(lambda e, ts=beam.DoFn.TimestampParam: (e, ts))) # Deduplication should happen per window. expect_unique_keys_per_window = { window.IntervalWindow(0, 30): [('k1', Timestamp(0)), ('k2', Timestamp(10)), ('k3', Timestamp(20))], window.IntervalWindow(30, 60): [('k1', Timestamp(30)), ('k2', Timestamp(40)), ('k3', Timestamp(50))], window.IntervalWindow(60, 90): [('k4', Timestamp(60)), ('k5', Timestamp(70)), ('k6', Timestamp(80))], } assert_that( res, equal_to_per_window(expect_unique_keys_per_window), use_global_window=False, label='assert per window') @unittest.skip('TestStream not yet supported') def test_deduplication_with_event_time(self): deduplicate_duration = 60 with self.create_pipeline() as p: test_stream = ( TestStream(coder=coders.StrUtf8Coder()).with_output_types( str).advance_watermark_to(0).add_elements([ window.TimestampedValue('k1', 0), window.TimestampedValue('k2', 20), window.TimestampedValue('k3', 30) ]).advance_watermark_to(30).add_elements([ window.TimestampedValue('k1', 40), window.TimestampedValue('k2', 50), window.TimestampedValue('k3', 60) ]).advance_watermark_to(deduplicate_duration).add_elements([ window.TimestampedValue('k1', 70) ]).advance_watermark_to_infinity()) res = ( p | test_stream | deduplicate.Deduplicate( event_time_duration=Duration(deduplicate_duration)) | beam.Map(lambda e, ts=beam.DoFn.TimestampParam: (e, ts))) assert_that( res, equal_to([('k1', Timestamp(0)), ('k2', Timestamp(20)), ('k3', Timestamp(30)), ('k1', Timestamp(70))])) @unittest.skip('TestStream not yet supported') def test_deduplication_with_processing_time(self): deduplicate_duration = 60 with self.create_pipeline() as p: test_stream = ( TestStream(coder=coders.StrUtf8Coder()).with_output_types( str).advance_watermark_to(0).add_elements([ window.TimestampedValue('k1', 0), window.TimestampedValue('k2', 20), window.TimestampedValue('k3', 30) ]).advance_processing_time(30).add_elements([ window.TimestampedValue('k1', 40), window.TimestampedValue('k2', 50), window.TimestampedValue('k3', 60) ]).advance_processing_time(deduplicate_duration).add_elements([ window.TimestampedValue('k1', 70) ]).advance_watermark_to_infinity()) res = ( p | test_stream | deduplicate.Deduplicate( processing_time_duration=Duration(deduplicate_duration)) | beam.Map(lambda e, ts=beam.DoFn.TimestampParam: (e, ts))) assert_that( res, equal_to([('k1', Timestamp(0)), ('k2', Timestamp(20)), ('k3', Timestamp(30)), ('k1', Timestamp(70))])) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/display.py0000644000175100017510000003505215206113033022100 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ :class:`DisplayData`, its classes, interfaces and methods. The classes in this module allow users and transform developers to define static display data to be displayed when a pipeline runs. :class:`~apache_beam.transforms.ptransform.PTransform` s, :class:`~apache_beam.transforms.core.DoFn` s and other pipeline components are subclasses of the :class:`HasDisplayData` mixin. To add static display data to a component, you can override the :meth:`HasDisplayData.display_data()` method. Available classes: * :class:`HasDisplayData` - Components that inherit from this class can have static display data shown in the UI. * :class:`DisplayDataItem` - This class represents static display data elements. * :class:`DisplayData` - Internal class that is used to create display data and communicate it to the API. """ # pytype: skip-file import calendar import inspect import json from datetime import datetime from datetime import timedelta from typing import TYPE_CHECKING from typing import List from typing import Optional from typing import Union from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 if TYPE_CHECKING: from apache_beam.options.pipeline_options import PipelineOptions __all__ = ['HasDisplayData', 'DisplayDataItem', 'DisplayData'] class HasDisplayData(object): """ Basic mixin for elements that contain display data. It implements only the display_data method and a _get_display_data_namespace method. """ def display_data(self): # type: () -> dict """ Returns the display data associated to a pipeline component. It should be reimplemented in pipeline components that wish to have static display data. Returns: Dict[str, Any]: A dictionary containing ``key:value`` pairs. The value might be an integer, float or string value; a :class:`DisplayDataItem` for values that have more data (e.g. short value, label, url); or a :class:`HasDisplayData` instance that has more display data that should be picked up. For example:: { 'key1': 'string_value', 'key2': 1234, 'key3': 3.14159265, 'key4': DisplayDataItem('apache.org', url='http://apache.org'), 'key5': subComponent } """ return {} def _get_display_data_namespace(self): # type: () -> str return '{}.{}'.format(self.__module__, self.__class__.__name__) class DisplayData(object): """ Static display data associated with a pipeline component. """ def __init__( self, namespace, # type: str display_data_dict # type: dict ): # type: (...) -> None self.namespace = namespace self.items = [ ] # type: List[Union[DisplayDataItem, beam_runner_api_pb2.DisplayData]] self._populate_items(display_data_dict) def _populate_items(self, display_data_dict): """ Populates the list of display data items. """ for key, element in display_data_dict.items(): if isinstance(element, HasDisplayData): subcomponent_display_data = DisplayData( element._get_display_data_namespace(), element.display_data()) self.items += subcomponent_display_data.items elif isinstance(element, DisplayDataItem): if element.should_drop(): continue element.key = key element.namespace = self.namespace self.items.append(element) elif isinstance(element, beam_runner_api_pb2.DisplayData): self.items.append(element) else: # If it's not a HasDisplayData element, # nor a dictionary, then it's a simple value self.items.append( DisplayDataItem(element, namespace=self.namespace, key=key)) def to_proto(self): # type: (...) -> List[beam_runner_api_pb2.DisplayData] """Returns a List of Beam proto representation of Display data.""" def create_payload(dd) -> Optional[beam_runner_api_pb2.LabelledPayload]: try: display_data_dict = dd.get_dict() except ValueError: # Skip if the display data is invalid. return None # We use 'label' or 'key' properties to populate the 'label' attribute of # 'LabelledPayload'. 'label' is a better choice since it's expected to be # more human readable but some transforms, sources, etc. may not set a # 'label' property when configuring DisplayData. label = ( display_data_dict['label'] if 'label' in display_data_dict else display_data_dict['key']) value = display_data_dict['value'] if isinstance(value, str): return beam_runner_api_pb2.LabelledPayload( label=label, string_value=value, key=display_data_dict['key'], namespace=display_data_dict.get('namespace', '')) elif isinstance(value, bool): return beam_runner_api_pb2.LabelledPayload( label=label, bool_value=value, key=display_data_dict['key'], namespace=display_data_dict.get('namespace', '')) elif isinstance(value, int): return beam_runner_api_pb2.LabelledPayload( label=label, int_value=value, key=display_data_dict['key'], namespace=display_data_dict.get('namespace', '')) elif isinstance(value, (float, complex)): return beam_runner_api_pb2.LabelledPayload( label=label, double_value=value, # type: ignore[arg-type] key=display_data_dict['key'], namespace=display_data_dict.get('namespace', '')) else: raise ValueError( 'Unsupported type %s for value of display data %s' % (type(value), label)) dd_protos = [] for dd in self.items: if isinstance(dd, beam_runner_api_pb2.DisplayData): dd_protos.append(dd) else: dd_payload = create_payload(dd) if dd_payload: dd_protos.append( beam_runner_api_pb2.DisplayData( urn=common_urns.StandardDisplayData.DisplayData.LABELLED.urn, payload=dd_payload.SerializeToString())) return dd_protos @classmethod def create_from_options(cls, pipeline_options): """ Creates :class:`~apache_beam.transforms.display.DisplayData` from a :class:`~apache_beam.options.pipeline_options.PipelineOptions` instance. When creating :class:`~apache_beam.transforms.display.DisplayData`, this method will convert the value of any item of a non-supported type to its string representation. The normal :meth:`.create_from()` method rejects those items. Returns: ~apache_beam.transforms.display.DisplayData: A :class:`~apache_beam.transforms.display.DisplayData` instance with populated items. Raises: ValueError: If the **has_display_data** argument is not an instance of :class:`HasDisplayData`. """ from apache_beam.options.pipeline_options import PipelineOptions if not isinstance(pipeline_options, PipelineOptions): raise ValueError( 'Element of class {}.{} does not subclass PipelineOptions'.format( pipeline_options.__module__, pipeline_options.__class__.__name__)) items = { k: (v if DisplayDataItem._get_value_type(v) is not None else str(v)) for k, v in pipeline_options.display_data().items() } return cls(pipeline_options._get_display_data_namespace(), items) @classmethod def create_from(cls, has_display_data, extra_items=None): """ Creates :class:`~apache_beam.transforms.display.DisplayData` from a :class:`HasDisplayData` instance. Returns: ~apache_beam.transforms.display.DisplayData: A :class:`~apache_beam.transforms.display.DisplayData` instance with populated items. Raises: ValueError: If the **has_display_data** argument is not an instance of :class:`HasDisplayData`. """ if not isinstance(has_display_data, HasDisplayData): raise ValueError( 'Element of class {}.{} does not subclass HasDisplayData'.format( has_display_data.__module__, has_display_data.__class__.__name__)) if extra_items is None: extra_items = {} return cls( has_display_data._get_display_data_namespace(), dict(**has_display_data.display_data(), **extra_items)) class DisplayDataItem(object): """ A DisplayDataItem represents a unit of static display data. Each item is identified by a key and the namespace of the component the display item belongs to. """ typeDict = { str: 'STRING', int: 'INTEGER', float: 'FLOAT', bool: 'BOOLEAN', timedelta: 'DURATION', datetime: 'TIMESTAMP' } def __init__( self, value, url=None, label=None, namespace=None, key=None, shortValue=None): self.namespace = namespace self.key = key self.type = self._get_value_type(value) self.shortValue = ( shortValue if shortValue is not None else self._get_short_value( value, self.type)) self.value = value self.url = url self.label = label self._drop_if_none = False self._drop_if_default = False def drop_if_none(self): # type: () -> DisplayDataItem """ The item should be dropped if its value is None. Returns: Returns self. """ self._drop_if_none = True return self def drop_if_default(self, default): # type: (...) -> DisplayDataItem """ The item should be dropped if its value is equal to its default. Returns: Returns self. """ self._default = default self._drop_if_default = True return self def should_drop(self): # type: () -> bool """ Return True if the item should be dropped, or False if it should not be dropped. This depends on the drop_if_none, and drop_if_default calls. Returns: True or False; depending on whether the item should be dropped or kept. """ if self._drop_if_none and self.value is None: return True if self._drop_if_default and self.value == self._default: return True return False def is_valid(self): # type: () -> None """ Checks that all the necessary fields of the :class:`DisplayDataItem` are filled in. It checks that neither key, namespace, value or type are :data:`None`. Raises: ValueError: If the item does not have a key, namespace, value or type. """ if self.key is None: raise ValueError( 'Invalid DisplayDataItem %s. Key must not be None.' % self) if self.namespace is None: raise ValueError( 'Invalid DisplayDataItem %s. Namespace must not be None' % self) if self.value is None: raise ValueError( 'Invalid DisplayDataItem %s. Value must not be None' % self) if self.type is None: raise ValueError( 'Invalid DisplayDataItem. Value {} is of an unsupported type.'.format( self.value)) def _get_dict(self): res = { 'key': self.key, 'namespace': self.namespace, 'type': self.type if self.type != 'CLASS' else 'STRING' } # TODO: Python Class types should not be special-cased once # the Fn API is in. if self.url is not None: res['url'] = self.url if self.shortValue is not None: res['shortValue'] = self.shortValue if self.label is not None: res['label'] = self.label res['value'] = self._format_value(self.value, self.type) return res def get_dict(self): # type: () -> dict """ Returns the internal-API dictionary representing the :class:`DisplayDataItem`. Returns: Dict[str, Any]: A dictionary. The internal-API dictionary representing the :class:`DisplayDataItem`. Raises: ValueError: if the item is not valid. """ self.is_valid() return self._get_dict() def __repr__(self): return 'DisplayDataItem({})'.format(json.dumps(self._get_dict())) def __eq__(self, other): if isinstance(other, self.__class__): return self._get_dict() == other._get_dict() return False def __hash__(self): return hash(tuple(sorted(self._get_dict().items()))) @classmethod def _format_value(cls, value, type_): """ Returns the API representation of a value given its type. Args: value: The value of the item that needs to be shortened. type_(string): The type of the value. Returns: A formatted value in the form of a float, int, or string. """ res = value if type_ == 'CLASS': res = '{}.{}'.format(value.__module__, value.__name__) elif type_ == 'DURATION': res = value.total_seconds() * 1000 elif type_ == 'TIMESTAMP': res = calendar.timegm( value.timetuple()) * 1000 + value.microsecond // 1000 return res @classmethod def _get_short_value(cls, value, type_): """ Calculates the short value for an item. Args: value: The value of the item that needs to be shortened. type_(string): The type of the value. Returns: The unqualified name of a class if type_ is 'CLASS'. None otherwise. """ if type_ == 'CLASS': return value.__name__ return None @classmethod def _get_value_type(cls, value): """ Infers the type of a given value. Args: value: The value whose type needs to be inferred. For 'DURATION' and 'TIMESTAMP', the corresponding Python type is datetime.timedelta and datetime.datetime respectively. For Python classes, the API type is just 'STRING' at the moment. Returns: One of 'STRING', 'INTEGER', 'FLOAT', 'CLASS', 'DURATION', or 'TIMESTAMP', depending on the type of the value. """ #TODO: Fix Args: documentation once the Python classes handling has changed type_ = cls.typeDict.get(type(value)) if type_ is None: type_ = 'CLASS' if inspect.isclass(value) else None if type_ is None and value is None: type_ = 'STRING' return type_ ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/display_test.py0000644000175100017510000002176715206113033023147 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the DisplayData API.""" # pytype: skip-file import unittest from datetime import datetime # pylint: disable=ungrouped-imports import hamcrest as hc from hamcrest.core.base_matcher import BaseMatcher import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.transforms.display import DisplayData from apache_beam.transforms.display import DisplayDataItem from apache_beam.transforms.display import HasDisplayData # pylint: enable=ungrouped-imports class DisplayDataItemMatcher(BaseMatcher): """ Matcher class for DisplayDataItems in unit tests. """ IGNORED = object() def __init__( self, key=IGNORED, value=IGNORED, namespace=IGNORED, label=IGNORED, shortValue=IGNORED): if all(member == DisplayDataItemMatcher.IGNORED for member in [key, value, namespace, label, shortValue]): raise ValueError('Must receive at least one item attribute to match') self.key = key self.value = value self.namespace = namespace self.label = label self.shortValue = shortValue def _matches(self, item): if self.key != DisplayDataItemMatcher.IGNORED and item.key != self.key: return False if (self.namespace != DisplayDataItemMatcher.IGNORED and item.namespace != self.namespace): return False if (self.value != DisplayDataItemMatcher.IGNORED and item.value != self.value): return False if (self.label != DisplayDataItemMatcher.IGNORED and item.label != self.label): return False if (self.shortValue != DisplayDataItemMatcher.IGNORED and item.shortValue != self.shortValue): return False return True def describe_to(self, description): descriptors = [] if self.key != DisplayDataItemMatcher.IGNORED: descriptors.append('key is {}'.format(self.key)) if self.value != DisplayDataItemMatcher.IGNORED: descriptors.append('value is {}'.format(self.value)) if self.namespace != DisplayDataItemMatcher.IGNORED: descriptors.append('namespace is {}'.format(self.namespace)) if self.label != DisplayDataItemMatcher.IGNORED: descriptors.append('label is {}'.format(self.label)) if self.shortValue != DisplayDataItemMatcher.IGNORED: descriptors.append('shortValue is {}'.format(self.shortValue)) item_description = '{}'.format(' and '.join(descriptors)) description.append(item_description) class DisplayDataTest(unittest.TestCase): def test_display_data_item_matcher(self): with self.assertRaises(ValueError): DisplayDataItemMatcher() def test_inheritance_ptransform(self): class MyTransform(beam.PTransform): pass display_pt = MyTransform() # PTransform inherits from HasDisplayData. self.assertTrue(isinstance(display_pt, HasDisplayData)) self.assertEqual(display_pt.display_data(), {}) def test_inheritance_dofn(self): class MyDoFn(beam.DoFn): pass display_dofn = MyDoFn() self.assertTrue(isinstance(display_dofn, HasDisplayData)) self.assertEqual(display_dofn.display_data(), {}) def test_unsupported_type_display_data(self): class MyDisplayComponent(HasDisplayData): def display_data(self): return {'item_key': 'item_value'} with self.assertRaises(ValueError): DisplayData.create_from_options(MyDisplayComponent()) def test_value_provider_display_data(self): class TestOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_value_provider_argument( '--int_flag', type=int, help='int_flag description') parser.add_value_provider_argument( '--str_flag', type=str, default='hello', help='str_flag description') parser.add_value_provider_argument( '--float_flag', type=float, help='float_flag description') options = TestOptions(['--int_flag', '1']) items = DisplayData.create_from_options(options).items expected_items = [ DisplayDataItemMatcher('int_flag', '1'), DisplayDataItemMatcher( 'str_flag', 'RuntimeValueProvider(option: str_flag,' ' type: str, default_value: \'hello\')'), DisplayDataItemMatcher( 'float_flag', 'RuntimeValueProvider(option: float_flag,' ' type: float, default_value: None)') ] hc.assert_that(items, hc.has_items(*expected_items)) def test_create_list_display_data(self): flags = ['--extra_package', 'package1', '--extra_package', 'package2'] pipeline_options = PipelineOptions(flags=flags) items = DisplayData.create_from_options(pipeline_options).items hc.assert_that( items, hc.has_items( DisplayDataItemMatcher( 'extra_packages', str(['package1', 'package2'])))) def test_unicode_type_display_data(self): class MyDoFn(beam.DoFn): def display_data(self): return { 'unicode_string': 'my string', 'unicode_literal_string': 'my literal string' } fn = MyDoFn() dd = DisplayData.create_from(fn) for item in dd.items: self.assertEqual(item.type, 'STRING') def test_base_cases(self): """ Tests basic display data cases (key:value, key:dict) It does not test subcomponent inclusion """ class MyDoFn(beam.DoFn): def __init__(self, my_display_data=None): self.my_display_data = my_display_data def process(self, element): yield element + 1 def display_data(self): return { 'static_integer': 120, 'static_string': 'static me!', 'complex_url': DisplayDataItem( 'github.com', url='http://github.com', label='The URL'), 'python_class': HasDisplayData, 'my_dd': self.my_display_data } now = datetime.now() fn = MyDoFn(my_display_data=now) dd = DisplayData.create_from(fn) nspace = '{}.{}'.format(fn.__module__, fn.__class__.__name__) expected_items = [ DisplayDataItemMatcher( key='complex_url', value='github.com', namespace=nspace, label='The URL'), DisplayDataItemMatcher(key='my_dd', value=now, namespace=nspace), DisplayDataItemMatcher( key='python_class', value=HasDisplayData, namespace=nspace, shortValue='HasDisplayData'), DisplayDataItemMatcher( key='static_integer', value=120, namespace=nspace), DisplayDataItemMatcher( key='static_string', value='static me!', namespace=nspace) ] hc.assert_that(dd.items, hc.has_items(*expected_items)) expected_items.append( DisplayDataItemMatcher( key='extra_key', value='extra_value', namespace=nspace)) hc.assert_that( DisplayData.create_from(fn, extra_items={ 'extra_key': 'extra_value' }).items, hc.has_items(*expected_items)) def test_drop_if_none(self): class MyDoFn(beam.DoFn): def display_data(self): return { 'some_val': DisplayDataItem('something').drop_if_none(), 'non_val': DisplayDataItem(None).drop_if_none(), 'def_val': DisplayDataItem(True).drop_if_default(True), 'nodef_val': DisplayDataItem(True).drop_if_default(False) } dd = DisplayData.create_from(MyDoFn()) expected_items = [ DisplayDataItemMatcher('some_val', 'something'), DisplayDataItemMatcher('nodef_val', True) ] hc.assert_that(dd.items, hc.has_items(*expected_items)) def test_subcomponent(self): class SpecialDoFn(beam.DoFn): def display_data(self): return {'dofn_value': 42} dofn = SpecialDoFn() pardo = beam.ParDo(dofn) dd = DisplayData.create_from(pardo) dofn_nspace = '{}.{}'.format(dofn.__module__, dofn.__class__.__name__) pardo_nspace = '{}.{}'.format(pardo.__module__, pardo.__class__.__name__) expected_items = [ DisplayDataItemMatcher('dofn_value', 42, dofn_nspace), DisplayDataItemMatcher('fn', SpecialDoFn, pardo_nspace) ] hc.assert_that(dd.items, hc.has_items(*expected_items)) # TODO: Test __repr__ function # TODO: Test PATH when added by swegner@ if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/dofn_lifecycle_test.py0000644000175100017510000000775315206113033024446 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """UnitTests for DoFn lifecycle and bundle methods""" # pytype: skip-file import unittest import pytest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline class CallSequenceEnforcingDoFn(beam.DoFn): def __init__(self): self._setup_called = False self._start_bundle_calls = 0 self._finish_bundle_calls = 0 self._teardown_called = False def setup(self): assert not self._setup_called, 'setup should not be called twice' assert self._start_bundle_calls == 0, \ 'setup should be called before start_bundle' assert self._finish_bundle_calls == 0, \ 'setup should be called before finish_bundle' assert not self._teardown_called, 'setup should be called before teardown' self._setup_called = True def start_bundle(self): assert self._setup_called, 'setup should have been called' assert self._start_bundle_calls == self._finish_bundle_calls, \ 'there should be as many start_bundle calls as finish_bundle calls' assert not self._teardown_called, 'teardown should not have been called' self._start_bundle_calls += 1 def process(self, element): assert self._setup_called, 'setup should have been called' assert self._start_bundle_calls > 0, 'start_bundle should have been called' assert self._start_bundle_calls == self._finish_bundle_calls + 1, \ 'there should be one start_bundle call with no call to finish_bundle' assert not self._teardown_called, 'teardown should not have been called' return [element * element] def finish_bundle(self): assert self._setup_called, 'setup should have been called' assert self._start_bundle_calls > 0, 'start_bundle should have been called' assert self._start_bundle_calls == self._finish_bundle_calls + 1, \ 'there should be one start_bundle call with no call to finish_bundle' assert not self._teardown_called, 'teardown should not have been called' self._finish_bundle_calls += 1 def teardown(self): assert self._setup_called, 'setup should have been called' assert self._start_bundle_calls == self._finish_bundle_calls, \ 'there should be as many start_bundle calls as finish_bundle calls' assert not self._teardown_called, 'teardown should not be called twice' self._teardown_called = True @pytest.mark.it_validatesrunner class DoFnLifecycleTest(unittest.TestCase): def test_dofn_lifecycle(self): with TestPipeline() as p: _ = ( p | 'Start' >> beam.Create([1, 2, 3]) | 'Do' >> beam.ParDo(CallSequenceEnforcingDoFn())) # Assumes that the worker is run in the same process as the test. class LocalDoFnLifecycleTest(unittest.TestCase): def test_dofn_lifecycle(self): from apache_beam.runners.direct import direct_runner from apache_beam.runners.portability import fn_api_runner runners = [ direct_runner.BundleBasedDirectRunner(), fn_api_runner.FnApiRunner() ] for r in runners: with TestPipeline(runner=r) as p: _ = ( p | 'Start' >> beam.Create([1, 2, 3]) | 'Do' >> beam.ParDo(CallSequenceEnforcingDoFn())) # Assumes that the worker is run in the same process as the test. if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment.py0000644000175100017510000001754215206113033022573 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging from collections.abc import Callable from datetime import timedelta from typing import Any from typing import Optional from typing import TypeVar from typing import Union import apache_beam as beam from apache_beam.coders import coders from apache_beam.io.requestresponse import DEFAULT_CACHE_ENTRY_TTL_SEC from apache_beam.io.requestresponse import DEFAULT_TIMEOUT_SECS from apache_beam.io.requestresponse import Caller from apache_beam.io.requestresponse import DefaultThrottler from apache_beam.io.requestresponse import ExponentialBackOffRepeater from apache_beam.io.requestresponse import PreCallThrottler from apache_beam.io.requestresponse import RedisCache from apache_beam.io.requestresponse import Repeater from apache_beam.io.requestresponse import RequestResponseIO __all__ = [ "EnrichmentSourceHandler", "Enrichment", "cross_join", ] InputT = TypeVar('InputT') OutputT = TypeVar('OutputT') JoinFn = Callable[[dict[str, Any], dict[str, Any]], beam.Row] _LOGGER = logging.getLogger(__name__) def has_valid_redis_address(host: str, port: int) -> bool: """returns `True` if both host and port are not `None`.""" if host and port: return True return False def cross_join(left: dict[str, Any], right: dict[str, Any]) -> beam.Row: """performs a cross join on two `dict` objects. Joins the columns of the right row onto the left row. Args: left (dict[str, Any]): input request dictionary. right (dict[str, Any]): response dictionary from the API. Returns: `beam.Row` containing the merged columns. """ for k, v in right.items(): if k not in left: # Don't override the values in left. left[k] = v elif left[k] != v: _LOGGER.warning( '%s exists in the input row as well the row fetched ' 'from API but have different values - %s and %s. Using the input ' 'value (%s) for the enriched row. You can override this behavior by ' 'passing a custom `join_fn` to Enrichment transform.' % (k, left[k], v, left[k])) return beam.Row(**left) class EnrichmentSourceHandler(Caller[InputT, OutputT]): """Wrapper class for `apache_beam.io.requestresponse.Caller`. Ensure that the implementation of ``__call__`` method returns a tuple of `beam.Row` objects. """ def get_cache_key(self, request: InputT) -> str: """Returns the request to be cached. This is how the response will be looked up in the cache as well. Implement this method to provide the key for the cache. By default, the entire request is stored as the cache key. For example, in `BigTableEnrichmentHandler`, the row key for the element is returned here. """ return "request: %s" % request class Enrichment(beam.PTransform[beam.PCollection[InputT], beam.PCollection[OutputT]]): """A :class:`apache_beam.transforms.enrichment.Enrichment` transform to enrich elements in a PCollection. Uses the :class:`apache_beam.transforms.enrichment.EnrichmentSourceHandler` to enrich elements by joining the metadata from external source. Processes an input :class:`~apache_beam.pvalue.PCollection` of `beam.Row` by applying a :class:`apache_beam.transforms.enrichment.EnrichmentSourceHandler` to each element and returning the enriched :class:`~apache_beam.pvalue.PCollection`. Args: source_handler: Handles source lookup and metadata retrieval. Implements the :class:`apache_beam.transforms.enrichment.EnrichmentSourceHandler` join_fn: A lambda function to join original element with lookup metadata. Defaults to `CROSS_JOIN`. timeout: (Optional) timeout for source requests. Defaults to 30 seconds. repeater: provides method to repeat failed requests to API due to service errors. Defaults to :class:`apache_beam.io.requestresponse.ExponentialBackOffRepeater` to repeat requests with exponential backoff. throttler: provides methods to pre-throttle a request. Defaults to :class:`apache_beam.io.requestresponse.DefaultThrottler` for client-side adaptive throttling using :class:`apache_beam.io.components.adaptive_throttler.AdaptiveThrottler`. """ def __init__( self, source_handler: EnrichmentSourceHandler, join_fn: Optional[JoinFn] = None, timeout: Optional[float] = DEFAULT_TIMEOUT_SECS, repeater: Repeater = ExponentialBackOffRepeater(), throttler: PreCallThrottler = DefaultThrottler(), use_custom_types: bool = False): self._cache = None self._source_handler = source_handler self._join_fn = ( join_fn if join_fn else source_handler.join_fn if hasattr( source_handler, 'join_fn') else cross_join) self._timeout = timeout self._repeater = repeater self._throttler = throttler self._use_custom_types = ( source_handler.use_custom_types if hasattr( source_handler, 'use_custom_types') else use_custom_types) def expand(self, input_row: beam.PCollection[InputT]) -> beam.PCollection[OutputT]: # For caching with enrichment transform, enrichment handlers provide a # get_cache_key() method that returns a unique string formatted # request for that row. request_coder = coders.StrUtf8Coder() if self._cache: self._cache.request_coder = request_coder fetched_data = ( input_row | "Enrichment-RRIO" >> RequestResponseIO( caller=self._source_handler, timeout=self._timeout, repeater=self._repeater, cache=self._cache, throttler=self._throttler)) # EnrichmentSourceHandler returns a tuple of (request,response). return ( fetched_data | "enrichment_join" >> beam.Map( lambda x: self._join_fn(x[0]._asdict(), x[1]._asdict()) if not self._use_custom_types else self._join_fn(x[0], x[1]))) def with_redis_cache( self, host: str, port: int, time_to_live: Union[int, timedelta] = DEFAULT_CACHE_ENTRY_TTL_SEC, *, request_coder: Optional[coders.Coder] = None, response_coder: Optional[coders.Coder] = None, **kwargs, ): """Configure the Redis cache to use with enrichment transform. Args: host (str): The hostname or IP address of the Redis server. port (int): The port number of the Redis server. time_to_live: `(Union[int, timedelta])` The time-to-live (TTL) for records stored in Redis. Provide an integer (in seconds) or a `datetime.timedelta` object. request_coder: (Optional[`coders.Coder`]) coder for requests stored in Redis. response_coder: (Optional[`coders.Coder`]) coder for decoding responses received from Redis. kwargs: Optional additional keyword arguments that are required to connect to your redis server. Same as `redis.Redis()`. """ if has_valid_redis_address(host, port): self._cache = RedisCache( # type: ignore[assignment] host=host, port=port, time_to_live=time_to_live, request_coder=request_coder, response_coder=response_coder, **kwargs) return self ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.760448 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/0000755000175100017510000000000015206113160024071 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/__init__.py0000644000175100017510000000142015206113033026176 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/bigquery.py0000644000175100017510000003032115206113033026270 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging from collections import defaultdict from collections.abc import Callable from collections.abc import Mapping from typing import Any from typing import Optional from typing import Union from google.api_core.exceptions import BadRequest from google.cloud import bigquery import apache_beam as beam from apache_beam.pvalue import Row from apache_beam.transforms.enrichment import EnrichmentSourceHandler QueryFn = Callable[[beam.Row], str] ConditionValueFn = Callable[[beam.Row], list[Any]] _LOGGER = logging.getLogger(__name__) def _validate_bigquery_metadata( table_name, row_restriction_template, fields, condition_value_fn, query_fn): if query_fn: if bool(table_name or row_restriction_template or fields or condition_value_fn): raise ValueError( "Please provide either `query_fn` or the parameters `table_name`, " "`row_restriction_template`, and `fields/condition_value_fn` " "together.") else: if not (table_name and row_restriction_template): raise ValueError( "Please provide either `query_fn` or the parameters " "`table_name`, `row_restriction_template` together.") if ((fields and condition_value_fn) or (not fields and not condition_value_fn)): raise ValueError( "Please provide exactly one of `fields` or " "`condition_value_fn`") class BigQueryEnrichmentHandler(EnrichmentSourceHandler[Union[Row, list[Row]], Union[Row, list[Row]]]): """Enrichment handler for Google Cloud BigQuery. Use this handler with :class:`apache_beam.transforms.enrichment.Enrichment` transform. To use this handler you need either of the following combinations: * `table_name`, `row_restriction_template`, `fields` * `table_name`, `row_restriction_template`, `condition_value_fn` * `query_fn` By default, the handler pulls all columns from the BigQuery table. To override this, use the `column_name` parameter to specify a list of column names to fetch. This handler pulls data from BigQuery per element by default. To change this behavior, set the `min_batch_size`, `max_batch_size`, and `max_batch_duration_secs` parameters. These batching values are sent to the :class:`apache_beam.transforms.utils.BatchElements` transform. NOTE: Elements cannot be batched when using the `query_fn` parameter. """ def __init__( self, project: str, *, table_name: str = "", row_restriction_template: str = "", fields: Optional[list[str]] = None, column_names: Optional[list[str]] = None, condition_value_fn: Optional[ConditionValueFn] = None, query_fn: Optional[QueryFn] = None, min_batch_size: int = 1, max_batch_size: int = 10000, max_batch_duration_secs: Optional[float] = None, throw_exception_on_empty_results: bool = True, **kwargs, ): """ Example Usage: handler = BigQueryEnrichmentHandler(project=project_name, row_restriction="id='{}'", table_name='project.dataset.table', fields=fields, min_batch_size=2, max_batch_size=100) Args: project: Google Cloud project ID for the BigQuery table. table_name (str): Fully qualified BigQuery table name in the format `project.dataset.table`. row_restriction_template (str): A template string for the `WHERE` clause in the BigQuery query with placeholders (`{}`) to dynamically filter rows based on input data. fields: (Optional[list[str]]) List of field names present in the input `beam.Row`. These are used to construct the WHERE clause (if `condition_value_fn` is not provided). column_names: (Optional[list[str]]) Names of columns to select from the BigQuery table. If not provided, all columns (`*`) are selected. condition_value_fn: (Optional[Callable[[beam.Row], Any]]) A function that takes a `beam.Row` and returns a list of value to populate in the placeholder `{}` of `WHERE` clause in the query. query_fn: (Optional[Callable[[beam.Row], str]]) A function that takes a `beam.Row` and returns a complete BigQuery SQL query string. min_batch_size (int): Minimum number of rows to batch together when querying BigQuery. Defaults to 1 if `query_fn` is not specified. max_batch_size (int): Maximum number of rows to batch together. Defaults to 10,000 if `query_fn` is not specified. max_batch_duration_secs (float): Maximum amount of time in seconds to buffer a batch before emitting it. If not provided, batching duration is determined by `BatchElements` defaults. **kwargs: Additional keyword arguments to pass to `bigquery.Client`. Note: * `min_batch_size`, `max_batch_size`, and `max_batch_duration_secs` are not used if `query_fn` is provided. * Either `fields` or `condition_value_fn` must be provided for query construction if `query_fn` is not provided. * Ensure appropriate permissions are granted for BigQuery access. """ _validate_bigquery_metadata( table_name, row_restriction_template, fields, condition_value_fn, query_fn) self.project = project self.column_names = column_names self.select_fields = ",".join(column_names) if column_names else '*' self.row_restriction_template = row_restriction_template self.table_name = table_name self.fields = fields if fields else [] self.condition_value_fn = condition_value_fn self.query_fn = query_fn self.query_template = ( "SELECT %s FROM %s WHERE %s" % (self.select_fields, self.table_name, self.row_restriction_template)) self.throw_exception_on_empty_results = throw_exception_on_empty_results self.kwargs = kwargs self._batching_kwargs = {} if not query_fn: self._batching_kwargs['min_batch_size'] = min_batch_size self._batching_kwargs['max_batch_size'] = max_batch_size if max_batch_duration_secs is not None: self._batching_kwargs[ 'max_batch_duration_secs'] = max_batch_duration_secs def __enter__(self): self.client = bigquery.Client(project=self.project, **self.kwargs) def _execute_query(self, query: str): try: results = self.client.query(query=query).result() row_list = [dict(row.items()) for row in results] if not row_list: return None if self._batching_kwargs: return row_list else: return row_list[0] except BadRequest as e: raise BadRequest( f'Could not execute the query: {query}. Please check if ' f'the query is properly formatted and the BigQuery ' f'table exists. {e}') except RuntimeError as e: raise RuntimeError(f"Could not complete the query request: {query}. {e}") def create_row_key(self, row: beam.Row): if self.condition_value_fn: return tuple(self.condition_value_fn(row)) if self.fields: row_dict = row._asdict() return (tuple(row_dict[field] for field in self.fields)) raise ValueError("Either fields or condition_value_fn must be specified") def __call__(self, request: Union[beam.Row, list[beam.Row]], *args, **kwargs): if isinstance(request, list): values = [] responses = [] requests_map: dict[Any, list[beam.Row]] = defaultdict(list) batch_size = len(request) raw_query = self.query_template if batch_size > 1: batched_condition_template = ' or '.join( [fr'({self.row_restriction_template})'] * batch_size) raw_query = self.query_template.replace( self.row_restriction_template, batched_condition_template) for req in request: request_dict = req._asdict() try: current_values = ( self.condition_value_fn(req) if self.condition_value_fn else [request_dict[field] for field in self.fields]) except KeyError as e: raise KeyError( "Make sure the values passed in `fields` are the " "keys in the input `beam.Row`." + str(e)) values.extend(current_values) requests_map[self.create_row_key(req)].append(req) query = raw_query.format(*values) responses_dict = self._execute_query(query) unmatched_requests = { key: list(reqs) for key, reqs in requests_map.items() } if responses_dict: for response in responses_dict: response_row = beam.Row(**response) response_key = self.create_row_key(response_row) if response_key in unmatched_requests: for req in unmatched_requests.pop(response_key): responses.append((req, response_row)) if unmatched_requests: if self.throw_exception_on_empty_results: raise ValueError(f"no matching row found for query: {query}") else: _LOGGER.warning('no matching row found for query: %s', query) for reqs in unmatched_requests.values(): for req in reqs: responses.append((req, beam.Row())) return responses else: request_dict = request._asdict() if self.query_fn: # if a query_fn is provided then it return a list of values # that should be populated into the query template string. query = self.query_fn(request) else: values = ( self.condition_value_fn(request) if self.condition_value_fn else list(map(request_dict.get, self.fields))) # construct the query. query = self.query_template.format(*values) response_dict = self._execute_query(query) if response_dict is None: if self.throw_exception_on_empty_results: raise ValueError(f"no matching row found for query: {query}") else: _LOGGER.warning('no matching row found for query: %s', query) return request, beam.Row() return request, beam.Row(**response_dict) def __exit__(self, exc_type, exc_val, exc_tb): self.client.close() def get_cache_key(self, request: Union[beam.Row, list[beam.Row]]): if isinstance(request, list): cache_keys = [] for req in request: req_dict = req._asdict() try: current_values = ( self.condition_value_fn(req) if self.condition_value_fn else [req_dict[field] for field in self.fields]) key = ";".join(["%s"] * len(current_values)) cache_keys.extend([key % tuple(current_values)]) except KeyError as e: raise KeyError( "Make sure the values passed in `fields` are the " "keys in the input `beam.Row`." + str(e)) return cache_keys else: req_dict = request._asdict() try: current_values = ( self.condition_value_fn(request) if self.condition_value_fn else [req_dict[field] for field in self.fields]) key = ";".join(["%s"] * len(current_values)) cache_key = key % tuple(current_values) except KeyError as e: raise KeyError( "Make sure the values passed in `fields` are the " "keys in the input `beam.Row`." + str(e)) return cache_key def batch_elements_kwargs(self) -> Mapping[str, Any]: """Returns a kwargs suitable for `beam.BatchElements`.""" return self._batching_kwargs ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/bigquery_it_test.py0000644000175100017510000004255015206113033030032 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import functools import logging import secrets import time import unittest from unittest.mock import MagicMock import pytest import apache_beam as beam from apache_beam.coders import coders from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # pylint: disable=ungrouped-imports try: from apitools.base.py.exceptions import HttpError from testcontainers.redis import RedisContainer from apache_beam.transforms.enrichment import Enrichment from apache_beam.transforms.enrichment_handlers.bigquery import BigQueryEnrichmentHandler except ImportError: raise unittest.SkipTest( 'Google Cloud BigQuery dependencies are not installed.') _LOGGER = logging.getLogger(__name__) def condition_value_fn(row: beam.Row): return [row.id] # type: ignore[attr-defined] def query_fn(table, row: beam.Row): return f"SELECT * FROM `{table}` WHERE id = {row.id}" # type: ignore[attr-defined] @pytest.mark.uses_testcontainer class BigQueryEnrichmentIT(unittest.TestCase): bigquery_dataset_id = 'python_enrichment_transform_read_table_' project = "apache-beam-testing" @classmethod def setUpClass(cls): cls.bigquery_client = BigQueryWrapper() cls.dataset_id = '%s%d%s' % ( cls.bigquery_dataset_id, int(time.time()), secrets.token_hex(3)) cls.bigquery_client.get_or_create_dataset(cls.project, cls.dataset_id) _LOGGER.info( "Created dataset %s in project %s", cls.dataset_id, cls.project) @classmethod def tearDownClass(cls): request = bigquery.BigqueryDatasetsDeleteRequest( projectId=cls.project, datasetId=cls.dataset_id, deleteContents=True) try: _LOGGER.debug( "Deleting dataset %s in project %s", cls.dataset_id, cls.project) cls.bigquery_client.client.datasets.Delete(request) except HttpError: _LOGGER.warning( 'Failed to clean up dataset %s in project %s', cls.dataset_id, cls.project) @pytest.mark.uses_testcontainer class TestBigQueryEnrichmentIT(BigQueryEnrichmentIT): table_data = [ { "id": 1, "name": "A", 'quantity': 2, 'distribution_center_id': 3 }, { "id": 2, "name": "B", 'quantity': 3, 'distribution_center_id': 1 }, { "id": 3, "name": "C", 'quantity': 10, 'distribution_center_id': 4 }, { "id": 4, "name": "D", 'quantity': 1, 'distribution_center_id': 3 }, { "id": 5, "name": "C", 'quantity': 100, 'distribution_center_id': 4 }, { "id": 6, "name": "D", 'quantity': 11, 'distribution_center_id': 3 }, { "id": 7, "name": "C", 'quantity': 7, 'distribution_center_id': 1 }, { "id": 8, "name": "D", 'quantity': 4, 'distribution_center_id': 1 }, ] @classmethod def create_table(cls, table_name): fields = [('id', 'INTEGER'), ('name', 'STRING'), ('quantity', 'INTEGER'), ('distribution_center_id', 'INTEGER')] table_schema = bigquery.TableSchema() for name, field_type in fields: table_field = bigquery.TableFieldSchema() table_field.name = name table_field.type = field_type table_schema.fields.append(table_field) table = bigquery.Table( tableReference=bigquery.TableReference( projectId=cls.project, datasetId=cls.dataset_id, tableId=table_name), schema=table_schema) request = bigquery.BigqueryTablesInsertRequest( projectId=cls.project, datasetId=cls.dataset_id, table=table) cls.bigquery_client.client.tables.Insert(request) cls.bigquery_client.insert_rows( cls.project, cls.dataset_id, table_name, cls.table_data) cls.table_name = f"{cls.project}.{cls.dataset_id}.{table_name}" @classmethod def setUpClass(cls): super(TestBigQueryEnrichmentIT, cls).setUpClass() cls.create_table('product_details') def setUp(self) -> None: self.condition_template = "id = {}" self.retries = 5 self._start_container() def _start_container(self): for i in range(self.retries): try: self.container = RedisContainer(image='redis:7.2.4') self.container.start() self.host = self.container.get_container_host_ip() self.port = self.container.get_exposed_port(6379) self.client = self.container.get_client() break except Exception as e: if i == self.retries - 1: _LOGGER.error( 'Unable to start redis container for BigQuery ' ' enrichment tests.') raise e # Add a small delay between retries to avoid rapid successive failures time.sleep(2) def tearDown(self) -> None: self.container.stop() self.client = None def test_bigquery_enrichment(self): expected_rows = [ beam.Row(id=1, name="A", quantity=2, distribution_center_id=3), beam.Row(id=2, name="B", quantity=3, distribution_center_id=1) ] fields = ['id'] requests = [ beam.Row(id=1, name='A'), beam.Row(id=2, name='B'), ] handler = BigQueryEnrichmentHandler( project=self.project, row_restriction_template="id = {}", table_name=self.table_name, fields=fields, min_batch_size=1, max_batch_size=100, ) with TestPipeline(is_integration_test=True) as test_pipeline: pcoll = (test_pipeline | beam.Create(requests) | Enrichment(handler)) assert_that(pcoll, equal_to(expected_rows)) def test_bigquery_enrichment_batched(self): expected_rows = [ beam.Row(id=1, name="A", quantity=2, distribution_center_id=3), beam.Row(id=2, name="B", quantity=3, distribution_center_id=1) ] fields = ['id'] requests = [ beam.Row(id=1, name='A'), beam.Row(id=2, name='B'), ] handler = BigQueryEnrichmentHandler( project=self.project, row_restriction_template="id = {}", table_name=self.table_name, fields=fields, min_batch_size=2, max_batch_size=100, ) with TestPipeline(is_integration_test=True) as test_pipeline: pcoll = (test_pipeline | beam.Create(requests) | Enrichment(handler)) assert_that(pcoll, equal_to(expected_rows)) def test_bigquery_enrichment_batched_multiple_fields(self): expected_rows = [ beam.Row(id=1, distribution_center_id=3, name="A", quantity=2), beam.Row(id=2, distribution_center_id=1, name="B", quantity=3) ] fields = ['id', 'distribution_center_id'] requests = [ beam.Row(id=1, distribution_center_id=3), beam.Row(id=2, distribution_center_id=1), ] handler = BigQueryEnrichmentHandler( project=self.project, row_restriction_template="id = {} AND distribution_center_id = {}", table_name=self.table_name, fields=fields, min_batch_size=8, max_batch_size=100, ) with TestPipeline(is_integration_test=True) as test_pipeline: pcoll = (test_pipeline | beam.Create(requests) | Enrichment(handler)) assert_that(pcoll, equal_to(expected_rows)) def test_bigquery_enrichment_with_query_fn(self): expected_rows = [ beam.Row(id=1, name="A", quantity=2, distribution_center_id=3), beam.Row(id=2, name="B", quantity=3, distribution_center_id=1) ] requests = [ beam.Row(id=1, name='A'), beam.Row(id=2, name='B'), ] fn = functools.partial(query_fn, self.table_name) handler = BigQueryEnrichmentHandler(project=self.project, query_fn=fn) with TestPipeline(is_integration_test=True) as test_pipeline: pcoll = (test_pipeline | beam.Create(requests) | Enrichment(handler)) assert_that(pcoll, equal_to(expected_rows)) def test_bigquery_enrichment_with_condition_value_fn(self): expected_rows = [ beam.Row(id=1, name="A", quantity=2, distribution_center_id=3), beam.Row(id=2, name="B", quantity=3, distribution_center_id=1) ] requests = [ beam.Row(id=1, name='A'), beam.Row(id=2, name='B'), ] handler = BigQueryEnrichmentHandler( project=self.project, row_restriction_template=self.condition_template, table_name=self.table_name, condition_value_fn=condition_value_fn, min_batch_size=2, max_batch_size=100, ) with TestPipeline(is_integration_test=True) as test_pipeline: pcoll = (test_pipeline | beam.Create(requests) | Enrichment(handler)) assert_that(pcoll, equal_to(expected_rows)) def test_bigquery_enrichment_bad_request(self): requests = [ beam.Row(id=1, name='A'), beam.Row(id=2, name='B'), ] handler = BigQueryEnrichmentHandler( project=self.project, row_restriction_template=self.condition_template, table_name=self.table_name, column_names=['wrong_column'], condition_value_fn=condition_value_fn, ) with self.assertRaises(Exception): test_pipeline = beam.Pipeline() _ = ( test_pipeline | "Create" >> beam.Create(requests) | "Enrichment" >> Enrichment(handler)) res = test_pipeline.run() res.wait_until_finish() def test_bigquery_enrichment_with_redis(self): """ In this test, we run two pipelines back to back. In the first pipeline, we run a simple BigQuery enrichment pipeline with zero cache records. Therefore, it makes call to the source and ultimately writes to the cache with a TTL of 300 seconds. For the second pipeline, we mock the `BigQueryEnrichmentHandler`'s `__call__` method to always return a `None` response. However, this change won't impact the second pipeline because the Enrichment transform first checks the cache to fulfill requests. Since all requests are cached, it will return from there without making calls to the BigQuery service. """ requests = [ beam.Row(id=1, name='A'), beam.Row(id=2, name='B'), ] expected_rows = [ beam.Row(id=1, name="A", quantity=2, distribution_center_id=3), beam.Row(id=2, name="B", quantity=3, distribution_center_id=1) ] handler = BigQueryEnrichmentHandler( project=self.project, row_restriction_template=self.condition_template, table_name=self.table_name, condition_value_fn=condition_value_fn, min_batch_size=2, max_batch_size=100, ) with TestPipeline(is_integration_test=True) as test_pipeline: pcoll_populate_cache = ( test_pipeline | beam.Create(requests) | Enrichment(handler).with_redis_cache(self.host, self.port)) assert_that(pcoll_populate_cache, equal_to(expected_rows)) # manually check cache entry c = coders.StrUtf8Coder() for req in requests: key = handler.get_cache_key(req) response = self.client.get(c.encode(key)) if not response: raise ValueError("No cache entry found for %s" % key) actual = BigQueryEnrichmentHandler.__call__ BigQueryEnrichmentHandler.__call__ = MagicMock(return_value=(beam.Row())) with TestPipeline(is_integration_test=True) as test_pipeline: pcoll_cached = ( test_pipeline | beam.Create(requests) | Enrichment(handler).with_redis_cache(self.host, self.port)) assert_that(pcoll_cached, equal_to(expected_rows)) BigQueryEnrichmentHandler.__call__ = actual def test_bigquery_enrichment_no_results_throws_exception(self): requests = [ beam.Row(id=999, name='X'), # This ID does not exist ] handler = BigQueryEnrichmentHandler( project=self.project, row_restriction_template="id = {}", table_name=self.table_name, fields=['id'], throw_exception_on_empty_results=True, ) with self.assertRaisesRegex(ValueError, "no matching row found for query"): with TestPipeline(is_integration_test=True) as test_pipeline: _ = (test_pipeline | beam.Create(requests) | Enrichment(handler)) def test_bigquery_enrichment_no_results_graceful(self): requests = [ beam.Row(id=999, name='X'), # This ID does not exist beam.Row(id=1000, name='Y'), # This ID does not exist ] # When no results are found and not throwing, Enrichment yields original. expected_rows = requests handler = BigQueryEnrichmentHandler( project=self.project, row_restriction_template="id = {}", table_name=self.table_name, fields=['id'], min_batch_size=1, max_batch_size=100, throw_exception_on_empty_results=False, ) with TestPipeline(is_integration_test=True) as test_pipeline: pcoll = (test_pipeline | beam.Create(requests) | Enrichment(handler)) assert_that(pcoll, equal_to(expected_rows)) def test_bigquery_enrichment_no_results_partial_graceful_batched(self): requests = [ beam.Row(id=1, name='A'), # This ID exists beam.Row(id=1000, name='Y'), # This ID does not exist ] # When no results are found and not throwing, Enrichment yields original. expected_rows = [ beam.Row(id=1, name='A', quantity=2, distribution_center_id=3), beam.Row(id=1000, name='Y'), # This ID does not exist so remains unchanged ] handler = BigQueryEnrichmentHandler( project=self.project, row_restriction_template="id = {}", table_name=self.table_name, fields=['id'], min_batch_size=2, max_batch_size=100, throw_exception_on_empty_results=False, ) with TestPipeline(is_integration_test=True) as test_pipeline: pcoll = (test_pipeline | beam.Create(requests) | Enrichment(handler)) assert_that(pcoll, equal_to(expected_rows)) def test_bigquery_enrichment_no_results_graceful_batched(self): requests = [ beam.Row(id=999, name='X'), # This ID does not exist beam.Row(id=1000, name='Y'), # This ID does not exist ] # When no results are found and not throwing, Enrichment yields original. expected_rows = requests handler = BigQueryEnrichmentHandler( project=self.project, row_restriction_template="id = {}", table_name=self.table_name, fields=['id'], min_batch_size=2, max_batch_size=100, throw_exception_on_empty_results=False, ) with TestPipeline(is_integration_test=True) as test_pipeline: pcoll = (test_pipeline | beam.Create(requests) | Enrichment(handler)) assert_that(pcoll, equal_to(expected_rows)) def test_bigquery_enrichment_no_results_with_query_fn_throws_exception(self): requests = [ beam.Row(id=999, name='X'), # This ID does not exist ] # This query_fn will return no results fn = functools.partial(query_fn, self.table_name) handler = BigQueryEnrichmentHandler( project=self.project, query_fn=fn, throw_exception_on_empty_results=True, ) with self.assertRaisesRegex(ValueError, "no matching row found for query"): with TestPipeline(is_integration_test=True) as test_pipeline: _ = (test_pipeline | beam.Create(requests) | Enrichment(handler)) def test_bigquery_enrichment_no_results_with_query_fn_graceful(self): requests = [ beam.Row(id=999, name='X'), # This ID does not exist beam.Row(id=1000, name='Y'), # This ID does not exist ] # When no results are found and not throwing, Enrichment yields original. expected_rows = requests # This query_fn will return no results fn = functools.partial(query_fn, self.table_name) handler = BigQueryEnrichmentHandler( project=self.project, query_fn=fn, throw_exception_on_empty_results=False, ) with TestPipeline(is_integration_test=True) as test_pipeline: pcoll = (test_pipeline | beam.Create(requests) | Enrichment(handler)) assert_that(pcoll, equal_to(expected_rows)) def test_bigquery_enrichment_partial_results_throws_exception_batched(self): requests = [ beam.Row(id=1, name='A'), # This ID exists beam.Row(id=1000, name='Y'), # This ID does not exist ] handler = BigQueryEnrichmentHandler( project=self.project, row_restriction_template="id = {}", table_name=self.table_name, fields=['id'], min_batch_size=2, max_batch_size=100, throw_exception_on_empty_results=True, ) with self.assertRaisesRegex(ValueError, "no matching row found for query"): with TestPipeline(is_integration_test=True) as test_pipeline: _ = (test_pipeline | beam.Create(requests) | Enrichment(handler)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/bigquery_test.py0000644000175100017510000001173615206113033027340 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest from unittest import mock from parameterized import parameterized import apache_beam as beam # pylint: disable=ungrouped-imports try: from apache_beam.transforms.enrichment_handlers.bigquery import BigQueryEnrichmentHandler from apache_beam.transforms.enrichment_handlers.bigquery_it_test import condition_value_fn from apache_beam.transforms.enrichment_handlers.bigquery_it_test import query_fn except ImportError: raise unittest.SkipTest( 'Google Cloud BigQuery dependencies are not installed.') class TestBigQueryEnrichment(unittest.TestCase): def setUp(self) -> None: self.project = 'apache-beam-testing' @parameterized.expand([ ("", "", [], None, None, 1, 2), ("table", "", ["id"], condition_value_fn, None, 2, 10), ("table", "id='{}'", ["id"], condition_value_fn, None, 2, 10), ("table", "id='{}'", ["id"], None, query_fn, 2, 10), ]) def test_valid_params( self, table_name, row_restriction_template, fields, condition_value_fn, query_fn, min_batch_size, max_batch_size): """ TC 1: Only batch size are provided. It should raise an error. TC 2: Either of `row_restriction template` or `query_fn` is not provided. TC 3: Both `fields` and `condition_value_fn` are provided. TC 4: Query construction details are provided along with `query_fn`. """ with self.assertRaises(ValueError): _ = BigQueryEnrichmentHandler( project=self.project, table_name=table_name, row_restriction_template=row_restriction_template, fields=fields, condition_value_fn=condition_value_fn, query_fn=query_fn, min_batch_size=min_batch_size, max_batch_size=max_batch_size, ) def test_batch_mode_fans_out_response_for_duplicate_keys(self): handler = BigQueryEnrichmentHandler( project=self.project, table_name='project.dataset.table', row_restriction_template="id='{}'", fields=['id'], min_batch_size=2, max_batch_size=2, ) requests = [beam.Row(id='1', name='first'), beam.Row(id='1', name='second')] with mock.patch.object(handler, '_execute_query', return_value=[{'id': '1', 'value': 'enriched'}]): responses = handler(requests) self.assertEqual( responses, [ (requests[0], beam.Row(id='1', value='enriched')), (requests[1], beam.Row(id='1', value='enriched')), ], ) def test_batch_mode_emits_empty_rows_for_all_unmatched_duplicate_keys(self): handler = BigQueryEnrichmentHandler( project=self.project, table_name='project.dataset.table', row_restriction_template="id='{}'", fields=['id'], min_batch_size=2, max_batch_size=2, throw_exception_on_empty_results=False, ) requests = [beam.Row(id='1', name='first'), beam.Row(id='1', name='second')] with mock.patch.object(handler, '_execute_query', return_value=None): responses = handler(requests) self.assertEqual( responses, [(requests[0], beam.Row()), (requests[1], beam.Row())], ) def test_batch_elements_kwargs_include_max_batch_duration_secs(self): handler = BigQueryEnrichmentHandler( project=self.project, table_name='project.dataset.table', row_restriction_template="id='{}'", fields=['id'], min_batch_size=2, max_batch_size=10, max_batch_duration_secs=0.75, ) self.assertEqual( handler.batch_elements_kwargs(), { 'min_batch_size': 2, 'max_batch_size': 10, 'max_batch_duration_secs': 0.75, }) def test_batch_elements_kwargs_omit_max_batch_duration_secs_by_default(self): handler = BigQueryEnrichmentHandler( project=self.project, table_name='project.dataset.table', row_restriction_template="id='{}'", fields=['id'], min_batch_size=2, max_batch_size=10, ) self.assertEqual( handler.batch_elements_kwargs(), { 'min_batch_size': 2, 'max_batch_size': 10, }) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/bigtable.py0000644000175100017510000001515515206113033026222 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging from collections.abc import Callable from typing import Any from typing import Optional from google.api_core.exceptions import NotFound from google.cloud import bigtable from google.cloud.bigtable import Client from google.cloud.bigtable.row_filters import CellsColumnLimitFilter from google.cloud.bigtable.row_filters import RowFilter import apache_beam as beam from apache_beam.transforms.enrichment import EnrichmentSourceHandler from apache_beam.transforms.enrichment_handlers.utils import ExceptionLevel __all__ = [ 'BigTableEnrichmentHandler', ] RowKeyFn = Callable[[beam.Row], bytes] _LOGGER = logging.getLogger(__name__) class BigTableEnrichmentHandler(EnrichmentSourceHandler[beam.Row, beam.Row]): """A handler for :class:`apache_beam.transforms.enrichment.Enrichment` transform to interact with GCP BigTable. Args: project_id (str): GCP project-id of the BigTable cluster. instance_id (str): GCP instance-id of the BigTable cluster. table_id (str): GCP table-id of the BigTable. row_key (str): unique row-key field name from the input `beam.Row` object to use as `row_key` for BigTable querying. row_filter: a ``:class:`google.cloud.bigtable.row_filters.RowFilter``` to filter data read with ``read_row()``. Defaults to `CellsColumnLimitFilter(1)`. app_profile_id (str): App profile ID to use for BigTable. See https://cloud.google.com/bigtable/docs/app-profiles for more details. encoding (str): encoding type to convert the string to bytes and vice-versa from BigTable. Default is `utf-8`. row_key_fn: a lambda function that returns a string row key from the input row. It is used to build/extract the row key for Bigtable. exception_level: a `enum.Enum` value from ``apache_beam.transforms.enrichment_handlers.utils.ExceptionLevel`` to set the level when an empty row is returned from the BigTable query. Defaults to ``ExceptionLevel.WARN``. include_timestamp (bool): If enabled, the timestamp associated with the value is returned as `(value, timestamp)` for each `row_key`. Defaults to `False` - only the latest value without the timestamp is returned. """ def __init__( self, project_id: str, instance_id: str, table_id: str, row_key: str = "", row_filter: Optional[RowFilter] = CellsColumnLimitFilter(1), *, app_profile_id: str = None, # type: ignore[assignment] encoding: str = 'utf-8', row_key_fn: Optional[RowKeyFn] = None, exception_level: ExceptionLevel = ExceptionLevel.WARN, include_timestamp: bool = False, ): self._project_id = project_id self._instance_id = instance_id self._table_id = table_id self._row_key = row_key self._row_filter = row_filter self._app_profile_id = app_profile_id self._encoding = encoding self._row_key_fn = row_key_fn self._exception_level = exception_level self._include_timestamp = include_timestamp if ((not self._row_key_fn and not self._row_key) or bool(self._row_key_fn and self._row_key)): raise ValueError( "Please specify exactly one of `row_key` or a lambda " "function with `row_key_fn` to extract the row key " "from the input row.") def __enter__(self): """connect to the Google BigTable cluster.""" self.client = Client(project=self._project_id) self.instance = self.client.instance(self._instance_id) self._table = bigtable.table.Table( table_id=self._table_id, instance=self.instance, app_profile_id=self._app_profile_id) def __call__(self, request: beam.Row, *args, **kwargs): """ Reads a row from the GCP BigTable and returns a `Tuple` of request and response. Args: request: the input `beam.Row` to enrich. """ response_dict: dict[str, Any] = {} row_key_str: str = "" try: if self._row_key_fn: row_key = self._row_key_fn(request) else: request_dict = request._asdict() row_key_str = str(request_dict[self._row_key]) row_key = row_key_str.encode(self._encoding) row = self._table.read_row(row_key, filter_=self._row_filter) if row: for cf_id, cf_v in row.cells.items(): response_dict[cf_id] = {} for col_id, col_v in cf_v.items(): if self._include_timestamp: response_dict[cf_id][col_id.decode(self._encoding)] = [ (v.value.decode(self._encoding), v.timestamp) for v in col_v ] else: response_dict[cf_id][col_id.decode( self._encoding)] = col_v[0].value.decode(self._encoding) elif self._exception_level == ExceptionLevel.WARN: _LOGGER.warning( 'no matching row found for row_key: %s ' 'with row_filter: %s' % (row_key_str, self._row_filter)) elif self._exception_level == ExceptionLevel.RAISE: raise ValueError( 'no matching row found for row_key: %s ' 'with row_filter=%s' % (row_key_str, self._row_filter)) except KeyError: raise KeyError('row_key %s not found in input PCollection.' % row_key_str) except NotFound: raise NotFound( 'GCP BigTable cluster `%s:%s:%s` not found.' % (self._project_id, self._instance_id, self._table_id)) except Exception as e: raise e return request, beam.Row(**response_dict) def __exit__(self, exc_type, exc_val, exc_tb): """Clean the instantiated BigTable client.""" self.client = None self.instance = None self._table = None def get_cache_key(self, request: beam.Row) -> str: """Returns a string formatted with row key since it is unique to a request made to `Bigtable`.""" if self._row_key_fn: return "row_key: %s" % str(self._row_key_fn(request)) return "%s: %s" % (self._row_key, request._asdict()[self._row_key]) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/bigtable_it_test.py0000644000175100017510000004071315206113033027753 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import datetime import logging import time import unittest from typing import NamedTuple from unittest.mock import MagicMock import pytest import apache_beam as beam from apache_beam.coders import coders from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import BeamAssertException # pylint: disable=ungrouped-imports try: from google.cloud.bigtable import Client from google.cloud.bigtable.row_filters import ColumnRangeFilter from testcontainers.redis import RedisContainer from apache_beam.transforms.enrichment import Enrichment from apache_beam.transforms.enrichment_handlers.bigtable import BigTableEnrichmentHandler from apache_beam.transforms.enrichment_handlers.bigtable import ExceptionLevel except ImportError: raise unittest.SkipTest('Bigtable test dependencies are not installed.') _LOGGER = logging.getLogger(__name__) def _row_key_fn(request: beam.Row) -> bytes: row_key = str(request.product_id) # type: ignore[attr-defined] return row_key.encode(encoding='utf-8') class ValidateResponse(beam.DoFn): """ValidateResponse validates if a PCollection of `beam.Row` has the required fields.""" def __init__( self, n_fields: int, fields: list[str], enriched_fields: dict[str, list[str]], include_timestamp: bool = False, ): self.n_fields = n_fields self._fields = fields self._enriched_fields = enriched_fields self._include_timestamp = include_timestamp def process(self, element: beam.Row, *args, **kwargs): element_dict = element.as_dict() if len(element_dict.keys()) != self.n_fields: raise BeamAssertException( "Expected %d fields in enriched PCollection:" % self.n_fields) for field in self._fields: if field not in element_dict or element_dict[field] is None: raise BeamAssertException(f"Expected a not None field: {field}") for column_family, columns in self._enriched_fields.items(): if len(element_dict[column_family]) != len(columns): raise BeamAssertException( "Response from bigtable should contain a %s column_family with " "%s keys." % (column_family, columns)) for key in columns: if key not in element_dict[column_family]: raise BeamAssertException( "Response from bigtable should contain a %s column_family with " "%s columns." % (column_family, columns)) if (self._include_timestamp and not isinstance(element_dict[column_family][key][0], tuple)): raise BeamAssertException( "Response from bigtable should contain timestamp associated with " "its value.") class _Currency(NamedTuple): s_id: int id: str def create_rows(table): product_id = 'product_id' product_name = 'product_name' product_stock = 'product_stock' column_family_id = "product" products = [ { 'product_id': 1, 'product_name': 'pixel 5', 'product_stock': 2 }, { 'product_id': 2, 'product_name': 'pixel 6', 'product_stock': 4 }, { 'product_id': 3, 'product_name': 'pixel 7', 'product_stock': 20 }, { 'product_id': 4, 'product_name': 'pixel 8', 'product_stock': 10 }, { 'product_id': 5, 'product_name': 'iphone 11', 'product_stock': 3 }, { 'product_id': 6, 'product_name': 'iphone 12', 'product_stock': 7 }, { 'product_id': 7, 'product_name': 'iphone 13', 'product_stock': 8 }, { 'product_id': 8, 'product_name': 'iphone 14', 'product_stock': 3 }, ] for item in products: row_key = str(item[product_id]).encode() row = table.direct_row(row_key) row.set_cell( column_family_id, product_id.encode(), str(item[product_id]), timestamp=datetime.datetime.now(datetime.timezone.utc)) row.set_cell( column_family_id, product_name.encode(), item[product_name], timestamp=datetime.datetime.now(datetime.timezone.utc)) row.set_cell( column_family_id, product_stock.encode(), str(item[product_stock]), timestamp=datetime.datetime.now(datetime.timezone.utc)) row.commit() @pytest.mark.uses_testcontainer class TestBigTableEnrichment(unittest.TestCase): def setUp(self): self.project_id = 'apache-beam-testing' self.instance_id = 'beam-test' self.table_id = 'bigtable-enrichment-test' self.req = [ beam.Row(sale_id=1, customer_id=1, product_id=1, quantity=1), beam.Row(sale_id=3, customer_id=3, product_id=2, quantity=3), beam.Row(sale_id=5, customer_id=5, product_id=4, quantity=2), beam.Row(sale_id=7, customer_id=7, product_id=1, quantity=1), ] self.row_key = 'product_id' self.column_family_id = 'product' client = Client(project=self.project_id) instance = client.instance(self.instance_id) self.table = instance.table(self.table_id) create_rows(self.table) self.retries = 5 self._start_container() def _start_container(self): for i in range(self.retries): try: self.container = RedisContainer(image='redis:7.2.4') self.container.start() self.host = self.container.get_container_host_ip() self.port = self.container.get_exposed_port(6379) self.client = self.container.get_client() break except Exception as e: if i == self.retries - 1: _LOGGER.error('Unable to start redis container for RRIO tests.') raise e # Add a small delay between retries to avoid rapid successive failures time.sleep(2) def tearDown(self) -> None: self.container.stop() self.table = None def test_enrichment_with_bigtable(self): expected_fields = [ 'sale_id', 'customer_id', 'product_id', 'quantity', 'product' ] expected_enriched_fields = { 'product': ['product_id', 'product_name', 'product_stock'], } bigtable = BigTableEnrichmentHandler( project_id=self.project_id, instance_id=self.instance_id, table_id=self.table_id, row_key=self.row_key) with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | "Create" >> beam.Create(self.req) | "Enrich W/ BigTable" >> Enrichment(bigtable) | "Validate Response" >> beam.ParDo( ValidateResponse( len(expected_fields), expected_fields, expected_enriched_fields))) def test_enrichment_with_bigtable_row_filter(self): expected_fields = [ 'sale_id', 'customer_id', 'product_id', 'quantity', 'product' ] expected_enriched_fields = { 'product': ['product_name', 'product_stock'], } start_column = 'product_name'.encode() column_filter = ColumnRangeFilter(self.column_family_id, start_column) bigtable = BigTableEnrichmentHandler( project_id=self.project_id, instance_id=self.instance_id, table_id=self.table_id, row_key=self.row_key, row_filter=column_filter) with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | "Create" >> beam.Create(self.req) | "Enrich W/ BigTable" >> Enrichment(bigtable) | "Validate Response" >> beam.ParDo( ValidateResponse( len(expected_fields), expected_fields, expected_enriched_fields))) def test_enrichment_with_bigtable_no_enrichment(self): # row_key which is product_id=11 doesn't exist, so the enriched field # won't be added. Hence, the response is same as the request. expected_fields = ['sale_id', 'customer_id', 'product_id', 'quantity'] expected_enriched_fields = {} bigtable = BigTableEnrichmentHandler( project_id=self.project_id, instance_id=self.instance_id, table_id=self.table_id, row_key=self.row_key) req = [beam.Row(sale_id=1, customer_id=1, product_id=11, quantity=1)] with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | "Create" >> beam.Create(req) | "Enrich W/ BigTable" >> Enrichment(bigtable) | "Validate Response" >> beam.ParDo( ValidateResponse( len(expected_fields), expected_fields, expected_enriched_fields))) def test_enrichment_with_bigtable_bad_row_filter(self): # in case of a bad column filter, that is, incorrect column_family_id and # columns, no enrichment is done. If the column_family is correct but not # column names then all columns in that column_family are returned. start_column = 'car_name'.encode() column_filter = ColumnRangeFilter('car_name', start_column) bigtable = BigTableEnrichmentHandler( project_id=self.project_id, instance_id=self.instance_id, table_id=self.table_id, row_key=self.row_key, row_filter=column_filter) with self.assertRaises(Exception): test_pipeline = beam.Pipeline() _ = ( test_pipeline | "Create" >> beam.Create(self.req) | "Enrich W/ BigTable" >> Enrichment(bigtable)) res = test_pipeline.run() res.wait_until_finish() def test_enrichment_with_bigtable_raises_key_error(self): """raises a `KeyError` when the row_key doesn't exist in the input PCollection.""" bigtable = BigTableEnrichmentHandler( project_id=self.project_id, instance_id=self.instance_id, table_id=self.table_id, row_key='car_name') with self.assertRaisesRegex(Exception, "not found in input"): test_pipeline = beam.Pipeline() _ = ( test_pipeline | "Create" >> beam.Create(self.req) | "Enrich W/ BigTable" >> Enrichment(bigtable)) res = test_pipeline.run() res.wait_until_finish() def test_enrichment_with_bigtable_raises_not_found(self): """raises a `NotFound` exception when the GCP BigTable Cluster doesn't exist.""" bigtable = BigTableEnrichmentHandler( project_id=self.project_id, instance_id=self.instance_id, table_id='invalid_table', row_key=self.row_key) with self.assertRaises(Exception): test_pipeline = beam.Pipeline() _ = ( test_pipeline | "Create" >> beam.Create(self.req) | "Enrich W/ BigTable" >> Enrichment(bigtable)) res = test_pipeline.run() res.wait_until_finish() def test_enrichment_with_bigtable_exception_level(self): """raises a `ValueError` exception when the GCP BigTable query returns an empty row.""" bigtable = BigTableEnrichmentHandler( project_id=self.project_id, instance_id=self.instance_id, table_id=self.table_id, row_key=self.row_key, exception_level=ExceptionLevel.RAISE) req = [beam.Row(sale_id=1, customer_id=1, product_id=11, quantity=1)] with self.assertRaisesRegex(Exception, "no matching row"): test_pipeline = beam.Pipeline() _ = ( test_pipeline | "Create" >> beam.Create(req) | "Enrich W/ BigTable" >> Enrichment(bigtable)) res = test_pipeline.run() res.wait_until_finish() def test_enrichment_with_bigtable_with_timestamp(self): """test whether the `(value,timestamp)` is returned when the `include_timestamp` is enabled.""" expected_fields = [ 'sale_id', 'customer_id', 'product_id', 'quantity', 'product' ] expected_enriched_fields = { 'product': ['product_id', 'product_name', 'product_stock'], } bigtable = BigTableEnrichmentHandler( project_id=self.project_id, instance_id=self.instance_id, table_id=self.table_id, row_key=self.row_key, include_timestamp=True) with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | "Create" >> beam.Create(self.req) | "Enrich W/ BigTable" >> Enrichment(bigtable) | "Validate Response" >> beam.ParDo( ValidateResponse( len(expected_fields), expected_fields, expected_enriched_fields, include_timestamp=True))) def test_bigtable_enrichment_with_redis(self): """ In this test, we run two pipelines back to back. In the first pipeline, we run a simple bigtable enrichment pipeline with zero cache records. Therefore, it makes call to the Bigtable source and ultimately writes to the cache with a TTL of 300 seconds. For the second pipeline, we mock the `BigTableEnrichmentHandler`'s `__call__` method to always return a `None` response. However, this change won't impact the second pipeline because the Enrichment transform first checks the cache to fulfill requests. Since all requests are cached, it will return from there without making calls to the Bigtable source. """ expected_fields = [ 'sale_id', 'customer_id', 'product_id', 'quantity', 'product' ] expected_enriched_fields = { 'product': ['product_name', 'product_stock'], } start_column = 'product_name'.encode() column_filter = ColumnRangeFilter(self.column_family_id, start_column) bigtable = BigTableEnrichmentHandler( project_id=self.project_id, instance_id=self.instance_id, table_id=self.table_id, row_key=self.row_key, row_filter=column_filter) with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | "Create1" >> beam.Create(self.req) | "Enrich W/ BigTable1" >> Enrichment(bigtable).with_redis_cache( self.host, self.port, 300) | "Validate Response" >> beam.ParDo( ValidateResponse( len(expected_fields), expected_fields, expected_enriched_fields))) # manually check cache entry c = coders.StrUtf8Coder() for req in self.req: key = bigtable.get_cache_key(req) response = self.client.get(c.encode(key)) if not response: raise ValueError("No cache entry found for %s" % key) actual = BigTableEnrichmentHandler.__call__ BigTableEnrichmentHandler.__call__ = MagicMock( return_value=( beam.Row(sale_id=1, customer_id=1, product_id=1, quantity=1), beam.Row())) with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | "Create2" >> beam.Create(self.req) | "Enrich W/ BigTable2" >> Enrichment(bigtable).with_redis_cache( self.host, self.port) | "Validate Response" >> beam.ParDo( ValidateResponse( len(expected_fields), expected_fields, expected_enriched_fields))) BigTableEnrichmentHandler.__call__ = actual def test_bigtable_enrichment_with_lambda(self): expected_fields = [ 'sale_id', 'customer_id', 'product_id', 'quantity', 'product' ] expected_enriched_fields = { 'product': ['product_id', 'product_name', 'product_stock'], } bigtable = BigTableEnrichmentHandler( project_id=self.project_id, instance_id=self.instance_id, table_id=self.table_id, row_key_fn=_row_key_fn) with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | "Create" >> beam.Create(self.req) | "Enrich W/ BigTable" >> Enrichment(bigtable) | "Validate Response" >> beam.ParDo( ValidateResponse( len(expected_fields), expected_fields, expected_enriched_fields))) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/bigtable_test.py0000644000175100017510000000313315206113033027252 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest from parameterized import parameterized try: from apache_beam.transforms.enrichment_handlers.bigtable import BigTableEnrichmentHandler from apache_beam.transforms.enrichment_handlers.bigtable_it_test import _row_key_fn except ImportError: raise unittest.SkipTest('Bigtable test dependencies are not installed.') class TestBigTableEnrichmentHandler(unittest.TestCase): @parameterized.expand([('product_id', _row_key_fn), ('', None)]) def test_bigtable_enrichment_invalid_args(self, row_key, row_key_fn): with self.assertRaises(ValueError): _ = BigTableEnrichmentHandler( project_id='apache-beam-testing', instance_id='beam-test', table_id='bigtable-enrichment-test', row_key=row_key, row_key_fn=row_key_fn) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/cloudsql.py0000644000175100017510000006225315206113033026300 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import re from abc import ABC from abc import abstractmethod from collections.abc import Callable from collections.abc import Mapping from dataclasses import dataclass from dataclasses import field from enum import Enum from typing import Any from typing import Optional from typing import Union import pg8000 import pymysql import pytds from google.cloud.sql.connector import Connector as CloudSQLConnector from google.cloud.sql.connector.enums import RefreshStrategy from sqlalchemy import create_engine from sqlalchemy import text from sqlalchemy.engine import Connection as DBAPIConnection import apache_beam as beam from apache_beam.transforms.enrichment import EnrichmentSourceHandler QueryFn = Callable[[beam.Row], str] ConditionValueFn = Callable[[beam.Row], list[Any]] @dataclass class CustomQueryConfig: """Configuration for using a custom query function.""" query_fn: QueryFn def __post_init__(self): if not self.query_fn: raise ValueError("CustomQueryConfig must provide a valid query_fn") @dataclass class TableFieldsQueryConfig: """Configuration for using table name, where clause, and field names.""" table_id: str where_clause_template: str where_clause_fields: list[str] def __post_init__(self): if not self.table_id or not self.where_clause_template: raise ValueError( "TableFieldsQueryConfig must provide table_id and " + "where_clause_template") if not self.where_clause_fields: raise ValueError( "TableFieldsQueryConfig must provide non-empty " + "where_clause_fields") @dataclass class TableFunctionQueryConfig: """Configuration for using table name, where clause, and a value function.""" table_id: str where_clause_template: str where_clause_value_fn: ConditionValueFn def __post_init__(self): if not self.table_id or not self.where_clause_template: raise ValueError( "TableFunctionQueryConfig must provide table_id and " + "where_clause_template") if not self.where_clause_value_fn: raise ValueError( "TableFunctionQueryConfig must provide " + "where_clause_value_fn") class DatabaseTypeAdapter(Enum): POSTGRESQL = "pg8000" MYSQL = "pymysql" SQLSERVER = "pytds" def to_sqlalchemy_dialect(self): """Map the adapter type to its corresponding SQLAlchemy dialect. Returns: str: SQLAlchemy dialect string. """ if self == DatabaseTypeAdapter.POSTGRESQL: return f"postgresql+{self.value}" elif self == DatabaseTypeAdapter.MYSQL: return f"mysql+{self.value}" elif self == DatabaseTypeAdapter.SQLSERVER: return f"mssql+{self.value}" else: raise ValueError(f"Unsupported database adapter type: {self.name}") class ConnectionConfig(ABC): @abstractmethod def get_connector_handler(self) -> Callable[[], DBAPIConnection]: pass @abstractmethod def get_db_url(self) -> str: pass @dataclass class CloudSQLConnectionConfig(ConnectionConfig): """Connects to Google Cloud SQL using Cloud SQL Python Connector. Args: db_adapter: The database adapter type (PostgreSQL, MySQL, SQL Server). instance_connection_uri: URI for connecting to the Cloud SQL instance. user: Username for authentication. password: Password for authentication. Defaults to None. db_id: Database identifier/name. refresh_strategy: Strategy for refreshing connection (default: LAZY). connector_kwargs: Additional keyword arguments for the Cloud SQL Python Connector. Enables forward compatibility. connect_kwargs: Additional keyword arguments for the client connect method. Enables forward compatibility. """ db_adapter: DatabaseTypeAdapter instance_connection_uri: str user: str = field(default_factory=str) password: str = field(default_factory=str) db_id: str = field(default_factory=str) refresh_strategy: RefreshStrategy = RefreshStrategy.LAZY connector_kwargs: dict[str, Any] = field(default_factory=dict) connect_kwargs: dict[str, Any] = field(default_factory=dict) def __post_init__(self): if not self.instance_connection_uri: raise ValueError("Instance connection URI cannot be empty") def get_connector_handler(self) -> Callable[[], DBAPIConnection]: """Returns a function that creates a new database connection. The returned connector function creates database connections that should be properly closed by the caller when no longer needed. """ cloudsql_client = CloudSQLConnector( refresh_strategy=self.refresh_strategy, **self.connector_kwargs) cloudsql_connector = lambda: cloudsql_client.connect( instance_connection_string=self.instance_connection_uri, driver=self. db_adapter.value, user=self.user, password=self.password, db=self.db_id, **self.connect_kwargs) return cloudsql_connector def get_db_url(self) -> str: return self.db_adapter.to_sqlalchemy_dialect() + "://" @dataclass class ExternalSQLDBConnectionConfig(ConnectionConfig): """Connects to External SQL DBs (PostgreSQL, MySQL, SQL Server) over TCP. Args: db_adapter: The database adapter type (PostgreSQL, MySQL, SQL Server). host: Hostname or IP address of the database server. port: Port number for the database connection. user: Username for authentication. password: Password for authentication. db_id: Database identifier/name. connect_kwargs: Additional keyword arguments for the client connect method. Enables forward compatibility. """ db_adapter: DatabaseTypeAdapter host: str port: int user: str = field(default_factory=str) password: str = field(default_factory=str) db_id: str = field(default_factory=str) connect_kwargs: dict[str, Any] = field(default_factory=dict) def __post_init__(self): if not self.host: raise ValueError("Database host cannot be empty") def get_connector_handler(self) -> Callable[[], DBAPIConnection]: """Returns a function that creates a new database connection. The returned connector function creates database connections that should be properly closed by the caller when no longer needed. """ if self.db_adapter == DatabaseTypeAdapter.POSTGRESQL: return lambda: pg8000.connect( host=self.host, port=self.port, database=self.db_id, user=self.user, password=self.password, **self.connect_kwargs) elif self.db_adapter == DatabaseTypeAdapter.MYSQL: return lambda: pymysql.connect( host=self.host, port=self.port, database=self.db_id, user=self.user, password=self.password, **self.connect_kwargs) elif self.db_adapter == DatabaseTypeAdapter.SQLSERVER: return lambda: pytds.connect( dsn=self.host, port=self.port, database=self.db_id, user=self.user, password=self.password, **self.connect_kwargs) else: raise ValueError(f"Unsupported database adapter: {self.db_adapter}") def get_db_url(self) -> str: return self.db_adapter.to_sqlalchemy_dialect() + "://" QueryConfig = Union[CustomQueryConfig, TableFieldsQueryConfig, TableFunctionQueryConfig] class CloudSQLEnrichmentHandler(EnrichmentSourceHandler[beam.Row, beam.Row]): """Enrichment handler for Cloud SQL databases. This handler is designed to work with the :class:`apache_beam.transforms.enrichment.Enrichment` transform. To use this handler, you need to provide one of the following query configs: * CustomQueryConfig - For providing a custom query function * TableFieldsQueryConfig - For specifying table, where clause, and fields * TableFunctionQueryConfig - For specifying table, where clause, and val fn By default, the handler retrieves all columns from the specified table. To limit the columns, use the `column_names` parameter to specify the desired column names. This handler queries the Cloud SQL database per element by default. To enable batching, set the `min_batch_size`, `max_batch_size`, and `max_batch_duration_secs` parameters. These values control batching behavior in the :class:`apache_beam.transforms.utils.BatchElements` transform. NOTE: Batching is not supported when using the CustomQueryConfig. """ def __init__( self, connection_config: ConnectionConfig, *, query_config: QueryConfig, column_names: Optional[list[str]] = None, min_batch_size: int = 1, max_batch_size: int = 10000, max_batch_duration_secs: Optional[float] = None, **kwargs, ): """ Example usage:: connection_config = CloudSQLConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, instance_connection_uri="apache-beam-testing:us-central1:itests", user='postgres', password= os.getenv("CLOUDSQL_PG_PASSWORD")) query_config=TableFieldsQueryConfig('my_table',"id = :param0",['id']), cloudsql_handler = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=query_config, min_batch_size=2, max_batch_size=100) Args: connection_config (ConnectionConfig): Configuration for connecting to the SQL database. Must be an instance of a subclass of `ConnectionConfig`, such as `CloudSQLConnectionConfig` or `ExternalSQLDBConnectionConfig`. This determines how the handler connects to the target SQL database. query_config: Configuration for database queries. Must be one of: * CustomQueryConfig: For providing a custom query function * TableFieldsQueryConfig: specifies table, where clause, and field names * TableFunctionQueryConfig: specifies table, where clause, and val func column_names (Optional[list[str]]): List of column names to select from the Cloud SQL table. If not provided, all columns (`*`) are selected. min_batch_size (int): Minimum number of rows to batch together when querying the database. Defaults to 1 if `query_fn` is not used. max_batch_size (int): Maximum number of rows to batch together. Defaults to 10,000 if `query_fn` is not used. max_batch_duration_secs (float): Maximum amount of time in seconds to buffer a batch before emitting it. If not provided, batching duration is determined by `BatchElements` defaults. **kwargs: Additional keyword arguments for database connection or query handling. Note: * `min_batch_size`, `max_batch_size`, and `max_batch_duration_secs` are not used with `query_fn`. * Either `where_clause_fields` or `where_clause_value_fn` must be provided for query construction if `query_fn` is not provided. * Ensure that the database user has the necessary permissions to query the specified table. """ self._connection_config = connection_config self._query_config = query_config self._column_names = ",".join(column_names) if column_names else "*" self.kwargs = kwargs self._batching_kwargs = {} table_query_configs = (TableFieldsQueryConfig, TableFunctionQueryConfig) if isinstance(query_config, table_query_configs): self.query_template = ( f"SELECT {self._column_names} " f"FROM {query_config.table_id} " f"WHERE {query_config.where_clause_template}") self._batching_kwargs['min_batch_size'] = min_batch_size self._batching_kwargs['max_batch_size'] = max_batch_size if max_batch_duration_secs is not None: self._batching_kwargs[ 'max_batch_duration_secs'] = max_batch_duration_secs def __enter__(self): connector = self._connection_config.get_connector_handler() self._engine = create_engine( url=self._connection_config.get_db_url(), creator=connector) def __call__( self, request: Union[beam.Row, list[beam.Row]], *_args, **_kwargs): """Handle requests by delegating to single or batch processing.""" if isinstance(request, list): return self._process_batch_request(request) else: return self._process_single_request(request) def _process_single_request(self, request: beam.Row): """Process a single request and return with its response.""" response: Union[list[dict[str, Any]], dict[str, Any]] if isinstance(self._query_config, CustomQueryConfig): query = self._query_config.query_fn(request) response = self._execute_query(query, is_batch=False) else: values = self._extract_values_from_request(request) param_dict = self._build_single_param_dict(values) response = self._execute_query( self.query_template, params=param_dict, is_batch=False) return request, beam.Row(**response) # type: ignore[arg-type] def _process_batch_request(self, requests: list[beam.Row]): """Process batch requests and match responses to original requests.""" values, responses = [], [] requests_map: dict[Any, Any] = {} batch_size = len(requests) # Build the appropriate query (single or batched). raw_query = self._build_batch_query(requests, batch_size) # Extract where_clause_fields values and map the generated request key to # the original request object.. for req in requests: current_values = self._extract_values_from_request(req) values.extend(current_values) requests_map[self.create_row_key(req)] = req # Build named parameters dictionary for parameterized query. param_dict = self._build_parameters_dict(requests, batch_size) # Execute the parameterized query with validated parameters. result: Union[list[dict[str, Any]], dict[str, Any]] = self._execute_query( raw_query, params=param_dict, is_batch=True) for response in result: response_row = beam.Row(**response) # type: ignore[arg-type] response_key = self.create_row_key(response_row) if response_key in requests_map: responses.append((requests_map[response_key], response_row)) return responses def _execute_query( self, query: str, params: Optional[dict] = None, is_batch: bool = False) -> Union[list[dict[str, Any]], dict[str, Any]]: connection = None try: connection = self._engine.connect() transaction = connection.begin() try: if params: result = connection.execute(text(query), params) else: result = connection.execute(text(query)) # Materialize results while transaction is active. data: Union[list[dict[str, Any]], dict[str, Any]] if is_batch: data = [row._asdict() for row in result] else: result_row = result.first() data = result_row._asdict() if result_row else {} # Explicitly commit the transaction. transaction.commit() return data except Exception as e: transaction.rollback() raise RuntimeError(f"Database operation failed: {e}") from e except Exception as e: raise Exception( f'Could not execute the query. Please check if the query is properly ' f'formatted and the table exists. {e}') from e finally: if connection: connection.close() def _build_batch_query( self, requests: list[beam.Row], batch_size: int) -> str: """Build batched query with unique parameter names for multiple requests. This method extracts parameter placeholders from the where_clause_template using regex and creates unique parameter names for each batch item. The parameter names in the template can be any valid identifiers (e.g., :id, :param_0, :user_name) and don't need to match field names exactly. For batch queries, placeholders are replaced with unique names like :batch_0_id, :batch_1_param_0, etc., based on the actual parameter names found in the template. Args: requests: List of beam.Row requests to process batch_size: Number of requests in the batch Returns: SQL query string with batched WHERE clauses using unique parameter names """ # Single request - return original query. if batch_size <= 1: return self.query_template # Only batch table-based query configs. table_query_configs = (TableFieldsQueryConfig, TableFunctionQueryConfig) if not isinstance(self._query_config, table_query_configs): return self.query_template # Build batched WHERE clauses. where_clauses = [self._create_batch_clause(i) for i in range(batch_size)] # Combine clauses and update query. where_clause_batched = ' OR '.join(where_clauses) # We know this is a table-based config from the check above. assert isinstance(self._query_config, table_query_configs) return self.query_template.replace( self._query_config.where_clause_template, where_clause_batched) def _create_batch_clause(self, batch_index: int) -> str: """Create a WHERE clause for a single batch item with unique parameter names.""" # This method is only called for table-based query configs table_query_configs = (TableFieldsQueryConfig, TableFunctionQueryConfig) assert isinstance(self._query_config, table_query_configs) clause = self._query_config.where_clause_template # Extract parameter names from the template using regex. param_names = self._extract_parameter_names( self._query_config.where_clause_template) for param_name in param_names: old_param = f':{param_name}' new_param = f':batch_{batch_index}_{param_name}' clause = clause.replace(old_param, new_param) return f'({clause})' def _build_parameters_dict( self, requests: list[beam.Row], batch_size: int) -> dict: """Build named parameters dictionary for parameterized queries. Args: requests: List of beam.Row requests to process batch_size: Number of requests in the batch Returns: Dictionary mapping parameter names to validated values """ param_dict = {} for i, req in enumerate(requests): current_values = self._extract_values_from_request(req) # For batched queries, use unique parameter names per batch item. if batch_size > 1: # Batching is only used with table-based query configs. table_query_configs = (TableFieldsQueryConfig, TableFunctionQueryConfig) assert isinstance(self._query_config, table_query_configs) batch_param_dict = self._build_single_param_dict(current_values) # Prefix batch parameters to make them globally unique. for param_name, val in batch_param_dict.items(): param_dict[f'batch_{i}_{param_name}'] = val else: single_param_dict = self._build_single_param_dict(current_values) param_dict.update(single_param_dict) return param_dict def _build_single_param_dict(self, values: list[Any]) -> dict[str, Any]: """Build parameter dictionary for single request processing. Args: values: List of parameter values Returns: Dictionary mapping parameter names to values """ table_query_configs = (TableFieldsQueryConfig, TableFunctionQueryConfig) if not isinstance(self._query_config, table_query_configs): raise ValueError( f"Parameter binding not supported for " f"{type(self._query_config).__name__}") _, param_dict = self._get_unique_template_and_params( self._query_config.where_clause_template, values) return param_dict def _get_unique_template_and_params( self, template: str, values: list[Any]) -> tuple[str, dict[str, Any]]: """Generate unique binding parameter names for duplicate templates. Args: template: SQL template with potentially duplicate binding parameter names values: List of parameter values Returns: Tuple of (updated_template, param_dict) with unique binding names. """ param_names = self._extract_parameter_names(template) unique_param_names = [ f"{param_name}_{i}" if param_names.count(param_name) > 1 else param_name for i, param_name in enumerate(param_names) ] # Update template by replacing each parameter occurrence in order. updated_template = template param_positions = [] # Find all parameter positions. for match in re.finditer(r':(\w+)', template): param_positions.append((match.start(), match.end(), match.group(1))) # Replace parameters from right to left to avoid position shifts. for i in reversed(range(len(param_positions))): start, end, _ = param_positions[i] unique_name = unique_param_names[i] updated_template = ( updated_template[:start] + f':{unique_name}' + updated_template[end:]) # Build parameter dictionary. param_dict = { unique_name: val for unique_name, val in zip(unique_param_names, values) } return updated_template, param_dict def _extract_values_from_request(self, request: beam.Row) -> list[Any]: """Extract parameter values from a request based on query configuration. Args: request: The beam.Row request to extract values from Returns: List of parameter values Raises: KeyError: If required fields are missing from the request """ try: if isinstance(self._query_config, TableFunctionQueryConfig): return [ val for val in self._query_config.where_clause_value_fn(request) ] elif isinstance(self._query_config, TableFieldsQueryConfig): request_dict = request._asdict() return [ request_dict[field] for field in self._query_config.where_clause_fields ] else: raise ValueError("Unsupported query configuration type") except KeyError as e: raise KeyError( "Make sure the values passed in `where_clause_fields` are " "the keys in the input `beam.Row`." + str(e)) def _extract_parameter_names(self, template: str) -> list[str]: """Extract parameter names from a SQL template string. Args: template: SQL template string with named parameters (e.g., "id = :id") Returns: List of parameter names found in the template (e.g., ["id"]) """ return re.findall(r':(\w+)', template) def create_row_key(self, row: beam.Row): if isinstance(self._query_config, TableFunctionQueryConfig): return tuple(self._query_config.where_clause_value_fn(row)) if isinstance(self._query_config, TableFieldsQueryConfig): row_dict = row._asdict() return ( tuple( row_dict[where_clause_field] for where_clause_field in self._query_config.where_clause_fields)) raise ValueError( "Either where_clause_fields or where_clause_value_fn must be specified") def get_cache_key(self, request: Union[beam.Row, list[beam.Row]]): if isinstance(self._query_config, CustomQueryConfig): raise NotImplementedError( "Caching is not supported for CustomQueryConfig. " "Consider using TableFieldsQueryConfig or " + "TableFunctionQueryConfig instead.") if isinstance(request, list): cache_keys = [] for req in request: req_dict = req._asdict() try: if isinstance(self._query_config, TableFunctionQueryConfig): current_values = self._query_config.where_clause_value_fn(req) elif isinstance(self._query_config, TableFieldsQueryConfig): current_values = [ req_dict[field] for field in self._query_config.where_clause_fields ] key = ';'.join(map(repr, current_values)) cache_keys.append(key) except KeyError as e: raise KeyError( "Make sure the values passed in `where_clause_fields` are the " "keys in the input `beam.Row`." + str(e)) return cache_keys else: req_dict = request._asdict() try: if isinstance(self._query_config, TableFunctionQueryConfig): current_values = self._query_config.where_clause_value_fn(request) else: # TableFieldsQueryConfig. current_values = [ req_dict[field] for field in self._query_config.where_clause_fields ] key = ";".join(["%s"] * len(current_values)) cache_key = key % tuple(current_values) except KeyError as e: raise KeyError( "Make sure the values passed in `where_clause_fields` are the " "keys in the input `beam.Row`." + str(e)) return cache_key def __exit__(self, _exc_type, _exc_val, _exc_tb): self._engine.dispose(close=True) self._engine = None def batch_elements_kwargs(self) -> Mapping[str, Any]: """Returns a kwargs suitable for `beam.BatchElements`.""" return self._batching_kwargs ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/cloudsql_it_test.py0000644000175100017510000005177415206113033030041 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import functools import logging import os import unittest import uuid from dataclasses import dataclass from typing import Optional from unittest.mock import MagicMock import pytest import apache_beam as beam from apache_beam.coders import coders from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to # pylint: disable=ungrouped-imports try: from sqlalchemy import VARCHAR from sqlalchemy import Column from sqlalchemy import Engine from sqlalchemy import Integer from sqlalchemy import MetaData from sqlalchemy import Table from sqlalchemy import create_engine from testcontainers.core.generic import DbContainer from testcontainers.mssql import SqlServerContainer from testcontainers.mysql import MySqlContainer from testcontainers.postgres import PostgresContainer from testcontainers.redis import RedisContainer from apache_beam.transforms.enrichment import Enrichment from apache_beam.transforms.enrichment_handlers.cloudsql import CloudSQLConnectionConfig from apache_beam.transforms.enrichment_handlers.cloudsql import CloudSQLEnrichmentHandler from apache_beam.transforms.enrichment_handlers.cloudsql import ConnectionConfig from apache_beam.transforms.enrichment_handlers.cloudsql import CustomQueryConfig from apache_beam.transforms.enrichment_handlers.cloudsql import DatabaseTypeAdapter from apache_beam.transforms.enrichment_handlers.cloudsql import ExternalSQLDBConnectionConfig from apache_beam.transforms.enrichment_handlers.cloudsql import TableFieldsQueryConfig from apache_beam.transforms.enrichment_handlers.cloudsql import TableFunctionQueryConfig except ImportError as e: raise unittest.SkipTest(f'CloudSQL dependencies not installed: {str(e)}') _LOGGER = logging.getLogger(__name__) def where_clause_value_fn(row: beam.Row): return [row.id] # type: ignore[attr-defined] def query_fn(table, row: beam.Row): return f"SELECT * FROM {table} WHERE id = {row.id}" # type: ignore[attr-defined] @dataclass class SQLDBContainerInfo: adapter: DatabaseTypeAdapter container: DbContainer host: str port: int user: str password: str id: str @property def address(self) -> str: return f"{self.host}:{self.port}" @property def url(self) -> str: return self.adapter.to_sqlalchemy_dialect() + "://" class SQLEnrichmentTestHelper: @staticmethod def start_sql_db_container( database_type: DatabaseTypeAdapter, sql_client_retries=3) -> Optional[SQLDBContainerInfo]: info = None for i in range(sql_client_retries): sql_db_container = DbContainer("") try: if database_type == DatabaseTypeAdapter.POSTGRESQL: user, password, db_id = "test", "test", "test" sql_db_container = PostgresContainer( image="postgres:16", username=user, password=password, dbname=db_id, driver=database_type.value) sql_db_container.start() host = sql_db_container.get_container_host_ip() port = int(sql_db_container.get_exposed_port(5432)) elif database_type == DatabaseTypeAdapter.MYSQL: user, password, db_id = "test", "test", "test" sql_db_container = MySqlContainer( image="mysql:8.0", username=user, root_password=password, password=password, dbname=db_id) sql_db_container.start() host = sql_db_container.get_container_host_ip() port = int(sql_db_container.get_exposed_port(3306)) elif database_type == DatabaseTypeAdapter.SQLSERVER: user, password, db_id = "SA", "A_Str0ng_Required_Password", "tempdb" sql_db_container = SqlServerContainer( image="mcr.microsoft.com/mssql/server:2022-latest", username=user, password=password, dbname=db_id, dialect=database_type.to_sqlalchemy_dialect()) sql_db_container.start() host = sql_db_container.get_container_host_ip() port = int(sql_db_container.get_exposed_port(1433)) else: raise ValueError(f"Unsupported database type: {database_type}") info = SQLDBContainerInfo( adapter=database_type, container=sql_db_container, host=host, port=port, user=user, password=password, id=db_id) _LOGGER.info( "%s container started successfully on %s.", database_type.name, info.address) break except Exception as e: stdout_logs, stderr_logs = sql_db_container.get_logs() stdout_logs = stdout_logs.decode("utf-8") stderr_logs = stderr_logs.decode("utf-8") _LOGGER.warning( "Retry %d/%d: Failed to start %s container. Reason: %s. " "STDOUT logs:\n%s\nSTDERR logs:\n%s", i + 1, sql_client_retries, database_type.name, e, stdout_logs, stderr_logs) if i == sql_client_retries - 1: _LOGGER.error( "Unable to start %s container for I/O tests after %d " "retries. Tests cannot proceed. STDOUT logs:\n%s\n" "STDERR logs:\n%s", database_type.name, sql_client_retries, stdout_logs, stderr_logs) raise e return info @staticmethod def stop_sql_db_container(db_info: SQLDBContainerInfo): try: _LOGGER.debug("Stopping %s container.", db_info.adapter.name) db_info.container.stop() _LOGGER.info("%s container stopped successfully.", db_info.adapter.name) except Exception as e: _LOGGER.warning( "Error encountered while stopping %s container: %s", db_info.adapter.name, e) @staticmethod def create_table( table_id: str, engine: Engine, columns: list[Column], table_data: list[dict], metadata: MetaData): # Create table metadata. table = Table(table_id, metadata, *columns) # Create contextual connection for schema creation. with engine.connect() as schema_connection: try: metadata.create_all(schema_connection) schema_connection.commit() except Exception as e: schema_connection.rollback() raise RuntimeError(f"Failed to create table schema: {e}") # Now create a separate contextual connection for data insertion. with engine.connect() as connection: try: connection.execute(table.insert(), table_data) connection.commit() except Exception as e: connection.rollback() raise Exception(f"Failed to insert table data: {e}") class BaseTestSQLEnrichment(unittest.TestCase): _cache_client_retries = 3 _table_data = [ { "id": 1, "name": "A", 'quantity': 2, 'distribution_center_id': 3 }, { "id": 2, "name": "B", 'quantity': 3, 'distribution_center_id': 1 }, { "id": 3, "name": "C", 'quantity': 10, 'distribution_center_id': 4 }, { "id": 4, "name": "D", 'quantity': 1, 'distribution_center_id': 3 }, { "id": 5, "name": "C", 'quantity': 100, 'distribution_center_id': 4 }, { "id": 6, "name": "D", 'quantity': 11, 'distribution_center_id': 3 }, { "id": 7, "name": "C", 'quantity': 7, 'distribution_center_id': 1 }, { "id": 8, "name": "D", 'quantity': 4, 'distribution_center_id': 1 }, ] @classmethod def setUpClass(cls): if not hasattr(cls, '_connection_config') or not hasattr(cls, '_metadata'): # Skip setup for the base class. raise unittest.SkipTest( "Base class - no connection_config or metadata defined") # Type hint data from subclasses. cls._table_id: str cls._connection_config: ConnectionConfig cls._metadata: MetaData connector = cls._connection_config.get_connector_handler() cls._engine = create_engine( url=cls._connection_config.get_db_url(), creator=connector) SQLEnrichmentTestHelper.create_table( table_id=cls._table_id, engine=cls._engine, columns=cls.get_columns(), table_data=cls._table_data, metadata=cls._metadata) @classmethod def get_columns(cls): """Returns fresh column objects each time it's called.""" return [ Column("id", Integer, nullable=False), Column("name", VARCHAR(255), nullable=False), Column("quantity", Integer, nullable=False), Column("distribution_center_id", Integer, nullable=False), ] @pytest.fixture def cache_container(self): self._start_cache_container() # Hand control to the test. yield self._cache_container.stop() self._cache_container = None def _start_cache_container(self): for i in range(self._cache_client_retries): try: self._cache_container = RedisContainer(image="redis:7.2.4") self._cache_container.start() host = self._cache_container.get_container_host_ip() port = self._cache_container.get_exposed_port(6379) self._cache_container_host = host self._cache_container_port = port self._cache_client = self._cache_container.get_client() break except Exception as e: if i == self._cache_client_retries - 1: _LOGGER.error( "Unable to start redis container for RRIO tests after " "%d retries.", self._cache_client_retries) raise e @classmethod def tearDownClass(cls): # Drop all tables using metadata as the primary approach. cls._metadata.drop_all(cls._engine) # Fallback to raw SQL drop if needed. try: with cls._engine.connect() as conn: conn.execute(f"DROP TABLE IF EXISTS {cls._table_id}") conn.commit() _LOGGER.info("Dropped table %s", cls._table_id) except Exception as e: _LOGGER.warning("Failed to drop table %s: %s", cls._table_id, e) cls._engine.dispose(close=True) cls._engine = None def test_sql_enrichment(self): expected_rows = [ beam.Row(id=1, name="A", quantity=2, distribution_center_id=3), beam.Row(id=2, name="B", quantity=3, distribution_center_id=1) ] fields = ['id'] requests = [ beam.Row(id=1, name='A'), beam.Row(id=2, name='B'), ] query_config = TableFieldsQueryConfig( table_id=self._table_id, where_clause_template="id = :id_param", where_clause_fields=fields) handler = CloudSQLEnrichmentHandler( connection_config=self._connection_config, query_config=query_config, min_batch_size=1, max_batch_size=100, ) with TestPipeline() as test_pipeline: pcoll = (test_pipeline | beam.Create(requests) | Enrichment(handler)) assert_that(pcoll, equal_to(expected_rows)) def test_sql_enrichment_batched(self): expected_rows = [ beam.Row(id=1, name="A", quantity=2, distribution_center_id=3), beam.Row(id=2, name="B", quantity=3, distribution_center_id=1) ] fields = ['id'] requests = [ beam.Row(id=1, name='A'), beam.Row(id=2, name='B'), ] query_config = TableFieldsQueryConfig( table_id=self._table_id, where_clause_template="id = :id", where_clause_fields=fields) handler = CloudSQLEnrichmentHandler( connection_config=self._connection_config, query_config=query_config, min_batch_size=2, max_batch_size=100, ) with TestPipeline() as test_pipeline: pcoll = (test_pipeline | beam.Create(requests) | Enrichment(handler)) assert_that(pcoll, equal_to(expected_rows)) def test_sql_enrichment_batched_multiple_fields(self): expected_rows = [ beam.Row(id=1, distribution_center_id=3, name="A", quantity=2), beam.Row(id=2, distribution_center_id=1, name="B", quantity=3) ] fields = ['id', 'distribution_center_id'] requests = [ beam.Row(id=1, distribution_center_id=3), beam.Row(id=2, distribution_center_id=1), ] query_config = TableFieldsQueryConfig( table_id=self._table_id, where_clause_template="id = :id AND distribution_center_id = :param_1", where_clause_fields=fields) handler = CloudSQLEnrichmentHandler( connection_config=self._connection_config, query_config=query_config, min_batch_size=8, max_batch_size=100, ) with TestPipeline() as test_pipeline: pcoll = (test_pipeline | beam.Create(requests) | Enrichment(handler)) assert_that(pcoll, equal_to(expected_rows)) def test_sql_enrichment_with_query_fn(self): expected_rows = [ beam.Row(id=1, name="A", quantity=2, distribution_center_id=3), beam.Row(id=2, name="B", quantity=3, distribution_center_id=1) ] requests = [ beam.Row(id=1, name='A'), beam.Row(id=2, name='B'), ] fn = functools.partial(query_fn, self._table_id) query_config = CustomQueryConfig(query_fn=fn) handler = CloudSQLEnrichmentHandler( connection_config=self._connection_config, query_config=query_config) with TestPipeline() as test_pipeline: pcoll = (test_pipeline | beam.Create(requests) | Enrichment(handler)) assert_that(pcoll, equal_to(expected_rows)) def test_sql_enrichment_with_condition_value_fn(self): expected_rows = [ beam.Row(id=1, name="A", quantity=2, distribution_center_id=3), beam.Row(id=2, name="B", quantity=3, distribution_center_id=1) ] requests = [ beam.Row(id=1, name='A'), beam.Row(id=2, name='B'), ] query_config = TableFunctionQueryConfig( table_id=self._table_id, where_clause_template="id = :param_0", where_clause_value_fn=where_clause_value_fn) handler = CloudSQLEnrichmentHandler( connection_config=self._connection_config, query_config=query_config, min_batch_size=2, max_batch_size=100) with TestPipeline() as test_pipeline: pcoll = (test_pipeline | beam.Create(requests) | Enrichment(handler)) assert_that(pcoll, equal_to(expected_rows)) def test_sql_enrichment_on_non_existent_table(self): requests = [ beam.Row(id=1, name='A'), beam.Row(id=2, name='B'), ] query_config = TableFunctionQueryConfig( table_id=self._table_id, where_clause_template="id = :id", where_clause_value_fn=where_clause_value_fn) handler = CloudSQLEnrichmentHandler( connection_config=self._connection_config, query_config=query_config, column_names=["wrong_column"], ) with self.assertRaises(Exception) as context: with TestPipeline() as p: _ = (p | beam.Create(requests) | Enrichment(handler)) expect_err_msg_contains = ( "Could not execute the query. Please check if the query is properly " "formatted and the table exists.") self.assertIn(expect_err_msg_contains, str(context.exception)) @pytest.mark.usefixtures("cache_container") def test_sql_enrichment_with_redis(self): requests = [ beam.Row(id=1, name='A'), beam.Row(id=2, name='B'), ] expected_rows = [ beam.Row(id=1, name="A", quantity=2, distribution_center_id=3), beam.Row(id=2, name="B", quantity=3, distribution_center_id=1) ] query_config = TableFunctionQueryConfig( table_id=self._table_id, where_clause_template="id = :param_0", where_clause_value_fn=where_clause_value_fn) handler = CloudSQLEnrichmentHandler( connection_config=self._connection_config, query_config=query_config, min_batch_size=2, max_batch_size=100) with TestPipeline() as test_pipeline: pcoll_populate_cache = ( test_pipeline | beam.Create(requests) | Enrichment(handler).with_redis_cache( self._cache_container_host, self._cache_container_port)) assert_that(pcoll_populate_cache, equal_to(expected_rows)) # Manually check cache entry to verify entries were correctly stored. c = coders.StrUtf8Coder() for req in requests: key = handler.get_cache_key(req) response = self._cache_client.get(c.encode(key)) if not response: raise ValueError("No cache entry found for %s" % key) # Mocks the CloudSQL enrichment handler to prevent actual database calls. # This ensures that a cache hit scenario does not trigger any database # interaction, raising an exception if an unexpected call occurs. actual = CloudSQLEnrichmentHandler.__call__ CloudSQLEnrichmentHandler.__call__ = MagicMock( side_effect=Exception("Database should not be called on a cache hit.")) # Run a second pipeline to verify cache is being used. with TestPipeline() as test_pipeline: pcoll_cached = ( test_pipeline | beam.Create(requests) | Enrichment(handler).with_redis_cache( self._cache_container_host, self._cache_container_port)) assert_that(pcoll_cached, equal_to(expected_rows)) # Restore the original CloudSQL enrichment handler implementation. CloudSQLEnrichmentHandler.__call__ = actual class BaseCloudSQLDBEnrichment(BaseTestSQLEnrichment): @classmethod def setUpClass(cls): if not hasattr(cls, '_db_adapter'): # Skip setup for the base class. raise unittest.SkipTest("Base class - no db_adapter defined") # Type hint data from subclasses. cls._db_adapter: DatabaseTypeAdapter cls._instance_connection_uri: str cls._user: str cls._password: str cls._db_id: str cls._connection_config = CloudSQLConnectionConfig( db_adapter=cls._db_adapter, instance_connection_uri=cls._instance_connection_uri, user=cls._user, password=cls._password, db_id=cls._db_id) super().setUpClass() @classmethod def tearDownClass(cls): super().tearDownClass() @unittest.skipUnless( os.environ.get('ALLOYDB_PASSWORD'), "ALLOYDB_PASSWORD environment var is not provided") class TestCloudSQLPostgresEnrichment(BaseCloudSQLDBEnrichment): _db_adapter = DatabaseTypeAdapter.POSTGRESQL # Configuration required for locating the CloudSQL instance. _unique_suffix = str(uuid.uuid4())[:8] _table_id = f"product_details_cloudsql_pg_enrichment_{_unique_suffix}" _gcp_project_id = "apache-beam-testing" _region = "us-central1" _instance_name = "beam-integration-tests" _instance_connection_uri = f"{_gcp_project_id}:{_region}:{_instance_name}" # Configuration required for authenticating to the CloudSQL instance. _user = "postgres" _password = os.getenv("ALLOYDB_PASSWORD") _db_id = "postgres" _metadata = MetaData() class BaseExternalSQLDBEnrichment(BaseTestSQLEnrichment): @classmethod def setUpClass(cls): if not hasattr(cls, '_db_adapter'): # Skip setup for the base class. raise unittest.SkipTest("Base class - no db_adapter defined") # Type hint data from subclasses. cls._db_adapter: DatabaseTypeAdapter cls._db = SQLEnrichmentTestHelper.start_sql_db_container(cls._db_adapter) cls._connection_config = ExternalSQLDBConnectionConfig( db_adapter=cls._db_adapter, host=cls._db.host, user=cls._db.user, password=cls._db.password, db_id=cls._db.id, port=cls._db.port) super().setUpClass() @classmethod def tearDownClass(cls): super().tearDownClass() SQLEnrichmentTestHelper.stop_sql_db_container(cls._db) cls._db = None class TestExternalPostgresEnrichment(BaseExternalSQLDBEnrichment): _db_adapter = DatabaseTypeAdapter.POSTGRESQL _unique_suffix = str(uuid.uuid4())[:8] _table_id = f"product_details_external_pg_enrichment_{_unique_suffix}" _metadata = MetaData() class TestExternalMySQLEnrichment(BaseExternalSQLDBEnrichment): _db_adapter = DatabaseTypeAdapter.MYSQL _unique_suffix = str(uuid.uuid4())[:8] _table_id = f"product_details_external_mysql_enrichment_{_unique_suffix}" _metadata = MetaData() class TestExternalSQLServerEnrichment(BaseExternalSQLDBEnrichment): _db_adapter = DatabaseTypeAdapter.SQLSERVER _unique_suffix = str(uuid.uuid4())[:8] _table_id = f"product_details_external_mssql_enrichment_{_unique_suffix}" _metadata = MetaData() if __name__ == "__main__": unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/cloudsql_test.py0000644000175100017510000005264515206113033027343 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest from parameterized import parameterized # pylint: disable=ungrouped-imports try: from apache_beam.transforms.enrichment_handlers.cloudsql import CloudSQLConnectionConfig from apache_beam.transforms.enrichment_handlers.cloudsql import CloudSQLEnrichmentHandler from apache_beam.transforms.enrichment_handlers.cloudsql import CustomQueryConfig from apache_beam.transforms.enrichment_handlers.cloudsql import DatabaseTypeAdapter from apache_beam.transforms.enrichment_handlers.cloudsql import ExternalSQLDBConnectionConfig from apache_beam.transforms.enrichment_handlers.cloudsql import TableFieldsQueryConfig from apache_beam.transforms.enrichment_handlers.cloudsql import TableFunctionQueryConfig from apache_beam.transforms.enrichment_handlers.cloudsql_it_test import query_fn from apache_beam.transforms.enrichment_handlers.cloudsql_it_test import where_clause_value_fn except ImportError as e: raise unittest.SkipTest(f'CloudSQL dependencies not installed: {str(e)}') class TestCloudSQLEnrichment(unittest.TestCase): def test_invalid_external_db_connection_params(self): with self.assertRaises(ValueError) as context: _ = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='', port=5432, user='', password='', db_id='') self.assertIn("Database host cannot be empty", str(context.exception)) def test_invalid_cloudsql_db_connection_params(self): with self.assertRaises(ValueError) as context: _ = CloudSQLConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, instance_connection_uri='', user='', password='', db_id='') self.assertIn( "Instance connection URI cannot be empty", str(context.exception)) @parameterized.expand([ # Empty TableFieldsQueryConfig. ( lambda: TableFieldsQueryConfig( table_id="", where_clause_template="", where_clause_fields=[]), 1, 2, "must provide table_id and where_clause_template" ), # Missing where_clause_template in TableFieldsQueryConfig. ( lambda: TableFieldsQueryConfig( table_id="table", where_clause_template="", where_clause_fields=["id"]), 2, 10, "must provide table_id and where_clause_template" ), # Invalid CustomQueryConfig with None query_fn. ( lambda: CustomQueryConfig(query_fn=None), # type: ignore[arg-type] 2, 10, "must provide a valid query_fn" ), # Missing table_id in TableFunctionQueryConfig. ( lambda: TableFunctionQueryConfig( table_id="", where_clause_template="id='{}'", where_clause_value_fn=where_clause_value_fn), 2, 10, "must provide table_id and where_clause_template" ), # Missing where_clause_fields in TableFieldsQueryConfig. ( lambda: TableFieldsQueryConfig( table_id="table", where_clause_template="id = '{}'", where_clause_fields=[]), 1, 10, "must provide non-empty where_clause_fields" ), # Missing where_clause_value_fn in TableFunctionQueryConfig. ( lambda: TableFunctionQueryConfig( table_id="table", where_clause_template="id = '{}'", where_clause_value_fn=None), # type: ignore[arg-type] 1, 10, "must provide where_clause_value_fn" ), ]) def test_invalid_query_config( self, create_config, min_batch_size, max_batch_size, expected_error_msg): """Test that validation errors are raised for invalid query configs. The test verifies both that the appropriate ValueError is raised and that the error message contains the expected text. """ with self.assertRaises(ValueError) as context: # Call the lambda to create the config. query_config = create_config() connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='', password='', db_id='') _ = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=query_config, min_batch_size=min_batch_size, max_batch_size=max_batch_size, ) # Verify the error message contains the expected text. self.assertIn(expected_error_msg, str(context.exception)) def test_valid_query_configs(self): """Test valid query configuration cases.""" # Valid TableFieldsQueryConfig. table_fields_config = TableFieldsQueryConfig( table_id="my_table", where_clause_template="id = :id", where_clause_fields=["id"]) connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') handler1 = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=table_fields_config, min_batch_size=1, max_batch_size=10) self.assertEqual( handler1.query_template, "SELECT * FROM my_table WHERE id = :id") # Valid TableFunctionQueryConfig. table_function_config = TableFunctionQueryConfig( table_id="my_table", where_clause_template="id = :id", where_clause_value_fn=where_clause_value_fn) handler2 = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=table_function_config, min_batch_size=1, max_batch_size=10) self.assertEqual( handler2.query_template, "SELECT * FROM my_table WHERE id = :id") # Valid CustomQueryConfig. custom_config = CustomQueryConfig(query_fn=query_fn) handler3 = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=custom_config) # Verify that batching kwargs are empty for CustomQueryConfig. self.assertEqual(handler3.batch_elements_kwargs(), {}) def test_custom_query_config_cache_key_error(self): """Test get_cache_key raises NotImplementedError with CustomQueryConfig.""" custom_config = CustomQueryConfig(query_fn=query_fn) connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') handler = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=custom_config) # Create a dummy request. import apache_beam as beam request = beam.Row(id=1) # Verify that get_cache_key raises NotImplementedError. with self.assertRaises(NotImplementedError): handler.get_cache_key(request) def test_batch_elements_kwargs_include_max_batch_duration_secs(self): connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') query_config = TableFieldsQueryConfig( table_id='my_table', where_clause_template='id = :id', where_clause_fields=['id']) handler = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=query_config, min_batch_size=2, max_batch_size=10, max_batch_duration_secs=0.5) self.assertEqual( handler.batch_elements_kwargs(), { 'min_batch_size': 2, 'max_batch_size': 10, 'max_batch_duration_secs': 0.5, }) def test_batch_elements_kwargs_omit_max_batch_duration_secs_by_default(self): connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') query_config = TableFieldsQueryConfig( table_id='my_table', where_clause_template='id = :id', where_clause_fields=['id']) handler = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=query_config, min_batch_size=2, max_batch_size=10) self.assertEqual( handler.batch_elements_kwargs(), { 'min_batch_size': 2, 'max_batch_size': 10, }) def test_extract_parameter_names(self): """Test parameter extraction from SQL templates.""" connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') config = TableFieldsQueryConfig( table_id="users", where_clause_template="id = :id", where_clause_fields=["id"]) handler = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=config) # Test simple parameter extraction. self.assertEqual(handler._extract_parameter_names("id = :id"), ["id"]) # Test multiple parameters. self.assertEqual( handler._extract_parameter_names("id = :id AND name = :name"), ["id", "name"]) # Test no parameters. self.assertEqual( handler._extract_parameter_names("SELECT * FROM users"), []) # Test complex query. self.assertEqual( handler._extract_parameter_names( "age > :min_age AND city = :city AND status = :status"), ["min_age", "city", "status"]) def test_build_single_param_dict(self): """Test building parameter dictionaries for single requests.""" connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') # Test TableFieldsQueryConfig. table_config = TableFieldsQueryConfig( table_id="users", where_clause_template="id = :id AND name = :name", where_clause_fields=["id", "name"]) handler1 = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=table_config) result1 = handler1._build_single_param_dict([123, "John"]) self.assertEqual(result1, {"id": 123, "name": "John"}) # Test TableFunctionQueryConfig. table_func_config = TableFunctionQueryConfig( table_id="users", where_clause_template="age > :min_age AND city = :city", where_clause_value_fn=lambda row: [row.min_age, row.city]) handler2 = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=table_func_config) result2 = handler2._build_single_param_dict([3, "NYC"]) self.assertEqual(result2, {"min_age": 3, "city": "NYC"}) def test_extract_values_from_request(self): """Test extracting values from requests based on query configuration.""" connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') # Test TableFieldsQueryConfig. table_config = TableFieldsQueryConfig( table_id="users", where_clause_template="id = :id AND name = :name", where_clause_fields=["id", "name"]) handler1 = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=table_config) import apache_beam as beam request1 = beam.Row(id=123, name="John", age=30) result1 = handler1._extract_values_from_request(request1) self.assertEqual(result1, [123, "John"]) # Test TableFunctionQueryConfig. def test_value_fn(row): return [row.age, row.city] table_func_config = TableFunctionQueryConfig( table_id="users", where_clause_template="age > :min_age AND city = :city", where_clause_value_fn=test_value_fn) handler2 = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=table_func_config) request2 = beam.Row(age=25, city="NYC", name="Jane") result2 = handler2._extract_values_from_request(request2) self.assertEqual(result2, [25, "NYC"]) # Test missing field error. request3 = beam.Row(age=30) # Missing name field. with self.assertRaises(KeyError) as context: handler1._extract_values_from_request(request3) self.assertIn("where_clause_fields", str(context.exception)) def test_build_batch_query_single_request(self): """Test batch query building with single request returns original query.""" connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') config = TableFieldsQueryConfig( table_id="users", where_clause_template="id = :id", where_clause_fields=["id"]) handler = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=config) import apache_beam as beam requests = [beam.Row(id=1)] # Single request should return original query template. result = handler._build_batch_query(requests, batch_size=1) self.assertEqual(result, "SELECT * FROM users WHERE id = :id") def test_build_batch_query_multiple_requests(self): """Test batch query building with multiple requests creates OR clauses.""" connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') config = TableFieldsQueryConfig( table_id="users", where_clause_template="id = :id AND name = :name", where_clause_fields=["id", "name"]) handler = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=config) import apache_beam as beam requests = [beam.Row(id=1, name="Alice"), beam.Row(id=2, name="Bob")] result = handler._build_batch_query(requests, batch_size=2) expected = ( "SELECT * FROM users WHERE " "(id = :batch_0_id AND name = :batch_0_name) OR " "(id = :batch_1_id AND name = :batch_1_name)") self.assertEqual(result, expected) def test_build_parameters_dict_batch(self): """Test parameter dictionary building for batch requests.""" connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') config = TableFieldsQueryConfig( table_id="users", where_clause_template="id = :id", where_clause_fields=["id"]) handler = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=config) import apache_beam as beam requests = [beam.Row(id=1), beam.Row(id=2)] result = handler._build_parameters_dict(requests, batch_size=2) expected = {"batch_0_id": 1, "batch_1_id": 2} self.assertEqual(result, expected) def test_create_batch_clause(self): """Test batch clause creation with unique parameter names.""" connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') config = TableFieldsQueryConfig( table_id="users", where_clause_template="id = :id AND name = :name", where_clause_fields=["id", "name"]) handler = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=config) result = handler._create_batch_clause(batch_index=0) expected = "(id = :batch_0_id AND name = :batch_0_name)" self.assertEqual(result, expected) result2 = handler._create_batch_clause(batch_index=1) expected2 = "(id = :batch_1_id AND name = :batch_1_name)" self.assertEqual(result2, expected2) def test_security_parameter_extraction_edge_cases(self): """Test parameter extraction with edge cases and security issues.""" connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') config = TableFieldsQueryConfig( table_id="users", where_clause_template="id = :id", where_clause_fields=["id"]) handler = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=config) # Test empty template. self.assertEqual(handler._extract_parameter_names(""), []) # Test template with no parameters. self.assertEqual( handler._extract_parameter_names("SELECT * FROM users"), []) # Test template with malformed parameters (should not match). self.assertEqual(handler._extract_parameter_names("id = :"), []) # Test template with numeric parameter names. self.assertEqual( handler._extract_parameter_names("col = :param_123"), ["param_123"]) # Test template with underscore parameter names. self.assertEqual( handler._extract_parameter_names("col = :user_id"), ["user_id"]) # Test duplicate parameter names. result = handler._extract_parameter_names("id = :id OR id = :id") # Note: re.findall returns all matches, so we'd get ["id", "id"]. self.assertEqual(result, ["id", "id"]) def test_build_single_param_dict_with_generic_names(self): """Test parameter dict building with generic names.""" connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') # Test TableFunctionQueryConfig with generic parameter names. config = TableFunctionQueryConfig( table_id="users", where_clause_template="id = :param_0 AND name = :param_1", where_clause_value_fn=lambda row: [row.id, row.name]) handler = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=config) result = handler._build_single_param_dict([123, "John"]) expected = {"param_0": 123, "param_1": "John"} self.assertEqual(result, expected) def test_duplicate_binding_parameter_names_handling(self): """Test handling of duplicate binding parameter names.""" connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') # Test TableFunctionQueryConfig with duplicate parameter names. config = TableFunctionQueryConfig( table_id="users", where_clause_template="id = :param AND name = :param", where_clause_value_fn=lambda row: [row.id, row.name]) handler = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=config) # Test that duplicate parameter names are made unique. result = handler._build_single_param_dict([123, "John"]) expected = {"param_0": 123, "param_1": "John"} self.assertEqual(result, expected) # Test the helper method for template and params. template = "id = :param AND name = :param" values = [123, "John"] updated_template, param_dict = handler._get_unique_template_and_params( template, values) expected_template = "id = :param_0 AND name = :param_1" expected_params = {"param_0": 123, "param_1": "John"} self.assertEqual(updated_template, expected_template) self.assertEqual(param_dict, expected_params) def test_unsupported_query_config_type_error(self): """Test that unsupported query config types raise ValueError.""" connection_config = ExternalSQLDBConnectionConfig( db_adapter=DatabaseTypeAdapter.POSTGRESQL, host='localhost', port=5432, user='user', password='password', db_id='db') config = TableFieldsQueryConfig( table_id="users", where_clause_template="id = :id", where_clause_fields=["id"]) handler = CloudSQLEnrichmentHandler( connection_config=connection_config, query_config=config) # Temporarily replace the config with an unsupported type. handler._query_config = "unsupported_type" import apache_beam as beam request = beam.Row(id=1) with self.assertRaises(ValueError) as context: handler._extract_values_from_request(request) self.assertIn( "Unsupported query configuration type", str(context.exception)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/feast_feature_store.py0000644000175100017510000002046515206113033030502 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import tempfile from collections.abc import Callable from collections.abc import Mapping from pathlib import Path from typing import Any from typing import Optional from feast import FeatureStore import apache_beam as beam from apache_beam.io.filesystems import FileSystems from apache_beam.transforms.enrichment import EnrichmentSourceHandler from apache_beam.transforms.enrichment_handlers.utils import ExceptionLevel __all__ = [ 'FeastFeatureStoreEnrichmentHandler', ] EntityRowFn = Callable[[beam.Row], Mapping[str, Any]] _LOGGER = logging.getLogger(__name__) LOCAL_FEATURE_STORE_YAML_FILENAME = 'fs_yaml_file.yaml' def download_fs_yaml_file(gcs_fs_yaml_file: str): """Download the feature store config file for Feast.""" try: with FileSystems.open(gcs_fs_yaml_file, 'r') as gcs_file: with tempfile.NamedTemporaryFile(suffix=LOCAL_FEATURE_STORE_YAML_FILENAME, delete=False) as local_file: local_file.write(gcs_file.read()) return Path(local_file.name) except Exception: raise RuntimeError( 'error downloading the file %s locally to load the ' 'Feast feature store.' % gcs_fs_yaml_file) def _validate_feature_names(feature_names, feature_service_name): """Check if one of `feature_names` or `feature_service_name` is provided.""" if ((not feature_names and not feature_service_name) or bool(feature_names and feature_service_name)): raise ValueError( 'Please provide exactly one of a list of feature names to fetch ' 'from online store (`feature_names`) or a feature service name for ' 'the Feast online feature store (`feature_service_name`).') def _validate_feature_store_yaml_path_exists(fs_yaml_file): """Check if the feature store yaml path exists.""" if not FileSystems.exists(fs_yaml_file): raise ValueError( 'The feature store yaml path (%s) does not exist.' % fs_yaml_file) def _validate_entity_key_exists(entity_id, entity_row_fn): """Checks if the entity key or a lambda to build entity key exists.""" if ((not entity_row_fn and not entity_id) or bool(entity_row_fn and entity_id)): raise ValueError( "Please specify exactly one of a `entity_id` or a lambda " "function with `entity_row_fn` to extract the entity id " "from the input row.") class FeastFeatureStoreEnrichmentHandler(EnrichmentSourceHandler[beam.Row, beam.Row]): """Enrichment handler to interact with the Feast feature store. To specify the features to fetch from Feast online store, please specify exactly one of `feature_names` or `feature_service_name`. Use this handler with :class:`apache_beam.transforms.enrichment.Enrichment` transform. To filter the features to enrich, use the `join_fn` param in :class:`apache_beam.transforms.enrichment.Enrichment`. """ def __init__( self, feature_store_yaml_path: str, feature_names: Optional[list[str]] = None, feature_service_name: Optional[str] = "", full_feature_names: Optional[bool] = False, entity_id: str = "", *, entity_row_fn: Optional[EntityRowFn] = None, exception_level: ExceptionLevel = ExceptionLevel.WARN, ): """Initializes an instance of `FeastFeatureStoreEnrichmentHandler`. Args: feature_store_yaml_path (str): The path to a YAML configuration file for the Feast feature store. See https://docs.feast.dev/reference/feature-repository/feature-store-yaml for configuration options supported by Feast. feature_names: A list of feature names to be retrieved from the online Feast feature store. feature_service_name (str): The name of the feature service containing the features to fetch from the online Feast feature store. full_feature_names (bool): Whether to use full feature names (including namespaces, etc.). Defaults to False. entity_id (str): entity name for the entity associated with the features. The `entity_id` is used to extract the entity value from the input row. Please provide exactly one of `entity_id` or `entity_row_fn`. entity_row_fn: a lambda function that takes an input `beam.Row` and returns a dictionary with a mapping from the entity key column name to entity key value. It is used to build/extract the entity dict for feature retrieval. Please provide exactly one of `entity_id` or `entity_row_fn`. See https://docs.feast.dev/getting-started/concepts/feature-retrieval for more information. exception_level: a `enum.Enum` value from `apache_beam.transforms.enrichment_handlers.utils.ExceptionLevel` to set the level when `None` feature values are fetched from the online Feast store. Defaults to `ExceptionLevel.WARN`. """ self.entity_id = entity_id self.feature_store_yaml_path = feature_store_yaml_path self.feature_names = feature_names self.feature_service_name = feature_service_name self.full_feature_names = full_feature_names self.entity_row_fn = entity_row_fn self._exception_level = exception_level _validate_entity_key_exists(self.entity_id, self.entity_row_fn) _validate_feature_store_yaml_path_exists(self.feature_store_yaml_path) _validate_feature_names(self.feature_names, self.feature_service_name) def __enter__(self): """Connect with the Feast feature store.""" local_repo_path = download_fs_yaml_file(self.feature_store_yaml_path) try: self.store = FeatureStore(fs_yaml_file=local_repo_path) except Exception: raise RuntimeError( 'Invalid feature store yaml file provided. Make sure ' 'the %s contains the valid configuration for Feast feature store.' % self.feature_store_yaml_path) if self.feature_service_name: try: self.features = self.store.get_feature_service( self.feature_service_name) except Exception: raise RuntimeError( 'Could not find the feature service %s for the feature ' 'store configured in %s.' % (self.feature_service_name, self.feature_store_yaml_path)) else: self.features = self.feature_names def __call__(self, request: beam.Row, *args, **kwargs): """Fetches feature values for an entity-id from the Feast feature store. Args: request: the input `beam.Row` to enrich. """ if self.entity_row_fn: entity_dict = self.entity_row_fn(request) else: request_dict = request._asdict() entity_dict = {self.entity_id: request_dict[self.entity_id]} feature_values = self.store.get_online_features( features=self.features, entity_rows=[entity_dict], full_feature_names=self.full_feature_names).to_dict() # get_online_features() returns a list of feature values per entity-id. # Since we do this per entity, the list of feature values only contain # a single element at position 0. response_dict = {k: v[0] for k, v in feature_values.items()} return request, beam.Row(**response_dict) def __exit__(self, exc_type, exc_val, exc_tb): """Clean the instantiated Feast feature store client.""" self.store = None def get_cache_key(self, request: beam.Row) -> str: """Returns a string formatted with unique entity-id for the feature values. """ if self.entity_row_fn: entity_dict = self.entity_row_fn(request) entity_id = list(entity_dict.keys())[0] else: entity_id = self.entity_id return 'entity_id: %s' % request._asdict()[entity_id] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/feast_feature_store_it_test.py0000644000175100017510000001016115206113033032225 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests Feast feature store enrichment handler for enrichment transform. See https://s.apache.org/feast-enrichment-test-setup to set up test feast feature repository. """ import unittest from collections.abc import Mapping from typing import Any import pytest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline # pylint: disable=ungrouped-imports try: from apache_beam.transforms.enrichment import Enrichment from apache_beam.transforms.enrichment_handlers.feast_feature_store import FeastFeatureStoreEnrichmentHandler from apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store_it_test import \ ValidateResponse # pylint: disable=line-too-long except ImportError: raise unittest.SkipTest( 'Feast feature store test dependencies are not installed.') def _entity_row_fn(request: beam.Row) -> Mapping[str, Any]: entity_value = request.user_id # type: ignore[attr-defined] return {'user_id': entity_value} @pytest.mark.uses_feast class TestFeastEnrichmentHandler(unittest.TestCase): def setUp(self) -> None: self.feature_store_yaml_file = ( 'gs://apache-beam-testing-enrichment/' 'feast-feature-store/repos/ecommerce/' 'feature_repo/feature_store.yaml') self.feature_service_name = 'demograph_service' def test_feast_enrichment(self): requests = [ beam.Row(user_id=2, product_id=1), beam.Row(user_id=6, product_id=2), beam.Row(user_id=9, product_id=3), ] expected_fields = [ 'user_id', 'product_id', 'state', 'country', 'gender', 'age' ] handler = FeastFeatureStoreEnrichmentHandler( entity_id='user_id', feature_store_yaml_path=self.feature_store_yaml_file, feature_service_name=self.feature_service_name, ) with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | beam.Create(requests) | Enrichment(handler) | beam.ParDo(ValidateResponse(expected_fields))) def test_feast_enrichment_bad_feature_service_name(self): """Test raising an error when a bad feature service name is given.""" requests = [ beam.Row(user_id=1, product_id=1), ] handler = FeastFeatureStoreEnrichmentHandler( entity_id='user_id', feature_store_yaml_path=self.feature_store_yaml_file, feature_service_name="bad_name", ) with self.assertRaises(RuntimeError): test_pipeline = beam.Pipeline() _ = (test_pipeline | beam.Create(requests) | Enrichment(handler)) res = test_pipeline.run() res.wait_until_finish() def test_feast_enrichment_with_lambda(self): requests = [ beam.Row(user_id=2, product_id=1), beam.Row(user_id=6, product_id=2), beam.Row(user_id=9, product_id=3), ] expected_fields = [ 'user_id', 'product_id', 'state', 'country', 'gender', 'age' ] handler = FeastFeatureStoreEnrichmentHandler( feature_store_yaml_path=self.feature_store_yaml_file, feature_service_name=self.feature_service_name, entity_row_fn=_entity_row_fn, ) with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | beam.Create(requests) | Enrichment(handler) | beam.ParDo(ValidateResponse(expected_fields))) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/feast_feature_store_test.py0000644000175100017510000000503015206113033031530 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest from parameterized import parameterized try: from apache_beam.transforms.enrichment_handlers.feast_feature_store import FeastFeatureStoreEnrichmentHandler from apache_beam.transforms.enrichment_handlers.feast_feature_store_it_test import _entity_row_fn except ImportError: raise unittest.SkipTest( 'Feast feature store test dependencies are not installed.') class TestFeastFeatureStoreHandler(unittest.TestCase): def setUp(self) -> None: self.feature_store_yaml_file = ( 'gs://apache-beam-testing-enrichment/' 'feast-feature-store/repos/ecommerce/' 'feature_repo/feature_store.yaml') self.feature_service_name = 'demograph_service' def test_feature_store_yaml_path_exists(self): feature_store_yaml_path = 'gs://apache-beam-testing-enrichment/invalid.yaml' with self.assertRaises(ValueError): _ = FeastFeatureStoreEnrichmentHandler( entity_id='user_id', feature_store_yaml_path=feature_store_yaml_path, feature_service_name=self.feature_service_name, ) def test_feast_enrichment_no_feature_service(self): """Test raising an error in case of no feature service name.""" with self.assertRaises(ValueError): _ = FeastFeatureStoreEnrichmentHandler( entity_id='user_id', feature_store_yaml_path=self.feature_store_yaml_file, ) @parameterized.expand([('user_id', _entity_row_fn), ('', None)]) def test_feast_enrichment_invalid_args(self, entity_id, entity_row_fn): with self.assertRaises(ValueError): _ = FeastFeatureStoreEnrichmentHandler( feature_store_yaml_path=self.feature_store_yaml_file, entity_id=entity_id, entity_row_fn=entity_row_fn, ) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/utils.py0000644000175100017510000000240215206113033025600 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from enum import Enum __all__ = [ 'ExceptionLevel', ] class ExceptionLevel(Enum): """Options to set the severity of exceptions. You can set the exception level option to either log a warning, or raise an exception, or do nothing when an empty row is fetched from the external service. Members: - RAISE: Raise the exception. - WARN: Log a warning for exception without raising it. - QUIET: Neither log nor raise the exception. """ RAISE = 0 WARN = 1 QUIET = 2 ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store.py0000644000175100017510000003222115206113033031357 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import proto from google.api_core.exceptions import NotFound from google.cloud import aiplatform import apache_beam as beam from apache_beam.transforms.enrichment import EnrichmentSourceHandler from apache_beam.transforms.enrichment_handlers.utils import ExceptionLevel __all__ = [ 'VertexAIFeatureStoreEnrichmentHandler', 'VertexAIFeatureStoreLegacyEnrichmentHandler' ] _LOGGER = logging.getLogger(__name__) def _not_found_err_message( feature_store_name: str, feature_view_name: str, entity_id: str) -> str: """returns a string formatted with given parameters""" return ( "make sure the Feature Store: %s with Feature View " "%s has entity_id: %s" % (feature_store_name, feature_view_name, entity_id)) class VertexAIFeatureStoreEnrichmentHandler(EnrichmentSourceHandler[beam.Row, beam.Row]): """Enrichment handler to interact with Vertex AI Feature Store. Use this handler with :class:`apache_beam.transforms.enrichment.Enrichment` transform when the Vertex AI Feature Store is set up for Bigtable Online serving. With the Bigtable Online serving approach, the client fetches all the available features for an entity-id. The entity-id is extracted from the `row_key` field in the input `beam.Row` object. To filter the features to enrich, use the `join_fn` param in :class:`apache_beam.transforms.enrichment.Enrichment`. **NOTE:** The default severity to report exceptions is logging a warning. For this handler, Vertex AI client returns the same exception `Requested entity was not found` even though the feature store doesn't exist. So make sure the feature store instance exists or set `exception_level` as `ExceptionLevel.RAISE`. """ def __init__( self, project: str, location: str, api_endpoint: str, feature_store_name: str, feature_view_name: str, row_key: str, *, exception_level: ExceptionLevel = ExceptionLevel.WARN, **kwargs, ): """Initializes an instance of `VertexAIFeatureStoreEnrichmentHandler`. Args: project (str): The GCP project-id for the Vertex AI Feature Store. location (str): The region for the Vertex AI Feature Store. api_endpoint (str): The API endpoint for the Vertex AI Feature Store. feature_store_name (str): The name of the Vertex AI Feature Store. feature_view_name (str): The name of the feature view within the Feature Store. row_key (str): The row key field name containing the unique id for the feature values. exception_level: a `enum.Enum` value from `apache_beam.transforms.enrichment_handlers.utils.ExceptionLevel` to set the level when an empty row is returned from the BigTable query. Defaults to `ExceptionLevel.WARN`. kwargs: Optional keyword arguments to configure the `aiplatform.gapic.FeatureOnlineStoreServiceClient`. """ self.project = project self.location = location self.api_endpoint = api_endpoint self.feature_store_name = feature_store_name self.feature_view_name = feature_view_name self.row_key = row_key self.exception_level = exception_level self.kwargs = kwargs if kwargs else {} if 'client_options' in self.kwargs: if not self.kwargs['client_options']['api_endpoint']: self.kwargs['client_options']['api_endpoint'] = self.api_endpoint elif self.kwargs['client_options']['api_endpoint'] != self.api_endpoint: raise ValueError( 'Multiple values received for api_endpoint in ' 'api_endpoint and client_options parameters.') else: self.kwargs['client_options'] = {"api_endpoint": self.api_endpoint} # check if the feature store exists try: admin_client = aiplatform.gapic.FeatureOnlineStoreAdminServiceClient( **self.kwargs) except Exception: _LOGGER.warning( 'Due to insufficient admin permission, could not verify ' 'the existence of feature store. If the `exception_level` ' 'is set to WARN then make sure the feature store exists ' 'otherwise the data enrichment will not happen without ' 'throwing an error.') else: location_path = admin_client.common_location_path( project=self.project, location=self.location) feature_store_path = admin_client.feature_online_store_path( project=self.project, location=self.location, feature_online_store=self.feature_store_name) feature_store = admin_client.get_feature_online_store( name=feature_store_path) if not feature_store: raise NotFound( 'Vertex AI Feature Store %s does not exists in %s' % (self.feature_store_name, location_path)) def __enter__(self): """Connect with the Vertex AI Feature Store.""" self.client = aiplatform.gapic.FeatureOnlineStoreServiceClient( **self.kwargs) self.feature_view_path = self.client.feature_view_path( self.project, self.location, self.feature_store_name, self.feature_view_name) def __call__(self, request: beam.Row, *args, **kwargs): """Fetches feature value for an entity-id from Vertex AI Feature Store. Args: request: the input `beam.Row` to enrich. """ try: entity_id = request._asdict()[self.row_key] except KeyError: raise KeyError( "Enrichment requests to Vertex AI Feature Store should " "contain a field: %s in the input `beam.Row` to join " "the input with fetched response. This is used as the " "`FeatureViewDataKey` to fetch feature values " "corresponding to this key." % self.row_key) try: response = self.client.fetch_feature_values( request=aiplatform.gapic.FetchFeatureValuesRequest( data_key=aiplatform.gapic.FeatureViewDataKey(key=entity_id), feature_view=self.feature_view_path, data_format=aiplatform.gapic.FeatureViewDataFormat.PROTO_STRUCT, )) except NotFound: if self.exception_level == ExceptionLevel.WARN: _LOGGER.warning( _not_found_err_message( self.feature_store_name, self.feature_view_name, entity_id)) return request, beam.Row() elif self.exception_level == ExceptionLevel.RAISE: raise ValueError( _not_found_err_message( self.feature_store_name, self.feature_view_name, entity_id)) response_dict = dict(response.proto_struct) return request, beam.Row(**response_dict) def __exit__(self, exc_type, exc_val, exc_tb): """Clean the instantiated Vertex AI client.""" self.client = None def get_cache_key(self, request: beam.Row) -> str: """Returns a string formatted with unique entity-id for the feature values. """ return 'entity_id: %s' % request._asdict()[self.row_key] class VertexAIFeatureStoreLegacyEnrichmentHandler(EnrichmentSourceHandler): """Enrichment handler to interact with Vertex AI Feature Store (Legacy). Use this handler with :class:`apache_beam.transforms.enrichment.Enrichment` transform for the Vertex AI Feature Store (Legacy). By default, it fetches all the features values for an entity-id. The entity-id is extracted from the `row_key` field in the input `beam.Row` object.You can specify the features names using `feature_ids` to fetch specific features. """ def __init__( self, project: str, location: str, api_endpoint: str, feature_store_id: str, entity_type_id: str, feature_ids: list[str], row_key: str, *, exception_level: ExceptionLevel = ExceptionLevel.WARN, **kwargs, ): """Initializes an instance of `VertexAIFeatureStoreLegacyEnrichmentHandler`. Args: project (str): The GCP project for the Vertex AI Feature Store (Legacy). location (str): The region for the Vertex AI Feature Store (Legacy). api_endpoint (str): The API endpoint for the Vertex AI Feature Store (Legacy). feature_store_id (str): The id of the Vertex AI Feature Store (Legacy). entity_type_id (str): The entity type of the feature store. feature_ids (list[str]): A list of feature-ids to fetch from the Feature Store. row_key (str): The row key field name containing the entity id for the feature values. exception_level: a `enum.Enum` value from `apache_beam.transforms.enrichment_handlers.utils.ExceptionLevel` to set the level when an empty row is returned from the BigTable query. Defaults to `ExceptionLevel.WARN`. kwargs: Optional keyword arguments to configure the `aiplatform.gapic.FeaturestoreOnlineServingServiceClient`. """ self.project = project self.location = location self.api_endpoint = api_endpoint self.feature_store_id = feature_store_id self.entity_type_id = entity_type_id self.feature_ids = feature_ids self.row_key = row_key self.exception_level = exception_level self.kwargs = kwargs if kwargs else {} if 'client_options' in self.kwargs: if not self.kwargs['client_options']['api_endpoint']: self.kwargs['client_options']['api_endpoint'] = self.api_endpoint elif self.kwargs['client_options']['api_endpoint'] != self.api_endpoint: raise ValueError( 'Multiple values received for api_endpoint in ' 'api_endpoint and client_options parameters.') else: self.kwargs['client_options'] = {"api_endpoint": self.api_endpoint} # checks if feature store exists try: _ = aiplatform.Featurestore( featurestore_name=self.feature_store_id, project=self.project, location=self.location, credentials=self.kwargs.get('credentials'), ) except NotFound: raise NotFound( 'Vertex AI Feature Store (Legacy) %s does not exist' % self.feature_store_id) def __enter__(self): """Connect with the Vertex AI Feature Store (Legacy).""" self.client = aiplatform.gapic.FeaturestoreOnlineServingServiceClient( **self.kwargs) self.entity_type_path = self.client.entity_type_path( self.project, self.location, self.feature_store_id, self.entity_type_id) def __call__(self, request: beam.Row, *args, **kwargs): """Fetches feature value for an entity-id from Vertex AI Feature Store (Legacy). Args: request: the input `beam.Row` to enrich. """ try: entity_id = request._asdict()[self.row_key] except KeyError: raise KeyError( "Enrichment requests to Vertex AI Feature Store should " "contain a field: %s in the input `beam.Row` to join " "the input with fetched response. This is used as the " "`FeatureViewDataKey` to fetch feature values " "corresponding to this key." % self.row_key) try: selector = aiplatform.gapic.FeatureSelector( id_matcher=aiplatform.gapic.IdMatcher(ids=self.feature_ids)) response = self.client.read_feature_values( request=aiplatform.gapic.ReadFeatureValuesRequest( entity_type=self.entity_type_path, entity_id=entity_id, feature_selector=selector)) except NotFound: raise ValueError( _not_found_err_message( self.feature_store_id, self.entity_type_id, entity_id)) response_dict = {} proto_to_dict = proto.Message.to_dict(response.entity_view) for key, msg in zip(response.header.feature_descriptors, proto_to_dict['data']): if msg and 'value' in msg: response_dict[key.id] = list(msg['value'].values())[0] # skip fetching the metadata elif self.exception_level == ExceptionLevel.RAISE: raise ValueError( _not_found_err_message( self.feature_store_id, self.entity_type_id, entity_id)) elif self.exception_level == ExceptionLevel.WARN: _LOGGER.warning( _not_found_err_message( self.feature_store_id, self.entity_type_id, entity_id)) return request, beam.Row(**response_dict) def __exit__(self, exc_type, exc_val, exc_tb): """Clean the instantiated Vertex AI client.""" self.client = None def get_cache_key(self, request: beam.Row) -> str: """Returns a string formatted with unique entity-id for the feature values. """ return 'entity_id: %s' % request._asdict()[self.row_key] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_it_test.py0000644000175100017510000002576415206113033033130 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import time import unittest from unittest.mock import MagicMock import pytest import apache_beam as beam from apache_beam.coders import coders from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import BeamAssertException # pylint: disable=ungrouped-imports try: from testcontainers.redis import RedisContainer from apache_beam.transforms.enrichment import Enrichment from apache_beam.transforms.enrichment_handlers.utils import ExceptionLevel from apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store import VertexAIFeatureStoreEnrichmentHandler from apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store import \ VertexAIFeatureStoreLegacyEnrichmentHandler except ImportError: raise unittest.SkipTest( 'VertexAI Feature Store test dependencies ' 'are not installed.') _LOGGER = logging.getLogger(__name__) class ValidateResponse(beam.DoFn): """ValidateResponse validates if a PCollection of `beam.Row` has the required fields.""" def __init__(self, expected_fields): self.expected_fields = expected_fields def process(self, element: beam.Row, *args, **kwargs): element_dict = element.as_dict() if len(self.expected_fields) != len(element_dict.keys()): raise BeamAssertException( "Expected %d fields in enriched PCollection:" % len(self.expected_fields)) for field in self.expected_fields: if field not in element_dict: raise BeamAssertException( f"Expected to fetch field: {field}" f"from feature store") @pytest.mark.uses_testcontainer class TestVertexAIFeatureStoreHandler(unittest.TestCase): def setUp(self) -> None: self.project = 'apache-beam-testing' self.location = 'us-central1' self.feature_store_name = "the_look_demo_unique" self.feature_view_name = "registry_product" self.entity_type_name = "entity_id" self.api_endpoint = "us-central1-aiplatform.googleapis.com" self.feature_ids = ['title', 'genres'] self.retries = 5 self._start_container() def _start_container(self): for i in range(self.retries): try: self.container = RedisContainer(image='redis:7.2.4') self.container.start() self.host = self.container.get_container_host_ip() self.port = self.container.get_exposed_port(6379) self.client = self.container.get_client() break except Exception as e: if i == self.retries - 1: _LOGGER.error('Unable to start redis container for RRIO tests.') raise e # Add a small delay between retries to avoid rapid successive failures time.sleep(2) def tearDown(self) -> None: self.container.stop() self.client = None def test_vertex_ai_feature_store_bigtable_serving_enrichment(self): requests = [ beam.Row(entity_id="847", name='cardigan jacket'), beam.Row(entity_id="16050", name='stripe t-shirt'), ] expected_fields = [ 'entity_id', 'bad_order_count', 'good_order_count', 'feature_timestamp', 'category', 'cost', 'brand', 'retail_price', 'name' ] handler = VertexAIFeatureStoreEnrichmentHandler( project=self.project, location=self.location, api_endpoint=self.api_endpoint, feature_store_name=self.feature_store_name, feature_view_name=self.feature_view_name, row_key=self.entity_type_name, ) with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | beam.Create(requests) | Enrichment(handler) | beam.ParDo(ValidateResponse(expected_fields))) def test_vertex_ai_feature_store_wrong_name(self): requests = [ beam.Row(entity_id="847", name='cardigan jacket'), beam.Row(entity_id="16050", name='stripe t-shirt'), ] with self.assertRaisesRegex(Exception, "does not exist"): handler = VertexAIFeatureStoreEnrichmentHandler( project=self.project, location=self.location, api_endpoint=self.api_endpoint, feature_store_name="incorrect_name", feature_view_name=self.feature_view_name, row_key=self.entity_type_name, ) test_pipeline = beam.Pipeline() _ = (test_pipeline | beam.Create(requests) | Enrichment(handler)) res = test_pipeline.run() res.wait_until_finish() def test_vertex_ai_feature_store_bigtable_serving_enrichment_bad(self): requests = [ beam.Row(entity_id="ui", name="fred perry men\'s sharp stripe t-shirt") ] handler = VertexAIFeatureStoreEnrichmentHandler( project=self.project, location=self.location, api_endpoint=self.api_endpoint, feature_store_name=self.feature_store_name, feature_view_name=self.feature_view_name, row_key=self.entity_type_name, exception_level=ExceptionLevel.RAISE, ) with self.assertRaises(Exception): test_pipeline = beam.Pipeline() _ = ( test_pipeline | "Create" >> beam.Create(requests) | "Enrich w/ VertexAI" >> Enrichment(handler)) res = test_pipeline.run() res.wait_until_finish() def test_vertex_ai_legacy_feature_store_enrichment(self): requests = [ beam.Row(entity_id="movie_02", title="The Shining"), beam.Row(entity_id="movie_04", title='The Dark Knight'), ] expected_fields = ['entity_id', 'title', 'genres'] feature_store_id = "movie_prediction_unique" entity_type_id = "movies" handler = VertexAIFeatureStoreLegacyEnrichmentHandler( project=self.project, location=self.location, api_endpoint=self.api_endpoint, feature_store_id=feature_store_id, entity_type_id=entity_type_id, feature_ids=self.feature_ids, row_key=self.entity_type_name, ) with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | beam.Create(requests) | Enrichment(handler) | beam.ParDo(ValidateResponse(expected_fields))) def test_vertex_ai_legacy_feature_store_enrichment_bad(self): requests = [ beam.Row(entity_id="12345", title="The Shining"), ] feature_store_id = "movie_prediction_unique" entity_type_id = "movies" handler = VertexAIFeatureStoreLegacyEnrichmentHandler( project=self.project, location=self.location, api_endpoint=self.api_endpoint, feature_store_id=feature_store_id, entity_type_id=entity_type_id, feature_ids=self.feature_ids, row_key=self.entity_type_name, exception_level=ExceptionLevel.RAISE, ) with self.assertRaises(Exception): test_pipeline = beam.Pipeline() _ = ( test_pipeline | "Create" >> beam.Create(requests) | "Enrichment" >> Enrichment(handler)) res = test_pipeline.run() res.wait_until_finish() def test_vertex_ai_legacy_feature_store_invalid_featurestore(self): requests = [ beam.Row(entity_id="movie_02", title="The Shining"), ] feature_store_id = "invalid_name" entity_type_id = "movies" with self.assertRaisesRegex(Exception, "does not exist"): handler = VertexAIFeatureStoreLegacyEnrichmentHandler( project=self.project, location=self.location, api_endpoint=self.api_endpoint, feature_store_id=feature_store_id, entity_type_id=entity_type_id, feature_ids=self.feature_ids, row_key=self.entity_type_name, exception_level=ExceptionLevel.RAISE, ) test_pipeline = beam.Pipeline() _ = ( test_pipeline | "Create" >> beam.Create(requests) | "Enrichment" >> Enrichment(handler)) res = test_pipeline.run() res.wait_until_finish() def test_feature_store_enrichment_with_redis(self): """ In this test, we run two pipelines back to back. In the first pipeline, we run a simple feature store enrichment pipeline with zero cache records. Therefore, it makes call to the source and ultimately writes to the cache with a TTL of 300 seconds. For the second pipeline, we mock the `VertexAIFeatureStoreEnrichmentHandler`'s `__call__` method to always return a `None` response. However, this change won't impact the second pipeline because the Enrichment transform first checks the cache to fulfill requests. Since all requests are cached, it will return from there without making calls to the feature store instance. """ expected_fields = ['entity_id', 'title', 'genres'] requests = [ beam.Row(entity_id="movie_02", title="The Shining"), beam.Row(entity_id="movie_04", title="The Dark Knight"), ] feature_store_id = "movie_prediction_unique" entity_type_id = "movies" handler = VertexAIFeatureStoreLegacyEnrichmentHandler( project=self.project, location=self.location, api_endpoint=self.api_endpoint, feature_store_id=feature_store_id, entity_type_id=entity_type_id, feature_ids=self.feature_ids, row_key=self.entity_type_name, ) with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | beam.Create(requests) | Enrichment(handler).with_redis_cache(self.host, self.port) | beam.ParDo(ValidateResponse(expected_fields))) # manually check cache entry c = coders.StrUtf8Coder() for req in requests: key = handler.get_cache_key(req) response = self.client.get(c.encode(key)) if not response: raise ValueError("No cache entry found for %s" % key) actual = VertexAIFeatureStoreLegacyEnrichmentHandler.__call__ VertexAIFeatureStoreLegacyEnrichmentHandler.__call__ = MagicMock( return_value=( beam.Row(entity_id="movie_02", title="The Shining"), beam.Row())) with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | beam.Create(requests) | Enrichment(handler).with_redis_cache(self.host, self.port) | beam.ParDo(ValidateResponse(expected_fields))) VertexAIFeatureStoreLegacyEnrichmentHandler.__call__ = actual if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_test.py0000644000175100017510000000430515206113033032420 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest try: from apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store import VertexAIFeatureStoreEnrichmentHandler from apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store import \ VertexAIFeatureStoreLegacyEnrichmentHandler except ImportError: raise unittest.SkipTest( 'VertexAI Feature Store test dependencies ' 'are not installed.') class TestVertexAIFeatureStoreHandlerInit(unittest.TestCase): def test_raise_error_duplicate_api_endpoint_online_store(self): with self.assertRaises(ValueError): _ = VertexAIFeatureStoreEnrichmentHandler( project='project', location='location', api_endpoint='location@google.com', feature_store_name='feature_store', feature_view_name='feature_view', row_key='row_key', client_options={'api_endpoint': 'region@google.com'}, ) def test_raise_error_duplicate_api_endpoint_legacy_store(self): with self.assertRaises(ValueError): _ = VertexAIFeatureStoreLegacyEnrichmentHandler( project='project', location='location', api_endpoint='location@google.com', feature_store_id='feature_store', entity_type_id='entity_id', feature_ids=['feature1', 'feature2'], row_key='row_key', client_options={'api_endpoint': 'region@google.com'}, ) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_it_test.py0000644000175100017510000001345315206113033024323 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import time import unittest from typing import NamedTuple from typing import Union import pytest import urllib3 import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import BeamAssertException # pylint: disable=ungrouped-imports try: from apache_beam.io.requestresponse import UserCodeExecutionException from apache_beam.io.requestresponse import UserCodeQuotaException from apache_beam.io.requestresponse_it_test import _PAYLOAD from apache_beam.io.requestresponse_it_test import EchoITOptions from apache_beam.transforms.enrichment import Enrichment from apache_beam.transforms.enrichment import EnrichmentSourceHandler except ImportError: raise unittest.SkipTest('RequestResponseIO dependencies are not installed.') class Request(NamedTuple): id: str payload: bytes def _custom_join(left, right): """custom_join returns the id and resp_payload along with a timestamp""" right['timestamp'] = time.time() return beam.Row(**right) class SampleHTTPEnrichment(EnrichmentSourceHandler[Request, beam.Row]): """Implements ``EnrichmentSourceHandler`` to call the ``EchoServiceGrpc``'s HTTP handler. """ def __init__(self, url: str): self.url = url + '/v1/echo' # append path to the mock API. def __call__(self, request: Request, *args, **kwargs): """Overrides ``Caller``'s call method invoking the ``EchoServiceGrpc``'s HTTP handler with an `dict`, returning either a successful ``tuple[dict,dict]`` or throwing either a ``UserCodeExecutionException``, ``UserCodeTimeoutException``, or a ``UserCodeQuotaException``. """ try: resp = urllib3.request( "POST", self.url, json={ "id": request.id, "payload": str(request.payload, 'utf-8') }, retries=False) if resp.status < 300: resp_body = resp.json() resp_id = resp_body['id'] payload = resp_body['payload'] return ( request, beam.Row(id=resp_id, resp_payload=bytes(payload, 'utf-8'))) if resp.status == 429: # Too Many Requests raise UserCodeQuotaException(resp.reason) elif resp.status != 200: raise UserCodeExecutionException(resp.status, resp.reason, request) except urllib3.exceptions.HTTPError as e: raise UserCodeExecutionException(e) class ValidateFields(beam.DoFn): """ValidateFields validates if a PCollection of `beam.Row` has certain fields.""" def __init__(self, n_fields: int, fields: list[str]): self.n_fields = n_fields self._fields = fields def process(self, element: beam.Row, *args, **kwargs): element_dict = element.as_dict() if len(element_dict.keys()) != self.n_fields: raise BeamAssertException( "Expected %d fields in enriched PCollection:" " id, payload and resp_payload" % self.n_fields) for field in self._fields: if field not in element_dict or element_dict[field] is None: raise BeamAssertException(f"Expected a not None field: {field}") @pytest.mark.uses_mock_api class TestEnrichment(unittest.TestCase): options: Union[EchoITOptions, None] = None client: Union[SampleHTTPEnrichment, None] = None @classmethod def setUpClass(cls) -> None: cls.options = EchoITOptions() http_endpoint_address = cls.options.http_endpoint_address if not http_endpoint_address or http_endpoint_address == '': raise unittest.SkipTest('HTTP_ENDPOINT_ADDRESS is required.') cls.client = SampleHTTPEnrichment(http_endpoint_address) @classmethod def _get_client_and_options( cls) -> tuple[SampleHTTPEnrichment, EchoITOptions]: assert cls.options is not None assert cls.client is not None return cls.client, cls.options def test_http_enrichment(self): """Tests Enrichment Transform against the Mock-API HTTP endpoint with the default cross join.""" client, options = TestEnrichment._get_client_and_options() req = Request(id=options.never_exceed_quota_id, payload=_PAYLOAD) fields = ['id', 'payload', 'resp_payload'] with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | 'Create PCollection' >> beam.Create([req]) | 'Enrichment Transform' >> Enrichment(client) | 'Assert Fields' >> beam.ParDo( ValidateFields(len(fields), fields=fields))) def test_http_enrichment_custom_join(self): """Tests Enrichment Transform against the Mock-API HTTP endpoint with a custom join function.""" client, options = TestEnrichment._get_client_and_options() req = Request(id=options.never_exceed_quota_id, payload=_PAYLOAD) fields = ['id', 'resp_payload', 'timestamp'] with TestPipeline(is_integration_test=True) as test_pipeline: _ = ( test_pipeline | 'Create PCollection' >> beam.Create([req]) | 'Enrichment Transform' >> Enrichment(client, join_fn=_custom_join) | 'Assert Fields' >> beam.ParDo( ValidateFields(len(fields), fields=fields))) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/enrichment_test.py0000644000175100017510000000262415206113033023625 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest import apache_beam as beam # pylint: disable=ungrouped-imports try: from apache_beam.transforms.enrichment import cross_join except ImportError: raise unittest.SkipTest('RequestResponseIO dependencies are not installed.') class TestEnrichmentTransform(unittest.TestCase): def test_cross_join(self): left = {'id': 1, 'key': 'city'} right = {'id': 1, 'value': 'durham'} expected = beam.Row(id=1, key='city', value='durham') output = cross_join(left, right) self.assertEqual(expected, output) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/environments.py0000644000175100017510000007702515206113033023170 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Environments concepts. For internal use only. No backwards compatibility guarantees.""" # pytype: skip-file import json import logging import sys import tempfile from collections.abc import Callable from collections.abc import Iterable from collections.abc import Iterator from collections.abc import Mapping from types import MappingProxyType from typing import TYPE_CHECKING from typing import Any from typing import Optional from typing import TypeVar from typing import Union from typing import overload from google.protobuf import message from apache_beam import coders from apache_beam.options.pipeline_options import PortableOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.portability.api import endpoints_pb2 from apache_beam.runners.portability import stager from apache_beam.transforms.resources import resource_hints_from_options from apache_beam.utils import proto_utils if TYPE_CHECKING: from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.runners.pipeline_context import PipelineContext __all__ = [ 'Environment', 'AnyOfEnvironment', 'DefaultEnvironment', 'DockerEnvironment', 'ProcessEnvironment', 'ExternalEnvironment', 'EmbeddedPythonEnvironment', 'EmbeddedPythonGrpcEnvironment', 'SubprocessSDKEnvironment', 'PyPIArtifactRegistry' ] T = TypeVar('T') EnvironmentT = TypeVar('EnvironmentT', bound='Environment') ConstructorFn = Callable[[ Optional[Any], Iterable[str], Iterable[beam_runner_api_pb2.ArtifactInformation], Mapping[str, bytes], 'PipelineContext' ], Any] def looks_like_json(s): import re return re.match(r'\s*\{.*\}\s*$', s) APACHE_BEAM_DOCKER_IMAGE_PREFIX = 'apache/beam' APACHE_BEAM_JAVA_CONTAINER_NAME_PREFIX = 'beam_java' SDK_VERSION_CAPABILITY_PREFIX = 'beam:version:sdk_base:' def is_apache_beam_container(container_image): return container_image and container_image.startswith( APACHE_BEAM_DOCKER_IMAGE_PREFIX) class Environment(object): """Abstract base class for environments. Represents a type and configuration of environment. Each type of Environment should have a unique urn. For internal use only. No backwards compatibility guarantees. """ _known_urns = {} # type: dict[str, tuple[Optional[type], ConstructorFn]] _urn_to_env_cls = {} # type: dict[str, type] def __init__( self, capabilities=(), # type: Iterable[str] artifacts=(), # type: Iterable[beam_runner_api_pb2.ArtifactInformation] resource_hints=None, # type: Optional[Mapping[str, bytes]] ): # type: (...) -> None self._capabilities = capabilities self._artifacts = sorted(artifacts, key=lambda x: x.SerializeToString()) # Hints on created environments should be immutable since pipeline context # stores environments in hash maps and we use hints to compute the hash. self._resource_hints = MappingProxyType( dict(resource_hints) if resource_hints else {}) def __eq__(self, other): return ( self.__class__ == other.__class__ and self._artifacts == other._artifacts # Assuming that we don't have instances of the same Environment subclass # with different set of capabilities. and self._resource_hints == other._resource_hints) def __hash__(self): # type: () -> int return hash((self.__class__, frozenset(self._resource_hints.items()))) def artifacts(self): # type: () -> Iterable[beam_runner_api_pb2.ArtifactInformation] return self._artifacts def to_runner_api_parameter(self, context): # type: (PipelineContext) -> tuple[str, Optional[Union[message.Message, bytes, str]]] raise NotImplementedError def capabilities(self): # type: () -> Iterable[str] return self._capabilities def resource_hints(self): # type: () -> Mapping[str, bytes] return self._resource_hints @classmethod @overload def register_urn( cls, urn, # type: str parameter_type, # type: type[T] ): # type: (...) -> Callable[[Union[type, Callable[[T, Iterable[str], PipelineContext], Any]]], Callable[[T, Iterable[str], PipelineContext], Any]] pass @classmethod @overload def register_urn( cls, urn, # type: str parameter_type, # type: None ): # type: (...) -> Callable[[Union[type, Callable[[bytes, Iterable[str], Iterable[beam_runner_api_pb2.ArtifactInformation], PipelineContext], Any]]], Callable[[bytes, Iterable[str], PipelineContext], Any]] pass @classmethod @overload def register_urn( cls, urn, # type: str parameter_type, # type: type[T] constructor # type: Callable[[T, Iterable[str], Iterable[beam_runner_api_pb2.ArtifactInformation], PipelineContext], Any] ): # type: (...) -> None pass @classmethod @overload def register_urn( cls, urn, # type: str parameter_type, # type: None constructor # type: Callable[[bytes, Iterable[str], Iterable[beam_runner_api_pb2.ArtifactInformation], PipelineContext], Any] ): # type: (...) -> None pass @classmethod def register_urn(cls, urn, parameter_type, constructor=None): def register(constructor): if isinstance(constructor, type): constructor.from_runner_api_parameter = register( constructor.from_runner_api_parameter) # register environment urn to environment class cls._urn_to_env_cls[urn] = constructor return constructor else: cls._known_urns[urn] = parameter_type, constructor return staticmethod(constructor) if constructor: # Used as a statement. register(constructor) else: # Used as a decorator. return register @classmethod def get_env_cls_from_urn(cls, urn): # type: (str) -> type[Environment] return cls._urn_to_env_cls[urn] def to_runner_api(self, context): # type: (PipelineContext) -> beam_runner_api_pb2.Environment urn, typed_param = self.to_runner_api_parameter(context) return beam_runner_api_pb2.Environment( urn=urn, payload=typed_param.SerializeToString() if isinstance( typed_param, message.Message) else typed_param if (isinstance(typed_param, bytes) or typed_param is None) else typed_param.encode('utf-8'), capabilities=self.capabilities(), dependencies=self.artifacts(), resource_hints=self.resource_hints()) @classmethod def from_runner_api( cls, proto, # type: Optional[beam_runner_api_pb2.Environment] context # type: PipelineContext ): # type: (...) -> Optional[Environment] if proto is None or not proto.urn: return None parameter_type, constructor = cls._known_urns[proto.urn] return constructor( proto_utils.parse_Bytes(proto.payload, parameter_type), proto.capabilities, proto.dependencies, proto.resource_hints, context) @classmethod def from_options(cls, options): # type: (type[EnvironmentT], PortableOptions) -> EnvironmentT """Creates an Environment object from PortableOptions. Args: options: The PortableOptions object. """ if cls != Environment: raise NotImplementedError portable_options = options.view_as(PortableOptions) environment_type = portable_options.environment_type if not environment_type: environment_urn = common_urns.environments.DOCKER.urn elif environment_type.startswith('beam:env:'): environment_urn = environment_type elif environment_type == 'LOOPBACK': environment_urn = python_urns.EMBEDDED_PYTHON_LOOPBACK else: try: environment_urn = getattr( common_urns.environments, environment_type).urn except AttributeError: raise ValueError('Unknown environment type: %s' % environment_type) env_class = Environment.get_env_cls_from_urn(environment_urn) return env_class.from_options(portable_options) # type: ignore @Environment.register_urn(common_urns.environments.DEFAULT.urn, None) class DefaultEnvironment(Environment): """Used as a stub when context is missing a default environment.""" def to_runner_api_parameter(self, context): return common_urns.environments.DEFAULT.urn, None @staticmethod def from_runner_api_parameter( payload, # type: beam_runner_api_pb2.DockerPayload capabilities, # type: Iterable[str] artifacts, # type: Iterable[beam_runner_api_pb2.ArtifactInformation] resource_hints, # type: Mapping[str, bytes] context # type: PipelineContext ): # type: (...) -> DefaultEnvironment return DefaultEnvironment( capabilities=capabilities, artifacts=artifacts, resource_hints=resource_hints) @Environment.register_urn( common_urns.environments.DOCKER.urn, beam_runner_api_pb2.DockerPayload) class DockerEnvironment(Environment): def __init__( self, container_image=None, # type: Optional[str] capabilities=(), # type: Iterable[str] artifacts=(), # type: Iterable[beam_runner_api_pb2.ArtifactInformation] resource_hints=None, # type: Optional[Mapping[str, bytes]] ): super().__init__(capabilities, artifacts, resource_hints) if container_image: logging.debug( 'Using provided Python SDK container image: %s' % (container_image)) self.container_image = container_image else: logging.debug('No image given, using default Python SDK image') self.container_image = self.default_docker_image() logging.debug( 'Python SDK container image set to "%s" for Docker environment' % (self.container_image)) def __eq__(self, other): return ( super().__eq__(other) and self.container_image == other.container_image) def __hash__(self): return hash((super().__hash__(), self.container_image)) def __repr__(self): return 'DockerEnvironment(container_image=%s)' % self.container_image def to_runner_api_parameter(self, context): # type: (PipelineContext) -> tuple[str, beam_runner_api_pb2.DockerPayload] return ( common_urns.environments.DOCKER.urn, beam_runner_api_pb2.DockerPayload(container_image=self.container_image)) @staticmethod def from_runner_api_parameter( payload, # type: beam_runner_api_pb2.DockerPayload capabilities, # type: Iterable[str] artifacts, # type: Iterable[beam_runner_api_pb2.ArtifactInformation] resource_hints, # type: Mapping[str, bytes] context # type: PipelineContext ): # type: (...) -> DockerEnvironment return DockerEnvironment( container_image=payload.container_image, capabilities=capabilities, artifacts=artifacts, resource_hints=resource_hints) @classmethod def from_options(cls, options): # type: (PortableOptions) -> DockerEnvironment if options.view_as(SetupOptions).prebuild_sdk_container_engine: # Imported here to avoid circular dependencies. # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.runners.portability.sdk_container_builder import SdkContainerImageBuilder prebuilt_container_image = SdkContainerImageBuilder.build_container_image( options) return cls.from_container_image( container_image=prebuilt_container_image, artifacts=python_sdk_dependencies(options), resource_hints=resource_hints_from_options(options), ) return cls.from_container_image( container_image=options.lookup_environment_option( 'docker_container_image') or options.environment_config, artifacts=python_sdk_dependencies(options), resource_hints=resource_hints_from_options(options), ) @classmethod def from_container_image( cls, container_image, artifacts=(), resource_hints=None): # type: (str, Iterable[beam_runner_api_pb2.ArtifactInformation], Optional[Mapping[str, bytes]]) -> DockerEnvironment return cls( container_image=container_image, capabilities=python_sdk_docker_capabilities(), artifacts=artifacts, resource_hints=resource_hints) @staticmethod def default_docker_image(): # type: () -> str from apache_beam import version as beam_version sdk_version = beam_version.__version__ version_suffix = '.'.join([str(i) for i in sys.version_info[0:2]]) image = ( APACHE_BEAM_DOCKER_IMAGE_PREFIX + '_python{version_suffix}_sdk:{tag}'.format( version_suffix=version_suffix, tag=sdk_version)) return image @Environment.register_urn( common_urns.environments.PROCESS.urn, beam_runner_api_pb2.ProcessPayload) class ProcessEnvironment(Environment): def __init__( self, command, # type: str os='', # type: str arch='', # type: str env=None, # type: Optional[Mapping[str, str]] capabilities=(), # type: Iterable[str] artifacts=(), # type: Iterable[beam_runner_api_pb2.ArtifactInformation] resource_hints=None, # type: Optional[Mapping[str, bytes]] ): # type: (...) -> None super().__init__(capabilities, artifacts, resource_hints) self.command = command self.os = os self.arch = arch self.env = env or {} def __eq__(self, other): return ( super().__eq__(other) and self.command == other.command and self.os == other.os and self.arch == other.arch and self.env == other.env) def __hash__(self): # type: () -> int return hash(( super().__hash__(), self.command, self.os, self.arch, frozenset(self.env.items()))) def __repr__(self): # type: () -> str repr_parts = ['command=%s' % self.command] if self.os: repr_parts.append('os=%s' % self.os) if self.arch: repr_parts.append('arch=%s' % self.arch) repr_parts.append('env=%s' % self.env) return 'ProcessEnvironment(%s)' % ','.join(repr_parts) def to_runner_api_parameter(self, context): # type: (PipelineContext) -> tuple[str, beam_runner_api_pb2.ProcessPayload] return ( common_urns.environments.PROCESS.urn, beam_runner_api_pb2.ProcessPayload( os=self.os, arch=self.arch, command=self.command, env=self.env)) @staticmethod def from_runner_api_parameter( payload, capabilities, # type: Iterable[str] artifacts, # type: Iterable[beam_runner_api_pb2.ArtifactInformation] resource_hints, # type: Mapping[str, bytes] context # type: PipelineContext ): # type: (...) -> ProcessEnvironment return ProcessEnvironment( command=payload.command, os=payload.os, arch=payload.arch, env=payload.env, capabilities=capabilities, artifacts=artifacts, resource_hints=resource_hints, ) @staticmethod def parse_environment_variables(variables): env = {} for var in variables: try: name, value = var.split('=', 1) env[name] = value except ValueError: raise ValueError( 'Invalid process_variables "%s" (expected assignment in the ' 'form "FOO=bar").' % var) return env @classmethod def from_options(cls, options): # type: (PortableOptions) -> ProcessEnvironment if options.environment_config: config = json.loads(options.environment_config) return cls( config.get('command'), os=config.get('os', ''), arch=config.get('arch', ''), env=config.get('env', ''), capabilities=python_sdk_capabilities(), artifacts=python_sdk_dependencies(options), resource_hints=resource_hints_from_options(options), ) env = cls.parse_environment_variables( options.lookup_environment_option('process_variables').split(',') if options.lookup_environment_option('process_variables') else []) return cls( options.lookup_environment_option('process_command'), env=env, capabilities=python_sdk_capabilities(), artifacts=python_sdk_dependencies(options), resource_hints=resource_hints_from_options(options), ) @Environment.register_urn( common_urns.environments.EXTERNAL.urn, beam_runner_api_pb2.ExternalPayload) class ExternalEnvironment(Environment): def __init__( self, url, # type: str params=None, # type: Optional[Mapping[str, str]] capabilities=(), # type: Iterable[str] artifacts=(), # type: Iterable[beam_runner_api_pb2.ArtifactInformation] resource_hints=None, # type: Optional[Mapping[str, bytes]] ): super().__init__(capabilities, artifacts, resource_hints) self.url = url self.params = params def __eq__(self, other): return ( super().__eq__(other) and self.url == other.url and self.params == other.params) def __hash__(self): # type: () -> int return hash(( super().__hash__(), self.url, frozenset(self.params.items()) if self.params is not None else None)) def __repr__(self): # type: () -> str return 'ExternalEnvironment(url=%s,params=%s)' % (self.url, self.params) def to_runner_api_parameter(self, context): # type: (PipelineContext) -> tuple[str, beam_runner_api_pb2.ExternalPayload] return ( common_urns.environments.EXTERNAL.urn, beam_runner_api_pb2.ExternalPayload( endpoint=endpoints_pb2.ApiServiceDescriptor(url=self.url), params=self.params)) @staticmethod def from_runner_api_parameter( payload, # type: beam_runner_api_pb2.ExternalPayload capabilities, # type: Iterable[str] artifacts, # type: Iterable[beam_runner_api_pb2.ArtifactInformation] resource_hints, # type: Mapping[str, bytes] context # type: PipelineContext ): # type: (...) -> ExternalEnvironment return ExternalEnvironment( payload.endpoint.url, params=payload.params or None, capabilities=capabilities, artifacts=artifacts, resource_hints=resource_hints) @classmethod def from_options(cls, options): # type: (PortableOptions) -> ExternalEnvironment if looks_like_json(options.environment_config): config = json.loads(options.environment_config) url = config.get('url') if not url: raise ValueError('External environment endpoint must be set.') params = config.get('params') elif options.environment_config: url = options.environment_config params = None else: url = options.lookup_environment_option('external_service_address') params = None return cls( url, params=params, capabilities=python_sdk_capabilities(), artifacts=python_sdk_dependencies(options), resource_hints=resource_hints_from_options(options)) def expand_anyof_environments(env_proto): if env_proto.urn == common_urns.environments.ANYOF.urn: for alt in beam_runner_api_pb2.AnyOfEnvironmentPayload.FromString( env_proto.payload).environments: yield from expand_anyof_environments(alt) else: yield env_proto def resolve_anyof_environment(env_proto, *preferred_types): all_environments = list(expand_anyof_environments(env_proto)) for preferred_type in preferred_types: for env in all_environments: if env.urn == preferred_type: return env return all_environments[0] @Environment.register_urn(python_urns.EMBEDDED_PYTHON, None) class EmbeddedPythonEnvironment(Environment): def to_runner_api_parameter(self, context): # type: (PipelineContext) -> tuple[str, None] return python_urns.EMBEDDED_PYTHON, None @staticmethod def from_runner_api_parameter( unused_payload, # type: None capabilities, # type: Iterable[str] artifacts, # type: Iterable[beam_runner_api_pb2.ArtifactInformation] resource_hints, # type: Mapping[str, bytes] context # type: PipelineContext ): # type: (...) -> EmbeddedPythonEnvironment return EmbeddedPythonEnvironment(capabilities, artifacts, resource_hints) @classmethod def from_options(cls, options): # type: (PortableOptions) -> EmbeddedPythonEnvironment return cls( capabilities=python_sdk_capabilities(), artifacts=python_sdk_dependencies(options), resource_hints=resource_hints_from_options(options), ) @classmethod def default(cls): # type: () -> EmbeddedPythonEnvironment return cls(capabilities=python_sdk_capabilities(), artifacts=()) @Environment.register_urn(python_urns.EMBEDDED_PYTHON_GRPC, bytes) class EmbeddedPythonGrpcEnvironment(Environment): def __init__( self, state_cache_size=None, data_buffer_time_limit_ms=None, capabilities=(), artifacts=(), resource_hints=None, ): super().__init__(capabilities, artifacts, resource_hints) self.state_cache_size = state_cache_size self.data_buffer_time_limit_ms = data_buffer_time_limit_ms def __eq__(self, other): return ( super().__eq__(other) and self.state_cache_size == other.state_cache_size and self.data_buffer_time_limit_ms == other.data_buffer_time_limit_ms) def __hash__(self): # type: () -> int return hash(( super().__hash__(), self.state_cache_size, self.data_buffer_time_limit_ms)) def __repr__(self): # type: () -> str repr_parts = [] if not self.state_cache_size is None: repr_parts.append('state_cache_size=%d' % self.state_cache_size) if not self.data_buffer_time_limit_ms is None: repr_parts.append( 'data_buffer_time_limit_ms=%d' % self.data_buffer_time_limit_ms) return 'EmbeddedPythonGrpcEnvironment(%s)' % ','.join(repr_parts) def to_runner_api_parameter(self, context): # type: (PipelineContext) -> tuple[str, bytes] params = {} if self.state_cache_size is not None: params['state_cache_size'] = self.state_cache_size if self.data_buffer_time_limit_ms is not None: params['data_buffer_time_limit_ms'] = self.data_buffer_time_limit_ms payload = json.dumps(params).encode('utf-8') return python_urns.EMBEDDED_PYTHON_GRPC, payload @staticmethod def from_runner_api_parameter( payload, # type: bytes capabilities, # type: Iterable[str] artifacts, # type: Iterable[beam_runner_api_pb2.ArtifactInformation] resource_hints, # type: Mapping[str, bytes] context # type: PipelineContext ): # type: (...) -> EmbeddedPythonGrpcEnvironment if payload: config = EmbeddedPythonGrpcEnvironment.parse_config( payload.decode('utf-8')) return EmbeddedPythonGrpcEnvironment( state_cache_size=config.get('state_cache_size'), data_buffer_time_limit_ms=config.get('data_buffer_time_limit_ms'), capabilities=capabilities, artifacts=artifacts, resource_hints=resource_hints) else: return EmbeddedPythonGrpcEnvironment() @classmethod def from_options(cls, options): # type: (PortableOptions) -> EmbeddedPythonGrpcEnvironment if options.environment_config: config = EmbeddedPythonGrpcEnvironment.parse_config( options.environment_config) return cls( state_cache_size=config.get('state_cache_size'), data_buffer_time_limit_ms=config.get('data_buffer_time_limit_ms'), capabilities=python_sdk_capabilities(), artifacts=python_sdk_dependencies(options)) else: return cls( capabilities=python_sdk_capabilities(), artifacts=python_sdk_dependencies(options), resource_hints=resource_hints_from_options(options)) @staticmethod def parse_config(s): # type: (str) -> dict[str, Any] if looks_like_json(s): config_dict = json.loads(s) if 'state_cache_size' in config_dict: config_dict['state_cache_size'] = int(config_dict['state_cache_size']) if 'data_buffer_time_limit_ms' in config_dict: config_dict['data_buffer_time_limit_ms'] = \ int(config_dict['data_buffer_time_limit_ms']) return config_dict else: return {'state_cache_size': int(s)} @classmethod def default(cls): # type: () -> EmbeddedPythonGrpcEnvironment return cls(capabilities=python_sdk_capabilities(), artifacts=()) @Environment.register_urn(python_urns.EMBEDDED_PYTHON_LOOPBACK, None) class PythonLoopbackEnvironment(EmbeddedPythonEnvironment): """Used as a stub when the loopback worker has not yet been started.""" def to_runner_api_parameter(self, context): # type: (PipelineContext) -> tuple[str, None] return python_urns.EMBEDDED_PYTHON_LOOPBACK, None @staticmethod def from_runner_api_parameter( unused_payload, # type: None capabilities, # type: Iterable[str] artifacts, # type: Iterable[beam_runner_api_pb2.ArtifactInformation] resource_hints, # type: Mapping[str, bytes] context # type: PipelineContext ): # type: (...) -> PythonLoopbackEnvironment return PythonLoopbackEnvironment( capabilities=capabilities, artifacts=artifacts, resource_hints=resource_hints) @Environment.register_urn(python_urns.SUBPROCESS_SDK, bytes) class SubprocessSDKEnvironment(Environment): def __init__( self, command_string, # type: str capabilities=(), # type: Iterable[str] artifacts=(), # type: Iterable[beam_runner_api_pb2.ArtifactInformation] resource_hints=None, # type: Optional[Mapping[str, bytes]] ): super().__init__(capabilities, artifacts, resource_hints) self.command_string = command_string def __eq__(self, other): return ( super().__eq__(other) and self.command_string == other.command_string) def __hash__(self): # type: () -> int return hash((super().__hash__(), self.command_string)) def __repr__(self): # type: () -> str return 'SubprocessSDKEnvironment(command_string=%s)' % self.command_string def to_runner_api_parameter(self, context): # type: (PipelineContext) -> tuple[str, bytes] return python_urns.SUBPROCESS_SDK, self.command_string.encode('utf-8') @staticmethod def from_runner_api_parameter( payload, # type: bytes capabilities, # type: Iterable[str] artifacts, # type: Iterable[beam_runner_api_pb2.ArtifactInformation] resource_hints, # type: Mapping[str, bytes] context # type: PipelineContext ): # type: (...) -> SubprocessSDKEnvironment return SubprocessSDKEnvironment( payload.decode('utf-8'), capabilities, artifacts, resource_hints) @classmethod def from_options(cls, options): # type: (PortableOptions) -> SubprocessSDKEnvironment return cls( options.environment_config, capabilities=python_sdk_capabilities(), artifacts=python_sdk_dependencies(options), resource_hints=resource_hints_from_options(options)) @classmethod def from_command_string(cls, command_string): # type: (str) -> SubprocessSDKEnvironment return cls( command_string, capabilities=python_sdk_capabilities(), artifacts=()) @Environment.register_urn( common_urns.environments.ANYOF.urn, beam_runner_api_pb2.AnyOfEnvironmentPayload) class AnyOfEnvironment(Environment): def __init__(self, environments): self._environments = environments def to_runner_api_parameter(self, context): # type: (PipelineContext) -> tuple[str, beam_runner_api_pb2.AnyOfEnvironmentPayload] return ( common_urns.environments.ANYOF.urn, beam_runner_api_pb2.AnyOfEnvironmentPayload( environments=[ env.to_runner_api(context) for env in self._environments ])) @staticmethod def from_runner_api_parameter( payload, # type: beam_runner_api_pb2.AnyOfEnvironmentPayload capabilities, # type: Iterable[str] artifacts, # type: Iterable[beam_runner_api_pb2.ArtifactInformation] resource_hints, # type: Mapping[str, bytes] context # type: PipelineContext ): # type: (...) -> AnyOfEnvironment return AnyOfEnvironment([ Environment.from_runner_api(env, context) for env in payload.environments ]) @staticmethod def create_proto( environments: Iterable[beam_runner_api_pb2.Environment] ) -> beam_runner_api_pb2.Environment: return beam_runner_api_pb2.Environment( urn=common_urns.environments.ANYOF.urn, payload=beam_runner_api_pb2.AnyOfEnvironmentPayload( environments=environments).SerializeToString()) def __hash__(self): return sum(hash(env) for env in self._environments) def __eq__(self, other): return ( isinstance(other, AnyOfEnvironment) and set(self._environments) == set(other._environments)) def to_runner_api(self, context): return self.create_proto( [env.to_runner_api(context) for env in self._environments]) class PyPIArtifactRegistry(object): _registered_artifacts = set() # type: set[tuple[str, str]] @classmethod def register_artifact(cls, name, version): cls._registered_artifacts.add((name, version)) @classmethod def get_artifacts(cls): for artifact in cls._registered_artifacts: yield artifact def python_sdk_capabilities(): # type: () -> list[str] return list(_python_sdk_capabilities_iter()) def python_sdk_docker_capabilities(): return python_sdk_capabilities() + [common_urns.protocols.SIBLING_WORKERS.urn] def sdk_base_version_capability(): return ( SDK_VERSION_CAPABILITY_PREFIX + DockerEnvironment.default_docker_image()) def _python_sdk_capabilities_iter(): # type: () -> Iterator[str] for urn_spec in common_urns.coders.__dict__.values(): if getattr(urn_spec, 'urn', None) in coders.Coder._known_urns: yield urn_spec.urn yield common_urns.protocols.LEGACY_PROGRESS_REPORTING.urn yield common_urns.protocols.HARNESS_MONITORING_INFOS.urn yield common_urns.protocols.WORKER_STATUS.urn yield python_urns.PACKED_COMBINE_FN yield sdk_base_version_capability() yield common_urns.sdf_components.TRUNCATE_SIZED_RESTRICTION.urn yield common_urns.primitives.TO_STRING.urn yield common_urns.protocols.DATA_SAMPLING.urn yield common_urns.protocols.SDK_CONSUMING_RECEIVED_DATA.urn yield common_urns.protocols.ORDERED_LIST_STATE.urn def python_sdk_dependencies(options, tmp_dir=None): if tmp_dir is None: tmp_dir = tempfile.mkdtemp() skip_prestaged_dependencies = options.view_as( SetupOptions).prebuild_sdk_container_engine is not None runner = options.view_as( StandardOptions).runner or StandardOptions.DEFAULT_RUNNER log_submission_env_dependencies = runner.split( '.')[-1] not in StandardOptions.LOCAL_RUNNERS return stager.Stager.create_job_resources( options, tmp_dir, pypi_requirements=[ artifact[0] + artifact[1] for artifact in PyPIArtifactRegistry.get_artifacts() ], skip_prestaged_dependencies=skip_prestaged_dependencies, log_submission_env_dependencies=log_submission_env_dependencies) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/environments_test.py0000644000175100017510000001366415206113033024226 0ustar00runnerrunner# -- coding: utf-8 -- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the transform.environments classes.""" # pytype: skip-file import logging import tempfile import unittest from unittest.mock import MagicMock from apache_beam.options.pipeline_options import PortableOptions from apache_beam.portability import common_urns from apache_beam.runners import pipeline_context from apache_beam.transforms import environments from apache_beam.transforms.environments import DockerEnvironment from apache_beam.transforms.environments import EmbeddedPythonEnvironment from apache_beam.transforms.environments import EmbeddedPythonGrpcEnvironment from apache_beam.transforms.environments import Environment from apache_beam.transforms.environments import ExternalEnvironment from apache_beam.transforms.environments import ProcessEnvironment from apache_beam.transforms.environments import SubprocessSDKEnvironment class RunnerApiTest(unittest.TestCase): def test_environment_encoding(self): for environment in (DockerEnvironment(), DockerEnvironment(container_image='img'), DockerEnvironment(capabilities=['x, y, z']), ProcessEnvironment('run.sh'), ProcessEnvironment('run.sh', os='linux', arch='amd64', env={'k1': 'v1'}), ExternalEnvironment('localhost:8080'), ExternalEnvironment('localhost:8080', params={'k1': 'v1'}), EmbeddedPythonEnvironment(), EmbeddedPythonGrpcEnvironment(), EmbeddedPythonGrpcEnvironment( state_cache_size=0, data_buffer_time_limit_ms=0), SubprocessSDKEnvironment(command_string='foö')): context = pipeline_context.PipelineContext() proto = environment.to_runner_api(context) reconstructed = Environment.from_runner_api(proto, context) self.assertEqual(environment, reconstructed) self.assertEqual(proto, reconstructed.to_runner_api(context)) def test_sdk_capabilities(self): sdk_capabilities = environments.python_sdk_capabilities() self.assertIn(common_urns.coders.LENGTH_PREFIX.urn, sdk_capabilities) self.assertIn( common_urns.protocols.HARNESS_MONITORING_INFOS.urn, sdk_capabilities) self.assertIn(common_urns.protocols.WORKER_STATUS.urn, sdk_capabilities) self.assertIn( common_urns.sdf_components.TRUNCATE_SIZED_RESTRICTION.urn, sdk_capabilities) self.assertIn(common_urns.primitives.TO_STRING.urn, sdk_capabilities) def test_default_capabilities(self): environment = DockerEnvironment.from_options( PortableOptions(sdk_location='container')) context = pipeline_context.PipelineContext() proto = environment.to_runner_api(context) self.assertEqual( set(proto.capabilities), set(environments.python_sdk_docker_capabilities())) class EnvironmentOptionsTest(unittest.TestCase): def setUp(self) -> None: self.tmp_dir = tempfile.TemporaryDirectory() self.actual_mkdtemp = tempfile.mkdtemp tempfile.mkdtemp = MagicMock(return_value=self.tmp_dir.name) def tearDown(self) -> None: tempfile.mkdtemp = self.actual_mkdtemp self.tmp_dir.cleanup() def test_process_variables_empty(self): options = PortableOptions([ '--environment_type=PROCESS', '--environment_option=process_command=foo', '--sdk_location=container' ]) environment = ProcessEnvironment.from_options(options) self.assertEqual(environment.command, 'foo') self.assertEqual(environment.env, {}) def test_process_variables_set(self): options = PortableOptions([ '--environment_type=PROCESS', '--environment_option=process_command=foo', '--environment_option=process_variables=' 'BASH_VARIABLE_ONE=spam,BASH_VARIABLE_TWO=ham', '--sdk_location=container' ]) environment = ProcessEnvironment.from_options(options) self.assertEqual(environment.command, 'foo') self.assertEqual( environment.env, { 'BASH_VARIABLE_ONE': 'spam', 'BASH_VARIABLE_TWO': 'ham' }) def test_process_variables_missing_rvalue(self): with self.assertRaises(ValueError): options = PortableOptions([ '--environment_type=PROCESS', '--environment_option=process_command=foo', '--environment_option=process_variables=' 'BASH_VARIABLE_ONE=spam,MISSING_RVALUE', '--sdk_location=container' ]) ProcessEnvironment.from_options(options) def test_environments_with_same_hints_are_equal(self): options = PortableOptions([ '--environment_type=PROCESS', '--environment_option=process_command=foo', '--sdk_location=container', '--resource_hint=accelerator=gpu', ]) environment1 = ProcessEnvironment.from_options(options) environment2 = ProcessEnvironment.from_options(options) self.assertEqual(environment1, environment2) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/error_handling.py0000644000175100017510000001067115206113033023430 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utilities for gracefully handling errors and excluding bad elements.""" import traceback from apache_beam import transforms class ErrorHandler: """ErrorHandlers are used to skip and otherwise process bad records. Error handlers allow one to implement the "dead letter queue" pattern in a fluent manner, disaggregating the error processing specification from the main processing chain. This is typically used as follows:: with error_handling.ErrorHandler(WriteToSomewhere(...)) as error_handler: result = pcoll | SomeTransform().with_error_handler(error_handler) in which case errors encountered by `SomeTransform()`` in processing pcoll will be written by the PTransform `WriteToSomewhere(...)` and excluded from `result` rather than failing the pipeline. To implement `with_error_handling` on a PTransform, one caches the provided error handler for use in `expand`. During `expand()` one can invoke `error_handler.add_error_pcollection(...)` any number of times with PCollections containing error records to be processed by the given error handler, or (if applicable) simply invoke `with_error_handling(...)` on any subtransforms. The `with_error_handling` should accept `None` to indicate that error handling is not enabled (and make implementation-by-forwarding-error-handlers easier). In this case, any non-recoverable errors should fail the pipeline (e.g. propagate exceptions in `process` methods) rather than silently ignore errors. """ def __init__(self, consumer): self._consumer = consumer self._creation_traceback = traceback.format_stack()[-2] self._error_pcolls = [] self._closed = False def __enter__(self): self._error_pcolls = [] self._closed = False return self def __exit__(self, *exec_info): if exec_info[0] is None: self.close() def close(self): """Indicates all error-producing operations have reported any errors. Invokes the provided error consuming PTransform on any provided error PCollections. """ self._output = ( tuple(self._error_pcolls) | transforms.Flatten() | self._consumer) self._closed = True def output(self): """Returns result of applying the error consumer to the error pcollections. """ if not self._closed: raise RuntimeError( "Cannot access the output of an error handler " "until it has been closed.") return self._output def add_error_pcollection(self, pcoll): """Called by a class implementing error handling on the error records. """ pcoll.pipeline._register_error_handler(self) self._error_pcolls.append(pcoll) def verify_closed(self): """Called at end of pipeline construction to ensure errors are not ignored. """ if not self._closed: raise RuntimeError( "Unclosed error handler initialized at %s" % self._creation_traceback) class _IdentityPTransform(transforms.PTransform): def expand(self, pcoll): return pcoll class CollectingErrorHandler(ErrorHandler): """An ErrorHandler that simply collects all errors for further processing. This ErrorHandler requires the set of errors be retrieved via `output()` and consumed (or explicitly discarded). """ def __init__(self): super().__init__(_IdentityPTransform()) self._creation_traceback = traceback.format_stack()[-2] self._output_accessed = False def output(self): self._output_accessed = True return super().output() def verify_closed(self): if not self._output_accessed: raise RuntimeError( "CollectingErrorHandler requires the output to be retrieved. " "Initialized at %s" % self._creation_traceback) return super().verify_closed() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/error_handling_test.py0000644000175100017510000001257515206113033024474 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest import apache_beam as beam from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import error_handling class PTransformWithErrors(beam.PTransform): def __init__(self, limit): self._limit = limit self._error_handler = None def with_error_handler(self, error_handler): self._error_handler = error_handler return self def expand(self, pcoll): limit = self._limit def process(element): if len(element) < limit: return element.title() else: return beam.pvalue.TaggedOutput('bad', element) def raise_on_everything(element): raise ValueError(element) good, bad = pcoll | beam.Map(process).with_outputs('bad', main='good') if self._error_handler: self._error_handler.add_error_pcollection(bad) else: # Will throw an exception if there are any bad elements. _ = bad | beam.Map(raise_on_everything) return good def exception_throwing_map(x, limit): if len(x) > limit: raise ValueError(x) else: return x.title() class ErrorHandlingTest(unittest.TestCase): def test_error_handling(self): with beam.Pipeline() as p: pcoll = p | beam.Create(['a', 'bb', 'cccc']) with error_handling.ErrorHandler( beam.Map(lambda x: "error: %s" % x)) as error_handler: result = pcoll | PTransformWithErrors(3).with_error_handler( error_handler) error_pcoll = error_handler.output() assert_that(result, equal_to(['A', 'Bb']), label='CheckGood') assert_that(error_pcoll, equal_to(['error: cccc']), label='CheckBad') def test_error_handling_pardo(self): with beam.Pipeline() as p: pcoll = p | beam.Create(['a', 'bb', 'cccc']) with error_handling.ErrorHandler( beam.Map(lambda x: "error: %s" % x[0])) as error_handler: result = pcoll | beam.Map( exception_throwing_map, limit=3).with_error_handler(error_handler) error_pcoll = error_handler.output() assert_that(result, equal_to(['A', 'Bb']), label='CheckGood') assert_that(error_pcoll, equal_to(['error: cccc']), label='CheckBad') def test_error_handling_pardo_with_exception_handling_kwargs(self): def side_effect(*args): beam._test_error_handling_pardo_with_exception_handling_kwargs_val = True def check_side_effect(): return getattr( beam, '_test_error_handling_pardo_with_exception_handling_kwargs_val', False) self.assertFalse(check_side_effect()) with beam.Pipeline() as p: pcoll = p | beam.Create(['a', 'bb', 'cccc']) with error_handling.ErrorHandler( beam.Map(lambda x: "error: %s" % x[0])) as error_handler: result = pcoll | beam.Map( exception_throwing_map, limit=3).with_error_handler( error_handler, on_failure_callback=side_effect) error_pcoll = error_handler.output() assert_that(result, equal_to(['A', 'Bb']), label='CheckGood') assert_that(error_pcoll, equal_to(['error: cccc']), label='CheckBad') self.assertTrue(check_side_effect()) def test_error_on_unclosed_error_handler(self): with self.assertRaisesRegex(RuntimeError, r'.*Unclosed error handler.*'): with beam.Pipeline() as p: pcoll = p | beam.Create(['a', 'bb', 'cccc']) # Use this outside of a context to allow it to remain unclosed. error_handler = error_handling.ErrorHandler(beam.Map(lambda x: x)) _ = pcoll | PTransformWithErrors(3).with_error_handler(error_handler) def test_collecting_error_handler(self): with beam.Pipeline() as p: pcoll = p | beam.Create(['a', 'bb', 'cccc']) with error_handling.CollectingErrorHandler() as error_handler: result = pcoll | beam.Map( exception_throwing_map, limit=3).with_error_handler(error_handler) error_pcoll = error_handler.output() | beam.Map(lambda x: x[0]) assert_that(result, equal_to(['A', 'Bb']), label='CheckGood') assert_that(error_pcoll, equal_to(['cccc']), label='CheckBad') def test_error_on_collecting_error_handler_without_output_retrieval(self): with self.assertRaisesRegex( RuntimeError, r'.*CollectingErrorHandler requires the output to be retrieved.*'): with beam.Pipeline() as p: pcoll = p | beam.Create(['a', 'bb', 'cccc']) with error_handling.CollectingErrorHandler() as error_handler: _ = pcoll | beam.Map( exception_throwing_map, limit=3).with_error_handler(error_handler) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/external.py0000644000175100017510000014052315206113033022255 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Defines Transform whose expansion is implemented elsewhere.""" # pytype: skip-file import contextlib import copy import functools import glob import logging import re import subprocess import threading import uuid from collections import OrderedDict from collections import namedtuple import grpc import yaml from apache_beam import pvalue from apache_beam.coders import RowCoder from apache_beam.options import pipeline_options from apache_beam.portability import common_urns from apache_beam.portability.api import beam_artifact_api_pb2_grpc from apache_beam.portability.api import beam_expansion_api_pb2 from apache_beam.portability.api import beam_expansion_api_pb2_grpc from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.portability.api import external_transforms_pb2 from apache_beam.portability.api import schema_pb2 from apache_beam.portability.common_urns import ManagedTransforms from apache_beam.runners import pipeline_context from apache_beam.runners.portability import artifact_service from apache_beam.transforms import environments from apache_beam.transforms import ptransform from apache_beam.typehints import WithTypeHints from apache_beam.typehints import native_type_compatibility from apache_beam.typehints import row_type from apache_beam.typehints.schemas import named_fields_to_schema from apache_beam.typehints.schemas import named_tuple_from_schema from apache_beam.typehints.schemas import named_tuple_to_schema from apache_beam.typehints.schemas import typing_from_runner_api from apache_beam.typehints.trivial_inference import instance_to_type from apache_beam.typehints.typehints import Union from apache_beam.typehints.typehints import UnionConstraint from apache_beam.utils import subprocess_server from apache_beam.utils import transform_service_launcher DEFAULT_EXPANSION_SERVICE = 'localhost:8097' MANAGED_SCHEMA_TRANSFORM_IDENTIFIER = "beam:transform:managed:v1" _IO_EXPANSION_SERVICE_JAR_TARGET = "sdks:java:io:expansion-service:shadowJar" _GCP_EXPANSION_SERVICE_JAR_TARGET = ( "sdks:java:io:google-cloud-platform:expansion-service:shadowJar") # A mapping from supported managed transforms URNs to expansion service jars # that include the corresponding transforms. MANAGED_TRANSFORM_URN_TO_JAR_TARGET_MAPPING = { ManagedTransforms.Urns.ICEBERG_READ.urn: _IO_EXPANSION_SERVICE_JAR_TARGET, ManagedTransforms.Urns.ICEBERG_WRITE.urn: _IO_EXPANSION_SERVICE_JAR_TARGET, ManagedTransforms.Urns.ICEBERG_CDC_READ.urn: _IO_EXPANSION_SERVICE_JAR_TARGET, # pylint: disable=line-too-long ManagedTransforms.Urns.KAFKA_READ.urn: _IO_EXPANSION_SERVICE_JAR_TARGET, ManagedTransforms.Urns.KAFKA_WRITE.urn: _IO_EXPANSION_SERVICE_JAR_TARGET, ManagedTransforms.Urns.BIGQUERY_READ.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET, ManagedTransforms.Urns.BIGQUERY_WRITE.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET, # pylint: disable=line-too-long ManagedTransforms.Urns.POSTGRES_READ.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET, ManagedTransforms.Urns.POSTGRES_WRITE.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET, # pylint: disable=line-too-long ManagedTransforms.Urns.MYSQL_READ.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET, ManagedTransforms.Urns.MYSQL_WRITE.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET, ManagedTransforms.Urns.SQL_SERVER_READ.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET, # pylint: disable=line-too-long ManagedTransforms.Urns.SQL_SERVER_WRITE.urn: _GCP_EXPANSION_SERVICE_JAR_TARGET, # pylint: disable=line-too-long } def convert_to_typing_type(type_): if isinstance(type_, row_type.RowTypeConstraint): return named_tuple_from_schema(named_fields_to_schema(type_._fields)) else: return native_type_compatibility.convert_to_python_type(type_) def _is_optional_or_none(typehint): return ( type(None) in typehint.union_types if isinstance( typehint, UnionConstraint) else typehint is type(None)) def _strip_optional(typehint): if not _is_optional_or_none(typehint): return typehint new_types = typehint.union_types.difference({type(None)}) if len(new_types) == 1: return list(new_types)[0] return Union[new_types] def iter_urns(coder, context=None): yield coder.to_runner_api_parameter(context)[0] for child in coder._get_component_coders(): for urn in iter_urns(child, context): yield urn class PayloadBuilder(object): """ Abstract base class for building payloads to pass to ExternalTransform. """ def build(self): """ :return: ExternalConfigurationPayload """ raise NotImplementedError def payload(self): """ The serialized ExternalConfigurationPayload :return: bytes """ return self.build().SerializeToString() def _get_schema_proto_and_payload(self, **kwargs): named_fields = [] fields_to_values = OrderedDict() for key, value in kwargs.items(): if not key: raise ValueError('Parameter name cannot be empty') if value is None: raise ValueError( 'Received value None for key %s. None values are currently not ' 'supported' % key) named_fields.append( (key, convert_to_typing_type(instance_to_type(value)))) fields_to_values[key] = value schema_proto = named_fields_to_schema(named_fields) row = named_tuple_from_schema(schema_proto)(**fields_to_values) schema = named_tuple_to_schema(type(row)) payload = RowCoder(schema).encode(row) return (schema_proto, payload) class SchemaBasedPayloadBuilder(PayloadBuilder): """ Base class for building payloads based on a schema that provides type information for each configuration value to encode. """ def _get_named_tuple_instance(self): raise NotImplementedError() def build(self): row = self._get_named_tuple_instance() schema = named_tuple_to_schema(type(row)) return external_transforms_pb2.ExternalConfigurationPayload( schema=schema, payload=RowCoder(schema).encode(row)) class ImplicitSchemaPayloadBuilder(SchemaBasedPayloadBuilder): """ Build a payload that generates a schema from the provided values. """ def __init__(self, values): self._values = values def _get_named_tuple_instance(self): # omit fields with value=None since we can't infer their type values = { key: value for key, value in self._values.items() if value is not None } schema = named_fields_to_schema([ (key, convert_to_typing_type(instance_to_type(value))) for key, value in values.items() ]) return named_tuple_from_schema(schema)(**values) class NamedTupleBasedPayloadBuilder(SchemaBasedPayloadBuilder): """ Build a payload based on a NamedTuple schema. """ def __init__(self, tuple_instance): """ :param tuple_instance: an instance of a typing.NamedTuple """ super().__init__() self._tuple_instance = tuple_instance def _get_named_tuple_instance(self): return self._tuple_instance class SchemaTransformPayloadBuilder(PayloadBuilder): def __init__(self, identifier, **kwargs): self._identifier = identifier self._kwargs = kwargs def identifier(self): """ The URN referencing this SchemaTransform :return: str """ return self._identifier def build(self): schema_proto, payload = self._get_schema_proto_and_payload(**self._kwargs) payload = external_transforms_pb2.SchemaTransformPayload( identifier=self._identifier, configuration_schema=schema_proto, configuration_row=payload) return payload class ExplicitSchemaTransformPayloadBuilder(SchemaTransformPayloadBuilder): def __init__(self, identifier, schema_proto, **kwargs): self._identifier = identifier self._schema_proto = schema_proto self._kwargs = kwargs def build(self): def dict_to_row_recursive(field_type, py_value): if py_value is None: return None type_info = field_type.WhichOneof('type_info') if type_info == 'row_type': return dict_to_row(field_type.row_type.schema, py_value) elif type_info == 'array_type': return [ dict_to_row_recursive(field_type.array_type.element_type, value) for value in py_value ] elif type_info == 'map_type': return { key: dict_to_row_recursive(field_type.map_type.value_type, value) for key, value in py_value.items() } else: return py_value def dict_to_row(schema_proto, py_value): row_type = named_tuple_from_schema(schema_proto) if isinstance(py_value, dict): extra = set(py_value.keys()) - set(row_type._fields) if extra: raise ValueError( f"Transform '{self.identifier()}' was configured with unknown " f"fields: {extra}. Valid fields: {set(row_type._fields)}") return row_type( *[ dict_to_row_recursive( field.type, py_value.get(field.name, None)) for field in schema_proto.fields ]) else: return row_type(py_value) return external_transforms_pb2.SchemaTransformPayload( identifier=self._identifier, configuration_schema=self._schema_proto, configuration_row=RowCoder(self._schema_proto).encode( dict_to_row(self._schema_proto, self._kwargs))) class JavaClassLookupPayloadBuilder(PayloadBuilder): """ Builds a payload for directly instantiating a Java transform using a constructor and builder methods. """ IGNORED_ARG_FORMAT = 'ignore%d' def __init__(self, class_name): """ :param class_name: fully qualified name of the transform class. """ if not class_name: raise ValueError('Class name must not be empty') self._class_name = class_name self._constructor_method = None self._constructor_param_args = None self._constructor_param_kwargs = None self._builder_methods_and_params = OrderedDict() def _args_to_named_fields(self, args): next_field_id = 0 named_fields = OrderedDict() for value in args: if value is None: raise ValueError( 'Received value None. None values are currently not supported') named_fields[( JavaClassLookupPayloadBuilder.IGNORED_ARG_FORMAT % next_field_id)] = value next_field_id += 1 return named_fields def build(self): all_constructor_param_kwargs = self._args_to_named_fields( self._constructor_param_args) if self._constructor_param_kwargs: all_constructor_param_kwargs.update(self._constructor_param_kwargs) constructor_schema, constructor_payload = ( self._get_schema_proto_and_payload(**all_constructor_param_kwargs)) payload = external_transforms_pb2.JavaClassLookupPayload( class_name=self._class_name, constructor_schema=constructor_schema, constructor_payload=constructor_payload) if self._constructor_method: payload.constructor_method = self._constructor_method for builder_method_name, params in self._builder_methods_and_params.items(): builder_method_args, builder_method_kwargs = params all_builder_method_kwargs = self._args_to_named_fields( builder_method_args) if builder_method_kwargs: all_builder_method_kwargs.update(builder_method_kwargs) builder_method_schema, builder_method_payload = ( self._get_schema_proto_and_payload(**all_builder_method_kwargs)) builder_method = external_transforms_pb2.BuilderMethod( name=builder_method_name, schema=builder_method_schema, payload=builder_method_payload) builder_method.name = builder_method_name payload.builder_methods.append(builder_method) return payload def with_constructor(self, *args, **kwargs): """ Specifies the Java constructor to use. Arguments provided using args and kwargs will be applied to the Java transform constructor in the specified order. :param args: parameter values of the constructor. :param kwargs: parameter names and values of the constructor. """ if self._has_constructor(): raise ValueError( 'Constructor or constructor method can only be specified once') self._constructor_param_args = args self._constructor_param_kwargs = kwargs def with_constructor_method(self, method_name, *args, **kwargs): """ Specifies the Java constructor method to use. Arguments provided using args and kwargs will be applied to the Java transform constructor method in the specified order. :param method_name: name of the constructor method. :param args: parameter values of the constructor method. :param kwargs: parameter names and values of the constructor method. """ if self._has_constructor(): raise ValueError( 'Constructor or constructor method can only be specified once') self._constructor_method = method_name self._constructor_param_args = args self._constructor_param_kwargs = kwargs def add_builder_method(self, method_name, *args, **kwargs): """ Specifies a Java builder method to be invoked after instantiating the Java transform class. Specified builder method will be applied in order. Arguments provided using args and kwargs will be applied to the Java transform builder method in the specified order. :param method_name: name of the builder method. :param args: parameter values of the builder method. :param kwargs: parameter names and values of the builder method. """ self._builder_methods_and_params[method_name] = (args, kwargs) def _has_constructor(self): return ( self._constructor_method or self._constructor_param_args or self._constructor_param_kwargs) # Information regarding a SchemaTransform available in an external SDK. SchemaTransformsConfig = namedtuple( 'SchemaTransformsConfig', ['identifier', 'configuration_schema', 'inputs', 'outputs', 'description']) ManagedReplacement = namedtuple( 'ManagedReplacement', ['underlying_transform_identifier', 'update_compatibility_version']) class SchemaAwareExternalTransform(ptransform.PTransform): """A proxy transform for SchemaTransforms implemented in external SDKs. This allows Python pipelines to directly use existing SchemaTransforms available to the expansion service without adding additional code in external SDKs. :param identifier: unique identifier of the SchemaTransform. :param expansion_service: an expansion service to use. This should already be available and the Schema-aware transforms to be used must already be deployed. :param rearrange_based_on_discovery: if this flag is set, the input kwargs will be rearranged to match the order of fields in the external SchemaTransform configuration. A discovery call will be made to fetch the configuration. :param classpath: (Optional) A list paths to additional jars to place on the expansion service classpath. :param managed_replacement: (Optional) a 'ManagedReplacement' namedtuple that defines information needed to replace the transform with an equivalent managed transform during the expansion. If an 'updateCompatibilityBeamVersion' pipeline option is provided, we will only replace if the managed transform is update compatible with the provided version. :kwargs: field name to value mapping for configuring the schema transform. keys map to the field names of the schema of the SchemaTransform (in-order). """ def __init__( self, identifier, expansion_service, rearrange_based_on_discovery=False, classpath=None, managed_replacement=None, **kwargs): self._expansion_service = expansion_service self._kwargs = kwargs self._classpath = classpath if managed_replacement: assert isinstance(managed_replacement, ManagedReplacement) self._managed_replacement = managed_replacement _kwargs = kwargs if rearrange_based_on_discovery: config = SchemaAwareExternalTransform.discover_config( self._expansion_service, identifier) self._payload_builder = ExplicitSchemaTransformPayloadBuilder( identifier, named_tuple_to_schema(config.configuration_schema), **_kwargs) if self._managed_replacement: # We have to do the replacement at the expansion instead of at # construction # since we don't have access to the PipelineOptions object at the # construction. underlying_transform_id = ( self._managed_replacement.underlying_transform_identifier) if not (underlying_transform_id in MANAGED_TRANSFORM_URN_TO_JAR_TARGET_MAPPING): raise ValueError( 'Could not find an expansion service jar for the managed ' + 'transform ' + underlying_transform_id) managed_expansion_service_jar = ( MANAGED_TRANSFORM_URN_TO_JAR_TARGET_MAPPING )[underlying_transform_id] self._managed_expansion_service = BeamJarExpansionService( managed_expansion_service_jar) managed_config = SchemaAwareExternalTransform.discover_config( self._managed_expansion_service, MANAGED_SCHEMA_TRANSFORM_IDENTIFIER) yaml_config = yaml.dump(kwargs) self._managed_payload_builder = ( ExplicitSchemaTransformPayloadBuilder( MANAGED_SCHEMA_TRANSFORM_IDENTIFIER, named_tuple_to_schema(managed_config.configuration_schema), transform_identifier=underlying_transform_id, config=yaml_config)) else: self._payload_builder = SchemaTransformPayloadBuilder( identifier, **_kwargs) def expand(self, pcolls): # Expand the transform using the expansion service. payload_builder = self._payload_builder expansion_service = self._expansion_service if self._managed_replacement: compat_version_prior_to_current = ( pcolls.pipeline._options.is_compat_version_prior_to( self._managed_replacement.update_compatibility_version)) if not compat_version_prior_to_current: payload_builder = self._managed_payload_builder expansion_service = self._managed_expansion_service return pcolls | self._payload_builder.identifier() >> ExternalTransform( common_urns.schematransform_based_expand.urn, payload_builder, expansion_service) @classmethod @functools.lru_cache def discover(cls, expansion_service, ignore_errors=False): """Discover all SchemaTransforms available to the given expansion service. :return: a list of SchemaTransformsConfigs that represent the discovered SchemaTransforms. """ return list(cls.discover_iter(expansion_service, ignore_errors)) @staticmethod def discover_iter(expansion_service, ignore_errors=True): with ExternalTransform.service(expansion_service) as service: discover_response = service.DiscoverSchemaTransform( beam_expansion_api_pb2.DiscoverSchemaTransformRequest()) for identifier in discover_response.schema_transform_configs: proto_config = discover_response.schema_transform_configs[identifier] try: schema = named_tuple_from_schema(proto_config.config_schema) except Exception as exn: if ignore_errors: truncated_schema = schema_pb2.Schema() truncated_schema.CopyFrom(proto_config.config_schema) for field in truncated_schema.fields: try: typing_from_runner_api(field.type) except Exception: if field.type.nullable: # Set it to an empty placeholder type. field.type.CopyFrom( schema_pb2.FieldType( nullable=True, row_type=schema_pb2.RowType( schema=schema_pb2.Schema()))) try: schema = named_tuple_from_schema(truncated_schema) except Exception as exn: logging.info("Bad schema for %s: %s", identifier, str(exn)[:250]) continue else: raise yield SchemaTransformsConfig( identifier=identifier, configuration_schema=schema, inputs=proto_config.input_pcollection_names, outputs=proto_config.output_pcollection_names, description=proto_config.description) @staticmethod def discover_config(expansion_service, name): """Discover one SchemaTransform by name in the given expansion service. :return: one SchemaTransformsConfig that represents the discovered SchemaTransform :raises: ValueError: if more than one SchemaTransform is discovered, or if none are discovered """ schematransforms = SchemaAwareExternalTransform.discover( expansion_service, ignore_errors=True) matched = [] for st in schematransforms: if name in st.identifier: matched.append(st) if not matched: raise ValueError( "Did not discover any SchemaTransforms resembling the name '%s'" % name) elif len(matched) > 1: raise ValueError( "Found multiple SchemaTransforms with the name '%s':\n%s\n" % (name, [st.identifier for st in matched])) return matched[0] class JavaExternalTransform(ptransform.PTransform): """A proxy for Java-implemented external transforms. One builds these transforms just as one would in Java, e.g.:: transform = JavaExternalTransform('fully.qualified.ClassName' )(contructorArg, ... ).builderMethod(...) or:: JavaExternalTransform('fully.qualified.ClassName').staticConstructor( ...).builderMethod1(...).builderMethod2(...) :param class_name: fully qualified name of the java class :param expansion_service: (Optional) an expansion service to use. If none is provided, a default expansion service will be started. :param classpath: (Optional) A list paths to additional jars to place on the expansion service classpath. """ def __init__(self, class_name, expansion_service=None, classpath=None): if expansion_service and classpath: raise ValueError( f'Only one of expansion_service ({expansion_service}) ' f'or classpath ({classpath}) may be provided.') self._payload_builder = JavaClassLookupPayloadBuilder(class_name) self._classpath = classpath self._expansion_service = expansion_service # Beam explicitly looks for following attributes. Hence adding # 'None' values here to prevent '__getattr__' from being called. self.inputs = None self._fn_api_payload = None def __call__(self, *args, **kwargs): self._payload_builder.with_constructor(*args, **kwargs) return self def __getattr__(self, name): # Don't try to emulate special methods. if name.startswith('__') and name.endswith('__'): return super().__getattr__(name) else: return self[name] def __getitem__(self, name): # Use directly for keywords or attribute conflicts. def construct(*args, **kwargs): if self._payload_builder._has_constructor(): builder_method = self._payload_builder.add_builder_method else: builder_method = self._payload_builder.with_constructor_method builder_method(name, *args, **kwargs) return self return construct def expand(self, pcolls): if self._expansion_service is None: self._expansion_service = BeamJarExpansionService( ':sdks:java:expansion-service:app:shadowJar', extra_args=['{{PORT}}', '--javaClassLookupAllowlistFile=*'], classpath=self._classpath) return pcolls | ExternalTransform( common_urns.java_class_lookup.urn, self._payload_builder, self._expansion_service) class AnnotationBasedPayloadBuilder(SchemaBasedPayloadBuilder): """ Build a payload based on an external transform's type annotations. """ def __init__(self, transform, **values): """ :param transform: a PTransform instance or class. type annotations will be gathered from its __init__ method :param values: values to encode """ self._transform = transform self._values = values def _get_named_tuple_instance(self): schema = named_fields_to_schema([ (k, convert_to_typing_type(v)) for k, v in self._transform.__init__.__annotations__.items() if k in self._values ]) return named_tuple_from_schema(schema)(**self._values) class DataclassBasedPayloadBuilder(SchemaBasedPayloadBuilder): """ Build a payload based on an external transform that uses dataclasses. """ def __init__(self, transform): """ :param transform: a dataclass-decorated PTransform instance from which to gather type annotations and values """ self._transform = transform def _get_named_tuple_instance(self): import dataclasses schema = named_fields_to_schema([ (field.name, convert_to_typing_type(field.type)) for field in dataclasses.fields(self._transform) ]) return named_tuple_from_schema(schema)( **dataclasses.asdict(self._transform)) class ExternalTransform(ptransform.PTransform): """ External provides a cross-language transform via expansion services in foreign SDKs. """ _namespace_counter = 0 # Variable name _namespace conflicts with DisplayData._namespace so we use # name _external_namespace here. _external_namespace = threading.local() _IMPULSE_PREFIX = 'impulse' def __init__(self, urn, payload, expansion_service=None): """Wrapper for an external transform with the given urn and payload. :param urn: the unique beam identifier for this transform :param payload: the payload, either as a byte string or a PayloadBuilder :param expansion_service: an expansion service implementing the beam ExpansionService protocol, either as an object with an Expand method or an address (as a str) to a grpc server that provides this method. """ expansion_service = expansion_service or DEFAULT_EXPANSION_SERVICE if not urn and isinstance(payload, JavaClassLookupPayloadBuilder): urn = common_urns.java_class_lookup.urn self._urn = urn self._payload = ( payload.payload() if isinstance(payload, PayloadBuilder) else payload) self._expansion_service = expansion_service self._external_namespace = self._fresh_namespace() self._inputs: dict[str, pvalue.PCollection] = {} self._outputs: dict[str, pvalue.PCollection] = {} def with_output_types(self, *args, **kwargs): return WithTypeHints.with_output_types(self, *args, **kwargs) def replace_named_inputs(self, named_inputs): self._inputs = named_inputs def replace_named_outputs(self, named_outputs): self._outputs = named_outputs def __post_init__(self, expansion_service): """ This will only be invoked if ExternalTransform is used as a base class for a class decorated with dataclasses.dataclass """ ExternalTransform.__init__( self, self.URN, DataclassBasedPayloadBuilder(self), expansion_service) def default_label(self): return '%s(%s)' % (self.__class__.__name__, self._urn) @classmethod def get_local_namespace(cls): return getattr(cls._external_namespace, 'value', 'external') @classmethod @contextlib.contextmanager def outer_namespace(cls, namespace): prev = cls.get_local_namespace() try: cls._external_namespace.value = namespace yield finally: cls._external_namespace.value = prev @classmethod def _fresh_namespace(cls) -> str: ExternalTransform._namespace_counter += 1 return '%s_%d' % (cls.get_local_namespace(), cls._namespace_counter) def expand(self, pvalueish: pvalue.PCollection) -> pvalue.PCollection: if isinstance(pvalueish, pvalue.PBegin): self._inputs = {} elif isinstance(pvalueish, (list, tuple)): self._inputs = {str(ix): pvalue for ix, pvalue in enumerate(pvalueish)} elif isinstance(pvalueish, dict): self._inputs = pvalueish else: self._inputs = {'input': pvalueish} pipeline = ( next(iter(self._inputs.values())).pipeline if self._inputs else pvalueish.pipeline) context = pipeline_context.PipelineContext( component_id_map=pipeline.component_id_map) transform_proto = beam_runner_api_pb2.PTransform( unique_name=pipeline._current_transform().full_label, spec=beam_runner_api_pb2.FunctionSpec( urn=self._urn, payload=self._payload)) for tag, pcoll in self._inputs.items(): transform_proto.inputs[tag] = context.pcollections.get_id(pcoll) # Conversion to/from proto assumes producers. # TODO: Possibly loosen this. context.transforms.put_proto( '%s_%s' % (self._IMPULSE_PREFIX, tag), beam_runner_api_pb2.PTransform( unique_name='%s_%s' % (self._IMPULSE_PREFIX, tag), spec=beam_runner_api_pb2.FunctionSpec( urn=common_urns.primitives.IMPULSE.urn), outputs={'out': transform_proto.inputs[tag]})) # Retrieve type hints and store them in variables # to avoid duplicate calls and AttributeError hints = self.get_type_hints() output_coders = None if hints.output_types: if hints.output_types[0]: output_coders = dict((str(k), context.coder_id_from_element_type(v)) for (k, v) in enumerate(hints.output_types[0])) elif hints.output_types[1]: output_coders = { k: context.coder_id_from_element_type(v) for (k, v) in hints.output_types[1].items() } components = context.to_runner_api() request = beam_expansion_api_pb2.ExpansionRequest( components=components, namespace=self._external_namespace, # type: ignore[arg-type] transform=transform_proto, output_coder_requests=output_coders, pipeline_options=pipeline._options.to_runner_api()) expansion_service = _maybe_use_transform_service( self._expansion_service, pipeline.options) with ExternalTransform.service(expansion_service) as service: response = service.Expand(request) if response.error: raise RuntimeError(_sanitize_java_traceback(response.error)) self._expanded_components = response.components if any(e.dependencies for env in self._expanded_components.environments.values() for e in environments.expand_anyof_environments(env)): self._expanded_components = self._resolve_artifacts( self._expanded_components, service.artifact_service(), pipeline.local_tempdir) self._expanded_transform = response.transform self._expanded_requirements = response.requirements result_context = pipeline_context.PipelineContext(response.components) def fix_output(pcoll, tag): pcoll.pipeline = pipeline pcoll.tag = tag return pcoll self._outputs = { tag: fix_output(result_context.pcollections.get_by_id(pcoll_id), tag) for tag, pcoll_id in self._expanded_transform.outputs.items() } return self._output_to_pvalueish(self._outputs) @staticmethod @contextlib.contextmanager def service(expansion_service): if isinstance(expansion_service, str): channel_options = [("grpc.max_receive_message_length", -1), ("grpc.max_send_message_length", -1)] if hasattr(grpc, 'local_channel_credentials'): # Some environments may not support insecure channels. Hence use a # secure channel with local credentials here. # TODO: update this to support secure non-local channels. channel_factory_fn = functools.partial( grpc.secure_channel, expansion_service, grpc.local_channel_credentials(), options=channel_options) else: # local_channel_credentials is an experimental API which is unsupported # by older versions of grpc which may be pulled in due to other project # dependencies. channel_factory_fn = functools.partial( grpc.insecure_channel, expansion_service, options=channel_options) with channel_factory_fn() as channel: yield ExpansionAndArtifactRetrievalStub(channel) elif hasattr(expansion_service, 'Expand'): yield expansion_service else: with expansion_service as stub: yield stub def _resolve_artifacts(self, components, service, dest): def _resolve_artifacts_for(env): if env.urn == common_urns.environments.ANYOF.urn: env.CopyFrom( environments.AnyOfEnvironment.create_proto([ _resolve_artifacts_for(e) for e in environments.expand_anyof_environments(env) ])) elif env.dependencies: resolved = list( artifact_service.resolve_artifacts(env.dependencies, service, dest)) del env.dependencies[:] env.dependencies.extend(resolved) return env for env in components.environments.values(): _resolve_artifacts_for(env) return components def _output_to_pvalueish(self, output_dict): if len(output_dict) == 1: return next(iter(output_dict.values())) else: return output_dict def to_runner_api_transform(self, context, full_label): pcoll_renames = {} renamed_tag_seen = False for tag, pcoll in self._inputs.items(): if tag not in self._expanded_transform.inputs: if renamed_tag_seen: raise RuntimeError( 'Ambiguity due to non-preserved tags: %s vs %s' % ( sorted(self._expanded_transform.inputs.keys()), sorted(self._inputs.keys()))) else: renamed_tag_seen = True tag, = self._expanded_transform.inputs.keys() pcoll_renames[self._expanded_transform.inputs[tag]] = ( context.pcollections.get_id(pcoll)) for tag, pcoll in self._outputs.items(): pcoll_renames[self._expanded_transform.outputs[tag]] = ( context.pcollections.get_id(pcoll)) def _equivalent(coder1, coder2): return coder1 == coder2 or _normalize(coder1) == _normalize(coder2) def _normalize(coder_proto): normalized = copy.copy(coder_proto) normalized.spec.environment_id = '' # TODO(robertwb): Normalize components as well. return normalized for id, proto in self._expanded_components.coders.items(): if id.startswith(self._external_namespace): context.coders.put_proto(id, proto) elif id in context.coders: if not _equivalent(context.coders._id_to_proto[id], proto): raise RuntimeError( 'Re-used coder id: %s\n%s\n%s' % (id, context.coders._id_to_proto[id], proto)) else: context.coders.put_proto(id, proto) for id, proto in self._expanded_components.windowing_strategies.items(): if id.startswith(self._external_namespace): context.windowing_strategies.put_proto(id, proto) for id, proto in self._expanded_components.environments.items(): if id.startswith(self._external_namespace): context.environments.put_proto(id, proto) for id, proto in self._expanded_components.pcollections.items(): id = pcoll_renames.get(id, id) if id not in context.pcollections._id_to_obj.keys(): context.pcollections.put_proto(id, proto) for id, proto in self._expanded_components.transforms.items(): if id.startswith(self._IMPULSE_PREFIX): # Our fake inputs. continue assert id.startswith( self._external_namespace), (id, self._external_namespace) new_proto = beam_runner_api_pb2.PTransform( unique_name=proto.unique_name, # If URN is not set this is an empty spec. spec=proto.spec if proto.spec.urn else None, subtransforms=proto.subtransforms, inputs={ tag: pcoll_renames.get(pcoll, pcoll) for tag, pcoll in proto.inputs.items() }, outputs={ tag: pcoll_renames.get(pcoll, pcoll) for tag, pcoll in proto.outputs.items() }, display_data=proto.display_data, environment_id=proto.environment_id) context.transforms.put_proto(id, new_proto) for requirement in self._expanded_requirements: context.add_requirement(requirement) return beam_runner_api_pb2.PTransform( unique_name=full_label, spec=self._expanded_transform.spec, subtransforms=self._expanded_transform.subtransforms, inputs={ tag: pcoll_renames.get(pcoll, pcoll) for tag, pcoll in self._expanded_transform.inputs.items() }, outputs={ tag: pcoll_renames.get(pcoll, pcoll) for tag, pcoll in self._expanded_transform.outputs.items() }, annotations=self._expanded_transform.annotations, environment_id=self._expanded_transform.environment_id) class ExpansionAndArtifactRetrievalStub( beam_expansion_api_pb2_grpc.ExpansionServiceStub): def __init__(self, channel, **kwargs): self._channel = channel self._kwargs = kwargs super().__init__(channel, **kwargs) def artifact_service(self): return beam_artifact_api_pb2_grpc.ArtifactRetrievalServiceStub( self._channel, **self._kwargs) def ready(self, timeout_sec): grpc.channel_ready_future(self._channel).result(timeout=timeout_sec) class JavaJarExpansionService(object): """An expansion service based on an Java Jar file. This can be passed into an ExternalTransform as the expansion_service argument which will spawn a subprocess using this jar to expand the transform. Args: path_to_jar: the path to a locally available executable jar file to be used to start up the expansion service. extra_args: arguments to be provided when starting up the expansion service using the jar file. These arguments will replace the default arguments. classpath: Additional dependencies to be added to the classpath. append_args: arguments to be provided when starting up the expansion service using the jar file. These arguments will be appended to the default arguments. user_agent: HTTP user agent string used when downloading jars via `JavaJarServer.local_jar`, including the main jar and any classpath dependencies. maven_repository_url: Maven repository base URL to resolve artifacts when classpath entries or jars are specified as Maven coordinates (`group:artifact:version`). Defaults to Maven Central if not provided. """ def __init__( self, path_to_jar, extra_args=None, classpath=None, append_args=None, user_agent=None, maven_repository_url=None): if extra_args and append_args: raise ValueError('Only one of extra_args or append_args may be provided') self.path_to_jar = path_to_jar self._extra_args = extra_args self._classpath = classpath or [] self._service_count = 0 self._append_args = append_args or [] self._user_agent = user_agent self._maven_repository_url = maven_repository_url def is_existing_service(self): return subprocess_server.is_service_endpoint(self.path_to_jar) @staticmethod def _expand_jars(jar, user_agent=None, maven_repository_url=None): if glob.glob(jar): return glob.glob(jar) elif isinstance(jar, str) and (jar.startswith('http://') or jar.startswith('https://')): return [subprocess_server.JavaJarServer.local_jar(jar)] else: # If the input JAR is not a local glob, nor an http/https URL, then # we assume that it's a gradle-style Java artifact in Maven Central, # in the form group:artifact:version, so we attempt to parse that way. try: group_id, artifact_id, version = jar.split(':') except ValueError: # If we are not able to find a JAR, nor a JAR artifact, nor a URL for # a JAR path, we still choose to include it in the path. logging.warning('Unable to parse %s into group:artifact:version.', jar) return [jar] path = subprocess_server.JavaJarServer.local_jar( subprocess_server.JavaJarServer.path_to_maven_jar( artifact_id, group_id, version, repository=( maven_repository_url or subprocess_server.JavaJarServer.MAVEN_CENTRAL_REPOSITORY)), user_agent=user_agent) return [path] def _default_args(self): """Default arguments to be used by `JavaJarExpansionService`.""" to_stage = ','.join([self.path_to_jar] + sum(( JavaJarExpansionService._expand_jars( jar, self._user_agent, self._maven_repository_url) for jar in self._classpath or []), [])) args = ['{{PORT}}', f'--filesToStage={to_stage}'] # TODO(robertwb): See if it's possible to scope this per pipeline. # Checks to see if the cache is being used for this server. if subprocess_server.SubprocessServer._cache._live_owners: args.append('--alsoStartLoopbackWorker') return args def with_user_agent(self, user_agent: str): self._user_agent = user_agent return self def __enter__(self): if self._service_count == 0: self.path_to_jar = subprocess_server.JavaJarServer.local_jar( self.path_to_jar, user_agent=self._user_agent) if self._extra_args is None: self._extra_args = self._default_args() + self._append_args # Consider memoizing these servers (with some timeout). logging.info( 'Starting a JAR-based expansion service from JAR %s ' + ( 'and with classpath: %s' % self._classpath if self._classpath else ''), self.path_to_jar) classpath_urls = [ subprocess_server.JavaJarServer.local_jar(path) for jar in self._classpath for path in JavaJarExpansionService._expand_jars( jar, user_agent=self._user_agent, maven_repository_url=self._maven_repository_url) ] self._service_provider = subprocess_server.JavaJarServer( ExpansionAndArtifactRetrievalStub, self.path_to_jar, self._extra_args, classpath=classpath_urls, logger="ExpansionService") self._service = self._service_provider.__enter__() self._service_count += 1 return self._service def __exit__(self, *args): self._service_count -= 1 if self._service_count == 0: self._service_provider.__exit__(*args) class BeamJarExpansionService(JavaJarExpansionService): """An expansion service based on an Beam Java Jar file. Attempts to use a locally-built copy of the jar based on the gradle target, if it exists, otherwise attempts to download and cache the released artifact corresponding to this version of Beam from the apache maven repository. Args: gradle_target: Beam Gradle target for building an executable jar which will be used to start the expansion service. extra_args: arguments to be provided when starting up the expansion service using the jar file. These arguments will replace the default arguments. gradle_appendix: Gradle appendix of the artifact. classpath: Additional dependencies to be added to the classpath. append_args: arguments to be provided when starting up the expansion service using the jar file. These arguments will be appended to the default arguments. user_agent: HTTP user agent string used when downloading the Beam jar and any classpath dependencies. maven_repository_url: Maven repository base URL to resolve the Beam jar for the provided Gradle target. Defaults to Maven Central if not provided. """ def __init__( self, gradle_target, extra_args=None, gradle_appendix=None, classpath=None, append_args=None, user_agent=None, maven_repository_url=None): path_to_jar = subprocess_server.JavaJarServer.path_to_beam_jar( gradle_target, gradle_appendix, maven_repository_url=maven_repository_url) self.gradle_target = gradle_target super().__init__( path_to_jar, extra_args, classpath=classpath, append_args=append_args, user_agent=user_agent, maven_repository_url=maven_repository_url) def _maybe_use_transform_service(provided_service=None, options=None): # For anything other than 'JavaJarExpansionService' we just use the # provided service. For example, string address of an already available # service. if not isinstance(provided_service, JavaJarExpansionService): return provided_service if provided_service.is_existing_service(): # This is an existing service supported through the 'beam_services' # PipelineOption. return provided_service def is_java_available(): java = subprocess_server.JavaHelper.get_java() cmd = [java, '--version'] try: subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) except: # pylint: disable=bare-except return False return True def is_docker_available(): cmd = ['docker', '--version'] try: subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) except: # pylint: disable=bare-except return False return True # We try java and docker based expansion services in that order. java_available = is_java_available() docker_available = is_docker_available() use_transform_service = options.view_as( pipeline_options.CrossLanguageOptions).use_transform_service user_agent = options.view_as(pipeline_options.SetupOptions).user_agent if (java_available and provided_service and not use_transform_service): return provided_service.with_user_agent(user_agent) elif docker_available: if use_transform_service: error_append = 'it was explicitly requested' elif not java_available: error_append = 'the Java executable is not available in the system' else: error_append = 'a Java expansion service was not provided.' project_name = str(uuid.uuid4()) port = subprocess_server.pick_port(None)[0] logging.info( 'Trying to expand the external transform using the Docker Compose ' 'based transform service since %s. Transform service will be under ' 'Docker Compose project name %s and will be made available at port %r.' % (error_append, project_name, str(port))) from apache_beam import version as beam_version beam_version = beam_version.__version__ return transform_service_launcher.TransformServiceLauncher( project_name, port, beam_version, user_agent) else: raise ValueError( 'Cannot start an expansion service since neither Java nor ' 'Docker executables are available in the system.') def _sanitize_java_traceback(s): """Attempts to highlight the root cause in the error string. Java tracebacks read bottom to top, while Python tracebacks read top to bottom, resulting in the actual error message getting sandwiched between two walls of text. This may result in the error being duplicated (as we don't want to remove relevant information) but should be clearer in most cases. Best-effort but non-destructive. """ # We delete non-java-traceback lines. traceback_lines = [ r'\tat \S+\(\S+\.java:\d+\)', r'\t\.\.\. \d+ more', # A bit more restrictive to avoid accidentally capturing non-java lines. r'Caused by: [a-z]+(\.\S+)?\.[A-Z][A-Za-z0-9_$]+(Error|Exception):[^\n]*' ] without_java_traceback = s + '\n' for p in traceback_lines: without_java_traceback = re.sub( fr'\n{p}$', '', without_java_traceback, flags=re.M) # If what's left is substantially smaller, duplicate it at the end for better # visibility. if len(without_java_traceback) < len(s) / 2: return s + '\n\n' + without_java_traceback.strip() else: return s def memoize(func): cache = {} def wrapper(*args): if args not in cache: cache[args] = func(*args) return cache[args] return wrapper ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/external_it_test.py0000644000175100017510000000437315206113033024012 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration tests for cross-language transform expansion.""" # pytype: skip-file import unittest import pytest import apache_beam as beam from apache_beam import Pipeline from apache_beam.runners.portability import expansion_service from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import ptransform class ExternalTransformIT(unittest.TestCase): @pytest.mark.it_postcommit def test_job_python_from_python_it(self): @ptransform.PTransform.register_urn('simple', None) class SimpleTransform(ptransform.PTransform): def expand(self, pcoll): return pcoll | beam.Map(lambda x: 'Simple(%s)' % x) def to_runner_api_parameter(self, unused_context): return 'simple', None @staticmethod def from_runner_api_parameter(_0, _1, _2): return SimpleTransform() pipeline = TestPipeline(is_integration_test=True) res = ( pipeline | beam.Create(['a', 'b']) | beam.ExternalTransform( 'simple', None, expansion_service.ExpansionServiceServicer())) assert_that(res, equal_to(['Simple(a)', 'Simple(b)'])) proto_pipeline, _ = pipeline.to_runner_api(return_context=True) pipeline_from_proto = Pipeline.from_runner_api( proto_pipeline, pipeline.runner, pipeline._options) pipeline_from_proto.run().wait_until_finish() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/external_java.py0000644000175100017510000001561415206113033023260 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for the Java external transforms.""" import argparse import logging import subprocess import sys from typing import Optional import grpc from mock import patch import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.external import ImplicitSchemaPayloadBuilder from apache_beam.utils import subprocess_server # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.runners.dataflow.internal import apiclient as _apiclient except ImportError: apiclient = None else: apiclient = _apiclient # pylint: enable=wrong-import-order, wrong-import-position class JavaExternalTransformTest(object): # This will be overwritten if set via a flag. expansion_service_jar: Optional[str] = None expansion_service_port: Optional[int] = None class _RunWithExpansion(object): def __init__(self): self._server = None def __enter__(self): if not (JavaExternalTransformTest.expansion_service_jar or JavaExternalTransformTest.expansion_service_port): raise RuntimeError('No expansion service jar or port provided.') JavaExternalTransformTest.expansion_service_port = ( JavaExternalTransformTest.expansion_service_port or 8091) jar = JavaExternalTransformTest.expansion_service_jar port = JavaExternalTransformTest.expansion_service_port # Start the java server and wait for it to be ready. if jar: java = subprocess_server.JavaHelper.get_java() self._server = subprocess.Popen([java, '-jar', jar, str(port)]) address = 'localhost:%s' % str(port) with grpc.insecure_channel(address) as channel: grpc.channel_ready_future(channel).result() def __exit__(self, type, value, traceback): if self._server: self._server.kill() self._server = None @staticmethod def test_java_expansion_dataflow(): if apiclient is None: return # This test does not actually running the pipeline in Dataflow. It just # tests the translation to a Dataflow job request. with patch.object(apiclient.DataflowApplicationClient, 'create_job') as mock_create_job: with JavaExternalTransformTest._RunWithExpansion(): pipeline_options = PipelineOptions([ '--runner=DataflowRunner', '--project=dummyproject', '--region=some-region1', '--experiments=beam_fn_api', '--temp_location=gs://dummybucket/' ]) # Run a simple count-filtered-letters pipeline. JavaExternalTransformTest.run_pipeline( pipeline_options, JavaExternalTransformTest.expansion_service_port, False) mock_args = mock_create_job.call_args_list assert mock_args args, kwargs = mock_args[0] job = args[0] job_str = '%s' % job assert 'beam:transforms:xlang:filter_less_than_eq' in job_str @staticmethod def run_pipeline_with_expansion_service(pipeline_options): with JavaExternalTransformTest._RunWithExpansion(): # Run a simple count-filtered-letters pipeline. JavaExternalTransformTest.run_pipeline( pipeline_options, JavaExternalTransformTest.expansion_service_port, True) @staticmethod def run_pipeline(pipeline_options, expansion_service, wait_until_finish=True): # The actual definitions of these transforms is in # org.apache.beam.sdk.util.construction.TestExpansionService. TEST_COUNT_URN = "beam:transforms:xlang:count" TEST_FILTER_URN = "beam:transforms:xlang:filter_less_than_eq" # Run a simple count-filtered-letters pipeline. p = TestPipeline(options=pipeline_options) if isinstance(expansion_service, int): # Only the port was specified. expansion_service = 'localhost:%s' % str(expansion_service) res = ( p | beam.Create(list('aaabccxyyzzz')) | beam.Map(str) | beam.ExternalTransform( TEST_FILTER_URN, ImplicitSchemaPayloadBuilder({'data': 'middle'}), expansion_service) | beam.ExternalTransform(TEST_COUNT_URN, None, expansion_service) # Map(lambda) produces a label formatted like this, but it cannot be # changed without breaking update compat. Here, we pin to the transform # name used in the 2.68 release to avoid breaking changes when the line # number changes. Context: https://github.com/apache/beam/pull/36381 | "Map()" >> beam.Map(lambda kv: '%s: %s' % kv)) assert_that(res, equal_to(['a: 3', 'b: 1', 'c: 2'])) result = p.run() if wait_until_finish: result.wait_until_finish() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) parser = argparse.ArgumentParser() parser.add_argument('--expansion_service_jar') parser.add_argument('--expansion_service_port') parser.add_argument('--expansion_service_target') parser.add_argument('--expansion_service_target_appendix') known_args, pipeline_args = parser.parse_known_args(sys.argv) if known_args.expansion_service_jar: JavaExternalTransformTest.expansion_service_jar = ( known_args.expansion_service_jar) JavaExternalTransformTest.expansion_service_port = int( known_args.expansion_service_port) pipeline_options = PipelineOptions(pipeline_args) JavaExternalTransformTest.run_pipeline_with_expansion_service( pipeline_options) elif known_args.expansion_service_target: pipeline_options = PipelineOptions(pipeline_args) JavaExternalTransformTest.run_pipeline( pipeline_options, beam.transforms.external.BeamJarExpansionService( known_args.expansion_service_target, gradle_appendix=known_args.expansion_service_target_appendix)) else: raise RuntimeError( "--expansion_service_jar or --expansion_service_target " "should be provided.") ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/external_test.py0000644000175100017510000010530215206113033023310 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the transform.external classes.""" # pytype: skip-file import dataclasses import logging import os import tempfile import typing import unittest import mock import apache_beam as beam from apache_beam import ManagedReplacement from apache_beam import Pipeline from apache_beam.coders import RowCoder from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.portability.api import beam_expansion_api_pb2 from apache_beam.portability.api import external_transforms_pb2 from apache_beam.portability.api import schema_pb2 from apache_beam.portability.common_urns import ManagedTransforms from apache_beam.runners import pipeline_context from apache_beam.runners.portability import expansion_service from apache_beam.runners.portability.expansion_service_test import FibTransform from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import external from apache_beam.transforms.external import MANAGED_SCHEMA_TRANSFORM_IDENTIFIER from apache_beam.transforms.external import AnnotationBasedPayloadBuilder from apache_beam.transforms.external import ImplicitSchemaPayloadBuilder from apache_beam.transforms.external import JavaClassLookupPayloadBuilder from apache_beam.transforms.external import JavaExternalTransform from apache_beam.transforms.external import JavaJarExpansionService from apache_beam.transforms.external import NamedTupleBasedPayloadBuilder from apache_beam.transforms.external import SchemaTransformPayloadBuilder from apache_beam.typehints import typehints from apache_beam.typehints.native_type_compatibility import convert_to_beam_type from apache_beam.utils import proto_utils from apache_beam.utils.subprocess_server import JavaJarServer from apache_beam.utils.subprocess_server import SubprocessServer # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position try: from apache_beam.runners.dataflow.internal import apiclient except ImportError: apiclient = None # type: ignore # pylint: enable=wrong-import-order, wrong-import-position def get_payload(cls): payload = external_transforms_pb2.ExternalConfigurationPayload() payload.ParseFromString(cls._payload) return payload class PayloadBase(object): values = { 'integer_example': 1, 'boolean': True, 'string_example': 'thing', 'list_of_strings': ['foo', 'bar'], 'mapping': { 'key': 1.1 }, 'optional_integer': None, } bytes_values = { 'integer_example': 1, 'boolean': True, 'string_example': 'thing', 'list_of_strings': ['foo', 'bar'], 'mapping': { 'key': 1.1 }, 'optional_integer': None, } def get_payload_from_typing_hints(self, values): """Return ExternalConfigurationPayload based on python typing hints""" raise NotImplementedError def get_payload_from_beam_typehints(self, values): """Return ExternalConfigurationPayload based on beam typehints""" raise NotImplementedError def test_typing_payload_builder(self): result = self.get_payload_from_typing_hints(self.values) decoded = RowCoder(result.schema).decode(result.payload) for key, value in self.values.items(): self.assertEqual(getattr(decoded, key), value) def test_typehints_payload_builder(self): result = self.get_payload_from_typing_hints(self.values) decoded = RowCoder(result.schema).decode(result.payload) for key, value in self.values.items(): self.assertEqual(getattr(decoded, key), value) def test_optional_error(self): """ value can only be None if typehint is Optional """ with self.assertRaises(ValueError): self.get_payload_from_typing_hints({k: None for k in self.values}) class ExternalTuplePayloadTest(PayloadBase, unittest.TestCase): def get_payload_from_typing_hints(self, values): TestSchema = typing.NamedTuple( 'TestSchema', [ ('integer_example', int), ('boolean', bool), ('string_example', str), ('list_of_strings', list[str]), ('mapping', typing.Mapping[str, float]), ('optional_integer', typing.Optional[int]), ]) builder = NamedTupleBasedPayloadBuilder(TestSchema(**values)) return builder.build() def get_payload_from_beam_typehints(self, values): raise unittest.SkipTest( "Beam typehints cannot be used with " "typing.NamedTuple") class ExternalImplicitPayloadTest(unittest.TestCase): """ ImplicitSchemaPayloadBuilder works very differently than the other payload builders """ def test_implicit_payload_builder(self): builder = ImplicitSchemaPayloadBuilder(PayloadBase.values) result = builder.build() decoded = RowCoder(result.schema).decode(result.payload) for key, value in PayloadBase.values.items(): # Note the default value in the getattr call. # ImplicitSchemaPayloadBuilder omits fields with valu=None since their # type cannot be inferred. self.assertEqual(getattr(decoded, key, None), value) def test_implicit_payload_builder_with_bytes(self): values = PayloadBase.bytes_values builder = ImplicitSchemaPayloadBuilder(values) result = builder.build() decoded = RowCoder(result.schema).decode(result.payload) for key, value in PayloadBase.values.items(): # Note the default value in the getattr call. # ImplicitSchemaPayloadBuilder omits fields with valu=None since their # type cannot be inferred. self.assertEqual(getattr(decoded, key, None), value) # Verify we have not modified a cached type (BEAM-10766) # TODO(BEAM-7372): Remove when bytes coercion code is removed. self.assertEqual(typehints.List[bytes], convert_to_beam_type(list[bytes])) class ExternalTransformTest(unittest.TestCase): def test_pipeline_generation(self): pipeline = beam.Pipeline() _ = ( pipeline | beam.Create(['a', 'b']) | beam.ExternalTransform( 'beam:transforms:xlang:test:prefix', ImplicitSchemaPayloadBuilder({'data': '0'}), expansion_service.ExpansionServiceServicer())) proto, _ = pipeline.to_runner_api(return_context=True) pipeline_from_proto = Pipeline.from_runner_api( proto, pipeline.runner, pipeline._options) # Original pipeline has the un-expanded external transform self.assertEqual([], pipeline.transforms_stack[0].parts[1].parts) # new pipeline has the expanded external transform self.assertNotEqual([], pipeline_from_proto.transforms_stack[0].parts[1].parts) self.assertEqual( 'ExternalTransform(beam:transforms:xlang:test:prefix)/TestLabel', pipeline_from_proto.transforms_stack[0].parts[1].parts[0].full_label) @unittest.skipIf(apiclient is None, 'GCP dependencies are not installed') def test_pipeline_generation_with_runner_overrides(self): pipeline_properties = [ '--job_name=test-job', '--project=test-project', '--temp_location=gs://beam/tmp', '--no_auth', '--dry_run=True', '--sdk_location=container', '--runner=DataflowRunner', '--streaming', '--region=us-central1' ] with beam.Pipeline(options=PipelineOptions(pipeline_properties)) as p: _ = ( p | beam.io.ReadFromPubSub( subscription= 'projects/dummy-project/subscriptions/dummy-subscription') | beam.ExternalTransform( 'beam:transforms:xlang:test:prefix', ImplicitSchemaPayloadBuilder({'data': '0'}), expansion_service.ExpansionServiceServicer())) pipeline_proto, _ = p.to_runner_api(return_context=True) pubsub_read_transform = None external_transform = None proto_transforms = pipeline_proto.components.transforms for id in proto_transforms: if 'beam:transforms:xlang:test:prefix' in proto_transforms[ id].unique_name: external_transform = proto_transforms[id] if 'ReadFromPubSub' in proto_transforms[id].unique_name: pubsub_read_transform = proto_transforms[id] if not (pubsub_read_transform and external_transform): raise ValueError( 'Could not find an external transform and the PubSub read transform ' 'in the pipeline') self.assertEqual(1, len(list(pubsub_read_transform.outputs.values()))) self.assertEqual(1, len(list(external_transform.inputs.values()))) # Verify that the PubSub read transform output is connected to the # external transform input. Instead of comparing exact PCollection # reference IDs (which can be non-deterministic), we verify that both # transforms reference valid PCollections in the pipeline components pubsub_output_id = list(pubsub_read_transform.outputs.values())[0] external_input_id = list(external_transform.inputs.values())[0] # Both should reference valid PCollections in the pipeline components self.assertIn(pubsub_output_id, pipeline_proto.components.pcollections) self.assertIn(external_input_id, pipeline_proto.components.pcollections) # Verify that the pipeline structure is correct by checking that # we have exactly 2 PCollections total (the intermediate one between # the transforms, and the final output from external transform) total_pcollections = len(pipeline_proto.components.pcollections) self.assertGreaterEqual( total_pcollections, 1, "Pipeline should have at least 1 PCollection") def test_payload(self): with beam.Pipeline() as p: res = ( p | beam.Create(['a', 'bb'], reshuffle=False) | beam.ExternalTransform( 'payload', b's', expansion_service.ExpansionServiceServicer())) assert_that(res, equal_to(['as', 'bbs'])) def test_output_coder(self): external_transform = beam.ExternalTransform( 'map_to_union_types', None, expansion_service.ExpansionServiceServicer()).with_output_types(int) with beam.Pipeline() as p: res = (p | beam.Create([2, 2], reshuffle=False) | external_transform) assert_that(res, equal_to([2, 2])) context = pipeline_context.PipelineContext( external_transform._expanded_components) self.assertEqual(len(external_transform._expanded_transform.outputs), 1) for _, pcol_id in external_transform._expanded_transform.outputs.items(): pcol = context.pcollections.get_by_id(pcol_id) self.assertEqual(pcol.element_type, int) def test_no_output_coder(self): external_transform = beam.ExternalTransform( 'map_to_union_types', None, expansion_service.ExpansionServiceServicer()) with beam.Pipeline() as p: res = (p | beam.Create([2, 2], reshuffle=False) | external_transform) assert_that(res, equal_to([2, 2])) context = pipeline_context.PipelineContext( external_transform._expanded_components) self.assertEqual(len(external_transform._expanded_transform.outputs), 1) for _, pcol_id in external_transform._expanded_transform.outputs.items(): pcol = context.pcollections.get_by_id(pcol_id) self.assertEqual(pcol.element_type, typehints.Any) def test_nested(self): with beam.Pipeline() as p: assert_that(p | FibTransform(6), equal_to([8])) def test_external_empty_spec_translation(self): pipeline = beam.Pipeline() external_transform = beam.ExternalTransform( 'beam:transforms:xlang:test:prefix', ImplicitSchemaPayloadBuilder({'data': '0'}), expansion_service.ExpansionServiceServicer()) _ = (pipeline | beam.Create(['a', 'b']) | external_transform) pipeline.run().wait_until_finish() external_transform_label = ( 'ExternalTransform(beam:transforms:xlang:test:prefix)/TestLabel') for transform in external_transform._expanded_components.transforms.values( ): # We clear the spec of one of the external transforms. if transform.unique_name == external_transform_label: transform.spec.Clear() context = pipeline_context.PipelineContext() proto_pipeline = pipeline.to_runner_api(context=context) proto_transform = None for transform in proto_pipeline.components.transforms.values(): if (transform.unique_name == 'ExternalTransform(beam:transforms:xlang:test:prefix)/TestLabel'): proto_transform = transform self.assertIsNotNone(proto_transform) self.assertTrue(str(proto_transform).strip().find('spec {') == -1) def test_unique_name(self): p = beam.Pipeline() _ = p | FibTransform(6) proto = p.to_runner_api() xforms = [x.unique_name for x in proto.components.transforms.values()] self.assertEqual( len(set(xforms)), len(xforms), msg='Transform names are not unique.') pcolls = [x.unique_name for x in proto.components.pcollections.values()] self.assertEqual( len(set(pcolls)), len(pcolls), msg='PCollection names are not unique.') def test_external_transform_finder_non_leaf(self): pipeline = beam.Pipeline() _ = ( pipeline | beam.Create(['a', 'b']) | beam.ExternalTransform( 'beam:transforms:xlang:test:prefix', ImplicitSchemaPayloadBuilder({'data': '0'}), expansion_service.ExpansionServiceServicer()) | beam.Map(lambda x: x)) pipeline.run().wait_until_finish() self.assertTrue(pipeline.contains_external_transforms) def test_external_transform_finder_leaf(self): pipeline = beam.Pipeline() _ = ( pipeline | beam.Create(['a', 'b']) | beam.ExternalTransform( 'beam:transforms:xlang:test:nooutput', ImplicitSchemaPayloadBuilder({'data': '0'}), expansion_service.ExpansionServiceServicer())) pipeline.run().wait_until_finish() self.assertTrue(pipeline.contains_external_transforms) def test_sanitize_java_traceback(self): error_string = ''' java.lang.RuntimeException: ACTUAL \n MULTILINE \n ERROR \tat org.apache.beam.sdk.expansion.service.ExpansionService$TransformProviderForBuilder.getTransform(ExpansionService.java:308) \tat org.apache.beam.sdk.expansion.service.TransformProvider.apply(TransformProvider.java:121) \tat org.apache.beam.sdk.expansion.service.ExpansionService.expand(ExpansionService.java:627) \tat org.apache.beam.sdk.expansion.service.ExpansionService.expand(ExpansionService.java:729) \tat org.apache.beam.model.expansion.v1.ExpansionServiceGrpc$MethodHandlers.invoke(ExpansionServiceGrpc.java:306) \tat org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ServerCalls$UnaryServerCallHandler$UnaryServerCallListener.onHalfClose(ServerCalls.java:182) \tat org.apache.beam.vendor.grpc.v1p69p0.io.grpc.internal.ServerCallImpl$ServerStreamListenerImpl.halfClosed(ServerCallImpl.java:351) \tat org.apache.beam.vendor.grpc.v1p69p0.io.grpc.internal.ServerImpl$JumpToApplicationThreadServerStreamListener$1HalfClosed.runInContext(ServerImpl.java:861) \tat org.apache.beam.vendor.grpc.v1p69p0.io.grpc.internal.ContextRunnable.run(ContextRunnable.java:37) \tat org.apache.beam.vendor.grpc.v1p69p0.io.grpc.internal.SerializingExecutor.run(SerializingExecutor.java:133) \tat java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1149) \tat java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:624) \tat java.lang.Thread.run(Thread.java:748) Caused by: java.lang.IllegalArgumentException: Received unknown SQL Dialect 'X'. Known dialects: [calcite] \tat org.apache.beam.sdk.extensions.sql.expansion.ExternalSqlTransformRegistrar$Builder.buildExternal(ExternalSqlTransformRegistrar.java:73) \tat org.apache.beam.sdk.extensions.sql.expansion.ExternalSqlTransformRegistrar$Builder.buildExternal(ExternalSqlTransformRegistrar.java:63) \tat org.apache.beam.sdk.expansion.service.ExpansionService$TransformProviderForBuilder.getTransform(ExpansionService.java:303) \t... 12 more '''.strip() core_msg = 'java.lang.RuntimeException: ACTUAL \n MULTILINE \n ERROR' self.assertEqual( f"{error_string}\n\n{core_msg}", external._sanitize_java_traceback(error_string)) class ExternalAnnotationPayloadTest(PayloadBase, unittest.TestCase): def get_payload_from_typing_hints(self, values): class AnnotatedTransform(beam.ExternalTransform): URN = 'beam:external:fakeurn:v1' def __init__( self, integer_example: int, boolean: bool, string_example: str, list_of_strings: list[str], mapping: typing.Mapping[str, float], optional_integer: typing.Optional[int] = None, expansion_service=None): super().__init__( self.URN, AnnotationBasedPayloadBuilder( self, integer_example=integer_example, boolean=boolean, string_example=string_example, list_of_strings=list_of_strings, mapping=mapping, optional_integer=optional_integer, ), expansion_service) return get_payload(AnnotatedTransform(**values)) def get_payload_from_beam_typehints(self, values): class AnnotatedTransform(beam.ExternalTransform): URN = 'beam:external:fakeurn:v1' def __init__( self, integer_example: int, boolean: bool, string_example: str, list_of_strings: typehints.List[str], mapping: typehints.Dict[str, float], optional_integer: typehints.Optional[int] = None, expansion_service=None): super().__init__( self.URN, AnnotationBasedPayloadBuilder( self, integer_example=integer_example, boolean=boolean, string_example=string_example, list_of_strings=list_of_strings, mapping=mapping, optional_integer=optional_integer, ), expansion_service) return get_payload(AnnotatedTransform(**values)) class ExternalDataclassesPayloadTest(PayloadBase, unittest.TestCase): def get_payload_from_typing_hints(self, values): @dataclasses.dataclass class DataclassTransform(beam.ExternalTransform): URN = 'beam:external:fakeurn:v1' integer_example: int boolean: bool string_example: str list_of_strings: list[str] mapping: typing.Mapping[str, float] = dataclasses.field(default=dict) optional_integer: typing.Optional[int] = None expansion_service: dataclasses.InitVar[typing.Optional[str]] = None return get_payload(DataclassTransform(**values)) def get_payload_from_beam_typehints(self, values): @dataclasses.dataclass class DataclassTransform(beam.ExternalTransform): URN = 'beam:external:fakeurn:v1' integer_example: int boolean: bool string_example: str list_of_strings: typehints.List[str] mapping: typehints.Dict[str, float] = {} optional_integer: typehints.Optional[int] = None expansion_service: dataclasses.InitVar[typehints.Optional[str]] = None return get_payload(DataclassTransform(**values)) class SchemaTransformPayloadBuilderTest(unittest.TestCase): def test_build_payload(self): ComplexType = typing.NamedTuple( "ComplexType", [ ("str_sub_field", str), ("int_sub_field", int), ]) payload_builder = SchemaTransformPayloadBuilder( identifier='dummy_id', str_field='aaa', int_field=123, object_field=ComplexType(str_sub_field="bbb", int_sub_field=456)) payload_bytes = payload_builder.payload() payload_from_bytes = proto_utils.parse_Bytes( payload_bytes, external_transforms_pb2.SchemaTransformPayload) self.assertEqual('dummy_id', payload_from_bytes.identifier) expected_coder = RowCoder(payload_from_bytes.configuration_schema) schema_transform_config = expected_coder.decode( payload_from_bytes.configuration_row) self.assertEqual('aaa', schema_transform_config.str_field) self.assertEqual(123, schema_transform_config.int_field) self.assertEqual('bbb', schema_transform_config.object_field.str_sub_field) self.assertEqual(456, schema_transform_config.object_field.int_sub_field) class SchemaAwareExternalTransformTest(unittest.TestCase): class MockDiscoveryService: # define context manager enter and exit functions def __enter__(self): return self def __exit__(self, unusued1, unused2, unused3): pass def DiscoverSchemaTransform(self, unused_request=None): test_config = beam_expansion_api_pb2.SchemaTransformConfig( config_schema=schema_pb2.Schema( fields=[ schema_pb2.Field( name="str_field", type=schema_pb2.FieldType(atomic_type="STRING")), schema_pb2.Field( name="int_field", type=schema_pb2.FieldType(atomic_type="INT64")) ], id="test-id"), input_pcollection_names=["input"], output_pcollection_names=["output"]) test_managed_config = beam_expansion_api_pb2.SchemaTransformConfig( config_schema=schema_pb2.Schema( fields=[ schema_pb2.Field( name="transform_identifier", type=schema_pb2.FieldType(atomic_type="STRING")), schema_pb2.Field( name="config_url", type=schema_pb2.FieldType(atomic_type="STRING")), schema_pb2.Field( name="config", type=schema_pb2.FieldType(atomic_type="STRING")) ], id="test-id1"), input_pcollection_names=["input"], output_pcollection_names=["output"]) return beam_expansion_api_pb2.DiscoverSchemaTransformResponse( schema_transform_configs={ "test_schematransform": test_config, MANAGED_SCHEMA_TRANSFORM_IDENTIFIER: test_managed_config }) @mock.patch("apache_beam.transforms.external.ExternalTransform.service") def test_discover_one_config(self, mock_service): _mock = self.MockDiscoveryService() mock_service.return_value = _mock config = beam.SchemaAwareExternalTransform.discover_config( "test_service", name="test_schematransform") self.assertEqual(config.outputs[0], "output") self.assertEqual(config.inputs[0], "input") self.assertEqual(config.identifier, "test_schematransform") @mock.patch("apache_beam.transforms.external.ExternalTransform.service") def test_discover_one_config_fails_with_no_configs_found(self, mock_service): mock_service.return_value = self.MockDiscoveryService() with self.assertRaises(ValueError): beam.SchemaAwareExternalTransform.discover_config( "test_service", name="non_existent") @mock.patch("apache_beam.transforms.external.ExternalTransform.service") def test_rearrange_kwargs_based_on_discovery(self, mock_service): mock_service.return_value = self.MockDiscoveryService() identifier = "test_schematransform" expansion_service = "test_service" kwargs = {"int_field": 0, "str_field": "str"} transform = beam.SchemaAwareExternalTransform( identifier=identifier, expansion_service=expansion_service, rearrange_based_on_discovery=True, **kwargs) payload = transform._payload_builder.build() ordered_fields = [f.name for f in payload.configuration_schema.fields] schematransform_config = beam.SchemaAwareExternalTransform.discover_config( expansion_service, identifier) external_config_fields = schematransform_config.configuration_schema._fields self.assertNotEqual(tuple(kwargs.keys()), external_config_fields) self.assertEqual(tuple(ordered_fields), external_config_fields) @mock.patch("apache_beam.transforms.external.ExternalTransform.service") def test_managed_replacement_unknown_id(self, mock_service): mock_service.return_value = self.MockDiscoveryService() identifier = "test_schematransform" kwargs = {"int_field": 0, "str_field": "str"} managed_replacement = ManagedReplacement( underlying_transform_identifier="unknown_id", update_compatibility_version="2.50.0") with self.assertRaises(ValueError): beam.SchemaAwareExternalTransform( identifier=identifier, expansion_service=expansion_service, rearrange_based_on_discovery=True, managed_replacement=managed_replacement, **kwargs) @mock.patch("apache_beam.transforms.external.ExternalTransform.service") @mock.patch("apache_beam.transforms.external.BeamJarExpansionService") def test_managed_replacement_known_id( self, mock_service, mock_beam_jar_service): mock_service.return_value = self.MockDiscoveryService() mock_beam_jar_service.return_value = self.MockDiscoveryService() identifier = "test_schematransform" kwargs = {"int_field": 0, "str_field": "str"} managed_replacement = ManagedReplacement( underlying_transform_identifier=ManagedTransforms.Urns.ICEBERG_READ.urn, update_compatibility_version="2.50.0") external_transform = beam.SchemaAwareExternalTransform( identifier=identifier, expansion_service=expansion_service, rearrange_based_on_discovery=True, managed_replacement=managed_replacement, **kwargs) self.assertIsNotNone(external_transform._managed_payload_builder) class JavaClassLookupPayloadBuilderTest(unittest.TestCase): def _verify_row(self, schema, row_payload, expected_values): row = RowCoder(schema).decode(row_payload) for attr_name, expected_value in expected_values.items(): self.assertTrue(hasattr(row, attr_name)) value = getattr(row, attr_name) self.assertEqual(expected_value, value) def test_build_payload_with_constructor(self): payload_builder = JavaClassLookupPayloadBuilder('dummy_class_name') payload_builder.with_constructor('abc', 123, str_field='def', int_field=456) payload_bytes = payload_builder.payload() payload_from_bytes = proto_utils.parse_Bytes( payload_bytes, external_transforms_pb2.JavaClassLookupPayload) self.assertTrue( isinstance( payload_from_bytes, external_transforms_pb2.JavaClassLookupPayload)) self.assertFalse(payload_from_bytes.constructor_method) self._verify_row( payload_from_bytes.constructor_schema, payload_from_bytes.constructor_payload, { 'ignore0': 'abc', 'ignore1': 123, 'str_field': 'def', 'int_field': 456 }) def test_build_payload_with_constructor_method(self): payload_builder = JavaClassLookupPayloadBuilder('dummy_class_name') payload_builder.with_constructor_method( 'dummy_constructor_method', 'abc', 123, str_field='def', int_field=456) payload_bytes = payload_builder.payload() payload_from_bytes = proto_utils.parse_Bytes( payload_bytes, external_transforms_pb2.JavaClassLookupPayload) self.assertTrue( isinstance( payload_from_bytes, external_transforms_pb2.JavaClassLookupPayload)) self.assertEqual( 'dummy_constructor_method', payload_from_bytes.constructor_method) self._verify_row( payload_from_bytes.constructor_schema, payload_from_bytes.constructor_payload, { 'ignore0': 'abc', 'ignore1': 123, 'str_field': 'def', 'int_field': 456 }) def test_build_payload_with_builder_methods(self): payload_builder = JavaClassLookupPayloadBuilder('dummy_class_name') payload_builder.with_constructor('abc', 123, str_field='def', int_field=456) payload_builder.add_builder_method( 'builder_method1', 'abc1', 1234, str_field1='abc2', int_field1=2345) payload_builder.add_builder_method( 'builder_method2', 'abc3', 3456, str_field2='abc4', int_field2=4567) payload_bytes = payload_builder.payload() payload_from_bytes = proto_utils.parse_Bytes( payload_bytes, external_transforms_pb2.JavaClassLookupPayload) self.assertTrue( isinstance( payload_from_bytes, external_transforms_pb2.JavaClassLookupPayload)) self._verify_row( payload_from_bytes.constructor_schema, payload_from_bytes.constructor_payload, { 'ignore0': 'abc', 'ignore1': 123, 'str_field': 'def', 'int_field': 456 }) self.assertEqual(2, len(payload_from_bytes.builder_methods)) builder_method = payload_from_bytes.builder_methods[0] self.assertTrue( isinstance(builder_method, external_transforms_pb2.BuilderMethod)) self.assertEqual('builder_method1', builder_method.name) self._verify_row( builder_method.schema, builder_method.payload, { 'ignore0': 'abc1', 'ignore1': 1234, 'str_field1': 'abc2', 'int_field1': 2345 }) builder_method = payload_from_bytes.builder_methods[1] self.assertTrue( isinstance(builder_method, external_transforms_pb2.BuilderMethod)) self.assertEqual('builder_method2', builder_method.name) self._verify_row( builder_method.schema, builder_method.payload, { 'ignore0': 'abc3', 'ignore1': 3456, 'str_field2': 'abc4', 'int_field2': 4567 }) def test_build_payload_with_constructor_twice_fails(self): payload_builder = JavaClassLookupPayloadBuilder('dummy_class_name') payload_builder.with_constructor('abc') with self.assertRaises(ValueError): payload_builder.with_constructor('def') def test_implicit_builder_with_constructor(self): constructor_transform = ( JavaExternalTransform('org.pkg.MyTransform')('abc').withIntProperty(5)) payload_bytes = constructor_transform._payload_builder.payload() payload_from_bytes = proto_utils.parse_Bytes( payload_bytes, external_transforms_pb2.JavaClassLookupPayload) self.assertEqual('org.pkg.MyTransform', payload_from_bytes.class_name) self._verify_row( payload_from_bytes.constructor_schema, payload_from_bytes.constructor_payload, {'ignore0': 'abc'}) builder_method = payload_from_bytes.builder_methods[0] self.assertEqual('withIntProperty', builder_method.name) self._verify_row( builder_method.schema, builder_method.payload, {'ignore0': 5}) def test_implicit_builder_with_constructor_method(self): constructor_transform = JavaExternalTransform('org.pkg.MyTransform').of( str_field='abc').withProperty(int_field=1234).build() payload_bytes = constructor_transform._payload_builder.payload() payload_from_bytes = proto_utils.parse_Bytes( payload_bytes, external_transforms_pb2.JavaClassLookupPayload) self.assertEqual('of', payload_from_bytes.constructor_method) self._verify_row( payload_from_bytes.constructor_schema, payload_from_bytes.constructor_payload, {'str_field': 'abc'}) with_property_method = payload_from_bytes.builder_methods[0] self.assertEqual('withProperty', with_property_method.name) self._verify_row( with_property_method.schema, with_property_method.payload, {'int_field': 1234}) build_method = payload_from_bytes.builder_methods[1] self.assertEqual('build', build_method.name) self._verify_row(build_method.schema, build_method.payload, {}) class JavaJarExpansionServiceTest(unittest.TestCase): def setUp(self): SubprocessServer._cache._live_owners = set() def test_classpath(self): with tempfile.TemporaryDirectory() as temp_dir: try: # Avoid having to prefix everything in our test strings. oldwd = os.getcwd() os.chdir(temp_dir) # Touch some files for globing. with open('a1.jar', 'w') as _: pass service = JavaJarExpansionService( 'main.jar', classpath=['a*.jar', 'b.jar']) self.assertEqual( service._default_args(), ['{{PORT}}', '--filesToStage=main.jar,a1.jar,b.jar']) finally: os.chdir(oldwd) @mock.patch.object(JavaJarServer, 'local_jar') def test_classpath_with_url(self, local_jar): def _side_effect_fn(path): return path[path.rindex('/') + 1:] local_jar.side_effect = _side_effect_fn with tempfile.TemporaryDirectory() as temp_dir: try: # Avoid having to prefix everything in our test strings. oldwd = os.getcwd() os.chdir(temp_dir) service = JavaJarExpansionService( 'main.jar', classpath=['https://dummy_path/dummyjar.jar']) self.assertEqual( service._default_args(), ['{{PORT}}', '--filesToStage=main.jar,dummyjar.jar']) finally: os.chdir(oldwd) @mock.patch.object(JavaJarServer, 'local_jar') def test_classpath_with_gradle_artifact(self, local_jar): def _side_effect_fn(path, user_agent=None): return path[path.rindex('/') + 1:] local_jar.side_effect = _side_effect_fn with tempfile.TemporaryDirectory() as temp_dir: try: # Avoid having to prefix everything in our test strings. oldwd = os.getcwd() os.chdir(temp_dir) service = JavaJarExpansionService( 'main.jar', classpath=['dummy_group:dummy_artifact:dummy_version']) self.assertEqual( service._default_args(), [ '{{PORT}}', '--filesToStage=main.jar,dummy_artifact-dummy_version.jar' ]) finally: os.chdir(oldwd) def test_classpath_with_glob(self): with tempfile.TemporaryDirectory() as temp_dir: try: # Avoid having to prefix everything in our test strings. oldwd = os.getcwd() os.chdir(temp_dir) # Touch some files for globing. with open('a1.jar', 'w') as _: pass service = JavaJarExpansionService( 'main.jar', classpath=['a*.jar', 'b.jar']) self.assertEqual( service._default_args(), ['{{PORT}}', '--filesToStage=main.jar,a1.jar,b.jar']) finally: os.chdir(oldwd) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/external_transform_provider.py0000644000175100017510000002615515206113033026266 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import re from collections import namedtuple from inspect import Parameter from inspect import Signature from apache_beam.transforms import PTransform from apache_beam.transforms.external import BeamJarExpansionService from apache_beam.transforms.external import JavaJarExpansionService from apache_beam.transforms.external import SchemaAwareExternalTransform from apache_beam.transforms.external import SchemaTransformsConfig from apache_beam.typehints.schemas import named_tuple_to_schema from apache_beam.typehints.schemas import typing_from_runner_api __all__ = ['ExternalTransform', 'ExternalTransformProvider'] def snake_case_to_upper_camel_case(string): """Convert snake_case to UpperCamelCase""" components = string.split('_') output = ''.join(n.capitalize() for n in components) return output # Information regarding a Wrapper parameter. ParamInfo = namedtuple('ParamInfo', ['type', 'description', 'original_name']) def get_config_with_descriptions( schematransform: SchemaTransformsConfig) -> dict[str, ParamInfo]: # Prepare a configuration schema that includes types and descriptions schema = named_tuple_to_schema(schematransform.configuration_schema) descriptions = schematransform.configuration_schema._field_descriptions fields_with_descriptions = {} for field in schema.fields: fields_with_descriptions[field.name] = ParamInfo( typing_from_runner_api(field.type), descriptions[field.name], field.name) return fields_with_descriptions def _generate_signature(schematransform: SchemaTransformsConfig) -> Signature: schema = named_tuple_to_schema(schematransform.configuration_schema) descriptions = schematransform.configuration_schema._field_descriptions params: list[Parameter] = [] for field in schema.fields: annotation = str(typing_from_runner_api(field.type)) description = descriptions[field.name] if description: annotation = annotation + f": {description}" params.append( Parameter( field.name, Parameter.POSITIONAL_OR_KEYWORD, annotation=annotation)) return Signature(params) class ExternalTransform(PTransform): """Template for a wrapper class of an external SchemaTransform This is a superclass for dynamically generated SchemaTransform wrappers and is not meant to be manually instantiated. Experimental; no backwards compatibility guarantees.""" # These attributes need to be set when # creating an ExternalTransform type default_expansion_service = None identifier: str = "" configuration_schema: dict[str, ParamInfo] = {} def __init__(self, expansion_service=None, **kwargs): self._kwargs = kwargs self._expansion_service = \ expansion_service or self.default_expansion_service def expand(self, input): external_schematransform = SchemaAwareExternalTransform( identifier=self.identifier, expansion_service=self._expansion_service, rearrange_based_on_discovery=True, **self._kwargs) return input | external_schematransform STANDARD_URN_PATTERN = r"^beam:schematransform:org.apache.beam:([\w-]+):(\w+)$" def infer_name_from_identifier(identifier: str, pattern: str): """Infer a class name from an identifier, adhering to the input pattern""" match = re.match(pattern, identifier) if not match: return None groups = match.groups() components = [snake_case_to_upper_camel_case(n) for n in groups] # Special handling for standard SchemaTransform identifiers: # We don't include the version number if it's the first version if (pattern == STANDARD_URN_PATTERN and components[1].lower() == 'v1'): return components[0] else: return ''.join(components) class ExternalTransformProvider: """Dynamically discovers Schema-aware external transforms from a given list of expansion services and provides them as ready PTransforms. A :class:`ExternalTransform` subclass is generated for each external transform, and is named based on what can be inferred from the URN (see the `urn_pattern` parameter). These classes are generated when :class:`ExternalTransformProvider` is initialized. You can give it an expansion service address that is already up and running: >>> provider = ExternalTransformProvider("localhost:12345") Or you can give it the path to an expansion service Jar file: >>> provider = ExternalTransformProvider(JavaJarExpansionService( "path/to/expansion-service.jar")) Or you can give it the gradle target of a standard Beam expansion service: >>> provider = ExternalTransformProvider(BeamJarExpansionService( "sdks:java:io:google-cloud-platform:expansion-service:shadowJar")) Note that you can provide a list of these services: >>> provider = ExternalTransformProvider([ "localhost:12345", JavaJarExpansionService("path/to/expansion-service.jar"), BeamJarExpansionService( "sdks:java:io:google-cloud-platform:expansion-service:shadowJar")]) The output of :func:`get_available()` provides a list of available transforms in the provided expansion service(s): >>> provider.get_available() [('JdbcWrite', 'beam:schematransform:org.apache.beam:jdbc_write:v1'), ('BigtableRead', 'beam:schematransform:org.apache.beam:bigtable_read:v1'), ...] You can retrieve a transform with :func:`get()`, :func:`get_urn()`, or by directly accessing it as an attribute. The following lines all do the same thing: >>> provider.get('BigqueryStorageRead') >>> provider.get_urn( 'beam:schematransform:org.apache.beam:bigquery_storage_read:v1') >>> provider.BigqueryStorageRead You can inspect the transform's documentation for more details. The following returns the documentation provided by the underlying SchemaTransform. If no such documentation is provided, this will be empty. >>> import inspect >>> inspect.getdoc(provider.BigqueryStorageRead) Similarly, you can inspect the transform's signature to know more about its parameters, including their names, types, and any documentation that the underlying SchemaTransform may provide: >>> inspect.signature(provider.BigqueryStorageRead) (query: 'typing.Union[str, NoneType]: The SQL query to be executed to...', row_restriction: 'typing.Union[str, NoneType]: Read only rows that match...', selected_fields: 'typing.Union[typing.Sequence[str], NoneType]: Read ...', table_spec: 'typing.Union[str, NoneType]: The fully-qualified name of ...') The retrieved external transform can be used as a normal PTransform like so:: with Pipeline() as p: _ = (p | 'Read from BigQuery` >> provider.BigqueryStorageRead( query=query, row_restriction=restriction) | 'Some processing' >> beam.Map(...)) """ def __init__(self, expansion_services, urn_pattern=STANDARD_URN_PATTERN): f"""Initialize an ExternalTransformProvider :param expansion_services: A list of expansion services to discover transforms from. Supported forms: * a string representing the expansion service address * a :attr:`JavaJarExpansionService` pointing to the path of a Java Jar * a :attr:`BeamJarExpansionService` pointing to a gradle target :param urn_pattern: The regular expression used to match valid transforms. In addition to validating, the captured groups are used to infer a name for each class. By default, the following pattern is used: [{STANDARD_URN_PATTERN}] """ self._urn_pattern = urn_pattern self._transforms: dict[str, type(ExternalTransform)] = {} self._name_to_urn: dict[str, str] = {} if isinstance(expansion_services, set): expansion_services = list(expansion_services) if not isinstance(expansion_services, list): expansion_services = [expansion_services] self.expansion_services = expansion_services self._create_wrappers() def _create_wrappers(self): # multiple services can overlap and include the same URNs. If this happens, # we prioritize by the order of services in the list identifiers = set() for service in self.expansion_services: target = service if isinstance(service, BeamJarExpansionService): target = service.gradle_target if isinstance(service, JavaJarExpansionService): target = service.path_to_jar try: schematransform_configs = SchemaAwareExternalTransform.discover(service) except Exception as e: logging.exception( "Encountered an error while discovering " "expansion service at '%s':\n%s", target, e) continue skipped_urns = [] for config in schematransform_configs: identifier = config.identifier if identifier not in identifiers: identifiers.add(identifier) name = infer_name_from_identifier(identifier, self._urn_pattern) if name is None: skipped_urns.append(identifier) continue transform = type( name, (ExternalTransform, ), dict( identifier=identifier, default_expansion_service=service, schematransform=config, # configuration_schema is used by the auto-wrapper generator configuration_schema=get_config_with_descriptions(config))) transform.__doc__ = config.description transform.__signature__ = _generate_signature(config) self._transforms[identifier] = transform self._name_to_urn[name] = identifier if skipped_urns: logging.info( "Skipped URN(s) in '%s' that don't follow the pattern \"%s\": %s", target, self._urn_pattern, skipped_urns) for transform in self._transforms.values(): setattr(self, transform.__name__, transform) def get_available(self) -> list[tuple[str, str]]: """Get a list of available ExternalTransform names and identifiers""" return list(self._name_to_urn.items()) def get_all(self) -> dict[str, ExternalTransform]: """Get all ExternalTransforms""" return self._transforms def get(self, name) -> ExternalTransform: """Get an ExternalTransform by its inferred class name""" return self._transforms[self._name_to_urn[name]] def get_urn(self, identifier) -> ExternalTransform: """Get an ExternalTransform by its SchemaTransform identifier""" return self._transforms[identifier] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/external_transform_provider_it_test.py0000644000175100017510000003754415206113033030025 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import importlib import inspect import logging import os import secrets import shutil import time import typing import unittest from os.path import dirname import numpy import pytest import yaml import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.external import BeamJarExpansionService from apache_beam.transforms.external_transform_provider import STANDARD_URN_PATTERN from apache_beam.transforms.external_transform_provider import ExternalTransform from apache_beam.transforms.external_transform_provider import ExternalTransformProvider from apache_beam.transforms.external_transform_provider import infer_name_from_identifier from apache_beam.transforms.external_transform_provider import snake_case_to_upper_camel_case from apache_beam.transforms.xlang.io import GenerateSequence class NameAndTypeUtilsTest(unittest.TestCase): def test_snake_case_to_upper_camel_case(self): test_cases = [("", ""), ("test", "Test"), ("test_name", "TestName"), ("test_double_underscore", "TestDoubleUnderscore"), ("TEST_CAPITALIZED", "TestCapitalized"), ("_prepended_underscore", "PrependedUnderscore"), ("appended_underscore_", "AppendedUnderscore")] for case in test_cases: self.assertEqual(case[1], snake_case_to_upper_camel_case(case[0])) def test_infer_name_from_identifier(self): standard_test_cases = [ ("beam:schematransform:org.apache.beam:transform:v1", "Transform"), ("beam:schematransform:org.apache.beam:my_transform:v1", "MyTransform"), ( "beam:schematransform:org.apache.beam:my_transform:v2", "MyTransformV2"), ("beam:schematransform:org.apache.beam:fe_fi_fo_fum:v2", "FeFiFoFumV2"), ("beam:schematransform:bad_match:my_transform:v1", None) ] for case in standard_test_cases: self.assertEqual( case[1], infer_name_from_identifier(case[0], STANDARD_URN_PATTERN)) custom_pattern_cases = [ # (, , ) ( r"^custom:transform:([\w-]+):(\w+)$", "custom:transform:my_transform:v1", "MyTransformV1"), ( r"^org.user:([\w-]+):([\w-]+):([\w-]+):external$", "org.user:some:custom_transform:we_made:external", "SomeCustomTransformWeMade"), ( r"^([\w-]+):user.transforms", "my_eXTErnal:user.transforms", "MyExternal"), (r"^([\w-]+):user.transforms", "my_external:badinput.transforms", None), ] for case in custom_pattern_cases: self.assertEqual(case[2], infer_name_from_identifier(case[1], case[0])) @pytest.mark.uses_io_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_JARS'), "EXPANSION_JARS environment var is not provided, " "indicating that jars have not been built") class ExternalTransformProviderIT(unittest.TestCase): def test_generate_sequence_signature_and_doc(self): provider = ExternalTransformProvider( BeamJarExpansionService(":sdks:java:io:expansion-service:shadowJar")) self.assertTrue(( 'GenerateSequence', 'beam:schematransform:org.apache.beam:generate_sequence:v1' ) in provider.get_available()) GenerateSequence = provider.get('GenerateSequence') signature = inspect.signature(GenerateSequence) for param in ['start', 'end', 'rate']: self.assertTrue(param in signature.parameters.keys()) doc_substring = ( "Outputs a PCollection of Beam Rows, each " "containing a single INT64") self.assertTrue(doc_substring in inspect.getdoc(GenerateSequence)) def test_run_generate_sequence(self): provider = ExternalTransformProvider( BeamJarExpansionService(":sdks:java:io:expansion-service:shadowJar")) with beam.Pipeline() as p: numbers = p | provider.GenerateSequence( start=0, end=10) | beam.Map(lambda row: row.value) assert_that(numbers, equal_to([i for i in range(10)])) def test_run_generate_sequence_with_rate(self): provider = ExternalTransformProvider( BeamJarExpansionService(":sdks:java:io:expansion-service:shadowJar")) with beam.Pipeline() as p: numbers = p | provider.GenerateSequence( start=0, end=3, rate={ 'elements': 1, 'seconds': 1 }) | beam.Map(lambda row: row.value) assert_that(numbers, equal_to([0, 1, 2])) @pytest.mark.xlang_wrapper_generation @unittest.skipUnless( os.environ.get('EXPANSION_JARS'), "EXPANSION_JARS environment var is not provided, " "indicating that jars have not been built") class AutoGenerationScriptIT(unittest.TestCase): """ This class tests the generation and regeneration operations in `gen_xlang_wrappers.py`. """ # tests cases will use GenerateSequence GEN_SEQ_IDENTIFIER = \ 'beam:schematransform:org.apache.beam:generate_sequence:v1' def setUp(self): # import script from top-level sdks/python directory self.sdk_dir = os.path.abspath(dirname(dirname(dirname(__file__)))) xlang_spec = importlib.util.spec_from_file_location( 'gen_xlang_wrappers', os.path.join(self.sdk_dir, 'gen_xlang_wrappers.py')) self.xlang_script = importlib.util.module_from_spec(xlang_spec) xlang_spec.loader.exec_module(self.xlang_script) args = TestPipeline(is_integration_test=True).get_full_options_as_args() runner = PipelineOptions(args).get_all_options()['runner'] if runner and "direct" not in runner.lower(): self.skipTest( "It is sufficient to run this test in the DirectRunner " "test suite only.") self.test_dir_name = 'test_gen_script_%d_%s' % ( int(time.time()), secrets.token_hex(3)) self.test_dir = os.path.join( os.path.abspath(dirname(__file__)), self.test_dir_name) self.service_config_path = os.path.join( self.test_dir, "test_expansion_service_config.yaml") self.transform_config_path = os.path.join( self.test_dir, "test_transform_config.yaml") os.mkdir(self.test_dir) def tearDown(self): shutil.rmtree(self.test_dir, ignore_errors=False) def delete_and_validate(self): self.xlang_script.delete_generated_files(self.test_dir) self.assertEqual(len(os.listdir(self.test_dir)), 0) def test_script_fails_with_invalid_destinations(self): expansion_service_config = { "gradle_target": 'sdks:java:io:expansion-service:shadowJar', 'destinations': { 'python': 'apache_beam/some_nonexistent_dir' } } with self.assertRaises(ValueError): self.create_and_check_transforms_config_exists(expansion_service_config) def test_pretty_types(self): types = [ typing.Optional[list[str]], numpy.int16, str, dict[str, numpy.float64], typing.Optional[dict[str, list[numpy.int64]]], dict[int, typing.Optional[str]] ] expected_type_names = [('list[str]', True), ('int16', False), ('str', False), ('dict[str, float64]', False), ('dict[str, list[int64]]', True), ('dict[int, Optional[str]]', False)] for i in range(len(types)): self.assertEqual( self.xlang_script.pretty_type(types[i]), expected_type_names[i]) def create_and_check_transforms_config_exists(self, expansion_service_config): with open(self.service_config_path, 'w') as f: yaml.dump([expansion_service_config], f) self.xlang_script.generate_transforms_config( self.service_config_path, self.transform_config_path) self.assertTrue(os.path.exists(self.transform_config_path)) def create_and_validate_transforms_config( self, expansion_service_config, expected_name, expected_destination): self.create_and_check_transforms_config_exists(expansion_service_config) with open(self.transform_config_path) as f: configs = yaml.safe_load(f) gen_seq_config = None for config in configs: if config['identifier'] == self.GEN_SEQ_IDENTIFIER: gen_seq_config = config self.assertIsNotNone(gen_seq_config) self.assertEqual( gen_seq_config['default_service'], expansion_service_config['gradle_target']) self.assertEqual(gen_seq_config['name'], expected_name) self.assertEqual( gen_seq_config['destinations']['python'], expected_destination) field_names = [field['name'] for field in gen_seq_config['fields']] self.assertIn("end", field_names) self.assertIn("start", field_names) self.assertIn("rate", field_names) def get_module(self, dest): module_name = dest.replace('apache_beam/', '').replace('/', '_') module = 'apache_beam.transforms.%s.%s' % (self.test_dir_name, module_name) return importlib.import_module(module) def write_wrappers_to_destinations_and_validate( self, destinations: list[str]): """ Generate wrappers from the config path and validate all destinations are included. Then write wrappers to destinations and validate all destination paths exist. :return: Generated wrappers grouped by destination """ grouped_wrappers = self.xlang_script.get_wrappers_from_transform_configs( self.transform_config_path) for dest in destinations: self.assertIn(dest, grouped_wrappers) # write to our test directory to avoid messing with other files self.xlang_script.write_wrappers_to_destinations( grouped_wrappers, self.test_dir, format_code=False) for dest in destinations: self.assertTrue( os.path.exists( os.path.join( self.test_dir, dest.replace('apache_beam/', '').replace('/', '_') + ".py"))) return grouped_wrappers def test_script_workflow(self): expected_destination = 'apache_beam/transforms' expansion_service_config = { "gradle_target": 'sdks:java:io:expansion-service:shadowJar', 'destinations': { 'python': expected_destination } } self.create_and_validate_transforms_config( expansion_service_config, 'GenerateSequence', expected_destination) grouped_wrappers = self.write_wrappers_to_destinations_and_validate( [expected_destination]) # at least the GenerateSequence wrapper is set to this destination self.assertGreaterEqual(len(grouped_wrappers[expected_destination]), 1) # check the wrapper exists in this destination and has correct properties output_module = self.get_module(expected_destination) self.assertTrue(hasattr(output_module, 'GenerateSequence')) # Since our config isn't skipping any transforms, # it should include these two Kafka IOs as well self.assertTrue(hasattr(output_module, 'KafkaWrite')) self.assertTrue(hasattr(output_module, 'KafkaRead')) self.assertTrue( isinstance(output_module.GenerateSequence(start=0), ExternalTransform)) self.assertEqual( output_module.GenerateSequence.identifier, self.GEN_SEQ_IDENTIFIER) self.delete_and_validate() def test_script_workflow_with_modified_transforms(self): modified_name = 'ModifiedSequence' modified_dest = 'apache_beam/io/gcp' expansion_service_config = { "gradle_target": 'sdks:java:io:expansion-service:shadowJar', 'destinations': { 'python': 'apache_beam/transforms' }, 'transforms': { 'beam:schematransform:org.apache.beam:generate_sequence:v1': { 'name': modified_name, 'destinations': { 'python': modified_dest } } } } self.create_and_validate_transforms_config( expansion_service_config, modified_name, modified_dest) grouped_wrappers = self.write_wrappers_to_destinations_and_validate( [modified_dest]) self.assertIn(modified_name, grouped_wrappers[modified_dest][0]) self.assertEqual(len(grouped_wrappers[modified_dest]), 1) # check the modified wrapper exists in the modified destination # and check it has the correct properties output_module = self.get_module(modified_dest) self.assertTrue( isinstance(output_module.ModifiedSequence(start=0), ExternalTransform)) self.assertEqual( output_module.ModifiedSequence.identifier, self.GEN_SEQ_IDENTIFIER) self.delete_and_validate() def test_script_workflow_with_skipped_transform(self): expansion_service_config = { "gradle_target": 'sdks:java:io:expansion-service:shadowJar', 'destinations': { 'python': f'apache_beam/transforms/{self.test_dir_name}' }, 'skip_transforms': [ 'beam:schematransform:org.apache.beam:generate_sequence:v1' ] } with open(self.service_config_path, 'w') as f: yaml.dump([expansion_service_config], f) self.xlang_script.generate_transforms_config( self.service_config_path, self.transform_config_path) # gen sequence shouldn't exist in the transform config with open(self.transform_config_path) as f: transforms = yaml.safe_load(f) gen_seq_config = None for transform in transforms: if transform['identifier'] == self.GEN_SEQ_IDENTIFIER: gen_seq_config = transform self.assertIsNone(gen_seq_config) def test_run_pipeline_with_generated_transform(self): with beam.Pipeline() as p: numbers = ( p | GenerateSequence(start=0, end=10) | beam.Map(lambda row: row.value)) assert_that(numbers, equal_to([i for i in range(10)])) def test_check_standard_external_transforms_config_in_sync(self): """ This test creates a transforms config file and checks it against `sdks/standard_external_transforms.yaml`. Fails if the two configs don't match. Fix by running `./gradlew generateExternalTransformsConfig` and committing the changes. """ sdks_dir = os.path.abspath(dirname(self.sdk_dir)) self.xlang_script.generate_transforms_config( os.path.join(sdks_dir, 'standard_expansion_services.yaml'), self.transform_config_path) with open(self.transform_config_path) as f: test_config = yaml.safe_load(f) with open(os.path.join(sdks_dir, 'standard_external_transforms.yaml'), 'r') as f: standard_config = yaml.safe_load(f) def _normalize(cfg): # Sort each transform's fields deterministically for t in cfg: if "fields" in t and isinstance(t["fields"], list): t["fields"] = sorted(t["fields"], key=lambda f: f.get("name", "")) # Sort transforms deterministically return sorted(cfg, key=lambda t: t.get("identifier", "")) self.assertEqual( _normalize(test_config), _normalize(standard_config), "The standard xlang transforms config file " "\"standard_external_transforms.yaml\" is out of sync! Please update " "by running './gradlew generateExternalTransformsConfig' " "and committing the changes.") if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/fully_qualified_named_transform.py0000644000175100017510000001163015206113033027044 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import contextlib import fnmatch import importlib from apache_beam import coders from apache_beam.portability.api import external_transforms_pb2 from apache_beam.pvalue import Row from apache_beam.transforms import ptransform from apache_beam.typehints.native_type_compatibility import convert_to_python_type from apache_beam.typehints.schemas import named_fields_to_schema from apache_beam.typehints.trivial_inference import instance_to_type from apache_beam.utils import python_callable PYTHON_FULLY_QUALIFIED_NAMED_TRANSFORM_URN = ( 'beam:transforms:python:fully_qualified_named') @ptransform.PTransform.register_urn( PYTHON_FULLY_QUALIFIED_NAMED_TRANSFORM_URN, external_transforms_pb2.ExternalConfigurationPayload) class FullyQualifiedNamedTransform(ptransform.PTransform): _FILTER_GLOB = None @classmethod @contextlib.contextmanager def with_filter(cls, filter): old_filter, cls._FILTER_GLOB = cls._FILTER_GLOB, filter try: yield finally: cls._FILTER_GLOB = old_filter def __init__(self, constructor, args, kwargs): self._constructor = constructor self._args = args self._kwargs = kwargs def expand(self, pinput): if self._constructor in ('__callable__', '__constructor__'): self._check_allowed(self._constructor) if self._args: source, *args = tuple(self._args) kwargs = self._kwargs else: args = self._args kwargs = dict(self._kwargs) source = kwargs.pop('source') if isinstance(source, str): source = python_callable.PythonCallableWithSource(source) if self._constructor == '__constructor__': transform = source(*args, **kwargs) else: transform = ptransform._PTransformFnPTransform(source, *args, **kwargs) else: transform = self._resolve(self._constructor)(*self._args, **self._kwargs) return pinput | transform @classmethod def _check_allowed(cls, fully_qualified_name): if not cls._FILTER_GLOB or not fnmatch.fnmatchcase(fully_qualified_name, cls._FILTER_GLOB): raise ValueError( f'Fully qualifed name "{fully_qualified_name}" ' f'not allowed by filter {cls._FILTER_GLOB}.') @classmethod def _resolve(cls, fully_qualified_name): cls._check_allowed(fully_qualified_name) o = None path = '' for segment in fully_qualified_name.split('.'): path = '.'.join([path, segment]) if path else segment if o is not None and hasattr(o, segment): o = getattr(o, segment) else: o = importlib.import_module(path) return o def to_runner_api_parameter(self, unused_context): _args_schema = named_fields_to_schema([ (f'arg{ix}', convert_to_python_type(instance_to_type(value))) for (ix, value) in enumerate(self._args) ]) _kwargs_schema = named_fields_to_schema([ (key, convert_to_python_type(instance_to_type(value))) for (key, value) in self._kwargs.items() ]) payload_schema = named_fields_to_schema({ 'constructor': str, 'args': _args_schema, 'kwargs': _kwargs_schema, }) return ( PYTHON_FULLY_QUALIFIED_NAMED_TRANSFORM_URN, external_transforms_pb2.ExternalConfigurationPayload( schema=payload_schema, payload=coders.RowCoder(payload_schema).encode( Row( constructor=self._constructor, args=Row( **{ f'arg{ix}': arg for (ix, arg) in enumerate(self._args) }), kwargs=Row(**self._kwargs)), ))) @staticmethod def from_runner_api_parameter(unused_ptransform, payload, unused_context): row = coders.RowCoder(payload.schema).decode(payload.payload) def maybe_as_dict(x): if isinstance(x, dict): return x elif x: return x._asdict() else: return {} return FullyQualifiedNamedTransform( row.constructor, tuple(getattr(row, 'args', ())), maybe_as_dict(getattr(row, 'kwargs', None))) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/fully_qualified_named_transform_test.py0000644000175100017510000002004415206113033030102 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest from mock import call from mock import patch import apache_beam as beam from apache_beam.runners.portability import expansion_service from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.external import ImplicitSchemaPayloadBuilder from apache_beam.transforms.fully_qualified_named_transform import PYTHON_FULLY_QUALIFIED_NAMED_TRANSFORM_URN from apache_beam.transforms.fully_qualified_named_transform import FullyQualifiedNamedTransform from apache_beam.utils import python_callable class FullyQualifiedNamedTransformTest(unittest.TestCase): def test_test_transform(self): with beam.Pipeline() as p: assert_that( p | beam.Create(['a', 'b', 'c']) | _TestTransform('x', 'y'), equal_to(['xay', 'xby', 'xcy'])) def test_expand(self): with FullyQualifiedNamedTransform.with_filter('*'): with beam.Pipeline() as p: assert_that( p | beam.Create(['a', 'b', 'c']) | FullyQualifiedNamedTransform( 'apache_beam.transforms.fully_qualified_named_transform_test' '._TestTransform', ('x', ), {'suffix': 'y'}), equal_to(['xay', 'xby', 'xcy'])) def test_static_constructor(self): with FullyQualifiedNamedTransform.with_filter('*'): with beam.Pipeline() as p: assert_that( p | beam.Create(['a', 'b', 'c']) | FullyQualifiedNamedTransform( 'apache_beam.transforms.fully_qualified_named_transform_test' '._TestTransform.create', ('x', ), {'suffix': 'y'}), equal_to(['xay', 'xby', 'xcy'])) def test_as_external_transform(self): with FullyQualifiedNamedTransform.with_filter('*'): with beam.Pipeline() as p: assert_that( p | beam.Create(['a', 'b', 'c']) | beam.ExternalTransform( PYTHON_FULLY_QUALIFIED_NAMED_TRANSFORM_URN, ImplicitSchemaPayloadBuilder({ 'constructor': 'apache_beam.transforms' '.fully_qualified_named_transform_test._TestTransform', 'args': beam.Row(arg0='x'), 'kwargs': beam.Row(suffix='y'), }), expansion_service.ExpansionServiceServicer()), equal_to(['xay', 'xby', 'xcy'])) def test_as_external_transform_no_args(self): with FullyQualifiedNamedTransform.with_filter('*'): with beam.Pipeline() as p: assert_that( p | beam.Create(['a', 'b', 'c']) | beam.ExternalTransform( PYTHON_FULLY_QUALIFIED_NAMED_TRANSFORM_URN, ImplicitSchemaPayloadBuilder({ 'constructor': 'apache_beam.transforms' '.fully_qualified_named_transform_test._TestTransform', 'kwargs': beam.Row(prefix='x', suffix='y'), }), expansion_service.ExpansionServiceServicer()), equal_to(['xay', 'xby', 'xcy'])) def test_as_external_transform_no_kwargs(self): with FullyQualifiedNamedTransform.with_filter('*'): with beam.Pipeline() as p: assert_that( p | beam.Create(['a', 'b', 'c']) | beam.ExternalTransform( PYTHON_FULLY_QUALIFIED_NAMED_TRANSFORM_URN, ImplicitSchemaPayloadBuilder({ 'constructor': 'apache_beam.transforms' '.fully_qualified_named_transform_test._TestTransform', 'args': beam.Row(arg0='x', arg1='y'), }), expansion_service.ExpansionServiceServicer()), equal_to(['xay', 'xby', 'xcy'])) def test_callable_transform(self): with FullyQualifiedNamedTransform.with_filter('*'): with beam.Pipeline() as p: assert_that( p | beam.Create(['a', 'b', 'c']) | FullyQualifiedNamedTransform( '__callable__', # the next argument is a callable to be applied ( python_callable.PythonCallableWithSource( """ def func(pcoll, x): return pcoll | beam.Map(lambda e: e + x) """), 'x' # arguments passed to the callable ), {}), equal_to(['ax', 'bx', 'cx'])) def test_constructor_transform(self): with FullyQualifiedNamedTransform.with_filter('*'): with beam.Pipeline() as p: assert_that( p | beam.Create(['a', 'b', 'c']) | FullyQualifiedNamedTransform( '__constructor__', # the next argument constructs a PTransform (), { 'source': python_callable.PythonCallableWithSource( """ class MyTransform(beam.PTransform): def __init__(self, x): self._x = x def expand(self, pcoll): return pcoll | beam.Map(lambda e: e + self._x) """), 'x': 'x' # arguments passed to the above constructor }), equal_to(['ax', 'bx', 'cx'])) def test_glob_filter(self): with FullyQualifiedNamedTransform.with_filter('*'): self.assertIs( FullyQualifiedNamedTransform._resolve('apache_beam.Row'), beam.Row) with FullyQualifiedNamedTransform.with_filter('apache_beam.*'): self.assertIs( FullyQualifiedNamedTransform._resolve('apache_beam.Row'), beam.Row) with FullyQualifiedNamedTransform.with_filter('apache_beam.foo.*'): with self.assertRaises(ValueError): FullyQualifiedNamedTransform._resolve('apache_beam.Row') @patch('importlib.import_module') def test_resolve_by_path_segment(self, mock_import_module): mock_import_module.return_value = None with FullyQualifiedNamedTransform.with_filter('*'): FullyQualifiedNamedTransform._resolve('a.b.c.d') mock_import_module.assert_has_calls( [call('a'), call('a.b'), call('a.b.c'), call('a.b.c.d')]) def test_resolve(self): # test _resolve with the module that is not exposed to the top level with FullyQualifiedNamedTransform.with_filter('*'): dataframe_transform = FullyQualifiedNamedTransform._resolve( 'apache_beam.dataframe.transforms.DataframeTransform') from apache_beam.dataframe.transforms import DataframeTransform self.assertIs(dataframe_transform, DataframeTransform) # test _resolve with the module that will never be exposed # to the top level in the future with FullyQualifiedNamedTransform.with_filter('*'): argument_placeholder = FullyQualifiedNamedTransform._resolve( 'apache_beam.internal.util.ArgumentPlaceholder') from apache_beam.internal.util import ArgumentPlaceholder self.assertIs(argument_placeholder, ArgumentPlaceholder) class _TestTransform(beam.PTransform): @classmethod def create(cls, *args, **kwargs): return cls(*args, **kwargs) def __init__(self, prefix, suffix): self._prefix = prefix self._suffix = suffix def expand(self, pcoll): return pcoll | beam.Map(lambda s: self._prefix + s + self._suffix) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/managed.py0000644000175100017510000002034715206113033022030 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Managed Transforms. This module builds and instantiates turnkey transforms that can be managed by the underlying runner. This means the runner can upgrade the transform to a more optimal/updated version without requiring the user to do anything. It may also replace the transform with something entirely different if it chooses to. By default, however, the specified transform will remain unchanged. Available transforms ==================== Please check the Managed IO configuration page: https://beam.apache.org/documentation/io/managed-io/ Using Managed Transforms ======================== Managed turnkey transforms have a defined configuration and can be built using an inline :class:`dict` like so:: results = p | beam.managed.Read( beam.managed.ICEBERG, config={"table": "foo", "catalog_name": "bar", "catalog_properties": { "warehouse": "path/to/warehouse", "catalog-impl": "org.apache.my.CatalogImpl"}}) A YAML configuration file can also be used to build a Managed transform. Say we have the following `config.yaml` file:: topic: "foo" bootstrap_servers: "localhost:1234" format: "AVRO" Simply provide the location to the file like so:: input_rows = p | beam.Create(...) input_rows | beam.managed.Write( beam.managed.KAFKA, config_url="path/to/config.yaml") **Note:** inputs and outputs need to be PCollection(s) of Beam :py:class:`apache_beam.pvalue.Row` elements. Runner specific features ======================== Google Cloud Dataflow supports additional management features for `managed` including automatically upgrading transforms to the latest supported version. For more details and examples, please see Dataflow managed I/O https://cloud.google.com/dataflow/docs/guides/managed-io. """ from typing import Any from typing import Optional import yaml from apache_beam.portability.common_urns import ManagedTransforms from apache_beam.transforms.external import MANAGED_SCHEMA_TRANSFORM_IDENTIFIER from apache_beam.transforms.external import MANAGED_TRANSFORM_URN_TO_JAR_TARGET_MAPPING from apache_beam.transforms.external import BeamJarExpansionService from apache_beam.transforms.external import SchemaAwareExternalTransform from apache_beam.transforms.ptransform import PTransform ICEBERG = "iceberg" # TODO(https://github.com/apache/beam/issues/34212): keep ICEBERG_CDC private # until we vet it with integration tests _ICEBERG_CDC = "iceberg_cdc" KAFKA = "kafka" BIGQUERY = "bigquery" POSTGRES = "postgres" MYSQL = "mysql" SQL_SERVER = "sqlserver" __all__ = ["ICEBERG", "KAFKA", "BIGQUERY", "Read", "Write"] class Read(PTransform): """Read using Managed Transforms""" _READ_TRANSFORMS = { ICEBERG: ManagedTransforms.Urns.ICEBERG_READ.urn, _ICEBERG_CDC: ManagedTransforms.Urns.ICEBERG_CDC_READ.urn, KAFKA: ManagedTransforms.Urns.KAFKA_READ.urn, BIGQUERY: ManagedTransforms.Urns.BIGQUERY_READ.urn, POSTGRES: ManagedTransforms.Urns.POSTGRES_READ.urn, MYSQL: ManagedTransforms.Urns.MYSQL_READ.urn, SQL_SERVER: ManagedTransforms.Urns.SQL_SERVER_READ.urn, } def __init__( self, source: str, config: Optional[dict[str, Any]] = None, config_url: Optional[str] = None, expansion_service=None): super().__init__() self._source = source identifier = self._READ_TRANSFORMS.get(source.lower()) if not identifier: raise ValueError( f"An unsupported source was specified: '{source}'. Please specify " f"one of the following sources: {list(self._READ_TRANSFORMS.keys())}") # Store parameters for deferred expansion service creation self._identifier = identifier self._provided_expansion_service = expansion_service self._underlying_identifier = identifier self._yaml_config = yaml.dump(config) self._config_url = config_url def expand(self, input): # Create expansion service with access to pipeline options expansion_service = _resolve_expansion_service( self._source, self._identifier, self._provided_expansion_service, pipeline_options=input.pipeline._options) return input | SchemaAwareExternalTransform( identifier=MANAGED_SCHEMA_TRANSFORM_IDENTIFIER, expansion_service=expansion_service, rearrange_based_on_discovery=True, transform_identifier=self._underlying_identifier, config=self._yaml_config, config_url=self._config_url) def default_label(self) -> str: return "Managed Read(%s)" % self._source.upper() class Write(PTransform): """Write using Managed Transforms""" _WRITE_TRANSFORMS = { ICEBERG: ManagedTransforms.Urns.ICEBERG_WRITE.urn, KAFKA: ManagedTransforms.Urns.KAFKA_WRITE.urn, BIGQUERY: ManagedTransforms.Urns.BIGQUERY_WRITE.urn, POSTGRES: ManagedTransforms.Urns.POSTGRES_WRITE.urn, MYSQL: ManagedTransforms.Urns.MYSQL_WRITE.urn, SQL_SERVER: ManagedTransforms.Urns.SQL_SERVER_WRITE.urn } def __init__( self, sink: str, config: Optional[dict[str, Any]] = None, config_url: Optional[str] = None, expansion_service=None): super().__init__() self._sink = sink identifier = self._WRITE_TRANSFORMS.get(sink.lower()) if not identifier: raise ValueError( f"An unsupported sink was specified: '{sink}'. Please specify " f"one of the following sinks: {list(self._WRITE_TRANSFORMS.keys())}") # Store parameters for deferred expansion service creation self._identifier = identifier self._provided_expansion_service = expansion_service self._underlying_identifier = identifier self._yaml_config = yaml.dump(config) self._config_url = config_url def expand(self, input): # Create expansion service with access to pipeline options expansion_service = _resolve_expansion_service( self._sink, self._identifier, self._provided_expansion_service, pipeline_options=input.pipeline._options) return input | SchemaAwareExternalTransform( identifier=MANAGED_SCHEMA_TRANSFORM_IDENTIFIER, expansion_service=expansion_service, rearrange_based_on_discovery=True, transform_identifier=self._underlying_identifier, config=self._yaml_config, config_url=self._config_url) def default_label(self) -> str: return "Managed Write(%s)" % self._sink.upper() def _resolve_expansion_service( transform_name: str, identifier: str, expansion_service, pipeline_options=None): if expansion_service: return expansion_service gradle_target = None if identifier in MANAGED_TRANSFORM_URN_TO_JAR_TARGET_MAPPING: gradle_target = MANAGED_TRANSFORM_URN_TO_JAR_TARGET_MAPPING.get(identifier) if not gradle_target: raise ValueError( "No expansion service was specified and could not find a " f"default expansion service for {transform_name}: '{identifier}'.") # Extract maven_repository_url and user_agent from pipeline options if # available maven_repository_url = None user_agent = None if pipeline_options: from apache_beam.options import pipeline_options as po setup_options = pipeline_options.view_as(po.SetupOptions) maven_repository_url = setup_options.maven_repository_url user_agent = setup_options.user_agent return BeamJarExpansionService( gradle_target, maven_repository_url=maven_repository_url, user_agent=user_agent) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/managed_iceberg_it_test.py0000644000175100017510000000615215206113033025241 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import datetime import os import unittest import uuid import pytest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to @pytest.mark.uses_io_java_expansion_service @unittest.skipUnless( os.environ.get('EXPANSION_JARS'), "EXPANSION_JARS environment var is not provided, " "indicating that jars have not been built") class ManagedIcebergIT(unittest.TestCase): WAREHOUSE = "gs://temp-storage-for-end-to-end-tests" def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) self.args = self.test_pipeline.get_full_options_as_args() self.args.extend([ # '--experiments=enable_managed_transforms', ]) def _create_row(self, num: int): return beam.Row( int_=num, str_=str(num), bytes_=bytes(num), bool_=(num % 2 == 0), float_=(num + float(num) / 100), arr_=[num, num, num], date_=datetime.date.today() - datetime.timedelta(days=num)) def test_write_read_pipeline(self): biglake_catalog_props = { 'type': 'rest', 'uri': 'https://biglake.googleapis.com/iceberg/v1/restcatalog', 'warehouse': self.WAREHOUSE, 'header.x-goog-user-project': 'apache-beam-testing', 'rest.auth.type': 'google', 'io-impl': 'org.apache.iceberg.gcp.gcs.GCSFileIO', 'header.X-Iceberg-Access-Delegation': 'vended-credentials', 'rest-metrics-reporting-enabled': 'false' } iceberg_config = { "table": "test_iceberg_write_read.test_" + uuid.uuid4().hex, "catalog_name": "default", "catalog_properties": biglake_catalog_props } rows = [self._create_row(i) for i in range(100)] expected_dicts = [row.as_dict() for row in rows] with beam.Pipeline(argv=self.args) as write_pipeline: _ = ( write_pipeline | beam.Create(rows) | beam.managed.Write(beam.managed.ICEBERG, config=iceberg_config)) with beam.Pipeline(argv=self.args) as read_pipeline: output_dicts = ( read_pipeline | beam.managed.Read(beam.managed.ICEBERG, config=iceberg_config) | beam.Map(lambda row: row._asdict())) assert_that(output_dicts, equal_to(expected_dicts)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/maven_repository_url_test.py0000644000175100017510000002223315206113033025756 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the maven_repository_url functionality.""" import unittest from unittest import mock from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.transforms.external import MANAGED_TRANSFORM_URN_TO_JAR_TARGET_MAPPING # pylint: disable=line-too-long from apache_beam.transforms.external import BeamJarExpansionService from apache_beam.transforms.external import JavaJarExpansionService from apache_beam.transforms.managed import _resolve_expansion_service from apache_beam.utils.subprocess_server import JavaJarServer class MavenRepositoryUrlTest(unittest.TestCase): """Test cases for maven_repository_url functionality.""" def test_beam_jar_expansion_service_with_maven_repository_url(self): """Test that BeamJarExpansionService accepts and uses maven_repository_url.""" custom_repo_url = "https://custom.maven.repo/" custom_user_agent = "test-user-agent/1.0" with mock.patch.object(JavaJarServer, 'path_to_beam_jar') as \ mock_path_to_beam_jar: mock_path_to_beam_jar.return_value = "/path/to/beam.jar" service = BeamJarExpansionService( gradle_target="dummy:target", maven_repository_url=custom_repo_url, user_agent=custom_user_agent) # Verify that path_to_beam_jar was called with the custom repository # URL mock_path_to_beam_jar.assert_called_once() call_args = mock_path_to_beam_jar.call_args self.assertEqual(call_args[1]['maven_repository_url'], custom_repo_url) # Verify that the user_agent is stored self.assertEqual(service._user_agent, custom_user_agent) def test_java_jar_expansion_service_with_maven_repository_url(self): """Test that JavaJarExpansionService accepts and uses maven_repository_url.""" custom_repo_url = "https://custom.maven.repo/" custom_user_agent = "test-user-agent/1.0" service = JavaJarExpansionService( "dummy.jar", maven_repository_url=custom_repo_url, user_agent=custom_user_agent) # Verify that the maven_repository_url is stored self.assertEqual(service._maven_repository_url, custom_repo_url) # Verify that the user_agent is stored self.assertEqual(service._user_agent, custom_user_agent) def test_expand_jars_with_maven_repository_url(self): """Test that JavaJarExpansionService passes maven_repository_url to _expand_jars.""" custom_repo_url = "https://custom.maven.repo/" custom_user_agent = "test-user-agent/1.0" # Test with a Maven artifact format in classpath with mock.patch( 'apache_beam.transforms.external.JavaJarExpansionService._expand_jars' ) as mock_expand_jars: mock_expand_jars.return_value = ["/path/to/expanded.jar"] # Create service with maven_repository_url and user_agent service = JavaJarExpansionService( "dummy.jar", classpath=["group:artifact:1.0"], maven_repository_url=custom_repo_url, user_agent=custom_user_agent) # Call _default_args which should trigger _expand_jars service._default_args() # Verify that _expand_jars was called with the custom repository URL # and user_agent # Note: The actual call uses positional arguments, not keyword # arguments mock_expand_jars.assert_called_with( 'group:artifact:1.0', custom_user_agent, # user_agent custom_repo_url # maven_repository_url ) @mock.patch.dict( MANAGED_TRANSFORM_URN_TO_JAR_TARGET_MAPPING, {'test:identifier': 'test:gradle:target'}) def test_resolve_expansion_service_with_pipeline_options(self): """Test that _resolve_expansion_service uses maven_repository_url and user_agent from pipeline options.""" custom_repo_url = "https://custom.maven.repo/" custom_user_agent = "test-user-agent/1.0" # Create pipeline options with maven_repository_url and user_agent options = PipelineOptions() setup_options = options.view_as(SetupOptions) setup_options.maven_repository_url = custom_repo_url setup_options.user_agent = custom_user_agent with mock.patch.object(JavaJarServer, 'path_to_beam_jar') as \ mock_path_to_beam_jar: mock_path_to_beam_jar.return_value = "/path/to/beam.jar" # Call _resolve_expansion_service with pipeline options service = _resolve_expansion_service( "test_source", "test:identifier", None, pipeline_options=options) # Verify that the returned service has the correct parameters self.assertIsInstance(service, BeamJarExpansionService) self.assertEqual(service._maven_repository_url, custom_repo_url) self.assertEqual(service._user_agent, custom_user_agent) # Verify that path_to_beam_jar was called with the custom repository # URL mock_path_to_beam_jar.assert_called_once() call_args = mock_path_to_beam_jar.call_args self.assertEqual(call_args[1]['maven_repository_url'], custom_repo_url) @mock.patch.dict( MANAGED_TRANSFORM_URN_TO_JAR_TARGET_MAPPING, {'test:identifier': 'test:gradle:target'}) def test_resolve_expansion_service_without_maven_repository_url(self): """Test that _resolve_expansion_service works without maven_repository_url.""" # Create pipeline options without maven_repository_url options = PipelineOptions() with mock.patch.object(JavaJarServer, 'path_to_beam_jar') as \ mock_path_to_beam_jar: mock_path_to_beam_jar.return_value = "/path/to/beam.jar" # Call _resolve_expansion_service with pipeline options _ = _resolve_expansion_service( "test_source", "test:identifier", None, pipeline_options=options) # Verify that path_to_beam_jar was called without maven_repository_url mock_path_to_beam_jar.assert_called_once() call_args = mock_path_to_beam_jar.call_args self.assertIsNone(call_args[1].get('maven_repository_url')) @mock.patch.dict( MANAGED_TRANSFORM_URN_TO_JAR_TARGET_MAPPING, {'test:identifier': 'test:gradle:target'}) def test_resolve_expansion_service_without_pipeline_options(self): """Test that _resolve_expansion_service works without pipeline options.""" with mock.patch.object(JavaJarServer, 'path_to_beam_jar') as \ mock_path_to_beam_jar: mock_path_to_beam_jar.return_value = "/path/to/beam.jar" # Call _resolve_expansion_service without pipeline options _ = _resolve_expansion_service( "test_source", "test:identifier", None, pipeline_options=None) # Verify that path_to_beam_jar was called without maven_repository_url mock_path_to_beam_jar.assert_called_once() call_args = mock_path_to_beam_jar.call_args self.assertIsNone(call_args[1].get('maven_repository_url')) def test_user_agent_only_beam_jar_expansion_service(self): """Test BeamJarExpansionService with only user_agent parameter.""" custom_user_agent = "test-user-agent/1.0" with mock.patch.object(JavaJarServer, 'path_to_beam_jar') as \ mock_path_to_beam_jar: mock_path_to_beam_jar.return_value = "/path/to/beam.jar" service = BeamJarExpansionService( "dummy.jar", user_agent=custom_user_agent) # Verify that the user_agent is stored self.assertEqual(service._user_agent, custom_user_agent) # Verify that maven_repository_url is None (default) self.assertIsNone(service._maven_repository_url) def test_user_agent_only_java_jar_expansion_service(self): """Test JavaJarExpansionService with only user_agent parameter.""" custom_user_agent = "test-user-agent/1.0" service = JavaJarExpansionService("dummy.jar", user_agent=custom_user_agent) # Verify that the user_agent is stored self.assertEqual(service._user_agent, custom_user_agent) # Verify that maven_repository_url is None (default) self.assertIsNone(service._maven_repository_url) def test_default_user_agent_values(self): """Test that services have None as default user_agent.""" with mock.patch.object(JavaJarServer, 'path_to_beam_jar') as \ mock_path_to_beam_jar: mock_path_to_beam_jar.return_value = "/path/to/beam.jar" beam_service = BeamJarExpansionService("dummy.jar") java_service = JavaJarExpansionService("dummy.jar") # Verify that user_agent defaults to None self.assertIsNone(beam_service._user_agent) self.assertIsNone(java_service._user_agent) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/periodicsequence.py0000644000175100017510000003404015206113033023756 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import enum import math import time import warnings from typing import Any from typing import Optional from typing import Sequence import apache_beam as beam from apache_beam.io.restriction_trackers import OffsetRange from apache_beam.io.restriction_trackers import OffsetRestrictionTracker from apache_beam.io.watermark_estimators import ManualWatermarkEstimator from apache_beam.runners import sdf_utils from apache_beam.transforms import core from apache_beam.transforms import window from apache_beam.transforms.ptransform import PTransform from apache_beam.transforms.window import TimestampedValue from apache_beam.utils import timestamp from apache_beam.utils.timestamp import MAX_TIMESTAMP from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.timestamp import TimestampTypes class ImpulseSeqGenRestrictionProvider(core.RestrictionProvider): def initial_restriction(self, element): start, end, interval = element if not isinstance(start, Timestamp): start = Timestamp.of(start) if not isinstance(end, Timestamp): end = Timestamp.of(end) interval_duration = Duration(interval) assert start <= end assert interval > 0 total_duration = end - start total_outputs = math.ceil(total_duration.micros / interval_duration.micros) return OffsetRange(0, total_outputs) def create_tracker(self, restriction): return OffsetRestrictionTracker(restriction) def restriction_size(self, element, restriction): return _sequence_backlog_bytes(element, time.time(), restriction) # On drain, immediately stop emitting new elements def truncate(self, unused_element, unused_restriction): return None def _sequence_backlog_bytes(element, now, offset_range): ''' Calculates size of the output that the sequence should have emitted up to now. ''' start, _, interval = element if isinstance(start, Timestamp): start = start.micros / 1000000 assert interval > 0 now_index = math.floor((now - start) / interval) if now_index < offset_range.start: return 0 # We attempt to be precise as some runners scale based upon bytes and # output byte throughput. return 8 * (min(offset_range.stop, now_index) - offset_range.start) class ImpulseSeqGenDoFn(beam.DoFn): ''' ImpulseSeqGenDoFn fn receives tuple elements with three parts: * first_timestamp = The timestamp of the first element to be generated (inclusive). * last_timestamp = The timestamp marking the end of the generation period (exclusive). No elements will be generated at or after this time. * fire_interval = how often to fire an element. For each input element received, ImpulseSeqGenDoFn fn will start generating output elements in following pattern: * if element timestamp is less than current runtime then output element. * if element timestamp is greater than current runtime, wait until next element timestamp. ImpulseSeqGenDoFn can't guarantee that each element is output at exact time. ImpulseSeqGenDoFn guarantees that elements would not be output prior to given runtime timestamp. The output mode of the DoFn is based on the input `data`: - **None**: If `data` is None (by default), the output element will be the timestamp. - **Non-Timestamped Data**: If `data` is a sequence of arbitrary values (e.g., `[v1, v2, ...]`), the DoFn will assign a timestamp to each emitted element. - **Pre-Timestamped Data**: If `data` is a sequence of tuples, where each tuple is `(apache_beam.utils.timestamp.Timestamp, value)`, the DoFn will use the provided timestamp for the emitted element. See the parameter description of `PeriodicImpulse` for more information. ''' def __init__(self, data: Optional[Sequence[Any]] = None): self._data = data assert self._data is None or len(self._data) > 0 self._len = len(self._data) if self._data is not None else 0 self._is_pre_timestamped = self._data is not None and self._len > 0 and \ isinstance(self._data[0], tuple) and \ isinstance(self._data[0][0], timestamp.Timestamp) def _get_output(self, index, current_output_timestamp): if self._data is None: return TimestampedValue( current_output_timestamp, current_output_timestamp) if self._is_pre_timestamped: event_time, value = self._data[index % self._len] return TimestampedValue(value, event_time) else: value = self._data[index % self._len] return TimestampedValue(value, current_output_timestamp) @beam.DoFn.unbounded_per_element() def process( self, element, restriction_tracker=beam.DoFn.RestrictionParam( ImpulseSeqGenRestrictionProvider()), watermark_estimator=beam.DoFn.WatermarkEstimatorParam( ManualWatermarkEstimator.default_provider())): ''' :param element: (start_timestamp, end_timestamp, interval) :param restriction_tracker: :return: yields elements at processing real-time intervals with value of target output timestamp for the element. ''' start, _, interval = element if isinstance(start, Timestamp): start = start.micros / 1000000 assert isinstance(restriction_tracker, sdf_utils.RestrictionTrackerView) current_output_index = restriction_tracker.current_restriction().start while True: current_output_timestamp = start + interval * current_output_index if current_output_timestamp > time.time(): # we are too ahead of time, let's wait. restriction_tracker.defer_remainder( timestamp.Timestamp(current_output_timestamp)) break if not restriction_tracker.try_claim(current_output_index): # nothing to claim, just stop break output = self._get_output(current_output_index, current_output_timestamp) current_watermark = watermark_estimator.current_watermark() if current_watermark is None or output.timestamp > current_watermark: # ensure watermark is monotonic watermark_estimator.set_watermark(output.timestamp) yield output current_output_index += 1 # Don't yield any values here so that the generator # raises StopIteration when we break out of the while loop. class PeriodicSequence(PTransform): ''' PeriodicSequence transform receives tuple elements with three parts: * first_timestamp = first timestamp to output element for. * last_timestamp = last timestamp/time to output element for. * fire_interval = how often to fire an element. For each input element received, PeriodicSequence transform will start generating output elements in following pattern: * if element timestamp is less than current runtime then output element. * if element timestamp is greater than current runtime, wait until next element timestamp. PeriodicSequence can't guarantee that each element is output at exact time. PeriodicSequence guarantees that elements would not be output prior to given runtime timestamp. The PCollection generated by PeriodicSequence is unbounded. ''' def __init__(self): pass def expand(self, pcoll): return ( pcoll | 'GenSequence' >> beam.ParDo(ImpulseSeqGenDoFn()) | 'MapToTimestamped' >> beam.Map(lambda tt: TimestampedValue(tt, tt))) class RebaseMode(enum.Enum): '''Controls how the start and stop timestamps are rebased to execution time. Attributes: REBASE_NONE: Timestamps are not changed. REBASE_ALL: Both start and stop timestamps are rebased, preserving the original duration. REBASE_START: Only the start timestamp is rebased; the stop timestamp is unchanged. ''' REBASE_NONE = 0 REBASE_ALL = 1 REBASE_START = 2 class PeriodicImpulse(PTransform): ''' PeriodicImpulse transform generates an infinite sequence of elements with given runtime interval. PeriodicImpulse transform behaves same as {@link PeriodicSequence} transform, but can be used as first transform in pipeline. The PCollection generated by PeriodicImpulse is unbounded. ''' def _validate_and_adjust_duration(self): assert self.data # The total time we need to impulse all the data. data_duration = (len(self.data) - 1) * Duration(self.interval) is_pre_timestamped = isinstance(self.data[0], tuple) and \ isinstance(self.data[0][0], timestamp.Timestamp) start_ts = Timestamp.of(self.start_ts) stop_ts = Timestamp.of(self.stop_ts) if stop_ts == MAX_TIMESTAMP: # When the stop timestamp is unbounded (MAX_TIMESTAMP), set it to the # data's actual end time plus an extra fire interval, because the # impulse duration's upper bound is exclusive. self.stop_ts = start_ts + data_duration + Duration(self.interval) stop_ts = self.stop_ts # The total time for the impulse signal which occurs in [start, end). impulse_duration = stop_ts - start_ts if data_duration + Duration(self.interval) < impulse_duration: # We don't have enough data for the impulse. # If we can fit at least one more data point in the impulse duration, # then we will be in the repeat mode. message = 'The number of elements in the provided pre-timestamped ' \ 'data sequence is not enough to span the full impulse duration. ' \ f'Expected duration: {impulse_duration}, ' \ f'actual data duration: {data_duration}.' if is_pre_timestamped: raise ValueError( f'{message} Please either provide more data or decrease ' '`stop_timestamp`.') else: warnings.warn( f'{message} As a result, the data sequence will be repeated to ' 'generate elements for the entire duration.') def __init__( self, start_timestamp: TimestampTypes = Timestamp.now(), stop_timestamp: TimestampTypes = MAX_TIMESTAMP, fire_interval: float = 360.0, apply_windowing: bool = False, data: Optional[Sequence[Any]] = None, rebase: RebaseMode = RebaseMode.REBASE_NONE): ''' :param start_timestamp: Timestamp for first element. :param stop_timestamp: Timestamp at or after which no elements will be output. :param fire_interval: Interval in seconds at which to output elements. :param apply_windowing: Whether each element should be assigned to individual window. If false, all elements will reside in global window. :param data: A sequence of elements to emit. The behavior depends on the content: - **None (default):** The transform emits the event timestamps as the element values, starting from start_timestamp and incrementing by `fire_interval` up to the `stop_timestamp` (exclusive) - **Sequence of raw values (e.g., `['a', 'b']`)**: The transform emits each value in the sequence, assigning it an event timestamp that is calculated in the same manner as the default scenario. The sequence is repeated if the impulse duration requires more elements than are in the sequence (a warning will be given in this case). - **Sequence of pre-timestamped tuples (e.g., `[(t1, v1), (t2, v2)]`)**: The transform emits each value with its explicitly provided event time. The format must be `(apache_beam.utils.timestamp.Timestamp, value)`. The provided timestamps are used directly, overriding the calculated ones. Note that the elements in the sequence is NOT required to be ordered by event time; an element with a timestamp earlier than a preceding one will be treated as a potential late event. **Important**: In this mode, the number of elements in `data` must be sufficient to cover the duration defined by `start_timestamp`, `stop_timestamp`, and `fire_interval`; otherwise, a `ValueError` is raised. :param rebase: Controls how the start and stop timestamps are rebased to execution time. See `RebaseMode` for more details. Defaults to `REBASE_NONE`. ''' self.start_ts = start_timestamp self.stop_ts = stop_timestamp self.interval = fire_interval self.apply_windowing = apply_windowing self.data = data self.rebase = rebase if self.data: self._validate_and_adjust_duration() def expand(self, pbegin): if self.rebase == RebaseMode.REBASE_ALL: duration = Timestamp.of(self.stop_ts) - Timestamp.of(self.start_ts) impulse_element = pbegin | beam.Impulse() | beam.Map( lambda _: [now := Timestamp.now(), now + duration, self.interval]) elif self.rebase == RebaseMode.REBASE_START: impulse_element = pbegin | beam.Impulse() | beam.Map( lambda _: [Timestamp.now(), self.stop_ts, self.interval]) else: impulse_element = pbegin | 'ImpulseElement' >> beam.Create( [(self.start_ts, self.stop_ts, self.interval)]) result = ( impulse_element | 'GenSequence' >> beam.ParDo(ImpulseSeqGenDoFn(self.data))) if not self.data: # This step is actually an identity transform, because the Timestamped # values have already been generated in `ImpulseSeqGenDoFn`. # We keep this step here to prevent the current PeriodicImpulse from # breaking the compatibility. result = (result | 'MapToTimestamped' >> beam.Map(lambda tt: tt)) if self.apply_windowing: result = result | 'ApplyWindowing' >> beam.WindowInto( window.FixedWindows(self.interval)) return result ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/periodicsequence_it_test.py0000644000175100017510000000622315206113033025513 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Integration tests for cross-language transform expansion.""" # pytype: skip-file import time import unittest import pytest import apache_beam as beam from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import is_empty from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.transforms.core import DoFn from apache_beam.transforms.periodicsequence import PeriodicSequence @unittest.skipIf( not TestPipeline().get_pipeline_options().view_as( StandardOptions).streaming, "Watermark tests are only valid for streaming jobs.") class PeriodicSequenceIT(unittest.TestCase): def setUp(self): self.test_pipeline = TestPipeline(is_integration_test=True) @pytest.mark.it_postcommit @pytest.mark.sickbay_direct @pytest.mark.sickbay_spark @pytest.mark.timeout( 1800) # Timeout after 30 minutes to give Dataflow some extra time def test_periodicsequence_outputs_valid_watermarks_it(self): """Tests periodic sequence with watermarks on dataflow. For testing that watermarks are being correctly emitted, we make sure that there's not a long gap between an element being emitted and being correctly aggregated. """ class FindLongGaps(DoFn): def process(self, element): emitted_at, unused_count = element processed_at = time.time() if processed_at - emitted_at > 25: yield ( 'Elements emitted took too long to process.', emitted_at, processed_at) start_time = time.time() # Run long enough for Dataflow to start up duration_sec = 540 end_time = start_time + duration_sec interval = 1 res = ( self.test_pipeline | 'ImpulseElement' >> beam.Create([(start_time, end_time, interval)]) | 'ImpulseSeqGen' >> PeriodicSequence() | 'MapToCurrentTime' >> beam.Map(lambda element: time.time()) | 'window_into' >> beam.WindowInto( window.FixedWindows(2), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.combiners.Count.PerElement() | beam.ParDo(FindLongGaps())) assert_that(res, is_empty()) self.test_pipeline.run().wait_until_finish() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/periodicsequence_test.py0000644000175100017510000003173315206113033025023 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the PTransform and descendants.""" # pytype: skip-file import inspect import logging import random import time import unittest from parameterized import parameterized import apache_beam as beam from apache_beam.io.restriction_trackers import OffsetRange from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.testing.util import is_empty from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.transforms.periodicsequence import PeriodicImpulse from apache_beam.transforms.periodicsequence import PeriodicSequence from apache_beam.transforms.periodicsequence import RebaseMode from apache_beam.transforms.periodicsequence import _sequence_backlog_bytes from apache_beam.transforms.window import FixedWindows from apache_beam.utils.timestamp import Timestamp # Disable frequent lint warning due to pipe operator for chaining transforms. # pylint: disable=expression-not-assigned class PeriodicSequenceTest(unittest.TestCase): def test_periodicsequence_outputs_valid_sequence(self): start_offset = 1 start_time = time.time() + start_offset duration = 1 end_time = start_time + duration interval = 0.25 with TestPipeline() as p: result = ( p | 'ImpulseElement' >> beam.Create([(start_time, end_time, interval)]) | 'ImpulseSeqGen' >> PeriodicSequence()) k = [ start_time + x * interval for x in range(0, int(duration / interval), 1) ] self.assertEqual(result.is_bounded, False) assert_that(result, equal_to(k)) def test_periodicsequence_outputs_valid_sequence_in_past(self): start_offset = -10000 it = time.time() + start_offset duration = 5 et = it + duration interval = 1 with TestPipeline() as p: result = ( p | 'ImpulseElement' >> beam.Create([(it, et, interval)]) | 'ImpulseSeqGen' >> PeriodicSequence()) k = [it + x * interval for x in range(0, int(duration / interval), 1)] self.assertEqual(result.is_bounded, False) assert_that(result, equal_to(k)) def test_periodicsequence_output_size(self): element = [0, 1000000000, 10] self.assertEqual( _sequence_backlog_bytes(element, 100, OffsetRange(10, 100000000)), 0) self.assertEqual( _sequence_backlog_bytes(element, 100, OffsetRange(9, 100000000)), 8) self.assertEqual( _sequence_backlog_bytes(element, 100, OffsetRange(8, 100000000)), 16) self.assertEqual( _sequence_backlog_bytes(element, 101, OffsetRange(9, 100000000)), 8) self.assertEqual( _sequence_backlog_bytes(element, 10000, OffsetRange(0, 100000000)), 8 * 10000 / 10) self.assertEqual( _sequence_backlog_bytes(element, 10000, OffsetRange(1002, 1003)), 0) self.assertEqual( _sequence_backlog_bytes(element, 10100, OffsetRange(1002, 1003)), 8) class PeriodicImpulseTest(unittest.TestCase): def test_windowing_on_si(self): start_offset = -15 it = time.time() + start_offset duration = 15 et = it + duration interval = 5 with TestPipeline() as p: si = ( p | 'PeriodicImpulse' >> PeriodicImpulse(it, et, interval, True) | 'AddKey' >> beam.Map(lambda v: ('key', v)) | 'GBK' >> beam.GroupByKey() | 'SortGBK' >> beam.MapTuple(lambda k, vs: (k, sorted(vs)))) actual = si k = [('key', [it + x * interval]) for x in range(0, int(duration / interval), 1)] assert_that(actual, equal_to(k)) def test_default_start(self): default_parameters = inspect.signature(PeriodicImpulse.__init__).parameters it = default_parameters["start_timestamp"].default duration = 1 et = it + duration interval = 0.5 # Check default `stop_timestamp` is the same type `start_timestamp` is_same_type = isinstance( it, type(default_parameters["stop_timestamp"].default)) error = "'start_timestamp' and 'stop_timestamp' have different type" assert is_same_type, error with TestPipeline() as p: result = p | 'PeriodicImpulse' >> PeriodicImpulse(it, et, interval) k = [it + x * interval for x in range(0, int(duration / interval))] self.assertEqual(result.is_bounded, False) assert_that(result, equal_to(k)) @unittest.skip("hard to determine warm-up time and threshold for runners.") def test_processing_time(self): warmup_time = 3 threshold = 0.5 with TestPipeline() as p: ret = ( p | PeriodicImpulse( start_timestamp=Timestamp.now() + warmup_time, data=[10, 20, 30], fire_interval=2) | beam.Map(lambda _: time.time()) | beam.WindowInto( window.GlobalWindows(), trigger=trigger.Repeatedly(trigger.AfterCount(3)), accumulation_mode=trigger.AccumulationMode.DISCARDING, ) | beam.GroupBy() | beam.FlatMap(lambda x: [v - min(x[1]) for v in x[1]])) expected = [0, 2, 4] assert_that(ret, equal_to(expected, lambda x, y: abs(x - y) < threshold)) @parameterized.expand([0.5, 1, 2, 10]) def test_stop_over_by_epsilon(self, interval): with TestPipeline() as p: ret = ( p | PeriodicImpulse( start_timestamp=Timestamp(seconds=1), stop_timestamp=Timestamp(seconds=1, micros=1), data=[1, 2], fire_interval=interval) | beam.WindowInto(FixedWindows(interval)) | beam.WithKeys(0) | beam.GroupByKey()) expected = [ (0, [1]), ] assert_that(ret, equal_to(expected)) @parameterized.expand([1, 2]) def test_stop_over_by_interval(self, interval): with TestPipeline() as p: ret = ( p | PeriodicImpulse( start_timestamp=Timestamp(seconds=1), stop_timestamp=Timestamp(seconds=1 + interval), data=[1, 2], fire_interval=interval) | beam.WindowInto(FixedWindows(interval)) | beam.WithKeys(0) | beam.GroupByKey()) expected = [(0, [1])] assert_that(ret, equal_to(expected)) @parameterized.expand([1, 2]) def test_stop_over_by_interval_and_epsilon(self, interval): with TestPipeline() as p: ret = ( p | PeriodicImpulse( start_timestamp=Timestamp(seconds=1), stop_timestamp=Timestamp(seconds=1 + interval, micros=1), data=[1, 2], fire_interval=interval) | beam.WindowInto(FixedWindows(interval)) | beam.WithKeys(0) | beam.GroupByKey()) expected = [(0, [1]), (0, [2])] assert_that(ret, equal_to(expected)) def test_interval(self): with TestPipeline() as p: ret = ( p | PeriodicImpulse(data=[1, 2, 3, 4], fire_interval=0.5) | beam.WindowInto(FixedWindows(0.5)) | beam.WithKeys(0) | beam.GroupByKey()) expected = [(0, [1]), (0, [2]), (0, [3]), (0, [4])] assert_that(ret, equal_to(expected)) def test_repeat(self): now = Timestamp.now() with self.assertWarnsRegex(UserWarning, "not enough to span"): with TestPipeline() as p: ret = ( p | PeriodicImpulse( start_timestamp=now, stop_timestamp=now + 2.6, data=[1, 2, 3, 4], fire_interval=0.5) | beam.WindowInto(FixedWindows(0.5)) | beam.WithKeys(0) | beam.GroupByKey()) expected = [(0, [1]), (0, [2]), (0, [3]), (0, [4]), (0, [1]), (0, [2])] assert_that(ret, equal_to(expected)) def test_timestamped_value(self): with TestPipeline() as p: ret = ( p | PeriodicImpulse( data=[(Timestamp(1), 1), (Timestamp(3), 2), (Timestamp(2), 3), (Timestamp(1), 4)], fire_interval=0.5) | beam.WindowInto(FixedWindows(0.5)) | beam.WithKeys(0) | beam.GroupByKey()) expected = [(0, [1, 4]), (0, [2]), (0, [3])] assert_that(ret, equal_to(expected)) def test_not_enough_timestamped_value(self): now = Timestamp.now() data = [(Timestamp(1), 1), (Timestamp(2), 2), (Timestamp(3), 3)] with self.assertRaisesRegex(ValueError, "not enough to span"): with TestPipeline() as p: _ = ( p | PeriodicImpulse( start_timestamp=now, stop_timestamp=now + 2.6, data=data, fire_interval=0.5)) def test_fuzzy_length_and_interval(self): times = 30 for _ in range(times): seed = int(time.time() * 1000) random.seed(seed) n = int(random.randint(1, 100)) data = list(range(n)) m = random.randint(1, 1000) interval = m / 1e6 now = Timestamp.now() try: with TestPipeline() as p: ret = ( p | PeriodicImpulse( start_timestamp=now, data=data, fire_interval=interval)) assert_that(ret, equal_to(data)) except Exception as e: # pylint: disable=broad-except logging.error("Error occurred at random seed=%d", seed) raise e def test_fuzzy_length_at_minimal_interval(self): times = 30 for _ in range(times): seed = int(time.time() * 1000) random.seed(seed) n = int(random.randint(1, 100)) data = list(range(n)) interval = 1e-6 now = Timestamp.now() try: with TestPipeline() as p: ret = ( p | PeriodicImpulse( start_timestamp=now, data=data, fire_interval=interval)) assert_that(ret, equal_to(data)) except Exception as e: # pylint: disable=broad-except logging.error("Error occurred at random seed=%d", seed) raise e def test_int_type_input(self): # This test is to verify that if input timestamps and interval are integers, # the generated timestamped values are also integers. # This is necessary for the following test to pass: # apache_beam.examples.snippets.snippets_test.SlowlyChangingSideInputsTest with TestPipeline() as p: ret = ( p | PeriodicImpulse( start_timestamp=1, stop_timestamp=5, fire_interval=1)) expected = [1, 2, 3, 4] assert_that( ret, equal_to(expected, lambda x, y: type(x) is type(y) and x == y)) def test_float_type_input(self): with TestPipeline() as p: ret = ( p | PeriodicImpulse( start_timestamp=1.0, stop_timestamp=5.0, fire_interval=1)) expected = [1.0, 2.0, 3.0, 4.0] assert_that( ret, equal_to(expected, lambda x, y: type(x) is type(y) and x == y)) def test_timestamp_type_input(self): with TestPipeline() as p: ret = ( p | PeriodicImpulse( start_timestamp=Timestamp.of(1), stop_timestamp=Timestamp.of(5), fire_interval=1)) expected = [1.0, 2.0, 3.0, 4.0] assert_that( ret, equal_to(expected, lambda x, y: type(x) is type(y) and x == y)) def test_rebase_timestamp(self): class CheckTimeStamp(beam.DoFn): def process(self, elem): ts = Timestamp.of(elem) now = Timestamp.now() # When rebase is enabled, the timestamp should be closer to now than the # original start. if (ts - Timestamp.of(1)) < (now - ts): yield "wrong" with TestPipeline() as p: ret = ( p | PeriodicImpulse( start_timestamp=Timestamp.of(1), stop_timestamp=Timestamp.of(5), fire_interval=1, rebase=RebaseMode.REBASE_ALL) | beam.ParDo(CheckTimeStamp())) assert_that(ret, is_empty()) def test_rebase_timestamp_with_wrong_setting(self): with self.assertRaises(Exception): # exception is raised because start_timestamp is rebased to the pipeline # execution time, but the stop_timestamp is 5 seconds after unix epoch. with TestPipeline() as p: _ = ( p | PeriodicImpulse( start_timestamp=Timestamp.of(1), stop_timestamp=Timestamp.of(5), fire_interval=1, rebase=RebaseMode.REBASE_START)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/ptransform.py0000644000175100017510000013162015206113033022624 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """PTransform and descendants. A PTransform is an object describing (not executing) a computation. The actual execution semantics for a transform is captured by a runner object. A transform object always belongs to a pipeline object. A PTransform derived class needs to define the expand() method that describes how one or more PValues are created by the transform. The module defines a few standard transforms: FlatMap (parallel do), GroupByKey (group by key), etc. Note that the expand() methods for these classes contain code that will add nodes to the processing graph associated with a pipeline. As support for the FlatMap transform, the module also defines a DoFn class and wrapper class that allows lambda functions to be used as FlatMap processing functions. """ # pytype: skip-file import copy import functools import inspect import itertools import json import logging import operator import os import sys import threading import warnings from collections.abc import Callable from collections.abc import Mapping from collections.abc import Sequence from functools import reduce from functools import wraps from typing import TYPE_CHECKING from typing import Any from typing import Generic from typing import Optional from typing import TypeVar from typing import Union from typing import overload from google.protobuf import message from apache_beam import error from apache_beam import pvalue from apache_beam.internal import pickler from apache_beam.internal import util from apache_beam.portability import python_urns from apache_beam.pvalue import DoOutputsTuple from apache_beam.transforms import resources from apache_beam.transforms.display import DisplayDataItem from apache_beam.transforms.display import HasDisplayData from apache_beam.transforms.sideinputs import SIDE_INPUT_PREFIX from apache_beam.typehints import native_type_compatibility from apache_beam.typehints import typehints from apache_beam.typehints.decorators import IOTypeHints from apache_beam.typehints.decorators import TypeCheckError from apache_beam.typehints.decorators import WithTypeHints from apache_beam.typehints.decorators import get_signature from apache_beam.typehints.decorators import get_type_hints from apache_beam.typehints.decorators import getcallargs_forhints from apache_beam.typehints.trivial_inference import instance_to_type from apache_beam.typehints.typehints import validate_composite_type_param from apache_beam.utils import proto_utils from apache_beam.utils import python_callable if TYPE_CHECKING: from apache_beam import coders from apache_beam.pipeline import Pipeline from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners.pipeline_context import PipelineContext from apache_beam.transforms.core import Windowing __all__ = [ 'PTransform', 'ptransform_fn', 'label_from_callable', 'annotate_yaml', ] _LOGGER = logging.getLogger(__name__) T = TypeVar('T') InputT = TypeVar('InputT') OutputT = TypeVar('OutputT') PTransformT = TypeVar('PTransformT', bound='PTransform') ConstructorFn = Callable[ ['beam_runner_api_pb2.PTransform', Optional[Any], 'PipelineContext'], Any] ptransform_fn_typehints_enabled = False class _PValueishTransform(object): """Visitor for PValueish objects. A PValueish is a PValue, or list, tuple, dict of PValuesish objects. This visits a PValueish, contstructing a (possibly mutated) copy. """ def visit_nested(self, node, *args): if isinstance(node, (tuple, list)): args = [self.visit(x, *args) for x in node] if isinstance(node, tuple) and hasattr(node.__class__, '_make'): # namedtuples require unpacked arguments in their constructor return node.__class__(*args) else: return node.__class__(args) elif isinstance(node, dict): return node.__class__( {key: self.visit(value, *args) for (key, value) in node.items()}) else: return node class _SetInputPValues(_PValueishTransform): def visit(self, node, replacements): if id(node) in replacements: return replacements[id(node)] else: return self.visit_nested(node, replacements) # Caches to allow for materialization of values when executing a pipeline # in-process, in eager mode. This cache allows the same _MaterializedResult # object to be accessed and used despite Runner API round-trip serialization. _pipeline_materialization_cache = { } # type: dict[tuple[int, int], dict[int, _MaterializedResult]] _pipeline_materialization_lock = threading.Lock() def _allocate_materialized_pipeline(pipeline): # type: (Pipeline) -> None pid = os.getpid() with _pipeline_materialization_lock: pipeline_id = id(pipeline) _pipeline_materialization_cache[(pid, pipeline_id)] = {} def _allocate_materialized_result(pipeline): # type: (Pipeline) -> _MaterializedResult pid = os.getpid() with _pipeline_materialization_lock: pipeline_id = id(pipeline) if (pid, pipeline_id) not in _pipeline_materialization_cache: raise ValueError( 'Materialized pipeline is not allocated for result ' 'cache.') result_id = len(_pipeline_materialization_cache[(pid, pipeline_id)]) result = _MaterializedResult(pipeline_id, result_id) _pipeline_materialization_cache[(pid, pipeline_id)][result_id] = result return result def _get_materialized_result(pipeline_id, result_id): # type: (int, int) -> _MaterializedResult pid = os.getpid() with _pipeline_materialization_lock: if (pid, pipeline_id) not in _pipeline_materialization_cache: raise Exception( 'Materialization in out-of-process and remote runners is not yet ' 'supported.') return _pipeline_materialization_cache[(pid, pipeline_id)][result_id] def _release_materialized_pipeline(pipeline): # type: (Pipeline) -> None pid = os.getpid() with _pipeline_materialization_lock: pipeline_id = id(pipeline) del _pipeline_materialization_cache[(pid, pipeline_id)] class _MaterializedResult(object): def __init__(self, pipeline_id, result_id): # type: (int, int) -> None self._pipeline_id = pipeline_id self._result_id = result_id self.elements = [] # type: list[Any] def __reduce__(self): # When unpickled (during Runner API roundtrip serailization), get the # _MaterializedResult object from the cache so that values are written # to the original _MaterializedResult when run in eager mode. return (_get_materialized_result, (self._pipeline_id, self._result_id)) class _MaterializedDoOutputsTuple(pvalue.DoOutputsTuple): def __init__(self, deferred, results_by_tag): super().__init__(None, None, deferred._tags, deferred._main_tag) self._deferred = deferred self._results_by_tag = results_by_tag def __getitem__(self, tag): if tag not in self._results_by_tag: raise KeyError( 'Tag %r is not a defined output tag of %s.' % (tag, self._deferred)) return self._results_by_tag[tag].elements class _AddMaterializationTransforms(_PValueishTransform): def _materialize_transform(self, pipeline): result = _allocate_materialized_result(pipeline) # Need to define _MaterializeValuesDoFn here to avoid circular # dependencies. from apache_beam import DoFn from apache_beam import ParDo class _MaterializeValuesDoFn(DoFn): def __init__(self): self.is_materialize_values_do_fn = True def process(self, element): result.elements.append(element) materialization_label = '_MaterializeValues%d' % result._result_id return (materialization_label >> ParDo(_MaterializeValuesDoFn()), result) def visit(self, node): if isinstance(node, pvalue.PValue): transform, result = self._materialize_transform(node.pipeline) node | transform return result elif isinstance(node, pvalue.DoOutputsTuple): results_by_tag = {} for tag in itertools.chain([node._main_tag], node._tags): results_by_tag[tag] = self.visit(node[tag]) return _MaterializedDoOutputsTuple(node, results_by_tag) else: return self.visit_nested(node) class _FinalizeMaterialization(_PValueishTransform): def visit(self, node): if isinstance(node, _MaterializedResult): return node.elements elif isinstance(node, _MaterializedDoOutputsTuple): return node else: return self.visit_nested(node) def get_named_nested_pvalues(pvalueish, as_inputs=False): if isinstance(pvalueish, tuple): # Check to see if it's a named tuple. fields = getattr(pvalueish, '_fields', None) if fields and len(fields) == len(pvalueish): tagged_values = zip(fields, pvalueish) else: tagged_values = enumerate(pvalueish) elif isinstance(pvalueish, list): if as_inputs: # Full list treated as a list of value for eager evaluation. yield None, pvalueish return tagged_values = enumerate(pvalueish) elif isinstance(pvalueish, dict): tagged_values = pvalueish.items() else: if as_inputs or isinstance(pvalueish, (pvalue.PValue, pvalue.DoOutputsTuple)): yield None, pvalueish return for tag, subvalue in tagged_values: for subtag, subsubvalue in get_named_nested_pvalues( subvalue, as_inputs=as_inputs): if subtag is None: yield tag, subsubvalue else: yield '%s.%s' % (tag, subtag), subsubvalue class _ZipPValues(object): """Pairs each PValue in a pvalueish with a value in a parallel out sibling. Sibling should have the same nested structure as pvalueish. Leaves in sibling are expanded across nested pvalueish lists, tuples, and dicts. For example ZipPValues().visit({'a': pc1, 'b': (pc2, pc3)}, {'a': 'A', 'b', 'B'}) will return [('a', pc1, 'A'), ('b', pc2, 'B'), ('b', pc3, 'B')] """ def visit(self, pvalueish, sibling, pairs=None, context=None): if pairs is None: pairs = [] self.visit(pvalueish, sibling, pairs, context) return pairs elif isinstance(pvalueish, (pvalue.PValue, pvalue.DoOutputsTuple)): pairs.append((context, pvalueish, sibling)) elif isinstance(pvalueish, (list, tuple)): self.visit_sequence(pvalueish, sibling, pairs, context) elif isinstance(pvalueish, dict): self.visit_dict(pvalueish, sibling, pairs, context) def visit_sequence(self, pvalueish, sibling, pairs, context): if isinstance(sibling, (list, tuple)): for ix, (p, s) in enumerate(zip(pvalueish, list(sibling) + [None] * len(pvalueish))): self.visit(p, s, pairs, 'position %s' % ix) else: for p in pvalueish: self.visit(p, sibling, pairs, context) def visit_dict(self, pvalueish, sibling, pairs, context): if isinstance(sibling, dict): for key, p in pvalueish.items(): self.visit(p, sibling.get(key), pairs, key) else: for p in pvalueish.values(): self.visit(p, sibling, pairs, context) class PTransform(WithTypeHints, HasDisplayData, Generic[InputT, OutputT]): """A transform object used to modify one or more PCollections. Subclasses must define an expand() method that will be used when the transform is applied to some arguments. Typical usage pattern will be: input | CustomTransform(...) The expand() method of the CustomTransform object passed in will be called with input as an argument. """ # By default, transforms don't have any side inputs. side_inputs = () # type: Sequence[pvalue.AsSideInput] # Used for nullary transforms. pipeline = None # type: Optional[Pipeline] # Default is unset. _user_label = None # type: Optional[str] def __init__(self, label=None): # type: (Optional[str]) -> None super().__init__() self.label = label # type: ignore # https://github.com/python/mypy/issues/3004 @property def label(self): # type: () -> str return self._user_label or self.default_label() @label.setter def label(self, value): # type: (Optional[str]) -> None self._user_label = value def default_label(self): # type: () -> str return self.__class__.__name__ def annotations(self) -> dict[str, Union[bytes, str, message.Message]]: return { 'python_type': # f'{self.__class__.__module__}.{self.__class__.__qualname__}' } def default_type_hints(self): fn_type_hints = IOTypeHints.from_callable(self.expand) if fn_type_hints is not None: fn_type_hints = fn_type_hints.strip_pcoll() # Prefer class decorator type hints for backwards compatibility. return get_type_hints(self.__class__).with_defaults(fn_type_hints) def with_input_types(self, input_type_hint): """Annotates the input type of a :class:`PTransform` with a type-hint. Args: input_type_hint (type): An instance of an allowed built-in type, a custom class, or an instance of a :class:`~apache_beam.typehints.typehints.TypeConstraint`. Raises: TypeError: If **input_type_hint** is not a valid type-hint. See :obj:`apache_beam.typehints.typehints.validate_composite_type_param()` for further details. Returns: PTransform: A reference to the instance of this particular :class:`PTransform` object. This allows chaining type-hinting related methods. """ input_type_hint = native_type_compatibility.convert_to_beam_type( input_type_hint) validate_composite_type_param( input_type_hint, 'Type hints for a PTransform') return super().with_input_types(input_type_hint) def with_output_types(self, type_hint, **tagged_type_hints): """Annotates the output type of a :class:`PTransform` with a type-hint. Args: type_hint (type): An instance of an allowed built-in type, a custom class, or a :class:`~apache_beam.typehints.typehints.TypeConstraint`. This is the type hint for the main output. **tagged_type_hints: Type hints for tagged outputs. Each keyword argument specifies the type for a tagged output e.g., ``errors=str``. Raises: TypeError: If **type_hint** is not a valid type-hint. See :obj:`~apache_beam.typehints.typehints.validate_composite_type_param()` for further details. Returns: PTransform: A reference to the instance of this particular :class:`PTransform` object. This allows chaining type-hinting related methods. Example:: result = pcoll | beam.ParDo(MyDoFn()).with_output_types( int, # main output type errors=str, # 'errors' tagged output type warnings=str # 'warnings' tagged output type ).with_outputs('errors', 'warnings', main='main') """ type_hint = native_type_compatibility.convert_to_beam_type(type_hint) validate_composite_type_param(type_hint, 'Type hints for a PTransform') for tag, hint in tagged_type_hints.items(): tagged_type_hints[tag] = native_type_compatibility.convert_to_beam_type( hint) validate_composite_type_param( tagged_type_hints[tag], f'Tagged output type hint for {tag!r}') return super().with_output_types(type_hint, **tagged_type_hints) def with_resource_hints(self, **kwargs): # type: (...) -> PTransform """Adds resource hints to the :class:`PTransform`. Resource hints allow users to express constraints on the environment where the transform should be executed. Interpretation of the resource hints is defined by Beam Runners. Runners may ignore the unsupported hints. Args: **kwargs: key-value pairs describing hints and their values. Raises: ValueError: if provided hints are unknown to the SDK. See :mod:`apache_beam.transforms.resources` for a list of known hints. Returns: PTransform: A reference to the instance of this particular :class:`PTransform` object. """ self.get_resource_hints().update(resources.parse_resource_hints(kwargs)) return self def get_resource_hints(self): # type: () -> dict[str, bytes] if '_resource_hints' not in self.__dict__: # PTransform subclasses don't always call super(), so prefer lazy # initialization. By default, transforms don't have any resource hints. self._resource_hints = {} # type: dict[str, bytes] return self._resource_hints def type_check_inputs(self, pvalueish): self.type_check_inputs_or_outputs(pvalueish, 'input') def infer_output_type(self, unused_input_type): return self.get_type_hints().simple_output_type(self.label) or typehints.Any def type_check_outputs(self, pvalueish): self.type_check_inputs_or_outputs(pvalueish, 'output') def type_check_inputs_or_outputs(self, pvalueish, input_or_output): type_hints = self.get_type_hints() hints = getattr(type_hints, input_or_output + '_types') if hints is None or not any(hints): return arg_hints, kwarg_hints = hints # Output types can have kwargs for tagged output types. if arg_hints and kwarg_hints and input_or_output != 'output': raise TypeCheckError( 'PTransform cannot have both positional and keyword input type hints' ' without overriding %s._type_check_%s()' % (self.__class__, input_or_output)) root_hint = ( arg_hints[0] if len(arg_hints) == 1 else arg_hints or kwarg_hints) for context, pvalue_, hint in _ZipPValues().visit(pvalueish, root_hint): if isinstance(pvalue_, DoOutputsTuple): continue if pvalue_.element_type is None: # TODO(robertwb): It's a bug that we ever get here. (typecheck) continue if hint and not typehints.is_consistent_with(pvalue_.element_type, hint): at_context = ' %s %s' % (input_or_output, context) if context else '' raise TypeCheckError( '{type} type hint violation at {label}{context}: expected {hint}, ' 'got {actual_type}'.format( type=input_or_output.title(), label=self.label, context=at_context, hint=hint, actual_type=pvalue_.element_type)) def _infer_output_coder(self, input_type=None, input_coder=None): # type: (...) -> Optional[coders.Coder] """Returns the output coder to use for output of this transform. The Coder returned here should not be wrapped in a WindowedValueCoder wrapper. Args: input_type: An instance of an allowed built-in type, a custom class, or a typehints.TypeConstraint for the input type, or None if not available. input_coder: Coder object for encoding input to this PTransform, or None if not available. Returns: Coder object for encoding output of this PTransform or None if unknown. """ # TODO(ccy): further refine this API. return None def _clone(self, new_label): """Clones the current transform instance under a new label.""" transform = copy.copy(self) transform.label = new_label return transform def expand(self, input_or_inputs: InputT) -> OutputT: raise NotImplementedError def __str__(self): return '<%s>' % self._str_internal() def __repr__(self): return '<%s at %s>' % (self._str_internal(), hex(id(self))) def _str_internal(self): return '%s(PTransform)%s%s%s' % ( self.__class__.__name__, ' label=[%s]' % self.label if (hasattr(self, 'label') and self.label) else '', ' inputs=%s' % str(self.inputs) if (hasattr(self, 'inputs') and self.inputs) else '', ' side_inputs=%s' % str(self.side_inputs) if self.side_inputs else '') def _check_pcollection(self, pcoll): # type: (pvalue.PCollection) -> None if not isinstance(pcoll, pvalue.PCollection): raise error.TransformError('Expecting a PCollection argument.') if not pcoll.pipeline: raise error.TransformError('PCollection not part of a pipeline.') def get_windowing(self, inputs): # type: (Any) -> Windowing """Returns the window function to be associated with transform's output. By default most transforms just return the windowing function associated with the input PCollection (or the first input if several). """ if inputs: return inputs[0].windowing else: from apache_beam.transforms.core import Windowing from apache_beam.transforms.window import GlobalWindows # TODO(robertwb): Return something compatible with every windowing? return Windowing(GlobalWindows()) def __rrshift__(self, label): return _NamedPTransform(self, label) def __or__(self, right): """Used to compose PTransforms, e.g., ptransform1 | ptransform2.""" if isinstance(right, PTransform): return _ChainedPTransform(self, right) return NotImplemented def __ror__(self, left, label=None): """Used to apply this PTransform to non-PValues, e.g., a tuple.""" pvalueish, pvalues = self._extract_input_pvalues(left) if isinstance(pvalues, dict): pvalues = tuple(pvalues.values()) pipelines = [v.pipeline for v in pvalues if isinstance(v, pvalue.PValue)] if pvalues and not pipelines: deferred = False # pylint: disable=wrong-import-order, wrong-import-position from apache_beam import pipeline from apache_beam.options.pipeline_options import PipelineOptions # pylint: enable=wrong-import-order, wrong-import-position p = pipeline.Pipeline('DirectRunner', PipelineOptions(sys.argv)) else: if not pipelines: if self.pipeline is not None: p = self.pipeline else: raise ValueError( '"%s" requires a pipeline to be specified ' 'as there are no deferred inputs.' % self.label) else: p = self.pipeline or pipelines[0] for pp in pipelines: if p != pp: raise ValueError( 'Mixing values in different pipelines is not allowed.' '\n{%r} != {%r}' % (p, pp)) deferred = not getattr(p.runner, 'is_eager', False) # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.transforms.core import Create # pylint: enable=wrong-import-order, wrong-import-position replacements = { id(v): p | 'CreatePInput%s' % ix >> Create(v, reshuffle=False) for (ix, v) in enumerate(pvalues) if not isinstance(v, pvalue.PValue) and v is not None } pvalueish = _SetInputPValues().visit(pvalueish, replacements) self.pipeline = p result = p.apply(self, pvalueish, label) if deferred: return result _allocate_materialized_pipeline(p) materialized_result = _AddMaterializationTransforms().visit(result) p.run().wait_until_finish() _release_materialized_pipeline(p) return _FinalizeMaterialization().visit(materialized_result) def _extract_input_pvalues(self, pvalueish): """Extract all the pvalues contained in the input pvalueish. Returns pvalueish as well as the flat inputs list as the input may have to be copied as inspection may be destructive. By default, recursively extracts tuple components and dict values. Generally only needs to be overriden for multi-input PTransforms. """ # pylint: disable=wrong-import-order from apache_beam import pipeline # pylint: enable=wrong-import-order if isinstance(pvalueish, pipeline.Pipeline): pvalueish = pvalue.PBegin(pvalueish) return pvalueish, { str(tag): value for (tag, value) in get_named_nested_pvalues( pvalueish, as_inputs=True) } def _pvaluish_from_dict(self, input_dict): if len(input_dict) == 1: return next(iter(input_dict.values())) else: return input_dict def _named_inputs(self, main_inputs, side_inputs): # type: (Mapping[str, pvalue.PValue], Sequence[Any]) -> dict[str, pvalue.PValue] """Returns the dictionary of named inputs (including side inputs) as they should be named in the beam proto. """ main_inputs = { tag: input for (tag, input) in main_inputs.items() if isinstance(input, pvalue.PCollection) } named_side_inputs = {(SIDE_INPUT_PREFIX + '%s') % ix: si.pvalue for (ix, si) in enumerate(side_inputs)} return dict(main_inputs, **named_side_inputs) def _named_outputs(self, outputs): # type: (dict[object, pvalue.PCollection]) -> dict[str, pvalue.PCollection] """Returns the dictionary of named outputs as they should be named in the beam proto. """ # TODO(BEAM-1833): Push names up into the sdk construction. return { str(tag): output for (tag, output) in outputs.items() if isinstance(output, pvalue.PCollection) } _known_urns = {} # type: dict[str, tuple[Optional[type], ConstructorFn]] @classmethod @overload def register_urn( cls, urn, # type: str parameter_type, # type: type[T] ): # type: (...) -> Callable[[Union[type, Callable[[beam_runner_api_pb2.PTransform, T, PipelineContext], Any]]], Callable[[T, PipelineContext], Any]] pass @classmethod @overload def register_urn( cls, urn, # type: str parameter_type, # type: None ): # type: (...) -> Callable[[Union[type, Callable[[beam_runner_api_pb2.PTransform, bytes, PipelineContext], Any]]], Callable[[bytes, PipelineContext], Any]] pass @classmethod @overload def register_urn( cls, urn, # type: str parameter_type, # type: type[T] constructor # type: Callable[[beam_runner_api_pb2.PTransform, T, PipelineContext], Any] ): # type: (...) -> None pass @classmethod @overload def register_urn( cls, urn, # type: str parameter_type, # type: None constructor # type: Callable[[beam_runner_api_pb2.PTransform, bytes, PipelineContext], Any] ): # type: (...) -> None pass @classmethod def register_urn(cls, urn, parameter_type, constructor=None): def register(constructor): if isinstance(constructor, type): constructor.from_runner_api_parameter = register( constructor.from_runner_api_parameter) else: cls._known_urns[urn] = parameter_type, constructor return constructor if constructor: # Used as a statement. register(constructor) else: # Used as a decorator. return register def to_runner_api(self, context, has_parts=False, **extra_kwargs): # type: (PipelineContext, bool, Any) -> beam_runner_api_pb2.FunctionSpec from apache_beam.portability.api import beam_runner_api_pb2 # typing: only ParDo supports extra_kwargs urn, typed_param = self.to_runner_api_parameter(context, **extra_kwargs) if urn == python_urns.GENERIC_COMPOSITE_TRANSFORM and not has_parts: # TODO(https://github.com/apache/beam/issues/18713): Remove this fallback. urn, typed_param = self.to_runner_api_pickled(context) return beam_runner_api_pb2.FunctionSpec( urn=urn, payload=typed_param.SerializeToString() if isinstance( typed_param, message.Message) else typed_param.encode('utf-8') if isinstance(typed_param, str) else typed_param) @classmethod def from_runner_api( cls, proto, # type: Optional[beam_runner_api_pb2.PTransform] context # type: PipelineContext ): # type: (...) -> Optional[PTransform] if proto is None or proto.spec is None or not proto.spec.urn: return None parameter_type, constructor = cls._known_urns[proto.spec.urn] return constructor( proto, proto_utils.parse_Bytes(proto.spec.payload, parameter_type), context) def to_runner_api_parameter( self, unused_context # type: PipelineContext ): # type: (...) -> tuple[str, Optional[Union[message.Message, bytes, str]]] # The payload here is just to ease debugging. return ( python_urns.GENERIC_COMPOSITE_TRANSFORM, getattr(self, '_fn_api_payload', str(self))) def to_runner_api_pickled(self, context): # type: (PipelineContext) -> tuple[str, bytes] return ( python_urns.PICKLED_TRANSFORM, pickler.dumps( self, enable_best_effort_determinism=context. enable_best_effort_deterministic_pickling, enable_stable_code_identifier_pickling=context. enable_stable_code_identifier_pickling, ), ) def runner_api_requires_keyed_input(self): return False def _add_type_constraint_from_consumer(self, full_label, input_type_hints): # type: (str, tuple[str, Any]) -> None """Adds a consumer transform's input type hints to our output type constraints, which is used during performance runtime type-checking. """ pass @PTransform.register_urn(python_urns.GENERIC_COMPOSITE_TRANSFORM, None) def _create_transform(unused_ptransform, payload, unused_context): empty_transform = PTransform() empty_transform._fn_api_payload = payload return empty_transform @PTransform.register_urn(python_urns.PICKLED_TRANSFORM, None) def _unpickle_transform(unused_ptransform, pickled_bytes, unused_context): return pickler.loads(pickled_bytes) class _ChainedPTransform(PTransform): def __init__(self, *parts): # type: (*PTransform) -> None super().__init__(label=self._chain_label(parts)) self._parts = parts def _chain_label(self, parts): return '|'.join(p.label for p in parts) def __or__(self, right): if isinstance(right, PTransform): # Create a flat list rather than a nested tree of composite # transforms for better monitoring, etc. return _ChainedPTransform(*(self._parts + (right, ))) return NotImplemented def expand(self, pval): return reduce(operator.or_, self._parts, pval) class PTransformWithSideInputs(PTransform): """A superclass for any :class:`PTransform` (e.g. :func:`~apache_beam.transforms.core.FlatMap` or :class:`~apache_beam.transforms.core.CombineFn`) invoking user code. :class:`PTransform` s like :func:`~apache_beam.transforms.core.FlatMap` invoke user-supplied code in some kind of package (e.g. a :class:`~apache_beam.transforms.core.DoFn`) and optionally provide arguments and side inputs to that code. This internal-use-only class contains common functionality for :class:`PTransform` s that fit this model. """ def __init__(self, fn, *args, **kwargs): # type: (WithTypeHints, *Any, **Any) -> None if isinstance(fn, type) and issubclass(fn, WithTypeHints): # Don't treat Fn class objects as callables. raise ValueError('Use %s() not %s.' % (fn.__name__, fn.__name__)) self.fn = self.make_fn(fn, bool(args or kwargs)) # Now that we figure out the label, initialize the super-class. super().__init__() if (any(isinstance(v, pvalue.PCollection) for v in args) or any(isinstance(v, pvalue.PCollection) for v in kwargs.values())): raise error.SideInputError( 'PCollection used directly as side input argument. Specify ' 'AsIter(pcollection) or AsSingleton(pcollection) to indicate how the ' 'PCollection is to be used.') self.args, self.kwargs, self.side_inputs = util.remove_objects_from_args( args, kwargs, pvalue.AsSideInput) self.raw_side_inputs = args, kwargs # Prevent name collisions with fns of the form ' at ...>' self._cached_fn = self.fn # Ensure fn and side inputs are picklable for remote execution. try: self.fn = pickler.roundtrip(self.fn) except (RuntimeError, TypeError, Exception) as e: raise RuntimeError( 'Unable to pickle fn %s: %s. ' 'User code must be serializable (picklable) for distributed ' 'execution. This usually happens when lambdas or closures capture ' 'non-serializable objects like file handles, database connections, ' 'or thread locks. Try: (1) using module-level functions instead of ' 'lambdas, (2) initializing resources in setup() methods, ' '(3) checking what your closure captures.' % (self.fn, e)) from e self.args = pickler.roundtrip(self.args) self.kwargs = pickler.roundtrip(self.kwargs) # For type hints, because loads(dumps(class)) != class. self.fn = self._cached_fn def with_input_types( self, input_type_hint, *side_inputs_arg_hints, **side_input_kwarg_hints): """Annotates the types of main inputs and side inputs for the PTransform. Args: input_type_hint: An instance of an allowed built-in type, a custom class, or an instance of a typehints.TypeConstraint. *side_inputs_arg_hints: A variable length argument composed of of an allowed built-in type, a custom class, or a typehints.TypeConstraint. **side_input_kwarg_hints: A dictionary argument composed of of an allowed built-in type, a custom class, or a typehints.TypeConstraint. Example of annotating the types of side-inputs:: FlatMap().with_input_types(int, int, bool) Raises: :class:`TypeError`: If **type_hint** is not a valid type-hint. See :func:`~apache_beam.typehints.typehints.validate_composite_type_param` for further details. Returns: :class:`PTransform`: A reference to the instance of this particular :class:`PTransform` object. This allows chaining type-hinting related methods. """ super().with_input_types(input_type_hint) side_inputs_arg_hints = native_type_compatibility.convert_to_beam_types( side_inputs_arg_hints) side_input_kwarg_hints = native_type_compatibility.convert_to_beam_types( side_input_kwarg_hints) for si in side_inputs_arg_hints: validate_composite_type_param(si, 'Type hints for a PTransform') for si in side_input_kwarg_hints.values(): validate_composite_type_param(si, 'Type hints for a PTransform') self.side_inputs_types = side_inputs_arg_hints return WithTypeHints.with_input_types( self, input_type_hint, *side_inputs_arg_hints, **side_input_kwarg_hints) def type_check_inputs(self, pvalueish): type_hints = self.get_type_hints() input_types = type_hints.input_types if input_types: args, kwargs = self.raw_side_inputs def element_type(side_input): if isinstance(side_input, pvalue.AsSideInput): return side_input.element_type return instance_to_type(side_input) arg_types = [pvalueish.element_type] + [element_type(v) for v in args] kwargs_types = {k: element_type(v) for (k, v) in kwargs.items()} argspec_fn = self._process_argspec_fn() bindings = getcallargs_forhints(argspec_fn, *arg_types, **kwargs_types) hints = getcallargs_forhints( argspec_fn, *input_types[0], **input_types[1]) # First check the main input. arg_hints = iter(hints.items()) element_arg, element_hint = next(arg_hints) if not typehints.is_consistent_with( bindings.get(element_arg, typehints.Any), element_hint): transform_nest_level = self.label.count("/") split_producer_label = pvalueish.producer.full_label.split("/") producer_label = "/".join( split_producer_label[:transform_nest_level + 1]) raise TypeCheckError( f"The transform '{self.label}' requires " f"PCollections of type '{element_hint}' " f"but was applied to a PCollection of type" f" '{bindings[element_arg]}' " f"(produced by the transform '{producer_label}'). ") # Now check the side inputs. for arg, hint in arg_hints: if arg.startswith('__unknown__'): continue if hint is None: continue if not typehints.is_consistent_with(bindings.get(arg, typehints.Any), hint): raise TypeCheckError( 'Type hint violation for \'{label}\': requires {hint} but got ' '{actual_type} for {arg}\nFull type hint:\n{debug_str}'.format( label=self.label, hint=hint, actual_type=bindings[arg], arg=arg, debug_str=type_hints.debug_str())) def _process_argspec_fn(self): """Returns an argspec of the function actually consuming the data. """ raise NotImplementedError def make_fn(self, fn, has_side_inputs): # TODO(silviuc): Add comment describing that this is meant to be overriden # by methods detecting callables and wrapping them in DoFns. return fn def default_label(self): return '%s(%s)' % (self.__class__.__name__, self.fn.default_label()) class _PTransformFnPTransform(PTransform): """A class wrapper for a function-based transform.""" def __init__(self, fn, *args, **kwargs): super().__init__() self._fn = fn self._args = args self._kwargs = kwargs self._use_backwards_compatible_label = True def display_data(self): res = { 'fn': ( self._fn.__name__ if hasattr(self._fn, '__name__') else self._fn.__class__), 'args': DisplayDataItem(str(self._args)).drop_if_default('()'), 'kwargs': DisplayDataItem(str(self._kwargs)).drop_if_default('{}') } return res def expand(self, pcoll): # Since the PTransform will be implemented entirely as a function # (once called), we need to pass through any type-hinting information that # may have been annotated via the .with_input_types() and # .with_output_types() methods. kwargs = dict(self._kwargs) args = tuple(self._args) # TODO(BEAM-5878) Support keyword-only arguments. try: if 'type_hints' in get_signature(self._fn).parameters: args = (self.get_type_hints(), ) + args except TypeError: # Might not be a function. pass return self._fn(pcoll, *args, **kwargs) def set_options(self, options): self._use_backwards_compatible_label = options.is_compat_version_prior_to( '2.68.0') def default_label(self) -> str: # Attempt to give a reasonable name to this transform. # We want it to be reasonably unique, but also not sensitive to # irrelevent parameters to minimize pipeline-to-pipeline variance. # For now, use only the first argument (if any), iff it would not make # the name unwieldy. if self._args: first_arg_string = label_from_callable(self._args[0]) if (self._use_backwards_compatible_label or not isinstance(first_arg_string, str) or len(first_arg_string) <= 19): suffix = '(%s)' % first_arg_string else: suffix = ('(%s...%s)' % (first_arg_string[:10], first_arg_string[-6:])).replace( '\n', ' ') else: suffix = '' return label_from_callable(self._fn) + suffix def ptransform_fn(fn): # type: (Callable) -> Callable[..., _PTransformFnPTransform] """A decorator for a function-based PTransform. Args: fn: A function implementing a custom PTransform. Returns: A CallablePTransform instance wrapping the function-based PTransform. This wrapper provides an alternative, simpler way to define a PTransform. The standard method is to subclass from PTransform and override the expand() method. An equivalent effect can be obtained by defining a function that accepts an input PCollection and additional optional arguments and returns a resulting PCollection. For example:: @ptransform_fn @beam.typehints.with_input_types(..) @beam.typehints.with_output_types(..) def CustomMapper(pcoll, mapfn): return pcoll | ParDo(mapfn) The equivalent approach using PTransform subclassing:: @beam.typehints.with_input_types(..) @beam.typehints.with_output_types(..) class CustomMapper(PTransform): def __init__(self, mapfn): super().__init__() self.mapfn = mapfn def expand(self, pcoll): return pcoll | ParDo(self.mapfn) With either method the custom PTransform can be used in pipelines as if it were one of the "native" PTransforms:: result_pcoll = input_pcoll | 'Label' >> CustomMapper(somefn) Note that for both solutions the underlying implementation of the pipe operator (i.e., `|`) will inject the pcoll argument in its proper place (first argument if no label was specified and second argument otherwise). Type hint support needs to be enabled via the --type_check_additional=ptransform_fn flag in Beam 2. If CustomMapper is a Cython function, you can still specify input and output types provided the decorators appear before @ptransform_fn. """ # TODO(robertwb): Consider removing staticmethod to allow for self parameter. @wraps(fn) def callable_ptransform_factory(*args, **kwargs): res = _PTransformFnPTransform(fn, *args, **kwargs) if ptransform_fn_typehints_enabled: # Apply type hints applied before or after the ptransform_fn decorator, # falling back on PTransform defaults. # If the @with_{input,output}_types decorator comes before ptransform_fn, # the type hints get applied to this function. If it comes after they will # get applied to fn, and @wraps will copy the _type_hints attribute to # this function. type_hints = get_type_hints(callable_ptransform_factory) res._set_type_hints(type_hints.with_defaults(res.get_type_hints())) _LOGGER.debug( 'type hints for %s: %s', res.default_label(), res.get_type_hints()) return res # The signature of this PTransform constructor is that of fn minus the first # argument (which is where the pvalue is passed during expand). try: callable_ptransform_factory.__signature__ = inspect.signature( # type: ignore functools.partial(fn, None)) except Exception: # Sometimes we can't get the original signature. pass return callable_ptransform_factory def label_from_callable(fn): if hasattr(fn, 'default_label'): return fn.default_label() elif hasattr(fn, '__name__'): if fn.__name__ == '': return '' % ( os.path.basename(fn.__code__.co_filename), fn.__code__.co_firstlineno) return fn.__name__ return str(fn) class _NamedPTransform(PTransform): def __init__(self, transform, label): super().__init__(label) self.transform = transform def __ror__(self, pvalueish, _unused=None): return self.transform.__ror__(pvalueish, self.label) def expand(self, pvalue): raise RuntimeError("Should never be expanded directly.") def annotations(self): return self.transform.annotations() def __rrshift__(self, label): return _NamedPTransform(self.transform, label) def with_resource_hints(self, **kwargs): self.transform.with_resource_hints(**kwargs) return self def __getattr__(self, attr): transform_attr = getattr(self.transform, attr) if callable(transform_attr): @wraps(transform_attr) def wrapper(*args, **kwargs): result = transform_attr(*args, **kwargs) if isinstance(result, PTransform): return _NamedPTransform(result, self.label) else: return result return wrapper else: return transform_attr def __setattr__(self, attr, value): if attr == 'annotations': self.transform.annotations = value else: super().__setattr__(attr, value) # Defined here to avoid circular import issues for Beam library transforms. def annotate_yaml(constructor): """Causes instances of this transform to be annotated with their yaml syntax. Should only be used for transforms that are fully defined by their constructor arguments. """ @wraps(constructor) def wrapper(*args, **kwargs): transform = constructor(*args, **kwargs) fully_qualified_name = ( f'{constructor.__module__}.{constructor.__qualname__}') try: imported_constructor = ( python_callable.PythonCallableWithSource. load_from_fully_qualified_name(fully_qualified_name)) if imported_constructor != wrapper: raise ImportError('Different object.') except ImportError: warnings.warn(f'Cannot import {constructor} as {fully_qualified_name}.') return transform try: config = json.dumps({ 'constructor': fully_qualified_name, 'args': args, 'kwargs': kwargs, }) except TypeError as exn: warnings.warn( f'Cannot serialize arguments for {constructor} as json: {exn}') return transform original_annotations = transform.annotations transform.annotations = lambda: { **original_annotations(), # These override whatever may have been provided earlier. # The outermost call is expected to be the most specific. 'yaml_provider': 'python', 'yaml_type': 'PyTransform', 'yaml_args': config, } return transform return wrapper ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/ptransform_test.py0000644000175100017510000033554715206113033023701 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the PTransform and descendants.""" # pytype: skip-file import collections import operator import os import pickle import random import re import sys import typing import unittest from functools import reduce from typing import Optional from unittest.mock import patch import hamcrest as hc import numpy as np import pytest from parameterized import param from parameterized import parameterized from parameterized import parameterized_class import apache_beam as beam import apache_beam.transforms.combiners as combine from apache_beam import pvalue from apache_beam import typehints from apache_beam.coders import coders_test_common from apache_beam.io.iobase import Read from apache_beam.metrics import Metrics from apache_beam.metrics.metric import MetricsFilter from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import StreamingOptions from apache_beam.options.pipeline_options import TypeOptions from apache_beam.portability import common_urns from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import SortLists from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import WindowInto from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.transforms.display import DisplayData from apache_beam.transforms.display import DisplayDataItem from apache_beam.transforms.ptransform import PTransform from apache_beam.transforms.trigger import AccumulationMode from apache_beam.transforms.trigger import AfterProcessingTime from apache_beam.transforms.trigger import _AfterSynchronizedProcessingTime from apache_beam.transforms.window import TimestampedValue from apache_beam.typehints import with_input_types from apache_beam.typehints import with_output_types from apache_beam.typehints.typehints_test import TypeHintTestCase from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.windowed_value import WindowedValue # Disable frequent lint warning due to pipe operator for chaining transforms. # pylint: disable=expression-not-assigned class PTransformTest(unittest.TestCase): def assertStartswith(self, msg, prefix): self.assertTrue( msg.startswith(prefix), '"%s" does not start with "%s"' % (msg, prefix)) def test_str(self): self.assertEqual( '', str(PTransform())) pa = TestPipeline() res = pa | 'ALabel' >> beam.Impulse() self.assertEqual('AppliedPTransform(ALabel, Impulse)', str(res.producer)) pc = TestPipeline() res = pc | beam.Impulse() inputs_tr = res.producer.transform inputs_tr.inputs = ('ci', ) self.assertEqual( "", str(inputs_tr)) pd = TestPipeline() res = pd | beam.Impulse() side_tr = res.producer.transform side_tr.side_inputs = (4, ) self.assertEqual( '', str(side_tr)) inputs_tr.side_inputs = ('cs', ) self.assertEqual( """""", str(inputs_tr)) def test_named_annotations(self): t = beam.Impulse() t.annotations = lambda: {'test': 'value'} named_t = 'Name' >> t self.assertEqual(named_t.annotations(), {'test': 'value'}) original_annotations = named_t.annotations() named_t.annotations = lambda: {'another': 'value', **original_annotations} # Verify this is reflected on the original transform, # which is what gets used in apply. self.assertEqual(t.annotations(), {'test': 'value', 'another': 'value'}) def test_do_with_do_fn(self): class AddNDoFn(beam.DoFn): def process(self, element, addon): return [element + addon] with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3]) result = pcoll | 'Do' >> beam.ParDo(AddNDoFn(), 10) assert_that(result, equal_to([11, 12, 13])) def test_do_with_unconstructed_do_fn(self): class MyDoFn(beam.DoFn): def process(self): pass with self.assertRaises(ValueError): with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3]) pcoll | 'Do' >> beam.ParDo(MyDoFn) # Note the lack of ()'s def test_do_with_callable(self): with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3]) result = pcoll | 'Do' >> beam.FlatMap(lambda x, addon: [x + addon], 10) assert_that(result, equal_to([11, 12, 13])) def test_do_with_side_input_as_arg(self): with TestPipeline() as pipeline: side = pipeline | 'Side' >> beam.Create([10]) pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3]) result = pcoll | 'Do' >> beam.FlatMap( lambda x, addon: [x + addon], pvalue.AsSingleton(side)) assert_that(result, equal_to([11, 12, 13])) def test_do_with_side_input_as_keyword_arg(self): with TestPipeline() as pipeline: side = pipeline | 'Side' >> beam.Create([10]) pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3]) result = pcoll | 'Do' >> beam.FlatMap( lambda x, addon: [x + addon], addon=pvalue.AsSingleton(side)) assert_that(result, equal_to([11, 12, 13])) def test_callable_non_serializable_error_message(self): class NonSerializable: def __getstate__(self): raise RuntimeError('nope') bad = NonSerializable() with self.assertRaises(RuntimeError) as context: _ = beam.Map(lambda x: bad) message = str(context.exception) self.assertIn('Unable to pickle fn', message) self.assertIn( 'User code must be serializable (picklable) for distributed execution.', message) self.assertIn('non-serializable objects like file handles', message) self.assertIn( 'Try: (1) using module-level functions instead of lambdas', message) def test_do_with_do_fn_returning_string_raises_warning(self): ex_details = r'.*Returning a str from a ParDo or FlatMap is discouraged.' with self.assertRaisesRegex(Exception, ex_details): with TestPipeline() as pipeline: pipeline._options.view_as(TypeOptions).runtime_type_check = True pcoll = pipeline | 'Start' >> beam.Create(['2', '9', '3']) pcoll | 'Do' >> beam.FlatMap(lambda x: x + '1') # Since the DoFn directly returns a string we should get an # error warning us when the pipeliene runs. def test_do_with_do_fn_returning_dict_raises_warning(self): ex_details = r'.*Returning a dict from a ParDo or FlatMap is discouraged.' with self.assertRaisesRegex(Exception, ex_details): with TestPipeline() as pipeline: pipeline._options.view_as(TypeOptions).runtime_type_check = True pcoll = pipeline | 'Start' >> beam.Create(['2', '9', '3']) pcoll | 'Do' >> beam.FlatMap(lambda x: {x: '1'}) # Since the DoFn directly returns a dict we should get an error warning # us when the pipeliene runs. def test_do_with_multiple_outputs_maintains_unique_name(self): with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3]) r1 = pcoll | 'A' >> beam.FlatMap(lambda x: [x + 1]).with_outputs(main='m') r2 = pcoll | 'B' >> beam.FlatMap(lambda x: [x + 2]).with_outputs(main='m') assert_that(r1.m, equal_to([2, 3, 4]), label='r1') assert_that(r2.m, equal_to([3, 4, 5]), label='r2') @pytest.mark.it_validatesrunner def test_impulse(self): with TestPipeline() as pipeline: result = pipeline | beam.Impulse() | beam.Map(lambda _: 0) assert_that(result, equal_to([0])) # TODO(BEAM-3544): Disable this test in streaming temporarily. # Remove sickbay-streaming tag after it's resolved. @pytest.mark.no_sickbay_streaming @pytest.mark.it_validatesrunner def test_read_metrics(self): from apache_beam.io.utils import CountingSource class CounterDoFn(beam.DoFn): def __init__(self): # This counter is unused. self.received_records = Metrics.counter( self.__class__, 'receivedRecords') def process(self, element): self.received_records.inc() pipeline = TestPipeline() (pipeline | Read(CountingSource(100)) | beam.ParDo(CounterDoFn())) res = pipeline.run() res.wait_until_finish() # This counter is defined in utils.CountingSource. metric_results = res.metrics().query( MetricsFilter().with_name('recordsRead')) outputs_counter = metric_results['counters'][0] msg = outputs_counter.key.step cont = 'SDFBoundedSourceReader' self.assertTrue(cont in msg, '"%s" does not contain "%s"' % (msg, cont)) self.assertEqual(outputs_counter.key.metric.name, 'recordsRead') self.assertEqual(outputs_counter.committed, 100) @pytest.mark.it_validatesrunner def test_par_do_with_multiple_outputs_and_using_yield(self): class SomeDoFn(beam.DoFn): """A custom DoFn using yield.""" def process(self, element): yield element if element % 2 == 0: yield pvalue.TaggedOutput('even', element) else: yield pvalue.TaggedOutput('odd', element) with TestPipeline() as pipeline: nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4]) results = nums | 'ClassifyNumbers' >> beam.ParDo(SomeDoFn()).with_outputs( 'odd', 'even', main='main') assert_that(results.main, equal_to([1, 2, 3, 4])) assert_that(results.odd, equal_to([1, 3]), label='assert:odd') assert_that(results.even, equal_to([2, 4]), label='assert:even') @pytest.mark.it_validatesrunner def test_par_do_with_multiple_outputs_and_using_return(self): def some_fn(v): if v % 2 == 0: return [v, pvalue.TaggedOutput('even', v)] return [v, pvalue.TaggedOutput('odd', v)] with TestPipeline() as pipeline: nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4]) results = nums | 'ClassifyNumbers' >> beam.FlatMap(some_fn).with_outputs( 'odd', 'even', main='main') assert_that(results.main, equal_to([1, 2, 3, 4])) assert_that(results.odd, equal_to([1, 3]), label='assert:odd') assert_that(results.even, equal_to([2, 4]), label='assert:even') @pytest.mark.it_validatesrunner def test_undeclared_outputs(self): with TestPipeline() as pipeline: nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4]) results = nums | 'ClassifyNumbers' >> beam.FlatMap( lambda x: [ x, pvalue.TaggedOutput('even' if x % 2 == 0 else 'odd', x), pvalue .TaggedOutput('extra', x) ]).with_outputs() assert_that(results[None], equal_to([1, 2, 3, 4])) assert_that(results.odd, equal_to([1, 3]), label='assert:odd') assert_that(results.even, equal_to([2, 4]), label='assert:even') @pytest.mark.it_validatesrunner def test_multiple_empty_outputs(self): with TestPipeline() as pipeline: nums = pipeline | 'Some Numbers' >> beam.Create([1, 3, 5]) results = nums | 'ClassifyNumbers' >> beam.FlatMap( lambda x: [x, pvalue.TaggedOutput('even' if x % 2 == 0 else 'odd', x)]).with_outputs() assert_that(results[None], equal_to([1, 3, 5])) assert_that(results.odd, equal_to([1, 3, 5]), label='assert:odd') assert_that(results.even, equal_to([]), label='assert:even') def test_do_requires_do_fn_returning_iterable(self): # This function is incorrect because it returns an object that isn't an # iterable. def incorrect_par_do_fn(x): return x + 5 ex_details = r'.*FlatMap and ParDo must return an iterable.' with self.assertRaisesRegex(Exception, ex_details): with TestPipeline() as pipeline: pipeline._options.view_as(TypeOptions).runtime_type_check = True pcoll = pipeline | 'Start' >> beam.Create([2, 9, 3]) pcoll | 'Do' >> beam.FlatMap(incorrect_par_do_fn) # It's a requirement that all user-defined functions to a ParDo return # an iterable. def test_do_fn_with_finish(self): class MyDoFn(beam.DoFn): def process(self, element): pass def finish_bundle(self): yield WindowedValue('finish', -1, [window.GlobalWindow()]) with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3]) result = pcoll | 'Do' >> beam.ParDo(MyDoFn()) # May have many bundles, but each has a start and finish. def matcher(): def match(actual): equal_to(['finish'])(list(set(actual))) equal_to([1])([actual.count('finish')]) return match assert_that(result, matcher()) def test_do_fn_with_windowing_in_finish_bundle(self): windowfn = window.FixedWindows(2) class MyDoFn(beam.DoFn): def process(self, element): yield TimestampedValue('process' + str(element), 5) def finish_bundle(self): yield WindowedValue('finish', 1, [windowfn]) with TestPipeline() as pipeline: result = ( pipeline | 'Start' >> beam.Create([1]) | beam.ParDo(MyDoFn()) | WindowInto(windowfn) | 'create tuple' >> beam.Map( lambda v, t=beam.DoFn.TimestampParam, w=beam.DoFn.WindowParam: (v, t, w.start, w.end))) expected_process = [ ('process1', Timestamp(5), Timestamp(4), Timestamp(6)) ] expected_finish = [('finish', Timestamp(1), Timestamp(0), Timestamp(2))] assert_that(result, equal_to(expected_process + expected_finish)) def test_do_fn_with_start(self): class MyDoFn(beam.DoFn): def __init__(self): self.state = 'init' def start_bundle(self): self.state = 'started' def process(self, element): if self.state == 'started': yield 'started' self.state = 'process' with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3]) result = pcoll | 'Do' >> beam.ParDo(MyDoFn()) # May have many bundles, but each has a start and finish. def matcher(): def match(actual): equal_to(['started'])(list(set(actual))) equal_to([1])([actual.count('started')]) return match assert_that(result, matcher()) def test_do_fn_with_start_error(self): class MyDoFn(beam.DoFn): def start_bundle(self): return [1] def process(self, element): pass with self.assertRaises(RuntimeError): with TestPipeline() as p: p | 'Start' >> beam.Create([1, 2, 3]) | 'Do' >> beam.ParDo(MyDoFn()) def test_map_builtin(self): with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([[1, 2], [1], [1, 2, 3]]) result = pcoll | beam.Map(len) assert_that(result, equal_to([1, 2, 3])) def test_flatmap_builtin(self): with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([ [np.array([1, 2, 3])] * 3, [np.array([5, 4, 3]), np.array([5, 6, 7])] ]) result = pcoll | beam.FlatMap(sum) assert_that(result, equal_to([3, 6, 9, 10, 10, 10])) def test_filter_builtin(self): with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([[], [2], [], [4]]) result = pcoll | 'Filter' >> beam.Filter(len) assert_that(result, equal_to([[2], [4]])) def test_filter(self): with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3, 4]) result = pcoll | 'Filter' >> beam.Filter(lambda x: x % 2 == 0) assert_that(result, equal_to([2, 4])) class _MeanCombineFn(beam.CombineFn): def create_accumulator(self): return (0, 0) def add_input(self, sum_count, element): (sum_, count) = sum_count return sum_ + element, count + 1 def merge_accumulators(self, accumulators): sums, counts = zip(*accumulators) return sum(sums), sum(counts) def extract_output(self, sum_count): (sum_, count) = sum_count if not count: return float('nan') return sum_ / float(count) def test_combine_with_combine_fn(self): vals = [1, 2, 3, 4, 5, 6, 7] with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create(vals) result = pcoll | 'Mean' >> beam.CombineGlobally(self._MeanCombineFn()) assert_that(result, equal_to([sum(vals) // len(vals)])) def test_combine_with_callable(self): vals = [1, 2, 3, 4, 5, 6, 7] with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create(vals) result = pcoll | beam.CombineGlobally(sum) assert_that(result, equal_to([sum(vals)])) def test_combine_with_side_input_as_arg(self): values = [1, 2, 3, 4, 5, 6, 7] with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create(values) divisor = pipeline | 'Divisor' >> beam.Create([2]) result = pcoll | 'Max' >> beam.CombineGlobally( # Multiples of divisor only. lambda vals, d: max(v for v in vals if v % d == 0), pvalue.AsSingleton(divisor)).without_defaults() filt_vals = [v for v in values if v % 2 == 0] assert_that(result, equal_to([max(filt_vals)])) def test_combine_per_key_with_combine_fn(self): vals_1 = [1, 2, 3, 4, 5, 6, 7] vals_2 = [2, 4, 6, 8, 10, 12, 14] with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create( ([('a', x) for x in vals_1] + [('b', x) for x in vals_2])) result = pcoll | 'Mean' >> beam.CombinePerKey(self._MeanCombineFn()) assert_that( result, equal_to([('a', sum(vals_1) // len(vals_1)), ('b', sum(vals_2) // len(vals_2))])) def test_combine_per_key_with_callable(self): vals_1 = [1, 2, 3, 4, 5, 6, 7] vals_2 = [2, 4, 6, 8, 10, 12, 14] with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create( ([('a', x) for x in vals_1] + [('b', x) for x in vals_2])) result = pcoll | beam.CombinePerKey(sum) assert_that(result, equal_to([('a', sum(vals_1)), ('b', sum(vals_2))])) def test_combine_per_key_with_side_input_as_arg(self): vals_1 = [1, 2, 3, 4, 5, 6, 7] vals_2 = [2, 4, 6, 8, 10, 12, 14] with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create( ([('a', x) for x in vals_1] + [('b', x) for x in vals_2])) divisor = pipeline | 'Divisor' >> beam.Create([2]) result = pcoll | beam.CombinePerKey( lambda vals, d: max(v for v in vals if v % d == 0), pvalue.AsSingleton(divisor)) # Multiples of divisor only. m_1 = max(v for v in vals_1 if v % 2 == 0) m_2 = max(v for v in vals_2 if v % 2 == 0) assert_that(result, equal_to([('a', m_1), ('b', m_2)])) def test_group_by_key(self): with TestPipeline() as pipeline: pcoll = pipeline | 'start' >> beam.Create([(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 3)]) result = pcoll | 'Group' >> beam.GroupByKey() | SortLists assert_that(result, equal_to([(1, [1, 2, 3]), (2, [1, 2]), (3, [1])])) def test_group_by_key_unbounded_global_default_trigger(self): test_options = PipelineOptions() test_options.view_as(TypeOptions).allow_unsafe_triggers = False with self.assertRaisesRegex( ValueError, 'GroupByKey cannot be applied to an unbounded PCollection with ' + 'global windowing and a default trigger'): with TestPipeline(options=test_options) as pipeline: pipeline | TestStream() | beam.GroupByKey() def test_group_by_key_trigger(self): options = PipelineOptions(['--allow_unsafe_triggers']) options.view_as(StandardOptions).streaming = True with TestPipeline(runner='BundleBasedDirectRunner', options=options) as pipeline: pcoll = pipeline | 'Start' >> beam.Create([(0, 0)]) triggered = pcoll | 'Trigger' >> beam.WindowInto( window.GlobalWindows(), trigger=AfterProcessingTime(1), accumulation_mode=AccumulationMode.DISCARDING) output = triggered | 'Gbk' >> beam.GroupByKey() self.assertTrue( isinstance( output.windowing.triggerfn, _AfterSynchronizedProcessingTime)) def test_group_by_key_unsafe_trigger(self): test_options = PipelineOptions() test_options.view_as(TypeOptions).allow_unsafe_triggers = False with self.assertRaisesRegex(ValueError, 'Unsafe trigger'): with TestPipeline(options=test_options) as pipeline: _ = ( pipeline | beam.Create([(None, None)]) | WindowInto( window.GlobalWindows(), trigger=trigger.AfterCount(5), accumulation_mode=trigger.AccumulationMode.ACCUMULATING) | beam.GroupByKey()) def test_group_by_key_allow_unsafe_triggers(self): test_options = PipelineOptions(flags=['--allow_unsafe_triggers']) with TestPipeline(options=test_options) as pipeline: pcoll = ( pipeline | beam.Create([(1, 1), (1, 2), (1, 3), (1, 4)]) | WindowInto( window.GlobalWindows(), trigger=trigger.AfterCount(4), accumulation_mode=trigger.AccumulationMode.ACCUMULATING) | beam.GroupByKey()) assert_that(pcoll, equal_to([(1, [1, 2, 3, 4])])) def test_group_by_key_reiteration(self): class MyDoFn(beam.DoFn): def process(self, gbk_result): key, value_list = gbk_result sum_val = 0 # Iterate the GBK result for multiple times. for _ in range(0, 17): sum_val += sum(value_list) return [(key, sum_val)] with TestPipeline() as pipeline: pcoll = pipeline | 'start' >> beam.Create([(1, 1), (1, 2), (1, 3), (1, 4)]) result = ( pcoll | 'Group' >> beam.GroupByKey() | 'Reiteration-Sum' >> beam.ParDo(MyDoFn())) assert_that(result, equal_to([(1, 170)])) def test_group_by_key_deterministic_coder(self): # pylint: disable=global-variable-not-assigned global MyObject # for pickling of the class instance class MyObject: def __init__(self, value): self.value = value def __eq__(self, other): return self.value == other.value def __hash__(self): return hash(self.value) class MyObjectCoder(beam.coders.Coder): def encode(self, o): return pickle.dumps((o.value, random.random())) def decode(self, encoded): return MyObject(pickle.loads(encoded)[0]) def as_deterministic_coder(self, *args, **kwargs): return MydeterministicObjectCoder() def to_type_hint(self): return MyObject class MydeterministicObjectCoder(beam.coders.Coder): def encode(self, o): return pickle.dumps(o.value) def decode(self, encoded): return MyObject(pickle.loads(encoded)) def is_deterministic(self): return True beam.coders.registry.register_coder(MyObject, MyObjectCoder) with TestPipeline() as pipeline: pcoll = pipeline | beam.Create([(MyObject(k % 2), k) for k in range(10)]) grouped = pcoll | beam.GroupByKey() | beam.MapTuple( lambda k, vs: (k.value, sorted(vs))) combined = pcoll | beam.CombinePerKey(sum) | beam.MapTuple( lambda k, v: (k.value, v)) assert_that( grouped, equal_to([(0, [0, 2, 4, 6, 8]), (1, [1, 3, 5, 7, 9])]), 'CheckGrouped') assert_that(combined, equal_to([(0, 20), (1, 25)]), 'CheckCombined') def test_group_by_key_non_deterministic_coder(self): with self.assertRaisesRegex(Exception, r'deterministic'): with TestPipeline() as pipeline: _ = ( pipeline | beam.Create([(PickledObject(10), None)]) | beam.GroupByKey() | beam.MapTuple(lambda k, v: list(v))) def test_group_by_key_allow_non_deterministic_coder(self): with TestPipeline() as pipeline: # The GroupByKey below would fail without this option. pipeline._options.view_as( TypeOptions).allow_non_deterministic_key_coders = True grouped = ( pipeline | beam.Create([(PickledObject(10), None)]) | beam.GroupByKey() | beam.MapTuple(lambda k, v: list(v))) assert_that(grouped, equal_to([[None]])) def test_group_by_key_fake_deterministic_coder(self): fresh_registry = beam.coders.typecoders.CoderRegistry() with patch.object( beam.coders, 'registry', fresh_registry), patch.object( beam.coders.typecoders, 'registry', fresh_registry): with TestPipeline() as pipeline: # The GroupByKey below would fail without this registration. beam.coders.registry.register_fallback_coder( beam.coders.coders.FakeDeterministicFastPrimitivesCoder()) grouped = ( pipeline | beam.Create([(PickledObject(10), None)]) | beam.GroupByKey() | beam.MapTuple(lambda k, v: list(v))) assert_that(grouped, equal_to([[None]])) def test_partition_with_partition_fn(self): class SomePartitionFn(beam.PartitionFn): def partition_for(self, element, num_partitions, offset): return (element % 3) + offset with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8]) # Attempt nominal partition operation. partitions = pcoll | 'Part 1' >> beam.Partition(SomePartitionFn(), 4, 1) assert_that(partitions[0], equal_to([])) assert_that(partitions[1], equal_to([0, 3, 6]), label='p1') assert_that(partitions[2], equal_to([1, 4, 7]), label='p2') assert_that(partitions[3], equal_to([2, 5, 8]), label='p3') # Check that a bad partition label will yield an error. For the # DirectRunner, this error manifests as an exception. with self.assertRaises(Exception): with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8]) partitions = pcoll | beam.Partition(SomePartitionFn(), 4, 10000) def test_partition_with_callable(self): with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8]) partitions = ( pcoll | 'part' >> beam.Partition(lambda e, n, offset: (e % 3) + offset, 4, 1)) assert_that(partitions[0], equal_to([])) assert_that(partitions[1], equal_to([0, 3, 6]), label='p1') assert_that(partitions[2], equal_to([1, 4, 7]), label='p2') assert_that(partitions[3], equal_to([2, 5, 8]), label='p3') def test_partition_with_callable_and_side_input(self): with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8]) side_input = pipeline | 'Side Input' >> beam.Create([100, 1000]) partitions = ( pcoll | 'part' >> beam.Partition( lambda e, n, offset, si_list: ((e + len(si_list)) % 3) + offset, 4, 1, pvalue.AsList(side_input))) assert_that(partitions[0], equal_to([])) assert_that(partitions[1], equal_to([1, 4, 7]), label='p1') assert_that(partitions[2], equal_to([2, 5, 8]), label='p2') assert_that(partitions[3], equal_to([0, 3, 6]), label='p3') def test_partition_followed_by_flatten_and_groupbykey(self): """Regression test for an issue with how partitions are handled.""" with TestPipeline() as pipeline: contents = [('aa', 1), ('bb', 2), ('aa', 2)] created = pipeline | 'A' >> beam.Create(contents) partitioned = created | 'B' >> beam.Partition(lambda x, n: len(x) % n, 3) flattened = partitioned | 'C' >> beam.Flatten() grouped = flattened | 'D' >> beam.GroupByKey() | SortLists assert_that(grouped, equal_to([('aa', [1, 2]), ('bb', [2])])) @pytest.mark.it_validatesrunner def test_flatten_pcollections(self): with TestPipeline() as pipeline: pcoll_1 = pipeline | 'Start 1' >> beam.Create([0, 1, 2, 3]) pcoll_2 = pipeline | 'Start 2' >> beam.Create([4, 5, 6, 7]) result = (pcoll_1, pcoll_2) | 'Flatten' >> beam.Flatten() assert_that(result, equal_to([0, 1, 2, 3, 4, 5, 6, 7])) def test_flatten_no_pcollections(self): with TestPipeline() as pipeline: with self.assertRaises(ValueError): () | 'PipelineArgMissing' >> beam.Flatten() result = () | 'Empty' >> beam.Flatten(pipeline=pipeline) assert_that(result, equal_to([])) @pytest.mark.it_validatesrunner def test_flatten_one_single_pcollection(self): with TestPipeline() as pipeline: input = [0, 1, 2, 3] pcoll = pipeline | 'Input' >> beam.Create(input) result = (pcoll, ) | 'Single Flatten' >> beam.Flatten() assert_that(result, equal_to(input)) @parameterized.expand([ param(compat_version=None), param(compat_version="2.66.0"), ]) @pytest.mark.it_validatesrunner @pytest.mark.uses_dill def test_group_by_key_importable_special_types(self, compat_version): def generate(_): for _ in range(100): yield (coders_test_common.MyTypedNamedTuple(1, 'a'), 1) pipeline = TestPipeline(is_integration_test=True) if compat_version: pytest.importorskip("dill") pipeline.get_pipeline_options().view_as( StreamingOptions).update_compatibility_version = compat_version with pipeline as p: result = ( p | 'Create' >> beam.Create([i for i in range(100)]) | 'Generate' >> beam.ParDo(generate) | 'Reshuffle' >> beam.Reshuffle() | 'GBK' >> beam.GroupByKey()) assert_that( result, equal_to([( coders_test_common.MyTypedNamedTuple(1, 'a'), [1 for i in range(10000)])])) @pytest.mark.it_validatesrunner def test_group_by_key_dynamic_special_types(self): def create_dynamic_named_tuple(): return collections.namedtuple('DynamicNamedTuple', ['x', 'y']) dynamic_named_tuple = create_dynamic_named_tuple() # Standard FastPrimitivesCoder falls back to python PickleCoder which # cannot serialize dynamic types or types defined in __main__. Use # CloudPickleCoder as fallback coder for non-deterministic steps. class FastPrimitivesCoderV2(beam.coders.FastPrimitivesCoder): def __init__(self): super().__init__(fallback_coder=beam.coders.CloudpickleCoder()) beam.coders.typecoders.registry.register_coder( dynamic_named_tuple, FastPrimitivesCoderV2) def generate(_): for _ in range(100): yield (dynamic_named_tuple(1, 'a'), 1) pipeline = TestPipeline(is_integration_test=True) with pipeline as p: result = ( p | 'Create' >> beam.Create([i for i in range(100)]) | 'Reshuffle' >> beam.Reshuffle() | 'Generate' >> beam.ParDo(generate).with_output_types( tuple[dynamic_named_tuple, int]) | 'GBK' >> beam.GroupByKey() | 'Count Elements' >> beam.Map(lambda x: len(x[1]))) assert_that(result, equal_to([10000])) # TODO(https://github.com/apache/beam/issues/20067): Does not work in # streaming mode on Dataflow. @pytest.mark.no_sickbay_streaming @pytest.mark.it_validatesrunner def test_flatten_same_pcollections(self): with TestPipeline() as pipeline: pc = pipeline | beam.Create(['a', 'b']) assert_that((pc, pc, pc) | beam.Flatten(), equal_to(['a', 'b'] * 3)) def test_flatten_pcollections_in_iterable(self): with TestPipeline() as pipeline: pcoll_1 = pipeline | 'Start 1' >> beam.Create([0, 1, 2, 3]) pcoll_2 = pipeline | 'Start 2' >> beam.Create([4, 5, 6, 7]) result = [pcoll for pcoll in (pcoll_1, pcoll_2)] | beam.Flatten() assert_that(result, equal_to([0, 1, 2, 3, 4, 5, 6, 7])) @pytest.mark.it_validatesrunner def test_flatten_a_flattened_pcollection(self): with TestPipeline() as pipeline: pcoll_1 = pipeline | 'Start 1' >> beam.Create([0, 1, 2, 3]) pcoll_2 = pipeline | 'Start 2' >> beam.Create([4, 5, 6, 7]) pcoll_3 = pipeline | 'Start 3' >> beam.Create([8, 9]) pcoll_12 = (pcoll_1, pcoll_2) | 'Flatten' >> beam.Flatten() pcoll_123 = (pcoll_12, pcoll_3) | 'Flatten again' >> beam.Flatten() assert_that(pcoll_123, equal_to([x for x in range(10)])) def test_flatten_input_type_must_be_iterable(self): # Inputs to flatten *must* be an iterable. with self.assertRaises(ValueError): 4 | beam.Flatten() def test_flatten_input_type_must_be_iterable_of_pcolls(self): # Inputs to flatten *must* be an iterable of PCollections. with self.assertRaises(TypeError): {'l': 'test'} | beam.Flatten() with self.assertRaises(TypeError): set([1, 2, 3]) | beam.Flatten() @pytest.mark.it_validatesrunner def test_flatten_multiple_pcollections_having_multiple_consumers(self): with TestPipeline() as pipeline: input = pipeline | 'Start' >> beam.Create(['AA', 'BBB', 'CC']) def split_even_odd(element): tag = 'even_length' if len(element) % 2 == 0 else 'odd_length' return pvalue.TaggedOutput(tag, element) even_length, odd_length = (input | beam.Map(split_even_odd) .with_outputs('even_length', 'odd_length')) merged = (even_length, odd_length) | 'Flatten' >> beam.Flatten() assert_that(merged, equal_to(['AA', 'BBB', 'CC'])) assert_that(even_length, equal_to(['AA', 'CC']), label='assert:even') assert_that(odd_length, equal_to(['BBB']), label='assert:odd') def test_flatten_with(self): with TestPipeline() as pipeline: input = pipeline | 'Start' >> beam.Create(['AA', 'BBB', 'CC']) result = ( input | 'WithPCollection' >> beam.FlattenWith(input | beam.Map(str.lower)) | 'WithPTransform' >> beam.FlattenWith(beam.Create(['x', 'y']))) assert_that( result, equal_to(['AA', 'BBB', 'CC', 'aa', 'bbb', 'cc', 'x', 'y'])) def test_group_by_key_input_must_be_kv_pairs(self): with self.assertRaises(typehints.TypeCheckError) as e: with TestPipeline() as pipeline: pcolls = pipeline | 'A' >> beam.Create([1, 2, 3, 4, 5]) pcolls | 'D' >> beam.GroupByKey() self.assertStartswith( e.exception.args[0], 'Input type hint violation at D: expected ' 'Tuple[TypeVariable[K], TypeVariable[V]]') def test_group_by_key_only_input_must_be_kv_pairs(self): with self.assertRaises(typehints.TypeCheckError) as cm: with TestPipeline() as pipeline: pcolls = pipeline | 'A' >> beam.Create(['a', 'b', 'f']) pcolls | 'D' >> beam.GroupByKey() expected_error_prefix = ( 'Input type hint violation at D: expected ' 'Tuple[TypeVariable[K], TypeVariable[V]]') self.assertStartswith(cm.exception.args[0], expected_error_prefix) def test_keys_and_values(self): with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([(3, 1), (2, 1), (1, 1), (3, 2), (2, 2), (3, 3)]) keys = pcoll.apply(beam.Keys('keys')) vals = pcoll.apply(beam.Values('vals')) assert_that(keys, equal_to([1, 2, 2, 3, 3, 3]), label='assert:keys') assert_that(vals, equal_to([1, 1, 1, 2, 2, 3]), label='assert:vals') def test_kv_swap(self): with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create([(6, 3), (1, 2), (7, 1), (5, 2), (3, 2)]) result = pcoll.apply(beam.KvSwap(), label='swap') assert_that(result, equal_to([(1, 7), (2, 1), (2, 3), (2, 5), (3, 6)])) def test_distinct(self): with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create( [6, 3, 1, 1, 9, 'pleat', 'pleat', 'kazoo', 'navel']) result = pcoll.apply(beam.Distinct()) assert_that(result, equal_to([1, 3, 6, 9, 'pleat', 'kazoo', 'navel'])) def test_chained_ptransforms(self): with TestPipeline() as pipeline: t = ( beam.Map(lambda x: (x, 1)) | beam.GroupByKey() | beam.Map(lambda x_ones: (x_ones[0], sum(x_ones[1])))) result = pipeline | 'Start' >> beam.Create(['a', 'a', 'b']) | t assert_that(result, equal_to([('a', 2), ('b', 1)])) def test_apply_to_list(self): self.assertCountEqual([1, 2, 3], [0, 1, 2] | 'AddOne' >> beam.Map(lambda x: x + 1)) self.assertCountEqual([1], [0, 1, 2] | 'Odd' >> beam.Filter(lambda x: x % 2)) self.assertCountEqual([1, 2, 100, 3], ([1, 2, 3], [100]) | beam.Flatten()) join_input = ([('k', 'a')], [('k', 'b'), ('k', 'c')]) self.assertCountEqual([('k', (['a'], ['b', 'c']))], join_input | beam.CoGroupByKey() | SortLists) def test_multi_input_ptransform(self): class DisjointUnion(PTransform): def expand(self, pcollections): return ( pcollections | beam.Flatten() | beam.Map(lambda x: (x, None)) | beam.GroupByKey() | beam.Map(lambda kv: kv[0])) self.assertEqual([1, 2, 3], sorted(([1, 2], [2, 3]) | DisjointUnion())) def test_apply_to_crazy_pvaluish(self): class NestedFlatten(PTransform): """A PTransform taking and returning nested PValueish. Takes as input a list of dicts, and returns a dict with the corresponding values flattened. """ def _extract_input_pvalues(self, pvalueish): pvalueish = list(pvalueish) return pvalueish, sum([list(p.values()) for p in pvalueish], []) def expand(self, pcoll_dicts): keys = reduce(operator.or_, [set(p.keys()) for p in pcoll_dicts]) res = {} for k in keys: res[k] = [p[k] for p in pcoll_dicts if k in p] | k >> beam.Flatten() return res res = [{ 'a': [1, 2, 3] }, { 'a': [4, 5, 6], 'b': ['x', 'y', 'z'] }, { 'a': [7, 8], 'b': ['x', 'y'], 'c': [] }] | NestedFlatten() self.assertEqual(3, len(res)) self.assertEqual([1, 2, 3, 4, 5, 6, 7, 8], sorted(res['a'])) self.assertEqual(['x', 'x', 'y', 'y', 'z'], sorted(res['b'])) self.assertEqual([], sorted(res['c'])) def test_named_tuple(self): MinMax = collections.namedtuple('MinMax', ['min', 'max']) class MinMaxTransform(PTransform): def expand(self, pcoll): return MinMax( min=pcoll | beam.CombineGlobally(min).without_defaults(), max=pcoll | beam.CombineGlobally(max).without_defaults()) res = [1, 2, 4, 8] | MinMaxTransform() self.assertIsInstance(res, MinMax) self.assertEqual(res, MinMax(min=[1], max=[8])) flat = res | beam.Flatten() self.assertEqual(sorted(flat), [1, 8]) def test_tuple_twice(self): class Duplicate(PTransform): def expand(self, pcoll): return pcoll, pcoll res1, res2 = [1, 2, 4, 8] | Duplicate() self.assertEqual(sorted(res1), [1, 2, 4, 8]) self.assertEqual(sorted(res2), [1, 2, 4, 8]) def test_resource_hint_application_is_additive(self): t = beam.Map(lambda x: x + 1).with_resource_hints( accelerator='gpu').with_resource_hints(min_ram=1).with_resource_hints( accelerator='tpu') self.assertEqual( t.get_resource_hints(), { common_urns.resource_hints.ACCELERATOR.urn: b'tpu', common_urns.resource_hints.MIN_RAM_BYTES.urn: b'1' }) class TestGroupBy(unittest.TestCase): def test_lambdas(self): def normalize(key, values): return tuple(key) if isinstance(key, tuple) else key, sorted(values) with TestPipeline() as p: pcoll = p | beam.Create(range(6)) assert_that( pcoll | beam.GroupBy() | beam.MapTuple(normalize), equal_to([((), [0, 1, 2, 3, 4, 5])]), 'GroupAll') assert_that( pcoll | beam.GroupBy(lambda x: x % 2) | 'n2' >> beam.MapTuple(normalize), equal_to([(0, [0, 2, 4]), (1, [1, 3, 5])]), 'GroupOne') assert_that( pcoll | 'G2' >> beam.GroupBy(lambda x: x % 2).force_tuple_keys() | 'n3' >> beam.MapTuple(normalize), equal_to([((0, ), [0, 2, 4]), ((1, ), [1, 3, 5])]), 'GroupOneTuple') assert_that( pcoll | beam.GroupBy(a=lambda x: x % 2, b=lambda x: x < 4) | 'n4' >> beam.MapTuple(normalize), equal_to([((0, True), [0, 2]), ((1, True), [1, 3]), ((0, False), [4]), ((1, False), [5])]), 'GroupTwo') def test_fields(self): def normalize(key, values): if isinstance(key, tuple): key = beam.Row( **{ name: value for name, value in zip(type(key)._fields, key) }) return key, sorted(v.value for v in values) with TestPipeline() as p: pcoll = p | beam.Create(range(-2, 3)) | beam.Map(int) | beam.Map( lambda x: beam.Row( value=x, square=x * x, sign=x // abs(x) if x else 0)) assert_that( pcoll | beam.GroupBy('square') | beam.MapTuple(normalize), equal_to([ (0, [0]), (1, [-1, 1]), (4, [-2, 2]), ]), 'GroupSquare') assert_that( pcoll | 'G2' >> beam.GroupBy('square').force_tuple_keys() | 'n2' >> beam.MapTuple(normalize), equal_to([ (beam.Row(square=0), [0]), (beam.Row(square=1), [-1, 1]), (beam.Row(square=4), [-2, 2]), ]), 'GroupSquareTupleKey') assert_that( pcoll | beam.GroupBy('square', 'sign') | 'n3' >> beam.MapTuple(normalize), equal_to([ (beam.Row(square=0, sign=0), [0]), (beam.Row(square=1, sign=1), [1]), (beam.Row(square=4, sign=1), [2]), (beam.Row(square=1, sign=-1), [-1]), (beam.Row(square=4, sign=-1), [-2]), ]), 'GroupSquareSign') assert_that( pcoll | beam.GroupBy('square', big=lambda x: x.value > 1) | 'n4' >> beam.MapTuple(normalize), equal_to([ (beam.Row(square=0, big=False), [0]), (beam.Row(square=1, big=False), [-1, 1]), (beam.Row(square=4, big=False), [-2]), (beam.Row(square=4, big=True), [2]), ]), 'GroupSquareNonzero') def test_aggregate(self): def named_tuple_to_row(t): return beam.Row( **{ name: value for name, value in zip(type(t)._fields, t) }) with TestPipeline() as p: pcoll = p | beam.Create(range(-2, 3)) | beam.Map( lambda x: beam.Row( value=x, square=x * x, sign=x // abs(x) if x else 0)) assert_that( pcoll | beam.GroupBy('square', big=lambda x: x.value > 1).aggregate_field( 'value', sum, 'sum').aggregate_field( lambda x: x.sign == 1, all, 'positive') | beam.Map(named_tuple_to_row), equal_to([ beam.Row(square=0, big=False, sum=0, positive=False), # [0], beam.Row(square=1, big=False, sum=0, positive=False), # [-1, 1] beam.Row(square=4, big=False, sum=-2, positive=False), # [-2] beam.Row(square=4, big=True, sum=2, positive=True), # [2] ])) def test_pickled_field(self): with TestPipeline() as p: assert_that( p | beam.Create(['a', 'a', 'b']) | beam.Map( lambda s: beam.Row( key1=PickledObject(s), key2=s.upper(), value=0)) | beam.GroupBy('key1', 'key2') | beam.MapTuple(lambda k, vs: (k.key1.value, k.key2, len(list(vs)))), equal_to([('a', 'A', 2), ('b', 'B', 1)])) class SelectTest(unittest.TestCase): def test_simple(self): with TestPipeline() as p: rows = ( p | beam.Create([1, 2, 10]) | beam.Select(a=lambda x: x * x, b=lambda x: -x)) assert_that( rows, equal_to([ beam.Row(a=1, b=-1), beam.Row(a=4, b=-2), beam.Row(a=100, b=-10), ]), label='CheckFromLambdas') from_attr = rows | beam.Select('b', z='a') assert_that( from_attr, equal_to([ beam.Row(b=-1, z=1), beam.Row(b=-2, z=4), beam.Row( b=-10, z=100, ), ]), label='CheckFromAttrs') def test_type_inference(self): with TestPipeline() as p: input_rows = p | beam.Create([beam.Row(s='abc', i=1)]) output_rows = input_rows | beam.Select( 's', 'i', s_again='s', expr=lambda x: x.i + 1) field_types = dict(output_rows.element_type._fields) self.assertEqual(field_types['s'], str) self.assertEqual(field_types['i'], int) self.assertEqual(field_types['s_again'], str) self.assertEqual(field_types['expr'], int) @beam.ptransform_fn def SamplePTransform(pcoll): """Sample transform using the @ptransform_fn decorator.""" map_transform = 'ToPairs' >> beam.Map(lambda v: (v, None)) combine_transform = 'Group' >> beam.CombinePerKey(lambda vs: None) keys_transform = 'Distinct' >> beam.Keys() return pcoll | map_transform | combine_transform | keys_transform class PTransformLabelsTest(unittest.TestCase): class CustomTransform(beam.PTransform): pardo: Optional[beam.PTransform] = None def expand(self, pcoll): self.pardo = '*Do*' >> beam.FlatMap(lambda x: [x + 1]) return pcoll | self.pardo def test_chained_ptransforms(self): """Tests that chaining gets proper nesting.""" with TestPipeline() as pipeline: map1 = 'Map1' >> beam.Map(lambda x: (x, 1)) gbk = 'Gbk' >> beam.GroupByKey() map2 = 'Map2' >> beam.Map(lambda x_ones2: (x_ones2[0], sum(x_ones2[1]))) t = (map1 | gbk | map2) result = pipeline | 'Start' >> beam.Create(['a', 'a', 'b']) | t self.assertTrue('Map1|Gbk|Map2/Map1' in pipeline.applied_labels) self.assertTrue('Map1|Gbk|Map2/Gbk' in pipeline.applied_labels) self.assertTrue('Map1|Gbk|Map2/Map2' in pipeline.applied_labels) assert_that(result, equal_to([('a', 2), ('b', 1)])) def test_apply_custom_transform_without_label(self): with TestPipeline() as pipeline: pcoll = pipeline | 'PColl' >> beam.Create([1, 2, 3]) custom = PTransformLabelsTest.CustomTransform() result = pipeline.apply(custom, pcoll) self.assertTrue('CustomTransform' in pipeline.applied_labels) self.assertTrue('CustomTransform/*Do*' in pipeline.applied_labels) assert_that(result, equal_to([2, 3, 4])) def test_apply_custom_transform_with_label(self): with TestPipeline() as pipeline: pcoll = pipeline | 'PColl' >> beam.Create([1, 2, 3]) custom = PTransformLabelsTest.CustomTransform('*Custom*') result = pipeline.apply(custom, pcoll) self.assertTrue('*Custom*' in pipeline.applied_labels) self.assertTrue('*Custom*/*Do*' in pipeline.applied_labels) assert_that(result, equal_to([2, 3, 4])) def test_combine_without_label(self): vals = [1, 2, 3, 4, 5, 6, 7] with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create(vals) combine = beam.CombineGlobally(sum) result = pcoll | combine self.assertTrue('CombineGlobally(sum)' in pipeline.applied_labels) assert_that(result, equal_to([sum(vals)])) def test_apply_ptransform_using_decorator(self): pipeline = TestPipeline() pcoll = pipeline | 'PColl' >> beam.Create([1, 2, 3]) _ = pcoll | '*Sample*' >> SamplePTransform() self.assertTrue('*Sample*' in pipeline.applied_labels) self.assertTrue('*Sample*/ToPairs' in pipeline.applied_labels) self.assertTrue('*Sample*/Group' in pipeline.applied_labels) self.assertTrue('*Sample*/Distinct' in pipeline.applied_labels) def test_ptransformfn_default_label(self): @beam.ptransform_fn def MyTransform(self, suffix="xyz"): return pcoll | beam.Map(lambda s: s + suffix) pipeline = TestPipeline() pcoll = pipeline | beam.Create(['a', 'b', 'c']) _ = pcoll | MyTransform() self.assertIn('MyTransform', pipeline.applied_labels) _ = pcoll | MyTransform("suffix") self.assertIn('MyTransform(suffix)', pipeline.applied_labels) _ = pcoll | MyTransform("looooooooooooooooooooooooooooooooooooooooong") self.assertIn('MyTransform(looooooooo...oooong)', pipeline.applied_labels) def test_ptransformfn_legacy_default_label(self): @beam.ptransform_fn def MyTransform(self, suffix="xyz"): return pcoll | beam.Map(lambda s: s + suffix) pipeline = TestPipeline( options=PipelineOptions(update_compatibility_version='2.67.0')) pcoll = pipeline | beam.Create(['a', 'b', 'c']) _ = pcoll | MyTransform() self.assertIn('MyTransform', pipeline.applied_labels) _ = pcoll | MyTransform("suffix") self.assertIn('MyTransform(suffix)', pipeline.applied_labels) _ = pcoll | MyTransform("looooooooooooooooooooooooooooooooooooooooong") self.assertIn( 'MyTransform(looooooooooooooooooooooooooooooooooooooooong)', pipeline.applied_labels) def test_combine_with_label(self): vals = [1, 2, 3, 4, 5, 6, 7] with TestPipeline() as pipeline: pcoll = pipeline | 'Start' >> beam.Create(vals) combine = '*Sum*' >> beam.CombineGlobally(sum) result = pcoll | combine self.assertTrue('*Sum*' in pipeline.applied_labels) assert_that(result, equal_to([sum(vals)])) def check_label(self, ptransform, expected_label): pipeline = TestPipeline() pipeline | 'Start' >> beam.Create([('a', 1)]) | ptransform actual_label = sorted( label for label in pipeline.applied_labels if not label.startswith('Start'))[0] self.assertEqual(expected_label, re.sub(r'\d{3,}', '#', actual_label)) def test_default_labels(self): def my_function(*args): pass self.check_label(beam.Map(len), 'Map(len)') self.check_label(beam.Map(my_function), 'Map(my_function)') self.check_label( beam.Map(lambda x: x), 'Map()') self.check_label(beam.FlatMap(list), 'FlatMap(list)') self.check_label(beam.FlatMap(my_function), 'FlatMap(my_function)') self.check_label(beam.Filter(sum), 'Filter(sum)') self.check_label(beam.CombineGlobally(sum), 'CombineGlobally(sum)') self.check_label(beam.CombinePerKey(sum), 'CombinePerKey(sum)') class MyDoFn(beam.DoFn): def process(self, unused_element): pass self.check_label(beam.ParDo(MyDoFn()), 'ParDo(MyDoFn)') def test_label_propogation(self): self.check_label('TestMap' >> beam.Map(len), 'TestMap') self.check_label('TestLambda' >> beam.Map(lambda x: x), 'TestLambda') self.check_label('TestFlatMap' >> beam.FlatMap(list), 'TestFlatMap') self.check_label('TestFilter' >> beam.Filter(sum), 'TestFilter') self.check_label('TestCG' >> beam.CombineGlobally(sum), 'TestCG') self.check_label('TestCPK' >> beam.CombinePerKey(sum), 'TestCPK') class MyDoFn(beam.DoFn): def process(self, unused_element): pass self.check_label('TestParDo' >> beam.ParDo(MyDoFn()), 'TestParDo') class PTransformTestDisplayData(unittest.TestCase): def test_map_named_function(self): tr = beam.Map(len) dd = DisplayData.create_from(tr) nspace = 'apache_beam.transforms.core.CallableWrapperDoFn' expected_item = DisplayDataItem( 'len', key='fn', label='Transform Function', namespace=nspace) hc.assert_that(dd.items, hc.has_item(expected_item)) def test_map_anonymous_function(self): tr = beam.Map(lambda x: x) dd = DisplayData.create_from(tr) nspace = 'apache_beam.transforms.core.CallableWrapperDoFn' expected_item = DisplayDataItem( '', key='fn', label='Transform Function', namespace=nspace) hc.assert_that(dd.items, hc.has_item(expected_item)) def test_flatmap_named_function(self): tr = beam.FlatMap(list) dd = DisplayData.create_from(tr) nspace = 'apache_beam.transforms.core.CallableWrapperDoFn' expected_item = DisplayDataItem( 'list', key='fn', label='Transform Function', namespace=nspace) hc.assert_that(dd.items, hc.has_item(expected_item)) def test_flatmap_anonymous_function(self): tr = beam.FlatMap(lambda x: [x]) dd = DisplayData.create_from(tr) nspace = 'apache_beam.transforms.core.CallableWrapperDoFn' expected_item = DisplayDataItem( '', key='fn', label='Transform Function', namespace=nspace) hc.assert_that(dd.items, hc.has_item(expected_item)) def test_filter_named_function(self): tr = beam.Filter(sum) dd = DisplayData.create_from(tr) nspace = 'apache_beam.transforms.core.CallableWrapperDoFn' expected_item = DisplayDataItem( 'sum', key='fn', label='Transform Function', namespace=nspace) hc.assert_that(dd.items, hc.has_item(expected_item)) def test_filter_anonymous_function(self): tr = beam.Filter(lambda x: x // 30) dd = DisplayData.create_from(tr) nspace = 'apache_beam.transforms.core.CallableWrapperDoFn' expected_item = DisplayDataItem( '', key='fn', label='Transform Function', namespace=nspace) hc.assert_that(dd.items, hc.has_item(expected_item)) class PTransformTypeCheckTestCase(TypeHintTestCase): def assertStartswith(self, msg, prefix): self.assertTrue( msg.startswith(prefix), '"%s" does not start with "%s"' % (msg, prefix)) def setUp(self): self.p = TestPipeline() def test_do_fn_pipeline_pipeline_type_check_satisfied(self): @with_input_types(int, int) @with_output_types(int) class AddWithFive(beam.DoFn): def process(self, element, five): return [element + five] d = ( self.p | 'T' >> beam.Create([1, 2, 3]).with_output_types(int) | 'Add' >> beam.ParDo(AddWithFive(), 5)) assert_that(d, equal_to([6, 7, 8])) self.p.run() def test_child_with_both_input_and_output_hints_binds_typevars_correctly( self): """ When a child transform has both input and output type hints with type variables, those variables bind correctly from the actual input data. Example: Child with .with_input_types(Tuple[K, V]) .with_output_types(Tuple[K, V]) receiving Tuple['a', 'hello'] will bind K=str, V=str correctly. """ K = typehints.TypeVariable('K') V = typehints.TypeVariable('V') @typehints.with_input_types(typehints.Tuple[K, V]) @typehints.with_output_types(typehints.Tuple[K, V]) class TransformWithoutChildHints(beam.PTransform): class MyDoFn(beam.DoFn): def process(self, element): k, v = element yield (k, v.upper()) def expand(self, pcoll): return ( pcoll | beam.ParDo(self.MyDoFn()).with_input_types( tuple[K, V]).with_output_types(tuple[K, V])) with TestPipeline() as p: result = ( p | beam.Create([('a', 'hello'), ('b', 'world')]) | TransformWithoutChildHints()) self.assertEqual(result.element_type, typehints.Tuple[str, str]) def test_child_without_input_hints_fails_to_bind_typevars(self): """ When a child transform lacks input type hints, type variables in its output hints cannot bind and default to Any, even when parent composite has decorated type hints. This test demonstrates the current limitation: without explicit input hints on the child, the type variable K in .with_output_types(Tuple[K, str]) remains unbound, resulting in Tuple[Any, str] instead of the expected Tuple[str, str]. """ K = typehints.TypeVariable('K') @typehints.with_input_types(typehints.Tuple[K, str]) @typehints.with_output_types(typehints.Tuple[K, str]) class TransformWithoutChildHints(beam.PTransform): class MyDoFn(beam.DoFn): def process(self, element): k, v = element yield (k, v.upper()) def expand(self, pcoll): return ( pcoll | beam.ParDo(self.MyDoFn()).with_output_types(tuple[K, str])) with TestPipeline() as p: result = ( p | beam.Create([('a', 'hello'), ('b', 'world')]) | TransformWithoutChildHints()) self.assertEqual(result.element_type, typehints.Tuple[typehints.Any, str]) def test_child_without_output_hints_infers_partial_types_from_dofn(self): """ When a child transform has input hints but no output hints, type inference from the DoFn's process method produces partially inferred types. Type inference is able to infer the first element of the tuple as str, but not the v.upper() and falls back to any. """ K = typehints.TypeVariable('K') V = typehints.TypeVariable('V') @typehints.with_input_types(typehints.Tuple[K, V]) @typehints.with_output_types(typehints.Tuple[K, V]) class TransformWithoutChildHints(beam.PTransform): class MyDoFn(beam.DoFn): def process(self, element): k, v = element yield (k, v.upper()) def expand(self, pcoll): return (pcoll | beam.ParDo(self.MyDoFn()).with_input_types(tuple[K, V])) with TestPipeline() as p: result = ( p | beam.Create([('a', 'hello'), ('b', 'world')]) | TransformWithoutChildHints()) self.assertEqual(result.element_type, typehints.Tuple[str, typing.Any]) def test_do_fn_pipeline_pipeline_type_check_violated(self): @with_input_types(str, str) @with_output_types(str) class ToUpperCaseWithPrefix(beam.DoFn): def process(self, element, prefix): return [prefix + element.upper()] with self.assertRaisesRegex(typehints.TypeCheckError, r'Upper.*requires.*str.*applied.*int'): ( self.p | 'T' >> beam.Create([1, 2, 3]).with_output_types(int) | 'Upper' >> beam.ParDo(ToUpperCaseWithPrefix(), 'hello')) def test_do_fn_pipeline_runtime_type_check_satisfied(self): self.p._options.view_as(TypeOptions).runtime_type_check = True @with_input_types(int, int) @with_output_types(int) class AddWithNum(beam.DoFn): def process(self, element, num): return [element + num] d = ( self.p | 'T' >> beam.Create([1, 2, 3]).with_output_types(int) | 'Add' >> beam.ParDo(AddWithNum(), 5)) assert_that(d, equal_to([6, 7, 8])) self.p.run() def test_do_fn_pipeline_runtime_type_check_violated(self): self.p._options.view_as(TypeOptions).runtime_type_check = True @with_input_types(int, int) @with_output_types(int) class AddWithNum(beam.DoFn): def process(self, element, num): return [element + num] with self.assertRaisesRegex(typehints.TypeCheckError, r'Add.*requires.*int.*applied.*str'): ( self.p | 'T' >> beam.Create(['1', '2', '3']).with_output_types(str) | 'Add' >> beam.ParDo(AddWithNum(), 5)) self.p.run() def test_pardo_does_not_type_check_using_type_hint_decorators(self): @with_input_types(a=int) @with_output_types(list[str]) def int_to_str(a): return [str(a)] # The function above is expecting an int for its only parameter. However, it # will receive a str instead, which should result in a raised exception. with self.assertRaisesRegex(typehints.TypeCheckError, r'ToStr.*requires.*int.*applied.*str'): ( self.p | 'S' >> beam.Create(['b', 'a', 'r']).with_output_types(str) | 'ToStr' >> beam.FlatMap(int_to_str)) def test_pardo_properly_type_checks_using_type_hint_decorators(self): @with_input_types(a=str) @with_output_types(list[str]) def to_all_upper_case(a): return [a.upper()] # If this type-checks than no error should be raised. d = ( self.p | 'T' >> beam.Create(['t', 'e', 's', 't']).with_output_types(str) | 'Case' >> beam.FlatMap(to_all_upper_case)) assert_that(d, equal_to(['T', 'E', 'S', 'T'])) self.p.run() # Output type should have been recognized as 'str' rather than List[str] to # do the flatten part of FlatMap. self.assertEqual(str, d.element_type) def test_pardo_does_not_type_check_using_type_hint_methods(self): # The first ParDo outputs pcoll's of type int, however the second ParDo is # expecting pcoll's of type str instead. with self.assertRaisesRegex(typehints.TypeCheckError, r'Upper.*requires.*str.*applied.*int'): ( self.p | 'S' >> beam.Create(['t', 'e', 's', 't']).with_output_types(str) | ( 'Score' >> beam.FlatMap(lambda x: [1] if x == 't' else [2]). with_input_types(str).with_output_types(int)) | ( 'Upper' >> beam.FlatMap(lambda x: [x.upper()]).with_input_types( str).with_output_types(str))) def test_pardo_properly_type_checks_using_type_hint_methods(self): # Pipeline should be created successfully without an error d = ( self.p | 'S' >> beam.Create(['t', 'e', 's', 't']).with_output_types(str) | 'Dup' >> beam.FlatMap(lambda x: [x + x]).with_input_types( str).with_output_types(str) | 'Upper' >> beam.FlatMap(lambda x: [x.upper()]).with_input_types( str).with_output_types(str)) assert_that(d, equal_to(['TT', 'EE', 'SS', 'TT'])) self.p.run() def test_map_does_not_type_check_using_type_hints_methods(self): # The transform before 'Map' has indicated that it outputs PCollections with # int's, while Map is expecting one of str. with self.assertRaisesRegex(typehints.TypeCheckError, r'Upper.*requires.*str.*applied.*int'): ( self.p | 'S' >> beam.Create([1, 2, 3, 4]).with_output_types(int) | 'Upper' >> beam.Map(lambda x: x.upper()).with_input_types( str).with_output_types(str)) def test_map_properly_type_checks_using_type_hints_methods(self): # No error should be raised if this type-checks properly. d = ( self.p | 'S' >> beam.Create([1, 2, 3, 4]).with_output_types(int) | 'ToStr' >> beam.Map(lambda x: str(x)).with_input_types(int).with_output_types(str)) assert_that(d, equal_to(['1', '2', '3', '4'])) self.p.run() def test_map_does_not_type_check_using_type_hints_decorator(self): @with_input_types(s=str) @with_output_types(str) def upper(s): return s.upper() # Hinted function above expects a str at pipeline construction. # However, 'Map' should detect that Create has hinted an int instead. with self.assertRaisesRegex(typehints.TypeCheckError, r'Upper.*requires.*str.*applied.*int'): ( self.p | 'S' >> beam.Create([1, 2, 3, 4]).with_output_types(int) | 'Upper' >> beam.Map(upper)) def test_map_properly_type_checks_using_type_hints_decorator(self): @with_input_types(a=bool) @with_output_types(int) def bool_to_int(a): return int(a) # If this type-checks than no error should be raised. d = ( self.p | 'Bools' >> beam.Create([True, False, True]).with_output_types(bool) | 'ToInts' >> beam.Map(bool_to_int)) assert_that(d, equal_to([1, 0, 1])) self.p.run() def test_filter_does_not_type_check_using_type_hints_method(self): # Filter is expecting an int but instead looks to the 'left' and sees a str # incoming. with self.assertRaisesRegex(typehints.TypeCheckError, r'Below 3.*requires.*int.*applied.*str'): ( self.p | 'Strs' >> beam.Create(['1', '2', '3', '4', '5' ]).with_output_types(str) | 'Lower' >> beam.Map(lambda x: x.lower()).with_input_types( str).with_output_types(str) | 'Below 3' >> beam.Filter(lambda x: x < 3).with_input_types(int)) def test_filter_type_checks_using_type_hints_method(self): # No error should be raised if this type-checks properly. d = ( self.p | beam.Create(['1', '2', '3', '4', '5']).with_output_types(str) | 'ToInt' >> beam.Map(lambda x: int(x)).with_input_types(str).with_output_types(int) | 'Below 3' >> beam.Filter(lambda x: x < 3).with_input_types(int)) assert_that(d, equal_to([1, 2])) self.p.run() def test_filter_does_not_type_check_using_type_hints_decorator(self): @with_input_types(a=float) def more_than_half(a): return a > 0.50 # Func above was hinted to only take a float, yet a str will be passed. with self.assertRaisesRegex(typehints.TypeCheckError, r'Half.*requires.*float.*applied.*str'): ( self.p | 'Ints' >> beam.Create(['1', '2', '3', '4']).with_output_types(str) | 'Half' >> beam.Filter(more_than_half)) def test_filter_type_checks_using_type_hints_decorator(self): @with_input_types(b=int) def half(b): return bool(random.choice([0, 1])) # Filter should deduce that it returns the same type that it takes. ( self.p | 'Str' >> beam.Create(range(5)).with_output_types(int) | 'Half' >> beam.Filter(half) | 'ToBool' >> beam.Map(lambda x: bool(x)).with_input_types( int).with_output_types(bool)) def test_pardo_like_inheriting_output_types_from_annotation(self): def fn1(x: str) -> int: return 1 def fn1_flat(x: str) -> list[int]: return [1] def fn2(x: int, y: str) -> str: return y def fn2_flat(x: int, y: str) -> list[str]: return [y] # We only need the args section of the hints. def output_hints(transform): return transform.default_type_hints().output_types[0][0] self.assertEqual(int, output_hints(beam.Map(fn1))) self.assertEqual(int, output_hints(beam.FlatMap(fn1_flat))) self.assertEqual(str, output_hints(beam.MapTuple(fn2))) self.assertEqual(str, output_hints(beam.FlatMapTuple(fn2_flat))) def add(a: typing.Iterable[int]) -> int: return sum(a) self.assertCompatible( tuple[typing.TypeVar('K'), int], output_hints(beam.CombinePerKey(add))) def test_group_by_key_only_output_type_deduction(self): d = ( self.p | 'Str' >> beam.Create(['t', 'e', 's', 't']).with_output_types(str) | ( 'Pair' >> beam.Map(lambda x: (x, ord(x))).with_output_types(tuple[str, str])) | beam.GroupByKey()) # Output type should correctly be deduced. # GBK-only should deduce that Tuple[A, B] is turned into # Tuple[A, Iterable[B]]. self.assertCompatible(tuple[str, typing.Iterable[str]], d.element_type) def test_group_by_key_output_type_deduction(self): d = ( self.p | 'Str' >> beam.Create(range(20)).with_output_types(int) | ( 'PairNegative' >> beam.Map(lambda x: (x % 5, -x)).with_output_types( tuple[int, int])) | beam.GroupByKey()) # Output type should correctly be deduced. # GBK should deduce that Tuple[A, B] is turned into Tuple[A, Iterable[B]]. self.assertCompatible(tuple[int, typing.Iterable[int]], d.element_type) def test_group_by_key_only_does_not_type_check(self): # GBK will be passed raw int's here instead of some form of Tuple[A, B]. with self.assertRaises(typehints.TypeCheckError) as e: ( self.p | beam.Create([1, 2, 3]).with_output_types(int) | 'F' >> beam.GroupByKey()) self.assertStartswith( e.exception.args[0], "Input type hint violation at F: " "expected Tuple[TypeVariable[K], TypeVariable[V]], " "got {}".format(int)) def test_group_by_does_not_type_check(self): # Create is returning a List[int, str], rather than a Tuple[int, str] # that is aliased to Tuple[int, str]. with self.assertRaises(typehints.TypeCheckError) as e: ( self.p | (beam.Create([[1], [2]]).with_output_types(typing.Iterable[int])) | 'T' >> beam.GroupByKey()) self.assertStartswith( e.exception.args[0], "Input type hint violation at T: " "expected Tuple[TypeVariable[K], TypeVariable[V]], " "got Iterable[]") def test_pipeline_checking_pardo_insufficient_type_information(self): self.p._options.view_as(TypeOptions).type_check_strictness = 'ALL_REQUIRED' # Type checking is enabled, but 'Create' doesn't pass on any relevant type # information to the ParDo. with self.assertRaises(typehints.TypeCheckError) as e: ( self.p | 'Nums' >> beam.Create(range(5)) | 'ModDup' >> beam.FlatMap(lambda x: (x % 2, x))) self.assertEqual( 'Pipeline type checking is enabled, however no output ' 'type-hint was found for the PTransform Create(Nums)', e.exception.args[0]) def test_pipeline_checking_gbk_insufficient_type_information(self): self.p._options.view_as(TypeOptions).type_check_strictness = 'ALL_REQUIRED' # Type checking is enabled, but 'Map' doesn't pass on any relevant type # information to GBK-only. with self.assertRaises(typehints.TypeCheckError) as e: ( self.p | 'Nums' >> beam.Create(range(5)).with_output_types(int) | 'ModDup' >> beam.Map(lambda x: (x % 2, x)) | beam.GroupByKey()) self.assertEqual( 'Pipeline type checking is enabled, however no output ' 'type-hint was found for the PTransform ' 'ParDo(ModDup)', e.exception.args[0]) def test_disable_pipeline_type_check(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False # The pipeline below should raise a TypeError, however pipeline type # checking was disabled above. ( self.p | 'T' >> beam.Create([1, 2, 3]).with_output_types(int) | 'Lower' >> beam.Map(lambda x: x.lower()).with_input_types( str).with_output_types(str)) def test_run_time_type_checking_enabled_type_violation(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False self.p._options.view_as(TypeOptions).runtime_type_check = True @with_output_types(str) @with_input_types(x=int) def int_to_string(x): return str(x) # Function above has been type-hinted to only accept an int. But in the # pipeline execution it'll be passed a string due to the output of Create. ( self.p | 'T' >> beam.Create(['some_string']) | 'ToStr' >> beam.Map(int_to_string)) error_regex = "Type-hint for argument: 'x' violated. Expected an instance " "of {}, instead found some_string, an instance of {}.".format(int, str) with self.assertRaisesRegex(Exception, error_regex): self.p.run() def test_run_time_type_checking_enabled_types_satisfied(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False self.p._options.view_as(TypeOptions).runtime_type_check = True @with_output_types(tuple[int, str]) @with_input_types(x=str) def group_with_upper_ord(x): return (ord(x.upper()) % 5, x) # Pipeline checking is off, but the above function should satisfy types at # run-time. result = ( self.p | 'T' >> beam.Create(['t', 'e', 's', 't', 'i', 'n', 'g' ]).with_output_types(str) | 'GenKeys' >> beam.Map(group_with_upper_ord) | 'O' >> beam.GroupByKey() | SortLists) assert_that( result, equal_to([(1, ['g']), (3, ['i', 'n', 's']), (4, ['e', 't', 't'])])) self.p.run() def test_pipeline_checking_satisfied_but_run_time_types_violate(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False self.p._options.view_as(TypeOptions).runtime_type_check = True @with_output_types(tuple[bool, int]) @with_input_types(a=int) def is_even_as_key(a): # Simulate a programming error, should be: return (a % 2 == 0, a) # However this returns Tuple[int, int] return (a % 2, a) ( self.p | 'Nums' >> beam.Create(range(5)).with_output_types(int) | 'IsEven' >> beam.Map(is_even_as_key) | 'Parity' >> beam.GroupByKey()) # Although all the types appear to be correct when checked at pipeline # construction. Runtime type-checking should detect the 'is_even_as_key' is # returning Tuple[int, int], instead of Tuple[bool, int]. error_regex = "Runtime type violation detected" with self.assertRaisesRegex(Exception, error_regex): self.p.run() def test_pipeline_checking_satisfied_run_time_checking_satisfied(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False @with_output_types(tuple[bool, int]) @with_input_types(a=int) def is_even_as_key(a): # The programming error in the above test-case has now been fixed. # Everything should properly type-check. return (a % 2 == 0, a) result = ( self.p | 'Nums' >> beam.Create(range(5)).with_output_types(int) | 'IsEven' >> beam.Map(is_even_as_key) | 'Parity' >> beam.GroupByKey() | SortLists) assert_that(result, equal_to([(False, [1, 3]), (True, [0, 2, 4])])) self.p.run() def test_pipeline_runtime_checking_violation_simple_type_input(self): self.p._options.view_as(TypeOptions).runtime_type_check = True self.p._options.view_as(TypeOptions).pipeline_type_check = False # The type-hinted applied via the 'with_input_types()' method indicates the # ParDo should receive an instance of type 'str', however an 'int' will be # passed instead. error_regex = "Runtime type violation detected" with self.assertRaisesRegex(Exception, error_regex): ( self.p | beam.Create([1, 1, 1]) | ( 'ToInt' >> beam.FlatMap(lambda x: [int(x)]).with_input_types( str).with_output_types(int))) self.p.run() def test_pipeline_runtime_checking_violation_composite_type_input(self): self.p._options.view_as(TypeOptions).runtime_type_check = True self.p._options.view_as(TypeOptions).pipeline_type_check = False error_regex = "Runtime type violation detected" with self.assertRaisesRegex(Exception, error_regex): ( self.p | beam.Create([(1, 3.0), (2, 4.9), (3, 9.5)]) | ( 'Add' >> beam.FlatMap(lambda x_y: [x_y[0] + x_y[1]]).with_input_types( tuple[int, int]).with_output_types(int))) self.p.run() def test_pipeline_runtime_checking_violation_simple_type_output(self): self.p._options.view_as(TypeOptions).runtime_type_check = True self.p._options.view_as(TypeOptions).pipeline_type_check = False # The type-hinted applied via the 'returns()' method indicates the ParDo # should output an instance of type 'int', however a 'float' will be # generated instead. print( "HINTS", ( 'ToInt' >> beam.FlatMap(lambda x: [float(x)]).with_input_types( int).with_output_types(int)).get_type_hints()) error_regex = "" if self.p._options.view_as(TypeOptions).runtime_type_check: error_regex = ( "Runtime type violation detected within ParDo\\(ToInt\\):" + " According to type-hint expected output should be of type . Instead, received '1.0', an instance of type " ) if self.p._options.view_as(TypeOptions).performance_runtime_type_check: error_regex = ( "Runtime type violation detected within ToInt: Type-hint " + "for argument: 'x' violated. Expected an instance of , " + "instead found 1.0, an instance of ") with self.assertRaisesRegex(Exception, error_regex): ( self.p | beam.Create([1, 1, 1]) | ( 'ToInt' >> beam.FlatMap(lambda x: [float(x)]).with_input_types( int).with_output_types(int))) self.p.run().wait_until_finish() def test_pipeline_runtime_checking_violation_composite_type_output(self): self.p._options.view_as(TypeOptions).runtime_type_check = True self.p._options.view_as(TypeOptions).pipeline_type_check = False # The type-hinted applied via the 'returns()' method indicates the ParDo # should return an instance of type: Tuple[float, int]. However, an instance # of 'int' will be generated instead. error_regex = ( "Runtime type violation detected within " + "ParDo\\(Swap\\): Tuple type constraint violated. " + "Valid object instance must be of type 'tuple'. Instead, " + "an instance of 'float' was received.") with self.assertRaisesRegex(Exception, error_regex): ( self.p | beam.Create([(1, 3.0), (2, 4.9), (3, 9.5)]) | ( 'Swap' >> beam.FlatMap(lambda x_y1: [x_y1[0] + x_y1[1]]).with_input_types( tuple[int, float]).with_output_types(tuple[float, int]))) self.p.run() def test_pipeline_runtime_checking_violation_with_side_inputs_decorator(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False self.p._options.view_as(TypeOptions).runtime_type_check = True @with_output_types(int) @with_input_types(a=int, b=int) def add(a, b): return a + b error_regex = "Runtime type violation detected" with self.assertRaisesRegex(Exception, error_regex): (self.p | beam.Create([1, 2, 3, 4]) | 'Add 1' >> beam.Map(add, 1.0)) self.p.run() def test_pipeline_runtime_checking_violation_with_side_inputs_via_method(self): # pylint: disable=line-too-long self.p._options.view_as(TypeOptions).runtime_type_check = True self.p._options.view_as(TypeOptions).pipeline_type_check = False error_regex = "Runtime type violation detected" with self.assertRaisesRegex(Exception, error_regex): ( self.p | beam.Create([1, 2, 3, 4]) | ( 'Add 1' >> beam.Map(lambda x, one: x + one, 1.0).with_input_types( int, int).with_output_types(float))) self.p.run() def test_combine_properly_pipeline_type_checks_using_decorator(self): @with_output_types(int) @with_input_types(ints=typing.Iterable[int]) def sum_ints(ints): return sum(ints) d = ( self.p | 'T' >> beam.Create([1, 2, 3]).with_output_types(int) | 'Sum' >> beam.CombineGlobally(sum_ints)) self.assertEqual(int, d.element_type) assert_that(d, equal_to([6])) self.p.run() def test_combine_properly_pipeline_type_checks_without_decorator(self): def sum_ints(ints): return sum(ints) d = ( self.p | beam.Create([1, 2, 3]) | beam.Map(lambda x: ('key', x)) | beam.CombinePerKey(sum_ints)) self.assertEqual(typehints.Tuple[str, typehints.Any], d.element_type) self.p.run() def test_combine_func_type_hint_does_not_take_iterable_using_decorator(self): @with_output_types(int) @with_input_types(a=int) def bad_combine(a): 5 + a with self.assertRaises(typehints.TypeCheckError) as e: ( self.p | 'M' >> beam.Create([1, 2, 3]).with_output_types(int) | 'Add' >> beam.CombineGlobally(bad_combine)) self.assertEqual( "All functions for a Combine PTransform must accept a " "single argument compatible with: Iterable[Any]. " "Instead a function with input type: {} was received.".format(int), e.exception.args[0]) def test_combine_pipeline_type_propagation_using_decorators(self): @with_output_types(int) @with_input_types(ints=typing.Iterable[int]) def sum_ints(ints): return sum(ints) @with_output_types(list[int]) @with_input_types(n=int) def range_from_zero(n): return list(range(n + 1)) d = ( self.p | 'T' >> beam.Create([1, 2, 3]).with_output_types(int) | 'Sum' >> beam.CombineGlobally(sum_ints) | 'Range' >> beam.ParDo(range_from_zero)) self.assertEqual(int, d.element_type) assert_that(d, equal_to([0, 1, 2, 3, 4, 5, 6])) self.p.run() def test_combine_runtime_type_check_satisfied_using_decorators(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False @with_output_types(int) @with_input_types(ints=typing.Iterable[int]) def iter_mul(ints): return reduce(operator.mul, ints, 1) d = ( self.p | 'K' >> beam.Create([5, 5, 5, 5]).with_output_types(int) | 'Mul' >> beam.CombineGlobally(iter_mul)) assert_that(d, equal_to([625])) self.p.run() def test_combine_runtime_type_check_violation_using_decorators(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False self.p._options.view_as(TypeOptions).runtime_type_check = True # Combine fn is returning the incorrect type @with_output_types(int) @with_input_types(ints=typing.Iterable[int]) def iter_mul(ints): return str(reduce(operator.mul, ints, 1)) error_regex = ( "Runtime type violation detected within " + "Mul/CombinePerKey: Type-hint for return type violated. " + "Expected an instance of {}, instead found".format(int)) with self.assertRaisesRegex(Exception, error_regex): ( self.p | 'K' >> beam.Create([5, 5, 5, 5]).with_output_types(int) | 'Mul' >> beam.CombineGlobally(iter_mul)) self.p.run() def test_combine_pipeline_type_check_using_methods(self): d = ( self.p | beam.Create(['t', 'e', 's', 't']).with_output_types(str) | ( 'concat' >> beam.CombineGlobally(lambda s: ''.join(s)). with_input_types(str).with_output_types(str))) def matcher(expected): def match(actual): equal_to(expected)(list(actual[0])) return match assert_that(d, matcher('estt')) self.p.run() def test_combine_runtime_type_check_using_methods(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False self.p._options.view_as(TypeOptions).runtime_type_check = True d = ( self.p | beam.Create(range(5)).with_output_types(int) | ( 'Sum' >> beam.CombineGlobally(lambda s: sum(s)).with_input_types( int).with_output_types(int))) assert_that(d, equal_to([10])) self.p.run() def test_combine_pipeline_type_check_violation_using_methods(self): with self.assertRaises(typehints.TypeCheckError) as e: ( self.p | beam.Create(range(3)).with_output_types(int) | ( 'SortJoin' >> beam.CombineGlobally(lambda s: ''.join(sorted(s))). with_input_types(str).with_output_types(str))) self.assertStartswith( e.exception.args[0], "Input type hint violation at SortJoin: " "expected {}, got {}".format(str, int)) def test_combine_runtime_type_check_violation_using_methods(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False self.p._options.view_as(TypeOptions).runtime_type_check = True error_regex = ( "Runtime type violation detected within " + "ParDo\\(SortJoin/KeyWithVoid\\): " + "Type-hint for argument: 'v' violated. Expected an instance of " + ", instead found 0, an instance of .") with self.assertRaisesRegex(Exception, error_regex): ( self.p | beam.Create([0]).with_output_types(int) | ( 'SortJoin' >> beam.CombineGlobally(lambda s: ''.join(sorted(s))). with_input_types(str).with_output_types(str))) self.p.run() def test_combine_insufficient_type_hint_information(self): self.p._options.view_as(TypeOptions).type_check_strictness = 'ALL_REQUIRED' with self.assertRaises(typehints.TypeCheckError) as e: ( self.p | 'E' >> beam.Create(range(3)).with_output_types(int) | 'SortJoin' >> beam.CombineGlobally(lambda s: ''.join(sorted(s))) | 'F' >> beam.Map(lambda x: x + 1)) self.assertStartswith( e.exception.args[0], 'Pipeline type checking is enabled, ' 'however no output type-hint was found for the PTransform ' 'ParDo(' 'SortJoin/CombinePerKey/') def test_mean_globally_pipeline_checking_satisfied(self): d = ( self.p | 'C' >> beam.Create(range(5)).with_output_types(int) | 'Mean' >> combine.Mean.Globally()) self.assertEqual(float, d.element_type) assert_that(d, equal_to([2.0])) self.p.run() def test_mean_globally_pipeline_checking_violated(self): with self.assertRaises(typehints.TypeCheckError) as e: ( self.p | 'C' >> beam.Create(['test']).with_output_types(str) | 'Mean' >> combine.Mean.Globally()) err_msg = e.exception.args[0] assert "CombinePerKey" in err_msg assert "Tuple[TypeVariable[K]" in err_msg assert "Tuple[None, " in err_msg def test_mean_globally_runtime_checking_satisfied(self): self.p._options.view_as(TypeOptions).runtime_type_check = True d = ( self.p | 'C' >> beam.Create(range(5)).with_output_types(int) | 'Mean' >> combine.Mean.Globally()) self.assertEqual(float, d.element_type) assert_that(d, equal_to([2.0])) self.p.run() def test_mean_globally_runtime_checking_violated(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False self.p._options.view_as(TypeOptions).runtime_type_check = True error_regex = "Runtime type violation detected" with self.assertRaisesRegex(Exception, error_regex): ( self.p | 'C' >> beam.Create(['t', 'e', 's', 't']).with_output_types(str) | 'Mean' >> combine.Mean.Globally()) self.p.run() def test_mean_per_key_pipeline_checking_satisfied(self): d = ( self.p | beam.Create(range(5)).with_output_types(int) | ( 'EvenGroup' >> beam.Map(lambda x: (not x % 2, x)).with_output_types( tuple[bool, int])) | 'EvenMean' >> combine.Mean.PerKey()) self.assertCompatible(tuple[bool, float], d.element_type) assert_that(d, equal_to([(False, 2.0), (True, 2.0)])) self.p.run() def test_mean_per_key_pipeline_checking_violated(self): with self.assertRaises(typehints.TypeCheckError) as e: ( self.p | beam.Create(map(str, range(5))).with_output_types(str) | ( 'UpperPair' >> beam.Map(lambda x: (x.upper(), x)).with_output_types( tuple[str, str])) | 'EvenMean' >> combine.Mean.PerKey()) self.p.run() err_msg = e.exception.args[0] assert "CombinePerKey(MeanCombineFn)" in err_msg assert "requires" in err_msg assert "Tuple[TypeVariable[K]" in err_msg assert "applied" in err_msg assert "Tuple[, ]" in err_msg def test_mean_per_key_runtime_checking_satisfied(self): self.p._options.view_as(TypeOptions).runtime_type_check = True d = ( self.p | beam.Create(range(5)).with_output_types(int) | ( 'OddGroup' >> beam.Map(lambda x: (bool(x % 2), x)).with_output_types( tuple[bool, int])) | 'OddMean' >> combine.Mean.PerKey()) self.assertCompatible(tuple[bool, float], d.element_type) assert_that(d, equal_to([(False, 2.0), (True, 2.0)])) self.p.run() def test_mean_per_key_runtime_checking_violated(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False self.p._options.view_as(TypeOptions).runtime_type_check = True error_regex = "Runtime type violation detected" with self.assertRaisesRegex(Exception, error_regex): ( self.p | beam.Create(range(5)).with_output_types(int) | ( 'OddGroup' >> beam.Map(lambda x: (x, str(bool(x % 2)))).with_output_types( tuple[int, str])) | 'OddMean' >> combine.Mean.PerKey()) self.p.run() def test_count_globally_pipeline_type_checking_satisfied(self): d = ( self.p | 'P' >> beam.Create(range(5)).with_output_types(int) | 'CountInt' >> combine.Count.Globally()) self.assertEqual(int, d.element_type) assert_that(d, equal_to([5])) self.p.run() def test_count_globally_runtime_type_checking_satisfied(self): self.p._options.view_as(TypeOptions).runtime_type_check = True d = ( self.p | 'P' >> beam.Create(range(5)).with_output_types(int) | 'CountInt' >> combine.Count.Globally()) self.assertEqual(int, d.element_type) assert_that(d, equal_to([5])) self.p.run() def test_count_perkey_pipeline_type_checking_satisfied(self): d = ( self.p | beam.Create(range(5)).with_output_types(int) | 'EvenGroup' >> beam.Map(lambda x: (not x % 2, x)).with_output_types( tuple[bool, int]) | 'CountInt' >> combine.Count.PerKey()) self.assertCompatible(tuple[bool, int], d.element_type) assert_that(d, equal_to([(False, 2), (True, 3)])) self.p.run() def test_count_perkey_pipeline_type_checking_violated(self): with self.assertRaises(typehints.TypeCheckError) as e: ( self.p | beam.Create(range(5)).with_output_types(int) | 'CountInt' >> combine.Count.PerKey()) self.assertStartswith( e.exception.args[0], 'Input type hint violation at CountInt') def test_count_perkey_runtime_type_checking_satisfied(self): self.p._options.view_as(TypeOptions).runtime_type_check = True d = ( self.p | beam.Create(['t', 'e', 's', 't']).with_output_types(str) | 'DupKey' >> beam.Map(lambda x: (x, x)).with_output_types(tuple[str, str]) | 'CountDups' >> combine.Count.PerKey()) self.assertCompatible(tuple[str, int], d.element_type) assert_that(d, equal_to([('e', 1), ('s', 1), ('t', 2)])) self.p.run() def test_count_perelement_pipeline_type_checking_satisfied(self): d = ( self.p | beam.Create([1, 1, 2, 3]).with_output_types(int) | 'CountElems' >> combine.Count.PerElement()) self.assertCompatible(tuple[int, int], d.element_type) assert_that(d, equal_to([(1, 2), (2, 1), (3, 1)])) self.p.run() def test_count_perelement_pipeline_type_checking_violated(self): self.p._options.view_as(TypeOptions).type_check_strictness = 'ALL_REQUIRED' with self.assertRaises(typehints.TypeCheckError) as e: ( self.p | 'f' >> beam.Create([1, 1, 2, 3]) | 'CountElems' >> combine.Count.PerElement()) self.assertEqual( 'Pipeline type checking is enabled, however no output ' 'type-hint was found for the PTransform ' 'Create(f)', e.exception.args[0]) def test_count_perelement_runtime_type_checking_satisfied(self): self.p._options.view_as(TypeOptions).runtime_type_check = True d = ( self.p | beam.Create([True, True, False, True, True]).with_output_types(bool) | 'CountElems' >> combine.Count.PerElement()) self.assertCompatible(tuple[bool, int], d.element_type) assert_that(d, equal_to([(False, 1), (True, 4)])) self.p.run() def test_top_of_pipeline_checking_satisfied(self): d = ( self.p | beam.Create(range(5, 11)).with_output_types(int) | 'Top 3' >> combine.Top.Of(3)) self.assertCompatible(typing.Iterable[int], d.element_type) assert_that(d, equal_to([[10, 9, 8]])) self.p.run() def test_top_of_runtime_checking_satisfied(self): self.p._options.view_as(TypeOptions).runtime_type_check = True d = ( self.p | beam.Create(list('testing')).with_output_types(str) | 'AciiTop' >> combine.Top.Of(3)) self.assertCompatible(typing.Iterable[str], d.element_type) assert_that(d, equal_to([['t', 't', 's']])) self.p.run() def test_per_key_pipeline_checking_violated(self): with self.assertRaises(typehints.TypeCheckError) as e: ( self.p | beam.Create(range(100)).with_output_types(int) | 'Num + 1' >> beam.Map(lambda x: x + 1).with_output_types(int) | 'TopMod' >> combine.Top.PerKey(1)) self.assertStartswith( e.exception.args[0], "Input type hint violation at TopMod: expected Tuple[TypeVariable[K], " "TypeVariable[V]], got {}".format(int)) def test_per_key_pipeline_checking_satisfied(self): d = ( self.p | beam.Create(range(100)).with_output_types(int) | ( 'GroupMod 3' >> beam.Map(lambda x: (x % 3, x)).with_output_types( tuple[int, int])) | 'TopMod' >> combine.Top.PerKey(1)) self.assertCompatible(tuple[int, typing.Iterable[int]], d.element_type) assert_that(d, equal_to([(0, [99]), (1, [97]), (2, [98])])) self.p.run() def test_per_key_runtime_checking_satisfied(self): self.p._options.view_as(TypeOptions).runtime_type_check = True d = ( self.p | beam.Create(range(21)) | ( 'GroupMod 3' >> beam.Map(lambda x: (x % 3, x)).with_output_types( tuple[int, int])) | 'TopMod' >> combine.Top.PerKey(1)) self.assertCompatible(tuple[int, typing.Iterable[int]], d.element_type) assert_that(d, equal_to([(0, [18]), (1, [19]), (2, [20])])) self.p.run() def test_sample_globally_pipeline_satisfied(self): d = ( self.p | beam.Create([2, 2, 3, 3]).with_output_types(int) | 'Sample' >> combine.Sample.FixedSizeGlobally(3)) self.assertCompatible(typing.Iterable[int], d.element_type) def matcher(expected_len): def match(actual): equal_to([expected_len])([len(actual[0])]) return match assert_that(d, matcher(3)) self.p.run() def test_sample_globally_runtime_satisfied(self): self.p._options.view_as(TypeOptions).runtime_type_check = True d = ( self.p | beam.Create([2, 2, 3, 3]).with_output_types(int) | 'Sample' >> combine.Sample.FixedSizeGlobally(2)) self.assertCompatible(typing.Iterable[int], d.element_type) def matcher(expected_len): def match(actual): equal_to([expected_len])([len(actual[0])]) return match assert_that(d, matcher(2)) self.p.run() def test_sample_per_key_pipeline_satisfied(self): d = ( self.p | ( beam.Create([(1, 2), (1, 2), (2, 3), (2, 3)]).with_output_types(tuple[int, int])) | 'Sample' >> combine.Sample.FixedSizePerKey(2)) self.assertCompatible(tuple[int, typing.Iterable[int]], d.element_type) def matcher(expected_len): def match(actual): for _, sample in actual: equal_to([expected_len])([len(sample)]) return match assert_that(d, matcher(2)) self.p.run() def test_sample_per_key_runtime_satisfied(self): self.p._options.view_as(TypeOptions).runtime_type_check = True d = ( self.p | ( beam.Create([(1, 2), (1, 2), (2, 3), (2, 3)]).with_output_types(tuple[int, int])) | 'Sample' >> combine.Sample.FixedSizePerKey(1)) self.assertCompatible(tuple[int, typing.Iterable[int]], d.element_type) def matcher(expected_len): def match(actual): for _, sample in actual: equal_to([expected_len])([len(sample)]) return match assert_that(d, matcher(1)) self.p.run() def test_to_list_pipeline_check_satisfied(self): d = ( self.p | beam.Create((1, 2, 3, 4)).with_output_types(int) | combine.ToList()) self.assertCompatible(list[int], d.element_type) def matcher(expected): def match(actual): equal_to(expected)(actual[0]) return match assert_that(d, matcher([1, 2, 3, 4])) self.p.run() def test_to_list_runtime_check_satisfied(self): self.p._options.view_as(TypeOptions).runtime_type_check = True d = ( self.p | beam.Create(list('test')).with_output_types(str) | combine.ToList()) self.assertCompatible(list[str], d.element_type) def matcher(expected): def match(actual): equal_to(expected)(actual[0]) return match assert_that(d, matcher(['e', 's', 't', 't'])) self.p.run() def test_to_dict_pipeline_check_violated(self): with self.assertRaises(typehints.TypeCheckError) as e: ( self.p | beam.Create([1, 2, 3, 4]).with_output_types(int) | combine.ToDict()) self.assertStartswith( e.exception.args[0], "Input type hint violation at ToDict: expected Tuple[TypeVariable[K], " "TypeVariable[V]], got {}".format(int)) def test_to_dict_pipeline_check_satisfied(self): d = ( self.p | beam.Create([(1, 2), (3, 4)]).with_output_types(tuple[int, int]) | combine.ToDict()) self.assertCompatible(dict[int, int], d.element_type) assert_that(d, equal_to([{1: 2, 3: 4}])) self.p.run() def test_to_dict_runtime_check_satisfied(self): self.p._options.view_as(TypeOptions).runtime_type_check = True d = ( self.p | (beam.Create([('1', 2), ('3', 4)]).with_output_types(tuple[str, int])) | combine.ToDict()) self.assertCompatible(dict[str, int], d.element_type) assert_that(d, equal_to([{'1': 2, '3': 4}])) self.p.run() def test_runtime_type_check_python_type_error(self): self.p._options.view_as(TypeOptions).runtime_type_check = True error_regex = "object of type 'int' has no len()" with self.assertRaisesRegex(Exception, error_regex): ( self.p | beam.Create([1, 2, 3]).with_output_types(int) | 'Len' >> beam.Map(lambda x: len(x)).with_output_types(int)) self.p.run() def test_pardo_type_inference(self): self.assertEqual(int, beam.Filter(lambda x: False).infer_output_type(int)) self.assertEqual( typehints.Tuple[str, int], beam.Map(lambda x: (x, 1)).infer_output_type(str)) def test_gbk_type_inference(self): self.assertEqual( typehints.Tuple[str, typehints.Iterable[int]], beam.GroupByKey().infer_output_type(typehints.KV[str, int])) def test_pipeline_inference(self): created = self.p | beam.Create(['a', 'b', 'c']) mapped = created | 'pair with 1' >> beam.Map(lambda x: (x, 1)) grouped = mapped | beam.GroupByKey() self.assertEqual(str, created.element_type) self.assertEqual(typehints.KV[str, int], mapped.element_type) self.assertEqual( typehints.KV[str, typehints.Iterable[int]], grouped.element_type) def test_inferred_bad_kv_type(self): with self.assertRaises(typehints.TypeCheckError) as e: _ = ( self.p | beam.Create(['a', 'b', 'c']) | 'Ungroupable' >> beam.Map(lambda x: (x, 0, 1.0)) | beam.GroupByKey()) self.assertStartswith( e.exception.args[0], "Input type hint violation at GroupByKey: " "expected Tuple[TypeVariable[K], TypeVariable[V]], " "got Tuple[, , ]") def test_type_inference_command_line_flag_toggle(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False x = self.p | 'C1' >> beam.Create([1, 2, 3, 4]) self.assertIsNone(x.element_type) self.p._options.view_as(TypeOptions).pipeline_type_check = True x = self.p | 'C2' >> beam.Create([1, 2, 3, 4]) self.assertEqual(int, x.element_type) def test_eager_execution(self): doubled = [1, 2, 3, 4] | beam.Map(lambda x: 2 * x) self.assertEqual([2, 4, 6, 8], doubled) def test_eager_execution_tagged_outputs(self): result = [1, 2, 3, 4] | beam.Map( lambda x: pvalue.TaggedOutput('bar', 2 * x)).with_outputs('bar') self.assertEqual([2, 4, 6, 8], result.bar) with self.assertRaises(KeyError, msg='Tag \'foo\' is not a defined output tag'): result.foo @parameterized_class([{'use_subprocess': False}, {'use_subprocess': True}]) class DeadLettersTest(unittest.TestCase): @classmethod def die(cls, x): if cls.use_subprocess: os._exit(x) else: raise ValueError(x) @classmethod def die_if_negative(cls, x): if x < 0: cls.die(x) else: return x @classmethod def exception_if_negative(cls, x): if x < 0: raise ValueError(x) else: return x @classmethod def die_if_less(cls, x, bound=0): if x < bound: cls.die(x) else: return x, bound def test_error_messages(self): with TestPipeline() as p: good, bad = ( p | beam.Create([-1, 10, -100, 2, 0]) | beam.Map(self.exception_if_negative).with_exception_handling()) assert_that(good, equal_to([0, 2, 10]), label='CheckGood') assert_that( bad | beam.MapTuple(lambda e, exc_info: (e, exc_info[1].replace(',', ''))), equal_to([(-1, 'ValueError(-1)'), (-100, 'ValueError(-100)')]), label='CheckBad') def test_filters_exceptions(self): with TestPipeline() as p: good, _ = ( p | beam.Create([-1, 10, -100, 2, 0]) | beam.Map(self.exception_if_negative).with_exception_handling( use_subprocess=self.use_subprocess, exc_class=(ValueError, TypeError))) assert_that(good, equal_to([0, 2, 10]), label='CheckGood') with self.assertRaises(Exception): with TestPipeline() as p: good, _ = ( p | beam.Create([-1, 10, -100, 2, 0]) | beam.Map(self.die_if_negative).with_exception_handling( use_subprocess=self.use_subprocess, exc_class=TypeError)) def test_tuples(self): with TestPipeline() as p: good, _ = ( p | beam.Create([(1, 2), (3, 2), (1, -10)]) | beam.MapTuple(self.die_if_less).with_exception_handling( use_subprocess=self.use_subprocess)) assert_that(good, equal_to([(3, 2), (1, -10)]), label='CheckGood') def test_side_inputs(self): with TestPipeline() as p: input = p | beam.Create([-1, 10, 100]) assert_that(( input | 'Default' >> beam.Map(self.die_if_less).with_exception_handling( use_subprocess=self.use_subprocess)).good, equal_to([(10, 0), (100, 0)]), label='CheckDefault') assert_that(( input | 'Pos' >> beam.Map(self.die_if_less, 20).with_exception_handling( use_subprocess=self.use_subprocess)).good, equal_to([(100, 20)]), label='PosSideInput') assert_that(( input | 'Key' >> beam.Map(self.die_if_less, bound=30).with_exception_handling( use_subprocess=self.use_subprocess)).good, equal_to([(100, 30)]), label='KeySideInput') def test_multiple_outputs(self): die = type(self).die def die_on_negative_even_odd(x): if x < 0: die(x) elif x % 2 == 0: return pvalue.TaggedOutput('even', x) elif x % 2 == 1: return pvalue.TaggedOutput('odd', x) with TestPipeline() as p: results = ( p | beam.Create([1, -1, 2, -2, 3]) | beam.Map(die_on_negative_even_odd).with_exception_handling( use_subprocess=self.use_subprocess)) assert_that(results.even, equal_to([2]), label='CheckEven') assert_that(results.odd, equal_to([1, 3]), label='CheckOdd') def test_params(self): die = type(self).die def die_if_negative_with_timestamp(x, ts=beam.DoFn.TimestampParam): if x < 0: die(x) else: return x, ts with TestPipeline() as p: good, _ = ( p | beam.Create([-1, 0, 1]) | beam.Map(lambda x: TimestampedValue(x, x)) | beam.Map(die_if_negative_with_timestamp).with_exception_handling( use_subprocess=self.use_subprocess)) assert_that(good, equal_to([(0, Timestamp(0)), (1, Timestamp(1))])) def test_timeout(self): import time timeout = 1 if self.use_subprocess else .1 with TestPipeline() as p: good, bad = ( p | beam.Create('records starting with lowercase S are slow'.split()) | beam.Map( lambda x: time.sleep(2.5 * timeout) if x.startswith('s') else x) .with_exception_handling( use_subprocess=self.use_subprocess, timeout=timeout)) assert_that( good, equal_to(['records', 'with', 'lowercase', 'S', 'are']), label='CheckGood') assert_that( bad | beam.MapTuple(lambda e, exc_info: (e, exc_info[1].replace(',', ''))), equal_to([('starting', 'TimeoutError()'), ('slow', 'TimeoutError()')]), label='CheckBad') def test_increment_counter(self): # Counters are not currently supported for # ParDo#with_exception_handling(use_subprocess=True). if (self.use_subprocess): return class CounterDoFn(beam.DoFn): def __init__(self): self.records_counter = Metrics.counter(self.__class__, 'recordsCounter') def process(self, element): self.records_counter.inc() with TestPipeline() as p: _, _ = ( (p | beam.Create([1,2,3])) | beam.ParDo(CounterDoFn()) .with_exception_handling( use_subprocess=self.use_subprocess, timeout=1)) results = p.result metric_results = results.metrics().query( MetricsFilter().with_name("recordsCounter")) records_counter = metric_results['counters'][0] self.assertEqual(records_counter.key.metric.name, 'recordsCounter') self.assertEqual(records_counter.result, 3) def test_lifecycle(self): die = type(self).die class MyDoFn(beam.DoFn): state = None def setup(self): assert self.state is None self.state = 'setup' def start_bundle(self): assert self.state in ('setup', 'finish_bundle'), self.state self.state = 'start_bundle' def finish_bundle(self): assert self.state in ('start_bundle', ), self.state self.state = 'finish_bundle' def teardown(self): assert self.state in ('setup', 'finish_bundle'), self.state self.state = 'teardown' def process(self, x): if x < 0: die(x) else: yield self.state with TestPipeline() as p: good, _ = ( p | beam.Create([-1, 0, 1, -10, 10]) | beam.ParDo(MyDoFn()).with_exception_handling( use_subprocess=self.use_subprocess)) assert_that(good, equal_to(['start_bundle'] * 3)) def test_partial(self): if self.use_subprocess: self.skipTest('Subprocess and partial mutally exclusive.') def die_if_negative_iter(elements): for element in elements: if element < 0: raise ValueError(element) yield element with TestPipeline() as p: input = p | beam.Create([(-1, 1, 11), (2, -2, 22), (3, 33, -3), (4, 44)]) assert_that(( input | 'Partial' >> beam.FlatMap( die_if_negative_iter).with_exception_handling(partial=True)).good, equal_to([2, 3, 33, 4, 44]), 'CheckPartial') assert_that(( input | 'Complete' >> beam.FlatMap(die_if_negative_iter). with_exception_handling(partial=False)).good, equal_to([4, 44]), 'CheckComplete') def test_threshold(self): # The threshold is high enough. with TestPipeline() as p: _ = ( p | beam.Create([-1, -2, 0, 1, 2, 3, 4, 5]) | beam.Map(self.die_if_negative).with_exception_handling( threshold=0.5, use_subprocess=self.use_subprocess)) # The threshold is too low enough. with self.assertRaisesRegex(Exception, "2 / 8 = 0.25 > 0.1"): with TestPipeline() as p: _ = ( p | beam.Create([-1, -2, 0, 1, 2, 3, 4, 5]) | beam.Map(self.die_if_negative).with_exception_handling( threshold=0.1, use_subprocess=self.use_subprocess)) # The threshold is too low per window. with self.assertRaisesRegex(Exception, "2 / 2 = 1.0 > 0.5"): with TestPipeline() as p: _ = ( p | beam.Create([-1, -2, 0, 1, 2, 3, 4, 5]) | beam.Map(lambda x: TimestampedValue(x, x)) | beam.Map(self.die_if_negative).with_exception_handling( threshold=0.5, threshold_windowing=window.FixedWindows(10), use_subprocess=self.use_subprocess)) class PTransformTypeAliasTest(unittest.TestCase): @unittest.skipIf(sys.version_info < (3, 12), "Python 3.12 required") def test_type_alias_statement_supported_in_with_output_types(self): ns = {} exec("type InputType = tuple[int, ...]", ns) # pylint: disable=exec-used InputType = ns["InputType"] def print_element(element: InputType) -> InputType: return element with beam.Pipeline() as p: _ = ( p | beam.Create([(1, 2)]) | beam.Map(lambda x: x) | beam.Map(print_element)) @unittest.skipIf(sys.version_info < (3, 12), "Python 3.12 required") def test_type_alias_supported_in_ptransform_with_output_types(self): ns = {} exec("type OutputType = tuple[int, int]", ns) # pylint: disable=exec-used OutputType = ns["OutputType"] with beam.Pipeline() as p: _ = ( p | beam.Create([(1, 2)]) | beam.Map(lambda x: x) | beam.Map(lambda x: x).with_output_types(OutputType)) class TestPTransformFn(TypeHintTestCase): def test_type_checking_fail(self): @beam.ptransform_fn def MyTransform(pcoll): return pcoll | beam.ParDo(lambda x: [x]).with_output_types(str) p = TestPipeline() with self.assertRaisesRegex(beam.typehints.TypeCheckError, r'expected.*int.*got.*str'): _ = (p | beam.Create([1, 2]) | MyTransform().with_output_types(int)) def test_type_checking_success(self): @beam.ptransform_fn def MyTransform(pcoll): return pcoll | beam.ParDo(lambda x: [x]).with_output_types(int) with TestPipeline() as p: _ = (p | beam.Create([1, 2]) | MyTransform().with_output_types(int)) def test_type_hints_arg(self): # Tests passing type hints via the magic 'type_hints' argument name. @beam.ptransform_fn def MyTransform(pcoll, type_hints, test_arg): self.assertEqual(test_arg, 'test') return ( pcoll | beam.ParDo(lambda x: [x]).with_output_types( type_hints.output_types[0][0])) with TestPipeline() as p: _ = (p | beam.Create([1, 2]) | MyTransform('test').with_output_types(int)) class PickledObject(object): def __init__(self, value): self.value = value if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/py_dataflow_distribution_counter.py0000644000175100017510000001136315206113033027301 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """For internal use only; no backwards-compatibility guarantees.""" # pytype: skip-file globals()['INT64_MAX'] = 2**63 - 1 globals()['INT64_MIN'] = -2**63 POWER_TEN = [ 10e-1, 10e0, 10e1, 10e2, 10e3, 10e4, 10e5, 10e6, 10e7, 10e8, 10e9, 10e10, 10e11, 10e12, 10e13, 10e14, 10e15, 10e16, 10e17, 10e18 ] def get_log10_round_to_floor(element): power = 0 while element >= POWER_TEN[power]: power += 1 return power - 1 class DataflowDistributionCounter(object): """Pure python DataflowDistributionCounter in case Cython not available. Please avoid using python mode if possible, since it's super slow Cythonized DatadflowDistributionCounter defined in apache_beam.transforms.cy_dataflow_distribution_counter. Currently using special bucketing strategy suitable for Dataflow Attributes: min: minimum value of all inputs. max: maximum value of all inputs. count: total count of all inputs. sum: sum of all inputs. buckets: histogram buckets of value counts for a distribution(1,2,5 bucketing). Max bucket_index is 58( sys.maxint as input). is_cythonized: mark whether DataflowDistributionCounter cythonized. """ # Assume the max input is sys.maxint, then the possible max bucket size is 59 MAX_BUCKET_SIZE = 59 # 3 buckets for every power of ten -> 1, 2, 5 BUCKET_PER_TEN = 3 def __init__(self): global INT64_MAX # pylint: disable=global-variable-not-assigned self.min = INT64_MAX self.max = 0 self.count = 0 self.sum = 0 self.buckets = [0] * self.MAX_BUCKET_SIZE self.is_cythonized = False def add_input(self, element): if element < 0: raise ValueError('Distribution counters support only non-negative value') self.min = min(self.min, element) self.max = max(self.max, element) self.count += 1 self.sum += element bucket_index = self.calculate_bucket_index(element) self.buckets[bucket_index] += 1 def add_input_n(self, element, n): if element < 0: raise ValueError('Distribution counters support only non-negative value') self.min = min(self.min, element) self.max = max(self.max, element) self.count += n self.sum += element * n bucket_index = self.calculate_bucket_index(element) self.buckets[bucket_index] += n def calculate_bucket_index(self, element): """Calculate the bucket index for the given element.""" if element == 0: return 0 log10_floor = get_log10_round_to_floor(element) power_of_ten = POWER_TEN[log10_floor] if element < power_of_ten * 2: bucket_offset = 0 elif element < power_of_ten * 5: bucket_offset = 1 else: bucket_offset = 2 return 1 + log10_floor * self.BUCKET_PER_TEN + bucket_offset def translate_to_histogram(self, histogram): """Translate buckets into Histogram. Args: histogram: apache_beam.runners.dataflow.internal.clents.dataflow.Histogram Ideally, only call this function when reporting counter to dataflow service. """ first_bucket_offset = 0 last_bucket_offset = 0 for index in range(0, self.MAX_BUCKET_SIZE): if self.buckets[index] != 0: first_bucket_offset = index break for index in range(self.MAX_BUCKET_SIZE - 1, -1, -1): if self.buckets[index] != 0: last_bucket_offset = index break histogram.firstBucketOffset = first_bucket_offset histogram.bucketCounts = ( self.buckets[first_bucket_offset:last_bucket_offset + 1]) def extract_output(self): global INT64_MIN # pylint: disable=global-variable-not-assigned global INT64_MAX # pylint: disable=global-variable-not-assigned if not INT64_MIN <= self.sum <= INT64_MAX: self.sum %= 2**64 if self.sum >= INT64_MAX: self.sum -= 2**64 mean = self.sum // self.count if self.count else float('nan') return mean, self.sum, self.count, self.min, self.max def merge(self, accumulators): raise NotImplementedError() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/resources.py0000644000175100017510000002034615206113033022445 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A module for defining resource requirements for execution of transforms. Pipeline authors can use resource hints to provide additional information to runners about the desired aspects of the execution environment. Resource hints can be specified on a transform level for parts of the pipeline, or globally via --resource_hint pipeline option. See also: PTransforms.with_resource_hints(). """ import re from collections.abc import Mapping from typing import Any from typing import Optional from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.portability.common_urns import resource_hints __all__ = [ 'ResourceHint', 'AcceleratorHint', 'MinRamHint', 'CpuCountHint', 'MaxActiveBundlesPerWorkerHint', 'merge_resource_hints', 'parse_resource_hints', 'resource_hints_from_options', ] class ResourceHint: """A superclass to define resource hints.""" # A unique URN, one per Resource Hint class. urn: Optional[str] = None _urn_to_known_hints: dict[str, type] = {} _name_to_known_hints: dict[str, type] = {} @classmethod def parse(cls, value: str) -> dict[str, bytes]: """Describes how to parse the hint. Override to specify a custom parsing logic.""" assert cls.urn is not None # Override this method to have a custom parsing logic. return {cls.urn: ResourceHint._parse_str(value)} @classmethod def get_merged_value(cls, outer_value: bytes, inner_value: bytes) -> bytes: """Reconciles values of a hint when the hint specified on a transform is also defined in an outer context, for example on a composite transform, or specified in the transform's execution environment. Override to specify a custom merging logic. """ # Defaults to the inner value as it is the most specific one. return inner_value @staticmethod def get_by_urn(urn): return ResourceHint._urn_to_known_hints[urn] @staticmethod def get_by_name(name): return ResourceHint._name_to_known_hints[name] @staticmethod def is_registered(name): return name in ResourceHint._name_to_known_hints @staticmethod def register_resource_hint(hint_name: str, hint_class: type) -> None: assert issubclass(hint_class, ResourceHint) assert hint_class.urn is not None ResourceHint._name_to_known_hints[hint_name] = hint_class ResourceHint._urn_to_known_hints[hint_class.urn] = hint_class @staticmethod def _parse_str(value): if not isinstance(value, str): raise ValueError("Input must be a string.") return value.encode('ascii') @staticmethod def _parse_int(value): if isinstance(value, str): value = int(value) if not isinstance(value, int): raise ValueError("Input must be an integer.") return str(value).encode('ascii') @staticmethod def _parse_storage_size_str(value): """Parses a human-friendly storage size string into a number of bytes. """ if isinstance(value, int): return ResourceHint._parse_int(value) if not isinstance(value, str): raise ValueError("Input must be a string or integer.") value = value.strip().replace(" ", "") units = { 'PiB': 2**50, 'TiB': 2**40, 'GiB': 2**30, 'MiB': 2**20, 'KiB': 2**10, 'PB': 10**15, 'TB': 10**12, 'GB': 10**9, 'MB': 10**6, 'KB': 10**3, 'B': 1, } match = re.match(r'.*?(\D+)$', value) if not match: raise ValueError("Unrecognized value pattern.") suffix = match.group(1) if suffix not in units: raise ValueError("Unrecognized unit.") multiplier = units[suffix] value = value[:-len(suffix)] return str(round(float(value) * multiplier)).encode('ascii') @staticmethod def _use_max(v1, v2): return str(max(int(v1), int(v2))).encode('ascii') @staticmethod def _use_sum(v1, v2): return str(int(v1) + int(v2)).encode('ascii') class AcceleratorHint(ResourceHint): """Describes desired hardware accelerators in execution environment.""" urn = resource_hints.ACCELERATOR.urn ResourceHint.register_resource_hint('accelerator', AcceleratorHint) class MinRamHint(ResourceHint): """Describes min RAM requirements for transform's execution environment.""" urn = resource_hints.MIN_RAM_BYTES.urn @classmethod def parse(cls, value: str) -> dict[str, bytes]: return {cls.urn: ResourceHint._parse_storage_size_str(value)} @classmethod def get_merged_value(cls, outer_value: bytes, inner_value: bytes) -> bytes: return ResourceHint._use_max(outer_value, inner_value) ResourceHint.register_resource_hint('min_ram', MinRamHint) # Alias for interoperability with SDKs preferring camelCase. ResourceHint.register_resource_hint('minRam', MinRamHint) class CpuCountHint(ResourceHint): """Describes number of CPUs available in transform's execution environment.""" urn = resource_hints.CPU_COUNT.urn @classmethod def get_merged_value(cls, outer_value: bytes, inner_value: bytes) -> bytes: return ResourceHint._use_max(outer_value, inner_value) ResourceHint.register_resource_hint('cpu_count', CpuCountHint) # Alias for interoperability with SDKs preferring camelCase. ResourceHint.register_resource_hint('cpuCount', CpuCountHint) class MaxActiveBundlesPerWorkerHint(ResourceHint): """ Describes max active bundles processed in parallel in transform's execution environment. """ urn = resource_hints.MAX_ACTIVE_BUNDLES_PER_WORKER.urn @classmethod def get_merged_value(cls, outer_value: bytes, inner_value: bytes) -> bytes: return ResourceHint._use_sum(outer_value, inner_value) ResourceHint.register_resource_hint( 'max_active_bundles_per_worker', MaxActiveBundlesPerWorkerHint) # Alias for interoperability with SDKs preferring camelCase. ResourceHint.register_resource_hint( 'MaxActiveBundlesPerWorker', MaxActiveBundlesPerWorkerHint) # Alias for common typo. ResourceHint.register_resource_hint( 'max_active_bundle_per_worker', MaxActiveBundlesPerWorkerHint) ResourceHint.register_resource_hint( 'MaxActiveBundlePerWorker', MaxActiveBundlesPerWorkerHint) def parse_resource_hints(hints: dict[Any, Any]) -> dict[str, bytes]: parsed_hints = {} for hint, value in hints.items(): try: hint_cls = ResourceHint.get_by_name(hint) try: parsed_hints.update(hint_cls.parse(value)) except ValueError: raise ValueError(f"Resource hint {hint} has invalid value {value}.") except KeyError: raise ValueError(f"Unknown resource hint: {hint}.") return parsed_hints def resource_hints_from_options( options: Optional[PipelineOptions]) -> dict[str, bytes]: if options is None: return {} hints = {} option_specified_hints = options.view_as(StandardOptions).resource_hints if isinstance(option_specified_hints, dict): return parse_resource_hints(option_specified_hints) for hint in option_specified_hints: if '=' in hint: k, v = hint.split('=', maxsplit=1) hints[k] = v else: hints[hint] = None return parse_resource_hints(hints) def merge_resource_hints( outer_hints: Mapping[str, bytes], inner_hints: Mapping[str, bytes]) -> dict[str, bytes]: merged_hints = dict(inner_hints) for urn, outer_value in outer_hints.items(): if urn in inner_hints: merged_value = ResourceHint.get_by_urn(urn).get_merged_value( outer_value=outer_value, inner_value=inner_hints[urn]) else: merged_value = outer_value merged_hints[urn] = merged_value return merged_hints ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/resources_test.py0000644000175100017510000000753415206113033023510 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest from parameterized import param from parameterized import parameterized from apache_beam import PTransform from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.transforms.resources import ResourceHint class ResourcesTest(unittest.TestCase): @parameterized.expand([ param( name='min_ram', val='100 MiB', urn='beam:resources:min_ram_bytes:v1', bytestr=b'104857600'), param( name='minRam', val='100MB', urn='beam:resources:min_ram_bytes:v1', bytestr=b'100000000'), param( name='min_ram', val='6.5 GiB', urn='beam:resources:min_ram_bytes:v1', bytestr=b'6979321856'), param( name='accelerator', val='gpu', urn='beam:resources:accelerator:v1', bytestr=b'gpu'), param( name='cpu_count', val='4', urn='beam:resources:cpu_count:v1', bytestr=b'4'), param( name='max_active_bundles_per_worker', val='2', urn='beam:resources:max_active_bundles_per_worker:v1', bytestr=b'2'), param( name='max_active_bundle_per_worker', val='20', urn='beam:resources:max_active_bundles_per_worker:v1', bytestr=b'20'), param( name='MaxActiveBundlePerWorker', val='30', urn='beam:resources:max_active_bundles_per_worker:v1', bytestr=b'30'), param( name='MaxActiveBundlesPerWorker', val='3', urn='beam:resources:max_active_bundles_per_worker:v1', bytestr=b'3'), ]) def test_known_resource_hints(self, name, val, urn, bytestr): t = PTransform() t = t.with_resource_hints(**{name: val}) self.assertTrue(ResourceHint.is_registered(name)) self.assertEqual(t.get_resource_hints(), {urn: bytestr}) @parameterized.expand([ param(name='min_ram', val='3,500G'), param(name='accelerator', val=1), param(name='cpu_count', val=1), param(name='unknown_hint', val=1) ]) def test_resource_hint_parsing_fails_early(self, name, val): t = PTransform() with self.assertRaises(ValueError): _ = t.with_resource_hints(**{name: val}) def test_resource_hints_from_options(self): options = PipelineOptions() standard_options = options.view_as(StandardOptions) standard_options.resource_hints = { "min_ram": "16GB", "accelerator": "gpu", "cpu_count": "4", "max_active_bundles_per_worker": "2" } p = TestPipeline(options=options) self.assertEqual( p._root_transform().resource_hints, { 'beam:resources:min_ram_bytes:v1': b'16000000000', 'beam:resources:accelerator:v1': b'gpu', 'beam:resources:cpu_count:v1': b'4', 'beam:resources:max_active_bundles_per_worker:v1': b'2' }) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/sideinputs.py0000644000175100017510000000727115206113033022624 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Internal side input transforms and implementations. For internal use only; no backwards-compatibility guarantees. Important: this module is an implementation detail and should not be used directly by pipeline writers. Instead, users should use the helper methods AsSingleton, AsIter, AsList and AsDict in apache_beam.pvalue. """ # pytype: skip-file import re from collections.abc import Callable from typing import TYPE_CHECKING from typing import Any from apache_beam.transforms import window if TYPE_CHECKING: from apache_beam import pvalue WindowMappingFn = Callable[[window.BoundedWindow], window.BoundedWindow] SIDE_INPUT_PREFIX = 'python_side_input' SIDE_INPUT_REGEX = SIDE_INPUT_PREFIX + '([0-9]+)(-.*)?$' # Top-level function so we can identify it later. def _global_window_mapping_fn( w, global_window=window.GlobalWindow()) -> window.GlobalWindow: return global_window def default_window_mapping_fn( target_window_fn: window.WindowFn) -> WindowMappingFn: if target_window_fn == window.GlobalWindows(): return _global_window_mapping_fn if isinstance(target_window_fn, window.Sessions): raise RuntimeError("Sessions is not allowed in side inputs") def map_via_end(source_window: window.BoundedWindow) -> window.BoundedWindow: return list( target_window_fn.assign( window.WindowFn.AssignContext( source_window.max_timestamp(), window=source_window)))[-1] return map_via_end def get_sideinput_index(tag: str) -> int: match = re.match(SIDE_INPUT_REGEX, tag, re.DOTALL) if match: return int(match.group(1)) else: raise RuntimeError("Invalid tag %r" % tag) class SideInputMap(object): """Represents a mapping of windows to side input values.""" def __init__(self, view_class: 'pvalue.AsSideInput', view_options, iterable): self._window_mapping_fn = view_options.get( 'window_mapping_fn', _global_window_mapping_fn) self._view_class = view_class self._view_options = view_options self._iterable = iterable self._cache: dict[window.BoundedWindow, Any] = {} def __getitem__(self, window: window.BoundedWindow) -> Any: if window not in self._cache: target_window = self._window_mapping_fn(window) self._cache[window] = self._view_class._from_runtime_iterable( _FilteringIterable(self._iterable, target_window), self._view_options) return self._cache[window] def is_globally_windowed(self) -> bool: return self._window_mapping_fn == _global_window_mapping_fn class _FilteringIterable(object): """An iterable containing only those values in the given window. """ def __init__(self, iterable, target_window): self._iterable = iterable self._target_window = target_window def __iter__(self): for wv in self._iterable: if self._target_window in wv.windows: yield wv.value def __reduce__(self): # Pickle self as an already filtered list. return list, (list(self), ) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/sideinputs_test.py0000644000175100017510000005040415206113033023657 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for side inputs.""" # pytype: skip-file import hashlib import itertools import logging import unittest from typing import Any from typing import Iterable from typing import Union import pytest import apache_beam as beam from apache_beam.testing.synthetic_pipeline import SyntheticSDFAsSource from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.testing.util import equal_to_per_window from apache_beam.transforms import Map from apache_beam.transforms import sideinputs from apache_beam.transforms import trigger from apache_beam.transforms import window from apache_beam.utils.timestamp import Timestamp class SideInputsTest(unittest.TestCase): def create_pipeline(self): return TestPipeline() def run_windowed_side_inputs( self, elements, main_window_fn, side_window_fn=None, side_input_type=beam.pvalue.AsList, combine_fn=None, expected=None): with self.create_pipeline() as p: pcoll = p | beam.Create(elements) | beam.Map( lambda t: window.TimestampedValue(t, t)) main = pcoll | 'WindowMain' >> beam.WindowInto(main_window_fn) side = pcoll | 'WindowSide' >> beam.WindowInto( side_window_fn or main_window_fn) kw = {} if combine_fn is not None: side |= beam.CombineGlobally(combine_fn).without_defaults() kw['default_value'] = 0 elif side_input_type == beam.pvalue.AsDict: side |= beam.Map(lambda x: ('k%s' % x, 'v%s' % x)) res = main | beam.Map(lambda x, s: (x, s), side_input_type(side, **kw)) if side_input_type in (beam.pvalue.AsIter, beam.pvalue.AsList): res |= beam.Map(lambda x_s: (x_s[0], sorted(x_s[1]))) assert_that(res, equal_to(expected)) def test_global_global_windows(self): self.run_windowed_side_inputs([1, 2, 3], window.GlobalWindows(), expected=[(1, [1, 2, 3]), (2, [1, 2, 3]), (3, [1, 2, 3])]) def test_same_fixed_windows(self): self.run_windowed_side_inputs([1, 2, 11], window.FixedWindows(10), expected=[(1, [1, 2]), (2, [1, 2]), (11, [11])]) def test_different_fixed_windows(self): self.run_windowed_side_inputs([1, 2, 11, 21, 31], window.FixedWindows(10), window.FixedWindows(20), expected=[(1, [1, 2, 11]), (2, [1, 2, 11]), (11, [1, 2, 11]), (21, [21, 31]), (31, [21, 31])]) def test_fixed_global_window(self): self.run_windowed_side_inputs([1, 2, 11], window.FixedWindows(10), window.GlobalWindows(), expected=[(1, [1, 2, 11]), (2, [1, 2, 11]), (11, [1, 2, 11])]) def test_sliding_windows(self): self.run_windowed_side_inputs( [1, 2, 4], window.SlidingWindows(size=6, period=2), window.SlidingWindows(size=6, period=2), expected=[ # Element 1 falls in three windows (1, [1]), # [-4, 2) (1, [1, 2]), # [-2, 4) (1, [1, 2, 4]), # [0, 6) # as does 2, (2, [1, 2]), # [-2, 4) (2, [1, 2, 4]), # [0, 6) (2, [2, 4]), # [2, 8) # and 4. (4, [1, 2, 4]), # [0, 6) (4, [2, 4]), # [2, 8) (4, [4]), # [4, 10) ]) def test_windowed_iter(self): self.run_windowed_side_inputs([1, 2, 11], window.FixedWindows(10), side_input_type=beam.pvalue.AsIter, expected=[(1, [1, 2]), (2, [1, 2]), (11, [11])]) def test_windowed_singleton(self): self.run_windowed_side_inputs([1, 2, 11], window.FixedWindows(10), side_input_type=beam.pvalue.AsSingleton, combine_fn=sum, expected=[(1, 3), (2, 3), (11, 11)]) def test_windowed_dict(self): self.run_windowed_side_inputs([1, 2, 11], window.FixedWindows(10), side_input_type=beam.pvalue.AsDict, expected=[ (1, { 'k1': 'v1', 'k2': 'v2' }), (2, { 'k1': 'v1', 'k2': 'v2' }), (11, { 'k11': 'v11' }), ]) @pytest.mark.it_validatesrunner def test_empty_singleton_side_input(self): pipeline = self.create_pipeline() pcol = pipeline | 'start' >> beam.Create([1, 2]) side = pipeline | 'side' >> beam.Create([]) # Empty side input. def my_fn(k, s): # TODO(robertwb): Should this be an error as in Java? v = ('empty' if isinstance(s, beam.pvalue.EmptySideInput) else 'full') return [(k, v)] result = pcol | 'compute' >> beam.FlatMap( my_fn, beam.pvalue.AsSingleton(side)) assert_that(result, equal_to([(1, 'empty'), (2, 'empty')])) pipeline.run() # TODO(https://github.com/apache/beam/issues/19012): Disable this test in # streaming temporarily. Remove sickbay-streaming tag after it's fixed. @pytest.mark.no_sickbay_streaming @pytest.mark.it_validatesrunner def test_multi_valued_singleton_side_input(self): pipeline = self.create_pipeline() pcol = pipeline | 'start' >> beam.Create([1, 2]) side = pipeline | 'side' >> beam.Create([3, 4]) # 2 values in side input. pcol | 'compute' >> beam.FlatMap( # pylint: disable=expression-not-assigned lambda x, s: [x * s], beam.pvalue.AsSingleton(side)) with self.assertRaises(Exception): pipeline.run() @pytest.mark.it_validatesrunner def test_default_value_singleton_side_input(self): pipeline = self.create_pipeline() pcol = pipeline | 'start' >> beam.Create([1, 2]) side = pipeline | 'side' >> beam.Create([]) # 0 values in side input. result = pcol | beam.FlatMap( lambda x, s: [x * s], beam.pvalue.AsSingleton(side, 10)) assert_that(result, equal_to([10, 20])) pipeline.run() @pytest.mark.it_validatesrunner def test_iterable_side_input(self): pipeline = self.create_pipeline() pcol = pipeline | 'start' >> beam.Create([1, 2]) side = pipeline | 'side' >> beam.Create([3, 4]) # 2 values in side input. result = pcol | 'compute' >> beam.FlatMap( lambda x, s: [x * y for y in s], beam.pvalue.AsIter(side)) assert_that(result, equal_to([3, 4, 6, 8])) pipeline.run() @pytest.mark.it_validatesrunner def test_reiterable_side_input(self): expected_side = frozenset(range(100)) def check_reiteration(main, side): assert expected_side == set(side), side # Iterate a second time. assert expected_side == set(side), side # Iterate over two copies of the input at the same time. both = zip(side, side) first, second = zip(*both) assert expected_side == set(first), first assert expected_side == set(second), second # This will iterate over two copies of the side input, but offset. offset = [None] * (len(expected_side) // 2) both = zip(itertools.chain(side, offset), itertools.chain(offset, side)) first, second = zip(*both) expected_and_none = frozenset.union(expected_side, [None]) assert expected_and_none == set(first), first assert expected_and_none == set(second), second pipeline = self.create_pipeline() pcol = pipeline | 'start' >> beam.Create(['A', 'B']) side = pipeline | 'side' >> beam.Create(expected_side) _ = pcol | 'check' >> beam.Map(check_reiteration, beam.pvalue.AsIter(side)) pipeline.run() @pytest.mark.it_validatesrunner def test_as_list_and_as_dict_side_inputs(self): a_list = [5, 1, 3, 2, 9] some_pairs = [('crouton', 17), ('supreme', None)] pipeline = self.create_pipeline() main_input = pipeline | 'main input' >> beam.Create([1]) side_list = pipeline | 'side list' >> beam.Create(a_list) side_pairs = pipeline | 'side pairs' >> beam.Create(some_pairs) results = main_input | 'concatenate' >> beam.Map( lambda x, the_list, the_dict: [x, the_list, the_dict], beam.pvalue.AsList(side_list), beam.pvalue.AsDict(side_pairs)) def matcher(expected_elem, expected_list, expected_pairs): def match(actual): [[actual_elem, actual_list, actual_dict]] = actual equal_to([expected_elem])([actual_elem]) equal_to(expected_list)(actual_list) equal_to(expected_pairs)(actual_dict.items()) return match assert_that(results, matcher(1, a_list, some_pairs)) pipeline.run() @pytest.mark.it_validatesrunner def test_as_singleton_without_unique_labels(self): # This should succeed as calling beam.pvalue.AsSingleton on the same # PCollection twice with the same defaults will return the same # view. a_list = [2] pipeline = self.create_pipeline() main_input = pipeline | 'main input' >> beam.Create([1]) side_list = pipeline | 'side list' >> beam.Create(a_list) results = main_input | beam.Map( lambda x, s1, s2: [x, s1, s2], beam.pvalue.AsSingleton(side_list), beam.pvalue.AsSingleton(side_list)) def matcher(expected_elem, expected_singleton): def match(actual): [[actual_elem, actual_singleton1, actual_singleton2]] = actual equal_to([expected_elem])([actual_elem]) equal_to([expected_singleton])([actual_singleton1]) equal_to([expected_singleton])([actual_singleton2]) return match assert_that(results, matcher(1, 2)) pipeline.run() @pytest.mark.it_validatesrunner def test_as_singleton_with_different_defaults(self): a_list = [] pipeline = self.create_pipeline() main_input = pipeline | 'main input' >> beam.Create([1]) side_list = pipeline | 'side list' >> beam.Create(a_list) results = main_input | beam.Map( lambda x, s1, s2: [x, s1, s2], beam.pvalue.AsSingleton(side_list, default_value=2), beam.pvalue.AsSingleton(side_list, default_value=3)) def matcher(expected_elem, expected_singleton1, expected_singleton2): def match(actual): [[actual_elem, actual_singleton1, actual_singleton2]] = actual equal_to([expected_elem])([actual_elem]) equal_to([expected_singleton1])([actual_singleton1]) equal_to([expected_singleton2])([actual_singleton2]) return match assert_that(results, matcher(1, 2, 3)) pipeline.run() @pytest.mark.it_validatesrunner def test_as_list_twice(self): # This should succeed as calling beam.pvalue.AsList on the same # PCollection twice will return the same view. a_list = [1, 2, 3] pipeline = self.create_pipeline() main_input = pipeline | 'main input' >> beam.Create([1]) side_list = pipeline | 'side list' >> beam.Create(a_list) results = main_input | beam.Map( lambda x, ls1, ls2: [x, ls1, ls2], beam.pvalue.AsList(side_list), beam.pvalue.AsList(side_list)) def matcher(expected_elem, expected_list): def match(actual): [[actual_elem, actual_list1, actual_list2]] = actual equal_to([expected_elem])([actual_elem]) equal_to(expected_list)(actual_list1) equal_to(expected_list)(actual_list2) return match assert_that(results, matcher(1, [1, 2, 3])) pipeline.run() @pytest.mark.it_validatesrunner def test_as_dict_twice(self): some_kvs = [('a', 1), ('b', 2)] pipeline = self.create_pipeline() main_input = pipeline | 'main input' >> beam.Create([1]) side_kvs = pipeline | 'side kvs' >> beam.Create(some_kvs) results = main_input | beam.Map( lambda x, dct1, dct2: [x, dct1, dct2], beam.pvalue.AsDict(side_kvs), beam.pvalue.AsDict(side_kvs)) def matcher(expected_elem, expected_kvs): def match(actual): [[actual_elem, actual_dict1, actual_dict2]] = actual equal_to([expected_elem])([actual_elem]) equal_to(expected_kvs)(actual_dict1.items()) equal_to(expected_kvs)(actual_dict2.items()) return match assert_that(results, matcher(1, some_kvs)) pipeline.run() @pytest.mark.it_validatesrunner def test_flattened_side_input(self): pipeline = self.create_pipeline() main_input = pipeline | 'main input' >> beam.Create([None]) side_input = ( pipeline | 'side1' >> beam.Create(['a']), pipeline | 'side2' >> beam.Create(['b'])) | beam.Flatten() results = main_input | beam.FlatMap( lambda _, ab: ab, beam.pvalue.AsList(side_input)) assert_that(results, equal_to(['a', 'b'])) pipeline.run() # TODO(BEAM-9499): Disable this test in streaming temporarily. @pytest.mark.no_sickbay_batch @pytest.mark.no_sickbay_streaming @pytest.mark.it_validatesrunner def test_multi_triggered_gbk_side_input(self): """Test a GBK sideinput, with multiple triggering.""" # TODO(https://github.com/apache/beam/issues/20065): Remove use of this # experiment. This flag is only necessary when using the multi-output # TestStream b/c it relies on using the PCollection output tags as the # PCollection output ids. with TestPipeline() as p: test_stream = ( p | 'Mixed TestStream' >> TestStream().advance_watermark_to( 3, tag='main').add_elements(['a1'], tag='main').advance_watermark_to( 8, tag='main').add_elements(['a2'], tag='main').add_elements( [window.TimestampedValue(('k', 100), 2)], tag='side'). add_elements([window.TimestampedValue( ('k', 400), 7)], tag='side').advance_watermark_to_infinity( tag='main').advance_watermark_to_infinity(tag='side')) main_data = ( test_stream['main'] | 'Main windowInto' >> beam.WindowInto( window.FixedWindows(5), accumulation_mode=trigger.AccumulationMode.DISCARDING)) side_data = ( test_stream['side'] | 'Side windowInto' >> beam.WindowInto( window.FixedWindows(5), trigger=trigger.AfterWatermark(early=trigger.AfterCount(1)), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.CombinePerKey(sum) | 'Values' >> Map(lambda k_vs: k_vs[1])) class RecordFn(beam.DoFn): def process( self, elm=beam.DoFn.ElementParam, ts=beam.DoFn.TimestampParam, side=beam.DoFn.SideInputParam): yield (elm, ts, side) records = ( main_data | beam.ParDo(RecordFn(), beam.pvalue.AsList(side_data))) expected_window_to_elements = { window.IntervalWindow(0, 5): [ ('a1', Timestamp(3), [100, 0]), ], window.IntervalWindow(5, 10): [('a2', Timestamp(8), [400, 0])], } assert_that( records, equal_to_per_window(expected_window_to_elements), use_global_window=False, label='assert per window') @pytest.mark.it_validatesrunner def test_side_input_with_sdf(self): """Test a side input with SDF. This test verifies consisency of side input when it is split due to SDF (Splittable DoFns). The consistency is verified by checking the size and fingerprint of the side input. This test needs to run with at least 2 workers (--num_workers=2) and autoscaling disabled (--autoscaling_algorithm=NONE). Otherwise it might provide false positives (i.e. not fail on bad state). """ initial_elements = 1000 num_records = 10000 key_size = 10 value_size = 100 expected_fingerprint = '00f7eeac8514721e2683d14a504b33d1' class GetSyntheticSDFOptions(beam.DoFn): """A DoFn that emits elements for genenrating SDF.""" def process(self, element: Any) -> Iterable[dict[str, Union[int, str]]]: yield { 'num_records': num_records // initial_elements, 'key_size': key_size, 'value_size': value_size, 'initial_splitting_num_bundles': 0, 'initial_splitting_desired_bundle_size': 0, 'sleep_per_input_record_sec': 0, 'initial_splitting': 'const', } class SideInputTrackingDoFn(beam.DoFn): """A DoFn that emits the size and fingerprint of the side input. In this context, the size is the number of elements and the fingerprint is the hash of the sorted serialized elements. """ def process( self, element: Any, side_input: Iterable[tuple[bytes, bytes]]) -> Iterable[tuple[int, str]]: # Sort for consistent hashing. sorted_side_input = sorted(side_input) size = len(sorted_side_input) m = hashlib.md5() for key, value in sorted_side_input: m.update(key) m.update(value) yield (size, m.hexdigest()) pipeline = self.create_pipeline() main_input = pipeline | 'Main input: Create' >> beam.Create([0]) side_input = pipeline | 'Side input: Create' >> beam.Create( range(initial_elements)) side_input |= 'Side input: Get synthetic SDF options' >> beam.ParDo( GetSyntheticSDFOptions()) side_input |= 'Side input: Process and split' >> beam.ParDo( SyntheticSDFAsSource()) results = main_input | 'Emit side input' >> beam.ParDo( SideInputTrackingDoFn(), beam.pvalue.AsIter(side_input)) assert_that(results, equal_to([(num_records, expected_fingerprint)])) pipeline.run() def test_default_window_mapping_fn_source_window(self): """Test that the default window mapping function will propagate the source window when attempting to assign context. """ class StringIDWindow(window.BoundedWindow): """A window defined by an arbitrary string ID.""" def __init__(self, window_id: str): super().__init__(self._getTimestampFromProto()) self.id = window_id @staticmethod def _getTimestampFromProto() -> Timestamp: return Timestamp(micros=0) class StringIDWindows(window.NonMergingWindowFn): """ A windowing function that assigns each element a window with ID.""" def assign( self, assign_context: window.WindowFn.AssignContext ) -> Iterable[window.BoundedWindow]: if assign_context.element is None: assert assign_context.window is not None return [assign_context.window] return [StringIDWindow(str(assign_context.element))] def get_window_coder(self): return None mapping_fn = sideinputs.default_window_mapping_fn(StringIDWindows()) source_window = StringIDWindows().assign( window.WindowFn.AssignContext(Timestamp(10), element='element'))[0] bounded_window = mapping_fn(source_window) assert bounded_window is not None assert bounded_window.id == 'element' if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/sql.py0000644000175100017510000000754415206113033021237 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Package for SqlTransform and related classes.""" # pytype: skip-file import typing from apache_beam.transforms.external import BeamJarExpansionService from apache_beam.transforms.external import ExternalTransform from apache_beam.transforms.external import NamedTupleBasedPayloadBuilder __all__ = ['SqlTransform'] SqlTransformSchema = typing.NamedTuple( 'SqlTransformSchema', [('query', str), ('dialect', typing.Optional[str]), ('ddl', typing.Optional[str])]) class SqlTransform(ExternalTransform): """A transform that can translate a SQL query into PTransforms. Input PCollections must have a schema. Currently, there are two ways to define a schema for a PCollection: 1) Register a `typing.NamedTuple` type to use RowCoder, and specify it as the output type. For example:: Purchase = typing.NamedTuple('Purchase', [('item_name', unicode), ('price', float)]) coders.registry.register_coder(Purchase, coders.RowCoder) with Pipeline() as p: purchases = (p | beam.io... | beam.Map(..).with_output_types(Purchase)) 2) Produce `beam.Row` instances. Note this option will fail if Beam is unable to infer data types for any of the fields. For example:: with Pipeline() as p: purchases = (p | beam.io... | beam.Map(lambda x: beam.Row(item_name=unicode(..), price=float(..)))) Similarly, the output of SqlTransform is a PCollection with a schema. The columns produced by the query can be accessed as attributes. For example:: purchases | SqlTransform(\"\"\" SELECT item_name, COUNT(*) AS `count` FROM PCOLLECTION GROUP BY item_name\"\"\") | beam.Map(lambda row: "We've sold %d %ss!" % (row.count, row.item_name)) Additional examples can be found in `apache_beam.examples.wordcount_xlang_sql`, `apache_beam.examples.sql_taxi`, and `apache_beam.transforms.sql_test`. For more details about Beam SQL in general, see the `Java transform `_, and the `documentation `_. """ URN = 'beam:external:java:sql:v1' def __init__(self, query, dialect=None, ddl=None, expansion_service=None): """ Creates a SqlTransform which will be expanded to Java's SqlTransform. (See class docs). :param query: The SQL query. :param dialect: (optional, deprecated) The dialect. :param ddl: (optional) The DDL statement. :param expansion_service: (optional) The URL of the expansion service to use """ expansion_service = expansion_service or BeamJarExpansionService( ':sdks:java:extensions:sql:expansion-service:shadowJar') super().__init__( self.URN, NamedTupleBasedPayloadBuilder( SqlTransformSchema(query=query, dialect=dialect, ddl=ddl)), expansion_service=expansion_service) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/sql_test.py0000644000175100017510000002002215206113033022260 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for transforms that use the SQL Expansion service.""" # pytype: skip-file import logging import typing import unittest import pytest import apache_beam as beam from apache_beam import coders from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.sql import SqlTransform SimpleRow = typing.NamedTuple( "SimpleRow", [("id", int), ("str", str), ("flt", float)]) coders.registry.register_coder(SimpleRow, coders.RowCoder) Enrich = typing.NamedTuple("Enrich", [("id", int), ("metadata", str)]) coders.registry.register_coder(Enrich, coders.RowCoder) Shopper = typing.NamedTuple( "Shopper", [("shopper", str), ("cart", typing.Mapping[str, int])]) coders.registry.register_coder(Shopper, coders.RowCoder) class Aribitrary: def __init__(self, obj): self.obj = obj def __eq__(self, other): return self.obj == other.obj UserTypeRow = typing.NamedTuple( "UserTypeRow", [("id", int), ("arb", Aribitrary), ("complex", complex)]) coders.registry.register_coder(UserTypeRow, coders.RowCoder) @pytest.mark.xlang_sql_expansion_service @unittest.skipIf( TestPipeline().get_pipeline_options().view_as(StandardOptions).runner is None, "Must be run with a runner that supports staging java artifacts.") class SqlTransformTest(unittest.TestCase): """Tests that exercise the cross-language SqlTransform (implemented in java). Note this test must be executed with pipeline options that run jobs on a local job server. The easiest way to accomplish this is to run the `validatesCrossLanguageRunnerPythonUsingSql` gradle target for a particular job server, which will start the runner and job server for you. For example, `:runners:flink:1.13:job-server:validatesCrossLanguageRunnerPythonUsingSql` to test on Flink 1.13. Alternatively, you may be able to iterate faster if you run the tests directly using a runner like `FlinkRunner`, which can start a local Flink cluster and job server for you: $ pip install -e './sdks/python[gcp,test]' $ pytest apache_beam/transforms/sql_test.py \\ --test-pipeline-options="--runner=FlinkRunner" """ _multiprocess_can_split_ = True def test_generate_data(self): with TestPipeline() as p: out = p | SqlTransform( """SELECT CAST(1 AS INT) AS `id`, CAST('foo' AS VARCHAR) AS `str`, CAST(3.14 AS DOUBLE) AS `flt`""") assert_that(out, equal_to([(1, "foo", 3.14)])) def test_project(self): with TestPipeline() as p: out = ( p | beam.Create([SimpleRow(1, "foo", 3.14)]) | SqlTransform("SELECT `id`, `flt` FROM PCOLLECTION")) assert_that(out, equal_to([(1, 3.14)])) def test_filter(self): with TestPipeline() as p: out = ( p | beam.Create([SimpleRow(1, "foo", 3.14), SimpleRow(2, "bar", 1.414)]) | SqlTransform("SELECT * FROM PCOLLECTION WHERE `str` = 'bar'")) assert_that(out, equal_to([(2, "bar", 1.414)])) def test_agg(self): with TestPipeline() as p: out = ( p | beam.Create([ SimpleRow(1, "foo", 1.), SimpleRow(1, "foo", 2.), SimpleRow(1, "foo", 3.), SimpleRow(2, "bar", 1.414), SimpleRow(2, "bar", 1.414), SimpleRow(2, "bar", 1.414), SimpleRow(2, "bar", 1.414), ]) | SqlTransform( """ SELECT `str`, COUNT(*) AS `count`, SUM(`id`) AS `sum`, AVG(`flt`) AS `avg` FROM PCOLLECTION GROUP BY `str`""")) assert_that(out, equal_to([("foo", 3, 3, 2), ("bar", 4, 8, 1.414)])) def test_tagged_join(self): with TestPipeline() as p: enrich = ( p | "Create enrich" >> beam.Create( [Enrich(1, "a"), Enrich(2, "b"), Enrich(26, "z")])) simple = ( p | "Create simple" >> beam.Create([ SimpleRow(1, "foo", 3.14), SimpleRow(26, "bar", 1.11), SimpleRow(1, "baz", 2.34) ])) out = ({ 'simple': simple, 'enrich': enrich } | SqlTransform( """ SELECT simple.`id` AS `id`, enrich.metadata AS metadata FROM simple JOIN enrich ON simple.`id` = enrich.`id`""")) assert_that(out, equal_to([(1, "a"), (26, "z"), (1, "a")])) def test_row(self): with TestPipeline() as p: out = ( p | beam.Create([1, 2, 10]) | beam.Map(lambda x: beam.Row(a=x, b=str(x))) | SqlTransform("SELECT a*a as s, LENGTH(b) AS c FROM PCOLLECTION")) assert_that(out, equal_to([(1, 1), (4, 1), (100, 2)])) def test_row_user_type(self): with TestPipeline() as p: out = ( p | beam.Create([ UserTypeRow(1, Aribitrary(1.0), 1 + 2.5j), UserTypeRow(1, Aribitrary("abc"), -1j), ]) | SqlTransform("SELECT arb, complex FROM PCOLLECTION") | beam.Map(tuple)) assert_that( out, equal_to([(Aribitrary(1.0), 1 + 2.5j), (Aribitrary("abc"), -1j)])) def test_windowing_before_sql(self): with TestPipeline() as p: out = ( p | beam.Create([ SimpleRow(5, "foo", 1.), SimpleRow(15, "bar", 2.), SimpleRow(25, "baz", 3.) ]) | beam.Map(lambda v: beam.window.TimestampedValue(v, v.id)). with_output_types(SimpleRow) | beam.WindowInto( beam.window.FixedWindows(10)).with_output_types(SimpleRow) | SqlTransform("SELECT COUNT(*) as `count` FROM PCOLLECTION")) assert_that(out, equal_to([(1, ), (1, ), (1, )])) def test_map(self): with TestPipeline() as p: out = ( p | beam.Create([ Shopper('bob', { 'bananas': 6, 'cherries': 3 }), Shopper('alice', { 'apples': 2, 'bananas': 3 }) ]).with_output_types(Shopper) | SqlTransform("SELECT * FROM PCOLLECTION WHERE shopper = 'alice'")) assert_that(out, equal_to([('alice', {'apples': 2, 'bananas': 3})])) def test_sql_ddl_set_option(self): with TestPipeline() as p: input_data = [ beam.Row(id=1, value=10), beam.Row(id=2, value=20), beam.Row(id=3, value=30) ] # DDL uses SET to modify a session option (tests DDL parsing) # Using a known Calcite option like sqlConformance ddl_statement = """ SET sqlConformance = 'LENIENT' """ # Query still operates on the implicit PCOLLECTION query_statement = "SELECT * FROM PCOLLECTION WHERE id > 2" # Input PCollection is piped directly out = ( p | beam.Create(input_data) # Pass both the query and the DDL | SqlTransform(query=query_statement, ddl=ddl_statement)) # Verify the output matches the query (unaffected by the SET DDL) assert_that(out, equal_to([(3, 30)])) if __name__ == "__main__": logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/stats.pxd0000644000175100017510000000410715206113033021731 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # cimport cython from libc.stdint cimport int64_t cdef class _QuantileSpec(object): cdef readonly int64_t buffer_size cdef readonly int64_t num_buffers cdef readonly bint weighted cdef readonly key cdef readonly bint reverse cdef readonly weighted_key cdef readonly less_than cdef class _QuantileBuffer(object): cdef readonly elements cdef readonly weights cdef readonly bint weighted cdef readonly int64_t level cdef readonly min_val cdef readonly max_val cdef readonly _iter cdef class _QuantileState(object): cdef readonly _QuantileSpec spec cdef public buffers cdef public unbuffered_elements cdef public unbuffered_weights cdef public add_unbuffered cpdef bint is_empty(self) @cython.locals(num_new_buffers=int64_t, idx=int64_t) cpdef _add_unbuffered(self, elements, offset_fn) @cython.locals(num_new_buffers=int64_t, idx=int64_t) cpdef _add_unbuffered_weighted(self, elements, offset_fn) cpdef finalize(self) @cython.locals(min_level=int64_t) cpdef collapse_if_needed(self, offset_fn) @cython.locals(new_level=int64_t, new_weight=double, step=double, offset=double) cdef _QuantileBuffer _collapse(buffers, offset_fn, _QuantileSpec spec) @cython.locals(j=int64_t) cdef _interpolate(buffers, int64_t count, double step, double offset, _QuantileSpec spec)././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/stats.py0000644000175100017510000010635515206113033021576 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """This module has all statistic related transforms. This ApproximateUnique class will be deprecated [1]. PLease look into using HLLCount in the zetasketch extension module [2]. [1] https://lists.apache.org/thread.html/501605df5027567099b81f18c080469661fb426 4a002615fa1510502%40%3Cdev.beam.apache.org%3E [2] https://beam.apache.org/releases/javadoc/2.16.0/org/apache/beam/sdk/extensio ns/zetasketch/HllCount.html """ # pytype: skip-file # ruff: noqa: UP006 import hashlib import heapq import itertools import logging import math import typing from collections.abc import Callable from typing import Any from typing import List from typing import Tuple from apache_beam import coders from apache_beam import typehints from apache_beam.transforms.core import * from apache_beam.transforms.display import DisplayDataItem from apache_beam.transforms.ptransform import PTransform __all__ = [ 'ApproximateQuantiles', 'ApproximateUnique', ] # Type variables T = typing.TypeVar('T') K = typing.TypeVar('K') V = typing.TypeVar('V') try: import mmh3 # pylint: disable=import-error def _mmh3_hash(value): # mmh3.hash64 returns two 64-bit unsigned integers return mmh3.hash64(value, seed=0, signed=False)[0] _default_hash_fn = _mmh3_hash _default_hash_fn_type = 'mmh3' except ImportError: def _md5_hash(value): # md5 is a 128-bit hash, so we truncate the hexdigest (string of 32 # hexadecimal digits) to 16 digits and convert to int to get the 64-bit # integer fingerprint. return int(hashlib.md5(value).hexdigest()[:16], 16) _default_hash_fn = _md5_hash _default_hash_fn_type = 'md5' def _get_default_hash_fn(): """Returns either murmurhash or md5 based on installation.""" if _default_hash_fn_type == 'md5': logging.warning( 'Couldn\'t find murmurhash. Install mmh3 for a faster implementation of' 'ApproximateUnique.') return _default_hash_fn class ApproximateUnique(object): """ Hashes input elements and uses those to extrapolate the size of the entire set of hash values by assuming the rest of the hash values are as densely distributed as the sample space. """ _NO_VALUE_ERR_MSG = 'Either size or error should be set. Received {}.' _MULTI_VALUE_ERR_MSG = 'Either size or error should be set. ' \ 'Received {size = %s, error = %s}.' _INPUT_SIZE_ERR_MSG = 'ApproximateUnique needs a size >= 16 for an error ' \ '<= 0.50. In general, the estimation error is about ' \ '2 / sqrt(sample_size). Received {size = %s}.' _INPUT_ERROR_ERR_MSG = 'ApproximateUnique needs an estimation error ' \ 'between 0.01 and 0.50. Received {error = %s}.' @staticmethod def parse_input_params(size=None, error=None): """ Check if input params are valid and return sample size. :param size: an int not smaller than 16, which we would use to estimate number of unique values. :param error: max estimation error, which is a float between 0.01 and 0.50. If error is given, sample size will be calculated from error with _get_sample_size_from_est_error function. :return: sample size :raises: ValueError: If both size and error are given, or neither is given, or values are out of range. """ if None not in (size, error): raise ValueError(ApproximateUnique._MULTI_VALUE_ERR_MSG % (size, error)) elif size is None and error is None: raise ValueError(ApproximateUnique._NO_VALUE_ERR_MSG) elif size is not None: if not isinstance(size, int) or size < 16: raise ValueError(ApproximateUnique._INPUT_SIZE_ERR_MSG % (size)) else: return size else: if error < 0.01 or error > 0.5: raise ValueError(ApproximateUnique._INPUT_ERROR_ERR_MSG % (error)) else: return ApproximateUnique._get_sample_size_from_est_error(error) @staticmethod def _get_sample_size_from_est_error(est_err): """ :return: sample size Calculate sample size from estimation error """ return math.ceil(4.0 / math.pow(est_err, 2.0)) @typehints.with_input_types(T) @typehints.with_output_types(int) class Globally(PTransform): """ Approximate.Globally approximate number of unique values""" def __init__(self, size=None, error=None): self._sample_size = ApproximateUnique.parse_input_params(size, error) def expand(self, pcoll): coder = coders.registry.get_coder(pcoll) return pcoll \ | 'CountGlobalUniqueValues' \ >> (CombineGlobally(ApproximateUniqueCombineFn(self._sample_size, coder))) @typehints.with_input_types(typing.Tuple[K, V]) @typehints.with_output_types(typing.Tuple[K, int]) class PerKey(PTransform): """ Approximate.PerKey approximate number of unique values per key""" def __init__(self, size=None, error=None): self._sample_size = ApproximateUnique.parse_input_params(size, error) def expand(self, pcoll): coder = coders.registry.get_coder(pcoll) return pcoll \ | 'CountPerKeyUniqueValues' \ >> (CombinePerKey(ApproximateUniqueCombineFn(self._sample_size, coder))) class _LargestUnique(object): """ An object to keep samples and calculate sample hash space. It is an accumulator of a combine function. """ # We use unsigned 64-bit integer hashes. _HASH_SPACE_SIZE = 2.0**64 def __init__(self, sample_size): self._sample_size = sample_size self._min_hash = 2.0**64 self._sample_heap = [] self._sample_set = set() def add(self, element): """ :param an element from pcoll. :return: boolean type whether the value is in the heap Adds a value to the heap, returning whether the value is (large enough to be) in the heap. """ if len(self._sample_heap) >= self._sample_size and element < self._min_hash: return False if element not in self._sample_set: self._sample_set.add(element) heapq.heappush(self._sample_heap, element) if len(self._sample_heap) > self._sample_size: temp = heapq.heappop(self._sample_heap) self._sample_set.remove(temp) self._min_hash = self._sample_heap[0] elif element < self._min_hash: self._min_hash = element return True def get_estimate(self): """ :return: estimation count of unique values If heap size is smaller than sample size, just return heap size. Otherwise, takes into account the possibility of hash collisions, which become more likely than not for 2^32 distinct elements. Note that log(1+x) ~ x for small x, so for sampleSize << maxHash log(1 - sample_size/sample_space) / log(1 - 1/sample_space) ~ sample_size and hence estimate ~ sample_size * hash_space / sample_space as one would expect. Given sample_size / sample_space = est / hash_space est = sample_size * hash_space / sample_space Given above sample_size approximate, est = log1p(-sample_size/sample_space) / log1p(-1/sample_space) * hash_space / sample_space """ if len(self._sample_heap) < self._sample_size: return len(self._sample_heap) else: sample_space_size = self._HASH_SPACE_SIZE - 1.0 * self._min_hash est = ( math.log1p(-self._sample_size / sample_space_size) / math.log1p(-1 / sample_space_size) * self._HASH_SPACE_SIZE / sample_space_size) return round(est) class ApproximateUniqueCombineFn(CombineFn): """ ApproximateUniqueCombineFn computes an estimate of the number of unique values that were combined. """ def __init__(self, sample_size, coder): self._sample_size = sample_size coder = coders.typecoders.registry.verify_deterministic( coder, 'ApproximateUniqueCombineFn') self._coder = coder self._hash_fn = _get_default_hash_fn() def create_accumulator(self, *args, **kwargs): return _LargestUnique(self._sample_size) def add_input(self, accumulator, element, *args, **kwargs): try: hashed_value = self._hash_fn(self._coder.encode(element)) accumulator.add(hashed_value) return accumulator except Exception as e: raise RuntimeError("Runtime exception: %s" % e) # created an issue https://github.com/apache/beam/issues/19459 to speed up # merge process. def merge_accumulators(self, accumulators, *args, **kwargs): merged_accumulator = self.create_accumulator() for accumulator in accumulators: for i in accumulator._sample_heap: merged_accumulator.add(i) return merged_accumulator @staticmethod def extract_output(accumulator): return accumulator.get_estimate() def display_data(self): return {'sample_size': self._sample_size} class ApproximateQuantiles(object): """ PTransform for getting the idea of data distribution using approximate N-tile (e.g. quartiles, percentiles etc.) either globally or per-key. Examples: in: list(range(101)), num_quantiles=5 out: [0, 25, 50, 75, 100] in: [(i, 1 if i<10 else 1e-5) for i in range(101)], num_quantiles=5, weighted=True out: [0, 2, 5, 7, 100] in: [list(range(10)), ..., list(range(90, 101))], num_quantiles=5, input_batched=True out: [0, 25, 50, 75, 100] in: [(list(range(10)), [1]*10), (list(range(10)), [0]*10), ..., (list(range(90, 101)), [0]*11)], num_quantiles=5, input_batched=True, weighted=True out: [0, 2, 5, 7, 100] """ @staticmethod def _display_data(num_quantiles, key, reverse, weighted, input_batched): return { 'num_quantiles': DisplayDataItem(num_quantiles, label='Quantile Count'), 'key': DisplayDataItem( key.__name__ if hasattr(key, '__name__') else key.__class__.__name__, label='Record Comparer Key'), 'reverse': DisplayDataItem(str(reverse), label='Is Reversed'), 'weighted': DisplayDataItem(str(weighted), label='Is Weighted'), 'input_batched': DisplayDataItem( str(input_batched), label='Is Input Batched'), } @typehints.with_input_types( typing.Union[T, typing.Sequence[T], Tuple[T, float]]) @typehints.with_output_types(List[T]) class Globally(PTransform): """ PTransform takes PCollection and returns a list whose single value is approximate N-tiles of the input collection globally. Args: num_quantiles: number of elements in the resulting quantiles values list. key: (optional) Key is a mapping of elements to a comparable key, similar to the key argument of Python's sorting methods. reverse: (optional) whether to order things smallest to largest, rather than largest to smallest. weighted: (optional) if set to True, the transform returns weighted quantiles. The input PCollection is then expected to contain tuples of input values with the corresponding weight. input_batched: (optional) if set to True, the transform expects each element of input PCollection to be a batch, which is a list of elements for non-weighted case and a tuple of lists of elements and weights for weighted. Provides a way to accumulate multiple elements at a time more efficiently. """ def __init__( self, num_quantiles, key=None, reverse=False, weighted=False, input_batched=False): self._num_quantiles = num_quantiles self._key = key self._reverse = reverse self._weighted = weighted self._input_batched = input_batched def expand(self, pcoll): return pcoll | CombineGlobally( ApproximateQuantilesCombineFn.create( num_quantiles=self._num_quantiles, key=self._key, reverse=self._reverse, weighted=self._weighted, input_batched=self._input_batched)) def display_data(self): return ApproximateQuantiles._display_data( num_quantiles=self._num_quantiles, key=self._key, reverse=self._reverse, weighted=self._weighted, input_batched=self._input_batched) @typehints.with_input_types( typehints.Union[Tuple[K, V], Tuple[K, Tuple[V, float]]]) @typehints.with_output_types(Tuple[K, List[V]]) class PerKey(PTransform): """ PTransform takes PCollection of KV and returns a list based on each key whose single value is list of approximate N-tiles of the input element of the key. Args: num_quantiles: number of elements in the resulting quantiles values list. key: (optional) Key is a mapping of elements to a comparable key, similar to the key argument of Python's sorting methods. reverse: (optional) whether to order things smallest to largest, rather than largest to smallest. weighted: (optional) if set to True, the transform returns weighted quantiles. The input PCollection is then expected to contain tuples of input values with the corresponding weight. input_batched: (optional) if set to True, the transform expects each element of input PCollection to be a batch, which is a list of elements for non-weighted case and a tuple of lists of elements and weights for weighted. Provides a way to accumulate multiple elements at a time more efficiently. """ def __init__( self, num_quantiles, key=None, reverse=False, weighted=False, input_batched=False): self._num_quantiles = num_quantiles self._key = key self._reverse = reverse self._weighted = weighted self._input_batched = input_batched def expand(self, pcoll): return pcoll | CombinePerKey( ApproximateQuantilesCombineFn.create( num_quantiles=self._num_quantiles, key=self._key, reverse=self._reverse, weighted=self._weighted, input_batched=self._input_batched)) def display_data(self): return ApproximateQuantiles._display_data( num_quantiles=self._num_quantiles, key=self._key, reverse=self._reverse, weighted=self._weighted, input_batched=self._input_batched) class _QuantileSpec(object): """Quantiles computation specifications.""" def __init__(self, buffer_size, num_buffers, weighted, key, reverse): # type: (int, int, bool, Any, bool) -> None self.buffer_size = buffer_size self.num_buffers = num_buffers self.weighted = weighted self.key = key self.reverse = reverse # Used to sort tuples of values and weights. self.weighted_key = None if key is None else (lambda x: key(x[0])) # Used to compare values. if reverse and key is None: self.less_than = lambda a, b: a > b elif reverse: self.less_than = lambda a, b: key(a) > key(b) elif key is None: self.less_than = lambda a, b: a < b else: self.less_than = lambda a, b: key(a) < key(b) def get_argsort_key(self, elements): # type: (List) -> Callable[[int], Any] """Returns a key for sorting indices of elements by element's value.""" if self.key is None: return elements.__getitem__ else: return lambda idx: self.key(elements[idx]) def __reduce__(self): return ( self.__class__, ( self.buffer_size, self.num_buffers, self.weighted, self.key, self.reverse)) class _QuantileBuffer(object): """A single buffer in the sense of the referenced algorithm. (see http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.6.6513&rep=rep1 &type=pdf and ApproximateQuantilesCombineFn for further information)""" def __init__( self, elements, weights, weighted, level=0, min_val=None, max_val=None): # type: (list, list, bool, int, Any, Any) -> None self.elements = elements # In non-weighted case weights contains a single element representing weight # of the buffer in the sense of the original algorithm. In weighted case, # it stores weights of individual elements. self.weights = weights self.weighted = weighted self.level = level if min_val is None or max_val is None: # Buffer is always initialized with sorted elements. self.min_val = elements[0] self.max_val = elements[-1] else: # Note that collapsed buffer may not contain min and max in the list of # elements. self.min_val = min_val self.max_val = max_val def __iter__(self): return zip( self.elements, self.weights if self.weighted else itertools.repeat(self.weights[0])) def __lt__(self, other): return self.level < other.level class _QuantileState(object): """ Compact summarization of a collection on which quantiles can be estimated. """ def __init__(self, unbuffered_elements, unbuffered_weights, buffers, spec): # type: (List, List, List[_QuantileBuffer], _QuantileSpec) -> None self.buffers = buffers self.spec = spec if spec.weighted: self.add_unbuffered = self._add_unbuffered_weighted else: self.add_unbuffered = self._add_unbuffered # The algorithm requires that the manipulated buffers always be filled to # capacity to perform the collapse operation. This operation can be extended # to buffers of varying sizes by introducing the notion of fractional # weights, but it's easier to simply combine the remainders from all shards # into new, full buffers and then take them into account when computing the # final output. self.unbuffered_elements = unbuffered_elements self.unbuffered_weights = unbuffered_weights # This is needed for pickling to work when Cythonization is enabled. def __reduce__(self): return ( self.__class__, ( self.unbuffered_elements, self.unbuffered_weights, self.buffers, self.spec)) def is_empty(self): # type: () -> bool """Check if the buffered & unbuffered elements are empty or not.""" return not self.unbuffered_elements and not self.buffers def _add_unbuffered(self, elements, offset_fn): # type: (List, Any) -> None """ Add elements to the unbuffered list, creating new buffers and collapsing if needed. """ self.unbuffered_elements.extend(elements) num_new_buffers = len(self.unbuffered_elements) // self.spec.buffer_size for idx in range(num_new_buffers): to_buffer = sorted( self.unbuffered_elements[idx * self.spec.buffer_size:(idx + 1) * self.spec.buffer_size], key=self.spec.key, reverse=self.spec.reverse) heapq.heappush( self.buffers, _QuantileBuffer(elements=to_buffer, weights=[1], weighted=False)) if num_new_buffers > 0: self.unbuffered_elements = self.unbuffered_elements[num_new_buffers * self.spec. buffer_size:] self.collapse_if_needed(offset_fn) def _add_unbuffered_weighted(self, elements, offset_fn): # type: (List, Any) -> None """ Add elements with weights to the unbuffered list, creating new buffers and collapsing if needed. """ if len(elements) == 1: self.unbuffered_elements.append(elements[0][0]) self.unbuffered_weights.append(elements[0][1]) else: self.unbuffered_elements.extend(elements[0]) self.unbuffered_weights.extend(elements[1]) num_new_buffers = len(self.unbuffered_elements) // self.spec.buffer_size argsort_key = self.spec.get_argsort_key(self.unbuffered_elements) for idx in range(num_new_buffers): argsort = sorted( range(idx * self.spec.buffer_size, (idx + 1) * self.spec.buffer_size), key=argsort_key, reverse=self.spec.reverse) elements_to_buffer = [self.unbuffered_elements[idx] for idx in argsort] weights_to_buffer = [self.unbuffered_weights[idx] for idx in argsort] heapq.heappush( self.buffers, _QuantileBuffer( elements=elements_to_buffer, weights=weights_to_buffer, weighted=True)) if num_new_buffers > 0: self.unbuffered_elements = self.unbuffered_elements[num_new_buffers * self.spec. buffer_size:] self.unbuffered_weights = self.unbuffered_weights[num_new_buffers * self.spec.buffer_size:] self.collapse_if_needed(offset_fn) def finalize(self): # type: () -> None """ Creates a new buffer using all unbuffered elements. Called before extracting an output. Note that the buffer doesn't have to be put in a proper position since _collapse is not going to be called after. """ if self.unbuffered_elements and self.spec.weighted: argsort_key = self.spec.get_argsort_key(self.unbuffered_elements) argsort = sorted( range(len(self.unbuffered_elements)), key=argsort_key, reverse=self.spec.reverse) self.unbuffered_elements = [ self.unbuffered_elements[idx] for idx in argsort ] self.unbuffered_weights = [ self.unbuffered_weights[idx] for idx in argsort ] self.buffers.append( _QuantileBuffer( self.unbuffered_elements, self.unbuffered_weights, weighted=True)) self.unbuffered_weights = [] elif self.unbuffered_elements: self.unbuffered_elements.sort( key=self.spec.key, reverse=self.spec.reverse) self.buffers.append( _QuantileBuffer( self.unbuffered_elements, weights=[1], weighted=False)) self.unbuffered_elements = [] def collapse_if_needed(self, offset_fn): # type: (Any) -> None """ Checks if summary has too many buffers and collapses some of them until the limit is restored. """ while len(self.buffers) > self.spec.num_buffers: to_collapse = [heapq.heappop(self.buffers), heapq.heappop(self.buffers)] min_level = to_collapse[1].level while self.buffers and self.buffers[0].level <= min_level: to_collapse.append(heapq.heappop(self.buffers)) heapq.heappush(self.buffers, _collapse(to_collapse, offset_fn, self.spec)) def _collapse(buffers, offset_fn, spec): # type: (List[_QuantileBuffer], Any, _QuantileSpec) -> _QuantileBuffer """ Approximates elements from multiple buffers and produces a single buffer. """ new_level = 0 new_weight = 0 for buffer in buffers: # As presented in the paper, there should always be at least two # buffers of the same (minimal) level to collapse, but it is possible # to violate this condition when combining buffers from independently # computed shards. If they differ we take the max. new_level = max([new_level, buffer.level + 1]) new_weight = new_weight + sum(buffer.weights) if spec.weighted: step = new_weight / (spec.buffer_size - 1) offset = new_weight / (2 * spec.buffer_size) else: step = new_weight offset = offset_fn(new_weight) new_elements, new_weights, min_val, max_val = \ _interpolate(buffers, spec.buffer_size, step, offset, spec) if not spec.weighted: new_weights = [new_weight] return _QuantileBuffer( new_elements, new_weights, spec.weighted, new_level, min_val, max_val) def _interpolate(buffers, count, step, offset, spec): # type: (List[_QuantileBuffer], int, float, float, _QuantileSpec) -> Tuple[List, List, Any, Any] """ Emulates taking the ordered union of all elements in buffers, repeated according to their weight, and picking out the (k * step + offset)-th elements of this list for `0 <= k < count`. """ buffer_iterators = [] min_val = buffers[0].min_val max_val = buffers[0].max_val for buffer in buffers: # Calculate extreme values for the union of buffers. min_val = buffer.min_val if spec.less_than( buffer.min_val, min_val) else min_val max_val = buffer.max_val if spec.less_than( max_val, buffer.max_val) else max_val buffer_iterators.append(iter(buffer)) # Note that `heapq.merge` can also be used here since the buffers are sorted. # In practice, however, `sorted` uses natural order in the union and # significantly outperforms `heapq.merge`. sorted_elements = sorted( itertools.chain.from_iterable(buffer_iterators), key=spec.weighted_key, reverse=spec.reverse) if not spec.weighted: # If all buffers have the same weight, then quantiles' indices are evenly # distributed over a range [0, len(sorted_elements)]. buffers_have_same_weight = True weight = buffers[0].weights[0] for buffer in buffers: if buffer.weights[0] != weight: buffers_have_same_weight = False break if buffers_have_same_weight: offset = offset / weight step = step / weight max_idx = len(sorted_elements) - 1 result = [ sorted_elements[min(int(j * step + offset), max_idx)][0] for j in range(count) ] return result, [], min_val, max_val sorted_elements_iter = iter(sorted_elements) weighted_element = next(sorted_elements_iter) new_elements = [] new_weights = [] j = 0 current_weight = weighted_element[1] previous_weight = 0 while j < count: target_weight = j * step + offset j += 1 try: while current_weight <= target_weight: weighted_element = next(sorted_elements_iter) current_weight += weighted_element[1] except StopIteration: pass new_elements.append(weighted_element[0]) if spec.weighted: new_weights.append(current_weight - previous_weight) previous_weight = current_weight return new_elements, new_weights, min_val, max_val class ApproximateQuantilesCombineFn(CombineFn): """ This combiner gives an idea of the distribution of a collection of values using approximate N-tiles. The output of this combiner is the list of size of the number of quantiles (num_quantiles), containing the input values of the minimum value item of the list, the intermediate values (n-tiles) and the maximum value item of the list, in the sort order provided via key (similar to the key argument of Python's sorting methods). If there are fewer values to combine than the number of quantile (num_quantiles), then the resulting list will contain all the values being combined, in sorted order. If no `key` is provided, then the results are sorted in the natural order. To evaluate the quantiles, we use the "New Algorithm" described here: [MRL98] Manku, Rajagopalan & Lindsay, "Approximate Medians and other Quantiles in One Pass and with Limited Memory", Proc. 1998 ACM SIGMOD, Vol 27, No 2, p 426-435, June 1998. http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.6.6513&rep=rep1 &type=pdf Note that the weighted quantiles are evaluated using a generalized version of the algorithm referenced in the paper. The default error bound is (1 / num_quantiles) for uniformly distributed data and min(1e-2, 1 / num_quantiles) for weighted case, though in practice the accuracy tends to be much better. Args: num_quantiles: Number of quantiles to produce. It is the size of the final output list, including the mininum and maximum value items. buffer_size: The size of the buffers, corresponding to k in the referenced paper. num_buffers: The number of buffers, corresponding to b in the referenced paper. key: (optional) Key is a mapping of elements to a comparable key, similar to the key argument of Python's sorting methods. reverse: (optional) whether to order things smallest to largest, rather than largest to smallest. weighted: (optional) if set to True, the combiner produces weighted quantiles. The input elements are then expected to be tuples of input values with the corresponding weight. input_batched: (optional) if set to True, inputs are expected to be batches of elements. """ # For alternating between biasing up and down in the above even weight # collapse operation. _offset_jitter = 0 # The cost (in time and space) to compute quantiles to a given accuracy is a # function of the total number of elements in the data set. If an estimate is # not known or specified, we use this as an upper bound. If this is too low, # errors may exceed the requested tolerance; if too high, efficiency may be # non-optimal. The impact is logarithmic with respect to this value, so this # default should be fine for most uses. _MAX_NUM_ELEMENTS = 1e9 _qs = None # type: _QuantileState def __init__( self, num_quantiles, # type: int buffer_size, # type: int num_buffers, # type: int key=None, reverse=False, weighted=False, input_batched=False): self._num_quantiles = num_quantiles self._spec = _QuantileSpec(buffer_size, num_buffers, weighted, key, reverse) self._input_batched = input_batched if self._input_batched: setattr(self, 'add_input', self._add_inputs) def __reduce__(self): return ( self.__class__, ( self._num_quantiles, self._spec.buffer_size, self._spec.num_buffers, self._spec.key, self._spec.reverse, self._spec.weighted, self._input_batched)) @classmethod def create( cls, num_quantiles, # type: int epsilon=None, max_num_elements=None, key=None, reverse=False, weighted=False, input_batched=False): # type: (...) -> ApproximateQuantilesCombineFn """ Creates an approximate quantiles combiner with the given key and desired number of quantiles. Args: num_quantiles: Number of quantiles to produce. It is the size of the final output list, including the mininum and maximum value items. epsilon: (optional) The default error bound is `epsilon`, which holds as long as the number of elements is less than `_MAX_NUM_ELEMENTS`. Specifically, if one considers the input as a sorted list x_1, ..., x_N, then the distance between each exact quantile x_c and its approximation x_c' is bounded by `|c - c'| < epsilon * N`. Note that these errors are worst-case scenarios. In practice the accuracy tends to be much better. max_num_elements: (optional) The cost (in time and space) to compute quantiles to a given accuracy is a function of the total number of elements in the data set. key: (optional) Key is a mapping of elements to a comparable key, similar to the key argument of Python's sorting methods. reverse: (optional) whether to order things smallest to largest, rather than largest to smallest. weighted: (optional) if set to True, the combiner produces weighted quantiles. The input elements are then expected to be tuples of values with the corresponding weight. input_batched: (optional) if set to True, inputs are expected to be batches of elements. """ max_num_elements = max_num_elements or cls._MAX_NUM_ELEMENTS if not epsilon: epsilon = min(1e-2, 1.0 / num_quantiles) \ if weighted else (1.0 / num_quantiles) # Note that calculation of the buffer size and the number of buffers here # is based on technique used in the Munro-Paterson algorithm. Switching to # the logic used in the "New Algorithm" may result in memory savings since # it results in lower values for b and k in practice. b = 2 while (b - 2) * (1 << (b - 2)) < epsilon * max_num_elements: b = b + 1 b = b - 1 k = max(2, int(math.ceil(max_num_elements / float(1 << (b - 1))))) return cls( num_quantiles=num_quantiles, buffer_size=k, num_buffers=b, key=key, reverse=reverse, weighted=weighted, input_batched=input_batched) def _offset(self, new_weight): # type: (int) -> float """ If the weight is even, we must round up or down. Alternate between these two options to avoid a bias. """ if new_weight % 2 == 1: return (new_weight + 1) / 2 else: self._offset_jitter = 2 - self._offset_jitter return (new_weight + self._offset_jitter) / 2 # TODO(https://github.com/apache/beam/issues/19737): Signature incompatible # with supertype def create_accumulator(self): # type: () -> _QuantileState self._qs = _QuantileState( unbuffered_elements=[], unbuffered_weights=[], buffers=[], spec=self._spec) return self._qs def add_input(self, quantile_state, element): """ Add a new element to the collection being summarized by quantile state. """ quantile_state.add_unbuffered([element], self._offset) return quantile_state def _add_inputs(self, quantile_state, elements): # type: (_QuantileState, List) -> _QuantileState """ Add a batch of elements to the collection being summarized by quantile state. """ if len(elements) == 0: return quantile_state quantile_state.add_unbuffered(elements, self._offset) return quantile_state def merge_accumulators(self, accumulators): """Merges all the accumulators (quantile state) as one.""" qs = self.create_accumulator() for accumulator in accumulators: if accumulator.is_empty(): continue if self._spec.weighted: qs.add_unbuffered( [accumulator.unbuffered_elements, accumulator.unbuffered_weights], self._offset) else: qs.add_unbuffered(accumulator.unbuffered_elements, self._offset) qs.buffers.extend(accumulator.buffers) heapq.heapify(qs.buffers) qs.collapse_if_needed(self._offset) return qs def extract_output(self, accumulator): """ Outputs num_quantiles elements consisting of the minimum, maximum and num_quantiles - 2 evenly spaced intermediate elements. Returns the empty list if no elements have been added. """ if accumulator.is_empty(): return [] accumulator.finalize() all_elems = accumulator.buffers total_weight = 0 if self._spec.weighted: for buffer_elem in all_elems: total_weight += sum(buffer_elem.weights) else: for buffer_elem in all_elems: total_weight += len(buffer_elem.elements) * buffer_elem.weights[0] step = total_weight / (self._num_quantiles - 1) offset = (total_weight - 1) / (self._num_quantiles - 1) quantiles, _, min_val, max_val = \ _interpolate(all_elems, self._num_quantiles - 2, step, offset, self._spec) return [min_val] + quantiles + [max_val] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/stats_test.py0000644000175100017510000006076215206113033022636 0ustar00runnerrunner# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import math import random import sys import unittest from collections import defaultdict import hamcrest as hc from parameterized import parameterized from parameterized import parameterized_class import apache_beam as beam from apache_beam.coders import coders from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import BeamAssertException from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.core import Create from apache_beam.transforms.display import DisplayData from apache_beam.transforms.display_test import DisplayDataItemMatcher from apache_beam.transforms.stats import ApproximateQuantilesCombineFn from apache_beam.transforms.stats import ApproximateUniqueCombineFn try: import mmh3 mmh3_options = [(mmh3, ), (None, )] except ImportError: mmh3_options = [(None, )] @parameterized_class(('mmh3_option', ), mmh3_options) class ApproximateUniqueTest(unittest.TestCase): """Unit tests for ApproximateUnique.Globally and ApproximateUnique.PerKey.""" random.seed(0) def setUp(self): sys.modules['mmh3'] = self.mmh3_option @parameterized.expand([ ( 'small_population_by_size', list(range(30)), 32, None, 'assert:global_by_sample_size_with_small_population'), ( 'large_population_by_size', list(range(100)), 16, None, 'assert:global_by_sample_size_with_large_population'), ( 'with_duplicates_by_size', [10] * 50 + [20] * 50, 30, None, 'assert:global_by_sample_size_with_duplicates'), ( 'small_population_by_error', list(range(30)), None, 0.3, 'assert:global_by_error_with_small_population'), ( 'large_population_by_error', [random.randint(1, 1000) for _ in range(500)], None, 0.1, 'assert:global_by_error_with_large_population'), ]) def test_approximate_unique_global( self, name, test_input, sample_size, est_error, label): # check that only either sample_size or est_error is not None assert bool(sample_size) != bool(est_error) if sample_size: error = 2 / math.sqrt(sample_size) else: error = est_error random.shuffle(test_input) actual_count = len(set(test_input)) with TestPipeline() as pipeline: result = ( pipeline | 'create' >> beam.Create(test_input) | 'get_estimate' >> beam.ApproximateUnique.Globally( size=sample_size, error=est_error) | 'compare' >> beam.FlatMap( lambda x: [abs(x - actual_count) * 1.0 / actual_count <= error])) assert_that(result, equal_to([True]), label=label) @parameterized.expand([ ('by_size', 20, None, 'assert:unique_perkey_by_sample_size'), ('by_error', None, 0.02, 'assert:unique_perkey_by_error') ]) def test_approximate_unique_perkey(self, name, sample_size, est_error, label): # check that only either sample_size or est_error is set assert bool(sample_size) != bool(est_error) if sample_size: error = 2 / math.sqrt(sample_size) else: error = est_error test_input = [(8, 73), (6, 724), (7, 70), (1, 576), (10, 120), (2, 662), (7, 115), (3, 731), (6, 340), (6, 623), (1, 74), (9, 280), (8, 298), (6, 440), (10, 243), (1, 125), (9, 754), (8, 833), (9, 751), (4, 818), (6, 176), (9, 253), (2, 721), (8, 936), (3, 691), (10, 685), (1, 69), (3, 155), (8, 86), (5, 693), (2, 809), (4, 723), (8, 102), (9, 707), (8, 558), (4, 537), (5, 371), (7, 432), (2, 51), (10, 397)] actual_count_dict = defaultdict(set) for (x, y) in test_input: actual_count_dict[x].add(y) with TestPipeline() as pipeline: result = ( pipeline | 'create' >> beam.Create(test_input) | 'get_estimate' >> beam.ApproximateUnique.PerKey( size=sample_size, error=est_error) | 'compare' >> beam.FlatMap( lambda x: [ abs(x[1] - len(actual_count_dict[x[0]])) * 1.0 / len( actual_count_dict[x[0]]) <= error ])) assert_that( result, equal_to([True] * len(actual_count_dict)), label=label) @parameterized.expand([ ( 'invalid_input_size', list(range(30)), 10, None, beam.ApproximateUnique._INPUT_SIZE_ERR_MSG % 10), ( 'invalid_type_size', list(range(30)), 100.0, None, beam.ApproximateUnique._INPUT_SIZE_ERR_MSG % 100.0), ( 'invalid_small_error', list(range(30)), None, 0.0, beam.ApproximateUnique._INPUT_ERROR_ERR_MSG % 0.0), ( 'invalid_big_error', list(range(30)), None, 0.6, beam.ApproximateUnique._INPUT_ERROR_ERR_MSG % 0.6), ( 'no_input', list(range(30)), None, None, beam.ApproximateUnique._NO_VALUE_ERR_MSG), ( 'both_input', list(range(30)), 30, 0.2, beam.ApproximateUnique._MULTI_VALUE_ERR_MSG % (30, 0.2)), ]) def test_approximate_unique_global_value_error( self, name, test_input, sample_size, est_error, expected_msg): with self.assertRaises(ValueError) as e: with TestPipeline() as pipeline: _ = ( pipeline | 'create' >> beam.Create(test_input) | 'get_estimate' >> beam.ApproximateUnique.Globally( size=sample_size, error=est_error)) assert e.exception.args[0] == expected_msg def test_approximate_unique_combine_fn_requires_nondeterministic_coder(self): sample_size = 30 coder = coders.Base64PickleCoder() with self.assertRaises(ValueError) as e: _ = ApproximateUniqueCombineFn(sample_size, coder) self.assertRegex( e.exception.args[0], 'The key coder "Base64PickleCoder" ' 'for ApproximateUniqueCombineFn is not deterministic.') def test_approximate_unique_combine_fn_requires_compatible_coder(self): test_input = 'a' sample_size = 30 coder = coders.FloatCoder() combine_fn = ApproximateUniqueCombineFn(sample_size, coder) accumulator = combine_fn.create_accumulator() with self.assertRaises(RuntimeError) as e: accumulator = combine_fn.add_input(accumulator, test_input) self.assertRegex(e.exception.args[0], 'Runtime exception') def test_get_sample_size_from_est_error(self): # test if get correct sample size from input error. assert beam.ApproximateUnique._get_sample_size_from_est_error(0.5) == 16 assert beam.ApproximateUnique._get_sample_size_from_est_error(0.4) == 25 assert beam.ApproximateUnique._get_sample_size_from_est_error(0.2) == 100 assert beam.ApproximateUnique._get_sample_size_from_est_error(0.1) == 400 assert beam.ApproximateUnique._get_sample_size_from_est_error(0.05) == 1600 assert beam.ApproximateUnique._get_sample_size_from_est_error(0.01) == 40000 class ApproximateQuantilesTest(unittest.TestCase): _kv_data = [("a", 1), ("a", 2), ("a", 3), ("b", 1), ("b", 10), ("b", 10), ("b", 100)] _kv_str_data = [("a", "a"), ("a", "a" * 2), ("a", "a" * 3), ("b", "b"), ("b", "b" * 10), ("b", "b" * 10), ("b", "b" * 100)] @staticmethod def _quantiles_matcher(expected): l = len(expected) def assert_true(exp): if not exp: raise BeamAssertException('%s Failed assert True' % repr(exp)) def match(actual): actual = actual[0] for i in range(l): if isinstance(expected[i], list): assert_true(expected[i][0] <= actual[i] <= expected[i][1]) else: equal_to([expected[i]])([actual[i]]) return match @staticmethod def _approx_quantile_generator(size, num_of_quantiles, absoluteError): quantiles = [0] k = 1 while k < num_of_quantiles - 1: expected = (size - 1) * k / (num_of_quantiles - 1) quantiles.append([expected - absoluteError, expected + absoluteError]) k = k + 1 quantiles.append(size - 1) return quantiles def test_quantiles_globaly(self): with TestPipeline() as p: pc = p | Create(list(range(101))) quantiles = pc | 'Quantiles globally' >> \ beam.ApproximateQuantiles.Globally(5) quantiles_reversed = pc | 'Quantiles globally reversed' >> \ beam.ApproximateQuantiles.Globally(5, reverse=True) assert_that( quantiles, equal_to([[0, 25, 50, 75, 100]]), label='checkQuantilesGlobally') assert_that( quantiles_reversed, equal_to([[100, 75, 50, 25, 0]]), label='checkReversedQuantiles') def test_quantiles_globally_weighted(self): num_inputs = 1e3 a = -3 b = 3 # Weighting function coincides with the pdf of the standard normal # distribution up to a constant. Since 99.7% of the probability mass for # this pdf is concentrated in the interval [a, b] = [-3, 3], the quantiles # for a sample from this interval with the given weight function are # expected to be close to the quantiles of the standard normal distribution. def weight(x): return math.exp(-(x**2) / 2) input_data = [ (a + (b - a) * i / num_inputs, weight(a + (b - a) * i / num_inputs)) for i in range(int(num_inputs) + 1) ] with TestPipeline() as p: pc = p | Create(input_data) weighted_quantiles = pc | "Quantiles globally weighted" >> \ beam.ApproximateQuantiles.Globally(5, weighted=True) reversed_weighted_quantiles = ( pc | 'Quantiles globally weighted reversed' >> beam.ApproximateQuantiles.Globally(5, reverse=True, weighted=True)) assert_that( weighted_quantiles, equal_to([[-3., -0.6720000000000002, 0., 0.6720000000000002, 3.]]), label="checkWeightedQuantilesGlobally") assert_that( reversed_weighted_quantiles, equal_to([[3., 0.6720000000000002, 0., -0.6720000000000002, -3.]]), label="checkWeightedReversedQuantilesGlobally") def test_quantiles_per_key(self): with TestPipeline() as p: data = self._kv_data pc = p | Create(data) per_key = pc | 'Quantiles PerKey' >> beam.ApproximateQuantiles.PerKey(2) per_key_reversed = ( pc | 'Quantiles PerKey Reversed' >> beam.ApproximateQuantiles.PerKey( 2, reverse=True)) assert_that( per_key, equal_to([('a', [1, 3]), ('b', [1, 100])]), label='checkQuantilePerKey') assert_that( per_key_reversed, equal_to([('a', [3, 1]), ('b', [100, 1])]), label='checkReversedQuantilesPerKey') def test_quantiles_per_key_weighted(self): with TestPipeline() as p: data = [(k, (v, 2.)) for k, v in self._kv_data] pc = p | Create(data) per_key = pc | 'Weighted Quantiles PerKey' >> \ beam.ApproximateQuantiles.PerKey(2, weighted=True) per_key_reversed = pc | 'Weighted Quantiles PerKey Reversed' >> \ beam.ApproximateQuantiles.PerKey( 2, reverse=True, weighted=True) assert_that( per_key, equal_to([('a', [1, 3]), ('b', [1, 100])]), label='checkWeightedQuantilesPerKey') assert_that( per_key_reversed, equal_to([('a', [3, 1]), ('b', [100, 1])]), label='checkWeightedReversedQuantilesPerKey') def test_quantiles_per_key_with_key_argument(self): with TestPipeline() as p: data = self._kv_str_data pc = p | Create(data) per_key = pc | 'Per Key' >> beam.ApproximateQuantiles.PerKey(2, key=len) per_key_reversed = ( pc | 'Per Key Reversed' >> beam.ApproximateQuantiles.PerKey( 2, key=len, reverse=True)) assert_that( per_key, equal_to([('a', ['a', 'a' * 3]), ('b', ['b', 'b' * 100])]), label='checkPerKey') assert_that( per_key_reversed, equal_to([('a', ['a' * 3, 'a']), ('b', ['b' * 100, 'b'])]), label='checkPerKeyReversed') def test_singleton(self): with TestPipeline() as p: data = [389] pc = p | Create(data) quantiles = pc | beam.ApproximateQuantiles.Globally(5) assert_that(quantiles, equal_to([[389, 389, 389, 389, 389]])) def test_uneven_quantiles(self): with TestPipeline() as p: pc = p | Create(list(range(5000))) quantiles = pc | beam.ApproximateQuantiles.Globally(37) approx_quantiles = self._approx_quantile_generator( size=5000, num_of_quantiles=37, absoluteError=20) assert_that(quantiles, self._quantiles_matcher(approx_quantiles)) def test_large_quantiles(self): with TestPipeline() as p: pc = p | Create(list(range(10001))) quantiles = pc | beam.ApproximateQuantiles.Globally(50) approx_quantiles = self._approx_quantile_generator( size=10001, num_of_quantiles=50, absoluteError=20) assert_that(quantiles, self._quantiles_matcher(approx_quantiles)) def test_random_quantiles(self): with TestPipeline() as p: data = list(range(101)) random.shuffle(data) pc = p | Create(data) quantiles = pc | beam.ApproximateQuantiles.Globally(5) assert_that(quantiles, equal_to([[0, 25, 50, 75, 100]])) def test_duplicates(self): y = list(range(101)) data = [] for _ in range(10): data.extend(y) with TestPipeline() as p: pc = p | Create(data) quantiles = ( pc | 'Quantiles Globally' >> beam.ApproximateQuantiles.Globally(5)) quantiles_reversed = ( pc | 'Quantiles Reversed' >> beam.ApproximateQuantiles.Globally( 5, reverse=True)) assert_that( quantiles, equal_to([[0, 25, 50, 75, 100]]), label="checkQuantilesGlobally") assert_that( quantiles_reversed, equal_to([[100, 75, 50, 25, 0]]), label="checkQuantileReversed") def test_lots_of_duplicates(self): with TestPipeline() as p: data = [1] data.extend([2 for _ in range(299)]) data.extend([3 for _ in range(799)]) pc = p | Create(data) quantiles = pc | beam.ApproximateQuantiles.Globally(5) assert_that(quantiles, equal_to([[1, 2, 3, 3, 3]])) def test_log_distribution(self): with TestPipeline() as p: data = [int(math.log(x)) for x in range(1, 1000)] pc = p | Create(data) quantiles = pc | beam.ApproximateQuantiles.Globally(5) assert_that(quantiles, equal_to([[0, 5, 6, 6, 6]])) def test_zipfian_distribution(self): with TestPipeline() as p: data = [] for i in range(1, 1000): data.append(int(1000 / i)) pc = p | Create(data) quantiles = pc | beam.ApproximateQuantiles.Globally(5) assert_that(quantiles, equal_to([[1, 1, 2, 4, 1000]])) def test_alternate_quantiles(self): data = ["aa", "aaa", "aaaa", "b", "ccccc", "dddd", "zz"] with TestPipeline() as p: pc = p | Create(data) globally = pc | 'Globally' >> beam.ApproximateQuantiles.Globally(3) with_key = ( pc | 'Globally with key' >> beam.ApproximateQuantiles.Globally(3, key=len)) key_with_reversed = ( pc | 'Globally with key and reversed' >> beam.ApproximateQuantiles.Globally(3, key=len, reverse=True)) assert_that( globally, equal_to([["aa", "b", "zz"]]), label='checkGlobally') assert_that( with_key, equal_to([["b", "aaa", "ccccc"]]), label='checkGloballyWithKey') assert_that( key_with_reversed, equal_to([["ccccc", "aaa", "b"]]), label='checkWithKeyAndReversed') def test_batched_quantiles(self): with TestPipeline() as p: data = [] for i in range(100): data.append([(j / 10, abs(j - 500)) for j in range(i * 10, (i + 1) * 10)]) pc = p | Create(data) globally = ( pc | 'Globally' >> beam.ApproximateQuantiles.Globally( 3, input_batched=True)) with_key = ( pc | 'Globally with key' >> beam.ApproximateQuantiles.Globally( 3, key=sum, input_batched=True)) key_with_reversed = ( pc | 'Globally with key and reversed' >> beam.ApproximateQuantiles.Globally( 3, key=sum, reverse=True, input_batched=True)) assert_that( globally, equal_to([[(0.0, 500), (49.9, 1), (99.9, 499)]]), label='checkGlobally') assert_that( with_key, equal_to([[(50.0, 0), (72.5, 225), (99.9, 499)]]), label='checkGloballyWithKey') assert_that( key_with_reversed, equal_to([[(99.9, 499), (72.5, 225), (50.0, 0)]]), label='checkGloballyWithKeyAndReversed') def test_batched_weighted_quantiles(self): with TestPipeline() as p: data = [] for i in range(100): data.append([[(i / 10, abs(i - 500)) for i in range(i * 10, (i + 1) * 10)], [i] * 10]) pc = p | Create(data) globally = ( pc | 'Globally' >> beam.ApproximateQuantiles.Globally( 3, weighted=True, input_batched=True)) with_key = ( pc | 'Globally with key' >> beam.ApproximateQuantiles.Globally( 3, key=sum, weighted=True, input_batched=True)) key_with_reversed = ( pc | 'Globally with key and reversed' >> beam.ApproximateQuantiles.Globally( 3, key=sum, reverse=True, weighted=True, input_batched=True)) assert_that( globally, equal_to([[(0.0, 500), (70.8, 208), (99.9, 499)]]), label='checkGlobally') assert_that( with_key, equal_to([[(50.0, 0), (21.0, 290), (99.9, 499)]]), label='checkGloballyWithKey') assert_that( key_with_reversed, equal_to([[(99.9, 499), (21.0, 290), (50.0, 0)]]), label='checkGloballyWithKeyAndReversed') def test_quantiles_merge_accumulators(self): # This test exercises merging multiple buffers and approximation accuracy. # The max_num_elements is set to a small value to trigger buffers collapse # and interpolation. Under the conditions below, buffer_size=125 and # num_buffers=4, so we're only allowed to keep half of the input values. num_accumulators = 100 num_quantiles = 5 eps = 0.01 max_num_elements = 1000 combine_fn = ApproximateQuantilesCombineFn.create( num_quantiles, eps, max_num_elements) combine_fn_weighted = ApproximateQuantilesCombineFn.create( num_quantiles, eps, max_num_elements, weighted=True) data = list(range(1000)) weights = list(reversed(range(1000))) step = math.ceil(len(data) / num_accumulators) accumulators = [] accumulators_weighted = [] for i in range(num_accumulators): accumulator = combine_fn.create_accumulator() accumulator_weighted = combine_fn_weighted.create_accumulator() for element, weight in zip(data[i*step:(i+1)*step], weights[i*step:(i+1)*step]): accumulator = combine_fn.add_input(accumulator, element) accumulator_weighted = combine_fn_weighted.add_input( accumulator_weighted, (element, weight)) accumulators.append(accumulator) accumulators_weighted.append(accumulator_weighted) accumulator = combine_fn.merge_accumulators(accumulators) accumulator_weighted = combine_fn_weighted.merge_accumulators( accumulators_weighted) quantiles = combine_fn.extract_output(accumulator) quantiles_weighted = combine_fn_weighted.extract_output( accumulator_weighted) # In fact, the final accuracy is much higher than eps, but we test for a # minimal accuracy here. for q, actual_q in zip(quantiles, [0, 249, 499, 749, 999]): self.assertAlmostEqual(q, actual_q, delta=max_num_elements * eps) for q, actual_q in zip(quantiles_weighted, [0, 133, 292, 499, 999]): self.assertAlmostEqual(q, actual_q, delta=max_num_elements * eps) @staticmethod def _display_data_matcher(instance): expected_items = [ DisplayDataItemMatcher('num_quantiles', instance._num_quantiles), DisplayDataItemMatcher('weighted', str(instance._weighted)), DisplayDataItemMatcher('key', str(instance._key.__name__)), DisplayDataItemMatcher('reverse', str(instance._reverse)), DisplayDataItemMatcher('input_batched', str(instance._input_batched)), ] return expected_items def test_global_display_data(self): transform = beam.ApproximateQuantiles.Globally( 3, weighted=True, key=len, reverse=True) data = DisplayData.create_from(transform) expected_items = self._display_data_matcher(transform) hc.assert_that(data.items, hc.contains_inanyorder(*expected_items)) def test_perkey_display_data(self): transform = beam.ApproximateQuantiles.PerKey( 3, weighted=True, key=len, reverse=True) data = DisplayData.create_from(transform) expected_items = self._display_data_matcher(transform) hc.assert_that(data.items, hc.contains_inanyorder(*expected_items)) def _build_quantilebuffer_test_data(): """ Test data taken from "Munro-Paterson Algorithm" reference values table of "Approximate Medians and other Quantiles in One Pass and with Limited Memory" paper. See ApproximateQuantilesCombineFn for paper reference. """ epsilons = [0.1, 0.05, 0.01, 0.005, 0.001] maxElementExponents = [5, 6, 7, 8, 9] expectedNumBuffersValues = [[11, 14, 17, 21, 24], [11, 14, 17, 20, 23], [9, 11, 14, 17, 21], [8, 11, 14, 17, 20], [6, 9, 11, 14, 17]] expectedBufferSizeValues = [[98, 123, 153, 96, 120], [98, 123, 153, 191, 239], [391, 977, 1221, 1526, 954], [782, 977, 1221, 1526, 1908], [3125, 3907, 9766, 12208, 15259]] test_data = [] i = 0 for epsilon in epsilons: j = 0 for maxElementExponent in maxElementExponents: test_data.append([ epsilon, (10**maxElementExponent), expectedNumBuffersValues[i][j], expectedBufferSizeValues[i][j] ]) j += 1 i += 1 return test_data class ApproximateQuantilesBufferTest(unittest.TestCase): """ Approximate Quantiles Buffer Tests to ensure we are calculating the optimal buffers.""" @parameterized.expand(_build_quantilebuffer_test_data) def test_efficiency( self, epsilon, maxInputSize, expectedNumBuffers, expectedBufferSize): """ Verify the buffers are efficiently calculated according to the reference table values. """ combine_fn = ApproximateQuantilesCombineFn.create( num_quantiles=10, max_num_elements=maxInputSize, epsilon=epsilon) self.assertEqual( expectedNumBuffers, combine_fn._spec.num_buffers, "Number of buffers") self.assertEqual( expectedBufferSize, combine_fn._spec.buffer_size, "Buffer size") @parameterized.expand(_build_quantilebuffer_test_data) def test_correctness(self, epsilon, maxInputSize, *args): """ Verify that buffers are correct according to the two constraint equations. """ combine_fn = ApproximateQuantilesCombineFn.create( num_quantiles=10, max_num_elements=maxInputSize, epsilon=epsilon) b = combine_fn._spec.num_buffers k = combine_fn._spec.buffer_size n = maxInputSize self.assertLessEqual((b - 2) * (1 << (b - 2)) + 0.5, (epsilon * n), '(b-2)2^(b-2) + 1/2 <= eN') self.assertGreaterEqual((k * 2)**(b - 1), n, 'k2^(b-1) >= N') if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/timestamped_value_type_test.py0000644000175100017510000001054215206113033026240 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest from typing import Any from typing import TypeVar import apache_beam as beam from apache_beam.transforms.window import TimestampedValue T = TypeVar("T") def ConvertToTimestampedValue(plant: dict[str, Any]) -> TimestampedValue[str]: return TimestampedValue[str](plant["name"], plant["season"]) def ConvertToTimestampedValue_1(plant: dict[str, Any]) -> TimestampedValue: return TimestampedValue(plant["name"], plant["season"]) def ConvertToTimestampedValue_2( plant: dict[str, Any]) -> TimestampedValue[list[str]]: return TimestampedValue[list[str]](plant["name"], plant["season"]) def ConvertToTimestampedValue_3(plant: dict[str, Any]) -> TimestampedValue[T]: return TimestampedValue[T](plant["name"], plant["season"]) class TypeCheckTimestampedValueTestCase(unittest.TestCase): def setUp(self): self.opts = beam.options.pipeline_options.PipelineOptions( runtime_type_check=True) self.data = [ { "name": "Strawberry", "season": 1585699200 }, # April, 2020 ] self.data_1 = [ { "name": 1234, "season": 1585699200 }, # April, 2020 ] self.data_2 = [ { "name": ["abc", "cde"], "season": 1585699200 }, # April, 2020 ] self.data_3 = [ { "name": [123, "cde"], "season": 1585699200 }, # April, 2020 ] def test_pcoll_default_hints(self): for fn in (ConvertToTimestampedValue, ConvertToTimestampedValue_1): pc = beam.Map(fn) ht = pc.default_type_hints() assert len(ht) == 3 assert ht.output_types[0][0] def test_pcoll_with_output_hints(self): pc = beam.Map(ConvertToTimestampedValue).with_output_types(str) ht = pc.get_type_hints() assert len(ht) == 3 assert ht.output_types[0][0] == str def test_opts_with_check(self): with beam.Pipeline(options=self.opts) as p: _ = ( p | "Garden plants" >> beam.Create(self.data) | "With timestamps" >> beam.Map(ConvertToTimestampedValue) | beam.Map(print)) def test_opts_with_check_list_str(self): with beam.Pipeline(options=self.opts) as p: _ = ( p | "Garden plants" >> beam.Create(self.data_2) | "With timestamps" >> beam.Map(ConvertToTimestampedValue_2) | beam.Map(print)) def test_opts_with_check_wrong_data(self): with self.assertRaises(Exception): with beam.Pipeline(options=self.opts) as p: _ = ( p | "Garden plants" >> beam.Create(self.data_1) | "With timestamps" >> beam.Map(ConvertToTimestampedValue) | beam.Map(print)) def test_opts_with_check_wrong_data_list_str(self): with self.assertRaises(Exception): with beam.Pipeline(options=self.opts) as p: _ = ( p | "Garden plants" >> beam.Create(self.data_1) | "With timestamps" >> beam.Map(ConvertToTimestampedValue_2) | beam.Map(print)) with self.assertRaises(Exception): with beam.Pipeline(options=self.opts) as p: _ = ( p | "Garden plants" >> beam.Create(self.data_3) | "With timestamps" >> beam.Map(ConvertToTimestampedValue_2) | beam.Map(print)) def test_opts_with_check_typevar(self): with self.assertRaises(Exception): with beam.Pipeline(options=self.opts) as p: _ = ( p | "Garden plants" >> beam.Create(self.data_2) | "With timestamps" >> beam.Map(ConvertToTimestampedValue_3) | beam.Map(print)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/timeutil.py0000644000175100017510000001070215206113033022262 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Timestamp utilities.""" # pytype: skip-file from abc import ABCMeta from abc import abstractmethod from apache_beam.portability.api import beam_runner_api_pb2 __all__ = [ 'TimeDomain', ] class TimeDomain(object): """Time domain for streaming timers.""" WATERMARK = 'WATERMARK' REAL_TIME = 'REAL_TIME' DEPENDENT_REAL_TIME = 'DEPENDENT_REAL_TIME' _RUNNER_API_MAPPING = { WATERMARK: beam_runner_api_pb2.TimeDomain.EVENT_TIME, REAL_TIME: beam_runner_api_pb2.TimeDomain.PROCESSING_TIME, } @staticmethod def from_string(domain): if domain in (TimeDomain.WATERMARK, TimeDomain.REAL_TIME, TimeDomain.DEPENDENT_REAL_TIME): return domain raise ValueError('Unknown time domain: %s' % domain) @staticmethod def to_runner_api(domain): return TimeDomain._RUNNER_API_MAPPING[domain] @staticmethod def is_event_time(domain): return TimeDomain.from_string(domain) == TimeDomain.WATERMARK class TimestampCombinerImpl(metaclass=ABCMeta): """Implementation of TimestampCombiner.""" @abstractmethod def assign_output_time(self, window, input_timestamp): raise NotImplementedError @abstractmethod def combine(self, output_timestamp, other_output_timestamp): raise NotImplementedError def combine_all(self, merging_timestamps): """Apply combine to list of timestamps.""" combined_output_time = None for output_time in merging_timestamps: if combined_output_time is None: combined_output_time = output_time elif output_time is not None: combined_output_time = self.combine(combined_output_time, output_time) return combined_output_time def merge(self, unused_result_window, merging_timestamps): """Default to returning the result of combine_all.""" return self.combine_all(merging_timestamps) class DependsOnlyOnWindow(TimestampCombinerImpl, metaclass=ABCMeta): """TimestampCombinerImpl that only depends on the window.""" def merge(self, result_window, unused_merging_timestamps): # Since we know that the result only depends on the window, we can ignore # the given timestamps. return self.assign_output_time(result_window, None) class OutputAtEarliestInputTimestampImpl(TimestampCombinerImpl): """TimestampCombinerImpl outputting at earliest input timestamp.""" def assign_output_time(self, window, input_timestamp): return input_timestamp def combine(self, output_timestamp, other_output_timestamp): """Default to returning the earlier of two timestamps.""" return min(output_timestamp, other_output_timestamp) class OutputAtEarliestTransformedInputTimestampImpl(TimestampCombinerImpl): """TimestampCombinerImpl outputting at earliest input timestamp.""" def __init__(self, window_fn): self.window_fn = window_fn def assign_output_time(self, window, input_timestamp): return self.window_fn.get_transformed_output_time(window, input_timestamp) def combine(self, output_timestamp, other_output_timestamp): return min(output_timestamp, other_output_timestamp) class OutputAtLatestInputTimestampImpl(TimestampCombinerImpl): """TimestampCombinerImpl outputting at latest input timestamp.""" def assign_output_time(self, window, input_timestamp): return input_timestamp def combine(self, output_timestamp, other_output_timestamp): return max(output_timestamp, other_output_timestamp) class OutputAtEndOfWindowImpl(DependsOnlyOnWindow): """TimestampCombinerImpl outputting at end of window.""" def assign_output_time(self, window, unused_input_timestamp): return window.max_timestamp() def combine(self, output_timestamp, other_output_timestamp): return max(output_timestamp, other_output_timestamp) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/transforms_keyword_only_args_test.py0000644000175100017510000001275415206113033027515 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for side inputs.""" # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to class KeywordOnlyArgsTests(unittest.TestCase): def test_side_input_keyword_only_args(self): with TestPipeline() as pipeline: def sort_with_side_inputs(x, *s, reverse=False): for y in s: yield sorted([x] + y, reverse=reverse) def sort_with_side_inputs_without_default_values(x, *s, reverse): for y in s: yield sorted([x] + y, reverse=reverse) pcol = pipeline | 'start' >> beam.Create([1, 2]) side = pipeline | 'side' >> beam.Create([3, 4]) # 2 values in side input. result1 = pcol | 'compute1' >> beam.FlatMap( sort_with_side_inputs, beam.pvalue.AsList(side), reverse=True) assert_that(result1, equal_to([[4, 3, 1], [4, 3, 2]]), label='assert1') result2 = pcol | 'compute2' >> beam.FlatMap( sort_with_side_inputs, beam.pvalue.AsList(side)) assert_that(result2, equal_to([[1, 3, 4], [2, 3, 4]]), label='assert2') result3 = pcol | 'compute3' >> beam.FlatMap(sort_with_side_inputs) assert_that(result3, equal_to([]), label='assert3') result4 = pcol | 'compute4' >> beam.FlatMap( sort_with_side_inputs, reverse=True) assert_that(result4, equal_to([]), label='assert4') result5 = pcol | 'compute5' >> beam.FlatMap( sort_with_side_inputs_without_default_values, beam.pvalue.AsList(side), reverse=True) assert_that(result5, equal_to([[4, 3, 1], [4, 3, 2]]), label='assert5') result6 = pcol | 'compute6' >> beam.FlatMap( sort_with_side_inputs_without_default_values, beam.pvalue.AsList(side), reverse=False) assert_that(result6, equal_to([[1, 3, 4], [2, 3, 4]]), label='assert6') result7 = pcol | 'compute7' >> beam.FlatMap( sort_with_side_inputs_without_default_values, reverse=False) assert_that(result7, equal_to([]), label='assert7') result8 = pcol | 'compute8' >> beam.FlatMap( sort_with_side_inputs_without_default_values, reverse=True) assert_that(result8, equal_to([]), label='assert8') def test_combine_keyword_only_args(self): with TestPipeline() as pipeline: def bounded_sum(values, *s, bound=500): return min(sum(values) + sum(s), bound) def bounded_sum_without_default_values(values, *s, bound): return min(sum(values) + sum(s), bound) pcoll = pipeline | 'start' >> beam.Create([6, 3, 1]) result1 = pcoll | 'sum1' >> beam.CombineGlobally( bounded_sum, 5, 8, bound=20) result2 = pcoll | 'sum2' >> beam.CombineGlobally(bounded_sum, 0, 0) result3 = pcoll | 'sum3' >> beam.CombineGlobally(bounded_sum) result4 = pcoll | 'sum4' >> beam.CombineGlobally(bounded_sum, bound=5) result5 = pcoll | 'sum5' >> beam.CombineGlobally( bounded_sum_without_default_values, 5, 8, bound=20) result6 = pcoll | 'sum6' >> beam.CombineGlobally( bounded_sum_without_default_values, 0, 0, bound=500) result7 = pcoll | 'sum7' >> beam.CombineGlobally( bounded_sum_without_default_values, bound=500) result8 = pcoll | 'sum8' >> beam.CombineGlobally( bounded_sum_without_default_values, bound=5) assert_that(result1, equal_to([20]), label='assert1') assert_that(result2, equal_to([10]), label='assert2') assert_that(result3, equal_to([10]), label='assert3') assert_that(result4, equal_to([5]), label='assert4') assert_that(result5, equal_to([20]), label='assert5') assert_that(result6, equal_to([10]), label='assert6') assert_that(result7, equal_to([10]), label='assert7') assert_that(result8, equal_to([5]), label='assert8') def test_do_fn_keyword_only_args(self): with TestPipeline() as pipeline: class MyDoFn(beam.DoFn): def process(self, element, *s, bound=500): return [min(sum(s) + element, bound)] pcoll = pipeline | 'start' >> beam.Create([6, 3, 1]) result1 = pcoll | 'sum1' >> beam.ParDo(MyDoFn(), 5, 8, bound=15) result2 = pcoll | 'sum2' >> beam.ParDo(MyDoFn(), 5, 8) result3 = pcoll | 'sum3' >> beam.ParDo(MyDoFn()) result4 = pcoll | 'sum4' >> beam.ParDo(MyDoFn(), bound=5) assert_that(result1, equal_to([15, 15, 14]), label='assert1') assert_that(result2, equal_to([19, 16, 14]), label='assert2') assert_that(result3, equal_to([6, 3, 1]), label='assert3') assert_that(result4, equal_to([5, 3, 1]), label='assert4') if __name__ == '__main__': logging.getLogger().setLevel(logging.DEBUG) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/trigger.py0000644000175100017510000016040715206113033022101 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Support for Apache Beam triggers. Triggers control when in processing time windows get emitted. """ # pytype: skip-file import collections import copy import logging import numbers from abc import ABCMeta from abc import abstractmethod from collections import abc as collections_abc # ambiguty with direct abc from enum import Flag from enum import auto from itertools import zip_longest from apache_beam.coders import coder_impl from apache_beam.coders import observable from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.transforms import combiners from apache_beam.transforms import core from apache_beam.transforms.timeutil import TimeDomain from apache_beam.transforms.window import GlobalWindow from apache_beam.transforms.window import GlobalWindows from apache_beam.transforms.window import TimestampCombiner from apache_beam.transforms.window import WindowedValue from apache_beam.transforms.window import WindowFn from apache_beam.utils import windowed_value from apache_beam.utils.timestamp import MAX_TIMESTAMP from apache_beam.utils.timestamp import MIN_TIMESTAMP from apache_beam.utils.timestamp import TIME_GRANULARITY __all__ = [ 'AccumulationMode', 'TriggerFn', 'DefaultTrigger', 'AfterWatermark', 'AfterProcessingTime', 'AfterCount', 'Repeatedly', 'AfterAny', 'AfterAll', 'AfterEach', 'OrFinally', ] _LOGGER = logging.getLogger(__name__) class AccumulationMode(object): """Controls what to do with data when a trigger fires multiple times.""" DISCARDING = beam_runner_api_pb2.AccumulationMode.DISCARDING ACCUMULATING = beam_runner_api_pb2.AccumulationMode.ACCUMULATING # TODO(robertwb): Provide retractions of previous outputs. # RETRACTING = 3 class _StateTag(metaclass=ABCMeta): """An identifier used to store and retrieve typed, combinable state. The given tag must be unique for this step.""" def __init__(self, tag): self.tag = tag class _ReadModifyWriteStateTag(_StateTag): """StateTag pointing to an element.""" def __repr__(self): return 'ValueStateTag(%s)' % (self.tag) def with_prefix(self, prefix): return _ReadModifyWriteStateTag(prefix + self.tag) class _SetStateTag(_StateTag): """StateTag pointing to an element.""" def __repr__(self): return 'SetStateTag({tag})'.format(tag=self.tag) def with_prefix(self, prefix): return _SetStateTag(prefix + self.tag) class _CombiningValueStateTag(_StateTag): """StateTag pointing to an element, accumulated with a combiner. The given tag must be unique for this step. The given CombineFn will be applied (possibly incrementally and eagerly) when adding elements.""" # TODO(robertwb): Also store the coder (perhaps extracted from the combine_fn) def __init__(self, tag, combine_fn): super().__init__(tag) if not combine_fn: raise ValueError('combine_fn must be specified.') if not isinstance(combine_fn, core.CombineFn): combine_fn = core.CombineFn.from_callable(combine_fn) self.combine_fn = combine_fn def __repr__(self): return 'CombiningValueStateTag(%s, %s)' % (self.tag, self.combine_fn) def with_prefix(self, prefix): return _CombiningValueStateTag(prefix + self.tag, self.combine_fn) def without_extraction(self): class NoExtractionCombineFn(core.CombineFn): setup = self.combine_fn.setup create_accumulator = self.combine_fn.create_accumulator add_input = self.combine_fn.add_input merge_accumulators = self.combine_fn.merge_accumulators compact = self.combine_fn.compact extract_output = staticmethod(lambda x: x) teardown = self.combine_fn.teardown return _CombiningValueStateTag(self.tag, NoExtractionCombineFn()) class _ListStateTag(_StateTag): """StateTag pointing to a list of elements.""" def __repr__(self): return 'ListStateTag(%s)' % self.tag def with_prefix(self, prefix): return _ListStateTag(prefix + self.tag) class _WatermarkHoldStateTag(_StateTag): def __init__(self, tag, timestamp_combiner_impl): super().__init__(tag) self.timestamp_combiner_impl = timestamp_combiner_impl def __repr__(self): return 'WatermarkHoldStateTag(%s, %s)' % ( self.tag, self.timestamp_combiner_impl) def with_prefix(self, prefix): return _WatermarkHoldStateTag( prefix + self.tag, self.timestamp_combiner_impl) class DataLossReason(Flag): """Enum defining potential reasons that a trigger may cause data loss. These flags should only cover when the trigger is the cause, though windowing can be taken into account. For instance, AfterWatermark may not flag itself as finishing if the windowing doesn't allow lateness. """ # Trigger will never be the source of data loss. NO_POTENTIAL_LOSS = 0 # Trigger may finish. In this case, data that comes in after the trigger may # be lost. Example: AfterCount(1) will stop firing after the first element. MAY_FINISH = auto() # Deprecated: Beam will emit buffered data at GC time. Any other behavior # should be treated as a bug with the runner used. CONDITION_NOT_GUARANTEED = auto() # Convenience functions for checking if a flag is included. Each is equivalent # to `reason & flag == flag` def _IncludesMayFinish(reason: DataLossReason) -> bool: return reason & DataLossReason.MAY_FINISH == DataLossReason.MAY_FINISH # pylint: disable=unused-argument # TODO(robertwb): Provisional API, Java likely to change as well. class TriggerFn(metaclass=ABCMeta): """A TriggerFn determines when window (panes) are emitted. See https://beam.apache.org/documentation/programming-guide/#triggers """ @abstractmethod def on_element(self, element, window, context): """Called when a new element arrives in a window. Args: element: the element being added window: the window to which the element is being added context: a context (e.g. a TriggerContext instance) for managing state and setting timers """ pass @abstractmethod def on_merge(self, to_be_merged, merge_result, context): """Called when multiple windows are merged. Args: to_be_merged: the set of windows to be merged merge_result: the window into which the windows are being merged context: a context (e.g. a TriggerContext instance) for managing state and setting timers """ pass @abstractmethod def should_fire(self, time_domain, timestamp, window, context): """Whether this trigger should cause the window to fire. Args: time_domain: WATERMARK for event-time timers and REAL_TIME for processing-time timers. timestamp: for time_domain WATERMARK, it represents the watermark: (a lower bound on) the watermark of the system and for time_domain REAL_TIME, it represents the trigger: timestamp of the processing-time timer. window: the window whose trigger is being considered context: a context (e.g. a TriggerContext instance) for managing state and setting timers Returns: whether this trigger should cause a firing """ pass @abstractmethod def has_ontime_pane(self): """Whether this trigger creates an empty pane even if there are no elements. Returns: True if this trigger guarantees that there will always be an ON_TIME pane even if there are no elements in that pane. """ pass @abstractmethod def on_fire(self, watermark, window, context): """Called when a trigger actually fires. Args: watermark: (a lower bound on) the watermark of the system window: the window whose trigger is being fired context: a context (e.g. a TriggerContext instance) for managing state and setting timers Returns: whether this trigger is finished """ pass @abstractmethod def reset(self, window, context): """Clear any state and timers used by this TriggerFn.""" pass def may_lose_data(self, unused_windowing: core.Windowing) -> DataLossReason: """Returns whether or not this trigger could cause data loss. A trigger can cause data loss in the following scenarios: * The trigger has a chance to finish. For instance, AfterWatermark() without a late trigger would cause all late data to be lost. This scenario is only accounted for if the windowing strategy allows late data. Otherwise, the trigger is not responsible for the data loss. Note that this only returns the potential for loss. It does not mean that there will be data loss. It also only accounts for loss related to the trigger, not other potential causes. Args: windowing: The Windowing that this trigger belongs to. It does not need to be the top-level trigger. Returns: The DataLossReason. If there is no potential loss, DataLossReason.NO_POTENTIAL_LOSS is returned. Otherwise, all the potential reasons are returned as a single value. """ # For backwards compatibility's sake, we're assuming the trigger is safe. return DataLossReason.NO_POTENTIAL_LOSS # pylint: enable=unused-argument @staticmethod def from_runner_api(proto, context): return { 'after_all': AfterAll, 'after_any': AfterAny, 'after_each': AfterEach, 'after_end_of_window': AfterWatermark, 'after_processing_time': AfterProcessingTime, 'after_synchronized_processing_time': _AfterSynchronizedProcessingTime, 'always': Always, 'default': DefaultTrigger, 'element_count': AfterCount, 'never': _Never, 'or_finally': OrFinally, 'repeat': Repeatedly, }[proto.WhichOneof('trigger')].from_runner_api(proto, context) @abstractmethod def to_runner_api(self, unused_context): pass @abstractmethod def get_continuation_trigger(self): """Returns: Trigger to use after a GroupBy to preserve the intention of this trigger. Specifically, triggers that are time based and intended to provide speculative results should continue providing speculative results. Triggers that fire once (or multiple times) should continue firing once (or multiple times). """ pass class DefaultTrigger(TriggerFn): """Semantically Repeatedly(AfterWatermark()), but more optimized.""" def __init__(self): pass def __repr__(self): return 'DefaultTrigger()' def on_element(self, element, window, context): context.set_timer(str(window), TimeDomain.WATERMARK, window.end) def on_merge(self, to_be_merged, merge_result, context): for window in to_be_merged: context.clear_timer(str(window), TimeDomain.WATERMARK) def should_fire(self, time_domain, watermark, window, context): if watermark >= window.end: # Explicitly clear the timer so that late elements are not emitted again # when the timer is fired. context.clear_timer(str(window), TimeDomain.WATERMARK) return watermark >= window.end def on_fire(self, watermark, window, context): return False def reset(self, window, context): context.clear_timer(str(window), TimeDomain.WATERMARK) def may_lose_data(self, unused_windowing): return DataLossReason.NO_POTENTIAL_LOSS def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) @staticmethod def from_runner_api(proto, context): return DefaultTrigger() def to_runner_api(self, unused_context): return beam_runner_api_pb2.Trigger( default=beam_runner_api_pb2.Trigger.Default()) def has_ontime_pane(self): return True def get_continuation_trigger(self): return self class AfterProcessingTime(TriggerFn): """Fire exactly once after a specified delay from processing time.""" STATE_TAG = _SetStateTag('has_timer') def __init__(self, delay=0): """Initialize a processing time trigger with a delay in seconds.""" self.delay = delay def __repr__(self): return 'AfterProcessingTime(delay=%d)' % self.delay def on_element(self, element, window, context): if not context.get_state(self.STATE_TAG): context.set_timer( '', TimeDomain.REAL_TIME, context.get_current_time() + self.delay) context.add_state(self.STATE_TAG, True) def on_merge(self, to_be_merged, merge_result, context): # timers will be kept through merging pass def should_fire(self, time_domain, timestamp, window, context): if time_domain == TimeDomain.REAL_TIME: return True def on_fire(self, timestamp, window, context): return True def reset(self, window, context): context.clear_state(self.STATE_TAG) def may_lose_data(self, unused_windowing): """AfterProcessingTime may finish.""" return DataLossReason.MAY_FINISH @staticmethod def from_runner_api(proto, context): return AfterProcessingTime( delay=( proto.after_processing_time.timestamp_transforms[0].delay. delay_millis) // 1000) def to_runner_api(self, context): delay_proto = beam_runner_api_pb2.TimestampTransform( delay=beam_runner_api_pb2.TimestampTransform.Delay( delay_millis=self.delay * 1000)) return beam_runner_api_pb2.Trigger( after_processing_time=beam_runner_api_pb2.Trigger.AfterProcessingTime( timestamp_transforms=[delay_proto])) def has_ontime_pane(self): return False def get_continuation_trigger(self): # The continuation of an AfterProcessingTime trigger is an # _AfterSynchronizedProcessingTime trigger. return _AfterSynchronizedProcessingTime() class Always(TriggerFn): """Repeatedly invoke the given trigger, never finishing.""" def __init__(self): pass def __repr__(self): return 'Always' def __eq__(self, other): return type(self) == type(other) def __hash__(self): return 1 def on_element(self, element, window, context): pass def on_merge(self, to_be_merged, merge_result, context): pass def has_ontime_pane(self): return False def reset(self, window, context): pass def should_fire(self, time_domain, watermark, window, context): return True def on_fire(self, watermark, window, context): return False def may_lose_data(self, unused_windowing): """No potential loss, since the trigger always fires.""" return DataLossReason.NO_POTENTIAL_LOSS @staticmethod def from_runner_api(proto, context): return Always() def to_runner_api(self, context): return beam_runner_api_pb2.Trigger( always=beam_runner_api_pb2.Trigger.Always()) def get_continuation_trigger(self): return self class _Never(TriggerFn): """A trigger that never fires. Data may still be released at window closing. """ def __init__(self): pass def __repr__(self): return 'Never' def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) def on_element(self, element, window, context): pass def on_merge(self, to_be_merged, merge_result, context): pass def has_ontime_pane(self): False def reset(self, window, context): pass def should_fire(self, time_domain, watermark, window, context): return False def on_fire(self, watermark, window, context): return True def may_lose_data(self, unused_windowing): """No potential data loss. Though Never doesn't explicitly trigger, it still collects data on windowing closing. """ return DataLossReason.NO_POTENTIAL_LOSS @staticmethod def from_runner_api(proto, context): return _Never() def to_runner_api(self, context): return beam_runner_api_pb2.Trigger( never=beam_runner_api_pb2.Trigger.Never()) def get_continuation_trigger(self): return self class AfterWatermark(TriggerFn): """Fire exactly once when the watermark passes the end of the window. Args: early: if not None, a speculative trigger to repeatedly evaluate before the watermark passes the end of the window late: if not None, a speculative trigger to repeatedly evaluate after the watermark passes the end of the window """ LATE_TAG = _CombiningValueStateTag('is_late', any) def __init__(self, early=None, late=None): self.early = self._wrap_if_not_repeatedly(early) self.late = self._wrap_if_not_repeatedly(late) @staticmethod def _wrap_if_not_repeatedly(trigger): if trigger and not isinstance(trigger, Repeatedly): return Repeatedly(trigger) return trigger def get_continuation_trigger(self): return AfterWatermark( self.early.get_continuation_trigger() if self.early else None, self.late.get_continuation_trigger() if self.late else None) def __repr__(self): qualifiers = [] if self.early: qualifiers.append('early=%s' % self.early.underlying) if self.late: qualifiers.append('late=%s' % self.late.underlying) return 'AfterWatermark(%s)' % ', '.join(qualifiers) def is_late(self, context): return self.late and context.get_state(self.LATE_TAG) def on_element(self, element, window, context): if self.is_late(context): self.late.on_element(element, window, NestedContext(context, 'late')) else: context.set_timer('', TimeDomain.WATERMARK, window.end) if self.early: self.early.on_element(element, window, NestedContext(context, 'early')) def on_merge(self, to_be_merged, merge_result, context): # TODO(robertwb): Figure out whether the 'rewind' semantics could be used # here. if self.is_late(context): self.late.on_merge( to_be_merged, merge_result, NestedContext(context, 'late')) else: # Note: Timer clearing solely an optimization. for window in to_be_merged: if window.end != merge_result.end: context.clear_timer('', TimeDomain.WATERMARK) if self.early: self.early.on_merge( to_be_merged, merge_result, NestedContext(context, 'early')) def should_fire(self, time_domain, watermark, window, context): if self.is_late(context): return self.late.should_fire( time_domain, watermark, window, NestedContext(context, 'late')) elif watermark >= window.end: # Explicitly clear the timer so that late elements are not emitted again # when the timer is fired. context.clear_timer('', TimeDomain.WATERMARK) return True elif self.early: return self.early.should_fire( time_domain, watermark, window, NestedContext(context, 'early')) return False def on_fire(self, watermark, window, context): if self.is_late(context): return self.late.on_fire( watermark, window, NestedContext(context, 'late')) elif watermark >= window.end: context.add_state(self.LATE_TAG, True) return not self.late elif self.early: self.early.on_fire(watermark, window, NestedContext(context, 'early')) return False def reset(self, window, context): if self.late: context.clear_state(self.LATE_TAG) if self.early: self.early.reset(window, NestedContext(context, 'early')) if self.late: self.late.reset(window, NestedContext(context, 'late')) def may_lose_data(self, windowing): """May cause data loss if lateness allowed and no late trigger set.""" if windowing.allowed_lateness == 0: return DataLossReason.NO_POTENTIAL_LOSS if self.late is None: return DataLossReason.MAY_FINISH return self.late.may_lose_data(windowing) def __eq__(self, other): return ( type(self) == type(other) and self.early == other.early and self.late == other.late) def __hash__(self): return hash((type(self), self.early, self.late)) @staticmethod def from_runner_api(proto, context): return AfterWatermark( early=TriggerFn.from_runner_api( proto.after_end_of_window.early_firings, context) if proto.after_end_of_window.HasField('early_firings') else None, late=TriggerFn.from_runner_api( proto.after_end_of_window.late_firings, context) if proto.after_end_of_window.HasField('late_firings') else None) def to_runner_api(self, context): early_proto = self.early.underlying.to_runner_api( context) if self.early else None late_proto = self.late.underlying.to_runner_api( context) if self.late else None return beam_runner_api_pb2.Trigger( after_end_of_window=beam_runner_api_pb2.Trigger.AfterEndOfWindow( early_firings=early_proto, late_firings=late_proto)) def has_ontime_pane(self): return True class AfterCount(TriggerFn): """Fire when there are at least count elements in this window pane.""" COUNT_TAG = _CombiningValueStateTag('count', combiners.CountCombineFn()) def __init__(self, count): if not isinstance(count, numbers.Integral) or count < 1: raise ValueError("count (%d) must be a positive integer." % count) self.count = count def __repr__(self): return 'AfterCount(%s)' % self.count def __eq__(self, other): return type(self) == type(other) and self.count == other.count def __hash__(self): return hash(self.count) def on_element(self, element, window, context): context.add_state(self.COUNT_TAG, 1) def on_merge(self, to_be_merged, merge_result, context): # states automatically merged pass def should_fire(self, time_domain, watermark, window, context): return context.get_state(self.COUNT_TAG) >= self.count def on_fire(self, watermark, window, context): return True def reset(self, window, context): context.clear_state(self.COUNT_TAG) def may_lose_data(self, unused_windowing): """AfterCount may finish.""" return DataLossReason.MAY_FINISH @staticmethod def from_runner_api(proto, unused_context): return AfterCount(proto.element_count.element_count) def to_runner_api(self, unused_context): return beam_runner_api_pb2.Trigger( element_count=beam_runner_api_pb2.Trigger.ElementCount( element_count=self.count)) def has_ontime_pane(self): return False def get_continuation_trigger(self): return AfterCount(1) class Repeatedly(TriggerFn): """Repeatedly invoke the given trigger, never finishing.""" def __init__(self, underlying): self.underlying = underlying def __repr__(self): return 'Repeatedly(%s)' % self.underlying def __eq__(self, other): return type(self) == type(other) and self.underlying == other.underlying def __hash__(self): return hash(self.underlying) def on_element(self, element, window, context): self.underlying.on_element(element, window, context) def on_merge(self, to_be_merged, merge_result, context): self.underlying.on_merge(to_be_merged, merge_result, context) def should_fire(self, time_domain, watermark, window, context): return self.underlying.should_fire(time_domain, watermark, window, context) def on_fire(self, watermark, window, context): if self.underlying.on_fire(watermark, window, context): self.underlying.reset(window, context) return False def reset(self, window, context): self.underlying.reset(window, context) def may_lose_data(self, windowing): """Repeatedly will run in a loop and pick up whatever is left at GC.""" return DataLossReason.NO_POTENTIAL_LOSS @staticmethod def from_runner_api(proto, context): return Repeatedly( TriggerFn.from_runner_api(proto.repeat.subtrigger, context)) def to_runner_api(self, context): return beam_runner_api_pb2.Trigger( repeat=beam_runner_api_pb2.Trigger.Repeat( subtrigger=self.underlying.to_runner_api(context))) def has_ontime_pane(self): return self.underlying.has_ontime_pane() def get_continuation_trigger(self): return Repeatedly(self.underlying.get_continuation_trigger()) class _ParallelTriggerFn(TriggerFn, metaclass=ABCMeta): def __init__(self, *triggers): self.triggers = triggers def __repr__(self): return '%s(%s)' % ( self.__class__.__name__, ', '.join(str(t) for t in self.triggers)) def __eq__(self, other): return type(self) == type(other) and self.triggers == other.triggers def __hash__(self): return hash(self.triggers) @abstractmethod def combine_op(self, trigger_results): pass def on_element(self, element, window, context): for ix, trigger in enumerate(self.triggers): trigger.on_element(element, window, self._sub_context(context, ix)) def on_merge(self, to_be_merged, merge_result, context): for ix, trigger in enumerate(self.triggers): trigger.on_merge( to_be_merged, merge_result, self._sub_context(context, ix)) def should_fire(self, time_domain, watermark, window, context): self._time_domain = time_domain return self.combine_op( trigger.should_fire( time_domain, watermark, window, self._sub_context(context, ix)) for ix, trigger in enumerate(self.triggers)) def on_fire(self, watermark, window, context): finished = [] for ix, trigger in enumerate(self.triggers): nested_context = self._sub_context(context, ix) if trigger.should_fire(TimeDomain.WATERMARK, watermark, window, nested_context): finished.append(trigger.on_fire(watermark, window, nested_context)) return self.combine_op(finished) def may_lose_data(self, windowing): may_finish = self.combine_op( _IncludesMayFinish(t.may_lose_data(windowing)) for t in self.triggers) return ( DataLossReason.MAY_FINISH if may_finish else DataLossReason.NO_POTENTIAL_LOSS) def reset(self, window, context): for ix, trigger in enumerate(self.triggers): trigger.reset(window, self._sub_context(context, ix)) @staticmethod def _sub_context(context, index): return NestedContext(context, '%d/' % index) @staticmethod def from_runner_api(proto, context): subtriggers = [ TriggerFn.from_runner_api(subtrigger, context) for subtrigger in proto.after_all.subtriggers or proto.after_any.subtriggers ] if proto.after_all.subtriggers: return AfterAll(*subtriggers) else: return AfterAny(*subtriggers) def to_runner_api(self, context): subtriggers = [ subtrigger.to_runner_api(context) for subtrigger in self.triggers ] if self.combine_op == all: return beam_runner_api_pb2.Trigger( after_all=beam_runner_api_pb2.Trigger.AfterAll( subtriggers=subtriggers)) elif self.combine_op == any: return beam_runner_api_pb2.Trigger( after_any=beam_runner_api_pb2.Trigger.AfterAny( subtriggers=subtriggers)) else: raise NotImplementedError(self) def has_ontime_pane(self): return any(t.has_ontime_pane() for t in self.triggers) def get_continuation_trigger(self): return self.__class__( *( subtrigger.get_continuation_trigger() for subtrigger in self.triggers)) class AfterAny(_ParallelTriggerFn): """Fires when any subtrigger fires. Also finishes when any subtrigger finishes. """ combine_op = any class AfterAll(_ParallelTriggerFn): """Fires when all subtriggers have fired. Also finishes when all subtriggers have finished. """ combine_op = all class AfterEach(TriggerFn): INDEX_TAG = _CombiningValueStateTag( 'index', (lambda indices: 0 if not indices else max(indices))) def __init__(self, *triggers): self.triggers = triggers def __repr__(self): return '%s(%s)' % ( self.__class__.__name__, ', '.join(str(t) for t in self.triggers)) def __eq__(self, other): return type(self) == type(other) and self.triggers == other.triggers def __hash__(self): return hash(self.triggers) def on_element(self, element, window, context): ix = context.get_state(self.INDEX_TAG) if ix < len(self.triggers): self.triggers[ix].on_element( element, window, self._sub_context(context, ix)) def on_merge(self, to_be_merged, merge_result, context): # This takes the furthest window on merging. # TODO(robertwb): Revisit this when merging windows logic is settled for # all possible merging situations. ix = context.get_state(self.INDEX_TAG) if ix < len(self.triggers): self.triggers[ix].on_merge( to_be_merged, merge_result, self._sub_context(context, ix)) def should_fire(self, time_domain, watermark, window, context): ix = context.get_state(self.INDEX_TAG) if ix < len(self.triggers): return self.triggers[ix].should_fire( time_domain, watermark, window, self._sub_context(context, ix)) def on_fire(self, watermark, window, context): ix = context.get_state(self.INDEX_TAG) if ix < len(self.triggers): if self.triggers[ix].on_fire(watermark, window, self._sub_context(context, ix)): ix += 1 context.add_state(self.INDEX_TAG, ix) return ix == len(self.triggers) def reset(self, window, context): context.clear_state(self.INDEX_TAG) for ix, trigger in enumerate(self.triggers): trigger.reset(window, self._sub_context(context, ix)) def may_lose_data(self, windowing): """If all sub-triggers may finish, this may finish.""" may_finish = all( _IncludesMayFinish(t.may_lose_data(windowing)) for t in self.triggers) return ( DataLossReason.MAY_FINISH if may_finish else DataLossReason.NO_POTENTIAL_LOSS) @staticmethod def _sub_context(context, index): return NestedContext(context, '%d/' % index) @staticmethod def from_runner_api(proto, context): return AfterEach( *[ TriggerFn.from_runner_api(subtrigger, context) for subtrigger in proto.after_each.subtriggers ]) def to_runner_api(self, context): return beam_runner_api_pb2.Trigger( after_each=beam_runner_api_pb2.Trigger.AfterEach( subtriggers=[ subtrigger.to_runner_api(context) for subtrigger in self.triggers ])) def has_ontime_pane(self): return any(t.has_ontime_pane() for t in self.triggers) def get_continuation_trigger(self): return Repeatedly( AfterAny( *( subtrigger.get_continuation_trigger() for subtrigger in self.triggers))) class OrFinally(AfterAny): @staticmethod def from_runner_api(proto, context): return OrFinally( TriggerFn.from_runner_api(proto.or_finally.main, context), # getattr is used as finally is a keyword in Python TriggerFn.from_runner_api( getattr(proto.or_finally, 'finally'), context)) def to_runner_api(self, context): return beam_runner_api_pb2.Trigger( or_finally=beam_runner_api_pb2.Trigger.OrFinally( main=self.triggers[0].to_runner_api(context), # dict keyword argument is used as finally is a keyword in Python **{'finally': self.triggers[1].to_runner_api(context)})) class TriggerContext(object): def __init__(self, outer, window, clock): self._outer = outer self._window = window self._clock = clock def get_current_time(self): return self._clock.time() def set_timer(self, name, time_domain, timestamp): self._outer.set_timer(self._window, name, time_domain, timestamp) def clear_timer(self, name, time_domain): self._outer.clear_timer(self._window, name, time_domain) def add_state(self, tag, value): self._outer.add_state(self._window, tag, value) def get_state(self, tag): return self._outer.get_state(self._window, tag) def clear_state(self, tag): return self._outer.clear_state(self._window, tag) class NestedContext(object): """Namespaced context useful for defining composite triggers.""" def __init__(self, outer, prefix): self._outer = outer self._prefix = prefix def get_current_time(self): return self._outer.get_current_time() def set_timer(self, name, time_domain, timestamp): self._outer.set_timer(self._prefix + name, time_domain, timestamp) def clear_timer(self, name, time_domain): self._outer.clear_timer(self._prefix + name, time_domain) def add_state(self, tag, value): self._outer.add_state(tag.with_prefix(self._prefix), value) def get_state(self, tag): return self._outer.get_state(tag.with_prefix(self._prefix)) def clear_state(self, tag): self._outer.clear_state(tag.with_prefix(self._prefix)) # pylint: disable=unused-argument class SimpleState(metaclass=ABCMeta): """Basic state storage interface used for triggering. Only timers must hold the watermark (by their timestamp). """ @abstractmethod def set_timer( self, window, name, time_domain, timestamp, dynamic_timer_tag=''): pass @abstractmethod def get_window(self, window_id): pass @abstractmethod def clear_timer(self, window, name, time_domain, dynamic_timer_tag=''): pass @abstractmethod def add_state(self, window, tag, value): pass @abstractmethod def get_state(self, window, tag): pass @abstractmethod def clear_state(self, window, tag): pass def at(self, window, clock): return NestedContext(TriggerContext(self, window, clock), 'trigger') class UnmergedState(SimpleState): """State suitable for use in TriggerDriver. This class must be implemented by each backend. """ @abstractmethod def set_global_state(self, tag, value): pass @abstractmethod def get_global_state(self, tag, default=None): pass # pylint: enable=unused-argument class MergeableStateAdapter(SimpleState): """Wraps an UnmergedState, tracking merged windows.""" # TODO(robertwb): A similar indirection could be used for sliding windows # or other window_fns when a single element typically belongs to many windows. WINDOW_IDS = _ReadModifyWriteStateTag('window_ids') def __init__(self, raw_state): self.raw_state = raw_state self.window_ids = self.raw_state.get_global_state(self.WINDOW_IDS, {}) self.counter = None def set_timer( self, window, name, time_domain, timestamp, dynamic_timer_tag=''): self.raw_state.set_timer( self._get_id(window), name, time_domain, timestamp, dynamic_timer_tag=dynamic_timer_tag) def clear_timer(self, window, name, time_domain, dynamic_timer_tag=''): for window_id in self._get_ids(window): self.raw_state.clear_timer( window_id, name, time_domain, dynamic_timer_tag=dynamic_timer_tag) def add_state(self, window, tag, value): if isinstance(tag, _ReadModifyWriteStateTag): raise ValueError( 'Merging requested for non-mergeable state tag: %r.' % tag) elif isinstance(tag, _CombiningValueStateTag): tag = tag.without_extraction() self.raw_state.add_state(self._get_id(window), tag, value) def get_state(self, window, tag): if isinstance(tag, _CombiningValueStateTag): original_tag, tag = tag, tag.without_extraction() values = [ self.raw_state.get_state(window_id, tag) for window_id in self._get_ids(window) ] if isinstance(tag, _ReadModifyWriteStateTag): raise ValueError( 'Merging requested for non-mergeable state tag: %r.' % tag) elif isinstance(tag, _CombiningValueStateTag): return original_tag.combine_fn.extract_output( original_tag.combine_fn.merge_accumulators(values)) elif isinstance(tag, _ListStateTag): return [v for vs in values for v in vs] elif isinstance(tag, _SetStateTag): return {v for vs in values for v in vs} elif isinstance(tag, _WatermarkHoldStateTag): return tag.timestamp_combiner_impl.combine_all(values) else: raise ValueError('Invalid tag.', tag) def clear_state(self, window, tag): for window_id in self._get_ids(window): self.raw_state.clear_state(window_id, tag) if tag is None: del self.window_ids[window] self._persist_window_ids() def merge(self, to_be_merged, merge_result): for window in to_be_merged: if window != merge_result: if window in self.window_ids: if merge_result in self.window_ids: merge_window_ids = self.window_ids[merge_result] else: merge_window_ids = self.window_ids[merge_result] = [] merge_window_ids.extend(self.window_ids.pop(window)) self._persist_window_ids() def known_windows(self): return list(self.window_ids) def get_window(self, window_id): for window, ids in self.window_ids.items(): if window_id in ids: return window raise ValueError('No window for %s' % window_id) def _get_id(self, window): if window in self.window_ids: return self.window_ids[window][0] window_id = self._get_next_counter() self.window_ids[window] = [window_id] self._persist_window_ids() return window_id def _get_ids(self, window): return self.window_ids.get(window, []) def _get_next_counter(self): if not self.window_ids: self.counter = 0 elif self.counter is None: self.counter = max(k for ids in self.window_ids.values() for k in ids) self.counter += 1 return self.counter def _persist_window_ids(self): self.raw_state.set_global_state(self.WINDOW_IDS, self.window_ids) def __repr__(self): return '\n\t'.join([repr(self.window_ids)] + repr(self.raw_state).split('\n')) def create_trigger_driver( windowing, is_batch=False, phased_combine_fn=None, clock=None): """Create the TriggerDriver for the given windowing and options.""" # TODO(https://github.com/apache/beam/issues/20165): Respect closing and # on-time behaviors. For batch, we should always fire once, no matter what. if is_batch and windowing.triggerfn == _Never(): windowing = copy.copy(windowing) windowing.triggerfn = Always() # TODO(robertwb): We can do more if we know elements are in timestamp # sorted order. if windowing.is_default() and is_batch: driver = BatchGlobalTriggerDriver() elif (windowing.windowfn == GlobalWindows() and (windowing.triggerfn in [AfterCount(1), Always()]) and is_batch): # Here we also just pass through all the values exactly once. driver = BatchGlobalTriggerDriver() else: driver = GeneralTriggerDriver(windowing, clock) if phased_combine_fn: # TODO(ccy): Refactor GeneralTriggerDriver to combine values eagerly using # the known phased_combine_fn here. driver = CombiningTriggerDriver(phased_combine_fn, driver) return driver class TriggerDriver(metaclass=ABCMeta): """Breaks a series of bundle and timer firings into window (pane)s.""" @abstractmethod def process_elements( self, state, windowed_values, output_watermark, input_watermark=MIN_TIMESTAMP): pass @abstractmethod def process_timer( self, window_id, name, time_domain, timestamp, state, input_watermark=None): pass def process_entire_key(self, key, windowed_values): # This state holds per-key, multi-window state. state = InMemoryUnmergedState() for wvalue in self.process_elements(state, windowed_values, MIN_TIMESTAMP, MIN_TIMESTAMP): yield wvalue.with_value((key, wvalue.value)) while state.timers: fired = state.get_and_clear_timers() for timer_window, (name, time_domain, fire_time, _) in fired: for wvalue in self.process_timer(timer_window, name, time_domain, fire_time, state): yield wvalue.with_value((key, wvalue.value)) class _UnwindowedValues(observable.ObservableMixin): """Exposes iterable of windowed values as iterable of unwindowed values.""" def __init__(self, windowed_values): super().__init__() self._windowed_values = windowed_values def __iter__(self): for wv in self._windowed_values: unwindowed_value = wv.value self.notify_observers(unwindowed_value) yield unwindowed_value def __repr__(self): return '<_UnwindowedValues of %s>' % self._windowed_values def __reduce__(self): return list, (list(self), ) def __eq__(self, other): if isinstance(other, collections_abc.Iterable): return all( a == b for a, b in zip_longest(self, other, fillvalue=object())) else: return NotImplemented def __hash__(self): return hash(tuple(self)) coder_impl.FastPrimitivesCoderImpl.register_iterable_like_type( _UnwindowedValues) class BatchGlobalTriggerDriver(TriggerDriver): """Groups all received values together. """ GLOBAL_WINDOW_TUPLE = (GlobalWindow(), ) ONLY_FIRING = windowed_value.PaneInfo( is_first=True, is_last=True, timing=windowed_value.PaneInfoTiming.ON_TIME, index=0, nonspeculative_index=0) def process_elements( self, state, windowed_values, unused_output_watermark, unused_input_watermark=MIN_TIMESTAMP): yield WindowedValue( _UnwindowedValues(windowed_values), MIN_TIMESTAMP, self.GLOBAL_WINDOW_TUPLE, self.ONLY_FIRING) def process_timer( self, window_id, name, time_domain, timestamp, state, input_watermark=None): raise TypeError('Triggers never set or called for batch default windowing.') class CombiningTriggerDriver(TriggerDriver): """Uses a phased_combine_fn to process output of wrapped TriggerDriver.""" def __init__(self, phased_combine_fn, underlying): self.phased_combine_fn = phased_combine_fn self.underlying = underlying def process_elements( self, state, windowed_values, output_watermark, input_watermark=MIN_TIMESTAMP): uncombined = self.underlying.process_elements( state, windowed_values, output_watermark, input_watermark) for output in uncombined: yield output.with_value(self.phased_combine_fn.apply(output.value)) def process_timer( self, window_id, name, time_domain, timestamp, state, input_watermark=None): uncombined = self.underlying.process_timer( window_id, name, time_domain, timestamp, state, input_watermark) for output in uncombined: yield output.with_value(self.phased_combine_fn.apply(output.value)) class GeneralTriggerDriver(TriggerDriver): """Breaks a series of bundle and timer firings into window (pane)s. Suitable for all variants of Windowing. """ ELEMENTS = _ListStateTag('elements') TOMBSTONE = _CombiningValueStateTag('tombstone', combiners.CountCombineFn()) INDEX = _CombiningValueStateTag('index', combiners.CountCombineFn()) NONSPECULATIVE_INDEX = _CombiningValueStateTag( 'nonspeculative_index', combiners.CountCombineFn()) def __init__(self, windowing, clock): self.clock = clock self.allowed_lateness = windowing.allowed_lateness self.window_fn = windowing.windowfn self.timestamp_combiner_impl = TimestampCombiner.get_impl( windowing.timestamp_combiner, self.window_fn) # pylint: disable=invalid-name self.WATERMARK_HOLD = _WatermarkHoldStateTag( 'watermark', self.timestamp_combiner_impl) # pylint: enable=invalid-name self.trigger_fn = windowing.triggerfn self.accumulation_mode = windowing.accumulation_mode self.is_merging = True def process_elements( self, state, windowed_values, output_watermark, input_watermark=MIN_TIMESTAMP): if self.is_merging: state = MergeableStateAdapter(state) windows_to_elements = collections.defaultdict(list) for wv in windowed_values: for window in wv.windows: # ignore expired windows if input_watermark > window.end + self.allowed_lateness: continue windows_to_elements[window].append((wv.value, wv.timestamp)) # First handle merging. if self.is_merging: old_windows = set(state.known_windows()) all_windows = old_windows.union(list(windows_to_elements)) if all_windows != old_windows: merged_away = {} class TriggerMergeContext(WindowFn.MergeContext): def merge(_, to_be_merged, merge_result): # pylint: disable=no-self-argument for window in to_be_merged: if window != merge_result: merged_away[window] = merge_result # Clear state associated with PaneInfo since it is # not preserved across merges. state.clear_state(window, self.INDEX) state.clear_state(window, self.NONSPECULATIVE_INDEX) state.merge(to_be_merged, merge_result) # using the outer self argument. self.trigger_fn.on_merge( to_be_merged, merge_result, state.at(merge_result, self.clock)) self.window_fn.merge(TriggerMergeContext(all_windows)) merged_windows_to_elements = collections.defaultdict(list) for window, values in windows_to_elements.items(): while window in merged_away: window = merged_away[window] merged_windows_to_elements[window].extend(values) windows_to_elements = merged_windows_to_elements for window in merged_away: state.clear_state(window, self.WATERMARK_HOLD) # Next handle element adding. for window, elements in windows_to_elements.items(): if state.get_state(window, self.TOMBSTONE): continue # Add watermark hold. # TODO(ccy): Add late data and garbage-collection hold support. output_time = self.timestamp_combiner_impl.merge( window, ( element_output_time for element_output_time in ( self.timestamp_combiner_impl.assign_output_time( window, timestamp) for unused_value, timestamp in elements) if element_output_time >= output_watermark)) if output_time is not None: state.add_state(window, self.WATERMARK_HOLD, output_time) context = state.at(window, self.clock) for value, unused_timestamp in elements: state.add_state(window, self.ELEMENTS, value) self.trigger_fn.on_element(value, window, context) # Maybe fire this window. if self.trigger_fn.should_fire(TimeDomain.WATERMARK, input_watermark, window, context): finished = self.trigger_fn.on_fire(input_watermark, window, context) yield self._output(window, finished, state, output_watermark, False) def process_timer( self, window_id, unused_name, time_domain, timestamp, state, input_watermark=None): if input_watermark is None: input_watermark = timestamp if self.is_merging: state = MergeableStateAdapter(state) window = state.get_window(window_id) if state.get_state(window, self.TOMBSTONE): return if time_domain in (TimeDomain.WATERMARK, TimeDomain.REAL_TIME): if not self.is_merging or window in state.known_windows(): context = state.at(window, self.clock) if self.trigger_fn.should_fire(time_domain, timestamp, window, context): finished = self.trigger_fn.on_fire(timestamp, window, context) yield self._output( window, finished, state, timestamp, time_domain == TimeDomain.WATERMARK) else: raise Exception('Unexpected time domain: %s' % time_domain) def _output(self, window, finished, state, output_watermark, maybe_ontime): """Output window and clean up if appropriate.""" index = state.get_state(window, self.INDEX) state.add_state(window, self.INDEX, 1) if output_watermark <= window.max_timestamp(): nonspeculative_index = -1 timing = windowed_value.PaneInfoTiming.EARLY if state.get_state(window, self.NONSPECULATIVE_INDEX): nonspeculative_index = state.get_state( window, self.NONSPECULATIVE_INDEX) state.add_state(window, self.NONSPECULATIVE_INDEX, 1) _LOGGER.warning( 'Watermark moved backwards in time ' 'or late data moved window end forward.') else: nonspeculative_index = state.get_state(window, self.NONSPECULATIVE_INDEX) state.add_state(window, self.NONSPECULATIVE_INDEX, 1) timing = ( windowed_value.PaneInfoTiming.ON_TIME if maybe_ontime and nonspeculative_index == 0 else windowed_value.PaneInfoTiming.LATE) pane_info = windowed_value.PaneInfo( index == 0, finished, timing, index, nonspeculative_index) values = state.get_state(window, self.ELEMENTS) if finished: # TODO(robertwb): allowed lateness state.clear_state(window, self.ELEMENTS) state.add_state(window, self.TOMBSTONE, 1) elif self.accumulation_mode == AccumulationMode.DISCARDING: state.clear_state(window, self.ELEMENTS) timestamp = state.get_state(window, self.WATERMARK_HOLD) if timestamp is None: # If no watermark hold was set, output at end of window. timestamp = window.max_timestamp() elif output_watermark < window.end and self.trigger_fn.has_ontime_pane(): # Hold the watermark in case there is an empty pane that needs to be fired # at the end of the window. pass else: state.clear_state(window, self.WATERMARK_HOLD) return WindowedValue(values, timestamp, (window, ), pane_info) class InMemoryUnmergedState(UnmergedState): """In-memory implementation of UnmergedState. Used for batch and testing. """ def __init__(self, defensive_copy=False): # TODO(robertwb): Clean defensive_copy. It is too expensive in production. self.timers = collections.defaultdict(dict) self.state = collections.defaultdict(lambda: collections.defaultdict(list)) self.global_state = {} self.defensive_copy = defensive_copy def copy(self): cloned_object = InMemoryUnmergedState(defensive_copy=self.defensive_copy) cloned_object.timers = copy.deepcopy(self.timers) cloned_object.global_state = copy.deepcopy(self.global_state) for window in self.state: for tag in self.state[window]: cloned_object.state[window][tag] = copy.copy(self.state[window][tag]) return cloned_object def set_global_state(self, tag, value): assert isinstance(tag, _ReadModifyWriteStateTag) if self.defensive_copy: value = copy.deepcopy(value) self.global_state[tag.tag] = value def get_global_state(self, tag, default=None): return self.global_state.get(tag.tag, default) def set_timer( self, window, name, time_domain, timestamp, dynamic_timer_tag=''): self.timers[window][(name, time_domain, dynamic_timer_tag)] = timestamp def clear_timer(self, window, name, time_domain, dynamic_timer_tag=''): self.timers[window].pop((name, time_domain, dynamic_timer_tag), None) if not self.timers[window]: del self.timers[window] def get_window(self, window_id): return window_id def add_state(self, window, tag, value): if self.defensive_copy: value = copy.deepcopy(value) if isinstance(tag, _ReadModifyWriteStateTag): self.state[window][tag.tag] = value elif isinstance(tag, _CombiningValueStateTag): # TODO(robertwb): Store merged accumulators. self.state[window][tag.tag].append(value) elif isinstance(tag, _ListStateTag): self.state[window][tag.tag].append(value) elif isinstance(tag, _SetStateTag): self.state[window][tag.tag].append(value) elif isinstance(tag, _WatermarkHoldStateTag): self.state[window][tag.tag].append(value) else: raise ValueError('Invalid tag.', tag) def get_state(self, window, tag): values = self.state[window][tag.tag] if isinstance(tag, _ReadModifyWriteStateTag): return values elif isinstance(tag, _CombiningValueStateTag): return tag.combine_fn.apply(values) elif isinstance(tag, _ListStateTag): return values elif isinstance(tag, _SetStateTag): return values elif isinstance(tag, _WatermarkHoldStateTag): return tag.timestamp_combiner_impl.combine_all(values) else: raise ValueError('Invalid tag.', tag) def clear_state(self, window, tag): self.state[window].pop(tag.tag, None) if not self.state[window]: self.state.pop(window, None) def get_timers( self, clear=False, watermark=MAX_TIMESTAMP, processing_time=None): """Gets expired timers and reports if there are any realtime timers set per state. Expiration is measured against the watermark for event-time timers, and against a wall clock for processing-time timers. """ expired = [] has_realtime_timer = False for window, timers in list(self.timers.items()): for (name, time_domain, dynamic_timer_tag), timestamp in list( timers.items()): if time_domain == TimeDomain.REAL_TIME: time_marker = processing_time has_realtime_timer = True elif time_domain == TimeDomain.WATERMARK: time_marker = watermark else: _LOGGER.error( 'TimeDomain error: No timers defined for time domain %s.', time_domain) if timestamp <= time_marker: expired.append( (window, (name, time_domain, timestamp, dynamic_timer_tag))) if clear: del timers[(name, time_domain, dynamic_timer_tag)] if not timers and clear: del self.timers[window] return expired, has_realtime_timer def get_and_clear_timers(self, watermark=MAX_TIMESTAMP): return self.get_timers(clear=True, watermark=watermark)[0] def get_earliest_hold(self): earliest_hold = MAX_TIMESTAMP for unused_window, tagged_states in self.state.items(): # TODO(https://github.com/apache/beam/issues/18441): currently, this # assumes that the watermark hold tag is named "watermark". This is # currently only true because the only place watermark holds are set is # in the GeneralTriggerDriver, where we use this name. We should fix # this by allowing enumeration of the tag types used in adding state. if 'watermark' in tagged_states and tagged_states['watermark']: hold = min(tagged_states['watermark']) - TIME_GRANULARITY earliest_hold = min(earliest_hold, hold) return earliest_hold def __repr__(self): state_str = '\n'.join( '%s: %s' % (key, dict(state)) for key, state in self.state.items()) return 'timers: %s\nstate: %s' % (dict(self.timers), state_str) class _AfterSynchronizedProcessingTime(TriggerFn): """A "runner's-discretion" trigger downstream of a GroupByKey with AfterProcessingTime trigger. In runners that directly execute this Python code, the trigger currently always fires, but this behavior is neither guaranteed nor required by runners, regardless of whether they execute triggers via Python. _AfterSynchronizedProcessingTime is experimental and internal-only. No backwards compatibility guarantees. """ def __init__(self): pass def __repr__(self): return '_AfterSynchronizedProcessingTime()' def __eq__(self, other): return type(self) == type(other) def __hash__(self): return hash(type(self)) def on_element(self, _element, _window, _context): pass def on_merge(self, _to_be_merged, _merge_result, _context): pass def should_fire(self, _time_domain, _timestamp, _window, _context): return True def on_fire(self, _timestamp, _window, _context): return False def reset(self, _window, _context): pass @staticmethod def from_runner_api(_proto, _context): return _AfterSynchronizedProcessingTime() def to_runner_api(self, _context): return beam_runner_api_pb2.Trigger( after_synchronized_processing_time=beam_runner_api_pb2.Trigger. AfterSynchronizedProcessingTime()) def has_ontime_pane(self): return False def get_continuation_trigger(self): return self ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/trigger_test.py0000644000175100017510000014246115206113033023140 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the triggering classes.""" # pytype: skip-file import collections import json import os.path import pickle import random import unittest import yaml import apache_beam as beam from apache_beam import coders from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import TypeOptions from apache_beam.portability import common_urns from apache_beam.runners import pipeline_context from apache_beam.runners.direct.clock import TestClock from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import WindowInto from apache_beam.transforms import ptransform from apache_beam.transforms import trigger from apache_beam.transforms.core import Windowing from apache_beam.transforms.trigger import AccumulationMode from apache_beam.transforms.trigger import AfterAll from apache_beam.transforms.trigger import AfterAny from apache_beam.transforms.trigger import AfterCount from apache_beam.transforms.trigger import AfterEach from apache_beam.transforms.trigger import AfterProcessingTime from apache_beam.transforms.trigger import AfterWatermark from apache_beam.transforms.trigger import Always from apache_beam.transforms.trigger import DataLossReason from apache_beam.transforms.trigger import DefaultTrigger from apache_beam.transforms.trigger import GeneralTriggerDriver from apache_beam.transforms.trigger import InMemoryUnmergedState from apache_beam.transforms.trigger import Repeatedly from apache_beam.transforms.trigger import TriggerFn from apache_beam.transforms.trigger import _Never from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import GlobalWindows from apache_beam.transforms.window import IntervalWindow from apache_beam.transforms.window import Sessions from apache_beam.transforms.window import TimestampCombiner from apache_beam.transforms.window import TimestampedValue from apache_beam.transforms.window import WindowedValue from apache_beam.transforms.window import WindowFn from apache_beam.utils.timestamp import MAX_TIMESTAMP from apache_beam.utils.timestamp import MIN_TIMESTAMP from apache_beam.utils.timestamp import Duration from apache_beam.utils.windowed_value import PaneInfoTiming class CustomTimestampingFixedWindowsWindowFn(FixedWindows): """WindowFn for testing custom timestamping.""" def get_transformed_output_time(self, unused_window, input_timestamp): return input_timestamp + 100 class TriggerTest(unittest.TestCase): def run_trigger_simple( self, window_fn, trigger_fn, accumulation_mode, timestamped_data, expected_panes, *groupings, **kwargs): # Groupings is a list of integers indicating the (uniform) size of bundles # to try. For example, if timestamped_data has elements [a, b, c, d, e] # then groupings=(5, 2) would first run the test with everything in the same # bundle, and then re-run the test with bundling [a, b], [c, d], [e]. # A negative value will reverse the order, e.g. -2 would result in bundles # [e, d], [c, b], [a]. This is useful for deterministic triggers in testing # that the output is not a function of ordering or bundling. # If empty, defaults to bundles of size 1 in the given order. late_data = kwargs.pop('late_data', []) assert not kwargs def bundle_data(data, size): if size < 0: data = list(data)[::-1] size = -size bundle = [] for timestamp, elem in data: windows = window_fn.assign(WindowFn.AssignContext(timestamp, elem)) bundle.append(WindowedValue(elem, timestamp, windows)) if len(bundle) == size: yield bundle bundle = [] if bundle: yield bundle if not groupings: groupings = [1] for group_by in groupings: self.run_trigger( window_fn, trigger_fn, accumulation_mode, bundle_data(timestamped_data, group_by), bundle_data(late_data, group_by), expected_panes) def run_trigger( self, window_fn, trigger_fn, accumulation_mode, bundles, late_bundles, expected_panes): actual_panes = collections.defaultdict(list) allowed_lateness = Duration( micros=int(common_urns.constants.MAX_TIMESTAMP_MILLIS.constant) * 1000) driver = GeneralTriggerDriver( Windowing( window_fn, trigger_fn, accumulation_mode, allowed_lateness=allowed_lateness), TestClock()) state = InMemoryUnmergedState() for bundle in bundles: for wvalue in driver.process_elements(state, bundle, MIN_TIMESTAMP, MIN_TIMESTAMP): window, = wvalue.windows self.assertEqual(window.max_timestamp(), wvalue.timestamp) actual_panes[window].append(set(wvalue.value)) while state.timers: for timer_window, (name, time_domain, timestamp, _) in state.get_and_clear_timers(): for wvalue in driver.process_timer(timer_window, name, time_domain, timestamp, state, MIN_TIMESTAMP): window, = wvalue.windows self.assertEqual(window.max_timestamp(), wvalue.timestamp) actual_panes[window].append(set(wvalue.value)) for bundle in late_bundles: for wvalue in driver.process_elements(state, bundle, MAX_TIMESTAMP, MAX_TIMESTAMP): window, = wvalue.windows self.assertEqual(window.max_timestamp(), wvalue.timestamp) actual_panes[window].append(set(wvalue.value)) while state.timers: for timer_window, (name, time_domain, timestamp, _) in state.get_and_clear_timers(): for wvalue in driver.process_timer(timer_window, name, time_domain, timestamp, state, MAX_TIMESTAMP): window, = wvalue.windows self.assertEqual(window.max_timestamp(), wvalue.timestamp) actual_panes[window].append(set(wvalue.value)) self.assertEqual(expected_panes, actual_panes) def test_fixed_watermark(self): self.run_trigger_simple( FixedWindows(10), # pyformat break AfterWatermark(), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (13, 'c')], { IntervalWindow(0, 10): [set('ab')], IntervalWindow(10, 20): [set('c')] }, 1, 2, 3, -3, -2, -1) def test_fixed_watermark_with_early(self): self.run_trigger_simple( FixedWindows(10), # pyformat break AfterWatermark(early=AfterCount(2)), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c')], {IntervalWindow(0, 10): [set('ab'), set('abc')]}, 2) self.run_trigger_simple( FixedWindows(10), # pyformat break AfterWatermark(early=AfterCount(2)), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c')], {IntervalWindow(0, 10): [set('abc'), set('abc')]}, 3) def test_fixed_watermark_with_early_late(self): self.run_trigger_simple( FixedWindows(100), # pyformat break AfterWatermark(early=AfterCount(3), late=AfterCount(2)), AccumulationMode.DISCARDING, zip(range(9), 'abcdefghi'), { IntervalWindow(0, 100): [ set('abcd'), set('efgh'), # early set('i'), # on time set('vw'), set('xy') # late ] }, 2, late_data=zip(range(5), 'vwxyz')) def test_sessions_watermark_with_early_late(self): self.run_trigger_simple( Sessions(10), # pyformat break AfterWatermark(early=AfterCount(2), late=AfterCount(1)), AccumulationMode.ACCUMULATING, [(1, 'a'), (15, 'b'), (7, 'c'), (30, 'd')], { IntervalWindow(1, 25): [ set('abc'), # early set('abc'), # on time set('abcxy') # late ], IntervalWindow(30, 40): [ set('d'), # on time ], IntervalWindow(1, 40): [ set('abcdxyz') # late ], }, 2, late_data=[(1, 'x'), (2, 'y'), (21, 'z')]) def test_fixed_after_count(self): self.run_trigger_simple( FixedWindows(10), # pyformat break AfterCount(2), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c'), (11, 'z')], {IntervalWindow(0, 10): [set('ab')]}, 1, 2) self.run_trigger_simple( FixedWindows(10), # pyformat break AfterCount(2), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c'), (11, 'z')], {IntervalWindow(0, 10): [set('abc')]}, 3, 4) def test_fixed_after_first(self): self.run_trigger_simple( FixedWindows(10), # pyformat break AfterAny(AfterCount(2), AfterWatermark()), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c')], {IntervalWindow(0, 10): [set('ab')]}, 1, 2) self.run_trigger_simple( FixedWindows(10), # pyformat break AfterAny(AfterCount(5), AfterWatermark()), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c')], {IntervalWindow(0, 10): [set('abc')]}, 1, 2, late_data=[(1, 'x'), (2, 'y'), (3, 'z')]) def test_repeatedly_after_first(self): self.run_trigger_simple( FixedWindows(100), # pyformat break Repeatedly(AfterAny(AfterCount(3), AfterWatermark())), AccumulationMode.ACCUMULATING, zip(range(7), 'abcdefg'), { IntervalWindow(0, 100): [ set('abc'), set('abcdef'), set('abcdefg'), set('abcdefgx'), set('abcdefgxy'), set('abcdefgxyz') ] }, 1, late_data=zip(range(3), 'xyz')) def test_sessions_after_all(self): self.run_trigger_simple( Sessions(10), # pyformat break AfterAll(AfterCount(2), AfterWatermark()), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c')], {IntervalWindow(1, 13): [set('abc')]}, 1, 2) self.run_trigger_simple( Sessions(10), # pyformat break AfterAll(AfterCount(5), AfterWatermark()), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (3, 'c')], {IntervalWindow(1, 13): [set('abcxy')]}, 1, 2, late_data=[(1, 'x'), (2, 'y'), (3, 'z')]) def test_sessions_default(self): self.run_trigger_simple( Sessions(10), # pyformat break DefaultTrigger(), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b')], {IntervalWindow(1, 12): [set('ab')]}, 1, 2, -2, -1) self.run_trigger_simple( Sessions(10), # pyformat break AfterWatermark(), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b'), (15, 'c'), (16, 'd'), (30, 'z'), (9, 'e'), (10, 'f'), (30, 'y')], { IntervalWindow(1, 26): [set('abcdef')], IntervalWindow(30, 40): [set('yz')] }, 1, 2, 3, 4, 5, 6, -4, -2, -1) def test_sessions_watermark(self): self.run_trigger_simple( Sessions(10), # pyformat break AfterWatermark(), AccumulationMode.ACCUMULATING, [(1, 'a'), (2, 'b')], {IntervalWindow(1, 12): [set('ab')]}, 1, 2, -2, -1) def test_sessions_after_count(self): self.run_trigger_simple( Sessions(10), # pyformat break AfterCount(2), AccumulationMode.ACCUMULATING, [(1, 'a'), (15, 'b'), (6, 'c'), (30, 's'), (31, 't'), (50, 'z'), (50, 'y')], { IntervalWindow(1, 25): [set('abc')], IntervalWindow(30, 41): [set('st')], IntervalWindow(50, 60): [set('yz')] }, 1, 2, 3) def test_sessions_repeatedly_after_count(self): self.run_trigger_simple( Sessions(10), # pyformat break Repeatedly(AfterCount(2)), AccumulationMode.ACCUMULATING, [(1, 'a'), (15, 'b'), (6, 'c'), (2, 'd'), (7, 'e')], {IntervalWindow(1, 25): [set('abc'), set('abcde')]}, 1, 3) self.run_trigger_simple( Sessions(10), # pyformat break Repeatedly(AfterCount(2)), AccumulationMode.DISCARDING, [(1, 'a'), (15, 'b'), (6, 'c'), (2, 'd'), (7, 'e')], {IntervalWindow(1, 25): [set('abc'), set('de')]}, 1, 3) def test_sessions_after_each(self): self.run_trigger_simple( Sessions(10), # pyformat break AfterEach(AfterCount(2), AfterCount(3)), AccumulationMode.ACCUMULATING, zip(range(10), 'abcdefghij'), { IntervalWindow(0, 11): [set('ab')], IntervalWindow(0, 15): [set('abcdef')] }, 2) self.run_trigger_simple( Sessions(10), # pyformat break Repeatedly(AfterEach(AfterCount(2), AfterCount(3))), AccumulationMode.ACCUMULATING, zip(range(10), 'abcdefghij'), { IntervalWindow(0, 11): [set('ab')], IntervalWindow(0, 15): [set('abcdef')], IntervalWindow(0, 17): [set('abcdefgh')] }, 2) def test_picklable_output(self): global_window = (trigger.GlobalWindow(), ) driver = trigger.BatchGlobalTriggerDriver() unpicklable = (WindowedValue(k, 0, global_window) for k in range(10)) with self.assertRaises(TypeError): pickle.dumps(unpicklable) for unwindowed in driver.process_elements(None, unpicklable, None, None): self.assertEqual( pickle.loads(pickle.dumps(unwindowed)).value, list(range(10))) class MayLoseDataTest(unittest.TestCase): def _test(self, trigger, lateness, expected): windowing = WindowInto( GlobalWindows(), trigger=trigger, accumulation_mode=AccumulationMode.ACCUMULATING, allowed_lateness=lateness).windowing self.assertEqual(trigger.may_lose_data(windowing), expected) def test_default_trigger(self): self._test(DefaultTrigger(), 0, DataLossReason.NO_POTENTIAL_LOSS) def test_after_processing(self): self._test(AfterProcessingTime(42), 0, DataLossReason.MAY_FINISH) def test_always(self): self._test(Always(), 0, DataLossReason.NO_POTENTIAL_LOSS) def test_never(self): self._test(_Never(), 0, DataLossReason.NO_POTENTIAL_LOSS) def test_after_watermark_no_allowed_lateness(self): self._test(AfterWatermark(), 0, DataLossReason.NO_POTENTIAL_LOSS) def test_after_watermark_no_late_trigger(self): self._test(AfterWatermark(), 60, DataLossReason.MAY_FINISH) def test_after_watermark_no_allowed_lateness_safe_late(self): self._test( AfterWatermark(late=DefaultTrigger()), 0, DataLossReason.NO_POTENTIAL_LOSS) def test_after_watermark_allowed_lateness_safe_late(self): self._test( AfterWatermark(late=DefaultTrigger()), 60, DataLossReason.NO_POTENTIAL_LOSS) def test_after_count(self): self._test(AfterCount(42), 0, DataLossReason.MAY_FINISH) def test_repeatedly_safe_underlying(self): self._test( Repeatedly(DefaultTrigger()), 0, DataLossReason.NO_POTENTIAL_LOSS) def test_repeatedly_unsafe_underlying(self): self._test(Repeatedly(AfterCount(42)), 0, DataLossReason.NO_POTENTIAL_LOSS) def test_after_any_one_may_finish(self): self._test( AfterAny(AfterCount(42), DefaultTrigger()), 0, DataLossReason.MAY_FINISH) def test_after_any_all_safe(self): self._test( AfterAny(Repeatedly(AfterCount(42)), DefaultTrigger()), 0, DataLossReason.NO_POTENTIAL_LOSS) def test_after_all_some_may_finish(self): self._test( AfterAll(AfterCount(1), DefaultTrigger()), 0, DataLossReason.NO_POTENTIAL_LOSS) def test_afer_all_all_may_finish(self): self._test( AfterAll(AfterCount(42), AfterProcessingTime(42)), 0, DataLossReason.MAY_FINISH) def test_after_each_at_least_one_safe(self): self._test( AfterEach(AfterCount(1), DefaultTrigger(), AfterCount(2)), 0, DataLossReason.NO_POTENTIAL_LOSS) def test_after_each_all_may_finish(self): self._test( AfterEach(AfterCount(1), AfterCount(2), AfterCount(3)), 0, DataLossReason.MAY_FINISH) class RunnerApiTest(unittest.TestCase): def test_trigger_encoding(self): for trigger_fn in (DefaultTrigger(), AfterAll(AfterCount(1), AfterCount(10)), AfterAny(AfterCount(10), AfterCount(100)), AfterWatermark(early=AfterCount(1000)), AfterWatermark(early=AfterCount(1000), late=AfterCount(1)), Repeatedly(AfterCount(100)), trigger.OrFinally(AfterCount(3), AfterCount(10))): context = pipeline_context.PipelineContext() self.assertEqual( trigger_fn, TriggerFn.from_runner_api(trigger_fn.to_runner_api(context), context)) class ContinuationTriggerTest(unittest.TestCase): def test_after_all(self): self.assertEqual( AfterAll(AfterCount(2), AfterCount(5)).get_continuation_trigger(), AfterAll(AfterCount(1), AfterCount(1))) def test_after_any(self): self.assertEqual( AfterAny(AfterCount(2), AfterCount(5)).get_continuation_trigger(), AfterAny(AfterCount(1), AfterCount(1))) def test_after_count(self): self.assertEqual(AfterCount(1).get_continuation_trigger(), AfterCount(1)) self.assertEqual(AfterCount(100).get_continuation_trigger(), AfterCount(1)) def test_after_each(self): self.assertEqual( AfterEach(AfterCount(2), AfterCount(5)).get_continuation_trigger(), Repeatedly(AfterAny(AfterCount(1), AfterCount(1)))) def test_after_processing_time(self): from apache_beam.transforms.trigger import _AfterSynchronizedProcessingTime self.assertEqual( AfterProcessingTime(10).get_continuation_trigger(), _AfterSynchronizedProcessingTime()) def test_after_watermark(self): self.assertEqual( AfterWatermark().get_continuation_trigger(), AfterWatermark()) self.assertEqual( AfterWatermark(early=AfterCount(10), late=AfterCount(20)).get_continuation_trigger(), AfterWatermark(early=AfterCount(1), late=AfterCount(1))) def test_always(self): self.assertEqual(Always().get_continuation_trigger(), Always()) def test_default(self): self.assertEqual( DefaultTrigger().get_continuation_trigger(), DefaultTrigger()) def test_never(self): self.assertEqual(_Never().get_continuation_trigger(), _Never()) def test_repeatedly(self): self.assertEqual( Repeatedly(AfterCount(10)).get_continuation_trigger(), Repeatedly(AfterCount(1))) class TriggerPipelineTest(unittest.TestCase): def test_after_processing_time(self): test_options = PipelineOptions( flags=['--allow_unsafe_triggers', '--streaming']) with TestPipeline(options=test_options) as p: total_elements_in_trigger = 4 processing_time_delay = 2 window_size = 10 # yapf: disable test_stream = TestStream() for i in range(total_elements_in_trigger): (test_stream .advance_processing_time( processing_time_delay / total_elements_in_trigger) .add_elements([('key', i)]) ) test_stream.advance_processing_time(processing_time_delay) # Add dropped elements (test_stream .advance_processing_time(0.1) .add_elements([('key', "dropped-1")]) .advance_processing_time(0.1) .add_elements([('key', "dropped-2")]) ) (test_stream .advance_processing_time(processing_time_delay) .advance_watermark_to_infinity() ) # yapf: enable results = ( p | test_stream | beam.WindowInto( FixedWindows(window_size), trigger=AfterProcessingTime(processing_time_delay), accumulation_mode=AccumulationMode.DISCARDING) | beam.GroupByKey() | beam.Map(lambda x: x[1])) assert_that(results, equal_to([list(range(total_elements_in_trigger))])) def test_repeatedly_after_processing_time(self): test_options = PipelineOptions(flags=['--streaming']) with TestPipeline(options=test_options) as p: total_elements = 7 processing_time_delay = 2 window_size = 10 # yapf: disable test_stream = TestStream() for i in range(total_elements): (test_stream .advance_processing_time(processing_time_delay - 0.01) .add_elements([('key', i)]) ) (test_stream .advance_processing_time(processing_time_delay) .advance_watermark_to_infinity() ) # yapf: enable results = ( p | test_stream | beam.WindowInto( FixedWindows(window_size), trigger=Repeatedly(AfterProcessingTime(processing_time_delay)), accumulation_mode=AccumulationMode.DISCARDING) | beam.GroupByKey() | beam.Map(lambda x: x[1])) expected = [[i, i + 1] for i in range(total_elements - total_elements % 2) if i % 2 == 0] expected += [] if total_elements % 2 == 0 else [[total_elements - 1]] assert_that(results, equal_to(expected)) def test_after_count(self): test_options = PipelineOptions(flags=['--allow_unsafe_triggers']) with TestPipeline(options=test_options) as p: def construct_timestamped(k, t): return TimestampedValue((k, t), t) def format_result(k, vs): return ('%s-%s' % (k, len(list(vs))), set(vs)) result = ( p | beam.Create([1, 2, 3, 4, 5, 10, 11]) | beam.FlatMap(lambda t: [('A', t), ('B', t + 5)]) | beam.MapTuple(construct_timestamped) | beam.WindowInto( FixedWindows(10), trigger=AfterCount(3), accumulation_mode=AccumulationMode.DISCARDING) | beam.GroupByKey() | beam.MapTuple(format_result)) assert_that( result, equal_to( list({ 'A-5': {1, 2, 3, 4, 5}, # A-10, A-11 never emitted due to AfterCount(3) never firing. 'B-4': {6, 7, 8, 9}, 'B-3': {10, 15, 16}, }.items()))) def test_after_count_streaming(self): test_options = PipelineOptions( flags=['--allow_unsafe_triggers', '--streaming']) with TestPipeline(options=test_options) as p: # yapf: disable test_stream = ( TestStream() .advance_watermark_to(0) .add_elements([('A', 1), ('A', 2), ('A', 3)]) .add_elements([('A', 4), ('A', 5), ('A', 6)]) .add_elements([('B', 1), ('B', 2), ('B', 3)]) .advance_watermark_to_infinity()) # yapf: enable results = ( p | test_stream | beam.WindowInto( FixedWindows(10), trigger=AfterCount(3), accumulation_mode=AccumulationMode.ACCUMULATING) | beam.GroupByKey()) assert_that( results, equal_to( list({ 'A': [1, 2, 3], # 4 - 6 discarded because trigger finished 'B': [1, 2, 3] }.items()))) def test_always(self): # Pin to FnApiRunner since portable runner could trigger differently if # using bundle sizes of greater than 1. with TestPipeline('FnApiRunner') as p: def construct_timestamped(k, t): return TimestampedValue((k, t), t) def format_result(k, vs): return ('%s-%s' % (k, len(list(vs))), set(vs)) result = ( p | beam.Create([1, 1, 2, 3, 4, 5, 10, 11]) | beam.FlatMap(lambda t: [('A', t), ('B', t + 5)]) | beam.MapTuple(construct_timestamped) | beam.WindowInto( FixedWindows(10), trigger=Always(), accumulation_mode=AccumulationMode.DISCARDING) | beam.GroupByKey() | beam.MapTuple(format_result)) assert_that( result, equal_to( list({ 'A-2': {10, 11}, # Elements out of windows are also emitted. 'A-6': {1, 2, 3, 4, 5}, # A,1 is emitted twice. 'B-5': {6, 7, 8, 9}, # B,6 is emitted twice. 'B-3': {10, 15, 16}, }.items()))) def test_never(self): with TestPipeline() as p: def construct_timestamped(k, t): return TimestampedValue((k, t), t) def format_result(k, vs): return ('%s-%s' % (k, len(list(vs))), set(vs)) result = ( p | beam.Create([1, 1, 2, 3, 4, 5, 10, 11]) | beam.FlatMap(lambda t: [('A', t), ('B', t + 5)]) | beam.MapTuple(construct_timestamped) | beam.WindowInto( FixedWindows(10), trigger=_Never(), accumulation_mode=AccumulationMode.DISCARDING) | beam.GroupByKey() | beam.MapTuple(format_result)) assert_that( result, equal_to( list({ 'A-2': {10, 11}, 'A-6': {1, 2, 3, 4, 5}, 'B-5': {6, 7, 8, 9}, 'B-3': {10, 15, 16}, }.items()))) def test_multiple_accumulating_firings(self): # PCollection will contain elements from 1 to 10. elements = [i for i in range(1, 11)] ts = TestStream().advance_watermark_to(0) for i in elements: ts.add_elements([('key', str(i))]) if i % 5 == 0: ts.advance_watermark_to(i) ts.advance_processing_time(5) ts.advance_watermark_to_infinity() options = PipelineOptions() options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as p: records = ( p | ts | beam.WindowInto( FixedWindows(10), accumulation_mode=trigger.AccumulationMode.ACCUMULATING, trigger=AfterWatermark( early=AfterAll(AfterCount(1), AfterProcessingTime(5)))) | beam.GroupByKey() | beam.FlatMap(lambda x: x[1])) # The trigger should fire twice. Once after 5 seconds, and once after 10. # The firings should accumulate the output. first_firing = [str(i) for i in elements if i <= 5] second_firing = [str(i) for i in elements] assert_that(records, equal_to(first_firing + second_firing)) def test_on_pane_watermark_hold_no_pipeline_stall(self): """A regression test added for ttps://issues.apache.org/jira/browse/BEAM-10054.""" START_TIMESTAMP = 1534842000 test_stream = TestStream() test_stream.add_elements(['a']) test_stream.advance_processing_time(START_TIMESTAMP + 1) test_stream.advance_watermark_to(START_TIMESTAMP + 1) test_stream.add_elements(['b']) test_stream.advance_processing_time(START_TIMESTAMP + 2) test_stream.advance_watermark_to(START_TIMESTAMP + 2) with TestPipeline(options=PipelineOptions( ['--streaming', '--allow_unsafe_triggers'])) as p: # pylint: disable=expression-not-assigned ( p | 'TestStream' >> test_stream | 'timestamp' >> beam.Map(lambda x: beam.window.TimestampedValue(x, START_TIMESTAMP)) | 'kv' >> beam.Map(lambda x: (x, x)) | 'window_1m' >> beam.WindowInto( beam.window.FixedWindows(60), trigger=trigger.AfterAny( trigger.AfterProcessingTime(3600), trigger.AfterWatermark()), accumulation_mode=trigger.AccumulationMode.DISCARDING) | 'group_by_key' >> beam.GroupByKey() | 'filter' >> beam.Map(lambda x: x)) class TranscriptTest(unittest.TestCase): # We must prepend an underscore to this name so that the open-source unittest # runner does not execute this method directly as a test. @classmethod def _create_test(cls, spec): counter = 0 name = spec.get('name', 'unnamed') unique_name = 'test_' + name while hasattr(cls, unique_name): counter += 1 unique_name = 'test_%s_%d' % (name, counter) test_method = lambda self: self._run_log_test(spec) test_method.__name__ = unique_name test_method.__test__ = True setattr(cls, unique_name, test_method) # We must prepend an underscore to this name so that the open-source unittest # runner does not execute this method directly as a test. @classmethod def _create_tests(cls, transcript_filename): for spec in yaml.load_all(open(transcript_filename), Loader=yaml.SafeLoader): cls._create_test(spec) def _run_log_test(self, spec): if 'error' in spec: self.assertRaisesRegex(Exception, spec['error'], self._run_log, spec) else: self._run_log(spec) def _run_log(self, spec): def parse_int_list(s): """Parses strings like '[1, 2, 3]'.""" s = s.strip() assert s[0] == '[' and s[-1] == ']', s if not s[1:-1].strip(): return [] return [int(x) for x in s[1:-1].split(',')] def split_args(s): """Splits 'a, b, [c, d]' into ['a', 'b', '[c, d]'].""" args = [] start = 0 depth = 0 for ix in range(len(s)): c = s[ix] if c in '({[': depth += 1 elif c in ')}]': depth -= 1 elif c == ',' and depth == 0: args.append(s[start:ix].strip()) start = ix + 1 assert depth == 0, s args.append(s[start:].strip()) return args def parse(s, names): """Parse (recursive) 'Foo(arg, kw=arg)' for Foo in the names dict.""" s = s.strip() if s in names: return names[s] elif s[0] == '[': return parse_int_list(s) elif '(' in s: assert s[-1] == ')', s callee = parse(s[:s.index('(')], names) posargs = [] kwargs = {} for arg in split_args(s[s.index('(') + 1:-1]): if '=' in arg: kw, value = arg.split('=', 1) kwargs[kw] = parse(value, names) else: posargs.append(parse(arg, names)) return callee(*posargs, **kwargs) else: try: return int(s) except ValueError: raise ValueError('Unknown function: %s' % s) def parse_fn(s, names): """Like parse(), but implicitly calls no-arg constructors.""" fn = parse(s, names) if isinstance(fn, type): return fn() return fn # pylint: disable=wrong-import-order, wrong-import-position from apache_beam.transforms import window as window_module # pylint: enable=wrong-import-order, wrong-import-position window_fn_names = dict(window_module.__dict__) # yapf: disable window_fn_names.update({ 'CustomTimestampingFixedWindowsWindowFn': CustomTimestampingFixedWindowsWindowFn }) # yapf: enable trigger_names = {'Default': DefaultTrigger} trigger_names.update(trigger.__dict__) window_fn = parse_fn( spec.get('window_fn', 'GlobalWindows'), window_fn_names) trigger_fn = parse_fn(spec.get('trigger_fn', 'Default'), trigger_names) accumulation_mode = getattr( AccumulationMode, spec.get('accumulation_mode', 'ACCUMULATING').upper()) timestamp_combiner = getattr( TimestampCombiner, spec.get('timestamp_combiner', 'OUTPUT_AT_EOW').upper()) allowed_lateness = spec.get('allowed_lateness', 0.000) def only_element(xs): x, = list(xs) return x transcript = [only_element(line.items()) for line in spec['transcript']] self._execute( window_fn, trigger_fn, accumulation_mode, timestamp_combiner, allowed_lateness, transcript, spec) def _windowed_value_info(windowed_value): # Currently some runners operate at the millisecond level, and some at the # microsecond level. Trigger transcript timestamps are expressed as # integral units of the finest granularity, whatever that may be. # In these tests we interpret them as integral seconds and then truncate # the results to integral seconds to allow for portability across # different sub-second resolutions. window, = windowed_value.windows return { 'window': [int(window.start), int(window.max_timestamp())], 'values': sorted(windowed_value.value), 'timestamp': int(windowed_value.timestamp), 'index': windowed_value.pane_info.index, 'nonspeculative_index': windowed_value.pane_info.nonspeculative_index, 'early': windowed_value.pane_info.timing == PaneInfoTiming.EARLY, 'late': windowed_value.pane_info.timing == PaneInfoTiming.LATE, 'final': windowed_value.pane_info.is_last, } def _windowed_value_info_map_fn( k, vs, window=beam.DoFn.WindowParam, t=beam.DoFn.TimestampParam, p=beam.DoFn.PaneInfoParam): return ( k, _windowed_value_info( WindowedValue(vs, windows=[window], timestamp=t, pane_info=p))) def _windowed_value_info_check(actual, expected, key=None): key_string = ' for %s' % key if key else '' def format(panes): return '\n[%s]\n' % '\n '.join( str(pane) for pane in sorted(panes, key=lambda pane: pane.get('timestamp', None))) if len(actual) > len(expected): raise AssertionError( 'Unexpected output%s: expected %s but got %s' % (key_string, format(expected), format(actual))) elif len(expected) > len(actual): raise AssertionError( 'Unmatched output%s: expected %s but got %s' % (key_string, format(expected), format(actual))) else: def diff(actual, expected): for key in sorted(expected.keys(), reverse=True): if key in actual: if actual[key] != expected[key]: return key for output in actual: diffs = [diff(output, pane) for pane in expected] if all(diffs): raise AssertionError( 'Unmatched output%s: %s not found in %s (diffs in %s)' % (key_string, output, format(expected), diffs)) class _ConcatCombineFn(beam.CombineFn): create_accumulator = lambda self: [] # type: ignore[var-annotated] add_input = lambda self, acc, element: acc.append(element) or acc merge_accumulators = lambda self, accs: sum(accs, []) # type: ignore[var-annotated] extract_output = lambda self, acc: acc class TriggerDriverTranscriptTest(TranscriptTest): def _execute( self, window_fn, trigger_fn, accumulation_mode, timestamp_combiner, allowed_lateness, transcript, unused_spec): driver = GeneralTriggerDriver( Windowing( window_fn, trigger_fn, accumulation_mode, timestamp_combiner, allowed_lateness), TestClock()) state = InMemoryUnmergedState() output = [] watermark = MIN_TIMESTAMP def fire_timers(): to_fire = state.get_and_clear_timers(watermark) while to_fire: for timer_window, (name, time_domain, t_timestamp, _) in to_fire: for wvalue in driver.process_timer(timer_window, name, time_domain, t_timestamp, state): output.append(_windowed_value_info(wvalue)) to_fire = state.get_and_clear_timers(watermark) for action, params in transcript: if action != 'expect': # Fail if we have output that was not expected in the transcript. self.assertEqual([], output, msg='Unexpected output: %s before %s: %s' % (output, action, params)) if action == 'input': bundle = [ WindowedValue(t, t, window_fn.assign(WindowFn.AssignContext(t, t))) for t in params ] output = [ _windowed_value_info(wv) for wv in driver.process_elements( state, bundle, watermark, watermark) ] fire_timers() elif action == 'watermark': watermark = params fire_timers() elif action == 'expect': for expected_output in params: for candidate in output: if all(candidate[k] == expected_output[k] for k in candidate if k in expected_output): output.remove(candidate) break else: self.fail('Unmatched output %s in %s' % (expected_output, output)) elif action == 'state': # TODO(robertwb): Implement once we support allowed lateness. pass else: self.fail('Unknown action: ' + action) # Fail if we have output that was not expected in the transcript. self.assertEqual([], output, msg='Unexpected output: %s' % output) class BaseTestStreamTranscriptTest(TranscriptTest): """A suite of TestStream-based tests based on trigger transcript entries. """ def _execute( self, window_fn, trigger_fn, accumulation_mode, timestamp_combiner, allowed_lateness, transcript, spec): runner_name = TestPipeline().runner.__class__.__name__ if runner_name in spec.get('broken_on', ()): self.skipTest('Known to be broken on %s' % runner_name) is_order_agnostic = ( isinstance(trigger_fn, DefaultTrigger) and accumulation_mode == AccumulationMode.ACCUMULATING) if is_order_agnostic: reshuffle_seed = random.randrange(1 << 20) keys = [ 'original', 'reversed', 'reshuffled(%s)' % reshuffle_seed, 'one-element-bundles', 'one-element-bundles-reversed', 'two-element-bundles' ] else: keys = ['key1', 'key2'] # Elements are encoded as a json strings to allow other languages to # decode elements while executing the test stream. # TODO(https://github.com/apache/beam/issues/19934): Eliminate these # gymnastics. test_stream = TestStream(coder=coders.StrUtf8Coder()).with_output_types(str) for action, params in transcript: if action == 'expect': test_stream.add_elements([json.dumps(('expect', params))]) else: test_stream.add_elements([json.dumps(('expect', []))]) if action == 'input': def keyed(key, values): return [json.dumps(('input', (key, v))) for v in values] if is_order_agnostic: # Must match keys above. test_stream.add_elements(keyed('original', params)) test_stream.add_elements(keyed('reversed', reversed(params))) r = random.Random(reshuffle_seed) reshuffled = list(params) r.shuffle(reshuffled) test_stream.add_elements( keyed('reshuffled(%s)' % reshuffle_seed, reshuffled)) for v in params: test_stream.add_elements(keyed('one-element-bundles', [v])) for v in reversed(params): test_stream.add_elements( keyed('one-element-bundles-reversed', [v])) for ix in range(0, len(params), 2): test_stream.add_elements( keyed('two-element-bundles', params[ix:ix + 2])) else: for key in keys: test_stream.add_elements(keyed(key, params)) elif action == 'watermark': test_stream.advance_watermark_to(params) elif action == 'clock': test_stream.advance_processing_time(params) elif action == 'state': pass # Requires inspection of implementation details. else: raise ValueError('Unexpected action: %s' % action) test_stream.add_elements([json.dumps(('expect', []))]) test_stream.advance_watermark_to_infinity() read_test_stream = test_stream | beam.Map(json.loads) class Check(beam.DoFn): """A StatefulDoFn that verifies outputs are produced as expected. This DoFn takes in two kinds of inputs, actual outputs and expected outputs. When an actual output is received, it is buffered into state, and when an expected output is received, this buffered state is retrieved and compared against the expected value(s) to ensure they match. The key is ignored, but all items must be on the same key to share state. """ def __init__(self, allow_out_of_order=True): # Some runners don't support cross-stage TestStream semantics. self.allow_out_of_order = allow_out_of_order def process( self, element, seen=beam.DoFn.StateParam( beam.transforms.userstate.BagStateSpec( 'seen', beam.coders.FastPrimitivesCoder())), expected=beam.DoFn.StateParam( beam.transforms.userstate.BagStateSpec( 'expected', beam.coders.FastPrimitivesCoder()))): key, (action, data) = element if self.allow_out_of_order: if action == 'expect' and not list(seen.read()): if data: expected.add(data) return elif action == 'actual' and list(expected.read()): seen.add(data) all_data = list(seen.read()) all_expected = list(expected.read()) if len(all_data) == len(all_expected[0]): expected.clear() for expect in all_expected[1:]: expected.add(expect) action, data = 'expect', all_expected[0] else: return if action == 'actual': seen.add(data) elif action == 'expect': actual = list(seen.read()) seen.clear() _windowed_value_info_check(actual, data, key) else: raise ValueError('Unexpected action: %s' % action) @ptransform.ptransform_fn def CheckAggregation(inputs_and_expected, aggregation): # Split the test stream into a branch of to-be-processed elements, and # a branch of expected results. inputs, expected = ( inputs_and_expected | beam.MapTuple( lambda tag, value: beam.pvalue.TaggedOutput(tag, value), ).with_outputs('input', 'expect')) # Process the inputs with the given windowing to produce actual outputs. outputs = ( inputs | beam.MapTuple( lambda key, value: TimestampedValue((key, value), value)) | beam.WindowInto( window_fn, trigger=trigger_fn, accumulation_mode=accumulation_mode, timestamp_combiner=timestamp_combiner, allowed_lateness=allowed_lateness) | aggregation | beam.MapTuple(_windowed_value_info_map_fn) # Place outputs back into the global window to allow flattening # and share a single state in Check. | 'Global' >> beam.WindowInto(beam.transforms.window.GlobalWindows())) # Feed both the expected and actual outputs to Check() for comparison. tagged_expected = ( expected | beam.FlatMap( lambda value: [(key, ('expect', value)) for key in keys])) tagged_outputs = ( outputs | beam.MapTuple(lambda key, value: (key, ('actual', value)))) # pylint: disable=expression-not-assigned ([tagged_expected, tagged_outputs] | beam.Flatten() | beam.ParDo(Check(self.allow_out_of_order))) with TestPipeline() as p: # TODO(https://github.com/apache/beam/issues/19933): Pass this during # pipeline construction. p._options.view_as(StandardOptions).streaming = True p._options.view_as(TypeOptions).allow_unsafe_triggers = True # We can have at most one test stream per pipeline, so we share it. inputs_and_expected = p | read_test_stream _ = inputs_and_expected | CheckAggregation(beam.GroupByKey()) _ = inputs_and_expected | CheckAggregation( beam.CombinePerKey(_ConcatCombineFn())) class TestStreamTranscriptTest(BaseTestStreamTranscriptTest): allow_out_of_order = False class WeakTestStreamTranscriptTest(BaseTestStreamTranscriptTest): allow_out_of_order = True class BatchTranscriptTest(TranscriptTest): def _execute( self, window_fn, trigger_fn, accumulation_mode, timestamp_combiner, allowed_lateness, transcript, spec): if timestamp_combiner == TimestampCombiner.OUTPUT_AT_EARLIEST_TRANSFORMED: self.skipTest( 'Non-fnapi timestamp combiner: %s' % spec.get('timestamp_combiner')) if accumulation_mode != AccumulationMode.ACCUMULATING: self.skipTest('Batch mode only makes sense for accumulating.') watermark = MIN_TIMESTAMP for action, params in transcript: if action == 'watermark': watermark = params elif action == 'input': if any(t <= watermark for t in params): self.skipTest('Batch mode never has late data.') inputs = sum([vs for action, vs in transcript if action == 'input'], []) final_panes_by_window = {} for action, params in transcript: if action == 'expect': for expected in params: trimmed = {} for field in ('window', 'values', 'timestamp'): if field in expected: trimmed[field] = expected[field] final_panes_by_window[tuple(expected['window'])] = trimmed final_panes = list(final_panes_by_window.values()) if window_fn.is_merging(): merged_away = set() class MergeContext(WindowFn.MergeContext): def merge(_, to_be_merged, merge_result): for window in to_be_merged: if window != merge_result: merged_away.add(window) all_windows = [IntervalWindow(*pane['window']) for pane in final_panes] window_fn.merge(MergeContext(all_windows)) final_panes = [ pane for pane in final_panes if IntervalWindow(*pane['window']) not in merged_away ] # TODO(https://github.com/apache/beam/issues/34549): This test relies on # the correct error message being reported, but because of how it is # structured Prism often fails on splitting at the same time as it fails # on an element, leading to good logs but potentially inconsistent error # messages. Its not a bug, but it does mess with regex matching. with TestPipeline('FnApiRunner') as p: input_pc = ( p | beam.Create(inputs) | beam.Map(lambda t: TimestampedValue(('key', t), t)) | beam.WindowInto( window_fn, trigger=trigger_fn, accumulation_mode=accumulation_mode, timestamp_combiner=timestamp_combiner, allowed_lateness=allowed_lateness)) grouped = input_pc | 'Grouped' >> ( beam.GroupByKey() | beam.MapTuple(_windowed_value_info_map_fn) | beam.MapTuple(lambda _, value: value)) combined = input_pc | 'Combined' >> ( beam.CombinePerKey(_ConcatCombineFn()) | beam.MapTuple(_windowed_value_info_map_fn) | beam.MapTuple(lambda _, value: value)) assert_that( grouped, lambda actual: _windowed_value_info_check(actual, final_panes), label='CheckGrouped') assert_that( combined, lambda actual: _windowed_value_info_check(actual, final_panes), label='CheckCombined') TRANSCRIPT_TEST_FILE = os.path.join( os.path.dirname(__file__), '..', 'testing', 'data', 'trigger_transcripts.yaml') if os.path.exists(TRANSCRIPT_TEST_FILE): TriggerDriverTranscriptTest._create_tests(TRANSCRIPT_TEST_FILE) TestStreamTranscriptTest._create_tests(TRANSCRIPT_TEST_FILE) WeakTestStreamTranscriptTest._create_tests(TRANSCRIPT_TEST_FILE) BatchTranscriptTest._create_tests(TRANSCRIPT_TEST_FILE) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/userstate.py0000644000175100017510000003712315206113033022453 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """User-facing interfaces for the Beam State and Timer APIs.""" # pytype: skip-file # mypy: disallow-untyped-defs import collections import types from collections.abc import Callable from collections.abc import Iterable from typing import TYPE_CHECKING from typing import Any from typing import NamedTuple from typing import Optional from typing import TypeVar from apache_beam.coders import Coder from apache_beam.coders import coders from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.transforms.timeutil import TimeDomain from apache_beam.utils import windowed_value from apache_beam.utils.timestamp import Timestamp if TYPE_CHECKING: from apache_beam.runners.pipeline_context import PipelineContext from apache_beam.transforms.core import DoFn CallableT = TypeVar('CallableT', bound=Callable) class StateSpec(object): """Specification for a user DoFn state cell.""" def __init__(self, name: str, coder: Coder) -> None: if not isinstance(name, str): raise TypeError("name is not a string") if not isinstance(coder, Coder): raise TypeError("coder is not of type Coder") self.name = name self.coder = coder def __repr__(self) -> str: return '%s(%s)' % (self.__class__.__name__, self.name) def to_runner_api( self, context: 'PipelineContext') -> beam_runner_api_pb2.StateSpec: raise NotImplementedError class ReadModifyWriteStateSpec(StateSpec): """Specification for a user DoFn value state cell. Read more about ReadModifyWriteState (ValueState) here: https://beam.apache.org/documentation/programming-guide/#valuestate """ def to_runner_api( self, context: 'PipelineContext') -> beam_runner_api_pb2.StateSpec: return beam_runner_api_pb2.StateSpec( read_modify_write_spec=beam_runner_api_pb2.ReadModifyWriteStateSpec( coder_id=context.coders.get_id(self.coder)), protocol=beam_runner_api_pb2.FunctionSpec( urn=common_urns.user_state.BAG.urn)) class BagStateSpec(StateSpec): """Specification for a user DoFn bag state cell. Read more about BagState here: https://beam.apache.org/documentation/programming-guide/#bagstate """ def to_runner_api( self, context: 'PipelineContext') -> beam_runner_api_pb2.StateSpec: return beam_runner_api_pb2.StateSpec( bag_spec=beam_runner_api_pb2.BagStateSpec( element_coder_id=context.coders.get_id(self.coder)), protocol=beam_runner_api_pb2.FunctionSpec( urn=common_urns.user_state.BAG.urn)) class SetStateSpec(StateSpec): """Specification for a user DoFn Set State cell. Read more about SetState here: https://beam.apache.org/documentation/programming-guide/#setstate """ def to_runner_api( self, context: 'PipelineContext') -> beam_runner_api_pb2.StateSpec: return beam_runner_api_pb2.StateSpec( set_spec=beam_runner_api_pb2.SetStateSpec( element_coder_id=context.coders.get_id(self.coder)), protocol=beam_runner_api_pb2.FunctionSpec( urn=common_urns.user_state.BAG.urn)) class CombiningValueStateSpec(StateSpec): """Specification for a user DoFn combining value state cell. Read more about CombiningValueState here: https://beam.apache.org/documentation/programming-guide/#combiningstate """ def __init__( self, name: str, coder: Optional[Coder] = None, combine_fn: Any = None) -> None: """Initialize the specification for CombiningValue state. CombiningValueStateSpec(name, combine_fn) -> Coder-inferred combining value state spec. CombiningValueStateSpec(name, coder, combine_fn) -> Combining value state spec with coder and combine_fn specified. Args: name (str): The name by which the state is identified. coder (Coder): Coder specifying how to encode the values to be combined. May be inferred. combine_fn (``CombineFn`` or ``callable``): Function specifying how to combine the values passed to state. """ # Avoid circular import. from apache_beam.transforms.core import CombineFn # We want the coder to be optional, but unfortunately it comes # before the non-optional combine_fn parameter, which we can't # change for backwards compatibility reasons. # # Instead, allow it to be omitted (by either passing two arguments # or combine_fn by keyword.) if combine_fn is None: if coder is None: raise ValueError('combine_fn must be provided') else: coder, combine_fn = None, coder self.combine_fn = CombineFn.maybe_from_callable(combine_fn) # The coder here should be for the accumulator type of the given CombineFn. if coder is None: coder = self.combine_fn.get_accumulator_coder() super().__init__(name, coder) def to_runner_api( self, context: 'PipelineContext') -> beam_runner_api_pb2.StateSpec: return beam_runner_api_pb2.StateSpec( combining_spec=beam_runner_api_pb2.CombiningStateSpec( combine_fn=self.combine_fn.to_runner_api(context), accumulator_coder_id=context.coders.get_id(self.coder)), protocol=beam_runner_api_pb2.FunctionSpec( urn=common_urns.user_state.BAG.urn)) class OrderedListStateSpec(StateSpec): """Specification for a user DoFn ordered list state cell. Read more about OrderedListState here: https://beam.apache.org/documentation/programming-guide/#orderliststate """ def to_runner_api( self, context: 'PipelineContext') -> beam_runner_api_pb2.StateSpec: return beam_runner_api_pb2.StateSpec( ordered_list_spec=beam_runner_api_pb2.OrderedListStateSpec( element_coder_id=context.coders.get_id(self.coder)), protocol=beam_runner_api_pb2.FunctionSpec( urn=common_urns.user_state.ORDERED_LIST.urn)) # TODO(BEAM-9562): Update Timer to have of() and clear() APIs. Timer = NamedTuple( 'Timer', [ ('user_key', Any), ('dynamic_timer_tag', str), ('windows', tuple['windowed_value.BoundedWindow', ...]), ('clear_bit', bool), ('fire_timestamp', Optional['Timestamp']), ('hold_timestamp', Optional['Timestamp']), ('paneinfo', Optional['windowed_value.PaneInfo']), ]) # TODO(BEAM-9562): Plumb through actual key_coder and window_coder. class TimerSpec(object): """Specification for a user stateful DoFn timer. Read more about Timers here: https://beam.apache.org/documentation/programming-guide/#timers """ prefix = "ts-" def __init__(self, name: str, time_domain: str) -> None: self.name = self.prefix + name if time_domain not in (TimeDomain.WATERMARK, TimeDomain.REAL_TIME): raise ValueError('Unsupported TimeDomain: %r.' % (time_domain, )) self.time_domain = time_domain self._attached_callback: Optional[Callable] = None def __repr__(self) -> str: return '%s(%s)' % (self.__class__.__name__, self.name) def to_runner_api( self, context: 'PipelineContext', key_coder: Coder, window_coder: Coder) -> beam_runner_api_pb2.TimerFamilySpec: return beam_runner_api_pb2.TimerFamilySpec( time_domain=TimeDomain.to_runner_api(self.time_domain), timer_family_coder_id=context.coders.get_id( coders._TimerCoder(key_coder, window_coder))) def on_timer(timer_spec: TimerSpec) -> Callable[[CallableT], CallableT]: """Decorator for timer firing DoFn method. This decorator allows a user to specify an on_timer processing method in a stateful DoFn. Sample usage:: class MyDoFn(DoFn): TIMER_SPEC = TimerSpec('timer', TimeDomain.WATERMARK) @on_timer(TIMER_SPEC) def my_timer_expiry_callback(self): logging.info('Timer expired!') """ if not isinstance(timer_spec, TimerSpec): raise ValueError('@on_timer decorator expected TimerSpec.') def _inner(method: CallableT) -> CallableT: if not callable(method): raise ValueError('@on_timer decorator expected callable.') if timer_spec._attached_callback: raise ValueError( 'Multiple on_timer callbacks registered for %r.' % timer_spec) timer_spec._attached_callback = method return method return _inner def get_dofn_specs(dofn: 'DoFn') -> tuple[set[StateSpec], set[TimerSpec]]: """Gets the state and timer specs for a DoFn, if any. Args: dofn (apache_beam.transforms.core.DoFn): The DoFn instance to introspect for timer and state specs. """ # Avoid circular import. from apache_beam.runners.common import MethodWrapper from apache_beam.transforms.core import _DoFnParam from apache_beam.transforms.core import _StateDoFnParam from apache_beam.transforms.core import _TimerDoFnParam all_state_specs = set() all_timer_specs = set() # Validate params to process(), start_bundle(), finish_bundle() and to # any on_timer callbacks. for method_name in dir(dofn): if not isinstance(getattr(dofn, method_name, None), types.MethodType): continue method = MethodWrapper(dofn, method_name) param_ids = [ d.param_id for d in method.defaults if isinstance(d, _DoFnParam) ] if len(param_ids) != len(set(param_ids)): raise ValueError( 'DoFn %r has duplicate %s method parameters: %s.' % (dofn, method_name, param_ids)) for d in method.defaults: if isinstance(d, _StateDoFnParam): all_state_specs.add(d.state_spec) elif isinstance(d, _TimerDoFnParam): all_timer_specs.add(d.timer_spec) return all_state_specs, all_timer_specs def is_stateful_dofn(dofn: 'DoFn') -> bool: """Determines whether a given DoFn is a stateful DoFn.""" # A Stateful DoFn is a DoFn that uses user state or timers. all_state_specs, all_timer_specs = get_dofn_specs(dofn) return bool(all_state_specs or all_timer_specs) def validate_stateful_dofn(dofn: 'DoFn') -> None: """Validates the proper specification of a stateful DoFn.""" # Get state and timer specs. all_state_specs, all_timer_specs = get_dofn_specs(dofn) # Reject DoFns that have multiple state or timer specs with the same name. if len(all_state_specs) != len(set(s.name for s in all_state_specs)): raise ValueError( 'DoFn %r has multiple StateSpecs with the same name: %s.' % (dofn, all_state_specs)) if len(all_timer_specs) != len(set(s.name for s in all_timer_specs)): raise ValueError( 'DoFn %r has multiple TimerSpecs with the same name: %s.' % (dofn, all_timer_specs)) # Reject DoFns that use timer specs without corresponding timer callbacks. for timer_spec in all_timer_specs: if not timer_spec._attached_callback: raise ValueError(( 'DoFn %r has a TimerSpec without an associated on_timer ' 'callback: %s.') % (dofn, timer_spec)) method_name = timer_spec._attached_callback.__name__ if (timer_spec._attached_callback != getattr(dofn, method_name, None).__func__): # type: ignore[union-attr] raise ValueError(( 'The on_timer callback for %s is not the specified .%s method ' 'for DoFn %r (perhaps it was overwritten?).') % (timer_spec, method_name, dofn)) class BaseTimer(object): def clear(self, dynamic_timer_tag: str = '') -> None: raise NotImplementedError def set(self, timestamp: Timestamp, dynamic_timer_tag: str = '') -> None: raise NotImplementedError _TimerTuple = collections.namedtuple('timer_tuple', ('cleared', 'timestamp')) # type: ignore[name-match] class RuntimeTimer(BaseTimer): """Timer interface object passed to user code.""" def __init__(self) -> None: self._timer_recordings: dict[str, _TimerTuple] = {} self._cleared = False self._new_timestamp: Optional[Timestamp] = None def clear(self, dynamic_timer_tag: str = '') -> None: self._timer_recordings[dynamic_timer_tag] = _TimerTuple( cleared=True, timestamp=None) def set(self, timestamp: Timestamp, dynamic_timer_tag: str = '') -> None: self._timer_recordings[dynamic_timer_tag] = _TimerTuple( cleared=False, timestamp=timestamp) class RuntimeState(object): """State interface object passed to user code.""" def prefetch(self) -> None: # The default implementation here does nothing. pass def finalize(self) -> None: pass class ReadModifyWriteRuntimeState(RuntimeState): def read(self) -> Any: raise NotImplementedError(type(self)) def write(self, value: Any) -> None: raise NotImplementedError(type(self)) def clear(self) -> None: raise NotImplementedError(type(self)) def commit(self) -> None: raise NotImplementedError(type(self)) class AccumulatingRuntimeState(RuntimeState): def read(self) -> Iterable[Any]: raise NotImplementedError(type(self)) def add(self, value: Any) -> None: raise NotImplementedError(type(self)) def clear(self) -> None: raise NotImplementedError(type(self)) def commit(self) -> None: raise NotImplementedError(type(self)) class BagRuntimeState(AccumulatingRuntimeState): """Bag state interface object passed to user code. Read more about BagState here: https://beam.apache.org/documentation/programming-guide/#bagstate """ class SetRuntimeState(AccumulatingRuntimeState): """Set state interface object passed to user code. Read more about SetState here: https://beam.apache.org/documentation/programming-guide/#setstate """ class CombiningValueRuntimeState(AccumulatingRuntimeState): """Combining value state interface object passed to user code. Read more about CombiningValueState here: https://beam.apache.org/documentation/programming-guide/#combiningstate """ class OrderedListRuntimeState(AccumulatingRuntimeState): """Ordered list state interface object passed to user code. Read more about OrderedListState here: https://beam.apache.org/documentation/programming-guide/#orderliststate """ def read(self) -> Iterable[tuple[Timestamp, Any]]: raise NotImplementedError(type(self)) def add(self, value: tuple[Timestamp, Any]) -> None: raise NotImplementedError(type(self)) def read_range( self, min_time_stamp: Timestamp, limit_time_stamp: Timestamp) -> Iterable[tuple[Timestamp, Any]]: raise NotImplementedError(type(self)) def clear_range( self, min_time_stamp: Timestamp, limit_time_stamp: Timestamp) -> None: raise NotImplementedError(type(self)) class UserStateContext(object): """Wrapper allowing user state and timers to be accessed by a DoFnInvoker.""" def get_timer( self, timer_spec: TimerSpec, key: Any, window: 'windowed_value.BoundedWindow', timestamp: Timestamp, pane: windowed_value.PaneInfo, ) -> BaseTimer: raise NotImplementedError(type(self)) def get_state( self, state_spec: StateSpec, key: Any, window: 'windowed_value.BoundedWindow', ) -> RuntimeState: raise NotImplementedError(type(self)) def commit(self) -> None: raise NotImplementedError(type(self)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/userstate_test.py0000644000175100017510000012547615206113033023523 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the Beam State and Timer API interfaces.""" # pytype: skip-file import unittest from typing import Any import mock import pytest import apache_beam as beam from apache_beam.coders import BytesCoder from apache_beam.coders import ListCoder from apache_beam.coders import StrUtf8Coder from apache_beam.coders import VarIntCoder from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.runners import pipeline_context from apache_beam.runners.common import DoFnSignature from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import trigger from apache_beam.transforms import userstate from apache_beam.transforms import window from apache_beam.transforms.combiners import ToListCombineFn from apache_beam.transforms.combiners import TopCombineFn from apache_beam.transforms.core import DoFn from apache_beam.transforms.timeutil import TimeDomain from apache_beam.transforms.userstate import BagStateSpec from apache_beam.transforms.userstate import CombiningValueStateSpec from apache_beam.transforms.userstate import ReadModifyWriteStateSpec from apache_beam.transforms.userstate import SetStateSpec from apache_beam.transforms.userstate import TimerSpec from apache_beam.transforms.userstate import get_dofn_specs from apache_beam.transforms.userstate import is_stateful_dofn from apache_beam.transforms.userstate import on_timer from apache_beam.transforms.userstate import validate_stateful_dofn class TestStatefulDoFn(DoFn): """An example stateful DoFn with state and timers.""" BUFFER_STATE_1 = BagStateSpec('buffer', BytesCoder()) BUFFER_STATE_2 = BagStateSpec('buffer2', VarIntCoder()) EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK) EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK) EXPIRY_TIMER_3 = TimerSpec('expiry3', TimeDomain.WATERMARK) EXPIRY_TIMER_FAMILY = TimerSpec('expiry_family', TimeDomain.WATERMARK) def process( self, element, t=DoFn.TimestampParam, buffer_1=DoFn.StateParam(BUFFER_STATE_1), buffer_2=DoFn.StateParam(BUFFER_STATE_2), timer_1=DoFn.TimerParam(EXPIRY_TIMER_1), timer_2=DoFn.TimerParam(EXPIRY_TIMER_2), dynamic_timer=DoFn.TimerParam(EXPIRY_TIMER_FAMILY)): yield element @on_timer(EXPIRY_TIMER_1) def on_expiry_1( self, window=DoFn.WindowParam, timestamp=DoFn.TimestampParam, key=DoFn.KeyParam, buffer=DoFn.StateParam(BUFFER_STATE_1), timer_1=DoFn.TimerParam(EXPIRY_TIMER_1), timer_2=DoFn.TimerParam(EXPIRY_TIMER_2), timer_3=DoFn.TimerParam(EXPIRY_TIMER_3)): yield 'expired1' @on_timer(EXPIRY_TIMER_2) def on_expiry_2( self, buffer=DoFn.StateParam(BUFFER_STATE_2), timer_2=DoFn.TimerParam(EXPIRY_TIMER_2), timer_3=DoFn.TimerParam(EXPIRY_TIMER_3)): yield 'expired2' @on_timer(EXPIRY_TIMER_3) def on_expiry_3( self, buffer_1=DoFn.StateParam(BUFFER_STATE_1), buffer_2=DoFn.StateParam(BUFFER_STATE_2), timer_3=DoFn.TimerParam(EXPIRY_TIMER_3)): yield 'expired3' @on_timer(EXPIRY_TIMER_FAMILY) def on_expiry_family( self, dynamic_timer=DoFn.TimerParam(EXPIRY_TIMER_FAMILY), dynamic_timer_tag=DoFn.DynamicTimerTagParam): yield (dynamic_timer_tag, 'expired_dynamic_timer') class InterfaceTest(unittest.TestCase): def _validate_dofn(self, dofn): # Construction of DoFnSignature performs validation of the given DoFn. # In particular, it ends up calling userstate._validate_stateful_dofn. # That behavior is explicitly tested below in test_validate_dofn() return DoFnSignature(dofn) @mock.patch('apache_beam.transforms.userstate.validate_stateful_dofn') def test_validate_dofn(self, unused_mock): dofn = TestStatefulDoFn() self._validate_dofn(dofn) userstate.validate_stateful_dofn.assert_called_with(dofn) def test_spec_construction(self): BagStateSpec('statename', VarIntCoder()) with self.assertRaises(TypeError): BagStateSpec(123, VarIntCoder()) CombiningValueStateSpec('statename', VarIntCoder(), TopCombineFn(10)) with self.assertRaises(TypeError): CombiningValueStateSpec(123, VarIntCoder(), TopCombineFn(10)) with self.assertRaises(TypeError): CombiningValueStateSpec('statename', VarIntCoder(), object()) SetStateSpec('setstatename', VarIntCoder()) with self.assertRaises(TypeError): SetStateSpec(123, VarIntCoder()) with self.assertRaises(TypeError): SetStateSpec('setstatename', object()) ReadModifyWriteStateSpec('valuestatename', VarIntCoder()) with self.assertRaises(TypeError): ReadModifyWriteStateSpec(123, VarIntCoder()) with self.assertRaises(TypeError): ReadModifyWriteStateSpec('valuestatename', object()) # TODO: add more spec tests with self.assertRaises(ValueError): DoFn.TimerParam(BagStateSpec('elements', BytesCoder())) TimerSpec('timer', TimeDomain.WATERMARK) TimerSpec('timer', TimeDomain.REAL_TIME) with self.assertRaises(ValueError): TimerSpec('timer', 'bogus_time_domain') with self.assertRaises(ValueError): DoFn.StateParam(TimerSpec('timer', TimeDomain.WATERMARK)) def test_state_spec_proto_conversion(self): context = pipeline_context.PipelineContext() state = BagStateSpec('statename', VarIntCoder()) state_proto = state.to_runner_api(context) self.assertEqual( beam_runner_api_pb2.FunctionSpec(urn=common_urns.user_state.BAG.urn), state_proto.protocol) context = pipeline_context.PipelineContext() state = CombiningValueStateSpec( 'statename', VarIntCoder(), TopCombineFn(10)) state_proto = state.to_runner_api(context) self.assertEqual( beam_runner_api_pb2.FunctionSpec(urn=common_urns.user_state.BAG.urn), state_proto.protocol) context = pipeline_context.PipelineContext() state = SetStateSpec('setstatename', VarIntCoder()) state_proto = state.to_runner_api(context) self.assertEqual( beam_runner_api_pb2.FunctionSpec(urn=common_urns.user_state.BAG.urn), state_proto.protocol) context = pipeline_context.PipelineContext() state = ReadModifyWriteStateSpec('valuestatename', VarIntCoder()) state_proto = state.to_runner_api(context) self.assertEqual( beam_runner_api_pb2.FunctionSpec(urn=common_urns.user_state.BAG.urn), state_proto.protocol) def test_param_construction(self): with self.assertRaises(ValueError): DoFn.StateParam(TimerSpec('timer', TimeDomain.WATERMARK)) with self.assertRaises(ValueError): DoFn.TimerParam(BagStateSpec('elements', BytesCoder())) def test_stateful_dofn_detection(self): self.assertFalse(is_stateful_dofn(DoFn())) self.assertTrue(is_stateful_dofn(TestStatefulDoFn())) def test_good_signatures(self): class BasicStatefulDoFn(DoFn): BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) EXPIRY_TIMER = TimerSpec('expiry1', TimeDomain.WATERMARK) EXPIRY_TIMER_FAMILY = TimerSpec('expiry_family_1', TimeDomain.WATERMARK) def process( self, element, buffer=DoFn.StateParam(BUFFER_STATE), timer1=DoFn.TimerParam(EXPIRY_TIMER), dynamic_timer=DoFn.TimerParam(EXPIRY_TIMER_FAMILY)): yield element @on_timer(EXPIRY_TIMER) def expiry_callback(self, element, timer=DoFn.TimerParam(EXPIRY_TIMER)): yield element @on_timer(EXPIRY_TIMER_FAMILY) def expiry_family_callback( self, element, dynamic_timer=DoFn.TimerParam(EXPIRY_TIMER_FAMILY)): yield element # Validate get_dofn_specs() and timer callbacks in # DoFnSignature. stateful_dofn = BasicStatefulDoFn() signature = self._validate_dofn(stateful_dofn) expected_specs = ( set([BasicStatefulDoFn.BUFFER_STATE]), set([ BasicStatefulDoFn.EXPIRY_TIMER, BasicStatefulDoFn.EXPIRY_TIMER_FAMILY ]), ) self.assertEqual(expected_specs, get_dofn_specs(stateful_dofn)) self.assertEqual( stateful_dofn.expiry_callback, signature.timer_methods[BasicStatefulDoFn.EXPIRY_TIMER].method_value) self.assertEqual( stateful_dofn.expiry_family_callback, signature.timer_methods[ BasicStatefulDoFn.EXPIRY_TIMER_FAMILY].method_value) stateful_dofn = TestStatefulDoFn() signature = self._validate_dofn(stateful_dofn) expected_specs = ( set([TestStatefulDoFn.BUFFER_STATE_1, TestStatefulDoFn.BUFFER_STATE_2]), set([ TestStatefulDoFn.EXPIRY_TIMER_1, TestStatefulDoFn.EXPIRY_TIMER_2, TestStatefulDoFn.EXPIRY_TIMER_3, TestStatefulDoFn.EXPIRY_TIMER_FAMILY ])) self.assertEqual(expected_specs, get_dofn_specs(stateful_dofn)) self.assertEqual( stateful_dofn.on_expiry_1, signature.timer_methods[TestStatefulDoFn.EXPIRY_TIMER_1].method_value) self.assertEqual( stateful_dofn.on_expiry_2, signature.timer_methods[TestStatefulDoFn.EXPIRY_TIMER_2].method_value) self.assertEqual( stateful_dofn.on_expiry_3, signature.timer_methods[TestStatefulDoFn.EXPIRY_TIMER_3].method_value) self.assertEqual( stateful_dofn.on_expiry_family, signature.timer_methods[ TestStatefulDoFn.EXPIRY_TIMER_FAMILY].method_value) def test_bad_signatures(self): # (1) The same state parameter is duplicated on the process method. class BadStatefulDoFn1(DoFn): BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) def process( self, element, b1=DoFn.StateParam(BUFFER_STATE), b2=DoFn.StateParam(BUFFER_STATE)): yield element with self.assertRaises(ValueError): self._validate_dofn(BadStatefulDoFn1()) # (2) The same timer parameter is duplicated on the process method. class BadStatefulDoFn2(DoFn): TIMER = TimerSpec('timer', TimeDomain.WATERMARK) def process( self, element, t1=DoFn.TimerParam(TIMER), t2=DoFn.TimerParam(TIMER)): yield element with self.assertRaises(ValueError): self._validate_dofn(BadStatefulDoFn2()) # (3) The same state parameter is duplicated on the on_timer method. class BadStatefulDoFn3(DoFn): BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK) EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK) @on_timer(EXPIRY_TIMER_1) def expiry_callback( self, element, b1=DoFn.StateParam(BUFFER_STATE), b2=DoFn.StateParam(BUFFER_STATE)): yield element with self.assertRaises(ValueError): self._validate_dofn(BadStatefulDoFn3()) # (4) The same timer parameter is duplicated on the on_timer method. class BadStatefulDoFn4(DoFn): BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK) EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK) @on_timer(EXPIRY_TIMER_1) def expiry_callback( self, element, t1=DoFn.TimerParam(EXPIRY_TIMER_2), t2=DoFn.TimerParam(EXPIRY_TIMER_2)): yield element with self.assertRaises(ValueError): self._validate_dofn(BadStatefulDoFn4()) # (5) The same timer family parameter is duplicated on the process method. class BadStatefulDoFn5(DoFn): EXPIRY_TIMER_FAMILY = TimerSpec('dynamic_timer', TimeDomain.WATERMARK) def process( self, element, dynamic_timer_1=DoFn.TimerParam(EXPIRY_TIMER_FAMILY), dynamic_timer_2=DoFn.TimerParam(EXPIRY_TIMER_FAMILY)): yield element with self.assertRaises(ValueError): self._validate_dofn(BadStatefulDoFn5()) def test_validation_typos(self): # (1) Here, the user mistakenly used the same timer spec twice for two # different timer callbacks. with self.assertRaisesRegex( ValueError, r'Multiple on_timer callbacks registered for TimerSpec\(.*expiry1\).'): class StatefulDoFnWithTimerWithTypo1(DoFn): # pylint: disable=unused-variable BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK) EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK) def process(self, element): pass @on_timer(EXPIRY_TIMER_1) def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)): yield 'expired1' # Note that we mistakenly associate this with the first timer. @on_timer(EXPIRY_TIMER_1) def on_expiry_2(self, buffer_state=DoFn.StateParam(BUFFER_STATE)): yield 'expired2' # (2) Here, the user mistakenly used the same callback name and overwrote # the first on_expiry_1 callback. class StatefulDoFnWithTimerWithTypo2(DoFn): BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK) EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK) def process( self, element, timer1=DoFn.TimerParam(EXPIRY_TIMER_1), timer2=DoFn.TimerParam(EXPIRY_TIMER_2)): pass @on_timer(EXPIRY_TIMER_1) def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)): yield 'expired1' # Note that we mistakenly reuse the "on_expiry_1" name; this is valid # syntactically in Python. @on_timer(EXPIRY_TIMER_2) def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)): # pylint: disable=function-redefined yield 'expired2' # Use a stable string value for matching. def __repr__(self): return 'StatefulDoFnWithTimerWithTypo2' dofn = StatefulDoFnWithTimerWithTypo2() with self.assertRaisesRegex( ValueError, (r'The on_timer callback for TimerSpec\(.*expiry1\) is not the ' r'specified .on_expiry_1 method for DoFn ' r'StatefulDoFnWithTimerWithTypo2 \(perhaps it was overwritten\?\).')): validate_stateful_dofn(dofn) # (2) Here, the user forgot to add an on_timer decorator for 'expiry2' class StatefulDoFnWithTimerWithTypo3(DoFn): BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK) EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK) def process( self, element, timer1=DoFn.TimerParam(EXPIRY_TIMER_1), timer2=DoFn.TimerParam(EXPIRY_TIMER_2)): pass @on_timer(EXPIRY_TIMER_1) def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)): yield 'expired1' def on_expiry_2(self, buffer_state=DoFn.StateParam(BUFFER_STATE)): yield 'expired2' # Use a stable string value for matching. def __repr__(self): return 'StatefulDoFnWithTimerWithTypo3' dofn = StatefulDoFnWithTimerWithTypo3() with self.assertRaisesRegex( ValueError, (r'DoFn StatefulDoFnWithTimerWithTypo3 has a TimerSpec without an ' r'associated on_timer callback: TimerSpec\(.*expiry2\).')): validate_stateful_dofn(dofn) class StatefulDoFnOnDirectRunnerTest(unittest.TestCase): # pylint: disable=expression-not-assigned all_records: list[Any] def setUp(self): # Use state on the TestCase class, since other references would be pickled # into a closure and not have the desired side effects. # # TODO(https://github.com/apache/beam/issues/18987): Use assert_that after # it works for the cases here in streaming mode. StatefulDoFnOnDirectRunnerTest.all_records = [] def record_dofn(self): class RecordDoFn(DoFn): def process(self, element): StatefulDoFnOnDirectRunnerTest.all_records.append(element) return RecordDoFn() def test_simple_stateful_dofn(self): class SimpleTestStatefulDoFn(DoFn): BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) EXPIRY_TIMER = TimerSpec('expiry', TimeDomain.WATERMARK) def process( self, element, buffer=DoFn.StateParam(BUFFER_STATE), timer1=DoFn.TimerParam(EXPIRY_TIMER)): unused_key, value = element buffer.add(b'A' + str(value).encode('latin1')) timer1.set(20) @on_timer(EXPIRY_TIMER) def expiry_callback( self, buffer=DoFn.StateParam(BUFFER_STATE), timer=DoFn.TimerParam(EXPIRY_TIMER)): yield b''.join(sorted(buffer.read())) with TestPipeline() as p: test_stream = ( TestStream().advance_watermark_to(10).add_elements( [1, 2]).add_elements([3]).advance_watermark_to(25).add_elements([4])) ( p | test_stream | beam.Map(lambda x: ('mykey', x)) | beam.ParDo(SimpleTestStatefulDoFn()) | beam.ParDo(self.record_dofn())) # Two firings should occur: once after element 3 since the timer should # fire after the watermark passes time 20, and another time after element # 4, since the timer issued at that point should fire immediately. self.assertEqual([b'A1A2A3', b'A1A2A3A4'], StatefulDoFnOnDirectRunnerTest.all_records) def test_clearing_bag_state(self): class BagStateClearingStatefulDoFn(beam.DoFn): BAG_STATE = BagStateSpec('bag_state', StrUtf8Coder()) EMIT_TIMER = TimerSpec('emit_timer', TimeDomain.WATERMARK) CLEAR_TIMER = TimerSpec('clear_timer', TimeDomain.WATERMARK) def process( self, element, bag_state=beam.DoFn.StateParam(BAG_STATE), emit_timer=beam.DoFn.TimerParam(EMIT_TIMER), clear_timer=beam.DoFn.TimerParam(CLEAR_TIMER)): value = element[1] bag_state.add(value) clear_timer.set(100) emit_timer.set(1000) @on_timer(EMIT_TIMER) def emit_values(self, bag_state=beam.DoFn.StateParam(BAG_STATE)): for value in bag_state.read(): yield value yield 'extra' @on_timer(CLEAR_TIMER) def clear_values(self, bag_state=beam.DoFn.StateParam(BAG_STATE)): bag_state.clear() with TestPipeline() as p: test_stream = ( TestStream().advance_watermark_to(0).add_elements([ ('key', 'value') ]).advance_watermark_to(100)) _ = ( p | test_stream | beam.ParDo(BagStateClearingStatefulDoFn()) | beam.ParDo(self.record_dofn())) self.assertEqual(['extra'], StatefulDoFnOnDirectRunnerTest.all_records) def test_two_timers_one_function(self): class BagStateClearingStatefulDoFn(beam.DoFn): BAG_STATE = BagStateSpec('bag_state', StrUtf8Coder()) EMIT_TIMER = TimerSpec('emit_timer', TimeDomain.WATERMARK) EMIT_TWICE_TIMER = TimerSpec('clear_timer', TimeDomain.WATERMARK) def process( self, element, bag_state=beam.DoFn.StateParam(BAG_STATE), emit_timer=beam.DoFn.TimerParam(EMIT_TIMER), emit_twice_timer=beam.DoFn.TimerParam(EMIT_TWICE_TIMER)): value = element[1] bag_state.add(value) emit_twice_timer.set(100) emit_timer.set(1000) @on_timer(EMIT_TWICE_TIMER) @on_timer(EMIT_TIMER) def emit_values(self, bag_state=beam.DoFn.StateParam(BAG_STATE)): for value in bag_state.read(): yield value with TestPipeline() as p: test_stream = ( TestStream().advance_watermark_to(0).add_elements([ ('key', 'value') ]).advance_watermark_to(100)) _ = ( p | test_stream | beam.ParDo(BagStateClearingStatefulDoFn()) | beam.ParDo(self.record_dofn())) self.assertEqual(['value', 'value'], StatefulDoFnOnDirectRunnerTest.all_records) def test_simple_read_modify_write_stateful_dofn(self): class SimpleTestReadModifyWriteStatefulDoFn(DoFn): VALUE_STATE = ReadModifyWriteStateSpec('value', StrUtf8Coder()) def process(self, element, last_element=DoFn.StateParam(VALUE_STATE)): last_element.write('%s:%s' % element) yield last_element.read() with TestPipeline() as p: test_stream = ( TestStream().advance_watermark_to(0).add_elements([ ('a', 1) ]).advance_watermark_to(10).add_elements([ ('a', 3) ]).advance_watermark_to(20).add_elements([('a', 5)])) ( p | test_stream | beam.ParDo(SimpleTestReadModifyWriteStatefulDoFn()) | beam.ParDo(self.record_dofn())) self.assertEqual(['a:1', 'a:3', 'a:5'], StatefulDoFnOnDirectRunnerTest.all_records) def test_clearing_read_modify_write_state(self): class SimpleClearingReadModifyWriteStatefulDoFn(DoFn): VALUE_STATE = ReadModifyWriteStateSpec('value', StrUtf8Coder()) def process(self, element, last_element=DoFn.StateParam(VALUE_STATE)): value = last_element.read() if value is not None: yield value last_element.clear() last_element.write("%s:%s" % (last_element.read(), element[1])) if element[1] == 5: yield last_element.read() with TestPipeline() as p: test_stream = ( TestStream().advance_watermark_to(0).add_elements([ ('a', 1) ]).advance_watermark_to(10).add_elements([ ('a', 3) ]).advance_watermark_to(20).add_elements([('a', 5)])) ( p | test_stream | beam.ParDo(SimpleClearingReadModifyWriteStatefulDoFn()) | beam.ParDo(self.record_dofn())) self.assertEqual(['None:1', 'None:3', 'None:5'], StatefulDoFnOnDirectRunnerTest.all_records) def test_simple_set_stateful_dofn(self): class SimpleTestSetStatefulDoFn(DoFn): BUFFER_STATE = SetStateSpec('buffer', VarIntCoder()) EXPIRY_TIMER = TimerSpec('expiry', TimeDomain.WATERMARK) def process( self, element, buffer=DoFn.StateParam(BUFFER_STATE), timer1=DoFn.TimerParam(EXPIRY_TIMER)): unused_key, value = element buffer.add(value) timer1.set(20) @on_timer(EXPIRY_TIMER) def expiry_callback(self, buffer=DoFn.StateParam(BUFFER_STATE)): yield sorted(buffer.read()) with TestPipeline() as p: test_stream = ( TestStream().advance_watermark_to(10).add_elements( [1, 2, 3]).add_elements([2]).advance_watermark_to(24)) ( p | test_stream | beam.Map(lambda x: ('mykey', x)) | beam.ParDo(SimpleTestSetStatefulDoFn()) | beam.ParDo(self.record_dofn())) # Two firings should occur: once after element 3 since the timer should # fire after the watermark passes time 20, and another time after element # 4, since the timer issued at that point should fire immediately. self.assertEqual([[1, 2, 3]], StatefulDoFnOnDirectRunnerTest.all_records) def test_clearing_set_state(self): class SetStateClearingStatefulDoFn(beam.DoFn): SET_STATE = SetStateSpec('buffer', StrUtf8Coder()) EMIT_TIMER = TimerSpec('emit_timer', TimeDomain.WATERMARK) CLEAR_TIMER = TimerSpec('clear_timer', TimeDomain.WATERMARK) def process( self, element, set_state=beam.DoFn.StateParam(SET_STATE), emit_timer=beam.DoFn.TimerParam(EMIT_TIMER), clear_timer=beam.DoFn.TimerParam(CLEAR_TIMER)): value = element[1] set_state.add(value) clear_timer.set(100) emit_timer.set(1000) @on_timer(EMIT_TIMER) def emit_values(self, set_state=beam.DoFn.StateParam(SET_STATE)): for value in set_state.read(): yield value @on_timer(CLEAR_TIMER) def clear_values(self, set_state=beam.DoFn.StateParam(SET_STATE)): set_state.clear() set_state.add('different-value') with TestPipeline() as p: test_stream = ( TestStream().advance_watermark_to(0).add_elements([ ('key1', 'value1') ]).advance_watermark_to(100)) _ = ( p | test_stream | beam.ParDo(SetStateClearingStatefulDoFn()) | beam.ParDo(self.record_dofn())) self.assertEqual(['different-value'], StatefulDoFnOnDirectRunnerTest.all_records) def test_stateful_set_state_portably(self): class SetStatefulDoFn(beam.DoFn): SET_STATE = SetStateSpec('buffer', VarIntCoder()) def process(self, element, set_state=beam.DoFn.StateParam(SET_STATE)): _, value = element aggregated_value = 0 set_state.add(value) for saved_value in set_state.read(): aggregated_value += saved_value yield aggregated_value with TestPipeline() as p: values = p | beam.Create([('key', 1), ('key', 2), ('key', 3), ('key', 4), ('key', 3)], reshuffle=False) actual_values = (values | beam.ParDo(SetStatefulDoFn())) assert_that(actual_values, equal_to([1, 3, 6, 10, 10])) def test_stateful_set_state_clean_portably(self): class SetStateClearingStatefulDoFn(beam.DoFn): SET_STATE = SetStateSpec('buffer', VarIntCoder()) EMIT_TIMER = TimerSpec('emit_timer', TimeDomain.WATERMARK) def process( self, element, set_state=beam.DoFn.StateParam(SET_STATE), emit_timer=beam.DoFn.TimerParam(EMIT_TIMER)): _, value = element set_state.add(value) all_elements = [element for element in set_state.read()] if len(all_elements) == 5: set_state.clear() set_state.add(100) emit_timer.set(1) @on_timer(EMIT_TIMER) def emit_values(self, set_state=beam.DoFn.StateParam(SET_STATE)): yield sorted(set_state.read()) # Pin to FnApiRunner since this assumes a large bundle size to contain # all elements with TestPipeline('FnApiRunner') as p: values = p | beam.Create([('key', 1), ('key', 2), ('key', 3), ('key', 4), ('key', 5)]) actual_values = ( values | beam.Map(lambda t: window.TimestampedValue(t, 1)) | beam.WindowInto(window.FixedWindows(1)) | beam.ParDo(SetStateClearingStatefulDoFn())) assert_that(actual_values, equal_to([[100]])) def test_stateful_dofn_nonkeyed_input(self): p = TestPipeline() values = p | beam.Create([1, 2, 3]) with self.assertRaisesRegex( ValueError, ('Input elements to the transform .* with stateful DoFn must be ' 'key-value pairs.')): values | beam.ParDo(TestStatefulDoFn()) def test_generate_sequence_with_realtime_timer(self): from apache_beam.transforms.combiners import CountCombineFn class GenerateRecords(beam.DoFn): EMIT_TIMER = TimerSpec('emit_timer', TimeDomain.REAL_TIME) COUNT_STATE = CombiningValueStateSpec( 'count_state', VarIntCoder(), CountCombineFn()) def __init__(self, frequency, total_records): self.total_records = total_records self.frequency = frequency def process(self, element, emit_timer=beam.DoFn.TimerParam(EMIT_TIMER)): # Processing time timers should be set on ABSOLUTE TIME. emit_timer.set(self.frequency) yield element[1] @on_timer(EMIT_TIMER) def emit_values( self, emit_timer=beam.DoFn.TimerParam(EMIT_TIMER), count_state=beam.DoFn.StateParam(COUNT_STATE)): count = count_state.read() or 0 if self.total_records == count: return count_state.add(1) # Processing time timers should be set on ABSOLUTE TIME. emit_timer.set(count + 1 + self.frequency) yield 'value' TOTAL_RECORDS = 3 FREQUENCY = 1 test_stream = ( TestStream().advance_watermark_to(0).add_elements([ ('key', 0) ]).advance_processing_time(1) # Timestamp: 1 .add_elements([('key', 1)]).advance_processing_time(1) # Timestamp: 2 .add_elements([('key', 2)]).advance_processing_time(1) # Timestamp: 3 .add_elements([('key', 3)])) with beam.Pipeline(argv=['--streaming', '--runner=DirectRunner']) as p: _ = ( p | test_stream | beam.ParDo(GenerateRecords(FREQUENCY, TOTAL_RECORDS)) | beam.ParDo(self.record_dofn())) self.assertEqual( # 4 RECORDS go through process # 3 values are emitted from timer # Timestamp moves gradually. [0, 'value', 1, 'value', 2, 'value', 3], StatefulDoFnOnDirectRunnerTest.all_records) def test_simple_stateful_dofn_combining(self): class SimpleTestStatefulDoFn(DoFn): BUFFER_STATE = CombiningValueStateSpec( 'buffer', ListCoder(VarIntCoder()), ToListCombineFn()) EXPIRY_TIMER = TimerSpec('expiry1', TimeDomain.WATERMARK) def process( self, element, buffer=DoFn.StateParam(BUFFER_STATE), timer1=DoFn.TimerParam(EXPIRY_TIMER)): unused_key, value = element buffer.add(value) timer1.set(20) @on_timer(EXPIRY_TIMER) def expiry_callback( self, buffer=DoFn.StateParam(BUFFER_STATE), timer=DoFn.TimerParam(EXPIRY_TIMER)): yield ''.join(str(x) for x in sorted(buffer.read())) with TestPipeline() as p: test_stream = ( TestStream().advance_watermark_to(10).add_elements( [1, 2]).add_elements([3]).advance_watermark_to(25).add_elements([4])) ( p | test_stream | beam.Map(lambda x: ('mykey', x)) | beam.ParDo(SimpleTestStatefulDoFn()) | beam.ParDo(self.record_dofn())) self.assertEqual(['123', '1234'], StatefulDoFnOnDirectRunnerTest.all_records) def test_timer_output_timestamp(self): class TimerEmittingStatefulDoFn(DoFn): EMIT_TIMER_1 = TimerSpec('emit1', TimeDomain.WATERMARK) EMIT_TIMER_2 = TimerSpec('emit2', TimeDomain.WATERMARK) EMIT_TIMER_3 = TimerSpec('emit3', TimeDomain.WATERMARK) def process( self, element, timer1=DoFn.TimerParam(EMIT_TIMER_1), timer2=DoFn.TimerParam(EMIT_TIMER_2), timer3=DoFn.TimerParam(EMIT_TIMER_3)): timer1.set(10) timer2.set(20) timer3.set(30) @on_timer(EMIT_TIMER_1) def emit_callback_1(self): yield 'timer1' @on_timer(EMIT_TIMER_2) def emit_callback_2(self): yield 'timer2' @on_timer(EMIT_TIMER_3) def emit_callback_3(self): yield 'timer3' class TimestampReifyingDoFn(DoFn): def process(self, element, ts=DoFn.TimestampParam): yield (element, int(ts)) with TestPipeline() as p: test_stream = (TestStream().advance_watermark_to(10).add_elements([1])) ( p | test_stream | beam.Map(lambda x: ('mykey', x)) | beam.ParDo(TimerEmittingStatefulDoFn()) | beam.ParDo(TimestampReifyingDoFn()) | beam.ParDo(self.record_dofn())) self.assertEqual([('timer1', 10), ('timer2', 20), ('timer3', 30)], sorted(StatefulDoFnOnDirectRunnerTest.all_records)) def test_timer_output_timestamp_and_window(self): class TimerEmittingStatefulDoFn(DoFn): EMIT_TIMER_1 = TimerSpec('emit1', TimeDomain.WATERMARK) def process(self, element, timer1=DoFn.TimerParam(EMIT_TIMER_1)): timer1.set(10) @on_timer(EMIT_TIMER_1) def emit_callback_1( self, window=DoFn.WindowParam, ts=DoFn.TimestampParam, key=DoFn.KeyParam): yield ( 'timer1-{key}'.format(key=key), int(ts), int(window.start), int(window.end)) pipeline_options = PipelineOptions() with TestPipeline(options=pipeline_options) as p: test_stream = (TestStream().advance_watermark_to(10).add_elements([1])) ( p | test_stream | beam.Map(lambda x: ('mykey', x)) | "window_into" >> beam.WindowInto( window.FixedWindows(5), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.ParDo(TimerEmittingStatefulDoFn()) | beam.ParDo(self.record_dofn())) self.assertEqual([('timer1-mykey', 10, 10, 15)], sorted(StatefulDoFnOnDirectRunnerTest.all_records)) def test_timer_default_tag(self): class DynamicTimerDoFn(DoFn): EMIT_TIMER_FAMILY = TimerSpec('emit', TimeDomain.WATERMARK) def process(self, element, emit=DoFn.TimerParam(EMIT_TIMER_FAMILY)): emit.set(10) emit.set(20, dynamic_timer_tag='') @on_timer(EMIT_TIMER_FAMILY) def emit_callback( self, ts=DoFn.TimestampParam, tag=DoFn.DynamicTimerTagParam): yield (tag, ts) with TestPipeline() as p: test_stream = (TestStream().advance_watermark_to(10).add_elements( [1])).advance_watermark_to_infinity() ( p | test_stream | beam.Map(lambda x: ('mykey', x)) | beam.ParDo(DynamicTimerDoFn()) | beam.ParDo(self.record_dofn())) self.assertEqual([('', 20)], sorted(StatefulDoFnOnDirectRunnerTest.all_records)) def test_dynamic_timer_simple_dofn(self): class DynamicTimerDoFn(DoFn): EMIT_TIMER_FAMILY = TimerSpec('emit', TimeDomain.WATERMARK) def process(self, element, emit=DoFn.TimerParam(EMIT_TIMER_FAMILY)): emit.set(10, dynamic_timer_tag='emit1') emit.set(20, dynamic_timer_tag='emit2') emit.set(30, dynamic_timer_tag='emit3') @on_timer(EMIT_TIMER_FAMILY) def emit_callback( self, ts=DoFn.TimestampParam, tag=DoFn.DynamicTimerTagParam): yield (tag, ts) with TestPipeline() as p: test_stream = (TestStream().advance_watermark_to(10).add_elements( [1])).advance_watermark_to_infinity() ( p | test_stream | beam.Map(lambda x: ('mykey', x)) | beam.ParDo(DynamicTimerDoFn()) | beam.ParDo(self.record_dofn())) self.assertEqual([('emit1', 10), ('emit2', 20), ('emit3', 30)], sorted(StatefulDoFnOnDirectRunnerTest.all_records)) @pytest.mark.no_xdist @pytest.mark.timeout(60) def test_dynamic_timer_clear_then_set_timer(self): class EmitTwoEvents(DoFn): EMIT_CLEAR_SET_TIMER = TimerSpec('emitclear', TimeDomain.WATERMARK) def process(self, element, emit=DoFn.TimerParam(EMIT_CLEAR_SET_TIMER)): yield ('1', 'set') emit.set(1) @on_timer(EMIT_CLEAR_SET_TIMER) def emit_clear(self): yield ('1', 'clear') class DynamicTimerDoFn(DoFn): EMIT_TIMER_FAMILY = TimerSpec('emit', TimeDomain.WATERMARK) def process(self, element, emit=DoFn.TimerParam(EMIT_TIMER_FAMILY)): if element[1] == 'set': emit.set(10, dynamic_timer_tag='emit1') emit.set(20, dynamic_timer_tag='emit2') if element[1] == 'clear': emit.set(30, dynamic_timer_tag='emit3') emit.clear(dynamic_timer_tag='emit3') emit.set(40, dynamic_timer_tag='emit3') return [] @on_timer(EMIT_TIMER_FAMILY) def emit_callback( self, ts=DoFn.TimestampParam, tag=DoFn.DynamicTimerTagParam): yield (tag, ts) with TestPipeline() as p: res = ( p | beam.Create([('1', 'impulse')]) | beam.ParDo(EmitTwoEvents()) | beam.ParDo(DynamicTimerDoFn())) assert_that(res, equal_to([('emit1', 10), ('emit2', 20), ('emit3', 40)])) def test_dynamic_timer_clear_timer(self): class DynamicTimerDoFn(DoFn): EMIT_TIMER_FAMILY = TimerSpec('emit', TimeDomain.WATERMARK) def process(self, element, emit=DoFn.TimerParam(EMIT_TIMER_FAMILY)): if element[1] == 'set': emit.set(10, dynamic_timer_tag='emit1') emit.set(20, dynamic_timer_tag='emit2') emit.set(30, dynamic_timer_tag='emit3') if element[1] == 'clear': emit.clear(dynamic_timer_tag='emit3') @on_timer(EMIT_TIMER_FAMILY) def emit_callback( self, ts=DoFn.TimestampParam, tag=DoFn.DynamicTimerTagParam): yield (tag, ts) with TestPipeline() as p: test_stream = ( TestStream().advance_watermark_to(5).add_elements( ['set']).advance_watermark_to(10).add_elements( ['clear']).advance_watermark_to_infinity()) ( p | test_stream | beam.Map(lambda x: ('mykey', x)) | beam.ParDo(DynamicTimerDoFn()) | beam.ParDo(self.record_dofn())) self.assertEqual([('emit1', 10), ('emit2', 20)], sorted(StatefulDoFnOnDirectRunnerTest.all_records)) def test_dynamic_timer_multiple(self): class DynamicTimerDoFn(DoFn): EMIT_TIMER_FAMILY1 = TimerSpec('emit_family_1', TimeDomain.WATERMARK) EMIT_TIMER_FAMILY2 = TimerSpec('emit_family_2', TimeDomain.WATERMARK) def process( self, element, emit1=DoFn.TimerParam(EMIT_TIMER_FAMILY1), emit2=DoFn.TimerParam(EMIT_TIMER_FAMILY2)): emit1.set(10, dynamic_timer_tag='emit11') emit1.set(20, dynamic_timer_tag='emit12') emit1.set(30, dynamic_timer_tag='emit13') emit2.set(30, dynamic_timer_tag='emit21') emit2.set(20, dynamic_timer_tag='emit22') emit2.set(10, dynamic_timer_tag='emit23') @on_timer(EMIT_TIMER_FAMILY1) def emit_callback( self, ts=DoFn.TimestampParam, tag=DoFn.DynamicTimerTagParam): yield (tag, ts) @on_timer(EMIT_TIMER_FAMILY2) def emit_callback_2( self, ts=DoFn.TimestampParam, tag=DoFn.DynamicTimerTagParam): yield (tag, ts) with TestPipeline() as p: test_stream = ( TestStream().advance_watermark_to(5).add_elements( ['1']).advance_watermark_to_infinity()) ( p | test_stream | beam.Map(lambda x: ('mykey', x)) | beam.ParDo(DynamicTimerDoFn()) | beam.ParDo(self.record_dofn())) self.assertEqual([('emit11', 10), ('emit12', 20), ('emit13', 30), ('emit21', 30), ('emit22', 20), ('emit23', 10)], sorted(StatefulDoFnOnDirectRunnerTest.all_records)) def test_dynamic_timer_and_simple_timer(self): class DynamicTimerDoFn(DoFn): EMIT_TIMER_FAMILY = TimerSpec('emit', TimeDomain.WATERMARK) GC_TIMER = TimerSpec('gc', TimeDomain.WATERMARK) def process( self, element, emit=DoFn.TimerParam(EMIT_TIMER_FAMILY), gc=DoFn.TimerParam(GC_TIMER)): emit.set(10, dynamic_timer_tag='emit1') emit.set(20, dynamic_timer_tag='emit2') emit.set(30, dynamic_timer_tag='emit3') gc.set(40) @on_timer(EMIT_TIMER_FAMILY) def emit_callback( self, ts=DoFn.TimestampParam, tag=DoFn.DynamicTimerTagParam): yield (tag, ts) @on_timer(GC_TIMER) def gc(self, ts=DoFn.TimestampParam): yield ('gc', ts) with TestPipeline() as p: test_stream = ( TestStream().advance_watermark_to(5).add_elements( ['1']).advance_watermark_to_infinity()) ( p | test_stream | beam.Map(lambda x: ('mykey', x)) | beam.ParDo(DynamicTimerDoFn()) | beam.ParDo(self.record_dofn())) self.assertEqual([('emit1', 10), ('emit2', 20), ('emit3', 30), ('gc', 40)], sorted(StatefulDoFnOnDirectRunnerTest.all_records)) def test_index_assignment(self): class IndexAssigningStatefulDoFn(DoFn): INDEX_STATE = CombiningValueStateSpec('index', sum) def process(self, element, state=DoFn.StateParam(INDEX_STATE)): unused_key, value = element current_index = state.read() yield (value, current_index) state.add(1) with TestPipeline() as p: test_stream = ( TestStream().advance_watermark_to(10).add_elements([ 'A', 'B' ]).add_elements(['C']).advance_watermark_to(25).add_elements(['D'])) ( p | test_stream | beam.Map(lambda x: ('mykey', x)) | beam.ParDo(IndexAssigningStatefulDoFn()) | beam.ParDo(self.record_dofn())) self.assertEqual([('A', 0), ('B', 1), ('C', 2), ('D', 3)], StatefulDoFnOnDirectRunnerTest.all_records) def test_hash_join(self): class HashJoinStatefulDoFn(DoFn): BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) UNMATCHED_TIMER = TimerSpec('unmatched', TimeDomain.WATERMARK) def process( self, element, state=DoFn.StateParam(BUFFER_STATE), timer=DoFn.TimerParam(UNMATCHED_TIMER)): key, value = element existing_values = list(state.read()) if not existing_values: state.add(value) timer.set(100) else: yield b'Record<%s,%s,%s>' % (key, existing_values[0], value) state.clear() timer.clear() @on_timer(UNMATCHED_TIMER) def expiry_callback(self, state=DoFn.StateParam(BUFFER_STATE)): buffered = list(state.read()) assert len(buffered) == 1, buffered state.clear() yield b'Unmatched<%s>' % (buffered[0], ) with TestPipeline() as p: test_stream = ( TestStream().advance_watermark_to(10).add_elements([ (b'A', b'a'), (b'B', b'b') ]).add_elements([ (b'A', b'aa'), (b'C', b'c') ]).advance_watermark_to(25).add_elements([ (b'A', b'aaa'), (b'B', b'bb') ]).add_elements([ (b'D', b'd'), (b'D', b'dd'), (b'D', b'ddd'), (b'D', b'dddd') ]).advance_watermark_to(125).add_elements([(b'C', b'cc')])) ( p | test_stream | beam.ParDo(HashJoinStatefulDoFn()) | beam.ParDo(self.record_dofn())) equal_to(StatefulDoFnOnDirectRunnerTest.all_records)([ b'Record', b'Record', b'Record', b'Record', b'Unmatched', b'Unmatched', b'Unmatched' ]) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/util.py0000644000175100017510000027537215206113033021423 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Simple utility PTransforms. """ # pytype: skip-file import bisect import collections import contextlib import hashlib import hmac import logging import random import re import threading import time import uuid from collections.abc import Callable from collections.abc import Iterable from typing import TYPE_CHECKING from typing import Any from typing import Optional from typing import TypeVar from typing import Union from cryptography.fernet import Fernet import apache_beam as beam from apache_beam import coders from apache_beam import pvalue from apache_beam import typehints from apache_beam.metrics import Metrics from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.pvalue import AsSideInput from apache_beam.pvalue import PCollection from apache_beam.transforms import window from apache_beam.transforms.combiners import CountCombineFn from apache_beam.transforms.core import CombinePerKey from apache_beam.transforms.core import Create from apache_beam.transforms.core import DoFn from apache_beam.transforms.core import FlatMap from apache_beam.transforms.core import Flatten from apache_beam.transforms.core import GroupByKey from apache_beam.transforms.core import Map from apache_beam.transforms.core import MapTuple from apache_beam.transforms.core import ParDo from apache_beam.transforms.core import Windowing from apache_beam.transforms.ptransform import PTransform from apache_beam.transforms.ptransform import ptransform_fn from apache_beam.transforms.timeutil import TimeDomain from apache_beam.transforms.trigger import AccumulationMode from apache_beam.transforms.trigger import Always from apache_beam.transforms.userstate import BagStateSpec from apache_beam.transforms.userstate import CombiningValueStateSpec from apache_beam.transforms.userstate import ReadModifyWriteStateSpec from apache_beam.transforms.userstate import TimerSpec from apache_beam.transforms.userstate import on_timer from apache_beam.transforms.window import NonMergingWindowFn from apache_beam.transforms.window import TimestampCombiner from apache_beam.transforms.window import TimestampedValue from apache_beam.typehints import trivial_inference from apache_beam.typehints.decorators import get_signature from apache_beam.typehints.native_type_compatibility import TypedWindowedValue from apache_beam.typehints.sharded_key_type import ShardedKeyType from apache_beam.utils import shared from apache_beam.utils import windowed_value from apache_beam.utils.annotations import deprecated from apache_beam.utils.sharded_key import ShardedKey from apache_beam.utils.timestamp import Timestamp if TYPE_CHECKING: from apache_beam.runners.pipeline_context import PipelineContext _LOGGER = logging.getLogger(__name__) __all__ = [ 'BatchElements', 'CoGroupByKey', 'Distinct', 'GcpSecret', 'GroupByEncryptedKey', 'Keys', 'KvSwap', 'LogElements', 'Regex', 'Reify', 'RemoveDuplicates', 'Reshuffle', 'Secret', 'SortAndBatchElements', 'ToString', 'Tee', 'Values', 'WithKeys', 'GroupIntoBatches', 'WaitOn' ] K = TypeVar('K') V = TypeVar('V') T = TypeVar('T') U = TypeVar('U') RESHUFFLE_TYPEHINT_BREAKING_CHANGE_VERSION = "2.64.0" class CoGroupByKey(PTransform): """Groups results across several PCollections by key. Given an input dict of serializable keys (called "tags") to 0 or more PCollections of (key, value) tuples, it creates a single output PCollection of (key, value) tuples whose keys are the unique input keys from all inputs, and whose values are dicts mapping each tag to an iterable of whatever values were under the key in the corresponding PCollection, in this manner:: ('some key', {'tag1': ['value 1 under "some key" in pcoll1', 'value 2 under "some key" in pcoll1', ...], 'tag2': ... , ... }) where `[]` refers to an iterable, not a list. For example, given:: {'tag1': pc1, 'tag2': pc2, 333: pc3} where:: pc1 = beam.Create([(k1, v1)])) pc2 = beam.Create([]) pc3 = beam.Create([(k1, v31), (k1, v32), (k2, v33)]) The output PCollection would consist of items:: [(k1, {'tag1': [v1], 'tag2': [], 333: [v31, v32]}), (k2, {'tag1': [], 'tag2': [], 333: [v33]})] where `[]` refers to an iterable, not a list. CoGroupByKey also works for tuples, lists, or other flat iterables of PCollections, in which case the values of the resulting PCollections will be tuples whose nth value is the iterable of values from the nth PCollection---conceptually, the "tags" are the indices into the input. Thus, for this input:: (pc1, pc2, pc3) the output would be:: [(k1, ([v1], [], [v31, v32]), (k2, ([], [], [v33]))] where, again, `[]` refers to an iterable, not a list. Attributes: **kwargs: Accepts a single named argument "pipeline", which specifies the pipeline that "owns" this PTransform. Ordinarily CoGroupByKey can obtain this information from one of the input PCollections, but if there are none (or if there's a chance there may be none), this argument is the only way to provide pipeline information, and should be considered mandatory. """ def __init__(self, *, pipeline=None): self.pipeline = pipeline def _extract_input_pvalues(self, pvalueish): try: # If this works, it's a dict. return pvalueish, tuple(pvalueish.values()) except AttributeError: # Cast iterables a tuple so we can do re-iteration. pcolls = tuple(pvalueish) return pcolls, pcolls def expand(self, pcolls): if not pcolls: pcolls = (self.pipeline | Create([]), ) if isinstance(pcolls, dict): tags = list(pcolls.keys()) if all(isinstance(tag, str) and len(tag) < 10 for tag in tags): # Small, string tags. Pass them as data. pcolls_dict = pcolls restore_tags = None else: # Pass the tags in the restore_tags closure. tags = list(pcolls.keys()) pcolls_dict = {str(ix): pcolls[tag] for (ix, tag) in enumerate(tags)} restore_tags = lambda vs: { tag: vs[str(ix)] for (ix, tag) in enumerate(tags) } else: # Tags are tuple indices. tags = [str(ix) for ix in range(len(pcolls))] pcolls_dict = dict(zip(tags, pcolls)) restore_tags = lambda vs: tuple(vs[tag] for tag in tags) input_key_types = [] input_value_types = [] for pcoll in pcolls_dict.values(): key_type, value_type = typehints.trivial_inference.key_value_types( pcoll.element_type) input_key_types.append(key_type) input_value_types.append(value_type) output_key_type = typehints.Union[tuple(input_key_types)] iterable_input_value_types = tuple( typehints.Iterable[t] for t in input_value_types) output_value_type = typehints.Dict[ str, typehints.Union[iterable_input_value_types or [typehints.Any]]] result = ( pcolls_dict | 'CoGroupByKeyImpl' >> _CoGBKImpl(pipeline=self.pipeline).with_output_types( typehints.Tuple[output_key_type, output_value_type])) if restore_tags: if isinstance(pcolls, dict): dict_key_type = typehints.Union[tuple( trivial_inference.instance_to_type(tag) for tag in tags)] output_value_type = typehints.Dict[ dict_key_type, typehints.Union[iterable_input_value_types]] else: output_value_type = typehints.Tuple[iterable_input_value_types] result |= 'RestoreTags' >> MapTuple( lambda k, vs: (k, restore_tags(vs))).with_output_types( typehints.Tuple[output_key_type, output_value_type]) return result class _CoGBKImpl(PTransform): def __init__(self, *, pipeline=None): self.pipeline = pipeline def expand(self, pcolls): # Check input PCollections for PCollection-ness, and that they all belong # to the same pipeline. for pcoll in pcolls.values(): self._check_pcollection(pcoll) if self.pipeline: assert pcoll.pipeline == self.pipeline, ( 'All input PCollections must belong to the same pipeline.') tags = list(pcolls.keys()) def add_tag(tag): return lambda k, v: (k, (tag, v)) def collect_values(key, tagged_values): grouped_values = {tag: [] for tag in tags} for tag, value in tagged_values: grouped_values[tag].append(value) return key, grouped_values return ([ pcoll | 'Tag[%s]' % tag >> MapTuple(add_tag(tag)) for (tag, pcoll) in pcolls.items() ] | Flatten(pipeline=self.pipeline) | GroupByKey() | MapTuple(collect_values).with_input_types( tuple[K, Iterable[tuple[U, V]]]).with_output_types( tuple[K, dict[U, list[V]]])) @ptransform_fn @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types(K) def Keys(pcoll, label='Keys'): # pylint: disable=invalid-name """Produces a PCollection of first elements of 2-tuples in a PCollection.""" return pcoll | label >> MapTuple(lambda k, _: k) @ptransform_fn @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types(V) def Values(pcoll, label='Values'): # pylint: disable=invalid-name """Produces a PCollection of second elements of 2-tuples in a PCollection.""" return pcoll | label >> MapTuple(lambda _, v: v) @ptransform_fn @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types(tuple[V, K]) def KvSwap(pcoll, label='KvSwap'): # pylint: disable=invalid-name """Produces a PCollection reversing 2-tuples in a PCollection.""" return pcoll | label >> MapTuple(lambda k, v: (v, k)) @ptransform_fn @typehints.with_input_types(T) @typehints.with_output_types(T) def Distinct(pcoll): # pylint: disable=invalid-name """Produces a PCollection containing distinct elements of a PCollection.""" return ( pcoll | 'ToPairs' >> Map(lambda v: (v, None)) | 'Group' >> CombinePerKey(lambda vs: None) | 'Distinct' >> Keys()) @deprecated(since='2.12', current='Distinct') @ptransform_fn @typehints.with_input_types(T) @typehints.with_output_types(T) def RemoveDuplicates(pcoll): """Produces a PCollection containing distinct elements of a PCollection.""" return pcoll | 'RemoveDuplicates' >> Distinct() class Secret(): """A secret management class used for handling sensitive data. This class provides a generic interface for secret management. Implementations of this class should handle fetching secrets from a secret management system. """ def get_secret_bytes(self) -> bytes: """Returns the secret as a byte string.""" raise NotImplementedError() @staticmethod def generate_secret_bytes() -> bytes: """Generates a new secret key.""" return Fernet.generate_key() @staticmethod def parse_secret_option(secret) -> 'Secret': """Parses a secret string and returns the appropriate secret type. The secret string should be formatted like: 'type:;:' For example, 'type:GcpSecret;version_name:my_secret/versions/latest' would return a GcpSecret initialized with 'my_secret/versions/latest'. """ param_map = {} for param in secret.split(';'): parts = param.split(':') param_map[parts[0]] = parts[1] if 'type' not in param_map: raise ValueError('Secret string must contain a valid type parameter') secret_type = param_map['type'].lower() del param_map['type'] secret_class = Secret secret_params = None if secret_type == 'gcpsecret': secret_class = GcpSecret # type: ignore[assignment] secret_params = ['version_name'] elif secret_type == 'gcphsmgeneratedsecret': secret_class = GcpHsmGeneratedSecret # type: ignore[assignment] secret_params = [ 'project_id', 'location_id', 'key_ring_id', 'key_id', 'job_name' ] else: raise ValueError( f'Invalid secret type {secret_type}, currently only ' 'GcpSecret and GcpHsmGeneratedSecret are supported') for param_name in param_map.keys(): if param_name not in secret_params: raise ValueError( f'Invalid secret parameter {param_name}, ' f'{secret_type} only supports the following ' f'parameters: {secret_params}') return secret_class(**param_map) class GcpSecret(Secret): """A secret manager implementation that retrieves secrets from Google Cloud Secret Manager. """ def __init__(self, version_name: str): """Initializes a GcpSecret object. Args: version_name: The full version name of the secret in Google Cloud Secret Manager. For example: projects//secrets//versions/1. For more info, see https://cloud.google.com/python/docs/reference/secretmanager/latest/google.cloud.secretmanager_v1beta1.services.secret_manager_service.SecretManagerServiceClient#google_cloud_secretmanager_v1beta1_services_secret_manager_service_SecretManagerServiceClient_access_secret_version """ self._version_name = version_name def get_secret_bytes(self) -> bytes: try: from google.cloud import secretmanager client = secretmanager.SecretManagerServiceClient() response = client.access_secret_version( request={"name": self._version_name}) secret = response.payload.data return secret except Exception as e: raise RuntimeError( 'Failed to retrieve secret bytes for secret ' f'{self._version_name} with exception {e}') def __eq__(self, secret): return self._version_name == getattr(secret, '_version_name', None) class GcpHsmGeneratedSecret(Secret): """A secret manager implementation that generates a secret using a GCP HSM key and stores it in Google Cloud Secret Manager. If the secret already exists, it will be retrieved. """ def __init__( self, project_id: str, location_id: str, key_ring_id: str, key_id: str, job_name: str): """Initializes a GcpHsmGeneratedSecret object. Args: project_id: The GCP project ID. location_id: The GCP location ID for the HSM key. key_ring_id: The ID of the KMS key ring. key_id: The ID of the KMS key. job_name: The name of the job, used to generate a unique secret name. """ self._project_id = project_id self._location_id = location_id self._key_ring_id = key_ring_id self._key_id = key_id self._secret_version_name = f'HsmGeneratedSecret_{job_name}' def get_secret_bytes(self) -> bytes: """Retrieves the secret bytes. If the secret version already exists in Secret Manager, it is retrieved. Otherwise, a new secret and version are created. The new secret is generated using the HSM key. Returns: The secret as a byte string. """ try: from google.api_core import exceptions as api_exceptions from google.cloud import secretmanager client = secretmanager.SecretManagerServiceClient() project_path = f"projects/{self._project_id}" secret_path = f"{project_path}/secrets/{self._secret_version_name}" # Since we may generate multiple versions when doing this on workers, # just always take the first version added to maintain consistency. secret_version_path = f"{secret_path}/versions/1" try: response = client.access_secret_version( request={"name": secret_version_path}) return response.payload.data except api_exceptions.NotFound: # Don't bother logging yet, we'll only log if we actually add the # secret version below pass try: client.create_secret( request={ "parent": project_path, "secret_id": self._secret_version_name, "secret": { "replication": { "automatic": {} } }, }) except api_exceptions.AlreadyExists: # Don't bother logging yet, we'll only log if we actually add the # secret version below pass new_key = self.generate_dek() try: # Try one more time in case it was created while we were generating the # DEK. response = client.access_secret_version( request={"name": secret_version_path}) return response.payload.data except api_exceptions.NotFound: _LOGGER.info( "Secret version %s not found. " "Creating new secret and version.", secret_version_path) client.add_secret_version( request={ "parent": secret_path, "payload": { "data": new_key } }) response = client.access_secret_version( request={"name": secret_version_path}) return response.payload.data except Exception as e: raise RuntimeError( f'Failed to retrieve or create secret bytes for secret ' f'{self._secret_version_name} with exception {e}') def generate_dek(self, dek_size: int = 32) -> bytes: """Generates a new Data Encryption Key (DEK) using an HSM-backed key. This function follows a key derivation process that incorporates entropy from the HSM-backed key into the nonce used for key derivation. Args: dek_size: The size of the DEK to generate. Returns: A new DEK of the specified size, url-safe base64-encoded. """ try: import base64 import os from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.hkdf import HKDF from google.cloud import kms # 1. Generate a random nonce (nonce_one) nonce_one = os.urandom(dek_size) # 2. Use the HSM-backed key to encrypt nonce_one to create nonce_two kms_client = kms.KeyManagementServiceClient() key_path = kms_client.crypto_key_path( self._project_id, self._location_id, self._key_ring_id, self._key_id) response = kms_client.encrypt( request={ 'name': key_path, 'plaintext': nonce_one }) nonce_two = response.ciphertext # 3. Generate a Derivation Key (DK) dk = os.urandom(dek_size) # 4. Use a KDF to derive the DEK using DK and nonce_two hkdf = HKDF( algorithm=hashes.SHA256(), length=dek_size, salt=nonce_two, info=None, ) dek = hkdf.derive(dk) return base64.urlsafe_b64encode(dek) except Exception as e: raise RuntimeError(f'Failed to generate DEK with exception {e}') class _EncryptMessage(DoFn): """A DoFn that encrypts the key and value of each element.""" def __init__( self, hmac_key_secret: Secret, key_coder: coders.Coder, value_coder: coders.Coder): self.hmac_key_secret = hmac_key_secret self.key_coder = key_coder self.value_coder = value_coder def setup(self): self._hmac_key = self.hmac_key_secret.get_secret_bytes() self.fernet = Fernet(self._hmac_key) def process(self, element: Any) -> Iterable[tuple[bytes, tuple[bytes, bytes]]]: """Encrypts the key and value of an element. Args: element: A tuple containing the key and value to be encrypted. Yields: A tuple containing the HMAC of the encoded key, and a tuple of the encrypted key and value. """ k, v = element encoded_key = self.key_coder.encode(k) encoded_value = self.value_coder.encode(v) hmac_encoded_key = hmac.new(self._hmac_key, encoded_key, hashlib.sha256).digest() out_element = ( hmac_encoded_key, (self.fernet.encrypt(encoded_key), self.fernet.encrypt(encoded_value))) yield out_element class _DecryptMessage(DoFn): """A DoFn that decrypts the key and value of each element.""" def __init__( self, hmac_key_secret: Secret, key_coder: coders.Coder, value_coder: coders.Coder): self.hmac_key_secret = hmac_key_secret self.key_coder = key_coder self.value_coder = value_coder def setup(self): hmac_key = self.hmac_key_secret.get_secret_bytes() self.fernet = Fernet(hmac_key) def decode_value(self, encoded_element: tuple[bytes, bytes]) -> Any: encrypted_value = encoded_element[1] encoded_value = self.fernet.decrypt(encrypted_value) real_val = self.value_coder.decode(encoded_value) return real_val def filter_elements_by_key( self, encrypted_key: bytes, encoded_elements: Iterable[tuple[bytes, bytes]]) -> Iterable[Any]: for e in encoded_elements: if encrypted_key == self.fernet.decrypt(e[0]): yield self.decode_value(e) # Right now, GBK always returns a list of elements, so we match this behavior # here. This does mean that the whole list will be materialized every time, # but passing an Iterable containing an Iterable breaks when pickling happens def process( self, element: tuple[bytes, Iterable[tuple[bytes, bytes]]] ) -> Iterable[tuple[Any, list[Any]]]: """Decrypts the key and values of an element. Args: element: A tuple containing the HMAC of the encoded key and an iterable of tuples of encrypted keys and values. Yields: A tuple containing the decrypted key and a list of decrypted values. """ unused_hmac_encoded_key, encoded_elements = element seen_keys = set() # Since there could be hmac collisions, we will use the fernet encrypted # key to confirm that the mapping is actually correct. for e in encoded_elements: encrypted_key, unused_encrypted_value = e encoded_key = self.fernet.decrypt(encrypted_key) if encoded_key in seen_keys: continue seen_keys.add(encoded_key) real_key = self.key_coder.decode(encoded_key) yield ( real_key, list(self.filter_elements_by_key(encoded_key, encoded_elements))) @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types(tuple[K, Iterable[V]]) class GroupByEncryptedKey(PTransform): """A PTransform that provides a secure alternative to GroupByKey. This transform encrypts the keys of the input PCollection, performs a GroupByKey on the encrypted keys, and then decrypts the keys in the output. This is useful when the keys contain sensitive data that should not be stored at rest by the runner. Note the following caveats: 1) Runners can implement arbitrary materialization steps, so this does not guarantee that the whole pipeline will not have unencrypted data at rest by itself. 2) If using this transform in streaming mode, this transform may not properly handle update compatibility checks around coders. This means that an improper update could lead to invalid coders, causing pipeline failure or data corruption. If you need to update, make sure that the input type passed into this transform does not change. """ def __init__(self, hmac_key: Secret): """Initializes a GroupByEncryptedKey transform. Args: hmac_key: A Secret object that provides the secret key for HMAC and encryption. For example, a GcpSecret can be used to access a secret stored in GCP Secret Manager """ self._hmac_key = hmac_key def expand(self, pcoll): key_type, value_type = (typehints.typehints.coerce_to_kv_type( pcoll.element_type).tuple_types) kv_type_hint = typehints.KV[key_type, value_type] if kv_type_hint and kv_type_hint != typehints.Any: coder = coders.registry.get_coder(kv_type_hint) try: coder = coder.as_deterministic_coder(self.label) except ValueError: _LOGGER.warning( 'GroupByEncryptedKey %s: ' 'The key coder is not deterministic. This may result in incorrect ' 'pipeline output. This can be fixed by adding a type hint to the ' 'operation preceding the GroupByKey step, and for custom key ' 'classes, by writing a deterministic custom Coder. Please see the ' 'documentation for more details.', self.label) if not coder.is_kv_coder(): raise ValueError( 'Input elements to the transform %s with stateful DoFn must be ' 'key-value pairs.' % self) key_coder = coder.key_coder() value_coder = coder.value_coder() else: key_coder = coders.registry.get_coder(typehints.Any) value_coder = key_coder gbk = beam.GroupByKey() gbk._inside_gbek = True output_type = tuple[key_type, Iterable[value_type]] return ( pcoll | beam.ParDo(_EncryptMessage(self._hmac_key, key_coder, value_coder)) | gbk | beam.ParDo( _DecryptMessage(self._hmac_key, key_coder, value_coder)).with_output_types(output_type)) class _BatchSizeEstimator(object): """Estimates the best size for batches given historical timing. """ _MAX_DATA_POINTS = 100 _MAX_GROWTH_FACTOR = 2 def __init__( self, min_batch_size=1, max_batch_size=10000, target_batch_overhead=.05, target_batch_duration_secs=10, target_batch_duration_secs_including_fixed_cost=None, variance=0.25, clock=time.time, ignore_first_n_seen_per_batch_size=0, record_metrics=True): if min_batch_size > max_batch_size: raise ValueError( "Minimum (%s) must not be greater than maximum (%s)" % (min_batch_size, max_batch_size)) if target_batch_overhead and not 0 < target_batch_overhead <= 1: raise ValueError( "target_batch_overhead (%s) must be between 0 and 1" % (target_batch_overhead)) if target_batch_duration_secs and target_batch_duration_secs <= 0: raise ValueError( "target_batch_duration_secs (%s) must be positive" % (target_batch_duration_secs)) if (target_batch_duration_secs_including_fixed_cost and target_batch_duration_secs_including_fixed_cost <= 0): raise ValueError( "target_batch_duration_secs_including_fixed_cost " "(%s) must be positive" % (target_batch_duration_secs_including_fixed_cost)) if not (target_batch_overhead or target_batch_duration_secs or target_batch_duration_secs_including_fixed_cost): raise ValueError( "At least one of target_batch_overhead or " "target_batch_duration_secs or " "target_batch_duration_secs_including_fixed_cost must be positive.") if ignore_first_n_seen_per_batch_size < 0: raise ValueError( 'ignore_first_n_seen_per_batch_size (%s) must be non ' 'negative' % (ignore_first_n_seen_per_batch_size)) self._min_batch_size = min_batch_size self._max_batch_size = max_batch_size self._target_batch_overhead = target_batch_overhead self._target_batch_duration_secs = target_batch_duration_secs self._target_batch_duration_secs_including_fixed_cost = ( target_batch_duration_secs_including_fixed_cost) self._variance = variance self._clock = clock self._data = [] self._ignore_next_timing = False self._ignore_first_n_seen_per_batch_size = ( ignore_first_n_seen_per_batch_size) self._batch_size_num_seen = {} self._replay_last_batch_size = None self._record_metrics = record_metrics self._element_count = 0 self._batch_count = 0 if record_metrics: self._size_distribution = Metrics.distribution( 'BatchElements', 'batch_size') self._time_distribution = Metrics.distribution( 'BatchElements', 'msec_per_batch') else: self._size_distribution = self._time_distribution = None # Beam distributions only accept integer values, so we use this to # accumulate under-reported values until they add up to whole milliseconds. # (Milliseconds are chosen because that's conventionally used elsewhere in # profiling-style counters.) self._remainder_msecs = 0 def ignore_next_timing(self): """Call to indicate the next timing should be ignored. For example, the first emit of a ParDo operation is known to be anomalous due to setup that may occur. """ self._ignore_next_timing = True @contextlib.contextmanager def record_time(self, batch_size): start = self._clock() yield elapsed = float(self._clock() - start) elapsed_msec = 1e3 * elapsed + self._remainder_msecs if self._record_metrics: self._size_distribution.update(batch_size) self._time_distribution.update(int(elapsed_msec)) self._element_count += batch_size self._batch_count += 1 self._remainder_msecs = elapsed_msec - int(elapsed_msec) # If we ignore the next timing, replay the batch size to get accurate # timing. if self._ignore_next_timing: self._ignore_next_timing = False self._replay_last_batch_size = min(batch_size, self._max_batch_size) else: self._data.append((batch_size, elapsed)) if len(self._data) >= self._MAX_DATA_POINTS: self._thin_data() def _thin_data(self): # Make sure we don't change the parity of len(self._data) # As it's used below to alternate jitter. self._data.pop(random.randrange(len(self._data) // 4)) self._data.pop(random.randrange(len(self._data) // 2)) @staticmethod def linear_regression_no_numpy(xs, ys): # Least squares fit for y = a + bx over all points. n = float(len(xs)) xbar = sum(xs) / n ybar = sum(ys) / n if xbar == 0: return ybar, 0 if all(xs[0] == x for x in xs): # Simply use the mean if all values in xs are same. return 0, ybar / xbar b = ( sum([(x - xbar) * (y - ybar) for x, y in zip(xs, ys)]) / sum([(x - xbar)**2 for x in xs])) a = ybar - b * xbar return a, b @staticmethod def linear_regression_numpy(xs, ys): # pylint: disable=wrong-import-order, wrong-import-position import numpy as np from numpy import sum n = len(xs) if all(xs[0] == x for x in xs): # If all values of xs are same then fallback to linear_regression_no_numpy return _BatchSizeEstimator.linear_regression_no_numpy(xs, ys) xs = np.asarray(xs, dtype=float) ys = np.asarray(ys, dtype=float) # First do a simple least squares fit for y = a + bx over all points. b, a = np.polyfit(xs, ys, 1) if n < 10: return a, b else: # Refine this by throwing out outliers, according to Cook's distance. # https://en.wikipedia.org/wiki/Cook%27s_distance sum_x = sum(xs) sum_x2 = sum(xs**2) errs = a + b * xs - ys s2 = sum(errs**2) / (n - 2) if s2 == 0: # It's an exact fit! return a, b h = (sum_x2 - 2 * sum_x * xs + n * xs**2) / (n * sum_x2 - sum_x**2) cook_ds = 0.5 / s2 * errs**2 * (h / (1 - h)**2) # Re-compute the regression, excluding those points with Cook's distance # greater than 0.5, and weighting by the inverse of x to give a more # stable y-intercept (as small batches have relatively more information # about the fixed overhead). weight = (cook_ds <= 0.5) / xs b, a = np.polyfit(xs, ys, 1, w=weight) return a, b try: # pylint: disable=wrong-import-order, wrong-import-position import numpy as np linear_regression = linear_regression_numpy except ImportError: linear_regression = linear_regression_no_numpy def _calculate_next_batch_size(self): if self._min_batch_size == self._max_batch_size: return self._min_batch_size elif len(self._data) < 1: return self._min_batch_size elif len(self._data) < 2: # Force some variety so we have distinct batch sizes on which to do # linear regression below. return int( max( min( self._max_batch_size, self._min_batch_size * self._MAX_GROWTH_FACTOR), self._min_batch_size + 1)) # There tends to be a lot of noise in the top quantile, which also # has outsided influence in the regression. If we have enough data, # Simply declare the top 20% to be outliers. trimmed_data = sorted(self._data)[:max(20, len(self._data) * 4 // 5)] # Linear regression for y = a + bx, where x is batch size and y is time. xs, ys = zip(*trimmed_data) a, b = self.linear_regression(xs, ys) # Avoid nonsensical or division-by-zero errors below due to noise. a = max(a, 1e-10) b = max(b, 1e-20) last_batch_size = self._data[-1][0] cap = min(last_batch_size * self._MAX_GROWTH_FACTOR, self._max_batch_size) target = self._max_batch_size if self._target_batch_duration_secs_including_fixed_cost: # Solution to # a + b*x = self._target_batch_duration_secs_including_fixed_cost. target = min( target, (self._target_batch_duration_secs_including_fixed_cost - a) / b) if self._target_batch_duration_secs: # Solution to b*x = self._target_batch_duration_secs. # We ignore the fixed cost in this computation as it has negligeabel # impact when it is small and unhelpfully forces the minimum batch size # when it is large. target = min(target, self._target_batch_duration_secs / b) if self._target_batch_overhead: # Solution to a / (a + b*x) = self._target_batch_overhead. target = min(target, (a / b) * (1 / self._target_batch_overhead - 1)) # Avoid getting stuck at a single batch size (especially the minimal # batch size) which would not allow us to extrapolate to other batch # sizes. # Jitter alternates between 0 and 1. jitter = len(self._data) % 2 # Smear our samples across a range centered at the target. if len(self._data) > 10: target += int(target * self._variance * 2 * (random.random() - .5)) return int(max(self._min_batch_size + jitter, min(target, cap))) def next_batch_size(self): # Check if we should replay a previous batch size due to it not being # recorded. if self._replay_last_batch_size: result = self._replay_last_batch_size self._replay_last_batch_size = None else: result = self._calculate_next_batch_size() seen_count = self._batch_size_num_seen.get(result, 0) + 1 if seen_count <= self._ignore_first_n_seen_per_batch_size: self.ignore_next_timing() self._batch_size_num_seen[result] = seen_count return result def stats(self): return "element_count=%s batch_count=%s next_batch_size=%s timings=%s" % ( self._element_count, self._batch_count, self._calculate_next_batch_size(), self._data) class _GlobalWindowsBatchingDoFn(DoFn): def __init__(self, batch_size_estimator, element_size_fn): self._batch_size_estimator = batch_size_estimator self._element_size_fn = element_size_fn def start_bundle(self): self._batch = [] self._running_batch_size = 0 self._target_batch_size = self._batch_size_estimator.next_batch_size() # The first emit often involves non-trivial setup. self._batch_size_estimator.ignore_next_timing() def process(self, element): element_size = self._element_size_fn(element) if self._running_batch_size + element_size > self._target_batch_size: with self._batch_size_estimator.record_time(self._running_batch_size): yield window.GlobalWindows.windowed_value_at_end_of_window(self._batch) self._batch = [] self._running_batch_size = 0 self._target_batch_size = self._batch_size_estimator.next_batch_size() self._batch.append(element) self._running_batch_size += element_size def finish_bundle(self): if self._batch: with self._batch_size_estimator.record_time(self._running_batch_size): yield window.GlobalWindows.windowed_value_at_end_of_window(self._batch) self._batch = None self._running_batch_size = 0 self._target_batch_size = self._batch_size_estimator.next_batch_size() _LOGGER.info( "BatchElements statistics: " + self._batch_size_estimator.stats()) class _SizedBatch(): def __init__(self): self.elements = [] self.size = 0 class _WindowAwareBatchingDoFn(DoFn): _MAX_LIVE_WINDOWS = 10 def __init__(self, batch_size_estimator, element_size_fn): self._batch_size_estimator = batch_size_estimator self._element_size_fn = element_size_fn def start_bundle(self): self._batches = collections.defaultdict(_SizedBatch) self._target_batch_size = self._batch_size_estimator.next_batch_size() # The first emit often involves non-trivial setup. self._batch_size_estimator.ignore_next_timing() def process(self, element, window=DoFn.WindowParam): batch = self._batches[window] element_size = self._element_size_fn(element) if batch.size + element_size > self._target_batch_size: with self._batch_size_estimator.record_time(batch.size): yield windowed_value.WindowedValue( batch.elements, window.max_timestamp(), (window, )) del self._batches[window] self._target_batch_size = self._batch_size_estimator.next_batch_size() self._batches[window].elements.append(element) self._batches[window].size += element_size if len(self._batches) > self._MAX_LIVE_WINDOWS: window, batch = max( self._batches.items(), key=lambda window_batch: window_batch[1].size) with self._batch_size_estimator.record_time(batch.size): yield windowed_value.WindowedValue( batch.elements, window.max_timestamp(), (window, )) del self._batches[window] self._target_batch_size = self._batch_size_estimator.next_batch_size() def finish_bundle(self): for window, batch in self._batches.items(): if batch: with self._batch_size_estimator.record_time(batch.size): yield windowed_value.WindowedValue( batch.elements, window.max_timestamp(), (window, )) self._batches = None self._target_batch_size = self._batch_size_estimator.next_batch_size() def _pardo_stateful_batch_elements( input_coder: coders.Coder, batch_size_estimator: _BatchSizeEstimator, max_buffering_duration_secs: int, clock=time.time): ELEMENT_STATE = BagStateSpec('values', input_coder) COUNT_STATE = CombiningValueStateSpec('count', input_coder, CountCombineFn()) BATCH_SIZE_STATE = ReadModifyWriteStateSpec('batch_size', input_coder) WINDOW_TIMER = TimerSpec('window_end', TimeDomain.WATERMARK) BUFFERING_TIMER = TimerSpec('buffering_end', TimeDomain.REAL_TIME) BATCH_ESTIMATOR_STATE = ReadModifyWriteStateSpec( 'batch_estimator', coders.PickleCoder()) class _StatefulBatchElementsDoFn(DoFn): def process( self, element, window=DoFn.WindowParam, element_state=DoFn.StateParam(ELEMENT_STATE), count_state=DoFn.StateParam(COUNT_STATE), batch_size_state=DoFn.StateParam(BATCH_SIZE_STATE), batch_estimator_state=DoFn.StateParam(BATCH_ESTIMATOR_STATE), window_timer=DoFn.TimerParam(WINDOW_TIMER), buffering_timer=DoFn.TimerParam(BUFFERING_TIMER)): window_timer.set(window.end) # Drop the fixed key since we don't care about it element_state.add(element[1]) count_state.add(1) count = count_state.read() target_size = batch_size_state.read() # Should only happen on the first element if target_size is None: batch_estimator = batch_size_estimator target_size = batch_estimator.next_batch_size() batch_size_state.write(target_size) batch_estimator_state.write(batch_estimator) if count == 1 and max_buffering_duration_secs > 0: # First element in batch, start buffering timer buffering_timer.set(clock() + max_buffering_duration_secs) if count >= target_size: return self.flush_batch( element_state, count_state, batch_size_state, batch_estimator_state, buffering_timer) @on_timer(WINDOW_TIMER) def on_window_timer( self, element_state=DoFn.StateParam(ELEMENT_STATE), count_state=DoFn.StateParam(COUNT_STATE), batch_size_state=DoFn.StateParam(BATCH_SIZE_STATE), batch_estimator_state=DoFn.StateParam(BATCH_ESTIMATOR_STATE), buffering_timer=DoFn.TimerParam(BUFFERING_TIMER)): return self.flush_batch( element_state, count_state, batch_size_state, batch_estimator_state, buffering_timer) @on_timer(BUFFERING_TIMER) def on_buffering_timer( self, element_state=DoFn.StateParam(ELEMENT_STATE), count_state=DoFn.StateParam(COUNT_STATE), batch_size_state=DoFn.StateParam(BATCH_SIZE_STATE), batch_estimator_state=DoFn.StateParam(BATCH_ESTIMATOR_STATE), buffering_timer=DoFn.TimerParam(BUFFERING_TIMER)): return self.flush_batch( element_state, count_state, batch_size_state, batch_estimator_state, buffering_timer) def flush_batch( self, element_state, count_state, batch_size_state, batch_estimator_state, buffering_timer): batch = [element for element in element_state.read()] if not batch: return element_state.clear() count_state.clear() batch_estimator = batch_estimator_state.read() with batch_estimator.record_time(len(batch)): yield batch batch_size_state.write(batch_estimator.next_batch_size()) batch_estimator_state.write(batch_estimator) buffering_timer.clear() return _StatefulBatchElementsDoFn() class SharedKey(): """A class that holds a per-process UUID used to key elements for streaming BatchElements. """ def __init__(self): self.key = uuid.uuid4().hex def load_shared_key(): return SharedKey() class WithSharedKey(DoFn): """A DoFn that keys elements with a per-process UUID. Used in streaming BatchElements. """ def __init__(self): self.shared_handle = shared.Shared() def setup(self): self.key = self.shared_handle.acquire(load_shared_key, "WithSharedKey").key def process(self, element): yield (self.key, element) class WithLengthBucketKey(DoFn): """Keys elements with (worker_uuid, length_bucket) for length-aware stateful batching. Elements of similar length are routed to the same state partition, reducing padding waste.""" def __init__(self, length_fn, bucket_boundaries): self.shared_handle = shared.Shared() self._length_fn = length_fn self._bucket_boundaries = bucket_boundaries def setup(self): self.key = self.shared_handle.acquire( load_shared_key, "WithLengthBucketKey").key def _get_bucket(self, length): # bisect_right: boundaries are lower-inclusive. # e.g., for boundaries [10, 50], buckets are (-inf, 10), [10, 50), [50, inf) return bisect.bisect_right(self._bucket_boundaries, length) def process(self, element): length = self._length_fn(element) bucket = self._get_bucket(length) yield ((self.key, bucket), element) @typehints.with_input_types(T) @typehints.with_output_types(list[T]) class BatchElements(PTransform): """A Transform that batches elements for amortized processing. This transform is designed to precede operations whose processing cost is of the form time = fixed_cost + num_elements * per_element_cost where the per element cost is (often significantly) smaller than the fixed cost and could be amortized over multiple elements. It consumes a PCollection of element type T and produces a PCollection of element type list[T]. This transform attempts to find the best batch size between the minimim and maximum parameters by profiling the time taken by (fused) downstream operations. For a fixed batch size, set the min and max to be equal. Elements are batched per-window and batches emitted in the window corresponding to its contents. Each batch is emitted with a timestamp at the end of their window. When the max_batch_duration_secs arg is provided, a stateful implementation of BatchElements is used to batch elements across bundles. This is most impactful in streaming applications where many bundles only contain one element. Larger max_batch_duration_secs values `might` reduce the throughput of the transform, while smaller values might improve the throughput but make it more likely that batches are smaller than the target batch size. As a general recommendation, start with low values (e.g. 0.005 aka 5ms) and increase as needed to get the desired tradeoff between target batch size and latency or throughput. For more information on tuning parameters to this transform, see https://beam.apache.org/documentation/patterns/batch-elements Args: min_batch_size: (optional) the smallest size of a batch max_batch_size: (optional) the largest size of a batch target_batch_overhead: (optional) a target for fixed_cost / time, as used in the formula above target_batch_duration_secs: (optional) a target for total time per bundle, in seconds, excluding fixed cost target_batch_duration_secs_including_fixed_cost: (optional) a target for total time per bundle, in seconds, including fixed cost max_batch_duration_secs: (optional) the maximum amount of time to buffer a batch before emitting. Setting this argument to be non-none uses the stateful implementation of BatchElements. element_size_fn: (optional) A mapping of an element to its contribution to batch size, defaulting to every element having size 1. When provided, attempts to provide batches of optimal total size which may consist of a varying number of elements. variance: (optional) the permitted (relative) amount of deviation from the (estimated) ideal batch size used to produce a wider base for linear interpolation clock: (optional) an alternative to time.time for measuring the cost of donwstream operations (mostly for testing) record_metrics: (optional) whether or not to record beam metrics on distributions of the batch size. Defaults to True. length_fn: (optional) a callable mapping an element to its length (int). When set together with bucket_boundaries, enables length-aware bucketed keying on the stateful path so that elements of similar length are routed to the same batch, reducing padding waste. bucket_boundaries: (optional) a sorted list of positive boundary values for length bucketing. Boundaries are lower-inclusive (bisect_right semantics): e.g., for boundaries [10, 50], buckets are (-inf, 10), [10, 50), [50, inf). Defaults to [16, 32, 64, 128, 256, 512] when length_fn is set. Requires length_fn. """ _DEFAULT_BUCKET_BOUNDARIES = [16, 32, 64, 128, 256, 512] def __init__( self, min_batch_size=1, max_batch_size=10000, target_batch_overhead=.05, target_batch_duration_secs=10, target_batch_duration_secs_including_fixed_cost=None, max_batch_duration_secs=None, *, element_size_fn=lambda x: 1, variance=0.25, clock=time.time, record_metrics=True, length_fn=None, bucket_boundaries=None): if bucket_boundaries is not None and length_fn is None: raise ValueError('bucket_boundaries requires length_fn to be set.') if bucket_boundaries is not None: if (not bucket_boundaries or any(b <= 0 for b in bucket_boundaries) or bucket_boundaries != sorted(bucket_boundaries)): raise ValueError( 'bucket_boundaries must be a non-empty sorted list of ' 'positive values.') self._batch_size_estimator = _BatchSizeEstimator( min_batch_size=min_batch_size, max_batch_size=max_batch_size, target_batch_overhead=target_batch_overhead, target_batch_duration_secs=target_batch_duration_secs, target_batch_duration_secs_including_fixed_cost=( target_batch_duration_secs_including_fixed_cost), variance=variance, clock=clock, record_metrics=record_metrics) self._element_size_fn = element_size_fn self._max_batch_dur = max_batch_duration_secs self._clock = clock self._length_fn = length_fn if length_fn is not None and bucket_boundaries is None: self._bucket_boundaries = self._DEFAULT_BUCKET_BOUNDARIES else: self._bucket_boundaries = bucket_boundaries def expand(self, pcoll): if getattr(pcoll.pipeline.runner, 'is_streaming', False): raise NotImplementedError("Requires stateful processing (BEAM-2687)") elif self._max_batch_dur is not None: coder = coders.registry.get_coder(pcoll) if self._length_fn is not None: keying_dofn = WithLengthBucketKey( self._length_fn, self._bucket_boundaries) else: keying_dofn = WithSharedKey() return pcoll | ParDo(keying_dofn) | ParDo( _pardo_stateful_batch_elements( coder, self._batch_size_estimator, self._max_batch_dur, self._clock)) elif pcoll.windowing.is_default(): # This is the same logic as _GlobalWindowsBatchingDoFn, but optimized # for that simpler case. return pcoll | ParDo( _GlobalWindowsBatchingDoFn( self._batch_size_estimator, self._element_size_fn)) else: return pcoll | ParDo( _WindowAwareBatchingDoFn( self._batch_size_estimator, self._element_size_fn)) class _SortAndBatchElementsDoFn(DoFn): """DoFn that buffers, sorts by element size, and batches elements. This DoFn is used internally by ``SortAndBatchElements`` for PCollections with the default (global) window. It accumulates all elements in the current bundle, sorts them by size in ascending order, and emits optimally-sized batches on ``finish_bundle``. Args: min_batch_size: The minimum number of elements per batch. Must be >= 1. max_batch_size: The maximum number of elements per batch. Must be >= ``min_batch_size``. max_batch_weight: The maximum total weight of elements in a batch, where weight is computed by ``element_size_fn``. Must be >= 1. element_size_fn: An optional callable mapping an element to its integer size/weight. """ def __init__( self, min_batch_size: int, max_batch_size: int, max_batch_weight: int, element_size_fn: Optional[Callable[[Any], int]]): self._min_batch_size = min_batch_size self._max_batch_size = max_batch_size self._max_batch_weight = max_batch_weight self._element_size_fn = element_size_fn or self._default_element_size self._has_warned_type_error = False self._buffer = [] def _default_element_size(self, element): try: return len(element) except TypeError: if not self._has_warned_type_error: _LOGGER.warning( 'Element of type %s does not support len(). Falling back to ' 'size 1. Consider providing a custom element_size_fn to ' 'SortAndBatchElements for meaningful size-based batching.', type(element).__name__) self._has_warned_type_error = True return 1 def start_bundle(self): self._buffer = [] def process(self, element): self._buffer.append(element) def finish_bundle(self): if not self._buffer: return # Sort elements by size (ascending) for optimal batching # Elements of similar sizes will be grouped together sorted_elements = sorted(self._buffer, key=self._element_size_fn) batch = [] batch_weight = 0 for element in sorted_elements: element_size = self._element_size_fn(element) # Check if adding this element would exceed limits would_exceed_count = len(batch) >= self._max_batch_size would_exceed_weight = ( batch_weight + element_size >= self._max_batch_weight and batch) if would_exceed_count or would_exceed_weight: # Emit current batch yield window.GlobalWindows.windowed_value_at_end_of_window(batch) batch = [] batch_weight = 0 batch.append(element) batch_weight += element_size # Emit remaining elements if batch: yield window.GlobalWindows.windowed_value_at_end_of_window(batch) self._buffer = None class _WindowAwareSortAndBatchElementsDoFn(DoFn): """DoFn that buffers, sorts by element size, and batches elements per window. This DoFn is used internally by ``SortAndBatchElements`` for PCollections with non-default (e.g. fixed, sliding, or session) windows. Elements are buffered per window and each window is flushed independently. To prevent a single bundle from retaining too many per-window buffers at once, when the number of live windows exceeds ``_MAX_LIVE_WINDOWS`` the largest window buffer is flushed early. This DoFn reuses ``_WindowAwareBatchingDoFn._MAX_LIVE_WINDOWS`` so it follows the same existing window-aware batching behavior already used in this module. Args: min_batch_size: The minimum number of elements per batch. Must be >= 1. max_batch_size: The maximum number of elements per batch. Must be >= ``min_batch_size``. max_batch_weight: The maximum total weight of elements in a batch, where weight is computed by ``element_size_fn``. Must be >= 1. element_size_fn: An optional callable mapping an element to its integer size/weight. """ _MAX_LIVE_WINDOWS = _WindowAwareBatchingDoFn._MAX_LIVE_WINDOWS def __init__( self, min_batch_size: int, max_batch_size: int, max_batch_weight: int, element_size_fn: Optional[Callable[[Any], int]]): self._min_batch_size = min_batch_size self._max_batch_size = max_batch_size self._max_batch_weight = max_batch_weight self._element_size_fn = element_size_fn or self._default_element_size self._has_warned_type_error = False self._buffers = collections.defaultdict(list) def _default_element_size(self, element): try: return len(element) except TypeError: if not self._has_warned_type_error: _LOGGER.warning( 'Element of type %s does not support len(). Falling back to ' 'size 1. Consider providing a custom element_size_fn to ' 'SortAndBatchElements for meaningful size-based batching.', type(element).__name__) self._has_warned_type_error = True return 1 def start_bundle(self): self._buffers = collections.defaultdict(list) def process(self, element, window=DoFn.WindowParam): self._buffers[window].append(element) # If we have too many live windows, flush the largest one if len(self._buffers) > self._MAX_LIVE_WINDOWS: largest_window = max( self._buffers.keys(), key=lambda w: len(self._buffers[w])) yield from self._flush_window(largest_window) def _flush_window(self, win): """Flush all elements for a given window.""" buffer = self._buffers.pop(win, []) if not buffer: return # Sort elements by size (ascending) sorted_elements = sorted(buffer, key=self._element_size_fn) batch = [] batch_weight = 0 for element in sorted_elements: element_size = self._element_size_fn(element) would_exceed_count = len(batch) >= self._max_batch_size would_exceed_weight = ( batch_weight + element_size >= self._max_batch_weight and batch) if would_exceed_count or would_exceed_weight: yield windowed_value.WindowedValue(batch, win.max_timestamp(), (win, )) batch = [] batch_weight = 0 batch.append(element) batch_weight += element_size if batch: yield windowed_value.WindowedValue(batch, win.max_timestamp(), (win, )) def finish_bundle(self): for win in list(self._buffers.keys()): yield from self._flush_window(win) self._buffers = None @typehints.with_input_types(T) @typehints.with_output_types(list[T]) class SortAndBatchElements(PTransform): """A Transform that sorts elements by size before batching. This transform is designed to optimize batch processing by grouping elements of similar sizes together. This is particularly useful for ML inference workloads where input sequences of varying lengths need to be padded to the maximum length in the batch - by sorting elements by size before batching, padding overhead is minimized. The transform consumes a PCollection of element type T and produces a PCollection of element type list[T], where elements within each batch are sorted by their size (as determined by element_size_fn). Elements are batched per-window and batches emitted in the window corresponding to its contents. Each batch is emitted with a timestamp at the end of their window. Unlike BatchElements which emits batches as soon as size limits are reached, SortAndBatchElements buffers all elements in a bundle, sorts them by size, and then creates optimally-sized batches. This trade-off of increased memory usage for better batch homogeneity can significantly reduce padding overhead. Args: min_batch_size: The minimum number of elements in a batch. Must be >= 1. max_batch_size: The maximum number of elements in a batch. Must be >= min_batch_size. max_batch_weight: The maximum total weight of elements in a batch, where weight is computed by element_size_fn. Must be >= 1. element_size_fn: (optional) A function mapping an element to its size/weight. If not provided, defaults to trying len(element) and falling back to 1 if the element doesn't support len(). This default allows sorting to work for common types like strings, lists, and arrays. Example usage:: # Batch strings by total character count strings = ['a', 'bb', 'ccc', 'dddd', 'eeeee'] batched = strings | SortAndBatchElements( min_batch_size=1, max_batch_size=3, max_batch_weight=10) # Possible output: [['a', 'bb', 'ccc'], ['dddd', 'eeeee']] # Elements are sorted by length and batched optimally # Batch with custom size function data = [{'text': 'short'}, {'text': 'medium text'}, {'text': 'long text here'}] batched = data | SortAndBatchElements( min_batch_size=1, max_batch_size=10, max_batch_weight=100, element_size_fn=lambda x: len(x['text'])) """ def __init__( self, min_batch_size: int, max_batch_size: int, max_batch_weight: int, element_size_fn: Optional[Callable[[Any], int]] = None): if min_batch_size < 1: raise ValueError(f'min_batch_size must be >= 1, got {min_batch_size}') if max_batch_size < min_batch_size: raise ValueError( f'max_batch_size ({max_batch_size}) must be >= ' f'min_batch_size ({min_batch_size})') if max_batch_weight < 1: raise ValueError(f'max_batch_weight must be >= 1, got {max_batch_weight}') if element_size_fn is not None and not callable(element_size_fn): raise TypeError('element_size_fn must be callable') self._min_batch_size = min_batch_size self._max_batch_size = max_batch_size self._max_batch_weight = max_batch_weight # None means the DoFn will use its own _default_element_size method, # which tries len() and warns once on TypeError before falling back to 1. self._element_size_fn = element_size_fn def expand(self, pcoll): if pcoll.windowing.is_default(): return pcoll | ParDo( _SortAndBatchElementsDoFn( self._min_batch_size, self._max_batch_size, self._max_batch_weight, self._element_size_fn)) return pcoll | ParDo( _WindowAwareSortAndBatchElementsDoFn( self._min_batch_size, self._max_batch_size, self._max_batch_weight, self._element_size_fn)) class _IdentityWindowFn(NonMergingWindowFn): """Windowing function that preserves existing windows. To be used internally with the Reshuffle transform. Will raise an exception when used after DoFns that return TimestampedValue elements. """ def __init__(self, window_coder): """Create a new WindowFn with compatible coder. To be applied to PCollections with windows that are compatible with the given coder. Arguments: window_coder: coders.Coder object to be used on windows. """ super().__init__() if window_coder is None: raise ValueError('window_coder should not be None') self._window_coder = window_coder def assign(self, assign_context): if assign_context.window is None: raise ValueError( 'assign_context.window should not be None. ' 'This might be due to a DoFn returning a TimestampedValue.') return [assign_context.window] def get_window_coder(self): return self._window_coder def reify_metadata_default_window( element, timestamp=DoFn.TimestampParam, pane_info=DoFn.PaneInfoParam): key, value = element if timestamp == window.MIN_TIMESTAMP: timestamp = None return key, (value, timestamp, pane_info) def restore_metadata_default_window(element): key, values = element return [ window.GlobalWindows.windowed_value(None).with_value((key, value)) if timestamp is None else window.GlobalWindows.windowed_value( value=(key, value), timestamp=timestamp, pane_info=pane_info) for (value, timestamp, pane_info) in values ] def reify_metadata_custom_window( element, timestamp=DoFn.TimestampParam, window=DoFn.WindowParam, pane_info=DoFn.PaneInfoParam): key, value = element return key, windowed_value.WindowedValue( value, timestamp, [window], pane_info) def restore_metadata_custom_window(element): key, windowed_values = element return [wv.with_value((key, wv.value)) for wv in windowed_values] def _reify_restore_metadata(is_default_windowing): if is_default_windowing: return reify_metadata_default_window, restore_metadata_default_window return reify_metadata_custom_window, restore_metadata_custom_window def _add_pre_map_gkb_types(pre_gbk_map, is_default_windowing): if is_default_windowing: return pre_gbk_map.with_input_types(tuple[K, V]).with_output_types( tuple[K, tuple[V, Optional[Timestamp], windowed_value.PaneInfo]]) return pre_gbk_map.with_input_types(tuple[K, V]).with_output_types( tuple[K, TypedWindowedValue[V]]) @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types(tuple[K, V]) class ReshufflePerKey(PTransform): """PTransform that returns a PCollection equivalent to its input, but operationally provides some of the side effects of a GroupByKey, in particular checkpointing, and preventing fusion of the surrounding transforms. """ def expand_2_64_0(self, pcoll): windowing_saved = pcoll.windowing if windowing_saved.is_default(): # In this (common) case we can use a trivial trigger driver # and avoid the (expensive) window param. globally_windowed = window.GlobalWindows.windowed_value(None) MIN_TIMESTAMP = window.MIN_TIMESTAMP def reify_timestamps(element, timestamp=DoFn.TimestampParam): key, value = element if timestamp == MIN_TIMESTAMP: timestamp = None return key, (value, timestamp) def restore_timestamps(element): key, values = element return [ globally_windowed.with_value((key, value)) if timestamp is None else window.GlobalWindows.windowed_value((key, value), timestamp) for (value, timestamp) in values ] if pcoll.pipeline.options.is_compat_version_prior_to( RESHUFFLE_TYPEHINT_BREAKING_CHANGE_VERSION): pre_gbk_map = Map(reify_timestamps).with_output_types(Any) else: pre_gbk_map = Map(reify_timestamps).with_input_types( tuple[K, V]).with_output_types( tuple[K, tuple[V, Optional[Timestamp]]]) else: # typing: All conditional function variants must have identical signatures def reify_timestamps( # type: ignore[misc] element, timestamp=DoFn.TimestampParam, window=DoFn.WindowParam): key, value = element # Transport the window as part of the value and restore it later. return key, windowed_value.WindowedValue(value, timestamp, [window]) def restore_timestamps(element): key, windowed_values = element return [wv.with_value((key, wv.value)) for wv in windowed_values] if pcoll.pipeline.options.is_compat_version_prior_to( RESHUFFLE_TYPEHINT_BREAKING_CHANGE_VERSION): pre_gbk_map = Map(reify_timestamps).with_output_types(Any) else: pre_gbk_map = Map(reify_timestamps).with_input_types( tuple[K, V]).with_output_types(tuple[K, TypedWindowedValue[V]]) ungrouped = pcoll | pre_gbk_map # TODO(https://github.com/apache/beam/issues/19785) Using global window as # one of the standard window. This is to mitigate the Dataflow Java Runner # Harness limitation to accept only standard coders. ungrouped._windowing = Windowing( window.GlobalWindows(), triggerfn=Always(), accumulation_mode=AccumulationMode.DISCARDING, timestamp_combiner=TimestampCombiner.OUTPUT_AT_EARLIEST) result = ( ungrouped | GroupByKey() | FlatMap(restore_timestamps).with_output_types(Any)) result._windowing = windowing_saved return result def expand(self, pcoll): if pcoll.pipeline.options.is_compat_version_prior_to("2.65.0"): return self.expand_2_64_0(pcoll) windowing_saved = pcoll.windowing is_default_windowing = windowing_saved.is_default() reify_fn, restore_fn = _reify_restore_metadata(is_default_windowing) pre_gbk_map = _add_pre_map_gkb_types(Map(reify_fn), is_default_windowing) ungrouped = pcoll | pre_gbk_map # TODO(https://github.com/apache/beam/issues/19785) Using global window as # one of the standard window. This is to mitigate the Dataflow Java Runner # Harness limitation to accept only standard coders. ungrouped._windowing = Windowing( window.GlobalWindows(), triggerfn=Always(), accumulation_mode=AccumulationMode.DISCARDING, timestamp_combiner=TimestampCombiner.OUTPUT_AT_EARLIEST) result = ( ungrouped | GroupByKey() | FlatMap(restore_fn).with_output_types(Any)) result._windowing = windowing_saved return result @typehints.with_input_types(T) @typehints.with_output_types(T) class Reshuffle(PTransform): """PTransform that returns a PCollection equivalent to its input, but operationally provides some of the side effects of a GroupByKey, in particular checkpointing, and preventing fusion of the surrounding transforms. Reshuffle adds a temporary random key to each element, performs a ReshufflePerKey, and finally removes the temporary key. """ # We use 32-bit integer as the default number of buckets. _DEFAULT_NUM_BUCKETS = 1 << 32 def __init__(self, num_buckets=None): """ :param num_buckets: If set, specifies the maximum random keys that would be generated. """ self.num_buckets = num_buckets if num_buckets else self._DEFAULT_NUM_BUCKETS valid_buckets = isinstance(num_buckets, int) and num_buckets > 0 if not (num_buckets is None or valid_buckets): raise ValueError( 'If `num_buckets` is set, it has to be an ' 'integer greater than 0, got %s' % num_buckets) def expand(self, pcoll): # type: (pvalue.PValue) -> pvalue.PCollection if pcoll.pipeline.options.is_compat_version_prior_to( RESHUFFLE_TYPEHINT_BREAKING_CHANGE_VERSION): reshuffle_step = ReshufflePerKey() else: reshuffle_step = ReshufflePerKey().with_input_types( tuple[int, T]).with_output_types(tuple[int, T]) return ( pcoll | 'AddRandomKeys' >> Map(lambda t: (random.randrange(0, self.num_buckets), t) ).with_input_types(T).with_output_types(tuple[int, T]) | reshuffle_step | 'RemoveRandomKeys' >> Map(lambda t: t[1]).with_input_types( tuple[int, T]).with_output_types(T)) def to_runner_api_parameter(self, unused_context): # type: (PipelineContext) -> tuple[str, None] return common_urns.composites.RESHUFFLE.urn, None @staticmethod @PTransform.register_urn(common_urns.composites.RESHUFFLE.urn, None) def from_runner_api_parameter( unused_ptransform, unused_parameter, unused_context): return Reshuffle() def fn_takes_side_inputs(fn): fn = getattr(fn, '_argspec_fn', fn) try: signature = get_signature(fn) except TypeError: # We can't tell; maybe it does. return True return ( len(signature.parameters) > 1 or any( p.kind == p.VAR_POSITIONAL or p.kind == p.VAR_KEYWORD for p in signature.parameters.values())) @ptransform_fn def WithKeys(pcoll, k, *args, **kwargs): """PTransform that takes a PCollection, and either a constant key or a callable, and returns a PCollection of (K, V), where each of the values in the input PCollection has been paired with either the constant key or a key computed from the value. The callable may optionally accept positional or keyword arguments, which should be passed to WithKeys directly. These may be either SideInputs or static (non-PCollection) values, such as ints. """ if callable(k): if fn_takes_side_inputs(k): if all(isinstance(arg, AsSideInput) for arg in args) and all(isinstance(kwarg, AsSideInput) for kwarg in kwargs.values()): # Map(lambda) produces a label formatted like this, but it cannot be # changed without breaking update compat. Here, we pin to the transform # name used in the 2.68 release to avoid breaking changes when the line # number changes. Context: https://github.com/apache/beam/pull/36381 return pcoll | "Map()" >> Map( lambda v, *args, **kwargs: (k(v, *args, **kwargs), v), *args, **kwargs) return pcoll | "Map()" >> Map( lambda v: (k(v, *args, **kwargs), v)) return pcoll | "Map()" >> Map(lambda v: (k(v), v)) return pcoll | "Map()" >> Map(lambda v: (k, v)) @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types(tuple[K, Iterable[V]]) class GroupIntoBatches(PTransform): """PTransform that batches the input into desired batch size. Elements are buffered until they are equal to batch size provided in the argument at which point they are output to the output Pcollection. Windows are preserved (batches will contain elements from the same window) """ def __init__( self, batch_size, max_buffering_duration_secs=None, clock=time.time): """Create a new GroupIntoBatches. Arguments: batch_size: (required) How many elements should be in a batch max_buffering_duration_secs: (optional) How long in seconds at most an incomplete batch of elements is allowed to be buffered in the states. The duration must be a positive second duration and should be given as an int or float. Setting this parameter to zero effectively means no buffering limit. clock: (optional) an alternative to time.time (mostly for testing) """ self.params = _GroupIntoBatchesParams( batch_size, max_buffering_duration_secs) self.clock = clock def expand(self, pcoll): input_coder = coders.registry.get_coder(pcoll) return pcoll | ParDo( _pardo_group_into_batches( input_coder, self.params.batch_size, self.params.max_buffering_duration_secs, self.clock)) def to_runner_api_parameter( self, unused_context # type: PipelineContext ): # type: (...) -> tuple[str, beam_runner_api_pb2.GroupIntoBatchesPayload] return ( common_urns.group_into_batches_components.GROUP_INTO_BATCHES.urn, self.params.get_payload()) @staticmethod @PTransform.register_urn( common_urns.group_into_batches_components.GROUP_INTO_BATCHES.urn, beam_runner_api_pb2.GroupIntoBatchesPayload) def from_runner_api_parameter(unused_ptransform, proto, unused_context): return GroupIntoBatches(*_GroupIntoBatchesParams.parse_payload(proto)) @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types( typehints.Tuple[ ShardedKeyType[typehints.TypeVariable(K)], # type: ignore[misc] typehints.Iterable[typehints.TypeVariable(V)]]) class WithShardedKey(PTransform): """A GroupIntoBatches transform that outputs batched elements associated with sharded input keys. By default, keys are sharded to such that the input elements with the same key are spread to all available threads executing the transform. Runners may override the default sharding to do a better load balancing during the execution time. """ def __init__( self, batch_size, max_buffering_duration_secs=None, clock=time.time): """Create a new GroupIntoBatches with sharded output. See ``GroupIntoBatches`` transform for a description of input parameters. """ self.params = _GroupIntoBatchesParams( batch_size, max_buffering_duration_secs) self.clock = clock _shard_id_prefix = uuid.uuid4().bytes def expand(self, pcoll): key_type, value_type = pcoll.element_type.tuple_types # Map(lambda) produces a label formatted like this, but it cannot be # changed without breaking update compat. Here, we pin to the transform # name used in the 2.68 release to avoid breaking changes when the line # number changes. Context: https://github.com/apache/beam/pull/36381 sharded_pcoll = pcoll | "Map()" >> Map( lambda key_value: ( ShardedKey( key_value[0], # Use [uuid, thread id] as the shard id. GroupIntoBatches.WithShardedKey._shard_id_prefix + bytes( threading.get_ident().to_bytes(8, 'big'))), key_value[1])).with_output_types( typehints.Tuple[ ShardedKeyType[key_type], # type: ignore[misc] value_type]) return ( sharded_pcoll | GroupIntoBatches( self.params.batch_size, self.params.max_buffering_duration_secs, self.clock)) def to_runner_api_parameter( self, unused_context # type: PipelineContext ): # type: (...) -> tuple[str, beam_runner_api_pb2.GroupIntoBatchesPayload] return ( common_urns.composites.GROUP_INTO_BATCHES_WITH_SHARDED_KEY.urn, self.params.get_payload()) @staticmethod @PTransform.register_urn( common_urns.composites.GROUP_INTO_BATCHES_WITH_SHARDED_KEY.urn, beam_runner_api_pb2.GroupIntoBatchesPayload) def from_runner_api_parameter(unused_ptransform, proto, unused_context): return GroupIntoBatches.WithShardedKey( *_GroupIntoBatchesParams.parse_payload(proto)) class _GroupIntoBatchesParams: """This class represents the parameters for :class:`apache_beam.utils.GroupIntoBatches` transform, used to define how elements should be batched. """ def __init__(self, batch_size, max_buffering_duration_secs): self.batch_size = batch_size self.max_buffering_duration_secs = ( 0 if max_buffering_duration_secs is None else max_buffering_duration_secs) self._validate() def __eq__(self, other): if other is None or not isinstance(other, _GroupIntoBatchesParams): return False return ( self.batch_size == other.batch_size and self.max_buffering_duration_secs == other.max_buffering_duration_secs) def _validate(self): assert self.batch_size is not None and self.batch_size > 0, ( 'batch_size must be a positive value') assert ( self.max_buffering_duration_secs is not None and self.max_buffering_duration_secs >= 0), ('max_buffering_duration must be a non-negative value') def get_payload(self): return beam_runner_api_pb2.GroupIntoBatchesPayload( batch_size=self.batch_size, max_buffering_duration_millis=int( self.max_buffering_duration_secs * 1000)) @staticmethod def parse_payload( proto # type: beam_runner_api_pb2.GroupIntoBatchesPayload ): return proto.batch_size, proto.max_buffering_duration_millis / 1000 def _pardo_group_into_batches( input_coder, batch_size, max_buffering_duration_secs, clock=time.time): ELEMENT_STATE = BagStateSpec('values', input_coder) COUNT_STATE = CombiningValueStateSpec('count', input_coder, CountCombineFn()) WINDOW_TIMER = TimerSpec('window_end', TimeDomain.WATERMARK) BUFFERING_TIMER = TimerSpec('buffering_end', TimeDomain.REAL_TIME) class _GroupIntoBatchesDoFn(DoFn): def process( self, element, window=DoFn.WindowParam, element_state=DoFn.StateParam(ELEMENT_STATE), count_state=DoFn.StateParam(COUNT_STATE), window_timer=DoFn.TimerParam(WINDOW_TIMER), buffering_timer=DoFn.TimerParam(BUFFERING_TIMER)): # Allowed lateness not supported in Python SDK # https://beam.apache.org/documentation/programming-guide/#watermarks-and-late-data window_timer.set(window.end) element_state.add(element) count_state.add(1) count = count_state.read() if count == 1 and max_buffering_duration_secs > 0: # This is the first element in batch. Start counting buffering time if a # limit was set. # pylint: disable=deprecated-method buffering_timer.set(clock() + max_buffering_duration_secs) if count >= batch_size: return self.flush_batch(element_state, count_state, buffering_timer) @on_timer(WINDOW_TIMER) def on_window_timer( self, element_state=DoFn.StateParam(ELEMENT_STATE), count_state=DoFn.StateParam(COUNT_STATE), buffering_timer=DoFn.TimerParam(BUFFERING_TIMER)): return self.flush_batch(element_state, count_state, buffering_timer) @on_timer(BUFFERING_TIMER) def on_buffering_timer( self, element_state=DoFn.StateParam(ELEMENT_STATE), count_state=DoFn.StateParam(COUNT_STATE), buffering_timer=DoFn.TimerParam(BUFFERING_TIMER)): return self.flush_batch(element_state, count_state, buffering_timer) def flush_batch(self, element_state, count_state, buffering_timer): batch = [element for element in element_state.read()] if not batch: return key, _ = batch[0] batch_values = [v for (k, v) in batch] element_state.clear() count_state.clear() buffering_timer.clear() yield key, batch_values return _GroupIntoBatchesDoFn() class ToString(object): """ PTransform for converting a PCollection element, KV or PCollection Iterable to string. """ # pylint: disable=invalid-name @staticmethod def Element(): """ Transforms each element of the PCollection to a string. """ return 'ElementToString' >> Map(str) @staticmethod def Iterables(delimiter=None): """ Transforms each item in the iterable of the input of PCollection to a string. There is no trailing delimiter. """ if delimiter is None: delimiter = ',' return ( 'IterablesToString' >> Map(lambda xs: delimiter.join(str(x) for x in xs)).with_input_types( Iterable[Any]).with_output_types(str)) # An alias for Iterables. Kvs = Iterables @typehints.with_input_types(T) @typehints.with_output_types(T) class LogElements(PTransform): """ PTransform for printing the elements of a PCollection. Args: label (str): (optional) A custom label for the transform. prefix (str): (optional) A prefix string to prepend to each logged element. with_timestamp (bool): (optional) Whether to include element's timestamp. with_window (bool): (optional) Whether to include element's window. level: (optional) The logging level for the output (e.g. `logging.DEBUG`, `logging.INFO`, `logging.WARNING`, `logging.ERROR`). If not specified, the log is printed to stdout. with_pane_info (bool): (optional) Whether to include element's pane info. use_epoch_time (bool): (optional) Whether to display epoch timestamps. """ class _LoggingFn(DoFn): def __init__( self, prefix='', with_timestamp=False, with_window=False, level=None, with_pane_info=False, use_epoch_time=False): super().__init__() self.prefix = prefix self.with_timestamp = with_timestamp self.with_window = with_window self.level = level self.with_pane_info = with_pane_info self.use_epoch_time = use_epoch_time def format_timestamp(self, timestamp): if self.use_epoch_time: return timestamp.seconds() return timestamp.to_rfc3339() def process( self, element, timestamp=DoFn.TimestampParam, window=DoFn.WindowParam, pane_info=DoFn.PaneInfoParam, **kwargs): log_line = self.prefix + str(element) if self.with_timestamp: log_line += ', timestamp=' + repr(self.format_timestamp(timestamp)) if self.with_window: log_line += ', window(start=' + str(self.format_timestamp(window.start)) log_line += ', end=' + str(self.format_timestamp(window.end)) + ')' if self.with_pane_info: log_line += ', pane_info=' + repr(pane_info) if self.level == logging.DEBUG: _LOGGER.debug(log_line) elif self.level == logging.INFO: _LOGGER.info(log_line) elif self.level == logging.WARNING: _LOGGER.warning(log_line) elif self.level == logging.ERROR: _LOGGER.error(log_line) elif self.level == logging.CRITICAL: _LOGGER.critical(log_line) else: print(log_line) yield element def __init__( self, label=None, prefix='', with_timestamp=False, with_window=False, level=None, with_pane_info=False, use_epoch_time=False, ): super().__init__(label) self.prefix = prefix self.with_timestamp = with_timestamp self.with_window = with_window self.with_pane_info = with_pane_info self.use_epoch_time = use_epoch_time self.level = level def expand(self, input): return input | ParDo( self._LoggingFn( self.prefix, self.with_timestamp, self.with_window, self.level, self.with_pane_info, self.use_epoch_time, )) class Reify(object): """PTransforms for converting between explicit and implicit form of various Beam values.""" @typehints.with_input_types(T) @typehints.with_output_types(T) class Timestamp(PTransform): """PTransform to wrap a value in a TimestampedValue with it's associated timestamp.""" @staticmethod def add_timestamp_info(element, timestamp=DoFn.TimestampParam): yield TimestampedValue(element, timestamp) def expand(self, pcoll): return pcoll | ParDo(self.add_timestamp_info) @typehints.with_input_types(T) @typehints.with_output_types(T) class Window(PTransform): """PTransform to convert an element in a PCollection into a tuple of (element, timestamp, window), wrapped in a TimestampedValue with it's associated timestamp.""" @staticmethod def add_window_info( element, timestamp=DoFn.TimestampParam, window=DoFn.WindowParam): yield TimestampedValue((element, timestamp, window), timestamp) def expand(self, pcoll): return pcoll | ParDo(self.add_window_info) @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types(tuple[K, V]) class TimestampInValue(PTransform): """PTransform to wrap the Value in a KV pair in a TimestampedValue with the element's associated timestamp.""" @staticmethod def add_timestamp_info(element, timestamp=DoFn.TimestampParam): key, value = element yield (key, TimestampedValue(value, timestamp)) def expand(self, pcoll): return pcoll | ParDo(self.add_timestamp_info) @typehints.with_input_types(tuple[K, V]) @typehints.with_output_types(tuple[K, V]) class WindowInValue(PTransform): """PTransform to convert the Value in a KV pair into a tuple of (value, timestamp, window), with the whole element being wrapped inside a TimestampedValue.""" @staticmethod def add_window_info( element, timestamp=DoFn.TimestampParam, window=DoFn.WindowParam): key, value = element yield TimestampedValue((key, (value, timestamp, window)), timestamp) def expand(self, pcoll): return pcoll | ParDo(self.add_window_info) class Regex(object): """ PTransform to use Regular Expression to process the elements in a PCollection. """ ALL = "__regex_all_groups" @staticmethod def _regex_compile(regex): """Return re.compile if the regex has a string value""" if isinstance(regex, str): regex = re.compile(regex) return regex @staticmethod @typehints.with_input_types(str) @typehints.with_output_types(str) @ptransform_fn def matches(pcoll, regex, group=0): """ Returns the matches (group 0 by default) if zero or more characters at the beginning of string match the regular expression. To match the entire string, add "$" sign at the end of regex expression. Group can be integer value or a string value. Args: regex: the regular expression string or (re.compile) pattern. group: (optional) name/number of the group, it can be integer or a string value. Defaults to 0, meaning the entire matched string will be returned. """ regex = Regex._regex_compile(regex) def _process(element): m = regex.match(element) if m: yield m.group(group) return pcoll | FlatMap(_process) @staticmethod @typehints.with_input_types(str) @typehints.with_output_types(list[str]) @ptransform_fn def all_matches(pcoll, regex): """ Returns all matches (groups) if zero or more characters at the beginning of string match the regular expression. Args: regex: the regular expression string or (re.compile) pattern. """ regex = Regex._regex_compile(regex) def _process(element): m = regex.match(element) if m: yield [m.group(ix) for ix in range(m.lastindex + 1)] return pcoll | FlatMap(_process) @staticmethod @typehints.with_input_types(str) @typehints.with_output_types(tuple[str, str]) @ptransform_fn def matches_kv(pcoll, regex, keyGroup, valueGroup=0): """ Returns the KV pairs if the string matches the regular expression, deriving the key & value from the specified group of the regular expression. Args: regex: the regular expression string or (re.compile) pattern. keyGroup: The Regex group to use as the key. Can be int or str. valueGroup: (optional) Regex group to use the value. Can be int or str. The default value "0" returns entire matched string. """ regex = Regex._regex_compile(regex) def _process(element): match = regex.match(element) if match: yield (match.group(keyGroup), match.group(valueGroup)) return pcoll | FlatMap(_process) @staticmethod @typehints.with_input_types(str) @typehints.with_output_types(str) @ptransform_fn def find(pcoll, regex, group=0): """ Returns the matches if a portion of the line matches the Regex. Returns the entire group (group 0 by default). Group can be integer value or a string value. Args: regex: the regular expression string or (re.compile) pattern. group: (optional) name of the group, it can be integer or a string value. """ regex = Regex._regex_compile(regex) def _process(element): r = regex.search(element) if r: yield r.group(group) return pcoll | FlatMap(_process) @staticmethod @typehints.with_input_types(str) @typehints.with_output_types(Union[list[str], list[tuple[str, str]]]) @ptransform_fn def find_all(pcoll, regex, group=0, outputEmpty=True): """ Returns the matches if a portion of the line matches the Regex. By default, list of group 0 will return with empty items. To get all groups, pass the `Regex.ALL` flag in the `group` parameter which returns all the groups in the tuple format. Args: regex: the regular expression string or (re.compile) pattern. group: (optional) name of the group, it can be integer or a string value. outputEmpty: (optional) Should empty be output. True to output empties and false if not. """ regex = Regex._regex_compile(regex) def _process(element): matches = regex.finditer(element) if group == Regex.ALL: yield [(m.group(), m.groups()[0]) for m in matches if outputEmpty or m.groups()[0]] else: yield [m.group(group) for m in matches if outputEmpty or m.group(group)] return pcoll | FlatMap(_process) @staticmethod @typehints.with_input_types(str) @typehints.with_output_types(tuple[str, str]) @ptransform_fn def find_kv(pcoll, regex, keyGroup, valueGroup=0): """ Returns the matches if a portion of the line matches the Regex. Returns the specified groups as the key and value pair. Args: regex: the regular expression string or (re.compile) pattern. keyGroup: The Regex group to use as the key. Can be int or str. valueGroup: (optional) Regex group to use the value. Can be int or str. The default value "0" returns entire matched string. """ regex = Regex._regex_compile(regex) def _process(element): matches = regex.finditer(element) if matches: for match in matches: yield (match.group(keyGroup), match.group(valueGroup)) return pcoll | FlatMap(_process) @staticmethod @typehints.with_input_types(str) @typehints.with_output_types(str) @ptransform_fn def replace_all(pcoll, regex, replacement): """ Returns the matches if a portion of the line matches the regex and replaces all matches with the replacement string. Args: regex: the regular expression string or (re.compile) pattern. replacement: the string to be substituted for each match. """ regex = Regex._regex_compile(regex) # Map(lambda) produces a label formatted like this, but it cannot be # changed without breaking update compat. Here, we pin to the transform # name used in the 2.68 release to avoid breaking changes when the line # number changes. Context: https://github.com/apache/beam/pull/36381 return pcoll | "Map()" >> Map( lambda elem: regex.sub(replacement, elem)) @staticmethod @typehints.with_input_types(str) @typehints.with_output_types(str) @ptransform_fn def replace_first(pcoll, regex, replacement): """ Returns the matches if a portion of the line matches the regex and replaces the first match with the replacement string. Args: regex: the regular expression string or (re.compile) pattern. replacement: the string to be substituted for each match. """ regex = Regex._regex_compile(regex) # Map(lambda) produces a label formatted like this, but it cannot be # changed without breaking update compat. Here, we pin to the transform # name used in the 2.68 release to avoid breaking changes when the line # number changes. Context: https://github.com/apache/beam/pull/36381 return pcoll | "Map()" >> Map( lambda elem: regex.sub(replacement, elem, 1)) @staticmethod @typehints.with_input_types(str) @typehints.with_output_types(list[str]) @ptransform_fn def split(pcoll, regex, outputEmpty=False): """ Returns the list string which was splitted on the basis of regular expression. It will not output empty items (by defaults). Args: regex: the regular expression string or (re.compile) pattern. outputEmpty: (optional) Should empty be output. True to output empties and false if not. """ regex = Regex._regex_compile(regex) outputEmpty = bool(outputEmpty) def _process(element): r = regex.split(element) if r and not outputEmpty: r = list(filter(None, r)) yield r return pcoll | FlatMap(_process) @typehints.with_input_types(T) @typehints.with_output_types(T) class Tee(PTransform): """A PTransform that returns its input, but also applies its input elsewhere. Similar to the shell {@code tee} command. This can be useful to write out or otherwise process an intermediate transform without breaking the linear flow of a chain of transforms, e.g.:: (input | SomePTransform() | ... | Tee(SomeSideTransform()) | ...) """ def __init__( self, *consumers: Union[PTransform[PCollection[T], Any], Callable[[PCollection[T]], Any]]): self._consumers = consumers def expand(self, input): for consumer in self._consumers: if callable(consumer): _ = input | ptransform_fn(consumer)() else: _ = input | consumer return input @typehints.with_input_types(T) @typehints.with_output_types(T) class WaitOn(PTransform): """Delays processing of a {@link PCollection} until another set of PCollections has finished being processed. For example:: X | WaitOn(Y, Z) | SomeTransform() would ensure that PCollections Y and Z (and hence their producing transforms) are complete before SomeTransform gets executed on the elements of X. This can be especially useful the waited-on PCollections are the outputs of transforms that interact with external systems (such as writing to a database or other sink). For streaming, this delay is done on a per-window basis, i.e. the corresponding window of each waited-on PCollection is computed before elements are passed through the main collection. This barrier often induces a fusion break. """ def __init__(self, *to_be_waited_on): self._to_be_waited_on = to_be_waited_on def expand(self, pcoll): # All we care about is the watermark, not the data itself. # The GroupByKey avoids writing empty files for each shard, and also # ensures the respective window finishes before advancing the timestamp. sides = [ pvalue.AsIter( side | f"WaitOn{ix}" >> (beam.FlatMap(lambda x: ()) | GroupByKey())) for (ix, side) in enumerate(self._to_be_waited_on) ] # Map(lambda) produces a label formatted like this, but it cannot be # changed without breaking update compat. Here, we pin to the transform # name used in the 2.68 release to avoid breaking changes when the line # number changes. Context: https://github.com/apache/beam/pull/36381 return pcoll | "Map()" >> beam.Map( lambda x, *unused_sides: x, *sides) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/util_test.py0000644000175100017510000035005015206113033022445 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the transform.util classes.""" # pytype: skip-file # pylint: disable=too-many-function-args import collections import hashlib import hmac import importlib import logging import math import random import re import time import unittest import warnings from collections.abc import Mapping from datetime import datetime import mock import pytest import pytz from cryptography.fernet import Fernet from cryptography.fernet import InvalidToken from parameterized import param from parameterized import parameterized import apache_beam as beam from apache_beam import GroupByKey from apache_beam import Map from apache_beam import WindowInto from apache_beam.coders import coders from apache_beam.metrics import MetricsFilter from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import SetupOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.pipeline_options import TypeOptions from apache_beam.portability import common_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.pvalue import AsList from apache_beam.pvalue import AsSingleton from apache_beam.runners import pipeline_context from apache_beam.testing.synthetic_pipeline import SyntheticSource from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.test_stream import TestStream from apache_beam.testing.util import SortLists from apache_beam.testing.util import TestWindowedValue from apache_beam.testing.util import assert_that from apache_beam.testing.util import contains_in_any_order from apache_beam.testing.util import equal_to from apache_beam.testing.util import is_not_empty from apache_beam.transforms import trigger from apache_beam.transforms import util from apache_beam.transforms import window from apache_beam.transforms.core import FlatMapTuple from apache_beam.transforms.trigger import AfterCount from apache_beam.transforms.trigger import Repeatedly from apache_beam.transforms.util import GcpHsmGeneratedSecret from apache_beam.transforms.util import GcpSecret from apache_beam.transforms.util import Secret from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import GlobalWindow from apache_beam.transforms.window import GlobalWindows from apache_beam.transforms.window import IntervalWindow from apache_beam.transforms.window import Sessions from apache_beam.transforms.window import SlidingWindows from apache_beam.transforms.window import TimestampedValue from apache_beam.typehints import typehints from apache_beam.typehints.sharded_key_type import ShardedKeyType from apache_beam.utils import proto_utils from apache_beam.utils import timestamp from apache_beam.utils.timestamp import MAX_TIMESTAMP from apache_beam.utils.timestamp import MIN_TIMESTAMP from apache_beam.utils.windowed_value import PANE_INFO_UNKNOWN from apache_beam.utils.windowed_value import PaneInfo from apache_beam.utils.windowed_value import PaneInfoTiming from apache_beam.utils.windowed_value import WindowedValue try: from google.cloud import secretmanager except ImportError: secretmanager = None # type: ignore[assignment] warnings.filterwarnings( 'ignore', category=FutureWarning, module='apache_beam.transform.util_test') class _Unpicklable(object): def __init__(self, value): self.value = value def __getstate__(self): raise NotImplementedError() def __setstate__(self, state): raise NotImplementedError() class _UnpicklableCoder(beam.coders.Coder): def encode(self, value): return str(value.value).encode() def decode(self, encoded): return _Unpicklable(int(encoded.decode())) def to_type_hint(self): return _Unpicklable def is_deterministic(self): return True class CoGroupByKeyTest(unittest.TestCase): def test_co_group_by_key_on_tuple(self): with TestPipeline() as pipeline: pcoll_1 = pipeline | 'Start 1' >> beam.Create([('a', 1), ('a', 2), ('b', 3), ('c', 4)]) pcoll_2 = pipeline | 'Start 2' >> beam.Create([('a', 5), ('a', 6), ('c', 7), ('c', 8)]) result = (pcoll_1, pcoll_2) | beam.CoGroupByKey() | SortLists assert_that( result, equal_to([('a', ([1, 2], [5, 6])), ('b', ([3], [])), ('c', ([4], [7, 8]))])) def test_co_group_by_key_on_iterable(self): with TestPipeline() as pipeline: pcoll_1 = pipeline | 'Start 1' >> beam.Create([('a', 1), ('a', 2), ('b', 3), ('c', 4)]) pcoll_2 = pipeline | 'Start 2' >> beam.Create([('a', 5), ('a', 6), ('c', 7), ('c', 8)]) result = iter([pcoll_1, pcoll_2]) | beam.CoGroupByKey() | SortLists assert_that( result, equal_to([('a', ([1, 2], [5, 6])), ('b', ([3], [])), ('c', ([4], [7, 8]))])) def test_co_group_by_key_on_list(self): with TestPipeline() as pipeline: pcoll_1 = pipeline | 'Start 1' >> beam.Create([('a', 1), ('a', 2), ('b', 3), ('c', 4)]) pcoll_2 = pipeline | 'Start 2' >> beam.Create([('a', 5), ('a', 6), ('c', 7), ('c', 8)]) result = [pcoll_1, pcoll_2] | beam.CoGroupByKey() | SortLists assert_that( result, equal_to([('a', ([1, 2], [5, 6])), ('b', ([3], [])), ('c', ([4], [7, 8]))])) def test_co_group_by_key_on_dict(self): with TestPipeline() as pipeline: pcoll_1 = pipeline | 'Start 1' >> beam.Create([('a', 1), ('a', 2), ('b', 3), ('c', 4)]) pcoll_2 = pipeline | 'Start 2' >> beam.Create([('a', 5), ('a', 6), ('c', 7), ('c', 8)]) result = {'X': pcoll_1, 'Y': pcoll_2} | beam.CoGroupByKey() | SortLists assert_that( result, equal_to([('a', { 'X': [1, 2], 'Y': [5, 6] }), ('b', { 'X': [3], 'Y': [] }), ('c', { 'X': [4], 'Y': [7, 8] })])) def test_co_group_by_key_on_dict_with_tuple_keys(self): with TestPipeline() as pipeline: key = ('a', ('b', 'c')) pcoll_1 = pipeline | 'Start 1' >> beam.Create([(key, 1)]) pcoll_2 = pipeline | 'Start 2' >> beam.Create([(key, 2)]) result = {'X': pcoll_1, 'Y': pcoll_2} | beam.CoGroupByKey() | SortLists assert_that(result, equal_to([(key, {'X': [1], 'Y': [2]})])) def test_co_group_by_key_on_empty(self): with TestPipeline() as pipeline: assert_that( tuple() | 'EmptyTuple' >> beam.CoGroupByKey(pipeline=pipeline), equal_to([]), label='AssertEmptyTuple') assert_that([] | 'EmptyList' >> beam.CoGroupByKey(pipeline=pipeline), equal_to([]), label='AssertEmptyList') assert_that( iter([]) | 'EmptyIterable' >> beam.CoGroupByKey(pipeline=pipeline), equal_to([]), label='AssertEmptyIterable') assert_that({} | 'EmptyDict' >> beam.CoGroupByKey(pipeline=pipeline), equal_to([]), label='AssertEmptyDict') def test_co_group_by_key_on_one(self): with TestPipeline() as pipeline: pcoll = pipeline | beam.Create([('a', 1), ('b', 2)]) expected = [('a', ([1], )), ('b', ([2], ))] assert_that((pcoll, ) | 'OneTuple' >> beam.CoGroupByKey(), equal_to(expected), label='AssertOneTuple') assert_that([pcoll] | 'OneList' >> beam.CoGroupByKey(), equal_to(expected), label='AssertOneList') assert_that( iter([pcoll]) | 'OneIterable' >> beam.CoGroupByKey(), equal_to(expected), label='AssertOneIterable') assert_that({'tag': pcoll} | 'OneDict' >> beam.CoGroupByKey() | beam.MapTuple(lambda k, v: (k, (v['tag'], ))), equal_to(expected), label='AssertOneDict') def test_co_group_by_key_on_unpickled(self): beam.coders.registry.register_coder(_Unpicklable, _UnpicklableCoder) values = [_Unpicklable(i) for i in range(5)] with TestPipeline() as pipeline: xs = pipeline | beam.Create(values) | beam.WithKeys(lambda x: x) pcoll = ({ 'x': xs } | beam.CoGroupByKey() | beam.FlatMapTuple( lambda k, tagged: (k.value, tagged['x'][0].value * 2))) expected = [0, 0, 1, 2, 2, 4, 3, 6, 4, 8] assert_that(pcoll, equal_to(expected)) class FakeSecret(beam.Secret): def __init__(self, version_name=None, should_throw=False): self._secret = b'aKwI2PmqYFt2p5tNKCyBS5qYmHhHsGZcyZrnZQiQ-uE=' self._should_throw = should_throw def get_secret_bytes(self) -> bytes: if self._should_throw: raise RuntimeError('Exception retrieving secret') return self._secret class MockNoOpDecrypt(beam.transforms.util._DecryptMessage): def __init__(self, hmac_key_secret, key_coder, value_coder): hmac_key = hmac_key_secret.get_secret_bytes() self.fernet_tester = Fernet(hmac_key) self.known_hmacs = [] for key in ['a', 'b', 'c']: self.known_hmacs.append( hmac.new(hmac_key, key_coder.encode(key), hashlib.sha256).digest()) super().__init__(hmac_key_secret, key_coder, value_coder) def process(self, element): final_elements = list(super().process(element)) # Check if we're looking at the actual elements being encoded/decoded # There is also a gbk on assertEqual, which uses None as the key type. final_element_keys = [e for e in final_elements if e[0] in ['a', 'b', 'c']] if len(final_element_keys) == 0: return final_elements hmac_key, actual_elements = element if hmac_key not in self.known_hmacs: raise ValueError(f'GBK produced unencrypted value {hmac_key}') for e in actual_elements: try: self.fernet_tester.decrypt(e[0], None) except InvalidToken: raise ValueError(f'GBK produced unencrypted value {e[0]}') try: self.fernet_tester.decrypt(e[1], None) except InvalidToken: raise ValueError(f'GBK produced unencrypted value {e[1]}') return final_elements class SecretTest(unittest.TestCase): @parameterized.expand([ param( secret_string='type:GcpSecret;version_name:my_secret/versions/latest', secret=GcpSecret('my_secret/versions/latest')), param( secret_string='type:GcpSecret;version_name:foo', secret=GcpSecret('foo')), param( secret_string='type:gcpsecreT;version_name:my_secret/versions/latest', secret=GcpSecret('my_secret/versions/latest')), ]) def test_secret_manager_parses_correctly(self, secret_string, secret): self.assertEqual(secret, Secret.parse_secret_option(secret_string)) @parameterized.expand([ param( secret_string='version_name:foo', exception_str='must contain a valid type parameter'), param( secret_string='type:gcpsecreT', exception_str='missing 1 required positional argument'), param( secret_string='type:gcpsecreT;version_name:foo;extra:val', exception_str='Invalid secret parameter extra'), ]) def test_secret_manager_throws_on_invalid(self, secret_string, exception_str): with self.assertRaisesRegex(Exception, exception_str): Secret.parse_secret_option(secret_string) class GroupByEncryptedKeyTest(unittest.TestCase): @classmethod def setUpClass(cls): if secretmanager is not None: cls.project_id = 'apache-beam-testing' cls.secret_id = 'gbek_util_secret_tests' cls.client = secretmanager.SecretManagerServiceClient() cls.project_path = f'projects/{cls.project_id}' cls.secret_path = f'{cls.project_path}/secrets/{cls.secret_id}' try: cls.client.get_secret(request={'name': cls.secret_path}) except Exception: cls.client.create_secret( request={ 'parent': cls.project_path, 'secret_id': cls.secret_id, 'secret': { 'replication': { 'automatic': {} } } }) cls.client.add_secret_version( request={ 'parent': cls.secret_path, 'payload': { 'data': Secret.generate_secret_bytes() } }) version_name = f'{cls.secret_path}/versions/latest' cls.gcp_secret = GcpSecret(version_name) cls.secret_option = f'type:GcpSecret;version_name:{version_name}' def test_gbek_fake_secret_manager_roundtrips(self): fakeSecret = FakeSecret() with TestPipeline() as pipeline: pcoll_1 = pipeline | 'Start 1' >> beam.Create([('a', 1), ('a', 2), ('b', 3), ('c', 4)]) result = (pcoll_1) | beam.GroupByEncryptedKey(fakeSecret) assert_that( result, equal_to([('a', ([1, 2])), ('b', ([3])), ('c', ([4]))])) @unittest.skipIf(secretmanager is None, 'GCP dependencies are not installed') def test_gbk_with_gbek_option_fake_secret_manager_roundtrips(self): options = PipelineOptions() options.view_as(SetupOptions).gbek = self.secret_option with beam.Pipeline(options=options) as pipeline: pcoll_1 = pipeline | 'Start 1' >> beam.Create([('a', 1), ('a', 2), ('b', 3), ('c', 4)]) result = (pcoll_1) | beam.GroupByKey() sorted_result = result | beam.Map(lambda x: (x[0], sorted(x[1]))) assert_that( sorted_result, equal_to([('a', ([1, 2])), ('b', ([3])), ('c', ([4]))])) @mock.patch('apache_beam.transforms.util._DecryptMessage', MockNoOpDecrypt) def test_gbek_fake_secret_manager_actually_does_encryption(self): fakeSecret = FakeSecret() with TestPipeline('FnApiRunner') as pipeline: pcoll_1 = pipeline | 'Start 1' >> beam.Create([('a', 1), ('a', 2), ('b', 3), ('c', 4)]) result = (pcoll_1) | beam.GroupByEncryptedKey(fakeSecret) assert_that( result, equal_to([('a', ([1, 2])), ('b', ([3])), ('c', ([4]))])) @mock.patch('apache_beam.transforms.util._DecryptMessage', MockNoOpDecrypt) @mock.patch('apache_beam.transforms.util.GcpSecret', FakeSecret) def test_gbk_actually_does_encryption(self): options = PipelineOptions() # Version of GcpSecret doesn't matter since it is replaced by FakeSecret options.view_as(SetupOptions).gbek = 'type:GcpSecret;version_name:Foo' with TestPipeline('FnApiRunner', options=options) as pipeline: pcoll_1 = pipeline | 'Start 1' >> beam.Create([('a', 1), ('a', 2), ('b', 3), ('c', 4)], reshuffle=False) result = pcoll_1 | beam.GroupByKey() assert_that( result, equal_to([('a', ([1, 2])), ('b', ([3])), ('c', ([4]))])) def test_gbek_fake_secret_manager_throws(self): fakeSecret = FakeSecret(None, True) with self.assertRaisesRegex(RuntimeError, r'Exception retrieving secret'): with TestPipeline() as pipeline: pcoll_1 = pipeline | 'Start 1' >> beam.Create([('a', 1), ('a', 2), ('b', 3), ('c', 4)]) result = (pcoll_1) | beam.GroupByEncryptedKey(fakeSecret) assert_that( result, equal_to([('a', ([1, 2])), ('b', ([3])), ('c', ([4]))])) @unittest.skipIf(secretmanager is None, 'GCP dependencies are not installed') def test_gbek_gcp_secret_manager_roundtrips(self): with TestPipeline() as pipeline: pcoll_1 = pipeline | 'Start 1' >> beam.Create([('a', 1), ('a', 2), ('b', 3), ('c', 4)]) result = (pcoll_1) | beam.GroupByEncryptedKey(self.gcp_secret) assert_that( result, equal_to([('a', ([1, 2])), ('b', ([3])), ('c', ([4]))])) @unittest.skipIf(secretmanager is None, 'GCP dependencies are not installed') def test_gbek_gcp_secret_manager_throws(self): gcp_secret = GcpSecret('bad_path/versions/latest') with self.assertRaisesRegex(RuntimeError, r'Failed to retrieve secret bytes'): with TestPipeline() as pipeline: pcoll_1 = pipeline | 'Start 1' >> beam.Create([('a', 1), ('a', 2), ('b', 3), ('c', 4)]) result = (pcoll_1) | beam.GroupByEncryptedKey(gcp_secret) assert_that( result, equal_to([('a', ([1, 2])), ('b', ([3])), ('c', ([4]))])) @unittest.skipIf(secretmanager is None, 'GCP dependencies are not installed') class GcpHsmGeneratedSecretTest(unittest.TestCase): def setUp(self): self.mock_secret_manager_client = mock.MagicMock() self.mock_kms_client = mock.MagicMock() # Patch the clients self.secretmanager_patcher = mock.patch( 'google.cloud.secretmanager.SecretManagerServiceClient', return_value=self.mock_secret_manager_client) self.kms_patcher = mock.patch( 'google.cloud.kms.KeyManagementServiceClient', return_value=self.mock_kms_client) self.os_urandom_patcher = mock.patch('os.urandom', return_value=b'0' * 32) self.hkdf_patcher = mock.patch( 'cryptography.hazmat.primitives.kdf.hkdf.HKDF.derive', return_value=b'derived_key') self.secretmanager_patcher.start() self.kms_patcher.start() self.os_urandom_patcher.start() self.hkdf_patcher.start() def tearDown(self): self.secretmanager_patcher.stop() self.kms_patcher.stop() self.os_urandom_patcher.stop() self.hkdf_patcher.stop() def test_happy_path_secret_creation(self): from google.api_core import exceptions as api_exceptions project_id = 'test-project' location_id = 'global' key_ring_id = 'test-key-ring' key_id = 'test-key' job_name = 'test-job' secret = GcpHsmGeneratedSecret( project_id, location_id, key_ring_id, key_id, job_name) # Mock responses for secret creation path self.mock_secret_manager_client.access_secret_version.side_effect = [ api_exceptions.NotFound('not found'), # first check api_exceptions.NotFound('not found'), # second check mock.MagicMock(payload=mock.MagicMock(data=b'derived_key')) ] self.mock_kms_client.encrypt.return_value = mock.MagicMock( ciphertext=b'encrypted_nonce') secret_bytes = secret.get_secret_bytes() self.assertEqual(secret_bytes, b'derived_key') # Assertions on mocks secret_version_path = ( f'projects/{project_id}/secrets/{secret._secret_version_name}' '/versions/1') self.mock_secret_manager_client.access_secret_version.assert_any_call( request={'name': secret_version_path}) self.assertEqual( self.mock_secret_manager_client.access_secret_version.call_count, 3) self.mock_secret_manager_client.create_secret.assert_called_once() self.mock_kms_client.encrypt.assert_called_once() self.mock_secret_manager_client.add_secret_version.assert_called_once() def test_secret_already_exists(self): from google.api_core import exceptions as api_exceptions project_id = 'test-project' location_id = 'global' key_ring_id = 'test-key-ring' key_id = 'test-key' job_name = 'test-job' secret = GcpHsmGeneratedSecret( project_id, location_id, key_ring_id, key_id, job_name) # Mock responses for secret creation path self.mock_secret_manager_client.access_secret_version.side_effect = [ api_exceptions.NotFound('not found'), api_exceptions.NotFound('not found'), mock.MagicMock(payload=mock.MagicMock(data=b'derived_key')) ] self.mock_secret_manager_client.create_secret.side_effect = ( api_exceptions.AlreadyExists('exists')) self.mock_kms_client.encrypt.return_value = mock.MagicMock( ciphertext=b'encrypted_nonce') secret_bytes = secret.get_secret_bytes() self.assertEqual(secret_bytes, b'derived_key') # Assertions on mocks self.mock_secret_manager_client.create_secret.assert_called_once() self.mock_secret_manager_client.add_secret_version.assert_called_once() def test_secret_version_already_exists(self): project_id = 'test-project' location_id = 'global' key_ring_id = 'test-key-ring' key_id = 'test-key' job_name = 'test-job' secret = GcpHsmGeneratedSecret( project_id, location_id, key_ring_id, key_id, job_name) self.mock_secret_manager_client.access_secret_version.return_value = ( mock.MagicMock(payload=mock.MagicMock(data=b'existing_dek'))) secret_bytes = secret.get_secret_bytes() self.assertEqual(secret_bytes, b'existing_dek') # Assertions self.mock_secret_manager_client.access_secret_version.assert_called_once() self.mock_secret_manager_client.create_secret.assert_not_called() self.mock_secret_manager_client.add_secret_version.assert_not_called() self.mock_kms_client.encrypt.assert_not_called() class FakeClock(object): def __init__(self, now=time.time()): self._now = now def __call__(self): return self._now def sleep(self, duration): self._now += duration class BatchElementsTest(unittest.TestCase): NUM_ELEMENTS = 10 BATCH_SIZE = 5 @staticmethod def _create_test_data(): scientists = [ "Einstein", "Darwin", "Copernicus", "Pasteur", "Curie", "Faraday", "Newton", "Bohr", "Galilei", "Maxwell" ] data = [] for i in range(BatchElementsTest.NUM_ELEMENTS): index = i % len(scientists) data.append(scientists[index]) return data def test_constant_batch(self): # Assumes a single bundle... p = TestPipeline() output = ( p | beam.Create(range(35)) | util.BatchElements(min_batch_size=10, max_batch_size=10) | beam.Map(len)) assert_that(output, equal_to([10, 10, 10, 5])) res = p.run() res.wait_until_finish() metrics = res.metrics() results = metrics.query(MetricsFilter().with_name("batch_size")) self.assertEqual(len(results["distributions"]), 1) def test_constant_batch_no_metrics(self): p = TestPipeline() output = ( p | beam.Create(range(35)) | util.BatchElements( min_batch_size=10, max_batch_size=10, record_metrics=False) | beam.Map(len)) assert_that(output, equal_to([10, 10, 10, 5])) res = p.run() res.wait_until_finish() metrics = res.metrics() results = metrics.query(MetricsFilter().with_name("batch_size")) self.assertEqual(len(results["distributions"]), 0) def test_grows_to_max_batch(self): # Assumes a single bundle, so we pin to the FnApiRunner with TestPipeline('FnApiRunner') as p: res = ( p | beam.Create(range(164)) | util.BatchElements( min_batch_size=1, max_batch_size=50, clock=FakeClock()) | beam.Map(len)) assert_that(res, equal_to([1, 1, 2, 4, 8, 16, 32, 50, 50])) def test_windowed_batches(self): # Assumes a single bundle in order, so we pin to the FnApiRunner with TestPipeline('FnApiRunner') as p: res = ( p | beam.Create(range(47), reshuffle=False) | beam.Map(lambda t: window.TimestampedValue(t, t)) | beam.WindowInto(window.FixedWindows(30)) | util.BatchElements( min_batch_size=5, max_batch_size=10, clock=FakeClock()) | beam.Map(len)) assert_that( res, equal_to([ 5, 5, 10, 10, # elements in [0, 30) 10, 7, # elements in [30, 47) ])) def test_global_batch_timestamps(self): # Assumes a single bundle, so we pin to the FnApiRunner with TestPipeline('FnApiRunner') as p: res = ( p | beam.Create(range(3), reshuffle=False) | util.BatchElements(min_batch_size=2, max_batch_size=2) | beam.Map( lambda batch, timestamp=beam.DoFn.TimestampParam: (len(batch), timestamp))) assert_that( res, equal_to([ (2, GlobalWindow().max_timestamp()), (1, GlobalWindow().max_timestamp()), ])) def test_sized_batches(self): with TestPipeline() as p: res = ( p | beam.Create( [ 'a', 'a', # First batch. 'aaaaaaaaaa', # Second batch. 'aaaaa', 'aaaaa', # Third batch. 'a', 'aaaaaaa', 'a', 'a' # Fourth batch. ], reshuffle=False) | util.BatchElements( min_batch_size=10, max_batch_size=10, element_size_fn=len) | beam.Map(lambda batch: ''.join(batch)) | beam.Map(len)) assert_that(res, equal_to([2, 10, 10, 10])) def test_sized_windowed_batches(self): # Assumes a single bundle, in order so we pin to the FnApiRunner with TestPipeline('FnApiRunner') as p: res = ( p | beam.Create(range(1, 8), reshuffle=False) | beam.Map(lambda t: window.TimestampedValue('a' * t, t)) | beam.WindowInto(window.FixedWindows(3)) | util.BatchElements( min_batch_size=11, max_batch_size=11, element_size_fn=len, clock=FakeClock()) | beam.Map(lambda batch: ''.join(batch))) assert_that( res, equal_to([ 'a' * (1 + 2), # Elements in [1, 3) 'a' * (3 + 4), # Elements in [3, 6) 'a' * 5, 'a' * 6, # Elements in [6, 9) 'a' * 7, ])) def test_target_duration(self): clock = FakeClock() batch_estimator = util._BatchSizeEstimator( target_batch_overhead=None, target_batch_duration_secs=10, clock=clock) batch_duration = lambda batch_size: 1 + .7 * batch_size # 14 * .7 is as close as we can get to 10 as possible. expected_sizes = [1, 2, 4, 8, 14, 14, 14] actual_sizes = [] for _ in range(len(expected_sizes)): actual_sizes.append(batch_estimator.next_batch_size()) with batch_estimator.record_time(actual_sizes[-1]): clock.sleep(batch_duration(actual_sizes[-1])) self.assertEqual(expected_sizes, actual_sizes) def test_target_duration_including_fixed_cost(self): clock = FakeClock() batch_estimator = util._BatchSizeEstimator( target_batch_overhead=None, target_batch_duration_secs_including_fixed_cost=10, clock=clock) batch_duration = lambda batch_size: 1 + .7 * batch_size # 1 + 14 * .7 is as close as we can get to 10 as possible. expected_sizes = [1, 2, 4, 8, 12, 12, 12] actual_sizes = [] for _ in range(len(expected_sizes)): actual_sizes.append(batch_estimator.next_batch_size()) with batch_estimator.record_time(actual_sizes[-1]): clock.sleep(batch_duration(actual_sizes[-1])) self.assertEqual(expected_sizes, actual_sizes) def test_target_overhead(self): clock = FakeClock() batch_estimator = util._BatchSizeEstimator( target_batch_overhead=.05, target_batch_duration_secs=None, clock=clock) batch_duration = lambda batch_size: 1 + .7 * batch_size # At 27 items, a batch takes ~20 seconds with 5% (~1 second) overhead. expected_sizes = [1, 2, 4, 8, 16, 27, 27, 27] actual_sizes = [] for _ in range(len(expected_sizes)): actual_sizes.append(batch_estimator.next_batch_size()) with batch_estimator.record_time(actual_sizes[-1]): clock.sleep(batch_duration(actual_sizes[-1])) self.assertEqual(expected_sizes, actual_sizes) def test_variance(self): clock = FakeClock() variance = 0.25 batch_estimator = util._BatchSizeEstimator( target_batch_overhead=.05, target_batch_duration_secs=None, variance=variance, clock=clock) batch_duration = lambda batch_size: 1 + .7 * batch_size expected_target = 27 actual_sizes = [] for _ in range(util._BatchSizeEstimator._MAX_DATA_POINTS - 1): actual_sizes.append(batch_estimator.next_batch_size()) with batch_estimator.record_time(actual_sizes[-1]): clock.sleep(batch_duration(actual_sizes[-1])) # Check that we're testing a good range of values. stable_set = set(actual_sizes[-20:]) self.assertGreater(len(stable_set), 3) self.assertGreater( min(stable_set), expected_target - expected_target * variance) self.assertLess( max(stable_set), expected_target + expected_target * variance) def test_ignore_first_n_batch_size(self): clock = FakeClock() batch_estimator = util._BatchSizeEstimator( clock=clock, ignore_first_n_seen_per_batch_size=2) expected_sizes = [ 1, 1, 1, 2, 2, 2, 4, 4, 4, 8, 8, 8, 16, 16, 16, 32, 32, 32, 64, 64, 64 ] actual_sizes = [] for i in range(len(expected_sizes)): actual_sizes.append(batch_estimator.next_batch_size()) with batch_estimator.record_time(actual_sizes[-1]): if i % 3 == 2: clock.sleep(0.01) else: clock.sleep(1) self.assertEqual(expected_sizes, actual_sizes) # Check we only record the third timing. expected_data_batch_sizes = [1, 2, 4, 8, 16, 32, 64] actual_data_batch_sizes = [x[0] for x in batch_estimator._data] self.assertEqual(expected_data_batch_sizes, actual_data_batch_sizes) expected_data_timing = [0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01] for i in range(len(expected_data_timing)): self.assertAlmostEqual( expected_data_timing[i], batch_estimator._data[i][1]) def test_ignore_next_timing(self): clock = FakeClock() batch_estimator = util._BatchSizeEstimator(clock=clock) batch_estimator.ignore_next_timing() expected_sizes = [1, 1, 2, 4, 8, 16] actual_sizes = [] for i in range(len(expected_sizes)): actual_sizes.append(batch_estimator.next_batch_size()) with batch_estimator.record_time(actual_sizes[-1]): if i == 0: clock.sleep(1) else: clock.sleep(0.01) self.assertEqual(expected_sizes, actual_sizes) # Check the first record_time was skipped. expected_data_batch_sizes = [1, 2, 4, 8, 16] actual_data_batch_sizes = [x[0] for x in batch_estimator._data] self.assertEqual(expected_data_batch_sizes, actual_data_batch_sizes) expected_data_timing = [0.01, 0.01, 0.01, 0.01, 0.01] for i in range(len(expected_data_timing)): self.assertAlmostEqual( expected_data_timing[i], batch_estimator._data[i][1]) def _run_regression_test(self, linear_regression_fn, test_outliers): xs = [random.random() for _ in range(10)] ys = [2 * x + 1 for x in xs] a, b = linear_regression_fn(xs, ys) self.assertAlmostEqual(a, 1) self.assertAlmostEqual(b, 2) xs = [1 + random.random() for _ in range(100)] ys = [7 * x + 5 + 0.01 * random.random() for x in xs] a, b = linear_regression_fn(xs, ys) self.assertAlmostEqual(a, 5, delta=0.02) self.assertAlmostEqual(b, 7, delta=0.02) # Test repeated xs xs = [1 + random.random()] * 100 ys = [7 * x + 5 + 0.01 * random.random() for x in xs] a, b = linear_regression_fn(xs, ys) self.assertAlmostEqual(a, 0, delta=0.02) self.assertAlmostEqual(b, sum(ys) / (len(ys) * xs[0]), delta=0.02) if test_outliers: xs = [1 + random.random() for _ in range(100)] ys = [2 * x + 1 for x in xs] a, b = linear_regression_fn(xs, ys) self.assertAlmostEqual(a, 1) self.assertAlmostEqual(b, 2) # An outlier or two doesn't affect the result. for _ in range(2): xs += [10] ys += [30] a, b = linear_regression_fn(xs, ys) self.assertAlmostEqual(a, 1) self.assertAlmostEqual(b, 2) # But enough of them, and they're no longer outliers. xs += [10] * 10 ys += [30] * 10 a, b = linear_regression_fn(xs, ys) self.assertLess(a, 0.5) self.assertGreater(b, 2.5) def test_no_numpy_regression(self): self._run_regression_test( util._BatchSizeEstimator.linear_regression_no_numpy, False) def test_numpy_regression(self): try: # pylint: disable=wrong-import-order, wrong-import-position import numpy as _ except ImportError: self.skipTest('numpy not available') self._run_regression_test( util._BatchSizeEstimator.linear_regression_numpy, True) def test_stateful_constant_batch(self): # Assumes a single bundle, so we pin to the FnApiRunner p = TestPipeline('FnApiRunner') output = ( p | beam.Create(range(35)) | util.BatchElements( min_batch_size=10, max_batch_size=10, max_batch_duration_secs=100) | beam.Map(len)) assert_that(output, equal_to([10, 10, 10, 5])) res = p.run() res.wait_until_finish() def test_stateful_in_global_window(self): with TestPipeline() as pipeline: collection = pipeline \ | beam.Create( BatchElementsTest._create_test_data()) \ | util.BatchElements( min_batch_size=BatchElementsTest.BATCH_SIZE, max_batch_size=BatchElementsTest.BATCH_SIZE, max_batch_duration_secs=100) num_batches = collection | beam.combiners.Count.Globally() assert_that( num_batches, equal_to([ int( math.ceil( BatchElementsTest.NUM_ELEMENTS / BatchElementsTest.BATCH_SIZE)) ])) def test_stateful_buffering_timer_in_fixed_window_streaming(self): window_duration = 6 max_buffering_duration_secs = 100 start_time = timestamp.Timestamp(0) test_stream = ( TestStream().add_elements([ TimestampedValue(value, start_time + i) for i, value in enumerate(BatchElementsTest._create_test_data()) ]).advance_processing_time(150).advance_watermark_to( start_time + window_duration).advance_watermark_to( start_time + window_duration + 1).advance_watermark_to_infinity()) with TestPipeline(options=StandardOptions(streaming=True)) as pipeline: # To trigger the processing time timer, use a fake clock with start time # being Timestamp(0). fake_clock = FakeClock(now=start_time) num_elements_per_batch = ( pipeline | test_stream | "fixed window" >> WindowInto(FixedWindows(window_duration)) | util.BatchElements( min_batch_size=BatchElementsTest.BATCH_SIZE, max_batch_size=BatchElementsTest.BATCH_SIZE, max_batch_duration_secs=max_buffering_duration_secs, clock=fake_clock) | "count elements in batch" >> Map(lambda x: (None, len(x))) | GroupByKey() | "global window" >> WindowInto(GlobalWindows()) | FlatMapTuple(lambda k, vs: vs)) # Window duration is 6 and batch size is 5, so output batch size # should be 5 (flush because of batch size reached). expected_0 = 5 # There is only one element left in the window so batch size # should be 1 (flush because of max buffering duration reached). expected_1 = 1 # Collection has 10 elements, there are only 4 left, so batch size should # be 4 (flush because of end of window reached). expected_2 = 4 assert_that( num_elements_per_batch, equal_to([expected_0, expected_1, expected_2]), "assert2") def test_stateful_buffering_timer_in_global_window_streaming(self): max_buffering_duration_secs = 42 start_time = timestamp.Timestamp(0) test_stream = TestStream().advance_watermark_to(start_time) for i, value in enumerate(BatchElementsTest._create_test_data()): test_stream.add_elements( [TimestampedValue(value, start_time + i)]) \ .advance_processing_time(5) test_stream.advance_watermark_to( start_time + BatchElementsTest.NUM_ELEMENTS + 1) \ .advance_watermark_to_infinity() with TestPipeline(options=StandardOptions(streaming=True)) as pipeline: # Set a batch size larger than the total number of elements. # Since we're in a global window, we would have been waiting # for all the elements to arrive without the buffering time limit. batch_size = BatchElementsTest.NUM_ELEMENTS * 2 # To trigger the processing time timer, use a fake clock with start time # being Timestamp(0). Since the fake clock never really advances during # the pipeline execution, meaning that the timer is always set to the same # value, the timer will be fired on every element after the first firing. fake_clock = FakeClock(now=start_time) num_elements_per_batch = ( pipeline | test_stream | WindowInto( GlobalWindows(), trigger=Repeatedly(AfterCount(1)), accumulation_mode=trigger.AccumulationMode.DISCARDING) | util.BatchElements( min_batch_size=batch_size, max_batch_size=batch_size, max_batch_duration_secs=max_buffering_duration_secs, clock=fake_clock) | 'count elements in batch' >> Map(lambda x: (None, len(x))) | GroupByKey() | FlatMapTuple(lambda k, vs: vs)) # We will flush twice when the max buffering duration is reached and when # the global window ends. assert_that(num_elements_per_batch, equal_to([9, 1])) def test_stateful_grows_to_max_batch(self): # Assumes a single bundle, so we pin to the FnApiRunner with TestPipeline('FnApiRunner') as p: res = ( p | beam.Create(range(164)) | util.BatchElements( min_batch_size=1, max_batch_size=50, max_batch_duration_secs=100, clock=FakeClock()) | beam.Map(len)) assert_that(res, equal_to([1, 1, 2, 4, 8, 16, 32, 50, 50])) def test_length_bucket_assignment(self): """WithLengthBucketKey assigns correct bucket indices.""" boundaries = [10, 50, 100] dofn = util.WithLengthBucketKey(length_fn=len, bucket_boundaries=boundaries) # bisect_right: boundaries are lower-inclusive. # e.g., for boundaries [10, 50, 100], buckets are: # (-inf, 10), [10, 50), [50, 100), [100, inf) self.assertEqual(dofn._get_bucket(5), 0) self.assertEqual(dofn._get_bucket(10), 1) self.assertEqual(dofn._get_bucket(11), 1) self.assertEqual(dofn._get_bucket(50), 2) self.assertEqual(dofn._get_bucket(51), 2) self.assertEqual(dofn._get_bucket(100), 3) self.assertEqual(dofn._get_bucket(101), 3) self.assertEqual(dofn._get_bucket(999), 3) def test_stateful_length_aware_constant_batch(self): """Elements in distinct length groups produce separate batches.""" # Create short strings (len 1-5) and long strings (len 50-55) short = ['x' * i for i in range(1, 6)] * 4 # 20 short strings long = ['y' * i for i in range(50, 56)] * 4 # 24 long strings elements = short + long p = TestPipeline('FnApiRunner') batches = ( p | beam.Create(elements) | util.BatchElements( min_batch_size=5, max_batch_size=10, max_batch_duration_secs=100, length_fn=len, bucket_boundaries=[10, 50])) # Verify that no batch mixes short and long elements def check_no_mixing(batch): lengths = [len(s) for s in batch] min_len, max_len = min(lengths), max(lengths) # Within a bucket, all elements should have similar length assert max_len - min_len < 50, ( f'Batch mixed short and long: lengths {lengths}') return True checks = batches | beam.Map(check_no_mixing) assert_that(checks, is_not_empty()) res = p.run() res.wait_until_finish() def test_stateful_length_aware_default_boundaries(self): """Default boundaries [16, 32, 64, 128, 256, 512] are applied.""" be = util.BatchElements(max_batch_duration_secs=100, length_fn=len) self.assertEqual(be._bucket_boundaries, [16, 32, 64, 128, 256, 512]) def test_length_aware_requires_length_fn(self): """bucket_boundaries without length_fn raises ValueError.""" with self.assertRaises(ValueError): util.BatchElements( max_batch_duration_secs=100, bucket_boundaries=[10, 20]) def test_bucket_boundaries_must_be_sorted(self): """Unsorted boundaries raise ValueError.""" with self.assertRaises(ValueError): util.BatchElements( max_batch_duration_secs=100, length_fn=len, bucket_boundaries=[50, 10, 100]) def test_bucket_boundaries_must_be_positive(self): """Non-positive boundaries raise ValueError.""" with self.assertRaises(ValueError): util.BatchElements( max_batch_duration_secs=100, length_fn=len, bucket_boundaries=[0, 10, 100]) def test_length_fn_without_stateful_is_ignored(self): """length_fn without max_batch_duration_secs uses non-stateful path.""" with TestPipeline() as p: res = ( p | beam.Create(['a', 'bb', 'ccc']) | util.BatchElements( min_batch_size=3, max_batch_size=3, length_fn=len) | beam.Map(len)) assert_that(res, equal_to([3])) def test_padding_efficiency_bimodal(self): """Benchmark: length-aware bucketing yields better padding efficiency than unbucketed batching on a bimodal length distribution. Padding efficiency per batch = sum(lengths) / (max_len * batch_size). With bucketing, short and long elements land in separate batches, so each batch pads to a smaller max, improving efficiency. """ random.seed(42) short = ['x' * random.randint(5, 30) for _ in range(500)] long = ['y' * random.randint(200, 512) for _ in range(500)] elements = short + long batch_size = 32 def batch_efficiency(batch): """Returns (useful_tokens, padded_tokens) for one batch.""" lengths = [len(s) for s in batch] return (sum(lengths), max(lengths) * len(lengths)) # Run WITH bucketing — collect (useful, padded) per batch p_bucketed = TestPipeline('FnApiRunner') bucketed_eff = ( p_bucketed | 'CreateBucketed' >> beam.Create(elements) | 'BatchBucketed' >> util.BatchElements( min_batch_size=batch_size, max_batch_size=batch_size, max_batch_duration_secs=100, length_fn=len, bucket_boundaries=[16, 32, 64, 128, 256, 512]) | 'EffBucketed' >> beam.Map(batch_efficiency) | 'SumBucketed' >> beam.CombineGlobally( lambda pairs: (sum(p[0] for p in pairs), sum(p[1] for p in pairs)))) # Run WITHOUT bucketing p_unbucketed = TestPipeline('FnApiRunner') unbucketed_eff = ( p_unbucketed | 'CreateUnbucketed' >> beam.Create(elements) | 'BatchUnbucketed' >> util.BatchElements( min_batch_size=batch_size, max_batch_size=batch_size, max_batch_duration_secs=100) | 'EffUnbucketed' >> beam.Map(batch_efficiency) | 'SumUnbucketed' >> beam.CombineGlobally( lambda pairs: (sum(p[0] for p in pairs), sum(p[1] for p in pairs)))) def check_bucketed_above_threshold(totals): useful, padded = totals[0] eff = useful / padded if padded else 0 assert eff > 0.70, ( f'Bucketed padding efficiency {eff:.2%} should be > 70%') def check_unbucketed_below_bucketed(totals): useful, padded = totals[0] eff = useful / padded if padded else 0 # With bimodal data in a single key, short elements get padded # to the max of each batch which often includes long elements. assert eff < 0.70, ( f'Unbucketed efficiency {eff:.2%} expected < 70% for ' f'bimodal distribution (sanity check)') assert_that(bucketed_eff, check_bucketed_above_threshold) res = p_bucketed.run() res.wait_until_finish() assert_that(unbucketed_eff, check_unbucketed_below_bucketed) res = p_unbucketed.run() res.wait_until_finish() def test_with_length_bucket_key_setup_and_process(self): """WithLengthBucketKey.setup() and process() work correctly in pipeline.""" boundaries = [10, 50] elements = ['short', 'x' * 30, 'y' * 60] with TestPipeline('FnApiRunner') as p: result = ( p | beam.Create(elements) | beam.ParDo(util.WithLengthBucketKey(len, boundaries))) def check_keys(keyed_elements): # Each element should have format ((worker_key, bucket), element) for (key, bucket), elem in keyed_elements: # Verify key is a UUID string assert isinstance(key, str) and len(key) > 0 # Verify bucket is correct if len(elem) < 10: assert bucket == 0, f'Expected bucket 0 for {elem}' elif len(elem) < 50: assert bucket == 1, f'Expected bucket 1 for {elem}' else: assert bucket == 2, f'Expected bucket 2 for {elem}' assert_that(result, check_keys) def test_bucket_boundaries_empty_list(self): """Empty bucket_boundaries list raises ValueError.""" with self.assertRaises(ValueError): util.BatchElements( max_batch_duration_secs=100, length_fn=len, bucket_boundaries=[]) def test_with_custom_bucket_boundaries(self): """Custom bucket_boundaries are used instead of defaults.""" custom_boundaries = [5, 15, 25] be = util.BatchElements( max_batch_duration_secs=100, length_fn=len, bucket_boundaries=custom_boundaries) self.assertEqual(be._bucket_boundaries, custom_boundaries) def test_length_fn_applied_in_pipeline(self): """Verify length_fn is used for bucketing in stateful batching.""" # Create strings of different lengths that should go to different buckets short_strings = ['x' * i for i in range(1, 5)] # lengths 1-4, bucket 0 medium_strings = ['y' * i for i in range(20, 24)] # lengths 20-23, bucket 1 elements = short_strings + medium_strings with TestPipeline('FnApiRunner') as p: batches = ( p | beam.Create(elements) | util.BatchElements( min_batch_size=2, max_batch_size=10, max_batch_duration_secs=100, length_fn=len, bucket_boundaries=[10, 30])) def check_batch_homogeneity(batch): """Batches should contain elements of similar length.""" lengths = [len(s) for s in batch] # If bucketing works, all elements should be in same bucket # (either all < 10 or all between 10 and 30) min_len, max_len = min(lengths), max(lengths) if min_len < 10: # Short bucket: all should be < 10 assert max_len < 10, f'Mixed batch: {lengths}' else: # Medium bucket: all should be >= 10 assert min_len >= 10, f'Mixed batch: {lengths}' return True checks = batches | beam.Map(check_batch_homogeneity) assert_that(checks, is_not_empty()) class SortAndBatchElementsTest(unittest.TestCase): """Tests for SortAndBatchElements transform.""" def test_elements_are_sorted_by_size(self): """Test that elements are sorted by size within batches.""" with TestPipeline() as p: # Create elements with varying sizes data = ['aaaaa', 'bb', 'cccc', 'a', 'ddd'] expected = [['a', 'bb', 'ddd', 'cccc', 'aaaaa']] res = ( p | beam.Create(data, reshuffle=False) | util.SortAndBatchElements( min_batch_size=1, max_batch_size=5, max_batch_weight=100)) # All elements fit in one batch, so the expected order is explicit. assert_that(res, equal_to(expected)) def test_batch_respects_max_batch_size(self): """Test that batches do not exceed max_batch_size.""" with TestPipeline() as p: res = ( p | beam.Create(['a'] * 10, reshuffle=False) | util.SortAndBatchElements( min_batch_size=1, max_batch_size=3, max_batch_weight=100) | beam.Map(len)) assert_that(res, equal_to([3, 3, 3, 1])) def test_batch_respects_max_batch_weight(self): """Test that batches do not exceed max_batch_weight.""" with TestPipeline() as p: # Each element has size 5, max_batch_weight is 12 # So we can fit at most 2 elements per batch data = ['aaaaa', 'bbbbb', 'ccccc', 'ddddd'] res = ( p | beam.Create(data, reshuffle=False) | util.SortAndBatchElements( min_batch_size=1, max_batch_size=10, max_batch_weight=12) | beam.Map(len)) assert_that(res, equal_to([2, 2])) def test_default_element_size_fn_with_strings(self): """Test default element_size_fn works with strings.""" with TestPipeline() as p: data = ['a', 'bbb', 'cc'] res = ( p | beam.Create(data, reshuffle=False) | util.SortAndBatchElements( min_batch_size=1, max_batch_size=3, max_batch_weight=100) | beam.FlatMap(lambda batch: [len(s) for s in batch])) # Elements should be sorted by length: 'a'(1), 'cc'(2), 'bbb'(3) assert_that(res, equal_to([1, 2, 3])) def test_default_element_size_fn_with_integers(self): """Test default element_size_fn falls back to 1 for integers.""" with TestPipeline() as p: data = [10, 20, 30, 40, 50] res = ( p | beam.Create(data, reshuffle=False) | util.SortAndBatchElements( min_batch_size=1, max_batch_size=3, max_batch_weight=100) | beam.Map(len)) # With size=1 for all, should batch by max_batch_size assert_that(res, equal_to([3, 2])) def test_custom_element_size_fn(self): """Test using a custom element_size_fn.""" with TestPipeline() as p: data = [{'text': 'a'}, {'text': 'bbb'}, {'text': 'cc'}] res = ( p | beam.Create(data, reshuffle=False) | util.SortAndBatchElements( min_batch_size=1, max_batch_size=3, max_batch_weight=100, element_size_fn=lambda x: len(x['text'])) | beam.FlatMap(lambda batch: [len(e['text']) for e in batch])) # Should be sorted by text length assert_that(res, equal_to([1, 2, 3])) def test_empty_input(self): """Test with empty input produces no output.""" with TestPipeline() as p: res = ( p | beam.Create([], reshuffle=False) | util.SortAndBatchElements( min_batch_size=1, max_batch_size=10, max_batch_weight=100) | beam.Map(len)) assert_that(res, equal_to([])) def test_single_element(self): """Test with a single element.""" with TestPipeline() as p: res = ( p | beam.Create(['hello'], reshuffle=False) | util.SortAndBatchElements( min_batch_size=1, max_batch_size=10, max_batch_weight=100)) assert_that(res, equal_to([['hello']])) def test_windowed_batches(self): """Test that windowed elements are batched per window.""" with TestPipeline('FnApiRunner') as p: res = ( p | beam.Create(range(1, 8), reshuffle=False) | beam.Map(lambda t: window.TimestampedValue('a' * t, t)) | beam.WindowInto(window.FixedWindows(3)) | util.SortAndBatchElements( min_batch_size=1, max_batch_size=10, max_batch_weight=100) | beam.Map(lambda batch: ''.join(batch))) # FixedWindows(3) with default offset 0 produces: # Window [0, 3): elements at t=1,2 with sizes 1,2 # Window [3, 6): elements at t=3,4,5 with sizes 3,4,5 # Window [6, 9): elements at t=6,7 with sizes 6,7 assert_that( res, equal_to([ 'a' * (1 + 2), # Window [0, 3) 'a' * (3 + 4 + 5), # Window [3, 6) 'a' * (6 + 7), # Window [6, 9) ])) def test_validation_min_batch_size(self): """Test that min_batch_size validation raises ValueError.""" with self.assertRaises(ValueError) as cm: util.SortAndBatchElements( min_batch_size=0, max_batch_size=10, max_batch_weight=100) self.assertIn('min_batch_size must be >= 1', str(cm.exception)) def test_validation_max_batch_size(self): """Test that max_batch_size < min_batch_size raises ValueError.""" with self.assertRaises(ValueError) as cm: util.SortAndBatchElements( min_batch_size=10, max_batch_size=5, max_batch_weight=100) self.assertIn('max_batch_size', str(cm.exception)) self.assertIn('min_batch_size', str(cm.exception)) def test_validation_max_batch_weight(self): """Test that max_batch_weight validation raises ValueError.""" with self.assertRaises(ValueError) as cm: util.SortAndBatchElements( min_batch_size=1, max_batch_size=10, max_batch_weight=0) self.assertIn('max_batch_weight must be >= 1', str(cm.exception)) def test_validation_element_size_fn_callable(self): """Test that a non-callable element_size_fn raises TypeError.""" with self.assertRaises(TypeError) as cm: util.SortAndBatchElements( min_batch_size=1, max_batch_size=10, max_batch_weight=100, element_size_fn=123) self.assertIn('element_size_fn must be callable', str(cm.exception)) def test_batch_timestamps(self): """Test that batches have correct timestamps.""" with TestPipeline('FnApiRunner') as p: res = ( p | beam.Create(['a', 'bb', 'ccc'], reshuffle=False) | util.SortAndBatchElements( min_batch_size=1, max_batch_size=10, max_batch_weight=100) | beam.Map(lambda batch, ts=beam.DoFn.TimestampParam: (len(batch), ts))) # The single global-window batch is emitted at end-of-window. expected = [(3, GlobalWindow().max_timestamp())] assert_that(res, equal_to(expected)) class SortAndBatchElementsDoFnDirectTest(unittest.TestCase): """Direct unit tests for DoFn internals to ensure coverage. Beam's FnApiRunner executes DoFns in a separate SDK harness process, so coverage tools in the main process cannot capture DoFn code paths. These tests exercise the DoFn methods directly in-process. """ def test_default_element_size_len(self): from apache_beam.transforms.util import _SortAndBatchElementsDoFn dofn = _SortAndBatchElementsDoFn( min_batch_size=1, max_batch_size=10, max_batch_weight=100, element_size_fn=None) self.assertEqual(dofn._element_size_fn('abc'), 3) self.assertEqual(dofn._element_size_fn([1, 2]), 2) def test_default_element_size_fallback_warns_once(self): from apache_beam.transforms.util import _SortAndBatchElementsDoFn dofn = _SortAndBatchElementsDoFn( min_batch_size=1, max_batch_size=10, max_batch_weight=100, element_size_fn=None) with self.assertLogs('apache_beam.transforms.util', level='WARNING') as cm: self.assertEqual(dofn._element_size_fn(42), 1) self.assertIn('does not support len()', cm.output[0]) # Second call should not warn again self.assertEqual(dofn._element_size_fn(3.14), 1) self.assertTrue(dofn._has_warned_type_error) def test_global_dofn_sort_and_batch(self): """Test _SortAndBatchElementsDoFn directly.""" from apache_beam.transforms.util import _SortAndBatchElementsDoFn dofn = _SortAndBatchElementsDoFn( min_batch_size=1, max_batch_size=3, max_batch_weight=100, element_size_fn=len) dofn.start_bundle() for elem in ['ccccc', 'bb', 'dddd', 'a', 'eee']: dofn.process(elem) batches = [wv.value for wv in dofn.finish_bundle()] # All elements emitted self.assertEqual(sum(len(b) for b in batches), 5) # Each batch respects max_batch_size=3 for batch in batches: self.assertLessEqual(len(batch), 3) # Elements within each batch are sorted by size for batch in batches: lengths = [len(s) for s in batch] self.assertEqual(lengths, sorted(lengths)) def test_global_dofn_empty_bundle(self): """Test finish_bundle with no elements returns nothing.""" from apache_beam.transforms.util import _SortAndBatchElementsDoFn dofn = _SortAndBatchElementsDoFn( min_batch_size=1, max_batch_size=10, max_batch_weight=100, element_size_fn=len) dofn.start_bundle() result = list(dofn.finish_bundle() or []) self.assertEqual(result, []) def test_global_dofn_weight_splitting(self): """Test weight-based splitting in the global DoFn.""" from apache_beam.transforms.util import _SortAndBatchElementsDoFn # Each element has size 5, max_batch_weight=12 -> 2 per batch dofn = _SortAndBatchElementsDoFn( min_batch_size=1, max_batch_size=100, max_batch_weight=12, element_size_fn=len) dofn.start_bundle() for elem in ['aaaaa', 'bbbbb', 'ccccc', 'ddddd']: dofn.process(elem) batches = [wv.value for wv in dofn.finish_bundle()] self.assertEqual(len(batches), 2) for batch in batches: self.assertEqual(len(batch), 2) def test_windowed_dofn_flush_and_finish(self): """Test _WindowAwareSortAndBatchElementsDoFn directly.""" from apache_beam.transforms.util import _WindowAwareSortAndBatchElementsDoFn dofn = _WindowAwareSortAndBatchElementsDoFn( min_batch_size=1, max_batch_size=10, max_batch_weight=100, element_size_fn=len) dofn.start_bundle() win1 = IntervalWindow(0, 3) win2 = IntervalWindow(3, 6) # Manually add to buffers (bypass process() to avoid DoFn.WindowParam) dofn._buffers[win1].extend(['aa', 'b', 'ccc']) dofn._buffers[win2].extend(['dddd', 'ee']) batches = list(dofn.finish_bundle()) # All elements across both windows emitted total_elements = sum(len(wv.value) for wv in batches) self.assertEqual(total_elements, 5) # Each batch has the correct window for wv in batches: self.assertIn(wv.windows[0], (win1, win2)) def test_windowed_dofn_overflow_flush(self): """Test that exceeding _MAX_LIVE_WINDOWS triggers early flush.""" from apache_beam.transforms.util import _WindowAwareSortAndBatchElementsDoFn dofn = _WindowAwareSortAndBatchElementsDoFn( min_batch_size=1, max_batch_size=10, max_batch_weight=100, element_size_fn=len) dofn.start_bundle() # Fill up to _MAX_LIVE_WINDOWS for i in range(dofn._MAX_LIVE_WINDOWS): win = IntervalWindow(i * 10, (i + 1) * 10) dofn._buffers[win].append('x' * (i + 1)) self.assertEqual(len(dofn._buffers), dofn._MAX_LIVE_WINDOWS) # Adding one more window should trigger overflow flush overflow_win = IntervalWindow(100, 110) results = list(dofn.process('overflow', overflow_win)) # One window was flushed, so buffer count stays at _MAX_LIVE_WINDOWS self.assertLessEqual(len(dofn._buffers), dofn._MAX_LIVE_WINDOWS) # The flushed window produced output self.assertGreater(len(results), 0) def test_windowed_dofn_flush_empty_window(self): """Test _flush_window with a non-existent window returns nothing.""" from apache_beam.transforms.util import _WindowAwareSortAndBatchElementsDoFn dofn = _WindowAwareSortAndBatchElementsDoFn( min_batch_size=1, max_batch_size=10, max_batch_weight=100, element_size_fn=len) dofn.start_bundle() result = list(dofn._flush_window(IntervalWindow(0, 10))) self.assertEqual(result, []) def test_windowed_dofn_weight_splitting(self): """Test weight-based splitting in the windowed DoFn.""" from apache_beam.transforms.util import _WindowAwareSortAndBatchElementsDoFn dofn = _WindowAwareSortAndBatchElementsDoFn( min_batch_size=1, max_batch_size=100, max_batch_weight=12, element_size_fn=len) dofn.start_bundle() win = IntervalWindow(0, 10) dofn._buffers[win].extend(['aaaaa', 'bbbbb', 'ccccc', 'ddddd']) batches = list(dofn._flush_window(win)) self.assertEqual(len(batches), 2) for wv in batches: self.assertEqual(len(wv.value), 2) self.assertEqual(wv.windows[0], win) class IdentityWindowTest(unittest.TestCase): def test_window_preserved(self): expected_timestamp = timestamp.Timestamp(5) expected_window = window.IntervalWindow(1.0, 2.0) class AddWindowDoFn(beam.DoFn): def process(self, element): yield WindowedValue(element, expected_timestamp, [expected_window]) with TestPipeline() as pipeline: data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)] expected_windows = [ TestWindowedValue(kv, expected_timestamp, [expected_window]) for kv in data ] before_identity = ( pipeline | 'start' >> beam.Create(data) | 'add_windows' >> beam.ParDo(AddWindowDoFn())) assert_that( before_identity, equal_to(expected_windows), label='before_identity', reify_windows=True) after_identity = ( before_identity | 'window' >> beam.WindowInto( beam.transforms.util._IdentityWindowFn( coders.IntervalWindowCoder()))) assert_that( after_identity, equal_to(expected_windows), label='after_identity', reify_windows=True) def test_no_window_context_fails(self): expected_timestamp = timestamp.Timestamp(5) # Assuming the default window function is window.GlobalWindows. expected_window = window.GlobalWindow() class AddTimestampDoFn(beam.DoFn): def process(self, element): yield window.TimestampedValue(element, expected_timestamp) with self.assertRaisesRegex(Exception, r'.*window.*None.*add_timestamps2'): with TestPipeline() as pipeline: data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)] expected_windows = [ TestWindowedValue(kv, expected_timestamp, [expected_window]) for kv in data ] before_identity = ( pipeline | 'start' >> beam.Create(data) | 'add_timestamps' >> beam.ParDo(AddTimestampDoFn())) assert_that( before_identity, equal_to(expected_windows), label='before_identity', reify_windows=True) _ = ( before_identity | 'window' >> beam.WindowInto( beam.transforms.util._IdentityWindowFn( coders.GlobalWindowCoder())) # This DoFn will return TimestampedValues, making # WindowFn.AssignContext passed to IdentityWindowFn # contain a window of None. IdentityWindowFn should # raise an exception. | 'add_timestamps2' >> beam.ParDo(AddTimestampDoFn())) class ReshuffleTest(unittest.TestCase): def test_reshuffle_contents_unchanged(self): with TestPipeline() as pipeline: data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 3)] result = (pipeline | beam.Create(data) | beam.Reshuffle()) assert_that(result, equal_to(data)) def test_reshuffle_contents_unchanged_with_buckets(self): with TestPipeline() as pipeline: data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 3)] buckets = 2 result = (pipeline | beam.Create(data) | beam.Reshuffle(buckets)) assert_that(result, equal_to(data)) def test_reshuffle_contents_unchanged_with_wrong_buckets(self): wrong_buckets = [0, -1, "wrong", 2.5] for wrong_bucket in wrong_buckets: with self.assertRaisesRegex(ValueError, 'If `num_buckets` is set, it has to be an ' 'integer greater than 0, got %s' % wrong_bucket): beam.Reshuffle(wrong_bucket) def test_reshuffle_after_gbk_contents_unchanged(self): with TestPipeline() as pipeline: data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 3)] expected_result = [(1, [1, 2, 3]), (2, [1, 2]), (3, [1])] after_gbk = ( pipeline | beam.Create(data) | beam.GroupByKey() | beam.MapTuple(lambda k, vs: (k, sorted(vs)))) assert_that(after_gbk, equal_to(expected_result), label='after_gbk') after_reshuffle = after_gbk | beam.Reshuffle() assert_that( after_reshuffle, equal_to(expected_result), label='after_reshuffle') def test_reshuffle_timestamps_unchanged(self): with TestPipeline() as pipeline: timestamp = 5 data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 3)] expected_result = [ TestWindowedValue(v, timestamp, [GlobalWindow()]) for v in data ] before_reshuffle = ( pipeline | 'start' >> beam.Create(data) | 'add_timestamp' >> beam.Map(lambda v: beam.window.TimestampedValue(v, timestamp))) assert_that( before_reshuffle, equal_to(expected_result), label='before_reshuffle', reify_windows=True) after_reshuffle = before_reshuffle | beam.Reshuffle() assert_that( after_reshuffle, equal_to(expected_result), label='after_reshuffle', reify_windows=True) def test_reshuffle_windows_unchanged(self): with TestPipeline() as pipeline: data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)] expected_data = [ TestWindowedValue( v, t - .001, [w], pane_info=PaneInfo(True, False, PaneInfoTiming.ON_TIME, 0, 0)) for (v, t, w) in [((1, contains_in_any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)), ( (2, contains_in_any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)), (( 3, [1]), 3.0, IntervalWindow(1.0, 3.0)), (( 1, [4]), 6.0, IntervalWindow(4.0, 6.0))] ] before_reshuffle = ( pipeline | 'start' >> beam.Create(data) | 'add_timestamp' >> beam.Map(lambda v: beam.window.TimestampedValue(v, v[1])) | 'window' >> beam.WindowInto(Sessions(gap_size=2)) | 'group_by_key' >> beam.GroupByKey()) assert_that( before_reshuffle, equal_to(expected_data), label='before_reshuffle', reify_windows=True) after_reshuffle = before_reshuffle | beam.Reshuffle() assert_that( after_reshuffle, equal_to(expected_data), label='after reshuffle', reify_windows=True) def test_reshuffle_window_fn_preserved(self): any_order = contains_in_any_order with TestPipeline() as pipeline: data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)] expected_windows = [ TestWindowedValue(v, t, [w]) for (v, t, w) in [((1, 1), 1.0, IntervalWindow(1.0, 3.0)), (( 2, 1), 1.0, IntervalWindow(1.0, 3.0)), (( 3, 1), 1.0, IntervalWindow(1.0, 3.0)), (( 1, 2), 2.0, IntervalWindow(2.0, 4.0)), ( (2, 2), 2.0, IntervalWindow(2.0, 4.0)), ((1, 4), 4.0, IntervalWindow(4.0, 6.0))] ] expected_merged_windows = [ TestWindowedValue( v, t - .001, [w], pane_info=PaneInfo(True, False, PaneInfoTiming.ON_TIME, 0, 0)) for (v, t, w) in [((1, any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)), ( (2, any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)), ( (3, [1]), 3.0, IntervalWindow(1.0, 3.0)), ((1, [4]), 6.0, IntervalWindow(4.0, 6.0))] ] before_reshuffle = ( pipeline | 'start' >> beam.Create(data) | 'add_timestamp' >> beam.Map(lambda v: TimestampedValue(v, v[1])) | 'window' >> beam.WindowInto(Sessions(gap_size=2))) assert_that( before_reshuffle, equal_to(expected_windows), label='before_reshuffle', reify_windows=True) after_reshuffle = before_reshuffle | beam.Reshuffle() assert_that( after_reshuffle, equal_to(expected_windows), label='after_reshuffle', reify_windows=True) after_group = after_reshuffle | beam.GroupByKey() assert_that( after_group, equal_to(expected_merged_windows), label='after_group', reify_windows=True) def test_reshuffle_global_window(self): with TestPipeline() as pipeline: data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)] expected_data = [(1, [1, 2, 4]), (2, [1, 2]), (3, [1])] before_reshuffle = ( pipeline | beam.Create(data) | beam.WindowInto(GlobalWindows()) | beam.GroupByKey() | beam.MapTuple(lambda k, vs: (k, sorted(vs)))) assert_that( before_reshuffle, equal_to(expected_data), label='before_reshuffle') after_reshuffle = before_reshuffle | beam.Reshuffle() assert_that( after_reshuffle, equal_to(expected_data), label='after reshuffle') def test_reshuffle_sliding_window(self): with TestPipeline() as pipeline: data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)] window_size = 2 expected_data = [(1, [1, 2, 4]), (2, [1, 2]), (3, [1])] * window_size before_reshuffle = ( pipeline | beam.Create(data) | beam.WindowInto(SlidingWindows(size=window_size, period=1)) | beam.GroupByKey() | beam.MapTuple(lambda k, vs: (k, sorted(vs)))) assert_that( before_reshuffle, equal_to(expected_data), label='before_reshuffle') after_reshuffle = before_reshuffle | beam.Reshuffle() # If Reshuffle applies the sliding window function a second time there # should be extra values for each key. assert_that( after_reshuffle, equal_to(expected_data), label='after reshuffle') def test_reshuffle_streaming_global_window(self): options = PipelineOptions() options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as pipeline: data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)] expected_data = [(1, [1, 2, 4]), (2, [1, 2]), (3, [1])] before_reshuffle = ( pipeline | beam.Create(data) | beam.WindowInto(GlobalWindows()) | beam.GroupByKey() | beam.MapTuple(lambda k, vs: (k, sorted(vs)))) assert_that( before_reshuffle, equal_to(expected_data), label='before_reshuffle') after_reshuffle = before_reshuffle | beam.Reshuffle() assert_that( after_reshuffle, equal_to(expected_data), label='after reshuffle') def test_reshuffle_streaming_global_window_with_buckets(self): options = PipelineOptions() options.view_as(StandardOptions).streaming = True with TestPipeline(options=options) as pipeline: data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)] expected_data = [(1, [1, 2, 4]), (2, [1, 2]), (3, [1])] buckets = 2 before_reshuffle = ( pipeline | beam.Create(data) | beam.WindowInto(GlobalWindows()) | beam.GroupByKey() | beam.MapTuple(lambda k, vs: (k, sorted(vs)))) assert_that( before_reshuffle, equal_to(expected_data), label='before_reshuffle') after_reshuffle = before_reshuffle | beam.Reshuffle(buckets) assert_that( after_reshuffle, equal_to(expected_data), label='after reshuffle') @parameterized.expand([ param(compat_version=None), param(compat_version="2.64.0"), ]) @mock.patch( 'apache_beam.coders.coders._should_force_use_dill', return_value=True) def test_reshuffle_custom_window_preserves_metadata( self, mock_force_dill, compat_version): """Tests that Reshuffle preserves pane info.""" element_count = 12 timestamp_value = timestamp.Timestamp(0) l = [ TimestampedValue(("key", i), timestamp_value) for i in range(element_count) ] expected_timestamp = GlobalWindow().max_timestamp() expected = [ TestWindowedValue( ('key', [0, 1, 2]), expected_timestamp, [GlobalWindow()], pane_info=PaneInfo( is_first=True, is_last=False, timing=PaneInfoTiming.EARLY, # 0 index=0, nonspeculative_index=-1)), TestWindowedValue( ('key', [3, 4, 5]), expected_timestamp, [GlobalWindow()], pane_info=PaneInfo( is_first=False, is_last=False, timing=PaneInfoTiming.EARLY, # 0 index=1, nonspeculative_index=-1)), TestWindowedValue( ('key', [6, 7, 8]), expected_timestamp, [GlobalWindow()], pane_info=PaneInfo( is_first=False, is_last=False, timing=PaneInfoTiming.EARLY, # 0 index=2, nonspeculative_index=-1)), TestWindowedValue( ('key', [9, 10, 11]), expected_timestamp, [GlobalWindow()], pane_info=PaneInfo( is_first=False, is_last=False, timing=PaneInfoTiming.EARLY, # 0 index=3, nonspeculative_index=-1)) ] if compat_version is None else ([ TestWindowedValue(('key', [0, 1, 2]), expected_timestamp, [GlobalWindow()], PANE_INFO_UNKNOWN), TestWindowedValue(('key', [3, 4, 5]), expected_timestamp, [GlobalWindow()], PANE_INFO_UNKNOWN), TestWindowedValue(('key', [6, 7, 8]), expected_timestamp, [GlobalWindow()], PANE_INFO_UNKNOWN), TestWindowedValue(('key', [9, 10, 11]), expected_timestamp, [GlobalWindow()], PANE_INFO_UNKNOWN) ]) options = PipelineOptions(update_compatibility_version=compat_version) options.view_as(StandardOptions).streaming = True with beam.Pipeline(options=options) as p: stream_source = ( TestStream().advance_watermark_to(0).advance_processing_time( 100).add_elements(l[:element_count // 4]).advance_processing_time( 100).advance_watermark_to(100).add_elements( l[element_count // 4:2 * element_count // 4]). advance_processing_time(100).advance_watermark_to(200).add_elements( l[2 * element_count // 4:3 * element_count // 4]).advance_processing_time( 100).advance_watermark_to(300).add_elements( l[3 * element_count // 4:]).advance_processing_time( 100).advance_watermark_to_infinity()) grouped = ( p | stream_source | "Rewindow" >> beam.WindowInto( beam.window.GlobalWindows(), trigger=trigger.Repeatedly(trigger.AfterProcessingTime(1)), accumulation_mode=trigger.AccumulationMode.DISCARDING) | beam.GroupByKey()) after_reshuffle = (grouped | 'Reshuffle' >> beam.Reshuffle()) assert_that( after_reshuffle, equal_to(expected), label='CheckMetadataPreserved', reify_windows=True) @parameterized.expand([ param(compat_version=None), param(compat_version="2.64.0"), ]) @mock.patch( 'apache_beam.coders.coders._should_force_use_dill', return_value=True) def test_reshuffle_default_window_preserves_metadata( self, mock_force_dill, compat_version): """Tests that Reshuffle preserves timestamp, window, and pane info metadata.""" no_firing = PaneInfo( is_first=True, is_last=True, timing=PaneInfoTiming.UNKNOWN, index=0, nonspeculative_index=0) on_time_only = PaneInfo( is_first=True, is_last=True, timing=PaneInfoTiming.ON_TIME, index=0, nonspeculative_index=0) late_firing = PaneInfo( is_first=False, is_last=False, timing=PaneInfoTiming.LATE, index=1, nonspeculative_index=1) # Portable runners may not have the same level of precision on timestamps - # this gets the largest supported timestamp with the extra non-supported # bits truncated gt = GlobalWindow().max_timestamp() truncated_gt = gt - (gt % 0.001) expected_preserved = [ TestWindowedValue('a', MIN_TIMESTAMP, [GlobalWindow()], no_firing), TestWindowedValue( 'b', timestamp.Timestamp(0), [GlobalWindow()], on_time_only), TestWindowedValue( 'c', timestamp.Timestamp(33), [GlobalWindow()], late_firing), TestWindowedValue('d', truncated_gt, [GlobalWindow()], no_firing) ] expected_not_preserved = [ TestWindowedValue( 'a', MIN_TIMESTAMP, [GlobalWindow()], PANE_INFO_UNKNOWN), TestWindowedValue( 'b', timestamp.Timestamp(0), [GlobalWindow()], PANE_INFO_UNKNOWN), TestWindowedValue( 'c', timestamp.Timestamp(33), [GlobalWindow()], PANE_INFO_UNKNOWN), TestWindowedValue( 'd', truncated_gt, [GlobalWindow()], PANE_INFO_UNKNOWN) ] expected = ( expected_preserved if compat_version is None else expected_not_preserved) options = PipelineOptions(update_compatibility_version=compat_version) with TestPipeline(options=options) as pipeline: # Create windowed values with specific metadata elements = [ WindowedValue('a', MIN_TIMESTAMP, [GlobalWindow()], no_firing), WindowedValue( 'b', timestamp.Timestamp(0), [GlobalWindow()], on_time_only), WindowedValue( 'c', timestamp.Timestamp(33), [GlobalWindow()], late_firing), WindowedValue('d', truncated_gt, [GlobalWindow()], no_firing) ] after_reshuffle = ( pipeline | 'Create' >> beam.Create(elements) | 'Reshuffle' >> beam.Reshuffle()) assert_that( after_reshuffle, equal_to(expected), label='CheckMetadataPreserved', reify_windows=True) @pytest.mark.it_validatesrunner def test_reshuffle_preserves_timestamps(self): with TestPipeline() as pipeline: # Create a PCollection and assign each element with a different timestamp. before_reshuffle = ( pipeline | beam.Create([ { 'name': 'foo', 'timestamp': MIN_TIMESTAMP }, { 'name': 'foo', 'timestamp': 0 }, { 'name': 'bar', 'timestamp': 33 }, { 'name': 'bar', 'timestamp': 0 }, ]) | beam.Map( lambda element: beam.window.TimestampedValue( element, element['timestamp']))) # Reshuffle the PCollection above and assign the timestamp of an element # to that element again. after_reshuffle = before_reshuffle | beam.Reshuffle() # Given an element, emits a string which contains the timestamp and the # name field of the element. def format_with_timestamp(element, timestamp=beam.DoFn.TimestampParam): t = str(timestamp) if timestamp == MIN_TIMESTAMP: t = 'MIN_TIMESTAMP' elif timestamp == MAX_TIMESTAMP: t = 'MAX_TIMESTAMP' return '{} - {}'.format(t, element['name']) # Combine each element in before_reshuffle with its timestamp. formatted_before_reshuffle = ( before_reshuffle | "Get before_reshuffle timestamp" >> beam.Map(format_with_timestamp)) # Combine each element in after_reshuffle with its timestamp. formatted_after_reshuffle = ( after_reshuffle | "Get after_reshuffle timestamp" >> beam.Map(format_with_timestamp)) expected_data = [ 'MIN_TIMESTAMP - foo', 'Timestamp(0) - foo', 'Timestamp(33) - bar', 'Timestamp(0) - bar' ] # Can't compare formatted_before_reshuffle and formatted_after_reshuffle # directly, because they are deferred PCollections while equal_to only # takes a concrete argument. assert_that( formatted_before_reshuffle, equal_to(expected_data), label="formatted_before_reshuffle") assert_that( formatted_after_reshuffle, equal_to(expected_data), label="formatted_after_reshuffle") def reshuffle_unpicklable_in_global_window_helper( self, update_compatibility_version=None): with TestPipeline(options=PipelineOptions( update_compatibility_version=update_compatibility_version)) as pipeline: data = [_Unpicklable(i) for i in range(5)] expected_data = [0, 10, 20, 30, 40] result = ( pipeline | beam.Create(data) | beam.WindowInto(GlobalWindows()) | beam.Reshuffle() | beam.Map(lambda u: u.value * 10)) assert_that(result, equal_to(expected_data)) def test_reshuffle_unpicklable_in_global_window(self): beam.coders.registry.register_coder(_Unpicklable, _UnpicklableCoder) self.reshuffle_unpicklable_in_global_window_helper() # An exception is raised when running reshuffle on unpicklable objects # prior to 2.64.0 self.assertRaises( RuntimeError, self.reshuffle_unpicklable_in_global_window_helper, "2.63.0") def reshuffle_unpicklable_in_non_global_window_helper( self, update_compatibility_version=None): with TestPipeline(options=PipelineOptions( update_compatibility_version=update_compatibility_version)) as pipeline: data = [_Unpicklable(i) for i in range(5)] expected_data = [0, 0, 0, 10, 10, 10, 20, 20, 20, 30, 30, 30, 40, 40, 40] result = ( pipeline | beam.Create(data) | beam.WindowInto(window.SlidingWindows(size=3, period=1)) | beam.Reshuffle() | beam.Map(lambda u: u.value * 10)) assert_that(result, equal_to(expected_data)) def test_reshuffle_unpicklable_in_non_global_window(self): beam.coders.registry.register_coder(_Unpicklable, _UnpicklableCoder) self.reshuffle_unpicklable_in_non_global_window_helper() # An exception is raised when running reshuffle on unpicklable objects # prior to 2.64.0 self.assertRaises( RuntimeError, self.reshuffle_unpicklable_in_non_global_window_helper, "2.63.0") class WithKeysTest(unittest.TestCase): def setUp(self): self.l = [1, 2, 3] def test_constant_k(self): with TestPipeline() as p: pc = p | beam.Create(self.l) with_keys = pc | util.WithKeys('k') assert_that(with_keys, equal_to([('k', 1), ('k', 2), ('k', 3)], )) def test_callable_k(self): with TestPipeline() as p: pc = p | beam.Create(self.l) with_keys = pc | util.WithKeys(lambda x: x * x) assert_that(with_keys, equal_to([(1, 1), (4, 2), (9, 3)])) @staticmethod def _test_args_kwargs_fn(x, multiply, subtract): return x * multiply - subtract def test_args_kwargs_k(self): with TestPipeline() as p: pc = p | beam.Create(self.l) with_keys = pc | util.WithKeys( WithKeysTest._test_args_kwargs_fn, 2, subtract=1) assert_that(with_keys, equal_to([(1, 1), (3, 2), (5, 3)])) def test_sideinputs(self): with TestPipeline() as p: pc = p | beam.Create(self.l) si1 = AsList(p | "side input 1" >> beam.Create([1, 2, 3])) si2 = AsSingleton(p | "side input 2" >> beam.Create([10])) with_keys = pc | util.WithKeys( lambda x, the_list, the_singleton: x + sum(the_list) + the_singleton, si1, the_singleton=si2) assert_that(with_keys, equal_to([(17, 1), (18, 2), (19, 3)])) class GroupIntoBatchesTest(unittest.TestCase): NUM_ELEMENTS = 10 BATCH_SIZE = 5 @staticmethod def _create_test_data(): scientists = [ "Einstein", "Darwin", "Copernicus", "Pasteur", "Curie", "Faraday", "Newton", "Bohr", "Galilei", "Maxwell" ] data = [] for i in range(GroupIntoBatchesTest.NUM_ELEMENTS): index = i % len(scientists) data.append(("key", scientists[index])) return data def test_in_global_window(self): with TestPipeline() as pipeline: collection = pipeline \ | beam.Create(GroupIntoBatchesTest._create_test_data()) \ | util.GroupIntoBatches(GroupIntoBatchesTest.BATCH_SIZE) num_batches = collection | beam.combiners.Count.Globally() assert_that( num_batches, equal_to([ int( math.ceil( GroupIntoBatchesTest.NUM_ELEMENTS / GroupIntoBatchesTest.BATCH_SIZE)) ])) def test_in_global_window_with_synthetic_source(self): with beam.Pipeline() as pipeline: collection = ( pipeline | beam.io.Read( SyntheticSource({ "numRecords": 10, "keySizeBytes": 1, "valueSizeBytes": 1 })) | "identical keys" >> beam.Map(lambda x: (None, x[1])) | "Group key" >> beam.GroupIntoBatches(2) | "count size" >> beam.Map(lambda x: len(x[1]))) assert_that(collection, equal_to([2, 2, 2, 2, 2])) def test_with_sharded_key_in_global_window(self): with TestPipeline() as pipeline: collection = ( pipeline | beam.Create(GroupIntoBatchesTest._create_test_data()) | util.GroupIntoBatches.WithShardedKey( GroupIntoBatchesTest.BATCH_SIZE)) num_batches = collection | beam.combiners.Count.Globally() assert_that( num_batches, equal_to([ int( math.ceil( GroupIntoBatchesTest.NUM_ELEMENTS / GroupIntoBatchesTest.BATCH_SIZE)) ])) def test_buffering_timer_in_fixed_window_streaming(self): window_duration = 6 max_buffering_duration_secs = 100 start_time = timestamp.Timestamp(0) test_stream = ( TestStream().add_elements([ TimestampedValue(value, start_time + i) for i, value in enumerate(GroupIntoBatchesTest._create_test_data()) ]).advance_processing_time(150).advance_watermark_to( start_time + window_duration).advance_watermark_to( start_time + window_duration + 1).advance_watermark_to_infinity()) with TestPipeline(options=StandardOptions(streaming=True)) as pipeline: # To trigger the processing time timer, use a fake clock with start time # being Timestamp(0). fake_clock = FakeClock(now=start_time) num_elements_per_batch = ( pipeline | test_stream | "fixed window" >> WindowInto(FixedWindows(window_duration)) | util.GroupIntoBatches( GroupIntoBatchesTest.BATCH_SIZE, max_buffering_duration_secs, fake_clock) | "count elements in batch" >> Map(lambda x: (None, len(x[1]))) | GroupByKey() | "global window" >> WindowInto(GlobalWindows()) | FlatMapTuple(lambda k, vs: vs)) # Window duration is 6 and batch size is 5, so output batch size # should be 5 (flush because of batch size reached). expected_0 = 5 # There is only one element left in the window so batch size # should be 1 (flush because of max buffering duration reached). expected_1 = 1 # Collection has 10 elements, there are only 4 left, so batch size should # be 4 (flush because of end of window reached). expected_2 = 4 assert_that( num_elements_per_batch, equal_to([expected_0, expected_1, expected_2]), "assert2") def test_buffering_timer_in_global_window_streaming(self): max_buffering_duration_secs = 42 start_time = timestamp.Timestamp(0) test_stream = TestStream().advance_watermark_to(start_time) for i, value in enumerate(GroupIntoBatchesTest._create_test_data()): test_stream.add_elements( [TimestampedValue(value, start_time + i)]) \ .advance_processing_time(5) test_stream.advance_watermark_to( start_time + GroupIntoBatchesTest.NUM_ELEMENTS + 1) \ .advance_watermark_to_infinity() with TestPipeline(options=StandardOptions(streaming=True)) as pipeline: # Set a batch size larger than the total number of elements. # Since we're in a global window, we would have been waiting # for all the elements to arrive without the buffering time limit. batch_size = GroupIntoBatchesTest.NUM_ELEMENTS * 2 # To trigger the processing time timer, use a fake clock with start time # being Timestamp(0). Since the fake clock never really advances during # the pipeline execution, meaning that the timer is always set to the same # value, the timer will be fired on every element after the first firing. fake_clock = FakeClock(now=start_time) num_elements_per_batch = ( pipeline | test_stream | WindowInto( GlobalWindows(), trigger=Repeatedly(AfterCount(1)), accumulation_mode=trigger.AccumulationMode.DISCARDING) | util.GroupIntoBatches( batch_size, max_buffering_duration_secs, fake_clock) | 'count elements in batch' >> Map(lambda x: (None, len(x[1]))) | GroupByKey() | FlatMapTuple(lambda k, vs: vs)) # We will flush twice when the max buffering duration is reached and when # the global window ends. assert_that(num_elements_per_batch, equal_to([9, 1])) def test_output_typehints(self): transform = util.GroupIntoBatches.WithShardedKey( GroupIntoBatchesTest.BATCH_SIZE) unused_input_type = typehints.Tuple[str, str] output_type = transform.infer_output_type(unused_input_type) self.assertTrue(isinstance(output_type, typehints.TupleConstraint)) k, v = output_type.tuple_types self.assertTrue(isinstance(k, ShardedKeyType)) self.assertTrue(isinstance(v, typehints.IterableTypeConstraint)) with TestPipeline() as pipeline: collection = ( pipeline | beam.Create([((1, 2), 'a'), ((2, 3), 'b')]) | util.GroupIntoBatches.WithShardedKey( GroupIntoBatchesTest.BATCH_SIZE)) self.assertTrue( collection.element_type, typehints.Tuple[ ShardedKeyType[typehints.Tuple[int, int]], # type: ignore[misc] typehints.Iterable[str]]) def test_runtime_type_check(self): options = PipelineOptions() options.view_as(TypeOptions).runtime_type_check = True with TestPipeline(options=options) as pipeline: collection = ( pipeline | beam.Create(GroupIntoBatchesTest._create_test_data()) | util.GroupIntoBatches(GroupIntoBatchesTest.BATCH_SIZE)) num_batches = collection | beam.combiners.Count.Globally() assert_that( num_batches, equal_to([ int( math.ceil( GroupIntoBatchesTest.NUM_ELEMENTS / GroupIntoBatchesTest.BATCH_SIZE)) ])) def _test_runner_api_round_trip(self, transform, urn): context = pipeline_context.PipelineContext() proto = transform.to_runner_api(context) self.assertEqual(urn, proto.urn) payload = ( proto_utils.parse_Bytes( proto.payload, beam_runner_api_pb2.GroupIntoBatchesPayload)) self.assertEqual(transform.params.batch_size, payload.batch_size) self.assertEqual( transform.params.max_buffering_duration_secs * 1000, payload.max_buffering_duration_millis) transform_from_proto = ( transform.__class__.from_runner_api_parameter(None, payload, None)) self.assertIsInstance(transform_from_proto, transform.__class__) self.assertEqual(transform.params, transform_from_proto.params) def test_runner_api(self): batch_size = 10 max_buffering_duration_secs = [None, 0, 5] for duration in max_buffering_duration_secs: self._test_runner_api_round_trip( util.GroupIntoBatches(batch_size, duration), common_urns.group_into_batches_components.GROUP_INTO_BATCHES.urn) self._test_runner_api_round_trip( util.GroupIntoBatches(batch_size), common_urns.group_into_batches_components.GROUP_INTO_BATCHES.urn) for duration in max_buffering_duration_secs: self._test_runner_api_round_trip( util.GroupIntoBatches.WithShardedKey(batch_size, duration), common_urns.composites.GROUP_INTO_BATCHES_WITH_SHARDED_KEY.urn) self._test_runner_api_round_trip( util.GroupIntoBatches.WithShardedKey(batch_size), common_urns.composites.GROUP_INTO_BATCHES_WITH_SHARDED_KEY.urn) class ToStringTest(unittest.TestCase): def test_tostring_elements(self): with TestPipeline() as p: result = (p | beam.Create([1, 1, 2, 3]) | util.ToString.Element()) assert_that(result, equal_to(["1", "1", "2", "3"])) def test_tostring_iterables(self): with TestPipeline() as p: result = ( p | beam.Create([("one", "two", "three"), ("four", "five", "six")]) | util.ToString.Iterables()) assert_that(result, equal_to(["one,two,three", "four,five,six"])) def test_tostring_iterables_with_delimeter(self): with TestPipeline() as p: data = [("one", "two", "three"), ("four", "five", "six")] result = (p | beam.Create(data) | util.ToString.Iterables("\t")) assert_that(result, equal_to(["one\ttwo\tthree", "four\tfive\tsix"])) def test_tostring_kvs(self): with TestPipeline() as p: result = (p | beam.Create([("one", 1), ("two", 2)]) | util.ToString.Kvs()) assert_that(result, equal_to(["one,1", "two,2"])) def test_tostring_kvs_delimeter(self): with TestPipeline() as p: result = ( p | beam.Create([("one", 1), ("two", 2)]) | util.ToString.Kvs("\t")) assert_that(result, equal_to(["one\t1", "two\t2"])) def test_tostring_kvs_empty_delimeter(self): with TestPipeline() as p: result = ( p | beam.Create([("one", 1), ("two", 2)]) | util.ToString.Kvs("")) assert_that(result, equal_to(["one1", "two2"])) class LogElementsTest(unittest.TestCase): @pytest.fixture(scope="function") def _capture_stdout_log(request, capsys): with TestPipeline() as p: result = ( p | beam.Create([ TimestampedValue( "event", datetime(2022, 10, 1, 0, 0, 0, 0, tzinfo=pytz.UTC).timestamp()), TimestampedValue( "event", datetime(2022, 10, 2, 0, 0, 0, 0, tzinfo=pytz.UTC).timestamp()), ]) | beam.WindowInto(FixedWindows(60)) | util.LogElements( prefix='prefix_', with_window=True, with_timestamp=True, with_pane_info=True)) request.captured_stdout = capsys.readouterr().out return result @pytest.mark.usefixtures("_capture_stdout_log") def test_stdout_logs(self): assert self.captured_stdout == \ ("prefix_event, timestamp='2022-10-01T00:00:00Z', " "window(start=2022-10-01T00:00:00Z, end=2022-10-01T00:01:00Z), " "pane_info=PaneInfo(first: True, last: True, timing: UNKNOWN, " "index: 0, nonspeculative_index: 0)\n" "prefix_event, timestamp='2022-10-02T00:00:00Z', " "window(start=2022-10-02T00:00:00Z, end=2022-10-02T00:01:00Z), " "pane_info=PaneInfo(first: True, last: True, timing: UNKNOWN, " "index: 0, nonspeculative_index: 0)\n"), \ f'Received from stdout: {self.captured_stdout}' @pytest.fixture(scope="function") def _capture_stdout_log_without_rfc3339(request, capsys): with TestPipeline() as p: result = ( p | beam.Create([ TimestampedValue( "event", datetime(2022, 10, 1, 0, 0, 0, 0, tzinfo=pytz.UTC).timestamp()), TimestampedValue( "event", datetime(2022, 10, 2, 0, 0, 0, 0, tzinfo=pytz.UTC).timestamp()), ]) | beam.WindowInto(FixedWindows(60)) | util.LogElements( prefix='prefix_', with_window=True, with_timestamp=True, use_epoch_time=True)) request.captured_stdout = capsys.readouterr().out return result @pytest.mark.usefixtures("_capture_stdout_log_without_rfc3339") def test_stdout_logs_without_rfc3339(self): assert self.captured_stdout == \ ("prefix_event, timestamp=1664582400, " "window(start=1664582400, end=1664582460)\n" "prefix_event, timestamp=1664668800, " "window(start=1664668800, end=1664668860)\n"), \ f'Received from stdout: {self.captured_stdout}' def test_ptransform_output(self): with TestPipeline() as p: result = ( p | beam.Create(['a', 'b', 'c']) | util.LogElements(prefix='prefix_')) assert_that(result, equal_to(['a', 'b', 'c'])) @pytest.fixture(scope="function") def _capture_logs(request, caplog): with caplog.at_level(logging.INFO): with TestPipeline() as p: _ = ( p | "info" >> beam.Create(["element"]) | "I" >> beam.LogElements(prefix='info_', level=logging.INFO)) _ = ( p | "warning" >> beam.Create(["element"]) | "W" >> beam.LogElements(prefix='warning_', level=logging.WARNING)) _ = ( p | "error" >> beam.Create(["element"]) | "E" >> beam.LogElements(prefix='error_', level=logging.ERROR)) request.captured_log = caplog.text @pytest.mark.usefixtures("_capture_logs") def test_setting_level_uses_appropriate_log_channel(self): self.assertTrue( re.compile('INFO(.*)info_element').search(self.captured_log)) self.assertTrue( re.compile('WARNING(.*)warning_element').search(self.captured_log)) self.assertTrue( re.compile('ERROR(.*)error_element').search(self.captured_log)) class ReifyTest(unittest.TestCase): def test_timestamp(self): l = [ TimestampedValue('a', 100), TimestampedValue('b', 200), TimestampedValue('c', 300) ] expected = [ TestWindowedValue('a', 100, [GlobalWindow()]), TestWindowedValue('b', 200, [GlobalWindow()]), TestWindowedValue('c', 300, [GlobalWindow()]) ] with TestPipeline() as p: # Map(lambda x: x) PTransform is added after Create here, because when # a PCollection of TimestampedValues is created with Create PTransform, # the timestamps are not assigned to it. Adding a Map forces the # PCollection to go through a DoFn so that the PCollection consists of # the elements with timestamps assigned to them instead of a PCollection # of TimestampedValue(element, timestamp). pc = p | beam.Create(l) | beam.Map(lambda x: x) reified_pc = pc | util.Reify.Timestamp() assert_that(reified_pc, equal_to(expected), reify_windows=True) def test_window(self): l = [ GlobalWindows.windowed_value('a', 100), GlobalWindows.windowed_value('b', 200), GlobalWindows.windowed_value('c', 300) ] expected = [ TestWindowedValue(('a', 100, GlobalWindow()), 100, [GlobalWindow()]), TestWindowedValue(('b', 200, GlobalWindow()), 200, [GlobalWindow()]), TestWindowedValue(('c', 300, GlobalWindow()), 300, [GlobalWindow()]) ] with TestPipeline() as p: pc = p | beam.Create(l) # Map(lambda x: x) PTransform is added after Create here, because when # a PCollection of WindowedValues is created with Create PTransform, # the windows are not assigned to it. Adding a Map forces the # PCollection to go through a DoFn so that the PCollection consists of # the elements with timestamps assigned to them instead of a PCollection # of WindowedValue(element, timestamp, window). pc = pc | beam.Map(lambda x: x) reified_pc = pc | util.Reify.Window() assert_that(reified_pc, equal_to(expected), reify_windows=True) def test_timestamp_in_value(self): l = [ TimestampedValue(('a', 1), 100), TimestampedValue(('b', 2), 200), TimestampedValue(('c', 3), 300) ] expected = [ TestWindowedValue(('a', TimestampedValue(1, 100)), 100, [GlobalWindow()]), TestWindowedValue(('b', TimestampedValue(2, 200)), 200, [GlobalWindow()]), TestWindowedValue(('c', TimestampedValue(3, 300)), 300, [GlobalWindow()]) ] with TestPipeline() as p: pc = p | beam.Create(l) | beam.Map(lambda x: x) reified_pc = pc | util.Reify.TimestampInValue() assert_that(reified_pc, equal_to(expected), reify_windows=True) def test_window_in_value(self): l = [ GlobalWindows.windowed_value(('a', 1), 100), GlobalWindows.windowed_value(('b', 2), 200), GlobalWindows.windowed_value(('c', 3), 300) ] expected = [ TestWindowedValue(('a', (1, 100, GlobalWindow())), 100, [GlobalWindow()]), TestWindowedValue(('b', (2, 200, GlobalWindow())), 200, [GlobalWindow()]), TestWindowedValue(('c', (3, 300, GlobalWindow())), 300, [GlobalWindow()]) ] with TestPipeline() as p: # Map(lambda x: x) hack is used for the same reason here. # Also, this makes the typehint on Reify.WindowInValue work. pc = p | beam.Create(l) | beam.Map(lambda x: x) reified_pc = pc | util.Reify.WindowInValue() assert_that(reified_pc, equal_to(expected), reify_windows=True) class RegexTest(unittest.TestCase): def test_find(self): with TestPipeline() as p: result = ( p | beam.Create(["aj", "xj", "yj", "zj"]) | util.Regex.find("[xyz]")) assert_that(result, equal_to(["x", "y", "z"])) def test_find_pattern(self): with TestPipeline() as p: rc = re.compile("[xyz]") result = (p | beam.Create(["aj", "xj", "yj", "zj"]) | util.Regex.find(rc)) assert_that(result, equal_to(["x", "y", "z"])) def test_find_group(self): with TestPipeline() as p: result = ( p | beam.Create(["aj", "xj", "yj", "zj"]) | util.Regex.find("([xyz])j", group=1)) assert_that(result, equal_to(["x", "y", "z"])) def test_find_empty(self): with TestPipeline() as p: result = ( p | beam.Create(["a", "b", "c", "d"]) | util.Regex.find("[xyz]")) assert_that(result, equal_to([])) def test_find_group_name(self): with TestPipeline() as p: result = ( p | beam.Create(["aj", "xj", "yj", "zj"]) | util.Regex.find("(?P[xyz])j", group="namedgroup")) assert_that(result, equal_to(["x", "y", "z"])) def test_find_group_name_pattern(self): with TestPipeline() as p: rc = re.compile("(?P[xyz])j") result = ( p | beam.Create(["aj", "xj", "yj", "zj"]) | util.Regex.find(rc, group="namedgroup")) assert_that(result, equal_to(["x", "y", "z"])) def test_find_all_groups(self): data = ["abb ax abbb", "abc qwerty abcabcd xyz"] with TestPipeline() as p: pcol = (p | beam.Create(data)) assert_that( pcol | 'with default values' >> util.Regex.find_all('a(b*)'), equal_to([['abb', 'a', 'abbb'], ['ab', 'ab', 'ab']]), label='CheckWithDefaultValues') assert_that( pcol | 'group 1' >> util.Regex.find_all('a(b*)', 1), equal_to([['b', 'b', 'b'], ['bb', '', 'bbb']]), label='CheckWithGroup1') assert_that( pcol | 'group 1 non empty' >> util.Regex.find_all( 'a(b*)', 1, outputEmpty=False), equal_to([['b', 'b', 'b'], ['bb', 'bbb']]), label='CheckGroup1NonEmpty') assert_that( pcol | 'named group' >> util.Regex.find_all( 'a(?Pb*)', 'namedgroup'), equal_to([['b', 'b', 'b'], ['bb', '', 'bbb']]), label='CheckNamedGroup') assert_that( pcol | 'all groups' >> util.Regex.find_all( 'a(?Pb*)', util.Regex.ALL), equal_to([[('ab', 'b'), ('ab', 'b'), ('ab', 'b')], [('abb', 'bb'), ('a', ''), ('abbb', 'bbb')]]), label='CheckAllGroups') assert_that( pcol | 'all non empty groups' >> util.Regex.find_all( 'a(b*)', util.Regex.ALL, outputEmpty=False), equal_to([[('ab', 'b'), ('ab', 'b'), ('ab', 'b')], [('abb', 'bb'), ('abbb', 'bbb')]]), label='CheckAllNonEmptyGroups') def test_find_kv(self): with TestPipeline() as p: pcol = (p | beam.Create(['a b c d'])) assert_that( pcol | 'key 1' >> util.Regex.find_kv( 'a (b) (c)', 1, ), equal_to([('b', 'a b c')]), label='CheckKey1') assert_that( pcol | 'key 1 group 1' >> util.Regex.find_kv('a (b) (c)', 1, 2), equal_to([('b', 'c')]), label='CheckKey1Group1') def test_find_kv_pattern(self): with TestPipeline() as p: rc = re.compile("a (b) (c)") result = (p | beam.Create(["a b c"]) | util.Regex.find_kv(rc, 1, 2)) assert_that(result, equal_to([("b", "c")])) def test_find_kv_none(self): with TestPipeline() as p: result = ( p | beam.Create(["x y z"]) | util.Regex.find_kv("a (b) (c)", 1, 2)) assert_that(result, equal_to([])) def test_match(self): with TestPipeline() as p: result = ( p | beam.Create(["a", "x", "y", "z"]) | util.Regex.matches("[xyz]")) assert_that(result, equal_to(["x", "y", "z"])) with TestPipeline() as p: result = ( p | beam.Create(["a", "ax", "yby", "zzc"]) | util.Regex.matches("[xyz]")) assert_that(result, equal_to(["y", "z"])) def test_match_entire_line(self): with TestPipeline() as p: result = ( p | beam.Create(["a", "x", "y", "ay", "zz"]) | util.Regex.matches("[xyz]$")) assert_that(result, equal_to(["x", "y"])) def test_match_pattern(self): with TestPipeline() as p: rc = re.compile("[xyz]") result = (p | beam.Create(["a", "x", "y", "z"]) | util.Regex.matches(rc)) assert_that(result, equal_to(["x", "y", "z"])) def test_match_none(self): with TestPipeline() as p: result = ( p | beam.Create(["a", "b", "c", "d"]) | util.Regex.matches("[xyz]")) assert_that(result, equal_to([])) def test_match_group(self): with TestPipeline() as p: result = ( p | beam.Create(["a", "x xxx", "x yyy", "x zzz"]) | util.Regex.matches("x ([xyz]*)", 1)) assert_that(result, equal_to(("xxx", "yyy", "zzz"))) def test_match_group_name(self): with TestPipeline() as p: result = ( p | beam.Create(["a", "x xxx", "x yyy", "x zzz"]) | util.Regex.matches("x (?P[xyz]*)", 'namedgroup')) assert_that(result, equal_to(("xxx", "yyy", "zzz"))) def test_match_group_name_pattern(self): with TestPipeline() as p: rc = re.compile("x (?P[xyz]*)") result = ( p | beam.Create(["a", "x xxx", "x yyy", "x zzz"]) | util.Regex.matches(rc, 'namedgroup')) assert_that(result, equal_to(("xxx", "yyy", "zzz"))) def test_match_group_empty(self): with TestPipeline() as p: result = ( p | beam.Create(["a", "b", "c", "d"]) | util.Regex.matches("x (?P[xyz]*)", 'namedgroup')) assert_that(result, equal_to([])) def test_all_matched(self): with TestPipeline() as p: result = ( p | beam.Create(["a x", "x x", "y y", "z z"]) | util.Regex.all_matches("([xyz]) ([xyz])")) expected_result = [["x x", "x", "x"], ["y y", "y", "y"], ["z z", "z", "z"]] assert_that(result, equal_to(expected_result)) def test_all_matched_pattern(self): with TestPipeline() as p: rc = re.compile("([xyz]) ([xyz])") result = ( p | beam.Create(["a x", "x x", "y y", "z z"]) | util.Regex.all_matches(rc)) expected_result = [["x x", "x", "x"], ["y y", "y", "y"], ["z z", "z", "z"]] assert_that(result, equal_to(expected_result)) def test_match_group_kv(self): with TestPipeline() as p: result = ( p | beam.Create(["a b c"]) | util.Regex.matches_kv("a (b) (c)", 1, 2)) assert_that(result, equal_to([("b", "c")])) def test_match_group_kv_pattern(self): with TestPipeline() as p: rc = re.compile("a (b) (c)") pcol = (p | beam.Create(["a b c"])) assert_that( pcol | 'key 1' >> util.Regex.matches_kv(rc, 1), equal_to([("b", "a b c")]), label="CheckKey1") assert_that( pcol | 'key 1 group 2' >> util.Regex.matches_kv(rc, 1, 2), equal_to([("b", "c")]), label="CheckKey1Group2") def test_match_group_kv_none(self): with TestPipeline() as p: result = ( p | beam.Create(["x y z"]) | util.Regex.matches_kv("a (b) (c)", 1, 2)) assert_that(result, equal_to([])) def test_match_kv_group_names(self): with TestPipeline() as p: result = ( p | beam.Create(["a b c"]) | util.Regex.matches_kv( "a (?Pb) (?Pc)", 'keyname', 'valuename')) assert_that(result, equal_to([("b", "c")])) def test_match_kv_group_names_pattern(self): with TestPipeline() as p: rc = re.compile("a (?Pb) (?Pc)") result = ( p | beam.Create(["a b c"]) | util.Regex.matches_kv(rc, 'keyname', 'valuename')) assert_that(result, equal_to([("b", "c")])) def test_match_kv_group_name_none(self): with TestPipeline() as p: result = ( p | beam.Create(["x y z"]) | util.Regex.matches_kv( "a (?Pb) (?Pc)", 'keyname', 'valuename')) assert_that(result, equal_to([])) def test_replace_all(self): with TestPipeline() as p: result = ( p | beam.Create(["xj", "yj", "zj"]) | util.Regex.replace_all("[xyz]", "new")) assert_that(result, equal_to(["newj", "newj", "newj"])) def test_replace_all_mixed(self): with TestPipeline() as p: result = ( p | beam.Create(["abc", "xj", "yj", "zj", "def"]) | util.Regex.replace_all("[xyz]", 'new')) assert_that(result, equal_to(["abc", "newj", "newj", "newj", "def"])) def test_replace_all_mixed_pattern(self): with TestPipeline() as p: rc = re.compile("[xyz]") result = ( p | beam.Create(["abc", "xj", "yj", "zj", "def"]) | util.Regex.replace_all(rc, 'new')) assert_that(result, equal_to(["abc", "newj", "newj", "newj", "def"])) def test_replace_first(self): with TestPipeline() as p: result = ( p | beam.Create(["xjx", "yjy", "zjz"]) | util.Regex.replace_first("[xyz]", 'new')) assert_that(result, equal_to(["newjx", "newjy", "newjz"])) def test_replace_first_mixed(self): with TestPipeline() as p: result = ( p | beam.Create(["abc", "xjx", "yjy", "zjz", "def"]) | util.Regex.replace_first("[xyz]", 'new')) assert_that(result, equal_to(["abc", "newjx", "newjy", "newjz", "def"])) def test_replace_first_mixed_pattern(self): with TestPipeline() as p: rc = re.compile("[xyz]") result = ( p | beam.Create(["abc", "xjx", "yjy", "zjz", "def"]) | util.Regex.replace_first(rc, 'new')) assert_that(result, equal_to(["abc", "newjx", "newjy", "newjz", "def"])) def test_split(self): with TestPipeline() as p: data = ["The quick brown fox jumps over the lazy dog"] result = (p | beam.Create(data) | util.Regex.split("\\W+")) expected_result = [[ "The", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog" ]] assert_that(result, equal_to(expected_result)) def test_split_pattern(self): with TestPipeline() as p: data = ["The quick brown fox jumps over the lazy dog"] rc = re.compile("\\W+") result = (p | beam.Create(data) | util.Regex.split(rc)) expected_result = [[ "The", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog" ]] assert_that(result, equal_to(expected_result)) def test_split_with_empty(self): with TestPipeline() as p: data = ["The quick brown fox jumps over the lazy dog"] result = (p | beam.Create(data) | util.Regex.split("\\s", True)) expected_result = [[ 'The', '', 'quick', '', '', 'brown', 'fox', 'jumps', 'over', '', '', '', 'the', 'lazy', 'dog' ]] assert_that(result, equal_to(expected_result)) def test_split_without_empty(self): with TestPipeline() as p: data = ["The quick brown fox jumps over the lazy dog"] result = (p | beam.Create(data) | util.Regex.split("\\s", False)) expected_result = [[ "The", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog" ]] assert_that(result, equal_to(expected_result)) class TeeTest(unittest.TestCase): _side_effects: Mapping[str, int] = collections.defaultdict(int) def test_tee(self): # The imports here are to avoid issues with the class (and its attributes) # possibly being pickled rather than referenced. def cause_side_effect(element): importlib.import_module(__name__).TeeTest._side_effects[element] += 1 def count_side_effects(element): return importlib.import_module(__name__).TeeTest._side_effects[element] with TestPipeline() as p: result = ( p | beam.Create(['a', 'b', 'c']) | 'TeePTransform' >> beam.Tee(beam.Map(cause_side_effect)) | 'TeeCallable' >> beam.Tee( lambda pcoll: pcoll | beam.Map( lambda element: cause_side_effect('X' + element)))) assert_that(result, equal_to(['a', 'b', 'c'])) self.assertEqual(count_side_effects('a'), 1) self.assertEqual(count_side_effects('Xa'), 1) class WaitOnTest(unittest.TestCase): def test_find(self): # We need shared reference that survives pickling. def increment_global_counter(): try: value = getattr(beam, '_WAIT_ON_TEST_COUNTER', 0) return value finally: setattr(beam, '_WAIT_ON_TEST_COUNTER', value + 1) def record(tag): return f'Record({tag})' >> beam.Map( lambda x: (x[0], tag, increment_global_counter())) with TestPipeline() as p: start = p | beam.Create([(None, ), (None, )]) x = start | record('x') y = start | 'WaitForX' >> util.WaitOn(x) | record('y') z = start | 'WaitForY' >> util.WaitOn(y) | record('z') result = x | 'WaitForYZ' >> util.WaitOn(y, z) | record('result') assert_that(x, equal_to([(None, 'x', 0), (None, 'x', 1)]), label='x') assert_that(y, equal_to([(None, 'y', 2), (None, 'y', 3)]), label='y') assert_that(z, equal_to([(None, 'z', 4), (None, 'z', 5)]), label='z') assert_that( result, equal_to([(None, 'result', 6), (None, 'result', 7)]), label='result') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/validate_runner_xlang_test.py0000644000175100017510000003713615206113033026052 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ ########################################################### Runner Validation Test Suite for Cross-language Transforms ########################################################### As per Beams's Portability Framework design, Cross-language transforms should work out of the box. In spite of this, there always exists a possibility of rough edges existing. It could be caused due to unpolished implementation of any part of the execution code path, for example: - Transform expansion [SDK] - Pipeline construction [SDK] - Cross-language artifact staging [Runner] - Language specific serialization/deserialization of PCollection (and other data types) [Runner/SDK] In an effort to improve developer visibility into potential problems, this test suite validates correct execution of 5 Core Beam transforms when used as cross-language transforms within the Python SDK from any foreign SDK: - ParDo (https://beam.apache.org/documentation/programming-guide/#pardo) - GroupByKey (https://beam.apache.org/documentation/programming-guide/#groupbykey) - CoGroupByKey (https://beam.apache.org/documentation/programming-guide/#cogroupbykey) - Combine (https://beam.apache.org/documentation/programming-guide/#combine) - Flatten (https://beam.apache.org/documentation/programming-guide/#flatten) - Partition (https://beam.apache.org/documentation/programming-guide/#partition) See Runner Validation Test Plan for Cross-language transforms at https://docs.google.com/document/d/1xQp0ElIV84b8OCVz8CD2hvbiWdR8w4BvWxPTZJZA6NA for further details. """ import logging import os import sys import unittest from datetime import datetime import pytest import apache_beam as beam from apache_beam.options.pipeline_options import SetupOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.external import ImplicitSchemaPayloadBuilder from apache_beam.transforms.util import GcpSecret from apache_beam.transforms.util import Secret try: from google.cloud import secretmanager except ImportError: secretmanager = None # type: ignore[assignment] TEST_PREFIX_URN = "beam:transforms:xlang:test:prefix" TEST_MULTI_URN = "beam:transforms:xlang:test:multi" TEST_GBK_URN = "beam:transforms:xlang:test:gbk" TEST_CGBK_URN = "beam:transforms:xlang:test:cgbk" TEST_COMGL_URN = "beam:transforms:xlang:test:comgl" TEST_COMPK_URN = "beam:transforms:xlang:test:compk" TEST_FLATTEN_URN = "beam:transforms:xlang:test:flatten" TEST_PARTITION_URN = "beam:transforms:xlang:test:partition" class CrossLanguageTestPipelines(object): def __init__(self, expansion_service=None): self.expansion_service = expansion_service or ( 'localhost:%s' % os.environ.get('EXPANSION_PORT')) def run_prefix(self, pipeline): """ Target transform - ParDo (https://beam.apache.org/documentation/programming-guide/#pardo) Test scenario - Mapping elements from a single input collection to a single output collection Boundary conditions checked - - PCollection to external transforms - PCollection from external transforms """ with pipeline as p: res = ( p | beam.Create(['a', 'b']).with_output_types(str) | beam.ExternalTransform( TEST_PREFIX_URN, ImplicitSchemaPayloadBuilder({'data': '0'}), self.expansion_service)) assert_that(res, equal_to(['0a', '0b'])) def run_multi_input_output_with_sideinput(self, pipeline): """ Target transform - ParDo (https://beam.apache.org/documentation/programming-guide/#pardo) Test scenario - Mapping elements from multiple input collections (main and side) to multiple output collections (main and side) Boundary conditions checked - - PCollectionTuple to external transforms - PCollectionTuple from external transforms """ with pipeline as p: main1 = p | 'Main1' >> beam.Create( ['a', 'bb'], reshuffle=False).with_output_types(str) main2 = p | 'Main2' >> beam.Create( ['x', 'yy', 'zzz'], reshuffle=False).with_output_types(str) side = p | 'Side' >> beam.Create(['s']).with_output_types(str) res = dict( main1=main1, main2=main2, side=side) | beam.ExternalTransform( TEST_MULTI_URN, None, self.expansion_service) assert_that(res['main'], equal_to(['as', 'bbs', 'xs', 'yys', 'zzzs'])) assert_that(res['side'], equal_to(['ss']), label='CheckSide') def run_group_by_key(self, pipeline): """ Target transform - GroupByKey (https://beam.apache.org/documentation/programming-guide/#groupbykey) Test scenario - Grouping a collection of KV to a collection of KV> by key Boundary conditions checked - - PCollection> to external transforms - PCollection>> from external transforms """ with pipeline as p: res = ( p | beam.Create([(0, "1"), (0, "2"), (1, "3")], reshuffle=False).with_output_types( tuple[int, str]) | beam.ExternalTransform(TEST_GBK_URN, None, self.expansion_service) | beam.Map(lambda x: "{}:{}".format(x[0], ','.join(sorted(x[1]))))) assert_that(res, equal_to(['0:1,2', '1:3'])) def run_group_by_key_no_assert(self, pipeline): """ Target transform - GroupByKey, with no assertion for checking errors (https://beam.apache.org/documentation/programming-guide/#groupbykey) Test scenario - Grouping a collection of KV to a collection of KV> by key Boundary conditions checked - - PCollection> to external transforms - PCollection>> from external transforms """ with pipeline as p: _ = ( p | beam.Create([(0, "1"), (0, "2"), (1, "3")], reshuffle=False).with_output_types( tuple[int, str]) | beam.ExternalTransform(TEST_GBK_URN, None, self.expansion_service)) def run_cogroup_by_key(self, pipeline): """ Target transform - CoGroupByKey (https://beam.apache.org/documentation/programming-guide/#cogroupbykey) Test scenario - Grouping multiple input collections with keys to a collection of KV by key Boundary conditions checked - - KeyedPCollectionTuple to external transforms - PCollection>> from external transforms """ with pipeline as p: col1 = p | 'create_col1' >> beam.Create( [(0, "1"), (0, "2"), (1, "3")], reshuffle=False).with_output_types(tuple[int, str]) col2 = p | 'create_col2' >> beam.Create( [(0, "4"), (1, "5"), (1, "6")], reshuffle=False).with_output_types(tuple[int, str]) res = ( dict(col1=col1, col2=col2) | beam.ExternalTransform(TEST_CGBK_URN, None, self.expansion_service) | beam.Map(lambda x: "{}:{}".format(x[0], ','.join(sorted(x[1]))))) assert_that(res, equal_to(['0:1,2,4', '1:3,5,6'])) def run_combine_globally(self, pipeline): """ Target transform - Combine (https://beam.apache.org/documentation/programming-guide/#combine) Test scenario - Combining elements globally with a predefined simple CombineFn Boundary conditions checked - - PCollection to external transforms - PCollection from external transforms """ with pipeline as p: res = ( p | beam.Create([1, 2, 3]).with_output_types(int) | beam.ExternalTransform( TEST_COMGL_URN, None, self.expansion_service)) assert_that(res, equal_to([6])) def run_combine_per_key(self, pipeline): """ Target transform - Combine (https://beam.apache.org/documentation/programming-guide/#combine) Test scenario - Combining elements per key with a predefined simple merging function Boundary conditions checked - - PCollection to external transforms - PCollection from external transforms """ with pipeline as p: res = ( p | beam.Create([('a', 1), ('a', 2), ('b', 3)]).with_output_types(tuple[str, int]) | beam.ExternalTransform( TEST_COMPK_URN, None, self.expansion_service)) assert_that(res, equal_to([('a', 3), ('b', 3)])) def run_flatten(self, pipeline): """ Target transform - Flatten (https://beam.apache.org/documentation/programming-guide/#flatten) Test scenario - Merging multiple collections into a single collection Boundary conditions checked - - PCollectionList to external transforms - PCollection from external transforms """ with pipeline as p: col1 = p | 'col1' >> beam.Create([1, 2, 3]).with_output_types(int) col2 = p | 'col2' >> beam.Create([4, 5, 6]).with_output_types(int) res = ((col1, col2) | beam.ExternalTransform( TEST_FLATTEN_URN, None, self.expansion_service)) assert_that(res, equal_to([1, 2, 3, 4, 5, 6])) def run_partition(self, pipeline): """ Target transform - Partition (https://beam.apache.org/documentation/programming-guide/#partition) Test scenario - Splitting a single collection into multiple collections with a predefined simple PartitionFn Boundary conditions checked - - PCollection to external transforms - PCollectionList from external transforms """ with pipeline as p: res = ( p | beam.Create([1, 2, 3, 4, 5, 6]).with_output_types(int) | beam.ExternalTransform( TEST_PARTITION_URN, None, self.expansion_service)) assert_that(res['0'], equal_to([2, 4, 6]), label='check_even') assert_that(res['1'], equal_to([1, 3, 5]), label='check_odd') @unittest.skipUnless( os.environ.get('EXPANSION_PORT'), "EXPANSION_PORT environment var is not provided.") class ValidateRunnerXlangTest(unittest.TestCase): _multiprocess_can_split_ = True def create_pipeline(self): test_pipeline = TestPipeline() test_pipeline.not_use_test_runner_api = True return test_pipeline @pytest.mark.uses_java_expansion_service @pytest.mark.uses_python_expansion_service def test_prefix(self, test_pipeline=None): CrossLanguageTestPipelines().run_prefix( test_pipeline or self.create_pipeline()) @pytest.mark.uses_java_expansion_service @pytest.mark.uses_python_expansion_service def test_multi_input_output_with_sideinput(self, test_pipeline=None): CrossLanguageTestPipelines().run_multi_input_output_with_sideinput( test_pipeline or self.create_pipeline()) @pytest.mark.uses_java_expansion_service @pytest.mark.uses_python_expansion_service def test_group_by_key(self, test_pipeline=None): CrossLanguageTestPipelines().run_group_by_key( test_pipeline or self.create_pipeline()) @pytest.mark.uses_java_expansion_service @pytest.mark.uses_python_expansion_service def test_cogroup_by_key(self, test_pipeline=None): CrossLanguageTestPipelines().run_cogroup_by_key( test_pipeline or self.create_pipeline()) @pytest.mark.uses_java_expansion_service @pytest.mark.uses_python_expansion_service def test_combine_globally(self, test_pipeline=None): CrossLanguageTestPipelines().run_combine_globally( test_pipeline or self.create_pipeline()) @pytest.mark.uses_java_expansion_service @pytest.mark.uses_python_expansion_service def test_combine_per_key(self, test_pipeline=None): CrossLanguageTestPipelines().run_combine_per_key( test_pipeline or self.create_pipeline()) # TODO: enable after fixing BEAM-10507 # @pytest.mark.uses_java_expansion_service # @pytest.mark.uses_python_expansion_service def test_flatten(self, test_pipeline=None): CrossLanguageTestPipelines().run_flatten( test_pipeline or self.create_pipeline()) @pytest.mark.uses_java_expansion_service @pytest.mark.uses_python_expansion_service def test_partition(self, test_pipeline=None): CrossLanguageTestPipelines().run_partition( test_pipeline or self.create_pipeline()) @unittest.skipUnless( os.environ.get('EXPANSION_PORT'), "EXPANSION_PORT environment var is not provided.") @unittest.skipIf(secretmanager is None, 'secretmanager not installed') class ValidateRunnerGBEKTest(unittest.TestCase): @classmethod def setUpClass(cls): if secretmanager is not None: cls.project_id = 'apache-beam-testing' py_version = f'_py{sys.version_info.major}{sys.version_info.minor}' secret_postfix = datetime.now().strftime('%m%d_%H%M%S') + py_version cls.secret_id = 'gbek_validaterunner_secret_tests_' + secret_postfix cls.client = secretmanager.SecretManagerServiceClient() cls.project_path = f'projects/{cls.project_id}' cls.secret_path = f'{cls.project_path}/secrets/{cls.secret_id}' try: cls.client.get_secret(request={'name': cls.secret_path}) except Exception: cls.client.create_secret( request={ 'parent': cls.project_path, 'secret_id': cls.secret_id, 'secret': { 'replication': { 'automatic': {} } } }) cls.client.add_secret_version( request={ 'parent': cls.secret_path, 'payload': { 'data': Secret.generate_secret_bytes() } }) version_name = f'{cls.secret_path}/versions/latest' cls.gcp_secret = GcpSecret(version_name) cls.secret_option = f'type:GcpSecret;version_name:{version_name}' @classmethod def tearDownClass(cls): if secretmanager is not None: cls.client.delete_secret(request={'name': cls.secret_path}) def create_pipeline(self): test_pipeline = TestPipeline() test_pipeline.not_use_test_runner_api = True return test_pipeline # This test and test_group_by_key_gbek_bad_secret validate that the gbek # pipeline option is correctly passed through @pytest.mark.uses_java_expansion_service @pytest.mark.uses_python_expansion_service def test_group_by_key_gbek(self, test_pipeline=None): test_pipeline = test_pipeline or self.create_pipeline() good_secret = self.secret_option test_pipeline.options.view_as(SetupOptions).gbek = good_secret CrossLanguageTestPipelines().run_group_by_key(test_pipeline) # Verify actually using secret manager test_pipeline = self.create_pipeline() nonexistent_secret = 'version_name:nonexistent_secret' test_pipeline.options.view_as(SetupOptions).gbek = nonexistent_secret with self.assertRaisesRegex( Exception, 'Secret string must contain a valid type parameter'): CrossLanguageTestPipelines().run_group_by_key_no_assert(test_pipeline) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/window.py0000644000175100017510000004735115206113033021747 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Windowing concepts. A WindowInto transform logically divides up or groups the elements of a PCollection into finite windows according to a windowing function (derived from WindowFn). The output of WindowInto contains the same elements as input, but they have been logically assigned to windows. The next GroupByKey(s) transforms, including one within a composite transform, will group by the combination of keys and windows. Windowing a PCollection allows chunks of it to be processed individually, before the entire PCollection is available. This is especially important for PCollection(s) with unbounded size, since the full PCollection is never available at once, since more data is continually arriving. For PCollection(s) with a bounded size (aka. conventional batch mode), by default, all data is implicitly in a single window (see GlobalWindows), unless WindowInto is applied. For example, a simple form of windowing divides up the data into fixed-width time intervals, using FixedWindows. Seconds are used as the time unit for the built-in windowing primitives here. Integer or floating point seconds can be passed to these primitives. Internally, seconds, with microsecond granularity, are stored as timeutil.Timestamp and timeutil.Duration objects. This is done to avoid precision errors that would occur with floating point representations. Custom windowing function classes can be created, by subclassing from WindowFn. """ # pytype: skip-file import abc from collections.abc import Iterable from functools import total_ordering from typing import Any from typing import Generic from typing import Optional from typing import TypeVar from google.protobuf import duration_pb2 from google.protobuf import timestamp_pb2 from apache_beam.coders import coders from apache_beam.portability import common_urns from apache_beam.portability import python_urns from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.portability.api import standard_window_fns_pb2 from apache_beam.transforms import timeutil from apache_beam.utils import proto_utils from apache_beam.utils import urns from apache_beam.utils import windowed_value from apache_beam.utils.timestamp import MIN_TIMESTAMP from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import DurationTypes # pylint: disable=unused-import from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.timestamp import TimestampTypes # pylint: disable=unused-import from apache_beam.utils.windowed_value import WindowedValue __all__ = [ 'TimestampCombiner', 'WindowFn', 'BoundedWindow', 'IntervalWindow', 'TimestampedValue', 'GlobalWindow', 'NonMergingWindowFn', 'GlobalWindows', 'FixedWindows', 'SlidingWindows', 'Sessions', ] # TODO(ccy): revisit naming and semantics once Java Apache Beam finalizes their # behavior. class TimestampCombiner(object): """Determines how output timestamps of grouping operations are assigned.""" OUTPUT_AT_EOW = beam_runner_api_pb2.OutputTime.END_OF_WINDOW OUTPUT_AT_EARLIEST = beam_runner_api_pb2.OutputTime.EARLIEST_IN_PANE OUTPUT_AT_LATEST = beam_runner_api_pb2.OutputTime.LATEST_IN_PANE # TODO(robertwb): Add this to the runner API or remove it. OUTPUT_AT_EARLIEST_TRANSFORMED = 'OUTPUT_AT_EARLIEST_TRANSFORMED' @staticmethod def get_impl( timestamp_combiner: beam_runner_api_pb2.OutputTime.Enum, window_fn: 'WindowFn') -> timeutil.TimestampCombinerImpl: if timestamp_combiner == TimestampCombiner.OUTPUT_AT_EOW: return timeutil.OutputAtEndOfWindowImpl() elif timestamp_combiner == TimestampCombiner.OUTPUT_AT_EARLIEST: return timeutil.OutputAtEarliestInputTimestampImpl() elif timestamp_combiner == TimestampCombiner.OUTPUT_AT_LATEST: return timeutil.OutputAtLatestInputTimestampImpl() elif timestamp_combiner == TimestampCombiner.OUTPUT_AT_EARLIEST_TRANSFORMED: return timeutil.OutputAtEarliestTransformedInputTimestampImpl(window_fn) else: raise ValueError('Invalid TimestampCombiner: %s.' % timestamp_combiner) class WindowFn(urns.RunnerApiFn, metaclass=abc.ABCMeta): """An abstract windowing function defining a basic assign and merge.""" class AssignContext(object): """Context passed to WindowFn.assign().""" def __init__( self, timestamp: TimestampTypes, element: Optional[Any] = None, window: Optional['BoundedWindow'] = None) -> None: self.timestamp = Timestamp.of(timestamp) self.element = element self.window = window @abc.abstractmethod def assign(self, assign_context: 'AssignContext') -> Iterable['BoundedWindow']: # noqa: F821 """Associates windows to an element. Arguments: assign_context: Instance of AssignContext. Returns: An iterable of BoundedWindow. """ raise NotImplementedError class MergeContext(object): """Context passed to WindowFn.merge() to perform merging, if any.""" def __init__(self, windows: Iterable['BoundedWindow']) -> None: self.windows = list(windows) def merge( self, to_be_merged: Iterable['BoundedWindow'], merge_result: 'BoundedWindow') -> None: raise NotImplementedError @abc.abstractmethod def merge(self, merge_context: 'WindowFn.MergeContext') -> None: """Returns a window that is the result of merging a set of windows.""" raise NotImplementedError def is_merging(self) -> bool: """Returns whether this WindowFn merges windows.""" return True @abc.abstractmethod def get_window_coder(self) -> coders.Coder: raise NotImplementedError def get_transformed_output_time( self, window: 'BoundedWindow', input_timestamp: Timestamp) -> Timestamp: # pylint: disable=unused-argument """Given input time and output window, returns output time for window. If TimestampCombiner.OUTPUT_AT_EARLIEST_TRANSFORMED is used in the Windowing, the output timestamp for the given window will be the earliest of the timestamps returned by get_transformed_output_time() for elements of the window. Arguments: window: Output window of element. input_timestamp: Input timestamp of element as a timeutil.Timestamp object. Returns: Transformed timestamp. """ # By default, just return the input timestamp. return input_timestamp urns.RunnerApiFn.register_pickle_urn(python_urns.PICKLED_WINDOWFN) class BoundedWindow(object): """A window for timestamps in range (-infinity, end). Attributes: end: End of window. """ def __init__(self, end: TimestampTypes) -> None: self._end = Timestamp.of(end) @property def start(self) -> Timestamp: raise NotImplementedError @property def end(self) -> Timestamp: return self._end def max_timestamp(self) -> Timestamp: return self.end.predecessor() def __eq__(self, other): raise NotImplementedError def __ne__(self, other): # Order first by endpoint, then arbitrarily return self.end != other.end or hash(self) != hash(other) def __lt__(self, other): if self.end != other.end: return self.end < other.end return hash(self) < hash(other) def __le__(self, other): if self.end != other.end: return self.end <= other.end return hash(self) <= hash(other) def __gt__(self, other): if self.end != other.end: return self.end > other.end return hash(self) > hash(other) def __ge__(self, other): if self.end != other.end: return self.end >= other.end return hash(self) >= hash(other) def __hash__(self): raise NotImplementedError def __repr__(self): return '[?, %s)' % float(self.end) @total_ordering class IntervalWindow(windowed_value._IntervalWindowBase, BoundedWindow): """A window for timestamps in range [start, end). Attributes: start: Start of window as seconds since Unix epoch. end: End of window as seconds since Unix epoch. """ def __lt__(self, other): if self.end != other.end: return self.end < other.end return hash(self) < hash(other) def intersects(self, other: 'IntervalWindow') -> bool: return other.start < self.end or self.start < other.end def union(self, other: 'IntervalWindow') -> 'IntervalWindow': return IntervalWindow( min(self.start, other.start), max(self.end, other.end)) V = TypeVar("V") @total_ordering class TimestampedValue(Generic[V]): """A timestamped value having a value and a timestamp. Attributes: value: The underlying value. timestamp: Timestamp associated with the value as seconds since Unix epoch. """ def __init__(self, value: V, timestamp: TimestampTypes) -> None: self.value = value self.timestamp = Timestamp.of(timestamp) def __eq__(self, other): return ( type(self) == type(other) and self.value == other.value and self.timestamp == other.timestamp) def __hash__(self): return hash((self.value, self.timestamp)) def __lt__(self, other): if type(self) != type(other): return type(self).__name__ < type(other).__name__ if self.value != other.value: return self.value < other.value return self.timestamp < other.timestamp class GlobalWindow(BoundedWindow): """The default window into which all data is placed (via GlobalWindows).""" _instance: Optional['GlobalWindow'] = None def __new__(cls): if cls._instance is None: cls._instance = super(GlobalWindow, cls).__new__(cls) return cls._instance def __init__(self) -> None: super().__init__(GlobalWindow._getTimestampFromProto()) def __repr__(self): return 'GlobalWindow' def __hash__(self): return hash(type(self)) def __eq__(self, other): # Global windows are always and only equal to each other. return self is other or type(self) is type(other) @property def start(self) -> Timestamp: return MIN_TIMESTAMP @staticmethod def _getTimestampFromProto() -> Timestamp: ts_millis = int( common_urns.constants.GLOBAL_WINDOW_MAX_TIMESTAMP_MILLIS.constant) return Timestamp(micros=ts_millis * 1000) class NonMergingWindowFn(WindowFn): def is_merging(self) -> bool: return False def merge(self, merge_context: WindowFn.MergeContext) -> None: pass # No merging. class GlobalWindows(NonMergingWindowFn): """A windowing function that assigns everything to one global window.""" @classmethod def windowed_batch( cls, batch: Any, timestamp: Timestamp = MIN_TIMESTAMP, pane_info: windowed_value.PaneInfo = windowed_value.PANE_INFO_UNKNOWN ) -> windowed_value.WindowedBatch: return windowed_value.HomogeneousWindowedBatch.of( batch, timestamp, (GlobalWindow(), ), pane_info) @classmethod def windowed_value( cls, value: Any, timestamp: Timestamp = MIN_TIMESTAMP, pane_info: windowed_value.PaneInfo = windowed_value.PANE_INFO_UNKNOWN ) -> WindowedValue: return WindowedValue(value, timestamp, (GlobalWindow(), ), pane_info) @classmethod def windowed_value_at_end_of_window(cls, value): return cls.windowed_value(value, GlobalWindow().max_timestamp()) def assign(self, assign_context: WindowFn.AssignContext) -> list[GlobalWindow]: return [GlobalWindow()] def get_window_coder(self) -> coders.GlobalWindowCoder: return coders.GlobalWindowCoder() def __hash__(self): return hash(type(self)) def __eq__(self, other): # Global windowfn is always and only equal to each other. return self is other or type(self) is type(other) def to_runner_api_parameter(self, context): return common_urns.global_windows.urn, None @staticmethod @urns.RunnerApiFn.register_urn(common_urns.global_windows.urn, None) def from_runner_api_parameter( unused_fn_parameter, unused_context) -> 'GlobalWindows': return GlobalWindows() class FixedWindows(NonMergingWindowFn): """A windowing function that assigns each element to one time interval. The attributes size and offset determine in what time interval a timestamp will be slotted. The time intervals have the following formula: [N * size + offset, (N + 1) * size + offset) Attributes: size: Size of the window as seconds. offset: Offset of this window as seconds. Windows start at t=N * size + offset where t=0 is the UNIX epoch. The offset must be a value in range [0, size). If it is not it will be normalized to this range. """ def __init__(self, size: DurationTypes, offset: TimestampTypes = 0): """Initialize a ``FixedWindows`` function for a given size and offset. Args: size (int): Size of the window in seconds. offset(int): Offset of this window as seconds. Windows start at t=N * size + offset where t=0 is the UNIX epoch. The offset must be a value in range [0, size). If it is not it will be normalized to this range. """ if size <= 0: raise ValueError('The size parameter must be strictly positive.') self.size = Duration.of(size) self.offset = Timestamp.of(offset) % self.size def assign(self, context: WindowFn.AssignContext) -> list[IntervalWindow]: timestamp = context.timestamp start = timestamp - (timestamp - self.offset) % self.size return [IntervalWindow(start, start + self.size)] def get_window_coder(self) -> coders.IntervalWindowCoder: return coders.IntervalWindowCoder() def __eq__(self, other): if type(self) == type(other) == FixedWindows: return self.size == other.size and self.offset == other.offset def __hash__(self): return hash((self.size, self.offset)) def to_runner_api_parameter(self, context): return ( common_urns.fixed_windows.urn, standard_window_fns_pb2.FixedWindowsPayload( size=proto_utils.from_micros( duration_pb2.Duration, self.size.micros), offset=proto_utils.from_micros( timestamp_pb2.Timestamp, self.offset.micros))) @staticmethod @urns.RunnerApiFn.register_urn( common_urns.fixed_windows.urn, standard_window_fns_pb2.FixedWindowsPayload) def from_runner_api_parameter(fn_parameter, unused_context) -> 'FixedWindows': return FixedWindows( size=Duration(micros=proto_utils.to_micros(fn_parameter.size)), offset=Timestamp(micros=proto_utils.to_micros(fn_parameter.offset))) class SlidingWindows(NonMergingWindowFn): """A windowing function that assigns each element to a set of sliding windows. The attributes size and offset determine in what time interval a timestamp will be slotted. The time intervals have the following formula: [N * period + offset, N * period + offset + size) Attributes: size: Size of the window as seconds. period: Period of the windows as seconds. offset: Offset of this window as seconds since Unix epoch. Windows start at t=N * period + offset where t=0 is the epoch. The offset must be a value in range [0, period). If it is not it will be normalized to this range. """ def __init__( self, size: DurationTypes, period: DurationTypes, offset: TimestampTypes = 0, ): if size <= 0: raise ValueError('The size parameter must be strictly positive.') self.size = Duration.of(size) self.period = Duration.of(period) self.offset = Timestamp.of(offset) % period def assign(self, context: WindowFn.AssignContext) -> list[IntervalWindow]: timestamp = context.timestamp start = timestamp - ((timestamp - self.offset) % self.period) return [ IntervalWindow( (interval_start := Timestamp(micros=s)), interval_start + self.size, ) for s in range( start.micros, timestamp.micros - self.size.micros, -self.period.micros) ] def get_window_coder(self) -> coders.IntervalWindowCoder: return coders.IntervalWindowCoder() def __eq__(self, other): if type(self) == type(other) == SlidingWindows: return ( self.size == other.size and self.offset == other.offset and self.period == other.period) def __hash__(self): return hash((self.offset, self.period)) def to_runner_api_parameter(self, context): return ( common_urns.sliding_windows.urn, standard_window_fns_pb2.SlidingWindowsPayload( size=proto_utils.from_micros( duration_pb2.Duration, self.size.micros), offset=proto_utils.from_micros( timestamp_pb2.Timestamp, self.offset.micros), period=proto_utils.from_micros( duration_pb2.Duration, self.period.micros))) @staticmethod @urns.RunnerApiFn.register_urn( common_urns.sliding_windows.urn, standard_window_fns_pb2.SlidingWindowsPayload) def from_runner_api_parameter( fn_parameter, unused_context) -> 'SlidingWindows': return SlidingWindows( size=Duration(micros=proto_utils.to_micros(fn_parameter.size)), offset=Timestamp(micros=proto_utils.to_micros(fn_parameter.offset)), period=Duration(micros=proto_utils.to_micros(fn_parameter.period))) class Sessions(WindowFn): """A windowing function that groups elements into sessions. A session is defined as a series of consecutive events separated by a specified gap size. Attributes: gap_size: Size of the gap between windows as floating-point seconds. """ def __init__(self, gap_size: DurationTypes) -> None: if gap_size <= 0: raise ValueError('The size parameter must be strictly positive.') self.gap_size = Duration.of(gap_size) def assign(self, context: WindowFn.AssignContext) -> list[IntervalWindow]: timestamp = context.timestamp return [IntervalWindow(timestamp, timestamp + self.gap_size)] def get_window_coder(self) -> coders.IntervalWindowCoder: return coders.IntervalWindowCoder() def merge(self, merge_context: WindowFn.MergeContext) -> None: to_merge: list[BoundedWindow] = [] end = MIN_TIMESTAMP for w in sorted(merge_context.windows, key=lambda w: w.start): if to_merge: if end > w.start: to_merge.append(w) if w.end > end: end = w.end else: if len(to_merge) > 1: merge_context.merge( to_merge, IntervalWindow(to_merge[0].start, end)) to_merge = [w] end = w.end else: to_merge = [w] end = w.end if len(to_merge) > 1: merge_context.merge(to_merge, IntervalWindow(to_merge[0].start, end)) def __eq__(self, other): if type(self) == type(other) == Sessions: return self.gap_size == other.gap_size def __hash__(self): return hash(self.gap_size) def to_runner_api_parameter(self, context): return ( common_urns.session_windows.urn, standard_window_fns_pb2.SessionWindowsPayload( gap_size=proto_utils.from_micros( duration_pb2.Duration, self.gap_size.micros))) @staticmethod @urns.RunnerApiFn.register_urn( common_urns.session_windows.urn, standard_window_fns_pb2.SessionWindowsPayload) def from_runner_api_parameter(fn_parameter, unused_context) -> 'Sessions': return Sessions( gap_size=Duration(micros=proto_utils.to_micros(fn_parameter.gap_size))) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/window_test.py0000644000175100017510000003727715206113033023014 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the windowing classes.""" # pytype: skip-file import unittest import apache_beam as beam from apache_beam.coders import coders from apache_beam.runners import pipeline_context from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import CombinePerKey from apache_beam.transforms import Create from apache_beam.transforms import FlatMapTuple from apache_beam.transforms import GroupByKey from apache_beam.transforms import Map from apache_beam.transforms import MapTuple from apache_beam.transforms import WindowInto from apache_beam.transforms import combiners from apache_beam.transforms import core from apache_beam.transforms.core import Windowing from apache_beam.transforms.trigger import AccumulationMode from apache_beam.transforms.trigger import AfterCount from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import GlobalWindow from apache_beam.transforms.window import GlobalWindows from apache_beam.transforms.window import IntervalWindow from apache_beam.transforms.window import NonMergingWindowFn from apache_beam.transforms.window import Sessions from apache_beam.transforms.window import SlidingWindows from apache_beam.transforms.window import TimestampCombiner from apache_beam.transforms.window import TimestampedValue from apache_beam.transforms.window import WindowedValue from apache_beam.transforms.window import WindowFn from apache_beam.utils.timestamp import MAX_TIMESTAMP from apache_beam.utils.timestamp import MIN_TIMESTAMP def context(element, timestamp): return WindowFn.AssignContext(timestamp, element) class ReifyWindowsFn(core.DoFn): def process(self, element, window=core.DoFn.WindowParam): key, values = element yield "%s @ %s" % (key, window), values class TestCustomWindows(NonMergingWindowFn): """A custom non merging window fn which assigns elements into interval windows [0, 3), [3, 5) and [5, element timestamp) based on the element timestamps. """ def assign(self, context): timestamp = context.timestamp if timestamp < 3: return [IntervalWindow(0, 3)] elif timestamp < 5: return [IntervalWindow(3, 5)] else: return [IntervalWindow(5, timestamp)] def get_window_coder(self): return coders.IntervalWindowCoder() class WindowTest(unittest.TestCase): def test_timestamped_value_cmp(self): self.assertEqual(TimestampedValue('a', 2), TimestampedValue('a', 2)) self.assertEqual(TimestampedValue('a', 2), TimestampedValue('a', 2.0)) self.assertNotEqual(TimestampedValue('a', 2), TimestampedValue('a', 2.1)) self.assertNotEqual(TimestampedValue('a', 2), TimestampedValue('b', 2)) def test_global_window(self): self.assertEqual(GlobalWindow(), GlobalWindow()) self.assertNotEqual( GlobalWindow(), IntervalWindow(MIN_TIMESTAMP, MAX_TIMESTAMP)) self.assertNotEqual( IntervalWindow(MIN_TIMESTAMP, MAX_TIMESTAMP), GlobalWindow()) self.assertTrue(GlobalWindow().max_timestamp() < MAX_TIMESTAMP) def test_fixed_windows(self): # Test windows with offset: 2, 7, 12, 17, ... windowfn = FixedWindows(size=5, offset=2) self.assertEqual([IntervalWindow(7, 12)], windowfn.assign(context('v', 7))) self.assertEqual([IntervalWindow(7, 12)], windowfn.assign(context('v', 11))) self.assertEqual([IntervalWindow(12, 17)], windowfn.assign(context('v', 12))) # Test windows without offset: 0, 5, 10, 15, ... windowfn = FixedWindows(size=5) self.assertEqual([IntervalWindow(5, 10)], windowfn.assign(context('v', 5))) self.assertEqual([IntervalWindow(5, 10)], windowfn.assign(context('v', 9))) self.assertEqual([IntervalWindow(10, 15)], windowfn.assign(context('v', 10))) # Test windows with offset out of range. windowfn = FixedWindows(size=5, offset=12) self.assertEqual([IntervalWindow(7, 12)], windowfn.assign(context('v', 11))) def test_sliding_windows_assignment(self): windowfn = SlidingWindows(size=15, period=5, offset=2) expected = [ IntervalWindow(7, 22), IntervalWindow(2, 17), IntervalWindow(-3, 12) ] self.assertEqual(expected, windowfn.assign(context('v', 7))) self.assertEqual(expected, windowfn.assign(context('v', 8))) self.assertEqual(expected, windowfn.assign(context('v', 11))) def test_sliding_windows_assignment_fraction(self): windowfn = SlidingWindows(size=3.5, period=2.5, offset=1.5) self.assertEqual([IntervalWindow(1.5, 5.0), IntervalWindow(-1.0, 2.5)], windowfn.assign(context('v', 1.7))) self.assertEqual([IntervalWindow(1.5, 5.0)], windowfn.assign(context('v', 3))) def test_sliding_windows_assignment_fraction_large_offset(self): windowfn = SlidingWindows(size=3.5, period=2.5, offset=4.0) self.assertEqual([IntervalWindow(1.5, 5.0), IntervalWindow(-1.0, 2.5)], windowfn.assign(context('v', 1.7))) self.assertEqual([IntervalWindow(4.0, 7.5), IntervalWindow(1.5, 5.0)], windowfn.assign(context('v', 4.5))) def test_sessions_merging(self): windowfn = Sessions(10) def merge(*timestamps): windows = [windowfn.assign(context(None, t)) for t in timestamps] running = set() class TestMergeContext(WindowFn.MergeContext): def __init__(self): super().__init__(running) def merge(self, to_be_merged, merge_result): for w in to_be_merged: if w in running: running.remove(w) running.add(merge_result) for ws in windows: running.update(ws) windowfn.merge(TestMergeContext()) windowfn.merge(TestMergeContext()) return sorted(running) self.assertEqual([IntervalWindow(2, 12)], merge(2)) self.assertEqual([IntervalWindow(2, 12), IntervalWindow(19, 29)], merge(2, 19)) self.assertEqual([IntervalWindow(2, 19)], merge(2, 9)) self.assertEqual([IntervalWindow(2, 19)], merge(9, 2)) self.assertEqual([IntervalWindow(2, 19), IntervalWindow(19, 29)], merge(2, 9, 19)) self.assertEqual([IntervalWindow(2, 19), IntervalWindow(19, 29)], merge(19, 9, 2)) self.assertEqual([IntervalWindow(2, 25)], merge(2, 15, 10)) def timestamped_key_values(self, pipeline, key, *timestamps): return ( pipeline | 'start' >> Create(timestamps) | Map(lambda x: WindowedValue((key, x), x, [GlobalWindow()]))) def test_sliding_windows(self): with TestPipeline() as p: pcoll = self.timestamped_key_values(p, 'key', 1, 2, 3) result = ( pcoll | 'w' >> WindowInto(SlidingWindows(period=2, size=4)) | GroupByKey() | beam.MapTuple(lambda k, vs: (k, sorted(vs))) | beam.ParDo(ReifyWindowsFn())) expected = [('key @ [-2.0, 2.0)', [1]), ('key @ [0.0, 4.0)', [1, 2, 3]), ('key @ [2.0, 6.0)', [2, 3])] assert_that(result, equal_to(expected)) def test_sliding_windows_period_longer_than_size(self): with TestPipeline() as p: pcoll = self.timestamped_key_values(p, 'key', 1, 2, 3, 4, 5, 6, 7, 8) result = ( pcoll | 'w' >> WindowInto(SlidingWindows(period=4, size=2)) | GroupByKey() | beam.MapTuple(lambda k, vs: (k, sorted(vs))) | beam.ParDo(ReifyWindowsFn())) expected = [('key @ [0.0, 2.0)', [1]), ('key @ [4.0, 6.0)', [4, 5]), ('key @ [8.0, 10.0)', [8])] assert_that(result, equal_to(expected)) def test_sessions(self): with TestPipeline() as p: pcoll = self.timestamped_key_values(p, 'key', 1, 2, 3, 20, 35, 27) sort_values = Map(lambda k_vs: (k_vs[0], sorted(k_vs[1]))) result = ( pcoll | 'w' >> WindowInto(Sessions(10)) | GroupByKey() | sort_values | beam.ParDo(ReifyWindowsFn())) expected = [('key @ [1.0, 13.0)', [1, 2, 3]), ('key @ [20.0, 45.0)', [20, 27, 35])] assert_that(result, equal_to(expected)) def test_timestamped_value(self): with TestPipeline() as p: result = ( p | 'start' >> Create([(k, k) for k in range(10)]) | Map(lambda x_t: TimestampedValue(x_t[0], x_t[1])) | 'w' >> WindowInto(FixedWindows(5)) | Map(lambda v: ('key', v)) | GroupByKey() | beam.MapTuple(lambda k, vs: (k, sorted(vs)))) assert_that( result, equal_to([('key', [0, 1, 2, 3, 4]), ('key', [5, 6, 7, 8, 9])])) def test_rewindow(self): with TestPipeline() as p: result = ( p | Create([(k, k) for k in range(10)]) | Map(lambda x_t1: TimestampedValue(x_t1[0], x_t1[1])) | 'window' >> WindowInto(SlidingWindows(period=2, size=6)) # Per the model, each element is now duplicated across # three windows. Rewindowing must preserve this duplication. | 'rewindow' >> WindowInto(FixedWindows(5)) | 'rewindow2' >> WindowInto(FixedWindows(5)) | Map(lambda v: ('key', v)) | GroupByKey() | beam.MapTuple(lambda k, vs: (k, sorted(vs)))) assert_that( result, equal_to([('key', sorted([0, 1, 2, 3, 4] * 3)), ('key', sorted([5, 6, 7, 8, 9] * 3))])) def test_rewindow_regroup(self): with TestPipeline() as p: grouped = ( p | Create(range(5)) | Map(lambda t: TimestampedValue(('key', t), t)) | 'window' >> WindowInto(FixedWindows(5, offset=3)) | GroupByKey() | MapTuple(lambda k, vs: (k, sorted(vs)))) # Both of these group-and-ungroup sequences should be idempotent. regrouped1 = ( grouped | 'w1' >> WindowInto(FixedWindows(5, offset=3)) | 'g1' >> GroupByKey() | FlatMapTuple(lambda k, vs: [(k, v) for v in vs])) regrouped2 = ( grouped | FlatMapTuple(lambda k, vs: [(k, v) for v in vs]) | 'w2' >> WindowInto(FixedWindows(5, offset=3)) | 'g2' >> GroupByKey() | MapTuple(lambda k, vs: (k, sorted(vs)))) with_windows = Map(lambda e, w=beam.DoFn.WindowParam: (e, w)) expected = [(('key', [0, 1, 2]), IntervalWindow(-2, 3)), (('key', [3, 4]), IntervalWindow(3, 8))] assert_that(grouped | 'ww' >> with_windows, equal_to(expected)) assert_that( regrouped1 | 'ww1' >> with_windows, equal_to(expected), label='r1') assert_that( regrouped2 | 'ww2' >> with_windows, equal_to(expected), label='r2') def test_timestamped_with_combiners(self): with TestPipeline() as p: result = ( p # Create some initial test values. | 'start' >> Create([(k, k) for k in range(10)]) # The purpose of the WindowInto transform is to establish a # FixedWindows windowing function for the PCollection. # It does not bucket elements into windows since the timestamps # from Create are not spaced 5 ms apart and very likely they all # fall into the same window. | 'w' >> WindowInto(FixedWindows(5)) # Generate timestamped values using the values as timestamps. # Now there are values 5 ms apart and since Map propagates the # windowing function from input to output the output PCollection # will have elements falling into different 5ms windows. | Map(lambda x_t2: TimestampedValue(x_t2[0], x_t2[1])) # We add a 'key' to each value representing the index of the # window. This is important since there is no guarantee of # order for the elements of a PCollection. | Map(lambda v: (v // 5, v))) # Sum all elements associated with a key and window. Although it # is called CombinePerKey it is really CombinePerKeyAndWindow the # same way GroupByKey is really GroupByKeyAndWindow. sum_per_window = result | CombinePerKey(sum) # Compute mean per key and window. mean_per_window = result | combiners.Mean.PerKey() assert_that( sum_per_window, equal_to([(0, 10), (1, 35)]), label='assert:sum') assert_that( mean_per_window, equal_to([(0, 2.0), (1, 7.0)]), label='assert:mean') def test_custom_windows(self): with TestPipeline() as p: pcoll = self.timestamped_key_values(p, 'key', 0, 1, 2, 3, 4, 5, 6) # pylint: disable=abstract-class-instantiated result = ( pcoll | 'custom window' >> WindowInto(TestCustomWindows()) | GroupByKey() | 'sort values' >> MapTuple(lambda k, vs: (k, sorted(vs)))) assert_that( result, equal_to([('key', [0, 1, 2]), ('key', [3, 4]), ('key', [5]), ('key', [6])])) def test_window_assignment_idempotency(self): with TestPipeline() as p: pcoll = self.timestamped_key_values(p, 'key', 0, 2, 4) result = ( pcoll | 'window' >> WindowInto(FixedWindows(2)) | 'same window' >> WindowInto(FixedWindows(2)) | 'same window again' >> WindowInto(FixedWindows(2)) | GroupByKey()) assert_that(result, equal_to([('key', [0]), ('key', [2]), ('key', [4])])) def test_window_assignment_through_multiple_gbk_idempotency(self): with TestPipeline() as p: pcoll = self.timestamped_key_values(p, 'key', 0, 2, 4) result = ( pcoll | 'window' >> WindowInto(FixedWindows(2)) | 'gbk' >> GroupByKey() | 'same window' >> WindowInto(FixedWindows(2)) | 'another gbk' >> GroupByKey() | 'same window again' >> WindowInto(FixedWindows(2)) | 'gbk again' >> GroupByKey()) assert_that( result, equal_to([('key', [[[0]]]), ('key', [[[2]]]), ('key', [[[4]]])])) class RunnerApiTest(unittest.TestCase): def test_windowfn_encoding(self): for window_fn in (GlobalWindows(), FixedWindows(37), SlidingWindows(2, 389), Sessions(5077)): context = pipeline_context.PipelineContext() self.assertEqual( window_fn, WindowFn.from_runner_api(window_fn.to_runner_api(context), context)) def test_windowing_encoding(self): for windowing in (Windowing(GlobalWindows()), Windowing( FixedWindows(1, 3), AfterCount(6), accumulation_mode=AccumulationMode.ACCUMULATING), Windowing( SlidingWindows(10, 15, 21), AfterCount(28), timestamp_combiner=TimestampCombiner.OUTPUT_AT_LATEST, accumulation_mode=AccumulationMode.DISCARDING)): context = pipeline_context.PipelineContext() self.assertEqual( windowing, Windowing.from_runner_api(windowing.to_runner_api(context), context)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/write_ptransform_test.py0000644000175100017510000001021115206113033025065 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the write transform.""" # pytype: skip-file import logging import unittest import apache_beam as beam from apache_beam.io import iobase from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import is_empty from apache_beam.transforms.ptransform import PTransform class _TestSink(iobase.Sink): TEST_INIT_RESULT = 'test_init_result' def __init__(self, return_init_result=True, return_write_results=True): self.return_init_result = return_init_result self.return_write_results = return_write_results def initialize_write(self): if self.return_init_result: return _TestSink.TEST_INIT_RESULT def pre_finalize(self, init_result, writer_results): pass def finalize_write( self, init_result, writer_results, unused_pre_finalize_result): self.init_result_at_finalize = init_result self.write_results_at_finalize = writer_results def open_writer(self, init_result, uid): writer = _TestWriter(init_result, uid, self.return_write_results) return writer class _TestWriter(iobase.Writer): STATE_UNSTARTED, STATE_WRITTEN, STATE_CLOSED = 0, 1, 2 TEST_WRITE_RESULT = 'test_write_result' def __init__(self, init_result, uid, return_write_results=True): self.state = _TestWriter.STATE_UNSTARTED self.init_result = init_result self.uid = uid self.write_output = [] self.return_write_results = return_write_results def close(self): assert self.state in ( _TestWriter.STATE_WRITTEN, _TestWriter.STATE_UNSTARTED) self.state = _TestWriter.STATE_CLOSED if self.return_write_results: return _TestWriter.TEST_WRITE_RESULT def write(self, value): if self.write_output: assert self.state == _TestWriter.STATE_WRITTEN else: assert self.state == _TestWriter.STATE_UNSTARTED self.state = _TestWriter.STATE_WRITTEN self.write_output.append(value) class WriteToTestSink(PTransform): def __init__(self, return_init_result=True, return_write_results=True): self.return_init_result = return_init_result self.return_write_results = return_write_results self.last_sink = None self.label = 'write_to_test_sink' def expand(self, pcoll): self.last_sink = _TestSink( return_init_result=self.return_init_result, return_write_results=self.return_write_results) return pcoll | beam.io.Write(self.last_sink) class WriteTest(unittest.TestCase): DATA = ['some data', 'more data', 'another data', 'yet another data'] def _run_write_test( self, data, return_init_result=True, return_write_results=True): write_to_test_sink = WriteToTestSink( return_init_result, return_write_results) with TestPipeline() as p: result = p | beam.Create(data) | write_to_test_sink | beam.Map(list) assert_that(result, is_empty()) sink = write_to_test_sink.last_sink self.assertIsNotNone(sink) def test_write(self): self._run_write_test(WriteTest.DATA) def test_write_with_empty_pcollection(self): data = [] self._run_write_test(data) def test_write_no_init_result(self): self._run_write_test(WriteTest.DATA, return_init_result=False) def test_write_no_write_results(self): self._run_write_test(WriteTest.DATA, return_write_results=False) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.760448 apache_beam-2.74.0/apache_beam/transforms/xlang/0000755000175100017510000000000015206113160021166 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/transforms/xlang/__init__.py0000644000175100017510000000221215206113033023273 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ This package contains autogenerated Python wrappers for cross-language transforms available in other Beam SDKs. For documentation on creating and using cross-language transforms, see: https://beam.apache.org/documentation/programming-guide/#use-x-lang-transforms For more information on transform wrapper generation, see the top level script: `gen_xlang_wrappers.py` """ ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996269.0 apache_beam-2.74.0/apache_beam/transforms/xlang/io.py0000644000175100017510000001227415206113155022161 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # NOTE: This file contains autogenerated external transform(s) # and should not be edited by hand. # Refer to gen_xlang_wrappers.py for more info. """Cross-language transforms in this module can be imported from the :py:mod:`apache_beam.io` package.""" # pylint:disable=line-too-long from apache_beam.transforms.external import BeamJarExpansionService from apache_beam.transforms.external_transform_provider import ExternalTransform class GenerateSequence(ExternalTransform): """ Outputs a PCollection of Beam Rows, each containing a single INT64 number called "value". The count is produced from the given "start" value and either up to the given "end" or until 2^63 - 1. To produce an unbounded PCollection, simply do not specify an "end" value. Unbounded sequences can specify a "rate" for output elements. In all cases, the sequence of numbers is generated in parallel, so there is no inherent ordering between the generated values """ identifier = "beam:schematransform:org.apache.beam:generate_sequence:v1" def __init__(self, start, end=None, rate=None, expansion_service=None): """ :param start: (int64) The minimum number to generate (inclusive). :param end: (int64) The maximum number to generate (exclusive). Will be an unbounded sequence if left unspecified. :param rate: (Row(elements=, seconds=typing.Optional[int64])) Specifies the rate to generate a given number of elements per a given number of seconds. Applicable only to unbounded sequences. """ self.default_expansion_service = BeamJarExpansionService( "sdks:java:io:expansion-service:shadowJar") super().__init__( start=start, end=end, rate=rate, expansion_service=expansion_service) class TfrecordRead(ExternalTransform): identifier = "beam:schematransform:org.apache.beam:tfrecord_read:v1" def __init__( self, compression, file_pattern, validate, error_handling=None, expansion_service=None): """ :param compression: (str) Decompression type to use when reading input files. :param file_pattern: (str) Filename or file pattern used to find input files. :param validate: (boolean) Validate file pattern. :param error_handling: (Row(output=)) This option specifies whether and where to output unwritable rows. """ self.default_expansion_service = BeamJarExpansionService( "sdks:java:io:expansion-service:shadowJar") super().__init__( compression=compression, file_pattern=file_pattern, validate=validate, error_handling=error_handling, expansion_service=expansion_service) class TfrecordWrite(ExternalTransform): identifier = "beam:schematransform:org.apache.beam:tfrecord_write:v1" def __init__( self, compression, num_shards, output_prefix, error_handling=None, filename_suffix=None, max_num_writers_per_bundle=None, no_spilling=None, shard_template=None, expansion_service=None): """ :param compression: (str) Option to indicate the output sink's compression type. Default is NONE. :param num_shards: (int32) The number of shards to use, or 0 for automatic. :param output_prefix: (str) The directory to which files will be written. :param error_handling: (Row(output=)) This option specifies whether and where to output unwritable rows. :param filename_suffix: (str) The suffix of each file written, combined with prefix and shardTemplate. :param max_num_writers_per_bundle: (int32) Maximum number of writers created in a bundle before spilling to shuffle. :param no_spilling: (boolean) Whether to skip the spilling of data caused by having maxNumWritersPerBundle. :param shard_template: (str) The shard template of each file written, combined with prefix and suffix. """ self.default_expansion_service = BeamJarExpansionService( "sdks:java:io:expansion-service:shadowJar") super().__init__( compression=compression, num_shards=num_shards, output_prefix=output_prefix, error_handling=error_handling, filename_suffix=filename_suffix, max_num_writers_per_bundle=max_num_writers_per_bundle, no_spilling=no_spilling, shard_template=shard_template, expansion_service=expansion_service) ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.7724493 apache_beam-2.74.0/apache_beam/typehints/0000755000175100017510000000000015206113160017706 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/__init__.py0000644000175100017510000000244315206113033022021 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A package defining the syntax and decorator semantics for type-hints.""" # pylint: disable=wildcard-import from apache_beam.typehints.batch import * from apache_beam.typehints.decorators import * from apache_beam.typehints.typehints import * # pylint: disable=ungrouped-imports try: import pandas as _ except ImportError: pass else: from apache_beam.typehints.pandas_type_compatibility import * try: import pyarrow as _ except ImportError: pass else: from apache_beam.typehints.arrow_type_compatibility import * ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/arrow_batching_microbenchmark.py0000644000175100017510000000533515206113033026322 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A microbenchmark for pyarrow batch creation. This microbenchmark exercises the PyarrowBatchConverter.produce_batch method for different batch sizes. """ import argparse import logging import pyarrow as pa from apache_beam.portability.api import schema_pb2 from apache_beam.tools import utils from apache_beam.typehints.arrow_type_compatibility import PyarrowBatchConverter from apache_beam.typehints.arrow_type_compatibility import beam_schema_from_arrow_schema from apache_beam.typehints.schemas import typing_from_runner_api def benchmark_produce_batch(size): batch = pa.Table.from_pydict({ 'foo': pa.array(range(size), type=pa.int64()), 'bar': pa.array([i / size for i in range(size)], type=pa.float64()), 'baz': pa.array([str(i) for i in range(size)], type=pa.string()), }) beam_schema = beam_schema_from_arrow_schema(batch.schema) element_type = typing_from_runner_api( schema_pb2.FieldType(row_type=schema_pb2.RowType(schema=beam_schema))) batch_converter = PyarrowBatchConverter.from_typehints(element_type, pa.Table) elements = list(batch_converter.explode_batch(batch)) def _do_benchmark(): _ = batch_converter.produce_batch(elements) return _do_benchmark def run_benchmark( starting_point=1, num_runs=10, num_elements_step=300, verbose=True): suite = [ utils.LinearRegressionBenchmarkConfig( benchmark_produce_batch, starting_point, num_elements_step, num_runs) ] return utils.run_benchmarks(suite, verbose=verbose) if __name__ == '__main__': logging.basicConfig() parser = argparse.ArgumentParser() parser.add_argument('--num_runs', default=10, type=int) parser.add_argument('--starting_point', default=50, type=int) parser.add_argument('--increment', default=1000, type=int) parser.add_argument('--verbose', default=True, type=bool) options = parser.parse_args() run_benchmark( options.starting_point, options.num_runs, options.increment, options.verbose) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/arrow_type_compatibility.py0000644000175100017510000003336115206113033025411 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utilities for converting between Beam and Arrow schemas. For internal use only, no backward compatibility guarantees. """ # ruff: noqa: UP006 from functools import partial from typing import Dict from typing import List from typing import Optional from typing import Sequence from typing import Tuple import pyarrow as pa from apache_beam.portability.api import schema_pb2 from apache_beam.typehints.batch import BatchConverter from apache_beam.typehints.row_type import RowTypeConstraint from apache_beam.typehints.schemas import typing_from_runner_api from apache_beam.typehints.schemas import typing_to_runner_api from apache_beam.utils import proto_utils __all__ = [] # Get major, minor version PYARROW_VERSION = tuple(map(int, pa.__version__.split('.')[0:2])) BEAM_SCHEMA_ID_KEY = b'beam:schema_id' # We distinguish between schema and field options, because they have to be # combined into arrow Field-level metadata for nested structs. BEAM_SCHEMA_OPTION_KEY_PREFIX = b'beam:schema_option:' BEAM_FIELD_OPTION_KEY_PREFIX = b'beam:field_option:' def _hydrate_beam_option(encoded_option: bytes) -> schema_pb2.Option: return proto_utils.parse_Bytes(encoded_option, schema_pb2.Option) def beam_schema_from_arrow_schema(arrow_schema: pa.Schema) -> schema_pb2.Schema: if arrow_schema.metadata: schema_id = arrow_schema.metadata.get(BEAM_SCHEMA_ID_KEY, None) schema_options = [ _hydrate_beam_option(value) for key, value in arrow_schema.metadata.items() if key.startswith(BEAM_SCHEMA_OPTION_KEY_PREFIX) ] else: schema_id = None schema_options = [] return schema_pb2.Schema( fields=[ _beam_field_from_arrow_field(arrow_schema.field(i)) for i in range(len(arrow_schema.types)) ], options=schema_options, id=schema_id) def _beam_field_from_arrow_field(arrow_field: pa.Field) -> schema_pb2.Field: beam_fieldtype = _beam_fieldtype_from_arrow_field(arrow_field) if arrow_field.metadata: field_options = [ _hydrate_beam_option(value) for key, value in arrow_field.metadata.items() if key.startswith(BEAM_FIELD_OPTION_KEY_PREFIX) ] if isinstance(arrow_field.type, pa.StructType): beam_fieldtype.row_type.schema.options.extend([ _hydrate_beam_option(value) for key, value in arrow_field.metadata.items() if key.startswith(BEAM_SCHEMA_OPTION_KEY_PREFIX) ]) if BEAM_SCHEMA_ID_KEY in arrow_field.metadata: beam_fieldtype.row_type.schema.id = arrow_field.metadata[ BEAM_SCHEMA_ID_KEY] else: field_options = None return schema_pb2.Field( name=arrow_field.name, type=beam_fieldtype, options=field_options, ) def _beam_fieldtype_from_arrow_field( arrow_field: pa.Field) -> schema_pb2.FieldType: beam_fieldtype = _beam_fieldtype_from_arrow_type(arrow_field.type) beam_fieldtype.nullable = arrow_field.nullable return beam_fieldtype def _beam_fieldtype_from_arrow_type( arrow_type: pa.DataType) -> schema_pb2.FieldType: if arrow_type in PYARROW_TO_ATOMIC_TYPE: return schema_pb2.FieldType(atomic_type=PYARROW_TO_ATOMIC_TYPE[arrow_type]) elif isinstance(arrow_type, pa.ListType): return schema_pb2.FieldType( array_type=schema_pb2.ArrayType( element_type=_beam_fieldtype_from_arrow_field( arrow_type.value_field))) elif isinstance(arrow_type, pa.MapType): return schema_pb2.FieldType(map_type=_arrow_map_to_beam_map(arrow_type)) elif isinstance(arrow_type, pa.StructType): return schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( fields=[ _beam_field_from_arrow_field(arrow_type[i]) for i in range(len(arrow_type)) ], ))) else: raise ValueError(f"Unrecognized arrow type: {arrow_type!r}") def _option_as_arrow_metadata(beam_option: schema_pb2.Option, *, prefix: bytes) -> Tuple[bytes, bytes]: return ( prefix + beam_option.name.encode('UTF-8'), beam_option.SerializeToString()) _field_option_as_arrow_metadata = partial( _option_as_arrow_metadata, prefix=BEAM_FIELD_OPTION_KEY_PREFIX) _schema_option_as_arrow_metadata = partial( _option_as_arrow_metadata, prefix=BEAM_SCHEMA_OPTION_KEY_PREFIX) def arrow_schema_from_beam_schema(beam_schema: schema_pb2.Schema) -> pa.Schema: return pa.schema( [_arrow_field_from_beam_field(field) for field in beam_schema.fields], { BEAM_SCHEMA_ID_KEY: beam_schema.id, **dict( _schema_option_as_arrow_metadata(option) for option in beam_schema.options) # pylint: disable=line-too-long }, ) def _arrow_field_from_beam_field(beam_field: schema_pb2.Field) -> pa.Field: return _arrow_field_from_beam_fieldtype( beam_field.type, name=beam_field.name, field_options=beam_field.options) _ARROW_PRIMITIVE_MAPPING = [ # TODO(https://github.com/apache/beam/issues/23816): Support unsigned ints # and float16 (schema_pb2.BYTE, pa.int8()), (schema_pb2.INT16, pa.int16()), (schema_pb2.INT32, pa.int32()), (schema_pb2.INT64, pa.int64()), (schema_pb2.FLOAT, pa.float32()), (schema_pb2.DOUBLE, pa.float64()), (schema_pb2.BOOLEAN, pa.bool_()), (schema_pb2.STRING, pa.string()), (schema_pb2.BYTES, pa.binary()), ] ATOMIC_TYPE_TO_PYARROW = { beam: arrow for beam, arrow in _ARROW_PRIMITIVE_MAPPING } PYARROW_TO_ATOMIC_TYPE = { arrow: beam for beam, arrow in _ARROW_PRIMITIVE_MAPPING } def _arrow_field_from_beam_fieldtype( beam_fieldtype: schema_pb2.FieldType, name=b'', field_options: Sequence[schema_pb2.Option] = None) -> pa.DataType: arrow_type = _arrow_type_from_beam_fieldtype(beam_fieldtype) if field_options is not None: metadata = dict( _field_option_as_arrow_metadata(field_option) for field_option in field_options) else: metadata = {} type_info = beam_fieldtype.WhichOneof("type_info") if type_info == "row_type": schema = beam_fieldtype.row_type.schema metadata.update( dict( _schema_option_as_arrow_metadata(schema_option) for schema_option in schema.options)) if schema.id: metadata[BEAM_SCHEMA_ID_KEY] = schema.id return pa.field( name=name, type=arrow_type, nullable=beam_fieldtype.nullable, metadata=metadata, ) if PYARROW_VERSION < (6, 0): # In pyarrow < 6.0.0 we cannot construct a MapType object from Field # instances, pa.map_ will only accept DataType instances. This makes it # impossible to propagate nullability. # # Note this was changed in: # https://github.com/apache/arrow/commit/64bef2ad8d9cd2fea122cfa079f8ca3fea8cdf5d # # Here we define a custom arrow map conversion function to handle these cases # and error as appropriate. def _make_arrow_map(beam_map_type: schema_pb2.MapType): if beam_map_type.key_type.nullable: raise TypeError('Arrow map key field cannot be nullable') elif beam_map_type.value_type.nullable: raise TypeError( "pyarrow<6 does not support creating maps with nullable " "values. Please use pyarrow>=6.0.0") return pa.map_( _arrow_type_from_beam_fieldtype(beam_map_type.key_type), _arrow_type_from_beam_fieldtype(beam_map_type.value_type)) def _arrow_map_to_beam_map(arrow_map_type): return schema_pb2.MapType( key_type=_beam_fieldtype_from_arrow_type(arrow_map_type.key_type), value_type=_beam_fieldtype_from_arrow_type(arrow_map_type.item_type)) else: def _make_arrow_map(beam_map_type: schema_pb2.MapType): return pa.map_( _arrow_field_from_beam_fieldtype(beam_map_type.key_type), _arrow_field_from_beam_fieldtype(beam_map_type.value_type)) def _arrow_map_to_beam_map(arrow_map_type): return schema_pb2.MapType( key_type=_beam_fieldtype_from_arrow_field(arrow_map_type.key_field), value_type=_beam_fieldtype_from_arrow_field(arrow_map_type.item_field)) def _arrow_type_from_beam_fieldtype( beam_fieldtype: schema_pb2.FieldType, ) -> Tuple[pa.DataType, Optional[Dict[bytes, bytes]]]: # Note this function is not concerned with beam_fieldtype.nullable, as # nullability is a property of the Field in Arrow. type_info = beam_fieldtype.WhichOneof("type_info") if type_info == 'atomic_type': try: output_arrow_type = ATOMIC_TYPE_TO_PYARROW[beam_fieldtype.atomic_type] except KeyError: raise ValueError( "Unsupported atomic type: {0}".format(beam_fieldtype.atomic_type)) elif type_info == "array_type": output_arrow_type = pa.list_( _arrow_field_from_beam_fieldtype( beam_fieldtype.array_type.element_type)) elif type_info == "map_type": output_arrow_type = _make_arrow_map(beam_fieldtype.map_type) elif type_info == "row_type": schema = beam_fieldtype.row_type.schema # Note schema id and options are handled at the arrow field level, they are # added at field-level metadata. output_arrow_type = pa.struct( [_arrow_field_from_beam_field(field) for field in schema.fields]) elif type_info == "logical_type": # TODO(https://github.com/apache/beam/issues/23817): Add support for logical # types. raise NotImplementedError( "Beam logical types are not currently supported " "in arrow_type_compatibility.") else: raise ValueError(f"Unrecognized type_info: {type_info!r}") return output_arrow_type class PyarrowBatchConverter(BatchConverter): def __init__(self, element_type: RowTypeConstraint): super().__init__(pa.Table, element_type) self._beam_schema = typing_to_runner_api(element_type).row_type.schema arrow_schema = arrow_schema_from_beam_schema(self._beam_schema) self._arrow_schema = arrow_schema @staticmethod def from_typehints(element_type, batch_type) -> Optional['PyarrowBatchConverter']: assert batch_type == pa.Table if not isinstance(element_type, RowTypeConstraint): element_type = RowTypeConstraint.from_user_type(element_type) if element_type is None: raise TypeError( f"Element type {element_type} must be compatible with Beam Schemas " "(https://beam.apache.org/documentation/programming-guide/#schemas)" " for batch type pa.Table.") return PyarrowBatchConverter(element_type) def produce_batch(self, elements): arrays = [ pa.array([getattr(el, name) for el in elements], type=self._arrow_schema.field(name).type) for name, _ in self._element_type._fields ] return pa.Table.from_arrays(arrays, schema=self._arrow_schema) def explode_batch(self, batch: pa.Table): """Convert an instance of B to Generator[E].""" for row_values in zip(*batch.columns): yield self._element_type.user_type( **{ name: val.as_py() for name, val in zip(self._arrow_schema.names, row_values) }) def combine_batches(self, batches: List[pa.Table]): return pa.concat_tables(batches) def get_length(self, batch: pa.Table): return batch.num_rows def estimate_byte_size(self, batch: pa.Table): return batch.nbytes @staticmethod def _from_serialized_schema(serialized_schema): beam_schema = proto_utils.parse_Bytes(serialized_schema, schema_pb2.Schema) element_type = typing_from_runner_api( schema_pb2.FieldType(row_type=schema_pb2.RowType(schema=beam_schema))) return PyarrowBatchConverter(element_type) def __reduce__(self): return self._from_serialized_schema, ( self._beam_schema.SerializeToString(), ) class PyarrowArrayBatchConverter(BatchConverter): def __init__(self, element_type: type): super().__init__(pa.Array, element_type) self._element_type = element_type beam_fieldtype = typing_to_runner_api(element_type) self._arrow_type = _arrow_type_from_beam_fieldtype(beam_fieldtype) @staticmethod def from_typehints(element_type, batch_type) -> Optional['PyarrowArrayBatchConverter']: assert batch_type == pa.Array return PyarrowArrayBatchConverter(element_type) def produce_batch(self, elements): return pa.array(list(elements), type=self._arrow_type) def explode_batch(self, batch: pa.Array): """Convert an instance of B to Generator[E].""" for val in batch: yield val.as_py() def combine_batches(self, batches: List[pa.Array]): return pa.concat_arrays(batches) def get_length(self, batch: pa.Array): return batch.num_rows def estimate_byte_size(self, batch: pa.Array): return batch.nbytes @BatchConverter.register(name="pyarrow") def create_pyarrow_batch_converter( element_type: type, batch_type: type) -> BatchConverter: if batch_type == pa.Table: return PyarrowBatchConverter.from_typehints( element_type=element_type, batch_type=batch_type) elif batch_type == pa.Array: return PyarrowArrayBatchConverter.from_typehints( element_type=element_type, batch_type=batch_type) raise TypeError("batch type must be pa.Table or pa.Array") ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/arrow_type_compatibility_test.py0000644000175100017510000001716615206113033026455 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for arrow_type_compatibility.""" import logging import unittest from typing import Any from typing import Optional import pyarrow as pa import pytest from parameterized import parameterized from parameterized import parameterized_class from apache_beam.typehints import row_type from apache_beam.typehints import typehints from apache_beam.typehints.arrow_type_compatibility import arrow_schema_from_beam_schema from apache_beam.typehints.arrow_type_compatibility import beam_schema_from_arrow_schema from apache_beam.typehints.batch import BatchConverter from apache_beam.typehints.batch_test import temp_seed from apache_beam.typehints.schemas_test import get_test_beam_schemas_protos @pytest.mark.uses_pyarrow class ArrowTypeCompatibilityTest(unittest.TestCase): @parameterized.expand([(beam_schema, ) for beam_schema in get_test_beam_schemas_protos()]) def test_beam_schema_survives_roundtrip(self, beam_schema): roundtripped = beam_schema_from_arrow_schema( arrow_schema_from_beam_schema(beam_schema)) self.assertEqual(beam_schema, roundtripped) @parameterized_class([ { 'batch_typehint': pa.Table, 'element_typehint': row_type.RowTypeConstraint.from_fields([ ('foo', Optional[int]), ('bar', Optional[float]), ('baz', Optional[str]), ]), 'batch': pa.Table.from_pydict({ 'foo': pa.array(range(100), type=pa.int64()), 'bar': pa.array([i / 100 for i in range(100)], type=pa.float64()), 'baz': pa.array([str(i) for i in range(100)], type=pa.string()), }), }, { 'batch_typehint': pa.Table, 'element_typehint': row_type.RowTypeConstraint.from_fields([ ('foo', Optional[int]), ( 'nested', Optional[row_type.RowTypeConstraint.from_fields([ ("bar", Optional[float]), # noqa: F821 ("baz", Optional[str]), # noqa: F821 ])]), ]), 'batch': pa.Table.from_pydict({ 'foo': pa.array(range(100), type=pa.int64()), 'nested': pa.array([ None if i % 11 else { 'bar': i / 100, 'baz': str(i) } for i in range(100) ]), }), }, { 'batch_typehint': pa.Array, 'element_typehint': int, 'batch': pa.array(range(100), type=pa.int64()), }, { 'batch_typehint': pa.Array, 'element_typehint': row_type.RowTypeConstraint.from_fields([ ("bar", Optional[float]), # noqa: F821 ("baz", Optional[str]), # noqa: F821 ]), 'batch': pa.array([{ 'bar': i / 100, 'baz': str(i) } if i % 7 else None for i in range(100)]), } ]) @pytest.mark.uses_pyarrow class ArrowBatchConverterTest(unittest.TestCase): def create_batch_converter(self): return BatchConverter.from_typehints( element_type=self.element_typehint, batch_type=self.batch_typehint) def setUp(self): self.converter = self.create_batch_converter() self.normalized_batch_typehint = typehints.normalize(self.batch_typehint) self.normalized_element_typehint = typehints.normalize( self.element_typehint) def equality_check(self, left, right): if isinstance(left, pa.Array): self.assertTrue(left.equals(right)) else: self.assertEqual(left, right) def test_typehint_validates(self): typehints.validate_composite_type_param(self.batch_typehint, '') typehints.validate_composite_type_param(self.element_typehint, '') def test_type_check(self): typehints.check_constraint(self.normalized_batch_typehint, self.batch) def test_type_check_element(self): for element in self.converter.explode_batch(self.batch): typehints.check_constraint(self.normalized_element_typehint, element) def test_explode_rebatch(self): exploded = list(self.converter.explode_batch(self.batch)) rebatched = self.converter.produce_batch(exploded) typehints.check_constraint(self.normalized_batch_typehint, rebatched) self.equality_check(self.batch, rebatched) def test_estimate_byte_size_implemented(self): # Just verify that we can call byte size self.assertGreater(self.converter.estimate_byte_size(self.batch), 0) @parameterized.expand([ (2, ), (3, ), (10, ), ]) def test_estimate_byte_size_partitions(self, N): elements = list(self.converter.explode_batch(self.batch)) # Split elements into N contiguous partitions, create a batch out of each batches = [ self.converter.produce_batch( elements[len(elements) * i // N:len(elements) * (i + 1) // N]) for i in range(N) ] # Some estimate_byte_size implementations use random samples, # set a seed temporarily to make this test deterministic with temp_seed(12345): partitioned_size_estimate = sum( self.converter.estimate_byte_size(batch) for batch in batches) size_estimate = self.converter.estimate_byte_size(self.batch) # Assert that size estimate for partitions is within 10% of size estimate # for the whole partition. self.assertLessEqual( abs(partitioned_size_estimate / size_estimate - 1), 0.1) @parameterized.expand([ (2, ), (3, ), (10, ), ]) def test_combine_batches(self, N): elements = list(self.converter.explode_batch(self.batch)) # Split elements into N contiguous partitions, create a batch out of each batches = [ self.converter.produce_batch( elements[len(elements) * i // N:len(elements) * (i + 1) // N]) for i in range(N) ] # Combine the batches, output should be equivalent to the original batch combined = self.converter.combine_batches(batches) self.equality_check(self.batch, combined) def test_equals(self): self.assertTrue(self.converter == self.create_batch_converter()) self.assertTrue(self.create_batch_converter() == self.converter) def test_hash(self): self.assertEqual(hash(self.create_batch_converter()), hash(self.converter)) class ArrowBatchConverterErrorsTest(unittest.TestCase): @parameterized.expand([ ( pa.RecordBatch, row_type.RowTypeConstraint.from_fields([ ("bar", Optional[float]), # noqa: F821 ("baz", Optional[str]), # noqa: F821 ]), r'batch type must be pa\.Table or pa\.Array', ), ( pa.Table, Any, r'Element type .* must be compatible with Beam Schemas', ), ]) def test_construction_errors( self, batch_typehint, element_typehint, error_regex): with self.assertRaisesRegex(TypeError, error_regex): BatchConverter.from_typehints( element_type=element_typehint, batch_type=batch_typehint) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/batch.py0000644000175100017510000002245215206113033021345 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utilities for working with batched types in the Beam SDK. A batched type is a type B that is logically equivalent to Sequence[E], where E is some other type. Typically B has a different physical representation than Sequence[E] for performance reasons. A trivial example is B=np.array(dtype=np.int64), E=int. """ import random from math import ceil from typing import Callable from typing import Generic from typing import Iterator from typing import Mapping from typing import Optional from typing import Sequence from typing import TypeVar import numpy as np from apache_beam import coders from apache_beam.typehints import typehints __all__ = ['BatchConverter'] B = TypeVar('B') E = TypeVar('E') BatchConverterConstructor = Callable[[type, type], 'BatchConverter'] BATCH_CONVERTER_REGISTRY: Mapping[str, BatchConverterConstructor] = {} __all__ = ['BatchConverter'] class BatchConverter(Generic[B, E]): def __init__(self, batch_type, element_type): self._batch_type = batch_type self._element_type = element_type def produce_batch(self, elements: Sequence[E]) -> B: """Convert an instance of List[E] to a single instance of B.""" raise NotImplementedError def explode_batch(self, batch: B) -> Iterator[E]: """Convert an instance of B to Iterator[E].""" raise NotImplementedError def combine_batches(self, batches: Sequence[B]) -> B: raise NotImplementedError def get_length(self, batch: B) -> int: raise NotImplementedError def estimate_byte_size(self, batch): raise NotImplementedError @staticmethod def register(*, name: str): def do_registration( batch_converter_constructor: Callable[[type, type], 'BatchConverter']): if name in BATCH_CONVERTER_REGISTRY: raise AssertionError( f"Attempted to register two batch converters with name {name}") BATCH_CONVERTER_REGISTRY[name] = batch_converter_constructor return batch_converter_constructor return do_registration @staticmethod def from_typehints(*, element_type, batch_type) -> 'BatchConverter': element_type = typehints.normalize(element_type) batch_type = typehints.normalize(batch_type) errors = {} for name, constructor in BATCH_CONVERTER_REGISTRY.items(): try: return constructor(element_type, batch_type) except TypeError as e: errors[name] = e.args[0] error_summaries = '\n\n'.join( f"{name}:\n\t{msg}" for name, msg in errors.items()) raise TypeError( f"Unable to find BatchConverter for element_type={element_type!r} and " f"batch_type={batch_type!r}. Error summaries:\n\n{error_summaries}") @property def batch_type(self): return self._batch_type @property def element_type(self): return self._element_type def __key(self): return (self._element_type, self._batch_type) def __eq__(self, other: 'BatchConverter') -> bool: if isinstance(other, BatchConverter): return self.__key() == other.__key() return NotImplemented def __hash__(self) -> int: return hash(self.__key()) class ListBatchConverter(BatchConverter): SAMPLE_FRACTION = 0.2 MAX_SAMPLES = 100 SAMPLED_BATCH_SIZE = MAX_SAMPLES / SAMPLE_FRACTION def __init__(self, batch_type, element_type): super().__init__(batch_type, element_type) self.element_coder = coders.registry.get_coder(element_type) @staticmethod @BatchConverter.register(name="list") def from_typehints(element_type, batch_type): if (not isinstance(batch_type, typehints.ListConstraint) or batch_type.inner_type != element_type): raise TypeError("batch type must be List[T] for element type T") return ListBatchConverter(batch_type, element_type) def produce_batch(self, elements): return list(elements) def explode_batch(self, batch): return iter(batch) def combine_batches(self, batches): return sum(batches, []) def get_length(self, batch): return len(batch) def estimate_byte_size(self, batch): # randomly sample a fraction of the elements and use the element_coder to # estimate the size of each nsampled = ( ceil(len(batch) * self.SAMPLE_FRACTION) if len(batch) < self.SAMPLED_BATCH_SIZE else self.MAX_SAMPLES) mean_byte_size = sum( self.element_coder.estimate_size(element) for element in random.sample(batch, nsampled)) / nsampled return ceil(mean_byte_size * len(batch)) N = "ARBITRARY LENGTH DIMENSION" class NumpyBatchConverter(BatchConverter): def __init__( self, batch_type, element_type, dtype, element_shape=(), partition_dimension=0): super().__init__(batch_type, element_type) self.dtype = np.dtype(dtype) self.element_shape = element_shape self.partition_dimension = partition_dimension @staticmethod @BatchConverter.register(name="numpy") def from_typehints(element_type, batch_type) -> Optional['NumpyBatchConverter']: if not isinstance(element_type, NumpyTypeHint.NumpyTypeConstraint): try: element_type = NumpyArray[element_type, ()] except TypeError as e: raise TypeError("Element type is not a dtype") from e if not isinstance(batch_type, NumpyTypeHint.NumpyTypeConstraint): if not batch_type == np.ndarray: raise TypeError( "batch type must be np.ndarray or " "beam.typehints.batch.NumpyArray[..]") batch_type = NumpyArray[element_type.dtype, (N, )] if not batch_type.dtype == element_type.dtype: raise TypeError( "batch type and element type must have equivalent dtypes " f"(batch={batch_type.dtype}, element={element_type.dtype})") computed_element_shape = list(batch_type.shape) partition_dimension = computed_element_shape.index(N) computed_element_shape.pop(partition_dimension) if not tuple(computed_element_shape) == element_type.shape: raise TypeError( "Failed to align batch type's batch dimension with element type. " f"(batch type dimensions: {batch_type.shape}, element type " f"dimenstions: {element_type.shape}") return NumpyBatchConverter( batch_type, element_type, batch_type.dtype, element_type.shape, partition_dimension) def produce_batch(self, elements): return np.stack(elements, axis=self.partition_dimension) def explode_batch(self, batch): """Convert an instance of B to Generator[E].""" yield from batch.swapaxes(self.partition_dimension, 0) def combine_batches(self, batches): return np.concatenate(batches, axis=self.partition_dimension) def get_length(self, batch): return np.size(batch, axis=self.partition_dimension) def estimate_byte_size(self, batch): return batch.nbytes # numpy is starting to add typehints, which we should support # https://numpy.org/doc/stable/reference/typing.html for now they don't allow # specifying shape, seems to be coming after # https://www.python.org/dev/peps/pep-0646/ class NumpyTypeHint(): class NumpyTypeConstraint(typehints.TypeConstraint): def __init__(self, dtype, shape=()): self.dtype = np.dtype(dtype) self.shape = shape def type_check(self, batch): if not isinstance(batch, np.ndarray): raise TypeError(f"Batch {batch!r} is not an instance of ndarray") if not np.issubdtype(batch.dtype, self.dtype): raise TypeError( f"Batch {batch!r} does not have expected dtype: {self.dtype!r}") for dim in range(len(self.shape)): if not self.shape[dim] == N and not batch.shape[dim] == self.shape[dim]: raise TypeError( f"Batch {batch!r} does not have expected shape: {self.shape!r}") def _consistent_with_check_(self, sub): # TODO Check sub against batch type, and element type return True def __key(self): return (self.dtype, self.shape) def __eq__(self, other) -> bool: if isinstance(other, NumpyTypeHint.NumpyTypeConstraint): return self.__key() == other.__key() return NotImplemented def __hash__(self) -> int: return hash(self.__key()) def __repr__(self): if self.shape == (N, ): return f'NumpyArray[{self.dtype!r}]' else: return f'NumpyArray[{self.dtype!r}, {self.shape!r}]' def __getitem__(self, value): if isinstance(value, tuple): if len(value) == 2: dtype, shape = value return self.NumpyTypeConstraint(dtype, shape=shape) else: raise ValueError else: dtype = value return self.NumpyTypeConstraint( dtype, shape=(N, )) NumpyArray = NumpyTypeHint() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/batch_test.py0000644000175100017510000001442415206113033022404 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for batched type converters.""" # ruff: noqa: UP006 import contextlib import random import typing import unittest import numpy as np from parameterized import parameterized from parameterized import parameterized_class from apache_beam.typehints import typehints from apache_beam.typehints.batch import BatchConverter from apache_beam.typehints.batch import N from apache_beam.typehints.batch import NumpyArray @parameterized_class([ { 'batch_typehint': np.ndarray, 'element_typehint': np.int32, 'batch': np.array(range(100), np.int32) }, { 'batch_typehint': NumpyArray[np.int64, (N, 10)], 'element_typehint': NumpyArray[np.int64, (10, )], 'batch': np.array([list(range(i, i + 10)) for i in range(100)], np.int64), }, { 'batch_typehint': typehints.List[str], 'element_typehint': str, 'batch': ["foo" * (i % 5) + str(i) for i in range(1000)], }, { 'batch_typehint': typing.List[str], 'element_typehint': str, 'batch': ["foo" * (i % 5) + str(i) for i in range(1000)], }, ]) class BatchConverterTest(unittest.TestCase): def create_batch_converter(self): return BatchConverter.from_typehints( element_type=self.element_typehint, batch_type=self.batch_typehint) def setUp(self): self.converter = self.create_batch_converter() self.normalized_batch_typehint = typehints.normalize(self.batch_typehint) self.normalized_element_typehint = typehints.normalize( self.element_typehint) def equality_check(self, left, right): if isinstance(left, np.ndarray) and isinstance(right, np.ndarray): return np.array_equal(left, right) else: return left == right def test_typehint_validates(self): typehints.validate_composite_type_param(self.batch_typehint, '') typehints.validate_composite_type_param(self.element_typehint, '') def test_type_check(self): typehints.check_constraint(self.normalized_batch_typehint, self.batch) def test_type_check_element(self): for element in self.converter.explode_batch(self.batch): typehints.check_constraint(self.normalized_element_typehint, element) def test_explode_rebatch(self): exploded = list(self.converter.explode_batch(self.batch)) rebatched = self.converter.produce_batch(exploded) typehints.check_constraint(self.normalized_batch_typehint, rebatched) self.assertTrue(self.equality_check(self.batch, rebatched)) def test_estimate_byte_size_implemented(self): # Just verify that we can call byte size self.assertGreater(self.converter.estimate_byte_size(self.batch), 0) @parameterized.expand([ (2, ), (3, ), (10, ), ]) def test_estimate_byte_size_partitions(self, N): elements = list(self.converter.explode_batch(self.batch)) # Split elements into N contiguous partitions, create a batch out of each batches = [ self.converter.produce_batch( elements[len(elements) * i // N:len(elements) * (i + 1) // N]) for i in range(N) ] # Some estimate_byte_size implementations use random samples, # set a seed temporarily to make this test deterministic with temp_seed(12345): partitioned_size_estimate = sum( self.converter.estimate_byte_size(batch) for batch in batches) size_estimate = self.converter.estimate_byte_size(self.batch) # Assert that size estimate for partitions is within 10% of size estimate # for the whole partition. self.assertLessEqual( abs(partitioned_size_estimate / size_estimate - 1), 0.1) @parameterized.expand([ (2, ), (3, ), (10, ), ]) def test_combine_batches(self, N): elements = list(self.converter.explode_batch(self.batch)) # Split elements into N contiguous partitions, create a batch out of each batches = [ self.converter.produce_batch( elements[len(elements) * i // N:len(elements) * (i + 1) // N]) for i in range(N) ] # Combine the batches, output should be equivalent to the original batch combined = self.converter.combine_batches(batches) self.assertTrue(self.equality_check(self.batch, combined)) def test_equals(self): self.assertTrue(self.converter == self.create_batch_converter()) self.assertTrue(self.create_batch_converter() == self.converter) def test_hash(self): self.assertEqual(hash(self.create_batch_converter()), hash(self.converter)) class BatchConverterErrorsTest(unittest.TestCase): @parameterized.expand([ ( typing.List[int], str, r'batch type must be List\[T\] for element type T', ), ( np.ndarray, typing.Any, r'Element type is not a dtype', ), ( np.array, np.int64, ( r'batch type must be np\.ndarray or ' r'beam\.typehints\.batch\.NumpyArray\[\.\.\]'), ), ( NumpyArray[np.int64, (3, N, 2)], NumpyArray[np.int64, (3, 7)], r'Failed to align batch type\'s batch dimension', ), ]) def test_construction_errors( self, batch_typehint, element_typehint, error_regex): with self.assertRaisesRegex(TypeError, error_regex): BatchConverter.from_typehints( element_type=element_typehint, batch_type=batch_typehint) @contextlib.contextmanager def temp_seed(seed): state = random.getstate() random.seed(seed) try: yield finally: random.setstate(state) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/decorators.py0000644000175100017510000011556115206113033022435 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Type hinting decorators allowing static or runtime type-checking for the SDK. This module defines decorators which utilize the type-hints defined in 'type_hints.py' to allow annotation of the types of function arguments and return values. Type-hints for functions are annotated using two separate decorators. One is for type-hinting the types of function arguments, the other for type-hinting the function return value. Type-hints can either be specified in the form of positional arguments:: @with_input_types(int, int) def add(a, b): return a + b Keyword arguments:: @with_input_types(a=int, b=int) def add(a, b): return a + b Or even a mix of both:: @with_input_types(int, b=int) def add(a, b): return a + b Example usage for type-hinting arguments only:: @with_input_types(s=str) def to_lower(a): return a.lower() Example usage for type-hinting return values only:: @with_output_types(Tuple[int, bool]) def compress_point(ec_point): return ec_point.x, ec_point.y < 0 Example usage for type-hinting both arguments and return values:: @with_input_types(a=int) @with_output_types(str) def int_to_str(a): return str(a) The valid type-hint for such as function looks like the following:: @with_input_types(a=int, b=int) def foo((a, b)): ... Notice that we hint the type of each unpacked argument independently, rather than hinting the type of the tuple as a whole (Tuple[int, int]). Optionally, type-hints can be type-checked at runtime. To toggle this behavior this module defines two functions: 'enable_run_time_type_checking' and 'disable_run_time_type_checking'. NOTE: for this toggle behavior to work properly it must appear at the top of the module where all functions are defined, or before importing a module containing type-hinted functions. """ # pytype: skip-file # ruff: noqa: UP006 import inspect import itertools import logging import traceback import types from typing import Any from typing import Callable from typing import Dict from typing import Iterable from typing import List from typing import Literal from typing import NamedTuple from typing import Optional from typing import Tuple from typing import TypeVar from typing import Union from typing import get_args from typing import get_origin from apache_beam.pvalue import TaggedOutput from apache_beam.typehints import native_type_compatibility from apache_beam.typehints import typehints from apache_beam.typehints.native_type_compatibility import convert_to_beam_type from apache_beam.typehints.typehints import CompositeTypeHintError from apache_beam.typehints.typehints import SimpleTypeHintError from apache_beam.typehints.typehints import check_constraint from apache_beam.typehints.typehints import validate_composite_type_param __all__ = [ 'disable_type_annotations', 'no_annotations', 'with_input_types', 'with_output_types', 'WithTypeHints', 'TypeCheckError', ] T = TypeVar('T') WithTypeHintsT = TypeVar('WithTypeHintsT', bound='WithTypeHints') # pylint: disable=invalid-name # This is missing in the builtin types module. str.upper is arbitrary, any # method on a C-implemented type will do. # pylint: disable=invalid-name _MethodDescriptorType = type(str.upper) # pylint: enable=invalid-name _ANY_VAR_POSITIONAL = typehints.Tuple[typehints.Any, ...] _ANY_VAR_KEYWORD = typehints.Dict[typehints.Any, typehints.Any] _disable_from_callable = False # pylint: disable=invalid-name def get_signature(func): """Like inspect.signature(), but supports Py2 as well. This module uses inspect.signature instead of getfullargspec since in the latter: 'the "self" parameter is always reported, even for bound methods' https://github.com/python/cpython/blob/44f91c388a6f4da9ed3300df32ca290b8aa104ea/Lib/inspect.py#L1103 """ try: signature = inspect.signature(func) except ValueError: # Fall back on a catch-all signature. params = [ inspect.Parameter('_', inspect.Parameter.POSITIONAL_OR_KEYWORD), inspect.Parameter( '__unknown__varargs', inspect.Parameter.VAR_POSITIONAL), inspect.Parameter('__unknown__keywords', inspect.Parameter.VAR_KEYWORD) ] signature = inspect.Signature(params) # This is a specialization to hint the first argument of certain builtins, # such as str.strip. if isinstance(func, _MethodDescriptorType): params = list(signature.parameters.values()) if params[0].annotation == params[0].empty: params[0] = params[0].replace(annotation=func.__objclass__) signature = signature.replace(parameters=params) # This is a specialization to hint the return value of type callables. if (signature.return_annotation == signature.empty and isinstance(func, type)): signature = signature.replace(return_annotation=typehints.normalize(func)) return signature def no_annotations(fn): """Decorator that prevents Beam from using type hint annotations on a callable.""" setattr(fn, '_beam_no_annotations', True) return fn def disable_type_annotations(): """Prevent Beam from using type hint annotations to determine input and output types of transforms. This setting applies globally. """ global _disable_from_callable _disable_from_callable = True TRACEBACK_LIMIT = 5 _NO_MAIN_TYPE = object() def _tag_and_type(t): """Extract tag name and value type from TaggedOutput[Literal['tag'], Type]. Returns raw Python types - conversion to beam types happens in _extract_output_types. """ args = get_args(t) if len(args) != 2: raise TypeError( f"TaggedOutput expects 2 type parameters, got {len(args)}: {t}") literal_type, value_type = args if get_origin(literal_type) is not Literal: raise TypeError( f"First type parameter of TaggedOutput must be Literal['tag_name'], " f"got {literal_type}. Example: TaggedOutput[Literal['errors'], str]") tag_string = get_args(literal_type)[0] return tag_string, value_type def _extract_tagged_from_type(beam_type): """Extract tagged output types from a Beam type (post-convert_to_beam_type). Called after the Iterable wrapper has been removed. At this point, the type has already been through convert_to_beam_type, so unions are typehints.UnionConstraint (not typing.Union), but TaggedOutput[Literal['tag'], T] passes through unchanged as a typing generic alias. Returns: (clean_type, tagged_dict) where clean_type is the type without TaggedOutput members (or _NO_MAIN_TYPE if no main type), and tagged_dict maps tag names to their Beam types. """ # Single TaggedOutput[Literal['tag'], Type] if get_origin(beam_type) is TaggedOutput: tag, typ = _tag_and_type(beam_type) return _NO_MAIN_TYPE, {tag: convert_to_beam_type(typ)} # Bare TaggedOutput (unparameterized) if beam_type is TaggedOutput: logging.warning( "TaggedOutput in return type must include type parameters: " "TaggedOutput[Literal['tag_name'], ValueType]. " "Bare TaggedOutput will be ignored.") return _NO_MAIN_TYPE, {} if not isinstance(beam_type, typehints.UnionHint.UnionConstraint): return beam_type, {} # UnionConstraint containing TaggedOutput members main_types = [] tagged = {} for member in beam_type.union_types: if get_origin(member) is TaggedOutput: tag, typ = _tag_and_type(member) tagged[tag] = convert_to_beam_type(typ) elif member is TaggedOutput: logging.warning( "TaggedOutput in return type must include type parameters: " "TaggedOutput[Literal['tag_name'], ValueType]. " "Bare TaggedOutput will be ignored.") else: main_types.append(member) if not tagged and len(main_types) == len(beam_type.union_types): return beam_type, {} if not main_types: return _NO_MAIN_TYPE, tagged elif len(main_types) == 1: return main_types[0], tagged else: return typehints.Union[tuple(main_types)], tagged class IOTypeHints(NamedTuple): """Encapsulates all type hint information about a Beam construct. This should primarily be used via the WithTypeHints mixin class, though may also be attached to other objects (such as Python functions). Attributes: input_types: (tuple, dict) List of typing types, and an optional dictionary. May be None. The list and dict correspond to args and kwargs. output_types: (tuple, dict) List of typing types, and an optional dictionary (unused). Only the first element of the list is used. May be None. origin: (List[str]) Stack of tracebacks of method calls used to create this instance. """ input_types: Optional[Tuple[Tuple[Any, ...], Dict[str, Any]]] output_types: Optional[Tuple[Tuple[Any, ...], Dict[str, Any]]] origin: List[str] @classmethod def _make_origin( cls, bases: List['IOTypeHints'], tb: bool = True, msg: Iterable[str] = ()) -> List[str]: if msg: res = list(msg) else: res = [] if tb: # Omit this method and the IOTypeHints method that called it. num_frames_skip = 2 tbs = traceback.format_stack(limit=TRACEBACK_LIMIT + num_frames_skip)[:-num_frames_skip] # tb is a list of strings in the form of 'File ...\n[code]\n'. Split into # single lines and flatten. res += list( itertools.chain.from_iterable(s.strip().split('\n') for s in tbs)) bases = [base for base in bases if base.origin] if bases: res += ['', 'based on:'] for i, base in enumerate(bases): if i > 0: res += ['', 'and:'] res += [' ' + str(base)] res += [' ' + s for s in base.origin] return res @classmethod def empty(cls) -> 'IOTypeHints': """Construct a base IOTypeHints object with no hints.""" return IOTypeHints(None, None, []) @classmethod def from_callable(cls, fn: Callable) -> Optional['IOTypeHints']: """Construct an IOTypeHints object from a callable's signature. Supports Python 3 annotations. For partial annotations, sets unknown types to Any, _ANY_VAR_POSITIONAL, or _ANY_VAR_KEYWORD. Returns: A new IOTypeHints or None if no annotations found. """ if _disable_from_callable or getattr(fn, '_beam_no_annotations', False): return None signature = get_signature(fn) if (all(param.annotation == param.empty for param in signature.parameters.values()) and signature.return_annotation == signature.empty): return None input_args = [] input_kwargs = {} for param in signature.parameters.values(): if param.annotation == param.empty: if param.kind == param.VAR_POSITIONAL: input_args.append(_ANY_VAR_POSITIONAL) elif param.kind == param.VAR_KEYWORD: input_kwargs[param.name] = _ANY_VAR_KEYWORD elif param.kind == param.KEYWORD_ONLY: input_kwargs[param.name] = typehints.Any else: input_args.append(typehints.Any) else: if param.kind in [param.KEYWORD_ONLY, param.VAR_KEYWORD]: input_kwargs[param.name] = convert_to_beam_type(param.annotation) else: assert param.kind in [param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD, param.VAR_POSITIONAL], \ 'Unsupported Parameter kind: %s' % param.kind input_args.append(convert_to_beam_type(param.annotation)) output_args = [] if signature.return_annotation != signature.empty: output_args.append(convert_to_beam_type(signature.return_annotation)) else: output_args.append(typehints.Any) name = getattr(fn, '__name__', '') msg = ['from_callable(%s)' % name, ' signature: %s' % signature] if hasattr(fn, '__code__'): msg.append( ' File "%s", line %d' % (fn.__code__.co_filename, fn.__code__.co_firstlineno)) return IOTypeHints( input_types=(tuple(input_args), input_kwargs), output_types=(tuple(output_args), {}), origin=cls._make_origin([], tb=False, msg=msg)) def with_input_types(self, *args, **kwargs) -> 'IOTypeHints': return self._replace( input_types=(args, kwargs), origin=self._make_origin([self])) def with_output_types(self, *args, **kwargs) -> 'IOTypeHints': return self._replace( output_types=(args, kwargs), origin=self._make_origin([self])) def with_input_types_from(self, other: 'IOTypeHints') -> 'IOTypeHints': return self._replace( input_types=other.input_types, origin=self._make_origin([self])) def with_output_types_from(self, other: 'IOTypeHints') -> 'IOTypeHints': return self._replace( output_types=other.output_types, origin=self._make_origin([self])) def simple_output_type(self, context): if self._has_output_types(): args, _ = self.output_types # Note: kwargs may contain tagged output types, which are ignored here. # Use tagged_output_types() to access those. if len(args) != 1: raise TypeError( 'Expected single output type hint for %s but got: %s' % (context, self.output_types)) return args[0] def tagged_output_types(self): if not self._has_output_types(): return {} _, tagged_output_types = self.output_types return tagged_output_types def has_simple_output_type(self): """Whether there's a single positional output type.""" return (self.output_types and len(self.output_types[0]) == 1) def strip_pcoll(self): from apache_beam.pipeline import Pipeline from apache_beam.pvalue import PBegin from apache_beam.pvalue import PDone return self.strip_pcoll_helper(self.input_types, self._has_input_types, 'input_types', [Pipeline, PBegin], 'This input type hint will be ignored ' 'and not used for type-checking purposes. ' 'Typically, input type hints for a ' 'PTransform are single (or nested) types ' 'wrapped by a PCollection, or PBegin.', 'strip_pcoll_input()').\ strip_pcoll_helper(self.output_types, self.has_simple_output_type, 'output_types', [PDone, None], 'This output type hint will be ignored ' 'and not used for type-checking purposes. ' 'Typically, output type hints for a ' 'PTransform are single (or nested) types ' 'wrapped by a PCollection, PDone, or None.', 'strip_pcoll_output()') def strip_pcoll_helper( self, my_type: any, has_my_type: Callable[[], bool], my_key: str, special_containers: List[Union[ 'PBegin', # noqa: F821 'PDone', # noqa: F821 'PCollection']], # noqa: F821 error_str: str, source_str: str) -> 'IOTypeHints': from apache_beam.pvalue import PCollection if not has_my_type() or not my_type or len(my_type[0]) != 1: return self my_type = my_type[0][0] if isinstance(my_type, typehints.AnyTypeConstraint): return self special_containers += [PCollection] kwarg_dict = {} if (my_type not in special_containers and getattr(my_type, '__origin__', None) != PCollection): logging.warning(error_str + ' Got: %s instead.' % my_type) kwarg_dict[my_key] = None return self._replace( origin=self._make_origin([self], tb=False, msg=[source_str]), **kwarg_dict) if (getattr(my_type, '__args__', -1) in [-1, None] or len(my_type.__args__) == 0): # e.g. PCollection (or PBegin/PDone) kwarg_dict[my_key] = ((typehints.Any, ), {}) else: # e.g. PCollection[type] kwarg_dict[my_key] = ((convert_to_beam_type(my_type.__args__[0]), ), {}) return self._replace( origin=self._make_origin([self], tb=False, msg=[source_str]), **kwarg_dict) def strip_iterable(self) -> 'IOTypeHints': """Removes outer Iterable (or equivalent) from output type. Only affects instances with simple output types, otherwise is a no-op. Does not modify self. Designed to be used with type hints from callables of ParDo, FlatMap, DoFn. Output type may be Optional[T], in which case the result of stripping T is used as the output type. Output type may be None/NoneType, in which case nothing is done. Example: Generator[Tuple(int, int)] becomes Tuple(int, int) Returns: A copy of this instance with a possibly different output type. Raises: ValueError if output type is simple and not iterable. """ if self.output_types is None or not self.has_simple_output_type(): return self output_type = self.output_types[0][0] tagged_output_types = self.output_types[1] if output_type is None or isinstance(output_type, type(None)): return self # If output_type == Optional[T]: output_type = T. if isinstance(output_type, typehints.UnionConstraint): types = list(output_type.union_types) if len(types) == 2: try: types.remove(type(None)) output_type = types[0] except ValueError: pass if isinstance(output_type, typehints.TypeVariable): # We don't know what T yields, so we just assume Any. return self._replace( output_types=((typehints.Any, ), tagged_output_types), origin=self._make_origin([self], tb=False, msg=['strip_iterable()'])) yielded_type = typehints.get_yielded_type(output_type) # Also strip Iterable from tagged output types (e.g. from Map/MapTuple # which wrap both main and tagged types in Iterable). stripped_tags = { tag: typehints.get_yielded_type(hint) for tag, hint in tagged_output_types.items() } return self._replace( output_types=((yielded_type, ), stripped_tags), origin=self._make_origin([self], tb=False, msg=['strip_iterable()'])) def extract_tagged_outputs(self): """Extract TaggedOutput types from the main output type into kwargs. For annotation style (e.g. -> Iterable[int | TaggedOutput[...]]), TaggedOutput stays embedded in the main type through convert_to_beam_type and strip_iterable. This method extracts those TaggedOutput members into the tagged output kwargs dict. Should be called after strip_iterable(). Returns: A copy of this instance with TaggedOutput members moved from the main output type into the output kwargs dict. """ if self.output_types is None or not self.has_simple_output_type(): return self output_type = self.output_types[0][0] clean_type, extracted_tags = _extract_tagged_from_type(output_type) if not extracted_tags: return self if clean_type is _NO_MAIN_TYPE: clean_type = typehints.Any return self._replace( output_types=((clean_type, ), extracted_tags), origin=self._make_origin([self], tb=False, msg=['extract_tagged_outputs()'])) def with_defaults(self, hints: Optional['IOTypeHints']) -> 'IOTypeHints': if not hints: return self if not self: return hints if self._has_input_types(): input_types = self.input_types else: input_types = hints.input_types if self._has_output_types(): output_types = self.output_types else: output_types = hints.output_types res = IOTypeHints( input_types, output_types, self._make_origin([self, hints], tb=False, msg=['with_defaults()'])) if res == self: return self # Don't needlessly increase origin traceback length. else: return res def _has_input_types(self): return self.input_types is not None and any(self.input_types) def _has_output_types(self): return self.output_types is not None and any(self.output_types) def __bool__(self): return self._has_input_types() or self._has_output_types() def __repr__(self): return 'IOTypeHints[inputs=%s, outputs=%s]' % ( self.input_types, self.output_types) def debug_str(self): return '\n'.join([self.__repr__()] + self.origin) def __eq__(self, other): def same(a, b): if a is None or not any(a): return b is None or not any(b) else: return a == b return ( same(self.input_types, other.input_types) and same(self.output_types, other.output_types)) def __hash__(self): return hash(str(self)) def __reduce__(self): # Don't include "origin" debug information in pickled form. return (IOTypeHints, (self.input_types, self.output_types, [])) class WithTypeHints(object): """A mixin class that provides the ability to set and retrieve type hints. """ def __init__(self, *unused_args, **unused_kwargs): self._type_hints = IOTypeHints.empty() def _get_or_create_type_hints(self) -> IOTypeHints: # __init__ may have not been called try: # Only return an instance bound to self (see BEAM-8629). return self.__dict__['_type_hints'] except KeyError: self._type_hints = IOTypeHints.empty() return self._type_hints def get_type_hints(self): """Gets and/or initializes type hints for this object. If type hints have not been set, attempts to initialize type hints in this order: - Using self.default_type_hints(). - Using self.__class__ type hints. """ return ( self._get_or_create_type_hints().with_defaults( self.default_type_hints()).with_defaults( get_type_hints(self.__class__))) def _set_type_hints(self, type_hints: IOTypeHints) -> None: self._type_hints = type_hints def default_type_hints(self): return None def with_input_types( self: WithTypeHintsT, *arg_hints: Any, **kwarg_hints: Any) -> WithTypeHintsT: arg_hints = native_type_compatibility.convert_to_beam_types(arg_hints) kwarg_hints = native_type_compatibility.convert_to_beam_types(kwarg_hints) self._type_hints = self._get_or_create_type_hints().with_input_types( *arg_hints, **kwarg_hints) return self def with_output_types( self: WithTypeHintsT, *arg_hints: Any, **kwarg_hints: Any) -> WithTypeHintsT: arg_hints = native_type_compatibility.convert_to_beam_types(arg_hints) kwarg_hints = native_type_compatibility.convert_to_beam_types(kwarg_hints) self._type_hints = self._get_or_create_type_hints().with_output_types( *arg_hints, **kwarg_hints) return self class TypeCheckError(Exception): pass def _positional_arg_hints(arg, hints): """Returns the type of a (possibly tuple-packed) positional argument. E.g. for lambda ((a, b), c): None the single positional argument is (as returned by inspect) [[a, b], c] which should have type Tuple[Tuple[Int, Any], float] when applied to the type hints {a: int, b: Any, c: float}. """ if isinstance(arg, list): return typehints.Tuple[[_positional_arg_hints(a, hints) for a in arg]] return hints.get(arg, typehints.Any) def _unpack_positional_arg_hints(arg, hint): """Unpacks the given hint according to the nested structure of arg. For example, if arg is [[a, b], c] and hint is Tuple[Any, int], then this function would return ((Any, Any), int) so it can be used in conjunction with inspect.getcallargs. """ if isinstance(arg, list): tuple_constraint = typehints.Tuple[[typehints.Any] * len(arg)] if not typehints.is_consistent_with(hint, tuple_constraint): raise TypeCheckError( 'Bad tuple arguments for %s: expected %s, got %s' % (arg, tuple_constraint, hint)) if isinstance(hint, typehints.TupleConstraint): return tuple( _unpack_positional_arg_hints(a, t) for a, t in zip(arg, hint.tuple_types)) return (typehints.Any, ) * len(arg) return hint def _normalize_var_positional_hint(hint): """Converts a var_positional hint into Tuple[Union[], ...] form. Args: hint: (tuple) Should be either a tuple of one or more types, or a single Tuple[, ...]. Raises: TypeCheckError if hint does not have the right form. """ if not hint or type(hint) != tuple: raise TypeCheckError('Unexpected VAR_POSITIONAL value: %s' % hint) if len(hint) == 1 and isinstance(hint[0], typehints.TupleSequenceConstraint): # Example: tuple(Tuple[Any, ...]) -> Tuple[Any, ...] return hint[0] else: # Example: tuple(int, str) -> Tuple[Union[int, str], ...] return typehints.Tuple[typehints.Union[hint], ...] def _normalize_var_keyword_hint(hint, arg_name): """Converts a var_keyword hint into Dict[, ] form. Args: hint: (dict) Should either contain a pair (arg_name, Dict[, ]), or one or more possible types for the value. arg_name: (str) The keyword receiving this hint. Raises: TypeCheckError if hint does not have the right form. """ if not hint or type(hint) != dict: raise TypeCheckError('Unexpected VAR_KEYWORD value: %s' % hint) keys = list(hint.keys()) values = list(hint.values()) if (len(values) == 1 and keys[0] == arg_name and isinstance(values[0], typehints.DictConstraint)): # Example: dict(kwargs=Dict[str, Any]) -> Dict[str, Any] return values[0] else: # Example: dict(k1=str, k2=int) -> Dict[str, Union[str,int]] return typehints.Dict[str, typehints.Union[values]] def getcallargs_forhints(func, *type_args, **type_kwargs): """Bind type_args and type_kwargs to func. Works like inspect.getcallargs, with some modifications to support type hint checks. For unbound args, will use annotations and fall back to Any (or variants of Any). Returns: A mapping from parameter name to argument. """ try: signature = get_signature(func) except ValueError as e: logging.warning('Could not get signature for function: %s: %s', func, e) return {} try: bindings = signature.bind(*type_args, **type_kwargs) except TypeError as e: # Might be raised due to too few or too many arguments. raise TypeCheckError(e) bound_args = bindings.arguments for param in signature.parameters.values(): if param.name in bound_args: # Bound: unpack/convert variadic arguments. if param.kind == param.VAR_POSITIONAL: bound_args[param.name] = _normalize_var_positional_hint( bound_args[param.name]) elif param.kind == param.VAR_KEYWORD: bound_args[param.name] = _normalize_var_keyword_hint( bound_args[param.name], param.name) else: # Unbound: must have a default or be variadic. if param.annotation != param.empty: bound_args[param.name] = param.annotation elif param.kind == param.VAR_POSITIONAL: bound_args[param.name] = _ANY_VAR_POSITIONAL elif param.kind == param.VAR_KEYWORD: bound_args[param.name] = _ANY_VAR_KEYWORD elif param.default is not param.empty: # Declare unbound parameters with defaults to be Any. bound_args[param.name] = typehints.Any else: # This case should be caught by signature.bind() above. raise ValueError('Unexpected unbound parameter: %s' % param.name) return dict(bound_args) def get_type_hints(fn: Any) -> IOTypeHints: """Gets the type hint associated with an arbitrary object fn. Always returns a valid IOTypeHints object, creating one if necessary. """ # pylint: disable=protected-access if not hasattr(fn, '_type_hints'): try: fn._type_hints = IOTypeHints.empty() except (AttributeError, TypeError): # Can't add arbitrary attributes to this object, # but might have some restrictions anyways... hints = IOTypeHints.empty() return hints return fn._type_hints # pylint: enable=protected-access def with_input_types(*positional_hints: Any, **keyword_hints: Any) -> Callable[[T], T]: """A decorator that type-checks defined type-hints with passed func arguments. All type-hinted arguments can be specified using positional arguments, keyword arguments, or a mix of both. Additionaly, all function arguments must be type-hinted in totality if even one parameter is type-hinted. Once fully decorated, if the arguments passed to the resulting function violate the type-hint constraints defined, a :class:`TypeCheckError` detailing the error will be raised. To be used as: .. testcode:: from apache_beam.typehints import with_input_types @with_input_types(str) def upper(s): return s.upper() Or: .. testcode:: from apache_beam.typehints import with_input_types from apache_beam.typehints import List from apache_beam.typehints import Tuple @with_input_types(ls=List[Tuple[int, int]]) def increment(ls): [(i + 1, j + 1) for (i,j) in ls] Args: *positional_hints: Positional type-hints having identical order as the function's formal arguments. Values for this argument must either be a built-in Python type or an instance of a :class:`~apache_beam.typehints.typehints.TypeConstraint` created by 'indexing' a :class:`~apache_beam.typehints.typehints.CompositeTypeHint` instance with a type parameter. **keyword_hints: Keyword arguments mirroring the names of the parameters to the decorated functions. The value of each keyword argument must either be one of the allowed built-in Python types, a custom class, or an instance of a :class:`~apache_beam.typehints.typehints.TypeConstraint` created by 'indexing' a :class:`~apache_beam.typehints.typehints.CompositeTypeHint` instance with a type parameter. Raises: :class:`ValueError`: If not all function arguments have corresponding type-hints specified. Or if the inner wrapper function isn't passed a function object. :class:`TypeCheckError`: If the any of the passed type-hint constraints are not a type or :class:`~apache_beam.typehints.typehints.TypeConstraint` instance. Returns: The original function decorated such that it enforces type-hint constraints for all received function arguments. """ converted_positional_hints = ( native_type_compatibility.convert_to_beam_types(positional_hints)) converted_keyword_hints = ( native_type_compatibility.convert_to_beam_types(keyword_hints)) del positional_hints del keyword_hints def annotate_input_types(f): if isinstance(f, types.FunctionType): for t in (list(converted_positional_hints) + list(converted_keyword_hints.values())): validate_composite_type_param( t, error_msg_prefix='All type hint arguments') th = getattr(f, '_type_hints', IOTypeHints.empty()).with_input_types( *converted_positional_hints, **converted_keyword_hints) f._type_hints = th # pylint: disable=protected-access return f return annotate_input_types def with_output_types(*return_type_hint: Any, **tagged_type_hints: Any) -> Callable[[T], T]: """A decorator that type-checks defined type-hints for return values(s). This decorator will type-check the return value(s) of the decorated function. Only a single type-hint is accepted to specify the return type of the return value. If the function to be decorated has multiple return values, then one should use: ``Tuple[type_1, type_2]`` to annotate the types of the return values. If the ultimate return value for the function violates the specified type-hint a :class:`TypeCheckError` will be raised detailing the type-constraint violation. This decorator is intended to be used like: .. testcode:: from apache_beam.typehints import with_output_types from apache_beam.typehints import Set class Coordinate(object): def __init__(self, x, y): self.x = x self.y = y @with_output_types(Set[Coordinate]) def parse_ints(ints): return {Coordinate(i, i) for i in ints} Or with a simple type-hint: .. testcode:: from apache_beam.typehints import with_output_types @with_output_types(bool) def negate(p): return not p if p else p For DoFns with tagged outputs, you can specify type hints for each tag: .. testcode:: from apache_beam.typehints import with_input_types, with_output_types @with_output_types(int, errors=str, warnings=str) class MyDoFn(beam.DoFn): def process(self, element): if element < 0: yield beam.pvalue.TaggedOutput('errors', 'Negative value') elif element == 0: yield beam.pvalue.TaggedOutput('warnings', 'Zero value') else: yield element Args: *return_type_hint: A type-hint specifying the proper return type of the function. This argument should either be a built-in Python type or an instance of a :class:`~apache_beam.typehints.typehints.TypeConstraint` created by 'indexing' a :class:`~apache_beam.typehints.typehints.CompositeTypeHint`. **tagged_type_hints: Type hints for tagged outputs. Each keyword argument specifies the type for a tagged output, e.g., ``errors=str``. Raises: :class:`ValueError`: If the length of **return_type_hint** is greater than ``1``. Or if the inner wrapper function isn't passed a function object. :class:`TypeCheckError`: If the **return_type_hint** object is in invalid type-hint. Returns: The original function decorated such that it enforces type-hint constraints for all return values. """ if len(return_type_hint) != 1: raise ValueError( "'returns' accepts only a single positional argument. In " "order to specify multiple return types, use the 'Tuple' " "type-hint.") return_type_hint = native_type_compatibility.convert_to_beam_type( return_type_hint[0]) validate_composite_type_param( return_type_hint, error_msg_prefix='All type hint arguments') converted_tag_hints = {} for tag, hint in tagged_type_hints.items(): converted_hint = native_type_compatibility.convert_to_beam_type(hint) validate_composite_type_param( converted_hint, 'Tagged output type hint for %r' % tag) converted_tag_hints[tag] = converted_hint def annotate_output_types(f): th = getattr(f, '_type_hints', IOTypeHints.empty()) f._type_hints = th.with_output_types( # pylint: disable=protected-access return_type_hint, **converted_tag_hints) return f return annotate_output_types def _check_instance_type( type_constraint, instance, var_name=None, verbose=False): """A helper function to report type-hint constraint violations. Args: type_constraint: An instance of a 'TypeConstraint' or a built-in Python type. instance: The candidate object which will be checked by to satisfy 'type_constraint'. var_name: If 'instance' is an argument, then the actual name for the parameter in the original function definition. Raises: TypeCheckError: If 'instance' fails to meet the type-constraint of 'type_constraint'. """ hint_type = ( "argument: '%s'" % var_name if var_name is not None else 'return type') try: check_constraint(type_constraint, instance) except SimpleTypeHintError: if verbose: verbose_instance = '%s, ' % instance else: verbose_instance = '' raise TypeCheckError( 'Type-hint for %s violated. Expected an ' 'instance of %s, instead found %san instance of %s.' % (hint_type, type_constraint, verbose_instance, type(instance))) except CompositeTypeHintError as e: raise TypeCheckError('Type-hint for %s violated: %s' % (hint_type, e)) def _interleave_type_check(type_constraint, var_name=None): """Lazily type-check the type-hint for a lazily generated sequence type. This function can be applied as a decorator or called manually in a curried manner: * @_interleave_type_check(List[int]) def gen(): yield 5 or * gen = _interleave_type_check(Tuple[int, int], 'coord_gen')(gen) As a result, all type-checking for the passed generator will occur at 'yield' time. This way, we avoid having to depleat the generator in order to type-check it. Args: type_constraint: An instance of a TypeConstraint. The output yielded of 'gen' will be type-checked according to this type constraint. var_name: The variable name binded to 'gen' if type-checking a function argument. Used solely for templating in error message generation. Returns: A function which takes a generator as an argument and returns a wrapped version of the generator that interleaves type-checking at 'yield' iteration. If the generator received is already wrapped, then it is simply returned to avoid nested wrapping. """ def wrapper(gen): if isinstance(gen, GeneratorWrapper): return gen return GeneratorWrapper( gen, lambda x: _check_instance_type(type_constraint, x, var_name)) return wrapper class GeneratorWrapper(object): """A wrapper around a generator, allows execution of a callback per yield. Additionally, wrapping a generator with this class allows one to assign arbitary attributes to a generator object just as with a function object. Attributes: internal_gen: A instance of a generator object. As part of 'step' of the generator, the yielded object will be passed to 'interleave_func'. interleave_func: A callback accepting a single argument. This function will be called with the result of each yielded 'step' in the internal generator. """ def __init__(self, gen, interleave_func): self.internal_gen = gen self.interleave_func = interleave_func def __getattr__(self, attr): # TODO(laolu): May also want to intercept 'send' in the future if we move to # a GeneratorHint with 3 type-params: # * Generator[send_type, return_type, yield_type] if attr == '__next__': return self.__next__() elif attr == '__iter__': return self.__iter__() return getattr(self.internal_gen, attr) def __next__(self): next_val = next(self.internal_gen) self.interleave_func(next_val) return next_val next = __next__ def __iter__(self): for x in self.internal_gen: self.interleave_func(x) yield x ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/decorators_test.py0000644000175100017510000004721115206113033023470 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for decorators module.""" # pytype: skip-file # ruff: noqa: UP006 import functools import typing import unittest from apache_beam import Map from apache_beam.pvalue import TaggedOutput from apache_beam.typehints import Any from apache_beam.typehints import Dict from apache_beam.typehints import List from apache_beam.typehints import Tuple from apache_beam.typehints import TypeCheckError from apache_beam.typehints import TypeVariable from apache_beam.typehints import WithTypeHints from apache_beam.typehints import decorators from apache_beam.typehints import typehints from apache_beam.typehints.native_type_compatibility import convert_to_beam_type T = TypeVariable('T') # Name is 'T' so it converts to a beam type with the same name. # mypy requires that the name of the variable match, so we must ignore this. # pylint: disable=typevar-name-mismatch T_typing = typing.TypeVar('T') # type: ignore class IOTypeHintsTest(unittest.TestCase): def test_get_signature(self): # Basic coverage only to make sure function works. def fn(a, b=1, *c, **d): return a, b, c, d s = decorators.get_signature(fn) self.assertListEqual(list(s.parameters), ['a', 'b', 'c', 'd']) def test_get_signature_builtin(self): s = decorators.get_signature(list) self.assertListEqual(list(s.parameters), ['iterable']) self.assertEqual(s.return_annotation, List[Any]) def test_from_callable_without_annotations(self): def fn(a, b=None, *args, **kwargs): return a, b, args, kwargs th = decorators.IOTypeHints.from_callable(fn) self.assertIsNone(th) def test_from_callable_builtin(self): th = decorators.IOTypeHints.from_callable(len) self.assertIsNone(th) def test_from_callable_method_descriptor(self): # from_callable() injects an annotation in this special type of builtin. th = decorators.IOTypeHints.from_callable(str.strip) self.assertEqual(th.input_types, ((str, Any), {})) self.assertEqual(th.output_types, ((Any, ), {})) def test_strip_iterable_not_simple_output_noop(self): th = decorators.IOTypeHints( input_types=None, output_types=((int, str), {}), origin=[]) th = th.strip_iterable() self.assertEqual(((int, str), {}), th.output_types) def _test_strip_iterable(self, before, expected_after): th = decorators.IOTypeHints( input_types=None, output_types=((before, ), {}), origin=[]) after = th.strip_iterable() self.assertEqual(((expected_after, ), {}), after.output_types) def test_with_output_types_from(self): th = decorators.IOTypeHints( input_types=((int), { 'foo': str }), output_types=((int, str), {}), origin=[]) self.assertEqual( th.with_output_types_from(decorators.IOTypeHints.empty()), decorators.IOTypeHints( input_types=((int), { 'foo': str }), output_types=None, origin=[])) self.assertEqual( decorators.IOTypeHints.empty().with_output_types_from(th), decorators.IOTypeHints( input_types=None, output_types=((int, str), {}), origin=[])) def test_with_input_types_from(self): th = decorators.IOTypeHints( input_types=((int), { 'foo': str }), output_types=((int, str), {}), origin=[]) self.assertEqual( th.with_input_types_from(decorators.IOTypeHints.empty()), decorators.IOTypeHints( input_types=None, output_types=((int, str), {}), origin=[])) self.assertEqual( decorators.IOTypeHints.empty().with_input_types_from(th), decorators.IOTypeHints( input_types=((int), { 'foo': str }), output_types=None, origin=[])) def _test_strip_iterable_fail(self, before): with self.assertRaisesRegex(ValueError, r'not iterable'): self._test_strip_iterable(before, None) def test_strip_iterable(self): self._test_strip_iterable(None, None) self._test_strip_iterable(typehints.Any, typehints.Any) self._test_strip_iterable(T, typehints.Any) self._test_strip_iterable(typehints.Iterable[str], str) self._test_strip_iterable(typehints.List[str], str) self._test_strip_iterable(typehints.Iterator[str], str) self._test_strip_iterable(typehints.Generator[str], str) self._test_strip_iterable(typehints.Tuple[str], str) self._test_strip_iterable( typehints.Tuple[str, int], typehints.Union[str, int]) self._test_strip_iterable(typehints.Tuple[str, ...], str) self._test_strip_iterable(typehints.KV[str, int], typehints.Union[str, int]) self._test_strip_iterable(typehints.Set[str], str) self._test_strip_iterable(typehints.FrozenSet[str], str) self._test_strip_iterable_fail(typehints.Union[str, int]) self._test_strip_iterable_fail(typehints.Optional[str]) self._test_strip_iterable_fail( typehints.WindowedValue[str]) # type: ignore[misc] self._test_strip_iterable_fail(typehints.Dict[str, int]) def test_make_traceback(self): origin = ''.join( decorators.IOTypeHints.empty().with_input_types(str).origin) self.assertRegex(origin, __name__) self.assertNotRegex(origin, r'\b_make_traceback') def test_origin(self): th = decorators.IOTypeHints.empty() self.assertListEqual([], th.origin) th = th.with_input_types(str) self.assertRegex(th.debug_str(), r'with_input_types') th = th.with_output_types(str) self.assertRegex(th.debug_str(), r'(?s)with_output_types.*with_input_types') th = decorators.IOTypeHints.empty().with_output_types(str) th2 = decorators.IOTypeHints.empty().with_input_types(int) th = th.with_defaults(th2) self.assertRegex(th.debug_str(), r'(?s)based on:.*\'str\'.*and:.*\'int\'') def test_with_defaults_noop_does_not_grow_origin(self): th = decorators.IOTypeHints.empty() expected_id = id(th) th = th.with_defaults(None) self.assertEqual(expected_id, id(th)) th = th.with_defaults(decorators.IOTypeHints.empty()) self.assertEqual(expected_id, id(th)) th = th.with_input_types(str) expected_id = id(th) th = th.with_defaults(th) self.assertEqual(expected_id, id(th)) th2 = th.with_output_types(int) th = th.with_defaults(th2) self.assertNotEqual(expected_id, id(th)) def test_from_callable(self): def fn( a: int, b: str = '', *args: Tuple[T], foo: List[int], **kwargs: Dict[str, str]) -> Tuple[Any, ...]: return a, b, args, foo, kwargs th = decorators.IOTypeHints.from_callable(fn) self.assertEqual( th.input_types, ((int, str, Tuple[T]), { 'foo': List[int], 'kwargs': Dict[str, str] })) self.assertEqual(th.output_types, ((Tuple[Any, ...], ), {})) def test_from_callable_partial_annotations(self): def fn(a: int, b=None, *args, foo: List[int], **kwargs): return a, b, args, foo, kwargs th = decorators.IOTypeHints.from_callable(fn) self.assertEqual( th.input_types, ((int, Any, Tuple[Any, ...]), { 'foo': List[int], 'kwargs': Dict[Any, Any] })) self.assertEqual(th.output_types, ((Any, ), {})) def test_from_callable_class(self): class Class(object): def __init__(self, unused_arg: int): pass th = decorators.IOTypeHints.from_callable(Class) self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((Class, ), {})) def test_from_callable_method(self): class Class(object): def method(self, arg: T = None) -> None: pass th = decorators.IOTypeHints.from_callable(Class.method) self.assertEqual(th.input_types, ((Any, T), {})) self.assertEqual(th.output_types, ((None, ), {})) th = decorators.IOTypeHints.from_callable(Class().method) self.assertEqual(th.input_types, ((T, ), {})) self.assertEqual(th.output_types, ((None, ), {})) def test_from_callable_convert_to_beam_types(self): def fn( a: typing.List[int], b: str = '', *args: typing.Tuple[T_typing], foo: typing.List[int], **kwargs: typing.Dict[str, str]) -> typing.Tuple[typing.Any, ...]: return a, b, args, foo, kwargs th = decorators.IOTypeHints.from_callable(fn) self.assertEqual( th.input_types, ((List[int], str, Tuple[T]), { 'foo': List[int], 'kwargs': Dict[str, str] })) self.assertEqual(th.output_types, ((Tuple[Any, ...], ), {})) def test_from_callable_partial(self): def fn(a: int) -> int: return a # functools.partial objects don't have __name__ attributes by default. fn = functools.partial(fn, 1) th = decorators.IOTypeHints.from_callable(fn) self.assertRegex(th.debug_str(), r'unknown') def test_from_callable_no_tagged_output(self): def fn(x: int) -> str: return str(x) th = decorators.IOTypeHints.from_callable(fn) self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((str, ), {})) def fn2(x: int) -> typing.Iterable[str]: yield str(x) th = decorators.IOTypeHints.from_callable(fn2) self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((typehints.Iterable[str], ), {})) def test_from_callable_tagged_output_union(self): """Tagged types are NOT extracted in from_callable. They stay embedded in the main type and are extracted later in strip_iterable().""" def fn( x: int ) -> int | str | TaggedOutput[typing.Literal['errors'], float | str] | TaggedOutput[ typing.Literal['warnings'], str]: return x th = decorators.IOTypeHints.from_callable(fn) self.assertEqual(th.input_types, ((int, ), {})) # TaggedOutput members are preserved in the union no extraction yet. output_type = th.output_types[0][0] self.assertIsInstance(output_type, typehints.UnionConstraint) self.assertEqual(th.output_types[1], {}) def test_from_callable_tagged_output_iterable(self): """Tagged types inside Iterable are preserved until strip_iterable.""" def fn( x: int ) -> typing.Iterable[int | TaggedOutput[typing.Literal['errors'], str]]: yield x th = decorators.IOTypeHints.from_callable(fn) self.assertEqual(th.input_types, ((int, ), {})) # The full Iterable[Union[int, TaggedOutput[...]]] is preserved. output_type = th.output_types[0][0] self.assertIsInstance(output_type, typehints.IterableTypeConstraint) self.assertEqual(th.output_types[1], {}) def test_from_callable_tagged_output_only(self): """A standalone TaggedOutput annotation passes through from_callable.""" def fn(x: int) -> TaggedOutput[typing.Literal['errors'], str]: pass th = decorators.IOTypeHints.from_callable(fn) self.assertEqual(th.input_types, ((int, ), {})) # TaggedOutput[...] passes through convert_to_beam_type unchanged. self.assertIs(typing.get_origin(th.output_types[0][0]), TaggedOutput) self.assertEqual(th.output_types[1], {}) def test_getcallargs_forhints(self): def fn( a: int, b: str = '', *args: Tuple[T], foo: List[int], **kwargs: Dict[str, str]) -> Tuple[Any, ...]: return a, b, args, foo, kwargs callargs = decorators.getcallargs_forhints(fn, float, foo=List[str]) self.assertDictEqual( callargs, { 'a': float, 'b': str, 'args': Tuple[T], 'foo': List[str], 'kwargs': Dict[str, str] }) def test_getcallargs_forhints_default_arg(self): # Default args are not necessarily types, so they should be ignored. def fn(a=List[int], b=None, *args, foo=(), **kwargs) -> Tuple[Any, ...]: return a, b, args, foo, kwargs callargs = decorators.getcallargs_forhints(fn) self.assertDictEqual( callargs, { 'a': Any, 'b': Any, 'args': Tuple[Any, ...], 'foo': Any, 'kwargs': Dict[Any, Any] }) def test_getcallargs_forhints_missing_arg(self): def fn(a, b=None, *args, foo, **kwargs): return a, b, args, foo, kwargs with self.assertRaisesRegex(decorators.TypeCheckError, "missing.*'a'"): decorators.getcallargs_forhints(fn, foo=List[int]) with self.assertRaisesRegex(decorators.TypeCheckError, "missing.*'foo'"): decorators.getcallargs_forhints(fn, 5) def test_origin_annotated(self): def annotated(e: str) -> str: return e t = Map(annotated) th = t.get_type_hints() th = th.with_input_types(str) self.assertRegex(th.debug_str(), r'with_input_types') th = th.with_output_types(str) self.assertRegex( th.debug_str(), r'(?s)with_output_types.*with_input_types.*Map.annotated') class WithTypeHintsTest(unittest.TestCase): def test_get_type_hints_no_settings(self): class Base(WithTypeHints): pass th = Base().get_type_hints() self.assertEqual(th.input_types, None) self.assertEqual(th.output_types, None) def test_get_type_hints_class_decorators(self): @decorators.with_input_types(int, str) @decorators.with_output_types(int) class Base(WithTypeHints): pass th = Base().get_type_hints() self.assertEqual(th.input_types, ((int, str), {})) self.assertEqual(th.output_types, ((int, ), {})) def test_get_type_hints_class_defaults(self): class Base(WithTypeHints): def default_type_hints(self): return decorators.IOTypeHints( input_types=((int, str), {}), output_types=((int, ), {}), origin=[]) th = Base().get_type_hints() self.assertEqual(th.input_types, ((int, str), {})) self.assertEqual(th.output_types, ((int, ), {})) def test_get_type_hints_precedence_defaults_over_decorators(self): @decorators.with_input_types(int) @decorators.with_output_types(str) class Base(WithTypeHints): def default_type_hints(self): return decorators.IOTypeHints( input_types=((float, ), {}), output_types=None, origin=[]) th = Base().get_type_hints() self.assertEqual(th.input_types, ((float, ), {})) self.assertEqual(th.output_types, ((str, ), {})) def test_get_type_hints_precedence_instance_over_defaults(self): class Base(WithTypeHints): def default_type_hints(self): return decorators.IOTypeHints( input_types=((float, ), {}), output_types=((str, ), {}), origin=[]) th = Base().with_input_types(int).get_type_hints() self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((str, ), {})) def test_inherits_does_not_modify(self): # See BEAM-8629. @decorators.with_output_types(int) class Subclass(WithTypeHints): def __init__(self): pass # intentionally avoiding super call # These should be equal, but not the same object lest mutating the instance # mutates the class. self.assertIsNot( Subclass()._get_or_create_type_hints(), Subclass._type_hints) self.assertEqual(Subclass().get_type_hints(), Subclass._type_hints) self.assertNotEqual( Subclass().with_input_types(str)._type_hints, Subclass._type_hints) class DecoratorsTest(unittest.TestCase): def tearDown(self): decorators._disable_from_callable = False def test_disable_type_annotations(self): self.assertFalse(decorators._disable_from_callable) decorators.disable_type_annotations() self.assertTrue(decorators._disable_from_callable) def test_no_annotations_on_same_function(self): def fn(a: int) -> int: return a with self.assertRaisesRegex(TypeCheckError, r'requires .*int.* but was applied .*str'): _ = ['a', 'b', 'c'] | Map(fn) # Same pipeline doesn't raise without annotations on fn. fn = decorators.no_annotations(fn) _ = ['a', 'b', 'c'] | Map(fn) def test_no_annotations_on_diff_function(self): def fn(a: int) -> int: return a _ = [1, 2, 3] | Map(fn) # Doesn't raise - correct types. with self.assertRaisesRegex(TypeCheckError, r'requires .*int.* but was applied .*str'): _ = ['a', 'b', 'c'] | Map(fn) @decorators.no_annotations def fn2(a: int) -> int: return a _ = ['a', 'b', 'c'] | Map(fn2) # Doesn't raise - no input type hints. class ExtractTaggedFromTypeTest(unittest.TestCase): """Tests for _extract_tagged_from_type (Beam-level type extraction).""" def test_simple_type_no_extraction(self): main, tagged = decorators._extract_tagged_from_type(int) self.assertEqual(main, int) self.assertEqual(tagged, {}) def test_beam_union_no_tagged(self): t = typehints.Union[int, str] main, tagged = decorators._extract_tagged_from_type(t) self.assertEqual(main, t) self.assertEqual(tagged, {}) def test_standalone_tagged_output(self): t = TaggedOutput[typing.Literal['errors'], str] main, tagged = decorators._extract_tagged_from_type(t) self.assertIs(main, decorators._NO_MAIN_TYPE) self.assertEqual(tagged, {'errors': str}) def test_beam_union_with_tagged(self): t = convert_to_beam_type(int | TaggedOutput[typing.Literal['errors'], str]) main, tagged = decorators._extract_tagged_from_type(t) self.assertEqual(main, int) self.assertEqual(tagged, {'errors': str}) def test_beam_union_multiple_tagged(self): t = convert_to_beam_type( int | TaggedOutput[typing.Literal['errors'], str] | TaggedOutput[typing.Literal['warnings'], str]) main, tagged = decorators._extract_tagged_from_type(t) self.assertEqual(main, int) self.assertEqual(tagged, {'errors': str, 'warnings': str}) def test_beam_union_multiple_main_types(self): t = convert_to_beam_type( int | str | TaggedOutput[typing.Literal['errors'], bytes]) main, tagged = decorators._extract_tagged_from_type(t) self.assertIsInstance(main, typehints.UnionConstraint) self.assertIn(int, main.union_types) self.assertIn(str, main.union_types) self.assertEqual(tagged, {'errors': bytes}) def test_beam_union_tagged_only(self): t = convert_to_beam_type( TaggedOutput[typing.Literal['errors'], str] | TaggedOutput[typing.Literal['warnings'], int]) main, tagged = decorators._extract_tagged_from_type(t) self.assertIs(main, decorators._NO_MAIN_TYPE) self.assertEqual(tagged, {'errors': str, 'warnings': int}) def test_bare_tagged_output_standalone(self): with self.assertLogs(level='WARNING') as cm: main, tagged = decorators._extract_tagged_from_type(TaggedOutput) self.assertIn('Bare TaggedOutput will be ignored', cm.output[0]) self.assertIs(main, decorators._NO_MAIN_TYPE) self.assertEqual(tagged, {}) def test_bare_tagged_output_in_union(self): with self.assertLogs(level='WARNING') as cm: t = convert_to_beam_type(str | TaggedOutput) main, tagged = decorators._extract_tagged_from_type(t) self.assertIn('Bare TaggedOutput will be ignored', cm.output[0]) self.assertEqual(main, str) self.assertEqual(tagged, {}) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/intrinsic_one_ops.py0000644000175100017510000000465715206113033024017 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Defines the actions intrinsic bytecodes have on the frame. Each function here corresponds to a bytecode documented in https://docs.python.org/3/library/dis.html . The first argument is a (mutable) FrameState object, the second the integer opcode argument. Bytecodes with more complicated behavior (e.g. modifying the program counter) are handled inline rather than here. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file from . import opcodes def intrinsic_1_invalid(state, arg): pass def intrinsic_print(state, arg): pass def intrinsic_import_star(state, arg): pass def intrinsic_stopiteration_error(state, arg): pass def intrinsic_async_gen_wrap(state, arg): pass def intrinsic_unary_positive(state, arg): opcodes.unary_positive(state, arg) pass def intrinsic_list_to_tuple(state, arg): opcodes.list_to_tuple(state, arg) pass def intrinsic_typevar(state, arg): pass def intrinsic_paramspec(state, arg): pass def intrinsic_typevartuple(state, arg): pass def intrinsic_subscript_generic(state, arg): pass def intrinsic_typealias(state, arg): pass # The order of operations in the table of the intrinsic one operations is # defined in https://docs.python.org/3/library/dis.html#opcode-CALL_INTRINSIC_1 # and may change between minor versions. INT_ONE_OPS = tuple([ intrinsic_1_invalid, intrinsic_print, intrinsic_import_star, intrinsic_stopiteration_error, intrinsic_async_gen_wrap, intrinsic_unary_positive, intrinsic_list_to_tuple, intrinsic_typevar, intrinsic_paramspec, intrinsic_typevartuple, intrinsic_subscript_generic, intrinsic_typealias ]) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/intrinsic_one_ops_test.py0000644000175100017510000000247415206113033025051 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.typehints.intrinsic_one_ops.""" # pytype: skip-file import dis import sys import unittest from apache_beam.typehints import intrinsic_one_ops class IntrinsicOneOpsTest(unittest.TestCase): def test_unary_intrinsic_ops_are_in_the_same_order_as_in_cpython(self): if sys.version_info >= (3, 12): dis_order = dis.__dict__['_intrinsic_1_descs'] beam_ops = [fn.__name__.upper() for fn in intrinsic_one_ops.INT_ONE_OPS] self.assertListEqual(dis_order, beam_ops) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/native_type_compatibility.py0000644000175100017510000005345015206113033025546 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module to convert Python's native typing types to Beam types.""" # pytype: skip-file # ruff: noqa: UP006 import collections import collections.abc import dataclasses import logging import sys import types import typing from typing import Generic from typing import TypeVar from apache_beam.typehints import typehints try: from typing import is_typeddict except ImportError: from typing_extensions import is_typeddict # Python 3.12 adds TypeAliasType for `type` statements; keep optional import. # pylint: disable=ungrouped-imports # isort: off try: from typing import TypeAliasType # type: ignore[attr-defined] except Exception: # pragma: no cover - pre-3.12 TypeAliasType = None # type: ignore[assignment] T = TypeVar('T') _LOGGER = logging.getLogger(__name__) # Describes an entry in the type map in convert_to_beam_type. # match is a function that takes a user type and returns whether the conversion # should trigger. # arity is the expected arity of the user type. -1 means it's variadic. # beam_type is the Beam type the user type should map to. _TypeMapEntry = collections.namedtuple( '_TypeMapEntry', ['match', 'arity', 'beam_type']) _BUILTINS_TO_TYPING = { dict: typing.Dict, list: typing.List, tuple: typing.Tuple, set: typing.Set, frozenset: typing.FrozenSet, } _BUILTINS = [ dict, list, tuple, set, frozenset, ] _CONVERTED_COLLECTIONS = [ collections.abc.Iterable, collections.abc.Iterator, collections.abc.Generator, collections.abc.Set, collections.abc.MutableSet, collections.abc.Collection, collections.abc.Sequence, collections.abc.Mapping, ] _CONVERTED_MODULES = ( 'typing', 'collections', 'collections.abc', 'annotationlib', ) def _get_args(typ): """Returns a list of arguments to the given type. Args: typ: A typing module typing type. Returns: A tuple of args. """ try: if typ.__args__ is None or not isinstance(typ.__args__, tuple): return () return typ.__args__ except AttributeError: if isinstance(typ, typing.TypeVar): return (typ.__name__, ) return () def _safe_issubclass(derived, parent): """Like issubclass, but swallows TypeErrors. This is useful for when either parameter might not actually be a class, e.g. typing.Union isn't actually a class. Args: derived: As in issubclass. parent: As in issubclass. Returns: issubclass(derived, parent), or False if a TypeError was raised. """ try: return issubclass(derived, parent) except (TypeError, AttributeError): if hasattr(derived, '__origin__'): try: return issubclass(derived.__origin__, parent) except TypeError: pass return False def _match_issubclass(match_against): return lambda user_type: _safe_issubclass(user_type, match_against) def _is_primitive(user_type, primitive): # catch bare primitives if user_type is primitive: return True return getattr(user_type, '__origin__', None) is primitive def _match_is_primitive(match_against): return lambda user_type: _is_primitive(user_type, match_against) def _match_is_dict(user_type): return _is_primitive(user_type, dict) or _safe_issubclass(user_type, dict) def _match_is_exactly_mapping(user_type): # Avoid unintentionally catching all subtypes (e.g. strings and mappings). expected_origin = collections.abc.Mapping return getattr(user_type, '__origin__', None) is expected_origin def _match_is_exactly_iterable(user_type): if user_type is typing.Iterable or user_type is collections.abc.Iterable: return True # Avoid unintentionally catching all subtypes (e.g. strings and mappings). expected_origin = collections.abc.Iterable return getattr(user_type, '__origin__', None) is expected_origin def _match_is_exactly_collection(user_type): return getattr(user_type, '__origin__', None) is collections.abc.Collection def _match_is_exactly_sequence(user_type): return getattr(user_type, '__origin__', None) is collections.abc.Sequence def match_is_named_tuple(user_type): return ( _safe_issubclass(user_type, typing.Tuple) and hasattr(user_type, '__annotations__') and hasattr(user_type, '_fields')) def match_dataclass_for_row(user_type): """Match whether the type is a dataclass handled by row coder. For frozen dataclasses, only true when explicitly registered with row coder: beam.coders.typecoders.registry.register_coder( MyDataClass, beam.coders.RowCoder) (for backward-compatibility reason). For non-frozen dataclasses, default to true otherwise explicitly registered with a coder other than the row coder. """ if not dataclasses.is_dataclass(user_type): return False # pylint: disable=wrong-import-position try: from apache_beam.options.pipeline_options_context import get_pipeline_options # pylint: disable=line-too-long except AttributeError: pass else: opts = get_pipeline_options() if opts and opts.is_compat_version_prior_to("2.73.0"): return False is_frozen = user_type.__dataclass_params__.frozen # avoid circular import try: from apache_beam.coders.typecoders import registry as coders_registry from apache_beam.coders import RowCoder except AttributeError: # coder registery not yet initialized so it must be absent return not is_frozen if is_frozen: return ( user_type in coders_registry._coders and coders_registry._coders[user_type] == RowCoder) else: return ( user_type not in coders_registry._coders or coders_registry._coders[user_type] == RowCoder) def _match_is_optional(user_type): return _match_is_union(user_type) and sum( tp is type(None) for tp in _get_args(user_type)) == 1 def extract_optional_type(user_type): """Extracts the non-None type from Optional type user_type. If user_type is not Optional, returns None """ if not _match_is_optional(user_type): return None else: return next(tp for tp in _get_args(user_type) if tp is not type(None)) def _match_is_union(user_type): # For non-subscripted unions (Python 2.7.14+ with typing 3.64) if user_type is typing.Union: return True try: # Python 3.5.4+, or Python 2.7.14+ with typing 3.64 return user_type.__origin__ is typing.Union except AttributeError: pass return False def _match_is_set(user_type): if _safe_issubclass(user_type, typing.Set) or _is_primitive(user_type, set): return True elif getattr(user_type, '__origin__', None) is not None: return _safe_issubclass( user_type.__origin__, collections.abc.Set) or _safe_issubclass( user_type.__origin__, collections.abc.MutableSet) else: return False def is_any(typ): return typ is typing.Any def is_new_type(typ): return hasattr(typ, '__supertype__') _ForwardRef = typing.ForwardRef # Python 3.7+ def is_forward_ref(typ): return isinstance(typ, _ForwardRef) # Mapping from typing.TypeVar/typehints.TypeVariable ids to an object of the # other type. Bidirectional mapping preserves typing.TypeVar instances. _type_var_cache: typing.Dict[int, typehints.TypeVariable] = {} def convert_builtin_to_typing(typ): """Convert recursively a given builtin to a typing object. Args: typ (`builtins`): builtin object that exist in _BUILTINS_TO_TYPING. Returns: type: The given builtins converted to a type. """ if getattr(typ, '__origin__', None) in _BUILTINS_TO_TYPING: args = map(convert_builtin_to_typing, typ.__args__) typ = _BUILTINS_TO_TYPING[typ.__origin__].copy_with(tuple(args)) return typ def convert_typing_to_builtin(typ): """Converts a given typing collections type to its builtin counterpart. Args: typ: A typing type (e.g., typing.List[int]). Returns: type: The corresponding builtin type (e.g., list[int]). """ origin = getattr(typ, '__origin__', None) args = getattr(typ, '__args__', None) # Typing types return the primitive type as the origin from 3.9 on if origin not in _BUILTINS: return typ # Early return for bare types if not args: return origin if origin is list: return list[convert_typing_to_builtin(args[0])] elif origin is dict: return dict[convert_typing_to_builtin(args[0]), convert_typing_to_builtin(args[1])] elif origin is tuple: return tuple[tuple(convert_typing_to_builtin(args))] elif origin is set: return set[convert_typing_to_builtin(args)] elif origin is frozenset: return frozenset[convert_typing_to_builtin(args)] def convert_collections_to_typing(typ): """Converts a given collections.abc type to a typing object. Args: typ: an object inheriting from a collections.abc object Returns: type: The corresponding typing object. """ if hasattr(typ, '__iter__'): if hasattr(typ, '__next__'): typ = typing.Iterator[typ.__args__] elif hasattr(typ, 'send') and hasattr(typ, 'throw'): typ = typing.Generator[typ.__args__] elif _match_is_exactly_iterable(typ): typ = typing.Iterable[typ.__args__] return typ def is_builtin(typ): if typ in _BUILTINS: return True return getattr(typ, '__origin__', None) in _BUILTINS # During type inference of WindowedValue, we need to pass in the inner value # type. This cannot be achieved immediately with WindowedValue class because it # is not parameterized. Changing it to a generic class (e.g. WindowedValue[T]) # could work in theory. However, the class is cythonized and it seems that # cython does not handle generic classes well. # The workaround here is to create a separate class solely for the type # inference purpose. This class should never be used for creating instances. class TypedWindowedValue(Generic[T]): def __init__(self, *args, **kwargs): raise NotImplementedError("This class is solely for type inference") def convert_to_beam_type(typ): """Convert a given typing type to a Beam type. Args: typ (`typing.Union[type, str]`): typing type or string literal representing a type. Returns: type: The given type converted to a Beam type as far as we can do the conversion. Raises: ValueError: The type was malformed. """ # Convert `int | float` to typing.Union[int, float] # pipe operator as Union and types.UnionType are introduced # in Python 3.10. # GH issue: https://github.com/apache/beam/issues/21972 if isinstance(typ, types.UnionType): typ = typing.Union[typ] # Unwrap Python 3.12 `type` aliases (TypeAliasType) to their underlying value. # This ensures Beam sees the actual aliased type (e.g., tuple[int, ...]). if sys.version_info >= (3, 12) and TypeAliasType is not None: if isinstance(typ, TypeAliasType): # pylint: disable=isinstance-second-argument-not-valid-type underlying = getattr(typ, '__value__', None) if underlying is not None: typ = underlying if getattr(typ, '__module__', None) == 'typing': typ = convert_typing_to_builtin(typ) typ_module = getattr(typ, '__module__', None) if isinstance(typ, typing.TypeVar): # This is a special case, as it's not parameterized by types. # Also, identity must be preserved through conversion (i.e. the same # TypeVar instance must get converted into the same TypeVariable instance). # A global cache should be OK as the number of distinct type variables # is generally small. if id(typ) not in _type_var_cache: new_type_variable = typehints.TypeVariable(typ.__name__) _type_var_cache[id(typ)] = new_type_variable _type_var_cache[id(new_type_variable)] = typ return _type_var_cache[id(typ)] elif isinstance(typ, str): # Special case for forward references. # TODO(https://github.com/apache/beam/issues/19954): Currently unhandled. _LOGGER.info('Converting string literal type hint to Any: "%s"', typ) return typehints.Any elif isinstance(typ, typing.NewType): # pylint: disable=isinstance-second-argument-not-valid-type # Special case for NewType, where, since Python 3.10, NewType is now a class # rather than a function. # TODO(https://github.com/apache/beam/issues/20076): Currently unhandled. _LOGGER.info('Converting NewType type hint to Any: "%s"', typ) return typehints.Any elif typ_module == 'apache_beam.typehints.native_type_compatibility' and \ getattr(typ, "__name__", typ.__origin__.__name__) == 'TypedWindowedValue': # Need to pass through WindowedValue class so that it can be converted # to the correct type constraint in Beam # This is needed to fix https://github.com/apache/beam/issues/33356 pass elif is_typeddict(typ): # Special-case for the TypedDict constructor, which is not actually a type, # and therefore fails to be recognised as compatible with Dict or Mapping. return typehints.Dict[str, typehints.Any] elif typ_module not in _CONVERTED_MODULES and not is_builtin(typ): # Only translate primitives and types from collections.abc and typing. return typ if (typ_module == 'collections.abc' and getattr(typ, '__origin__', typ) not in _CONVERTED_COLLECTIONS): # TODO(https://github.com/apache/beam/issues/29135): # Support more collections types return typ type_map = [ # TODO(https://github.com/apache/beam/issues/20076): Currently # unsupported. _TypeMapEntry(match=is_new_type, arity=0, beam_type=typehints.Any), # TODO(https://github.com/apache/beam/issues/19954): Currently # unsupported. _TypeMapEntry(match=is_forward_ref, arity=0, beam_type=typehints.Any), _TypeMapEntry(match=is_any, arity=0, beam_type=typehints.Any), _TypeMapEntry(match=_match_is_dict, arity=2, beam_type=typehints.Dict), _TypeMapEntry( match=_match_is_exactly_iterable, arity=1, beam_type=typehints.Iterable), _TypeMapEntry( match=_match_is_primitive(list), arity=1, beam_type=typehints.List), # FrozenSets are a specific instance of a set, so we check this first. _TypeMapEntry( match=_match_is_primitive(frozenset), arity=1, beam_type=typehints.FrozenSet), _TypeMapEntry(match=_match_is_set, arity=1, beam_type=typehints.Set), # NamedTuple is a subclass of Tuple, but it needs special handling. # We just convert it to Any for now. # This MUST appear before the entry for the normal Tuple. _TypeMapEntry( match=match_is_named_tuple, arity=0, beam_type=typehints.Any), _TypeMapEntry( match=_match_is_primitive(tuple), arity=-1, beam_type=typehints.Tuple), _TypeMapEntry(match=_match_is_union, arity=-1, beam_type=typehints.Union), _TypeMapEntry( match=_match_issubclass(collections.abc.Generator), arity=3, beam_type=typehints.Generator), _TypeMapEntry( match=_match_issubclass(collections.abc.Iterator), arity=1, beam_type=typehints.Iterator), _TypeMapEntry( match=_match_is_exactly_collection, arity=1, beam_type=typehints.Collection), _TypeMapEntry( match=_match_issubclass(TypedWindowedValue), arity=1, beam_type=typehints.WindowedValue), _TypeMapEntry( match=_match_is_exactly_sequence, arity=1, beam_type=typehints.Sequence), _TypeMapEntry( match=_match_is_exactly_mapping, arity=2, beam_type=typehints.Mapping), ] # Find the first matching entry. matched_entry = next((entry for entry in type_map if entry.match(typ)), None) if not matched_entry: # Please add missing type support if you see this message. _LOGGER.info('Using Any for unsupported type: %s', typ) return typehints.Any args = _get_args(typ) len_args = len(args) if len_args == 0 and len_args != matched_entry.arity: arity = matched_entry.arity # Handle unsubscripted types. if _match_issubclass(typing.Tuple)(typ): args = (typehints.TypeVariable('T'), Ellipsis) elif _match_is_union(typ): raise ValueError('Unsupported Union with no arguments.') elif _match_issubclass(typing.Generator)(typ): # Assume a simple generator. args = (typehints.TypeVariable('T_co'), type(None), type(None)) elif _match_issubclass(typing.Dict)(typ): args = (typehints.TypeVariable('KT'), typehints.TypeVariable('VT')) elif (_match_issubclass(typing.Iterator)(typ) or _match_is_exactly_iterable(typ)): args = (typehints.TypeVariable('T_co'), ) else: args = (typehints.TypeVariable('T'), ) * arity elif matched_entry.arity == -1: arity = len_args # Counters are special dict types that are implicitly parameterized to # [T, int], so we fix cases where they only have one argument to match # a more traditional dict hint. elif len_args == 1 and _safe_issubclass(getattr(typ, '__origin__', typ), collections.Counter): args = (args[0], int) len_args = 2 arity = matched_entry.arity else: arity = matched_entry.arity if len_args != arity: raise ValueError( 'expecting type %s to have arity %d, had arity %d ' 'instead' % (str(typ), arity, len_args)) typs = convert_to_beam_types(args) if arity == 0: # Nullary types (e.g. Any) don't accept empty tuples as arguments. return matched_entry.beam_type elif arity == 1: # Unary types (e.g. Set) don't accept 1-tuples as arguments return matched_entry.beam_type[typs[0]] else: return matched_entry.beam_type[tuple(typs)] def convert_to_beam_types(args): """Convert the given list or dictionary of args to Beam types. Args: args: Either an iterable of types, or a dictionary where the values are types. Returns: If given an iterable, a list of converted types. If given a dictionary, a dictionary with the same keys, and values which have been converted. """ if isinstance(args, dict): return {k: convert_to_beam_type(v) for k, v in args.items()} else: return [convert_to_beam_type(v) for v in args] def convert_to_python_type(typ): """Converts a given Beam type to a python type. This is the reverse of convert_to_beam_type. Args: typ: If a typehints.TypeConstraint, the type to convert. Otherwise, typ will be unchanged. Returns: Converted version of typ, or unchanged. Raises: ValueError: The type was malformed or could not be converted. """ if isinstance(typ, typehints.TypeVariable): # This is a special case, as it's not parameterized by types. # Also, identity must be preserved through conversion (i.e. the same # TypeVariable instance must get converted into the same TypeVar instance). # A global cache should be OK as the number of distinct type variables # is generally small. if id(typ) not in _type_var_cache: new_type_variable = typing.TypeVar(typ.name) _type_var_cache[id(typ)] = new_type_variable _type_var_cache[id(new_type_variable)] = typ return _type_var_cache[id(typ)] elif not getattr(typ, '__module__', None).endswith('typehints'): # Only translate types from the typehints module. return typ if isinstance(typ, typehints.AnyTypeConstraint): return typing.Any if isinstance(typ, typehints.DictConstraint): return dict[convert_to_python_type(typ.key_type), convert_to_python_type(typ.value_type)] if isinstance(typ, typehints.ListConstraint): return list[convert_to_python_type(typ.inner_type)] if isinstance(typ, typehints.IterableTypeConstraint): return collections.abc.Iterable[convert_to_python_type(typ.inner_type)] if isinstance(typ, typehints.UnionConstraint): if not typ.union_types: # Gracefully handle the empty union type. return typing.Any return typing.Union[tuple(convert_to_python_types(typ.union_types))] if isinstance(typ, typehints.SetTypeConstraint): return set[convert_to_python_type(typ.inner_type)] if isinstance(typ, typehints.FrozenSetTypeConstraint): return frozenset[convert_to_python_type(typ.inner_type)] if isinstance(typ, typehints.TupleConstraint): return tuple[tuple(convert_to_python_types(typ.tuple_types))] if isinstance(typ, typehints.TupleSequenceConstraint): return tuple[convert_to_python_type(typ.inner_type), ...] if isinstance(typ, typehints.ABCSequenceTypeConstraint): return collections.abc.Sequence[convert_to_python_type(typ.inner_type)] if isinstance(typ, typehints.IteratorTypeConstraint): return collections.abc.Iterator[convert_to_python_type(typ.yielded_type)] if isinstance(typ, typehints.MappingTypeConstraint): return collections.abc.Mapping[convert_to_python_type(typ.key_type), convert_to_python_type(typ.value_type)] raise ValueError('Failed to convert Beam type: %s' % typ) def convert_to_python_types(args): """Convert the given list or dictionary of args to python types. Args: args: Either an iterable of types, or a dictionary where the values are types. Returns: If given an iterable, a list of converted types. If given a dictionary, a dictionary with the same keys, and values which have been converted. """ if isinstance(args, dict): return {k: convert_to_python_type(v) for k, v in args.items()} else: return [convert_to_python_type(v) for v in args] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/native_type_compatibility_test.py0000644000175100017510000005304615206113033026606 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for Beam type compatibility library.""" # pytype: skip-file # ruff: noqa: UP006 import collections.abc import dataclasses import enum import re import typing import unittest from parameterized import param from parameterized import parameterized from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options_context import scoped_pipeline_options from apache_beam.typehints import typehints from apache_beam.typehints.native_type_compatibility import convert_builtin_to_typing from apache_beam.typehints.native_type_compatibility import convert_to_beam_type from apache_beam.typehints.native_type_compatibility import convert_to_beam_types from apache_beam.typehints.native_type_compatibility import convert_to_python_type from apache_beam.typehints.native_type_compatibility import convert_to_python_types from apache_beam.typehints.native_type_compatibility import convert_typing_to_builtin from apache_beam.typehints.native_type_compatibility import is_any from apache_beam.typehints.native_type_compatibility import match_dataclass_for_row _TestNamedTuple = typing.NamedTuple( '_TestNamedTuple', [('age', int), ('name', bytes)]) _TestFlatAlias = tuple[bytes, float] _TestNestedAlias = list[_TestFlatAlias] _TestFlatAliasTyping = typing.Tuple[bytes, float] _TestNestedAliasTyping = typing.List[_TestFlatAliasTyping] class _TestClass(object): pass T = typing.TypeVar('T') U = typing.TypeVar('U') class _TestGeneric(typing.Generic[T]): pass class _TestPair(typing.NamedTuple('TestTuple', [('first', T), ('second', T)]), typing.Generic[T]): pass class _TestEnum(enum.Enum): FOO = enum.auto() BAR = enum.auto() class _TestTypedDict(typing.TypedDict): foo: int bar: str class NativeTypeCompatibilityTest(unittest.TestCase): def test_convert_to_beam_type(self): test_cases = [ ('raw bytes', bytes, bytes), ('raw int', int, int), ('raw float', float, float), ('any', typing.Any, typehints.Any), ('simple dict', dict[bytes, int], typehints.Dict[bytes, int]), ('simple list', list[int], typehints.List[int]), ('simple iterable', collections.abc.Iterable[int], typehints.Iterable[int]), ('simple optional', typing.Optional[int], typehints.Optional[int]), ('simple set', set[float], typehints.Set[float]), ('simple frozenset', frozenset[float], typehints.FrozenSet[float]), ('simple unary tuple', tuple[bytes], typehints.Tuple[bytes]), ('simple union', typing.Union[int, bytes, float], typehints.Union[int, bytes, float]), ('namedtuple', _TestNamedTuple, _TestNamedTuple), ('test class', _TestClass, _TestClass), ('test class in list', list[_TestClass], typehints.List[_TestClass]), ('generic bare', _TestGeneric, _TestGeneric), ('generic subscripted', _TestGeneric[int], _TestGeneric[int]), ('complex tuple', tuple[bytes, list[tuple[ bytes, typing.Union[int, bytes, float]]]], typehints.Tuple[bytes, typehints.List[typehints.Tuple[ bytes, typehints.Union[int, bytes, float]]]]), ('arbitrary-length tuple', tuple[int, ...], typehints.Tuple[int, ...]), ('flat alias', _TestFlatAlias, typehints.Tuple[bytes, float]), # type: ignore[misc] ('nested alias', _TestNestedAlias, typehints.List[typehints.Tuple[bytes, float]]), ('complex dict', dict[bytes, list[tuple[bytes, _TestClass]]], typehints.Dict[bytes, typehints.List[typehints.Tuple[ bytes, _TestClass]]]), ('type var', typing.TypeVar('T'), typehints.TypeVariable('T')), ('nested type var', tuple[typing.TypeVar('K'), typing.TypeVar('V')], typehints.Tuple[typehints.TypeVariable('K'), typehints.TypeVariable('V')]), ('iterator', collections.abc.Iterator[typing.Any], typehints.Iterator[typehints.Any]), ('nested generic bare', list[_TestGeneric], typehints.List[_TestGeneric]), ('nested generic subscripted', list[_TestGeneric[int]], typehints.List[_TestGeneric[int]]), ('nested generic with any', list[_TestPair[typing.Any]], typehints.List[_TestPair[typing.Any]]), ('raw enum', _TestEnum, _TestEnum), ] for test_case in test_cases: if test_case is None: continue # Unlike typing types, Beam types are guaranteed to compare equal. description = test_case[0] typing_type = test_case[1] expected_beam_type = test_case[2] converted_beam_type = convert_to_beam_type(typing_type) self.assertEqual(converted_beam_type, expected_beam_type, description) converted_typing_type = convert_to_python_type(converted_beam_type) self.assertEqual(converted_typing_type, typing_type, description) def test_convert_to_beam_type_with_typing_types(self): test_cases = [ ('simple dict', typing.Dict[bytes, int], typehints.Dict[bytes, int]), ('simple list', typing.List[int], typehints.List[int]), ('simple iterable', typing.Iterable[int], typehints.Iterable[int]), ('simple optional', typing.Optional[int], typehints.Optional[int]), ('simple set', typing.Set[float], typehints.Set[float]), ('simple frozenset', typing.FrozenSet[float], typehints.FrozenSet[float]), ('simple unary tuple', typing.Tuple[bytes], typehints.Tuple[bytes]), ('test class in list', typing.List[_TestClass], typehints.List[_TestClass]), ('complex tuple', typing.Tuple[bytes, typing.List[typing.Tuple[ bytes, typing.Union[int, bytes, float]]]], typehints.Tuple[bytes, typehints.List[typehints.Tuple[ bytes, typehints.Union[int, bytes, float]]]]), ('arbitrary-length tuple', typing.Tuple[int, ...], typehints.Tuple[int, ...]), ('flat alias', _TestFlatAliasTyping, typehints.Tuple[bytes, float]), # type: ignore[misc] ('nested alias', _TestNestedAliasTyping, typehints.List[typehints.Tuple[bytes, float]]), ('complex dict', typing.Dict[bytes, typing.List[typing.Tuple[bytes, _TestClass]]], typehints.Dict[bytes, typehints.List[typehints.Tuple[ bytes, _TestClass]]]), ('nested type var', typing.Tuple[typing.TypeVar('K'), typing.TypeVar('V')], typehints.Tuple[typehints.TypeVariable('K'), typehints.TypeVariable('V')]), ('iterator', typing.Iterator[typing.Any], typehints.Iterator[typehints.Any]), ('nested generic bare', typing.List[_TestGeneric], typehints.List[_TestGeneric]), ('nested generic subscripted', typing.List[_TestGeneric[int]], typehints.List[_TestGeneric[int]]), ('nested generic with any', typing.List[_TestPair[typing.Any]], typehints.List[_TestPair[typing.Any]]), ] for test_case in test_cases: description = test_case[0] builtins_type = test_case[1] expected_beam_type = test_case[2] converted_beam_type = convert_to_beam_type(builtins_type) self.assertEqual(converted_beam_type, expected_beam_type, description) def test_convert_to_beam_type_with_builtin_types(self): test_cases = [ ('builtin dict', dict[str, int], typehints.Dict[str, int]), ('builtin list', list[str], typehints.List[str]), ('builtin tuple', tuple[str], typehints.Tuple[str]), ('builtin set', set[str], typehints.Set[str]), ('builtin frozenset', frozenset[int], typehints.FrozenSet[int]), ( 'nested builtin', dict[str, list[tuple[float]]], typehints.Dict[str, typehints.List[typehints.Tuple[float]]]), ( 'builtin nested tuple', tuple[str, list], typehints.Tuple[str, typehints.List[typehints.TypeVariable('T')]], ) ] for test_case in test_cases: description = test_case[0] builtins_type = test_case[1] expected_beam_type = test_case[2] converted_beam_type = convert_to_beam_type(builtins_type) self.assertEqual(converted_beam_type, expected_beam_type, description) def test_convert_to_beam_type_with_collections_types(self): test_cases = [ ( 'collection iterable', collections.abc.Iterable[int], typehints.Iterable[int]), ( 'collection generator', collections.abc.Generator[int, None, None], typehints.Generator[int]), ( 'collection iterator', collections.abc.Iterator[int], typehints.Iterator[int]), ( 'nested iterable', tuple[bytes, collections.abc.Iterable[int]], typehints.Tuple[bytes, typehints.Iterable[int]]), ( 'iterable over tuple', collections.abc.Iterable[tuple[str, int]], typehints.Iterable[typehints.Tuple[str, int]]), ( 'mapping', collections.abc.Mapping[str, int], typehints.Mapping[str, int]), ('set', collections.abc.Set[int], typehints.Set[int]), ('mutable set', collections.abc.MutableSet[int], typehints.Set[int]), ( 'enum mutable set', collections.abc.MutableSet[_TestEnum], typehints.Set[_TestEnum]), ( 'collection enum', collections.abc.Collection[_TestEnum], typehints.Collection[_TestEnum]), ( 'collection of tuples', collections.abc.Collection[tuple[str, int]], typehints.Collection[typehints.Tuple[str, int]]), ( 'nested sequence', tuple[collections.abc.Sequence[str], int], typehints.Tuple[typehints.Sequence[str], int]), ( 'sequence of tuples', collections.abc.Sequence[tuple[str, int]], typehints.Sequence[typehints.Tuple[str, int]]), ( 'ordered dict', collections.OrderedDict[str, int], typehints.Dict[str, int]), ( 'default dict', collections.defaultdict[str, int], typehints.Dict[str, int]), ('typed dict', _TestTypedDict, typehints.Dict[str, typehints.Any]), ('count', collections.Counter[str, int], typehints.Dict[str, int]), ( 'single param counter', collections.Counter[str], typehints.Dict[str, int]), ( 'bare callable', collections.abc.Callable, collections.abc.Callable, ), ( 'parameterized callable', collections.abc.Callable[[str], int], collections.abc.Callable[[str], int], ), ] for test_case in test_cases: description = test_case[0] builtins_type = test_case[1] expected_beam_type = test_case[2] converted_beam_type = convert_to_beam_type(builtins_type) self.assertEqual(converted_beam_type, expected_beam_type, description) def test_convert_builtin_to_typing(self): test_cases = [ ('dict', dict[str, int], typing.Dict[str, int]), ('list', list[str], typing.List[str]), ('tuple', tuple[str], typing.Tuple[str]), ('set', set[str], typing.Set[str]), ( 'nested', dict[str, list[tuple[float]]], typing.Dict[str, typing.List[typing.Tuple[float]]]), ] for test_case in test_cases: description = test_case[0] builtin_type = test_case[1] expected_typing_type = test_case[2] converted_typing_type = convert_builtin_to_typing(builtin_type) self.assertEqual(converted_typing_type, expected_typing_type, description) def test_generator_converted_to_iterator(self): self.assertEqual( typehints.Iterator[int], convert_to_beam_type(typing.Generator[int, None, None])) def test_newtype(self): self.assertEqual( typehints.Any, convert_to_beam_type(typing.NewType('Number', int))) def test_pattern(self): self.assertEqual(re.Pattern, convert_to_beam_type(re.Pattern)) self.assertEqual(re.Pattern[str], convert_to_beam_type(re.Pattern[str])) self.assertEqual(re.Pattern[bytes], convert_to_beam_type(re.Pattern[bytes])) self.assertNotEqual( re.Pattern[bytes], convert_to_beam_type(re.Pattern[str])) def test_match(self): self.assertEqual(re.Match, convert_to_beam_type(re.Match)) self.assertEqual(re.Match[str], convert_to_beam_type(re.Match[str])) self.assertEqual(re.Match[bytes], convert_to_beam_type(re.Match[bytes])) self.assertNotEqual(re.Match[bytes], convert_to_beam_type(re.Match[str])) def test_forward_reference(self): self.assertEqual(typehints.Any, convert_to_beam_type('int')) self.assertEqual(typehints.Any, convert_to_beam_type('typing.List[int]')) self.assertEqual( typehints.List[typehints.Any], convert_to_beam_type(list['int'])) def test_annotationlib_forward_ref(self): beam_type = convert_to_beam_type(dict[str, typing.ForwardRef('foo.Bar')]) self.assertEqual(typehints.Dict[str, typehints.Any], beam_type) def test_convert_nested_to_beam_type(self): self.assertEqual(typehints.List[typing.Any], typehints.List[typehints.Any]) self.assertEqual( typehints.List[typing.Dict[int, str]], typehints.List[typehints.Dict[int, str]]) def test_convert_bare_types(self): # Conversions for unsubscripted types that have implicit subscripts. test_cases = [ ('bare list', typing.List, typehints.List[typehints.TypeVariable('T')]), ( 'bare dict', typing.Dict, typehints.Dict[typehints.TypeVariable('KT'), typehints.TypeVariable('VT')]), ( 'bare tuple', typing.Tuple, typehints.Tuple[typehints.TypeVariable('T'), ...]), ('bare set', typing.Set, typehints.Set[typehints.TypeVariable('T')]), ( 'bare frozenset', typing.FrozenSet, typehints.FrozenSet[typehints.TypeVariable( 'T', use_name_in_eq=False)]), ( 'bare iterator', typing.Iterator, typehints.Iterator[typehints.TypeVariable('T_co')]), ( 'bare iterable', typing.Iterable, typehints.Iterable[typehints.TypeVariable('T_co')]), ( 'nested bare', typing.Tuple[typing.Iterator], typehints.Tuple[typehints.Iterator[typehints.TypeVariable('T_co')]] ), ( 'bare generator', typing.Generator, typehints.Generator[typehints.TypeVariable('T_co')]), ] for test_case in test_cases: description = test_case[0] typing_type = test_case[1] expected_beam_type = test_case[2] converted_beam_type = convert_to_beam_type(typing_type) self.assertEqual(expected_beam_type, converted_beam_type, description) def test_convert_bare_collections_types(self): # Conversions for unsubscripted types that have implicit subscripts. test_cases = [ ('bare list', list, typehints.List[typehints.TypeVariable('T')]), ( 'bare dict', dict, typehints.Dict[typehints.TypeVariable('KT'), typehints.TypeVariable('VT')]), ( 'bare tuple', tuple, typehints.Tuple[typehints.TypeVariable('T'), ...]), ('bare set', typing.Set, typehints.Set[typehints.TypeVariable('T')]), ( 'bare frozenset', frozenset, typehints.FrozenSet[typehints.TypeVariable( 'T', use_name_in_eq=False)]), ( 'bare iterator', collections.abc.Iterator, typehints.Iterator[typehints.TypeVariable('T_co')]), ( 'bare iterable', collections.abc.Iterable, typehints.Iterable[typehints.TypeVariable('T_co')]), ( 'nested bare', tuple[collections.abc.Iterator], typehints.Tuple[typehints.Iterator[typehints.TypeVariable('T_co')]] ), ( 'bare generator', collections.abc.Generator, typehints.Generator[typehints.TypeVariable('T_co')]), ] for test_case in test_cases: description = test_case[0] typing_type = test_case[1] expected_beam_type = test_case[2] converted_beam_type = convert_to_beam_type(typing_type) self.assertEqual(expected_beam_type, converted_beam_type, description) def test_convert_bare_types_fail(self): # These conversions should fail. test_cases = [ ('bare union', typing.Union), ] for test_case in test_cases: description = test_case[0] typing_type = test_case[1] with self.assertRaises(ValueError, msg=description): convert_to_beam_type(typing_type) def test_convert_to_beam_types(self): typing_types = [ bytes, list[bytes], list[tuple[bytes, int]], typing.Union[int, list[int]] ] beam_types = [ bytes, typehints.List[bytes], typehints.List[typehints.Tuple[bytes, int]], typehints.Union[int, typehints.List[int]] ] converted_beam_types = convert_to_beam_types(typing_types) self.assertEqual(converted_beam_types, beam_types) converted_typing_types = convert_to_python_types(converted_beam_types) self.assertEqual(converted_typing_types, typing_types) def test_is_any(self): test_cases = [ (True, typing.Any), (False, typing.List[int]), (False, typing.Union), (False, 1), (False, 'a'), ] for expected, typ in test_cases: self.assertEqual(expected, is_any(typ), msg='%s' % typ) def test_convert_typing_to_builtin(self): test_cases = [ ('list', typing.List[int], list[int]), ('dict', typing.Dict[str, int], dict[str, int]), ('tuple', typing.Tuple[str, int], tuple[str, int]), ('set', typing.Set[str], set[str]), ('frozenset', typing.FrozenSet[int], frozenset[int]), ( 'nested', typing.List[typing.Dict[str, typing.Tuple[int]]], list[dict[str, tuple[int]]]), ('typevar', typing.List[T], list[T]), ('nested_typevar', typing.Dict[T, typing.List[U]], dict[T, list[U]]) ] for description, typing_type, expected_builtin_type in test_cases: builtin_type = convert_typing_to_builtin(typing_type) self.assertEqual(builtin_type, expected_builtin_type, description) def test_type_alias_type_unwrapped(self): # Only applicable on Python 3.12+, where typing.TypeAliasType exists # and the `type` statement is available. TypeAliasType = getattr(typing, 'TypeAliasType', None) if TypeAliasType is None: self.skipTest('TypeAliasType not available') ns = {} try: exec('type AliasTuple = tuple[int, ...]', {}, ns) # pylint: disable=exec-used except SyntaxError: self.skipTest('type statement not supported') AliasTuple = ns['AliasTuple'] self.assertTrue(isinstance(AliasTuple, TypeAliasType)) # pylint: disable=isinstance-second-argument-not-valid-type self.assertEqual( typehints.Tuple[int, ...], convert_to_beam_type(AliasTuple)) def test_dataclass_default(self): @dataclasses.dataclass(frozen=True) class FrozenDC: foo: int @dataclasses.dataclass class NonFrozenDC: foo: int self.assertFalse(match_dataclass_for_row(FrozenDC)) self.assertTrue(match_dataclass_for_row(NonFrozenDC)) def test_dataclass_registered(self): @dataclasses.dataclass(frozen=True) class FrozenRegisteredDC: foo: int @dataclasses.dataclass class NonFrozenRegisteredDC: foo: int # pylint: disable=wrong-import-position from apache_beam.coders import RowCoder from apache_beam.coders import typecoders from apache_beam.coders.coders import FastPrimitivesCoder typecoders.registry.register_coder(FrozenRegisteredDC, RowCoder) typecoders.registry.register_coder( NonFrozenRegisteredDC, FastPrimitivesCoder) self.assertTrue(match_dataclass_for_row(FrozenRegisteredDC)) self.assertFalse(match_dataclass_for_row(NonFrozenRegisteredDC)) @parameterized.expand([ param(compat_version="2.72.0"), param(compat_version="2.73.0"), ]) def test_dataclass_update_compatibility(self, compat_version): @dataclasses.dataclass(frozen=True) class FrozenDC: foo: int @dataclasses.dataclass class NonFrozenDC: foo: int with scoped_pipeline_options( PipelineOptions(update_compatibility_version=compat_version)): self.assertFalse(match_dataclass_for_row(FrozenDC)) self.assertEqual( compat_version == "2.73.0", match_dataclass_for_row(NonFrozenDC)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/opcodes.py0000644000175100017510000004721315206113033021722 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Defines the actions various bytecodes have on the frame. Each function here corresponds to a bytecode documented in https://docs.python.org/2/library/dis.html or https://docs.python.org/3/library/dis.html. The first argument is a (mutable) FrameState object, the second the integer opcode argument. Bytecodes with more complicated behavior (e.g. modifying the program counter) are handled inline rather than here. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file import inspect import logging import sys import types from functools import reduce from apache_beam.typehints import row_type from apache_beam.typehints import typehints from apache_beam.typehints.trivial_inference import BoundMethod from apache_beam.typehints.trivial_inference import Const from apache_beam.typehints.trivial_inference import element_type from apache_beam.typehints.trivial_inference import key_value_types from apache_beam.typehints.trivial_inference import union from apache_beam.typehints.typehints import Any from apache_beam.typehints.typehints import Dict from apache_beam.typehints.typehints import Iterable from apache_beam.typehints.typehints import List from apache_beam.typehints.typehints import Set from apache_beam.typehints.typehints import Tuple from apache_beam.typehints.typehints import Union # This is missing in the builtin types module. str.upper is arbitrary, any # method on a C-implemented type will do. _MethodDescriptorType = type(str.upper) if sys.version_info >= (3, 11): import opcode _div_binop_args = frozenset([ ix for (ix, (argname, _)) in enumerate(opcode._nb_ops) if 'TRUE_DIVIDE' in argname ]) else: _div_binop_args = frozenset() if sys.version_info >= (3, 14): _NB_SUBSCR_OPCODE = [op[0] for op in opcode._nb_ops].index('NB_SUBSCR') else: _NB_SUBSCR_OPCODE = -1 def pop_one(state, unused_arg): del state.stack[-1:] def pop_two(state, unused_arg): del state.stack[-2:] def pop_three(state, unused_arg): del state.stack[-3:] def push_value(v): def pusher(state, unused_arg): state.stack.append(v) return pusher def nop(unused_state, unused_arg): pass resume = nop def pop_top(state, unused_arg): state.stack.pop() def end_for(state, unused_arg): state.stack.pop() state.stack.pop() def end_send(state, unused_arg): del state.stack[-2] def copy(state, arg): state.stack.append(state.stack[-arg]) def rot_n(state, n): state.stack[-n:] = [state.stack[-1]] + state.stack[-n:-1] def rot_two(state, unused_arg): rot_n(state, 2) def rot_three(state, unused_arg): rot_n(state, 3) def rot_four(state, unused_arg): rot_n(state, 4) def dup_top(state, unused_arg): state.stack.append(state.stack[-1]) def unary(state, unused_arg): state.stack[-1] = Const.unwrap(state.stack[-1]) unary_positive = unary_negative = unary_invert = unary def unary_not(state, unused_arg): state.stack[-1] = bool def unary_convert(state, unused_arg): state.stack[-1] = str def get_iter(state, unused_arg): state.stack.append(Iterable[element_type(state.stack.pop())]) _NUMERIC_PROMOTION_LADDER = [bool, int, float, complex] def symmetric_binary_op(state, arg, is_true_div=None): # TODO(robertwb): This may not be entirely correct... # BINARY_SUBSCR was rolled into BINARY_OP in 3.14. if arg == _NB_SUBSCR_OPCODE: return binary_subscr(state, arg) b, a = Const.unwrap(state.stack.pop()), Const.unwrap(state.stack.pop()) if a == b: if a is int and b is int and (arg in _div_binop_args or is_true_div): state.stack.append(float) else: state.stack.append(a) elif type(a) == type(b) and isinstance(a, typehints.SequenceTypeConstraint): state.stack.append(type(a)(union(element_type(a), element_type(b)))) # Technically these next two will be errors for anything but multiplication, # but that's OK. elif a is int and (b in (bytes, str) or isinstance(b, typehints.SequenceTypeConstraint)): state.stack.append(b) elif b is int and (a in (bytes, str) or isinstance(a, typehints.SequenceTypeConstraint)): state.stack.append(a) elif a in _NUMERIC_PROMOTION_LADDER and b in _NUMERIC_PROMOTION_LADDER: state.stack.append( _NUMERIC_PROMOTION_LADDER[max( _NUMERIC_PROMOTION_LADDER.index(a), _NUMERIC_PROMOTION_LADDER.index(b))]) else: state.stack.append(Any) # Except for int ** -int binary_power = inplace_power = symmetric_binary_op binary_multiply = inplace_multiply = symmetric_binary_op binary_divide = inplace_divide = symmetric_binary_op binary_floor_divide = inplace_floor_divide = symmetric_binary_op def binary_true_divide(state, arg): return symmetric_binary_op(state, arg, True) inplace_true_divide = binary_true_divide binary_modulo = inplace_modulo = symmetric_binary_op # TODO(robertwb): Tuple add. binary_add = inplace_add = symmetric_binary_op binary_subtract = inplace_subtract = symmetric_binary_op def binary_subscr(state, unused_arg): index = state.stack.pop() base = Const.unwrap(state.stack.pop()) if base is str: out = base elif (isinstance(index, Const) and isinstance(index.value, int) and isinstance(base, typehints.IndexableTypeConstraint)): try: out = base._constraint_for_index(index.value) except IndexError: out = element_type(base) elif (index == slice or getattr(index, 'type', None) == slice) and isinstance( base, typehints.IndexableTypeConstraint): # The slice is treated as a const in 3.14, using this instead of # BINARY_SLICE out = base else: out = element_type(base) state.stack.append(out) # As far as types are concerned. binary_lshift = inplace_lshift = binary_rshift = inplace_rshift = pop_top binary_and = inplace_and = symmetric_binary_op binary_xor = inplace_xor = symmetric_binary_op binary_or = inplace_or = symmetric_binary_op binary_op = symmetric_binary_op def store_subscr(unused_state, unused_args): # TODO(robertwb): Update element/value type of iterable/dict. pass def binary_slice(state, args): _ = state.stack.pop() _ = state.stack.pop() base = Const.unwrap(state.stack.pop()) if base is str: out = base elif isinstance(base, typehints.IndexableTypeConstraint): out = base else: out = element_type(base) state.stack.append(out) def store_slice(state, args): """Clears elements off the stack like it was constructing a container, but leaves the container type back at stack[-1] since that's all that is relevant for type checking. """ _ = state.stack.pop() # End _ = state.stack.pop() # Start container = state.stack.pop() # Container type _ = state.stack.pop() # Values that would go in container state.stack.append(container) print_item = pop_top print_newline = nop def list_append(state, arg): new_element_type = Const.unwrap(state.stack.pop()) state.stack[-arg] = List[union( element_type(state.stack[-arg]), new_element_type)] def set_add(state, arg): new_element_type = Const.unwrap(state.stack.pop()) state.stack[-arg] = Set[union( element_type(state.stack[-arg]), new_element_type)] def map_add(state, arg): # PEP 572 The MAP_ADD expects the value as the first element in the stack # and the key as the second element. new_value_type = Const.unwrap(state.stack.pop()) new_key_type = Const.unwrap(state.stack.pop()) state.stack[-arg] = Dict[Union[state.stack[-arg].key_type, new_key_type], Union[state.stack[-arg].value_type, new_value_type]] load_locals = push_value(Dict[str, Any]) exec_stmt = pop_three build_class = pop_three def unpack_sequence(state, arg): t = state.stack.pop() if isinstance(t, Const): try: unpacked = [Const(ti) for ti in t.value] if len(unpacked) != arg: unpacked = [Any] * arg except TypeError: unpacked = [Any] * arg elif (isinstance(t, typehints.TupleHint.TupleConstraint) and len(t.tuple_types) == arg): unpacked = list(t.tuple_types) else: unpacked = [element_type(t)] * arg state.stack += reversed(unpacked) def dup_topx(state, arg): state.stack += state[-arg:] store_attr = pop_two delete_attr = pop_top store_global = pop_top delete_global = nop def load_const(state, arg): state.stack.append(state.const_type(arg)) load_name = push_value(Any) def build_tuple(state, arg): if arg == 0: state.stack.append(Tuple[()]) else: state.stack[-arg:] = [Tuple[[Const.unwrap(t) for t in state.stack[-arg:]]]] def build_list(state, arg): if arg == 0: state.stack.append(List[Union[()]]) else: state.stack[-arg:] = [List[reduce(union, state.stack[-arg:], Union[()])]] def build_set(state, arg): if arg == 0: state.stack.append(Set[Union[()]]) else: state.stack[-arg:] = [Set[reduce(union, state.stack[-arg:], Union[()])]] # A Dict[Union[], Union[]] is the type of an empty dict. def build_map(state, arg): if arg == 0: state.stack.append(Dict[Union[()], Union[()]]) else: state.stack[-2 * arg:] = [ Dict[reduce(union, state.stack[-2 * arg::2], Union[()]), reduce(union, state.stack[-2 * arg + 1::2], Union[()])] ] def build_const_key_map(state, arg): key_tuple = state.stack.pop() if isinstance(key_tuple, typehints.TupleHint.TupleConstraint): key_types = key_tuple.tuple_types elif isinstance(key_tuple, Const): key_types = [Const(v) for v in key_tuple.value] else: key_types = [Any] state.stack[-arg:] = [ Dict[reduce(union, key_types, Union[()]), reduce(union, state.stack[-arg:], Union[()])] ] def list_to_tuple(state, arg): base = state.stack.pop() state.stack.append(Tuple[element_type(base), ...]) def build_string(state, arg): state.stack[-arg:] = [str] def list_extend(state, arg): tail = state.stack.pop() base = state.stack[-arg] state.stack[-arg] = List[union(element_type(base), element_type(tail))] def set_update(state, arg): other = state.stack.pop() base = state.stack[-arg] state.stack[-arg] = Set[union(element_type(base), element_type(other))] def dict_update(state, arg): other = state.stack.pop() base = state.stack[-arg] if isinstance(base, typehints.Dict.DictConstraint): base_key_type = base.key_type base_value_type = base.value_type else: base_key_type = Any base_value_type = Any if isinstance(other, typehints.Dict.DictConstraint): other_key_type = other.key_type other_value_type = other.value_type else: other_key_type, other_value_type = key_value_types(element_type(other)) state.stack[-arg] = Dict[union(base_key_type, other_key_type), union(base_value_type, other_value_type)] dict_merge = dict_update def load_attr(state, arg): """Replaces the top of the stack, TOS, with getattr(TOS, co_names[arg]) Will replace with Any for builtin methods, but these don't have bytecode in CPython so that's okay. """ if (sys.version_info.major, sys.version_info.minor) >= (3, 12): # Load attribute's arg was bit-shifted in 3.12 to also allow for # adding extra information to the stack based on the lower byte, # so we have to adjust it back. # # See https://docs.python.org/3/library/dis.html#opcode-LOAD_ATTR # for more information. arg = arg >> 1 o = state.stack.pop() name = state.get_name(arg) state.stack.append(_getattr(o, name)) def _getattr(o, name): if isinstance(o, Const) and hasattr(o.value, name): return Const(getattr(o.value, name)) elif (inspect.isclass(o) and isinstance(getattr(o, name, None), (types.MethodType, types.FunctionType))): # TODO(luke-zhu): Support other callable objects func = getattr(o, name) # Python 3 has no unbound methods return Const(BoundMethod(func, o)) elif isinstance(o, row_type.RowTypeConstraint): return o.get_type_for(name) else: return Any def load_method(state, arg): """Like load_attr. Replaces TOS object with method and TOS.""" o = state.stack.pop() name = state.get_name(arg) if isinstance(o, Const): method = Const(getattr(o.value, name)) elif isinstance(o, typehints.AnyTypeConstraint): method = typehints.Any elif hasattr(o, name): attr = getattr(o, name) if isinstance(attr, _MethodDescriptorType): # Skip builtins since they don't disassemble. method = typehints.Any else: method = Const(BoundMethod(attr, o)) else: method = typehints.Any state.stack.append(method) def compare_op(state, unused_arg): # Could really be anything... state.stack[-2:] = [bool] is_op = compare_op contains_op = compare_op def import_name(state, unused_arg): state.stack[-2:] = [Any] import_from = push_value(Any) def load_global(state, arg): if (sys.version_info.major, sys.version_info.minor) >= (3, 11): arg = arg >> 1 state.stack.append(state.get_global(arg)) def load_small_int(state, arg): state.stack.append(Const(arg)) store_map = pop_two def load_fast(state, arg): state.stack.append(state.vars[arg]) load_fast_borrow = load_fast def load_fast_load_fast(state, arg): arg1 = arg >> 4 arg2 = arg & 15 state.stack.append(state.vars[arg1]) state.stack.append(state.vars[arg2]) load_fast_borrow_load_fast_borrow = load_fast_load_fast load_fast_check = load_fast def load_fast_and_clear(state, arg): state.stack.append(state.vars[arg]) state.vars[arg] = None def store_fast(state, arg): state.vars[arg] = state.stack.pop() def store_fast_store_fast(state, arg): arg1 = arg >> 4 arg2 = arg & 15 state.vars[arg1] = state.stack.pop() state.vars[arg2] = state.stack.pop() def store_fast_load_fast(state, arg): arg1 = arg >> 4 arg2 = arg & 15 state.vars[arg1] = state.stack.pop() state.stack.append(state.vars[arg2]) def delete_fast(state, arg): state.vars[arg] = Any # really an error def swap(state, arg): state.stack[-arg], state.stack[-1] = state.stack[-1], state.stack[-arg] def reraise(state, arg): pass # bpo-43683 Adds GEN_START in Python 3.10, but removed in Python 3.11 # https://github.com/python/cpython/pull/25138 def gen_start(state, arg): assert len(state.stack) == 0 def load_closure(state, arg): # The arg is no longer offset by len(covar_names) as of 3.11 # See https://docs.python.org/3/library/dis.html#opcode-LOAD_CLOSURE if (sys.version_info.major, sys.version_info.minor) >= (3, 11): arg -= len(state.co.co_varnames) state.stack.append(state.closure_type(arg)) def load_deref(state, arg): # The arg is no longer offset by len(covar_names) as of 3.11 # See https://docs.python.org/3/library/dis.html#opcode-LOAD_DEREF if (sys.version_info.major, sys.version_info.minor) >= (3, 11): arg -= len(state.co.co_varnames) state.stack.append(state.closure_type(arg)) def make_function(state, arg): """Creates a function with the arguments at the top of the stack. """ # TODO(luke-zhu): Handle default argument types globals = state.f.__globals__ # Inherits globals from the current frame tos = state.stack[-1].value # In Python 3.11 lambdas no longer have fully qualified names on the stack, # so we check for this case (AKA the code is top of stack.) if isinstance(tos, types.CodeType): func_name = None func_code = tos pop_count = 1 is_lambda = True else: func_name = tos func_code = state.stack[-2].value pop_count = 2 is_lambda = False closure = None if (sys.version_info.major, sys.version_info.minor) < (3, 13): # arg contains flags, with corresponding stack values if positive. # https://docs.python.org/3.6/library/dis.html#opcode-MAKE_FUNCTION pop_count += bin(arg).count('1') if arg & 0x08: # Convert types in Tuple constraint to a tuple of CPython cells. # https://stackoverflow.com/a/44670295 if is_lambda: closureTuplePos = -2 else: closureTuplePos = -3 closure = tuple((lambda _: lambda: _)(t).__closure__[0] for t in state.stack[closureTuplePos].tuple_types) func = types.FunctionType(func_code, globals, name=func_name, closure=closure) assert pop_count <= len(state.stack) state.stack[-pop_count:] = [Const(func)] def set_function_attribute(state, arg): func = state.stack.pop().value attr = state.stack.pop().value closure = None if arg & 0x08: closure = tuple((lambda _: lambda: _)(t).__closure__[0] for t in state.stack[attr].tuple_types) new_func = types.FunctionType( func.code, func.globals, name=func.name, closure=closure) if arg & 0x10: new_func.__annotate__ = attr state.stack.append(Const(new_func)) def make_closure(state, arg): state.stack[-arg - 2:] = [Any] # a callable def build_slice(state, arg): state.stack[-arg:] = [slice] # a slice object def to_bool(state, arg): state.stack[-1] = bool def format_value(state, arg): if arg & 0x04: state.stack.pop() state.stack.pop() state.stack.append(str) def convert_value(state, arg): state.stack.pop() state.stack.append(str) def format_simple(state, arg): state.stack.pop() state.stack.append(str) def format_with_spec(state, arg): state.stack.pop() state.stack.pop() state.stack.append(str) def _unpack_lists(state, arg): """Extract inner types of Lists and Tuples. Pops arg count items from the stack, concatenates their inner types into 1 list, and returns that list. Example: if stack[-arg:] == [[i1, i2], [i3]], the output is [i1, i2, i3] """ types = [] for i in range(arg, 0, -1): type_constraint = state.stack[-i] if isinstance(type_constraint, typehints.IndexableTypeConstraint): types.extend(type_constraint._inner_types()) elif type_constraint == Union[()]: continue else: logging.debug('Unhandled type_constraint: %r', type_constraint) types.append(typehints.Any) state.stack[-arg:] = [] return types def build_list_unpack(state, arg): """Joins arg count iterables from the stack into a single list.""" state.stack.append(List[Union[_unpack_lists(state, arg)]]) def build_set_unpack(state, arg): """Joins arg count iterables from the stack into a single set.""" state.stack.append(Set[Union[_unpack_lists(state, arg)]]) def build_tuple_unpack(state, arg): """Joins arg count iterables from the stack into a single tuple.""" state.stack.append(Tuple[Union[_unpack_lists(state, arg)], ...]) def build_tuple_unpack_with_call(state, arg): """Same as build_tuple_unpack, with an extra fn argument at the bottom of the stack, which remains untouched.""" build_tuple_unpack(state, arg) def build_map_unpack(state, arg): """Joins arg count maps from the stack into a single dict.""" key_types = [] value_types = [] for _ in range(arg): type_constraint = state.stack.pop() if isinstance(type_constraint, typehints.Dict.DictConstraint): key_types.append(type_constraint.key_type) value_types.append(type_constraint.value_type) else: key_type, value_type = key_value_types(element_type(type_constraint)) key_types.append(key_type) value_types.append(value_type) state.stack.append(Dict[Union[key_types], Union[value_types]]) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/pandas_type_compatibility.py0000644000175100017510000002422515206113033025524 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # r"""Utilities for converting between Beam schemas and pandas DataFrames. Imposes a mapping between native Python typings (specifically those compatible with :mod:`apache_beam.typehints.schemas`), and common pandas dtypes:: pandas dtype Python typing np.int{8,16,32,64} <-----> np.int{8,16,32,64}* pd.Int{8,16,32,64}Dtype <-----> Optional[np.int{8,16,32,64}]* np.float{32,64} <-----> Optional[np.float{32,64}] \--- np.float{32,64} Not supported <------ Optional[bytes] np.bool <-----> np.bool np.dtype('S') <-----> bytes pd.BooleanDType() <-----> Optional[bool] pd.StringDType() <-----> Optional[str] \--- str np.object <-----> Any * int, float, bool are treated the same as np.int64, np.float64, np.bool Note that when converting to pandas dtypes, any types not specified here are shunted to ``np.object``. Similarly when converting from pandas to Python types, types that aren't otherwise specified here are shunted to ``Any``. Notably, this includes ``np.datetime64``. Pandas does not support hierarchical data natively. Currently, all structured types (``Sequence``, ``Mapping``, nested ``NamedTuple`` types), are shunted to ``np.object`` like all other unknown types. In the future these types may be given special consideration. Note utilities in this package are for internal use only, we make no backward compatibility guarantees, except for the type mapping itself. """ # ruff: noqa: UP006 from typing import Any from typing import List from typing import Optional import numpy as np import pandas as pd from apache_beam.typehints.batch import BatchConverter from apache_beam.typehints.row_type import RowTypeConstraint from apache_beam.typehints.typehints import is_nullable from apache_beam.typehints.typehints import normalize # No public API currently, this just exists to register BatchConverter # implementations. __all__ = [] # Name for a valueless field-level option which, when present, indicates that # a field should map to an index in the Beam DataFrame API. INDEX_OPTION_NAME = 'beam:dataframe:index' # Generate type map (presented visually in the docstring) _BIDIRECTIONAL = [ (bool, bool), (np.int8, np.int8), (np.int16, np.int16), (np.int32, np.int32), (np.int64, np.int64), (pd.Int8Dtype(), Optional[np.int8]), (pd.Int16Dtype(), Optional[np.int16]), (pd.Int32Dtype(), Optional[np.int32]), (pd.Int64Dtype(), Optional[np.int64]), (np.float32, Optional[np.float32]), (np.float64, Optional[np.float64]), (object, Any), (pd.StringDtype(), Optional[str]), (pd.BooleanDtype(), Optional[bool]), ] PANDAS_TO_BEAM = { pd.Series([], dtype=dtype).dtype: fieldtype for dtype, fieldtype in _BIDIRECTIONAL } BEAM_TO_PANDAS = {fieldtype: dtype for dtype, fieldtype in _BIDIRECTIONAL} # Shunt non-nullable Beam types to the same pandas types as their non-nullable # equivalents for FLOATs, DOUBLEs, and STRINGs. pandas has no non-nullable dtype # for these. OPTIONAL_SHUNTS = [np.float32, np.float64, str] for typehint in OPTIONAL_SHUNTS: BEAM_TO_PANDAS[typehint] = BEAM_TO_PANDAS[Optional[typehint]] # int, float -> int64, np.float64 BEAM_TO_PANDAS[int] = BEAM_TO_PANDAS[np.int64] BEAM_TO_PANDAS[Optional[int]] = BEAM_TO_PANDAS[Optional[np.int64]] BEAM_TO_PANDAS[float] = BEAM_TO_PANDAS[np.float64] BEAM_TO_PANDAS[Optional[float]] = BEAM_TO_PANDAS[Optional[np.float64]] BEAM_TO_PANDAS[bytes] = 'bytes' # Add shunts for normalized (Beam) typehints as well BEAM_TO_PANDAS.update({ normalize(typehint): pandas_dtype for (typehint, pandas_dtype) in BEAM_TO_PANDAS.items() }) def dtype_from_typehint(typehint): # Default to np.object. This is lossy, we won't be able to recover # the type at the output. return BEAM_TO_PANDAS.get(typehint, object) def dtype_to_fieldtype(dtype): fieldtype = PANDAS_TO_BEAM.get(dtype) if fieldtype is not None: return fieldtype elif dtype.kind == 'S': return bytes else: return Any @BatchConverter.register(name="pandas") def create_pandas_batch_converter( element_type: type, batch_type: type) -> BatchConverter: if batch_type == pd.DataFrame: return DataFrameBatchConverter.from_typehints( element_type=element_type, batch_type=batch_type) elif batch_type == pd.Series: return SeriesBatchConverter.from_typehints( element_type=element_type, batch_type=batch_type) raise TypeError("batch type must be pd.Series or pd.DataFrame") class DataFrameBatchConverter(BatchConverter): def __init__( self, element_type: RowTypeConstraint, ): super().__init__(pd.DataFrame, element_type) self._columns = [name for name, _ in element_type._fields] @staticmethod def from_typehints(element_type, batch_type) -> Optional['DataFrameBatchConverter']: assert batch_type == pd.DataFrame if not isinstance(element_type, RowTypeConstraint): element_type = RowTypeConstraint.from_user_type(element_type) if element_type is None: raise TypeError( "Element type must be compatible with Beam Schemas (" "https://beam.apache.org/documentation/programming-guide/#schemas) " "for batch type pd.DataFrame") index_columns = [ field_name for (field_name, field_options) in element_type._field_options.items() if any(key == INDEX_OPTION_NAME for key, value in field_options) ] if index_columns: return DataFrameBatchConverterKeepIndex(element_type, index_columns) else: return DataFrameBatchConverterDropIndex(element_type) def _get_series(self, batch: pd.DataFrame): raise NotImplementedError def explode_batch(self, batch: pd.DataFrame): # TODO(https://github.com/apache/beam/issues/22948): Only do null checks for # nullable types def make_null_checking_generator(series): nulls = pd.isnull(series) return (None if isnull else value for isnull, value in zip(nulls, series)) all_series = self._get_series(batch) iterators = [make_null_checking_generator(series) for series in all_series] for values in zip(*iterators): yield self._element_type.user_type( **{ column: value for column, value in zip(self._columns, values) }) def combine_batches(self, batches: List[pd.DataFrame]): return pd.concat(batches) def estimate_byte_size(self, batch: pd.DataFrame): return batch.memory_usage().sum() def get_length(self, batch: pd.DataFrame): return len(batch) class DataFrameBatchConverterDropIndex(DataFrameBatchConverter): """A DataFrameBatchConverter that assumes the DataFrame index has no meaning. When producing a DataFrame from Rows, a meaningless index will be generated. When exploding a DataFrame into Rows, the index will be dropped. """ def _get_series(self, batch: pd.DataFrame): return [batch[column] for column in batch.columns] def produce_batch(self, elements): batch = pd.DataFrame.from_records(elements, columns=self._columns) for column, typehint in self._element_type._fields: batch[column] = batch[column].astype(dtype_from_typehint(typehint)) return batch class DataFrameBatchConverterKeepIndex(DataFrameBatchConverter): """A DataFrameBatchConverter that preserves the DataFrame index. This is tracked via options on the Beam schema. Each field in the schema that should map to the index is tagged in an option with name 'dataframe:index'. """ def __init__(self, element_type: RowTypeConstraint, index_columns: List[Any]): super().__init__(element_type) self._index_columns = index_columns def _get_series(self, batch: pd.DataFrame): assert list(batch.index.names) == self._index_columns return [ batch.index.get_level_values(i) for i in range(len(batch.index.names)) ] + [batch[column] for column in batch.columns] def produce_batch(self, elements): # Note from_records has an index= parameter batch = pd.DataFrame.from_records(elements, columns=self._columns) for column, typehint in self._element_type._fields: batch[column] = batch[column].astype(dtype_from_typehint(typehint)) return batch.set_index(self._index_columns) class SeriesBatchConverter(BatchConverter): def __init__( self, element_type: type, dtype, ): super().__init__(pd.DataFrame, element_type) self._dtype = dtype if is_nullable(element_type): def unbatch(series): for isnull, value in zip(pd.isnull(series), series): yield None if isnull else value else: def unbatch(series): yield from series self.explode_batch = unbatch @staticmethod def from_typehints(element_type, batch_type) -> Optional['SeriesBatchConverter']: assert batch_type == pd.Series dtype = dtype_from_typehint(element_type) return SeriesBatchConverter(element_type, dtype) def produce_batch(self, elements: List[Any]) -> pd.Series: return pd.Series(elements, dtype=self._dtype) def explode_batch(self, batch: pd.Series): raise NotImplementedError( "explode_batch should be generated in SeriesBatchConverter.__init__") def combine_batches(self, batches: List[pd.Series]): return pd.concat(batches) def estimate_byte_size(self, batch: pd.Series): return batch.memory_usage() def get_length(self, batch: pd.Series): return len(batch) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/pandas_type_compatibility_test.py0000644000175100017510000002020315206113033026553 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for pandas batched type converters.""" import unittest from typing import Any from typing import Optional import numpy as np import pandas as pd from parameterized import parameterized from parameterized import parameterized_class from apache_beam.typehints import row_type from apache_beam.typehints import typehints from apache_beam.typehints.batch import BatchConverter @parameterized_class([ { 'batch_typehint': pd.DataFrame, 'element_typehint': row_type.RowTypeConstraint.from_fields([ ('foo', int), ('bar', float), ('baz', str), ]), 'match_index': False, 'batch': pd.DataFrame({ 'foo': pd.Series(range(100), dtype='int64'), 'bar': pd.Series([i / 100 for i in range(100)], dtype='float64'), 'baz': pd.Series([str(i) for i in range(100)], dtype=pd.StringDtype()), }), }, { 'batch_typehint': pd.DataFrame, 'element_typehint': row_type.RowTypeConstraint.from_fields( [ ('an_index', int), ('foo', int), ('bar', float), ('baz', str), ], field_options={'an_index': [('beam:dataframe:index', None)]}, ), 'match_index': True, 'batch': pd.DataFrame({ 'foo': pd.Series(range(100), dtype='int64'), 'bar': pd.Series([i / 100 for i in range(100)], dtype='float64'), 'baz': pd.Series([str(i) for i in range(100)], dtype=pd.StringDtype()), }).set_index(pd.Index(range(123, 223), dtype='int64', name='an_index')), }, { 'batch_typehint': pd.DataFrame, 'element_typehint': row_type.RowTypeConstraint.from_fields( [ ('an_index', int), ('another_index', int), ('foo', int), ('bar', float), ('baz', str), ], field_options={ 'an_index': [('beam:dataframe:index', None)], 'another_index': [('beam:dataframe:index', None)], }), 'match_index': True, 'batch': pd.DataFrame({ 'foo': pd.Series(range(100), dtype='int64'), 'bar': pd.Series([i / 100 for i in range(100)], dtype='float64'), 'baz': pd.Series([str(i) for i in range(100)], dtype=pd.StringDtype()), }).set_index([ pd.Index(range(123, 223), dtype='int64', name='an_index'), pd.Index(range(475, 575), dtype='int64', name='another_index'), ]), }, { 'batch_typehint': pd.Series, 'element_typehint': int, 'match_index': False, 'batch': pd.Series(range(500)), }, { 'batch_typehint': pd.Series, 'element_typehint': str, 'batch': pd.Series(['foo', 'bar', 'baz', 'def', 'ghi', 'abc'] * 10, dtype=pd.StringDtype()), }, { 'batch_typehint': pd.Series, 'element_typehint': Optional[np.int64], 'batch': pd.Series((i if i % 11 else None for i in range(500)), dtype=pd.Int64Dtype()), }, { 'batch_typehint': pd.Series, 'element_typehint': Optional[str], 'batch': pd.Series(['foo', None, 'bar', 'baz', None, 'def', 'ghi'] * 10, dtype=pd.StringDtype()), }, ]) class PandasBatchConverterTest(unittest.TestCase): def create_batch_converter(self): return BatchConverter.from_typehints( element_type=self.element_typehint, batch_type=self.batch_typehint) def setUp(self): self.converter = self.create_batch_converter() self.normalized_batch_typehint = typehints.normalize(self.batch_typehint) self.normalized_element_typehint = typehints.normalize( self.element_typehint) def equality_check(self, left, right): if isinstance(left, pd.DataFrame): if self.match_index: pd.testing.assert_frame_equal(left.sort_index(), right.sort_index()) else: pd.testing.assert_frame_equal( left.sort_values(by=list(left.columns)).reset_index(drop=True), right.sort_values(by=list(right.columns)).reset_index(drop=True)) elif isinstance(left, pd.Series): pd.testing.assert_series_equal( left.sort_values().reset_index(drop=True), right.sort_values().reset_index(drop=True)) else: raise TypeError(f"Encountered unexpected type, left is a {type(left)!r}") def test_typehint_validates(self): typehints.validate_composite_type_param(self.batch_typehint, '') typehints.validate_composite_type_param(self.element_typehint, '') def test_type_check_batch(self): typehints.check_constraint(self.normalized_batch_typehint, self.batch) def test_type_check_element(self): for element in self.converter.explode_batch(self.batch): typehints.check_constraint(self.normalized_element_typehint, element) def test_explode_rebatch(self): exploded = list(self.converter.explode_batch(self.batch)) rebatched = self.converter.produce_batch(exploded) typehints.check_constraint(self.normalized_batch_typehint, rebatched) self.equality_check(self.batch, rebatched) def _split_batch_into_n_partitions(self, N): elements = list(self.converter.explode_batch(self.batch)) # Split elements into N contiguous partitions element_batches = [ elements[len(elements) * i // N:len(elements) * (i + 1) // N] for i in range(N) ] lengths = [len(element_batch) for element_batch in element_batches] batches = [ self.converter.produce_batch(element_batch) for element_batch in element_batches ] return batches, lengths @parameterized.expand([ (2, ), (3, ), (10, ), ]) def test_combine_batches(self, N): batches, _ = self._split_batch_into_n_partitions(N) # Combine the batches, output should be equivalent to the original batch combined = self.converter.combine_batches(batches) self.equality_check(self.batch, combined) @parameterized.expand([ (2, ), (3, ), (10, ), ]) def test_get_length(self, N): batches, lengths = self._split_batch_into_n_partitions(N) for batch, expected_length in zip(batches, lengths): self.assertEqual(self.converter.get_length(batch), expected_length) def test_equals(self): self.assertTrue(self.converter == self.create_batch_converter()) self.assertTrue(self.create_batch_converter() == self.converter) def test_hash(self): self.assertEqual(hash(self.create_batch_converter()), hash(self.converter)) class PandasBatchConverterErrorsTest(unittest.TestCase): @parameterized.expand([ ( Any, row_type.RowTypeConstraint.from_fields([ ("bar", Optional[float]), # noqa: F821 ("baz", Optional[str]), # noqa: F821 ]), r'batch type must be pd\.Series or pd\.DataFrame', ), ( pd.DataFrame, Any, r'Element type must be compatible with Beam Schemas', ), ]) def test_construction_errors( self, batch_typehint, element_typehint, error_regex): with self.assertRaisesRegex(TypeError, error_regex): BatchConverter.from_typehints( element_type=element_typehint, batch_type=batch_typehint) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/pytorch_type_compatibility.py0000644000175100017510000001152115206113033025741 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from typing import Optional import torch from apache_beam.typehints import typehints from apache_beam.typehints.batch import BatchConverter from apache_beam.typehints.batch import N class PytorchBatchConverter(BatchConverter): def __init__( self, batch_type, element_type, dtype, element_shape=(), partition_dimension=0): super().__init__(batch_type, element_type) self.dtype = dtype self.element_shape = element_shape self.partition_dimension = partition_dimension @staticmethod @BatchConverter.register(name="pytorch") def from_typehints(element_type, batch_type) -> Optional['PytorchBatchConverter']: if not isinstance(element_type, PytorchTypeHint.PytorchTypeConstraint): element_type = PytorchTensor[element_type, ()] if not isinstance(batch_type, PytorchTypeHint.PytorchTypeConstraint): if not batch_type == torch.Tensor: raise TypeError( "batch type must be torch.Tensor or " "beam.typehints.pytorch_type_compatibility.PytorchTensor[..]") batch_type = PytorchTensor[element_type.dtype, (N, )] if not batch_type.dtype == element_type.dtype: raise TypeError( "batch type and element type must have equivalent dtypes " f"(batch={batch_type.dtype}, element={element_type.dtype})") computed_element_shape = list(batch_type.shape) partition_dimension = computed_element_shape.index(N) computed_element_shape.pop(partition_dimension) if not tuple(computed_element_shape) == element_type.shape: raise TypeError( "Could not align batch type's batch dimension with element type. " f"(batch type dimensions: {batch_type.shape}, element type " f"dimenstions: {element_type.shape}") return PytorchBatchConverter( batch_type, element_type, batch_type.dtype, element_type.shape, partition_dimension) def produce_batch(self, elements): return torch.stack(elements, dim=self.partition_dimension) def explode_batch(self, batch): """Convert an instance of B to Generator[E].""" yield from torch.swapaxes(batch, self.partition_dimension, 0) def combine_batches(self, batches): return torch.cat(batches, dim=self.partition_dimension) def get_length(self, batch): return batch.size(dim=self.partition_dimension) def estimate_byte_size(self, batch): return batch.nelement() * batch.element_size() class PytorchTypeHint(): class PytorchTypeConstraint(typehints.TypeConstraint): def __init__(self, dtype, shape=()): self.dtype = dtype self.shape = shape def type_check(self, batch): if not isinstance(batch, torch.Tensor): raise TypeError(f"Batch {batch!r} is not an instance of torch.Tensor") if not batch.dtype == self.dtype: raise TypeError( f"Batch {batch!r} does not have expected dtype: {self.dtype!r}") for dim in range(len(self.shape)): if not self.shape[dim] == N and not batch.shape[dim] == self.shape[dim]: raise TypeError( f"Batch {batch!r} does not have expected shape: {self.shape!r}") def _consistent_with_check_(self, sub): # TODO Check sub against batch type, and element type return True def __key(self): return (self.dtype, self.shape) def __eq__(self, other) -> bool: if isinstance(other, PytorchTypeHint.PytorchTypeConstraint): return self.__key() == other.__key() return NotImplemented def __hash__(self) -> int: return hash(self.__key()) def __repr__(self): if self.shape == (N, ): return f'PytorchTensor[{self.dtype!r}]' else: return f'PytorchTensor[{self.dtype!r}, {self.shape!r}]' def __getitem__(self, value): if isinstance(value, tuple): if len(value) == 2: dtype, shape = value return self.PytorchTypeConstraint(dtype, shape=shape) else: raise ValueError else: dtype = value return self.PytorchTypeConstraint( dtype, shape=(N, )) PytorchTensor = PytorchTypeHint() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/pytorch_type_compatibility_test.py0000644000175100017510000001314515206113033027004 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for pytorch_type_compabitility.""" import unittest from typing import Any import pytest from parameterized import parameterized from parameterized import parameterized_class from apache_beam.typehints import typehints from apache_beam.typehints.batch import BatchConverter from apache_beam.typehints.batch import N # Protect against environments where pytorch library is not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: import torch from apache_beam.typehints.pytorch_type_compatibility import PytorchTensor except ImportError: raise unittest.SkipTest('PyTorch dependencies are not installed') @parameterized_class([ { 'batch_typehint': torch.Tensor, 'element_typehint': PytorchTensor[torch.int32, ()], 'batch': torch.tensor(range(100), dtype=torch.int32) }, { 'batch_typehint': PytorchTensor[torch.int64, (N, 10)], 'element_typehint': PytorchTensor[torch.int64, (10, )], 'batch': torch.tensor([list(range(i, i + 10)) for i in range(100)], dtype=torch.int64), }, ]) @pytest.mark.uses_pytorch class PytorchBatchConverterTest(unittest.TestCase): def create_batch_converter(self): return BatchConverter.from_typehints( element_type=self.element_typehint, batch_type=self.batch_typehint) def setUp(self): self.converter = self.create_batch_converter() self.normalized_batch_typehint = typehints.normalize(self.batch_typehint) self.normalized_element_typehint = typehints.normalize( self.element_typehint) def equality_check(self, left, right): if isinstance(left, torch.Tensor): self.assertTrue(torch.equal(left, right)) else: raise TypeError(f"Encountered unexpected type, left is a {type(left)!r}") def test_typehint_validates(self): typehints.validate_composite_type_param(self.batch_typehint, '') typehints.validate_composite_type_param(self.element_typehint, '') def test_type_check_batch(self): typehints.check_constraint(self.normalized_batch_typehint, self.batch) def test_type_check_element(self): for element in self.converter.explode_batch(self.batch): typehints.check_constraint(self.normalized_element_typehint, element) def test_explode_rebatch(self): exploded = list(self.converter.explode_batch(self.batch)) rebatched = self.converter.produce_batch(exploded) typehints.check_constraint(self.normalized_batch_typehint, rebatched) self.equality_check(self.batch, rebatched) def _split_batch_into_n_partitions(self, N): elements = list(self.converter.explode_batch(self.batch)) # Split elements into N contiguous partitions element_batches = [ elements[len(elements) * i // N:len(elements) * (i + 1) // N] for i in range(N) ] lengths = [len(element_batch) for element_batch in element_batches] batches = [ self.converter.produce_batch(element_batch) for element_batch in element_batches ] return batches, lengths @parameterized.expand([ (2, ), (3, ), (10, ), ]) def test_combine_batches(self, N): batches, _ = self._split_batch_into_n_partitions(N) # Combine the batches, output should be equivalent to the original batch combined = self.converter.combine_batches(batches) self.equality_check(self.batch, combined) @parameterized.expand([ (2, ), (3, ), (10, ), ]) def test_get_length(self, N): batches, lengths = self._split_batch_into_n_partitions(N) for batch, expected_length in zip(batches, lengths): self.assertEqual(self.converter.get_length(batch), expected_length) def test_equals(self): self.assertTrue(self.converter == self.create_batch_converter()) self.assertTrue(self.create_batch_converter() == self.converter) def test_hash(self): self.assertEqual(hash(self.create_batch_converter()), hash(self.converter)) class PytorchBatchConverterErrorsTest(unittest.TestCase): @parameterized.expand([ ( Any, PytorchTensor[torch.int64, ()], ( r'batch type must be torch\.Tensor or ' r'beam\.typehints\.pytorch_type_compatibility.PytorchTensor'), ), ( PytorchTensor[torch.int64, (3, N, 2)], PytorchTensor[torch.int64, (3, 7)], r'Could not align batch type\'s batch dimension', ), ( PytorchTensor[torch.int64, (N, 10)], PytorchTensor[torch.float32, (10, )], r'batch type and element type must have equivalent dtypes', ), ]) def test_construction_errors( self, batch_typehint, element_typehint, error_regex): with self.assertRaisesRegex(TypeError, error_regex): BatchConverter.from_typehints( element_type=element_typehint, batch_type=batch_typehint) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/row_type.py0000644000175100017510000002235315206113033022134 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file # ruff: noqa: UP006 from __future__ import annotations import dataclasses from typing import Any from typing import Dict from typing import Optional from typing import Sequence from typing import Tuple from apache_beam.typehints import typehints from apache_beam.typehints.native_type_compatibility import match_dataclass_for_row from apache_beam.typehints.native_type_compatibility import match_is_named_tuple from apache_beam.typehints.schema_registry import SchemaTypeRegistry # Name of the attribute added to user types (existing and generated) to store # the corresponding schema ID _BEAM_SCHEMA_ID = "_beam_schema_id" def _user_type_is_generated(user_type: type) -> bool: if not hasattr(user_type, _BEAM_SCHEMA_ID): return False schema_id = getattr(user_type, _BEAM_SCHEMA_ID) type_name = 'BeamSchema_{}'.format(schema_id.replace('-', '_')) return user_type.__name__ == type_name class RowTypeConstraint(typehints.TypeConstraint): def __init__( self, fields: Sequence[Tuple[str, type]], user_type, schema_options: Optional[Sequence[Tuple[str, Any]]] = None, field_options: Optional[Dict[str, Sequence[Tuple[str, Any]]]] = None, field_descriptions: Optional[Dict[str, str]] = None): """For internal use only, no backwards comatibility guaratees. See https://beam.apache.org/documentation/programming-guide/#schemas-for-pl-types for guidance on creating PCollections with inferred schemas. Note RowTypeConstraint does not currently store arbitrary functions for converting to/from the user type. Instead, we support ``NamedTuple`` and ``dataclasses`` user types and make the follow assumptions: - The user type can be constructed with field values as arguments in order (i.e. ``constructor(*field_values)``). - Field values can be accessed from instances of the user type by attribute (i.e. with ``getattr(obj, field_name)``). The RowTypeConstraint constructor should not be called directly (even internally to Beam). Prefer static methods ``from_user_type`` or ``from_fields``. Parameters: fields: a list of (name, type) tuples, representing the schema inferred from user_type. user_type: constructor for a user type (e.g. NamedTuple class) that is used to represent this schema in user code. schema_options: A list of (key, value) tuples representing schema-level options. field_options: A dictionary representing field-level options. Dictionary keys are field names, and dictionary values are lists of (key, value) tuples representing field-level options for that field. """ # Recursively wrap row types in a RowTypeConstraint self._fields = tuple((name, RowTypeConstraint.from_user_type(typ) or typ) for name, typ in fields) self._user_type = user_type # Note schema ID can be None if the schema is not registered yet. # Currently registration happens when converting to schema protos, in # apache_beam.typehints.schemas self._schema_id = getattr(self._user_type, _BEAM_SCHEMA_ID, None) if self._schema_id and _BEAM_SCHEMA_ID not in self._user_type.__dict__: # schema id does not inherit. Unset if schema id is from base class self._schema_id = None self._schema_options = schema_options or [] self._field_options = field_options or {} self._field_descriptions = field_descriptions or {} @staticmethod def from_user_type( user_type: type, schema_options: Optional[Sequence[Tuple[str, Any]]] = None, field_options: Optional[Dict[str, Sequence[Tuple[str, Any]]]] = None ) -> Optional[RowTypeConstraint]: if match_is_named_tuple(user_type): fields = [(name, user_type.__annotations__[name]) for name in user_type._fields] elif match_dataclass_for_row(user_type): fields = [(field.name, field.type) for field in dataclasses.fields(user_type)] else: return None field_descriptions = getattr(user_type, '_field_descriptions', None) if _user_type_is_generated(user_type): return RowTypeConstraint.from_fields( fields, schema_id=getattr(user_type, _BEAM_SCHEMA_ID), schema_options=schema_options, field_options=field_options, field_descriptions=field_descriptions) # TODO(https://github.com/apache/beam/issues/22125): Add user API for # specifying schema/field options return RowTypeConstraint( fields=fields, user_type=user_type, schema_options=schema_options, field_options=field_options, field_descriptions=field_descriptions) @staticmethod def from_fields( fields: Sequence[Tuple[str, type]], schema_id: Optional[str] = None, schema_options: Optional[Sequence[Tuple[str, Any]]] = None, field_options: Optional[Dict[str, Sequence[Tuple[str, Any]]]] = None, schema_registry: Optional[SchemaTypeRegistry] = None, field_descriptions: Optional[Dict[str, str]] = None, ) -> RowTypeConstraint: return GeneratedClassRowTypeConstraint( fields, schema_id=schema_id, schema_options=schema_options, field_options=field_options, schema_registry=schema_registry, field_descriptions=field_descriptions) def __call__(self, *args, **kwargs): # We make RowTypeConstraint callable (defers to constructing the user type) # so that Python will recognize it as a type. This allows RowTypeConstraint # to be used in conjunction with native typehints, like Optional. # CPython (prior to 3.11) considers anything callable to be a type: # https://github.com/python/cpython/blob/d348afa15d5a997e7a8e51c0f789f41cb15cc651/Lib/typing.py#L137-L167 return self._user_type(*args, **kwargs) @property def user_type(self): return self._user_type def set_schema_id(self, schema_id): self._schema_id = schema_id if self._user_type is not None: setattr(self._user_type, _BEAM_SCHEMA_ID, self._schema_id) @property def schema_id(self): return self._schema_id @property def schema_options(self): return self._schema_options def field_options(self, field_name): # Raise if field_name is not one of the fields? return self._field_options.get(field_name, []) def _consistent_with_check_(self, sub): return self == sub def type_check(self, instance): from apache_beam import Row return isinstance(instance, (Row, self._user_type)) def _inner_types(self): """Iterates over the inner types of the composite type.""" return [field[1] for field in self._fields] def __eq__(self, other): return type(self) == type(other) and self._fields == other._fields def __hash__(self): return hash(self._fields) def __repr__(self): return 'Row(%s)' % ', '.join( '%s=%s' % (name, repr(t)) for name, t in self._fields) def get_type_for(self, name): try: return dict(self._fields)[name] except KeyError: return typehints.Any class GeneratedClassRowTypeConstraint(RowTypeConstraint): """Specialization of RowTypeConstraint which relies on a generated user_type. Since the generated user_type cannot be pickled, we supply a custom __reduce__ function that will regenerate the user_type. """ def __init__( self, fields, schema_id: Optional[str] = None, schema_options: Optional[Sequence[Tuple[str, Any]]] = None, field_options: Optional[Dict[str, Sequence[Tuple[str, Any]]]] = None, schema_registry: Optional[SchemaTypeRegistry] = None, field_descriptions: Optional[Dict[str, str]] = None, ): from apache_beam.typehints.schemas import named_fields_to_schema from apache_beam.typehints.schemas import named_tuple_from_schema kwargs = {'schema_registry': schema_registry} if schema_registry else {} schema = named_fields_to_schema( fields, schema_id=schema_id, schema_options=schema_options, field_options=field_options, field_descriptions=field_descriptions, **kwargs) user_type = named_tuple_from_schema(schema, **kwargs) super().__init__( fields, user_type, schema_options=schema_options, field_options=field_options, field_descriptions=field_descriptions) def __reduce__(self): return ( RowTypeConstraint.from_fields, ( self._fields, self._schema_id, self._schema_options, self._field_options, None, )) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/row_type_test.py0000644000175100017510000001370515206113033023174 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the Beam Row typing functionality.""" import typing import unittest from dataclasses import dataclass import apache_beam as beam from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.typehints import row_type from apache_beam.typehints import schemas class RowTypeTest(unittest.TestCase): @staticmethod def _check_key_type_and_count(x) -> int: key_type = type(x[0]) if row_type._user_type_is_generated(key_type): raise RuntimeError("Type after GBK not preserved, get generated type") if not hasattr(key_type, row_type._BEAM_SCHEMA_ID): raise RuntimeError("Type after GBK missing Beam schema ID") return len(x[1]) def test_group_by_key_namedtuple(self): MyNamedTuple = typing.NamedTuple( "MyNamedTuple", [("id", int), ("name", str)]) beam.coders.typecoders.registry.register_row(MyNamedTuple) def generate(num: int): for i in range(100): yield (MyNamedTuple(i, 'a'), num) pipeline = TestPipeline(is_integration_test=False) with pipeline as p: result = ( p | 'Create' >> beam.Create([i for i in range(10)]) | 'Generate' >> beam.ParDo(generate).with_output_types( tuple[MyNamedTuple, int]) | 'GBK' >> beam.GroupByKey() | 'Count Elements' >> beam.Map(self._check_key_type_and_count)) assert_that(result, equal_to([10] * 100)) def test_group_by_key_dataclass(self): @dataclass class MyDataClass: id: int name: str beam.coders.typecoders.registry.register_row(MyDataClass) def generate(num: int): for i in range(100): yield (MyDataClass(i, 'a'), num) pipeline = TestPipeline(is_integration_test=False) with pipeline as p: result = ( p | 'Create' >> beam.Create([i for i in range(10)]) | 'Generate' >> beam.ParDo(generate).with_output_types( tuple[MyDataClass, int]) | 'GBK' >> beam.GroupByKey() | 'Count Elements' >> beam.Map(self._check_key_type_and_count)) assert_that(result, equal_to([10] * 100)) def test_group_by_key_namedtuple_union(self): Tuple1 = typing.NamedTuple("Tuple1", [("id", int)]) Tuple2 = typing.NamedTuple("Tuple2", [("id", int), ("name", str)]) def generate(num: int): for i in range(2): yield (Tuple1(i), num) yield (Tuple2(i, 'a'), num) pipeline = TestPipeline(is_integration_test=False) with pipeline as p: result = ( p | 'Create' >> beam.Create([i for i in range(2)]) | 'Generate' >> beam.ParDo(generate).with_output_types( tuple[(Tuple1 | Tuple2), int]) | 'GBK' >> beam.GroupByKey() | 'Count' >> beam.Map(lambda x: len(x[1]))) assert_that(result, equal_to([2] * 4)) # Union of dataclasses as type hint currently result in FastPrimitiveCoder # fails at GBK @unittest.skip("https://github.com/apache/beam/issues/22085") def test_group_by_key_inherited_dataclass_union(self): @dataclass class DataClassInt: id: int @dataclass class DataClassStr(DataClassInt): name: str beam.coders.typecoders.registry.register_row(DataClassInt) beam.coders.typecoders.registry.register_row(DataClassStr) def generate(num: int): for i in range(10): yield (DataClassInt(i), num) yield (DataClassStr(i, 'a'), num) pipeline = TestPipeline(is_integration_test=False) with pipeline as p: result = ( p | 'Create' >> beam.Create([i for i in range(2)]) | 'Generate' >> beam.ParDo(generate).with_output_types( tuple[(DataClassInt | DataClassStr), int]) | 'GBK' >> beam.GroupByKey() | 'Count Elements' >> beam.Map(self._check_key_type_and_count)) assert_that(result, equal_to([2] * 4)) def test_derived_dataclass_schema_id(self): @dataclass class BaseDataClass: id: int @dataclass class DerivedDataClass(BaseDataClass): name: str self.assertFalse(hasattr(BaseDataClass, row_type._BEAM_SCHEMA_ID)) schema_for_base = schemas.schema_from_element_type(BaseDataClass) self.assertTrue(hasattr(BaseDataClass, row_type._BEAM_SCHEMA_ID)) self.assertEqual( schema_for_base.id, getattr(BaseDataClass, row_type._BEAM_SCHEMA_ID)) # Getting the schema for BaseDataClass sets the _beam_schema_id schemas.typing_to_runner_api( BaseDataClass, schema_registry=schemas.SchemaTypeRegistry()) # We create a RowTypeConstraint from DerivedDataClass. # It should not inherit the _beam_schema_id from BaseDataClass! derived_row_type = row_type.RowTypeConstraint.from_user_type( DerivedDataClass) self.assertIsNone(derived_row_type._schema_id) schema_for_derived = schemas.schema_from_element_type(DerivedDataClass) self.assertTrue(hasattr(DerivedDataClass, row_type._BEAM_SCHEMA_ID)) self.assertEqual( schema_for_derived.id, getattr(DerivedDataClass, row_type._BEAM_SCHEMA_ID)) self.assertNotEqual(schema_for_derived.id, schema_for_base.id) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/schema_registry.py0000644000175100017510000000425415206113033023454 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """This module is intended for internal use only. Nothing defined here provides any backwards-compatibility guarantee. """ from uuid import uuid4 # Registry of typings for a schema by UUID class SchemaTypeRegistry(object): def __init__(self): self.by_id = {} self.by_typing = {} # currently not used def generate_new_id(self): for _ in range(100): schema_id = str(uuid4()) if schema_id not in self.by_id: return schema_id raise AssertionError( "Failed to generate a unique UUID for schema after " f"100 tries! Registry contains {len(self.by_id)} " "schemas.") def add(self, typing, schema): if schema.id: self.by_id[schema.id] = (typing, schema) def load_registered_typings(self, by_id): for id, typing in by_id.items(): if id not in self.by_id: self.by_id[id] = (typing, None) def get_registered_typings(self): # Used by save_main_session, as pb2.schema isn't picklable return {k: v[0] for k, v in self.by_id.items()} def get_typing_by_id(self, unique_id): if not unique_id: return None result = self.by_id.get(unique_id, None) return result[0] if result is not None else None def get_schema_by_id(self, unique_id): if not unique_id: return None result = self.by_id.get(unique_id, None) return result[1] if result is not None else None SCHEMA_REGISTRY = SchemaTypeRegistry() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/schemas.py0000644000175100017510000013332215206113033021706 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Support for mapping python types to proto Schemas and back again. Imposes a mapping between common Python types and Beam portable schemas (https://s.apache.org/beam-schemas):: Python Schema np.int8 <-----> BYTE np.int16 <-----> INT16 np.int32 <-----> INT32 np.int64 <-----> INT64 int ------> INT64 np.float32 <-----> FLOAT np.float64 <-----> DOUBLE float ------> DOUBLE bool <-----> BOOLEAN str <-----> STRING bytes <-----> BYTES ByteString ------> BYTES Timestamp <-----> LogicalType(urn="beam:logical_type:micros_instant:v1") datetime.date <---> LogicalType(urn="beam:logical_type:date:v1") Decimal <-----> LogicalType(urn="beam:logical_type:fixed_decimal:v1") Mapping <-----> MapType Sequence <-----> ArrayType NamedTuple <-----> RowType beam.Row ------> RowType One direction mapping of Python types from Beam portable schemas: bytes <------ LogicalType(urn="beam:logical_type:fixed_bytes:v1") <------ LogicalType(urn="beam:logical_type:var_bytes:v1") str <------ LogicalType(urn="beam:logical_type:fixed_char:v1") <------ LogicalType(urn="beam:logical_type:var_char:v1") Timestamp <------ LogicalType(urn="beam:logical_type:millis_instant:v1") Note that some of these mappings are provided as conveniences, but they are lossy and will not survive a roundtrip from python to Beam schemas and back. For example, the Python type :code:`int` will map to :code:`INT64` in Beam schemas but converting that back to a Python type will yield :code:`np.int64`. :code:`nullable=True` on a Beam :code:`FieldType` is represented in Python by wrapping the type in :code:`Optional`. This module is intended for internal use only. Nothing defined here provides any backwards-compatibility guarantee. """ # pytype: skip-file # ruff: noqa: UP006 import datetime import decimal import logging from typing import Any from typing import ByteString from typing import Dict from typing import Generic from typing import Iterable from typing import List from typing import Mapping from typing import NamedTuple from typing import Optional from typing import Sequence from typing import Tuple from typing import TypeVar from typing import Union import numpy as np from google.protobuf import text_format from apache_beam.portability import common_urns from apache_beam.portability.api import schema_pb2 from apache_beam.typehints import row_type from apache_beam.typehints import typehints from apache_beam.typehints.native_type_compatibility import _get_args from apache_beam.typehints.native_type_compatibility import _match_is_exactly_mapping from apache_beam.typehints.native_type_compatibility import _match_is_optional from apache_beam.typehints.native_type_compatibility import _safe_issubclass from apache_beam.typehints.native_type_compatibility import convert_to_python_type from apache_beam.typehints.native_type_compatibility import extract_optional_type from apache_beam.typehints.native_type_compatibility import match_dataclass_for_row from apache_beam.typehints.native_type_compatibility import match_is_named_tuple from apache_beam.typehints.schema_registry import SCHEMA_REGISTRY from apache_beam.typehints.schema_registry import SchemaTypeRegistry from apache_beam.utils import proto_utils from apache_beam.utils.python_callable import PythonCallableWithSource from apache_beam.utils.timestamp import Timestamp PYTHON_ANY_URN = "beam:logical:pythonsdk_any:v1" _PYTHON_ANY_FIELD_TYPE_BYTE = "_pythonsdk_any_type_byte" _PYTHON_ANY_FIELD_PAYLOAD = "payload" _SCHEMA_OPTION_STATIC_ENCODING = "beam:option:row:static_encoding" # Bi-directional mappings _PRIMITIVES = ( (np.int8, schema_pb2.BYTE), (np.int16, schema_pb2.INT16), (np.int32, schema_pb2.INT32), (np.int64, schema_pb2.INT64), (np.float32, schema_pb2.FLOAT), (np.float64, schema_pb2.DOUBLE), (str, schema_pb2.STRING), (bool, schema_pb2.BOOLEAN), (bytes, schema_pb2.BYTES), ) PRIMITIVE_TO_ATOMIC_TYPE = dict((typ, atomic) for typ, atomic in _PRIMITIVES) ATOMIC_TYPE_TO_PRIMITIVE = dict((atomic, typ) for typ, atomic in _PRIMITIVES) # One-way mappings PRIMITIVE_TO_ATOMIC_TYPE.update({ ByteString: schema_pb2.BYTES, # Allow users to specify a native int, and use INT64 as the cross-language # representation. Technically ints have unlimited precision, but RowCoder # should throw an error if it sees one with a bit width > 64 when encoding. int: schema_pb2.INT64, float: schema_pb2.DOUBLE, }) _LOGGER = logging.getLogger(__name__) # Serialized schema_pb2.Schema w/o id to id. _SCHEMA_ID_CACHE = {} def named_fields_to_schema( names_and_types: Union[Dict[str, type], Sequence[Tuple[str, type]]], schema_id: Optional[str] = None, schema_options: Optional[Sequence[Tuple[str, Any]]] = None, field_options: Optional[Dict[str, Sequence[Tuple[str, Any]]]] = None, schema_registry: SchemaTypeRegistry = SCHEMA_REGISTRY, field_descriptions: Optional[Dict[str, str]] = None, ): schema_options = schema_options or [] field_options = field_options or {} field_descriptions = field_descriptions or {} if isinstance(names_and_types, dict): names_and_types = names_and_types.items() _, cached_schema = schema_registry.by_id.get(schema_id, (None, None)) if cached_schema: type_by_name_from_schema = { field.name: field.type for field in cached_schema.fields } else: type_by_name_from_schema = {} schema = schema_pb2.Schema( fields=[ schema_pb2.Field( name=name, type=type_by_name_from_schema.get( name, typing_to_runner_api(type)), options=[ option_to_runner_api(option_tuple) for option_tuple in field_options.get(name, []) ], description=field_descriptions.get(name, None)) for (name, type) in names_and_types ], options=[ option_to_runner_api(option_tuple) for option_tuple in schema_options ]) if schema_id is None: key = schema.SerializeToString() if key not in _SCHEMA_ID_CACHE: _SCHEMA_ID_CACHE[key] = schema_registry.generate_new_id() schema_id = _SCHEMA_ID_CACHE[key] schema.id = schema_id return schema def named_fields_from_schema( schema): # (schema_pb2.Schema) -> typing.List[typing.Tuple[str, type]] return [(field.name, typing_from_runner_api(field.type)) for field in schema.fields] def typing_to_runner_api( type_: type, schema_registry: SchemaTypeRegistry = SCHEMA_REGISTRY ) -> schema_pb2.FieldType: return SchemaTranslation( schema_registry=schema_registry).typing_to_runner_api(type_) def typing_from_runner_api( fieldtype_proto: schema_pb2.FieldType, schema_registry: SchemaTypeRegistry = SCHEMA_REGISTRY) -> type: return SchemaTranslation( schema_registry=schema_registry).typing_from_runner_api(fieldtype_proto) def value_to_runner_api( type_proto: schema_pb2.FieldType, value, schema_registry: SchemaTypeRegistry = SCHEMA_REGISTRY ) -> schema_pb2.FieldValue: return SchemaTranslation(schema_registry=schema_registry).value_to_runner_api( type_proto, value) def value_from_runner_api( type_proto: schema_pb2.FieldType, value_proto: schema_pb2.FieldValue, schema_registry: SchemaTypeRegistry = SCHEMA_REGISTRY ) -> schema_pb2.FieldValue: return SchemaTranslation( schema_registry=schema_registry).value_from_runner_api( type_proto, value_proto) def option_to_runner_api( option: Tuple[str, Any], schema_registry: SchemaTypeRegistry = SCHEMA_REGISTRY) -> schema_pb2.Option: return SchemaTranslation( schema_registry=schema_registry).option_to_runner_api(option) def option_from_runner_api( option_proto: schema_pb2.Option, schema_registry: SchemaTypeRegistry = SCHEMA_REGISTRY) -> Tuple[str, Any]: return SchemaTranslation( schema_registry=schema_registry).option_from_runner_api(option_proto) def schema_field( name: str, field_type: Union[schema_pb2.FieldType, type], description: Optional[str] = None) -> schema_pb2.Field: return schema_pb2.Field( name=name, type=field_type if isinstance(field_type, schema_pb2.FieldType) else typing_to_runner_api(field_type), description=description) def _python_any_schema_pb2(has_repr): # A portable schema matches FastPrimitivesCoder encoded values if has_repr: representation = schema_pb2.FieldType( nullable=False, row_type=schema_pb2.RowType( schema=schema_pb2.Schema( fields=[ schema_pb2.Field( name=_PYTHON_ANY_FIELD_TYPE_BYTE, type=schema_pb2.FieldType( atomic_type=schema_pb2.BYTE, nullable=False)), schema_pb2.Field( name=_PYTHON_ANY_FIELD_PAYLOAD, type=schema_pb2.FieldType( atomic_type=schema_pb2.BYTES, nullable=False)) ], options=[ schema_pb2.Option( name=_SCHEMA_OPTION_STATIC_ENCODING, type=schema_pb2.FieldType( atomic_type=schema_pb2.BOOLEAN), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue( boolean=True))) ]))) if has_repr else None else: representation = None return schema_pb2.FieldType( logical_type=schema_pb2.LogicalType( urn=PYTHON_ANY_URN, representation=representation), nullable=True) class SchemaTranslation(object): def __init__(self, schema_registry: SchemaTypeRegistry = SCHEMA_REGISTRY): self.schema_registry = schema_registry def typing_to_runner_api(self, type_: type) -> schema_pb2.FieldType: if isinstance(type_, schema_pb2.Schema): return schema_pb2.FieldType(row_type=schema_pb2.RowType(schema=type_)) if hasattr(type_, '_beam_schema_proto') and type_._beam_schema_proto.obj: return schema_pb2.FieldType( row_type=schema_pb2.RowType(schema=type_._beam_schema_proto.obj)) if isinstance(type_, row_type.RowTypeConstraint): if type_.schema_id is None: schema_id = SCHEMA_REGISTRY.generate_new_id() type_.set_schema_id(schema_id) schema = None else: schema_id = type_.schema_id schema = self.schema_registry.get_schema_by_id(schema_id) if schema is None: # Either user_type was not annotated with a schema id, or there was # no schema in the registry with the id. The latter should only happen # in tests. # Either way, we need to generate a new schema proto. schema = schema_pb2.Schema( fields=[ schema_pb2.Field( name=field_name, type=self.typing_to_runner_api(field_type), options=[ self.option_to_runner_api(option_tuple) for option_tuple in type_.field_options(field_name) ], description=type_._field_descriptions.get(field_name, None), ) for (field_name, field_type) in type_._fields ], id=schema_id, options=[ self.option_to_runner_api(option_tuple) for option_tuple in type_.schema_options ], ) self.schema_registry.add(type_.user_type, schema) return schema_pb2.FieldType(row_type=schema_pb2.RowType(schema=schema)) else: # See if this is coercible to a RowTypeConstraint (e.g. a NamedTuple or # dataclass) row_type_constraint = row_type.RowTypeConstraint.from_user_type(type_) if row_type_constraint is not None: return self.typing_to_runner_api(row_type_constraint) if isinstance(type_, typehints.TypeConstraint): type_ = convert_to_python_type(type_) # All concrete types (other than NamedTuple sub-classes) should map to # a supported primitive type. if type_ in PRIMITIVE_TO_ATOMIC_TYPE: return schema_pb2.FieldType(atomic_type=PRIMITIVE_TO_ATOMIC_TYPE[type_]) elif _match_is_exactly_mapping(type_): key_type, value_type = map(self.typing_to_runner_api, _get_args(type_)) return schema_pb2.FieldType( map_type=schema_pb2.MapType(key_type=key_type, value_type=value_type)) elif _match_is_optional(type_): # It's possible that a user passes us Optional[Optional[T]], but in python # typing this is indistinguishable from Optional[T] - both resolve to # Union[T, None] - so there's no need to check for that case here. result = self.typing_to_runner_api(extract_optional_type(type_)) result.nullable = True return result elif type_ == range: return schema_pb2.FieldType( array_type=schema_pb2.ArrayType( element_type=schema_pb2.FieldType( atomic_type=PRIMITIVE_TO_ATOMIC_TYPE[int]))) elif _safe_issubclass(type_, Sequence) and not _safe_issubclass(type_, str): arg_types = _get_args(type_) if len(arg_types) > 0: element_type = self.typing_to_runner_api(arg_types[0]) return schema_pb2.FieldType( array_type=schema_pb2.ArrayType(element_type=element_type)) elif _safe_issubclass(type_, Mapping): key_type, value_type = map(self.typing_to_runner_api, _get_args(type_)) return schema_pb2.FieldType( map_type=schema_pb2.MapType(key_type=key_type, value_type=value_type)) elif _safe_issubclass(type_, Iterable) and not _safe_issubclass(type_, str): arg_types = _get_args(type_) if len(arg_types) > 0: element_type = self.typing_to_runner_api(arg_types[0]) return schema_pb2.FieldType( array_type=schema_pb2.ArrayType(element_type=element_type)) elif type_ == Any: return _python_any_schema_pb2(has_repr=False) try: if LogicalType.is_known_logical_type(type_): logical_type = type_ else: logical_type = LogicalType.from_typing(type_) except ValueError: # Unknown type, use pythonsdk_any with a representation compatible with # FastPrimitiveCoder encoded UNKNOWN type return _python_any_schema_pb2(has_repr=True) else: argument_type = None argument = None if logical_type.argument_type() is not None: argument_type = self.typing_to_runner_api(logical_type.argument_type()) try: argument = self.value_to_runner_api( argument_type, logical_type.argument()) except ValueError: # TODO(https://github.com/apache/beam/issues/23373): Complete support # for logical types that require arguments beyond atomic type. # For now, skip arguments. argument = None return schema_pb2.FieldType( logical_type=schema_pb2.LogicalType( urn=logical_type.urn(), representation=self.typing_to_runner_api( logical_type.representation_type()), argument_type=argument_type, argument=argument)) def atomic_value_from_runner_api( self, atomic_type: schema_pb2.AtomicType, atomic_value: schema_pb2.AtomicTypeValue): if atomic_type == schema_pb2.BYTE: value = np.int8(atomic_value.byte) elif atomic_type == schema_pb2.INT16: value = np.int16(atomic_value.int16) elif atomic_type == schema_pb2.INT32: value = np.int32(atomic_value.int32) elif atomic_type == schema_pb2.INT64: value = np.int64(atomic_value.int64) elif atomic_type == schema_pb2.FLOAT: value = np.float32(atomic_value.float) elif atomic_type == schema_pb2.DOUBLE: value = np.float64(atomic_value.double) elif atomic_type == schema_pb2.STRING: value = atomic_value.string elif atomic_type == schema_pb2.BOOLEAN: value = atomic_value.boolean elif atomic_type == schema_pb2.BYTES: value = atomic_value.bytes else: raise ValueError( f"Unrecognized atomic_type ({atomic_type}) " f"when decoding value {atomic_value!r}") return value def atomic_value_to_runner_api( self, atomic_type: schema_pb2.AtomicType, value) -> schema_pb2.AtomicTypeValue: if atomic_type == schema_pb2.BYTE: atomic_value = schema_pb2.AtomicTypeValue(byte=value) elif atomic_type == schema_pb2.INT16: atomic_value = schema_pb2.AtomicTypeValue(int16=value) elif atomic_type == schema_pb2.INT32: atomic_value = schema_pb2.AtomicTypeValue(int32=value) elif atomic_type == schema_pb2.INT64: atomic_value = schema_pb2.AtomicTypeValue(int64=value) elif atomic_type == schema_pb2.FLOAT: atomic_value = schema_pb2.AtomicTypeValue(float=value) elif atomic_type == schema_pb2.DOUBLE: atomic_value = schema_pb2.AtomicTypeValue(double=value) elif atomic_type == schema_pb2.STRING: atomic_value = schema_pb2.AtomicTypeValue(string=value) elif atomic_type == schema_pb2.BOOLEAN: atomic_value = schema_pb2.AtomicTypeValue(boolean=value) elif atomic_type == schema_pb2.BYTES: atomic_value = schema_pb2.AtomicTypeValue(bytes=value) else: raise ValueError( "Unrecognized atomic_type {atomic_type} when encoding value {value}") return atomic_value def value_from_runner_api( self, type_proto: schema_pb2.FieldType, value_proto: schema_pb2.FieldValue): if type_proto.WhichOneof("type_info") != "atomic_type": # TODO: Allow other value types raise ValueError( "Encounterd option with unsupported type. Only " f"atomic_type options are supported: {type_proto}") value = self.atomic_value_from_runner_api( type_proto.atomic_type, value_proto.atomic_value) return value def value_to_runner_api(self, typing_proto: schema_pb2.FieldType, value): if typing_proto.WhichOneof("type_info") != "atomic_type": # TODO: Allow other value types raise ValueError( "Only atomic_type option values are currently supported in Python. " f"Got {value!r}, which maps to fieldtype {typing_proto!r}.") atomic_value = self.atomic_value_to_runner_api( typing_proto.atomic_type, value) value_proto = schema_pb2.FieldValue(atomic_value=atomic_value) return value_proto def option_from_runner_api( self, option_proto: schema_pb2.Option) -> Tuple[str, Any]: if not option_proto.HasField('type'): return option_proto.name, None value = self.value_from_runner_api(option_proto.type, option_proto.value) return option_proto.name, value def option_to_runner_api(self, option: Tuple[str, Any]) -> schema_pb2.Option: name, value = option if value is None: # a value of None indicates the option is just a flag. # Don't set type, value return schema_pb2.Option(name=name) type_proto = self.typing_to_runner_api(type(value)) value_proto = self.value_to_runner_api(type_proto, value) return schema_pb2.Option(name=name, type=type_proto, value=value_proto) def typing_from_runner_api( self, fieldtype_proto: schema_pb2.FieldType) -> type: if fieldtype_proto.nullable: # In order to determine the inner type, create a copy of fieldtype_proto # with nullable=False and pass back to typing_from_runner_api base_type = schema_pb2.FieldType() base_type.CopyFrom(fieldtype_proto) base_type.nullable = False base = self.typing_from_runner_api(base_type) if base == Any: return base else: return Optional[base] type_info = fieldtype_proto.WhichOneof("type_info") if type_info == "atomic_type": try: return ATOMIC_TYPE_TO_PRIMITIVE[fieldtype_proto.atomic_type] except KeyError: raise ValueError( "Unsupported atomic type: {0}".format(fieldtype_proto.atomic_type)) elif type_info == "array_type": return Sequence[self.typing_from_runner_api( fieldtype_proto.array_type.element_type)] elif type_info == "map_type": return Mapping[ self.typing_from_runner_api(fieldtype_proto.map_type.key_type), self.typing_from_runner_api(fieldtype_proto.map_type.value_type)] elif type_info == "row_type": schema = fieldtype_proto.row_type.schema schema_options = [ self.option_from_runner_api(option_proto) for option_proto in schema.options ] field_options = { field.name: [ self.option_from_runner_api(option_proto) for option_proto in field.options ] for field in schema.fields if field.options } # First look for user type in the registry user_type = self.schema_registry.get_typing_by_id(schema.id) if user_type is None: # If not in SDK options (the coder likely came from another SDK), # generate a NamedTuple type to use. fields = named_fields_from_schema(schema) descriptions = { field.name: field.description for field in schema.fields } result = row_type.RowTypeConstraint.from_fields( fields=fields, schema_id=schema.id, schema_options=schema_options, field_options=field_options, schema_registry=self.schema_registry, field_descriptions=descriptions or None) return result else: return row_type.RowTypeConstraint.from_user_type( user_type, schema_options=schema_options, field_options=field_options) elif type_info == "logical_type": if fieldtype_proto.logical_type.urn == PYTHON_ANY_URN: return Any else: return LogicalType.from_runner_api( fieldtype_proto.logical_type).language_type() elif type_info == "iterable_type": return Sequence[self.typing_from_runner_api( fieldtype_proto.iterable_type.element_type)] else: raise ValueError(f"Unrecognized type_info: {type_info!r}") def named_tuple_from_schema(self, schema: schema_pb2.Schema) -> type: from apache_beam import coders type_name = 'BeamSchema_{}'.format(schema.id.replace('-', '_')) subfields = [] descriptions = {} for field in schema.fields: try: field_py_type = self.typing_from_runner_api(field.type) if isinstance(field_py_type, row_type.RowTypeConstraint): field_py_type = field_py_type.user_type except ValueError as e: raise ValueError( "Failed to decode schema due to an issue with Field proto:\n\n" f"{text_format.MessageToString(field)}") from e descriptions[field.name] = field.description subfields.append((field.name, field_py_type)) if schema.id in self.schema_registry.by_id: user_type = self.schema_registry.by_id[schema.id][0] else: user_type = NamedTuple(type_name, subfields) # Define a reduce function, otherwise these types can't be pickled # (See BEAM-9574) setattr( user_type, '__reduce__', _named_tuple_reduce_method(schema.SerializeToString())) setattr(user_type, "_field_descriptions", descriptions) setattr(user_type, row_type._BEAM_SCHEMA_ID, schema.id) self.schema_registry.add(user_type, schema) coders.registry.register_coder(user_type, coders.RowCoder) return user_type def _named_tuple_reduce_method(serialized_schema): def __reduce__(self): return _hydrate_namedtuple_instance, (serialized_schema, tuple(self)) return __reduce__ def _hydrate_namedtuple_instance(encoded_schema, values): return named_tuple_from_schema( proto_utils.parse_Bytes(encoded_schema, schema_pb2.Schema))(*values) def named_tuple_from_schema( schema, schema_registry: SchemaTypeRegistry = SCHEMA_REGISTRY) -> type: return SchemaTranslation( schema_registry=schema_registry).named_tuple_from_schema(schema) def named_tuple_to_schema( named_tuple, schema_registry: SchemaTypeRegistry = SCHEMA_REGISTRY) -> schema_pb2.Schema: return typing_to_runner_api(named_tuple, schema_registry).row_type.schema def schema_from_element_type(element_type: type) -> schema_pb2.Schema: """Get a schema for the given PCollection element_type. Returns schema as a list of (name, python_type) tuples""" if isinstance(element_type, row_type.RowTypeConstraint): return named_fields_to_schema(element_type._fields) elif match_is_named_tuple(element_type) or match_dataclass_for_row( element_type): # schema id does not inherit from base classes if row_type._BEAM_SCHEMA_ID in element_type.__dict__: # if the named tuple's schema is in registry, we just use it instead of # regenerating one. schema_id = getattr(element_type, row_type._BEAM_SCHEMA_ID) schema = SCHEMA_REGISTRY.get_schema_by_id(schema_id) if schema is not None: return schema return named_tuple_to_schema(element_type) else: raise TypeError( f"Could not determine schema for type hint {element_type!r}. Did you " "mean to create a schema-aware PCollection? See " "https://s.apache.org/beam-python-schemas") def named_fields_from_element_type( element_type: type) -> List[Tuple[str, type]]: return named_fields_from_schema(schema_from_element_type(element_type)) def union_schema_type(element_types): """Returns a schema whose fields are the union of each corresponding field. element_types must be a set of schema-aware types whose fields have the same naming and ordering. """ named_fields_and_types = [] for t in element_types: n = named_fields_from_element_type(t) if named_fields_and_types and len(named_fields_and_types[-1]) != len(n): raise TypeError("element types has different number of fields") named_fields_and_types.append(n) union_fields_and_types = [] for field in zip(*named_fields_and_types): names, types = zip(*field) name_set = set(names) if len(name_set) != 1: raise TypeError( f"Could not determine schema for type hints {element_types!r}: " f"Inconsistent names: {name_set}") union_fields_and_types.append( (next(iter(name_set)), typehints.Union[types])) return named_tuple_from_schema(named_fields_to_schema(union_fields_and_types)) class _Ephemeral: """Helper class for wrapping unpicklable objects.""" def __init__(self, obj): self.obj = obj def __reduce__(self): return _Ephemeral, (None, ) # Registry of typings for a schema by UUID class LogicalTypeRegistry(object): def __init__(self): self.by_urn = {} self.by_logical_type = {} self.by_language_type = {} self._custom_urns = set() def _add_internal(self, urn, logical_type): self.by_urn[urn] = logical_type self.by_logical_type[logical_type] = urn self.by_language_type[logical_type.language_type()] = logical_type def add(self, urn, logical_type): self._add_internal(urn, logical_type) self._custom_urns.add(urn) def get_logical_type_by_urn(self, urn): return self.by_urn.get(urn, None) def get_urn_by_logical_type(self, logical_type): return self.by_logical_type.get(logical_type, None) def get_logical_type_by_language_type(self, representation_type): return self.by_language_type.get(representation_type, None) def copy(self): copy = LogicalTypeRegistry() copy.by_urn.update(self.by_urn) copy.by_logical_type.update(self.by_logical_type) copy.by_language_type.update(self.by_language_type) copy._custom_urns.update(self._custom_urns) return copy def copy_custom(self): copy = LogicalTypeRegistry() for urn in self._custom_urns: logical_type = self.by_urn[urn] copy.by_urn[urn] = logical_type copy.by_logical_type[logical_type] = urn copy.by_language_type[logical_type.language_type()] = logical_type copy._custom_urns.add(urn) return copy def load(self, another): self.by_urn.update(another.by_urn) self.by_logical_type.update(another.by_logical_type) self.by_language_type.update(another.by_language_type) self._custom_urns.update(another._custom_urns) LanguageT = TypeVar('LanguageT') RepresentationT = TypeVar('RepresentationT') ArgT = TypeVar('ArgT') class LogicalType(Generic[LanguageT, RepresentationT, ArgT]): _known_logical_types = LogicalTypeRegistry() @classmethod def urn(cls): # type: () -> str """Return the URN used to identify this logical type""" raise NotImplementedError() @classmethod def language_type(cls): # type: () -> type """Return the language type this LogicalType encodes. The returned type should match LanguageT""" raise NotImplementedError() @classmethod def representation_type(cls): # type: () -> type """Return the type of the representation this LogicalType uses to encode the language type. The returned type should match RepresentationT""" raise NotImplementedError() @classmethod def argument_type(cls): # type: () -> type """Return the type of the argument used for variations of this LogicalType. The returned type should match ArgT""" raise NotImplementedError(cls) def argument(self): # type: () -> ArgT """Return the argument for this instance of the LogicalType.""" raise NotImplementedError() def to_representation_type(self, value): # type: (LanguageT) -> RepresentationT """Convert an instance of LanguageT to RepresentationT.""" raise NotImplementedError() def to_language_type(self, value): # type: (RepresentationT) -> LanguageT """Convert an instance of RepresentationT to LanguageT.""" raise NotImplementedError() @classmethod def _register_internal(cls, logical_type_cls): """ Register an implementation of LogicalType. The types registered using this decorator are not pickled on pipeline submission, as it relies module import to be registered on worker initialization. Should be used within schemas module and static context. """ cls._known_logical_types._add_internal( logical_type_cls.urn(), logical_type_cls) return logical_type_cls @classmethod def register_logical_type(cls, logical_type_cls): """Register an implementation of LogicalType.""" cls._known_logical_types.add(logical_type_cls.urn(), logical_type_cls) return logical_type_cls @classmethod def from_typing(cls, typ): # type: (type) -> LogicalType """Construct an instance of a registered LogicalType implementation given a typing. Raises ValueError if no registered LogicalType implementation can encode the given typing.""" logical_type = cls._known_logical_types.get_logical_type_by_language_type( typ) if logical_type is None: raise ValueError("No logical type registered for typing '%s'" % typ) return logical_type._from_typing(typ) @classmethod def _from_typing(cls, typ): # type: (type) -> LogicalType """Construct an instance of this LogicalType implementation given a typing. """ raise NotImplementedError() @classmethod def from_runner_api(cls, logical_type_proto): # type: (schema_pb2.LogicalType) -> LogicalType """Construct an instance of a registered LogicalType implementation given a proto LogicalType. Raises ValueError if no LogicalType registered for the given URN. """ logical_type = cls._known_logical_types.get_logical_type_by_urn( logical_type_proto.urn) if logical_type is None: raise ValueError( "No logical type registered for URN '%s'" % logical_type_proto.urn) if not logical_type_proto.HasField( "argument_type") or not logical_type_proto.HasField("argument"): # logical type_proto without argument return logical_type() else: try: argument = value_from_runner_api( logical_type_proto.argument_type, logical_type_proto.argument) except ValueError: # TODO(https://github.com/apache/beam/issues/23373): Complete support # for logical types that require arguments beyond atomic type. # For now, skip arguments. _LOGGER.warning( 'Logical type %s with argument is currently unsupported. ' 'Argument values are omitted', logical_type_proto.urn) return logical_type() return logical_type(argument) @classmethod def is_known_logical_type(cls, logical_type): return cls._known_logical_types.get_urn_by_logical_type( logical_type) is not None class NoArgumentLogicalType(LogicalType[LanguageT, RepresentationT, None]): @classmethod def argument_type(cls): # type: () -> type return None def argument(self): # type: () -> ArgT return None @classmethod def _from_typing(cls, typ): # type: (type) -> LogicalType # Since there's no argument, there can be no additional information encoded # in the typing. Just construct an instance. return cls() class PassThroughLogicalType(LogicalType[LanguageT, LanguageT, ArgT]): """A base class for LogicalTypes that use the same type as the underlying representation type. """ def to_language_type(self, value): return value @classmethod def representation_type(cls): # type: () -> type return cls.language_type() def to_representation_type(self, value): return value @classmethod def _from_typing(cls, typ): # type: (type) -> LogicalType # TODO(https://github.com/apache/beam/issues/23373): enable argument return cls() MicrosInstantRepresentation = NamedTuple( 'MicrosInstantRepresentation', [('seconds', np.int64), ('micros', np.int64)]) @LogicalType._register_internal class MillisInstant(NoArgumentLogicalType[Timestamp, np.int64]): """Millisecond-precision instant logical type handles values consistent with that encoded by ``InstantCoder`` in the Java SDK. This class handles :class:`apache_beam.utils.timestamp.Timestamp` language type as :class:`MicrosInstant`, but it only provides millisecond precision, because it is aimed to handle data encoded by Java sdk's InstantCoder which has same precision level. Timestamp is handled by `MicrosInstant` by default. In some scenario, such as read from cross-language transform with rows containing InstantCoder encoded timestamps, one may need to override the mapping of Timetamp to MillisInstant. To do this, re-register this class with :func:`~LogicalType.register_logical_type`. """ @classmethod def representation_type(cls): # type: () -> type return np.int64 @classmethod def urn(cls): return common_urns.millis_instant.urn @classmethod def language_type(cls): return Timestamp def to_language_type(self, value): # type: (np.int64) -> Timestamp # value shifted as in apache_beams.coders.coder_impl.TimestampCoderImpl if value < 0: millis = int(value) + (1 << 63) else: millis = int(value) - (1 << 63) return Timestamp(micros=millis * 1000) # Make sure MicrosInstant is registered after MillisInstant so that it # overwrites the mapping of Timestamp language type representation choice and # thus does not lose microsecond precision inside python sdk. @LogicalType._register_internal class MicrosInstant(NoArgumentLogicalType[Timestamp, MicrosInstantRepresentation]): """Microsecond-precision instant logical type that handles ``Timestamp``.""" @classmethod def urn(cls): return common_urns.micros_instant.urn @classmethod def representation_type(cls): # type: () -> type return MicrosInstantRepresentation @classmethod def language_type(cls): return Timestamp def to_representation_type(self, value): # type: (Timestamp) -> MicrosInstantRepresentation return MicrosInstantRepresentation( value.micros // 1000000, value.micros % 1000000) def to_language_type(self, value): # type: (MicrosInstantRepresentation) -> Timestamp return Timestamp(seconds=int(value.seconds), micros=int(value.micros)) @LogicalType._register_internal class Date(NoArgumentLogicalType[datetime.date, np.int64]): """Date logical type that handles ``datetime.date``, days since epoch.""" EPOCH = datetime.date(1970, 1, 1) @classmethod def urn(cls): return common_urns.date.urn @classmethod def representation_type(cls): # type: () -> type return np.int64 @classmethod def language_type(cls): return datetime.date def to_representation_type(self, value): # type: (datetime.date) -> np.int64 return (value - self.EPOCH).days def to_language_type(self, value): # type: (np.int64) -> datetime.date return self.EPOCH + datetime.timedelta(days=value) @LogicalType._register_internal class PythonCallable(NoArgumentLogicalType[PythonCallableWithSource, str]): """A logical type for PythonCallableSource objects.""" @classmethod def urn(cls): return common_urns.python_callable.urn @classmethod def representation_type(cls): # type: () -> type return str @classmethod def language_type(cls): return PythonCallableWithSource def to_representation_type(self, value): # type: (PythonCallableWithSource) -> str return value.get_source() def to_language_type(self, value): # type: (str) -> PythonCallableWithSource return PythonCallableWithSource(value) FixedPrecisionDecimalArgumentRepresentation = NamedTuple( 'FixedPrecisionDecimalArgumentRepresentation', [('precision', np.int32), ('scale', np.int32)]) class DecimalLogicalType(NoArgumentLogicalType[decimal.Decimal, bytes]): """A logical type for decimal objects handling values consistent with that encoded by ``BigDecimalCoder`` in the Java SDK. """ @classmethod def urn(cls): return common_urns.decimal.urn @classmethod def representation_type(cls): # type: () -> type return bytes @classmethod def language_type(cls): return decimal.Decimal def to_representation_type(self, value): # type: (decimal.Decimal) -> bytes return str(value).encode() def to_language_type(self, value): # type: (bytes) -> decimal.Decimal return decimal.Decimal(value.decode()) @LogicalType._register_internal class FixedPrecisionDecimalLogicalType( LogicalType[decimal.Decimal, DecimalLogicalType, FixedPrecisionDecimalArgumentRepresentation]): """A wrapper of DecimalLogicalType that contains the precision value. """ def __init__(self, precision=-1, scale=0): self.precision = precision self.scale = scale @classmethod def urn(cls): # TODO(https://github.com/apache/beam/issues/23373) promote this URN to # schema.proto once logical types with argument are fully supported and the # implementation of this logical type can thus be considered standardized. return "beam:logical_type:fixed_decimal:v1" @classmethod def representation_type(cls): # type: () -> type return DecimalLogicalType @classmethod def language_type(cls): return decimal.Decimal # from language type (decimal.Decimal) to representation type # (the type corresponding to the coder used in DecimalLogicalType) def to_representation_type(self, value: decimal.Decimal) -> decimal.Decimal: return value # from representation type (the type corresponding to the coder used in # DecimalLogicalType) to language type def to_language_type(self, value: decimal.Decimal) -> decimal.Decimal: return value @classmethod def argument_type(cls): return FixedPrecisionDecimalArgumentRepresentation def argument(self): return FixedPrecisionDecimalArgumentRepresentation( precision=self.precision, scale=self.scale) @classmethod def _from_typing(cls, typ): return cls() # TODO(yathu,BEAM-10722): Investigate and resolve conflicts in logical type # registration when more than one logical types sharing the same language type LogicalType._register_internal(DecimalLogicalType) @LogicalType._register_internal class FixedBytes(PassThroughLogicalType[bytes, np.int32]): """A logical type for fixed-length bytes.""" @classmethod def urn(cls): return common_urns.fixed_bytes.urn def __init__(self, length: np.int32): self.length = length @classmethod def language_type(cls) -> type: return bytes def to_language_type(self, value: bytes): length = len(value) if length > self.length: raise ValueError( "value length {} > allowed length {}".format(length, self.length)) elif length < self.length: # padding at the end value = value + b'\0' * (self.length - length) return value @classmethod def argument_type(cls): return np.int32 def argument(self): return self.length @LogicalType._register_internal class VariableBytes(PassThroughLogicalType[bytes, np.int32]): """A logical type for variable-length bytes with specified maximum length.""" @classmethod def urn(cls): return common_urns.var_bytes.urn def __init__(self, max_length: np.int32 = np.iinfo(np.int32).max): self.max_length = max_length @classmethod def language_type(cls) -> type: return bytes def to_language_type(self, value: bytes): length = len(value) if length > self.max_length: raise ValueError( "value length {} > allowed length {}".format(length, self.max_length)) return value @classmethod def argument_type(cls): return np.int32 def argument(self): return self.max_length @LogicalType._register_internal class FixedString(PassThroughLogicalType[str, np.int32]): """A logical type for fixed-length string.""" @classmethod def urn(cls): return common_urns.fixed_char.urn def __init__(self, length: np.int32): self.length = length @classmethod def language_type(cls) -> type: return str def to_language_type(self, value: str): length = len(value) if length > self.length: raise ValueError( "value length {} > allowed length {}".format(length, self.length)) elif length < self.length: # padding at the end value = value + ' ' * (self.length - length) return value @classmethod def argument_type(cls): return np.int32 def argument(self): return self.length @LogicalType._register_internal class VariableString(PassThroughLogicalType[str, np.int32]): """A logical type for variable-length string with specified maximum length.""" @classmethod def urn(cls): return common_urns.var_char.urn def __init__(self, max_length: np.int32 = np.iinfo(np.int32).max): self.max_length = max_length @classmethod def language_type(cls) -> type: return str def to_language_type(self, value: str): length = len(value) if length > self.max_length: raise ValueError( "value length {} > allowed length {}".format(length, self.max_length)) return value @classmethod def argument_type(cls): return np.int32 def argument(self): return self.max_length # TODO: A temporary fix for missing jdbc logical types. # See the discussion in https://github.com/apache/beam/issues/35738 for # more detail. class JdbcDateType(LogicalType[datetime.date, MillisInstant, str]): """ For internal use only; no backwards-compatibility guarantees. Support of Legacy JdbcIO DATE logical type. Deemed to change when Java JDBCIO has been migrated to Beam portable logical types. """ def __init__(self, argument=""): pass @classmethod def representation_type(cls) -> type: return MillisInstant @classmethod def urn(cls): return "beam:logical_type:javasdk_date:v1" @classmethod def language_type(cls): return datetime.date def to_representation_type(self, value: datetime.date) -> Timestamp: return Timestamp.from_utc_datetime( datetime.datetime.combine( value, datetime.datetime.min.time(), tzinfo=datetime.timezone.utc)) def to_language_type(self, value: Timestamp) -> datetime.date: return value.to_utc_datetime().date() @classmethod def argument_type(cls): return str def argument(self): return "" @classmethod def _from_typing(cls, typ): return cls() @LogicalType._register_internal class JdbcTimeType(LogicalType[datetime.time, MillisInstant, str]): """ For internal use only; no backwards-compatibility guarantees. Support of Legacy JdbcIO TIME logical type. . Deemed to change when Java JDBCIO has been migrated to Beam portable logical types. """ def __init__(self, argument=""): pass @classmethod def representation_type(cls) -> type: return MillisInstant @classmethod def urn(cls): return "beam:logical_type:javasdk_time:v1" @classmethod def language_type(cls): return datetime.time def to_representation_type(self, value: datetime.time) -> Timestamp: return Timestamp.from_utc_datetime( datetime.datetime.combine( datetime.datetime.utcfromtimestamp(0), value, tzinfo=datetime.timezone.utc)) def to_language_type(self, value: Timestamp) -> datetime.time: return value.to_utc_datetime().time() @classmethod def argument_type(cls): return str def argument(self): return "" @classmethod def _from_typing(cls, typ): return cls() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/schemas_test.py0000644000175100017510000010241415206113033022743 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for schemas.""" # pytype: skip-file # ruff: noqa: UP006 import dataclasses import datetime import itertools import pickle import unittest from typing import Any from typing import ByteString from typing import List from typing import Mapping from typing import NamedTuple from typing import Optional from typing import Sequence import numpy as np import pytest from hypothesis import given from hypothesis import settings from parameterized import parameterized from parameterized import parameterized_class from apache_beam.internal.cloudpickle import cloudpickle from apache_beam.portability import common_urns from apache_beam.portability.api import schema_pb2 from apache_beam.typehints import row_type from apache_beam.typehints import schemas from apache_beam.typehints import typehints from apache_beam.typehints.native_type_compatibility import match_is_named_tuple from apache_beam.typehints.schemas import SchemaTypeRegistry from apache_beam.typehints.schemas import named_fields_from_element_type from apache_beam.typehints.schemas import named_tuple_from_schema from apache_beam.typehints.schemas import named_tuple_to_schema from apache_beam.typehints.schemas import typing_from_runner_api from apache_beam.typehints.schemas import typing_to_runner_api from apache_beam.typehints.testing.strategies import named_fields from apache_beam.utils.timestamp import Timestamp all_nonoptional_primitives = [ np.int8, np.int16, np.int32, np.int64, np.float32, np.float64, bool, bytes, str, ] all_optional_primitives = [Optional[typ] for typ in all_nonoptional_primitives] all_primitives = all_nonoptional_primitives + all_optional_primitives basic_array_types = [Sequence[typ] for typ in all_primitives] basic_map_types = [ Mapping[key_type, value_type] for key_type, value_type in itertools.product( all_primitives, all_primitives) ] class AllPrimitives(NamedTuple): field_int8: np.int8 field_int16: np.int16 field_int32: np.int32 field_int64: np.int64 field_float32: np.float32 field_float64: np.float64 field_bool: bool field_bytes: bytes field_str: str field_optional_int8: Optional[np.int8] field_optional_int16: Optional[np.int16] field_optional_int32: Optional[np.int32] field_optional_int64: Optional[np.int64] field_optional_float32: Optional[np.float32] field_optional_float64: Optional[np.float64] field_optional_bool: Optional[bool] field_optional_bytes: Optional[bytes] field_optional_str: Optional[str] class ComplexSchema(NamedTuple): id: np.int64 name: str optional_map: Optional[Mapping[str, Optional[np.float64]]] optional_array: Optional[Sequence[np.float32]] array_optional: Sequence[Optional[bool]] timestamp: Timestamp date: datetime.date def get_test_beam_fieldtype_protos(): all_nonoptional_primitives = [ schema_pb2.FieldType(atomic_type=typ) for typ in schema_pb2.AtomicType.values() if typ is not schema_pb2.UNSPECIFIED ] all_optional_primitives = [ schema_pb2.FieldType(nullable=True, atomic_type=typ) for typ in schema_pb2.AtomicType.values() if typ is not schema_pb2.UNSPECIFIED ] all_primitives = all_nonoptional_primitives + all_optional_primitives basic_array_types = [ schema_pb2.FieldType(array_type=schema_pb2.ArrayType(element_type=typ)) for typ in all_primitives ] basic_map_types = [ schema_pb2.FieldType( map_type=schema_pb2.MapType(key_type=key_type, value_type=value_type)) for key_type, value_type in itertools.product( all_primitives, all_primitives) ] selected_schemas = [ schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( id='32497414-85e8-46b7-9c90-9a9cc62fe390', fields=[ schema_pb2.Field(name='field%d' % i, type=typ) for i, typ in enumerate(all_primitives) ]))), schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( id='dead1637-3204-4bcb-acf8-99675f338600', fields=[ schema_pb2.Field( name='id', type=schema_pb2.FieldType( atomic_type=schema_pb2.INT64)), schema_pb2.Field( name='name', type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING)), schema_pb2.Field( name='optional_map', type=schema_pb2.FieldType( nullable=True, map_type=schema_pb2.MapType( key_type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING), value_type=schema_pb2.FieldType( atomic_type=schema_pb2.DOUBLE)))), schema_pb2.Field( name='optional_array', type=schema_pb2.FieldType( nullable=True, array_type=schema_pb2.ArrayType( element_type=schema_pb2.FieldType( atomic_type=schema_pb2.FLOAT)))), schema_pb2.Field( name='array_optional', type=schema_pb2.FieldType( array_type=schema_pb2.ArrayType( element_type=schema_pb2.FieldType( nullable=True, atomic_type=schema_pb2.BYTES)))), ]))), schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( id='a-schema-with-options', fields=[ schema_pb2.Field(name='field%d' % i, type=typ) for i, typ in enumerate(all_primitives) ], options=[ schema_pb2.Option(name='a_flag'), schema_pb2.Option( name='a_byte', type=schema_pb2.FieldType( atomic_type=schema_pb2.BYTE), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue( byte=127))), schema_pb2.Option( name='a_int16', type=schema_pb2.FieldType( atomic_type=schema_pb2.INT16), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue( int16=255))), schema_pb2.Option( name='a_int32', type=schema_pb2.FieldType( atomic_type=schema_pb2.INT32), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue( int32=255))), schema_pb2.Option( name='a_int64', type=schema_pb2.FieldType( atomic_type=schema_pb2.INT64), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue( int64=255))), schema_pb2.Option( name='a_float', type=schema_pb2.FieldType( atomic_type=schema_pb2.FLOAT), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue( float=3.14))), schema_pb2.Option( name='a_double', type=schema_pb2.FieldType( atomic_type=schema_pb2.DOUBLE), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue( double=2.718))), schema_pb2.Option( name='a_str', type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue( string='str'))), schema_pb2.Option( name='a_bool', type=schema_pb2.FieldType( atomic_type=schema_pb2.BOOLEAN), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue( boolean=True))), schema_pb2.Option( name='a_bytes', type=schema_pb2.FieldType( atomic_type=schema_pb2.BYTES), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue( bytes=b'bytes!'))), ]))), schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( id='a-schema-with-field-options', fields=[ schema_pb2.Field( name='field%d' % i, type=typ, options=[ schema_pb2.Option(name='a_flag'), schema_pb2.Option( name='a_str', type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue( string='str'))), ]) for i, typ in enumerate(all_primitives) ] + [ schema_pb2.Field( name='nested', type=schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( fields=[ schema_pb2.Field( name='nested_field', type=schema_pb2.FieldType( atomic_type=schema_pb2.INT64, ), options=[ schema_pb2.Option( name='a_nested_field_flag' ), ]), ], options=[ schema_pb2.Option( name='a_nested_schema_flag'), schema_pb2.Option( name='a_str', type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING ), value=schema_pb2.FieldValue( atomic_value=schema_pb2. AtomicTypeValue( string='str'))), ], ))), ), ]))), schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( id='a-schema-with-optional-nested-struct', fields=[ schema_pb2.Field( name='id', type=schema_pb2.FieldType( atomic_type=schema_pb2.INT64)), schema_pb2.Field( name='nested_row', type=schema_pb2.FieldType( nullable=True, row_type=schema_pb2.RowType( schema=schema_pb2.Schema( id='the-nested-schema', fields=[ schema_pb2.Field( name='name', type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING) ), schema_pb2.Field( name='optional_map', type=schema_pb2.FieldType( nullable=True, map_type=schema_pb2.MapType( key_type=schema_pb2. FieldType( atomic_type=schema_pb2 .STRING), value_type=schema_pb2. FieldType( atomic_type=schema_pb2 .DOUBLE)))), ])))) ]))), ] return all_primitives + \ basic_array_types + \ basic_map_types + \ selected_schemas def get_test_beam_schemas_protos(): return [ fieldtype.row_type.schema for fieldtype in get_test_beam_fieldtype_protos() if fieldtype.WhichOneof('type_info') == 'row_type' ] class SchemaTest(unittest.TestCase): """ Tests for Runner API Schema proto to/from typing conversions There are two main tests: test_typing_survives_proto_roundtrip, and test_proto_survives_typing_roundtrip. These are both necessary because Schemas are cached by ID, so performing just one of them wouldn't necessarily exercise all code paths. """ @parameterized.expand([(user_type,) for user_type in all_primitives + \ basic_array_types + \ basic_map_types] ) def test_typing_survives_proto_roundtrip(self, user_type): self.assertEqual( user_type, typing_from_runner_api( typing_to_runner_api( user_type, schema_registry=SchemaTypeRegistry()), schema_registry=SchemaTypeRegistry())) @parameterized.expand([(AllPrimitives, ), (ComplexSchema, )]) def test_namedtuple_roundtrip(self, user_type): roundtripped = typing_from_runner_api( typing_to_runner_api(user_type, schema_registry=SchemaTypeRegistry()), schema_registry=SchemaTypeRegistry()) self.assertIsInstance(roundtripped, row_type.RowTypeConstraint) self.assert_namedtuple_equivalent(roundtripped.user_type, user_type) def test_dataclass_roundtrip(self): @dataclasses.dataclass class SimpleDataclass: id: np.int64 name: str roundtripped = typing_from_runner_api( typing_to_runner_api( SimpleDataclass, schema_registry=SchemaTypeRegistry()), schema_registry=SchemaTypeRegistry()) self.assertIsInstance(roundtripped, row_type.RowTypeConstraint) # The roundtripped user_type is generated as a NamedTuple, so we can't test # equivalence directly with the dataclass. # Instead, let's verify annotations. self.assertEqual( roundtripped.user_type.__annotations__, SimpleDataclass.__annotations__) def test_row_type_constraint_to_schema(self): result_type = typing_to_runner_api( row_type.RowTypeConstraint.from_fields([ ('foo', np.int8), ('bar', float), ('baz', bytes), ])) self.assertIsInstance(result_type, schema_pb2.FieldType) self.assertEqual(result_type.WhichOneof("type_info"), "row_type") schema = result_type.row_type.schema self.assertIsNotNone(schema.id) expected = [ schema_pb2.Field( name='foo', type=schema_pb2.FieldType(atomic_type=schema_pb2.BYTE)), schema_pb2.Field( name='bar', type=schema_pb2.FieldType(atomic_type=schema_pb2.DOUBLE)), schema_pb2.Field( name='baz', type=schema_pb2.FieldType(atomic_type=schema_pb2.BYTES)), ] self.assertEqual(list(schema.fields), expected) def test_row_type_constraint_to_schema_with_options(self): row_type_with_options = row_type.RowTypeConstraint.from_fields( [ ('foo', np.int8), ('bar', float), ('baz', bytes), ], schema_options=[ ('some_metadata', 'foo'), ('some_other_metadata', 'baz'), ('some_metadata', 'bar'), ('an_integer_option', np.int32(123456)), ]) result_type = typing_to_runner_api(row_type_with_options) self.assertIsInstance(result_type, schema_pb2.FieldType) self.assertEqual(result_type.WhichOneof("type_info"), "row_type") schema = result_type.row_type.schema expected = [ schema_pb2.Option( name='some_metadata', type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue(string='foo'))), schema_pb2.Option( name='some_other_metadata', type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue(string='baz'))), schema_pb2.Option( name='some_metadata', type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue(string='bar'))), schema_pb2.Option( name='an_integer_option', type=schema_pb2.FieldType(atomic_type=schema_pb2.INT32), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue(int32=123456))), ] self.assertEqual(list(schema.options), expected) def test_row_type_constraint_to_schema_with_field_options(self): result_type = typing_to_runner_api( row_type.RowTypeConstraint.from_fields([ ('foo', np.int8), ('bar', float), ('baz', bytes), ], field_options={ 'foo': [ ('some_metadata', 123), ('some_flag', None) ] })) self.assertIsInstance(result_type, schema_pb2.FieldType) self.assertEqual(result_type.WhichOneof("type_info"), "row_type") field = result_type.row_type.schema.fields[0] expected = [ schema_pb2.Option( name='some_metadata', type=schema_pb2.FieldType(atomic_type=schema_pb2.INT64), value=schema_pb2.FieldValue( atomic_value=schema_pb2.AtomicTypeValue(int64=123)), ), schema_pb2.Option(name='some_flag') ] self.assertEqual(list(field.options), expected) def test_row_type_constraint_to_schema_with_field_descriptions(self): row_type_with_options = row_type.RowTypeConstraint.from_fields( [ ('foo', np.int8), ('bar', float), ('baz', bytes), ], field_descriptions={ 'foo': 'foo description', 'bar': 'bar description', 'baz': 'baz description', }) result_type = typing_to_runner_api(row_type_with_options) self.assertIsInstance(result_type, schema_pb2.FieldType) self.assertEqual(result_type.WhichOneof("type_info"), "row_type") fields = result_type.row_type.schema.fields expected = [ schema_pb2.Field( name='foo', description='foo description', type=schema_pb2.FieldType(atomic_type=schema_pb2.BYTE), ), schema_pb2.Field( name='bar', description='bar description', type=schema_pb2.FieldType(atomic_type=schema_pb2.DOUBLE), ), schema_pb2.Field( name='baz', description='baz description', type=schema_pb2.FieldType(atomic_type=schema_pb2.BYTES), ), ] self.assertEqual(list(fields), expected) def assert_namedtuple_equivalent(self, actual, expected): # Two types are only considered equal if they are literally the same # object (i.e. `actual == expected` is the same as `actual is expected` in # this case). # That's a much stricter check than we need, and it's necessarily not true # if types are pickled/unpickled. Here we just verify the features of the # types that actually matter to us. self.assertTrue(match_is_named_tuple(expected)) self.assertTrue(match_is_named_tuple(actual)) # TODO(https://github.com/apache/beam/issues/22082): This will break for # nested complex types. self.assertEqual(actual.__annotations__, expected.__annotations__) # TODO(https://github.com/apache/beam/issues/22082): Serialize user_type and # re-hydrate with sdk_options to make these checks pass. #self.assertEqual(dir(actual), dir(expected)) # #for attr in dir(expected): # self.assertEqual(getattr(actual, attr), getattr(expected, attr)) @parameterized.expand([ (fieldtype_proto, ) for fieldtype_proto in get_test_beam_fieldtype_protos() ]) def test_proto_survives_typing_roundtrip(self, fieldtype_proto): self.assertEqual( fieldtype_proto, typing_to_runner_api( typing_from_runner_api( fieldtype_proto, schema_registry=SchemaTypeRegistry()), schema_registry=SchemaTypeRegistry())) def test_any_maps_to_any(self): # python_any for typing.Any logical type's representation is delibrately set # absent to prevent the usage crossing language boundary, as its encoded # form isn't predictable from foreign SDK. self.assertEqual( typing_to_runner_api(Any), schemas._python_any_schema_pb2(has_repr=False)) def test_unknown_primitive_maps_to_any(self): self.assertEqual( typing_to_runner_api(np.uint32), schemas._python_any_schema_pb2(has_repr=True)) def test_unknown_user_type_maps_to_any(self): class MyUnknownType: pass self.assertEqual( typing_to_runner_api(MyUnknownType), schemas._python_any_schema_pb2(has_repr=True)) def test_unknown_atomic_raise_valueerror(self): self.assertRaises( ValueError, lambda: typing_from_runner_api( schema_pb2.FieldType(atomic_type=schema_pb2.UNSPECIFIED))) def test_int_maps_to_int64(self): self.assertEqual( schema_pb2.FieldType(atomic_type=schema_pb2.INT64), typing_to_runner_api(int)) def test_float_maps_to_float64(self): self.assertEqual( schema_pb2.FieldType(atomic_type=schema_pb2.DOUBLE), typing_to_runner_api(float)) def test_python_callable_maps_to_logical_type(self): from apache_beam.utils.python_callable import PythonCallableWithSource self.assertEqual( schema_pb2.FieldType( logical_type=schema_pb2.LogicalType( urn=common_urns.python_callable.urn, representation=typing_to_runner_api(str))), typing_to_runner_api(PythonCallableWithSource)) self.assertEqual( typing_from_runner_api( schema_pb2.FieldType( logical_type=schema_pb2.LogicalType( urn=common_urns.python_callable.urn, representation=typing_to_runner_api(str))), schema_registry=SchemaTypeRegistry()), PythonCallableWithSource) def test_trivial_example(self): MyCuteClass = NamedTuple( 'MyCuteClass', [ ('name', str), ('age', Optional[int]), ('interests', List[str]), ('height', float), ('blob', ByteString), ]) expected = schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( fields=[ schema_pb2.Field( name='name', type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING), ), schema_pb2.Field( name='age', type=schema_pb2.FieldType( nullable=True, atomic_type=schema_pb2.INT64)), schema_pb2.Field( name='interests', type=schema_pb2.FieldType( array_type=schema_pb2.ArrayType( element_type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING)))), schema_pb2.Field( name='height', type=schema_pb2.FieldType( atomic_type=schema_pb2.DOUBLE)), schema_pb2.Field( name='blob', type=schema_pb2.FieldType( atomic_type=schema_pb2.BYTES)), ]))) # Only test that the fields are equal. If we attempt to test the entire type # or the entire schema, the generated id will break equality. self.assertEqual( expected.row_type.schema.fields, typing_to_runner_api(MyCuteClass).row_type.schema.fields) def test_trivial_example_dataclass(self): @dataclasses.dataclass class MyCuteDataclass: name: str age: Optional[int] interests: List[str] height: float blob: ByteString expected = schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( fields=[ schema_pb2.Field( name='name', type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING), ), schema_pb2.Field( name='age', type=schema_pb2.FieldType( nullable=True, atomic_type=schema_pb2.INT64)), schema_pb2.Field( name='interests', type=schema_pb2.FieldType( array_type=schema_pb2.ArrayType( element_type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING)))), schema_pb2.Field( name='height', type=schema_pb2.FieldType( atomic_type=schema_pb2.DOUBLE)), schema_pb2.Field( name='blob', type=schema_pb2.FieldType( atomic_type=schema_pb2.BYTES)), ]))) self.assertEqual( expected.row_type.schema.fields, typing_to_runner_api(MyCuteDataclass).row_type.schema.fields) def test_user_type_annotated_with_id_after_conversion(self): MyCuteClass = NamedTuple('MyCuteClass', [ ('name', str), ]) self.assertFalse(hasattr(MyCuteClass, '_beam_schema_id')) schema = named_tuple_to_schema(MyCuteClass) self.assertTrue(hasattr(MyCuteClass, '_beam_schema_id')) self.assertEqual(MyCuteClass._beam_schema_id, schema.id) def test_schema_with_bad_field_raises_helpful_error(self): schema_proto = schema_pb2.Schema( fields=[ schema_pb2.Field( name="type_with_no_typeinfo", type=schema_pb2.FieldType()) ], id="helpful-error-uuid", ) # Should raise an exception referencing the problem field self.assertRaisesRegex( ValueError, "type_with_no_typeinfo", lambda: named_tuple_from_schema( schema_proto, # bypass schema cache schema_registry=SchemaTypeRegistry())) def test_row_type_is_callable(self): simple_row_type = row_type.RowTypeConstraint.from_fields([('foo', np.int64), ('bar', str)]) instance = simple_row_type(np.int64(35), 'baz') self.assertIsInstance(instance, simple_row_type.user_type) self.assertEqual(instance, (np.int64(35), 'baz')) def test_union(self): with_int = row_type.RowTypeConstraint.from_fields([('common', str), ('unique', int)]) with_any = row_type.RowTypeConstraint.from_fields([('common', str), ('unique', Any)]) union_type = typehints.Union[with_int, with_any] self.assertEqual( named_fields_from_element_type(union_type), [('common', str), ('unique', Any)]) class HypothesisTest(unittest.TestCase): # There is considerable variablility in runtime for this test, disable # deadline. @settings(deadline=None) @given(named_fields()) def test_named_fields_roundtrip(self, named_fields): typehint = row_type.RowTypeConstraint.from_fields(named_fields) roundtripped = typing_from_runner_api( typing_to_runner_api(typehint, schema_registry=SchemaTypeRegistry()), schema_registry=SchemaTypeRegistry()) self.assertEqual(typehint, roundtripped) @parameterized_class([ { 'pickler': pickle, }, { 'pickler': 'dill', }, { 'pickler': cloudpickle, }, ]) @pytest.mark.uses_dill class PickleTest(unittest.TestCase): def setUp(self): # pylint: disable=access-member-before-definition if self.pickler == 'dill': self.pickler = pytest.importorskip("dill") def test_generated_class_pickle_instance(self): schema = schema_pb2.Schema( id="some-uuid", fields=[ schema_pb2.Field( name='name', type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING), ) ]) user_type = named_tuple_from_schema(schema) instance = user_type(name="test") self.assertEqual(instance, self.pickler.loads(self.pickler.dumps(instance))) def test_generated_class_pickle(self): if self.pickler in [pickle, pytest.importorskip("dill")]: self.skipTest('https://github.com/apache/beam/issues/22714') schema = schema_pb2.Schema( id="some-uuid", fields=[ schema_pb2.Field( name='name', type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING), ) ]) user_type = named_tuple_from_schema(schema) self.assertEqual( user_type, self.pickler.loads(self.pickler.dumps(user_type))) def test_generated_class_row_type_pickle(self): row_proto = schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( id="some-other-uuid", fields=[ schema_pb2.Field( name='name', type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING), ) ]))) row_type_constraint = typing_from_runner_api( row_proto, schema_registry=SchemaTypeRegistry()) self.assertIsInstance(row_type_constraint, row_type.RowTypeConstraint) self.assertEqual( row_type_constraint, self.pickler.loads(self.pickler.dumps(row_type_constraint))) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/sharded_key_type.py0000644000175100017510000000612315206113033023604 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Type constraint for `ShardedKey`. Can be used like a type-hint, for instance, 'ShardedKeyType[int]' 'ShardedKeyType[Tuple[T]]' The type constraint is registered to be associated with :class:`apache_beam.coders.coders.ShardedKeyCoder`. Mostly for internal use. """ # pytype: skip-file from apache_beam import coders from apache_beam.typehints import typehints from apache_beam.utils.sharded_key import ShardedKey class ShardedKeyTypeConstraint(typehints.TypeConstraint, metaclass=typehints.GetitemConstructor): def __init__(self, key_type): typehints.validate_composite_type_param( key_type, error_msg_prefix='Parameter to ShardedKeyType hint') self.key_type = typehints.normalize(key_type) def _inner_types(self): yield self.key_type def _consistent_with_check_(self, sub): return ( isinstance(sub, self.__class__) and typehints.is_consistent_with(sub.key_type, self.key_type)) def type_check(self, instance): if not isinstance(instance, ShardedKey): raise typehints.CompositeTypeHintError( "ShardedKey type-constraint violated. Valid object instance " "must be of type 'ShardedKey'. Instead, an instance of '%s' " "was received." % (instance.__class__.__name__)) try: typehints.check_constraint(self.key_type, instance.key) except (typehints.CompositeTypeHintError, typehints.SimpleTypeHintError): raise typehints.CompositeTypeHintError( "%s type-constraint violated. The type of key in 'ShardedKey' " "is incorrect. Expected an instance of type '%s', " "instead received an instance of type '%s'." % (repr(self), repr(self.key_type), instance.key.__class__.__name__)) def match_type_variables(self, concrete_type): if isinstance(concrete_type, ShardedKeyTypeConstraint): return typehints.match_type_variables( self.key_type, concrete_type.key_type) return {} def __eq__(self, other): return isinstance( other, ShardedKeyTypeConstraint) and self.key_type == other.key_type def __hash__(self): return hash(self.key_type) def __repr__(self): return 'ShardedKey[%s]' % repr(self.key_type) ShardedKeyType = ShardedKeyTypeConstraint coders.typecoders.registry.register_coder( ShardedKeyType, coders.ShardedKeyCoder) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/sharded_key_type_test.py0000644000175100017510000000673715206113033024656 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the ShardedKeyTypeConstraint.""" # pytype: skip-file from apache_beam.typehints import Tuple from apache_beam.typehints import typehints from apache_beam.typehints.sharded_key_type import ShardedKeyType from apache_beam.typehints.typehints_test import TypeHintTestCase from apache_beam.utils.sharded_key import ShardedKey class ShardedKeyTypeConstraintTest(TypeHintTestCase): def test_compatibility(self): constraint1 = ShardedKeyType[int] constraint2 = ShardedKeyType[str] self.assertCompatible(constraint1, constraint1) self.assertCompatible(constraint2, constraint2) self.assertNotCompatible(constraint1, constraint2) def test_repr(self): constraint = ShardedKeyType[int] self.assertEqual('ShardedKey[]', repr(constraint)) def test_type_check_not_sharded_key(self): constraint = ShardedKeyType[int] obj = 5 with self.assertRaises(TypeError) as e: constraint.type_check(obj) self.assertEqual( "ShardedKey type-constraint violated. Valid object instance must be of " "type 'ShardedKey'. Instead, an instance of 'int' was received.", e.exception.args[0]) def test_type_check_invalid_key_type(self): constraint = ShardedKeyType[int] obj = ShardedKey(key='abc', shard_id=b'123') with self.assertRaises((TypeError, TypeError)) as e: constraint.type_check(obj) self.assertEqual( "ShardedKey[] type-constraint violated. The type of key " "in 'ShardedKey' is incorrect. Expected an instance of type \'\', instead received an instance of type 'str'.", e.exception.args[0]) def test_type_check_valid_simple_type(self): constraint = ShardedKeyType[str] obj = ShardedKey(key='abc', shard_id=b'123') self.assertIsNone(constraint.type_check(obj)) def test_type_check_valid_composite_type(self): constraint = ShardedKeyType[Tuple[int, str]] obj = ShardedKey(key=(1, 'a'), shard_id=b'123') self.assertIsNone(constraint.type_check(obj)) def test_match_type_variables(self): K = typehints.TypeVariable('K') # pylint: disable=invalid-name constraint = ShardedKeyType[K] self.assertEqual({K: int}, constraint.match_type_variables(ShardedKeyType[int])) def test_getitem(self): K = typehints.TypeVariable('K') # pylint: disable=invalid-name T = typehints.TypeVariable('T') # pylint: disable=invalid-name with self.assertRaisesRegex(TypeError, 'Parameter to ShardedKeyType hint.*'): _ = ShardedKeyType[K, T] with self.assertRaisesRegex(TypeError, 'Parameter to ShardedKeyType hint.*'): _ = ShardedKeyType[(K, T)] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/tagged_output_typehints_test.py0000644000175100017510000003641115206113033026305 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for tagged output type hints. This tests the implementation of type hints for tagged outputs via three styles: 1. Decorator style: @with_output_types(int, errors=str, warnings=str) class MyDoFn(beam.DoFn): ... 2. Method chain style: beam.ParDo(MyDoFn()).with_output_types(int, errors=str) 3. Function annotation style: def fn(element) -> int | TaggedOutput[Literal['errors'], str]: ... """ # pytype: skip-file import unittest from typing import Iterable from typing import Literal from typing import Union import apache_beam as beam from apache_beam.pvalue import TaggedOutput from apache_beam.typehints import Any from apache_beam.typehints import with_output_types from apache_beam.typehints.decorators import IOTypeHints class IOTypeHintsTaggedOutputTest(unittest.TestCase): """Tests for IOTypeHints.tagged_output_types() accessor.""" def test_empty_hints_returns_empty_dict(self): empty = IOTypeHints.empty() self.assertEqual(empty.tagged_output_types(), {}) def test_with_tagged_types(self): hints = IOTypeHints.empty().with_output_types(int, errors=str, warnings=str) self.assertEqual( hints.tagged_output_types(), { 'errors': str, 'warnings': str }) def test_simple_output_type_with_tagged_types(self): """simple_output_type() should still return main type when tags present.""" hints = IOTypeHints.empty().with_output_types(int, errors=str, warnings=str) self.assertEqual(hints.simple_output_type('test'), int) hints = IOTypeHints.empty().with_output_types( Union[int, str], errors=str, warnings=str) self.assertEqual(hints.simple_output_type('test'), Union[int, str]) def test_without_tagged_types(self): """Without tagged types, tagged_output_types() returns empty dict.""" hints = IOTypeHints.empty().with_output_types(int) self.assertEqual(hints.tagged_output_types(), {}) self.assertEqual(hints.simple_output_type('test'), int) class DecoratorStyleTaggedOutputTest(unittest.TestCase): """Tests for @with_output_types decorator style across all transforms.""" def test_pardo_decorator_pipeline(self): """Test that tagged types propagate through ParDo pipeline.""" @with_output_types(int, errors=str) class MyDoFn(beam.DoFn): def process(self, element): if element < 0: yield beam.pvalue.TaggedOutput('errors', f'Negative: {element}') else: yield element * 2 with beam.Pipeline() as p: results = ( p | beam.Create([-1, 0, 1, 2]) | beam.ParDo(MyDoFn()).with_outputs('errors', main='main')) self.assertEqual(results.main.element_type, int) self.assertEqual(results.errors.element_type, str) def test_map_decorator_pipeline(self): """Test that tagged types propagate through Map.""" @with_output_types(int, errors=str) def mapfn(element): if element < 0: return beam.pvalue.TaggedOutput('errors', f'Negative: {element}') else: return element * 2 with beam.Pipeline() as p: results = ( p | beam.Create([-1, 0, 1, 2]) | beam.Map(mapfn).with_outputs('errors', main='main')) self.assertEqual(results.main.element_type, int) self.assertEqual(results.errors.element_type, str) def test_flatmap_decorator_pipeline(self): """Test that tagged types propagate through FlatMap.""" @with_output_types(Iterable[int], errors=Iterable[str]) def flatmapfn(element): if element < 0: yield beam.pvalue.TaggedOutput('errors', f'Negative: {element}') else: yield element * 2 with beam.Pipeline() as p: results = ( p | beam.Create([-1, 0, 1, 2]) | beam.FlatMap(flatmapfn).with_outputs('errors', main='main')) self.assertEqual(results.main.element_type, int) self.assertEqual(results.errors.element_type, str) def test_maptuple_decorator_pipeline(self): """Test that tagged types propagate through MapTuple.""" @with_output_types(int, errors=str) def maptuplefn(key, value): if value < 0: return beam.pvalue.TaggedOutput('errors', f'Negative: {key}={value}') else: return value * 2 with beam.Pipeline() as p: results = ( p | beam.Create([('a', -1), ('b', 2), ('c', 3)]) | beam.MapTuple(maptuplefn).with_outputs('errors', main='main')) self.assertEqual(results.main.element_type, int) self.assertEqual(results.errors.element_type, str) def test_flatmaptuple_decorator_pipeline(self): """Test that tagged types propagate through FlatMapTuple.""" @with_output_types(Iterable[int], errors=Iterable[str]) def flatmaptuplefn(key, value): if value < 0: yield beam.pvalue.TaggedOutput('errors', f'Negative: {key}={value}') else: yield value * 2 with beam.Pipeline() as p: results = ( p | beam.Create([('a', -1), ('b', 2), ('c', 3)]) | beam.FlatMapTuple(flatmaptuplefn).with_outputs( 'errors', main='main')) self.assertEqual(results.main.element_type, int) self.assertEqual(results.errors.element_type, str) class ChainStyleTaggedOutputTest(unittest.TestCase): """Tests for .with_output_types() method chain style across all transforms.""" def test_pardo_chain_pipeline(self): """Test ParDo with chained type hints.""" class SimpleDoFn(beam.DoFn): def process(self, element): if element < 0: yield beam.pvalue.TaggedOutput('errors', f'Negative: {element}') else: yield element * 2 with beam.Pipeline() as p: results = ( p | beam.Create([-1, 0, 1, 2]) | beam.ParDo(SimpleDoFn()).with_output_types( int, errors=str).with_outputs('errors', main='main')) self.assertEqual(results.main.element_type, int) self.assertEqual(results.errors.element_type, str) def test_map_chain_pipeline(self): """Test Map with chained type hints.""" def mapfn(element): if element < 0: return beam.pvalue.TaggedOutput('errors', f'Negative: {element}') else: return element * 2 with beam.Pipeline() as p: results = ( p | beam.Create([-1, 0, 1, 2]) | beam.Map(mapfn).with_output_types(int, errors=str).with_outputs( 'errors', main='main')) self.assertEqual(results.main.element_type, int) self.assertEqual(results.errors.element_type, str) def test_flatmap_chain_pipeline(self): """Test FlatMap with chained type hints. Note: For FlatMap.with_output_types(), specify the element type directly (int), not wrapped in Iterable. The transform handles iteration internally. """ def flatmapfn(element): if element < 0: yield beam.pvalue.TaggedOutput('errors', f'Negative: {element}') else: yield element * 2 with beam.Pipeline() as p: results = ( p | beam.Create([-1, 0, 1, 2]) | beam.FlatMap(flatmapfn).with_output_types( int, errors=str).with_outputs('errors', main='main')) self.assertEqual(results.main.element_type, int) self.assertEqual(results.errors.element_type, str) def test_maptuple_chain_pipeline(self): """Test MapTuple with chained type hints.""" def maptuplefn(key, value): if value < 0: return beam.pvalue.TaggedOutput('errors', f'Negative: {key}={value}') else: return value * 2 with beam.Pipeline() as p: results = ( p | beam.Create([('a', -1), ('b', 2), ('c', 3)]) | beam.MapTuple(maptuplefn).with_output_types( int, errors=str).with_outputs('errors', main='main')) self.assertEqual(results.main.element_type, int) self.assertEqual(results.errors.element_type, str) def test_flatmaptuple_chain_pipeline(self): """Test FlatMapTuple with chained type hints. Note: For FlatMapTuple.with_output_types(), specify the element type directly (int), not wrapped in Iterable. """ def flatmaptuplefn(key, value): if value < 0: yield beam.pvalue.TaggedOutput('errors', f'Negative: {key}={value}') else: yield value * 2 with beam.Pipeline() as p: results = ( p | beam.Create([('a', -1), ('b', 2), ('c', 3)]) | beam.FlatMapTuple(flatmaptuplefn).with_output_types( int, errors=str).with_outputs('errors', main='main')) self.assertEqual(results.main.element_type, int) self.assertEqual(results.errors.element_type, str) class AnnotationStyleTaggedOutputTest(unittest.TestCase): """Tests for function annotation style across all transforms.""" def test_map_annotation_union(self): """Test Map with Union[int, TaggedOutput[...]] annotation.""" def mapfn(element: int) -> int | TaggedOutput[Literal['errors'], str]: if element < 0: return beam.pvalue.TaggedOutput('errors', f'Negative: {element}') else: return element * 2 with beam.Pipeline() as p: results = ( p | beam.Create([-1, 0, 1, 2]) | beam.Map(mapfn).with_outputs('errors', main='main')) self.assertEqual(results.main.element_type, int) self.assertEqual(results.errors.element_type, str) def test_map_annotation_multiple_tags(self): """Test Map with multiple TaggedOutput types in annotation.""" def mapfn( element: int ) -> int | TaggedOutput[Literal['errors'], str] | TaggedOutput[Literal['warnings'], str]: if element < 0: return beam.pvalue.TaggedOutput('errors', f'Negative: {element}') elif element == 0: return beam.pvalue.TaggedOutput('warnings', 'Zero value') else: return element * 2 with beam.Pipeline() as p: results = ( p | beam.Create([-1, 0, 1, 2]) | beam.Map(mapfn).with_outputs('errors', 'warnings', main='main')) self.assertEqual(results.main.element_type, int) self.assertEqual(results.errors.element_type, str) self.assertEqual(results.warnings.element_type, str) def test_flatmap_annotation_iterable(self): """Test FlatMap with Iterable[int | TaggedOutput[...]] annotation.""" def flatmapfn( element: int) -> Iterable[int | TaggedOutput[Literal['errors'], str]]: if element < 0: yield beam.pvalue.TaggedOutput('errors', f'Negative: {element}') else: yield element * 2 with beam.Pipeline() as p: results = ( p | beam.Create([-1, 0, 1, 2]) | beam.FlatMap(flatmapfn).with_outputs('errors', main='main')) self.assertEqual(results.main.element_type, int) self.assertEqual(results.errors.element_type, str) def test_pardo_annotation_process_method(self): """Test DoFn with process method annotation.""" class AnnotatedDoFn(beam.DoFn): def process( self, element: int) -> Iterable[int | TaggedOutput[Literal['errors'], str]]: if element < 0: yield beam.pvalue.TaggedOutput('errors', f'Negative: {element}') else: yield element * 2 with beam.Pipeline() as p: results = ( p | beam.Create([-1, 0, 1, 2]) | beam.ParDo(AnnotatedDoFn()).with_outputs('errors', main='main')) self.assertEqual(results.main.element_type, int) self.assertEqual(results.errors.element_type, str) class CompositeTaggedOutputInferenceTest(unittest.TestCase): """Tests for _infer_result_type when a composite PTransform returns tagged outputs as a dict of fresh PCollections. """ def test_composite_returning_tagged_dict_preserves_existing_types(self): """A composite that returns a dict of PCollections already typed by DoOutputsTuple.__getitem__ should preserve those types through _infer_result_type (the guard skips inference when element_type is set).""" class MyComposite(beam.PTransform): def expand(self, pcoll): results = ( pcoll | beam.ParDo(self._MyDoFn()).with_outputs('errors', main='main')) # Return a dict of the DoOutputsTuple's PCollections. # These already have types set via __getitem__. return {'main': results.main, 'errors': results.errors} @beam.typehints.with_output_types(int, errors=str) class _MyDoFn(beam.DoFn): def process(self, element): if element < 0: yield beam.pvalue.TaggedOutput('errors', f'Negative: {element}') else: yield element with beam.Pipeline() as p: result = (p | beam.Create([-1, 0, 1, 2]) | MyComposite()) self.assertEqual(result['main'].element_type, int) self.assertEqual(result['errors'].element_type, str) def test_composite_returning_tagged_dict_without_dofn_hints_is_any(self): class MyComposite(beam.PTransform): def expand(self, pcoll): results = ( pcoll | beam.ParDo(self._MyDoFn()).with_outputs('errors', main='main')) return {'main': results.main, 'errors': results.errors} class _MyDoFn(beam.DoFn): def process(self, element): if element < 0: yield beam.pvalue.TaggedOutput('errors', f'Negative: {element}') else: yield element with beam.Pipeline() as p: result = (p | beam.Create([-1, 0, 1, 2]) | MyComposite()) self.assertEqual(result['errors'].element_type, Any) def test_composite_pcollections_uses_tagged_type_hints(self): """A composite that creates new PCollections (element_type=None) and returns them as a dict should still get correct tagged types from the type hints.""" @beam.typehints.with_output_types(int, errors=str) class MyComposite(beam.PTransform): def expand(self, pcoll): results = ( pcoll | beam.ParDo(self._MyDoFn()).with_outputs('errors', main='main')) return { p.tag if p.tag else results._main_tag: beam.pvalue.PCollection( pcoll.pipeline, tag=p.tag) for p in results } class _MyDoFn(beam.DoFn): def process( self, element ) -> Iterable[int | beam.TaggedOutput[Literal['errors'], str]]: if element < 0: yield beam.pvalue.TaggedOutput('errors', f'Negative: {element}') else: yield element with beam.Pipeline() as p: result = (p | beam.Create([-1, 0, 1, 2]) | MyComposite()) self.assertEqual(result['main'].element_type, int) self.assertEqual(result['errors'].element_type, str) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.7744493 apache_beam-2.74.0/apache_beam/typehints/testing/0000755000175100017510000000000015206113160021363 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/testing/__init__.py0000644000175100017510000000150415206113033023473 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Testing utilities for apache_beam.typehints.""" ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/testing/strategies.py0000644000175100017510000000721515206113033024113 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """hypothesis strategies for generating schema types. Intended for internal use only, no backward-compatibility guarantees.""" import keyword import unicodedata from typing import Mapping from typing import Optional from typing import Sequence from hypothesis import strategies as st from apache_beam.typehints import row_type from apache_beam.typehints.schemas import _PRIMITIVES PRIMITIVES = [p[0] for p in _PRIMITIVES] def field_names(): @st.composite def field_name_candidates(draw): """Strategy to produce valid field names for Beam schema types.""" # unicode categories that cannot be used in Python identifiers identifer_denylist = ( 'Lo', 'Lm', 'C', 'P', 'Sm', 'Sc', 'So', 'Sk', 'M', 'No', 'Z') # First character can't be numeric (Nd). # It also can't start with '_' in a NamedTuple. field_first_character = draw( st.text( alphabet=st.characters( blacklist_categories=('Nd', ) + identifer_denylist, blacklist_characters=('_')), min_size=1, max_size=1)) field_remainder = draw( st.text( alphabet=st.characters(blacklist_categories=identifer_denylist))) return field_first_character + field_remainder return field_name_candidates().filter( lambda s: s.isidentifier() and not keyword.iskeyword(s)) def _named_fields_from_types(types): return st.lists( st.tuples(field_names(), types), min_size=1, # Python identifiers are normalized with form NFKC (see # https://peps.python.org/pep-0672/#normalizing-identifiers). We use the # same normalization here to avoid name collisions. unique_by=lambda name_and_type: unicodedata.normalize( 'NFKC', name_and_type[0]), ) def types(): """Strategy to produce types that are convertible to Beam schema FieldType instances.""" def _extend_types(types): optionals = types.map(lambda typ: Optional[typ]) sequences = types.map(lambda typ: Sequence[typ]) mappings = st.tuples(types, types).map(lambda typs: Mapping[typs[0], typs[1]]) rows = _named_fields_from_types(types).map( row_type.RowTypeConstraint.from_fields) return st.one_of(optionals, sequences, mappings, rows) # TODO: Currently this will only draw from the primitive types that can be # roundtripped faithfully (e.g. np.int64, not int). We should add support for # other types: # - Logical Types (e.g. Timestamp) # - Shunted primitive types (e.g. int) # We'll need to provide support for limiting the types that are drawn. This # could be similar to the allowlist[_categories]/denylist[_categories] pattern # used in st.characters. return st.recursive(st.sampled_from(PRIMITIVES), _extend_types) def named_fields(): """Strategy to produce a set of named fields (type ``List[Tuple[str, type]]``).""" return _named_fields_from_types(types()) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/trivial_inference.py0000644000175100017510000006350115206113033023754 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Trivial type inference for simple functions. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file import builtins import collections import dis import inspect import pprint import sys import traceback import types from functools import reduce from apache_beam import pvalue from apache_beam.typehints import Any from apache_beam.typehints import row_type from apache_beam.typehints import typehints from apache_beam.utils import python_callable class TypeInferenceError(ValueError): """Error to raise when type inference failed.""" pass def instance_to_type(o): """Given a Python object o, return the corresponding type hint. """ t = type(o) if o is None: return type(None) elif t == pvalue.Row: return row_type.RowTypeConstraint.from_fields([ (name, instance_to_type(value)) for name, value in o.as_dict().items() ]) elif t not in typehints.DISALLOWED_PRIMITIVE_TYPES: # pylint: disable=bad-option-value if t == BoundMethod: return types.MethodType return t elif t == tuple: return typehints.Tuple[[instance_to_type(item) for item in o]] elif t == list: if len(o) > 0: return typehints.List[typehints.Union[[ instance_to_type(item) for item in o ]]] else: return typehints.List[typehints.Any] elif t == set: if len(o) > 0: return typehints.Set[typehints.Union[[ instance_to_type(item) for item in o ]]] else: return typehints.Set[typehints.Any] elif t == frozenset: if len(o) > 0: return typehints.FrozenSet[typehints.Union[[ instance_to_type(item) for item in o ]]] else: return typehints.FrozenSet[typehints.Any] elif t == dict: if len(o) > 0: return typehints.Dict[ typehints.Union[[instance_to_type(k) for k, v in o.items()]], typehints.Union[[instance_to_type(v) for k, v in o.items()]], ] else: return typehints.Dict[typehints.Any, typehints.Any] else: raise TypeInferenceError('Unknown forbidden type: %s' % t) def union_list(xs, ys): assert len(xs) == len(ys) return [union(x, y) for x, y in zip(xs, ys)] class Const(object): def __init__(self, value): self.value = value self.type = instance_to_type(value) def __eq__(self, other): return isinstance(other, Const) and self.value == other.value def __hash__(self): return hash(self.value) def __repr__(self): return 'Const[%s]' % str(self.value)[:100] @staticmethod def unwrap(x): if isinstance(x, Const): return x.type return x @staticmethod def unwrap_all(xs): return [Const.unwrap(x) for x in xs] class FrameState(object): """Stores the state of the frame at a particular point of execution. """ def __init__(self, f, local_vars=None, stack=(), kw_names=None): self.f = f self.co = f.__code__ self.vars = list(local_vars) self.stack = list(stack) self.kw_names = kw_names def __eq__(self, other): return isinstance(other, FrameState) and self.__dict__ == other.__dict__ def __hash__(self): return hash(tuple(sorted(self.__dict__.items()))) def copy(self): return FrameState(self.f, self.vars, self.stack, self.kw_names) def const_type(self, i): return Const(self.co.co_consts[i]) def get_closure(self, i): num_cellvars = len(self.co.co_cellvars) if i < num_cellvars: return self.vars[i] else: return self.f.__closure__[i - num_cellvars].cell_contents def closure_type(self, i): """Returns a TypeConstraint or Const.""" val = self.get_closure(i) if isinstance(val, typehints.TypeConstraint): return val else: return Const(val) def get_global(self, i): name = self.get_name(i) if name in self.f.__globals__: return Const(self.f.__globals__[name]) if name in builtins.__dict__: return Const(builtins.__dict__[name]) return Any def get_name(self, i): return self.co.co_names[i] def __repr__(self): return 'Stack: %s Vars: %s' % (self.stack, self.vars) def __or__(self, other): if self is None: return other.copy() elif other is None: return self.copy() return FrameState( self.f, union_list(self.vars, other.vars), union_list(self.stack, other.stack)) def __ror__(self, left): return self | left def union(a, b): """Returns the union of two types or Const values. """ if a == b: return a elif not a: return b elif not b: return a a = Const.unwrap(a) b = Const.unwrap(b) # TODO(robertwb): Work this into the Union code in a more generic way. if type(a) == type(b) and element_type(a) == typehints.Union[()]: return b elif type(a) == type(b) and element_type(b) == typehints.Union[()]: return a return typehints.Union[a, b] def finalize_hints(type_hint): """Sets type hint for empty data structures to Any.""" def visitor(tc, unused_arg): if isinstance(tc, typehints.DictConstraint): empty_union = typehints.Union[()] if tc.key_type == empty_union: tc.key_type = Any if tc.value_type == empty_union: tc.value_type = Any if isinstance(type_hint, typehints.TypeConstraint): type_hint.visit(visitor, None) def element_type(hint): """Returns the element type of a composite type. """ hint = Const.unwrap(hint) if isinstance(hint, typehints.SequenceTypeConstraint): return hint.inner_type elif isinstance(hint, typehints.TupleHint.TupleConstraint): return typehints.Union[hint.tuple_types] elif isinstance(hint, typehints.UnionHint.UnionConstraint) and not hint.union_types: return hint return Any def key_value_types(kv_type): """Returns the key and value type of a KV type. """ # TODO(robertwb): Unions of tuples, etc. # TODO(robertwb): Assert? if (isinstance(kv_type, typehints.TupleHint.TupleConstraint) and len(kv_type.tuple_types) == 2): return kv_type.tuple_types elif isinstance( kv_type, typehints.UnionHint.UnionConstraint) and not kv_type.union_types: return kv_type, kv_type return Any, Any known_return_types = { len: int, hash: int, } class BoundMethod(object): """Used to create a bound method when we only know the type of the instance. """ def __init__(self, func, type): """Instantiates a bound method object. Args: func (types.FunctionType): The method's underlying function type (type): The class of the method. """ self.func = func self.type = type def hashable(c): try: hash(c) return True except TypeError: return False def infer_return_type(c, input_types, debug=False, depth=5): """Analyses a callable to deduce its return type. Args: c: A Python callable to infer the return type of. input_types: A sequence of inputs corresponding to the input types. debug: Whether to print verbose debugging information. depth: Maximum inspection depth during type inference. Returns: A TypeConstraint that that the return value of this function will (likely) satisfy given the specified inputs. """ try: if hashable(c) and c in known_return_types: return known_return_types[c] elif isinstance(c, types.FunctionType): return infer_return_type_func(c, input_types, debug, depth) elif isinstance(c, types.MethodType): if c.__self__ is not None: input_types = [Const(c.__self__)] + input_types return infer_return_type_func(c.__func__, input_types, debug, depth) elif isinstance(c, BoundMethod): input_types = [c.type] + input_types return infer_return_type_func(c.func, input_types, debug, depth) elif inspect.isclass(c): if c in typehints.DISALLOWED_PRIMITIVE_TYPES: return { list: typehints.List[Any], set: typehints.Set[Any], frozenset: typehints.FrozenSet[Any], tuple: typehints.Tuple[Any, ...], dict: typehints.Dict[Any, Any] }[c] return c elif (c == getattr and len(input_types) == 2 and isinstance(input_types[1], Const)): from apache_beam.typehints import opcodes return opcodes._getattr(input_types[0], input_types[1].value) elif isinstance(c, python_callable.PythonCallableWithSource): # TODO(https://github.com/apache/beam/issues/24755): This can be removed # once support for inference across *args and **kwargs is implemented. return infer_return_type(c._callable, input_types, debug, depth) else: return Any except TypeInferenceError: if debug: traceback.print_exc() return Any except Exception: if debug: sys.stdout.flush() raise else: return Any def infer_return_type_func(f, input_types, debug=False, depth=0): """Analyses a function to deduce its return type. Args: f: A Python function object to infer the return type of. input_types: A sequence of inputs corresponding to the input types. debug: Whether to print verbose debugging information. depth: Maximum inspection depth during type inference. Returns: A TypeConstraint that that the return value of this function will (likely) satisfy given the specified inputs. Raises: TypeInferenceError: if no type can be inferred. """ if debug: print() print(f, id(f), input_types) ver = (sys.version_info.major, sys.version_info.minor) if ver >= (3, 13): dis.dis(f, show_caches=True, show_offsets=True) elif ver >= (3, 11): dis.dis(f, show_caches=True) else: dis.dis(f) from . import opcodes simple_ops = dict((k.upper(), v) for k, v in opcodes.__dict__.items()) from . import intrinsic_one_ops co = f.__code__ code = co.co_code end = len(code) pc = 0 free = None yields = set() returns = set() # TODO(robertwb): Default args via inspect module. local_vars = list(input_types) + [typehints.Union[()]] * ( len(co.co_varnames) - len(input_types)) state = FrameState(f, local_vars) states = collections.defaultdict(lambda: None) jumps = collections.defaultdict(int) # In Python 3, use dis library functions to disassemble bytecode and handle # EXTENDED_ARGs. ofs_table = {} # offset -> instruction if (sys.version_info.major, sys.version_info.minor) >= (3, 11): dis_ints = dis.get_instructions(f, show_caches=True) else: dis_ints = dis.get_instructions(f) for instruction in dis_ints: ofs_table[instruction.offset] = instruction # Python 3.6+: 1 byte opcode + 1 byte arg (2 bytes, arg may be ignored). inst_size = 2 opt_arg_size = 0 jump_multiplier = 2 # Python 3.14+ push nulls are used to signal kwargs for CALL_FUNCTION_EX # so there must be a little extra bookkeeping even if we don't care about # the nulls themselves. last_op_push_null = 0 last_pc = -1 last_real_opname = opname = None while pc < end: # pylint: disable=too-many-nested-blocks if opname not in ('PRECALL', 'CACHE'): last_real_opname = opname start = pc instruction = ofs_table[pc] op = instruction.opcode if debug: print('-->' if pc == last_pc else ' ', end=' ') print(repr(pc).rjust(4), end=' ') print(dis.opname[op].ljust(20), end=' ') pc += inst_size # Python 3.13 deprecated show_caches and removed cache instruction # outputs, instead putting the cache instructions nested within the # Instruction object. if (sys.version_info.major, sys.version_info.minor) >= (3, 13): cache = instruction.cache_info if cache is not None: # Each entry in cache_info is a tuple with the name of the cached # object, the number of cache calls it produces, and the byte # representation of the cached item. for cache_entry in cache: pc += cache_entry[1] * inst_size arg = None if op >= dis.HAVE_ARGUMENT: arg = instruction.arg pc += opt_arg_size if debug: print(str(arg).rjust(5), end=' ') if op in dis.hasconst: print('(' + repr(co.co_consts[arg]) + ')', end=' ') elif op in dis.hasname: if (sys.version_info.major, sys.version_info.minor) >= (3, 11): # Pre-emptively bit-shift so the print doesn't go out of index print_arg = arg >> 1 else: print_arg = arg print('(' + co.co_names[print_arg] + ')', end=' ') elif op in dis.hasjrel: print('(to ' + repr(pc + (arg * jump_multiplier)) + ')', end=' ') elif op in dis.haslocal: # Args to double-fast opcodes are bit manipulated, correct the arg # for printing + avoid the out-of-index if dis.opname[op] == 'LOAD_FAST_LOAD_FAST' or dis.opname[ op] == "LOAD_FAST_BORROW_LOAD_FAST_BORROW": print( '(' + co.co_varnames[arg >> 4] + ', ' + co.co_varnames[arg & 15] + ')', end=' ') elif dis.opname[op] == 'STORE_FAST_LOAD_FAST': print('(' + co.co_varnames[arg & 15] + ')', end=' ') elif dis.opname[op] == 'STORE_FAST_STORE_FAST': pass elif dis.opname[op] == 'LOAD_DEREF': pass else: print('(' + co.co_varnames[arg] + ')', end=' ') elif op in dis.hascompare: if (sys.version_info.major, sys.version_info.minor) >= (3, 12): # In 3.12 this arg was bit-shifted. Shifting it back avoids an # out-of-index. arg = arg >> 4 print('(' + dis.cmp_op[arg] + ')', end=' ') elif op in dis.hasfree: if free is None: free = co.co_cellvars + co.co_freevars # From 3.11 on the arg is no longer offset by len(co_varnames) # so we adjust it back print_arg = arg if (sys.version_info.major, sys.version_info.minor) >= (3, 11): print_arg = arg - len(co.co_varnames) print('(' + free[print_arg] + ')', end=' ') # Actually emulate the op. if state is None and states[start] is None: # No control reaches here (yet). if debug: print() continue state |= states[start] opname = dis.opname[op] jmp = jmp_state = None if opname.startswith('CALL_FUNCTION'): if opname == 'CALL_FUNCTION': pop_count = arg + 1 if depth <= 0: return_type = Any elif isinstance(state.stack[-pop_count], Const): return_type = infer_return_type( state.stack[-pop_count].value, state.stack[1 - pop_count:], debug=debug, depth=depth - 1) else: return_type = Any elif opname == 'CALL_FUNCTION_KW': # TODO(BEAM-24755): Handle keyword arguments. Requires passing them by # name to infer_return_type. pop_count = arg + 2 if isinstance(state.stack[-pop_count], Const): from apache_beam.pvalue import Row if state.stack[-pop_count].value == Row: fields = state.stack[-1].value return_type = row_type.RowTypeConstraint.from_fields( list( zip( fields, Const.unwrap_all(state.stack[-pop_count + 1:-1])))) else: return_type = Any else: return_type = Any elif opname == 'CALL_FUNCTION_EX': # stack[-has_kwargs]: Map of keyword args. # stack[-1 - has_kwargs]: Iterable of positional args. # stack[-2 - has_kwargs]: Function to call. if arg is None: # CALL_FUNCTION_EX does not take an arg in 3.14, instead the # signaling for kwargs is done via a PUSH_NULL instruction # right before CALL_FUNCTION_EX. A PUSH_NULL indicates that # there are no kwargs. arg = ~last_op_push_null has_kwargs: int = arg & 1 pop_count = has_kwargs + 2 if has_kwargs: # TODO(BEAM-24755): Unimplemented. Requires same functionality as a # CALL_FUNCTION_KW implementation. return_type = Any else: args = state.stack[-1] _callable = state.stack[-2] if isinstance(args, typehints.ListConstraint): # Case where there's a single var_arg argument. args = [args] elif isinstance(args, typehints.TupleConstraint): args = list(args._inner_types()) elif isinstance(args, typehints.SequenceTypeConstraint): args = [element_type(args)] * len( inspect.getfullargspec(_callable.value).args) return_type = infer_return_type( _callable.value, args, debug=debug, depth=depth - 1) else: raise TypeInferenceError('unable to handle %s' % opname) state.stack[-pop_count:] = [return_type] elif opname == 'CALL_METHOD': pop_count = 1 + arg # LOAD_METHOD will return a non-Const (Any) if loading from an Any. if isinstance(state.stack[-pop_count], Const) and depth > 0: return_type = infer_return_type( state.stack[-pop_count].value, state.stack[1 - pop_count:], debug=debug, depth=depth - 1) else: return_type = typehints.Any state.stack[-pop_count:] = [return_type] elif opname == 'CALL': pop_count = 1 + arg # Keyword Args case if state.kw_names is not None: if isinstance(state.stack[-pop_count], Const): from apache_beam.pvalue import Row if state.stack[-pop_count].value == Row: fields = state.kw_names return_type = row_type.RowTypeConstraint.from_fields( list( zip(fields, Const.unwrap_all(state.stack[-pop_count + 1:])))) else: return_type = Any state.kw_names = None else: # Handle comprehensions always having an arg of 0 for CALL # See https://github.com/python/cpython/issues/102403 for context. if (pop_count == 1 and last_real_opname == 'GET_ITER' and len(state.stack) > 1 and isinstance(state.stack[-2], Const) and getattr(state.stack[-2].value, '__name__', None) in ('', '', '', '')): pop_count += 1 if depth <= 0 or pop_count > len(state.stack): return_type = Any elif isinstance(state.stack[-pop_count], Const): return_type = infer_return_type( state.stack[-pop_count].value, state.stack[1 - pop_count:], debug=debug, depth=depth - 1) else: return_type = Any state.stack[-pop_count:] = [return_type] # CALL_KW handles all calls with kwargs post-3.13, have to maintain # both paths for now. This call replaces state.kw_names with a tuple # of keyword names at state.stack[-1] elif opname == 'CALL_KW': pop_count = 2 + arg if isinstance(state.stack[-pop_count], Const): from apache_beam.pvalue import Row if state.stack[-pop_count].value == Row: fields = state.stack[-1].value return_type = row_type.RowTypeConstraint.from_fields( list( zip(fields, Const.unwrap_all(state.stack[-pop_count + 1:-1])))) else: return_type = Any else: # Handle comprehensions always having an arg of 0 for CALL # See https://github.com/python/cpython/issues/102403 for context. if (pop_count == 1 and last_real_opname == 'GET_ITER' and len(state.stack) > 1 and isinstance(state.stack[-2], Const) and getattr(state.stack[-2].value, '__name__', None) in ('', '', '', '')): pop_count += 1 if depth <= 0 or pop_count > len(state.stack): return_type = Any elif isinstance(state.stack[-pop_count], Const): return_type = infer_return_type( state.stack[-pop_count].value, state.stack[1 - pop_count:], debug=debug, depth=depth - 1) else: return_type = Any state.stack[-pop_count:] = [return_type] elif opname in simple_ops: if debug: print("Executing simple op " + opname) simple_ops[opname](state, arg) elif opname == 'RETURN_VALUE': returns.add(state.stack[-1]) state = None elif opname == 'YIELD_VALUE': yields.add(state.stack[-1]) elif opname == 'JUMP_FORWARD': jmp = pc + arg * jump_multiplier jmp_state = state state = None elif opname in ('JUMP_BACKWARD', 'JUMP_BACKWARD_NO_INTERRUPT'): jmp = pc - (arg * jump_multiplier) jmp_state = state state = None elif opname == 'JUMP_ABSOLUTE': jmp = arg * jump_multiplier jmp_state = state state = None elif opname in ('POP_JUMP_IF_TRUE', 'POP_JUMP_IF_FALSE'): state.stack.pop() # The arg was changed to be a relative delta instead of an absolute # in 3.11, and became a full instruction instead of a # pseudo-instruction in 3.12 if (sys.version_info.major, sys.version_info.minor) >= (3, 12): jmp = pc + arg * jump_multiplier else: jmp = arg * jump_multiplier jmp_state = state.copy() elif opname in ('POP_JUMP_FORWARD_IF_TRUE', 'POP_JUMP_FORWARD_IF_FALSE'): state.stack.pop() jmp = pc + arg * jump_multiplier jmp_state = state.copy() elif opname in ('POP_JUMP_BACKWARD_IF_TRUE', 'POP_JUMP_BACKWARD_IF_FALSE'): state.stack.pop() jmp = pc - (arg * jump_multiplier) jmp_state = state.copy() elif opname in ('POP_JUMP_FORWARD_IF_NONE', 'POP_JUMP_FORWARD_IF_NOT_NONE'): state.stack.pop() jmp = pc + arg * jump_multiplier jmp_state = state.copy() elif opname in ('POP_JUMP_BACKWARD_IF_NONE', 'POP_JUMP_BACKWARD_IF_NOT_NONE'): state.stack.pop() jmp = pc - (arg * jump_multiplier) jmp_state = state.copy() elif opname in ('JUMP_IF_TRUE_OR_POP', 'JUMP_IF_FALSE_OR_POP'): # The arg was changed to be a relative delta instead of an absolute # in 3.11 if (sys.version_info.major, sys.version_info.minor) >= (3, 11): jmp = pc + arg * jump_multiplier else: jmp = arg * jump_multiplier jmp_state = state.copy() state.stack.pop() elif opname == 'FOR_ITER': jmp = pc + arg * jump_multiplier if sys.version_info >= (3, 12): # The jump is relative to the next instruction after a cache call, # so jump 4 more bytes. jmp += 4 jmp_state = state.copy() jmp_state.stack.pop() state.stack.append(element_type(state.stack[-1])) elif opname == 'POP_ITER': # Introduced in 3.14. state.stack.pop() elif opname == 'COPY_FREE_VARS': # Helps with calling closures, but since we aren't executing # them we can treat this as a no-op pass elif opname == 'KW_NAMES': tup = co.co_consts[arg] state.kw_names = tup elif opname == 'RESUME': # RESUME is a no-op pass elif opname == 'PUSH_NULL': # We're treating this as a no-op to avoid having to check # for extra None values on the stack when we extract return # values last_op_push_null = 1 pass elif opname == 'NOT_TAKEN': # NOT_TAKEN is a no-op introduced in 3.14. pass elif opname == 'PRECALL': # PRECALL is a no-op. pass elif opname == 'MAKE_CELL': # TODO: see if we need to implement cells like this pass elif opname == 'RETURN_GENERATOR': # TODO: see what this behavior is supposed to be beyond # putting something on the stack to be popped off state.stack.append(None) pass elif opname == 'CACHE': # No-op introduced in 3.11. Without handling this some # instructions have functionally > 2 byte size. pass elif opname == 'RETURN_CONST': # Introduced in 3.12. Handles returning constants directly # instead of having a LOAD_CONST before a RETURN_VALUE. returns.add(state.const_type(arg)) state = None elif opname == 'CALL_INTRINSIC_1': # Introduced in 3.12. The arg is an index into a table of # operations reproduced in INT_ONE_OPS. Not all ops are # relevant for our type checking infrastructure. int_op = intrinsic_one_ops.INT_ONE_OPS[arg] if debug: print("Executing intrinsic one op", int_op.__name__.upper()) int_op(state, arg) else: raise TypeInferenceError('unable to handle %s' % opname) # Clear check for previous push_null. if opname != 'PUSH_NULL' and last_op_push_null == 1: last_op_push_null = 0 if jmp is not None: # TODO(robertwb): Is this guaranteed to converge? new_state = states[jmp] | jmp_state if jmp < pc and new_state != states[jmp] and jumps[pc] < 5: jumps[pc] += 1 pc = jmp states[jmp] = new_state if debug: print() print(state) pprint.pprint(dict(item for item in states.items() if item[1])) if yields: result = typehints.Iterable[reduce(union, Const.unwrap_all(yields))] else: result = reduce(union, Const.unwrap_all(returns)) finalize_hints(result) if debug: print(f, id(f), input_types, '->', result) return result ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/trivial_inference_test.py0000644000175100017510000004130215206113033025006 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for apache_beam.typehints.trivial_inference.""" # pytype: skip-file import types import unittest import apache_beam as beam from apache_beam.typehints import row_type from apache_beam.typehints import trivial_inference from apache_beam.typehints import typehints from apache_beam.utils import python_callable global_int = 1 class TrivialInferenceTest(unittest.TestCase): def assertReturnType(self, expected, f, inputs=(), depth=5): self.assertEqual( expected, trivial_inference.infer_return_type(f, inputs, debug=True, depth=depth)) # The meaning of Jump Offsets in Python 3.10 was changed. # https://github.com/python/cpython/issues/71316 # Reported as a bug in Beam https://github.com/apache/beam/issues/21671 def testJumpOffsets(self): fn = lambda x: False wrapper = lambda x, *args, **kwargs: [x] if fn(x, *args, **kwargs) else [] self.assertReturnType(typehints.List[int], wrapper, [int]) def testBuildListUnpack(self): # Lambda uses BUILD_LIST_UNPACK opcode in Python 3. self.assertReturnType( typehints.List[int], lambda _list: [*_list, *_list, *_list], [typehints.List[int]]) def testBuildTupleUnpack(self): # Lambda uses BUILD_TUPLE_UNPACK opcode in Python 3. # yapf: disable self.assertReturnType( typehints.Tuple[typehints.Union[int, str], ...], lambda _list1, _list2: (*_list1, *_list2, *_list2), [typehints.List[int], typehints.List[str]]) # yapf: enable def testBuildSetUnpackOrUpdate(self): self.assertReturnType( typehints.Set[typehints.Union[int, str]], lambda _list1, _list2: {*_list1, *_list2, *_list2}, [typehints.List[int], typehints.List[str]]) def testBuildMapUnpackOrUpdate(self): self.assertReturnType( typehints.Dict[str, typehints.Union[int, str, float]], lambda a, b, c: { **a, **b, **c }, [ typehints.Dict[str, int], typehints.Dict[str, str], typehints.List[typehints.Tuple[str, float]] ]) def testIdentity(self): self.assertReturnType(int, lambda x: x, [int]) def testIndexing(self): self.assertReturnType(int, lambda x: x[0], [typehints.Tuple[int, str]]) self.assertReturnType(str, lambda x: x[1], [typehints.Tuple[int, str]]) self.assertReturnType(str, lambda x: x[1], [typehints.List[str]]) def testTuples(self): self.assertReturnType( typehints.Tuple[typehints.Tuple[()], int], lambda x: ((), x), [int]) self.assertReturnType( typehints.Tuple[str, int, float], lambda x: (x, 0, 1.0), [str]) def testGetItem(self): def reverse(ab): return ab[-1], ab[0] self.assertReturnType( typehints.Tuple[typehints.Any, typehints.Any], reverse, [typehints.Any]) self.assertReturnType( typehints.Tuple[int, float], reverse, [typehints.Tuple[float, int]]) self.assertReturnType( typehints.Tuple[int, str], reverse, [typehints.Tuple[str, float, int]]) self.assertReturnType( typehints.Tuple[int, int], reverse, [typehints.List[int]]) def testGetItemSlice(self): self.assertReturnType( typehints.List[int], lambda v: v[::-1], [typehints.List[int]]) self.assertReturnType( typehints.Tuple[int], lambda v: v[::-1], [typehints.Tuple[int]]) self.assertReturnType(str, lambda v: v[::-1], [str]) self.assertReturnType(typehints.Any, lambda v: v[::-1], [typehints.Any]) self.assertReturnType(typehints.Any, lambda v: v[::-1], [object]) # Test binary_subscr on a slice of a Const. test_list = ['a', 'b'] self.assertReturnType(typehints.List[str], lambda: test_list[:], []) def testUnpack(self): def reverse(a_b): (a, b) = a_b return b, a any_tuple = typehints.Tuple[typehints.Any, typehints.Any] self.assertReturnType( typehints.Tuple[int, float], reverse, [typehints.Tuple[float, int]]) self.assertReturnType( typehints.Tuple[int, int], reverse, [typehints.Tuple[int, ...]]) self.assertReturnType( typehints.Tuple[int, int], reverse, [typehints.List[int]]) self.assertReturnType( typehints.Tuple[typehints.Union[int, float, str], typehints.Union[int, float, str]], reverse, [typehints.Tuple[int, float, str]]) self.assertReturnType(any_tuple, reverse, [typehints.Any]) self.assertReturnType( typehints.Tuple[int, float], reverse, [trivial_inference.Const((1.0, 1))]) self.assertReturnType( any_tuple, reverse, [trivial_inference.Const((1, 2, 3))]) def testBuildMap(self): self.assertReturnType( typehints.Dict[typehints.Any, typehints.Any], lambda k, v: {}, [int, float]) self.assertReturnType( typehints.Dict[int, float], lambda k, v: {k: v}, [int, float]) self.assertReturnType( typehints.Tuple[str, typehints.Dict[int, float]], lambda k, v: ('s', { k: v }), [int, float]) self.assertReturnType( typehints.Dict[int, typehints.Union[float, str]], lambda k1, v1, k2, v2: { k1: v1, k2: v2 }, [int, float, int, str]) # Constant map. self.assertReturnType( typehints.Dict[str, typehints.Union[int, float]], lambda a, b: { 'a': a, 'b': b }, [int, float]) self.assertReturnType( typehints.Tuple[int, typehints.Dict[str, typehints.Union[int, float]]], lambda a, b: (4, { 'a': a, 'b': b }), [int, float]) def testNoneReturn(self): def func(a): if a == 5: return a return None self.assertReturnType(typehints.Union[int, type(None)], func, [int]) def testSimpleList(self): self.assertReturnType( typehints.List[int], lambda xs: [1, 2], [typehints.Tuple[int, ...]]) self.assertReturnType( typehints.List[typehints.Any], lambda xs: list(xs), # List is a disallowed builtin [typehints.Tuple[int, ...]]) def testListComprehension(self): self.assertReturnType( typehints.List[int], lambda xs: [x for x in xs], [typehints.Tuple[int, ...]]) def testTupleListComprehension(self): self.assertReturnType( typehints.List[int], lambda xs: [x for x in xs], [typehints.Tuple[int, int, int]]) self.assertReturnType( typehints.List[typehints.Union[int, float]], lambda xs: [x for x in xs], [typehints.Tuple[int, float]]) expected = typehints.List[typehints.Tuple[str, int]] self.assertReturnType( expected, lambda kvs: [(kvs[0], v) for v in kvs[1]], [typehints.Tuple[str, typehints.Iterable[int]]]) self.assertReturnType( typehints.List[typehints.Tuple[str, typehints.Union[str, int], int]], lambda L: [(a, a or b, b) for a, b in L], [typehints.Iterable[typehints.Tuple[str, int]]]) def testGenerator(self): def foo(x, y): yield x yield y self.assertReturnType(typehints.Iterable[int], foo, [int, int]) self.assertReturnType( typehints.Iterable[typehints.Union[int, float]], foo, [int, float]) def testGeneratorComprehension(self): self.assertReturnType( typehints.Iterable[int], lambda xs: (x for x in xs), [typehints.Tuple[int, ...]]) def testBinOp(self): self.assertReturnType(int, lambda a, b: a + b, [int, int]) self.assertReturnType(int, lambda a: a + 1, [int]) self.assertReturnType( typehints.Any, lambda a, b: a + b, [int, typehints.Any]) self.assertReturnType( typehints.List[typehints.Union[int, str]], lambda a, b: a + b, [typehints.List[int], typehints.List[str]]) def testBinOpPromotion(self): self.assertReturnType(int, lambda a, b: a + b, [int, bool]) self.assertReturnType(float, lambda a, b: a + b, [int, float]) self.assertReturnType(complex, lambda a, b: a + b, [int, complex]) def testBinOpSequenceMul(self): self.assertReturnType(str, lambda a, b: a * b, [int, str]) self.assertReturnType(bytes, lambda a, b: a * b, [bytes, int]) self.assertReturnType( typehints.List[float], lambda a, b: a * b, [int, typehints.List[float]]) def testDiv(self): # We only support Python 3 now. self.assertReturnType(float, lambda a, b: a / b, [int, int]) self.assertReturnType(int, lambda a, b: a // b, [int, int]) def testCall(self): f = lambda x, *args: x self.assertReturnType( typehints.Tuple[int, float], lambda: (f(1), f(2.0, 3))) # We could do better here, but this is at least correct. self.assertReturnType( typehints.Tuple[int, typehints.Any], lambda: (1, f(x=1.0))) def testCallNullaryMethod(self): class Foo: pass self.assertReturnType( typehints.Tuple[Foo, typehints.Any], lambda x: (x, x.unknown()), [Foo]) def testCallNestedLambda(self): class Foo: pass self.assertReturnType( typehints.Tuple[Foo, int], lambda x: (x, (lambda: 3)()), [Foo]) def testClosure(self): x = 1 y = 1.0 self.assertReturnType(typehints.Tuple[int, float], lambda: (x, y)) @unittest.skip("https://github.com/apache/beam/issues/28420") def testLocalClosure(self): self.assertReturnType( typehints.Tuple[int, int], lambda x: (x, (lambda: x)()), [int]) def testGlobals(self): self.assertReturnType(int, lambda: global_int) def testBuiltins(self): self.assertReturnType(int, lambda x: len(x), [typehints.Any]) def testGetAttr(self): self.assertReturnType( typehints.Tuple[str, typehints.Any], lambda: (typehints.__doc__, typehints.fake)) def testSetAttr(self): def fn(obj, flag): if flag == 1: obj.attr = 1 res = 1 elif flag == 2: obj.attr = 2 res = 1.5 return res self.assertReturnType(typehints.Union[int, float], fn, [int]) def testSetDeleteGlobal(self): def fn(flag): # pylint: disable=global-variable-undefined global global_var if flag == 1: global_var = 3 res = 1 elif flag == 4: del global_var res = "str" return res self.assertReturnType(typehints.Union[int, str], fn, [int]) def testMethod(self): class A(object): def m(self, x): return x self.assertReturnType(int, lambda: A().m(3)) self.assertReturnType(float, lambda: A.m(A(), 3.0)) def testCallFunctionOnAny(self): # Tests inference when CALL_FUNCTION/CALL_METHOD's function argument is Any. # The function cannot be called but inference should continue. Also tests # that LOAD_ATTR/LOAD_METHOD implementations don't load builtin functions, # which also break inference since they don't disassemble. def call_function_on_any(s): # str.split is a builtin so opcodes.load_attr (load_method in Py3.7+) # should put Any on the stack. # If infer_return_type_func raises while trying to simulate CALL_FUNCTION # on Any, the result will be Any instead of int. s.split() return 0 self.assertReturnType(int, call_function_on_any, [str]) def testAlwaysReturnsEarly(self): def some_fn(v): if v: return 1 return 2 self.assertReturnType(int, some_fn) def testDict(self): self.assertReturnType( typehints.Dict[typehints.Any, typehints.Any], lambda: {}) # yapf: disable def testDictComprehension(self): fields = [] expected_type = typehints.Dict[typehints.Any, typehints.Any] self.assertReturnType( expected_type, lambda row: {f: row[f] for f in fields}, [typehints.Any]) def testDictComprehensionSimple(self): self.assertReturnType( typehints.Dict[str, int], lambda _list: {'a': 1 for _ in _list}, []) def testSet(self): self.assertReturnType( typehints.Set[typehints.Union[()]], lambda: {x for x in ()}) self.assertReturnType( typehints.Set[int], lambda xs: {x for x in xs}, [typehints.List[int]]) # yapf: enable def testDepthFunction(self): def f(i): return i self.assertReturnType(typehints.Any, lambda i: f(i), [int], depth=0) self.assertReturnType(int, lambda i: f(i), [int], depth=1) def testDepthMethod(self): class A(object): def m(self, x): return x self.assertReturnType(typehints.Any, lambda: A().m(3), depth=0) self.assertReturnType(int, lambda: A().m(3), depth=1) self.assertReturnType(typehints.Any, lambda: A.m(A(), 3.0), depth=0) self.assertReturnType(float, lambda: A.m(A(), 3.0), depth=1) def testBuildTupleUnpackWithCall(self): # Lambda uses BUILD_TUPLE_UNPACK_WITH_CALL opcode in Python 3.6, 3.7. def fn(x1, x2, *unused_args): return x1, x2 self.assertReturnType( typehints.Tuple[typehints.Union[str, float, int], typehints.Union[str, float, int]], lambda x1, x2, _list: fn(x1, x2, *_list), [str, float, typehints.List[int]]) # No *args self.assertReturnType( typehints.Tuple[typehints.Union[str, typehints.List[int]], typehints.Union[str, typehints.List[int]]], lambda x1, x2, _list: fn(x1, x2, *_list), [str, typehints.List[int]]) def testCallFunctionEx(self): # Test when fn arguments are built using BUiLD_LIST. def fn(*args): return args self.assertReturnType( typehints.List[typehints.Union[str, float]], lambda x1, x2: fn(*[x1, x2]), [str, float]) def testCallFunctionExKwargs(self): def fn(x1, x2, **unused_kwargs): return x1, x2 # Keyword args are currently unsupported for CALL_FUNCTION_EX. self.assertReturnType( typehints.Any, lambda x1, x2, _dict: fn(x1, x2, **_dict), [str, float, typehints.List[int]]) def testInstanceToType(self): class MyClass(object): def method(self): pass test_cases = [ (typehints.Dict[str, int], { 'a': 1 }), (typehints.Dict[str, typehints.Union[str, int]], { 'a': 1, 'b': 'c' }), (typehints.Dict[typehints.Any, typehints.Any], {}), (typehints.Set[str], {'a'}), (typehints.Set[typehints.Union[str, float]], {'a', 0.4}), (typehints.Set[typehints.Any], set()), (typehints.FrozenSet[str], frozenset(['a'])), ( typehints.FrozenSet[typehints.Union[str, float]], frozenset(['a', 0.4])), (typehints.FrozenSet[typehints.Any], frozenset()), (typehints.Tuple[int], (1, )), (typehints.Tuple[int, int, str], (1, 2, '3')), (typehints.Tuple[()], ()), (typehints.List[int], [1]), (typehints.List[typehints.Union[int, str]], [1, 'a']), (typehints.List[typehints.Any], []), (type(None), None), (type(MyClass), MyClass), (MyClass, MyClass()), (type(MyClass.method), MyClass.method), (types.MethodType, MyClass().method), (row_type.RowTypeConstraint.from_fields([('x', int)]), beam.Row(x=37)), ] for expected_type, instance in test_cases: self.assertEqual( expected_type, trivial_inference.instance_to_type(instance), msg=instance) def testRow(self): self.assertReturnType( row_type.RowTypeConstraint.from_fields([('x', int), ('y', str)]), lambda x, y: beam.Row(x=x + 1, y=y), [int, str]) self.assertReturnType( row_type.RowTypeConstraint.from_fields([('x', int), ('y', str)]), lambda x: beam.Row(x=x, y=str(x)), [int]) def testRowAttr(self): self.assertReturnType( typehints.Tuple[int, str], lambda row: (row.x, getattr(row, 'y')), [row_type.RowTypeConstraint.from_fields([('x', int), ('y', str)])]) def testRowMissingAttr(self): self.assertReturnType( typehints.Any, lambda row: getattr(row, '_asdict'), [row_type.RowTypeConstraint.from_fields([('x', int), ('y', str)])]) def testFString(self): self.assertReturnType(str, lambda x, y: f'{x}: {y:0.2}', [str, float]) def testPyCallable(self): self.assertReturnType( typehints.Tuple[int, str], python_callable.PythonCallableWithSource("lambda x: (x, str(x))"), [int]) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/typecheck.py0000644000175100017510000003170015206113033022237 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Runtime type checking support. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file import inspect import sys import types from collections import abc from apache_beam import pipeline from apache_beam.pvalue import TaggedOutput from apache_beam.transforms import core from apache_beam.transforms.core import DoFn from apache_beam.transforms.window import TimestampedValue from apache_beam.transforms.window import WindowedValue from apache_beam.typehints.decorators import GeneratorWrapper from apache_beam.typehints.decorators import TypeCheckError from apache_beam.typehints.decorators import _check_instance_type from apache_beam.typehints.decorators import getcallargs_forhints from apache_beam.typehints.typehints import CompositeTypeHintError from apache_beam.typehints.typehints import SimpleTypeHintError from apache_beam.typehints.typehints import check_constraint from apache_beam.typehints.typehints import normalize class AbstractDoFnWrapper(DoFn): """An abstract class to create wrapper around DoFn""" def __init__(self, dofn): super().__init__() self.dofn = dofn def __getattribute__(self, name): if (name.startswith('_') or name in self.__dict__ or hasattr(type(self), name)): return object.__getattribute__(self, name) else: return getattr(self.dofn, name) def _inspect_start_bundle(self): return self.dofn.get_function_arguments('start_bundle') def _inspect_process(self): return self.dofn.get_function_arguments('process') def _inspect_finish_bundle(self): return self.dofn.get_function_arguments('finish_bundle') def wrapper(self, method, args, kwargs): return method(*args, **kwargs) def setup(self): return self.dofn.setup() def start_bundle(self, *args, **kwargs): return self.wrapper(self.dofn.start_bundle, args, kwargs) def process(self, *args, **kwargs): return self.wrapper(self.dofn.process, args, kwargs) def finish_bundle(self, *args, **kwargs): return self.wrapper(self.dofn.finish_bundle, args, kwargs) def teardown(self): return self.dofn.teardown() class OutputCheckWrapperDoFn(AbstractDoFnWrapper): """A DoFn that verifies against common errors in the output type.""" def __init__(self, dofn, full_label): super().__init__(dofn) self.full_label = full_label def wrapper(self, method, args, kwargs): try: result = method(*args, **kwargs) except TypeCheckError as e: # TODO(BEAM-10710): Remove the 'ParDo' prefix for the label name error_msg = ( 'Runtime type violation detected within ParDo(%s): ' '%s' % (self.full_label, e)) _, _, tb = sys.exc_info() raise TypeCheckError(error_msg).with_traceback(tb) else: return self._check_type(result) @staticmethod def _check_type(output): if output is None: return output elif isinstance(output, (dict, bytes, str)): object_type = type(output).__name__ raise TypeCheckError( 'Returning a %s from a ParDo or FlatMap is ' 'discouraged. Please use list("%s") if you really ' 'want this behavior.' % (object_type, output)) elif not isinstance(output, abc.Iterable): raise TypeCheckError( 'FlatMap and ParDo must return an ' 'iterable. %s was returned instead.' % type(output)) return output class TypeCheckWrapperDoFn(AbstractDoFnWrapper): """A wrapper around a DoFn which performs type-checking of input and output. """ def __init__(self, dofn, type_hints, label=None): super().__init__(dofn) self._process_fn = self.dofn._process_argspec_fn() if type_hints.input_types: input_args, input_kwargs = type_hints.input_types self._input_hints = getcallargs_forhints( self._process_fn, *input_args, **input_kwargs) else: self._input_hints = None # TODO(robertwb): Multi-output. self._output_type_hint = type_hints.simple_output_type(label) def wrapper(self, method, args, kwargs): result = method(*args, **kwargs) return self._type_check_result(result) def process(self, *args, **kwargs): if self._input_hints: actual_inputs = inspect.getcallargs(self._process_fn, *args, **kwargs) # pylint: disable=deprecated-method for var, hint in self._input_hints.items(): if hint is actual_inputs[var]: # self parameter continue _check_instance_type(hint, actual_inputs[var], var, True) return self._type_check_result(self.dofn.process(*args, **kwargs)) def _type_check_result(self, transform_results): if self._output_type_hint is None or transform_results is None: return transform_results def type_check_output(o): if isinstance(o, TimestampedValue) and hasattr(o, "__orig_class__"): # when a typed TimestampedValue is set, check the value type x = o.value # per https://stackoverflow.com/questions/57706180/, # __orig_class__ is te safe way to obtain the actual type # from from Generic[T], supported since Python 3.5.3 beam_type = normalize(o.__orig_class__.__args__[0]) self.type_check(beam_type, x, is_input=False) else: # TODO(robertwb): Multi-output. x = o.value if isinstance(o, (TaggedOutput, WindowedValue)) else o self.type_check(self._output_type_hint, x, is_input=False) # If the return type is a generator, then we will need to interleave our # type-checking with its normal iteration so we don't deplete the # generator initially just by type-checking its yielded contents. if isinstance(transform_results, types.GeneratorType): return GeneratorWrapper(transform_results, type_check_output) for o in transform_results: type_check_output(o) return transform_results @staticmethod def type_check(type_constraint, datum, is_input): """Typecheck a PTransform related datum according to a type constraint. This function is used to optionally type-check either an input or an output to a PTransform. Args: type_constraint: An instance of a typehints.TypeContraint, one of the white-listed builtin Python types, or a custom user class. datum: An instance of a Python object. is_input: True if 'datum' is an input to a PTransform's DoFn. False otherwise. Raises: TypeError: If 'datum' fails to type-check according to 'type_constraint'. """ datum_type = 'input' if is_input else 'output' try: check_constraint(type_constraint, datum) except CompositeTypeHintError as e: _, _, tb = sys.exc_info() raise TypeCheckError(e.args[0]).with_traceback(tb) except SimpleTypeHintError: error_msg = ( "According to type-hint expected %s should be of type %s. " "Instead, received '%s', an instance of type %s." % (datum_type, type_constraint, datum, type(datum))) _, _, tb = sys.exc_info() raise TypeCheckError(error_msg).with_traceback(tb) class TypeCheckCombineFn(core.CombineFn): """A wrapper around a CombineFn performing type-checking of input and output. """ def __init__(self, combinefn, type_hints, label=None): self._combinefn = combinefn self._input_type_hint = type_hints.input_types self._output_type_hint = type_hints.simple_output_type(label) self._label = label def setup(self, *args, **kwargs): self._combinefn.setup(*args, **kwargs) def create_accumulator(self, *args, **kwargs): return self._combinefn.create_accumulator(*args, **kwargs) def add_input(self, accumulator, element, *args, **kwargs): if self._input_type_hint: try: _check_instance_type( self._input_type_hint[0][0].tuple_types[1], element, 'element', True) except TypeCheckError as e: error_msg = ( 'Runtime type violation detected within %s: ' '%s' % (self._label, e)) _, _, tb = sys.exc_info() raise TypeCheckError(error_msg).with_traceback(tb) return self._combinefn.add_input(accumulator, element, *args, **kwargs) def merge_accumulators(self, accumulators, *args, **kwargs): return self._combinefn.merge_accumulators(accumulators, *args, **kwargs) def compact(self, accumulator, *args, **kwargs): return self._combinefn.compact(accumulator, *args, **kwargs) def extract_output(self, accumulator, *args, **kwargs): result = self._combinefn.extract_output(accumulator, *args, **kwargs) if self._output_type_hint: try: _check_instance_type( self._output_type_hint.tuple_types[1], result, None, True) except TypeCheckError as e: error_msg = ( 'Runtime type violation detected within %s: ' '%s' % (self._label, e)) _, _, tb = sys.exc_info() raise TypeCheckError(error_msg).with_traceback(tb) return result def teardown(self, *args, **kwargs): self._combinefn.teardown(*args, **kwargs) class TypeCheckVisitor(pipeline.PipelineVisitor): _in_combine = False def enter_composite_transform(self, applied_transform): if isinstance(applied_transform.transform, core.CombinePerKey): self._in_combine = True self._wrapped_fn = applied_transform.transform.fn = TypeCheckCombineFn( applied_transform.transform.fn, applied_transform.transform.get_type_hints(), applied_transform.full_label) def leave_composite_transform(self, applied_transform): if isinstance(applied_transform.transform, core.CombinePerKey): self._in_combine = False def visit_transform(self, applied_transform): transform = applied_transform.transform if isinstance(transform, core.ParDo): if self._in_combine: if isinstance(transform.fn, core.CombineValuesDoFn): transform.fn.combinefn = self._wrapped_fn else: transform.fn = transform.dofn = OutputCheckWrapperDoFn( TypeCheckWrapperDoFn( transform.fn, transform.get_type_hints(), applied_transform.full_label), applied_transform.full_label) class PerformanceTypeCheckVisitor(pipeline.PipelineVisitor): def visit_transform(self, applied_transform): transform = applied_transform.transform full_label = applied_transform.full_label # Store output type hints in current transform output_type_hints = self.get_output_type_hints(transform) if output_type_hints: transform._add_type_constraint_from_consumer( full_label, output_type_hints) # Store input type hints in producer transform input_type_hints = self.get_input_type_hints(transform) if input_type_hints and len(applied_transform.inputs): producer = applied_transform.inputs[0].producer if producer: producer.transform._add_type_constraint_from_consumer( full_label, input_type_hints) def get_input_type_hints(self, transform): type_hints = transform.get_type_hints() input_types = None if type_hints.input_types: normal_hints, kwarg_hints = type_hints.input_types if kwarg_hints: input_types = kwarg_hints if normal_hints: input_types = normal_hints parameter_name = 'Unknown Parameter' if hasattr(transform, 'fn'): try: argspec = inspect.getfullargspec(transform.fn._process_argspec_fn()) except TypeError: # An unsupported callable was passed to getfullargspec pass else: if len(argspec.args): arg_index = 0 if argspec.args[0] == 'self' and len(argspec.args) > 1: arg_index = 1 parameter_name = argspec.args[arg_index] if isinstance(input_types, dict): input_types = (input_types[argspec.args[arg_index]], ) if input_types and len(input_types): input_types = input_types[0] return parameter_name, input_types def get_output_type_hints(self, transform): type_hints = transform.get_type_hints() output_types = None if type_hints.output_types: normal_hints, kwarg_hints = type_hints.output_types if kwarg_hints: output_types = kwarg_hints if normal_hints: output_types = normal_hints if output_types and len(output_types): output_types = output_types[0] return None, output_types ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/typecheck_test.py0000644000175100017510000002522715206113033023305 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Unit tests for typecheck. See additional runtime_type_check=True tests in ptransform_test.py. """ # pytype: skip-file # ruff: noqa: UP006 import os import tempfile import unittest from typing import Iterable from typing import Tuple import apache_beam as beam from apache_beam import Pipeline from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import TypeOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.typehints import decorators from apache_beam.typehints import with_input_types from apache_beam.typehints import with_output_types decorators._enable_from_callable = True # Disable frequent lint warning due to pipe operator for chaining transforms. # pylint: disable=expression-not-assigned class MyDoFn(beam.DoFn): def __init__(self, output_filename): super().__init__() self.output_filename = output_filename def _output(self): """Returns a file used to record function calls.""" if not hasattr(self, 'output_file'): self._output_file = open(self.output_filename, 'at', buffering=1) return self._output_file def start_bundle(self): self._output().write('start_bundle\n') def finish_bundle(self): self._output().write('finish_bundle\n') def setup(self): self._output().write('setup\n') def teardown(self): self._output().write('teardown\n') self._output().close() def process(self, element: int, *args, **kwargs) -> Iterable[int]: self._output().write('process\n') yield element class MyDoFnBadAnnotation(MyDoFn): def process(self, element: int, *args, **kwargs) -> int: # Should raise an exception about return type not being iterable. return super().process() class RuntimeTypeCheckTest(unittest.TestCase): def setUp(self): # Use FnApiRunner since it guarantees all lifecycle methods will be called. self.p = TestPipeline( 'FnApiRunner', options=PipelineOptions( runtime_type_check=True, performance_runtime_type_check=False)) def test_setup(self): # Verifies that runtime type checking is enabled for test cases. def fn(e: int) -> int: return str(e) # type: ignore with self.assertRaisesRegex(Exception, r'output should be.*int.*received.*str'): _ = self.p | beam.Create([1, 2, 3]) | beam.Map(fn) self.p.run() def test_wrapper_pass_through(self): # We use a file to check the result because the MyDoFn instance passed is # not the same one that actually runs in the pipeline (it is serialized # here and deserialized in the worker). with tempfile.TemporaryDirectory() as tmp_dirname: path = os.path.join(tmp_dirname + "tmp_filename") dofn = MyDoFn(path) result = self.p | beam.Create([1, 2, 3]) | beam.ParDo(dofn) assert_that(result, equal_to([1, 2, 3])) self.p.run() with open(path, mode="r") as ft: lines = [line.strip() for line in ft] self.assertListEqual([ 'setup', 'start_bundle', 'process', 'process', 'process', 'finish_bundle', 'teardown', ], lines) def test_wrapper_pipeline_type_check(self): # Verifies that type hints are not masked by the wrapper. What actually # happens is that the wrapper is applied during self.p.run() (not invoked # in this case), while pipeline type checks happen during pipeline creation. # Thus, the wrapper does not have to implement: default_type_hints, # infer_output_type, get_type_hints. with tempfile.NamedTemporaryFile(mode='w+t') as f: dofn = MyDoFnBadAnnotation(f.name) with self.assertRaisesRegex(ValueError, r'int.*is not iterable'): _ = self.p | beam.Create([1, 2, 3]) | beam.ParDo(dofn) class PerformanceRuntimeTypeCheckTest(unittest.TestCase): def setUp(self): self.p = Pipeline( options=PipelineOptions( performance_runtime_type_check=True, pipeline_type_check=False)) def assertStartswith(self, msg, prefix): self.assertTrue( msg.startswith(prefix), '"%s" does not start with "%s"' % (msg, prefix)) def test_simple_input_error(self): with self.assertRaisesRegex(Exception, "Type-hint for argument: 'x' violated. " "Expected an instance of {}, " "instead found 1, an instance of {}".format( str, int)): ( self.p | beam.Create([1, 1]) | beam.FlatMap(lambda x: [int(x)]).with_input_types( str).with_output_types(int)) self.p.run().wait_until_finish() def test_simple_output_error(self): with self.assertRaisesRegex(Exception, "Type-hint for argument: 'x' violated. "): ( self.p | beam.Create(['1', '1']) | beam.FlatMap(lambda x: [int(x)]).with_input_types( int).with_output_types(int)) self.p.run().wait_until_finish() def test_simple_input_error_with_kwarg_typehints(self): @with_input_types(element=int) @with_output_types(int) class ToInt(beam.DoFn): def process(self, element, *args, **kwargs): yield int(element) with self.assertRaisesRegex(Exception, "Type-hint for argument: 'element' violated"): (self.p | beam.Create(['1', '1']) | beam.ParDo(ToInt())) self.p.run().wait_until_finish() def test_bad_flatten_input(self): with self.assertRaisesRegex( TypeError, "Inputs to Flatten cannot include an iterable of PCollections. "): with beam.Pipeline() as p: pc = p | beam.Create([1, 1]) flatten_inputs = [pc, (pc, )] flatten_inputs | beam.Flatten() def test_do_fn_returning_non_iterable_throws_error(self): # This function is incorrect because it returns a non-iterable object def incorrect_par_do_fn(x): return x + 5 with self.assertRaisesRegex(Exception, "'int' object is not iterable "): (self.p | beam.Create([1, 1]) | beam.FlatMap(incorrect_par_do_fn)) self.p.run().wait_until_finish() def test_simple_type_satisfied(self): @with_input_types(int, int) @with_output_types(int) class AddWithNum(beam.DoFn): def process(self, element, num): return [element + num] results = ( self.p | 'T' >> beam.Create([1, 2, 3]).with_output_types(int) | 'Add' >> beam.ParDo(AddWithNum(), 1)) assert_that(results, equal_to([2, 3, 4])) self.p.run() def test_simple_type_violation(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False @with_output_types(str) @with_input_types(x=int) def int_to_string(x): return str(x) ( self.p | 'Create' >> beam.Create(['some_string']) | 'ToStr' >> beam.Map(int_to_string)) with self.assertRaisesRegex(Exception, "Type-hint for argument: 'x' violated. "): self.p.run().wait_until_finish() def test_pipeline_checking_satisfied_but_run_time_types_violate(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False @with_output_types(Tuple[bool, int]) @with_input_types(a=int) def is_even_as_key(a): # Simulate a programming error, should be: return (a % 2 == 0, a) # However this returns Tuple[int, int] return (a % 2, a) ( self.p | 'Nums' >> beam.Create(range(1)).with_output_types(int) | 'IsEven' >> beam.Map(is_even_as_key) | 'Parity' >> beam.GroupByKey()) with self.assertRaisesRegex( Exception, ("Type-hint for return type violated: Tuple\\[, \\] hint type-constraint violated. The type of element #0 in " + "the passed tuple is incorrect. Expected an instance of type , instead received an instance of type int.")): self.p.run().wait_until_finish() def test_pipeline_runtime_checking_violation_composite_type_output(self): self.p._options.view_as(TypeOptions).pipeline_type_check = False # The type-hinted applied via the 'returns()' method indicates the ParDo # should return an instance of type: Tuple[float, int]. However, an instance # of 'int' will be generated instead. with self.assertRaisesRegex( Exception, ("Type-hint for return type violated. Expected an instance of {}, " + "instead found 4.0, an instance of {}.").format(int, float)): ( self.p | beam.Create([(1, 3.0)]) | ( 'Swap' >> beam.FlatMap(lambda x_y1: [x_y1[0] + x_y1[1]]).with_input_types( Tuple[int, float]).with_output_types(int))) self.p.run().wait_until_finish() def test_downstream_input_type_hint_error_has_descriptive_error_msg(self): @with_input_types(int) @with_output_types(int) class IntToInt(beam.DoFn): def process(self, element, *args, **kwargs): yield element @with_input_types(str) @with_output_types(int) class StrToInt(beam.DoFn): def process(self, element, *args, **kwargs): yield int(element) # This will raise a type check error in IntToInt even though the actual # type check error won't happen until StrToInt. The user will be told that # StrToInt's input type hints were not satisfied while running IntToInt. with self.assertRaisesRegex( Exception, ("Type-hint for argument: 'element' violated. Expected an instance of " + "{}, instead found 9, an instance of {}.").format(str, int)): ( self.p | beam.Create([9]) | beam.ParDo(IntToInt()) | beam.ParDo(StrToInt())) self.p.run().wait_until_finish() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/typed_pipeline_test.py0000644000175100017510000011104415206113033024331 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the type-hint objects and decorators.""" # pytype: skip-file # ruff: noqa: UP006 import typing import unittest from typing import Tuple import apache_beam as beam from apache_beam import pvalue from apache_beam import typehints from apache_beam.options.pipeline_options import OptionsContext from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.typehints import WithTypeHints from apache_beam.typehints.decorators import get_signature # These test often construct a pipeline as value | PTransform to test side # effects (e.g. errors). # pylint: disable=expression-not-assigned class MainInputTest(unittest.TestCase): def assertStartswith(self, msg, prefix): self.assertTrue( msg.startswith(prefix), '"%s" does not start with "%s"' % (msg, prefix)) def test_bad_main_input(self): @typehints.with_input_types(str, int) def repeat(s, times): return s * times with self.assertRaises(typehints.TypeCheckError): [1, 2, 3] | beam.Map(repeat, 3) def test_non_function(self): result = ['a', 'bb', 'c'] | beam.Map(str.upper) self.assertEqual(['A', 'BB', 'C'], sorted(result)) result = ['xa', 'bbx', 'xcx'] | beam.Map(str.strip, 'x') self.assertEqual(['a', 'bb', 'c'], sorted(result)) result = ['1', '10', '100'] | beam.Map(int) self.assertEqual([1, 10, 100], sorted(result)) result = ['1', '10', '100'] | beam.Map(int, 16) self.assertEqual([1, 16, 256], sorted(result)) def test_non_function_fails(self): with self.assertRaises(typehints.TypeCheckError): [1, 2, 3] | beam.Map(str.upper) def test_loose_bounds(self): @typehints.with_input_types(typing.Union[int, float]) @typehints.with_output_types(str) def format_number(x): return '%g' % x result = [1, 2, 3] | beam.Map(format_number) self.assertEqual(['1', '2', '3'], sorted(result)) def test_typed_dofn_class(self): @typehints.with_input_types(int) @typehints.with_output_types(str) class MyDoFn(beam.DoFn): def process(self, element): return [str(element)] result = [1, 2, 3] | beam.ParDo(MyDoFn()) self.assertEqual(['1', '2', '3'], sorted(result)) with self.assertRaisesRegex(typehints.TypeCheckError, r'requires.*int.*applied.*str'): ['a', 'b', 'c'] | beam.ParDo(MyDoFn()) with self.assertRaisesRegex(typehints.TypeCheckError, r'requires.*int.*applied.*str'): [1, 2, 3] | (beam.ParDo(MyDoFn()) | 'again' >> beam.ParDo(MyDoFn())) def test_typed_dofn_method(self): class MyDoFn(beam.DoFn): def process(self, element: int) -> typehints.Tuple[str]: return tuple(str(element)) result = [1, 2, 3] | beam.ParDo(MyDoFn()) self.assertEqual(['1', '2', '3'], sorted(result)) with self.assertRaisesRegex(typehints.TypeCheckError, r'requires.*int.*applied.*str'): _ = ['a', 'b', 'c'] | beam.ParDo(MyDoFn()) with self.assertRaisesRegex(typehints.TypeCheckError, r'requires.*int.*applied.*str'): _ = [1, 2, 3] | (beam.ParDo(MyDoFn()) | 'again' >> beam.ParDo(MyDoFn())) def test_typed_dofn_method_with_class_decorators(self): # Class decorators take precedence over PEP 484 hints. @typehints.with_input_types(typehints.Tuple[int, int]) @typehints.with_output_types(int) class MyDoFn(beam.DoFn): def process(self, element: int) -> typehints.Tuple[str]: yield element[0] result = [(1, 2)] | beam.ParDo(MyDoFn()) self.assertEqual([1], sorted(result)) with self.assertRaisesRegex( typehints.TypeCheckError, r'requires.*Tuple\[, \].*applied.*str'): _ = ['a', 'b', 'c'] | beam.ParDo(MyDoFn()) with self.assertRaisesRegex( typehints.TypeCheckError, r'requires.*Tuple\[, \].*applied.*int'): _ = [1, 2, 3] | (beam.ParDo(MyDoFn()) | 'again' >> beam.ParDo(MyDoFn())) def test_typed_callable_iterable_output(self): # Only the outer Iterable should be stripped. def do_fn(element: int) -> typehints.Iterable[typehints.Iterable[str]]: return [[str(element)] * 2] result = [1, 2] | beam.ParDo(do_fn) self.assertEqual([['1', '1'], ['2', '2']], sorted(result)) def test_typed_dofn_instance(self): # Type hints applied to DoFn instance take precedence over decorators and # process annotations. @typehints.with_input_types(typehints.Tuple[int, int]) @typehints.with_output_types(int) class MyDoFn(beam.DoFn): def process(self, element: typehints.Tuple[int, int]) -> \ typehints.List[int]: return [str(element)] my_do_fn = MyDoFn().with_input_types(int).with_output_types(str) result = [1, 2, 3] | beam.ParDo(my_do_fn) self.assertEqual(['1', '2', '3'], sorted(result)) with self.assertRaisesRegex(typehints.TypeCheckError, r'requires.*int.*applied.*str'): _ = ['a', 'b', 'c'] | beam.ParDo(my_do_fn) with self.assertRaisesRegex(typehints.TypeCheckError, r'requires.*int.*applied.*str'): _ = [1, 2, 3] | (beam.ParDo(my_do_fn) | 'again' >> beam.ParDo(my_do_fn)) def test_typed_callable_instance(self): # Type hints applied to ParDo instance take precedence over callable # decorators and annotations. @typehints.with_input_types(typehints.Tuple[int, int]) @typehints.with_output_types(typehints.Generator[int]) def do_fn(element: typehints.Tuple[int, int]) -> typehints.Generator[str]: yield str(element) pardo = beam.ParDo(do_fn).with_input_types(int).with_output_types(str) result = [1, 2, 3] | pardo self.assertEqual(['1', '2', '3'], sorted(result)) with self.assertRaisesRegex(typehints.TypeCheckError, r'requires.*int.*applied.*str'): _ = ['a', 'b', 'c'] | pardo with self.assertRaisesRegex(typehints.TypeCheckError, r'requires.*int.*applied.*str'): _ = [1, 2, 3] | (pardo | 'again' >> pardo) def test_filter_type_hint(self): @typehints.with_input_types(int) def filter_fn(data): return data % 2 self.assertEqual([1, 3], [1, 2, 3] | beam.Filter(filter_fn)) def test_partition(self): with TestPipeline() as p: even, odd = (p | beam.Create([1, 2, 3]) | 'even_odd' >> beam.Partition(lambda e, _: e % 2, 2)) self.assertIsNotNone(even.element_type) self.assertIsNotNone(odd.element_type) res_even = ( even | 'IdEven' >> beam.ParDo(lambda e: [e]).with_input_types(int)) res_odd = ( odd | 'IdOdd' >> beam.ParDo(lambda e: [e]).with_input_types(int)) assert_that(res_even, equal_to([2]), label='even_check') assert_that(res_odd, equal_to([1, 3]), label='odd_check') def test_typed_dofn_multi_output(self): class MyDoFn(beam.DoFn): def process(self, element): if element % 2: yield beam.pvalue.TaggedOutput('odd', element) else: yield beam.pvalue.TaggedOutput('even', element) with TestPipeline() as p: res = ( p | beam.Create([1, 2, 3]) | beam.ParDo(MyDoFn()).with_outputs('odd', 'even')) self.assertIsNotNone(res[None].element_type) self.assertIsNotNone(res['even'].element_type) self.assertIsNotNone(res['odd'].element_type) res_main = ( res[None] | 'id_none' >> beam.ParDo(lambda e: [e]).with_input_types(int)) res_even = ( res['even'] | 'id_even' >> beam.ParDo(lambda e: [e]).with_input_types(int)) res_odd = ( res['odd'] | 'id_odd' >> beam.ParDo(lambda e: [e]).with_input_types(int)) assert_that(res_main, equal_to([]), label='none_check') assert_that(res_even, equal_to([2]), label='even_check') assert_that(res_odd, equal_to([1, 3]), label='odd_check') with self.assertRaises(ValueError): _ = res['undeclared tag'] def test_typed_dofn_multi_output_no_tags(self): class MyDoFn(beam.DoFn): def process(self, element): if element % 2: yield beam.pvalue.TaggedOutput('odd', element) else: yield beam.pvalue.TaggedOutput('even', element) with TestPipeline() as p: res = (p | beam.Create([1, 2, 3]) | beam.ParDo(MyDoFn()).with_outputs()) self.assertIsNotNone(res[None].element_type) self.assertIsNotNone(res['even'].element_type) self.assertIsNotNone(res['odd'].element_type) res_main = ( res[None] | 'id_none' >> beam.ParDo(lambda e: [e]).with_input_types(int)) res_even = ( res['even'] | 'id_even' >> beam.ParDo(lambda e: [e]).with_input_types(int)) res_odd = ( res['odd'] | 'id_odd' >> beam.ParDo(lambda e: [e]).with_input_types(int)) assert_that(res_main, equal_to([]), label='none_check') assert_that(res_even, equal_to([2]), label='even_check') assert_that(res_odd, equal_to([1, 3]), label='odd_check') def test_typed_ptransform_fn_pre_hints(self): # Test that type hints are propagated to the created PTransform. # Decorator appears before type hints. This is the more common style. @beam.ptransform_fn @typehints.with_input_types(int) def MyMap(pcoll): return pcoll | beam.ParDo(lambda x: [x]) self.assertListEqual([1, 2, 3], [1, 2, 3] | MyMap()) with self.assertRaises(typehints.TypeCheckError): _ = ['a'] | MyMap() def test_typed_ptransform_fn_post_hints(self): # Test that type hints are propagated to the created PTransform. # Decorator appears after type hints. This style is required for Cython # functions, since they don't accept assigning attributes to them. @typehints.with_input_types(int) @beam.ptransform_fn def MyMap(pcoll): return pcoll | beam.ParDo(lambda x: [x]) self.assertListEqual([1, 2, 3], [1, 2, 3] | MyMap()) with self.assertRaises(typehints.TypeCheckError): _ = ['a'] | MyMap() def test_typed_ptransform_fn_multi_input_types_pos(self): @beam.ptransform_fn @beam.typehints.with_input_types(str, int) def multi_input(pcoll_tuple, additional_arg): _, _ = pcoll_tuple assert additional_arg == 'additional_arg' with TestPipeline() as p: pcoll1 = p | 'c1' >> beam.Create(['a']) pcoll2 = p | 'c2' >> beam.Create([1]) _ = (pcoll1, pcoll2) | multi_input('additional_arg') with self.assertRaises(typehints.TypeCheckError): _ = (pcoll2, pcoll1) | 'fails' >> multi_input('additional_arg') def test_typed_ptransform_fn_multi_input_types_kw(self): @beam.ptransform_fn @beam.typehints.with_input_types(strings=str, integers=int) def multi_input(pcoll_dict, additional_arg): _ = pcoll_dict['strings'] _ = pcoll_dict['integers'] assert additional_arg == 'additional_arg' with TestPipeline() as p: pcoll1 = p | 'c1' >> beam.Create(['a']) pcoll2 = p | 'c2' >> beam.Create([1]) _ = { 'strings': pcoll1, 'integers': pcoll2 } | multi_input('additional_arg') with self.assertRaises(typehints.TypeCheckError): _ = { 'strings': pcoll2, 'integers': pcoll1 } | 'fails' >> multi_input('additional_arg') def test_typed_dofn_method_not_iterable(self): class MyDoFn(beam.DoFn): def process(self, element: int) -> str: return str(element) with self.assertRaisesRegex(ValueError, r'str.*is not iterable'): _ = [1, 2, 3] | beam.ParDo(MyDoFn()) def test_typed_dofn_method_return_none(self): class MyDoFn(beam.DoFn): def process(self, unused_element: int) -> None: pass result = [1, 2, 3] | beam.ParDo(MyDoFn()) self.assertListEqual([], result) def test_typed_dofn_method_return_optional(self): class MyDoFn(beam.DoFn): def process( self, unused_element: int) -> typehints.Optional[typehints.Iterable[int]]: pass result = [1, 2, 3] | beam.ParDo(MyDoFn()) self.assertListEqual([], result) def test_typed_dofn_method_return_optional_not_iterable(self): class MyDoFn(beam.DoFn): def process(self, unused_element: int) -> typehints.Optional[int]: pass with self.assertRaisesRegex(ValueError, r'int.*is not iterable'): _ = [1, 2, 3] | beam.ParDo(MyDoFn()) def test_typed_callable_not_iterable(self): def do_fn(element: int) -> int: return element with self.assertRaisesRegex(typehints.TypeCheckError, r'int.*is not iterable'): _ = [1, 2, 3] | beam.ParDo(do_fn) def test_typed_dofn_kwonly(self): class MyDoFn(beam.DoFn): # TODO(BEAM-5878): A kwonly argument like # timestamp=beam.DoFn.TimestampParam would not work here. def process(self, element: int, *, side_input: str) -> \ typehints.Generator[typehints.Optional[str]]: yield str(element) if side_input else None my_do_fn = MyDoFn() result = [1, 2, 3] | beam.ParDo(my_do_fn, side_input='abc') self.assertEqual(['1', '2', '3'], sorted(result)) with self.assertRaisesRegex(typehints.TypeCheckError, r'requires.*str.*got.*int.*side_input'): _ = [1, 2, 3] | beam.ParDo(my_do_fn, side_input=1) def test_typed_dofn_var_kwargs(self): class MyDoFn(beam.DoFn): def process(self, element: int, **side_inputs: typehints.Dict[str, str]) \ -> typehints.Generator[typehints.Optional[str]]: yield str(element) if side_inputs else None my_do_fn = MyDoFn() result = [1, 2, 3] | beam.ParDo(my_do_fn, foo='abc', bar='def') self.assertEqual(['1', '2', '3'], sorted(result)) with self.assertRaisesRegex(typehints.TypeCheckError, r'requires.*str.*got.*int.*side_inputs'): _ = [1, 2, 3] | beam.ParDo(my_do_fn, a=1) def test_typed_callable_string_literals(self): def do_fn(element: 'int') -> 'typehints.List[str]': return [[str(element)] * 2] result = [1, 2] | beam.ParDo(do_fn) self.assertEqual([['1', '1'], ['2', '2']], sorted(result)) def test_typed_ptransform_fn(self): # Test that type hints are propagated to the created PTransform. @beam.ptransform_fn @typehints.with_input_types(int) def MyMap(pcoll): def fn(element: int): yield element return pcoll | beam.ParDo(fn) self.assertListEqual([1, 2, 3], [1, 2, 3] | MyMap()) with self.assertRaisesRegex(typehints.TypeCheckError, r'int.*got.*str'): _ = ['a'] | MyMap() def test_typed_ptransform_fn_conflicting_hints(self): # In this case, both MyMap and its contained ParDo have separate type # checks (that disagree with each other). @beam.ptransform_fn @typehints.with_input_types(str) def MyMap(pcoll): def fn(element: float): yield element return pcoll | beam.ParDo(fn) with self.assertRaisesRegex(typehints.TypeCheckError, r'ParDo.*requires.*float.*applied.*str'): _ = ['1', '2', '3'] | MyMap() with self.assertRaisesRegex(typehints.TypeCheckError, r'MyMap.*expected.*str.*got.*bytes'): _ = [b'a'] | MyMap() def test_typed_dofn_string_literals(self): class MyDoFn(beam.DoFn): def process(self, element: 'int') -> 'typehints.List[str]': return [[str(element)] * 2] result = [1, 2] | beam.ParDo(MyDoFn()) self.assertEqual([['1', '1'], ['2', '2']], sorted(result)) def test_typed_map(self): def fn(element: int) -> int: return element * 2 result = [1, 2, 3] | beam.Map(fn) self.assertEqual([2, 4, 6], sorted(result)) def test_typed_map_return_optional(self): # None is a valid element value for Map. def fn(element: int) -> typehints.Optional[int]: if element > 1: return element result = [1, 2, 3] | beam.Map(fn) self.assertCountEqual([None, 2, 3], result) def test_typed_flatmap(self): def fn(element: int) -> typehints.Iterable[int]: yield element * 2 result = [1, 2, 3] | beam.FlatMap(fn) self.assertCountEqual([2, 4, 6], result) def test_typed_flatmap_output_hint_not_iterable(self): def fn(element: int) -> int: return element * 2 # This is raised (originally) in strip_iterable. with self.assertRaisesRegex(typehints.TypeCheckError, r'int.*is not iterable'): _ = [1, 2, 3] | beam.FlatMap(fn) def test_typed_flatmap_output_value_not_iterable(self): def fn(element: int) -> typehints.Iterable[int]: return element * 2 # This is raised in runners/common.py (process_outputs). with self.assertRaisesRegex(TypeError, r'int.*is not iterable'): _ = [1, 2, 3] | beam.FlatMap(fn) def test_typed_flatmap_optional(self): def fn(element: int) -> typehints.Optional[typehints.Iterable[int]]: if element > 1: yield element * 2 # Verify that the output type of fn is int and not Optional[int]. def fn2(element: int) -> int: return element result = [1, 2, 3] | beam.FlatMap(fn) | beam.Map(fn2) self.assertCountEqual([4, 6], result) def test_typed_ptransform_with_no_error(self): class StrToInt(beam.PTransform): def expand( self, pcoll: beam.pvalue.PCollection[str]) -> beam.pvalue.PCollection[int]: return pcoll | beam.Map(lambda x: int(x)) class IntToStr(beam.PTransform): def expand( self, pcoll: beam.pvalue.PCollection[int]) -> beam.pvalue.PCollection[str]: return pcoll | beam.Map(lambda x: str(x)) _ = ['1', '2', '3'] | StrToInt() | IntToStr() def test_typed_ptransform_with_bad_typehints(self): class StrToInt(beam.PTransform): def expand( self, pcoll: beam.pvalue.PCollection[str]) -> beam.pvalue.PCollection[int]: return pcoll | beam.Map(lambda x: int(x)) class IntToStr(beam.PTransform): def expand( self, pcoll: beam.pvalue.PCollection[str]) -> beam.pvalue.PCollection[str]: return pcoll | beam.Map(lambda x: str(x)) with self.assertRaisesRegex(typehints.TypeCheckError, "Input type hint violation at IntToStr: " "expected , got "): _ = ['1', '2', '3'] | StrToInt() | IntToStr() def test_typed_ptransform_with_bad_input(self): class StrToInt(beam.PTransform): def expand( self, pcoll: beam.pvalue.PCollection[str]) -> beam.pvalue.PCollection[int]: return pcoll | beam.Map(lambda x: int(x)) class IntToStr(beam.PTransform): def expand( self, pcoll: beam.pvalue.PCollection[int]) -> beam.pvalue.PCollection[str]: return pcoll | beam.Map(lambda x: str(x)) with self.assertRaisesRegex(typehints.TypeCheckError, "Input type hint violation at StrToInt: " "expected , got "): # Feed integers to a PTransform that expects strings _ = [1, 2, 3] | StrToInt() | IntToStr() def test_typed_ptransform_with_partial_typehints(self): class StrToInt(beam.PTransform): def expand(self, pcoll) -> beam.pvalue.PCollection[int]: return pcoll | beam.Map(lambda x: int(x)) class IntToStr(beam.PTransform): def expand( self, pcoll: beam.pvalue.PCollection[int]) -> beam.pvalue.PCollection[str]: return pcoll | beam.Map(lambda x: str(x)) # Feed integers to a PTransform that should expect strings # but has no typehints so it expects any _ = [1, 2, 3] | StrToInt() | IntToStr() def test_typed_ptransform_with_bare_wrappers(self): class StrToInt(beam.PTransform): def expand( self, pcoll: beam.pvalue.PCollection) -> beam.pvalue.PCollection: return pcoll | beam.Map(lambda x: int(x)) class IntToStr(beam.PTransform): def expand( self, pcoll: beam.pvalue.PCollection[int]) -> beam.pvalue.PCollection[str]: return pcoll | beam.Map(lambda x: str(x)) _ = [1, 2, 3] | StrToInt() | IntToStr() def test_typed_ptransform_with_no_typehints(self): class StrToInt(beam.PTransform): def expand(self, pcoll): return pcoll | beam.Map(lambda x: int(x)) class IntToStr(beam.PTransform): def expand( self, pcoll: beam.pvalue.PCollection[int]) -> beam.pvalue.PCollection[str]: return pcoll | beam.Map(lambda x: str(x)) # Feed integers to a PTransform that should expect strings # but has no typehints so it expects any _ = [1, 2, 3] | StrToInt() | IntToStr() def test_typed_ptransform_with_generic_annotations(self): T = typing.TypeVar('T') class IntToInt(beam.PTransform): def expand( self, pcoll: beam.pvalue.PCollection[T]) -> beam.pvalue.PCollection[T]: return pcoll | beam.Map(lambda x: x) class IntToStr(beam.PTransform): def expand( self, pcoll: beam.pvalue.PCollection[T]) -> beam.pvalue.PCollection[str]: return pcoll | beam.Map(lambda x: str(x)) _ = [1, 2, 3] | IntToInt() | IntToStr() def test_typed_ptransform_with_do_outputs_tuple_compiles(self): class MyDoFn(beam.DoFn): def process(self, element: int, *args, **kwargs): if element % 2: yield beam.pvalue.TaggedOutput('odd', 1) else: yield beam.pvalue.TaggedOutput('even', 1) class MyPTransform(beam.PTransform): def expand(self, pcoll: beam.pvalue.PCollection[int]): return pcoll | beam.ParDo(MyDoFn()).with_outputs('odd', 'even') # This test fails if you remove the following line from ptransform.py # if isinstance(pvalue_, DoOutputsTuple): continue _ = [1, 2, 3] | MyPTransform() def test_typed_ptransform_with_unknown_type_vars_tuple_compiles(self): @typehints.with_input_types(typing.TypeVar('T')) @typehints.with_output_types(typing.TypeVar('U')) def produces_unkown(e): return e @typehints.with_input_types(int) def accepts_int(e): return e class MyPTransform(beam.PTransform): def expand(self, pcoll): unknowns = pcoll | beam.Map(produces_unkown) ints = pcoll | beam.Map(int) return (unknowns, ints) | beam.Flatten() | beam.Map(accepts_int) _ = [1, 2, 3] | MyPTransform() class NativeTypesTest(unittest.TestCase): def test_good_main_input(self): @typehints.with_input_types(typing.Tuple[str, int]) def munge(s_i): (s, i) = s_i return (s + 's', i * 2) result = [('apple', 5), ('pear', 3)] | beam.Map(munge) self.assertEqual([('apples', 10), ('pears', 6)], sorted(result)) def test_bad_main_input(self): @typehints.with_input_types(typing.Tuple[str, str]) def munge(s_i): (s, i) = s_i return (s + 's', i * 2) with self.assertRaises(typehints.TypeCheckError): [('apple', 5), ('pear', 3)] | beam.Map(munge) def test_bad_main_output(self): @typehints.with_input_types(typing.Tuple[int, int]) @typehints.with_output_types(typing.Tuple[str, str]) def munge(a_b): (a, b) = a_b return (str(a), str(b)) with self.assertRaises(typehints.TypeCheckError): [(5, 4), (3, 2)] | beam.Map(munge) | 'Again' >> beam.Map(munge) class SideInputTest(unittest.TestCase): def _run_repeat_test(self, repeat): self._run_repeat_test_good(repeat) self._run_repeat_test_bad(repeat) @OptionsContext(pipeline_type_check=True) def _run_repeat_test_good(self, repeat): # As a positional argument. result = ['a', 'bb', 'c'] | beam.Map(repeat, 3) self.assertEqual(['aaa', 'bbbbbb', 'ccc'], sorted(result)) # As a keyword argument. result = ['a', 'bb', 'c'] | beam.Map(repeat, times=3) self.assertEqual(['aaa', 'bbbbbb', 'ccc'], sorted(result)) def _run_repeat_test_bad(self, repeat): # Various mismatches. with self.assertRaises(typehints.TypeCheckError): ['a', 'bb', 'c'] | beam.Map(repeat, 'z') with self.assertRaises(typehints.TypeCheckError): ['a', 'bb', 'c'] | beam.Map(repeat, times='z') with self.assertRaises(typehints.TypeCheckError): ['a', 'bb', 'c'] | beam.Map(repeat, 3, 4) if all(param.default == param.empty for param in get_signature(repeat).parameters.values()): with self.assertRaisesRegex(typehints.TypeCheckError, r'(takes exactly|missing a required)'): ['a', 'bb', 'c'] | beam.Map(repeat) def test_basic_side_input_hint(self): @typehints.with_input_types(str, int) def repeat(s, times): return s * times self._run_repeat_test(repeat) def test_keyword_side_input_hint(self): @typehints.with_input_types(str, times=int) def repeat(s, times): return s * times self._run_repeat_test(repeat) def test_default_typed_hint(self): @typehints.with_input_types(str, int) def repeat(s, times=3): return s * times self._run_repeat_test(repeat) def test_default_untyped_hint(self): @typehints.with_input_types(str) def repeat(s, times=3): return s * times # No type checking on default arg. self._run_repeat_test_good(repeat) @OptionsContext(pipeline_type_check=True) def test_varargs_side_input_hint(self): @typehints.with_input_types(str, int) def repeat(s, *times): return s * times[0] result = ['a', 'bb', 'c'] | beam.Map(repeat, 3) self.assertEqual(['aaa', 'bbbbbb', 'ccc'], sorted(result)) with self.assertRaisesRegex( typehints.TypeCheckError, (r'requires Tuple\[, ...\] but got ' r'Tuple\[, ...\]')): ['a', 'bb', 'c'] | beam.Map(repeat, 'z') def test_var_positional_only_side_input_hint(self): # Test that a lambda that accepts only a VAR_POSITIONAL can accept # side-inputs. # TODO(https://github.com/apache/beam/issues/19824): There's a bug with # trivial_inference inferring the output type when side-inputs are used # (their type hints are not passed). Remove with_output_types(...) when # this bug is fixed. result = (['a', 'b', 'c'] | beam.Map(lambda *args: args, 5).with_input_types( str, int).with_output_types(typehints.Tuple[str, int])) self.assertEqual([('a', 5), ('b', 5), ('c', 5)], sorted(result)) with self.assertRaisesRegex( typehints.TypeCheckError, r'requires.*Tuple\[Union\[, \], ...\].*' r'applied.*Tuple\[Union\[, \], ...\]'): _ = [1.2] | beam.Map(lambda *_: 'a', 5).with_input_types(int, str) def test_var_keyword_side_input_hint(self): # Test that a lambda that accepts a VAR_KEYWORD can accept # side-inputs. result = (['a', 'b', 'c'] | beam.Map(lambda e, **kwargs: (e, kwargs), kw=5).with_input_types(str, ignored=int)) self.assertEqual([('a', { 'kw': 5 }), ('b', { 'kw': 5 }), ('c', { 'kw': 5 })], sorted(result)) with self.assertRaisesRegex( typehints.TypeCheckError, r'requires Dict\[, \] but got ' r'Dict\[, \]'): _ = (['a', 'b', 'c'] | beam.Map(lambda e, **_: 'a', kw=5).with_input_types( str, ignored=str)) def test_deferred_side_inputs(self): @typehints.with_input_types(str, int) def repeat(s, times): return s * times with TestPipeline() as p: main_input = p | beam.Create(['a', 'bb', 'c']) side_input = p | 'side' >> beam.Create([3]) result = main_input | beam.Map(repeat, pvalue.AsSingleton(side_input)) assert_that(result, equal_to(['aaa', 'bbbbbb', 'ccc'])) bad_side_input = p | 'bad_side' >> beam.Create(['z']) with self.assertRaises(typehints.TypeCheckError): main_input | 'bis' >> beam.Map(repeat, pvalue.AsSingleton(bad_side_input)) def test_deferred_side_input_iterable(self): @typehints.with_input_types(str, typing.Iterable[str]) def concat(glue, items): return glue.join(sorted(items)) with TestPipeline() as p: main_input = p | beam.Create(['a', 'bb', 'c']) side_input = p | 'side' >> beam.Create(['x', 'y', 'z']) result = main_input | beam.Map(concat, pvalue.AsIter(side_input)) assert_that(result, equal_to(['xayaz', 'xbbybbz', 'xcycz'])) bad_side_input = p | 'bad_side' >> beam.Create([1, 2, 3]) with self.assertRaises(typehints.TypeCheckError): main_input | 'fail' >> beam.Map(concat, pvalue.AsIter(bad_side_input)) class CustomTransformTest(unittest.TestCase): class CustomTransform(beam.PTransform): def _extract_input_pvalues(self, pvalueish): return pvalueish, (pvalueish['in0'], pvalueish['in1']) def expand(self, pvalueish): return {'out0': pvalueish['in0'], 'out1': pvalueish['in1']} # TODO(robertwb): (typecheck) Make these the default? def with_input_types(self, *args, **kwargs): return WithTypeHints.with_input_types(self, *args, **kwargs) def with_output_types(self, *args, **kwargs): return WithTypeHints.with_output_types(self, *args, **kwargs) test_input = {'in0': ['a', 'b', 'c'], 'in1': [1, 2, 3]} def check_output(self, result): self.assertEqual(['a', 'b', 'c'], sorted(result['out0'])) self.assertEqual([1, 2, 3], sorted(result['out1'])) def test_custom_transform(self): self.check_output(self.test_input | self.CustomTransform()) def test_keyword_type_hints(self): self.check_output( self.test_input | self.CustomTransform().with_input_types(in0=str, in1=int)) self.check_output( self.test_input | self.CustomTransform().with_input_types(in0=str)) self.check_output( self.test_input | self.CustomTransform().with_output_types(out0=str, out1=int)) with self.assertRaises(typehints.TypeCheckError): self.test_input | self.CustomTransform().with_input_types(in0=int) with self.assertRaises(typehints.TypeCheckError): self.test_input | self.CustomTransform().with_output_types(out0=int) def test_flat_type_hint(self): # Type hint is applied to both. ({ 'in0': ['a', 'b', 'c'], 'in1': ['x', 'y', 'z'] } | self.CustomTransform().with_input_types(str)) with self.assertRaises(typehints.TypeCheckError): self.test_input | self.CustomTransform().with_input_types(str) with self.assertRaises(typehints.TypeCheckError): self.test_input | self.CustomTransform().with_input_types(int) with self.assertRaises(typehints.TypeCheckError): self.test_input | self.CustomTransform().with_output_types(int) class AnnotationsTest(unittest.TestCase): def test_pardo_wrapper_builtin_method(self): th = beam.ParDo(str.strip).get_type_hints() self.assertEqual(th.input_types, ((str, typehints.Any), {})) self.assertEqual(th.output_types, ((typehints.Any, ), {})) def test_pardo_wrapper_builtin_type(self): th = beam.ParDo(list).get_type_hints() self.assertEqual(th.input_types, ((typehints.Any, ), {})) self.assertEqual(th.output_types, ((typehints.Any, ), {})) def test_pardo_wrapper_builtin_func(self): th = beam.ParDo(len).get_type_hints() self.assertIsNone(th.input_types) self.assertIsNone(th.output_types) def test_pardo_dofn(self): class MyDoFn(beam.DoFn): def process(self, element: int) -> typehints.Generator[str]: yield str(element) th = beam.ParDo(MyDoFn()).get_type_hints() self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((str, ), {})) def test_pardo_dofn_not_iterable(self): class MyDoFn(beam.DoFn): def process(self, element: int) -> str: return str(element) with self.assertRaisesRegex(ValueError, r'str.*is not iterable'): _ = beam.ParDo(MyDoFn()).get_type_hints() def test_pardo_wrapper(self): def do_fn(element: int) -> typehints.Iterable[str]: return [str(element)] th = beam.ParDo(do_fn).get_type_hints() self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((str, ), {})) def test_pardo_wrapper_tuple(self): # Test case for callables that return key-value pairs for GBK. The outer # Iterable should be stripped but the inner Tuple left intact. def do_fn(element: int) -> typehints.Iterable[typehints.Tuple[str, int]]: return [(str(element), element)] th = beam.ParDo(do_fn).get_type_hints() self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((typehints.Tuple[str, int], ), {})) def test_pardo_wrapper_not_iterable(self): def do_fn(element: int) -> str: return str(element) with self.assertRaisesRegex(typehints.TypeCheckError, r'str.*is not iterable'): _ = beam.ParDo(do_fn).get_type_hints() def test_flat_map_wrapper(self): def map_fn(element: int) -> typehints.Iterable[int]: return [element, element + 1] th = beam.FlatMap(map_fn).get_type_hints() self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((int, ), {})) def test_flat_map_wrapper_optional_output(self): # Optional should not affect output type (Nones are ignored). def map_fn(element: int) -> typehints.Optional[typehints.Iterable[int]]: return [element, element + 1] th = beam.FlatMap(map_fn).get_type_hints() self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((int, ), {})) @unittest.skip( 'https://github.com/apache/beam/issues/19961: Py3 annotations not yet ' 'supported for MapTuple') def test_flat_map_tuple_wrapper(self): # TODO(https://github.com/apache/beam/issues/19961): Also test with a fn # that accepts default arguments. def tuple_map_fn(a: str, b: str, c: str) -> typehints.Iterable[str]: return [a, b, c] th = beam.FlatMapTuple(tuple_map_fn).get_type_hints() self.assertEqual(th.input_types, ((str, str, str), {})) self.assertEqual(th.output_types, ((str, ), {})) def test_map_wrapper(self): def map_fn(unused_element: int) -> int: return 1 th = beam.Map(map_fn).get_type_hints() self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((int, ), {})) def test_map_wrapper_optional_output(self): # Optional does affect output type (Nones are NOT ignored). def map_fn(unused_element: int) -> typehints.Optional[int]: return 1 th = beam.Map(map_fn).get_type_hints() self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((typehints.Optional[int], ), {})) @unittest.skip( 'https://github.com/apache/beam/issues/19961: Py3 annotations not yet ' 'supported for MapTuple') def test_map_tuple(self): # TODO(https://github.com/apache/beam/issues/19961): Also test with a fn # that accepts default arguments. def tuple_map_fn(a: str, b: str, c: str) -> str: return a + b + c th = beam.MapTuple(tuple_map_fn).get_type_hints() self.assertEqual(th.input_types, ((str, str, str), {})) self.assertEqual(th.output_types, ((str, ), {})) def test_filter_wrapper(self): def filter_fn(element: int) -> bool: return bool(element % 2) th = beam.Filter(filter_fn).get_type_hints() self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((int, ), {})) class TestFlatMapTuple(unittest.TestCase): def test_flatmaptuple(self): # Regression test. See # https://github.com/apache/beam/issues/33014 def identity(x: Tuple[str, int]) -> Tuple[str, int]: return x with beam.Pipeline() as p: # Just checking that this doesn't raise an exception. ( p | "Generate input" >> beam.Create([('P1', [2])]) | "Flat" >> beam.FlatMapTuple(lambda k, vs: [(k, v) for v in vs]) | "Identity" >> beam.Map(identity)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/typehints.py0000644000175100017510000015703215206113033022316 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Syntax & semantics for type-hinting custom-functions/PTransforms in the SDK. This module defines type-hinting objects and the corresponding syntax for type-hinting function arguments, function return types, or PTransform object themselves. TypeHint's defined in the module can be used to implement either static or run-time type-checking in regular Python code. Type-hints are defined by 'indexing' a type-parameter into a defined CompositeTypeHint instance: * 'List[int]'. Valid type-hints are partitioned into two categories: simple, and composite. Simple type hints are type hints based on a subset of Python primitive types: int, bool, float, str, object, None, and bytes. No other primitive types are allowed. Composite type-hints are reserved for hinting the types of container-like Python objects such as 'list'. Composite type-hints can be parameterized by an inner simple or composite type-hint, using the 'indexing' syntax. In order to avoid conflicting with the namespace of the built-in container types, when specifying this category of type-hints, the first letter should be capitalized. The following composite type-hints are permitted. NOTE: 'T' can be any of the type-hints listed or a simple Python type: * Any * Union[T, T, T] * Optional[T] * Tuple[T, T] * Tuple[T, ...] * List[T] * KV[T, T] * Dict[T, T] * Set[T] * Iterable[T] * Iterator[T] * Generator[T] Type-hints can be nested, allowing one to define type-hints for complex types: * 'List[Tuple[int, int, str]] In addition, type-hints can be used to implement run-time type-checking via the 'type_check' method on each TypeConstraint. """ # pytype: skip-file # ruff: noqa: UP006 import copy import logging import types import typing from collections import abc from beartype.door import is_subhint from beartype.roar import BeartypeDoorException __all__ = [ 'Any', 'Union', 'Optional', 'Tuple', 'List', 'KV', 'Dict', 'Set', 'FrozenSet', 'Collection', 'Sequence', 'Iterable', 'Iterator', 'Generator', 'WindowedValue', 'TypeVariable', ] # A set of the built-in Python types we don't support, guiding the users # to templated (upper-case) versions instead. DISALLOWED_PRIMITIVE_TYPES = (list, set, frozenset, tuple, dict) _LOGGER = logging.getLogger(__name__) class SimpleTypeHintError(TypeError): pass class CompositeTypeHintError(TypeError): pass class GetitemConstructor(type): """A metaclass that makes Cls[arg] an alias for Cls(arg).""" def __getitem__(cls, arg): return cls(arg) class TypeConstraint(object): """The base-class for all created type-constraints defined below. A :class:`TypeConstraint` is the result of parameterizing a :class:`CompositeTypeHint` with with one of the allowed Python types or another :class:`CompositeTypeHint`. It binds and enforces a specific version of a generalized TypeHint. """ def _consistent_with_check_(self, sub): """Returns whether sub is consistent with self. Has the same relationship to is_consistent_with() as __subclasscheck__ does for issubclass(). Not meant to be called directly; call is_consistent_with(sub, self) instead. Implementation may assume that maybe_sub_type is not Any and has been normalized. """ raise NotImplementedError def type_check(self, instance): """Determines if the type of 'instance' satisfies this type constraint. Args: instance: An instance of a Python object. Raises: :class:`TypeError`: The passed **instance** doesn't satisfy this :class:`TypeConstraint`. Subclasses of :class:`TypeConstraint` are free to raise any of the subclasses of :class:`TypeError` defined above, depending on the manner of the type hint error. All :class:`TypeConstraint` sub-classes must define this method in other for the class object to be created. """ raise NotImplementedError def match_type_variables(self, unused_concrete_type): return {} def bind_type_variables(self, unused_bindings): return self def _inner_types(self): """Iterates over the inner types of the composite type.""" return [] def visit(self, visitor, visitor_arg): """Visitor method to visit all inner types of a composite type. Args: visitor: A callable invoked for all nodes in the type tree comprising a composite type. The visitor will be called with the node visited and the visitor argument specified here. visitor_arg: Visitor callback second argument. """ visitor(self, visitor_arg) for t in self._inner_types(): if isinstance(t, TypeConstraint): t.visit(visitor, visitor_arg) else: visitor(t, visitor_arg) def visit_inner_types(type_constraint, visitor, visitor_arg): """Visitor pattern to visit all inner types of a type constraint. Args: type_constraint: A type constraint or a type. visitor: A callable invoked for all nodes in the type tree comprising a composite type. The visitor will be called with the node visited and the visitor argument specified here. visitor_arg: Visitor callback second argument. Note: Raise and capture a StopIteration to terminate the visit, e.g. ``` def visitor(type_constraint, visitor_arg): if ...: raise StopIteration try: visit_inner_types(type_constraint, visitor, visitor_arg) except StopIteration: pass ``` """ if isinstance(type_constraint, TypeConstraint): return type_constraint.visit(visitor, visitor_arg) return visitor(type_constraint, visitor_arg) def match_type_variables(type_constraint, concrete_type): if isinstance(type_constraint, TypeConstraint): return type_constraint.match_type_variables(concrete_type) return {} def bind_type_variables(type_constraint, bindings): if isinstance(type_constraint, TypeConstraint): return type_constraint.bind_type_variables(bindings) return type_constraint class IndexableTypeConstraint(TypeConstraint): """An internal common base-class for all type constraints with indexing. E.G. SequenceTypeConstraint + Tuple's of fixed size. """ def _constraint_for_index(self, idx): """Returns the type at the given index. This is used to allow type inference to determine the correct type for a specific index. On lists this will also be the same, however for tuples the value will depend on the position. This was added as part of the futurize changes since more of the expressions now index into tuples.""" raise NotImplementedError class SequenceTypeConstraint(IndexableTypeConstraint): """A common base-class for all sequence related type-constraint classes. A sequence is defined as an arbitrary length homogeneous container type. Type hints which fall under this category include: List[T], Set[T], Iterable[T], and Tuple[T, ...]. Sub-classes may need to override '_consistent_with_check_' if a particular sequence requires special handling with respect to type compatibility. Attributes: inner_type: The type which every element in the sequence should be an instance of. """ def __init__(self, inner_type, sequence_type): self.inner_type = normalize(inner_type) self._sequence_type = sequence_type def __eq__(self, other): return ( isinstance(other, SequenceTypeConstraint) and type(self) == type(other) and self.inner_type == other.inner_type) def __hash__(self): return hash(self.inner_type) ^ 13 * hash(type(self)) def _inner_types(self): yield self.inner_type def _constraint_for_index(self, idx): """Returns the type at the given index.""" return self.inner_type def _consistent_with_check_(self, sub): return ( isinstance(sub, self.__class__) and is_consistent_with(sub.inner_type, self.inner_type)) def type_check(self, sequence_instance): if not isinstance(sequence_instance, self._sequence_type): raise CompositeTypeHintError( "%s type-constraint violated. Valid object instance " "must be of type '%s'. Instead, an instance of '%s' " "was received." % ( self._sequence_type.__name__.title(), self._sequence_type.__name__.lower(), sequence_instance.__class__.__name__)) for index, elem in enumerate(sequence_instance): try: check_constraint(self.inner_type, elem) except SimpleTypeHintError: raise CompositeTypeHintError( '%s hint type-constraint violated. The type of element #%s in ' 'the passed %s is incorrect. Expected an instance of type %s, ' 'instead received an instance of type %s.' % ( repr(self), index, repr(self._sequence_type), repr(self.inner_type), elem.__class__.__name__)) except CompositeTypeHintError as e: raise CompositeTypeHintError( '%s hint type-constraint violated. The type of element #%s in ' 'the passed %s is incorrect: %s' % (repr(self), index, self._sequence_type.__name__, e)) def match_type_variables(self, concrete_type): if isinstance(concrete_type, SequenceTypeConstraint): return match_type_variables(self.inner_type, concrete_type.inner_type) return {} def bind_type_variables(self, bindings): bound_inner_type = bind_type_variables(self.inner_type, bindings) if bound_inner_type == self.inner_type: return self bound_self = copy.copy(self) bound_self.inner_type = bound_inner_type return bound_self class CompositeTypeHint(object): """The base-class for all created type-hint classes defined below. CompositeTypeHint's serve primarily as TypeConstraint factories. They are only required to define a single method: '__getitem__' which should return a parameterized TypeConstraint, that can be used to enforce static or run-time type-checking. '__getitem__' is used as a factory function in order to provide a familiar API for defining type-hints. The ultimate result is that one will be able to use: CompositeTypeHint[type_parameter] to create a type-hint object that behaves like any other Python object. This allows one to create 'type-aliases' by assigning the returned type-hints to a variable. * Example: 'Coordinates = List[Tuple[int, int]]' """ def __getitem___(self, py_type): """Given a type creates a TypeConstraint instance parameterized by the type. This function serves as a factory function which creates TypeConstraint instances. Additionally, implementations by sub-classes should perform any sanity checking of the passed types in this method in order to rule-out disallowed behavior. Such as, attempting to create a TypeConstraint whose parameterized type is actually an object instance. Args: py_type: An instance of a Python type or TypeConstraint. Returns: An instance of a custom TypeConstraint for this CompositeTypeHint. Raises: TypeError: If the passed type violates any contraints for this particular TypeHint. """ raise NotImplementedError def is_typing_generic(type_param): """Determines if an object is a subscripted typing.Generic type, such as PCollection[int]. Such objects are considered valid type parameters. For Python versions 3.9 and above, also permits types.GenericAlias. """ if hasattr(types, "GenericAlias") and isinstance(type_param, types.GenericAlias): return True return isinstance(type_param, typing._GenericAlias) def validate_composite_type_param(type_param, error_msg_prefix): """Determines if an object is a valid type parameter to a :class:`CompositeTypeHint`. Implements sanity checking to disallow things like:: List[1, 2, 3] or Dict[5]. Args: type_param: An object instance. error_msg_prefix (:class:`str`): A string prefix used to format an error message in the case of an exception. Raises: TypeError: If the passed **type_param** is not a valid type parameter for a :class:`CompositeTypeHint`. """ # Must either be a TypeConstraint instance or a basic Python type. possible_classes = [type, TypeConstraint] is_not_type_constraint = ( not is_typing_generic(type_param) and not isinstance(type_param, tuple(possible_classes)) and type_param is not None and getattr(type_param, '__module__', None) != 'typing') if isinstance(type_param, types.UnionType): is_not_type_constraint = False if is_not_type_constraint: raise TypeError( '%s must be a non-sequence, a type, or a TypeConstraint. %s' ' is an instance of %s.' % (error_msg_prefix, type_param, type_param.__class__.__name__)) def check_constraint(type_constraint, object_instance): """Determine if the passed type instance satisfies the TypeConstraint. When examining a candidate type for constraint satisfaction in 'type_check', all CompositeTypeHint's eventually call this function. This function may end up being called recursively if the hinted type of a CompositeTypeHint is another CompositeTypeHint. Args: type_constraint: An instance of a TypeConstraint or a built-in Python type. object_instance: An object instance. Raises: SimpleTypeHintError: If 'type_constraint' is a one of the allowed primitive Python types and 'object_instance' isn't an instance of this type. CompositeTypeHintError: If 'type_constraint' is a TypeConstraint object and 'object_instance' does not satisfy its constraint. """ if type_constraint is None and object_instance is None: return elif isinstance(type_constraint, TypeConstraint): type_constraint.type_check(object_instance) elif type_constraint is None: # TODO(robertwb): Fix uses of None for Any. pass elif not isinstance(type_constraint, type): raise RuntimeError("bad type: %s" % (type_constraint, )) elif not isinstance(object_instance, type_constraint): raise SimpleTypeHintError class AnyTypeConstraint(TypeConstraint): """An Any type-hint. Any is intended to be used as a "don't care" when hinting the types of function arguments or return types. All other TypeConstraint's are equivalent to 'Any', and its 'type_check' method is a no-op. """ def __eq__(self, other): return type(self) == type(other) def __repr__(self): return 'Any' def __hash__(self): # TODO(https://github.com/apache/beam/issues/18633): Fix # typehints.TypeVariable issues with __hash__. return hash(id(self)) def type_check(self, instance): pass class TypeVariable(AnyTypeConstraint): def __init__(self, name, use_name_in_eq=True): self.name = name self.use_name_in_eq = use_name_in_eq def __eq__(self, other): # The "other" may be an Ellipsis object # so we have to check if it has use_name_in_eq first if self.use_name_in_eq and (hasattr(other, 'use_name_in_eq') and other.use_name_in_eq): return type(self) == type(other) and self.name == other.name return type(self) == type(other) def __hash__(self): # TODO(https://github.com/apache/beam/issues/18633): Fix # typehints.TypeVariable issues with __hash__. return hash(id(self)) def __repr__(self): return 'TypeVariable[%s]' % self.name def match_type_variables(self, concrete_type): return {self: concrete_type} def bind_type_variables(self, bindings): return bindings.get( self, # Star matches all type variables. bindings.get('*', self)) class UnionHint(CompositeTypeHint): """A Union type-hint. Union[X, Y] accepts instances of type X OR type Y. Duplicate type parameters are ignored. Additonally, Nested Union hints will be flattened out. For example: * Union[Union[str, int], bool] -> Union[str, int, bool] A candidate type instance satisfies a UnionConstraint if it is an instance of any of the parameterized 'union_types' for a Union. Union[X] is disallowed, and all type parameters will be sanity checked to ensure compatibility with nested type-hints. When comparing two Union hints, ordering is enforced before comparison. * Union[int, str] == Union[str, int] """ class UnionConstraint(TypeConstraint): def __init__(self, union_types): # Use a list instead of a set to make the pickle serialization # deterministic. self.union_types = list(set(normalize(t) for t in union_types)) # Sorting the type name strings simplifies unit tests. self.union_types.sort(key=repr) def __eq__(self, other): return ( isinstance(other, UnionHint.UnionConstraint) and # The union types represent a set even thought they are stored as a # list. set(self.union_types) == set(other.union_types)) def __hash__(self): return 1 + sum(hash(t) for t in self.union_types) def __repr__(self): return 'Union[%s]' % (', '.join(repr(t) for t in self.union_types)) def inner_types(self): for t in self.union_types: yield t def contains_type(self, maybe_type): return maybe_type in self.union_types def _consistent_with_check_(self, sub): if isinstance(sub, UnionConstraint): # A union type is compatible if every possible type is compatible. # E.g. Union[A, B, C] > Union[A, B]. return all(is_consistent_with(elem, self) for elem in sub.union_types) # Other must be compatible with at least one of this union's subtypes. # E.g. Union[A, B, C] > T if T > A or T > B or T > C. return any(is_consistent_with(sub, elem) for elem in self.union_types) def type_check(self, instance): error_msg = '' for t in self.union_types: try: check_constraint(t, instance) return except TypeError as e: error_msg = str(e) continue raise CompositeTypeHintError( '%s type-constraint violated. Expected an instance of one of: %s, ' 'received %s instead.%s' % ( repr(self), tuple(repr(t) for t in self.union_types), instance.__class__.__name__, error_msg)) def match_type_variables(self, concrete_type): sub_bindings = [ match_type_variables(t, concrete_type) for t in self.union_types if is_consistent_with(concrete_type, t) ] if sub_bindings: return { var: Union[(sub[var] for sub in sub_bindings)] for var in set.intersection( *[set(sub.keys()) for sub in sub_bindings]) } else: return {} def bind_type_variables(self, bindings): return Union[(bind_type_variables(t, bindings) for t in self.union_types)] def __getitem__(self, type_params): if not isinstance(type_params, (abc.Iterable, set)): raise TypeError('Cannot create Union without a sequence of types.') # Flatten nested Union's and duplicated repeated type hints. params = set() dict_union = None for t in type_params: validate_composite_type_param( t, error_msg_prefix='All parameters to a Union hint') if isinstance(t, self.UnionConstraint): params |= set(t.union_types) elif isinstance(t, DictConstraint): if dict_union is None: dict_union = t else: dict_union.key_type = Union[dict_union.key_type, t.key_type] dict_union.value_type = Union[dict_union.value_type, t.value_type] else: params.add(t) if dict_union is not None: params.add(dict_union) if Any in params: return Any elif len(params) == 1: return next(iter(params)) if len(params) > 1: from apache_beam.typehints import schemas try: return schemas.union_schema_type(params) except (TypeError, KeyError): # Not a union of compatible schema types. pass return self.UnionConstraint(params) UnionConstraint = UnionHint.UnionConstraint class OptionalHint(UnionHint): """An Option type-hint. Optional[X] accepts instances of X or None. The Optional[X] factory function proxies to Union[X, type(None)] """ def __getitem__(self, py_type): # A single type must have been passed. if isinstance(py_type, abc.Sequence): raise TypeError( 'An Option type-hint only accepts a single type ' 'parameter.') return Union[py_type, type(None)] def is_nullable(typehint): return ( isinstance(typehint, UnionConstraint) and typehint.contains_type(type(None)) and len(list(typehint.inner_types())) == 2) def get_concrete_type_from_nullable(typehint): if is_nullable(typehint): for inner_type in typehint.inner_types(): if not type(None) == inner_type: return inner_type else: raise TypeError('Typehint is not of nullable type', typehint) class TupleHint(CompositeTypeHint): """A Tuple type-hint. Tuple can accept 1 or more type-hint parameters. Tuple[X, Y] represents a tuple of *exactly* two elements, with the first being of type 'X' and the second an instance of type 'Y'. * (1, 2) satisfies Tuple[int, int] Additionally, one is able to type-hint an arbitary length, homogeneous tuple by passing the Ellipsis (...) object as the second parameter. As an example, Tuple[str, ...] indicates a tuple of any length with each element being an instance of 'str'. """ class TupleSequenceConstraint(SequenceTypeConstraint): def __init__(self, type_param): super().__init__(type_param, tuple) def __repr__(self): return 'Tuple[%s, ...]' % repr(self.inner_type) def _consistent_with_check_(self, sub): if isinstance(sub, TupleConstraint): # E.g. Tuple[A, B] < Tuple[C, ...] iff A < C and B < C. return all( is_consistent_with(elem, self.inner_type) for elem in sub.tuple_types) return super()._consistent_with_check_(sub) class TupleConstraint(IndexableTypeConstraint): def __init__(self, type_params): self.tuple_types = tuple(normalize(t) for t in type_params) def __eq__(self, other): return ( isinstance(other, TupleHint.TupleConstraint) and self.tuple_types == other.tuple_types) def __hash__(self): return hash(self.tuple_types) def __repr__(self): return 'Tuple[%s]' % (', '.join(repr(t) for t in self.tuple_types)) def _inner_types(self): for t in self.tuple_types: yield t def _constraint_for_index(self, idx): """Returns the type at the given index.""" return self.tuple_types[idx] def _consistent_with_check_(self, sub): return ( isinstance(sub, self.__class__) and len(sub.tuple_types) == len(self.tuple_types) and all( is_consistent_with(sub_elem, elem) for sub_elem, elem in zip(sub.tuple_types, self.tuple_types))) def type_check(self, tuple_instance): if not isinstance(tuple_instance, tuple): raise CompositeTypeHintError( "Tuple type constraint violated. Valid object instance must be of " "type 'tuple'. Instead, an instance of '%s' was received." % tuple_instance.__class__.__name__) if len(tuple_instance) != len(self.tuple_types): raise CompositeTypeHintError( 'Passed object instance is of the proper type, but differs in ' 'length from the hinted type. Expected a tuple of length %s, ' 'received a tuple of length %s.' % (len(self.tuple_types), len(tuple_instance))) for type_pos, (expected, actual) in enumerate(zip(self.tuple_types, tuple_instance)): try: check_constraint(expected, actual) continue except SimpleTypeHintError: raise CompositeTypeHintError( '%s hint type-constraint violated. The type of element #%s in ' 'the passed tuple is incorrect. Expected an instance of ' 'type %s, instead received an instance of type %s.' % (repr(self), type_pos, repr(expected), actual.__class__.__name__)) except CompositeTypeHintError as e: raise CompositeTypeHintError( '%s hint type-constraint violated. The type of element #%s in ' 'the passed tuple is incorrect. %s' % (repr(self), type_pos, e)) def match_type_variables(self, concrete_type): bindings = {} if isinstance(concrete_type, TupleConstraint): for a, b in zip(self.tuple_types, concrete_type.tuple_types): bindings.update(match_type_variables(a, b)) return bindings def bind_type_variables(self, bindings): bound_tuple_types = tuple( bind_type_variables(t, bindings) for t in self.tuple_types) if bound_tuple_types == self.tuple_types: return self return Tuple[bound_tuple_types] def __getitem__(self, type_params): ellipsis = False if not isinstance(type_params, abc.Iterable): # Special case for hinting tuples with arity-1. type_params = (type_params, ) if type_params and type_params[-1] == Ellipsis: if len(type_params) != 2: raise TypeError( 'Ellipsis can only be used to type-hint an arbitrary ' 'length tuple of containing a single type: ' 'Tuple[A, ...].') # Tuple[A, ...] indicates an arbitary length homogeneous tuple. type_params = type_params[:1] ellipsis = True for t in type_params: validate_composite_type_param( t, error_msg_prefix='All parameters to a Tuple hint') if ellipsis: return self.TupleSequenceConstraint(type_params[0]) return self.TupleConstraint(type_params) TupleConstraint = TupleHint.TupleConstraint TupleSequenceConstraint = TupleHint.TupleSequenceConstraint class ListHint(CompositeTypeHint): """A List type-hint. List[X] represents an instance of a list populated by a single homogeneous type. The parameterized type 'X' can either be a built-in Python type or an instance of another TypeConstraint. * ['1', '2', '3'] satisfies List[str] """ class ListConstraint(SequenceTypeConstraint): def __init__(self, list_type): super().__init__(list_type, list) def __repr__(self): return 'List[%s]' % repr(self.inner_type) def __getitem__(self, t): validate_composite_type_param(t, error_msg_prefix='Parameter to List hint') return self.ListConstraint(t) ListConstraint = ListHint.ListConstraint class KVHint(CompositeTypeHint): """A KV type-hint, represents a Key-Value pair of a particular type. Internally, KV[X, Y] proxies to Tuple[X, Y]. A KV type-hint accepts only accepts exactly two type-parameters. The first represents the required key-type and the second the required value-type. """ def __getitem__(self, type_params): if not isinstance(type_params, tuple): raise TypeError( 'Parameter to KV type-hint must be a tuple of types: ' 'KV[.., ..].') if len(type_params) != 2: raise TypeError( 'Length of parameters to a KV type-hint must be exactly 2. Passed ' 'parameters: %s, have a length of %s.' % (type_params, len(type_params))) return Tuple[type_params] def key_value_types(kv): """Returns the key and value type of a KV type-hint. Args: kv: An instance of a TypeConstraint sub-class. Returns: A tuple: (key_type, value_type) if the passed type-hint is an instance of a KV type-hint, and (Any, Any) otherwise. """ if isinstance(kv, TupleHint.TupleConstraint): return kv.tuple_types return Any, Any class DictHint(CompositeTypeHint): """A Dict type-hint. Dict[K, V] Represents a dictionary where all keys are of a particular type and all values are of another (possible the same) type. """ class DictConstraint(TypeConstraint): def __init__(self, key_type, value_type): self.key_type = normalize(key_type) self.value_type = normalize(value_type) def __repr__(self): return 'Dict[%s, %s]' % (repr(self.key_type), repr(self.value_type)) def __eq__(self, other): return ( type(self) == type(other) and self.key_type == other.key_type and self.value_type == other.value_type) def __hash__(self): return hash((type(self), self.key_type, self.value_type)) def _inner_types(self): yield self.key_type yield self.value_type def _consistent_with_check_(self, sub): return ( isinstance(sub, self.__class__) and is_consistent_with(sub.key_type, self.key_type) and is_consistent_with(sub.value_type, self.value_type)) def _raise_hint_exception_or_inner_exception( self, is_key, incorrect_instance, inner_error_message=''): incorrect_type = 'values' if not is_key else 'keys' hinted_type = self.value_type if not is_key else self.key_type if inner_error_message: raise CompositeTypeHintError( '%s hint %s-type constraint violated. All %s should be of type ' '%s. Instead: %s' % ( repr(self), incorrect_type[:-1], incorrect_type, repr(hinted_type), inner_error_message)) else: raise CompositeTypeHintError( '%s hint %s-type constraint violated. All %s should be of ' 'type %s. Instead, %s is of type %s.' % ( repr(self), incorrect_type[:-1], incorrect_type, repr(hinted_type), incorrect_instance, incorrect_instance.__class__.__name__)) def type_check(self, dict_instance): if not isinstance(dict_instance, dict): raise CompositeTypeHintError( 'Dict type-constraint violated. All passed instances must be of ' 'type dict. %s is of type %s.' % (dict_instance, dict_instance.__class__.__name__)) for key, value in dict_instance.items(): try: check_constraint(self.key_type, key) except CompositeTypeHintError as e: self._raise_hint_exception_or_inner_exception(True, key, str(e)) except SimpleTypeHintError: self._raise_hint_exception_or_inner_exception(True, key) try: check_constraint(self.value_type, value) except CompositeTypeHintError as e: self._raise_hint_exception_or_inner_exception(False, value, str(e)) except SimpleTypeHintError: self._raise_hint_exception_or_inner_exception(False, value) def match_type_variables(self, concrete_type): if isinstance(concrete_type, DictConstraint): bindings = {} bindings.update( match_type_variables(self.key_type, concrete_type.key_type)) bindings.update( match_type_variables(self.value_type, concrete_type.value_type)) return bindings return {} def bind_type_variables(self, bindings): bound_key_type = bind_type_variables(self.key_type, bindings) bound_value_type = bind_type_variables(self.value_type, bindings) if (bound_key_type, self.key_type) == (bound_value_type, self.value_type): return self return Dict[bound_key_type, bound_value_type] def __getitem__(self, type_params): # Type param must be a (k, v) pair. if not isinstance(type_params, tuple): raise TypeError( 'Parameter to Dict type-hint must be a tuple of types: ' 'Dict[.., ..].') if len(type_params) != 2: raise TypeError( 'Length of parameters to a Dict type-hint must be exactly 2. Passed ' 'parameters: %s, have a length of %s.' % (type_params, len(type_params))) key_type, value_type = type_params validate_composite_type_param( key_type, error_msg_prefix='Key-type parameter to a Dict hint') validate_composite_type_param( value_type, error_msg_prefix='Value-type parameter to a Dict hint') return self.DictConstraint(key_type, value_type) DictConstraint = DictHint.DictConstraint class SetHint(CompositeTypeHint): """A Set type-hint. Set[X] defines a type-hint for a set of homogeneous types. 'X' may be either a built-in Python type or a another nested TypeConstraint. """ class SetTypeConstraint(SequenceTypeConstraint): def __init__(self, type_param): super().__init__(type_param, set) def __repr__(self): return 'Set[%s]' % repr(self.inner_type) def __getitem__(self, type_param): validate_composite_type_param( type_param, error_msg_prefix='Parameter to a Set hint') return self.SetTypeConstraint(type_param) SetTypeConstraint = SetHint.SetTypeConstraint class FrozenSetHint(CompositeTypeHint): """A FrozenSet type-hint. FrozenSet[X] defines a type-hint for a set of homogeneous types. 'X' may be either a built-in Python type or a another nested TypeConstraint. This is a mirror copy of SetHint - consider refactoring common functionality. """ class FrozenSetTypeConstraint(SequenceTypeConstraint): def __init__(self, type_param): super(FrozenSetHint.FrozenSetTypeConstraint, self).__init__(type_param, frozenset) def __repr__(self): return 'FrozenSet[%s]' % repr(self.inner_type) def __getitem__(self, type_param): validate_composite_type_param( type_param, error_msg_prefix='Parameter to a FrozenSet hint') return self.FrozenSetTypeConstraint(type_param) FrozenSetTypeConstraint = FrozenSetHint.FrozenSetTypeConstraint class CollectionHint(CompositeTypeHint): """ A Collection type-hint. Collection[X] defines a type-hint for a collection of homogenous types. 'X' may be either a built-in Python type or another nested TypeConstraint. This represents a collections.abc.Collection type, which implements __contains__, __iter__, and __len__. This acts as a parent type for sets but has fewer guarantees for mixins. """ class CollectionTypeConstraint(SequenceTypeConstraint): def __init__(self, type_param): super().__init__(type_param, abc.Collection) def __repr__(self): return 'Collection[%s]' % repr(self.inner_type) @staticmethod def _is_subclass_constraint(sub): return isinstance( sub, ( CollectionTypeConstraint, ABCSequenceTypeConstraint, FrozenSetTypeConstraint, SetTypeConstraint, ListConstraint)) def _consistent_with_check_(self, sub): if self._is_subclass_constraint(sub): return is_consistent_with(sub.inner_type, self.inner_type) elif isinstance(sub, TupleConstraint): if not sub.tuple_types: # The empty tuple is consistent with Iterator[T] for any T. return True # Each element in the hetrogenious tuple must be consistent with # the collection type. # E.g. Tuple[A, B] < Collection[C] if A < C and B < C. return all( is_consistent_with(elem, self.inner_type) for elem in sub.tuple_types) # TODO(https://github.com/apache/beam/issues/29135): allow for # consistency checks with Mapping types elif isinstance(sub, (MappingTypeConstraint, DictConstraint)): return True elif not isinstance(sub, TypeConstraint): if getattr(sub, '__origin__', None) is not None and getattr( sub, '__args__', None) is not None: return issubclass(sub, abc.Collection) and is_consistent_with( sub.__args__, self.inner_type) return False def __getitem__(self, type_param): validate_composite_type_param( type_param, error_msg_prefix='Parameter to a Collection hint') return self.CollectionTypeConstraint(type_param) CollectionTypeConstraint = CollectionHint.CollectionTypeConstraint class SequenceHint(CompositeTypeHint): """A Sequence type-hint. Sequence[X] defines a type-hint for a sequence of homogeneous types. 'X' may be either a built-in Python type or another nested TypeConstraint. This represents collections.abc.Sequence type, which implements __getitem__, __len__, and __contains__. This is more specific than Iterable but less restrictive than List, providing a good middle ground for sequence-like types. """ class ABCSequenceTypeConstraint(SequenceTypeConstraint): def __init__(self, type_param): super().__init__(type_param, abc.Sequence) def __repr__(self): return 'Sequence[%s]' % repr(self.inner_type) def _consistent_with_check_(self, sub): if isinstance(sub, (ListConstraint, TupleConstraint)): # Lists and Tuples are Sequences if isinstance(sub, TupleConstraint): # For tuples, all elements must be consistent with the sequence type return all( is_consistent_with(elem, self.inner_type) for elem in sub.tuple_types) return is_consistent_with(sub.inner_type, self.inner_type) return super()._consistent_with_check_(sub) def __getitem__(self, type_param): validate_composite_type_param( type_param, error_msg_prefix='Parameter to a Sequence hint') return self.ABCSequenceTypeConstraint(type_param) ABCSequenceTypeConstraint = SequenceHint.ABCSequenceTypeConstraint class MappingHint(CompositeTypeHint): """A Mapping type-hint. Mapping[K, V] represents any mapping (dict-like object) where all keys are of type K and all values are of type V. This is more general than Dict as it supports any object implementing the Mapping ABC. Examples of valid mappings include: - dict - collections.defaultdict - collections.OrderedDict - types.MappingProxyType """ class MappingTypeConstraint(TypeConstraint): def __init__(self, key_type, value_type): self.key_type = normalize(key_type) self.value_type = normalize(value_type) def __repr__(self): return 'Mapping[%s, %s]' % (repr(self.key_type), repr(self.value_type)) def __eq__(self, other): return ( type(self) == type(other) and self.key_type == other.key_type and self.value_type == other.value_type) def __hash__(self): return hash((type(self), self.key_type, self.value_type)) def _inner_types(self): yield self.key_type yield self.value_type def _consistent_with_check_(self, sub): if isinstance(sub, (self.__class__, DictConstraint)): # A Dict is a Mapping of the same types return ( is_consistent_with(sub.key_type, self.key_type) and is_consistent_with(sub.value_type, self.value_type)) elif hasattr(sub, '__origin__'): # Handle collection subtypes using ABC if issubclass(sub.__origin__, abc.Mapping): args = getattr(sub, '__args__', None) if args and len(args) == 2: return ( is_consistent_with(args[0], self.key_type) and is_consistent_with(args[1], self.value_type)) return True return False def _raise_type_error(self, is_key, instance, inner_error_message=''): type_desc = 'key' if is_key else 'value' expected_type = self.key_type if is_key else self.value_type if inner_error_message: raise CompositeTypeHintError( '%s hint %s-type constraint violated. All %ss should be of type ' '%s. Instead: %s' % ( repr(self), type_desc, type_desc, repr(expected_type), inner_error_message, )) else: raise CompositeTypeHintError( '%s hint %s-type constraint violated. All %ss should be of ' 'type %s. Instead, %s is of type %s.' % ( repr(self), type_desc, type_desc, repr(expected_type), instance, instance.__class__.__name__, )) def type_check(self, instance): if not isinstance(instance, abc.Mapping): raise CompositeTypeHintError( 'Mapping type-constraint violated. All passed instances must be of ' 'type Mapping. %s is of type %s.' % (instance, instance.__class__.__name__)) for key, value in instance.items(): try: check_constraint(self.key_type, key) except CompositeTypeHintError as e: self._raise_type_error(True, key, str(e)) except SimpleTypeHintError: self._raise_type_error(True, key) try: check_constraint(self.value_type, value) except CompositeTypeHintError as e: self._raise_type_error(False, value, str(e)) except SimpleTypeHintError: self._raise_type_error(False, value) def match_type_variables(self, concrete_type): if isinstance(concrete_type, (MappingTypeConstraint, DictConstraint)): bindings = {} bindings.update( match_type_variables(self.key_type, concrete_type.key_type)) bindings.update( match_type_variables(self.value_type, concrete_type.value_type)) return bindings return {} def bind_type_variables(self, bindings): bound_key_type = bind_type_variables(self.key_type, bindings) bound_value_type = bind_type_variables(self.value_type, bindings) if (bound_key_type, bound_value_type) == (self.key_type, self.value_type): return self return Mapping[bound_key_type, bound_value_type] def __getitem__(self, type_params): if not isinstance(type_params, tuple): raise TypeError( 'Parameter to Mapping type-hint must be a tuple of types: ' 'Mapping[.., ..].') if len(type_params) != 2: raise TypeError( 'Length of parameters to a Mapping type-hint must be exactly 2. ' 'Passed parameters: %s, have a length of %s.' % (type_params, len(type_params))) key_type, value_type = type_params validate_composite_type_param( key_type, error_msg_prefix='Key-type parameter to a Mapping hint') validate_composite_type_param( value_type, error_msg_prefix='Value-type parameter to a Mapping hint') return self.MappingTypeConstraint(key_type, value_type) MappingTypeConstraint = MappingHint.MappingTypeConstraint class IterableHint(CompositeTypeHint): """An Iterable type-hint. Iterable[X] defines a type-hint for an object implementing an '__iter__' method which yields objects which are all of the same type. """ class IterableTypeConstraint(SequenceTypeConstraint): def __init__(self, iter_type): super(IterableHint.IterableTypeConstraint, self).__init__(iter_type, abc.Iterable) def __repr__(self): return 'Iterable[%s]' % repr(self.inner_type) def _consistent_with_check_(self, sub): if isinstance(sub, SequenceTypeConstraint): return is_consistent_with(sub.inner_type, self.inner_type) elif isinstance(sub, TupleConstraint): if not sub.tuple_types: # The empty tuple is consistent with Iterator[T] for any T. return True # Each element in the hetrogenious tuple must be consistent with # the iterator type. # E.g. Tuple[A, B] < Iterable[C] if A < C and B < C. return all( is_consistent_with(elem, self.inner_type) for elem in sub.tuple_types) elif isinstance(sub, MappingTypeConstraint): return is_consistent_with(sub.key_type, self.inner_type) return False def __getitem__(self, type_param): validate_composite_type_param( type_param, error_msg_prefix='Parameter to an Iterable hint') return self.IterableTypeConstraint(type_param) IterableTypeConstraint = IterableHint.IterableTypeConstraint class IteratorHint(CompositeTypeHint): """An Iterator type-hint. Iterator[X] defines a type-hint for an object implementing both '__iter__' and a 'next' method which yields objects which are all of the same type. Type checking a type-hint of this type is deferred in order to avoid depleting the underlying lazily generated sequence. See decorators.interleave_type_check for further information. """ class IteratorTypeConstraint(TypeConstraint): def __init__(self, t): self.yielded_type = normalize(t) def __repr__(self): return 'Iterator[%s]' % repr(self.yielded_type) def __eq__(self, other): return ( type(self) == type(other) and self.yielded_type == other.yielded_type) def __hash__(self): return hash(self.yielded_type) def _inner_types(self): yield self.yielded_type def _consistent_with_check_(self, sub): return ( isinstance(sub, self.__class__) and is_consistent_with(sub.yielded_type, self.yielded_type)) def type_check(self, instance): # Special case for lazy types, we only need to enforce the underlying # type. This avoid having to compute the entirety of the generator/iter. try: check_constraint(self.yielded_type, instance) return except CompositeTypeHintError as e: raise CompositeTypeHintError( '%s hint type-constraint violated: %s' % (repr(self), str(e))) except SimpleTypeHintError: raise CompositeTypeHintError( '%s hint type-constraint violated. Expected a iterator of type %s. ' 'Instead received a iterator of type %s.' % (repr(self), repr(self.yielded_type), instance.__class__.__name__)) def __getitem__(self, type_param): validate_composite_type_param( type_param, error_msg_prefix='Parameter to an Iterator hint') return self.IteratorTypeConstraint(type_param) IteratorTypeConstraint = IteratorHint.IteratorTypeConstraint class WindowedTypeConstraint(TypeConstraint, metaclass=GetitemConstructor): """A type constraint for WindowedValue objects. Mostly for internal use. Attributes: inner_type: The type which the element should be an instance of. """ def __init__(self, inner_type): self.inner_type = normalize(inner_type) def __eq__(self, other): return ( isinstance(other, WindowedTypeConstraint) and self.inner_type == other.inner_type) def __hash__(self): return hash(self.inner_type) ^ 13 * hash(type(self)) def _inner_types(self): yield self.inner_type def _consistent_with_check_(self, sub): return ( isinstance(sub, self.__class__) and is_consistent_with(sub.inner_type, self.inner_type)) def type_check(self, instance): from apache_beam.transforms import window if not isinstance(instance, window.WindowedValue): raise CompositeTypeHintError( "Window type-constraint violated. Valid object instance " "must be of type 'WindowedValue'. Instead, an instance of '%s' " "was received." % (instance.__class__.__name__)) try: check_constraint(self.inner_type, instance.value) except (CompositeTypeHintError, SimpleTypeHintError): raise CompositeTypeHintError( '%s hint type-constraint violated. The type of element in ' 'is incorrect. Expected an instance of type %s, ' 'instead received an instance of type %s.' % ( repr(self), repr(self.inner_type), instance.value.__class__.__name__)) def bind_type_variables(self, bindings): bound_inner_type = bind_type_variables(self.inner_type, bindings) if bound_inner_type == self.inner_type: return self return WindowedValue[bound_inner_type] def __repr__(self): return 'WindowedValue[%s]' % repr(self.inner_type) class GeneratorHint(IteratorHint): """A Generator type hint. Subscriptor is in the form [yield_type, send_type, return_type], however only yield_type is supported. The 2 others are expected to be None. """ def __getitem__(self, type_params): if isinstance(type_params, tuple) and len(type_params) == 3: yield_type, send_type, return_type = type_params if send_type is not type(None): _LOGGER.warning('Ignoring send_type hint: %s' % send_type) if return_type is not type(None): _LOGGER.warning('Ignoring return_type hint: %s' % return_type) else: yield_type = type_params return self.IteratorTypeConstraint(yield_type) # Create the actual instances for all defined type-hints above. Any = AnyTypeConstraint() Union = UnionHint() Optional = OptionalHint() Tuple = TupleHint() List = ListHint() KV = KVHint() Dict = DictHint() Set = SetHint() FrozenSet = FrozenSetHint() Collection = CollectionHint() Sequence = SequenceHint() Mapping = MappingHint() Iterable = IterableHint() Iterator = IteratorHint() Generator = GeneratorHint() WindowedValue = WindowedTypeConstraint # There is a circular dependency between defining this mapping # and using it in normalize(). Initialize it here and populate # it below. _KNOWN_PRIMITIVE_TYPES: typing.Dict[type, typing.Any] = {} def normalize(x, none_as_type=False): # None is inconsistantly used for Any, unknown, or NoneType. # Avoid circular imports from apache_beam.typehints import native_type_compatibility if none_as_type and x is None: return type(None) # Convert bare builtin types to correct type hints directly elif x in _KNOWN_PRIMITIVE_TYPES: return _KNOWN_PRIMITIVE_TYPES[x] elif isinstance(x, types.UnionType) or getattr( x, '__module__', None) in ('typing', 'collections', 'collections.abc') or getattr( x, '__origin__', None) in _KNOWN_PRIMITIVE_TYPES: beam_type = native_type_compatibility.convert_to_beam_type(x) if beam_type != x: # We were able to do the conversion. return beam_type else: # It might be a compatible type we don't understand. return Any return x _KNOWN_PRIMITIVE_TYPES.update({ dict: Dict[Any, Any], list: List[Any], tuple: Tuple[Any, ...], set: Set[Any], frozenset: FrozenSet[Any], }) def is_consistent_with( sub, base, use_beartype: typing.Optional[bool] = None) -> bool: """Checks whether sub a is consistent with base. This is according to the terminology of PEP 483/484. This relationship is neither symmetric nor transitive, but a good mnemonic to keep in mind is that is_consistent_with(a, b) is roughly equivalent to the issubclass(a, b) relation, but also handles the special Any type as well as type parameterization. """ if use_beartype is None: from apache_beam.options.pipeline_options_context import get_pipeline_options options = get_pipeline_options() if options: from apache_beam.options.pipeline_options import TypeOptions use_beartype = not options.view_as(TypeOptions).disable_beartype else: use_beartype = True from apache_beam.pvalue import Row from apache_beam.typehints.row_type import RowTypeConstraint if sub == base: # Common special case. return True # Per PEP484, ints are considered floats and complexes and # floats are considered complexes. if sub is int and base in (float, complex): return True if sub is float and base is complex: return True sub = normalize(sub, none_as_type=True) base = normalize(base, none_as_type=True) if isinstance(sub, AnyTypeConstraint) or isinstance(base, AnyTypeConstraint): return True elif isinstance(sub, UnionConstraint): return all(is_consistent_with(c, base) for c in sub.union_types) elif isinstance(base, TypeConstraint): return base._consistent_with_check_(sub) elif isinstance(sub, RowTypeConstraint): return base == Row elif isinstance(sub, TypeConstraint): # Nothing but object lives above any type constraints. return base == object elif getattr(base, '__module__', None) == 're': return regex_consistency(sub, base) elif is_typing_generic(base): # Cannot check unsupported parameterized generic which will cause issubclass # to fail with an exception. return False if use_beartype: try: return is_subhint(sub, base) except (BeartypeDoorException): return False else: return issubclass(sub, base) def regex_consistency(sub, base) -> bool: """Checks whether two regular expression (re) type hints are consistent with each other. Either the sub or base hint can be parameterized generics since the set of possible parameters is restricted to str | bytes. A base hint without a parameter is treated as re.Class[str|bytes] so any sub param with a matching base class is consistent. On the flip side, a sub hint without a parameter is treated as inconsistent with a parameterized base hint. """ base_generic = getattr(base, '__origin__', None) sub_class = getattr(sub, '__origin__', sub) if base_generic: if sub_class == sub: # if the sub hint is not parameterized but the base hint is we # auto-fail return False return issubclass(sub_class, base_generic) and (sub.__args__ == base.__args__) else: return issubclass(sub_class, base) def get_yielded_type(type_hint): """Obtains the type of elements yielded by an iterable.s Note that "iterable" here means: can be iterated over in a for loop, excluding strings and dicts. Args: type_hint: (TypeConstraint) The iterable in question. Must be normalize()-d. Returns: Yielded type of the iterable. Raises: ValueError if not iterable. """ if isinstance(type_hint, typing.TypeVar): return typing.Any if isinstance(type_hint, AnyTypeConstraint): return type_hint if is_consistent_with(type_hint, Iterator[Any]): return type_hint.yielded_type if is_consistent_with(type_hint, Tuple[Any, ...]): if isinstance(type_hint, TupleConstraint): return Union[type_hint.tuple_types] else: # TupleSequenceConstraint return type_hint.inner_type if is_consistent_with(type_hint, Iterable[Any]): if isinstance(type_hint, UnionConstraint): yielded_types = set() for typ in type_hint.inner_types(): yielded_types.add(get_yielded_type(typ)) return Union[yielded_types] return type_hint.inner_type raise ValueError('%s is not iterable' % type_hint) def coerce_to_kv_type(element_type, label=None, side_input_producer=None): """Attempts to coerce element_type to a compatible kv type. Raises an error on failure. """ if side_input_producer: consumer = 'side-input of %r (producer: %r)' % (label, side_input_producer) else: consumer = '%r' % label # If element_type is not specified, then treat it as `Any`. if not element_type: return KV[Any, Any] elif isinstance(element_type, TupleHint.TupleConstraint): if len(element_type.tuple_types) == 2: return element_type else: raise ValueError( "Tuple input to %s must have two components. " "Found %s." % (consumer, element_type)) elif isinstance(element_type, AnyTypeConstraint): # `Any` type needs to be replaced with a KV[Any, Any] to # satisfy the KV form. return KV[Any, Any] elif isinstance(element_type, UnionConstraint): union_types = [coerce_to_kv_type(t) for t in element_type.union_types] return KV[Union[tuple(t.tuple_types[0] for t in union_types)], Union[tuple(t.tuple_types[1] for t in union_types)]] else: # TODO: Possibly handle other valid types. raise ValueError( "Input to %s must be compatible with KV[Any, Any]. " "Found %s." % (consumer, element_type)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/typehints/typehints_test.py0000644000175100017510000022034215206113033023350 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the type-hint objects and decorators.""" # pytype: skip-file # ruff: noqa: UP006 import collections.abc import functools import re import typing import unittest from apache_beam import Map from apache_beam import PTransform from apache_beam.pvalue import PBegin from apache_beam.pvalue import PCollection from apache_beam.pvalue import PDone from apache_beam.transforms.core import DoFn from apache_beam.typehints import KV from apache_beam.typehints import Any from apache_beam.typehints import Iterable from apache_beam.typehints import Tuple from apache_beam.typehints import TypeCheckError from apache_beam.typehints import Union from apache_beam.typehints import native_type_compatibility from apache_beam.typehints import typehints from apache_beam.typehints import with_input_types from apache_beam.typehints import with_output_types from apache_beam.typehints.decorators import GeneratorWrapper from apache_beam.typehints.decorators import _check_instance_type from apache_beam.typehints.decorators import _interleave_type_check from apache_beam.typehints.decorators import _positional_arg_hints from apache_beam.typehints.decorators import get_signature from apache_beam.typehints.decorators import get_type_hints from apache_beam.typehints.decorators import getcallargs_forhints from apache_beam.typehints.typehints import is_consistent_with from apache_beam.typehints.typehints import visit_inner_types def check_or_interleave(hint, value, var): if hint is None: return value elif isinstance(hint, typehints.IteratorHint.IteratorTypeConstraint): return _interleave_type_check(hint, var)(value) _check_instance_type(hint, value, var) return value def check_type_hints(f): @functools.wraps(f) def wrapper(*args, **kwargs): hints = get_type_hints(f) if hints.input_types: # pylint: disable=too-many-nested-blocks input_hints = getcallargs_forhints( f, *hints.input_types[0], **hints.input_types[1]) inputs = get_signature(f).bind(*args, **kwargs).arguments for var, hint in input_hints.items(): value = inputs[var] new_value = check_or_interleave(hint, value, var) if new_value is not value: if var in kwargs: kwargs[var] = new_value else: args = list(args) for ix, pvar in enumerate(get_signature(f).parameters): if pvar == var: args[ix] = new_value break else: raise NotImplementedError('Iterable in nested argument %s' % var) res = f(*args, **kwargs) return check_or_interleave(hints.simple_output_type('typecheck'), res, None) return wrapper class DummyTestClass1(object): pass class DummyTestClass2(object): pass class SuperClass(object): pass class SubClass(SuperClass): pass T = typing.TypeVar('T') class NonBuiltInGeneric(typing.NamedTuple('Entry', [('Field1', T), ('Field2', T)]), typing.Generic[T]): pass class TypeHintTestCase(unittest.TestCase): def assertCompatible(self, base, sub): # pylint: disable=invalid-name base, sub = native_type_compatibility.convert_to_beam_types([base, sub]) self.assertTrue( is_consistent_with(sub, base), '%s is not consistent with %s' % (sub, base)) def assertNotCompatible(self, base, sub): # pylint: disable=invalid-name base, sub = native_type_compatibility.convert_to_beam_type([base, sub]) self.assertFalse( is_consistent_with(sub, base), '%s is consistent with %s' % (sub, base)) class TypeVariableTestCase(TypeHintTestCase): def test_eq_with_name_check(self): use_name_in_eq = True self.assertNotEqual( typehints.Set[typehints.TypeVariable('T', use_name_in_eq)], typehints.Set[typehints.TypeVariable('T2', use_name_in_eq)]) def test_eq_without_name_check(self): use_name_in_eq = False self.assertEqual( typehints.Set[typehints.TypeVariable('T', use_name_in_eq)], typehints.Set[typehints.TypeVariable('T2', use_name_in_eq)]) def test_eq_with_different_name_checks(self): self.assertEqual( typehints.Set[typehints.TypeVariable('T', True)], typehints.Set[typehints.TypeVariable('T2', False)]) self.assertEqual( typehints.Set[typehints.TypeVariable('T', False)], typehints.Set[typehints.TypeVariable('T2', True)]) class AnyTypeConstraintTestCase(TypeHintTestCase): def test_any_compatibility(self): self.assertCompatible(typehints.Any, typehints.List[int]) self.assertCompatible(typehints.Any, DummyTestClass1) self.assertCompatible(typehints.Union[int, bool], typehints.Any) self.assertCompatible(typehints.Optional[int], typehints.Any) self.assertCompatible(typehints.Tuple[int], typehints.Any) self.assertCompatible(typehints.KV[int, str], typehints.Any) self.assertCompatible(typehints.Dict[int, bool], typehints.Any) self.assertCompatible(typehints.Set[int], typehints.Any) self.assertCompatible(typehints.FrozenSet[int], typehints.Any) self.assertCompatible(typehints.Iterable[int], typehints.Any) self.assertCompatible(typehints.Iterator[int], typehints.Any) self.assertCompatible(typehints.Generator[int], typehints.Any) self.assertCompatible(object, typehints.Any) self.assertCompatible(typehints.Any, object) def test_int_float_complex_compatibility(self): self.assertCompatible(float, int) self.assertCompatible(complex, int) self.assertCompatible(complex, float) self.assertNotCompatible(int, float) self.assertNotCompatible(int, complex) self.assertNotCompatible(float, complex) def test_repr(self): self.assertEqual('Any', repr(typehints.Any)) def test_type_check(self): # This test passes if the type_check call does not raise any exception. typehints.Any.type_check(1) typehints.Any.type_check([1, 2, 3]) typehints.Any.type_check(DummyTestClass1()) class UnionHintTestCase(TypeHintTestCase): def test_getitem_must_be_valid_type_param_cant_be_object_instance(self): with self.assertRaises(TypeError) as e: typehints.Union[5] self.assertEqual( 'Cannot create Union without a sequence of types.', e.exception.args[0]) def test_getitem_must_be_valid_type_param(self): t = [2, 3] with self.assertRaises(TypeError) as e: typehints.Union[t] self.assertEqual( 'All parameters to a Union hint must be a non-sequence, ' 'a type, or a TypeConstraint. 2 is an instance of int.', e.exception.args[0]) def test_getitem_duplicates_ignored(self): # Types should be de-duplicated. hint = typehints.Union[int, int, str] self.assertEqual(len(hint.union_types), 2) def test_getitem_nested_unions_flattened(self): # The two Union's should be merged into 1. hint = typehints.Union[typehints.Union[int, str], typehints.Union[float, bool]] self.assertTrue(len(hint.union_types) == 4) self.assertTrue(all(t in hint.union_types for t in (int, str, float, bool))) def test_union_hint_compatibility(self): self.assertCompatible(typehints.Union[int, float], int) self.assertCompatible(typehints.Union[int, str], typehints.Union[str, int]) self.assertCompatible( typehints.Union[int, float, str], typehints.Union[str, int]) self.assertCompatible( typehints.Union[DummyTestClass1, str], typehints.Union[DummyTestClass1, str]) self.assertCompatible( typehints.Union[int, str], typehints.Union[str, typehints.Union[int, str]]) self.assertCompatible( typehints.Union[float, bool], typehints.Union[int, bool]) self.assertNotCompatible( typehints.Union[bool, str], typehints.Union[float, bool, int]) def test_nested_compatibility(self): self.assertCompatible(Union[int, Tuple[Any, int]], Tuple[int, int]) self.assertCompatible( Union[int, Tuple[Any, Any]], Union[Tuple[int, Any], Tuple[Any, int]]) self.assertCompatible(Union[int, SuperClass], SubClass) self.assertCompatible(Union[int, float, SuperClass], Union[int, SubClass]) self.assertCompatible(int, Union[()]) self.assertNotCompatible(Union[int, SubClass], SuperClass) self.assertNotCompatible( Union[int, float, SubClass], Union[int, SuperClass]) self.assertNotCompatible( Union[int, SuperClass], Union[int, float, SubClass]) self.assertCompatible( Tuple[Any, Any], Union[Tuple[str, int], Tuple[str, float]]) def test_union_hint_repr_ordered_by_type(self): hint = typehints.Union[DummyTestClass1, str, int, bool] self.assertEqual( str(hint), "Union[, " ", , ]") def test_union_hint_enforcement_composite_type_in_union(self): o = DummyTestClass1() hint = typehints.Union[int, DummyTestClass1] self.assertIsNone(hint.type_check(4)) self.assertIsNone(hint.type_check(o)) def test_union_hint_enforcement_part_of_union(self): hint = typehints.Union[int, str] self.assertIsNone(hint.type_check(5)) self.assertIsNone(hint.type_check('test')) def test_union_hint_enforcement_not_part_of_union(self): hint = typehints.Union[int, float] with self.assertRaises(TypeError) as e: hint.type_check('test') self.assertEqual( "Union[, ] type-constraint violated. " "Expected an instance of one of: (\"\", \"\"), received str instead.", e.exception.args[0]) def test_dict_union(self): hint = Union[typehints.Dict[Any, int], typehints.Dict[Union[()], Union[()]]] self.assertEqual(typehints.Dict[Any, int], hint) def test_empty_union(self): self.assertEqual( typehints.Union[()], typehints.Union[typehints.Union[()], typehints.Union[()]]) self.assertEqual(int, typehints.Union[typehints.Union[()], int]) def test_match_type_variables(self): A = typehints.TypeVariable('A') # pylint: disable=invalid-name B = typehints.TypeVariable('B') # pylint: disable=invalid-name self.assertEqual( typehints.Union[A, int].match_type_variables(str), {A: str}) self.assertEqual(typehints.Union[A, int].match_type_variables(int), {}) self.assertEqual(typehints.Union[A, B, int].match_type_variables(str), {}) # We could do better here, but most importantly we don't want to # incorrectly infer A is a float. self.assertEqual( typehints.Tuple[A, typehints.Union[A, B, int]].match_type_variables( typehints.Tuple[str, float]), {A: str}) self.assertEqual( typehints.Union[Tuple[str, A], Tuple[float, B]].match_type_variables( typehints.Tuple[Any, int]), {}) self.assertEqual( typehints.Union[Tuple[str, A], Tuple[float, A]].match_type_variables( typehints.Tuple[Any, int]), {A: int}) self.assertEqual( typehints.Union[Tuple[str, A], Tuple[float, B]].match_type_variables( typehints.Tuple[str, int]), {A: int}) def test_bind_type_variables(self): A = typehints.TypeVariable('A') # pylint: disable=invalid-name B = typehints.TypeVariable('B') # pylint: disable=invalid-name hint = typehints.Union[A, B, int] self.assertEqual( hint.bind_type_variables({ A: str, B: float }), typehints.Union[str, float, int]) self.assertEqual( hint.bind_type_variables({ A: str, B: int }), typehints.Union[str, int]) self.assertEqual(hint.bind_type_variables({A: int, B: int}), int) def test_visit_inner_types(self): A = typehints.TypeVariable('A') # pylint: disable=invalid-name B = typehints.TypeVariable('B') # pylint: disable=invalid-name hint = typehints.Tuple[Tuple[A, A], B, int] user_data = object() nodes = [] def visitor(hint, arg): self.assertIs(arg, user_data) nodes.append(hint) visit_inner_types(hint, visitor, user_data) self.assertEqual(nodes, [hint, Tuple[A, A], A, A, B, int]) class OptionalHintTestCase(TypeHintTestCase): def test_getitem_sequence_not_allowed(self): with self.assertRaises(TypeError) as e: typehints.Optional[int, str] self.assertTrue( e.exception.args[0].startswith( 'An Option type-hint only accepts a single type parameter.')) def test_getitem_proxy_to_union(self): hint = typehints.Optional[int] self.assertTrue(isinstance(hint, typehints.UnionHint.UnionConstraint)) def test_is_optional(self): hint1 = typehints.Optional[int] self.assertTrue(typehints.is_nullable(hint1)) hint2 = typehints.UnionConstraint({int, bytes}) self.assertFalse(typehints.is_nullable(hint2)) hint3 = typehints.UnionConstraint({int, bytes, type(None)}) self.assertFalse(typehints.is_nullable(hint3)) class TupleHintTestCase(TypeHintTestCase): def test_getitem_invalid_ellipsis_type_param(self): error_msg = ( 'Ellipsis can only be used to type-hint an arbitrary length ' 'tuple of containing a single type: Tuple[A, ...].') with self.assertRaises(TypeError) as e: typehints.Tuple[int, int, ...] self.assertEqual(error_msg, e.exception.args[0]) with self.assertRaises(TypeError) as e: typehints.Tuple[...] self.assertEqual(error_msg, e.exception.args[0]) def test_getitem_params_must_be_type_or_constraint(self): expected_error_prefix = 'All parameters to a Tuple hint must be' with self.assertRaises(TypeError) as e: typehints.Tuple[5, [1, 3]] self.assertTrue(e.exception.args[0].startswith(expected_error_prefix)) try: typehints.Tuple[list, dict] except TypeError: self.fail("built-in composite raised TypeError unexpectedly") def test_compatibility_arbitrary_length(self): self.assertNotCompatible( typehints.Tuple[int, int], typehints.Tuple[int, ...]) self.assertCompatible(typehints.Tuple[int, ...], typehints.Tuple[int, int]) self.assertCompatible( typehints.Tuple[Any, ...], typehints.Tuple[int, float]) self.assertCompatible( typehints.Tuple[SuperClass, ...], typehints.Tuple[SubClass, SuperClass]) self.assertCompatible(typehints.Iterable[int], typehints.Tuple[int, ...]) self.assertCompatible( typehints.Iterable[SuperClass], typehints.Tuple[SubClass, ...]) def test_compatibility(self): self.assertCompatible(typehints.Tuple[int, str], typehints.Tuple[int, str]) self.assertCompatible(typehints.Tuple[int, Any], typehints.Tuple[int, str]) self.assertCompatible(typehints.Tuple[int, str], typehints.Tuple[int, Any]) self.assertCompatible( typehints.Tuple[typehints.Union[int, str], bool], typehints.Tuple[typehints.Union[int, str], bool]) self.assertCompatible( typehints.Tuple[typehints.Union[str, int], int], typehints.Tuple[typehints.Union[int, str], int]) self.assertCompatible( typehints.Tuple[SuperClass, int], typehints.Tuple[SubClass, int]) self.assertNotCompatible( typehints.Tuple[int, int], typehints.Tuple[int, int, int]) def test_raw_tuple(self): self.assertCompatible(tuple, typehints.Tuple[int]) self.assertCompatible(tuple, typehints.Tuple[int, float]) self.assertCompatible(tuple, typehints.Tuple[int, ...]) def test_repr(self): hint = typehints.Tuple[int, str, float] self.assertEqual( 'Tuple[, , ]', str(hint)) hint = typehints.Tuple[DummyTestClass1, DummyTestClass2] self.assertEqual( 'Tuple[, ]', str(hint)) hint = typehints.Tuple[float, ...] self.assertEqual('Tuple[, ...]', str(hint)) def test_type_check_must_be_tuple(self): hint = typehints.Tuple[int, str] expected_error_prefix = 'Tuple type constraint violated. Valid object' invalid_instances = ([1, 2, 3], {4: 'f'}, 9, 'test', None) for t in invalid_instances: with self.assertRaises(TypeError) as e: hint.type_check(t) self.assertTrue(e.exception.args[0].startswith(expected_error_prefix)) def test_type_check_must_have_same_arity(self): # A 2-tuple of ints. hint = typehints.Tuple[int, int] t = (1, 2, 3) with self.assertRaises(TypeError) as e: hint.type_check(t) self.assertEqual( 'Passed object instance is of the proper type, but ' 'differs in length from the hinted type. Expected a ' 'tuple of length 2, received a tuple of length 3.', e.exception.args[0]) def test_type_check_invalid_simple_types(self): hint = typehints.Tuple[str, bool] with self.assertRaises(TypeError) as e: hint.type_check((4, False)) self.assertEqual( 'Tuple[, ] hint type-constraint ' 'violated. The type of element #0 in the passed tuple is incorrect.' ' Expected an instance of type , instead received ' 'an instance of type int.', e.exception.args[0]) def test_type_check_invalid_composite_type(self): hint = typehints.Tuple[DummyTestClass1, DummyTestClass2] t = (DummyTestClass2(), DummyTestClass1()) with self.assertRaises(TypeError) as e: hint.type_check(t) self.assertEqual( 'Tuple[, ]' ' hint type-constraint violated. The type of element #0 in the ' 'passed tuple is incorrect. Expected an instance of type ' ', ' 'instead received an instance of type ' 'DummyTestClass2.', e.exception.args[0]) def test_type_check_valid_simple_types(self): hint = typehints.Tuple[float, bool] self.assertIsNone(hint.type_check((4.3, True))) hint = typehints.Tuple[int] self.assertIsNone(hint.type_check((1, ))) def test_type_check_valid_composite_types(self): hint = typehints.Tuple[typehints.Tuple[int, str], typehints.Tuple[int, bool]] self.assertIsNone(hint.type_check(((4, 'test'), (4, True)))) def test_type_check_valid_simple_type_arbitrary_length(self): hint = typehints.Tuple[int, ...] t = (1, 2, 3, 4) self.assertIsNone(hint.type_check(t)) def test_type_check_valid_composite_type_arbitrary_length(self): hint = typehints.Tuple[typehints.List[str], ...] t = (['h', 'e'], ['l', 'l'], ['o']) self.assertIsNone(hint.type_check(t)) def test_type_check_invalid_simple_type_arbitrary_length(self): hint = typehints.Tuple[str, ...] t = ('t', 'e', 5, 't') with self.assertRaises(TypeError) as e: hint.type_check(t) self.assertEqual( 'Tuple[, ...] hint type-constraint violated. The type ' 'of element #2 in the passed is incorrect. Expected ' 'an instance of type , instead received an instance of ' 'type int.', e.exception.args[0]) def test_type_check_invalid_composite_type_arbitrary_length(self): hint = typehints.Tuple[typehints.List[int], ...] t = ([1, 2], 'e', 's', 't') with self.assertRaises(TypeError) as e: hint.type_check(t) self.assertEqual( "Tuple[List[], ...] hint type-constraint violated. The " "type of element #1 in the passed tuple is incorrect: " "List type-constraint violated. Valid object instance " "must be of type 'list'. Instead, an instance of 'str' " "was received.", e.exception.args[0]) def test_normalize_with_builtin_tuple(self): expected_beam_type = typehints.Tuple[int, int] converted_beam_type = typehints.normalize(tuple[int, int], False) self.assertEqual(converted_beam_type, expected_beam_type) def test_builtin_and_type_compatibility(self): self.assertCompatible(tuple, typing.Tuple) self.assertCompatible(tuple[int, int], typing.Tuple[int, int]) class ListHintTestCase(TypeHintTestCase): def test_getitem_invalid_composite_type_param(self): with self.assertRaises(TypeError): typehints.List[4] def test_list_constraint_compatibility(self): hint1 = typehints.List[typehints.Tuple[int, str]] hint2 = typehints.List[typehints.Tuple[float, bool]] self.assertCompatible(hint1, hint1) self.assertNotCompatible(hint1, hint2) self.assertCompatible(typehints.List[SuperClass], typehints.List[SubClass]) def test_list_repr(self): hint = (typehints.List[typehints.Tuple[DummyTestClass1, DummyTestClass2]]) self.assertEqual( 'List[Tuple[, ]]', repr(hint)) def test_enforce_list_type_constraint_valid_simple_type(self): hint = typehints.List[int] self.assertIsNone(hint.type_check([1, 2, 3])) def test_enforce_list_type_constraint_valid_composite_type(self): hint = typehints.List[DummyTestClass1] l = [DummyTestClass1(), DummyTestClass1()] self.assertIsNone(hint.type_check(l)) def test_enforce_list_type_constraint_invalid_simple_type(self): hint = typehints.List[int] l = ['f', 'd', 'm'] with self.assertRaises(TypeError) as e: hint.type_check(l) self.assertEqual( 'List[] hint type-constraint violated. The type of ' 'element #0 in the passed is incorrect. Expected an ' 'instance of type , instead received an instance of ' 'type str.', e.exception.args[0]) def test_enforce_list_type_constraint_invalid_composite_type(self): hint = typehints.List[typehints.Tuple[int, int]] l = [('f', 1), ('m', 5)] with self.assertRaises(TypeError) as e: hint.type_check(l) self.assertEqual( 'List[Tuple[, ]] hint type-constraint ' 'violated. The type of element #0 in the passed list is ' 'incorrect: Tuple[, ] hint ' 'type-constraint violated. The type of element #0 in the passed tuple' ' is incorrect. Expected an instance of type , ' 'instead received an instance of type str.', e.exception.args[0]) def test_normalize_with_builtin_list(self): expected_beam_type = typehints.List[int] converted_beam_type = typehints.normalize(list[int], False) self.assertEqual(converted_beam_type, expected_beam_type) def test_builtin_and_type_compatibility(self): self.assertCompatible(list, typing.List) self.assertCompatible(list[int], typing.List[int]) def test_is_typing_generic(self): self.assertTrue(typehints.is_typing_generic(typing.List[str])) def test_builtin_is_typing_generic(self): self.assertTrue(typehints.is_typing_generic(list[str])) class KVHintTestCase(TypeHintTestCase): def test_getitem_param_must_be_tuple(self): with self.assertRaises(TypeError) as e: typehints.KV[4] self.assertEqual( 'Parameter to KV type-hint must be a tuple of types: ' 'KV[.., ..].', e.exception.args[0]) def test_getitem_param_must_have_length_2(self): with self.assertRaises(TypeError) as e: typehints.KV[int, str, bool] self.assertEqual( "Length of parameters to a KV type-hint must be " "exactly 2. Passed parameters: ({}, {}, {}), have a " "length of 3.".format(int, str, bool), e.exception.args[0]) def test_getitem_proxy_to_tuple(self): hint = typehints.KV[int, str] self.assertTrue(isinstance(hint, typehints.Tuple.TupleConstraint)) def test_enforce_kv_type_constraint(self): hint = typehints.KV[str, typehints.Tuple[int, int, int]] t = ('test', (1, 2, 3)) self.assertIsNone(hint.type_check(t)) class DictHintTestCase(TypeHintTestCase): def test_getitem_param_must_be_tuple(self): with self.assertRaises(TypeError) as e: typehints.Dict[4] self.assertEqual( 'Parameter to Dict type-hint must be a tuple of ' 'types: Dict[.., ..].', e.exception.args[0]) def test_getitem_param_must_have_length_2(self): with self.assertRaises(TypeError) as e: typehints.Dict[float, int, bool] self.assertEqual( "Length of parameters to a Dict type-hint must be " "exactly 2. Passed parameters: ({}, {}, {}), have a " "length of 3.".format(float, int, bool), e.exception.args[0]) def test_key_type_must_be_valid_composite_param(self): try: typehints.Tuple[list, int] except TypeError: self.fail("built-in composite raised TypeError unexpectedly") def test_value_type_must_be_valid_composite_param(self): with self.assertRaises(TypeError): typehints.Dict[str, 5] def test_compatibility(self): hint1 = typehints.Dict[int, str] hint2 = typehints.Dict[bool, int] hint3 = typehints.Dict[int, typehints.List[typehints.Tuple[str, str, str]]] hint4 = typehints.Dict[int, int] self.assertCompatible(hint1, hint1) self.assertCompatible(hint3, hint3) self.assertNotCompatible(hint3, 4) self.assertNotCompatible(hint2, hint1) # Key incompatibility. self.assertNotCompatible(hint1, hint4) # Value incompatibility. def test_repr(self): hint3 = typehints.Dict[int, typehints.List[typehints.Tuple[str, str, str]]] self.assertEqual( 'Dict[, List[Tuple[, ' \ ', ]]]', repr(hint3)) def test_type_checks_not_dict(self): hint = typehints.Dict[int, str] l = [1, 2] with self.assertRaises(TypeError) as e: hint.type_check(l) self.assertEqual( 'Dict type-constraint violated. All passed instances ' 'must be of type dict. [1, 2] is of type list.', e.exception.args[0]) def test_type_check_collection(self): hint = typehints.Dict[str, int] element = ("blue", 2) l = collections.defaultdict(list[element]) self.assertIsNone(hint.type_check(l)) def test_type_check_invalid_key_type(self): hint = typehints.Dict[typehints.Tuple[int, int, int], typehints.List[str]] d = {(1, 2): ['m', '1', '2', '3']} with self.assertRaises((TypeError, TypeError)) as e: hint.type_check(d) self.assertEqual( 'Dict[Tuple[, , ], ' 'List[]] hint key-type ' 'constraint violated. All keys should be of type ' 'Tuple[, , ]. Instead: ' 'Passed object instance is of the proper type, but differs in ' 'length from the hinted type. Expected a tuple of ' 'length 3, received a tuple of length 2.', e.exception.args[0]) def test_type_check_invalid_value_type(self): hint = typehints.Dict[str, typehints.Dict[int, str]] d = {'f': [1, 2, 3]} with self.assertRaises(TypeError) as e: hint.type_check(d) self.assertEqual( 'Dict[, Dict[, ]] hint ' 'value-type constraint violated. All values should be of type ' 'Dict[, ]. Instead: Dict type-constraint ' 'violated. All passed instances must be of type dict.' ' [1, 2, 3] is of type list.', e.exception.args[0]) def test_type_check_valid_simple_type(self): hint = typehints.Dict[int, str] d = {4: 'f', 9: 'k'} self.assertIsNone(hint.type_check(d)) def test_type_check_valid_composite_type(self): hint = typehints.Dict[typehints.Tuple[str, str], typehints.List[int]] d = {('f', 'k'): [1, 2, 3], ('m', 'r'): [4, 6, 9]} self.assertIsNone(hint.type_check(d)) def test_match_type_variables(self): S = typehints.TypeVariable('S') # pylint: disable=invalid-name T = typehints.TypeVariable('T') # pylint: disable=invalid-name hint = typehints.Dict[S, T] self.assertEqual({ S: int, T: str }, hint.match_type_variables(typehints.Dict[int, str])) def test_normalize_with_builtin_dict(self): expected_beam_type = typehints.Dict[str, int] converted_beam_type = typehints.normalize(dict[str, int], False) self.assertEqual(converted_beam_type, expected_beam_type) def test_normalize_with_collections_dicts(self): test_cases = [ ( 'default dict', collections.defaultdict[str, bool], typehints.Dict[str, bool]), ( 'ordered dict', collections.OrderedDict[str, bool], typehints.Dict[str, bool]), ('counter', collections.Counter[str, int], typehints.Dict[str, int]), ] for test_case in test_cases: description = test_case[0] collections_type = test_case[1] expected_beam_type = test_case[2] converted_beam_type = typehints.normalize(collections_type) self.assertEqual(converted_beam_type, expected_beam_type, description) def test_builtin_and_type_compatibility(self): self.assertCompatible(dict, typing.Dict) self.assertCompatible(dict[str, int], typing.Dict[str, int]) self.assertCompatible( dict[str, list[int]], typing.Dict[str, typing.List[int]]) def test_collections_subclass_compatibility(self): self.assertCompatible( collections.defaultdict[str, bool], typing.Dict[str, bool]) self.assertCompatible( collections.OrderedDict[str, int], typing.Dict[str, int]) self.assertCompatible(collections.Counter[str, int], typing.Dict[str, int]) class MappingHintTestCase(TypeHintTestCase): def test_getitem_param_must_be_tuple(self): with self.assertRaises(TypeError) as e: typehints.Mapping[4] self.assertEqual( 'Parameter to Mapping type-hint must be a tuple of ' 'types: Mapping[.., ..].', e.exception.args[0]) def test_getitem_param_must_have_length_2(self): with self.assertRaises(TypeError) as e: typehints.Mapping[float, int, bool] self.assertEqual( "Length of parameters to a Mapping type-hint must be " "exactly 2. Passed parameters: ({}, {}, {}), have a " "length of 3.".format(float, int, bool), e.exception.args[0]) def test_key_type_must_be_valid_composite_param(self): try: typehints.Mapping[list, int] except TypeError: self.fail("built-in composite raised TypeError unexpectedly") def test_value_type_must_be_valid_composite_param(self): with self.assertRaises(TypeError): typehints.Mapping[str, 5] def test_compatibility(self): hint1 = typehints.Mapping[int, str] hint2 = typehints.Mapping[bool, int] hint3 = typehints.Mapping[int, typehints.List[typehints.Tuple[str, str, str]]] hint4 = typehints.Mapping[int, int] self.assertCompatible(hint1, hint1) self.assertCompatible(hint3, hint3) self.assertNotCompatible(hint3, 4) self.assertNotCompatible(hint2, hint1) # Key incompatibility. self.assertNotCompatible(hint1, hint4) # Value incompatibility. def test_repr(self): hint3 = typehints.Mapping[int, typehints.List[typehints.Tuple[str, str, str]]] self.assertEqual( 'Mapping[, List[Tuple[, ' \ ', ]]]', repr(hint3)) def test_type_checks_not_dict(self): hint = typehints.Mapping[int, str] l = [1, 2] with self.assertRaises(TypeError) as e: hint.type_check(l) self.assertEqual( 'Mapping type-constraint violated. All passed instances ' 'must be of type Mapping. [1, 2] is of type list.', e.exception.args[0]) def test_type_check_invalid_key_type(self): hint = typehints.Mapping[typehints.Tuple[int, int, int], typehints.List[str]] d = {(1, 2): ['m', '1', '2', '3']} with self.assertRaises((TypeError, TypeError)) as e: hint.type_check(d) self.assertEqual( 'Mapping[Tuple[, , ], ' 'List[]] hint key-type ' 'constraint violated. All keys should be of type ' 'Tuple[, , ]. Instead: ' 'Passed object instance is of the proper type, but differs in ' 'length from the hinted type. Expected a tuple of ' 'length 3, received a tuple of length 2.', e.exception.args[0]) def test_type_check_invalid_value_type(self): hint = typehints.Mapping[str, typehints.Mapping[int, str]] d = {'f': [1, 2, 3]} with self.assertRaises(TypeError) as e: hint.type_check(d) self.assertEqual( "Mapping[, Mapping[, ]] hint" ' value-type constraint violated. All values should be of type' " Mapping[, ]. Instead: Mapping" ' type-constraint violated. All passed instances must be of type' ' Mapping. [1, 2, 3] is of type list.', e.exception.args[0], ) def test_type_check_valid_simple_type(self): hint = typehints.Mapping[int, str] d = {4: 'f', 9: 'k'} self.assertIsNone(hint.type_check(d)) def test_type_check_valid_composite_type(self): hint = typehints.Mapping[typehints.Tuple[str, str], typehints.List[int]] d = {('f', 'k'): [1, 2, 3], ('m', 'r'): [4, 6, 9]} self.assertIsNone(hint.type_check(d)) def test_match_type_variables(self): S = typehints.TypeVariable('S') # pylint: disable=invalid-name T = typehints.TypeVariable('T') # pylint: disable=invalid-name hint = typehints.Mapping[S, T] self.assertEqual({ S: int, T: str }, hint.match_type_variables(typehints.Mapping[int, str])) def test_builtin_and_type_compatibility(self): self.assertCompatible(typing.Mapping, dict) self.assertCompatible(typing.Mapping[str, int], dict[str, int]) self.assertCompatible( typing.Mapping[str, typing.List[int]], dict[str, list[int]]) self.assertCompatible(typing.Iterable[str], typing.Mapping[str, int]) self.assertNotCompatible(typing.Mapping[str, int], typing.Iterable[str]) self.assertCompatible(typing.Mapping[str, int], typing.Mapping[str, int]) def test_collections_compatibility(self): self.assertCompatible(typing.Mapping, collections.defaultdict) self.assertCompatible(typing.Mapping, collections.defaultdict[Any, Any]) self.assertCompatible( typing.Mapping[str, int], collections.defaultdict[str, int]) self.assertCompatible( typing.Mapping[str, int], collections.OrderedDict[str, int]) class BaseSetHintTest: class CommonTests(TypeHintTestCase): def test_getitem_invalid_composite_type_param(self): try: self.beam_type[list] except TypeError: self.fail("built-in composite raised TypeError unexpectedly") def test_non_typing_generic(self): testCase = DummyTestClass1() self.assertFalse(typehints.is_typing_generic(testCase)) def test_compatibility(self): hint1 = self.beam_type[typehints.List[str]] hint2 = self.beam_type[typehints.Tuple[int, int]] self.assertCompatible(hint1, hint1) self.assertNotCompatible(hint2, hint1) def test_repr(self): hint = self.beam_type[typehints.List[bool]] self.assertEqual( '{}[List[]]'.format(self.string_type), repr(hint)) def test_type_check_must_be_set(self): hint = self.beam_type[str] with self.assertRaises(TypeError) as e: hint.type_check(4) self.assertEqual( "{} type-constraint violated. Valid object instance " "must be of type '{}'. Instead, an instance of 'int'" " was received.".format(self.py_type, self.py_type.__name__), e.exception.args[0]) def test_type_check_invalid_elem_type(self): hint = self.beam_type[float] with self.assertRaises(TypeError): hint.type_check('f') def test_type_check_valid_elem_simple_type(self): hint = self.beam_type[str] s = self.py_type(['f', 'm', 'k']) self.assertIsNone(hint.type_check(s)) def test_type_check_valid_elem_composite_type(self): hint = self.beam_type[typehints.Union[int, str]] s = self.py_type([9, 'm', 'k']) self.assertIsNone(hint.type_check(s)) class SetHintTestCase(BaseSetHintTest.CommonTests): py_type = set beam_type = typehints.Set string_type = 'Set' def test_builtin_compatibility(self): self.assertCompatible(set[int], collections.abc.Set[int]) self.assertCompatible(set[int], collections.abc.MutableSet[int]) def test_collections_compatibility(self): self.assertCompatible( collections.abc.Set[int], collections.abc.MutableSet[int]) self.assertCompatible( collections.abc.MutableSet[int], collections.abc.Set[int]) class FrozenSetHintTestCase(BaseSetHintTest.CommonTests): py_type = frozenset beam_type = typehints.FrozenSet string_type = 'FrozenSet' class CollectionHintTestCase(TypeHintTestCase): def test_type_constraint_compatibility(self): self.assertCompatible(typehints.Collection[int], typehints.Set[int]) self.assertCompatible(typehints.Iterable[int], typehints.Collection[int]) self.assertCompatible(typehints.Collection[int], typehints.FrozenSet[int]) self.assertCompatible( typehints.Collection[typehints.Any], typehints.Collection[int]) self.assertCompatible(typehints.Collection[int], typehints.Tuple[int]) self.assertCompatible(typehints.Any, typehints.Collection[str]) self.assertCompatible(typehints.Collection[str], typehints.List[str]) def test_one_way_compatibility(self): self.assertNotCompatible(typehints.Set[int], typehints.Collection[int]) self.assertNotCompatible( typehints.FrozenSet[int], typehints.Collection[int]) self.assertNotCompatible(typehints.Tuple[int], typehints.Collection[int]) self.assertNotCompatible(typehints.Collection[int], typehints.Iterable[int]) self.assertNotCompatible(typehints.List[str], typehints.Collection[str]) def test_getitem_invalid_composite_type_param(self): with self.assertRaises(TypeError) as e: typehints.Collection[5] self.assertEqual( 'Parameter to a Collection hint must be a ' 'non-sequence, a type, or a TypeConstraint. 5 is ' 'an instance of int.', e.exception.args[0]) class SequenceHintTestCase(TypeHintTestCase): def test_type_constraint_compatibility(self): self.assertCompatible(typehints.Sequence[str], typehints.List[str]) self.assertCompatible(typehints.Sequence[str], typehints.Tuple[str]) self.assertCompatible( typehints.Sequence[typehints.Any], typehints.Sequence[str]) self.assertCompatible( typehints.Sequence[str], typehints.Sequence[typehints.Any]) self.assertCompatible(typehints.Any, typehints.Sequence[str]) def test_one_way_compatibility(self): self.assertNotCompatible(typehints.List[str], typehints.Sequence[str]) self.assertNotCompatible(typehints.Tuple[str], typehints.Sequence[str]) def test_getitem_invalid_composite_type_param(self): with self.assertRaises(TypeError) as e: typehints.Sequence[5] self.assertEqual( 'Parameter to a Sequence hint must be a ' 'non-sequence, a type, or a TypeConstraint. 5 is ' 'an instance of int.', e.exception.args[0]) def test_normalize(self): expected_beam_type = typehints.Sequence[int] converted_beam_type = typehints.normalize( collections.abc.Sequence[int], False) self.assertEqual(converted_beam_type, expected_beam_type) class IterableHintTestCase(TypeHintTestCase): def test_getitem_invalid_composite_type_param(self): with self.assertRaises(TypeError) as e: typehints.Iterable[5] self.assertEqual( 'Parameter to an Iterable hint must be a ' 'non-sequence, a type, or a TypeConstraint. 5 is ' 'an instance of int.', e.exception.args[0]) def test_compatibility(self): self.assertCompatible(typehints.Iterable[int], typehints.List[int]) self.assertCompatible(typehints.Iterable[int], typehints.Set[int]) self.assertCompatible( typehints.Iterable[typehints.Any], typehints.List[typehints.Tuple[int, bool]]) self.assertCompatible(typehints.Iterable[str], typehints.Sequence[str]) self.assertCompatible(typehints.Iterable[int], typehints.Iterable[int]) self.assertCompatible( typehints.Iterable[typehints.Union[int, str]], typehints.Iterable[typehints.Union[int, str]]) self.assertNotCompatible(typehints.Iterable[str], typehints.Iterable[bool]) self.assertCompatible(typehints.Iterable[int], typehints.List[int]) self.assertCompatible(typehints.Iterable[int], typehints.Set[int]) self.assertCompatible( typehints.Iterable[typehints.Any], typehints.List[typehints.Tuple[int, bool]]) self.assertCompatible(typehints.Iterable[int], typehints.Collection[int]) def test_tuple_compatibility(self): self.assertCompatible(typehints.Iterable[int], typehints.Tuple[int, ...]) self.assertCompatible( typehints.Iterable[SuperClass], typehints.Tuple[SubClass, ...]) self.assertCompatible(typehints.Iterable[int], typehints.Tuple[int, int]) self.assertCompatible(typehints.Iterable[Any], typehints.Tuple[int, float]) self.assertCompatible( typehints.Iterable[typehints.Union[int, float]], typehints.Tuple[int, ...]) self.assertCompatible( typehints.Iterable[typehints.Union[int, float]], typehints.Tuple[int, float]) self.assertCompatible( typehints.Iterable[typehints.Union[int, float]], typehints.Tuple[int, float, int]) def test_repr(self): hint = typehints.Iterable[typehints.Set[str]] self.assertEqual('Iterable[Set[]]', repr(hint)) def test_type_check_must_be_iterable(self): with self.assertRaises(TypeError) as e: hint = typehints.Iterable[int] hint.type_check(5) self.assertEqual( "Iterable type-constraint violated. Valid object " "instance must be of type 'iterable'. Instead, an " "instance of 'int' was received.", e.exception.args[0]) def test_type_check_violation_invalid_simple_type(self): hint = typehints.Iterable[float] l = set([1, 2, 3, 4]) with self.assertRaises(TypeError): hint.type_check(l) def test_type_check_violation_valid_simple_type(self): hint = typehints.Iterable[str] l = ('t', 'e', 's', 't') self.assertIsNone(hint.type_check(l)) def test_type_check_violation_invalid_composite_type(self): hint = typehints.Iterable[typehints.List[int]] l = ([['t', 'e'], ['s', 't']]) with self.assertRaises(TypeError): hint.type_check(l) def test_type_check_violation_valid_composite_type(self): hint = typehints.Iterable[typehints.List[int]] l = ([[1, 2], [3, 4, 5]]) self.assertIsNone(hint.type_check(l)) def test_normalize(self): expected_beam_type = typehints.Iterable[str] converted_beam_type = typehints.normalize( collections.abc.Iterable[str], False) self.assertEqual(converted_beam_type, expected_beam_type) class TestGeneratorWrapper(TypeHintTestCase): def test_functions_as_regular_generator(self): def count(n): for i in range(n): yield i l = [] interleave_func = lambda x: l.append(x) wrapped_gen = GeneratorWrapper(count(4), interleave_func) # Should function as a normal generator. self.assertEqual(0, next(wrapped_gen)) self.assertEqual((1, 2, 3), tuple(wrapped_gen)) # Interleave function should have been called each time. self.assertEqual([0, 1, 2, 3], l) class GeneratorHintTestCase(TypeHintTestCase): def test_repr(self): hint = typehints.Iterator[typehints.Set[str]] self.assertEqual('Iterator[Set[]]', repr(hint)) def test_compatibility(self): self.assertCompatible(typehints.Iterator[int], typehints.Iterator[int]) self.assertNotCompatible(typehints.Iterator[str], typehints.Iterator[float]) def test_conversion(self): self.assertCompatible(typehints.Iterator[int], typehints.Generator[int]) self.assertCompatible( typehints.Iterator[int], typehints.Generator[int, None, None]) def test_generator_return_hint_invalid_yield_type(self): @check_type_hints @with_output_types(typehints.Iterator[int]) def all_upper(s): for e in s: yield e.upper() with self.assertRaises(TypeCheckError) as e: next(all_upper('hello')) self.assertEqual( 'Type-hint for return type violated: Iterator[] ' 'hint type-constraint violated. Expected a iterator ' 'of type . Instead received a iterator of type ' 'str.', e.exception.args[0]) def test_generator_argument_hint_invalid_yield_type(self): def wrong_yield_gen(): for e in ['a', 'b']: yield e @check_type_hints @with_input_types(a=typehints.Iterator[int]) def increment(a): return [e + 1 for e in a] with self.assertRaises(TypeCheckError) as e: increment(wrong_yield_gen()) self.assertEqual( "Type-hint for argument: 'a' violated: Iterator[] " "hint type-constraint violated. Expected a iterator " "of type . Instead received a iterator of type " "str.", e.exception.args[0]) class RegexTestCase(TypeHintTestCase): def test_pattern(self): self.assertCompatible(re.Pattern, re.Pattern) self.assertCompatible(re.Pattern, re.Pattern[str]) self.assertCompatible(re.Pattern, re.Pattern[bytes]) self.assertCompatible(re.Pattern[str], re.Pattern[str]) self.assertCompatible(re.Pattern[bytes], re.Pattern[bytes]) self.assertNotCompatible(re.Pattern[str], re.Pattern) self.assertNotCompatible(re.Pattern[bytes], re.Pattern) self.assertNotCompatible(re.Pattern[str], re.Pattern[bytes]) self.assertNotCompatible(re.Pattern[bytes], re.Pattern[str]) def test_match(self): self.assertCompatible(re.Match, re.Match) self.assertCompatible(re.Match, re.Match[str]) self.assertCompatible(re.Match, re.Match[bytes]) self.assertCompatible(re.Match[str], re.Match[str]) self.assertCompatible(re.Match[bytes], re.Match[bytes]) self.assertNotCompatible(re.Match[str], re.Match) self.assertNotCompatible(re.Match[bytes], re.Match) self.assertNotCompatible(re.Match[str], re.Match[bytes]) self.assertNotCompatible(re.Match[bytes], re.Match[str]) def test_mix_fails(self): self.assertNotCompatible(re.Pattern, re.Match) self.assertNotCompatible(re.Match, re.Pattern) self.assertNotCompatible(re.Pattern[str], re.Match[str]) # TODO(https://github.com/apache/beam/issues/34644): implement a # CallableTypeHint class and implement test suites for inner type checking class CallableTestCase(TypeHintTestCase): def test_callable(self): self.assertCompatible(collections.abc.Callable, collections.abc.Callable) self.assertCompatible( collections.abc.Callable[[str], int], collections.abc.Callable[[str], int]) def test_normalize_to_any_bare(self): input = collections.abc.Callable output = typehints.normalize(input) self.assertEqual(Any, output) def test_normalize_to_any_parameterized(self): input = collections.abc.Callable[[str], int] output = typehints.normalize(input) self.assertEqual(Any, output) class TakesDecoratorTestCase(TypeHintTestCase): def test_must_be_primitive_type_or_constraint(self): with self.assertRaises(TypeError) as e: t = [1, 2] @with_input_types(a=t) def unused_foo(a): pass self.assertEqual( 'All type hint arguments must be a non-sequence, a ' 'type, or a TypeConstraint. [1, 2] is an instance of ' 'list.', e.exception.args[0]) with self.assertRaises(TypeError) as e: t = 5 @check_type_hints @with_input_types(a=t) def unused_foo(a): pass self.assertEqual( 'All type hint arguments must be a non-sequence, a type, ' 'or a TypeConstraint. 5 is an instance of int.', e.exception.args[0]) def test_basic_type_assertion(self): @check_type_hints @with_input_types(a=int) def foo(a): return a + 1 with self.assertRaises(TypeCheckError) as e: m = 'a' foo(m) self.assertEqual( "Type-hint for argument: 'a' violated. Expected an " "instance of {}, instead found an instance of " "{}.".format(int, type(m)), e.exception.args[0]) def test_composite_type_assertion(self): @check_type_hints @with_input_types(a=typehints.List[int]) def foo(a): a.append(1) return a with self.assertRaises(TypeCheckError) as e: m = ['f', 'f'] foo(m) self.assertEqual( "Type-hint for argument: 'a' violated: List[] hint " "type-constraint violated. The type of element #0 in " "the passed list is incorrect. Expected an instance of " "type int, instead received an instance of type str.", e.exception.args[0]) def test_valid_simple_type_arguments(self): @with_input_types(a=str) def upper(a): return a.upper() # Type constraints should pass, and function will be evaluated as normal. self.assertEqual('M', upper('m')) def test_any_argument_type_hint(self): @check_type_hints @with_input_types(a=typehints.Any) def foo(a): return 4 self.assertEqual(4, foo('m')) def test_valid_mix_positional_and_keyword_arguments(self): @check_type_hints @with_input_types(typehints.List[int], elem=typehints.List[int]) def combine(container, elem): return container + elem self.assertEqual([1, 2, 3], combine([1, 2], [3])) def test_invalid_only_positional_arguments(self): @check_type_hints @with_input_types(int, int) def sub(a, b): return a - b with self.assertRaises(TypeCheckError) as e: m = 'two' sub(1, m) self.assertEqual( "Type-hint for argument: 'b' violated. Expected an " "instance of {}, instead found an instance of " "{}.".format(int, type(m)), e.exception.args[0]) def test_valid_only_positional_arguments(self): @with_input_types(int, int) def add(a, b): return a + b self.assertEqual(3, add(1, 2)) class InputDecoratorTestCase(TypeHintTestCase): def test_valid_hint(self): @with_input_types(int, int) def unused_add(a, b): return a + b @with_input_types(int, b=int) def unused_add2(a, b): return a + b @with_input_types(a=int, b=int) def unused_add3(a, b): return a + b def test_invalid_kw_hint(self): with self.assertRaisesRegex(TypeError, r'\[1, 2\]'): @with_input_types(a=[1, 2]) def unused_foo(a): pass def test_invalid_pos_hint(self): with self.assertRaisesRegex(TypeError, r'\[1, 2\]'): @with_input_types([1, 2]) def unused_foo(a): pass class OutputDecoratorTestCase(TypeHintTestCase): def test_valid_hint(self): @with_output_types(int) def unused_foo(): return 5 @with_output_types(None) def unused_foo(): return 5 @with_output_types(Tuple[int, str]) def unused_foo(): return 5, 'bar' def test_no_kwargs_accepted(self): with self.assertRaisesRegex(ValueError, r'single positional argument'): @with_output_types(m=int) def unused_foo(): return 5 def test_must_be_primitive_type_or_type_constraint(self): with self.assertRaises(TypeError): @with_output_types(5) def unused_foo(): pass with self.assertRaises(TypeError): @with_output_types([1, 2]) def unused_foo(): pass def test_must_be_single_return_type(self): with self.assertRaises(ValueError): @with_output_types(int, str) def unused_foo(): return 4, 'f' def test_type_check_violation(self): @check_type_hints @with_output_types(int) def foo(a): return 'test' with self.assertRaises(TypeCheckError) as e: m = 4 foo(m) self.assertEqual( "Type-hint for return type violated. Expected an " "instance of {}, instead found an instance of " "{}.".format(int, type('test')), e.exception.args[0]) def test_type_check_simple_type(self): @check_type_hints @with_output_types(str) def upper(a): return a.upper() self.assertEqual('TEST', upper('test')) def test_type_check_composite_type(self): @check_type_hints @with_output_types(typehints.List[typehints.Tuple[int, int]]) def bar(): return [(i, i + 1) for i in range(5)] self.assertEqual([(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)], bar()) def test_any_return_type_hint(self): @check_type_hints @with_output_types(typehints.Any) def bar(): return 'foo' self.assertEqual('foo', bar()) class CombinedReturnsAndTakesTestCase(TypeHintTestCase): def test_enable_and_disable_type_checking_takes(self): @with_input_types(a=int) def int_to_str(a): return str(a) # The function call below violates the argument type-hint above, but won't # result in an exception since run-time type-checking was disabled above. self.assertEqual('a', int_to_str('a')) # Must re-define since the conditional is in the (maybe)wrapper. @check_type_hints @with_input_types(a=int) def int_to_str(a): # pylint: disable=function-redefined return str(a) # With run-time type checking enabled once again the same call-atttempt # should result in a TypeCheckError. with self.assertRaises(TypeCheckError): int_to_str('a') def test_enable_and_disable_type_checking_returns(self): @with_output_types(str) def int_to_str(a): return a # The return value of the function above violates the return-type # type-hint above, but won't result in an exception since run-time # type-checking was disabled above. self.assertEqual(9, int_to_str(9)) # Must re-define since the conditional is in the (maybe)wrapper. @check_type_hints @with_output_types(str) def int_to_str(a): # pylint: disable=function-redefined return a # With type-checking enabled once again we should get a TypeCheckError here. with self.assertRaises(TypeCheckError): int_to_str(9) def test_valid_mix_pos_and_keyword_with_both_orders(self): @with_input_types(str, start=int) @with_output_types(str) def to_upper_with_slice(string, start): return string.upper()[start:] self.assertEqual('ELLO', to_upper_with_slice('hello', 1)) def test_simple_takes_and_returns_hints(self): @check_type_hints @with_output_types(str) @with_input_types(a=str) def to_lower(a): return a.lower() # Return type and argument type satisfied, should work as normal. self.assertEqual('m', to_lower('M')) # Invalid argument type should raise a TypeCheckError with self.assertRaises(TypeCheckError): to_lower(5) @check_type_hints @with_output_types(str) @with_input_types(a=str) def to_lower(a): # pylint: disable=function-redefined return 9 # Modified function now has an invalid return type. with self.assertRaises(TypeCheckError): to_lower('a') def test_composite_takes_and_returns_hints(self): @check_type_hints @with_input_types(it=typehints.List[int]) @with_output_types(typehints.List[typehints.Tuple[int, int]]) def expand_ints(it): return [(i, i + 1) for i in it] # Return type and argument type satisfied, should work as normal. self.assertEqual([(0, 1), (1, 2), (2, 3)], expand_ints(list(range(3)))) # Invalid argument, list of str instead of int. with self.assertRaises(TypeCheckError): expand_ints('t e s t'.split()) @check_type_hints @with_output_types(typehints.List[typehints.Tuple[int, int]]) @with_input_types(it=typehints.List[int]) def expand_ints(it): # pylint: disable=function-redefined return [str(i) for i in it] # Modified function now has invalid return type. with self.assertRaises(TypeCheckError): expand_ints(list(range(2))) class DecoratorHelpers(TypeHintTestCase): def test_hint_helper(self): self.assertTrue(is_consistent_with(Any, int)) self.assertTrue(is_consistent_with(int, Any)) self.assertTrue(is_consistent_with(str, object)) self.assertFalse(is_consistent_with(object, str)) self.assertTrue(is_consistent_with(str, Union[str, int])) self.assertFalse(is_consistent_with(Union[str, int], str)) self.assertFalse(is_consistent_with(str, NonBuiltInGeneric[str])) def test_hint_helper_pipe_union(self): pipe_union = int | None # pylint: disable=unsupported-binary-operation typing_union = Union[int, None] self.assertTrue(is_consistent_with(int, pipe_union)) self.assertTrue(is_consistent_with(type(None), pipe_union)) self.assertFalse(is_consistent_with(str, pipe_union)) self.assertTrue( is_consistent_with(int, pipe_union) == is_consistent_with( int, typing_union)) self.assertTrue( is_consistent_with(str, pipe_union) == is_consistent_with( str, typing_union)) pipe_union_2 = int | float # pylint: disable=unsupported-binary-operation self.assertTrue(is_consistent_with(int, pipe_union_2)) self.assertTrue(is_consistent_with(float, pipe_union_2)) def test_is_consistent_with_disable_beartype(self): import unittest.mock from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options_context import scoped_pipeline_options with unittest.mock.patch( 'apache_beam.typehints.typehints.is_subhint') as mock_is_subhint: mock_is_subhint.return_value = True class A: pass class B(A): pass options = PipelineOptions([]) with scoped_pipeline_options(options): typehints.is_consistent_with(B, A) self.assertTrue(mock_is_subhint.called) mock_is_subhint.reset_mock() options = PipelineOptions(['--disable_beartype']) with scoped_pipeline_options(options): typehints.is_consistent_with(B, A) self.assertFalse(mock_is_subhint.called) def test_positional_arg_hints(self): self.assertEqual(typehints.Any, _positional_arg_hints('x', {})) self.assertEqual(int, _positional_arg_hints('x', {'x': int})) self.assertEqual( typehints.Tuple[int, typehints.Any], _positional_arg_hints(['x', 'y'], {'x': int})) def test_getcallargs_forhints(self): def func(a, b_c, *d): return a, b_c, d self.assertEqual({ 'a': Any, 'b_c': Any, 'd': Tuple[Any, ...] }, getcallargs_forhints(func, *[Any, Any])) self.assertEqual({ 'a': Any, 'b_c': Any, 'd': Tuple[Union[int, str], ...] }, getcallargs_forhints(func, *[Any, Any, str, int])) self.assertEqual({ 'a': int, 'b_c': Tuple[str, Any], 'd': Tuple[Any, ...] }, getcallargs_forhints(func, *[int, Tuple[str, Any]])) self.assertEqual({ 'a': Any, 'b_c': Any, 'd': Tuple[str, ...] }, getcallargs_forhints(func, *[Any, Any, Tuple[str, ...]])) self.assertEqual({ 'a': Any, 'b_c': Any, 'd': Tuple[Union[Tuple[str, ...], int], ...] }, getcallargs_forhints( func, *[Any, Any, Tuple[str, ...], int])) def test_getcallargs_forhints_builtins(self): self.assertEqual({'self': str}, getcallargs_forhints(str.upper, str)) # str.strip has an optional second argument. self.assertEqual({ 'self': str, 'chars': Any }, getcallargs_forhints(str.strip, str)) self.assertEqual({ 'self': str, 'iterable': typehints.List[int] }, getcallargs_forhints(str.join, str, typehints.List[int])) class TestGetYieldedType(unittest.TestCase): def test_iterables(self): self.assertEqual(int, typehints.get_yielded_type(typehints.Iterable[int])) self.assertEqual(int, typehints.get_yielded_type(typehints.Iterator[int])) self.assertEqual(int, typehints.get_yielded_type(typehints.Generator[int])) self.assertEqual(int, typehints.get_yielded_type(typehints.List[int])) self.assertEqual( typehints.Union[int, str], typehints.get_yielded_type(typehints.Tuple[int, str])) self.assertEqual(int, typehints.get_yielded_type(typehints.Set[int])) self.assertEqual(int, typehints.get_yielded_type(typehints.FrozenSet[int])) self.assertEqual(typing.Any, typehints.get_yielded_type(T)) self.assertEqual( typehints.Union[int, str], typehints.get_yielded_type( typehints.Union[typehints.List[int], typehints.List[str]])) def test_not_iterable(self): with self.assertRaisesRegex(ValueError, r'not iterable'): typehints.get_yielded_type(int) def test_union_not_iterable(self): with self.assertRaisesRegex(ValueError, r'not iterable'): typehints.get_yielded_type(typehints.Union[int, typehints.List[int]]) class TestCoerceToKvType(TypeHintTestCase): def test_coercion_success(self): cases = [ ((Any, ), typehints.KV[Any, Any]), ((typehints.KV[Any, Any], ), typehints.KV[Any, Any]), ((typehints.Tuple[str, int], ), typehints.KV[str, int]), ] for args, expected in cases: self.assertEqual(typehints.coerce_to_kv_type(*args), expected) self.assertCompatible(args[0], expected) def test_coercion_fail(self): cases = [ ((str, 'label', 'producer'), r'producer.*compatible'), ((Tuple[str], ), r'two components'), # It seems that the only Unions that may be successfully coerced are not # Unions but Any (e.g. Union[Any, Tuple[Any, Any]] is Any). ((Union[str, int], ), r'compatible'), ((Union, ), r'compatible'), ((typehints.List[Any], ), r'compatible'), ] for args, regex in cases: with self.assertRaisesRegex(ValueError, regex): typehints.coerce_to_kv_type(*args) class TestParDoAnnotations(unittest.TestCase): def test_with_side_input(self): class MyDoFn(DoFn): def process(self, element: float, side_input: str) -> \ Iterable[KV[str, float]]: pass th = MyDoFn().get_type_hints() self.assertEqual(th.input_types, ((float, str), {})) self.assertEqual(th.output_types, ((KV[str, float], ), {})) def test_pep484_annotations(self): class MyDoFn(DoFn): def process(self, element: int) -> Iterable[str]: pass th = MyDoFn().get_type_hints() self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((str, ), {})) class TestPTransformAnnotations(unittest.TestCase): def test_pep484_annotations(self): class MyPTransform(PTransform): def expand(self, pcoll: PCollection[int]) -> PCollection[str]: return pcoll | Map(lambda num: str(num)) th = MyPTransform().get_type_hints() self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((str, ), {})) def test_annotations_without_input_pcollection_wrapper(self): class MyPTransform(PTransform): def expand(self, pcoll: int) -> PCollection[str]: return pcoll | Map(lambda num: str(num)) error_str = ( r'This input type hint will be ignored and not used for ' r'type-checking purposes. Typically, input type hints for a ' r'PTransform are single (or nested) types wrapped by a ' r'PCollection, or PBegin. Got: {} instead.'.format(int)) with self.assertLogs(level='WARN') as log: MyPTransform().get_type_hints() self.assertIn(error_str, log.output[0]) def test_annotations_without_output_pcollection_wrapper(self): class MyPTransform(PTransform): def expand(self, pcoll: PCollection[int]) -> str: return pcoll | Map(lambda num: str(num)) error_str = ( r'This output type hint will be ignored and not used for ' r'type-checking purposes. Typically, output type hints for a ' r'PTransform are single (or nested) types wrapped by a ' r'PCollection, PDone, or None. Got: {} instead.'.format(str)) with self.assertLogs(level='WARN') as log: th = MyPTransform().get_type_hints() self.assertIn(error_str, log.output[0]) self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, None) def test_annotations_without_input_internal_type(self): class MyPTransform(PTransform): def expand(self, pcoll: PCollection) -> PCollection[str]: return pcoll | Map(lambda num: str(num)) th = MyPTransform().get_type_hints() self.assertEqual(th.input_types, ((Any, ), {})) self.assertEqual(th.output_types, ((str, ), {})) def test_annotations_without_output_internal_type(self): class MyPTransform(PTransform): def expand(self, pcoll: PCollection[int]) -> PCollection: return pcoll | Map(lambda num: str(num)) th = MyPTransform().get_type_hints() self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((Any, ), {})) def test_annotations_without_any_internal_type(self): class MyPTransform(PTransform): def expand(self, pcoll: PCollection) -> PCollection: return pcoll | Map(lambda num: str(num)) th = MyPTransform().get_type_hints() self.assertEqual(th.input_types, ((Any, ), {})) self.assertEqual(th.output_types, ((Any, ), {})) def test_annotations_without_input_typehint(self): class MyPTransform(PTransform): def expand(self, pcoll) -> PCollection[str]: return pcoll | Map(lambda num: str(num)) th = MyPTransform().get_type_hints() self.assertEqual(th.input_types, ((Any, ), {})) self.assertEqual(th.output_types, ((str, ), {})) def test_annotations_without_output_typehint(self): class MyPTransform(PTransform): def expand(self, pcoll: PCollection[int]): return pcoll | Map(lambda num: str(num)) th = MyPTransform().get_type_hints() self.assertEqual(th.input_types, ((int, ), {})) self.assertEqual(th.output_types, ((Any, ), {})) def test_annotations_without_any_typehints(self): class MyPTransform(PTransform): def expand(self, pcoll): return pcoll | Map(lambda num: str(num)) th = MyPTransform().get_type_hints() self.assertEqual(th.input_types, None) self.assertEqual(th.output_types, None) def test_annotations_with_pbegin(self): class MyPTransform(PTransform): def expand(self, pcoll: PBegin): return pcoll | Map(lambda num: str(num)) th = MyPTransform().get_type_hints() self.assertEqual(th.input_types, ((Any, ), {})) self.assertEqual(th.output_types, ((Any, ), {})) def test_annotations_with_pdone(self): class MyPTransform(PTransform): def expand(self, pcoll) -> PDone: return pcoll | Map(lambda num: str(num)) th = MyPTransform().get_type_hints() self.assertEqual(th.input_types, ((Any, ), {})) self.assertEqual(th.output_types, ((Any, ), {})) def test_annotations_with_none_input(self): class MyPTransform(PTransform): def expand(self, pcoll: None) -> PCollection[str]: return pcoll | Map(lambda num: str(num)) error_str = ( r'This input type hint will be ignored and not used for ' r'type-checking purposes. Typically, input type hints for a ' r'PTransform are single (or nested) types wrapped by a ' r'PCollection, or PBegin. Got: {} instead.'.format(None)) with self.assertLogs(level='WARN') as log: th = MyPTransform().get_type_hints() self.assertIn(error_str, log.output[0]) self.assertEqual(th.input_types, None) self.assertEqual(th.output_types, ((str, ), {})) def test_annotations_with_none_output(self): class MyPTransform(PTransform): def expand(self, pcoll) -> None: return pcoll | Map(lambda num: str(num)) th = MyPTransform().get_type_hints() self.assertEqual(th.input_types, ((Any, ), {})) self.assertEqual(th.output_types, ((Any, ), {})) def test_annotations_with_arbitrary_output(self): class MyPTransform(PTransform): def expand(self, pcoll) -> str: return pcoll | Map(lambda num: str(num)) th = MyPTransform().get_type_hints() self.assertEqual(th.input_types, ((Any, ), {})) self.assertEqual(th.output_types, None) def test_annotations_with_arbitrary_input_and_output(self): class MyPTransform(PTransform): def expand(self, pcoll: int) -> str: return pcoll | Map(lambda num: str(num)) input_error_str = ( r'This input type hint will be ignored and not used for ' r'type-checking purposes. Typically, input type hints for a ' r'PTransform are single (or nested) types wrapped by a ' r'PCollection, or PBegin. Got: {} instead.'.format(int)) output_error_str = ( r'This output type hint will be ignored and not used for ' r'type-checking purposes. Typically, output type hints for a ' r'PTransform are single (or nested) types wrapped by a ' r'PCollection, PDone, or None. Got: {} instead.'.format(str)) with self.assertLogs(level='WARN') as log: th = MyPTransform().get_type_hints() self.assertIn(input_error_str, log.output[0]) self.assertIn(output_error_str, log.output[1]) self.assertEqual(th.input_types, None) self.assertEqual(th.output_types, None) def test_typing_module_annotations_are_converted_to_beam_annotations(self): class MyPTransform(PTransform): def expand( self, pcoll: PCollection[typing.Dict[str, str]] ) -> PCollection[typing.Dict[str, str]]: return pcoll th = MyPTransform().get_type_hints() self.assertEqual(th.input_types, ((typehints.Dict[str, str], ), {})) self.assertEqual(th.input_types, ((typehints.Dict[str, str], ), {})) def test_nested_typing_annotations_are_converted_to_beam_annotations(self): class MyPTransform(PTransform): def expand(self, pcoll: PCollection[typing.Union[int, typing.Any, typing.Dict[str, float]]]) \ -> PCollection[typing.Union[int, typing.Any, typing.Dict[str, float]]]: return pcoll th = MyPTransform().get_type_hints() self.assertEqual( th.input_types, ((typehints.Union[int, typehints.Any, typehints.Dict[str, float]], ), {})) self.assertEqual( th.input_types, ((typehints.Union[int, typehints.Any, typehints.Dict[str, float]], ), {})) def test_mixed_annotations_are_converted_to_beam_annotations(self): class MyPTransform(PTransform): def expand(self, pcoll: typing.Any) -> typehints.Any: return pcoll th = MyPTransform().get_type_hints() self.assertEqual(th.input_types, ((typehints.Any, ), {})) self.assertEqual(th.input_types, ((typehints.Any, ), {})) def test_pipe_operator_as_union(self): # union types can be written using pipe operator from Python 3.10. # https://peps.python.org/pep-0604/ type_a = int | float # pylint: disable=unsupported-binary-operation type_b = typing.Union[int, float] self.assertEqual( native_type_compatibility.convert_to_beam_type(type_a), native_type_compatibility.convert_to_beam_type(type_b)) def test_normalize_pipe_union(self): pipe_union = int | None # pylint: disable=unsupported-binary-operation normalized = typehints.normalize(pipe_union) self.assertIsInstance(normalized, typehints.UnionConstraint) pipe_union_2 = int | float # pylint: disable=unsupported-binary-operation normalized_2 = typehints.normalize(pipe_union_2) self.assertIsInstance(normalized_2, typehints.UnionConstraint) class TestNonBuiltInGenerics(unittest.TestCase): def test_no_error_thrown(self): input = NonBuiltInGeneric[str] output = typehints.normalize(input) self.assertEqual(input, output) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.790451 apache_beam-2.74.0/apache_beam/utils/0000755000175100017510000000000015206113160017017 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/__init__.py0000644000175100017510000000160015206113033021124 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A package containing internal utilities. For internal use only; no backwards-compatibility guarantees. """ ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/annotations.py0000644000175100017510000001450115206113033021726 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Deprecated annotations. Deprecated: Signifies that users are discouraged from using a public API typically because a better alternative exists, and the current form might be removed in a future version. Usage: For internal use only; no backwards-compatibility guarantees. Annotations come in two flavors: deprecated and experimental The 'deprecated' annotation requires a 'since" parameter to specify what version deprecated it. Both 'deprecated' and 'experimental' annotations can specify the current recommended version to use by means of a 'current' parameter. The following example illustrates how to annotate coexisting versions of the same function 'multiply'.:: def multiply(arg1, arg2): print(arg1, '*', arg2, '=', end=' ') return arg1*arg2 # This annotation marks 'old_multiply' as deprecated since 'v.1' and suggests # using 'multiply' instead.:: @deprecated(since='v.1', current='multiply') def old_multiply(arg1, arg2): result = 0 for i in xrange(arg1): result += arg2 print(arg1, '*', arg2, '(the old way)=', end=' ') return result # Set a warning filter to control how often warnings are produced.:: warnings.simplefilter("always") print(multiply(5, 6)) print(old_multiply(5,6)) """ # pytype: skip-file import inspect import warnings from functools import partial from functools import wraps def _add_deprecation_notice_to_docstring(docstring, message): """Adds a deprecation notice to a docstring. Args: docstring: The original docstring (can be None or empty). message: The deprecation message to add. Returns: The modified docstring. """ if docstring: return f"{docstring}\n\n.. deprecated:: {message}" else: return f".. deprecated:: {message}" class BeamDeprecationWarning(DeprecationWarning): """Beam-specific deprecation warnings.""" # Don't ignore BeamDeprecationWarnings. warnings.simplefilter('once', BeamDeprecationWarning) class _WarningMessage: """Utility class for assembling the warning message.""" def __init__(self, label, since, current, extra_message, custom_message): """Initialize message, leave only name as placeholder.""" if custom_message is None: message = '%name% is ' + label if label == 'deprecated': message += ' since %s' % since message += '. Use %s instead.' % current if current else '.' if extra_message: message += ' ' + extra_message else: if label == 'deprecated' and '%since%' not in custom_message: raise TypeError( "Replacement string %since% not found on \ custom message") emptyArg = lambda x: '' if x is None else x message = custom_message\ .replace('%since%', emptyArg(since))\ .replace('%current%', emptyArg(current))\ .replace('%extra%', emptyArg(extra_message)) self.label = label self.message = message def emit_warning(self, fnc_name): if self.label == 'deprecated': warning_type = BeamDeprecationWarning else: warning_type = FutureWarning warnings.warn( self.message.replace('%name%', fnc_name), warning_type, stacklevel=3) def annotate(label, since, current, extra_message, custom_message=None): """Decorates an API with a deprecated or experimental annotation. Args: label: the kind of annotation ('deprecated' or 'experimental'). since: the version that causes the annotation. current: the suggested replacement function. extra_message: an optional additional message. custom_message: if the default message does not suffice, the message can be changed using this argument. A string whit replacement tokens. A replecement string is were the previus args will be located on the custom message. The following replacement strings can be used: %name% -> API.__name__ %since% -> since (Mandatory for the decapreted annotation) %current% -> current %extra% -> extra_message Returns: The decorator for the API. """ warning_message = _WarningMessage( label=label, since=since, current=current, extra_message=extra_message, custom_message=custom_message) def _annotate(fnc): if inspect.isclass(fnc): # Wrapping class into function causes documentation rendering issue. # Patch class's __new__ method instead. old_new = fnc.__new__ def wrapped_new(cls, *args, **kwargs): # Emit a warning when class instance is created. warning_message.emit_warning(fnc.__name__) if old_new is object.__new__: # object.__new__ takes no extra argument for python>=3 return old_new(cls) return old_new(cls, *args, **kwargs) fnc.__new__ = staticmethod(wrapped_new) if label == 'deprecated': fnc.__doc__ = _add_deprecation_notice_to_docstring( fnc.__doc__, warning_message.message.replace('%name%', fnc.__name__)) return fnc else: @wraps(fnc) def inner(*args, **kwargs): # Emit a warning when the function is called. warning_message.emit_warning(fnc.__name__) return fnc(*args, **kwargs) if label == 'deprecated': inner.__doc__ = _add_deprecation_notice_to_docstring( fnc.__doc__, warning_message.message.replace('%name%', fnc.__name__)) return inner return _annotate # Use partial application to customize each annotation. # 'current' will be optional in both deprecated and experimental # while 'since' will be mandatory for deprecated. deprecated = partial( annotate, label='deprecated', current=None, extra_message=None) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/annotations_test.py0000644000175100017510000002032715206113033022770 0ustar00runnerrunner# Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import unittest import warnings from apache_beam.utils.annotations import BeamDeprecationWarning from apache_beam.utils.annotations import deprecated class AnnotationTests(unittest.TestCase): # Note: use different names for each of the functions decorated # so that a warning is produced for each of them. def setUp(self): # This is necessary to ensure warnings are captured in all # platforms (Windows/Unix). warnings.resetwarnings() def test_deprecated_with_since_current_message(self): with warnings.catch_warnings(record=True) as w: @deprecated(since='v.1', current='multiply', extra_message='Do this') def fnc_test_deprecated_with_since_current_message(): return 'lol' fnc_test_deprecated_with_since_current_message() self.check_annotation( warning=w, warning_type=BeamDeprecationWarning, obj_name='fnc_test_deprecated_with_since_current_message', annotation_type='deprecated', label_check_list=[('since', True), ('instead', True), ('Do this', True)]) self.assertIn( '.. deprecated::', fnc_test_deprecated_with_since_current_message.__doc__) def test_deprecated_with_class(self): @deprecated(since='v.1', current='multiply', extra_message='Do this') class MyClass(): """This is a test.""" def __init__(self): pass def foo(self): pass self.assertIn('.. deprecated::', MyClass.__doc__) self.assertIn('This is a test', MyClass.__doc__) def test_deprecated_with_since_current(self): with warnings.catch_warnings(record=True) as w: @deprecated(since='v.1', current='multiply') def fnc_test_deprecated_with_since_current(): return 'lol' fnc_test_deprecated_with_since_current() self.check_annotation( warning=w, warning_type=BeamDeprecationWarning, obj_name='fnc_test_deprecated_with_since_current', annotation_type='deprecated', label_check_list=[('since', True), ('instead', True)]) def test_deprecated_without_current(self): with warnings.catch_warnings(record=True) as w: @deprecated(since='v.1') def fnc_test_deprecated_without_current(): return 'lol' fnc_test_deprecated_without_current() self.check_annotation( warning=w, warning_type=BeamDeprecationWarning, obj_name='fnc_test_deprecated_without_current', annotation_type='deprecated', label_check_list=[('since', True), ('instead', False)]) def test_deprecated_without_since_should_fail(self): with warnings.catch_warnings(record=True) as w: with self.assertRaises(TypeError): @deprecated() def fnc_test_deprecated_without_since_should_fail(): return 'lol' fnc_test_deprecated_without_since_should_fail() assert not w def test_deprecated_without_since_custom_should_fail(self): with warnings.catch_warnings(record=True) as w: with self.assertRaises(TypeError): @deprecated(custom_message='Test %since%') def fnc_test_deprecated_without_since_custom_should_fail(): return 'lol' fnc_test_deprecated_without_since_custom_should_fail() assert not w def test_deprecated_custom_no_replacements(self): """Tests if custom message prints an empty string for each replacement token when only the custom_message and since parameter are given.""" with warnings.catch_warnings(record=True) as w: strSince = 'v1' strCustom = 'Replacement:%since%%current%%extra%' @deprecated(since=strSince, custom_message=strCustom) def fnc_test_experimental_custom_no_replacements(): return 'lol' fnc_test_experimental_custom_no_replacements() self.check_custom_annotation( warning=w, warning_type=BeamDeprecationWarning, obj_name='fnc_test_experimental_custom_no_\ replacements', annotation_type='experimental', intended_message=strCustom.replace('%since%', strSince).replace( '%current%', '').replace('%extra%', '')) def test_enforce_custom_since_deprecated_must_fail(self): """Tests since replacement token inclusion on the custom message for the decapreted string. If no since replacement token is given, the annotation must fail""" with warnings.catch_warnings(record=True) as w: with self.assertRaises(TypeError): strSince = 'v1' strCustom = 'Replacement:' @deprecated(since=strSince, custom_message=strCustom) def fnc_test_experimental_custom_no_replacements(): return 'lol' fnc_test_experimental_custom_no_replacements() assert not w def test_deprecated_with_since_current_message_custom(self): with warnings.catch_warnings(record=True) as w: strSince = 'v.1' strCurrent = 'multiply' strExtra = 'Do this' strCustom = "%name% Will be deprecated from %since%. \ Please use %current% insted. Will %extra%" @deprecated( since=strSince, current=strCurrent, extra_message=strExtra, custom_message=strCustom) def fnc_test_deprecated_with_since_current_message_custom(): return 'lol' strName = fnc_test_deprecated_with_since_current_message_custom.__name__ fnc_test_deprecated_with_since_current_message_custom() self.check_custom_annotation( warning=w, warning_type=BeamDeprecationWarning, obj_name='fnc_test_deprecated_with_since_\ current_message_custom', annotation_type='deprecated', intended_message=strCustom.replace('%name%', strName).replace( '%since%', strSince).replace('%current%', strCurrent).replace('%extra%', strExtra)) def test_deprecated_with_since_current_message_class(self): with warnings.catch_warnings(record=True) as w: @deprecated(since='v.1', current='multiply', extra_message='Do this') class Class_test_deprecated_with_since_current_message(object): fooo = 'lol' def __init__(self): pass def foo(self): return 'lol' foo = Class_test_deprecated_with_since_current_message() strName = Class_test_deprecated_with_since_current_message.__name__ foo.foo() self.check_annotation( warning=w, warning_type=BeamDeprecationWarning, obj_name=strName, annotation_type='deprecated', label_check_list=[('since', True), ('instead', True), ('Do this', True)]) # helper function def check_annotation( self, warning, warning_type, obj_name, annotation_type, label_check_list): self.assertTrue(issubclass(warning[-1].category, warning_type)) self.assertIn(obj_name + ' is ' + annotation_type, str(warning[-1].message)) for label in label_check_list: if label[1] is True: self.assertIn(label[0], str(warning[-1].message)) else: self.assertNotIn(label[0], str(warning[-1].message)) # Helper function for custom messages def check_custom_annotation( self, warning, warning_type, obj_name, annotation_type, intended_message): self.assertTrue(issubclass(warning[-1].category, warning_type)) self.assertIn(intended_message, str(warning[-1].message)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/byte_limited_queue.pxd0000644000175100017510000000233215206113033023411 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # cython: overflowcheck=True cdef class ByteLimitedQueue(object): cdef readonly Py_ssize_t max_elements cdef readonly Py_ssize_t max_bytes cdef readonly Py_ssize_t _byte_size cdef readonly object _mutex cdef readonly object _not_empty cdef readonly object _waiting_writers cdef readonly list _condition_pool cdef readonly object _queue cdef readonly Py_ssize_t _blocked_bytes cpdef bint _can_fit(self, Py_ssize_t item_bytes) except -1 ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/byte_limited_queue.py0000644000175100017510000001474615206113033023262 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A thread-safe queue that limits capacity by total byte size.""" import collections import queue import threading import time import types class ByteLimitedQueue(object): """A fair queue that limits by both element count and total byte size. A single element is allowed to exceed the maxbytes to avoid deadlock. """ __class_getitem__ = classmethod(types.GenericAlias) def __init__( self, maxsize=0, # type: int maxbytes=0, # type: int ): # type: (...) -> None """Initializes a ByteLimitedQueue. Args: maxsize: The maximum number of items allowed in the queue. If 0 or negative, there is no limit on the number of elements. maxbytes: The maximum accumulated bytes allowed in the queue. If 0 or negative, there is no limit on the total bytes of the elements. """ self.max_elements = maxsize self.max_bytes = maxbytes self._byte_size = 0 self._blocked_bytes = 0 self._mutex = threading.Lock() self._not_empty = threading.Condition(self._mutex) self._waiting_writers = collections.deque() self._condition_pool = [] self._queue = collections.deque() def put(self, item, item_bytes, *, block=True, timeout=None): """Put an item into the queue. If the queue is full, block until a free slot is available, unless `block` is false or a timeout occurs. Args: item: The item to put into the queue. item_bytes: The size of the item. block: If True, block until space is available. If False, raise queue.Full immediately if the queue is full. timeout: If block is True, wait for at most `timeout` seconds. If None, block indefinitely. Raises: ValueError: If timeout or item_bytes is negative. queue.Full: If the queue is full and block is False or the timeout occurs. """ if timeout is not None and timeout < 0: raise ValueError("'timeout' must be a non-negative number") if item_bytes < 0: raise ValueError("'item_bytes' must be a non-negative number") with self._mutex: if not self._waiting_writers and self._can_fit(item_bytes): self._queue.append((item, item_bytes)) self._byte_size += item_bytes self._not_empty.notify() return if not block: raise queue.Full # Reuse or create a condition my_cond = ( self._condition_pool.pop() if self._condition_pool else threading.Condition(self._mutex)) endtime = time.monotonic() + timeout if timeout is not None else None try: self._blocked_bytes += item_bytes self._waiting_writers.append(my_cond) while True: if timeout is None: my_cond.wait() else: remaining = endtime - time.monotonic() if remaining <= 0.0: raise queue.Full my_cond.wait(remaining) if self._waiting_writers[0] is my_cond and self._can_fit(item_bytes): break self._queue.append((item, item_bytes)) self._byte_size += item_bytes self._not_empty.notify() finally: self._blocked_bytes -= item_bytes if self._waiting_writers: was_first = (self._waiting_writers[0] is my_cond) if was_first: self._waiting_writers.popleft() else: self._waiting_writers.remove(my_cond) self._condition_pool.append(my_cond) if was_first and self._waiting_writers: self._waiting_writers[0].notify() def get(self, *, block=True, timeout=None): """Remove and return an item from the queue. If the queue is empty, block until an item is available, unless `block` is false or a timeout occurs. Args: block: If True, block until an item is available. If False, raise queue.Empty immediately if the queue is empty. timeout: If block is True, wait for at most `timeout` seconds. If None, block indefinitely. Returns: The item removed from the queue. Raises: ValueError: If timeout is negative. queue.Empty: If the queue is empty and block is False or the timeout occurs. """ if timeout is not None and timeout < 0: raise ValueError("'timeout' must be a non-negative number") with self._mutex: if not block: if not self._queue: raise queue.Empty elif timeout is None: while not self._queue: self._not_empty.wait() else: endtime = time.monotonic() + timeout while not self._queue: remaining = endtime - time.monotonic() if remaining <= 0.0: raise queue.Empty self._not_empty.wait(remaining) item, item_bytes = self._queue.popleft() self._byte_size -= item_bytes if self._waiting_writers: self._waiting_writers[0].notify() return item def get_nowait(self): """Remove and return an item from the queue without blocking.""" return self.get(block=False) def byte_size(self): """Return the total byte size of elements in the queue.""" with self._mutex: return self._byte_size def blocked_byte_size(self): """Return the total byte size of elements in the queue that are blocked.""" with self._mutex: return self._blocked_bytes def qsize(self): """Return the total number of elements in the queue.""" with self._mutex: return len(self._queue) def _can_fit(self, item_bytes): # Always let in a single element, regardless of size. if not self._queue: return True if self.max_elements > 0 and len(self._queue) >= self.max_elements: return False if self.max_bytes > 0 and self._byte_size + item_bytes > self.max_bytes: return False return True ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/byte_limited_queue_test.py0000644000175100017510000001757315206113033024322 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for byte-limited queue.""" import queue import threading import time import unittest from apache_beam.utils.byte_limited_queue import ByteLimitedQueue class ByteLimitedQueueTest(unittest.TestCase): def test_unbounded(self): bq = ByteLimitedQueue() for i in range(201): bq.put(str(i), i) self.assertEqual(bq.byte_size(), sum(range(201))) self.assertEqual(bq.qsize(), 201) def test_put_and_get(self): bq = ByteLimitedQueue(maxbytes=200) bq.put('50', 50) bq.put('140', 140) self.assertEqual(bq.byte_size(), 190) self.assertEqual(bq.qsize(), 2) # Putting another would exceed 200. with self.assertRaises(queue.Full): bq.put('20', 20, block=False) bq.put('10', 10, block=False) self.assertEqual(bq.byte_size(), 200) self.assertEqual(bq.qsize(), 3) self.assertEqual(bq.get(), '50') self.assertEqual(bq.byte_size(), 150) self.assertEqual(bq.qsize(), 2) bq.put('20', 20, block=False) def test_dual_limit(self): # Queue limits: at most 3 items, OR at most 100 item bytes. bq = ByteLimitedQueue(maxsize=3, maxbytes=100) bq.put('30', 30) bq.put('40', 40) bq.put('20', 20) self.assertEqual(bq.byte_size(), 90) self.assertEqual(bq.qsize(), 3) # Full on element count (size=3). with self.assertRaises(queue.Full): bq.put('10', 10, block=False) self.assertEqual(bq.get(), '30') self.assertEqual(bq.get(), '40') bq.put('10', 10) # Full on byte count with self.assertRaises(queue.Full): bq.put('90', 90, block=False) self.assertEqual(bq.get(), '20') bq.put('90', 90, block=False) def test_multithreading(self): bq = ByteLimitedQueue(maxsize=0, maxbytes=100) received = [] def producer(): for i in range(101): bq.put(str(i), i) poison_pill = 'POISON' def consumer(): while True: item = bq.get() if item == poison_pill: break received.append(int(item)) t1 = threading.Thread(target=producer) t2 = threading.Thread(target=producer) t3 = threading.Thread(target=consumer) t1.start() t2.start() t3.start() t1.join() t2.join() bq.put(poison_pill, 0) t3.join() self.assertEqual(len(received), 202) self.assertEqual(sum(received), 2 * sum(range(101))) def test_put_timeout(self): bq = ByteLimitedQueue(maxsize=0, maxbytes=10) bq.put('10', 10) # The queue is completely full. A timeout put should raise queue.Full. with self.assertRaises(queue.Full): bq.put('5', 5, timeout=0.01) def delayed_consumer(): time.sleep(0.05) bq.get() # Start a thread that will free up space after 50ms. t = threading.Thread(target=delayed_consumer) t.start() # The put should succeed once the consumer runs, use a high timeout to # flakiness. bq.put('item', 5, timeout=60) t.join() def test_get_timeout(self): bq = ByteLimitedQueue(maxsize=0, maxbytes=100) with self.assertRaises(queue.Empty): bq.get(block=False) with self.assertRaises(queue.Empty): bq.get(timeout=0.0) with self.assertRaises(queue.Empty): bq.get(timeout=.01) bq.put('1', 1) self.assertEqual('1', bq.get(timeout=0)) bq.put('2', 2) self.assertEqual('2', bq.get(timeout=0.1)) def delayed_producer(): time.sleep(0.05) bq.put('3', 3) # Start a thread that will produce soon t = threading.Thread(target=delayed_producer) t.start() # The get should succeed once the produer runs, use a high timeout to # flakiness. self.assertEqual('3', bq.get(timeout=60)) t.join() def test_negative_timeout(self): bq = ByteLimitedQueue() # Putting an item with a negative timeout should raise ValueError. with self.assertRaises(ValueError): bq.put('5', 5, timeout=-1) with self.assertRaises(ValueError): bq.get(timeout=-1) def test_single_element_override(self): bq = ByteLimitedQueue(maxbytes=10) # An item of size 50 exceeds maxbytes 10, but should be admitted # immediately without blocking since the queue is currently empty! bq.put('50', 50, block=False) self.assertEqual(bq.qsize(), 1) self.assertEqual(bq.byte_size(), 50) def test_fairness(self): bq = ByteLimitedQueue(maxbytes=10) # Put an initial item so that the queue is not empty, # causing the subsequent large item to block. bq.put('first', 2) self.assertEqual(bq.blocked_byte_size(), 0) def producer(item, size): bq.put(item, size) # Add an item in a background thread that should block due to exceeding # the limit. t1 = threading.Thread(target=producer, args=('too_large', 9)) t1.start() # Wait until the background write is queued. while bq.blocked_byte_size() < 1: time.sleep(0.005) self.assertEqual(bq.blocked_byte_size(), 9) # Add smaller items afterwards. t2 = threading.Thread(target=producer, args=('small1', 1)) t2.start() while bq.blocked_byte_size() < 10: time.sleep(0.005) self.assertEqual(bq.blocked_byte_size(), 10) t3 = threading.Thread(target=producer, args=('small2', 1)) t3.start() while bq.blocked_byte_size() < 11: time.sleep(0.005) self.assertEqual(bq.blocked_byte_size(), 11) # Verify all items are received in order. self.assertEqual(bq.get(), 'first') t1.join() t2.join() self.assertEqual(bq.get(), 'too_large') t3.join() self.assertEqual(bq.get(), 'small1') self.assertEqual(bq.get(), 'small2') def test_blocked_waiter_timeout_multiple(self): bq = ByteLimitedQueue(maxbytes=10) bq.put('initial', 5) status = [] lock = threading.Lock() def producer(name, size, timeout_val): try: bq.put(name, size, timeout=timeout_val) with lock: status.append((name, 'success')) except queue.Full: with lock: status.append((name, 'timeout')) threads = [] threads.append(threading.Thread(target=producer, args=('t1', 8, 0.2))) threads.append(threading.Thread(target=producer, args=('t2', 8, 60.0))) threads.append(threading.Thread(target=producer, args=('t3', 3, 0.1))) threads.append(threading.Thread(target=producer, args=('t4', 3, 60.0))) threads.append(threading.Thread(target=producer, args=('t5', 3, 0.1))) for t in threads: t.start() # Wait for the short-timeout threads. threads[4].join() threads[2].join() threads[0].join() # Now waiting writers should just be t1 and t3 self.assertEqual(bq.blocked_byte_size(), 11) self.assertEqual(bq.get(), 'initial') threads[1].join() self.assertGreater(bq.blocked_byte_size(), 0) elem = bq.get() self.assertTrue(elem == 't2' or elem == 't4') threads[3].join() self.assertEqual(bq.blocked_byte_size(), 0) elem = bq.get() self.assertTrue(elem == 't2' or elem == 't4') with lock: self.assertIn(('t1', 'timeout'), status) self.assertIn(('t2', 'success'), status) self.assertIn(('t3', 'timeout'), status) self.assertIn(('t4', 'success'), status) self.assertIn(('t5', 'timeout'), status) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/counters.pxd0000644000175100017510000000217415206113033021401 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # cython: overflowcheck=True cdef class Counter(object): cdef readonly object name cdef readonly object combine_fn cdef readonly object accumulator cdef readonly object _add_input cpdef bint update(self, value) except -1 cdef class AccumulatorCombineFnCounter(Counter): cdef readonly object _fast_add_input cdef readonly object _fast_add_input_n ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/counters.py0000644000175100017510000002116515206113033021237 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # cython: profile=False # cython: overflowcheck=True """Counters collect the progress of the Worker for reporting to the service. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file # ruff: noqa: UP006 import threading from collections import namedtuple from typing import TYPE_CHECKING from typing import Dict from apache_beam.transforms import cy_combiners if TYPE_CHECKING: from apache_beam.transforms import core # Information identifying the IO being measured by a counter. # # A CounterName with IOTarget helps identify the IO being measured by a # counter. # # It may represent the consumption of Shuffle IO, or the consumption of # side inputs. The way in which each is represented is explained in the # documentation of the side_input_id, and shuffle_id functions. IOTargetName = namedtuple( 'IOTargetName', ['requesting_step_name', 'input_index']) def side_input_id(step_name, input_index): # type: (str, int) -> IOTargetName """Create an IOTargetName that identifies the reading of a side input. Given a step "s4" that receives two side inputs, then the CounterName that represents the consumption of side input number 2 is: * step_name: s4 <---| * input_index: 2 <---|-- Identifying the side input itself * requesting_step_name: s4 <-- Identifying the step that reads from it. If "s4" emits the whole AsIter of the side input, down to a step, say "s5", then the requesting_step_name of the subsequent consumption will be "s5". """ return IOTargetName(step_name, input_index) def shuffle_id(step_name): # type: (str) -> IOTargetName """Create an IOTargetName that identifies a GBK step. Given a step "s6" that is downstream from a GBK "s5", then "s6" will read from shuffle. The CounterName that quantifies the consumption of data from shuffle has: * step_name: s5 * requesting_step_name: s6 If "s6" emits the whole iterable down to a step, say "s7", and "s7" continues to consume data from the iterable, then a new CounterName will be: * step_name: s5 <--- Identifying the GBK * requesting_step_name: s6 """ return IOTargetName(step_name, None) _CounterName = namedtuple( '_CounterName', [ 'name', 'stage_name', 'step_name', 'system_name', 'namespace', 'origin', 'output_index', 'io_target' ]) class CounterName(_CounterName): """Naming information for a counter.""" SYSTEM = object() USER = object() def __new__( cls, name, stage_name=None, step_name=None, system_name=None, namespace=None, origin=None, output_index=None, io_target=None): origin = origin or CounterName.SYSTEM return super().__new__( cls, name, stage_name, step_name, system_name, namespace, origin, output_index, io_target) def __repr__(self): return ' at %s>' % (self._str_internal(), hex(id(self))) def __str__(self): return self._str_internal() def _str_internal(self): if self.origin == CounterName.USER: return 'user-%s-%s' % (self.step_name, self.name) elif self.origin == CounterName.SYSTEM and self.output_index: return '%s-out%s-%s' % (self.step_name, self.output_index, self.name) else: return '%s-%s-%s' % (self.stage_name, self.step_name, self.name) class Counter(object): """A counter aggregates a series of values. The aggregation kind of the Counter is specified when the Counter is created. The values aggregated must be of an appropriate for the aggregation used. Aggregations supported are listed in the code. (The aggregated value will be reported to the Dataflow service.) Do not create directly; call CounterFactory.get_counter instead. Attributes: name: the name of the counter, a string combine_fn: the CombineFn to use for aggregation accumulator: the accumulator created for the combine_fn """ # Handy references to common counters. SUM = cy_combiners.SumInt64Fn() MEAN = cy_combiners.MeanInt64Fn() BEAM_DISTRIBUTION = cy_combiners.DistributionInt64Fn() # Dataflow Distribution Accumulator Fn. # TODO(https://github.com/apache/beam/issues/18843): Generalize distribution # counter if necessary. DATAFLOW_DISTRIBUTION = cy_combiners.DataflowDistributionCounterFn() def __init__(self, name, combine_fn): # type: (CounterName, core.CombineFn) -> None """Creates a Counter object. Args: name: the name of this counter. It may be a string, or a CounterName object. combine_fn: the CombineFn to use for aggregation """ self.name = name self.combine_fn = combine_fn self.accumulator = combine_fn.create_accumulator() self._add_input = self.combine_fn.add_input def update(self, value): self.accumulator = self._add_input(self.accumulator, value) def update_n(self, value, n): """Update the counter with the same value N times""" for _ in range(n): self.accumulator = self._add_input(self, value) def reset(self, value): self.accumulator = self.combine_fn.create_accumulator() def value(self): return self.combine_fn.extract_output(self.accumulator) def __str__(self): return '<%s>' % self._str_internal() def __repr__(self): return '<%s at %s>' % (self._str_internal(), hex(id(self))) def _str_internal(self): return '%s %s %s' % ( self.name, self.combine_fn.__class__.__name__, self.value()) class AccumulatorCombineFnCounter(Counter): """Counter optimized for a mutating accumulator that holds all the logic.""" def __init__(self, name, combine_fn): # type: (CounterName, cy_combiners.AccumulatorCombineFn) -> None assert isinstance(combine_fn, cy_combiners.AccumulatorCombineFn) super().__init__(name, combine_fn) self.reset() def update(self, value): self._fast_add_input(value) def update_n(self, value, n): self._fast_add_input_n(value, n) def reset(self): self.accumulator = self.combine_fn.create_accumulator() self._fast_add_input = self.accumulator.add_input self._fast_add_input_n = self.accumulator.add_input_n class CounterFactory(object): """Keeps track of unique counters.""" def __init__(self): self.counters = {} # type: Dict[CounterName, Counter] # Lock to be acquired when accessing the counters map. self._lock = threading.Lock() def get_counter(self, name, combine_fn): # type: (CounterName, core.CombineFn) -> Counter """Returns a counter with the requested name. Passing in the same name will return the same counter; the combine_fn must agree. Args: name: the name of this counter. Typically has three parts: "step-output-counter". combine_fn: the CombineFn to use for aggregation Returns: A new or existing counter with the requested name. """ with self._lock: counter = self.counters.get(name, None) if counter: assert counter.combine_fn == combine_fn else: if isinstance(combine_fn, cy_combiners.AccumulatorCombineFn): counter = AccumulatorCombineFnCounter(name, combine_fn) else: counter = Counter(name, combine_fn) self.counters[name] = counter return counter def reset(self): # Counters are cached in state sampler states. with self._lock: for counter in self.counters.values(): counter.reset() def get_counters(self): """Returns the current set of counters. Returns: An iterable that contains the current set of counters. To make sure that multiple threads can iterate over the set of counters, we return a new iterable here. Note that the actual set of counters may get modified after this method returns hence the returned iterable may be stale. """ with self._lock: return self.counters.values() # pylint: disable=bad-option-value ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/counters_test.py0000644000175100017510000001212215206113033022267 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for counters and counter names.""" # pytype: skip-file import unittest from parameterized import parameterized_class from apache_beam.utils import counters from apache_beam.utils.counters import CounterName class CounterNameTest(unittest.TestCase): def test_name_string_representation(self): counter_name = CounterName('counter_name', 'stage_name', 'step_name') # This string representation is utilized by the worker to report progress. # Change only if the worker code has also been changed. self.assertEqual('stage_name-step_name-counter_name', str(counter_name)) self.assertIn( ' at 0x', repr(counter_name)) def test_equal_objects(self): self.assertEqual( CounterName('counter_name', 'stage_name', 'step_name'), CounterName('counter_name', 'stage_name', 'step_name')) self.assertNotEqual( CounterName('counter_name', 'stage_name', 'step_name'), CounterName('counter_name', 'stage_name', 'step_nam')) # Testing objects with an IOTarget. self.assertEqual( CounterName( 'counter_name', 'stage_name', 'step_name', io_target=counters.side_input_id(1, 's9')), CounterName( 'counter_name', 'stage_name', 'step_name', io_target=counters.side_input_id(1, 's9'))) self.assertNotEqual( CounterName( 'counter_name', 'stage_name', 'step_name', io_target=counters.side_input_id(1, 's')), CounterName( 'counter_name', 'stage_name', 'step_name', io_target=counters.side_input_id(1, 's9'))) def test_hash_two_objects(self): self.assertEqual( hash(CounterName('counter_name', 'stage_name', 'step_name')), hash(CounterName('counter_name', 'stage_name', 'step_name'))) self.assertNotEqual( hash(CounterName('counter_name', 'stage_name', 'step_name')), hash(CounterName('counter_name', 'stage_name', 'step_nam'))) class CounterTest(unittest.TestCase): def setUp(self): self.counter_factory = counters.CounterFactory() def test_sum_counter(self): sum_counter = self.counter_factory.get_counter( CounterName('sum', 'stage_foo', 'step_bar'), counters.Counter.SUM) for i in range(100): sum_counter.update(i) self.assertEqual(99 * 50, sum_counter.value()) def test_mean_counter(self): mean_counter = self.counter_factory.get_counter( CounterName('mean', 'stage_foo', 'step_bar'), counters.Counter.MEAN) for i in range(100): mean_counter.update(i) self.assertEqual(49, mean_counter.value()) def test_distribution_counter(self): distribution_counter = self.counter_factory.get_counter( CounterName('distribution', 'stage_foo', 'step_bar'), counters.Counter.BEAM_DISTRIBUTION) for i in range(100): distribution_counter.update(i) self.assertEqual((49, 4950, 100, 0, 99), distribution_counter.value()) @parameterized_class([ { 'combiner': counters.Counter.SUM }, { 'combiner': counters.Counter.MEAN }, { 'combiner': counters.Counter.BEAM_DISTRIBUTION }, { 'combiner': counters.Counter.DATAFLOW_DISTRIBUTION }, ]) class GeneralCounterTest(unittest.TestCase): def setUp(self): self.counter_factory = counters.CounterFactory() def test_reset(self): counter = self.counter_factory.get_counter( CounterName(self.combiner.default_label, 'stage_foo', 'reset'), self.combiner) for value in range(100): counter.update(value) expected = counter.value() counter.reset() for value in range(100): counter.update(value) self.assertEqual(expected, counter.value()) def test_update_n(self): counter = self.counter_factory.get_counter( CounterName(self.combiner.default_label, 'stage_foo', 'update_n'), self.combiner) for i in range(100): value = i n = 100 - i for _ in range(n): counter.update(value) expected = counter.value() counter.reset() for i in range(100): value = i n = 100 - i counter.update_n(value, n) self.assertEqual(expected, counter.value()) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/histogram.py0000644000175100017510000002364515206113033021377 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import math import threading from collections import Counter from apache_beam.portability.api import metrics_pb2 _LOGGER = logging.getLogger(__name__) class Histogram(object): """A histogram that supports estimated percentile with linear interpolation. """ def __init__(self, bucket_type): self._lock = threading.Lock() self._bucket_type = bucket_type self._buckets = Counter() self._num_records = 0 self._num_top_records = 0 self._num_bot_records = 0 def clear(self): with self._lock: self._buckets = Counter() self._num_records = 0 self._num_top_records = 0 self._num_bot_records = 0 def copy(self): with self._lock: histogram = Histogram(self._bucket_type) histogram._num_records = self._num_records histogram._num_top_records = self._num_top_records histogram._num_bot_records = self._num_bot_records histogram._buckets = self._buckets.copy() return histogram def combine(self, other): if not isinstance(other, Histogram) or self._bucket_type != other._bucket_type: raise RuntimeError('failed to combine histogram.') other_histogram = other.copy() with self._lock: histogram = Histogram(self._bucket_type) histogram._num_records = self._num_records + other_histogram._num_records histogram._num_top_records = ( self._num_top_records + other_histogram._num_top_records) histogram._num_bot_records = ( self._num_bot_records + other_histogram._num_bot_records) histogram._buckets = self._buckets + other_histogram._buckets return histogram def record(self, *args): for arg in args: self._record(arg) def _record(self, value): range_from = self._bucket_type.range_from() range_to = self._bucket_type.range_to() with self._lock: if value >= range_to: _LOGGER.warning('record is out of upper bound %s: %s', range_to, value) self._num_top_records += 1 elif value < range_from: _LOGGER.warning( 'record is out of lower bound %s: %s', range_from, value) self._num_bot_records += 1 else: index = self._bucket_type.bucket_index(value) self._buckets[index] = self._buckets.get(index, 0) + 1 self._num_records += 1 def total_count(self): return self._num_records + self._num_top_records + self._num_bot_records def p99(self): return self.get_linear_interpolation(0.99) def p90(self): return self.get_linear_interpolation(0.90) def p50(self): return self.get_linear_interpolation(0.50) def get_percentile_info(self): def _format(f): if f == float('-inf'): return '<%s' % self._bucket_type.range_from() elif f == float('inf'): return '>=%s' % self._bucket_type.range_to() else: return str(int(round(f))) # pylint: disable=bad-option-value with self._lock: if self.total_count(): return ( 'Total count: %s, ' 'P99: %s, P90: %s, P50: %s' % ( self.total_count(), _format(self._get_linear_interpolation(0.99)), _format(self._get_linear_interpolation(0.90)), _format(self._get_linear_interpolation(0.50)))) else: return ('Total count: %s' % (self.total_count(), )) def get_linear_interpolation(self, percentile): """Calculate percentile estimation based on linear interpolation. It first finds the bucket which includes the target percentile and projects the estimated point in the bucket by assuming all the elements in the bucket are uniformly distributed. Args: percentile: The target percentile of the value returning from this method. Should be a floating point number greater than 0 and less than 1. """ if percentile > 1 or percentile < 0: raise ValueError('percentile should be between 0 and 1.') with self._lock: return self._get_linear_interpolation(percentile) def _get_linear_interpolation(self, percentile): total_num_records = self.total_count() if total_num_records == 0: raise RuntimeError('histogram has no record.') index = 0 record_sum = self._num_bot_records if record_sum / total_num_records >= percentile: return float('-inf') while index < self._bucket_type.num_buckets(): record_sum += self._buckets.get(index, 0) if record_sum / total_num_records >= percentile: break index += 1 if index == self._bucket_type.num_buckets(): return float('inf') frac_percentile = percentile - ( record_sum - self._buckets[index]) / total_num_records bucket_percentile = self._buckets[index] / total_num_records frac_bucket_size = frac_percentile * self._bucket_type.bucket_size( index) / bucket_percentile return ( self._bucket_type.range_from() + self._bucket_type.accumulated_bucket_size(index) + frac_bucket_size) def __eq__(self, other): if not isinstance(other, Histogram): return False def nonzero_buckets(buckets): return {k: v for k, v in buckets.items() if v != 0} return ( self._bucket_type == other._bucket_type and self._num_records == other._num_records and self._num_top_records == other._num_top_records and self._num_bot_records == other._num_bot_records and nonzero_buckets(self._buckets) == nonzero_buckets(other._buckets)) def __hash__(self): return hash(( self._bucket_type, self._num_records, self._num_top_records, self._num_bot_records, frozenset(self._buckets.items()))) def to_runner_api(self) -> metrics_pb2.HistogramValue: return metrics_pb2.HistogramValue( count=self.total_count(), bucket_counts=[ self._buckets.get(idx, 0) for idx in range(self._bucket_type.num_buckets()) ], bucket_options=self._bucket_type.to_runner_api()) @classmethod def from_runner_api(cls, proto: metrics_pb2.HistogramValue): bucket_options_proto = proto.bucket_options if bucket_options_proto.linear is not None: bucket_options = LinearBucket.from_runner_api(bucket_options_proto) else: raise NotImplementedError histogram = cls(bucket_options) with histogram._lock: for bucket_index, count in enumerate(proto.bucket_counts): histogram._buckets[bucket_index] = count histogram._num_records = sum(proto.bucket_counts) return histogram class BucketType(object): def range_from(self): """Lower bound of a starting bucket.""" raise NotImplementedError def range_to(self): """Upper bound of an ending bucket.""" raise NotImplementedError def num_buckets(self): """The number of buckets.""" raise NotImplementedError def bucket_index(self, value): """Get the bucket array index for the given value.""" raise NotImplementedError def bucket_size(self, index): """Get the bucket size for the given bucket array index.""" raise NotImplementedError def accumulated_bucket_size(self, end_index): """Get the accumulated bucket size from bucket index 0 until endIndex. Generally, this can be calculated as `sigma(0 <= i < endIndex) getBucketSize(i)`. However, a child class could provide better optimized calculation. """ raise NotImplementedError def to_runner_api(self): """Convert to the runner API representation.""" raise NotImplementedError @classmethod def from_runner_api(cls, proto): raise NotImplementedError class LinearBucket(BucketType): def __init__(self, start, width, num_buckets): """Create a histogram with linear buckets. Args: start: Lower bound of a starting bucket. width: Bucket width. Smaller width implies a better resolution for percentile estimation. num_buckets: The number of buckets. Upper bound of an ending bucket is defined by start + width * numBuckets. """ self._start = start self._width = width self._num_buckets = num_buckets def range_from(self): return self._start def range_to(self): return self._start + self._width * self._num_buckets def num_buckets(self): return self._num_buckets def bucket_index(self, value): return math.floor((value - self._start) / self._width) def bucket_size(self, index): return self._width def accumulated_bucket_size(self, end_index): return self._width * end_index def __eq__(self, other): if not isinstance(other, LinearBucket): return False return ( self._start == other._start and self._width == other._width and self._num_buckets == other._num_buckets) def __hash__(self): return hash((self._start, self._width, self._num_buckets)) def to_runner_api(self): return metrics_pb2.HistogramValue.BucketOptions( linear=metrics_pb2.HistogramValue.BucketOptions.Linear( number_of_buckets=self._num_buckets, width=self._width, start=self._start)) @classmethod def from_runner_api(cls, proto): return LinearBucket( start=proto.linear.start, width=proto.linear.width, num_buckets=proto.linear.number_of_buckets) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/histogram_test.py0000644000175100017510000001416515206113033022433 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for Histogram.""" import unittest from mock import patch from apache_beam.utils.histogram import Histogram from apache_beam.utils.histogram import LinearBucket class HistogramTest(unittest.TestCase): @patch('apache_beam.utils.histogram._LOGGER') def test_out_of_range(self, mock_logger): histogram = Histogram(LinearBucket(0, 20, 5)) histogram.record(100) mock_logger.warning.assert_called_with( 'record is out of upper bound %s: %s', 100, 100) self.assertEqual(histogram.total_count(), 1) def test_boundary_buckets(self): histogram = Histogram(LinearBucket(0, 20, 5)) histogram.record(0) histogram.record(99.9) self.assertEqual(histogram._buckets[0], 1) self.assertEqual(histogram._buckets[4], 1) def test_fractional_buckets(self): histogram1 = Histogram(LinearBucket(0, 10 / 3, 3)) histogram1.record(3.33) histogram1.record(6.66) self.assertEqual(histogram1._buckets[0], 1) self.assertEqual(histogram1._buckets[1], 1) histogram2 = Histogram(LinearBucket(0, 10 / 3, 3)) histogram2.record(3.34) histogram2.record(6.67) self.assertEqual(histogram2._buckets[1], 1) self.assertEqual(histogram2._buckets[2], 1) def test_p50(self): histogram1 = Histogram(LinearBucket(0, 0.2, 50)) histogram1.record(0, 1, 2, 3, 4, 5, 6, 7, 8, 9) self.assertEqual(histogram1.p50(), 4.2) histogram2 = Histogram(LinearBucket(0, 0.02, 50)) histogram2.record(0, 0, 0) self.assertEqual(histogram2.p50(), 0.01) def test_p90(self): histogram1 = Histogram(LinearBucket(0, 0.2, 50)) histogram1.record(0, 1, 2, 3, 4, 5, 6, 7, 8, 9) self.assertEqual(histogram1.p90(), 8.2) histogram2 = Histogram(LinearBucket(0, 0.02, 50)) histogram2.record(0, 0, 0) self.assertAlmostEqual(histogram2.p90(), 0.018, places=3) def test_p99(self): histogram1 = Histogram(LinearBucket(0, 0.2, 50)) histogram1.record(0, 1, 2, 3, 4, 5, 6, 7, 8, 9) self.assertEqual(histogram1.p99(), 9.18) histogram2 = Histogram(LinearBucket(0, 0.02, 50)) histogram2.record(0, 0, 0) self.assertAlmostEqual(histogram2.p99(), 0.02, places=3) def test_p90_negative(self): histogram1 = Histogram(LinearBucket(-10, 0.2, 50)) histogram1.record(-1, -2, -3, -4, -5, -6, -7, -8, -9, -10) self.assertEqual(histogram1.p90(), -1.8) histogram2 = Histogram(LinearBucket(-1, 0.02, 50)) histogram2.record(-1, -1, -1) self.assertEqual(histogram2.p90(), -0.982) def test_p90_negative_to_positive(self): histogram1 = Histogram(LinearBucket(-5, 0.2, 50)) histogram1.record(-1, -2, -3, -4, -5, 0, 1, 2, 3, 4) self.assertEqual(histogram1.p90(), 3.2) histogram2 = Histogram(LinearBucket(-0.5, 0.02, 50)) histogram2.record(-0.5, -0.5, -0.5) self.assertEqual(histogram2.p90(), -0.482) def test_p50_negative_infinity(self): histogram = Histogram(LinearBucket(0, 0.2, 50)) histogram.record(-1, -2, -3, -4, -5, 0, 1, 2, 3, 4) self.assertEqual(histogram.p50(), float('-inf')) self.assertIn('P50: <0', histogram.get_percentile_info()) def test_p50_positive_infinity(self): histogram = Histogram(LinearBucket(0, 0.2, 50)) histogram.record(6, 7, 8, 9, 10, 11, 12, 13, 14, 15) self.assertEqual(histogram.p50(), float('inf')) self.assertIn('P50: >=10', histogram.get_percentile_info()) def test_empty_p99(self): histogram = Histogram(LinearBucket(0, 0.2, 50)) with self.assertRaises(RuntimeError) as cm: histogram.p99() self.assertEqual(str(cm.exception), 'histogram has no record.') def test_clear(self): histogram = Histogram(LinearBucket(0, 0.2, 50)) histogram.record(-1, 1, 2, 3) self.assertEqual(histogram.total_count(), 4) self.assertEqual(histogram._buckets[5], 1) histogram.clear() self.assertEqual(histogram.total_count(), 0) self.assertEqual(histogram._buckets.get(5, 0), 0) def test_equal(self): histogram1 = Histogram(LinearBucket(0, 0.2, 50)) histogram2 = Histogram(LinearBucket(0, 0.2, 50)) self.assertEqual(histogram1, histogram2) self.assertEqual(hash(histogram1), hash(histogram2)) histogram1 = Histogram(LinearBucket(0, 0.2, 5)) histogram2 = Histogram(LinearBucket(0, 0.2, 50)) self.assertNotEqual(histogram1, histogram2) histogram1 = Histogram(LinearBucket(0, 0.2, 50)) histogram2 = Histogram(LinearBucket(0, 0.2, 50)) histogram1.record(1) histogram2.record(1) self.assertEqual(histogram1, histogram2) self.assertEqual(hash(histogram1), hash(histogram2)) histogram1 = Histogram(LinearBucket(0, 0.2, 50)) histogram2 = Histogram(LinearBucket(0, 0.2, 50)) histogram1.record(1) histogram2.record(10) self.assertNotEqual(histogram1, histogram2) def test_copy(self): histogram1 = Histogram(LinearBucket(0, 0.2, 50)) histogram1.record(0, 1, 2, 3, 4) histogram2 = histogram1.copy() self.assertEqual(histogram1, histogram2) histogram1.record(5, 6, 7, 8, 9) self.assertNotEqual(histogram1, histogram2) def test_combine(self): histogram1 = Histogram(LinearBucket(0, 0.2, 50)) histogram1.record(0, 1, 2, 3, 4) histogram2 = Histogram(LinearBucket(0, 0.2, 50)) histogram2.record(5, 6, 7, 8, 9) histogram = Histogram(LinearBucket(0, 0.2, 50)) histogram.record(0, 1, 2, 3, 4, 5, 6, 7, 8, 9) self.assertEqual(histogram1.combine(histogram2), histogram) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/interactive_utils.py0000644000175100017510000000723215206113033023131 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Common interactive utility module.""" # pytype: skip-file import logging _LOGGER = logging.getLogger(__name__) def is_in_ipython(): """Determines if current code is executed within an ipython session.""" try: from IPython import get_ipython # pylint: disable=import-error if get_ipython(): return True return False except ImportError: # If dependencies are not available, then not interactive for sure. return False except (KeyboardInterrupt, SystemExit): raise except: # pylint: disable=bare-except _LOGGER.info( 'Unexpected error occurred, treated as not in IPython.', exc_info=True) return False def is_in_notebook(): """Determines if current code is executed from an ipython notebook. If is_in_notebook() is True, then is_in_ipython() must also be True. """ is_in_notebook = False if is_in_ipython(): # The import and usage must be valid under the execution path. from IPython import get_ipython if 'IPKernelApp' in get_ipython().config: is_in_notebook = True return is_in_notebook def alter_label_if_ipython(transform, pvalueish): """Alters the label to an interactive label with ipython prompt metadata prefixed for the given transform if the given pvalueish belongs to a user-defined pipeline and current code execution is within an ipython kernel. Otherwise, noop. A label is either a user-defined or auto-generated str name of a PTransform that is unique within a pipeline. If current environment is_in_ipython(), Beam can implicitly create interactive labels to replace labels of top-level PTransforms to be applied. The label is formatted as: `Cell {prompt}: {original_label}`. """ if is_in_ipython(): from apache_beam.runners.interactive import interactive_environment as ie # Tracks user defined pipeline instances in watched scopes so that we only # alter labels for any transform to pvalueish belonging to those pipeline # instances, excluding any transform to be applied in other pipeline # instances the Beam SDK creates implicitly. ie.current_env().track_user_pipelines() from IPython import get_ipython prompt = get_ipython().execution_count pipeline = _extract_pipeline_of_pvalueish(pvalueish) if (pipeline # We only alter for transforms to be applied to user-defined pipelines # at pipeline construction time. and pipeline in ie.current_env().tracked_user_pipelines): transform.label = '[{}]: {}'.format(prompt, transform.label) def _extract_pipeline_of_pvalueish(pvalueish): """Extracts the pipeline that the given pvalueish belongs to.""" if isinstance(pvalueish, tuple) and len(pvalueish) > 0: pvalue = pvalueish[0] elif isinstance(pvalueish, dict) and len(pvalueish) > 0: pvalue = next(iter(pvalueish.values())) else: pvalue = pvalueish if hasattr(pvalue, 'pipeline'): return pvalue.pipeline return None ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/interactive_utils_test.py0000644000175100017510000000435715206113033024175 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Tests for interactive_utils module.""" # pytype: skip-file import unittest from unittest.mock import patch from apache_beam.runners.interactive import interactive_environment as ie from apache_beam.runners.interactive.testing.mock_ipython import mock_get_ipython from apache_beam.utils.interactive_utils import is_in_ipython def unavailable_ipython(): # ModuleNotFoundError since Py3.6 is sub class of ImportError. An example, # 'import apache_beam.hahaha' raises a ModuleNotFoundError while # 'from apache_beam import hahaha' raises an ImportError. # For simplicity, we only use super class ImportError here. raise ImportError('Module IPython is not found.') def corrupted_ipython(): raise AttributeError('Module IPython does not contain get_ipython.') @unittest.skipIf( not ie.current_env().is_interactive_ready, '[interactive] dependency is not installed.') class IPythonTest(unittest.TestCase): @patch('IPython.get_ipython', new_callable=mock_get_ipython) def test_is_in_ipython_when_in_ipython_kernel(self, kernel): self.assertTrue(is_in_ipython()) @patch('IPython.get_ipython', new_callable=lambda: unavailable_ipython) def test_is_not_in_ipython_when_no_ipython_dep(self, unavailable): self.assertFalse(is_in_ipython()) @patch('IPython.get_ipython', new_callable=lambda: corrupted_ipython) def test_is_not_ipython_when_ipython_errors_out(self, corrupted): self.assertFalse(is_in_ipython()) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/logger.py0000644000175100017510000001061415206113033020651 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Helper functions for easier logging. This module provides a few convenient logging methods, some of which were adopted from https://github.com/abseil/abseil-py/blob/master/absl/logging/__init__.py in https://github.com/facebookresearch/detectron2/blob/main/detectron2/utils/logger.py """ import logging import os import sys import time from collections import Counter from types import FrameType from typing import Optional from typing import Union def _find_caller() -> tuple[str, tuple]: """ Returns: str: module name of the caller tuple: a hashable key to be used to identify different callers """ frame: Optional[FrameType] = sys._getframe(2) while frame: code = frame.f_code if os.path.join("utils", "logger.") not in code.co_filename: mod_name = frame.f_globals["__name__"] if mod_name == "__main__": mod_name = "apache_beam" return mod_name, (code.co_filename, frame.f_lineno, code.co_name) frame = frame.f_back # To appease mypy. Code returns earlier in practice. return "unknown", ("unknown", 0, "unknown") _LOG_COUNTER = Counter() _LOG_TIMER = {} def log_first_n( lvl: int, msg: str, *args, n: int = 1, name: Optional[str] = None, key: Union[str, tuple[str]] = "caller") -> None: """ Log only for the first n times. Args: lvl (int): the logging level msg (str): n (int): name (str): name of the logger to use. Will use the caller's module by default. key (str or tuple[str]): the string(s) can be one of "caller" or "message", which defines how to identify duplicated logs. For example, if called with `n=1, key="caller"`, this function will only log the first call from the same caller, regardless of the message content. If called with `n=1, key="message"`, this function will log the same content only once, even if they are called from different places. If called with `n=1, key=("caller", "message")`, this function will not log only if the same caller has logged the same message before. """ key_tuple = (key, ) if isinstance(key, str) else key assert len(key_tuple) > 0 caller_module, caller_key = _find_caller() hash_key: tuple = () if "caller" in key_tuple: hash_key = hash_key + caller_key if "message" in key_tuple: hash_key = hash_key + (msg, ) _LOG_COUNTER[hash_key] += 1 if _LOG_COUNTER[hash_key] <= n: logging.getLogger(name or caller_module).log(lvl, msg, *args) def log_every_n( lvl: int, msg: str, *args, n: int = 1, name: Optional[str] = None) -> None: """ Log once per n times. Args: lvl (int): the logging level msg (str): n (int): name (str): name of the logger to use. Will use the caller's module by default. """ caller_module, key = _find_caller() _LOG_COUNTER[key] += 1 if n == 1 or _LOG_COUNTER[key] % n == 1: logging.getLogger(name or caller_module).log(lvl, msg, *args) def log_every_n_seconds( lvl: int, msg: str, *args, n: int = 1, name: Optional[str] = None) -> None: """ Log no more than once per n seconds. Args: lvl (int): the logging level msg (str): n (int): name (str): name of the logger to use. Will use the caller's module by default. """ caller_module, key = _find_caller() last_logged = _LOG_TIMER.get(key, None) current_time = time.time() if last_logged is None or current_time - last_logged >= n: logging.getLogger(name or caller_module).log(lvl, msg, *args) _LOG_TIMER[key] = current_time ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/logger_test.py0000644000175100017510000001022315206113033021704 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest from unittest.mock import patch import pytest from apache_beam.utils.logger import _LOG_COUNTER from apache_beam.utils.logger import _LOG_TIMER from apache_beam.utils.logger import log_every_n from apache_beam.utils.logger import log_every_n_seconds from apache_beam.utils.logger import log_first_n @pytest.mark.no_xdist class TestLogFirstN(unittest.TestCase): def setUp(self): _LOG_COUNTER.clear() _LOG_TIMER.clear() @patch('apache_beam.utils.logger.logging.getLogger') def test_log_first_n_once(self, mock_get_logger): mock_logger = mock_get_logger.return_value for _ in range(5): log_first_n(logging.INFO, "Test message %s", "arg", n=1) mock_logger.log.assert_called_once_with( logging.INFO, "Test message %s", "arg") @patch('apache_beam.utils.logger.logging.getLogger') def test_log_first_n_multiple(self, mock_get_logger): mock_logger = mock_get_logger.return_value for _ in range(5): log_first_n(logging.INFO, "Test message %s", "arg", n=3) self.assertEqual(mock_logger.log.call_count, 3) mock_logger.log.assert_called_with(logging.INFO, "Test message %s", "arg") @patch('apache_beam.utils.logger.logging.getLogger') def test_log_first_n_with_different_callers(self, mock_get_logger): mock_logger = mock_get_logger.return_value for _ in range(5): log_first_n(logging.INFO, "Test message", n=2) # call from another "caller" (another line) for _ in range(5): log_first_n(logging.INFO, "Test message", n=2) self.assertEqual(mock_logger.log.call_count, 4) @patch('apache_beam.utils.logger.logging.getLogger') def test_log_first_n_with_message_key(self, mock_get_logger): mock_logger = mock_get_logger.return_value log_first_n(logging.INFO, "Test message", n=1, key="message") log_first_n(logging.INFO, "Test message", n=1, key="message") self.assertEqual(mock_logger.log.call_count, 1) @patch('apache_beam.utils.logger.logging.getLogger') def test_log_first_n_with_caller_and_message_key(self, mock_get_logger): mock_logger = mock_get_logger.return_value for message in ["Test message", "Another message"]: for _ in range(5): log_first_n(logging.INFO, message, n=1, key=("caller", "message")) self.assertEqual(mock_logger.log.call_count, 2) @patch('apache_beam.utils.logger.logging.getLogger') def test_log_every_n_multiple(self, mock_get_logger): mock_logger = mock_get_logger.return_value for _ in range(9): log_every_n(logging.INFO, "Test message", n=2) self.assertEqual(mock_logger.log.call_count, 5) @patch('apache_beam.utils.logger.logging.getLogger') @patch('apache_beam.utils.logger.time.time') def test_log_every_n_seconds_always(self, mock_time, mock_get_logger): mock_logger = mock_get_logger.return_value for i in range(3): mock_time.return_value = i log_every_n_seconds(logging.INFO, "Test message", n=0) self.assertEqual(mock_logger.log.call_count, 3) @patch('apache_beam.utils.logger.logging.getLogger') @patch('apache_beam.utils.logger.time.time') def test_log_every_n_seconds_multiple(self, mock_time, mock_get_logger): mock_logger = mock_get_logger.return_value for i in range(4): mock_time.return_value = i log_every_n_seconds(logging.INFO, "Test message", n=2) self.assertEqual(mock_logger.log.call_count, 2) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/multi_process_shared.py0000644000175100017510000004362415206113033023617 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Implements a shared object that spans processes. This object will be instanciated once per VM and methods will be invoked on it via rpc. """ # pytype: skip-file import atexit import logging import multiprocessing.managers import os import tempfile import threading import time import traceback from typing import Any from typing import Callable from typing import Generic from typing import Optional from typing import TypeVar import fasteners from apache_beam.utils import retry # In some python versions, there is a bug where AutoProxy doesn't handle # the kwarg 'manager_owned'. We implement our own backup here to make sure # we avoid this problem. More info here: # https://stackoverflow.com/questions/46779860/multiprocessing-managers-and-custom-classes autoproxy = multiprocessing.managers.AutoProxy # type: ignore[attr-defined] def patched_autoproxy( token, serializer, manager=None, authkey=None, exposed=None, incref=True, manager_owned=True): return autoproxy(token, serializer, manager, authkey, exposed, incref) multiprocessing.managers.AutoProxy = patched_autoproxy # type: ignore[attr-defined] T = TypeVar('T') AUTH_KEY = b'mps' class _SingletonProxy: """Proxies the shared object so we can release it with better errors and no risk of dangling references in the multiprocessing manager infrastructure. """ def __init__(self, entry): # Guard names so as to not conflict with names of underlying object. self._SingletonProxy_entry = entry self._SingletonProxy_valid = True # Used to make the shared object callable (see _AutoProxyWrapper below) def singletonProxy_call__(self, *args, **kwargs): if not self._SingletonProxy_valid: raise RuntimeError('Entry was released.') return self._SingletonProxy_entry.obj.__call__(*args, **kwargs) def singletonProxy_release(self): assert self._SingletonProxy_valid self._SingletonProxy_valid = False def singletonProxy_unsafe_hard_delete(self): assert self._SingletonProxy_valid self._SingletonProxy_entry.unsafe_hard_delete() def __getattr__(self, name): if not self._SingletonProxy_valid: raise RuntimeError('Entry was released.') try: return getattr(self._SingletonProxy_entry.obj, name) except AttributeError as e: # Swallow AttributeError exceptions so that they are ignored when # calculating public functions. These can occur if __getattr__ is # overriden, for example to only support some platforms. This will mean # that these functions will be silently unavailable to the # MultiProcessShared object, leading to worse errors when someone tries # to use them, but it will keep them from breaking the whole object for # functions which are unusable anyways. logging.info( 'Attribute %s is unavailable as a public function because ' 'its __getattr__ function raised the following exception ' '%s', name, e) return None def __dir__(self): # Needed for multiprocessing.managers's proxying. dir = self._SingletonProxy_entry.obj.__dir__() dir.append('singletonProxy_call__') dir.append('singletonProxy_release') dir.append('singletonProxy_unsafe_hard_delete') return dir class _SingletonEntry: """Represents a single, refcounted entry in this process.""" def __init__( self, constructor, initialize_eagerly=True, hard_delete_callback=None): self.constructor = constructor self._hard_delete_callback = hard_delete_callback self.refcount = 0 self.lock = threading.Lock() if initialize_eagerly: self.obj = constructor() self.initialied = True else: self.initialied = False def acquire(self): with self.lock: if not self.initialied: self.obj = self.constructor() self.initialied = True self.refcount += 1 return _SingletonProxy(self) def release(self, proxy): proxy.singletonProxy_release() with self.lock: self.refcount -= 1 if self.refcount == 0: del self.obj self.initialied = False def unsafe_hard_delete(self): with self.lock: if self.initialied: del self.obj self.initialied = False if self._hard_delete_callback: self._hard_delete_callback() class _SingletonManager: entries: dict[Any, Any] = {} def __init__(self): self._hard_delete_callback = None def set_hard_delete_callback(self, callback): self._hard_delete_callback = callback def register_singleton( self, constructor, tag, initialize_eagerly=True, hard_delete_callback=None): assert tag not in self.entries, tag self.entries[tag] = _SingletonEntry( constructor, initialize_eagerly, hard_delete_callback) def has_singleton(self, tag): return tag in self.entries def acquire_singleton(self, tag): return self.entries[tag].acquire() def release_singleton(self, tag, obj): return self.entries[tag].release(obj) def unsafe_hard_delete_singleton(self, tag): self.entries[tag].unsafe_hard_delete() _process_level_singleton_manager = _SingletonManager() _process_local_lock = threading.Lock() class _SingletonRegistrar(multiprocessing.managers.BaseManager): pass _SingletonRegistrar.register( 'acquire_singleton', callable=_process_level_singleton_manager.acquire_singleton) _SingletonRegistrar.register( 'release_singleton', callable=_process_level_singleton_manager.release_singleton) _SingletonRegistrar.register( 'unsafe_hard_delete_singleton', callable=_process_level_singleton_manager.unsafe_hard_delete_singleton) # By default, objects registered with BaseManager.register will have only # public methods available (excluding __call__). If you know the functions # you would like to expose, you can do so at register time with the `exposed` # attribute. Since we don't, we will add a wrapper around the returned AutoProxy # object to handle __call__ function calls and turn them into # singletonProxy_call__ calls (which is a wrapper around the underlying # object's __call__ function) class _AutoProxyWrapper: def __init__(self, proxyObject: multiprocessing.managers.BaseProxy): self._proxyObject = proxyObject def __call__(self, *args, **kwargs): return self._proxyObject.singletonProxy_call__(*args, **kwargs) def __getattr__(self, name): return getattr(self._proxyObject, name) def get_auto_proxy_object(self): return self._proxyObject def unsafe_hard_delete(self): self._proxyObject.unsafe_hard_delete() def _run_server_process(address_file, tag, constructor, authkey, life_line): """ Runs in a separate process. Includes a 'Suicide Pact' monitor: If parent dies, I die. """ parent_pid = os.getppid() def cleanup_files(): logging.info("Server process exiting. Deleting files for %s", tag) try: if os.path.exists(address_file): os.remove(address_file) if os.path.exists(address_file + ".error"): os.remove(address_file + ".error") except Exception as e: logging.warning('Failed to cleanup files for tag %s: %s', tag, e) def handle_unsafe_hard_delete(): cleanup_files() os._exit(0) def _monitor_parent(): """Checks if parent is alive every second.""" while True: try: # This will break if parent dies. life_line.recv_bytes() except (EOFError, OSError, BrokenPipeError): logging.warning( "Process %s detected Parent %s died. Self-destructing.", os.getpid(), parent_pid) cleanup_files() os._exit(0) time.sleep(0.5) atexit.register(cleanup_files) try: t = threading.Thread(target=_monitor_parent, daemon=True) logging.getLogger().setLevel(logging.INFO) multiprocessing.current_process().authkey = authkey serving_manager = _SingletonRegistrar( address=('localhost', 0), authkey=authkey) _process_level_singleton_manager.set_hard_delete_callback( handle_unsafe_hard_delete) _process_level_singleton_manager.register_singleton( constructor, tag, initialize_eagerly=True, hard_delete_callback=handle_unsafe_hard_delete) # Start monitoring parent after initialisation is done to avoid # potential race conditions. t.start() server = serving_manager.get_server() logging.info( 'Process %s: Proxy serving %s at %s', os.getpid(), tag, server.address) with open(address_file + '.tmp', 'w') as fout: fout.write('%s:%d' % server.address) os.rename(address_file + '.tmp', address_file) server.serve_forever() except Exception: tb = traceback.format_exc() try: with open(address_file + ".error.tmp", 'w') as fout: fout.write(tb) os.rename(address_file + ".error.tmp", address_file + ".error") except Exception: logging.error("CRITICAL ERROR IN SHARED SERVER:\n%s", tb) os._exit(1) class MultiProcessShared(Generic[T]): """MultiProcessShared is used to share a single object across processes. For example, one could have the class:: class MyExpensiveObject(object): def __init__(self, args): [expensive initialization and memory allocation] def method(self, arg): ... One could share a single instance of this class by wrapping it as:: shared_ptr = MultiProcessShared(lambda: MyExpensiveObject(...)) my_expensive_object = shared_ptr.acquire() which could then be invoked as:: my_expensive_object.method(arg) This can then be released with:: shared_ptr.release(my_expensive_object) but care should be taken to avoid releasing the object too soon or expensive re-initialization may be required, defeating the point of using a shared object. Args: constructor: function that initialises / constructs the object if not present in the cache. This function should take no arguments. It should return an initialised object, or raise an exception if the object could not be initialised / constructed. tag: an indentifier to store with the cached object. If multiple MultiProcessShared instances are created with the same tag, they will all share the same proxied object. path: a temporary path in which to create the inter-process lock always_proxy: whether to direct all calls through the proxy, rather than call the object directly for the process that created it """ def __init__( self, constructor: Callable[[], T], tag: Any, *, path: str = tempfile.gettempdir(), always_proxy: Optional[bool] = None, spawn_process: bool = False): self._constructor = constructor self._tag = tag self._path = path self._always_proxy = False if always_proxy is None else always_proxy self._proxy = None self._manager = None self._rpc_address = None self._cross_process_lock = fasteners.InterProcessLock( os.path.join(self._path, self._tag) + '.lock') self._spawn_process = spawn_process def _get_manager(self): if self._manager is None: address_file = os.path.join(self._path, self._tag) + ".address" while self._manager is None: with _process_local_lock: with self._cross_process_lock: if not os.path.exists(address_file): self._create_server(address_file) if _process_level_singleton_manager.has_singleton( self._tag) and not self._always_proxy: self._manager = _process_level_singleton_manager else: with open(address_file) as fin: address = fin.read() logging.info('Connecting to remote proxy at %s', address) host, port = address.split(':') # We need to be able to authenticate with both the manager and # the process. manager = _SingletonRegistrar( address=(host, int(port)), authkey=AUTH_KEY) multiprocessing.current_process().authkey = AUTH_KEY retryable_exceptions = (ConnectionError, EOFError) @retry.with_exponential_backoff( num_retries=5, initial_delay_secs=0.1, retry_filter=lambda exn: isinstance( exn, retryable_exceptions)) def connect_manager(): manager.connect() try: connect_manager() self._manager = manager except retryable_exceptions: # The server is no longer good, assume it died. os.unlink(address_file) return self._manager def acquire(self): # TODO: Allow passing/parameterizing the callable here, in case they are # not available at MultiProcessShared construction time (e.g. from side # inputs) # Caveat: They must always agree, as they will be ignored if the object # is already constructed. singleton = self._get_manager().acquire_singleton(self._tag) # Trigger a sweep of zombie processes. # calling active_children() has the side-effect of joining any finished # processes, effectively reaping zombies from previous unsafe_hard_deletes. if self._spawn_process: multiprocessing.active_children() return _AutoProxyWrapper(singleton) def release(self, obj): self._manager.release_singleton(self._tag, obj.get_auto_proxy_object()) def unsafe_hard_delete(self): """Force deletes the underlying object This function should be used with great care since any other references to this object will now be invalid and may lead to strange errors. Only call unsafe_hard_delete if either (a) you are sure no other references to this object exist, or (b) you are ok with all existing references to this object throwing strange errors when derefrenced. """ self._get_manager().unsafe_hard_delete_singleton(self._tag) def _create_server(self, address_file): if self._spawn_process: error_file = address_file + ".error" if os.path.exists(error_file): try: os.remove(error_file) except OSError: pass # Create a pipe to connect with child process # used to clean up child process if parent dies reader, writer = multiprocessing.Pipe(duplex=False) self._life_line = writer ctx = multiprocessing.get_context('spawn') p = ctx.Process( target=_run_server_process, args=(address_file, self._tag, self._constructor, AUTH_KEY, reader), daemon=False # Must be False for nested proxies ) p.start() logging.info("Parent: Waiting for %s to write address file...", self._tag) def cleanup_process(): if self._life_line: self._life_line.close() if p.is_alive(): logging.info( "Parent: Terminating server process %s for %s", p.pid, self._tag) p.terminate() p.join() try: if os.path.exists(address_file): os.remove(address_file) if os.path.exists(error_file): os.remove(error_file) except Exception as e: logging.warning( 'Failed to cleanup files for tag %s in atexit handler: %s', self._tag, e) atexit.register(cleanup_process) start_time = time.time() last_log = start_time while True: if os.path.exists(address_file): break if os.path.exists(error_file): with open(error_file, 'r') as f: error_msg = f.read() try: os.remove(error_file) except OSError: pass if p.is_alive(): p.terminate() raise RuntimeError(f"Shared Server Process crashed:\n{error_msg}") if not p.is_alive(): exit_code = p.exitcode raise RuntimeError( "Shared Server Process died unexpectedly" f" with exit code {exit_code}") if time.time() - last_log > 300: logging.warning( "Still waiting for %s to initialize... %ss elapsed)", self._tag, int(time.time() - start_time)) last_log = time.time() time.sleep(0.05) logging.info('External process successfully started for %s', self._tag) else: # We need to be able to authenticate with both the manager # and the process. self._serving_manager = _SingletonRegistrar( address=('localhost', 0), authkey=AUTH_KEY) multiprocessing.current_process().authkey = AUTH_KEY # Initialize eagerly to avoid acting as the server if there are issues. # Note, however, that _create_server itself is called lazily. _process_level_singleton_manager.register_singleton( self._constructor, self._tag, initialize_eagerly=True) self._server = self._serving_manager.get_server() logging.info( 'Starting proxy server at %s for shared %s', self._server.address, self._tag) with open(address_file + '.tmp', 'w') as fout: fout.write('%s:%d' % self._server.address) os.rename(address_file + '.tmp', address_file) t = threading.Thread(target=self._server.serve_forever, daemon=True) t.start() logging.info('Done starting server') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/multi_process_shared_test.py0000644000175100017510000003632115206113033024652 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import logging import multiprocessing import os import tempfile import threading import unittest from typing import Any from unittest.mock import patch from apache_beam.utils import multi_process_shared class CallableCounter(object): def __init__(self, start=0): self.running = start self.lock = threading.Lock() def __call__(self): return self.running def increment(self, value=1): with self.lock: self.running += value return self.running def error(self, msg): raise RuntimeError(msg) class Counter(object): def __init__(self, start=0): self.running = start self.lock = threading.Lock() def get(self): return self.running def increment(self, value=1): with self.lock: self.running += value return self.running def error(self, msg): raise RuntimeError(msg) class CounterWithBadAttr(object): def __init__(self, start=0): self.running = start self.lock = threading.Lock() def get(self): return self.running def increment(self, value=1): with self.lock: self.running += value return self.running def error(self, msg): raise RuntimeError(msg) def __getattribute__(self, __name: str) -> Any: if __name == 'error': raise AttributeError('error is not actually supported on this platform') else: # Default behaviour return object.__getattribute__(self, __name) class MultiProcessSharedTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.shared = multi_process_shared.MultiProcessShared( Counter, tag='basic', always_proxy=True).acquire() cls.sharedCallable = multi_process_shared.MultiProcessShared( CallableCounter, tag='callable', always_proxy=True).acquire() def test_call(self): self.assertEqual(self.shared.get(), 0) self.assertEqual(self.shared.increment(), 1) self.assertEqual(self.shared.increment(10), 11) self.assertEqual(self.shared.increment(value=10), 21) self.assertEqual(self.shared.get(), 21) def test_call_illegal_attr(self): shared_handle = multi_process_shared.MultiProcessShared( CounterWithBadAttr, tag='test_call_illegal_attr', always_proxy=True) shared = shared_handle.acquire() self.assertEqual(shared.get(), 0) self.assertEqual(shared.increment(), 1) self.assertEqual(shared.get(), 1) def test_call_callable(self): self.assertEqual(self.sharedCallable(), 0) self.assertEqual(self.sharedCallable.increment(), 1) self.assertEqual(self.sharedCallable.increment(10), 11) self.assertEqual(self.sharedCallable.increment(value=10), 21) self.assertEqual(self.sharedCallable(), 21) def test_error(self): with self.assertRaisesRegex(Exception, 'something bad'): self.shared.error('something bad') def test_no_method(self): with self.assertRaisesRegex(Exception, 'no_such_method'): self.shared.no_such_method() def test_connect(self): first = multi_process_shared.MultiProcessShared( Counter, tag='counter').acquire() second = multi_process_shared.MultiProcessShared( Counter, tag='counter').acquire() self.assertEqual(first.get(), 0) self.assertEqual(first.increment(), 1) self.assertEqual(second.get(), 1) self.assertEqual(second.increment(), 2) self.assertEqual(first.get(), 2) self.assertEqual(first.increment(), 3) def test_release(self): shared1 = multi_process_shared.MultiProcessShared( Counter, tag='test_release') shared2 = multi_process_shared.MultiProcessShared( Counter, tag='test_release') counter1 = shared1.acquire() counter2 = shared2.acquire() self.assertEqual(counter1.increment(), 1) self.assertEqual(counter2.increment(), 2) counter1again = shared1.acquire() self.assertEqual(counter1again.increment(), 3) shared1.release(counter1) shared2.release(counter2) with self.assertRaisesRegex(Exception, 'released'): counter1.get() with self.assertRaisesRegex(Exception, 'released'): counter2.get() self.assertEqual(counter1again.get(), 3) shared1.release(counter1again) counter1New = shared1.acquire() self.assertEqual(counter1New.get(), 0) with self.assertRaisesRegex(Exception, 'released'): counter1.get() def test_unsafe_hard_delete(self): shared1 = multi_process_shared.MultiProcessShared( Counter, tag='test_unsafe_hard_delete', always_proxy=True) shared2 = multi_process_shared.MultiProcessShared( Counter, tag='test_unsafe_hard_delete', always_proxy=True) counter1 = shared1.acquire() counter2 = shared2.acquire() self.assertEqual(counter1.increment(), 1) self.assertEqual(counter2.increment(), 2) try: multi_process_shared.MultiProcessShared( Counter, tag='test_unsafe_hard_delete').unsafe_hard_delete() except Exception: pass with self.assertRaises(Exception): counter1.get() with self.assertRaises(Exception): counter2.get() shared3 = multi_process_shared.MultiProcessShared( Counter, tag='test_unsafe_hard_delete', always_proxy=True) counter3 = shared3.acquire() self.assertEqual(counter3.increment(), 1) def test_unsafe_hard_delete_autoproxywrapper(self): shared1 = multi_process_shared.MultiProcessShared( Counter, tag='test_unsafe_hard_delete_autoproxywrapper', always_proxy=True) shared2 = multi_process_shared.MultiProcessShared( Counter, tag='test_unsafe_hard_delete_autoproxywrapper', always_proxy=True) counter1 = shared1.acquire() counter2 = shared2.acquire() self.assertEqual(counter1.increment(), 1) self.assertEqual(counter2.increment(), 2) try: counter2.singletonProxy_unsafe_hard_delete() except Exception: pass with self.assertRaises(Exception): counter1.get() with self.assertRaises(Exception): counter2.get() counter3 = multi_process_shared.MultiProcessShared( Counter, tag='test_unsafe_hard_delete_autoproxywrapper', always_proxy=True).acquire() self.assertEqual(counter3.increment(), 1) def test_unsafe_hard_delete_no_op(self): shared1 = multi_process_shared.MultiProcessShared( Counter, tag='test_unsafe_hard_delete_no_op', always_proxy=True) shared2 = multi_process_shared.MultiProcessShared( Counter, tag='test_unsafe_hard_delete_no_op', always_proxy=True) counter1 = shared1.acquire() counter2 = shared2.acquire() self.assertEqual(counter1.increment(), 1) self.assertEqual(counter2.increment(), 2) try: multi_process_shared.MultiProcessShared( Counter, tag='no_tag_to_delete').unsafe_hard_delete() except Exception: pass self.assertEqual(counter1.increment(), 3) self.assertEqual(counter2.increment(), 4) def test_release_always_proxy(self): shared1 = multi_process_shared.MultiProcessShared( Counter, tag='test_release_always_proxy', always_proxy=True) shared2 = multi_process_shared.MultiProcessShared( Counter, tag='test_release_always_proxy', always_proxy=True) counter1 = shared1.acquire() counter2 = shared2.acquire() self.assertEqual(counter1.increment(), 1) self.assertEqual(counter2.increment(), 2) counter1again = shared1.acquire() self.assertEqual(counter1again.increment(), 3) shared1.release(counter1) shared2.release(counter2) with self.assertRaisesRegex(Exception, 'released'): counter1.get() with self.assertRaisesRegex(Exception, 'released'): counter2.get() self.assertEqual(counter1again.get(), 3) shared1.release(counter1again) counter1New = shared1.acquire() self.assertEqual(counter1New.get(), 0) with self.assertRaisesRegex(Exception, 'released'): counter1.get() class MultiProcessSharedSpawnProcessTest(unittest.TestCase): def setUp(self): tempdir = tempfile.gettempdir() for tag in ['basic', 'main', 'to_delete', 'to_keep', 'mix1', 'mix2', 'test_process_exit', 'thundering_herd_test', 'transient_test']: for ext in ['', '.address', '.address.error']: try: os.remove(os.path.join(tempdir, tag + ext)) except OSError: pass def tearDown(self): for p in multiprocessing.active_children(): if p.is_alive(): try: p.kill() p.join(timeout=1.0) except Exception: pass def test_call(self): shared = multi_process_shared.MultiProcessShared( Counter, tag='main', always_proxy=True, spawn_process=True).acquire() self.assertEqual(shared.get(), 0) self.assertEqual(shared.increment(), 1) self.assertEqual(shared.increment(10), 11) self.assertEqual(shared.increment(value=10), 21) self.assertEqual(shared.get(), 21) def test_unsafe_hard_delete_autoproxywrapper(self): shared1 = multi_process_shared.MultiProcessShared( Counter, tag='to_delete', always_proxy=True, spawn_process=True) shared2 = multi_process_shared.MultiProcessShared( Counter, tag='to_delete', always_proxy=True, spawn_process=True) counter3 = multi_process_shared.MultiProcessShared( Counter, tag='to_keep', always_proxy=True, spawn_process=True).acquire() counter1 = shared1.acquire() counter2 = shared2.acquire() self.assertEqual(counter1.increment(), 1) self.assertEqual(counter2.increment(), 2) try: counter2.singletonProxy_unsafe_hard_delete() except Exception: pass with self.assertRaises(Exception): counter1.get() with self.assertRaises(Exception): counter2.get() counter4 = multi_process_shared.MultiProcessShared( Counter, tag='to_delete', always_proxy=True, spawn_process=True).acquire() self.assertEqual(counter3.increment(), 1) self.assertEqual(counter4.increment(), 1) def test_mix_usage(self): shared1 = multi_process_shared.MultiProcessShared( Counter, tag='mix1', always_proxy=True, spawn_process=False).acquire() shared2 = multi_process_shared.MultiProcessShared( Counter, tag='mix2', always_proxy=True, spawn_process=True).acquire() self.assertEqual(shared1.get(), 0) self.assertEqual(shared1.increment(), 1) self.assertEqual(shared2.get(), 0) self.assertEqual(shared2.increment(), 1) def test_process_exits_on_unsafe_hard_delete(self): shared = multi_process_shared.MultiProcessShared( Counter, tag='test_process_exit', always_proxy=True, spawn_process=True) obj = shared.acquire() self.assertEqual(obj.increment(), 1) children = multiprocessing.active_children() server_process = None for p in children: if p.pid != os.getpid() and p.is_alive(): server_process = p break self.assertIsNotNone( server_process, "Could not find spawned server process") try: obj.singletonProxy_unsafe_hard_delete() except Exception: pass server_process.join(timeout=5) self.assertFalse( server_process.is_alive(), f"Server process {server_process.pid} is still alive after hard delete") self.assertIsNotNone( server_process.exitcode, "Process has no exit code (did not exit)") with self.assertRaises(Exception): obj.get() def test_process_exits_on_unsafe_hard_delete_with_manager(self): shared = multi_process_shared.MultiProcessShared( Counter, tag='test_process_exit', always_proxy=True, spawn_process=True) obj = shared.acquire() self.assertEqual(obj.increment(), 1) children = multiprocessing.active_children() server_process = None for p in children: if p.pid != os.getpid() and p.is_alive(): server_process = p break self.assertIsNotNone( server_process, "Could not find spawned server process") try: shared.unsafe_hard_delete() except Exception: pass server_process.join(timeout=5) self.assertFalse( server_process.is_alive(), f"Server process {server_process.pid} is still alive after hard delete") self.assertIsNotNone( server_process.exitcode, "Process has no exit code (did not exit)") with self.assertRaises(Exception): obj.get() def test_zombie_reaping_on_acquire(self): shared1 = multi_process_shared.MultiProcessShared( Counter, tag='test_zombie_reap', always_proxy=True, spawn_process=True) obj = shared1.acquire() children = multiprocessing.active_children() server_pid = next( p.pid for p in children if p.is_alive() and p.pid != os.getpid()) try: obj.singletonProxy_unsafe_hard_delete() except Exception: pass try: os.kill(server_pid, 0) is_zombie = True except OSError: is_zombie = False self.assertTrue( is_zombie, f"Server process {server_pid} was reaped too early before acquire()") shared2 = multi_process_shared.MultiProcessShared( Counter, tag='unrelated_tag', always_proxy=True, spawn_process=True) _ = shared2.acquire() # If reaping worked, our old server_pid should NOT be in this list. current_children_pids = [p.pid for p in multiprocessing.active_children()] self.assertNotIn( server_pid, current_children_pids, f"Old server process {server_pid} was not reaped by acquire() sweep") try: shared2.unsafe_hard_delete() except Exception: pass def test_transient_connection_error_recovery(self): shared1 = multi_process_shared.MultiProcessShared( Counter, tag='transient_test', always_proxy=True, spawn_process=True) shared2 = multi_process_shared.MultiProcessShared( Counter, tag='transient_test', always_proxy=True, spawn_process=True) counter1 = shared1.acquire() orig_connect = multi_process_shared._SingletonRegistrar.connect connect_calls = [0] def side_effect_connect(self_mgr, *args, **kwargs): connect_calls[0] += 1 if connect_calls[0] == 1: raise ConnectionError("Simulated transient connection failure") return orig_connect(self_mgr, *args, **kwargs) with patch.object(multi_process_shared._SingletonRegistrar, 'connect', autospec=True, side_effect=side_effect_connect): counter2 = shared2.acquire() self.assertEqual(counter1.increment(), 1) self.assertEqual(counter2.increment(), 2) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/plugin.py0000644000175100017510000000300615206113033020665 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A BeamPlugin base class.""" # pytype: skip-file class BeamPlugin(object): """Plugin base class to be extended by dependent users such as FileSystem. Any instantiated subclass will be imported at worker startup time.""" @classmethod def get_all_subclasses(cls): """Get all the subclasses of the BeamPlugin class.""" all_subclasses = [] for subclass in cls.__subclasses__(): all_subclasses.append(subclass) all_subclasses.extend(subclass.get_all_subclasses()) return all_subclasses @classmethod def get_all_plugin_paths(cls): """Get full import paths of the BeamPlugin subclass.""" def fullname(o): return o.__module__ + "." + o.__name__ return [fullname(o) for o in cls.get_all_subclasses()] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/processes.py0000644000175100017510000000760515206113033021406 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Cross-platform utilities for creating subprocesses. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file import platform import subprocess import traceback from typing import TYPE_CHECKING # On Windows, we need to use shell=True when creating subprocesses for binary # paths to be resolved correctly. force_shell = platform.system() == 'Windows' # We mimic the interface of the standard Python subprocess module. PIPE = subprocess.PIPE STDOUT = subprocess.STDOUT CalledProcessError = subprocess.CalledProcessError if TYPE_CHECKING: call = subprocess.call check_call = subprocess.check_call check_output = subprocess.check_output Popen = subprocess.Popen else: def call(*args, **kwargs): if force_shell: kwargs['shell'] = True try: out = subprocess.call(*args, **kwargs) except OSError as e: raise RuntimeError("Executable {} not found".format(args[0])) from e except subprocess.CalledProcessError as error: if isinstance(args, tuple) and (args[0][2] == "pip"): raise RuntimeError( \ "Full traceback: {}\n Pip install failed for package: {} \ \n Output from execution of subprocess: {}" \ .format(traceback.format_exc(), args[0][6], error. output)) from error else: raise RuntimeError("Full trace: {}\ \n Output of the failed child process: {} " \ .format(traceback.format_exc(), error.output)) from error return out def check_call(*args, **kwargs): if force_shell: kwargs['shell'] = True try: out = subprocess.check_call(*args, **kwargs) except OSError as e: raise RuntimeError("Executable {} not found".format(args[0])) from e except subprocess.CalledProcessError as error: if isinstance(args, tuple) and (args[0][2] == "pip"): raise RuntimeError( \ "Full traceback: {} \n Pip install failed for package: {} \ \n Output from execution of subprocess: {}" \ .format(traceback.format_exc(), args[0][6], error.output)) from error else: raise RuntimeError("Full trace: {} \ \n Output of the failed child process: {}" \ .format(traceback.format_exc(), error.output)) from error return out def check_output(*args, **kwargs): if force_shell: kwargs['shell'] = True try: out = subprocess.check_output(*args, **kwargs) except OSError as e: raise RuntimeError("Executable {} not found".format(args[0])) from e except subprocess.CalledProcessError as error: if isinstance(args, tuple) and (args[0][2] == "pip"): raise RuntimeError( \ "Full traceback: {} \n Pip install failed for package: {} \ \n Output from execution of subprocess: {}" \ .format(traceback.format_exc(), args[0][6], error.output)) from error else: raise RuntimeError("Full trace: {}, \ output of the failed child process {} "\ .format(traceback.format_exc(), error.output)) from error return out def Popen(*args, **kwargs): if force_shell: kwargs['shell'] = True return subprocess.Popen(*args, **kwargs) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/processes_test.py0000644000175100017510000001772515206113033022451 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the processes module.""" # pytype: skip-file import subprocess import unittest import mock from apache_beam.utils import processes class Exec(unittest.TestCase): def setUp(self): pass @mock.patch('apache_beam.utils.processes.subprocess') def test_method_forwarding_not_windows(self, *unused_mocks): # Test that the correct calls are being forwarded to the subprocess module # when we are not on Windows. processes.force_shell = False processes.call(['subprocess', 'call'], shell=False, other_arg=True) processes.subprocess.call.assert_called_once_with(['subprocess', 'call'], shell=False, other_arg=True) processes.check_call(['subprocess', 'check_call'], shell=False, other_arg=True) processes.subprocess.check_call.assert_called_once_with( ['subprocess', 'check_call'], shell=False, other_arg=True) processes.check_output(['subprocess', 'check_output'], shell=False) processes.subprocess.check_output.assert_called_once_with( ['subprocess', 'check_output'], shell=False) processes.Popen(['subprocess', 'Popen'], shell=False, other_arg=True) processes.subprocess.Popen.assert_called_once_with(['subprocess', 'Popen'], shell=False, other_arg=True) @mock.patch('apache_beam.utils.processes.subprocess') def test_method_forwarding_windows(self, *unused_mocks): # Test that the correct calls are being forwarded to the subprocess module # and that the shell=True flag is added when we are on Windows. processes.force_shell = True processes.call(['subprocess', 'call'], shell=False, other_arg=True) processes.subprocess.call.assert_called_once_with(['subprocess', 'call'], shell=True, other_arg=True) processes.check_call(['subprocess', 'check_call'], shell=False, other_arg=True) processes.subprocess.check_call.assert_called_once_with( ['subprocess', 'check_call'], shell=True, other_arg=True) processes.check_output(['subprocess', 'check_output'], shell=False) processes.subprocess.check_output.assert_called_once_with( ['subprocess', 'check_output'], shell=True) processes.Popen(['subprocess', 'Popen'], shell=False, other_arg=True) processes.subprocess.Popen.assert_called_once_with(['subprocess', 'Popen'], shell=True, other_arg=True) class TestErrorHandlingCheckCall(unittest.TestCase): @classmethod def setUpClass(cls): cls.mock_get_patcher = mock.patch(\ 'apache_beam.utils.processes.subprocess.check_call') cls.mock_get = cls.mock_get_patcher.start() @classmethod def tearDownClass(cls): cls.mock_get_patcher.stop() def test_oserror_check_call(self): # Configure the mock to return a response with an OK status code. self.mock_get.side_effect = OSError("Test OSError") with self.assertRaises(RuntimeError): processes.check_call(["lls"]) def test_oserror_check_call_message(self): # Configure the mock to return a response with an OK status code. self.mock_get.side_effect = OSError() cmd = ["lls"] try: processes.check_call(cmd) except RuntimeError as error: self.assertIn('Executable {} not found'.format(str(cmd)),\ error.args[0]) def test_check_call_pip_install_non_existing_package(self): returncode = 1 package = "non-exsisting-package" cmd = ['python', '-m', 'pip', 'download', '--dest', '/var',\ '{}'.format(package),\ '--no-deps', '--no-binary', ':all:'] output = "Collecting {}".format(package) self.mock_get.side_effect = subprocess.CalledProcessError(returncode,\ cmd, output=output) try: output = processes.check_call(cmd) self.fail( "The test failed due to that\ no error was raised when calling process.check_call") except RuntimeError as error: self.assertIn("Output from execution of subprocess: {}".format(output),\ error.args[0]) self.assertIn("Pip install failed for package: {}".format(package),\ error.args[0]) class TestErrorHandlingCheckOutput(unittest.TestCase): @classmethod def setUpClass(cls): cls.mock_get_patcher = mock.patch(\ 'apache_beam.utils.processes.subprocess.check_output') cls.mock_get = cls.mock_get_patcher.start() @classmethod def tearDownClass(cls): cls.mock_get_patcher.stop() def test_oserror_check_output_message(self): self.mock_get.side_effect = OSError() cmd = ["lls"] try: processes.check_output(cmd) except RuntimeError as error: self.assertIn('Executable {} not found'.format(str(cmd)),\ error.args[0]) def test_check_output_pip_install_non_existing_package(self): returncode = 1 package = "non-exsisting-package" cmd = ['python', '-m', 'pip', 'download', '--dest', '/var',\ '{}'.format(package),\ '--no-deps', '--no-binary', ':all:'] output = "Collecting {}".format(package) self.mock_get.side_effect = subprocess.CalledProcessError(returncode,\ cmd, output=output) try: output = processes.check_output(cmd) self.fail( "The test failed due to that\ no error was raised when calling process.check_call") except RuntimeError as error: self.assertIn("Output from execution of subprocess: {}".format(output),\ error.args[0]) self.assertIn("Pip install failed for package: {}".format(package),\ error.args[0]) class TestErrorHandlingCall(unittest.TestCase): @classmethod def setUpClass(cls): cls.mock_get_patcher = mock.patch(\ 'apache_beam.utils.processes.subprocess.call') cls.mock_get = cls.mock_get_patcher.start() @classmethod def tearDownClass(cls): cls.mock_get_patcher.stop() def test_oserror_check_output_message(self): self.mock_get.side_effect = OSError() cmd = ["lls"] try: processes.call(cmd) except RuntimeError as error: self.assertIn('Executable {} not found'.format(str(cmd)),\ error.args[0]) def test_check_output_pip_install_non_existing_package(self): returncode = 1 package = "non-exsisting-package" cmd = ['python', '-m', 'pip', 'download', '--dest', '/var',\ '{}'.format(package),\ '--no-deps', '--no-binary', ':all:'] output = "Collecting {}".format(package) self.mock_get.side_effect = subprocess.CalledProcessError(returncode,\ cmd, output=output) try: output = processes.call(cmd) self.fail( "The test failed due to that\ no error was raised when calling process.check_call") except RuntimeError as error: self.assertIn("Output from execution of subprocess: {}".format(output),\ error.args[0]) self.assertIn("Pip install failed for package: {}".format(package),\ error.args[0]) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/profiler.py0000644000175100017510000001333515206113033021217 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """A profiler context manager based on cProfile.Profile and guppy.hpy objects. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file # mypy: check-untyped-defs import cProfile import io import logging import os import pstats import random import tempfile import time from typing import Callable from typing import Optional from apache_beam.io import filesystems _LOGGER = logging.getLogger(__name__) class Profile(object): """cProfile and Heapy wrapper context for saving and logging profiler results.""" SORTBY = 'cumulative' profile_output: str stats: pstats.Stats def __init__( self, profile_id: str, profile_location: Optional[str] = None, log_results: bool = False, file_copy_fn: Optional[Callable[[str, str], None]] = None, time_prefix: str = '%Y-%m-%d_%H_%M_%S-', enable_cpu_profiling: bool = False, enable_memory_profiling: bool = False, ): """Creates a Profile object. Args: profile_id: Unique id of the profiling session. profile_location: The file location where the profiling results will be stored. log_results: Log the result to console if true. file_copy_fn: Lambda function for copying files. time_prefix: Format of the timestamp prefix in profiling result files. enable_cpu_profiling: CPU profiler will be enabled during the profiling session. enable_memory_profiling: Memory profiler will be enabled during the profiling session, the profiler only records the newly allocated objects in this session. """ self.profile_id = str(profile_id) self.profile_location = profile_location self.log_results = log_results self.file_copy_fn = file_copy_fn or self.default_file_copy_fn self.time_prefix = time_prefix self.enable_cpu_profiling = enable_cpu_profiling self.enable_memory_profiling = enable_memory_profiling def __enter__(self): _LOGGER.info('Start profiling: %s', self.profile_id) if self.enable_cpu_profiling: self.profile = cProfile.Profile() self.profile.enable() if self.enable_memory_profiling: try: from guppy import hpy self.hpy = hpy() self.hpy.setrelheap() except ImportError: _LOGGER.info("Unable to import guppy for memory profiling") self.hpy = None return self def __exit__(self, *args): _LOGGER.info('Stop profiling: %s', self.profile_id) if self.profile_location: if self.enable_cpu_profiling: self.profile.create_stats() self.profile_output = self._upload_profile_data( # typing: seems stats attr is missing from typeshed self.profile_location, 'cpu_profile', self.profile.stats) if self.enable_memory_profiling: if not self.hpy: pass else: h = self.hpy.heap() heap_dump_data = '%s\n%s' % (h, h.more) self._upload_profile_data( self.profile_location, 'memory_profile', heap_dump_data, write_binary=False) if self.log_results: if self.enable_cpu_profiling: s = io.StringIO() self.stats = pstats.Stats( self.profile, stream=s).sort_stats(Profile.SORTBY) self.stats.print_stats() _LOGGER.info('Cpu profiler data: [%s]', s.getvalue()) if self.enable_memory_profiling and self.hpy: _LOGGER.info('Memory profiler data: \n%s' % self.hpy.heap()) @staticmethod def default_file_copy_fn(src, dest): dest_handle = filesystems.FileSystems.create(dest + '.tmp') try: with open(src, 'rb') as src_handle: dest_handle.write(src_handle.read()) finally: dest_handle.close() filesystems.FileSystems.rename([dest + '.tmp'], [dest]) @staticmethod def factory_from_options(options) -> Optional[Callable[..., 'Profile']]: if options.profile_cpu or options.profile_memory: def create_profiler(profile_id, **kwargs): if random.random() < options.profile_sample_rate: return Profile( profile_id, options.profile_location, enable_cpu_profiling=options.profile_cpu, enable_memory_profiling=options.profile_memory, **kwargs) return create_profiler return None def _upload_profile_data( self, profile_location, dir, data, write_binary=True) -> str: dump_location = os.path.join( profile_location, dir, time.strftime(self.time_prefix + self.profile_id)) fd, filename = tempfile.mkstemp() try: os.close(fd) if write_binary: with open(filename, 'wb') as fb: import marshal marshal.dump(data, fb) else: with open(filename, 'w') as f: f.write(data) _LOGGER.info('Copying profiler data to: [%s]', dump_location) self.file_copy_fn(filename, dump_location) finally: os.remove(filename) return dump_location ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/profiler_test.py0000644000175100017510000000405115206113033022251 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import os import tempfile import unittest from parameterized import param from parameterized import parameterized from apache_beam.utils.profiler import Profile class ProfilerTest(unittest.TestCase): @parameterized.expand([ param(enable_cpu_memory=(True, True)), param(enable_cpu_memory=(True, False)), param(enable_cpu_memory=(False, True)), param(enable_cpu_memory=(False, False)), ]) def test_profiler(self, enable_cpu_memory): try: from guppy import hpy # pylint: disable=unused-import guppy_imported = True except ImportError: guppy_imported = False enable_cpu, enable_memory = enable_cpu_memory with tempfile.TemporaryDirectory() as tmp: with Profile('id', profile_location=tmp, enable_cpu_profiling=enable_cpu, enable_memory_profiling=enable_memory, log_results=True): hash('test') files = os.listdir(tmp) if enable_cpu: self.assertIn('cpu_profile', files) else: self.assertNotIn('cpu_profile', files) if enable_memory and guppy_imported: self.assertIn('memory_profile', files) else: self.assertNotIn('memory_profile', files) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/proto_utils.py0000644000175100017510000001072115206113033021754 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """For internal use only; no backwards-compatibility guarantees.""" # pytype: skip-file from typing import TypeVar from typing import Union from typing import overload from google.protobuf import any_pb2 from google.protobuf import duration_pb2 from google.protobuf import message from google.protobuf import struct_pb2 from google.protobuf import timestamp_pb2 MessageT = TypeVar('MessageT', bound=message.Message) TimeMessageT = TypeVar( 'TimeMessageT', duration_pb2.Duration, timestamp_pb2.Timestamp) message_types = (message.Message, ) _SECONDS_TO_MICROS = 10**6 _MICROS_TO_NANOS = 10**3 @overload def pack_Any(msg: message.Message) -> any_pb2.Any: pass @overload def pack_Any(msg: None) -> None: # type: ignore[overload-cannot-match] pass def pack_Any(msg): """Creates a protobuf Any with msg as its content. Returns None if msg is None. """ if msg is None: return None result = any_pb2.Any() result.Pack(msg) return result @overload def unpack_Any(any_msg: any_pb2.Any, msg_class: type[MessageT]) -> MessageT: pass @overload def unpack_Any(any_msg: any_pb2.Any, msg_class: None) -> None: pass def unpack_Any(any_msg, msg_class): """Unpacks any_msg into msg_class. Returns None if msg_class is None. """ if msg_class is None: return None msg = msg_class() any_msg.Unpack(msg) return msg @overload def parse_Bytes(serialized_bytes: bytes, msg_class: type[MessageT]) -> MessageT: pass @overload def parse_Bytes( serialized_bytes: bytes, msg_class: Union[type[bytes], None]) -> bytes: pass def parse_Bytes(serialized_bytes, msg_class): """Parses the String of bytes into msg_class. Returns the input bytes if msg_class is None.""" if msg_class is None or msg_class is bytes: return serialized_bytes msg = msg_class() msg.ParseFromString(serialized_bytes) return msg def pack_Struct(**kwargs) -> struct_pb2.Struct: """Returns a struct containing the values indicated by kwargs. """ msg = struct_pb2.Struct() for key, value in kwargs.items(): msg[key] = value # pylint: disable=unsubscriptable-object, unsupported-assignment-operation return msg def from_micros(cls: type[TimeMessageT], micros: int) -> TimeMessageT: result = cls() if isinstance(result, duration_pb2.Duration): result.FromMicroseconds(micros) return result # Protobuf 5.x enforces a maximum timestamp value less than the Beam # maximum allowable timestamp, so we cannot use the built-in conversion. elif isinstance(result, timestamp_pb2.Timestamp): result.seconds = micros // _SECONDS_TO_MICROS result.nanos = (micros % _SECONDS_TO_MICROS) * _MICROS_TO_NANOS return result else: raise RuntimeError('cannot convert the micro seconds to %s' % cls) def to_micros(value: Union[duration_pb2.Duration, timestamp_pb2.Timestamp]): if isinstance(value, duration_pb2.Duration): return value.ToMicroseconds() # Protobuf 5.x enforces a maximum timestamp value less than the Beam # maximum allowable timestamp, so we cannot use the built-in conversion. elif isinstance(value, timestamp_pb2.Timestamp): micros = value.seconds * _SECONDS_TO_MICROS return micros + (value.nanos // _MICROS_TO_NANOS) else: raise RuntimeError('cannot convert %s to micro seconds' % value) def to_Timestamp(time: Union[int, float]) -> timestamp_pb2.Timestamp: """Convert a float returned by time.time() to a Timestamp. """ seconds = int(time) nanos = int((time - seconds) * 10**9) return timestamp_pb2.Timestamp(seconds=seconds, nanos=nanos) def from_Timestamp(timestamp: timestamp_pb2.Timestamp) -> float: """Convert a Timestamp to a float expressed as seconds since the epoch. """ return timestamp.seconds + float(timestamp.nanos) / 10**9 ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/proto_utils_test.py0000644000175100017510000000462515206113033023021 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest from google.protobuf import duration_pb2 from google.protobuf import timestamp_pb2 from apache_beam.utils import proto_utils from apache_beam.utils.timestamp import MAX_TIMESTAMP class TestProtoUtils(unittest.TestCase): def test_from_micros_duration(self): ts = proto_utils.from_micros(duration_pb2.Duration, MAX_TIMESTAMP.micros) expected = duration_pb2.Duration( seconds=MAX_TIMESTAMP.seconds(), nanos=775000000) self.assertEqual(ts, expected) def test_from_micros_timestamp(self): ts = proto_utils.from_micros(timestamp_pb2.Timestamp, MAX_TIMESTAMP.micros) expected = timestamp_pb2.Timestamp( seconds=MAX_TIMESTAMP.seconds(), nanos=775000000) self.assertEqual(ts, expected) def test_to_micros_duration(self): dur = duration_pb2.Duration( seconds=MAX_TIMESTAMP.seconds(), nanos=775000000) ts = proto_utils.to_micros(dur) expected = MAX_TIMESTAMP.micros self.assertEqual(ts, expected) def test_to_micros_timestamp(self): dur = timestamp_pb2.Timestamp( seconds=MAX_TIMESTAMP.seconds(), nanos=775000000) ts = proto_utils.to_micros(dur) expected = MAX_TIMESTAMP.micros self.assertEqual(ts, expected) def test_round_trip_duration(self): expected = 919336704 dur = proto_utils.from_micros(duration_pb2.Duration, expected) ms = proto_utils.to_micros(dur) self.assertEqual(ms, expected) def test_round_trip_timestamp(self): expected = 919336704 ts = proto_utils.from_micros(timestamp_pb2.Timestamp, expected) ms = proto_utils.to_micros(ts) self.assertEqual(ms, expected) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/python_callable.py0000644000175100017510000001036415206113033022534 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Python Callable utilities. For internal use only; no backwards-compatibility guarantees. """ import importlib class PythonCallableWithSource(object): """Represents a Python callable object with source codes before evaluated. Proxy object to Store a callable object with its string form (source code). The string form is used when the object is encoded and transferred to foreign SDKs (non-Python SDKs). Supported formats include fully-qualified names such as `math.sin`, expressions such as `lambda x: x * x` or `str.upper`, and multi-line function definitions such as `def foo(x): ...` or class definitions like `class Foo(...): ...`. If the source string contains multiple lines then lines prior to the last will be evaluated to provide the context in which to evaluate the expression, for example:: import math lambda x: x - math.sin(x) is a valid chunk of source code. """ def __init__(self, source: str) -> None: self._source = source self._callable = self.load_from_source(source) @classmethod def load_from_source(cls, source): if source in __builtins__: return cls.load_from_expression(source) elif all(s.isidentifier() for s in source.split('.')): if source.split('.')[0] in __builtins__: return cls.load_from_expression(source) else: return cls.load_from_fully_qualified_name(source) else: return cls.load_from_script(source) @staticmethod def load_from_expression(source): return eval(source) # pylint: disable=eval-used @staticmethod def load_from_fully_qualified_name(fully_qualified_name): o = None path = '' for segment in fully_qualified_name.split('.'): path = '.'.join([path, segment]) if path else segment if o is not None and hasattr(o, segment): o = getattr(o, segment) else: o = importlib.import_module(path) return o @staticmethod def load_from_script(source, method_name=None): lines = [ line for line in source.split('\n') if line.strip() and line.strip()[0] != '#' ] common_indent = min(len(line) - len(line.lstrip()) for line in lines) lines = [line[common_indent:] for line in lines] if method_name is None: for ix, line in reversed(list(enumerate(lines))): if line[0] != ' ': if line.startswith('def '): method_name = line[4:line.index('(')].strip() elif line.startswith('class '): method_name = line[5:line.index('(') if '(' in line else line.index(':')].strip() else: method_name = '__python_callable__' lines[ix] = method_name + ' = ' + line break else: raise ValueError("Unable to identify callable from %r" % source) # pylint: disable=exec-used # pylint: disable=ungrouped-imports import apache_beam as beam exec_globals = {'beam': beam} exec('\n'.join(lines), exec_globals) return exec_globals[method_name] def default_label(self): src = self._source.strip() last_line = src.split('\n')[-1] if last_line[0] != ' ' and len(last_line) < 72: return last_line # Avoid circular import. from apache_beam.transforms.ptransform import label_from_callable return label_from_callable(self._callable) @property def _argspec_fn(self): return self._callable def get_source(self) -> str: return self._source def __call__(self, *args, **kwargs): return self._callable(*args, **kwargs) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/python_callable_test.py0000644000175100017510000000531615206113033023574 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os import unittest import apache_beam as beam from apache_beam.utils.python_callable import PythonCallableWithSource class PythonCallableWithSourceTest(unittest.TestCase): def test_builtin(self): self.assertEqual(PythonCallableWithSource.load_from_source('str'), str) def test_builtin_attribute(self): self.assertEqual( PythonCallableWithSource.load_from_source('str.lower'), str.lower) def test_fully_qualified_name(self): self.assertEqual( PythonCallableWithSource.load_from_source('os.path.abspath'), os.path.abspath) def test_expression(self): self.assertEqual(PythonCallableWithSource('lambda x: x*x')(10), 100) def test_expression_with_dependency(self): self.assertEqual( PythonCallableWithSource('import math\nlambda x: math.sqrt(x) + x')( 100), 110) def test_def(self): self.assertEqual( PythonCallableWithSource( """ def foo(x): return x * x """)(10), 100) def test_def_with_preamble(self): self.assertEqual( PythonCallableWithSource( """ def bar(x): return x + 1 def foo(x): return bar(x) * x """)(10), 110) def test_class(self): self.assertEqual( PythonCallableWithSource( """ class BareClass: def __init__(self, x): self.x = x """)(10).x, 10) self.assertEqual( PythonCallableWithSource( """ class SubClass(object): def __init__(self, x): self.x = x """)(10).x, 10) def test_pycallable_map(self): p = beam.Pipeline() result = ( p | beam.Create([1, 2, 3]) | beam.Map(PythonCallableWithSource("lambda x: x"))) self.assertEqual(result.element_type, int) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/retry.py0000644000175100017510000003165315206113033020545 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Retry decorators for calls raising exceptions. For internal use only; no backwards-compatibility guarantees. This module is used mostly to decorate all integration points where the code makes calls to remote services. Searching through the code base for @retry should find all such places. For this reason even places where retry is not needed right now use a @retry.no_retries decorator. """ # pytype: skip-file import functools import json import logging import random import sys import time import traceback from apache_beam.io.filesystem import BeamIOError # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position # TODO(sourabhbajaj): Remove the GCP specific error code to a submodule try: from apitools.base.py.exceptions import HttpError from google.api_core.exceptions import GoogleAPICallError except ImportError as e: HttpError = None GoogleAPICallError = None # type: ignore # Protect against environments where aws tools are not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: from apache_beam.io.aws.clients.s3 import messages as _s3messages except ImportError: S3ClientError = None else: S3ClientError = _s3messages.S3ClientError # pylint: enable=wrong-import-order, wrong-import-position _LOGGER = logging.getLogger(__name__) _RETRYABLE_REASONS = ["rateLimitExceeded", "internalError", "backendError"] class PermanentException(Exception): """Base class for exceptions that should not be retried.""" pass class FuzzedExponentialIntervals(object): """Iterable for intervals that are exponentially spaced, with fuzzing. On iteration, yields retry interval lengths, in seconds. Every iteration over this iterable will yield differently fuzzed interval lengths, as long as fuzz is nonzero. Args: initial_delay_secs: The delay before the first retry, in seconds. num_retries: The total number of times to retry. factor: The exponential factor to use on subsequent retries. Default is 2 (doubling). fuzz: A value between 0 and 1, indicating the fraction of fuzz. For a given delay d, the fuzzed delay is randomly chosen between [(1 - fuzz) * d, d]. max_delay_secs: Maximum delay (in seconds). After this limit is reached, further tries use max_delay_sec instead of exponentially increasing the time. Defaults to 1 hour. stop_after_secs: Places a limit on the sum of intervals returned (in seconds), such that the sum is <= stop_after_secs. Defaults to disabled (None). You may need to increase num_retries to effectively use this feature. """ def __init__( self, initial_delay_secs, num_retries, factor=2, fuzz=0.5, max_delay_secs=60 * 60 * 1, stop_after_secs=None): self._initial_delay_secs = initial_delay_secs if num_retries > 10000: raise ValueError('num_retries parameter cannot exceed 10000.') self._num_retries = num_retries self._factor = factor if not 0 <= fuzz <= 1: raise ValueError('fuzz parameter expected to be in [0, 1] range.') self._fuzz = fuzz self._max_delay_secs = max_delay_secs self._stop_after_secs = stop_after_secs def __iter__(self): current_delay_secs = min(self._max_delay_secs, self._initial_delay_secs) total_delay_secs = 0 for _ in range(self._num_retries): fuzz_multiplier = 1 - self._fuzz + random.random() * self._fuzz delay_secs = current_delay_secs * fuzz_multiplier total_delay_secs += delay_secs if (self._stop_after_secs is not None and total_delay_secs > self._stop_after_secs): break yield delay_secs current_delay_secs = min( self._max_delay_secs, current_delay_secs * self._factor) def retry_on_server_errors_filter(exception): """Filter allowing retries on server errors and non-HttpErrors.""" if (HttpError is not None) and isinstance(exception, HttpError): return exception.status_code >= 500 if GoogleAPICallError is not None and isinstance(exception, GoogleAPICallError): if exception.code >= 500: # 500 are internal server errors return True else: # If we have a GoogleAPICallError with a code that doesn't # indicate a server error, we do not need to retry. return False if (S3ClientError is not None) and isinstance(exception, S3ClientError): return exception.code is None or exception.code >= 500 return not isinstance(exception, PermanentException) # TODO(https://github.com/apache/beam/issues/19350): Dataflow returns 404 for # job ids that actually exist. Retry on those errors. def retry_on_server_errors_and_notfound_filter(exception): if HttpError is not None and isinstance(exception, HttpError): if exception.status_code == 404: # 404 Not Found return True if GoogleAPICallError is not None and isinstance(exception, GoogleAPICallError): if exception.code == 404: # 404 Not found return True return retry_on_server_errors_filter(exception) def retry_on_server_errors_and_timeout_filter(exception): if HttpError is not None and isinstance(exception, HttpError): if exception.status_code == 408: # 408 Request Timeout return True if GoogleAPICallError is not None and isinstance(exception, GoogleAPICallError): if exception.code == 408: # 408 Request Timeout return True if S3ClientError is not None and isinstance(exception, S3ClientError): if exception.code == 408: # 408 Request Timeout return True return retry_on_server_errors_filter(exception) def retry_on_server_errors_timeout_or_quota_issues_filter(exception): """Retry on server, timeout, 429, and some 403 errors. 403 errors from BigQuery include both non-transient (accessDenied, billingNotEnabled) and transient errors (rateLimitExceeded). Only retry transient errors.""" if HttpError is not None and isinstance(exception, HttpError): if exception.status_code == 429: return True if exception.status_code == 403: try: # attempt to extract the reason and check if it's retryable content = exception.content if not isinstance(content, dict): content = json.loads(exception.content) return content["error"]["errors"][0]["reason"] in _RETRYABLE_REASONS except (KeyError, IndexError, TypeError) as e: _LOGGER.warning( "Could not determine if HttpError is transient. " "Will not retry: %s", e) return False if GoogleAPICallError is not None and isinstance(exception, GoogleAPICallError): if exception.code == 429: return True if exception.code == 403: if not hasattr(exception, "errors") or len(exception.errors) == 0: # default to not retrying return False reason = exception.errors[0]["reason"] return reason in _RETRYABLE_REASONS if S3ClientError is not None and isinstance(exception, S3ClientError): if exception.code == 403: return True return retry_on_server_errors_and_timeout_filter(exception) def retry_on_beam_io_error_filter(exception): """Filter allowing retries on Beam IO errors.""" return isinstance(exception, BeamIOError) def retry_if_valid_input_but_server_error_and_timeout_filter(exception): if isinstance(exception, ValueError): return False return retry_on_server_errors_and_timeout_filter(exception) SERVER_ERROR_OR_TIMEOUT_CODES = [408, 500, 502, 503, 504, 598, 599] class Clock(object): """A simple clock implementing sleep().""" def sleep(self, value): time.sleep(value) def no_retries(fun): """A retry decorator for places where we do not want retries.""" return with_exponential_backoff(retry_filter=lambda _: False, clock=None)(fun) def with_exponential_backoff( num_retries=7, initial_delay_secs=5.0, logger=_LOGGER.warning, retry_filter=retry_on_server_errors_filter, clock=Clock(), fuzz=True, factor=2, max_delay_secs=60 * 60, stop_after_secs=None): """Decorator with arguments that control the retry logic. Args: num_retries: The total number of times to retry. initial_delay_secs: The delay before the first retry, in seconds. logger: A callable used to report an exception. Must have the same signature as functions in the standard logging module. The default is _LOGGER.warning. retry_filter: A callable getting the exception raised and returning True if the retry should happen. For instance we do not want to retry on 404 Http errors most of the time. The default value will return true for server errors (HTTP status code >= 500) and non Http errors. clock: A clock object implementing a sleep method. The default clock will use time.sleep(). fuzz: True if the delay should be fuzzed (default). During testing False can be used so that the delays are not randomized. factor: The exponential factor to use on subsequent retries. Default is 2 (doubling). max_delay_secs: Maximum delay (in seconds). After this limit is reached, further tries use max_delay_sec instead of exponentially increasing the time. Defaults to 1 hour. stop_after_secs: Places a limit on the sum of delays between retries, such that the sum is <= stop_after_secs. Retries will stop after the limit is reached. Defaults to disabled (None). You may need to increase num_retries to effectively use this feature. Returns: As per Python decorators with arguments pattern returns a decorator for the function which in turn will return the wrapped (decorated) function. The decorator is intended to be used on callables that make HTTP or RPC requests that can temporarily timeout or have transient errors. For instance the make_http_request() call below will be retried 16 times with exponential backoff and fuzzing of the delay interval (default settings). The callable should return values directly instead of yielding them, as generators are not evaluated within the try-catch block and will not be retried on exception. from apache_beam.utils import retry # ... @retry.with_exponential_backoff() make_http_request(args) """ def real_decorator(fun): """The real decorator whose purpose is to return the wrapped function.""" @functools.wraps(fun) def wrapper(*args, **kwargs): retry_intervals = iter( FuzzedExponentialIntervals( initial_delay_secs, num_retries, factor, fuzz=0.5 if fuzz else 0, max_delay_secs=max_delay_secs, stop_after_secs=stop_after_secs)) while True: try: return fun(*args, **kwargs) except Exception as exn: # pylint: disable=broad-except if not retry_filter(exn): raise # Get the traceback object for the current exception. The # sys.exc_info() function returns a tuple with three elements: # exception type, exception value, and exception traceback. exn_traceback = sys.exc_info()[2] try: try: sleep_interval = next(retry_intervals) except StopIteration: # Re-raise the original exception since we finished the retries. raise exn.with_traceback(exn_traceback) logger( 'Retry with exponential backoff: waiting for %s seconds before ' 'retrying %s because we caught exception: %s ' 'Traceback for above exception (most recent call last):\n%s', sleep_interval, getattr(fun, '__name__', str(fun)), ''.join(traceback.format_exception_only(exn.__class__, exn)), ''.join(traceback.format_tb(exn_traceback))) clock.sleep(sleep_interval) finally: # Traceback objects in locals can cause reference cycles that will # prevent garbage collection. Clear it now since we do not need # it anymore. exn_traceback = None return wrapper return real_decorator ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/retry_test.py0000644000175100017510000001772315206113033021606 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the retry module.""" # pytype: skip-file import unittest from parameterized import parameterized from apache_beam.utils import retry # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position # TODO(sourabhbajaj): Remove the GCP specific error code to a submodule try: from apitools.base.py.exceptions import HttpError except ImportError: HttpError = None # pylint: enable=wrong-import-order, wrong-import-position class FakeClock(object): """A fake clock object implementing sleep() and recording calls.""" def __init__(self): self.calls = [] def sleep(self, value): self.calls.append(value) class FakeLogger(object): """A fake logger object implementing log() and recording calls.""" def __init__(self): self.calls = [] def log(self, message, interval, func_name, exn_name, exn_traceback): _ = interval, exn_traceback self.calls.append((message, func_name, exn_name)) @retry.with_exponential_backoff(clock=FakeClock()) def _test_function(a, b): _ = a, b raise NotImplementedError @retry.with_exponential_backoff(initial_delay_secs=0.1, num_retries=1) def _test_function_with_real_clock(a, b): _ = a, b raise NotImplementedError @retry.no_retries def _test_no_retry_function(a, b): _ = a, b raise NotImplementedError class RetryTest(unittest.TestCase): def setUp(self): self.clock = FakeClock() self.logger = FakeLogger() self.calls = 0 def permanent_failure(self, a, b): raise NotImplementedError def transient_failure(self, a, b): self.calls += 1 if self.calls > 4: return a + b raise NotImplementedError def http_error(self, code): if HttpError is None: raise RuntimeError("This is not a valid test as GCP is not enabled") raise HttpError({'status': str(code)}, '', '') def test_with_explicit_decorator(self): # We pass one argument as positional argument and one as keyword argument # so that we cover both code paths for argument handling. self.assertRaises(NotImplementedError, _test_function, 10, b=20) def test_with_no_retry_decorator(self): self.assertRaises(NotImplementedError, _test_no_retry_function, 1, 2) def test_with_real_clock(self): self.assertRaises( NotImplementedError, _test_function_with_real_clock, 10, b=20) def test_with_default_number_of_retries(self): self.assertRaises( NotImplementedError, retry.with_exponential_backoff(clock=self.clock)( self.permanent_failure), 10, b=20) self.assertEqual(len(self.clock.calls), 7) def test_with_explicit_number_of_retries(self): self.assertRaises( NotImplementedError, retry.with_exponential_backoff(clock=self.clock, num_retries=10)(self.permanent_failure), 10, b=20) self.assertEqual(len(self.clock.calls), 10) @unittest.skipIf(HttpError is None, 'google-apitools is not installed') def test_with_http_error_that_should_not_be_retried(self): self.assertRaises( HttpError, retry.with_exponential_backoff(clock=self.clock, num_retries=10)(self.http_error), 404) # Make sure just one call was made. self.assertEqual(len(self.clock.calls), 0) @unittest.skipIf(HttpError is None, 'google-apitools is not installed') def test_with_http_error_that_should_be_retried(self): self.assertRaises( HttpError, retry.with_exponential_backoff(clock=self.clock, num_retries=10)(self.http_error), 500) self.assertEqual(len(self.clock.calls), 10) def test_with_explicit_initial_delay(self): self.assertRaises( NotImplementedError, retry.with_exponential_backoff( initial_delay_secs=10.0, clock=self.clock, fuzz=False)(self.permanent_failure), 10, b=20) self.assertEqual(len(self.clock.calls), 7) self.assertEqual(self.clock.calls[0], 10.0) @parameterized.expand([(str(i), i) for i in range(0, 1000, 47)]) def test_with_stop_after_secs(self, _, stop_after_secs): max_delay_secs = 10 self.assertRaises( NotImplementedError, retry.with_exponential_backoff( num_retries=10000, initial_delay_secs=10.0, clock=self.clock, fuzz=False, max_delay_secs=max_delay_secs, stop_after_secs=stop_after_secs)(self.permanent_failure), 10, b=20) total_delay = sum(self.clock.calls) self.assertLessEqual(total_delay, stop_after_secs) self.assertGreaterEqual(total_delay, stop_after_secs - max_delay_secs) def test_log_calls_for_permanent_failure(self): self.assertRaises( NotImplementedError, retry.with_exponential_backoff( clock=self.clock, logger=self.logger.log)(self.permanent_failure), 10, b=20) self.assertEqual(len(self.logger.calls), 7) for message, func_name, exn_name in self.logger.calls: self.assertTrue(message.startswith('Retry with exponential backoff:')) self.assertEqual(exn_name, 'NotImplementedError\n') self.assertEqual(func_name, 'permanent_failure') def test_log_calls_for_transient_failure(self): result = retry.with_exponential_backoff( clock=self.clock, logger=self.logger.log, fuzz=False)(self.transient_failure)( 10, b=20) self.assertEqual(result, 30) self.assertEqual(len(self.clock.calls), 4) self.assertEqual(self.clock.calls, [ 5.0 * 1, 5.0 * 2, 5.0 * 4, 5.0 * 8, ]) self.assertEqual(len(self.logger.calls), 4) for message, func_name, exn_name in self.logger.calls: self.assertTrue(message.startswith('Retry with exponential backoff:')) self.assertEqual(exn_name, 'NotImplementedError\n') self.assertEqual(func_name, 'transient_failure') class DummyClass(object): def __init__(self, results): self.index = 0 self.results = results @retry.with_exponential_backoff(num_retries=2, initial_delay_secs=0.1) def func(self): self.index += 1 if self.index > len(self.results) or \ self.results[self.index - 1] == "Error": raise ValueError("Error") return self.results[self.index - 1] class RetryStateTest(unittest.TestCase): """The test_two_failures and test_single_failure would fail if we have any shared state for the retry decorator. This test tries to prevent a bug we found where the state in the decorator was shared across objects and retries were not available correctly. The test_call_two_objects would test this inside the same test. """ def test_two_failures(self): dummy = DummyClass(["Error", "Error", "Success"]) dummy.func() self.assertEqual(3, dummy.index) def test_single_failure(self): dummy = DummyClass(["Error", "Success"]) dummy.func() self.assertEqual(2, dummy.index) def test_call_two_objects(self): dummy = DummyClass(["Error", "Error", "Success"]) dummy.func() self.assertEqual(3, dummy.index) dummy2 = DummyClass(["Error", "Success"]) dummy2.func() self.assertEqual(2, dummy2.index) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/sentinel.py0000644000175100017510000000161715206113033021216 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import enum class Sentinel(enum.Enum): """ A type-safe sentinel class """ sentinel = object() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/sharded_key.py0000644000175100017510000000313615206113033021655 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file class ShardedKey(object): """ A sharded key consisting of a user key and an opaque shard id represented by bytes. Attributes: key: The user key. shard_id: An opaque byte string that uniquely represents a shard of the key. """ def __init__( self, key, shard_id: bytes, ) -> None: assert shard_id is not None self._key = key self._shard_id = shard_id @property def key(self): return self._key def __repr__(self): return '(%s, %s)' % (repr(self.key), self._shard_id) def __eq__(self, other): return ( type(self) == type(other) and self.key == other.key and self._shard_id == other._shard_id) def __hash__(self): return hash((self.key, self._shard_id)) def __reduce__(self): return ShardedKey, (self.key, self._shard_id) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/shared.py0000644000175100017510000002636715206113033020654 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Shared class. Shared is a helper class for managing a single instance of an object shared by multiple threads within the same process. Instances of Shared are serializable objects that can be shared by all threads of each worker process. A Shared object encapsulates a weak reference to a singleton instance of the shared resource. The singleton is lazily initialized by calls to Shared.acquire(). Example usage: To share a very large list across all threads of each worker in a DoFn:: # Several built-in types such as list and dict do not directly support weak # references but can add support through subclassing: # https://docs.python.org/3/library/weakref.html class WeakRefList(list): pass class GetNthStringFn(beam.DoFn): def __init__(self): self._shared_handle = shared.Shared() def setup(self): # setup is a good place to initialize transient in-memory resources. def initialize_list(): # Build the giant initial list. return WeakRefList([str(i) for i in range(1000000)]) self._giant_list = self._shared_handle.acquire(initialize_list) def process(self, element): yield self._giant_list[element] p = beam.Pipeline() (p | beam.Create([2, 4, 6, 8]) | beam.ParDo(GetNthStringFn())) Real-world uses will typically involve using a side-input to a DoFn to initialize the shared resource in a way that can't be done with just its constructor:: class RainbowTableLookupFn(beam.DoFn): def __init__(self): self._shared_handle = shared.Shared() def process(self, element, table_elements): def construct_table(): # Construct the rainbow table from the table elements. # The table contains lines in the form "string::hash" result = {} for key, value in table_elements: result[value] = key return result rainbow_table = self._shared_handle.acquire(construct_table) unhashed_str = rainbow_table.get(element) if unhashed_str is not None: yield unhashed_str p = beam.Pipeline() reverse_hash_table = p | "ReverseHashTable" >> beam.Create([ ('a', '0cc175b9c0f1b6a831c399e269772661'), ('b', '92eb5ffee6ae2fec3ad71c777531578f'), ('c', '4a8a08f09d37b73795649038408b5f33'), ('d', '8277e0910d750195b448797616e091ad')]) unhashed = (p | 'Hashes' >> beam.Create([ '0cc175b9c0f1b6a831c399e269772661', '8277e0910d750195b448797616e091ad']) | 'Unhash' >> beam.ParDo( RainbowTableLookupFn(), reverse_hash_table)) """ import threading import uuid import weakref from typing import Any from typing import Callable from typing import Text class _SharedControlBlock(object): """Wrapper class for holding objects in the SharedMap. We need this so we can call constructors for distinct Shared elements in the SharedMap concurrently. """ def __init__(self): self._lock = threading.Lock() self._ref = None self._tag = None def acquire(self, constructor_fn: Callable[[], Any], tag: Any = None) -> Any: """Acquire a reference to the object this shared control block manages. Args: constructor_fn: function that initialises / constructs the object if not present in the cache. This function should take no arguments. It should return an initialised object, or None if the object could not be initialised / constructed. tag: an optional indentifier to store with the cached object. If subsequent calls to acquire use different tags, the object will be reloaded rather than returned from cache. Returns: An initialised object, either from a previous initialisation, or newly-constructed. """ with self._lock: # self._ref is None if this is a new control block. # self._ref() is None if the weak reference was GCed. # self._tag != tag if user specifies a new identifier if self._ref is None or self._ref() is None or self._tag != tag: result = constructor_fn() if result is None: return None self._ref = weakref.ref(result) self._tag = tag else: result = self._ref() return result class _SharedMap(object): """Map for storing objects pointed to by Shared. The behaviour of SharedMap is as follows: when acquire is called, if the Shared object has already been initialised, we return the already-initialised copy. If not, we call the constructor_fn to construct it, and store it in the cache. One big caveat is this: we want to support cases where there is some delay between reacquistion of Shared objects, i.e. there may be a short period of time in which there are no references to the object before it is reacquired. This happens in various Beam runners (e.g. Dataflow runner): if we use a single thread for doing predictions with a large model, when the thread finishes its workitem, it will release the reference to the model. Since there's only a single thread, the model will have zero references to it and will be garbage collected. Shortly after this, the process receives a new workitem, creates a new thread, and attempts to reacquire the model. If we don't keep the model alive in between, the new thread will have to reinitialise the model from scratch. As such, we need to do some extra work to manage cached objects' lifetime. Ideally we would want to release the shared objects once the stage is complete, but we don't have information about that. As such, we work around this limitation as follows: when an object is first initialised, we create and maintain an explicit reference to it. This means that it will always have one reference to it from within _SharedMap. When acquire is called for a *different* object, we delete explicit references to *all other objects*. This means that if there are no external references to these objects, they will be garbage collected. This has the following implications: * A shared object won't be GC'ed if there isn't another acquire called for a different shared object. This is okay for our use-cases. This means that the shared object will be kept alive for all stages fused with the stage that works with the shared object. However, all these stages would be allocated the same memory anyway, even if the shared object were released after the stage that uses it was done with it. * Each stage can only use exactly one Shared token, otherwise only one Shared token, *NOT NECESSARILY THE LATEST*, will be "kept-alive" (using multiple shared tokens per-stage won't affect correctness, but will have no performance benefit either) * If there are two different stages using separate Shared tokens, but which get fused together, only one Shared token will be "kept-alive". This effectively means that the Shared tokens do nothing: since S2 displaces S1, and after S2 executes a new thread is created starting with S1 again, which displaces S2. Related issues: BEAM-562 - DoFn reuse """ def __init__(self): # Lock that protects cache_map self._lock = threading.Lock() # Dictionary of references to shared control blocks self._cache_map = {} # Tuple of (key, obj), where obj is an object we explicitly hold a reference # to keep it alive self._keepalive = (None, None) def make_key(self) -> Text: return str(uuid.uuid1()) def acquire( self, key: Text, constructor_fn: Callable[[], Any], tag: Any = None) -> Any: """Acquire a reference to a Shared object. Args: key: the key to the shared object constructor_fn: function that initialises / constructs the object if not present in the cache. This function should take no arguments. It should return an initialised object, or None if the object could not be initialised / constructed. tag: an optional indentifier to store with the cached object. If subsequent calls to acquire use different tags, the object will be reloaded rather than returned from cache. Returns: A reference to the initialised object, either from the cache, or newly-constructed. """ with self._lock: control_block = self._cache_map.get(key) if control_block is None: control_block = _SharedControlBlock() self._cache_map[key] = control_block result = control_block.acquire(constructor_fn, tag) # Because we release the lock in between, if we acquire multiple Shareds # in a short time, there's no guarantee as to which one will be kept alive. with self._lock: self._keepalive = (key, result) return result # Instance of the shared map to be used with Shared objects. _shared_map = _SharedMap() class Shared(object): """Handle for managing shared per-process objects. Each instance of a Shared object represents a distinct handle to a distinct object. Example usage is described in the file comment of shared.py. This object has the following limitations: * A shared object won't be GC'ed if there isn't another acquire called for a different shared object. * Each stage can only use exactly one Shared token, otherwise only one Shared token, *NOT NECESSARILY THE LATEST*, will be "kept-alive". * If there are two different stages using separate Shared tokens, but which get fused together, only one Shared token will be "kept-alive". (See documentation of _SharedMap for details.) """ # TODO(altay): Consider allowing users to also pass in a key (GUID) # for more easily sharing of identifiable expensive objects. User would be # responsible for handling collisions. def __init__(self): self._key = _shared_map.make_key() def acquire(self, constructor_fn: Callable[[], Any], tag: Any = None) -> Any: """Acquire a reference to the object associated with this Shared handle. Args: constructor_fn: function that initialises / constructs the object if not present in the cache. This function should take no arguments. It should return an initialised object, or None if the object could not be initialised / constructed. tag: an optional indentifier to store with the cached object. If subsequent calls to acquire use different tags, the object will be reloaded rather than returned from cache. Returns: A reference to an initialised object, either from the cache, or newly-constructed. """ return _shared_map.acquire(self._key, constructor_fn, tag) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/shared_test.py0000644000175100017510000001615615206113033021706 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Test for Shared class.""" import gc import threading import time import unittest from apache_beam.utils import shared class Count(object): def __init__(self): self._lock = threading.Lock() self._total = 0 self._active = 0 def add_ref(self): with self._lock: self._total += 1 self._active += 1 def release_ref(self): with self._lock: self._active -= 1 def get_active(self): with self._lock: return self._active def get_total(self): with self._lock: return self._total class Marker(object): def __init__(self, count): self._count = count self._count.add_ref() def __del__(self): self._count.release_ref() class NamedObject(object): def __init__(self, name): self._name = name def get_name(self): return self._name class Sequence(object): def __init__(self): self._sequence = 0 def make_acquire_fn(self): # Every time acquire_fn is called, increases the sequence number and returns # a NamedObject with that sequenece number. def acquire_fn(): self._sequence += 1 return NamedObject('sequence%d' % self._sequence) return acquire_fn class SharedTest(unittest.TestCase): def testKeepalive(self): count = Count() shared_handle = shared.Shared() other_shared_handle = shared.Shared() def dummy_acquire_fn(): return None def acquire_fn(): return Marker(count) p1 = shared_handle.acquire(acquire_fn) self.assertEqual(1, count.get_total()) self.assertEqual(1, count.get_active()) del p1 gc.collect() # Won't be garbage collected, because of the keep-alive self.assertEqual(1, count.get_active()) # Reacquire. p2 = shared_handle.acquire(acquire_fn) self.assertEqual(1, count.get_total()) # No reinitialisation. self.assertEqual(1, count.get_active()) # Get rid of the keepalive other_shared_handle.acquire(dummy_acquire_fn) del p2 gc.collect() self.assertEqual(0, count.get_active()) def testMultiple(self): count = Count() shared_handle = shared.Shared() other_shared_handle = shared.Shared() def dummy_acquire_fn(): return None def acquire_fn(): return Marker(count) p = shared_handle.acquire(acquire_fn) other_shared_handle.acquire(dummy_acquire_fn) # Get rid of the keepalive self.assertEqual(1, count.get_total()) self.assertEqual(1, count.get_active()) del p gc.collect() self.assertEqual(0, count.get_active()) # Shared value should be garbage collected. # Acquiring multiple times only results in one initialisation p1 = shared_handle.acquire(acquire_fn) # Since shared value was released, expect a reinitialisation. self.assertEqual(2, count.get_total()) self.assertEqual(1, count.get_active()) p2 = shared_handle.acquire(acquire_fn) self.assertEqual(2, count.get_total()) self.assertEqual(1, count.get_active()) other_shared_handle.acquire(dummy_acquire_fn) # Get rid of the keepalive # Check that shared object isn't destroyed if there's still a reference to # it. del p2 gc.collect() self.assertEqual(1, count.get_active()) del p1 gc.collect() self.assertEqual(0, count.get_active()) def testConcurrentCallsDeduped(self): # Test that only one among many calls to acquire will actually run the # initialisation function. count = Count() shared_handle = shared.Shared() other_shared_handle = shared.Shared() refs = [] ref_lock = threading.Lock() def dummy_acquire_fn(): return None def acquire_fn(): time.sleep(1) return Marker(count) def thread_fn(): p = shared_handle.acquire(acquire_fn) with ref_lock: refs.append(p) threads = [] for _ in range(100): t = threading.Thread(target=thread_fn) threads.append(t) t.start() for t in threads: t.join() self.assertEqual(1, count.get_total()) self.assertEqual(1, count.get_active()) other_shared_handle.acquire(dummy_acquire_fn) # Get rid of the keepalive with ref_lock: del refs[:] gc.collect() self.assertEqual(0, count.get_active()) def testDifferentObjects(self): sequence = Sequence() def dummy_acquire_fn(): return None first_handle = shared.Shared() second_handle = shared.Shared() dummy_handle = shared.Shared() f1 = first_handle.acquire(sequence.make_acquire_fn()) s1 = second_handle.acquire(sequence.make_acquire_fn()) self.assertEqual('sequence1', f1.get_name()) self.assertEqual('sequence2', s1.get_name()) f2 = first_handle.acquire(sequence.make_acquire_fn()) s2 = second_handle.acquire(sequence.make_acquire_fn()) # Check that the repeated acquisitions return the earlier objects self.assertEqual('sequence1', f2.get_name()) self.assertEqual('sequence2', s2.get_name()) # Release all references and force garbage-collection del f1 del f2 del s1 del s2 dummy_handle.acquire(dummy_acquire_fn) # Get rid of the keepalive gc.collect() # Check that acquiring again after they're released gives new objects f3 = first_handle.acquire(sequence.make_acquire_fn()) s3 = second_handle.acquire(sequence.make_acquire_fn()) self.assertEqual('sequence3', f3.get_name()) self.assertEqual('sequence4', s3.get_name()) def testTagCacheEviction(self): shared1 = shared.Shared() shared2 = shared.Shared() def acquire_fn_1(): return NamedObject('obj_1') def acquire_fn_2(): return NamedObject('obj_2') # with no tag, shared handle does not know when to evict objects p1 = shared1.acquire(acquire_fn_1) assert p1.get_name() == 'obj_1' p2 = shared1.acquire(acquire_fn_2) assert p2.get_name() == 'obj_1' # cache eviction can be forced by specifying different tags p1 = shared2.acquire(acquire_fn_1, tag='1') assert p1.get_name() == 'obj_1' p2 = shared2.acquire(acquire_fn_2, tag='2') assert p2.get_name() == 'obj_2' def testTagReturnsCached(self): sequence = Sequence() handle = shared.Shared() f1 = handle.acquire(sequence.make_acquire_fn(), tag='1') self.assertEqual('sequence1', f1.get_name()) # should return cached f1 = handle.acquire(sequence.make_acquire_fn(), tag='1') self.assertEqual('sequence1', f1.get_name()) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/subprocess_server.py0000644000175100017510000005262015206113033023153 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import contextlib import dataclasses import glob import hashlib import logging import os import re import shutil import signal import socket import subprocess import threading import time import zipfile from typing import Any from urllib.error import URLError from urllib.request import Request from urllib.request import urlopen import grpc from apache_beam.io.filesystems import FileSystems from apache_beam.runners.internal.names import BEAM_SDK_NAME from apache_beam.version import __version__ as beam_version _LOGGER = logging.getLogger(__name__) @dataclasses.dataclass class _SharedCacheEntry: obj: Any owners: set[str] class _SharedCache: """A cache that keeps objects alive (and repeatedly returns the same instance) until the last user indicates that they're done. The typical usage is as follows:: try: token = cache.register() # All objects retrieved from the cache from this point on will be memoized # and kept alive (including across other threads and callers) at least # until the purge is called below (and possibly longer, if other calls # to register were made). obj = cache.get(...) another_obj = cache.get(...) ... finally: cache.purge(token) """ def __init__(self, constructor, destructor): self._constructor = constructor self._destructor = destructor self._live_owners = set() self._cache = {} self._lock = threading.Lock() self._counter = 0 def _next_id(self): with self._lock: self._counter += 1 return self._counter def register(self): owner = self._next_id() self._live_owners.add(owner) return owner def purge(self, owner): to_delete = [] with self._lock: if owner not in self._live_owners: _LOGGER.warning( "Subprocess owner %s already purged. If this occurs during atexit " "shutdown, the subprocess was already cleaned up earlier.", owner) return self._live_owners.remove(owner) for key, entry in list(self._cache.items()): if owner in entry.owners: entry.owners.remove(owner) if not entry.owners: to_delete.append(entry.obj) del self._cache[key] # Actually call the destructors outside of the lock. for value in to_delete: self._destructor(value) def get(self, *key): if not self._live_owners: raise RuntimeError("At least one owner must be registered.") with self._lock: if key not in self._cache: self._cache[key] = _SharedCacheEntry(self._constructor(*key), set()) for owner in self._live_owners: self._cache[key].owners.add(owner) return self._cache[key].obj class JavaHelper: @classmethod def get_java(cls): java_path = 'java' java_home = os.environ.get('JAVA_HOME') if java_home: java_path = os.path.join(java_home, 'bin', 'java') return java_path class SubprocessServer(object): """An abstract base class for running GRPC Servers as an external process. This class acts as a context which will start up a server, provides a stub to connect to it, and then shuts the server down. For example:: with SubprocessServer(GrpcStubClass, [executable, arg, ...]) as stub: stub.CallService(...) """ def __init__(self, stub_class, cmd, port=None, logger=None): """Creates the server object. :param stub_class: the auto-generated GRPC client stub class used for connecting to the GRPC service :param cmd: command (including arguments) for starting up the server, suitable for passing to `subprocess.POpen`. :param port: (optional) the port at which the subprocess will serve its service. If not given, one will be randomly chosen and the special string "{{PORT}}" will be substituted in the command line arguments with the chosen port. :param logger: (optional) The logger or logger name to use for the subprocess's stderr and stdout. If not given, the current module logger would be used. """ self._owner_id = None self._stub_class = stub_class self._cmd = [str(arg) for arg in cmd] self._port = port self._grpc_channel = None if isinstance(logger, str): self._logger = logging.getLogger(logger) elif isinstance(logger, logging.Logger): self._logger = logger else: self._logger = _LOGGER @classmethod @contextlib.contextmanager def cache_subprocesses(cls): """A context that ensures any subprocess created or used in its duration stay alive for at least the duration of this context. These subprocesses may be shared with other contexts as well. """ try: unique_id = cls._cache.register() yield finally: cls._cache.purge(unique_id) def __enter__(self): return self.start() def __exit__(self, *unused_args): self.stop() def start(self): try: process, endpoint = self.start_process() wait_secs = .1 channel_options = [ ("grpc.max_receive_message_length", -1), ("grpc.max_send_message_length", -1), # Default: 20000ms (20s), increased to 10 minutes for stability ("grpc.keepalive_timeout_ms", 600_000), # Default: 2, set to 0 to allow unlimited pings without data ("grpc.http2.max_pings_without_data", 0), # Default: False, set to True to allow keepalive pings when no calls ("grpc.keepalive_permit_without_calls", True), # Default: 2, set to 0 to allow unlimited ping strikes ("grpc.http2.max_ping_strikes", 0), # Default: 0 (disabled), enable socket reuse for better handling ("grpc.so_reuseport", 1), ] self._grpc_channel = grpc.insecure_channel( endpoint, options=channel_options) channel_ready = grpc.channel_ready_future(self._grpc_channel) while True: if process is not None and process.poll() is not None: _LOGGER.error("Started job service with %s", process.args) raise RuntimeError( 'Service failed to start up with error %s' % process.poll()) try: channel_ready.result(timeout=wait_secs) break except (grpc.FutureTimeoutError, grpc.RpcError): wait_secs *= 1.2 logging.log( logging.WARNING if wait_secs > 1 else logging.DEBUG, 'Waiting for grpc channel to be ready at %s.', endpoint) return self._stub_class(self._grpc_channel) except: # pylint: disable=bare-except _LOGGER.exception("Error bringing up service") self.stop() raise def start_process(self): if self._owner_id is not None: self._cache.purge(self._owner_id) self._owner_id = self._cache.register() return self._cache.get(tuple(self._cmd), self._port, self._logger) def _really_start_process(cmd, port, logger): if not port: port, = pick_port(None) cmd = [arg.replace('{{PORT}}', str(port)) for arg in cmd] # pylint: disable=not-an-iterable endpoint = 'localhost:%s' % port _LOGGER.info("Starting service with %s", str(cmd).replace("',", "'")) process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) # Emit the output of this command as info level logging. def log_stdout(): line = process.stdout.readline() while line: # The log obtained from stdout is bytes, decode it into string. # Remove newline via rstrip() to not print an empty line. logger.info(line.decode(errors='backslashreplace').rstrip()) line = process.stdout.readline() t = threading.Thread(target=log_stdout) t.daemon = True t.start() return process, endpoint def stop(self): self.stop_process() def stop_process(self): if self._owner_id is not None: try: self._cache.purge(self._owner_id) finally: # Make sure _owner_id is set to None even if purge fails. self._owner_id = None if self._grpc_channel: try: self._grpc_channel.close() except: # pylint: disable=bare-except _LOGGER.error( "Could not close the gRPC channel started with cmd %s", self._cmd) finally: self._grpc_channel = None def _really_stop_process(process_and_endpoint): process, _ = process_and_endpoint # pylint: disable=unpacking-non-sequence if not process: return for _ in range(5): if process.poll() is not None: break logging.debug("Sending SIGINT to process") try: process.send_signal(signal.SIGINT) except ValueError: # process.send_signal raises a ValueError on Windows. process.terminate() time.sleep(1) if process.poll() is None: process.kill() _cache = _SharedCache( constructor=_really_start_process, destructor=_really_stop_process) class JavaJarServer(SubprocessServer): MAVEN_CENTRAL_REPOSITORY = 'https://repo.maven.apache.org/maven2' MAVEN_STAGING_REPOSITORY = ( 'https://repository.apache.org/content/groups/staging') GOOGLE_MAVEN_MIRROR = ( 'https://maven-central.storage-download.googleapis.com/maven2') BEAM_GROUP_ID = 'org.apache.beam' JAR_CACHE = os.path.expanduser("~/.apache_beam/cache/jars") _BEAM_SERVICES = type( 'local', (threading.local, ), dict(__init__=lambda self: setattr(self, 'replacements', {})))() _DEFAULT_USER_AGENT = f'{BEAM_SDK_NAME}/{beam_version}' def __init__( self, stub_class, path_to_jar, java_arguments, classpath=None, cache_dir=None, logger=None): self._java_path = JavaHelper.get_java() if classpath: # java -jar ignores the classpath, so we make a new jar that embeds # the requested classpath. path_to_jar = self.make_classpath_jar(path_to_jar, classpath, cache_dir) super().__init__( stub_class, [self._java_path, '-jar', path_to_jar] + list(java_arguments), logger=logger) self._existing_service = path_to_jar if is_service_endpoint( path_to_jar) else None def start_process(self): if self._existing_service: return None, self._existing_service else: if not shutil.which(self._java_path): java_home = os.environ.get('JAVA_HOME') if java_home: raise RuntimeError( 'Java is not correctly installed in JAVA_HOME=%s to use this ' 'transform/runner. Please check if JAVA_HOME is correctly set and' ' points to your Java installation directory.' % java_home) else: raise RuntimeError( 'Java must be installed on this system to use this ' 'transform/runner.') return super().start_process() def stop_process(self): if self._existing_service: pass else: return super().stop_process() @classmethod def jar_name(cls, artifact_id, version, classifier=None, appendix=None): return '-'.join( filter(None, [artifact_id, appendix, version, classifier])) + '.jar' @classmethod def path_to_maven_jar( cls, artifact_id, group_id, version, repository=MAVEN_CENTRAL_REPOSITORY, classifier=None, appendix=None): return '/'.join([ repository, group_id.replace('.', '/'), artifact_id, version, cls.jar_name(artifact_id, version, classifier, appendix) ]) @classmethod def parse_gradle_target(cls, gradle_target, artifact_id=None): gradle_package = gradle_target.strip(':').rsplit(':', 1)[0] if not artifact_id: artifact_id = 'beam-' + gradle_package.replace(':', '-') return gradle_package, artifact_id @classmethod def path_to_dev_beam_jar( cls, gradle_target, appendix=None, version=beam_version, artifact_id=None): gradle_package, artifact_id = cls.parse_gradle_target( gradle_target, artifact_id) project_root = os.path.sep.join( os.path.abspath(__file__).split(os.path.sep)[:-5]) return os.path.join( project_root, gradle_package.replace(':', os.path.sep), 'build', 'libs', cls.jar_name( artifact_id, version.replace('.dev', ''), classifier='SNAPSHOT', appendix=appendix)) @classmethod def path_to_beam_jar( cls, gradle_target, appendix=None, version=beam_version, artifact_id=None, maven_repository_url=None): if gradle_target in cls._BEAM_SERVICES.replacements: return cls._BEAM_SERVICES.replacements[gradle_target] _, artifact_id = cls.parse_gradle_target(gradle_target, artifact_id) project_root = os.path.sep.join( os.path.abspath(__file__).split(os.path.sep)[:-5]) local_path = cls.path_to_dev_beam_jar( gradle_target, appendix, version, artifact_id) if os.path.exists(local_path): _LOGGER.info('Using pre-built snapshot at %s', local_path) return local_path maven_repo = maven_repository_url or cls.MAVEN_CENTRAL_REPOSITORY if 'rc' in version: # Release candidate version = version.split('rc')[0] maven_repo = cls.MAVEN_STAGING_REPOSITORY elif '.dev' in version: # TODO: Attempt to use nightly snapshots? raise RuntimeError( ( '%s not found. ' 'Please build the server with \n cd %s; ./gradlew %s') % (local_path, os.path.abspath(project_root), gradle_target)) return cls.path_to_maven_jar( artifact_id, cls.BEAM_GROUP_ID, version, maven_repo, appendix=appendix) @classmethod def _download_jar_to_cache( cls, download_url, cached_jar_path, user_agent=None): """Downloads a jar from the given URL to the specified cache path. Args: download_url (str): The URL to download from. cached_jar_path (str): The local path where the jar should be cached. user_agent (str): The user agent to use when downloading. """ # Issue warning when downloading from public repositories public_repos = [ cls.MAVEN_CENTRAL_REPOSITORY, cls.GOOGLE_MAVEN_MIRROR, ] if any(download_url.startswith(repo) for repo in public_repos): _LOGGER.warning( " WARNING: Apache Beam is downloading dependencies from a " "public repository at runtime.\n" " This may pose security risks or cause instability due to " "repository availability.\n" " URL: %s\n" " Destination: %s\n" " Consider pre-staging dependencies or using a private repository " "mirror.\n" " For more information, see: " "https://beam.apache.org/documentation/sdks/python-dependencies/", download_url, cached_jar_path) try: url_read = FileSystems.open(download_url) except ValueError: if user_agent is None: user_agent = cls._DEFAULT_USER_AGENT url_request = Request(download_url, headers={'User-Agent': user_agent}) url_read = urlopen(url_request) with open(cached_jar_path + '.tmp', 'wb') as jar_write: shutil.copyfileobj(url_read, jar_write, length=1 << 20) try: os.rename(cached_jar_path + '.tmp', cached_jar_path) except FileNotFoundError: # A race when multiple programs run in parallel and the cached_jar # is already moved. Safe to ignore. pass @classmethod def local_jar(cls, url, cache_dir=None, user_agent=None): """Returns a local path to the given jar, downloading it if necessary. Args: url (str): A URL or local path to a jar file. cache_dir (str): The directory to use for caching downloaded jars. If not specified, a default temporary directory will be used. user_agent (str): The user agent to use when downloading the jar. Returns: str: The local path to the jar file. """ if cache_dir is None: cache_dir = cls.JAR_CACHE # TODO: Verify checksum? if is_service_endpoint(url): return url elif os.path.exists(url): return url else: cached_jar = os.path.join(cache_dir, os.path.basename(url)) if os.path.exists(cached_jar): _LOGGER.info('Using cached job server jar from %s' % url) else: _LOGGER.info('Downloading job server jar from %s' % url) if not os.path.exists(cache_dir): os.makedirs(cache_dir) # TODO: Clean up this cache according to some policy. try: cls._download_jar_to_cache(url, cached_jar, user_agent) except URLError as e: # Try Google Maven mirror as fallback if the original URL is from # Maven Central if url.startswith(cls.MAVEN_CENTRAL_REPOSITORY): fallback_url = url.replace( cls.MAVEN_CENTRAL_REPOSITORY, cls.GOOGLE_MAVEN_MIRROR) _LOGGER.info( 'Trying Google Maven mirror fallback: %s' % fallback_url) try: cls._download_jar_to_cache(fallback_url, cached_jar, user_agent) _LOGGER.info( 'Successfully downloaded from Google Maven mirror: %s' % fallback_url) except URLError as fallback_e: raise RuntimeError( f'Unable to fetch remote job server jar at {url}: {e}. ' f'Also failed to fetch from Google Maven mirror at ' f'{fallback_url}: {fallback_e}. ' f'If no Internet access at runtime, stage the jar at ' f'{cached_jar}') else: raise RuntimeError( f'Unable to fetch remote job server jar at {url}: {e}. If no ' f'Internet access at runtime, stage the jar at {cached_jar}') return cached_jar @classmethod @contextlib.contextmanager def beam_services(cls, replacements): try: old = cls._BEAM_SERVICES.replacements cls._BEAM_SERVICES.replacements = dict(old, **replacements) yield finally: cls._BEAM_SERVICES.replacements = old @classmethod def make_classpath_jar(cls, main_jar, extra_jars, cache_dir=None): if cache_dir is None: cache_dir = cls.JAR_CACHE composite_jar_dir = os.path.join(cache_dir, 'composite-jars') os.makedirs(composite_jar_dir, exist_ok=True) classpath = [] # Class-Path references from a jar must be relative, so we create # a relatively-addressable subdirectory with symlinks to all the # required jars. for pattern in [main_jar] + list(extra_jars): for path in glob.glob(pattern) or [pattern]: path = os.path.abspath(path) rel_path = hashlib.sha256( path.encode('utf-8')).hexdigest() + os.path.splitext(path)[1] classpath.append(rel_path) if not os.path.lexists(os.path.join(composite_jar_dir, rel_path)): os.symlink(path, os.path.join(composite_jar_dir, rel_path)) # Now create a single jar that simply references the rest and has the same # main class as main_jar. composite_jar = os.path.join( composite_jar_dir, hashlib.sha256(' '.join(sorted(classpath)).encode('ascii')).hexdigest() + '.jar') if not os.path.exists(composite_jar): with zipfile.ZipFile(main_jar) as main: with main.open('META-INF/MANIFEST.MF') as manifest: main_class = next( filter(lambda line: line.startswith(b'Main-Class: '), manifest)) with zipfile.ZipFile(composite_jar + '.tmp', 'w') as composite: with composite.open('META-INF/MANIFEST.MF', 'w') as manifest: manifest.write(b'Manifest-Version: 1.0\n') manifest.write(main_class) manifest.write( b'Class-Path: ' + '\n '.join(classpath).encode('ascii') + b'\n') os.rename(composite_jar + '.tmp', composite_jar) return composite_jar def is_service_endpoint(path): """Checks whether the path conforms to the 'beam_services' PipelineOption.""" return re.match(r'^[a-zA-Z0-9.-]+:\d+$', path) def pick_port(*ports): """ Returns a list of ports, same length as input ports list, but replaces all None or 0 ports with a random free port. """ sockets = [] def find_free_port(port): if port: return port else: try: s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) except OSError as e: # [Errno 97] Address family not supported by protocol # Likely indicates we are in an IPv6-only environment (BEAM-10618). Try # again with AF_INET6. if e.errno == 97: s = socket.socket(socket.AF_INET6, socket.SOCK_STREAM) else: raise e sockets.append(s) s.bind(('localhost', 0)) return s.getsockname()[1] ports = list(map(find_free_port, ports)) # Close sockets only now to avoid the same port to be chosen twice for s in sockets: s.close() return ports ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/subprocess_server_test.py0000644000175100017510000003647215206113033024221 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the processes module.""" # pytype: skip-file import atexit import glob import os import random import re import shutil import socketserver import subprocess import tempfile import threading import unittest from unittest.mock import patch from apache_beam.runners.portability import job_server from apache_beam.utils import subprocess_server class JavaJarServerTest(unittest.TestCase): def test_gradle_jar_release(self): self.assertEqual( 'https://repo.maven.apache.org/maven2/org/apache/beam/' 'beam-sdks-java-fake/VERSION/beam-sdks-java-fake-VERSION.jar', subprocess_server.JavaJarServer.path_to_beam_jar( ':sdks:java:fake:fatJar', version='VERSION')) self.assertEqual( 'https://repo.maven.apache.org/maven2/org/apache/beam/' 'beam-sdks-java-fake/VERSION/beam-sdks-java-fake-A-VERSION.jar', subprocess_server.JavaJarServer.path_to_beam_jar( ':sdks:java:fake:fatJar', appendix='A', version='VERSION')) self.assertEqual( 'https://repo.maven.apache.org/maven2/org/apache/beam/' 'beam-sdks-java-fake/VERSION/beam-sdks-java-fake-A-VERSION.jar', subprocess_server.JavaJarServer.path_to_beam_jar( ':gradle:target:doesnt:matter', appendix='A', version='VERSION', artifact_id='beam-sdks-java-fake')) self.assertEqual( 'https://repository.apache.org/content/groups/staging/org/apache/beam/' 'beam-sdks-java-fake/2.99.9/beam-sdks-java-fake-2.99.9.jar', subprocess_server.JavaJarServer.path_to_beam_jar( ':sdks:java:fake:fatJar', version='2.99.9rc2')) def test_gradle_jar_dev(self): with self.assertRaisesRegex( Exception, re.escape(os.path.join('sdks', 'java', 'fake', 'build', 'libs', 'beam-sdks-java-fake-VERSION-SNAPSHOT.jar')) + ' not found.'): subprocess_server.JavaJarServer.path_to_beam_jar( ':sdks:java:fake:fatJar', version='VERSION.dev') with self.assertRaisesRegex( Exception, re.escape(os.path.join('sdks', 'java', 'fake', 'build', 'libs', 'beam-sdks-java-fake-A-VERSION-SNAPSHOT.jar')) + ' not found.'): subprocess_server.JavaJarServer.path_to_beam_jar( ':sdks:java:fake:fatJar', appendix='A', version='VERSION.dev') with self.assertRaisesRegex( Exception, re.escape(os.path.join('sdks', 'java', 'fake', 'build', 'libs', 'fake-artifact-id-A-VERSION-SNAPSHOT.jar')) + ' not found.'): subprocess_server.JavaJarServer.path_to_beam_jar( ':sdks:java:fake:fatJar', appendix='A', version='VERSION.dev', artifact_id='fake-artifact-id') def test_beam_services(self): with subprocess_server.JavaJarServer.beam_services({':some:target': 'foo'}): self.assertEqual( 'foo', subprocess_server.JavaJarServer.path_to_beam_jar(':some:target')) def test_local_jar(self): class Handler(socketserver.BaseRequestHandler): timeout = 1 def handle(self): data = self.request.recv(1024) if 'User-Agent: Apache Beam SDK for Python' in str(data): self.request.sendall(b'HTTP/1.1 200 OK\n\ndata') else: self.request.sendall(b'HTTP/1.1 400 BAD REQUEST\n\n') port, = subprocess_server.pick_port(None) server = socketserver.TCPServer(('localhost', port), Handler) t = threading.Thread(target=server.handle_request) t.daemon = True t.start() with tempfile.TemporaryDirectory() as temp_dir: subprocess_server.JavaJarServer.local_jar( 'http://localhost:%s/path/to/file.jar' % port, temp_dir) with open(os.path.join(temp_dir, 'file.jar')) as fin: self.assertEqual(fin.read(), 'data') def test_local_jar_fallback_to_google_maven_mirror(self): """Test that Google Maven mirror is used as fallback when Maven Central fails.""" class MavenCentralHandler(socketserver.BaseRequestHandler): timeout = 1 def handle(self): # Simulate Maven Central returning 403 Forbidden self.request.sendall(b'HTTP/1.1 403 Forbidden\n\n') # Set up Maven Central server (will return 403) maven_port, = subprocess_server.pick_port(None) maven_server = socketserver.TCPServer(('localhost', maven_port), MavenCentralHandler) maven_thread = threading.Thread(target=maven_server.handle_request) maven_thread.daemon = True maven_thread.start() # Temporarily replace the Maven Central constant to use our test server original_maven_central = ( subprocess_server.JavaJarServer.MAVEN_CENTRAL_REPOSITORY) try: subprocess_server.JavaJarServer.MAVEN_CENTRAL_REPOSITORY = ( f'http://localhost:{maven_port}/maven2') with tempfile.TemporaryDirectory() as temp_dir: # Use a Maven Central URL that will trigger the fallback to real # Google mirror maven_url = ( f'http://localhost:{maven_port}/maven2/org/apache/beam/' f'beam-sdks-java-extensions-schemaio-expansion-service/2.63.0/' f'beam-sdks-java-extensions-schemaio-expansion-service-2.63.0.jar') # This should fail on our mock Maven Central and fallback to the # real Google mirror jar_path = subprocess_server.JavaJarServer.local_jar( maven_url, temp_dir) # Verify the file was downloaded successfully (from the real Google # mirror) self.assertTrue(os.path.exists(jar_path)) jar_size = os.path.getsize(jar_path) self.assertTrue(jar_size > 0) # Should have actual content finally: # Restore original constants subprocess_server.JavaJarServer.MAVEN_CENTRAL_REPOSITORY = ( original_maven_central) @unittest.skipUnless(shutil.which('javac'), 'missing java jdk') def test_classpath_jar(self): with tempfile.TemporaryDirectory() as temp_dir: try: # Avoid having to prefix everything in our test strings. oldwd = os.getcwd() os.chdir(temp_dir) with open('Main.java', 'w') as fout: fout.write( """ public class Main { public static void main(String[] args) { Other.greet(); } } """) with open('Other.java', 'w') as fout: fout.write( """ public class Other { public static void greet() { System.out.println("You got me!"); } } """) os.mkdir('jars') # Using split just for readability/copyability. subprocess.check_call('javac Main.java Other.java'.split()) subprocess.check_call('jar cfe jars/Main.jar Main Main.class'.split()) subprocess.check_call('jar cf jars/Other.jar Other.class'.split()) # Make sure the java and class files don't get picked up. for path in glob.glob('*.*'): os.unlink(path) # These should fail. self.assertNotEqual( subprocess.call('java -jar jars/Main.jar'.split()), 0) self.assertNotEqual( subprocess.call('java -jar jars/Other.jar'.split()), 0) os.mkdir('beam_temp') composite_jar = subprocess_server.JavaJarServer.make_classpath_jar( 'jars/Main.jar', ['jars/Other.jar'], cache_dir='beam_temp') # This, however, should work. subprocess.check_call(f'java -jar {composite_jar}'.split()) finally: os.chdir(oldwd) class CacheTest(unittest.TestCase): @staticmethod def with_prefix(prefix): return '%s-%s' % (prefix, random.random()) def test_memoization(self): cache = subprocess_server._SharedCache(self.with_prefix, lambda x: None) try: token = cache.register() a = cache.get('a') self.assertEqual(a[0], 'a') self.assertEqual(cache.get('a'), a) b = cache.get('b') self.assertEqual(b[0], 'b') self.assertEqual(cache.get('b'), b) finally: cache.purge(token) def test_purged(self): cache = subprocess_server._SharedCache(self.with_prefix, lambda x: None) try: token = cache.register() a = cache.get('a') self.assertEqual(cache.get('a'), a) finally: cache.purge(token) try: token = cache.register() new_a = cache.get('a') self.assertNotEqual(new_a, a) finally: cache.purge(token) def test_multiple_owners(self): cache = subprocess_server._SharedCache(self.with_prefix, lambda x: None) try: owner1 = cache.register() a = cache.get('a') try: self.assertEqual(cache.get('a'), a) owner2 = cache.register() b = cache.get('b') self.assertEqual(cache.get('b'), b) finally: cache.purge(owner2) self.assertEqual(cache.get('a'), a) self.assertEqual(cache.get('b'), b) finally: cache.purge(owner1) try: owner3 = cache.register() self.assertNotEqual(cache.get('a'), a) self.assertNotEqual(cache.get('b'), b) finally: cache.purge(owner3) def test_interleaved_owners(self): cache = subprocess_server._SharedCache(self.with_prefix, lambda x: None) owner1 = cache.register() a = cache.get('a') self.assertEqual(cache.get('a'), a) owner2 = cache.register() b = cache.get('b') self.assertEqual(cache.get('b'), b) cache.purge(owner1) self.assertNotEqual(cache.get('a'), a) self.assertEqual(cache.get('b'), b) cache.purge(owner2) owner3 = cache.register() self.assertNotEqual(cache.get('b'), b) cache.purge(owner3) def test_destructor_exception_partial_state(self): # In SubprocessServer.stop_process(), we need to make sure self._owner_id is always # set to None if it is not already set, even if a destructor exception happens # during purge(owner_id). destructor_calls = [] def faulty_destructor(obj): destructor_calls.append(obj) raise RuntimeError("Destructor failed") custom_cache = subprocess_server._SharedCache( lambda *args: "process_obj", faulty_destructor) class CustomServer(subprocess_server.SubprocessServer): _cache = custom_cache def __init__(self): super().__init__(lambda channel: None, ["dummy_cmd"], port=12345) server = CustomServer() server.start_process() owner_id = server._owner_id self.assertIsNotNone(owner_id) self.assertIn(owner_id, custom_cache._live_owners) # First stop attempt fails in the destructor with self.assertRaises(RuntimeError): server.stop_process() # Verify fixed state: owner is purged from cache set, AND self._owner_id is successfully cleared to None self.assertNotIn(owner_id, custom_cache._live_owners) self.assertIsNone(server._owner_id) # Second stop attempt safely does nothing (no ValueError raised) try: server.stop_process() except ValueError: self.fail("ValueError should not be raised here.") def test_duplicate_atexit_registration_on_restart(self): # Make sure we don't have duplicate atexit registration when reusing a # StopOnExistJobServer instance. class DummyJobServer(job_server.JobServer): def start(self): return "localhost:8080" def stop(self): pass wrapper = job_server.StopOnExitJobServer(DummyJobServer()) registered_callbacks = [] def mock_register(cb): registered_callbacks.append(cb) def mock_unregister(cb): if cb in registered_callbacks: registered_callbacks.remove(cb) with patch('atexit.register', side_effect=mock_register), \ patch('atexit.unregister', side_effect=mock_unregister, create=True): # First start registers stop callback wrapper.start() self.assertTrue(wrapper._started) self.assertEqual(len(registered_callbacks), 1) # Explicit stop clears _started AND unregisters the callback wrapper.stop() self.assertFalse(wrapper._started) self.assertEqual(len(registered_callbacks), 0) # Re-starting registers the callback again, leaving exactly 1 active callback wrapper.start() self.assertTrue(wrapper._started) self.assertEqual(len(registered_callbacks), 1) def test_concurrent_purge_race_condition(self): # Concurrent threads attempting to check memebership and call purge for the same owner. # Here we explicitly define a synchronized set to mimic the behavior of _live_owners. # This set will block two threads on __contains__, allowing us to test the race condition. cache = subprocess_server._SharedCache(lambda x: "obj", lambda x: None) owner = cache.register() barrier = threading.Barrier(2) exceptions = [] class SynchronizedSet(set): def __contains__(self, item): res = super().__contains__(item) try: # Force both threads to align right after checking membership but before removal barrier.wait(timeout=0.2) except threading.BrokenBarrierError: pass return res cache._live_owners = SynchronizedSet(cache._live_owners) def purge_worker(): try: cache.purge(owner) except Exception as e: exceptions.append(e) t1 = threading.Thread(target=purge_worker) t2 = threading.Thread(target=purge_worker) t1.start() t2.start() t1.join() t2.join() # Both threads should succeed cleanly without raising an exception under idempotent purging. self.assertEqual(len(exceptions), 0) def test_stop_process_after_cache_purged(self): # Reproduce the ValueError when stop_process() (called by atexit) # runs after the cache/owner was already purged during test teardown. cache = subprocess_server._SharedCache( lambda *args: "dummy_process", lambda obj: None) class DummySubprocessServer(subprocess_server.SubprocessServer): _cache = cache def __init__(self): super().__init__(lambda channel: None, ["dummy_cmd"], port=12345) server = DummySubprocessServer() server.start_process() owner_id = server._owner_id # Simulate pipeline context exit or test teardown purging the cache directly cache.purge(owner_id) # Calling stop_process() (which happens during atexit) should succeed cleanly # without raising ValueError. server.stop_process() if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/thread_pool_executor.py0000644000175100017510000000774315206113033023621 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import queue import threading import weakref from concurrent.futures import _base class _WorkItem(object): def __init__(self, future, fn, args, kwargs): self._future = future self._fn = fn self._fn_args = args self._fn_kwargs = kwargs def run(self): if self._future.set_running_or_notify_cancel(): # If the future wasn't cancelled, then attempt to execute it. try: self._future.set_result(self._fn(*self._fn_args, **self._fn_kwargs)) except BaseException as exc: self._future.set_exception(exc) class _Worker(threading.Thread): def __init__(self, idle_worker_queue, work_item): super().__init__() self._idle_worker_queue = idle_worker_queue self._work_item = work_item self._wake_semaphore = threading.Semaphore(0) self._lock = threading.Lock() self._shutdown = False def run(self): while True: self._work_item.run() self._work_item = None self._idle_worker_queue.put(self) self._wake_semaphore.acquire() if self._work_item is None: return def assign_work(self, work_item): """Assigns the work item and wakes up the thread. This method must only be called while the worker is idle. """ self._work_item = work_item self._wake_semaphore.release() def shutdown(self): """Wakes up this thread with a 'None' work item signalling to shutdown.""" self._wake_semaphore.release() class UnboundedThreadPoolExecutor(_base.Executor): def __init__(self): self._idle_worker_queue = queue.Queue() self._max_idle_threads = 16 self._workers = weakref.WeakSet() self._shutdown = False self._lock = threading.Lock() # Guards access to _workers and _shutdown def submit(self, fn, *args, **kwargs): """Attempts to submit the work item. A runtime error is raised if the pool has been shutdown. """ future = _base.Future() work_item = _WorkItem(future, fn, args, kwargs) with self._lock: if self._shutdown: raise RuntimeError( 'Cannot schedule new tasks after thread pool has been shutdown.') try: self._idle_worker_queue.get(block=False).assign_work(work_item) # If we have more idle threads then the max allowed, shutdown a thread. if self._idle_worker_queue.qsize() > self._max_idle_threads: try: self._idle_worker_queue.get(block=False).shutdown() except queue.Empty: pass except queue.Empty: worker = _Worker(self._idle_worker_queue, work_item) worker.daemon = True worker.start() self._workers.add(worker) return future def shutdown(self, wait=True): with self._lock: if self._shutdown: return self._shutdown = True for worker in self._workers: worker.shutdown() if wait: for worker in self._workers: worker.join() class _SharedUnboundedThreadPoolExecutor(UnboundedThreadPoolExecutor): def shutdown(self, wait=True): # Prevent shutting down the shared thread pool pass _SHARED_UNBOUNDED_THREAD_POOL_EXECUTOR = _SharedUnboundedThreadPoolExecutor() def shared_unbounded_instance(): return _SHARED_UNBOUNDED_THREAD_POOL_EXECUTOR ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/thread_pool_executor_test.py0000644000175100017510000000735215206113033024654 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for UnboundedThreadPoolExecutor.""" # pytype: skip-file import itertools import threading import time import traceback import unittest from apache_beam.utils import thread_pool_executor from apache_beam.utils.thread_pool_executor import UnboundedThreadPoolExecutor class UnboundedThreadPoolExecutorTest(unittest.TestCase): def setUp(self): self._lock = threading.Lock() self._worker_idents = [] def append_and_sleep(self, sleep_time): with self._lock: self._worker_idents.append(threading.current_thread().ident) time.sleep(sleep_time) def raise_error(self, message): raise ValueError(message) def test_shutdown_with_no_workers(self): with UnboundedThreadPoolExecutor(): pass def test_shutdown_with_fast_workers(self): futures = [] with UnboundedThreadPoolExecutor() as executor: for _ in range(0, 5): futures.append(executor.submit(self.append_and_sleep, 0.01)) for future in futures: future.result(timeout=10) with self._lock: self.assertEqual(5, len(self._worker_idents)) def test_shutdown_with_slow_workers(self): futures = [] with UnboundedThreadPoolExecutor() as executor: for _ in range(0, 5): futures.append(executor.submit(self.append_and_sleep, 1)) for future in futures: future.result(timeout=10) with self._lock: self.assertEqual(5, len(self._worker_idents)) def test_worker_reuse(self): futures = [] with UnboundedThreadPoolExecutor() as executor: for _ in range(0, 5): futures.append(executor.submit(self.append_and_sleep, 0.01)) time.sleep(3) for _ in range(0, 5): futures.append(executor.submit(self.append_and_sleep, 0.01)) for future in futures: future.result(timeout=10) with self._lock: self.assertEqual(10, len(self._worker_idents)) self.assertTrue(len(set(self._worker_idents)) < 10) def test_exception_propagation(self): with UnboundedThreadPoolExecutor() as executor: future = executor.submit(self.raise_error, 'footest') try: future.result() except Exception: message = traceback.format_exc() else: raise AssertionError('expected exception not raised') self.assertIn('footest', message) self.assertIn('raise_error', message) def test_map(self): with UnboundedThreadPoolExecutor() as executor: executor.map(self.append_and_sleep, itertools.repeat(0.01, 5)) with self._lock: self.assertEqual(5, len(self._worker_idents)) def test_shared_shutdown_does_nothing(self): thread_pool_executor.shared_unbounded_instance().shutdown() futures = [] with thread_pool_executor.shared_unbounded_instance() as executor: for _ in range(0, 5): futures.append(executor.submit(self.append_and_sleep, 0.01)) for future in futures: future.result(timeout=10) with self._lock: self.assertEqual(5, len(self._worker_idents)) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/timestamp.py0000644000175100017510000003441615206113033021403 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Timestamp utilities. For internal use only; no backwards-compatibility guarantees. """ # pytype: skip-file # mypy: disallow-untyped-defs import datetime import time from typing import Union from typing import overload import dateutil.parser import pytz from google.protobuf import duration_pb2 from google.protobuf import timestamp_pb2 from apache_beam.portability import common_urns # types compatible with Timestamp.of() TimestampTypes = Union[int, float, 'Timestamp'] # types compatible with Duration.of() DurationTypes = Union[int, float, 'Duration'] TimestampDurationTypes = Union[int, float, 'Duration', 'Timestamp'] class Timestamp(object): """Represents a Unix second timestamp with microsecond granularity. Can be treated in common timestamp arithmetic operations as a numeric type. Internally stores a time interval as an int of microseconds. This strategy is necessary since floating point values lose precision when storing values, especially after arithmetic operations (for example, 10000000 % 0.1 evaluates to 0.0999999994448885). """ def __init__( self, seconds: Union[int, float] = 0, micros: Union[int, float] = 0) -> None: if not isinstance(seconds, (int, float)): raise TypeError( 'Cannot interpret %s %s as seconds.' % (seconds, type(seconds))) if not isinstance(micros, (int, float)): raise TypeError( 'Cannot interpret %s %s as micros.' % (micros, type(micros))) self.micros = int(seconds * 1000000) + int(micros) @staticmethod def of(seconds: TimestampTypes) -> 'Timestamp': """Return the Timestamp for the given number of seconds. If the input is already a Timestamp, the input itself will be returned. Args: seconds: Number of seconds as int, float, long, or Timestamp. Returns: Corresponding Timestamp object. """ if isinstance(seconds, Timestamp): return seconds elif isinstance(seconds, (int, float)): return Timestamp(seconds) elif isinstance(seconds, datetime.datetime): return Timestamp.from_utc_datetime(seconds) else: raise TypeError( 'Cannot interpret %s %s as Timestamp.' % (seconds, type(seconds))) @staticmethod def now() -> 'Timestamp': return Timestamp(seconds=time.time()) @staticmethod def _epoch_datetime_utc() -> datetime.datetime: return datetime.datetime.fromtimestamp(0, pytz.utc) @classmethod def from_utc_datetime(cls, dt: datetime.datetime) -> 'Timestamp': """Create a ``Timestamp`` instance from a ``datetime.datetime`` object. Args: dt: A ``datetime.datetime`` object in UTC (offset-aware). """ if dt.tzinfo is None: raise ValueError( "dt has no timezone info " + "(https://docs.python.org/3/library/datetime.html" + "#aware-and-naive-objects): %s" % dt) if dt.tzinfo != pytz.utc and dt.tzinfo != datetime.timezone.utc: raise ValueError('dt not in UTC: %s' % dt) duration = dt - cls._epoch_datetime_utc() return Timestamp(duration.total_seconds()) @classmethod def from_rfc3339(cls, rfc3339: str) -> 'Timestamp': """Create a ``Timestamp`` instance from an RFC 3339 compliant string. .. note:: All timezones are implicitly converted to UTC. Args: rfc3339: String in RFC 3339 form. """ try: dt = dateutil.parser.isoparse(rfc3339).astimezone(pytz.UTC) except ValueError as e: raise ValueError( "Could not parse RFC 3339 string '{}' due to error: '{}'.".format( rfc3339, e)) return cls.from_utc_datetime(dt) def seconds(self) -> int: """Returns the timestamp in seconds.""" return self.micros // 1000000 def predecessor(self) -> 'Timestamp': """Returns the largest timestamp smaller than self.""" return Timestamp(micros=self.micros - 1) def successor(self) -> 'Timestamp': """Returns the smallest timestamp larger than self.""" return Timestamp(micros=self.micros + 1) def __repr__(self) -> str: micros = self.micros sign = '' if micros < 0: sign = '-' micros = -micros int_part = micros // 1000000 frac_part = micros % 1000000 if frac_part: return 'Timestamp(%s%d.%06d)' % (sign, int_part, frac_part) return 'Timestamp(%s%d)' % (sign, int_part) def to_utc_datetime(self, has_tz: bool = False) -> datetime.datetime: """Returns a ``datetime.datetime`` object of UTC for this Timestamp. Note that this method returns a ``datetime.datetime`` object without a timezone info by default, as builtin `datetime.datetime.utcnow` method. If this is used as part of the processed data, one should set has_tz=True to avoid offset due to default timezone mismatch. Args: has_tz: whether the timezone info is attached, default to False. Returns: a ``datetime.datetime`` object of UTC for this Timestamp. """ # We can't easily construct a datetime object from microseconds, so we # create one at the epoch and add an appropriate timedelta interval. epoch = self._epoch_datetime_utc() if not has_tz: epoch = epoch.replace(tzinfo=None) return epoch + datetime.timedelta(microseconds=self.micros) def to_rfc3339(self) -> str: # Append 'Z' for UTC timezone. return self.to_utc_datetime().isoformat() + 'Z' def to_proto(self) -> timestamp_pb2.Timestamp: """Returns the `google.protobuf.timestamp_pb2` representation.""" secs = self.micros // 1000000 nanos = (self.micros % 1000000) * 1000 return timestamp_pb2.Timestamp(seconds=secs, nanos=nanos) @staticmethod def from_proto(timestamp_proto: timestamp_pb2.Timestamp) -> 'Timestamp': """Creates a Timestamp from a `google.protobuf.timestamp_pb2`. Note that the google has a sub-second resolution of nanoseconds whereas this class has a resolution of microsends. This class will truncate the nanosecond resolution down to the microsecond. """ if timestamp_proto.nanos % 1000 != 0: # TODO(https://github.com/apache/beam/issues/19922): Better define # timestamps. raise ValueError( "Cannot convert from nanoseconds to microseconds " + "because this loses precision. Please make sure that " + "this is the correct behavior you want and manually " + "truncate the precision to the nearest microseconds. " + "See [https://github.com/apache/beam/issues/19922] for " + "more information.") return Timestamp( seconds=timestamp_proto.seconds, micros=timestamp_proto.nanos // 1000) def __float__(self) -> float: # Note that the returned value may have lost precision. return self.micros / 1000000 def __int__(self) -> int: # Note that the returned value may have lost precision. return self.micros // 1000000 def __eq__(self, other: object) -> bool: # Allow comparisons between Duration and Timestamp values. if isinstance(other, (Duration, Timestamp)): return self.micros == other.micros elif isinstance(other, (int, float)): return self.micros == Timestamp.of(other).micros else: # Support equality with other types return NotImplemented def __lt__(self, other: TimestampDurationTypes) -> bool: # Allow comparisons between Duration and Timestamp values. if not isinstance(other, Duration): other = Timestamp.of(other) return self.micros < other.micros def __gt__(self, other: TimestampDurationTypes) -> bool: return not (self < other or self == other) def __le__(self, other: TimestampDurationTypes) -> bool: return self < other or self == other def __ge__(self, other: TimestampDurationTypes) -> bool: return not self < other def __hash__(self) -> int: return hash(self.micros) def __add__(self, other: DurationTypes) -> 'Timestamp': other = Duration.of(other) return Timestamp(micros=self.micros + other.micros) def __radd__(self, other: DurationTypes) -> 'Timestamp': return self + other @overload def __sub__(self, other: DurationTypes) -> 'Timestamp': pass @overload def __sub__(self, other: 'Timestamp') -> 'Duration': pass def __sub__( self, other: Union[DurationTypes, 'Timestamp']) -> Union['Timestamp', 'Duration']: if isinstance(other, Timestamp): return Duration(micros=self.micros - other.micros) other = Duration.of(other) return Timestamp(micros=self.micros - other.micros) def __mod__(self, other: DurationTypes) -> 'Duration': other = Duration.of(other) return Duration(micros=self.micros % other.micros) MIN_TIMESTAMP = Timestamp( micros=int(common_urns.constants.MIN_TIMESTAMP_MILLIS.constant) * 1000) MAX_TIMESTAMP = Timestamp( micros=int(common_urns.constants.MAX_TIMESTAMP_MILLIS.constant) * 1000) class Duration(object): """Represents a second duration with microsecond granularity. Can be treated in common arithmetic operations as a numeric type. Internally stores a time interval as an int of microseconds. This strategy is necessary since floating point values lose precision when storing values, especially after arithmetic operations (for example, 10000000 % 0.1 evaluates to 0.0999999994448885). """ def __init__( self, seconds: Union[int, float] = 0, micros: Union[int, float] = 0) -> None: self.micros = int(seconds * 1000000) + int(micros) @staticmethod def of(seconds: DurationTypes) -> 'Duration': """Return the Duration for the given number of seconds since Unix epoch. If the input is already a Duration, the input itself will be returned. Args: seconds: Number of seconds as int, float or Duration. Returns: Corresponding Duration object. """ if isinstance(seconds, Timestamp): raise TypeError('Cannot interpret %s as Duration.' % seconds) if isinstance(seconds, Duration): return seconds return Duration(seconds) def to_proto(self) -> duration_pb2.Duration: """Returns the `google.protobuf.duration_pb2` representation.""" secs = self.micros // 1000000 nanos = (self.micros % 1000000) * 1000 return duration_pb2.Duration(seconds=secs, nanos=nanos) @staticmethod def from_proto(duration_proto: duration_pb2.Duration) -> 'Duration': """Creates a Duration from a `google.protobuf.duration_pb2`. Note that the google has a sub-second resolution of nanoseconds whereas this class has a resolution of microsends. This class will truncate the nanosecond resolution down to the microsecond. """ if duration_proto.nanos % 1000 != 0: # TODO(https://github.com/apache/beam/issues/19922): Better define # durations. raise ValueError( "Cannot convert from nanoseconds to microseconds " + "because this loses precision. Please make sure that " + "this is the correct behavior you want and manually " + "truncate the precision to the nearest microseconds. " + "See [https://github.com/apache/beam/issues/19922] for " + "more information.") return Duration( seconds=duration_proto.seconds, micros=duration_proto.nanos // 1000) def __repr__(self) -> str: micros = self.micros sign = '' if micros < 0: sign = '-' micros = -micros int_part = micros // 1000000 frac_part = micros % 1000000 if frac_part: return 'Duration(%s%d.%06d)' % (sign, int_part, frac_part) return 'Duration(%s%d)' % (sign, int_part) def __float__(self) -> float: # Note that the returned value may have lost precision. return self.micros / 1000000 def __eq__(self, other: object) -> bool: # Allow comparisons between Duration and Timestamp values. if isinstance(other, (Duration, Timestamp)): return self.micros == other.micros elif isinstance(other, (int, float)): return self.micros == Duration.of(other).micros else: # Support equality with other types return NotImplemented def __lt__(self, other: TimestampDurationTypes) -> bool: # Allow comparisons between Duration and Timestamp values. if not isinstance(other, Timestamp): other = Duration.of(other) return self.micros < other.micros def __gt__(self, other: TimestampDurationTypes) -> bool: return not (self < other or self == other) def __le__(self, other: TimestampDurationTypes) -> bool: return self < other or self == other def __ge__(self, other: TimestampDurationTypes) -> bool: return not self < other def __hash__(self) -> int: return hash(self.micros) def __neg__(self) -> 'Duration': return Duration(micros=-self.micros) def __add__(self, other: DurationTypes) -> 'Duration': if isinstance(other, Timestamp): # defer to Timestamp.__add__ return NotImplemented other = Duration.of(other) return Duration(micros=self.micros + other.micros) def __radd__(self, other: DurationTypes) -> 'Duration': return self + other def __sub__(self, other: DurationTypes) -> 'Duration': other = Duration.of(other) return Duration(micros=self.micros - other.micros) def __rsub__(self, other: DurationTypes) -> 'Duration': return -(self - other) def __mul__(self, other: DurationTypes) -> 'Duration': other = Duration.of(other) return Duration(micros=self.micros * other.micros // 1000000) def __rmul__(self, other: DurationTypes) -> 'Duration': return self * other def __mod__(self, other: DurationTypes) -> 'Duration': other = Duration.of(other) return Duration(micros=self.micros % other.micros) # The minimum granularity / interval expressible in a Timestamp / Duration # object. TIME_GRANULARITY = Duration(micros=1) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/timestamp_test.py0000644000175100017510000002475315206113033022445 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for time utilities.""" # pytype: skip-file import datetime import unittest import pytz from google.protobuf import duration_pb2 from google.protobuf import timestamp_pb2 from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import Timestamp class TimestampTest(unittest.TestCase): def test_of(self): interval = Timestamp(123) self.assertEqual(id(interval), id(Timestamp.of(interval))) self.assertEqual(interval, Timestamp.of(123.0)) with self.assertRaises(TypeError): Timestamp.of(Duration(10)) def test_precision(self): self.assertEqual(Timestamp(10000000) % 0.1, 0) self.assertEqual(Timestamp(10000000) % 0.05, 0) self.assertEqual(Timestamp(10000000) % 0.000005, 0) self.assertEqual(Timestamp(10000000) % Duration(0.1), 0) self.assertEqual(Timestamp(10000000) % Duration(0.05), 0) self.assertEqual(Timestamp(10000000) % Duration(0.000005), 0) def test_utc_timestamp(self): self.assertEqual(Timestamp(10000000).to_rfc3339(), '1970-04-26T17:46:40Z') self.assertEqual( Timestamp(10000000.000001).to_rfc3339(), '1970-04-26T17:46:40.000001Z') self.assertEqual( Timestamp(1458343379.123456).to_rfc3339(), '2016-03-18T23:22:59.123456Z') def test_from_rfc3339(self): test_cases = [ (10000000, '1970-04-26T17:46:40Z'), (10000000.000001, '1970-04-26T17:46:40.000001Z'), (1458343379.123456, '2016-03-18T23:22:59.123456Z'), ] for seconds_float, rfc3339_str in test_cases: self.assertEqual( Timestamp(seconds_float), Timestamp.from_rfc3339(rfc3339_str)) self.assertEqual( rfc3339_str, Timestamp.from_rfc3339(rfc3339_str).to_rfc3339()) def test_from_rfc3339_with_timezone(self): test_cases = [ (1458328979.123456, '2016-03-18T23:22:59.123456+04:00'), (1458357779.123456, '2016-03-18T23:22:59.123456-04:00'), ] for seconds_float, rfc3339_str in test_cases: self.assertEqual( Timestamp(seconds_float), Timestamp.from_rfc3339(rfc3339_str)) def test_from_rfc3339_failure(self): with self.assertRaisesRegex(ValueError, 'parse'): Timestamp.from_rfc3339('not rfc3339') with self.assertRaisesRegex(ValueError, 'parse'): Timestamp.from_rfc3339('2016-03-18T23:22:59.123456Z unparseable') def test_from_utc_datetime(self): self.assertEqual( Timestamp.from_utc_datetime( datetime.datetime(1970, 1, 1, tzinfo=pytz.utc)), Timestamp(0)) with self.assertRaisesRegex(ValueError, r'UTC'): Timestamp.from_utc_datetime( datetime.datetime(1970, 1, 1, tzinfo=pytz.timezone('US/Eastern'))) with self.assertRaisesRegex(ValueError, r'dt has no timezone info'): Timestamp.from_utc_datetime(datetime.datetime(1970, 1, 1, tzinfo=None)) def test_from_to_utc_datetime(self): timestamp = Timestamp(seconds=1458343379.123456) dt = timestamp.to_utc_datetime(has_tz=True) self.assertEqual(timestamp, Timestamp.from_utc_datetime(dt)) def test_arithmetic(self): # Supported operations. self.assertEqual(Timestamp(123) + 456, 579) self.assertEqual(Timestamp(123) + Duration(456), 579) self.assertEqual(456 + Timestamp(123), 579) self.assertEqual(Duration(456) + Timestamp(123), 579) self.assertEqual(Timestamp(123) - 456, -333) self.assertEqual(Timestamp(123) - Duration(456), -333) self.assertEqual(Timestamp(1230) % 456, 318) self.assertEqual(Timestamp(1230) % Duration(456), 318) self.assertEqual(Timestamp(123) - Timestamp(100), 23) # Check that direct comparison of Timestamp and Duration is allowed. self.assertTrue(Duration(123) == Timestamp(123)) self.assertTrue(Timestamp(123) == Duration(123)) self.assertFalse(Duration(123) == Timestamp(1230)) self.assertFalse(Timestamp(123) == Duration(1230)) # Check return types. self.assertEqual((Timestamp(123) + 456).__class__, Timestamp) self.assertEqual((Timestamp(123) + Duration(456)).__class__, Timestamp) self.assertEqual((456 + Timestamp(123)).__class__, Timestamp) self.assertEqual((Duration(456) + Timestamp(123)).__class__, Timestamp) self.assertEqual((Timestamp(123) - 456).__class__, Timestamp) self.assertEqual((Timestamp(123) - Duration(456)).__class__, Timestamp) self.assertEqual((Timestamp(1230) % 456).__class__, Duration) self.assertEqual((Timestamp(1230) % Duration(456)).__class__, Duration) self.assertEqual((Timestamp(123) - Timestamp(100)).__class__, Duration) # Unsupported operations. with self.assertRaises(TypeError): self.assertEqual(Timestamp(123) * 456, 56088) with self.assertRaises(TypeError): self.assertEqual(Timestamp(123) * Duration(456), 56088) with self.assertRaises(TypeError): self.assertEqual(456 * Timestamp(123), 56088) with self.assertRaises(TypeError): self.assertEqual(Duration(456) * Timestamp(123), 56088) with self.assertRaises(TypeError): self.assertEqual(456 - Timestamp(123), 333) with self.assertRaises(TypeError): self.assertEqual(Duration(456) - Timestamp(123), 333) with self.assertRaises(TypeError): self.assertEqual(-Timestamp(123), -123) # pylint: disable=invalid-unary-operand-type with self.assertRaises(TypeError): self.assertEqual(-Timestamp(123), -Duration(123)) # pylint: disable=invalid-unary-operand-type with self.assertRaises(TypeError): self.assertEqual(1230 % Timestamp(456), 318) with self.assertRaises(TypeError): self.assertEqual(Duration(1230) % Timestamp(456), 318) def test_sort_order(self): self.assertEqual([-63, Timestamp(-3), 2, 9, Timestamp(292.3), 500], sorted([9, 2, Timestamp(-3), Timestamp(292.3), -63, 500])) self.assertEqual([4, 5, Timestamp(6), Timestamp(7), 8, 9], sorted([9, 8, Timestamp(7), Timestamp(6), 5, 4])) def test_str(self): self.assertEqual('Timestamp(1.234567)', str(Timestamp(1.234567))) self.assertEqual('Timestamp(-1.234567)', str(Timestamp(-1.234567))) self.assertEqual( 'Timestamp(-999999999.900000)', str(Timestamp(-999999999.9))) self.assertEqual('Timestamp(999999999)', str(Timestamp(999999999))) self.assertEqual('Timestamp(-999999999)', str(Timestamp(-999999999))) def test_now(self): now = Timestamp.now() self.assertTrue(isinstance(now, Timestamp)) def test_from_proto(self): ts_proto = timestamp_pb2.Timestamp(seconds=1234, nanos=56000) actual_ts = Timestamp.from_proto(ts_proto) expected_ts = Timestamp(seconds=1234, micros=56) self.assertEqual(actual_ts, expected_ts) def test_from_proto_fails_with_truncation(self): # TODO(https://github.com/apache/beam/issues/19922): Better define # timestamps. with self.assertRaises(ValueError): Timestamp.from_proto(timestamp_pb2.Timestamp(seconds=1234, nanos=56789)) def test_to_proto(self): ts = Timestamp(seconds=1234, micros=56) actual_ts_proto = Timestamp.to_proto(ts) expected_ts_proto = timestamp_pb2.Timestamp(seconds=1234, nanos=56000) self.assertEqual(actual_ts_proto, expected_ts_proto) def test_equality(self): for min_val in (Timestamp(1), Duration(1), 1, 1.1): for max_val in (Timestamp(123), Duration(123), 123, 123.4): self.assertTrue(min_val < max_val, "%s < %s" % (min_val, max_val)) self.assertTrue(min_val <= max_val, "%s <= %s" % (min_val, max_val)) self.assertTrue(max_val > min_val, "%s > %s" % (max_val, min_val)) self.assertTrue(max_val >= min_val, "%s >= %s" % (max_val, min_val)) class DurationTest(unittest.TestCase): def test_of(self): interval = Duration(123) self.assertEqual(id(interval), id(Duration.of(interval))) self.assertEqual(interval, Duration.of(123.0)) with self.assertRaises(TypeError): Duration.of(Timestamp(10)) def test_precision(self): self.assertEqual(Duration(10000000) % 0.1, 0) self.assertEqual(Duration(10000000) % 0.05, 0) self.assertEqual(Duration(10000000) % 0.000005, 0) def test_arithmetic(self): self.assertEqual(Duration(123) + 456, 579) self.assertEqual(456 + Duration(123), 579) self.assertEqual(Duration(123) * 456, 56088) self.assertEqual(456 * Duration(123), 56088) self.assertEqual(Duration(123) - 456, -333) self.assertEqual(456 - Duration(123), 333) self.assertEqual(-Duration(123), -123) def test_sort_order(self): self.assertEqual([-63, Duration(-3), 2, 9, Duration(292.3), 500], sorted([9, 2, Duration(-3), Duration(292.3), -63, 500])) self.assertEqual([4, 5, Duration(6), Duration(7), 8, 9], sorted([9, 8, Duration(7), Duration(6), 5, 4])) def test_str(self): self.assertEqual('Duration(1.234567)', str(Duration(1.234567))) self.assertEqual('Duration(-1.234567)', str(Duration(-1.234567))) self.assertEqual('Duration(-999999999.900000)', str(Duration(-999999999.9))) self.assertEqual('Duration(999999999)', str(Duration(999999999))) self.assertEqual('Duration(-999999999)', str(Duration(-999999999))) def test_from_proto(self): dur_proto = duration_pb2.Duration(seconds=1234, nanos=56000) actual_dur = Duration.from_proto(dur_proto) expected_dur = Duration(seconds=1234, micros=56) self.assertEqual(actual_dur, expected_dur) def test_from_proto_fails_with_truncation(self): # TODO(https://github.com/apache/beam/issues/19922): Better define # durations. with self.assertRaises(ValueError): Duration.from_proto(duration_pb2.Duration(seconds=1234, nanos=56789)) def test_to_proto(self): dur = Duration(seconds=1234, micros=56) actual_dur_proto = Duration.to_proto(dur) expected_dur_proto = duration_pb2.Duration(seconds=1234, nanos=56000) self.assertEqual(actual_dur_proto, expected_dur_proto) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/transform_service_launcher.py0000644000175100017510000002317615206113033025015 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import argparse import logging import os import shutil import subprocess import sys import tempfile import threading import time import zipfile from pathlib import Path import grpc from apache_beam.utils import subprocess_server _LOGGER = logging.getLogger(__name__) _COMMAND_POSSIBLE_VALUES = ['up', 'down', 'ps'] _EXPANSION_SERVICE_LAUNCHER_JAR = ':sdks:java:transform-service:app:build' def is_docker_compose_v2_available(): cmd = ['docker', 'compose', 'version'] try: subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) except: # pylint: disable=bare-except return False return True class TransformServiceLauncher(object): _DEFAULT_PROJECT_NAME = 'apache.beam.transform.service' _DEFAULT_START_WAIT_TIMEOUT = 50000 _launchers = {} # type: ignore # Maintaining a static list of launchers to prevent temporary resources # from being created unnecessarily. def __new__(cls, project_name, port, beam_version=None, user_agent=None): if project_name not in TransformServiceLauncher._launchers: TransformServiceLauncher._launchers[project_name] = super( TransformServiceLauncher, cls).__new__(cls) return TransformServiceLauncher._launchers[project_name] def __init__(self, project_name, port, beam_version=None, user_agent=None): logging.info('Initializing the Beam Transform Service %s.' % project_name) self._project_name = project_name self._port = port self._address = 'localhost:' + str(self._port) self._launcher_lock = threading.RLock() # Setting up Docker Compose configuration. # We use Docker Compose project name as the name of the temporary directory # to isolate different transform service instances that may be running in # the same machine. temp_dir = os.path.join(tempfile.gettempdir(), project_name) if not os.path.exists(temp_dir): os.mkdir(temp_dir) # Get the jar with configs path_to_local_jar = subprocess_server.JavaJarServer.local_jar( subprocess_server.JavaJarServer.path_to_beam_jar( _EXPANSION_SERVICE_LAUNCHER_JAR), user_agent=user_agent) with zipfile.ZipFile(path_to_local_jar) as launcher_jar: launcher_jar.extract('docker-compose.yml', path=temp_dir) launcher_jar.extract('.env', path=temp_dir) compose_file = os.path.join(temp_dir, 'docker-compose.yml') # Creating the credentials volume. credentials_dir = os.path.join(temp_dir, 'credentials_dir') if not os.path.exists(credentials_dir): os.mkdir(credentials_dir) logging.info('Copying the Google Application Default Credentials file.') is_windows = 'windows' in os.name.lower() application_default_path_suffix = ( '\\gcloud\\application_default_credentials.json' if is_windows else '.config/gcloud/application_default_credentials.json') application_default_path_file = os.path.join( str(Path.home()), application_default_path_suffix) application_default_path_copied = os.path.join( credentials_dir, 'application_default_credentials.json') if os.path.exists(application_default_path_file): shutil.copyfile( application_default_path_file, application_default_path_copied) else: logging.info( 'GCP credentials will not be available for the transform service ' 'since could not find the Google Cloud application default ' 'credentials file at the expected location %s.' % application_default_path_file) # Creating the dependencies volume. dependencies_dir = os.path.join(temp_dir, 'dependencies_dir') if not os.path.exists(dependencies_dir): os.mkdir(dependencies_dir) self._environmental_variables = {} self._environmental_variables['CREDENTIALS_VOLUME'] = credentials_dir self._environmental_variables['DEPENDENCIES_VOLUME'] = dependencies_dir self._environmental_variables['TRANSFORM_SERVICE_PORT'] = str(port) self._environmental_variables['BEAM_VERSION'] = beam_version # Setting an empty requirements file requirements_file_name = os.path.join(dependencies_dir, 'requirements.txt') with open(requirements_file_name, 'w') as _: pass self._environmental_variables['PYTHON_REQUIREMENTS_FILE_NAME'] = ( 'requirements.txt') # Building docker compose start command if is_docker_compose_v2_available(): self._docker_compose_start_command_prefix = ['docker', 'compose'] else: self._docker_compose_start_command_prefix = ['docker-compose'] self._docker_compose_start_command_prefix.append('-p') self._docker_compose_start_command_prefix.append(project_name) self._docker_compose_start_command_prefix.append('-f') self._docker_compose_start_command_prefix.append(compose_file) def _get_channel(self): channel_options = [("grpc.max_receive_message_length", -1), ("grpc.max_send_message_length", -1)] if hasattr(grpc, 'local_channel_credentials'): # TODO: update this to support secure non-local channels. return grpc.secure_channel( self._address, grpc.local_channel_credentials(), options=channel_options) else: return grpc.insecure_channel(self._address, options=channel_options) def __enter__(self): self.start() self.wait_till_up(-1) self._channel = self._get_channel() from apache_beam import external return external.ExpansionAndArtifactRetrievalStub(self._channel.__enter__()) def __exit__(self, *args): self.shutdown() self._channel.__exit__(*args) def _run_docker_compose_command(self, command, output_override=None): cmd = [] cmd.extend(self._docker_compose_start_command_prefix) cmd.extend(command) myenv = os.environ.copy() myenv.update(self._environmental_variables) process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=myenv) std_out, _ = process.communicate() if output_override: output_override.write(std_out) else: print(std_out.decode(errors='backslashreplace')) def start(self): with self._launcher_lock: self._run_docker_compose_command(['up', '-d']) def shutdown(self): with self._launcher_lock: self._run_docker_compose_command(['down']) def status(self): with self._launcher_lock: self._run_docker_compose_command(['ps']) def wait_till_up(self, timeout_ms): channel = self._get_channel() timeout_ms = ( TransformServiceLauncher._DEFAULT_START_WAIT_TIMEOUT if timeout_ms <= 0 else timeout_ms) # Waiting till the service is up. channel_ready = grpc.channel_ready_future(channel) wait_secs = .1 start_time = time.time() while True: if (time.time() - start_time) * 1000 > timeout_ms > 0: raise ValueError( 'Transform service did not start in %s seconds.' % (timeout_ms / 1000)) try: channel_ready.result(timeout=wait_secs) break except (grpc.FutureTimeoutError, grpc.RpcError): wait_secs *= 1.2 logging.log( logging.WARNING if wait_secs > 1 else logging.DEBUG, 'Waiting for the transform service to be ready at %s.', self._address) logging.info('Transform service ' + self._project_name + ' started.') def _get_status(self): tmp = tempfile.NamedTemporaryFile(delete=False) self._run_docker_compose_command(['ps'], tmp) tmp.close() return tmp.name def main(argv): parser = argparse.ArgumentParser() parser.add_argument('--project_name', help='Docker Compose project name.') parser.add_argument( '--command', required=True, choices=_COMMAND_POSSIBLE_VALUES, help='Command to run. Possible values are ' + ', '.join(_COMMAND_POSSIBLE_VALUES)) parser.add_argument( '--port', type=int, default=-1, help='External visible port of the transform service.') parser.add_argument( '--beam_version', required=True, help='Beam version of the expansion service containers to be used.') known_args, _ = parser.parse_known_args(argv) project_name = ( TransformServiceLauncher._DEFAULT_PROJECT_NAME if known_args.project_name is None else known_args.project_name) logging.info( 'Starting the Beam Transform Service at %s.' % ( 'the default port' if known_args.port < 0 else (' port ' + str(known_args.port)))) launcher = TransformServiceLauncher( project_name, known_args.port, known_args.beam_version) if known_args.command == 'up': launcher.start() launcher.wait_till_up(-1) elif known_args.command == 'down': launcher.shutdown() elif known_args.command == 'ps': launcher.status() else: raise ValueError( 'Unknown command %s possible values are %s' % (known_args.command, ', '.join(_COMMAND_POSSIBLE_VALUES))) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) main(sys.argv) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/urns.py0000644000175100017510000001347515206113033020371 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """For internal use only; no backwards-compatibility guarantees.""" # pytype: skip-file # TODO(https://github.com/apache/beam/issues/18399): Issue with dill + local # classes + abc metaclass # import abc import inspect from typing import TYPE_CHECKING from typing import Any from typing import Callable from typing import Optional from typing import TypeVar from typing import Union from typing import overload from google.protobuf import message from google.protobuf import wrappers_pb2 from apache_beam.internal import pickler from apache_beam.portability.api import beam_runner_api_pb2 from apache_beam.utils import proto_utils if TYPE_CHECKING: from apache_beam.runners.pipeline_context import PipelineContext T = TypeVar('T') RunnerApiFnT = TypeVar('RunnerApiFnT', bound='RunnerApiFn') ConstructorFn = Callable[[Union['message.Message', bytes], 'PipelineContext'], Any] class RunnerApiFn(object): """Abstract base class that provides urn registration utilities. A class that inherits from this class will get a registration-based from_runner_api and to_runner_api method that convert to and from beam_runner_api_pb2.FunctionSpec. Additionally, register_pickle_urn can be called from the body of a class to register serialization via pickling. """ # TODO(https://github.com/apache/beam/issues/18399): Issue with dill + local # classes + abc metaclass # __metaclass__ = abc.ABCMeta _known_urns: dict[str, tuple[Optional[type], ConstructorFn]] = {} # @abc.abstractmethod is disabled here to avoid an error with mypy. mypy # performs abc.abtractmethod/property checks even if a class does # not use abc.ABCMeta, however, functions like `register_pickle_urn` # dynamically patch `to_runner_api_parameter`, which mypy cannot track, so # mypy incorrectly infers that this method has not been overridden with a # concrete implementation. # @abc.abstractmethod def to_runner_api_parameter( self, unused_context: 'PipelineContext') -> tuple[str, Any]: """Returns the urn and payload for this Fn. The returned urn(s) should be registered with `register_urn`. """ raise NotImplementedError @classmethod @overload def register_urn( cls, urn: str, parameter_type: type[T], ) -> Callable[[Callable[[T, 'PipelineContext'], Any]], Callable[[T, 'PipelineContext'], Any]]: pass @classmethod @overload def register_urn( cls, urn: str, parameter_type: None, ) -> Callable[[Callable[[bytes, 'PipelineContext'], Any]], Callable[[bytes, 'PipelineContext'], Any]]: pass @classmethod @overload def register_urn( cls, urn: str, parameter_type: type[T], fn: Callable[[T, 'PipelineContext'], Any]) -> None: pass @classmethod @overload def register_urn( cls, urn: str, parameter_type: None, fn: Callable[[bytes, 'PipelineContext'], Any]) -> None: pass @classmethod def register_urn(cls, urn, parameter_type, fn=None): """Registers a urn with a constructor. For example, if 'beam:fn:foo' had parameter type FooPayload, one could write `RunnerApiFn.register_urn('bean:fn:foo', FooPayload, foo_from_proto)` where foo_from_proto took as arguments a FooPayload and a PipelineContext. This function can also be used as a decorator rather than passing the callable in as the final parameter. A corresponding to_runner_api_parameter method would be expected that returns the tuple ('beam:fn:foo', FooPayload) """ def register(fn): cls._known_urns[urn] = parameter_type, fn return fn if fn: # Used as a statement. register(fn) else: # Used as a decorator. return register @classmethod def register_pickle_urn(cls, pickle_urn): """Registers and implements the given urn via pickling. """ inspect.currentframe().f_back.f_locals['to_runner_api_parameter'] = ( lambda self, context: (pickle_urn, wrappers_pb2.BytesValue(value=pickler.dumps(self)))) cls.register_urn( pickle_urn, wrappers_pb2.BytesValue, lambda proto, unused_context: pickler.loads(proto.value)) def to_runner_api( self, context: 'PipelineContext') -> beam_runner_api_pb2.FunctionSpec: """Returns an FunctionSpec encoding this Fn. Prefer overriding self.to_runner_api_parameter. """ urn, typed_param = self.to_runner_api_parameter(context) return beam_runner_api_pb2.FunctionSpec( urn=urn, payload=typed_param.SerializeToString() if isinstance( typed_param, message.Message) else typed_param) @classmethod def from_runner_api( cls: type[RunnerApiFnT], fn_proto: beam_runner_api_pb2.FunctionSpec, context: 'PipelineContext') -> RunnerApiFnT: """Converts from an FunctionSpec to a Fn object. Prefer registering a urn with its parameter type and constructor. """ parameter_type, constructor = cls._known_urns[fn_proto.urn] return constructor( proto_utils.parse_Bytes(fn_proto.payload, parameter_type), context) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/windowed_value.pxd0000644000175100017510000000362215206113033022552 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # cimport cython from libc.stdint cimport int64_t cdef type Timestamp cdef list _BYTE_TO_PANE_INFO @cython.final cdef class PaneInfo(object): cdef readonly bint _is_first cdef readonly bint _is_last cdef readonly int _timing cdef readonly int _index cdef readonly int _nonspeculative_index cdef readonly unsigned char _encoded_byte @cython.final cdef class WindowedValue(object): cdef public object value cdef public object windows cdef public object pane_info cdef public int64_t timestamp_micros cdef object timestamp_object cpdef WindowedValue with_value(self, new_value) cdef class WindowedBatch(object): cpdef WindowedBatch with_values(self, object new_values) cdef class HomogeneousWindowedBatch(WindowedBatch): cdef public WindowedValue _wv cpdef WindowedBatch with_values(self, object new_values) @cython.locals(wv=WindowedValue) cpdef WindowedValue create( object value, int64_t timestamp_micros, object windows, object pane_info=*) cdef class _IntervalWindowBase(object): cdef object _start_object cdef readonly int64_t _start_micros cdef object _end_object cdef readonly int64_t _end_micros ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/windowed_value.py0000644000175100017510000003432715206113033022415 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Core windowing data structures.""" # This module is carefully crafted to have optimal performance when # compiled while still being valid Python. Care needs to be taken when # editing this file as WindowedValues are created for every element for # every step in a Beam pipeline. # pytype: skip-file # ruff: noqa: UP006 import collections from typing import TYPE_CHECKING from typing import Any from typing import Callable from typing import Iterable from typing import List from typing import Optional from typing import Sequence from typing import Tuple from apache_beam.utils.timestamp import MAX_TIMESTAMP from apache_beam.utils.timestamp import MIN_TIMESTAMP from apache_beam.utils.timestamp import Timestamp from apache_beam.utils.timestamp import TimestampTypes # pylint: disable=unused-import if TYPE_CHECKING: from apache_beam.transforms.window import BoundedWindow class PaneInfoTiming(object): """The timing of a PaneInfo.""" EARLY = 0 ON_TIME = 1 LATE = 2 UNKNOWN = 3 @classmethod def to_string(cls, value): return { cls.EARLY: 'EARLY', cls.ON_TIME: 'ON_TIME', cls.LATE: 'LATE', cls.UNKNOWN: 'UNKNOWN', }[value] @classmethod def from_string(cls, value): return { 'EARLY': cls.EARLY, 'ON_TIME': cls.ON_TIME, 'LATE': cls.LATE, 'UNKNOWN': cls.UNKNOWN }[value] class PaneInfo(object): """Describes the trigger firing information for a given WindowedValue. "Panes" represent individual firings on a single window. ``PaneInfo``s are passed downstream after trigger firings. They contain information about whether it's an early/on time/late firing, if it's the last or first firing from a window, and the index of the firing. """ def __init__(self, is_first, is_last, timing, index, nonspeculative_index): self._is_first = is_first self._is_last = is_last self._timing = timing self._index = index self._nonspeculative_index = nonspeculative_index self._encoded_byte = self._get_encoded_byte() def _get_encoded_byte(self): byte = 0 if self._is_first: byte |= 1 if self._is_last: byte |= 2 byte |= self._timing << 2 return byte @staticmethod def from_encoded_byte(encoded_byte): assert encoded_byte in _BYTE_TO_PANE_INFO return _BYTE_TO_PANE_INFO[encoded_byte] # Because common PaneInfo objects are cached, it is important that the value # is immutable. We therefore explicitly enforce this here with read-only # properties. @property def is_first(self): return self._is_first @property def is_last(self): return self._is_last @property def timing(self): return self._timing @property def index(self): # type: () -> int return self._index @property def nonspeculative_index(self): # type: () -> int return self._nonspeculative_index @property def encoded_byte(self): # type: () -> int return self._encoded_byte def __repr__(self): return ( 'PaneInfo(first: %r, last: %r, timing: %s, index: %d, ' 'nonspeculative_index: %d)') % ( self.is_first, self.is_last, PaneInfoTiming.to_string(self.timing), self.index, self.nonspeculative_index) def __eq__(self, other): if self is other: return True if isinstance(other, PaneInfo): return ( self.is_first == other.is_first and self.is_last == other.is_last and self.timing == other.timing and self.index == other.index and self.nonspeculative_index == other.nonspeculative_index) return NotImplemented def __hash__(self): return hash(( self.is_first, self.is_last, self.timing, self.index, self.nonspeculative_index)) def __reduce__(self): return PaneInfo, (self._is_first, self._is_last, self._timing, self._index, self._nonspeculative_index) def _construct_well_known_pane_infos(): # type: () -> List[PaneInfo] pane_infos = [] for timing in (PaneInfoTiming.EARLY, PaneInfoTiming.ON_TIME, PaneInfoTiming.LATE, PaneInfoTiming.UNKNOWN): nonspeculative_index = -1 if timing == PaneInfoTiming.EARLY else 0 pane_infos.append(PaneInfo(True, True, timing, 0, nonspeculative_index)) pane_infos.append(PaneInfo(True, False, timing, 0, nonspeculative_index)) pane_infos.append(PaneInfo(False, True, timing, -1, nonspeculative_index)) pane_infos.append(PaneInfo(False, False, timing, -1, nonspeculative_index)) result = [None] * ( max(p.encoded_byte for p in pane_infos) + 1 ) # type: List[PaneInfo] # type: ignore[list-item] for pane_info in pane_infos: result[pane_info.encoded_byte] = pane_info return result # Cache of well-known PaneInfo objects. _BYTE_TO_PANE_INFO = _construct_well_known_pane_infos() # Default PaneInfo descriptor for when a value is not the output of triggering. PANE_INFO_UNKNOWN = _BYTE_TO_PANE_INFO[0xF] class WindowedValue(object): """A windowed value having a value, a timestamp and set of windows. Attributes: value: The underlying value of a windowed value. timestamp: Timestamp associated with the value as seconds since Unix epoch. windows: A set (iterable) of window objects for the value. The window object are descendants of the BoundedWindow class. pane_info: A PaneInfo descriptor describing the triggering information for the pane that contained this value. If None, will be set to PANE_INFO_UNKNOWN. """ def __init__( self, value, timestamp, # type: TimestampTypes windows, # type: Tuple[BoundedWindow, ...] pane_info=PANE_INFO_UNKNOWN # type: PaneInfo ): # type: (...) -> None # For performance reasons, only timestamp_micros is stored by default # (as a C int). The Timestamp object is created on demand below. self.value = value if isinstance(timestamp, int): self.timestamp_micros = timestamp * 1000000 if TYPE_CHECKING: self.timestamp_object = None # type: Optional[Timestamp] else: self.timestamp_object = ( timestamp if isinstance(timestamp, Timestamp) else Timestamp.of(timestamp)) self.timestamp_micros = self.timestamp_object.micros self.windows = windows self.pane_info = pane_info @property def timestamp(self): # type: () -> Timestamp if self.timestamp_object is None: self.timestamp_object = Timestamp(0, self.timestamp_micros) return self.timestamp_object def __repr__(self): return '(%s, %s, %s, %s)' % ( repr(self.value), 'MIN_TIMESTAMP' if self.timestamp == MIN_TIMESTAMP else 'MAX_TIMESTAMP' if self.timestamp == MAX_TIMESTAMP else float(self.timestamp), self.windows, self.pane_info) def __eq__(self, other): if isinstance(other, WindowedValue): return ( type(self) == type(other) and self.timestamp_micros == other.timestamp_micros and self.value == other.value and self.windows == other.windows and self.pane_info == other.pane_info) return NotImplemented def __hash__(self): return ((hash(self.value) & 0xFFFFFFFFFFFFFFF) + 3 * (self.timestamp_micros & 0xFFFFFFFFFFFFFF) + 7 * (hash(tuple(self.windows)) & 0xFFFFFFFFFFFFF) + 11 * (hash(self.pane_info) & 0xFFFFFFFFFFFFF)) def with_value(self, new_value): # type: (Any) -> WindowedValue """Creates a new WindowedValue with the same timestamps and windows as this. This is the fasted way to create a new WindowedValue. """ return create( new_value, self.timestamp_micros, self.windows, self.pane_info) def __reduce__(self): return WindowedValue, ( self.value, self.timestamp, self.windows, self.pane_info) # TODO(robertwb): Move this to a static method. def create(value, timestamp_micros, windows, pane_info=PANE_INFO_UNKNOWN): wv = WindowedValue.__new__(WindowedValue) wv.value = value wv.timestamp_micros = timestamp_micros wv.windows = windows wv.pane_info = pane_info return wv class WindowedBatch(object): """A batch of N windowed values, each having a value, a timestamp and set of windows.""" def with_values(self, new_values): # type: (Any) -> WindowedBatch """Creates a new WindowedBatch with the same timestamps and windows as this. This is the fasted way to create a new WindowedValue. """ raise NotImplementedError def as_windowed_values(self, explode_fn: Callable) -> Iterable[WindowedValue]: raise NotImplementedError @staticmethod def from_windowed_values( windowed_values: Sequence[WindowedValue], *, produce_fn: Callable) -> Iterable['WindowedBatch']: return HomogeneousWindowedBatch.from_windowed_values( windowed_values, produce_fn=produce_fn) class HomogeneousWindowedBatch(WindowedBatch): """A WindowedBatch with Homogeneous event-time information, represented internally as a WindowedValue. """ def __init__(self, wv): self._wv = wv @staticmethod def of(values, timestamp, windows, pane_info): return HomogeneousWindowedBatch( WindowedValue(values, timestamp, windows, pane_info)) @property def values(self): return self._wv.value @property def timestamp(self): return self._wv.timestamp @property def pane_info(self): return self._wv.pane_info @property def windows(self): return self._wv.windows @windows.setter def windows(self, value): self._wv.windows = value def with_values(self, new_values): # type: (Any) -> WindowedBatch return HomogeneousWindowedBatch(self._wv.with_value(new_values)) def as_windowed_values(self, explode_fn: Callable) -> Iterable[WindowedValue]: for value in explode_fn(self._wv.value): yield self._wv.with_value(value) def as_empty_windowed_value(self): """Get a single WindowedValue with identical windowing information to this HomogeneousWindowedBatch, but with value=None. Useful for re-using APIs that pull windowing information from a WindowedValue.""" return self._wv.with_value(None) def __eq__(self, other): if isinstance(other, HomogeneousWindowedBatch): return self._wv == other._wv return NotImplemented def __hash__(self): return hash(self._wv) @staticmethod def from_batch_and_windowed_value( *, batch, windowed_value: WindowedValue) -> 'WindowedBatch': return HomogeneousWindowedBatch(windowed_value.with_value(batch)) @staticmethod def from_windowed_values( windowed_values: Sequence[WindowedValue], *, produce_fn: Callable) -> Iterable['WindowedBatch']: grouped = collections.defaultdict(lambda: []) for wv in windowed_values: grouped[wv.with_value(None)].append(wv.value) for key, values in grouped.items(): yield HomogeneousWindowedBatch(key.with_value(produce_fn(values))) try: setattr(WindowedValue, 'timestamp_object', None) except TypeError: # When we're compiled, we can't dynamically add attributes to # the cdef class, but in this case it's OK as it's already present # on each instance. pass class _IntervalWindowBase(object): """Optimized form of IntervalWindow storing only microseconds for endpoints. """ def __init__(self, start, end): # type: (TimestampTypes, TimestampTypes) -> None if start is not None: self._start_object = Timestamp.of(start) # type: Optional[Timestamp] try: self._start_micros = self._start_object.micros except OverflowError: self._start_micros = ( MIN_TIMESTAMP.micros if self._start_object.micros < 0 else MAX_TIMESTAMP.micros) else: # Micros must be populated elsewhere. self._start_object = None if end is not None: self._end_object = Timestamp.of(end) # type: Optional[Timestamp] try: self._end_micros = self._end_object.micros except OverflowError: self._end_micros = ( MIN_TIMESTAMP.micros if self._end_object.micros < 0 else MAX_TIMESTAMP.micros) else: # Micros must be populated elsewhere. self._end_object = None @property def start(self): # type: () -> Timestamp if self._start_object is None: self._start_object = Timestamp(0, self._start_micros) return self._start_object @property def end(self): # type: () -> Timestamp if self._end_object is None: self._end_object = Timestamp(0, self._end_micros) return self._end_object def __hash__(self): # Cut off window at start/end timestamps for comparison purposes since some # portable runners do this already, and timestamps outside of the bands of # Min/Max timestamps are functionally equal to Min/Max. start = max(self._start_micros, MIN_TIMESTAMP.micros) end = min(self._end_micros, MAX_TIMESTAMP.micros) return hash((start, end)) def __eq__(self, other): if type(self) != type(other): return False # Cut off window at start/end timestamps for comparison purposes since some # portable runners do this already, and timestamps outside of the bands of # Min/Max timestamps are functionally equal to Min/Max. self_start = max(self._start_micros, MIN_TIMESTAMP.micros) self_end = min(self._end_micros, MAX_TIMESTAMP.micros) other_start = max(other._start_micros, MIN_TIMESTAMP.micros) other_end = min(other._end_micros, MAX_TIMESTAMP.micros) return (self_start == other_start and self_end == other_end) def __repr__(self): return '[%s, %s)' % (float(self.start), float(self.end)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/utils/windowed_value_test.py0000644000175100017510000001441215206113033023445 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Unit tests for the windowed_value.""" # pytype: skip-file import copy import itertools import pickle import unittest from parameterized import parameterized from parameterized import parameterized_class from apache_beam.utils import windowed_value from apache_beam.utils.timestamp import Timestamp class WindowedValueTest(unittest.TestCase): def test_timestamps(self): wv = windowed_value.WindowedValue(None, 3, ()) self.assertEqual(wv.timestamp, Timestamp.of(3)) self.assertTrue(wv.timestamp is wv.timestamp) self.assertEqual( windowed_value.WindowedValue(None, -2.5, ()).timestamp, Timestamp.of(-2.5)) def test_with_value(self): pane_info = windowed_value.PaneInfo( True, True, windowed_value.PaneInfoTiming.ON_TIME, 0, 0) wv = windowed_value.WindowedValue(1, 3, (), pane_info) self.assertEqual( wv.with_value(10), windowed_value.WindowedValue(10, 3, (), pane_info)) def test_equality(self): self.assertEqual( windowed_value.WindowedValue(1, 3, ()), windowed_value.WindowedValue(1, 3, ())) self.assertNotEqual( windowed_value.WindowedValue(1, 3, ()), windowed_value.WindowedValue(100, 3, ())) self.assertNotEqual( windowed_value.WindowedValue(1, 3, ()), windowed_value.WindowedValue(1, 300, ())) self.assertNotEqual( windowed_value.WindowedValue(1, 3, ()), windowed_value.WindowedValue(1, 300, ((), ))) self.assertNotEqual(windowed_value.WindowedValue(1, 3, ()), object()) def test_hash(self): wv = windowed_value.WindowedValue(1, 3, ()) wv_copy = copy.copy(wv) self.assertFalse(wv is wv_copy) self.assertEqual({wv: 100}.get(wv_copy), 100) def test_pickle(self): pane_info = windowed_value.PaneInfo( True, True, windowed_value.PaneInfoTiming.ON_TIME, 0, 0) wv = windowed_value.WindowedValue(1, 3, (), pane_info) self.assertTrue(pickle.loads(pickle.dumps(wv)) == wv) WINDOWED_BATCH_INSTANCES = [ windowed_value.HomogeneousWindowedBatch.of( None, 3, (), windowed_value.PANE_INFO_UNKNOWN), windowed_value.HomogeneousWindowedBatch.of( None, 3, (), windowed_value.PaneInfo( True, False, windowed_value.PaneInfoTiming.ON_TIME, 0, 0)), ] class WindowedBatchTest(unittest.TestCase): def test_homogeneous_windowed_batch_with_values(self): pane_info = windowed_value.PaneInfo( True, True, windowed_value.PaneInfoTiming.ON_TIME, 0, 0) wb = windowed_value.HomogeneousWindowedBatch.of(['foo', 'bar'], 6, (), pane_info) self.assertEqual( wb.with_values(['baz', 'foo']), windowed_value.HomogeneousWindowedBatch.of(['baz', 'foo'], 6, (), pane_info)) def test_homogeneous_windowed_batch_as_windowed_values(self): pane_info = windowed_value.PaneInfo( True, True, windowed_value.PaneInfoTiming.ON_TIME, 0, 0) wb = windowed_value.HomogeneousWindowedBatch.of(['foo', 'bar'], 3, (), pane_info) self.assertEqual( list(wb.as_windowed_values(iter)), [ windowed_value.WindowedValue('foo', 3, (), pane_info), windowed_value.WindowedValue('bar', 3, (), pane_info) ]) @parameterized.expand(itertools.combinations(WINDOWED_BATCH_INSTANCES, 2)) def test_inequality(self, left_wb, right_wb): self.assertNotEqual(left_wb, right_wb) def test_equals_different_type(self): wb = windowed_value.HomogeneousWindowedBatch.of( None, 3, (), windowed_value.PANE_INFO_UNKNOWN) self.assertNotEqual(wb, object()) def test_homogeneous_from_windowed_values(self): pane_info = windowed_value.PaneInfo( True, True, windowed_value.PaneInfoTiming.ON_TIME, 0, 0) windowed_values = [ windowed_value.WindowedValue('foofoo', 3, (), pane_info), windowed_value.WindowedValue('foobar', 6, (), pane_info), windowed_value.WindowedValue('foobaz', 9, (), pane_info), windowed_value.WindowedValue('barfoo', 3, (), pane_info), windowed_value.WindowedValue('barbar', 6, (), pane_info), windowed_value.WindowedValue('barbaz', 9, (), pane_info), windowed_value.WindowedValue('bazfoo', 3, (), pane_info), windowed_value.WindowedValue('bazbar', 6, (), pane_info), windowed_value.WindowedValue('bazbaz', 9, (), pane_info), ] self.assertEqual( list( windowed_value.WindowedBatch.from_windowed_values( windowed_values, produce_fn=list)), [ windowed_value.HomogeneousWindowedBatch.of( ['foofoo', 'barfoo', 'bazfoo'], 3, (), pane_info), windowed_value.HomogeneousWindowedBatch.of( ['foobar', 'barbar', 'bazbar'], 6, (), pane_info), windowed_value.HomogeneousWindowedBatch.of( ['foobaz', 'barbaz', 'bazbaz'], 9, (), pane_info) ]) @parameterized_class(('wb', ), [(wb, ) for wb in WINDOWED_BATCH_INSTANCES]) class WindowedBatchUtilitiesTest(unittest.TestCase): def test_hash(self): wb_copy = copy.copy(self.wb) self.assertFalse(self.wb is wb_copy) self.assertEqual({self.wb: 100}.get(wb_copy), 100) def test_pickle(self): self.assertTrue(pickle.loads(pickle.dumps(self.wb)) == self.wb) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/version.py0000644000175100017510000000154215206113033017717 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Apache Beam SDK version information and utilities.""" __version__ = '2.74.0' ././@PaxHeader0000000000000000000000000000003300000000000010211 xustar0027 mtime=1779996271.807453 apache_beam-2.74.0/apache_beam/yaml/0000755000175100017510000000000015206113160016621 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/__init__.py0000644000175100017510000000147715206113033020742 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from apache_beam.yaml.yaml_transform import * ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/cache_provider_artifacts.py0000644000175100017510000000335315206113033024213 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import sys import time from apache_beam.version import __version__ as beam_version from apache_beam.yaml import yaml_provider def cache_provider_artifacts(): providers_by_id = {} for providers in yaml_provider.standard_providers().values(): for provider in providers: # Dedup for better logging. providers_by_id[id(provider)] = provider for provider in providers_by_id.values(): t = time.time() artifacts = provider.cache_artifacts() if artifacts: logging.info( 'Cached %s in %0.03f seconds.', ', '.join(artifacts), time.time() - t) if '.dev' not in beam_version: # Also cache a base python venv for fast cloning. t = time.time() artifacts = yaml_provider.PypiExpansionService._create_venv_to_clone( sys.executable) logging.info('Cached %s in %0.03f seconds.', artifacts, time.time() - t) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) cache_provider_artifacts() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/conftest.py0000644000175100017510000000237515206113033021026 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # This global will be set by pytest_configure # and can be imported by other modules. # Initialize with a default that matches the default. yaml_test_files_dir = "tests" def pytest_addoption(parser): parser.addoption( "--test_files_dir", action="store", default="tests", help="Directory with YAML test files, relative to integration_tests.py") def pytest_configure(config): global yaml_test_files_dir yaml_test_files_dir = config.getoption("test_files_dir") ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.8134534 apache_beam-2.74.0/apache_beam/yaml/docs/0000755000175100017510000000000015206113160017551 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996269.0 apache_beam-2.74.0/apache_beam/yaml/docs/yaml-combine.md0000644000175100017510000000631615206113155022461 0ustar00runnerrunner--- type: languages title: "Apache Beam YAML Aggregations" --- # Beam YAML Aggregations Beam YAML has the ability to do aggregations to group and combine values across records. The is accomplished via the `Combine` transform type. For example, one can write ``` - type: Combine config: group_by: col1 combine: total: value: col2 fn: type: sum ``` If the function has no configuration requirements, it can be provided directly as a string ``` - type: Combine config: group_by: col1 combine: total: value: col2 fn: sum ``` This can be simplified further if the output field name is the same as the input field name ``` - type: Combine config: group_by: col1 combine: col2: sum ``` One can aggregate over many fields at once ``` - type: Combine config: group_by: col1 combine: col2: sum col3: max ``` and/or group by more than one field ``` - type: Combine config: group_by: [col1, col2] combine: col3: sum ``` or none at all (which will result in a global combine with a single output) ``` - type: Combine config: group_by: [] combine: col2: sum col3: max ``` ## Windowed aggregation As with all transforms, `Combine` can take a windowing parameter ``` - type: Combine windowing: type: fixed size: 60s config: group_by: col1 combine: col2: sum col3: max ``` If no windowing specification is provided, it inherits the windowing parameters from upstream, e.g. ``` - type: WindowInto windowing: type: fixed size: 60s - type: Combine config: group_by: col1 combine: col2: sum col3: max ``` is equivalent to the previous example. ## Custom aggregation functions One can use aggregation functions defined in Python by setting the language parameter. ``` - type: Combine config: language: python group_by: col1 combine: biggest: value: "col2 + col2" fn: type: 'apache_beam.transforms.combiners.TopCombineFn' config: n: 10 ``` ## SQL-style aggregations By setting the language to SQL, one can provide full SQL snippets as the combine fn. ``` - type: Combine config: language: sql group_by: col1 combine: num_values: "count(*)" total: "sum(col2)" ``` One can of course also use the `Sql` transform type and provide a query directly. ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996269.0 apache_beam-2.74.0/apache_beam/yaml/docs/yaml-errors.md0000644000175100017510000002135115206113155022355 0ustar00runnerrunner--- type: languages title: "Apache Beam YAML Error Handling" --- # Beam YAML Error Handling The larger one's pipeline gets, the more common it is to encounter "exceptional" data that is malformatted, doesn't handle the proper preconditions, or otherwise breaks during processing. Generally any such record will cause the pipeline to permanently fail, but often it is desirable to allow the pipeline to continue, re-directing bad records to another path for special handling or simply recording them for later off-line analysis. This is often called the "dead letter queue" pattern. Beam YAML has special support for this pattern if the transform supports a `error_handling` config parameter with an `output` field. The `output` parameter is a name that must referenced as an input to another transform that will process the errors (e.g. by writing them out). For example, the following code will write all "good" processed records to one file and any "bad" records, along with metadata about what error was encountered, to a separate file. ```yaml pipeline: transforms: - type: ReadFromCsv config: path: /path/to/input*.csv - type: MapToFields input: ReadFromCsv config: language: python fields: col1: col1 # This could raise a divide-by-zero error. ratio: col2 / col3 error_handling: output: my_error_output - type: WriteToJson input: MapToFields config: path: /path/to/output.json - type: WriteToJson name: WriteErrorsToJson input: MapToFields.my_error_output config: path: /path/to/errors.json ``` Note that with `error_handling` declared, `MapToFields.my_error_output` **must** be consumed; to ignore it will be an error. Any use is fine, e.g. logging the bad records to stdout would be sufficient (though not recommended for a robust pipeline). Note also that the exact format of the error outputs is still being finalized. They can be safely printed and written to outputs, but their precise schema may change in a future version of Beam and should not yet be depended on. It generally contains the failed record itself as well as information about the error that was encountered (e.g. error messages and tracebacks). To recover the bad record alone one can process the error output with the `StripErrorMetadata` transformation. Some transforms allow for extra arguments in their error_handling config, e.g. for Python functions one can give a `threshold` which limits the relative number of records that can be bad before considering the entire pipeline a failure ```yaml pipeline: transforms: - type: ReadFromCsv config: path: /path/to/input*.csv - type: MapToFields input: ReadFromCsv config: language: python fields: col1: col1 # This could raise a divide-by-zero error. ratio: col2 / col3 error_handling: output: my_error_output # If more than 10% of records throw an error, stop the pipeline. threshold: 0.1 - type: WriteToJson input: MapToFields config: path: /path/to/output.json - type: WriteToJson name: WriteErrorsToJson input: MapToFields.my_error_output config: path: /path/to/errors.json ``` One can do arbitrary further processing on these failed records if desired, e.g. ```yaml pipeline: transforms: - type: ReadFromCsv config: path: /path/to/input*.csv - type: MapToFields name: ComputeRatio input: ReadFromCsv config: language: python fields: col1: col1 # This could raise a divide-by-zero error. ratio: col2 / col3 error_handling: output: my_error_output - type: StripErrorMetadata name: FailedRecordsWithoutMetadata # Takes the error information from ComputeRatio and returns just the # failing records themselves for another attempt with a different # transform. input: ComputeRatio.my_error_output - type: MapToFields name: ComputeRatioForBadRecords input: FailedRecordsWithoutMetadata config: language: python fields: col1: col1 ratio: col2 / (col3 + 1) error_handling: output: still_bad - type: WriteToJson # Takes as input everything from the "success" path of both transforms. input: [ComputeRatio, ComputeRatioForBadRecords] config: path: /path/to/output.json - type: WriteToJson name: WriteErrorsToJson # These failed the first and the second transform. input: ComputeRatioForBadRecords.still_bad config: path: /path/to/errors.json ``` When using the `chain` syntax, the required error consumption can happen in an `extra_transforms` block. ```yaml pipeline: type: chain transforms: - type: ReadFromCsv config: path: /path/to/input*.csv - type: MapToFields name: SomeStep config: language: python fields: col1: col1 # This could raise a divide-by-zero error. ratio: col2 / col3 error_handling: output: errors - type: MapToFields name: AnotherStep config: language: python fields: col1: col1 # This could raise a divide-by-zero error. inverse_ratio: 1 / ratio error_handling: output: errors - type: WriteToJson config: path: /path/to/output.json extra_transforms: - type: WriteToJson name: WriteErrors input: [SomeStep.errors, AnotherStep.errors] config: path: /path/to/errors.json ``` ## Error Handling with Custom Providers Custom transforms, such as those defined in separate YAML files via a `YamlProvider`, can also expose error outputs from their underlying transforms. Consider a file `my_transforms.yaml` that defines a `RaiseElementToPower` transform: ```yaml # my_transforms.yaml - type: yaml transforms: RaiseElementToPower: config_schema: properties: n: {type: integer} body: type: MapToFields config: language: python append: true fields: power: "element ** {{n}}" # This transform internally defines and exposes an error output. error_handling: output: my_error ``` This transform takes a numeric element and raises it to the power of `n`. If the element is not a number, it will produce an error. The error output from the internal `MapToFields` is named `my_error`. This error output is automatically exposed by the `RaiseElementToPower` transform. When using this transform in a pipeline, you can access this error output and handle it. The main output of the transform will contain only the successfully processed elements. ```yaml pipeline: transforms: - type: Create config: elements: [2, 'bad', 3] - type: RaiseElementToPower input: Create config: n: 2 - type: WriteToJson name: WriteGood # The main output contains successfully processed elements. input: RaiseElementToPower config: path: /path/to/good - type: WriteToJson name: WriteBad # The error output is accessed by its name. input: RaiseElementToPower.my_error config: path: /path/to/bad providers: - include: my_transforms.yaml ``` In this example, the pipeline separates the good and bad records coming from the custom `RaiseElementToPower` transform. The good records are written to one location, and the error records are written to another. A pipeline will fail at construction time if an error output is declared (either in a built-in transform or a custom one) but not consumed. This helps ensure that all error paths are considered. See YAML schema [info](https://beam.apache.org/documentation/sdks/yaml-schema/) for another use of error_handling in a schema context. ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996269.0 apache_beam-2.74.0/apache_beam/yaml/docs/yaml-inline-python.md0000644000175100017510000001230015206113155023630 0ustar00runnerrunner--- type: languages title: "Apache Beam YAML Inline Python" --- # Using PyTransform from YAML Beam YAML provides the ability to easily invoke Python transforms via the `PyTransform` type, simply referencing them by fully qualified name. For example, ``` - type: PyTransform config: constructor: apache_beam.pkg.module.SomeTransform args: [1, 'foo'] kwargs: baz: 3 ``` will invoke the transform `apache_beam.pkg.mod.SomeTransform(1, 'foo', baz=3)`. This fully qualified name can be any PTransform class or other callable that returns a PTransform. Note, however, that PTransforms that do not accept or return schema'd data may not be as useable to use from YAML. Restoring the schema-ness after a non-schema returning transform can be done by using the `callable` option on `MapToFields` which takes the entire element as an input, e.g. ``` - type: PyTransform config: constructor: apache_beam.pkg.module.SomeTransform args: [1, 'foo'] kwargs: baz: 3 - type: MapToFields config: language: python fields: col1: callable: 'lambda element: element.col1' output_type: string col2: callable: 'lambda element: element.col2' output_type: integer ``` This can be used to call arbitrary transforms in the Beam SDK, e.g. ``` pipeline: transforms: - type: PyTransform name: ReadFromTsv input: {} config: constructor: apache_beam.io.ReadFromCsv kwargs: path: '/path/to/*.tsv' sep: '\t' skip_blank_lines: True true_values: ['yes'] false_values: ['no'] comment: '#' on_bad_lines: 'skip' binary: False splittable: False ``` ## Defining a transform inline using `__constructor__` If the desired transform does not exist, one can define it inline as well. This is done with the special `__constructor__` keywords, similar to how cross-language transforms are done. With the `__constuctor__` keyword, one defines a Python callable that, on invocation, *returns* the desired transform. The first argument (or `source` keyword argument, if there are no positional arguments) is interpreted as the Python code. For example ``` - type: PyTransform config: constructor: __constructor__ kwargs: source: | def create_my_transform(inc): return beam.Map(lambda x: beam.Row(a=x.col2 + inc)) inc: 10 ``` will apply `beam.Map(lambda x: beam.Row(a=x.col2 + 10))` to the incoming PCollection. As a class object can be invoked as its own constructor, this allows one to define a `beam.PTransform` inline, e.g. ``` - type: PyTransform config: constructor: __constructor__ kwargs: source: | class MyPTransform(beam.PTransform): def __init__(self, inc): self._inc = inc def expand(self, pcoll): return pcoll | beam.Map(lambda x: beam.Row(a=x.col2 + self._inc)) inc: 10 ``` which works exactly as one would expect. ## Defining a transform inline using `__callable__` The `__callable__` keyword works similarly, but instead of defining a callable that returns an applicable `PTransform` one simply defines the expansion to be performed as a callable. This is analogous to BeamPython's `ptransform.ptransform_fn` decorator. In this case one can simply write ``` - type: PyTransform config: constructor: __callable__ kwargs: source: | def my_ptransform(pcoll, inc): return pcoll | beam.Map(lambda x: beam.Row(a=x.col2 + inc)) inc: 10 ``` # External transforms One can also invoke PTransforms define elsewhere via a `python` provider, for example ``` pipeline: transforms: - ... - type: MyTransform config: kwarg: whatever providers: - ... - type: python input: ... config: packages: - 'some_pypi_package>=version' transforms: MyTransform: 'pkg.module.MyTransform' ``` These can be defined inline as well, with or without dependencies, e.g. ``` pipeline: transforms: - ... - type: ToCase input: ... config: upper: True providers: - type: python config: {} transforms: 'ToCase': | @beam.ptransform_fn def ToCase(pcoll, upper): if upper: return pcoll | beam.Map(lambda x: str(x).upper()) else: return pcoll | beam.Map(lambda x: str(x).lower()) ``` ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996269.0 apache_beam-2.74.0/apache_beam/yaml/docs/yaml-join.md0000644000175100017510000001117315206113155022001 0ustar00runnerrunner--- type: languages title: "Apache Beam YAML Join" --- # Beam YAML Join Beam YAML can join two or more inputs on specified columns. For example, the following pipeline joins the First Input pcollection and Second Input pcollection when col1 in First Input is equal to col2 in Second Input. ``` - type: Join input: input1: First Input input2: Second Input config: equalities: - input1: col1 input2: col2 ``` When joining multiple inputs on one column that is named the same across all the inputs, one can use the following shorthand syntax: ``` - type: Join input: input1: First Input input2: Second Input input3: Third Input config: equalities: col1 ``` ## Join Types When using the Join transform, one can specify the type of join to perform on the inputs. If no join type is specified, the inputs are all joined using an inner join. The supported join types are: | Join Type | YAML Keyword | | -------- | ------- | | Inner Join | inner | | Full Outer Join | left | | Right Outer Join | right | The following example joins two inputs using an inner join on the specified equalities: ``` - type: Join input: input1: First Input input2: Second Input config: type: inner equalities: - input1: col1 input2: col1 ``` The following example joins two inputs using a left outer join on the specified equalities. In this case, all rows from input1 will be kept because input1 is the left input. Order of joins follows the sequence as specified in equalities. ``` - type: Join input: input1: First Input input2: Second Input config: type: left equalities: - input1: col1 input2: col1 ``` The following example joins three inputs using an full outer join on the specified equalities: ``` - type: Join input: input1: First Input input2: Second Input input3: Third Input config: type: outer equalities: - input1: col1 input2: col1 - input2: col2 input3: col2 ``` If you want a combination of join types, you can specify the inputs to be outer joined. The following example joins input1 with input2 using a right outer join since input2 is on the right side and will join input2 with input 3 using a left outer join since input2 is on the left side. ``` - type: Join input: input1: First Input input2: Second Input input3: Third Input config: type: outer: - input2 equalities: - input1: col1 input2: col1 - input2: col2 input3: col2 ``` ## Fields By default, the join transform includes all columns from all input tables. If column names clash, it's best to rename them explicitly. Otherwise, the system will deduplicate names by adding a numeric suffix To choose which columns to output, or to customize the output column names, use the "fields" configuration. To specify which columns to output from an input, use the input reference as the configuration key and a list of desired columns as the configuration value. The following example outputs col1 from input1, col2 and col3 from input2, and all the columns from input 3. If there is a name clash, it appends a numeric suffix to avoid duplicate naming. ``` - type: Join input: input1: First Input input2: Second Input input3: Third Input config: equalities: col1 fields: input1: [col1] input2: [col2, col3] ``` To rename a column in the output, create a mapping for the input with the key as the new column name and the value as the original column name. The following example maps col1 from input3 to the column name "renamed_col1": ``` - type: Join input: input1: First Input input2: Second Input input3: Third Input config: equalities: col1 fields: input1: [col1] input2: [col2, col3] input3: renamed_col1: col1 ``` ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996269.0 apache_beam-2.74.0/apache_beam/yaml/docs/yaml-providers.md0000644000175100017510000001717115206113155023063 0ustar00runnerrunner--- type: languages title: "Apache Beam YAML Providers" --- # Providers Though we aim to offer a large suite of built-in transforms, it is inevitable that people will need to author their own. This is made possible through the notion of Providers which leverage expansion services and vend catalogues of schema transforms. ## Java Exposing transform in Java that can be used in a YAML pipeline consists of four main steps: 1. Defining the transformation itself as a [PTransform](https://beam.apache.org/documentation/programming-guide/#composite-transforms) that consumes and produces zero or more [schema'd PCollections](https://beam.apache.org/documentation/programming-guide/#creating-schemas). 2. Exposing this transform via a [SchemaTransformProvider](https://beam.apache.org/releases/javadoc/current/org/apache/beam/sdk/schemas/transforms/SchemaTransformProvider.html) which provides an identifier used to refer to this transform later as well as metadata like a human-readable description and its configuration parameters. 3. Building a Jar that contains these classes and vends them via the [Service Loader](https://github.com/apache/beam-starter-java-provider/blob/main/src/main/java/org/example/ToUpperCaseTransformProvider.java#L47) infrastructure. 4. Writing a [provider specification](https://beam.apache.org/documentation/sdks/yaml/#providers) that tells Beam YAML where to find this jar and what it contains. If the transform is already exposed as a [cross language transform](https://beam.apache.org/documentation/sdks/python-multi-language-pipelines/) or [schema transform](https://beam.apache.org/releases/javadoc/current/org/apache/beam/sdk/schemas/transforms/SchemaTransformProvider.html) then steps 1-3 have been done for you. One then uses this transform as follows: ``` pipeline: type: chain source: type: ReadFromCsv config: path: /path/to/input*.csv transforms: - type: MyCustomTransform config: arg: whatever sink: type: WriteToJson config: path: /path/to/output.json providers: - type: javaJar config: jar: /path/or/url/to/myExpansionService.jar transforms: MyCustomTransform: "urn:registered:in:expansion:service" ``` We provide a [full cloneable example of how to build a java provider](https://github.com/apache/beam-starter-java-provider) that can be used to get started. ## Python Arbitrary Python transforms can be provided as well, using the syntax ``` providers: - type: pythonPackage config: packages: - my_pypi_package>=version - /path/to/local/package.zip transforms: MyCustomTransform: "pkg.module.PTransformClassOrCallable" ``` which can then be used as ``` - type: MyCustomTransform config: num: 3 arg: whatever ``` This will cause the dependencies to be installed before the transform is imported (via its given fully qualified name) and instantiated with the config values passed as keyword arguments (e.g. in this case `pkg.module.PTransformClassOrCallable(num=3, arg="whatever")`). We offer a [python provider starter project](https://github.com/apache/beam-starter-python-provider) that serves as a complete example for how to do this. ## YAML New, re-usable transforms can be defined in YAML as well. This type of provider simply has a mapping of names to their YAML definitions. Jinja2 templatization of their string representations is used to parameterize them. The `config_schema` section of the transform definition specifies what parameters are required (with their types) and the `body` section gives the implementation in terms of other YAML transforms. ``` - type: yaml transforms: # Define the first transform of type "RaiseElementToPower" RaiseElementToPower: config_schema: properties: n: {type: integer} body: type: chain transforms: - type: MapToFields config: language: python append: true fields: power: "element ** {{n}}" # Define a second transform that produces consecutive integers. Range: config_schema: properties: end: {type: integer} # Setting this parameter lets this transform type be used as a source. requires_inputs: false body: | type: Create config: elements: {% for ix in range(end) %} - {{ix}} {% endfor %} ``` Note that in this second example the `body` of Range is defined as a [block string literal](https://yaml-multiline.info/) to prevent any attempt by the system to parse the `{%` and `%}` pragmas used for control statements before a specialization with a concrete value for `end` is instantiated and the loop is expanded. These could then be used in a pipeline as ``` transforms: - type: Range config: end: 10 - type: RaiseElementToPower input: Range config: n: 3 ... ``` One can define composite transforms as well, e.g. in a provider listing one could have ``` - type: yaml transforms: ConsecutivePowers: # This takes two parameters. config_schema: properties: end: {type: integer} n: {type: integer} # It can be used as a source transform. requires_inputs: false # The body uses the transforms defined above linked together in a chain. body: | type: chain transforms: - type: Range config: end: {{end}} - type: RaiseElementToPower config: n: {{n}} ``` which allows one to use this whole fragment as ``` type: ConsecutivePowers config: end: 10 n: 3 ``` Note that YAML-defined transforms work better in a listing file than directly in the `providers` block of a pipeline file as pipeline files are always pre-processed with Jinja2 themselves which would necessitate double escaping. ## YAML Provider listing files One can reference an external listings of providers in the yaml pipeline file via the syntax ``` providers: - include: "file:///path/to/local/providers.yaml" - include: "gs://path/to/remote/providers.yaml" - include: "https://example.com/hosted/providers.yaml" ... ``` where `providers.yaml` is simply a yaml file containing a list of providers in the same format as those inlined in this providers block. See, for example, the provider listing [here]( https://github.com/apache/beam-starter-python-provider/blob/main/examples/provider_listing.yaml). In fact, this is how many of the the built in transforms are declared, see for example the [builtin io listing file]( https://github.com/apache/beam/blob/master/sdks/python/apache_beam/yaml/standard_io.yaml). Hosting these listing files (together with their required artifacts) allows one to easily share catalogues of transforms that can be directly used by others in their YAML pipelines. ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996269.0 apache_beam-2.74.0/apache_beam/yaml/docs/yaml-schema.md0000644000175100017510000001270615206113155022305 0ustar00runnerrunner--- type: languages title: "Apache Beam YAML Schema" --- # Beam YAML Schema As pipelines grow in size and complexity, it becomes more common to encounter data that is malformed, doesn't meet preconditions, or otherwise causes issues during processing. Beam YAML helps the user detect and capture these issues by using the optional `output_schema` configuration, which is available for any transform in the YAML SDK. For example, the following code creates a few "good" records and specifies that the output schema from the `Create` transform should have records that follow the expected schema: `sdk` as a string and `year` as an integer. ```yaml pipeline: type: chain transforms: - type: Create config: elements: - {sdk: MapReduce, year: 2004} - {sdk: MillWheel, year: 2008} output_schema: type: object properties: sdk: type: string year: type: integer - type: AssertEqual config: elements: - {sdk: MapReduce, year: 2004} - {sdk: MillWheel, year: 2008} ``` However, a user will more likely want to detect and handle schema errors. If a transform has a built-in error_handling configuration, the user can specify that error_handling configuration and any errors found will be appended to the transform error_handling output. For example, the following code will create a few "good" and "bad" records with a specified schema of `user` as a string and `timestamp` as a boolean. The `alice` row will fail in the standard way because of not being an integer for the AssignTimestamps transform, while the `bob` row will fail because after the AssignTimestamp transformation, the output row will have the timestamp as an integer when it should be a boolean. ```yaml pipeline: type: composite transforms: - type: Create name: CreateVisits config: elements: - {user: alice, timestamp: "not-valid"} - {user: bob, timestamp: 3} - type: AssignTimestamps input: CreateVisits config: timestamp: timestamp error_handling: output: invalid_rows output_schema: type: object properties: user: type: string timestamp: type: boolean - type: MapToFields name: ExtractInvalidTimestamp input: AssignTimestamps.invalid_rows config: language: python fields: user: "element.user" timestamp: "element.timestamp" - type: AssertEqual input: ExtractInvalidTimestamp config: elements: - {user: "alice", timestamp: "not-valid"} - {user: bob, timestamp: 3} - type: AssertEqual input: AssignTimestamps config: elements: [] ``` ## Implicit Input Chaining for Composite Transforms When using `type: composite`, you can now omit explicit `input` specifications on sub-transforms. The sub-transforms will automatically chain together, similar to how `type: chain` works. For example, this pipeline: ```yaml pipeline: type: composite input: SomeInput transforms: - type: MapToFields config: language: python fields: x: "element.x * 2" - type: Filter config: language: python keep: "element.x > 5" ``` Is equivalent to: ```yaml pipeline: type: composite input: SomeInput transforms: - type: MapToFields input: SomeInput config: language: python fields: x: "element.x * 2" - type: Filter input: MapToFields config: language: python keep: "element.x > 5" ``` Note: Implicit chaining only works when: 1. The composite transform has an input (e.g., `input: SomeInput`) 2. Sub-transforms have no explicit `input` or `output` specifications WARNING: If a transform doesn't have the error_handling configuration available and a user chooses to use this optional output_schema feature, any failures found will result in the entire pipeline failing. If the user would still like to have some kind of output schema validation, please use the ValidateWithSchema transform instead. NOTE: When using the output_schema config, the main output key to validate on will be determined based on these criteria: 1. An output with the key 'output'. 2. An output with the key 'good'. 3. The single output if there is only one. Failures will result if the main output cannot be determined because there are multiple outputs and none are named 'output' or 'good'. For more detailed information on error handling, see this [page](https://beam.apache.org/documentation/sdks/yaml-errors/). ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996269.0 apache_beam-2.74.0/apache_beam/yaml/docs/yaml-testing.md0000644000175100017510000002023715206113155022520 0ustar00runnerrunner--- type: languages title: "Apache Beam YAML Testing" --- # Beam YAML Tests A robust testing story is an important part of any production setup. Though the various built-in (and externally provided) transform in a Beam YAML pipeline can be expected to be well tested, it can be important to have tests that ensure the pipeline as a whole behaves as expected. This is particularly true for transforms that contain non-trivial UDF logic. ## Whole pipeline tests For example, consider the example word count pipeline. ``` pipeline: transforms: - type: ReadFromText name: Read from GCS config: path: gs://dataflow-samples/shakespeare/kinglear.txt - type: MapToFields name: Split words config: language: python fields: word: callable: | import re def all_words(row): return re.findall(r'[a-z]+', row.line.lower()) value: 1 input: Read from GCS - type: Explode name: Explode word arrays config: fields: [word] input: Split words - type: Combine name: Count words config: group_by: [word] combine: value: sum input: Explode word arrays - type: MapToFields name: Format output config: language: python fields: output: "word + ': ' + str(value)" input: Count words - type: WriteToText name: Write to GCS config: path: gs://bucket/counts.txt input: Format output tests: [] ``` To write tests for this pipeline, one creates a `tests` section that enumerates a number of tests, each of which provide example input and assert the expected output is produced. An example test might be as follows ``` tests: - name: MyRegressionTest mock_outputs: - name: Read from GCS elements: - line: "Nothing can come of nothing" expected_inputs: - name: Write to GCS elements: - output: 'nothing: 2' - output: 'can: 1' - output: 'come: 1' - output: 'of: 1' ``` The `mock_outputs` section designates that the transform named `Read from GCS` should produce the single row `{line: "Nothing can come of nothing"}` for the purposes of this test, and the `expected_inputs` section indicates that the transform `Write to GCS` should expect to receive exactly the given elements. Neither the actual Read transform nor Write transform from the original pipelines are executed when running the test, but all intermediate transforms are. This test can then be executed by running ``` python -m apache_beam.yaml.main \ --yaml_pipeline_file=wordcount.yaml \ --tests ``` Alternatively, the a `tests:` block may be placed in a separate file and be validated by running ``` python -m apache_beam.yaml.main \ --yaml_pipeline_file=wordcount.yaml \ --tests \ --test_suite=test_file.yaml ``` For hermeticity, we require that all inputs (with the exception of `Create`) that are needed to compute the expected outputs are explicitly mocked; to explicitly allow a sources to be executed as part of a test their names or types can be enumerated in an `allowed_sources` attribute of the test specification. ## Pipeline fragment tests One can also test a portion of a pipeline using the `mock_inputs` and `expected_outputs` section of a test, for example ``` tests: - name: TestSplittingWithPunctuation mock_inputs: - name: Split words elements: - line: "lots-of-words" - line: "...and more" expected_outputs: - name: Explode elements: - word: lots value: 1 - word: of value: 1 - word: words value: 1 - word: and value: 1 - word: more value: 1 - name: TestCombineAndFormat mock_inputs: - name: Count words elements: - word: more value: 1 - word: and value: 1 - word: more value: 1 expected_outputs: - name: Format output elements: - output: "more: 2" - output: "and: 1" ``` As before, each test only executes the portion of the pipeline between the mock inputs and expected outputs. Note that the named transform in a `mock_inputs` specification *is* executed, while the named transform of a `mock_outputs` specification is not. Similarly, the named transform of a `expected_inputs` specification is *not* executed, while the named transform of an `expected_outputs` necessarily is. ## Automatically generating tests. In an effort to make tests as easy to write and maintain as possible, Beam YAML provides utilities to compute the expected outputs for your tests. Running ``` python -m apache_beam.yaml.main \ --yaml_pipeline_file=wordcount.yaml \ --tests \ [--test_suite=...] \ --create_test ``` will create an entirely new test by sampling all the sources and constructing a test accordingly. One can also keep tests up to date by running ``` python -m apache_beam.yaml.main \ --yaml_pipeline_file=wordcount.yaml \ --tests \ [--test_suite=...] \ --fix_tests ``` which will update any existing `expected_input` and `expected_output` blocks of your pipeline to contain the actual values computed during the test. This can be useful in authoring tests as well--one can simply specify a nonsensical or empty elements block in the expectation and the `--fix_tests` flag will populate it for you. (Of course, it is on any user of these flags to verify that the produced values are meaningful and as expected.) ## Branching pipelines For complex, branching pipelines, any number of `mock_inputs` and `mock_outputs` may be enumerated to provide the input data, and any number of `expected_inputs` and `expected_outputs` validations may be specified as well. In both the `mock_outputs` and `expected_outputs` block, multiple outputs can be disambiguated with the `TransformName.output_name` notation just as when authoring a yaml pipeline. ``` pipeline: transforms: - type: Create name: Abc config: elements: [a, b, ccc] - type: Create name: Xyz config: elements: [x, y, zzz] - type: MapToFields name: Upper input: [Abc, Xyz] config: language: python fields: element: element.upper() - type: Partition input: Upper config: language: python by: '"big" if len(element) > 1 else "small"' outputs: ["big", "small"] - type: MapToFields name: MaybeHasErrors input: Abc config: language: python fields: inverse_size: 1 / len(element) error_handling: output: errors - type: StripErrorMetadata input: MaybeHasErrors.errors tests: - name: MockMultipleInputs mock_outputs: - name: Abc elements: [element: a] - name: Xyz elements: [element: z] expected_outputs: - name: Upper elements: [element: A, element: Z] - name: TestMultipelOuptuts mock_inputs: - name: Upper elements: [element: m, element: nnn] expected_outputs: - name: Partition.big elements: [element: NNN] - name: Partition.small elements: [element: M] - name: TestErrorHandling mock_outputs: - name: Abc elements: [element: 'Aaaa', element: ''] expected_outputs: - name: MaybeHasErrors elements: [inverse_size: 0.25] - name: StripErrorMetadata elements: [element: ''] ``` ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996269.0 apache_beam-2.74.0/apache_beam/yaml/docs/yaml-udf.md0000644000175100017510000003355715206113155021632 0ustar00runnerrunner--- type: languages title: "Apache Beam YAML UDFs" --- # Beam YAML mappings Beam YAML has the ability to do simple transformations which can be used to get data into the correct shape. The simplest of these is `MapToFields` which creates records with new fields defined in terms of the input fields. ## Field renames To rename fields one can write ``` - type: MapToFields config: fields: new_col1: col1 new_col2: col2 ``` will result in an output where each record has two fields, `new_col1` and `new_col2`, whose values are those of `col1` and `col2` respectively (which are the names of two fields from the input schema). One can specify the append parameter which indicates the original fields should be retained similar to the use of `*` in an SQL select statement. For example ``` - type: MapToFields config: append: true fields: new_col1: col1 new_col2: col2 ``` will output records that have `new_col1` and `new_col2` as *additional* fields. When the append field is specified, one can drop fields as well, e.g. ``` - type: MapToFields config: append: true drop: - col3 fields: new_col1: col1 new_col2: col2 ``` which includes all original fields *except* col3 in addition to outputting the two new ones. ## Mapping functions Of course one may want to do transformations beyond just dropping and renaming fields. Beam YAML has the ability to inline simple UDFs. This requires a language specification. For example, we can provide a Python expression referencing the input fields ``` - type: MapToFields config: language: python fields: new_col: "col1.upper()" another_col: "col2 + col3" ``` In addition, one can provide a full Python callable that takes the row as an argument to do more complex mappings (see [PythonCallableSource](https://beam.apache.org/releases/pydoc/current/apache_beam.utils.python_callable.html#apache_beam.utils.python_callable.PythonCallableWithSource) for acceptable formats). Thus, one can write ``` - type: MapToFields config: language: python fields: new_col: callable: | import re def my_mapping(row): if re.match("[0-9]+", row.col1) and row.col2 > 0: return "good" else: return "bad" ``` Once one reaches a certain level of complexity, it may be preferable to package this up as a dependency and simply refer to it by fully qualified name, e.g. ``` - type: MapToFields config: language: python fields: new_col: callable: pkg.module.fn ``` It is also possible to store the function logic in a file and point to the function name, e.g. ``` - type: MapToFields config: language: python fields: new_col: path: /path/to/some/udf.py name: my_mapping ``` Currently, in addition to Python, Java, SQL, and JavaScript (experimental) expressions are supported as well ### Java When using Java mappings, the UDF type must be declared, even for simple expressions, e.g. ``` - type: MapToFields config: language: java fields: new_col: expression: col1.toUpperCase() ``` For callable UDFs, Java requires that the function be declared as a class that implements [`java.util.function.Function`](https://docs.oracle.com/javase/8/docs/api/java/util/function/Function.html), e.g. ``` - type: MapToFields config: language: java fields: new_col: callable: | import org.apache.beam.sdk.values.Row; import java.util.function.Function; public class MyFunction implements Function { public String apply(Row row) { return row.getString("col1").toUpperCase(); } } ``` ### SQL When SQL is used for a `MapToFields` UDF, it is essentially the SQL `SELECT` statement. For example, the query `SELECT UPPER(col1) AS new_col, "col2 + col3" AS another_col FROM PCOLLECTION` would look like: ``` - type: MapToFields config: language: sql fields: new_col: "UPPER(col1)" another_col: "col2 + col3" ``` keeping in mind that any fields not intended to be included in the output should be added to the `drop` field. If one wanted to select a field that collides with a [reserved SQL keyword](https://calcite.apache.org/docs/reference.html#keywords), the field(s) must be surrounded in backticks. For example, say the incoming PCollection has a field "timestamp", one would have to write: ``` - type: MapToFields config: language: sql fields: new_col: "`timestamp`" ``` **Note**: the field mapping tags and fields defined in `drop` do not need to be escaped. Only the UDF itself needs to be a valid SQL statement. ### Generic If a language is not specified the set of expressions is limited to pre-existing fields and integer, floating point, or string literals. For example ``` - type: MapToFields config: fields: new_col: col1 int_literal: 389 float_litera: 1.90216 str_literal: '"example"' # note the double quoting ``` ## FlatMap Sometimes it may be desirable to emit more (or less) than one record for each input record. This can be accomplished by mapping to an iterable type and following the mapping with an `Explode` operation, e.g. ``` - type: MapToFields config: language: python fields: new_col: "[col1.upper(), col1.lower(), col1.title()]" another_col: "col2 + col3" - type: Explode config: fields: new_col ``` will result in three output records for every input record. If more than one record is to be exploded, one must specify whether the cross product over all fields should be taken. For example ``` - type: MapToFields config: language: python fields: new_col: "[col1.upper(), col1.lower(), col1.title()]" another_col: "[col2 - 1, col2, col2 + 1]" - type: Explode config: fields: [new_col, another_col] cross_product: true ``` will emit nine records whereas ``` - type: MapToFields config: language: python fields: new_col: "[col1.upper(), col1.lower(), col1.title()]" another_col: "[col2 - 1, col2, col2 + 1]" - type: Explode config: fields: [new_col, another_col] cross_product: false ``` will only emit three. The `Explode` operation can be used on its own if the field in question is already an iterable type. ``` - type: Explode config: fields: [col1] ``` ## Filtering Sometimes it can be desirable to only keep records that satisfy a certain criteria. This can be accomplished with a `Filter` transform, e.g. ``` - type: Filter config: keep: "col2 > 0" ``` For anything more complicated than a simple comparison between existing fields and numeric literals a `language` parameter must be provided, e.g. ``` - type: Filter config: language: python keep: "col2 + col3 > 0" ``` For more complicated filtering functions, one can provide a full Python callable that takes the row as an argument to do more complex mappings (see [PythonCallableSource](https://beam.apache.org/releases/pydoc/current/apache_beam.utils.python_callable.html#apache_beam.utils.python_callable.PythonCallableWithSource) for acceptable formats). Thus, one can write ``` - type: Filter config: language: python keep: callable: | import re def my_filter(row): return re.match("[0-9]+", row.col1) and row.col2 > 0 ``` Once one reaches a certain level of complexity, it may be preferable to package this up as a dependency and simply refer to it by fully qualified name, e.g. ``` - type: Filter config: language: python keep: callable: pkg.module.fn ``` It is also possible to store the function logic in a file and point to the function name, e.g. ``` - type: Filter config: language: python keep: path: /path/to/some/udf.py name: my_filter ``` This allows the logic of the UDF itself to be more easily developed and tested using standard software engineering practices. Currently, in addition to Python, Java, SQL, and JavaScript (experimental) expressions are supported as well ### Java When using Java filtering, the UDF type must be declared, even for simple expressions, e.g. ``` - type: Filter config: language: java keep: expression: col2 > 0 ``` For callable UDFs, Java requires that the function be declared as a class that implements [`java.util.function.Function`](https://docs.oracle.com/javase/8/docs/api/java/util/function/Function.html), e.g. ``` - type: Filter config: language: java keep: callable: | import org.apache.beam.sdk.values.Row; import java.util.function.Function; import java.util.regex.Pattern; public class MyFunction implements Function { public Boolean apply(Row row) { Pattern pattern = Pattern.compile("[0-9]+"); return pattern.matcher(row.getString("col1")).matches() && row.getInt64("col2") > 0; } } ``` ### SQL Similar to [Mapping Functions](#mapping-functions), when SQL is used for a `MapToFields` UDF, it is essentially theSQL `WHERE` statement. For example, the query `SELECT * FROM PCOLLECTION WHERE col2 > 0` would look like: ``` - type: Filter config: language: sql keep: "col2 > 0" ``` If one wanted to filter on a field that collides with a [reserved SQL keyword](https://calcite.apache.org/docs/reference.html#keywords), the field(s) must be surrounded in backticks. For example, say the incoming PCollection has a field "timestamp", one would have to write: ``` - type: Filter config: language: sql keep: "`timestamp` > 0" ``` ## Partitioning It can also be useful to send different elements to different places (similar to what is done with side outputs in other SDKs). While this can be done with a set of `Filter` operations, if every element has a single destination it can be more natural to use a `Partition` transform instead which sends every element to a unique output. For example, this will send all elements where `col1` is equal to `"a"` to the output `Partition.a`. ``` - type: Partition input: input config: by: col1 outputs: ['a', 'b', 'c'] - type: SomeTransform input: Partition.a config: param: ... - type: AnotherTransform input: Partition.b config: param: ... ``` One can also specify the destination as a function, e.g. ``` - type: Partition input: input config: by: "'even' if col2 % 2 == 0 else 'odd'" language: python outputs: ['even', 'odd'] ``` One can optionally provide a catch-all output which will capture all elements that are not in the named outputs (which would otherwise be an error): ``` - type: Partition input: input config: by: col1 outputs: ['a', 'b', 'c'] unknown_output: 'other' ``` Sometimes one wants to split a PCollection into multiple PCollections that aren't necessarily disjoint. To send elements to multiple (or no) outputs, one could use an iterable column and precede the `Partition` with an `Explode`. ``` - type: Explode input: input config: fields: col1 - type: Partition input: Explode config: by: col1 outputs: ['a', 'b', 'c'] ``` ## Types Beam will try to infer the types involved in the mappings, but sometimes this is not possible. In these cases one can explicitly denote the expected output type, e.g. ``` - type: MapToFields config: language: python fields: new_col: expression: "col1.upper()" output_type: string ``` The expected type is given in json schema notation, with the addition that a top-level basic types may be given as a literal string rather than requiring a `{type: 'basic_type_name'}` nesting. ``` - type: MapToFields config: language: python fields: new_col: expression: "col1.upper()" output_type: string another_col: expression: "beam.Row(a=col1, b=[col2])" output_type: type: 'object' properties: a: type: 'string' b: type: 'array' items: type: 'number' ``` This can be especially useful to resolve errors involving the inability to handle the `beam:logical:pythonsdk_any:v1` type. ## Dependencies Often user defined functions need to rely on external dependencies. These can be provided with a `dependencies` attribute in the transform config. For example ``` - type: MapToFields config: language: python dependencies: - 'scipy>=1.15' fields: new_col: callable: | import scipy.special def func(t): return scipy.special.zeta(complex(1/2, t)) ``` The dependencies are interpreted according to the language, e.g. for Java one provides a list of maven identifiers and/or jars, and for Python one provides a list of pypi package specifiers and/or sdist tarballs. See also the full examples using [java dependencies](https://github.com/apache/beam/blob/master/sdks/python/apache_beam/yaml/examples/transforms/elementwise/map_to_fields_with_java_deps.yaml) and [python dependencies](https://github.com/apache/beam/blob/master/sdks/python/apache_beam/yaml/examples/transforms/elementwise/map_to_fields_with_deps.yaml). ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996269.0 apache_beam-2.74.0/apache_beam/yaml/docs/yaml.md0000644000175100017510000005226315206113155021051 0ustar00runnerrunner--- type: languages title: "Apache Beam YAML" --- # Beam YAML Beam YAML is a declarative syntax for describing Apache Beam pipelines by using YAML files. You can use Beam YAML to author and run a Beam pipeline without writing any code. For a full transform API reference, see https://beam.apache.org/releases/yamldoc/current/ ## Overview Beam provides a powerful model for creating sophisticated data processing pipelines. However, getting started with Beam programming can be challenging because it requires writing code in one of the supported Beam SDK languages. You need to understand the APIs, set up a project, manage dependencies, and perform other programming tasks. Beam YAML makes it easier to get started with creating Beam pipelines. Instead of writing code, you create a YAML file using any text editor. Then you submit the YAML file to be executed by a runner. The Beam YAML syntax is designed to be human-readable but also suitable as an intermediate representation for tools. For example, a pipeline authoring GUI could output YAML, or a lineage analysis tool could consume the YAML pipeline specifications. Beam YAML is still under development, but any features already included are considered stable. Feedback is welcome at dev@apache.beam.org. ## Prerequisites The Beam YAML parser is currently included as part of the [Apache Beam Python SDK](../python/). You don't need to write Python code to use Beam YAML, but you need the SDK to run pipelines locally. We recommend creating a [virtual environment](../../../get-started/quickstart/python/#create-and-activate-a-virtual-environment) so that all packages are installed in an isolated and self-contained environment. After you set up your Python environment, install the SDK as follows: ``` pip install apache_beam[yaml,gcp] ``` In addition, several of the provided transforms, such as the SQL transform, are implemented in Java and require a working Java interpeter. When you a run a pipeline with these transforms, the required artifacts are automatically downloaded from the Apache Maven repositories. ## Getting started Use a text editor to create a file named `pipeline.yaml`. Paste the following text into the file and save: ```yaml pipeline: transforms: - type: Create config: elements: [1, 2, 3] - type: LogForTesting input: Create ``` This file defines a simple pipeline with two transforms: - The `Create` transform creates a collection. The value of `config` is a dictionary of configuration settings. In this case, `elements` specifies the members of the collection. Other transform types have other configuration settings. - The `LogForTesting` transform logs each input element. This transform doesn't require a `config` setting. The `input` key specifies that `LogForTesting` receives input from the `Create` transform. ### Run the pipeline To execute the pipeline, run the following Python command: ```sh python -m apache_beam.yaml.main --yaml_pipeline_file=pipeline.yaml ``` The output should contain log statements similar to the following: ```sh INFO:root:{"element": 1} INFO:root:{"element": 2} INFO:root:{"element": 3} ``` ### Run the pipeline in Dataflow You can submit a YAML pipeline to Dataflow by using the [gcloud CLI](https://cloud.google.com/sdk/gcloud). To create a Dataflow job from the YAML file, use the [`gcloud dataflow yaml run`](https://cloud.google.com/sdk/gcloud/reference/dataflow/yaml/run) command: ``` gcloud dataflow yaml run $JOB_NAME \ --yaml-pipeline-file=pipeline.yaml \ --region=$REGION ``` When you use the `gcloud` CLI, you don't need to install the Beam SDKs locally. ### Visualize the pipeline You can use the [`apache_beam.runners.render`](https://beam.apache.org/releases/pydoc/current/apache_beam.runners.render.html) module to render the pipeline execution graph as a PNG file, as follows: 1. Install [Graphviz](https://graphviz.org/download/). 1. Run the following command: ``` python -m apache_beam.yaml.main --yaml_pipeline_file=pipeline.yaml \ --runner=apache_beam.runners.render.RenderRunner \ --render_output=out.png ``` ## Example: Reading CSV data The following pipeline reads data from a set of CSV files and writes the data in JSON format. This pipeline assumes the CSV files have a header row. The column names become JSON field names. ``` pipeline: transforms: - type: ReadFromCsv config: path: /path/to/input*.csv - type: WriteToJson config: path: /path/to/output.json input: ReadFromCsv ``` ### Add a filter The [`Filter`](../yaml-udf/#filtering) transform filters records. It keeps input records that satisfy a Boolean predicate and discards records that don't satisify the predicate. The following example keeps records where the value of `col3` is greater than 100: ``` pipeline: transforms: - type: ReadFromCsv config: path: /path/to/input*.csv - type: Filter config: language: python keep: "col3 > 100" input: ReadFromCsv - type: WriteToJson config: path: /path/to/output.json input: Filter ``` ### Add a mapping function Beam YAML supports various [mapping functions](../yaml-udf/#mapping-functions). The following example uses the `Sql` transform to group by `col1` and output the counts for each key. ``` pipeline: transforms: - type: ReadFromCsv config: path: /path/to/input*.csv - type: Filter config: language: python keep: "col3 > 100" input: ReadFromCsv - type: Sql config: query: "select col1, count(*) as cnt from PCOLLECTION group by col1" input: Filter - type: WriteToJson config: path: /path/to/output.json input: Sql ``` ## Patterns This section describes some common patterns in Beam YAML. ### Named transforms You can name the transforms in your pipeline to help with monitoring and debugging. Names are also used to disambiguate transforms if the pipeline contains more than one transform of the same type. ``` pipeline: transforms: - type: ReadFromCsv name: ReadMyData config: path: /path/to/input*.csv - type: Filter name: KeepBigRecords config: language: python keep: "col3 > 100" input: ReadMyData - type: Sql name: MySqlTransform config: query: "select col1, count(*) as cnt from PCOLLECTION group by col1" input: KeepBigRecords - type: WriteToJson name: WriteTheOutput config: path: /path/to/output.json input: MySqlTransform ``` ### Chaining transforms If a pipeline is linear (no branching or merging), you can designate the pipeline as a `chain` type. In a `chain`-type pipeline, you don't need to specify the inputs. The inputs are implicit from the order they appear in the YAML file: ``` pipeline: type: chain transforms: - type: ReadFromCsv config: path: /path/to/input*.csv - type: Filter config: language: python keep: "col3 > 100" - type: Sql name: MySqlTransform config: query: "select col1, count(*) as cnt from PCOLLECTION group by col1" - type: WriteToJson config: path: /path/to/output.json ``` If a `chain` pipeline has required error consumption or needs additional transforms not supported in a typical `chain` context, use an `extra_transforms` block. ``` pipeline: type: chain transforms: - type: ReadFromCsv config: path: /path/to/input*.csv - type: MapToFields name: SomeStep config: language: python fields: col1: col1 # This could raise a divide-by-zero error. ratio: col2 / col3 error_handling: output: errors - type: MapToFields name: AnotherStep config: language: python fields: col1: col1 # This could raise a divide-by-zero error. inverse_ratio: 1 / ratio error_handling: output: errors - type: WriteToJson config: path: /path/to/output.json extra_transforms: - type: WriteToJson name: WriteErrors input: [SomeStep.errors, AnotherStep.errors] config: path: /path/to/errors.json ``` ### Source and sink transforms As syntactic sugar, you can name the first and last transforms in your pipeline as `source` and `sink`. This convention does not change the resulting pipeline, but it signals the intent of the source and sink transforms. ``` pipeline: type: chain source: type: ReadFromCsv config: path: /path/to/input*.csv transforms: - type: Filter config: language: python keep: "col3 > 100" - type: Sql name: MySqlTransform config: query: "select col1, count(*) as cnt from PCOLLECTION group by col1" sink: type: WriteToJson config: path: /path/to/output.json ``` ### Non-linear pipelines Beam YAML supports arbitrary non-linear pipelines. The following pipeline reads two sources, joins them, and writes two outputs: ``` pipeline: transforms: - type: ReadFromCsv name: ReadLeft config: path: /path/to/left*.csv - type: ReadFromCsv name: ReadRight config: path: /path/to/right*.csv - type: Sql config: query: select A.col1, B.col2 from A join B using (col3) input: A: ReadLeft B: ReadRight - type: WriteToJson name: WriteAll input: Sql config: path: /path/to/all.json - type: Filter name: FilterToBig input: Sql config: language: python keep: "col2 > 100" - type: WriteToCsv name: WriteBig input: FilterToBig config: path: /path/to/big.csv ``` Because the pipeline is not linear, you must explicitly declare the inputs for each transform. However, you can nest a `chain` within a non-linear pipeline. The chain is a linear sub-path within the pipeline. The following example creates a chain named `ExtraProcessingForBigRows`. The chain takes input from the `Sql` transform and applies several additional filters plus a sink. Notice that within the chain, the inputs don't need to be specified. ``` pipeline: transforms: - type: ReadFromCsv name: ReadLeft config: path: /path/to/left*.csv - type: ReadFromCsv name: ReadRight config: path: /path/to/right*.csv - type: Sql config: query: select A.col1, B.col2 from A join B using (col3) input: A: ReadLeft B: ReadRight - type: WriteToJson name: WriteAll input: Sql config: path: /path/to/all.json - type: chain name: ExtraProcessingForBigRows input: Sql transforms: - type: Filter config: language: python keep: "col2 > 100" - type: Filter config: language: python keep: "len(col1) > 10" - type: Filter config: language: python keep: "col1 > 'z'" sink: type: WriteToCsv config: path: /path/to/big.csv ``` ## Windowing This API can be used to define both streaming and batch pipelines. In order to meaningfully aggregate elements in a streaming pipeline, some kind of windowing is typically required. Beam's [windowing](https://beam.apache.org/documentation/programming-guide/#windowing) and [triggering](https://beam.apache.org/documentation/programming-guide/#triggers) can be declared using the same `WindowInto` transform available in all other Beam SDKs. ``` pipeline: type: chain transforms: - type: ReadFromPubSub config: topic: myPubSubTopic format: JSON schema: type: object properties: col1: {type: string} col2: {type: integer} col3: {type: number} - type: WindowInto windowing: type: fixed size: 60s - type: SomeGroupingTransform config: arg: ... - type: WriteToPubSub config: topic: anotherPubSubTopic format: JSON options: streaming: true ``` Rather than using an explicit `WindowInto` operation, you can tag a transform with a specified windowing, which causes its inputs (and hence the transform itself) to be applied with that windowing. ``` pipeline: type: chain transforms: - type: ReadFromPubSub config: topic: myPubSubTopic format: ... schema: ... - type: SomeGroupingTransform config: arg: ... windowing: type: sliding size: 60s period: 10s - type: WriteToPubSub config: topic: anotherPubSubTopic format: JSON options: streaming: true ``` Note that the `Sql` operation itself is often a from of aggregation, and applying a windowing (or consuming an already windowed input) causes all grouping to be done per window. ``` pipeline: type: chain transforms: - type: ReadFromPubSub config: topic: myPubSubTopic format: ... schema: ... - type: Sql config: query: "select col1, count(*) as c from PCOLLECTION" windowing: type: sessions gap: 60s - type: WriteToPubSub config: topic: anotherPubSubTopic format: JSON options: streaming: true ``` The specified windowing is applied to all inputs, in this case resulting in a join per window. ``` pipeline: transforms: - type: ReadFromPubSub name: ReadLeft config: topic: leftTopic format: ... schema: ... - type: ReadFromPubSub name: ReadRight config: topic: rightTopic format: ... schema: ... - type: Sql config: query: select A.col1, B.col2 from A join B using (col3) input: A: ReadLeft B: ReadRight windowing: type: fixed size: 60s options: streaming: true ``` For a transform with no inputs, the specified windowing is instead applied to its output(s). As per the Beam model, the windowing is then inherited by all consuming operations. This is especially useful for root operations like Read. ``` pipeline: type: chain transforms: - type: ReadFromPubSub config: topic: myPubSubTopic format: ... schema: ... windowing: type: fixed size: 60s - type: Sql config: query: "select col1, count(*) as c from PCOLLECTION" - type: WriteToPubSub config: topic: anotherPubSubTopic format: JSON options: streaming: true ``` One can also specify windowing at the top level of a pipeline (or composite), which is a shorthand for applying this same windowing to all root operations that don't otherwise specify their own windowing. This approach is effective way to apply a window everywhere in the pipeline. ``` pipeline: type: chain transforms: - type: ReadFromPubSub config: topic: myPubSubTopic format: ... schema: ... - type: Sql config: query: "select col1, count(*) as c from PCOLLECTION" - type: WriteToPubSub config: topic: anotherPubSubTopic format: JSON windowing: type: fixed size: 60 options: streaming: true ``` Note that all these windowing specifications are compatible with the `source` and `sink` syntax as well: ``` pipeline: type: chain source: type: ReadFromPubSub config: topic: myPubSubTopic format: ... schema: ... windowing: type: fixed size: 10s transforms: - type: Sql config: query: "select col1, count(*) as c from PCOLLECTION" sink: type: WriteToCsv config: path: /path/to/output.json windowing: type: fixed size: 5m options: streaming: true ``` ## Pipeline options [Pipeline options](https://beam.apache.org/documentation/programming-guide/#configuring-pipeline-options) are used to configure different aspects of your pipeline, such as the pipeline runner that will execute your pipeline and any runner-specific configuration required by the chosen runner. To set pipeline options, append an options block at the end of your yaml file. For example: ``` pipeline: type: chain transforms: - type: ReadFromPubSub config: topic: myPubSubTopic format: ... schema: ... windowing: type: fixed size: 60s - type: Sql config: query: "select col1, count(*) as c from PCOLLECTION" - type: WriteToPubSub config: topic: anotherPubSubTopic format: JSON options: streaming: true ``` ## Jinja Templatization It is a common to want to run a single Beam pipeline in different contexts and/or with different configurations. When running a YAML pipeline using `apache_beam.yaml.main` or via gcloud, the yaml file can be parameterized with externally provided variables using the [jinja variable syntax](https://jinja.palletsprojects.com/en/stable/templates/#variables). The values are then passed via a `--jinja_variables` command line flag. For example, one could start a pipeline with: ``` pipeline: transforms: - type: ReadFromCsv config: path: {{input_pattern}} ``` and then run it with ```sh python -m apache_beam.yaml.main \ --yaml_pipeline_file=pipeline.yaml \ --jinja_variables='{"input_pattern": "gs://path/to/this/runs/files*.csv"}' ``` Arbitrary [jinja control structures](https://jinja.palletsprojects.com/en/stable/templates/#list-of-control-structures), such as looping and conditionals, can be used as well if desired as long as the output results in a valid Beam YAML pipeline. We also expose the [`datetime`](https://docs.python.org/3/library/datetime.html) module as a variable by default, which can be particularly useful in reading or writing dated sources and sinks, e.g. ``` - type: WriteToJson config: path: "gs://path/to/{{ datetime.datetime.now().strftime('%Y/%m/%d') }}/dated-output.json" ``` would write to files like `gs://path/to/2016/08/04/dated-output*.json`. A user can also use the `% include` directive to pull in other common templates: /pipeline.yaml ```yaml pipeline: transforms: - name: Read from GCS type: ReadFromText config: # NOTE: For include, the indentation has to line up correctly for it to be # parsed correctly. So in this example the included readFromText.yaml has # already indented yaml lines to line up correctly when including into this # pipeline here. {% include '/submodules/readFromText.yaml' %} - name: Write to GCS type: WriteToText input: Read from GCS config: path: "gs://MY-BUCKET/wordCounts/" ``` /submodules/readFromText.yaml ```yaml path: {{readFromText.path}} ``` This pipeline can be run like this: ```sh python -m apache_beam.yaml.main \ --yaml_pipeline_file=pipeline.yaml \ --jinja_variables='{"readFromText": {"path": "gs://dataflow-samples/shakespeare/kinglear.txt"}}' ``` The `% import` jinja directive can also be used to pull in macros: /pipeline.yaml ```yaml {% import '/macros.yaml' as macros %} pipeline: type: chain transforms: # Read in text file {{ macros.readFromText(readFromText) | indent(4, true) }} # Write to text file on GCS, locally, etc - name: Write to GCS type: WriteToText input: Read from GCS config: path: "gs://MY-BUCKET/wordCounts/" ``` /macros.yaml ```yaml {%- macro readFromText(params) -%} - name: Read from GCS type: ReadFromText config: path: "{{ params.path }}" {%- endmacro -%} ``` This pipeline can be run with the same command as in the `% include` example above. There are many more ways to import and even use template inheritance using Jinja as seen [here](https://jinja.palletsprojects.com/en/stable/templates/#import) and [here](https://jinja.palletsprojects.com/en/stable/templates/#inheritance). Note that for large chunks of functionality, we recommend packaging them up via more reusable [yaml providers](../yaml-providers) rather than using textual `%includes`. Full jinja pipeline examples can be found [here](https://github.com/apache/beam/tree/master/sdks/python/apache_beam/yaml/examples/transforms/jinja). ## Other Resources * [Example pipeline](https://github.com/apache/beam/tree/master/sdks/python/apache_beam/yaml/examples) * [More examples](https://gist.github.com/robertwb/2cb26973f1b1203e8f5f8f88c5764da0) ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.8134534 apache_beam-2.74.0/apache_beam/yaml/examples/0000755000175100017510000000000015206113160020437 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/examples/__init__.py0000644000175100017510000000142015206113033022544 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.8154538 apache_beam-2.74.0/apache_beam/yaml/examples/testing/0000755000175100017510000000000015206113160022114 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/examples/testing/__init__.py0000644000175100017510000000142015206113033024221 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/examples/testing/examples_test.py0000644000175100017510000014125515206113033025352 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # pytype: skip-file import glob import json import logging import os import random import re import sys import unittest from typing import Any from typing import Callable from typing import Optional from typing import Union from unittest import mock import pytest import yaml from jinja2 import DictLoader from jinja2 import Environment from jinja2 import StrictUndefined import apache_beam as beam from apache_beam import PCollection from apache_beam.examples.snippets.util import assert_matches_stdout from apache_beam.ml.inference.base import PredictionResult from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.typehints.row_type import RowTypeConstraint from apache_beam.utils import subprocess_server from apache_beam.yaml import yaml_provider from apache_beam.yaml import yaml_transform from apache_beam.yaml.readme_test import TestEnvironment from apache_beam.yaml.readme_test import replace_recursive from . import input_data @beam.ptransform.ptransform_fn def test_enrichment( pcoll, enrichment_handler: str, handler_config: dict[str, Any], timeout: Optional[float] = 30): """ Mocks the Enrichment transform for testing purposes. This PTransform simulates the behavior of the Enrichment transform by looking up data from predefined in-memory tables based on the provided `enrichment_handler` and `handler_config`. Note: The Github action that invokes these tests does not have gcp dependencies installed which is a prerequisite to apache_beam.transforms.enrichment.Enrichment as a top-level import. Args: pcoll: The input PCollection. enrichment_handler: A string indicating the type of enrichment handler to simulate (e.g., 'BigTable', 'BigQuery'). handler_config: A dictionary containing configuration details for the simulated handler (e.g., table names, row keys, fields). timeout: An optional timeout value (ignored in this mock). Returns: A PCollection containing the enriched data. """ if enrichment_handler == 'BigTable': row_key = handler_config['row_key'] bt_data = INPUT_TABLES[( 'BigTable', handler_config['instance_id'], handler_config['table_id'])] products = {str(data[row_key]): data for data in bt_data} def _fn(row): left = row._asdict() right = products[str(left[row_key])] left['product'] = left.get('product', None) or right return beam.Row(**left) elif enrichment_handler == 'BigQuery': row_key = handler_config['fields'] dataset, table = handler_config['table_name'].split('.')[-2:] bq_data = INPUT_TABLES[('BigQuery', str(dataset), str(table))] bq_data = { tuple(str(data[key]) for key in row_key): data for data in bq_data } def _fn(row): left = row._asdict() right = bq_data[tuple(str(left[k]) for k in row_key)] row = { key: left.get(key, None) or right[key] for key in {*left.keys(), *right.keys()} } return beam.Row(**row) else: raise ValueError(f'{enrichment_handler} is not a valid enrichment_handler.') return pcoll | beam.Map(_fn) @beam.ptransform.ptransform_fn def test_kafka_read( pcoll, topic: Optional[str] = None, format: Optional[str] = None, schema: Optional[Any] = None, bootstrap_servers: Optional[str] = None, auto_offset_reset_config: Optional[str] = None, consumer_config: Optional[Any] = None): """ Mocks the ReadFromKafka transform for testing purposes. This PTransform simulates the behavior of the ReadFromKafka transform by reading from predefined in-memory data based on the Kafka topic argument. Args: pcoll: The input PCollection. topic: The name of Kafka topic to read from. format: The format of the Kafka messages (e.g., 'RAW'). schema: The schema of the Kafka messages. bootstrap_servers: A list of Kafka bootstrap servers to connect to. auto_offset_reset_config: A configuration for the auto offset reset. consumer_config: A map for additional consumer configuration parameters. Returns: A PCollection containing the sample data. """ if topic == 'test-topic': kafka_byte_messages = KAFKA_TOPICS['test-topic'] return ( pcoll | beam.Create([msg.decode('utf-8') for msg in kafka_byte_messages]) | beam.Map(lambda element: beam.Row(payload=element.encode('utf-8')))) return None @beam.ptransform.ptransform_fn def test_pubsub_read( pcoll, topic: Optional[str] = None, subscription: Optional[str] = None, format: Optional[str] = None, schema: Optional[Any] = None, attributes: Optional[list[str]] = None, attributes_map: Optional[str] = None, id_attribute: Optional[str] = None, timestamp_attribute: Optional[str] = None): """ Mocks the ReadFromPubSub transform for testing purposes. This PTransform simulates the behavior of the ReadFromPubSub transform by reading from predefined in-memory data based on the Pub/Sub topic argument. Args: pcoll: The input PCollection. topic: The name of Pub/Sub topic to read from. subscription: The name of Pub/Sub subscription to read from. format: The format of the Pub/Sub messages (e.g., 'JSON'). schema: The schema of the Pub/Sub messages. attributes: A list of attributes to include in the output. attributes_map: A string representing a mapping of attributes. id_attribute: The attribute to use as the ID for the message. timestamp_attribute: The attribute to use as the timestamp for the message. Returns: A PCollection containing the sample data. """ if topic == 'test-topic': pubsub_messages = PUBSUB_TOPICS['test-topic'] return ( pcoll | beam.Create([json.loads(msg.data) for msg in pubsub_messages]) | beam.Map(lambda element: beam.Row(**element))) elif topic == 'taxi-ride-topic': pubsub_messages = PUBSUB_TOPICS['taxi-ride-topic'] schema = input_data.TaxiRideEventSchema return ( pcoll | beam.Create([json.loads(msg.data) for msg in pubsub_messages]) | beam.Map(lambda element: beam.Row(**element)).with_output_types(schema)) return None @beam.ptransform.ptransform_fn def test_run_inference_taxi_fare(pcoll, inference_tag, model_handler): """ This PTransform simulates the behavior of the RunInference transform. Args: pcoll: The input PCollection. inference_tag: The tag to use for the returned inference. model_handler: A configuration for the respective ML model handler Returns: A PCollection containing the enriched data. """ def _fn(row): input = row._asdict() row = {inference_tag: PredictionResult(input, 10.0), **input} return beam.Row(**row) schema = _format_predicition_result_ouput(pcoll, inference_tag) return pcoll | beam.Map(_fn).with_output_types(schema) @beam.ptransform.ptransform_fn def test_run_inference_youtube_comments(pcoll, inference_tag, model_handler): """ This PTransform simulates the behavior of the RunInference transform. Args: pcoll: The input PCollection. inference_tag: The tag to use for the returned inference. model_handler: A configuration for the respective ML model handler Returns: A PCollection containing the enriched data. """ def _fn(row): input = row._asdict() row = { inference_tag: PredictionResult( input['comment_text'], [{ 'label': 'POSITIVE' if 'happy' in input['comment_text'] else 'NEGATIVE', 'score': 0.95 }]), **input } return beam.Row(**row) schema = _format_predicition_result_ouput(pcoll, inference_tag) return pcoll | beam.Map(_fn).with_output_types(schema) def _format_predicition_result_ouput(pcoll, inference_tag): user_type = RowTypeConstraint.from_user_type(pcoll.element_type.user_type) user_schema_fields = [(name, type(typ) if not isinstance(typ, type) else typ) for (name, typ) in user_type._fields] if user_type else [] inference_output_type = RowTypeConstraint.from_fields([ ('example', Any), ('inference', Any), ('model_id', Optional[str]) ]) return RowTypeConstraint.from_fields( user_schema_fields + [(str(inference_tag), inference_output_type)]) TEST_PROVIDERS = { 'TestEnrichment': test_enrichment, 'TestReadFromKafka': test_kafka_read, 'TestReadFromPubSub': test_pubsub_read, 'TestRunInferenceYouTubeComments': test_run_inference_youtube_comments, 'TestRunInferenceTaxiFare': test_run_inference_taxi_fare, } """ Transforms not requiring inputs. """ INPUT_TRANSFORM_TEST_PROVIDERS = ['TestReadFromKafka', 'TestReadFromPubSub'] def check_output(expected: list[str]): """ Helper function to check the output of a pipeline against expected values. This function takes a list of expected output strings and returns a callable that can be used within a Beam pipeline to assert that the actual output matches the expected output. Args: expected: A list of strings representing the expected output elements. Returns: A callable that takes a list of PCollections and asserts their combined elements match the expected output. """ def _check_inner(actual: list[PCollection[str]]): formatted_actual = actual | beam.Flatten() | beam.Map( lambda row: str(beam.Row(**row._asdict()))) assert_matches_stdout(formatted_actual, expected) return _check_inner def create_test_method( pipeline_spec_file: str, custom_preprocessors: list[Callable[..., Union[dict, list]]]): """ Generates a test method for a given YAML pipeline specification file. This function reads the YAML file, extracts the expected output (if present), and creates a test function that uses `TestPipeline` to run the pipeline defined in the YAML file. It also applies any custom preprocessors registered for this test. Args: pipeline_spec_file: The path to the YAML file containing the pipeline specification. custom_preprocessors: A list of preprocessor functions to apply before running the test. Returns: A test method (Callable) that can be added to a unittest.TestCase class. """ @mock.patch('apache_beam.Pipeline', TestPipeline) def test_yaml_example(self): with open(pipeline_spec_file, encoding="utf-8") as f: lines = f.readlines() expected_key = '# Expected:\n' if expected_key in lines: expected = lines[lines.index('# Expected:\n') + 1:] else: raise ValueError( f"Missing '# Expected:' tag in example file '{pipeline_spec_file}'") for i, line in enumerate(expected): expected[i] = line.replace('# ', '').replace('\n', '') expected = [line for line in expected if line] raw_spec_string = ''.join(lines) # Filter for any jinja preprocessor - this has to be done before other # preprocessors. jinja_preprocessor = [ preprocessor for preprocessor in custom_preprocessors if 'jinja_preprocessor' in preprocessor.__name__ ] if jinja_preprocessor: jinja_preprocessor = jinja_preprocessor[0] raw_spec_string = jinja_preprocessor( raw_spec_string, self._testMethodName) custom_preprocessors.remove(jinja_preprocessor) pipeline_spec = yaml.load( raw_spec_string, Loader=yaml_transform.SafeLineLoader) with TestEnvironment() as env: for fn in custom_preprocessors: pipeline_spec = fn(pipeline_spec, expected, env) with beam.Pipeline(options=PipelineOptions( pickle_library='cloudpickle', **yaml_transform.SafeLineLoader.strip_metadata(pipeline_spec.get( 'options', {})))) as p: actual = [ yaml_transform.expand_pipeline( p, pipeline_spec, [ yaml_provider.InlineProvider( TEST_PROVIDERS, INPUT_TRANSFORM_TEST_PROVIDERS) ]) ] if not actual[0]: actual = list(p.transforms_stack[0].parts[-1].outputs.values()) for transform in p.transforms_stack[0].parts[:-1]: if transform.transform.label == 'log_for_testing': actual += list(transform.outputs.values()) check_output(expected)(actual) def _python_deps_involved(spec_filename): return any( substr in spec_filename for substr in ['deps', 'streaming_sentiment_analysis', 'ml_preprocessing']) def _java_deps_involved(spec_filename): return any( substr in spec_filename for substr in ['java_deps', 'streaming_taxifare_prediction']) if _python_deps_involved(pipeline_spec_file): test_yaml_example = unittest.skipIf( sys.platform == 'win32', "Github virtualenv permissions issues.")( test_yaml_example) # This test fails, with an import error, for some (but not all) cloud # tox environments when run as a github action (not reproducible locally). # Adding debugging makes the failure go away. All indications are that # this is some testing environmental issue. test_yaml_example = unittest.skipIf( '-cloud' in os.environ.get('TOX_ENV_NAME', ''), 'Github actions environment issue.')( test_yaml_example) test_yaml_example = pytest.mark.no_xdist(test_yaml_example) if _java_deps_involved(pipeline_spec_file): test_yaml_example = pytest.mark.xlang_sql_expansion_service( test_yaml_example) test_yaml_example = unittest.skipIf( not os.path.exists( subprocess_server.JavaJarServer.path_to_dev_beam_jar( 'sdks:java:extensions:sql:expansion-service:shadowJar')), "Requires expansion service jars.")( test_yaml_example) return test_yaml_example class YamlExamplesTestSuite: """ YamlExamplesTestSuites class is used to scan specified directories for .yaml files and dynamically generate a Python test method. Additionally, it creates a method to complete some preprocessing for mocking IO. """ _test_preprocessor: dict[str, list[Callable[..., Union[dict, list]]]] = {} def __init__(self, name: str, path: str): """ Initializes the YamlExamplesTestSuite. Args: name: The name of the test suite. This will be used as the class name for the dynamically generated test suite. path: A string representing the path or glob pattern to search for YAML example files. """ self._test_suite = self.create_test_suite(name, path) def run(self): """ Runs the dynamically generated test suite. This method simply returns the test suite class created during initialization. The test runner (e.g., unittest.main()) can then be used to discover and run the tests within this suite. Returns: The dynamically created unittest.TestCase subclass. """ return self._test_suite @classmethod def parse_test_methods(cls, path: str): """Scans a given path for YAML files and generates test methods. This method uses glob to find files matching the provided path. For each YAML file found, it constructs a unique test name and then calls `create_test_method` to generate the actual test function. It also retrieves any registered preprocessors for that specific test. Args: path: A string representing the path or glob pattern to search for YAML example files. Yields: A tuple containing the generated test name (str) and the corresponding test method (Callable). """ files = glob.glob(path) if not files and os.path.exists(path) and os.path.isfile(path): files = [path] for file in files: test_name = f'test_{file.split(os.sep)[-1].replace(".", "_")}' custom_preprocessors = cls._test_preprocessor.get(test_name, []) yield test_name, create_test_method(file, custom_preprocessors) @classmethod def create_test_suite(cls, name: str, path: str): """Dynamically creates a unittest.TestCase subclass with generated tests. This method takes a suite name and a path (or glob pattern). It uses `parse_test_methods` to find YAML files at the given path and generate individual test methods for each. These generated test methods are then added as attributes to a new class, which is a subclass of `unittest.TestCase`. Args: name: The desired name for the dynamically created test suite class. path: A string representing the path or glob pattern to search for YAML example files, which will be used to generate test methods. Returns: A new class, subclass of `unittest.TestCase`, containing dynamically generated test methods based on the YAML files found at the given path. """ return type( name, (unittest.TestCase, ), dict(cls.parse_test_methods(path))) @classmethod def register_test_preprocessor(cls, test_names: Union[str, list]): """Decorator to register a preprocessor function for specific tests. This decorator is used to associate a preprocessor function with one or more test names. The preprocessor function will be called before the corresponding test is executed, allowing for modification of the test specification or environment setup. Args: test_names: A string or a list of strings representing the names of the tests for which the preprocessor should be registered. The test names should match the names generated by `parse_test_methods`. Returns: A decorator function that takes the preprocessor function as an argument and registers it. """ if isinstance(test_names, str): test_names = [test_names] def apply(preprocessor): for test_name in test_names: if test_name not in cls._test_preprocessor: cls._test_preprocessor[test_name] = [] cls._test_preprocessor[test_name].append(preprocessor) return preprocessor return apply @YamlExamplesTestSuite.register_test_preprocessor( ['test_wordcount_minimal_yaml']) def _wordcount_minimal_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for the wordcount_minimal.yaml test. This preprocessor generates a random input file based on the expected output of the wordcount example. This allows the test to verify the pipeline's correctness without relying on a fixed input file. Based on this expected output: # Row(word='king', count=311) Args: test_spec: The dictionary representation of the YAML pipeline specification. expected: A list of strings representing the expected output of the pipeline. env: The TestEnvironment object providing utilities for creating temporary files. Returns: The modified test_spec dictionary with the input file path replaced. """ all_words = [] for element in expected: word = element.split('=')[1].split(',')[0].replace("'", '') count = int(element.split('=')[2].replace(')', '')) all_words += [word] * count return _wordcount_random_shuffler(test_spec, all_words, env) @YamlExamplesTestSuite.register_test_preprocessor([ 'test_wordCountInclude_yaml', 'test_wordCountImport_yaml', 'test_wordCountInheritance_yaml' ]) def _wordcount_jinja_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for the wordcount Jinja tests. This preprocessor generates a random input file based on the expected output of the wordcount example. This allows the test to verify the pipeline's correctness without relying on a fixed input file. Based on this expected output: # Row(output='king - 311') Args: test_spec: The dictionary representation of the YAML pipeline specification. expected: A list of strings representing the expected output of the pipeline. env: The TestEnvironment object providing utilities for creating temporary files. Returns: The modified test_spec dictionary with the input file path replaced. """ all_words = [] for element in expected: match = re.search(r"output='(.*) - (\d+)'", element) if match: word, count_str = match.groups() all_words += [word] * int(count_str) return _wordcount_random_shuffler(test_spec, all_words, env) def _wordcount_random_shuffler( test_spec: dict, all_words: list[str], env: TestEnvironment): """ Helper function to create a randomized input file for wordcount-style tests. This function takes a list of words, shuffles them, and arranges them into randomly sized lines. It then creates a temporary input file with this content and updates the provided test specification to use this file as the input for a 'ReadFromText' transform. Args: test_spec: The dictionary representation of the YAML pipeline specification. all_words: A list of strings, where each string is a word to be included in the generated input file. env: The TestEnvironment object providing utilities for creating temporary files. Returns: The modified test_spec dictionary with the input file path for 'ReadFromText' replaced with the path to the newly generated file. """ random.shuffle(all_words) lines = [] while all_words: line_length = random.randint(1, min(10, len(all_words))) line = " ".join( all_words.pop(random.randrange(len(all_words))) for _ in range(line_length)) lines.append(line) return replace_recursive( test_spec, 'ReadFromText', 'path', env.input_file('kinglear.txt', '\n'.join(lines))) @YamlExamplesTestSuite.register_test_preprocessor( ['test_kafka_yaml', 'test_kafka_to_iceberg_yaml']) def _kafka_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): test_spec = replace_recursive( test_spec, 'ReadFromText', 'path', env.input_file('kinglear.txt', input_data.text_data())) if pipeline := test_spec.get('pipeline', None): for transform in pipeline.get('transforms', []): if transform.get('type', '') == 'ReadFromKafka': transform['type'] = 'TestReadFromKafka' transform['config']['topic'] = 'test-topic' return test_spec @YamlExamplesTestSuite.register_test_preprocessor([ 'test_simple_filter_yaml', 'test_simple_filter_and_combine_yaml', 'test_iceberg_read_yaml', 'test_iceberg_write_yaml', 'test_kafka_yaml', 'test_spanner_read_yaml', 'test_spanner_write_yaml', 'test_enrich_spanner_with_bigquery_yaml', 'test_pubsub_topic_to_bigquery_yaml', 'test_pubsub_subscription_to_bigquery_yaml', 'test_jdbc_to_bigquery_yaml', 'test_spanner_to_avro_yaml', 'test_gcs_text_to_bigquery_yaml', 'test_sqlserver_to_bigquery_yaml', 'test_postgres_to_bigquery_yaml', 'test_kafka_to_iceberg_yaml', 'test_pubsub_to_iceberg_yaml', 'test_oracle_to_bigquery_yaml', 'test_mysql_to_bigquery_yaml', 'test_spanner_to_bigquery_yaml', 'test_streaming_sentiment_analysis_yaml', 'test_iceberg_migration_yaml', 'test_ml_preprocessing_yaml', 'test_anomaly_scoring_yaml', 'test_wordCountInclude_yaml', 'test_wordCountImport_yaml', 'test_wordCountInheritance_yaml', 'test_iceberg_to_alloydb_yaml', 'test_bigquery_write_yaml' ]) def _io_write_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for tests that involve writing to IO. This preprocessor replaces any WriteTo transform with a LogForTesting transform. This allows the test to verify the data being written without actually writing to an external system. Args: test_spec: The dictionary representation of the YAML pipeline specification. expected: A list of strings representing the expected output of the pipeline. env: The TestEnvironment object providing utilities for creating temporary files. Returns: The modified test_spec dictionary with WriteTo transforms replaced. """ if pipeline := test_spec.get('pipeline', None): for transform in pipeline.get('transforms', []): if transform.get('type', '').startswith('WriteTo'): transform['type'] = 'LogForTesting' transform['config'] = { k: v for (k, v) in transform.get('config', {}).items() if (k.startswith('__') or k == 'error_handling') } return test_spec @YamlExamplesTestSuite.register_test_preprocessor([ 'test_simple_filter_yaml', 'test_simple_filter_and_combine_yaml', 'test_gcs_text_to_bigquery_yaml' ]) def _file_io_read_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ This preprocessor replaces any file IO ReadFrom transform with a Create transform that reads from a predefined in-memory dictionary. This allows the test to verify the pipeline's correctness without relying on external files. Args: test_spec: The dictionary representation of the YAML pipeline specification. expected: A list of strings representing the expected output of the pipeline. env: The TestEnvironment object providing utilities for creating temporary files. Returns: The modified test_spec dictionary with ReadFrom transforms replaced. """ if pipeline := test_spec.get('pipeline', None): for transform in pipeline.get('transforms', []): if transform.get('type', '').startswith('ReadFrom'): file_name = transform['config']['path'].split('/')[-1] return replace_recursive( test_spec, transform['type'], 'path', env.input_file(file_name, INPUT_FILES[file_name])) return test_spec @YamlExamplesTestSuite.register_test_preprocessor( ['test_iceberg_read_yaml', 'test_iceberg_to_alloydb_yaml']) def _iceberg_io_read_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for tests that involve reading from Iceberg. This preprocessor replaces any ReadFromIceberg transform with a Create transform that reads from a predefined in-memory dictionary. This allows the test to verify the pipeline's correctness without relying on Iceberg tables stored externally. Args: test_spec: The dictionary representation of the YAML pipeline specification. expected: A list of strings representing the expected output of the pipeline. env: The TestEnvironment object providing utilities for creating temporary files. Returns: The modified test_spec dictionary with ReadFromIceberg transforms replaced. """ if pipeline := test_spec.get('pipeline', None): for transform in pipeline.get('transforms', []): if transform.get('type', '') == 'ReadFromIceberg': config = transform['config'] (db_name, table_name, field_value_dynamic_destinations) = config['table'].split('.') transform['type'] = 'Create' transform['config'] = { k: v for k, v in config.items() if k.startswith('__') } transform['config']['elements'] = INPUT_TABLES[( str(db_name), str(table_name), str(field_value_dynamic_destinations))] return test_spec @YamlExamplesTestSuite.register_test_preprocessor([ 'test_spanner_read_yaml', 'test_enrich_spanner_with_bigquery_yaml', 'test_spanner_to_avro_yaml', 'test_spanner_to_bigquery_yaml' ]) def _spanner_io_read_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for tests that involve reading from Spanner. This preprocessor replaces any ReadFromSpanner transform with a Create transform that reads from a predefined in-memory dictionary. This allows the test to verify the pipeline's correctness without relying on external Spanner instances. Args: test_spec: The dictionary representation of the YAML pipeline specification. expected: A list of strings representing the expected output of the pipeline. env: The TestEnvironment object providing utilities for creating temporary files. Returns: The modified test_spec dictionary with ReadFromSpanner transforms replaced. """ if pipeline := test_spec.get('pipeline', None): for transform in pipeline.get('transforms', []): if transform.get('type', '').startswith('ReadFromSpanner'): config = transform['config'] instance, database = config['instance_id'], config['database_id'] if table := config.get('table', None) is None: table = config.get('query', '').split('FROM')[-1].strip() transform['type'] = 'Create' transform['config'] = { k: v for k, v in config.items() if k.startswith('__') } elements = INPUT_TABLES[(str(instance), str(database), str(table))] if config.get('query', None): config['query'].replace('select ', 'SELECT ').replace(' from ', ' FROM ') columns = set( ''.join(config['query'].split('SELECT ')[1:]).split( ' FROM', maxsplit=1)[0].split(', ')) if columns != {'*'}: elements = [{ column: element[column] for column in element if column in columns } for element in elements] transform['config']['elements'] = elements return test_spec @YamlExamplesTestSuite.register_test_preprocessor( ['test_bigtable_enrichment_yaml', 'test_enrich_spanner_with_bigquery_yaml']) def _enrichment_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for tests that involve the Enrichment transform. This preprocessor replaces the actual Enrichment transform with a mock `TestEnrichment` transform. This allows the test to verify the pipeline's correctness without requiring external services like BigTable or BigQuery. Args: test_spec: The dictionary representation of the YAML pipeline specification. expected: A list of strings representing the expected output of the pipeline. env: The TestEnvironment object providing utilities for creating temporary files. Returns: The modified test_spec dictionary with Enrichment transforms replaced. """ if pipeline := test_spec.get('pipeline', None): for transform in pipeline.get('transforms', []): if transform.get('type', '').startswith('Enrichment'): transform['type'] = 'TestEnrichment' return test_spec @YamlExamplesTestSuite.register_test_preprocessor([ 'test_pubsub_topic_to_bigquery_yaml', 'test_pubsub_subscription_to_bigquery_yaml', 'test_pubsub_to_iceberg_yaml' ]) def _pubsub_io_read_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for tests that involve reading from Pub/Sub. This preprocessor replaces any ReadFromPubSub transform with a Create transform that reads from a predefined in-memory list of messages. This allows the test to verify the pipeline's correctness without relying on an active Pub/Sub subscription or topic. """ if pipeline := test_spec.get('pipeline', None): for transform in pipeline.get('transforms', []): if transform.get('type', '') == 'ReadFromPubSub': transform['type'] = 'TestReadFromPubSub' transform['config']['topic'] = 'test-topic' return test_spec @YamlExamplesTestSuite.register_test_preprocessor([ 'test_jdbc_to_bigquery_yaml', ]) def _jdbc_io_read_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for tests that involve reading from generic Jdbc. url syntax: 'jdbc:://:/' """ return _db_io_read_test_processor( test_spec, lambda url: url.split('/')[-1], 'Jdbc') @YamlExamplesTestSuite.register_test_preprocessor([ 'test_sqlserver_to_bigquery_yaml', ]) def __sqlserver_io_read_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for tests that involve reading from SqlServer. url syntax: 'jdbc:sqlserver://:;databaseName=; user=;password=;encrypt=false;trustServerCertificate=true' """ return _db_io_read_test_processor( test_spec, lambda url: url.split(';')[1].split('=')[-1], 'SqlServer') @YamlExamplesTestSuite.register_test_preprocessor([ 'test_postgres_to_bigquery_yaml', ]) def __postgres_io_read_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for tests that involve reading from Postgres. url syntax: 'jdbc:postgresql://:/shipment?user=& password=' """ return _db_io_read_test_processor( test_spec, lambda url: url.split('/')[3].split('?')[0], 'Postgres') @YamlExamplesTestSuite.register_test_preprocessor([ 'test_oracle_to_bigquery_yaml', ]) def __oracle_io_read_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for tests that involve reading from Oracle. url syntax: 'jdbc:oracle:thin:system/oracle@:{port}/' """ return _db_io_read_test_processor( test_spec, lambda url: url.split('/')[2], 'Oracle') @YamlExamplesTestSuite.register_test_preprocessor([ 'test_mysql_to_bigquery_yaml', ]) def __mysql_io_read_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for tests that involve reading from MySql. url syntax: 'jdbc:mysql://:/?user=& password=' """ return _db_io_read_test_processor( test_spec, lambda url: url.split('/')[3].split('?')[0], 'MySql') def _db_io_read_test_processor( test_spec: dict, database_url_fn: Callable, database_type: str): """ This preprocessor replaces any ReadFrom transform with a Create transform that reads from a predefined in-memory list of records. This allows the test to verify the pipeline's correctness without relying on an active database. """ if pipeline := test_spec.get('pipeline', None): for transform in pipeline.get('transforms', []): transform_name = f"ReadFrom{database_type}" if transform.get('type', '').startswith(transform_name): config = transform['config'] url = config['url'] database = database_url_fn(url) if (table := config.get('table', None)) is None: table = config.get('query', '').split('FROM')[-1].strip() transform['type'] = 'Create' transform['config'] = { k: v for k, v in config.items() if k.startswith('__') } elements = INPUT_TABLES[(database_type, database, table)] if config.get('query', None): config['query'].replace('select ', 'SELECT ').replace(' from ', ' FROM ') columns = set( ''.join(config['query'].split('SELECT ')[1:]).split( ' FROM', maxsplit=1)[0].split(', ')) if columns != {'*'}: elements = [{ column: element[column] for column in element if column in columns } for element in elements] transform['config']['elements'] = elements return test_spec @YamlExamplesTestSuite.register_test_preprocessor( 'test_streaming_sentiment_analysis_yaml') def _streaming_sentiment_analysis_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for tests that involve the streaming sentiment analysis example. This preprocessor replaces several IO transforms and the RunInference transform. This allows the test to verify the pipeline's correctness without relying on external data sources and the model hosted on VertexAI. Args: test_spec: The dictionary representation of the YAML pipeline specification. expected: A list of strings representing the expected output of the pipeline. env: The TestEnvironment object providing utilities for creating temporary files. Returns: The modified test_spec dictionary with ... transforms replaced. """ if pipeline := test_spec.get('pipeline', None): for transform in pipeline.get('transforms', []): if transform.get('type', '') == 'PyTransform' and transform.get( 'name', '') == 'ReadFromGCS': transform['windowing'] = {'type': 'fixed', 'size': '30s'} file_name = 'youtube-comments.csv' local_path = env.input_file(file_name, INPUT_FILES[file_name]) transform['config']['kwargs']['file_pattern'] = local_path if pipeline := test_spec.get('pipeline', None): for transform in pipeline.get('transforms', []): if transform.get('type', '') == 'ReadFromKafka': config = transform['config'] transform['type'] = 'ReadFromCsv' transform['config'] = { k: v for k, v in config.items() if k.startswith('__') } transform['config']['path'] = "" file_name = 'youtube-comments.csv' test_spec = replace_recursive( test_spec, transform['type'], 'path', env.input_file(file_name, INPUT_FILES[file_name])) if pipeline := test_spec.get('pipeline', None): for transform in pipeline.get('transforms', []): if transform.get('type', '') == 'RunInference': transform['type'] = 'TestRunInferenceYouTubeComments' return test_spec @YamlExamplesTestSuite.register_test_preprocessor( 'test_streaming_taxifare_prediction_yaml') def _streaming_taxifare_prediction_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for tests that involve the streaming taxi fare prediction example. This preprocessor replaces several IO transforms and the RunInference transform. This allows the test to verify the pipeline's correctness without relying on external data sources and the model hosted on VertexAI. It also turns this non-linear pipeline into a linear pipeline by replacing the ReadFromKafka and WriteToKafka transforms with MapToFields and linking the two disconnected pipeline components together. The pipeline logic, however, remains the same and is still being tested accordingly. Args: test_spec: The dictionary representation of the YAML pipeline specification. expected: A list of strings representing the expected output of the pipeline. env: The TestEnvironment object providing utilities for creating temporary files. Returns: The modified test_spec dictionary with several involved IO transforms and the RunInference transform replaced. """ if pipeline := test_spec.get('pipeline', None): for transform in pipeline.get('transforms', []): if transform.get('type', '') == 'ReadFromPubSub': transform['type'] = 'TestReadFromPubSub' transform['config']['topic'] = 'taxi-ride-topic' elif transform.get('type', '') == 'WriteToKafka': transform['type'] = 'MapToFields' transform['config'] = { k: v for (k, v) in transform.get('config', {}).items() if k.startswith('__') } transform['config']['fields'] = { 'ride_id': 'ride_id', 'pickup_longitude': 'pickup_longitude', 'pickup_latitude': 'pickup_latitude', 'pickup_datetime': 'pickup_datetime', 'dropoff_longitude': 'dropoff_longitude', 'dropoff_latitude': 'dropoff_latitude', 'passenger_count': 'passenger_count', } elif transform.get('type', '') == 'ReadFromKafka': transform['type'] = 'MapToFields' transform['config'] = { k: v for (k, v) in transform.get('config', {}).items() if k.startswith('__') } transform['input'] = 'WriteKafka' transform['config']['fields'] = { 'ride_id': 'ride_id', 'pickup_longitude': 'pickup_longitude', 'pickup_latitude': 'pickup_latitude', 'pickup_datetime': 'pickup_datetime', 'dropoff_longitude': 'dropoff_longitude', 'dropoff_latitude': 'dropoff_latitude', 'passenger_count': 'passenger_count', } elif transform.get('type', '') == 'WriteToBigQuery': transform['type'] = 'LogForTesting' transform['config'] = { k: v for (k, v) in transform.get('config', {}).items() if (k.startswith('__') or k == 'error_handling') } elif transform.get('type', '') == 'RunInference': transform['type'] = 'TestRunInferenceTaxiFare' return test_spec @YamlExamplesTestSuite.register_test_preprocessor([ 'test_iceberg_migration_yaml', 'test_ml_preprocessing_yaml', 'test_anomaly_scoring_yaml' ]) def _batch_log_analysis_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): """ Preprocessor for tests that involve the batch log analysis example. This preprocessor replaces several IO transforms and the MLTransform. This allows the test to verify the pipeline's correctness without relying on external data sources or MLTransform's many dependencies. Args: test_spec: The dictionary representation of the YAML pipeline specification. expected: A list of strings representing the expected output of the pipeline. env: The TestEnvironment object providing utilities for creating temporary files. Returns: The modified test_spec dictionary with ReadFromText transforms replaced. """ if pipeline := test_spec.get('pipeline', None): for transform in pipeline.get('transforms', []): # Mock ReadFromCsv in iceberg_migration.yaml pipeline if transform.get('type', '') == 'ReadFromCsv': file_name = 'system-logs.csv' local_path = env.input_file(file_name, INPUT_FILES[file_name]) transform['config']['path'] = local_path # Mock ReadFromIceberg in ml_preprocessing.yaml pipeline elif transform.get('type', '') == 'ReadFromIceberg': transform['type'] = 'Create' transform['config'] = { k: v for (k, v) in transform.get('config', {}).items() if (k.startswith('__')) } transform['config']['elements'] = input_data.system_logs_data() # Mock MLTransform in ml_preprocessing.yaml pipeline elif transform.get('type', '') == 'MLTransform': transform['type'] = 'MapToFields' transform['config'] = { k: v for (k, v) in transform.get('config', {}).items() if k.startswith('__') } transform['config']['language'] = 'python' transform['config']['fields'] = { 'LineId': 'LineId', 'Date': 'Date', 'Time': 'Time', 'Level': 'Level', 'Process': 'Process', 'Component': 'Component', 'Content': 'Content', 'embedding': { 'callable': f"lambda row: {input_data.embedding_data()}", } } # Mock MapToFields in ml_preprocessing.yaml pipeline elif transform.get('type', '') == 'MapToFields' and \ transform.get('name', '') == 'Normalize': transform['config']['dependencies'] = ['numpy'] # Mock ReadFromBigQuery in anomaly_scoring.yaml pipeline elif transform.get('type', '') == 'ReadFromBigQuery': transform['type'] = 'Create' transform['config'] = { k: v for (k, v) in transform.get('config', {}).items() if (k.startswith('__')) } transform['config']['elements'] = ( input_data.system_logs_embedding_data()) # Mock PyTransform in anomaly_scoring.yaml pipeline elif transform.get('type', '') == 'PyTransform' and \ transform.get('name', '') == 'AnomalyScoring': transform['type'] = 'MapToFields' transform['config'] = { k: v for (k, v) in transform.get('config', {}).items() if k.startswith('__') } transform['config']['language'] = 'python' transform['config']['fields'] = { 'example': 'embedding', 'predictions': { 'callable': """lambda row: [{ 'score': 0.65, 'label': 0, 'threshold': 0.8 }]""", } } return test_spec @YamlExamplesTestSuite.register_test_preprocessor([ 'test_wordCountInclude_yaml', 'test_wordCountImport_yaml', 'test_wordCountInheritance_yaml' ]) def _jinja_preprocessor(raw_spec_string: str, test_name: str): """ Preprocessor for Jinja-based YAML tests. This function takes a raw YAML string, which is treated as a Jinja2 template, and renders it to produce the final pipeline specification. It specifically handles templates that use the `{% include ... %}` directive by manually loading the content of the included files from the filesystem. The Jinja variables required for rendering are loaded from a predefined data source. Args: raw_spec_string: A string containing the raw YAML content, which is a Jinja2 template. Returns: A string containing the fully rendered YAML pipeline specification. """ jinja_variables = json.loads(input_data.word_count_jinja_parameter_data()) test_file_dir = os.path.dirname(__file__) sdk_root = os.path.abspath(os.path.join(test_file_dir, '../../../..')) include_files = input_data.word_count_jinja_template_data(test_name) mock_templates = {'main_template': raw_spec_string} for file_path in include_files: full_path = os.path.join(sdk_root, file_path) with open(full_path, 'r', encoding='utf-8') as f: mock_templates[file_path] = f.read() # Can't use the standard expand_jinja method due to it not supporting # `% include` jinja templization. # TODO(#35936): Maybe update expand_jinja to handle this case. jinja_env = Environment( loader=DictLoader(mock_templates), undefined=StrictUndefined) template = jinja_env.get_template('main_template') rendered_yaml_string = template.render(jinja_variables) return rendered_yaml_string INPUT_FILES = { 'products.csv': input_data.products_csv(), 'kinglear.txt': input_data.text_data(), 'youtube-comments.csv': input_data.youtube_comments_csv(), 'system-logs.csv': input_data.system_logs_csv() } KAFKA_TOPICS = {'test-topic': input_data.kafka_messages_data()} PUBSUB_TOPICS = { 'test-topic': input_data.pubsub_messages_data(), 'taxi-ride-topic': input_data.pubsub_taxi_ride_events_data() } INPUT_TABLES = { ('shipment-test', 'shipment', 'shipments'): input_data.shipments_data(), ('orders-test', 'order-database', 'orders'): input_data. spanner_orders_data(), ('db', 'users', 'NY'): input_data.iceberg_dynamic_destinations_users_data(), ('BigTable', 'beam-test', 'bigtable-enrichment-test'): input_data. bigtable_data(), ('BigQuery', 'ALL_TEST', 'customers'): input_data.bigquery_data(), ('Jdbc', 'shipment', 'shipments'): input_data.shipments_data(), ('SqlServer', 'shipment', 'shipments'): input_data.shipments_data(), ('Postgres', 'shipment', 'shipments'): input_data.shipments_data(), ('Oracle', 'shipment', 'shipments'): input_data.shipments_data(), ('MySql', 'shipment', 'shipments'): input_data.shipments_data() } YAML_DOCS_DIR = os.path.join(os.path.dirname(__file__)) AggregationTest = YamlExamplesTestSuite( 'AggregationExamplesTest', os.path.join(YAML_DOCS_DIR, '../transforms/aggregation/*.yaml')).run() BlueprintTest = YamlExamplesTestSuite( 'BlueprintExamplesTest', os.path.join(YAML_DOCS_DIR, '../transforms/blueprint/*.yaml')).run() ElementWiseTest = YamlExamplesTestSuite( 'ElementwiseExamplesTest', os.path.join(YAML_DOCS_DIR, '../transforms/elementwise/*.yaml')).run() ExamplesTest = YamlExamplesTestSuite( 'ExamplesTest', os.path.join(YAML_DOCS_DIR, '../*.yaml')).run() JinjaTest = YamlExamplesTestSuite( 'JinjaExamplesTest', os.path.join(YAML_DOCS_DIR, '../transforms/jinja/**/*.yaml')).run() IOTest = YamlExamplesTestSuite( 'IOExamplesTest', os.path.join(YAML_DOCS_DIR, '../transforms/io/*.yaml')).run() MLTest = YamlExamplesTestSuite( 'MLExamplesTest', os.path.join(YAML_DOCS_DIR, '../transforms/ml/**/*.yaml')).run() if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/examples/testing/input_data.py0000644000175100017510000002635615206113033024631 0ustar00runnerrunner# coding=utf-8 # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import json import typing from apache_beam.io.gcp.pubsub import PubsubMessage # This file contains the input data to be requested by the example tests, if # needed. def text_data(): return '\n'.join([ "Fool\tThou shouldst not have been old till thou hadst", "\tbeen wise.", "KING LEAR\tNothing will come of nothing: speak again.", "\tNever, never, never, never, never!" ]) def word_count_jinja_parameter_data(): params = { "readFromTextTransform": { "path": "gs://dataflow-samples/shakespeare/kinglear.txt" }, "mapToFieldsSplitConfig": { "language": "python", "fields": { "value": "1" } }, "explodeTransform": { "fields": "word" }, "combineTransform": { "group_by": "word", "combine": { "value": "sum" } }, "mapToFieldsCountConfig": { "language": "python", "fields": { "output": "word + \" - \" + str(value)" } }, "writeToTextTransform": { "path": "gs://apache-beam-testing-derrickaw/wordCounts/" } } return json.dumps(params) def word_count_jinja_template_data(test_name: str) -> list[str]: if test_name == 'test_wordCountInclude_yaml': return [ 'apache_beam/yaml/examples/transforms/jinja/' 'include/submodules/readFromTextTransform.yaml', 'apache_beam/yaml/examples/transforms/jinja/' 'include/submodules/mapToFieldsSplitConfig.yaml', 'apache_beam/yaml/examples/transforms/jinja/' 'include/submodules/explodeTransform.yaml', 'apache_beam/yaml/examples/transforms/jinja/' 'include/submodules/combineTransform.yaml', 'apache_beam/yaml/examples/transforms/jinja/' 'include/submodules/mapToFieldsCountConfig.yaml', 'apache_beam/yaml/examples/transforms/jinja/' 'include/submodules/writeToTextTransform.yaml' ] elif test_name == 'test_wordCountImport_yaml': return [ 'apache_beam/yaml/examples/transforms/jinja/' 'import/macros/wordCountMacros.yaml' ] elif test_name == 'test_wordCountInheritance_yaml': return [ 'apache_beam/yaml/examples/transforms/jinja/' 'inheritance/base/base_pipeline.yaml' ] return [] def iceberg_dynamic_destinations_users_data(): return [{ 'id': 3, 'name': 'Smith', 'email': 'smith@example.com', 'zip': 'NY' }, { 'id': 4, 'name': 'Beamberg', 'email': 'beamberg@example.com', 'zip': 'NY' }] def products_csv(): return '\n'.join([ 'transaction_id,product_name,category,price', 'T0012,Headphones,Electronics,59.99', 'T5034,Leather Jacket,Apparel,109.99', 'T0024,Aluminum Mug,Kitchen,29.99', 'T0104,Headphones,Electronics,59.99', 'T0302,Monitor,Electronics,249.99' ]) def youtube_comments_csv(): return '\n'.join([ 'video_id,comment_text,likes,replies', 'XpVt6Z1Gjjo,I AM HAPPY,1,1', 'XpVt6Z1Gjjo,I AM SAD,1,1', 'XpVt6Z1Gjjo,§ÁĐ,1,1' ]) def spanner_orders_data(): return [{ 'order_id': 1, 'customer_id': 1001, 'product_id': 2001, 'order_date': '24-03-24', 'order_amount': 150, }, { 'order_id': 2, 'customer_id': 1002, 'product_id': 2002, 'order_date': '19-04-24', 'order_amount': 90, }, { 'order_id': 3, 'customer_id': 1003, 'product_id': 2003, 'order_date': '7-05-24', 'order_amount': 110, }] def shipments_data(): return [{ 'shipment_id': 'S1', 'customer_id': 'C1', 'shipment_date': '2023-05-01', 'shipment_cost': 150.0, 'customer_name': 'Alice', 'customer_email': 'alice@example.com' }, { 'shipment_id': 'S2', 'customer_id': 'C2', 'shipment_date': '2023-06-12', 'shipment_cost': 300.0, 'customer_name': 'Bob', 'customer_email': 'bob@example.com' }, { 'shipment_id': 'S3', 'customer_id': 'C1', 'shipment_date': '2023-05-10', 'shipment_cost': 20.0, 'customer_name': 'Alice', 'customer_email': 'alice@example.com' }, { 'shipment_id': 'S4', 'customer_id': 'C4', 'shipment_date': '2024-07-01', 'shipment_cost': 150.0, 'customer_name': 'Derek', 'customer_email': 'derek@example.com' }, { 'shipment_id': 'S5', 'customer_id': 'C5', 'shipment_date': '2023-05-09', 'shipment_cost': 300.0, 'customer_name': 'Erin', 'customer_email': 'erin@example.com' }, { 'shipment_id': 'S6', 'customer_id': 'C4', 'shipment_date': '2024-07-02', 'shipment_cost': 150.0, 'customer_name': 'Derek', 'customer_email': 'derek@example.com' }] def bigtable_data(): return [{ 'product_id': '1', 'product_name': 'pixel 5', 'product_stock': '2' }, { 'product_id': '2', 'product_name': 'pixel 6', 'product_stock': '4' }, { 'product_id': '3', 'product_name': 'pixel 7', 'product_stock': '20' }, { 'product_id': '4', 'product_name': 'pixel 8', 'product_stock': '10' }, { 'product_id': '5', 'product_name': 'pixel 11', 'product_stock': '3' }, { 'product_id': '6', 'product_name': 'pixel 12', 'product_stock': '7' }, { 'product_id': '7', 'product_name': 'pixel 13', 'product_stock': '8' }, { 'product_id': '8', 'product_name': 'pixel 14', 'product_stock': '3' }] def bigquery_data(): return [{ 'customer_id': 1001, 'customer_name': 'Alice', 'customer_email': 'alice@gmail.com' }, { 'customer_id': 1002, 'customer_name': 'Bob', 'customer_email': 'bob@gmail.com' }, { 'customer_id': 1003, 'customer_name': 'Claire', 'customer_email': 'claire@gmail.com' }] def pubsub_messages_data(): """ Provides a list of PubsubMessage objects for testing. """ return [ PubsubMessage(data=b"{\"label\": \"37a\", \"rank\": 1}", attributes={}), PubsubMessage(data=b"{\"label\": \"37b\", \"rank\": 4}", attributes={}), PubsubMessage(data=b"{\"label\": \"37c\", \"rank\": 3}", attributes={}), PubsubMessage(data=b"{\"label\": \"37d\", \"rank\": 2}", attributes={}), ] def pubsub_taxi_ride_events_data(): """ Provides a list of PubsubMessage objects for testing taxi ride events. """ return [ PubsubMessage( data=b"{\"ride_id\": \"1\", \"longitude\": 11.0, \"latitude\": -11.0," b"\"passenger_count\": 1, \"meter_reading\": 100.0, \"timestamp\": " b"\"2025-01-01T00:29:00.00000-04:00\", \"ride_status\": \"pickup\"}", attributes={}), PubsubMessage( data=b"{\"ride_id\": \"2\", \"longitude\": 22.0, \"latitude\": -22.0," b"\"passenger_count\": 2, \"meter_reading\": 100.0, \"timestamp\": " b"\"2025-01-01T00:30:00.00000-04:00\", \"ride_status\": \"pickup\"}", attributes={}), PubsubMessage( data=b"{\"ride_id\": \"1\", \"longitude\": 13.0, \"latitude\": -13.0," b"\"passenger_count\": 1, \"meter_reading\": 100.0, \"timestamp\": " b"\"2025-01-01T00:31:00.00000-04:00\", \"ride_status\": \"enroute\"}", attributes={}), PubsubMessage( data=b"{\"ride_id\": \"2\", \"longitude\": 24.0, \"latitude\": -24.0," b"\"passenger_count\": 2, \"meter_reading\": 100.0, \"timestamp\": " b"\"2025-01-01T00:32:00.00000-04:00\", \"ride_status\": \"enroute\"}", attributes={}), PubsubMessage( data=b"{\"ride_id\": \"3\", \"longitude\": 33.0, \"latitude\": -33.0," b"\"passenger_count\": 3, \"meter_reading\": 100.0, \"timestamp\": " b"\"2025-01-01T00:35:00.00000-04:00\", \"ride_status\": \"enroute\"}", attributes={}), PubsubMessage( data=b"{\"ride_id\": \"4\", \"longitude\": 44.0, \"latitude\": -44.0," b"\"passenger_count\": 4, \"meter_reading\": 100.0, \"timestamp\": " b"\"2025-01-01T00:35:00.00000-04:00\", \"ride_status\": \"dropoff\"}", attributes={}), PubsubMessage( data=b"{\"ride_id\": \"1\", \"longitude\": 15.0, \"latitude\": -15.0," b"\"passenger_count\": 1, \"meter_reading\": 100.0, \"timestamp\": " b"\"2025-01-01T00:33:00.00000-04:00\", \"ride_status\": \"dropoff\"}", attributes={}), PubsubMessage( data=b"{\"ride_id\": \"2\", \"longitude\": 26.0, \"latitude\": -26.0," b"\"passenger_count\": 2, \"meter_reading\": 100.0, \"timestamp\": " b"\"2025-01-01T00:34:00.00000-04:00\", \"ride_status\": \"dropoff\"}", attributes={}), ] def kafka_messages_data(): """ Provides a list of Kafka messages for testing. """ return [data.encode('utf-8') for data in text_data().split('\n')] class TaxiRideEventSchema(typing.NamedTuple): ride_id: str longitude: float latitude: float passenger_count: int meter_reading: float timestamp: str ride_status: str def system_logs_csv(): return '\n'.join([ 'LineId,Date,Time,Level,Process,Component,Content', '1,2024-10-01,12:00:00,INFO,Main,ComponentA,System started successfully', '2,2024-10-01,12:00:05,WARN,Main,ComponentA,Memory usage is high', '3,2024-10-01,12:00:10,ERROR,Main,ComponentA,Task failed due to timeout', ]) def system_logs_data(): csv_data = system_logs_csv() lines = csv_data.strip().split('\n') headers = lines[0].split(',') logs = [] for row in lines[1:]: values = row.split(',') log = dict(zip(headers, values)) log['LineId'] = int(log['LineId']) logs.append(log) return logs def embedding_data(): return [0.1, 0.2, 0.3, 0.4, 0.5] def system_logs_embedding_data(): csv_data = system_logs_csv() lines = csv_data.strip().split('\n') headers = lines[0].split(',') headers.append('embedding') logs = [] for row in lines[1:]: values = row.split(',') values.append(embedding_data()) log = dict(zip(headers, values)) log['LineId'] = int(log['LineId']) logs.append(log) return logs ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/generate_yaml_docs.py0000644000175100017510000004644015206113033023026 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import argparse import glob import io import itertools import os import re import docstring_parser import yaml from apache_beam.portability.api import schema_pb2 from apache_beam.typehints import schemas from apache_beam.utils import subprocess_server from apache_beam.utils.python_callable import PythonCallableWithSource from apache_beam.version import __version__ as beam_version from apache_beam.yaml import json_utils from apache_beam.yaml import yaml_provider from apache_beam.yaml import yaml_utils from apache_beam.yaml.yaml_errors import ErrorHandlingConfig def _singular(name): # Simply removing an 's' (or 'es', or 'ies', ...) may result in surprising # manglings. Better to play it safe and leave a correctly-spelled plural # than a botched singular in our examples configs. return { 'args': 'arg', 'attributes': 'attribute', 'elements': 'element', 'fields': 'field', }.get(name, name) def _fake_value(name, beam_type): type_info = beam_type.WhichOneof("type_info") if type_info == "atomic_type": if beam_type.atomic_type == schema_pb2.STRING: return f'"{name}"' elif beam_type.atomic_type == schema_pb2.BOOLEAN: return "true|false" else: return name elif type_info == "array_type": return [ _fake_value(_singular(name), beam_type.array_type.element_type), _fake_value(_singular(name), beam_type.array_type.element_type), '...' ] elif type_info == "iterable_type": return [ _fake_value(_singular(name), beam_type.iterable_type.element_type), _fake_value(_singular(name), beam_type.iterable_type.element_type), '...' ] elif type_info == "map_type": if beam_type.map_type.key_type.atomic_type == schema_pb2.STRING: return { 'a': _fake_value(name + '_value_a', beam_type.map_type.value_type), 'b': _fake_value(name + '_value_b', beam_type.map_type.value_type), 'c': '...', } else: return { _fake_value(name + '_key', beam_type.map_type.key_type): _fake_value( name + '_value', beam_type.map_type.value_type) } elif type_info == "row_type": return _fake_row(beam_type.row_type.schema) elif type_info == "logical_type": return name else: raise ValueError(f"Unrecognized type_info: {type_info!r}") def _fake_row(schema): if schema is None: return '...' return {f.name: _fake_value(f.name, f.type) for f in schema.fields} def pretty_example(provider, t, base_t=None): spec = {'type': base_t or t} try: requires_inputs = provider.requires_inputs(t, {}) except Exception: requires_inputs = False if requires_inputs: spec['input'] = '...' config_schema = provider.config_schema(t) if config_schema is None or config_schema.fields: spec['config'] = _fake_row(config_schema) s = yaml.dump(spec, sort_keys=False) return s.replace("'", "") def config_docs(schema): if schema is None: return '' elif not schema.fields: return 'No configuration parameters.' def pretty_type(beam_type): type_info = beam_type.WhichOneof("type_info") if type_info == "atomic_type": return schema_pb2.AtomicType.Name(beam_type.atomic_type).lower() elif type_info == "array_type": return f'Array[{pretty_type(beam_type.array_type.element_type)}]' elif type_info == "iterable_type": return f'Iterable[{pretty_type(beam_type.iterable_type.element_type)}]' elif type_info == "map_type": return ( f'Map[{pretty_type(beam_type.map_type.key_type)}, ' f'{pretty_type(beam_type.map_type.value_type)}]') elif type_info == "row_type": return 'Row' else: return '?' def maybe_row_parameters(t): if t.WhichOneof("type_info") == "row_type": return indent('\n\nRow fields:\n\n' + config_docs(t.row_type.schema), 4) else: return '' def maybe_optional(t): return " (Optional)" if t.nullable else "" def normalize_error_handling(f): doc = docstring_parser.parse( ErrorHandlingConfig.__doc__, docstring_parser.DocstringStyle.GOOGLE) if f.name == "error_handling": f = schema_pb2.Field( name="error_handling", type=schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( fields=[ schemas.schema_field( param.arg_name, PythonCallableWithSource.load_from_expression( param.type_name), param.description) for param in doc.params ])), nullable=True), description=f.description or doc.short_description) return f def lines(): for f in schema.fields: f = normalize_error_handling(f) yield ''.join([ f'**{f.name}** `{pretty_type(f.type)}`', maybe_optional(f.type), indent(': ' + f.description if f.description else '', 2), maybe_row_parameters(f.type), ]) return '\n\n'.join('*' + indent(line, 2) for line in lines()).strip() def indent(lines, size): return '\n'.join(' ' * size + line for line in lines.split('\n')) def longest(func, xs): return max([func(x) or '' for x in xs], key=len) def io_grouping_key(transform_name): """Place reads and writes next to each other, after all other transforms.""" if transform_name.startswith('ReadFrom'): return 1, transform_name[8:], 0 elif transform_name.startswith('WriteTo'): return 1, transform_name[7:], 1 else: return 0, transform_name # Exclude providers SKIP = {} def add_transform_links(transform, description, provider_list): """ Convert references of Providers to urls that link to their respective pages. For example, "Some description talking about MyTransform." would be converted to "Some description talking about MyTransform" meanwhile:: type: MyTransform config: ... Would remain unchanged. Avoid self-linking within a Transform page. """ for p in provider_list: # Match all instances of built-in transforms within the description # excluding the transform whose description is currently being evaluated. # Match the entire word boundary so that partial matches do not count. # (i.e. OtherTransform should not match Transform) description = re.sub( rf"(?{p}', description or '') return description def transform_docs(transform_base, transforms, providers, extra_docs=''): return '\n'.join([ f'## {transform_base}', '', longest( lambda t: longest( lambda p: add_transform_links( t, p.description(t), providers.keys()), providers[t]), transforms).replace('::\n', '\n\n :::yaml\n'), '', extra_docs, '', '### Configuration', '', longest( lambda t: longest( lambda p: config_docs(p.config_schema(t)), providers[t]), transforms), '', '### Usage', '', ' :::yaml', '', indent( longest( lambda t: longest( lambda p: pretty_example(p, t, transform_base), providers[t]), transforms), 4), ]) def create_index(include, exclude, options): with subprocess_server.SubprocessServer.cache_subprocesses(): json_config_schemas = [] markdown_out = io.StringIO() if options.provider_source_file: providers = yaml_provider.merge_providers( yaml_provider.load_providers( yaml_utils.locate_data_file(options.provider_source_file))) else: providers = yaml_provider.standard_providers() for transform_base, transforms in itertools.groupby( sorted(providers.keys(), key=io_grouping_key), key=lambda s: s.split('-')[0]): transforms = list(transforms) if include(transform_base) and not exclude(transform_base): print(transform_base) if options.markdown_file or options.html_file: if '-' in transforms[0]: extra_docs = 'Supported languages: ' + ', '.join( t.split('-')[-1] for t in sorted(transforms)) + '.' else: extra_docs = '' markdown_out.write( transform_docs(transform_base, transforms, providers, extra_docs)) markdown_out.write('\n\n') if options.schema_file: for transform in transforms: schema = providers[transform][0].config_schema(transform) if schema: json_config_schemas.append({ 'if': { 'properties': { 'type': { 'const': transform } } }, 'then': { 'properties': { 'config': { 'type': 'object', 'properties': { '__line__': { 'type': 'integer' }, '__uuid__': {}, **{ f.name: # json_utils.beam_type_to_json_type(f.type) for f in schema.fields } }, 'additionalProperties': False, } } } }) return json_config_schemas, markdown_out.getvalue() def create_examples_markdown(): markdown_out = io.StringIO() base = os.path.join(os.path.dirname(__file__), 'examples') section = last_section = '' for path in sorted(glob.glob(os.path.join(base, '**', '*.yaml'), recursive=True), key=lambda path: (path.count(os.sep), path)): short_path = path[len(base):].replace('transforms', '').strip(os.sep) def to_title(path): base, _ = os.path.splitext(path) nice = base.replace('_', ' ').replace(os.sep, ' ').title() # These acronyms should be upper, not title. nice = re.sub(r'\bMl\b', 'ML', nice) nice = re.sub(r'\bIo\b', 'IO', nice) return nice def clean_yaml(content): content = re.sub( '# Licensed to the Apache Software Foundation.*' '# limitations under the License.', '', content, flags=re.MULTILINE | re.DOTALL) content = re.sub('# coding=.*', '', content) return content def split_header(yaml): lines = yaml.split('\n') ix = 0 # make lint happy for ix, line in enumerate(lines): if not line.strip(): continue if not line.startswith('#'): break return ( '\n'.join([line[1:].strip() for line in lines[:ix]]), '\n'.join(lines[ix:])) if os.sep in short_path: section = to_title(short_path.split(os.sep)[0]) if section != last_section: markdown_out.write(f'# {section}\n\n') last_section = section title = to_title(short_path)[len(section):] markdown_out.write(f'## {title}\n\n') with open(path) as fin: content = fin.read() header, body = split_header(clean_yaml(content)) markdown_out.write(header) markdown_out.write('\n\n :::yaml\n\n') markdown_out.write(' ' + body.replace('\n', '\n ')) markdown_out.write('\n') return markdown_out.getvalue() def markdown_to_html(title, markdown_content, header=''): import markdown import markdown.extensions.toc import pygments.formatters md = markdown.Markdown( extensions=[ markdown.extensions.toc.TocExtension(toc_depth=2), 'codehilite', ]) pygments_style = pygments.formatters.HtmlFormatter().get_style_defs( '.codehilite') extra_style = ''' * { box-sizing: border-box; } body { font-family: 'Roboto', sans-serif; font-weight: normal; color: #404040; background: #edf0f2; } .body-for-nav { background: #fcfcfc; } .grid-for-nav { width: 100%; } .nav-side { position: fixed; top: 0; left: 0; width: 300px; height: 100%; padding-bottom: 2em; color: #9b9b9b; background: #343131; } .nav-header { display: block; width: 300px; padding: 1em; background-color: #2980B9; text-align: center; color: #fcfcfc; } .nav-header a { color: #fcfcfc; font-weight: bold; display: inline-block; padding: 4px 6px; margin-bottom: 1em; text-decoration:none; } .nav-header>div.version { margin-top: -.5em; margin-bottom: 1em; font-weight: normal; color: rgba(255, 255, 255, 0.3); } .toc { width: 300px; text-align: left; overflow-y: auto; max-height: calc(100% - 4.3em); scrollbar-width: thin; scrollbar-color: #9b9b9b #343131; } .toc ul { margin: 0; padding: 0; list-style: none; } .toc li { border-bottom: 1px solid #4e4a4a; margin-left: 1em; } .toc a { display: block; line-height: 36px; font-size: 90%; color: #d9d9d9; padding: .1em 0.6em; text-decoration: none; transition: background-color 0.3s ease, color 0.3s ease; } .toc a:hover { background-color: #4e4a4a; color: #ffffff; } .transform-content-wrap { margin-left: 300px; background: #fcfcfc; } .transform-content { padding: 1.5em 3em; margin: 20px; padding-bottom: 2em; } .transform-content li::marker { display: inline-block; width: 0.5em; } .transform-content h1 { font-size: 40px; } .transform-content ul { margin-left: 0.75em; text-align: left; list-style-type: disc; } hr { color: gray; display: block; height: 1px; border: 0; border-top: 1px solid #e1e4e5; margin-bottom: 3em; margin-top: 3em; padding: 0; } .codehilite { background: #f5f5f5; border: 1px solid #ccc; border-radius: 4px; padding: 0.2em 1em; overflow: auto; font-family: monospace; font-size: 14px; line-height: 1.5; } p code, li code { white-space: nowrap; max-width: 100%; background: #fff; border: solid 1px #e1e4e5; padding: 0 5px; font-family: monospace; color: #404040; font-weight: bold; padding: 2px 5px; } ''' html = md.convert(markdown_content) return f''' {title}

{title}

{header} {html.replace('
''' def main(): parser = argparse.ArgumentParser() parser.add_argument('--examples_file') parser.add_argument('--markdown_file') parser.add_argument('--html_file') parser.add_argument('--schema_file') parser.add_argument('--include', default='.*') parser.add_argument('--exclude', default='') parser.add_argument( "--provider_source_file", help="Path to a YAML file containing custom providers. " "If not provided, uses standard Beam providers.", ) options = parser.parse_args() include = re.compile(options.include).match exclude = ( re.compile(options.exclude).match if options.exclude else lambda x: x in SKIP) json_config_schemas, markdown_content = create_index( include, exclude, options) if options.schema_file: with open(options.schema_file, 'w') as fout: yaml.dump(json_config_schemas, fout, sort_keys=False) if options.markdown_file: with open(options.markdown_file, 'w') as fout: fout.write(markdown_content) if options.html_file: with open(options.html_file, 'w') as html_out: html_out.write( markdown_to_html('Beam YAML Transform Index', markdown_content)) if options.examples_file: with open(options.examples_file, 'w') as html_out: html_out.write( markdown_to_html( 'Beam YAML Examples', create_examples_markdown(), header='''

Example pipelines using the Beam YAML API. These examples can also be found on github.

Note: These examples below are automatically tested for correctness and may be used as a starting point for your own pipelines.

''')) if __name__ == '__main__': main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/integration_tests.py0000644000175100017510000007130115206113033022741 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Runs integration tests in the tests directory.""" import contextlib import copy import glob import itertools import logging import os import random import secrets import sqlite3 import string import unittest import uuid from datetime import datetime from datetime import timezone import mock import psycopg2 import pytds import sqlalchemy import yaml from apitools.base.py.exceptions import HttpError from google.cloud import pubsub_v1 from google.cloud.bigtable import client from google.cloud.bigtable_admin_v2.types import instance from testcontainers.core.container import DockerContainer from testcontainers.core.waiting_utils import wait_for_logs from testcontainers.google import PubSubContainer from testcontainers.kafka import KafkaContainer from testcontainers.mssql import SqlServerContainer from testcontainers.mysql import MySqlContainer from testcontainers.postgres import PostgresContainer import apache_beam as beam from apache_beam.io import filesystems from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.io.gcp.spanner_wrapper import SpannerWrapper from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.utils import python_callable from apache_beam.yaml import yaml_provider from apache_beam.yaml import yaml_transform from apache_beam.yaml.conftest import yaml_test_files_dir _LOGGER = logging.getLogger(__name__) @contextlib.contextmanager def gcs_temp_dir(bucket): """Context manager to create and clean up a temporary GCS directory. Creates a unique temporary directory within the specified GCS bucket and yields the path. Upon exiting the context, the directory and its contents are deleted. Args: bucket (str): The GCS bucket name (e.g., 'gs://my-bucket'). Yields: str: The full path to the created temporary GCS directory. Example: 'gs://my-bucket/yaml-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx' """ gcs_tempdir = bucket + '/yaml-' + str(uuid.uuid4()) yield gcs_tempdir filesystems.FileSystems.delete([gcs_tempdir]) @contextlib.contextmanager def temp_spanner_table(project, prefix='temp_spanner_db_'): """Context manager to create and clean up a temporary Spanner database and table. Creates a unique temporary Spanner database within the specified project and a predefined table named 'tmp_table' with columns ['UserId', 'Key']. It yields connection details for the created resources. Upon exiting the context, the temporary database (and its table) is deleted. Args: project (str): The Google Cloud project ID. prefix (str): A prefix to use for the temporary database name. Defaults to 'temp_spanner_db_'. Yields: list[str]: A list containing connection details: [project_id, instance_id, database_id, table_name, list_of_columns]. Example: ['my-project', 'beam-test', 'temp_spanner_db_...', 'tmp_table', ['UserId', 'Key']] """ spanner_client = SpannerWrapper(project) spanner_client._create_database() instance = "beam-test" database = spanner_client._test_database table = 'tmp_table' columns = ['UserId', 'Key'] logging.info("Created Spanner database: %s", database) try: yield [f'{project}', f'{instance}', f'{database}', f'{table}', columns] finally: logging.info("Deleting Spanner database: %s", database) spanner_client._delete_database() @contextlib.contextmanager def temp_bigquery_table(project, prefix='yaml_bq_it_'): """Context manager to create and clean up a temporary BigQuery dataset. Creates a unique temporary BigQuery dataset within the specified project. It yields a placeholder table name string within that dataset (e.g., 'project.dataset_id.tmp_table'). The actual table is expected to be created by the test using this context. Upon exiting the context, the temporary dataset and all its contents (including any tables created within it) are deleted. Args: project (str): The Google Cloud project ID. prefix (str): A prefix to use for the temporary dataset name. Defaults to 'yaml_bq_it_'. Yields: str: The full path for a temporary BigQuery table within the created dataset. Example: 'my-project.yaml_bq_it_a1b2c3d4e5f6...tmp_table' """ bigquery_client = BigQueryWrapper() dataset_id = '%s_%s' % (prefix, uuid.uuid4().hex) bigquery_client.get_or_create_dataset(project, dataset_id) logging.info("Created dataset %s in project %s", dataset_id, project) yield f'{project}.{dataset_id}.tmp_table' request = bigquery.BigqueryDatasetsDeleteRequest( projectId=project, datasetId=dataset_id, deleteContents=True) logging.info("Deleting dataset %s in project %s", dataset_id, project) bigquery_client.client.datasets.Delete(request) def instance_prefix(instance): datestr = "".join(filter(str.isdigit, str(datetime.now(timezone.utc).date()))) instance_id = '%s-%s-%s' % (instance, datestr, secrets.token_hex(4)) assert len(instance_id) < 34, "instance id length needs to be within [6, 33]" return instance_id @contextlib.contextmanager def temp_bigtable_table(project, prefix='yaml_bt_it_'): INSTANCE = "bt-write-tests" TABLE_ID = "test-table" instance_id = instance_prefix(INSTANCE) clientT = client.Client(admin=True, project=project) # create cluster and instance instanceT = clientT.instance( instance_id, display_name=INSTANCE, instance_type=instance.Instance.Type.DEVELOPMENT) cluster = instanceT.cluster("test-cluster", "us-central1-a") operation = instanceT.create(clusters=[cluster]) operation.result(timeout=500) _LOGGER.info("Created instance [%s] in project [%s]", instance_id, project) # create table inside instance table = instanceT.table(TABLE_ID) table.create() _LOGGER.info("Created table [%s]", table.table_id) # in the table that is created, make a new family called cf1 col_fam = table.column_family('cf1') col_fam.create() # another family called cf2 col_fam = table.column_family('cf2') col_fam.create() #yielding the tmp table for all the bigTable tests yield instance_id #try catch for deleting table and instance after all tests are ran try: _LOGGER.info("Deleting table [%s]", table.table_id) table.delete() instanceT.delete() except HttpError: _LOGGER.warning("Failed to clean up instance") @contextlib.contextmanager def temp_sqlite_database(prefix='yaml_jdbc_it_'): """Context manager to provide a temporary SQLite database via JDBC for testing. This function creates a temporary SQLite database file on the local filesystem. It establishes a connection using 'sqlite3', creates a predefined 'tmp_table', and then yields a JDBC connection string suitable for use in tests that require a generic JDBC connection (specifically configured for SQLite in this case). The SQLite database file is automatically cleaned up (closed and deleted) when the context manager exits. Args: prefix (str): A prefix to use for the temporary database file name. Yields: str: A JDBC connection string for the temporary SQLite database. Example format: "jdbc:sqlite:" Raises: sqlite3.Error: If there's an error connecting to or interacting with the SQLite database during setup. Exception: Any other exception encountered during the setup or cleanup process. """ conn = cursor = None try: # Establish connection to the temp file db_name = f'{prefix}{uuid.uuid4().hex}.db' conn = sqlite3.connect(db_name) cursor = conn.cursor() # Create a temp table for tests cursor.execute( ''' CREATE TABLE tmp_table ( value INTEGER PRIMARY KEY, rank INTEGER ) ''') conn.commit() yield f'jdbc:sqlite:{db_name}' except (sqlite3.Error, Exception) as err: logging.error("Error interacting with temporary SQLite DB: %s", err) raise err finally: # Close connections if cursor: cursor.close() if conn: conn.close() try: if os.path.exists(db_name): os.remove(db_name) except Exception as err: logging.error("Error deleting temporary SQLite DB: %s", err) raise err @contextlib.contextmanager def temp_mysql_database(): """Context manager to provide a temporary MySQL database for testing. This function utilizes the 'testcontainers' library to spin up a MySQL instance within a Docker container. It then connects to this temporary database using 'mysql.connector', creates a predefined 'tmp_table', and yields a JDBC connection string suitable for use in tests. The Docker container and the database instance are automatically managed and torn down when the context manager exits. Yields: str: A JDBC connection string for the temporary MySQL database. Example format: "jdbc:mysql://:/? user=&password=" Raises: mysql.connector.Error: If there's an error connecting to or interacting with the MySQL database during setup. Exception: Any other exception encountered during the setup process. """ with MySqlContainer(init=True, dialect='pymysql') as mysql_container: # Make connection to temp database and create tmp table engine = sqlalchemy.create_engine(mysql_container.get_connection_url()) with engine.begin() as connection: connection.execute( sqlalchemy.text( "CREATE TABLE tmp_table (value INTEGER, `rank` INTEGER);")) # Construct the JDBC url for connections later on by tests jdbc_url = ( f"jdbc:mysql://{mysql_container.get_container_host_ip()}:" f"{mysql_container.get_exposed_port(mysql_container.port)}/" f"{mysql_container.dbname}?" f"user={mysql_container.username}&" f"password={mysql_container.password}") yield jdbc_url @contextlib.contextmanager def temp_postgres_database(): """Context manager to provide a temporary PostgreSQL database for testing. This function utilizes the 'testcontainers' library to spin up a PostgreSQL instance within a Docker container. It then connects to this temporary database using 'psycopg2', creates a predefined 'tmp_table', and yields a JDBC connection string suitable for use in tests. The Docker container and the database instance are automatically managed and torn down when the context manager exits. Yields: str: A JDBC connection string for the temporary PostgreSQL database. Example format: "jdbc:postgresql://:/? user=&password=" Raises: psycopg2.Error: If there's an error connecting to or interacting with the PostgreSQL database during setup. Exception: Any other exception encountered during the setup process. """ default_port = 5432 # Start the postgress container using testcontainers with PostgresContainer(port=default_port) as postgres_container: try: # Make connection to temp database and create tmp table engine = sqlalchemy.create_engine(postgres_container.get_connection_url()) with engine.begin() as connection: connection.execute( sqlalchemy.text( "CREATE TABLE tmp_table (value INTEGER, rank INTEGER);")) # Construct the JDBC url for connections later on by tests jdbc_url = ( f"jdbc:postgresql://{postgres_container.get_container_host_ip()}:" f"{postgres_container.get_exposed_port(default_port)}/" f"{postgres_container.dbname}?" f"user={postgres_container.username}&" f"password={postgres_container.password}") yield jdbc_url except (psycopg2.Error, Exception) as err: logging.error("Error interacting with temporary Postgres DB: %s", err) raise err @contextlib.contextmanager def temp_sqlserver_database(): """Context manager to provide a temporary SQL Server database for testing. This function utilizes the 'testcontainers' library to spin up a Microsoft SQL Server instance within a Docker container. It then connects to this temporary database using 'pytds', creates a predefined 'tmp_table', and yields a JDBC connection string suitable for use in tests. The Docker container and the database instance are automatically managed and torn down when the context manager exits. Yields: str: A JDBC connection string for the temporary SQL Server database. Example format: "jdbc:sqlserver://:; databaseName=; user=; password=; encrypt=false; trustServerCertificate=true" Raises: pytds.Error: If there's an error connecting to or interacting with the SQL Server database during setup. Exception: Any other exception encountered during the setup process. """ default_port = 1433 # Start the sql server using testcontainers with SqlServerContainer(port=default_port, dialect='mssql+pytds') as sqlserver_container: try: # Make connection to temp database and create tmp table engine = sqlalchemy.create_engine( sqlserver_container.get_connection_url()) with engine.begin() as connection: connection.execute( sqlalchemy.text( "CREATE TABLE tmp_table (value INTEGER, rank INTEGER);")) # Construct the JDBC url for connections later on by tests # NOTE: encrypt=false and trustServerCertificate=true is generally # needed for test container connections without proper certificates setup jdbc_url = ( f"jdbc:sqlserver://{sqlserver_container.get_container_host_ip()}:" f"{int(sqlserver_container.get_exposed_port(default_port))};" f"databaseName={sqlserver_container.dbname};" f"user={sqlserver_container.username};" f"password={sqlserver_container.password};" f"encrypt=true;" f"trustServerCertificate=true") yield jdbc_url except (pytds.Error, Exception) as err: logging.error("Error interacting with temporary SQL Server DB: %s", err) raise err class OracleTestContainer(DockerContainer): """ OracleTestContainer is an updated version of OracleDBContainer that goes ahead and sets the oracle password, waits for logs to establish that the container is ready before calling get_exposed_port, and uses a more modern oracle driver. """ def __init__(self): super().__init__("gvenzl/oracle-xe:21-slim") self.with_env("ORACLE_PASSWORD", "oracle") self.with_exposed_ports(1521) def start(self): super().start() wait_for_logs(self, "DATABASE IS READY TO USE!", timeout=300) return self def get_connection_url(self): port = self.get_exposed_port(1521) return \ f"oracle+oracledb://system:oracle@localhost:{port}/?service_name=XEPDB1" @contextlib.contextmanager def temp_oracle_database(): """Context manager to provide a temporary Oracle database for testing. This function utilizes the 'testcontainers' library to spin up an Oracle Database instance within a Docker container. It then connects to this temporary database using 'oracledb', creates a predefined NOTE: A custom OracleTestContainer class was created due to the current version (OracleDBContainer) that calls get_exposed_port too soon causing the service to hang until timeout. Yields: str: A JDBC connection string for the temporary Oracle database. Example format: "jdbc:oracle:thin:system/oracle@localhost:{port}/XEPDB1" Raises: oracledb.Error: If there's an error connecting to or interacting with the Oracle database during setup. Exception: Any other exception encountered during the setup process. """ with OracleTestContainer() as oracle: engine = sqlalchemy.create_engine(oracle.get_connection_url()) with engine.connect() as connection: connection.execute( sqlalchemy.text( """ CREATE TABLE tmp_table ( value NUMBER PRIMARY KEY, rank NUMBER ) """)) connection.commit() port = oracle.get_exposed_port(1521) yield f"jdbc:oracle:thin:system/oracle@localhost:{port}/XEPDB1" @contextlib.contextmanager def temp_kafka_server(): """Context manager to provide a temporary Kafka server for testing. This function utilizes the 'testcontainers' library to spin up a Kafka instance within a Docker container. It then yields the bootstrap server string, which can be used by Kafka clients to connect to this temporary server. The Docker container and the Kafka instance are automatically managed and torn down when the context manager exits. Yields: str: The bootstrap server string for the temporary Kafka instance. Example format: "localhost:XXXXX" or "PLAINTEXT://localhost:XXXXX" Raises: Exception: If there's an error starting the Kafka container or interacting with the temporary Kafka server. """ with KafkaContainer() as kafka_container: try: yield kafka_container.get_bootstrap_server() except Exception as err: logging.error("Error interacting with temporary Kakfa Server: %s", err) raise err @contextlib.contextmanager def temp_pubsub_emulator(project_id="apache-beam-testing"): """ Context manager to provide a temporary Pub/Sub emulator for testing. This function uses 'testcontainers' to spin up a Google Cloud SDK container running the Pub/Sub emulator. It yields the emulator host string (e.g., "localhost:xxxxx") that can be used to configure Pub/Sub clients. The Docker container is automatically managed and torn down when the context manager exits. Args: project_id (str): The GCP project ID to be used by the emulator. This doesn't need to be a real project for the emulator. Yields: str: The host and port for the Pub/Sub emulator (e.g., "localhost:xxxx"). This will be the address to point your Pub/Sub client to. Raises: Exception: If the container fails to start or the emulator isn't ready. """ with PubSubContainer(project=project_id) as pubsub_container: publisher = pubsub_v1.PublisherClient() random_front_charactor = random.choice(string.ascii_lowercase) topic_id = f"{random_front_charactor}{uuid.uuid4().hex[:8]}" topic_name_to_create = \ f"projects/{pubsub_container.project}/topics/{topic_id}" created_topic_object = publisher.create_topic(name=topic_name_to_create) yield created_topic_object.name def replace_recursive(spec, vars): """Recursively replaces string placeholders in a spec with values from vars. Traverses a nested structure (dicts, lists, or other types). If a string is encountered and contains placeholders in the format '{key}', it attempts to replace them using the `vars` dictionary. Args: spec: The (potentially nested) structure to process. vars: A dictionary of variable names to their replacement values. Returns: The spec with placeholders replaced. Raises: ValueError: If a string formatting error occurs. """ if isinstance(spec, dict): return { key: replace_recursive(value, vars) for (key, value) in spec.items() } elif isinstance(spec, list): return [replace_recursive(value, vars) for value in spec] # TODO(https://github.com/apache/beam/issues/35067): Consider checking for # callable in the if branch above instead of checking lambda here. elif isinstance( spec, str) and '{' in spec and '{\n' not in spec and 'lambda' not in spec: try: return spec.format(**vars) except Exception as exn: raise ValueError(f"Error evaluating {spec}: {exn}") from exn else: return spec def transform_types(spec): """Recursively extracts all transform types from a pipeline specification. This generator function traverses a nested pipeline specification (likely parsed from YAML). It identifies and yields the 'type' string for each transform defined within the specification, including those within 'composite' or 'chain' structures. Args: spec (dict): A dictionary representing a pipeline or transform specification. Yields: str: The 'type' of each transform found in the specification. """ if spec.get('type', None) in (None, 'composite', 'chain'): if 'source' in spec: yield from transform_types(spec['source']) for t in spec.get('transforms', []): yield from transform_types(t) if 'sink' in spec: yield from transform_types(spec['sink']) else: yield spec['type'] def provider_sets(spec, require_available=False): """ Generates all relevant combinations of providers for a given pipeline spec. This function analyzes a pipeline specification to identify transforms that can be implemented by multiple underlying providers (e.g., a generic transform vs. a SQL-backed one). It then yields different "provider sets," each representing a unique combination of choices for these multi-provider transforms. If no transforms have multiple available providers, it yields a single provider set using the standard defaults. Args: spec (dict): The pipeline specification, typically loaded from YAML. require_available (bool): If True, raises an error if a provider needed for a transform is not available. If False (default), unavailable providers are skipped, potentially reducing the number of yielded combinations. Yields: tuple: A tuple where the first element is a string suffix uniquely identifying the provider combination (e.g., "MyTransform_SqlProvider_0"), and the second element is a dictionary mapping transform types to a list containing the selected provider(s) for that combination. """ try: for p in spec['pipelines']: _ = yaml_transform.preprocess(copy.deepcopy(p['pipeline'])) except Exception as exn: print(exn) all_transform_types = [] else: all_transform_types = set.union( *( set( transform_types( yaml_transform.preprocess(copy.deepcopy(p['pipeline'])))) for p in spec['pipelines'])) def filter_to_available(t, providers): if t == 'LogForTesting': # Don't fan out to all the (many) possibilities for this one... return [ p for p in providers if isinstance(p, yaml_provider.InlineProvider) ][:1] elif require_available: for p in providers: if not p.available(): raise ValueError("Provider {p} required for {t} is not available.") return providers else: return [p for p in providers if p.available()] standard_providers = yaml_provider.standard_providers() multiple_providers = { t: filter_to_available(t, standard_providers[t]) for t in all_transform_types if len(filter_to_available(t, standard_providers[t])) > 1 } if not multiple_providers: yield 'only', standard_providers else: names, provider_lists = zip(*sorted(multiple_providers.items())) for ix, c in enumerate(itertools.product(*provider_lists)): yield ( '_'.join( n + '_' + type(p.underlying_provider()).__name__ for (n, p) in zip(names, c)) + f'_{ix}', dict(standard_providers, **{n: [p] for (n, p) in zip(names, c)})) def create_test_methods(spec): """Dynamically creates test methods based on a YAML specification. This function takes a YAML specification (`spec`) which defines pipelines, fixtures, and potentially options. It iterates through different combinations of "providers" (which determine how YAML transforms are implemented, e.g., using Python or SQL). For each combination of providers: 1. It constructs a unique test method name (e.g., `test_only`). 2. It defines a test method that: a. Sets up any specified fixtures, making their values available as variables. b. Mocks the standard YAML providers to use the current combination of providers for this test run. c. For each pipeline defined in the `spec`: i. Creates a `beam.Pipeline` instance with specified options. ii. Expands the YAML pipeline definition using `yaml_transform.expand_pipeline`, substituting any fixture variables. iii. Runs the Beam pipeline. The function yields tuples of (test_method_name, test_method_function), which can then be used to populate a `unittest.TestCase` class. Args: spec (dict): A dictionary parsed from a YAML test specification file. It's expected to have keys like 'fixtures' (optional) and 'pipelines'. Yields: tuple: A tuple containing: - str: The generated name for the test method (e.g., "test_only"). - function: The dynamically generated test method. """ for suffix, providers in provider_sets(spec): def test(self, providers=providers): # default arg to capture loop value vars = {} with contextlib.ExitStack() as stack: stack.enter_context( mock.patch( 'apache_beam.yaml.yaml_provider.standard_providers', lambda: providers)) for fixture in spec.get('fixtures', []): vars[fixture['name']] = stack.enter_context( python_callable.PythonCallableWithSource. load_from_fully_qualified_name(fixture['type'])( **yaml_transform.SafeLineLoader.strip_metadata( fixture.get('config', {})))) for pipeline_spec in spec['pipelines']: try: with beam.Pipeline(options=PipelineOptions( pickle_library='cloudpickle', **replace_recursive( yaml_transform.SafeLineLoader.strip_metadata( pipeline_spec.get('options', {})), vars))) as p: yaml_transform.expand_pipeline( p, replace_recursive(pipeline_spec, vars)) except ValueError as exn: # FnApiRunner currently does not support this requirement in # some xlang scenarios (e.g. Iceberg YAML pipelines). if 'beam:requirement:pardo:on_window_expiration:v1' in str(exn): self.skipTest( 'Runner does not support ' 'beam:requirement:pardo:on_window_expiration:v1') raise yield f'test_{suffix}', test def parse_test_files(filepattern): """Parses YAML test files and dynamically creates test cases. This function iterates through all files matching the given glob pattern. For each YAML file found, it: 1. Reads the file content. 2. Determines a test suite name based on the file name. 3. Calls `create_test_methods` to generate test methods from the YAML specification. 4. Dynamically creates a new TestCase class (inheriting from `unittest.TestCase`) and populates it with the generated test methods. 5. Adds this newly created TestCase class to the global scope, making it discoverable by the unittest framework. Args: filepattern (str): A glob pattern specifying the YAML test files to parse. For example, 'path/to/tests/*.yaml'. """ for path in glob.glob(filepattern): with open(path) as fin: suite_name = os.path.splitext(os.path.basename(path))[0].title().replace( '-', '') + 'Test' print(path, suite_name) methods = dict( create_test_methods( yaml.load(fin, Loader=yaml_transform.SafeLineLoader))) globals()[suite_name] = type(suite_name, (unittest.TestCase, ), methods) # Logging setups logging.getLogger().setLevel(logging.INFO) # Dynamically create test methods from the tests directory. # yaml_test_files_dir comes from conftest.py and set by pytest_configure. _test_files_dir = yaml_test_files_dir _file_pattern = os.path.join( os.path.dirname(__file__), _test_files_dir, '*.yaml') parse_test_files(_file_pattern) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/json_utils.py0000644000175100017510000003213315206113033021365 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Utilities for converting between JSON and Beam Schema'd data. For internal use, no backward compatibility guarantees. """ import json from collections.abc import Callable from typing import Any from typing import Optional import apache_beam as beam from apache_beam.portability.api import schema_pb2 from apache_beam.typehints import schemas try: import jsonschema except ImportError: pass JSON_ATOMIC_TYPES_TO_BEAM = { 'boolean': schema_pb2.BOOLEAN, 'integer': schema_pb2.INT64, 'number': schema_pb2.DOUBLE, 'string': schema_pb2.STRING, 'bytes': schema_pb2.BYTES } BEAM_ATOMIC_TYPES_TO_JSON = { schema_pb2.INT16: 'integer', schema_pb2.INT32: 'integer', schema_pb2.FLOAT: 'number', **{ v: k for k, v in JSON_ATOMIC_TYPES_TO_BEAM.items() } } def json_schema_to_beam_schema( json_schema: dict[str, Any]) -> schema_pb2.Schema: """Returns a Beam schema equivalent for the given Json schema.""" def maybe_nullable(beam_type, nullable): if nullable: beam_type.nullable = True return beam_type json_type = json_schema.get('type', None) if json_type != 'object': raise ValueError(f'Expected object type, got {json_type}.') if 'properties' not in json_schema: # Technically this is a valid (vacuous) schema, but as it's not generally # meaningful, throw an informative error instead. # (We could add a flag to allow this degenerate case.) raise ValueError('Missing properties for {json_schema}.') required = set(json_schema.get('required', [])) return schema_pb2.Schema( fields=[ schemas.schema_field( name, maybe_nullable(json_type_to_beam_type(t), name not in required), description=t.get('description') if isinstance(t, dict) else None) for (name, t) in json_schema['properties'].items() ]) def json_type_to_beam_type(json_type: dict[str, Any]) -> schema_pb2.FieldType: """Returns a Beam schema type for the given Json (schema) type.""" if not isinstance(json_type, dict) or 'type' not in json_type: raise ValueError(f'Malformed type {json_type}.') type_name = json_type['type'] if type_name in JSON_ATOMIC_TYPES_TO_BEAM: return schema_pb2.FieldType( atomic_type=JSON_ATOMIC_TYPES_TO_BEAM[type_name]) elif type_name == 'array': return schema_pb2.FieldType( array_type=schema_pb2.ArrayType( element_type=json_type_to_beam_type(json_type['items']))) elif type_name == 'object': if 'properties' in json_type: return schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=json_schema_to_beam_schema(json_type))) elif 'additionalProperties' in json_type: return schema_pb2.FieldType( map_type=schema_pb2.MapType( key_type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING), value_type=json_type_to_beam_type( json_type['additionalProperties']))) else: raise ValueError( f'Object type must have either properties or additionalProperties, ' f'got {json_type}.') else: raise ValueError(f'Unable to convert {json_type} to a Beam schema.') def beam_schema_to_json_schema( beam_schema: schema_pb2.Schema) -> dict[str, Any]: return { 'type': 'object', 'properties': { field.name: beam_type_to_json_type(field.type) for field in beam_schema.fields }, 'additionalProperties': False } def beam_type_to_json_type(beam_type: schema_pb2.FieldType) -> dict[str, Any]: type_info = beam_type.WhichOneof("type_info") if type_info == "atomic_type": if beam_type.atomic_type in BEAM_ATOMIC_TYPES_TO_JSON: return {'type': BEAM_ATOMIC_TYPES_TO_JSON[beam_type.atomic_type]} else: return {} elif type_info == "array_type": return { 'type': 'array', 'items': beam_type_to_json_type(beam_type.array_type.element_type) } elif type_info == "iterable_type": return { 'type': 'array', 'items': beam_type_to_json_type(beam_type.iterable_type.element_type) } elif type_info == "map_type": return { 'type': 'object', 'properties': { '__line__': { 'type': 'integer' }, '__uuid__': {} }, 'additionalProperties': beam_type_to_json_type( beam_type.map_type.value_type) } elif type_info == "row_type": return { 'type': 'object', 'properties': { field.name: beam_type_to_json_type(field.type) for field in beam_type.row_type.schema.fields }, 'additionalProperties': False } else: return {} def json_to_row(beam_type: schema_pb2.FieldType) -> Callable[[Any], Any]: """Returns a callable converting Json objects to Beam rows of the given type. The input to the returned callable is expected to conform to the Json schema corresponding to this Beam type. """ if beam_type.nullable: non_null_type = schema_pb2.FieldType() non_null_type.CopyFrom(beam_type) non_null_type.nullable = False non_null_converter = json_to_row(non_null_type) return lambda value: None if value is None else non_null_converter(value) type_info = beam_type.WhichOneof("type_info") if type_info == "atomic_type": return lambda value: value elif type_info == "array_type": element_converter = json_to_row(beam_type.array_type.element_type) return lambda value: [element_converter(e) for e in value] elif type_info == "iterable_type": element_converter = json_to_row(beam_type.iterable_type.element_type) return lambda value: [element_converter(e) for e in value] elif type_info == "map_type": if beam_type.map_type.key_type.atomic_type != schema_pb2.STRING: raise TypeError( f'Only strings allowd as map keys when converting from JSON, ' f'found {beam_type}') value_converter = json_to_row(beam_type.map_type.value_type) return lambda value: {k: value_converter(v) for (k, v) in value.items()} elif type_info == "row_type": field_nullable_status = { field.name: field.type.nullable for field in beam_type.row_type.schema.fields } converters = { field.name: json_to_row(field.type) for field in beam_type.row_type.schema.fields } def convert_row(value): kwargs = {} for name, convert in converters.items(): if name in value: kwargs[name] = convert(value[name]) elif field_nullable_status[name]: kwargs[name] = convert(None) else: raise KeyError(f"Missing required field: {name}") return beam.Row(**kwargs) return convert_row elif type_info == "logical_type": return lambda value: value else: raise ValueError(f"Unrecognized type_info: {type_info!r}") def json_parser( beam_schema: schema_pb2.Schema, json_schema: Optional[dict[str, Any]] = None) -> Callable[[bytes], beam.Row]: """Returns a callable converting Json strings to Beam rows of the given type. The input to the returned callable is expected to conform to the Json schema corresponding to this Beam type. """ if json_schema is None: validate_fn = None else: cls = jsonschema.validators.validator_for(json_schema) cls.check_schema(json_schema) validate_fn = _PicklableFromConstructor( lambda: jsonschema.validators.validator_for(json_schema) (json_schema).validate) to_row = json_to_row( schema_pb2.FieldType(row_type=schema_pb2.RowType(schema=beam_schema))) def parse(s: bytes): o = json.loads(s) if validate_fn is not None: validate_fn(o) return to_row(o) return parse class _PicklableFromConstructor: def __init__(self, constructor): self._constructor = constructor self._value = None def __call__(self, o): if self._value is None: self._value = self._constructor() return self._value(o) def __getstate__(self): return {'_constructor': self._constructor, '_value': None} def row_to_json(beam_type: schema_pb2.FieldType) -> Callable[[Any], Any]: """Returns a callable converting rows of the given type to Json objects.""" type_info = beam_type.WhichOneof("type_info") if type_info == "atomic_type": return lambda value: value elif type_info == "array_type": element_converter = row_to_json(beam_type.array_type.element_type) return lambda value: [element_converter(e) for e in value] elif type_info == "iterable_type": element_converter = row_to_json(beam_type.iterable_type.element_type) return lambda value: [element_converter(e) for e in value] elif type_info == "map_type": if beam_type.map_type.key_type.atomic_type != schema_pb2.STRING: raise TypeError( f'Only strings allowd as map keys when converting to JSON, ' f'found {beam_type}') value_converter = row_to_json(beam_type.map_type.value_type) return lambda value: {k: value_converter(v) for (k, v) in value.items()} elif type_info == "row_type": converters = { field.name: row_to_json(field.type) for field in beam_type.row_type.schema.fields } return lambda row: { name: converted for (name, convert) in converters.items() # To filter out nullable fields in rows if (converted := convert(getattr(row, name, None))) is not None } elif type_info == "logical_type": return lambda value: value else: raise ValueError(f"Unrecognized type_info: {type_info!r}") def json_formater( beam_schema: schema_pb2.Schema) -> Callable[[beam.Row], bytes]: """Returns a callable converting rows of the given schema to Json strings.""" convert = row_to_json( schema_pb2.FieldType(row_type=schema_pb2.RowType(schema=beam_schema))) return lambda row: json.dumps(convert(row), sort_keys=True).encode('utf-8') def _validate_compatible(weak_schema, strong_schema): if not weak_schema: return if weak_schema['type'] != strong_schema['type']: raise ValueError( f"Incompatible types: {weak_schema['type']} vs {strong_schema['type']}") if weak_schema['type'] == 'array': _validate_compatible(weak_schema['items'], strong_schema['items']) elif weak_schema['type'] == 'object': # If the weak schema allows for arbitrary keys (is a map), # the strong schema must also allow for arbitrary keys. if weak_schema.get('additionalProperties'): if not strong_schema.get('additionalProperties', True): raise ValueError('Incompatible types: map vs object') _validate_compatible( weak_schema['additionalProperties'], strong_schema['additionalProperties']) for required in strong_schema.get('required', []): if required not in weak_schema['properties']: raise ValueError(f"Missing or unknown property '{required}'") for name, spec in weak_schema.get('properties', {}).items(): if name in strong_schema['properties']: try: _validate_compatible(spec, strong_schema['properties'][name]) except Exception as exn: raise ValueError(f"Incompatible schema for '{name}'") from exn elif not strong_schema.get('additionalProperties', True): # The property is not explicitly in the strong schema, and the strong # schema does not allow for extra properties. raise ValueError( f"Prohibited property: '{name}'; " "perhaps additionalProperties: False is missing?") def row_validator(beam_schema: schema_pb2.Schema, json_schema: dict[str, Any]) -> Callable[[Any], Any]: """Returns a callable that will fail on elements not respecting json_schema. """ if not json_schema: return lambda x: None # Validate that this compiles, but avoid pickling the validator itself. _ = jsonschema.validators.validator_for(json_schema)(json_schema) _validate_compatible(beam_schema_to_json_schema(beam_schema), json_schema) validator = None convert = row_to_json( schema_pb2.FieldType(row_type=schema_pb2.RowType(schema=beam_schema))) def validate(row): nonlocal validator if validator is None: validator = jsonschema.validators.validator_for(json_schema)(json_schema) # NOTE: A row like BeamSchema_...(name='Bob', score=None, age=25) needs to # have any fields that are None to be filtered out or the validator will # fail (e.g. {'age': 25, 'name': 'Bob'}). validator.validate(convert(row)) return validate ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/json_utils_test.py0000644000175100017510000002275615206113033022436 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import unittest import apache_beam as beam from apache_beam.portability.api import schema_pb2 from apache_beam.yaml import json_utils class JsonUtilsTest(unittest.TestCase): def test_json_to_row_with_missing_optional_field(self): beam_schema = schema_pb2.Schema( fields=[ schema_pb2.Field( name='id', type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), schema_pb2.Field( name='event_subtype', type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING, nullable=True)), ]) beam_type = schema_pb2.FieldType( row_type=schema_pb2.RowType(schema=beam_schema)) converter = json_utils.json_to_row(beam_type) json_data = {'id': '123'} beam_row = converter(json_data) self.assertEqual(beam_row, beam.Row(id='123', event_subtype=None)) def test_json_to_row_with_missing_optional_object(self): nested_schema = schema_pb2.Schema( fields=[ schema_pb2.Field( name='nested_id', type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), ]) beam_schema = schema_pb2.Schema( fields=[ schema_pb2.Field( name='id', type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), schema_pb2.Field( name='nested', type=schema_pb2.FieldType( row_type=schema_pb2.RowType(schema=nested_schema), nullable=True)), ]) beam_type = schema_pb2.FieldType( row_type=schema_pb2.RowType(schema=beam_schema)) converter = json_utils.json_to_row(beam_type) json_data = {'id': '123'} beam_row = converter(json_data) self.assertEqual(beam_row, beam.Row(id='123', nested=None)) def test_json_to_row_with_missing_optional_array(self): beam_schema = schema_pb2.Schema( fields=[ schema_pb2.Field( name='id', type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), schema_pb2.Field( name='items', type=schema_pb2.FieldType( array_type=schema_pb2.ArrayType( element_type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING)), nullable=True)), ]) beam_type = schema_pb2.FieldType( row_type=schema_pb2.RowType(schema=beam_schema)) converter = json_utils.json_to_row(beam_type) json_data = {'id': '123'} beam_row = converter(json_data) self.assertEqual(beam_row, beam.Row(id='123', items=None)) def test_json_to_row_with_all_fields(self): nested_schema = schema_pb2.Schema( fields=[ schema_pb2.Field( name='nested_id', type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), ]) beam_schema = schema_pb2.Schema( fields=[ schema_pb2.Field( name='id', type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), schema_pb2.Field( name='event_subtype', type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING, nullable=True)), schema_pb2.Field( name='nested', type=schema_pb2.FieldType( row_type=schema_pb2.RowType(schema=nested_schema), nullable=True)), schema_pb2.Field( name='items', type=schema_pb2.FieldType( array_type=schema_pb2.ArrayType( element_type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING)), nullable=True)), ]) beam_type = schema_pb2.FieldType( row_type=schema_pb2.RowType(schema=beam_schema)) converter = json_utils.json_to_row(beam_type) json_data = { 'id': '123', 'event_subtype': 'subtype_val', 'nested': { 'nested_id': 'nested_123' }, 'items': ['a', 'b', 'c'] } beam_row = converter(json_data) expected_row = beam.Row( id='123', event_subtype='subtype_val', nested=beam.Row(nested_id='nested_123'), items=['a', 'b', 'c']) self.assertEqual(beam_row, expected_row) def test_json_to_row_with_missing_required_field(self): beam_schema = schema_pb2.Schema( fields=[ schema_pb2.Field( name='id', type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), schema_pb2.Field( name='required_field', type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), ]) beam_type = schema_pb2.FieldType( row_type=schema_pb2.RowType(schema=beam_schema)) converter = json_utils.json_to_row(beam_type) json_data = {'id': '123'} with self.assertRaises(KeyError): converter(json_data) def test_validate_compatible(self): from apache_beam.yaml.json_utils import _validate_compatible # Compatible cases _validate_compatible({'type': 'string'}, {'type': 'string'}) _validate_compatible( { 'type': 'object', 'properties': { 'f': { 'type': 'string' } } }, { 'type': 'object', 'properties': { 'f': { 'type': 'string' } } }) # Incompatible types with self.assertRaisesRegex(ValueError, 'Incompatible types'): _validate_compatible({'type': 'string'}, {'type': 'integer'}) # Missing property with self.assertRaisesRegex(ValueError, 'Missing or unknown property'): _validate_compatible({ 'type': 'object', 'properties': {} }, { 'type': 'object', 'properties': { 'f': { 'type': 'string' } }, 'required': ['f'] }) # Incompatible property type with self.assertRaisesRegex(ValueError, 'Incompatible schema for \'f\''): _validate_compatible( { 'type': 'object', 'properties': { 'f': { 'type': 'integer' } } }, { 'type': 'object', 'properties': { 'f': { 'type': 'string' } } }) def test_validate_compatible_map(self): from apache_beam.yaml.json_utils import _validate_compatible # Compatible maps _validate_compatible( { 'type': 'object', 'additionalProperties': { 'type': 'string' } }, { 'type': 'object', 'additionalProperties': { 'type': 'string' } }) # Incompatible map values with self.assertRaisesRegex(ValueError, 'Incompatible types'): _validate_compatible( { 'type': 'object', 'additionalProperties': { 'type': 'string' } }, { 'type': 'object', 'additionalProperties': { 'type': 'integer' } }) # Map vs Object with self.assertRaisesRegex(ValueError, 'Incompatible types: map vs object'): _validate_compatible( { 'type': 'object', 'additionalProperties': { 'type': 'string' } }, { 'type': 'object', 'properties': {}, 'additionalProperties': False }) def test_validate_compatible_extra_properties(self): from apache_beam.yaml.json_utils import _validate_compatible # Extra properties in weak_schema should be allowed if strong_schema # doesn't explicitly forbid them (default additionalProperties=True). _validate_compatible({ 'type': 'object', 'properties': { 'extra': { 'type': 'string' } } }, { 'type': 'object', 'properties': {} }) # But if strong_schema says additionalProperties: False, it should raise. with self.assertRaisesRegex(ValueError, 'Prohibited property'): _validate_compatible( { 'type': 'object', 'properties': { 'extra': { 'type': 'string' } } }, { 'type': 'object', 'properties': {}, 'additionalProperties': False }) if __name__ == '__main__': unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/main.py0000644000175100017510000002517715206113033020132 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import argparse import json import os import sys import unittest import yaml import apache_beam as beam from apache_beam.io.filesystems import FileSystems # The following imports force the registration of JDBC logical types. # When running a Beam YAML pipeline, the expansion service handles JDBCIO using # Java transforms, bypassing the Python module (`apache_beam.io.jdbc`) that # registers these types. These imports load the module, preventing a # "logical type not found" error. from apache_beam.io.jdbc import JdbcDateType # pylint: disable=unused-import from apache_beam.io.jdbc import JdbcTimeType # pylint: disable=unused-import from apache_beam.transforms import resources from apache_beam.yaml import yaml_testing from apache_beam.yaml import yaml_transform from apache_beam.yaml import yaml_utils def _preparse_jinja_flags(argv): """Promotes any flags to --jinja_variables based on --jinja_variable_flags. This is to facilitate tools (such as dataflow templates) that must pass options as un-nested flags. """ parser = argparse.ArgumentParser(allow_abbrev=False) parser.add_argument( '--jinja_variable_flags', default=[], type=lambda s: s.split(','), help='A list of flag names that should be used as jinja variables.') parser.add_argument( '--jinja_variables', default={}, type=json.loads, help='A json dict of variables used when invoking the jinja preprocessor ' 'on the provided yaml pipeline.') jinja_args, other_args = parser.parse_known_args(argv) if not jinja_args.jinja_variable_flags: return argv jinja_variable_parser = argparse.ArgumentParser(allow_abbrev=False) for flag_name in jinja_args.jinja_variable_flags: jinja_variable_parser.add_argument('--' + flag_name) jinja_flag_variables, pipeline_args = jinja_variable_parser.parse_known_args( other_args) jinja_args.jinja_variables.update( ** {k: v for (k, v) in vars(jinja_flag_variables).items() if v is not None}) if jinja_args.jinja_variables: pipeline_args = pipeline_args + [ '--jinja_variables=' + json.dumps(jinja_args.jinja_variables) ] return pipeline_args def _parse_arguments(argv): parser = argparse.ArgumentParser(allow_abbrev=False) parser.add_argument( '--yaml_pipeline', '--pipeline_spec', help='A yaml description of the pipeline to run.') parser.add_argument( '--yaml_pipeline_file', '--pipeline_spec_file', help='A file containing a yaml description of the pipeline to run.') parser.add_argument( '--json_schema_validation', default='generic', help='none: do no pipeline validation against the schema; ' 'generic: validate the pipeline shape, but not individual transforms; ' 'per_transform: also validate the config of known transforms') parser.add_argument( '--jinja_variables', default=None, type=json.loads, help='A json dict of variables used when invoking the jinja preprocessor ' 'on the provided yaml pipeline.') parser.add_argument( '--tests', '--test', action=argparse.BooleanOptionalAction, help='Run the tests associated with the given pipeline, rather than the ' 'pipeline itself.') parser.add_argument( '--fix_tests', action=argparse.BooleanOptionalAction, help='Update failing test expectations to match the actual ouput. ' 'Requires --test_suite if the pipeline uses jinja formatting.') parser.add_argument( '--create_test', action=argparse.BooleanOptionalAction, help='Automatically creates a regression test for the given pipeline, ' 'adding it to the pipeline spec or test suite dependon on whether ' '--test_suite is given. ' 'Requires --test_suite if the pipeline uses jinja formatting.') parser.add_argument( '--test_suite', help='Run the given tests against the given pipeline, rather than the ' 'pipeline itself. ' 'Should be a file containing a list of yaml test specifications.') return parser.parse_known_args(argv) def _pipeline_spec_from_args(known_args): if known_args.yaml_pipeline_file and known_args.yaml_pipeline: raise ValueError( "Exactly one of yaml_pipeline or yaml_pipeline_file must be set.") elif known_args.yaml_pipeline_file: with FileSystems.open(known_args.yaml_pipeline_file) as fin: pipeline_yaml = fin.read().decode() elif known_args.yaml_pipeline: pipeline_yaml = known_args.yaml_pipeline else: raise ValueError( "Exactly one of yaml_pipeline or yaml_pipeline_file must be set.") return pipeline_yaml def run(argv=None): options, constructor, display_data = build_pipeline_components_from_argv(argv) with beam.Pipeline(options=options, display_data=display_data) as p: print('Building pipeline...') constructor(p) print('Running pipeline...') def run_tests(argv=None, exit=True): known_args, pipeline_args, _, pipeline_yaml = _build_pipeline_yaml_from_argv( argv) pipeline_spec = yaml.load(pipeline_yaml, Loader=yaml_transform.SafeLineLoader) options = _build_pipeline_options(pipeline_spec, pipeline_args) if known_args.create_test and known_args.fix_tests: raise ValueError( 'At most one of --create_test and --fix_tests may be specified.') elif known_args.create_test: result = unittest.TestResult() tests = [] else: if known_args.test_suite: with open(known_args.test_suite) as fin: test_suite_holder = yaml.load( fin, Loader=yaml_transform.SafeLineLoader) or {} else: test_suite_holder = pipeline_spec test_specs = test_suite_holder.get('tests', []) if not isinstance(test_specs, list): raise TypeError('tests attribute must be a list of test specifications.') elif not test_specs: raise RuntimeError( 'No tests found. ' "If you haven't added a set of tests yet, you can get started by " 'running your pipeline with the --create_test flag enabled.') tests = [ yaml_testing.YamlTestCase( pipeline_spec, test_spec, options, known_args.fix_tests) for test_spec in test_specs ] suite = unittest.TestSuite(tests) result = unittest.TextTestRunner().run(suite) if known_args.fix_tests or known_args.create_test: update_tests(known_args, pipeline_yaml, pipeline_spec, options, tests) if exit: # emulates unittest.main() sys.exit(0 if result.wasSuccessful() else 1) else: if not result.wasSuccessful(): raise RuntimeError(result) def update_tests(known_args, pipeline_yaml, pipeline_spec, options, tests): if known_args.test_suite: path = known_args.test_suite if not os.path.exists(path) and known_args.create_test: with open(path, 'w') as fout: fout.write('tests: []') elif known_args.yaml_pipeline_file: path = known_args.yaml_pipeline_file else: raise RuntimeError( 'Test fixing only supported for file-backed tests. ' 'Please use the --test_suite flag.') with open(path) as fin: original_yaml = fin.read() if path == known_args.yaml_pipeline_file and pipeline_yaml.strip( ) != original_yaml.strip(): raise RuntimeError( 'In-file test fixing not yet supported for templated pipelines. ' 'Please use the --test_suite flag.') updated_spec = yaml.load(original_yaml, Loader=yaml.SafeLoader) or {} if known_args.fix_tests: updated_spec['tests'] = [test.fixed_test() for test in tests] if known_args.create_test: if 'tests' not in updated_spec: updated_spec['tests'] = [] updated_spec['tests'].append( yaml_testing.create_test(pipeline_spec, options)) updated_yaml = yaml_utils.patch_yaml(original_yaml, updated_spec) with open(path, 'w') as fout: fout.write(updated_yaml) def _build_pipeline_yaml_from_argv(argv): argv = _preparse_jinja_flags(argv) known_args, pipeline_args = _parse_arguments(argv) pipeline_template = _pipeline_spec_from_args(known_args) pipeline_yaml = yaml_transform.expand_jinja( pipeline_template, known_args.jinja_variables or {}) return known_args, pipeline_args, pipeline_template, pipeline_yaml def _build_pipeline_options(pipeline_spec, pipeline_args): return beam.options.pipeline_options.PipelineOptions( pipeline_args, pickle_library='cloudpickle', **yaml_transform.SafeLineLoader.strip_metadata( pipeline_spec.get('options', {}))) def build_pipeline_components_from_argv(argv): (known_args, pipeline_args, pipeline_template, pipeline_yaml) = _build_pipeline_yaml_from_argv(argv) display_data = { 'yaml': pipeline_yaml, 'yaml_jinja_template': pipeline_template, 'yaml_jinja_variables': json.dumps(known_args.jinja_variables), } options, constructor = build_pipeline_components_from_yaml( pipeline_yaml, pipeline_args, known_args.json_schema_validation, known_args.yaml_pipeline_file, ) return options, constructor, display_data def build_pipeline_components_from_yaml( pipeline_yaml, pipeline_args, validate_schema='generic', pipeline_path=''): pipeline_spec = yaml.load(pipeline_yaml, Loader=yaml_transform.SafeLineLoader) options = _build_pipeline_options(pipeline_spec, pipeline_args) def constructor(root): if 'resource_hints' in pipeline_spec.get('pipeline', {}): # Add the declared resource hints to the "root" spec. root._current_transform().resource_hints.update( resources.parse_resource_hints( yaml_transform.SafeLineLoader.strip_metadata( pipeline_spec['pipeline']['resource_hints']))) yaml_transform.expand_pipeline( root, pipeline_spec, validate_schema=validate_schema, pipeline_path=pipeline_path, ) return options, constructor if __name__ == '__main__': import logging logging.getLogger().setLevel(logging.INFO) if '--tests' in sys.argv or '--test' in sys.argv: run_tests() else: run() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/main_test.py0000644000175100017510000001603115206113033021156 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import datetime import glob import logging import os import tempfile import unittest from apache_beam.yaml import main try: import jsonschema except ImportError: jsonschema = None TEST_PIPELINE = ''' pipeline: type: chain transforms: - type: Create config: elements: [ELEMENT] # Writing to an actual file here rather than just using AssertThat # because this is an integration test and above all we want to ensure # the pipeline actually runs (and asserts may not fail if there's a # bug in the invocation logic). - type: WriteToText config: path: PATH tests: - name: InlineTest mock_outputs: - name: Create elements: ['a', 'b', 'c'] expected_inputs: - name: WriteToText elements: - {element: a} - {element: b} - {element: c} ''' PASSING_TEST_SUITE = ''' tests: - name: ExternalTest # comment mock_outputs: - name: Create elements: ['a', 'b', 'c'] expected_inputs: - name: WriteToText elements: - element: a - element: b - element: c ''' FAILING_TEST_SUITE = ''' tests: - name: ExternalTest # comment mock_outputs: - name: Create elements: ['a', 'b', 'c'] expected_inputs: - name: WriteToText elements: - element: x ''' @unittest.skipIf(jsonschema is None, "Yaml dependencies not installed") class MainTest(unittest.TestCase): def test_pipeline_spec_from_file(self): with tempfile.TemporaryDirectory() as tmpdir: yaml_path = os.path.join(tmpdir, 'test.yaml') out_path = os.path.join(tmpdir, 'out.txt') with open(yaml_path, 'wt') as fout: fout.write(TEST_PIPELINE.replace('PATH', out_path)) main.run(['--yaml_pipeline_file', yaml_path]) with open(glob.glob(out_path + '*')[0], 'rt') as fin: self.assertEqual(fin.read().strip(), 'ELEMENT') def test_pipeline_spec_from_flag(self): with tempfile.TemporaryDirectory() as tmpdir: out_path = os.path.join(tmpdir, 'out.txt') main.run(['--yaml_pipeline', TEST_PIPELINE.replace('PATH', out_path)]) with open(glob.glob(out_path + '*')[0], 'rt') as fin: self.assertEqual(fin.read().strip(), 'ELEMENT') def test_jinja_variables(self): with tempfile.TemporaryDirectory() as tmpdir: out_path = os.path.join(tmpdir, 'out.txt') main.run([ '--yaml_pipeline', TEST_PIPELINE.replace('PATH', out_path).replace('ELEMENT', '{{var}}'), '--jinja_variables', '{"var": "my_line"}' ]) with open(glob.glob(out_path + '*')[0], 'rt') as fin: self.assertEqual(fin.read().strip(), 'my_line') def test_jinja_variable_flags(self): with tempfile.TemporaryDirectory() as tmpdir: out_path = os.path.join(tmpdir, 'out.txt') main.run([ '--yaml_pipeline', TEST_PIPELINE.replace('PATH', out_path).replace('ELEMENT', '{{var}}'), '--jinja_variable_flags=var', '--var=my_line' ]) with open(glob.glob(out_path + '*')[0], 'rt') as fin: self.assertEqual(fin.read().strip(), 'my_line') def test_preparse_jinja_flags(self): argv = [ '--jinja_variables={"from_vars": 1, "from_both": 2}', '--jinja_variable_flags=from_both,from_flag,from_missing_flag', '--from_both=30', '--from_flag=40', '--another_arg=foo', 'pos_arg', ] self.assertCountEqual( main._preparse_jinja_flags(argv), [ '--jinja_variables=' '{"from_vars": 1, "from_both": "30", "from_flag": "40"}', '--another_arg=foo', 'pos_arg', ]) def test_jinja_datetime(self): with tempfile.TemporaryDirectory() as tmpdir: out_path = os.path.join(tmpdir, 'out.txt') main.run([ '--yaml_pipeline', TEST_PIPELINE.replace('PATH', out_path).replace( 'ELEMENT', '"{{datetime.datetime.now().strftime("%Y-%m-%d")}}"'), ]) with open(glob.glob(out_path + '*')[0], 'rt') as fin: self.assertEqual( fin.read().strip(), datetime.datetime.now().strftime("%Y-%m-%d")) def test_inline_test_specs(self): main.run_tests(['--yaml_pipeline', TEST_PIPELINE, '--test'], exit=False) def test_external_test_specs(self): with tempfile.TemporaryDirectory() as tmpdir: good_suite = os.path.join(tmpdir, 'good.yaml') with open(good_suite, 'w') as fout: fout.write(PASSING_TEST_SUITE) bad_suite = os.path.join(tmpdir, 'bad.yaml') with open(bad_suite, 'w') as fout: fout.write(FAILING_TEST_SUITE) # Must pass. main.run_tests([ '--yaml_pipeline', TEST_PIPELINE, '--test_suite', good_suite, ], exit=False) # Must fail. (Ensures testing is not a no-op.) with self.assertRaisesRegex(Exception, 'errors=1 failures=0'): main.run_tests([ '--yaml_pipeline', TEST_PIPELINE, '--test_suite', bad_suite, ], exit=False) def test_fix_suite(self): with tempfile.TemporaryDirectory() as tmpdir: test_suite = os.path.join(tmpdir, 'tests.yaml') with open(test_suite, 'w') as fout: fout.write(FAILING_TEST_SUITE) main.run_tests([ '--yaml_pipeline', TEST_PIPELINE, '--test_suite', test_suite, '--fix_tests' ], exit=False) with open(test_suite) as fin: self.assertEqual(fin.read(), PASSING_TEST_SUITE) def test_create_test(self): with tempfile.TemporaryDirectory() as tmpdir: test_suite = os.path.join(tmpdir, 'tests.yaml') with open(test_suite, 'w') as fout: fout.write('') main.run_tests([ '--yaml_pipeline', TEST_PIPELINE.replace('ELEMENT', 'x'), '--test_suite', test_suite, '--create_test' ], exit=False) with open(test_suite) as fin: self.assertEqual( fin.read(), ''' tests: - mock_outputs: [] expected_inputs: - name: WriteToText elements: - element: x '''.lstrip()) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/options.py0000644000175100017510000000272215206113033020670 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from apache_beam.options import pipeline_options class YamlOptions(pipeline_options.PipelineOptions): @classmethod def _add_argparse_args(cls, parser): parser.add_argument( '--yaml_experimental_features', dest='yaml_experimental_features', action='append', default=[], help=('Enable yaml features ahead of them being declared stable.')) @classmethod def check_enabled(cls, pipeline, feature, description=None): if feature not in pipeline._options.view_as(cls).yaml_experimental_features: raise ValueError( f'{description or feature} unsupported because ' f'{feature} is not set in --yaml_experimental_features option.') ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/pipeline.schema.yaml0000644000175100017510000001545015206113033022555 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # $schema: 'http://json-schema.org/draft-07/schema#' $id: https://github.com/apache/beam/tree/master/sdks/python/apache_beam/yaml/pipeline.schema.yaml $defs: transformBase: type: object properties: type: { type: string } name: { type: string } windowing: $ref: '#/$defs/windowing' resource_hints: $ref: '#/$defs/resource_hints' input: oneOf: - type: string - type: array items: { type: string } - type: object properties: { __line__: {}} additionalProperties: type: string output: oneOf: - type: string - type: object properties: { __line__: {}} additionalProperties: type: string output_schema: type: object additionalProperties: true required: - type nestedTransform: allOf: - type: object properties: type: {} name: {} input: {} output: {} windowing: {} resource_hints: {} transforms: type: array items: $ref: '#/$defs/transform' extra_transforms: type: array items: $ref: '#/$defs/transform' source: $ref: '#/$defs/transform' sink: $ref: '#/$defs/transform' __line__: {} __uuid__: {} additionalProperties: false - anyOf: - required: - transforms - required: - source - sink implicitInputOutputs: properties: transforms: type: array items: properties: # Must be implicit. input: { not: {} } output: { not: {} } source: properties: # Must be implicit. input: { not: {} } output: { not: {} } sink: properties: # Must be implicit. input: { not: {} } output: { not: {} } transform: allOf: - $ref: '#/$defs/transformBase' - if: properties: { type: { const: composite }} then: allOf: - $ref: '#/$defs/nestedTransform' - properties: extra_transforms: { not: {} } - if: properties: { type: { const: chain }} then: allOf: - $ref: '#/$defs/nestedTransform' - $ref: '#/$defs/implicitInputOutputs' - if: not: anyOf: - properties: { type: { const: composite }} - properties: { type: { const: chain }} then: properties: type: {} name: {} input: {} output: {} output_schema: { type: object } windowing: {} resource_hints: {} config: { type: object } __line__: {} __uuid__: {} additionalProperties: false windowing: {} # TODO resource_hints: type: object provider: # TODO(robertwb): Consider enumerating the provider types along with # the arguments they accept/expect (possibly in a separate schema file). type: object properties: type: { type: string } transforms: type: object properties: { __line__: {}} additionalProperties: type: string config: { type: object } __line__: {} __uuid__: {} additionalProperties: false required: - type - transforms - config providerInclude: # TODO(robertwb): Consider enumerating the provider types along with # the arguments they accept/expect (possibly in a separate schema file). type: object properties: include: { type: string } __line__: {} __uuid__: {} additionalProperties: false required: - include providerOrProviderInclude: if: allOf: [ { properties: { include: { type: string }}}, { required: [ "include" ] } ] then: $ref: '#/$defs/providerInclude' else: $ref: '#/$defs/provider' testElementsSpec: type: object properties: name: { type: string } elements: type: array item: type: object __line__: {} __uuid__: {} additionalProperties: false testSpec: type: object properties: name: { type: string } mock_inputs: type: array items: $ref: '#/$defs/testElementsSpec' mock_outputs: type: array items: $ref: '#/$defs/testElementsSpec' expected_inputs: type: array items: $ref: '#/$defs/testElementsSpec' expected_outputs: type: array items: $ref: '#/$defs/testElementsSpec' allowed_sources: type: array items: { type: string } __line__: {} __uuid__: {} additionalProperties: false type: object properties: pipeline: allOf: # These are the only top-level properties defined in pipeline. - type: object properties: type: oneOf: - { const: composite } - { const: chain } windowing: $ref: '#/$defs/windowing' resource_hints: $ref: '#/$defs/resource_hints' transforms: {} extra_transforms: {} sink: {} source: {} __line__: {} __uuid__: {} additionalProperties: false # This defines the allowable contents of the attributes above. - $ref: '#/$defs/nestedTransform' # A chain-type transform, like a chain composite, must have implicit io. - if: properties: { type: { const: chain }} required: [type] then: $ref: '#/$defs/implicitInputOutputs' providers: type: array items: $ref: '#/$defs/providerOrProviderInclude' options: type: object tests: type: array items: $ref: '#/$defs/testSpec' required: - pipeline ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/programming_guide_test.py0000644000175100017510000003245015206113033023734 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest import apache_beam as beam from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.yaml.readme_test import create_test_method from apache_beam.yaml.yaml_transform import YamlTransform DATA = [ beam.Row(animal='cat', weight=1), beam.Row(animal='cat', weight=5), beam.Row(animal='dog', weight=10), ] class ProgrammingGuideTest(unittest.TestCase): test_pipelines_constructing_reading = create_test_method( 'RUN', 'test_pipelines_constructing_reading', ''' # [START pipelines_constructing_reading] pipeline: source: type: ReadFromText config: path: ... # [END pipelines_constructing_reading] transforms: [] ''') test_create_pcollection = create_test_method( 'RUN', 'test_create_pcollection', ''' # [START create_pcollection] pipeline: transforms: - type: Create config: elements: - A - B - ... # [END create_pcollection] transforms: [] ''') def test_group_by(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' # [START group_by] type: Combine config: group_by: animal combine: weight: group # [END group_by] ''') assert_that( result | beam.Map(lambda x: beam.Row(**x._asdict())), equal_to([ beam.Row(animal='cat', weight=[1, 5]), beam.Row(animal='dog', weight=[10]), ])) def test_co_group_by(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: composite transforms: # [START cogroupbykey_inputs] - type: Create name: CreateEmails config: elements: - { name: "amy", email: "amy@example.com" } - { name: "carl", email: "carl@example.com" } - { name: "julia", email: "julia@example.com" } - { name: "carl", email: "carl@email.com" } - type: Create name: CreatePhones config: elements: - { name: "amy", phone: "111-222-3333" } - { name: "james", phone: "222-333-4444" } - { name: "amy", phone: "333-444-5555" } - { name: "carl", phone: "444-555-6666" } # [END cogroupbykey_inputs] # [START cogroupbykey] - type: MapToFields name: PrepareEmails input: CreateEmails config: language: python fields: name: name email: "[email]" phone: "[]" - type: MapToFields name: PreparePhones input: CreatePhones config: language: python fields: name: name email: "[]" phone: "[phone]" - type: Combine name: CoGropuBy input: [PrepareEmails, PreparePhones] config: group_by: [name] combine: email: concat phone: concat - type: MapToFields name: FormatResults input: CoGropuBy config: language: python fields: formatted: "'%s; %s; %s' % (name, sorted(email), sorted(phone))" # [END cogroupbykey] output: FormatResults ''') assert_that( result | beam.Map(lambda x: x.formatted), equal_to([ # [START cogroupbykey_formatted_outputs] "amy; ['amy@example.com']; ['111-222-3333', '333-444-5555']", "carl; ['carl@email.com', 'carl@example.com']; ['444-555-6666']", "james; []; ['222-333-4444']", "julia; ['julia@example.com']; []", # [END cogroupbykey_formatted_outputs] ])) def test_combine_ref(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' # [START combine_simple_py] type: Combine config: language: python group_by: animal combine: biggest: fn: type: 'apache_beam.transforms.combiners.TopCombineFn' config: n: 2 value: weight # [END combine_simple_py] ''') assert_that( result | beam.Map(lambda x: beam.Row(**x._asdict())), equal_to([ beam.Row(animal='cat', biggest=[5, 1]), beam.Row(animal='dog', biggest=[10]), ])) def test_combine_globally(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' # [START combine_global_sum] type: Combine config: group_by: [] combine: weight: sum # [END combine_global_sum] ''') assert_that( result | beam.Map(lambda x: beam.Row(**x._asdict())), equal_to([ beam.Row(weight=16), ])) def test_combine_per_key(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' # [START combine_per_key] type: Combine config: group_by: [animal] combine: total_weight: fn: sum value: weight average_weight: fn: mean value: weight # [END combine_per_key] ''') assert_that( result | beam.Map(lambda x: beam.Row(**x._asdict())), equal_to([ beam.Row(animal='cat', total_weight=6, average_weight=3), beam.Row(animal='dog', total_weight=10, average_weight=10), ])) def test_flatten(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: composite transforms: - type: Create name: SomeProducingTransform config: elements: [a, b, c] - type: Create name: AnotherProducingTransform config: elements: [d, e, f] # [START model_multiple_pcollections_flatten] - type: Flatten input: [SomeProducingTransform, AnotherProducingTransform] # [END model_multiple_pcollections_flatten] output: Flatten ''') assert_that( result | beam.Map(lambda x: x.element), equal_to(['a', 'b', 'c', 'd', 'e', 'f'])) def test_expode(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: composite transforms: - type: Create name: SomeProducingTransform config: elements: - {line: "a aa"} - {line: "b"} # [START model_multiple_output_dofn] - type: MapToFields input: SomeProducingTransform config: language: python fields: word: "line.split()" - type: Explode input: MapToFields config: fields: word # [END model_multiple_output_dofn] output: Explode ''') assert_that( result | beam.Map(lambda x: x.word), equal_to(['a', 'aa', 'b'])) def test_schema_output_type(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' # [START schema_output_type] type: MapToFields config: language: python fields: new_field: expression: "hex(weight)" output_type: { "type": "string" } # [END schema_output_type] ''') assert_that( result | beam.Map(lambda x: x.new_field), equal_to(['0x1', '0x5', '0xa'])) def test_fixed_windows(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) _ = elements | YamlTransform( ''' # [START setting_fixed_windows] type: WindowInto windowing: type: fixed size: 60s # [END setting_fixed_windows] ''') def test_sliding_windows(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) _ = elements | YamlTransform( ''' # [START setting_sliding_windows] type: WindowInto windowing: type: sliding size: 5m period: 30s # [END setting_sliding_windows] ''') def test_session_windows(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) _ = elements | YamlTransform( ''' # [START setting_session_windows] type: WindowInto windowing: type: sessions gap: 60s # [END setting_session_windows] ''') def test_global_windows(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) _ = elements | YamlTransform( ''' # [START setting_global_window] type: WindowInto windowing: type: global # [END setting_global_window] ''') def test_assign_timestamps(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([ beam.Row(external_timestamp_field='2024-01-01', value=1), beam.Row(external_timestamp_field='2024-01-01', value=2), beam.Row(external_timestamp_field='2024-01-02', value=10), ]) _ = elements | YamlTransform( ''' # [START setting_timestamp] type: AssignTimestamps config: language: python timestamp: callable: | import datetime def extract_timestamp(x): raw = datetime.datetime.strptime( x.external_timestamp_field, "%Y-%m-%d") return raw.astimezone(datetime.timezone.utc) # [END setting_timestamp] ''') def test_partition(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([ beam.Row(percentile=1), beam.Row(percentile=20), beam.Row(percentile=90), ]) _ = elements | YamlTransform( ''' # [START model_multiple_pcollections_partition] type: Partition config: by: str(percentile // 10) language: python outputs: ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10"] # [END model_multiple_pcollections_partition] ''') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/readme_test.py0000644000175100017510000003204215206113033021467 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Runs the examples from the README.md file.""" import argparse import logging import os import random import re import sys import tempfile import unittest import mock import yaml from yaml.loader import SafeLoader import apache_beam as beam from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.typehints import trivial_inference from apache_beam.yaml import yaml_provider from apache_beam.yaml import yaml_testing from apache_beam.yaml import yaml_transform from apache_beam.yaml import yaml_utils class FakeSql(beam.PTransform): def __init__(self, query): self.query = query def default_label(self): return 'Sql' def expand(self, inputs): if isinstance(inputs, beam.PCollection): inputs = {'PCOLLECTION': inputs} # This only handles the most basic of queries, trying to infer the output # schema... m = re.match('select (.*?) from', self.query, flags=re.IGNORECASE) if not m: raise ValueError(self.query) def guess_name_and_type(expr): expr = expr.strip().replace('`', '') if expr.endswith('*'): return 'unknown', str parts = expr.split() if len(parts) >= 2 and parts[-2].lower() == 'as': name = parts[-1] elif re.match(r'[\w.]+', parts[0]): name = parts[0].split('.')[-1] else: name = f'expr{hash(expr)}' if '(' in expr: expr = expr.lower() if expr.startswith('count'): typ = int elif expr.startswith('avg'): typ = float else: typ = str elif '+' in expr: typ = float else: part = parts[0] if '.' in part: table, field = part.split('.') typ = inputs[table].element_type.get_type_for(field) else: typ = next(iter(inputs.values())).element_type.get_type_for(name) # Handle optionals more gracefully. if (str(typ).startswith('typing.Union[') or str(typ).startswith('typing.Optional[') or '|' in str(typ)): if len(typ.__args__) == 2 and type(None) in typ.__args__: typ, = [t for t in typ.__args__ if t is not type(None)] return name, typ if m.group(1) == '*': return next(iter(inputs.values())) | beam.Filter(lambda _: True) else: output_schema = [ guess_name_and_type(expr) for expr in m.group(1).split(',') ] output_element = beam.Row(**{name: typ() for name, typ in output_schema}) return next(iter(inputs.values())) | beam.Map( lambda _: output_element).with_output_types( trivial_inference.instance_to_type(output_element)) class FakeReadFromPubSub(beam.PTransform): def __init__(self, topic, format, schema): pass def expand(self, p): data = p | beam.Create([beam.Row(col1='a', col2=1, col3=0.5)]) result = data | beam.Map( lambda row: beam.transforms.window.TimestampedValue(row, 0)) # TODO(robertwb): Allow this to be inferred. result.element_type = data.element_type return result class FakeWriteToPubSub(beam.PTransform): def __init__(self, topic, format): pass def expand(self, pcoll): return pcoll class FakeAggregation(beam.PTransform): def __init__(self, **unused_kwargs): pass def expand(self, pcoll): return pcoll | beam.GroupBy(lambda _: 'key').aggregate_field( lambda _: 1, sum, 'count') class _Fakes: fn = str class SomeTransform(beam.PTransform): def __init__(*args, **kwargs): pass def expand(self, pcoll): return pcoll RENDER_DIR = None TEST_TRANSFORMS = { 'Sql': FakeSql, 'ReadFromPubSub': FakeReadFromPubSub, 'WriteToPubSub': FakeWriteToPubSub, 'SomeGroupingTransform': FakeAggregation, 'SomeTransform': _Fakes.SomeTransform, 'AnotherTransform': _Fakes.SomeTransform, } class TestProvider(yaml_provider.InlineProvider): def _affinity(self, other): # Always try to choose this one. return float('inf') class TestEnvironment: def __enter__(self): self.tempdir = tempfile.TemporaryDirectory() return self def input_file(self, name, content): path = os.path.join(self.tempdir.name, name) with open(path, 'w') as fout: fout.write(content) return path def input_csv(self): return self.input_file('input.csv', 'col1,col2,col3\na,1,2.5\n') def input_tsv(self): return self.input_file('input.tsv', 'col1\tcol2\tcol3\nabc\t1\t2.5\n') def input_json(self): return self.input_file( 'input.json', '{"col1": "abc", "col2": 1, "col3": 2.5"}\n') def output_file(self): return os.path.join( self.tempdir.name, str(random.randint(0, 1000)) + '.out') def udf_file(self, name): if name == 'my_mapping': lines = '\n'.join(['def my_mapping(row):', '\treturn "good"']) else: lines = '\n'.join(['def my_filter(row):', '\treturn True']) return self.input_file('udf.py', lines) def __exit__(self, *args): self.tempdir.cleanup() def replace_recursive(spec, transform_type, arg_name, arg_value): if isinstance(spec, dict): spec = { key: replace_recursive(value, transform_type, arg_name, arg_value) for (key, value) in spec.items() } if spec.get('type', None) == transform_type: spec['config'][arg_name] = arg_value return spec elif isinstance(spec, list): return [ replace_recursive(value, transform_type, arg_name, arg_value) for value in spec ] else: return spec def create_test_method(test_type, test_name, test_yaml): test_yaml = test_yaml.replace( 'apache_beam.pkg.module.', 'apache_beam.yaml.readme_test._Fakes.') test_yaml = test_yaml.replace( 'pkg.module.', 'apache_beam.yaml.readme_test._Fakes.') def test(self): with TestEnvironment() as env: nonlocal test_yaml test_yaml = test_yaml.replace('/path/to/*.tsv', env.input_tsv()) if 'MapToFields' in test_yaml or 'Filter' in test_yaml: if 'my_mapping' in test_yaml: test_yaml = test_yaml.replace( '/path/to/some/udf.py', env.udf_file('my_mapping')) elif 'my_filter' in test_yaml: test_yaml = test_yaml.replace( '/path/to/some/udf.py', env.udf_file('my_filter')) spec = yaml.load(test_yaml, Loader=SafeLoader) if test_type == 'PARSE': return if 'ReadFromCsv' in test_yaml: spec = replace_recursive(spec, 'ReadFromCsv', 'path', env.input_csv()) if 'ReadFromText' in test_yaml: spec = replace_recursive(spec, 'ReadFromText', 'path', env.input_csv()) if 'ReadFromJson' in test_yaml: spec = replace_recursive(spec, 'ReadFromJson', 'path', env.input_json()) for write in ['WriteToText', 'WriteToCsv', 'WriteToJson']: if write in test_yaml: spec = replace_recursive(spec, write, 'path', env.output_file()) modified_yaml = yaml.dump(spec) options = {'pickle_library': 'cloudpickle'} if RENDER_DIR is not None: options['runner'] = 'apache_beam.runners.render.RenderRunner' options['render_output'] = [ os.path.join(RENDER_DIR, test_name + '.png') ] options['render_leaf_composite_nodes'] = ['.*'] test_provider = TestProvider(TEST_TRANSFORMS) with mock.patch( 'apache_beam.yaml.yaml_provider.SqlBackedProvider.sql_provider', lambda self: test_provider): # TODO(polber) - remove once there is support for ExternalTransforms # in precommits with mock.patch( 'apache_beam.yaml.yaml_provider.ExternalProvider.create_transform', lambda *args, **kwargs: _Fakes.SomeTransform(*args, **kwargs)): # Uses the FnApiRunner to ensure errors are mocked/passed through # correctly p = beam.Pipeline('FnApiRunner', options=PipelineOptions(**options)) yaml_transform.expand_pipeline( p, modified_yaml, yaml_provider.merge_providers([test_provider])) if test_type == 'BUILD': return p.run().wait_until_finish() return test def parse_test_methods(markdown_lines): # pylint: disable=too-many-nested-blocks def extract_inputs(input_spec): if not input_spec: return set() elif isinstance(input_spec, str): return set([input_spec.split('.')[0]]) elif isinstance(input_spec, list): return set.union(*[extract_inputs(v) for v in input_spec]) elif isinstance(input_spec, dict): return set.union(*[extract_inputs(v) for v in input_spec.values()]) else: raise ValueError("Misformed inputs: " + input_spec) def extract_name(input_spec): return input_spec.get('name', input_spec.get('type')) code_lines = None last_pipeline = None for ix, line in enumerate(markdown_lines): line = line.rstrip() if line == '```': if code_lines is None: code_lines = [] test_type = 'RUN' test_name = f'test_line_{ix + 2}' else: if code_lines: if code_lines[0].startswith('- type:'): specs = yaml.load('\n'.join(code_lines), Loader=SafeLoader) if 'dependencies:' in specs: test_type = 'PARSE' is_chain = not any('input' in spec for spec in specs) if is_chain: undefined_inputs = set(['input']) else: undefined_inputs = set.union( *[extract_inputs(spec.get('input')) for spec in specs]) - set( extract_name(spec) for spec in specs) # Treat this as a fragment of a larger pipeline. # pylint: disable=not-an-iterable code_lines = [ 'pipeline:', ' type: chain' if is_chain else '', ' transforms:', ] + [ ' - {type: ReadFromCsv, name: "%s", config: {path: x}}' % undefined_input for undefined_input in undefined_inputs ] + [' ' + line for line in code_lines] if code_lines[0] == 'pipeline:': yaml_pipeline = '\n'.join(code_lines) last_pipeline = yaml_pipeline if 'providers:' in yaml_pipeline or 'tests:' in yaml_pipeline: test_type = 'PARSE' yield test_name, create_test_method( test_type, test_name, yaml_pipeline) if 'tests:' in code_lines: test_spec = '\n'.join(code_lines) if code_lines[0] == 'pipeline:': yaml_pipeline = '\n'.join(code_lines) else: yaml_pipeline = last_pipeline for sub_ix, test_spec in enumerate(yaml.load( '\n'.join(code_lines), Loader=yaml_utils.SafeLineLoader)['tests']): suffix = test_spec.get('name', str(sub_ix)) yield ( test_name + '_' + suffix, # The yp=... ts=... is to capture the looped closure values. lambda _, yp=yaml_pipeline, ts=test_spec: yaml_testing. run_test(yp, ts)) code_lines = None elif code_lines is not None: code_lines.append(line) def createTestSuite(name, path): with open(path) as readme: return type(name, (unittest.TestCase, ), dict(parse_test_methods(readme))) # These are copied from $ROOT/website/www/site/content/en/documentation/sdks # at build time. YAML_DOCS_DIR = os.path.join(os.path.join(os.path.dirname(__file__), 'docs')) ReadMeTest = createTestSuite( 'ReadMeTest', os.path.join(YAML_DOCS_DIR, 'yaml.md')) ErrorHandlingTest = createTestSuite( 'ErrorHandlingTest', os.path.join(YAML_DOCS_DIR, 'yaml-errors.md')) MappingTest = createTestSuite( 'MappingTest', os.path.join(YAML_DOCS_DIR, 'yaml-udf.md')) CombineTest = createTestSuite( 'CombineTest', os.path.join(YAML_DOCS_DIR, 'yaml-combine.md')) InlinePythonTest = createTestSuite( 'InlinePythonTest', os.path.join(YAML_DOCS_DIR, 'yaml-inline-python.md')) JoinTest = createTestSuite( 'JoinTest', os.path.join(YAML_DOCS_DIR, 'yaml-join.md')) TestingTest = createTestSuite( 'TestingTest', os.path.join(YAML_DOCS_DIR, 'yaml-testing.md')) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--render_dir', default=None) known_args, unknown_args = parser.parse_known_args(sys.argv) if known_args.render_dir: RENDER_DIR = known_args.render_dir logging.getLogger().setLevel(logging.INFO) unittest.main(argv=unknown_args) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/standard_io.yaml0000644000175100017510000003606415206113033022004 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # This file enumerates the various IOs that are available by default as # top-level transforms in Beam's YAML. # # Note that there may be redundant implementations. In these cases the specs # should be kept in sync. # TODO(yaml): See if this can be enforced programmatically. # BigQuery - type: renaming transforms: 'ReadFromBigQuery': 'ReadFromBigQuery' 'WriteToBigQuery': 'WriteToBigQuery' config: mappings: 'ReadFromBigQuery': query: 'query' table: 'table_spec' fields: 'selected_fields' row_restriction: 'row_restriction' 'WriteToBigQuery': table: 'table' create_disposition: 'create_disposition' write_disposition: 'write_disposition' error_handling: 'error_handling' # TODO(https://github.com/apache/beam/issues/30058): Required until autosharding support is fixed num_streams: 'num_streams' underlying_provider: type: beamJar transforms: 'ReadFromBigQuery': 'beam:schematransform:org.apache.beam:bigquery_storage_read:v1' 'WriteToBigQuery': 'beam:schematransform:org.apache.beam:bigquery_write:v1' config: gradle_target: 'sdks:java:io:google-cloud-platform:expansion-service:shadowJar' managed_replacement: # Following transforms may be replaced with equivalent managed transforms, # if the pipelines 'updateCompatibilityBeamVersion' match the provided # version. 'ReadFromBigQuery': '2.69.0' 'WriteToBigQuery': '2.69.0' # Kafka - type: renaming transforms: 'ReadFromKafka': 'ReadFromKafka' 'WriteToKafka': 'WriteToKafka' config: mappings: 'ReadFromKafka': 'schema': 'schema' 'consumer_config': 'consumer_config_updates' 'format': 'format' 'topic': 'topic' 'bootstrap_servers': 'bootstrap_servers' 'confluent_schema_registry_url': 'confluent_schema_registry_url' 'confluent_schema_registry_subject': 'confluent_schema_registry_subject' 'auto_offset_reset_config': 'auto_offset_reset_config' 'error_handling': 'error_handling' 'file_descriptor_path': 'file_descriptor_path' 'message_name': 'message_name' 'max_read_time_seconds': 'max_read_time_seconds' 'WriteToKafka': 'format': 'format' 'topic': 'topic' 'bootstrap_servers': 'bootstrap_servers' 'producer_config_updates': 'producer_config_updates' 'error_handling': 'error_handling' 'file_descriptor_path': 'file_descriptor_path' 'message_name': 'message_name' 'schema': 'schema' underlying_provider: type: beamJar transforms: 'ReadFromKafka': 'beam:schematransform:org.apache.beam:kafka_read:v1' 'WriteToKafka': 'beam:schematransform:org.apache.beam:kafka_write:v1' config: gradle_target: 'sdks:java:io:expansion-service:shadowJar' managed_replacement: # Following transforms may be replaced with equivalent managed transforms, # if the pipelines 'updateCompatibilityBeamVersion' match the provided # version. 'ReadFromKafka': '2.65.0' 'WriteToKafka': '2.65.0' # TODO(yaml): Tests are assuming python providers are before java ones, hence # the order below. This should be fixed in the future. # Python Providers - type: python transforms: 'ReadFromBigQuery': 'apache_beam.yaml.yaml_io.read_from_bigquery' # Disable until https://github.com/apache/beam/issues/28162 is resolved. # 'WriteToBigQuery': 'apache_beam.yaml.yaml_io.write_to_bigquery' 'ReadFromText': 'apache_beam.yaml.yaml_io.read_from_text' 'WriteToText': 'apache_beam.yaml.yaml_io.write_to_text' 'ReadFromPubSub': 'apache_beam.yaml.yaml_io.read_from_pubsub' 'WriteToPubSub': 'apache_beam.yaml.yaml_io.write_to_pubsub' 'ReadFromIceberg': 'apache_beam.yaml.yaml_io.read_from_iceberg' 'WriteToIceberg': 'apache_beam.yaml.yaml_io.write_to_iceberg' 'ReadFromTFRecord': 'apache_beam.yaml.yaml_io.read_from_tfrecord' 'WriteToTFRecord': 'apache_beam.yaml.yaml_io.write_to_tfrecord' # General File Formats # Declared as a renaming transform to avoid exposing all # (implementation-specific) pandas arguments and aligning with possible Java # implementation. # Invoking these directly as a PyTransform is still an option for anyone wanting # to use these power-features in a language-dependent manner. - type: renaming transforms: 'ReadFromCsv': 'ReadFromCsv' 'WriteToCsv': 'WriteToCsv' 'ReadFromJson': 'ReadFromJson' 'WriteToJson': 'WriteToJson' 'ReadFromParquet': 'ReadFromParquet' 'WriteToParquet': 'WriteToParquet' 'ReadFromAvro': 'ReadFromAvro' 'WriteToAvro': 'WriteToAvro' config: mappings: 'ReadFromCsv': path: 'path' delimiter: 'sep' comment: 'comment' filename_column: 'filename_column' 'WriteToCsv': path: 'path' delimiter: 'sep' 'ReadFromJson': path: 'path' 'WriteToJson': path: 'path' 'ReadFromParquet': path: 'file_pattern' 'WriteToParquet': path: 'file_path_prefix' 'ReadFromAvro': path: 'file_pattern' 'WriteToAvro': path: 'file_path_prefix' defaults: 'ReadFromParquet': as_rows: True 'ReadFromAvro': as_rows: True underlying_provider: type: python transforms: 'ReadFromCsv': 'apache_beam.io.ReadFromCsv' 'WriteToCsv': 'apache_beam.io.WriteToCsv' 'ReadFromJson': 'apache_beam.io.ReadFromJson' 'WriteToJson': 'apache_beam.io.WriteToJson' 'ReadFromParquet': 'apache_beam.io.ReadFromParquet' 'WriteToParquet': 'apache_beam.io.WriteToParquet' 'ReadFromAvro': 'apache_beam.io.ReadFromAvro' 'WriteToAvro': 'apache_beam.io.WriteToAvro' # BeamJar Providers - type: beamJar transforms: 'WriteToCsv': 'beam:schematransform:org.apache.beam:csv_write:v1' 'WriteToJson': 'beam:schematransform:org.apache.beam:json_write:v1' config: gradle_target: 'sdks:java:extensions:schemaio-expansion-service:shadowJar' # Databases - type: renaming transforms: 'ReadFromJdbc': 'ReadFromJdbc' 'WriteToJdbc': 'WriteToJdbc' 'ReadFromMySql': 'ReadFromMySql' 'WriteToMySql': 'WriteToMySql' 'ReadFromPostgres': 'ReadFromPostgres' 'WriteToPostgres': 'WriteToPostgres' 'ReadFromOracle': 'ReadFromOracle' 'WriteToOracle': 'WriteToOracle' 'ReadFromSqlServer': 'ReadFromSqlServer' 'WriteToSqlServer': 'WriteToSqlServer' config: mappings: 'ReadFromJdbc': url: 'jdbc_url' connection_init_sql: 'connection_init_sql' connection_properties: 'connection_properties' disable_auto_commit: 'disable_auto_commit' driver_class_name: 'driver_class_name' driver_jars: 'driver_jars' fetch_size: 'fetch_size' output_parallelization: 'output_parallelization' password: 'password' query: 'read_query' table: 'location' partition_column : 'partition_column' num_partitions: 'num_partitions' type: 'jdbc_type' username: 'username' 'WriteToJdbc': url: 'jdbc_url' auto_sharding: 'autosharding' connection_init_sql: 'connection_init_sql' connection_properties: 'connection_properties' driver_class_name: 'driver_class_name' driver_jars: 'driver_jars' password: 'password' table: 'location' batch_size: 'batch_size' type: 'jdbc_type' username: 'username' query: 'write_statement' 'ReadFromMySql': 'ReadFromJdbc' 'WriteToMySql': 'WriteToJdbc' 'ReadFromPostgres': 'ReadFromJdbc' 'WriteToPostgres': 'WriteToJdbc' 'ReadFromOracle': 'ReadFromJdbc' 'WriteToOracle': 'WriteToJdbc' 'ReadFromSqlServer': 'ReadFromJdbc' 'WriteToSqlServer': 'WriteToJdbc' defaults: 'ReadFromMySql': driver_class_name: '' driver_jars: '' jdbc_type: '' 'WriteToMySql': driver_class_name: '' driver_jars: '' jdbc_type: '' 'ReadFromPostgres': connection_init_sql: [] driver_class_name: '' driver_jars: '' jdbc_type: '' 'WriteToPostgres': connection_init_sql: [] driver_class_name: '' driver_jars: '' jdbc_type: '' 'ReadFromOracle': connection_init_sql: [] driver_class_name: '' driver_jars: '' jdbc_type: '' 'WriteToOracle': connection_init_sql: [] driver_class_name: '' driver_jars: '' jdbc_type: '' 'ReadFromSqlServer': connection_init_sql: [] driver_class_name: '' driver_jars: '' jdbc_type: '' 'WriteToSqlServer': connection_init_sql: [] driver_class_name: '' driver_jars: '' jdbc_type: '' underlying_provider: type: beamJar transforms: 'ReadFromJdbc': 'beam:schematransform:org.apache.beam:jdbc_read:v1' 'ReadFromMySql': 'beam:schematransform:org.apache.beam:mysql_read:v1' 'ReadFromPostgres': 'beam:schematransform:org.apache.beam:postgres_read:v1' 'ReadFromOracle': 'beam:schematransform:org.apache.beam:oracle_read:v1' 'ReadFromSqlServer': 'beam:schematransform:org.apache.beam:sql_server_read:v1' 'WriteToJdbc': 'beam:schematransform:org.apache.beam:jdbc_write:v1' 'WriteToMySql': 'beam:schematransform:org.apache.beam:mysql_write:v1' 'WriteToPostgres': 'beam:schematransform:org.apache.beam:postgres_write:v1' 'WriteToOracle': 'beam:schematransform:org.apache.beam:oracle_write:v1' 'WriteToSqlServer': 'beam:schematransform:org.apache.beam:sql_server_write:v1' config: gradle_target: 'sdks:java:extensions:schemaio-expansion-service:shadowJar' managed_replacement: # Following transforms may be replaced with equivalent managed transforms, # if the pipelines 'updateCompatibilityBeamVersion' match the provided # version. 'ReadFromPostgres': '2.73.0' 'WriteToPostgres': '2.73.0' 'ReadFromMySql': '2.73.0' 'WriteToMySql': '2.73.0' 'ReadFromSqlServer': '2.73.0' 'WriteToSqlServer': '2.73.0' # Spanner - type: renaming transforms: 'ReadFromSpanner': 'ReadFromSpanner' 'WriteToSpanner': 'WriteToSpanner' config: mappings: 'ReadFromSpanner': project: 'project_id' instance: 'instance_id' database: 'database_id' table: 'table_id' query: 'query' columns: 'columns' index: 'index' batching: 'batching' error_handling: 'error_handling' 'WriteToSpanner': project: 'project_id' instance: 'instance_id' database: 'database_id' table: 'table_id' error_handling: 'error_handling' underlying_provider: type: beamJar transforms: 'ReadFromSpanner': 'beam:schematransform:org.apache.beam:spanner_read:v1' 'WriteToSpanner': 'beam:schematransform:org.apache.beam:spanner_write:v1' config: gradle_target: 'sdks:java:io:google-cloud-platform:expansion-service:shadowJar' # TFRecord - type: renaming transforms: 'ReadFromTFRecord': 'ReadFromTFRecord' 'WriteToTFRecord': 'WriteToTFRecord' config: mappings: 'ReadFromTFRecord': file_pattern: 'file_pattern' compression_type: 'compression' validate: 'validate' 'WriteToTFRecord': file_path_prefix: 'output_prefix' shard_name_template: 'shard_template' file_name_suffix: 'filename_suffix' num_shards: 'num_shards' compression_type: 'compression' underlying_provider: type: beamJar transforms: 'ReadFromTFRecord': 'beam:schematransform:org.apache.beam:tfrecord_read:v1' 'WriteToTFRecord': 'beam:schematransform:org.apache.beam:tfrecord_write:v1' config: gradle_target: 'sdks:java:io:expansion-service:shadowJar' #BigTable - type: renaming transforms: 'ReadFromBigTable': 'ReadFromBigTable' 'WriteToBigTable': 'WriteToBigTable' config: mappings: #Temp removing read from bigTable IO 'ReadFromBigTable': project: 'project_id' instance: 'instance_id' table: 'table_id' flatten: "flatten" 'WriteToBigTable': project: 'project_id' instance: 'instance_id' table: 'table_id' underlying_provider: type: beamJar transforms: 'ReadFromBigTable': 'beam:schematransform:org.apache.beam:bigtable_read:v1' 'WriteToBigTable': 'beam:schematransform:org.apache.beam:bigtable_write:v1' config: gradle_target: 'sdks:java:io:google-cloud-platform:expansion-service:shadowJar' #IcebergCDC - type: renaming transforms: 'ReadFromIcebergCDC': 'ReadFromIcebergCDC' config: mappings: 'ReadFromIcebergCDC': table: 'table' catalog_name: 'catalog_name' catalog_properties: 'catalog_properties' config_properties: 'config_properties' drop: 'drop' filter: 'filter' from_snapshot: 'from_snapshot' from_timestamp: 'from_timestamp' keep: 'keep' poll_interval_seconds: 'poll_interval_seconds' starting_strategy: 'starting_strategy' streaming: 'streaming' to_snapshot: 'to_snapshot' to_timestamp: 'to_timestamp' underlying_provider: type: beamJar transforms: 'ReadFromIcebergCDC': 'beam:schematransform:org.apache.beam:iceberg_cdc_read:v1' config: gradle_target: 'sdks:java:io:expansion-service:shadowJar' #IcebergAddFiles - type: renaming transforms: 'IcebergAddFiles': 'IcebergAddFiles' config: mappings: 'IcebergAddFiles': table: 'table' catalog_properties: 'catalog_properties' config_properties: 'config_properties' triggering_frequency_seconds: 'triggering_frequency_seconds' append_batch_size: 'append_batch_size' location_prefix: 'location_prefix' partition_fields: 'partition_fields' table_properties: 'table_properties' error_handling: 'error_handling' underlying_provider: type: beamJar transforms: 'IcebergAddFiles': 'beam:schematransform:iceberg_add_files:v1' config: gradle_target: 'sdks:java:io:expansion-service:shadowJar' ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/standard_providers.yaml0000644000175100017510000000746615206113033023416 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # TODO(robertwb): Add more providers. # TODO(robertwb): Perhaps auto-generate this file? - type: 'beamJar' config: gradle_target: 'sdks:java:extensions:sql:expansion-service:shadowJar' version: BEAM_VERSION transforms: Sql: 'beam:external:java:sql:v1' Flatten: 'beam:schematransform:org.apache.beam:yaml:flatten:v1' LogForTesting: 'beam:schematransform:org.apache.beam:yaml:log_for_testing:v1' # TODO(robertwb): Auto-detect redundantly provided transforms for maximal fusion. - type: 'beamJar' config: gradle_target: 'sdks:java:io:expansion-service:shadowJar' version: BEAM_VERSION transforms: Flatten: 'beam:schematransform:org.apache.beam:yaml:flatten:v1' LogForTesting: 'beam:schematransform:org.apache.beam:yaml:log_for_testing:v1' - type: 'beamJar' config: gradle_target: 'sdks:java:io:google-cloud-platform:expansion-service:shadowJar' version: BEAM_VERSION transforms: Flatten: 'beam:schematransform:org.apache.beam:yaml:flatten:v1' LogForTesting: 'beam:schematransform:org.apache.beam:yaml:log_for_testing:v1' - type: 'beamJar' config: gradle_target: 'sdks:java:extensions:schemaio-expansion-service:shadowJar' version: BEAM_VERSION transforms: Flatten: 'beam:schematransform:org.apache.beam:yaml:flatten:v1' LogForTesting: 'beam:schematransform:org.apache.beam:yaml:log_for_testing:v1' - type: 'python' config: {} transforms: MLTransform: 'apache_beam.yaml.yaml_ml.ml_transform' RunInference: 'apache_beam.yaml.yaml_ml.run_inference' - type: renaming transforms: 'MapToFields-java': 'MapToFields-java' 'MapToFields-generic': 'MapToFields-java' 'Filter-java': 'Filter-java' 'Filter-generic': 'Filter-java' 'Explode': 'Explode' config: mappings: 'Sql': query: 'query' # Unfortunately dialect is a java logical type. 'MapToFields-generic': language: 'language' append: 'append' drop: 'drop' fields: 'fields' error_handling: 'error_handling' 'MapToFields-java': language: 'language' append: 'append' drop: 'drop' fields: 'fields' error_handling: 'error_handling' 'Filter-generic': language: 'language' keep: 'keep' error_handling: 'error_handling' 'Filter-java': language: 'language' keep: 'keep' error_handling: 'error_handling' 'Explode': fields: 'fields' cross_product: 'cross_product' underlying_provider: type: beamJar transforms: Sql: "schematransform:org.apache.beam:sql_transform:v1" MapToFields-java: "beam:schematransform:org.apache.beam:yaml:map_to_fields-java:v1" Filter-java: "beam:schematransform:org.apache.beam:yaml:filter-java:v1" Explode: "beam:schematransform:org.apache.beam:yaml:explode:v1" config: gradle_target: 'sdks:java:extensions:sql:expansion-service:shadowJar' - type: 'python' config: {} transforms: Enrichment: 'apache_beam.yaml.yaml_enrichment.enrichment_transform' ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_combine.py0000644000175100017510000001710715206113033021636 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """This module defines the basic Combine operation.""" from collections.abc import Iterable from collections.abc import Mapping from typing import Any from typing import Optional import apache_beam as beam from apache_beam import typehints from apache_beam.typehints import row_type from apache_beam.typehints import trivial_inference from apache_beam.typehints.decorators import get_type_hints from apache_beam.typehints.schemas import named_fields_from_element_type from apache_beam.utils import python_callable from apache_beam.yaml import yaml_mapping from apache_beam.yaml import yaml_provider BUILTIN_COMBINE_FNS = { 'sum': sum, 'max': max, 'min': min, 'all': all, 'any': any, 'mean': beam.transforms.combiners.MeanCombineFn(), 'count': beam.transforms.combiners.CountCombineFn(), 'group': beam.transforms.combiners.ToListCombineFn(), 'concat': beam.transforms.combiners.ConcatListCombineFn(), } def normalize_combine(spec): """Expands various shorthand specs for combine (which can otherwise be quite verbose for simple cases.) We do this here so that it doesn't need to be done per language. The following are all equivalent:: dest: fn_type dest: value: dest fn: fn_type dest: value: dest fn: type: fn_type """ from apache_beam.yaml.yaml_transform import SafeLineLoader if spec['type'] == 'Combine': config = spec.get('config') if isinstance(config.get('group_by'), str): config['group_by'] = [config['group_by']] def normalize_agg(dest, agg): if isinstance(agg, str): agg = {'fn': agg} if 'value' not in agg and config.get('language') != 'sql': agg['value'] = dest if isinstance(agg['fn'], str): agg['fn'] = {'type': agg['fn']} return agg if 'combine' not in config: raise ValueError('Missing combine parameter in Combine config.') config['combine'] = { dest: normalize_agg(dest, agg) for (dest, agg) in SafeLineLoader.strip_metadata(config['combine']).items() } return spec class PyJsYamlCombine(beam.PTransform): """Groups and combines records sharing common fields. Built-in combine functions are BUILTIN_COMBINE_FNS but custom aggregation functions can be used as well. See also the documentation on [YAML Aggregation](https://beam.apache.org/documentation/sdks/yaml-combine/). Args: group_by: The field(s) to aggregate on. combine: The aggregation function to use. language: The language used to define (and execute) the custom callables in `combine`. Defaults to generic. """ def __init__( self, group_by: Iterable[str], combine: Mapping[str, Mapping[str, Any]], language: Optional[str] = None): self._group_by = group_by self._combine = combine self._language = language def expand(self, pcoll): input_types = dict(named_fields_from_element_type(pcoll.element_type)) all_fields = list(input_types.keys()) unknown_keys = set(self._group_by) - set(all_fields) if unknown_keys: raise ValueError(f'Unknown grouping columns: {list(unknown_keys)}') def create_combine_fn(fn_spec): if 'type' not in fn_spec: raise ValueError(f'CombineFn spec missing type: {fn_spec}') elif fn_spec['type'] in BUILTIN_COMBINE_FNS: return BUILTIN_COMBINE_FNS[fn_spec['type']] elif self._language == 'python': # TODO(yaml): Support output_type here as well. fn = python_callable.PythonCallableWithSource.load_from_source( fn_spec['type']) if 'config' in fn_spec: fn = fn(**fn_spec['config']) return fn else: raise TypeError('Unknown CombineFn: {fn_spec}') def extract_return_type(expr): if isinstance(expr, str) and expr in input_types: return input_types[expr] expr_hints = get_type_hints(expr) if (expr_hints and expr_hints.has_simple_output_type() and expr_hints.simple_output_type(None) != typehints.Any): return expr_hints.simple_output_type(None) elif callable(expr): return trivial_inference.infer_return_type(expr, [pcoll.element_type]) else: return Any # TODO(yaml): Support error handling. transform = beam.GroupBy(*self._group_by) output_types = [(k, input_types[k]) for k in self._group_by] for output, agg in self._combine.items(): expr = yaml_mapping._as_callable( all_fields, agg['value'], 'Combine', self._language, input_types) fn = create_combine_fn(agg['fn']) transform = transform.aggregate_field(expr, fn, output) # TODO(yaml): See if this logic can be pushed into GroupBy itself. expr_type = extract_return_type(expr) if isinstance(fn, beam.CombineFn): # TODO(yaml): Better inference on CombineFns whose outputs types are # functions of their input types combined_type = extract_return_type(fn) elif fn in (sum, min, max): combined_type = expr_type elif fn in (any, all): combined_type = bool else: combined_type = Any output_types.append((output, combined_type)) return pcoll | transform.with_output_types( row_type.RowTypeConstraint.from_fields(output_types)) if PyJsYamlCombine.__doc__: # make mypy happy PyJsYamlCombine.__doc__ = PyJsYamlCombine.__doc__.replace( 'BUILTIN_COMBINE_FNS', ', '.join('`%s`' % k for k in BUILTIN_COMBINE_FNS.keys())) @beam.ptransform.ptransform_fn def _SqlCombineTransform( pcoll, sql_transform_constructor, group_by, combine, language=None): all_fields = [ x for x, _ in named_fields_from_element_type(pcoll.element_type) ] unknown_keys = set(group_by) - set(all_fields) if unknown_keys: raise ValueError(f'Unknown grouping columns: {list(unknown_keys)}') def combine_col(dest, fn_spec): if 'value' in fn_spec or 'config' in fn_spec['fn']: expr = '%s(%s)' % ( fn_spec['fn']['type'], ', '.join([fn_spec['value']] + list(fn_spec['fn'].get('config', {}).values()))) else: expr = fn_spec['fn']['type'] return f'{expr} as {dest}' return pcoll | sql_transform_constructor( 'SELECT %s FROM PCOLLECTION GROUP BY %s' % ( ', '.join( list(group_by) + [combine_col(dest, fn_spec) for dest, fn_spec in combine.items()]), ', '.join(group_by), )) def create_combine_providers(): return [ yaml_provider.InlineProvider({ 'Combine-generic': PyJsYamlCombine, 'Combine-python': PyJsYamlCombine, 'Combine-javascript': PyJsYamlCombine, }), yaml_provider.SqlBackedProvider({ 'Combine-generic': _SqlCombineTransform, 'Combine-sql': _SqlCombineTransform, 'Combine-calcite': _SqlCombineTransform, }), ] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_combine_test.py0000644000175100017510000001340515206113033022672 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest import apache_beam as beam from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.yaml.yaml_transform import YamlTransform DATA = [ beam.Row(a='x', b=1, c=101), beam.Row(a='x', b=1, c=102), beam.Row(a='y', b=1, c=103), beam.Row(a='y', b=2, c=104), ] class YamlCombineTest(unittest.TestCase): def test_multiple_aggregations(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' type: Combine config: group_by: a combine: b: sum c: max ''') assert_that( result | beam.Map(lambda x: beam.Row(**x._asdict())), equal_to([ beam.Row(a='x', b=2, c=102), beam.Row(a='y', b=3, c=104), ])) def test_multiple_keys(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' type: Combine config: group_by: [a, b] combine: c: sum ''') assert_that( result | beam.Map(lambda x: beam.Row(**x._asdict())), equal_to([ beam.Row(a='x', b=1, c=203), beam.Row(a='y', b=1, c=103), beam.Row(a='y', b=2, c=104), ])) def test_no_keys(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' type: Combine config: group_by: [] combine: c: sum ''') assert_that( result | beam.Map(lambda x: beam.Row(**x._asdict())), equal_to([ beam.Row(c=410), ])) def test_multiple_combines(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' type: Combine config: group_by: a combine: min_c: fn: min value: c max_c: fn: max value: c ''') assert_that( result | beam.Map(lambda x: beam.Row(**x._asdict())), equal_to([ beam.Row(a='x', min_c=101, max_c=102), beam.Row(a='y', min_c=103, max_c=104), ])) def test_group(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' type: Combine config: language: python group_by: a combine: b: fn: sum value: b c: fn: group value: c ''') assert_that( result | beam.Map(lambda x: beam.Row(**x._asdict())), equal_to([ beam.Row(a='x', b=2, c=[101, 102]), beam.Row(a='y', b=3, c=[103, 104]), ])) def test_expression(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' type: Combine config: language: python group_by: a combine: max: fn: max value: b + c ''') assert_that( result | beam.Map(lambda x: beam.Row(**x._asdict())), equal_to([ beam.Row(a='x', max=103), beam.Row(a='y', max=106), ])) def test_config(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' type: Combine config: language: python group_by: b combine: biggest: fn: type: 'apache_beam.transforms.combiners.TopCombineFn' config: n: 2 value: c ''') assert_that( result | beam.Map(lambda x: beam.Row(**x._asdict())), equal_to([ beam.Row(b=1, biggest=[103, 102]), beam.Row(b=2, biggest=[104]), ])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_enrichment.py0000644000175100017510000001205315206113033022351 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from typing import Any from typing import Optional import apache_beam as beam from apache_beam.yaml import options try: from apache_beam.transforms.enrichment import Enrichment from apache_beam.transforms.enrichment_handlers.bigquery import BigQueryEnrichmentHandler from apache_beam.transforms.enrichment_handlers.bigtable import BigTableEnrichmentHandler from apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store import VertexAIFeatureStoreEnrichmentHandler except ImportError: Enrichment = None # type: ignore BigQueryEnrichmentHandler = None # type: ignore BigTableEnrichmentHandler = None # type: ignore VertexAIFeatureStoreEnrichmentHandler = None # type: ignore try: from apache_beam.transforms.enrichment_handlers.feast_feature_store import FeastFeatureStoreEnrichmentHandler except ImportError: FeastFeatureStoreEnrichmentHandler = None # type: ignore @beam.ptransform.ptransform_fn def enrichment_transform( pcoll, enrichment_handler: str, handler_config: dict[str, Any], timeout: Optional[float] = 30): # pylint: disable=line-too-long """ The Enrichment transform allows one to dynamically enhance elements in a pipeline by performing key-value lookups against external services like APIs or databases. Example using BigTable: :: - type: Enrichment config: enrichment_handler: 'BigTable' handler_config: project_id: 'apache-beam-testing' instance_id: 'beam-test' table_id: 'bigtable-enrichment-test' row_key: 'product_id' timeout: 30 For more information on Enrichment, see the [Beam docs]( https://beam.apache.org/documentation/transforms/python/elementwise/enrichment/). Args: enrichment_handler (str): Specifies the source from where data needs to be extracted into the pipeline for enriching data. One of "BigQuery", "BigTable", "FeastFeatureStore" or "VertexAIFeatureStore". handler_config (str): Specifies the parameters for the respective enrichment_handler in a YAML/JSON format. To see the full set of handler_config parameters, see their corresponding doc pages: - [BigQueryEnrichmentHandler](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.bigquery.html#apache_beam.transforms.enrichment_handlers.bigquery.BigQueryEnrichmentHandler) - [BigTableEnrichmentHandler](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.bigtable.html#apache_beam.transforms.enrichment_handlers.bigtable.BigTableEnrichmentHandler) - [FeastFeatureStoreEnrichmentHandler](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.feast_feature_store.html#apache_beam.transforms.enrichment_handlers.feast_feature_store.FeastFeatureStoreEnrichmentHandler) - [VertexAIFeatureStoreEnrichmentHandler](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store.html#apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store.VertexAIFeatureStoreEnrichmentHandler) timeout (float): Timeout for source requests in seconds. Defaults to 30 seconds. """ options.YamlOptions.check_enabled(pcoll.pipeline, 'Enrichment') if not Enrichment: raise ValueError( f"gcp dependencies not installed. Cannot use {enrichment_handler} " f"handler. Please install using 'pip install apache-beam[gcp]'.") if (enrichment_handler == 'FeastFeatureStore' and not FeastFeatureStoreEnrichmentHandler): raise ValueError( "FeastFeatureStore handler requires 'feast' package to be installed. " + "Please install using 'pip install feast[gcp]' and try again.") handler_map = { 'BigQuery': BigQueryEnrichmentHandler, 'BigTable': BigTableEnrichmentHandler, 'FeastFeatureStore': FeastFeatureStoreEnrichmentHandler, 'VertexAIFeatureStore': VertexAIFeatureStoreEnrichmentHandler } if enrichment_handler not in handler_map: raise ValueError(f"Unknown enrichment source: {enrichment_handler}") handler = handler_map[enrichment_handler](**handler_config) return pcoll | Enrichment(source_handler=handler, timeout=timeout) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_enrichment_test.py0000644000175100017510000000541215206113033023411 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest import mock import apache_beam as beam from apache_beam import Row from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.yaml import yaml_provider from apache_beam.yaml.yaml_transform import YamlTransform class FakeEnrichmentTransform: def __init__(self, enrichment_handler, handler_config, timeout=30): self._enrichment_handler = enrichment_handler self._handler_config = handler_config self._timeout = timeout def __call__(self, enrichment_handler, *, handler_config, timeout=30): assert enrichment_handler == self._enrichment_handler assert handler_config == self._handler_config assert timeout == self._timeout return beam.Map(lambda x: beam.Row(**x._asdict())) class EnrichmentTransformTest(unittest.TestCase): def test_enrichment_with_bigquery(self): input_data = [ Row(label="item1", rank=0), Row(label="item2", rank=1), ] handler = 'BigQuery' config = { "project": "apache-beam-testing", "table_name": "project.database.table", "row_restriction_template": "label='item1' or label='item2'", "fields": ["label"] } with beam.Pipeline() as p: with mock.patch('apache_beam.yaml.yaml_enrichment.enrichment_transform', FakeEnrichmentTransform(enrichment_handler=handler, handler_config=config)): # Force a reload to respect our mock. yaml_provider.standard_providers.cache_clear() input_pcoll = p | 'CreateInput' >> beam.Create(input_data) result = input_pcoll | YamlTransform( f''' type: Enrichment config: enrichment_handler: {handler} handler_config: {config} ''') assert_that(result, equal_to(input_data)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_errors.py0000644000175100017510000000627515206113033021542 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import functools import inspect from typing import NamedTuple import apache_beam as beam from apache_beam.typehints.row_type import RowTypeConstraint from apache_beam.yaml.yaml_utils import SafeLineLoader class ErrorHandlingConfig(NamedTuple): """This option specifies whether and where to output error rows. Args: output (str): Name to use for the output error collection """ output: str # TODO: Other parameters are valid here too, but not common to Java. def exception_handling_args(error_handling_spec): if error_handling_spec: # error_handling_spec may have come from a yaml file and have metadata. clean_spec = SafeLineLoader.strip_metadata(error_handling_spec) return { 'dead_letter_tag' if k == 'output' else k: v for (k, v) in clean_spec.items() } else: return None def map_errors_to_standard_format(input_type): # TODO(https://github.com/apache/beam/issues/24755): Switch to MapTuple. return beam.Map( lambda x: beam.Row( element=x[0], msg=str(x[1][1]), stack=''.join(x[1][2])) ).with_output_types( RowTypeConstraint.from_fields([("element", input_type), ("msg", str), ("stack", str)])) def maybe_with_exception_handling(inner_expand): def expand(self, pcoll): wrapped_pcoll = beam.core._MaybePValueWithErrors( pcoll, self._exception_handling_args) return inner_expand(self, wrapped_pcoll).as_result( map_errors_to_standard_format(pcoll.element_type)) return expand def maybe_with_exception_handling_transform_fn(transform_fn): @functools.wraps(transform_fn) def expand(pcoll, error_handling=None, **kwargs): wrapped_pcoll = beam.core._MaybePValueWithErrors( pcoll, exception_handling_args(error_handling)) return transform_fn(wrapped_pcoll, **kwargs).as_result( map_errors_to_standard_format(pcoll.element_type)) original_signature = inspect.signature(transform_fn) new_parameters = list(original_signature.parameters.values()) error_handling_param = inspect.Parameter( 'error_handling', inspect.Parameter.KEYWORD_ONLY, default=None, annotation=ErrorHandlingConfig) if new_parameters[-1].kind == inspect.Parameter.VAR_KEYWORD: new_parameters.insert(-1, error_handling_param) else: new_parameters.append(error_handling_param) expand.__signature__ = original_signature.replace(parameters=new_parameters) return expand ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_io.py0000644000175100017510000007030615206113033020631 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """This module contains the Python implementations for the builtin IOs. They are referenced from standard_io.py. Note that in the case that they overlap with other (likely Java) implementations of the same transforms, the configs must be kept in sync. """ import io from collections.abc import Callable from collections.abc import Iterable from collections.abc import Mapping from typing import Any from typing import Optional import fastavro import apache_beam as beam import apache_beam.io as beam_io from apache_beam import coders from apache_beam.coders.row_coder import RowCoder from apache_beam.io import ReadFromBigQuery from apache_beam.io import ReadFromTFRecord from apache_beam.io import WriteToBigQuery from apache_beam.io import WriteToTFRecord from apache_beam.io import avroio from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.gcp.bigquery import BigQueryDisposition from apache_beam.portability.api import schema_pb2 from apache_beam.typehints import schemas from apache_beam.yaml import json_utils from apache_beam.yaml import yaml_errors from apache_beam.yaml import yaml_provider from apache_beam.yaml import yaml_utils def read_from_text(path: str): # TODO(yaml): Consider passing the filename and offset, possibly even # by default. """Reads lines from a text files. The resulting PCollection consists of rows with a single string field named "line." Args: path (str): The file path to read from. The path can contain glob characters such as ``*`` and ``?``. """ return beam_io.ReadFromText(path) | beam.Map(lambda s: beam.Row(line=s)) @beam.ptransform_fn def write_to_text(pcoll, path: str): """Writes a PCollection to a (set of) text files(s). The input must be a PCollection whose schema has exactly one field. Args: path (str): The file path to write to. The files written will begin with this prefix, followed by a shard identifier. """ try: field_names = [ name for name, _ in schemas.named_fields_from_element_type( pcoll.element_type) ] except Exception as exn: raise ValueError( "WriteToText requires an input schema with exactly one field.") from exn if len(field_names) != 1: raise ValueError( "WriteToText requires an input schema with exactly one field, got %s" % field_names) sole_field_name, = field_names return pcoll | beam.Map( lambda x: str(getattr(x, sole_field_name))) | beam.io.WriteToText(path) def read_from_bigquery( *, table: Optional[str] = None, query: Optional[str] = None, row_restriction: Optional[str] = None, fields: Optional[Iterable[str]] = None): """Reads data from BigQuery. Exactly one of table or query must be set. If query is set, neither row_restriction nor fields should be set. Args: table (str): The table to read from, specified as `DATASET.TABLE` or `PROJECT:DATASET.TABLE`. query (str): A query to be used instead of the table argument. row_restriction (str): Optional SQL text filtering statement, similar to a WHERE clause in a query. Aggregates are not supported. Restricted to a maximum length for 1 MB. selected_fields (list[str]): Optional List of names of the fields in the table that should be read. If empty, all fields will be read. If the specified field is a nested field, all the sub-fields in the field will be selected. The output field order is unrelated to the order of fields given here. """ if query is None: assert table is not None else: assert table is None and row_restriction is None and fields is None return ReadFromBigQuery( query=query, table=table, row_restriction=row_restriction, selected_fields=fields, method='DIRECT_READ', output_type='BEAM_ROW') def write_to_bigquery( table: str, *, create_disposition: Optional[str] = BigQueryDisposition.CREATE_IF_NEEDED, write_disposition: Optional[str] = BigQueryDisposition.WRITE_APPEND, error_handling=None): f"""Writes data to a BigQuery table. Args: table (str): The table to read from, specified as `DATASET.TABLE` or `PROJECT:DATASET.TABLE`. create_disposition (BigQueryDisposition): A string describing what happens if the table does not exist. Possible values are: * :attr:`{BigQueryDisposition.CREATE_IF_NEEDED}`: create if does not exist. * :attr:`{BigQueryDisposition.CREATE_NEVER}`: fail the write if does not exist. Defaults to `{BigQueryDisposition.CREATE_IF_NEEDED}`. write_disposition (BigQueryDisposition): A string describing what happens if the table has already some data. Possible values are: * :attr:`{BigQueryDisposition.WRITE_TRUNCATE}`: delete existing rows. * :attr:`{BigQueryDisposition.WRITE_APPEND}`: add to existing rows. * :attr:`{BigQueryDisposition.WRITE_EMPTY}`: fail the write if table not empty. For streaming pipelines WriteTruncate can not be used. Defaults to `{BigQueryDisposition.WRITE_APPEND}`. error_handling: If specified, should be a mapping giving an output into which to emit records that failed to be written to BigQuery, as described at https://beam.apache.org/documentation/sdks/yaml-errors/ Otherwise permanently failing records will cause pipeline failure. Note: error_handling requires the Storage Write API method and is not supported with File Loads (FILE_LOADS method). When error_handling is specified, the transform will automatically use STORAGE_WRITE_API. """ class WriteToBigQueryHandlingErrors(beam.PTransform): def default_label(self): return 'WriteToBigQuery' def expand(self, pcoll): write_result = pcoll | WriteToBigQuery( table, method=WriteToBigQuery.Method.STORAGE_WRITE_API if error_handling else None, create_disposition=create_disposition, write_disposition=write_disposition, temp_file_format='AVRO') if error_handling and 'output' in error_handling: # TODO: Support error rates. return { 'post_write': write_result.failed_rows_with_errors | beam.FlatMap(lambda x: None), error_handling['output']: write_result.failed_rows_with_errors } else: if write_result._method == WriteToBigQuery.Method.FILE_LOADS: # Never returns errors, just fails. return { 'post_write': write_result.destination_load_jobid_pairs | beam.FlatMap(lambda x: None) } else: # This should likely be pushed into the BQ read itself to avoid # the possibility of silently ignoring errors. def raise_exception(failed_row_with_error): raise RuntimeError(failed_row_with_error.error_message) _ = write_result.failed_rows_with_errors | beam.Map(raise_exception) return { 'post_write': write_result.failed_rows_with_errors | beam.FlatMap(lambda x: None) } return WriteToBigQueryHandlingErrors() def _create_parser( format, schema: Any) -> tuple[schema_pb2.Schema, Callable[[bytes], beam.Row]]: format = format.upper() def _validate_schema(): if not schema: raise ValueError( f'{format} format requires valid {format} schema to be passed to ' f'schema parameter.') if format == 'RAW': if schema: raise ValueError('RAW format does not take a schema') return ( schema_pb2.Schema(fields=[schemas.schema_field('payload', bytes)]), lambda payload: beam.Row(payload=payload)) if format == 'STRING': if schema: raise ValueError('STRING format does not take a schema') return ( schema_pb2.Schema(fields=[schemas.schema_field('payload', str)]), lambda payload: beam.Row(payload=payload.decode('utf-8'))) elif format == 'JSON': _validate_schema() beam_schema = json_utils.json_schema_to_beam_schema(schema) return beam_schema, json_utils.json_parser(beam_schema, schema) elif format == 'AVRO': _validate_schema() beam_schema = avroio.avro_schema_to_beam_schema(schema) covert_to_row = avroio.avro_dict_to_beam_row(schema, beam_schema) # pylint: disable=line-too-long return ( beam_schema, lambda record: covert_to_row( fastavro.schemaless_reader(io.BytesIO(record), schema))) # type: ignore[call-arg] elif format == 'PROTO': _validate_schema() beam_schema = json_utils.json_schema_to_beam_schema(schema) return beam_schema, RowCoder(beam_schema).decode else: raise ValueError(f'Unknown format: {format}') def _create_formatter( format, schema: Any, beam_schema: schema_pb2.Schema) -> Callable[[beam.Row], bytes]: if format.islower(): format = format.upper() if format == 'RAW': if schema: raise ValueError('RAW format does not take a schema') field_names = [field.name for field in beam_schema.fields] if len(field_names) != 1: raise ValueError(f'Expecting exactly one field, found {field_names}') def convert_to_bytes(row): output = getattr(row, field_names[0]) if isinstance(output, bytes): return output elif isinstance(output, str): return output.encode('utf-8') else: raise ValueError( f"Cannot encode payload for WriteToPubSub. " f"Expected valid string or bytes object, " f"got {repr(output)} of type {type(output)}.") return convert_to_bytes elif format == 'JSON': return json_utils.json_formater(beam_schema) elif format == 'AVRO': avro_schema = schema or avroio.beam_schema_to_avro_schema(beam_schema) from_row = avroio.beam_row_to_avro_dict(avro_schema, beam_schema) def formatter(row): buffer = io.BytesIO() fastavro.schemaless_writer(buffer, avro_schema, from_row(row)) buffer.seek(0) return buffer.read() return formatter elif format == 'PROTO': return RowCoder(beam_schema).encode else: raise ValueError(f'Unknown format: {format}') @beam.ptransform_fn @yaml_errors.maybe_with_exception_handling_transform_fn def read_from_pubsub( root, *, topic: Optional[str] = None, subscription: Optional[str] = None, format: str, schema: Optional[Any] = None, attributes: Optional[Iterable[str]] = None, attributes_map: Optional[str] = None, id_attribute: Optional[str] = None, timestamp_attribute: Optional[str] = None): """Reads messages from Cloud Pub/Sub. Args: topic: Cloud Pub/Sub topic in the form "projects//topics/". If provided, subscription must be None. subscription: Existing Cloud Pub/Sub subscription to use in the form "projects//subscriptions/". If not specified, a temporary subscription will be created from the specified topic. If provided, topic must be None. format: The expected format of the message payload. Currently suported formats are - RAW: Produces records with a single `payload` field whose contents are the raw bytes of the pubsub message. - STRING: Like RAW, but the bytes are decoded as a UTF-8 string. - AVRO: Parses records with a given Avro schema. - JSON: Parses records with a given JSON schema. schema: Schema specification for the given format. attributes: List of attribute keys whose values will be flattened into the output message as additional fields. For example, if the format is `raw` and attributes is `["a", "b"]` then this read will produce elements of the form `Row(payload=..., a=..., b=...)`. attribute_map: Name of a field in which to store the full set of attributes associated with this message. For example, if the format is `raw` and `attribute_map` is set to `"attrs"` then this read will produce elements of the form `Row(payload=..., attrs=...)` where `attrs` is a Map type of string to string. If both `attributes` and `attribute_map` are set, the overlapping attribute values will be present in both the flattened structure and the attribute map. id_attribute: The attribute on incoming Pub/Sub messages to use as a unique record identifier. When specified, the value of this attribute (which can be any string that uniquely identifies the record) will be used for deduplication of messages. If not provided, we cannot guarantee that no duplicate data will be delivered on the Pub/Sub stream. In this case, deduplication of the stream will be strictly best effort. timestamp_attribute: Message value to use as element timestamp. If None, uses message publishing time as the timestamp. Timestamp values should be in one of two formats: - A numerical value representing the number of milliseconds since the Unix epoch. - A string in RFC 3339 format, UTC timezone. Example: ``2015-10-29T23:41:41.123Z``. The sub-second component of the timestamp is optional, and digits beyond the first three (i.e., time units smaller than milliseconds) may be ignored. """ if topic and subscription: raise TypeError('Only one of topic and subscription may be specified.') elif not topic and not subscription: raise TypeError('One of topic or subscription may be specified.') payload_schema, parser = _create_parser(format, schema) extra_fields: list[schema_pb2.Field] = [] if not attributes and not attributes_map: mapper = lambda msg: parser(msg) else: if isinstance(attributes, str): attributes = [attributes] if attributes: extra_fields.extend( [schemas.schema_field(attr, str) for attr in attributes]) if attributes_map: extra_fields.append( schemas.schema_field(attributes_map, Mapping[str, str])) def mapper(msg): values = parser(msg.data).as_dict() if attributes: # Should missing attributes be optional or parse errors? for attr in attributes: values[attr] = msg.attributes[attr] if attributes_map: values[attributes_map] = msg.attributes return beam.Row(**values) output = ( root | beam.io.ReadFromPubSub( topic=topic, subscription=subscription, with_attributes=bool(attributes or attributes_map), id_label=id_attribute, timestamp_attribute=timestamp_attribute) | 'ParseMessage' >> beam.Map(mapper)) output.element_type = schemas.named_tuple_from_schema( schema_pb2.Schema(fields=list(payload_schema.fields) + extra_fields)) return output @beam.ptransform_fn @yaml_errors.maybe_with_exception_handling_transform_fn def write_to_pubsub( pcoll, *, topic: str, format: str, schema: Optional[Any] = None, attributes: Optional[Iterable[str]] = None, attributes_map: Optional[str] = None, id_attribute: Optional[str] = None, timestamp_attribute: Optional[str] = None): """Writes messages to Cloud Pub/Sub. Args: topic: Cloud Pub/Sub topic in the form "/topics//". format: How to format the message payload. Currently supported formats are - RAW: Expects a message with a single field (excluding attribute-related fields) whose contents are used as the raw bytes of the pubsub message. - AVRO: Encodes records with a given Avro schema, which may be inferred from the input PCollection schema. - JSON: Formats records with a given JSON schema, which may be inferred from the input PCollection schema. - PROTO: Encodes records with a given Protobuf schema, which may be inferred from the input PCollection schema. schema: Schema specification for the given format. attributes: List of attribute keys whose values will be pulled out as PubSub message attributes. For example, if the format is `raw` and attributes is `["a", "b"]` then elements of the form `Row(any_field=..., a=..., b=...)` will result in PubSub messages whose payload has the contents of any_field and whose attribute will be populated with the values of `a` and `b`. attribute_map: Name of a string-to-string map field in which to pull a set of attributes associated with this message. For example, if the format is `raw` and `attribute_map` is set to `"attrs"` then elements of the form `Row(any_field=..., attrs=...)` will result in PubSub messages whose payload has the contents of any_field and whose attribute will be populated with the values from attrs. If both `attributes` and `attribute_map` are set, the union of attributes from these two sources will be used to populate the PubSub message attributes. id_attribute: If set, will set an attribute for each Cloud Pub/Sub message with the given name and a unique value. This attribute can then be used in a ReadFromPubSub PTransform to deduplicate messages. timestamp_attribute: If set, will set an attribute for each Cloud Pub/Sub message with the given name and the message's publish time as the value. """ input_schema = schemas.schema_from_element_type(pcoll.element_type) extra_fields: list[str] = [] if isinstance(attributes, str): attributes = [attributes] if attributes: extra_fields.extend(attributes) if attributes_map: extra_fields.append(attributes_map) def attributes_extractor(row): if attributes_map: attribute_values = dict(getattr(row, attributes_map)) else: attribute_values = {} if attributes: attribute_values.update({attr: getattr(row, attr) for attr in attributes}) return attribute_values schema_names = set(f.name for f in input_schema.fields) missing_attribute_names = set(extra_fields) - schema_names if missing_attribute_names: raise ValueError( f'Attribute fields {missing_attribute_names} ' f'not found in schema fields {schema_names}') payload_schema = schema_pb2.Schema( fields=[ field for field in input_schema.fields if field.name not in extra_fields ]) formatter = _create_formatter(format, schema, payload_schema) return ( pcoll | beam.Map( lambda row: beam.io.gcp.pubsub.PubsubMessage( formatter(row), attributes_extractor(row))) | beam.io.WriteToPubSub( topic, with_attributes=True, id_label=id_attribute, timestamp_attribute=timestamp_attribute)) def read_from_iceberg( table: str, filter: Optional[str] = None, keep: Optional[list[str]] = None, drop: Optional[list[str]] = None, catalog_name: Optional[str] = None, catalog_properties: Optional[Mapping[str, str]] = None, config_properties: Optional[Mapping[str, str]] = None, ): # TODO(robertwb): It'd be nice to derive this list of parameters, along with # their types and docs, programmatically from the iceberg (or managed) # schemas. """Reads an Apache Iceberg table. See also the [Apache Iceberg Beam documentation]( https://cloud.google.com/dataflow/docs/guides/managed-io#iceberg). Args: table: The identifier of the Apache Iceberg table. Example: "db.table1". filter: SQL-like predicate to filter data at scan time. Example: "id > 5 AND status = 'ACTIVE'". Uses Apache Calcite syntax: https://calcite.apache.org/docs/reference.html keep: A subset of column names to read exclusively. If null or empty, all columns will be read. drop: A subset of column names to exclude from reading. If null or empty, all columns will be read. catalog_name: The name of the catalog. Example: "local". catalog_properties: A map of configuration properties for the Apache Iceberg catalog. The required properties depend on the catalog. For more information, see CatalogUtil in the Apache Iceberg documentation. config_properties: An optional set of Hadoop configuration properties. For more information, see CatalogUtil in the Apache Iceberg documentation. """ return beam.managed.Read( "iceberg", config=dict( table=table, catalog_name=catalog_name, filter=filter, keep=keep, drop=drop, catalog_properties=catalog_properties, config_properties=config_properties)) def write_to_iceberg( table: str, catalog_name: Optional[str] = None, catalog_properties: Optional[Mapping[str, str]] = None, config_properties: Optional[Mapping[str, str]] = None, partition_fields: Optional[Iterable[str]] = None, table_properties: Optional[Mapping[str, str]] = None, triggering_frequency_seconds: Optional[int] = None, keep: Optional[Iterable[str]] = None, drop: Optional[Iterable[str]] = None, only: Optional[str] = None, distribution_mode: Optional[str] = None, autosharding: Optional[bool] = None, ): # TODO(robertwb): It'd be nice to derive this list of parameters, along with # their types and docs, programmatically from the iceberg (or managed) # schemas. """Writes to an Apache Iceberg table. See also the [Apache Iceberg Beam documentation]( https://cloud.google.com/dataflow/docs/guides/managed-io#iceberg) including the [dynamic destinations section]( https://cloud.google.com/dataflow/docs/guides/managed-io#dynamic-destinations) for use of the keep, drop, and only parameters. Args: table: The identifier of the Apache Iceberg table. Example: "db.table1". catalog_name: The name of the catalog. Example: "local". catalog_properties: A map of configuration properties for the Apache Iceberg catalog. The required properties depend on the catalog. For more information, see CatalogUtil in the Apache Iceberg documentation. config_properties: An optional set of Hadoop configuration properties. For more information, see CatalogUtil in the Apache Iceberg documentation. partition_fields: Fields used to create a partition spec that is applied when tables are created. For a field 'foo', the available partition transforms are: - foo - truncate(foo, N) - bucket(foo, N) - hour(foo) - day(foo) - month(foo) - year(foo) - void(foo) For more information on partition transforms, please visit https://iceberg.apache.org/spec/#partition-transforms. table_properties: Iceberg table properties to be set on the table when it is created. For more information on table properties, please visit https://iceberg.apache.org/docs/latest/configuration/#table-properties. triggering_frequency_seconds: For streaming write pipelines, the frequency at which the sink attempts to produce snapshots, in seconds. keep: An optional list of field names to keep when writing to the destination. Other fields are dropped. Mutually exclusive with drop and only. drop: An optional list of field names to drop before writing to the destination. Mutually exclusive with keep and only. only: The name of exactly one field to keep as the top level record when writing to the destination. All other fields are dropped. This field must be of row type. Mutually exclusive with drop and keep. distribution_mode: Defines distribution of write data. Supported distributions: - none: don't shuffle rows (default) - hash: shuffle rows by partition key before writing data autosharding: Enables dynamic sharding to automatically adjust the number of parallel writers based on data volume. It handles data skew by further sub-dividing partitions into multiple shards to prevent bottlenecks during high-throughput writes. Only available with 'hash' distribution mode. """ return beam.managed.Write( "iceberg", config=dict( table=table, catalog_name=catalog_name, catalog_properties=catalog_properties, config_properties=config_properties, partition_fields=partition_fields, table_properties=table_properties, triggering_frequency_seconds=triggering_frequency_seconds, keep=keep, drop=drop, only=only, distribution_mode=distribution_mode, autosharding=autosharding)) def io_providers(): return yaml_provider.load_providers( yaml_utils.locate_data_file('standard_io.yaml')) def read_from_tfrecord( file_pattern: str, coder: Optional[coders.BytesCoder] = coders.BytesCoder(), compression_type: str = "AUTO", validate: Optional[bool] = True): """Reads data from TFRecord. Args: file_pattern (str): A file glob pattern to read TFRecords from. coder (coders.BytesCoder): Coder used to decode each record. compression_type (CompressionTypes): Used to handle compressed input files. Default value is CompressionTypes.AUTO, in which case the file_path's extension will be used to detect the compression. validate (bool): Boolean flag to verify that the files exist during the pipeline creation time. """ return ReadFromTFRecord( file_pattern=file_pattern, compression_type=getattr(CompressionTypes, compression_type), validate=validate) | beam.Map(lambda s: beam.Row(record=s)) @beam.ptransform_fn def write_to_tfrecord( pcoll, file_path_prefix: str, coder: Optional[coders.BytesCoder] = coders.BytesCoder(), file_name_suffix: Optional[str] = "", num_shards: Optional[int] = 0, shard_name_template: Optional[str] = None, compression_type: str = "AUTO"): """Writes data to TFRecord. Args: file_path_prefix: The file path to write to. The files written will begin with this prefix, followed by a shard identifier (see num_shards), and end in a common extension, if given by file_name_suffix. coder: Coder used to encode each record. file_name_suffix: Suffix for the files written. num_shards: The number of files (shards) used for output. If not set, the default value will be used. shard_name_template: A template string containing placeholders for the shard number and shard count. When constructing a filename for a particular shard number, the upper-case letters 'S' and 'N' are replaced with the 0-padded shard number and shard count respectively. This argument can be '' in which case it behaves as if num_shards was set to 1 and only one file will be generated. The default pattern used is '-SSSSS-of-NNNNN' if None is passed as the shard_name_template. compression_type: Used to handle compressed output files. Typical value is CompressionTypes.AUTO, in which case the file_path's extension will be used to detect the compression. Returns: A WriteToTFRecord transform object. """ try: field_names = [ name for name, _ in schemas.named_fields_from_element_type( pcoll.element_type) ] except Exception as exn: raise ValueError( "WriteToTFRecord requires an input schema with exactly one field." ) from exn if len(field_names) != 1: raise ValueError( "WriteToTFRecord requires an input schema with exactly one field,got %s" % field_names) sole_field_name, = field_names return pcoll | beam.Map( lambda x: getattr(x, sole_field_name)) | WriteToTFRecord( file_path_prefix=file_path_prefix, coder=coder, file_name_suffix=file_name_suffix, num_shards=num_shards, shard_name_template=shard_name_template, compression_type=getattr(CompressionTypes, compression_type)) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_io_test.py0000644000175100017510000005013415206113033021665 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import io import json import logging import unittest import fastavro import mock import apache_beam as beam from apache_beam.coders.row_coder import RowCoder from apache_beam.io.gcp.pubsub import PubsubMessage from apache_beam.testing.util import AssertThat from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.typehints import schemas as schema_utils from apache_beam.yaml.yaml_transform import YamlTransform try: import jsonschema except ImportError: jsonschema = None class FakeReadFromPubSub: def __init__( self, topic, messages, subscription=None, id_attribute=None, timestamp_attribute=None): self._topic = topic self._subscription = subscription self._messages = messages self._id_attribute = id_attribute self._timestamp_attribute = timestamp_attribute def __call__( self, *, topic, subscription, with_attributes, id_label, timestamp_attribute): assert topic == self._topic assert id_label == self._id_attribute assert timestamp_attribute == self._timestamp_attribute assert subscription == self._subscription if with_attributes: data = self._messages else: data = [x.data for x in self._messages] return beam.Create(data) class FakeWriteToPubSub: def __init__( self, topic, messages, id_attribute=None, timestamp_attribute=None): self._topic = topic self._messages = messages self._id_attribute = id_attribute self._timestamp_attribute = timestamp_attribute def __call__(self, topic, *, with_attributes, id_label, timestamp_attribute): assert topic == self._topic assert with_attributes is True assert id_label == self._id_attribute assert timestamp_attribute == self._timestamp_attribute return AssertThat(equal_to(self._messages)) @unittest.skipIf(jsonschema is None, "Yaml dependencies not installed") class YamlPubSubTest(unittest.TestCase): def test_simple_read(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch('apache_beam.io.ReadFromPubSub', FakeReadFromPubSub( topic='my_topic', messages=[PubsubMessage(b'msg1', {'attr': 'value1'}), PubsubMessage(b'msg2', {'attr': 'value2'})])): result = p | YamlTransform( ''' type: ReadFromPubSub config: topic: my_topic format: RAW ''') assert_that( result, equal_to([beam.Row(payload=b'msg1'), beam.Row(payload=b'msg2')])) def test_simple_read_string(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch('apache_beam.io.ReadFromPubSub', FakeReadFromPubSub( topic='my_topic', messages=[PubsubMessage('äter'.encode('utf-8'), {'attr': 'value1'}), PubsubMessage('köttbullar'.encode('utf-8'), {'attr': 'value2'})])): result = p | YamlTransform( ''' type: ReadFromPubSub config: topic: my_topic format: STRING ''') assert_that( result, equal_to([beam.Row(payload='äter'), beam.Row(payload='köttbullar')])) def test_read_with_attribute(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch('apache_beam.io.ReadFromPubSub', FakeReadFromPubSub( topic='my_topic', messages=[PubsubMessage(b'msg1', {'attr': 'value1'}), PubsubMessage(b'msg2', {'attr': 'value2'})])): result = p | YamlTransform( ''' type: ReadFromPubSub config: topic: my_topic format: RAW attributes: [attr] ''') assert_that( result, equal_to([ beam.Row(payload=b'msg1', attr='value1'), beam.Row(payload=b'msg2', attr='value2') ])) def test_read_with_attribute_map(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch('apache_beam.io.ReadFromPubSub', FakeReadFromPubSub( topic='my_topic', messages=[PubsubMessage(b'msg1', {'attr': 'value1'}), PubsubMessage(b'msg2', {'attr': 'value2'})])): result = p | YamlTransform( ''' type: ReadFromPubSub config: topic: my_topic format: RAW attributes_map: attrMap ''') assert_that( result, equal_to([ beam.Row(payload=b'msg1', attrMap={'attr': 'value1'}), beam.Row(payload=b'msg2', attrMap={'attr': 'value2'}) ])) def test_read_with_id_attribute(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch('apache_beam.io.ReadFromPubSub', FakeReadFromPubSub( topic='my_topic', messages=[PubsubMessage(b'msg1', {'attr': 'value1'}), PubsubMessage(b'msg2', {'attr': 'value2'})], id_attribute='some_attr')): result = p | YamlTransform( ''' type: ReadFromPubSub config: topic: my_topic format: RAW id_attribute: some_attr ''') assert_that( result, equal_to([beam.Row(payload=b'msg1'), beam.Row(payload=b'msg2')])) _avro_schema = { 'type': 'record', 'name': 'ec', 'fields': [{ 'name': 'label', 'type': 'string' }, { 'name': 'rank', 'type': 'int' }] } def _encode_avro(self, data): buffer = io.BytesIO() fastavro.schemaless_writer(buffer, self._avro_schema, data) buffer.seek(0) return buffer.read() def test_read_avro(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch( 'apache_beam.io.ReadFromPubSub', FakeReadFromPubSub( topic='my_topic', messages=[PubsubMessage(self._encode_avro({'label': '37a', 'rank': 1}), {}), PubsubMessage(self._encode_avro({'label': '389a', 'rank': 2}), {})])): result = p | YamlTransform( ''' type: ReadFromPubSub config: topic: my_topic format: AVRO schema: %s ''' % json.dumps(self._avro_schema)) assert_that( result, equal_to([ beam.Row(label='37a', rank=1), # linebreak beam.Row(label='389a', rank=2) ])) def test_read_json(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch('apache_beam.io.ReadFromPubSub', FakeReadFromPubSub( topic='my_topic', messages=[PubsubMessage( b'{"generator": {"x": 0, "y": 0}, "rank": 1}', {'weierstrass': 'y^2+y=x^3-x', 'label': '37a'}) ])): result = p | YamlTransform( ''' type: ReadFromPubSub config: topic: my_topic format: JSON schema: type: object properties: generator: type: object properties: x: {type: integer} y: {type: integer} rank: {type: integer} attributes: [label] attributes_map: other ''') assert_that( result, equal_to([ beam.Row( generator=beam.Row(x=0, y=0), rank=1, label='37a', other={ 'label': '37a', 'weierstrass': 'y^2+y=x^3-x' }) ])) def test_read_json_with_error_handling(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch( 'apache_beam.io.ReadFromPubSub', FakeReadFromPubSub(topic='my_topic', messages=[PubsubMessage('{"some_int": 123}', attributes={}), PubsubMessage('unparsable', attributes={})])): result = p | YamlTransform( ''' type: ReadFromPubSub config: topic: my_topic format: JSON schema: type: object properties: some_int: {type: integer} error_handling: output: errors ''') assert_that( result['good'], equal_to([beam.Row(some_int=123)]), label='CheckGood') assert_that( result['errors'] | beam.Map(lambda error: error.element), equal_to(['unparsable']), label='CheckErrors') def test_read_json_without_error_handling(self): with self.assertRaises(Exception): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch( 'apache_beam.io.ReadFromPubSub', FakeReadFromPubSub(topic='my_topic', messages=[PubsubMessage('{"some_int": 123}', attributes={}), PubsubMessage('unparsable', attributes={})])): _ = p | YamlTransform( ''' type: ReadFromPubSub config: topic: my_topic format: JSON schema: type: object properties: some_int: {type: integer} ''') def test_read_json_with_bad_schema(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch('apache_beam.io.ReadFromPubSub', FakeReadFromPubSub( topic='my_topic', messages=[PubsubMessage('{"some_int": 123}', attributes={}), PubsubMessage('{"some_int": "NOT"}', attributes={})])): result = p | YamlTransform( ''' type: ReadFromPubSub config: topic: my_topic format: JSON schema: type: object properties: some_int: {type: integer} error_handling: output: errors ''') assert_that( result['good'], equal_to([beam.Row(some_int=123)]), label='CheckGood') assert_that( result['errors'] | beam.Map(lambda error: error.element), equal_to(['{"some_int": "NOT"}']), label='CheckErrors') def test_simple_write(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch('apache_beam.io.WriteToPubSub', FakeWriteToPubSub(topic='my_topic', messages=[PubsubMessage(b'msg1', {}), PubsubMessage(b'msg2', {})])): _ = ( p | beam.Create([beam.Row(a=b'msg1'), beam.Row(a=b'msg2')]) | YamlTransform( ''' type: WriteToPubSub config: topic: my_topic format: RAW ''')) def test_write_with_attribute(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch('apache_beam.io.WriteToPubSub', FakeWriteToPubSub( topic='my_topic', messages=[PubsubMessage(b'msg1', {'attr': 'value1'}), PubsubMessage(b'msg2', {'attr': 'value2'})])): _ = ( p | beam.Create([ beam.Row(a=b'msg1', attr='value1'), beam.Row(a=b'msg2', attr='value2') ]) | YamlTransform( ''' type: WriteToPubSub config: topic: my_topic format: RAW attributes: [attr] ''')) def test_write_with_attribute_map(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch('apache_beam.io.WriteToPubSub', FakeWriteToPubSub(topic='my_topic', messages=[PubsubMessage(b'msg1', {'a': 'b'}), PubsubMessage(b'msg2', {'c': 'd'})])): _ = ( p | beam.Create([ beam.Row(a=b'msg1', attrMap={'a': 'b'}), beam.Row(a=b'msg2', attrMap={'c': 'd'}) ]) | YamlTransform( ''' type: WriteToPubSub config: topic: my_topic format: RAW attributes_map: attrMap ''')) def test_write_with_id_attribute(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch('apache_beam.io.WriteToPubSub', FakeWriteToPubSub(topic='my_topic', messages=[PubsubMessage(b'msg1', {}), PubsubMessage(b'msg2', {})], id_attribute='some_attr')): _ = ( p | beam.Create([beam.Row(a=b'msg1'), beam.Row(a=b'msg2')]) | YamlTransform( ''' type: WriteToPubSub config: topic: my_topic format: RAW id_attribute: some_attr ''')) def test_write_avro(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch( 'apache_beam.io.WriteToPubSub', FakeWriteToPubSub( topic='my_topic', messages=[PubsubMessage(self._encode_avro({'label': '37a', 'rank': 1}), {}), PubsubMessage(self._encode_avro({'label': '389a', 'rank': 2}), {})])): _ = ( p | beam.Create( [beam.Row(label='37a', rank=1), beam.Row(label='389a', rank=2)]) | YamlTransform( ''' type: WriteToPubSub config: topic: my_topic format: AVRO ''')) def test_write_json(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with mock.patch('apache_beam.io.WriteToPubSub', FakeWriteToPubSub( topic='my_topic', messages=[PubsubMessage( b'{"generator": {"x": 0, "y": 0}, "rank": 1}', {'weierstrass': 'y^2+y=x^3-x', 'label': '37a'}) ])): _ = ( p | beam.Create([ beam.Row( label='37a', generator=beam.Row(x=0, y=0), rank=1, other={'weierstrass': 'y^2+y=x^3-x'}) ]) | YamlTransform( ''' type: WriteToPubSub config: topic: my_topic format: JSON attributes: [label] attributes_map: other ''')) def test_write_proto(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: data = [beam.Row(label='37a', rank=1), beam.Row(label='389a', rank=2)] coder = RowCoder( schema_utils.named_fields_to_schema([('label', str), ('rank', int)])) expected_messages = [PubsubMessage(coder.encode(r), {}) for r in data] with mock.patch('apache_beam.io.WriteToPubSub', FakeWriteToPubSub(topic='my_topic', messages=expected_messages)): _ = ( p | beam.Create(data) | YamlTransform( ''' type: WriteToPubSub config: topic: my_topic format: PROTO ''')) def test_read_proto(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: data = [beam.Row(label='37a', rank=1), beam.Row(label='389a', rank=2)] coder = RowCoder( schema_utils.named_fields_to_schema([('label', str), ('rank', int)])) expected_messages = [PubsubMessage(coder.encode(r), {}) for r in data] with mock.patch('apache_beam.io.ReadFromPubSub', FakeReadFromPubSub(topic='my_topic', messages=expected_messages)): result = p | YamlTransform( ''' type: ReadFromPubSub config: topic: my_topic format: PROTO schema: type: object properties: label: {type: string} rank: {type: integer} ''') assert_that(result, equal_to(data)) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_join.py0000644000175100017510000002533515206113033021163 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """This module defines the Join operation.""" from typing import Any from typing import Optional from typing import Union import apache_beam as beam from apache_beam.yaml import yaml_provider def _validate_input(pcolls): error_prefix = f'Invalid input {pcolls} specified.' if not isinstance(pcolls, dict): raise ValueError(f'{error_prefix} It must be a dict.') if len(pcolls) < 2: raise ValueError( f'{error_prefix} There should be at least 2 inputs to join.') def _validate_type(type, pcolls): error_prefix = f'Invalid value "{type}" for "type".' if not isinstance(type, dict) and not isinstance(type, str): raise ValueError(f'{error_prefix} It must be a dict or a str.') if isinstance(type, dict): error = ValueError( f'{error_prefix} When specifying a dict for type, ' f'it must follow this format: ' f'{{"outer": [list of inputs to outer join]}}. ' f'Example: {{"outer": ["input1", "input2"]}}') if (len(type) != 1 or next(iter(type)) != 'outer' or not isinstance(type['outer'], list)): raise error for input in type['outer']: if input not in list(pcolls.keys()): raise ValueError( f'{error_prefix} An invalid input "{input}" was specified.') if isinstance(type, str) and type not in ('inner', 'outer', 'left', 'right'): raise ValueError( f'{error_prefix} When specifying the value for type as a str, ' f'it must be one of the following: "inner", "outer", "left", "right"') def _validate_equalities(equalities, pcolls): error_prefix = f'Invalid value "{equalities}" for "equalities".' valid_cols = { name: set( dict(fields).keys() if fields and all( isinstance(field, tuple) for field in fields) else fields) for (name, pcoll) in pcolls.items() for fields in [getattr(pcoll.element_type, '_fields', [])] } if isinstance(equalities, str): for cols in valid_cols.values(): if equalities not in cols: raise ValueError( f'{error_prefix} When "equalities" is a str, ' f'it must be a field name that exists in all the specified inputs.') equality = {pcoll_tag: equalities for pcoll_tag in pcolls} return [equality] if not isinstance(equalities, list): raise ValueError(f'{error_prefix} It should be a str or a list.') input_edge_list = [] for equality in equalities: invalid_dict_error = ValueError( f'{error_prefix} {equality} ' f'should be a dict[str, str] containing at least 2 items.') if not isinstance(equality, dict): raise invalid_dict_error if len(equality) < 2: raise invalid_dict_error for pcoll_tag, col in equality.items(): if pcoll_tag not in pcolls: raise ValueError( f'{error_prefix} "{pcoll_tag}" is not a specified alias in "input"') if col not in valid_cols[pcoll_tag]: raise ValueError( f'{error_prefix} "{col}" is not a valid field in "{pcoll_tag}".') input_edge_list.append(tuple(equality.keys())) if not _is_connected(input_edge_list, len(pcolls)): raise ValueError( f'{error_prefix} ' f'The provided equalities do not connect all of {list(pcolls.keys())}.') return equalities def _parse_fields(tables, fields): error_prefix = f'Invalid value "{fields}" for "fields".' if not isinstance(fields, dict): raise ValueError(f'{error_prefix} Fields must be a dict.') output_fields = [] named_columns = set() for input, cols in fields.items(): if input not in tables: raise ValueError(f'An invalid input "{input}" was specified in "fields".') if isinstance(cols, list): for col in cols: if not isinstance(col, str): raise ValueError( f'Invalid column "{col}" in "fields". Column name must be a str.') if col in named_columns: raise ValueError( f'The field name "{col}" was specified more than once.') output_fields.append(f'{input}.{col} AS {col}') named_columns.add(col) elif isinstance(cols, dict): for k, v in cols.items(): if k in named_columns: raise ValueError( f'The field name "{k}" was specified more than once.') if not isinstance(v, str): raise ValueError( f'Invalid column "{v}" in "fields". Column name must be a str.') output_fields.append(f'{input}.{v} AS {k}') named_columns.add(k) else: raise ValueError( f'{error_prefix} ' f'For every input key in fields, ' f'the value must either be a list or dict.') for table in tables: if table not in fields.keys(): output_fields.append(f'{table}.*') return output_fields def _is_connected(edge_list, expected_node_count): graph = {} for edge_set in edge_list: for u in edge_set: if u not in graph: graph[u] = set() for v in edge_set: if u != v: graph[u].add(v) visited = set() stack = [next(iter(graph))] while stack: node = stack.pop() visited.add(node) for neighbor in graph[node]: if neighbor not in visited: stack.append(neighbor) return len(visited) == len(graph) == expected_node_count @beam.ptransform.ptransform_fn def _SqlJoinTransform( pcolls, sql_transform_constructor, *, equalities: Union[str, list[dict[str, str]]], type: Union[str, dict[str, list]] = 'inner', fields: Optional[dict[str, Any]] = None): """Joins two or more inputs using a specified condition. For example:: type: Join input: input1: SomeTransform input2: AnotherTransform input3: YetAnotherTransform config: type: inner equalities: - input1: colA input2: colB - input2: colX input3: colY fields: input1: [colA, colB, colC] input2: {new_name: colB} would perform an inner join on the three inputs satisfying the constraints that `input1.colA = input2.colB` and `input2.colX = input3.colY` emitting rows with `colA`, `colB` and `colC` from `input1`, the values of `input2.colB` as a field called `new_name`, and all the fields from `input3`. Args: type: The type of join. Could be a string value in ["inner", "left", "right", "outer"] that specifies the type of join to be performed. For scenarios with multiple inputs to join where different join types are desired, specify the inputs to be outer joined. For example, ``{outer: [input1, input2]}`` means that `input1` and `input2` will be outer joined using the conditions specified, while other inputs will be inner joined. equalities: The condition to join on. A list of sets of columns that should be equal to fulfill the join condition. For the simple scenario of joining on the same column across all inputs where the column name is the same, one can specify the column name as the equality rather than having to list it for every input. fields: The fields to be outputted. A mapping with the input alias as the key and the list of fields in the input to be outputted. The value in the map can either be a dictionary with the new field name as the key and the original field name as the value (e.g new_field_name: field_name), or a list of the fields to be outputted with their original names (e.g ``[col1, col2, col3]``), or an '*' indicating all fields in the input will be outputted. If not specified, all fields from all inputs will be outputted. """ _validate_input(pcolls) _validate_type(type, pcolls) validate_equalities = _validate_equalities(equalities, pcolls) equalities_in_pairs = [] for equality in validate_equalities: inputs = list(equality.keys()) first_input = inputs[0] for input in inputs[1:]: equalities_in_pairs.append({ first_input: equality[first_input], input: equality[input] }) tables = list(pcolls.keys()) if isinstance(type, dict): outer = type['outer'] elif type == 'outer': outer = tables else: outer = [] first_table = tables[0] conditioned = [first_table] def generate_join_type(left, right): if left in outer and right in outer: return 'FULL' if left in outer: return 'LEFT' if right in outer: return 'RIGHT' if not outer: return type.upper() return 'INNER' prev_table = tables[0] join_conditions = {} for i in range(1, len(tables)): curr_table = tables[i] join_type = generate_join_type(prev_table, curr_table) join_conditions[curr_table] = f' {join_type} JOIN {curr_table}' prev_table = curr_table for equality in equalities_in_pairs: left, right = equality.keys() if left in conditioned and right in conditioned: t = tables[max(tables.index(left), tables.index(right))] join_conditions[t] = ( f'{join_conditions[t]} ' f'AND {left}.{equality[left]} = {right}.{equality[right]}') elif left in conditioned: join_conditions[right] = ( f'{join_conditions[right]} ' f'ON {left}.{equality[left]} = {right}.{equality[right]}') conditioned.append(right) elif right in conditioned: join_conditions[left] = ( f'{join_conditions[left]} ' f'ON {left}.{equality[left]} = {right}.{equality[right]}') conditioned.append(left) else: t = tables[max(tables.index(left), tables.index(right))] join_conditions[t] = ( f'{join_conditions[t]} ' f'ON {left}.{equality[left]} = {right}.{equality[right]}') conditioned.append(t) if fields: selects = ', '.join(_parse_fields(tables, fields)) else: selects = '*' query = f'SELECT {selects} FROM {first_table}' query += ' '.join(condition for condition in join_conditions.values()) return pcolls | sql_transform_constructor(query) def create_join_providers(): return [ yaml_provider.SqlBackedProvider({'Join': _SqlJoinTransform}), ] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_join_test.py0000644000175100017510000001467215206113033022224 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest import apache_beam as beam from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.yaml.yaml_transform import YamlTransform class ToRow(beam.PTransform): def expand(self, pcoll): return pcoll | beam.Map(lambda row: beam.Row(**row._asdict())) FRUITS = [ beam.Row(id=1, name='raspberry'), beam.Row(id=2, name='blackberry'), ] QUANTITIES = [ beam.Row(name='raspberry', quantity=1), beam.Row(name='blackberry', quantity=2), beam.Row(name='blueberry', quantity=3), ] CATEGORIES = [ beam.Row(name='raspberry', category='juicy'), beam.Row(name='blackberry', category='dry'), beam.Row(name='blueberry', category='dry'), beam.Row(name='blueberry', category='juicy'), ] @unittest.skipIf( TestPipeline().get_pipeline_options().view_as(StandardOptions).runner is None, 'Do not run this test on precommit suites.') class YamlJoinTest(unittest.TestCase): def test_basic_join(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: fruits = p | "fruits" >> beam.Create(FRUITS) quantities = p | "quantities" >> beam.Create(QUANTITIES) result = { "fruits": fruits, "quantities": quantities } | YamlTransform( ''' type: Join input: f: fruits q: quantities config: type: inner equalities: - f: name q: name ''') | ToRow() assert_that( result, equal_to([ beam.Row(id=1, name='raspberry', name0='raspberry', quantity=1), beam.Row(id=2, name='blackberry', name0='blackberry', quantity=2) ])) def test_join_three_inputs(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: fruits = p | "fruits" >> beam.Create(FRUITS) quantities = p | "quantities" >> beam.Create(QUANTITIES) categories = p | "categories" >> beam.Create(CATEGORIES) result = { "fruits": fruits, "quantities": quantities, "categories": categories } | YamlTransform( ''' type: Join input: f: fruits q: quantities c: categories config: type: inner equalities: - f: name q: name - f: name c: name ''') | ToRow() assert_that( result, equal_to([ beam.Row( id=1, name='raspberry', name0='raspberry', quantity=1, name1='raspberry', category='juicy'), beam.Row( id=2, name='blackberry', name0='blackberry', quantity=2, name1='blackberry', category='dry') ])) def test_join_with_fields(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: fruits = p | "fruits" >> beam.Create(FRUITS) quantities = p | "quantities" >> beam.Create(QUANTITIES) categories = p | "categories" >> beam.Create(CATEGORIES) result = { "fruits": fruits, "quantities": quantities, "categories": categories } | YamlTransform( ''' type: Join input: f: fruits q: quantities c: categories config: type: inner equalities: - f: name q: name - f: name c: name fields: f: f_id: id q: - name - quantity ''') | ToRow() assert_that( result, equal_to([ beam.Row( f_id=1, name='raspberry', quantity=1, name0='raspberry', category='juicy'), beam.Row( f_id=2, name='blackberry', quantity=2, name0='blackberry', category='dry') ])) def test_join_with_equalities_shorthand(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: fruits = p | "fruits" >> beam.Create(FRUITS) quantities = p | "quantities" >> beam.Create(QUANTITIES) categories = p | "categories" >> beam.Create(CATEGORIES) result = { "fruits": fruits, "quantities": quantities, "categories": categories } | YamlTransform( ''' type: Join input: f: fruits q: quantities c: categories config: type: inner equalities: name ''') | ToRow() assert_that( result, equal_to([ beam.Row( id=1, name='raspberry', name0='raspberry', quantity=1, name1='raspberry', category='juicy'), beam.Row( id=2, name='blackberry', name0='blackberry', quantity=2, name1='blackberry', category='dry') ])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_mapping.py0000644000175100017510000010654515206113033021662 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """This module defines the basic MapToFields operation.""" import itertools import re from collections import abc from collections.abc import Callable from collections.abc import Collection from collections.abc import Iterable from collections.abc import Mapping from typing import Any from typing import NamedTuple from typing import Optional from typing import TypeVar from typing import Union import apache_beam as beam from apache_beam.io.filesystems import FileSystems from apache_beam.portability.api import schema_pb2 from apache_beam.transforms.window import TimestampedValue from apache_beam.typehints import row_type from apache_beam.typehints import schemas from apache_beam.typehints import trivial_inference from apache_beam.typehints import typehints from apache_beam.typehints.native_type_compatibility import convert_to_beam_type from apache_beam.typehints.schemas import named_fields_from_element_type from apache_beam.typehints.schemas import schema_from_element_type from apache_beam.typehints.schemas import typing_from_runner_api from apache_beam.utils import python_callable from apache_beam.utils import windowed_value from apache_beam.utils.timestamp import Timestamp from apache_beam.yaml import json_utils from apache_beam.yaml import options from apache_beam.yaml import yaml_provider from apache_beam.yaml.yaml_errors import exception_handling_args from apache_beam.yaml.yaml_errors import map_errors_to_standard_format from apache_beam.yaml.yaml_errors import maybe_with_exception_handling from apache_beam.yaml.yaml_errors import maybe_with_exception_handling_transform_fn from apache_beam.yaml.yaml_provider import dicts_to_rows # Import js2py package if it exists try: import js2py from js2py.base import JsObjectWrapper except ImportError: js2py = None JsObjectWrapper = object _str_expression_fields = { 'AssignTimestamps': 'timestamp', 'Filter': 'keep', 'Partition': 'by', } def normalize_mapping(spec): """ Normalizes various fields for mapping transforms. """ if spec['type'] == 'MapToFields': config = spec.get('config') if isinstance(config.get('drop'), str): config['drop'] = [config['drop']] for field, value in list(config.get('fields', {}).items()): if isinstance(value, (str, int, float)): config['fields'][field] = {'expression': str(value)} elif spec['type'] in _str_expression_fields: param = _str_expression_fields[spec['type']] config = spec.get('config', {}) if isinstance(config.get(param), (str, int, float)): config[param] = {'expression': str(config.get(param))} return spec def is_literal(expr: str) -> bool: # Some languages have limited integer literal ranges. if re.fullmatch(r'-?\d+?', expr) and -1 << 31 < int(expr) < 1 << 31: return True elif re.fullmatch(r'-?\d+\.\d*', expr): return True elif re.fullmatch(r'"[^\\"]*"', expr): return True else: return False def validate_generic_expression( expr_dict: dict, input_fields: Collection[str], allow_cmp: bool, error_field: str) -> None: if not isinstance(expr_dict, dict): raise ValueError( f"Ambiguous expression type (perhaps missing quoting?): {expr_dict}") if len(expr_dict) != 1 or 'expression' not in expr_dict: raise ValueError( "Missing language specification. " "Must specify a language when using a map with custom logic for %s" % error_field) expr = str(expr_dict['expression']) def is_atomic(expr: str): return is_literal(expr) or expr in input_fields if is_atomic(expr): return if allow_cmp: maybe_cmp = re.fullmatch('(.*)([<>=!]+)(.*)', expr) if maybe_cmp: left, cmp, right = maybe_cmp.groups() if (is_atomic(left.strip()) and is_atomic(right.strip()) and cmp in {'==', '<=', '>=', '<', '>', '!='}): return raise ValueError( "Missing language specification, unknown input fields, " f"or invalid generic expression: {expr}. " f"The given input fields are {input_fields}. " "See https://beam.apache.org/documentation/sdks/yaml-udf/#generic") def validate_generic_expressions(base_type, config, input_pcolls) -> None: if not input_pcolls: return try: input_fields = [ name for (name, _) in named_fields_from_element_type( next(iter(input_pcolls)).element_type) ] except (TypeError, ValueError): input_fields = [] if base_type == 'MapToFields': for field, value in list(config.get('fields', {}).items()): validate_generic_expression(value, input_fields, True, field) elif base_type in _str_expression_fields: param = _str_expression_fields[base_type] validate_generic_expression( config.get(param), input_fields, base_type == 'Filter', param) def _check_mapping_arguments( transform_name, expression=None, callable=None, name=None, path=None): # Argument checking if not expression and not callable and not path and not name: raise ValueError( f'{transform_name} must specify either "expression", "callable", ' f'or both "path" and "name"') if expression and callable: raise ValueError( f'{transform_name} cannot specify both "expression" and "callable"') if (expression or callable) and (path or name): raise ValueError( f'{transform_name} cannot specify "expression" or "callable" with ' f'"path" or "name"') if path and not name: raise ValueError(f'{transform_name} cannot specify "path" without "name"') if name and not path: raise ValueError(f'{transform_name} cannot specify "name" without "path"') # js2py's JsObjectWrapper object has a self-referencing __dict__ property # that cannot be pickled without implementing the __getstate__ and # __setstate__ methods. class _CustomJsObjectWrapper(JsObjectWrapper): def __init__(self, js_obj): super().__init__(js_obj.__dict__['_obj']) def __getstate__(self): return self.__dict__.copy() def __setstate__(self, state): self.__dict__.update(state) # TODO(yaml) Improve type inferencing for JS UDF's def py_value_to_js_dict(py_value): if ((isinstance(py_value, tuple) and hasattr(py_value, '_asdict')) or isinstance(py_value, beam.Row)): py_value = py_value._asdict() if isinstance(py_value, dict): return {key: py_value_to_js_dict(value) for key, value in py_value.items()} elif not isinstance(py_value, str) and isinstance(py_value, abc.Iterable): return [py_value_to_js_dict(value) for value in list(py_value)] else: return py_value # TODO(yaml) Consider adding optional language version parameter to support # ECMAScript 5 and 6 def _expand_javascript_mapping_func( original_fields, expression=None, callable=None, path=None, name=None): # Check for installed js2py package if js2py is None: raise ValueError( "Javascript mapping functions are not supported on" " Python 3.12 or later.") # import remaining js2py objects from js2py import base from js2py.constructors import jsdate from js2py.internals import simplex js_array_type = ( base.PyJsArray, base.PyJsArrayBuffer, base.PyJsInt8Array, base.PyJsUint8Array, base.PyJsUint8ClampedArray, base.PyJsInt16Array, base.PyJsUint16Array, base.PyJsInt32Array, base.PyJsUint32Array, base.PyJsFloat32Array, base.PyJsFloat64Array) def _js_object_to_py_object(obj): if isinstance(obj, (base.PyJsNumber, base.PyJsString, base.PyJsBoolean)): return base.to_python(obj) elif isinstance(obj, js_array_type): return [_js_object_to_py_object(value) for value in obj.to_list()] elif isinstance(obj, jsdate.PyJsDate): return obj.to_utc_dt() elif isinstance(obj, (base.PyJsNull, base.PyJsUndefined)): return None elif isinstance(obj, base.PyJsError): raise RuntimeError(obj['message']) elif isinstance(obj, base.PyJsObject): return { key: _js_object_to_py_object(value['value']) for (key, value) in obj.own.items() } elif isinstance(obj, base.JsObjectWrapper): return _js_object_to_py_object(obj._obj) return obj if expression: source = '\n'.join(['function(__row__) {'] + [ f' {name} = __row__.{name}' for name in original_fields if name in expression ] + [' return (' + expression + ')'] + ['}']) js_func = _CustomJsObjectWrapper(js2py.eval_js(source)) elif callable: js_func = _CustomJsObjectWrapper(js2py.eval_js(callable)) else: if not path.endswith('.js'): raise ValueError(f'File "{path}" is not a valid .js file.') udf_code = FileSystems.open(path).read().decode() js = js2py.EvalJs() js.eval(udf_code) js_func = _CustomJsObjectWrapper(getattr(js, name)) def js_wrapper(row): row_as_dict = py_value_to_js_dict(row) try: js_result = js_func(row_as_dict) except simplex.JsException as exn: raise RuntimeError( f"Error evaluating javascript expression: " f"{exn.mes['message']}") from exn return dicts_to_rows(_js_object_to_py_object(js_result)) return js_wrapper def _expand_python_mapping_func( original_fields, expression=None, callable=None, path=None, name=None): if path and name: if not path.endswith('.py'): raise ValueError(f'File "{path}" is not a valid .py file.') py_file = FileSystems.open(path).read().decode() return python_callable.PythonCallableWithSource.load_from_script( py_file, name) elif expression: # TODO(robertwb): Consider constructing a single callable that takes # the row and returns the new row, rather than invoking (and unpacking) # for each field individually. source = '\n'.join(['def fn(__row__):'] + [ f' {name} = __row__.{name}' for name in original_fields if name in expression ] + [' return (' + expression + ')']) else: source = callable return python_callable.PythonCallableWithSource(source) def _validator(beam_type: schema_pb2.FieldType) -> Callable[[Any], bool]: """Returns a callable converting rows of the given type to Json objects.""" type_info = beam_type.WhichOneof("type_info") if type_info == "atomic_type": if beam_type.atomic_type == schema_pb2.BOOLEAN: return lambda x: isinstance(x, bool) elif beam_type.atomic_type == schema_pb2.INT64: return lambda x: isinstance(x, int) elif beam_type.atomic_type == schema_pb2.DOUBLE: return lambda x: isinstance(x, (int, float)) elif beam_type.atomic_type == schema_pb2.STRING: return lambda x: isinstance(x, str) elif beam_type.atomic_type == schema_pb2.BYTES: return lambda x: isinstance(x, bytes) else: raise ValueError( f'Unknown or unsupported atomic type: {beam_type.atomic_type}') elif type_info == "array_type": element_validator = _validator(beam_type.array_type.element_type) return lambda value: all(element_validator(e) for e in value) elif type_info == "iterable_type": element_validator = _validator(beam_type.iterable_type.element_type) return lambda value: all(element_validator(e) for e in value) elif type_info == "map_type": key_validator = _validator(beam_type.map_type.key_type) value_validator = _validator(beam_type.map_type.value_type) return lambda value: all( key_validator(k) and value_validator(v) for (k, v) in value.items()) elif type_info == "row_type": validators = { field.name: _validator(field.type) for field in beam_type.row_type.schema.fields } return lambda row: all( validator(getattr(row, name)) for (name, validator) in validators.items()) else: raise ValueError(f"Unrecognized type_info: {type_info!r}") def _as_callable_for_pcoll( pcoll, fn_spec: Union[str, dict[str, str]], msg: str, language: Optional[str]): if language == 'javascript': options.YamlOptions.check_enabled(pcoll.pipeline, 'javascript') try: input_schema = dict(named_fields_from_element_type(pcoll.element_type)) except (TypeError, ValueError) as exn: if is_expr(fn_spec): raise ValueError("Can only use expressions on a schema'd input.") from exn input_schema = {} # unused if isinstance(fn_spec, str) and fn_spec in input_schema: return lambda row: getattr(row, fn_spec) else: return _as_callable( list(input_schema.keys()), fn_spec, msg, language, input_schema) def _as_callable(original_fields, expr, transform_name, language, input_schema): if isinstance(expr, str): expr = {'expression': expr} # Extract original type from upstream pcoll when doing simple mappings original_type = input_schema.get(expr.get('expression'), None) if expr in original_fields: language = "python" # TODO(yaml): support an imports parameter # TODO(yaml): support a requirements parameter (possibly at a higher level) if not isinstance(expr, dict): raise ValueError( f"Ambiguous expression type (perhaps missing quoting?): {expr}") explicit_type = expr.pop('output_type', None) _check_mapping_arguments(transform_name, **expr) if language == "javascript": func = _expand_javascript_mapping_func(original_fields, **expr) elif language in ("python", "generic", None): func = _expand_python_mapping_func(original_fields, **expr) else: raise ValueError( f'Unknown language for mapping transform: {language}. ' 'Supported languages are "javascript" and "python."') if explicit_type: if isinstance(explicit_type, str): explicit_type = {'type': explicit_type} beam_type = json_utils.json_type_to_beam_type(explicit_type) validator = _validator(beam_type) @beam.typehints.with_output_types(schemas.typing_from_runner_api(beam_type)) def checking_func(row): result = func(row) if not validator(result): raise TypeError(f'{result} violates schema {explicit_type}') return result return checking_func elif original_type: return beam.typehints.with_output_types( convert_to_beam_type(original_type))( func) else: return func class _StripErrorMetadata(beam.PTransform): """Strips error metadata from outputs returned via error handling. Generally the error outputs for transformations return information about the error encountered (e.g. error messages and tracebacks) in addition to the failing element itself. This transformation attempts to remove that metadata and returns the bad element alone which can be useful for re-processing. For example, in the following pipeline snippet:: - name: MyMappingTransform type: MapToFields input: SomeInput config: language: python fields: ... error_handling: output: errors - name: RecoverOriginalElements type: StripErrorMetadata input: MyMappingTransform.errors the output of `RecoverOriginalElements` will contain exactly those elements from SomeInput that failed to processes (whereas `MyMappingTransform.errors` would contain those elements paired with error information). Note that this relies on the preceding transform actually returning the failing input in a schema'd way. Most built-in transformation follow the correct conventions. """ _ERROR_FIELD_NAMES = ('failed_row', 'element', 'record') def __init__(self): super().__init__(label=None) def expand(self, pcoll): try: existing_fields = { fld.name: fld.type for fld in schema_from_element_type(pcoll.element_type).fields } except TypeError: fld = None else: for fld in self._ERROR_FIELD_NAMES: if fld in existing_fields: break else: raise ValueError( 'The input to this transform does not appear to be an error ' + "output. Expected a schema'd input with a field named " + ' or '.join(repr(fld) for fld in self._ERROR_FIELD_NAMES)) if fld is None: # This handles with_exception_handling() that returns bare tuples. return pcoll | beam.Map(lambda x: x[0]) else: return pcoll | beam.Map(lambda x: getattr(x, fld)).with_output_types( typing_from_runner_api(existing_fields[fld])) class Validate(beam.PTransform): """Validates each element of a PCollection against a json schema. Args: schema: A json schema against which to validate each element. error_handling: Whether and how to handle errors during iteration. If this is not set, invalid elements will fail the pipeline, otherwise invalid elements will be passed to the specified error output along with information about how the schema was invalidated. """ def __init__( self, schema: dict[str, Any], error_handling: Optional[Mapping[str, Any]] = None): self._schema = schema self._exception_handling_args = exception_handling_args(error_handling) @maybe_with_exception_handling def expand(self, pcoll): validator = json_utils.row_validator( schema_from_element_type(pcoll.element_type), self._schema) def invoke_validator(x): validator(x) return x return pcoll | beam.Map(invoke_validator) def with_exception_handling(self, **kwargs): # It's possible there's an error in iteration... self._exception_handling_args = kwargs return self class _Explode(beam.PTransform): """Explodes (aka unnest/flatten) one or more fields producing multiple rows. Given one or more fields of iterable type, produces multiple rows, one for each value of that field. For example, a row of the form `('a', [1, 2, 3])` would expand to `('a', 1)`, `('a', 2')`, and `('a', 3)` when exploded on the second field. This is akin to a `FlatMap` when paired with the MapToFields transform. See more complete documentation on [YAML Mapping Functions](https://beam.apache.org/documentation/sdks/yaml-udf/#flatmap). Args: fields: The list of fields to expand. cross_product: If multiple fields are specified, indicates whether the full cross-product of combinations should be produced, or if the first element of the first field corresponds to the first element of the second field, etc. For example, the row `(['a', 'b'], [1, 2])` would expand to the four rows `('a', 1)`, `('a', 2)`, `('b', 1)`, and `('b', 2)` when `cross_product` is set to `true` but only the two rows `('a', 1)` and `('b', 2)` when it is set to `false`. Only meaningful (and required) if multiple rows are specified. error_handling: Whether and how to handle errors during iteration. """ # pylint: disable=line-too-long def __init__( self, fields: Union[str, Collection[str]], cross_product: Optional[bool] = None, error_handling: Optional[Mapping[str, Any]] = None): if isinstance(fields, str): fields = [fields] if cross_product is None: if len(fields) > 1: raise ValueError( 'cross_product must be specified true or false ' 'when exploding multiple fields') else: # Doesn't matter. cross_product = True self._fields = fields self._cross_product = cross_product # TODO(yaml): # 1. Support standard error handling argument. # 2. Supposedly error_handling parameter is not an accepted parameter when # executing. Needs further investigation. self._exception_handling_args = exception_handling_args(error_handling) @maybe_with_exception_handling def expand(self, pcoll): all_fields = [ x for x, _ in named_fields_from_element_type(pcoll.element_type) ] for field in self._fields: if field not in all_fields: raise ValueError(f'Exploding unknown field "{field}"') to_explode = self._fields def explode_cross_product(base, fields): if fields: copy = dict(base) for value in base[fields[0]]: copy[fields[0]] = value yield from explode_cross_product(copy, fields[1:]) else: yield beam.Row(**base) def explode_zip(base, fields): to_zip = [base[field] for field in fields] copy = dict(base) for values in itertools.zip_longest(*to_zip, fillvalue=None): for ix, field in enumerate(fields): copy[field] = values[ix] yield beam.Row(**copy) cross_product = self._cross_product return ( pcoll | beam.FlatMap( lambda row: (explode_cross_product if cross_product else explode_zip)({ name: getattr(row, name) for name in all_fields }, to_explode))) def infer_output_type(self, input_type): return row_type.RowTypeConstraint.from_fields([( name, trivial_inference.element_type(typ) if name in self._fields else typ) for (name, typ) in named_fields_from_element_type(input_type)]) def with_exception_handling(self, **kwargs): # It's possible there's an error in iteration... self._exception_handling_args = kwargs return self @beam.ptransform.ptransform_fn @maybe_with_exception_handling_transform_fn def _PyJsFilter( pcoll, keep: Union[str, dict[str, str]], language: Optional[str] = None): """Keeps only records that satisfy the given criteria. See more complete documentation on [YAML Filtering](https://beam.apache.org/documentation/sdks/yaml-udf/#filtering). Args: keep: An expression evaluating to true for those records that should be kept. language: The language of the above expression. Defaults to generic. error_handling: Whether and where to output records that throw errors when the above expressions are evaluated. """ # pylint: disable=line-too-long keep_fn = _as_callable_for_pcoll(pcoll, keep, "keep", language or 'generic') return pcoll | beam.Filter(keep_fn) def is_expr(v): return isinstance(v, str) or (isinstance(v, dict) and 'expression' in v) def normalize_fields(pcoll, fields, drop=(), append=False, language='generic'): try: input_schema = dict(named_fields_from_element_type(pcoll.element_type)) except (TypeError, ValueError) as exn: if drop: raise ValueError("Can only drop fields on a schema'd input.") from exn if append: raise ValueError("Can only append fields on a schema'd input.") from exn elif any(is_expr(x) for x in fields.values()): raise ValueError("Can only use expressions on a schema'd input.") from exn input_schema = {} if drop and not append: raise ValueError("Can only drop fields if append is true.") for name in drop: if name not in input_schema: raise ValueError(f'Dropping unknown field "{name}"') if append: for name in fields: if name in input_schema and name not in drop: raise ValueError( f'Redefinition of field "{name}". ' 'Cannot append a field that already exists in original input.') if append: return input_schema, { **{name: f'`{name}`' if language in ['sql', 'calcite'] else name for name in input_schema.keys() if name not in drop}, **fields } else: return input_schema, fields @beam.ptransform.ptransform_fn @maybe_with_exception_handling_transform_fn def _PyJsMapToFields( pcoll, fields: Mapping[str, Union[str, Mapping[str, str]]], append: Optional[bool] = False, drop: Optional[Iterable[str]] = None, language: Optional[str] = None, dependencies: Optional[Iterable[str]] = None): """Creates records with new fields defined in terms of the input fields. See more complete documentation on [YAML Mapping Functions](https://beam.apache.org/documentation/sdks/yaml-udf/#mapping-functions). Args: fields: The output fields to compute, each mapping to the expression or callable that creates them. append: Whether to append the created fields to the set of fields already present, outputting a union of both the new fields and the original fields for each record. Defaults to False. drop: If `append` is true, enumerates a subset of fields from the original record that should not be kept language: The language used to define (and execute) the expressions and/or callables in `fields`. Defaults to generic. dependencies: An optional list of extra dependencies that are needed for these UDFs. The interpretation of these strings is language-dependent. error_handling: Whether and where to output records that throw errors when the above expressions are evaluated. """ # pylint: disable=line-too-long input_schema, fields = normalize_fields( pcoll, fields, drop or (), append, language=language or 'generic') if language == 'javascript': options.YamlOptions.check_enabled(pcoll.pipeline, 'javascript') original_fields = list(input_schema.keys()) return pcoll | beam.Select( **{ name: _as_callable( original_fields, expr, name, language, input_schema) for (name, expr) in fields.items() }) @beam.ptransform.ptransform_fn def _SqlFilterTransform(pcoll, sql_transform_constructor, keep, language): return pcoll | sql_transform_constructor( f"SELECT * FROM PCOLLECTION WHERE {keep.get('expression')}") @beam.ptransform.ptransform_fn def _SqlMapToFieldsTransform(pcoll, sql_transform_constructor, **mapping_args): _, fields = normalize_fields(pcoll, **mapping_args) def extract_expr(name, v): if isinstance(v, str): return v elif 'expression' in v: return v['expression'] else: raise ValueError(f"Only expressions allowed in SQL at {name}.") selects = [ f'({extract_expr(name, expr)}) AS `{name}`' for (name, expr) in fields.items() ] query = "SELECT " + ", ".join(selects) + " FROM PCOLLECTION" return pcoll | sql_transform_constructor(query) @beam.ptransform.ptransform_fn def _Partition( pcoll, by: Union[str, dict[str, str]], outputs: list[str], unknown_output: Optional[str] = None, error_handling: Optional[Mapping[str, Any]] = None, language: str = 'generic'): """Splits an input into several distinct outputs. Each input element will go to a distinct output based on the field or function given in the `by` configuration parameter. Args: by: A field, callable, or expression giving the destination output for this element. Should return a string that is a member of the `outputs` parameter. If `unknown_output` is also set, other returns values are accepted as well, otherwise an error will be raised. outputs: The set of outputs into which this input is being partitioned. unknown_output: (Optional) If set, indicates a destination output for any elements that are not assigned an output listed in the `outputs` parameter. error_handling: (Optional) Whether and how to handle errors during partitioning. language: The language of the `by` expression. """ split_fn = _as_callable_for_pcoll(pcoll, by, 'by', language) try: split_fn_output_type = trivial_inference.infer_return_type( split_fn, [pcoll.element_type]) except (TypeError, ValueError): pass else: if not typehints.is_consistent_with(split_fn_output_type, typehints.Optional[str]): raise ValueError( f'Partition function "{by}" must return a string type ' f'not {split_fn_output_type}') error_output = error_handling['output'] if error_handling else None if error_output in outputs: raise ValueError( f'Error handling output "{error_output}" ' f'cannot be among the listed outputs {outputs}') T = TypeVar('T') def split(element): tag = split_fn(element) if tag is None: tag = unknown_output if not isinstance(tag, str): raise ValueError( f'Returned output name "{tag}" of type {type(tag)} ' f'from "{by}" must be a string.') if tag not in outputs: if unknown_output: tag = unknown_output else: raise ValueError(f'Unknown output name "{tag}" from {by}') return beam.pvalue.TaggedOutput(tag, element) output_set = set(outputs) if unknown_output: output_set.add(unknown_output) if error_output: output_set.add(error_output) mapping_transform = beam.Map(split) if error_output: mapping_transform = mapping_transform.with_exception_handling( **exception_handling_args(error_handling)) else: mapping_transform = mapping_transform.with_outputs(*output_set) splits = pcoll | mapping_transform.with_input_types(T).with_output_types(T) result = {out: getattr(splits, out) for out in output_set} for tag, out in result.items(): if tag != error_output: out.element_type = pcoll.element_type if error_output: result[error_output] = result[error_output] | map_errors_to_standard_format( pcoll.element_type) return result @beam.ptransform.ptransform_fn @maybe_with_exception_handling_transform_fn def _AssignTimestamps( pcoll, timestamp: Union[str, dict[str, str]], language: Optional[str] = None): """Assigns a new timestamp each element of its input. This can be useful when reading records that have the timestamp embedded in them, for example with various file types or other sources that by default set all timestamps to the infinite past. Note that the timestamp should only be set forward, as setting it backwards may not cause it to hold back an already advanced watermark and the data could become droppably late. Args: timestamp: A field, callable, or expression giving the new timestamp. language: The language of the timestamp expression. error_handling: Whether and how to handle errors during timestamp evaluation. """ timestamp_fn = _as_callable_for_pcoll(pcoll, timestamp, 'timestamp', language) T = TypeVar('T') return pcoll | beam.Map(lambda x: TimestampedValue(x, timestamp_fn(x)) ).with_input_types(T).with_output_types(T) class PaneInfoTuple(NamedTuple): is_first: bool is_last: bool timing: str index: int # type: ignore[assignment] nonspeculative_index: int @classmethod def from_pane_info(cls, pane_info): return cls( pane_info.is_first, pane_info.is_last, windowed_value.PaneInfoTiming.to_string(pane_info.timing), pane_info.index, pane_info.nonspeculative_index) _WINDOWING_INFO_TYPES = { 'timestamp': Timestamp, 'window_start': Optional[Timestamp], 'window_end': Timestamp, 'window_string': str, 'window_type': str, 'window_object': Any, 'pane_info': PaneInfoTuple, } _WINDOWING_INFO_EXTRACTORS = { 'timestamp': lambda locals: locals['timestamp'], 'window_start': lambda locals: getattr(locals['window'], 'start', None), 'window_end': lambda locals: locals['window'].end, 'window_string': lambda locals: str(locals['window']), 'window_type': lambda locals: type(locals['window']).__name__, 'window_object': lambda locals: locals['window'], 'pane_info': lambda locals: PaneInfoTuple.from_pane_info( locals['pane_info']), } assert set(_WINDOWING_INFO_TYPES.keys()) == set( _WINDOWING_INFO_EXTRACTORS.keys()) @beam.ptransform.ptransform_fn def _ExtractWindowingInfo( pcoll, fields: Optional[Union[Mapping[str, str], Iterable[str]]] = None): """ Extracts the implicit windowing information from an element and makes it explicit as field(s) in the element itself. The following windowing parameter values are supported: * `timestamp`: The event timestamp of the current element. * `window_start`: The start of the window iff it is an interval window. * `window_end`: The (exclusive) end of the window. * `window_string`: The string representation of the window. * `window_type`: The type of the window as a string. * `winodw_object`: The actual window object itself, as a Java or Python object. * `pane_info`: A schema'd representation of the current pane info, including its index, whether it was the last firing, etc. As a convenience, a list rather than a mapping of fields may be provided, in which case the fields will be named according to the requested values. Args: fields: A mapping of new field names to various windowing parameters, as documented above. If omitted, defaults to `[timestamp, window_start, window_end]`. """ if fields is None: fields = ['timestamp', 'window_start', 'window_end'] if not isinstance(fields, Mapping): if isinstance(fields, Iterable) and not isinstance(fields, str): fields = {fld: fld for fld in fields} else: raise TypeError( 'Fields must be a mapping or iterable of strings, got {fields}') existing_fields = named_fields_from_element_type(pcoll.element_type) new_fields = [] for field, value in fields.items(): if value not in _WINDOWING_INFO_TYPES: raise ValueError( f'{value} is not a valid windowing parameter; ' f'must be one of {list(_WINDOWING_INFO_TYPES.keys())}') elif field in existing_fields: raise ValueError(f'Input schema already has a field named {field}.') else: new_fields.append((field, _WINDOWING_INFO_TYPES[value])) def augment_row( row, timestamp=beam.DoFn.TimestampParam, window=beam.DoFn.WindowParam, pane_info=beam.DoFn.PaneInfoParam): as_dict = row._asdict() for field, value in fields.items(): as_dict[field] = _WINDOWING_INFO_EXTRACTORS[value](locals()) return beam.Row(**as_dict) return pcoll | beam.Map(augment_row).with_output_types( row_type.RowTypeConstraint.from_fields( existing_fields + new_fields)) # type: ignore[operator] def create_mapping_providers(): # These are MetaInlineProviders because their expansion is in terms of other # YamlTransforms, but in a way that needs to be deferred until the input # schema is known. return [ yaml_provider.InlineProvider({ 'AssignTimestamps-python': _AssignTimestamps, 'AssignTimestamps-javascript': _AssignTimestamps, 'AssignTimestamps-generic': _AssignTimestamps, 'Explode': _Explode, 'ExtractWindowingInfo': _ExtractWindowingInfo, 'Filter-python': _PyJsFilter, 'Filter-javascript': _PyJsFilter, 'Filter-generic': _PyJsFilter, 'MapToFields-python': _PyJsMapToFields, 'MapToFields-javascript': _PyJsMapToFields, 'MapToFields-generic': _PyJsMapToFields, 'Partition-python': _Partition, 'Partition-javascript': _Partition, 'Partition-generic': _Partition, 'StripErrorMetadata': _StripErrorMetadata, 'ValidateWithSchema': Validate, }), yaml_provider.SqlBackedProvider({ 'Filter-sql': _SqlFilterTransform, 'Filter-calcite': _SqlFilterTransform, 'MapToFields-sql': _SqlMapToFieldsTransform, 'MapToFields-calcite': _SqlMapToFieldsTransform, }), ] ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_mapping_test.py0000644000175100017510000004610015206113033022707 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import typing import unittest import numpy as np import apache_beam as beam from apache_beam import schema_pb2 from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.typehints import schemas from apache_beam.utils.timestamp import Timestamp from apache_beam.yaml import yaml_mapping from apache_beam.yaml.yaml_transform import YamlTransform try: import jsonschema except ImportError: jsonschema = None DATA = [ beam.Row(label='11a', conductor=11, rank=0), beam.Row(label='37a', conductor=37, rank=1), beam.Row(label='389a', conductor=389, rank=2), ] @unittest.skipIf(jsonschema is None, "Yaml dependencies not installed") class YamlMappingTest(unittest.TestCase): def test_basic(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' type: MapToFields config: language: python fields: label: label isogeny: "label[-1]" ''') assert_that( result, equal_to([ beam.Row(label='11a', isogeny='a'), beam.Row(label='37a', isogeny='a'), beam.Row(label='389a', isogeny='a'), ])) def test_drop(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' type: MapToFields config: fields: {} append: true drop: [conductor] ''') assert_that( result, equal_to([ beam.Row(label='11a', rank=0), beam.Row(label='37a', rank=1), beam.Row(label='389a', rank=2), ])) def test_filter(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( ''' type: Filter config: language: python keep: "rank > 0" ''') assert_that( result | beam.Map(lambda named_tuple: beam.Row(**named_tuple._asdict())), equal_to([ beam.Row(label='37a', conductor=37, rank=1), beam.Row(label='389a', conductor=389, rank=2), ])) def test_explode(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([ beam.Row(a=2, b='abc', c=.25), beam.Row(a=3, b='xy', c=.125), ]) result = elements | YamlTransform( ''' type: chain transforms: - type: MapToFields config: language: python append: true fields: range: "range(a)" - type: Explode config: fields: [range, b] cross_product: true ''') assert_that( result, equal_to([ beam.Row(a=2, b='a', c=.25, range=0), beam.Row(a=2, b='a', c=.25, range=1), beam.Row(a=2, b='b', c=.25, range=0), beam.Row(a=2, b='b', c=.25, range=1), beam.Row(a=2, b='c', c=.25, range=0), beam.Row(a=2, b='c', c=.25, range=1), beam.Row(a=3, b='x', c=.125, range=0), beam.Row(a=3, b='x', c=.125, range=1), beam.Row(a=3, b='x', c=.125, range=2), beam.Row(a=3, b='y', c=.125, range=0), beam.Row(a=3, b='y', c=.125, range=1), beam.Row(a=3, b='y', c=.125, range=2), ])) def test_validate(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([ beam.Row(key='good', small=[5], nested=beam.Row(big=100)), beam.Row(key='bad1', small=[500], nested=beam.Row(big=100)), beam.Row(key='bad2', small=[5], nested=beam.Row(big=1)), ]) result = elements | YamlTransform( ''' type: ValidateWithSchema config: schema: type: object properties: small: type: array items: type: integer maximum: 10 nested: type: object properties: big: type: integer minimum: 10 error_handling: output: bad ''') assert_that( result['good'] | beam.Map(lambda x: x.key), equal_to(['good'])) assert_that( result['bad'] | beam.Map(lambda x: x.element.key), equal_to(['bad1', 'bad2']), label='Errors') def test_validate_explicit_types(self): with self.assertRaisesRegex(Exception, r'.*violates schema.*'): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([ beam.Row(a=2, b='abc', c=.25), beam.Row(a=3, b='xy', c=.125), ]) result = elements | YamlTransform( ''' type: MapToFields input: input config: language: python fields: bad: expression: "a + c" output_type: string # This is a lie. ''') self.assertEqual(result.element_type._fields[0][1], str) def test_partition(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([ beam.Row(element='apple'), beam.Row(element='banana'), beam.Row(element='orange'), ]) result = elements | YamlTransform( ''' type: Partition input: input config: by: "'even' if len(element) % 2 == 0 else 'odd'" language: python outputs: [even, odd] ''') self.assertEqual(result['even'].element_type, elements.element_type) assert_that( result['even'] | beam.Map(lambda x: x.element), equal_to(['banana', 'orange']), label='Even') assert_that( result['odd'] | beam.Map(lambda x: x.element), equal_to(['apple']), label='Odd') def test_partition_callable(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([ beam.Row(element='apple'), beam.Row(element='banana'), beam.Row(element='orange'), ]) result = elements | YamlTransform( ''' type: Partition input: input config: by: callable: "lambda row: 'even' if len(row.element) % 2 == 0 else 'odd'" language: python outputs: [even, odd] ''') assert_that( result['even'] | beam.Map(lambda x: x.element), equal_to(['banana', 'orange']), label='Even') assert_that( result['odd'] | beam.Map(lambda x: x.element), equal_to(['apple']), label='Odd') def test_partition_with_unknown(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([ beam.Row(element='apple'), beam.Row(element='banana'), beam.Row(element='orange'), ]) result = elements | YamlTransform( ''' type: Partition input: input config: by: "element.lower()[0]" language: python outputs: [a, b, c] unknown_output: other ''') assert_that( result['a'] | beam.Map(lambda x: x.element), equal_to(['apple']), label='A') assert_that( result['b'] | beam.Map(lambda x: x.element), equal_to(['banana']), label='B') assert_that( result['c'] | beam.Map(lambda x: x.element), equal_to([]), label='C') assert_that( result['other'] | beam.Map(lambda x: x.element), equal_to(['orange']), label='Other') def test_partition_without_unknown(self): with self.assertRaisesRegex(Exception, r'.*Unknown output name.*"o".*'): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([ beam.Row(element='apple'), beam.Row(element='banana'), beam.Row(element='orange'), ]) _ = elements | YamlTransform( ''' type: Partition input: input config: by: "element.lower()[0]" language: python outputs: [a, b, c] ''') def test_partition_without_unknown_with_error(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([ beam.Row(element='apple'), beam.Row(element='banana'), beam.Row(element='orange'), ]) result = elements | YamlTransform( ''' type: Partition input: input config: by: "element.lower()[0]" language: python outputs: [a, b, c] error_handling: output: unknown ''') assert_that( result['a'] | beam.Map(lambda x: x.element), equal_to(['apple']), label='A') assert_that( result['b'] | beam.Map(lambda x: x.element), equal_to(['banana']), label='B') assert_that( result['c'] | beam.Map(lambda x: x.element), equal_to([]), label='C') assert_that( result['unknown'] | beam.Map(lambda x: x.element.element), equal_to(['orange']), label='Errors') def test_partition_with_actual_error(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([ beam.Row(element='apple'), beam.Row(element='banana'), beam.Row(element='orange'), ]) result = elements | YamlTransform( ''' type: Partition input: input config: by: "element.lower()[5]" language: python outputs: [a, b, c] unknown_output: other error_handling: output: errors ''') assert_that( result['a'] | beam.Map(lambda x: x.element), equal_to(['banana']), label='B') assert_that( result['other'] | beam.Map(lambda x: x.element), equal_to(['orange']), label='Other') # Apple only has 5 letters, resulting in an index error. assert_that( result['errors'] | beam.Map(lambda x: x.element.element), equal_to(['apple']), label='Errors') def test_partition_no_language(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([ beam.Row(element='apple', texture='smooth'), beam.Row(element='banana', texture='smooth'), beam.Row(element='orange', texture='bumpy'), ]) result = elements | YamlTransform( ''' type: Partition input: input config: by: texture outputs: [bumpy, smooth] ''') assert_that( result['bumpy'] | beam.Map(lambda x: x.element), equal_to(['orange']), label='Bumpy') assert_that( result['smooth'] | beam.Map(lambda x: x.element), equal_to(['apple', 'banana']), label='Smooth') def test_partition_bad_static_type(self): with self.assertRaisesRegex( ValueError, r'.*Partition function .*must return a string.*'): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([ beam.Row(element='apple', texture='smooth'), beam.Row(element='banana', texture='smooth'), beam.Row(element='orange', texture='bumpy'), ]) _ = elements | YamlTransform( ''' type: Partition input: input config: by: len(texture) outputs: [bumpy, smooth] language: python ''') def test_partition_bad_runtime_type(self): with self.assertRaisesRegex(Exception, r'Returned output name.*must be a string.*'): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([ beam.Row(element='apple', texture='smooth'), beam.Row(element='banana', texture='smooth'), beam.Row(element='orange', texture='bumpy'), ]) _ = elements | YamlTransform( ''' type: Partition input: input config: by: print(texture) outputs: [bumpy, smooth] language: python ''') def test_append_type_inference(self): p = beam.Pipeline( options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) elements = p | beam.Create(DATA) elements.element_type = schemas.named_tuple_from_schema( schema_pb2.Schema( fields=[ schemas.schema_field('label', str), schemas.schema_field('conductor', int), schemas.schema_field('rank', int) ])) result = elements | YamlTransform( ''' type: MapToFields config: language: python append: true fields: new_label: label ''') self.assertSequenceEqual( result.element_type._fields, (('label', str), ('conductor', np.int64), ('rank', np.int64), ('new_label', str))) def test_extract_windowing_info(self): T = typing.TypeVar('T') with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = ( p | beam.Create( [beam.Row(value=1), beam.Row(value=2), beam.Row(value=11)]) | beam.Map( lambda x: beam.transforms.window.TimestampedValue( x, timestamp=x.value)).with_input_types(T).with_output_types( T) | beam.WindowInto(beam.transforms.window.FixedWindows(10))) result = elements | YamlTransform( ''' type: ExtractWindowingInfo config: fields: timestamp: timestamp window_start: window_start window_end: window_end window_string: window_string window_type: window_type window_object: window_object pane_info_field: pane_info ''') assert_that( result, equal_to([ beam.Row( value=1, timestamp=Timestamp(1), window_start=Timestamp(0), window_end=Timestamp(10), window_string='[0.0, 10.0)', window_type='IntervalWindow', window_object=beam.transforms.window.IntervalWindow(0, 10), pane_info_field=yaml_mapping.PaneInfoTuple( True, True, 'UNKNOWN', 0, 0)), beam.Row( value=2, timestamp=Timestamp(2), window_start=Timestamp(0), window_end=Timestamp(10), window_string='[0.0, 10.0)', window_type='IntervalWindow', window_object=beam.transforms.window.IntervalWindow(0, 10), pane_info_field=yaml_mapping.PaneInfoTuple( True, True, 'UNKNOWN', 0, 0)), beam.Row( value=11, timestamp=Timestamp(11), window_start=Timestamp(10), window_end=Timestamp(20), window_string='[10.0, 20.0)', window_type='IntervalWindow', window_object=beam.transforms.window.IntervalWindow(10, 20), pane_info_field=yaml_mapping.PaneInfoTuple( True, True, 'UNKNOWN', 0, 0)), ])) def test_extract_windowing_info_iterable(self): T = typing.TypeVar('T') with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = ( p | beam.Create( [beam.Row(value=1), beam.Row(value=2), beam.Row(value=11)]) | beam.Map( lambda x: beam.transforms.window.TimestampedValue( x, timestamp=x.value)).with_input_types(T).with_output_types( T)) result = elements | YamlTransform( ''' type: ExtractWindowingInfo config: fields: [timestamp, window_type] ''') assert_that( result, equal_to([ beam.Row( value=1, timestamp=Timestamp(1), window_type='GlobalWindow'), beam.Row( value=2, timestamp=Timestamp(2), window_type='GlobalWindow'), beam.Row( value=11, timestamp=Timestamp(11), window_type='GlobalWindow'), ])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_ml.py0000644000175100017510000005117215206113033020632 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """This module defines yaml wrappings for some ML transforms.""" import logging import pkgutil from collections.abc import Callable from importlib import import_module from typing import Any from typing import Optional import apache_beam as beam from apache_beam.io.filesystems import FileSystems from apache_beam.ml.inference import RunInference from apache_beam.ml.inference.base import KeyedModelHandler from apache_beam.typehints.row_type import RowTypeConstraint from apache_beam.typehints.schemas import named_fields_from_element_type from apache_beam.utils import python_callable from apache_beam.yaml import options from apache_beam.yaml.yaml_utils import SafeLineLoader def _list_submodules(package): """ Lists all submodules within a given package. """ submodules = [] skip_modules = ['base', 'handlers', 'test', 'utils'] for _, module_name, _ in pkgutil.walk_packages( package.__path__, package.__name__ + '.'): if any(skip_name in module_name for skip_name in skip_modules): continue submodules.append(module_name) return submodules _transform_constructors = {} try: from apache_beam.ml.transforms.base import MLTransform # Load all available ML Transform modules for module_name in _list_submodules(beam.ml.transforms): try: module = import_module(module_name) _transform_constructors |= module.__dict__ except ImportError as e: logging.warning( 'Could not load ML transform module %s: %s. Please ' 'install the necessary module dependencies', module_name, e) except ImportError: MLTransform = None # type: ignore class ModelHandlerProvider: handler_types: dict[str, Callable[..., "ModelHandlerProvider"]] = {} def __init__( self, handler, preprocess: Optional[dict[str, str]] = None, postprocess: Optional[dict[str, str]] = None): self._handler = handler self._preprocess_fn = self.parse_processing_transform( preprocess, 'preprocess') or self.default_preprocess_fn() self._postprocess_fn = self.parse_processing_transform( postprocess, 'postprocess') or self.default_postprocess_fn() def inference_output_type(self): return Any @staticmethod def parse_processing_transform(processing_transform, typ): def _parse_config(callable=None, path=None, name=None): if callable and (path or name): raise ValueError( f"Cannot specify 'callable' with 'path' and 'name' for {typ} " f"function.") if path and name: return python_callable.PythonCallableWithSource.load_from_script( FileSystems.open(path).read().decode(), name) elif callable: return python_callable.PythonCallableWithSource(callable) else: raise ValueError( f"Must specify one of 'callable' or 'path' and 'name' for {typ} " f"function.") if processing_transform: if isinstance(processing_transform, dict): return _parse_config(**processing_transform) else: raise ValueError("Invalid model_handler specification.") def underlying_handler(self): return self._handler @staticmethod def default_preprocess_fn(): raise ValueError( 'Model Handler does not implement a default preprocess ' 'method. Please define a preprocessing method using the ' '\'preprocess\' tag. This is required in most cases because ' 'most models will have a different input shape, so the model ' 'cannot generalize how the input Row should be transformed. For ' 'an example preprocess method, see VertexAIModelHandlerJSONProvider') def _preprocess_fn_internal(self): return lambda row: (row, self._preprocess_fn(row)) @staticmethod def default_postprocess_fn(): return lambda x: x def _postprocess_fn_internal(self): return lambda result: (result[0], self._postprocess_fn(result[1])) @staticmethod def validate(model_handler_spec): raise NotImplementedError(type(ModelHandlerProvider)) @classmethod def register_handler_type(cls, type_name): def apply(constructor): cls.handler_types[type_name] = constructor return constructor return apply @classmethod def create_handler(cls, model_handler_spec) -> "ModelHandlerProvider": typ = model_handler_spec['type'] config = model_handler_spec['config'] try: result = cls.handler_types[typ](**config) if not hasattr(result, 'to_json'): result.to_json = lambda: model_handler_spec return result except Exception as exn: raise ValueError( f'Unable to instantiate model handler of type {typ}. {exn}') @ModelHandlerProvider.register_handler_type('VertexAIModelHandlerJSON') class VertexAIModelHandlerJSONProvider(ModelHandlerProvider): def __init__( self, endpoint_id: str, project: str, location: str, preprocess: dict[str, str], postprocess: Optional[dict[str, str]] = None, experiment: Optional[str] = None, network: Optional[str] = None, private: bool = False, invoke_route: Optional[str] = None, min_batch_size: Optional[int] = None, max_batch_size: Optional[int] = None, max_batch_duration_secs: Optional[int] = None): """ ModelHandler for Vertex AI. This Model Handler can be used with RunInference to load a model hosted on VertexAI. Every model that is hosted on VertexAI should have three distinct, required, parameters - `endpoint_id`, `project` and `location`. These parameters tell the Model Handler how to access the model's endpoint so that input data can be sent using an API request, and inferences can be received as a response. This Model Handler also requires a `preprocess` function to be defined. Preprocessing and Postprocessing are described in more detail in the RunInference docs: https://beam.apache.org/releases/yamldoc/current/#runinference Every model will have a unique input, but all requests should be JSON-formatted. For example, most language models such as Llama and Gemma expect a JSON with the key "prompt" (among other optional keys). In Python, JSON can be expressed as a dictionary. For example: :: - type: RunInference config: inference_tag: 'my_inference' model_handler: type: VertexAIModelHandlerJSON config: endpoint_id: 9876543210 project: my-project location: us-east1 preprocess: callable: 'lambda x: {"prompt": x.prompt, "max_tokens": 50}' In the above example, which mimics a call to a Llama 3 model hosted on VertexAI, the preprocess function (in this case a lambda) takes in a Beam Row with a single field, "prompt", and maps it to a dict with the same field. It also specifies an optional parameter, "max_tokens", that tells the model the allowed token size (in this case input + output token size). Args: endpoint_id: the numerical ID of the Vertex AI endpoint to query. project: the GCP project name where the endpoint is deployed. location: the GCP location where the endpoint is deployed. preprocess: A python callable, defined either inline, or using a file, that is invoked on the input row before sending to the model to be loaded by this ModelHandler. This parameter is required by the `VertexAIModelHandlerJSON` ModelHandler. postprocess: A python callable, defined either inline, or using a file, that is invoked on the PredictionResult output by the ModelHandler before parsing into the output Beam Row under the field name defined by the inference_tag. experiment: Experiment label to apply to the queries. See https://cloud.google.com/vertex-ai/docs/experiments/intro-vertex-ai-experiments for more information. network: The full name of the Compute Engine network the endpoint is deployed on; used for private endpoints. The network or subnetwork Dataflow pipeline option must be set and match this network for pipeline execution. Ex: "projects/12345/global/networks/myVPC" private: If the deployed Vertex AI endpoint is private, set to true. Requires a network to be provided as well. invoke_route: The custom route path to use when invoking endpoints with arbitrary prediction routes. When specified, uses `Endpoint.invoke()` instead of `Endpoint.predict()`. The route should start with a forward slash, e.g., "/predict/v1". See https://cloud.google.com/vertex-ai/docs/predictions/use-arbitrary-custom-routes for more information. min_batch_size: The minimum batch size to use when batching inputs. max_batch_size: The maximum batch size to use when batching inputs. max_batch_duration_secs: The maximum amount of time to buffer a batch before emitting; used in streaming contexts. """ try: from apache_beam.ml.inference.vertex_ai_inference import VertexAIModelHandlerJSON except ImportError: raise ValueError( 'Unable to import VertexAIModelHandlerJSON. Please ' 'install gcp dependencies: `pip install apache_beam[gcp]`') _handler = VertexAIModelHandlerJSON( endpoint_id=str(endpoint_id), project=project, location=location, experiment=experiment, network=network, private=private, invoke_route=invoke_route, min_batch_size=min_batch_size, max_batch_size=max_batch_size, max_batch_duration_secs=max_batch_duration_secs) super().__init__(_handler, preprocess, postprocess) @staticmethod def validate(model_handler_spec): pass def inference_output_type(self): return RowTypeConstraint.from_fields([('example', Any), ('inference', Any), ('model_id', Optional[str])]) @beam.ptransform.ptransform_fn def run_inference( pcoll, model_handler: dict[str, Any], inference_tag: Optional[str] = 'inference', inference_args: Optional[dict[str, Any]] = None) -> beam.PCollection[beam.Row]: # pylint: disable=line-too-long """ A transform that takes the input rows, containing examples (or features), for use on an ML model. The transform then appends the inferences (or predictions) for those examples to the input row. A ModelHandler must be passed to the `model_handler` parameter. The ModelHandler is responsible for configuring how the ML model will be loaded and how input data will be passed to it. Every ModelHandler has a config tag, similar to how a transform is defined, where the parameters are defined. For example: :: - type: RunInference config: model_handler: type: ModelHandler config: param_1: arg1 param_2: arg2 ... By default, the RunInference transform will return the input row with a single field appended named by the `inference_tag` parameter ("inference" by default) that contains the inference directly returned by the underlying ModelHandler, after any optional postprocessing. For example, if the input had the following: :: Row(question="What is a car?") The output row would look like: :: Row(question="What is a car?", inference=...) where the `inference` tag can be overridden with the `inference_tag` parameter. However, if one specified the following transform config: :: - type: RunInference config: inference_tag: my_inference model_handler: ... The output row would look like: :: Row(question="What is a car?", my_inference=...) See more complete documentation on the underlying [RunInference](https://beam.apache.org/documentation/ml/inference-overview/) transform. ### Preprocessing input data In most cases, the model will be expecting data in a particular data format, whether it be a Python Dict, PyTorch tensor, etc. However, the outputs of all built-in Beam YAML transforms are Beam Rows. To allow for transforming the Beam Row into a data format the model recognizes, each ModelHandler is equipped with a `preprocessing` parameter for performing necessary data preprocessing. It is possible for a ModelHandler to define a default preprocessing function, but in most cases, one will need to be specified by the caller. For example, using `callable`: :: pipeline: type: chain transforms: - type: Create config: elements: - question: "What is a car?" - question: "Where is the Eiffel Tower located?" - type: RunInference config: model_handler: type: ModelHandler config: param_1: arg1 param_2: arg2 preprocess: callable: 'lambda row: {"prompt": row.question}' ... In the above example, the Create transform generates a collection of two Beam Row elements, each with a single field - "question". The model, however, expects a Python Dict with a single key, "prompt". In this case, we can specify a simple Lambda function (alternatively could define a full function), to map the data. ### Postprocessing predictions It is also possible to define a postprocessing function to postprocess the data output by the ModelHandler. See the documentation for the ModelHandler you intend to use (list defined below under `model_handler` parameter doc). In many cases, before postprocessing, the object will be a [PredictionResult](https://beam.apache.org/releases/pydoc/BEAM_VERSION/apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.PredictionResult). # pylint: disable=line-too-long This type behaves very similarly to a Beam Row and fields can be accessed using dot notation. However, make sure to check the docs for your ModelHandler to see which fields its PredictionResult contains or if it returns a different object altogether. For example: :: - type: RunInference config: model_handler: type: ModelHandler config: param_1: arg1 param_2: arg2 postprocess: callable: | def fn(x: PredictionResult): return beam.Row(x.example, x.inference, x.model_id) ... The above example demonstrates converting the original output data type (in this case it is PredictionResult), and converts to a Beam Row, which allows for easier mapping in a later transform. ### File-based pre/postprocessing functions For both preprocessing and postprocessing, it is also possible to specify a Python UDF (User-defined function) file that contains the function. This is possible by specifying the `path` to the file (local file or GCS path) and the `name` of the function in the file. For example: :: - type: RunInference config: model_handler: type: ModelHandler config: param_1: arg1 param_2: arg2 preprocess: path: gs://my-bucket/path/to/preprocess.py name: my_preprocess_fn postprocess: path: gs://my-bucket/path/to/postprocess.py name: my_postprocess_fn ... Args: model_handler: Specifies the parameters for the respective enrichment_handler in a YAML/JSON format. To see the full set of handler_config parameters, see their corresponding doc pages: - [VertexAIModelHandlerJSON](https://beam.apache.org/releases/pydoc/current/apache_beam.yaml.yaml_ml.VertexAIModelHandlerJSONProvider) # pylint: disable=line-too-long inference_tag: The tag to use for the returned inference. Default is 'inference'. inference_args: Extra arguments for models whose inference call requires extra parameters. Make sure to check the underlying ModelHandler docs to see which args are allowed. """ options.YamlOptions.check_enabled(pcoll.pipeline, 'ML') if not isinstance(model_handler, dict): raise ValueError( 'Invalid model_handler specification. Expected dict but was ' f'{type(model_handler)}.') expected_model_handler_params = {'type', 'config'} given_model_handler_params = set( SafeLineLoader.strip_metadata(model_handler).keys()) extra_params = given_model_handler_params - expected_model_handler_params if extra_params: raise ValueError(f'Unexpected parameters in model_handler: {extra_params}') missing_params = expected_model_handler_params - given_model_handler_params if missing_params: raise ValueError(f'Missing parameters in model_handler: {missing_params}') typ = model_handler['type'] model_handler_provider_type = ModelHandlerProvider.handler_types.get( typ, None) if not model_handler_provider_type: raise NotImplementedError(f'Unknown model handler type: {typ}.') model_handler_provider = ModelHandlerProvider.create_handler(model_handler) model_handler_provider.validate(model_handler['config']) schema = RowTypeConstraint.from_fields( named_fields_from_element_type(pcoll.element_type) + [(str(inference_tag), model_handler_provider.inference_output_type())]) return ( pcoll | RunInference( model_handler=KeyedModelHandler( model_handler_provider.underlying_handler()).with_preprocess_fn( model_handler_provider._preprocess_fn_internal()). with_postprocess_fn( model_handler_provider._postprocess_fn_internal()), inference_args=inference_args) | beam.Map( lambda row: beam.Row( **{ **row[0]._asdict(), str(inference_tag): row[1] })).with_output_types(schema)) def _config_to_obj(spec): if 'type' not in spec: raise ValueError(f"Missing type in ML transform spec {spec}") if 'config' not in spec: raise ValueError(f"Missing config in ML transform spec {spec}") constructor = _transform_constructors.get(spec['type']) if constructor is None: raise ValueError("Unknown ML transform type: %r" % spec['type']) return constructor(**spec['config']) @beam.ptransform.ptransform_fn def ml_transform( pcoll, write_artifact_location: Optional[str] = None, read_artifact_location: Optional[str] = None, transforms: Optional[list[Any]] = None): if MLTransform is None: raise ValueError( 'No MLTransform found. Please install tensorflow-transform or ' 'sentence-transformers to use this transform.') options.YamlOptions.check_enabled(pcoll.pipeline, 'ML') result_ml_transform = MLTransform( write_artifact_location=write_artifact_location, read_artifact_location=read_artifact_location, transforms=[_config_to_obj(t) for t in transforms] if transforms else []) if transforms: embedding_transforms = [ t for t in transforms if t.get('type', '').endswith('Embeddings') ] if embedding_transforms: from apache_beam.typehints import List try: if pcoll.element_type: columns_to_change = { col for t_spec in embedding_transforms for col in t_spec.get('config', {}).get('columns', []) } new_fields = named_fields_from_element_type(pcoll.element_type) final_fields = [ (name, List[float] if name in columns_to_change else typ) for name, typ in new_fields ] output_schema = RowTypeConstraint.from_fields(final_fields) return pcoll | result_ml_transform.with_output_types(output_schema) except TypeError: # If we can't get a schema, just return the result. pass return pcoll | result_ml_transform if MLTransform is not None: ml_transform.__doc__ = MLTransform.__doc__ ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_ml_test.py0000644000175100017510000002133515206113033021667 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import tempfile import unittest import apache_beam as beam from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.yaml.yaml_transform import YamlTransform try: # pylint: disable=wrong-import-order, wrong-import-position, unused-import from apache_beam.ml.transforms import tft except ImportError: raise unittest.SkipTest('tensorflow_transform is not installed.') TRAIN_DATA = [ beam.Row(num=0, text='And God said, Let there be light,'), beam.Row(num=2, text='And there was light'), beam.Row(num=8, text='And God saw the light, that it was good'), ] TEST_DATA = [ beam.Row(num=6, text='And God divided the light from the darkness.'), ] class MLTransformTest(unittest.TestCase): def test_ml_transform(self): ml_opts = beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle', yaml_experimental_features=['ML']) with tempfile.TemporaryDirectory() as tempdir: with beam.Pipeline(options=ml_opts) as p: elements = p | beam.Create(TRAIN_DATA) result = elements | YamlTransform( f''' type: MLTransform config: write_artifact_location: {tempdir} transforms: - type: ScaleTo01 config: columns: [num] - type: ComputeAndApplyVocabulary config: columns: [text] split_string_by_delimiter: ' ,.' ''') assert_that( # Why is this an array, not a scalar? result | beam.Map(lambda x: x.num[0]), equal_to([0, .25, 1])) assert_that( result | beam.Map(lambda x: set(x.text)) | beam.CombineGlobally(lambda xs: set.union(*xs)), equal_to([set(range(13))]), label='CheckVocab') with beam.Pipeline(options=ml_opts) as p: elements = p | beam.Create(TEST_DATA) result = elements | YamlTransform( f''' type: MLTransform config: read_artifact_location: {tempdir} ''') assert_that(result | beam.Map(lambda x: x.num[0]), equal_to([.75])) assert_that( result | beam.Map(lambda x: len(set(x.text))), equal_to([5]), label='CheckVocab') def test_ml_transform_read_with_map_to_fields(self): ml_opts = beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle', yaml_experimental_features=['ML']) with tempfile.TemporaryDirectory() as tempdir: # First, write the artifacts. with beam.Pipeline(options=ml_opts) as p: elements = p | beam.Create(TRAIN_DATA) _ = elements | YamlTransform( f''' type: MLTransform config: write_artifact_location: {tempdir} transforms: - type: ScaleTo01 config: columns: [num] - type: ComputeAndApplyVocabulary config: columns: [text] split_string_by_delimiter: ' ,.' ''') # Now, read the artifacts and use MapToFields. with beam.Pipeline(options=ml_opts) as p: elements = p | beam.Create(TEST_DATA) result = elements | YamlTransform( f''' type: chain transforms: - type: MLTransform config: read_artifact_location: {tempdir} - type: MapToFields config: language: python fields: num_scaled: "num[0]" text_vocab: text ''') def check_row(row): assert row.num_scaled == 0.75 assert len(set(row.text_vocab)) == 5 return row.num_scaled assert_that(result | beam.Map(check_row), equal_to([0.75])) def test_sentence_transformer_embedding(self): SENTENCE_EMBEDDING_DIMENSION = 384 DATA = [{ 'id': 1, 'log_message': "Error in module A" }, { 'id': 2, 'log_message': "Warning in module B" }, { 'id': 3, 'log_message': "Info in module C" }] ml_opts = beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle', yaml_experimental_features=['ML']) with tempfile.TemporaryDirectory() as tempdir: with beam.Pipeline(options=ml_opts) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( f''' type: MLTransform config: write_artifact_location: {tempdir} transforms: - type: SentenceTransformerEmbeddings config: model_name: all-MiniLM-L6-v2 columns: [log_message] ''') # Perform a basic check to ensure that embeddings are generated # and that the dimension of those embeddings is correct. actual_output = result | beam.Map(lambda x: len(x['log_message'])) assert_that( actual_output, equal_to([SENTENCE_EMBEDDING_DIMENSION] * len(DATA))) def test_sentence_transformer_embedding_with_beam_rows(self): SENTENCE_EMBEDDING_DIMENSION = 384 DATA = [ beam.Row(id=1, log_message="Error in module A"), beam.Row(id=2, log_message="Warning in module B"), beam.Row(id=3, log_message="Info in module C"), ] ml_opts = beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle', yaml_experimental_features=['ML']) with tempfile.TemporaryDirectory() as tempdir: with beam.Pipeline(options=ml_opts) as p: elements = p | beam.Create(DATA) result = elements | YamlTransform( f''' type: MLTransform config: write_artifact_location: {tempdir} transforms: - type: SentenceTransformerEmbeddings config: model_name: all-MiniLM-L6-v2 columns: [log_message] ''') # Perform a basic check to ensure that embeddings are generated # and that the dimension of those embeddings is correct. actual_output = result | beam.Map(lambda x: len(x.log_message)) assert_that( actual_output, equal_to([SENTENCE_EMBEDDING_DIMENSION] * len(DATA))) def test_ml_transform_outputs_schema(self): SENTENCE_EMBEDDING_DIMENSION = 384 ml_opts = beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle', yaml_experimental_features=['ML']) with tempfile.TemporaryDirectory() as tempdir: with beam.Pipeline(options=ml_opts) as p: result = p | YamlTransform( f''' type: chain transforms: - type: Create config: elements: - {{id: 1, log_message: "Error in module A"}} - {{id: 2, log_message: "Warning in module B"}} - {{id: 3, log_message: "Info in module C"}} - type: MLTransform config: write_artifact_location: {tempdir} transforms: - type: SentenceTransformerEmbeddings config: model_name: all-MiniLM-L6-v2 columns: [log_message] - type: MapToFields config: language: python fields: id: id embedding: log_message ''') def check_row(row): assert isinstance(row.id, int) assert isinstance(row.embedding, list) assert len(row.embedding) == SENTENCE_EMBEDDING_DIMENSION return row.id assert_that(result | beam.Map(check_row), equal_to([1, 2, 3])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_provider.py0000755000175100017510000016632315206113033022064 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """This module defines Providers usable from yaml, which is a specification for where to find and how to invoke services that vend implementations of various PTransforms.""" import abc import collections import functools import hashlib import inspect import json import logging import os import re import shutil import subprocess import sys import tempfile import urllib.parse import warnings from collections.abc import Callable from collections.abc import Iterable from collections.abc import Mapping from typing import Any from typing import Optional from typing import Union import docstring_parser import yaml import apache_beam as beam import apache_beam.dataframe.io import apache_beam.io import apache_beam.transforms.util from apache_beam import ManagedReplacement from apache_beam.io.filesystems import FileSystems from apache_beam.portability.api import schema_pb2 from apache_beam.runners import pipeline_context from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms import external from apache_beam.transforms import window from apache_beam.transforms.fully_qualified_named_transform import FullyQualifiedNamedTransform from apache_beam.typehints import schemas from apache_beam.typehints import trivial_inference from apache_beam.typehints.schemas import named_tuple_to_schema from apache_beam.typehints.schemas import typing_to_runner_api from apache_beam.utils import python_callable from apache_beam.utils import subprocess_server from apache_beam.version import __version__ as beam_version from apache_beam.yaml import json_utils from apache_beam.yaml import yaml_utils from apache_beam.yaml.yaml_errors import maybe_with_exception_handling_transform_fn _LOGGER = logging.getLogger(__name__) class NotAvailableWithReason: """A False value that provides additional content. Primarily used to return a value from Provider.available(). """ def __init__(self, reason): self.reason = reason def __bool__(self): return False class Provider(abc.ABC): """Maps transform types names and args to concrete PTransform instances.""" @abc.abstractmethod def available(self) -> Union[bool, NotAvailableWithReason]: """Returns whether this provider is available to use in this environment.""" raise NotImplementedError(type(self)) @abc.abstractmethod def cache_artifacts(self) -> Optional[Iterable[str]]: raise NotImplementedError(type(self)) @abc.abstractmethod def provided_transforms(self) -> Iterable[str]: """Returns a list of transform type names this provider can handle.""" raise NotImplementedError(type(self)) def config_schema(self, type): return None def description(self, type): return None def requires_inputs(self, typ: str, args: Mapping[str, Any]) -> bool: """Returns whether this transform requires inputs. Specifically, if this returns True and inputs are not provided than an error will be thrown. This is best-effort, primarily for better and earlier error messages. """ return not typ.startswith('Read') @abc.abstractmethod def create_transform( self, typ: str, args: Mapping[str, Any], yaml_create_transform: Callable[ [Mapping[str, Any], Iterable[beam.PCollection]], beam.PTransform] ) -> beam.PTransform: """Creates a PTransform instance for the given transform type and arguments. """ raise NotImplementedError(type(self)) def underlying_provider(self): """If this provider is simply a proxy to another provider, return the provider that should actually be used for affinity checking. """ return self def affinity(self, other: "Provider"): """Returns a value approximating how good it would be for this provider to be used immediately following a transform from the other provider (e.g. to encourage fusion). """ # TODO(yaml): This is a very rough heuristic. Consider doing better. # E.g. we could look at the expected environments themselves. # Possibly, we could provide multiple expansions and have the runner itself # choose the actual implementation based on fusion (and other) criteria. a = self.underlying_provider() b = other.underlying_provider() return a._affinity(b) + b._affinity(a) def _affinity(self, other: "Provider"): if self is other or self == other: return 100 elif type(self) == type(other): return 10 else: return 0 @functools.cache # pylint: disable=method-cache-max-size-none def with_extra_dependencies(self, dependencies: Iterable[str]): result = self._with_extra_dependencies(dependencies) if not hasattr(result, 'to_json'): result.to_json = lambda: {'type': type(result).__name__} return result def _with_extra_dependencies(self, dependencies: Iterable[str]): raise ValueError( 'This provider of type %s does not support additional dependencies.' % type(self).__name__) def as_provider(name, provider_or_constructor): if isinstance(provider_or_constructor, Provider): return provider_or_constructor else: return InlineProvider({name: provider_or_constructor}) def as_provider_list(name, lst): if not isinstance(lst, list): return as_provider_list(name, [lst]) return [as_provider(name, x) for x in lst] class ExternalProvider(Provider): """A Provider implemented via the cross language transform service.""" _provider_types: dict[str, Callable[..., Provider]] = {} def __init__(self, urns, service, managed_replacement=None): """Initializes the ExternalProvider. Args: urns: a set of URNs that uniquely identify the transforms supported. service: the gradle target that identified the expansion service jar. managed_replacement (Optional): a map that defines the transform for which the SDK may replace the transform with an available managed transform. """ self._urns = urns self._service = service self._schema_transforms = None self._managed_replacement = managed_replacement def provided_transforms(self): return self._urns.keys() def schema_transforms(self): if callable(self._service): self._service = self._service() if self._schema_transforms is None: try: self._schema_transforms = { config.identifier: config for config in external.SchemaAwareExternalTransform.discover( self._service, ignore_errors=True) } except Exception: # It's possible this service doesn't vend schema transforms. self._schema_transforms = {} return self._schema_transforms def config_schema(self, type): if self._urns[type] in self.schema_transforms(): return named_tuple_to_schema( self.schema_transforms()[self._urns[type]].configuration_schema) def description(self, type): if self._urns[type] in self.schema_transforms(): return self.schema_transforms()[self._urns[type]].description def requires_inputs(self, typ, args): if self._urns[typ] in self.schema_transforms(): return bool(self.schema_transforms()[self._urns[typ]].inputs) else: return super().requires_inputs(typ, args) def create_transform(self, type, args, yaml_create_transform): if callable(self._service): self._service = self._service() urn = self._urns[type] if urn in self.schema_transforms(): managed_replacement = None if self._managed_replacement and type in self._managed_replacement: managed_replacement = ManagedReplacement( underlying_transform_identifier=urn, update_compatibility_version=self._managed_replacement[type]) return external.SchemaAwareExternalTransform( urn, self._service, rearrange_based_on_discovery=True, managed_replacement=managed_replacement, **args) else: return type >> self.create_external_transform(urn, args) def create_external_transform(self, urn, args): return external.ExternalTransform( urn, external.ImplicitSchemaPayloadBuilder(args).payload(), self._service) @classmethod def provider_from_spec(cls, source_path, spec): from apache_beam.yaml.yaml_transform import SafeLineLoader for required in ('type', 'transforms'): if required not in spec: raise ValueError( f'Missing {required} in provider ' f'at line {SafeLineLoader.get_line(spec)}') urns = SafeLineLoader.strip_metadata(spec['transforms']) type = spec['type'] config = SafeLineLoader.strip_metadata(spec.get('config', {})) extra_params = set(SafeLineLoader.strip_metadata(spec).keys()) - { 'transforms', 'type', 'config' } if extra_params: raise ValueError( f'Unexpected parameters in provider of type {type} ' f'at line {SafeLineLoader.get_line(spec)}: {extra_params}') if config.get('version', None) == 'BEAM_VERSION': config['version'] = beam_version if type in cls._provider_types: try: constructor = cls._provider_types[type] if 'provider_base_path' in inspect.signature(constructor).parameters: config['provider_base_path'] = source_path result = constructor(urns, **config) if not hasattr(result, 'to_json'): result.to_json = lambda: spec return result except Exception as exn: raise ValueError( f'Unable to instantiate provider of type {type} ' f'at line {SafeLineLoader.get_line(spec)}: {exn}') from exn else: raise NotImplementedError( f'Unknown provider type: {type} ' f'at line {SafeLineLoader.get_line(spec)}.') @classmethod def register_provider_type(cls, type_name): def apply(constructor): cls._provider_types[type_name] = constructor return constructor return apply @ExternalProvider.register_provider_type('javaJar') def java_jar(urns, provider_base_path, jar: str): if not os.path.exists(jar): parsed = urllib.parse.urlparse(jar) if not parsed.scheme or not parsed.netloc: raise ValueError(f'Invalid path or url: {jar}') return ExternalJavaProvider( urns, lambda: _join_url_or_filepath(provider_base_path, jar)) @ExternalProvider.register_provider_type('mavenJar') def maven_jar( urns, *, artifact_id, group_id, version, repository=subprocess_server.JavaJarServer.MAVEN_CENTRAL_REPOSITORY, classifier=None, appendix=None): return ExternalJavaProvider( urns, lambda: subprocess_server.JavaJarServer.path_to_maven_jar( artifact_id=artifact_id, group_id=group_id, version=version, repository=repository, classifier=classifier, appendix=appendix)) @ExternalProvider.register_provider_type('beamJar') def beam_jar( urns, *, gradle_target, managed_replacement=None, appendix=None, version=beam_version, artifact_id=None): return ExternalJavaProvider( urns, lambda: subprocess_server.JavaJarServer.path_to_beam_jar( gradle_target=gradle_target, version=version, artifact_id=artifact_id ), managed_replacement=managed_replacement) @ExternalProvider.register_provider_type('docker') def docker(urns, **config): raise NotImplementedError() @ExternalProvider.register_provider_type('remote') class RemoteProvider(ExternalProvider): _is_available = None def __init__(self, urns, address: str): super().__init__(urns, service=address) self._address = address def available(self): if self._is_available is None: try: with external.ExternalTransform.service(self._service) as service: service.ready(1) self._is_available = True except Exception: self._is_available = NotAvailableWithReason( f'Remote provider not reachable at {self._address}.') return self._is_available def cache_artifacts(self): pass class ExternalJavaProvider(ExternalProvider): def __init__( self, urns, jar_provider, managed_replacement=None, classpath=None): super().__init__( urns, lambda: external.JavaJarExpansionService( jar_provider(), classpath=classpath), managed_replacement) self._jar_provider = jar_provider self._classpath = classpath def available(self): # Directly use shutil.which to find the Java executable cross-platform java_path = shutil.which(subprocess_server.JavaHelper.get_java()) if java_path: return True # Return error message when not found return NotAvailableWithReason( 'Unable to locate java executable: java not found in PATH or JAVA_HOME') def cache_artifacts(self): return [self._jar_provider()] def _with_extra_dependencies(self, dependencies: Iterable[str]): jars = sum(( external.JavaJarExpansionService._expand_jars(dep) for dep in dependencies), []) return ExternalJavaProvider( self._urns, jar_provider=self._jar_provider, classpath=(list(self._classpath or []) + list(jars))) @ExternalProvider.register_provider_type('python') def python(urns, provider_base_path, packages=()): if packages: return ExternalPythonProvider(urns, provider_base_path, packages) else: return InlineProvider({ name: python_callable.PythonCallableWithSource.load_from_source( constructor) for (name, constructor) in urns.items() }) @ExternalProvider.register_provider_type('pythonPackage') class ExternalPythonProvider(ExternalProvider): def __init__(self, urns, provider_base_path, packages: Iterable[str]): def is_path_or_urn(package): return ( '/' in package or urllib.parse.urlparse(package).scheme or os.path.exists(package)) super().__init__( urns, PypiExpansionService([ _join_url_or_filepath(provider_base_path, package) if is_path_or_urn(package) else package for package in packages ])) self._packages = packages def available(self): return True # If we're running this script, we have Python installed. def cache_artifacts(self): return [self._service._venv()] def create_external_transform(self, urn, args): # Python transforms are "registered" by fully qualified name. if not re.match(r'^[\w.]*$', urn): # Treat it as source. args = {'source': urn, **args} urn = '__constructor__' return external.ExternalTransform( "beam:transforms:python:fully_qualified_named", external.ImplicitSchemaPayloadBuilder({ 'constructor': urn, 'kwargs': args, }).payload(), self._service) def _affinity(self, other: "Provider"): if isinstance(other, InlineProvider): return 50 else: return super()._affinity(other) def _with_extra_dependencies(self, dependencies: Iterable[str]): return ExternalPythonProvider( self._urns, None, set(self._packages).union(set(dependencies))) @ExternalProvider.register_provider_type('yaml') class YamlProvider(Provider): def __init__(self, transforms: Mapping[str, Mapping[str, Any]]): if not isinstance(transforms, dict): raise ValueError('Transform mapping must be a dict.') self._transforms = transforms def available(self): return True def cache_artifacts(self): pass def provided_transforms(self): return self._transforms.keys() def config_schema(self, type): return json_utils.json_schema_to_beam_schema(self.json_config_schema(type)) def json_config_schema(self, type): return dict( type='object', additionalProperties=False, **self._transforms[type].get('config_schema', {})) def description(self, type): return self._transforms[type].get('description') def requires_inputs(self, type, args): return self._transforms[type].get( 'requires_inputs', super().requires_inputs(type, args)) def create_transform( self, type: str, args: Mapping[str, Any], yaml_create_transform: Callable[ [Mapping[str, Any], Iterable[beam.PCollection]], beam.PTransform] ) -> beam.PTransform: from apache_beam.yaml.yaml_transform import SafeLineLoader from apache_beam.yaml.yaml_transform import expand_jinja from apache_beam.yaml.yaml_transform import preprocess spec = self._transforms[type] try: import jsonschema jsonschema.validate(args, self.json_config_schema(type)) except ImportError: warnings.warn( 'Please install jsonschema ' f'for better provider validation of "{type}"') body = spec['body'] # Stringify to apply jinja. if isinstance(body, str): body_str = body else: body_str = yaml.safe_dump(SafeLineLoader.strip_metadata(body)) # Now re-parse resolved templatization. body = yaml.load(expand_jinja(body_str, args), Loader=SafeLineLoader) if (body.get('type') == 'chain' and 'input' not in body and spec.get('requires_inputs', True)): body['input'] = 'input' return yaml_create_transform(preprocess(body)) # type: ignore # This is needed because type inference can't handle *args, **kwargs forwarding. # TODO(BEAM-24755): Add support for type inference of through kwargs calls. def fix_pycallable(): from apache_beam.transforms.ptransform import label_from_callable def default_label(self): src = self._source.strip() last_line = src.split('\n')[-1] if last_line[0] != ' ' and len(last_line) < 72: return last_line return label_from_callable(self._callable) def _argspec_fn(self): return self._callable python_callable.PythonCallableWithSource.default_label = default_label python_callable.PythonCallableWithSource._argspec_fn = property(_argspec_fn) original_infer_return_type = trivial_inference.infer_return_type def infer_return_type(fn, *args, **kwargs): if isinstance(fn, python_callable.PythonCallableWithSource): fn = fn._callable return original_infer_return_type(fn, *args, **kwargs) trivial_inference.infer_return_type = infer_return_type original_fn_takes_side_inputs = ( apache_beam.transforms.util.fn_takes_side_inputs) def fn_takes_side_inputs(fn): if isinstance(fn, python_callable.PythonCallableWithSource): fn = fn._callable return original_fn_takes_side_inputs(fn) apache_beam.transforms.util.fn_takes_side_inputs = fn_takes_side_inputs class InlineProvider(Provider): def __init__(self, transform_factories, no_input_transforms=()): self._transform_factories = transform_factories self._no_input_transforms = set(no_input_transforms) def available(self): return True def cache_artifacts(self): pass def provided_transforms(self): return self._transform_factories.keys() def config_schema(self, typ): return self.config_schema_from_callable(self._transform_factories[typ]) @classmethod def config_schema_from_callable(cls, factory): if isinstance(factory, type) and issubclass(factory, beam.PTransform): # https://bugs.python.org/issue40897 params = dict(inspect.signature(factory.__init__).parameters) if 'self' in params: del params['self'] else: params = inspect.signature(factory).parameters def type_of(p): t = p.annotation if t == p.empty: return Any else: return t docs = { param.arg_name: param.description for param in cls.get_docs(factory).params } names_and_types = [(name, typing_to_runner_api(type_of(p))) for name, p in params.items()] return schema_pb2.Schema( fields=[ schema_pb2.Field(name=name, type=type, description=docs.get(name)) for (name, type) in names_and_types ]) def description(self, typ): return self.description_from_callable(self._transform_factories[typ]) @classmethod def description_from_callable(cls, factory): def empty_if_none(s): return s or '' docs = cls.get_docs(factory) return ( empty_if_none(docs.short_description) + ('\n\n' if docs.blank_after_short_description else '\n') + empty_if_none(docs.long_description)).strip() or None @classmethod def get_docs(cls, factory): docstring = factory.__doc__ or '' # These "extra" docstring parameters are not relevant for YAML and mess # up the parsing. docstring = re.sub( r'Pandas Parameters\s+-----.*', '', docstring, flags=re.S) return docstring_parser.parse( docstring, docstring_parser.DocstringStyle.GOOGLE) def create_transform(self, type, args, yaml_create_transform): return self._transform_factories[type](**args) def to_json(self): return {'type': "InlineProvider"} def requires_inputs(self, typ, args): if typ in self._no_input_transforms: return False elif hasattr(self._transform_factories[typ], '_yaml_requires_inputs'): return self._transform_factories[typ]._yaml_requires_inputs else: return super().requires_inputs(typ, args) def _with_extra_dependencies(self, dependencies): external_provider = ExternalPythonProvider( # disable yapf { typ: 'apache_beam.yaml.yaml_provider.standard_inline_providers.' + typ.replace('-', '_') for typ in self._transform_factories.keys() }, '__inline__', dependencies) external_provider.to_json = self.to_json return external_provider class MetaInlineProvider(InlineProvider): def create_transform(self, type, args, yaml_create_transform): return self._transform_factories[type](yaml_create_transform, **args) # Note: This function is used to override the default provider by some # users, so a change here will be breaking to those users. Change with # caution. def get_default_sql_provider(): return beam_jar( urns={'Sql': 'beam:external:java:sql:v1'}, gradle_target='sdks:java:extensions:sql:expansion-service:shadowJar') class SqlBackedProvider(Provider): def __init__( self, transforms: Mapping[str, Callable[..., beam.PTransform]], sql_provider: Optional[Provider] = None): self._transforms = transforms if sql_provider is None: sql_provider = get_default_sql_provider() self._sql_provider = sql_provider def sql_provider(self): return self._sql_provider def provided_transforms(self): return self._transforms.keys() def config_schema(self, type): full_config = InlineProvider.config_schema_from_callable( self._transforms[type]) # Omit the (first) query -> transform parameter. return schema_pb2.Schema(fields=full_config.fields[1:]) def description(self, type): return InlineProvider.description_from_callable(self._transforms[type]) def available(self): return self.sql_provider().available() def cache_artifacts(self): return self.sql_provider().cache_artifacts() def underlying_provider(self): return self.sql_provider() def to_json(self): return {'type': "SqlBackedProvider"} def create_transform( self, typ: str, args: Mapping[str, Any], yaml_create_transform: Any) -> beam.PTransform: return self._transforms[typ]( lambda query: self.sql_provider().create_transform( 'Sql', {'query': query}, yaml_create_transform), **args) PRIMITIVE_NAMES_TO_ATOMIC_TYPE = { py_type.__name__: schema_type for (py_type, schema_type) in schemas.PRIMITIVE_TO_ATOMIC_TYPE.items() if py_type.__module__ != 'typing' } def element_to_rows(e): if isinstance(e, dict): return dicts_to_rows(e) else: return beam.Row(element=dicts_to_rows(e)) def dicts_to_rows(o): if isinstance(o, dict): return beam.Row(**{k: dicts_to_rows(v) for k, v in o.items()}) elif isinstance(o, list): return [dicts_to_rows(e) for e in o] else: return o def _unify_element_with_schema(element, target_schema): """Convert an element to match the target schema, preserving existing fields only.""" if target_schema is None: return element # If element is already a named tuple, convert to dict first if hasattr(element, '_asdict'): element_dict = element._asdict() elif isinstance(element, dict): element_dict = element else: # This element is not a row-like object. If the target schema has a single # field, assume this element is the value for that field. if len(target_schema._fields) == 1: return target_schema(**{target_schema._fields[0]: element}) else: return element # Create new element with only the fields that exist in the original # element plus None for fields that are expected but missing unified_dict = {} for field_name in target_schema._fields: if field_name in element_dict: value = element_dict[field_name] # Ensure the value matches the expected type # This is particularly important for list fields if value is not None and not isinstance(value, list) and hasattr( value, '__iter__') and not isinstance( value, (str, bytes)) and not hasattr(value, '_asdict'): # Convert iterables to lists if needed unified_dict[field_name] = list(value) else: unified_dict[field_name] = value else: unified_dict[field_name] = None return target_schema(**unified_dict) class YamlProviders: class AssertEqual(beam.PTransform): """Asserts that the input contains exactly the elements provided. This is primarily used for testing; it will cause the entire pipeline to fail if the input to this transform is not exactly the set of `elements` given in the config parameter. As with Create, YAML/JSON-style mappings are interpreted as Beam rows, e.g.:: type: AssertEqual input: SomeTransform config: elements: - {a: 0, b: "foo"} - {a: 1, b: "bar"} would ensure that `SomeTransform` produced exactly two elements with values `(a=0, b="foo")` and `(a=1, b="bar")` respectively. Args: elements: The set of elements that should belong to the PCollection. YAML/JSON-style mappings will be interpreted as Beam rows. """ def __init__(self, elements: Iterable[Any]): self._elements = elements def expand(self, pcoll): def to_dict(row): # filter None when comparing temp_dict = {k: v for k, v in row._asdict().items() if v is not None} return dict(temp_dict.items()) return assert_that( pcoll | beam.Map(to_dict), equal_to([to_dict(e) for e in dicts_to_rows(self._elements)])) @staticmethod def create(elements: Iterable[Any], reshuffle: Optional[bool] = True): """Creates a collection containing a specified set of elements. This transform always produces schema'd data. For example:: type: Create config: elements: [1, 2, 3] will result in an output with three elements with a schema of Row(element=int) whereas YAML/JSON-style mappings will be interpreted directly as Beam rows, e.g.:: type: Create config: elements: - {first: 0, second: {str: "foo", values: [1, 2, 3]}} - {first: 1, second: {str: "bar", values: [4, 5, 6]}} will result in a schema of the form (int, Row(string, list[int])). This can also be expressed as YAML:: type: Create config: elements: - first: 0 second: str: "foo" values: [1, 2, 3] - first: 1 second: str: "bar" values: [4, 5, 6] If the elements are a mix of dicts and non-dicts, the non-dict elements will be wrapped in a Row with a single field "element". For example:: type: Create config: elements: [1, {"a": 2}] will result in an output with two elements with a schema of Row(element=int, a=int) looking like: Row(element=1, a=None) Row(element=None, a=2) Args: elements: The set of elements that should belong to the PCollection. YAML/JSON-style mappings will be interpreted as Beam rows. Primitives will be mapped to rows with a single "element" field. reshuffle: (optional) Whether to introduce a reshuffle (to possibly redistribute the work) if there is more than one element in the collection. Defaults to True. """ # Though str and dict are technically iterable, we disallow them # as using the characters or keys respectively is almost certainly # not the intent. if not isinstance(elements, Iterable) or isinstance(elements, (dict, str)): raise TypeError('elements must be a list of elements') if elements: # Normalize elements to be all dicts or all primitives. has_dict = False has_non_dict = False for e in elements: if isinstance(e, dict): has_dict = True else: has_non_dict = True if has_dict and has_non_dict: break if has_dict and has_non_dict: elements = [ e if isinstance(e, dict) else { 'element': e } for e in elements ] # Check if elements have different keys updated_elements = elements if elements and all(isinstance(e, dict) for e in elements): keys = [set(e.keys()) for e in elements] if len(set.union(*keys)) > min(len(k) for k in keys): # Merge all dictionaries to get all possible keys all_keys = set() for element in elements: if isinstance(element, dict): all_keys.update(element.keys()) # Create a merged dictionary with all keys merged_dict = {} for key in all_keys: merged_dict[key] = None # Use None as a default value # Update each element with the merged dictionary updated_elements = [] for e in elements: if isinstance(e, dict): updated_elements.append({**merged_dict, **e}) else: updated_elements.append(e) return beam.Create([element_to_rows(e) for e in updated_elements], reshuffle=reshuffle is not False) # Or should this be posargs, args? # pylint: disable=dangerous-default-value @staticmethod def fully_qualified_named_transform( constructor: str, args: Optional[Iterable[Any]] = (), kwargs: Optional[Mapping[str, Any]] = {}): """A Python PTransform identified by fully qualified name. This allows one to import, construct, and apply any Beam Python transform. This can be useful for using transforms that have not yet been exposed via a YAML interface. Note, however, that conversion may be required if this transform does not accept or produce Beam Rows. For example:: type: PyTransform config: constructor: apache_beam.pkg.mod.SomeClass args: [1, 'foo'] kwargs: baz: 3 can be used to access the transform `apache_beam.pkg.mod.SomeClass(1, 'foo', baz=3)`. See also the documentation on [Inlining Python](https://beam.apache.org/documentation/sdks/yaml-inline-python/). Args: constructor: Fully qualified name of a callable used to construct the transform. Often this is a class such as `apache_beam.pkg.mod.SomeClass` but it can also be a function or any other callable that returns a PTransform. args: A list of parameters to pass to the callable as positional arguments. kwargs: A list of parameters to pass to the callable as keyword arguments. """ with FullyQualifiedNamedTransform.with_filter('*'): return constructor >> FullyQualifiedNamedTransform( constructor, args, kwargs) # This intermediate is needed because there is no way to specify a tuple of # exactly zero or one PCollection in yaml (as they would be interpreted as # PBegin and the PCollection itself respectively). class Flatten(beam.PTransform): """Flattens multiple PCollections into a single PCollection. The elements of the resulting PCollection will be the (disjoint) union of all the elements of all the inputs. Note that in YAML transforms can always take a list of inputs which will be implicitly flattened. """ def __init__(self): # Suppress the "label" argument from the superclass for better docs. # pylint: disable=useless-parent-delegation super().__init__() def _merge_schemas(self, pcolls): """Merge schemas from multiple PCollections to create a unified schema. This function creates a unified schema that contains all fields from all input PCollections. Fields are made optional to handle missing values. If fields have different types, they are unified to Optional[Any]. """ from apache_beam.typehints.schemas import named_fields_from_element_type # Collect all schemas schemas = [] for pcoll in pcolls: if hasattr(pcoll, 'element_type') and pcoll.element_type: try: fields = named_fields_from_element_type(pcoll.element_type) schemas.append(dict(fields)) except (ValueError, TypeError): # If we can't extract schema, skip this PCollection continue if not schemas: return None # Merge all field names and types. all_field_names = set().union(*(s.keys() for s in schemas)) unified_fields = {} for name in all_field_names: present_types = {s[name] for s in schemas if name in s} if len(present_types) > 1: unified_fields[name] = Optional[Any] else: unified_fields[name] = Optional[present_types.pop()] # Create unified schema if unified_fields: from apache_beam.typehints.schemas import named_fields_to_schema from apache_beam.typehints.schemas import named_tuple_from_schema unified_schema = named_fields_to_schema(list(unified_fields.items())) return named_tuple_from_schema(unified_schema) return None def expand(self, pcolls): if isinstance(pcolls, beam.PCollection): pipeline_arg = {} pcolls = (pcolls, ) elif isinstance(pcolls, dict): pipeline_arg = {} pcolls = tuple(pcolls.values()) else: pipeline_arg = {'pipeline': pcolls.pipeline} pcolls = () if not pcolls: return pcolls | beam.Flatten(**pipeline_arg) # Try to unify schemas unified_schema = self._merge_schemas(pcolls) if unified_schema is None: # No schema unification needed, use standard flatten return pcolls | beam.Flatten(**pipeline_arg) # Apply schema unification to each PCollection before flattening. unified_pcolls = [] for i, pcoll in enumerate(pcolls): unified_pcoll = pcoll | f'UnifySchema{i}' >> beam.Map( _unify_element_with_schema, target_schema=unified_schema).with_output_types(unified_schema) unified_pcolls.append(unified_pcoll) # Flatten the unified PCollections return unified_pcolls | beam.Flatten(**pipeline_arg) class WindowInto(beam.PTransform): # pylint: disable=line-too-long """A window transform assigning windows to each element of a PCollection. The assigned windows will affect all downstream aggregating operations, which will aggregate by window as well as by key. See [the Beam documentation on windowing](https://beam.apache.org/documentation/programming-guide/#windowing) for more details. Sizes, offsets, periods and gaps (where applicable) must be defined using a time unit suffix 'ms', 's', 'm', 'h' or 'd' for milliseconds, seconds, minutes, hours or days, respectively. If a time unit is not specified, it will default to 's'. For example:: windowing: type: fixed size: 30s Note that any Yaml transform can have a [windowing parameter](https://beam.apache.org/documentation/sdks/yaml/#windowing), which is applied to its inputs (if any) or outputs (if there are no inputs) which means that explicit WindowInto operations are not typically needed. Args: windowing: the type and parameters of the windowing to perform """ def __init__(self, windowing): self._window_transform = self._parse_window_spec(windowing) def expand(self, pcoll): return pcoll | self._window_transform @staticmethod def _parse_duration(value, name): time_units = { 'ms': 0.001, 's': 1, 'm': 60, 'h': 60 * 60, 'd': 60 * 60 * 12 } value, suffix = re.match(r'^(.*?)([^\d]*)$', str(value)).groups() # Default to seconds if time unit suffix is not defined if not suffix: suffix = 's' if not value: raise ValueError( f"Invalid windowing {name} value " f"'{suffix if not value else value}'. " f"Must provide numeric value.") if suffix not in time_units: raise ValueError(( "Invalid windowing {} time unit '{}'. " + "Valid time units are {}.").format( name, suffix, ', '.join("'{}'".format(k) for k in time_units.keys()))) return float(value) * time_units[suffix] @staticmethod def _parse_window_spec(spec): spec = dict(spec) window_type = spec.pop('type') # TODO: These are in seconds, perhaps parse duration strings meaningfully? if window_type == 'global': window_fn = window.GlobalWindows() elif window_type == 'fixed': window_fn = window.FixedWindows( YamlProviders.WindowInto._parse_duration(spec.pop('size'), 'size'), YamlProviders.WindowInto._parse_duration( spec.pop('offset', 0), 'offset')) elif window_type == 'sliding': window_fn = window.SlidingWindows( YamlProviders.WindowInto._parse_duration(spec.pop('size'), 'size'), YamlProviders.WindowInto._parse_duration( spec.pop('period'), 'period'), YamlProviders.WindowInto._parse_duration( spec.pop('offset', 0), 'offset')) elif window_type == 'sessions': window_fn = window.Sessions( YamlProviders.WindowInto._parse_duration(spec.pop('gap'), 'gap')) else: raise ValueError(f'Unknown window type {window_type}') if spec: raise ValueError(f'Unknown parameters {spec.keys()}') # TODO: Triggering, etc. return beam.WindowInto(window_fn) @staticmethod @beam.ptransform_fn @maybe_with_exception_handling_transform_fn def log_for_testing( pcoll, *, level: Optional[str] = 'INFO', prefix: Optional[str] = ''): """Logs each element of its input PCollection. The output of this transform is a copy of its input for ease of use in chain-style pipelines. Args: level: one of ERROR, INFO, or DEBUG, mapped to a corresponding language-specific logging level prefix: an optional identifier that will get prepended to the element being logged """ # Keeping this simple to be language agnostic. # The intent is not to develop a logging library (and users can always do) # their own mappings to get fancier output. log_levels = { 'ERROR': logging.error, 'INFO': logging.info, 'DEBUG': logging.debug, } if level not in log_levels: raise ValueError( f'Unknown log level {level} not in {list(log_levels.keys())}') logger = log_levels[level] def to_loggable_json_recursive(o): if isinstance(o, (str, bytes)): return str(o) elif callable(getattr(o, '_asdict', None)): return to_loggable_json_recursive(o._asdict()) elif isinstance(o, Mapping) and callable(getattr(o, 'items', None)): return {str(k): to_loggable_json_recursive(v) for k, v in o.items()} elif isinstance(o, Iterable): return [to_loggable_json_recursive(x) for x in o] else: return o def log_and_return(x): logger(prefix + json.dumps(to_loggable_json_recursive(x))) return x return pcoll | "LogForTesting" >> beam.Map(log_and_return) class Reshuffle(beam.PTransform): """Reshuffles the elements of a PCollection. Redistributes the elements of a PCollection to prevent fusion or balance load. Args: num_buckets: (optional) Specifies the maximum random keys that would be generated. If not set, a default value is used. """ def __init__(self, num_buckets: Optional[int] = None): self.num_buckets = num_buckets def expand(self, pcoll): return pcoll | beam.Reshuffle(num_buckets=self.num_buckets) @staticmethod def create_builtin_provider(): return InlineProvider({ 'AssertEqual': YamlProviders.AssertEqual, 'Create': YamlProviders.create, 'LogForTesting': YamlProviders.log_for_testing, 'PyTransform': YamlProviders.fully_qualified_named_transform, 'Flatten': YamlProviders.Flatten, 'WindowInto': YamlProviders.WindowInto, 'Reshuffle': YamlProviders.Reshuffle, }, no_input_transforms=('Create', )) class TranslatingProvider(Provider): def __init__( self, transforms: Mapping[str, Callable[..., beam.PTransform]], underlying_provider: Provider): self._transforms = transforms self._underlying_provider = underlying_provider def provided_transforms(self): return self._transforms.keys() def available(self): return self._underlying_provider.available() def cache_artifacts(self): return self._underlying_provider.cache_artifacts() def underlying_provider(self): return self._underlying_provider def to_json(self): return {'type': "TranslatingProvider"} def create_transform( self, typ: str, config: Mapping[str, Any], yaml_create_transform: Any) -> beam.PTransform: return self._transforms[typ](self._underlying_provider, **config) def _with_extra_dependencies(self, dependencies: Iterable[str]): return TranslatingProvider( self._transforms, self._underlying_provider._with_extra_dependencies(dependencies)) def create_java_builtin_provider(): """Exposes built-in transforms from Java as well as Python to maximize opportunities for fusion. This class holds those transforms that require pre-processing of the configs. For those Java transforms that can consume the user-provided configs directly (or only need a simple renaming of parameters) a direct or renaming provider is the simpler choice. """ # An alternative could be examining the capabilities of various environments # during (or as a pre-processing phase before) fusion to align environments # where possible. This would also require extra care in skipping these # common transforms when doing the provider affinity analysis. def java_window_into(java_provider, windowing): """Use the `windowing` WindowingStrategy and invokes the Java class. Though it would not be that difficult to implement this in Java as well, we prefer to implement it exactly once for consistency (especially as it evolves). """ windowing_strategy = YamlProviders.WindowInto._parse_window_spec( windowing).get_windowing(None) # No context needs to be preserved for the basic WindowFns. empty_context = pipeline_context.PipelineContext() return java_provider.create_transform( 'WindowIntoStrategy', { 'serialized_windowing_strategy': windowing_strategy.to_runner_api( empty_context).SerializeToString() }, None) return TranslatingProvider( transforms={'WindowInto': java_window_into}, underlying_provider=beam_jar( urns={ 'WindowIntoStrategy': ( 'beam:schematransform:' 'org.apache.beam:yaml:window_into_strategy:v1') }, gradle_target= 'sdks:java:extensions:schemaio-expansion-service:shadowJar')) class PypiExpansionService: """Expands transforms by fully qualified name in a virtual environment with the given dependencies. """ if 'TOX_WORK_DIR' in os.environ: VENV_CACHE = tempfile.mkdtemp( prefix='test-venv-cache-', dir=os.environ['TOX_WORK_DIR']) elif 'RUNNER_WORKDIR' in os.environ: VENV_CACHE = tempfile.mkdtemp( prefix='test-venv-cache-', dir=os.environ['RUNNER_WORKDIR']) else: VENV_CACHE = os.path.expanduser("~/.apache_beam/cache/venvs") def __init__( self, packages: Iterable[str], base_python: str = sys.executable): if not isinstance(packages, Iterable) or isinstance(packages, str): raise TypeError( "Packages must be an iterable of strings, got %r" % packages) self._packages = list(packages) self._base_python = base_python @classmethod def _key(cls, base_python: str, packages: list[str]) -> str: def normalize_package(package): if os.path.exists(package): # Ignore the exact path by which this package was referenced, # but do create a new environment if it changed. with open(package, 'rb') as fin: return os.path.basename(package) + '-' + _file_digest( fin, 'sha256').hexdigest() else: # Assume urls and pypi identifiers are immutable. return package return json.dumps({ 'binary': base_python, 'packages': sorted(normalize_package(p) for p in packages) }, sort_keys=True) @classmethod def _path(cls, base_python: str, packages: list[str]) -> str: return os.path.join( cls.VENV_CACHE, hashlib.sha256(cls._key(base_python, packages).encode('utf-8')).hexdigest()) @classmethod def _create_venv_from_scratch( cls, base_python: str, packages: list[str]) -> str: venv = cls._path(base_python, packages) if not os.path.exists(venv): try: subprocess.run([base_python, '-m', 'venv', venv], check=True) venv_python = os.path.join(venv, 'bin', 'python') venv_pip = os.path.join(venv, 'bin', 'pip') subprocess.run([venv_python, '-m', 'ensurepip'], check=True) # Issue warning when installing packages from PyPI _LOGGER.warning( " WARNING: Apache Beam is installing Python packages " "from PyPI at runtime.\n" " This may pose security risks or cause instability due to " "repository availability.\n" " Packages: %s\n" " Consider pre-staging dependencies or using a private " "repository mirror.\n" " For more information, see: " "https://beam.apache.org/documentation/sdks/python-dependencies/", ', '.join(packages)) subprocess.run([venv_pip, 'install'] + packages, check=True) with open(venv + '-requirements.txt', 'w') as fout: fout.write('\n'.join(packages)) except: # pylint: disable=bare-except if os.path.exists(venv): shutil.rmtree(venv, ignore_errors=True) raise return venv @classmethod def _create_venv_from_clone( cls, base_python: str, packages: list[str]) -> str: venv = cls._path(base_python, packages) if not os.path.exists(venv): try: # Avoid hard dependency for environments where this is never used. import clonevirtualenv clonable_venv = cls._create_venv_to_clone(base_python) clonevirtualenv.clone_virtualenv(clonable_venv, venv) venv_pip = os.path.join(venv, 'bin', 'pip') # Issue warning when installing packages from PyPI _LOGGER.warning( " WARNING: Apache Beam is installing Python packages " "from PyPI at runtime.\n" " This may pose security risks or cause instability due to " "repository availability.\n" " Packages: %s\n" " Consider pre-staging dependencies or using a private " "repository mirror.\n" " For more information, see: " "https://beam.apache.org/documentation/sdks/python-dependencies/", ', '.join(packages)) subprocess.run([venv_pip, 'install'] + packages, check=True) with open(venv + '-requirements.txt', 'w') as fout: fout.write('\n'.join(packages)) except: # pylint: disable=bare-except if os.path.exists(venv): shutil.rmtree(venv, ignore_errors=True) raise return venv @classmethod def _create_venv_to_clone(cls, base_python: str) -> str: # For '.dev', the default clone source is the venv that owns base_python. # In CI that is often the active tox/sandbox tree; clonevirtualenv can # race with ephemeral paths (tmp/, caches) under that tree. Use the # scratch clonable venv in CI instead. Locally, keep cloning the dev venv # for speed. _ci = os.environ.get('CI', '').lower() in ('true', '1', 'yes') if '.dev' in beam_version and not _ci: base_venv = os.path.dirname(os.path.dirname(base_python)) print('Cloning dev environment from', base_venv) return base_venv return cls._create_venv_from_scratch( base_python, [ 'apache_beam[dataframe,gcp,test,yaml]==' + beam_version, 'virtualenv-clone' ]) def _venv(self) -> str: return self._create_venv_from_clone(self._base_python, self._packages) def __enter__(self): venv = self._venv() self._service_provider = subprocess_server.SubprocessServer( external.ExpansionAndArtifactRetrievalStub, [ os.path.join(venv, 'bin', 'python'), '-m', 'apache_beam.runners.portability.expansion_service_main', '--port', '{{PORT}}', '--fully_qualified_name_glob=*', '--pickle_library=cloudpickle', '--requirements_file=' + os.path.join(venv + '-requirements.txt'), '--serve_loopback_worker', ]) self._service = self._service_provider.__enter__() return self._service def __exit__(self, *args): self._service_provider.__exit__(*args) self._service = None @ExternalProvider.register_provider_type('renaming') class RenamingProvider(Provider): def __init__( self, transforms, provider_base_path, mappings, underlying_provider, defaults=None): if isinstance(underlying_provider, dict): underlying_provider = ExternalProvider.provider_from_spec( provider_base_path, underlying_provider) self._transforms = transforms self._underlying_provider = underlying_provider for transform in transforms.keys(): if transform not in mappings: raise ValueError(f'Missing transform {transform} in mappings.') self._mappings = self.expand_mappings(mappings) self._defaults = defaults or {} @staticmethod def expand_mappings(mappings): if not isinstance(mappings, dict): raise ValueError( "RenamingProvider mappings must be dict of transform " "mappings.") for key, value in mappings.items(): if isinstance(value, str): if value not in mappings.keys(): raise ValueError( "RenamingProvider transform mappings must be dict or " "specify transform that has mappings within same " "provider.") mappings[key] = mappings[value] return mappings def available(self) -> bool: return self._underlying_provider.available() def provided_transforms(self) -> Iterable[str]: return self._transforms.keys() def config_schema(self, type): underlying_schema = self._underlying_provider.config_schema( self._transforms[type]) if underlying_schema is None: return None defaults = self._defaults.get(type, {}) underlying_schema_fields = {f.name: f for f in underlying_schema.fields} missing = set(self._mappings[type].values()) - set( underlying_schema_fields.keys()) if missing: if 'kwargs' in underlying_schema_fields.keys(): # These are likely passed by keyword argument dict rather than missing. for field_name in missing: underlying_schema_fields[field_name] = schema_pb2.Field( name=field_name, type=typing_to_runner_api(Any)) else: raise ValueError( f"Mapping destinations {missing} for {type} are not in the " f"underlying config schema {list(underlying_schema_fields.keys())}") def with_name( original: schema_pb2.Field, new_name: str) -> schema_pb2.Field: result = schema_pb2.Field() result.CopyFrom(original) result.name = new_name return result return schema_pb2.Schema( fields=[ with_name(underlying_schema_fields[dest], src) for (src, dest) in self._mappings[type].items() if dest not in defaults ]) def description(self, typ): return self._underlying_provider.description(self._transforms[typ]) def requires_inputs(self, typ, args): return self._underlying_provider.requires_inputs( self._transforms[typ], args) def create_transform( self, typ: str, args: Mapping[str, Any], yaml_create_transform: Callable[ [Mapping[str, Any], Iterable[beam.PCollection]], beam.PTransform] ) -> beam.PTransform: """Creates a PTransform instance for the given transform type and arguments. """ mappings = self._mappings[typ] # NOTE: If the `key` is not found in the mappings # (e.g. standard_io.yaml), the `key` is passed down # as is to the underlying transform. remapped_args = { mappings.get(key, key): value for key, value in args.items() } for key, value in self._defaults.get(typ, {}).items(): if key not in remapped_args: remapped_args[key] = value return self._underlying_provider.create_transform( self._transforms[typ], remapped_args, yaml_create_transform) def _affinity(self, other): raise NotImplementedError( 'Should not be calling _affinity directly on this provider.') def underlying_provider(self): return self._underlying_provider.underlying_provider() def cache_artifacts(self): self._underlying_provider.cache_artifacts() def _with_extra_dependencies(self, dependencies: Iterable[str]): return RenamingProvider( self._transforms, None, self._mappings, self._underlying_provider._with_extra_dependencies(dependencies), self._defaults) def _as_list(func): @functools.wraps(func) def wrapper(*args, **kwargs): return list(func(*args, **kwargs)) return wrapper def _join_url_or_filepath(base, path): if not base: return path if urllib.parse.urlparse(path).scheme: # path is an absolute path with scheme (whether it is the same as base or # not). return path # path is a relative path or an absolute path without scheme (e.g. /a/b/c) base_scheme = urllib.parse.urlparse(base, '').scheme if base_scheme and base_scheme in urllib.parse.uses_relative: return urllib.parse.urljoin(base, path) else: if FileSystems.join(base, "") == base: # base ends with a filesystem separator return FileSystems.join(base, path) return FileSystems.join(FileSystems.split(base)[0], path) def _read_url_or_filepath(path): scheme = urllib.parse.urlparse(path, '').scheme if scheme and scheme in urllib.parse.uses_netloc: with urllib.request.urlopen(path) as response: return response.read() else: with FileSystems.open(path) as fin: return fin.read() def load_providers(source_path: str) -> Iterable[Provider]: from apache_beam.yaml.yaml_transform import SafeLineLoader provider_specs = yaml.load( _read_url_or_filepath(source_path), Loader=SafeLineLoader) if not isinstance(provider_specs, list): raise ValueError(f"Provider file {source_path} must be a list of Providers") return parse_providers(source_path, provider_specs) @_as_list def parse_providers(source_path, provider_specs: Iterable[Mapping]) -> Iterable[Provider]: from apache_beam.yaml.yaml_transform import SafeLineLoader for provider_spec in provider_specs: if 'include' in provider_spec: if len(SafeLineLoader.strip_metadata(provider_spec)) != 1: raise ValueError( f"When using include, it must be the only parameter: " f"{provider_spec} " f"at {source_path}:{SafeLineLoader.get_line(provider_spec)}") include_path = _join_url_or_filepath( source_path, provider_spec['include']) try: yield from load_providers(include_path) except Exception as exn: raise ValueError( f"Error loading providers from {include_path} included at " f"{source_path}:{SafeLineLoader.get_line(provider_spec)}\n" + str(exn)) from exn else: yield ExternalProvider.provider_from_spec(source_path, provider_spec) def merge_providers(*provider_sets) -> Mapping[str, Iterable[Provider]]: result = collections.defaultdict(list) for provider_set in provider_sets: if isinstance(provider_set, Provider): provider = provider_set provider_set = { transform_type: [provider] for transform_type in provider.provided_transforms() } elif isinstance(provider_set, list): provider_set = merge_providers(*provider_set) for transform_type, providers in provider_set.items(): result[transform_type].extend(providers) return result @functools.cache def standard_providers(): from apache_beam.yaml.yaml_combine import create_combine_providers from apache_beam.yaml.yaml_io import io_providers from apache_beam.yaml.yaml_join import create_join_providers from apache_beam.yaml.yaml_mapping import create_mapping_providers from apache_beam.yaml.yaml_specifiable import create_spec_providers return merge_providers( YamlProviders.create_builtin_provider(), create_java_builtin_provider(), create_mapping_providers(), create_combine_providers(), create_join_providers(), io_providers(), create_spec_providers(), load_providers(yaml_utils.locate_data_file('standard_providers.yaml'))) def _file_digest(fileobj, digest): if hasattr(hashlib, 'file_digest'): # Python 3.11+ return hashlib.file_digest(fileobj, digest) else: hasher = hashlib.new(digest) data = fileobj.read(1 << 20) while data: hasher.update(data) data = fileobj.read(1 << 20) return hasher class _InlineProviderNamespace: """Gives fully qualified names to inline providers from standard_providers(). This is needed to upgrade InlineProvider to ExternalPythonProvider. """ def __getattr__(self, name): typ = name.replace('_', '-') for provider in standard_providers()[typ]: if isinstance(provider, InlineProvider): return provider._transform_factories[typ] raise ValueError(f"No inline provider found for {name}") standard_inline_providers = _InlineProviderNamespace() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_provider_unit_test.py0000644000175100017510000002765415206113033024162 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import os import sys import tempfile import unittest import mock import yaml import apache_beam as beam from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.yaml import yaml_provider from apache_beam.yaml.yaml_provider import YamlProviders from apache_beam.yaml.yaml_transform import SafeLineLoader from apache_beam.yaml.yaml_transform import YamlTransform class WindowIntoTest(unittest.TestCase): def __init__(self, methodName="runWindowIntoTest"): unittest.TestCase.__init__(self, methodName) self.parse_duration = YamlProviders.WindowInto._parse_duration def test_parse_duration_ms(self): value = self.parse_duration('1000ms', 'size') self.assertEqual(1, value) def test_parse_duration_sec(self): value = self.parse_duration('1s', 'size') self.assertEqual(1, value) def test_parse_duration_min(self): value = self.parse_duration('1m', 'size') self.assertEqual(60, value) def test_parse_duration_hour(self): value = self.parse_duration('1h', 'size') self.assertEqual(3600, value) def test_parse_duration_from_decimal(self): value = self.parse_duration('1.5m', 'size') self.assertEqual(90, value) def test_parse_duration_to_decimal(self): value = self.parse_duration('1ms', 'size') self.assertEqual(0.001, value) def test_parse_duration_with_missing_suffix(self): value = self.parse_duration('1', 'size') self.assertEqual(1, value) def test_parse_duration_with_invalid_suffix(self): with self.assertRaises(ValueError): self.parse_duration('1x', 'size') def test_parse_duration_with_missing_value(self): with self.assertRaises(ValueError): self.parse_duration('s', 'size') class ProviderParsingTest(unittest.TestCase): INLINE_PROVIDER = {'type': 'TEST', 'name': 'INLINED'} INCLUDED_PROVIDER = {'type': 'TEST', 'name': 'INCLUDED'} EXTRA_PROVIDER = {'type': 'TEST', 'name': 'EXTRA'} @classmethod def setUpClass(cls): cls.tempdir = tempfile.TemporaryDirectory() cls.to_include = os.path.join(cls.tempdir.name, 'providers.yaml') with open(cls.to_include, 'w') as fout: yaml.dump([cls.INCLUDED_PROVIDER], fout) cls.to_include_nested = os.path.join( cls.tempdir.name, 'nested_providers.yaml') with open(cls.to_include_nested, 'w') as fout: yaml.dump([{'include': './providers.yaml'}, cls.EXTRA_PROVIDER], fout) @classmethod def tearDownClass(cls): cls.tempdir.cleanup() @mock.patch( 'apache_beam.yaml.yaml_provider.ExternalProvider.provider_from_spec', lambda _, x: x) def test_include_file(self): flattened = [ SafeLineLoader.strip_metadata(spec) for spec in yaml_provider.parse_providers( '', [ self.INLINE_PROVIDER, { 'include': self.to_include }, ]) ] self.assertEqual([ self.INLINE_PROVIDER, self.INCLUDED_PROVIDER, ], flattened) @mock.patch( 'apache_beam.yaml.yaml_provider.ExternalProvider.provider_from_spec', lambda _, x: x) def test_include_url(self): flattened = [ SafeLineLoader.strip_metadata(spec) for spec in yaml_provider.parse_providers( '', [ self.INLINE_PROVIDER, { 'include': 'file:///' + self.to_include }, ]) ] self.assertEqual([ self.INLINE_PROVIDER, self.INCLUDED_PROVIDER, ], flattened) @mock.patch( 'apache_beam.yaml.yaml_provider.ExternalProvider.provider_from_spec', lambda _, x: x) def test_nested_include(self): flattened = [ SafeLineLoader.strip_metadata(spec) for spec in yaml_provider.parse_providers( '', [ self.INLINE_PROVIDER, { 'include': self.to_include_nested }, ]) ] self.assertEqual([ self.INLINE_PROVIDER, self.INCLUDED_PROVIDER, self.EXTRA_PROVIDER, ], flattened) class YamlDefinedProider(unittest.TestCase): def test_yaml_define_provider(self): providers = ''' - type: yaml transforms: Range: config_schema: properties: end: {type: integer} requires_inputs: false body: | type: Create config: elements: {% for ix in range(end) %} - {{ix}} {% endfor %} Power: config_schema: properties: n: {type: integer} body: type: MapToFields config: language: python append: true fields: power: "element**{{n}}" ''' pipeline = ''' type: chain transforms: - type: Range config: end: 4 - type: Power config: n: 2 ''' with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( pipeline, providers=yaml_provider.parse_providers( '', yaml.load(providers, Loader=SafeLineLoader))) assert_that( result | beam.Map(lambda x: (x.element, x.power)), equal_to([(0, 0), (1, 1), (2, 4), (3, 9)])) def test_recursive(self): providers = ''' - type: yaml transforms: Factorial: config_schema: properties: n: {type: integer} requires_inputs: false body: | {% if n <= 1 %} type: Create config: elements: - {value: 1} {% else %} type: chain transforms: - type: Factorial config: n: {{n-1}} - type: MapToFields name: Multiply config: language: python fields: value: value * {{n}} {% endif %} ''' pipeline = ''' type: Factorial config: n: 5 ''' with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( pipeline, providers=yaml_provider.parse_providers( '', yaml.load(providers, Loader=SafeLineLoader))) assert_that(result | beam.Map(lambda x: x.value), equal_to([120])) class PythonProviderDepsTest(unittest.TestCase): def test_env_package_sensitive(self): self.assertNotEqual( yaml_provider.PypiExpansionService._key('base', ['pkg1']), yaml_provider.PypiExpansionService._key('base', ['pkg2'])) def test_env_base_sensitive(self): self.assertNotEqual( yaml_provider.PypiExpansionService._key('base1', ['pkg']), yaml_provider.PypiExpansionService._key('base2', ['pkg'])) def test_env_order_invariant(self): self.assertEqual( yaml_provider.PypiExpansionService._key('base', ['pkg1', 'pkg2']), yaml_provider.PypiExpansionService._key('base', ['pkg2', 'pkg1'])) def test_env_path_invariant(self): with tempfile.TemporaryDirectory() as tmpdir: os.mkdir(os.path.join(tmpdir, 'a')) pkgA = os.path.join(tmpdir, 'a', 'pkg.tgz') os.mkdir(os.path.join(tmpdir, 'b')) pkgB = os.path.join(tmpdir, 'b', 'pkg.tgz') with open(pkgA, 'w') as fout: fout.write('content') with open(pkgB, 'w') as fout: fout.write('content') self.assertEqual( yaml_provider.PypiExpansionService._key('base', [pkgA]), yaml_provider.PypiExpansionService._key('base', [pkgB])) def test_env_content_sensitive(self): with tempfile.TemporaryDirectory() as tmpdir: pkg = os.path.join(tmpdir, 'pkg.tgz') with open(pkg, 'w') as fout: fout.write('content') before = yaml_provider.PypiExpansionService._key('base', [pkg]) with open(pkg, 'w') as fout: fout.write('new content') after = yaml_provider.PypiExpansionService._key('base', [pkg]) self.assertNotEqual(before, after) class JoinUrlOrFilepathTest(unittest.TestCase): def test_join_url_relative_path(self): self.assertEqual( yaml_provider._join_url_or_filepath('http://example.com/a', 'b/c.yaml'), 'http://example.com/b/c.yaml') self.assertEqual( yaml_provider._join_url_or_filepath( 'http://example.com/a/', 'b/c.yaml'), 'http://example.com/a/b/c.yaml') # use os.path.join to mock gcs filesystem split and join. with mock.patch('apache_beam.io.filesystems.FileSystems.split', new=lambda x: ("gs://bucket", x.removeprefix("gs://bucket/"))): with mock.patch('apache_beam.io.filesystems.FileSystems.join', new=lambda *args: '/'.join(args)): self.assertEqual( yaml_provider._join_url_or_filepath('gs://bucket', 'b/c.yaml'), 'gs://bucket/b/c.yaml') self.assertEqual( yaml_provider._join_url_or_filepath('gs://bucket/', 'b/c.yaml'), 'gs://bucket/b/c.yaml') self.assertEqual( yaml_provider._join_url_or_filepath('gs://bucket/a', 'b/c.yaml'), 'gs://bucket/b/c.yaml') def test_join_filepath_relative_path(self): if sys.platform != 'win32': self.assertEqual( yaml_provider._join_url_or_filepath('/a/b/', 'c/d.yaml'), '/a/b/c/d.yaml') self.assertEqual( yaml_provider._join_url_or_filepath('/a/b', 'c/d.yaml'), '/a/c/d.yaml') else: self.assertEqual( yaml_provider._join_url_or_filepath('C:\\a\\b\\', 'c\\d.yaml'), 'C:\\a\\b\\c\\d.yaml') self.assertEqual( yaml_provider._join_url_or_filepath('C:\\a\\b', 'c\\d.yaml'), 'C:\\a\\c\\d.yaml') def test_absolute_path(self): self.assertEqual( yaml_provider._join_url_or_filepath( 'gs://bucket/a', 'gs://bucket/b/c.yaml'), 'gs://bucket/b/c.yaml') if sys.platform != 'win32': self.assertEqual( yaml_provider._join_url_or_filepath('/a/b', '/c/d.yaml'), '/c/d.yaml') def test_different_scheme(self): self.assertEqual( yaml_provider._join_url_or_filepath( 'http://example.com/a', 'gs://bucket/b/c.yaml'), 'gs://bucket/b/c.yaml') def test_empty_base(self): self.assertEqual( yaml_provider._join_url_or_filepath('', 'a/b.yaml'), 'a/b.yaml') self.assertEqual( yaml_provider._join_url_or_filepath(None, 'a/b.yaml'), 'a/b.yaml') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() class YamlProvidersCreateTest(unittest.TestCase): def test_create_mixed_types(self): with beam.Pipeline() as p: # A mix of a primitive (Row(element=1)) and a dict (Row(a=2)) result = p | YamlProviders.create([1, {"a": 2}]) assert_that( result | beam.Map(lambda x: sorted(x._asdict().items())), equal_to([ [('a', None), ('element', 1)], [('a', 2), ('element', None)], ])) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_specifiable.py0000644000175100017510000000420315206113033022461 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # from apache_beam.io.filesystems import FileSystems from apache_beam.ml.anomaly.specifiable import Spec from apache_beam.ml.anomaly.transforms import AnomalyDetection from apache_beam.ml.anomaly.transforms import Specifiable from apache_beam.utils import python_callable from apache_beam.yaml.yaml_provider import InlineProvider def maybe_make_specifiable(v): if isinstance(v, dict): if "type" in v and "config" in v: return Specifiable.from_spec( Spec(type=v["type"], config=maybe_make_specifiable(v["config"]))) if "callable" in v: if "path" in v or "name" in v: raise ValueError( "Cannot specify 'callable' with 'path' and 'name' for function.") else: return python_callable.PythonCallableWithSource(v["callable"]) if "path" in v and "name" in v: return python_callable.PythonCallableWithSource.load_from_script( FileSystems.open(v["path"]).read().decode(), v["name"]) ret = {k: maybe_make_specifiable(v[k]) for k in v} return ret else: return v class SpecProvider(InlineProvider): def create_transform(self, type, args, yaml_create_transform): return self._transform_factories[type]( **{ k: maybe_make_specifiable(v) for k, v in args.items() }) def create_spec_providers(): return SpecProvider({"AnomalyDetection": AnomalyDetection}) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_specifiable_test.py0000644000175100017510000000757215206113033023534 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest from typing import Callable import apache_beam as beam from apache_beam.ml.anomaly.base import AnomalyDetector from apache_beam.ml.anomaly.specifiable import specifiable from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.utils import python_callable from apache_beam.yaml.yaml_transform import YamlTransform TEST_PROVIDERS = { 'PyMap': lambda fn: beam.Map(python_callable.PythonCallableWithSource(fn)), } @specifiable class FakeDetector(AnomalyDetector): # pylint: disable=unused-variable def __init__(self, fn: Callable): super().__init__() self._fn = fn def learn_one(self, x: beam.Row) -> None: pass def score_one(self, x: beam.Row) -> float: v = next(iter(x)) return self._fn(v) class YamlSpecifiableTransformTest(unittest.TestCase): def test_specifiable_transform(self): TRAIN_DATA = [ (0, beam.Row(x=1)), (0, beam.Row(x=2)), (0, beam.Row(x=2)), (0, beam.Row(x=4)), (0, beam.Row(x=9)), ] pipeline_options = beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle') # Pin to FnApiRunner since this requires data from create to be # ppassed to anomaly detection in a certain order with beam.Pipeline('FnApiRunner', options=pipeline_options) as p: result = p | beam.Create(TRAIN_DATA) | YamlTransform( ''' type: chain transforms: - type: AnomalyDetection config: detector: type: 'ZScore' config: sub_stat_tracker: type: 'IncSlidingMeanTracker' config: window_size: 5 stdev_tracker: type: 'IncSlidingStdevTracker' config: window_size: 5 - type: PyMap config: fn: "lambda x: (x[1].predictions[0].label)" ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([-2, -2, 0, 1, 1])) def test_specifiable_transform_with_callable(self): TRAIN_DATA = [ (0, beam.Row(x=1)), (0, beam.Row(x=2)), (0, beam.Row(x=2)), (0, beam.Row(x=4)), (0, beam.Row(x=9)), ] with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | beam.Create(TRAIN_DATA) | YamlTransform( ''' type: chain transforms: - type: AnomalyDetection config: detector: type: 'FakeDetector' config: fn: callable: "lambda x: x * 10.0" - type: PyMap config: fn: "lambda x: (x[1].predictions[0].score)" ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([10.0, 20.0, 20.0, 40.0, 90.0])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_testing.py0000644000175100017510000004423115206113033021675 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import collections import functools import json import random import unittest import uuid from typing import Mapping from typing import Optional from typing import TypeVar from typing import Union import yaml import apache_beam as beam from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.yaml import yaml_provider from apache_beam.yaml import yaml_transform from apache_beam.yaml import yaml_utils class YamlTestCase(unittest.TestCase): def __init__(self, pipeline_spec, test_spec, options, fix_tests): super().__init__() self._pipeline_spec = pipeline_spec self._test_spec = test_spec self._options = options self._fix_tests = fix_tests self._fixes = None def runTest(self): self._fixes = run_test( self._pipeline_spec, self._test_spec, self._options, self._fix_tests) def fixed_test(self): fixed_test_spec = yaml_transform.SafeLineLoader.strip_metadata( self._test_spec) if self._fixes: expectation_by_id = {(loc, expectation['name']): expectation for loc in ('expected_inputs', 'expected_outputs') for expectation in fixed_test_spec.get(loc, [])} for name_loc, values in self._fixes.items(): expectation_by_id[name_loc]['elements'] = sorted(values, key=json.dumps) return fixed_test_spec def id(self): return ( self._test_spec.get('name', 'unknown') + f' (line {yaml_transform.SafeLineLoader.get_line(self._test_spec)})') def __str__(self): return self.id() def run_test(pipeline_spec, test_spec, options=None, fix_failures=False): if isinstance(pipeline_spec, str): pipeline_spec_dict = yaml.load( pipeline_spec, Loader=yaml_utils.SafeLineLoader) else: pipeline_spec_dict = pipeline_spec processed_pipeline_spec = _preprocess_for_testing(pipeline_spec_dict) transform_spec, recording_ids = inject_test_tranforms( processed_pipeline_spec, test_spec, fix_failures) allowed_sources = set(test_spec.get('allowed_sources', []) + ['Create']) for transform in transform_spec['transforms']: name_or_type = transform.get('name', transform['type']) if (not yaml_transform.empty_if_explicitly_empty(transform.get('input', [])) and not transform.get('name') in allowed_sources and not transform['type'] in allowed_sources): raise ValueError( f'Non-mocked source {name_or_type} ' f'at line {yaml_transform.SafeLineLoader.get_line(transform)}') if options is None: options = beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle', **yaml_transform.SafeLineLoader.strip_metadata( pipeline_spec_dict.get('options', {}))) providers = yaml_provider.merge_providers( yaml_provider.parse_providers( '', pipeline_spec_dict.get('providers', [])), { 'AssertEqualAndRecord': yaml_provider.as_provider_list( 'AssertEqualAndRecord', AssertEqualAndRecord) }) with beam.Pipeline(options=options) as p: _ = p | yaml_transform.YamlTransform(transform_spec, providers=providers) if fix_failures: fixes = {} for recording_id in recording_ids: if AssertEqualAndRecord.has_recorded_result(recording_id): fixes[recording_id[1:]] = [ _try_row_as_dict(row) for row in AssertEqualAndRecord.get_recorded_result(recording_id) ] AssertEqualAndRecord.remove_recorded_result(recording_id) return fixes def _preprocess_for_testing(pipeline_spec): spec = yaml_transform.pipeline_as_composite(pipeline_spec['pipeline']) # These are idempotent, so it's OK to do them preemptively. for phase in [ yaml_transform.ensure_transforms_have_types, yaml_transform.preprocess_source_sink, yaml_transform.preprocess_chain, yaml_transform.tag_explicit_inputs, yaml_transform.normalize_inputs_outputs, ]: spec = yaml_transform.apply_phase(phase, spec) return spec def validate_test_spec(test_spec): if not isinstance(test_spec, dict): raise TypeError( f'Test specification must be an object, got {type(test_spec)}') identifier = ( test_spec.get('name', 'unknown') + f' at line {yaml_transform.SafeLineLoader.get_line(test_spec)}') if not isinstance(test_spec.get('allowed_sources', []), list): raise TypeError( f'allowed_sources of test specification {identifier} ' f'must be a list, got {type(test_spec["allowed_sources"])}') if (not test_spec.get('expected_outputs', []) and not test_spec.get('expected_inputs', [])): raise ValueError( f'test specification {identifier} ' f'must have at least one expected_outputs or expected_inputs') unknown_attrs = set( yaml_transform.SafeLineLoader.strip_metadata(test_spec).keys()) - set([ 'name', 'mock_inputs', 'mock_outputs', 'expected_outputs', 'expected_inputs', 'allowed_sources', ]) if unknown_attrs: raise ValueError( f'test specification {identifier} ' f'has unknown attributes {list(unknown_attrs)}') for attr_type in ('mock_inputs', 'mock_outputs', 'expected_outputs', 'expected_inputs'): attr = test_spec.get(attr_type, []) if not isinstance(attr, list): raise TypeError( f'{attr_type} of test specification {identifier} ' f'must be a list, got {type(attr_type)}') for ix, attr_item in enumerate(attr): if not isinstance(attr_item, dict): raise TypeError( f'{attr_type} {ix} of test specification {identifier} ' f'must be an object, got {type(attr_item)}') if 'name' not in attr_item: raise TypeError( f'{attr_type} {ix} of test specification {identifier} ' f'missing a name') if 'elements' not in attr_item: raise TypeError( f'{attr_type} {ix} of test specification {identifier} ' f'missing a elements') if not isinstance(attr_item['elements'], list): raise TypeError( f'{attr_type} {ix} of test specification {identifier} ' f'must be a list, got {type(attr_item["elements"])}') def inject_test_tranforms(spec, test_spec, fix_failures): validate_test_spec(test_spec) scope = yaml_transform.LightweightScope(spec['transforms']) mocked_inputs_by_id = { scope.get_transform_id(mock_input['name']): mock_input for mock_input in test_spec.get('mock_inputs', []) } mocked_outputs_by_id = _composite_key_to_nested({ scope.get_transform_id_and_output_name(mock_output['name']): mock_output for mock_output in test_spec.get('mock_outputs', []) }) recording_id_prefix = str(uuid.uuid4()) recording_ids = [] transforms = [] @functools.cache def create_inputs(transform_id: str) -> InputsType: def require_output_or_outputs(name_or_names): if isinstance(name_or_names, str): return require_output(name_or_names) else: return [require_output(name) for name in name_or_names] if transform_id in mocked_inputs_by_id: return create_mocked_input(transform_id) else: input_spec = scope.get_transform_spec(transform_id)['input'] return { tag: require_output_or_outputs(input_ref) for tag, input_ref in yaml_transform.empty_if_explicitly_empty( input_spec).items() } def require_output(name: str) -> str: # The same output may be referenced under different names. # Normalize before we cache. transform_id, tag = scope.get_transform_id_and_output_name(name) return _require_output(transform_id, tag) or name @functools.cache def _require_output(transform_id: str, tag: str) -> Optional[str]: if transform_id in mocked_outputs_by_id: if tag not in mocked_outputs_by_id[transform_id]: name = next(iter( mocked_outputs_by_id[transform_id].values()))['name'].split('.')[0] raise ValueError( f'Unmocked output {tag} of {name}.' 'If any used output is mocked all used outputs must be mocked.') return create_mocked_output(transform_id, tag) else: _use_transform(transform_id) return None # Use original name. @functools.cache def _use_transform(transform_id: str) -> None: transform_spec = dict(scope.get_transform_spec(transform_id)) transform_spec['input'] = create_inputs(transform_id) transforms.append(transform_spec) @functools.cache def create_mocked_input(transform_id: str) -> str: transform = create_create( f'MockInput[{mocked_inputs_by_id[transform_id]["name"]}]', mocked_inputs_by_id[transform_id]['elements'], mocked_inputs_by_id[transform_id]['name']) transforms.append(transform) return transform['__uuid__'] @functools.cache def create_mocked_output(transform_id: str, tag: str) -> str: transform = create_create( f'MockOutput[{mocked_outputs_by_id[transform_id][tag]["name"]}]', mocked_outputs_by_id[transform_id][tag]['elements'], mocked_outputs_by_id[transform_id][tag]['name']) transforms.append(transform) return transform['__uuid__'] def create_create(name, elements, line_source): return { '__uuid__': yaml_utils.SafeLineLoader.create_uuid(), '__line__': yaml_utils.SafeLineLoader.get_line(line_source), 'name': name, 'type': 'Create', 'config': { 'elements': elements, }, } def create_assertion(name, inputs, elements, recording_id, line_source): return { '__uuid__': yaml_utils.SafeLineLoader.create_uuid(), '__line__': yaml_utils.SafeLineLoader.get_line(line_source), 'name': name, 'input': inputs, 'type': 'AssertEqualAndRecord', 'config': { 'elements': elements, 'recording_id': recording_id, }, } for expected_output in test_spec.get('expected_outputs', []): if fix_failures: recording_id = ( recording_id_prefix, 'expected_outputs', expected_output['name']) recording_ids.append(recording_id) else: recording_id = None require_output(expected_output['name']) transforms.append( create_assertion( f'CheckExpectedOutput[{expected_output["name"]}]', expected_output['name'], expected_output['elements'], recording_id, expected_output['name'])) for expected_input in test_spec.get('expected_inputs', []): if fix_failures: recording_id = ( recording_id_prefix, 'expected_inputs', expected_input['name']) recording_ids.append(recording_id) else: recording_id = None transform_id = scope.get_transform_id(expected_input['name']) transforms.append( create_assertion( f'CheckExpectedInput[{expected_input["name"]}]', create_inputs(transform_id), expected_input['elements'], recording_id, expected_input['name'])) return { '__uuid__': yaml_utils.SafeLineLoader.create_uuid(), '__line__': 0, 'type': 'composite', 'transforms': transforms, }, recording_ids class AssertEqualAndRecord(beam.PTransform): _recorded_results = {} @classmethod def store_recorded_result(cls, recording_id, value): assert recording_id not in cls._recorded_results cls._recorded_results[recording_id] = value @classmethod def has_recorded_result(cls, recording_id): return recording_id in cls._recorded_results @classmethod def get_recorded_result(cls, recording_id): return cls._recorded_results[recording_id] @classmethod def remove_recorded_result(cls, recording_id): del cls._recorded_results[recording_id] def __init__(self, elements, recording_id): self._elements = elements self._recording_id = recording_id def expand(self, pcoll): # Convert elements to rows outside the matcher to avoid capturing # any grpc channels that might be created during the conversion expected_rows = yaml_provider.dicts_to_rows(self._elements) recording_id = self._recording_id # Create a serializable matcher function that doesn't capture # any external references that might contain grpc channels class SerializableMatcher: def __init__(self, expected_rows, recording_id): self.expected_rows = expected_rows self.recording_id = recording_id self.equal_to_matcher = equal_to(expected_rows) def __call__(self, actual): try: self.equal_to_matcher(actual) except Exception: if self.recording_id: AssertEqualAndRecord.store_recorded_result( tuple(self.recording_id), actual) else: raise matcher = SerializableMatcher(expected_rows, recording_id) return assert_that( pcoll | beam.Map(lambda row: beam.Row(**row._asdict())), matcher) def create_test( pipeline_spec, options=None, max_num_inputs=40, min_num_outputs=3): if isinstance(pipeline_spec, str): pipeline_spec = yaml.load(pipeline_spec, Loader=yaml_utils.SafeLineLoader) transform_spec = _preprocess_for_testing(pipeline_spec) if options is None: options = beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle', **yaml_transform.SafeLineLoader.strip_metadata( pipeline_spec.get('options', {}))) providers = yaml_provider.merge_providers( yaml_provider.parse_providers('', pipeline_spec.get('providers', [])), { 'AssertEqualAndRecord': yaml_provider.as_provider_list( 'AssertEqualAndRecord', AssertEqualAndRecord) }) def get_name(transform): if 'name' in transform: return str(transform['name']) else: if sum(1 for t in transform_spec['transforms'] if t['type'] == transform['type']) > 1: raise ValueError('Ambiguous unnamed transform {transform["type"]}') return str(transform['type']) input_transforms = [ t for t in transform_spec['transforms'] if t['type'] != 'Create' and not yaml_transform.empty_if_explicitly_empty(t.get('input', [])) ] mock_outputs = [{ 'name': get_name(t), 'elements': [ _try_row_as_dict(row) for row in _first_n(t, options, max_num_inputs, providers) ], } for t in input_transforms] output_transforms = [ t for t in transform_spec['transforms'] if t['type'] == 'LogForTesting' or yaml_transform.empty_if_explicitly_empty(t.get('output', [])) or t['type'].startswith('Write') ] expected_inputs = [{ 'name': get_name(t), 'elements': [], } for t in output_transforms] if not expected_inputs: # TODO: Optionally take this as a parameter. raise ValueError('No output transforms detected.') num_inputs = min_num_outputs while True: test_spec = { 'mock_outputs': [{ 'name': t['name'], 'elements': random.sample( t['elements'], min(len(t['elements']), num_inputs)), } for t in mock_outputs], 'expected_inputs': expected_inputs, } fixes = run_test(pipeline_spec, test_spec, options, fix_failures=True) if len(fixes) < len(output_transforms): actual_output_size = 0 else: actual_output_size = min(len(e) for e in fixes.values()) if actual_output_size >= min_num_outputs: break elif num_inputs == max_num_inputs: break else: num_inputs = min(2 * num_inputs, max_num_inputs) for expected_input in test_spec['expected_inputs']: if ('expected_inputs', expected_input['name']) in fixes: expected_input['elements'] = fixes['expected_inputs', expected_input['name']] return test_spec class _DoneException(Exception): pass class RecordElements(beam.PTransform): _recorded_results = collections.defaultdict(list) def __init__(self, n): self._n = n self._id = str(uuid.uuid4()) def get_and_remove(self): listing = RecordElements._recorded_results[self._id] del RecordElements._recorded_results[self._id] return listing def expand(self, pcoll): def record(element): listing = RecordElements._recorded_results[self._id] if len(listing) < self._n: listing.append(element) else: raise _DoneException() return pcoll | beam.Map(record) def _first_n(transform_spec, options, n, providers=None): recorder = RecordElements(n) if providers is None: providers = { 'AssertEqualAndRecord': yaml_provider.as_provider_list( 'AssertEqualAndRecord', AssertEqualAndRecord) } try: with beam.Pipeline(options=options) as p: _ = ( p | yaml_transform.YamlTransform(transform_spec, providers=providers) | recorder) except _DoneException: pass except Exception as exn: # Runners don't always raise a faithful exception type. if not '_DoneException' in str(exn): raise return recorder.get_and_remove() K1 = TypeVar('K1') K2 = TypeVar('K2') V = TypeVar('V') InputsType = dict[str, Union[str, list[str]]] def _composite_key_to_nested( d: Mapping[tuple[K1, K2], V]) -> Mapping[K1, Mapping[K2, V]]: nested = collections.defaultdict(dict) for (k1, k2), v in d.items(): nested[k1][k2] = v return nested def _try_row_as_dict(row): try: return row._asdict() except AttributeError: return row # Linter: No need for unittest.main here. ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_testing_test.py0000644000175100017510000002717315206113033022742 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import os import tempfile import unittest # isort is fighting with yapf here. # isort: off from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.yaml import yaml_testing # Note that executing this pipeline will actually fail if the untested # transforms are not in fact elided. SIMPLE_PIPELINE = ''' pipeline: type: chain transforms: - type: ReadFromText name: MyRead config: path: 'does_not_exist.txt' - type: MapToFields name: MyMapToFields config: append: True language: python fields: square: element**2 - type: MapToFields name: ToBeExcluded config: language: python fields: bad: callable: | def non_shall_pass(row): raise ValueError(row) ''' class YamlTestingTest(unittest.TestCase): def test_expected_outputs(self): yaml_testing.run_test( SIMPLE_PIPELINE, { 'mock_inputs': [{ 'name': 'MyMapToFields', 'elements': [1, 2, 3], }], 'expected_outputs': [{ 'name': 'MyMapToFields', 'elements': [ { 'element': 1, 'square': 1 }, { 'element': 2, 'square': 4 }, { 'element': 3, 'square': 9 }, ] }] }) def test_expected_outputs_fails_correctly(self): with self.assertRaisesRegex(Exception, 'unexpected.*9'): yaml_testing.run_test( SIMPLE_PIPELINE, { 'mock_inputs': [{ 'name': 'MyMapToFields', 'elements': [1, 2, 3], }], 'expected_outputs': [{ 'name': 'MyMapToFields', 'elements': [ { 'element': 1, 'square': 1 }, { 'element': 2, 'square': 4 }, ] }] }) def test_expected_inputs(self): yaml_testing.run_test( SIMPLE_PIPELINE, { 'mock_outputs': [{ 'name': 'MyRead', 'elements': [1, 2, 3], }], 'expected_inputs': [{ 'name': 'ToBeExcluded', 'elements': [ { 'element': 1, 'square': 1 }, { 'element': 2, 'square': 4 }, { 'element': 3, 'square': 9 }, ] }] }) def test_expected_inputs_fails_correctly(self): with self.assertRaisesRegex(Exception, 'unexpected.*9'): yaml_testing.run_test( SIMPLE_PIPELINE, { 'mock_outputs': [{ 'name': 'MyRead', 'elements': [1, 2, 3], }], 'expected_inputs': [{ 'name': 'ToBeExcluded', 'elements': [ { 'element': 1, 'square': 1 }, { 'element': 2, 'square': 4 }, ] }] }) def test_unmocked_inputs(self): with self.assertRaisesRegex(Exception, 'Non-mocked source MyRead'): yaml_testing.run_test( SIMPLE_PIPELINE, { 'expected_inputs': [{ 'name': 'ToBeExcluded', 'elements': [ { 'element': 1, 'square': 1 }, { 'element': 2, 'square': 4 }, { 'element': 3, 'square': 9 }, ] }] }) def test_fixes(self): fixes = yaml_testing.run_test( SIMPLE_PIPELINE, { 'mock_outputs': [{ 'name': 'MyRead', 'elements': [1, 2, 3], }], 'expected_inputs': [{ 'name': 'ToBeExcluded', 'elements': [ { 'element': 1, 'square': 1 }, { 'element': 2, 'square': 4 }, ] }] }, fix_failures=True) self.assertEqual( fixes, {('expected_inputs', 'ToBeExcluded'): [ dict(element=1, square=1), dict(element=2, square=4), dict(element=3, square=9), ]}) def test_create(self): with tempfile.TemporaryDirectory() as tmpdir: input_path = os.path.join(tmpdir, 'input.csv') with open(input_path, 'w') as fout: fout.write('a,b,c\n') for ix in range(1000): fout.write(f'{ix % 5},{ix},Ccc\n') pipeline = f''' pipeline: type: chain transforms: - type: ReadFromCsv config: path: {input_path} - type: Filter config: language: python keep: a == 1 - type: WriteToSink ''' test_spec = yaml_testing.create_test( pipeline, max_num_inputs=100, min_num_outputs=5) self.assertEqual(len(test_spec['mock_outputs']), 1) self.assertEqual(len(test_spec['expected_inputs']), 1) self.assertGreaterEqual(len(test_spec['expected_inputs'][0]['elements']), 5) yaml_testing.run_test(pipeline, test_spec) @unittest.skipIf( TestPipeline().get_pipeline_options().view_as(StandardOptions).runner is None, 'Do not run this test on precommit suites.') def test_join_transform_serialization(self): """Test that Join transforms work with YAML testing framework and cloudpickle. This test validates the fix for the grpc channel serialization issue that was causing TypeError: no default __reduce__ due to non-trivial __cinit__ when using Join transforms with the YAML testing framework. """ join_pipeline = ''' pipeline: transforms: - type: Create name: Create1 config: elements: - ride_id: "1" pickup_location: "downtown" - ride_id: "2" pickup_location: "airport" - type: Create name: Create2 config: elements: - ride_id: "1" dropoff_location: "mall" - ride_id: "2" dropoff_location: "hotel" - type: Join name: JoinRides input: pickup: Create1 dropoff: Create2 config: equalities: ride_id type: inner fields: pickup: [ride_id, pickup_location] dropoff: [dropoff_location] - type: LogForTesting name: LogResult input: JoinRides ''' # Test with expected_inputs to validate the Join transform output yaml_testing.run_test( join_pipeline, { 'expected_inputs': [{ 'name': 'LogResult', 'elements': [{ 'ride_id': '1', 'pickup_location': 'downtown', 'dropoff_location': 'mall' }, { 'ride_id': '2', 'pickup_location': 'airport', 'dropoff_location': 'hotel' }] }] }) # Test with mock_outputs to validate Join transform can handle mocked inputs yaml_testing.run_test( join_pipeline, { 'mock_outputs': [{ 'name': 'Create1', 'elements': [{ 'ride_id': '3', 'pickup_location': 'station' }] }, { 'name': 'Create2', 'elements': [{ 'ride_id': '3', 'dropoff_location': 'office' }] }], 'expected_inputs': [{ 'name': 'LogResult', 'elements': [{ 'ride_id': '3', 'pickup_location': 'station', 'dropoff_location': 'office' }] }] }) def test_toplevel_providers(self): yaml_testing.run_test( ''' pipeline: type: chain transforms: - type: Create config: elements: [1, 2, 3] - type: MyDoubler providers: - type: yaml transforms: MyDoubler: body: type: MapToFields config: language: python fields: doubled: element * 2 ''', { 'expected_outputs': [{ 'name': 'MyDoubler', 'elements': [{ 'doubled': 2 }, { 'doubled': 4 }, { 'doubled': 6 }] }] }) def test_create_with_external_providers(self): """Test that create_test works with external providers defined in the pipeline spec. This test validates the fix for issue #37136 where external providers defined in YAML files were not recognized when running tests. """ pipeline = ''' pipeline: type: chain transforms: - type: Create config: elements: - {a: 1, b: 2} - {a: 2, b: 3} - {a: 3, b: 4} - {a: 4, b: 5} - {a: 5, b: 6} - type: MyCustomTransform - type: LogForTesting providers: - type: yaml transforms: MyCustomTransform: body: type: MapToFields config: language: python fields: sum_ab: a + b ''' test_spec = yaml_testing.create_test( pipeline, max_num_inputs=10, min_num_outputs=3) self.assertEqual(len(test_spec['expected_inputs']), 1) self.assertGreaterEqual(len(test_spec['expected_inputs'][0]['elements']), 3) yaml_testing.run_test(pipeline, test_spec) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_transform.py0000644000175100017510000015124415206113033022236 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import collections import datetime import functools import json import logging import os import pprint import re from collections.abc import Iterable from collections.abc import Mapping from typing import Any import jinja2 import yaml import apache_beam as beam from apache_beam.io.filesystems import FileSystems from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.transforms.fully_qualified_named_transform import FullyQualifiedNamedTransform from apache_beam.typehints import schemas from apache_beam.typehints import typehints from apache_beam.yaml import json_utils from apache_beam.yaml import yaml_provider from apache_beam.yaml import yaml_utils from apache_beam.yaml.yaml_combine import normalize_combine from apache_beam.yaml.yaml_mapping import Validate from apache_beam.yaml.yaml_mapping import normalize_mapping from apache_beam.yaml.yaml_mapping import validate_generic_expressions from apache_beam.yaml.yaml_utils import SafeLineLoader __all__ = ["YamlTransform"] _LOGGER = logging.getLogger(__name__) yaml_provider.fix_pycallable() try: import jsonschema except ImportError: jsonschema = None @functools.lru_cache def pipeline_schema(strictness): with open(yaml_utils.locate_data_file('pipeline.schema.yaml')) as yaml_file: pipeline_schema = yaml.safe_load(yaml_file) if strictness == 'per_transform': transform_schemas_path = yaml_utils.locate_data_file( 'transforms.schema.yaml') if not os.path.exists(transform_schemas_path): raise RuntimeError( "Please run " "python -m apache_beam.yaml.generate_yaml_docs " f"--schema_file='{transform_schemas_path}' " "to run with transform-specific validation.") with open(transform_schemas_path) as fin: pipeline_schema['$defs']['transform']['allOf'].extend(yaml.safe_load(fin)) return pipeline_schema def _closest_line(o, path): best_line = SafeLineLoader.get_line(o) for step in path: o = o[step] maybe_line = SafeLineLoader.get_line(o) if maybe_line != 'unknown': best_line = maybe_line return best_line def validate_against_schema(pipeline, strictness): try: jsonschema.validate(pipeline, pipeline_schema(strictness)) except jsonschema.ValidationError as exn: exn.message += f" around line {_closest_line(pipeline, exn.path)}" # validation message for chain-type transform if (exn.schema_path[-1] == 'not' and exn.schema_path[-2] in ['input', 'output']): exn.message = ( f"'{exn.schema_path[-2]}' should not be used " "along with 'chain' type transforms. " + exn.message) raise exn def memoize_method(func): def wrapper(self, *args): if not hasattr(self, '_cache'): self._cache = {} key = func.__name__, args if key not in self._cache: self._cache[key] = func(self, *args) return self._cache[key] return wrapper def only_element(xs): x, = xs return x # These allow a user to explicitly pass no input to a transform (i.e. use it # as a root transform) without an error even if the transform is not known to # handle it. def explicitly_empty(): return {'__explicitly_empty__': None} def is_explicitly_empty(io): return io == explicitly_empty() def is_empty(io): return not io or is_explicitly_empty(io) def empty_if_explicitly_empty(io): if is_explicitly_empty(io): return {} else: return io class LightweightScope(object): def __init__(self, transforms): self._transforms = transforms self._transforms_by_uuid = {t['__uuid__']: t for t in self._transforms} self._uuid_by_name = collections.defaultdict(set) for spec in self._transforms: if 'name' in spec: self._uuid_by_name[spec['name']].add(spec['__uuid__']) if 'type' in spec: self._uuid_by_name[spec['type']].add(spec['__uuid__']) def get_transform_id_and_output_name(self, name): if '.' in name: transform_name, output = name.rsplit('.', 1) else: transform_name, output = name, None return self.get_transform_id(transform_name), output def get_transform_id(self, transform_name): if transform_name in self._transforms_by_uuid: return transform_name else: candidates = self._uuid_by_name[transform_name] if not candidates: raise ValueError( f'Unknown transform at line ' f'{SafeLineLoader.get_line(transform_name)}: {transform_name}') elif len(candidates) > 1: raise ValueError( f'Ambiguous transform at line ' f'{SafeLineLoader.get_line(transform_name)}: {transform_name}') else: return only_element(candidates) def get_transform_spec(self, transform_name_or_id): return self._transforms_by_uuid[self.get_transform_id(transform_name_or_id)] class Scope(LightweightScope): """To look up PCollections (typically outputs of prior transforms) by name.""" def __init__( self, root, inputs: Mapping[str, Any], transforms: Iterable[dict], providers: Mapping[str, Iterable[yaml_provider.Provider]], input_providers: Iterable[yaml_provider.Provider]): super().__init__(transforms) self.root = root self._inputs = inputs self.providers = providers self._seen_names: set[str] = set() self.input_providers = input_providers self._all_followers = None def followers(self, transform_name): if self._all_followers is None: self._all_followers = collections.defaultdict(list) # TODO(yaml): Also trace through outputs and composites. for transform in self._transforms: if transform['type'] != 'composite': for input in empty_if_explicitly_empty(transform['input']).values(): if input not in self._inputs: transform_id, _ = self.get_transform_id_and_output_name(input) self._all_followers[transform_id].append(transform['__uuid__']) return self._all_followers[self.get_transform_id(transform_name)] def compute_all(self): for transform_id in self._transforms_by_uuid.keys(): self.compute_outputs(transform_id) def get_pcollection(self, name): if name in self._inputs: return self._inputs[name] elif '.' in name: transform, output = name.rsplit('.', 1) outputs = self.get_outputs(transform) if output in outputs: return outputs[output] elif len(outputs) == 1 and outputs[next(iter(outputs))].tag == output: return outputs[next(iter(outputs))] else: raise ValueError( f'Unknown output {repr(output)} ' f'at line {SafeLineLoader.get_line(name)}: ' f'{transform} only has outputs {list(outputs.keys())}') else: outputs = self.get_outputs(name) if len(outputs) == 1: return only_element(outputs.values()) else: error_output = self._transforms_by_uuid[self.get_transform_id( name)]['config'].get('error_handling', {}).get('output') if error_output and error_output in outputs and len(outputs) == 2: return next( output for tag, output in outputs.items() if tag != error_output) raise ValueError( f'Ambiguous output at line {SafeLineLoader.get_line(name)}: ' f'{name} has outputs {list(outputs.keys())}') def get_outputs(self, transform_name): return self.compute_outputs(self.get_transform_id(transform_name)) @memoize_method def compute_outputs(self, transform_id): return expand_transform(self._transforms_by_uuid[transform_id], self) def best_provider( self, t, input_providers: yaml_provider.Iterable[yaml_provider.Provider]): if isinstance(t, dict): spec = t else: spec = self._transforms_by_uuid[self.get_transform_id(t)] possible_providers = [] unavailable_provider_messages = [] for p in self.providers[spec['type']]: is_available = p.available() if is_available: possible_providers.append(p) else: reason = getattr(is_available, 'reason', 'no reason given') unavailable_provider_messages.append( f'{p.__class__.__name__} ({reason})') if not possible_providers: if unavailable_provider_messages: unavailable_provider_message = ( '\nThe following providers were found but not available: ' + '\n'.join(unavailable_provider_messages)) else: unavailable_provider_message = '' raise ValueError( 'No available provider for type %r at %s%s' % (spec['type'], identify_object(spec), unavailable_provider_message)) # From here on, we have the invariant that possible_providers is not empty. # Only one possible provider, no need to rank further. if len(possible_providers) == 1: return possible_providers[0] def best_matches( possible_providers: Iterable[yaml_provider.Provider], adjacent_provider_options: Iterable[Iterable[yaml_provider.Provider]] ) -> list[yaml_provider.Provider]: """Given a set of possible providers, and a set of providers for each adjacent transform, returns the top possible providers as ranked by affinity to the adjacent transforms' providers. """ providers_by_score = collections.defaultdict(list) for p in possible_providers: # The sum of the affinity of the best provider # for each adjacent transform. providers_by_score[sum( max(p.affinity(ap) for ap in apo) for apo in adjacent_provider_options)].append(p) return providers_by_score[max(providers_by_score.keys())] # If there are any inputs, prefer to match them. if input_providers: possible_providers = best_matches( possible_providers, [[p] for p in input_providers]) # Without __uuid__ we can't find downstream operations. if '__uuid__' not in spec: return possible_providers[0] # Match against downstream transforms, continuing until there is no tie # or we run out of downstream transforms. if len(possible_providers) > 1: adjacent_transforms = list(self.followers(spec['__uuid__'])) while adjacent_transforms: # This is a list of all possible providers for each adjacent transform. adjacent_provider_options = [[ p for p in self.providers[self._transforms_by_uuid[t]['type']] if p.available() ] for t in adjacent_transforms] if any(not apo for apo in adjacent_provider_options): # One of the transforms had no available providers. # We will throw an error later, doesn't matter what we return. break # Filter down the set of possible providers to the best ones. possible_providers = best_matches( possible_providers, adjacent_provider_options) # If we are down to one option, no need to go further. if len(possible_providers) == 1: break # Go downstream one more step. adjacent_transforms = sum( [list(self.followers(t)) for t in adjacent_transforms], []) return possible_providers[0] # A method on scope as providers may be scoped... def create_ptransform(self, spec, input_pcolls): def maybe_with_resource_hints(transform): if 'resource_hints' in spec: return transform.with_resource_hints( **SafeLineLoader.strip_metadata(spec['resource_hints'])) else: return transform if 'type' not in spec: raise ValueError(f'Missing transform type: {identify_object(spec)}') if spec['type'] == 'composite': class _CompositeTransformStub(beam.PTransform): @staticmethod def expand(pcolls): if isinstance(pcolls, beam.PCollection): pcolls = {'input': pcolls} elif isinstance(pcolls, beam.pvalue.PBegin): pcolls = {} inner_scope = Scope( self.root, pcolls, spec['transforms'], self.providers, self.input_providers) inner_scope.compute_all() if '__implicit_outputs__' in spec['output']: return inner_scope.get_outputs( spec['output']['__implicit_outputs__']) else: return { key: inner_scope.get_pcollection(value) for (key, value) in spec['output'].items() } return maybe_with_resource_hints(_CompositeTransformStub()) if spec['type'] not in self.providers: raise ValueError( 'Unknown transform type %r at %s' % (spec['type'], identify_object(spec))) # TODO(yaml): Perhaps we can do better than a greedy choice here. # TODO(yaml): Figure out why this is needed. providers_by_input = {k: v for k, v in self.input_providers.items()} input_providers = [ providers_by_input[pcoll] for pcoll in input_pcolls if pcoll in providers_by_input ] provider = self.best_provider(spec, input_providers) extra_dependencies, spec = extract_extra_dependencies(spec) if extra_dependencies: provider = provider.with_extra_dependencies(frozenset(extra_dependencies)) config = SafeLineLoader.strip_metadata(spec.get('config', {})) if not isinstance(config, dict): raise ValueError( 'Config for transform at %s must be a mapping.' % identify_object(spec)) if (not input_pcolls and not is_explicitly_empty(spec.get('input', {})) and provider.requires_inputs(spec['type'], config)): raise ValueError( f'Missing inputs for transform at {identify_object(spec)}') try: if spec['type'].endswith('-generic'): # Centralize the validation rather than require every implementation # to do it. validate_generic_expressions( spec['type'].rsplit('-', 1)[0], config, input_pcolls) # pylint: disable=undefined-loop-variable ptransform = maybe_with_resource_hints( provider.create_transform( spec['type'], config, lambda config, input_pcolls=input_pcolls: self.create_ptransform( config, input_pcolls))) # TODO(robertwb): Should we have a better API for adding annotations # than this? annotations = { **{ 'yaml_type': spec['type'], 'yaml_args': json.dumps(config), 'yaml_provider': json.dumps(provider.to_json()) }, **ptransform.annotations() } ptransform.annotations = lambda: annotations original_expand = ptransform.expand def recording_expand(pvalue): result = original_expand(pvalue) def record_providers(pvalueish): if isinstance(pvalueish, (tuple, list)): for p in pvalueish: record_providers(p) elif isinstance(pvalueish, dict): for p in pvalueish.values(): record_providers(p) elif isinstance(pvalueish, beam.PCollection): if pvalueish not in self.input_providers: self.input_providers[pvalueish] = provider record_providers(result) return result ptransform.expand = recording_expand return ptransform except Exception as exn: if isinstance(exn, TypeError): # Create a slightly more generic error message for argument errors. msg = str(exn).replace('positional', '').replace('keyword', '') msg = re.sub(r'\S+lambda\S+', '', msg) msg = re.sub(' +', ' ', msg).strip() else: msg = str(exn) raise ValueError( f'Invalid transform specification at {identify_object(spec)}: {msg}' ) from exn def unique_name(self, spec, ptransform, strictness=0): if 'name' in spec: name = spec['name'] strictness += 1 elif ('ExternalTransform' not in ptransform.label and not ptransform.label.startswith('_')): # The label may have interesting information. name = ptransform.label else: name = spec['type'] if name in self._seen_names: if strictness >= 2: raise ValueError(f'Duplicate name at {identify_object(spec)}: {name}') else: name = f'{name}@{SafeLineLoader.get_line(spec)}' self._seen_names.add(name) return name def expand_transform(spec, scope): if 'type' not in spec: raise TypeError( f'Missing type parameter for transform at {identify_object(spec)}') type = spec['type'] if type == 'composite': return expand_composite_transform(spec, scope) else: return expand_leaf_transform(spec, scope) def expand_leaf_transform(spec, scope): spec = spec.copy() # Check for optional output_schema to verify on. # The idea is to pass this output_schema config to the ValidateWithSchema # transform. output_schema_spec = {} if 'output_schema' in spec.get('config', {}): output_schema_spec = spec.get('config').pop('output_schema') spec = normalize_inputs_outputs(spec) inputs_dict = { key: scope.get_pcollection(value) for (key, value) in empty_if_explicitly_empty(spec['input']).items() } input_type = spec.get('input_type', 'default') if input_type == 'list': inputs = tuple(inputs_dict.values()) elif input_type == 'map': inputs = inputs_dict else: if len(inputs_dict) == 0: inputs = scope.root elif len(inputs_dict) == 1: inputs = next(iter(inputs_dict.values())) else: inputs = inputs_dict _LOGGER.info("Expanding %s ", identify_object(spec)) ptransform = scope.create_ptransform(spec, inputs_dict.values()) try: # TODO: Move validation to construction? with FullyQualifiedNamedTransform.with_filter('*'): outputs = inputs | scope.unique_name(spec, ptransform) >> ptransform except Exception as exn: raise ValueError( f"Error applying transform {identify_object(spec)}: {exn}") from exn # Optional output_schema was found, so lets expand on that before returning. if output_schema_spec: error_handling_spec = {} # Obtain original transform error_handling_spec, so that all validate # schema errors use that. if 'error_handling' in spec.get('config', None): error_handling_spec = spec.get('config').get('error_handling', {}) outputs = expand_output_schema_transform( spec=output_schema_spec, outputs=outputs, error_handling_spec=error_handling_spec) if isinstance(outputs, dict): # TODO: Handle (or at least reject) nested case. return outputs elif isinstance(outputs, (tuple, list)): return {f'out{ix}': pcoll for (ix, pcoll) in enumerate(outputs)} elif isinstance(outputs, beam.PCollection): return {'out': outputs} elif outputs is None or isinstance(outputs, beam.pvalue.PDone): return {} else: raise ValueError( f'Transform {identify_object(spec)} returned an unexpected type ' f'{type(outputs)}') def expand_output_schema_transform(spec, outputs, error_handling_spec): """Applies a `Validate` transform to the output of another transform. This function is called when an `output_schema` is defined on a transform. It wraps the original transform's output(s) with a `Validate` transform to ensure the data conforms to the specified schema. If the original transform has error handling configured, validation errors will be routed to the specified error output. If not, validation failures will cause the pipeline to fail. Args: spec (dict): The `output_schema` specification from the YAML config. outputs (beam.PCollection or dict[str, beam.PCollection]): The output(s) from the transform to be validated. error_handling_spec (dict): The `error_handling` configuration from the original transform. Returns: The validated PCollection(s). If error handling is enabled, this will be a dictionary containing the 'good' output and any error outputs. Raises: ValueError: If `error_handling` is incorrectly specified within the `output_schema` spec itself, or if the main output of a multi-output transform cannot be determined. """ if 'error_handling' in spec: raise ValueError( 'error_handling config is not supported directly in ' 'the output_schema. Please use error_handling config in ' 'the transform, if possible, or use ValidateWithSchema transform ' 'instead.') # Strip metadata such as __line__ and __uuid__ as these will interfere with # the validation downstream. clean_schema = SafeLineLoader.strip_metadata(spec) # If no error handling is specified for the main transform, warn the user # that the pipeline may fail if any output data fails the output schema # validation. if not error_handling_spec: _LOGGER.warning("Output_schema config is attached to a transform that has "\ "no error_handling config specified. Any failures validating on output" \ "schema will fail the pipeline unless the user specifies an" \ "error_handling config on a capable transform. Alternatively, you can " \ "remove the output_schema config on this transform and add a " \ "ValidateWithSchema transform with separate error handling downstream of " \ "the current transform.") # The transform produced outputs with a single beam.PCollection if isinstance(outputs, beam.PCollection): outputs = _enforce_schema( outputs, 'EnforceOutputSchema', error_handling_spec, clean_schema) if isinstance(outputs, dict): main_tag = error_handling_spec.get('main_tag', 'good') main_output = outputs.pop(main_tag) if error_handling_spec: error_output_tag = error_handling_spec.get('output') if error_output_tag in outputs: return { 'output': main_output, error_output_tag: outputs.pop(error_output_tag) } return main_output # The transform produced outputs with many named PCollections and need to # determine which PCollection should be validated on. elif isinstance(outputs, dict): main_output_key = get_main_output_key(spec, outputs, error_handling_spec) validation_result = _enforce_schema( outputs[main_output_key], f'EnforceOutputSchema_{main_output_key}', error_handling_spec, clean_schema) outputs = _integrate_validation_results( outputs, validation_result, main_output_key, error_handling_spec) return outputs def get_main_output_key(spec, outputs, error_handling_spec): """Determines the main output key from a dictionary of PCollections. This is used to identify which output of a multi-output transform should be validated against an `output_schema`. The main output is determined using the following precedence: 1. An output with the key 'output'. 2. An output with the key 'good'. 3. The single output if there is only one. Args: spec: The transform specification, used for creating informative error messages. outputs: A dictionary mapping output tags to their corresponding PCollections. error_handling_spec (dict): The `error_handling` configuration from the original transform. Returns: The key of the main output PCollection. Raises: ValueError: If a main output cannot be determined because there are multiple outputs and none are named 'output' or 'good'. """ main_output_key = 'output' if main_output_key not in outputs: if 'good' in outputs: main_output_key = 'good' elif len(outputs) == 1: main_output_key = next(iter(outputs.keys())) else: raise ValueError( f"Transform {identify_object(spec)} has outputs " f"{list(outputs.keys())}, but none are named 'output' or 'good'. To " "apply an 'output_schema', please ensure the transform has exactly " "one output, or that the main output is named 'output' or 'good'.") if len(outputs) >= 3 or \ (len(outputs) == 2 and error_handling_spec.get('output') not in outputs): _LOGGER.warning( "There are currently %s outputs: %s. Only the main output will be " "validated.", len(outputs), outputs) return main_output_key def _integrate_validation_results( outputs, validation_result, main_output_key, error_handling_spec): """ Integrates the results of a validation transform back into the outputs of the original transform. This function handles merging the "good" and "bad" outputs from a `Validate` transform with the existing outputs of the transform that was validated. Args: outputs: The original dictionary of output PCollections from the transform. validation_result: The output of the `Validate` transform. This can be a single PCollection (if all elements passed) or a dictionary of PCollections (if error handling was enabled for validation). main_output_key: The key in the `outputs` dictionary corresponding to the PCollection that was validated. error_handling_spec: The error handling configuration of the original transform. Returns: The updated dictionary of output PCollections, with validation results integrated. Raises: ValueError: If the validation transform produces unexpected outputs. """ if not isinstance(validation_result, dict): outputs[main_output_key] = validation_result return outputs # The main output from validation is the good output. main_tag = error_handling_spec.get('main_tag', 'good') outputs[main_output_key] = validation_result.pop(main_tag) if error_handling_spec: error_output_tag = error_handling_spec['output'] if error_output_tag in validation_result: schema_error_pcoll = validation_result.pop(error_output_tag) # The original transform also had an error output. Merge them. outputs[error_output_tag] = ( (outputs[error_output_tag], schema_error_pcoll) | f'FlattenErrors_{main_output_key}' >> beam.Flatten()) # There should be no other outputs from validation. if validation_result: raise ValueError( "Unexpected outputs from validation: " f"{list(validation_result.keys())}") return outputs def _enforce_schema(pcoll, label, error_handling_spec, clean_schema): """Applies schema to PCollection elements if necessary, then validates. This function ensures that the input PCollection conforms to a specified schema. If the PCollection is schemaless (i.e., its element_type is Any), it attempts to convert its elements into schema-aware `beam.Row` objects based on the provided `clean_schema`. After ensuring the PCollection has a defined schema, it applies a `Validate` transform to perform the actual schema validation. Args: pcoll: The input PCollection to be schema-enforced and validated. label: A string label to be used for the Beam transforms created within this function. error_handling_spec: A dictionary specifying how to handle validation errors. clean_schema: A dictionary representing the schema to enforce and validate against. Returns: A PCollection (or PCollectionTuple if error handling is enabled) resulting from the `Validate` transform. """ if pcoll.element_type == typehints.Any: _LOGGER.info( "PCollection for %s has no schema (element_type=Any). " "Converting elements to beam.Row based on provided output_schema.", label) try: # Attempt to confer the schemaless elements into schema-aware beam.Row # objects beam_schema = json_utils.json_schema_to_beam_schema(clean_schema) row_type_constraint = schemas.named_tuple_from_schema(beam_schema) def to_row(element): """ Convert a single element into the row type constraint type. """ if isinstance(element, dict): return row_type_constraint(**element) elif hasattr(element, '_asdict'): # Handle NamedTuple, beam.Row return row_type_constraint(**element._asdict()) else: raise TypeError( f"Cannot convert element of type {type(element)} to beam.Row " f"for validation in {label}. Element: {element}") pcoll = pcoll | f'{label}_ConvertToRow' >> beam.Map( to_row).with_output_types(row_type_constraint) except Exception as e: raise ValueError( f"Failed to prepare schemaless PCollection for \ validation in {label}: {e}") from e # Add Validation step downstream of current transform return pcoll | label >> Validate( schema=clean_schema, error_handling=error_handling_spec) def expand_composite_transform(spec, scope): spec = normalize_inputs_outputs(normalize_source_sink(spec)) original_transforms = spec['transforms'] # Check if any transform has a NON-EMPTY explicit input or output. # Note: {} (empty dict) means "no explicit input specified" and should # NOT count as having explicit io. # However, if the composite has no input, we can't do implicit chaining. has_explicit_io = any( io is not None and not is_empty(t.get(io, {})) for t in original_transforms for io in ('input', 'output')) # If the composite has no input, we can't do implicit chaining composite_has_input = not is_empty(spec.get('input', {})) # Only do implicit chaining if: # 1. No transform has explicit io, AND # 2. The composite has an input to chain from if not has_explicit_io and composite_has_input: new_transforms = [] for ix, transform in enumerate(original_transforms): transform = dict(transform) if ix == 0: composite_input = spec.get('input', {}) if is_explicitly_empty(composite_input): transform['input'] = composite_input elif is_empty(composite_input): # No explicit input - the composite input IS the pipeline input. # Reference the 'input' key from the Scope's inputs. transform['input'] = 'input' else: transform['input'] = {key: key for key in composite_input.keys()} else: transform['input'] = new_transforms[-1]['__uuid__'] new_transforms.append(transform) if new_transforms: spec = dict(spec, transforms=new_transforms) # Check if output is empty, not just present (normalization sets it to {}) if is_empty(spec.get('output', {})): spec['output'] = { '__implicit_outputs__': new_transforms[-1]['__uuid__'] } # Compute the inputs for the inner scope. # If the composite has an empty input dict ({}), it means the composite # should use the parent scope's inputs directly. composite_input = spec.get('input', {}) if is_empty(composite_input): # No explicit input - use the parent scope's inputs directly inner_scope_inputs = dict(scope._inputs) else: # The composite has explicit input references # They can reference either: # 1. A parent scope input (e.g., 'input' key in scope._inputs) # 2. A transform output (e.g., 'uuid' -> the output of a transform) inner_scope_inputs = {} for key, value in composite_input.items(): if isinstance(value, str) and value in scope._inputs: # Reference to a parent scope input inner_scope_inputs[key] = scope._inputs[value] else: # Reference to a transform output inner_scope_inputs[key] = scope.get_pcollection(value) inner_scope = Scope( scope.root, inner_scope_inputs, spec['transforms'], # TODO(robertwb): Are scoped providers ever used? Worth supporting? yaml_provider.merge_providers( yaml_provider.parse_providers('', spec.get('providers', [])), scope.providers), scope.input_providers) class CompositePTransform(beam.PTransform): @staticmethod def expand(inputs): inner_scope.compute_all() if '__implicit_outputs__' in spec['output']: result = inner_scope.get_outputs(spec['output']['__implicit_outputs__']) return result else: return { key: inner_scope.get_pcollection(value) for (key, value) in spec['output'].items() } transform = CompositePTransform() if 'resource_hints' in spec: transform = transform.with_resource_hints( **SafeLineLoader.strip_metadata(spec['resource_hints'])) # Always set a name for the composite to ensure proper return value if 'name' not in spec: spec['name'] = 'Composite' if spec['name'] is None: # top-level pipeline, don't nest return transform.expand(None) else: _LOGGER.info("Expanding %s ", identify_object(spec)) # When the input references a scope input (not a transform output), # we need to use the scope's inputs directly input_dict = {} for key, value in empty_if_explicitly_empty(spec['input']).items(): if isinstance(value, str) and value in scope._inputs: # Reference to a scope input input_dict[key] = scope._inputs[value] else: # Reference to a transform output input_dict[key] = scope.get_pcollection(value) return (input_dict or scope.root) | scope.unique_name(spec, None) >> transform def expand_chain_transform(spec, scope): return expand_composite_transform(chain_as_composite(spec), scope) def chain_as_composite(spec): def is_not_output_of_last_transform(new_transforms, value): return ( ('name' in new_transforms[-1] and value != new_transforms[-1]['name']) or ('type' in new_transforms[-1] and value != new_transforms[-1]['type'])) # A chain is simply a composite transform where all inputs and outputs # are implicit. spec = normalize_source_sink(spec) if 'transforms' not in spec: raise TypeError( f"Chain at {identify_object(spec)} missing transforms property.") has_explicit_outputs = 'output' in spec composite_spec = dict(normalize_inputs_outputs(tag_explicit_inputs(spec))) new_transforms = [] for ix, transform in enumerate(composite_spec['transforms']): transform = dict(transform) if any(io in transform for io in ('input', 'output')): if (ix == 0 and 'input' in transform and 'output' not in transform and is_explicitly_empty(transform['input'])): # This is OK as source clause sets an explicitly empty input. pass else: raise ValueError( f'Transform {identify_object(transform)} is part of a chain. ' 'Cannot define explicit inputs on chain pipeline') if ix == 0: if is_explicitly_empty(transform.get('input', None)): pass elif is_explicitly_empty(composite_spec['input']): transform['input'] = composite_spec['input'] elif is_empty(composite_spec['input']): del composite_spec['input'] else: transform['input'] = { key: key for key in composite_spec['input'].keys() } else: transform['input'] = new_transforms[-1]['__uuid__'] new_transforms.append(transform) new_transforms.extend(spec.get('extra_transforms', [])) composite_spec['transforms'] = new_transforms last_transform = new_transforms[-1]['__uuid__'] if has_explicit_outputs: for (key, value) in composite_spec['output'].items(): if is_not_output_of_last_transform(new_transforms, value): raise ValueError( f"Explicit output {identify_object(value)} of the chain transform" f" is not an output of the last transform.") composite_spec['output'] = { key: f'{last_transform}.{value}' for (key, value) in composite_spec['output'].items() } else: composite_spec['output'] = {'__implicit_outputs__': last_transform} if 'name' not in composite_spec: composite_spec['name'] = 'Chain' composite_spec['type'] = 'composite' return composite_spec def preprocess_chain(spec): if spec['type'] == 'chain': return chain_as_composite(spec) else: return spec def pipeline_as_composite(spec): if isinstance(spec, list): return { 'type': 'composite', 'name': None, 'transforms': spec, '__line__': spec[0]['__line__'], '__uuid__': SafeLineLoader.create_uuid(), } else: return dict(spec, name=None, type=spec.get('type', 'composite')) def normalize_source_sink(spec): if 'source' not in spec and 'sink' not in spec: return spec spec = dict(spec) spec['transforms'] = list(spec.get('transforms', [])) if 'source' in spec: if 'input' not in spec['source']: spec['source']['input'] = explicitly_empty() spec['transforms'].insert(0, spec.pop('source')) if 'sink' in spec: spec['transforms'].append(spec.pop('sink')) return spec def preprocess_source_sink(spec): if spec['type'] in ('chain', 'composite'): return normalize_source_sink(spec) else: return spec def tag_explicit_inputs(spec): if 'input' in spec and not SafeLineLoader.strip_metadata(spec['input']): return dict(spec, input=explicitly_empty()) else: return spec def normalize_inputs_outputs(spec): spec = dict(spec) def normalize_io(tag): io = spec.get(tag, {}) if isinstance(io, (str, list)): return {tag: io} else: return SafeLineLoader.strip_metadata(io, tagged_str=False) return dict(spec, input=normalize_io('input'), output=normalize_io('output')) def identify_object(spec): line = SafeLineLoader.get_line(spec) name = extract_name(spec) if name: return f'"{name}" at line {line}' else: return f'at line {line}' def extract_name(spec): if isinstance(spec, dict): if 'name' in spec: return spec['name'] elif 'id' in spec: return spec['id'] elif 'type' in spec: return spec['type'] elif len(spec) == 1: return extract_name(next(iter(spec.values()))) else: return '' elif isinstance(spec, str): return spec else: return '' def extract_extra_dependencies(spec): deps = spec.get('config', {}).get('dependencies', []) if not deps: return [], spec if not isinstance(deps, list): raise TypeError(f'Dependencies must be a list of strings, got {deps}') return deps, dict( spec, config={k: v for k, v in spec['config'].items() if k != 'dependencies'}) def push_windowing_to_roots(spec): scope = LightweightScope(spec['transforms']) consumed_outputs_by_transform = collections.defaultdict(set) for transform in spec['transforms']: for _, input_ref in empty_if_explicitly_empty(transform['input']).items(): try: transform_id, output = scope.get_transform_id_and_output_name(input_ref) consumed_outputs_by_transform[transform_id].add(output) except ValueError: # Could be an input or an ambiguity we'll raise later. pass for transform in spec['transforms']: if is_empty(transform['input']) and 'windowing' not in transform: transform['windowing'] = spec['windowing'] transform['__consumed_outputs'] = consumed_outputs_by_transform[ transform['__uuid__']] return spec def preprocess_windowing(spec): if spec['type'] == 'WindowInto': # This is the transform where it is actually applied. if 'windowing' in spec: spec['config'] = spec.get('config', {}) spec['config']['windowing'] = spec.pop('windowing') if spec.get('config', {}).get('windowing'): windowing_config = spec['config']['windowing'] if isinstance(windowing_config, str): try: # PyYAML can load a JSON string - one-line and multi-line. # Without this code, multi-line is not supported. parsed_config = yaml.safe_load(windowing_config) if not isinstance(parsed_config, dict): raise TypeError('Windowing config string must be a YAML/JSON map.') spec['config']['windowing'] = parsed_config except Exception as e: raise ValueError( f'Error parsing windowing config string at \ {identify_object(spec)}: {e}') from e return spec elif 'windowing' not in spec: # Nothing to do. return spec if spec['type'] == 'composite': # Apply the windowing to any reads, creates, etc. in this transform # TODO(robertwb): Better handle the case where a read is followed by a # setting of the timestamps. We should be careful of sliding windows # in particular. spec = push_windowing_to_roots(spec) windowing = spec.pop('windowing') if not is_empty(spec['input']): # Apply the windowing to all inputs by wrapping it in a transform that # first applies windowing and then applies the original transform. original_inputs = spec['input'] windowing_transforms = [{ 'type': 'WindowInto', 'name': f'WindowInto[{key}]', 'windowing': windowing, 'input': { 'input': key }, '__line__': spec['__line__'], '__uuid__': SafeLineLoader.create_uuid(), } for key in original_inputs.keys()] windowed_inputs = { key: t['__uuid__'] for (key, t) in zip(original_inputs.keys(), windowing_transforms) } modified_spec = dict( spec, input=windowed_inputs, __uuid__=SafeLineLoader.create_uuid()) return { 'type': 'composite', 'name': spec.get('name', None) or spec['type'], 'transforms': [modified_spec] + windowing_transforms, 'input': spec['input'], 'output': modified_spec['__uuid__'], 'config': { 'error_handling': spec.get('config', {}).get('error_handling', {}) }, '__line__': spec['__line__'], '__uuid__': spec['__uuid__'], } elif spec['type'] == 'composite': # Pushing the windowing down was sufficient. return spec else: # No inputs, apply the windowing to all outputs. consumed_outputs = list(spec.pop('__consumed_outputs', {None})) modified_spec = dict(spec, __uuid__=SafeLineLoader.create_uuid()) windowing_transforms = [{ 'type': 'WindowInto', 'name': f'WindowInto[{out}]', 'windowing': windowing, 'input': { 'input': modified_spec['__uuid__'] + ('.' + out if out else '') }, '__line__': spec['__line__'], '__uuid__': SafeLineLoader.create_uuid(), } for out in consumed_outputs] if consumed_outputs == [None]: windowed_outputs = only_element(windowing_transforms)['__uuid__'] else: windowed_outputs = { out: t['__uuid__'] for (out, t) in zip(consumed_outputs, windowing_transforms) } return { 'type': 'composite', 'name': spec.get('name', None) or spec['type'], 'transforms': [modified_spec] + windowing_transforms, 'output': windowed_outputs, 'config': { 'error_handling': spec.get('config', {}).get('error_handling', {}) }, '__line__': spec['__line__'], '__uuid__': spec['__uuid__'], } def preprocess_flattened_inputs(spec): if spec['type'] != 'composite': return spec # Prefer to add the flattens as sibling operations rather than nesting # to keep graph shape consistent when the number of inputs goes from # one to multiple. new_transforms = [] for t in spec['transforms']: if t['type'] == 'Flatten': # Don't flatten before explicit flatten. # But we do have to expand list inputs into singleton inputs. def all_inputs(t): for key, values in t.get('input', {}).items(): if isinstance(values, list): for ix, value in enumerate(values): yield f'{key}{ix}', value else: yield key, values inputs_dict = {} for key, value in all_inputs(t): while key in inputs_dict: key += '_' inputs_dict[key] = value t = dict(t, input=inputs_dict) else: replaced_inputs = {} for key, values in t.get('input', {}).items(): if isinstance(values, list): flatten_id = SafeLineLoader.create_uuid() new_transforms.append({ 'type': 'Flatten', 'name': '%s-Flatten[%s]' % (t.get('name', t['type']), key), 'input': { f'input{ix}': value for (ix, value) in enumerate(values) }, '__line__': spec['__line__'], '__uuid__': flatten_id, }) replaced_inputs[key] = flatten_id if replaced_inputs: t = dict(t, input={**t['input'], **replaced_inputs}) new_transforms.append(t) return dict(spec, transforms=new_transforms) def ensure_transforms_have_types(spec): if 'type' not in spec: raise ValueError(f'Missing type specification in {identify_object(spec)}') return spec def ensure_errors_consumed(spec): if spec['type'] == 'composite': scope = LightweightScope(spec['transforms']) to_handle = {} consumed = set( scope.get_transform_id_and_output_name(output) for output in spec['output'].values()) for t in spec['transforms']: config = t.get('config', t) if 'error_handling' in config: if 'output' not in config['error_handling']: raise ValueError( f'Missing output in error_handling of {identify_object(t)}') to_handle[t['__uuid__'], config['error_handling']['output']] = t for _, input in empty_if_explicitly_empty(t['input']).items(): if input not in spec['input']: consumed.add(scope.get_transform_id_and_output_name(input)) for error_pcoll, t in to_handle.items(): if error_pcoll not in consumed: config = t.get('config', t) transform_name = t.get('name', t.get('type')) error_output_name = config['error_handling']['output'] raise ValueError( f'Unconsumed error output for {identify_object(t)}. ' f'The output named {transform_name}.{error_output_name} ' 'must be used as an input to some other transform. ' 'See https://beam.apache.org/documentation/sdks/yaml-errors') return spec def lift_config(spec): if 'config' not in spec: common_params = 'name', 'type', 'input', 'output', 'transforms' return { 'config': { k: v for (k, v) in spec.items() if k not in common_params }, **{ k: v for (k, v) in spec.items() # if k in common_params or k in ('__line__', '__uuid__') } } else: return spec def ensure_config(spec): if 'config' not in spec: spec['config'] = {} return spec def apply_phase(phase, spec): spec = phase(spec) if spec['type'] in {'composite', 'chain'} and 'transforms' in spec: spec = dict( spec, transforms=[apply_phase(phase, t) for t in spec['transforms']]) return spec def preprocess(spec, verbose=False, known_transforms=None): if verbose: pprint.pprint(spec) if known_transforms: known_transforms = set(known_transforms).union(['chain', 'composite']) def ensure_transforms_have_providers(spec): if known_transforms: if spec['type'] not in known_transforms: raise ValueError( 'Unknown type or missing provider ' f'for type {spec["type"]} for {identify_object(spec)}') return spec def preprocess_languages(spec): if spec['type'] in ('AssignTimestamps', 'Combine', 'Filter', 'MapToFields', 'Partition'): language = spec.get('config', {}).get('language', 'generic') new_type = spec['type'] + '-' + language if known_transforms and new_type not in known_transforms: if language == 'generic': raise ValueError(f'Missing language for {identify_object(spec)}') else: raise ValueError( f'Unknown language {language} for {identify_object(spec)}') return dict(spec, type=new_type, name=spec.get('name', spec['type'])) else: return spec def validate_transform_references(spec): name = spec.get('name', '') transform_type = spec.get('type') inputs = spec.get('input').get('input', []) if not is_empty(inputs): input_values = [inputs] if isinstance(inputs, str) else inputs for input_value in input_values: if input_value in (name, transform_type): raise ValueError( f"Circular reference detected: Transform {name} " f"references itself as input in {identify_object(spec)}") return spec for phase in [ ensure_transforms_have_types, normalize_mapping, normalize_combine, preprocess_languages, ensure_transforms_have_providers, preprocess_source_sink, preprocess_chain, tag_explicit_inputs, normalize_inputs_outputs, validate_transform_references, preprocess_flattened_inputs, ensure_errors_consumed, preprocess_windowing, # TODO(robertwb): Consider enabling this by default, or as an option. # lift_config, ensure_config, ]: spec = apply_phase(phase, spec) if verbose: print('=' * 20, phase, '=' * 20) pprint.pprint(spec) return spec class _BeamFileIOLoader(jinja2.BaseLoader): def get_source(self, environment, path): with FileSystems.open(path) as fin: source = fin.read().decode() return source, path, lambda: True def expand_jinja( jinja_template: str, jinja_variables: Mapping[str, Any]) -> str: return ( # keep formatting jinja2.Environment( undefined=jinja2.StrictUndefined, loader=_BeamFileIOLoader()) .from_string(jinja_template) .render(datetime=datetime, **jinja_variables)) class YamlTransform(beam.PTransform): def __init__(self, spec, providers={}): # pylint: disable=dangerous-default-value if isinstance(spec, str): spec = yaml.load(spec, Loader=SafeLineLoader) if isinstance(providers, dict): providers = { key: yaml_provider.as_provider_list(key, value) for (key, value) in providers.items() } # TODO(BEAM-26941): Validate as a transform. self._providers = yaml_provider.merge_providers( providers, yaml_provider.standard_providers()) self._spec = preprocess(spec, known_transforms=self._providers.keys()) self._was_chain = spec['type'] == 'chain' def expand(self, pcolls): if isinstance(pcolls, beam.pvalue.PBegin): root = pcolls pipeline = root.pipeline pcolls = {} elif isinstance(pcolls, beam.PCollection): root = pcolls.pipeline pipeline = root pcolls = {'input': pcolls} if not self._spec['input']: self._spec['input'] = {'input': 'input'} if self._was_chain and self._spec['transforms']: # This should have been copied as part of the composite-to-chain. self._spec['transforms'][0]['input'] = self._spec['input'] else: root = next(iter(pcolls.values())).pipeline pipeline = root if not self._spec['input']: self._spec['input'] = {name: name for name in pcolls.keys()} python_provider = yaml_provider.InlineProvider({}) # Label goog-dataflow-yaml if job is started using Beam YAML. options = pipeline.options.view_as(GoogleCloudOptions) yaml_version = ('beam-yaml=' + beam.version.__version__.replace('.', '_')) if not options.labels: options.labels = [] if yaml_version not in options.labels: options.labels.append(yaml_version) result = expand_transform( self._spec, Scope( root, pcolls, transforms=[self._spec], providers=self._providers, input_providers={ pcoll: python_provider for pcoll in pcolls.values() })) if len(result) == 1: return only_element(result.values()) else: return result def expand_pipeline( pipeline, pipeline_spec, providers=None, validate_schema='generic' if jsonschema is not None else None, pipeline_path=''): if isinstance(pipeline, beam.pvalue.PBegin): root = pipeline else: root = beam.pvalue.PBegin(pipeline) if isinstance(pipeline_spec, str): pipeline_spec = yaml.load(pipeline_spec, Loader=SafeLineLoader) # TODO(robertwb): It's unclear whether this gives as good of errors, but # this could certainly be handy as a first pass when Beam is not available. if validate_schema and validate_schema != 'none': validate_against_schema(pipeline_spec, validate_schema) # Calling expand directly to avoid outer layer of nesting. return YamlTransform( pipeline_as_composite(pipeline_spec['pipeline']), yaml_provider.merge_providers( yaml_provider.parse_providers( pipeline_path, pipeline_spec.get('providers', [])), providers or {})).expand(root) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_transform_scope_test.py0000644000175100017510000002442415206113033024465 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import collections import logging import unittest import yaml import apache_beam as beam from apache_beam.yaml import yaml_provider from apache_beam.yaml import yaml_transform from apache_beam.yaml.yaml_transform import LightweightScope from apache_beam.yaml.yaml_transform import SafeLineLoader from apache_beam.yaml.yaml_transform import Scope class ScopeTest(unittest.TestCase): def get_scope_by_spec(self, p, spec): spec = yaml.load(spec, Loader=SafeLineLoader) scope = Scope( beam.pvalue.PBegin(p), {}, spec['transforms'], yaml_provider.standard_providers(), {}) return scope, spec def test_get_pcollection_input(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(range(3)) scope = Scope( p, {'input': elements}, transforms=[], providers=yaml_provider.standard_providers(), input_providers={}) result = scope.get_pcollection('input') self.assertEqual("PCollection[Create/Map(decode).None]", str(result)) def test_get_pcollection_output(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: spec = ''' transforms: - type: Create config: elements: [0, 1, 3, 4] - type: LogForTesting name: Square input: Create ''' scope, spec = self.get_scope_by_spec(p, spec) self.assertEqual( "PCollection[Create/Map(decode).None]", str(scope.get_pcollection("Create"))) self.assertEqual( "PCollection[Square/LogForTesting.None]", str(scope.get_pcollection("Square"))) self.assertEqual( "PCollection[Square/LogForTesting.None]", str(scope.get_pcollection("LogForTesting"))) self.assertTrue( scope.get_pcollection("Square") == scope.get_pcollection( "LogForTesting")) def test_create_ptransform(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: spec = ''' transforms: - type: Create config: elements: [1, 2, 3] ''' scope, spec = self.get_scope_by_spec(p, spec) result = scope.create_ptransform(spec['transforms'][0], []) self.assertIsInstance(result, beam.transforms.Create) self.assertEqual(result.label, 'Create') result_annotations = {**result.annotations()} target_annotations = { 'yaml_type': 'Create', 'yaml_args': '{"elements": [1, 2, 3]}', 'yaml_provider': '{"type": "InlineProvider"}' } # Check if target_annotations is a subset of result_annotations self.assertDictEqual( result_annotations, { **result_annotations, **target_annotations }) class TestProvider(yaml_provider.InlineProvider): def __init__(self, transform, name): super().__init__({ name: lambda: beam.Map(lambda x: (x or ()) + (name, )), # or None transform: lambda: beam.Map(lambda x: (x or ()) + (name, )), }) self._transform = transform self._name = name def __repr__(self): return 'TestProvider(%r, %r)' % (self._transform, self._name) def _affinity(self, other): if isinstance(other, TestProvider): # Providers are closer based on how much their names match prefixes. affinity = 1 for x, y in zip(self._name, other._name): if x != y: break affinity *= 10 return affinity else: return -1000 class ProviderAffinityTest(unittest.TestCase): @staticmethod def create_scope(s, providers): providers_dict = collections.defaultdict(list) for provider in providers: for transform_type in provider.provided_transforms(): providers_dict[transform_type].append(provider) spec = yaml_transform.preprocess(yaml.load(s, Loader=SafeLineLoader)) return Scope( None, {}, transforms=spec['transforms'], providers=providers_dict, input_providers={}) def test_best_provider_based_on_input(self): provider_Ax = TestProvider('A', 'xxx') provider_Ay = TestProvider('A', 'yyy') provider_Bx = TestProvider('B', 'xxz') provider_By = TestProvider('B', 'yyz') scope = self.create_scope( ''' type: chain transforms: - type: A - type: B ''', [provider_Ax, provider_Ay, provider_Bx, provider_By]) self.assertEqual(scope.best_provider('B', [provider_Ax]), provider_Bx) self.assertEqual(scope.best_provider('B', [provider_Ay]), provider_By) def test_best_provider_based_on_followers(self): close_provider = TestProvider('A', 'xxy') far_provider = TestProvider('A', 'yyy') following_provider = TestProvider('B', 'xxx') scope = self.create_scope( ''' type: chain transforms: - type: A - type: B ''', [far_provider, close_provider, following_provider]) self.assertEqual(scope.best_provider('A', []), close_provider) def test_best_provider_based_on_multiple_followers(self): close_provider = TestProvider('A', 'xxy') provider_B = TestProvider('B', 'xxx') # These are not quite as close as the two above. far_provider = TestProvider('A', 'yyy') provider_C = TestProvider('C', 'yzz') scope = self.create_scope( ''' type: composite transforms: - type: A - type: B input: A - type: C input: A ''', [far_provider, close_provider, provider_B, provider_C]) self.assertEqual(scope.best_provider('A', []), close_provider) def test_best_provider_based_on_distant_follower(self): providers = [ # xxx and yyy vend both TestProvider('A', 'xxx'), TestProvider('A', 'yyy'), TestProvider('B', 'xxx'), TestProvider('B', 'yyy'), TestProvider('C', 'xxx'), TestProvider('C', 'yyy'), # D and E are only provided by a single provider each. TestProvider('D', 'xxx'), TestProvider('E', 'yyy') ] # If D is the eventual destination, pick the xxx one. scope = self.create_scope( ''' type: chain transforms: - type: A - type: B - type: C - type: D ''', providers) self.assertEqual(scope.best_provider('A', []), providers[0]) # If instead E is the eventual destination, pick the yyy one. scope = self.create_scope( ''' type: chain transforms: - type: A - type: B - type: C - type: E ''', providers) self.assertEqual(scope.best_provider('A', []), providers[1]) # If we have D then E, stay with xxx as long as possible to only switch once scope = self.create_scope( ''' type: chain transforms: - type: A - type: B - type: C - type: D - type: E ''', providers) self.assertEqual(scope.best_provider('A', []), providers[0]) class LightweightScopeTest(unittest.TestCase): @staticmethod def get_spec(): pipeline_yaml = ''' - type: PyMap name: Square input: elements fn: "lambda x: x * x" - type: PyMap name: PyMap input: elements fn: "lambda x: x * x * x" - type: Filter name: FilterOutBigNumbers input: PyMap keep: "lambda x: x<100" ''' return yaml.load(pipeline_yaml, Loader=SafeLineLoader) def test_init(self): spec = self.get_spec() scope = LightweightScope(spec) self.assertEqual(len(scope._transforms_by_uuid), 3) self.assertCountEqual( list(scope._uuid_by_name.keys()), ["PyMap", "Square", "Filter", "FilterOutBigNumbers"]) def test_get_transform_id_and_output_name(self): spec = self.get_spec() scope = LightweightScope(spec) transform_id, output = scope.get_transform_id_and_output_name("Square") self.assertEqual(transform_id, spec[0]['__uuid__']) self.assertEqual(output, None) def test_get_transform_id_and_output_name_with_dot(self): spec = self.get_spec() scope = LightweightScope(spec) transform_id, output = \ scope.get_transform_id_and_output_name("Square.OutputName") self.assertEqual(transform_id, spec[0]['__uuid__']) self.assertEqual(output, "OutputName") def test_get_transform_id_by_uuid(self): spec = self.get_spec() scope = LightweightScope(spec) transform_id = scope.get_transform_id(spec[0]['__uuid__']) self.assertEqual(spec[0]['__uuid__'], transform_id) def test_get_transform_id_by_unique_name(self): spec = self.get_spec() scope = LightweightScope(spec) transform_id = scope.get_transform_id("Square") self.assertEqual(transform_id, spec[0]['__uuid__']) def test_get_transform_id_by_ambiguous_name(self): spec = self.get_spec() scope = LightweightScope(spec) with self.assertRaisesRegex(ValueError, r'Ambiguous.*PyMap'): scope.get_transform_id(scope.get_transform_id(spec[1]['name'])) def test_get_transform_id_by_unknown_name(self): spec = self.get_spec() scope = LightweightScope(spec) with self.assertRaisesRegex(ValueError, r'Unknown.*NotExistingTransform'): scope.get_transform_id("NotExistingTransform") if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_transform_test.py0000644000175100017510000014170615206113033023277 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import collections import glob import logging import os import shutil import tempfile import unittest import apache_beam as beam from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.utils import python_callable from apache_beam.yaml import yaml_provider from apache_beam.yaml.yaml_transform import YamlTransform try: import jsonschema except ImportError: jsonschema = None _LOGGER = logging.getLogger(__name__) class CreateTimestamped(beam.PTransform): _yaml_requires_inputs = False def __init__(self, elements): self._elements = elements def expand(self, p): return ( p | beam.Create(self._elements) | beam.Map(lambda x: beam.transforms.window.TimestampedValue(x, x))) class CreateInts(beam.PTransform): _yaml_requires_inputs = False def __init__(self, elements): self._elements = elements def expand(self, p): return p | beam.Create(self._elements) class SumGlobally(beam.PTransform): def expand(self, pcoll): return pcoll | beam.CombineGlobally(sum).without_defaults() class SizeLimiter(beam.PTransform): def __init__(self, limit, error_handling): self._limit = limit self._error_handling = error_handling def expand(self, pcoll): def raise_on_big(row): if len(row.element) > self._limit: raise ValueError(row.element) else: return row.element good, bad = pcoll | beam.Map(raise_on_big).with_exception_handling() return {'small_elements': good, self._error_handling['output']: bad} TEST_PROVIDERS = { 'CreateInts': CreateInts, 'CreateTimestamped': CreateTimestamped, 'SumGlobally': SumGlobally, 'SizeLimiter': SizeLimiter, 'PyMap': lambda fn: beam.Map(python_callable.PythonCallableWithSource(fn)), } @unittest.skipIf(jsonschema is None, "Yaml dependencies not installed") class YamlTransformE2ETest(unittest.TestCase): def test_composite(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([1, 2, 3]) # TODO(robertwb): Consider making the input implicit (and below). result = elements | YamlTransform( ''' type: composite input: elements: input transforms: - type: PyMap name: Square input: elements config: fn: "lambda x: x * x" - type: PyMap name: Cube input: elements config: fn: "lambda x: x * x * x" - type: Flatten input: [Square, Cube] output: Flatten ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([1, 4, 9, 1, 8, 27])) def test_composite_implicit_input_chaining(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create([1, 2, 3]) result = elements | YamlTransform( ''' type: composite transforms: - type: PyMap name: Square config: fn: "lambda x: x * x" - type: PyMap name: Increment config: fn: "lambda x: x + 1" ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([2, 5, 10])) def test_chain_with_input(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(range(10)) result = elements | YamlTransform( ''' type: chain input: elements: input transforms: - type: PyMap config: fn: "lambda x: x * x + x" - type: PyMap config: fn: "lambda x: x + 41" ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([41, 43, 47, 53, 61, 71, 83, 97, 113, 131])) def test_chain_with_source_sink(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: chain source: type: CreateInts config: elements: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] transforms: - type: PyMap config: fn: "lambda x: x * x + x" sink: type: PyMap config: fn: "lambda x: x + 41" ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([41, 43, 47, 53, 61, 71, 83, 97, 113, 131])) def test_chain_with_root(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: chain transforms: - type: CreateInts config: elements: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] - type: PyMap config: fn: "lambda x: x * x + x" - type: PyMap config: fn: "lambda x: x + 41" ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([41, 43, 47, 53, 61, 71, 83, 97, 113, 131])) def create_has_schema(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: chain transforms: - type: Create config: elements: [{a: 1, b: 'x'}, {a: 2, b: 'y'}] - type: MapToFields config: language: python fields: repeated: a * b ''') | beam.Map(lambda x: x.repeated) assert_that(result, equal_to(['x', 'yy'])) def test_implicit_flatten(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: composite transforms: - type: Create name: CreateSmall config: elements: [1, 2, 3] - type: Create name: CreateBig config: elements: [100, 200] - type: PyMap input: [CreateBig, CreateSmall] config: fn: "lambda x: x.element * x.element" output: PyMap ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([1, 4, 9, 10000, 40000])) def test_csv_to_json(self): try: import pandas as pd except ImportError: raise unittest.SkipTest('Pandas not available.') with tempfile.TemporaryDirectory() as tmpdir: data = pd.DataFrame([{'label': f'{i}a', 'rank': i} for i in range(1024)]) input = os.path.join(tmpdir, 'input.csv') output = os.path.join(tmpdir, 'output.json') data.to_csv(input, index=False) with open(input, 'r') as f: lines = f.readlines() _LOGGER.debug("input.csv has these {lines} lines.") self.assertEqual(len(lines), len(data) + 1) # +1 for header with beam.Pipeline() as p: result = p | YamlTransform( ''' type: chain transforms: - type: ReadFromCsv config: path: %s - type: WriteToJson config: path: %s num_shards: 1 - type: LogForTesting ''' % (repr(input), repr(output))) all_output = list(glob.glob(output + "-*")) file_and_size = {f: os.path.getsize(f) for f in all_output} self.assertEqual( len(all_output), 1, msg=f"Expected 1 shard file, but found {len(all_output)}. " f"Files & sizes (bytes): {file_and_size}") output_shard = all_output[0] result = pd.read_json( output_shard, orient='records', lines=True).sort_values('rank').reindex() pd.testing.assert_frame_equal(data, result) def test_circular_reference_validation(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: # pylint: disable=expression-not-assigned with self.assertRaisesRegex(ValueError, r'Circular reference detected.*'): p | YamlTransform( ''' type: composite transforms: - type: Create name: Create config: elements: [0, 1, 3, 4] input: Create - type: PyMap name: PyMap config: fn: "lambda row: row.element * row.element" input: Create output: PyMap ''', providers=TEST_PROVIDERS) def test_circular_reference_multi_inputs_validation(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: # pylint: disable=expression-not-assigned with self.assertRaisesRegex(ValueError, r'Circular reference detected.*'): p | YamlTransform( ''' type: composite transforms: - type: Create name: Create config: elements: [0, 1, 3, 4] - type: PyMap name: PyMap config: fn: "lambda row: row.element * row.element" input: [Create, PyMap] output: PyMap ''', providers=TEST_PROVIDERS) def test_name_is_not_ambiguous(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: composite transforms: - type: Create name: Create config: elements: [0, 1, 3, 4] - type: PyMap name: PyMap config: fn: "lambda row: row.element * row.element" input: Create output: PyMap ''', providers=TEST_PROVIDERS) # No exception raised assert_that(result, equal_to([0, 1, 9, 16])) def test_name_is_ambiguous(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: # pylint: disable=expression-not-assigned with self.assertRaisesRegex(ValueError, r'Circular reference detected.*'): p | YamlTransform( ''' type: composite transforms: - type: Create name: CreateData config: elements: [0, 1, 3, 4] - type: PyMap name: PyMap config: fn: "lambda elem: elem + 2" input: CreateData - type: PyMap name: AnotherMap config: fn: "lambda elem: elem + 3" input: PyMap output: AnotherMap ''', providers=TEST_PROVIDERS) def test_empty_inputs_throws_error(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: with self.assertRaisesRegex(ValueError, 'Missing inputs for transform at ' '"EmptyInputOkButYamlDoesntKnow" at line .*'): _ = p | YamlTransform( ''' type: composite transforms: - type: PyTransform name: EmptyInputOkButYamlDoesntKnow config: constructor: apache_beam.Impulse ''') def test_empty_inputs_ok_in_source(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: # Does not throw an error like it does above. _ = p | YamlTransform( ''' type: composite source: type: PyTransform name: EmptyInputOkButYamlDoesntKnow config: constructor: apache_beam.Impulse ''') def test_empty_inputs_ok_if_explicit(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: # Does not throw an error like it does above. _ = p | YamlTransform( ''' type: composite transforms: - type: PyTransform name: EmptyInputOkButYamlDoesntKnow input: {} config: constructor: apache_beam.Impulse ''') def test_annotations(self): t = LinearTransform(5, b=100) annotations = t.annotations() with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: chain transforms: - type: Create config: elements: [0, 1, 2, 3] - type: %r config: %s ''' % (annotations['yaml_type'], annotations['yaml_args'])) assert_that(result, equal_to([100, 105, 110, 115])) def test_resource_hints(self): t = LinearTransform(5, b=100) with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: chain transforms: - type: Create config: elements: [0, 1, 2, 3] - type: MapToFields name: WithResourceHints config: language: python fields: square: element * element resource_hints: min_ram: 1GB ''') assert_that(result | beam.Map(lambda x: x.square), equal_to([0, 1, 4, 9])) proto = p.to_runner_api() transform, = [ t for t in proto.components.transforms.values() if t.unique_name == 'YamlTransform/Chain/WithResourceHints'] self.assertEqual( proto.components.environments[transform.environment_id]. resource_hints['beam:resources:min_ram_bytes:v1'], b'1000000000', proto) def test_composite_resource_hints(self): t = LinearTransform(5, b=100) with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: chain transforms: - type: Create config: elements: [0, 1, 2, 3] - type: MapToFields name: WithInheritedResourceHints config: language: python fields: square: element * element resource_hints: min_ram: 1GB ''') assert_that(result | beam.Map(lambda x: x.square), equal_to([0, 1, 4, 9])) proto = p.to_runner_api() transform, = [ t for t in proto.components.transforms.values() if t.unique_name == 'YamlTransform/Chain/WithInheritedResourceHints'] self.assertEqual( proto.components.environments[transform.environment_id]. resource_hints['beam:resources:min_ram_bytes:v1'], b'1000000000', proto) def test_flatten_unifies_schemas(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: _ = p | YamlTransform( ''' type: composite transforms: - type: Create name: Create1 config: elements: - {ride_id: '1', passenger_count: 1} - {ride_id: '2', passenger_count: 2} - type: Create name: Create2 config: elements: - {ride_id: '3'} - {ride_id: '4'} - type: Flatten input: [Create1, Create2] - type: AssertEqual input: Flatten config: elements: - {ride_id: '1', passenger_count: 1} - {ride_id: '2', passenger_count: 2} - {ride_id: '3'} - {ride_id: '4'} ''') def test_flatten_unifies_optional_fields(self): """Test that Flatten correctly unifies schemas with optional fields.""" with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: _ = p | YamlTransform( ''' type: composite transforms: - type: Create name: Create1 config: elements: - {id: '1', name: 'Alice', age: 30} - {id: '2', name: 'Bob', age: 25} - type: Create name: Create2 config: elements: - {id: '3', name: 'Charlie'} - {id: '4', name: 'Diana'} - type: Flatten input: [Create1, Create2] - type: AssertEqual input: Flatten config: elements: - {id: '1', name: 'Alice', age: 30} - {id: '2', name: 'Bob', age: 25} - {id: '3', name: 'Charlie'} - {id: '4', name: 'Diana'} ''') def test_flatten_unifies_different_types(self): """Test that Flatten correctly unifies schemas with different field types.""" with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: _ = p | YamlTransform( ''' type: composite transforms: - type: Create name: Create1 config: elements: - {id: 1, value: 100} - {id: 2, value: 200} - type: Create name: Create2 config: elements: - {id: '3', value: 'text'} - {id: '4', value: 'data'} - type: Flatten input: [Create1, Create2] - type: AssertEqual input: Flatten config: elements: - {id: 1, value: 100} - {id: 2, value: 200} - {id: '3', value: 'text'} - {id: '4', value: 'data'} ''') def test_flatten_unifies_list_fields(self): """Test that Flatten correctly unifies schemas with list fields.""" with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: _ = p | YamlTransform( ''' type: composite transforms: - type: Create name: Create1 config: elements: - {id: '1', tags: ['red', 'blue']} - {id: '2', tags: ['green']} - type: Create name: Create2 config: elements: - {id: '3', tags: ['yellow', 'purple', 'orange']} - {id: '4', tags: []} - type: Flatten input: [Create1, Create2] - type: AssertEqual input: Flatten config: elements: - {id: '1', tags: ['red', 'blue']} - {id: '2', tags: ['green']} - {id: '3', tags: ['yellow', 'purple', 'orange']} - {id: '4', tags: []} ''') def test_flatten_unifies_with_missing_fields(self): """Test that Flatten correctly unifies schemas when some inputs have missing fields.""" with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: _ = p | YamlTransform( ''' type: composite transforms: - type: Create name: Create1 config: elements: - {id: '1', name: 'Alice', department: 'Engineering', salary: 75000} - {id: '2', name: 'Bob', department: 'Marketing', salary: 65000} - type: Create name: Create2 config: elements: - {id: '3', name: 'Charlie', department: 'Sales'} - {id: '4', name: 'Diana'} - type: Flatten input: [Create1, Create2] - type: AssertEqual input: Flatten config: elements: - {id: '1', name: 'Alice', department: 'Engineering', salary: 75000} - {id: '2', name: 'Bob', department: 'Marketing', salary: 65000} - {id: '3', name: 'Charlie', department: 'Sales'} - {id: '4', name: 'Diana'} ''') def test_flatten_unifies_complex_mixed_schemas(self): """Test that Flatten correctly unifies complex mixed schemas.""" with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: _ = p | YamlTransform( ''' type: composite transforms: - type: Create name: Create1 config: elements: - {id: 1, name: 'Product A', price: 29.99, categories: ['electronics', 'gadgets']} - {id: 2, name: 'Product B', price: 15.50, categories: ['books']} - type: Create name: Create2 config: elements: - {id: 3, name: 'Product C', categories: ['clothing']} - {id: 4, name: 'Product D', price: 99.99} - type: Create name: Create3 config: elements: - {id: 5, name: 'Product E', price: 5.00, categories: []} - type: Flatten input: [Create1, Create2, Create3] - type: AssertEqual input: Flatten config: elements: - {id: 1, name: 'Product A', price: 29.99, categories: ['electronics', 'gadgets']} - {id: 2, name: 'Product B', price: 15.50, categories: ['books']} - {id: 3, name: 'Product C', categories: ['clothing']} - {id: 4, name: 'Product D', price: 99.99} - {id: 5, name: 'Product E', price: 5.00, categories: []} ''') def test_output_schema_success(self): """Test that optional output_schema works.""" with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: _ = p | YamlTransform( ''' type: composite transforms: - type: Create name: MyCreate config: elements: - {sdk: 'Beam', year: 2016} - {sdk: 'Flink', year: 2015} output_schema: type: object properties: sdk: type: string year: type: integer - type: AssertEqual name: CheckGood input: MyCreate config: elements: - {sdk: 'Beam', year: 2016} - {sdk: 'Flink', year: 2015} ''') def test_output_schema_fails(self): """ Test that optional output_schema works by failing the pipeline since main transform doesn't have error_handling config. """ with self.assertRaises(Exception) as e: with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: _ = p | YamlTransform( ''' type: composite transforms: - type: Create name: MyCreate config: elements: - {sdk: 'Beam', year: 2016} - {sdk: 'Spark', year: 'date'} - {sdk: 'Flink', year: 2015} output_schema: type: object properties: sdk: type: string year: type: integer - type: AssertEqual name: CheckGood input: MyCreate config: elements: - {sdk: 'Beam', year: 2016} - {sdk: 'Flink', year: 2015} ''') self.assertIn("'date' is not of type 'integer'", str(e.exception)) def test_output_schema_with_main_transform_error_handling_success(self): """Test that optional output_schema works in conjunction with main transform error handling.""" with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: _ = p | YamlTransform( ''' type: composite transforms: - type: Create name: CreateVisits config: elements: - {user: alice, timestamp: "not-valid"} - {user: bob, timestamp: 3} - type: AssignTimestamps input: CreateVisits config: timestamp: timestamp error_handling: output: invalid_rows output_schema: type: object properties: user: type: string timestamp: type: boolean - type: MapToFields name: ExtractInvalidTimestamp input: AssignTimestamps.invalid_rows config: language: python fields: user: "element.user" timestamp: "element.timestamp" - type: AssertEqual input: ExtractInvalidTimestamp config: elements: - {user: "alice", timestamp: "not-valid"} - {user: bob, timestamp: 3} - type: AssertEqual input: AssignTimestamps config: elements: [] ''') class ErrorHandlingTest(unittest.TestCase): def test_error_handling_outputs(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: composite transforms: - type: Create config: elements: ['a', 'b', 'biiiiig'] - type: SizeLimiter input: Create config: limit: 5 error_handling: output: errors - name: TrimErrors type: PyMap input: SizeLimiter.errors config: fn: "lambda x: x[1][1]" output: good: SizeLimiter bad: TrimErrors ''', providers=TEST_PROVIDERS) assert_that(result['good'], equal_to(['a', 'b']), label="CheckGood") assert_that(result['bad'], equal_to(["ValueError('biiiiig')"])) def test_strip_error_metadata(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: composite transforms: - type: Create config: elements: ['a', 'b', 'biiiiig'] - type: SizeLimiter input: Create config: limit: 5 error_handling: output: errors - type: StripErrorMetadata name: StripErrorMetadata1 input: SizeLimiter.errors - type: MapToFields input: Create config: language: python fields: out: "1.0/(1-len(element))" error_handling: output: errors - type: StripErrorMetadata name: StripErrorMetadata2 input: MapToFields.errors output: good: SizeLimiter bad1: StripErrorMetadata1 bad2: StripErrorMetadata2 ''', providers=TEST_PROVIDERS) assert_that(result['good'], equal_to(['a', 'b']), label="CheckGood") assert_that( result['bad1'] | beam.Map(lambda x: x.element), equal_to(['biiiiig'])) assert_that( result['bad2'] | beam.Map(lambda x: x.element), equal_to(['a', 'b'])) def test_must_handle_error_output(self): with self.assertRaisesRegex(Exception, 'Unconsumed error output .*line 7'): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: _ = p | YamlTransform( ''' type: composite transforms: - type: Create config: elements: ['a', 'b', 'biiiiig'] - type: SizeLimiter input: Create config: limit: 5 error_handling: output: errors ''', providers=TEST_PROVIDERS) def test_error_handling_log_combined_errors(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: composite transforms: - type: Create name: Input1 config: elements: [1, 2, 0] - type: Create name: Input2 config: elements: [3, 'a', 5] - type: MapToFields name: Inverse input: Input1 config: language: python fields: inverse: "1 / element" error_handling: output: errors - type: MapToFields name: Square input: Input2 config: language: python fields: square: "element * element" error_handling: output: errors - type: LogForTesting input: - Inverse.errors - Square.errors - type: Flatten name: GoodData input: - Inverse - Square output: GoodData ''', providers=TEST_PROVIDERS) assert_that( result, equal_to([ beam.Row(inverse=1.0, square=None), beam.Row(inverse=0.5, square=None), beam.Row(square=9, inverse=None), beam.Row(square=25, inverse=None) ])) def test_mapping_errors(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: composite transforms: - type: Create config: elements: [0, 1, 2, 4] - type: MapToFields name: ToRow input: Create config: language: python fields: num: element str: "'a' * element or 'bbb'" - type: Filter input: ToRow config: language: python keep: str[1] >= 'a' error_handling: output: errors - type: MapToFields name: MapWithErrorHandling input: Filter config: language: python fields: num: num inverse: float(1 / num) error_handling: output: errors - type: PyMap name: TrimErrors input: [MapWithErrorHandling.errors, Filter.errors] config: fn: "lambda x: x.msg" - type: MapToFields name: Sum input: MapWithErrorHandling config: language: python append: True fields: sum: num + inverse output: good: Sum bad: TrimErrors ''', providers=TEST_PROVIDERS) assert_that( result['good'], equal_to([ beam.Row(num=2, inverse=.5, sum=2.5), beam.Row(num=4, inverse=.25, sum=4.25) ]), label="CheckGood") assert_that( result['bad'], equal_to([ "IndexError('string index out of range')", # from the filter "ZeroDivisionError('division by zero')", # from the mapping ]), label='CheckErrors') class YamlWindowingTest(unittest.TestCase): def test_explicit_window_into(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: chain transforms: - type: CreateTimestamped config: elements: [0, 1, 2, 3, 4, 5] - type: WindowInto windowing: type: fixed size: 4s - type: SumGlobally ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([6, 9])) def test_explicit_window_into_with_json_string_config_one_line(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: chain transforms: - type: CreateTimestamped config: elements: [0, 1, 2, 3, 4, 5] - type: WindowInto config: windowing: {"type": "fixed", "size": "4s"} - type: SumGlobally ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([6, 9])) def test_explicit_window_into_with_json_string_config_multi_line(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: chain transforms: - type: CreateTimestamped config: elements: [0, 1, 2, 3, 4, 5] - type: WindowInto config: windowing: | {"type": "fixed", "size": "4s"} - type: SumGlobally ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([6, 9])) def test_explicit_window_into_with_string_config_fails(self): with self.assertRaisesRegex(ValueError, 'Error parsing windowing config'): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: _ = p | YamlTransform( ''' type: chain transforms: - type: CreateTimestamped config: elements: [0, 1, 2, 3, 4, 5] - type: WindowInto config: windowing: | 'not a valid yaml' ''', providers=TEST_PROVIDERS) def test_windowing_on_input(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: chain transforms: - type: CreateTimestamped config: elements: [0, 1, 2, 3, 4, 5] - type: SumGlobally windowing: type: fixed size: 4s ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([6, 9])) def test_windowing_multiple_inputs(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: composite transforms: - type: CreateTimestamped name: Create1 config: elements: [0, 2, 4] - type: CreateTimestamped name: Create2 config: elements: [1, 3, 5] - type: SumGlobally input: [Create1, Create2] windowing: type: fixed size: 4s output: SumGlobally ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([6, 9])) def test_windowing_on_output(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: chain transforms: - type: CreateTimestamped config: elements: [0, 1, 2, 3, 4, 5] windowing: type: fixed size: 4s - type: SumGlobally ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([6, 9])) def test_windowing_on_outer(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: chain transforms: - type: CreateTimestamped config: elements: [0, 1, 2, 3, 4, 5] - type: SumGlobally windowing: type: fixed size: 4s ''', providers=TEST_PROVIDERS) assert_that(result, equal_to([6, 9])) def test_assign_timestamps(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result = p | YamlTransform( ''' type: chain transforms: - type: Create config: elements: - {t: 1, v: 1} - {t: 10, v: 2} - {t: 11, v: 3} - type: AssignTimestamps config: timestamp: t - type: Combine config: group_by: [] combine: v: sum windowing: type: fixed size: 10s ''', providers=TEST_PROVIDERS) assert_that( result | beam.Map(lambda x: beam.Row(**x._asdict())), equal_to([ beam.Row(v=1), beam.Row(v=5), ])) class AnnotatingProvider(yaml_provider.InlineProvider): """A provider that vends transforms that do nothing but record that this provider (as identified by name) was used, along with any prior history of the given element. """ def __init__(self, name, transform_names): super().__init__({ transform_name: lambda: beam.Map( lambda x: (x if type(x) == tuple else ()) + (name, )) for transform_name in transform_names.strip().split() }) self._name = name def __repr__(self): return 'AnnotatingProvider(%r)' % self._name class AnotherAnnProvider(AnnotatingProvider): """A Provider that behaves exactly as AnnotatingProvider, but is not of the same type and so is considered "more distant" for matching purposes. """ pass class ProviderAffinityTest(unittest.TestCase): """These tests check that for a sequence of transforms, the "closest" proveders are chosen among multiple possible implementations. """ provider1 = AnnotatingProvider("provider1", "P1 A B C ") provider2 = AnnotatingProvider("provider2", "P2 A C D") provider3 = AnotherAnnProvider("provider3", "P3 A B ") provider4 = AnotherAnnProvider("provider4", "P4 A B D") providers_dict = collections.defaultdict(list) for provider in [provider1, provider2, provider3, provider4]: for transform_type in provider.provided_transforms(): providers_dict[transform_type].append(provider) def test_prefers_same_provider(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result1 = p | 'Yaml1' >> YamlTransform( ''' type: chain transforms: - type: Create config: elements: [0] - type: P1 - type: A - type: C ''', providers=self.providers_dict) assert_that( result1, equal_to([( # provider1 was chosen, as it is the only one vending P1 'provider1', # All of the providers vend A, but since the input was produced # by provider1, we prefer to use that again. 'provider1', # Similarly for C. 'provider1')]), label='StartWith1') result2 = p | 'Yaml2' >> YamlTransform( ''' type: chain transforms: - type: Create config: elements: [0] - type: P2 - type: A - type: C ''', providers=self.providers_dict) assert_that( result2, equal_to([( # provider2 was necessarily chosen for P2 'provider2', # Unlike above, we choose provider2 to implement A. 'provider2', # Likewise for C. 'provider2')]), label='StartWith2') def test_prefers_same_provider_class(self): # Like test_prefers_same_provider, but as we cannot choose the same # exact provider, we go with the next closest (which is of the same type) # over an implementation from a Provider of a different type. with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: result1 = p | 'Yaml1' >> YamlTransform( ''' type: chain transforms: - type: Create config: elements: [0] - type: P1 - type: A - type: D - type: A ''', providers=self.providers_dict) assert_that( result1, equal_to([('provider1', 'provider1', 'provider2', 'provider2')]), label='StartWith1') result3 = p | 'Yaml2' >> YamlTransform( ''' type: chain transforms: - type: Create config: elements: [0] - type: P3 - type: A - type: D - type: A ''', providers=self.providers_dict) assert_that( result3, equal_to([('provider3', 'provider3', 'provider4', 'provider4')]), label='StartWith3') class TestExternalYamlProvider(unittest.TestCase): def setUp(self): self.temp_dir = tempfile.mkdtemp() self.provider_path = os.path.join(self.temp_dir, 'power_provider.yaml') with open(self.provider_path, 'w') as f: f.write( """ - type: yaml transforms: RaiseElementToPower: config_schema: properties: n: {type: integer} body: type: MapToFields config: language: python append: true fields: power: "element ** {{n}}" error_handling: output: my_error """) def tearDown(self): shutil.rmtree(self.temp_dir) def test_provider_with_error_handling(self): loaded_providers = yaml_provider.load_providers(self.provider_path) test_providers = yaml_provider.InlineProvider(TEST_PROVIDERS) merged_providers = yaml_provider.merge_providers( loaded_providers, [test_providers]) with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: results = p | YamlTransform( ''' type: composite transforms: - type: Create config: elements: [2, 'bad', 3] - type: RaiseElementToPower input: Create config: n: 2 - type: PyMap name: TrimErrors input: RaiseElementToPower.my_error config: fn: "lambda x: x.msg" output: good: RaiseElementToPower.good bad: TrimErrors ''', providers=merged_providers) assert_that( results['good'], equal_to([beam.Row(element=2, power=4), beam.Row(element=3, power=9)]), label="CheckGood") assert_that( results['bad'], equal_to([ 'TypeError("unsupported operand type(s) for ** or pow(): ' + '\'str\' and \'int\'")' ]), label="CheckBad") def test_must_consume_error_output(self): # By adding a dummy error_handling block here, we signal to the static # checker that this transform has an error output that must be consumed. # The framework is able to handle the "nesting" where the provider for # RaiseElementToPower also defines error handling internally. loaded_providers = yaml_provider.load_providers(self.provider_path) test_providers = yaml_provider.InlineProvider(TEST_PROVIDERS) merged_providers = yaml_provider.merge_providers( loaded_providers, [test_providers]) with self.assertRaisesRegex(Exception, 'Unconsumed error output.*'): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: _ = p | YamlTransform( ''' type: composite transforms: - type: Create config: elements: [2, 'bad', 3] - type: RaiseElementToPower input: Create config: n: 2 error_handling: output: my_error ''', providers=merged_providers) @beam.transforms.ptransform.annotate_yaml class LinearTransform(beam.PTransform): """A transform used for testing annotate_yaml.""" def __init__(self, a, b): self._a = a self._b = b def expand(self, pcoll): a = self._a b = self._b return pcoll | beam.Map(lambda x: a * x.element + b) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_transform_unit_test.py0000644000175100017510000007552415206113033024342 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest import yaml import apache_beam as beam from apache_beam.yaml import YamlTransform from apache_beam.yaml import yaml_provider from apache_beam.yaml.yaml_provider import InlineProvider from apache_beam.yaml.yaml_transform import Scope from apache_beam.yaml.yaml_transform import chain_as_composite from apache_beam.yaml.yaml_transform import ensure_errors_consumed from apache_beam.yaml.yaml_transform import ensure_transforms_have_types from apache_beam.yaml.yaml_transform import expand_composite_transform from apache_beam.yaml.yaml_transform import expand_pipeline from apache_beam.yaml.yaml_transform import extract_name from apache_beam.yaml.yaml_transform import get_main_output_key from apache_beam.yaml.yaml_transform import identify_object from apache_beam.yaml.yaml_transform import normalize_inputs_outputs from apache_beam.yaml.yaml_transform import normalize_source_sink from apache_beam.yaml.yaml_transform import only_element from apache_beam.yaml.yaml_transform import pipeline_as_composite from apache_beam.yaml.yaml_transform import preprocess from apache_beam.yaml.yaml_transform import preprocess_flattened_inputs from apache_beam.yaml.yaml_transform import preprocess_windowing from apache_beam.yaml.yaml_transform import push_windowing_to_roots from apache_beam.yaml.yaml_transform import validate_against_schema from apache_beam.yaml.yaml_utils import SafeLineLoader try: import jsonschema except ImportError: jsonschema = None def new_pipeline(): return beam.Pipeline( options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) @unittest.skipIf(jsonschema is None, "Yaml dependencies not installed") class MainTest(unittest.TestCase): def assertYaml(self, expected, result): result = SafeLineLoader.strip_metadata(result) expected = yaml.load(expected, Loader=SafeLineLoader) expected = SafeLineLoader.strip_metadata(expected) self.assertEqual(expected, result) def get_scope_by_spec(self, p, spec, inputs=None): if inputs is None: inputs = {} spec = yaml.load(spec, Loader=SafeLineLoader) if 'transforms' in spec: spec['transforms'] = [ normalize_inputs_outputs(t) for t in spec['transforms'] ] scope = Scope( beam.pvalue.PBegin(p), inputs, spec['transforms'], yaml_provider.standard_providers(), {}) return scope, spec def test_pipeline_as_composite_with_type_transforms(self): spec = ''' type: composite transforms: - type: Create config: elements: [0,1,2] - type: PyMap config: fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) result = pipeline_as_composite(spec) self.assertEqual(result['type'], 'composite') self.assertEqual(result['name'], None) def test_pipeline_as_composite_with_transforms(self): spec = ''' transforms: - type: Create config: elements: [0,1,2] - type: PyMap config: fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) result = pipeline_as_composite(spec) self.assertEqual(result['type'], 'composite') self.assertEqual(result['name'], None) def test_pipeline_as_composite_list(self): spec = ''' - type: Create config: elements: [0,1,2] - type: PyMap config: fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) result = pipeline_as_composite(spec) expected = ''' type: composite name: null transforms: - type: Create config: elements: [0,1,2] - type: PyMap config: fn: 'lambda x: x*x' ''' self.assertYaml(expected, result) def test_expand_composite_transform_with_name(self): with new_pipeline() as p: spec = ''' type: composite name: Custom transforms: - type: Create config: elements: [0,1,2] output: Create ''' scope, spec = self.get_scope_by_spec(p, spec) self.assertRegex( str(expand_composite_transform(spec, scope)['output']), r"PCollection.*Custom/Create/Map.*") def test_expand_composite_transform_with_name_input(self): with new_pipeline() as p: spec = ''' type: composite input: elements transforms: - type: LogForTesting input: input output: LogForTesting ''' elements = p | beam.Create(range(3)) scope, spec = self.get_scope_by_spec(p, spec, inputs={'elements': elements}) self.assertRegex( str(expand_composite_transform(spec, scope)['output']), r"PCollection.*Composite/log_for_testing/LogForTesting.*") def test_expand_composite_transform_root(self): with new_pipeline() as p: spec = ''' type: composite transforms: - type: Create config: elements: [0,1,2] output: Create ''' scope, spec = self.get_scope_by_spec(p, spec) self.assertRegex( str(expand_composite_transform(spec, scope)['output']), r"PCollection.*Create/Map.*") def test_chain_as_composite(self): spec = ''' type: chain transforms: - type: Create config: elements: [0,1,2] - type: PyMap config: fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) result = chain_as_composite(spec) expected = f''' type: composite name: Chain transforms: - type: Create config: elements: [0,1,2] - type: PyMap config: fn: 'lambda x: x*x' input: {spec['transforms'][0]['__uuid__']} output: '__implicit_outputs__': {spec['transforms'][1]['__uuid__']} ''' self.assertYaml(expected, result) def test_chain_as_composite_with_wrong_output_type(self): spec = ''' type: chain transforms: - type: Create elements: [0,1,2] - type: PyMap fn: 'lambda x: x*x' output: Create ''' spec = yaml.load(spec, Loader=SafeLineLoader) with self.assertRaisesRegex(ValueError, r"Explicit output.*of the chain transform is " r"not an output of the last transform"): chain_as_composite(spec) def test_chain_as_composite_with_wrong_output_name(self): spec = ''' type: chain transforms: - type: Create name: elements elements: [0,1,2] - type: PyMap fn: 'lambda x: x*x' output: elements ''' spec = yaml.load(spec, Loader=SafeLineLoader) with self.assertRaisesRegex(ValueError, r"Explicit output.*of the chain transform is " r"not an output of the last transform"): chain_as_composite(spec) def test_chain_as_composite_with_outputs_override(self): spec = ''' type: chain transforms: - type: Create elements: [0,1,2] - type: PyMap fn: 'lambda x: x*x' output: PyMap ''' spec = yaml.load(spec, Loader=SafeLineLoader) self.assertEqual( chain_as_composite(spec)['output']['output'], f"{spec['transforms'][1]['__uuid__']}.PyMap") def test_chain_as_composite_with_input(self): spec = ''' type: chain input: elements transforms: - type: PyMap fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) self.assertEqual( chain_as_composite(spec)['transforms'][0]['input'], {"input": "input"}) def test_chain_as_composite_with_transform_input(self): spec = ''' type: chain transforms: - type: Create config: elements: [0,1,2] - type: LogForTesting input: Create ''' spec = yaml.load(spec, Loader=SafeLineLoader) with self.assertRaisesRegex( ValueError, r"Transform .* is part of a chain. " r"Cannot define explicit inputs on chain pipeline"): chain_as_composite(spec) def test_validate_against_schema_with_transform_input(self): spec = ''' pipeline: type: chain transforms: - type: Create config: elements: [1, 2, 3] - type: LogForTesting input: Create ''' spec = yaml.load(spec, Loader=SafeLineLoader) with self.assertRaisesRegex(jsonschema.ValidationError, r"'input' should not be used " r"along with 'chain' type transforms.*"): validate_against_schema(spec, 'generic') def test_normalize_source_sink(self): spec = ''' source: type: Create config: elements: [0,1,2] transforms: - type: PyMap config: fn: 'lambda x: x*x' sink: type: PyMap config: fn: "lambda x: x + 41" ''' spec = yaml.load(spec, Loader=SafeLineLoader) result = normalize_source_sink(spec) expected = ''' transforms: - type: Create input: {'__explicitly_empty__': null} config: elements: [0,1,2] - type: PyMap config: fn: 'lambda x: x*x' - type: PyMap config: fn: "lambda x: x + 41" ''' self.assertYaml(expected, result) def test_normalize_source_sink_only_source(self): spec = ''' source: type: Create config: elements: [0,1,2] transforms: - type: PyMap config: fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) result = normalize_source_sink(spec) expected = ''' transforms: - type: Create input: {'__explicitly_empty__': null} config: elements: [0,1,2] - type: PyMap config: fn: 'lambda x: x*x' ''' self.assertYaml(expected, result) def test_normalize_source_sink_only_sink(self): spec = ''' transforms: - type: PyMap config: fn: 'lambda x: x*x' sink: type: PyMap config: fn: "lambda x: x + 41" ''' spec = yaml.load(spec, Loader=SafeLineLoader) result = normalize_source_sink(spec) expected = ''' transforms: - type: PyMap config: fn: 'lambda x: x*x' - type: PyMap config: fn: "lambda x: x + 41" ''' self.assertYaml(expected, result) def test_normalize_source_sink_no_source_no_sink(self): spec = ''' transforms: - type: PyMap config: fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) result = normalize_source_sink(spec) expected = ''' transforms: - type: PyMap config: fn: 'lambda x: x*x' ''' self.assertYaml(expected, result) def test_preprocess_source_sink_composite(self): spec = ''' type: composite source: type: Create config: elements: [0,1,2] transforms: - type: PyMap config: fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) result = normalize_source_sink(spec) expected = ''' type: composite transforms: - type: Create input: {'__explicitly_empty__': null} config: elements: [0,1,2] - type: PyMap config: fn: 'lambda x: x*x' ''' self.assertYaml(expected, result) def test_preprocess_source_sink_chain(self): spec = ''' type: chain source: type: Create config: elements: [0,1,2] transforms: - type: PyMap config: fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) result = normalize_source_sink(spec) expected = ''' type: chain transforms: - type: Create input: {'__explicitly_empty__': null} config: elements: [0,1,2] - type: PyMap config: fn: 'lambda x: x*x' ''' self.assertYaml(expected, result) def test_preprocess_source_sink_other(self): spec = ''' - type: PyMap fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) self.assertEqual(normalize_source_sink(spec), spec) def test_normalize_inputs_outputs(self): spec = ''' type: PyMap input: [Create1, Create2] fn: 'lambda x: x*x' output: Squared ''' spec = yaml.load(spec, Loader=SafeLineLoader) result = normalize_inputs_outputs(spec) expected = ''' type: PyMap input: input: [Create1, Create2] fn: 'lambda x: x*x' output: output: Squared ''' self.assertYaml(expected, result) def test_normalize_inputs_outputs_dict(self): spec = ''' type: PyMap input: [Create1, Create2] fn: 'lambda x: x*x' output: out1: Squared1 out2: Squared2 ''' spec = yaml.load(spec, Loader=SafeLineLoader) result = normalize_inputs_outputs(spec) expected = ''' type: PyMap input: input: [Create1, Create2] fn: 'lambda x: x*x' output: out1: Squared1 out2: Squared2 ''' self.assertYaml(expected, result) def test_identify_object_with_name(self): spec = ''' type: PyMap fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) self.assertRegex(identify_object(spec), r"PyMap.*[0-9]") def test_identify_object(self): spec = ''' argument: PyMap fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) self.assertRegex(identify_object(spec), r"at.*[0-9]") def test_extract_name_by_type(self): spec = ''' type: PyMap fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) self.assertEqual(extract_name(spec), "PyMap") def test_extract_name_by_id(self): spec = ''' type: PyMap id: PyMapId fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) self.assertEqual(extract_name(spec), "PyMapId") def test_extract_name_by_name(self): spec = ''' type: PyMap name: PyMapName fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) self.assertEqual(extract_name(spec), "PyMapName") def test_extract_name_no_name(self): spec = ''' transforms: - arg: PyMap fn: 'lambda x: x*x' ''' spec = yaml.load(spec, Loader=SafeLineLoader) self.assertEqual(extract_name(spec), "") def test_push_windowing_to_roots(self): spec = ''' type: composite transforms: - type: Create elements: [0,1,2] - type: PyMap fn: 'lambda x: x*x' input: Create windowing: type: fixed size: 2 ''' spec = yaml.load(spec, Loader=SafeLineLoader) spec = normalize_inputs_outputs(spec) spec['transforms'] = [ normalize_inputs_outputs(t) for t in spec['transforms'] ] result = push_windowing_to_roots(spec) expected = ''' type: composite transforms: - type: Create elements: [0,1,2] windowing: type: fixed size: 2 __consumed_outputs: - null input: {} output: {} - type: PyMap fn: 'lambda x: x*x' input: input: Create output: {} windowing: type: fixed size: 2 input: {} output: {} ''' self.assertYaml(expected, result) def test_preprocess_windowing_custom_type(self): spec = ''' type: SumGlobally input: Create windowing: type: fixed size: 4 ''' spec = yaml.load(spec, Loader=SafeLineLoader) spec = normalize_inputs_outputs(spec) result = preprocess_windowing(spec) expected = f''' type: composite name: SumGlobally input: input: Create transforms: - type: SumGlobally input: input: {result['transforms'][1]['__uuid__']} output: {{}} - type: WindowInto name: WindowInto[input] windowing: type: fixed size: 4 input: {{input: input}} output: {result['transforms'][0]['__uuid__']} config: error_handling: {{}} ''' self.assertYaml(expected, result) def test_preprocess_windowing_composite_with_windowing_outer(self): spec = ''' type: composite transforms: - type: CreateTimestamped name: Create elements: [0, 2, 4] - type: SumGlobally input: Create windowing: type: fixed size: 4 output: SumGlobally ''' spec = yaml.load(spec, Loader=SafeLineLoader) spec = normalize_inputs_outputs(spec) spec['transforms'] = [ normalize_inputs_outputs(t) for t in spec['transforms'] ] result = preprocess_windowing(spec) expected = ''' type: composite input: {} transforms: - type: CreateTimestamped name: Create elements: [0, 2, 4] windowing: type: fixed size: 4 __consumed_outputs: - null input: {} output: {} - type: SumGlobally input: input: Create output: {} output: output: SumGlobally ''' self.assertYaml(expected, result) def test_preprocess_windowing_composite_with_windowing_on_input(self): spec = ''' type: composite transforms: - type: CreateTimestamped name: Create elements: [0, 2, 4] - type: SumGlobally input: Create windowing: type: fixed size: 4 output: SumGlobally ''' spec = yaml.load(spec, Loader=SafeLineLoader) spec = normalize_inputs_outputs(spec) spec['transforms'] = [ normalize_inputs_outputs(t) for t in spec['transforms'] ] result = preprocess_windowing(spec) expected = ''' type: composite input: {} transforms: - type: CreateTimestamped name: Create elements: [0, 2, 4] input: {} output: {} - type: SumGlobally input: input: Create windowing: type: fixed size: 4 output: {} output: output: SumGlobally ''' self.assertYaml(expected, result) def test_preprocess_windowing_other_type_with_no_inputs(self): spec = ''' type: CreateTimestamped name: Create elements: [0, 2, 4] windowing: type: fixed size: 4 ''' spec = yaml.load(spec, Loader=SafeLineLoader) spec = normalize_inputs_outputs(spec) result = preprocess_windowing(spec) expected = f''' type: composite name: Create transforms: - type: CreateTimestamped name: Create elements: [0, 2, 4] input: {{}} output: {{}} - type: WindowInto name: WindowInto[None] input: input: {result['transforms'][0]["__uuid__"]} windowing: type: fixed size: 4 output: {result['transforms'][1]["__uuid__"]} config: error_handling: {{}} ''' self.maxDiff = 1e9 self.assertYaml(expected, result) def test_preprocess_flattened_inputs_implicit(self): spec = ''' type: composite transforms: - type: PyMap fn: 'lambda x: x*x' input: [Create1, Create2] output: CreateTimestamped ''' spec = yaml.load(spec, Loader=SafeLineLoader) spec['transforms'] = [ normalize_inputs_outputs(t) for t in spec['transforms'] ] result = preprocess_flattened_inputs(spec) expected = f''' type: composite transforms: - type: Flatten name: PyMap-Flatten[input] input: input0: Create1 input1: Create2 - type: PyMap fn: 'lambda x: x*x' input: input: {result['transforms'][0]['__uuid__']} output: {{}} output: CreateTimestamped ''' self.assertYaml(expected, result) def test_preprocess_flattened_inputs_explicit_flatten(self): spec = ''' type: composite transforms: - type: Flatten input: [Create1, Create2] - type: PyMap fn: 'lambda x: x*x' input: Flatten output: CreateTimestamped ''' spec = yaml.load(spec, Loader=SafeLineLoader) spec['transforms'] = [ normalize_inputs_outputs(t) for t in spec['transforms'] ] result = preprocess_flattened_inputs(spec) expected = ''' type: composite transforms: - type: Flatten input: input0: Create1 input1: Create2 output: {} - type: PyMap fn: 'lambda x: x*x' input: input: Flatten output: {} output: CreateTimestamped ''' self.assertYaml(expected, result) def test_ensure_transforms_have_types(self): spec = ''' type: PyMap fn: 'lambda x: x*x' input: Flatten ''' spec = yaml.load(spec, Loader=SafeLineLoader) result = ensure_transforms_have_types(spec) self.assertEqual(result, spec) def test_ensure_transforms_have_types_error(self): spec = ''' name: PyMap fn: 'lambda x: x*x' input: Flatten ''' spec = yaml.load(spec, Loader=SafeLineLoader) with self.assertRaisesRegex(ValueError, r"Missing type .*"): ensure_transforms_have_types(spec) with self.assertRaisesRegex(ValueError, r"Missing type .*"): preprocess(spec) def test_ensure_transforms_have_providers_error(self): spec = ''' type: UnknownType ''' spec = yaml.load(spec, Loader=SafeLineLoader) with self.assertRaisesRegex(ValueError, r"Unknown type or missing provider .*"): preprocess(spec, known_transforms=['KnownType']) def test_ensure_errors_consumed_unconsumed(self): spec = ''' type: composite transforms: - type: Create elements: [1,2,3] - type: MyTransform input: Create error_handling: output: errors output: good: MyTransform ''' spec = yaml.load(spec, Loader=SafeLineLoader) spec = normalize_inputs_outputs(spec) spec['transforms'] = [ normalize_inputs_outputs(t) for t in spec['transforms'] ] with self.assertRaisesRegex(ValueError, r"Unconsumed error.*MyTransform.errors"): ensure_errors_consumed(spec) def test_ensure_errors_consumed_in_transform(self): spec = ''' type: composite transforms: - type: Create elements: [1,2,3] - type: MyTransform input: Create error_handling: output: errors - name: SaveToFile type: PyMap input: MyTransform.errors output: good: MyTransform ''' spec = yaml.load(spec, Loader=SafeLineLoader) spec = normalize_inputs_outputs(spec) spec['transforms'] = [ normalize_inputs_outputs(t) for t in spec['transforms'] ] result = ensure_errors_consumed(spec) self.assertEqual(result, spec) self.assertEqual(result['transforms'], spec['transforms']) def test_ensure_errors_consumed_no_output_in_error_handling(self): spec = ''' type: composite transforms: - type: Create elements: [1,2,3] - type: MyTransform input: Create error_handling: arg: errors - name: SaveToFile type: PyMap input: MyTransform.errors output: good: MyTransform ''' spec = yaml.load(spec, Loader=SafeLineLoader) spec = normalize_inputs_outputs(spec) spec['transforms'] = [ normalize_inputs_outputs(t) for t in spec['transforms'] ] with self.assertRaisesRegex(ValueError, r"Missing output.*"): ensure_errors_consumed(spec) def test_only_element(self): self.assertEqual(only_element((1, )), 1) def test_get_main_output_key(self): spec = {'type': 'TestTransform'} error_handling_spec = {'output': 'invalid_rows'} # Case 1: 'output' key exists outputs = {'output': 1, 'another': 2} with self.assertLogs('apache_beam.yaml.yaml_transform', level='WARNING') as al: self.assertEqual( get_main_output_key(spec, outputs, error_handling_spec), 'output') self.assertIn("Only the main output will be validated.", al.output[0]) # Case 2: 'good' key exists, 'output' does not outputs = {'good': 1, 'another': 2} with self.assertLogs('apache_beam.yaml.yaml_transform', level='WARNING') as al: self.assertEqual( get_main_output_key(spec, outputs, error_handling_spec), 'good') self.assertIn("Only the main output will be validated.", al.output[0]) # Case 3: Only one output outputs = {'single_output': 1} self.assertEqual( get_main_output_key(spec, outputs, error_handling_spec), 'single_output') # Case 4: Multiple outputs, no 'output' or 'good' outputs = {'another': 1, 'yet_another': 2} with self.assertRaisesRegex( ValueError, "Transform .* has outputs .* but none are named 'output'"): get_main_output_key(spec, outputs, error_handling_spec) # Case 5: Empty outputs outputs = {} with self.assertRaisesRegex( ValueError, "Transform .* has outputs .* but none are named 'output'"): get_main_output_key(spec, outputs, error_handling_spec) # Case 6: More than two outputs with 'good' present outputs = {'good': 1, 'bad': 2, 'something': 3} with self.assertLogs('apache_beam.yaml.yaml_transform', level='WARNING') as al: self.assertEqual( get_main_output_key(spec, outputs, error_handling_spec), 'good') self.assertIn("Only the main output will be validated.", al.output[0]) class YamlTransformTest(unittest.TestCase): def test_init_with_string(self): provider1 = InlineProvider({"MyTransform1": lambda: beam.Map(lambda x: x)}) provider2 = InlineProvider({"MyTransform2": lambda: beam.Map(lambda x: x)}) providers_dict = {"p1": [provider1], "p2": [provider2]} spec = ''' type: chain transforms: - type: Create elements: [1,2,3] - type: LogForTesting ''' result = YamlTransform(spec, providers_dict) self.assertIn('p1', result._providers) # check for custom providers self.assertIn('p2', result._providers) # check for custom providers self.assertIn( 'LogForTesting', result._providers) # check for standard provider self.assertEqual(result._spec['type'], "composite") # preprocessed spec def test_init_with_dict(self): spec = ''' type: chain transforms: - type: Create config: elements: [1,2,3] - type: LogForTesting ''' spec = yaml.load(spec, Loader=SafeLineLoader) result = YamlTransform(spec) self.assertIn( 'LogForTesting', result._providers) # check for standard provider self.assertEqual(result._spec['type'], "composite") # preprocessed spec class ExpandPipelineTest(unittest.TestCase): def test_expand_pipeline_with_pipeline_key_only(self): spec = ''' pipeline: type: chain transforms: - type: Create config: elements: [1,2,3] - type: LogForTesting ''' with new_pipeline() as p: expand_pipeline(p, spec, validate_schema=None) def test_expand_pipeline_with_pipeline_and_option_keys(self): spec = ''' pipeline: type: chain transforms: - type: Create config: elements: [1,2,3] - type: LogForTesting options: streaming: false ''' with new_pipeline() as p: expand_pipeline(p, spec, validate_schema=None) def test_expand_pipeline_with_extra_top_level_keys(self): spec = ''' template: version: "1.0" author: "test_user" pipeline: type: chain transforms: - type: Create config: elements: [1,2,3] - type: LogForTesting other_metadata: "This is an ignored comment." ''' with new_pipeline() as p: expand_pipeline(p, spec, validate_schema=None) def test_expand_pipeline_with_incorrect_pipelines_key_fails(self): spec = ''' pipelines: type: chain transforms: - type: Create config: elements: [1,2,3] - type: LogForTesting ''' with new_pipeline() as p: with self.assertRaises(KeyError): expand_pipeline(p, spec, validate_schema=None) @unittest.skipIf(jsonschema is None, "Yaml dependencies not installed") def test_expand_pipeline_with_valid_schema(self): spec = ''' pipeline: type: chain transforms: - type: Create config: elements: [1,2,3] - type: LogForTesting ''' with new_pipeline() as p: expand_pipeline(p, spec, validate_schema='generic') @unittest.skipIf(jsonschema is None, "Yaml dependencies not installed") def test_expand_pipeline_with_invalid_schema(self): spec = ''' pipeline: type: chain transforms: - name: Create config: elements: [1,2,3] - type: LogForTesting ''' with new_pipeline() as p: with self.assertRaises(jsonschema.ValidationError): expand_pipeline(p, spec, validate_schema='generic') if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_udf_test.py0000644000175100017510000002760515206113033022043 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import os import shutil import tempfile import unittest import apache_beam as beam from apache_beam.io import localfilesystem from apache_beam.options import pipeline_options from apache_beam.options.pipeline_options import StandardOptions from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.yaml.yaml_mapping import py_value_to_js_dict from apache_beam.yaml.yaml_provider import dicts_to_rows from apache_beam.yaml.yaml_transform import YamlTransform try: import js2py except ImportError: js2py = None logging.warning('js2py is not installed; some tests will be skipped.') def as_rows(): return beam.Map( lambda named_tuple: dicts_to_rows(py_value_to_js_dict(named_tuple))) class YamlUDFMappingTest(unittest.TestCase): def __init__(self, method_name='runYamlMappingTest'): super().__init__(method_name) self.data = [ beam.Row( label='11a', conductor=11, row=beam.Row(rank=0, values=[1, 2, 3])), beam.Row( label='37a', conductor=37, row=beam.Row(rank=1, values=[4, 5, 6])), beam.Row( label='389a', conductor=389, row=beam.Row(rank=2, values=[7, 8, 9])), ] def setUp(self): self.tmpdir = tempfile.mkdtemp() self.fs = localfilesystem.LocalFileSystem(pipeline_options) def tearDown(self): shutil.rmtree(self.tmpdir) @unittest.skipIf(js2py is None, 'js2py not installed.') def test_map_to_fields_filter_inline_js(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle', yaml_experimental_features=['javascript' ])) as p: elements = p | beam.Create(self.data) result = elements | YamlTransform( ''' type: MapToFields config: language: javascript fields: label: callable: | function label_map(x) { return x.label + 'x' } conductor: callable: | function conductor_map(x) { return x.conductor + 1 } row: callable: | function row_map(x) { x.row.values.push(x.row.rank + 10) return x.row } ''') assert_that( result, equal_to([ beam.Row( label='11ax', conductor=12, row=beam.Row(rank=0, values=[1, 2, 3, 10])), beam.Row( label='37ax', conductor=38, row=beam.Row(rank=1, values=[4, 5, 6, 11])), beam.Row( label='389ax', conductor=390, row=beam.Row(rank=2, values=[7, 8, 9, 12])), ])) def test_map_to_fields_filter_inline_py(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(self.data) result = elements | YamlTransform( ''' type: MapToFields config: language: python fields: label: callable: "lambda x: x.label + 'x'" conductor: callable: "lambda x: x.conductor + 1" sum: callable: "lambda x: sum(x.row.values)" ''') assert_that( result, equal_to([ beam.Row(label='11ax', conductor=12, sum=6), beam.Row(label='37ax', conductor=38, sum=15), beam.Row(label='389ax', conductor=390, sum=24), ])) @staticmethod @unittest.skipIf( TestPipeline().get_pipeline_options().view_as(StandardOptions).runner is None, 'Do not run this test on precommit suites.') def test_map_to_fields_sql_reserved_keyword(): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: data = [ beam.Row(label="11a", rank=0), beam.Row(label="37a", rank=1), beam.Row(label="389a", rank=2), ] elements = p | beam.Create(data) result = elements | YamlTransform( ''' type: MapToFields config: language: sql append: true drop: [ rank ] fields: timestamp: "`rank`" ''') assert_that( result | as_rows(), equal_to([ beam.Row(label='11a', timestamp=0), beam.Row(label='37a', timestamp=1), beam.Row(label='389a', timestamp=2), ])) @staticmethod @unittest.skipIf( TestPipeline().get_pipeline_options().view_as(StandardOptions).runner is None, 'Do not run this test on precommit suites.') def test_map_to_fields_sql_reserved_keyword_append(): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: data = [ beam.Row(label="11a", timestamp=0), beam.Row(label="37a", timestamp=1), beam.Row(label="389a", timestamp=2), ] elements = p | beam.Create(data) result = elements | YamlTransform( ''' type: MapToFields config: language: sql append: true fields: label_copy: label ''') assert_that( result | as_rows(), equal_to([ beam.Row(label='11a', timestamp=0, label_copy="11a"), beam.Row(label='37a', timestamp=1, label_copy="37a"), beam.Row(label='389a', timestamp=2, label_copy="389a"), ])) @unittest.skipIf(js2py is None, 'js2py not installed.') def test_filter_inline_js(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle', yaml_experimental_features=['javascript' ])) as p: elements = p | beam.Create(self.data) result = elements | YamlTransform( ''' type: Filter config: language: javascript keep: callable: | function filter(x) { return x.row.rank > 0 } ''') assert_that( result | as_rows(), equal_to([ beam.Row( label='37a', conductor=37, row=beam.Row(rank=1, values=[4, 5, 6])), beam.Row( label='389a', conductor=389, row=beam.Row(rank=2, values=[7, 8, 9])), ])) def test_filter_inline_py(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(self.data) result = elements | YamlTransform( ''' type: Filter config: language: python keep: callable: "lambda x: x.row.rank > 0" ''') assert_that( result | as_rows(), equal_to([ beam.Row( label='37a', conductor=37, row=beam.Row(rank=1, values=[4, 5, 6])), beam.Row( label='389a', conductor=389, row=beam.Row(rank=2, values=[7, 8, 9])), ])) @unittest.skipIf(js2py is None, 'js2py not installed.') def test_filter_expression_js(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle', yaml_experimental_features=['javascript' ])) as p: elements = p | beam.Create(self.data) result = elements | YamlTransform( ''' type: Filter config: language: javascript keep: expression: "label.toUpperCase().indexOf('3') == -1 && row.rank < 1" ''') assert_that( result | as_rows(), equal_to([ beam.Row( label='11a', conductor=11, row=beam.Row(rank=0, values=[1, 2, 3])), ])) def test_filter_expression_py(self): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(self.data) result = elements | YamlTransform( ''' type: Filter config: language: python keep: expression: "'3' not in label" ''') assert_that( result | as_rows(), equal_to([ beam.Row( label='11a', conductor=11, row=beam.Row(rank=0, values=[1, 2, 3])), ])) @unittest.skipIf(js2py is None, 'js2py not installed.') def test_filter_inline_js_file(self): data = ''' function f(x) { return x.row.rank > 0 } function g(x) { return x.row.rank > 1 } '''.replace(' ', '') path = os.path.join(self.tmpdir, 'udf.js') self.fs.create(path).write(data.encode('utf8')) with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle', yaml_experimental_features=['javascript' ])) as p: elements = p | beam.Create(self.data) result = elements | YamlTransform( f''' type: Filter config: language: javascript keep: path: {path} name: "f" ''') assert_that( result | as_rows(), equal_to([ beam.Row( label='37a', conductor=37, row=beam.Row(rank=1, values=[4, 5, 6])), beam.Row( label='389a', conductor=389, row=beam.Row(rank=2, values=[7, 8, 9])), ])) def test_filter_inline_py_file(self): data = ''' def f(x): return x.row.rank > 0 def g(x): return x.row.rank > 1 '''.replace(' ', '') path = os.path.join(self.tmpdir, 'udf.py') self.fs.create(path).write(data.encode('utf8')) with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: elements = p | beam.Create(self.data) result = elements | YamlTransform( f''' type: Filter config: language: python keep: path: {path} name: "f" ''') assert_that( result | as_rows(), equal_to([ beam.Row( label='37a', conductor=37, row=beam.Row(rank=1, values=[4, 5, 6])), beam.Row( label='389a', conductor=389, row=beam.Row(rank=2, values=[7, 8, 9])), ])) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_utils.py0000644000175100017510000002077015206113033021362 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import json import os import re import uuid from collections.abc import Iterable from collections.abc import Mapping from typing import Any import yaml from yaml import SafeLoader class SafeLineLoader(SafeLoader): """A yaml loader that attaches line information to mappings and strings.""" class TaggedString(str): """A string class to which we can attach metadata. This is primarily used to trace a string's origin back to its place in a yaml file. """ def __reduce__(self): # Pickle as an ordinary string. return str, (str(self), ) def construct_scalar(self, node): value = super().construct_scalar(node) if isinstance(value, str): value = SafeLineLoader.TaggedString(value) value._line_ = node.start_mark.line + 1 return value def construct_mapping(self, node, deep=False): mapping = super().construct_mapping(node, deep=deep) mapping['__line__'] = node.start_mark.line + 1 mapping['__uuid__'] = self.create_uuid() return mapping @classmethod def create_uuid(cls): return str(uuid.uuid4()) @classmethod def strip_metadata(cls, spec, tagged_str=True): if isinstance(spec, Mapping): return { cls.strip_metadata(key, tagged_str): cls.strip_metadata( value, tagged_str) for (key, value) in spec.items() if key not in ('__line__', '__uuid__') } elif isinstance(spec, Iterable) and not isinstance(spec, (str, bytes)): return [cls.strip_metadata(value, tagged_str) for value in spec] elif isinstance(spec, SafeLineLoader.TaggedString) and tagged_str: return str(spec) else: return spec @staticmethod def get_line(obj): if isinstance(obj, dict): return obj.get('__line__', 'unknown') else: return getattr(obj, '_line_', 'unknown') def patch_yaml(original_str: str, updated): """Updates a yaml string to match the updated with minimal changes. This only changes the portions of original_str that differ between original_str and updated in an attempt to preserve comments and formatting. """ if not original_str and updated: return yaml.dump(updated, sort_keys=False) if original_str[-1] != '\n': # Add a trialing newline to avoid having to constantly check this edge case. # (It's also a good idea generally...) original_str += '\n' # The yaml parser returns positions in terms of line and column numbers. # Here we construct the mapping between the two. line_starts = [0] ix = original_str.find('\n') while ix != -1: line_starts.append(ix + 1) ix = original_str.find('\n', ix + 1) def pos(line_or_mark, column=0): if isinstance(line_or_mark, yaml.Mark): line = line_or_mark.line column = line_or_mark.column else: line = line_or_mark return line_starts[line] + column # Here we define a custom loader with hooks that record where each element is # found so we can swap it out appropriately. spans = {} class SafeMarkLoader(SafeLoader): pass # We create special subclass types to ensure each returned node is # a distinct object. marked_types = {} def record_yaml_scalar(constructor): def wrapper(self, node): raw_data = constructor(self, node) typ = type(raw_data) if typ not in marked_types: marked_types[typ] = type(f'Marked_{typ}', (type(raw_data), ), {}) marked_data = marked_types[typ](raw_data) spans[id(marked_data)] = node.start_mark, node.end_mark return marked_data return wrapper SafeMarkLoader.add_constructor( 'tag:yaml.org,2002:seq', record_yaml_scalar(SafeMarkLoader.construct_sequence)) SafeMarkLoader.add_constructor( 'tag:yaml.org,2002:map', record_yaml_scalar(SafeMarkLoader.construct_mapping)) for typ in ('bool', 'int', 'float', 'binary', 'timestamp', 'str'): SafeMarkLoader.add_constructor( f'tag:yaml.org,2002:{typ}', record_yaml_scalar(getattr(SafeMarkLoader, f'construct_yaml_{typ}'))) # Now load the original yaml using our special parser. original = yaml.load(original_str, Loader=SafeMarkLoader) # This (recursively) finds the portion of the original string that must # be replaced with new content. def diff(a: Any, b: Any) -> Iterable[tuple[int, int, str]]: if a == b: return elif (isinstance(a, dict) and isinstance(b, dict) and set(a.keys()) == set(b.keys()) and all(id(v) in spans for v in a.values())): for k, v in a.items(): yield from diff(v, b[k]) elif (isinstance(a, list) and isinstance(b, list) and a and b and all(id(v) in spans for v in a)): # Diff the matching entries. for va, vb in zip(a, b): yield from diff(va, vb) if len(b) < len(a): # Remove extra entries yield ( # End of last preserved element. pos(spans[id(a[len(b) - 1])][1]), # End of last original element. pos(spans[id(a[-1])][1]), '') elif len(b) > len(a): # Add extra entries list_start, list_end = spans[id(a)] start_char = original_str[pos(list_start)] if start_char == '[': for v in b[len(a):]: yield pos(list_end) - 1, pos(list_end) - 1, ', ' + json.dumps(v) else: assert start_char == '-' indent = original_str[pos(list_start.line):pos(list_start)] + '- ' content = original_str[pos(list_start):pos(list_end)].rstrip() actual_end_pos = pos(list_start) + len(content) for v in b[len(a):]: if isinstance(v, (list, dict)): v_str = ( yaml.dump(v, sort_keys=False) # Indent. .replace('\n', '\n' + ' ' * len(indent)) # Remove blank line indents. .replace(' ' * len(indent) + '\n', '\n').rstrip()) else: v_str = json.dumps(v) yield actual_end_pos, actual_end_pos, '\n' + indent + v_str else: start, end = spans[id(a)] indent = original_str[pos(start.line):pos(start)] # We strip trailing whitespace as the "end" of an element is often on # a subsequent line where the subsequent element actually starts. content = original_str[pos(start):pos(end)].rstrip() actual_end_pos = pos(start) + len(content) trailing = original_str[actual_end_pos:original_str. find('\n', actual_end_pos)] if isinstance(b, (list, dict)): if indent.strip() in ('', '-') and not trailing.strip(): # This element wholly occupies its set of lines, so it is safe to use # a multi-line yaml representation (appropriately indented). yield ( pos(start), actual_end_pos, yaml.dump(b, sort_keys=False) # Indent. .replace('\n', '\n' + ' ' * len(indent)) # Remove blank line indents. .replace(' ' * len(indent) + '\n', '\n').rstrip()) else: # Force flow style. yield ( pos(start), actual_end_pos, yaml.dump(b, default_flow_style=True, line_break=False).strip()) elif isinstance(b, str) and re.match('^[A-Za-z0-9_]+$', b): # A simple string literal. yield pos(start), actual_end_pos, b else: # A scalar. yield pos(start), actual_end_pos, json.dumps(b) # Now stick it all together. last_end = 0 content = [] for start, end, new_content in sorted(diff(original, updated)): content.append(original_str[last_end:start]) content.append(new_content) last_end = end content.append(original_str[last_end:]) return ''.join(content) def locate_data_file(relpath): return os.path.join(os.path.dirname(__file__), relpath) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/apache_beam/yaml/yaml_utils_test.py0000644000175100017510000001324315206113033022416 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging import unittest import yaml from apache_beam.yaml.yaml_utils import SafeLineLoader from apache_beam.yaml.yaml_utils import patch_yaml class SafeLineLoaderTest(unittest.TestCase): def test_get_line(self): pipeline_yaml = ''' type: composite input: elements: input transforms: - type: PyMap name: Square input: elements config: fn: "lambda x: x * x" - type: PyMap name: Cube input: elements config: fn: "lambda x: x * x * x" output: Flatten ''' spec = yaml.load(pipeline_yaml, Loader=SafeLineLoader) self.assertEqual(SafeLineLoader.get_line(spec['type']), 2) self.assertEqual(SafeLineLoader.get_line(spec['input']), 4) self.assertEqual(SafeLineLoader.get_line(spec['transforms'][0]), 6) self.assertEqual(SafeLineLoader.get_line(spec['transforms'][0]['type']), 6) self.assertEqual(SafeLineLoader.get_line(spec['transforms'][0]['name']), 7) self.assertEqual(SafeLineLoader.get_line(spec['transforms'][1]), 11) self.assertEqual(SafeLineLoader.get_line(spec['output']), 17) self.assertEqual(SafeLineLoader.get_line(spec['transforms']), "unknown") def test_strip_metadata(self): spec_yaml = ''' transforms: - type: PyMap name: Square ''' spec = yaml.load(spec_yaml, Loader=SafeLineLoader) stripped = SafeLineLoader.strip_metadata(spec['transforms']) self.assertFalse(hasattr(stripped[0], '__line__')) self.assertFalse(hasattr(stripped[0], '__uuid__')) def test_strip_metadata_nothing_to_strip(self): spec_yaml = 'prop: 123' spec = yaml.load(spec_yaml, Loader=SafeLineLoader) stripped = SafeLineLoader.strip_metadata(spec['prop']) self.assertFalse(hasattr(stripped, '__line__')) self.assertFalse(hasattr(stripped, '__uuid__')) _BASIC_YAML = ''' a: [1, 2, 3] # comment b: - 1 - 2 - 3 outer: - k1: v1 k2: v2 k3: [1, 2, {x: y}] k0: out_of_order - item ''' class PatchYamlTest(unittest.TestCase): def test_unchanged(self): obj = yaml.load(_BASIC_YAML, Loader=yaml.SafeLoader) self.assertEqual(patch_yaml(_BASIC_YAML, obj), _BASIC_YAML) def test_inline_list(self): obj = yaml.load(_BASIC_YAML, Loader=yaml.SafeLoader) obj['a'] = [100, 200] self.assertEqual( patch_yaml(_BASIC_YAML, obj), _BASIC_YAML.replace('[1, 2, 3]', '[100, 200]')) def test_bullet_list(self): obj = yaml.load(_BASIC_YAML, Loader=yaml.SafeLoader) obj['b'] = [100, 200] self.assertEqual( patch_yaml(_BASIC_YAML, obj), _BASIC_YAML.replace(''' - 1 - 2 - 3 ''', ''' - 100 - 200 ''')) def test_inline_map(self): obj = yaml.load(_BASIC_YAML, Loader=yaml.SafeLoader) obj['outer'][0]['k3'][-1] = 500 self.assertEqual( patch_yaml(_BASIC_YAML, obj), _BASIC_YAML.replace('{x: y}', '500')) def test_multiline_map(self): obj = yaml.load(_BASIC_YAML, Loader=yaml.SafeLoader) obj['outer'][0]['k4'] = 500 self.assertEqual( patch_yaml(_BASIC_YAML, obj), _BASIC_YAML.replace( ''' - k1: v1 k2: v2 k3: [1, 2, {x: y}] k0: out_of_order '''.strip(), ''' - k1: v1 k2: v2 k3: - 1 - 2 - x: y k0: out_of_order k4: 500 '''.strip())) def test_inline_list_shorten(self): obj = yaml.load(_BASIC_YAML, Loader=yaml.SafeLoader) obj['a'] = [100, 2] self.assertEqual( patch_yaml(_BASIC_YAML, obj), _BASIC_YAML.replace('[1, 2, 3]', '[100, 2]')) def test_bullet_list_shorten(self): obj = yaml.load(_BASIC_YAML, Loader=yaml.SafeLoader) obj['b'] = [100, 2] self.assertEqual( patch_yaml(_BASIC_YAML, obj), _BASIC_YAML.replace( ''' - 1 - 2 - 3 '''.strip(), ''' - 100 - 2 '''.strip())) def test_inline_list_extend(self): obj = yaml.load(_BASIC_YAML, Loader=yaml.SafeLoader) obj['a'] = [100, 2, 3, 4] self.assertEqual( patch_yaml(_BASIC_YAML, obj), _BASIC_YAML.replace('[1, 2, 3]', '[100, 2, 3, 4]')) def test_bullet_list_extend(self): obj = yaml.load(_BASIC_YAML, Loader=yaml.SafeLoader) obj['b'] = [100, 2, 3, 4] self.assertEqual( patch_yaml(_BASIC_YAML, obj), _BASIC_YAML.replace( ''' - 1 - 2 - 3 '''.strip(), ''' - 100 - 2 - 3 - 4 '''.strip())) def test_bullet_list_map_extend(self): obj = yaml.load(_BASIC_YAML, Loader=yaml.SafeLoader) obj['outer'].append({'a': 'A', 'b': 'B'}) self.assertEqual( patch_yaml(_BASIC_YAML, obj), _BASIC_YAML.replace( ''' - item '''.strip(), ''' - item - a: A b: B '''.strip())) if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.8154538 apache_beam-2.74.0/apache_beam.egg-info/0000755000175100017510000000000015206113160017351 5ustar00runnerrunner././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996270.0 apache_beam-2.74.0/apache_beam.egg-info/PKG-INFO0000644000175100017510000006037715206113156020470 0ustar00runnerrunnerMetadata-Version: 2.4 Name: apache-beam Version: 2.74.0 Summary: Apache Beam SDK for Python Home-page: https://beam.apache.org Download-URL: https://pypi.python.org/pypi/apache-beam Author: Apache Software Foundation Author-email: dev@beam.apache.org License: Apache License, Version 2.0 Keywords: apache beam Classifier: Intended Audience :: End Users/Desktop Classifier: License :: OSI Approved :: Apache Software License Classifier: Operating System :: POSIX :: Linux Classifier: Programming Language :: Python :: 3.10 Classifier: Programming Language :: Python :: 3.11 Classifier: Programming Language :: Python :: 3.12 Classifier: Programming Language :: Python :: 3.13 Classifier: Programming Language :: Python :: 3.14 Classifier: Topic :: Software Development :: Libraries Classifier: Topic :: Software Development :: Libraries :: Python Modules Requires-Python: >=3.10 Description-Content-Type: text/markdown Requires-Dist: cryptography<48.0.0,>=39.0.0 Requires-Dist: envoy-data-plane<2,>=1.0.3; python_version >= "3.11" Requires-Dist: envoy-data-plane<0.3.0; python_version < "3.11" Requires-Dist: fastavro<2,>=0.23.6 Requires-Dist: fasteners<1.0,>=0.3 Requires-Dist: grpcio!=1.48.0,!=1.59.*,!=1.60.*,!=1.61.*,!=1.62.0,!=1.62.1,!=1.66.*,!=1.67.*,!=1.68.*,!=1.69.*,!=1.70.*,<2,>=1.33.1 Requires-Dist: httplib2<0.32.0,>=0.8 Requires-Dist: jsonpickle<4.0.0,>=3.0.0 Requires-Dist: numpy<2.5.0,>=1.14.3 Requires-Dist: objsize<0.8.0,>=0.6.1 Requires-Dist: packaging>=22.0 Requires-Dist: pillow<13,>=12.1.1 Requires-Dist: pymongo<5.0.0,>=3.8.0 Requires-Dist: proto-plus<2,>=1.7.1 Requires-Dist: protobuf!=4.0.*,!=4.21.*,!=4.22.0,!=4.23.*,!=4.24.*,<7.0.0.dev0,>=3.20.3 Requires-Dist: python-dateutil<3,>=2.8.0 Requires-Dist: pytz>=2018.3 Requires-Dist: requests<3.0.0,>=2.32.4 Requires-Dist: sortedcontainers>=2.4.0 Requires-Dist: typing-extensions>=3.7.0 Requires-Dist: zstandard<1,>=0.18.0 Requires-Dist: pyyaml<7.0.0,>=3.12 Requires-Dist: beartype<0.23.0,>=0.21.0 Requires-Dist: pyarrow<24.0.0,>=6.0.1 Requires-Dist: pyarrow-hotfix<1 Provides-Extra: dev Requires-Dist: isort==7.0.0; extra == "dev" Requires-Dist: pyrefly==0.54.0; extra == "dev" Requires-Dist: ruff==0.15.7; extra == "dev" Requires-Dist: yapf==0.43.0; extra == "dev" Provides-Extra: dill Requires-Dist: dill<0.3.2,>=0.3.1.1; extra == "dill" Provides-Extra: docs Requires-Dist: jinja2<3.2,>=3.0; extra == "docs" Requires-Dist: Sphinx<8.0,>=7.0.0; extra == "docs" Requires-Dist: docstring-parser<1.0,>=0.15; extra == "docs" Requires-Dist: docutils>=0.18.1; extra == "docs" Requires-Dist: markdown; extra == "docs" Requires-Dist: pandas<2.3.0; extra == "docs" Requires-Dist: openai; extra == "docs" Requires-Dist: virtualenv-clone<1.0,>=0.5; extra == "docs" Provides-Extra: test Requires-Dist: cloud-sql-python-connector[pg8000]<2.0.0,>=1.0.0; extra == "test" Requires-Dist: docstring-parser<1.0,>=0.15; extra == "test" Requires-Dist: freezegun>=0.3.12; extra == "test" Requires-Dist: jinja2<3.2,>=3.0; extra == "test" Requires-Dist: joblib>=1.0.1; extra == "test" Requires-Dist: mock<6.0.0,>=1.0.1; extra == "test" Requires-Dist: pandas<2.3.0; extra == "test" Requires-Dist: parameterized<0.10.0,>=0.7.1; extra == "test" Requires-Dist: pydot<2,>=1.2.0; extra == "test" Requires-Dist: pyhamcrest!=1.10.0,<3.0.0,>=1.9; extra == "test" Requires-Dist: requests_mock<2.0,>=1.7; extra == "test" Requires-Dist: tenacity<10,>=8.0.0; extra == "test" Requires-Dist: pytest<10.0,>=7.1.2; extra == "test" Requires-Dist: pytest-xdist<4,>=2.5.0; extra == "test" Requires-Dist: pytest-timeout<3,>=2.1.0; extra == "test" Requires-Dist: scikit-learn<1.8.0,>=0.20.0; extra == "test" Requires-Dist: sqlalchemy<3.0,>=1.3; extra == "test" Requires-Dist: psycopg2-binary<3.0,>=2.8.5; extra == "test" Requires-Dist: testcontainers[kafka,milvus,mysql]<5.0.0,>=4.0.0; extra == "test" Requires-Dist: cryptography>=41.0.2; extra == "test" Requires-Dist: hypothesis<6.148.4,>5.0.0; extra == "test" Requires-Dist: virtualenv-clone<1.0,>=0.5; extra == "test" Requires-Dist: python-tds>=1.16.1; extra == "test" Requires-Dist: sqlalchemy-pytds>=1.0.2; extra == "test" Requires-Dist: pg8000>=1.31.5; extra == "test" Requires-Dist: PyMySQL>=1.1.0; extra == "test" Requires-Dist: oracledb>=3.1.1; extra == "test" Provides-Extra: gcp Requires-Dist: cachetools<7,>=3.1.0; extra == "gcp" Requires-Dist: google-api-core<3,>=2.0.0; extra == "gcp" Requires-Dist: google-apitools<0.5.32,>=0.5.31; python_version < "3.13" and extra == "gcp" Requires-Dist: google-apitools>=0.5.35; python_version >= "3.13" and extra == "gcp" Requires-Dist: google-auth<3,>=1.18.0; extra == "gcp" Requires-Dist: google-auth-httplib2<0.3.0,>=0.1.0; extra == "gcp" Requires-Dist: google-cloud-datastore<3,>=2.0.0; extra == "gcp" Requires-Dist: google-cloud-pubsub<3,>=2.1.0; extra == "gcp" Requires-Dist: google-cloud-storage<4,>=2.18.2; extra == "gcp" Requires-Dist: google-cloud-bigquery<4,>=2.0.0; extra == "gcp" Requires-Dist: google-cloud-bigquery-storage<3,>=2.6.3; extra == "gcp" Requires-Dist: google-cloud-core<3,>=2.0.0; extra == "gcp" Requires-Dist: google-cloud-bigtable<3,>=2.19.0; extra == "gcp" Requires-Dist: google-cloud-build<4,>=3.35.0; extra == "gcp" Requires-Dist: google-cloud-spanner<4,>=3.0.0; extra == "gcp" Requires-Dist: google-cloud-dlp<4,>=3.0.0; extra == "gcp" Requires-Dist: google-cloud-kms<4,>=3.0.0; extra == "gcp" Requires-Dist: google-cloud-language<3,>=2.0; extra == "gcp" Requires-Dist: google-cloud-secret-manager<3,>=2.0; extra == "gcp" Requires-Dist: google-cloud-videointelligence<3,>=2.0; extra == "gcp" Requires-Dist: google-cloud-vision<4,>=2; extra == "gcp" Requires-Dist: google-cloud-recommendations-ai<0.11.0,>=0.1.0; extra == "gcp" Requires-Dist: google-cloud-aiplatform<2.0,>=1.26.0; extra == "gcp" Requires-Dist: cloud-sql-python-connector<2.0.0,>=1.18.2; extra == "gcp" Requires-Dist: python-tds>=1.16.1; extra == "gcp" Requires-Dist: pg8000>=1.31.5; extra == "gcp" Requires-Dist: PyMySQL>=1.1.0; extra == "gcp" Requires-Dist: keyrings.google-artifactregistry-auth; extra == "gcp" Requires-Dist: orjson<4,>=3.9.7; extra == "gcp" Requires-Dist: regex>=2020.6.8; extra == "gcp" Provides-Extra: interactive Requires-Dist: facets-overview<2,>=1.1.0; extra == "interactive" Requires-Dist: google-cloud-dataproc<6,>=5.0.0; extra == "interactive" Requires-Dist: ipython<9,>=7; extra == "interactive" Requires-Dist: ipykernel<7,>=6; extra == "interactive" Requires-Dist: ipywidgets<9,>=8; extra == "interactive" Requires-Dist: jupyter-client!=6.1.13,<8.2.1,>=6.1.11; extra == "interactive" Requires-Dist: pydot<2,>=1.2.0; extra == "interactive" Requires-Dist: timeloop<2,>=1.0.2; extra == "interactive" Requires-Dist: nbformat<6,>=5.0.5; extra == "interactive" Requires-Dist: nbconvert<8,>=6.2.0; extra == "interactive" Requires-Dist: pandas!=1.5.0,!=1.5.1,<2.3,>=1.4.3; extra == "interactive" Provides-Extra: interactive-test Requires-Dist: needle<1,>=0.5.0; extra == "interactive-test" Requires-Dist: chromedriver-binary<118,>=117; extra == "interactive-test" Requires-Dist: pillow<10,>=7.1.1; extra == "interactive-test" Requires-Dist: urllib3<2,>=1.21.1; extra == "interactive-test" Provides-Extra: ml-test Requires-Dist: datatable; extra == "ml-test" Requires-Dist: dill; extra == "ml-test" Requires-Dist: tensorflow_transform<1.15.0,>=1.14.0; extra == "ml-test" Requires-Dist: numpy<2; extra == "ml-test" Requires-Dist: tf2onnx<1.17,>=1.16.1; extra == "ml-test" Requires-Dist: embeddings>=0.0.4; extra == "ml-test" Requires-Dist: onnxruntime; extra == "ml-test" Requires-Dist: onnx<2,>=1.14.1; extra == "ml-test" Requires-Dist: langchain; extra == "ml-test" Requires-Dist: sentence-transformers>=2.2.2; extra == "ml-test" Requires-Dist: skl2onnx; extra == "ml-test" Requires-Dist: pyod>=0.7.6; extra == "ml-test" Requires-Dist: tensorflow; extra == "ml-test" Requires-Dist: absl-py>=0.12.0; extra == "ml-test" Requires-Dist: tensorflow-hub; extra == "ml-test" Requires-Dist: tokenizers<0.23; extra == "ml-test" Requires-Dist: torch<2.12.0; extra == "ml-test" Requires-Dist: transformers<4.56.0,>=4.28.0; extra == "ml-test" Requires-Dist: tokenizers<0.22.0,>=0.13.3; extra == "ml-test" Provides-Extra: p310-ml-test Requires-Dist: datatable; extra == "p310-ml-test" Requires-Dist: embeddings>=0.0.4; extra == "p310-ml-test" Requires-Dist: onnxruntime; extra == "p310-ml-test" Requires-Dist: onnx<2,>=1.14.1; extra == "p310-ml-test" Requires-Dist: langchain; extra == "p310-ml-test" Requires-Dist: sentence-transformers>=2.2.2; extra == "p310-ml-test" Requires-Dist: skl2onnx; extra == "p310-ml-test" Requires-Dist: pyod>=0.7.6; extra == "p310-ml-test" Requires-Dist: tensorflow; extra == "p310-ml-test" Requires-Dist: absl-py>=0.12.0; extra == "p310-ml-test" Requires-Dist: tensorflow-hub; extra == "p310-ml-test" Requires-Dist: tokenizers<0.23; extra == "p310-ml-test" Requires-Dist: torch<2.12.0; extra == "p310-ml-test" Requires-Dist: transformers<4.56.0,>=4.28.0; extra == "p310-ml-test" Requires-Dist: tokenizers<0.22.0,>=0.13.3; extra == "p310-ml-test" Requires-Dist: google-adk==1.28.1; extra == "p310-ml-test" Requires-Dist: proto-plus<2,>=1.26.1; extra == "p310-ml-test" Requires-Dist: opentelemetry-api==1.37.0; extra == "p310-ml-test" Requires-Dist: opentelemetry-sdk==1.37.0; extra == "p310-ml-test" Requires-Dist: opentelemetry-exporter-otlp-proto-http==1.37.0; extra == "p310-ml-test" Requires-Dist: tf2onnx<1.18,>=1.17.0; extra == "p310-ml-test" Provides-Extra: p312-ml-test Requires-Dist: datatable; extra == "p312-ml-test" Requires-Dist: embeddings>=0.0.4; extra == "p312-ml-test" Requires-Dist: onnxruntime; extra == "p312-ml-test" Requires-Dist: onnx<2,>=1.14.1; extra == "p312-ml-test" Requires-Dist: langchain; extra == "p312-ml-test" Requires-Dist: sentence-transformers>=2.2.2; extra == "p312-ml-test" Requires-Dist: skl2onnx; extra == "p312-ml-test" Requires-Dist: pyod>=0.7.6; extra == "p312-ml-test" Requires-Dist: tensorflow; extra == "p312-ml-test" Requires-Dist: absl-py>=0.12.0; extra == "p312-ml-test" Requires-Dist: tensorflow-hub; extra == "p312-ml-test" Requires-Dist: tokenizers<0.23; extra == "p312-ml-test" Requires-Dist: torch<2.12.0; extra == "p312-ml-test" Requires-Dist: transformers<4.56.0,>=4.28.0; extra == "p312-ml-test" Requires-Dist: tokenizers<0.22.0,>=0.13.3; extra == "p312-ml-test" Requires-Dist: google-adk==1.28.1; extra == "p312-ml-test" Requires-Dist: proto-plus<2,>=1.26.1; extra == "p312-ml-test" Requires-Dist: opentelemetry-api==1.37.0; extra == "p312-ml-test" Requires-Dist: opentelemetry-sdk==1.37.0; extra == "p312-ml-test" Requires-Dist: opentelemetry-exporter-otlp-proto-http==1.37.0; extra == "p312-ml-test" Requires-Dist: tf2onnx<1.18,>=1.17.0; extra == "p312-ml-test" Provides-Extra: p313-ml-test Requires-Dist: embeddings>=0.0.4; extra == "p313-ml-test" Requires-Dist: onnxruntime; extra == "p313-ml-test" Requires-Dist: onnx<2,>=1.14.1; extra == "p313-ml-test" Requires-Dist: langchain; extra == "p313-ml-test" Requires-Dist: sentence-transformers>=2.2.2; extra == "p313-ml-test" Requires-Dist: skl2onnx; extra == "p313-ml-test" Requires-Dist: pyod>=0.7.6; extra == "p313-ml-test" Requires-Dist: tensorflow; extra == "p313-ml-test" Requires-Dist: absl-py>=0.12.0; extra == "p313-ml-test" Requires-Dist: tensorflow-hub; extra == "p313-ml-test" Requires-Dist: tokenizers<0.23; extra == "p313-ml-test" Requires-Dist: torch<2.12.0; extra == "p313-ml-test" Requires-Dist: transformers<4.56.0,>=4.28.0; extra == "p313-ml-test" Requires-Dist: tokenizers<0.22.0,>=0.13.3; extra == "p313-ml-test" Requires-Dist: google-adk==1.28.1; extra == "p313-ml-test" Requires-Dist: proto-plus<2,>=1.26.1; extra == "p313-ml-test" Requires-Dist: opentelemetry-api==1.37.0; extra == "p313-ml-test" Requires-Dist: opentelemetry-sdk==1.37.0; extra == "p313-ml-test" Requires-Dist: opentelemetry-exporter-otlp-proto-http==1.37.0; extra == "p313-ml-test" Requires-Dist: tf2onnx<1.18,>=1.17.0; extra == "p313-ml-test" Requires-Dist: pymilvus<3.0.0,>=2.5.10; extra == "p313-ml-test" Provides-Extra: aws Requires-Dist: boto3<2,>=1.9; extra == "aws" Provides-Extra: azure Requires-Dist: azure-storage-blob<13,>=12.3.2; extra == "azure" Requires-Dist: azure-core<2,>=1.7.0; extra == "azure" Requires-Dist: azure-identity<2,>=1.12.0; extra == "azure" Provides-Extra: dataframe Requires-Dist: pandas!=1.5.0,!=1.5.1,<2.3,>=1.4.3; extra == "dataframe" Provides-Extra: dask Requires-Dist: distributed>=2024.4.2; extra == "dask" Requires-Dist: dask>=2024.4.2; extra == "dask" Provides-Extra: hadoop Requires-Dist: hdfs<3.0.0,>=2.1.0; extra == "hadoop" Provides-Extra: yaml Requires-Dist: docstring-parser<1.0,>=0.15; extra == "yaml" Requires-Dist: jinja2<3.2,>=3.0; extra == "yaml" Requires-Dist: virtualenv-clone<1.0,>=0.5; extra == "yaml" Requires-Dist: js2py<1,>=0.74; python_version < "3.12" and extra == "yaml" Requires-Dist: jsonschema<5.0.0,>=4.0.0; extra == "yaml" Requires-Dist: pandas!=1.5.0,!=1.5.1,<2.3,>=1.4.3; extra == "yaml" Provides-Extra: torch Requires-Dist: torch<2.8.0,>=1.9.0; extra == "torch" Provides-Extra: tensorflow Requires-Dist: tensorflow<2.21,>=2.12rc1; extra == "tensorflow" Requires-Dist: absl-py>=0.12.0; extra == "tensorflow" Provides-Extra: transformers Requires-Dist: transformers<4.56.0,>=4.28.0; extra == "transformers" Requires-Dist: tensorflow>=2.12.0; extra == "transformers" Requires-Dist: torch>=1.9.0; extra == "transformers" Provides-Extra: ml-cpu Requires-Dist: tensorflow>=2.12.0; extra == "ml-cpu" Requires-Dist: torch==2.8.0+cpu; extra == "ml-cpu" Requires-Dist: transformers<4.56.0,>=4.28.0; extra == "ml-cpu" Requires-Dist: absl-py>=0.12.0; extra == "ml-cpu" Provides-Extra: redis Requires-Dist: redis<6,>=5.0.0; extra == "redis" Provides-Extra: tft Requires-Dist: tensorflow_transform<1.15.0,>=1.14.0; extra == "tft" Requires-Dist: numpy<2; extra == "tft" Requires-Dist: dill; extra == "tft" Provides-Extra: tfrecord Requires-Dist: crcmod<2.0,>=1.7; extra == "tfrecord" Provides-Extra: onnx Requires-Dist: onnxruntime==1.13.1; extra == "onnx" Requires-Dist: torch==1.13.1; extra == "onnx" Requires-Dist: tensorflow==2.11.0; extra == "onnx" Requires-Dist: tf2onnx==1.13.0; extra == "onnx" Requires-Dist: skl2onnx==1.13; extra == "onnx" Requires-Dist: transformers==4.25.1; extra == "onnx" Requires-Dist: absl-py>=0.12.0; extra == "onnx" Provides-Extra: xgboost Requires-Dist: xgboost<2.1.3,>=1.6.0; extra == "xgboost" Requires-Dist: datatable==1.0.0; extra == "xgboost" Provides-Extra: tensorflow-hub Requires-Dist: tensorflow-hub<0.16.0,>=0.14.0; extra == "tensorflow-hub" Provides-Extra: milvus Requires-Dist: pymilvus<3.0.0,>=2.5.10; extra == "milvus" Provides-Extra: vllm Requires-Dist: openai==1.107.1; extra == "vllm" Requires-Dist: vllm==0.10.1.1; extra == "vllm" Requires-Dist: triton==3.3.1; extra == "vllm" Dynamic: author Dynamic: author-email Dynamic: classifier Dynamic: description Dynamic: description-content-type Dynamic: download-url Dynamic: home-page Dynamic: keywords Dynamic: license Dynamic: provides-extra Dynamic: requires-dist Dynamic: requires-python Dynamic: summary # Apache Beam [Apache Beam](http://beam.apache.org/) is a unified model for defining both batch and streaming data-parallel processing pipelines, as well as a set of language-specific SDKs for constructing pipelines and Runners for executing them on distributed processing backends, including [Apache Flink](http://flink.apache.org/), [Apache Spark](http://spark.apache.org/), [Google Cloud Dataflow](http://cloud.google.com/dataflow/), and [Hazelcast Jet](https://jet.hazelcast.org/). ## Overview Beam provides a general approach to expressing [embarrassingly parallel](https://en.wikipedia.org/wiki/Embarrassingly_parallel) data processing pipelines and supports three categories of users, each of which have relatively disparate backgrounds and needs. 1. _End Users_: Writing pipelines with an existing SDK, running it on an existing runner. These users want to focus on writing their application logic and have everything else just work. 2. _SDK Writers_: Developing a Beam SDK targeted at a specific user community (Java, Python, Scala, Go, R, graphical, etc). These users are language geeks and would prefer to be shielded from all the details of various runners and their implementations. 3. _Runner Writers_: Have an execution environment for distributed processing and would like to support programs written against the Beam Model. Would prefer to be shielded from details of multiple SDKs. ### The Beam Model The model behind Beam evolved from several internal Google data processing projects, including [MapReduce](http://research.google.com/archive/mapreduce.html), [FlumeJava](http://research.google.com/pubs/pub35650.html), and [Millwheel](http://research.google.com/pubs/pub41378.html). This model was originally known as the “[Dataflow Model](http://www.vldb.org/pvldb/vol8/p1792-Akidau.pdf)”. To learn more about the Beam Model (though still under the original name of Dataflow), see the World Beyond Batch: [Streaming 101](https://www.oreilly.com/ideas/the-world-beyond-batch-streaming-101) and [Streaming 102](https://www.oreilly.com/ideas/the-world-beyond-batch-streaming-102) posts on O’Reilly’s Radar site, and the [VLDB 2015 paper](http://www.vldb.org/pvldb/vol8/p1792-Akidau.pdf). The key concepts in the Beam programming model are: * `PCollection`: represents a collection of data, which could be bounded or unbounded in size. * `PTransform`: represents a computation that transforms input PCollections into output PCollections. * `Pipeline`: manages a directed acyclic graph of PTransforms and PCollections that is ready for execution. * `PipelineRunner`: specifies where and how the pipeline should execute. ### Runners Beam supports executing programs on multiple distributed processing backends through PipelineRunners. Currently, the following PipelineRunners are available: - The `DirectRunner` runs the pipeline on your local machine. - The `PrismRunner` runs the pipeline on your local machine using Beam Portability. - The `DataflowRunner` submits the pipeline to the [Google Cloud Dataflow](http://cloud.google.com/dataflow/). - The `FlinkRunner` runs the pipeline on an Apache Flink cluster. The code has been donated from [dataArtisans/flink-dataflow](https://github.com/dataArtisans/flink-dataflow) and is now part of Beam. - The `SparkRunner` runs the pipeline on an Apache Spark cluster. - The `JetRunner` runs the pipeline on a Hazelcast Jet cluster. The code has been donated from [hazelcast/hazelcast-jet](https://github.com/hazelcast/hazelcast-jet) and is now part of Beam. - The `Twister2Runner` runs the pipeline on a Twister2 cluster. The code has been donated from [DSC-SPIDAL/twister2](https://github.com/DSC-SPIDAL/twister2) and is now part of Beam. Have ideas for new Runners? See the [runner-ideas label](https://github.com/apache/beam/issues?q=is%3Aopen+is%3Aissue+label%3Arunner-ideas). ## Get started with the Python SDK Get started with the [Beam Python SDK quickstart](/get-started/quickstart-py) to set up your Python development environment, get the Beam SDK for Python, and run an example pipeline. Then, read through the [Beam programming guide](/documentation/programming-guide) to learn the basic concepts that apply to all SDKs in Beam. The [Python Tips document](../../contributor-docs/python-tips.md) is also a useful resource for setting up a development environment and performing common processes. See the [Python API reference](https://beam.apache.org/releases/pydoc/) for more information on individual APIs. ## Python streaming pipelines Python [streaming pipeline execution](https://beam.apache.org/documentation/sdks/python-streaming) is available (with some [limitations](https://beam.apache.org//documentation/sdks/python-streaming/#unsupported-features)) starting with Beam SDK version 2.5.0. ## Python type safety Python is a dynamically-typed language with no static type checking. The Beam SDK for Python uses type hints during pipeline construction and runtime to try to emulate the correctness guarantees achieved by true static typing. [Ensuring Python Type Safety](https://beam.apache.org/documentation/sdks/python-type-safety) walks through how to use type hints, which help you to catch potential bugs up front with the [Direct Runner](https://beam.apache.org//documentation/runners/direct/). ## Managing Python pipeline dependencies When you run your pipeline locally, the packages that your pipeline depends on are available because they are installed on your local machine. However, when you want to run your pipeline remotely, you must make sure these dependencies are available on the remote machines. [Managing Python Pipeline Dependencies](https://beam.apache.org/documentation/sdks/python-pipeline-dependencies) shows you how to make your dependencies available to the remote workers. ## Developing new I/O connectors for Python The Beam SDK for Python provides an extensible API that you can use to create new I/O connectors. See the [Developing I/O connectors overview](https://beam.apache.org/documentation/io/developing-io-overview) for information about developing new I/O connectors and links to language-specific implementation guidance. ## Making machine learning inferences with Python To integrate machine learning models into your pipelines for making inferences, use the RunInference API for PyTorch and Scikit-learn models. If you are using TensorFlow models, you can make use of the [library from `tfx_bsl`](https://github.com/tensorflow/tfx-bsl/tree/master/tfx_bsl/beam). You can create multiple types of transforms using the RunInference API: the API takes multiple types of setup parameters from model handlers, and the parameter type determines the model implementation. For more information, see [About Beam ML](https://beam.apache.org/documentation/ml/about-ml). [TensorFlow Extended (TFX)](https://www.tensorflow.org/tfx) is an end-to-end platform for deploying production ML pipelines. TFX is integrated with Beam. For more information, see [TFX user guide](https://www.tensorflow.org/tfx/guide). ## Python multi-language pipelines quickstart Apache Beam lets you combine transforms written in any supported SDK language and use them in one multi-language pipeline. To learn how to create a multi-language pipeline using the Python SDK, see the [Python multi-language pipelines quickstart](https://beam.apache.org/documentation/sdks/python-multi-language-pipelines). ## Unrecoverable Errors in Beam Python Some common errors can occur during worker start-up and prevent jobs from starting. To learn about these errors and how to troubleshoot them in the Python SDK, see [Unrecoverable Errors in Beam Python](https://beam.apache.org/documentation/sdks/python-unrecoverable-errors). ## 📚 Learn More Here are some resources actively maintained by the Beam community to help you get started:
Resource Details
Apache Beam Website Our website discussing the project, and it's specifics.
Python Quickstart A guide to getting started with the Python SDK.
Tour of Beam A comprehensive, interactive learning experience covering Beam concepts in depth.
Beam Quest A certification granted by Google Cloud, certifying proficiency in Beam.
Community Metrics Beam's Git Community Metrics.
## Contribution Instructions for building and testing Beam itself are in the [contribution guide](https://beam.apache.org/contribute). ## Contact Us To get involved with Apache Beam: * [Subscribe to](https://beam.apache.org/community/contact-us/#:~:text=Subscribe%20and%20Unsubscribe) or e-mail the [user@beam.apache.org](http://mail-archives.apache.org/mod_mbox/beam-user/) list. * [Subscribe to](https://beam.apache.org/community/contact-us/#:~:text=Subscribe%20and%20Unsubscribe) or e-mail the [dev@beam.apache.org](http://mail-archives.apache.org/mod_mbox/beam-dev/) list. * [Join ASF Slack](https://s.apache.org/slack-invite) on [#beam channel](https://s.apache.org/beam-slack-channel) * [Report an issue](https://github.com/apache/beam/issues/new/choose). ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996270.0 apache_beam-2.74.0/apache_beam.egg-info/SOURCES.txt0000644000175100017510000021305515206113156021250 0ustar00runnerrunnerMANIFEST.in README.md gen_protos.py gen_xlang_wrappers.py pyproject.toml setup.cfg setup.py apache_beam/__init__.py apache_beam/error.py apache_beam/pipeline.py apache_beam/pipeline_test.py apache_beam/programming_guide_test.py apache_beam/pvalue.py apache_beam/pvalue_test.py apache_beam/py.typed apache_beam/version.py apache_beam.egg-info/PKG-INFO apache_beam.egg-info/SOURCES.txt apache_beam.egg-info/dependency_links.txt apache_beam.egg-info/not-zip-safe apache_beam.egg-info/requires.txt apache_beam.egg-info/top_level.txt apache_beam/coders/__init__.py apache_beam/coders/avro_record.py apache_beam/coders/coder_impl.pxd apache_beam/coders/coder_impl.py apache_beam/coders/coder_impl_row_encoders.pyx apache_beam/coders/coders.py apache_beam/coders/coders_property_based_test.py apache_beam/coders/coders_test.py apache_beam/coders/coders_test_common.py apache_beam/coders/fast_coders_test.py apache_beam/coders/observable.py apache_beam/coders/observable_test.py apache_beam/coders/proto2_coder_test_messages_pb2.py apache_beam/coders/row_coder.py apache_beam/coders/row_coder_test.py apache_beam/coders/slow_coders_test.py apache_beam/coders/slow_stream.py apache_beam/coders/standard_coders_test.py apache_beam/coders/stream.pxd apache_beam/coders/stream.pyx apache_beam/coders/stream_test.py apache_beam/coders/typecoders.py apache_beam/coders/typecoders_test.py apache_beam/dataframe/__init__.py apache_beam/dataframe/convert.py apache_beam/dataframe/convert_test.py apache_beam/dataframe/doctests.py apache_beam/dataframe/doctests_test.py apache_beam/dataframe/expressions.py apache_beam/dataframe/expressions_test.py apache_beam/dataframe/frame_base.py apache_beam/dataframe/frame_base_test.py apache_beam/dataframe/frames.py apache_beam/dataframe/frames_test.py apache_beam/dataframe/io.py apache_beam/dataframe/io_it_test.py apache_beam/dataframe/io_test.py apache_beam/dataframe/pandas_docs_test.py apache_beam/dataframe/pandas_doctests_test.py apache_beam/dataframe/pandas_top_level_functions.py apache_beam/dataframe/partitionings.py apache_beam/dataframe/partitionings_test.py apache_beam/dataframe/schemas.py apache_beam/dataframe/schemas_test.py apache_beam/dataframe/transforms.py apache_beam/dataframe/transforms_test.py apache_beam/examples/__init__.py apache_beam/examples/avro_nyc_trips.py apache_beam/examples/avro_nyc_trips_it_test.py apache_beam/examples/avro_nyc_trips_test.py apache_beam/examples/california_housing_clustering.py apache_beam/examples/fastavro_it_test.py apache_beam/examples/matrix_power.py apache_beam/examples/matrix_power_test.py apache_beam/examples/online_clustering.py apache_beam/examples/per_entity_training.py apache_beam/examples/rate_limiter_simple.py apache_beam/examples/sql_taxi.py apache_beam/examples/streaming_wordcount.py apache_beam/examples/streaming_wordcount_debugging.py apache_beam/examples/streaming_wordcount_debugging_it_test.py apache_beam/examples/streaming_wordcount_debugging_test.py apache_beam/examples/streaming_wordcount_it_test.py apache_beam/examples/windowed_wordcount.py apache_beam/examples/wordcount.py apache_beam/examples/wordcount_dataframe.py apache_beam/examples/wordcount_debugging.py apache_beam/examples/wordcount_debugging_test.py apache_beam/examples/wordcount_it_test.py apache_beam/examples/wordcount_minimal.py apache_beam/examples/wordcount_minimal_test.py apache_beam/examples/wordcount_test.py apache_beam/examples/wordcount_with_metrics.py apache_beam/examples/wordcount_xlang.py apache_beam/examples/wordcount_xlang_sql.py apache_beam/examples/complete/__init__.py apache_beam/examples/complete/autocomplete.py apache_beam/examples/complete/autocomplete_it_test.py apache_beam/examples/complete/autocomplete_test.py apache_beam/examples/complete/distribopt.py apache_beam/examples/complete/distribopt_test.py apache_beam/examples/complete/estimate_pi.py apache_beam/examples/complete/estimate_pi_it_test.py apache_beam/examples/complete/estimate_pi_test.py apache_beam/examples/complete/tfidf.py apache_beam/examples/complete/tfidf_it_test.py apache_beam/examples/complete/tfidf_test.py apache_beam/examples/complete/top_wikipedia_sessions.py apache_beam/examples/complete/top_wikipedia_sessions_it_test.py apache_beam/examples/complete/top_wikipedia_sessions_test.py apache_beam/examples/complete/game/__init__.py apache_beam/examples/complete/game/game_stats.py apache_beam/examples/complete/game/game_stats_it_test.py apache_beam/examples/complete/game/game_stats_test.py apache_beam/examples/complete/game/hourly_team_score.py apache_beam/examples/complete/game/hourly_team_score_it_test.py apache_beam/examples/complete/game/hourly_team_score_test.py apache_beam/examples/complete/game/leader_board.py apache_beam/examples/complete/game/leader_board_it_test.py apache_beam/examples/complete/game/leader_board_test.py apache_beam/examples/complete/game/user_score.py apache_beam/examples/complete/game/user_score_it_test.py apache_beam/examples/complete/game/user_score_test.py apache_beam/examples/complete/juliaset/__init__.py apache_beam/examples/complete/juliaset/juliaset_main.py apache_beam/examples/complete/juliaset/setup.py apache_beam/examples/complete/juliaset/juliaset/__init__.py apache_beam/examples/complete/juliaset/juliaset/juliaset.py apache_beam/examples/complete/juliaset/juliaset/juliaset_test.py apache_beam/examples/complete/juliaset/juliaset/juliaset_test_it.py apache_beam/examples/cookbook/__init__.py apache_beam/examples/cookbook/bigquery_schema.py apache_beam/examples/cookbook/bigquery_side_input.py apache_beam/examples/cookbook/bigquery_side_input_it_test.py apache_beam/examples/cookbook/bigquery_side_input_test.py apache_beam/examples/cookbook/bigquery_tornadoes.py apache_beam/examples/cookbook/bigquery_tornadoes_it_test.py apache_beam/examples/cookbook/bigquery_tornadoes_test.py apache_beam/examples/cookbook/bigtableio_it_test.py apache_beam/examples/cookbook/coders.py apache_beam/examples/cookbook/coders_it_test.py apache_beam/examples/cookbook/coders_test.py apache_beam/examples/cookbook/combiners_test.py apache_beam/examples/cookbook/custom_ptransform.py apache_beam/examples/cookbook/custom_ptransform_it_test.py apache_beam/examples/cookbook/custom_ptransform_test.py apache_beam/examples/cookbook/datastore_wordcount.py apache_beam/examples/cookbook/datastore_wordcount_it_test.py apache_beam/examples/cookbook/filters.py apache_beam/examples/cookbook/filters_test.py apache_beam/examples/cookbook/group_with_coder.py apache_beam/examples/cookbook/group_with_coder_test.py apache_beam/examples/cookbook/mergecontacts.py apache_beam/examples/cookbook/mergecontacts_test.py apache_beam/examples/cookbook/multiple_output_pardo.py apache_beam/examples/cookbook/multiple_output_pardo_test.py apache_beam/examples/cookbook/ordered_window_elements/__init__.py apache_beam/examples/cookbook/ordered_window_elements/batch.py apache_beam/examples/cookbook/ordered_window_elements/batch_test.py apache_beam/examples/cookbook/ordered_window_elements/streaming.py apache_beam/examples/cookbook/ordered_window_elements/streaming_test.py apache_beam/examples/dataframe/__init__.py apache_beam/examples/dataframe/flight_delays.py apache_beam/examples/dataframe/flight_delays_it_test.py apache_beam/examples/dataframe/taxiride.py apache_beam/examples/dataframe/taxiride_it_test.py apache_beam/examples/dataframe/taxiride_test.py apache_beam/examples/dataframe/wordcount.py apache_beam/examples/dataframe/wordcount_test.py apache_beam/examples/flink/__init__.py apache_beam/examples/flink/flink_streaming_impulse.py apache_beam/examples/inference/__init__.py apache_beam/examples/inference/gemini_image_generation.py apache_beam/examples/inference/gemini_text_classification.py apache_beam/examples/inference/huggingface_language_modeling.py apache_beam/examples/inference/huggingface_question_answering.py apache_beam/examples/inference/milk_quality_prediction_windowing.py apache_beam/examples/inference/onnx_sentiment_classification.py apache_beam/examples/inference/pytorch_image_classification.py apache_beam/examples/inference/pytorch_image_classification_with_side_inputs.py apache_beam/examples/inference/pytorch_image_segmentation.py apache_beam/examples/inference/pytorch_language_modeling.py apache_beam/examples/inference/pytorch_model_per_key_image_segmentation.py apache_beam/examples/inference/pytorch_sentiment.py apache_beam/examples/inference/rate_limiter_vertex_ai.py apache_beam/examples/inference/run_inference_side_inputs.py apache_beam/examples/inference/sklearn_japanese_housing_regression.py apache_beam/examples/inference/sklearn_mnist_classification.py apache_beam/examples/inference/table_row_batch_example.py apache_beam/examples/inference/table_row_inference.py apache_beam/examples/inference/table_row_inference_test.py apache_beam/examples/inference/table_row_inference_utils.py apache_beam/examples/inference/tensorflow_imagenet_segmentation.py apache_beam/examples/inference/tensorflow_mnist_classification.py apache_beam/examples/inference/tensorflow_mnist_with_weights.py apache_beam/examples/inference/tensorrt_object_detection.py apache_beam/examples/inference/tensorrt_text_classification.py apache_beam/examples/inference/vertex_ai_image_classification.py apache_beam/examples/inference/vertex_ai_llm_text_classification.py apache_beam/examples/inference/vllm_gemma_batch.py apache_beam/examples/inference/vllm_text_completion.py apache_beam/examples/inference/xgboost_iris_classification.py apache_beam/examples/inference/runinference_metrics/__init__.py apache_beam/examples/inference/runinference_metrics/config.py apache_beam/examples/inference/runinference_metrics/main.py apache_beam/examples/inference/runinference_metrics/setup.py apache_beam/examples/inference/runinference_metrics/pipeline/__init__.py apache_beam/examples/inference/runinference_metrics/pipeline/options.py apache_beam/examples/inference/runinference_metrics/pipeline/transformations.py apache_beam/examples/inference/tfx_bsl/__init__.py apache_beam/examples/inference/tfx_bsl/build_tensorflow_model.py apache_beam/examples/inference/tfx_bsl/tensorflow_image_classification.py apache_beam/examples/inference/tfx_bsl/tfx_bsl_inference_it_test.py apache_beam/examples/kafkataxi/__init__.py apache_beam/examples/kafkataxi/kafka_taxi.py apache_beam/examples/ml_transform/__init__.py apache_beam/examples/ml_transform/ml_transform_basic.py apache_beam/examples/ml_transform/ml_transform_it_test.py apache_beam/examples/ml_transform/vocab_tfidf_processing.py apache_beam/examples/sinks/__init__.py apache_beam/examples/sinks/generate_event.py apache_beam/examples/sinks/test_periodicimpulse.py apache_beam/examples/sinks/test_write_bounded.py apache_beam/examples/sinks/test_write_unbounded.py apache_beam/examples/snippets/__init__.py apache_beam/examples/snippets/snippets.py apache_beam/examples/snippets/snippets_examples_wordcount_debugging.py apache_beam/examples/snippets/snippets_examples_wordcount_minimal.py apache_beam/examples/snippets/snippets_examples_wordcount_wordcount.py apache_beam/examples/snippets/snippets_test.py apache_beam/examples/snippets/util.py apache_beam/examples/snippets/util_test.py apache_beam/examples/snippets/transforms/__init__.py apache_beam/examples/snippets/transforms/aggregation/__init__.py apache_beam/examples/snippets/transforms/aggregation/approximatequantiles.py apache_beam/examples/snippets/transforms/aggregation/approximatequantiles_test.py apache_beam/examples/snippets/transforms/aggregation/approximateunique.py apache_beam/examples/snippets/transforms/aggregation/approximateunique_test.py apache_beam/examples/snippets/transforms/aggregation/batchelements.py apache_beam/examples/snippets/transforms/aggregation/batchelements_test.py apache_beam/examples/snippets/transforms/aggregation/cogroupbykey.py apache_beam/examples/snippets/transforms/aggregation/cogroupbykey_test.py apache_beam/examples/snippets/transforms/aggregation/combineglobally_combinefn.py apache_beam/examples/snippets/transforms/aggregation/combineglobally_function.py apache_beam/examples/snippets/transforms/aggregation/combineglobally_lambda.py apache_beam/examples/snippets/transforms/aggregation/combineglobally_multiple_arguments.py apache_beam/examples/snippets/transforms/aggregation/combineglobally_side_inputs_dict.py apache_beam/examples/snippets/transforms/aggregation/combineglobally_side_inputs_iter.py apache_beam/examples/snippets/transforms/aggregation/combineglobally_side_inputs_singleton.py apache_beam/examples/snippets/transforms/aggregation/combineglobally_test.py apache_beam/examples/snippets/transforms/aggregation/combineperkey_combinefn.py apache_beam/examples/snippets/transforms/aggregation/combineperkey_function.py apache_beam/examples/snippets/transforms/aggregation/combineperkey_lambda.py apache_beam/examples/snippets/transforms/aggregation/combineperkey_multiple_arguments.py apache_beam/examples/snippets/transforms/aggregation/combineperkey_side_inputs_dict.py apache_beam/examples/snippets/transforms/aggregation/combineperkey_side_inputs_iter.py apache_beam/examples/snippets/transforms/aggregation/combineperkey_side_inputs_singleton.py apache_beam/examples/snippets/transforms/aggregation/combineperkey_simple.py apache_beam/examples/snippets/transforms/aggregation/combineperkey_test.py apache_beam/examples/snippets/transforms/aggregation/combinevalues_combinefn.py apache_beam/examples/snippets/transforms/aggregation/combinevalues_function.py apache_beam/examples/snippets/transforms/aggregation/combinevalues_lambda.py apache_beam/examples/snippets/transforms/aggregation/combinevalues_multiple_arguments.py apache_beam/examples/snippets/transforms/aggregation/combinevalues_side_inputs_dict.py apache_beam/examples/snippets/transforms/aggregation/combinevalues_side_inputs_iter.py apache_beam/examples/snippets/transforms/aggregation/combinevalues_side_inputs_singleton.py apache_beam/examples/snippets/transforms/aggregation/combinevalues_simple.py apache_beam/examples/snippets/transforms/aggregation/combinevalues_test.py apache_beam/examples/snippets/transforms/aggregation/count_globally.py apache_beam/examples/snippets/transforms/aggregation/count_per_element.py apache_beam/examples/snippets/transforms/aggregation/count_per_key.py apache_beam/examples/snippets/transforms/aggregation/count_test.py apache_beam/examples/snippets/transforms/aggregation/distinct.py apache_beam/examples/snippets/transforms/aggregation/distinct_test.py apache_beam/examples/snippets/transforms/aggregation/groupby_attr.py apache_beam/examples/snippets/transforms/aggregation/groupby_attr_expr.py apache_beam/examples/snippets/transforms/aggregation/groupby_expr.py apache_beam/examples/snippets/transforms/aggregation/groupby_expr_aggregate.py apache_beam/examples/snippets/transforms/aggregation/groupby_global_aggregate.py apache_beam/examples/snippets/transforms/aggregation/groupby_simple_aggregate.py apache_beam/examples/snippets/transforms/aggregation/groupby_test.py apache_beam/examples/snippets/transforms/aggregation/groupby_two_exprs.py apache_beam/examples/snippets/transforms/aggregation/groupbykey.py apache_beam/examples/snippets/transforms/aggregation/groupbykey_test.py apache_beam/examples/snippets/transforms/aggregation/groupintobatches.py apache_beam/examples/snippets/transforms/aggregation/groupintobatches_test.py apache_beam/examples/snippets/transforms/aggregation/latest_globally.py apache_beam/examples/snippets/transforms/aggregation/latest_per_key.py apache_beam/examples/snippets/transforms/aggregation/latest_test.py apache_beam/examples/snippets/transforms/aggregation/max_globally.py apache_beam/examples/snippets/transforms/aggregation/max_per_key.py apache_beam/examples/snippets/transforms/aggregation/max_test.py apache_beam/examples/snippets/transforms/aggregation/mean_globally.py apache_beam/examples/snippets/transforms/aggregation/mean_per_key.py apache_beam/examples/snippets/transforms/aggregation/mean_test.py apache_beam/examples/snippets/transforms/aggregation/min_globally.py apache_beam/examples/snippets/transforms/aggregation/min_per_key.py apache_beam/examples/snippets/transforms/aggregation/min_test.py apache_beam/examples/snippets/transforms/aggregation/sample_fixed_size_globally.py apache_beam/examples/snippets/transforms/aggregation/sample_fixed_size_per_key.py apache_beam/examples/snippets/transforms/aggregation/sample_test.py apache_beam/examples/snippets/transforms/aggregation/sum_globally.py apache_beam/examples/snippets/transforms/aggregation/sum_per_key.py apache_beam/examples/snippets/transforms/aggregation/sum_test.py apache_beam/examples/snippets/transforms/aggregation/tolist.py apache_beam/examples/snippets/transforms/aggregation/tolist_test.py apache_beam/examples/snippets/transforms/aggregation/top_largest.py apache_beam/examples/snippets/transforms/aggregation/top_largest_per_key.py apache_beam/examples/snippets/transforms/aggregation/top_of.py apache_beam/examples/snippets/transforms/aggregation/top_per_key.py apache_beam/examples/snippets/transforms/aggregation/top_smallest.py apache_beam/examples/snippets/transforms/aggregation/top_smallest_per_key.py apache_beam/examples/snippets/transforms/aggregation/top_test.py apache_beam/examples/snippets/transforms/elementwise/__init__.py apache_beam/examples/snippets/transforms/elementwise/enrichment.py apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py apache_beam/examples/snippets/transforms/elementwise/filter_function.py apache_beam/examples/snippets/transforms/elementwise/filter_lambda.py apache_beam/examples/snippets/transforms/elementwise/filter_multiple_arguments.py apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_dict.py apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_iter.py apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_singleton.py apache_beam/examples/snippets/transforms/elementwise/filter_test.py apache_beam/examples/snippets/transforms/elementwise/flatmap_function.py apache_beam/examples/snippets/transforms/elementwise/flatmap_generator.py apache_beam/examples/snippets/transforms/elementwise/flatmap_lambda.py apache_beam/examples/snippets/transforms/elementwise/flatmap_multiple_arguments.py apache_beam/examples/snippets/transforms/elementwise/flatmap_nofunction.py apache_beam/examples/snippets/transforms/elementwise/flatmap_side_inputs_dict.py apache_beam/examples/snippets/transforms/elementwise/flatmap_side_inputs_iter.py apache_beam/examples/snippets/transforms/elementwise/flatmap_side_inputs_singleton.py apache_beam/examples/snippets/transforms/elementwise/flatmap_simple.py apache_beam/examples/snippets/transforms/elementwise/flatmap_test.py apache_beam/examples/snippets/transforms/elementwise/flatmap_tuple.py apache_beam/examples/snippets/transforms/elementwise/keys.py apache_beam/examples/snippets/transforms/elementwise/keys_test.py apache_beam/examples/snippets/transforms/elementwise/kvswap.py apache_beam/examples/snippets/transforms/elementwise/kvswap_test.py apache_beam/examples/snippets/transforms/elementwise/map_context.py apache_beam/examples/snippets/transforms/elementwise/map_function.py apache_beam/examples/snippets/transforms/elementwise/map_lambda.py apache_beam/examples/snippets/transforms/elementwise/map_multiple_arguments.py apache_beam/examples/snippets/transforms/elementwise/map_side_inputs_dict.py apache_beam/examples/snippets/transforms/elementwise/map_side_inputs_iter.py apache_beam/examples/snippets/transforms/elementwise/map_side_inputs_singleton.py apache_beam/examples/snippets/transforms/elementwise/map_simple.py apache_beam/examples/snippets/transforms/elementwise/map_test.py apache_beam/examples/snippets/transforms/elementwise/map_tuple.py apache_beam/examples/snippets/transforms/elementwise/mltransform.py apache_beam/examples/snippets/transforms/elementwise/mltransform_test.py apache_beam/examples/snippets/transforms/elementwise/pardo_dofn.py apache_beam/examples/snippets/transforms/elementwise/pardo_dofn_methods.py apache_beam/examples/snippets/transforms/elementwise/pardo_dofn_params.py apache_beam/examples/snippets/transforms/elementwise/pardo_test.py apache_beam/examples/snippets/transforms/elementwise/partition_function.py apache_beam/examples/snippets/transforms/elementwise/partition_lambda.py apache_beam/examples/snippets/transforms/elementwise/partition_multiple_arguments.py apache_beam/examples/snippets/transforms/elementwise/partition_test.py apache_beam/examples/snippets/transforms/elementwise/regex_all_matches.py apache_beam/examples/snippets/transforms/elementwise/regex_find.py apache_beam/examples/snippets/transforms/elementwise/regex_find_all.py apache_beam/examples/snippets/transforms/elementwise/regex_find_kv.py apache_beam/examples/snippets/transforms/elementwise/regex_matches.py apache_beam/examples/snippets/transforms/elementwise/regex_matches_kv.py apache_beam/examples/snippets/transforms/elementwise/regex_replace_all.py apache_beam/examples/snippets/transforms/elementwise/regex_replace_first.py apache_beam/examples/snippets/transforms/elementwise/regex_split.py apache_beam/examples/snippets/transforms/elementwise/regex_test.py apache_beam/examples/snippets/transforms/elementwise/runinference.py apache_beam/examples/snippets/transforms/elementwise/runinference_sklearn_keyed_model_handler.py apache_beam/examples/snippets/transforms/elementwise/runinference_sklearn_unkeyed_model_handler.py apache_beam/examples/snippets/transforms/elementwise/runinference_test.py apache_beam/examples/snippets/transforms/elementwise/tostring_element.py apache_beam/examples/snippets/transforms/elementwise/tostring_iterables.py apache_beam/examples/snippets/transforms/elementwise/tostring_kvs.py apache_beam/examples/snippets/transforms/elementwise/tostring_test.py apache_beam/examples/snippets/transforms/elementwise/values.py apache_beam/examples/snippets/transforms/elementwise/values_test.py apache_beam/examples/snippets/transforms/elementwise/withtimestamps.py apache_beam/examples/snippets/transforms/elementwise/withtimestamps_event_time.py apache_beam/examples/snippets/transforms/elementwise/withtimestamps_logical_clock.py apache_beam/examples/snippets/transforms/elementwise/withtimestamps_processing_time.py apache_beam/examples/snippets/transforms/elementwise/withtimestamps_test.py apache_beam/examples/snippets/transforms/other/__init__.py apache_beam/examples/snippets/transforms/other/create.py apache_beam/examples/snippets/transforms/other/create_test.py apache_beam/examples/snippets/transforms/other/flatten.py apache_beam/examples/snippets/transforms/other/flatten_test.py apache_beam/examples/snippets/transforms/other/window.py apache_beam/examples/snippets/transforms/other/window_test.py apache_beam/internal/__init__.py apache_beam/internal/cloudpickle_pickler.py apache_beam/internal/cloudpickle_pickler_test.py apache_beam/internal/code_object_pickler.py apache_beam/internal/code_object_pickler_test.py apache_beam/internal/dill_pickler.py apache_beam/internal/http_client.py apache_beam/internal/http_client_test.py apache_beam/internal/module_test.py apache_beam/internal/pickler.py apache_beam/internal/pickler_test.py apache_beam/internal/set_pickler.py apache_beam/internal/set_pickler_test.py apache_beam/internal/util.py apache_beam/internal/util_test.py apache_beam/internal/azure/__init__.py apache_beam/internal/azure/auth.py apache_beam/internal/cloudpickle/__init__.py apache_beam/internal/cloudpickle/cloudpickle.py apache_beam/internal/cloudpickle/cloudpickle_fast.py apache_beam/internal/gcp/__init__.py apache_beam/internal/gcp/auth.py apache_beam/internal/gcp/auth_test.py apache_beam/internal/gcp/json_value.py apache_beam/internal/gcp/json_value_test.py apache_beam/internal/metrics/__init__.py apache_beam/internal/metrics/metric.py apache_beam/internal/test_data/__init__.py apache_beam/internal/test_data/module_1.py apache_beam/internal/test_data/module_1_class_added.py apache_beam/internal/test_data/module_1_function_added.py apache_beam/internal/test_data/module_1_global_variable_added.py apache_beam/internal/test_data/module_1_lambda_variable_added.py apache_beam/internal/test_data/module_1_local_variable_added.py apache_beam/internal/test_data/module_1_local_variable_removed.py apache_beam/internal/test_data/module_1_nested_function_2_added.py apache_beam/internal/test_data/module_1_nested_function_added.py apache_beam/internal/test_data/module_2.py apache_beam/internal/test_data/module_2_modified.py apache_beam/internal/test_data/module_3.py apache_beam/internal/test_data/module_3_modified.py apache_beam/internal/test_data/module_with_default_argument.py apache_beam/io/__init__.py apache_beam/io/avroio.py apache_beam/io/avroio_test.py apache_beam/io/concat_source.py apache_beam/io/concat_source_test.py apache_beam/io/debezium.py apache_beam/io/filebasedio_perf_test.py apache_beam/io/filebasedsink.py apache_beam/io/filebasedsink_test.py apache_beam/io/filebasedsource.py apache_beam/io/filebasedsource_test.py apache_beam/io/fileio.py apache_beam/io/fileio_test.py apache_beam/io/filesystem.py apache_beam/io/filesystem_test.py apache_beam/io/filesystemio.py apache_beam/io/filesystemio_test.py apache_beam/io/filesystems.py apache_beam/io/filesystems_test.py apache_beam/io/hadoopfilesystem.py apache_beam/io/hadoopfilesystem_test.py apache_beam/io/iobase.py apache_beam/io/iobase_it_test.py apache_beam/io/iobase_test.py apache_beam/io/jdbc.py apache_beam/io/kafka.py apache_beam/io/kinesis.py apache_beam/io/localfilesystem.py apache_beam/io/localfilesystem_test.py apache_beam/io/mongodbio.py apache_beam/io/mongodbio_it_test.py apache_beam/io/mongodbio_test.py apache_beam/io/parquetio.py apache_beam/io/parquetio_it_test.py apache_beam/io/parquetio_test.py apache_beam/io/range_trackers.py apache_beam/io/range_trackers_test.py apache_beam/io/requestresponse.py apache_beam/io/requestresponse_it_test.py apache_beam/io/requestresponse_test.py apache_beam/io/restriction_trackers.py apache_beam/io/restriction_trackers_test.py apache_beam/io/snowflake.py apache_beam/io/source_test_utils.py apache_beam/io/source_test_utils_test.py apache_beam/io/sources_test.py apache_beam/io/textio.py apache_beam/io/textio_test.py apache_beam/io/tfrecordio.py apache_beam/io/tfrecordio_test.py apache_beam/io/utils.py apache_beam/io/utils_test.py apache_beam/io/watermark_estimators.py apache_beam/io/watermark_estimators_test.py apache_beam/io/aws/__init__.py apache_beam/io/aws/s3filesystem.py apache_beam/io/aws/s3filesystem_test.py apache_beam/io/aws/s3io.py apache_beam/io/aws/s3io_test.py apache_beam/io/aws/clients/__init__.py apache_beam/io/aws/clients/s3/__init__.py apache_beam/io/aws/clients/s3/boto3_client.py apache_beam/io/aws/clients/s3/client_test.py apache_beam/io/aws/clients/s3/fake_client.py apache_beam/io/aws/clients/s3/messages.py apache_beam/io/azure/__init__.py apache_beam/io/azure/blobstoragefilesystem.py apache_beam/io/azure/blobstoragefilesystem_test.py apache_beam/io/azure/blobstorageio.py apache_beam/io/azure/blobstorageio_test.py apache_beam/io/components/__init__.py apache_beam/io/components/adaptive_throttler.py apache_beam/io/components/adaptive_throttler_test.py apache_beam/io/components/rate_limiter.py apache_beam/io/components/rate_limiter_test.py apache_beam/io/components/util.py apache_beam/io/external/__init__.py apache_beam/io/external/generate_sequence.py apache_beam/io/external/generate_sequence_test.py apache_beam/io/external/kafka.py apache_beam/io/external/snowflake.py apache_beam/io/external/xlang_bigqueryio_it_test.py apache_beam/io/external/xlang_debeziumio_it_test.py apache_beam/io/external/xlang_jdbcio_it_test.py apache_beam/io/external/xlang_kafkaio_it_test.py apache_beam/io/external/xlang_kafkaio_perf_test.py apache_beam/io/external/xlang_kinesisio_it_test.py apache_beam/io/external/xlang_parquetio_test.py apache_beam/io/external/xlang_snowflakeio_it_test.py apache_beam/io/external/gcp/__init__.py apache_beam/io/external/gcp/pubsub.py apache_beam/io/flink/__init__.py apache_beam/io/flink/flink_streaming_impulse_source.py apache_beam/io/flink/flink_streaming_impulse_source_test.py apache_beam/io/gcp/__init__.py apache_beam/io/gcp/big_query_query_to_table_it_test.py apache_beam/io/gcp/big_query_query_to_table_pipeline.py apache_beam/io/gcp/bigquery.py apache_beam/io/gcp/bigquery_avro_tools.py apache_beam/io/gcp/bigquery_avro_tools_test.py apache_beam/io/gcp/bigquery_biglake_test.py apache_beam/io/gcp/bigquery_change_history.py apache_beam/io/gcp/bigquery_change_history_it_test.py apache_beam/io/gcp/bigquery_change_history_test.py apache_beam/io/gcp/bigquery_file_loads.py apache_beam/io/gcp/bigquery_file_loads_test.py apache_beam/io/gcp/bigquery_geography_it_test.py apache_beam/io/gcp/bigquery_io_metadata.py apache_beam/io/gcp/bigquery_io_metadata_test.py apache_beam/io/gcp/bigquery_io_read_it_test.py apache_beam/io/gcp/bigquery_io_read_pipeline.py apache_beam/io/gcp/bigquery_json_it_test.py apache_beam/io/gcp/bigquery_read_internal.py apache_beam/io/gcp/bigquery_read_internal_test.py apache_beam/io/gcp/bigquery_read_it_test.py apache_beam/io/gcp/bigquery_read_perf_test.py apache_beam/io/gcp/bigquery_schema_tools.py apache_beam/io/gcp/bigquery_schema_tools_test.py apache_beam/io/gcp/bigquery_test.py apache_beam/io/gcp/bigquery_tools.py apache_beam/io/gcp/bigquery_tools_test.py apache_beam/io/gcp/bigquery_write_it_test.py apache_beam/io/gcp/bigquery_write_perf_test.py apache_beam/io/gcp/bigtableio.py apache_beam/io/gcp/bigtableio_it_test.py apache_beam/io/gcp/bigtableio_test.py apache_beam/io/gcp/dicomclient.py apache_beam/io/gcp/dicomio.py apache_beam/io/gcp/gce_metadata_util.py apache_beam/io/gcp/gcsfilesystem.py apache_beam/io/gcp/gcsfilesystem_integration_test.py apache_beam/io/gcp/gcsfilesystem_test.py apache_beam/io/gcp/gcsio.py apache_beam/io/gcp/gcsio_integration_test.py apache_beam/io/gcp/gcsio_retry.py apache_beam/io/gcp/gcsio_retry_test.py apache_beam/io/gcp/gcsio_test.py apache_beam/io/gcp/pubsub.py apache_beam/io/gcp/pubsub_integration_test.py apache_beam/io/gcp/pubsub_io_perf_test.py apache_beam/io/gcp/pubsub_it_pipeline.py apache_beam/io/gcp/pubsub_test.py apache_beam/io/gcp/resource_identifiers.py apache_beam/io/gcp/spanner.py apache_beam/io/gcp/spanner_wrapper.py apache_beam/io/gcp/datastore/__init__.py apache_beam/io/gcp/datastore/v1new/__init__.py apache_beam/io/gcp/datastore/v1new/datastore_write_it_pipeline.py apache_beam/io/gcp/datastore/v1new/datastore_write_it_test.py apache_beam/io/gcp/datastore/v1new/datastoreio.py apache_beam/io/gcp/datastore/v1new/datastoreio_test.py apache_beam/io/gcp/datastore/v1new/helper.py apache_beam/io/gcp/datastore/v1new/query_splitter.py apache_beam/io/gcp/datastore/v1new/query_splitter_test.py apache_beam/io/gcp/datastore/v1new/rampup_throttling_fn.py apache_beam/io/gcp/datastore/v1new/rampup_throttling_fn_test.py apache_beam/io/gcp/datastore/v1new/types.py apache_beam/io/gcp/datastore/v1new/types_test.py apache_beam/io/gcp/datastore/v1new/util.py apache_beam/io/gcp/datastore/v1new/util_test.py apache_beam/io/gcp/experimental/__init__.py apache_beam/io/gcp/experimental/spannerio.py apache_beam/io/gcp/experimental/spannerio_read_it_test.py apache_beam/io/gcp/experimental/spannerio_read_perf_test.py apache_beam/io/gcp/experimental/spannerio_test.py apache_beam/io/gcp/experimental/spannerio_write_it_test.py apache_beam/io/gcp/experimental/spannerio_write_perf_test.py apache_beam/io/gcp/healthcare/__init__.py apache_beam/io/gcp/healthcare/dicomclient.py apache_beam/io/gcp/healthcare/dicomio.py apache_beam/io/gcp/healthcare/dicomio_integration_test.py apache_beam/io/gcp/healthcare/dicomio_test.py apache_beam/io/gcp/internal/__init__.py apache_beam/io/gcp/internal/clients/__init__.py apache_beam/io/gcp/internal/clients/bigquery/__init__.py apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_client.py apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py apache_beam/io/gcp/tests/__init__.py apache_beam/io/gcp/tests/bigquery_matcher.py apache_beam/io/gcp/tests/bigquery_matcher_test.py apache_beam/io/gcp/tests/pubsub_matcher.py apache_beam/io/gcp/tests/pubsub_matcher_test.py apache_beam/io/gcp/tests/utils.py apache_beam/io/gcp/tests/utils_test.py apache_beam/io/gcp/tests/xlang_spannerio_it_test.py apache_beam/metrics/__init__.py apache_beam/metrics/cells.pxd apache_beam/metrics/cells.py apache_beam/metrics/cells_test.py apache_beam/metrics/execution.pxd apache_beam/metrics/execution.py apache_beam/metrics/execution_test.py apache_beam/metrics/metric.py apache_beam/metrics/metric_test.py apache_beam/metrics/metricbase.py apache_beam/metrics/monitoring_infos.py apache_beam/metrics/monitoring_infos_test.py apache_beam/ml/__init__.py apache_beam/ml/anomaly/__init__.py apache_beam/ml/anomaly/aggregations.py apache_beam/ml/anomaly/aggregations_test.py apache_beam/ml/anomaly/base.py apache_beam/ml/anomaly/base_test.py apache_beam/ml/anomaly/specifiable.py apache_beam/ml/anomaly/specifiable_test.py apache_beam/ml/anomaly/thresholds.py apache_beam/ml/anomaly/thresholds_test.py apache_beam/ml/anomaly/transforms.py apache_beam/ml/anomaly/transforms_test.py apache_beam/ml/anomaly/detectors/__init__.py apache_beam/ml/anomaly/detectors/iqr.py apache_beam/ml/anomaly/detectors/iqr_test.py apache_beam/ml/anomaly/detectors/offline.py apache_beam/ml/anomaly/detectors/pyod_adapter.py apache_beam/ml/anomaly/detectors/pyod_adapter_test.py apache_beam/ml/anomaly/detectors/robust_zscore.py apache_beam/ml/anomaly/detectors/robust_zscore_test.py apache_beam/ml/anomaly/detectors/zscore.py apache_beam/ml/anomaly/detectors/zscore_test.py apache_beam/ml/anomaly/univariate/__init__.py apache_beam/ml/anomaly/univariate/base.py apache_beam/ml/anomaly/univariate/mad.py apache_beam/ml/anomaly/univariate/mad_test.py apache_beam/ml/anomaly/univariate/mean.py apache_beam/ml/anomaly/univariate/mean_test.py apache_beam/ml/anomaly/univariate/median.py apache_beam/ml/anomaly/univariate/median_test.py apache_beam/ml/anomaly/univariate/perf_test.py apache_beam/ml/anomaly/univariate/quantile.py apache_beam/ml/anomaly/univariate/quantile_test.py apache_beam/ml/anomaly/univariate/stdev.py apache_beam/ml/anomaly/univariate/stdev_test.py apache_beam/ml/gcp/__init__.py apache_beam/ml/gcp/cloud_dlp.py apache_beam/ml/gcp/cloud_dlp_it_test.py apache_beam/ml/gcp/cloud_dlp_test.py apache_beam/ml/gcp/naturallanguageml.py apache_beam/ml/gcp/naturallanguageml_test.py apache_beam/ml/gcp/naturallanguageml_test_it.py apache_beam/ml/gcp/recommendations_ai.py apache_beam/ml/gcp/recommendations_ai_test.py apache_beam/ml/gcp/recommendations_ai_test_it.py apache_beam/ml/gcp/videointelligenceml.py apache_beam/ml/gcp/videointelligenceml_test.py apache_beam/ml/gcp/videointelligenceml_test_it.py apache_beam/ml/gcp/visionml.py apache_beam/ml/gcp/visionml_test.py apache_beam/ml/gcp/visionml_test_it.py apache_beam/ml/inference/__init__.py apache_beam/ml/inference/agent_development_kit.py apache_beam/ml/inference/agent_development_kit_test.py apache_beam/ml/inference/anthropic_inference.py apache_beam/ml/inference/anthropic_inference_it_test.py apache_beam/ml/inference/anthropic_inference_test.py apache_beam/ml/inference/base.py apache_beam/ml/inference/base_test.py apache_beam/ml/inference/gemini_inference.py apache_beam/ml/inference/gemini_inference_it_test.py apache_beam/ml/inference/gemini_inference_test.py apache_beam/ml/inference/huggingface_inference.py apache_beam/ml/inference/huggingface_inference_it_test.py apache_beam/ml/inference/huggingface_inference_test.py apache_beam/ml/inference/model_manager.py apache_beam/ml/inference/model_manager_it_test.py apache_beam/ml/inference/model_manager_test.py apache_beam/ml/inference/onnx_inference.py apache_beam/ml/inference/onnx_inference_it_test.py apache_beam/ml/inference/onnx_inference_test.py apache_beam/ml/inference/pytorch_inference.py apache_beam/ml/inference/pytorch_inference_it_test.py apache_beam/ml/inference/pytorch_inference_test.py apache_beam/ml/inference/sklearn_inference.py apache_beam/ml/inference/sklearn_inference_it_test.py apache_beam/ml/inference/sklearn_inference_test.py apache_beam/ml/inference/tensorflow_inference.py apache_beam/ml/inference/tensorflow_inference_it_test.py apache_beam/ml/inference/tensorflow_inference_test.py apache_beam/ml/inference/tensorrt_inference.py apache_beam/ml/inference/tensorrt_inference_test.py apache_beam/ml/inference/utils.py apache_beam/ml/inference/utils_test.py apache_beam/ml/inference/vertex_ai_inference.py apache_beam/ml/inference/vertex_ai_inference_it_test.py apache_beam/ml/inference/vertex_ai_inference_test.py apache_beam/ml/inference/vllm_inference.py apache_beam/ml/inference/xgboost_inference.py apache_beam/ml/inference/xgboost_inference_it_test.py apache_beam/ml/inference/xgboost_inference_test.py apache_beam/ml/rag/__init__.py apache_beam/ml/rag/test_utils.py apache_beam/ml/rag/types.py apache_beam/ml/rag/utils.py apache_beam/ml/rag/chunking/__init__.py apache_beam/ml/rag/chunking/base.py apache_beam/ml/rag/chunking/base_test.py apache_beam/ml/rag/chunking/langchain.py apache_beam/ml/rag/chunking/langchain_test.py apache_beam/ml/rag/embeddings/__init__.py apache_beam/ml/rag/embeddings/base.py apache_beam/ml/rag/embeddings/base_test.py apache_beam/ml/rag/embeddings/huggingface.py apache_beam/ml/rag/embeddings/huggingface_test.py apache_beam/ml/rag/embeddings/vertex_ai.py apache_beam/ml/rag/embeddings/vertex_ai_test.py apache_beam/ml/rag/enrichment/__init__.py apache_beam/ml/rag/enrichment/bigquery_vector_search.py apache_beam/ml/rag/enrichment/bigquery_vector_search_it_test.py apache_beam/ml/rag/enrichment/milvus_search.py apache_beam/ml/rag/enrichment/milvus_search_it_test.py apache_beam/ml/rag/enrichment/milvus_search_test.py apache_beam/ml/rag/ingestion/__init__.py apache_beam/ml/rag/ingestion/alloydb.py apache_beam/ml/rag/ingestion/alloydb_it_test.py apache_beam/ml/rag/ingestion/base.py apache_beam/ml/rag/ingestion/base_test.py apache_beam/ml/rag/ingestion/bigquery.py apache_beam/ml/rag/ingestion/bigquery_it_test.py apache_beam/ml/rag/ingestion/cloudsql.py apache_beam/ml/rag/ingestion/cloudsql_it_test.py apache_beam/ml/rag/ingestion/jdbc_common.py apache_beam/ml/rag/ingestion/milvus_search.py apache_beam/ml/rag/ingestion/milvus_search_it_test.py apache_beam/ml/rag/ingestion/milvus_search_test.py apache_beam/ml/rag/ingestion/mysql.py apache_beam/ml/rag/ingestion/mysql_common.py apache_beam/ml/rag/ingestion/postgres.py apache_beam/ml/rag/ingestion/postgres_common.py apache_beam/ml/rag/ingestion/postgres_it_test.py apache_beam/ml/rag/ingestion/spanner.py apache_beam/ml/rag/ingestion/spanner_it_test.py apache_beam/ml/rag/ingestion/test_utils.py apache_beam/ml/transforms/__init__.py apache_beam/ml/transforms/base.py apache_beam/ml/transforms/base_test.py apache_beam/ml/transforms/handlers.py apache_beam/ml/transforms/handlers_test.py apache_beam/ml/transforms/tft.py apache_beam/ml/transforms/tft_test.py apache_beam/ml/transforms/utils.py apache_beam/ml/transforms/embeddings/__init__.py apache_beam/ml/transforms/embeddings/huggingface.py apache_beam/ml/transforms/embeddings/huggingface_test.py apache_beam/ml/transforms/embeddings/open_ai.py apache_beam/ml/transforms/embeddings/open_ai_it_test.py apache_beam/ml/transforms/embeddings/tensorflow_hub.py apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py apache_beam/ml/transforms/embeddings/vertex_ai.py apache_beam/ml/transforms/embeddings/vertex_ai_it_test.py apache_beam/ml/ts/__init__.py apache_beam/options/__init__.py apache_beam/options/pipeline_options.py apache_beam/options/pipeline_options_context.py apache_beam/options/pipeline_options_context_test.py apache_beam/options/pipeline_options_test.py apache_beam/options/pipeline_options_validator.py apache_beam/options/pipeline_options_validator_test.py apache_beam/options/value_provider.py apache_beam/options/value_provider_test.py apache_beam/portability/__init__.py apache_beam/portability/common_urns.py apache_beam/portability/python_urns.py apache_beam/portability/utils.py apache_beam/portability/api/__init__.py apache_beam/portability/api/standard_coders.yaml apache_beam/portability/api/org/__init__.py apache_beam/portability/api/org/apache/__init__.py apache_beam/portability/api/org/apache/beam/__init__.py apache_beam/portability/api/org/apache/beam/model/__init__.py apache_beam/portability/api/org/apache/beam/model/fn_execution/__init__.py apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/__init__.py apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_fn_api_pb2.py apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_fn_api_pb2.pyi apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_fn_api_pb2_grpc.py apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_fn_api_pb2_urns.py apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_provision_api_pb2.py apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_provision_api_pb2.pyi apache_beam/portability/api/org/apache/beam/model/fn_execution/v1/beam_provision_api_pb2_grpc.py apache_beam/portability/api/org/apache/beam/model/interactive/__init__.py apache_beam/portability/api/org/apache/beam/model/interactive/v1/__init__.py apache_beam/portability/api/org/apache/beam/model/interactive/v1/beam_interactive_api_pb2.py apache_beam/portability/api/org/apache/beam/model/interactive/v1/beam_interactive_api_pb2.pyi apache_beam/portability/api/org/apache/beam/model/interactive/v1/beam_interactive_api_pb2_grpc.py apache_beam/portability/api/org/apache/beam/model/job_management/__init__.py apache_beam/portability/api/org/apache/beam/model/job_management/v1/__init__.py apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_artifact_api_pb2.py apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_artifact_api_pb2.pyi apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_artifact_api_pb2_grpc.py apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_artifact_api_pb2_urns.py apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_expansion_api_pb2.py apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_expansion_api_pb2.pyi apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_expansion_api_pb2_grpc.py apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_job_api_pb2.py apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_job_api_pb2.pyi apache_beam/portability/api/org/apache/beam/model/job_management/v1/beam_job_api_pb2_grpc.py apache_beam/portability/api/org/apache/beam/model/pipeline/__init__.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/__init__.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/beam_runner_api_pb2.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/beam_runner_api_pb2.pyi apache_beam/portability/api/org/apache/beam/model/pipeline/v1/beam_runner_api_pb2_grpc.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/beam_runner_api_pb2_urns.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/endpoints_pb2.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/endpoints_pb2.pyi apache_beam/portability/api/org/apache/beam/model/pipeline/v1/endpoints_pb2_grpc.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/external_transforms_pb2.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/external_transforms_pb2.pyi apache_beam/portability/api/org/apache/beam/model/pipeline/v1/external_transforms_pb2_grpc.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/external_transforms_pb2_urns.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/metrics_pb2.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/metrics_pb2.pyi apache_beam/portability/api/org/apache/beam/model/pipeline/v1/metrics_pb2_grpc.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/metrics_pb2_urns.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/schema_pb2.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/schema_pb2.pyi apache_beam/portability/api/org/apache/beam/model/pipeline/v1/schema_pb2_grpc.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/schema_pb2_urns.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/standard_window_fns_pb2.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/standard_window_fns_pb2.pyi apache_beam/portability/api/org/apache/beam/model/pipeline/v1/standard_window_fns_pb2_grpc.py apache_beam/portability/api/org/apache/beam/model/pipeline/v1/standard_window_fns_pb2_urns.py apache_beam/runners/__init__.py apache_beam/runners/common.pxd apache_beam/runners/common.py apache_beam/runners/common_test.py apache_beam/runners/pipeline_context.py apache_beam/runners/pipeline_context_test.py apache_beam/runners/pipeline_utils.py apache_beam/runners/pipeline_utils_test.py apache_beam/runners/render.py apache_beam/runners/render_test.py apache_beam/runners/runner.py apache_beam/runners/runner_test.py apache_beam/runners/sdf_utils.py apache_beam/runners/sdf_utils_test.py apache_beam/runners/trivial_runner.py apache_beam/runners/trivial_runner_test.py apache_beam/runners/dask/__init__.py apache_beam/runners/dask/dask_runner.py apache_beam/runners/dask/dask_runner_test.py apache_beam/runners/dask/overrides.py apache_beam/runners/dask/transform_evaluator.py apache_beam/runners/dataflow/__init__.py apache_beam/runners/dataflow/dataflow_exercise_metrics_pipeline.py apache_beam/runners/dataflow/dataflow_exercise_metrics_pipeline_test.py apache_beam/runners/dataflow/dataflow_exercise_streaming_metrics_pipeline.py apache_beam/runners/dataflow/dataflow_exercise_streaming_metrics_pipeline_test.py apache_beam/runners/dataflow/dataflow_job_service.py apache_beam/runners/dataflow/dataflow_job_service_test.py apache_beam/runners/dataflow/dataflow_metrics.py apache_beam/runners/dataflow/dataflow_metrics_test.py apache_beam/runners/dataflow/dataflow_runner.py apache_beam/runners/dataflow/dataflow_runner_test.py apache_beam/runners/dataflow/ptransform_overrides.py apache_beam/runners/dataflow/template_runner_test.py apache_beam/runners/dataflow/test_dataflow_runner.py apache_beam/runners/dataflow/internal/__init__.py apache_beam/runners/dataflow/internal/apiclient.py apache_beam/runners/dataflow/internal/apiclient_test.py apache_beam/runners/dataflow/internal/names.py apache_beam/runners/dataflow/internal/clients/__init__.py apache_beam/runners/dataflow/internal/clients/dataflow/__init__.py apache_beam/runners/dataflow/internal/clients/dataflow/dataflow_v1b3_client.py apache_beam/runners/dataflow/internal/clients/dataflow/dataflow_v1b3_messages.py apache_beam/runners/dataflow/internal/clients/dataflow/message_matchers.py apache_beam/runners/dataflow/internal/clients/dataflow/message_matchers_test.py apache_beam/runners/direct/__init__.py apache_beam/runners/direct/bundle_factory.py apache_beam/runners/direct/clock.py apache_beam/runners/direct/consumer_tracking_pipeline_visitor.py apache_beam/runners/direct/consumer_tracking_pipeline_visitor_test.py apache_beam/runners/direct/direct_metrics.py apache_beam/runners/direct/direct_metrics_test.py apache_beam/runners/direct/direct_runner.py apache_beam/runners/direct/direct_runner_test.py apache_beam/runners/direct/direct_userstate.py apache_beam/runners/direct/evaluation_context.py apache_beam/runners/direct/executor.py apache_beam/runners/direct/sdf_direct_runner.py apache_beam/runners/direct/sdf_direct_runner_test.py apache_beam/runners/direct/test_direct_runner.py apache_beam/runners/direct/test_stream_impl.py apache_beam/runners/direct/transform_evaluator.py apache_beam/runners/direct/util.py apache_beam/runners/direct/watermark_manager.py apache_beam/runners/interactive/__init__.py apache_beam/runners/interactive/augmented_pipeline.py apache_beam/runners/interactive/augmented_pipeline_test.py apache_beam/runners/interactive/background_caching_job.py apache_beam/runners/interactive/background_caching_job_test.py apache_beam/runners/interactive/cache_manager.py apache_beam/runners/interactive/cache_manager_test.py apache_beam/runners/interactive/interactive_beam.py apache_beam/runners/interactive/interactive_beam_test.py apache_beam/runners/interactive/interactive_environment.py apache_beam/runners/interactive/interactive_environment_test.py apache_beam/runners/interactive/interactive_runner.py apache_beam/runners/interactive/interactive_runner_test.py apache_beam/runners/interactive/non_interactive_runner_test.py apache_beam/runners/interactive/pipeline_fragment.py apache_beam/runners/interactive/pipeline_fragment_test.py apache_beam/runners/interactive/pipeline_instrument.py apache_beam/runners/interactive/pipeline_instrument_test.py apache_beam/runners/interactive/recording_manager.py apache_beam/runners/interactive/recording_manager_test.py apache_beam/runners/interactive/user_pipeline_tracker.py apache_beam/runners/interactive/user_pipeline_tracker_test.py apache_beam/runners/interactive/utils.py apache_beam/runners/interactive/utils_test.py apache_beam/runners/interactive/caching/__init__.py apache_beam/runners/interactive/caching/cacheable.py apache_beam/runners/interactive/caching/expression_cache.py apache_beam/runners/interactive/caching/expression_cache_test.py apache_beam/runners/interactive/caching/read_cache.py apache_beam/runners/interactive/caching/read_cache_test.py apache_beam/runners/interactive/caching/reify.py apache_beam/runners/interactive/caching/streaming_cache.py apache_beam/runners/interactive/caching/streaming_cache_test.py apache_beam/runners/interactive/caching/write_cache.py apache_beam/runners/interactive/caching/write_cache_test.py apache_beam/runners/interactive/dataproc/__init__.py apache_beam/runners/interactive/dataproc/dataproc_cluster_manager.py apache_beam/runners/interactive/dataproc/dataproc_cluster_manager_test.py apache_beam/runners/interactive/dataproc/types.py apache_beam/runners/interactive/display/__init__.py apache_beam/runners/interactive/display/display_manager.py apache_beam/runners/interactive/display/interactive_pipeline_graph.py apache_beam/runners/interactive/display/pcoll_visualization.py apache_beam/runners/interactive/display/pcoll_visualization_test.py apache_beam/runners/interactive/display/pipeline_graph.py apache_beam/runners/interactive/display/pipeline_graph_renderer.py apache_beam/runners/interactive/display/pipeline_graph_test.py apache_beam/runners/interactive/messaging/__init__.py apache_beam/runners/interactive/messaging/interactive_environment_inspector.py apache_beam/runners/interactive/messaging/interactive_environment_inspector_test.py apache_beam/runners/interactive/options/__init__.py apache_beam/runners/interactive/options/capture_control.py apache_beam/runners/interactive/options/capture_control_test.py apache_beam/runners/interactive/options/capture_limiters.py apache_beam/runners/interactive/options/capture_limiters_test.py apache_beam/runners/interactive/options/interactive_options.py apache_beam/runners/interactive/sql/__init__.py apache_beam/runners/interactive/sql/beam_sql_magics.py apache_beam/runners/interactive/sql/beam_sql_magics_test.py apache_beam/runners/interactive/sql/sql_chain.py apache_beam/runners/interactive/sql/sql_chain_test.py apache_beam/runners/interactive/sql/utils.py apache_beam/runners/interactive/sql/utils_test.py apache_beam/runners/interactive/testing/__init__.py apache_beam/runners/interactive/testing/mock_env.py apache_beam/runners/interactive/testing/mock_ipython.py apache_beam/runners/interactive/testing/pipeline_assertion.py apache_beam/runners/interactive/testing/test_cache_manager.py apache_beam/runners/interactive/testing/integration/__init__.py apache_beam/runners/interactive/testing/integration/notebook_executor.py apache_beam/runners/interactive/testing/integration/screen_diff.py apache_beam/runners/internal/__init__.py apache_beam/runners/internal/names.py apache_beam/runners/job/__init__.py apache_beam/runners/job/manager.py apache_beam/runners/job/utils.py apache_beam/runners/portability/__init__.py apache_beam/runners/portability/abstract_job_service.py apache_beam/runners/portability/artifact_service.py apache_beam/runners/portability/artifact_service_test.py apache_beam/runners/portability/expansion_service.py apache_beam/runners/portability/expansion_service_main.py apache_beam/runners/portability/expansion_service_test.py apache_beam/runners/portability/flink_runner.py apache_beam/runners/portability/flink_runner_test.py apache_beam/runners/portability/flink_uber_jar_job_server.py apache_beam/runners/portability/flink_uber_jar_job_server_test.py apache_beam/runners/portability/job_server.py apache_beam/runners/portability/job_server_test.py apache_beam/runners/portability/local_job_service.py apache_beam/runners/portability/local_job_service_main.py apache_beam/runners/portability/local_job_service_test.py apache_beam/runners/portability/portable_metrics.py apache_beam/runners/portability/portable_runner.py apache_beam/runners/portability/portable_runner_test.py apache_beam/runners/portability/prism_runner.py apache_beam/runners/portability/prism_runner_test.py apache_beam/runners/portability/requirements_cache_it_test.py apache_beam/runners/portability/sdk_container_builder.py apache_beam/runners/portability/sdk_container_builder_test.py apache_beam/runners/portability/spark_java_job_server_test.py apache_beam/runners/portability/spark_runner.py apache_beam/runners/portability/spark_runner_test.py apache_beam/runners/portability/spark_uber_jar_job_server.py apache_beam/runners/portability/spark_uber_jar_job_server_test.py apache_beam/runners/portability/stager.py apache_beam/runners/portability/stager_test.py apache_beam/runners/portability/fn_api_runner/__init__.py apache_beam/runners/portability/fn_api_runner/execution.py apache_beam/runners/portability/fn_api_runner/fn_runner.py apache_beam/runners/portability/fn_api_runner/fn_runner_test.py apache_beam/runners/portability/fn_api_runner/translations.py apache_beam/runners/portability/fn_api_runner/translations_test.py apache_beam/runners/portability/fn_api_runner/trigger_manager.py apache_beam/runners/portability/fn_api_runner/trigger_manager_test.py apache_beam/runners/portability/fn_api_runner/visualization_tools.py apache_beam/runners/portability/fn_api_runner/watermark_manager.py apache_beam/runners/portability/fn_api_runner/worker_handlers.py apache_beam/runners/portability/fn_api_runner/worker_handlers_test.py apache_beam/runners/test/__init__.py apache_beam/runners/worker/__init__.py apache_beam/runners/worker/bundle_processor.py apache_beam/runners/worker/bundle_processor_test.py apache_beam/runners/worker/channel_factory.py apache_beam/runners/worker/crossplatform_time.h apache_beam/runners/worker/crossplatform_unistd.h apache_beam/runners/worker/data_plane.py apache_beam/runners/worker/data_plane_test.py apache_beam/runners/worker/data_sampler.py apache_beam/runners/worker/data_sampler_test.py apache_beam/runners/worker/log_handler.py apache_beam/runners/worker/log_handler_test.py apache_beam/runners/worker/logger.py apache_beam/runners/worker/logger_test.py apache_beam/runners/worker/opcounters.pxd apache_beam/runners/worker/opcounters.py apache_beam/runners/worker/opcounters_test.py apache_beam/runners/worker/operation_specs.py apache_beam/runners/worker/operations.pxd apache_beam/runners/worker/operations.py apache_beam/runners/worker/sdk_worker.py apache_beam/runners/worker/sdk_worker_main.py apache_beam/runners/worker/sdk_worker_main_test.py apache_beam/runners/worker/sdk_worker_test.py apache_beam/runners/worker/sideinputs.py apache_beam/runners/worker/sideinputs_test.py apache_beam/runners/worker/statecache.py apache_beam/runners/worker/statecache_test.py apache_beam/runners/worker/statesampler.py apache_beam/runners/worker/statesampler_fast.pxd apache_beam/runners/worker/statesampler_fast.pyx apache_beam/runners/worker/statesampler_slow.py apache_beam/runners/worker/statesampler_test.py apache_beam/runners/worker/worker_id_interceptor.py apache_beam/runners/worker/worker_id_interceptor_test.py apache_beam/runners/worker/worker_pool_main.py apache_beam/runners/worker/worker_status.py apache_beam/runners/worker/worker_status_test.py apache_beam/testing/__init__.py apache_beam/testing/datatype_inference.py apache_beam/testing/datatype_inference_test.py apache_beam/testing/extra_assertions.py apache_beam/testing/extra_assertions_test.py apache_beam/testing/fast_test_utils.pxd apache_beam/testing/fast_test_utils.pyx apache_beam/testing/metric_result_matchers.py apache_beam/testing/metric_result_matchers_test.py apache_beam/testing/pipeline_verifiers.py apache_beam/testing/pipeline_verifiers_test.py apache_beam/testing/synthetic_pipeline.py apache_beam/testing/synthetic_pipeline_test.py apache_beam/testing/test_pipeline.py apache_beam/testing/test_pipeline_test.py apache_beam/testing/test_stream.py apache_beam/testing/test_stream_it_test.py apache_beam/testing/test_stream_service.py apache_beam/testing/test_stream_service_test.py apache_beam/testing/test_stream_test.py apache_beam/testing/test_utils.py apache_beam/testing/test_utils_test.py apache_beam/testing/util.py apache_beam/testing/util_test.py apache_beam/testing/analyzers/__init__.py apache_beam/testing/analyzers/constants.py apache_beam/testing/analyzers/github_issues_utils.py apache_beam/testing/analyzers/load_test_perf_analysis.py apache_beam/testing/analyzers/perf_analysis.py apache_beam/testing/analyzers/perf_analysis_test.py apache_beam/testing/analyzers/perf_analysis_utils.py apache_beam/testing/benchmarks/__init__.py apache_beam/testing/benchmarks/sort_and_batch_benchmark.py apache_beam/testing/benchmarks/chicago_taxi/__init__.py apache_beam/testing/benchmarks/chicago_taxi/preprocess.py apache_beam/testing/benchmarks/chicago_taxi/process_tfma.py apache_beam/testing/benchmarks/chicago_taxi/setup.py apache_beam/testing/benchmarks/chicago_taxi/tfdv_analyze_and_validate.py apache_beam/testing/benchmarks/chicago_taxi/trainer/__init__.py apache_beam/testing/benchmarks/chicago_taxi/trainer/model.py apache_beam/testing/benchmarks/chicago_taxi/trainer/task.py apache_beam/testing/benchmarks/chicago_taxi/trainer/taxi.py apache_beam/testing/benchmarks/cloudml/__init__.py apache_beam/testing/benchmarks/cloudml/cloudml_benchmark_constants_lib.py apache_beam/testing/benchmarks/cloudml/cloudml_benchmark_test.py apache_beam/testing/benchmarks/cloudml/criteo_tft/__init__.py apache_beam/testing/benchmarks/cloudml/criteo_tft/criteo.py apache_beam/testing/benchmarks/cloudml/pipelines/__init__.py apache_beam/testing/benchmarks/cloudml/pipelines/workflow.py apache_beam/testing/benchmarks/inference/__init__.py apache_beam/testing/benchmarks/inference/pytorch_image_classification_benchmarks.py apache_beam/testing/benchmarks/inference/pytorch_language_modeling_benchmarks.py apache_beam/testing/benchmarks/inference/pytorch_sentiment_benchmarks.py apache_beam/testing/benchmarks/inference/table_row_inference_benchmark.py apache_beam/testing/benchmarks/inference/tensorflow_mnist_classification_cost_benchmark.py apache_beam/testing/benchmarks/inference/vllm_gemma_benchmarks.py apache_beam/testing/benchmarks/nexmark/__init__.py apache_beam/testing/benchmarks/nexmark/monitor.py apache_beam/testing/benchmarks/nexmark/nexmark_launcher.py apache_beam/testing/benchmarks/nexmark/nexmark_perf.py apache_beam/testing/benchmarks/nexmark/nexmark_util.py apache_beam/testing/benchmarks/nexmark/models/__init__.py apache_beam/testing/benchmarks/nexmark/models/auction_bid.py apache_beam/testing/benchmarks/nexmark/models/field_name.py apache_beam/testing/benchmarks/nexmark/models/nexmark_model.py apache_beam/testing/benchmarks/nexmark/queries/__init__.py apache_beam/testing/benchmarks/nexmark/queries/nexmark_query_util.py apache_beam/testing/benchmarks/nexmark/queries/query0.py apache_beam/testing/benchmarks/nexmark/queries/query1.py apache_beam/testing/benchmarks/nexmark/queries/query10.py apache_beam/testing/benchmarks/nexmark/queries/query11.py apache_beam/testing/benchmarks/nexmark/queries/query12.py apache_beam/testing/benchmarks/nexmark/queries/query2.py apache_beam/testing/benchmarks/nexmark/queries/query3.py apache_beam/testing/benchmarks/nexmark/queries/query4.py apache_beam/testing/benchmarks/nexmark/queries/query5.py apache_beam/testing/benchmarks/nexmark/queries/query6.py apache_beam/testing/benchmarks/nexmark/queries/query7.py apache_beam/testing/benchmarks/nexmark/queries/query8.py apache_beam/testing/benchmarks/nexmark/queries/query9.py apache_beam/testing/benchmarks/nexmark/queries/winning_bids.py apache_beam/testing/benchmarks/wordcount/__init__.py apache_beam/testing/benchmarks/wordcount/wordcount.py apache_beam/testing/data/trigger_transcripts.yaml apache_beam/testing/load_tests/__init__.py apache_beam/testing/load_tests/co_group_by_key_test.py apache_beam/testing/load_tests/combine_test.py apache_beam/testing/load_tests/dataflow_cost_benchmark.py apache_beam/testing/load_tests/dataflow_cost_consts.py apache_beam/testing/load_tests/group_by_key_test.py apache_beam/testing/load_tests/load_test.py apache_beam/testing/load_tests/load_test_metrics_utils.py apache_beam/testing/load_tests/microbenchmarks_test.py apache_beam/testing/load_tests/pardo_test.py apache_beam/testing/load_tests/sideinput_test.py apache_beam/tools/__init__.py apache_beam/tools/coders_microbenchmark.py apache_beam/tools/distribution_counter_microbenchmark.py apache_beam/tools/fn_api_runner_microbenchmark.py apache_beam/tools/map_fn_microbenchmark.py apache_beam/tools/microbenchmarks_test.py apache_beam/tools/runtime_type_check_microbenchmark.py apache_beam/tools/sideinput_microbenchmark.py apache_beam/tools/teststream_microbenchmark.py apache_beam/tools/utils.py apache_beam/transforms/__init__.py apache_beam/transforms/async_dofn.py apache_beam/transforms/async_dofn_test.py apache_beam/transforms/batch_dofn_test.py apache_beam/transforms/combinefn_lifecycle_pipeline.py apache_beam/transforms/combinefn_lifecycle_test.py apache_beam/transforms/combiners.py apache_beam/transforms/combiners_test.py apache_beam/transforms/core.py apache_beam/transforms/core_it_test.py apache_beam/transforms/core_test.py apache_beam/transforms/create_source.py apache_beam/transforms/create_test.py apache_beam/transforms/cy_combiners.pxd apache_beam/transforms/cy_combiners.py apache_beam/transforms/cy_dataflow_distribution_counter.pxd apache_beam/transforms/cy_dataflow_distribution_counter.pyx apache_beam/transforms/dataflow_distribution_counter_test.py apache_beam/transforms/deduplicate.py apache_beam/transforms/deduplicate_test.py apache_beam/transforms/display.py apache_beam/transforms/display_test.py apache_beam/transforms/dofn_lifecycle_test.py apache_beam/transforms/enrichment.py apache_beam/transforms/enrichment_it_test.py apache_beam/transforms/enrichment_test.py apache_beam/transforms/environments.py apache_beam/transforms/environments_test.py apache_beam/transforms/error_handling.py apache_beam/transforms/error_handling_test.py apache_beam/transforms/external.py apache_beam/transforms/external_it_test.py apache_beam/transforms/external_java.py apache_beam/transforms/external_test.py apache_beam/transforms/external_transform_provider.py apache_beam/transforms/external_transform_provider_it_test.py apache_beam/transforms/fully_qualified_named_transform.py apache_beam/transforms/fully_qualified_named_transform_test.py apache_beam/transforms/managed.py apache_beam/transforms/managed_iceberg_it_test.py apache_beam/transforms/maven_repository_url_test.py apache_beam/transforms/periodicsequence.py apache_beam/transforms/periodicsequence_it_test.py apache_beam/transforms/periodicsequence_test.py apache_beam/transforms/ptransform.py apache_beam/transforms/ptransform_test.py apache_beam/transforms/py_dataflow_distribution_counter.py apache_beam/transforms/resources.py apache_beam/transforms/resources_test.py apache_beam/transforms/sideinputs.py apache_beam/transforms/sideinputs_test.py apache_beam/transforms/sql.py apache_beam/transforms/sql_test.py apache_beam/transforms/stats.pxd apache_beam/transforms/stats.py apache_beam/transforms/stats_test.py apache_beam/transforms/timestamped_value_type_test.py apache_beam/transforms/timeutil.py apache_beam/transforms/transforms_keyword_only_args_test.py apache_beam/transforms/trigger.py apache_beam/transforms/trigger_test.py apache_beam/transforms/userstate.py apache_beam/transforms/userstate_test.py apache_beam/transforms/util.py apache_beam/transforms/util_test.py apache_beam/transforms/validate_runner_xlang_test.py apache_beam/transforms/window.py apache_beam/transforms/window_test.py apache_beam/transforms/write_ptransform_test.py apache_beam/transforms/enrichment_handlers/__init__.py apache_beam/transforms/enrichment_handlers/bigquery.py apache_beam/transforms/enrichment_handlers/bigquery_it_test.py apache_beam/transforms/enrichment_handlers/bigquery_test.py apache_beam/transforms/enrichment_handlers/bigtable.py apache_beam/transforms/enrichment_handlers/bigtable_it_test.py apache_beam/transforms/enrichment_handlers/bigtable_test.py apache_beam/transforms/enrichment_handlers/cloudsql.py apache_beam/transforms/enrichment_handlers/cloudsql_it_test.py apache_beam/transforms/enrichment_handlers/cloudsql_test.py apache_beam/transforms/enrichment_handlers/feast_feature_store.py apache_beam/transforms/enrichment_handlers/feast_feature_store_it_test.py apache_beam/transforms/enrichment_handlers/feast_feature_store_test.py apache_beam/transforms/enrichment_handlers/utils.py apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store.py apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_it_test.py apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_test.py apache_beam/transforms/xlang/__init__.py apache_beam/transforms/xlang/io.py apache_beam/typehints/__init__.py apache_beam/typehints/arrow_batching_microbenchmark.py apache_beam/typehints/arrow_type_compatibility.py apache_beam/typehints/arrow_type_compatibility_test.py apache_beam/typehints/batch.py apache_beam/typehints/batch_test.py apache_beam/typehints/decorators.py apache_beam/typehints/decorators_test.py apache_beam/typehints/intrinsic_one_ops.py apache_beam/typehints/intrinsic_one_ops_test.py apache_beam/typehints/native_type_compatibility.py apache_beam/typehints/native_type_compatibility_test.py apache_beam/typehints/opcodes.py apache_beam/typehints/pandas_type_compatibility.py apache_beam/typehints/pandas_type_compatibility_test.py apache_beam/typehints/pytorch_type_compatibility.py apache_beam/typehints/pytorch_type_compatibility_test.py apache_beam/typehints/row_type.py apache_beam/typehints/row_type_test.py apache_beam/typehints/schema_registry.py apache_beam/typehints/schemas.py apache_beam/typehints/schemas_test.py apache_beam/typehints/sharded_key_type.py apache_beam/typehints/sharded_key_type_test.py apache_beam/typehints/tagged_output_typehints_test.py apache_beam/typehints/trivial_inference.py apache_beam/typehints/trivial_inference_test.py apache_beam/typehints/typecheck.py apache_beam/typehints/typecheck_test.py apache_beam/typehints/typed_pipeline_test.py apache_beam/typehints/typehints.py apache_beam/typehints/typehints_test.py apache_beam/typehints/testing/__init__.py apache_beam/typehints/testing/strategies.py apache_beam/utils/__init__.py apache_beam/utils/annotations.py apache_beam/utils/annotations_test.py apache_beam/utils/byte_limited_queue.pxd apache_beam/utils/byte_limited_queue.py apache_beam/utils/byte_limited_queue_test.py apache_beam/utils/counters.pxd apache_beam/utils/counters.py apache_beam/utils/counters_test.py apache_beam/utils/histogram.py apache_beam/utils/histogram_test.py apache_beam/utils/interactive_utils.py apache_beam/utils/interactive_utils_test.py apache_beam/utils/logger.py apache_beam/utils/logger_test.py apache_beam/utils/multi_process_shared.py apache_beam/utils/multi_process_shared_test.py apache_beam/utils/plugin.py apache_beam/utils/processes.py apache_beam/utils/processes_test.py apache_beam/utils/profiler.py apache_beam/utils/profiler_test.py apache_beam/utils/proto_utils.py apache_beam/utils/proto_utils_test.py apache_beam/utils/python_callable.py apache_beam/utils/python_callable_test.py apache_beam/utils/retry.py apache_beam/utils/retry_test.py apache_beam/utils/sentinel.py apache_beam/utils/sharded_key.py apache_beam/utils/shared.py apache_beam/utils/shared_test.py apache_beam/utils/subprocess_server.py apache_beam/utils/subprocess_server_test.py apache_beam/utils/thread_pool_executor.py apache_beam/utils/thread_pool_executor_test.py apache_beam/utils/timestamp.py apache_beam/utils/timestamp_test.py apache_beam/utils/transform_service_launcher.py apache_beam/utils/urns.py apache_beam/utils/windowed_value.pxd apache_beam/utils/windowed_value.py apache_beam/utils/windowed_value_test.py apache_beam/yaml/__init__.py apache_beam/yaml/cache_provider_artifacts.py apache_beam/yaml/conftest.py apache_beam/yaml/generate_yaml_docs.py apache_beam/yaml/integration_tests.py apache_beam/yaml/json_utils.py apache_beam/yaml/json_utils_test.py apache_beam/yaml/main.py apache_beam/yaml/main_test.py apache_beam/yaml/options.py apache_beam/yaml/pipeline.schema.yaml apache_beam/yaml/programming_guide_test.py apache_beam/yaml/readme_test.py apache_beam/yaml/standard_io.yaml apache_beam/yaml/standard_providers.yaml apache_beam/yaml/yaml_combine.py apache_beam/yaml/yaml_combine_test.py apache_beam/yaml/yaml_enrichment.py apache_beam/yaml/yaml_enrichment_test.py apache_beam/yaml/yaml_errors.py apache_beam/yaml/yaml_io.py apache_beam/yaml/yaml_io_test.py apache_beam/yaml/yaml_join.py apache_beam/yaml/yaml_join_test.py apache_beam/yaml/yaml_mapping.py apache_beam/yaml/yaml_mapping_test.py apache_beam/yaml/yaml_ml.py apache_beam/yaml/yaml_ml_test.py apache_beam/yaml/yaml_provider.py apache_beam/yaml/yaml_provider_unit_test.py apache_beam/yaml/yaml_specifiable.py apache_beam/yaml/yaml_specifiable_test.py apache_beam/yaml/yaml_testing.py apache_beam/yaml/yaml_testing_test.py apache_beam/yaml/yaml_transform.py apache_beam/yaml/yaml_transform_scope_test.py apache_beam/yaml/yaml_transform_test.py apache_beam/yaml/yaml_transform_unit_test.py apache_beam/yaml/yaml_udf_test.py apache_beam/yaml/yaml_utils.py apache_beam/yaml/yaml_utils_test.py apache_beam/yaml/docs/yaml-combine.md apache_beam/yaml/docs/yaml-errors.md apache_beam/yaml/docs/yaml-inline-python.md apache_beam/yaml/docs/yaml-join.md apache_beam/yaml/docs/yaml-providers.md apache_beam/yaml/docs/yaml-schema.md apache_beam/yaml/docs/yaml-testing.md apache_beam/yaml/docs/yaml-udf.md apache_beam/yaml/docs/yaml.md apache_beam/yaml/examples/__init__.py apache_beam/yaml/examples/testing/__init__.py apache_beam/yaml/examples/testing/examples_test.py apache_beam/yaml/examples/testing/input_data.py././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996270.0 apache_beam-2.74.0/apache_beam.egg-info/dependency_links.txt0000644000175100017510000000000115206113156023424 0ustar00runnerrunner ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996267.0 apache_beam-2.74.0/apache_beam.egg-info/not-zip-safe0000644000175100017510000000000115206113153021601 0ustar00runnerrunner ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996270.0 apache_beam-2.74.0/apache_beam.egg-info/requires.txt0000644000175100017510000001302615206113156021760 0ustar00runnerrunnercryptography<48.0.0,>=39.0.0 fastavro<2,>=0.23.6 fasteners<1.0,>=0.3 grpcio!=1.48.0,!=1.59.*,!=1.60.*,!=1.61.*,!=1.62.0,!=1.62.1,!=1.66.*,!=1.67.*,!=1.68.*,!=1.69.*,!=1.70.*,<2,>=1.33.1 httplib2<0.32.0,>=0.8 jsonpickle<4.0.0,>=3.0.0 numpy<2.5.0,>=1.14.3 objsize<0.8.0,>=0.6.1 packaging>=22.0 pillow<13,>=12.1.1 pymongo<5.0.0,>=3.8.0 proto-plus<2,>=1.7.1 protobuf!=4.0.*,!=4.21.*,!=4.22.0,!=4.23.*,!=4.24.*,<7.0.0.dev0,>=3.20.3 python-dateutil<3,>=2.8.0 pytz>=2018.3 requests<3.0.0,>=2.32.4 sortedcontainers>=2.4.0 typing-extensions>=3.7.0 zstandard<1,>=0.18.0 pyyaml<7.0.0,>=3.12 beartype<0.23.0,>=0.21.0 pyarrow<24.0.0,>=6.0.1 pyarrow-hotfix<1 [:python_version < "3.11"] envoy-data-plane<0.3.0 [:python_version >= "3.11"] envoy-data-plane<2,>=1.0.3 [aws] boto3<2,>=1.9 [azure] azure-storage-blob<13,>=12.3.2 azure-core<2,>=1.7.0 azure-identity<2,>=1.12.0 [dask] distributed>=2024.4.2 dask>=2024.4.2 [dataframe] pandas!=1.5.0,!=1.5.1,<2.3,>=1.4.3 [dev] isort==7.0.0 pyrefly==0.54.0 ruff==0.15.7 yapf==0.43.0 [dill] dill<0.3.2,>=0.3.1.1 [docs] jinja2<3.2,>=3.0 Sphinx<8.0,>=7.0.0 docstring-parser<1.0,>=0.15 docutils>=0.18.1 markdown pandas<2.3.0 openai virtualenv-clone<1.0,>=0.5 [gcp] cachetools<7,>=3.1.0 google-api-core<3,>=2.0.0 google-auth<3,>=1.18.0 google-auth-httplib2<0.3.0,>=0.1.0 google-cloud-datastore<3,>=2.0.0 google-cloud-pubsub<3,>=2.1.0 google-cloud-storage<4,>=2.18.2 google-cloud-bigquery<4,>=2.0.0 google-cloud-bigquery-storage<3,>=2.6.3 google-cloud-core<3,>=2.0.0 google-cloud-bigtable<3,>=2.19.0 google-cloud-build<4,>=3.35.0 google-cloud-spanner<4,>=3.0.0 google-cloud-dlp<4,>=3.0.0 google-cloud-kms<4,>=3.0.0 google-cloud-language<3,>=2.0 google-cloud-secret-manager<3,>=2.0 google-cloud-videointelligence<3,>=2.0 google-cloud-vision<4,>=2 google-cloud-recommendations-ai<0.11.0,>=0.1.0 google-cloud-aiplatform<2.0,>=1.26.0 cloud-sql-python-connector<2.0.0,>=1.18.2 python-tds>=1.16.1 pg8000>=1.31.5 PyMySQL>=1.1.0 keyrings.google-artifactregistry-auth orjson<4,>=3.9.7 regex>=2020.6.8 [gcp:python_version < "3.13"] google-apitools<0.5.32,>=0.5.31 [gcp:python_version >= "3.13"] google-apitools>=0.5.35 [hadoop] hdfs<3.0.0,>=2.1.0 [interactive] facets-overview<2,>=1.1.0 google-cloud-dataproc<6,>=5.0.0 ipython<9,>=7 ipykernel<7,>=6 ipywidgets<9,>=8 jupyter-client!=6.1.13,<8.2.1,>=6.1.11 pydot<2,>=1.2.0 timeloop<2,>=1.0.2 nbformat<6,>=5.0.5 nbconvert<8,>=6.2.0 pandas!=1.5.0,!=1.5.1,<2.3,>=1.4.3 [interactive_test] needle<1,>=0.5.0 chromedriver-binary<118,>=117 pillow<10,>=7.1.1 urllib3<2,>=1.21.1 [milvus] pymilvus<3.0.0,>=2.5.10 [ml_cpu] tensorflow>=2.12.0 torch==2.8.0+cpu transformers<4.56.0,>=4.28.0 absl-py>=0.12.0 [ml_test] datatable dill tensorflow_transform<1.15.0,>=1.14.0 numpy<2 tf2onnx<1.17,>=1.16.1 embeddings>=0.0.4 onnxruntime onnx<2,>=1.14.1 langchain sentence-transformers>=2.2.2 skl2onnx pyod>=0.7.6 tensorflow absl-py>=0.12.0 tensorflow-hub tokenizers<0.23 torch<2.12.0 transformers<4.56.0,>=4.28.0 tokenizers<0.22.0,>=0.13.3 [onnx] onnxruntime==1.13.1 torch==1.13.1 tensorflow==2.11.0 tf2onnx==1.13.0 skl2onnx==1.13 transformers==4.25.1 absl-py>=0.12.0 [p310_ml_test] datatable embeddings>=0.0.4 onnxruntime onnx<2,>=1.14.1 langchain sentence-transformers>=2.2.2 skl2onnx pyod>=0.7.6 tensorflow absl-py>=0.12.0 tensorflow-hub tokenizers<0.23 torch<2.12.0 transformers<4.56.0,>=4.28.0 tokenizers<0.22.0,>=0.13.3 google-adk==1.28.1 proto-plus<2,>=1.26.1 opentelemetry-api==1.37.0 opentelemetry-sdk==1.37.0 opentelemetry-exporter-otlp-proto-http==1.37.0 tf2onnx<1.18,>=1.17.0 [p312_ml_test] datatable embeddings>=0.0.4 onnxruntime onnx<2,>=1.14.1 langchain sentence-transformers>=2.2.2 skl2onnx pyod>=0.7.6 tensorflow absl-py>=0.12.0 tensorflow-hub tokenizers<0.23 torch<2.12.0 transformers<4.56.0,>=4.28.0 tokenizers<0.22.0,>=0.13.3 google-adk==1.28.1 proto-plus<2,>=1.26.1 opentelemetry-api==1.37.0 opentelemetry-sdk==1.37.0 opentelemetry-exporter-otlp-proto-http==1.37.0 tf2onnx<1.18,>=1.17.0 [p313_ml_test] embeddings>=0.0.4 onnxruntime onnx<2,>=1.14.1 langchain sentence-transformers>=2.2.2 skl2onnx pyod>=0.7.6 tensorflow absl-py>=0.12.0 tensorflow-hub tokenizers<0.23 torch<2.12.0 transformers<4.56.0,>=4.28.0 tokenizers<0.22.0,>=0.13.3 google-adk==1.28.1 proto-plus<2,>=1.26.1 opentelemetry-api==1.37.0 opentelemetry-sdk==1.37.0 opentelemetry-exporter-otlp-proto-http==1.37.0 tf2onnx<1.18,>=1.17.0 pymilvus<3.0.0,>=2.5.10 [redis] redis<6,>=5.0.0 [tensorflow] tensorflow<2.21,>=2.12rc1 absl-py>=0.12.0 [tensorflow-hub] tensorflow-hub<0.16.0,>=0.14.0 [test] cloud-sql-python-connector[pg8000]<2.0.0,>=1.0.0 docstring-parser<1.0,>=0.15 freezegun>=0.3.12 jinja2<3.2,>=3.0 joblib>=1.0.1 mock<6.0.0,>=1.0.1 pandas<2.3.0 parameterized<0.10.0,>=0.7.1 pydot<2,>=1.2.0 pyhamcrest!=1.10.0,<3.0.0,>=1.9 requests_mock<2.0,>=1.7 tenacity<10,>=8.0.0 pytest<10.0,>=7.1.2 pytest-xdist<4,>=2.5.0 pytest-timeout<3,>=2.1.0 scikit-learn<1.8.0,>=0.20.0 sqlalchemy<3.0,>=1.3 psycopg2-binary<3.0,>=2.8.5 testcontainers[kafka,milvus,mysql]<5.0.0,>=4.0.0 cryptography>=41.0.2 hypothesis<6.148.4,>5.0.0 virtualenv-clone<1.0,>=0.5 python-tds>=1.16.1 sqlalchemy-pytds>=1.0.2 pg8000>=1.31.5 PyMySQL>=1.1.0 oracledb>=3.1.1 [tfrecord] crcmod<2.0,>=1.7 [tft] tensorflow_transform<1.15.0,>=1.14.0 numpy<2 dill [torch] torch<2.8.0,>=1.9.0 [transformers] transformers<4.56.0,>=4.28.0 tensorflow>=2.12.0 torch>=1.9.0 [vllm] openai==1.107.1 vllm==0.10.1.1 triton==3.3.1 [xgboost] xgboost<2.1.3,>=1.6.0 datatable==1.0.0 [yaml] docstring-parser<1.0,>=0.15 jinja2<3.2,>=3.0 virtualenv-clone<1.0,>=0.5 jsonschema<5.0.0,>=4.0.0 pandas!=1.5.0,!=1.5.1,<2.3,>=1.4.3 [yaml:python_version < "3.12"] js2py<1,>=0.74 ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996270.0 apache_beam-2.74.0/apache_beam.egg-info/top_level.txt0000644000175100017510000000001415206113156022103 0ustar00runnerrunnerapache_beam ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/gen_protos.py0000644000175100017510000004442215206113033016210 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Generates Python proto modules and grpc stubs for Beam protos. """ import argparse import contextlib import glob import importlib.resources import inspect import logging import os import platform import re import shutil import sys from collections import defaultdict from importlib import import_module LOG = logging.getLogger() LOG.setLevel(logging.INFO) LICENSE_HEADER = """ # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ NO_PROMISES_NOTICE = """ \"\"\" For internal use only; no backwards-compatibility guarantees. Automatically generated when running python -m build. \"\"\" """ def clean_path(path): return os.path.realpath(os.path.abspath(path)) # These paths are relative to the project root BEAM_PROTO_PATHS = [ os.path.join('model', 'pipeline', 'src', 'main', 'proto'), os.path.join('model', 'job-management', 'src', 'main', 'proto'), os.path.join('model', 'fn-execution', 'src', 'main', 'proto'), os.path.join('model', 'interactive', 'src', 'main', 'proto'), ] PYTHON_SDK_ROOT = os.path.dirname(clean_path(__file__)) PROJECT_ROOT = clean_path(os.path.join(PYTHON_SDK_ROOT, '..', '..')) PYTHON_OUTPUT_PATH = os.path.join( PYTHON_SDK_ROOT, 'apache_beam', 'portability', 'api') MODEL_RESOURCES = [ os.path.normpath(( 'model/fn-execution/src/main/resources/org/' 'apache/beam/model/fnexecution/v1/standard_coders.yaml')), ] class PythonPath(object): def __init__(self, path: str, front: bool = False): self._path = path self._front = front def __enter__(self): if not self._path: return self._sys_path = sys.path.copy() if self._front: sys.path.insert(0, self._path) else: sys.path.append(self._path) def __exit__(self, exc_type, exc_val, exc_tb): if not self._path: return sys.path = self._sys_path def generate_urn_files(out_dir, api_path): """ Create python files with statically defined URN constants. Creates a _pb2_urn.py file for each _pb2.py file that contains an enum type. This works by importing each api._pb2 module created by `protoc`, inspecting the module's contents, and generating a new side-car urn module. This is executed at build time rather than dynamically on import to ensure that it is compatible with static type checkers like mypy. """ from google.protobuf import message from google.protobuf.internal import api_implementation if api_implementation.Type() == 'python': from google.protobuf.internal import containers repeated_types = ( list, containers.RepeatedScalarFieldContainer, containers.RepeatedCompositeFieldContainer) elif api_implementation.Type() == 'upb': from google._upb import _message repeated_types = ( list, _message.RepeatedScalarContainer, _message.RepeatedCompositeContainer) elif api_implementation.Type() == 'cpp': from google.protobuf.pyext import _message repeated_types = ( list, _message.RepeatedScalarContainer, _message.RepeatedCompositeContainer) else: raise TypeError( "Unknown proto implementation: " + api_implementation.Type()) class Context(object): INDENT = ' ' CAP_SPLIT = re.compile('([A-Z][^A-Z]*|^[a-z]+)') def __init__(self, indent=0): self.lines = [] self.imports = set() self.empty_types = set() self._indent = indent @contextlib.contextmanager def indent(self): self._indent += 1 yield self._indent -= 1 def prepend(self, s): if s: self.lines.insert(0, (self.INDENT * self._indent) + s + '\n') else: self.lines.insert(0, '\n') def line(self, s): if s: self.lines.append((self.INDENT * self._indent) + s + '\n') else: self.lines.append('\n') def import_type(self, typ): modname = typ.__module__ if modname in ('__builtin__', 'builtin'): return typ.__name__ else: self.imports.add(modname) _, modname = modname.rsplit('.', 1) return modname + '.' + typ.__name__ @staticmethod def is_message_type(obj): return isinstance(obj, type) and \ issubclass(obj, message.Message) @staticmethod def is_enum_type(obj): return type(obj).__name__ == 'EnumTypeWrapper' def python_repr(self, obj): if isinstance(obj, message.Message): return self.message_repr(obj) elif isinstance(obj, repeated_types): return '[%s]' % ', '.join(self.python_repr(x) for x in obj) else: return repr(obj) def empty_type(self, typ): name = ( 'EMPTY_' + '_'.join(x.upper() for x in self.CAP_SPLIT.findall(typ.__name__))) self.empty_types.add('%s = %s()' % (name, self.import_type(typ))) return name def message_repr(self, msg): parts = [] for field, value in msg.ListFields(): parts.append('%s=%s' % (field.name, self.python_repr(value))) if parts: return '%s(%s)' % (self.import_type(type(msg)), ', '.join(parts)) else: return self.empty_type(type(msg)) def write_enum(self, enum_name, enum, indent): ctx = Context(indent=indent) with ctx.indent(): for enum_value_name in enum.values_by_name: enum_value_descriptor = enum.values_by_name[enum_value_name] extensions = enum_value_descriptor.GetOptions().Extensions prop = ( extensions[beam_runner_api_pb2.beam_urn], extensions[beam_runner_api_pb2.beam_constant], extensions[metrics_pb2.monitoring_info_spec], extensions[metrics_pb2.label_props], ) reprs = [self.python_repr(x) for x in prop] if all(x == "''" or x.startswith('EMPTY_') for x in reprs): continue ctx.line( '%s = PropertiesFromEnumValue(%s)' % (enum_value_name, ', '.join(self.python_repr(x) for x in prop))) if ctx.lines: ctx.prepend('class %s(object):' % enum_name) ctx.prepend('') ctx.line('') return ctx.lines def write_message(self, message_name, message, indent=0): ctx = Context(indent=indent) with ctx.indent(): for obj_name, obj in inspect.getmembers(message): if obj_name == 'DESCRIPTOR': for enum_name in obj.enum_types_by_name: enum = obj.enum_types_by_name[enum_name] ctx.lines += self.write_enum(enum_name, enum, ctx._indent) if ctx.lines: ctx.prepend('class %s(object):' % message_name) ctx.prepend('') return ctx.lines pb2_files = list(glob.glob(os.path.join(out_dir, '*_pb2.py'))) with PythonPath(os.path.dirname(api_path), front=True): beam_runner_api_pb2 = import_module( 'api.org.apache.beam.model.pipeline.v1.beam_runner_api_pb2') metrics_pb2 = import_module( 'api.org.apache.beam.model.pipeline.v1.metrics_pb2') for pb2_file in pb2_files: modname = os.path.splitext(pb2_file)[0] out_file = modname + '_urns.py' api_start_idx = modname.index(os.path.sep + 'api' + os.path.sep) import_path = modname[api_start_idx + 1:].replace(os.path.sep, '.') mod = import_module(import_path) ctx = Context() for obj_name, obj in inspect.getmembers(mod): if ctx.is_message_type(obj): ctx.lines += ctx.write_message(obj_name, obj) if ctx.lines: for line in reversed(sorted(ctx.empty_types)): ctx.prepend(line) for modname in reversed(sorted(ctx.imports)): pkg, target = modname.rsplit('.', 1) rel_import = build_relative_import(api_path, pkg, out_file) ctx.prepend('from %s import %s' % (rel_import, target)) rel_import = build_relative_import( os.path.dirname(api_path), 'utils', out_file) ctx.prepend('from %s import PropertiesFromEnumValue' % rel_import) LOG.info("Writing urn stubs: %s" % out_file) with open(out_file, 'w') as f: f.writelines(ctx.lines) def _find_protoc_gen_mypy(): # NOTE: this shouldn't be necessary if the virtualenv's environment # is passed to tasks below it, since protoc will search the PATH itself fname = 'protoc-gen-mypy' if platform.system() == 'Windows': fname += ".exe" pathstr = os.environ.get('PATH') search_paths = pathstr.split(os.pathsep) if pathstr else [] # should typically be installed into the venv's bin dir search_paths.insert(0, os.path.dirname(sys.executable)) for path in search_paths: fullpath = os.path.join(path, fname) if os.path.exists(fullpath): LOG.info('Found protoc_gen_mypy at %s' % fullpath) return fullpath raise RuntimeError( "Could not find %s in %s" % (fname, ', '.join(search_paths))) def find_by_ext(root_dir, ext): for root, _, files in os.walk(root_dir): for file in files: if file.endswith(ext): yield clean_path(os.path.join(root, file)) def build_relative_import(root_path, import_path, start_file_path): tail_path = import_path.replace('.', os.path.sep) source_path = os.path.join(root_path, tail_path) is_module = os.path.isfile(source_path + '.py') if is_module: source_path = os.path.dirname(source_path) rel_path = os.path.relpath( source_path, start=os.path.dirname(start_file_path)) if rel_path == '.': if is_module: rel_path += os.path.basename(tail_path) return rel_path if rel_path.endswith('..'): rel_path += os.path.sep # In a path that looks like ../../../foo, every double dot # after the right most double dot needs to be collapsed to # a single dot to look like ././../foo to which we can convert # to ....foo for the proper relative import. first_half_idx = rel_path.rfind('..' + os.path.sep) if first_half_idx == 0: return rel_path.replace(os.path.sep, '') first_half = rel_path[:first_half_idx].replace('..', '.') final_import = first_half.replace(os.path.sep, '') + '..' + \ rel_path[first_half_idx+3:].replace(os.path.sep, '.') if is_module: if final_import.count('.') == len(final_import): return final_import + os.path.basename(tail_path) return final_import + '.{}'.format(os.path.basename(tail_path)) return final_import def generate_init_files_lite(api_root): proto_root = os.path.join(api_root, 'org') for root, _, _ in os.walk(proto_root): init_file = os.path.join(root, '__init__.py') with open(init_file, 'w+'): pass def generate_init_files_full(api_root): proto_root = os.path.join(api_root, 'org') api_module_root = os.path.join(api_root, '__init__.py') modules = defaultdict(list) for root, _, files in os.walk(proto_root): init_file = os.path.join(root, '__init__.py') with open(init_file, 'w+') as f: f.write(LICENSE_HEADER.lstrip()) for file in files: if not file.endswith('.py') or file == '__init__.py': continue module_name = file.split('.')[0] f.write('from . import {}\n'.format(module_name)) modules[root].append(module_name) with open(api_module_root, 'w+') as f: f.write(LICENSE_HEADER.lstrip()) f.write(NO_PROMISES_NOTICE.lstrip()) remaining_lines = [] duplicate_modules = {} for module_root, modules in modules.items(): import_path = os.path.relpath(module_root, api_root).replace(os.path.sep, '.') import_root, imported_module = import_path.rsplit('.', 1) if imported_module not in duplicate_modules: f.write('from .{} import {}\n'.format(import_root, imported_module)) duplicate_modules[imported_module] = 1 else: duplicate_modules[imported_module] += 1 module_alias = '{}_{}'.format( imported_module, duplicate_modules[imported_module]) f.write( 'from .{} import {} as {}\n'.format( import_root, imported_module, module_alias)) imported_module = module_alias for module in modules: remaining_lines.append( '{module} = {}.{module}\n'.format(imported_module, module=module)) f.write('\n') f.writelines(remaining_lines) def generate_proto_files(force=False): """ Will compile proto files for python. If force is not true, then several heuristics are used to determine whether a compilation is necessary. If a compilation is not necessary, no compilation will be performed. :param force: Whether to force a recompilation of the proto files. """ proto_dirs = [ clean_path(os.path.join(PROJECT_ROOT, path)) for path in BEAM_PROTO_PATHS ] proto_files = [ proto_file for d in proto_dirs for proto_file in find_by_ext(d, '.proto') ] out_files = list(find_by_ext(PYTHON_OUTPUT_PATH, '_pb2.py')) if out_files and not proto_files and not force: # We have out_files but no protos; assume they're up-to-date. # This is actually the common case (e.g. installation from an sdist). LOG.info('No proto files; using existing generated files.') return elif not out_files and not proto_files: model = os.path.join(PROJECT_ROOT, 'model') if os.path.exists(model): error_msg = 'No proto files found in %s.' % proto_dirs else: error_msg = 'Not in apache git tree, unable to find proto definitions.' raise RuntimeError(error_msg) if force: regenerate_reason = 'forced' elif not out_files: regenerate_reason = 'no output files' elif len(out_files) < len(proto_files): regenerate_reason = 'not enough output files' elif (min(os.path.getmtime(path) for path in out_files) <= max( os.path.getmtime(path) for path in proto_files + [os.path.realpath(__file__)])): regenerate_reason = 'output files are out-of-date' elif len(out_files) > len(proto_files): regenerate_reason = 'output files without corresponding .proto files' # too many output files: probably due to switching between git branches. # remove them so they don't trigger constant regeneration. for out_file in out_files: os.remove(out_file) else: regenerate_reason = '' if not regenerate_reason: LOG.info('Skipping proto regeneration: all files up to date') return shutil.rmtree(PYTHON_OUTPUT_PATH, ignore_errors=True) if not os.path.exists(PYTHON_OUTPUT_PATH): os.mkdir(PYTHON_OUTPUT_PATH) protoc_gen_mypy = _find_protoc_gen_mypy() from grpc_tools import protoc builtin_protos = importlib.resources.files('grpc_tools') / '_proto' args = ( [sys.executable] + # expecting to be called from command line ['--proto_path=%s' % builtin_protos] + ['--proto_path=%s' % d for d in proto_dirs] + ['--python_out=%s' % PYTHON_OUTPUT_PATH] + ['--plugin=protoc-gen-mypy=%s' % protoc_gen_mypy] + # new version of mypy-protobuf converts None to zero default value # and remove Optional from the param type annotation. This causes # some mypy errors. So to mitigate and fall back to old behavior, # use `relax_strict_optional_primitives` flag. more at # https://github.com/nipunn1313/mypy-protobuf/tree/main#relax_strict_optional_primitives # pylint:disable=line-too-long ['--mypy_out=relax_strict_optional_primitives:%s' % PYTHON_OUTPUT_PATH ] + # TODO(robertwb): Remove the prefix once it's the default. ['--grpc_python_out=grpc_2_0:%s' % PYTHON_OUTPUT_PATH] + proto_files) LOG.info('Regenerating Python proto definitions (%s).' % regenerate_reason) ret_code = protoc.main(args) if ret_code: raise RuntimeError( 'Protoc returned non-zero status (see logs for details): ' '%s' % ret_code) # copy resource files for path in MODEL_RESOURCES: shutil.copy2(os.path.join(PROJECT_ROOT, path), PYTHON_OUTPUT_PATH) proto_packages = set() # see: https://github.com/protocolbuffers/protobuf/issues/1491 # force relative import paths for proto files compiled_import_re = re.compile('^from (.*) import (.*)$') for file_path in find_by_ext(PYTHON_OUTPUT_PATH, ('_pb2.py', '_pb2_grpc.py', '_pb2.pyi')): proto_packages.add(os.path.dirname(file_path)) lines = [] with open(file_path, encoding='utf-8') as f: for line in f: match_obj = compiled_import_re.match(line) if match_obj and \ match_obj.group(1).startswith('org.apache.beam.model'): new_import = build_relative_import( PYTHON_OUTPUT_PATH, match_obj.group(1), file_path) line = 'from %s import %s\n' % (new_import, match_obj.group(2)) lines.append(line) with open(file_path, 'w') as f: f.writelines(lines) generate_init_files_lite(PYTHON_OUTPUT_PATH) for proto_package in proto_packages: generate_urn_files(proto_package, PYTHON_OUTPUT_PATH) generate_init_files_full(PYTHON_OUTPUT_PATH) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--no-force', dest='force', action='store_false') args = parser.parse_args() generate_proto_files(force=args.force) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/gen_xlang_wrappers.py0000644000175100017510000004031415206113033017712 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """ Generates Python wrappers for external transforms (specifically, SchemaTransforms) """ import argparse import datetime import inspect import logging import os import shutil import subprocess import typing from typing import Any from typing import Union import yaml from gen_protos import LICENSE_HEADER from gen_protos import PROJECT_ROOT from gen_protos import PYTHON_SDK_ROOT SUPPORTED_SDK_DESTINATIONS = ['python'] PYTHON_SUFFIX = "_et.py" PY_WRAPPER_OUTPUT_DIR = os.path.join( PYTHON_SDK_ROOT, 'apache_beam', 'transforms', 'xlang') def generate_transforms_config(input_services, output_file): """ Generates a YAML file containing a list of transform configurations. Takes an input YAML file containing a list of expansion service gradle targets. Each service must provide a `destinations` field that specifies the default package (relative path) that generated wrappers should be imported to. A default destination package is specified for each SDK, like so:: - gradle_target: 'sdks:java:io:expansion-service:shadowJar' destinations: python: 'apache_beam/io' We use :class:`ExternalTransformProvider` to discover external transforms. Then, we extract the necessary details of each transform and compile them into a new YAML file, which is later used to generate wrappers. Importing generated transforms to an existing package ----------------------------------------------------- When running the script on the config above, a new module will be created at `apache_beam/transforms/xlang/io.py`. This contains all generated wrappers that are set to destination 'apache_beam/io'. Finally, to make these available to the `apache_beam.io` package (or any package really), just add the following line to the package's `__init__.py` file:: from apache_beam.transforms.xlang.io import * Modifying a transform's name and destination -------------------------------------------- Each service may also specify modifications for particular transform. Currently, one can modify the generated wrapper's **name** and **destination** package: - By default, the transform's identifier is used to generate the wrapper class name. This can be overriden by manually providing a name. - By default, generated wrappers are made available to the package provided by their respective expansion service. This can be overridden by providing a relative path to a different package. See the following example for what such modifications can look like:: - gradle_target: 'sdks:java:io:expansion-service:shadowJar' destinations: python: 'apache_beam/io' transforms: 'beam:schematransform:org.apache.beam:my_transform:v1': name: 'MyCustomTransformName' destinations: python: 'apache_beam/io/gcp' For the above example, we would take the transform with identifier `beam:schematransform:org.apache.beam:my_transform:v1` and by default infer a wrapper class name of `MyTransform` then write it to the module `apache_beam/transforms/xlang/io.py`. With these modifications however, we instead use the provided name `MyCustomTransformName` and write it to `apache_beam/transforms/xlang/io_gcp.py`. Similar to above, this can be made available by importing it in the `__init__.py` file like so:: from apache_beam.transforms.xlang.io_gcp import * Skipping transforms ------------------- To skip a particular transform, simply list its identifier in the `skip_transforms` field, like so:: - gradle_target: 'sdks:java:io:expansion-service:shadowJar' destinations: python: 'apache_beam/io' skip_transforms: - 'beam:schematransform:org.apache.beam:some_transform:v1' """ from apache_beam.transforms.external import BeamJarExpansionService from apache_beam.transforms.external_transform_provider import ExternalTransform from apache_beam.transforms.external_transform_provider import ExternalTransformProvider transform_list: list[dict[str, Any]] = [] with open(input_services) as f: services = yaml.safe_load(f) for service in services: target = service['gradle_target'] if "destinations" not in service: raise ValueError( f"Expansion service with target '{target}' does not " "specify any default destinations.") service_destinations: dict[str, str] = service['destinations'] for sdk, dest in service_destinations.items(): validate_sdks_destinations(sdk, dest, target) transforms_to_skip = service.get('skip_transforms', []) # use dynamic provider to discover and populate wrapper details provider = ExternalTransformProvider(BeamJarExpansionService(target)) discovered: dict[str, ExternalTransform] = provider.get_all() for identifier in sorted(discovered.keys()): wrapper = discovered[identifier] if identifier in transforms_to_skip: continue transform_destinations = service_destinations.copy() # apply any modifications modified_transform = {} if 'transforms' in service and identifier in service['transforms']: modified_transform = service['transforms'][identifier] for sdk, dest in modified_transform.get('destinations', {}).items(): validate_sdks_destinations(sdk, dest, target, identifier) transform_destinations[sdk] = dest # override the destination name = modified_transform.get('name', wrapper.__name__) fields = [] for param_name in sorted(wrapper.configuration_schema.keys()): param = wrapper.configuration_schema[param_name] (tp, nullable) = pretty_type(param.type) field_info = { 'name': param.original_name, 'type': str(tp), 'description': param.description, 'nullable': nullable } fields.append(field_info) transform = { 'identifier': identifier, 'name': name, 'destinations': transform_destinations, 'default_service': target, 'fields': fields, 'description': inspect.getdoc(wrapper) } transform_list.append(transform) with open(output_file, 'w') as f: f.write(LICENSE_HEADER.lstrip()) f.write( "# NOTE: This file is autogenerated and should " "not be edited by hand.\n") f.write( "# Configs are generated based on the expansion service\n" f"# configuration in {input_services.replace(PROJECT_ROOT, '')}.\n") f.write("# Refer to gen_xlang_wrappers.py for more info.\n") dt = datetime.datetime.now().date() f.write(f"#\n# Last updated on: {dt}\n\n") yaml.dump(transform_list, f) logging.info("Successfully wrote transform configs to file: %s", output_file) def validate_sdks_destinations(sdk, dest, service, identifier=None): if identifier: message = f"Identifier '{identifier}'" else: message = f"Service '{service}'" if sdk not in SUPPORTED_SDK_DESTINATIONS: raise ValueError( message + " specifies a destination for an invalid SDK:" f" '{sdk}'. The supported SDKs are {SUPPORTED_SDK_DESTINATIONS}") if not os.path.isdir(os.path.join(PYTHON_SDK_ROOT, *dest.split('/'))): raise ValueError( message + f" specifies an invalid destination '{dest}'." " Please make sure the destination is an existing directory.") def pretty_type(tp): """ Takes a type and returns a tuple containing a pretty string representing it and a bool signifying if it is nullable or not. For optional types, the contained type is unwrapped and returned. This does not recurse however, so inner Optional types are not affected. E.g. the input typing.Optional[typing.Dict[int, typing.Optional[str]]] will return (Dict[int, Union[str, NoneType]], True) """ nullable = False if (typing.get_origin(tp) is Union and type(None) in typing.get_args(tp)): nullable = True # only unwrap if it's a single nullable type. if the type is truly a union # of multiple types, leave it alone. args = typing.get_args(tp) if len(args) == 2: tp = list(filter(lambda t: not isinstance(t, type(None)), args))[0] # TODO(ahmedabu98): Make this more generic to support other remote SDKs # Potentially use Runner API types if tp.__module__ == 'builtins': if getattr(tp, '__origin__', None) is None: tp = tp.__name__ else: # Remove nested typing module name (like Optional) tp = str(tp).replace("typing.", "") elif tp.__module__ == 'typing': tp = str(tp).replace("typing.", "") tp = tp.replace("Sequence", "list") tp = tp.replace("Mapping", "map") elif tp.__module__ == 'numpy': tp = tp.__name__ tp = str(tp).replace("numpy.", "") if tp == "bool": tp = "boolean" return (tp, nullable) def get_wrappers_from_transform_configs(config_file) -> dict[str, list[str]]: """ Generates code for external transform wrapper classes (subclasses of :class:`ExternalTransform`). Takes a YAML file containing a list of SchemaTransform configurations. For each configuration, the code for a wrapper class is generated, along with any documentation that may be included. Each configuration must include a destination file that the generated class will be written to. Returns the generated classes, grouped by destination. """ from jinja2 import Environment from jinja2 import FileSystemLoader env = Environment(loader=FileSystemLoader(PYTHON_SDK_ROOT)) python_wrapper_template = env.get_template("python_xlang_wrapper.template") # maintain a list of wrappers to write in each file. if modified destinations # are used, we may end up with multiple wrappers in one file. destinations: dict[str, list[str]] = {} with open(config_file) as f: transforms = yaml.safe_load(f) for config in transforms: default_service = config['default_service'] description = config['description'] destination = config['destinations']['python'] name = config['name'] fields = config['fields'] identifier = config['identifier'] parameters = [] for field in fields: param_details = { "name": field['name'], "type": field['type'], "description": field['description'], } if field['nullable']: param_details["default"] = None parameters.append(param_details) # Python syntax requires function definitions to have # non-default parameters first parameters = sorted(parameters, key=lambda p: 'default' in p) default_service = f"BeamJarExpansionService(\"{default_service}\")" python_wrapper_class = python_wrapper_template.render( class_name=name, identifier=identifier, parameters=parameters, description=description, default_expansion_service=default_service) if destination not in destinations: destinations[destination] = [] destinations[destination].append(python_wrapper_class) return destinations def write_wrappers_to_destinations( grouped_wrappers: dict[str, list[str]], output_dir=PY_WRAPPER_OUTPUT_DIR, format_code=True): """ Takes a dictionary of generated wrapper code, grouped by destination. For each destination, create a new file containing the respective wrapper classes. Each file includes the Apache License header and relevant imports. Note: the Jinja template should already follow linting and formatting rules. """ written_files = [] for dest, wrappers in grouped_wrappers.items(): module_name = dest.replace('apache_beam/', '').replace('/', '_') module_path = os.path.join(output_dir, module_name) + ".py" with open(module_path, "w") as file: file.write(LICENSE_HEADER.lstrip()) file.write( "\n# NOTE: This file contains autogenerated external transform(s)\n" "# and should not be edited by hand.\n" "# Refer to gen_xlang_wrappers.py for more info.\n\n") file.write( "\"\"\"" "Cross-language transforms in this module can be imported from the\n" f":py:mod:`{dest.replace('/', '.')}` package." "\"\"\"\n\n") file.write( "# pylint:disable=line-too-long\n\n" "from apache_beam.transforms.external import " "BeamJarExpansionService\n" "from apache_beam.transforms.external_transform_provider " "import ExternalTransform\n") for wrapper in wrappers: file.write(wrapper + "\n") written_files.append(module_path) logging.info("Created external transform wrapper modules: %s", written_files) if format_code: formatting_cmd = ['yapf', '--in-place', *written_files] subprocess.run(formatting_cmd, capture_output=True, check=True) def delete_generated_files(root_dir): """Scans for and deletes generated wrapper files.""" logging.info("Deleting external transform wrappers from dir %s", root_dir) deleted_files = os.listdir(root_dir) for file in deleted_files: if file == '__init__.py': deleted_files.remove(file) continue path = os.path.join(root_dir, file) if os.path.isfile(path) or os.path.islink(path): os.unlink(os.path.join(root_dir, file)) else: shutil.rmtree(path) logging.info("Successfully deleted files: %s", deleted_files) def run_script( cleanup, generate_config_only, input_expansion_services, transforms_config_source): # Cleanup first if requested. This is needed to remove outdated wrappers. if cleanup: delete_generated_files(PY_WRAPPER_OUTPUT_DIR) # This step requires the expansion service. # Only generate a transforms config file if none are provided if not transforms_config_source: output_transforms_config = os.path.join( PROJECT_ROOT, 'sdks', 'standard_external_transforms.yaml') generate_transforms_config( input_services=input_expansion_services, output_file=output_transforms_config) transforms_config_source = output_transforms_config else: if not os.path.exists(transforms_config_source): raise RuntimeError( "Could not find the provided transforms config " f"source: {transforms_config_source}") if generate_config_only: return wrappers_grouped_by_destination = get_wrappers_from_transform_configs( transforms_config_source) write_wrappers_to_destinations(wrappers_grouped_by_destination) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument( '--cleanup', dest='cleanup', action='store_true', help="Whether to cleanup existing generated wrappers first.") parser.add_argument( '--generate-config-only', dest='generate_config_only', action='store_true', help="If set, will generate the transform config only without generating" "any wrappers.") parser.add_argument( '--input-expansion-services', dest='input_expansion_services', default=os.path.join( PROJECT_ROOT, 'sdks', 'standard_expansion_services.yaml'), help=( "Absolute path to the input YAML file that contains " "expansion service configs. Ignored if a transforms config" "source is provided.")) parser.add_argument( '--transforms-config-source', dest='transforms_config_source', help=( "Absolute path to a source transforms config YAML file to " "generate wrapper modules from. If not provided, one will be " "created by this script.")) args = parser.parse_args() run_script( args.cleanup, args.generate_config_only, args.input_expansion_services, args.transforms_config_source) ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/pyproject.toml0000644000175100017510000001443215206113033016371 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # since we rely on setuptools and according to https://peps.python.org/pep-0518/#build-system-table # this is the minimum requirements for the build system to execute. [build-system] requires = [ "setuptools", "wheel>=0.36.0", # Note: after upgrading to grpcio-tools, we must bump protobuf to >=5.0. # Currently Blocked on old dependencies of Beam tft extra: # https://github.com/apache/beam/issues/37854 "grpcio-tools==1.62.1; python_version <= '3.12'", "grpcio-tools==1.71.0; python_version >= '3.13'", "mypy-protobuf==3.5.0", # Avoid https://github.com/pypa/virtualenv/issues/2006 "distlib==0.3.9", # Numpy headers "numpy>=1.14.3,<2.5.0", # Update setup.py as well. # having cython here will create wheels that are platform dependent. "cython>=3.0,<4", ## deps for generating external transform wrappers: # also update PyYaml bounds in sdks:python:generateExternalTransformsConfig 'pyyaml>=3.12,<7.0.0', # also update Jinja2 bounds in test-suites/xlang/build.gradle (look for xlangWrapperValidation task) "jinja2>=2.7.1,<4.0.0", 'yapf==0.43.0' ] # legacy installation is needed to generate `apache_beam.portability.api` package. build-backend = "setuptools.build_meta" [tool.isort] py_version = 310 line_length = 120 old_finders = true order_by_type = true force_single_line = true combine_star = true src_paths = ["apache_beam"] extra_standard_library = ["dataclasses"] known_third_party = ["yaml"] skip = [ "apiclient.py", "avroio_test.py", "cloudpickle.py", "datastore_wordcount.py", "datastoreio_test.py", "doctests_test.py", "fast_coders_test.py", "hadoopfilesystem.py", "iobase_test.py", "main_test.py", "model.py", "preprocess.py", "process_tfma.py", "render_test.py", "slow_coders_test.py", "taxi.py", "tfdv_analyze_and_validate.py", "yaml/main.py", "yaml_testing_test.py", "bigquery_v2_client.py", "bigquery_v2_messages.py", "dataflow_v1b3_client.py", "dataflow_v1b3_messages.py", "storage_v1_client.py", "storage_v1_messages.py", "proto2_coder_test_messages_pb2.py", "cloudbuild_v1_client.py", "cloudbuild_v1_messages.py", "boto3_client.py", ] skip_glob = [ "*.pxd", "*.pyx", "*pb2*.py", "**/examples/**/*.py", "**/portability/api/**/*.py", "**/portability/api/__init__.py", ] [tool.ruff] exclude = [ ".bzr", ".direnv", ".eggs", ".git", ".git-rewrite", ".hg", ".ipynb_checkpoints", ".mypy_cache", ".nox", ".pants.d", ".pyenv", ".pytest_cache", ".pytype", ".ruff_cache", ".svn", ".tox", ".venv", ".vscode", "__pypackages__", "_build", "buck-out", "build", "dist", "node_modules", "site-packages", "venv", "*.pxd", "*.pyx", "*pb2*.py", "**/examples/**/*.py", "**/examples/**/*.ipynb", "**/portability/api/**/*.py", "**/portability/api/__init__.py", ] target-version = "py310" src = ["apache_beam"] [tool.ruff.lint] select = ["E9", "PL", "F821", "F822", "F823", "UP006"] ignore = [ # Ignored Pylint Checks "PLC0415", # import-outside-toplevel "PLR2004", # magic-value-comparison "PLR0913", # too-many-arguments "PLR0912", # too-many-branches "PLW0108", # unnecessary-lambda "PLW2901", # redefined-loop-name "PLR0915", # too-many-statements "PLR1714", # repeated-equality-comparison "PLR0911", # too-many-return-statements "PLR5501", # collapsible-else-if "PLW0603", # global-statement "PLR1730", # if-stmt-min-max "PLW1641", # eq-without-hash "PLW0602", # global-variable-not-assigned "PLC1802", # len-test "PLC3002", # unnecessary-direct-lambda-call "PLW0642", # self-or-cls-assignment "PLR1733", # unnecessary-dict-index-lookup "PLR0402", # manual-from-import "PLC0132", # type-param-name-mismatch "PLC0206", # dict-index-missing-items "PLC0207", # missing-maxsplit-arg "PLR1704", # redefined-argument-from-local "PLR1711", # useless-return "PLW0406", # import-self "PLW3301", # nested-min-max "PLR2044", # empty-comment ] # Allow fix for all enabled rules (when --fix) is provided. fixable = ["ALL"] unfixable = [] # Allow unused variables when underscore-prefixed. dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$" [tool.pyrefly] project-includes = ["apache_beam"] python-version = "3.10.0" ignore-missing-imports = ["*"] untyped-def-behavior = "skip-and-infer-return-any" [tool.pyrefly.errors] missing-attribute = "ignore" invalid-annotation = "ignore" redundant-condition = "ignore" untyped-import = "ignore" # Beam-specific supressions to keep pyrefly green # These are in descending order of instances. # # TODO(https://github.com/apache/beam/issues/37699): Reduce the number of # ignored checks. bad-param-name-override = "ignore" bad-override = "ignore" bad-argument-type = "ignore" unsupported-operation = "ignore" implicit-import = "ignore" bad-function-definition = "ignore" bad-return = "ignore" unbound-name = "ignore" no-matching-overload = "ignore" bad-assignment = "ignore" missing-argument = "ignore" bad-index = "ignore" invalid-type-var = "ignore" unknown-name = "ignore" not-a-type = "ignore" deprecated = "ignore" not-callable = "ignore" invalid-argument = "ignore" invalid-inheritance = "ignore" not-iterable = "ignore" unexpected-keyword = "ignore" bad-specialization = "ignore" bad-context-manager = "ignore" invalid-yield = "ignore" bad-argument-count = "ignore" bad-typed-dict-key = "ignore" ././@PaxHeader0000000000000000000000000000003400000000000010212 xustar0028 mtime=1779996271.8704596 apache_beam-2.74.0/setup.cfg0000644000175100017510000000175315206113160015301 0ustar00runnerrunner[nosetests] exe = True verbosity = 2 with-xunit = 1 [coverage:run] branch = True source = apache_beam omit = apache_beam/portability/api/*_pb2.py apache_beam/portability/api/*_pb2_grpc.py [coverage:report] exclude_lines = pragma: no cover abc.abstractmethod def __repr__ if self\.debug raise NotImplementedError if __name__ == .__main__.: [coverage:xml] output = target/site/cobertura/coverage.xml [yapf] indent_width = 2 continuation_indent_width = 4 column_limit = 80 allow_multiline_lambdas = False allow_split_before_dict_value = False blank_line_before_module_docstring = True blank_line_before_nested_class_or_def = False coalesce_brackets = True each_dict_entry_on_separate_line = True split_all_top_level_comma_separated_values = True split_arguments_when_comma_terminated = True split_before_expression_after_opening_paren = True split_before_first_argument = True split_before_logical_operator = False i18n_comment = # type: ignore.* [egg_info] tag_build = tag_date = 0 ././@PaxHeader0000000000000000000000000000002600000000000010213 xustar0022 mtime=1779996187.0 apache_beam-2.74.0/setup.py0000644000175100017510000006737215206113033015202 0ustar00runnerrunner# # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with # this work for additional information regarding copyright ownership. # The ASF licenses this file to You under the Apache License, Version 2.0 # (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Apache Beam SDK for Python setup file.""" import glob import logging import os import shutil import subprocess import sys import warnings # Pylint and isort disagree here. # pylint: disable=ungrouped-imports from importlib.metadata import PackageNotFoundError from importlib.metadata import distribution from pathlib import Path # pylint: disable=ungrouped-imports import setuptools from packaging.version import parse from setuptools import Command # pylint: disable=wrong-import-order # It is recommended to import setuptools prior to importing distutils to avoid # using legacy behavior from distutils. # https://setuptools.readthedocs.io/en/latest/history.html#v48-0-0 from distutils.errors import DistutilsError # isort:skip def to_filename(name: str) -> str: return name.replace('-', '_') def normalize_path(filename): return os.path.normcase(os.path.realpath(os.path.normpath(filename))) class pyrefly(Command): user_options = [] def initialize_options(self): """Abstract method that is required to be overwritten""" def finalize_options(self): """Abstract method that is required to be overwritten""" def get_project_path(self): self.run_command('egg_info') # Build extensions in-place self.reinitialize_command('build_ext', inplace=1) self.run_command('build_ext') ei_cmd = self.get_finalized_command("egg_info") project_path = normalize_path(ei_cmd.egg_base) return os.path.join(project_path, to_filename(ei_cmd.egg_name)) def run(self): args = ['pyrefly', 'check', self.get_project_path()] result = subprocess.call(args) if result != 0: raise DistutilsError("pyrefly exited with status %d" % result) def get_version(): global_names = {} exec( # pylint: disable=exec-used open(os.path.join( os.path.dirname(os.path.abspath(__file__)), 'apache_beam/version.py') ).read(), global_names ) return global_names['__version__'] PACKAGE_NAME = 'apache-beam' PACKAGE_VERSION = get_version() PACKAGE_DESCRIPTION = 'Apache Beam SDK for Python' PACKAGE_URL = 'https://beam.apache.org' PACKAGE_DOWNLOAD_URL = 'https://pypi.python.org/pypi/apache-beam' PACKAGE_AUTHOR = 'Apache Software Foundation' PACKAGE_EMAIL = 'dev@beam.apache.org' PACKAGE_KEYWORDS = 'apache beam' RECOMMENDED_MIN_PIP_VERSION = '19.3.0' try: _PIP_VERSION = distribution('pip').version if parse(_PIP_VERSION) < parse(RECOMMENDED_MIN_PIP_VERSION): warnings.warn( "You are using version {0} of pip. " \ "However, the recommended min version is {1}.".format( _PIP_VERSION, RECOMMENDED_MIN_PIP_VERSION ) ) except PackageNotFoundError: # Do nothing if pip is not found. This can happen when using `Poetry` or # `pipenv` package managers. pass REQUIRED_CYTHON_VERSION = '3.0.0' try: _CYTHON_VERSION = distribution('cython').version if parse(_CYTHON_VERSION) < parse(REQUIRED_CYTHON_VERSION): warnings.warn( "You are using version {0} of cython. " \ "However, version {1} is recommended.".format( _CYTHON_VERSION, REQUIRED_CYTHON_VERSION ) ) except PackageNotFoundError: # do nothing if Cython is not installed pass try: # pylint: disable=wrong-import-position from Cython.Build import cythonize as cythonize0 def cythonize(*args, **kwargs): import numpy extensions = cythonize0(*args, **kwargs) for e in extensions: e.include_dirs.append(numpy.get_include()) return extensions except ImportError: warnings.warn( "Cython not found, cython extensions will not be generated. " \ "To use cythonized extensions, pip install cython and run " \ "python setup.py build_ext --inplace" ) cythonize = lambda *args, **kwargs: [] # [BEAM-8181] pyarrow cannot be installed on 32-bit Windows platforms. if sys.platform == 'win32' and sys.maxsize <= 2**32: pyarrow_dependency = [''] else: pyarrow_dependency = [ 'pyarrow>=6.0.1,<24.0.0', # NOTE(https://github.com/apache/beam/issues/29392): We can remove this # once Beam increases the pyarrow lower bound to a version that fixes CVE. # (lower bound >= 14.0.1) 'pyarrow-hotfix<1' ] # Exclude pandas<=1.4.2 since it doesn't work with numpy 1.24.x. # Exclude 1.5.0 and 1.5.1 because of # https://github.com/pandas-dev/pandas/issues/45725 # must update the below "docs" and "test" for extras_require dataframe_dependency = [ 'pandas>=1.4.3,!=1.5.0,!=1.5.1,<2.3', ] milvus_dependency = ['pymilvus>=2.5.10,<3.0.0'] # google-adk / OpenTelemetry require protobuf>=5; tensorflow-transform in # ml_test is pinned to versions that require protobuf<5 on Python 3.10. Those # cannot be installed together, so ADK deps stay out of ml_test (use ml_base). ml_base_core = [ 'embeddings>=0.0.4', # 0.0.3 crashes setuptools 'onnxruntime', # onnx 1.12–1.13 cap protobuf in ways that trigger huge backtracking with # Beam[gcp]+ml_test; pip can fall back to onnx 1.11 sdist which needs cmake. # 1.14.1+ matches tf2onnx>=1.16 and ships manylinux wheels for py3.10. 'onnx>=1.14.1,<2', 'langchain', 'sentence-transformers>=2.2.2', 'skl2onnx', 'pyod>=0.7.6', # 0.7.5 crashes setuptools 'tensorflow', # tensorflow transitive dep, lower versions not compatible with Python3.10+ 'absl-py>=0.12.0', 'tensorflow-hub', # tokenizers 0.23.0rc0 renamed the PyO3 kwarg of # processors.RobertaProcessing (and BertProcessing) from `cls` to # `cls_token` -- the rename was a drive-by inside huggingface/tokenizers # https://github.com/huggingface/tokenizers/pull/1928. # transformers' slow CLIP tokenizer still calls # `processors.RobertaProcessing(sep=..., cls=..., ...)` at # transformers/models/clip/tokenization_clip.py, so model load fails with # "RobertaProcessing.__new__() got an unexpected keyword argument 'cls'". # The ml tox envs run with pip_pre=True (tox.ini:32), so even though no # 0.23 stable has shipped yet, the rc gets resolved. # Drop this cap once transformers updates the CLIP call site to # `cls_token=` or tokenizers reinstates `cls=` as a deprecation alias. 'tokenizers<0.23', # Avoid torch 2.12.0+ which fails to run unit tests with segfault. 'torch<2.12.0', # Match tested transformers range. 'transformers>=4.28.0,<4.56.0', # Keep tokenizers compatible with this transformers range. 'tokenizers>=0.13.3,<0.22.0', ] ml_adk_dependency = [ 'google-adk==1.28.1', # proto-plus<1.24 caps protobuf<5; opentelemetry-proto (via ADK) needs # protobuf>=5. Scoped here so the main dependency list stays broader. 'proto-plus>=1.26.1,<2', 'opentelemetry-api==1.37.0', 'opentelemetry-sdk==1.37.0', 'opentelemetry-exporter-otlp-proto-http==1.37.0', # protobuf>=5 (ADK/OTel); tf2onnx 1.16.x pins protobuf~=3.20 only. 'tf2onnx>=1.17.0,<1.18', ] ml_base = ml_base_core + ml_adk_dependency def find_by_ext(root_dir, ext): for root, _, files in os.walk(root_dir): for file in files: if file.endswith(ext): yield os.path.realpath(os.path.join(root, file)) # We must generate protos after setup_requires are installed. def generate_protos_first(): try: # Pyproject toml build happens in isolated environemnts. In those envs, # gen_protos is unable to get imported. so we run a subprocess call. cwd = os.path.abspath(os.path.dirname(__file__)) # when pip install <>.tar.gz gets called, if gen_protos.py is not available # in the sdist,then the proto files would have already been generated. So we # skip proto generation in that case. if not os.path.exists(os.path.join(cwd, 'gen_protos.py')): # make sure we already generated protos pb2_files = list( find_by_ext( os.path.join(cwd, 'apache_beam', 'portability', 'api'), '_pb2.py')) if not pb2_files: raise RuntimeError( 'protobuf files are not generated. ' 'Please generate pb2 files') warnings.warn('Skipping proto generation as they are already generated.') return out = subprocess.run( [sys.executable, os.path.join(cwd, 'gen_protos.py'), '--no-force'], capture_output=True, check=True) print(out.stdout) except subprocess.CalledProcessError as err: raise RuntimeError('Could not generate protos due to error: %s', err.stderr) def copy_tests_from_docs(): python_root = os.path.abspath(os.path.dirname(__file__)) docs_src = os.path.normpath( os.path.join( python_root, '../../website/www/site/content/en/documentation/sdks')) docs_dest = os.path.normpath( os.path.join(python_root, 'apache_beam/yaml/docs')) if os.path.exists(docs_src): shutil.rmtree(docs_dest, ignore_errors=True) os.mkdir(docs_dest) for path in glob.glob(os.path.join(docs_src, 'yaml*.md')): shutil.copy(path, docs_dest) else: warnings.warn( f'Could not locate yaml docs source directory {docs_src}. ' f'Skipping copying tests from docs.') def generate_external_transform_wrappers(): try: sdk_dir = os.path.abspath(os.path.dirname(__file__)) script_exists = os.path.exists( os.path.join(sdk_dir, 'gen_xlang_wrappers.py')) config_exists = os.path.exists( os.path.join( os.path.dirname(sdk_dir), 'standard_external_transforms.yaml')) # we need both the script and the standard transforms config file. # at build time, we don't have access to apache_beam to discover and # retrieve external transforms, so the config file has to already exist if not script_exists or not config_exists: generated_transforms_dir = os.path.join( sdk_dir, 'apache_beam', 'transforms', 'xlang') # if exists, this directory will have at least its __init__.py file if (not os.path.exists(generated_transforms_dir) or len(os.listdir(generated_transforms_dir)) <= 1): message = 'External transform wrappers have not been generated ' if not script_exists: message += 'and the generation script `gen_xlang_wrappers.py`' if not config_exists: message += 'and the standard external transforms config' message += ' could not be found' raise RuntimeError(message) else: logging.info( 'Skipping external transform wrapper generation as they ' 'are already generated.') return subprocess.run([ sys.executable, os.path.join(sdk_dir, 'gen_xlang_wrappers.py'), '--cleanup', '--transforms-config-source', os.path.join( os.path.dirname(sdk_dir), 'standard_external_transforms.yaml') ], capture_output=True, check=True) except subprocess.CalledProcessError as err: raise RuntimeError( 'Could not generate external transform wrappers due to ' 'error: {}'.format(err.stderr)) def get_portability_package_data(): files = [] portability_dir = Path(__file__).parent / 'apache_beam' / \ 'portability' / 'api' for ext in ['*.pyi', '*.yaml']: files.extend( str(p.relative_to(portability_dir.parent.parent)) for p in portability_dir.rglob(ext)) return files python_requires = '>=3.10' if sys.version_info.major == 3 and sys.version_info.minor >= 15: warnings.warn( 'This version of Apache Beam has not been sufficiently tested on ' 'Python %s.%s. You may encounter bugs or missing features.' % (sys.version_info.major, sys.version_info.minor)) if __name__ == '__main__': # In order to find the tree of proto packages, the directory # structure must exist before the call to setuptools.find_packages() # executes below. generate_protos_first() generate_external_transform_wrappers() # These data files live elsewhere in the full Beam repository. copy_tests_from_docs() # generate cythonize extensions only if we are building a wheel or # building an extension or running in editable mode. cythonize_cmds = ('bdist_wheel', 'build_ext', 'editable_wheel') if any(cmd in sys.argv for cmd in cythonize_cmds): extensions = cythonize([ 'apache_beam/**/*.pyx', 'apache_beam/coders/coder_impl.py', 'apache_beam/metrics/cells.py', 'apache_beam/metrics/execution.py', 'apache_beam/runners/common.py', 'apache_beam/runners/worker/logger.py', 'apache_beam/runners/worker/opcounters.py', 'apache_beam/runners/worker/operations.py', 'apache_beam/transforms/cy_combiners.py', 'apache_beam/transforms/stats.py', 'apache_beam/utils/byte_limited_queue.py', 'apache_beam/utils/counters.py', 'apache_beam/utils/windowed_value.py', ]) else: extensions = [] try: long_description = ((Path(__file__).parent / "README.md").read_text(encoding='utf-8')) except FileNotFoundError: long_description = ( 'Apache Beam is a unified programming model for both batch and ' 'streaming data processing, enabling efficient execution across ' 'diverse distributed execution engines and providing extensibility ' 'points for connecting to different technologies and user ' 'communities.') # Keep all dependencies inlined in the setup call, otherwise Dependabot won't # be able to parse it. setuptools.setup( name=PACKAGE_NAME, version=PACKAGE_VERSION, description=PACKAGE_DESCRIPTION, long_description=long_description, long_description_content_type='text/markdown', url=PACKAGE_URL, download_url=PACKAGE_DOWNLOAD_URL, author=PACKAGE_AUTHOR, author_email=PACKAGE_EMAIL, packages=setuptools.find_packages(), package_data={ 'apache_beam': [ '*/*.pyx', '*/*/*.pyx', '*/*.pxd', '*/*/*.pxd', '*/*.h', '*/*/*.h', 'testing/data/*.yaml', 'yaml/*.yaml', 'yaml/docs/*.md', *get_portability_package_data() ] }, ext_modules=extensions, install_requires=[ 'cryptography>=39.0.0,<48.0.0', 'envoy-data-plane>=1.0.3,<2; python_version >= "3.11"', # Newer version only work on Python 3.11. Versions 0.3 <= ver < 1.x # conflict with other GCP dependencies. 'envoy-data-plane<0.3.0; python_version < "3.11"', 'fastavro>=0.23.6,<2', 'fasteners>=0.3,<1.0', 'grpcio>=1.33.1,<2,!=1.48.0,!=1.59.*,!=1.60.*,!=1.61.*,!=1.62.0,!=1.62.1,!=1.66.*,!=1.67.*,!=1.68.*,!=1.69.*,!=1.70.*', # pylint: disable=line-too-long 'httplib2>=0.8,<0.32.0', 'jsonpickle>=3.0.0,<4.0.0', # numpy can have breaking changes in minor versions. # Use a strict upper bound. 'numpy>=1.14.3,<2.5.0', # Update pyproject.toml as well. 'objsize>=0.6.1,<0.8.0', 'packaging>=22.0', 'pillow>=12.1.1,<13', 'pymongo>=3.8.0,<5.0.0', # ADK / OpenTelemetry need proto-plus>=1.26.1 (protobuf>=5); that # floor is declared on ml_adk_dependency only so core installs stay # compatible with older proto-plus. 'proto-plus>=1.7.1,<2', # 1. Use a tighter upper bound in protobuf dependency to make sure # the minor version at job submission # does not exceed the minor version at runtime. # To avoid depending on an old dependency, update the minor version on # every Beam release, see: https://github.com/apache/beam/issues/25590 # 2. Allow latest protobuf 3 version as a courtesy to some customers. # # 3. Exclude protobuf 4 versions that leak memory, see: # https://github.com/apache/beam/issues/28246 'protobuf>=3.20.3,<7.0.0.dev0,!=4.0.*,!=4.21.*,!=4.22.0,!=4.23.*,!=4.24.*', # pylint: disable=line-too-long 'python-dateutil>=2.8.0,<3', 'pytz>=2018.3', 'requests>=2.32.4,<3.0.0', 'sortedcontainers>=2.4.0', 'typing-extensions>=3.7.0', 'zstandard>=0.18.0,<1', 'pyyaml>=3.12,<7.0.0', 'beartype>=0.21.0,<0.23.0', # Dynamic dependencies must be specified in a separate list, otherwise # Dependabot won't be able to parse the main list. Any dynamic # dependencies will not receive updates from Dependabot. ] + pyarrow_dependency, python_requires=python_requires, # BEAM-8840: Do NOT use tests_require or setup_requires. extras_require={ 'dev': [ 'isort==7.0.0', 'pyrefly==0.54.0', 'ruff==0.15.7', 'yapf==0.43.0', ], 'dill': [ # Dill doesn't have forwards-compatibility guarantees within minor # version. Pickles created with a new version of dill may not # unpickle using older version of dill. It is best to use the same # version of dill on client and server, therefore list of allowed # versions is very narrow. # See: https://github.com/uqfoundation/dill/issues/341. 'dill>=0.3.1.1,<0.3.2', ], 'docs': [ 'jinja2>=3.0,<3.2', 'Sphinx>=7.0.0,<8.0', 'docstring-parser>=0.15,<1.0', 'docutils>=0.18.1', 'markdown', 'pandas<2.3.0', 'openai', 'virtualenv-clone>=0.5,<1.0', ], 'test': [ 'cloud-sql-python-connector[pg8000]>=1.0.0,<2.0.0', 'docstring-parser>=0.15,<1.0', 'freezegun>=0.3.12', 'jinja2>=3.0,<3.2', 'joblib>=1.0.1', 'mock>=1.0.1,<6.0.0', 'pandas<2.3.0', 'parameterized>=0.7.1,<0.10.0', 'pydot>=1.2.0,<2', 'pyhamcrest>=1.9,!=1.10.0,<3.0.0', 'requests_mock>=1.7,<2.0', # google-adk 1.28+ requires tenacity>=9,<10 (conflicts with <9). 'tenacity>=8.0.0,<10', 'pytest>=7.1.2,<10.0', 'pytest-xdist>=2.5.0,<4', 'pytest-timeout>=2.1.0,<3', 'scikit-learn>=0.20.0,<1.8.0', 'sqlalchemy>=1.3,<3.0', 'psycopg2-binary>=2.8.5,<3.0', 'testcontainers[mysql,kafka,milvus]>=4.0.0,<5.0.0', 'cryptography>=41.0.2', # TODO(https://github.com/apache/beam/issues/36951): need to # further investigate the cause 'hypothesis>5.0.0,<6.148.4', 'virtualenv-clone>=0.5,<1.0', 'python-tds>=1.16.1', 'sqlalchemy-pytds>=1.0.2', 'pg8000>=1.31.5', "PyMySQL>=1.1.0", 'oracledb>=3.1.1' ], 'gcp': [ 'cachetools>=3.1.0,<7', 'google-api-core>=2.0.0,<3', 'google-apitools>=0.5.31,<0.5.32; python_version < "3.13"', 'google-apitools>=0.5.35; python_version >= "3.13"', # NOTE: Maintainers, please do not require google-auth>=2.x.x # Until this issue is closed # https://github.com/googleapis/google-cloud-python/issues/10566 'google-auth>=1.18.0,<3', 'google-auth-httplib2>=0.1.0,<0.3.0', 'google-cloud-datastore>=2.0.0,<3', 'google-cloud-pubsub>=2.1.0,<3', 'google-cloud-storage>=2.18.2,<4', # GCP packages required by tests 'google-cloud-bigquery>=2.0.0,<4', 'google-cloud-bigquery-storage>=2.6.3,<3', 'google-cloud-core>=2.0.0,<3', 'google-cloud-bigtable>=2.19.0,<3', 'google-cloud-build>=3.35.0,<4', 'google-cloud-spanner>=3.0.0,<4', # GCP Packages required by ML functionality 'google-cloud-dlp>=3.0.0,<4', 'google-cloud-kms>=3.0.0,<4', 'google-cloud-language>=2.0,<3', 'google-cloud-secret-manager>=2.0,<3', 'google-cloud-videointelligence>=2.0,<3', 'google-cloud-vision>=2,<4', 'google-cloud-recommendations-ai>=0.1.0,<0.11.0', 'google-cloud-aiplatform>=1.26.0, < 2.0', 'cloud-sql-python-connector>=1.18.2,<2.0.0', 'python-tds>=1.16.1', 'pg8000>=1.31.5', "PyMySQL>=1.1.0", # Authentication for Google Artifact Registry when using # --extra-index-url or --index-url in requirements.txt in # Dataflow, which allows installing python packages from private # Python repositories in GAR. 'keyrings.google-artifactregistry-auth', 'orjson>=3.9.7,<4', 'regex>=2020.6.8', ], 'interactive': [ 'facets-overview>=1.1.0,<2', 'google-cloud-dataproc>=5.0.0,<6', 'ipython>=7,<9', 'ipykernel>=6,<7', 'ipywidgets>=8,<9', # Skip version 6.1.13 due to # https://github.com/jupyter/jupyter_client/issues/637 'jupyter-client>=6.1.11,!=6.1.13,<8.2.1', 'pydot>=1.2.0,<2', 'timeloop>=1.0.2,<2', 'nbformat>=5.0.5,<6', 'nbconvert>=6.2.0,<8', ] + dataframe_dependency, 'interactive_test': [ # headless chrome based integration tests 'needle>=0.5.0,<1', 'chromedriver-binary>=117,<118', # use a fixed major version of PIL for different python versions 'pillow>=7.1.1,<10', # urllib 2.x is a breaking change for the headless chrome tests 'urllib3<2,>=1.21.1' ], # Optional dependencies to unit-test ML functionality. # We don't expect users to install this extra. Users should install # necessary dependencies individually, or we should create targeted # extras. Keeping the bounds open as much as possible so that we # can find out early when Beam doesn't work with new versions. 'ml_test': [ 'datatable', # tensorflow-transform requires dill, but doesn't set dill as a # hard requirement in setup.py. 'dill', # match tft extra. 'tensorflow_transform>=1.14.0,<1.15.0', # TFT->TFX-BSL require pandas 1.x, which is not compatible # with numpy 2.x 'numpy<2', # Comment out xgboost as it is breaking presubmit python ml # tests due to tag check introduced since pip 24.2 # https://github.com/apache/beam/issues/31285 # 'xgboost<2.0', # https://github.com/apache/beam/issues/31252 # tft needs protobuf<5; tf2onnx 1.17+ allows protobuf 5 on the # ADK-only path. 'tf2onnx>=1.16.1,<1.17', ] + ml_base_core, 'p310_ml_test': [ 'datatable', ] + ml_base, 'p312_ml_test': [ 'datatable', ] + ml_base, # maintainer: milvus tests only run with this extension. Make sure it # is covered by docker-in-docker test when changing py version 'p313_ml_test': ml_base + milvus_dependency, 'aws': ['boto3>=1.9,<2'], 'azure': [ 'azure-storage-blob>=12.3.2,<13', 'azure-core>=1.7.0,<2', 'azure-identity>=1.12.0,<2', ], 'dataframe': dataframe_dependency, 'dask': [ 'distributed >= 2024.4.2', 'dask >= 2024.4.2', # For development, 'distributed >= 2023.12.1' should work with # the above dask PR, however it can't be installed as part of # a single `pip` call, since distributed releases are pinned to # specific dask releases. As a workaround, distributed can be # installed first, and then `.[dask]` installed second, with the # `--update` / `-U` flag to replace the dask release brought in # by distributed. ], 'hadoop': ['hdfs>=2.1.0,<3.0.0'], 'yaml': [ 'docstring-parser>=0.15,<1.0', 'jinja2>=3.0,<3.2', 'virtualenv-clone>=0.5,<1.0', # https://github.com/PiotrDabkowski/Js2Py/issues/317 'js2py>=0.74,<1; python_version<"3.12"', 'jsonschema>=4.0.0,<5.0.0', ] + dataframe_dependency, # Keep the following dependencies in line with what we test against # in https://github.com/apache/beam/blob/master/sdks/python/tox.ini # For more info, see # https://docs.google.com/document/d/1c84Gc-cZRCfrU8f7kWGsNR2o8oSRjCM-dGHO9KvPWPw/edit?usp=sharing 'torch': ['torch>=1.9.0,<2.8.0'], 'tensorflow': [ 'tensorflow>=2.12rc1,<2.21', # tensorflow transitive dep 'absl-py>=0.12.0' ], 'transformers': [ 'transformers>=4.28.0,<4.56.0', 'tensorflow>=2.12.0', 'torch>=1.9.0' ], 'ml_cpu': [ 'tensorflow>=2.12.0', 'torch==2.8.0+cpu', 'transformers>=4.28.0,<4.56.0', # tensorflow transient dep 'absl-py>=0.12.0' ], 'redis': ['redis>=5.0.0,<6'], 'tft': [ 'tensorflow_transform>=1.14.0,<1.15.0', # TFT->TFX-BSL require pandas 1.x, which is not compatible # with numpy 2.x 'numpy<2', # tensorflow-transform requires dill, but doesn't set dill as a # hard requirement in setup.py. 'dill' ], 'tfrecord': ['crcmod>=1.7,<2.0'], 'onnx': [ 'onnxruntime==1.13.1', 'torch==1.13.1', 'tensorflow==2.11.0', 'tf2onnx==1.13.0', 'skl2onnx==1.13', 'transformers==4.25.1', # tensorflow transient dep 'absl-py>=0.12.0' ], 'xgboost': ['xgboost>=1.6.0,<2.1.3', 'datatable==1.0.0'], 'tensorflow-hub': ['tensorflow-hub>=0.14.0,<0.16.0'], 'milvus': milvus_dependency, 'vllm': ['openai==1.107.1', 'vllm==0.10.1.1', 'triton==3.3.1'] }, zip_safe=False, # PyPI package information. classifiers=[ 'Intended Audience :: End Users/Desktop', 'License :: OSI Approved :: Apache Software License', 'Operating System :: POSIX :: Linux', 'Programming Language :: Python :: 3.10', 'Programming Language :: Python :: 3.11', 'Programming Language :: Python :: 3.12', 'Programming Language :: Python :: 3.13', 'Programming Language :: Python :: 3.14', # When updating version classifiers, also update version warnings # above and in apache_beam/__init__.py. 'Topic :: Software Development :: Libraries', 'Topic :: Software Development :: Libraries :: Python Modules', ], license='Apache License, Version 2.0', keywords=PACKAGE_KEYWORDS, cmdclass={ 'pyrefly': pyrefly, }, )